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:
In
__init__
we added the extra pars and states needed for our modelWe 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 transmissionWe updated
update_pre
andupdate_death
to include changes to theexposed
stateWe rewrote
set_prognoses
to include the new exposed state.