Skip to content

anypinn.cli.scaffold.sir.ode_synthetic

SIR epidemic model — mathematical definition.

BETA_KEY = 'beta' module-attribute

C = 1000000.0 module-attribute

DELTA = 1 / 5 module-attribute

DELTA_KEY = 'delta' module-attribute

I_KEY = 'I' module-attribute

N_KEY = 'N' module-attribute

N_POP = 56000000.0 module-attribute

S_KEY = 'S' module-attribute

T = 90 module-attribute

TRUE_BETA = 0.6 module-attribute

validation: ValidationRegistry = {BETA_KEY: lambda x: torch.full_like(x, TRUE_BETA)} module-attribute

SIR(x: Tensor, y: Tensor, args: ArgsRegistry) -> Tensor

Scaled SIR ODE system.

Source code in src/anypinn/cli/scaffold/sir/ode_synthetic.py 
50
51
52
53
54
55
56
57
58
59
60
def SIR(x: Tensor, y: Tensor, args: ArgsRegistry) -> Tensor:
    """Scaled SIR ODE system."""
    S, I = y
    b, d, N = args[BETA_KEY], args[DELTA_KEY], args[N_KEY]

    dS = -b(x) * I * S * C / N(x)
    dI = b(x) * I * S * C / N(x) - d(x) * I

    dS = dS * T
    dI = dI * T
    return torch.stack([dS, dI])

create_data_module(hp: ODEHyperparameters)

Source code in src/anypinn/cli/scaffold/sir/ode_synthetic.py 
76
77
78
79
80
81
82
83
def create_data_module(hp: ODEHyperparameters):
    from anypinn.catalog.sir import SIRInvDataModule

    return SIRInvDataModule(
        hp=hp,
        validation=validation,
        callbacks=[DataScaling(y_scale=1 / C)],
    )

create_problem(hp: ODEHyperparameters) -> ODEInverseProblem

Source code in src/anypinn/cli/scaffold/sir/ode_synthetic.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def create_problem(hp: ODEHyperparameters) -> ODEInverseProblem:
    props = ODEProperties(
        ode=SIR,
        y0=torch.tensor([N_POP - 1, 1]) / C,
        args={
            DELTA_KEY: Argument(DELTA),
            N_KEY: Argument(N_POP),
        },
    )

    fields = FieldsRegistry(
        {
            S_KEY: Field(config=hp.fields_config),
            I_KEY: Field(config=hp.fields_config),
        }
    )
    params = ParamsRegistry(
        {
            BETA_KEY: Parameter(config=hp.params_config),
        }
    )

    def predict_data(x_data: Tensor, fields: FieldsRegistry, _params: ParamsRegistry) -> Tensor:
        I_pred = fields[I_KEY](x_data)
        return I_pred.unsqueeze(1)

    return ODEInverseProblem(
        props=props,
        hp=hp,
        fields=fields,
        params=params,
        predict_data=predict_data,
    )