Running a model in stochastic mode

CompartmentalModel runs are deterministic, because the ODE solvers we use are deterministic. For a given set of compartments and (deterministic) flows, the outputs will always be the same.

These deterministic models work better in scenarios where there is a large population and the law of large numbers has a strong effect.

Sometimes, however, we want to simulate scenarios where some compartments have very low numbers of people. For example, we might consider a case where only 1-3 people are infected/infectious. The the deterministic ODE-based approach can produce absurd results such as “half a person is infected” and does not capture the underlying randomness of this scenario. To address this issue, there is a stochastic approach that can be used instead of the deterministic, ODE-based appproach. In the stochastic model, flow rates are used to calculate entry/exit/transition probabilities, and the flow values are sampled from these probabilities at each timestep.

On this page:

See below for an example of using the stochastic approach. First, let’s define a function to build a model for these examples:

[1]
from summer import CompartmentalModel


def build_model(timestep, infectious_seed) -> CompartmentalModel:
    """
    Returns the SIR model, ready to run.
    """
    model = CompartmentalModel(
        times=[1990, 2020],
        compartments=["S", "I", "R"],
        infectious_compartments=["I"],
        timestep=timestep,
    )
    # Add 1000 susceptible people + some infectious people.
    model.set_initial_population(distribution={"S": 1000, "I": infectious_seed})

    # Add flows between the compartments.
    # Susceptible people get infected - contact rate is an arbitrary number.
    model.add_infection_frequency_flow(name="infection", contact_rate=2, source="S", dest="I")
    # Infectious people take 2 years, on average, to recover.
    model.add_transition_flow(name="recovery", fractional_rate=0.5, source="I", dest="R")
    # Add an infection-specific death flow to the I compartment.
    model.add_death_flow(name="infection_death", death_rate=0.05, source="I")
    return model

Comparison with the ODE solver

Here you can see the results of a stochastic model compared to a ODE-based model.

[2]
from matplotlib import pyplot

# Build and run the deterministic, ODE-based model.
model_ode = build_model(timestep=0.1, infectious_seed=20)
model_ode.run()

# Build and run the stochastic model.
model_st = build_model(timestep=0.1, infectious_seed=20)
model_st.run_stochastic()

# Plot the results for both.
pyplot.style.use("ggplot")
fig, (ax1, ax2) = pyplot.subplots(1, 2, figsize=(15, 5), dpi=100)
for ax, model, name in [(ax1, model_ode, "ODE"), (ax2, model_st, "Stochastic")]:
    values = {str(c): v for c, v in zip(model.compartments, model.outputs.T)}
    ax.set_title(name)
    ax.set_xlabel("Times")
    for plot_vals in values.values():
        ax.plot(model.times, plot_vals)

    ax.legend(list(values.keys()), loc="center right")
../_images/examples_5-stochastic-solver_3_0.png

A technical point to note:

  • The ODE-based model does not use the timestep parameter to perform the integration, rather it uses a timestep that is passed to the model.run() method, or an adaptive timstep. The timestep defined during model construction is only used for interpolating the results. So the results are not very sensitive to changes in the requested timestep.

  • The stochastic method does use the timestep parameter for running the integration, and the results are sensitive to the changes in the timestep.

Simulating disease extinction

Because the stochastic model deals with a discrete, rather than continuous, number of people, we can simulate the number of infected people going to 0, rather than 1/25th of a person.

[3]
NUM_TRIALS = 500

def simulate_extinction(timestep, infectious_seed):
    # Create a model.
    model = build_model(timestep, infectious_seed)

    # Run the model many times
    extinct_times = []  # Track when the number of infections goes to zero.
    infected_timeseries = []  # Track number of infections over time.
    for i in range(NUM_TRIALS):
        # Run the model
        model.run_stochastic()
        infected = model.outputs[:, 1].T
        infected_timeseries.append([model.times, infected])
        # Determine when infections went to zero.
        for s,t in zip(infected, model.times):
            if s == 0:
                extinct_times.append(t)
                break

    # Plot timeseries.
    fig, ax = pyplot.subplots(1, 1, figsize=(15, 5), dpi=100)
    ax.set_title("Number of infected people")
    ax.set_xlabel("Time")
    for t, v in infected_timeseries:
        ax.plot(t, v, alpha=0.3)

    # Plot number of simulated extinctions per year.
    fig, ax = pyplot.subplots(1, 1, figsize=(15, 5), dpi=100)
    ax.set_title("Number of extinctions")
    ax.set_xlabel("Time")
    ax.hist(extinct_times, bins=range(1990, 2026))

