| """Independent implementation of FluxNet-LAP (arXiv:2602.01941v2). |
| |
| The transport equations and epsilon follow Sections 3.2, 3.3 and 3.5; |
| the backbone dimensions follow Appendix C.2. See reproduction_spec.md |
| for the explicitly identified gaps filled by read-only official references. |
| """ |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class FeatureBlock(nn.Module): |
| """Two periodic convolutions, with a learned concatenation skip fusion.""" |
|
|
| def __init__(self, width, kernel_size): |
| super().__init__() |
| self.layers = nn.Sequential( |
| nn.Conv2d(width, width, kernel_size, padding=kernel_size // 2, |
| padding_mode="circular"), |
| nn.BatchNorm2d(width), nn.GELU(), |
| nn.Conv2d(width, width, kernel_size, padding=kernel_size // 2, |
| padding_mode="circular"), |
| nn.BatchNorm2d(width), nn.GELU(), |
| ) |
| self.fuse = nn.Conv2d(2 * width, width, 1) |
|
|
| def forward(self, features): |
| return self.fuse(torch.cat((self.layers(features), features), dim=1)) |
|
|
|
|
| class FluxNetShallowWater(nn.Module): |
| """Map physical states [B,3,H,W] ordered h,mx,my to the next state. |
| |
| This operator redistributes cumulative amounts, with no dt multiplier, |
| clipping, mass projection, or numerical-solver substeps. |
| """ |
|
|
| def __init__(self, base_channels=64, num_blocks=6, kernel_size=5, |
| neighborhood_size=3, epsilon=1e-6): |
| super().__init__() |
| if kernel_size % 2 != 1 or neighborhood_size % 2 != 1: |
| raise ValueError("Convolution and transport sizes must be odd") |
| if neighborhood_size < 3 or epsilon <= 0: |
| raise ValueError("Invalid transport size or epsilon") |
| self.epsilon = float(epsilon) |
| radius = neighborhood_size // 2 |
| self.offsets = tuple((dy, dx) for dy in range(-radius, radius + 1) |
| for dx in range(-radius, radius + 1) |
| if (dy, dx) != (0, 0)) |
| self.encoder = nn.Sequential( |
| nn.Conv2d(3, base_channels, kernel_size, padding=kernel_size // 2, |
| padding_mode="circular"), |
| nn.BatchNorm2d(base_channels), nn.GELU(), |
| *(FeatureBlock(base_channels, kernel_size) for _ in range(num_blocks)), |
| ) |
| directions = len(self.offsets) |
| self.depth_head = nn.Conv2d(base_channels, directions + 1, 1) |
| self.pressure_x = nn.Conv2d(base_channels, directions, 1) |
| self.pressure_y = nn.Conv2d(base_channels, directions, 1) |
|
|
| def transport(self, state): |
| features = self.encoder(state) |
| h = state[:, :1] |
| raw = self.depth_head(features) |
| outgoing_depth = h * raw[:, :1].sigmoid() |
| depth_transfers = outgoing_depth * raw[:, 1:].softmax(dim=1) |
| velocity = state[:, 1:] / (h + self.epsilon) |
| gate = h.square() |
| momentum_x = depth_transfers * velocity[:, :1] + gate * F.softplus(self.pressure_x(features)) |
| momentum_y = depth_transfers * velocity[:, 1:] + gate * F.softplus(self.pressure_y(features)) |
| transfers = torch.stack((depth_transfers, momentum_x, momentum_y), dim=1) |
| return outgoing_depth, transfers |
|
|
| def forward(self, state): |
| outgoing_depth, transfers = self.transport(state) |
| |
| |
| outgoing = torch.cat((outgoing_depth, transfers[:, 1:].sum(dim=2)), dim=1) |
| updated = state - outgoing |
| for index, offset in enumerate(self.offsets): |
| |
| |
| updated = updated + torch.roll(transfers[:, :, index], offset, dims=(-2, -1)) |
| return updated |
|
|
