anypinn.cli.scaffold.allen_cahn.ode

Allen-Cahn Equation — stiff reaction-diffusion problem definition.

EPSILON = TRUE_EPSILON module-attribute

GRID_SIZE = 50 module-attribute

AllenCahnResidualScorer

ResidualScorer protocol implementation for adaptive collocation.

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

class AllenCahnResidualScorer:
    """ResidualScorer protocol implementation for adaptive collocation."""

    def __init__(self, fields: FieldsRegistry, params: ParamsRegistry) -> None:
        self.fields = fields
        self.params = params

    def residual_score(self, x: Tensor) -> Tensor:
        device = next(iter(self.fields.values())).parameters().__next__().device
        x = x.detach().to(device).requires_grad_(True)
        with torch.enable_grad():
            res = allen_cahn_residual_forward(x, self.fields, self.params)
        return res.detach().cpu().abs().mean(dim=-1)

fields = fields instance-attribute

params = params instance-attribute

__init__(fields: FieldsRegistry, params: ParamsRegistry) -> None

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def __init__(self, fields: FieldsRegistry, params: ParamsRegistry) -> None:
    self.fields = fields
    self.params = params

residual_score(x: Tensor) -> Tensor

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def residual_score(self, x: Tensor) -> Tensor:
    device = next(iter(self.fields.values())).parameters().__next__().device
    x = x.detach().to(device).requires_grad_(True)
    with torch.enable_grad():
        res = allen_cahn_residual_forward(x, self.fields, self.params)
    return res.detach().cpu().abs().mean(dim=-1)

allen_cahn_residual_forward(x: Tensor, fields: FieldsRegistry, _params: ParamsRegistry) -> Tensor

PDE residual: du/dt - eps*d2u/dx2 - u + u^3 = 0 (eps known).

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def allen_cahn_residual_forward(
    x: Tensor, fields: FieldsRegistry, _params: ParamsRegistry
) -> Tensor:
    """PDE residual: du/dt - eps*d2u/dx2 - u + u^3 = 0 (eps known)."""
    u = fields[U_KEY](x)
    du_dt = partial(u, x, dim=1, order=1)
    d2u_dx2 = partial(u, x, dim=0, order=2)
    return du_dt - EPSILON * d2u_dx2 - u + u**3

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

PDE residual: du/dt - eps*d2u/dx2 - u + u^3 = 0 (eps learned).

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def allen_cahn_residual_inverse(
    x: Tensor, fields: FieldsRegistry, params: ParamsRegistry
) -> Tensor:
    """PDE residual: du/dt - eps*d2u/dx2 - u + u^3 = 0 (eps learned)."""
    u = fields[U_KEY](x)
    eps = params[EPSILON_KEY](x)
    du_dt = partial(u, x, dim=1, order=1)
    d2u_dx2 = partial(u, x, dim=0, order=2)
    return du_dt - eps * d2u_dx2 - u + u**3

create_data_module_csv(hp: PINNHyperparameters) -> AllenCahnDataModule

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def create_data_module_csv(hp: PINNHyperparameters) -> AllenCahnDataModule:
    # To enable adaptive collocation, pass residual_scorer=AllenCahnResidualScorer(fields, params)
    return AllenCahnDataModule(
        hp=hp,
        true_epsilon=TRUE_EPSILON,
        grid_size=GRID_SIZE,
        validation=_validation,
    )

create_data_module_synthetic(hp: PINNHyperparameters) -> AllenCahnDataModule

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def create_data_module_synthetic(hp: PINNHyperparameters) -> AllenCahnDataModule:
    # To enable adaptive collocation, pass residual_scorer=AllenCahnResidualScorer(fields, params)
    return AllenCahnDataModule(
        hp=hp,
        true_epsilon=TRUE_EPSILON,
        grid_size=GRID_SIZE,
        validation=_validation,
    )

create_problem_forward(hp: PINNHyperparameters) -> Problem

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def create_problem_forward(hp: PINNHyperparameters) -> Problem:
    rff = RandomFourierFeatures(in_dim=2, num_features=128, scale=5.0, seed=42)
    field_u = Field(
        config=MLPConfig(
            in_dim=rff.out_dim,
            out_dim=1,
            hidden_layers=hp.fields_config.hidden_layers,
            activation=hp.fields_config.activation,
            output_activation=hp.fields_config.output_activation,
            encode=rff,
        )
    )

    fields = FieldsRegistry({U_KEY: field_u})
    params = ParamsRegistry({})

    bcs = [
        PeriodicBCConstraint(
            bc_left=BoundaryCondition(sampler=_periodic_left, value=_dummy, n_pts=100),
            bc_right=BoundaryCondition(sampler=_periodic_right, value=_dummy, n_pts=100),
            field=field_u,
            match_dim=0,
            log_key="loss/bc_periodic",
            weight=10.0,
        ),
        DirichletBCConstraint(
            BoundaryCondition(sampler=_initial_condition, value=_ic_value, n_pts=200),
            field_u,
            log_key="loss/ic",
            weight=10.0,
        ),
    ]

    pde = PDEResidualConstraint(
        fields=fields,
        params=params,
        residual_fn=allen_cahn_residual_forward,
        log_key="loss/pde_residual",
        weight=1.0,
    )

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

create_problem_inverse(hp: PINNHyperparameters) -> Problem

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

def create_problem_inverse(hp: PINNHyperparameters) -> Problem:
    rff = RandomFourierFeatures(in_dim=2, num_features=128, scale=5.0, seed=42)
    field_u = Field(
        config=MLPConfig(
            in_dim=rff.out_dim,
            out_dim=1,
            hidden_layers=hp.fields_config.hidden_layers,
            activation=hp.fields_config.activation,
            output_activation=hp.fields_config.output_activation,
            encode=rff,
        )
    )
    param_epsilon = Parameter(
        config=ScalarConfig(init_value=hp.params_config.init_value),
    )

    fields = FieldsRegistry({U_KEY: field_u})
    params = ParamsRegistry({EPSILON_KEY: param_epsilon})

    bcs = [
        PeriodicBCConstraint(
            bc_left=BoundaryCondition(sampler=_periodic_left, value=_dummy, n_pts=100),
            bc_right=BoundaryCondition(sampler=_periodic_right, value=_dummy, n_pts=100),
            field=field_u,
            match_dim=0,
            log_key="loss/bc_periodic",
            weight=10.0,
        ),
        DirichletBCConstraint(
            BoundaryCondition(sampler=_initial_condition, value=_ic_value, n_pts=200),
            field_u,
            log_key="loss/ic",
            weight=10.0,
        ),
    ]

    pde = PDEResidualConstraint(
        fields=fields,
        params=params,
        residual_fn=allen_cahn_residual_inverse,
        log_key="loss/pde_residual",
        weight=1.0,
    )

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

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

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

Source code in src/anypinn/cli/scaffold/allen_cahn/ode.py

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