Here we simulate 1000 trials where we started with 1 infected person, who could die or recover before they infect anyone else. As you can see below, the disease dies out early in ~20-30% of cases.

[4]
simulate_extinction(timestep=0.1, infectious_seed=1)
../_images/examples_5-stochastic-solver_8_0.png
../_images/examples_5-stochastic-solver_8_1.png

Note that increasing the chosen timestep will reduce the probability that the disease will die out early.

[5]
simulate_extinction(timestep=1, infectious_seed=1)
../_images/examples_5-stochastic-solver_10_0.png
../_images/examples_5-stochastic-solver_10_1.png

Increasing the number of initially infectious people will also reduce the chance that the disease will die out.

[6]
simulate_extinction(timestep=0.1, infectious_seed=2)
../_images/examples_5-stochastic-solver_12_0.png
../_images/examples_5-stochastic-solver_12_1.png

How it works

At a high (and simplified) level, the stochastic model works by looping through each timestep and:

  • Calculating the flow rates, given the time and current compartment sizes, in the same manner as with an ODE-based model

  • Converting the flow rates into transition probabilities

  • Sampling from those probabilities to get the change in compartment sizes

In this approach, entry flows are treated differently to transition flows. To start, let’s look at how entry flows are sampled.

Entry flows

An entry flow is a flow that comes from outside of the system, such as a birth or a person entering the region from the outside. Entry flow rates are first calculated using the same approach as we would with the ODE solver.

For example we might have a crude birth rate of 0.02 births/person/year and a population of 1e6 people, giving us 2e4 people born per year. This calculation is the same for both the deterministic (ODE) and stochastic solver.

In the stochastic model we interpret this flow rate (2e4 people born per year) as the mean of a Poisson distributon, which is used to sample the number of entering individuals. We use np.random.poisson in practice.

If a smaller timestep is chosen, eg. 0.1, or 1/10th of a year, then the mean number of arrivals is scaled down by a factor of 10.

In practice, we sum up all the entry flow rates into a net flow. Eg. a birth flow of 2e4 people/year and an importation flow of 1e4 people/year into the same compartment will be summed into a net entry flow of 3e4 people/year. This net flow of 3e4 people/year is what is used as the mean of the Poisson distribution.

Let’s have a look at how this process looks at different population sizes and timesteps.

[7]
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(1337)

START_TIME = 0  # Year 0
END_TIME = 20  # Year 20

def sample_entry_flows(start_pop, timestep):
    """Returns sampled population growth over time period"""
    # Get times
    time_period = END_TIME - START_TIME + 1
    num_steps = time_period / timestep
    times = np.linspace(START_TIME, END_TIME, num=int(num_steps))

    # Define initial conditions
    initial_conditions = np.array([start_pop])

    # Define outputs
    outputs = np.zeros((int(num_steps)))
    outputs[0] = initial_conditions

    # Calculate outputs for each timestep
    for t_idx, t in enumerate(times):
        if t_idx == 0:
            continue

        population = outputs[t_idx - 1]

        # Get rate of births
        crude_birth_rate = 0.02  # Births / person / year
        birth_rate = population * crude_birth_rate

        # Sum entry flows to this compartment and sample number of entries.
        net_entry_rate = birth_rate  # ... + import_rate, etc.
        lam = net_entry_rate * timestep
        number_entries = np.random.poisson(lam)
        outputs[t_idx] = population + number_entries

    return times, outputs

# Try a bunch of different population sizes and timestep combinations.
scenarios = [
    {'timestep': 1, 'start_pop': 10},
    {'timestep': 0.1, 'start_pop': 10},
    {'timestep': 1, 'start_pop': 100},
    {'timestep': 0.1, 'start_pop': 100},
    {'timestep': 1, 'start_pop': 100000},
    {'timestep': 0.1, 'start_pop': 100000},

]

# Plot the results
fig, axes = pyplot.subplots(3, 2, figsize=(16, 12), dpi=100, tight_layout=True)
for ax, sc in zip(axes.flatten(), scenarios):
    times, outputs = sample_entry_flows(**sc)
    ax.plot(times, outputs)
    ax.set_title(f"Population starting from {sc['start_pop']}, timestep: {sc['timestep']}")
    start, end = ax.get_xlim()
    ax.xaxis.set_ticks(np.arange(start + 1, end, 1))

plt.show()
../_images/examples_5-stochastic-solver_14_0.png

Transition and exit flows

A transition flow is a person leaving one compartment and entering another. An exit flow is a person leaving compartment and exiting the system (ie. no destination). We calculate our transition and exit flows using a source population and a fractional rate.

For example, we might have 100 people infected with a disease that has a recovery rate of 0.1. This gives us a transition flow rate of 100 * 0.1 = 10 people/year. This is the rate of change used by the ODE solver to evaluate the model.

