anypinn.cli.scaffold.gray_scott_2d.ode_csv

Gray-Scott 2D Reaction-Diffusion — inverse PDE problem definition (CSV data).

validation: ValidationRegistry = {DU_KEY: lambda x: torch.full_like(x, TRUE_DU), DV_KEY: lambda x: torch.full_like(x, TRUE_DV), F_KEY: lambda x: torch.full_like(x, TRUE_F), K_KEY: lambda x: torch.full_like(x, TRUE_K)} module-attribute

create_data_module(hp: PINNHyperparameters) -> GrayScott2DDataModule

Source code in src/anypinn/cli/scaffold/gray_scott_2d/ode_csv.py

def create_data_module(hp: PINNHyperparameters) -> GrayScott2DDataModule:
    return GrayScott2DDataModule(
        hp=hp,
        true_du=TRUE_DU,
        true_dv=TRUE_DV,
        true_f=TRUE_F,
        true_k=TRUE_K,
        validation=validation,
    )

create_problem(hp: PINNHyperparameters) -> Problem

Source code in src/anypinn/cli/scaffold/gray_scott_2d/ode_csv.py

def create_problem(hp: PINNHyperparameters) -> Problem:
    encode = FourierEncoding(num_frequencies=6)
    field_u = Field(
        config=MLPConfig(
            in_dim=encode.out_dim(3),
            out_dim=1,
            hidden_layers=hp.fields_config.hidden_layers,
            activation=hp.fields_config.activation,
            output_activation=hp.fields_config.output_activation,
            encode=encode,
        )
    )
    field_v = Field(
        config=MLPConfig(
            in_dim=encode.out_dim(3),
            out_dim=1,
            hidden_layers=hp.fields_config.hidden_layers,
            activation=hp.fields_config.activation,
            output_activation=hp.fields_config.output_activation,
            encode=encode,
        )
    )
    param_du = Parameter(config=ScalarConfig(init_value=hp.params_config.init_value))
    param_dv = Parameter(config=ScalarConfig(init_value=hp.params_config.init_value))
    param_f = Parameter(config=ScalarConfig(init_value=hp.params_config.init_value))
    param_k = Parameter(config=ScalarConfig(init_value=hp.params_config.init_value))

    fields = FieldsRegistry({U_KEY: field_u, V_KEY: field_v})
    params = ParamsRegistry(
        {
            DU_KEY: param_du,
            DV_KEY: param_dv,
            F_KEY: param_f,
            K_KEY: param_k,
        }
    )

    bcs = [
        NeumannBCConstraint(
            BoundaryCondition(sampler=_left_edge, value=_zero, n_pts=100),
            field_u,
            normal_dim=0,
            log_key="loss/bc_left_u",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_left_edge, value=_zero, n_pts=100),
            field_v,
            normal_dim=0,
            log_key="loss/bc_left_v",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_right_edge, value=_zero, n_pts=100),
            field_u,
            normal_dim=0,
            log_key="loss/bc_right_u",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_right_edge, value=_zero, n_pts=100),
            field_v,
            normal_dim=0,
            log_key="loss/bc_right_v",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_bottom_edge, value=_zero, n_pts=100),
            field_u,
            normal_dim=1,
            log_key="loss/bc_bottom_u",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_bottom_edge, value=_zero, n_pts=100),
            field_v,
            normal_dim=1,
            log_key="loss/bc_bottom_v",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_top_edge, value=_zero, n_pts=100),
            field_u,
            normal_dim=1,
            log_key="loss/bc_top_u",
            weight=10.0,
        ),
        NeumannBCConstraint(
            BoundaryCondition(sampler=_top_edge, value=_zero, n_pts=100),
            field_v,
            normal_dim=1,
            log_key="loss/bc_top_v",
            weight=10.0,
        ),
        DirichletBCConstraint(
            BoundaryCondition(sampler=_initial_condition, value=_ic_u, n_pts=200),
            field_u,
            log_key="loss/ic_u",
            weight=10.0,
        ),
        DirichletBCConstraint(
            BoundaryCondition(sampler=_initial_condition, value=_ic_v, n_pts=200),
            field_v,
            log_key="loss/ic_v",
            weight=10.0,
        ),
    ]

    pde_u = PDEResidualConstraint(
        fields=fields,
        params=params,
        residual_fn=gs_residual_u,
        log_key="loss/pde_u",
        weight=1e-4,
    )
    pde_v = PDEResidualConstraint(
        fields=fields,
        params=params,
        residual_fn=gs_residual_v,
        log_key="loss/pde_v",
        weight=1e-4,
    )

    data = DataConstraint(
        fields=fields,
        params=params,
        predict_data=predict_data,
        weight=5.0,
    )

    return Problem(
        constraints=[pde_u, pde_v, *bcs, data],
        criterion=build_criterion(hp.criterion),
        fields=fields,
        params=params,
    )

gs_residual_u(x: Tensor, fields: FieldsRegistry, params: ParamsRegistry) -> Tensor

PDE residual for u: du/dt_norm - T * (D_u lap_u - uv^2 + F(1-u)) = 0.

Source code in src/anypinn/cli/scaffold/gray_scott_2d/ode_csv.py

def gs_residual_u(x: Tensor, fields: FieldsRegistry, params: ParamsRegistry) -> Tensor:
    """PDE residual for u: du/dt_norm - T * (D_u lap_u - uv^2 + F(1-u)) = 0."""
    u = fields[U_KEY](x)
    v = fields[V_KEY](x)
    du = torch.nn.functional.softplus(params[DU_KEY](x))
    f = torch.nn.functional.softplus(params[F_KEY](x))
    du_dt = partial(u, x, dim=2, order=1)
    lap_u = partial(u, x, dim=0, order=2) + partial(u, x, dim=1, order=2)
    return du_dt - (du * lap_u - u * v**2 + f * (1 - u)) * T_TOTAL

gs_residual_v(x: Tensor, fields: FieldsRegistry, params: ParamsRegistry) -> Tensor

PDE residual for v: dv/dt_norm - T * (D_v lap_v + uv^2 - (F+k)v) = 0.

Source code in src/anypinn/cli/scaffold/gray_scott_2d/ode_csv.py

def gs_residual_v(x: Tensor, fields: FieldsRegistry, params: ParamsRegistry) -> Tensor:
    """PDE residual for v: dv/dt_norm - T * (D_v lap_v + uv^2 - (F+k)v) = 0."""
    u = fields[U_KEY](x)
    v = fields[V_KEY](x)
    dv = torch.nn.functional.softplus(params[DV_KEY](x))
    f = torch.nn.functional.softplus(params[F_KEY](x))
    k = torch.nn.functional.softplus(params[K_KEY](x))
    dv_dt = partial(v, x, dim=2, order=1)
    lap_v = partial(v, x, dim=0, order=2) + partial(v, x, dim=1, order=2)
    return dv_dt - (dv * lap_v + u * v**2 - (f + k) * v) * T_TOTAL

predict_data(x_data: Tensor, fields: FieldsRegistry, _params: ParamsRegistry) -> Tensor

Source code in src/anypinn/cli/scaffold/gray_scott_2d/ode_csv.py

def predict_data(x_data: Tensor, fields: FieldsRegistry, _params: ParamsRegistry) -> Tensor:
    u_pred = fields[U_KEY](x_data)
    v_pred = fields[V_KEY](x_data)
    return torch.stack([u_pred, v_pred], dim=1)