notebooks/X_posterior_SHMF_animation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bayesian inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "import logging\n",
    "import numpy as np\n",
    "from scipy.stats import uniform, norm\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "%matplotlib inline\n",
    "\n",
    "sys.path.append('../')\n",
    "import paper_settings\n",
    "from simulation.units import *\n",
    "from simulation.prior import get_grid\n",
    "from inference.bayes import Posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "paper_settings.setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(62, 22.0, 88)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_samples = (\n",
    "    list(range(0,20)) + list(range(20,50,2))\n",
    "    + list(range(50,100, 5)) + list(range(100, 260, 10)) + [256]\n",
    ")\n",
    "\n",
    "frame_time = 0.25  # s\n",
    "frames = [8] + [1] * (len(n_samples) - 2) + [20]\n",
    "total_frames = sum(frames)\n",
    "len(n_samples), total_frames * frame_time, total_frames\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = \"calibrated_alices_full_sgd1e2_grid\"\n",
    "sample_seed = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flat_prior(thetas):\n",
    "    f_sub = thetas[:,0]\n",
    "    beta = thetas[:,1]\n",
    "    \n",
    "    return uniform(0.001, 0.199).pdf(f_sub).flatten() * uniform(-2.5, 1.0).pdf(beta).flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_per_frame = []\n",
    "for i, f in enumerate(frames):\n",
    "    samples_per_frame += [i] * f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_max = n_samples[-1]\n",
    "llr = np.load(\"../data/results/llr_{}.npy\".format(filename))[:,sample_seed:sample_seed+n_max]\n",
    "xs = np.load(\"../data/samples/x_test_full_point.npy\")[sample_seed:sample_seed+n_max]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "resolution = 25\n",
    "\n",
    "f_sub_1d = np.linspace(0.001, 0.200, resolution)\n",
    "beta_1d = np.linspace(-2.5, -1.5, resolution)\n",
    "\n",
    "theta0, theta1 = np.meshgrid(f_sub_1d, beta_1d)\n",
    "theta_grid = np.vstack((theta0.flatten(), theta1.flatten())).T\n",
    "\n",
    "bin_size = f_sub_1d[1] - f_sub_1d[0]\n",
    "alpha_edges = np.linspace(f_sub_1d[0] - bin_size/2, f_sub_1d[-1] + bin_size/2, resolution + 1)\n",
    "bin_size = beta_1d[1] - beta_1d[0]\n",
    "beta_edges = np.linspace(beta_1d[0] - bin_size/2, beta_1d[-1] + bin_size/2, resolution + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate posteriors on f_sub, beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "posteriors = []\n",
    "\n",
    "for n in n_samples:\n",
    "    if n == 0:\n",
    "        posteriors.append(flat_prior(theta_grid))\n",
    "    else:\n",
    "        bayes = Posterior(llr[:,:n], theta_grid)\n",
    "        posteriors.append(bayes.posterior_based_on_expected_llr(n, flat_prior))\n",
    "                          \n",
    "posteriors = np.array(posteriors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Drawing from posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_fsub_beta(posterior):\n",
    "    posterior /= np.sum(posterior)\n",
    "    \n",
    "    # Draw bin\n",
    "    i_bin = np.random.choice(list(range(625)), p=posterior)\n",
    "    \n",
    "    # Draw f_sub and theta uniformly within bin\n",
    "    f_sub = theta_grid[i_bin][0] + (f_sub_1d[1] - f_sub_1d[0]) * (np.random.random() - 0.5)\n",
    "    beta = theta_grid[i_bin][1] + (beta_edges[1] - beta_edges[0]) * (np.random.random() - 0.5)\n",
    "    \n",
    "    return f_sub, beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SHMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def M_200_sigma_v(sigma_v, scatter=False):\n",
    "    \"\"\"\n",
    "    Relate central velocity dispersion to halo virial mass\n",
    "    From https://arxiv.org/pdf/1804.04492.pdf\n",
    "    \"\"\"\n",
    "    a = 0.09\n",
    "    b = 3.48\n",
    "    if scatter:\n",
    "        sigma_log10_M_200 = 0.13  # Lognormal scatter\n",
    "        log10_M_200 = np.random.normal(a + b * np.log10(sigma_v / (100 * Kmps)), sigma_log10_M_200)\n",
    "    else:\n",
    "        log10_M_200 = a + b * np.log10(sigma_v / (100 * Kmps))\n",
    "    return (10 ** log10_M_200) * 1e12 * M_s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shmf_by_mhost(fsub, beta, msubs):\n",
    "    \"\"\" Evaluates dn / dm for a given set of m (msubs)\"\"\"\n",
    "    \n",
    "    sigma_v = np.random.normal(225, 50)\n",
    "    mhost = M_200_sigma_v(sigma_v * Kmps)\n",
    "    \n",
    "    m_min_calib=1e6 * M_s\n",
    "    m_max_calib=0.01 * mhost\n",
    "    M_0=M_MW\n",
    "    m_0 = 1.e9 * M_s\n",
    "                \n",
    "    alpha = fsub * ((2 + beta) * M_0 * m_0 ** beta) / (m_max_calib ** (beta + 2) - m_min_calib ** (beta + 2))\n",
    "    \n",
    "    return np.where(msubs > m_max_calib, 0., alpha * (mhost / M_0) * (msubs / m_0)**beta / mhost)\n",
    "                "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate SHMF data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_realizations = 10000\n",
    "fsub_true, beta_true = 0.05, -1.9\n",
    "msubs = np.geomspace(1.e7, 1e12, 200) * M_s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_and_ci(samples, levels=[0.68, 0.95]):\n",
    "    samples = samples[np.all(np.isfinite(samples),axis=1),:]\n",
    "    mean = np.mean(samples, axis=0)\n",
    "    all_down = [np.quantile(samples, (1. - level)/2, axis=0) for level in levels]\n",
    "    all_up = [np.quantile(samples, 1. - (1. - level)/2, axis=0) for level in levels]\n",
    "    all_cis = np.array([all_down, all_up])\n",
    "\n",
    "    return mean, all_cis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "shmfs = []\n",
    "for _ in range(n_realizations):\n",
    "    shmfs.append(shmf_by_mhost(fsub_true, beta_true, msubs))\n",
    "shmf_true, _ = mean_and_ci(np.asarray(shmfs), levels=[])    # (n_datapoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/johannbrehmer/anaconda3/envs/lensing/lib/python3.6/site-packages/ipykernel_launcher.py:12: RuntimeWarning: invalid value encountered in log10\n",
      "  if sys.path[0] == '':\n",
      "/Users/johannbrehmer/anaconda3/envs/lensing/lib/python3.6/site-packages/ipykernel_launcher.py:14: RuntimeWarning: invalid value encountered in greater\n",
      "  \n"
     ]
    }
   ],
   "source": [
    "shmf_posteriors = []\n",
    "shmf_posterior_cis = []\n",
    "\n",
    "for posterior_ in posteriors:\n",
    "    shmfs = []\n",
    "    for _ in range(n_realizations):\n",
    "        fsub_, beta_ = draw_fsub_beta(posterior_)\n",
    "        shmfs.append(shmf_by_mhost(fsub_, beta_, msubs))\n",
    "        \n",
    "    mean, ci = mean_and_ci(np.asarray(shmfs), levels=[0.68, 0.95])\n",
    "    \n",
    "    shmf_posteriors.append(mean)\n",
    "    shmf_posterior_cis.append(ci)\n",
    "    \n",
    "shmf_posteriors = np.asarray(shmf_posteriors)  # (n_frames, n_datapoints)\n",
    "shmf_posterior_cis = np.asarray(shmf_posterior_cis)  # (n_frames, 2, n_bands, n_datapoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 1./M_s * msubs\n",
    "y_mult = msubs**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animation with only SHMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_frame(frame, show=False):\n",
    "    step = samples_per_frame[frame]\n",
    "    n = n_samples[step]\n",
    " \n",
    "    # plt.clf()\n",
    "    \n",
    "    fillpost2 = plt.fill_between(x, y_mult * shmf_posterior_cis[step,0, 1], msubs**2 * shmf_posterior_cis[step,1, 1], facecolor=paper_settings.COLOR_FIX, edgecolor=paper_settings.COLOR_FIX)\n",
    "    fillpost1 = plt.fill_between(x, y_mult * shmf_posterior_cis[step,0, 0], msubs**2 * shmf_posterior_cis[step,1, 0], facecolor=paper_settings.COLOR_ALIGN, edgecolor=paper_settings.COLOR_ALIGN)\n",
    "    lpost, = plt.plot(x, y_mult * shmf_posteriors[step], alpha=1., c=paper_settings.COLOR_FULL, lw=1.0)\n",
    "    ltrue, = plt.plot(x, y_mult * shmf_true, alpha=1., c=\"black\", lw=1.0, ls=\":\")\n",
    "    text = plt.text(2.e7, 1.4e-2, r\"$n={}$\".format(n), ha=\"left\", va=\"bottom\")\n",
    "    \n",
    "    plt.xlabel(r\"$m_{200}$ [$\\mathrm{M}_{\\odot}$]\")\n",
    "    ax.xaxis.set_label_coords(0.5, -0.11)\n",
    "    plt.ylabel(r\"$\\frac {m_{200}^2}{M_{200}} \\frac{\\mathrm{d}n}{\\mathrm{d}m_{200}}$\")\n",
    "    ax.yaxis.set_label_coords(-0.115, 0.5)\n",
    "    ax.yaxis.set_ticks([0.,0.005,0.010])\n",
    "    ax.yaxis.set_tick_params(pad=1)\n",
    "\n",
    "    plt.xscale(\"log\")\n",
    "    plt.xlim(1.e7, 1.e12)\n",
    "    plt.yscale(\"log\")\n",
    "    plt.ylim(0.001,0.02)\n",
    "    \n",
    "    if show:\n",
    "        plt.show()\n",
    "        return\n",
    "    return [fillpost2, fillpost1, lpost, ltrue, text]\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAANYAAADTCAYAAAD9Lu2dAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAACz9JREFUeJzt3cFy29YVxvHvpOkyNaPUW09De9JtorAv0CrTlVeV47xA5HTfOpMX6IzT7JOofYA6drvxrmL6ApWYLLq10gdIZGrZReZ0gQsbhkAKpO4hCeL/m9EIgEgeUPbHC1xcXJm7C0Ber6x7B4BtRLCAAAQLCECwgAAECwhAsIAABAsIQLCAAAQLCPDqundgHjO7Len2a6+99uFbb7217t1Bj5ycnHzv7teXfb51YUjTaDTy4+Pjde8GesTMTtx9tOzzORQEAhAsIADBAgIQLCAAwQICECwgAMECAhAsIADBAgIQLCAAwQICECwgAMECAhAsIADBAgIQLCAAwQICECwgAMECAhAsIADBAgIQLCAAwQICECwgAMECAhAsIADBAgIQLCAAwQICECwgAMECAqz9D8+Z2VDSQNKepMfufrrmXQKuLFuLZWa7tfV9M9szs/uXPHVX0qmksaT9XPsDrFOWYJnZnqRHlfVdSXL3saRpPXRV7v7Y3adKLVaO/QHWLcuhoLuPzax6CHdX0lFaPlURmomZ1VuksbtPUzDHHAZiW0SdYw0knVXW35CK1qn+wBSqjyWdmtlR02OArll750U6XBzXt5vZgaQDSbpx48aqdwu4kqju9qmknbQ8kPTDoi/g7ofuPnL30fXr17PuHBAtKlgPJQ3T8lANLRKwzXL1Cu5LGpWdE+4+Sdv3JE3LdaAvcvUKPlatq9zdD3O8NtBFDGkCAmx0sMzstpkdnp+fr3tXgIVsdLDc/Ym7H1y7dm3duwIsZKODBXQVwQICECwgAMECAhAsIMBGB4vudnTVRgeL7nZ01UYHC+gqggUEIFhAAIIFBCBYQICNDhbd7eiqjQ4W3e3oqo0OFtBVSwXLzH7MvSPANlm2xbKsewFsmWWD5Vn3Atgy2WbCNbPfqZhD8Km7/yPX6wJdlHOK6Ym7/93M3sz4mkAn5ewVHJrZHyQRLPRezmANVHRqDHK9IBeI0VU5gzVx9z9L+ibXC3KBGF21ULDM7O05P+ZQEEhadV6Y2R9V/PE4l/Stmq9j7ar48z27kv6VaweBLmrbK3jk7t9W1puuY41V/FnUB1feK6Dj2gbrIzNzSSfu/lc1t1g77v6NmX2Zb/eAbmobrC9TaH6W1i+0WO7+tZm97e7ZOi+ArpobLDP7UNJNSW5mdyW9I+m3qrVY6RzspqRnKs7BgF67rMUau/tfzOyd1GL9ZsbjjlJXOwBd0t3u7t+lxdfLTTMe+oGZfZ6624Hea3uO9czMvpD0RVqvB+xvtV7DLMzstqTbt27dyv3SQChzX/wOEDP70d1/Uln/QkXYyl7DrEajkR8fH+d+WWAmMztx99Gyz881ur3eawj02pWCNafXEOi1q7ZYbXsNgV650uj2BXoNgV7JNZnMMzP7XNLZFfcH2ApLHQq6+yu19W8k/T7LHgFbgAk7gQBM2AkEYMJOIMCy3e0Xev+YVxB4oXWwUnDelPTdjIdkn1eQsYLoqkUOBSfu/pmkyYyfZ59Mhlma0FWLHAoOU6s16w7h7PMKAl21SLAuCw5TTAPJIsG6LDhlizbR7PMwoBcWOccqz6GGM36+K+k8fQd6rbHFSpPDPJM0SB0W0ssTcjZhXkEgmdVijSU9knSrtu0rFfdfNdlx93NJzCuI3psVrKaQvJ62HZYbzOzX5bK7f52+M68geq/xUHBGSM7StonZ8xFN75vZnbR8IumMURdAi17BeZNxuvtHqZfwTNKOpHuSCBZ6r013+9zJON39uzSJjEv6U7Y9AzqsTbA+MLN7KgbXflb9gZn9QkUv4U29+DM/n2TeR6Bz2gRr3mSch5L+KemRu/83214BHdcmWPU/4VN1x93PzewdM7spac/dabHQe22CVU7GWR1ibpKUut+f9x6aGdPVAmoxpKkMTRmitNz4vOpjcjCz22Z2eH6e9WWBcBs9mQz3Y6GrNjpYQFcRLCAAwQICECwgAMECAhAsIADBAgIQLCAAwQICECwgAMECAhAsIADBAgIQLCAAwQICECwgAMECAhAsIADBAgIQLCDARgeLWZrQVRsdLGZpQldtdLCAriJYQACCBQQgWEAAggUEIFhAAIIFBCBYQACCBQQgWEAAggUEIFhAAIIFBCBYQACCBQQgWEAAggUEIFhAAIIFBCBYQACCBQQgWEAAggUEIFhAAIIFBCBYQACCBQQgWEAAggUEIFhAAIIFBFh7sMxsYGZ7ZnbfzAbr3h8gh2zBMrPd2vp+GZhLnjqSdCxpKmmYa3+AdcoSLDPbk/Sosr4rSe4+ljSth64qPUaSBu4+ybE/wLplCVYKx2ll010VLZDS9j3peStW/RqY2YG7TyVNWrRuQCe8GvS6A0lnlfU3JMndH9cfaGbHqcUbSrrwc6CLooLVGod/2EZRwZpK2knLA0k/LPoCZnYg6SCt/s/M/tPiadcknV+yrbretHwtrZ9Xtv1c0vdL1J5Xr15z1vefZqo9b1/6+t7r21TZ/stL6s7n7lm+JB1VlnclHaTl+5J2r/jaxy0fd3jZtup607Kkw+py2/pNtefVq9ec9T1Xbd578/utblu09ryvXL2C+5JG6fvzw7t07jT11R3uPWmx7ckly09qy1epPa9evea87zlqz9uXvr73+rZl3nsjSwndaGZ27O6jPtbnvXez9to7L1o6LBdSq3hXRa/jU3f/dMX1Byp6MHdUHC5MZz4rf+2hisPsHUlfraD28/rpWuSOu4/T72Ck4vx5HLgfF2o3rUfWrtdL/wYDFZeQHrv7adOT1z6kqa5pBIek08o1rom731FxQfqw/vwV1B+pCHYZsFXWvidpnL4OlMm8UTPuXv6Oy/csSZ+k/9RX3o8lajetZ68/p96uimuzY0n7s153o4LVZgRH9RMi9ydlyxEkxyo+rR7o5Yviq6j9VEVrNVS6NriiupL0sPK0QXrMVNLNFdduWl9ZfXd/nN73nuZcd92oYPkCIzj08gXoVdZ/393flXRHGVuNNrXTJ/ggff17VXUbnjZNo2YGKsK+ytrZLFs/BXI86zBQ2rBgNWgcwaHiEztba7Fg/dP0i5ViR4pcqJ0+QYcqjvejas/6ne9J+lUK05dpfU95D8fb1G5aX1n99G//saR7ZS94k650XrxkRR0Ws2pHnjBfVnsiaS0jVWq/86lW88HWVHvl//61euW55Vyb3mJdeQRHh+uvq3Yf33P2+pserId60fM2VItPii2qv67afXzP2etvVLDWPYJjnfXXVbuP73kV9Tsx8gLomo1qsYBtQbCAAASrJ9IsWDPnHlngdfbnXb9BgWD1x9TdJ2a2a2ZH9flF0vi8p5WL343ShemdeY8Bweqd1NP1QMXwnbrTdV4A3yYEq79O0y0QpexjL/usk0Oatk0697kr6UjFf/A9FRcndyQNK7dP5PRQxW0Pn6b6Kxui1Ae0WJuhHFF9lg7V3lPxH/1Y0rtScZNjOj+6n5bLqbnL+RlfWr+sYDpXem/eY8qaV3tr/USwNkC6/WBYudI/Tff8jFS0YtLFG+zqNxsuc/PhNB0OXrivrbw1J3V47AWMJN9qBGvDpBaivNfqPRUzBO823GBXv9lwmZsPH0p6UL+vqFrPzAYpsMyrvwCCtQFSq1G2TNXBnz+omDquOoatvMGufrNhq5sP02s8MLP9dDhYhvhA0jD9fKeyX5/kfK99wVjBjqjcYHeqIoQTFYeH0ouOjufr9WkLrJgjv1UnSArnTtmSlRPoVALe+rX6il7BjqicP1VVD+Gy3XyYDgF3KheLz5ypwBdCsPpjkM6dWgUktVYXgho138i24VAQCEDnBRCAYAEBCBYQgGABAQgWEIBgAQH+D3wYk7WsKEI0AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 204.48x204.48 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = paper_settings.figure()\n",
    "ims = [draw_frame(n) for n in range(total_frames)]\n",
    "ani = animation.ArtistAnimation(fig, ims, interval=200)\n",
    "ani.save('../figures/live_inference_shmf.gif', dpi=300, writer='imagemagick')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animation of only population-level parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_frame(frame, show=False):\n",
    "    step = samples_per_frame[frame]\n",
    "    n = n_samples[step]\n",
    "    \n",
    "    pmin, pmax = 1.e-9, 0.05\n",
    "    z = np.clip(posteriors[step], pmin + 1.e-9, pmax)\n",
    "    z[~np.isfinite(z)] = pmin + 1.e-9\n",
    " \n",
    "    # plt.clf()\n",
    "\n",
    "    pcm = ax.imshow(\n",
    "        z.reshape((resolution, resolution)),\n",
    "        extent=(alpha_edges[0],alpha_edges[-1],beta_edges[0], beta_edges[-1]),\n",
    "        cmap=paper_settings.CMAP2,\n",
    "        origin=\"lower\",\n",
    "        aspect=\"auto\",\n",
    "        interpolation=\"nearest\",\n",
    "        norm=matplotlib.colors.LogNorm(vmin=pmin, vmax=pmax),\n",
    "    )\n",
    "    sc = ax.scatter(0.05, -1.9, s=15., color='black', marker='*')\n",
    "    text = plt.text(0.145, -2.45, r\"$n={}$\".format(n), color=\"white\", ha=\"left\", va=\"bottom\")\n",
    "    \n",
    "    plt.xlabel(r'$f_{\\mathrm{sub}}$')\n",
    "    plt.ylabel(r'$\\beta$')\n",
    "    plt.xlim(0.,0.2)\n",
    "    plt.ylim(-2.5,-1.5) \n",
    "    plt.xticks([0.,0.05,0.1,0.15,0.2])\n",
    "    plt.yticks([-1.6,-1.8,-2.0,-2.2,-2.4],[\"-0.6\", \"-0.8\", \"-1.0\", \"-1.2\", \"-1.4\"])\n",
    "    \n",
    "    if show:\n",
    "        plt.show()\n",
    "        return\n",
    "    return [pcm, sc, text]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAANgAAADUCAYAAAD+4u2WAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAACQNJREFUeJzt3UFyG8cZxfH3pbJNFSVKWxULWmRrU9AJTF0gxcQXiKkbSPEJHPICLsm5gCMmFzCVCxii1l4IqqxpyqjKAT4v0LBGEIABEbyZwej/q3KJw54B2oV57EFj0F9kpgB4/K7tDgB9RsAAIwIGGBEwwIiAAUYEDDAiYIARAQOMCBhg9Pu2O3ATd+7cyYODg7a7gU/Mq1evfs7Mu5scu1MBOzg40Gg0arsb+MRExH83PZZLRMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwst+LGBHHkiaSDjPzbEH7oaSBJGXmubs/QJOsI1gJjzLzQtJktj3n6xKswZJ2YGe5LxG/1HT0kqSxpKNqYxndfpSkzDzLzEtzf4BGuQO2J+ldZXt/rv2hpP2IOIyIJ+a+AI3rwiTH9WzkKiPaByLiJCJGETG6urpqvnfA/2HrAYuIZxHxooxIE0m3S9OepOu53a81vXRU2ffh/ONl5vPMHGbm8O7djb5UCrRm67OImfl49nOZtBiWzYGki/L7vcycSDqXNBu19lTejwF9Yb1ErFz6HUmaVCYxXpb2saazi8eS9pmmR9/YPwfLzOcLfvdgQTvhQu90YZID6C0CBhgRMMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwImCAEQEDjAgYYNSFAnyz9sGiRUqBXdZqAb6yPS7tYwrwoW9aLcBXnJZ/BxTgQ9+0WoCvBGocEb/M7Qf0QquTHBGxp+kI942k7yJisGAfCvBhZ219kiMinmladO9H1RfgO5H0TWZOImKsaa2wDyZCysTHc0kaDoe57f4CTm0X4Ksedx4RJ9vuD9Am6zR9Zl5GxHBJAb4HmXkWEU/K6HWbaXr0TRcK8H302RjQF9zJARgRMMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwImCAEQEDjAgYYNRIwFaVJYqI44g4iognTfQFaJI9YGVV3xdL2lbWDwN23dKARcQXEfH3iPhT2f4sIj676RPMiustaV6nfhiws1aNYJmZf5P0NiK+1bR4w8OI+DYiDrb0/CvrhwG7btXa9LckKTNfR8Q4M/89a4iIv0r6h7tzwK5bGrDM/FdEfCEpJT2ba3677LhqfbA1CjvU1Q9TKWl0Ikn37t2reTigW1ZWV8nMl5IUEZ9HxC1Joeml4kdBqBzzeFnbTKU+2PdaUD9s7vEowIedtdYsYma+zsz/ZObLzPxO0usyCVI76RERx5KG5d+Zl+VxL8s+8/XDgF7YqD5YZr7VisvEuX3PJZ3P/a5aH4yie+gt7uQAjAgYYETAACMCBhgRMMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwImCAEQEDjAgYYETAACMCBhgRMMCIgAFGBAwwImCA0UYLj95URBwuW7W3rD0vSfcz82kT/QGa0nYBviNJF2V130HZBnrDHrCaAnwDvS+6Ny7bQG80com4zNy69IeaVlsBeqMTkxylNvMl1VXQN1sPWEQ8i4gXEfHkBocdLZvgiIiTiBhFxOjq6mpLvQSasfVLxBsW4FNEnMwqYUbEUXnPVn08CvBhZzUxi7i0AF+ZNTyNiDcR8Yu7L0DT7JMcqwrwldHqlrsPQFs6MckB9BUBA4wIGGBEwAAjAgYYETDAiIABRgQMMCJggBEBA4wIGGBEwAAjAgYYETDAiIABRgQMMCJggBEBA4wIGGBEwAAjAgYYETDAiIABRgQMMGokYKW4Q90+N1nLHtgJrRbgm9vnkbsvQNPaLsAH9Frr78FK/eaL+j2B3dN6wCTdbrsDgEurBfjWGb0owIdd1nYBvkFEDDQdxW6XwH1QRpYCfNhlrRbgy8zzUj9MkvbcfQGaFpm7MygMh8McjUZtdwOfmIh4lZnDTY7twiQH0FsEDDAiYIARAQOMCBhgRMAAIwIGGBEwwIiAAUYEDDAiYIARAQOMCBhgtFN300fE/yT91HY/Vrgj6ee2O7EEfdvcHzPzD5scuPUvXJr9tOnXBpoQEaOu9o++bS4iNv6OFJeIgBEBA4x2LWDP2+5AjS73j75tbuP+7dQkB5Yra55MJB1m5tmSfT5YVGidY1rs22lmPo2Ik7Lw0U7q1AgWEccRcbRsybdF7XXHdKB/p+XfE2O/DqXfVlGeLKoFML+E+TrHtNW34iQi3qiBVaHXeF1Pyn+n6x4z05mA1b0Qi9qbOkk27V9pauJE+VLTEULleY7md1iwhHntMS32TZK+ysz77lWf13hdjyRdlFF0UEK19nnXmYCp/oVY1N7USbJp/6RmTpQ9Se8q2/umYzax6fMM1hkhtqDudR1Ufjcu22ufd10KWN0Lsai9qZNk2fOv097UidIrmXlW/ijtl1HEZeXrmpnPK+8BDyWN6o6p6lLAeqmhE2Wi92v870m6Nh2ziRs/T3m/M1uo9lrTUaNV5TLwcn7l6TpdCljdC7GovamTZKP+NXiifF957IGkC2m6RPlNj+lI30aV/twv2y7rnkNHmfn0hsd0KmB1L8Si9qZOkk3718iJMvurWkbISeWv7MvZPvNLmK84pit9+0vZfuPqW1H7B6B8VHBW+f9Y+7zr1OdgZSp7LGkwu+4tyxY/WNH+0e862L935Xe2z5qwuVWva+UjhHeajlp/zsyLdc+7TgUM6JsuXSICvUPAACMCBhgRsB4qt5GdzhU9RAt27RvNWM/Xkp7W7gU7AtZTmWm/Cx31uETsmXJZuGe+fw9rImD9M5b0Q93d++UGZN6jmRGw/hlKWufWorHe308HE96D9c8DSf+UfrtcvNT0axYTTW9MvZR0LOlc0v2IGGj6Vf7zdrrbb4xgPZSZ1S8DPi4/jzS9b26s6Y3H0vRG2rGkRw138ZNBwPpnUvl5qOmd37PQ7ZfvNc3uBN8r2z8028VPBzf79kSZNXwk6RlT9N1BwAAjLhEBIwIGGBEwwIiAAUYEDDAiYIARAQOMfgWN2oNobtzizgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 204.48x204.48 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = paper_settings.figure()\n",
    "ims = [draw_frame(n) for n in range(total_frames)]\n",
    "ani = animation.ArtistAnimation(fig, ims, interval=200)\n",
    "ani.save('../figures/live_inference_population.gif', dpi=300, writer='imagemagick')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combined animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_frame(frame, show=False):\n",
    "    step = samples_per_frame[frame]\n",
    "    n = n_samples[step]\n",
    " \n",
    "    # Left part\n",
    "    plt.sca(ax_left)\n",
    "    \n",
    "    pmin, pmax = 1.e-9, 0.05\n",
    "    z = np.clip(posteriors[step], pmin, pmax)\n",
    "    z[~np.isfinite(z)] = pmin\n",
    "\n",
    "    pcm = ax_left.imshow(\n",
    "        z.reshape((resolution, resolution)),\n",
    "        extent=(alpha_edges[0],alpha_edges[-1],beta_edges[0], beta_edges[-1]),\n",
    "        cmap=paper_settings.CMAP2,\n",
    "        origin=\"lower\",\n",
    "        aspect=\"auto\",\n",
    "        interpolation=\"nearest\",\n",
    "        norm=matplotlib.colors.LogNorm(vmin=pmin, vmax=pmax),\n",
    "    )\n",
    "    sc = ax_left.scatter(0.05, -1.9, s=15., color='black', marker='*')\n",
    "    # text2 = plt.text(0.145, -2.45, r\"$n={}$\".format(n), color=\"white\", ha=\"left\", va=\"bottom\")\n",
    "    \n",
    "    plt.xlabel(r'$f_{\\mathrm{sub}}$')\n",
    "    ax_left.xaxis.set_label_coords(0.5, -0.11)\n",
    "    plt.ylabel(r'$\\beta$')\n",
    "    plt.xlim(0.,0.2)\n",
    "    plt.ylim(-2.5,-1.5) \n",
    "    plt.xticks([0.,0.05,0.1,0.15,0.2])\n",
    "    plt.yticks([-1.6,-1.8,-2.0,-2.2,-2.4],[\"-0.6\", \"-0.8\", \"-1.0\", \"-1.2\", \"-1.4\"])\n",
    "    \n",
    "    \n",
    "    # Right part\n",
    "    plt.sca(ax_right)\n",
    "    \n",
    "    fillpost2 = ax_right.fill_between(x, y_mult * shmf_posterior_cis[step,0, 1], msubs**2 * shmf_posterior_cis[step,1, 1], facecolor=paper_settings.COLOR_FIX, edgecolor=paper_settings.COLOR_FIX)\n",
    "    fillpost1 = ax_right.fill_between(x, y_mult * shmf_posterior_cis[step,0, 0], msubs**2 * shmf_posterior_cis[step,1, 0], facecolor=paper_settings.COLOR_ALIGN, edgecolor=paper_settings.COLOR_ALIGN)\n",
    "    lpost, = ax_right.plot(x, y_mult * shmf_posteriors[step], alpha=1., c=paper_settings.COLOR_FULL, lw=1.0)\n",
    "    ltrue, = ax_right.plot(x, y_mult * shmf_true, alpha=1., c=\"black\", lw=1.0, ls=\":\")\n",
    "    text1 = plt.text(2.e7, 1.4e-2, r\"$n={}$\".format(n), ha=\"left\", va=\"bottom\")\n",
    "    \n",
    "    plt.xlabel(r\"$m_{200}$ [$\\mathrm{M}_{\\odot}$]\")\n",
    "    ax_right.xaxis.set_label_coords(0.5, -0.11)\n",
    "    plt.ylabel(r\"$\\frac {m_{200}^2}{M_{200}} \\frac{\\mathrm{d}n}{\\mathrm{d}m_{200}}$\")\n",
    "    ax_right.yaxis.set_label_coords(-0.115, 0.5)\n",
    "    ax_right.yaxis.set_ticks([0.,0.005,0.010])\n",
    "    ax_right.yaxis.set_tick_params(pad=1)\n",
    "\n",
    "    plt.xscale(\"log\")\n",
    "    plt.xlim(1.e7, 1.e12)\n",
    "    plt.yscale(\"log\")\n",
    "    plt.ylim(0.001,0.02)\n",
    "    \n",
    "    \n",
    "    # Done\n",
    "    if show:\n",
    "        plt.show()\n",
    "        return\n",
    "    return [pcm, sc, text1, fillpost2, fillpost1, lpost, ltrue]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaEAAADKCAYAAAAM7P/vAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAEnJJREFUeJzt3c1yG9eZxvHnyWSWiWDK2rpsSOXZxhR8AxOostIqspULGFPOfiKPb2BSUrJPRPsCIouTjXcjKhcwpmgvZjMLwdkzoqCq2cwi9c6iT0vtFr5IdvdpAP9fFYvoD+C8IAi8OKfPhyNCAADk8KPcAQAAthdJCACQDUkIAJANSQgAkA1JCACQDUkIAJANSQgAkA1JCACQDUkIAJDNj3MHcBZvv/12vPvuu7nDwJZ5+vTp3yLiSu44FrF9U9LNn/zkJ5+8//77ucPBFrno+8PrNG3PaDSKo6Oj3GFgy9h+GhGj3HGsgvcIunbR9wfNcQCAbEhCAIBsSEIAgGxIQgCAbEhCAIBsSEIAgGxIQgCAbEhCAIBsSEIAgGxIQgCAbFqfO872LUlTSbsRcX/G8V1JQ0mKiIO24wEA9EerNaGUYBQRh5Km5XbN5yn5DOccBwBsqLab426rqAVJ0kTSuHow1ZK+kaSIuB8Rxy3HAwDokbaT0EDSaWX7cu34h5Iu2961fbflWAAAPdOHjgnPyxpQqhn9gO0920e2j05OTrqPDgDQmsaTkO0Hth+lms1U0k46NJD0vHb6cxXNdErnflh/vIjYj4hRRIyuXOn1umIAgDNqvHdcRNwpb6eOBuViR0NJh2n/ICKmkg4klbWfgdL1IQDAdmi1Oa7SzDaWNK10PHiSjk9U9Jq7JekyXbQBYLu0Pk4oIvZn7Ls+4zgJCAC2TB86JgAAthRJCACQDUkIAJANSQgAkA1JCACQDUkIAJANSQgAkA1JCACQDUkIAJANSQgAkE3r0/YAaIbtoYqJfseSDtLci8BaoyYEdKy+jL3tW7bHKyzsuKti6ZNDvZ59HlhrJCGgQ2lG+UeV7V1JiohDFTPK7867b0QcpCVQxmLCX2wImuOADkXEoe1qM9ptSY/T7YmKBHM8Y5Xhw4iYpiR2SFMcNgVJCMhrIOm0sn1ZKmo99RNTAvpM0sT2Y9bfwiYgCQFrIjXZHdb3296TtCdJ77zzTtdhARfCNSEgr6mknXR7IOn5WR8gIvYjYhQRoytXrjQaHNA2khCQ10NJw3R7qBk1HWCTkYSADqUOB6Oy40FEHKf9Y0nTchvYFlwTAjqUOhMc1PbtZwoHyI6aEAAgG5IQsAFs37S9//Lly9yhAGfSehJaNiVJ5fhe27EAmyoivo6IvUuXLuUOBTiTVpPQsilJ0vYkHZ8smrIEALB52q4J3VYxDkJ6PSVJ3b30e0jPIADYLm0noZlTkpRS0pnYflE7DwCwBbJ2TLA9UFFT+q2kL9J6KfVz9mwf2T46OTnpPEYAQHsaHydk+4GKaUi+0fIpSfYk/TbNDjxRsUbK/eoJaQzFviSNRqNoOl5gE9i+KenmtWvXcocCnEnjNaGIuBMRH0XEfc2ZkiTVgOr3O9Dr60cAzoDecVhXrc6YEBHHtkczpiR5Iul6RNy3fTfVgnYYOQ4A26X1aXtmJZaIuF65fb9+HACwHZgxAQCQDUkIAJANSQgAkA1JCNgATGCKdUUSAjYAXbSxrkhCQEts/z13DEDfkYSA9jh3AEDfkYSA9jDNFLBE64NVAbxm+5cqprB6FhF/zh0PkBtJCOjWcUT8h+33cgcC9AHNcUC3hrb/VRJJCBBJCOjaQEWHhTdmkr8IxglhXZGEgG4dR8TvJH3b5IMyTgjriiQENMz2zxYcpjkOqKBjAtAQ27+RdFlF1+zvNHuc0K6KxRt3Jf2lu+iAfiIJAc15HBHfVbZnjRM6lDSRdK+bkIB+IwkBzfnUdkh6GhFfanZNaCcivrX9oOPYgF4iCQHNeZASzE/T9hs1oYh4YvtnEdFoxwRgXZGEgAuy/Ymkq5LC9m1JH0j6hWo1oXTN6KqkFyquGQFbjyQEXNxhRHxh+4NUE/r5nPMep+7ZABK6aAMXFBHfp5tvlbvmnPor239IXbQBiJoQ0KQXtv8o6Y9pu56M/lTrPdcY2zcl3bx27VobDw+0ppOakO3dBcdu2R7bvttFLEBbIuLbiPh0QaL5NNWE/qWFspkxAWup9SRkeyzp0Zxju5IUEYeSpouSFbABHkTEryV9lTsQoC/mNseli6s3JP1XRPy5nIrkrM0JEXFoezLn8G1Jj9PtiaSxpOOzPD7Qdwt6zwFbb9E1oYiIf7P9ge0/qEgWl23fkXQvIv7aQPkDSaeV7csNPCbQN6v2ngO2zqIk9JZUtHPbnlRXgUxt2l+2HRywCc7Qew7YOnOvCaXVH39u+58l1acY+X7WfSTJ9gPbj1bsaDCVtJNuDyQ9n/F4e7aPbB+dnJys8JBAb9Sn7XmRWhVOZ50MbKOFXbQj4okkpSa5t1S8qYaakSwq97mzrFDbg4iYSnooaZR2D1VM7lh/vH1J+5I0Go34Bom1ERE/qm1/K+nXmcIBemml3nGp6+lfIuJJRHwh6dtUS1q0boqkogu2pFH6XXqSHvc4nTOWNC23AQDb4VyDVVMb99wmudq5B5IOavuuV27vnycGoO9s/z0i/iF3HECfMW0P0J5ZSzkAqGDaHqA9b1zDtP1LFdc/n1V7nALbiiQENCglmfc0v7n6OPU8fa/hcpk7DmuJ5jigWccR8XvNn/ljmGbRbjQJMXcc1hU1IaBZw1Qbmrdy6kDFtaJBdyEB/UUSApq1LMm00hwHrCuSENCsZUmmrCkda8VhDsAm45oQ0Kzyms9wzvFdSS/Tb2DrURMCzsn2byS9kDRInRGkIrlMNT/JHKpYtuRe+xEC/UdNCDi/QxULNl6r7ftKxfpBs+xExEu9OSkwsJVIQsD5zUoob6V9r6ajSjPRS3o9KXCazBTYejTHAec0J6Gcpn3H9qtZez62/VG6/VTSKbMlAAWSENCAdH3oqoprRN9Vj0XEp6m33KmK9bPuSCIJASIJAU15HBG/m3cwIr63/VMV88n9e3dhAf1GEgKa8Svbd1RMTPr76gHb76roLXdV0mUViejzrgME+ogkBDTjTxHx3Zxj+5L+U9KjiPhrdyEB/UcSAprxqe2Q9DQivqwd+ygiXtr+wPZVSeOIoCYEiCQENOVBRHxruzqNtSUpddl+1YvO9lHThbOUA9YV44SABpQJpkw46fbM91f1nAbLZykHrCWSEAAgG5IQACAbkhAAIJtOOibY3o2Imcsd295LN69GxGddxAMA6IfWa0K2xypmGp537DAi9lWswzJuOx4AQH+0noQiolw/ZZahpDLxTDR/ITAAwAbKOk4o1YBKu5Ie5ooFANC9XnRMsL0r6XjedSMAwGZqPAnZfmD7ke27Z7jbeF6nBNt7to9sH52cnDQUJQCgDxpvjouIO8vOsT2IiGm6vRcR99PtcbqGVH28faVVKkejUTQdLwAgny56x92SNEq/S0/SsbGke7af2X7RdiwAgH5pvWNCRBxIOqjtu55+H0p6q+0YAAD91IuOCQCA7UQSAgBkQxICAGRDEgIAZEMSAgBkQxICAGRDEgIAZEMSAjaA7Zu291++fJk7FOBMSELABoiIryNi79KlS7lDAc6EJAQAyIYkBADIhiQEAMiGJAQAyIYkBADIhiQEAMiGJAQAyIYkBADIhiQEAMiGJAQAyIYkBADIhiQEAMiGJAQAyIYkBADIppMkZHt3hXPudhELAKA/Wk9CtseSHq1wzo22YwEA9EvrSSgiDiVN2i4HALB+sl8Tsr2bEhUAYMtkT0KSdnIHAADIo/EkZPuB7UerdDRYpRZke8/2ke2jk5OT5gIFAGT346YfMCLuLDvH9iAippKGtocqakM7KSkd1x5vX9K+JI1Go2g6XgBAPl30jrslaZR+l55IUkQcRMRB2jdoOxYAQL80XhOqS0nmoLbvem37VW0HALA9+tAxAQCwpUhCAIBsSEIAgGxIQgCAbEhCAIBsSEIAgGxIQsCasD2wPbZ91zbj6rARSEJAx+rra9m+VSaXJXcdSTqSNJU0bCs+oEskIaBD9fW1yoSU5lCcLloAsjLP4qA+vRWwrkhCQIdmrK91W0XNRmn/WHpVO6r+DGzvpTkXj1mJGJui9Wl7ACw0kHRa2b4svZru6gfSbPJjFU1xbxwH1hFJCFgTNMFhE61VEnr69On/2v6f3HEs8Lakv+UOYg5iO79/avGxp3q9sONA0vOzPoDtPUl7afP/bP/3krtckvRyyf76OeX2ot//qOWv4yplz4ulvk9nLHte+auUPSuObX/uZQwXe39ExNr8SDrKHcO6xkds/YlP0uPK7V1Je+n2XUm7bccqaX/Z/vo55fai302VPS+W+r6zlj2v/FXK5rnPfM7l7Qu9P+iYAHSovr5WpCa2dK1nGt00uX29wv76OV+v+LuJsufFUt931rLnnbtK2bPi2Pbnfpay53LKZGvB9lFEjHLHMU+f4yO28+t7fFU5Y93WsnOXv+7Pfa2uCan/C9/1OT5iO7++x1f1KtZU27qtovfds4i432HZAxW9+HZUNNdM596r4bJT+UMVTZ07kr7qqvw0zmsnIg7T32Ck4lrfYYsxvFH2rO0WzXruQxXPeyzpICIm8+68Vs1xUazA2ls541s26j4i9i8wUj9HbPfS771Z9+lSX//vZr2ekiaVv/NxRHykYnBso89hhbJHKhJgmYwaNWNQ72ntf+yOpMP008j/0KL3T+V/pHzOkvR5SgAXjuEcZc/a7rL8XRXj3g4l3Vr02L1KQss+rGYd7+qD9ALxtf5husqo+4uM1O86tmTP9jP9cGBnK1Z4XffSz71V79O2VV7P6rfPJr+Fr/i/dKTiW/A9Nfwarlj+MxW1oKHS2KsOypSkh5W7DdI5U0lXOy571nZn5UfEQXreYy0Z09abJLTsw2rW8a4+SM8bXzrUxYfpzFH3VbHiSP2exCZJn0TE1babElZ4XccqmlL2JQ1T4uns/26eVV/PVEM5VYNWLPvjiLgu6SM1VBM5S/np9Rqkn2+6KHPG3aYuZroYqEiKXZbdmPOWX3nvLPzs600S0vInNut4ly/GeeKTuvkwnTnqvoX7nMd5yyk/8NuuaSx7XYeVfZO03emHwIrm/Z2Har82OavsSfoQktqf3eGN8tMXg6GKaxRtlD/v7z2W9GFKPA/S9ljNNoeuUvas7c7KT6/9Z5LupC9Cc/WpY8KyD6tZx7v6IJ1X/irHh+kF2e3gwvDGKP9Wtm/YHreYxBe+rrXrQbsqmhyuL7pPn+T6n+vgYviy8o8ldT7DRO3vPVUHzclzyu78ta+VV14LW6pPNaGNFBH30xvycuWbYdPOM+r+wiP1V3TmctL1l/Lb03P1YNmC9M36uKNxPOfR1evZt7Jzlc/fu6Hy+5SElj2xWce7fDHOHF+HH6YPK489VPoGsqQaPvM+PYntqBLP1bTdllX/h8YR8dkZ79Olrl7PvpWdq3z+3g2V36cktOzDatbxLl+M88TXyYfpglH3T8pznGmk/gVi+zhtP2u59rE0SbpYQqFsHhzPu0+Xcr2eucvOVT5/7/bK79WMCakb80TSsGyLt/009bKZd/yNfT2M7zTt45pQDy16XSvdU09V1H4+SoPxOvu/AzZZr5IQAOSUemMeXvjb/etaA+s+LdGn5jg0JI2huresaySAN0wj4ji9hx7XhwikYQPPlnUySslnZ9E5KJCENtPnKsYo9LUnF9BrqSZ0T8WYsLpJ7i7om4QktKEiYrJspDKApSYuJuMsNTr7BEhCGyc1wQ2WNRcA66bSzDxOt++m32O3NzfjQ6UJONNYMb7YNYwktHkmKlbuXNhckN64XDPCOimnSjpNzWU3VPy/H6mYxUK2h5UENSynkHEx4eygvr2swHRt58aic8oyL/bUthdJaPOMtNq1oIm4cIo1kpqXh5Wea9M0U/NI0uO0r76EQH05hfMsrzBNTXJvzEReThCbOjOMW5inbeORhDbPdaVBsenb3tCvlx0ot8seP1fL4/nCBc4u1TzK2bFvSDp2sXxFfQmB+nIK51le4aGke/VrrNXybA9Scss+xdS6IQltoMraMRMVi3tJRWIapjdS+eZ7lrYXNjcAfZBqI2WNpzpTxXMVEwRXR/KXSwjUl1NYaXmF9Bj3bN9KTXJlwtvT60mJdypxfd7kc90mfZpFG82oNhmMlKaYSd/WqlPcS0UHhl29fmMDvZWSyn66fVDZ/2omksoSAhPbj/V6OQWl++7UtueVdah0nalaRvpdTuE0sF1+sfuMprjzYcaEDZHefDckPaBrNnA+aZ7AladhSrWg8kvdaeV61Zkfa1uRhAAgYdqe7pGEAADZ0DEBAJANSQgAkA1JCACQDUkIAJANSQgAkA1JCACQDUkIAJDN/wOuv2Wv+kyy4AAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 404.64x204.48 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax_left, ax_right = paper_settings.two_figures()\n",
    "ims = [draw_frame(n) for n in range(total_frames)]\n",
    "ani = animation.ArtistAnimation(fig, ims, interval=200)\n",
    "ani.save('../figures/live_inference_both.gif', dpi=300, writer='imagemagick')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Full animation including lens images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "pmin, pmax = 1.e-9, 0.05\n",
    "xmin, xmax = 2.3, 3.15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_frame(frame, show=False):\n",
    "    global ax_images\n",
    "    \n",
    "    step = samples_per_frame[frame]\n",
    "    n = n_samples[step]\n",
    " \n",
    "    # Left part\n",
    "    plt.sca(ax_left)\n",
    "    \n",
    "    z = np.clip(posteriors[step], pmin, pmax)\n",
    "    z[~np.isfinite(z)] = pmin\n",
    "\n",
    "    pcm = ax_left.imshow(\n",
    "        z.reshape((resolution, resolution)),\n",
    "        extent=(alpha_edges[0],alpha_edges[-1],beta_edges[0], beta_edges[-1]),\n",
    "        cmap=paper_settings.CMAP2,\n",
    "        origin=\"lower\",\n",
    "        aspect=\"auto\",\n",
    "        interpolation=\"nearest\",\n",
    "        norm=matplotlib.colors.LogNorm(vmin=pmin, vmax=pmax),\n",
    "    )\n",
    "    sc = ax_left.scatter(0.05, -1.9, s=15., color='black', marker='*')\n",
    "    \n",
    "    plt.xlabel(r'$f_{\\mathrm{sub}}$')\n",
    "    ax_left.xaxis.set_label_coords(0.5, -0.11)\n",
    "    plt.ylabel(r'$\\beta$')\n",
    "    plt.xlim(0.,0.2)\n",
    "    plt.ylim(-2.5,-1.5) \n",
    "    plt.xticks([0.,0.05,0.1,0.15,0.2])\n",
    "    plt.yticks([-1.6,-1.8,-2.0,-2.2,-2.4],[\"-0.6\", \"-0.8\", \"-1.0\", \"-1.2\", \"-1.4\"])\n",
    "    \n",
    "    \n",
    "    # Right part\n",
    "    plt.sca(ax_mid)\n",
    "    \n",
    "    fillpost2 = ax_mid.fill_between(x, y_mult * shmf_posterior_cis[step,0, 1], msubs**2 * shmf_posterior_cis[step,1, 1], facecolor=paper_settings.COLOR_FIX, edgecolor=paper_settings.COLOR_FIX)\n",
    "    fillpost1 = ax_mid.fill_between(x, y_mult * shmf_posterior_cis[step,0, 0], msubs**2 * shmf_posterior_cis[step,1, 0], facecolor=paper_settings.COLOR_ALIGN, edgecolor=paper_settings.COLOR_ALIGN)\n",
    "    lpost, = ax_mid.plot(x, y_mult * shmf_posteriors[step], alpha=1., c=paper_settings.COLOR_FULL, lw=1.0)\n",
    "    ltrue, = ax_mid.plot(x, y_mult * shmf_true, alpha=1., c=\"black\", lw=1.0, ls=\":\")\n",
    "    text1 = ax_mid.text(2.e7, 1.4e-2, r\"$n={}$\".format(n), ha=\"left\", va=\"bottom\")\n",
    "    \n",
    "    plt.xlabel(r\"$m_{200}$ [$\\mathrm{M}_{\\odot}$]\")\n",
    "    ax_mid.xaxis.set_label_coords(0.5, -0.11)\n",
    "    plt.ylabel(r\"$\\frac {m_{200}^2}{M_{200}} \\frac{\\mathrm{d}n}{\\mathrm{d}m_{200}}$\")\n",
    "    ax_mid.yaxis.set_label_coords(-0.115, 0.5)\n",
    "    ax_mid.yaxis.set_ticks([0.,0.005,0.010])\n",
    "    ax_mid.yaxis.set_tick_params(pad=1)\n",
    "\n",
    "    plt.xscale(\"log\")\n",
    "    plt.xlim(1.e7, 1.e12)\n",
    "    plt.yscale(\"log\")\n",
    "    plt.ylim(0.001,0.02)\n",
    "    \n",
    "    \n",
    "    # Lens images\n",
    "    while len(ax_images) < n:\n",
    "        ax_images = paper_settings.add_image_to_roster(fig, ax_images, images_coords, sep_fraction=0.08)\n",
    "        \n",
    "    images = []\n",
    "    for i, ax in enumerate(ax_images):\n",
    "        images.append(\n",
    "            ax.imshow(\n",
    "                np.log10(xs[i]),\n",
    "                vmin=xmin,\n",
    "                vmax=xmax,\n",
    "                cmap='gist_gray',\n",
    "                extent=(-3.2,3.2,-3.2,3.2),\n",
    "                origin=\"lower\",\n",
    "                alpha=1.\n",
    "            )\n",
    "        )\n",
    "    \n",
    "    # Done\n",
    "    if show:\n",
    "        plt.show()\n",
    "        return\n",
    "    return [pcm, sc, text1, fillpost2, fillpost1, lpost, ltrue] + images\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax_left, ax_mid, images_coords = paper_settings.animated_special()\n",
    "ax_images = []\n",
    "ims = [draw_frame(n) for n in range(total_frames)]\n",
    "ani = animation.ArtistAnimation(fig, ims, interval=200, blit=True)\n",
    "ani.save('../figures/live_inference_with_images.gif', dpi=300, writer='imagemagick')\n",
    "ani.save('../figures/live_inference_with_images_small.gif', dpi=120, writer='imagemagick')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax_left, ax_mid, images_coords = paper_settings.animated_special(flipped=True)\n",
    "ax_images = []\n",
    "ims = [draw_frame(n) for n in range(total_frames)]\n",
    "ani = animation.ArtistAnimation(fig, ims, interval=200, blit=True)\n",
    "ani.save('../figures/live_inference_with_images_reverse.gif', dpi=300, writer='imagemagick')\n",
    "ani.save('../figures/live_inference_with_images_reverse_small.gif', dpi=120, writer='imagemagick')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (lensing)",
   "language": "python",
   "name": "lensing"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}