Multi-strata Models with Multiple Movement Models

SCENARIO:

Exercise 1

Build a multi-strata RUME, with three strata based on age classes (0-19, 20-59, 60-79). Each strata will be build as a GPM object, which requires an IPM and a MM. This allows for unique IPM and/or MM per strata. At first, we’ll assume that each strata follows an SIRS IPM. For simplicity, we’ll do no movement (i.e., a single-node simulation) to start. And, we’ll initialize the infection in the 20-59 strata, but no initial infection in the other two strata.

# Construct a multistrata RUME

import matplotlib.pyplot as plt
import numpy as np
from sympy import Max

from epymorph.kit import *
from epymorph.adrio import acs5, us_tiger

class MyRume(MultiStrataRUMEBuilder):
    strata = [
        GPM(
            name="age_00-19",
            ipm=ipm.SIRS(),
            mm=mm.No(),
            init=init.NoInfection(),
        ),
        GPM(
            name="age_20-59",
            ipm=ipm.SIRS(),
            mm=mm.No(),
            init=init.SingleLocation(location=0, seed_size=100),
        ),
        GPM(
            name="age_60-79",
            ipm=ipm.SIRS(),
            mm=mm.No(),
            init=init.NoInfection(),
        ),
    ]

    meta_requirements = [
        AttributeDef("beta_12", float, Shapes.TxN),
        AttributeDef("beta_13", float, Shapes.TxN),
        AttributeDef("beta_21", float, Shapes.TxN),
        AttributeDef("beta_23", float, Shapes.TxN),
        AttributeDef("beta_31", float, Shapes.TxN),
        AttributeDef("beta_32", float, Shapes.TxN),
    ]

    def meta_edges(self, symbols: MultiStrataModelSymbols) -> list[TransitionDef]:
        # extract compartment symbols by strata
        S_1, I_1, R_1 = symbols.strata_compartments("age_00-19")
        S_2, I_2, R_2 = symbols.strata_compartments("age_20-59")
        S_3, I_3, R_3 = symbols.strata_compartments("age_60-79")

        # extract compartment totals by strata
        N_1 = Max(1, S_1 + I_1 + R_1)
        N_2 = Max(1, S_2 + I_2 + R_2)
        N_3 = Max(1, S_3 + I_3 + R_3)

        # extract meta attributes
        beta_12, beta_13, beta_21, beta_23, beta_31, beta_32 = (
            symbols.all_meta_requirements
        )

        return [
            edge(S_1, I_1, rate=S_1 * beta_12 * I_2 / N_2),  # 2 infects 1
            edge(S_1, I_1, rate=S_1 * beta_13 * I_3 / N_3),  # 3 infects 1
            edge(S_2, I_2, rate=S_2 * beta_21 * I_1 / N_1),  # 1 infects 2
            edge(S_2, I_2, rate=S_2 * beta_23 * I_3 / N_3),  # 3 infects 2
            edge(S_3, I_3, rate=S_3 * beta_31 * I_1 / N_1),  # 1 infects 3
            edge(S_3, I_3, rate=S_3 * beta_32 * I_2 / N_2),  # 2 infects 3
        ]

Assign parameter values, and use the ACS5 Population by Age ADRIO to specify the number of people in each population strata.

single_scope = StateScope.in_states(["AZ"], year=2020)

rume = MyRume().build(
    scope=single_scope,
    time_frame=TimeFrame.of("2020-01-01", 180),
    params={
        # IPM params
        "gpm:age_00-19::ipm::beta": 0.05,
        "gpm:age_20-59::ipm::beta": 0.20,
        "gpm:age_60-79::ipm::beta": 0.35,
        "*::ipm::gamma": 1 / 10,
        "*::ipm::xi": 1 / 90,
        "meta::ipm::beta_12": 0.05,
        "meta::ipm::beta_13": 0.05,
        "meta::ipm::beta_21": 0.20,
        "meta::ipm::beta_23": 0.20,
        "meta::ipm::beta_31": 0.35,
        "meta::ipm::beta_32": 0.35,

        # ADRIOs for population by age
        "*::*::population_by_age_table": acs5.PopulationByAgeTable(),
        "gpm:age_00-19::*::population": acs5.PopulationByAge(0, 19),
        "gpm:age_20-59::*::population": acs5.PopulationByAge(20, 59),
        "gpm:age_60-79::*::population": acs5.PopulationByAge(60, 79),
    },
)

