OneScience-Group/MassConservingCNN

1"""Periodic one-dimensional CNN for mass-aware data-assimilation correction."""2 3from __future__ import annotations4 5import torch6from torch import nn7from torch.nn import functional as F8 9 10class PeriodicConv1d(nn.Module):11    """Conv1d with explicit circular padding and unchanged spatial length."""12 13    def __init__(self, in_channels: int, out_channels: int, kernel_size: int):14        super().__init__()15        if kernel_size % 2 != 1:16            raise ValueError("kernel_size must be odd")17        self.pad = kernel_size // 218        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=0)19 20    def forward(self, inputs: torch.Tensor) -> torch.Tensor:21        return self.conv(F.pad(inputs, (self.pad, self.pad), mode="circular"))22 23 24class MassConservingCNN(nn.Module):25    """Four hidden SELU convolutions followed by the u/h/r output layer."""26 27    def __init__(self, input_channels: int = 4, hidden_channels: int = 32,28                 hidden_layers: int = 4, kernel_size: int = 3):29        super().__init__()30        if input_channels != 4 or hidden_layers != 4 or kernel_size != 3:31            raise ValueError("paper architecture requires 4 inputs, 4 hidden layers, kernel size 3")32        layers = []33        channels = input_channels34        for _ in range(hidden_layers):35            layers.extend((PeriodicConv1d(channels, hidden_channels, kernel_size), nn.SELU()))36            channels = hidden_channels37        self.hidden = nn.Sequential(*layers)38        self.output = PeriodicConv1d(hidden_channels, 3, kernel_size)39 40    @property41    def influence_radius(self) -> int:42        return 543 44    def forward(self, inputs: torch.Tensor) -> torch.Tensor:45        if inputs.ndim != 3 or inputs.shape[1] != 4 or inputs.shape[2] != 250:46            raise ValueError(f"expected float tensor [B,4,250], got {tuple(inputs.shape)}")47        if not inputs.is_floating_point():48            raise TypeError("inputs must have a floating-point dtype")49        raw = self.output(self.hidden(inputs))50        return torch.cat((raw[:, :2], F.relu(raw[:, 2:3])), dim=1)51