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
Example

Left boundary of a 1-D+time domain at x=0

bc_left = BoundaryCondition( ... sampler=lambda n: torch.stack([ ... torch.zeros(n), # x = 0 ... torch.rand(n) * T_max, # t in [0, T] ... ], dim=-1), ... value=lambda coords: torch.zeros(coords.shape[0], 1), ... n_pts=50, ... )

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
43
44
45
46
47
48
49
50
51
52
53
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.

    Example:
        >>> # Left boundary of a 1-D+time domain at x=0
        >>> bc_left = BoundaryCondition(
        ...     sampler=lambda n: torch.stack([
        ...         torch.zeros(n),            # x = 0
        ...         torch.rand(n) * T_max,     # t in [0, T]
        ...     ], dim=-1),
        ...     value=lambda coords: torch.zeros(coords.shape[0], 1),
        ...     n_pts=50,
        ... )
    """

    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 
45
46
47
48
49
50
51
52
53
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 
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
84
85
86
87
88
89
90
91
92
93
94
95
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:
        """Compute the Dirichlet boundary condition loss."""
        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 
68
69
70
71
72
73
74
75
76
77
78
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

Compute the Dirichlet boundary condition loss.

Source code in src/anypinn/problems/pde.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    """Compute the Dirichlet boundary condition loss."""
    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: du/dn(x_bc) = h(x_bc).

For a rectangular domain face whose outward normal is axis-aligned with dimension normal_dim, we have du/dn = du/dx[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 
 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
135
136
137
138
139
140
141
142
143
144
145
146
147
class NeumannBCConstraint(Constraint):
    """
    Enforces the Neumann boundary condition:
    ``du/dn(x_bc) = h(x_bc)``.

    For a rectangular domain face whose outward normal is axis-aligned with
    dimension ``normal_dim``, we have
    ``du/dn = du/dx[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:
        """Compute the Neumann boundary condition loss."""
        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 
117
118
119
120
121
122
123
124
125
126
127
128
129
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

Compute the Neumann boundary condition loss.

Source code in src/anypinn/problems/pde.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    """Compute the Neumann boundary condition loss."""
    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
Example

from anypinn.lib.diff import grad, partial def heat_residual(x, fields, params): ... u = fields["u"] ... u_pred = u(x) ... u_t = partial(u_pred, x, dim=1) # du/dt ... u_x = partial(u_pred, x, dim=0) # du/dx ... u_xx = partial(u_x, x, dim=0) # d2u/dx2 ... alpha = params"alpha" ... return u_t - alpha * u_xx # residual = 0 constraint = PDEResidualConstraint( ... fields=fields, params=params, ... residual_fn=heat_residual, ... )

Source code in src/anypinn/problems/pde.py 
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
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
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.

    Example:
        >>> from anypinn.lib.diff import grad, partial
        >>> def heat_residual(x, fields, params):
        ...     u = fields["u"]
        ...     u_pred = u(x)
        ...     u_t = partial(u_pred, x, dim=1)   # du/dt
        ...     u_x = partial(u_pred, x, dim=0)   # du/dx
        ...     u_xx = partial(u_x, x, dim=0)     # d2u/dx2
        ...     alpha = params["alpha"](x)
        ...     return u_t - alpha * u_xx          # residual = 0
        >>> constraint = PDEResidualConstraint(
        ...     fields=fields, params=params,
        ...     residual_fn=heat_residual,
        ... )
    """

    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:
        """Compute the PDE interior residual loss."""
        _, 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 
249
250
251
252
253
254
255
256
257
258
259
260
261
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

Compute the PDE interior residual loss.

Source code in src/anypinn/problems/pde.py 
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    """Compute the PDE interior residual loss."""
    _, 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 matching spatial derivatives.

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 
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
203
204
205
206
207
208
209
210
211
212
213
214
215
class PeriodicBCConstraint(Constraint):
    """
    Enforces periodic boundary conditions:
    ``u(x_left, t) = u(x_right, t)`` and matching spatial derivatives.

    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:
        """Compute the periodic boundary condition loss."""
        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 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
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

Compute the periodic boundary condition loss.

Source code in src/anypinn/problems/pde.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
@override
def loss(
    self,
    batch: TrainingBatch,
    criterion: nn.Module,
    log: LogFn | None = None,
) -> Tensor:
    """Compute the periodic boundary condition loss."""
    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