Skip to content

anypinn.problems.pde

Boundary condition constraints for PDE problems.

BCValueFn: TypeAlias = Callable[[Tensor], Tensor] module-attribute

A callable that maps boundary coordinates (n_pts, d) → target values (n_pts, out_dim).

PDEResidualFn: TypeAlias = Callable[[Tensor, FieldsRegistry, ParamsRegistry], Tensor] module-attribute

A callable (x, fields, params) → residual tensor, expected to be zero at the solution.

BoundaryCondition

Pairs a boundary region sampler with a prescribed value function.

Parameters:

Name Type Description Default
sampler Callable[[int], Tensor]

Callable (n_pts: int) -> Tensor of shape (n_pts, d). Called each training step to produce fresh boundary sample points.

 required
value BCValueFn

Callable Tensor -> Tensor giving the target value or normal derivative at boundary coordinates.

 required
n_pts int

Number of boundary points sampled per training step.

 100
Source code in src/anypinn/problems/pde.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class BoundaryCondition:
    """
    Pairs a boundary region sampler with a prescribed value function.

    Args:
        sampler: Callable ``(n_pts: int) -> Tensor`` of shape ``(n_pts, d)``.
            Called each training step to produce fresh boundary sample points.
        value: Callable ``Tensor -> Tensor`` giving the target value or normal
            derivative at boundary coordinates.
        n_pts: Number of boundary points sampled per training step.
    """

    def __init__(
        self,
        sampler: Callable[[int], Tensor],
        value: BCValueFn,
        n_pts: int = 100,
    ):
        self.sampler = sampler
        self.value = value
        self.n_pts = n_pts

n_pts = n_pts instance-attribute

sampler = sampler instance-attribute

value = value instance-attribute

__init__(sampler: Callable[[int], Tensor], value: BCValueFn, n_pts: int = 100)

Source code in src/anypinn/problems/pde.py 
34
35
36
37
38
39
40
41
42
def __init__(
    self,
    sampler: Callable[[int], Tensor],
    value: BCValueFn,
    n_pts: int = 100,
):
    self.sampler = sampler
    self.value = value
    self.n_pts = n_pts

DirichletBCConstraint

Bases: Constraint

Enforces the Dirichlet boundary condition: u(x_bc) = g(x_bc). Minimizes weight * criterion(u(x_bc), g(x_bc)).

Parameters:

Name Type Description Default
bc BoundaryCondition

Boundary condition (sampler + target value function).

 required
field Field

The neural field to enforce the condition on.

 required
log_key str

Key used when logging the loss value.

 'loss/bc_dirichlet'
weight float

Loss term weight.

 1.0
Source code in src/anypinn/problems/pde.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class DirichletBCConstraint(Constraint):
    """
    Enforces the Dirichlet boundary condition: u(x_bc) = g(x_bc).
    Minimizes ``weight * criterion(u(x_bc), g(x_bc))``.

    Args:
        bc: Boundary condition (sampler + target value function).
        field: The neural field to enforce the condition on.
        log_key: Key used when logging the loss value.
        weight: Loss term weight.
    """

    def __init__(
        self,
        bc: BoundaryCondition,
        field: Field,
        log_key: str = "loss/bc_dirichlet",
        weight: float = 1.0,
    ):
        self.bc = bc
        self.field = field
        self.log_key = log_key
        self.weight = weight

    @override
    def loss(
        self,
        batch: TrainingBatch,
        criterion: nn.Module,
        log: LogFn | None = None,
    ) -> Tensor:
        device = next(self.field.parameters()).device
        x_bc = self.bc.sampler(self.bc.n_pts).to(device)
        u_pred = self.field(x_bc)
        g = self.bc.value(x_bc).to(device)
        loss: Tensor = self.weight * criterion(u_pred, g)
        if log is not None:
            log(self.log_key, loss)
        return loss

bc = bc instance-attribute

field = field instance-attribute

log_key = log_key instance-attribute

weight = weight instance-attribute

__init__(bc: BoundaryCondition, field: Field, log_key: str = 'loss/bc_dirichlet', weight: float = 1.0)

Source code in src/anypinn/problems/pde.py 
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    bc: BoundaryCondition,
    field: Field,
    log_key: str = "loss/bc_dirichlet",
    weight: float = 1.0,
):
    self.bc = bc
    self.field = field
    self.log_key = log_key
    self.weight = weight

