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: gpl-3.0
+tags:
+  - OneScience
+  - KNO
+  - Koopman-operator
+  - neural-operator
+  - computational-fluid-dynamics
+  - Navier-Stokes
+tasks:
+  - time-series-prediction
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">KNO</span>
+  </strong>
+</p>
+
+# Model Overview
+
+KNO (Koopman Neural Operator) is a neural operator based on Koopman operator theory for learning the evolution of nonlinear dynamical systems. This model predicts Navier–Stokes time series on a two-dimensional regular grid, using the first 10 time steps to predict the subsequent 10 by default.
+
+Paper: Koopman Neural Operator as a Mesh-free Solver of Non-linear Partial Differential Equations  
+https://doi.org/10.1016/j.jcp.2024.113194
+
+# Model Description
+
+KNO uses an encoder to map historical physical fields into a latent Koopman space, learns an approximately linear evolution operator in the Fourier domain, and reconstructs future physical fields with a decoder. The model supports either linear or nonlinear latent-state propagation 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 |
+| Koopman operator research | Study approximately linear evolution of nonlinear dynamical systems in latent space |
+| CFD surrogate modeling | Learn mappings from historical to future physical fields on regular grids |
+| Pipeline validation | Validate training and inference using the bundled weights or 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/KNO --local_dir ./KNO
+cd KNO
+```
+
+### 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/kno_navier_stokes.pt`.
+
+### Model Weights
+
+This repository provides weights trained on the standard Navier–Stokes dataset in the `weight/` directory.
+
+### Inference, Evaluation, and Visualization
+
+The model package includes weights for pipeline validation. After preparing the data, run:
+
+```bash
+python scripts/inference.py
+```
+
+The script loads `weight/kno_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
+
+- Xiong, W. et al. Koopman Neural Operator as a Mesh-free Solver of Non-linear Partial Differential Equations. Journal of Computational Physics, 2024.
+- Xiong, W. et al. KoopmanLab: Machine Learning for Solving Complex Physics Equations. APL Machine Learning, 2023.
+- The model implementation is derived from KoopmanLab under the GPL-3.0 license. This model package retains the GPL-3.0 license and the corresponding source attribution.

conf/config.yaml

@@ -0,0 +1,52 @@
+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:
+    name: "KNO2DNavierStokes"
+    input_channels: 10
+    output_channels: 1
+    op_size: 32
+    modes_x: 12
+    modes_y: 12
+    decompose: 6
+    linear_type: true
+    normalization: false
+
+  training:
+    epochs: 500
+    lr: 0.001
+    weight_decay: 0.00001
+    step_size: 100
+    gamma: 0.5
+    eval_interval: 10
+    patience: 100
+    teacher_forcing: true
+    max_grad_norm: null
+    max_train_batches: null
+    max_eval_batches: null
+    weight_dir: "weight"
+    checkpoint_name: "kno_navier_stokes.pt"
+
+  inference:
+    result_dir: "result"
+    num_samples: 4

configuration.json

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

download.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+# Download files larger than 1MB, excluding .sh .py .md .yaml .yml
+# Total large files: 1
+modelscope download --model OneScience/KNO weight/kno_navier_stokes.pt --local_dir ./

model/kno.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class EncoderMLP(nn.Module):
+    def __init__(self, t_len: int, op_size: int):
+        super().__init__()
+        self.layer = nn.Linear(t_len, op_size)
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class DecoderMLP(nn.Module):
+    def __init__(self, t_len: int, op_size: int):
+        super().__init__()
+        self.layer = nn.Linear(op_size, t_len)
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class EncoderConv1D(nn.Module):
+    def __init__(self, t_len: int, op_size: int):
+        super().__init__()
+        self.layer = nn.Conv1d(t_len, op_size, 1)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 1)
+        x = self.layer(x)
+        return x.permute(0, 2, 1)
+
+
+class DecoderConv1D(nn.Module):
+    def __init__(self, t_len: int, op_size: int):
+        super().__init__()
+        self.layer = nn.Conv1d(op_size, t_len, 1)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 1)
+        x = self.layer(x)
+        return x.permute(0, 2, 1)
+
+
+class EncoderConv2D(nn.Module):
+    def __init__(self, t_len: int, op_size: int):
+        super().__init__()
+        self.layer = nn.Conv2d(t_len, op_size, 1)
+
+    def forward(self, x):
+        x = x.permute(0, 3, 1, 2)
+        x = self.layer(x)
+        return x.permute(0, 2, 3, 1)
+
+
+class DecoderConv2D(nn.Module):
+    def __init__(self, t_len: int, op_size: int):
+        super().__init__()
+        self.layer = nn.Conv2d(op_size, t_len, 1)
+
+    def forward(self, x):
+        x = x.permute(0, 3, 1, 2)
+        x = self.layer(x)
+        return x.permute(0, 2, 3, 1)
+
+
+class KoopmanOperator1D(nn.Module):
+    def __init__(self, op_size: int, modes_x: int = 16):
+        super().__init__()
+        self.op_size = int(op_size)
+        self.modes_x = int(modes_x)
+        if min(self.op_size, self.modes_x) < 1:
+            raise ValueError("op_size and modes_x must be positive")
+        scale = 1.0 / (self.op_size * self.op_size)
+        self.koopman_matrix = nn.Parameter(
+            scale * torch.rand(self.op_size, self.op_size, self.modes_x, dtype=torch.cfloat)
+        )
+
+    def time_marching(self, input_tensor, weights):
+        return torch.einsum("btx,tfx->bfx", input_tensor, weights)
+
+    def forward(self, x):
+        x_ft = torch.fft.rfft(x)
+        out_ft = torch.zeros(x_ft.shape, dtype=torch.cfloat, device=x.device)
+        modes_x = min(self.modes_x, x_ft.shape[-1])
+        out_ft[:, :, :modes_x] = self.time_marching(
+            x_ft[:, :, :modes_x],
+            self.koopman_matrix[:, :, :modes_x],
+        )
+        return torch.fft.irfft(out_ft, n=x.size(-1))
+
+
+class KNO1D(nn.Module):
+    def __init__(
+        self,
+        encoder: nn.Module,
+        decoder: nn.Module,
+        op_size: int,
+        modes_x: int = 16,
+        decompose: int = 4,
+        linear_type: bool = True,
+        normalization: bool = False,
+    ):
+        super().__init__()
+        self.op_size = int(op_size)
+        self.decompose = int(decompose)
+        if self.decompose < 1:
+            raise ValueError("decompose must be positive")
+        self.enc = encoder
+        self.dec = decoder
+        self.koopman_layer = KoopmanOperator1D(self.op_size, modes_x=modes_x)
+        self.w0 = nn.Conv1d(self.op_size, self.op_size, 1)
+        self.linear_type = bool(linear_type)
+        self.normalization = bool(normalization)
+        if self.normalization:
+            self.norm_layer = nn.BatchNorm1d(self.op_size)
+
+    def forward(self, x):
+        x_reconstruct = self.dec(torch.tanh(self.enc(x)))
+        x = torch.tanh(self.enc(x)).permute(0, 2, 1)
+        x_w = x
+        for _ in range(self.decompose):
+            x1 = self.koopman_layer(x)
+            x = x + x1 if self.linear_type else torch.tanh(x + x1)
+        shortcut = self.w0(x_w)
+        if self.normalization:
+            shortcut = self.norm_layer(shortcut)
+        x = torch.tanh(shortcut + x).permute(0, 2, 1)
+        return self.dec(x), x_reconstruct
+
+
+class KoopmanOperator2D(nn.Module):
+    def __init__(self, op_size: int, modes_x: int = 12, modes_y: int = 12):
+        super().__init__()
+        self.op_size = int(op_size)
+        self.modes_x = int(modes_x)
+        self.modes_y = int(modes_y)
+        if min(self.op_size, self.modes_x, self.modes_y) < 1:
+            raise ValueError("op_size, modes_x and modes_y must be positive")
+        scale = 1.0 / (self.op_size * self.op_size)
+        self.koopman_matrix = nn.Parameter(
+            scale * torch.rand(
+                self.op_size,
+                self.op_size,
+                self.modes_x,
+                self.modes_y,
+                dtype=torch.cfloat,
+            )
+        )
+
+    def time_marching(self, input_tensor, weights):
+        return torch.einsum("btxy,tfxy->bfxy", input_tensor, weights)
+
+    def forward(self, x):
+        x_ft = torch.fft.rfft2(x)
+        out_ft = torch.zeros(x_ft.shape, dtype=torch.cfloat, device=x.device)
+        modes_x = min(self.modes_x, max(1, x_ft.shape[-2] // 2))
+        modes_y = min(self.modes_y, x_ft.shape[-1])
+        weights = self.koopman_matrix[:, :, :modes_x, :modes_y]
+        out_ft[:, :, :modes_x, :modes_y] = self.time_marching(
+            x_ft[:, :, :modes_x, :modes_y],
+            weights,
+        )
+        out_ft[:, :, -modes_x:, :modes_y] = self.time_marching(
+            x_ft[:, :, -modes_x:, :modes_y],
+            weights,
+        )
+        return torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
+
+
+class KNO2D(nn.Module):
+    def __init__(
+        self,
+        encoder: nn.Module,
+        decoder: nn.Module,
+        op_size: int,
+        modes_x: int = 12,
+        modes_y: int = 12,
+        decompose: int = 6,
+        linear_type: bool = True,
+        normalization: bool = False,
+    ):
+        super().__init__()
+        self.op_size = int(op_size)
+        self.decompose = int(decompose)
+        if self.decompose < 1:
+            raise ValueError("decompose must be positive")
+        self.enc = encoder
+        self.dec = decoder
+        self.koopman_layer = KoopmanOperator2D(self.op_size, modes_x=modes_x, modes_y=modes_y)
+        self.w0 = nn.Conv2d(self.op_size, self.op_size, 1)
+        self.linear_type = bool(linear_type)
+        self.normalization = bool(normalization)
+        if self.normalization:
+            self.norm_layer = nn.BatchNorm2d(self.op_size)
+
+    def forward(self, x):
+        x_reconstruct = self.dec(torch.tanh(self.enc(x)))
+        x = torch.tanh(self.enc(x)).permute(0, 3, 1, 2)
+        x_w = x
+        for _ in range(self.decompose):
+            x1 = self.koopman_layer(x)
+            x = x + x1 if self.linear_type else torch.tanh(x + x1)
+        shortcut = self.w0(x_w)
+        if self.normalization:
+            shortcut = self.norm_layer(shortcut)
+        x = torch.tanh(shortcut + x).permute(0, 2, 3, 1)
+        return self.dec(x), x_reconstruct
+
+
+class KNO2DNavierStokes(nn.Module):
+    """KNO2D adapter for flattened Navier-Stokes batches.
+
+    ``pos`` is accepted for API parity with FNO but is not used by the Koopman
+    operator. ``fx`` is shaped ``(B, H*W, t_in*out_dim)`` and the output is
+    shaped ``(B, H*W, out_dim)`` for one-step autoregressive rollout.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        spatial_shape,
+        op_size: int = 32,
+        modes_x: int = 12,
+        modes_y: int = 12,
+        decompose: int = 6,
+        linear_type: bool = True,
+        normalization: bool = False,
+    ):
+        super().__init__()
+        if int(output_channels) != 1:
+            raise ValueError("KNO2DNavierStokes currently expects output_channels=1.")
+        self.input_channels = int(input_channels)
+        self.output_channels = int(output_channels)
+        self.spatial_shape = tuple(int(v) for v in spatial_shape)
+        if len(self.spatial_shape) != 2 or min(self.spatial_shape) < 1:
+            raise ValueError("spatial_shape must contain two positive dimensions")
+        if self.input_channels < 1:
+            raise ValueError("input_channels must be positive")
+        encoder = EncoderConv2D(self.input_channels, int(op_size))
+        decoder = DecoderConv2D(self.output_channels, int(op_size))
+        self.kno = KNO2D(
+            encoder=encoder,
+            decoder=decoder,
+            op_size=int(op_size),
+            modes_x=int(modes_x),
+            modes_y=int(modes_y),
+            decompose=int(decompose),
+            linear_type=bool(linear_type),
+            normalization=bool(normalization),
+        )
+
+    def forward(self, pos, fx):
+        batch_size, num_points, channels = fx.shape
+        height, width = self.spatial_shape
+        if pos.shape[:2] != (batch_size, num_points) or pos.shape[-1] != 2:
+            raise ValueError(
+                f"Expected pos [B, {num_points}, 2], got {tuple(pos.shape)}"
+            )
+        if num_points != height * width:
+            raise ValueError(f"Expected {height * width} points, got {num_points}.")
+        if channels != self.input_channels:
+            raise ValueError(f"Expected {self.input_channels} channels, got {channels}.")
+        x = fx.reshape(batch_size, height, width, channels)
+        pred, _ = self.kno(x)
+        return pred.reshape(batch_size, num_points, self.output_channels)

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,210 @@
+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.kno import KNO2DNavierStokes  # 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_datapipe = checkpoint.get(
+        "datapipe_config", checkpoint.get("data_config")
+    )
+    if checkpoint_datapipe is None:
+        raise KeyError("Checkpoint contains no datapipe configuration")
+    checkpoint_data = checkpoint_datapipe["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],
+    datapipe_config: dict[str, Any],
+) -> KNO2DNavierStokes:
+    data_config = datapipe_config["data"]
+    return KNO2DNavierStokes(
+        input_channels=int(
+            model_config.get(
+                "input_channels",
+                int(data_config["t_in"]) * int(data_config["out_dim"]),
+            )
+        ),
+        output_channels=int(
+            model_config.get("output_channels", data_config["out_dim"])
+        ),
+        spatial_shape=spatial_shape,
+        op_size=int(model_config["op_size"]),
+        modes_x=int(model_config["modes_x"]),
+        modes_y=int(model_config["modes_y"]),
+        decompose=int(model_config["decompose"]),
+        linear_type=bool(model_config["linear_type"]),
+        normalization=bool(model_config["normalization"]),
+    )
+
+
+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, datapipe_config).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,265 @@
+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.kno import KNO2DNavierStokes  # 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["input_channels"] = int(data["t_in"]) * int(data["out_dim"])
+    model["output_channels"] = 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"],
+        "op_size": model["op_size"],
+        "modes_x": model["modes_x"],
+        "modes_y": model["modes_y"],
+        "decompose": model["decompose"],
+        "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]
+) -> KNO2DNavierStokes:
+    return KNO2DNavierStokes(
+        input_channels=int(model_config["input_channels"]),
+        output_channels=int(model_config["output_channels"]),
+        spatial_shape=spatial_shape,
+        op_size=int(model_config["op_size"]),
+        modes_x=int(model_config["modes_x"]),
+        modes_y=int(model_config["modes_y"]),
+        decompose=int(model_config["decompose"]),
+        linear_type=bool(model_config["linear_type"]),
+        normalization=bool(model_config["normalization"]),
+    )
+
+
+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/kno_navier_stokes.pt

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