results.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config Completer.use_jedi = False\n",
    "import numpy as np\n",
    "import pickle\n",
    "import pathlib\n",
    "import shedding\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from scipy import stats, optimize, integrate\n",
    "import glob\n",
    "from tqdm.notebook import tqdm\n",
    "from pathlib import Path\n",
    "import re\n",
    "import collections\n",
    "import itertools as it\n",
    "import textwrap\n",
    "\n",
    "mpl.style.use('scrartcl.mplstyle')\n",
    "\n",
    "\n",
    "def style_violins(violins, **kwargs):\n",
    "    \"\"\"\n",
    "    Apply keyword arguments to a violinplot.\n",
    "    \"\"\"\n",
    "    kwargs.setdefault('zorder', 1)\n",
    "    [plt.setp(value, **kwargs) for value in violins.values()]\n",
    "    \n",
    "    \n",
    "def evaluate_pcolormesh_edges(x, scale='linear'):\n",
    "    \"\"\"\n",
    "    Evaluate the edges of cells for a `pcolormesh` visualisation.\n",
    "    \"\"\"\n",
    "    if scale == 'log':\n",
    "        x = np.log(x)\n",
    "    elif scale != 'linear':\n",
    "        raise ValueError(scale)\n",
    "        \n",
    "    # Find the (n - 1) midpoints\n",
    "    midpoints = (x[1:] + x[:-1]) / 2\n",
    "    # Find the endpoints\n",
    "    left = 2 * x[0] - midpoints[0]\n",
    "    right = 2 * x[-1] - midpoints[-1]\n",
    "    # Construct the edges\n",
    "    edges = np.concatenate([[left], midpoints, [right]])\n",
    "    \n",
    "    if scale == 'log':\n",
    "        edges = np.exp(edges)\n",
    "    return edges\n",
    "\n",
    "def evaluate_mode(x, lin=200):\n",
    "    \"\"\"\n",
    "    Evaluate the mode of a univariate distribution using kernel density estimation.\n",
    "    \"\"\"\n",
    "    kde = stats.gaussian_kde(x)\n",
    "    if isinstance(lin, int):\n",
    "        lin = np.linspace(np.min(x), np.max(x), lin)\n",
    "    y = kde(lin)\n",
    "    return lin[np.argmax(y)]\n",
    "\n",
    "\n",
    "def alpha_cmap(color):\n",
    "    \"\"\"\n",
    "    Create a colormap which interpolates between transparent and the given color.\n",
    "    \"\"\"\n",
    "    if isinstance(color, int):\n",
    "        color = f'C{color}'\n",
    "    return mpl.colors.LinearSegmentedColormap.from_list(\"\", [\n",
    "        mpl.colors.to_rgba(color, alpha=0),\n",
    "        mpl.colors.to_rgba(color, alpha=1),\n",
    "    ])\n",
    "\n",
    "\n",
    "def binary_cmap(color1, color2):\n",
    "    return mpl.colors.LinearSegmentedColormap.from_list(\"\", [\n",
    "        color1,\n",
    "        color2,\n",
    "    ])\n",
    "\n",
    "\n",
    "def evaluate_hpd_levels(pxf, pvals):\n",
    "    \"\"\"\n",
    "    Evaluate the levels for given highest-proability-density regions. \n",
    "    Works for multimodal distributions.\n",
    "    \"\"\"\n",
    "    if isinstance(pvals, int):\n",
    "        pvals = (pvals - np.arange(pvals)) / (pvals + 1)\n",
    "    pvals = np.atleast_1d(pvals)\n",
    "    idx = np.argsort(-pxf.ravel())\n",
    "    cum = np.cumsum(pxf.ravel()[idx])\n",
    "    cum /= cum[-1]\n",
    "    j = np.argmax(cum[:, None] > pvals, axis=0)\n",
    "    return pxf.ravel()[idx][j]\n",
    "\n",
    "\n",
    "def evaluate_hpd_levels(pdf, pvals):\n",
    "    \"\"\"\n",
    "    Evaluate the levels that include a given fraction of the the probability mass.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    pdf : array_like\n",
    "        Probability density function evaluated over a mesh.\n",
    "    pvals : array_like or int\n",
    "        Probability mass to be included within the corresponding level or the number of levels.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    levels : array_like\n",
    "        Contour levels of the probability density function that enclose the desired probability\n",
    "        mass.\n",
    "    \"\"\"\n",
    "    # Obtain equidistant levels if only the number is given\n",
    "    if isinstance(pvals, int):\n",
    "        pvals = (pvals - np.arange(pvals)) / (pvals + 1)\n",
    "    pvals = np.atleast_1d(pvals)\n",
    "    # Sort the probability density and evaluate the normalised cumulative distribution.\n",
    "    # We aggregate identical pdf values so we can interpolate.\n",
    "    pdf, weights = np.unique(-pdf, return_counts=True)\n",
    "    pdf = - pdf\n",
    "    cum = integrate.cumtrapz(pdf * weights)\n",
    "    cum = np.concatenate([np.zeros(1), cum])\n",
    "    cum /= cum[-1]\n",
    "    # Find the first index that encloses more than the desired mass\n",
    "    js = np.argmax(cum[:, None] > pvals, axis=0)\n",
    "    levels = []\n",
    "    for j, pval in zip(js, pvals):\n",
    "        i = j - 1\n",
    "        # Get the upper and lower bounds and interpolate\n",
    "        y2 = cum[j]\n",
    "        y1 = cum[i]\n",
    "        x2 = pdf[j]\n",
    "        x1 = pdf[i]\n",
    "        slope = (y2 - y1) / (x2 - x1)\n",
    "        offset = y1 - slope * x1\n",
    "        level = (pval - offset) / slope\n",
    "        levels.append(level)\n",
    "\n",
    "    return np.asarray(levels)\n",
    "\n",
    "\n",
    "\n",
    "    \n",
    "colors_by_key = {\n",
    "    ('standard', 'constant'): 'C0',\n",
    "    ('inflated', 'constant'): 'C2',\n",
    "    ('standard', 'temporal'): 'C1',\n",
    "    ('inflated', 'temporal'): 'C3',\n",
    "}\n",
    "# Add shorthand including the general key\n",
    "colors_by_key.update({\n",
    "    ('general',) + key: value for key, value in colors_by_key.items()\n",
    "})\n",
    "\n",
    "markers_by_key = {\n",
    "    ('standard', 'constant'): 'o',\n",
    "    ('inflated', 'constant'): 's',\n",
    "    ('standard', 'temporal'): 'v',\n",
    "    ('inflated', 'temporal'): 'D',\n",
    "}\n",
    "# Add shorthand including the general key\n",
    "markers_by_key.update({\n",
    "    ('general',) + key: value for key, value in markers_by_key.items()\n",
    "})\n",
    "\n",
    "violin_widths = 0.7\n",
    "log10formatter = mpl.ticker.FuncFormatter(lambda x, _: f'$10^{{{x if x % 1 else int(x)}}}$')\n",
    "\n",
    "\n",
    "# labels\n",
    "population_shape = 'Q'\n",
    "population_scale = 'S'\n",
    "patient_shape = 'q'\n",
    "patient_scale = r'\\sigma'\n",
    "\n",
    "shape_label = {\n",
    "    'population': population_shape,\n",
    "    'patient': patient_shape,\n",
    "}\n",
    "scale_label = {\n",
    "    'population': population_scale,\n",
    "    'patient': patient_scale,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Illustration for the model structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DistHandler:\n",
    "    \"\"\"\n",
    "    Generate a distribution legend handle.\n",
    "    \"\"\"\n",
    "    def legend_artist(self, legend, orig_handle, fontsize, handlebox):\n",
    "        color = orig_handle.get_color()\n",
    "        x0, y0 = handlebox.xdescent, handlebox.ydescent\n",
    "        width, height = handlebox.width, handlebox.height\n",
    "        scale = width / 4\n",
    "        lin = x0 + np.linspace(0, width)\n",
    "        z = (lin - (x0 + width / 2)) / scale\n",
    "        pdf = height * np.exp(- z ** 2 / 2)\n",
    "        \n",
    "        patch = mpl.lines.Line2D(lin, pdf, color=color, ls=orig_handle.get_linestyle())\n",
    "        handlebox.add_artist(patch)\n",
    "        \n",
    "        xy = [(x0, y0)]\n",
    "        xy.extend(zip(lin, pdf))\n",
    "        xy.append((x0 + width, y0))\n",
    "        patch = mpl.patches.Polygon(xy, color=color, alpha=alpha)\n",
    "        handlebox.add_artist(patch)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "np.random.seed(8)\n",
    "\n",
    "num_patients = 3\n",
    "time = np.linspace(0, 25, 200)\n",
    "profile_func = lambda x: 1e4 * x * np.exp(-x / 2)\n",
    "profile = profile_func(time)\n",
    "earliest = 6\n",
    "alpha = 0.25\n",
    "loq = 1e-1\n",
    "ax.axvspan(-1, earliest, facecolor='gray', alpha=.15)\n",
    "mpl.rcParams['hatch.color'] = mpl.colors.to_rgba('C0', alpha)\n",
    "mpl.rcParams['hatch.linewidth'] = 8\n",
    "\n",
    "offset_scale = 3\n",
    "offsets =  np.random.normal(0, offset_scale, num_patients)\n",
    "offsets = np.exp(offsets - np.mean(offsets))\n",
    "focus_patient = np.argmin(offsets)\n",
    "\n",
    "for i in np.arange(num_patients):\n",
    "    ls = '-' if i == focus_patient else ':'\n",
    "    offset = offset_scale * np.exp(i)\n",
    "    offset = offsets[i]\n",
    "    fltr = time <= earliest\n",
    "    ax.plot(time[fltr], profile[fltr] * offset, color='gray', ls=ls)\n",
    "    fltr = time >= earliest\n",
    "    ax.plot(time[fltr], profile[fltr] * offset, color='black', ls=ls)\n",
    "\n",
    "ax.scatter(np.ones(num_patients) * earliest, offsets * profile_func(earliest), zorder=9, marker='s')\n",
    "\n",
    "lin = np.linspace(-3 * offset_scale, 3 * offset_scale)\n",
    "y = np.exp(lin) * profile_func(earliest)\n",
    "left = earliest - 10 * offset_scale * stats.norm.pdf(lin, 0, offset_scale)\n",
    "right = earliest * np.ones_like(lin)\n",
    "fill = ax.fill_betweenx(y, left, right, alpha=alpha, hatch='/')\n",
    "fill.set_facecolor((1, 1, 1, 0))\n",
    "ax.plot(left, y, ls='--')\n",
    "\n",
    "sample_scale = 1\n",
    "fltr = time >= earliest\n",
    "y = offsets[focus_patient] * profile_func(time[fltr])\n",
    "ax.fill_between(time[fltr], y / np.exp(sample_scale), y * np.exp(sample_scale), color='C1', alpha=alpha)\n",
    "ax.fill_between(time[fltr], y / np.exp(2 * sample_scale), y * np.exp(2 * sample_scale), color='C1', alpha=alpha)\n",
    "\n",
    "lin = np.linspace(-3 * sample_scale, 3 * sample_scale)\n",
    "y = np.exp(lin) * profile_func(earliest) * offsets[focus_patient]\n",
    "right = earliest + 10 * sample_scale * stats.norm.pdf(lin, 0, sample_scale)\n",
    "left = earliest * np.ones_like(lin)\n",
    "fill = ax.fill_betweenx(y, left, right, color='C1', alpha=alpha)\n",
    "ax.plot(right, y)\n",
    "\n",
    "num_samples = 5\n",
    "sample_offsets = np.random.normal(0, sample_scale, num_samples)\n",
    "sample_offsets = np.exp(sample_offsets - np.mean(sample_offsets))\n",
    "sample_times = np.random.randint(earliest, time.max() + 1, num_samples)\n",
    "sample_values = profile_func(sample_times) * sample_offsets * offsets[focus_patient]\n",
    "fltr = sample_values > loq\n",
    "ax.scatter(sample_times[fltr], sample_values[fltr], zorder=9, marker='o', color='C1')\n",
    "ax.scatter(sample_times[~fltr], sample_values[~fltr], zorder=9, marker='X', color='C1')\n",
    "ax.axhline(loq, ls='--', color='k')\n",
    "\n",
    "ax.set_yscale('log')\n",
    "ax.set_xlim(-1)\n",
    "pop_handle = mpl.lines.Line2D([], [], color='C0', ls='--')\n",
    "patient_handle = mpl.lines.Line2D([], [], color='C1')\n",
    "handles_labels = [\n",
    "    [pop_handle, 'population-level distribution, describing\\nvariation in shedding between patients'],\n",
    "    [mpl.lines.Line2D([], [], color='C0', ls='none', marker='s'), \n",
    "     r'patient location parameter $\\mu$, describing the''\\namplitude of individual shedding profiles'],\n",
    "    \n",
    "    [patient_handle, 'patient-level distribution, describing variation\\nbetween samples from the same patient'],\n",
    "    [(mpl.lines.Line2D([], [], color='C1', ls='none', marker='o'),\n",
    "      mpl.lines.Line2D([], [], color='C1', ls='none', marker='X')), r'sample-level RNA load $y$ quantified by''\\nRT-qPCR assays'],\n",
    "    [mpl.lines.Line2D([], [], color='k'), r'shedding profile $g(t)$, describing the''\\nevolution of RNA loads'],\n",
    "    [mpl.lines.Line2D([], [], color='k', ls='--'), r'limit of quantification $\\theta$'],\n",
    "]\n",
    "\n",
    "ax.legend(*zip(*handles_labels), loc='best', fontsize='small', handler_map={\n",
    "    tuple: mpl.legend_handler.HandlerTuple(None),\n",
    "    pop_handle: DistHandler(),\n",
    "    patient_handle: DistHandler(),\n",
    "}, edgecolor='none')\n",
    "plt.setp(ax.xaxis.get_ticklabels(), visible=False)\n",
    "plt.setp(ax.yaxis.get_ticklabels(), visible=False)\n",
    "ax.set_xlabel('Days past symptom onset $t$')\n",
    "ax.set_ylabel(r'$\\log$ SARS-CoV-2 RNA load')\n",
    "\n",
    "pos = 5e-3\n",
    "kwargs = {\n",
    "    'arrowstyle': 'simple,tail_width=45,head_width=55,head_length=20',\n",
    "    'zorder': 9,\n",
    "    'shrinkA': 0,\n",
    "    'shrinkB': 0,\n",
    "    'edgecolor': 'none',\n",
    "}\n",
    "patch = mpl.patches.FancyArrowPatch((earliest, pos), (earliest - 6.5, pos),\n",
    "                                    facecolor='w', **kwargs)\n",
    "ax.add_artist(patch)\n",
    "patch = mpl.patches.FancyArrowPatch((earliest, pos), (earliest + 6.5, pos), facecolor='gray', \n",
    "                                    alpha=.15, **kwargs)\n",
    "ax.add_artist(patch)\n",
    "\n",
    "kwargs = dict(va='center', fontsize='small', zorder=10)\n",
    "ax.text(earliest - 0.5, pos, 'early shedding;\\ncurrently uncon-\\nstrained by data', ha='right', **kwargs)\n",
    "ax.text(earliest + 0.5, pos, 'late shedding;\\nconsistent with ex-\\nponential profile', ha='left', **kwargs)\n",
    "fig.tight_layout()\n",
    "fig.savefig('workspace/figures/model.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evidences for the different models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load all the results for a given seed.\n",
    "seed = 0\n",
    "filenames = glob.glob(f'workspace/*-*-*-{seed}*/polychord/result.pkl')\n",
    "print(f'found {len(filenames)} results\\n')\n",
    "results = {}\n",
    "evidences = {}\n",
    "\n",
    "for filename in sorted(filenames):\n",
    "    with open(filename, 'rb') as fp:\n",
    "        result = pickle.load(fp)\n",
    "        model = result['model']\n",
    "        key = (model.parametrisation.value, \n",
    "               'inflated' if model.inflated else 'standard', \n",
    "               'temporal' if model.temporal.value else 'constant')\n",
    "        \n",
    "        if result['args'].evidence:\n",
    "            x, err = result['evidence']\n",
    "            print(f'model: {key}; evidence: {x:.2f} +- {err:.2f}')\n",
    "            evidences[key] = result['evidence']\n",
    "        \n",
    "        if result['args'].evidence and model.temporal == shedding.Profile.CONSTANT:\n",
    "            key = key + ('evidence-calculation',)\n",
    "        results[key] = result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate summary statistics for the mean shedding rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for key, result in results.items():\n",
    "    lines = []\n",
    "    \n",
    "    samples = result['samples']\n",
    "    means = np.asarray([\n",
    "        shedding.gengamma_mean(x['population_shape'], x['population_loc'], x['population_scale'])\n",
    "        for x in shedding.transpose_samples(samples)\n",
    "    ])\n",
    "    \n",
    "    mode = evaluate_mode(np.ravel(np.log10(means)))\n",
    "    lines.append(f'mode: {mode:.3f}')\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    kde = stats.gaussian_kde(np.log10(means))\n",
    "    xmin = kde.dataset.min()\n",
    "    xmax = kde.dataset.max()\n",
    "    xrng = xmax - xmin\n",
    "    lin = np.linspace(xmin - 0.5 * xrng, xmax + 0.5 * xrng, 500)\n",
    "    pdf = kde(lin)\n",
    "    plt.plot(lin, pdf)\n",
    "    level = evaluate_hpd_levels(pdf, 0.95)\n",
    "    plt.axhline(level, color='C1', ls='--')\n",
    "    plt.scatter(mode, kde(mode), label='mode', marker='o')\n",
    "    plt.axvline(mode, ls='--')\n",
    "    plt.title(str(key))\n",
    "    \n",
    "    # Evaluate the two points at which things intersect\n",
    "    x0s = np.percentile(kde.dataset, [2.5, 97.5])\n",
    "    lims = [optimize.minimize(lambda x: (kde(x) - level) ** 2, x0)\n",
    "            for x0 in x0s]\n",
    "    lims = np.asarray([lim.x[0] for lim in lims])\n",
    "    \n",
    "    plt.axvspan(*lims, color='C1', alpha=.25, label='95% HPD')\n",
    "\n",
    "    lines.append(f'95% HPD: {lims[0]:.3f} -- {lims[1]:.3f}')\n",
    "    plt.text(0.95, 0.95, '\\n'.join(lines), transform=ax.transAxes, va='top', ha='right')\n",
    "    plt.legend(loc='upper left')\n",
    "    \n",
    "    print('\\n'.join([str(key)] + lines + ['']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Faecal shedding profile and halflife"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the datasets.\n",
    "datasets = shedding.load_datasets([\n",
    "    'Woelfel2020', \n",
    "    'Lui2020', \n",
    "    'Wang2020',\n",
    "    'Han2020',\n",
    "], 'publications/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax_days, ax) = plt.subplots(2, 1, sharex=True, gridspec_kw={'height_ratios': [1, 5]})\n",
    "\n",
    "key = ('general', 'standard', 'temporal')\n",
    "result = results[key]\n",
    "samples = result['samples']\n",
    "data = result['data']\n",
    "\n",
    "# Add the shaded region\n",
    "ax.axvspan(0, 5, facecolor='#eee', zorder=0)\n",
    "\n",
    "# Plot the fit\n",
    "lin = np.linspace(0, 40)\n",
    "ys = np.exp(samples['population_loc'][:, None] + samples['slope'][:, None] * lin)\n",
    "ax.fill_between(lin, *np.percentile(ys, [2.5, 97.5], axis=0), color='black', zorder=0.1, alpha=.1)\n",
    "ax.plot(lin, np.median(ys, axis=0), color='k', label=r'shedding profile', \n",
    "                zorder=0)\n",
    "\n",
    "x = data['day']\n",
    "y = np.where(data['positive'], data['load'], 50)\n",
    "colors_by_dataset = {}\n",
    "for i, key in enumerate(set(data['dataset'])):\n",
    "    color = f'C{i}'\n",
    "    colors_by_dataset[key] = color\n",
    "    ax.axhline(10 ** datasets[key]['loq'], color=color, ls='--', zorder=0)\n",
    "    for p, m in zip([True, False], 'ox'):\n",
    "        label = datasets[key]['key'] if p else None\n",
    "        fltr = (data['dataset'] == key) & (data['positive'] == p)\n",
    "        offset = (i - 1) / 3\n",
    "        offset = 0\n",
    "        pts = ax.scatter(x[fltr] + offset, y[fltr], alpha=.75, marker=m, color=color, label=label, s=10)\n",
    "\n",
    "if False:\n",
    "    for i in np.arange(data['num_patients']):\n",
    "        f = data['idx'] == i\n",
    "        key = data['dataset'][f][0]\n",
    "        ax.plot(x[f], y[f], color=colors_by_dataset[key], alpha=.5)\n",
    "        \n",
    "pos = 4e5\n",
    "kwargs = {\n",
    "    'arrowstyle': 'simple,tail_width=40,head_width=50,head_length=15',\n",
    "    'zorder': 0,\n",
    "    'shrinkA': 0,\n",
    "    'shrinkB': 0,\n",
    "    'edgecolor': 'none',\n",
    "}\n",
    "patch = mpl.patches.FancyArrowPatch((5, pos), (.25, pos),\n",
    "                                    facecolor='w', **kwargs)\n",
    "ax.add_artist(patch)\n",
    "\n",
    "kwargs = dict(va='center', fontsize='x-small', zorder=10)\n",
    "ax.text(4.5, pos, 'uncon-\\nstrained\\nearly\\nshedding', ha='right', **kwargs)\n",
    "\n",
    "\n",
    "ax.set_yscale('log')\n",
    "legend = ax.legend(ncol=2, loc='upper left', fontsize='x-small')\n",
    "ax.set_ylim(20, 1e9)\n",
    "ax.set_xlim(0, 36)\n",
    "ax.set_xlabel('Days past symptom onset $t$')\n",
    "ax.set_ylabel('Gene copies per mL')\n",
    "\n",
    "child = ax.inset_axes((.65, .5, .33, .48))\n",
    "halflife = - np.log(2) / samples['slope'].ravel()\n",
    "shedding.plot_kde(halflife, ax=child, xmax=50/24, numlin=100, color='k')\n",
    "child.yaxis.set_ticks([])\n",
    "child.set_ylabel(r'Posterior $P(\\tau_½)$')\n",
    "child.set_xlabel(r'Profile half-life $\\tau_½$ (days)')\n",
    "#child.xaxis.set_ticks([24, 36, 48])\n",
    "# child.axvline(evaluate_mode(halflife), ls='--', color='k')\n",
    "\n",
    "# Plot the histogram\n",
    "days_per_bin = 1\n",
    "ax_days.hist(data['day'], range=(-.5, 35.5), bins=36 // days_per_bin, \n",
    "             weights=np.ones_like(data['day']) / days_per_bin, color='#aaa')\n",
    "ax_days.set_ylabel('Samples\\nper day')\n",
    "\n",
    "\n",
    "child.text(0.05, 0.95, '(c)', transform=child.transAxes, ha='left', va='top')\n",
    "ax_days.text(0.02, 0.9, '(a)', transform=ax_days.transAxes, ha='left', va='top')\n",
    "ax.text(0.02, 0.03, '(b)', transform=ax.transAxes, ha='left', va='bottom')\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('workspace/figures/decay.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimates of halflife"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for inflated in ['standard', 'inflated']:\n",
    "    key = ('general', inflated, 'temporal')\n",
    "    result = results[key]\n",
    "    halflife = - np.log(2) / result['samples']['slope'] * 24\n",
    "    \n",
    "    lin = np.linspace(0, 3 * 24, 5001)\n",
    "    pdf = stats.gaussian_kde(halflife)(lin)\n",
    "    plt.figure()\n",
    "    plt.plot(lin, pdf)\n",
    "    \n",
    "    level = evaluate_hpd_levels(pdf, 0.95)\n",
    "    plt.axhline(level)\n",
    "    delta = np.diff(level > pdf)\n",
    "    idx, = np.nonzero(delta)\n",
    "    plt.title(inflated)\n",
    "    \n",
    "    \n",
    "    print('model', inflated)\n",
    "    print('mode', lin[np.argmax(pdf)])\n",
    "    print('95% hpd', lin[idx])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shedding prevalence estimates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for temporal in ['temporal', 'constant']:\n",
    "    plt.figure()\n",
    "    plt.title(temporal)\n",
    "    key = ('general', 'inflated', temporal)\n",
    "    result = results[key]\n",
    "    rho = result['samples']['rho']\n",
    "    \n",
    "    # Evaluate a \"reflected\" kernel density estimate to account for the boundary.\n",
    "    lin = np.linspace(0, 2, 2000)\n",
    "    pdf = stats.gaussian_kde(rho)(lin)\n",
    "    a, b = np.array_split(pdf, 2)\n",
    "    pdf = a + b[::-1]\n",
    "    plt.hist(rho, bins=20, density=True, alpha=.5)\n",
    "    plt.plot(lin[:len(pdf)], pdf)\n",
    "    \n",
    "    print('model', temporal)\n",
    "    print('mode', lin[np.argmax(pdf)])\n",
    "    \n",
    "    level = evaluate_hpd_levels(pdf, 0.95)\n",
    "    plt.axhline(level)\n",
    "    delta = np.diff(level > pdf)\n",
    "    idx, = np.nonzero(delta)\n",
    "    print('95% HPD', lin[idx])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shedding profile comparison (gamma/Teunis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparison for different shedding profiles\n",
    "lin = np.linspace(-15, 30, 200)\n",
    "np.random.seed(0)\n",
    "\n",
    "with open('workspace/profile-gamma-0/result.pkl', 'rb') as fp:\n",
    "    result = pickle.load(fp)\n",
    "    \n",
    "samples = samples_gamma = result['samples']\n",
    "offsets_gamma = samples['profile_offset'][:, None]\n",
    "dt = lin - offsets_gamma\n",
    "profiles_gamma = np.exp(samples['population_loc'][:, None] - samples['profile_scale'][:, None] * dt) * \\\n",
    "    dt ** samples['profile_shape'][:, None]\n",
    "    \n",
    "with open('workspace/profile-teunis-0/result.pkl', 'rb') as fp:\n",
    "    result = pickle.load(fp)\n",
    "    \n",
    "samples = samples_teunis = result['samples']\n",
    "offsets_teunis = samples['profile_offset'][:, None]\n",
    "dt = lin - offsets_teunis\n",
    "profiles_teunis = np.exp(samples['population_loc'][:, None] - samples['profile_decay'][:, None] * dt) * \\\n",
    "    (1 - np.exp(-samples['profile_rise'][:, None] * dt))\n",
    "profiles_teunis = np.where(profiles_teunis < .1, np.nan, profiles_teunis)\n",
    "\n",
    "result = results[('general', 'standard', 'temporal')]\n",
    "samples = result['samples']\n",
    "profiles_exponential = np.exp(samples['population_loc'][:, None] + samples['slope'][:, None] * lin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, sharex=True, sharey='row', gridspec_kw={'height_ratios': [2, 5]})\n",
    "\n",
    "top_axes, bottom_axes = axes\n",
    "\n",
    "for ax, key in zip(axes.ravel(), 'abcd'):\n",
    "    ax.axvspan(lin.min(), 5, facecolor='#eee')\n",
    "    ax.text(0.95, 0.95, f'({key})', ha='right', va='top', transform=ax.transAxes)\n",
    "\n",
    "labels = ['gamma', r'Teunis $et\\ al.$']\n",
    "for ax, profiles, label in zip(bottom_axes, [profiles_gamma, profiles_teunis], labels):\n",
    "    idx = np.random.choice(len(profiles), 500, False)\n",
    "    ax.plot(lin, profiles[idx].T, color='C0', alpha=.1)\n",
    "    line, = ax.plot(lin, np.median(profiles_exponential, axis=0), color='C1')\n",
    "    bounds = np.percentile(profiles_exponential, [2.5, 97.5], axis=0)\n",
    "    ax.plot(lin, np.transpose(bounds), color=line.get_color(), ls='--')\n",
    "    ax.fill_between(lin, *bounds, color=line.get_color(), alpha=.2, zorder=9)\n",
    "    ax.set_xlabel('Days past symptom onset $t$')\n",
    "    handles_labels = [\n",
    "        (mpl.lines.Line2D([], [], color='C0'), f'{label} profile'),\n",
    "        (mpl.lines.Line2D([], [], color='C1'), 'exponential profile'),\n",
    "    ]\n",
    "    ax.legend(*zip(*handles_labels), loc='lower left', fontsize=8)\n",
    "    \n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(10, 1e11)\n",
    "ax.set_xlim(lin.min(), lin.max())\n",
    "axes[0, 0].set_ylabel(r'Posterior $P(t_\\mathrm{peak})$')\n",
    "axes[1, 0].set_ylabel('Gene copies per mL')\n",
    "\n",
    "mode_gamma = samples_gamma['profile_offset'] + samples_gamma['profile_shape'] / samples_gamma['profile_scale']\n",
    "mode_teunis = samples_teunis['profile_offset'] + np.log((samples_teunis['profile_rise'] + samples_teunis['profile_decay']) /\n",
    "                                                samples_teunis['profile_decay']) / samples_teunis['profile_rise']\n",
    "\n",
    "for ax, offsets, label in zip(top_axes, [mode_gamma, mode_teunis], labels):\n",
    "    ax.hist(offsets.ravel(), range=(-14, 7), bins=21, density=True, \n",
    "            label='peak shedding\\ntime 'r'$t_\\mathrm{peak}$')\n",
    "    ax.legend(loc='center right', fontsize=8)\n",
    "    # sb.kdeplot(offsets.ravel(), ax=ax, legend=False)\n",
    "    \n",
    "axes[0, 0].set_title('Gamma profile')\n",
    "axes[0, 1].set_title(r'Teunis $et\\ al.$ profile')\n",
    "\n",
    "if False:\n",
    "    for ax in [ax1, ax2]:\n",
    "        pos = 50\n",
    "        length = 17\n",
    "        kwargs = {\n",
    "            'arrowstyle': 'simple,tail_width=30,head_width=35,head_length=20',\n",
    "            'zorder': 9,\n",
    "            'shrinkA': 0,\n",
    "            'shrinkB': 0,\n",
    "            'edgecolor': 'none',\n",
    "        }\n",
    "        patch = mpl.patches.FancyArrowPatch((5, pos), (5 - length, pos),\n",
    "                                            facecolor='w', **kwargs)\n",
    "        ax.add_artist(patch)\n",
    "        patch = mpl.patches.FancyArrowPatch((5, pos), (5 + length, pos), facecolor='gray', \n",
    "                                            alpha=.15, **kwargs)\n",
    "        ax.add_artist(patch)\n",
    "\n",
    "        kwargs = dict(va='center', fontsize='x-small', zorder=10)\n",
    "        ax.text(5 - 0.5, pos, 'early shedding;\\ncurrently uncon-\\nstrained by data', ha='right', **kwargs)\n",
    "        ax.text(5 + 0.5, pos, 'late shedding;\\nconsistent with ex-\\nponential profile', ha='left', **kwargs)\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('workspace/figures/profiles.pdf')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Replicates and prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "replicates_by_key = {}\n",
    "predictions_by_study_and_key = {}\n",
    "num_samples = None\n",
    "\n",
    "prediction_hyperparameters = {\n",
    "    'Kim et al.': {\n",
    "        'num_patients': 38,\n",
    "        'num_samples': 129,\n",
    "        'loq': 1,  # loq is irrelevant because we only care about the maximum\n",
    "        'min_samples_per_patient': 1,\n",
    "    },\n",
    "    'Ng et al.': {\n",
    "        'num_patients': 21,\n",
    "        'num_samples': 81,\n",
    "        'loq': 10 ** 2.5403294748,\n",
    "        'min_samples_per_patient': 1,\n",
    "    },\n",
    "}\n",
    "\n",
    "for key, result in tqdm(results.items()):\n",
    "    # Skip the one we just used for evidence calculations\n",
    "    if len(key) > 3:\n",
    "        continue\n",
    "        \n",
    "    # Potentially filter the samples\n",
    "    samples = {key: value for key, value in result['samples'].items()}\n",
    "    ys = shedding.transpose_samples(samples)\n",
    "    if num_samples:\n",
    "        ys = [ys[i] for i in np.random.choice(len(ys), num_samples, False)]\n",
    "    model = result['model']\n",
    "    \n",
    "    # Replicate the data we have fit to for posterior predictive checks\n",
    "    data = result['data']\n",
    "    replicates_by_key[key] = model.simulate(ys, data, 'existing_patients')\n",
    "    \n",
    "    if key[-1] == 'temporal':\n",
    "        continue\n",
    "        \n",
    "    # Predict data for unobserved studies for external validation\n",
    "    with tqdm(total=len(ys) * len(prediction_hyperparameters), desc=str(key)) as progress:\n",
    "        for study, hyperparameters in prediction_hyperparameters.items():\n",
    "            for y in ys:\n",
    "                # Sample the number of samples per patient (respecting the minimum number of samples per patient)\n",
    "                samples_to_distribute = hyperparameters['num_samples'] - \\\n",
    "                    hyperparameters['min_samples_per_patient'] * hyperparameters['num_patients']\n",
    "                if samples_to_distribute < 0:\n",
    "                    raise ValueError\n",
    "                p = np.ones(hyperparameters['num_patients']) / hyperparameters['num_patients']\n",
    "                num_samples_by_patient = np.ones(hyperparameters['num_patients'], int) * \\\n",
    "                    hyperparameters['min_samples_per_patient'] + np.random.multinomial(samples_to_distribute, p)\n",
    "                # Create a lookup index\n",
    "                idx = np.repeat(1 + np.arange(hyperparameters['num_patients']), num_samples_by_patient)\n",
    "                # Generate a prediction and add it to the list\n",
    "                prediction_data = {\n",
    "                    **hyperparameters, \n",
    "                    'num_samples_by_patient': num_samples_by_patient,\n",
    "                    'idx': idx,\n",
    "                    'day': 0,\n",
    "                }\n",
    "                prediction_data = model.simulate(y, prediction_data, 'new_patients')\n",
    "                predictions_by_study_and_key.setdefault(study, {}).setdefault(key, []).append(\n",
    "                    prediction_data\n",
    "                )\n",
    "                progress.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predictions for held-out-datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_violin(data_sequence, target, x, color, ax, \n",
    "                reference=None, pax=None, **kwargs):\n",
    "    kwargs.setdefault('showextrema', False)\n",
    "    kwargs.setdefault('widths', violin_widths)\n",
    "    label = kwargs.pop('label', None)\n",
    "    marker = kwargs.pop('marker', 'o')\n",
    "    \n",
    "    values = np.fromiter(map(target, data_sequence), float)\n",
    "    violins = ax.violinplot(values[:, None], [x], **kwargs)\n",
    "    style_violins(violins, color=color)\n",
    "    mode = evaluate_mode(values, kwargs.get('points', 200))\n",
    "    ax.scatter(x, mode, marker=marker, color=color, \n",
    "               zorder=9, label=label, edgecolor='w')\n",
    "\n",
    "    if isinstance(reference, dict):\n",
    "        reference = target(reference)\n",
    "    if reference and pax:\n",
    "        pval = min(\n",
    "            np.mean(reference < values),\n",
    "            np.mean(reference > values),\n",
    "        )\n",
    "        # Correct for equality\n",
    "        pval += np.mean(reference == values) / 2\n",
    "        pax.scatter(x, pval, marker=marker, color=color, edgecolor='w')\n",
    "        \n",
    "    # Restore the edges\n",
    "    for violin in violins['bodies']:\n",
    "        facecolor = violin.get_facecolor()\n",
    "        violin.set_alpha(None)\n",
    "        violin.set_facecolor(facecolor)\n",
    "        violin.set_edgecolor(facecolor[:, :3])\n",
    "        \n",
    "    return violins\n",
    "pad = 0.6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "gs = fig.add_gridspec(1, 3)\n",
    "ax1 = fig.add_subplot(gs[0, 0])\n",
    "ax2 = fig.add_subplot(gs[0, 1], sharey=ax1)\n",
    "ax3 = fig.add_subplot(gs[0, 2])\n",
    "\n",
    "# Plot the replicated sample means\n",
    "references_by_study = {\n",
    "    ('Ng et al.', 'max'): 7.1,\n",
    "    ('Kim et al.', 'max'): 7.4363217001,\n",
    "}\n",
    "studies = ['Kim et al.', 'Ng et al.']\n",
    "labels = [\n",
    "    r'$\\bar{x}^\\mathrm{rep}$''\\nconstant',\n",
    "    r'$\\bar{x}^\\mathrm{rep}$''\\ntemporal',\n",
    "    r'$\\max\\,x^\\mathrm{pred}$''\\nKim et al.',\n",
    "    r'$\\max\\,x^\\mathrm{pred}$''\\nNg et al.',\n",
    "]\n",
    "\n",
    "\n",
    "for i, inflated in enumerate(['standard', 'inflated']):\n",
    "    j = 2\n",
    "        \n",
    "    def target(data):\n",
    "        value = np.log10(data['load'].max())\n",
    "        return value if np.isfinite(value) else 1\n",
    "    \n",
    "    for ax, study in zip([ax1, ax2], studies):\n",
    "        j = 0\n",
    "        reference = references_by_study[(study, 'max')]\n",
    "        if i == 0:\n",
    "            ax.plot((2 * j - violin_widths * pad, 2 * j + 1 + violin_widths * pad), \n",
    "                     (reference, reference), color='k', ls='--')\n",
    "        \n",
    "        key = ('general', inflated, 'constant')\n",
    "        predictions_by_key = predictions_by_study_and_key[study]\n",
    "        color = colors_by_key[key]\n",
    "        plot_violin(predictions_by_key[key], target, i + 2 * j, color, ax, \n",
    "                    marker=markers_by_key[key])\n",
    "        \n",
    "        j += 1\n",
    "        \n",
    "    # Add the median number of positive samples for Ng et al.\n",
    "    predictions_by_key = predictions_by_study_and_key['Ng et al.']\n",
    "    target = lambda data: np.median(data['num_positives_by_patient'])\n",
    "    if i == 0:\n",
    "        reference = 2\n",
    "        ax3.plot((- violin_widths * pad, 1 + violin_widths * pad), \n",
    "                 (reference, reference), color='k', ls='--')\n",
    "\n",
    "    key = ('general', inflated, 'constant')\n",
    "    color = colors_by_key[key]\n",
    "    values = list(map(target, predictions_by_key[key]))\n",
    "    points = int((max(values) - min(values)) / .5 + 1)\n",
    "    plot_violin(predictions_by_key[key], target, i, color, ax3, \n",
    "                bw_method=1e-9, points=points, marker=markers_by_key[key],\n",
    "                label='subpopu-\\nlation' if inflated == 'inflated' else 'standard')\n",
    "         \n",
    "ax1.yaxis.set_major_formatter(log10formatter)\n",
    "ax1.set_ylim(5, 11)\n",
    "ax1.set_ylabel(r'Maximum $\\max\\,x^\\mathrm{pred}$ (gene copies per mL)')\n",
    "ax1.set_xlabel('Kim et al.')\n",
    "\n",
    "ax2.set_ylabel(r'Maximum $\\max\\,x^\\mathrm{pred}$ (gene copies per mL)')\n",
    "ax2.set_xlabel('Ng et al.')\n",
    "\n",
    "ax3.set_ylabel('Median number of positive samples\\nper patient'r' $\\mathrm{median}\\,m_{\\bullet(+)}^\\mathrm{pred}$')\n",
    "ax3.set_yticks(np.arange(5))\n",
    "ax3.set_xlabel('Ng et al.')\n",
    "ax3.legend(loc='lower center', handletextpad=.25)\n",
    "ax3.set_ylim(-.85)\n",
    "\n",
    "for ax, label in zip([ax1, ax2, ax3], '(a) (b) (c)'.split()):\n",
    "    ax.text(0.05, 0.975, label, ha='left', va='top', transform=ax.transAxes)\n",
    "    ax.xaxis.set_ticks([])\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('workspace/figures/prediction.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "step = 1\n",
    "# fig, axes = plt.subplots(2, 2, sharex='col', sharey=True)\n",
    "fig = plt.figure()\n",
    "gs = fig.add_gridspec(1, 2, width_ratios=[2, 1])\n",
    "gsl = gs[0, 0].subgridspec(2, 2, wspace=.2)\n",
    "ax1 = fig.add_subplot(gsl[0, 0])\n",
    "ax2 = fig.add_subplot(gsl[0, 1], sharey=ax1)\n",
    "ax3 = fig.add_subplot(gsl[1, 0], sharex=ax1, sharey=ax1)\n",
    "ax4 = fig.add_subplot(gsl[1, 1], sharex=ax2, sharey=ax1)\n",
    "axes = np.asarray([\n",
    "    [ax1, ax2],\n",
    "    [ax3, ax4],\n",
    "])\n",
    "\n",
    "i = 0\n",
    "\n",
    "keys = it.product(['standard', 'inflated'], ['constant', 'temporal'])\n",
    "for ax, key in zip(np.ravel(axes), keys):\n",
    "    key = ('general',) + key\n",
    "    replicates = replicates_by_key[key][::step]\n",
    "    shedders = [np.sum(replicate['num_positives_by_patient'] > 0) for replicate in replicates]\n",
    "    positive = [np.sum(replicate['positive']) for replicate in replicates]\n",
    "\n",
    "    if key[2] == 'temporal':\n",
    "        rng = [\n",
    "            [14.5, 23.5],\n",
    "            [82.5, 111.5],\n",
    "        ]\n",
    "    else:\n",
    "        rng = [\n",
    "            [20.5, 34.5],\n",
    "            [82.5, 123.5],\n",
    "        ]\n",
    "    steps = [1, 1]\n",
    "    bins = [int((b - a) / step) for (a, b), step in zip(rng, steps)]\n",
    "    z, x, y = np.histogram2d(shedders, positive, bins=bins, range=rng, density=True)\n",
    "    x = (x[1:] + x[:-1]) / 2\n",
    "    y = (y[1:] + y[:-1]) / 2\n",
    "    # ax.contour(x, y, z.T, cmap=alpha_cmap(colors_by_key[key]))\n",
    "    ax.imshow(z.T, aspect='auto', extent=np.ravel(rng), origin='lower', cmap=alpha_cmap(colors_by_key[key]))\n",
    "\n",
    "    # sb.kdeplot(x=shedders, y=positive, cmap=alpha_cmap(colors_by_key[key]), ax=ax)\n",
    "    \n",
    "    # Plot the actual value\n",
    "    data = results[key]['data']\n",
    "    ax.scatter(np.sum(data['num_positives_by_patient'] > 0), \n",
    "               np.sum(data['positive']), \n",
    "               marker='x', color='k', zorder=9)\n",
    "    \n",
    "    # Plot the 95% credible contour\n",
    "    alphas = [.95, -np.expm1(-1)]\n",
    "    levels = evaluate_hpd_levels(z, alphas)\n",
    "    cs = ax.contour(x, y, z.T, \n",
    "                    levels=levels, colors='k', linestyles=[':', '--'])\n",
    "    labels = ax.clabel(cs, fmt={level: f'{100 * alpha:.0f}%' for level, alpha in zip(levels, alphas)})\n",
    "    \n",
    "    label = key[2] + (\" sub-\\npopulation\" if key[1] == \"inflated\" else \"\\nstandard\")\n",
    "    handle = mpl.lines.Line2D([], [], color=colors_by_key[key], marker=markers_by_key[key], ls='none')\n",
    "    handle.set_markeredgecolor('w')\n",
    "    ax.scatter(0.9, 0.1, color=colors_by_key[key], marker=markers_by_key[key], transform=ax.transAxes,\n",
    "               edgecolor='w')\n",
    "    # ax.legend([handle], [label], loc='upper left', frameon=False, handletextpad=.25)\n",
    "    \n",
    "    i += 1\n",
    "    \n",
    "ax.set_xlim(15, 23)\n",
    "    \n",
    "ax = fig.add_subplot(gs[0, 0], frameon=False)\n",
    "plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)\n",
    "ax.set_xlabel('Number of positive patients $n_{(+)}^\\mathrm{rep}$')\n",
    "ax.set_ylabel('Number of positive samples $m_{(+)}^\\mathrm{rep}$')\n",
    "\n",
    "labels_ = [\n",
    "    '(a) constant\\nstandard',\n",
    "    '(b) temporal\\nstandard',\n",
    "    '(c) constant sub-\\npopulation',\n",
    "    '(d) temporal sub-\\npopulation',\n",
    "]\n",
    "for ax, key in zip(np.ravel(axes), labels_):\n",
    "    ax.text(0.05, 0.95, key, ha='left', va='top', transform=ax.transAxes, size='small')\n",
    "        \n",
    "[plt.setp(ax.yaxis.get_ticklabels(), visible=False) for ax in axes[:, 1]]\n",
    "[plt.setp(ax.xaxis.get_ticklabels(), visible=False) for ax in axes[0]]\n",
    "\n",
    "ax = fig.add_subplot(gs[0, 1])\n",
    "ax.text(0.04, 0.975, '(e)', ha='left', va='top', transform=ax.transAxes, size='small')\n",
    "ax.set_ylabel(r'Sample mean $\\bar{x}^\\mathrm{rep}$ (gene copies per mL)')\n",
    "\n",
    "for i, inflated in enumerate(['standard', 'inflated']):\n",
    "    # Plot constant and temporal means\n",
    "    j = 0\n",
    "    for temporal in ['constant', 'temporal']:\n",
    "        target = lambda data: np.log10(data['load'][data['positive']].mean())\n",
    "        \n",
    "        key = ('general', inflated, temporal)\n",
    "        color = colors_by_key[key]\n",
    "        result = results[key]\n",
    "        reference = target(result['data'])\n",
    "        if i == 0:\n",
    "            x = (\n",
    "                2 * j - violin_widths * pad, \n",
    "                2 * j + 1 + violin_widths * pad\n",
    "            )\n",
    "            ax.plot(x, (reference, reference), color='k', ls='--')\n",
    "        \n",
    "        plot_violin(replicates_by_key[key], target, i + 2 * j, color, ax, marker=markers_by_key[key],\n",
    "                    label=f'{temporal}\\n{\"subpopulation\" if inflated == \"inflated\" else inflated}')\n",
    "        j += 1\n",
    "        \n",
    "ax.yaxis.set_major_formatter(log10formatter)\n",
    "ax.set_ylim(5, 8.5)\n",
    "ax.set_xlim(-.75, 3.75)\n",
    "ax.set_xticks([])\n",
    "ax.yaxis.set_major_locator(mpl.ticker.MaxNLocator(integer=True))\n",
    "ax.set_xlabel('Models')\n",
    "        \n",
    "fig.tight_layout()\n",
    "fig.savefig('workspace/figures/replication.pdf')\n",
    "# gs.update(wspace=.35)\n",
    "#for ax in axes[1]:\n",
    "\n",
    "#for ax in axes[:, 0]:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Effect of shape and scale parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Precompute the density estimates to iterate faster on the plots\n",
    "zzs = {}\n",
    "for level in ['population', 'patient']:\n",
    "    xmin, xmax = 0, 3\n",
    "    if level == 'population':\n",
    "        ymin, ymax = 1.5, 6.5\n",
    "    else:\n",
    "        ymin, ymax = 1.5, 4.75\n",
    "    \n",
    "    for temporal in ['constant', 'temporal']:\n",
    "        key = ('general', 'standard', temporal)\n",
    "        result = results[key]\n",
    "        samples = result['samples']\n",
    "\n",
    "        x = samples[f'{level}_shape'][::step].ravel()\n",
    "        y = samples[f'{level}_scale'][::step].ravel()\n",
    "        kde = stats.gaussian_kde((x, y))\n",
    "        \n",
    "        nx = 200\n",
    "        linx = np.linspace(-xmax, xmax, nx)\n",
    "        liny = np.linspace(ymin, ymax, 101)\n",
    "        xx, yy = np.meshgrid(linx, liny)\n",
    "        zz = kde.evaluate((xx.ravel(), yy.ravel())).reshape(xx.shape)\n",
    "        # Flip and add\n",
    "        zz = zz[:, :nx // 2:][:, ::-1] + zz[:, nx // 2:]\n",
    "        linx = linx[nx // 2:]\n",
    "        zzs.setdefault(key, {})[level] = (linx, liny, zz)\n",
    "        \n",
    "        print(level)\n",
    "        print(key, 'Weibull', np.mean(samples[f'{level}_shape'] > 1))\n",
    "        print(key, 'Gamma', np.mean(samples[f'{level}_shape'] > samples[f'{level}_scale']))\n",
    "        \n",
    "        \n",
    "# Add the heatmap of the means\n",
    "key = ('general', 'standard', 'constant')\n",
    "result = results[key]\n",
    "samples = result['samples']\n",
    "q = samples['population_shape'].ravel()[::step]\n",
    "mu = samples['population_loc'].ravel()[::step]\n",
    "sigma = samples['population_scale'].ravel()[::step]\n",
    "mean = shedding.gengamma_mean(q, mu, sigma)\n",
    "\n",
    "ymin, ymax = 1.5e5, 1e12\n",
    "# sb.kdeplot(x=q, y=mean, log_scale=(False, True), cmap=alpha_cmap(colors_by_key[key]))\n",
    "kde = stats.gaussian_kde((q, np.log(mean)))\n",
    "nx = 200\n",
    "linx = np.linspace(-xmax, xmax, nx)\n",
    "liny = np.linspace(np.log(ymin), np.log(ymax), 101)\n",
    "xx, yy = np.meshgrid(linx, liny)\n",
    "zz = kde.evaluate((xx.ravel(), yy.ravel())).reshape(xx.shape)\n",
    "# Flip and add\n",
    "linx = linx[nx // 2:]\n",
    "liny = np.exp(liny)\n",
    "zz = zz[:, :nx // 2:][:, ::-1] + zz[:, nx // 2:]\n",
    "zzs['mean'] = (linx, liny, zz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_hpd_mass(pdf):\n",
    "    \"\"\"\n",
    "    Evaluate the highest posterior density mass excluded from isocontours.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    pdf : array_like\n",
    "    \"\"\"\n",
    "    shape = np.shape(pdf)\n",
    "    pdf = np.ravel(pdf)\n",
    "    idx = np.argsort(-pdf)\n",
    "    cum = np.cumsum(pdf[idx])\n",
    "    cum /= cum[-1]\n",
    "    return 1 - np.reshape(cum[np.argsort(idx)], shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "gs = mpl.gridspec.GridSpec(2, 2)\n",
    "step = 1\n",
    "\n",
    "ax1 = fig.add_subplot(gs[0, 1])\n",
    "ax2 = fig.add_subplot(gs[1, 1], sharex=ax1)\n",
    "axes = [ax1, ax2]\n",
    "\n",
    "contours95 = []\n",
    "\n",
    "for ax, level in zip(axes, ['population', 'patient']):\n",
    "    gamma = np.logspace(-2, 1)\n",
    "    ax.plot(gamma, gamma, color='k', ls=':', zorder=0)\n",
    "    \n",
    "    xmin, xmax = 0, 2.1\n",
    "    if level == 'population':\n",
    "        ymin, ymax = 1.5, 6.5\n",
    "    else:\n",
    "        ymin, ymax = 1.5, 4.75\n",
    "    \n",
    "    for temporal in ['constant', 'temporal']:\n",
    "        key = ('general', 'standard', temporal)\n",
    "        linx, liny, zz = zzs[key][level]\n",
    "\n",
    "        mass = evaluate_hpd_mass(zz)\n",
    "        # Flip and add\n",
    "        # levels = evaluate_hpd_levels(zz, 9)\n",
    "        nlevels = 9\n",
    "        levels = (1 + np.arange(nlevels)) / (nlevels + 1)\n",
    "        \n",
    "        ax.contour(linx, liny, mass, levels=levels, \n",
    "                   cmap=alpha_cmap(colors_by_key[key]))\n",
    "        contours95.append(\n",
    "            ax.contour(linx, liny, mass, levels=[.05], \n",
    "                       colors=colors_by_key[key], linestyles=':'))\n",
    "        \n",
    "        # ax.set_xscale('log')\n",
    "        \n",
    "    ax.set_xlim(xmin, xmax)\n",
    "    ax.set_ylim(ymin, ymax)\n",
    "    ax.yaxis.tick_right()\n",
    "    ax.yaxis.set_label_position('right')\n",
    "    ax.set_ylabel(f'{level.title()} scale ${scale_label[level]}$')\n",
    "    ax.set_xlabel(f'{level.title()} shape ${shape_label[level]}$')\n",
    "    \n",
    "    \n",
    "# Plot the mean\n",
    "ax = fig.add_subplot(gs[:, 0], sharex=ax1)\n",
    "# Add the heatmap of the means\n",
    "key = ('general', 'standard', 'constant')\n",
    "result = results[key]\n",
    "samples = result['samples']\n",
    "data = result['data']\n",
    "q = samples['population_shape'].ravel()[::step]\n",
    "mu = samples['population_loc'].ravel()[::step]\n",
    "sigma = samples['population_scale'].ravel()[::step]\n",
    "mean = shedding.gengamma_mean(q, mu, sigma)\n",
    "\n",
    "ymin, ymax = 1.5e5, 1e9\n",
    "\n",
    "linx, liny, zz = zzs['mean']\n",
    "\n",
    "if True:\n",
    "    # Evaluate the thresholds\n",
    "    cmap = alpha_cmap(colors_by_key[key])\n",
    "    pvals = (np.arange(9) + 1) / 10\n",
    "    for pval in pvals:\n",
    "        levels = np.squeeze([evaluate_hpd_levels(z, pval) for z in zz.T])\n",
    "        x = zz / levels\n",
    "        ax.contour(linx, liny, x, levels=[1], colors=[cmap(1 - pval)])\n",
    "\n",
    "    levels = np.squeeze([evaluate_hpd_levels(z, .95) for z in zz.T])\n",
    "    x = zz / levels\n",
    "    contours95.append(ax.contour(linx, liny, x, levels=[1], colors=colors_by_key[key], linestyles=':'))\n",
    "    \n",
    "else:\n",
    "    conditional = zz / np.sum(zz, axis=0, keepdims=True)\n",
    "    ax.pcolormesh(evaluate_pcolormesh_edges(linx), evaluate_pcolormesh_edges(liny, 'log'), conditional)\n",
    "\n",
    "    ax.contour(linx, liny, np.transpose([evaluate_hpd_mass(z) for z in zz.T]), levels=[.05, np.exp(-1)], \n",
    "               colors='w', linestyles=[':', '--'])\n",
    "    \n",
    "ax.set_ylim(ymin, ymax)\n",
    "ax.set_yscale('log')\n",
    "    \n",
    "ax.set_ylabel(r'Mean RNA copies per mL')\n",
    "ax.set_xlabel(f'Population shape ${shape_label[\"population\"]}$')\n",
    "ax.axhline(data['load'][data['positive']].mean(), color='k', ls='--')\n",
    "\n",
    "ax.legend([\n",
    "    mpl.lines.Line2D([], [], color=colors_by_key[('standard', 'constant')]),\n",
    "    mpl.lines.Line2D([], [], color=colors_by_key[('standard', 'temporal')]),\n",
    "    mpl.lines.Line2D([], [], color='k', ls='-.'),\n",
    "    mpl.lines.Line2D([], [], color='k', ls=':'),\n",
    "    mpl.lines.Line2D([], [], color='k', ls='--'),\n",
    "], [\n",
    "    'constant\\nstandard',\n",
    "    'temporal\\nstandard',\n",
    "    'Weibull',\n",
    "    'gamma',\n",
    "    'sample\\nmean',\n",
    "], ncol=1, loc='upper right')\n",
    "\n",
    "\n",
    "clabel_pos = [\n",
    "    (1.75, 4),\n",
    "    (0.75, 4.5),\n",
    "    (1.25, 3.5),\n",
    "    (0.75, 1.75),\n",
    "    (1.5, 2e7),\n",
    "]\n",
    "\n",
    "for cs, manual in zip(contours95, clabel_pos):\n",
    "    cs.axes.clabel(cs, fmt={cs.levels[0]: '95%'}, manual=[manual])\n",
    "\n",
    "y = 0.95\n",
    "for ax, label in zip([ax, ax1, ax2], ['(a)', '(b) population', '(c) patient']):\n",
    "    ax.text(0.05, y, label, transform=ax.transAxes, ha='left', va='top', size='small')\n",
    "    ax.axvline(1, color='k', ls='-.', zorder=0)\n",
    "    \n",
    "fig.tight_layout()\n",
    "fig.savefig('workspace/figures/shape-scale.pdf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}