def test_stratify_compartments(strata, proportions, to_stratify,
expected_names, expected_values):
"""
Ensure that `stratify_compartments` splits up compartment names and values correctly.
"""
model = StratifiedModel(**_get_model_kwargs())
model.stratify_compartments("test", strata, proportions, to_stratify)
assert model.compartment_names == expected_names
assert model.compartment_values == expected_values
def test_find_transition_indices_to_implement(back_one, include_change,
all_stratifications, flows,
expected_idxs):
"""
Ensure `find_transition_indices_to_implement` returns the correct list of indices.
"""
model = StratifiedModel(**_get_model_kwargs())
cols = [
"type", "parameter", "origin", "to", "implement", "strain",
"force_index"
]
model.transition_flows = pd.DataFrame(flows, columns=cols).astype(object)
model.all_stratifications = all_stratifications
actual_idxs = model.find_transition_indices_to_implement(
back_one, include_change)
assert expected_idxs == actual_idxs
def _build_model(*args, **kwargs):
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
)
return model
def test_set_ageing_rates(age_strata, expected_flows, expected_ageing):
"""
Ensure that `set_ageing_rates` adds ageing flows to the transition flows dataframe
"""
model = StratifiedModel(**_get_model_kwargs())
cols = [
"type", "parameter", "origin", "to", "implement", "strain",
"force_index"
]
# Ensure there are no initial flows
initial_df = pd.DataFrame([], columns=cols).astype(object)
assert_frame_equal(initial_df, model.transition_flows)
# Set ageing rates
model.set_ageing_rates(age_strata)
# Check ageing flows are set
expected_df = pd.DataFrame(expected_flows, columns=cols).astype(object)
assert_frame_equal(expected_df, model.transition_flows)
# Check ageing params are set
for k, v in expected_ageing.items():
assert model.parameters[k] == v
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_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
)
# 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_output = [
[250.0, 250.0, 250.0, 250.0, 0.0, 0.0, 0.0, 0.0],
[205.0, 269.0, 270.0, 256.0, 0.0, 0.0, 0.0, 0.0],
[168.0, 279.0, 291.0, 262.0, 0.0, 0.0, 0.0, 0.0],
[137.0, 281.0, 313.0, 269.0, 0.0, 0.0, 0.0, 0.0],
[112.0, 278.0, 334.0, 276.0, 0.0, 0.0, 0.0, 0.0],
[92.0, 271.0, 354.0, 284.0, 0.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
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_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic location stratification
model.stratify(
Stratification.LOCATION,
strata_request=["rural", "urban", "prison"],
compartment_types_to_stratify=[],
requested_proportions={"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_output = [
[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],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_stratify_transition_flows(
flows,
custom_func,
params,
adjustment,
comps,
expected_flows,
expected_custom_func,
expected_params,
):
"""
Ensure that `stratify_compartments` splits up transition flows correctly.
"""
model = StratifiedModel(**_get_model_kwargs())
cols = [
"type", "parameter", "origin", "to", "implement", "strain",
"force_index"
]
model.transition_flows = pd.DataFrame(flows, columns=cols).astype(object)
model.parameters = params
strat_name = "test"
strata_names = ["foo", "bar"]
model.customised_flow_functions = custom_func
model.all_stratifications = {strat_name: strata_names}
model.stratify_compartments(strat_name, strata_names, {
"foo": 0.5,
"bar": 0.5
}, comps)
model.stratify_transition_flows(strat_name, strata_names, adjustment,
comps)
# Check flows df stratified
expected_flows_df = pd.DataFrame(expected_flows,
columns=cols).astype(object)
assert_frame_equal(expected_flows_df, model.transition_flows)
# Check custom flow func is updated
assert model.customised_flow_functions == expected_custom_func
# Check params are stratified
for k, v in expected_params.items():
assert model.parameters[k] == v
def build_model(country: str, params: dict, update_params={}):
"""
Build the master function to run a simple SIR model
:param update_params: dict
Any parameters that need to be updated for the current run
:return: StratifiedModel
The final model with all parameters and stratifications
"""
params.update(update_params)
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.INFECTIOUS,
Compartment.RECOVERED,
]
flows = [
{
'type': 'infection_frequency',
'parameter': 'contact_rate',
'origin': Compartment.SUSCEPTIBLE,
'to': Compartment.INFECTIOUS
},
{
'type': 'standard_flows',
'parameter': 'recovery_rate',
'origin': Compartment.INFECTIOUS,
'to': Compartment.RECOVERED
},
]
integration_times = get_model_times_from_inputs(
round(params["start_time"]), params["end_time"], params["time_step"])
init_conditions = {Compartment.INFECTIOUS: 1}
sir_model = StratifiedModel(
integration_times,
compartments,
init_conditions,
params,
flows,
infectious_compartment=(Compartment.INFECTIOUS, ),
birth_approach=BirthApproach.NO_BIRTH,
starting_population=1000000,
)
return sir_model
def test_strat_model__with_locations_and_mixing__expect_varied_transmission():
"""
Ensure that location-based mixing works.
Expect urbanites to be highly infectious t other locations, but not the reverse.
"""
pop = 1000
model = StratifiedModel(
times=_get_integration_times(2000, 2005, 1),
compartment_types=[
Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS
],
initial_conditions={Compartment.EARLY_INFECTIOUS: 100},
parameters={"contact_rate": 3},
requested_flows=[{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.EARLY_INFECTIOUS,
}],
starting_population=pop,
)
# Add basic location stratification
model.stratify(
Stratification.LOCATION,
strata_request=["rural", "urban", "prison"],
compartment_types_to_stratify=[],
requested_proportions={},
mixing_matrix=np.array([
# Rural people catch disease from urbanites
[0.0, 1.0, 0.0],
# Urbanites cannot catch disease from anyone
[0.0, 0.0, 0.0],
# Prisoners catch disease from urbanites
[0.0, 1.0, 0.0],
]),
)
# 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_output = [
[300.0, 300.0, 300.0, 33.0, 33.0, 33.0],
[271.0, 300.0, 271.0, 62.0, 33.0, 62.0],
[246.0, 300.0, 246.0, 88.0, 33.0, 88.0],
[222.0, 300.0, 222.0, 111.0, 33.0, 111.0],
[201.0, 300.0, 201.0, 132.0, 33.0, 132.0],
[182.0, 300.0, 182.0, 151.0, 33.0, 151.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def build_model(params: dict) -> StratifiedModel:
"""
Build the master function to run a simple SIR model
"""
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.INFECTIOUS,
Compartment.RECOVERED,
]
flows = [
{
"type": "infection_frequency",
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.INFECTIOUS,
},
{
"type": "standard_flows",
"parameter": "recovery_rate",
"origin": Compartment.INFECTIOUS,
"to": Compartment.RECOVERED,
},
]
integration_times = get_model_times_from_inputs(
round(params["start_time"]), params["end_time"], params["time_step"])
init_conditions = {Compartment.INFECTIOUS: 1}
sir_model = StratifiedModel(
integration_times,
compartments,
init_conditions,
params,
flows,
infectious_compartment=(Compartment.INFECTIOUS, ),
birth_approach=BirthApproach.NO_BIRTH,
starting_population=1000000,
)
return sir_model
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_types=[
Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS
],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={
"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_output = [
[800.0, 100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[655.0, 228.0, 114.0, 2.0, 0.0, 0.0, 0.0, 0.0],
[536.0, 319.0, 139.0, 5.0, 0.0, 0.0, 0.0, 0.0],
[439.0, 381.0, 171.0, 9.0, 0.0, 0.0, 0.0, 0.0],
[360.0, 421.0, 207.0, 13.0, 0.0, 0.0, 0.0, 0.0],
[294.0, 442.0, 245.0, 18.0, 0.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_strat_model__with_age_and_infectiousness__expect_age_based_infectiousness():
"""
Ensure that a module with age stratification produces ageing flows,
and the correct output. Ensure that age-speific mortality is used.
"""
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={"contact_rate": None},
requested_flows=[
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.EARLY_INFECTIOUS,
}
],
starting_population=pop,
)
# Add age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
adjustment_requests={"contact_rate": {"0": 0.0, "5": 3.0, "15": 0.0, "60": 0.0}},
)
# 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_output = []
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def build_model(params: dict, update_params={}) -> StratifiedModel:
"""
Build the master function to run a simple SIR model
:param update_params: dict
Any parameters that need to be updated for the current run
:return: StratifiedModel
The final model with all parameters and stratifications
"""
params.update(update_params)
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.EARLY_LATENT,
Compartment.LATE_LATENT,
Compartment.INFECTIOUS,
Compartment.RECOVERED
]
params['delta'] = params['beta'] * params['rr_reinfection_once_recovered']
params['theta'] = params['beta'] * params['rr_reinfection_once_infected'] # not used at the moment
flows = [
{
'type': 'infection_frequency',
'parameter': 'beta',
'origin': Compartment.SUSCEPTIBLE,
'to': Compartment.EARLY_LATENT
},
{
'type': 'standard_flows',
'parameter': 'kappa',
'origin': Compartment.EARLY_LATENT,
'to': Compartment.LATE_LATENT
},
{
'type': 'standard_flows',
'parameter': 'epsilon',
'origin': Compartment.EARLY_LATENT,
'to': Compartment.INFECTIOUS
},
{
'type': 'standard_flows',
'parameter': 'nu',
'origin': Compartment.LATE_LATENT,
'to': Compartment.INFECTIOUS
},
{
'type': 'standard_flows',
'parameter': 'tau',
'origin': Compartment.INFECTIOUS,
'to': Compartment.RECOVERED
},
{
'type': 'standard_flows',
'parameter': 'gamma',
'origin': Compartment.INFECTIOUS,
'to': Compartment.LATE_LATENT
},
{
'type': 'standard_flows',
'parameter': 'delta',
'origin': Compartment.RECOVERED,
'to': Compartment.EARLY_LATENT
},
{
'type': 'compartment_death',
'parameter': 'infect_death',
'origin': Compartment.INFECTIOUS
}
]
integration_times = get_model_times_from_inputs(
round(params["start_time"]), params["end_time"], params["time_step"]
)
init_conditions = {Compartment.INFECTIOUS: 1}
tb_sir_model = StratifiedModel(
integration_times,
compartments,
init_conditions,
params,
requested_flows=flows,
birth_approach=BirthApproach.REPLACE_DEATHS,
entry_compartment='susceptible',
starting_population=100000,
infectious_compartment=(Compartment.INFECTIOUS,),
verbose=True,
)
tb_sir_model.adaptation_functions['universal_death_rateX'] = lambda x: 1./70
# # add time_variant parameters
def time_variant_CDR():
return scale_up_function(params['cdr'].keys(), params['cdr'].values(), smoothness=0.2, method=5)
def time_variant_TSR():
return scale_up_function(params['tsr'].keys(), params['tsr'].values(), smoothness=0.2, method=5)
def time_variant_DST():
return scale_up_function(params['dst'].keys(), params['dst'].values(), smoothness=0.2, method=5)
my_tv_CDR = time_variant_CDR()
my_tv_TSR = time_variant_TSR()
my_tv_DST = time_variant_DST()
def my_tv_tau(time):
return my_tv_detection_rate(time) * my_tv_TSR(time) * my_tv_DST(time)
tb_sir_model.adaptation_functions["tau"] = my_tv_tau
tb_sir_model.parameters["tau"] = "tau"
def my_tv_detection_rate(time):
return my_tv_CDR(time)/(1-my_tv_CDR(time)) * (params['gamma'] + params['universal_death_rate'] + params['infect_death']) #calculating the time varaint detection rate from tv CDR
#adding strain stratification
stratify_by = ['strain'] # ['strain', 'treatment_type']
if 'strain' in stratify_by:
tb_sir_model.stratify(
"strain", ['ds', 'inh_R', 'rif_R', 'mdr'],
compartment_types_to_stratify=[Compartment.EARLY_LATENT, Compartment.LATE_LATENT, Compartment.INFECTIOUS],
requested_proportions={'ds': 1., 'inh_R': 0., 'rif_R':0., 'mdr': 0.}, #adjustment_requests={'tau': tau_adjustment},
verbose=False,
adjustment_requests={'beta': {'ds': 1., 'inh_R': 0.9, 'rif_R': 0.7, 'mdr': 0.5}})
# set up some amplification flows
# INH_R
tb_sir_model.add_transition_flow(
{"type": "standard_flows", "parameter": "dr_amplification_inh",
"origin": "infectiousXstrain_ds", "to": "infectiousXstrain_inh_R",
"implement": len(tb_sir_model.all_stratifications)})
def tv_amplification_inh_rate(time):
return my_tv_detection_rate(time) * (1 - my_tv_TSR(time)) * params['prop_of_failures_developing_inh_R']
tb_sir_model.adaptation_functions["dr_amplification_inh"] = tv_amplification_inh_rate
tb_sir_model.parameters["dr_amplification_inh"] = "dr_amplification_inh"
# set up amplification flow for RIF_R
tb_sir_model.add_transition_flow(
{"type": "standard_flows", "parameter": "dr_amplification_rif",
"origin": "infectiousXstrain_ds", "to": "infectiousXstrain_rif_R",
"implement": len(tb_sir_model.all_stratifications)})
def tv_amplification_rif_rate(time):
return my_tv_detection_rate(time) * (1 - my_tv_TSR(time)) * params['prop_of_failures_developing_rif_R']
tb_sir_model.adaptation_functions["dr_amplification_rif"] = tv_amplification_rif_rate
tb_sir_model.parameters["dr_amplification_rif"] = "dr_amplification_rif"
# set up amplification flow for MDR
tb_sir_model.add_transition_flow(
{"type": "standard_flows", "parameter": "dr_amplification_rif",
"origin": "infectiousXstrain_inh_R", "to": "infectiousXstrain_mdr",
"implement": len(tb_sir_model.all_stratifications)})
tb_sir_model.add_transition_flow(
{"type": "standard_flows", "parameter": "dr_amplification_inh",
"origin": "infectiousXstrain_rif_R", "to": "infectiousXstrain_mdr",
"implement": len(tb_sir_model.all_stratifications)})
# tb_sir_model.transition_flows.to_csv("transitions.csv")
# create_flowchart(tb_sir_model, name="sir_model_diagram")
tb_sir_model.run_model()
return tb_sir_model
def build_model(params: dict) -> StratifiedModel:
external_params = deepcopy(params)
model_parameters = {
"contact_rate":
external_params["contact_rate"],
"contact_rate_recovered":
external_params["contact_rate"] *
external_params["rr_transmission_recovered"],
"contact_rate_late_latent":
external_params["contact_rate"] *
external_params["rr_transmission_late_latent"],
"recovery":
external_params["self_recovery_rate"],
"infect_death":
external_params["tb_mortality_rate"],
**external_params,
}
stratify_by = external_params["stratify_by"]
derived_output_types = external_params["derived_outputs"]
n_iter = (int(
round((external_params["end_time"] - external_params["start_time"]) /
external_params["time_step"])) + 1)
integration_times = numpy.linspace(external_params["start_time"],
external_params["end_time"],
n_iter).tolist()
model_parameters.update(
change_parameter_unit(provide_aggregated_latency_parameters(),
365.251))
# sequentially add groups of flows
flows = add_standard_infection_flows([])
flows = add_standard_latency_flows(flows)
flows = add_standard_natural_history_flows(flows)
# compartments
compartments = [
"susceptible",
"early_latent",
"late_latent",
"infectious",
"recovered",
]
# define model #replace_deaths add_crude_birth_rate
init_pop = {"infectious": 1000, "late_latent": 1000000}
tb_model = StratifiedModel(
integration_times,
compartments,
init_pop,
model_parameters,
flows,
birth_approach="replace_deaths",
starting_population=external_params["start_population"],
output_connections={},
derived_output_functions={},
death_output_categories=((), ("age_0", )),
)
# add crude birth rate from un estimates
birth_rates, years = get_crude_birth_rate("MNG")
# Provisional patch to birth rates
for i in range(len(birth_rates)):
if years[i] > 1990.0:
birth_rates[i] = 0.04
tb_model.time_variants["crude_birth_rate"] = scale_up_function(
years, birth_rates, smoothness=0.2, method=5)
# add case detection process to basic model
tb_model.add_transition_flow({
"type": "standard_flows",
"parameter": "case_detection",
"origin": "infectious",
"to": "recovered",
})
# Add IPT as a customised flow
def ipt_flow_func(model, n_flow, _time, _compartment_values):
"""
Work out the number of detected individuals from the relevant active TB compartments (with regard to the origin
latent compartment of n_flow) multiplied with the proportion of the relevant infected contacts that is from this
latent compartment.
"""
dict_flows = model.transition_flows_dict
origin_comp_name = dict_flows["origin"][n_flow]
components_latent_comp = find_name_components(origin_comp_name)
# find compulsory tags to be found in relevant infectious compartments
tags = []
for component in components_latent_comp:
if "location_" in component or "strain_" in component:
tags.append(component)
# loop through all relevant infectious compartments
total_tb_detected = 0.0
for comp_ind in model.infectious_indices["all_strains"]:
active_components = find_name_components(
model.compartment_names[comp_ind])
if all(elem in active_components for elem in tags):
infectious_pop = _compartment_values[comp_ind]
detection_indices = [
index for index, val in dict_flows["parameter"].items()
if "case_detection" in val
]
flow_index = [
index for index in detection_indices
if dict_flows["origin"][index] ==
model.compartment_names[comp_ind]
][0]
param_name = dict_flows["parameter"][flow_index]
detection_tx_rate = model.get_parameter_value(
param_name, _time)
tsr = mongolia_tsr(_time) + external_params[
"reduction_negative_tx_outcome"] * (1.0 -
mongolia_tsr(_time))
if "strain_mdr" in model.compartment_names[comp_ind]:
tsr = external_params["mdr_tsr"] * external_params[
"prop_mdr_detected_as_mdr"]
if tsr > 0.0:
total_tb_detected += infectious_pop * detection_tx_rate / tsr
# list all latent compartments relevant to the relevant infectious population
relevant_latent_compartments_indices = [
i for i, comp_name in enumerate(model.compartment_names)
if find_stem(comp_name) == "early_latent" and all(elem in comp_name
for elem in tags)
]
total_relevant_latent_size = sum(
_compartment_values[i]
for i in relevant_latent_compartments_indices)
current_latent_size = _compartment_values[
model.compartment_names.index(origin_comp_name)]
prop_of_relevant_latent = (current_latent_size /
total_relevant_latent_size if
total_relevant_latent_size > 0.0 else 0.0)
return total_tb_detected * prop_of_relevant_latent
tb_model.add_transition_flow({
"type": "customised_flows",
"parameter": "ipt_rate",
"origin": "early_latent",
"to": "recovered",
"function": ipt_flow_func,
})
# add ACF flow
tb_model.add_transition_flow({
"type": "standard_flows",
"parameter": "acf_rate",
"origin": "infectious",
"to": "recovered",
})
# load time-variant case detection rate
cdr_scaleup_overall = build_mongolia_timevariant_cdr(
external_params["cdr_multiplier"])
# targeted TB prevalence proportions by organ
prop_smearpos = 0.25
prop_smearneg = 0.40
prop_extrapul = 0.35
# disease duration by organ
overall_duration = prop_smearpos * 1.6 + 5.3 * (1 - prop_smearpos)
disease_duration = {
"smearpos": 1.6,
"smearneg": 5.3,
"extrapul": 5.3,
"overall": overall_duration,
}
# work out the CDR for smear-positive TB
def cdr_smearpos(time):
# Had to replace external_params['diagnostic_sensitivity_smearneg'] with its hard-coded value .7 to avoid
# cdr_smearpos to be affected when increasing diagnostic_sensitivity_smearneg in interventions (e.g. Xpert)
# return (cdr_scaleup_overall(time) /
# (prop_smearpos + prop_smearneg * external_params['diagnostic_sensitivity_smearneg'] +
# prop_extrapul * external_params['diagnostic_sensitivity_extrapul']))
return cdr_scaleup_overall(time) / (
prop_smearpos + prop_smearneg * 0.7 +
prop_extrapul * external_params["diagnostic_sensitivity_extrapul"])
def cdr_smearneg(time):
return cdr_smearpos(
time) * external_params["diagnostic_sensitivity_smearneg"]
def cdr_extrapul(time):
return cdr_smearpos(
time) * external_params["diagnostic_sensitivity_extrapul"]
cdr_by_organ = {
"smearpos": cdr_smearpos,
"smearneg": cdr_smearneg,
"extrapul": cdr_extrapul,
"overall": cdr_scaleup_overall,
}
detect_rate_by_organ = {}
for organ in ["smearpos", "smearneg", "extrapul", "overall"]:
prop_to_rate = convert_competing_proportion_to_rate(
1.0 / disease_duration[organ])
detect_rate_by_organ[organ] = return_function_of_function(
cdr_by_organ[organ], prop_to_rate)
# load time-variant treatment success rate
mongolia_tsr = build_mongolia_timevariant_tsr()
# create a treatment succes rate function adjusted for treatment support intervention
tsr_function = lambda t: mongolia_tsr(t) + external_params[
"reduction_negative_tx_outcome"] * (1.0 - mongolia_tsr(t))
# tb control recovery rate (detection and treatment) function set for overall if not organ-specific, smearpos otherwise
if "organ" not in stratify_by:
tb_control_recovery_rate = lambda t: tsr_function(
t) * detect_rate_by_organ["overall"](t)
else:
tb_control_recovery_rate = lambda t: tsr_function(
t) * detect_rate_by_organ["smearpos"](t)
# initialise ipt_rate function assuming coverage of 1.0 before age stratification
ipt_rate_function = (lambda t: 1.0 * external_params[
"yield_contact_ct_tstpos_per_detected_tb"] * external_params[
"ipt_efficacy"])
# initialise acf_rate function
acf_rate_function = (
lambda t: external_params["acf_coverage"] * external_params[
"acf_sensitivity"] *
(mongolia_tsr(t) + external_params["reduction_negative_tx_outcome"] *
(1.0 - mongolia_tsr(t))))
# assign newly created functions to model parameters
tb_model.adaptation_functions["case_detection"] = tb_control_recovery_rate
tb_model.parameters["case_detection"] = "case_detection"
tb_model.adaptation_functions["ipt_rate"] = ipt_rate_function
tb_model.parameters["ipt_rate"] = "ipt_rate"
tb_model.adaptation_functions["acf_rate"] = acf_rate_function
tb_model.parameters["acf_rate"] = "acf_rate"
if "strain" in stratify_by:
mdr_adjustment = (external_params["prop_mdr_detected_as_mdr"] *
external_params["mdr_tsr"] / 0.9
) # /.9 for last DS TSR
tb_model.stratify(
"strain",
["ds", "mdr"],
["early_latent", "late_latent", "infectious"],
verbose=False,
requested_proportions={"mdr": 0.0},
adjustment_requests={
"contact_rate": {
"ds": 1.0,
"mdr": 1.0
},
"case_detection": {
"mdr": mdr_adjustment
},
"ipt_rate": {
"ds": 1.0, # external_params['ds_ipt_switch'],
"mdr": external_params["mdr_ipt_switch"],
},
},
infectiousness_adjustments={
"ds": 1.0,
"mdr": external_params["mdr_infectiousness_multiplier"],
},
)
tb_model.add_transition_flow({
"type":
"standard_flows",
"parameter":
"dr_amplification",
"origin":
"infectiousXstrain_ds",
"to":
"infectiousXstrain_mdr",
"implement":
len(tb_model.all_stratifications),
})
dr_amplification_rate = (
lambda t: detect_rate_by_organ["overall"](t) *
(1.0 - mongolia_tsr(t)) *
(1.0 - external_params["reduction_negative_tx_outcome"]
) * external_params["dr_amplification_prop_among_nonsuccess"])
tb_model.adaptation_functions[
"dr_amplification"] = dr_amplification_rate
tb_model.parameters["dr_amplification"] = "dr_amplification"
if "age" in stratify_by:
age_breakpoints = [0, 5, 15, 60]
age_infectiousness = get_parameter_dict_from_function(
logistic_scaling_function(10.0), age_breakpoints)
age_params = get_adapted_age_parameters(age_breakpoints)
age_params.update(split_age_parameter(age_breakpoints, "contact_rate"))
# adjustment of latency parameters
for param in ["early_progression", "late_progression"]:
for age_break in age_breakpoints:
if age_break > 5:
age_params[param][
str(age_break) +
"W"] *= external_params["adult_latency_adjustment"]
death_rates_by_age, death_rate_years = get_death_rates_by_agegroup(
age_breakpoints, "MNG")
pop_morts = {}
for age_group in age_breakpoints:
pop_morts[age_group] = scale_up_function(
death_rate_years,
death_rates_by_age[age_group],
smoothness=0.2,
method=5)
age_params["universal_death_rate"] = {}
for age_break in age_breakpoints:
tb_model.time_variants["universal_death_rateXage_" +
str(age_break)] = pop_morts[age_break]
tb_model.parameters[
"universal_death_rateXage_" +
str(age_break)] = "universal_death_rateXage_" + str(age_break)
age_params["universal_death_rate"][
str(age_break) +
"W"] = "universal_death_rateXage_" + str(age_break)
tb_model.parameters["universal_death_rateX"] = 0.0
# age-specific IPT
ipt_by_age = {"ipt_rate": {}}
for age_break in age_breakpoints:
ipt_by_age["ipt_rate"][str(age_break)] = external_params[
"ipt_age_" + str(age_break) + "_ct_coverage"]
age_params.update(ipt_by_age)
# add BCG effect without stratification assuming constant 100% coverage
bcg_wane = create_sloping_step_function(15.0, 0.3, 30.0, 1.0)
age_bcg_efficacy_dict = get_parameter_dict_from_function(
lambda value: bcg_wane(value), age_breakpoints)
age_params.update({"contact_rate": age_bcg_efficacy_dict})
tb_model.stratify(
"age",
deepcopy(age_breakpoints),
[],
{},
adjustment_requests=age_params,
infectiousness_adjustments=age_infectiousness,
verbose=False,
)
# patch for IPT to overwrite parameters when ds_ipt has been turned off while we still need some coverage at baseline
if external_params["ds_ipt_switch"] == 0.0 and external_params[
"mdr_ipt_switch"] == 1.0:
tb_model.parameters["ipt_rateXstrain_dsXage_0"] = 0.17
for age_break in [5, 15, 60]:
tb_model.parameters["ipt_rateXstrain_dsXage_" +
str(age_break)] = 0.0
if "organ" in stratify_by:
props_smear = {
"smearpos": external_params["prop_smearpos"],
"smearneg": 1.0 - (external_params["prop_smearpos"] + 0.20),
"extrapul": 0.20,
}
mortality_adjustments = {
"smearpos": 1.0,
"smearneg": 0.064,
"extrapul": 0.064
}
recovery_adjustments = {
"smearpos": 1.0,
"smearneg": 0.56,
"extrapul": 0.56
}
# workout the detection rate adjustment by organ status
adjustment_smearneg = (detect_rate_by_organ["smearneg"](2015.0) /
detect_rate_by_organ["smearpos"](2015.0) if
detect_rate_by_organ["smearpos"](2015.0) > 0.0
else 1.0)
adjustment_extrapul = (detect_rate_by_organ["extrapul"](2015.0) /
detect_rate_by_organ["smearpos"](2015.0) if
detect_rate_by_organ["smearpos"](2015.0) > 0.0
else 1.0)
tb_model.stratify(
"organ",
["smearpos", "smearneg", "extrapul"],
["infectious"],
infectiousness_adjustments={
"smearpos": 1.0,
"smearneg": 0.25,
"extrapul": 0.0,
},
verbose=False,
requested_proportions=props_smear,
adjustment_requests={
"recovery": recovery_adjustments,
"infect_death": mortality_adjustments,
"case_detection": {
"smearpos": 1.0,
"smearneg": adjustment_smearneg,
"extrapul": adjustment_extrapul,
},
"early_progression": props_smear,
"late_progression": props_smear,
},
)
if "location" in stratify_by:
props_location = {
"rural_province": 0.48,
"urban_nonger": 0.368,
"urban_ger": 0.15,
"prison": 0.002,
}
raw_relative_risks_loc = {"rural_province": 1.0}
for stratum in ["urban_nonger", "urban_ger", "prison"]:
raw_relative_risks_loc[stratum] = external_params[
"rr_transmission_" + stratum]
scaled_relative_risks_loc = scale_relative_risks_for_equivalence(
props_location, raw_relative_risks_loc)
# dummy matrix for mixing by location
location_mixing = numpy.array([
0.899,
0.05,
0.05,
0.001,
0.049,
0.7,
0.25,
0.001,
0.049,
0.25,
0.7,
0.001,
0.1,
0.1,
0.1,
0.7,
]).reshape((4, 4))
location_mixing *= 3.0 # adjusted such that heterogeneous mixing yields similar overall burden as homogeneous
location_adjustments = {}
for beta_type in ["", "_late_latent", "_recovered"]:
location_adjustments["contact_rate" +
beta_type] = scaled_relative_risks_loc
location_adjustments["acf_rate"] = {}
for stratum in [
"rural_province", "urban_nonger", "urban_ger", "prison"
]:
location_adjustments["acf_rate"][stratum] = external_params[
"acf_" + stratum + "_switch"]
tb_model.stratify(
"location",
["rural_province", "urban_nonger", "urban_ger", "prison"],
[],
requested_proportions=props_location,
verbose=False,
entry_proportions=props_location,
adjustment_requests=location_adjustments,
mixing_matrix=location_mixing,
)
# tb_model.transition_flows.to_csv("transitions.csv")
# tb_model.death_flows.to_csv("deaths.csv")
# create some customised derived_outputs
if "notifications" in derived_output_types:
def notification_function_builder(stratum):
"""
example of stratum: "Xage_0Xstrain_mdr"
"""
def calculate_notifications(model, time):
total_notifications = 0.0
dict_flows = model.transition_flows_dict
comp_ind = model.compartment_names.index("infectious" +
stratum)
infectious_pop = model.compartment_values[comp_ind]
detection_indices = [
index for index, val in dict_flows["parameter"].items()
if "case_detection" in val
]
flow_index = [
index for index in detection_indices
if dict_flows["origin"][index] ==
model.compartment_names[comp_ind]
][0]
param_name = dict_flows["parameter"][flow_index]
detection_tx_rate = model.get_parameter_value(param_name, time)
tsr = mongolia_tsr(time) + external_params[
"reduction_negative_tx_outcome"] * (1.0 -
mongolia_tsr(time))
if "strain_mdr" in model.compartment_names[comp_ind]:
tsr = external_params["mdr_tsr"] * external_params[
"prop_mdr_detected_as_mdr"]
if tsr > 0.0:
total_notifications += infectious_pop * detection_tx_rate / tsr
return total_notifications
return calculate_notifications
for compartment in tb_model.compartment_names:
if "infectious" in compartment:
stratum = compartment.split("infectious")[1]
tb_model.derived_output_functions[
"notifications" +
stratum] = notification_function_builder(stratum)
# tb_model.derived_output_functions['popsize_treatment_support' + stratum] = notification_function_builder(stratum)
if "incidence" in derived_output_types:
# add output_connections for all stratum-specific incidence outputs
incidence_output_conns = create_output_connections_for_incidence_by_stratum(
tb_model.compartment_names)
tb_model.output_connections.update(incidence_output_conns)
# Create a 'combined incidence' derived output
early_names = [
k for k in incidence_output_conns.keys()
if k.startswith("incidence_early")
]
for early_name in early_names:
rootname = early_name[15:]
late_name = f"incidence_late{rootname}"
combined_name = f"incidence{rootname}"
def add_combined_incidence(model, time, e=early_name, l=late_name):
time_idx = model.times.index(time)
early_incidence = model.derived_outputs[e][time_idx]
late_incidence = model.derived_outputs[l][time_idx]
return early_incidence + late_incidence
tb_model.derived_output_functions[
combined_name] = add_combined_incidence
if "mortality" in derived_output_types:
# prepare death outputs for all strata
tb_model.death_output_categories = list_all_strata_for_mortality(
tb_model.compartment_names)
############################################
# population sizes for costing
############################################
if "popsizes" in derived_output_types:
# nb of detected individuals by strain:
def detected_popsize_function_builder(tag):
"""
example of tag: "starin_mdr" or "organ_smearpos"
"""
def calculate_nb_detected(model, time):
nb_treated = 0.0
for key, value in model.derived_outputs.items():
if "notifications" in key and tag in key:
this_time_index = model.times.index(time)
nb_treated += value[this_time_index]
return nb_treated
return calculate_nb_detected
for tag in [
"strain_mdr",
"strain_ds",
"organ_smearpos",
"organ_smearneg",
"organ_extrapul",
]:
tb_model.derived_output_functions[
"popsizeXnb_detectedX" +
tag] = detected_popsize_function_builder(tag)
# ACF popsize: number of people screened
def popsize_acf(model, time):
if external_params["acf_coverage"] == 0.0:
return 0.0
pop_urban_ger = sum([
model.compartment_values[i]
for i, c_name in enumerate(model.compartment_names)
if "location_urban_ger" in c_name
])
return external_params["acf_coverage"] * pop_urban_ger
tb_model.derived_output_functions[
"popsizeXnb_screened_acf"] = popsize_acf
return tb_model
def build_model(country: str, params: dict, update_params={}):
"""
Build the master function to run the TB model for Covid-19
:param update_params: dict
Any parameters that need to be updated for the current run
:return: StratifiedModel
The final model with all parameters and stratifications
"""
model_parameters = preprocess_params(params, update_params)
# Get population size (by age if age-stratified)
total_pops, model_parameters = get_population_size(model_parameters,
input_database)
# Replace with Victorian populations
# total_pops = load_population('31010DO001_201906.XLS', 'Table_6')
# total_pops = \
# [
# int(pop) for pop in
# total_pops.loc[
# (i_pop for i_pop in total_pops.index if 'Persons' in i_pop),
# 'Victoria'
# ]
# ]
# Define compartments with repeats as needed
all_compartments = [
Compartment.SUSCEPTIBLE,
Compartment.EXPOSED,
Compartment.PRESYMPTOMATIC,
Compartment.EARLY_INFECTIOUS,
Compartment.LATE_INFECTIOUS,
Compartment.RECOVERED,
]
final_compartments, replicated_compartments = [], []
for compartment in all_compartments:
if params["n_compartment_repeats"][compartment] == 1:
final_compartments.append(compartment)
else:
replicated_compartments.append(compartment)
is_infectious = {
Compartment.EXPOSED: False,
Compartment.PRESYMPTOMATIC: True,
Compartment.EARLY_INFECTIOUS: True,
Compartment.LATE_INFECTIOUS: True,
}
# Get progression rates from sojourn times, distinguishing to_infectious in order to split this parameter later
for compartment in params["compartment_periods"]:
model_parameters[
"within_" +
compartment] = 1.0 / params["compartment_periods"][compartment]
# Multiply the progression rates by the number of compartments to keep the average time in exposed the same
for compartment in is_infectious:
model_parameters["within_" + compartment] *= float(
model_parameters["n_compartment_repeats"][compartment])
for state in ["hospital_early", "icu_early"]:
model_parameters["within_" + state] *= float(
model_parameters["n_compartment_repeats"][
Compartment.EARLY_INFECTIOUS])
for state in ["hospital_late", "icu_late"]:
model_parameters["within_" + state] *= float(
model_parameters["n_compartment_repeats"][
Compartment.LATE_INFECTIOUS])
# Distribute infectious seed across infectious compartments
total_infectious_times = sum([
model_parameters["compartment_periods"][comp] for comp in is_infectious
])
init_pop = \
{comp: model_parameters["infectious_seed"] *
model_parameters["compartment_periods"][comp] /
total_infectious_times for
comp in is_infectious}
# force the remainder starting population to go to S compartment. Required as entry_compartment is late_infectious
init_pop[Compartment.SUSCEPTIBLE] = sum(total_pops) - sum(
init_pop.values())
# Set integration times
integration_times = get_model_times_from_inputs(
round(model_parameters["start_time"]),
model_parameters["end_time"],
model_parameters["time_step"],
)
# Add flows through replicated compartments
flows = []
for compartment in is_infectious:
flows = add_sequential_compartment_flows(
flows, model_parameters["n_compartment_repeats"][compartment],
compartment)
# Add other flows between compartment types
flows = add_infection_flows(
flows, model_parameters["n_compartment_repeats"][Compartment.EXPOSED])
flows = add_transition_flows(
flows,
model_parameters["n_compartment_repeats"][Compartment.EXPOSED],
model_parameters["n_compartment_repeats"][Compartment.PRESYMPTOMATIC],
Compartment.EXPOSED,
Compartment.PRESYMPTOMATIC,
"within_exposed",
)
# Distinguish to_infectious parameter, so that it can be split later
model_parameters["to_infectious"] = model_parameters["within_presympt"]
flows = add_transition_flows(
flows,
model_parameters["n_compartment_repeats"][Compartment.PRESYMPTOMATIC],
model_parameters["n_compartment_repeats"][
Compartment.EARLY_INFECTIOUS],
Compartment.PRESYMPTOMATIC,
Compartment.EARLY_INFECTIOUS,
"to_infectious",
)
flows = add_transition_flows(
flows,
model_parameters["n_compartment_repeats"][
Compartment.EARLY_INFECTIOUS],
model_parameters["n_compartment_repeats"][Compartment.LATE_INFECTIOUS],
Compartment.EARLY_INFECTIOUS,
Compartment.LATE_INFECTIOUS,
"within_" + Compartment.EARLY_INFECTIOUS,
)
flows = add_recovery_flows(
flows,
model_parameters["n_compartment_repeats"][Compartment.LATE_INFECTIOUS])
flows = add_infection_death_flows(
flows,
model_parameters["n_compartment_repeats"][Compartment.LATE_INFECTIOUS])
_birth_approach = 'no_birth' # may be changed if case importation implemented
# Add importation flows if requested
if model_parameters["implement_importation"] and not model_parameters[
"imported_cases_explict"]:
flows = add_transition_flows(
flows,
1,
model_parameters["n_compartment_repeats"][Compartment.EXPOSED],
Compartment.SUSCEPTIBLE,
Compartment.EXPOSED,
"import_secondary_rate",
)
elif model_parameters["implement_importation"] and model_parameters[
"imported_cases_explict"]:
_birth_approach = 'add_crude_birth_rate'
# Get mixing matrix
mixing_matrix = load_specific_prem_sheet("all_locations",
model_parameters["country"])
mixing_matrix_multipliers = model_parameters.get(
"mixing_matrix_multipliers")
if mixing_matrix_multipliers is not None:
mixing_matrix = update_mixing_with_multipliers(
mixing_matrix, mixing_matrix_multipliers)
# Define output connections to collate
output_connections = find_incidence_outputs(model_parameters)
# Define model
_covid_model = StratifiedModel(
integration_times,
final_compartments,
init_pop,
model_parameters,
flows,
birth_approach=_birth_approach,
entry_compartment=Compartment.
LATE_INFECTIOUS, # to model imported cases
starting_population=sum(total_pops),
infectious_compartment=[
i_comp for i_comp in is_infectious if is_infectious[i_comp]
],
)
# set time-variant importation rate
if model_parameters["implement_importation"] and not model_parameters[
"imported_cases_explict"]:
_covid_model = set_tv_importation_rate(
_covid_model, params["data"]["times_imported_cases"],
params["data"]["n_imported_cases"])
elif model_parameters["implement_importation"] and model_parameters[
"imported_cases_explict"]:
_covid_model = set_tv_importation_as_birth_rates(
_covid_model, params["data"]["times_imported_cases"],
params["data"]["n_imported_cases"])
# Stratify model by age
if "agegroup" in model_parameters["stratify_by"]:
age_strata = model_parameters["all_stratifications"]["agegroup"]
adjust_requests = split_multiple_parameters(
("to_infectious", "infect_death", "within_late"),
age_strata) # Split unchanged parameters for later adjustment
if model_parameters["implement_importation"] and not model_parameters[
"imported_cases_explict"]:
adjust_requests.update({
"import_secondary_rate":
get_total_contact_rates_by_age(mixing_matrix,
direction="horizontal")
})
# Adjust susceptibility for children
adjust_requests.update({
"contact_rate": {
key: value
for key, value in zip(
model_parameters["reduced_susceptibility_agegroups"],
[model_parameters["young_reduced_susceptibility"]] *
len(model_parameters["reduced_susceptibility_agegroups"]),
)
}
})
_covid_model.stratify(
"agegroup", # Don't use the string age, to avoid triggering automatic demography
convert_list_contents_to_int(age_strata),
[], # Apply to all compartments
{
i_break: prop
for i_break, prop in zip(age_strata,
normalise_sequence(total_pops))
}, # Distribute starting population
mixing_matrix=mixing_matrix,
adjustment_requests=adjust_requests,
verbose=False,
entry_proportions=importation_props_by_age)
# Stratify infectious compartment by clinical status
if "clinical" in model_parameters["stratify_by"] and model_parameters[
"clinical_strata"]:
_covid_model, model_parameters = stratify_by_clinical(
_covid_model, model_parameters, final_compartments)
# Add fully stratified incidence to output_connections
output_connections.update(
create_fully_stratified_incidence_covid(
model_parameters["stratify_by"],
model_parameters["all_stratifications"],
model_parameters,
))
output_connections.update(
create_fully_stratified_progress_covid(
model_parameters["stratify_by"],
model_parameters["all_stratifications"],
model_parameters,
))
_covid_model.output_connections = output_connections
# Add notifications to derived_outputs
_covid_model.derived_output_functions[
"notifications"] = calculate_notifications_covid
_covid_model.death_output_categories = list_all_strata_for_mortality(
_covid_model.compartment_names)
_covid_model.derived_output_functions[
"incidence_icu"] = calculate_incidence_icu_covid
# Do mixing matrix stuff
mixing_instructions = model_parameters.get("mixing")
if mixing_instructions:
if "npi_effectiveness" in model_parameters:
mixing_instructions = apply_npi_effectiveness(
mixing_instructions, model_parameters.get("npi_effectiveness"))
_covid_model.find_dynamic_mixing_matrix = build_covid_matrices(
model_parameters["country"], mixing_instructions)
_covid_model.dynamic_mixing_matrix = True
return _covid_model
def build_model(params: dict, update_params={}):
"""
Build the master function to run the TB model for the Republic of the Marshall Islands
:param update_params: dict
Any parameters that need to be updated for the current run
:return: StratifiedModel
The final model with all parameters and stratifications
"""
input_database = Database(database_name=INPUT_DB_PATH)
# Define compartments and initial conditions.
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.EARLY_LATENT,
Compartment.LATE_LATENT,
Compartment.EARLY_INFECTIOUS,
Compartment.ON_TREATMENT,
Compartment.RECOVERED,
# Compartment.LTBI_TREATED,
]
init_pop = {Compartment.EARLY_INFECTIOUS: 10, Compartment.LATE_LATENT: 100}
model_parameters = params
model_parameters.update(update_params)
# Update partial immunity/susceptibility parameters
model_parameters = update_transmission_parameters(
model_parameters, [Compartment.RECOVERED, Compartment.LATE_LATENT, Compartment.LTBI_TREATED]
)
# Set integration times
integration_times = get_model_times_from_inputs(
model_parameters["start_time"], model_parameters["end_time"], model_parameters["time_step"]
)
# Sequentially add groups of flows to flows list
flows = add_standard_infection_flows([])
flows = add_standard_latency_flows(flows)
flows = add_standard_natural_history_flows(flows)
# flows = add_latency_progression(flows)
flows = add_case_detection(flows, compartments)
flows = add_treatment_flows(flows)
# flows = add_acf(flows, compartments)
# flows = add_acf_ltbi(flows)
# Make sure incidence and notifications are tracked during integration
out_connections = {}
out_connections.update(
create_request_stratified_incidence(
model_parameters["incidence_stratification"], model_parameters["all_stratifications"]
)
)
out_connections.update(
create_request_stratified_notifications(
model_parameters["notification_stratifications"],
model_parameters["all_stratifications"],
)
)
# Define model
tb_model = StratifiedModel(
integration_times,
compartments,
init_pop,
model_parameters,
flows,
birth_approach="add_crude_birth_rate",
starting_population=model_parameters["start_population"],
output_connections=out_connections,
death_output_categories=list_all_strata_for_mortality(compartments),
)
# Add crude birth rate from UN estimates (using Federated States of Micronesia as a proxy as no data for RMI)
tb_model = add_birth_rate_functions(tb_model, input_database, "FSM")
# Find raw case detection rate with multiplier, which is 1 by default, and adjust for differences by organ status
cdr_scaleup_raw = build_scale_up_function(
model_parameters["cdr"], model_parameters["cdr_multiplier"]
)
detect_rate_by_organ = find_organ_specific_cdr(
cdr_scaleup_raw,
model_parameters,
model_parameters["all_stratifications"]["organ"],
target_organ_props=model_parameters["target_organ_props"],
)
# Find base case detection rate and time-variant treatment completion function
base_detection_rate = detect_rate_by_organ[
"smearpos" if "organ" in model_parameters["stratify_by"] else "overall"
]
treatment_success_rate = (
lambda time: build_scale_up_function(model_parameters["tsr"])(time)
/ model_parameters["treatment_duration"]
)
treatment_nonsuccess_rate = (
lambda time: (1.0 - build_scale_up_function(model_parameters["tsr"])(time))
/ model_parameters["treatment_duration"]
)
# Set acf screening rate using proportion of population reached and duration of intervention
# acf_screening_rate = -numpy.log(1 - 0.9) / 0.5
# acf_rate_over_time = progressive_step_function_maker(
# 2018.2, 2018.7, acf_screening_rate, scaling_time_fraction=0.3
# )
# Initialise acf_rate function
# acf_rate_function = (
# lambda t: model_parameters["acf_coverage"]
# * (acf_rate_over_time(t))
# * model_parameters["acf_sensitivity"]
# )
# acf_ltbi_rate_function = (
# lambda t: model_parameters["acf_coverage"]
# * (acf_rate_over_time(t))
# * model_parameters["acf_ltbi_sensitivity"]
# * model_parameters["acf_ltbi_efficacy"]
# )
# Assign newly created functions to model parameters
add_time_variant_parameter_to_model(
tb_model, "case_detection", base_detection_rate, len(model_parameters["stratify_by"])
)
add_time_variant_parameter_to_model(
tb_model, "treatment_success", treatment_success_rate, len(model_parameters["stratify_by"])
)
add_time_variant_parameter_to_model(
tb_model,
"treatment_nonsuccess",
treatment_nonsuccess_rate,
len(model_parameters["stratify_by"]),
)
# add_time_variant_parameter_to_model(
# tb_model, 'acf_rate', acf_rate_function, len(model_parameters['stratify_by']))
# add_time_variant_parameter_to_model(
# tb_model, 'acf_ltbi_rate', acf_ltbi_rate_function, len(model_parameters['stratify_by']))
# Stratification processes
if "age" in model_parameters["stratify_by"]:
age_specific_latency_parameters = manually_create_age_specific_latency_parameters(
model_parameters
)
tb_model = stratify_by_age(
tb_model,
age_specific_latency_parameters,
input_database,
model_parameters["all_stratifications"]["age"],
)
if "diabetes" in model_parameters["stratify_by"]:
tb_model = stratify_by_diabetes(
tb_model,
model_parameters,
model_parameters["all_stratifications"]["diabetes"],
model_parameters["diabetes_target_props"],
age_specific_prevalence=False,
)
if "organ" in model_parameters["stratify_by"]:
tb_model = stratify_by_organ(
tb_model,
model_parameters,
detect_rate_by_organ,
model_parameters["all_stratifications"]["organ"],
)
if "location" in model_parameters["stratify_by"]:
tb_model = stratify_by_location(
tb_model, model_parameters, model_parameters["all_stratifications"]["location"]
)
# Capture reported prevalence in Majuro assuming over-reporting (needed for calibration)
def calculate_reported_majuro_prevalence(model, time):
true_prev = 0.0
pop_majuro = 0.0
for i, compartment in enumerate(model.compartment_names):
if "majuro" in compartment:
pop_majuro += model.compartment_values[i]
if "infectious" in compartment:
true_prev += model.compartment_values[i]
return (
1.0e5
* true_prev
/ pop_majuro
* (1.0 + model_parameters["over_reporting_prevalence_proportion"])
)
tb_model.derived_output_functions.update(
{"reported_majuro_prevalence": calculate_reported_majuro_prevalence}
)
return tb_model
def build_model(params: dict) -> StratifiedModel:
"""
Build the master function to run the TB model for Covid-19
"""
validate_params(params)
# Get the agegroup strata breakpoints.
agegroup_max = params["agegroup_breaks"][0]
agegroup_step = params["agegroup_breaks"][1]
agegroup_strata = list(range(0, agegroup_max, agegroup_step))
# Look up the country population size by age-group, using UN data
country_iso3 = params["iso3"]
region = params["region"]
total_pops = inputs.get_population_by_agegroup(agegroup_strata,
country_iso3,
region,
year=2020)
life_expectancy = inputs.get_life_expectancy_by_agegroup(
agegroup_strata, country_iso3)[0]
life_expectancy_latest = [
life_expectancy[agegroup][-1] for agegroup in life_expectancy
]
# Define compartments
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.EXPOSED,
Compartment.PRESYMPTOMATIC,
Compartment.EARLY_INFECTIOUS,
Compartment.LATE_INFECTIOUS,
Compartment.RECOVERED,
]
# Indicate whether the compartments representing active disease are infectious
is_infectious = {
Compartment.EXPOSED: False,
Compartment.PRESYMPTOMATIC: True,
Compartment.EARLY_INFECTIOUS: True,
Compartment.LATE_INFECTIOUS: True,
}
# Calculate compartment periods
# FIXME: Needs tests.
base_compartment_periods = params["compartment_periods"]
compartment_periods_calc = params["compartment_periods_calculated"]
compartment_periods = preprocess.compartments.calc_compartment_periods(
base_compartment_periods, compartment_periods_calc)
# Get progression rates from sojourn times, distinguishing to_infectious in order to split this parameter later
compartment_exit_flow_rates = {}
for compartment in compartment_periods:
param_key = f"within_{compartment}"
compartment_exit_flow_rates[
param_key] = 1.0 / compartment_periods[compartment]
# Distribute infectious seed across infectious compartments
infectious_seed = params["infectious_seed"]
total_disease_time = sum([compartment_periods[c] for c in is_infectious])
init_pop = {
c: infectious_seed * compartment_periods[c] / total_disease_time
for c in is_infectious
}
# Force the remainder starting population to go to S compartment (Required as entry_compartment is late_infectious)
init_pop[Compartment.SUSCEPTIBLE] = sum(total_pops) - sum(
init_pop.values())
# Set integration times
start_time = params["start_time"]
end_time = params["end_time"]
time_step = params["time_step"]
integration_times = get_model_times_from_inputs(
round(start_time),
end_time,
time_step,
)
# Add inter-compartmental transition flows
flows = preprocess.flows.DEFAULT_FLOWS
# Choose a birth approach
is_importation_active = params["implement_importation"]
birth_approach = BirthApproach.ADD_CRUDE if is_importation_active else BirthApproach.NO_BIRTH
# Build mixing matrix.
static_mixing_matrix = preprocess.mixing_matrix.build_static(country_iso3)
dynamic_mixing_matrix = None
dynamic_location_mixing_params = params["mixing"]
dynamic_age_mixing_params = params["mixing_age_adjust"]
microdistancing = params["microdistancing"]
if dynamic_location_mixing_params or dynamic_age_mixing_params:
npi_effectiveness_params = params["npi_effectiveness"]
google_mobility_locations = params["google_mobility_locations"]
is_periodic_intervention = params.get("is_periodic_intervention")
periodic_int_params = params.get("periodic_intervention")
dynamic_mixing_matrix = preprocess.mixing_matrix.build_dynamic(
country_iso3,
region,
dynamic_location_mixing_params,
dynamic_age_mixing_params,
npi_effectiveness_params,
google_mobility_locations,
is_periodic_intervention,
periodic_int_params,
end_time,
microdistancing,
)
# FIXME: Remove params from model_parameters
model_parameters = {**params, **compartment_exit_flow_rates}
model_parameters["to_infectious"] = model_parameters["within_presympt"]
# Instantiate SUMMER model
model = StratifiedModel(
integration_times,
compartments,
init_pop,
model_parameters,
flows,
birth_approach=birth_approach,
entry_compartment=Compartment.
LATE_INFECTIOUS, # to model imported cases
starting_population=sum(total_pops),
infectious_compartment=[
i_comp for i_comp in is_infectious if is_infectious[i_comp]
],
)
if dynamic_mixing_matrix:
model.find_dynamic_mixing_matrix = dynamic_mixing_matrix
model.dynamic_mixing_matrix = True
# Implement seasonal forcing if requested, making contact rate a time-variant rather than constant
if model_parameters["seasonal_force"]:
seasonal_forcing_function = \
get_seasonal_forcing(
365., 173., model_parameters["seasonal_force"], model_parameters["contact_rate"]
)
model.time_variants["contact_rate"] = \
seasonal_forcing_function
model.adaptation_functions["contact_rate"] = \
seasonal_forcing_function
model.parameters["contact_rate"] = \
"contact_rate"
# Stratify model by age
# Coerce age breakpoint numbers into strings - all strata are represented as strings
agegroup_strata = [str(s) for s in agegroup_strata]
# Create parameter adjustment request for age stratifications
age_based_susceptibility = params["age_based_susceptibility"]
adjust_requests = {
# No change, but distinction is required for later stratification by clinical status
"to_infectious": {s: 1
for s in agegroup_strata},
"infect_death": {s: 1
for s in agegroup_strata},
"within_late": {s: 1
for s in agegroup_strata},
# Adjust susceptibility across age groups
"contact_rate": age_based_susceptibility,
}
if is_importation_active:
adjust_requests[
"import_secondary_rate"] = preprocess.mixing_matrix.get_total_contact_rates_by_age(
static_mixing_matrix, direction="horizontal")
# Distribute starting population over agegroups
requested_props = {
agegroup: prop
for agegroup, prop in zip(agegroup_strata,
normalise_sequence(total_pops))
}
# We use "agegroup" instead of "age" for this model, to avoid triggering automatic demography features
# (which work on the assumption that the time unit is years, so would be totally wrong)
model.stratify(
"agegroup",
agegroup_strata,
compartment_types_to_stratify=[], # Apply to all compartments
requested_proportions=requested_props,
mixing_matrix=static_mixing_matrix,
adjustment_requests=adjust_requests,
# FIXME: This seems awfully a lot like a parameter that should go in a YAML file.
entry_proportions=preprocess.importation.IMPORTATION_PROPS_BY_AGE,
)
model_parameters["all_stratifications"] = {"agegroup": agegroup_strata}
modelled_abs_detection_proportion_imported = stratify_by_clinical(
model, model_parameters, compartments)
# Set time-variant importation rate
if is_importation_active:
import_times = params["data"]["times_imported_cases"]
import_cases = params["data"]["n_imported_cases"]
import_rate_func = preprocess.importation.get_importation_rate_func_as_birth_rates(
import_times,
import_cases,
modelled_abs_detection_proportion_imported,
total_pops,
)
model.parameters["crude_birth_rate"] = "crude_birth_rate"
model.time_variants["crude_birth_rate"] = import_rate_func
# Define output connections to collate
# Track compartment output connections.
stratum_names = list(
set([find_all_strata(x) for x in model.compartment_names]))
incidence_connections = outputs.get_incidence_connections(stratum_names)
progress_connections = outputs.get_progress_connections(stratum_names)
model.output_connections = {
**incidence_connections,
**progress_connections,
}
# Add notifications to derived_outputs
implement_importation = model.parameters["implement_importation"]
model.derived_output_functions[
"notifications"] = outputs.get_calc_notifications_covid(
implement_importation,
modelled_abs_detection_proportion_imported,
)
model.derived_output_functions[
"incidence_icu"] = outputs.calculate_incidence_icu_covid
model.derived_output_functions[
"prevXlateXclinical_icuXamong"] = outputs.calculate_icu_prev
model.derived_output_functions[
"hospital_occupancy"] = outputs.calculate_hospital_occupancy
model.derived_output_functions[
"proportion_seropositive"] = outputs.calculate_proportion_seropositive
model.death_output_categories = list_all_strata_for_mortality(
model.compartment_names)
model.derived_output_functions[
"years_of_life_lost"] = outputs.get_calculate_years_of_life_lost(
life_expectancy_latest)
return model