Skip to content

anypinn.catalog.lorenz

BETA_KEY = 'beta' module-attribute

RHO_KEY = 'rho' module-attribute

SIGMA_KEY = 'sigma' module-attribute

X_KEY = 'x' module-attribute

Y_KEY = 'y' module-attribute

Z_KEY = 'z' module-attribute

LorenzDataModule

Bases: PINNDataModule

DataModule for Lorenz system inverse problem. Generates synthetic data via odeint.

Source code in src/anypinn/catalog/lorenz.py
class LorenzDataModule(PINNDataModule):
    """DataModule for Lorenz system inverse problem. Generates synthetic data via odeint."""

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

    @override
    def gen_data(self, config: GenerationConfig) -> tuple[Tensor, Tensor]:
        """Generate synthetic Lorenz data using odeint + additive Gaussian noise."""

        def lorenz_ode(t: Tensor, y: Tensor) -> Tensor:
            return self.gen_props.ode(t, y, self.gen_props.args)

        t = config.x

        sol = odeint(lorenz_ode, self.gen_props.y0, t)  # [T, 3]
        x_true = sol[:, 0]
        y_true = sol[:, 1]
        z_true = sol[:, 2]

        x_obs = x_true + self.noise_std * torch.randn_like(x_true)
        y_obs = y_true + self.noise_std * torch.randn_like(y_true)
        z_obs = z_true + self.noise_std * torch.randn_like(z_true)

        y_data = torch.stack([x_obs, y_obs, z_obs], dim=1).unsqueeze(-1)

        return t.unsqueeze(-1), y_data

gen_props = gen_props instance-attribute

noise_std = noise_std instance-attribute

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

Source code in src/anypinn/catalog/lorenz.py
def __init__(
    self,
    hp: ODEHyperparameters,
    gen_props: ODEProperties,
    noise_std: float = 0.0,
    validation: ValidationRegistry | None = None,
    callbacks: Sequence[DataCallback] | None = None,
):
    super().__init__(hp, validation, callbacks)
    self.gen_props = gen_props
    self.noise_std = noise_std

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

Generate synthetic Lorenz data using odeint + additive Gaussian noise.

Source code in src/anypinn/catalog/lorenz.py
@override
def gen_data(self, config: GenerationConfig) -> tuple[Tensor, Tensor]:
    """Generate synthetic Lorenz data using odeint + additive Gaussian noise."""

    def lorenz_ode(t: Tensor, y: Tensor) -> Tensor:
        return self.gen_props.ode(t, y, self.gen_props.args)

    t = config.x

    sol = odeint(lorenz_ode, self.gen_props.y0, t)  # [T, 3]
    x_true = sol[:, 0]
    y_true = sol[:, 1]
    z_true = sol[:, 2]

    x_obs = x_true + self.noise_std * torch.randn_like(x_true)
    y_obs = y_true + self.noise_std * torch.randn_like(y_true)
    z_obs = z_true + self.noise_std * torch.randn_like(z_true)

    y_data = torch.stack([x_obs, y_obs, z_obs], dim=1).unsqueeze(-1)

    return t.unsqueeze(-1), y_data