loss(batch: TrainingBatch, criterion: nn.Module, log: LogFn | None = None) -> Tensor

Source code in src/anypinn/problems/pde.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    device = next(self.field.parameters()).device
    x_bc = self.bc.sampler(self.bc.n_pts).to(device)
    u_pred = self.field(x_bc)
    g = self.bc.value(x_bc).to(device)
    loss: Tensor = self.weight * criterion(u_pred, g)
    if log is not None:
        log(self.log_key, loss)
    return loss

NeumannBCConstraint

Bases: Constraint

Enforces the Neumann boundary condition: \(\partial u / \partial n (x_{bc}) = h(x_{bc})\).

For a rectangular domain face whose outward normal is axis-aligned with dimension normal_dim, we have \(\partial u / \partial n = \partial u / \partial x_{\mathrm{normal\_dim}}\). Minimizes weight * criterion(du_dn(x_bc), h(x_bc)).

Parameters:

Name Type Description Default
bc BoundaryCondition

Boundary condition (sampler + target normal-derivative function).

 required
field Field

The neural field to enforce the condition on.

 required
normal_dim int

Index of the spatial dimension the boundary normal points along.

 required
log_key str

Key used when logging the loss value.

 'loss/bc_neumann'
weight float

Loss term weight.

 1.0
Source code in src/anypinn/problems/pde.py 
 86
 87
 88
 89
 90
 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
124
125
126
127
128
129
130
131
132
133
134
class NeumannBCConstraint(Constraint):
    """
    Enforces the Neumann boundary condition:
    $\\partial u / \\partial n (x_{bc}) = h(x_{bc})$.

    For a rectangular domain face whose outward normal is axis-aligned with
    dimension ``normal_dim``, we have
    $\\partial u / \\partial n = \\partial u / \\partial x_{\\mathrm{normal\\_dim}}$.
    Minimizes
    ``weight * criterion(du_dn(x_bc), h(x_bc))``.

    Args:
        bc: Boundary condition (sampler + target normal-derivative function).
        field: The neural field to enforce the condition on.
        normal_dim: Index of the spatial dimension the boundary normal points along.
        log_key: Key used when logging the loss value.
        weight: Loss term weight.
    """

    def __init__(
        self,
        bc: BoundaryCondition,
        field: Field,
        normal_dim: int,
        log_key: str = "loss/bc_neumann",
        weight: float = 1.0,
    ):
        self.bc = bc
        self.field = field
        self.normal_dim = normal_dim
        self.log_key = log_key
        self.weight = weight

    @override
    def loss(
        self,
        batch: TrainingBatch,
        criterion: nn.Module,
        log: LogFn | None = None,
    ) -> Tensor:
        device = next(self.field.parameters()).device
        x_bc = self.bc.sampler(self.bc.n_pts).to(device).detach().requires_grad_(True)
        u_pred = self.field(x_bc)
        du_dn = diff_partial(u_pred, x_bc, dim=self.normal_dim)
        h = self.bc.value(x_bc.detach()).to(device)
        loss: Tensor = self.weight * criterion(du_dn, h)
        if log is not None:
            log(self.log_key, loss)
        return loss

bc = bc instance-attribute

field = field instance-attribute

log_key = log_key instance-attribute

normal_dim = normal_dim instance-attribute

weight = weight instance-attribute

__init__(bc: BoundaryCondition, field: Field, normal_dim: int, log_key: str = 'loss/bc_neumann', weight: float = 1.0)

Source code in src/anypinn/problems/pde.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
def __init__(
    self,
    bc: BoundaryCondition,
    field: Field,
    normal_dim: int,
    log_key: str = "loss/bc_neumann",
    weight: float = 1.0,
):
    self.bc = bc
    self.field = field
    self.normal_dim = normal_dim
    self.log_key = log_key
    self.weight = weight

loss(batch: TrainingBatch, criterion: nn.Module, log: LogFn | None = None) -> Tensor

Source code in src/anypinn/problems/pde.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    device = next(self.field.parameters()).device
    x_bc = self.bc.sampler(self.bc.n_pts).to(device).detach().requires_grad_(True)
    u_pred = self.field(x_bc)
    du_dn = diff_partial(u_pred, x_bc, dim=self.normal_dim)
    h = self.bc.value(x_bc.detach()).to(device)
    loss: Tensor = self.weight * criterion(du_dn, h)
    if log is not None:
        log(self.log_key, loss)
    return loss

