File size: 2,085 Bytes
0fa8141 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 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 | """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)
|