{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os, sys\n",
"import numpy as np\n",
"import json\n",
"import matplotlib.pyplot as plt\n",
"import pandas as pd\n",
"from addict import Dict\n",
"import scipy.stats as stats\n",
"import corner\n",
"import lenstronomy.Util.param_util as param_util\n",
"from baobab import bnn_priors\n",
"from baobab.configs import BaobabConfig, tdlmc_diagonal_cosmo_config\n",
"%matplotlib inline\n",
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Visualizing the input prior PDF in the DiagonalCosmoBNNPrior and the resulting samples\n",
"__Author:__ Ji Won Park\n",
" \n",
"__Created:__ 11/20/19\n",
" \n",
"__Last run:__ 11/20/19\n",
"\n",
"__Goals:__\n",
"Plot the (marginal) distributions of the parameters sampled from the diagonal cosmology-aware BNN prior, in which parameters follow physically reasonable relations.\n",
"\n",
"__Before running this notebook:__\n",
"1. Generate some data. At the root of the `baobab` repo, run:\n",
"```\n",
"generate baobab/configs/tdlmc_diagonal_cosmo_config.py --n_data 1000\n",
"```\n",
"This generates 1000 samples using `DiagonalCosmoBNNPrior` at the current working directory (the repo root). \n",
"\n",
"2. The `generate` script you just ran also exported a log file in the end, as a json file, to the current working directory. The name follows the format `\"log_%m-%d-%Y_%H:%M_baobab.json\"` where the date and time are of those at which you ran the script. Modify `baobab_log_path` in the below cell to the correct log path."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"baobab_log_path = '/home/jwp/stage/sl/h0rton/log_12-10-2019_01:30_baobab.json'\n",
"with open(baobab_log_path, 'r') as f:\n",
" log_str = f.read()\n",
"cfg = Dict(json.loads(log_str))\n",
"meta = pd.read_csv(os.path.abspath(os.path.join(cfg.out_dir, 'metadata.csv')), index_col=None)\n",
"bnn_prior = getattr(bnn_priors, cfg.bnn_prior_class)(cfg.bnn_omega, cfg.components)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here are the parameters available. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"sorted(meta.columns.values)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Add shear and ellipticity modulus and angle\n",
"if 'external_shear_gamma_ext' in meta.columns.values:\n",
" gamma_ext = meta['external_shear_gamma_ext'].values\n",
" psi_ext = meta['external_shear_psi_ext'].values\n",
" gamma1, gamma2 = param_util.shear_polar2cartesian(phi=psi_ext, gamma=gamma_ext)\n",
" meta['external_shear_gamma1'] = gamma1\n",
" meta['external_shear_gamma2'] = gamma2\n",
"else:\n",
" gamma1 = meta['external_shear_gamma1'].values\n",
" gamma2 = meta['external_shear_gamma2'].values\n",
" psi_ext, gamma_ext = param_util.shear_cartesian2polar(gamma1, gamma2)\n",
" meta['external_shear_gamma_ext'] = gamma_ext\n",
" meta['external_shear_psi_ext'] = psi_ext\n",
"for comp in cfg.components:\n",
" if comp in ['lens_mass', 'src_light', 'lens_light']:\n",
" if '{:s}_e1'.format(comp) in meta.columns.values:\n",
" e1 = meta['{:s}_e1'.format(comp)].values\n",
" e2 = meta['{:s}_e2'.format(comp)].values\n",
" phi, q = param_util.ellipticity2phi_q(e1, e2)\n",
" meta['{:s}_q'.format(comp)] = q\n",
" meta['{:s}_phi'.format(comp)] = phi\n",
" else:\n",
" q = meta['{:s}_q'.format(comp)].values\n",
" phi = meta['{:s}_phi'.format(comp)].values\n",
" e1, e2 = param_util.phi_q2_ellipticity(phi, q)\n",
" meta['{:s}_e1'.format(comp)] = e1\n",
" meta['{:s}_e2'.format(comp)] = e2\n",
"meta['src_light_pos_offset_x'] = meta['src_light_center_x'] - meta['lens_mass_center_x']\n",
"meta['src_light_pos_offset_y'] = meta['src_light_center_y'] - meta['lens_mass_center_y']"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Add source gal positional offset\n",
"meta['src_pos_offset'] = np.sqrt(meta['src_light_center_x']**2.0 + meta['src_light_center_y']**2.0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def plot_prior_samples(eval_at, component, param, unit):\n",
" param_key = '{:s}_{:s}'.format(component, param)\n",
" if param_key == 'src_light_pos_offset_x':\n",
" hyperparams = cfg.bnn_omega['src_light']['center_x']\n",
" elif param_key == 'src_light_pos_offset_y':\n",
" hyperparams = cfg.bnn_omega['src_light']['center_y']\n",
" elif (param_key == 'src_light_center_x') or (param_key == 'src_light_center_y'):\n",
" raise NotImplementedError(\"Use `plot_derived_quantities` instead.\")\n",
" elif (component, param) in bnn_prior.params_to_exclude:\n",
" raise NotImplementedError(\"This parameter wasn't sampled independently. Please use `plot_derived_quantities` instead.\")\n",
" else:\n",
" hyperparams = cfg.bnn_omega[component][param].copy()\n",
" pdf_eval = bnn_prior.eval_param_pdf(eval_at, hyperparams)\n",
" plt.plot(eval_at, pdf_eval, 'r-', lw=2, alpha=0.6, label='PDF')\n",
" binning = np.linspace(eval_at[0], eval_at[-1], 50)\n",
" plt.hist(meta[param_key], bins=binning, edgecolor='k', density=True, align='mid', label='sampled')\n",
" print(hyperparams)\n",
" plt.xlabel(\"{:s} ({:s})\".format(param_key, unit))\n",
" plt.ylabel(\"density\")\n",
" plt.legend()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def plot_derived_quantities(param_key, unit, binning=None):\n",
" binning = 30 if binning is None else binning\n",
" _ = plt.hist(meta[param_key], bins=binning, edgecolor='k', density=True, align='mid', label='sampled')\n",
" plt.xlabel(\"{:s} ({:s})\".format(param_key, unit))\n",
" plt.ylabel(\"density\")\n",
" plt.legend()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Lens mass params"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.7, 1.3, 30), 'lens_mass', 'theta_E', 'arcsec')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.04, 0.04, 100), 'lens_mass', 'center_x', 'arcsec')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.04, 0.04, 100), 'lens_mass', 'center_y', 'arcsec')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(1.7, 2.2, 30), 'lens_mass', 'gamma', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.5, 1.0, 30), 'lens_mass', 'q', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.5*np.pi, 0.5*np.pi, 30), 'lens_mass', 'phi', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('lens_mass_e1', 'dimensionless', 50)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('lens_mass_e2', 'dimensionless', 50)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## External shear params"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0, 0.1, 100), 'external_shear', 'gamma_ext', 'no unit')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.5*np.pi, 0.5*np.pi, 30), 'external_shear', 'psi_ext', 'rad')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('external_shear_gamma1', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('external_shear_gamma2', 'dimensionless')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Lens light params"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(15.0, 20.0, 30), 'lens_light', 'magnitude', 'mag')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(1, 6, 100), 'lens_light', 'n_sersic', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.2, 1.5, 30), 'lens_light', 'R_sersic', 'arcsec')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.3, 1.2, 30), 'lens_light', 'q', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.5*np.pi, 0.5*np.pi, 30), 'lens_light', 'phi', 'rad')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('lens_light_e1', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('lens_light_e2', 'rad')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Source light params"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(15, 25, 100), 'src_light', 'magnitude', 'mag')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.5, 6.0, 100), 'src_light', 'n_sersic', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.1, 0.6, 100), 'src_light', 'R_sersic', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.1, 0.1, 100), 'src_light', 'pos_offset_x', 'arcsec')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.1, 0.1, 100), 'src_light', 'pos_offset_y', 'arcsec')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(0.2, 1.0, 100), 'src_light', 'q', 'dimensionless')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(-0.5*np.pi, 0.5*np.pi, 100), 'src_light', 'phi', 'rad')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('src_light_e1', 'dimensionless', 30)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('src_light_e2', 'dimensionless', 30)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## AGN light params"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_prior_samples(np.linspace(15, 25, 100), 'src_light', 'magnitude', 'mag')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Total magnification"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('total_magnification', 'dimensionless', binning=np.linspace(0, 300, 30))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Other quantities"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('z_lens', 'dimensionless', 20)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plot_derived_quantities('z_src', 'dimensionless', 20)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"_ = plt.hist(meta['z_src'] - meta['z_lens'], edgecolor='k', bins=30)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Pairwise distributions"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def plot_pairwise_dist(df, cols, fig=None):\n",
" n_params = len(cols)\n",
" plot = corner.corner(meta[cols],\n",
" color='tab:blue', \n",
" smooth=1.0, \n",
" labels=cols,\n",
" show_titles=True,\n",
" fill_contours=True,\n",
" levels=[0.68, 0.95, 0.997],\n",
" fig=fig,\n",
" range=[0.99]*n_params,\n",
" hist_kwargs=dict(density=True, ))\n",
" return plot"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cols = ['src_pos_offset', 'total_magnification',\n",
" 'external_shear_gamma_ext', 'external_shear_psi_ext',\n",
" 'lens_mass_q', 'lens_mass_theta_E',\n",
" 'src_light_q', 'src_light_R_sersic']\n",
"_ = plot_pairwise_dist(meta, cols)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cols = ['lens_mass_gamma', 'lens_light_n_sersic' ]\n",
"_ = plot_pairwise_dist(meta, cols)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python (baobab)",
"language": "python",
"name": "baobab"
},
"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.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}