SBI_NPE_2e6_h25_n10_galaxyspectra138_github.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eP_qNL7YazM0"
      },
      "outputs": [],
      "source": [
        "!pip install sbi --quiet\n",
        "!pip install deeplenstronomy --quiet\n",
        "!pip install lenstronomy --quiet\n",
        "!pip install deeplenstronomy --quiet\n",
        "!pip install getdist --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zUaoeRzFa9h_"
      },
      "outputs": [],
      "source": [
        "!pip uninstall matplotlib -y\n",
        "!pip install matplotlib==3.1.3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y49duZQWa9j5"
      },
      "outputs": [],
      "source": [
        "!pip install pandas"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "68rL7hXua9mE"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn \n",
        "import torch.nn.functional as F \n",
        "from sbi import utils, inference\n",
        "# from sbi import inference\n",
        "from sbi.inference import SNPE, simulate_for_sbi, prepare_for_sbi\n",
        "import numpy as np\n",
        "import scipy as sp\n",
        "\n",
        "import deeplenstronomy.deeplenstronomy as dl\n",
        "from deeplenstronomy.visualize import view_image\n",
        "from sbi.analysis import pairplot\n",
        "\n",
        "# import main simulation class of lenstronomy\n",
        "from lenstronomy.Util import util\n",
        "from lenstronomy.LightModel.light_model import LightModel\n",
        "from lenstronomy.LensModel.lens_model import LensModel\n",
        "from lenstronomy.Data.imaging_data import ImageData\n",
        "from lenstronomy.ImSim.image_model import ImageModel\n",
        "import lenstronomy.Util.image_util as image_util\n",
        "from lenstronomy.Data.psf import PSF\n",
        "\n",
        "import deeplenstronomy.deeplenstronomy as dl\n",
        "from deeplenstronomy.visualize import view_image\n",
        "import pickle\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-8VkYp6X2uBj"
      },
      "outputs": [],
      "source": [
        "!pip install astroML --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tBVx4iwO2w2l"
      },
      "outputs": [],
      "source": [
        "from matplotlib.font_manager import FontProperties\n",
        "from matplotlib import gridspec\n",
        "plt.rcParams.update({'xtick.major.pad': '7.0'})\n",
        "plt.rcParams.update({'xtick.major.size': '7.5'})\n",
        "plt.rcParams.update({'xtick.major.width': '1.5'})\n",
        "plt.rcParams.update({'xtick.minor.pad': '7.0'})\n",
        "plt.rcParams.update({'xtick.minor.size': '3.5'})\n",
        "plt.rcParams.update({'xtick.minor.width': '1.0'})\n",
        "plt.rcParams.update({'ytick.major.pad': '7.0'})\n",
        "plt.rcParams.update({'ytick.major.size': '7.5'})\n",
        "plt.rcParams.update({'ytick.major.width': '1.5'}) \n",
        "plt.rcParams.update({'ytick.minor.pad': '7.0'})\n",
        "plt.rcParams.update({'ytick.minor.size': '3.5'})\n",
        "plt.rcParams.update({'ytick.minor.width': '1.0'})\n",
        "plt.rcParams.update({'xtick.color': 'k'})\n",
        "plt.rcParams.update({'ytick.color': 'k'})\n",
        "plt.rcParams.update({'font.size': 20})\n",
        "plt.rcParams.update({'axes.linewidth':2})\n",
        "plt.rcParams.update({'patch.linewidth':2})\n",
        "#from astroML.plotting import setup_text_plots\n",
        "#setup_text_plots(fontsize=20, usetex=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kJHhqF4Ga9oR"
      },
      "outputs": [],
      "source": [
        "#fix seed\n",
        "torch.manual_seed(100)\n",
        "torch.cuda.manual_seed(100)\n",
        "torch.cuda.manual_seed_all(100)\n",
        "np.random.seed(100)\n",
        "#random.seed(100)\n",
        "torch.backends.cudnn.deterministic=True\n",
        "torch.backends.cudnn.enabled=False\n",
        "torch.backends.cudnn.benchmark = False\n",
        "#os.environ['PYTHONHASHSEED'] = str(100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V5xxz1-ta9qY"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x0Qq2YkIa9u0"
      },
      "outputs": [],
      "source": [
        "DATA_DIR = '/content/drive/MyDrive/data_sbi/ts5_10_s/'\n",
        "#PLOT_DIR = '../plots/ts1_100k/'\n",
        "\n",
        "num_sed = 10000\n",
        "\n",
        "df = pd.read_csv(DATA_DIR+'table.txt',header=None)\n",
        "mass = np.log10(df[1].values)\n",
        "met = df[4].values\n",
        "age = df[2].values\n",
        "tau = np.log10(df[3].values)\n",
        "dust = df[5].values\n",
        "\n",
        "data_j = []\n",
        "for e,element1 in enumerate(mass):\n",
        "    data_j.append(np.array([element1,age[e],tau[e],met[e],dust[e]]))\n",
        "    \n",
        "data_j = np.array(data_j)\n",
        "\n",
        "plt.hist(mass,bins=20)\n",
        "plt.xlabel('logM$_{*}$')\n",
        "plt.show()\n",
        "plt.hist(met,bins=20)\n",
        "plt.xlabel('log(Z/Z$_{\\odot})$')\n",
        "plt.show()\n",
        "plt.hist(tau,bins=20)\n",
        "plt.xlabel('$Tau$')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G2cAwUMsMgBx"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "start = time.time()\n",
        "sed_spline = pd.read_csv(DATA_DIR+'sed_resamp_all.txt',header=None,delimiter=' ').to_numpy()\n",
        "end = time.time()\n",
        "print(\"Time taken to load table:\", end - start)\n",
        "    \n",
        "#sed1 = np.log10(sed)\n",
        "sed1 = np.log10(sed_spline)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J9uE4epra91J"
      },
      "outputs": [],
      "source": [
        "unc_factor = np.log10(1.05)\n",
        "len_sed = len(sed1[2])\n",
        "\n",
        "def simulator2(parameter_set): #parameter_set = tensor --->\n",
        "\n",
        "    distances = np.linalg.norm(data_j - np.array(parameter_set), axis=1)\n",
        "    min_index = np.argmin(distances)\n",
        "    testsim = torch.tensor(np.random.normal(sed1[min_index],unc_factor,len_sed)) \n",
        "    return testsim\n",
        "   \n",
        "test_sim = simulator2(torch.tensor([8.8,  1.9,  0.15, -0.5,  0.7]))\n",
        "plt.plot(test_sim)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y9SkHNTua93Z"
      },
      "outputs": [],
      "source": [
        "from sbi.inference import SNPE, prepare_for_sbi, simulate_for_sbi\n",
        "import sbi.utils as utils\n",
        "from sbi.inference.base import infer\n",
        "from sbi.utils.get_nn_models import posterior_nn\n",
        "\n",
        "prior_min = [8,0.01,-1,-2.0,0.1]\n",
        "prior_max = [13,4,1,0.2,1.0]\n",
        "prior2 = utils.torchutils.BoxUniform(low=torch.as_tensor(prior_min), \n",
        "                                    high=torch.as_tensor(prior_max))\n",
        "\n",
        "newfac = 20\n",
        "num_sim=int(1e5) * newfac\n",
        "\n",
        "hfeat = 25\n",
        "ntrans = 10\n",
        "\n",
        "ndata = len(simulator2(torch.tensor([10.3,1.0,-0.1,0.0,0.3])))\n",
        "\n",
        "posterior_already_run = False\n",
        "#posterior_already_run = True\n",
        "\n",
        "unc_term = '5pct'\n",
        "pickle_filename = \"ts2_10k_%s_h%s_n%s_nsim%s_data%s.pkl\"%(unc_term,hfeat,ntrans,num_sim,ndata)\n",
        "\n",
        "if posterior_already_run:\n",
        "    with open(DATA_DIR + pickle_filename, 'rb') as handle:\n",
        "        posterior = pickle.load(handle)\n",
        "    \n",
        "    simulator, prior = prepare_for_sbi(simulator2, prior2)\n",
        "\n",
        "else:\n",
        "    simulator, prior = prepare_for_sbi(simulator2, prior2)\n",
        "    \n",
        "    density_estimator_build_fun = posterior_nn(model='maf', hidden_features=hfeat, num_transforms=ntrans) #model = maf or nsf\n",
        "    \n",
        "    inference = SNPE(prior=prior, density_estimator=density_estimator_build_fun)\n",
        "\n",
        "    theta, x = simulate_for_sbi(simulator, proposal=prior, num_simulations=num_sim)\n",
        "    inference = inference.append_simulations(theta, x)\n",
        "\n",
        "    start = time.time()\n",
        "    density_estimator = inference.train() #Train the neural density estimator\n",
        "    end = time.time()\n",
        "    print()\n",
        "    print(\"Time to train neural network:\", end - start)\n",
        "\n",
        "    posterior = inference.build_posterior(density_estimator) #Build posterior\n",
        "\n",
        "    with open(DATA_DIR + pickle_filename, \"wb\") as handle:\n",
        "        pickle.dump(posterior, handle)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iML9bfl3eh1j"
      },
      "source": [
        "# **Inference and Tests**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ua5IXfB5kG8e"
      },
      "outputs": [],
      "source": [
        "data_j[568]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IdBbv4m2a95r"
      },
      "outputs": [],
      "source": [
        "#OBSERVATIONS\n",
        "obs_par = torch.tensor([data_j[568]])\n",
        "\n",
        "\n",
        "def get_models(parameter_set):\n",
        "    distances = np.linalg.norm(data_j - np.array(parameter_set), axis=1)\n",
        "    ii = np.argmin(distances)\n",
        "    return torch.tensor(sed1[ii])\n",
        "\n",
        "obs1 = simulator2(obs_par) #with stochasticity\n",
        "observation = obs1\n",
        "print(len(observation))\n",
        "observation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DBExLhJypMmu"
      },
      "outputs": [],
      "source": [
        "from sbi import analysis as analysis\n",
        "from sbi.inference.base import infer\n",
        "\n",
        "samples1 = posterior.sample((10000,), x=observation)\n",
        "log_probability = posterior.log_prob(samples1, x=observation)\n",
        "\n",
        "_ = analysis.pairplot([samples1], figsize=(6,6),labels=['logM$_{*}$','age','tau','log(Z/Z$_{\\odot})$','dust'],points=obs_par,\n",
        "                     points_colors=\"red\")\n",
        "\n",
        "plt.savefig(DATA_DIR + 'posterior_5param.pdf',dpi=300)\n",
        "\n",
        "print(\"median logM:\",np.median(samples1[:,0]))\n",
        "print(\"16-50 percentile:\",np.percentile(samples1[:,0],16) - np.median(samples1[:,0]))\n",
        "print(\"84-50 percentile:\",np.percentile(samples1[:,0],84)- np.median(samples1[:,0]))\n",
        "print(\"median metallicity:\",np.median(samples1[:,3]))\n",
        "print(\"16-50 percentile:\",np.percentile(samples1[:,3],16) - np.median(samples1[:,3]))\n",
        "print(\"84-50 percentile:\",np.percentile(samples1[:,3],84)- np.median(samples1[:,3]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e_I3J2w4a-AB"
      },
      "outputs": [],
      "source": [
        "par1_rand = []\n",
        "par2_rand = []\n",
        "par3_rand = []\n",
        "par4_rand = []\n",
        "par5_rand = []\n",
        "\n",
        "ms1 = []\n",
        "ms2 = []\n",
        "ms3 = []\n",
        "ms4 = []\n",
        "ms5 = []\n",
        "\n",
        "us1 = []\n",
        "us2 = []\n",
        "us3 = []\n",
        "us4 = []\n",
        "us5 = []\n",
        "\n",
        "for t,elementt in enumerate(np.arange(500)):\n",
        "#for t,elementt in enumerate(np.arange(1000)):\n",
        "\n",
        "    mass_rand = np.random.uniform(8,13)\n",
        "    age_rand = np.random.uniform(0.01,4)\n",
        "    tau_rand = np.random.uniform(-1,1)\n",
        "    met_rand = np.random.uniform(-2.0,0.2)\n",
        "    dust_rand = np.random.uniform(0.1,1.0)\n",
        "    \n",
        "    samples = posterior.sample((100,), x=simulator2(torch.tensor([mass_rand,age_rand,tau_rand,met_rand,dust_rand])))\n",
        "    \n",
        "    median_mass = np.median(samples[:,0])\n",
        "    unc1 =  np.median(samples[:,0]) - np.percentile(samples[:,0],16)\n",
        "    unc2 = np.percentile(samples[:,0],84)- np.median(samples[:,0])\n",
        "    us1.append(np.sqrt(unc1**2 + unc2**2))\n",
        "\n",
        "    median_age = np.median(samples[:,1])\n",
        "    unc1 =  np.median(samples[:,1]) - np.percentile(samples[:,1],16)\n",
        "    unc2 = np.percentile(samples[:,1],84)- np.median(samples[:,1])\n",
        "    us2.append(np.sqrt(unc1**2 + unc2**2))\n",
        "\n",
        "    median_tau = np.median(samples[:,2])\n",
        "    unc1 =  np.median(samples[:,2]) - np.percentile(samples[:,2],16)\n",
        "    unc2 = np.percentile(samples[:,2],84)- np.median(samples[:,2])\n",
        "    us3.append(np.sqrt(unc1**2 + unc2**2))\n",
        "\n",
        "    median_met = np.median(samples[:,3])\n",
        "    unc1 =  np.median(samples[:,3]) - np.percentile(samples[:,3],16)\n",
        "    unc2 = np.percentile(samples[:,3],84)- np.median(samples[:,3])\n",
        "    us4.append(np.sqrt(unc1**2 + unc2**2))\n",
        "\n",
        "    median_dust = np.median(samples[:,4])\n",
        "    unc1 =  np.median(samples[:,4]) - np.percentile(samples[:,4],16)\n",
        "    unc2 = np.percentile(samples[:,4],84)- np.median(samples[:,4])\n",
        "    us5.append(np.sqrt(unc1**2 + unc2**2))\n",
        "    \n",
        "    par1_rand.append(mass_rand)\n",
        "    par2_rand.append(age_rand)\n",
        "    par3_rand.append(tau_rand)\n",
        "    par4_rand.append(met_rand)\n",
        "    par5_rand.append(dust_rand)\n",
        "    #print(\"median logM:\",median_scale)\n",
        "    ms1.append(median_mass)\n",
        "    ms2.append(median_age)\n",
        "    ms3.append(median_tau)\n",
        "    ms4.append(median_met)\n",
        "    ms5.append(median_dust)\n",
        "    \n",
        "ms1 = np.array(ms1)\n",
        "ms2 = np.array(ms2)\n",
        "ms3 = np.array(ms3)\n",
        "ms4 = np.array(ms4)\n",
        "ms5 = np.array(ms5)\n",
        "\n",
        "par1_rand = np.array(par1_rand)\n",
        "par2_rand = np.array(par2_rand)\n",
        "par3_rand = np.array(par3_rand)\n",
        "par4_rand = np.array(par4_rand)\n",
        "par5_rand = np.array(par5_rand)\n",
        "\n",
        "us1 = np.array(us1)\n",
        "us2 = np.array(us2)\n",
        "us3 = np.array(us3)\n",
        "us4 = np.array(us4)\n",
        "us5 = np.array(us5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8hSVXjAJa-CU"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(10,4))\n",
        "plt.errorbar(par1_rand,(ms1 - par1_rand)/us1,fmt='o',yerr = us1,markersize=2)\n",
        "plt.show()\n",
        "\n",
        "plt.hist((ms1 - par1_rand)/us1,bins=30)\n",
        "plt.show()\n",
        "######################################\n",
        "plt.figure(figsize=(10,4))\n",
        "plt.errorbar(par2_rand,(ms2 - par2_rand)/us2,fmt='o',yerr = us2,markersize=2)\n",
        "plt.show()\n",
        "\n",
        "plt.hist((ms2 - par2_rand)/us2,bins=30)\n",
        "plt.show()\n",
        "\n",
        "######################################\n",
        "plt.figure(figsize=(10,4))\n",
        "plt.errorbar(par4_rand,(ms4 - par4_rand)/us4,fmt='o',yerr = us4,markersize=2)\n",
        "plt.show()\n",
        "\n",
        "plt.hist((ms4 - par4_rand)/us4,bins=30)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q9Htc9yYa-Eb"
      },
      "outputs": [],
      "source": [
        "plt.scatter(par1_rand,ms1,s=50)\n",
        "plt.xlabel('True Value (logM)')\n",
        "plt.ylabel('Recovered Value (logM)')\n",
        "plt.plot(np.linspace(7,14,100),np.linspace(7,14,100),color='black',ls='--',alpha=0.75)\n",
        "plt.show()\n",
        "\n",
        "plt.scatter(par4_rand,ms4,s=50)\n",
        "plt.xlabel('True Value (Metallicity)')\n",
        "plt.ylabel('Recovered Value (Metallicity)')\n",
        "plt.plot(np.linspace(-2,0.2,100),np.linspace(-2,0.2,100),color='black',ls='--')\n",
        "plt.show()\n",
        "\n",
        "plt.scatter(par2_rand,ms2,s=50)\n",
        "plt.xlabel('True Value (Age)')\n",
        "plt.ylabel('Recovered Value (Age)')\n",
        "plt.plot(np.linspace(0.01,4,100),np.linspace(0.01,4,100),color='black',ls='--')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U7Ja_ldba-Gk"
      },
      "outputs": [],
      "source": [
        "#Posterior Predictive Checks (PPC)\n",
        "# A PPC is performed after we trained or neural posterior\n",
        "x_o = observation\n",
        "posterior.set_default_x(x_o)\n",
        "\n",
        "# We draw theta samples from the posterior. This part is not in the scope of SBI\n",
        "posterior_samples = posterior.sample((2000,))\n",
        "\n",
        "# We use posterior theta samples to generate x data\n",
        "x_pp = simulator(posterior_samples)\n",
        "\n",
        "#Reduce dimensionality of >4000 wavelength-SED\n",
        "x_pp_median = torch.median(x_pp,1,keepdim=True)[0]\n",
        "x_o_median = torch.median(x_o)\n",
        "\n",
        "# We verify if the observed data falls within the support of the generated data\n",
        "_ = analysis.pairplot(samples=x_pp_median, points=x_o_median,\n",
        "                      figsize=(6,6),\n",
        "                      points_colors=\"red\",\n",
        "                     labels=['Median(Splined SED)'])\n",
        "\n",
        "##################################################################################################################\n",
        "#Simulation-Based Calibration\n",
        "from torch import eye, ones, zeros\n",
        "from torch.distributions import MultivariateNormal\n",
        "\n",
        "from sbi.analysis import check_sbc,run_sbc, get_nltp, sbc_rank_plot\n",
        "from sbi.inference import SNPE, SNPE_C, prepare_for_sbi, simulate_for_sbi\n",
        "from sbi.simulators import linear_gaussian, diagonal_linear_gaussian\n",
        "\n",
        "num_sbc_runs = 1000 # choose a number of sbc runs, should be ~100s or ideally 1000\n",
        "# generate ground truth parameters and corresponding simulated observations for SBC.\n",
        "thetas = prior.sample((num_sbc_runs,))\n",
        "xs = simulator(thetas) \n",
        "\n",
        "# run SBC: for each inference we draw 1000 posterior samples.\n",
        "#num_posterior_samples = 1000\n",
        "num_posterior_samples = 100\n",
        "\n",
        "ranks, dap_samples = run_sbc(thetas, xs, posterior, num_posterior_samples=num_posterior_samples)\n",
        "check_stats = check_sbc(ranks, thetas, dap_samples, num_posterior_samples=num_posterior_samples)\n",
        "\n",
        "print(f\"kolmogorov-smirnov p-values \\ncheck_stats['ks_pvals'] = {check_stats['ks_pvals'].numpy()}\")\n",
        "print(f\"c2st accuracies \\ncheck_stats['c2st_ranks'] = {check_stats['c2st_ranks'].numpy()}\")\n",
        "print(f\"- c2st accuracies check_stats['c2st_dap'] = {check_stats['c2st_dap'].numpy()}\")\n",
        "\n",
        "f,ax=sbc_rank_plot(ranks=ranks,num_posterior_samples=num_posterior_samples,\n",
        "              plot_type=\"hist\",\n",
        "              num_bins=None)\n",
        "\n",
        "f, ax = sbc_rank_plot(ranks, 1000, plot_type = \"cdf\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cku3CmqDa-Io"
      },
      "outputs": [],
      "source": [
        "_ = analysis.pairplot(samples=x_pp[:,::18], \n",
        "                      points=x_o[::18],\n",
        "                      labels=['wave1','wave2','wave3','wave4','wave5','wave6','wave7','wave8'], \n",
        "                      points_colors=\"red\")\n",
        "\n",
        "#plt.savefig(DATA_DIR + 'bestfit_ppc_2.pdf',dpi=300)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}