Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,49 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
+ 
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,127 @@
+---
+frameworks:
+  - pytorch
+language:
+  - en
+license: apache-2.0
+tags:
+  - OneScience
+  - GFNO
+  - neural-operator
+  - group-equivariance
+  - Fourier-neural-operator
+  - Navier-Stokes
+tasks:
+  - time-series-prediction
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">GFNO</span>
+  </strong>
+</p>
+
+# Model Overview
+
+GFNO (Group Equivariant Fourier Neural Operator) is a Fourier neural operator augmented with group-equivariant structure for learning physical systems with geometric symmetries. This model predicts Navier–Stokes time series on a two-dimensional regular grid. By default, it uses the C4 rotation group and the first 10 time steps to predict the subsequent 10.
+
+Paper: Group Equivariant Fourier Neural Operators for Partial Differential Equations  
+https://openreview.net/forum?id=kgAOY5x4fi
+
+# Model Description
+
+GFNO adds group-equivariant weight sharing to the spectral convolutions of a Fourier Neural Operator, producing consistent responses under rotations or reflections and making more effective use of geometric symmetries in Navier–Stokes flow fields. The model supports the C4 rotation group and D4 dihedral group and performs multistep flow prediction autoregressively.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Flow-field time-series prediction | Predict future Navier–Stokes states from historical vorticity fields |
+| Geometric-equivariance research | Compare a standard FNO with C4- and D4-equivariant operators |
+| CFD surrogate modeling | Learn mappings from historical to future physical fields on regular grids |
+| Pipeline validation | Validate model training, inference, and result visualization with a small-scale configuration |
+
+# Usage
+
+## 1. OneCode
+
+Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
+
+[Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Setup
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended for training and inference.
+- A CPU can be used for small-scale pipeline validation, but full training will be slow.
+- DCU users must install DTK and a PyTorch environment compatible with the target cluster.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/GFNO --local_dir ./GFNO
+cd GFNO
+```
+
+### Set Up the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The model uses the standard Navier–Stokes dataset `NavierStokes_V1e-5_N1200_T20.mat`, in which the data variable `u` has shape `[1200, 64, 64, 20]`. Download the data with the following command and verify that the data path in `conf/config.yaml` is configured correctly:
+
+```bash
+modelscope download --dataset OneScience/cfd_benchmark data/ns/NavierStokes_V1e-5_N1200_T20.mat --local_dir ./data
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+The default weights are saved to `weight/gfno_navier_stokes.pt`.
+
+### Model Weights
+
+This repository provides weights trained on the standard Navier–Stokes dataset in the `weight/` directory.
+
+### Inference, Evaluation, and Visualization
+
+After training, run:
+
+```bash
+python scripts/inference.py
+```
+
+The script loads `weight/gfno_navier_stokes.pt` by default. Predicted tensors and visualizations are saved to the `result/` directory. Training and inference parameters can be modified in `conf/config.yaml`.
+
+# Official OneScience Resources
+
+| Platform | OneScience Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- Romero, D. W., Bekkers, E. J., Tomczak, J. M., and Hoogendoorn, M. Group Equivariant Fourier Neural Operators for Partial Differential Equations.
+- Li, Z. et al. Fourier Neural Operator for Parametric Partial Differential Equations. arXiv:2010.08895, 2020.
+- This model package is released under the Apache-2.0 license and retains attribution to the original paper and data sources.

conf/config.yaml

@@ -0,0 +1,53 @@
+root:
+  common:
+    device: "auto"
+    seed: 0
+
+  datapipe:
+    source:
+      data_dir: "/public/share/sugonhpcapp01/onestore/onedatasets/CFD_Benchmark/ns"
+      file_name: "NavierStokes_V1e-5_N1200_T20.mat"
+    data:
+      ntrain: 1000
+      ntest: 200
+      t_in: 10
+      t_out: 10
+      out_dim: 1
+      downsamplex: 1
+      downsampley: 1
+      normalize: false
+    dataloader:
+      batch_size: 8
+      num_workers: 4
+      pin_memory: true
+
+  model:
+    in_dim: 10
+    out_dim: 1
+    hidden_dim: 64
+    modes: 12
+    num_layers: 4
+    space_dim: 2
+    include_pos: true
+    activation: "gelu"
+    reflection: false
+    pad_to_multiple: 16
+
+  training:
+    epochs: 500
+    lr: 0.001
+    weight_decay: 0.00001
+    step_size: 100
+    gamma: 0.5
+    eval_interval: 1
+    patience: 100
+    teacher_forcing: false
+    max_grad_norm: null
+    max_train_batches: null
+    max_eval_batches: null
+    weight_dir: "weight"
+    checkpoint_name: "gfno_navier_stokes.pt"
+
+  inference:
+    result_dir: "result"
+    num_samples: 4

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/gfno.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _activation(name: str) -> nn.Module:
+    activations = {
+        "gelu": nn.GELU,
+        "relu": nn.ReLU,
+        "silu": nn.SiLU,
+        "tanh": nn.Tanh,
+    }
+    try:
+        return activations[name.lower()]()
+    except KeyError as error:
+        raise ValueError(f"Unsupported activation: {name}") from error
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        activation: str,
+    ) -> None:
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            _activation(activation),
+            nn.Linear(hidden_dim, output_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
+
+class GroupEquivariantConv2d(nn.Module):
+    """C4/D4 group convolution used by the GFNO lifting and channel mixers."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        bias: bool = True,
+        first_layer: bool = False,
+        last_layer: bool = False,
+        spectral: bool = False,
+        hermitian: bool = False,
+        reflection: bool = False,
+    ) -> None:
+        super().__init__()
+        if kernel_size % 2 != 1:
+            raise ValueError("kernel_size must be odd")
+        self.in_channels = int(in_channels)
+        self.out_channels = int(out_channels)
+        self.reflection = bool(reflection)
+        self.rotation_group_size = 4
+        self.group_size = 4 * (1 + int(self.reflection))
+        self.kernel_size_y = int(kernel_size)
+        self.kernel_size_x = kernel_size // 2 + 1 if hermitian else kernel_size
+        self.first_layer = bool(first_layer)
+        self.last_layer = bool(last_layer)
+        self.hermitian = bool(hermitian)
+        dtype = torch.cfloat if spectral else torch.float32
+
+        if self.first_layer or self.last_layer:
+            self.weight = nn.Parameter(
+                torch.empty(
+                    self.out_channels,
+                    1,
+                    self.in_channels,
+                    self.kernel_size_y,
+                    self.kernel_size_x,
+                    dtype=dtype,
+                )
+            )
+        elif self.hermitian:
+            self.weight = nn.ParameterDict(
+                {
+                    "vertical": nn.Parameter(
+                        torch.empty(
+                            self.out_channels,
+                            1,
+                            self.in_channels,
+                            self.group_size,
+                            self.kernel_size_x - 1,
+                            1,
+                            dtype=dtype,
+                        )
+                    ),
+                    "positive": nn.Parameter(
+                        torch.empty(
+                            self.out_channels,
+                            1,
+                            self.in_channels,
+                            self.group_size,
+                            self.kernel_size_y,
+                            self.kernel_size_x - 1,
+                            dtype=dtype,
+                        )
+                    ),
+                    "origin": nn.Parameter(
+                        torch.empty(
+                            self.out_channels,
+                            1,
+                            self.in_channels,
+                            self.group_size,
+                            1,
+                            1,
+                        )
+                    ),
+                }
+            )
+        else:
+            self.weight = nn.Parameter(
+                torch.empty(
+                    self.out_channels,
+                    1,
+                    self.in_channels,
+                    self.group_size,
+                    self.kernel_size_y,
+                    self.kernel_size_x,
+                    dtype=dtype,
+                )
+            )
+        self.base_bias = (
+            nn.Parameter(torch.empty(1, self.out_channels, 1, 1))
+            if bias
+            else None
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if isinstance(self.weight, nn.ParameterDict):
+            for parameter in self.weight.values():
+                nn.init.kaiming_uniform_(parameter, a=math.sqrt(5))
+        else:
+            nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.base_bias is not None:
+            nn.init.kaiming_uniform_(self.base_bias, a=math.sqrt(5))
+
+    def transformed_weight(self) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if self.hermitian:
+            vertical = self.weight["vertical"]
+            weights = torch.cat(
+                (
+                    vertical,
+                    self.weight["origin"].cfloat(),
+                    vertical.flip(dims=(-2,)).conj(),
+                ),
+                dim=-2,
+            )
+            weights = torch.cat((weights, self.weight["positive"]), dim=-1)
+            weights = torch.cat(
+                (weights[..., 1:].conj().rot90(k=2, dims=(-2, -1)), weights),
+                dim=-1,
+            )
+        else:
+            weights = self.weight
+
+        if self.first_layer or self.last_layer:
+            weights = weights.repeat(1, self.group_size, 1, 1, 1)
+            for rotation in range(1, self.rotation_group_size):
+                weights[:, rotation] = weights[:, rotation].rot90(
+                    k=rotation, dims=(-2, -1)
+                )
+            if self.reflection:
+                weights[:, self.rotation_group_size :] = weights[
+                    :, : self.rotation_group_size
+                ].flip(dims=(-2,))
+
+            if self.first_layer:
+                weights = weights.reshape(
+                    self.out_channels * self.group_size,
+                    self.in_channels,
+                    self.kernel_size_y,
+                    self.kernel_size_y,
+                )
+                bias = (
+                    self.base_bias.repeat_interleave(self.group_size, dim=1).reshape(-1)
+                    if self.base_bias is not None
+                    else None
+                )
+            else:
+                weights = weights.transpose(2, 1).reshape(
+                    self.out_channels,
+                    self.in_channels * self.group_size,
+                    self.kernel_size_y,
+                    self.kernel_size_y,
+                )
+                bias = self.base_bias.reshape(-1) if self.base_bias is not None else None
+        else:
+            weights = weights.repeat(1, self.group_size, 1, 1, 1, 1)
+            for rotation in range(1, self.rotation_group_size):
+                weights[:, rotation] = weights[:, rotation - 1].rot90(
+                    dims=(-2, -1)
+                )
+                if self.reflection:
+                    weights[:, rotation] = torch.cat(
+                        (
+                            weights[
+                                :, rotation, :, self.rotation_group_size - 1
+                            ].unsqueeze(2),
+                            weights[
+                                :, rotation, :, : self.rotation_group_size - 1
+                            ],
+                            weights[
+                                :, rotation, :, self.rotation_group_size + 1 :
+                            ],
+                            weights[
+                                :, rotation, :, self.rotation_group_size
+                            ].unsqueeze(2),
+                        ),
+                        dim=2,
+                    )
+                else:
+                    weights[:, rotation] = torch.cat(
+                        (
+                            weights[:, rotation, :, -1].unsqueeze(2),
+                            weights[:, rotation, :, :-1],
+                        ),
+                        dim=2,
+                    )
+            if self.reflection:
+                weights[:, self.rotation_group_size :] = torch.cat(
+                    (
+                        weights[
+                            :, : self.rotation_group_size, :, self.rotation_group_size :
+                        ],
+                        weights[
+                            :, : self.rotation_group_size, :, : self.rotation_group_size
+                        ],
+                    ),
+                    dim=3,
+                ).flip(dims=(-2,))
+
+            weights = weights.reshape(
+                self.out_channels * self.group_size,
+                self.in_channels * self.group_size,
+                self.kernel_size_y,
+                self.kernel_size_y,
+            )
+            bias = (
+                self.base_bias.repeat_interleave(self.group_size, dim=1).reshape(-1)
+                if self.base_bias is not None
+                else None
+            )
+
+        if self.hermitian:
+            weights = weights[..., -self.kernel_size_x :]
+        return weights, bias
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        weights, bias = self.transformed_weight()
+        return F.conv2d(x, weights, bias=bias)
+
+
+class GroupSpectralConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        modes: int,
+        reflection: bool,
+    ) -> None:
+        super().__init__()
+        self.modes = (int(modes), int(modes))
+        kernel_size = 2 * max(self.modes) - 1
+        self.group_conv = GroupEquivariantConv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            bias=False,
+            spectral=True,
+            hermitian=True,
+            reflection=reflection,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, _, height, width = x.shape
+        modes_y, modes_x = self.modes
+        zero_y = height // 2
+        start_y = zero_y - modes_y + 1
+        stop_y = zero_y + modes_y
+        if start_y < 0 or stop_y > height or modes_x > width // 2 + 1:
+            raise ValueError(
+                f"modes={self.modes} exceed the FFT grid {(height, width)}"
+            )
+
+        weights, _ = self.group_conv.transformed_weight()
+        weights = weights.transpose(0, 1)
+        x_ft = torch.fft.fftshift(torch.fft.rfft2(x), dim=-2)
+        filtered = x_ft[..., start_y:stop_y, :modes_x]
+        out_ft = torch.zeros(
+            batch_size,
+            weights.shape[1],
+            height,
+            width // 2 + 1,
+            dtype=torch.cfloat,
+            device=x.device,
+        )
+        out_ft[..., start_y:stop_y, :modes_x] = torch.einsum(
+            "bixy,ioxy->boxy", filtered, weights
+        )
+        return torch.fft.irfft2(
+            torch.fft.ifftshift(out_ft, dim=-2), s=(height, width)
+        )
+
+
+class GroupMLP2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mid_channels: int,
+        reflection: bool,
+        last_layer: bool = False,
+    ) -> None:
+        super().__init__()
+        self.first = GroupEquivariantConv2d(
+            in_channels, mid_channels, 1, reflection=reflection
+        )
+        self.second = GroupEquivariantConv2d(
+            mid_channels,
+            out_channels,
+            1,
+            reflection=reflection,
+            last_layer=last_layer,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.second(F.gelu(self.first(x)))
+
+
+class GroupNorm(nn.Module):
+    def __init__(self, width: int, group_size: int) -> None:
+        super().__init__()
+        self.group_size = group_size
+        self.norm = nn.InstanceNorm3d(width)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, _, height, width = x.shape
+        x = x.reshape(batch_size, -1, self.group_size, height, width)
+        x = self.norm(x)
+        return x.reshape(batch_size, -1, height, width)
+
+
+class GFNO(nn.Module):
+    """C4/D4 group-equivariant Fourier neural operator for regular 2D grids."""
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        spatial_shape: Sequence[int],
+        hidden_dim: int = 64,
+        modes: int = 12,
+        num_layers: int = 4,
+        space_dim: int = 2,
+        include_pos: bool = True,
+        activation: str = "gelu",
+        reflection: bool = False,
+        pad_to_multiple: int = 16,
+    ) -> None:
+        super().__init__()
+        self.in_dim = int(in_dim)
+        self.out_dim = int(out_dim)
+        self.spatial_shape = tuple(int(value) for value in spatial_shape)
+        self.hidden_dim = int(hidden_dim)
+        self.modes = int(modes)
+        self.num_layers = int(num_layers)
+        self.space_dim = int(space_dim)
+        self.include_pos = bool(include_pos)
+        self.reflection = bool(reflection)
+        self.pad_to_multiple = int(pad_to_multiple)
+        if len(self.spatial_shape) != 2 or self.space_dim != 2:
+            raise ValueError("GFNO supports two-dimensional structured grids")
+        if min(self.hidden_dim, self.modes, self.num_layers, self.pad_to_multiple) < 1:
+            raise ValueError(
+                "hidden_dim, modes, num_layers and pad_to_multiple must be positive"
+            )
+
+        feature_dim = self.in_dim + (self.space_dim if self.include_pos else 0)
+        self.preprocess = MLP(
+            feature_dim, self.hidden_dim * 2, self.hidden_dim, activation
+        )
+        self.group_size = 4 * (1 + int(self.reflection))
+        self.padding = tuple(
+            (self.pad_to_multiple - size % self.pad_to_multiple)
+            % self.pad_to_multiple
+            for size in self.spatial_shape
+        )
+        augmented_shape = tuple(
+            size + padding
+            for size, padding in zip(self.spatial_shape, self.padding)
+        )
+        if self.modes > augmented_shape[0] // 2:
+            raise ValueError(
+                f"modes={self.modes} exceed padded grid {augmented_shape}"
+            )
+
+        self.lifting = GroupEquivariantConv2d(
+            self.hidden_dim,
+            self.hidden_dim,
+            1,
+            reflection=self.reflection,
+            first_layer=True,
+        )
+        self.spectral_layers = nn.ModuleList(
+            GroupSpectralConv2d(
+                self.hidden_dim,
+                self.hidden_dim,
+                self.modes,
+                self.reflection,
+            )
+            for _ in range(self.num_layers)
+        )
+        self.mlp_layers = nn.ModuleList(
+            GroupMLP2d(
+                self.hidden_dim,
+                self.hidden_dim,
+                self.hidden_dim,
+                self.reflection,
+            )
+            for _ in range(self.num_layers)
+        )
+        self.residual_layers = nn.ModuleList(
+            GroupEquivariantConv2d(
+                self.hidden_dim,
+                self.hidden_dim,
+                1,
+                reflection=self.reflection,
+            )
+            for _ in range(self.num_layers)
+        )
+        self.norm = GroupNorm(self.hidden_dim, self.group_size)
+        self.output = GroupMLP2d(
+            self.hidden_dim,
+            self.out_dim,
+            self.hidden_dim * 4,
+            self.reflection,
+            last_layer=True,
+        )
+
+    def _pad(self, x: torch.Tensor) -> torch.Tensor:
+        pad_height, pad_width = self.padding
+        return F.pad(x, (0, pad_width, 0, pad_height)) if any(self.padding) else x
+
+    def _unpad(self, x: torch.Tensor) -> torch.Tensor:
+        height, width = self.spatial_shape
+        return x[..., :height, :width]
+
+    def forward(
+        self, pos: torch.Tensor, field: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        batch_size, point_count, coordinate_dim = pos.shape
+        expected_points = self.spatial_shape[0] * self.spatial_shape[1]
+        if point_count != expected_points or coordinate_dim != self.space_dim:
+            raise ValueError(
+                f"Expected pos [B, {expected_points}, {self.space_dim}], "
+                f"got {tuple(pos.shape)}"
+            )
+        if field is None:
+            if self.in_dim:
+                raise ValueError("field is required when in_dim > 0")
+            features = pos if self.include_pos else pos.new_empty(batch_size, point_count, 0)
+        else:
+            if field.shape != (batch_size, point_count, self.in_dim):
+                raise ValueError(
+                    f"Expected field [B, {point_count}, {self.in_dim}], "
+                    f"got {tuple(field.shape)}"
+                )
+            features = torch.cat((pos, field), dim=-1) if self.include_pos else field
+
+        x = self.preprocess(features)
+        x = x.permute(0, 2, 1).reshape(
+            batch_size, self.hidden_dim, *self.spatial_shape
+        )
+        x = self.lifting(self._pad(x))
+        for layer_index, (spectral, mlp, residual) in enumerate(
+            zip(self.spectral_layers, self.mlp_layers, self.residual_layers)
+        ):
+            update = self.norm(spectral(self.norm(x)))
+            x = mlp(update) + residual(x)
+            if layer_index != self.num_layers - 1:
+                x = F.gelu(x)
+        x = self.output(self._unpad(x))
+        return x.reshape(batch_size, self.out_dim, -1).permute(0, 2, 1)

scripts/common.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import random
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import yaml
+
+
+class AttrDict(dict):
+    """Dictionary with recursive attribute access for OneScience datapipes."""
+
+    def __getattr__(self, key: str) -> Any:
+        try:
+            return self[key]
+        except KeyError as error:
+            raise AttributeError(key) from error
+
+    def __setattr__(self, key: str, value: Any) -> None:
+        self[key] = value
+
+
+def to_attr_dict(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return AttrDict({key: to_attr_dict(item) for key, item in value.items()})
+    if isinstance(value, list):
+        return [to_attr_dict(item) for item in value]
+    return value
+
+
+def to_plain_dict(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return {key: to_plain_dict(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [to_plain_dict(item) for item in value]
+    return value
+
+
+def load_config(path: Path) -> dict[str, Any]:
+    with path.open("r", encoding="utf-8") as stream:
+        config = yaml.safe_load(stream)
+    if not isinstance(config, dict) or not isinstance(config.get("root"), dict):
+        raise ValueError(f"Config must contain a root mapping: {path}")
+    return config["root"]
+
+
+def project_path(value: str | Path, project_root: Path) -> Path:
+    path = Path(value).expanduser()
+    return path if path.is_absolute() else project_root / path
+
+
+def resolve_device(requested: str) -> torch.device:
+    if requested == "auto":
+        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    device = torch.device(requested)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
+    return device
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def rollout(
+    model: torch.nn.Module,
+    pos: torch.Tensor,
+    state: torch.Tensor,
+    t_out: int,
+    out_dim: int,
+    teacher_forcing: bool = False,
+    target: torch.Tensor | None = None,
+) -> torch.Tensor:
+    predictions = []
+    for step in range(t_out):
+        prediction = model(pos, state)
+        predictions.append(prediction)
+        next_frame = (
+            target[..., step * out_dim : (step + 1) * out_dim]
+            if teacher_forcing and target is not None
+            else prediction
+        )
+        state = torch.cat((state[..., out_dim:], next_frame), dim=-1)
+    return torch.cat(predictions, dim=-1)
+
+
+def relative_l2(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    prediction = prediction.reshape(prediction.shape[0], -1)
+    target = target.reshape(target.shape[0], -1)
+    numerator = torch.linalg.vector_norm(prediction - target, dim=1)
+    denominator = torch.linalg.vector_norm(target, dim=1).clamp_min(1.0e-8)
+    return numerator / denominator

scripts/inference.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import Any
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import torch
+from matplotlib.ticker import FormatStrFormatter
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+LOCAL_WORKSPACE = PROJECT_ROOT.parents[1]
+if (LOCAL_WORKSPACE / "onescience" / "src" / "onescience" / "datapipes").is_dir():
+    sys.path.insert(0, str(LOCAL_WORKSPACE))
+
+from common import (  # noqa: E402
+    load_config,
+    project_path,
+    relative_l2,
+    resolve_device,
+    rollout,
+    to_attr_dict,
+)
+from model.gfno import GFNO  # noqa: E402
+from onescience.datapipes.cfd import NavierStokesDatapipe  # noqa: E402
+from onescience.distributed.manager import DistributedManager  # noqa: E402
+
+
+DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
+
+
+def prepare_config(config: dict[str, Any]) -> None:
+    datapipe = config["datapipe"]
+    datapipe["source"]["data_dir"] = str(
+        project_path(datapipe["source"]["data_dir"], PROJECT_ROOT).resolve()
+    )
+
+
+def validate_checkpoint_config(
+    datapipe_config: dict[str, Any],
+    checkpoint: dict[str, Any],
+) -> None:
+    checkpoint_data = checkpoint["datapipe_config"]["data"]
+    current_data = datapipe_config["data"]
+    for key in ("t_in", "t_out", "out_dim", "downsamplex", "downsampley", "normalize"):
+        if current_data[key] != checkpoint_data[key]:
+            raise ValueError(
+                f"Config value datapipe.data.{key}={current_data[key]} does not match "
+                f"checkpoint value {checkpoint_data[key]}"
+            )
+
+
+def build_model(model_config: dict[str, Any], spatial_shape: tuple[int, int]) -> GFNO:
+    return GFNO(
+        in_dim=int(model_config["in_dim"]),
+        out_dim=int(model_config["out_dim"]),
+        spatial_shape=spatial_shape,
+        hidden_dim=int(model_config["hidden_dim"]),
+        modes=int(model_config["modes"]),
+        num_layers=int(model_config["num_layers"]),
+        space_dim=int(model_config["space_dim"]),
+        include_pos=bool(model_config["include_pos"]),
+        activation=str(model_config["activation"]),
+        reflection=bool(model_config["reflection"]),
+        pad_to_multiple=int(model_config["pad_to_multiple"]),
+    )
+
+
+def save_visualization(
+    prediction: torch.Tensor,
+    target: torch.Tensor,
+    spatial_shape: tuple[int, int],
+    output_path: Path,
+) -> None:
+    prediction = prediction.detach().cpu()
+    target = target.detach().cpu()
+    height, width = spatial_shape
+    frame_count = prediction.shape[-1]
+    frames = sorted({0, frame_count // 2, frame_count - 1})
+    sample_prediction = prediction[0].reshape(height, width, frame_count)
+    sample_target = target[0].reshape(height, width, frame_count)
+    rows = (
+        ("Target", sample_target),
+        ("Prediction", sample_prediction),
+        ("Abs Error", (sample_prediction - sample_target).abs()),
+    )
+    figure, axes = plt.subplots(3, len(frames), figsize=(4 * len(frames), 9), squeeze=False)
+    for row_index, (title, values) in enumerate(rows):
+        for column_index, frame in enumerate(frames):
+            axis = axes[row_index][column_index]
+            image = axis.imshow(values[..., frame], origin="lower", cmap="viridis")
+            axis.set_title(f"{title} t+{frame + 1}")
+            axis.set_xticks([])
+            axis.set_yticks([])
+            colorbar = figure.colorbar(image, ax=axis, fraction=0.046, pad=0.04)
+            colorbar.formatter = FormatStrFormatter("%.3g")
+            colorbar.update_ticks()
+    figure.tight_layout()
+    figure.savefig(output_path, dpi=150)
+    plt.close(figure)
+
+
+def main() -> None:
+    config_path = DEFAULT_CONFIG.resolve()
+    config = load_config(config_path)
+    prepare_config(config)
+    training_config = config["training"]
+    inference_config = config["inference"]
+    datapipe_config = config["datapipe"]
+    num_samples = int(inference_config["num_samples"])
+    if num_samples < 1:
+        raise ValueError("num_samples must be positive")
+    checkpoint_path = (
+        project_path(training_config["weight_dir"], PROJECT_ROOT).resolve()
+        / str(training_config["checkpoint_name"])
+    )
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model_config = dict(checkpoint["model_config"])
+    validate_checkpoint_config(datapipe_config, checkpoint)
+    data_file = (
+        Path(datapipe_config["source"]["data_dir"])
+        / datapipe_config["source"]["file_name"]
+    )
+    device = resolve_device(str(config["common"]["device"]))
+    result_dir = project_path(
+        inference_config["result_dir"], PROJECT_ROOT
+    ).resolve()
+
+    print(f"Checkpoint: {checkpoint_path}")
+    print(f"Data: {data_file}")
+    print(f"Device: {device}")
+    DistributedManager.initialize()
+    try:
+        datapipe = NavierStokesDatapipe(
+            to_attr_dict(datapipe_config),
+            distributed=False,
+            normalizer_state=checkpoint.get("normalizer"),
+        )
+        test_loader, _ = datapipe.test_dataloader()
+        spatial_shape = tuple(checkpoint["spatial_shape"])
+        if tuple(datapipe.spatial_shape) != spatial_shape:
+            raise ValueError(
+                f"Checkpoint grid {spatial_shape} does not match data grid {datapipe.spatial_shape}"
+            )
+        model = build_model(model_config, spatial_shape).to(device)
+        model.load_state_dict(checkpoint["model_state"])
+        model.eval()
+
+        batch = next(iter(test_loader))
+        pos = batch["pos"].to(device)
+        state = batch["x"].to(device)
+        target = batch["y"].to(device)
+        with torch.no_grad():
+            prediction = rollout(
+                model,
+                pos,
+                state,
+                int(datapipe_config["data"]["t_out"]),
+                int(datapipe_config["data"]["out_dim"]),
+            )
+        prediction = datapipe.decode_solution(prediction)
+        target = datapipe.decode_solution(target)
+        error = relative_l2(prediction, target).mean().item()
+        sample_count = min(num_samples, prediction.shape[0])
+        result_dir.mkdir(parents=True, exist_ok=True)
+        tensor_path = result_dir / "prediction_sample.pt"
+        image_path = result_dir / "prediction_sample.png"
+        torch.save(
+            {
+                "prediction": prediction[:sample_count].cpu(),
+                "target": target[:sample_count].cpu(),
+                "spatial_shape": spatial_shape,
+                "relative_l2": error,
+            },
+            tensor_path,
+        )
+        save_visualization(prediction[:1], target[:1], spatial_shape, image_path)
+        print(f"Relative L2: {error:.6e}")
+        print(f"Saved tensor: {tensor_path}")
+        print(f"Saved figure: {image_path}")
+    finally:
+        DistributedManager.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+LOCAL_WORKSPACE = PROJECT_ROOT.parents[1]
+if (LOCAL_WORKSPACE / "onescience" / "src" / "onescience" / "datapipes").is_dir():
+    sys.path.insert(0, str(LOCAL_WORKSPACE))
+
+from common import (  # noqa: E402
+    load_config,
+    project_path,
+    relative_l2,
+    resolve_device,
+    rollout,
+    seed_everything,
+    to_attr_dict,
+    to_plain_dict,
+)
+from model.gfno import GFNO  # noqa: E402
+from onescience.datapipes.cfd import NavierStokesDatapipe  # noqa: E402
+from onescience.distributed.manager import DistributedManager  # noqa: E402
+
+
+DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
+
+
+def prepare_config(config: dict[str, Any]) -> None:
+    datapipe = config["datapipe"]
+    data = datapipe["data"]
+    model = config["model"]
+    datapipe["source"]["data_dir"] = str(
+        project_path(datapipe["source"]["data_dir"], PROJECT_ROOT).resolve()
+    )
+    model["in_dim"] = int(data["t_in"]) * int(data["out_dim"])
+    model["out_dim"] = int(data["out_dim"])
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    data = config["datapipe"]["data"]
+    loader = config["datapipe"]["dataloader"]
+    model = config["model"]
+    training = config["training"]
+    positive_values = {
+        "ntrain": data["ntrain"],
+        "ntest": data["ntest"],
+        "t_in": data["t_in"],
+        "t_out": data["t_out"],
+        "downsamplex": data["downsamplex"],
+        "downsampley": data["downsampley"],
+        "batch_size": loader["batch_size"],
+        "hidden_dim": model["hidden_dim"],
+        "modes": model["modes"],
+        "num_layers": model["num_layers"],
+        "epochs": training["epochs"],
+        "eval_interval": training["eval_interval"],
+    }
+    for name, value in positive_values.items():
+        if int(value) < 1:
+            raise ValueError(f"{name} must be positive, got {value}")
+    if int(loader["num_workers"]) < 0:
+        raise ValueError("num_workers cannot be negative")
+    for name in ("max_train_batches", "max_eval_batches"):
+        value = training.get(name)
+        if value is not None and int(value) < 1:
+            raise ValueError(f"{name} must be positive when set")
+
+
+def build_model(model_config: dict[str, Any], spatial_shape: tuple[int, int]) -> GFNO:
+    return GFNO(
+        in_dim=int(model_config["in_dim"]),
+        out_dim=int(model_config["out_dim"]),
+        spatial_shape=spatial_shape,
+        hidden_dim=int(model_config["hidden_dim"]),
+        modes=int(model_config["modes"]),
+        num_layers=int(model_config["num_layers"]),
+        space_dim=int(model_config["space_dim"]),
+        include_pos=bool(model_config["include_pos"]),
+        activation=str(model_config["activation"]),
+        reflection=bool(model_config["reflection"]),
+        pad_to_multiple=int(model_config["pad_to_multiple"]),
+    )
+
+
+def evaluate(
+    model: torch.nn.Module,
+    loader: torch.utils.data.DataLoader,
+    device: torch.device,
+    t_out: int,
+    out_dim: int,
+    max_batches: int | None,
+) -> tuple[float, float]:
+    model.eval()
+    squared_error = 0.0
+    element_count = 0
+    relative_error = 0.0
+    sample_count = 0
+    with torch.no_grad():
+        for batch_index, batch in enumerate(tqdm(loader, desc="Evaluating")):
+            if max_batches is not None and batch_index >= max_batches:
+                break
+            pos = batch["pos"].to(device)
+            state = batch["x"].to(device)
+            target = batch["y"].to(device)
+            prediction = rollout(model, pos, state, t_out, out_dim)
+            squared_error += F.mse_loss(prediction, target, reduction="sum").item()
+            element_count += target.numel()
+            relative_error += relative_l2(prediction, target).sum().item()
+            sample_count += target.shape[0]
+    return (
+        squared_error / max(element_count, 1),
+        relative_error / max(sample_count, 1),
+    )
+
+
+def main() -> None:
+    config_path = DEFAULT_CONFIG.resolve()
+    config = load_config(config_path)
+    prepare_config(config)
+    validate_config(config)
+
+    common = config["common"]
+    data_config = config["datapipe"]
+    model_config = config["model"]
+    training = config["training"]
+    device = resolve_device(str(common["device"]))
+    seed_everything(int(common["seed"]))
+    weight_dir = project_path(training["weight_dir"], PROJECT_ROOT).resolve()
+    checkpoint_path = weight_dir / str(training["checkpoint_name"])
+
+    print(f"Config: {config_path}")
+    print(f"Data: {Path(data_config['source']['data_dir']) / data_config['source']['file_name']}")
+    print(f"Device: {device}")
+    print(
+        f"Samples: train={data_config['data']['ntrain']} test={data_config['data']['ntest']} "
+        f"t_in={data_config['data']['t_in']} t_out={data_config['data']['t_out']}"
+    )
+
+    DistributedManager.initialize()
+    started = time.time()
+    try:
+        datapipe = NavierStokesDatapipe(to_attr_dict(data_config), distributed=False)
+        train_loader, _ = datapipe.train_dataloader()
+        test_loader, _ = datapipe.test_dataloader()
+        model = build_model(model_config, datapipe.spatial_shape).to(device)
+        parameter_count = sum(parameter.numel() for parameter in model.parameters())
+        print(f"Spatial shape: {datapipe.spatial_shape}")
+        print(f"Parameters: {parameter_count:,}")
+
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=float(training["lr"]),
+            weight_decay=float(training["weight_decay"]),
+        )
+        scheduler = torch.optim.lr_scheduler.StepLR(
+            optimizer,
+            step_size=int(training["step_size"]),
+            gamma=float(training["gamma"]),
+        )
+        t_out = int(data_config["data"]["t_out"])
+        out_dim = int(data_config["data"]["out_dim"])
+        max_train_batches = training.get("max_train_batches")
+        max_train_batches = None if max_train_batches is None else int(max_train_batches)
+        max_eval_batches = training.get("max_eval_batches")
+        max_eval_batches = None if max_eval_batches is None else int(max_eval_batches)
+        best_relative_l2 = float("inf")
+        stale_evaluations = 0
+
+        for epoch in range(1, int(training["epochs"]) + 1):
+            model.train()
+            epoch_loss = 0.0
+            batch_count = 0
+            iterator = tqdm(train_loader, desc=f"Epoch {epoch}")
+            for batch_index, batch in enumerate(iterator):
+                if max_train_batches is not None and batch_index >= max_train_batches:
+                    break
+                pos = batch["pos"].to(device)
+                state = batch["x"].to(device)
+                target = batch["y"].to(device)
+                optimizer.zero_grad(set_to_none=True)
+                prediction = rollout(
+                    model,
+                    pos,
+                    state,
+                    t_out,
+                    out_dim,
+                    teacher_forcing=bool(training["teacher_forcing"]),
+                    target=target,
+                )
+                loss = F.mse_loss(prediction, target)
+                if not torch.isfinite(loss):
+                    raise FloatingPointError(f"Non-finite training loss at epoch {epoch}")
+                loss.backward()
+                if training.get("max_grad_norm") is not None:
+                    torch.nn.utils.clip_grad_norm_(
+                        model.parameters(), float(training["max_grad_norm"])
+                    )
+                optimizer.step()
+                epoch_loss += loss.item()
+                batch_count += 1
+                iterator.set_postfix(loss=f"{loss.item():.3e}")
+            scheduler.step()
+            if batch_count == 0:
+                raise RuntimeError("No training batches were processed")
+
+            should_evaluate = (
+                epoch % int(training["eval_interval"]) == 0
+                or epoch == int(training["epochs"])
+            )
+            if not should_evaluate:
+                continue
+            validation_mse, validation_relative_l2 = evaluate(
+                model,
+                test_loader,
+                device,
+                t_out,
+                out_dim,
+                max_eval_batches,
+            )
+            print(
+                f"epoch={epoch:4d} train_mse={epoch_loss / batch_count:.6e} "
+                f"val_mse={validation_mse:.6e} val_relative_l2={validation_relative_l2:.6e}"
+            )
+            if validation_relative_l2 < best_relative_l2:
+                best_relative_l2 = validation_relative_l2
+                stale_evaluations = 0
+                weight_dir.mkdir(parents=True, exist_ok=True)
+                torch.save(
+                    {
+                        "epoch": epoch,
+                        "model_state": model.state_dict(),
+                        "model_config": to_plain_dict(model_config),
+                        "datapipe_config": to_plain_dict(data_config),
+                        "training_config": to_plain_dict(training),
+                        "spatial_shape": tuple(datapipe.spatial_shape),
+                        "normalizer": datapipe.get_normalizer_state(),
+                        "best_relative_l2": best_relative_l2,
+                    },
+                    checkpoint_path,
+                )
+                print(f"Saved checkpoint: {checkpoint_path}")
+            else:
+                stale_evaluations += 1
+            if stale_evaluations >= int(training["patience"]):
+                print("Early stopping triggered")
+                break
+
+        print(f"Elapsed: {time.time() - started:.1f}s")
+    finally:
+        DistributedManager.cleanup()
+
+
+if __name__ == "__main__":
+    main()

weight/gfno_navier_stokes.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b1fb0e6c0dbfbd1767a78d2d65594b9a61397a0543554eb87c5758c08db5812
+size 139803378