File size: 1,539 Bytes
73d9e73 | 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 | """Thin validation wrapper around the official torch-harmonics SFNO class."""
from __future__ import annotations
import torch
from torch import nn
from torch_harmonics.examples.models.sfno import SphericalFourierNeuralOperator
from .config import SFNOConfig
class OfficialSFNOAdapter(nn.Module):
def __init__(self, config: SFNOConfig) -> None:
super().__init__()
config.validate()
self.expected_shape = (config.channels, config.nlat, config.nlon)
self.model = SphericalFourierNeuralOperator(
img_size=(config.nlat, config.nlon),
grid=config.grid,
grid_internal=config.grid_internal,
scale_factor=config.scale_factor,
in_chans=config.channels,
out_chans=config.channels,
embed_dim=config.embed_dim,
num_layers=config.num_layers,
use_mlp=True,
normalization_layer="none",
residual_prediction=False,
pos_embed="none",
)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
if inputs.ndim != 4:
raise ValueError(f"Expected [B, C, Nlat, Nlon], got {tuple(inputs.shape)}")
if tuple(inputs.shape[1:]) != self.expected_shape:
raise ValueError(
f"Expected trailing shape {self.expected_shape}, got {tuple(inputs.shape[1:])}"
)
if not inputs.is_floating_point():
raise TypeError("SFNO inputs must be floating point")
return self.model(inputs.float())
|