example_flu.py

"""An illustration of the use of the workbench for a model specified in Python itself.

The example is based on `Pruyt & Hamarat <https://www.systemdynamics.org/conferences/2010/proceed/papers/P1253.pdf>`_.
For comparison, run both this model and example_flu_vensim_no_policy.py.

"""

import matplotlib.pyplot as plt
import numpy as np
from numpy import exp, min, sin

from ema_workbench import (
    Model,
    RealParameter,
    SequentialEvaluator,
    TimeSeriesOutcome,
    ema_logging,
    perform_experiments,
)
from ema_workbench.analysis import Density, lines

# Created on 20 dec. 2010
# .. codeauthor:: jhkwakkel <j.h.kwakkel (at) tudelft (dot) nl>
#                 chamarat <c.hamarat  (at) tudelft (dot) nl>


# =============================================================================
#
#    the model itself
#
# =============================================================================

FINAL_TIME = 48
INITIAL_TIME = 0
TIME_STEP = 0.0078125

switch_regions = 1.0
switch_immunity = 1.0
switch_deaths = 1.0
switch_immunity_cap = 1.0


def lookup_function_x(variable, start, end, step, skew, growth, v=0.5):
    """A basic lookup function."""
    return start + (
        (end - start) / ((1 + skew * exp(-growth * (variable - step))) ** (1 / v))
    )


