How to plot infection rates by contact models

The following simulation is identical to the simulation in the tutorial on contact models. The heatmap shows the contribution of contact models to new infections for each day in the simulation. The values are shares in relation to the whole population.

import warnings

import holoviews as hv
import numpy as np
import pandas as pd

import sid
from sid.plotting import plot_infection_rates_by_contact_models

warnings.filterwarnings(
    "ignore", message="indexing past lexsort depth may impact performance."
)

Matplotlib is building the font cache; this may take a moment.

n_individuals = 10_000
available_ages = [
    "0-9",
    "10-19",
    "20-29",
    "30-39",
    "40-49",
    "50-59",
    "60-69",
    "70-79",
    "80-100",
]

ages = np.random.choice(available_ages, size=n_individuals)
regions = np.random.choice(["North", "East", "South", "West"], size=n_individuals)
hh_id = pd.Series(np.random.choice(int(n_individuals / 1.6), size=n_individuals))

initial_states = pd.DataFrame(
    {"age_group": ages, "region": regions, "hh_id": hh_id}
).astype("category")
initial_states.head(5)

	age_group	region	hh_id
0	20-29	East	4568
1	70-79	South	5077
2	60-69	South	2907
3	0-9	South	3679
4	10-19	West	2237

contact_models = {}


def random_encounters(states, params, seed):
    np.random.seed(seed)
    contacts = np.random.choice(np.arange(3), size=states.shape[0])
    return pd.Series(index=states.index, data=contacts)


contact_models["random_encounters"] = {
    "model": random_encounters,
    "assort_by": ["age_group", "region"],
    "is_recurrent": False,
}


def meet_household(states, params, seed):
    return pd.Series(index=states.index, data=True)


contact_models["meet_household"] = {
    "model": meet_household,
    "assort_by": "hh_id",
    "is_recurrent": True,
}

params = sid.load_epidemiological_parameters()
params.loc[("assortative_matching", "random_encounters", "age_group"), "value"] = 0.2
params.loc[("assortative_matching", "random_encounters", "region"), "value"] = 0.9
params.loc[("infection_prob", "random_encounters", "random_encounters"), "value"] = 0.1
params.loc[("infection_prob", "meet_household", "meet_household"), "value"] = 0.15

simulate = sid.get_simulate_func(
    params=params,
    initial_states=initial_states,
    contact_models=contact_models,
    duration={"start": "2020-03-01", "end": "2020-04-01"},
    saved_columns={"other": ["channel_infected_by_contact"]},
    seed=144,
)

result = simulate(params)

Start the simulation...
2020-04-01: 100%|██████████| 32/32 [00:04<00:00,  7.41it/s]

heatmap = plot_infection_rates_by_contact_models(result["time_series"], unit="share")
heatmap