polychord-sampling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import functools as ft\n",
    "import itertools as it\n",
    "from localscope import localscope\n",
    "import logging\n",
    "logging.basicConfig(level='WARNING')\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "import os\n",
    "import pypolychord as pc\n",
    "import pypolychord.settings\n",
    "import shedding\n",
    "from scipy import stats, special\n",
    "from tqdm.notebook import tqdm\n",
    "import sys\n",
    "import pickle\n",
    "\n",
    "mpl.rcParams['figure.dpi'] = 144\n",
    "\n",
    "args = os.environ.get('ARGS', \"-i -s 0 -l 1 -r 1 -f --day-noise 1 general workspace/test\").split()\n",
    "print(args)\n",
    "parser = argparse.ArgumentParser()\n",
    "# Positional\n",
    "parser.add_argument('parametrisation', type=shedding.Parametrisation, \n",
    "                    help='parametrisation to use')\n",
    "parser.add_argument('basedir')\n",
    "\n",
    "# Keyword\n",
    "parser.add_argument('--inflated', '-i', action='store_true', \n",
    "                    help='whether to use a zero-inflated model')\n",
    "parser.add_argument('--temporal', '-t', choices=['gamma', 'exponential', 'teunis'],\n",
    "                    help='whether to use a time-dependent model', default=False)\n",
    "parser.add_argument('--force', '-f', action='store_true', \n",
    "                    help='force regeneration of samples')\n",
    "parser.add_argument('--seed', '-s', type=int, help='random number generator seed')\n",
    "parser.add_argument('--nlive-factor', '-l', type=float, default=25, \n",
    "                    help='multiplicative factor for number of live points')\n",
    "parser.add_argument('--nrepeat-factor', '-r', type=float, default=5, \n",
    "                    help='multiplicative factor for number of monte carl steps')\n",
    "parser.add_argument('--evidence', '-e', action='store_true', help='focus on evaluating the evidence, '\n",
    "                    'make sure we use the same data for temporal and constant models')\n",
    "parser.add_argument('--day-noise', help='amount of noise to add to days past symptom onset',\n",
    "                    type=int, default=0)\n",
    "args = parser.parse_args(args)\n",
    "os.makedirs(args.basedir, exist_ok=True)\n",
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(args.seed)\n",
    "\n",
    "# Load and flatten some of the datasets\n",
    "datasets = [\n",
    "    'Woelfel2020', \n",
    "    'Lui2020', \n",
    "    'Han2020',\n",
    "]\n",
    "# Include Wang for non-temporal models if we just want samples, not evidence\n",
    "if not args.evidence and not args.temporal:\n",
    "    datasets.append('Wang2020')\n",
    "    \n",
    "datasets = shedding.load_datasets(datasets, 'publications/')\n",
    "\n",
    "# Add random bits of noise to the days to test whether there's sensitivity to misreporting\n",
    "# the number of days since symptom onset\n",
    "for dataset in datasets.values():\n",
    "    noise_by_patient = {}\n",
    "    for load in dataset['loads']:\n",
    "        noise = noise_by_patient.setdefault(\n",
    "            load['patient'], \n",
    "            np.random.randint(-args.day_noise, args.day_noise + 1)\n",
    "        )\n",
    "        if 'day' in load:\n",
    "            load['day'] += noise\n",
    "\n",
    "\n",
    "data = shedding.flatten_datasets(datasets, loq_fill_value=-99)\n",
    "print(f'Number of patients: {data[\"num_patients\"]}')\n",
    "print(f'Number of patients with one or more positive samples: '\n",
    "      f'{(data[\"num_positives_by_patient\"] > 0).sum()}')\n",
    "print(f'Number of samples: {data[\"num_samples\"]}')\n",
    "print(f'Number of positive samples: {data[\"positive\"].sum()}')\n",
    "\n",
    "model = shedding.Model(data['num_patients'], parametrisation=args.parametrisation, inflated=args.inflated,\n",
    "                       temporal=args.temporal)\n",
    "print(f'Number of parameters: {model.size}')\n",
    "\n",
    "with open(os.path.join(args.basedir, 'model.pkl'), 'wb') as fp:\n",
    "    pickle.dump(model, fp)\n",
    "    \n",
    "with open(os.path.join(args.basedir, 'data.pkl'), 'wb') as fp:\n",
    "    pickle.dump(data, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings = pc.settings.PolyChordSettings(model.size, 0)\n",
    "settings.base_dir = args.basedir\n",
    "settings.file_root = 'chain'\n",
    "settings.read_resume = False\n",
    "settings.feedback = 3\n",
    "settings.num_repeats = int(args.nrepeat_factor * model.size)\n",
    "settings.boost_posterior = min(settings.num_repeats, 10)\n",
    "settings.nlive = int(args.nlive_factor * model.size)\n",
    "settings.seed = -1 if args.seed is None else args.seed\n",
    "settings.write_resume = False\n",
    "\n",
    "filename = os.path.join(settings.base_dir, settings.file_root + '.paramnames')\n",
    "shedding.write_paramnames_file(model.parameters, filename)\n",
    "\n",
    "vars(settings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate that we can sample and evaluate\n",
    "log_likelihoods = []\n",
    "for _ in range(10_000):\n",
    "    x = np.random.uniform(size=model.size)\n",
    "    y = model.sample_params_from_vector(x)\n",
    "    log_likelihood, _ = model.evaluate_log_likelihood_from_vector(y, data)\n",
    "    assert np.isfinite(log_likelihood), log_likelihood\n",
    "    log_likelihoods.append(log_likelihood)\n",
    "log_likelihoods = np.asarray(log_likelihoods)\n",
    "\n",
    "# Verify that at least 10% of prior samples aren't unphysical.\n",
    "assert np.mean(log_likelihoods > settings.logzero) > .1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shedding.vector_to_values(model.parameters, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "filename = os.path.join(settings.base_dir, settings.file_root + '.txt')\n",
    "if os.path.isfile(filename) and not args.force:\n",
    "    print(f'{filename} already exists; remove it to regenerate samples')\n",
    "else:\n",
    "    log_likelihood = ft.partial(model.evaluate_log_likelihood_from_vector, data=data)\n",
    "    with tqdm() as progress: \n",
    "        output = pc.run_polychord(log_likelihood, model.size, 0, settings, model.sample_params_from_vector, \n",
    "                                  lambda *args: progress.update())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = np.loadtxt(os.path.join(settings.base_dir, settings.file_root + '_equal_weights.txt'))[:, 2:]\n",
    "num_samples = len(samples)\n",
    "print(f'Obtained {num_samples} samples.')\n",
    "samples = shedding.transpose_samples(samples, model.parameters)\n",
    "\n",
    "\n",
    "# Save the results in a format that makes our life easier\n",
    "with open(os.path.join(args.basedir, 'result.pkl'), 'wb') as fp:\n",
    "    pickle.dump({\n",
    "        'samples': samples,\n",
    "        'model': model,\n",
    "        'args': args,\n",
    "        'evidence': (output.logZ, output.logZerr),\n",
    "        'local_evidences': (output.logZs, output.logZerrs),\n",
    "        'data': data,\n",
    "    }, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a figure with the right number of columns\n",
    "model_pars = [par for par, value in samples.items() if not par.endswith('_') and\n",
    "              np.ndim(value) < 2]\n",
    "nrows = ncols = len(model_pars) - 1\n",
    "fig, axes = plt.subplots(nrows, ncols, sharex='col', sharey='row',\n",
    "                         figsize=(6, 6))\n",
    "\n",
    "# Iterate over the pairs\n",
    "step = max(1, num_samples // 500)\n",
    "for i in range(len(model_pars) - 1):\n",
    "    a = model_pars[i]\n",
    "    for j in range(1, len(model_pars)):\n",
    "        ax = axes[i, j - 1]\n",
    "        if j <= i:\n",
    "            ax.set_axis_off()\n",
    "            continue\n",
    "            \n",
    "        b = model_pars[j]\n",
    "        x = samples[a][::step]\n",
    "        y = samples[b][::step]\n",
    "\n",
    "        ax.scatter(y, x, marker='.', alpha=.1)\n",
    "        if i == 0:\n",
    "            ax.set_xlabel(b, size='small')\n",
    "            ax.xaxis.set_label_position('top')\n",
    "            ax.xaxis.tick_top()\n",
    "            ax.xaxis.set_tick_params(which='both', labeltop=True)\n",
    "            \n",
    "        if j == ncols:\n",
    "            ax.set_ylabel(a, size='small')\n",
    "            ax.yaxis.set_label_position('right')\n",
    "            ax.yaxis.tick_right()\n",
    "            ax.yaxis.set_tick_params(which='both', labelright=True)\n",
    "            \n",
    "        # Draw prior bounds if they exist\n",
    "        for k, line in zip([b, a], [ax.axvline, ax.axhline]):\n",
    "            prior = model.priors.get(k)\n",
    "            if not prior:\n",
    "                continue\n",
    "            bounds = prior.bounds\n",
    "            for bound in bounds:\n",
    "                if bound is not None:\n",
    "                    line(bound, color='k', ls=':')\n",
    "            \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "for ax, key in zip((ax1, ax2), ('population', 'patient')):\n",
    "    a = f'{key}_shape'\n",
    "    b = f'{key}_scale'\n",
    "    x = samples[a]\n",
    "    y = samples[b]\n",
    "    ax.scatter(x, y, marker='.', alpha=.5)\n",
    "    ax.set_xlabel(r'Shape $q$')\n",
    "    \n",
    "    xmin = max(x.min(), y.min())\n",
    "    xmax = min(x.max(), y.max())\n",
    "    lin = np.linspace(xmin, xmax)\n",
    "    ax.plot(lin, lin, color='k', ls='--')\n",
    "    \n",
    "    ax.axvline(1, color='k', ls=':')\n",
    "    # ax.set_aspect('equal')\n",
    "    if False:\n",
    "        ax.set_xscale('log')\n",
    "        ax.set_yscale('log')\n",
    "        \n",
    "ax1.set_ylabel(r'Scale $\\sigma$')\n",
    "ax1.set_title('Population')\n",
    "ax2.set_title('Patient')\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zipped = zip(samples['population_shape'], samples['population_loc'], samples['population_scale'])\n",
    "x = np.log10([shedding.gengamma_mean(q, mu, sigma) for q, mu, sigma in zipped])\n",
    "if model.inflated:\n",
    "    x += np.log10(samples['rho'])\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.hist(x, density=True, alpha=.5)\n",
    "shedding.plot_kde(x, color='C0', ax=ax, label=r'population mean $\\langle y\\rangle$')\n",
    "ax.axvline(np.log10(data['load'].max()), color='k', ls=':', label=r'$\\max x$')\n",
    "ax.axvline(np.log10(data['load'][data['positive']].mean()), color='k', ls='--', \n",
    "           label=r'sample mean $\\bar x$')\n",
    "ax.set_xlabel(r'$\\log_{10}$ gene copies per mL')\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model.inflated:\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.hist(samples['rho'], density=True, range=(samples['rho'].min(), 1), alpha=.5)\n",
    "    ax.axvline(np.mean(data['num_positives_by_patient'] > 0), color='k', ls='--')\n",
    "    ax.set_xlabel(r'$\\rho$')\n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model.temporal != shedding.Profile.CONSTANT:\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    keys = set(data['dataset'])\n",
    "    for key in keys:\n",
    "        fltr = data['positive'] & (data['dataset'] == key)\n",
    "        pts = ax.scatter(data['day'][fltr], data['load'][fltr], alpha=.5, \n",
    "                         label=key)\n",
    "        c = np.squeeze(pts.get_facecolor())\n",
    "        ax.axhline(data['loq'][fltr][0], color=c, ls=':')\n",
    "\n",
    "    lin = np.linspace(-21, data['day'][data['positive']].max(), 200)\n",
    "    profile = np.exp(samples['population_loc'] + \n",
    "                     model.temporal.evaluate_offset(lin[:, None], samples)).T\n",
    "    line, = ax.plot(lin, np.median(profile, axis=0), color='k')\n",
    "    ax.fill_between(lin, *np.percentile(profile, [2.5, 97.5], axis=0), alpha=.2, color=line.get_color())\n",
    "    ax.set_yscale('log')\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model.temporal != shedding.Profile.CONSTANT:\n",
    "    fig, ax = plt.subplots()\n",
    "    scale = None\n",
    "    for key in ['profile_decay', 'profile_scale']:\n",
    "        if key in samples:\n",
    "            scale = samples[key]\n",
    "            break\n",
    "    if scale is None:\n",
    "        scale = - samples['slope']\n",
    "    halflife = 24 * np.log(2) / scale\n",
    "    plt.hist(halflife, density=True, bins=20)"
   ]
  }
 ],
 "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
}