Visualize the meta-IPM diagram to understand the rates of flow between classes.

rume.ipm.diagram()

Simultate the model, as usual.

sim = BasicSimulator(rume)
with sim_messaging(live=False):
    out = sim.run(
        rng_factory=default_rng(2)
    )

Loading gpm:age_00-19::init::population_by_age_table (epymorph.adrio.acs5.PopulationByAgeTable):
  |####################| 100%  (11.126s)
Loading gpm:age_00-19::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Loading gpm:age_20-59::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Loading gpm:age_60-79::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Running simulation (BasicSimulator):
• 2020-01-01 to 2020-06-28 (180 days)
• 1 geo nodes
  |####################| 100% 
Runtime: 0.100s

Plot the infectious class dynamics for all age classes. You can use the wildcard ’*’ symbol in the quantity selection method.

out.plot.line(
    geo=out.rume.scope.select.all(),
    time=out.rume.time_frame.select.all(),
    quantity=out.rume.ipm.select.compartments("I*"),
    title="Compartment values per strata",
)

Plot dynamics of all compartments, but only for the 0-19 age class.

out.plot.line(
    geo=out.rume.scope.select.all(),
    time=out.rume.time_frame.select.all(),
    quantity=out.rume.ipm.select.compartments("*_age_00-19"),
    title="Compartment values for youngest strata",
)

Exercise 2

Now create a multi-strata RUME in which the age classes have different IPMs. Specifically, we’ll assume that the oldest age class (60-79) can go to the hospital due to infection, so we want an SIRH model for that strata. Not much needs to change with the model specification, so it’s mostly copy-and-paste.

class MyRume(MultiStrataRUMEBuilder):
    strata = [
        GPM(
            name="age_00-19",
            ipm=ipm.SIRS(),
            mm=mm.No(),
            init=init.NoInfection(),
        ),
        GPM(
            name="age_20-59",
            ipm=ipm.SIRS(),
            mm=mm.No(),
            init=init.SingleLocation(location=0, seed_size=100),
        ),
        GPM(
            name="age_60-79",
            ipm=ipm.SIRH(),
            mm=mm.No(),
            init=init.NoInfection(),
        ),
    ]

    meta_requirements = [
        AttributeDef("beta_12", float, Shapes.TxN),
        AttributeDef("beta_13", float, Shapes.TxN),
        AttributeDef("beta_21", float, Shapes.TxN),
        AttributeDef("beta_23", float, Shapes.TxN),
        AttributeDef("beta_31", float, Shapes.TxN),
        AttributeDef("beta_32", float, Shapes.TxN),
    ]

    def meta_edges(self, symbols: MultiStrataModelSymbols) -> list[TransitionDef]:
        # extract compartment symbols by strata
        S_1, I_1, R_1 = symbols.strata_compartments("age_00-19")
        S_2, I_2, R_2 = symbols.strata_compartments("age_20-59")
        S_3, I_3, R_3, H_3 = symbols.strata_compartments("age_60-79")

        # extract compartment totals by strata
        N_1 = Max(1, S_1 + I_1 + R_1)
        N_2 = Max(1, S_2 + I_2 + R_2)
        N_3 = Max(1, S_3 + I_3 + R_3 + H_3)

        # extract meta attributes
        beta_12, beta_13, beta_21, beta_23, beta_31, beta_32 = (
            symbols.all_meta_requirements
        )

        return [
            edge(S_1, I_1, rate=S_1 * beta_12 * I_2 / N_2),  # 2 infects 1
            edge(S_1, I_1, rate=S_1 * beta_13 * I_3 / N_3),  # 3 infects 1
            edge(S_2, I_2, rate=S_2 * beta_21 * I_1 / N_1),  # 1 infects 2
            edge(S_2, I_2, rate=S_2 * beta_23 * I_3 / N_3),  # 3 infects 2
            edge(S_3, I_3, rate=S_3 * beta_31 * I_1 / N_1),  # 1 infects 3
            edge(S_3, I_3, rate=S_3 * beta_32 * I_2 / N_2),  # 2 infects 3
        ]

Build the RUME by assigning the correct attributes. Notice we use gpm:age_60-79::ipm::... to assign parameter values relevant to that specific strata’s IPM.