def flu_model(
    x11=0,
    x12=0,
    x21=0,
    x22=0,
    x31=0,
    x32=0,
    x41=0,
    x51=0,
    x52=0,
    x61=0,
    x62=0,
    x81=0,
    x82=0,
    x91=0,
    x92=0,
    x101=0,
    x102=0,
):
    """The overall model."""
    # Assigning initial values
    additional_seasonal_immune_population_fraction_r1 = float(x11)
    additional_seasonal_immune_population_fraction_r2 = float(x12)

    fatality_rate_region_1 = float(x21)
    fatality_rate_region_2 = float(x22)

    initial_immune_fraction_of_the_population_of_region_1 = float(x31)
    initial_immune_fraction_of_the_population_of_region_2 = float(x32)

    normal_interregional_contact_rate = float(x41)
    interregional_contact_rate = switch_regions * normal_interregional_contact_rate

    permanent_immune_population_fraction_r1 = float(x51)
    permanent_immune_population_fraction_r2 = float(x52)

    recovery_time_region_1 = float(x61)
    recovery_time_region_2 = float(x62)

    susceptible_to_immune_population_delay_time_region_1 = 1
    susceptible_to_immune_population_delay_time_region_2 = 1

    root_contact_rate_region_1 = float(x81)
    root_contact_rate_region_2 = float(x82)

    infection_rate_region_1 = float(x91)
    infection_rate_region_2 = float(x92)

    normal_contact_rate_region_1 = float(x101)
    normal_contact_rate_region_2 = float(x102)

    ######
    susceptible_to_immune_population_flow_region_1 = 0.0
    susceptible_to_immune_population_flow_region_2 = 0.0
    ######

    initial_value_population_region_1 = 6.0 * 10**8
    initial_value_population_region_2 = 3.0 * 10**9

    initial_value_infected_population_region_1 = 10.0
    initial_value_infected_population_region_2 = 10.0

    initial_value_immune_population_region_1 = (
        switch_immunity
        * initial_immune_fraction_of_the_population_of_region_1
        * initial_value_population_region_1
    )
    initial_value_immune_population_region_2 = (
        switch_immunity
        * initial_immune_fraction_of_the_population_of_region_2
        * initial_value_population_region_2
    )

    initial_value_susceptible_population_region_1 = (
        initial_value_population_region_1 - initial_value_immune_population_region_1
    )
    initial_value_susceptible_population_region_2 = (
        initial_value_population_region_2 - initial_value_immune_population_region_2
    )

    recovered_population_region_1 = 0.0
    recovered_population_region_2 = 0.0

    infected_population_region_1 = initial_value_infected_population_region_1
    infected_population_region_2 = initial_value_infected_population_region_2

    susceptible_population_region_1 = initial_value_susceptible_population_region_1
    susceptible_population_region_2 = initial_value_susceptible_population_region_2

    immune_population_region_1 = initial_value_immune_population_region_1
    immune_population_region_2 = initial_value_immune_population_region_2

    deceased_population_region_1 = [0.0]
    deceased_population_region_2 = [0.0]
    run_time = [INITIAL_TIME]

    # --End of Initialization--

    max_infected = 0.0

    for time in range(int(INITIAL_TIME / TIME_STEP), int(FINAL_TIME / TIME_STEP)):
        run_time.append(run_time[-1] + TIME_STEP)
        total_population_region_1 = (
            infected_population_region_1
            + recovered_population_region_1
            + susceptible_population_region_1
            + immune_population_region_1
        )
        total_population_region_2 = (
            infected_population_region_2
            + recovered_population_region_2
            + susceptible_population_region_2
            + immune_population_region_2
        )

        infected_population_region_1 = max(0, infected_population_region_1)
        infected_population_region_2 = max(0, infected_population_region_2)

        infected_fraction_region_1 = (
            infected_population_region_1 / total_population_region_1
        )
        infected_fraction_region_2 = (
            infected_population_region_2 / total_population_region_2
        )

        impact_infected_population_on_contact_rate_region_1 = 1 - (
            infected_fraction_region_1 ** (1 / root_contact_rate_region_1)
        )
        impact_infected_population_on_contact_rate_region_2 = 1 - (
            infected_fraction_region_2 ** (1 / root_contact_rate_region_2)
        )

        #        if ((time*TIME_STEP) >= 4) and ((time*TIME_STEP)<=10):
        #            normal_contact_rate_region_1 = float(x101)*(1 - 0.5)
        #        else:normal_contact_rate_region_1 = float(x101)

        normal_contact_rate_region_1 = float(x101) * (
            1 - lookup_function_x(infected_fraction_region_1, 0, 1, 0.15, 0.75, 15)
        )

        contact_rate_region_1 = (
            normal_contact_rate_region_1
            * impact_infected_population_on_contact_rate_region_1
        )
        contact_rate_region_2 = (
            normal_contact_rate_region_2
            * impact_infected_population_on_contact_rate_region_2
        )

        recoveries_region_1 = (
            (1 - (fatality_rate_region_1 * switch_deaths))
            * infected_population_region_1
            / recovery_time_region_1
        )
        recoveries_region_2 = (
            (1 - (fatality_rate_region_2 * switch_deaths))
            * infected_population_region_2
            / recovery_time_region_2
        )

        flu_deaths_region_1 = (
            fatality_rate_region_1
            * switch_deaths
            * infected_population_region_1
            / recovery_time_region_1
        )
        flu_deaths_region_2 = (
            fatality_rate_region_2
            * switch_deaths
            * infected_population_region_2
            / recovery_time_region_2
        )

        infections_region_1 = (
            susceptible_population_region_1
            * contact_rate_region_1
            * infection_rate_region_1
            * infected_fraction_region_1
        ) + (
            susceptible_population_region_1
            * interregional_contact_rate
            * infection_rate_region_1
            * infected_fraction_region_2
        )
        infections_region_2 = (
            susceptible_population_region_2
            * contact_rate_region_2
            * infection_rate_region_2
            * infected_fraction_region_2
        ) + (
            susceptible_population_region_2
            * interregional_contact_rate
            * infection_rate_region_2
            * infected_fraction_region_1
        )

        infected_population_region_1_next = infected_population_region_1 + (
            TIME_STEP
            * (infections_region_1 - flu_deaths_region_1 - recoveries_region_1)
        )
        infected_population_region_2_next = infected_population_region_2 + (
            TIME_STEP
            * (infections_region_2 - flu_deaths_region_2 - recoveries_region_2)
        )

        if (
            infected_population_region_1_next < 0
            or infected_population_region_2_next < 0
        ):
            pass

        recovered_population_region_1_next = recovered_population_region_1 + (
            TIME_STEP * recoveries_region_1
        )
        recovered_population_region_2_next = recovered_population_region_2 + (
            TIME_STEP * recoveries_region_2
        )

        if fatality_rate_region_1 >= 0.025:
            qw = 1.0
        elif fatality_rate_region_1 >= 0.01:
            qw = 0.8
        elif fatality_rate_region_1 >= 0.001:
            qw = 0.6
        elif fatality_rate_region_1 >= 0.0001:
            qw = 0.4
        else:
            qw = 0.2

        if (time * TIME_STEP) <= 10:
            normal_immune_population_fraction_region_1 = (
                additional_seasonal_immune_population_fraction_r1 / 2
            ) * sin(4.5 + (time * TIME_STEP / 2)) + (
                (
                    (2 * permanent_immune_population_fraction_r1)
                    + additional_seasonal_immune_population_fraction_r1
                )
                / 2
            )
        else:
            normal_immune_population_fraction_region_1 = max(
                (
                    float(qw),
                    (additional_seasonal_immune_population_fraction_r1 / 2)
                    * sin(4.5 + (time * TIME_STEP / 2))
                    + (
                        (
                            (2 * permanent_immune_population_fraction_r1)
                            + additional_seasonal_immune_population_fraction_r1
                        )
                        / 2
                    ),
                )
            )

        normal_immune_population_fraction_region_2 = switch_immunity_cap * min(
            (
                (
                    sin((time * TIME_STEP / 2) + 1.5)
                    * additional_seasonal_immune_population_fraction_r2
                    / 2
                )
                + (
                    (
                        (2 * permanent_immune_population_fraction_r2)
                        + additional_seasonal_immune_population_fraction_r2
                    )
                    / 2
                ),
                (
                    permanent_immune_population_fraction_r1
                    + additional_seasonal_immune_population_fraction_r1
                ),
            ),
        ) + (
            (1 - switch_immunity_cap)
            * (
                (
                    sin((time * TIME_STEP / 2) + 1.5)
                    * additional_seasonal_immune_population_fraction_r2
                    / 2
                )
                + (
                    (
                        (2 * permanent_immune_population_fraction_r2)
                        + additional_seasonal_immune_population_fraction_r2
                    )
                    / 2
                )
            )
        )

        normal_immune_population_region_1 = (
            normal_immune_population_fraction_region_1 * total_population_region_1
        )
        normal_immune_population_region_2 = (
            normal_immune_population_fraction_region_2 * total_population_region_2
        )

        if switch_immunity == 1:
            susminreg1_1 = (
                normal_immune_population_region_1 - immune_population_region_1
            ) / susceptible_to_immune_population_delay_time_region_1
            susminreg1_2 = (
                susceptible_population_region_1
                / susceptible_to_immune_population_delay_time_region_1
            )
            susmaxreg1 = -(
                immune_population_region_1
                / susceptible_to_immune_population_delay_time_region_1
            )
            if (susmaxreg1 >= susminreg1_1) or (susmaxreg1 >= susminreg1_2):
                susceptible_to_immune_population_flow_region_1 = susmaxreg1
            elif (susminreg1_1 < susminreg1_2) and (susminreg1_1 > susmaxreg1):
                susceptible_to_immune_population_flow_region_1 = susminreg1_1
            elif (susminreg1_2 < susminreg1_1) and (susminreg1_2 > susmaxreg1):
                susceptible_to_immune_population_flow_region_1 = susminreg1_2
        else:
            susceptible_to_immune_population_flow_region_1 = 0

        if switch_immunity == 1:
            susminreg2_1 = (
                normal_immune_population_region_2 - immune_population_region_2
            ) / susceptible_to_immune_population_delay_time_region_2
            susminreg2_2 = (
                susceptible_population_region_2
                / susceptible_to_immune_population_delay_time_region_2
            )
            susmaxreg2 = -(
                immune_population_region_2
                / susceptible_to_immune_population_delay_time_region_2
            )
            if (susmaxreg2 >= susminreg2_1) or (susmaxreg2 >= susminreg2_2):
                susceptible_to_immune_population_flow_region_2 = susmaxreg2
            elif (susminreg2_1 < susminreg2_2) and (susminreg2_1 > susmaxreg2):
                susceptible_to_immune_population_flow_region_2 = susminreg2_1
            elif (susminreg2_2 < susminreg2_1) and (susminreg2_2 > susmaxreg2):
                susceptible_to_immune_population_flow_region_2 = susminreg2_2
        else:
            susceptible_to_immune_population_flow_region_2 = 0

        susceptible_population_region_1_next = susceptible_population_region_1 - (
            TIME_STEP
            * (infections_region_1 + susceptible_to_immune_population_flow_region_1)
        )
        susceptible_population_region_2_next = susceptible_population_region_2 - (
            TIME_STEP
            * (infections_region_2 + susceptible_to_immune_population_flow_region_2)
        )

        immune_population_region_1_next = immune_population_region_1 + (
            TIME_STEP * susceptible_to_immune_population_flow_region_1
        )
        immune_population_region_2_next = immune_population_region_2 + (
            TIME_STEP * susceptible_to_immune_population_flow_region_2
        )

        deceased_population_region_1_next = deceased_population_region_1[-1] + (
            TIME_STEP * flu_deaths_region_1
        )
        deceased_population_region_2_next = deceased_population_region_2[-1] + (
            TIME_STEP * flu_deaths_region_2
        )

        # Updating integral values
        if max_infected < (
            infected_population_region_1_next
            / (
                infected_population_region_1_next
                + recovered_population_region_1_next
                + susceptible_population_region_1_next
                + immune_population_region_1_next
            )
        ):
            max_infected = infected_population_region_1_next / (
                infected_population_region_1_next
                + recovered_population_region_1_next
                + susceptible_population_region_1_next
                + immune_population_region_1_next
            )

        recovered_population_region_1 = recovered_population_region_1_next
        recovered_population_region_2 = recovered_population_region_2_next

        infected_population_region_1 = infected_population_region_1_next
        infected_population_region_2 = infected_population_region_2_next

        susceptible_population_region_1 = susceptible_population_region_1_next
        susceptible_population_region_2 = susceptible_population_region_2_next

        immune_population_region_1 = immune_population_region_1_next
        immune_population_region_2 = immune_population_region_2_next

        deceased_population_region_1.append(deceased_population_region_1_next)
        deceased_population_region_2.append(deceased_population_region_2_next)

        # End of main code

    return {
        "TIME": run_time,
        "deceased_population_region_1": deceased_population_region_1,
    }


