| """Periodic one-dimensional CNN for mass-aware data-assimilation correction.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class PeriodicConv1d(nn.Module): |
| """Conv1d with explicit circular padding and unchanged spatial length.""" |
|
|
| def __init__(self, in_channels: int, out_channels: int, kernel_size: int): |
| super().__init__() |
| if kernel_size % 2 != 1: |
| raise ValueError("kernel_size must be odd") |
| self.pad = kernel_size // 2 |
| self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=0) |
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| return self.conv(F.pad(inputs, (self.pad, self.pad), mode="circular")) |
|
|
|
|
| class MassConservingCNN(nn.Module): |
| """Four hidden SELU convolutions followed by the u/h/r output layer.""" |
|
|
| def __init__(self, input_channels: int = 4, hidden_channels: int = 32, |
| hidden_layers: int = 4, kernel_size: int = 3): |
| super().__init__() |
| if input_channels != 4 or hidden_layers != 4 or kernel_size != 3: |
| raise ValueError("paper architecture requires 4 inputs, 4 hidden layers, kernel size 3") |
| layers = [] |
| channels = input_channels |
| for _ in range(hidden_layers): |
| layers.extend((PeriodicConv1d(channels, hidden_channels, kernel_size), nn.SELU())) |
| channels = hidden_channels |
| self.hidden = nn.Sequential(*layers) |
| self.output = PeriodicConv1d(hidden_channels, 3, kernel_size) |
|
|
| @property |
| def influence_radius(self) -> int: |
| return 5 |
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| if inputs.ndim != 3 or inputs.shape[1] != 4 or inputs.shape[2] != 250: |
| raise ValueError(f"expected float tensor [B,4,250], got {tuple(inputs.shape)}") |
| if not inputs.is_floating_point(): |
| raise TypeError("inputs must have a floating-point dtype") |
| raw = self.output(self.hidden(inputs)) |
| return torch.cat((raw[:, :2], F.relu(raw[:, 2:3])), dim=1) |
|
|
