{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# T6 - Using analyzers\n",
    "\n",
    "Analyzers are objects that do not change the behavior of a simulation, but just report on its internal state, almost always something to do with `sim.people`. This tutorial takes you through some of the built-in analyzers and gives a brief example of how to build your own."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "Click [here](https://mybinder.org/v2/gh/institutefordiseasemodeling/covasim/HEAD?urlpath=lab%2Ftree%2Fdocs%2Ftutorials%2Ftut_analyzers.ipynb) to open an interactive version of this notebook.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Age histograms and snapshots\n",
    "\n",
    "Age histograms and snapshots both take \"pictures\" of the `sim.people` object at specified points in time. This is because while most of the information from `sim.people` is retrievable at the end of the sim from the stored events, it's much easier to see what's going on at the time. While the snapshot literally just makes a copy of the people object, the age histogram calculates results for different age bins:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import covasim as cv\n",
    "cv.options(jupyter=True, verbose=0)\n",
    "\n",
    "sim = cv.Sim(interventions=cv.test_prob(0.5), analyzers=cv.age_histogram())\n",
    "sim.run()\n",
    "agehist = sim.get_analyzer() # Only one analyzer so we can retrieve it like this\n",
    "agehist.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transmission trees\n",
    "\n",
    "Another useful analysis to perform on the simulation is to calculate the *transmission tree* – i.e., who infected whom. Since all this information is stored in the sim, it doesn't have to be inserted at run time; it can be added to an already finished simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt = sim.make_transtree()\n",
    "fig1 = tt.plot()\n",
    "fig2 = tt.plot_histograms()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom analyzers\n",
    "\n",
    "Analyzers don't have to be complicated. While analyzers can derive from the `cv.Analyzer` class, they can also just be simple functions, unless they need to keep track of or update internal state. Here's an example of an analyzer that does keep track of internal state, namely by storing and plotting the S, E, I, and R compartments over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pylab as pl\n",
    "import sciris as sc\n",
    "import covasim as cv\n",
    "\n",
    "class store_seir(cv.Analyzer):\n",
    "\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs) # This is necessary to initialize the class properly\n",
    "        self.t = []\n",
    "        self.S = []\n",
    "        self.E = []\n",
    "        self.I = []\n",
    "        self.R = []\n",
    "        return\n",
    "\n",
    "    def apply(self, sim):\n",
    "        ppl = sim.people # Shorthand\n",
    "        self.t.append(sim.t)\n",
    "        self.S.append(ppl.susceptible.sum())\n",
    "        self.E.append(ppl.exposed.sum() - ppl.infectious.sum())\n",
    "        self.I.append(ppl.infectious.sum())\n",
    "        self.R.append(ppl.recovered.sum() + ppl.dead.sum())\n",
    "        return\n",
    "\n",
    "    def plot(self):\n",
    "        pl.figure()\n",
    "        pl.plot(self.t, self.S, label='S')\n",
    "        pl.plot(self.t, self.E, label='E')\n",
    "        pl.plot(self.t, self.I, label='I')\n",
    "        pl.plot(self.t, self.R, label='R')\n",
    "        pl.legend()\n",
    "        pl.xlabel('Day')\n",
    "        pl.ylabel('People')\n",
    "        sc.setylim() # Reset y-axis to start at 0\n",
    "        sc.commaticks() # Use commas in the y-axis labels\n",
    "        return\n",
    "\n",
    "sim = cv.Sim(n_days=180, analyzers=store_seir(label='seir'))\n",
    "sim.run()\n",
    "seir = sim.get_analyzer('seir') # Retrieve by label\n",
    "seir.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's an example of an analyzer that doesn't need any internal state, so it can just be a function: it simply reports if anyone aged 88 is currently infected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_88(sim):\n",
    "    people_who_are_88 = sim.people.age.round() == 88 # Find everyone who's aged 88 (to the nearest year)\n",
    "    people_exposed = sim.people.exposed # Find everyone who's infected with COVID\n",
    "    people_who_are_88_with_covid = cv.true(people_who_are_88 * people_exposed) # Multiplication is the same as logical \"and\"\n",
    "    n = len(people_who_are_88_with_covid) # Count how many people there are\n",
    "    if n:\n",
    "        print(f'Oh no! {n} people aged 88 have covid on timestep {sim.t} {\"🤯\"*n}')\n",
    "    return\n",
    "\n",
    "sim = cv.Sim(n_days=120, analyzers=check_88, verbose=0)\n",
    "sim.run()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}