rume = MyRume().build(
    scope=single_scope,
    time_frame=TimeFrame.of("2020-01-01", 180),
    params={
        # IPM params
        "gpm:age_00-19::ipm::beta": 0.05,
        "gpm:age_20-59::ipm::beta": 0.20,
        "gpm:age_60-79::ipm::beta": 0.35,
        "*::ipm::gamma": 1 / 10,
        "*::ipm::xi": 1 / 90,
        "meta::ipm::beta_12": 0.05,
        "meta::ipm::beta_13": 0.05,
        "meta::ipm::beta_21": 0.20,
        "meta::ipm::beta_23": 0.20,
        "meta::ipm::beta_31": 0.35,
        "meta::ipm::beta_32": 0.35,

        # SIRH IPM
        "gpm:age_60-79::ipm::hospitalization_prob": 0.1,
        "gpm:age_60-79::ipm::hospitalization_duration": 7.0,

        # ADRIO 
        "*::*::population_by_age_table": acs5.PopulationByAgeTable(),
        "gpm:age_00-19::*::population": acs5.PopulationByAge(0, 19),
        "gpm:age_20-59::*::population": acs5.PopulationByAge(20, 59),
        "gpm:age_60-79::*::population": acs5.PopulationByAge(60, 79),
    },
)

Visualize the meta-IPM diagram to understand the rates of flow between classes.

rume.ipm.diagram()

Simultate the model, as usual.

sim = BasicSimulator(rume)
with sim_messaging(live=False):
    out = sim.run(
        rng_factory=default_rng(2)
    )

Loading gpm:age_00-19::init::population_by_age_table (epymorph.adrio.acs5.PopulationByAgeTable):
  |####################| 100%  (0.551s)
Loading gpm:age_00-19::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Loading gpm:age_20-59::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Loading gpm:age_60-79::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Running simulation (BasicSimulator):
• 2020-01-01 to 2020-06-28 (180 days)
• 1 geo nodes
  |####################| 100% 
Runtime: 0.107s

Plot dynamics of all compartments, but only for the oldest age class, to see that it now has hospitalizations.

out.plot.line(
    geo=out.rume.scope.select.all(),
    time=out.rume.time_frame.select.all(),
    quantity=out.rume.ipm.select.compartments("*_age_60-79"),
    title="Compartment values for elder strata",
)

Exercise 3

In this most complicated example, we’re going to create a multi-strata simulation across multiple nodes, with multiple IPM, and with multiple movement models.

Create the multi-strata IPM, assigning the SIRH model to the oldest strata. For movement models, assume that the two youngest strata use a Centroids movement, while the oldest strata follows a commuter style movement.

from epymorph.data.mm.pei import Pei as CommuterMM

class MyRume(MultiStrataRUMEBuilder):
    strata = [
        GPM(
            name="age_00-19",
            ipm=ipm.SIRS(),
            mm=mm.Centroids(),
            init=init.NoInfection(),
        ),
        GPM(
            name="age_20-59",
            ipm=ipm.SIRS(),
            mm=mm.Centroids(),
            init=init.SingleLocation(location=0, seed_size=100),
        ),
        GPM(
            name="age_60-79",
            ipm=ipm.SIRH(),
            mm=CommuterMM(),
            init=init.NoInfection(),
        ),
    ]

    meta_requirements = [
        AttributeDef("beta_12", float, Shapes.TxN),
        AttributeDef("beta_13", float, Shapes.TxN),
        AttributeDef("beta_21", float, Shapes.TxN),
        AttributeDef("beta_23", float, Shapes.TxN),
        AttributeDef("beta_31", float, Shapes.TxN),
        AttributeDef("beta_32", float, Shapes.TxN),
    ]

    def meta_edges(self, symbols: MultiStrataModelSymbols) -> list[TransitionDef]:
        # extract compartment symbols by strata
        S_1, I_1, R_1 = symbols.strata_compartments("age_00-19")
        S_2, I_2, R_2 = symbols.strata_compartments("age_20-59")
        S_3, I_3, R_3, H_3 = symbols.strata_compartments("age_60-79")

        # extract compartment totals by strata
        N_1 = Max(1, S_1 + I_1 + R_1)
        N_2 = Max(1, S_2 + I_2 + R_2)
        N_3 = Max(1, S_3 + I_3 + R_3 + H_3)

        # extract meta attributes
        beta_12, beta_13, beta_21, beta_23, beta_31, beta_32 = (
            symbols.all_meta_requirements
        )

        return [
            edge(S_1, I_1, rate=S_1 * beta_12 * I_2 / N_2),  # 2 infects 1
            edge(S_1, I_1, rate=S_1 * beta_13 * I_3 / N_3),  # 3 infects 1
            edge(S_2, I_2, rate=S_2 * beta_21 * I_1 / N_1),  # 1 infects 2
            edge(S_2, I_2, rate=S_2 * beta_23 * I_3 / N_3),  # 3 infects 2
            edge(S_3, I_3, rate=S_3 * beta_31 * I_1 / N_1),  # 1 infects 3
            edge(S_3, I_3, rate=S_3 * beta_32 * I_2 / N_2),  # 2 infects 3
        ]

