Upload folder using huggingface_hub

README.md

@@ -0,0 +1,158 @@
+---
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- large-scale fluid dynamics benchmark
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">CFDBench</span>
+  </strong>
+</p>
+
+# Model Overview
+
+CFDBench is a large-scale benchmark developed by researchers at Tsinghua University to evaluate machine-learning methods for computational fluid dynamics. It focuses on model generalization across different boundary conditions, fluid properties, and geometric configurations.
+
+Paper: [CFDBench: A Large-Scale Benchmark for Machine Learning Methods in Fluid Dynamics](https://arxiv.org/abs/2310.05963)
+
+# Model Description
+CFDBench is built from several representative computational fluid dynamics datasets spanning diverse boundary conditions, fluid properties, and geometries. It evaluates both flow-field prediction performance and the generalization capabilities of machine-learning methods.
+
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :---: |
+| CFD surrogate benchmarking | Compare the flow-field prediction capabilities of neural operators and other deep-learning models using consistent data splits and metrics |
+| Autoregressive flow evolution | Predict subsequent two-dimensional velocity fields step by step from the current grid-based field and assess error accumulation over multiple steps |
+| Non-autoregressive field queries | Predict velocity directly at target locations from operating parameters and spatiotemporal coordinates, enabling evaluation over long time horizons |
+
+
+## Supported Models
+
+| Type | `root.model.name` | Training Entry Point |
+| :--- | :--- | :--- |
+| Non-autoregressive | `ffn` | `python scripts/train.py` |
+| Non-autoregressive | `deeponet` | `python scripts/train.py` |
+| Autoregressive | `auto_ffn` | `python scripts/train_auto.py` |
+| Autoregressive | `auto_deeponet` | `python scripts/train_auto.py` |
+| Autoregressive | `auto_edeeponet` | `python scripts/train_auto.py` |
+| Autoregressive | `auto_deeponet_cnn` | `python scripts/train_auto.py` |
+| Autoregressive | `resnet` | `python scripts/train_auto.py` |
+| Autoregressive | `unet` | `python scripts/train_auto.py` |
+| Autoregressive | `fno` | `python scripts/train_auto.py` |
+
+# 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.
+- A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.
+
+### Download the Model Package
+```
+modelscope download --model OneScience/CFDBench --local_dir ./CFDBench
+cd CFDBench
+```
+### 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
+```
+
+### Training Data
+
+The OneScience community provides the `cfdbench` dataset for training. Download it with the command below and verify that the data path in `config/config.yaml` is configured correctly:
+
+```bash
+modelscope download --dataset OneScience/cfdbench --local_dir ./data
+```
+
+### Training
+**Autoregressive Training**
+
+The default configuration sets `root.model.name: fno`, so use the autoregressive training entry point:
+
+```bash
+python scripts/train_auto.py
+```
+
+
+To select another autoregressive model, edit `conf/config.yaml`:
+
+```yaml
+root:
+  model:
+    name: "auto_ffn"
+```
+
+**Non-autoregressive Training**
+```bash
+python scripts/train.py --model deeponet
+```
+### Model Weights
+This repository will provide pretrained CFDBench model weights in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference and Visualization
+
+The inference script automatically selects the task type from the current model name and reads the following path by default:
+
+```text
+./weight/<model.name>.pt
+```
+
+Default FNO inference and visualization:
+
+```bash
+python scripts/inference.py
+python scripts/result.py
+```
+
+Non-autoregressive models can likewise be selected through a command-line argument:
+
+```bash
+python scripts/inference.py --model ffn
+python scripts/result.py --model ffn
+```
+# 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
+
+- Original CFDBench paper: [CFDBench: A Large-Scale Benchmark for Machine Learning Methods in Fluid Dynamics](https://arxiv.org/abs/2310.05963)
+- Original CFDBench code repository: https://github.com/luo-yining/CFDBench
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope. Before public redistribution, verify the applicable licensing requirements of the upstream project.

conf/config.yaml

@@ -0,0 +1,78 @@
+root:
+  datapipe:
+    source:
+      data_dir: "./data/cfdbench"
+      data_name: "tube_prop_bc_geo"
+    verbose: true
+
+    data:
+      task_type: "auto"
+      norm_props: true
+      norm_bc: true
+      delta_time: 0.1
+      stable_state_diff: 0.001
+      split_ratios: [0.7, 0.15, 0.15]
+      seed: 0
+      num_rows: 8
+      num_cols: 8
+
+    dataloader:
+      batch_size: 2
+      eval_batch_size: 2
+      num_workers: 0
+
+  model:
+    name: "fno"
+    num_label_samples: 64
+    in_chan: 2
+    out_chan: 2
+    ffn_depth: 2
+    ffn_width: 64
+    autoffn_depth: 2
+    autoffn_width: 64
+    deeponet_width: 64
+    branch_depth: 2
+    trunk_depth: 2
+    act_fn: "relu"
+    act_scale_invariant: true
+    act_on_output: false
+    autoedeeponet_width: 64
+    autoedeeponet_depth: 2
+    autoedeeponet_act_fn: "relu"
+    fno_depth: 2
+    fno_hidden_dim: 16
+    fno_modes_x: 4
+    fno_modes_y: 4
+    unet_dim: 8
+    unet_insert_case_params_at: "input"
+    resnet_depth: 2
+    resnet_hidden_chan: 16
+    resnet_kernel_size: 7
+    resnet_padding: 3
+
+  training:
+    output_dir: "./result"
+    mode: "train_test"
+    group_by_model: true
+    device: "cpu"
+    num_epochs: 1
+    lr: 1.0e-3
+    lr_step_size: 20
+    lr_gamma: 0.9
+    log_interval: 1
+    eval_interval: 1
+    loss_name: "nmse"
+    checkpoint_name: "auto"
+
+  inference:
+    checkpoint_path: "auto"
+    output_dir: "./result/inference"
+    group_by_model: true
+    device: "cpu"
+    max_batches: 2
+    plot_interval: 1
+
+  fake_data:
+    num_cases_per_subset: 4
+    num_frames: 4
+    seed: 123

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other","model":"CFDBench","domain":"cfd"}

model/__init__.py

@@ -0,0 +1,181 @@
+from typing import Tuple
+
+import torch
+
+from model.auto_deeponet import AutoDeepONet
+from model.auto_deeponet_cnn import AutoDeepONetCnn
+from model.auto_edeeponet import AutoEDeepONet
+from model.auto_ffn import AutoFfn
+from model.deeponet import DeepONet
+from model.ffn import FfnModel
+from model.fno.fno2d import Fno2d
+from model.loss import loss_name_to_fn
+from model.resnet import ResNet
+from model.unet import UNet
+
+
+AUTO_MODEL_NAMES = {
+    "auto_ffn",
+    "auto_deeponet",
+    "auto_edeeponet",
+    "auto_deeponet_cnn",
+    "resnet",
+    "unet",
+    "fno",
+}
+STATIC_MODEL_NAMES = {"ffn", "deeponet"}
+ALL_MODEL_NAMES = AUTO_MODEL_NAMES | STATIC_MODEL_NAMES
+
+
+def infer_task_type(model_name: str) -> str:
+    if model_name in AUTO_MODEL_NAMES:
+        return "auto"
+    if model_name in STATIC_MODEL_NAMES:
+        return "static"
+    raise ValueError(f"Unknown CFDBench model.name={model_name!r}. Available: {sorted(ALL_MODEL_NAMES)}")
+
+
+def get_input_shapes(data_name: str, num_rows: int, num_cols: int) -> Tuple[int, int, int]:
+    if any(name in data_name for name in ("tube", "dam", "cylinder")):
+        n_rows = num_rows + 2
+        n_cols = num_cols + 1
+    elif "cavity" in data_name:
+        n_rows = num_rows
+        n_cols = num_cols
+    else:
+        raise ValueError(f"Unknown CFDBench problem in data_name: {data_name}")
+
+    if "cylinder" in data_name:
+        n_case_params = 8
+    elif any(name in data_name for name in ("cavity", "tube", "dam")):
+        n_case_params = 5
+    else:
+        raise ValueError(f"Unknown CFDBench parameter set in data_name: {data_name}")
+
+    return n_rows, n_cols, n_case_params
+
+
+def build_model(cfg) -> torch.nn.Module:
+    model_cfg = cfg.model
+    data_cfg = cfg.datapipe.data
+    source_cfg = cfg.datapipe.source
+    train_cfg = cfg.training
+    loss_fn = loss_name_to_fn(train_cfg.loss_name)
+    n_rows, n_cols, n_case_params = get_input_shapes(
+        data_name=source_cfg.data_name,
+        num_rows=data_cfg.num_rows,
+        num_cols=data_cfg.num_cols,
+    )
+
+    name = model_cfg.name
+    if name == "ffn":
+        widths = [n_case_params + 3] + [model_cfg.ffn_width] * model_cfg.ffn_depth + [1]
+        return FfnModel(
+            widths=widths,
+            loss_fn=loss_fn,
+            act_name=model_cfg.act_fn,
+            act_norm=model_cfg.act_scale_invariant,
+            act_on_output=model_cfg.act_on_output,
+            num_label_samples=model_cfg.num_label_samples,
+        )
+    if name == "deeponet":
+        return DeepONet(
+            branch_dim=n_case_params,
+            trunk_dim=3,
+            loss_fn=loss_fn,
+            width=model_cfg.deeponet_width,
+            trunk_depth=model_cfg.trunk_depth,
+            branch_depth=model_cfg.branch_depth,
+            act_name=model_cfg.act_fn,
+            act_norm=model_cfg.act_scale_invariant,
+            act_on_output=model_cfg.act_on_output,
+            num_label_samples=model_cfg.num_label_samples,
+        )
+    if name == "auto_ffn":
+        return AutoFfn(
+            input_field_dim=n_rows * n_cols,
+            num_case_params=n_case_params,
+            query_dim=2,
+            loss_fn=loss_fn,
+            width=model_cfg.autoffn_width,
+            depth=model_cfg.autoffn_depth,
+            act_name=model_cfg.act_fn,
+            act_norm=model_cfg.act_scale_invariant,
+            num_label_samples=model_cfg.num_label_samples,
+        )
+    if name == "auto_deeponet":
+        return AutoDeepONet(
+            branch_dim=n_rows * n_cols + n_case_params,
+            trunk_dim=2,
+            loss_fn=loss_fn,
+            width=model_cfg.deeponet_width,
+            trunk_depth=model_cfg.trunk_depth,
+            branch_depth=model_cfg.branch_depth,
+            act_name=model_cfg.act_fn,
+            act_norm=model_cfg.act_scale_invariant,
+            act_on_output=model_cfg.act_on_output,
+            num_label_samples=model_cfg.num_label_samples,
+        )
+    if name == "auto_edeeponet":
+        return AutoEDeepONet(
+            dim_branch1=n_rows * n_cols,
+            dim_branch2=n_case_params,
+            trunk_dim=2,
+            loss_fn=loss_fn,
+            width=model_cfg.autoedeeponet_width,
+            trunk_depth=model_cfg.autoedeeponet_depth,
+            branch_depth=model_cfg.autoedeeponet_depth,
+            act_name=model_cfg.autoedeeponet_act_fn,
+            num_label_samples=model_cfg.num_label_samples,
+        )
+    if name == "auto_deeponet_cnn":
+        if n_rows // 16 != 4 or n_cols // 16 != 4:
+            raise ValueError(
+                "auto_deeponet_cnn expects the post-padding grid to reduce to 4x4 after four 2x pools. "
+                "For tube/dam/cylinder fake data, set datapipe.data.num_rows=64 and num_cols=64."
+            )
+        return AutoDeepONetCnn(
+            in_chan=model_cfg.in_chan,
+            height=n_rows,
+            width=n_cols,
+            num_case_params=n_case_params,
+            query_dim=2,
+            loss_fn=loss_fn,
+            trunk_depth=model_cfg.trunk_depth,
+            act_name=model_cfg.act_fn,
+            act_norm=model_cfg.act_scale_invariant,
+            act_on_output=model_cfg.act_on_output,
+        )
+    if name == "resnet":
+        return ResNet(
+            in_chan=model_cfg.in_chan,
+            out_chan=model_cfg.out_chan,
+            n_case_params=n_case_params,
+            loss_fn=loss_fn,
+            hidden_chan=model_cfg.resnet_hidden_chan,
+            num_blocks=model_cfg.resnet_depth,
+            kernel_size=model_cfg.resnet_kernel_size,
+            padding=model_cfg.resnet_padding,
+        )
+    if name == "unet":
+        return UNet(
+            in_chan=model_cfg.in_chan,
+            out_chan=model_cfg.out_chan,
+            loss_fn=loss_fn,
+            n_case_params=n_case_params,
+            insert_case_params_at=model_cfg.unet_insert_case_params_at,
+            dim=model_cfg.unet_dim,
+        )
+    if name == "fno":
+        return Fno2d(
+            in_chan=model_cfg.in_chan,
+            out_chan=model_cfg.out_chan,
+            n_case_params=n_case_params,
+            loss_fn=loss_fn,
+            num_layers=model_cfg.fno_depth,
+            hidden_dim=model_cfg.fno_hidden_dim,
+            modes1=model_cfg.fno_modes_x,
+            modes2=model_cfg.fno_modes_y,
+        )
+
+    raise ValueError(f"Invalid CFDBench model.name={name!r}. Available: {sorted(ALL_MODEL_NAMES)}")

model/act_fn.py

@@ -0,0 +1,47 @@
+from torch import nn
+from torch import Tensor
+
+
+def get_act_fn(name: str, norm: bool = False) -> nn.Module:
+    if name == "relu":
+        fn = nn.ReLU()
+    elif name == "tanh":
+        fn = nn.Tanh()
+    elif name == "gelu":
+        fn = nn.GELU()
+    elif name == "swish":
+        fn = nn.SiLU()
+    else:
+        raise ValueError(f"Unknown activation function: {name}")
+    if norm:
+        fn = NormAct(fn)
+    return fn
+
+
+class NormAct(nn.Module):
+    """
+    Normalized Activation Function.
+
+    A wrapper around any activation function that normalizes the input
+    before applying the activation function, and then transforms the
+    output back to the original scale.
+    """
+    def __init__(self, act_fn: nn.Module):
+        super().__init__()
+        self.act_fn = act_fn
+
+    def forward(self, x: Tensor) -> Tensor:
+        '''
+        x: (b, h, w)
+        '''
+        num_dims = len(x.shape)
+        dims = tuple(range(1, num_dims))
+        # find the mean and std of each example in the batch
+        mean = x.mean(dim=dims, keepdim=True)
+        std = x.std(dim=dims, keepdim=True)
+        # normalize
+        x = (x - mean) / std
+        x = self.act_fn(x)
+        # Transform back to the original scale
+        x = x * std + mean
+        return x

model/auto_deeponet.py

@@ -0,0 +1,200 @@
+# coding: utf8
+
+"""
+Contains the implementation of autoregressive DeepONet
+"""
+
+from itertools import product
+from typing import List, Optional
+
+import torch
+from torch import nn, Tensor
+
+from .ffn import Ffn
+from .base_model import AutoCfdModel
+from .act_fn import get_act_fn
+from .loss import MseLoss
+
+
+class AutoDeepONet(AutoCfdModel):
+    """
+    Auto-regressive DeepONet for CFD.
+
+    Our task is different from the one that the original DeepONet. In the
+    original DeepONet, the input function (input to the branch net)
+    is the initial condition (IC), but here, we have a fixed (zero) IC.
+    Instead, we have different boundary conditions (BCs), but we also
+    want the model to predict the next time step given the current time step.
+    Ideally, we should have two different branch nets, one accepting the
+    BCs, one accepting the current time step.
+
+    Here, we assume that the current time step includes the information about
+    BCs (which are the values on the bounds), so we just feed
+    the current time step to one branch net.
+    """
+
+    def __init__(
+        self,
+        branch_dim: int,
+        trunk_dim: int,
+        loss_fn: MseLoss,
+        num_label_samples: int = 1000,
+        branch_depth: int = 4,
+        trunk_depth: int = 4,
+        width: int = 100,
+        act_name="relu",
+        act_norm: bool = False,
+        act_on_output: bool = False,
+    ):
+        """
+        Args:
+        - branch_dim: int, the dimension of the branch net input.
+        - trunk_dim: int, the dimension of the trunk net input.
+        """
+        super().__init__(loss_fn)
+        self.branch_dim = branch_dim
+        self.trunk_dim = trunk_dim
+        self.branch_depth = branch_depth
+        self.trunk_depth = trunk_depth
+        self.width = width
+        self.act_name = act_name
+        self.act_norm = act_norm
+        self.act_on_output = act_on_output
+        self.num_label_samples = num_label_samples
+
+        act_fn = get_act_fn(act_name, act_norm)
+        self.branch_dims = [branch_dim] + [width] * branch_depth
+        self.trunk_dims = [trunk_dim] + [width] * trunk_depth
+        self.branch_net = Ffn(
+            self.branch_dims,
+            act_fn=act_fn,
+            act_on_output=act_on_output,
+        )
+        self.trunk_net = Ffn(self.trunk_dims, act_fn=act_fn)
+        self.bias = nn.Parameter(torch.zeros(1))  # type: ignore
+
+    def forward(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        label: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        query_idxs: Optional[Tensor] = None,
+    ):
+        """
+        Here, we just randomly sample some points, and use the label values on
+        those points as the label.
+
+        ### Args
+        - inputs: (b, c, h, w)
+        - case_params: (b, p)
+        - labels: (b, c, h, w)
+        - query_point: (k, 2), k is the number of query points, each is
+            an (x, y) coordinate.
+        - masks: For future use.
+
+        ### Returns
+            Output: Tensor, if query_points is not None, the shape is (b, k).
+                Else, the shape is (b, c, h, w).
+
+        Notations:
+        - b: batch size
+        - c: number of channels
+        - h: height
+        - w: width
+        - p: number of case parameters
+        - k: number of query points
+        """
+        batch_size, num_chan, height, width = inputs.shape
+        # Only use the u channel
+        inputs = inputs[:, 0]  # (B, h, w)
+        # Flatten
+        flat_inputs = inputs.view(batch_size, -1)  # (B, h * w)
+
+        # Simple prepend physical properties to the input field.
+        # (B, h * w + 2)
+        flat_inputs = torch.cat([flat_inputs, case_params], dim=1)
+        x_branch = self.branch_net(flat_inputs)
+
+        if query_idxs is None:
+            query_idxs = torch.tensor(
+                list(product(range(height), range(width))),
+                dtype=torch.long,
+                device=flat_inputs.device,
+            )  # (h * w, 2)
+
+        # Input to the trunk net
+        x_trunk = (query_idxs.float() - 50) / 100  # (k, 2)
+        x_trunk = self.trunk_net(x_trunk)  # (k, p)
+        x_trunk = x_trunk.unsqueeze(0)  # (1, k, p)
+        x_branch = x_branch.unsqueeze(1)  # (b, 1, p)
+        preds = torch.sum(x_branch * x_trunk, dim=-1) + self.bias  # (b, k)
+
+        # Use values of the input field at query points as residuals
+        residuals = inputs[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+        preds += residuals
+
+        if label is not None:
+            label = label[:, 0]  # (B, 1, h, w)  # Predict only u
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            loss = self.loss_fn(labels=labels, preds=preds)  # (b, k)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return dict(preds=preds)
+
+    def generate(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor
+    ) -> Tensor:
+        """
+        x: (c, h, w) or (B, c, h, w)
+
+        Returns:
+            (b, c, h, w)
+        """
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+        batch_size, num_chan, height, width = inputs.shape
+        query_idxs = torch.tensor(
+            list(product(range(height), range(width))),
+            dtype=torch.long,
+            device=inputs.device,
+        )  # (h * w, 2)
+        # query_points = query_points / 100
+        # (b, 1, h * w)
+        preds = self.forward(
+            inputs, case_params=case_params, query_idxs=query_idxs, mask=mask
+        )["preds"]
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return preds
+
+    def generate_many(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor, steps: int
+    ) -> List[Tensor]:
+        """
+        Args:
+            x: (c, h, w) or (B, c, h, w)
+            mask: (h, w). 1 for interior, 0 for boundaries.
+            steps: int, number of steps to generate.
+
+        Returns:
+            list of tensors, each of shape (b, c, h, w)
+        """
+        assert len(inputs.shape) == len(case_params.shape) + 2
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+            case_params = case_params.unsqueeze(0)  # (1, p)
+            mask = mask.unsqueeze(0)  # (1, h, w)
+        assert inputs.shape[0] == case_params.shape[0]
+        cur_frame = inputs
+        preds = []
+        for _ in range(steps):
+            # (b, c, h, w)
+            cur_frame = self.generate(
+                inputs=cur_frame, case_params=case_params, mask=mask
+            )
+            preds.append(cur_frame)
+        return preds

model/auto_deeponet_cnn.py

@@ -0,0 +1,238 @@
+from itertools import product
+from typing import List, Optional
+
+import torch
+from torch import nn, Tensor
+
+from .ffn import Ffn
+from .base_model import AutoCfdModel
+from .act_fn import get_act_fn
+from .loss import MseLoss
+
+
+class CnnBranch(nn.Module):
+    def __init__(
+        self, in_chan: int, kernel_size: int, padding: int, depth: int = 4
+    ):
+        super().__init__()
+        self.in_chan = in_chan
+        self.in_conv = nn.Conv2d(
+            in_chan, 32, kernel_size=kernel_size, padding=padding
+        )
+        self.out_conv = nn.Conv2d(
+            32, 32, kernel_size=kernel_size, padding=padding
+        )
+
+        blocks = []
+        for i in range(depth):
+            blocks += [
+                nn.Conv2d(32, 32, kernel_size, padding=padding),
+                nn.MaxPool2d(2),
+                nn.ReLU(),
+            ]
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.in_conv(x)  # (b, 16, h, w)
+        x = self.blocks(x)  # (b, 32, h/16=4, w/16=4)
+        x = self.out_conv(x)  # (b, 32, 4, 4)
+        return x
+
+
+class AutoDeepONetCnn(AutoCfdModel):
+    """
+    Auto-regressive DeepONet for CFD.
+
+    Our task is different from the one that the original DeepONet. In the
+    original DeepONet, the input function (input to the branch net)
+    is the initial condition (IC), but here, we have a fixed (zero) IC.
+    Instead, we have different boundary conditions (BCs), but we also
+    want the model to predict the next time step given the current time step.
+    Ideally, we should have two different branch nets, one accepting the
+    BCs, one accepting the current time step.
+
+    Here, we assume that the current time step includes the information about
+    BCs (which are the values on the bounds), so we just feed
+    the current time step to one branch net.
+    """
+
+    def __init__(
+        self,
+        in_chan: int,
+        query_dim: int,
+        loss_fn: MseLoss,
+        height: int = 100,
+        width: int = 100,
+        num_case_params: int = 5,
+        trunk_depth: int = 4,
+        act_name="relu",
+        act_norm: bool = False,
+        act_on_output: bool = False,
+    ):
+        """
+        Args:
+        - branch_dim: int, the dimension of the branch net input.
+        - trunk_dim: int, the dimension of the trunk net input.
+        """
+        super().__init__(loss_fn)
+        self.in_chan = in_chan
+        self.query_dim = query_dim
+        self.num_case_params = num_case_params
+        self.trunk_depth = trunk_depth
+        self.height = height
+        self.width = width
+        self.act_name = act_name
+        self.act_norm = act_norm
+        self.act_on_output = act_on_output
+
+        act_fn = get_act_fn(act_name, act_norm)
+        self.trunk_dims = [query_dim] + [100] * trunk_depth + [4 * 4 * 32]
+        # + 1 for mask
+        self.branch_net = CnnBranch(
+            in_chan + 1 + num_case_params, kernel_size=5, padding=2
+        )
+        self.trunk_net = Ffn(
+            self.trunk_dims, act_fn=act_fn, act_on_output=False
+        )
+        self.out_ffn = Ffn(
+            [32 * 4 * 4] * 3 + [1], act_fn=act_fn, act_on_output=False
+        )
+        self.bias = nn.Parameter(torch.zeros(1))  # type: ignore
+
+    def forward(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        label: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        query_idxs: Optional[Tensor] = None,
+    ):
+        """
+        Here, we just randomly sample some points, and use the label values on
+        those points as the label.
+
+        x: (b, c, h, w)
+        labels: (b, c, h, w)
+        query_point: (k, 2), k is the number of query points, each is an (x, y)
+            coordinate.
+
+        Goal:
+        Input: [b, branch_dim + trunk_dim]
+        Output: [b, 1]
+        """
+
+        # Add mask to input as additional channels
+        if mask is not None:
+            if mask.dim() == 3:
+                mask = mask.unsqueeze(1)  # (B, 1, h, w)
+            inputs = torch.cat([inputs, mask], dim=1)  # (B, c + 1, h, w)
+
+        batch_size, num_chan, height, width = inputs.shape
+
+        # Only use the u channel
+        residuals = inputs[:, :]  # (B, c, h, w)
+
+        # Add case params as additional channels
+        case_params = case_params.unsqueeze(-1).unsqueeze(-1)  # (B, c, 1, 1)
+        # (B, n_params, h, w)
+        case_params = case_params.expand(
+            -1, -1, inputs.shape[-2], inputs.shape[-1]
+        )
+        inputs = torch.cat(
+            [inputs, case_params], dim=1
+        )  # (B, c + n_params, h, w)
+
+        x_branch = self.branch_net(inputs)  # (b, 32, h/16=4, w/16=4)
+        x_branch = x_branch.view(batch_size, -1)  # (b, 32 * 4 * 4 = 512)
+
+        if query_idxs is None:
+            query_idxs = torch.tensor(
+                list(product(range(height), range(width))),
+                dtype=torch.long,
+                device=x_branch.device,
+            )  # (h * w, 2)
+
+        # Input to the trunk net
+        x_trunk = (query_idxs.float() - 50) / 100  # (k, 2)
+        x_trunk = self.trunk_net(x_trunk)  # (k, p)
+        x_trunk = x_trunk.unsqueeze(0)  # (1, k, p)
+        x_branch = x_branch.unsqueeze(1)  # (b, 1, 32, p)
+        # preds = torch.sum(x_branch * x_trunk, dim=-1) + self.bias  # (b, k)
+        preds = x_branch * x_trunk  # (b, k, h)
+        preds = self.out_ffn(preds)  # (b, k, 1)
+        preds = preds.squeeze(-1)  # (b, k)
+
+        # Use values of the input field at query points as residuals
+        residuals = residuals[
+            :, 0, query_idxs[:, 0], query_idxs[:, 1]
+        ]  # (b, c, k)
+        # print(residuals.shape, preds.shape)
+        preds += residuals
+
+        if label is not None:
+            label = label[:, 0]  # (B, 1, h, w)  # Predict only u
+            # we have labels[i, j] = label[
+            #     i, query_points[i, j, 0], query_points[i, j, 1]]
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            loss = self.loss_fn(labels=labels, preds=preds)  # (b, k)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return dict(preds=preds)
+
+    def generate(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor
+    ) -> Tensor:
+        """
+        x: (c, h, w) or (B, c, h, w)
+
+        Returns:
+            (b, c, h, w)
+        """
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+            case_params = case_params.unsqueeze(0)  # (1, p)
+            mask = mask.unsqueeze(0)  # (1, h, w)
+        batch_size, num_chan, height, width = inputs.shape
+        query_idxs = torch.tensor(
+            list(product(range(height), range(width))),
+            dtype=torch.long,
+            device=inputs.device,
+        )  # (h * w, 2)
+        # query_points = query_points / 100
+        # (b, 1, h * w)
+        preds = self.forward(
+            inputs, query_idxs=query_idxs, case_params=case_params, mask=mask
+        )["preds"]
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return preds
+
+    def generate_many(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor, steps: int
+    ) -> List[Tensor]:
+        """
+        x: (c, h, w) or (B, c, h, w)
+        mask: (h, w). 1 for interior, 0 for boundaries.
+        steps: int, number of steps to generate.
+
+        Returns:
+            list of tensors, each of shape (b, c, h, w)
+        """
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+            case_params = case_params.unsqueeze(0)  # (1, p)
+            mask = mask.unsqueeze(0)  # (1, h, w)
+        cur_frame = inputs
+        p = inputs[:, -1:]
+        preds = []
+        for _ in range(steps):
+            # (b, c, h, w)
+            cur_frame = self.generate(
+                cur_frame, case_params=case_params, mask=mask
+            )
+            preds.append(cur_frame)
+            cur_frame = torch.cat([cur_frame, p], dim=1)
+        return preds

model/auto_edeeponet.py

@@ -0,0 +1,185 @@
+from itertools import product
+from typing import Dict, Optional, List
+
+import torch
+from torch import nn, Tensor
+
+from .base_model import AutoCfdModel
+from .ffn import Ffn
+from .act_fn import get_act_fn
+from .loss import MseLoss
+
+
+class AutoEDeepONet(AutoCfdModel):
+    """
+    EDeepONet for autoregressive generation. The two input functions are
+    the previous field (flattened) and the case parameters.
+
+    Branch net accepts the boundary and physics properties as inputs.
+    Trunk net accepts the query location (t, x, y) as input.
+    """
+
+    def __init__(
+        self,
+        dim_branch1: int,
+        dim_branch2: int,
+        trunk_dim: int,  # (x, y)
+        loss_fn: MseLoss,
+        num_label_samples: int = 1000,
+        branch_depth: int = 4,
+        trunk_depth: int = 4,
+        width: int = 100,
+        act_name: str = "relu",
+        act_norm: bool = False,
+        act_on_output: bool = False,
+    ):
+        super().__init__(loss_fn)
+        self.dim_branch1 = dim_branch1
+        self.dim_branch2 = dim_branch2
+        self.trunk_dim = trunk_dim
+        self.loss_fn = loss_fn
+        self.branch_depth = branch_depth
+        self.trunk_depth = trunk_depth
+        self.width = width
+        self.act_name = act_name
+        self.act_norm = act_norm
+        self.act_on_output = act_on_output
+        self.num_label_samples = num_label_samples
+
+        self.branch1_dims = [dim_branch1] + [width] * branch_depth
+        self.branch2_dims = [dim_branch2] + [width] * branch_depth
+        self.trunk_dims = [trunk_dim] + [width] * trunk_depth
+        print(self.trunk_dims)
+
+        act_fn = get_act_fn(act_name, act_norm)
+        self.branch1 = Ffn(
+            self.branch1_dims, act_fn=act_fn, act_on_output=self.act_on_output
+        )
+        self.branch2 = Ffn(
+            self.branch2_dims, act_fn=act_fn, act_on_output=self.act_on_output
+        )
+        # we will be using an entire frame as label.
+        self.trunk_net = Ffn(self.trunk_dims, act_fn=act_fn)
+
+        self.bias = nn.Parameter(torch.zeros(1))  # type: ignore
+
+    def forward(
+        self,
+        inputs: Tensor,  # (b, d1), previous field u(t-1)
+        case_params: Tensor,  # (b, d2), physical properties
+        label: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,  # NOTE: not used
+        query_idxs: Optional[Tensor] = None,
+    ) -> Dict[str, Tensor]:
+        """
+        A faster forward by using all the points in the frame (`label`) at
+        time step `t` as training examples.
+
+        Args:
+        - x_branch: (b, branch_dim), input to the branch net.
+        - t: (b), input to the trunk net, a batch of t
+        - label: (b, w, h), the frame to be predicted.
+        - query_idxs: (b, k, 2), the query locations.
+        """
+        batch_size, num_chan, height, width = inputs.shape
+        # Only use the u channel, because using more channels is to expensive
+        inputs = inputs[:, 0]  # (B, h, w)
+        # Flatten
+        flat_inputs = inputs.view(batch_size, -1)  # (B, h * w)
+
+        b1 = self.branch1(flat_inputs)  # (b, p)
+        b2 = self.branch2(case_params)  # (b, p)
+        x_branch = b1 * b2
+
+        if query_idxs is None:
+            query_idxs = torch.tensor(
+                list(product(range(height), range(width))),
+                dtype=torch.long,
+                device=flat_inputs.device,
+            )  # (h * w, 2)
+
+        # Normalize query location
+        x_trunk = (query_idxs.float() - 50) / 100  # (k, 2)
+        x_trunk = self.trunk_net(x_trunk)  # (k, p)
+        x_trunk = x_trunk.unsqueeze(0)  # (1, k, p)
+        x_branch = x_branch.unsqueeze(1)  # (b, 1, p)
+        preds = torch.sum(x_branch * x_trunk, dim=-1) + self.bias  # (b, k)
+
+        # Use values of input field at query points as residuals
+        residuals = inputs[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+        preds = preds + residuals
+
+        if label is not None:
+            # Use only the u channel
+            label = label[:, 0]  # (B, w, h)
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            assert (
+                preds.shape == labels.shape
+            ), f"{preds.shape}, {labels.shape}"
+            loss = self.loss_fn(preds=preds, labels=labels)  # (b, k)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+        preds.view(-1, 1, height, width)
+        return dict(preds=preds)
+
+    def generate(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        """
+        Generate one frame at time t.
+
+        Args:
+        - x: Tensor, (b, field dim)
+        - case_params: Tensor, (b, case params dim)
+        - t: Tensor, (b)
+        - height: int
+        - width: int
+
+        Returns:
+            (b, c, h, w)
+        """
+        batch_size, num_chan, height, width = inputs.shape
+        # Create 2D lattice of query points to infer the frame.
+        query_idxs = torch.tensor(
+            list(product(range(height), range(width))),
+            dtype=torch.long,
+            device=inputs.device,
+        )  # (h * w, 2)
+        # (b, 1, h * w)
+        preds = self.forward(
+            inputs=inputs, case_params=case_params, query_idxs=query_idxs
+        )["preds"]
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return preds
+
+    def generate_many(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor, steps: int
+    ) -> List[Tensor]:
+        """
+                x: (c, h, w) or (B, c, h, w)
+                mask: (h, w). 1 for interior, 0 for boundaries.
+                steps: int, number of steps to generate.
+        F
+                Returns:
+                    list of tensors, each of shape (b, c, h, w)
+        """
+        assert len(inputs.shape) == len(case_params.shape) + 2
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+            case_params = case_params.unsqueeze(0)  # (1, p)
+            mask = mask.unsqueeze(0)  # (1, h, w)
+        assert inputs.shape[0] == case_params.shape[0]
+        cur_frame = inputs
+        preds = []
+        for _ in range(steps):
+            # (b, c, h, w)
+            cur_frame = self.generate(
+                inputs=cur_frame, case_params=case_params, mask=None
+            )
+            preds.append(cur_frame)
+        return preds

model/auto_ffn.py

@@ -0,0 +1,178 @@
+from itertools import product
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+from .ffn import Ffn
+from .base_model import AutoCfdModel
+from .act_fn import get_act_fn
+from .loss import MseLoss
+
+
+class AutoFfn(AutoCfdModel):
+    """
+    Equivalent to autoregressive data-driven PINN.
+    """
+
+    def __init__(
+        self,
+        input_field_dim: int,
+        num_case_params: int,
+        query_dim: int,
+        loss_fn: MseLoss,
+        num_label_samples: int = 1000,
+        depth: int = 8,
+        width: int = 100,
+        act_norm: bool = False,
+        act_name="relu",
+    ):
+        """
+        Args:
+        - branch_dim: int, the dimension of the branch net input.
+        - trunk_dim: int, the dimension of the trunk net input.
+        """
+        super().__init__(loss_fn)
+        self.input_field_dim = input_field_dim
+        self.num_case_params = num_case_params
+        self.query_dim = query_dim
+        self.depth = depth
+        self.width = width
+        self.act_name = act_name
+        self.act_norm = act_norm
+        self.num_label_samples = num_label_samples
+
+        self.in_dim = input_field_dim + num_case_params + query_dim
+        act_fn = get_act_fn(act_name, act_norm)
+        self.widths = [self.in_dim] + [width] * depth + [1]
+        self.ffn = Ffn(
+            self.widths,
+            act_fn=act_fn,
+            act_on_output=False,
+        )
+
+    def forward(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        label: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,  # NOTE: Not used
+        query_idxs: Optional[Tensor] = None,
+    ):
+        """
+        Here, we just randomly sample some points, and use the label values on
+        those points as the label.
+
+        ### Parameters
+        - `inputs: Tensor` -- (b, c, h, w)
+        - `labels: Tensor` -- (b, c, h, w)
+        - `query_idxs: Tensor` -- (k, 2), k is the number of query points,
+            each is an (x, y) coordinate.
+        - `mask: Tensor` -- Not used.
+
+        ### Function
+            Input: [b, branch_dim + trunk_dim]
+            Output: [b, 1]
+        """
+        batch_size, _num_chan, height, width = inputs.shape
+
+        # Only use the u channel
+        inputs = inputs[:, 0]  # (B, h, w)
+        # Flatten
+        flat_inputs = inputs.view(batch_size, -1)  # (B, h * w)
+        flat_inputs = torch.cat(
+            [flat_inputs, case_params], dim=1
+        )  # (B, h * w + 2)
+
+        if query_idxs is None:
+            query_idxs = torch.tensor(
+                list(product(range(height), range(width))),
+                dtype=torch.long,
+                device=flat_inputs.device,
+            )  # (k=h * w, 2)
+
+        n_queries = query_idxs.shape[0]
+
+        # For each combination of (input, query_point), we have a sample.
+        # Repeat tensors such that we get (b * k) samples
+        # (b, k, h * w)
+        flat_inputs = flat_inputs.repeat(n_queries, 1)  # (b * k, h * w)
+        batch_query_idxs = query_idxs.repeat(batch_size, 1)  # (b * k, 2)
+
+        # (b * k, h * w + 2)
+        flat_inputs = torch.cat([flat_inputs, batch_query_idxs.float()], dim=1)
+
+        preds = self.ffn(flat_inputs)  # (b * k, 1)
+        preds = preds.view(batch_size, -1)  # (b, k)
+
+        # Use values of the input field at query points as residuals
+        residuals = inputs[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+        preds += residuals
+
+        if label is not None:
+            label = label[:, 0]  # (B, 1, h, w)  # Predict only u
+            # we have labels[i, j] = label[
+            #     i, query_points[i, j, 0], query_points[i, j, 1]]
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            loss = self.loss_fn(labels=labels, preds=preds)  # (b, k)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return dict(preds=preds)
+
+    def generate(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor
+    ) -> Tensor:
+        """
+        x: (c, h, w) or (B, c, h, w)
+
+        Returns:
+            (b, c, h, w)
+        """
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+        batch_size, num_chan, height, width = inputs.shape
+        query_idxs = torch.tensor(
+            list(product(range(height), range(width))),
+            dtype=torch.long,
+            device=inputs.device,
+        )  # (h * w, 2)
+        # query_points = query_points / 100
+        # (b, 1, h * w)
+        preds = self.forward(
+            inputs, query_idxs=query_idxs, case_params=case_params, mask=mask
+        )["preds"]
+        preds = preds.view(-1, 1, height, width)  # (b, 1, h, w)
+        return preds
+
+    def generate_many(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Tensor,
+        steps: int,
+    ) -> List[Tensor]:
+        """
+        x: (c, h, w) or (B, c, h, w)
+        mask: (h, w). 1 for interior, 0 for boundaries.
+        steps: int, number of steps to generate.
+
+        Returns:
+            list of tensors, each of shape (b, c, h, w)
+        """
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+            case_params = case_params.unsqueeze(0)  # (1, p)
+            mask = mask.unsqueeze(0)  # (1, h, w)
+        cur_frame = inputs
+        preds = []
+        for _ in range(steps):
+            # (b, c, h, w)
+            cur_frame = self.generate(
+                cur_frame, case_params=case_params, mask=mask
+            )
+            preds.append(cur_frame)
+        return preds

model/base_model.py

@@ -0,0 +1,81 @@
+from torch import nn, Tensor
+from typing import Optional, List
+
+from .loss import MseLoss
+
+
+class CfdModel(nn.Module):
+    """
+    Base model for all data-driven NN for CFD, learns the mapping from
+    conditions (physics properties, boundary conditions and geometry)
+    to the solution at a later time.
+    """
+    def __init__(self, loss_fn: MseLoss):
+        super().__init__()
+        self.loss_fn = loss_fn
+
+    def forward(
+        self,
+        x: Tensor,
+        mask: Optional[Tensor] = None,
+        label: Optional[Tensor] = None,
+        case_params: Optional[dict] = None,
+    ) -> dict:
+        raise NotImplementedError
+
+    def generate_one(
+        self,
+        case_params,
+        t: Tensor,
+        height: int,
+        width: int,
+        **kwargs,
+    ) -> Tensor:
+        """
+        Generate the frame at time step `time`, given the case parameters.
+        `case_params`.
+        """
+        raise NotImplementedError
+
+
+class AutoCfdModel(nn.Module):
+    """
+    A CFD model that generates the solution auto-regressively, one frame
+    at a time.
+    """
+
+    def __init__(self, loss_fn: nn.Module):
+        super().__init__()
+        self.loss_fn = loss_fn
+
+    def forward(
+        self,
+        inputs: Tensor,
+        label: Optional[Tensor] = None,
+        case_params: Optional[dict] = None,
+        mask: Optional[Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        raise NotImplementedError
+
+    def generate(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Tensor,
+        **kwargs,
+    ) -> Tensor:
+        raise NotImplementedError
+
+    def generate_many(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Tensor,
+        steps: int,
+        **kwargs,
+    ) -> List[Tensor]:
+        """
+        Given a frame `inputs`, generate the next `steps` frames.
+        """
+        raise NotImplementedError

model/deeponet.py

@@ -0,0 +1,257 @@
+from itertools import product
+from typing import Dict, Optional
+
+import torch
+from torch import nn, Tensor
+
+from .base_model import CfdModel
+from .ffn import Ffn
+from .act_fn import get_act_fn
+from .loss import MseLoss
+
+
+class DeepONet(CfdModel):
+    """
+    DeepONet for CFD.
+
+    Branch net accepts the boundary and physics properties as inputs.
+    Trunk net accepts the query location (t, x, y) as input.
+    """
+
+    def __init__(
+        self,
+        branch_dim: int,
+        trunk_dim: int,  # (t, x, y)
+        loss_fn: MseLoss,
+        num_label_samples: int = 1000,
+        branch_depth: int = 4,
+        trunk_depth: int = 3,
+        width: int = 100,
+        act_name: str = "relu",
+        act_norm: bool = False,
+        act_on_output: bool = False,
+    ):
+        """
+        Args:
+        - branch_dim: int, the dimension of the branch net input.
+        - trunk_dim: int, the dimension of the trunk net input.
+        """
+        super().__init__(loss_fn)
+        self.branch_dim = branch_dim
+        self.trunk_dim = trunk_dim
+        self.loss_fn = loss_fn
+        self.branch_depth = branch_depth
+        self.trunk_depth = trunk_depth
+        self.width = width
+        self.act_name = act_name
+        self.act_norm = act_norm
+        self.act_on_output = act_on_output
+        self.num_label_samples = num_label_samples
+
+        self.branch_dims = [branch_dim] + [width] * branch_depth
+
+        act_fn = get_act_fn(act_name, act_norm)
+        self.branch_net = Ffn(
+            self.branch_dims, act_fn=act_fn, act_on_output=self.act_on_output
+        )
+        # t and x, y uses separate fc layers because during training
+        # we will be using an entire frame as label.
+        self.fc_trunk_t = nn.Linear(1, width)
+        self.fc_trunk_xy = nn.Linear(2, width)
+        self.trunk_dims = [width] * trunk_depth
+        self.trunk_net = Ffn(self.trunk_dims, act_fn=act_fn)
+
+        # self.params = self.__init_params()
+        self.bias = nn.Parameter(torch.zeros(1))  # type: ignore
+        # self.__initialize()
+
+    # NOTE: This is no long used!!
+    # This is the implementation of the ordinary DeepONet, but it is very slow.
+    # Our `forward` is much faster because it re-uses the output
+    # of the branch net for different queries in one batch.
+    def forward_vanilla(
+        self,
+        x_branch: Tensor,
+        x_trunk: Tensor,
+        label: Optional[Tensor] = None,
+        query_idxs: Optional[Tensor] = None,
+    ):
+        """
+        NOTE: This is no long used!!
+        This is the implementation of the ordinary DeepONet, but it is very
+        slow! Our `forward` is much faster because it re-uses the output
+        of the branch net for different queries in one batch.
+
+        Args:
+        - x_branch: (b, branch_dim), input to the branch net.
+        - x_trunk: (b), input to the trunk net, a batch of (t, x, y)
+        - label: (b), the label for the query location.
+
+        Return:
+        - if label is None:
+            - preds: (b, k), the prediction for the query location.
+        - else:
+            - preds: (b, k), the prediction for the query location.
+            - loss: (b, k), the loss for the query location.
+        """
+
+        # Create k query locations.
+        # (b, k, 3), where each element is (t, x, y)
+        if query_idxs is None:
+            assert label is not None
+            height, width = label.shape[-2:]
+            query_idxs = torch.stack(
+                [
+                    torch.randint(
+                        0,
+                        height,
+                        (self.num_label_samples,),
+                        device=x_trunk.device,
+                    ),
+                    torch.randint(
+                        0,
+                        width,
+                        (self.num_label_samples,),
+                        device=x_trunk.device,
+                    ),
+                ],
+                dim=-1,
+            )  # (k, 2)
+        t = x_trunk.unsqueeze(1).float()  # (b, 1)
+        x_trunk_t = self.fc_trunk_t(t)  # (b, p)
+        # Normalize query location
+        x_trunk_xy = (
+            query_idxs.float() - 32.0
+        ) / 64.0  # (k, 2)  # TODO: update this
+        x_trunk_xy = self.fc_trunk_xy(x_trunk_xy)  # (k, p)
+        x_trunk_t = x_trunk_t.unsqueeze(1)  # (b, 1, p)
+        x_trunk_xy = x_trunk_xy.unsqueeze(0)  # (1, k, p)
+        x_trunk = x_trunk_t + x_trunk_xy  # (b, k, p)
+
+        x_branch = self.branch_net(x_branch)  # (b, p)
+        x_trunk = self.trunk_net(x_trunk)  # (b, k, p)
+
+        # x_trunk = x_trunk.unsqueeze(0)  # (1, k, p)
+        x_branch = x_branch.unsqueeze(1)  # (b, 1, p)
+        preds = torch.sum(x_branch * x_trunk, dim=-1) + self.bias  # (b)
+        print(preds.dtype)
+
+        if label is not None:
+            print(label.dtype)
+            # Use only the u channel
+            label = label[:, 0]  # (B, w, h)
+            # we have labels[i, j] = label[
+            #     i, query_points[i, j, 0], query_points[i, j, 1]]
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            # assert preds.shape == label.shape, f"{preds.shape} {label.shape}"
+            loss = self.loss_fn(preds=preds, labels=labels)  # (b, k)
+            print(labels.dtype)
+            print(loss.dtype)
+            return preds, loss
+        return preds
+
+    def forward(
+        self,
+        case_params: Tensor,  # Case parameters
+        t: Tensor,
+        label: Optional[Tensor] = None,
+        query_idxs: Optional[Tensor] = None,
+    ) -> Dict[str, Tensor]:
+        """
+        A faster forward by using all the points in the frame (`label`) at
+        time step `t` as training examples.
+
+        Args:
+        - x_branch: (b, branch_dim), input to the branch net.
+        - t: (b), input to the trunk net, a batch of t
+        - label: (b, w, h), the frame to be predicted.
+        - query_idxs: (b, k, 2), the query locations.
+        """
+        if query_idxs is None:
+            # Create k query locations on a lattice.
+            # (b, k, 2), where each element is (x, y)
+            assert label is not None
+            height, width = label.shape[-2:]
+            query_idxs = torch.stack(
+                [
+                    torch.randint(
+                        0,
+                        height,
+                        (self.num_label_samples,),
+                        device=label.device,
+                    ),
+                    torch.randint(
+                        0,
+                        width,
+                        (self.num_label_samples,),
+                        device=label.device,
+                    ),
+                ],
+                dim=-1,
+            )  # (k, 2)
+
+        x_trunk_t = self.fc_trunk_t(t)  # (b, p)
+        # Normalize query location
+        x_trunk_xy = query_idxs.float()  # (k, 2)
+        x_trunk_xy = self.fc_trunk_xy(x_trunk_xy)  # (k, p)
+        x_trunk_t = x_trunk_t.unsqueeze(1)  # (b, 1, p)
+        x_trunk_xy = x_trunk_xy.unsqueeze(0)  # (1, k, p)
+        x_trunk = x_trunk_t + x_trunk_xy  # (b, k, p)
+
+        case_params = self.branch_net(case_params)  # (b, p)
+        x_trunk = self.trunk_net(x_trunk)  # (b, k, p)
+
+        case_params = case_params.unsqueeze(1)  # (b, 1, p)
+        preds = torch.sum(case_params * x_trunk, dim=-1) + self.bias  # (b, k)
+
+        if label is not None:
+            # Only predict u
+            label = label[:, 0]  # (B, w, h)
+            # we have labels[i, j] = label[
+            #     i, query_points[i, j, 0], query_points[i, j, 1]]
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            assert (
+                preds.shape == labels.shape
+            ), f"{preds.shape}, {labels.shape}"
+            loss = self.loss_fn(preds=preds, labels=labels)  # (b, k)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+        return dict(
+            preds=preds,
+        )
+
+    def generate_one(
+        self, case_params: Tensor, t: Tensor, height: int, width: int
+    ) -> Tensor:
+        """
+        Generate one frame at time t.
+
+        Args:
+        - x_branch: Tensor, (b, branch_dim)
+        - t: Tensor, (b)
+        - height: int
+        - width: int
+
+        Returns:
+            (b, c, h, w)
+        """
+        if len(case_params.shape) == 1:
+            case_params = case_params.unsqueeze(0)
+        if len(t.shape) == 0:
+            t = t.unsqueeze(0).unsqueeze(0)
+        elif len(t.shape) == 1:
+            t = t.unsqueeze(0)
+
+        # Create 2D lattice of query points to infer the frame.
+        query_idxs = torch.tensor(
+            list(product(range(height), range(width))),
+            # dtype=torch.long,
+            device=case_params.device,
+        )  # (h * w, 2)
+        # query_points = query_points / 100
+        # (b, 1, h * w)
+        output = self.forward(case_params, t=t, query_idxs=query_idxs)["preds"]
+        output = output.view(-1, 1, height, width)  # (b, 1, h, w)
+        return output

model/ffn.py

@@ -0,0 +1,181 @@
+from typing import Optional, Dict, List
+from itertools import product
+
+from torch import nn, Tensor
+import torch
+
+from .base_model import CfdModel
+from .loss import MseLoss
+from .act_fn import get_act_fn
+
+
+class Ffn(nn.Module):
+    """
+    A general fully connected multi-layer neural network.
+    """
+
+    def __init__(
+        self, dims: list, act_fn: nn.Module, act_on_output: bool = False
+    ):
+        super().__init__()
+        self.dims = dims
+
+        layers = []
+        for i in range(len(dims) - 2):
+            layers.append(nn.Linear(dims[i], dims[i + 1]))
+            layers.append(act_fn)
+            # layers.append(NormAct(nn.Tanh()))
+        layers.append(nn.Linear(dims[-2], dims[-1]))
+        if act_on_output:
+            layers.append(act_fn)
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.layers(x)
+        return x
+
+
+class FfnModel(CfdModel):
+    """
+    Non-autoregressive FNN for data-driven CFD.
+
+    Branch net accepts the boundary and physics properties as inputs.
+    Trunk net accepts the query location (t, x, y) as input.
+    """
+
+    def __init__(
+        self,
+        loss_fn: MseLoss,
+        widths: List[int],
+        act_name: str = "relu",
+        act_norm: bool = True,
+        act_on_output: bool = False,
+        num_label_samples: int = 1000,
+    ):
+        """
+        Args:
+        - branch_dim: int, the dimension of the branch net input.
+        - trunk_dim: int, the dimension of the trunk net input.
+        """
+        super().__init__(loss_fn)
+        self.loss_fn = loss_fn
+        self.widths = widths
+        self.act_name = act_name
+        self.act_norm = act_norm
+        self.act_on_output = act_on_output
+        self.num_label_samples = num_label_samples
+
+        act_fn = get_act_fn(act_name, act_norm)
+        self.ffn = Ffn(
+            self.widths, act_fn=act_fn, act_on_output=self.act_on_output
+        )
+
+    def forward(
+        self,
+        case_params: Tensor,  # Case parameters
+        t: Tensor,
+        label: Optional[Tensor] = None,
+        query_idxs: Optional[Tensor] = None,
+    ) -> Dict[str, Tensor]:
+        """
+        A faster forward by using all the points in the frame (`label`) at
+        time step `t` as training examples.
+
+        Args:
+        - x_branch: (b, branch_dim), input to the branch net.
+        - t: (b), input to the trunk net, a batch of t
+        - label: (b, w, h), the frame to be predicted.
+        - query_idxs: (b, k, 2), the query locations.
+        """
+
+        batch_size, dim_in = case_params.shape
+
+        if query_idxs is None:
+            # Create k query locations on a lattice.
+            # (b, k, 2), where each element is (x, y)
+            assert label is not None
+            height, width = label.shape[-2:]
+            query_idxs = torch.stack(
+                [
+                    torch.randint(
+                        0,
+                        height,
+                        (self.num_label_samples,),
+                        device=label.device,
+                    ),
+                    torch.randint(
+                        0,
+                        width,
+                        (self.num_label_samples,),
+                        device=label.device,
+                    ),
+                ],
+                dim=-1,
+            )  # (k, 2)
+
+        # Concatenate query spatial coordinates with time coordinate
+        coords = query_idxs.unsqueeze(0)  # (1, k, 2)
+        coords = coords.repeat(batch_size, 1, 1)  # (b, k, 2)
+        num_queries = coords.shape[1]
+        t = t.unsqueeze(-1)  # (b, 1)
+        t = t.repeat(1, num_queries, 1)  # (b, k)
+        coords = torch.cat([coords, t], dim=-1)  # (b, k, 3)
+
+        # Concatenate case params with queries
+        case_params = case_params.unsqueeze(1)  # (b, 1, p)
+        case_params = case_params.repeat(1, num_queries, 1)  # (b, k, p)
+        inp = torch.cat([case_params, coords], dim=-1)  # (b, k, p + 3)
+        inp = inp.view(batch_size * num_queries, -1)  # (bk, p + 3)
+        preds = self.ffn(inp)  # (bk, out_dim)
+        preds = preds.view(batch_size, num_queries)  # (b, k)
+
+        if label is not None:
+            # Use only the u channel
+            label = label[:, 0]  # (B, w, h)
+            labels = label[:, query_idxs[:, 0], query_idxs[:, 1]]  # (b, k)
+            assert (
+                preds.shape == labels.shape
+            ), f"{preds.shape}, {labels.shape}"
+            loss = self.loss_fn(preds=preds, labels=labels)  # (b, k)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+        return dict(
+            preds=preds,
+        )
+
+    def generate_one(
+        self, case_params: Tensor, t: Tensor, height: int, width: int
+    ) -> Tensor:
+        """
+        Generate one frame at time t.
+
+        Args:
+        - x_branch: Tensor, (branch_dim)
+        - t: Tensor, (1,)
+        - height: int
+        - width: int
+
+        Returns:
+            (b, c, h, w)
+        """
+        if len(case_params.shape) == 1:
+            case_params = case_params.unsqueeze(0)
+        if len(t.shape) == 0:
+            t = t.unsqueeze(0).unsqueeze(0)
+        elif len(t.shape) == 1:
+            t = t.unsqueeze(0)
+
+        # Create 2D lattice of query points to infer the frame.
+        query_idxs = torch.tensor(
+            list(product(range(height), range(width))),
+            # dtype=torch.long,
+            device=case_params.device,
+        )  # (h * w, 2)
+
+        # query_points = query_points / 100
+        # (b, 1, h * w)
+        output = self.forward(case_params, t=t, query_idxs=query_idxs)["preds"]
+        output = output.view(-1, 1, height, width)  # (b, 1, h, w)
+        return output

model/fno/Adam.py

@@ -0,0 +1,198 @@
+import math
+import torch
+from torch import Tensor
+from typing import List
+from torch.optim.optimizer import Optimizer
+
+
+def adam(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    max_exp_avg_sqs: List[Tensor],
+    state_steps: List[int],
+    *,
+    amsgrad: bool,
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+    eps: float
+):
+    r"""Functional API that performs Adam algorithm computation.
+    See :class:`~torch.optim.Adam` for details.
+    """
+
+    for i, param in enumerate(params):
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        step = state_steps[i]
+
+        bias_correction1 = 1 - beta1**step
+        bias_correction2 = 1 - beta2**step
+
+        if weight_decay != 0:
+            grad = grad.add(param, alpha=weight_decay)
+
+        # Decay the first and second moment running average coefficient
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
+        if amsgrad:
+            # Maintains the maximum of all 2nd moment running avg. till now
+            torch.maximum(
+                max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i]
+            )
+            # Use the max. for normalizing running avg. of gradient
+            denom = (
+                max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)
+            ).add_(eps)
+        else:
+            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
+
+        step_size = lr / bias_correction1
+
+        param.addcdiv_(exp_avg, denom, value=-step_size)
+
+
+class Adam(Optimizer):
+    r"""Implements Adam algorithm.
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+    The implementation of the L2 penalty follows changes proposed in
+    `Decoupled Weight Decay Regularization`_.
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0,
+        amsgrad=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(
+                "Invalid beta parameter at index 0: {}".format(betas[0])
+            )
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(
+                "Invalid beta parameter at index 1: {}".format(betas[1])
+            )
+        if not 0.0 <= weight_decay:
+            raise ValueError(
+                "Invalid weight_decay value: {}".format(weight_decay)
+            )
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+        )
+        super(Adam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Adam, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            max_exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            for p in group["params"]:
+                if p.grad is not None:
+                    params_with_grad.append(p)
+                    if p.grad.is_sparse:
+                        raise RuntimeError(
+                            "Adam does not support sparse gradients,"
+                            " please consider SparseAdam instead"
+                        )
+                    grads.append(p.grad)
+
+                    state = self.state[p]
+                    # Lazy state initialization
+                    if len(state) == 0:
+                        state["step"] = 0
+                        # Exponential moving average of gradient values
+                        state["exp_avg"] = torch.zeros_like(
+                            p, memory_format=torch.preserve_format
+                        )
+                        # Exponential moving average of squared gradient values
+                        state["exp_avg_sq"] = torch.zeros_like(
+                            p, memory_format=torch.preserve_format
+                        )
+                        if group["amsgrad"]:
+                            # Maintains max of all exp. moving avg. of
+                            # sq. grad. values
+                            state["max_exp_avg_sq"] = torch.zeros_like(
+                                p, memory_format=torch.preserve_format
+                            )
+
+                    exp_avgs.append(state["exp_avg"])
+                    exp_avg_sqs.append(state["exp_avg_sq"])
+
+                    if group["amsgrad"]:
+                        max_exp_avg_sqs.append(state["max_exp_avg_sq"])
+
+                    # update the steps for each param group update
+                    state["step"] += 1
+                    # record the step after step update
+                    state_steps.append(state["step"])
+
+            adam(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                max_exp_avg_sqs,
+                state_steps,
+                amsgrad=group["amsgrad"],
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+                eps=group["eps"],
+            )
+        return loss

model/fno/fno2d.py

@@ -0,0 +1,238 @@
+from typing import List, Dict, Optional
+
+import torch
+import numpy as np
+from torch import Tensor
+import torch.nn as nn
+import torch.nn.functional as F
+from ..base_model import AutoCfdModel
+
+from onescience.modules.fourier.fno_layers import SpectralConv2d
+torch.manual_seed(0)
+np.random.seed(0)
+
+class FnoBlock(nn.Module):
+    """
+    FNO 块 (Fourier Neural Operator Block)。
+    
+    包含一个 2D 频域卷积路径和一个 1x1 的空间卷积残差路径，
+    最后加上可选的激活函数。
+    """
+    def __init__(
+        self,
+        in_chan: int,
+        out_chan: int,
+        modes1: int,
+        modes2: int,
+        act_fn: Optional[nn.Module] = None,
+    ):
+        super().__init__()
+        self.in_chan = in_chan
+        self.out_chan = out_chan
+        self.modes1 = modes1
+        self.modes2 = modes2
+        self.act_fn = act_fn
+
+
+        self.conv0 = SpectralConv2d(
+            in_channels=self.in_chan,
+            out_channels=self.out_chan,
+            modes1=self.modes1,
+            modes2=self.modes2
+        )
+        
+        self.w0 = nn.Conv2d(self.in_chan, self.out_chan, 1)
+
+    def forward(self, x: Tensor) -> Tensor:
+
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        
+        if self.act_fn is not None:
+            x = self.act_fn(x)
+            
+        return x
+
+
+class Fno2d(AutoCfdModel):
+    def __init__(
+        self,
+        in_chan: int,
+        out_chan: int,
+        n_case_params: int,
+        loss_fn: nn.Module,
+        num_layers: int,
+        modes1: int = 12,
+        modes2: int = 12,
+        hidden_dim: int = 20,
+        padding: Optional[int] = None,
+    ):
+        super().__init__(loss_fn)
+
+        """
+        The overall network. It contains 4 layers of the Fourier layer.
+        1. Lift the input to the desire channel dimension by self.fc0 .
+        2. 4 layers of the integral operators u' = (W + K)(u).
+            W defined by self.w; K defined by self.conv .
+        3. Project from the channel space to the output space by self.fc1
+            and self.fc2 .
+        """
+        self.in_chan = in_chan
+        self.out_chan = out_chan
+        self.n_case_params = n_case_params
+        self.num_layers = num_layers
+        self.modes1 = modes1
+        self.modes2 = modes2
+        self.hidden_dim = hidden_dim
+        self.padding = padding  # pad the domain if input is non-periodic
+
+        self.act_fn = nn.GELU()
+        # Channel projection into `hidden_dim` channels
+        # +1 for mask, +2 for coordinates
+        self.fc0 = nn.Conv2d(
+            in_chan + 1 + 2 + n_case_params,
+            self.hidden_dim,
+            1,
+            1,
+            0,
+        )
+        # input channel is 12: the solution of the previous 10
+        # timesteps + 2 locations
+        # (u(t-10, x, y), ..., u(t-1, x, y),  x, y)
+
+        # FNO blocks
+        blocks = []
+        for _ in range(self.num_layers):
+            blocks.append(
+                FnoBlock(
+                    self.hidden_dim,
+                    self.hidden_dim,
+                    self.modes1,
+                    self.modes2,
+                    self.act_fn,
+                )
+            )
+        self.blocks = nn.Sequential(*blocks)
+
+        self.fc1 = nn.Conv2d(self.hidden_dim, 128, 1, 1, 0)
+        self.fc2 = nn.Conv2d(128, self.out_chan, 1, 1, 0)
+
+    def forward(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+        label: Optional[Tensor] = None,
+    ) -> Dict:
+        """
+        Args:
+        - input: (b, c, h, w)
+        - labels: (b, c, h, w)
+        - mask: (b, h, w), a binary mask for indicating the geometry
+                1 for interior, 0 for obstacles.
+
+        Returns: a tensor of shape (b, c, h, w), the solution at
+            the next timestep
+        """
+        batch_size, n_chan, height, width = inputs.shape
+
+        if mask is None:
+            # When there is no mask, we assume that there is no obstacles.
+            mask = torch.ones((batch_size, 1, height, width)).to(inputs.device)
+        else:
+            if mask.dim() == 3:  # (B, h, w)
+                mask = mask.unsqueeze(1)  # (B, 1, h, w)
+        inputs = torch.cat([inputs, mask], dim=1)  # (B, c + 1, h, w)
+
+        # Physical properties
+        props = case_params  # (B, p)
+        props = props.unsqueeze(-1).unsqueeze(-1)  # (B, p, 1, 1)
+        props = props.repeat(1, 1, height, width)  # (B, p, H, W)
+
+        # Append (x, y) coordinates to every location
+        grid = self.get_coords(inputs.shape, inputs.device)  # (b, 2, h, w)
+        inputs = torch.cat(
+            (inputs, grid, props), dim=1
+        )  # (b, c + 2 + 2, h, w)
+
+        # Project channels
+        inputs = self.fc0(inputs)  # (b, hidden_dim, h, w)
+        # x = x.permute(0, 3, 1, 2)  # (b, c, h, w)?
+        if self.padding is not None:
+            # pad the domain if input is non-periodic
+            inputs = F.pad(inputs, [0, self.padding, 0, self.padding])
+
+        inputs = self.blocks(inputs)  # (b, hidden_dim, h, w)
+        if self.padding is not None:
+            # pad the domain if inputis non-periodic
+            inputs = inputs[..., : -self.padding, : -self.padding]
+
+        inputs = self.fc1(inputs)  # (b, 128, h, w)
+        inputs = self.act_fn(inputs)
+        preds = self.fc2(inputs)  # (b, c_out, h, w)
+
+        # Masked locations are not prediction
+        preds = preds * mask
+
+        if label is not None:
+            label = label * mask
+            loss = self.loss_fn(preds=preds, labels=label)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+        return dict(preds=preds)
+
+    def get_coords(self, shape, device):
+        """
+        Return a tensor of shape (b, 2, h, w) such that the element at
+        [:, :, i, j] is the (x, y) coordinates at the grid location (i, j).
+        """
+        bsz, c, size_x, size_y = shape
+        grid_x = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float)
+        grid_x = grid_x.reshape(1, 1, size_x, 1).repeat([bsz, 1, 1, size_y])
+        grid_y = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float)
+        grid_y = grid_y.reshape(1, 1, 1, size_y).repeat([bsz, 1, size_x, 1])
+        coords = torch.cat([grid_x, grid_y], dim=1).to(device)  # (b, 2, h, w)
+        return coords
+
+    def generate(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        outputs = self.forward(
+            inputs=inputs, case_params=case_params, mask=mask
+        )  # (b, c, h, w)
+        preds = outputs["preds"]
+        return preds
+
+    def generate_many(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor, steps: int
+    ) -> List[Tensor]:
+        """
+        Args:
+            x (Tensor): (c, h, w)
+            case_params (Tensor): (p)
+            mask (Tensor): (h, w), 1 for interior, 0 for obstacles.
+        Returns:
+            output: (steps, c, h, w)
+        """
+        assert len(inputs.shape) == len(case_params.shape) + 2
+        if inputs.dim() == 3:
+            # Add a dimension for batch size of 1
+            inputs = inputs.unsqueeze(0)
+            case_params = case_params.unsqueeze(0)
+            mask = mask.unsqueeze(0)
+        assert inputs.shape[0] == case_params.shape[0] == mask.shape[0]
+
+        cur_frame = inputs  # (b, c, h, w)
+        preds = []
+        for _ in range(steps):
+            cur_frame = self.generate(
+                inputs=cur_frame, case_params=case_params, mask=mask
+            )
+            preds.append(cur_frame)
+        return preds

model/fno/utilities3.py

@@ -0,0 +1,364 @@
+import torch
+import numpy as np
+import scipy.io
+import h5py
+import torch.nn as nn
+
+import operator
+from functools import reduce
+
+#################################################
+#
+# Utilities
+#
+#################################################
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class MatReader(object):
+    """reading data"""
+
+    def __init__(self, file_path, to_torch=True, to_cuda=False, to_float=True):
+        super(MatReader, self).__init__()
+
+        self.to_torch = to_torch
+        self.to_cuda = to_cuda
+        self.to_float = to_float
+
+        self.file_path = file_path
+
+        self.data = None
+        self.old_mat = None
+        self._load_file()
+
+    def _load_file(self):
+        try:
+            self.data = scipy.io.loadmat(self.file_path)
+            self.old_mat = True
+        except Exception:
+            self.data = h5py.File(self.file_path)
+            self.old_mat = False
+
+    def load_file(self, file_path):
+        self.file_path = file_path
+        self._load_file()
+
+    def read_field(self, field):
+        x = self.data[field]  # type: ignore
+
+        if not self.old_mat:
+            x = x[()]  # type: ignore
+            x = np.transpose(x, axes=range(len(x.shape) - 1, -1, -1))  # type: ignore  # noqa
+
+        if self.to_float:
+            x = x.astype(np.float32)  # type: ignore
+
+        if self.to_torch:
+            x = torch.from_numpy(x)
+
+            if self.to_cuda:
+                x = x.cuda()
+
+        return x
+
+    def set_cuda(self, to_cuda):
+        self.to_cuda = to_cuda
+
+    def set_torch(self, to_torch):
+        self.to_torch = to_torch
+
+    def set_float(self, to_float):
+        self.to_float = to_float
+
+
+# normalization, pointwise gaussian
+class UnitGaussianNormalizer(object):
+    def __init__(self, x, eps=0.00001):
+        super(UnitGaussianNormalizer, self).__init__()
+
+        # x could be in shape of ntrain*n or ntrain*T*n or ntrain*n*T
+        self.mean = torch.mean(x, 0)
+        self.std = torch.std(x, 0)
+        self.eps = eps
+
+    def encode(self, x):
+        x = (x - self.mean) / (self.std + self.eps)
+        return x
+
+    def decode(self, x, sample_idx=None):
+        if sample_idx is None:
+            std = self.std + self.eps  # n
+            mean = self.mean
+        else:
+            if len(self.mean.shape) == len(sample_idx[0].shape):
+                std = self.std[sample_idx] + self.eps  # batch*n
+                mean = self.mean[sample_idx]
+            if len(self.mean.shape) > len(sample_idx[0].shape):
+                std = self.std[:, sample_idx] + self.eps  # T*batch*n
+                mean = self.mean[:, sample_idx]
+
+        # x is in shape of batch*n or T*batch*n
+        x = (x * std) + mean  # type: ignore
+        return x
+
+    def cuda(self):
+        self.mean = self.mean.cuda()
+        self.std = self.std.cuda()
+
+    def cpu(self):
+        self.mean = self.mean.cpu()
+        self.std = self.std.cpu()
+
+
+# normalization, Gaussian
+class GaussianNormalizer(object):
+    def __init__(self, x, eps=0.00001):
+        super(GaussianNormalizer, self).__init__()
+
+        self.mean = torch.mean(x)
+        self.std = torch.std(x)
+        self.eps = eps
+
+    def encode(self, x):
+        x = (x - self.mean) / (self.std + self.eps)
+        return x
+
+    def decode(self, x, sample_idx=None):
+        x = (x * (self.std + self.eps)) + self.mean
+        return x
+
+    def cuda(self):
+        self.mean = self.mean.cuda()
+        self.std = self.std.cuda()
+
+    def cpu(self):
+        self.mean = self.mean.cpu()
+        self.std = self.std.cpu()
+
+
+# normalization, scaling by range
+class RangeNormalizer(object):
+    def __init__(self, x, low=0.0, high=1.0):
+        super(RangeNormalizer, self).__init__()
+        mymin = torch.min(x, 0)[0].view(-1)
+        mymax = torch.max(x, 0)[0].view(-1)
+
+        self.a = (high - low) / (mymax - mymin)
+        self.b = -self.a * mymax + high
+
+    def encode(self, x):
+        s = x.size()
+        x = x.view(s[0], -1)
+        x = self.a * x + self.b
+        x = x.view(s)
+        return x
+
+    def decode(self, x):
+        s = x.size()
+        x = x.view(s[0], -1)
+        x = (x - self.b) / self.a
+        x = x.view(s)
+        return x
+
+
+# loss function with rel/abs Lp loss
+class LpLoss(object):
+    def __init__(self, d=2, p=2, size_average=True, reduction=True):
+        super(LpLoss, self).__init__()
+
+        # Dimension and Lp-norm type are postive
+        assert d > 0 and p > 0
+
+        self.d = d
+        self.p = p
+        self.reduction = reduction
+        self.size_average = size_average
+
+    def abs(self, x, y):
+        num_examples = x.size()[0]
+
+        # Assume uniform mesh
+        h = 1.0 / (x.size()[1] - 1.0)
+
+        all_norms = (h ** (self.d / self.p)) * torch.norm(
+            x.view(num_examples, -1) - y.view(num_examples, -1), self.p, 1
+        )
+
+        if self.reduction:
+            if self.size_average:
+                return torch.mean(all_norms)
+            else:
+                return torch.sum(all_norms)
+
+        return all_norms
+
+    def rel(self, x, y):
+        num_examples = x.size()[0]
+
+        diff_norms = torch.norm(
+            x.reshape(num_examples, -1) - y.reshape(num_examples, -1),
+            self.p,
+            1,
+        )
+        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)
+
+        if self.reduction:
+            if self.size_average:
+                return torch.mean(diff_norms / y_norms)
+            else:
+                return torch.sum(diff_norms / y_norms)
+
+        return diff_norms / y_norms
+
+    def __call__(self, x, y):
+        return self.rel(x, y)
+
+
+# Sobolev norm (HS norm)
+# where we also compare the numerical derivatives between the output and target
+class HsLoss(object):
+    def __init__(
+        self,
+        d=2,
+        p=2,
+        k=1,
+        a=None,
+        group=False,
+        size_average=True,
+        reduction=True,
+    ):
+        super(HsLoss, self).__init__()
+
+        # Dimension and Lp-norm type are postive
+        assert d > 0 and p > 0
+
+        self.d = d
+        self.p = p
+        self.k = k
+        self.balanced = group
+        self.reduction = reduction
+        self.size_average = size_average
+
+        if a is None:
+            a = [
+                1,
+            ] * k
+        self.a = a
+
+    def rel(self, x, y):
+        num_examples = x.size()[0]
+        diff_norms = torch.norm(
+            x.reshape(num_examples, -1) - y.reshape(num_examples, -1),
+            self.p,
+            1,
+        )
+        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)
+        if self.reduction:
+            if self.size_average:
+                return torch.mean(diff_norms / y_norms)
+            else:
+                return torch.sum(diff_norms / y_norms)
+        return diff_norms / y_norms
+
+    def __call__(self, x, y, a=None):
+        nx = x.size()[1]
+        ny = x.size()[2]
+        k = self.k
+        balanced = self.balanced
+        a = self.a
+        x = x.view(x.shape[0], nx, ny, -1)
+        y = y.view(y.shape[0], nx, ny, -1)
+
+        k_x = (
+            torch.cat(
+                (
+                    torch.arange(start=0, end=nx // 2, step=1),
+                    torch.arange(start=-nx // 2, end=0, step=1),
+                ),
+                0,
+            )
+            .reshape(nx, 1)
+            .repeat(1, ny)
+        )
+        k_y = (
+            torch.cat(
+                (
+                    torch.arange(start=0, end=ny // 2, step=1),
+                    torch.arange(start=-ny // 2, end=0, step=1),
+                ),
+                0,
+            )
+            .reshape(1, ny)
+            .repeat(nx, 1)
+        )
+        k_x = torch.abs(k_x).reshape(1, nx, ny, 1).to(x.device)
+        k_y = torch.abs(k_y).reshape(1, nx, ny, 1).to(x.device)
+
+        x = torch.fft.fftn(x, dim=[1, 2])
+        y = torch.fft.fftn(y, dim=[1, 2])
+
+        if not balanced:
+            weight = 1
+            if k >= 1:
+                weight += a[0] ** 2 * (k_x**2 + k_y**2)
+            if k >= 2:
+                weight += a[1] ** 2 * (
+                    k_x**4 + 2 * k_x**2 * k_y**2 + k_y**4
+                )
+            weight = torch.sqrt(weight)  # type: ignore
+            loss = self.rel(x * weight, y * weight)
+        else:
+            loss = self.rel(x, y)
+            if k >= 1:
+                weight = a[0] * torch.sqrt(k_x**2 + k_y**2)
+                loss += self.rel(x * weight, y * weight)
+            if k >= 2:
+                weight = a[1] * torch.sqrt(
+                    k_x**4 + 2 * k_x**2 * k_y**2 + k_y**4
+                )
+                loss += self.rel(x * weight, y * weight)
+            loss = loss / (k + 1)
+
+        return loss
+
+
+# A simple feedforward neural network
+class DenseNet(torch.nn.Module):
+    def __init__(
+        self, layers, nonlinearity, out_nonlinearity=None, normalize=False
+    ):
+        super(DenseNet, self).__init__()
+
+        self.n_layers = len(layers) - 1
+
+        assert self.n_layers >= 1
+
+        self.layers = nn.ModuleList()
+
+        for j in range(self.n_layers):
+            self.layers.append(nn.Linear(layers[j], layers[j + 1]))
+
+            if j != self.n_layers - 1:
+                if normalize:
+                    self.layers.append(nn.BatchNorm1d(layers[j + 1]))
+
+                self.layers.append(nonlinearity())
+
+        if out_nonlinearity is not None:
+            self.layers.append(out_nonlinearity())
+
+    def forward(self, x):
+        for _, l in enumerate(self.layers):
+            x = l(x)
+
+        return x
+
+
+# print the number of parameters
+def count_params(model):
+    c = 0
+    for p in list(model.parameters()):
+        c += reduce(
+            operator.mul, list(p.size() + (2,) if p.is_complex() else p.size())
+        )
+    return c

model/loss.py

@@ -0,0 +1,50 @@
+from typing import List
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+class MseLoss(nn.Module):
+    def __init__(self, normalize: bool, is_masked: bool = False):
+        super().__init__()
+        self.normalize = normalize
+        self.is_masked = is_masked
+
+    def get_score_names(self) -> List[str]:
+        names = [
+            'mse', 'rmse', 'mae'
+        ]
+        if self.normalize:
+            names += ['nmse']
+        return names
+
+    def forward(self, preds: Tensor, labels: Tensor) -> dict[str, Tensor]:
+        '''
+        Args:
+        - mask: 1 for valid pixels, 0 for invalid pixels.
+        '''
+        mse = F.mse_loss(input=preds, target=labels)
+        mae = F.l1_loss(input=preds, target=labels)
+        result = dict(
+            mse=mse,
+            rmse=torch.sqrt(mse),
+            mae=mae,
+        )
+        if self.normalize:
+            result['nmse'] = mse / torch.square(labels).mean()
+            # result['mre'] = mae / torch.abs(labels).mean()
+        return result
+
+
+def loss_name_to_fn(name: str, masked: bool = False) -> MseLoss:
+    name = name.lower()
+    if masked:
+        raise NotImplementedError
+    else:
+        if name == 'mse':
+            return MseLoss(normalize=False)
+        elif name == 'nmse':
+            return MseLoss(normalize=True)
+        else:
+            raise NotImplementedError

model/resnet.py

@@ -0,0 +1,238 @@
+from typing import Optional
+
+import torch
+from torch import nn, Tensor
+
+
+from .base_model import AutoCfdModel
+
+
+class ResidualBlock(nn.Module):
+    def __init__(
+        self,
+        in_chan: int,
+        out_chan: int,
+        hidden_chan: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dropout_rate: float = 0.2,
+        bias: bool = True,
+        use_1x1conv: bool = False,
+    ):
+        super().__init__()
+        if in_chan != out_chan:
+            assert use_1x1conv
+        self.in_chan = in_chan
+        self.out_chan = out_chan
+        self.hidden_chan = hidden_chan
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.bias = bias
+
+        # Sub-modules
+        self.conv1 = nn.Conv2d(
+            in_chan,
+            hidden_chan,
+            kernel_size,
+            stride,
+            padding,
+            bias=bias,
+            padding_mode="replicate",
+        )
+        self.bn1 = nn.BatchNorm2d(hidden_chan)
+        self.dropout = nn.Dropout(p=dropout_rate)
+        self.act = nn.GELU()
+        self.conv2 = nn.Conv2d(
+            hidden_chan,
+            out_chan,
+            kernel_size,
+            stride,
+            padding,
+            bias=bias,
+            padding_mode="replicate",
+        )
+        self.bn2 = nn.BatchNorm2d(out_chan)
+
+        if use_1x1conv:
+            self.res_conv = nn.Conv2d(
+                in_chan,
+                out_chan,
+                kernel_size=1,
+                stride=stride,
+                padding=0,
+                bias=bias,
+            )
+        else:
+            self.res_conv = None
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.res_conv is not None:
+            residual = self.res_conv(x)
+        else:
+            residual = x
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.dropout(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x += residual
+        return x
+
+
+class ResNet(AutoCfdModel):
+    def __init__(
+        self,
+        in_chan: int,
+        out_chan: int,
+        n_case_params: int,
+        loss_fn: nn.Module,
+        hidden_chan: int = 32,
+        num_blocks: int = 4,
+        kernel_size: int = 7,
+        padding: int = 3,
+        stride: int = 1,
+    ):
+        super().__init__(loss_fn)
+        assert in_chan == out_chan
+        self.in_chan = in_chan
+        self.out_chan = out_chan
+        self.n_case_params = n_case_params
+        self.hidden_chan = hidden_chan
+        self.num_blocks = num_blocks
+        self.kernel_size = kernel_size
+        self.padding = padding
+        self.stride = stride
+
+        # Create sequence of residual blocks
+        blocks = [
+            ResidualBlock(
+                in_chan + 1 + n_case_params,  # + 1 for mask
+                hidden_chan,
+                64,
+                kernel_size,
+                stride,
+                padding,
+                use_1x1conv=True,
+            ),
+        ]
+        for _ in range(num_blocks):
+            blocks.append(
+                ResidualBlock(
+                    hidden_chan,
+                    hidden_chan,
+                    64,
+                    kernel_size,
+                    stride,
+                    padding,
+                    use_1x1conv=False,
+                )
+            )
+        blocks.append(
+            ResidualBlock(
+                hidden_chan,
+                out_chan,
+                64,
+                kernel_size,
+                stride,
+                padding,
+                use_1x1conv=True,
+            )
+        )
+
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+        label: Optional[Tensor] = None,
+    ) -> dict:
+        """
+        Args:
+        - x: (B, in_chan, h, w)
+        - case_params: (B, 3)
+        - mask: (h, w) or (B, h, w). 1 for interior, 0 for boundaries.
+        - label: (B, out_chan, h, w)
+
+        Returns:
+            (B, out_chan, h, w) or (B, out_chan, h, w), loss
+        """
+        residual = inputs[:, : self.out_chan]
+        batch_size, n_chan, height, width = inputs.shape
+        if mask is None:
+            mask = torch.ones((batch_size, height, width)).to(inputs.device)
+        else:
+            if mask.dim() == 3:
+                mask = mask.unsqueeze(1)  # (B, 1, h, w)
+        inputs = torch.cat([inputs, mask], dim=1)  # (B, c + 1, h, w)
+
+        # Add case params as additional channels
+        case_params = case_params.unsqueeze(-1).unsqueeze(-1)  # (B, c, 1, 1)
+        # (B, n_params, h, w)
+        case_params = case_params.expand(
+            -1, -1, inputs.shape[-2], inputs.shape[-1]
+        )
+        inputs = torch.cat(
+            [inputs, case_params], dim=1
+        )  # (B, c + n_params, h, w)
+
+        inputs = self.blocks(inputs)  # (B, c, h, w)
+        preds = inputs + residual
+
+        if mask is not None:
+            # Mask out predictions.
+            preds = preds * mask
+
+        if label is not None:
+            if mask is not None:
+                # Mask out labels
+                label = label * mask
+            loss = self.loss_fn(preds=preds, labels=label)
+            return dict(
+                preds=preds,
+                loss=loss,
+            )
+        else:
+            return dict(preds=preds)
+
+    def generate(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+    ):
+        outputs = self.forward(inputs, case_params=case_params, mask=mask)
+        preds = outputs["preds"]
+        return preds
+
+    def generate_many(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        steps: int,
+        mask: Tensor,
+    ):
+        """
+        x: (c, h, w)
+        mask: (h, w). 1 for interior, 0 for boundaries.
+
+        Returns:
+            (steps, c, h, w)
+        """
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)  # (1, c, h, w)
+            case_params = case_params.unsqueeze(0)  # (1, p)
+            mask = mask.unsqueeze(0)  # (1, h, w)
+        cur_frame = inputs  # (1, c, h, w)
+        frames = [cur_frame]
+        # boundaries = (1 - mask) * cur_frame  # (1, c, h, w)
+        for _ in range(steps):
+            cur_frame = self.generate(
+                cur_frame, case_params=case_params, mask=mask
+            )
+            frames.append(cur_frame)
+        return frames

model/unet.py

@@ -0,0 +1,149 @@
+from typing import Optional, List
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from .base_model import AutoCfdModel
+from onescience.modules.decoder.unet_decoder import UNetDecoder2D
+from onescience.modules.encoder.unet_encoder import UNetEncoder2D
+from onescience.modules.head.unet_head import UNetHead2D
+
+class UNet(AutoCfdModel):
+    def __init__(
+        self,
+        in_chan: int,
+        out_chan: int,
+        loss_fn: nn.Module,
+        n_case_params: int,
+        insert_case_params_at: str = "hidden",
+        bilinear: bool = False,
+        dim: int = 8,
+    ):
+        assert insert_case_params_at in ["hidden", "input"]
+        super().__init__(loss_fn)
+        
+        self.in_chan = in_chan
+        self.out_chan = out_chan
+        self.n_case_params = n_case_params
+        self.insert_case_params_at = insert_case_params_at
+        self.dim = dim
+
+        # 计算 Encoder 输入通道
+        encoder_in_chan = in_chan + 1 # + Mask
+        if insert_case_params_at == "input":
+            encoder_in_chan += n_case_params
+
+        # 1. Encoder
+        self.encoder = UNetEncoder2D(
+            in_channels=encoder_in_chan,
+            base_channels=dim,
+            num_stages=4,
+            bilinear=bilinear,
+            normtype="bn"
+        )
+
+        # 2. Hidden Injection
+        self.case_params_fc = None
+        if insert_case_params_at == "hidden":
+            bottleneck_dim = dim * 16
+            self.case_params_fc = nn.Linear(n_case_params, bottleneck_dim)
+
+        # 3. Decoder
+        self.decoder = UNetDecoder2D(
+            base_channels=dim,
+            num_stages=4,
+            bilinear=bilinear,
+            normtype="bn"
+        )
+
+        # 4. Head
+        self.head = UNetHead2D(
+            in_channels=dim,
+            out_channels=out_chan
+        )
+
+    def forward(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+        label: Optional[Tensor] = None,
+    ):
+        batch_size, n_chan, height, width = inputs.shape
+        residual = inputs[:, : self.out_chan]
+
+        # 构造 Mask
+        if mask is None:
+            mask = torch.ones((batch_size, 1, height, width)).to(inputs.device)
+        else:
+            if mask.dim() == 3:
+                mask = mask.unsqueeze(1) # (B, H, W) -> (B, 1, H, W)
+        
+        # 拼接 Mask
+        x_in = torch.cat([inputs, mask], dim=1)
+
+        # 拼接 Case Params
+        if self.insert_case_params_at == "input":
+            cp_spatial = case_params.view(batch_size, self.n_case_params, 1, 1)
+            cp_spatial = cp_spatial.expand(-1, -1, height, width)
+            x_in = torch.cat([x_in, cp_spatial], dim=1)
+
+        # Encoder
+        features = self.encoder(x_in)
+        
+        # Hidden 注入
+        if self.insert_case_params_at == "hidden":
+            bottleneck = features[-1]
+            conds = self.case_params_fc(case_params)
+            conds = conds.view(batch_size, -1, 1, 1)
+            features[-1] = bottleneck + conds
+
+        # Decoder
+        decoded = self.decoder(features)
+
+        # Head
+        preds = self.head(decoded)
+        
+        # Residual & Mask
+        preds = preds + residual
+        preds = preds * mask
+
+        if label is not None:
+            label = label * mask
+            loss = self.loss_fn(labels=label, preds=preds)
+            return {"preds": preds, "loss": loss}
+        
+        return {"preds": preds}
+
+    def generate_many(
+        self, inputs: Tensor, case_params: Tensor, mask: Tensor, steps: int
+    ) -> List[Tensor]:
+        preds = []
+        
+        # 处理单样本输入 (增加 Batch 维)
+        if inputs.dim() == 3:
+            inputs = inputs.unsqueeze(0)
+            case_params = case_params.unsqueeze(0)
+            if mask.dim() == 2:
+                mask = mask.unsqueeze(0)
+
+        # 确保 Mask 是 (B, 1, H, W) 以匹配 forward 逻辑
+        if mask.dim() == 3:
+             mask = mask.unsqueeze(1) 
+
+        cur_frame = inputs
+        for _ in range(steps):
+            out_dict = self.forward(cur_frame, case_params=case_params, mask=mask)
+            cur_frame = out_dict["preds"]
+            preds.append(cur_frame)
+            
+        return preds
+
+    def generate(
+        self,
+        inputs: Tensor,
+        case_params: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        outputs = self.forward(inputs, case_params=case_params, mask=mask)
+        return outputs["preds"]

scripts/fake_data.py

@@ -0,0 +1,78 @@
+import json
+import shutil
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def make_case(case_dir, case_id, rows, cols, frames, rng):
+    case_dir.mkdir(parents=True, exist_ok=True)
+    vel_in = 1.0 + 0.05 * case_id
+    params = {
+        "vel_in": vel_in,
+        "density": 1.0 + 0.1 * (case_id % 3),
+        "viscosity": 0.001 + 0.0001 * (case_id % 5),
+        "height": 1.0,
+        "width": 1.0,
+    }
+
+    y = np.linspace(0.0, 1.0, rows, dtype=np.float32)
+    x = np.linspace(0.0, 1.0, cols, dtype=np.float32)
+    yy, xx = np.meshgrid(y, x, indexing="ij")
+    u = np.empty((frames, rows, cols), dtype=np.float32)
+    v = np.empty_like(u)
+
+    phase = 0.2 * case_id
+    for t in range(frames):
+        amp = vel_in * (1.0 + 0.03 * t)
+        u[t] = amp * np.sin(np.pi * yy) * (1.0 - 0.3 * xx) + 0.01 * rng.standard_normal((rows, cols))
+        v[t] = 0.2 * amp * np.cos(np.pi * xx + phase) * yy + 0.01 * rng.standard_normal((rows, cols))
+
+    (case_dir / "case.json").write_text(json.dumps(params, indent=2), encoding="utf-8")
+    np.save(case_dir / "u.npy", u)
+    np.save(case_dir / "v.npy", v)
+
+
+def main():
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    data_dir = resolve_path(cfg.datapipe.source.data_dir)
+    fake_cfg = cfg.fake_data
+    data_cfg = cfg.datapipe.data
+
+    if data_dir.exists():
+        shutil.rmtree(data_dir)
+
+    rng = np.random.default_rng(fake_cfg.seed)
+    problem_dir = data_dir / "tube"
+    case_id = 0
+    for subset in ("prop", "bc", "geo"):
+        for subset_idx in range(fake_cfg.num_cases_per_subset):
+            make_case(
+                problem_dir / subset / f"case{case_id}",
+                case_id=case_id + subset_idx,
+                rows=data_cfg.num_rows,
+                cols=data_cfg.num_cols,
+                frames=fake_cfg.num_frames,
+                rng=rng,
+            )
+            case_id += 1
+
+    total_cases = fake_cfg.num_cases_per_subset * 3
+    print(f"Fake CFDBench data written to {data_dir}")
+    print(f"cases={total_cases}, frames={fake_cfg.num_frames}, grid=({data_cfg.num_rows}, {data_cfg.num_cols})")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,145 @@
+import argparse
+import json
+import importlib.util
+import os
+import sys
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model import build_model, infer_task_type
+import onescience
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run CFDBench inference.")
+    parser.add_argument("--model", default=None, help="Override root.model.name and read weight/<model>.pt by default.")
+    return parser.parse_args()
+
+
+def load_config(model_name=None):
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    if model_name:
+        cfg.model.name = model_name
+    elif os.environ.get("CFDBENCH_MODEL_NAME"):
+        cfg.model.name = os.environ["CFDBENCH_MODEL_NAME"]
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.inference.output_dir = str(resolve_path(cfg.inference.output_dir))
+    return cfg
+
+
+def checkpoint_path(cfg):
+    path_value = cfg.inference.get("checkpoint_path", "auto")
+    if path_value == "auto":
+        return PROJECT_ROOT / "weight" / f"{cfg.model.name}.pt"
+    return resolve_path(path_value)
+
+
+def output_path(cfg, task_type):
+    output_dir = Path(cfg.inference.output_dir)
+    if cfg.inference.get("group_by_model", False):
+        output_dir = output_dir / task_type / cfg.model.name
+    return output_dir
+
+
+def select_device(requested, dist):
+    if requested == "auto":
+        return dist.device
+    if requested.startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError(f"Requested device {requested!r}, but CUDA is not available.")
+    return torch.device(requested)
+
+
+def load_cfdbench_datapipe_class():
+    runtime_root = Path(onescience.__file__).resolve().parent
+    datapipe_file = runtime_root / "datapipes" / "cfd" / "cfdbench.py"
+    spec = importlib.util.spec_from_file_location("_onescience_cfdbench_datapipe", datapipe_file)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load CFDBench datapipe from {datapipe_file}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module.CFDBenchDatapipe
+
+
+def predict_full_frame(model, batch, task_type):
+    if task_type == "auto":
+        return model.generate(
+            inputs=batch["inputs"],
+            case_params=batch["case_params"],
+            mask=batch.get("mask"),
+        )
+
+    labels = batch["label"]
+    return model.generate_one(
+        case_params=batch["case_params"],
+        t=batch["t"],
+        height=labels.shape[-2],
+        width=labels.shape[-1],
+    )
+
+
+def main():
+    args = parse_args()
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    cfg = load_config(args.model)
+    task_type = infer_task_type(cfg.model.name)
+    cfg.datapipe.data.task_type = task_type
+    device = select_device(cfg.inference.get("device", "auto"), dist)
+
+    ckpt_path = checkpoint_path(cfg)
+    if not ckpt_path.is_file():
+        raise FileNotFoundError(f"Missing checkpoint: {ckpt_path}. Run scripts/train.py first or update root.inference.checkpoint_path.")
+
+    torch.set_num_threads(1)
+    CFDBenchDatapipe = load_cfdbench_datapipe_class()
+    datapipe = CFDBenchDatapipe(cfg.datapipe, distributed=False)
+    model = build_model(cfg).to(device)
+    model.load_state_dict(torch.load(ckpt_path, map_location=device, weights_only=True))
+    model.eval()
+
+    output_dir = output_path(cfg, task_type)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    preds = []
+    score_lists = {name: [] for name in model.loss_fn.get_score_names()}
+    with torch.no_grad():
+        for step, batch in enumerate(tqdm(datapipe.test_dataloader(), desc="Inferring")):
+            if step >= cfg.inference.max_batches:
+                break
+            batch = {key: value.to(device) for key, value in batch.items()}
+            pred = predict_full_frame(model, batch, task_type)
+            label = batch["label"][:, : pred.shape[1]]
+            loss = model.loss_fn(preds=pred, labels=label)
+            preds.append(pred.detach().cpu())
+            for key, value in loss.items():
+                score_lists[key].append(float(value.detach().cpu()))
+
+    if not preds:
+        raise RuntimeError("Inference loader is empty. Check dataset split and data directory.")
+
+    torch.save(torch.cat(preds, dim=0), output_dir / "preds.pt")
+    means = {key: sum(values) / len(values) for key, values in score_lists.items() if values}
+    with open(output_dir / "scores.json", "w", encoding="utf-8") as f:
+        json.dump({"mean": means, "all": score_lists}, f, indent=2)
+    print(f"Saved predictions to {output_dir / 'preds.pt'}")
+    print(f"Scores: {means}")
+
+    DistributedManager.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,58 @@
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from onescience.utils.YParams import YParams
+from model import infer_task_type
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def output_path(cfg):
+    task_type = infer_task_type(cfg.model.name)
+    output_dir = resolve_path(cfg.inference.output_dir)
+    if cfg.inference.get("group_by_model", False):
+        output_dir = output_dir / task_type / cfg.model.name
+    return output_dir
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Summarize CFDBench inference outputs.")
+    parser.add_argument("--model", default=None, help="Override root.model.name and read that model's inference outputs.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    if args.model:
+        cfg.model.name = args.model
+    elif os.environ.get("CFDBENCH_MODEL_NAME"):
+        cfg.model.name = os.environ["CFDBENCH_MODEL_NAME"]
+    output_dir = output_path(cfg)
+    pred_path = output_dir / "preds.pt"
+    score_path = output_dir / "scores.json"
+
+    if not pred_path.is_file() or not score_path.is_file():
+        raise FileNotFoundError("Missing inference outputs. Run scripts/inference.py first.")
+
+    preds = torch.load(pred_path, map_location="cpu", weights_only=True)
+    scores = json.loads(score_path.read_text(encoding="utf-8"))
+    print(f"Prediction tensor: shape={tuple(preds.shape)}, dtype={preds.dtype}")
+    print(f"Prediction range: min={float(preds.min()):.4e}, max={float(preds.max()):.4e}")
+    print(f"Scores: {scores['mean']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,211 @@
+import argparse
+import json
+import importlib.util
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import Adam, lr_scheduler
+from tqdm import tqdm
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model import build_model, infer_task_type
+import onescience
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train CFDBench static or autoregressive models.")
+    parser.add_argument("--model", default=None, help="Override root.model.name, e.g. ffn, deeponet, fno, auto_ffn.")
+    return parser.parse_args()
+
+
+def load_config(model_name=None):
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    if model_name:
+        cfg.model.name = model_name
+    elif os.environ.get("CFDBENCH_MODEL_NAME"):
+        cfg.model.name = os.environ["CFDBENCH_MODEL_NAME"]
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.training.output_dir = str(resolve_path(cfg.training.output_dir))
+    return cfg
+
+
+def checkpoint_path(cfg):
+    name = cfg.training.get("checkpoint_name", "auto")
+    if name == "auto":
+        name = f"{cfg.model.name}.pt"
+    return PROJECT_ROOT / "weight" / name
+
+
+def output_path(cfg, task_type):
+    output_dir = Path(cfg.training.output_dir)
+    if cfg.training.get("group_by_model", False):
+        output_dir = output_dir / task_type / cfg.model.name
+    return output_dir
+
+
+def select_device(requested, dist):
+    if requested == "auto":
+        return dist.device
+    if requested.startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError(f"Requested device {requested!r}, but CUDA is not available.")
+    return torch.device(requested)
+
+
+def dump_json(data, path):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+
+def load_cfdbench_datapipe_class():
+    runtime_root = Path(onescience.__file__).resolve().parent
+    datapipe_file = runtime_root / "datapipes" / "cfd" / "cfdbench.py"
+    spec = importlib.util.spec_from_file_location("_onescience_cfdbench_datapipe", datapipe_file)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load CFDBench datapipe from {datapipe_file}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module.CFDBenchDatapipe
+
+
+def mean_scores(score_lists):
+    return {key: float(np.mean(values)) for key, values in score_lists.items() if values}
+
+
+def evaluate(model, loader, device, dist, desc="Evaluating"):
+    model_eval = model.module if hasattr(model, "module") else model
+    model_eval.eval()
+    score_lists = {name: [] for name in model_eval.loss_fn.get_score_names()}
+
+    with torch.no_grad():
+        iterator = tqdm(loader, desc=desc, disable=(dist.rank != 0))
+        for batch in iterator:
+            batch = {key: value.to(device) for key, value in batch.items()}
+            outputs = model_eval(**batch)
+            for key, value in outputs["loss"].items():
+                score_lists[key].append(float(value.detach().cpu()))
+
+    return {"mean": mean_scores(score_lists), "all": score_lists}
+
+
+def train(cfg, model, datapipe, output_dir, device, dist):
+    train_cfg = cfg.training
+    train_loader, train_sampler = datapipe.train_dataloader()
+    val_loader, _ = datapipe.val_dataloader()
+
+    if len(train_loader) == 0:
+        raise RuntimeError("Training loader is empty. Increase fake_data.num_cases_per_subset or reduce batch_size.")
+
+    optimizer = Adam(model.parameters(), lr=train_cfg.lr)
+    scheduler = lr_scheduler.StepLR(optimizer, step_size=train_cfg.lr_step_size, gamma=train_cfg.lr_gamma)
+    train_losses = []
+    best_nmse = float("inf")
+
+    for epoch in range(train_cfg.num_epochs):
+        if train_sampler:
+            train_sampler.set_epoch(epoch)
+        model.train()
+        start = time.time()
+
+        iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=(dist.rank != 0))
+        for step, batch in enumerate(iterator):
+            batch = {key: value.to(device) for key, value in batch.items()}
+            outputs = model(**batch)
+            loss = outputs["loss"][train_cfg.loss_name]
+
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+
+            loss_value = float(loss.detach().cpu())
+            train_losses.append(loss_value)
+            if dist.rank == 0 and (step + 1) % train_cfg.log_interval == 0:
+                iterator.set_postfix({"loss": f"{loss_value:.4e}"})
+
+        scheduler.step()
+
+        if dist.rank == 0 and (epoch + 1) % train_cfg.eval_interval == 0:
+            ckpt_dir = output_dir / f"ckpt-{epoch}"
+            ckpt_dir.mkdir(parents=True, exist_ok=True)
+            scores = evaluate(model, val_loader, device, dist, desc=f"Val {epoch}")
+            dump_json(scores, ckpt_dir / "dev_scores.json")
+            dump_json({"epoch": epoch, "train_loss": train_losses, "seconds": time.time() - start}, ckpt_dir / "scores.json")
+
+            nmse = scores["mean"].get(train_cfg.loss_name, float("inf"))
+            model_to_save = model.module if hasattr(model, "module") else model
+            torch.save(model_to_save.state_dict(), ckpt_dir / "model.pt")
+            if nmse <= best_nmse:
+                best_nmse = nmse
+                weight_path = checkpoint_path(cfg)
+                torch.save(model_to_save.state_dict(), weight_path)
+                print(f"Saved best checkpoint to {weight_path}")
+
+        if dist.world_size > 1:
+            torch.distributed.barrier()
+
+    if dist.rank == 0:
+        dump_json(train_losses, output_dir / "train_losses.json")
+
+
+def main(required_task_type="static", entry_name="scripts/train.py"):
+    args = parse_args()
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    cfg = load_config(args.model)
+    task_type = infer_task_type(cfg.model.name)
+    if required_task_type is not None and task_type != required_task_type:
+        expected = "ffn/deeponet" if required_task_type == "static" else "auto_* / resnet / unet / fno"
+        raise ValueError(
+            f"{entry_name} is the {required_task_type} entry, but model.name={cfg.model.name!r} is {task_type}. "
+            f"Use a {expected} model, or run scripts/train_auto.py for autoregressive models."
+        )
+    cfg.datapipe.data.task_type = task_type
+    device = select_device(cfg.training.get("device", "auto"), dist)
+
+    output_dir = output_path(cfg, task_type)
+    if dist.rank == 0:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        print(f"Config: {PROJECT_ROOT / 'conf' / 'config.yaml'}")
+        print(f"Model: {cfg.model.name} ({task_type})")
+        print(f"Data: {cfg.datapipe.source.data_dir}")
+        print(f"Output: {output_dir}")
+        print(f"Checkpoint: {checkpoint_path(cfg)}")
+
+    torch.set_num_threads(1)
+    CFDBenchDatapipe = load_cfdbench_datapipe_class()
+    datapipe = CFDBenchDatapipe(cfg.datapipe, distributed=(dist.world_size > 1))
+    model = build_model(cfg).to(device)
+
+    if dist.world_size > 1 and "train" in cfg.training.mode:
+        device_ids = [dist.local_rank] if device.type == "cuda" else None
+        model = DDP(model, device_ids=device_ids)
+
+    if "train" in cfg.training.mode:
+        train(cfg, model, datapipe, output_dir, device, dist)
+
+    if "test" in cfg.training.mode and dist.rank == 0:
+        scores = evaluate(model, datapipe.test_dataloader(), device, dist, desc="Test")
+        dump_json(scores, output_dir / "test_scores.json")
+        print(f"Test scores: {scores['mean']}")
+
+    DistributedManager.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/train_auto.py

@@ -0,0 +1,13 @@
+import sys
+from pathlib import Path
+
+
+SCRIPTS_DIR = Path(__file__).resolve().parent
+if str(SCRIPTS_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPTS_DIR))
+
+from train import main
+
+
+if __name__ == "__main__":
+    main(required_task_type="auto", entry_name="scripts/train_auto.py")

weight/.gitkeep

@@ -0,0 +1 @@
+