Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,65 @@
 *.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
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.m4a filter=lfs diff=lfs merge=lfs -text
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.opus filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.mov filter=lfs diff=lfs merge=lfs -text
+*.avi filter=lfs diff=lfs merge=lfs -text
+*.mkv filter=lfs diff=lfs merge=lfs -text
+*.webm filter=lfs diff=lfs merge=lfs -text
+*.obj filter=lfs diff=lfs merge=lfs -text
+*.glb filter=lfs diff=lfs merge=lfs -text
+*.gltf filter=lfs diff=lfs merge=lfs -text
+*.ply filter=lfs diff=lfs merge=lfs -text
+*.stl filter=lfs diff=lfs merge=lfs -text
+*.fbx filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,128 @@
+---
+frameworks:
+  - pytorch
+language:
+  - en
+license: apache-2.0
+tags:
+  - OneScience
+  - FactFormer
+  - neural-operator
+  - factorized-attention
+  - computational-fluid-dynamics
+  - Kolmogorov-flow
+tasks:
+  - time-series-prediction
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FactFormer</span>
+  </strong>
+</p>
+
+# Model Overview
+
+FactFormer (Factorized Transformer) is a Transformer architecture for PDE surrogate modeling. This model predicts the temporal evolution of the vorticity field in two-dimensional Kolmogorov flow, using the first 10 time steps to predict the subsequent 16 by default.
+
+Paper: [Scalable Transformer for PDE Surrogate Modeling](https://arxiv.org/abs/2305.17560)
+
+# Model Description
+
+FactFormer factorizes global attention on a two-dimensional regular grid into attention operations along the two spatial axes. This preserves long-range spatial interactions while reducing computational and memory costs. The model predicts future states in blocks through latent propagation and rolls out the complete forecast window autoregressively.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Vorticity time-series prediction | Predict future Kolmogorov-flow states from historical two-dimensional vorticity fields |
+| PDE surrogate modeling | Learn mappings from historical to future physical fields on regular grids |
+| Factorized-attention research | Evaluate axial attention for high-resolution physical-field modeling |
+| 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, although training with the default configuration 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/FactFormer --local_dir ./FactFormer
+cd FactFormer
+```
+
+### 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 two-dimensional Kolmogorov-flow data with an original array shape of `[120, 320, 256, 256]`, representing trajectories, time steps, and the two spatial dimensions, respectively. Because the data file is large, it is read on demand using memory mapping.
+
+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/Kolmogorov_flow_2d --local_dir ./data
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+The default weights are saved to `weight/factformer_kolmogorov.pt`.
+
+### Model Weights
+
+This repository provides weights trained on the two-dimensional Kolmogorov-flow 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/factformer_kolmogorov.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
+
+- Li, Shu, and Barati Farimani. [Scalable Transformer for PDE Surrogate Modeling](https://arxiv.org/abs/2305.17560).
+- The Kolmogorov-flow data is used for two-dimensional vorticity prediction at Re=1000.
+- 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,61 @@
+root:
+  common:
+    device: "auto"
+    seed: 1234
+
+  datapipe:
+    source:
+      data_dir: "/public/share/sugonhpcapp01/onestore/onedatasets/Kolmogorov_flow_2d"
+      file_name: "kf_2d_re1000_256_120seed.npy"
+    data:
+      train_num: 100
+      test_num: 20
+      resolution: 128
+      interval: 2
+      t_in: 10
+      t_out: 16
+      out_dim: 1
+      normalize: true
+      val_mode: false
+      split_seed: 1234
+      stats_file: "data/km2d_stat.npz"
+      stats_max_trajectories: 8
+    dataloader:
+      batch_size: 4
+      num_workers: 2
+      pin_memory: true
+
+  model:
+    in_dim: 10
+    out_dim: 1
+    hidden_dim: 128
+    depth: 4
+    heads: 8
+    mlp_ratio: 2
+    dropout: 0.0
+    activation: "gelu"
+    include_pos: true
+    space_dim: 2
+    latent_multiplier: 2.0
+    max_latent_steps: 4
+
+  training:
+    epochs: 50
+    lr: 0.0005
+    weight_decay: 0.0001
+    step_size: 20
+    gamma: 0.5
+    eval_interval: 1
+    patience: 20
+    train_latent_steps: 4
+    max_grad_norm: 1.0
+    max_train_batches: null
+    max_eval_batches: null
+    weight_dir: "weight"
+    checkpoint_name: "factformer_kolmogorov.pt"
+
+  inference:
+    result_dir: "result"
+    num_samples: 4
+    figure_name: "prediction_sample.png"
+    tensor_name: "prediction_sample.pt"

configuration.json

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

data/km2d_stat.npz

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

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/FactFormer weight/factformer_kolmogorov.pt --local_dir ./

model/factformer.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+
+def _activation(name: str) -> nn.Module:
+    activations = {
+        "gelu": nn.GELU,
+        "relu": nn.ReLU,
+        "silu": nn.SiLU,
+        "swish": nn.SiLU,
+        "tanh": nn.Tanh,
+    }
+    try:
+        return activations[name.lower()]()
+    except KeyError as error:
+        raise ValueError(f"Unsupported activation: {name}") from error
+
+
+class StandardMLP(nn.Module):
+    """Pointwise MLP with state-dict names compatible with OneScience."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        hidden_dims: Sequence[int],
+        activation: str = "gelu",
+        use_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        dimensions = [int(input_dim), *(int(value) for value in hidden_dims), int(output_dim)]
+        self.layers = nn.ModuleList(
+            nn.Linear(in_features, out_features, bias=use_bias)
+            for in_features, out_features in zip(dimensions[:-1], dimensions[1:])
+        )
+        self.activation = _activation(activation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers[:-1]:
+            x = self.activation(layer(x))
+        return self.layers[-1](x)
+
+
+class OneMlp(nn.Module):
+    """Small compatibility wrapper used by released OneScience checkpoints."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.mlp = StandardMLP(**kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class _PoolingReducer(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.to_in = nn.Linear(in_dim, hidden_dim, bias=False)
+        self.out_ffn = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.to_in(x)
+        if x.ndim > 3:
+            x = x.mean(dim=tuple(range(2, x.ndim - 1)))
+        return self.out_ffn(x)
+
+
+class _SwapGridAxes(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.permute(0, 2, 1, 3)
+
+
+class _FactAttnWeight(nn.Module):
+    def __init__(self, heads: int, dim_head: int, dropout: float) -> None:
+        super().__init__()
+        self.dim_head = int(dim_head)
+        self.heads = int(heads)
+        self.scale = self.dim_head**-0.5
+        self.softmax = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+        self.to_q = nn.Linear(self.dim_head, self.dim_head, bias=False)
+        self.to_k = nn.Linear(self.dim_head, self.dim_head, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, length, _ = x.shape
+        x = x.reshape(batch, length, self.heads, self.dim_head)
+        x = x.permute(0, 2, 1, 3).contiguous()
+        query = self.to_q(x)
+        key = self.to_k(x)
+        return self.softmax(torch.matmul(query, key.transpose(-1, -2)) * self.scale)
+
+
+class FactAttention2D(nn.Module):
+    """Factorized axial attention for a structured two-dimensional grid."""
+
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        dropout: float,
+        shapelist: Sequence[int],
+    ) -> None:
+        super().__init__()
+        if len(shapelist) != 2:
+            raise ValueError("FactAttention2D requires shapelist=(height, width)")
+        if dim != heads * dim_head:
+            raise ValueError("dim must equal heads * dim_head")
+        self.dim_head = int(dim_head)
+        self.heads = int(heads)
+        self.H, self.W = (int(value) for value in shapelist)
+        self.attn_x = _FactAttnWeight(self.heads, self.dim_head, dropout)
+        self.attn_y = _FactAttnWeight(self.heads, self.dim_head, dropout)
+        self.to_v = nn.Linear(self.dim_head, self.dim_head, bias=False)
+        self.to_x = nn.Sequential(_PoolingReducer(dim, dim, dim))
+        self.to_y = nn.Sequential(_SwapGridAxes(), _PoolingReducer(dim, dim, dim))
+        self.to_out = nn.Sequential(nn.Linear(dim, dim), nn.Dropout(dropout))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, length, channels = x.shape
+        if length != self.H * self.W:
+            raise ValueError(
+                f"Sequence length {length} does not match grid {self.H}x{self.W}"
+            )
+        grid = x.reshape(batch, self.H, self.W, channels).contiguous()
+        values = self.to_v(
+            grid.reshape(batch, self.H, self.W, self.heads, self.dim_head)
+        )
+        values = values.permute(0, 3, 1, 2, 4).contiguous()
+        result_x = torch.einsum(
+            "bhij,bhjmc->bhimc", self.attn_x(self.to_x(grid)), values
+        )
+        result_y = torch.einsum(
+            "bhlm,bhimc->bhilc", self.attn_y(self.to_y(grid)), result_x
+        )
+        result = result_y.permute(0, 2, 3, 1, 4).contiguous()
+        result = result.reshape(batch, length, channels)
+        return self.to_out(result)
+
+
+class OneAttention(nn.Module):
+    """Compatibility wrapper retaining the original checkpoint hierarchy."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.attentioner = FactAttention2D(**kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.attentioner(x)
+
+
+class FactformerBlock(nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        dropout: float,
+        activation: str,
+        mlp_ratio: int,
+        spatial_shape: Sequence[int],
+    ) -> None:
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(hidden_dim)
+        self.Attn = OneAttention(
+            dim=hidden_dim,
+            heads=num_heads,
+            dim_head=hidden_dim // num_heads,
+            dropout=dropout,
+            shapelist=spatial_shape,
+        )
+        self.ln_2 = nn.LayerNorm(hidden_dim)
+        self.mlp = OneMlp(
+            input_dim=hidden_dim,
+            output_dim=hidden_dim,
+            hidden_dims=[hidden_dim * mlp_ratio],
+            activation=activation,
+            use_bias=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.Attn(self.ln_1(x)) + x
+        return self.mlp(self.ln_2(x)) + x
+
+
+class OneTransformer(nn.Module):
+    """Compatibility wrapper retaining the original checkpoint hierarchy."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.transformer = FactformerBlock(**kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.transformer(x)
+
+
+class FactFormer2D(nn.Module):
+    """FactFormer neural operator for structured 2D flow fields.
+
+    ``pos`` has shape ``(B, H*W, 2)`` and ``fx`` has shape
+    ``(B, H*W, t_in*out_dim)``. The output contains ``latent_steps`` future
+    fields concatenated along the final dimension.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        spatial_shape: Sequence[int],
+        hidden_dim: int = 128,
+        depth: int = 4,
+        heads: int = 8,
+        mlp_ratio: int = 2,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        include_pos: bool = True,
+        space_dim: int = 2,
+        latent_multiplier: float = 2.0,
+        max_latent_steps: int = 4,
+    ) -> 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.depth = int(depth)
+        self.heads = int(heads)
+        self.include_pos = bool(include_pos)
+        self.space_dim = int(space_dim)
+        self.latent_dim = int(self.hidden_dim * float(latent_multiplier))
+        self.max_latent_steps = int(max_latent_steps)
+
+        if len(self.spatial_shape) != 2:
+            raise ValueError("FactFormer2D requires a two-dimensional spatial shape")
+        if self.hidden_dim % self.heads:
+            raise ValueError("hidden_dim must be divisible by heads")
+        if self.max_latent_steps < 1:
+            raise ValueError("max_latent_steps must be positive")
+
+        input_dim = self.in_dim + (self.space_dim if self.include_pos else 0)
+        self.preprocess = OneMlp(
+            input_dim=input_dim,
+            output_dim=self.hidden_dim,
+            hidden_dims=[self.hidden_dim * 2],
+            activation=activation,
+            use_bias=True,
+        )
+        self.blocks = nn.ModuleList(
+            OneTransformer(
+                num_heads=self.heads,
+                hidden_dim=self.hidden_dim,
+                dropout=float(dropout),
+                activation=activation,
+                mlp_ratio=int(mlp_ratio),
+                spatial_shape=self.spatial_shape,
+            )
+            for _ in range(self.depth)
+        )
+        self.expand_latent = nn.Linear(self.hidden_dim, self.latent_dim, bias=False)
+        self.latent_time_embedding = nn.Parameter(
+            torch.randn(1, self.max_latent_steps, 1, self.latent_dim) * 0.02
+        )
+        self.propagator = OneMlp(
+            input_dim=self.latent_dim,
+            output_dim=self.latent_dim,
+            hidden_dims=[self.hidden_dim],
+            activation=activation,
+            use_bias=True,
+        )
+        self.to_out = OneMlp(
+            input_dim=self.latent_dim,
+            output_dim=self.out_dim,
+            hidden_dims=[self.hidden_dim],
+            activation=activation,
+            use_bias=True,
+        )
+        self.initialize_weights()
+
+    def initialize_weights(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        pos: torch.Tensor,
+        fx: torch.Tensor,
+        latent_steps: int = 1,
+    ) -> torch.Tensor:
+        if fx.ndim != 3:
+            raise ValueError(f"fx must have shape (B, N, C), got {tuple(fx.shape)}")
+        if pos.ndim == 2:
+            pos = pos.unsqueeze(0)
+        if pos.shape[0] == 1 and fx.shape[0] > 1:
+            pos = pos.expand(fx.shape[0], -1, -1)
+        if pos.shape[:2] != fx.shape[:2]:
+            raise ValueError(f"Incompatible pos and fx shapes: {pos.shape}, {fx.shape}")
+        if not 1 <= int(latent_steps) <= self.max_latent_steps:
+            raise ValueError(
+                f"latent_steps must be in [1, {self.max_latent_steps}], got {latent_steps}"
+            )
+
+        hidden = torch.cat((pos, fx), dim=-1) if self.include_pos else fx
+        hidden = self.preprocess(hidden)
+        for block in self.blocks:
+            hidden = block(hidden)
+
+        latent = self.expand_latent(hidden)
+        outputs = []
+        for step in range(int(latent_steps)):
+            latent = latent + self.latent_time_embedding[:, step]
+            latent = self.propagator(latent) + latent
+            outputs.append(self.to_out(latent))
+        return torch.cat(outputs, dim=-1)
+
+
+__all__ = ["FactFormer2D"]

scripts/common.py

@@ -0,0 +1,97 @@
+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,
+    max_latent_steps: int,
+) -> torch.Tensor:
+    predictions = []
+    remaining = int(t_out)
+    while remaining > 0:
+        step_count = min(int(max_latent_steps), remaining)
+        prediction = model(pos, state, latent_steps=step_count)
+        predictions.append(prediction)
+        consumed = step_count * int(out_dim)
+        state = torch.cat((state[..., consumed:], prediction), dim=-1)
+        remaining -= step_count
+    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,221 @@
+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.factformer import FactFormer2D  # noqa: E402
+from onescience.datapipes.cfd import KolmogorovFlow2DDatapipe  # 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()
+    )
+    stats_file = datapipe["data"].get("stats_file")
+    if stats_file:
+        datapipe["data"]["stats_file"] = str(
+            project_path(stats_file, PROJECT_ROOT).resolve()
+        )
+
+
+def checkpoint_datapipe_config(checkpoint: dict[str, Any]) -> dict[str, Any]:
+    if "datapipe_config" in checkpoint:
+        return checkpoint["datapipe_config"]
+    if "data_config" in checkpoint:
+        return checkpoint["data_config"]
+    raise KeyError("Checkpoint does not contain a datapipe configuration")
+
+
+def validate_checkpoint_config(
+    datapipe_config: dict[str, Any], checkpoint: dict[str, Any]
+) -> None:
+    checkpoint_data = checkpoint_datapipe_config(checkpoint)["data"]
+    current_data = datapipe_config["data"]
+    for key in (
+        "resolution",
+        "interval",
+        "t_in",
+        "t_out",
+        "out_dim",
+        "normalize",
+    ):
+        if current_data[key] != checkpoint_data[key]:
+            raise ValueError(
+                f"Config 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]
+) -> FactFormer2D:
+    return FactFormer2D(
+        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"]),
+        depth=int(model_config["depth"]),
+        heads=int(model_config["heads"]),
+        mlp_ratio=int(model_config["mlp_ratio"]),
+        dropout=float(model_config["dropout"]),
+        activation=str(model_config["activation"]),
+        include_pos=bool(model_config["include_pos"]),
+        space_dim=int(model_config["space_dim"]),
+        latent_multiplier=float(model_config["latent_multiplier"]),
+        max_latent_steps=int(model_config["max_latent_steps"]),
+    )
+
+
+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["training"]
+    inference = config["inference"]
+    datapipe_config = config["datapipe"]
+    num_samples = int(inference["num_samples"])
+    if num_samples < 1:
+        raise ValueError("inference.num_samples must be positive")
+    checkpoint_path = (
+        project_path(training["weight_dir"], PROJECT_ROOT).resolve()
+        / str(training["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)
+    validate_checkpoint_config(datapipe_config, checkpoint)
+    model_config = dict(checkpoint["model_config"])
+    spatial_shape = tuple(int(value) for value in checkpoint["spatial_shape"])
+    device = resolve_device(str(config["common"]["device"]))
+    result_dir = project_path(inference["result_dir"], PROJECT_ROOT).resolve()
+
+    print(f"Config: {config_path}")
+    print(f"Checkpoint: {checkpoint_path}")
+    print(
+        "Data: "
+        f"{Path(datapipe_config['source']['data_dir']) / datapipe_config['source']['file_name']}"
+    )
+    print(f"Device: {device}")
+
+    DistributedManager.initialize()
+    try:
+        datapipe = KolmogorovFlow2DDatapipe(
+            to_attr_dict(datapipe_config),
+            distributed=False,
+            normalizer_state=checkpoint.get("normalizer"),
+        )
+        test_loader, _ = datapipe.test_dataloader()
+        if tuple(datapipe.spatial_shape) != spatial_shape:
+            raise ValueError(
+                f"Checkpoint grid {spatial_shape} does not match data grid "
+                f"{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"]),
+                int(model_config["max_latent_steps"]),
+            )
+        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 / str(inference["tensor_name"])
+        image_path = result_dir / str(inference["figure_name"])
+        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,319 @@
+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 torch.nn.parallel import DistributedDataParallel as DDP
+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.factformer import FactFormer2D  # noqa: E402
+from onescience.datapipes.cfd import KolmogorovFlow2DDatapipe  # 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()
+    )
+    stats_file = data.get("stats_file")
+    if stats_file:
+        data["stats_file"] = str(project_path(stats_file, 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 = {
+        "train_num": data["train_num"],
+        "test_num": data["test_num"],
+        "resolution": data["resolution"],
+        "interval": data["interval"],
+        "t_in": data["t_in"],
+        "t_out": data["t_out"],
+        "batch_size": loader["batch_size"],
+        "hidden_dim": model["hidden_dim"],
+        "depth": model["depth"],
+        "heads": model["heads"],
+        "mlp_ratio": model["mlp_ratio"],
+        "max_latent_steps": model["max_latent_steps"],
+        "epochs": training["epochs"],
+        "eval_interval": training["eval_interval"],
+        "train_latent_steps": training["train_latent_steps"],
+    }
+    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")
+    if int(model["hidden_dim"]) % int(model["heads"]):
+        raise ValueError("hidden_dim must be divisible by heads")
+    if int(training["train_latent_steps"]) > int(model["max_latent_steps"]):
+        raise ValueError("train_latent_steps cannot exceed model.max_latent_steps")
+    if int(training["train_latent_steps"]) > int(data["t_out"]):
+        raise ValueError("train_latent_steps cannot exceed datapipe.data.t_out")
+    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]
+) -> FactFormer2D:
+    return FactFormer2D(
+        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"]),
+        depth=int(model_config["depth"]),
+        heads=int(model_config["heads"]),
+        mlp_ratio=int(model_config["mlp_ratio"]),
+        dropout=float(model_config["dropout"]),
+        activation=str(model_config["activation"]),
+        include_pos=bool(model_config["include_pos"]),
+        space_dim=int(model_config["space_dim"]),
+        latent_multiplier=float(model_config["latent_multiplier"]),
+        max_latent_steps=int(model_config["max_latent_steps"]),
+    )
+
+
+def evaluate(
+    model: torch.nn.Module,
+    loader: torch.utils.data.DataLoader,
+    device: torch.device,
+    datapipe: KolmogorovFlow2DDatapipe,
+    t_out: int,
+    out_dim: int,
+    max_latent_steps: int,
+    max_batches: int | None,
+    distributed: bool,
+    rank: int,
+) -> tuple[float, float]:
+    model.eval()
+    totals = torch.zeros(4, dtype=torch.float64, device=device)
+    with torch.no_grad():
+        iterator = tqdm(loader, desc="Evaluating", disable=rank != 0)
+        for batch_index, batch in enumerate(iterator):
+            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, max_latent_steps
+            )
+            prediction = datapipe.decode_solution(prediction)
+            target = datapipe.decode_solution(target)
+            totals[0] += F.mse_loss(prediction, target, reduction="sum")
+            totals[1] += target.numel()
+            totals[2] += relative_l2(prediction, target).sum()
+            totals[3] += target.shape[0]
+    if distributed:
+        torch.distributed.all_reduce(totals)
+    return (
+        (totals[0] / totals[1].clamp_min(1)).item(),
+        (totals[2] / totals[3].clamp_min(1)).item(),
+    )
+
+
+def main() -> None:
+    config_path = DEFAULT_CONFIG.resolve()
+    config = load_config(config_path)
+    prepare_config(config)
+    validate_config(config)
+    common = config["common"]
+    datapipe_config = config["datapipe"]
+    model_config = config["model"]
+    training = config["training"]
+    seed_everything(int(common["seed"]))
+
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    distributed = dist.world_size > 1
+    device = dist.device if str(common["device"]) == "auto" else resolve_device(
+        str(common["device"])
+    )
+    weight_dir = project_path(training["weight_dir"], PROJECT_ROOT).resolve()
+    checkpoint_path = weight_dir / str(training["checkpoint_name"])
+
+    if dist.rank == 0:
+        print(f"Config: {config_path}")
+        print(
+            "Data: "
+            f"{Path(datapipe_config['source']['data_dir']) / datapipe_config['source']['file_name']}"
+        )
+        print(f"Device: {device}")
+        print(
+            f"Samples: train={datapipe_config['data']['train_num']} "
+            f"test={datapipe_config['data']['test_num']} "
+            f"t_in={datapipe_config['data']['t_in']} "
+            f"t_out={datapipe_config['data']['t_out']}"
+        )
+
+    started = time.time()
+    try:
+        datapipe = KolmogorovFlow2DDatapipe(
+            to_attr_dict(datapipe_config), distributed=distributed
+        )
+        train_loader, train_sampler = datapipe.train_dataloader()
+        test_loader, _ = datapipe.test_dataloader()
+        spatial_shape = tuple(datapipe.spatial_shape)
+        model = build_model(model_config, spatial_shape).to(device)
+        if distributed:
+            device_ids = [dist.local_rank] if device.type == "cuda" else None
+            model = DDP(model, device_ids=device_ids)
+        if dist.rank == 0:
+            parameter_count = sum(parameter.numel() for parameter in model.parameters())
+            print(f"Spatial shape: {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"]),
+        )
+        data = datapipe_config["data"]
+        t_out = int(data["t_out"])
+        out_dim = int(data["out_dim"])
+        max_latent_steps = int(model_config["max_latent_steps"])
+        train_latent_steps = int(training["train_latent_steps"])
+        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):
+            if train_sampler is not None:
+                train_sampler.set_epoch(epoch)
+            model.train()
+            epoch_loss = 0.0
+            batch_count = 0
+            iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=dist.rank != 0)
+            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)[..., : train_latent_steps * out_dim]
+                optimizer.zero_grad(set_to_none=True)
+                prediction = model(pos, state, latent_steps=train_latent_steps)
+                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
+                if dist.rank == 0:
+                    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,
+                datapipe,
+                t_out,
+                out_dim,
+                max_latent_steps,
+                max_eval_batches,
+                distributed,
+                dist.rank,
+            )
+            if dist.rank == 0:
+                print(
+                    f"epoch={epoch:4d} train_mse={epoch_loss / batch_count:.6e} "
+                    f"val_mse={validation_mse:.6e} "
+                    f"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)
+                    model_to_save = model.module if isinstance(model, DDP) else model
+                    torch.save(
+                        {
+                            "epoch": epoch,
+                            "model_state": model_to_save.state_dict(),
+                            "model_config": to_plain_dict(model_config),
+                            "datapipe_config": to_plain_dict(datapipe_config),
+                            "training_config": to_plain_dict(training),
+                            "spatial_shape": 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
+
+            stop = torch.tensor(
+                [int(stale_evaluations >= int(training["patience"]))], device=device
+            )
+            if distributed:
+                torch.distributed.broadcast(stop, src=0)
+            if stop.item():
+                if dist.rank == 0:
+                    print("Early stopping triggered")
+                break
+
+        if dist.rank == 0:
+            print(f"Elapsed: {time.time() - started:.1f}s")
+    finally:
+        DistributedManager.cleanup()
+
+
+if __name__ == "__main__":
+    main()

weight/factformer_kolmogorov.pt

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