Create the RUME with its requirements. Be sure to specify the strata-specific IPM parameters, as well as the strata-specific MM parameters. Use the LODES8 ADRIO for age-specific commuter data.

from epymorph.adrio import lodes

multi_scope = StateScope.in_states(["AZ", "NM", "CO", "UT"], year=2020)

rume = MyRume().build(
    scope=multi_scope,
    time_frame=TimeFrame.of("2020-01-01", 180),
    params={
        # IPM params
        "gpm:age_00-19::ipm::beta": 0.05,
        "gpm:age_20-59::ipm::beta": 0.20,
        "gpm:age_60-79::ipm::beta": 0.35,
        "*::ipm::gamma": 1 / 10,
        "*::ipm::xi": 1 / 90,
        "meta::ipm::beta_12": 0.05,
        "meta::ipm::beta_13": 0.05,
        "meta::ipm::beta_21": 0.20,
        "meta::ipm::beta_23": 0.20,
        "meta::ipm::beta_31": 0.35,
        "meta::ipm::beta_32": 0.35,

        # SIRH IPM
        "gpm:age_60-79::ipm::hospitalization_prob": 0.1,
        "gpm:age_60-79::ipm::hospitalization_duration": 7.0,

        # MM params
        "gpm:age_00-19::mm::phi": 20.0,
        "gpm:age_20-59::mm::phi": 40.0,
        "gpm:age_60-79::mm::move_control": 0.9,
        "gpm:age_60-79::mm::theta": 0.1,
        "gpm:age_60-79::mm::commuters": lodes.CommutersByAge("55 and Over"),

        # ADRIO 
        "*::*::centroid": us_tiger.InternalPoint(),
        "*::*::population_by_age_table": acs5.PopulationByAgeTable(),
        "gpm:age_00-19::*::population": acs5.PopulationByAge(0, 19),
        "gpm:age_20-59::*::population": acs5.PopulationByAge(20, 59),
        "gpm:age_60-79::*::population": acs5.PopulationByAge(60, 79),
    },
)

print(rume.params_description())

gpm:age_00-19::ipm::beta (type: float, shape: TxN)
    infectivity

gpm:age_00-19::ipm::gamma (type: float, shape: TxN)
    progression from infected to recovered

gpm:age_00-19::ipm::xi (type: float, shape: TxN)
    progression from recovered to susceptible

gpm:age_20-59::ipm::beta (type: float, shape: TxN)
    infectivity

gpm:age_20-59::ipm::gamma (type: float, shape: TxN)
    progression from infected to recovered

gpm:age_20-59::ipm::xi (type: float, shape: TxN)
    progression from recovered to susceptible

gpm:age_60-79::ipm::beta (type: float, shape: TxN)
    infectivity

gpm:age_60-79::ipm::gamma (type: float, shape: TxN)
    recovery rate

gpm:age_60-79::ipm::xi (type: float, shape: TxN)
    immune waning rate

gpm:age_60-79::ipm::hospitalization_prob (type: float, shape: TxN)
    a ratio of cases which are expected to require hospitalization

gpm:age_60-79::ipm::hospitalization_duration (type: float, shape: TxN)
    the mean duration of hospitalization, in days

meta::ipm::beta_12 (type: float, shape: TxN)

meta::ipm::beta_13 (type: float, shape: TxN)

meta::ipm::beta_21 (type: float, shape: TxN)

meta::ipm::beta_23 (type: float, shape: TxN)

meta::ipm::beta_31 (type: float, shape: TxN)

meta::ipm::beta_32 (type: float, shape: TxN)

gpm:age_00-19::mm::population (type: int, shape: N)
    The total population at each node.

gpm:age_00-19::mm::centroid (type: [(longitude, float), (latitude, float)], shape: N)
    The centroids for each node as (longitude, latitude) tuples.

