Developer tutorial: Diseases#

Overview of Starsim’s disease structure#

The basic template for modeling a disease in Starsim is the Disease class. Much like sims or networks, a Disease can be customized by passing in a pars dictionary containing parameters. The Disease module does lots of different things, but three of the model fundamental are:

  • set_initial_states, which initializes people into states (e.g. susceptible, infected, recovered)

  • make_new_cases, which makes new cases (e.g., by finding the susceptible contacts of all infected people and calculating the probability of transmission for each)

  • set_prognoses, which sets the outcomes for people who get infected (e.g., by setting their date of recovery or death).

Making your own disease#

If you want to make your own disease, you could either inherit from one of the templates in diseases.py, or you could copy the examples and extend them to capture features of the disease that you want to model. For example, suppose you wanted to change the SIR model to an SEIR model (i.e., add an ‘exposed’ state where people were transmissible but did not yet have symptoms. You might hope that this would be a relatively simple change to make. Here’s how it would look:

[ ]:



[3]:

import numpy as np
import starsim as ss


class SEIR(ss.SIR):
    def __init__(self, pars=None, *args, **kwargs):

        # Additional pars beyond the SIR model ones
        pars = ss.omerge({
            'dur_exp': 0.5,
        }, pars)

        # Initialize the SIR model, which will add all the parameters and states of that model.
        super().__init__(pars=pars, *args, **kwargs)

        # Additional states beyond the SIR ones
        self.exposed = ss.State('exposed', bool, False)
        self.ti_exposed = ss.State('ti_exposed', float, np.nan)

        return

    @property
    def infectious(self):
        return self.infected | self.exposed

    def update_pre(self, sim):
        # Make all the updates from the SIR model
        n_deaths = super().update_pre(sim)

        # Additional updates: progress exposed -> infected
        infected = ss.true(self.exposed & (self.ti_infected <= sim.ti))
        self.exposed[infected] = False
        self.infected[infected] = True

        return n_deaths

    def update_death(self, sim, uids):
        super().update_death(sim, uids)
        self.exposed[uids] = False
        return

    def set_prognoses(self, sim, uids, from_uids):
        # Carry out state changes associated with infection
        self.susceptible[uids] = False
        self.exposed[uids] = True
        self.ti_exposed[uids] = sim.ti

        # Calculate and schedule future outcomes
        dur_exp = self.pars['dur_exp'].rvs(uids)
        self.ti_infected[uids] = sim.ti + dur_exp / sim.dt
        dur_inf = self.pars['dur_inf'].rvs(uids)
        will_die = self.pars['p_death'].rvs(uids)
        self.ti_recovered[uids[~will_die]] = sim.ti + dur_inf[~will_die] / sim.dt
        self.ti_dead[uids[will_die]] = sim.ti + dur_inf[will_die] / sim.dt

        # Update result count of new infections
        self.results['new_infections'][sim.ti] += len(uids)
        return

The new class includes the following main changes:

  1. In __init__ we added the extra pars and states needed for our model

  2. We defined infectious to include both infected and exposed people - this means that we can just reuse the existing logic for how the SIR model handles transmission

  3. We updated update_pre and update_death to include changes to the exposed state

  4. We rewrote set_prognoses to include the new exposed state.