PDEResidualConstraint

Bases: Constraint

Enforces a PDE interior residual: residual_fn(x, fields, params) ≈ 0. Minimizes weight * criterion(residual_fn(x_coll, fields, params), 0).

Parameters:

Name Type Description Default
fields FieldsRegistry

Registry of neural fields the residual function operates on. Pass only the subset needed — other fields in the Problem are ignored.

 required
params ParamsRegistry

Registry of parameters the residual function uses.

 required
residual_fn PDEResidualFn

Callable (x, fields, params) → Tensor of residuals. Should use anypinn.lib.diff operators for derivatives. The returned tensor is compared against zeros.

 required
log_key str

Key used when logging the loss value.

 'loss/pde_residual'
weight float

Loss term weight.

 1.0
Source code in src/anypinn/problems/pde.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
class PDEResidualConstraint(Constraint):
    """
    Enforces a PDE interior residual: ``residual_fn(x, fields, params) ≈ 0``.
    Minimizes ``weight * criterion(residual_fn(x_coll, fields, params), 0)``.

    Args:
        fields: Registry of neural fields the residual function operates on.
            Pass only the subset needed — other fields in the Problem are ignored.
        params: Registry of parameters the residual function uses.
        residual_fn: Callable (x, fields, params) → Tensor of residuals.
            Should use ``anypinn.lib.diff`` operators for derivatives.
            The returned tensor is compared against zeros.
        log_key: Key used when logging the loss value.
        weight: Loss term weight.
    """

    def __init__(
        self,
        fields: FieldsRegistry,
        params: ParamsRegistry,
        residual_fn: PDEResidualFn,
        log_key: str = "loss/pde_residual",
        weight: float = 1.0,
    ):
        self.fields = fields
        self.params = params
        self.residual_fn = residual_fn
        self.log_key = log_key
        self.weight = weight

    @override
    def loss(
        self,
        batch: TrainingBatch,
        criterion: nn.Module,
        log: LogFn | None = None,
    ) -> Tensor:
        _, x_coll = batch
        x_coll = x_coll.detach().requires_grad_(True)
        residual = self.residual_fn(x_coll, self.fields, self.params)
        loss: Tensor = self.weight * criterion(residual, torch.zeros_like(residual))
        if log is not None:
            log(self.log_key, loss)
        return loss

fields = fields instance-attribute

log_key = log_key instance-attribute

params = params instance-attribute

residual_fn = residual_fn instance-attribute

weight = weight instance-attribute

__init__(fields: FieldsRegistry, params: ParamsRegistry, residual_fn: PDEResidualFn, log_key: str = 'loss/pde_residual', weight: float = 1.0)

Source code in src/anypinn/problems/pde.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
def __init__(
    self,
    fields: FieldsRegistry,
    params: ParamsRegistry,
    residual_fn: PDEResidualFn,
    log_key: str = "loss/pde_residual",
    weight: float = 1.0,
):
    self.fields = fields
    self.params = params
    self.residual_fn = residual_fn
    self.log_key = log_key
    self.weight = weight

loss(batch: TrainingBatch, criterion: nn.Module, log: LogFn | None = None) -> Tensor

Source code in src/anypinn/problems/pde.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    _, x_coll = batch
    x_coll = x_coll.detach().requires_grad_(True)
    residual = self.residual_fn(x_coll, self.fields, self.params)
    loss: Tensor = self.weight * criterion(residual, torch.zeros_like(residual))
    if log is not None:
        log(self.log_key, loss)
    return loss

PeriodicBCConstraint

Bases: Constraint

Enforces periodic boundary conditions: u(x_left, t) = u(x_right, t) and ∂u/∂x(x_left, t) = ∂u/∂x(x_right, t).

The two boundary samplers must produce paired points — identical coordinates in every dimension except the periodic one — so that the value- and derivative-matching losses are meaningful.

Minimizes weight * [criterion(u_left, u_right) + criterion(du_left, du_right)].

Parameters:

Name Type Description Default
bc_left BoundaryCondition

Left boundary sampler (sampler + dummy value function).

 required
bc_right BoundaryCondition

Right boundary sampler (sampler + dummy value function).

 required
field Field

The neural field to enforce the condition on.

 required
match_dim int

Spatial dimension index for the derivative matching.

 0
log_key str

