Esempio n. 1
0
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
Esempio n. 2
0
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
Esempio n. 3
0
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
Esempio n. 4
0
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
Esempio n. 5
0
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()
Esempio n. 6
0
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()
Esempio n. 7
0
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
Esempio n. 8
0
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
Esempio n. 9
0
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()
Esempio n. 10
0
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
Esempio n. 11
0
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()
Esempio n. 12
0
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()
Esempio n. 13
0
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
Esempio n. 14
0
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
Esempio n. 15
0
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
Esempio n. 16
0
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
Esempio n. 17
0
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