| """Independent PyTorch implementation of Appendix C.1, equations C.1--C.7. |
| |
| The jet path is forward automatic differentiation carrying value, dx, dy and |
| Laplacian. Parameter derivatives are reverse AD through these exact operations; |
| no finite differences, finite-volume targets or reference values enter the PDE. |
| """ |
| import math |
| from typing import Tuple |
|
|
| import torch |
| from torch import Tensor, nn |
| import torch.nn.functional as F |
|
|
|
|
| def jet_swish(z: Tensor) -> Tensor: |
| a, ax, ay, al = z.unbind(0) |
| s = torch.sigmoid(a) |
| sp = s * (1.0 - s) |
| first = s + a * sp |
| second = sp * (2.0 + a * (1.0 - 2.0 * s)) |
| return torch.stack((a * s, first * ax, first * ay, |
| first * al + second * (ax.square() + ay.square()))) |
|
|
|
|
| def jet_product(a: Tensor, b: Tensor) -> Tensor: |
| return torch.stack((a[0] * b[0], a[1] * b[0] + a[0] * b[1], |
| a[2] * b[0] + a[0] * b[2], |
| a[3] * b[0] + a[0] * b[3] |
| + 2.0 * (a[1] * b[1] + a[2] * b[2]))) |
|
|
|
|
| class Dense(nn.Module): |
| """LeCun truncated-normal initialization; layout is [in, out].""" |
| def __init__(self, nin: int, nout: int, bias: bool = True): |
| super().__init__() |
| self.weight = nn.Parameter(torch.empty(nin, nout)) |
| sigma = (1.0 / math.sqrt(nin)) / 0.87962566103423978 |
| nn.init.trunc_normal_(self.weight, std=sigma, a=-2 * sigma, b=2 * sigma) |
| self.bias = nn.Parameter(torch.zeros(nout)) if bias else None |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| y = x @ self.weight |
| if self.bias is not None: |
| y = y + self.bias |
| return y |
|
|
| @torch.jit.export |
| def jet(self, x: Tensor) -> Tensor: |
| z = x @ self.weight |
| if self.bias is not None: |
| z = torch.stack((z[0] + self.bias, z[1], z[2], z[3])) |
| return z |
|
|
|
|
| class ResidualBlock(nn.Module): |
| def __init__(self, width: int): |
| super().__init__() |
| self.first = Dense(width, width) |
| self.second = Dense(width, width) |
| self.third = Dense(width, width) |
| self.alpha = nn.Parameter(torch.zeros(1)) |
|
|
| def forward(self, h: Tensor, u: Tensor, v: Tensor) -> Tensor: |
| f = F.silu(self.first(h)) |
| z = v + f * (u - v) |
| g = F.silu(self.second(z)) |
| z = v + g * (u - v) |
| q = F.silu(self.third(z)) |
| return self.alpha * q + (1.0 - self.alpha) * h |
|
|
| @torch.jit.export |
| def jet(self, h: Tensor, u: Tensor, v: Tensor) -> Tensor: |
| f = jet_swish(self.first.jet(h)) |
| z = v + jet_product(f, u - v) |
| g = jet_swish(self.second.jet(z)) |
| z = v + jet_product(g, u - v) |
| q = jet_swish(self.third.jet(z)) |
| return self.alpha * q + (1.0 - self.alpha) * h |
|
|
|
|
| class PseudoTimePINN(nn.Module): |
| """BFS PirateNet mapping physical (x,y) to (u,v,p).""" |
| def __init__(self, width: int = 256, depth: int = 3, |
| rff_scale: float = 2.0, rff_trainable: bool = True, |
| activation: str = "swish", output_bias: bool = False, |
| length: float = 15.0, height: float = 1.0): |
| super().__init__() |
| if width % 2 or activation != "swish": |
| raise ValueError("Table 4 BFS requires an even RFF width and Swish") |
| self.register_buffer("coordinate_scale", torch.tensor([length, height])) |
| self.frequencies = nn.Parameter(torch.randn(2, width // 2) * rff_scale, |
| requires_grad=rff_trainable) |
| self.gate_u = Dense(width, width) |
| self.gate_v = Dense(width, width) |
| self.blocks = nn.ModuleList([ResidualBlock(width) for _ in range(depth)]) |
| self.readout = Dense(width, 3, output_bias) |
|
|
| def forward(self, xy: Tensor) -> Tensor: |
| phase = (xy / self.coordinate_scale) @ self.frequencies |
| h = torch.cat((torch.cos(phase), torch.sin(phase)), dim=-1) |
| u, v = F.silu(self.gate_u(h)), F.silu(self.gate_v(h)) |
| for block in self.blocks: |
| h = block(h, u, v) |
| return self.readout(h) |
|
|
| @torch.jit.export |
| def derivatives(self, xy: Tensor) -> Tensor: |
| b = self.frequencies / self.coordinate_scale[:, None] |
| phase = xy @ b |
| cs, sn = torch.cos(phase), torch.sin(phase) |
| h0 = torch.cat((cs, sn), -1) |
| hx = torch.cat((-sn * b[0], cs * b[0]), -1) |
| hy = torch.cat((-sn * b[1], cs * b[1]), -1) |
| bnorm = b.square().sum(0) |
| hl = torch.cat((-cs * bnorm, -sn * bnorm), -1) |
| h = torch.stack((h0, hx, hy, hl)) |
| u, v = jet_swish(self.gate_u.jet(h)), jet_swish(self.gate_v.jet(h)) |
| for block in self.blocks: |
| h = block.jet(h, u, v) |
| return self.readout.jet(h) |
|
|
| @torch.jit.export |
| def physics(self, xy: Tensor, viscosity: float) -> Tuple[Tensor, Tensor]: |
| d = self.derivatives(xy) |
| value, dx, dy, lap = d.unbind(0) |
| ru = value[:, 0] * dx[:, 0] + value[:, 1] * dy[:, 0] + dx[:, 2] - viscosity * lap[:, 0] |
| rv = value[:, 0] * dx[:, 1] + value[:, 1] * dy[:, 1] + dy[:, 2] - viscosity * lap[:, 1] |
| rc = dx[:, 0] + dy[:, 1] |
| return value, torch.stack((ru, rv, rc), dim=-1) |
|
|
