{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "dedc3576", "metadata": {}, "outputs": [], "source": [ "from dustbi_simulator import *\n", "from Functions import *" ] }, { "cell_type": "code", "execution_count": 2, "id": "084facff", "metadata": {}, "outputs": [], "source": [ "from astropy.cosmology import Planck18\n", "import astropy.units as u" ] }, { "cell_type": "code", "execution_count": 3, "id": "677cac15", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/var/folders/2p/hm6bd5n17d5g5kpsm6s7vf3r0002g2/T/ipykernel_5171/1437736737.py:14: PerformanceWarning: DataFrame is highly fragmented. This is usually the result of calling `frame.insert` many times, which has poor performance. Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n", " dfdata['MU'] = Planck18.distmod(dfdata.zHD.values).value\n", "/var/folders/2p/hm6bd5n17d5g5kpsm6s7vf3r0002g2/T/ipykernel_5171/1437736737.py:19: PerformanceWarning: DataFrame is highly fragmented. This is usually the result of calling `frame.insert` many times, which has poor performance. Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n", " dfdata['SIM_EBV'] = dfdata.SIM_AV/dfdata.SIM_RV\n" ] } ], "source": [ "import numpy as np\n", "import pandas as pd\n", "df = pd.read_csv(\"INPUT_DES5YR_D2D.FITRES\", comment=\"#\", sep='\\s+')\n", "\n", "df['SIM_EBV'] = df.SIM_AV/df.SIM_RV\n", "\n", "df['MU'] = Planck18.distmod(df.zHD.values).value\n", "\n", "\n", "\n", "dfdata = pd.read_csv(\"SIMS_FOR_TESTING/FITOPT000.FITRES.gz\", \n", " comment=\"#\", sep=r'\\s+')\n", "\n", "dfdata['MU'] = Planck18.distmod(dfdata.zHD.values).value\n", "\n", "#dfdata = pd.read_csv(\"../INVERSE_H0/D5YR_DATA/FITOPT000_MUOPT000.FITRES.gz\", comment=\"#\", sep=r'\\s+')\n", "\n", "try:\n", " dfdata['SIM_EBV'] = dfdata.SIM_AV/dfdata.SIM_RV\n", "except:\n", " print(\"eh.\")\n", "\n", "dfdata = dfdata.loc[dfdata.IDSURVEY == 10]\n", "dfdata = dfdata.loc[dfdata.PROB_SNNV19 >= 0.5]\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "221231f5", "metadata": {}, "outputs": [], "source": [ "bounds_dict = {\n", " \"SIM_c\" : (-0.5, 0.5),\n", " \"SIM_RV\" : (1.5, 5),\n", " \"SIM_EBV\" : (0,1),\n", " \"SIM_beta\": (0.5,4),\n", "}\n", "\n", "function_dict = {\n", " \"SIM_c\" : DistGaussian,\n", " \"SIM_RV\" : DistGaussian,\n", " \"SIM_EBV\" : DistExponential,\n", " \"SIM_beta\": DistGaussian,\n", "}\n", "\n", "split_dict = {\n", "# \"SIM_RV\":[\"HOST_LOGMASS\", 10],\n", " \"SIM_EBV\":['HOST_LOGMASS', 10],\n", "# 'SIM_c':['HOST_LOGMASS', 10]\n", "}\n", "\n", "\n", "#Prior dict is a weird one; it should be a tuple for each parameter and then a boolean statement.\n", "\n", "split_dict = {}\n", "\n", "\n", "priors_dict = {\n", " \n", " \"SIM_c\" : [(-0.2, 0), (0.0, 0.1), ],\n", " \"SIM_RV\" : [(1.5,4), (0.2,2), ],\n", " \"SIM_EBV\" : [(0.05, 0.3)],\n", " \"SIM_beta\": [(0,3), (0,1), ],\n", " \n", "}\n", "\n", "priors_dict = {\n", " 'SIM_c': [(-0.071, -0.069), (0.052, 0.053)],\n", " 'SIM_RV': [(1.5,4), (0.5,2)],\n", " 'SIM_EBV': [(0.138, 0.142)],\n", " 'SIM_beta': [(1,2), (0.1, 0.3)],\n", "\n", " \n", "}\n", "\n", "latex_dict = {\n", " \n", " 'DistExponential':[r'$\\tau$'],\n", " 'DistGaussian':[r'$\\mu$', r'$\\sigma$'],\n", " 'SIM_c':r\"$c_{\\rm int}$\",\n", " 'SIM_beta':r\"$\\beta_{\\rm int}$\",\n", " 'SIM_RV':r\"$R_V$\",\n", " 'SIM_EBV':r\"$EBV$\",\n", " \n", "}\n", "\n", "\n", "dicts = [bounds_dict, function_dict, split_dict, priors_dict]" ] }, { "cell_type": "code", "execution_count": null, "id": "38ca1051", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 5, "id": "3be81940", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Total priors added: 7\n", "[0] \n", "[1] \n", "[2] \n", "[3] \n", "[4] \n", "[5] \n", "[6] \n" ] } ], "source": [ "param_names = ['SIM_c', 'SIM_RV', 'SIM_beta', 'SIM_EBV']\n", "#param_names = ['SIM_RV']\n", "\n", "\n", "params_to_fit = parameter_generation(param_names, dicts)\n", "priors = prior_generator(param_names, dicts)" ] }, { "cell_type": "code", "execution_count": 6, "id": "04ee03af", "metadata": {}, "outputs": [], "source": [ "layout = build_layout(params_to_fit, dicts)" ] }, { "cell_type": "code", "execution_count": null, "id": "1334dbdc", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 7, "id": "9d707a5e", "metadata": {}, "outputs": [], "source": [ "parameters_to_condition_on = ['c', 'mB', 'x1', 'zHD', 'cERR', 'mBERR', 'x1ERR', 'MU']" ] }, { "cell_type": "code", "execution_count": 8, "id": "89250d87-83d2-4ce5-9e0d-106f1b5fce69", "metadata": {}, "outputs": [], "source": [ "df, dfdata = standardise_data(df, dfdata, parameters_to_condition_on)" ] }, { "cell_type": "code", "execution_count": 9, "id": "7c32dc0c", "metadata": {}, "outputs": [], "source": [ "simulatinator = make_simulator(layout, df, param_names, parameters_to_condition_on, dicts, dfdata, is_split=True)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6077b6c9", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 10, "id": "58111230", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "8\n" ] } ], "source": [ "ndim = len(parameters_to_condition_on)#+1\n", "\n", "if any(p in split_dict for p in param_names): #check early to see if we need to split anything. \n", " ndim *= 2\n", " \n", "print(ndim)" ] }, { "cell_type": "code", "execution_count": null, "id": "0e9266b8", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "1c4f0db6", "metadata": {}, "outputs": [], "source": [ "from joblib import Parallel, delayed" ] }, { "cell_type": "code", "execution_count": 11, "id": "54700df8", "metadata": {}, "outputs": [], "source": [ "def batched_simulator(theta_batch):\n", " return torch.stack([simulatinator(theta) for theta in theta_batch])" ] }, { "cell_type": "code", "execution_count": null, "id": "ecdc4e89", "metadata": {}, "outputs": [], "source": [ "def batched_simulator(theta_batch):\n", " results = Parallel(n_jobs=-1)(\n", " delayed(simulatinator)(theta)\n", " for theta in theta_batch\n", " )\n", " return torch.stack(results)" ] }, { "cell_type": "code", "execution_count": 12, "id": "fb92e743", "metadata": {}, "outputs": [], "source": [ "from sbi import analysis as analysis\n", "\n", "# sbi\n", "from sbi import utils as utils\n", "from sbi.inference import NPE, simulate_for_sbi\n", "from sbi.utils.user_input_checks import (\n", " check_sbi_inputs,\n", " process_prior,\n", " process_simulator,\n", ")" ] }, { "cell_type": "code", "execution_count": 13, "id": "04ec48cd", "metadata": {}, "outputs": [], "source": [ "# Check prior, simulator, consistency\n", "prior, num_parameters, prior_returns_numpy = process_prior(priors)\n", "simulation_wrapper = process_simulator(simulatinator, prior, prior_returns_numpy)\n", "check_sbi_inputs(simulation_wrapper, prior)" ] }, { "cell_type": "code", "execution_count": 14, "id": "63690751", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": 15, "id": "502d0f30", "metadata": {}, "outputs": [], "source": [ "from sbi.inference import SNPE\n", "from sbi.utils import MultipleIndependent\n", "\n", "from sbi.neural_nets import posterior_nn\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "f178d092", "metadata": {}, "source": [ "# Potentially Upgraded Version" ] }, { "cell_type": "code", "execution_count": 16, "id": "b23e2355", "metadata": {}, "outputs": [], "source": [ "from sbi import analysis as analysis\n", "from sbi.inference import SNPE\n", "from sbi.neural_nets import posterior_nn\n", "\n", "class PopulationEmbeddingFull(nn.Module):\n", " def __init__(self, input_dim=ndim, hidden_dim=64, output_dim=32):\n", " super().__init__()\n", " self.phi = nn.Sequential(\n", " nn.Linear(input_dim, hidden_dim),\n", " nn.ReLU(), #nn.Tanh()\n", " nn.Linear(hidden_dim, hidden_dim),\n", " nn.ReLU()\n", " )\n", " self.attention = nn.Linear(hidden_dim, 1)\n", " self.rho = nn.Sequential(\n", " nn.Linear(hidden_dim, hidden_dim),\n", " nn.ReLU(),\n", " nn.Linear(hidden_dim, output_dim)\n", " )\n", "\n", " def forward(self, x):\n", " h = self.phi(x) # (batch, N, hidden)\n", " w = torch.softmax(self.attention(h), dim=1) # May need to train this to only run on errors? \n", " # Right now runs on everything ... \n", " h = (h * w).sum(dim=1) # (batch, hidden)\n", " return self.rho(h)\n", " \n", "#might need to standardise errors and signal " ] }, { "cell_type": "code", "execution_count": 17, "id": "99d834cd", "metadata": {}, "outputs": [], "source": [ "from sbi.inference import SNPE\n", "from sbi.utils import MultipleIndependent\n", "\n", "from sbi.neural_nets import posterior_nn\n", "\n", "density_estimator = posterior_nn(\n", " model=\"nsf\", #switch to nsf if interested \n", " embedding_net=PopulationEmbeddingFull(input_dim=ndim)\n", ")\n", "\n", "inference = SNPE(\n", " prior=priors,\n", " density_estimator=density_estimator, \n", ")\n", "\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "aaf09d29", "metadata": {}, "source": [ "# Tanh Population Embedding" ] }, { "cell_type": "code", "execution_count": null, "id": "8bbd7b6e", "metadata": {}, "outputs": [], "source": [ "from sbi import analysis as analysis\n", "from sbi.inference import SNPE\n", "from sbi.neural_nets import posterior_nn\n", "\n", "class PopulationEmbeddingFull(nn.Module):\n", " def __init__(self, input_dim=ndim, hidden_dim=64, output_dim=32):\n", " super().__init__()\n", " self.phi = nn.Sequential(\n", " nn.Linear(input_dim, hidden_dim),\n", " nn.Tanh(), \n", " nn.Linear(hidden_dim, hidden_dim),\n", " nn.Tanh()\n", " )\n", " self.attention = nn.Linear(hidden_dim, 1)\n", " self.rho = nn.Sequential(\n", " nn.Linear(hidden_dim, hidden_dim),\n", " nn.Tanh(),\n", " nn.Linear(hidden_dim, output_dim)\n", " )\n", "\n", " def forward(self, x):\n", " h = self.phi(x) # (batch, N, hidden)\n", " w = torch.softmax(self.attention(h), dim=1) # May need to train this to only run on errors? \n", " # Right now runs on everything ... \n", " h = (h * w).sum(dim=1) # (batch, hidden)\n", " return self.rho(h)\n", " \n", "#might need to standardise errors and signal " ] }, { "cell_type": "code", "execution_count": null, "id": "684483ba", "metadata": {}, "outputs": [], "source": [ "from sbi.inference import SNPE\n", "from sbi.utils import MultipleIndependent\n", "\n", "from sbi.neural_nets import posterior_nn\n", "\n", "density_estimator = posterior_nn(\n", " model=\"nsf\", #switch to nsf if interested \n", " embedding_net=PopulationEmbeddingFull(input_dim=ndim)\n", ")\n", "\n", "inference = SNPE(\n", " prior=priors,\n", " density_estimator=density_estimator, \n", ")\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "93e2c393", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "89de3a01", "metadata": { "scrolled": true }, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "df484a29", "metadata": { "scrolled": true }, "outputs": [], "source": [ "import torch\n", "import os\n", "import time\n", "\n", "\n", "start_time = time.perf_counter()\n", "\n", "batch_size = 50\n", "num_simulations = 4000\n", "save_path = \"simulations_v2.pt\"\n", "\n", "# If the file already exists, start fresh\n", "if os.path.exists(save_path):\n", " os.remove(save_path)\n", "\n", "for start in range(0, num_simulations, batch_size):\n", " current_bs = min(batch_size, num_simulations - start)\n", "\n", " # Sample and simulate\n", " theta_batch = priors.sample((current_bs,))\n", " x_batch = batched_simulator(theta_batch)\n", "\n", " # Append to SBI inference\n", " inference.append_simulations(theta_batch, x_batch)\n", "\n", " # Save incrementally\n", " if start == 0:\n", " # First batch, create the file\n", " torch.save({'theta': theta_batch, 'x': x_batch}, save_path)\n", " else:\n", " # Load existing data\n", " data = torch.load(save_path)\n", " data['theta'] = torch.cat([data['theta'], theta_batch], dim=0)\n", " data['x'] = torch.cat([data['x'], x_batch], dim=0)\n", " torch.save(data, save_path)\n", "\n", " print(f\"Appended {start + current_bs}/{num_simulations} simulations and saved incrementally.\")\n", "\n", "print(f\"All simulations saved incrementally to '{save_path}'\")\n", "\n", "end = time.perf_counter()\n", "\n", "elapsed = end - start_time\n", "#print(f'Time taken: {elapsed:.6f} seconds')" ] }, { "cell_type": "code", "execution_count": null, "id": "f293c5e6", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "aadd26b1", "metadata": {}, "outputs": [], "source": [ "print(f'Time taken: {elapsed/60:.6f} minutes')" ] }, { "cell_type": "code", "execution_count": null, "id": "f811c12e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " Training neural network. Epochs trained: 3" ] } ], "source": [ "inference.append_simulations(theta_batch, x_batch)\n", "\n", "density_estimator = inference.train(validation_fraction=0.1,\n", " show_train_summary=True)\n", "#force_first_round_loss \n", "#when true, only compute standard NPE loss;\n", "#think about moving to the S of SNPE eventually... \n", "\n", "print(\"\\n inferred successfully\")\n", "\n", "posterior = inference.build_posterior(density_estimator, sample_with=\"mcmc\")\n", "\n", "torch.save(posterior, \"posterior.v2.pt\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ef048cd2", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "a83953b0", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 25, "id": "9fd5df40", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/var/folders/2p/hm6bd5n17d5g5kpsm6s7vf3r0002g2/T/ipykernel_5171/3416222729.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " data = torch.load(\"simulations_v2.pt\")\n" ] } ], "source": [ "data = torch.load(\"simulations_v2.pt\")\n", "theta_batch = data[\"theta\"]\n", "x_batch = data[\"x\"]\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ef060131-1c6c-4459-ab23-3de911021c66", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "3bbf70d5-7bc1-455f-9e34-716744b313f5", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "b232ab55-c1d7-4215-ad7e-21876769264a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 18, "id": "ef06b310", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/var/folders/2p/hm6bd5n17d5g5kpsm6s7vf3r0002g2/T/ipykernel_5171/175813564.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " posterior = torch.load(\"posterior.v2.pt\", )\n" ] } ], "source": [ "posterior = torch.load(\"posterior.v2.pt\", )" ] }, { "cell_type": "code", "execution_count": null, "id": "2d0e5e34", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "25bfcb0a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "e57d3562", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 19, "id": "937d74d4", "metadata": {}, "outputs": [], "source": [ "def preprocess_data(param_names, parameters_to_condition_on, split_dict, dfdata, ):\n", " \n", " #salt_mcmc = start_distance()\n", " \n", " output_distribution = preprocess_input_distribution(\n", " dfdata, parameters_to_condition_on+['x0', 'x0ERR', 'MU']\n", " )\n", " \n", " #salt_mcmc.run(\n", " # output_distribution['x0'],\n", " # output_distribution['x0ERR'],\n", " # output_distribution['x1'],\n", " # output_distribution['x1ERR'],\n", " # output_distribution['c'],\n", " # output_distribution['cERR'],\n", " # output_distribution['MU']\n", " # )\n", "\n", " #MURES = add_distance(salt_mcmc, output_distribution)\n", " #output_distribution['MURES'] = MURES\n", " \n", " is_split = False\n", " if any(p in split_dict for p in param_names):\n", " is_split = True#check early to see if we need to split anything. \n", " \n", " if is_split:\n", " \n", " matching = [p for p in param_names if p in split_dict]\n", " name = matching[0]\n", "\n", " split_param = split_dict[name][0]\n", " split_val = split_dict[name][1]\n", "\n", " split_tensor = torch.tensor(\n", " dfdata[split_param].to_numpy(),\n", " dtype=torch.float32,\n", " )\n", "\n", " x = split_outputs(\n", " output_distribution,\n", " split_tensor,\n", " split_val,\n", " parameters_to_condition_on#+['MURES']\n", " )\n", "\n", " else:\n", " x = torch.stack(\n", " [output_distribution[p] for p in parameters_to_condition_on],#+['MURES']],\n", " dim=-1\n", " )\n", " \n", " return x " ] }, { "cell_type": "code", "execution_count": 20, "id": "3acc80e7", "metadata": {}, "outputs": [], "source": [ "x = preprocess_data(param_names, parameters_to_condition_on, split_dict, dfdata)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6fdf79cb", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 21, "id": "b493a610", "metadata": {}, "outputs": [], "source": [ "true_vals = priors.sample()\n", "\n", "new_x = simulatinator(true_vals)" ] }, { "cell_type": "code", "execution_count": null, "id": "0714964a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 22, "id": "f4ff0cfc", "metadata": {}, "outputs": [], "source": [ "labels = unspool_labels(param_names, dicts, latex_dict, function_dict)" ] }, { "cell_type": "code", "execution_count": 23, "id": "e0c6f977", "metadata": { "scrolled": true }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "778eac44f16941528a5ed0277f825cbd", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/1000 [00:00 \u001b[39m\u001b[32m1\u001b[39m posterior_samples = \u001b[43mposterior\u001b[49m\u001b[43m.\u001b[49m\u001b[43msample\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m1000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m=\u001b[49m\u001b[43mnew_x\u001b[49m\u001b[43m)\u001b[49m\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/sbi/inference/posteriors/direct_posterior.py:180\u001b[39m, in \u001b[36mDirectPosterior.sample\u001b[39m\u001b[34m(self, sample_shape, x, max_sampling_batch_size, sample_with, show_progress_bars)\u001b[39m\n\u001b[32m 173\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m sample_with \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m 174\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[32m 175\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mYou set `sample_with=\u001b[39m\u001b[38;5;132;01m{\u001b[39;00msample_with\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m`. As of sbi v0.18.0, setting \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 176\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33m`sample_with` is no longer supported. You have to rerun \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 177\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33m`.build_posterior(sample_with=\u001b[39m\u001b[38;5;132;01m{\u001b[39;00msample_with\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m).`\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 178\u001b[39m )\n\u001b[32m--> \u001b[39m\u001b[32m180\u001b[39m samples = \u001b[43mrejection\u001b[49m\u001b[43m.\u001b[49m\u001b[43maccept_reject_sample\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 181\u001b[39m \u001b[43m \u001b[49m\u001b[43mproposal\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mposterior_estimator\u001b[49m\u001b[43m.\u001b[49m\u001b[43msample\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 182\u001b[39m \u001b[43m \u001b[49m\u001b[43maccept_reject_fn\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtheta\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mwithin_support\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mprior\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtheta\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 183\u001b[39m \u001b[43m \u001b[49m\u001b[43mnum_samples\u001b[49m\u001b[43m=\u001b[49m\u001b[43mnum_samples\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 184\u001b[39m \u001b[43m \u001b[49m\u001b[43mshow_progress_bars\u001b[49m\u001b[43m=\u001b[49m\u001b[43mshow_progress_bars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 185\u001b[39m \u001b[43m \u001b[49m\u001b[43mmax_sampling_batch_size\u001b[49m\u001b[43m=\u001b[49m\u001b[43mmax_sampling_batch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 186\u001b[39m \u001b[43m \u001b[49m\u001b[43mproposal_sampling_kwargs\u001b[49m\u001b[43m=\u001b[49m\u001b[43m{\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mcondition\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 187\u001b[39m \u001b[43m \u001b[49m\u001b[43malternative_method\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mbuild_posterior(..., sample_with=\u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[33;43mmcmc\u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[33;43m)\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m 188\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m[\u001b[32m0\u001b[39m] \u001b[38;5;66;03m# [0] to return only samples, not acceptance probabilities.\u001b[39;00m\n\u001b[32m 190\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m samples[:, \u001b[32m0\u001b[39m]\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/torch/utils/_contextlib.py:116\u001b[39m, in \u001b[36mcontext_decorator..decorate_context\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m 113\u001b[39m \u001b[38;5;129m@functools\u001b[39m.wraps(func)\n\u001b[32m 114\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mdecorate_context\u001b[39m(*args, **kwargs):\n\u001b[32m 115\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[32m--> \u001b[39m\u001b[32m116\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/sbi/samplers/rejection/rejection.py:288\u001b[39m, in \u001b[36maccept_reject_sample\u001b[39m\u001b[34m(proposal, accept_reject_fn, num_samples, num_xos, show_progress_bars, warn_acceptance, sample_for_correction_factor, max_sampling_batch_size, proposal_sampling_kwargs, alternative_method, **kwargs)\u001b[39m\n\u001b[32m 284\u001b[39m num_samples_possible = \u001b[32m0\u001b[39m\n\u001b[32m 286\u001b[39m \u001b[38;5;28;01mwhile\u001b[39;00m num_remaining > \u001b[32m0\u001b[39m:\n\u001b[32m 287\u001b[39m \u001b[38;5;66;03m# Sample and reject.\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m288\u001b[39m candidates = \u001b[43mproposal\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 289\u001b[39m \u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43msampling_batch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore\u001b[39;49;00m\n\u001b[32m 290\u001b[39m \u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mproposal_sampling_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 291\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 292\u001b[39m \u001b[38;5;66;03m# SNPE-style rejection-sampling when the proposal is the neural net.\u001b[39;00m\n\u001b[32m 293\u001b[39m are_accepted = accept_reject_fn(candidates)\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/sbi/neural_nets/estimators/nflows_flow.py:137\u001b[39m, in \u001b[36mNFlowsFlow.sample\u001b[39m\u001b[34m(self, sample_shape, condition)\u001b[39m\n\u001b[32m 134\u001b[39m condition_batch_dim = condition.shape[\u001b[32m0\u001b[39m]\n\u001b[32m 135\u001b[39m num_samples = torch.Size(sample_shape).numel()\n\u001b[32m--> \u001b[39m\u001b[32m137\u001b[39m samples = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mnet\u001b[49m\u001b[43m.\u001b[49m\u001b[43msample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnum_samples\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m=\u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 138\u001b[39m \u001b[38;5;66;03m# Change from Nflows' convention of (batch_dim, sample_dim, *event_shape) to\u001b[39;00m\n\u001b[32m 139\u001b[39m \u001b[38;5;66;03m# (sample_dim, batch_dim, *event_shape) (PyTorch + SBI).\u001b[39;00m\n\u001b[32m 140\u001b[39m samples = samples.transpose(\u001b[32m0\u001b[39m, \u001b[32m1\u001b[39m)\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/nflows/distributions/base.py:65\u001b[39m, in \u001b[36mDistribution.sample\u001b[39m\u001b[34m(self, num_samples, context, batch_size)\u001b[39m\n\u001b[32m 62\u001b[39m context = torch.as_tensor(context)\n\u001b[32m 64\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m batch_size \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m---> \u001b[39m\u001b[32m65\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_sample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnum_samples\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 67\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m 68\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m check.is_positive_int(batch_size):\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/nflows/flows/base.py:54\u001b[39m, in \u001b[36mFlow._sample\u001b[39m\u001b[34m(self, num_samples, context)\u001b[39m\n\u001b[32m 49\u001b[39m noise = torchutils.merge_leading_dims(noise, num_dims=\u001b[32m2\u001b[39m)\n\u001b[32m 50\u001b[39m embedded_context = torchutils.repeat_rows(\n\u001b[32m 51\u001b[39m embedded_context, num_reps=num_samples\n\u001b[32m 52\u001b[39m )\n\u001b[32m---> \u001b[39m\u001b[32m54\u001b[39m samples, _ = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_transform\u001b[49m\u001b[43m.\u001b[49m\u001b[43minverse\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnoise\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m=\u001b[49m\u001b[43membedded_context\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 56\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m embedded_context \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m 57\u001b[39m \u001b[38;5;66;03m# Split the context dimension from sample dimension.\u001b[39;00m\n\u001b[32m 58\u001b[39m samples = torchutils.split_leading_dim(samples, shape=[-\u001b[32m1\u001b[39m, num_samples])\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/nflows/transforms/base.py:60\u001b[39m, in \u001b[36mCompositeTransform.inverse\u001b[39m\u001b[34m(self, inputs, context)\u001b[39m\n\u001b[32m 58\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34minverse\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs, context=\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[32m 59\u001b[39m funcs = (transform.inverse \u001b[38;5;28;01mfor\u001b[39;00m transform \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m._transforms[::-\u001b[32m1\u001b[39m])\n\u001b[32m---> \u001b[39m\u001b[32m60\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_cascade\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfuncs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m)\u001b[49m\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/nflows/transforms/base.py:50\u001b[39m, in \u001b[36mCompositeTransform._cascade\u001b[39m\u001b[34m(inputs, funcs, context)\u001b[39m\n\u001b[32m 48\u001b[39m total_logabsdet = inputs.new_zeros(batch_size)\n\u001b[32m 49\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m func \u001b[38;5;129;01min\u001b[39;00m funcs:\n\u001b[32m---> \u001b[39m\u001b[32m50\u001b[39m outputs, logabsdet = \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 51\u001b[39m total_logabsdet += logabsdet\n\u001b[32m 52\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m outputs, total_logabsdet\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/nflows/transforms/coupling.py:118\u001b[39m, in \u001b[36mCouplingTransform.inverse\u001b[39m\u001b[34m(self, inputs, context)\u001b[39m\n\u001b[32m 113\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.unconditional_transform \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m 114\u001b[39m identity_split, logabsdet = \u001b[38;5;28mself\u001b[39m.unconditional_transform.inverse(\n\u001b[32m 115\u001b[39m identity_split, context\n\u001b[32m 116\u001b[39m )\n\u001b[32m--> \u001b[39m\u001b[32m118\u001b[39m transform_params = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mtransform_net\u001b[49m\u001b[43m(\u001b[49m\u001b[43midentity_split\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 119\u001b[39m transform_split, logabsdet_split = \u001b[38;5;28mself\u001b[39m._coupling_transform_inverse(\n\u001b[32m 120\u001b[39m inputs=transform_split, transform_params=transform_params\n\u001b[32m 121\u001b[39m )\n\u001b[32m 122\u001b[39m logabsdet += logabsdet_split\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/torch/nn/modules/module.py:1736\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 1734\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs) \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m 1735\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1736\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/torch/nn/modules/module.py:1747\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 1742\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m 1743\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m 1744\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m 1745\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m 1746\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1747\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 1749\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m 1750\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n", "\u001b[36mFile \u001b[39m\u001b[32m~/anaconda3/envs/SBI/lib/python3.11/site-packages/nflows/nn/nets/resnet.py:96\u001b[39m, in \u001b[36mResidualNet.forward\u001b[39m\u001b[34m(self, inputs, context)\u001b[39m\n\u001b[32m 94\u001b[39m temps = \u001b[38;5;28mself\u001b[39m.initial_layer(inputs)\n\u001b[32m 95\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m---> \u001b[39m\u001b[32m96\u001b[39m temps = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43minitial_layer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcat\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 97\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m block \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.blocks:\n\u001b[32m 98\u001b[39m temps = block(temps, context=context)\n", "\u001b[31mKeyboardInterrupt\u001b[39m: " ] } ], "source": [ "posterior_samples = posterior.sample((1000,), x=new_x)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "066fc7c0", "metadata": { "scrolled": true }, "outputs": [], "source": [ "fig, axes = analysis.pairplot(\n", " posterior_samples,\n", " labels=labels\n", "\n", ");" ] }, { "cell_type": "code", "execution_count": null, "id": "ee74790d", "metadata": {}, "outputs": [], "source": [ "priors.sample()" ] }, { "cell_type": "code", "execution_count": null, "id": "6f7711ac", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "4ca0f5bb", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "ca81edfb", "metadata": {}, "outputs": [], "source": [ "theta_hat = posterior_samples.mean(0)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f3987daf", "metadata": { "scrolled": true }, "outputs": [], "source": [ "theta_hat\n" ] }, { "cell_type": "code", "execution_count": null, "id": "da86542a", "metadata": {}, "outputs": [], "source": [ "posterior_samples.std(0)" ] }, { "cell_type": "code", "execution_count": null, "id": "3955ff57", "metadata": {}, "outputs": [], "source": [ "true_params = torch.tensor([-0.07, 0.053, 2, 0.95, 2.07, 0.22, 0.15, 0.12,])" ] }, { "cell_type": "code", "execution_count": null, "id": "e63b7aee", "metadata": {}, "outputs": [], "source": [ "from IPython.display import display, Math\n" ] }, { "cell_type": "code", "execution_count": null, "id": "27837a17", "metadata": {}, "outputs": [], "source": [ "for n in range(len(theta_hat)):\n", " \n", "\n", " delta = theta_hat[n] - true_params[n]\n", " sigma = delta/posterior_samples.std(0)[n]\n", " \n", " string = rf\"{labels[n]} = {theta_hat[n]:.3f} +/- {posterior_samples.std(0)[n]:.3f} which is {sigma:.3f} $\\sigma$\"\n", " \n", " display(Math(string))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "a0385aaf", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "2bb98263", "metadata": {}, "outputs": [], "source": [ "simulatinator = make_simulator(layout, df, param_names, parameters_to_condition_on, dicts, \n", " dfdata, is_split=True, debug=True)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "dc93b5f3", "metadata": {}, "outputs": [], "source": [ "#dft = simulatinator(theta_hat)\n", "\n", "dft = simulatinator(true_vals)" ] }, { "cell_type": "code", "execution_count": null, "id": "2afaa467", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt" ] }, { "cell_type": "code", "execution_count": null, "id": "40484790", "metadata": {}, "outputs": [], "source": [ "bins = np.linspace(-0.4, 0.4, 20)\n", "\n", "#plt.hist(dft.loc[dft.HOST_LOGMASS < 10].c.values, histtype='step', bins=bins, label=\"sim output, low mass\", density=True)\n", "#plt.hist(dft.loc[dft.HOST_LOGMASS > 10].c.values, histtype='step', bins=bins, label=\"sim output, high mass\", density=True)\n", "plt.hist(dft.c.values, histtype='step', bins=bins, label=\"sim output, all mass\", density=True)\n", "\n", "plt.hist(dfdata.c.values, histtype='step', bins=bins, label=\"data\", density=True)\n", "\n", "plt.legend()\n", "plt.xlabel(\"c\")" ] }, { "cell_type": "code", "execution_count": null, "id": "49d6737e", "metadata": {}, "outputs": [], "source": [ "bins = np.linspace(-4, 10, 20)\n", "\n", "#plt.hist(dft.loc[dft.HOST_LOGMASS < 10].c.values, histtype='step', bins=bins, label=\"sim output, low mass\", density=True)\n", "#plt.hist(dft.loc[dft.HOST_LOGMASS > 10].c.values, histtype='step', bins=bins, label=\"sim output, high mass\", density=True)\n", "plt.hist(dft.MURES.values, histtype='step', bins=bins, label=\"sim output, all mass\", density=True)\n", "\n", "plt.hist(dfdata.MURES.values, histtype='step', bins=bins, label=\"data\", density=True)\n", "\n", "plt.legend()\n", "plt.xlabel(\"MURES\")" ] }, { "cell_type": "code", "execution_count": null, "id": "ede6e6e3", "metadata": {}, "outputs": [], "source": [ "bins = np.linspace(18, 26, 20)\n", "\n", "plt.hist(dft.mB.values, histtype='step', bins=bins, label=\"sim output\", density=True)\n", "plt.hist(dfdata.mB.values, histtype='step', bins=bins, label=\"data\", density=True)\n", "\n", "plt.legend()\n", "plt.xlabel(\"mB\")" ] }, { "cell_type": "code", "execution_count": null, "id": "159335c8", "metadata": {}, "outputs": [], "source": [ "bins = np.linspace(0, 0.6, 20)\n", "\n", "plt.hist(dft.loc[dft.HOST_LOGMASS < 10].SIM_EBV.values, histtype='step', bins=bins, label=\"low mass output\", density=True)\n", "plt.hist(dft.loc[dft.HOST_LOGMASS > 10].SIM_EBV.values, histtype='step', bins=bins, label=\"high mass output\", density=True)\n", "\n", "\n", "plt.hist(dfdata.SIM_EBV.values, histtype='step', bins=bins, label=\"data\", density=True)\n", "\n", "\n", "plt.legend()\n", "plt.xlabel(\"E(B-V)\")" ] }, { "cell_type": "code", "execution_count": null, "id": "2fbe921b", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "a35099c0", "metadata": {}, "source": [ "# Calibrate some posteriors" ] }, { "cell_type": "code", "execution_count": null, "id": "7e411c8e", "metadata": { "scrolled": true }, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "7864a530", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt" ] }, { "cell_type": "code", "execution_count": null, "id": "7c9c9ee2", "metadata": {}, "outputs": [], "source": [ "from sbi.diagnostics import run_sbc\n", "from sbi.analysis.plot import sbc_rank_plot\n" ] }, { "cell_type": "code", "execution_count": null, "id": "7a347fd1", "metadata": {}, "outputs": [], "source": [ "# Obtain your `posterior_estimator` with NPE, NLE, NRE.\n", "#posterior = inference.build_posterior()\n", "\n", "num_sbc_samples = 200 # choose a number of sbc runs, should be ~100s\n", "prior_samples = prior.sample((num_sbc_samples,))\n", "prior_predictives = batched_simulator(prior_samples)\n", "\n", "num_posterior_samples = 4000" ] }, { "cell_type": "code", "execution_count": null, "id": "671a3203", "metadata": { "scrolled": true }, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "08417a58", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "79117cb6", "metadata": {}, "outputs": [], "source": [ "ranks, dap_samples = run_sbc(\n", " prior_samples,\n", " prior_predictives,\n", " posterior,\n", " num_posterior_samples=num_posterior_samples,\n", " use_batched_sampling=True, # `True` can give speed-ups, but can cause memory issues.\n", ")\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d9d9b45b", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "a3860f99", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "e622be2f", "metadata": {}, "outputs": [], "source": [ "\n", "fig, ax = sbc_rank_plot(\n", " ranks,\n", " num_posterior_samples,\n", " plot_type=\"cdf\",\n", " num_bins=20,\n", " figsize=(5, 3),\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "0ce150d7", "metadata": { "scrolled": true }, "outputs": [], "source": [ "f, ax = sbc_rank_plot(\n", " ranks=ranks,\n", " num_posterior_samples=num_posterior_samples,\n", " plot_type=\"hist\",\n", " num_bins=None, # by passing None we use a heuristic for the number of bins.\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "219abaa8", "metadata": {}, "outputs": [], "source": [ "labels" ] }, { "cell_type": "code", "execution_count": null, "id": "10f55b7a", "metadata": {}, "outputs": [], "source": [ "#Flat histogram → well-calibrated.\n", "\n", "#U-shaped → posteriors too narrow.\n", "\n", "#Bell-shaped → posteriors too wide." ] }, { "cell_type": "code", "execution_count": null, "id": "485e4253", "metadata": {}, "outputs": [], "source": [ "num_tarp_samples = 200 # choose a number of sbc runs, should be ~100s\n", "# generate ground truth parameters and corresponding simulated observations for SBC.\n", "thetas = prior.sample((num_tarp_samples,))\n", "xs = batched_simulator(thetas)" ] }, { "cell_type": "code", "execution_count": null, "id": "14497208", "metadata": {}, "outputs": [], "source": [ "from sbi.diagnostics import check_sbc, check_tarp, run_sbc, run_tarp\n", "# the tarp method returns the ECP values for a given set of alpha coverage levels.\n", "ecp, alpha = run_tarp(\n", " thetas,\n", " xs,\n", " posterior,\n", " references=None, # will be calculated automatically.\n", " num_posterior_samples=3000,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "efe4ad78", "metadata": {}, "outputs": [], "source": [ "atc, ks_pval = check_tarp(ecp, alpha)\n", "print(atc, \"Should be close to 0\")\n", "print(ks_pval, \"Should be larger than 0.05\")" ] }, { "cell_type": "code", "execution_count": null, "id": "2b3c0027", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "00154335", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "4b4f7c5a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "a7f40c8c", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "fb813e18", "metadata": {}, "outputs": [], "source": [ "from sbi.analysis.plot import plot_tarp\n", "\n", "plot_tarp(ecp, alpha)" ] }, { "cell_type": "code", "execution_count": null, "id": "79c1b464", "metadata": {}, "outputs": [], "source": [ "def plot_tarp(\n", " ecp, alpha, title,):\n", " \"\"\"\n", " Plots the expected coverage probability (ECP) against the credibility\n", " level,alpha, for a given alpha grid.\n", "\n", " Args:\n", " ecp : numpy.ndarray\n", " Array of expected coverage probabilities.\n", " alpha : numpy.ndarray\n", " Array of credibility levels.\n", " title : str, optional\n", " Title for the plot. The default is \"\".\n", "\n", " Returns\n", " fig : matplotlib.figure.Figure\n", " The figure object.\n", " ax : matplotlib.axes.Axes\n", " The axes object.\n", "\n", " \"\"\"\n", "\n", " fig = plt.figure(figsize=(6, 6))\n", " ax: Axes = plt.gca()\n", "\n", " ax.plot(alpha, ecp, color=\"blue\", label=\"TARP\")\n", " ax.plot(alpha, alpha, color=\"black\", linestyle=\"--\", label=\"ideal\")\n", " ax.set_xlabel(r\"Credibility Level $\\alpha$\")\n", " ax.set_ylabel(r\"Expected Coverage Probability\")\n", " ax.set_xlim(0.0, 1.0)\n", " ax.set_ylim(0.0, 1.0)\n", " ax.set_title(title or \"\")\n", " ax.legend()\n", " return fig, ax # type: ignore\n" ] }, { "cell_type": "code", "execution_count": null, "id": "13225d2e", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "7590371c", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "8cc73dd8", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "7e94b632", "metadata": {}, "outputs": [], "source": [ "losses = inference.summary[\"training_loss\"]\n", "val_losses = inference.summary[\"validation_loss\"]\n", "\n", "# Plot\n", "plt.plot(losses, label='Training Loss')\n", "plt.plot(val_losses, label=\"Validation Loss\")\n", "plt.xlabel(\"Epoch\")\n", "plt.ylabel(\"Loss\")\n", "#plt.ylim([-10,10])\n", "plt.title(\"Loss vs Epochs\")\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "fd775038", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "972ac514", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "b852dc49", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "d870fb56", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "9c7e85ee", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "3acbf612", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "1728e5b8", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "b096ba14", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "a1e3490e", "metadata": {}, "outputs": [], "source": [ "plt.hist(df.SIM_c.values)" ] }, { "cell_type": "code", "execution_count": null, "id": "271c0fd3", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "3cd40bb9", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "a0031554", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "cd4b0ee7", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "0ed9f94d", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "4eeb61f5", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "6f2a00ec", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "00768533", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "164b0364", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "NERSC Python", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.11" } }, "nbformat": 4, "nbformat_minor": 5 }