def test_stratify_compartments_partial():
"""Ensure compartment names and values get partially stratified properly"""
model = StratifiedModel(
times=[0, 1, 2, 3, 4, 5],
compartment_names=["sus", "inf"],
initial_conditions={"inf": 200},
parameters={},
requested_flows=[],
starting_population=1000,
infectious_compartments=["inf"],
birth_approach=BirthApproach.NO_BIRTH,
entry_compartment="sus",
)
assert model.compartment_names == ["sus", "inf"]
vals = np.array([800, 200])
assert_array_equal(model.compartment_values, vals)
model.stratify(
"location",
strata_request=["home", "work", "other"],
compartments_to_stratify=["sus"],
requested_proportions={"home": 0.6},
)
assert model.compartment_names == [
"susXlocation_home",
"susXlocation_work",
"susXlocation_other",
"inf",
]
vals = np.array([480, 160, 160, 200])
assert_array_equal(model.compartment_values, vals)
assert model.all_stratifications == {"location": ["home", "work", "other"]}
assert model.full_stratification_list == []
def test_no_mixing_matrix():
"""
Test that we are using the default 'null-op' mixing matrix when
we have a no user-specified mixing matrix
"""
model = StratifiedModel(**MODEL_KWARGS)
default_matrix = np.array([[2, 3], [5, 7]])
model.stratify(
stratification_name="agegroup",
strata_request=["child", "adult"],
compartments_to_stratify=["S", "I", "R"],
mixing_matrix=default_matrix,
)
# We should get the default mixing matrix
actual_mixing = model.get_mixing_matrix(0)
assert_array_equal(actual_mixing, default_matrix)
# Static matrices shouldn't change over time
actual_mixing = model.get_mixing_matrix(123)
assert_array_equal(actual_mixing, default_matrix)
model.stratify(
stratification_name="location",
strata_request=["work", "home"], # These kids have jobs.
compartments_to_stratify=["S", "I", "R"],
mixing_matrix=None,
)
# We should get the default mixing matrix
actual_mixing = model.get_mixing_matrix(0)
assert_array_equal(actual_mixing, default_matrix)
# Static matrices shouldn't change over time
actual_mixing = model.get_mixing_matrix(123)
assert_array_equal(actual_mixing, default_matrix)
def test_multiple_dynamic_mixing_matrices():
"""
Test that we are using the correct mixing matrix when
we have multiple dynamic mixing matrices
"""
model = StratifiedModel(**MODEL_KWARGS)
agegroup_matrix = lambda t: t * np.array([[2, 3], [5, 7]])
model.stratify(
stratification_name="agegroup",
strata_request=["child", "adult"],
compartments_to_stratify=["S", "I", "R"],
mixing_matrix=agegroup_matrix,
)
location_matrix = lambda t: t * np.array([[11, 13], [17, 19]])
model.stratify(
stratification_name="location",
strata_request=["work", "home"], # These kids have jobs.
compartments_to_stratify=["S", "I", "R"],
mixing_matrix=location_matrix,
)
expected_mixing = np.array([
[2 * 11, 2 * 13, 3 * 11, 3 * 13],
[2 * 17, 2 * 19, 3 * 17, 3 * 19],
[5 * 11, 5 * 13, 7 * 11, 7 * 13],
[5 * 17, 5 * 19, 7 * 17, 7 * 19],
])
# We should get the Kronecker product of the two matrices
actual_mixing = model.get_mixing_matrix(1)
assert_array_equal(actual_mixing, expected_mixing)
# Double check that we calculated the Kronecker product correctly
kron_mixing = np.kron(agegroup_matrix(1), location_matrix(1))
assert_array_equal(expected_mixing, kron_mixing)
# Dynamic matrices should change over time
actual_mixing = model.get_mixing_matrix(5)
assert_array_equal(actual_mixing, 25 * expected_mixing)
def _get_model():
pop = 1000
params = {
"contact_rate": 0.1,
"recovery_rate": 0.3,
"infect_death": 0.5,
"crude_birth_rate": 0.7,
"universal_death_rate": 0.11,
}
flows = [
{
"type": Flow.INFECTION_FREQUENCY,
"origin": "S",
"to": "I",
"parameter": "contact_rate"
},
{
"type": Flow.STANDARD,
"to": "R",
"origin": "I",
"parameter": "recovery_rate"
},
{
"type": Flow.DEATH,
"origin": "I",
"parameter": "infect_death"
},
]
model = StratifiedModel(
times=np.array([5.0, 6.0, 7.0, 8.0, 9.0, 10.0]),
compartment_names=["S", "I", "R"],
initial_conditions={"S": pop},
parameters=params,
requested_flows=flows,
starting_population=pop,
infectious_compartments=["I"],
birth_approach=BirthApproach.ADD_CRUDE,
entry_compartment="S",
)
# Add basic age stratification
model.stratify("agegroup",
strata_request=[0, 5, 15, 60],
compartments_to_stratify=["S", "I"])
# Pretend we ran the model
model.prepare_to_run()
model.outputs = np.array([
[250, 250, 250, 250, 0, 0, 0, 0, 0],
[150, 150, 150, 150, 100, 100, 100, 100, 0],
[100, 100, 100, 100, 100, 100, 100, 100, 200],
[50, 50, 50, 50, 150, 150, 150, 150, 200],
[50, 50, 50, 50, 100, 100, 100, 100, 400],
[0, 0, 0, 0, 150, 150, 150, 150, 400],
])
return model
def test_add_extra_flows__before_strat_then_stratify():
"""
Ensure that adding an extra flow without any stratifications creates a new flow.
Then check that stratifying the model creates the new stratified flows correctly.
"""
# Create the model
model = StratifiedModel(
times=np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
compartment_names=["S", "I", "R"],
initial_conditions={"S": 100},
parameters={},
requested_flows=[],
starting_population=1000,
infectious_compartments=["I"],
entry_compartment="S",
birth_approach=BirthApproach.NO_BIRTH,
)
# Add a new flow.
model.add_extra_flow(
flow={
"type": Flow.INFECTION_FREQUENCY,
"origin": "S",
"to": "I",
"parameter": "contact_rate",
},
source_strata={},
dest_strata={},
expected_flow_count=1,
)
# Ensure the flow was added correctly.
assert len(model.flows) == 1
flow = model.flows[0]
assert type(flow) is InfectionFrequencyFlow
assert flow.source == "S"
assert flow.dest == "I"
assert flow.adjustments == []
assert flow.param_name == "contact_rate"
# Stratify the model.
model.stratify("agegroup",
strata_request=["young", "old"],
compartments_to_stratify=["S", "I", "R"])
# Ensure the new flow was stratified correctly into two flows..
assert len(model.flows) == 2
flow_1 = model.flows[0]
assert type(flow_1) is InfectionFrequencyFlow
assert flow_1.source == "SXagegroup_young"
assert flow_1.dest == "IXagegroup_young"
assert flow_1.adjustments == []
assert flow_1.param_name == "contact_rate"
flow_2 = model.flows[1]
assert type(flow_2) is InfectionFrequencyFlow
assert flow_2.source == "SXagegroup_old"
assert flow_2.dest == "IXagegroup_old"
assert flow_2.adjustments == []
assert flow_2.param_name == "contact_rate"
def test_single_dynamic_mixing_matrix():
"""
Test that we are using the correct mixing matrix when
we have a single dynamic mixing matrix
"""
model = StratifiedModel(**MODEL_KWARGS)
agegroup_matrix = lambda t: t * np.array([[2, 3], [5, 7]])
model.stratify(
stratification_name="agegroup",
strata_request=["child", "adult"],
compartments_to_stratify=["S", "I", "R"],
mixing_matrix=agegroup_matrix,
)
# We should get the age mixing matrix
actual_mixing = model.get_mixing_matrix(1)
assert_array_equal(actual_mixing, agegroup_matrix(1))
# Dynamic matrices should change over time
actual_mixing = model.get_mixing_matrix(123)
assert_array_equal(actual_mixing, agegroup_matrix(123))
def test_strat_model__with_age_and_starting_proportion__expect_ageing():
"""
Ensure that a module with age stratification and starting proporptions
produces ageing flows, and the correct output.
"""
pop = 1000
model = StratifiedModel(
times=_get_integration_times(2000, 2005, 1),
compartment_names=["S", "I"],
initial_conditions={"S": pop},
parameters={},
requested_flows=[],
starting_population=pop,
infectious_compartments=["I"],
birth_approach=BirthApproach.NO_BIRTH,
entry_compartment="S",
)
# Add basic age stratification
model.stratify(
"age",
strata_request=[0, 5, 15, 60],
compartments_to_stratify=["S", "I"],
comp_split_props={
"0": 0.8,
"5": 0.1,
"15": 0.1
},
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to generally get older, but no one should die or get sick.
# Expect initial distribution of ages to be set according to "requested_proportions".
expected_arr = np.array([
[800.0, 100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[655.0, 228.3, 114.4, 2.4, 0.0, 0.0, 0.0, 0.0],
[536.3, 319.3, 139.3, 5.2, 0.0, 0.0, 0.0, 0.0],
[439.1, 381.3, 171.1, 8.6, 0.0, 0.0, 0.0, 0.0],
[359.5, 420.5, 207.2, 12.8, 0.0, 0.0, 0.0, 0.0],
[294.4, 442.4, 245.4, 17.8, 0.0, 0.0, 0.0, 0.0],
])
assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_strat_model__with_locations__expect_no_change():
"""
Ensure that a module with location stratification populates locations correctly.
"""
pop = 1000
model = StratifiedModel(
times=_get_integration_times(2000, 2005, 1),
compartment_names=["S", "I"],
initial_conditions={"S": pop},
parameters={},
requested_flows=[],
starting_population=pop,
infectious_compartments=["I"],
birth_approach=BirthApproach.NO_BIRTH,
entry_compartment="S",
)
# Add basic location stratification
model.stratify(
"location",
strata_request=["rural", "urban", "prison"],
compartments_to_stratify=["S", "I"],
comp_split_props={
"rural": 0.44,
"urban": 0.55,
"prison": 0.01
},
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to start in their locations, then nothing should change,
expected_arr = np.array([
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
])
assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_strat_model__with_age__expect_ageing():
"""
Ensure that a module with age stratification produces ageing flows,
and the correct output.
"""
pop = 1000
model = StratifiedModel(
times=_get_integration_times(2000, 2005, 1),
compartment_names=["S", "I"],
initial_conditions={"S": pop},
parameters={},
requested_flows=[],
starting_population=pop,
infectious_compartments=["I"],
birth_approach=BirthApproach.NO_BIRTH,
entry_compartment="S",
)
# Add basic age stratification
model.stratify("age",
strata_request=[0, 5, 15, 60],
compartments_to_stratify=["S", "I"])
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to generally get older, but no one should die or get sick
expected_arr = np.array([
[250.0, 250.0, 250.0, 250.0, 0.0, 0.0, 0.0, 0.0],
[204.8, 269.1, 270.3, 255.8, 0.0, 0.0, 0.0, 0.0],
[167.7, 278.8, 291.5, 262.0, 0.0, 0.0, 0.0, 0.0],
[137.3, 281.2, 312.8, 268.7, 0.0, 0.0, 0.0, 0.0],
[112.5, 277.9, 333.7, 275.9, 0.0, 0.0, 0.0, 0.0],
[92.1, 270.8, 353.5, 283.6, 0.0, 0.0, 0.0, 0.0],
])
assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_add_extra_flows__after_single_strat_with_cross_cut():
"""
Ensure that adding an extra flow with an existing stratification creates a new flow.
This flow cuts across strata.
"""
# Create the model
model = StratifiedModel(
times=np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
compartment_names=["S", "I", "R"],
initial_conditions={"S": 100},
parameters={},
requested_flows=[],
starting_population=1000,
infectious_compartments=["I"],
entry_compartment="S",
birth_approach=BirthApproach.NO_BIRTH,
)
# Ensure there are no flows yet.
assert len(model.flows) == 0
# Stratify the model.
model.stratify("agegroup",
strata_request=["young", "old"],
compartments_to_stratify=["S", "I", "R"])
# Add a new flow.
model.add_extra_flow(
flow={
"type": Flow.INFECTION_FREQUENCY,
"origin": "S",
"to": "I",
"parameter": "contact_rate",
},
source_strata={"agegroup": "young"},
dest_strata={"agegroup": "old"},
expected_flow_count=1,
)
# Ensure the new flow was added correctly.
assert len(model.flows) == 1
flow = model.flows[0]
assert type(flow) is InfectionFrequencyFlow
assert flow.source == "SXagegroup_young"
assert flow.dest == "IXagegroup_old"
assert flow.adjustments == []
assert flow.param_name == "contact_rate"
# Stratify the model again.
model.stratify("location",
strata_request=["urban", "rural"],
compartments_to_stratify=["S", "I", "R"])
# Ensure the new flow was stratified correctly into two flows.
assert len(model.flows) == 2
flow_1 = model.flows[0]
assert type(flow_1) is InfectionFrequencyFlow
assert flow_1.source == "SXagegroup_youngXlocation_urban"
assert flow_1.dest == "IXagegroup_oldXlocation_urban"
assert flow_1.adjustments == []
assert flow_1.param_name == "contact_rate"
flow_2 = model.flows[1]
assert type(flow_2) is InfectionFrequencyFlow
assert flow_2.source == "SXagegroup_youngXlocation_rural"
assert flow_2.dest == "IXagegroup_oldXlocation_rural"
assert flow_2.adjustments == []
assert flow_2.param_name == "contact_rate"
def test_stratify_flows_full__with_adjustment_requests():
"""Ensure flows get stratified properly in full strat"""
adjustment_requests = {
"contact_rate": {
"home": "home_contact_rate",
"work": 0.5,
},
"recovery_rate": {
"home": 1,
"work": 2,
},
"infect_death": {
"home": 1,
"work": 2,
},
}
requested_flows = [
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": "sus",
"to": "inf",
},
{
"type": Flow.STANDARD,
"parameter": "recovery_rate",
"origin": "inf",
"to": "sus",
},
{
"type": Flow.DEATH,
"parameter": "infect_death",
"origin": "inf",
},
]
parameters = {
"contact_rate": 1000,
"recovery_rate": "recovery_rate",
"infect_death": 10,
}
home_contact_rate_func = lambda t: t
recovery_rate_func = lambda t: 2 * t
model = StratifiedModel(
times=[0, 1, 2, 3, 4, 5],
compartment_names=["sus", "inf"],
initial_conditions={"inf": 200},
parameters=parameters,
requested_flows=requested_flows,
starting_population=1000,
infectious_compartments=["inf"],
birth_approach=BirthApproach.NO_BIRTH,
entry_compartment="sus",
)
model.time_variants["recovery_rate"] = recovery_rate_func
model.time_variants["home_contact_rate"] = home_contact_rate_func
assert model.compartment_names == ["sus", "inf"]
assert model.flows == requested_flows
assert model.parameters == parameters
assert model.time_variants == {
"recovery_rate": recovery_rate_func,
"home_contact_rate": home_contact_rate_func,
}
assert model.overwrite_parameters == []
assert model.adaptation_functions == {}
vals = np.array([800, 200])
assert_array_equal(model.compartment_values, vals)
model.stratify(
"location",
strata_request=["home", "work"],
compartments_to_stratify=[],
requested_proportions={"home": 0.6},
adjustment_requests=adjustment_requests,
)
assert model.flows == [
{
"origin": "susXlocation_home",
"parameter": "contact_rateXlocation_home",
"to": "infXlocation_home",
"type": Flow.INFECTION_FREQUENCY,
},
{
"origin": "susXlocation_work",
"parameter": "contact_rateXlocation_work",
"to": "infXlocation_work",
"type": Flow.INFECTION_FREQUENCY,
},
{
"origin": "infXlocation_home",
"parameter": "recovery_rateXlocation_home",
"to": "susXlocation_home",
"type": Flow.STANDARD,
},
{
"origin": "infXlocation_work",
"parameter": "recovery_rateXlocation_work",
"to": "susXlocation_work",
"type": Flow.STANDARD,
},
{
"origin": "infXlocation_home",
"parameter": "infect_deathXlocation_home",
"type": Flow.DEATH,
},
{
"origin": "infXlocation_work",
"parameter": "infect_deathXlocation_work",
"type": Flow.DEATH,
},
]
assert model.time_variants == {
"recovery_rate": recovery_rate_func,
"home_contact_rate": home_contact_rate_func,
}
assert model.parameters["contact_rate"] == 1000
assert model.parameters["infect_death"] == 10
assert model.parameters["recovery_rate"] == "recovery_rate"
assert model.parameters[
"contact_rateXlocation_home"] == "home_contact_rate"
assert model.parameters["contact_rateXlocation_work"] == 0.5
assert model.parameters["recovery_rateXlocation_home"] == 1
assert model.parameters["recovery_rateXlocation_work"] == 2
assert model.parameters["infect_deathXlocation_home"] == 1
assert model.parameters["infect_deathXlocation_work"] == 2
def build_model(params: dict) -> StratifiedModel:
"""
Build the master function to run a tuberculosis model
"""
validate_params(params) # perform validation of parameter format
check_param_values(params) # perform validation of some parameter values
# Define model times.
time = params["time"]
integration_times = get_model_times_from_inputs(round(time["start"]),
time["end"], time["step"],
time["critical_ranges"])
# Define model compartments.
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.EARLY_LATENT,
Compartment.LATE_LATENT,
Compartment.INFECTIOUS,
Compartment.ON_TREATMENT,
Compartment.RECOVERED,
]
infectious_comps = [
Compartment.INFECTIOUS,
Compartment.ON_TREATMENT,
]
# Define initial conditions - 1 infectious person.
init_conditions = {
Compartment.INFECTIOUS: 1,
}
# prepare infectiousness adjustment for individuals on treatment
treatment_infectiousness_adjustment = [{
"comp_name":
Compartment.ON_TREATMENT,
"comp_strata": {},
"value":
params["on_treatment_infect_multiplier"],
}]
# Define inter-compartmental flows.
flows = deepcopy(preprocess.flows.DEFAULT_FLOWS)
# is ACF implemented?
implement_acf = len(params["time_variant_acf"]) > 0
if implement_acf:
flows.append(preprocess.flows.ACF_FLOW)
# is ltbi screening implemented?
implement_ltbi_screening = len(params["time_variant_ltbi_screening"]) > 0
if implement_ltbi_screening:
flows += preprocess.flows.get_preventive_treatment_flows(
params["pt_destination_compartment"])
# Set some parameter values or parameters that require pre-processing
(
params,
treatment_recovery_func,
treatment_death_func,
relapse_func,
detection_rate_func,
acf_detection_rate_func,
preventive_treatment_func,
contact_rate_functions,
) = preprocess.flows.process_unstratified_parameter_values(
params, implement_acf, implement_ltbi_screening)
# Create the model.
tb_model = StratifiedModel(
times=integration_times,
compartment_names=compartments,
initial_conditions=init_conditions,
parameters=params,
requested_flows=flows,
infectious_compartments=infectious_comps,
birth_approach=BirthApproach.ADD_CRUDE,
entry_compartment=Compartment.SUSCEPTIBLE,
starting_population=int(params["start_population_size"]),
)
# register acf_detection_func
if acf_detection_rate_func is not None:
tb_model.time_variants["acf_detection_rate"] = acf_detection_rate_func
# register preventive_treatment_func
if preventive_treatment_func is not None:
tb_model.time_variants[
"preventive_treatment_rate"] = preventive_treatment_func
# register time-variant contact-rate functions:
for param_name, func in contact_rate_functions.items():
tb_model.time_variants[param_name] = func
# Apply infectiousness adjustment for individuals on treatment
tb_model.individual_infectiousness_adjustments = treatment_infectiousness_adjustment
# apply age stratification
if "age" in params["stratify_by"]:
stratify_by_age(tb_model, params, compartments)
else:
# set time-variant functions for treatment death and relapse rates
tb_model.time_variants[
"treatment_recovery_rate"] = treatment_recovery_func
tb_model.time_variants["treatment_death_rate"] = treatment_death_func
tb_model.time_variants["relapse_rate"] = relapse_func
# Load time-variant birth rates
birth_rates, years = inputs.get_crude_birth_rate(params["iso3"])
birth_rates = [b / 1000.0 for b in birth_rates
] # birth rates are provided / 1000 population
tb_model.time_variants["crude_birth_rate"] = scale_up_function(
years, birth_rates, smoothness=0.2, method=5)
tb_model.parameters["crude_birth_rate"] = "crude_birth_rate"
# apply user-defined stratifications
user_defined_stratifications = [
s for s in list(params["user_defined_stratifications"].keys())
if s in params["stratify_by"]
]
for stratification in user_defined_stratifications:
assert "_" not in stratification, "Stratification name should not include '_'"
stratification_details = params["user_defined_stratifications"][
stratification]
apply_user_defined_stratification(
tb_model,
compartments,
stratification,
stratification_details,
implement_acf,
implement_ltbi_screening,
)
if "organ" in params["stratify_by"]:
stratify_by_organ(tb_model, params)
else:
tb_model.time_variants["detection_rate"] = detection_rate_func
# Register derived output functions, which are calculations based on the model's compartment values or flows.
# These are calculated after the model is run.
outputs.get_all_derived_output_functions(params["calculated_outputs"],
params["outputs_stratification"],
tb_model)
return tb_model
def test_strains__with_two_symmetric_strains():
"""
Adding two strains with the same properties should yield the same infection dynamics and outputs as having no strains at all.
We expect the force of infection for each strain to be 1/2 of the unstratified model,
but the stratification process will not apply the usual conservation fraction to the fan out flows.
"""
params = {
"contact_rate": 0.2,
"recovery_rate": 0.1,
}
flows = (
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": "S",
"to": "I",
},
{
"type": Flow.STANDARD,
"parameter": "recovery_rate",
"origin": "I",
"to": "R",
},
)
# Create an unstratified model
kwargs = {**MODEL_KWARGS, "parameters": params, "requested_flows": flows}
model = StratifiedModel(**kwargs)
# Do pre-run force of infection calcs.
model.prepare_to_run()
model.prepare_time_step(0, model.compartment_values)
# Check infectiousness multipliers
susceptible = model.compartment_names[0]
infectious = model.compartment_names[1]
assert model.get_infection_density_multiplier(susceptible, infectious) == 100.0
assert model.get_infection_frequency_multiplier(susceptible, infectious) == 0.1
model.run_model()
# Create a stratified model where the two strains are symmetric
strain_model = StratifiedModel(**kwargs)
strain_model.stratify(
stratification_name="strain",
strata_request=["a", "b"],
compartments_to_stratify=["I"],
# Use defaults - equally split compartments, flows, etc.
comp_split_props={},
flow_adjustments={},
infectiousness_adjustments={},
mixing_matrix=None,
)
strain_model.run_model()
merged_outputs = np.zeros_like(model.outputs)
merged_outputs[:, 0] = strain_model.outputs[:, 0]
merged_outputs[:, 1] = strain_model.outputs[:, 1] + strain_model.outputs[:, 2]
merged_outputs[:, 2] = strain_model.outputs[:, 3]
assert_allclose(merged_outputs, model.outputs, atol=0.01, rtol=0.01, verbose=True)
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.
"""
contact_rate = 0.2
params = {
"contact_rate": contact_rate,
}
flows = (
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": "S",
"to": "I",
},
)
kwargs = {**MODEL_KWARGS, "parameters": params, "requested_flows": flows}
model = StratifiedModel(**kwargs)
model.stratify(
stratification_name="agegroup",
strata_request=["child", "adult"],
compartments_to_stratify=["S", "I", "R"],
comp_split_props={
"child": 0.6, # 600 people
"adult": 0.4, # 400 people
},
# Higher mixing among adults or children,
# than between adults or children.
mixing_matrix=np.array([[1.5, 0.5], [0.5, 1.5]]),
)
model.stratify(
stratification_name="strain",
strata_request=["a", "b", "c"],
compartments_to_stratify=["I"],
comp_split_props={
"a": 0.7, # 70 people
"b": 0.2, # 20 people
"c": 0.1, # 10 people
},
infectiousness_adjustments={
"a": 0.5, # 0.5x as infectious
"b": 3, # 3x as infectious
"c": 2, # 2x as infectious
},
mixing_matrix=None,
)
# 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.compartment_values)
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.compartment_names[0]
sus_adult = model.compartment_names[1]
inf_child_a = model.compartment_names[2]
inf_child_b = model.compartment_names[3]
inf_child_c = model.compartment_names[4]
inf_adult_a = model.compartment_names[5]
inf_adult_b = model.compartment_names[6]
inf_adult_c = model.compartment_names[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_flow_rates(model.compartment_values, 0)
assert_allclose(expected_flow_rates, flow_rates, verbose=True)
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.
"""
contact_rate = 0.2
params = {
"contact_rate": contact_rate,
}
flows = (
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": "S",
"to": "I",
},
)
kwargs = {**MODEL_KWARGS, "parameters": params, "requested_flows": flows}
model = StratifiedModel(**kwargs)
model.stratify(
stratification_name="strain",
strata_request=["a", "b", "c"],
compartments_to_stratify=["I"],
comp_split_props={
"a": 0.7, # 70 people
"b": 0.2, # 20 people
"c": 0.1, # 10 people
},
flow_adjustments={
"contact_rate": {
"a": 0.5, # 0.5x as infectious
"b": 3, # 3x as infectious
"c": 2, # 2x as infectious
}
},
mixing_matrix=None,
)
# 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.compartment_values)
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.compartment_names[0]
infectious_a = model.compartment_names[1]
infectious_b = model.compartment_names[2]
infectious_c = model.compartment_names[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_flow_rates(model.compartment_values, 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)