Key used when logging the loss value.

 'loss/bc_periodic'
weight float

Loss term weight.

 1.0
Source code in src/anypinn/problems/pde.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class PeriodicBCConstraint(Constraint):
    """
    Enforces periodic boundary conditions:
    ``u(x_left, t) = u(x_right, t)`` and
    ``∂u/∂x(x_left, t) = ∂u/∂x(x_right, t)``.

    The two boundary samplers must produce **paired** points — identical
    coordinates in every dimension except the periodic one — so that
    the value- and derivative-matching losses are meaningful.

    Minimizes
    ``weight * [criterion(u_left, u_right) + criterion(du_left, du_right)]``.

    Args:
        bc_left: Left boundary sampler (sampler + dummy value function).
        bc_right: Right boundary sampler (sampler + dummy value function).
        field: The neural field to enforce the condition on.
        match_dim: Spatial dimension index for the derivative matching.
        log_key: Key used when logging the loss value.
        weight: Loss term weight.
    """

    def __init__(
        self,
        bc_left: BoundaryCondition,
        bc_right: BoundaryCondition,
        field: Field,
        match_dim: int = 0,
        log_key: str = "loss/bc_periodic",
        weight: float = 1.0,
    ):
        self.bc_left = bc_left
        self.bc_right = bc_right
        self.field = field
        self.match_dim = match_dim
        self.log_key = log_key
        self.weight = weight

    @override
    def loss(
        self,
        batch: TrainingBatch,
        criterion: nn.Module,
        log: LogFn | None = None,
    ) -> Tensor:
        device = next(self.field.parameters()).device
        n_pts = self.bc_left.n_pts

        x_left = self.bc_left.sampler(n_pts).to(device).detach().requires_grad_(True)
        x_right = self.bc_right.sampler(n_pts).to(device).detach().requires_grad_(True)

        u_left = self.field(x_left)
        u_right = self.field(x_right)

        # Value matching: u(x_left, t) = u(x_right, t)
        loss_val: Tensor = criterion(u_left, u_right)

        # Derivative matching: du/dx(x_left, t) = du/dx(x_right, t)
        du_left = diff_partial(u_left, x_left, dim=self.match_dim)
        du_right = diff_partial(u_right, x_right, dim=self.match_dim)
        loss_deriv: Tensor = criterion(du_left, du_right)

        loss: Tensor = self.weight * (loss_val + loss_deriv)
        if log is not None:
            log(self.log_key, loss)
        return loss

bc_left = bc_left instance-attribute

bc_right = bc_right instance-attribute

field = field instance-attribute

log_key = log_key instance-attribute

match_dim = match_dim instance-attribute

weight = weight instance-attribute

__init__(bc_left: BoundaryCondition, bc_right: BoundaryCondition, field: Field, match_dim: int = 0, log_key: str = 'loss/bc_periodic', weight: float = 1.0)

Source code in src/anypinn/problems/pde.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def __init__(
    self,
    bc_left: BoundaryCondition,
    bc_right: BoundaryCondition,
    field: Field,
    match_dim: int = 0,
    log_key: str = "loss/bc_periodic",
    weight: float = 1.0,
):
    self.bc_left = bc_left
    self.bc_right = bc_right
    self.field = field
    self.match_dim = match_dim
    self.log_key = log_key
    self.weight = weight

loss(batch: TrainingBatch, criterion: nn.Module, log: LogFn | None = None) -> Tensor

Source code in src/anypinn/problems/pde.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    device = next(self.field.parameters()).device
    n_pts = self.bc_left.n_pts

    x_left = self.bc_left.sampler(n_pts).to(device).detach().requires_grad_(True)
    x_right = self.bc_right.sampler(n_pts).to(device).detach().requires_grad_(True)

    u_left = self.field(x_left)
    u_right = self.field(x_right)

    # Value matching: u(x_left, t) = u(x_right, t)
    loss_val: Tensor = criterion(u_left, u_right)

    # Derivative matching: du/dx(x_left, t) = du/dx(x_right, t)
    du_left = diff_partial(u_left, x_left, dim=self.match_dim)
    du_right = diff_partial(u_right, x_right, dim=self.match_dim)
    loss_deriv: Tensor = criterion(du_left, du_right)

    loss: Tensor = self.weight * (loss_val + loss_deriv)
    if log is not None:
        log(self.log_key, loss)
    return loss