if __name__ == "__main__":
    ema_logging.log_to_stderr(ema_logging.INFO)

    model = Model("mexicanFlu", function=flu_model)
    model.uncertainties = [
        RealParameter("x11", 0, 0.5),
        RealParameter("x12", 0, 0.5),
        RealParameter("x21", 0.0001, 0.1),
        RealParameter("x22", 0.0001, 0.1),
        RealParameter("x31", 0, 0.5),
        RealParameter("x32", 0, 0.5),
        RealParameter("x41", 0, 0.9),
        RealParameter("x51", 0, 0.5),
        RealParameter("x52", 0, 0.5),
        RealParameter("x61", 0, 0.8),
        RealParameter("x62", 0, 0.8),
        RealParameter("x81", 1, 10),
        RealParameter("x82", 1, 10),
        RealParameter("x91", 0, 0.1),
        RealParameter("x92", 0, 0.1),
        RealParameter("x101", 0, 200),
        RealParameter("x102", 0, 200),
    ]

    model.outcomes = [
        TimeSeriesOutcome("TIME"),
        TimeSeriesOutcome("deceased_population_region_1"),
    ]

    nr_experiments = 500

    with SequentialEvaluator(model) as evaluator:
        experiments, outcomes = perform_experiments(
            model, nr_experiments, evaluator=evaluator
        )

    lines(
        experiments,
        outcomes,
        outcomes_to_show="deceased_population_region_1",
        show_envelope=True,
        density=Density.KDE,
        titles=None,
        experiments_to_show=np.arange(0, nr_experiments, 10),
    )
    plt.show()