sim_plots.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c4522984",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/anasofiauzsoy/Desktop/School_Stuff/Mandel_Group/BayeSN-numpyro\r\n"
     ]
    }
   ],
   "source": [
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "27081062",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.lines as mlines\n",
    "from astropy.cosmology import FlatLambdaCDM\n",
    "import pickle\n",
    "import numpy as np\n",
    "import corner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20909830",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Input \u001b[0;32mIn [3]\u001b[0m, in \u001b[0;36m<cell line: 47>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     48\u001b[0m dataset \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124msim_population_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(dataset_number) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m/\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(i)\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m (\u001b[38;5;28mopen\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mresults/\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m dataset \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_vi/chains.pkl\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrb\u001b[39m\u001b[38;5;124m\"\u001b[39m)) \u001b[38;5;28;01mas\u001b[39;00m openfile:\n\u001b[0;32m---> 50\u001b[0m \tvi_objects \u001b[38;5;241m=\u001b[39m \u001b[43mpickle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mopenfile\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m (\u001b[38;5;28mopen\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mresults/\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m dataset \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_mcmc/chains.pkl\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrb\u001b[39m\u001b[38;5;124m\"\u001b[39m)) \u001b[38;5;28;01mas\u001b[39;00m openfile:\n\u001b[1;32m     53\u001b[0m \tmcmc_objects \u001b[38;5;241m=\u001b[39m pickle\u001b[38;5;241m.\u001b[39mload(openfile)\n",
      "File \u001b[0;32m~/anaconda/envs/py310_clone/lib/python3.10/site-packages/jax/__init__.py:35\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m _cloud_tpu_init\n\u001b[1;32m     32\u001b[0m \u001b[38;5;66;03m# Confusingly there are two things named \"config\": the module and the class.\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[38;5;66;03m# We want the exported object to be the class, so we first import the module\u001b[39;00m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;66;03m# to make sure a later import doesn't overwrite the class.\u001b[39;00m\n\u001b[0;32m---> 35\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config \u001b[38;5;28;01mas\u001b[39;00m _config_module\n\u001b[1;32m     36\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m _config_module\n\u001b[1;32m     38\u001b[0m \u001b[38;5;66;03m# Force early import, allowing use of `jax.core` after importing `jax`.\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda/envs/py310_clone/lib/python3.10/site-packages/jax/config.py:17\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Copyright 2018 The JAX Authors.\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     14\u001b[0m \n\u001b[1;32m     15\u001b[0m \u001b[38;5;66;03m# TODO(phawkins): fix users of this alias and delete this file.\u001b[39;00m\n\u001b[0;32m---> 17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconfig\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda/envs/py310_clone/lib/python3.10/site-packages/jax/_src/config.py:27\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mthreading\u001b[39;00m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, Generic, NamedTuple, Optional, TypeVar\n\u001b[0;32m---> 27\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m lib\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m jax_jit\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m transfer_guard_lib\n",
      "File \u001b[0;32m~/anaconda/envs/py310_clone/lib/python3.10/site-packages/jax/_src/lib/__init__.py:84\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;66;03m# Before importing any C compiled modules from jaxlib, first import the CPU\u001b[39;00m\n\u001b[1;32m     81\u001b[0m \u001b[38;5;66;03m# feature guard module to verify that jaxlib was compiled in a way that only\u001b[39;00m\n\u001b[1;32m     82\u001b[0m \u001b[38;5;66;03m# uses instructions that are present on this machine.\u001b[39;00m\n\u001b[1;32m     83\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mjaxlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcpu_feature_guard\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mcpu_feature_guard\u001b[39;00m\n\u001b[0;32m---> 84\u001b[0m \u001b[43mcpu_feature_guard\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheck_cpu_features\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mjaxlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mutils\u001b[39;00m\n\u001b[1;32m     87\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mjaxlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mxla_client\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mxla_client\u001b[39;00m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source."
     ]
    }
   ],
   "source": [
    "mcmc_point_estimates = {'AV':[], 'mu':[], 'theta':[]}\n",
    "# vi_medians = {'AV':[], 'mu':[], 'theta':[]}\n",
    "vi_point_estimates = {'AV':[], 'mu':[], 'theta':[]}\n",
    "\n",
    "fiducial_cosmology={\"H0\": 73.24, \"Om0\": 0.28}\n",
    "cosmo = FlatLambdaCDM(**fiducial_cosmology)\n",
    "\n",
    "mu_sample_diffs = []\n",
    "av_samples = []\n",
    "av_vi_samples = []\n",
    "\n",
    "mu_samples = []\n",
    "ds_samples = []\n",
    "delm_samples = []\n",
    "\n",
    "ds_sample_diffs = []\n",
    "ds_mcmc_sample_diffs = []\n",
    "\n",
    "ds_means = []\n",
    "ds_errs = []\n",
    "\n",
    "num_to_plot = 200\n",
    "\n",
    "dataset_number=20\n",
    "\n",
    "title_str = \"\"\n",
    "\n",
    "if dataset_number == 6:\n",
    "\ttitle_str = \"epsilons off for sims, on for fitting\"\n",
    "if dataset_number == 5 or dataset_number == 7:\n",
    "\ttitle_str = \"no epsilons for sim/fitting\"\n",
    "if dataset_number == 4:\n",
    "\ttitle_str = \"using epsilons for both sims and fitting\"\n",
    "\n",
    "true_avs = np.loadtxt(\"sim_population_AV_\" + str(dataset_number) + \".txt\")[:num_to_plot]\n",
    "true_thetas = np.loadtxt(\"sim_population_theta_\" + str(dataset_number) + \".txt\")[:num_to_plot]\n",
    "true_z = np.loadtxt(\"sim_population_z_\" + str(dataset_number) + \".txt\")[:num_to_plot]\n",
    "true_mus = np.array([cosmo.distmod(z).value for z in true_z])\n",
    "\n",
    "def get_mode_from_samples(samples):\n",
    "\thist, bin_edges = np.histogram(samples, bins=50)\n",
    "\tmax_index = np.argmax(hist)\n",
    "\tmode = (bin_edges[max_index] + bin_edges[max_index + 1])/2\n",
    "\treturn mode\n",
    "\n",
    "\n",
    "for i in range(num_to_plot):\n",
    "\tdataset = 'sim_population_' + str(dataset_number) + '/' + str(i)\n",
    "\twith (open(\"results/\" + dataset + \"_vi/chains.pkl\", \"rb\")) as openfile:\n",
    "\t\tvi_objects = pickle.load(openfile)\n",
    "\n",
    "\twith (open(\"results/\" + dataset + \"_mcmc/chains.pkl\", \"rb\")) as openfile:\n",
    "\t\tmcmc_objects = pickle.load(openfile)\n",
    "\n",
    "\tvi_params = np.load(\"results/\" + dataset + \"_vi/vi_params.npz\")\n",
    "\n",
    "\n",
    "\tfor var in ['AV', 'mu', 'theta']:\n",
    "\t\tmcmc_samples = mcmc_objects[var].reshape((1000,))\n",
    "\t\tvi_samples = np.squeeze(vi_objects[var])\n",
    "\n",
    "\t\tif var == \"AV\": # Use sample mode for AV point estimate\n",
    "\t\t\tmcmc_point_estimates[var].append(get_mode_from_samples(mcmc_samples))\n",
    "\t\t\tvi_point_estimates[var].append(get_mode_from_samples(vi_samples)) # mu means mean here\n",
    "\t\t\tav_samples.append(mcmc_samples)\n",
    "\n",
    "\t\telif var == 'theta': # Use sample median for theta point estimate\n",
    "\t\t\tvi_point_estimates[var].append(np.median(vi_samples))\n",
    "\t\t\tmcmc_point_estimates[var].append(np.median(mcmc_samples))\n",
    "\n",
    "\t\telse: # Use sample median for mu point estimate\n",
    "\t\t\tvi_point_estimates[var].append(np.median(vi_samples))\n",
    "\t\t\tmcmc_point_estimates[var].append(np.median(mcmc_samples))\n",
    "\t\t\tmu_sample_diffs.append(vi_samples - mcmc_samples)\t\n",
    "\t\t\tds_sample_diffs.append(vi_params['mu'][-1][0] - np.median(vi_samples))\n",
    "\t\t\tds_mcmc_sample_diffs.append(vi_params['mu'][-1][0] - mcmc_samples)\n",
    "\t\t\tds_means.append(vi_params['mu'][-1][0])\n",
    "\t\t\tds_errs.append(vi_params['cov'][-1][-1])\n",
    "\n",
    "\tmu_samples.append(np.squeeze(vi_objects['mu']))\n",
    "\tds_samples.append(np.squeeze(vi_objects['Ds']))\n",
    "\tdelm_samples.append(np.squeeze(vi_objects['delM']))\n",
    "\tav_vi_samples.append(np.squeeze(vi_objects['AV']))\n",
    "\n",
    "mu_sample_diffs = np.array(mu_sample_diffs)\n",
    "ds_mcmc_sample_diffs = np.array(ds_mcmc_sample_diffs)\n",
    "delm_samples = np.array(delm_samples)\n",
    "av_vi_samples = np.array(av_vi_samples)\n",
    "\n",
    "ds_means = np.array(ds_means)\n",
    "ds_errs = np.array(ds_errs)\n",
    "\n",
    "av_samples = np.array(av_samples)\n",
    "\n",
    "for var in ['AV', 'mu', 'theta']:\n",
    "\tmcmc_point_estimates[var] = np.array(mcmc_point_estimates[var])\n",
    "\tvi_point_estimates[var] = np.array(vi_point_estimates[var])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e12e53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VI vs true subplots\n",
    "alpha = 0.2\n",
    "fig, ax = plt.subplots(2,3, figsize = (16,9))\n",
    "ax[0][0].plot(true_mus, vi_point_estimates['mu'], 'o', alpha=alpha)\n",
    "ax[0][0].plot(np.linspace(min(true_mus), max(true_mus)),np.linspace(min(true_mus), max(true_mus)))\n",
    "ax[0][0].set_ylabel('VI $\\\\mu$', fontsize = 16)\n",
    "# ax[0][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[1][0].plot(true_mus, vi_point_estimates['mu'] - true_mus, 'o', alpha=alpha)\n",
    "ax[1][0].axhline(0, color = 'k')\n",
    "ax[1][0].set_ylabel('Residual $\\\\mu$ (VI - true)', fontsize = 16)\n",
    "ax[1][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[0][1].plot(true_thetas, vi_point_estimates['theta'], 'o', alpha=alpha)\n",
    "ax[0][1].plot(np.linspace(min(true_thetas), max(true_thetas)),np.linspace(min(true_thetas), max(true_thetas)))\n",
    "ax[0][1].set_ylabel('VI $\\\\theta$', fontsize = 16)\n",
    "# ax[0][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[1][1].plot(true_thetas, vi_point_estimates['theta'] - true_thetas, 'o', alpha=alpha)\n",
    "ax[1][1].axhline(0, color = 'k')\n",
    "ax[1][1].set_ylabel('Residual $\\\\theta$ (VI - true)', fontsize = 16)\n",
    "ax[1][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[0][2].plot(true_avs, vi_point_estimates['AV'], 'o', alpha=alpha)\n",
    "ax[0][2].plot(np.linspace(0, 2.1),np.linspace(0, 2.1))\n",
    "ax[0][2].set_ylabel('VI $A_V$', fontsize = 16)\n",
    "ax[0][2].set_xlim(0,1)\n",
    "ax[0][2].set_ylim(0,1)\n",
    "\n",
    "# ax[0][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "\n",
    "ax[1][2].plot(true_avs, vi_point_estimates['AV'] - true_avs, 'o', label='VI', alpha=alpha)\n",
    "ax[1][2].axhline(0, color = 'k')\n",
    "ax[1][2].set_ylabel('Residual $A_V$ (VI - true)', fontsize = 16)\n",
    "ax[1][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "# ax[1][2].set_xlim(0,1)\n",
    "\n",
    "# for axis in ax.flatten():\n",
    "#   axis.tick_params(axis='x', labelsize=12)\n",
    "#   axis.tick_params(axis='y', labelsize=12)\n",
    "\n",
    "# plt.tight_layout()\n",
    "# plt.show()\n",
    "\n",
    "# MCMC vs true subplots\n",
    "# fig, ax = plt.subplots(2,3, figsize = (16,9))\n",
    "\n",
    "mcmc_color = 'r'\n",
    "ax[0][0].plot(true_mus, mcmc_point_estimates['mu'], 'o', c=mcmc_color, alpha=alpha)\n",
    "ax[0][0].plot(np.linspace(min(true_mus), max(true_mus)),np.linspace(min(true_mus), max(true_mus)))\n",
    "ax[0][0].set_ylabel('MCMC $\\\\mu$', fontsize = 16)\n",
    "# ax[0][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[1][0].plot(true_mus, mcmc_point_estimates['mu'] - true_mus, 'o', c=mcmc_color, alpha=alpha)\n",
    "ax[1][0].axhline(0, color = 'k')\n",
    "ax[1][0].set_ylabel('Residual $\\\\mu$ (MCMC - true)', fontsize = 16)\n",
    "ax[1][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[0][1].plot(true_thetas, mcmc_point_estimates['theta'], 'o', c=mcmc_color, alpha=alpha)\n",
    "ax[0][1].plot(np.linspace(min(true_thetas), max(true_thetas)),np.linspace(min(true_thetas), max(true_thetas)))\n",
    "ax[0][1].set_ylabel('MCMC $\\\\theta$', fontsize = 16)\n",
    "# ax[0][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[1][1].plot(true_thetas, mcmc_point_estimates['theta'] - true_thetas, 'o', c=mcmc_color, alpha=alpha)\n",
    "ax[1][1].axhline(0, color = 'k')\n",
    "ax[1][1].set_ylabel('Residual $\\\\theta$ (MCMC - true)', fontsize = 16)\n",
    "ax[1][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[0][2].plot(true_avs, mcmc_point_estimates['AV'], 'o', c=mcmc_color, alpha=alpha)\n",
    "ax[0][2].plot(np.linspace(0, 2.1),np.linspace(0, 2.1))\n",
    "ax[0][2].set_ylabel('MCMC $A_V$', fontsize = 16)\n",
    "ax[0][2].set_xlim(0,1)\n",
    "ax[0][2].set_ylim(0,1)\n",
    "\n",
    "# ax[0][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "\n",
    "ax[1][2].plot(true_avs, mcmc_point_estimates['AV'] - true_avs, 'o', c=mcmc_color, label = 'MCMC', alpha=alpha)\n",
    "ax[1][2].axhline(0, color = 'k')\n",
    "ax[1][2].set_ylabel('Residual $A_V$ (MCMC - true)', fontsize = 16)\n",
    "ax[1][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "# ax[1][2].set_xlim(0,1)\n",
    "\n",
    "for axis in ax.flatten():\n",
    "  axis.tick_params(axis='x', labelsize=12)\n",
    "  axis.tick_params(axis='y', labelsize=12)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# VI vs MCMC subplots\n",
    "fig, ax = plt.subplots(2,3, figsize = (16,9))\n",
    "ax[0][0].plot(true_mus, vi_point_estimates['mu'], 'o')\n",
    "ax[0][0].plot(np.linspace(min(true_mus), max(true_mus)),np.linspace(min(true_mus), max(true_mus)))\n",
    "ax[0][0].set_ylabel('VI $\\\\mu$', fontsize = 16)\n",
    "# ax[0][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[1][0].plot(true_mus, vi_point_estimates['mu'] - mcmc_point_estimates['mu'], 'o')\n",
    "ax[1][0].axhline(0, color = 'k')\n",
    "ax[1][0].set_ylabel('Residual $\\\\mu$ (VI - MCMC)', fontsize = 16)\n",
    "ax[1][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[0][1].plot(true_thetas, vi_point_estimates['theta'], 'o')\n",
    "ax[0][1].plot(np.linspace(min(mcmc_point_estimates['theta']), max(mcmc_point_estimates['theta'])),np.linspace(min(mcmc_point_estimates['theta']), max(mcmc_point_estimates['theta'])))\n",
    "ax[0][1].set_ylabel('VI $\\\\theta$', fontsize = 16)\n",
    "# ax[0][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[1][1].plot(true_thetas, vi_point_estimates['theta'] - mcmc_point_estimates['theta'], 'o')\n",
    "ax[1][1].axhline(0, color = 'k')\n",
    "ax[1][1].set_ylabel('Residual $\\\\theta$ (VI - MCMC)', fontsize = 16)\n",
    "ax[1][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[0][2].plot(true_avs, vi_point_estimates['AV'], 'o')\n",
    "ax[0][2].plot(np.linspace(0, 2.1),np.linspace(0, 2.1))\n",
    "ax[0][2].set_ylabel('VI $A_V$', fontsize = 16)\n",
    "# ax[0][2].set_xlim(0,1)\n",
    "# ax[0][2].set_ylim(0,1)\n",
    "\n",
    "# ax[0][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "\n",
    "ax[1][2].plot(true_avs, vi_point_estimates['AV'] - mcmc_point_estimates['AV'], 'o')\n",
    "ax[1][2].axhline(0, color = 'k')\n",
    "ax[1][2].set_ylabel('Residual $A_V$ (VI - MCMC)', fontsize = 16)\n",
    "ax[1][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "# ax[1][2].set_xlim(0,1)\n",
    "\n",
    "for axis in ax.flatten():\n",
    "  axis.tick_params(axis='x', labelsize=12)\n",
    "  axis.tick_params(axis='y', labelsize=12)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# VI vs MCMC subplots\n",
    "fig, ax = plt.subplots(2,3, figsize = (16,9))\n",
    "ax[0][0].plot(mcmc_point_estimates['mu'], vi_point_estimates['mu'], 'o')\n",
    "ax[0][0].plot(np.linspace(min(mcmc_point_estimates['mu']), max(mcmc_point_estimates['mu'])),np.linspace(min(mcmc_point_estimates['mu']), max(mcmc_point_estimates['mu'])))\n",
    "ax[0][0].set_ylabel('VI $\\\\mu$', fontsize = 16)\n",
    "# ax[0][0].set_xlabel('True $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[1][0].plot(mcmc_point_estimates['mu'], vi_point_estimates['mu'] - mcmc_point_estimates['mu'], 'o')\n",
    "ax[1][0].axhline(0, color = 'k')\n",
    "ax[1][0].set_ylabel('Residual $\\\\mu$ (VI - MCMC)', fontsize = 16)\n",
    "ax[1][0].set_xlabel('MCMC $\\\\mu$', fontsize = 16)\n",
    "\n",
    "\n",
    "ax[0][1].plot(mcmc_point_estimates['theta'], vi_point_estimates['theta'], 'o')\n",
    "ax[0][1].plot(np.linspace(min(mcmc_point_estimates['theta']), max(mcmc_point_estimates['theta'])),np.linspace(min(mcmc_point_estimates['theta']), max(mcmc_point_estimates['theta'])))\n",
    "ax[0][1].set_ylabel('VI $\\\\theta$', fontsize = 16)\n",
    "# ax[0][1].set_xlabel('True $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[1][1].plot(mcmc_point_estimates['theta'], vi_point_estimates['theta'] - mcmc_point_estimates['theta'], 'o')\n",
    "ax[1][1].axhline(0, color = 'k')\n",
    "ax[1][1].set_ylabel('Residual $\\\\theta$ (VI - MCMC)', fontsize = 16)\n",
    "ax[1][1].set_xlabel('MCMC $\\\\theta$', fontsize = 16)\n",
    "\n",
    "ax[0][2].plot(mcmc_point_estimates['AV'], vi_point_estimates['AV'], 'o')\n",
    "ax[0][2].plot(np.linspace(0, 2.1),np.linspace(0, 2.1))\n",
    "ax[0][2].set_ylabel('VI $A_V$', fontsize = 16)\n",
    "# ax[0][2].set_xlim(0,1)\n",
    "# ax[0][2].set_ylim(0,1)\n",
    "\n",
    "# ax[0][2].set_xlabel('True $A_V$', fontsize = 16)\n",
    "\n",
    "ax[1][2].plot(mcmc_point_estimates['AV'], vi_point_estimates['AV'] - mcmc_point_estimates['AV'], 'o')\n",
    "ax[1][2].axhline(0, color = 'k')\n",
    "ax[1][2].set_ylabel('Residual $A_V$ (VI - MCMC)', fontsize = 16)\n",
    "ax[1][2].set_xlabel('MCMC $A_V$', fontsize = 16)\n",
    "# ax[1][2].set_xlim(0,1)\n",
    "\n",
    "for axis in ax.flatten():\n",
    "  axis.tick_params(axis='x', labelsize=12)\n",
    "  axis.tick_params(axis='y', labelsize=12)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "\n",
    "plt.plot(mcmc_point_estimates['AV'], vi_point_estimates['mu'] - mcmc_point_estimates['mu'], 'o')\n",
    "median_residual_vi_mcmc = np.median(vi_point_estimates['mu'] - mcmc_point_estimates['mu'])\n",
    "plt.axhline(median_residual_vi_mcmc, color='k', label='median residual = ' + str(round(median_residual_vi_mcmc, 4)))\n",
    "plt.xlabel('MCMC $A_V$', fontsize = 16)\n",
    "plt.ylabel('Residual $\\\\mu$ (VI - MCMC)', fontsize = 16)\n",
    "plt.legend()\n",
    "plt.title(title_str)\n",
    "plt.show()\n",
    "\n",
    "plt.plot(true_avs, vi_point_estimates['mu'] - mcmc_point_estimates['mu'], 'o')\n",
    "plt.axhline(median_residual_vi_mcmc, color='k', label='median residual = ' + str(round(median_residual_vi_mcmc, 4)))\n",
    "plt.xlabel('True $A_V$', fontsize = 16)\n",
    "plt.ylabel('Residual $\\\\mu$ (VI - MCMC)', fontsize = 16)\n",
    "plt.title(title_str)\n",
    "plt.legend()\n",
    "\n",
    "plt.show()\n",
    "\n",
    "plt.plot(true_avs, vi_point_estimates['mu'] - true_mus, 'o')\n",
    "median_residual_vi_true = np.median(vi_point_estimates['mu'] - true_mus)\n",
    "plt.axhline(median_residual_vi_true, color='k', label='median residual = ' + str(round(median_residual_vi_true, 4)))\n",
    "plt.xlabel('True $A_V$', fontsize = 16)\n",
    "plt.ylabel('Residual $\\\\mu$ (VI - true)', fontsize = 16)\n",
    "plt.title(title_str)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "plt.plot(true_avs, mcmc_point_estimates['mu'] - true_mus, 'o')\n",
    "median_residual_mcmc_true = np.median(mcmc_point_estimates['mu'] - true_mus)\n",
    "plt.axhline(median_residual_mcmc_true, color='k', label='median residual = ' + str(round(median_residual_mcmc_true, 4)))\n",
    "plt.xlabel('True $A_V$', fontsize = 16)\n",
    "plt.ylabel('Residual $\\\\mu$ (MCMC - true)', fontsize = 16)\n",
    "plt.title(title_str)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "av_residuals = vi_point_estimates['AV'] - mcmc_point_estimates['AV']\n",
    "for i, residual in enumerate(av_residuals):\n",
    "\tif abs(mcmc_point_estimates['AV'][i]) > 0.75:\n",
    "\t\tprint(i, mcmc_point_estimates['AV'][i], residual)\n",
    "\n",
    "plt.scatter(mcmc_point_estimates['AV'] - true_avs, vi_point_estimates['AV'] - true_avs, c = true_avs)\n",
    "plt.plot(np.linspace(-0.3,0.3), np.linspace(-0.3,0.3), color='goldenrod')\n",
    "cbar = plt.colorbar()\n",
    "cbar.set_label(\"True AV\")\n",
    "plt.xlabel(\"MCMC AV Residuals (MCMC - true)\")\n",
    "plt.ylabel(\"VI AV Residuals (VI - true)\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}