covid_19_modelling2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modelling and simulating COVID-19 using agent based models\n",
    "\n",
    "<p class=\"pm-node nj-authors\">Marc Sturrock, RCSI [marcsturrock@rcsi.com]</p>\n",
    "\n",
    "Here we expand upon the simple ODE based SIR model in the previous notebook to consider agent based models. We make use of the julia package \"Agents.jl\" for this. Note that these models are also stochastic, so repeated runs will yield different results. In general with stochastic models, we rely on summary statistics like mean and variance to judge the model behaviour. \n",
    "\n",
    "Here is what will be covered in this session:\n",
    "\n",
    "- [x] Agent based models description \n",
    "- [x] SIR model\n",
    "- [x] Model parameters\n",
    "- [x] Making the model in Agents.jl\n",
    "- [x] Example animation\n",
    "- [x] Simulating Delta vs Omicron"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Description: Agent Based Model (ABM)\n",
    "\n",
    "An agent-based (or individual-based) model is a computational simulation of autonomous agents that react to their environment (including other agents) given a predefined set of rules. ABMs have been adopted and studied in a variety of research disciplines. One reason for their popularity is that they enable a relaxation of many simplifying assumptions usually made by mathematical models and thus have greater flexibility.\n",
    "\n",
    "Agent-based models are increasingly recognised as a useful approach for studying complex systems. Complex systems cannot be fully understood using traditional mathematical tools which aggregate the behaviour of elements in a system. The behavior of a complex system depends on both the behavior of and interactions between its elements (agents). Small changes in the input to complex systems or the behaviour of its agents can lead to large changes in outcome. That is to say, a complex system's behaviour is nonlinear, and that it is not only the sum of the behaviour of its elements. Use of ABMs have become feasible after the availability of computers and has been growing ever since, especially in modelling biological and economical systems, and has extended to social studies and archaeology.\n",
    "\n",
    "An ABM consists of autonomous agents that behave given a set of rules. Relatively simple rules can give rise to complex emergent behaviour. ABMs of COVID-19 that have informed governments decisions for imposing restrictions - their flexibility means that they can account for census data - see e.g.  https://www.mdpi.com/2079-8954/9/2/41/htm\n",
    "\n",
    "# SIR model\n",
    "\n",
    "A SIR model tracks the ratio of Susceptible, Infected, and Recovered individuals within a population. Here we add one more category of individuals: those who are infected, but do not know it. Transmission rate for infected and diagnosed individuals is lower than infected and undetected. We also allow a fraction of recovered individuals to catch the disease again, meaning that recovering the disease does not bring full immunity.\n",
    "\n",
    "# Model parameters\n",
    "\n",
    "Here are the model parameters, some of which have default values.\n",
    "\n",
    "- Ns: a vector of population sizes per household. The amount of cities is just C=length(Ns).\n",
    "- β_und: a vector for transmission probabilities β of the infected but undetected per household. Transmission probability is how many susceptible are infected per day by an infected individual. If social - distancing is practiced, this number decreases.\n",
    "- β_det: an array for transmission probabilities β of the infected and detected per household. If hospitals are full, this number increases.\n",
    "- infection_period = 30: how many days before a person dies or recovers.\n",
    "- detection_time = 14: how many days before an infected person is detected.\n",
    "- death_rate = 0.02: the probability that the individual will die after the infection_period.\n",
    "- reinfection_probability = 0.05: The probability that a recovered person can get infected again.\n",
    "- migration_rates: A matrix of migration probability per individual per day from one household to another.\n",
    "- Is = [zeros(C-1)..., 1]: An array for initial number of infected but undetected people per household. This starts as only one infected individual in the last household.\n",
    "\n",
    "Notice that Ns, β, Is all need to have the same length, as they are numbers for each household. We've tried to add values to the infection parameters similar to the ones you would hear on the news about COVID-19.\n",
    "\n",
    "The good thing with Agent based models is that you could easily extend the model we implement here to also include age as an additional property of each agent. This makes ABMs flexible and suitable for research of virus spreading.\n",
    "\n",
    "# Making the model in Agents.jl\n",
    "\n",
    "We start by defining people in the simulation and the Agent based model.. (see comments following hash symbols for explanations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Agents, Random, DataFrames, Graphs\n",
    "using Distributions: Poisson, DiscreteNonParametric\n",
    "using DrWatson: @dict\n",
    "using CairoMakie\n",
    "\n",
    "mutable struct Person <: AbstractAgent\n",
    "    id::Int #A unique number associated with each individual\n",
    "    pos::Int #Current position of individual\n",
    "    location::Int #original location or home location\n",
    "    days_infected::Int  # number of days since is infected\n",
    "    status::Symbol  # 1: S, 2: I, 3:R\n",
    "end\n",
    "\n",
    "function model_initiation(;\n",
    "    Ns,\n",
    "    migration_rates,\n",
    "    β_und,\n",
    "    β_det,\n",
    "    infection_period = 30,\n",
    "    reinfection_probability = 0.05,\n",
    "    detection_time = 14,\n",
    "    death_rate = 0.02,\n",
    "    Is = [zeros(Int, length(Ns) - 1)..., 1],\n",
    "    seed = 0,\n",
    ")\n",
    "\n",
    "    rng = MersenneTwister(seed)\n",
    "    @assert length(Ns) ==\n",
    "    length(Is) ==\n",
    "    length(β_und) ==\n",
    "    length(β_det) ==\n",
    "    size(migration_rates, 1) \"length of Ns, Is, and B, and number of rows/columns in migration_rates should be the same \"\n",
    "    @assert size(migration_rates, 1) == size(migration_rates, 2) \"migration_rates rates should be a square matrix\"\n",
    "\n",
    "    C = length(Ns)\n",
    "    # normalize migration_rates\n",
    "    migration_rates_sum = sum(migration_rates, dims = 2)\n",
    "    for c in 1:C\n",
    "        migration_rates[c, :] ./= migration_rates_sum[c]\n",
    "    end\n",
    "\n",
    "    properties = @dict(\n",
    "        Ns,\n",
    "        Is,\n",
    "        β_und,\n",
    "        β_det,\n",
    "        β_det,\n",
    "        migration_rates,\n",
    "        infection_period,\n",
    "        infection_period,\n",
    "        reinfection_probability,\n",
    "        detection_time,\n",
    "        C,\n",
    "        death_rate\n",
    "    )\n",
    "    space = Agents.GraphSpace(Graphs.SimpleGraphs.barabasi_albert(C,1,1))\n",
    "    model = ABM(Person, space; properties, rng)\n",
    "\n",
    "    # Add initial individuals\n",
    "    for household in 1:C, n in 1:Ns[household]\n",
    "        ind = add_agent!(household, model,household, 0, :S) # Susceptible\n",
    "    end\n",
    "    # add infected individuals\n",
    "    for household in 1:C\n",
    "        inds = ids_in_position(household, model)\n",
    "        for n in 1:Is[household]\n",
    "            agent = model[inds[n]]\n",
    "            agent.status = :I # Infected\n",
    "            agent.days_infected = 1\n",
    "        end\n",
    "    end\n",
    "    return model\n",
    "end"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will make a function that starts a model with C number of households, and creates the other parameters automatically by attributing some random values to them. You could directly use the above constructor and specify all Ns, β, etc. for a given set of households.\n",
    "\n",
    "Households are connected to each other using a Barabasi-Albert graph which is a good model of social interactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra: diagind\n",
    "\n",
    "function create_params(;\n",
    "    C,\n",
    "    max_travel_rate,\n",
    "    infection_period = 30,\n",
    "    reinfection_probability = 0.1,\n",
    "    detection_time = 14,\n",
    "    death_rate = 0.01,\n",
    "    Is = [zeros(Int, C - 1)..., 1],\n",
    "    seed = 19,\n",
    ")\n",
    "\n",
    "    Random.seed!(seed)\n",
    "    Ns = rand(1:10, C)\n",
    "    β_und = 1 .* rand(0.3:0.02:0.6, C)\n",
    "    β_det = 1 .* β_und ./ 10\n",
    "\n",
    "    Random.seed!(seed)\n",
    "    migration_rates = zeros(C, C)\n",
    "    for c in 1:C\n",
    "        for c2 in 1:C\n",
    "            migration_rates[c, c2] = (Ns[c] + Ns[c2]) / Ns[c]\n",
    "        end\n",
    "    end\n",
    "    maxM = maximum(migration_rates)\n",
    "    migration_rates = (migration_rates .* max_travel_rate) ./ maxM\n",
    "    migration_rates[diagind(migration_rates)] .= 1.0\n",
    "\n",
    "    params = @dict(\n",
    "        Ns,\n",
    "        β_und,\n",
    "        β_det,\n",
    "        migration_rates,\n",
    "        infection_period,\n",
    "        reinfection_probability,\n",
    "        detection_time,\n",
    "        death_rate,\n",
    "        Is\n",
    "    )\n",
    "\n",
    "    return params\n",
    "end\n",
    "\n",
    "params = create_params(C = 100, max_travel_rate = 0.1,reinfection_probability = 0.001,death_rate = 0.02)\n",
    "model = model_initiation(; params...)\n",
    "\n",
    "function agent_step!(agent, model)\n",
    "    migrate!(agent, model)\n",
    "    transmit!(agent, model)\n",
    "    update!(agent, model)\n",
    "    recover_or_die!(agent, model)\n",
    "end\n",
    "\n",
    "function migrate!(agent, model)\n",
    "    current_location = agent.pos\n",
    "    home = agent.location\n",
    "    d = DiscreteNonParametric(1:(model.C), model.migration_rates[current_location, :])\n",
    "    m = rand(model.rng, d)\n",
    "    if current_location ≠ home\n",
    "        move_agent!(agent, home, model)\n",
    "    elseif m ≠ current_location && m ∈ nearby_positions(current_location,model)\n",
    "        move_agent!(agent, m, model)\n",
    "    end\n",
    "end\n",
    "\n",
    "function transmit!(agent, model)\n",
    "    agent.status == :S && return\n",
    "    rate = if agent.days_infected < model.detection_time\n",
    "        model.β_und[agent.pos]\n",
    "    else\n",
    "        model.β_det[agent.pos]\n",
    "    end\n",
    "\n",
    "    d = Poisson(rate)\n",
    "    n = rand(model.rng, d)\n",
    "    n == 0 && return\n",
    "\n",
    "    for contactID in ids_in_position(agent, model)\n",
    "        contact = model[contactID]\n",
    "        if contact.status == :S || (contact.status == :R && rand(model.rng) ≤ model.reinfection_probability)\n",
    "            contact.status = :I\n",
    "            n -= 1\n",
    "            n == 0 && return\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "update!(agent, model) = agent.status == :I && (agent.days_infected += 1)\n",
    "\n",
    "function recover_or_die!(agent, model)\n",
    "    if agent.days_infected ≥ model.infection_period\n",
    "        if rand(model.rng) ≤ model.death_rate\n",
    "            kill_agent!(agent, model)\n",
    "        else\n",
    "            agent.status = :R\n",
    "            agent.days_infected = 0\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using InteractiveDynamics\n",
    "abmobs = ABMObservable(model; agent_step!)\n",
    "\n",
    "alive_fractions(m) = [3 .* length(ids_in_position(p, m)) for p in positions(m)]\n",
    "infected_fraction(m, x) = count(m[id].status == :I for id in x) / length(x)\n",
    "infected_fractions(m) = [infected_fraction(m, ids_in_position(p, m)) for p in positions(m)]\n",
    "alives = lift(alive_fractions,abmobs.model)\n",
    "fracs = lift(infected_fractions, abmobs.model)\n",
    "color = lift(fs -> [cgrad(:inferno)[f] for f in fs], fracs)\n",
    "title = lift(\n",
    "    (s, m) -> \"step = $(s), infected = $(round(Int, 100infected_fraction(m, allids(m))))%\",\n",
    "    abmobs.s, abmobs.model\n",
    ")\n",
    "using Base64\n",
    "using GraphMakie\n",
    "fig = Figure(resolution = (1200, 800))\n",
    "ax2 = Axis(fig[1, 1]; title, xlabel = \"x\", ylabel = \"y\")\n",
    "graphplot!(ax2,model.space.graph,node_color=color,node_size=alives)\n",
    "Colorbar(fig[1, 2],label = \"percentage infected\",limits = (0, 0.35))\n",
    "fig\n",
    "record(fig, \"covid_evolution.mp4\"; framerate = 10) do io\n",
    "    for j in 1:100\n",
    "        recordframe!(io)\n",
    "        Agents.step!(abmobs, 1)\n",
    "    end\n",
    "    recordframe!(io)\n",
    "end\n",
    "function showanim(filename)\n",
    "    base64_video = base64encode(open(filename))\n",
    "    display(\"text/html\", \"\"\"<video controls src=\"data:video/x-m4v;base64,$base64_video\">\"\"\")\n",
    "end\n",
    "showanim(\"covid_evolution.mp4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infected(x) = count(i == :I for i in x)\n",
    "susceptible(x) = count(i == :S for i in x)\n",
    "recovered(x) = count(i == :R for i in x)\n",
    "\n",
    "model = model_initiation(; params...)\n",
    "\n",
    "to_collect = [(:status, f) for f in (susceptible, infected, recovered, length)]\n",
    "data, _ = run!(model, agent_step!, 1000; adata=to_collect)\n",
    "\n",
    "N = sum(model.Ns) # Total initial population\n",
    "x = data.step\n",
    "fig = Figure(resolution = (1200, 800),fontsize=24,grid=false)\n",
    "ax = fig[1, 1] = Axis(fig; xlabel = \"Time (days)\", ylabel = \"Number of Agents\")\n",
    "ls = lines!(ax, x, data[:, dataname((:status, susceptible))], color = :black,linewidth=5)\n",
    "li = lines!(ax, x, data[:, dataname((:status, infected))], color = :blue,linewidth=5)\n",
    "lr = lines!(ax, x, data[:, dataname((:status, recovered))], color = :red,linewidth=5)\n",
    "dead = N .- data[:, dataname((:status, length))]\n",
    "ld = lines!(ax, x, dead, color = :green,linewidth=5,grid=false)\n",
    "fig[1, 2][1,1] = Legend(fig, [ls,li, lr, ld], [\"susceptible\",\"infected\", \"recovered\", \"dead\"])\n",
    "fig"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating Delta vs Omicron variants\n",
    "\n",
    "Omicron is known to have a higher chance of reinfection than delta as well as a smaller death rate. Therefore to simulate a more omicron like variant try executing the above cells with the below lines changed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = create_params(C = 100, max_travel_rate = 0.01,reinfection_probability = 0.05,death_rate = 0.002)\n",
    "model = model_initiation(; params...)"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.9.0",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}