{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\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_cov_config, gamma_cov_config\n",
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the input prior PDF in the CovBNNPrior and the resulting samples\n",
    "__Author:__ Ji Won Park\n",
    "    \n",
    "__Created:__ 8/30/19\n",
    "    \n",
    "__Last run:__ 9/05/19\n",
    "\n",
    "__Goals:__\n",
    "Plot the (marginal) distributions of the parameters sampled from the covariate BNN prior, in which parameters take a multivariate (log)normal distribution.\n",
    "\n",
    "__Before running this notebook:__\n",
    "Generate some data. At the root of the `baobab` repo, run:\n",
    "```\n",
    "generate baobab/configs/tdlmc_cov_config.py --n_data 1000\n",
    "```\n",
    "This generates 1000 samples using `CovBNNPrior` at the location this notebook expects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO add description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg_path = tdlmc_cov_config.__file__\n",
    "#cfg_path = os.path.join('..', '..', 'time_delay_lens_modeling_challenge', 'data', 'baobab_configs', 'train_tdlmc_diagonal_config.py')\n",
    "cfg = BaobabConfig.from_file(cfg_path)\n",
    "#out_data_dir = os.path.join('..', '..', 'time_delay_lens_modeling_challenge', cfg.out_dir)\n",
    "out_data_dir = os.path.join('..', cfg.out_dir)\n",
    "print(out_data_dir)\n",
    "meta = pd.read_csv(os.path.join(out_data_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": {},
   "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.phi_gamma_ellipticity(psi_ext, 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.ellipticity2phi_gamma(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"
   ]
  },
  {
   "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.5, 1.5, 100), '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_derived_quantities('lens_mass_gamma', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1.0, 1.0, 100), 'lens_mass', 'e1', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1.0, 1.0, 100), 'lens_mass', 'e2', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('lens_mass_q', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('lens_mass_phi', 'rad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## External shear params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(0, 1.0, 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, 100), '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_derived_quantities('lens_light_magnitude', 'mag')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('lens_light_n_sersic', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(0.0, 2.0, 100), 'lens_light', 'R_sersic', 'arcsec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1.0, 1.0, 100), 'lens_light', 'e1', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1.0, 1.0, 100), 'lens_light', 'e2', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('lens_light_q', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('lens_light_phi', 'rad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Source light params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('src_light_magnitude', 'mag')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(0.0, 6.0, 100), 'src_light', 'n_sersic', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(0.0, 2.0, 100), 'src_light', 'R_sersic', 'arcsec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1, 1, 100), 'src_light', 'pos_offset_x', 'arcsec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1, 1, 100), 'src_light', 'pos_offset_y', 'arcsec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('src_light_center_x', 'arcsec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('src_light_center_y', 'arcsec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1.0, 1.0, 100), 'src_light', 'e1', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_prior_samples(np.linspace(-1.0, 1.0, 100), 'src_light', 'e2', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('src_light_q', 'dimensionless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('src_light_phi', 'rad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AGN light params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_derived_quantities('agn_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": [
    "## 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', ]\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
}