Beispiel #1
0
def test_strain__with_infectious_multipliers_and_heterogeneous_mixing():
    """
    Test infectious multiplier and flow rate calculations for
    3 strains which have different infectiousness levels plus a seperate
    stratification which has a mixing matrix.
    """
    model = CompartmentalModel(times=[0, 5],
                               compartments=["S", "I", "R"],
                               infectious_compartments=["I"])
    model.set_initial_population(distribution={"S": 900, "I": 100})
    contact_rate = 0.2
    model.add_infection_frequency_flow("infection", contact_rate, "S", "I")

    age_strat = Stratification("age", ["child", "adult"], ["S", "I", "R"])
    age_strat.set_population_split({
        "child": 0.6,  # 600 people
        "adult": 0.4,  # 400 people
    })
    # Higher mixing among adults or children,
    # than between adults or children.
    age_strat.set_mixing_matrix(np.array([[1.5, 0.5], [0.5, 1.5]]))
    model.stratify_with(age_strat)

    strain_strat = StrainStratification("strain", ["a", "b", "c"], ["I"])
    strain_strat.set_population_split({
        "a": 0.7,  # 70 people
        "b": 0.2,  # 20 people
        "c": 0.1,  # 10 people
    })
    strain_strat.add_infectiousness_adjustments(
        "I",
        {
            "a": adjust.Multiply(0.5),  # 0.5x as susceptible
            "b": adjust.Multiply(3),  # 3x as susceptible
            "c": adjust.Multiply(2),  # 2x as susceptible
        },
    )
    model.stratify_with(strain_strat)

    # Do pre-run force of infection calcs.
    model._prepare_to_run()
    assert_array_equal(model._compartment_infectiousness["a"],
                       np.array([0, 0, 0.5, 0, 0, 0.5, 0, 0, 0, 0]))
    assert_array_equal(model._compartment_infectiousness["b"],
                       np.array([0, 0, 0, 3, 0, 0, 3, 0, 0, 0]))
    assert_array_equal(model._compartment_infectiousness["c"],
                       np.array([0, 0, 0, 0, 2, 0, 0, 2, 0, 0]))
    # 0 for child, 1 for adult
    assert model._category_lookup == {
        0: 0,
        1: 1,
        2: 0,
        3: 0,
        4: 0,
        5: 1,
        6: 1,
        7: 1,
        8: 0,
        9: 1
    }
    assert_array_equal(
        model._category_matrix,
        np.array([
            [1, 0, 1, 1, 1, 0, 0, 0, 1, 0],
            [0, 1, 0, 0, 0, 1, 1, 1, 0, 1],
        ]),
    )

    # Do pre-iteration force of infection calcs
    model._prepare_time_step(0, model.initial_population)
    assert_array_equal(model._category_populations, np.array([600, 400]))
    assert_array_equal(
        model._infection_density["a"],
        np.array([0.5 * (42 * 1.5 + 28 * 0.5), 0.5 * (42 * 0.5 + 28 * 1.5)]),
    )
    assert_array_equal(
        model._infection_density["b"],
        np.array([
            3 * (12 * 1.5 + 8 * 0.5),
            3 * (8 * 1.5 + 12 * 0.5),
        ]),
    )
    assert_array_equal(
        model._infection_density["c"],
        np.array([2 * (6 * 1.5 + 4 * 0.5), 2 * (4 * 1.5 + 6 * 0.5)]),
    )
    assert_array_equal(
        model._infection_frequency["a"],
        np.array([
            0.5 * ((42 / 600) * 1.5 + (28 / 400) * 0.5),
            0.5 * ((42 / 600) * 0.5 + (28 / 400) * 1.5),
        ]),
    )
    assert_array_equal(
        model._infection_frequency["b"],
        np.array([
            3 * ((12 / 600) * 1.5 + (8 / 400) * 0.5),
            3 * ((8 / 400) * 1.5 + (12 / 600) * 0.5),
        ]),
    )
    assert_array_equal(
        model._infection_frequency["c"],
        np.array([
            2 * ((6 / 600) * 1.5 + (4 / 400) * 0.5),
            2 * ((4 / 400) * 1.5 + (6 / 600) * 0.5)
        ]),
    )

    # Get multipliers
    sus_child = model.compartments[0]
    sus_adult = model.compartments[1]
    inf_child_a = model.compartments[2]
    inf_child_b = model.compartments[3]
    inf_child_c = model.compartments[4]
    inf_adult_a = model.compartments[5]
    inf_adult_b = model.compartments[6]
    inf_adult_c = model.compartments[7]
    density = model._get_infection_density_multiplier
    freq = model._get_infection_frequency_multiplier
    assert density(sus_child, inf_child_a) == 0.5 * (42 * 1.5 + 28 * 0.5)
    assert density(sus_adult, inf_adult_a) == 0.5 * (42 * 0.5 + 28 * 1.5)
    assert density(sus_child, inf_child_b) == 3 * (12 * 1.5 + 8 * 0.5)
    assert density(sus_adult, inf_adult_b) == 3 * (8 * 1.5 + 12 * 0.5)
    assert density(sus_child, inf_child_c) == 2 * (6 * 1.5 + 4 * 0.5)
    assert density(sus_adult, inf_adult_c) == 2 * (4 * 1.5 + 6 * 0.5)
    assert freq(sus_child,
                inf_child_a) == 0.5 * ((42 / 600) * 1.5 + (28 / 400) * 0.5)
    assert freq(sus_adult,
                inf_adult_a) == 0.5 * ((42 / 600) * 0.5 + (28 / 400) * 1.5)
    assert freq(sus_child,
                inf_child_b) == 3 * ((12 / 600) * 1.5 + (8 / 400) * 0.5)
    assert freq(sus_adult,
                inf_adult_b) == 3 * ((8 / 400) * 1.5 + (12 / 600) * 0.5)
    assert freq(sus_child,
                inf_child_c) == 2 * ((6 / 600) * 1.5 + (4 / 400) * 0.5)
    assert freq(sus_adult,
                inf_adult_c) == 2 * ((4 / 400) * 1.5 + (6 / 600) * 0.5)

    # Get infection flow rates
    flow_to_inf_child_a = 540 * contact_rate * freq(sus_child, inf_child_a)
    flow_to_inf_adult_a = 360 * contact_rate * freq(sus_adult, inf_adult_a)
    flow_to_inf_child_b = 540 * contact_rate * freq(sus_child, inf_child_b)
    flow_to_inf_adult_b = 360 * contact_rate * freq(sus_adult, inf_adult_b)
    flow_to_inf_child_c = 540 * contact_rate * freq(sus_child, inf_child_c)
    flow_to_inf_adult_c = 360 * contact_rate * freq(sus_adult, inf_adult_c)
    expected_flow_rates = np.array([
        -flow_to_inf_child_a - flow_to_inf_child_b - flow_to_inf_child_c,
        -flow_to_inf_adult_a - flow_to_inf_adult_b - flow_to_inf_adult_c,
        flow_to_inf_child_a,
        flow_to_inf_child_b,
        flow_to_inf_child_c,
        flow_to_inf_adult_a,
        flow_to_inf_adult_b,
        flow_to_inf_adult_c,
        0.0,
        0.0,
    ])
    flow_rates = model._get_compartment_rates(model.initial_population, 0)
    assert_allclose(expected_flow_rates, flow_rates, verbose=True)
