Bases: Module
(x,t) -> u(x,t) with boundary conditioning for homogeneous Dirichlet on the unit cube.
The network output is multiplied by
Φ(x) = Π_i 4 x_i (1 - x_i)
which vanishes on every face of [0,1]^d for any network weights. The boundary
condition therefore holds exactly, by construction, and contributes no term to the
loss: there is nothing to trade off against the PDE residual.
The factor of 4 per dimension normalises the peak to 1. Without it Φ peaks at 4^-d --
0.25 in 1D but 0.0625 in 2D and 0.0156 in 3D -- so the network would have to grow
like 4^d to represent an O(1) solution, with gradients reaching it attenuated by the
same factor. Measured on the 2D heat equation, the unnormalised form plateaus at 55%
relative error where the normalised one reaches 5%; in 1D the two are equivalent,
since a constant factor is absorbed by the weights.
in_dim counts the spatial coordinates plus time, so in_dim=2 is the 1D
problem (x, t) and in_dim=3 is the 2D problem (x, y, t).
Source code in src\pinn_rk\model.py
| class MLP(nn.Module):
"""
(x,t) -> u(x,t) with boundary conditioning for homogeneous Dirichlet on the unit cube.
The network output is multiplied by
Φ(x) = Π_i 4 x_i (1 - x_i)
which vanishes on every face of [0,1]^d for any network weights. The boundary
condition therefore holds *exactly*, by construction, and contributes no term to the
loss: there is nothing to trade off against the PDE residual.
The factor of 4 per dimension normalises the peak to 1. Without it Φ peaks at 4^-d --
0.25 in 1D but 0.0625 in 2D and 0.0156 in 3D -- so the network would have to grow
like 4^d to represent an O(1) solution, with gradients reaching it attenuated by the
same factor. Measured on the 2D heat equation, the unnormalised form plateaus at 55%
relative error where the normalised one reaches 5%; in 1D the two are equivalent,
since a constant factor is absorbed by the weights.
``in_dim`` counts the spatial coordinates plus time, so ``in_dim=2`` is the 1D
problem (x, t) and ``in_dim=3`` is the 2D problem (x, y, t).
"""
def __init__(
self,
in_dim: int = 2,
width: int = 128,
depth: int = 4,
activation: str = "tanh",
dtype: torch.dtype | None = None,
) -> None:
super().__init__()
if in_dim < 2:
raise ValueError("MLP expects in_dim >= 2: d spatial coordinates plus time.")
self.space_dim = in_dim - 1
# RK-PINN residuals are ill-conditioned in float32; callers should pass
# torch.float64 explicitly rather than rely on the global default dtype.
dtype = torch.get_default_dtype() if dtype is None else dtype
act_cls = nn.Tanh if activation == "tanh" else nn.SiLU
layers: list[nn.Module] = []
d = in_dim
for _ in range(depth):
layers += [nn.Linear(d, width, dtype=dtype), act_cls()]
d = width
layers += [nn.Linear(d, 1, dtype=dtype)]
self.net = nn.Sequential(*layers)
def forward(self, x: Tensor, t: Tensor) -> Tensor:
"""
Evaluate u(x,t).
Parameters
----------
x : Tensor [B, d] spatial points
t : Tensor [B, 1] times. Time is a single coordinate whatever d is, so this
stays [B,1] rather than matching the shape of ``x``.
Returns
-------
Tensor [B, 1]
"""
if x.ndim != 2 or x.shape[1] != self.space_dim:
raise ValueError(
f"x must be [B,{self.space_dim}] for in_dim={self.space_dim + 1}; got "
f"{tuple(x.shape)}."
)
if t.ndim != 2 or t.shape[1] != 1 or t.shape[0] != x.shape[0]:
raise ValueError(f"t must be [B,1] matching x's batch; got {tuple(t.shape)}.")
xt = torch.cat([x, t], dim=1)
g = self.net(xt)
# Π_i 4 x_i(1-x_i): zero on every face of the unit cube, so homogeneous Dirichlet
# conditions hold identically, and peaking at 1 in any dimension so the network
# does not have to compensate for a 4^-d factor.
phi = (4.0 * x * (1.0 - x)).prod(dim=1, keepdim=True)
return cast(Tensor, phi * g)
|
Methods:
forward
forward(x: Tensor, t: Tensor) -> Tensor
Evaluate u(x,t).
Parameters:
| Name |
Type |
Description |
Default |
x
|
Tensor [B, d] spatial points
|
|
required
|
t
|
Tensor [B, 1] times. Time is a single coordinate whatever d is, so this
|
stays [B,1] rather than matching the shape of x.
|
required
|
Returns:
| Type |
Description |
Tensor[B, 1]
|
|
Source code in src\pinn_rk\model.py
| def forward(self, x: Tensor, t: Tensor) -> Tensor:
"""
Evaluate u(x,t).
Parameters
----------
x : Tensor [B, d] spatial points
t : Tensor [B, 1] times. Time is a single coordinate whatever d is, so this
stays [B,1] rather than matching the shape of ``x``.
Returns
-------
Tensor [B, 1]
"""
if x.ndim != 2 or x.shape[1] != self.space_dim:
raise ValueError(
f"x must be [B,{self.space_dim}] for in_dim={self.space_dim + 1}; got "
f"{tuple(x.shape)}."
)
if t.ndim != 2 or t.shape[1] != 1 or t.shape[0] != x.shape[0]:
raise ValueError(f"t must be [B,1] matching x's batch; got {tuple(t.shape)}.")
xt = torch.cat([x, t], dim=1)
g = self.net(xt)
# Π_i 4 x_i(1-x_i): zero on every face of the unit cube, so homogeneous Dirichlet
# conditions hold identically, and peaking at 1 in any dimension so the network
# does not have to compensate for a 4^-d factor.
phi = (4.0 * x * (1.0 - x)).prod(dim=1, keepdim=True)
return cast(Tensor, phi * g)
|