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,216 @@
+---
+frameworks:
+- ""
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- UNO
+- neural-operator
+- computational-fluid-dynamics
+- Navier-Stokes
+tasks: []
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">UNO</span>
+  </strong>
+</p>
+
+# Model Overview
+
+UNO (U-shaped Neural Operator) combines the multiscale encoder–decoder structure of U-Net with the spectral convolutions of a Fourier Neural Operator. It learns PDE solution operators across multiple spatial scales while preserving local flow-field details through skip connections.
+
+This model package targets Navier–Stokes time-series prediction on two-dimensional regular grids. By default, it takes the first 10 time steps as input and predicts the subsequent 10 autoregressively.
+
+Paper: U-NO: U-shaped Neural Operators  
+https://arxiv.org/abs/2204.11127
+
+# Repository Overview
+
+This repository is a minimal, self-contained, runnable UNO model package maintained by OneScience for ModelScope downloads, automated OneCode execution, and rapid local validation.
+
+Supported capabilities:
+
+- Train a two-dimensional UNO model from a YAML configuration
+- Perform multistep autoregressive prediction of Navier–Stokes flow fields
+- Compute relative L2 errors and save predicted tensors and visualizations
+- Override the sample count, temporal window, spatial downsampling, and model size from the command line
+- Run on a CPU, GPU, or DCU
+
+Unsupported capabilities:
+
+- Pretrained weights are not bundled
+- The approximately 394 MiB raw Navier–Stokes data file is not bundled
+- The standalone model includes only the two-dimensional UNO implementation required by this example; the general-purpose OneScience 1D and 3D components are not included
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Flow-field time-series prediction | Autoregressively predict future Navier–Stokes states from historical vorticity fields |
+| Neural operator training | Evaluate the combination of Fourier spectral convolutions and a U-shaped multiscale architecture |
+| CFD surrogate modeling | Learn mappings from historical to future fields on regular grids |
+| Pipeline validation | Validate training and inference with a small sample set and a single epoch |
+
+# File Structure
+
+| Path | Purpose | Notes |
+| :--- | :--- | :--- |
+| `README.md` | Project documentation | English |
+| `conf/config.yaml` | Data, model, training, and output configuration | Paths are resolved relative to the model package root |
+| `model/uno.py` | Standalone two-dimensional UNO model | Does not depend on `onescience.modules` |
+| `scripts/common.py` | Shared configuration, device, metric, and autoregressive utilities | Used by both training and inference |
+| `scripts/train.py` | Training and validation script | Saves the checkpoint with the best relative L2 error |
+| `scripts/inference.py` | Inference, evaluation, and visualization script | Loads `weight/*.pt` |
+| `data/` | Navier–Stokes data directory | Stores the benchmark `.mat` file |
+| `weight/` | Model weight directory | Checkpoints are written automatically during training |
+| `result/` | Inference result directory | Created automatically on first inference |
+
+# 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 full training.
+- A CPU can be used for pipeline validation with a reduced model and dataset, 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/UNO --local_dir ./UNO
+cd UNO
+```
+
+### Set Up the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported
+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
+# Installation with uv is also supported
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+## 3. Quick Start
+
+### Prepare the Data
+
+Download the following file from the Transolver `PDE-Solving-StandardBenchmark`:
+
+```text
+NavierStokes_V1e-5_N1200_T20.mat
+```
+
+Place the file in `data/`:
+
+```text
+data/
+  NavierStokes_V1e-5_N1200_T20.mat
+```
+
+Data downloads and documentation:
+
+- https://github.com/thuml/Transolver/tree/main/PDE-Solving-StandardBenchmark
+- https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-
+
+The OneScience community also provides the training data. Download it 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 checkpoint is saved to `weight/uno_navier_stokes.pt`. The training script is controlled entirely by `conf/config.yaml` and does not accept command-line configuration arguments.
+
+### Model Weights
+
+This repository provides weights trained on the standard Navier–Stokes dataset in the `weight/` directory.
+
+### Inference, Evaluation, and Visualization
+
+```bash
+python scripts/inference.py
+```
+
+The script loads `weight/uno_navier_stokes.pt` by default and generates the following files under `result/`:
+
+```text
+result/
+  prediction_sample.pt
+  prediction_sample.png
+```
+
+The inference script is likewise controlled entirely by `conf/config.yaml` and does not accept command-line configuration arguments. The weight path is determined jointly by `training.weight_dir` and `training.checkpoint_name`; the output directory and number of saved samples are controlled by `inference.result_dir` and `inference.num_samples`, respectively.
+
+# Configuration
+
+`conf/config.yaml` contains five sections:
+
+- `common`: device and random seed
+- `datapipe`: data file, sample splits, temporal windows, downsampling, and DataLoader settings
+- `model`: hidden channels, Fourier modes, normalization, and spatial padding
+- `training`: optimizer, learning-rate schedule, early stopping, and weight directory
+- `inference`: inference result directory and number of saved samples
+
+The model automatically updates `in_dim` from `t_in * out_dim` in the configuration, so the model input dimension does not need to be synchronized manually when the history window changes.
+
+# Data Format
+
+The `.mat` file must contain a variable named `u` with the following standard shape:
+
+```text
+[1200, 64, 64, 20]
+```
+
+| Dimension | Meaning |
+| --- | --- |
+| `1200` | Number of independent flow-field samples |
+| `64, 64` | Height and width of the two-dimensional regular grid |
+| `20` | Number of consecutive time steps |
+
+The data pipeline produces:
+
+- `pos`: `[H*W, 2]`, normalized two-dimensional coordinates
+- `x`: `[H*W, t_in*out_dim]`, historical states
+- `y`: `[H*W, t_out*out_dim]`, future states
+
+# 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
+
+- Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-NO: U-shaped Neural Operators. arXiv:2204.11127, 2022.
+- 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
+    space_dim: 2
+    include_pos: true
+    normtype: "in"
+    bilinear: true
+    activation: "gelu"
+    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: "uno_navier_stokes.pt"
+
+  inference:
+    result_dir: "result"
+    num_samples: 4

configuration.json

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

model/uno.py

@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+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 SpectralConv2d(nn.Module):
+    """Two-dimensional Fourier convolution over the lowest frequency modes."""
+
+    def __init__(self, in_channels: int, out_channels: int, modes1: int, modes2: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes1 = modes1
+        self.modes2 = modes2
+        scale = 1.0 / (in_channels * out_channels)
+        shape = (in_channels, out_channels, modes1, modes2)
+        self.weights1 = nn.Parameter(scale * torch.rand(*shape, dtype=torch.cfloat))
+        self.weights2 = nn.Parameter(scale * torch.rand(*shape, dtype=torch.cfloat))
+
+    @staticmethod
+    def _multiply(inputs: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        return torch.einsum("bixy,ioxy->boxy", inputs, weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, _, height, width = x.shape
+        x_ft = torch.fft.rfft2(x)
+        out_ft = torch.zeros(
+            batch_size,
+            self.out_channels,
+            height,
+            width // 2 + 1,
+            dtype=torch.cfloat,
+            device=x.device,
+        )
+        modes1 = min(self.modes1, max(1, height // 2))
+        modes2 = min(self.modes2, width // 2 + 1)
+        out_ft[:, :, :modes1, :modes2] = self._multiply(
+            x_ft[:, :, :modes1, :modes2],
+            self.weights1[:, :, :modes1, :modes2],
+        )
+        out_ft[:, :, -modes1:, :modes2] = self._multiply(
+            x_ft[:, :, -modes1:, :modes2],
+            self.weights2[:, :, :modes1, :modes2],
+        )
+        return torch.fft.irfft2(out_ft, s=(height, width))
+
+
+def _normalization(normtype: str, channels: int) -> nn.Module:
+    normalized = normtype.lower()
+    if normalized == "bn":
+        return nn.BatchNorm2d(channels)
+    if normalized == "in":
+        return nn.InstanceNorm2d(channels, affine=True)
+    if normalized in {"none", "identity"}:
+        return nn.Identity()
+    raise ValueError(f"Unsupported normtype: {normtype}")
+
+
+class DoubleConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mid_channels: int | None = None,
+        normtype: str = "in",
+    ) -> None:
+        super().__init__()
+        mid_channels = mid_channels or out_channels
+        self.layers = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            _normalization(normtype, mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            _normalization(normtype, out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
+
+class Down2d(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, normtype: str) -> None:
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv2d(in_channels, out_channels, normtype=normtype),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
+
+class Up2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        bilinear: bool,
+        normtype: str,
+    ) -> None:
+        super().__init__()
+        self.bilinear = bilinear
+        if bilinear:
+            self.up = None
+            mid_channels = in_channels // 2
+        else:
+            self.up = nn.ConvTranspose2d(
+                in_channels,
+                in_channels // 2,
+                kernel_size=2,
+                stride=2,
+            )
+            mid_channels = None
+        self.conv = DoubleConv2d(
+            in_channels,
+            out_channels,
+            mid_channels=mid_channels,
+            normtype=normtype,
+        )
+
+    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        if self.bilinear:
+            x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=True)
+        else:
+            x = self.up(x)
+            if x.shape[-2:] != skip.shape[-2:]:
+                x = F.interpolate(x, size=skip.shape[-2:], mode="nearest")
+        return self.conv(torch.cat((skip, x), dim=1))
+
+
+class UNO(nn.Module):
+    """Two-dimensional U-shaped neural operator for regular CFD grids.
+
+    ``pos`` has shape ``[B, H*W, 2]`` and ``fx`` has shape
+    ``[B, H*W, in_dim]``. The returned tensor has shape
+    ``[B, H*W, out_dim]``.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        spatial_shape: Sequence[int],
+        hidden_dim: int = 64,
+        modes: int = 12,
+        space_dim: int = 2,
+        include_pos: bool = True,
+        normtype: str = "in",
+        bilinear: bool = True,
+        activation: str = "gelu",
+        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.space_dim = int(space_dim)
+        self.include_pos = bool(include_pos)
+        self.pad_to_multiple = int(pad_to_multiple)
+        if len(self.spatial_shape) != 2:
+            raise ValueError(f"This standalone UNO supports 2D grids, got {self.spatial_shape}")
+        if self.space_dim != 2:
+            raise ValueError("space_dim must be 2 for the standalone 2D UNO")
+        if min(self.hidden_dim, self.modes, self.pad_to_multiple) < 1:
+            raise ValueError("hidden_dim, modes and pad_to_multiple must be positive")
+
+        input_dim = self.in_dim + (self.space_dim if self.include_pos else 0)
+        self.preprocess = MLP(input_dim, self.hidden_dim * 2, self.hidden_dim, activation)
+        factor = 2 if bilinear else 1
+
+        self.inc = DoubleConv2d(self.hidden_dim, self.hidden_dim, normtype=normtype)
+        self.down1 = Down2d(self.hidden_dim, self.hidden_dim * 2, normtype)
+        self.down2 = Down2d(self.hidden_dim * 2, self.hidden_dim * 4, normtype)
+        self.down3 = Down2d(self.hidden_dim * 4, self.hidden_dim * 8, normtype)
+        self.down4 = Down2d(self.hidden_dim * 8, self.hidden_dim * 16 // factor, normtype)
+
+        self.up1 = Up2d(self.hidden_dim * 16, self.hidden_dim * 8 // factor, bilinear, normtype)
+        self.up2 = Up2d(self.hidden_dim * 8, self.hidden_dim * 4 // factor, bilinear, normtype)
+        self.up3 = Up2d(self.hidden_dim * 4, self.hidden_dim * 2 // factor, bilinear, normtype)
+        self.up4 = Up2d(self.hidden_dim * 2, self.hidden_dim, bilinear, normtype)
+        self.outc = nn.Conv2d(self.hidden_dim, self.hidden_dim, kernel_size=1)
+
+        augmented = self._augmented_shape(self.spatial_shape)
+        self.padding = tuple(target - source for target, source in zip(augmented, self.spatial_shape))
+        down_channels = [
+            self.hidden_dim,
+            self.hidden_dim * 2,
+            self.hidden_dim * 4,
+            self.hidden_dim * 8,
+            self.hidden_dim * 16 // factor,
+        ]
+        up_channels = [
+            self.hidden_dim * 16 // factor,
+            self.hidden_dim * 8 // factor,
+            self.hidden_dim * 4 // factor,
+            self.hidden_dim * 2 // factor,
+            self.hidden_dim,
+        ]
+        divisors = [2, 4, 8, 16, 32]
+        self.spectral_down = nn.ModuleList(
+            self._spectral(channels, augmented, divisor)
+            for channels, divisor in zip(down_channels, divisors)
+        )
+        self.pointwise_down = nn.ModuleList(
+            nn.Conv2d(channels, channels, kernel_size=1) for channels in down_channels
+        )
+        self.spectral_up = nn.ModuleList(
+            self._spectral(channels, augmented, divisor)
+            for channels, divisor in zip(up_channels, reversed(divisors))
+        )
+        self.pointwise_up = nn.ModuleList(
+            nn.Conv2d(channels, channels, kernel_size=1) for channels in up_channels
+        )
+        self.fc1 = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
+        self.fc2 = nn.Linear(self.hidden_dim * 2, self.out_dim)
+
+    def _augmented_shape(self, shape: Sequence[int]) -> tuple[int, int]:
+        return tuple(
+            size + (self.pad_to_multiple - size % self.pad_to_multiple) % self.pad_to_multiple
+            for size in shape
+        )
+
+    def _spectral(
+        self,
+        channels: int,
+        augmented_shape: Sequence[int],
+        divisor: int,
+    ) -> SpectralConv2d:
+        modes = [
+            max(1, min(self.modes, max(1, size // divisor)))
+            for size in augmented_shape
+        ]
+        return SpectralConv2d(channels, channels, modes[0], modes[1])
+
+    @staticmethod
+    def _operator(
+        x: torch.Tensor,
+        spectral: nn.Module,
+        pointwise: nn.Module,
+    ) -> torch.Tensor:
+        return F.gelu(spectral(x) + pointwise(x))
+
+    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, fx: torch.Tensor | None = None) -> torch.Tensor:
+        batch_size, point_count, coordinates = pos.shape
+        expected_points = self.spatial_shape[0] * self.spatial_shape[1]
+        if point_count != expected_points or coordinates != self.space_dim:
+            raise ValueError(
+                f"Expected pos [B, {expected_points}, {self.space_dim}], got {tuple(pos.shape)}"
+            )
+        if fx is None:
+            if self.in_dim:
+                raise ValueError("fx is required when in_dim > 0")
+            features = pos if self.include_pos else pos.new_empty(batch_size, point_count, 0)
+        else:
+            if fx.shape != (batch_size, point_count, self.in_dim):
+                raise ValueError(
+                    f"Expected fx [B, {point_count}, {self.in_dim}], got {tuple(fx.shape)}"
+                )
+            features = torch.cat((pos, fx), dim=-1) if self.include_pos else fx
+
+        x = self.preprocess(features)
+        x = x.permute(0, 2, 1).reshape(batch_size, self.hidden_dim, *self.spatial_shape)
+        x = self._pad(x)
+
+        x1 = self._operator(self.inc(x), self.spectral_down[0], self.pointwise_down[0])
+        x2 = self._operator(self.down1(x1), self.spectral_down[1], self.pointwise_down[1])
+        x3 = self._operator(self.down2(x2), self.spectral_down[2], self.pointwise_down[2])
+        x4 = self._operator(self.down3(x3), self.spectral_down[3], self.pointwise_down[3])
+        x5 = self._operator(self.down4(x4), self.spectral_down[4], self.pointwise_down[4])
+        x = self._operator(x5, self.spectral_up[0], self.pointwise_up[0])
+        x = self._operator(self.up1(x, x4), self.spectral_up[1], self.pointwise_up[1])
+        x = self._operator(self.up2(x, x3), self.spectral_up[2], self.pointwise_up[2])
+        x = self._operator(self.up3(x, x2), self.spectral_up[3], self.pointwise_up[3])
+        x = self._operator(self.up4(x, x1), self.spectral_up[4], self.pointwise_up[4])
+        x = self._unpad(self.outc(x))
+
+        x = x.reshape(batch_size, self.hidden_dim, -1).permute(0, 2, 1)
+        return self.fc2(F.gelu(self.fc1(x)))

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,192 @@
+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
+
+
+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.uno import UNO  # 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]) -> UNO:
+    return UNO(
+        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"]),
+        space_dim=int(model_config["space_dim"]),
+        include_pos=bool(model_config["include_pos"]),
+        normtype=str(model_config["normtype"]),
+        bilinear=bool(model_config["bilinear"]),
+        activation=str(model_config["activation"]),
+        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([])
+            figure.colorbar(image, ax=axis, fraction=0.046, pad=0.04)
+    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,263 @@
+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.uno import UNO  # 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"],
+        "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]) -> UNO:
+    return UNO(
+        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"]),
+        space_dim=int(model_config["space_dim"]),
+        include_pos=bool(model_config["include_pos"]),
+        normtype=str(model_config["normtype"]),
+        bilinear=bool(model_config["bilinear"]),
+        activation=str(model_config["activation"]),
+        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/uno_navier_stokes.pt

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