---
title: T4 - Diseases
jupyter: python3
---


```{python}
#| tags: [hide-cell]
%pip install -q starsim # Install Starsim if required
```

Diseases are the cornerstone of almost any Starsim analysis. In this tutorial, you'll learn how to work with diseases in Starsim, starting with simple modifications and building up to creating your own custom disease models.

By the end of this tutorial, you'll understand how to:
- Modify parameters of existing diseases
- Run simulations with multiple diseases
- Create your own custom disease from scratch

## Step 1: Modifying disease parameters

The easiest way to customize a disease is by changing its parameters. Much like sims or networks, a `Disease` can be customized by passing in a `pars` dictionary containing parameters. Let's start with a simple SIR model and see how different parameters affect the simulation:

```{python}
import starsim as ss
ss.options(jupyter=True)
sir = ss.SIR(dur_inf=10, beta=0.2, init_prev=0.4, p_death=0.2)
sim = ss.Sim(n_agents=2_000, diseases=sir, networks='random')
sim.run().plot()
```

Notice how we modified several key parameters:
- `dur_inf=10`: How long people stay infectious (10 years)
- `beta=0.2`: Transmission probability per contact
- `init_prev=0.4`: Starting with 40% of the population infected
- `p_death=0.2`: 20% of infected people die

We already saw that this model creates results that are stored in `sim.results.sir`. The results can also be directly accessed via `sir.results`. 

For more detail on any of the diseases that are in the Starsim library of diseases, please refer to the docstrings and source code of the disease files. 

## Step 2: Simulating multiple diseases
You can add multiple diseases to the same simulation, like so. Here we are making use of a "connector". A connector is a module in Starsim that tells you how two things relate to one another - in this case, how HIV modifies a person's transmissibility and susceptibility to gonorrhea and vice versa. Unlike dieases, networks, interventions, etc., connectors don't have any pre-specified location in the sim. Instead, they can be placed wherever they make the most sense (for example, a connector that mediated how two networks behaved might be placed at the beginning or end of the list of networks; for diseases, it might be placed at the beginning or end of the list of diseases).

```{python}
#| jupyter: {outputs_hidden: false}
import starsim as ss
import starsim_examples as sse

class simple_hiv_ng(ss.Module):
    """ Simple connector whereby rel_sus to NG (Neisseria gonorrhoeae) doubles if CD4 count is <200"""
    def __init__(self, pars=None, label='HIV-Gonorrhea', **kwargs):
        super().__init__()
        self.define_pars(
            rel_trans_hiv  = 2,
            rel_trans_aids = 5,
            rel_sus_hiv    = 2,
            rel_sus_aids   = 5,
        )
        self.update_pars(pars, **kwargs)
        return

    def step(self):
        """ Specify how HIV increases NG rel_sus and rel_trans """
        ng = self.sim.people.gonorrhea
        hiv = self.sim.people.hiv
        p = self.pars
        ng.rel_sus[hiv.cd4 < 500] = p.rel_sus_hiv
        ng.rel_sus[hiv.cd4 < 200] = p.rel_sus_aids
        ng.rel_trans[hiv.cd4 < 500] = p.rel_trans_hiv
        ng.rel_trans[hiv.cd4 < 200] = p.rel_trans_aids
        return

# Make HIV
hiv = sse.HIV(
    beta = {'mf': [0.0008, 0.0004]},  # Specify transmissibility over the MF network
    init_prev = 0.05,
)

# Make gonorrhea
ng = sse.Gonorrhea(
    beta = {'mf': [0.05, 0.025]},  # Specify transmissibility over the MF network
    init_prev = 0.025,
)

# Make the sim, including a connector betweeh HIV and gonorrhea:
n_agents = 5_000
sim = ss.Sim(n_agents=n_agents, networks='mf', diseases=[simple_hiv_ng(), hiv, ng])
sim.run()
sim.plot('hiv')
sim.plot('gonorrhea')
```

You can see how the two diseases interact - HIV creates a vulnerable population that's more susceptible to gonorrhea infection.

## Step 3: Creating your own disease model

Now for the fun part - creating your own disease from scratch! Let's modify the SIR model to create an SEIR model, which adds an "Exposed" state where people are infected but not yet infectious.

This is a common pattern in epidemiology - think of it as the incubation period before someone becomes contagious.

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:

```{python}
import starsim as ss
import matplotlib.pyplot as plt

class SEIR(ss.SIR):
    def __init__(self, pars=None, *args, **kwargs):
        super().__init__()
        self.define_pars(
            dur_exp = ss.lognorm_ex(0.5),
        )
        self.update_pars(pars, **kwargs)

        # Additional states beyond the SIR ones 
        self.define_states(
            ss.BoolState('exposed', label='Exposed'),
            ss.FloatArr('ti_exposed', label='TIme of exposure'),
        )
        return

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

    def step_state(self):
        """ Make all the updates from the SIR model """
        # Perform SIR updates
        super().step_state()

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

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

    def set_prognoses(self, uids, sources=None):
        """ Carry out state changes associated with infection """
        super().set_prognoses(uids, sources)
        ti = self.ti
        self.susceptible[uids] = False
        self.exposed[uids] = True
        self.ti_exposed[uids] = ti

        # Calculate and schedule future outcomes
        p = self.pars # Shorten for convenience
        dur_exp = p.dur_exp.rvs(uids)
        self.ti_infected[uids] = ti + dur_exp
        dur_inf = p.dur_inf.rvs(uids)
        will_die = p.p_death.rvs(uids)        
        self.ti_recovered[uids[~will_die]] = ti + dur_inf[~will_die]
        self.ti_dead[uids[will_die]] = ti + dur_inf[will_die]
        return
    
    def plot(self):
        """ Update the plot with the exposed compartment """
        with ss.options.context(show=False): # Don't show yet since we're adding another line
            fig = super().plot()
            ax = plt.gca()
            res = self.results.n_exposed
            ax.plot(res.timevec, res, label=res.label)
            plt.legend()
        return ss.return_fig(fig)
```

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.

Here's how it would look in practice:

```{python}
seir = SEIR()
sim = ss.Sim(diseases=seir, networks='random')
sim.run()
sim.plot()
sim.diseases.seir.plot()
```

## Exercises
1. **Parameter exploration**: Try different values of `dur_exp` in the SEIR model - how does it affect the epidemic curve?
2. **SEIRS model**: Adapt the SEIR example above to be SEIRS (where recovered people can become susceptible again)
3. **Multi-strain model**: Can you create a model with two strains of the same disease that provide partial cross-immunity?