Beispiel #2
0
def test_strain__with_flow_adjustments():
    """
    Test infectious multiplier and flow rate calculations for
    3 strains which have different flow adjustments.

    These flow adjustments would correspond to some physical process that we're modelling,
    and they should be effectively the same as applying infectiousness multipliers.
    """
    model = CompartmentalModel(times=[0, 5],
                               compartments=["S", "I", "R"],
                               infectious_compartments=["I"])
    model.set_initial_population(distribution={"S": 900, "I": 100})
    contact_rate = 0.2
    model.add_infection_frequency_flow("infection", contact_rate, "S", "I")
    strat = StrainStratification("strain", ["a", "b", "c"], ["I"])
    strat.set_population_split({
        "a": 0.7,  # 70 people
        "b": 0.2,  # 20 people
        "c": 0.1,  # 10 people
    })
    strat.add_flow_adjustments(
        "infection",
        {
            "a": adjust.Multiply(0.5),  # 0.5x as susceptible
            "b": adjust.Multiply(3),  # 3x as susceptible
            "c": adjust.Multiply(2),  # 2x as susceptible
        },
    )
    model.stratify_with(strat)

    # Do pre-run force of infection calcs.
    model._prepare_to_run()
    assert_array_equal(model._compartment_infectiousness["a"],
                       np.array([0, 1, 0, 0, 0]))
    assert_array_equal(model._compartment_infectiousness["b"],
                       np.array([0, 0, 1, 0, 0]))
    assert_array_equal(model._compartment_infectiousness["c"],
                       np.array([0, 0, 0, 1, 0]))
    assert model._category_lookup == {0: 0, 1: 0, 2: 0, 3: 0, 4: 0}
    assert_array_equal(model._category_matrix, np.array([[1, 1, 1, 1, 1]]))

    # Do pre-iteration force of infection calcs
    model._prepare_time_step(0, model.initial_population)
    assert_array_equal(model._category_populations, np.array([1000]))
    assert_array_equal(model._infection_density["a"], np.array([70]))
    assert_array_equal(model._infection_density["b"], np.array([20]))
    assert_array_equal(model._infection_density["c"], np.array([10]))
    assert_array_equal(model._infection_frequency["a"], np.array([70 / 1000]))
    assert_array_equal(model._infection_frequency["b"], np.array([20 / 1000]))
    assert_array_equal(model._infection_frequency["c"], np.array([10 / 1000]))

    # Get multipliers
    susceptible = model.compartments[0]
    infectious_a = model.compartments[1]
    infectious_b = model.compartments[2]
    infectious_c = model.compartments[3]
    assert model._get_infection_density_multiplier(susceptible,
                                                   infectious_a) == 70
    assert model._get_infection_density_multiplier(susceptible,
                                                   infectious_b) == 20
    assert model._get_infection_density_multiplier(susceptible,
                                                   infectious_c) == 10
    assert model._get_infection_frequency_multiplier(susceptible,
                                                     infectious_a) == 70 / 1000
    assert model._get_infection_frequency_multiplier(susceptible,
                                                     infectious_b) == 20 / 1000
    assert model._get_infection_frequency_multiplier(susceptible,
                                                     infectious_c) == 10 / 1000

    # Get infection flow rates
    flow_rates = model._get_compartment_rates(model.initial_population, 0)
    sus_pop = 900
    flow_to_a = sus_pop * contact_rate * (70 * 0.5 / 1000)
    flow_to_b = sus_pop * contact_rate * (20 * 3 / 1000)
    flow_to_c = sus_pop * contact_rate * (10 * 2 / 1000)
    expected_flow_rates = np.array([
        -flow_to_a - flow_to_b - flow_to_c, flow_to_a, flow_to_b, flow_to_c,
        0.0
    ])
    assert_allclose(expected_flow_rates, flow_rates, verbose=True)