This approach assumes that the length of time spent by a person in the source compartment is a random variable with an exponential distribution. A recovery rate of 0.1 is the same as saying that the mean person spends 1 / 0.1 = 10 years in the infected compartment.

For transition and exit flows we use this implied exponential distribution to calculate the probability that an individual will leave their compartment via a given flow during a given timestep. You can read the details of how these exit probabilities are calculated here. This process is used to compute the transition matrix defining the underlying non-homogeneous Markov process.

Once we know the probabilites of all the outcomes (eg. stay, leave by flow A, leave by flow B… etc) then we can get a sample of how many people leave the compartment via various flows using np.random.multinomial,

Let’s have a look at a simplified example of this calculation, where infected people can either recover or die:

[8]
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(1337)

START_TIME = 0  # Year 0
END_TIME = 20  # Year 20

def sample_flows(start_pop, timestep):
    """Returns sampled flows over time period"""
    # Get times
    time_period = END_TIME - START_TIME + 1
    num_steps = time_period / timestep
    times = np.linspace(START_TIME, END_TIME, num=int(num_steps))

    # Define initial conditions, 2 compartments (I, R)
    initial_conditions = np.array([start_pop, 0])

    # Define outputs
    outputs = np.zeros((int(num_steps), 2))
    outputs[0] = initial_conditions

    # Calculate outputs for each timestep
    for t_idx, t in enumerate(times):
        if t_idx == 0:
            continue

        compartment_sizes = outputs[t_idx - 1]
        num_infected = compartment_sizes[0]

        # Get rate of recovery (in people / year)
        # This is the same process as we would use in an ODE-based calculation.
        recovery_rate = 0.1  # Mean person takes 10 years to recover.
        recovery_flow_rate = num_infected * recovery_rate  # People / year

        # Get rate of infection deaths (in people / year)
        # This is the same process as we would use in an ODE-based calculation.
        death_rate = 0.05  # 5% of infected population die per timestep.
        death_flow_rate = num_infected * death_rate  # People / year

        # Get a total exit flow rate
        total_flow_rates = recovery_flow_rate + death_flow_rate
        if total_flow_rates <= 0:
            # Handle case when everyone has recovered, no flows remain.
            outputs[t_idx] = compartment_sizes
            continue

        # This seems silly in this example, but due to some implementation details
        # we calculate the flow rates as people/year then renormalize them with
        # the source compartment sizes. This is so we can share single method to
        # calculate flow rates between the stochastic and ODE solvers.
        recovery_flow_rate /= num_infected
        death_flow_rate /= num_infected
        total_flow_rates /= num_infected

        # Calculate probability of staying or leaving compartment
        p_stay = np.exp(-1 * total_flow_rates * timestep)
        p_leave = 1 - p_stay
        p_recover = p_leave * (recovery_flow_rate / total_flow_rates)
        p_die = p_leave * (death_flow_rate / total_flow_rates)

        # Sample from probabilities of outcomes for each individual with multinomial
        flow_probs = np.array([p_recover, p_die, p_stay])
        num_recovered, num_dead, _ = np.random.multinomial(num_infected, flow_probs)

        # Apply changes to compartments to get next state
        comp_chages = np.zeros_like(compartment_sizes)
        # Apply recovery flow by moving people from I to R
        comp_chages[0] -= num_recovered
        comp_chages[1] += num_recovered
        # Apply death flow by removing people from I
        comp_chages[0] -= num_dead
        outputs[t_idx] = compartment_sizes + comp_chages

    return times, outputs

# Try a bunch of different population sizes and timestep combinations.
scenarios = [
    {'timestep': 1, 'start_pop': 10},
    {'timestep': 0.1, 'start_pop': 10},
    {'timestep': 1, 'start_pop': 100},
    {'timestep': 0.1, 'start_pop': 100},
    {'timestep': 1, 'start_pop': 100000},
    {'timestep': 0.1, 'start_pop': 100000},

]

# Plot the results
fig, axes = pyplot.subplots(3, 2, figsize=(16, 12), dpi=100, tight_layout=True)
for ax, sc in zip(axes.flatten(), scenarios):
    times, outputs = sample_flows(**sc)
    ax.plot(times, outputs)
    ax.set_title(f"Population starting from {sc['start_pop']}, timestep: {sc['timestep']}")
    start, end = ax.get_xlim()
    ax.xaxis.set_ticks(np.arange(start + 1, end, 1))
    ax.legend(['Infected', 'Recovered'],  loc='upper right')

plt.show()
../_images/examples_5-stochastic-solver_16_0.png
[ ]