gpm:age_00-19::mm::phi (type: float, shape: S, default: 40.0)
    Influences the distance that movers tend to travel.

gpm:age_00-19::mm::commuter_proportion (type: float, shape: S, default: 0.1)
    The proportion of the total population which commutes.

gpm:age_00-19::init::population (type: int, shape: N)
    The population at each geo node.

gpm:age_20-59::mm::population (type: int, shape: N)
    The total population at each node.

gpm:age_20-59::mm::centroid (type: [(longitude, float), (latitude, float)], shape: N)
    The centroids for each node as (longitude, latitude) tuples.

gpm:age_20-59::mm::phi (type: float, shape: S, default: 40.0)
    Influences the distance that movers tend to travel.

gpm:age_20-59::mm::commuter_proportion (type: float, shape: S, default: 0.1)
    The proportion of the total population which commutes.

gpm:age_20-59::init::population (type: int, shape: N)
    The population at each geo node.

gpm:age_60-79::mm::commuters (type: int, shape: NxN)
    A node-to-node commuters marix.

gpm:age_60-79::mm::move_control (type: float, shape: S, default: 0.9)
    A factor which modulates the number of commuters by conducting a
    binomial draw with this probability and the expected commuters
    from the commuters matrix.

gpm:age_60-79::mm::theta (type: float, shape: S, default: 0.1)
    A factor which allows for randomized movement by conducting a
    poisson draw with this factor times the average number of
    commuters between two nodes from the commuters matrix.

gpm:age_60-79::init::population (type: int, shape: N)
    The population at each geo node.

sim = BasicSimulator(rume)
with sim_messaging(live=False):
    out = sim.run()

Loading gpm:age_00-19::mm::population_by_age_table (epymorph.adrio.acs5.PopulationByAgeTable):
  |####################| 100%  (0.571s)
Loading gpm:age_00-19::mm::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Loading gpm:age_00-19::mm::centroid (epymorph.adrio.us_tiger.InternalPoint):
  |####################| 100%  (1.641s)
Loading gpm:age_20-59::mm::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Loading gpm:age_60-79::mm::commuters (epymorph.adrio.lodes.CommutersByAge):
  |####################| 100%  (40.104s)
Loading gpm:age_60-79::init::population (epymorph.adrio.acs5.PopulationByAge):
  |####################| 100%  (0.000s)
Running simulation (BasicSimulator):
• 2020-01-01 to 2020-06-28 (180 days)
• 4 geo nodes
  |####################| 100% 
Runtime: 0.771s

Create a (busy) plot with all of the compartments in the model, across all of the nodes.

out.plot.line(
    geo=out.rume.scope.select.all(),
    time=out.rume.time_frame.select.all().group("day").agg(),
    quantity=out.rume.ipm.select.compartments(),
    title="",
    legend="outside",
)

Create a table to summarize all the compartments and events with quantiles.

out.table.quantiles(
    quantiles=[0.025, 0.5, 0.975],
    geo=out.rume.scope.select.all(),
    time=out.rume.time_frame.select.all(),
    quantity=out.rume.ipm.select.all(),
)

	geo	quantity	other strata	0.025	0.5	0.975
0	AZ	S_age_00-19		619660.125	849743.0	1836807.325
1	AZ	I_age_00-19		41.825	94596.5	503575.450
2	AZ	R_age_00-19		7.850	889452.5	1058452.825
3	AZ	S_age_20-59		175365.450	669430.5	3646494.525
4	AZ	I_age_20-59		405.400	298161.5	1761381.350
...	...	...	...	...	...	...
103	UT	S_age_00-19 → I_age_00-19	age_60-79	0.000	869.5	9021.100
104	UT	S_age_20-59 → I_age_20-59	age_00-19	0.000	1124.0	7054.350
105	UT	S_age_20-59 → I_age_20-59	age_60-79	0.000	1934.0	28810.725
106	UT	S_age_60-79 → I_age_60-79	age_00-19	0.000	314.5	1849.100
107	UT	S_age_60-79 → I_age_60-79	age_20-59	0.000	493.5	5714.625

108 rows × 6 columns

Create a map of total new hospitalizations across the simulation per node.

out.map.choropleth(
    geo=out.rume.scope.select.all(),
    time=out.rume.time_frame.select.all().agg(events="sum"),
    quantity=out.rume.ipm.select.events("I*->H*"),
)