model/spherical_dyffusion.py

"""Compact trainable model used by the local Spherical DYffusion pipeline."""

from torch import nn


class SphericalDYffusion(nn.Module):
    """Small convolutional forecaster preserving the global-grid tensor contract."""

    def __init__(self, channels: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(channels, 32, 3, padding=1),
            nn.GELU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.GELU(),
            nn.Conv2d(32, channels, 1),
        )

    def forward(self, inputs):
        return self.net(inputs)