anypinn.catalog.sir

BETA_KEY = 'beta' module-attribute

DELTA_KEY = 'delta' module-attribute

I_KEY = 'I' module-attribute

N_KEY = 'N' module-attribute

Rt_KEY = 'Rt' module-attribute

S_KEY = 'S' module-attribute

SIRInvDataModule

Bases: PINNDataModule

DataModule for SIR Inverse problem.

Source code in src/anypinn/catalog/sir.py

class SIRInvDataModule(PINNDataModule):
    """
    DataModule for SIR Inverse problem.
    """

    def __init__(
        self,
        hp: ODEHyperparameters,
        gen_props: ODEProperties | None = None,
        validation: ValidationRegistry | None = None,
        callbacks: Sequence[DataCallback] | None = None,
    ):
        super().__init__(hp, validation, callbacks)
        self.gen_props = gen_props

    @override
    def gen_data(self, config: GenerationConfig) -> tuple[Tensor, Tensor]:
        """Generate synthetic data."""
        assert self.gen_props is not None, "SIR properties are required to generate data"
        gen_props = self.gen_props

        args = gen_props.args.copy()
        args.update(config.args_to_train)

        data = odeint(
            lambda x, y: gen_props.ode(x, y, args),
            gen_props.y0,
            config.x,
        )

        I_true = data[:, 1].clamp_min(0.0)

        I_obs = self._noise(I_true, config.noise_level)

        return config.x.unsqueeze(-1), I_obs.unsqueeze(-1).unsqueeze(1)

    def _noise(self, I_true: Tensor, noise_level: float) -> Tensor:
        if noise_level < 1.0:
            return I_true
        else:
            return torch.poisson(I_true / noise_level) * noise_level

gen_props = gen_props instance-attribute

__init__(hp: ODEHyperparameters, gen_props: ODEProperties | None = None, validation: ValidationRegistry | None = None, callbacks: Sequence[DataCallback] | None = None)

Source code in src/anypinn/catalog/sir.py

def __init__(
    self,
    hp: ODEHyperparameters,
    gen_props: ODEProperties | None = None,
    validation: ValidationRegistry | None = None,
    callbacks: Sequence[DataCallback] | None = None,
):
    super().__init__(hp, validation, callbacks)
    self.gen_props = gen_props

gen_data(config: GenerationConfig) -> tuple[Tensor, Tensor]

Generate synthetic data.

Source code in src/anypinn/catalog/sir.py

@override
def gen_data(self, config: GenerationConfig) -> tuple[Tensor, Tensor]:
    """Generate synthetic data."""
    assert self.gen_props is not None, "SIR properties are required to generate data"
    gen_props = self.gen_props

    args = gen_props.args.copy()
    args.update(config.args_to_train)

    data = odeint(
        lambda x, y: gen_props.ode(x, y, args),
        gen_props.y0,
        config.x,
    )

    I_true = data[:, 1].clamp_min(0.0)

    I_obs = self._noise(I_true, config.noise_level)

    return config.x.unsqueeze(-1), I_obs.unsqueeze(-1).unsqueeze(1)

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

The SIR ODE system. $$ \begin{align} \frac{dS}{dt} &= -\beta \frac{SI}{N} \ \frac{dI}{dt} &= \beta \frac{SI}{N} - \delta I \ \frac{dR}{dt} &= \delta I \ \end{align} $$

Parameters:

Name	Type	Description	Default
x	Tensor	Time variable.	required
y	Tensor	State variables [S, I].	required
args	ArgsRegistry	Arguments dictionary (beta, delta, N).	required

Returns:

Type	Description
Tensor	Derivatives [dS/dt, dI/dt].

Source code in src/anypinn/catalog/sir.py

def SIR(x: Tensor, y: Tensor, args: ArgsRegistry) -> Tensor:
    """
    The SIR ODE system.
    $$
    \\begin{align}
    \\frac{dS}{dt} &= -\\beta \\frac{SI}{N} \\\\
    \\frac{dI}{dt} &= \\beta \\frac{SI}{N} - \\delta I \\\\
    \\frac{dR}{dt} &= \\delta I \\\\
    \\end{align}
    $$

    Args:
        x: Time variable.
        y: State variables [S, I].
        args: Arguments dictionary (beta, delta, N).

    Returns:
        Derivatives [dS/dt, dI/dt].
    """
    S, I = y
    b, d, N = args[BETA_KEY], args[DELTA_KEY], args[N_KEY]

    dS = -b(x) * S * I / N(x)
    dI = b(x) * S * I / N(x) - d(x) * I
    return torch.stack([dS, dI])

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

The reduced SIR ODE system. $$ \begin{align} \frac{dI}{dt} &= \delta (R_t - 1) I \end{align} $$

Parameters:

Name	Type	Description	Default
x	Tensor	Time variable.	required
y	Tensor	State variables [I].	required
args	ArgsRegistry	Arguments dictionary (delta, Rt).	required

Returns:

Type	Description
Tensor	Derivatives [dI/dt].

Source code in src/anypinn/catalog/sir.py

def rSIR(x: Tensor, y: Tensor, args: ArgsRegistry) -> Tensor:
    """
    The reduced SIR ODE system.
    $$
    \\begin{align}
    \\frac{dI}{dt} &= \\delta (R_t - 1) I
    \\end{align}
    $$

    Args:
        x: Time variable.
        y: State variables [I].
        args: Arguments dictionary (delta, Rt).

    Returns:
        Derivatives [dI/dt].
    """
    I = y
    d, Rt = args[DELTA_KEY], args[Rt_KEY]

    dI = d(x) * (Rt(x) - 1) * I
    return dI