Upload folder using huggingface_hub

.gitattributes

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

README.md

@@ -4,37 +4,38 @@ language:
 - en
 tags:
 - OneScience
-- PDE
-- CFD
-- neural-operator
-- PDEBench
+- fluid dynamics
+- physics-informed neural networks
+- long-horizon physical prediction
 frameworks: PyTorch
-
-datasets:
- - OneScience/pdenneval
 ---
 <p align="center">
   <strong>
-    <span style="font-size: 30px;">PDENNEval</span>
+    <span style="font-size: 30px;">PINNsformer</span>
   </strong>
 </p>
 
 # Model Overview
-PDENNEval is a comprehensive benchmark developed by researchers at Sun Yat-sen University for evaluating neural-network methods for solving partial differential equations. It enables systematic comparison of function-learning and operator-learning approaches across a broad range of PDE tasks.
 
-Paper: [PDENNEval: A Comprehensive Evaluation of Neural Network Methods for Solving PDEs](https://www.ijcai.org/proceedings/2024/0573.pdf)
+PINNsFormer is a Transformer-based physics-informed neural network framework developed by researchers at the Georgia Institute of Technology and Carnegie Mellon University. It enables rapid prediction of solutions to time-dependent partial differential equations and their associated physical fields.
+
+Paper: [PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks](https://arxiv.org/abs/2307.11833).
 
 # Model Description
-PDENNEval provides a unified framework for evaluating neural-network PDE solvers. It contains 16 datasets generated with high-accuracy conventional scientific computing methods, covering 19 PDE problem classes from fluid dynamics, materials science, finance, electromagnetics, and other domains. The benchmark evaluates the solution accuracy, computational efficiency, and robustness of 12 neural-network methods.
 
-## Use Cases
+PINNsFormer uses a Transformer encoder–decoder with multi-head attention to numerically solve time-dependent partial differential equations, including convection, reaction, wave, and Navier–Stokes equations.
+
+
+
+# Use Cases
 
 | Use Case | Description |
-| --- | --- |
-| PDE neural-operator validation | Complete a minimal FNO training and inference workflow using PDEBench-style HDF5 data |
-| Neural PDE method development | Extend the model collection in `model/`, which already includes FNO, DeepONet, PINO-FNO, UNO, MPNN, and UNet definitions |
+| :--- | :--- |
+| Time-dependent PDE solving | Train a continuous-field surrogate constrained by physical residuals, boundary conditions, and initial conditions |
+| Physics-informed neural network validation | Rapidly validate the PINNsFormer network, loss functions, weight serialization, and inference pipeline |
+| One-dimensional reaction equation example | Generate target fields from an analytical solution for pipeline validation and error analysis |
 | ModelScope/OneCode execution | Download the standalone model package, install its dependencies, and run the provided scripts |
-| Rapid pipeline checks | Use `fake_data.py` and `smoke_models.py` to verify the data, models, and scripts |
+
 
 # Usage
 
@@ -48,16 +49,15 @@ Use the online OneCode environment for an intelligent, one-click AI for Science
 
 **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.
-
+* A CPU can be used for small-scale pipeline validation.
+* A GPU or DCU is recommended for training on larger grids or for more epochs.
+* 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
 
 ```bash
-modelscope download --model OneScience/PDENNEval --local_dir ./PDENNEval
-cd PDENNEval
+modelscope download --model OneScience/PINNsformer --local_dir ./PINNsformer
+cd PINNsformer
 ```
 
 ### Set Up the Runtime Environment
@@ -84,20 +84,30 @@ pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple
 
 ### Training Data
 
-The OneScience community provides the `pdenneval` dataset for training. Download it with the command below and verify that the data path in `conf/config.yaml` is configured correctly:
+To use real data, download it from the link below and set `data.data_dir` in `conf/config.yaml` to the correct path.
 
-```bash
-modelscope download --dataset OneScience/pdenneval  --local_dir ./data
-```
+| Source | Link | Extraction Code | Destination |
+|---|---|---|---|
+| Baidu Netdisk | https://pan.baidu.com/s/1pM4ICc6FJX5pLF7WEoozxQ?pwd=5gha | `5gha` | `convection/convection.mat` and `navier_stokes/cylinder_nektar_wake.mat` |
 
 ### Training
 
 ```bash
 python scripts/train.py
 ```
+The default configuration uses a smaller grid and fewer L-BFGS iterations for rapid end-to-end validation. To restore the scale of the original example, edit `conf/config.yaml`:
+
+```yaml
+data:
+  x_num: 101
+  t_num: 101
+
+training:
+  epochs: 500
+```
 
 ### Model Weights
-This repository will provide weights trained on the `pdenneval` dataset in the `weights/` directory. The weights will be uploaded soon.
+This repository will provide PINNsFormer model weights in the `weights/` directory. The weights will be uploaded soon.
 
 ### Inference
 
@@ -105,7 +115,11 @@ This repository will provide weights trained on the `pdenneval` dataset in the `
 python scripts/inference.py
 ```
 
-By default, inference loads `weight/best_model.pt` and writes `.npz` results to `result/output/`.
+Inference loads `weight/1dreaction_pinnsformer.pt` and saves:
+
+```text
+result/prediction.npz
+```
 
 ### Evaluation and Visualization
 
@@ -122,5 +136,5 @@ python scripts/result.py
 
 # Citations and License
 
-- Original PDENNEval paper: [PDENNEval: A Comprehensive Evaluation of Neural Network Methods for Solving PDEs](https://www.ijcai.org/proceedings/2024/0573.pdf).
-- This repository retains the relevant source and attribution notices. Follow all applicable license requirements when using, modifying, or distributing its contents.
+* Original PINNsformer paper: [PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks](https://arxiv.org/abs/2307.11833).
+* This repository retains the relevant source and attribution notices. Follow all applicable license requirements when using, modifying, or distributing its contents.

conf/config.yaml

@@ -1,52 +1,46 @@
-fno_config:
-  datapipe:
-    source:
-      data_dir: "./data"
-      file_name: "2D_DarcyFlow_beta0.1_Train.hdf5"
-    verbose: true
-    data:
-      single_file: true
-      reduced_resolution: 1
-      reduced_resolution_t: 1
-      reduced_batch: 1
-      initial_step: 1
-      test_ratio: 0.1
-      pde_name: "2D_Darcy_Flow"
-    dataloader:
-      batch_size: 2
-      num_workers: 0
-      pin_memory: false
+project:
+  name: PINNsformer
+  task: 1d_reaction
 
-  model:
-    name: "FNO"
-    num_channels: 1
-    modes: 4
-    width: 8
+runtime:
+  seed: 0
+  device: auto
 
-  training:
-    output_dir: "./weight"
-    if_training: true
-    continue_training: false
-    model_path: null
-    save_period: 1
-    seed: 0
-    training_type: "single"
-    t_train: 2
-    epochs: 1
-    optimizer:
-      name: "Adam"
-      lr: 1.0e-3
-      weight_decay: 1.0e-4
-    scheduler:
-      name: "StepLR"
-      step_size: 100
-      gamma: 0.5
+model:
+  name: PINNsformer1D
+  d_out: 1
+  d_hidden: 128
+  d_model: 32
+  num_layers: 1
+  heads: 2
 
-  inference:
-    checkpoint: "./weight/best_model.pt"
-    output_dir: "./result/output"
-    max_batches: 2
+data:
+  x_range: [0.0, 6.283185307179586]
+  t_range: [0.0, 1.0]
+  x_num: 31
+  t_num: 31
+  sequence:
+    num_step: 5
+    step: 0.0001
 
-  result:
-    prediction_dir: "./result/output"
-    output_dir: "./result"
+equation:
+  reaction_rate: 5.0
+  initial:
+    center: 3.141592653589793
+    sigma: 0.7853981633974483
+
+training:
+  epochs: 2
+  optimizer:
+    name: LBFGS
+    lr: 1.0
+    max_iter: 5
+    line_search_fn: strong_wolfe
+  checkpoint: weight/1dreaction_pinnsformer.pt
+
+paths:
+  fake_data: data/reaction_fake.npz
+  prediction: result/prediction.npz
+  metrics: result/metrics.json
+  loss: result/train_loss.npy
+  figure: result/1dreaction_pinnsformer.png

configuration.json

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

model/__init__.py

@@ -1,3 +1,3 @@
-from .fno import FNO1d, FNO2d, FNO3d, FNO_maxwell
+from .pinnsformer import PINNsformer1D, PINNsformer2D, WaveAct
 
-__all__ = ["FNO1d", "FNO2d", "FNO3d", "FNO_maxwell"]
+__all__ = ["PINNsformer1D", "PINNsformer2D", "WaveAct"]

model/pinnsformer.py

@@ -0,0 +1,146 @@
+# implementation of PINNsformer
+# paper: PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks
+# link: https://arxiv.org/abs/2307.11833
+
+import torch
+import torch.nn as nn
+
+from onescience.utils.pinnsformer_util import get_clones
+
+class WaveAct(nn.Module):
+    def __init__(self):
+        super(WaveAct, self).__init__() 
+        self.w1 = nn.Parameter(torch.ones(1), requires_grad=True)
+        self.w2 = nn.Parameter(torch.ones(1), requires_grad=True)
+
+    def forward(self, x):
+        return self.w1 * torch.sin(x)+ self.w2 * torch.cos(x)
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model, d_ff=256):
+        super(FeedForward, self).__init__() 
+        self.linear = nn.Sequential(*[
+            nn.Linear(d_model, d_ff),
+            WaveAct(),
+            nn.Linear(d_ff, d_ff),
+            WaveAct(),
+            nn.Linear(d_ff, d_model)
+        ])
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model, heads):
+        super(EncoderLayer, self).__init__()
+
+        self.attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=heads, batch_first=True)
+        self.ff = FeedForward(d_model)
+        self.act1 = WaveAct()
+        self.act2 = WaveAct()
+        
+    def forward(self, x):
+        x2 = self.act1(x)
+        # pdb.set_trace()
+        x = x + self.attn(x2,x2,x2)[0]
+        x2 = self.act2(x)
+        x = x + self.ff(x2)
+        return x
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model, heads):
+        super(DecoderLayer, self).__init__()
+
+        self.attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=heads, batch_first=True)
+        self.ff = FeedForward(d_model)
+        self.act1 = WaveAct()
+        self.act2 = WaveAct()
+
+    def forward(self, x, e_outputs): 
+        x2 = self.act1(x)
+        x = x + self.attn(x2, e_outputs, e_outputs)[0]
+        x2 = self.act2(x)
+        x = x + self.ff(x2)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, d_model, N, heads):
+        super(Encoder, self).__init__()
+        self.N = N
+        self.layers = get_clones(EncoderLayer(d_model, heads), N)
+        self.act = WaveAct()
+
+    def forward(self, x):
+        for i in range(self.N):
+            x = self.layers[i](x)
+        return self.act(x)
+
+
+class Decoder(nn.Module):
+    def __init__(self, d_model, N, heads):
+        super(Decoder, self).__init__()
+        self.N = N
+        self.layers = get_clones(DecoderLayer(d_model, heads), N)
+        self.act = WaveAct()
+        
+    def forward(self, x, e_outputs):
+        for i in range(self.N):
+            x = self.layers[i](x, e_outputs)
+        return self.act(x)
+
+
+
+class PINNsformer1D(nn.Module):
+    def __init__(self, d_out, d_model, d_hidden, N, heads):
+        super(PINNsformer1D, self).__init__()
+
+        self.linear_emb = nn.Linear(2, d_model)
+
+        self.encoder = Encoder(d_model, N, heads)
+        self.decoder = Decoder(d_model, N, heads)
+        self.linear_out = nn.Sequential(*[
+            nn.Linear(d_model, d_hidden),
+            WaveAct(),
+            nn.Linear(d_hidden, d_hidden),
+            WaveAct(),
+            nn.Linear(d_hidden, d_out)
+        ])
+
+    def forward(self, x, t):
+        src = torch.cat((x,t), dim=-1)
+        src = self.linear_emb(src)
+
+        e_outputs = self.encoder(src)
+        d_output = self.decoder(src, e_outputs)
+        output = self.linear_out(d_output)
+        # pdb.set_trace()
+        # raise Exception('stop')
+        return output
+
+
+class PINNsformer2D(nn.Module):
+    def __init__(self, d_out, d_model, d_hidden, N, heads):
+        super(PINNsformer2D, self).__init__()
+
+        self.linear_emb = nn.Linear(3, d_model)
+
+        self.encoder = Encoder(d_model, N, heads)
+        self.decoder = Decoder(d_model, N, heads)
+        self.linear_out = nn.Sequential(*[
+            nn.Linear(d_model, d_hidden),
+            WaveAct(),
+            nn.Linear(d_hidden, d_hidden),
+            WaveAct(),
+            nn.Linear(d_hidden, d_out)
+        ])
+
+    def forward(self, x, y, t):
+        src = torch.cat((x,y,t), dim=-1)
+        src = self.linear_emb(src)
+        e_outputs = self.encoder(src)
+        d_output = self.decoder(src, e_outputs)
+        output = self.linear_out(d_output)
+        return output

result/.gitkeep

@@ -0,0 +1 @@
+

scripts/common.py

@@ -1,337 +1,100 @@
 from __future__ import annotations
 
-import os
+import json
+import random
 import sys
-import warnings
 from pathlib import Path
-from types import SimpleNamespace
-from typing import Any, Optional
+from typing import Any
 
-import h5py
 import numpy as np
 import torch
-from torch.utils.data import DataLoader, Dataset
-
+import yaml
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
-DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
-
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-from model.fno import FNO1d, FNO2d, FNO3d, FNO_maxwell
-
-try:
-    from onescience.utils.YParams import YParams
-except ModuleNotFoundError as exc:
-    YParams = None
-    YPARAMS_IMPORT_ERROR = exc
-
-try:
-    from onescience.distributed.manager import DistributedManager
-except ModuleNotFoundError as exc:
-    DistributedManager = None
-    DISTRIBUTED_IMPORT_ERROR = exc
-
-try:
-    from onescience.datapipes.cfd.PDENNEval import PDEBenchFNODatapipe
-except (ImportError, ModuleNotFoundError, OSError) as exc:
-    PDEBenchFNODatapipe = None
-    DATAPIPE_IMPORT_ERROR = exc
-
-
-def get_attr(obj: Any, name: str, default: Any = None) -> Any:
-    if isinstance(obj, dict):
-        return obj.get(name, default)
-    return getattr(obj, name, default)
-
-
-def set_attr(obj: Any, name: str, value: Any) -> None:
-    if isinstance(obj, dict):
-        obj[name] = value
-    else:
-        setattr(obj, name, value)
-
-
-def to_namespace(value: Any) -> Any:
-    if isinstance(value, dict):
-        return SimpleNamespace(**{key: to_namespace(val) for key, val in value.items()})
-    if isinstance(value, list):
-        return [to_namespace(item) for item in value]
-    return value
-
-
-def load_config(config_path: Path | str = DEFAULT_CONFIG) -> Any:
-    config_path = Path(config_path).expanduser().resolve()
-    if YParams is not None:
-        return YParams(str(config_path), "fno_config")
-
-    import yaml
-
-    with config_path.open("r", encoding="utf-8") as handle:
-        raw = yaml.safe_load(handle)
-    if not isinstance(raw, dict) or "fno_config" not in raw:
-        raise ValueError(f"{config_path} must contain root key 'fno_config'")
-    return to_namespace(raw["fno_config"])
-
 
-def resolve_path(value: Any, root: Path = PROJECT_ROOT) -> Path:
-    if value is None:
-        raise ValueError("path value can not be null")
-    expanded = os.path.expandvars(os.path.expanduser(str(value)))
-    path = Path(expanded)
+def load_config(config_path: str | Path | None = None) -> dict[str, Any]:
+    path = Path(config_path) if config_path else PROJECT_ROOT / "conf" / "config.yaml"
     if not path.is_absolute():
-        path = root / path
-    return path.resolve()
+        path = PROJECT_ROOT / path
+    with path.open("r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
 
 
-def prepare_config(
-    cfg: Any,
-    data_dir: Optional[str] = None,
-    output_dir: Optional[str] = None,
-    checkpoint: Optional[str] = None,
-) -> Any:
-    source = cfg.datapipe.source
-    training = cfg.training
-    inference = get_attr(cfg, "inference", None)
+def project_path(path: str | Path) -> Path:
+    path = Path(path)
+    return path if path.is_absolute() else PROJECT_ROOT / path
 
-    resolved_data = resolve_path(data_dir or source.data_dir)
-    set_attr(source, "data_dir", str(resolved_data))
 
-    resolved_output = resolve_path(output_dir or training.output_dir)
-    set_attr(training, "output_dir", str(resolved_output))
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
 
-    model_path = get_attr(training, "model_path", None)
-    if model_path:
-        set_attr(training, "model_path", str(resolve_path(model_path)))
 
-    if inference is not None:
-        infer_output = get_attr(inference, "output_dir", "./result/output")
-        infer_checkpoint = checkpoint or get_attr(inference, "checkpoint", None)
-        set_attr(inference, "output_dir", str(resolve_path(infer_output)))
-        if infer_checkpoint:
-            set_attr(inference, "checkpoint", str(resolve_path(infer_checkpoint)))
+def select_device(value: str = "auto") -> torch.device:
+    if value == "auto":
+        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    return torch.device(value)
 
-    return cfg
 
+def ensure_parent(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
 
-class SingleProcessManager:
-    rank = 0
-    local_rank = 0
-    world_size = 1
-    distributed = False
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+def ensure_runtime_dirs(cfg: dict[str, Any]) -> None:
+    for key in ("fake_data", "prediction", "metrics", "loss", "figure"):
+        ensure_parent(project_path(cfg["paths"][key]))
+    ensure_parent(project_path(cfg["training"]["checkpoint"]))
 
-def initialize_distributed() -> Any:
-    if DistributedManager is None:
-        warnings.warn(
-            f"Falling back to a single-process manager: {DISTRIBUTED_IMPORT_ERROR}",
-            RuntimeWarning,
-        )
-        return SingleProcessManager()
-    DistributedManager.initialize()
-    return DistributedManager()
 
+def build_model(cfg: dict[str, Any]) -> torch.nn.Module:
+    from model import PINNsformer1D
 
-def cleanup_distributed() -> None:
-    if DistributedManager is not None and DistributedManager.is_initialized():
-        DistributedManager.cleanup()
+    model_cfg = cfg["model"]
+    return PINNsformer1D(
+        d_out=int(model_cfg["d_out"]),
+        d_hidden=int(model_cfg["d_hidden"]),
+        d_model=int(model_cfg["d_model"]),
+        N=int(model_cfg["num_layers"]),
+        heads=int(model_cfg["heads"]),
+    )
 
 
-class LocalPDEBenchFNODataset(Dataset):
-    """Small FNO-compatible reader used when optional OneScience graph deps are absent."""
+def initial_condition(x: np.ndarray | torch.Tensor, cfg: dict[str, Any]):
+    initial = cfg["equation"]["initial"]
+    center = initial["center"]
+    sigma = initial["sigma"]
+    if isinstance(x, torch.Tensor):
+        return torch.exp(-((x - center) ** 2) / (2 * sigma**2))
+    return np.exp(-((x - center) ** 2) / (2 * sigma**2))
 
-    def __init__(self, datapipe_cfg: Any, mode: str):
-        self.mode = mode
-        self.data_cfg = datapipe_cfg.data
-        self.source_cfg = datapipe_cfg.source
-        self.initial_step = int(self.data_cfg.initial_step)
-        self.reduced_resolution = int(self.data_cfg.reduced_resolution)
-        self.reduced_resolution_t = int(self.data_cfg.reduced_resolution_t)
-        self.reduced_batch = int(self.data_cfg.reduced_batch)
-        self.test_ratio = float(get_attr(self.data_cfg, "test_ratio", 0.1))
-        self.file_path = Path(self.source_cfg.data_dir) / self.source_cfg.file_name
-        self._load_single_file()
 
-    def _load_single_file(self) -> None:
-        if not self.file_path.is_file():
-            raise FileNotFoundError(f"HDF5 file not found: {self.file_path}")
-
-        with h5py.File(self.file_path, "r") as handle:
-            if "tensor" not in handle:
-                raise ValueError("local fallback datapipe requires HDF5 dataset 'tensor'")
-            tensor = np.asarray(handle["tensor"], dtype=np.float32)
-            if tensor.ndim == 3:
-                data = tensor[
-                    :: self.reduced_batch,
-                    :: self.reduced_resolution_t,
-                    :: self.reduced_resolution,
-                ]
-                data = np.transpose(data, (0, 2, 1))
-                self.data = data[:, :, :, None]
-                x = np.asarray(handle["x-coordinate"], dtype=np.float32)
-                self.grid = torch.tensor(
-                    x[:: self.reduced_resolution], dtype=torch.float32
-                ).unsqueeze(-1)
-            elif tensor.ndim == 4:
-                data = tensor[
-                    :: self.reduced_batch,
-                    :,
-                    :: self.reduced_resolution,
-                    :: self.reduced_resolution,
-                ]
-                data = np.transpose(data, (0, 2, 3, 1))
-                if "nu" in handle:
-                    nu = np.asarray(handle["nu"], dtype=np.float32)[
-                        :: self.reduced_batch,
-                        :: self.reduced_resolution,
-                        :: self.reduced_resolution,
-                    ]
-                    data = np.concatenate([nu[:, :, :, None], data], axis=-1)
-                self.data = data[:, :, :, :, None]
-                x = torch.tensor(np.asarray(handle["x-coordinate"], dtype=np.float32))
-                y = torch.tensor(np.asarray(handle["y-coordinate"], dtype=np.float32))
-                xx, yy = torch.meshgrid(x, y, indexing="ij")
-                self.grid = torch.stack((xx, yy), dim=-1)[
-                    :: self.reduced_resolution, :: self.reduced_resolution
-                ].float()
-            else:
-                raise ValueError(f"unsupported tensor ndim for FNO fallback: {tensor.ndim}")
-
-        sample_count = self.data.shape[0]
-        val_count = max(1, int(sample_count * self.test_ratio)) if sample_count > 1 else 0
-        if self.mode == "train":
-            self.data = self.data[val_count:]
-        else:
-            self.data = self.data[:val_count]
-        self.data = torch.tensor(self.data, dtype=torch.float32)
-        self.spatial_dim = len(self.data.shape) - 3
-
-    def __len__(self) -> int:
-        return len(self.data)
-
-    def __getitem__(self, idx: int):
-        sample = self.data[idx]
-        return sample[..., : self.initial_step, :], sample, self.grid
-
-
-class LocalPDEBenchFNODatapipe:
-    def __init__(self, cfg: Any, distributed: bool = False):
-        self.config = cfg
-        self.distributed = distributed
-        self.train_dataset = LocalPDEBenchFNODataset(cfg.datapipe, "train")
-        self.val_dataset = LocalPDEBenchFNODataset(cfg.datapipe, "val")
-        self.spatial_dim = self.train_dataset.spatial_dim
-
-    def train_dataloader(self):
-        loader_args = self.config.datapipe.dataloader
-        return DataLoader(
-            self.train_dataset,
-            batch_size=int(loader_args.batch_size),
-            num_workers=int(loader_args.num_workers),
-            pin_memory=bool(loader_args.pin_memory),
-            shuffle=True,
-            drop_last=True,
-        ), None
-
-    def val_dataloader(self):
-        loader_args = self.config.datapipe.dataloader
-        return DataLoader(
-            self.val_dataset,
-            batch_size=int(loader_args.batch_size),
-            num_workers=int(loader_args.num_workers),
-            pin_memory=bool(loader_args.pin_memory),
-            shuffle=False,
-            drop_last=False,
-        ), None
-
-
-def build_datapipe(cfg: Any, distributed: bool = False, force_local: bool = False) -> Any:
-    if PDEBenchFNODatapipe is not None and not force_local:
-        return PDEBenchFNODatapipe(cfg, distributed=distributed)
-    if PDEBenchFNODatapipe is None:
-        warnings.warn(
-            f"Using local FNO HDF5 datapipe because OneScience datapipe import failed: "
-            f"{DATAPIPE_IMPORT_ERROR}",
-            RuntimeWarning,
-        )
-    return LocalPDEBenchFNODatapipe(cfg, distributed=distributed)
-
-
-def build_model(spatial_dim: int, cfg: Any) -> torch.nn.Module:
-    model_args = cfg.model
-    data_cfg = cfg.datapipe.data
-    initial_step = int(data_cfg.initial_step)
-    pde_name = get_attr(data_cfg, "pde_name", "")
-    modes = int(model_args.modes)
-
-    if pde_name == "3D_Maxwell":
-        return FNO_maxwell(
-            num_channels=int(model_args.num_channels),
-            width=int(model_args.width),
-            modes1=modes,
-            modes2=modes,
-            modes3=modes,
-            initial_step=initial_step,
-        )
-    if spatial_dim == 1:
-        return FNO1d(
-            num_channels=int(model_args.num_channels),
-            width=int(model_args.width),
-            modes=modes,
-            initial_step=initial_step,
-        )
-    if spatial_dim == 2:
-        return FNO2d(
-            num_channels=int(model_args.num_channels),
-            width=int(model_args.width),
-            modes1=modes,
-            modes2=modes,
-            initial_step=initial_step,
-        )
-    if spatial_dim == 3:
-        return FNO3d(
-            num_channels=int(model_args.num_channels),
-            width=int(model_args.width),
-            modes1=modes,
-            modes2=modes,
-            modes3=modes,
-            initial_step=initial_step,
-        )
-    raise ValueError(f"unsupported spatial dimension: {spatial_dim}")
-
-
-def predict_batch(model: torch.nn.Module, x: torch.Tensor, y: torch.Tensor, grid: torch.Tensor, cfg: Any):
-    data_cfg = cfg.datapipe.data
-    train_cfg = cfg.training
-    initial_step = int(data_cfg.initial_step)
-    t_train = min(int(train_cfg.t_train), y.shape[-2])
-    input_shape = list(x.shape)[:-2] + [-1]
+def exact_reaction_solution(
+    x: np.ndarray | torch.Tensor,
+    t: np.ndarray | torch.Tensor,
+    cfg: dict[str, Any],
+):
+    h = initial_condition(x, cfg)
+    rate = cfg["equation"]["reaction_rate"]
+    if isinstance(h, torch.Tensor):
+        exp_term = torch.exp(rate * t)
+    else:
+        exp_term = np.exp(rate * t)
+    return h * exp_term / (h * exp_term + 1 - h)
 
-    if get_attr(train_cfg, "training_type", "single") == "autoregressive":
-        pred = y[..., :initial_step, :]
-        for _ in range(initial_step, t_train):
-            model_input = x.reshape(input_shape)
-            model_output = model(model_input, grid)
-            if model_output.dim() == pred.dim() - 1:
-                model_output = model_output.unsqueeze(-2)
-            pred = torch.cat((pred, model_output), dim=-2)
-            x = torch.cat((x[..., 1:, :], model_output), dim=-2)
-        return pred, y[..., :t_train, :]
 
-    model_input = x.reshape(input_shape)
-    target = y[..., t_train - 1 : t_train, :]
-    pred = model(model_input, grid)
-    if pred.dim() == target.dim() - 1:
-        pred = pred.unsqueeze(-2)
-    return pred, target
+def relative_errors(pred: np.ndarray, target: np.ndarray) -> dict[str, float]:
+    return {
+        "relative_l1": float(np.sum(np.abs(target - pred)) / np.sum(np.abs(target))),
+        "relative_l2": float(np.sqrt(np.sum((target - pred) ** 2) / np.sum(target**2))),
+    }
 
 
-def load_model_state(path: Path, device: torch.device) -> dict[str, torch.Tensor]:
-    checkpoint = torch.load(path, map_location=device, weights_only=False)
-    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
-        return checkpoint["model_state_dict"]
-    if isinstance(checkpoint, dict):
-        return checkpoint
-    raise ValueError(f"unsupported checkpoint format: {path}")
+def write_json(path: Path, payload: dict[str, Any]) -> None:
+    ensure_parent(path)
+    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")

scripts/fake_data.py

@@ -1,59 +1,31 @@
 from __future__ import annotations
 
 import argparse
-from pathlib import Path
 
-import h5py
 import numpy as np
 
-from common import DEFAULT_CONFIG, load_config, prepare_config
-
-
-def generate_darcy_file(path: Path, samples: int, grid_size: int, seed: int, overwrite: bool) -> None:
-    if path.exists() and not overwrite:
-        print(f"Fake data already exists: {path}")
-        return
-
-    path.parent.mkdir(parents=True, exist_ok=True)
-    rng = np.random.default_rng(seed)
-    x = np.linspace(0.0, 1.0, grid_size, dtype=np.float32)
-    y = np.linspace(0.0, 1.0, grid_size, dtype=np.float32)
-    xx, yy = np.meshgrid(x, y, indexing="ij")
-
-    nu = rng.uniform(0.05, 0.25, size=(samples, grid_size, grid_size)).astype(np.float32)
-    tensor = np.empty((samples, 2, grid_size, grid_size), dtype=np.float32)
-    base = np.sin(np.pi * xx) * np.sin(np.pi * yy)
-    for idx in range(samples):
-        tensor[idx, 0] = base + 0.05 * rng.standard_normal(base.shape)
-        tensor[idx, 1] = tensor[idx, 0] * (1.0 + nu[idx]) + 0.01 * rng.standard_normal(base.shape)
-
-    with h5py.File(path, "w") as handle:
-        handle.create_dataset("tensor", data=tensor)
-        handle.create_dataset("nu", data=nu)
-        handle.create_dataset("x-coordinate", data=x)
-        handle.create_dataset("y-coordinate", data=y)
-        handle.create_dataset("t-coordinate", data=np.array([0.0, 1.0], dtype=np.float32))
-
-    print(f"Fake Darcy HDF5 written: {path}")
-    print(f"tensor={tensor.shape}, nu={nu.shape}")
-
-
-def main() -> int:
-    parser = argparse.ArgumentParser(description="Generate a small PDENNEval FNO fake dataset.")
-    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
-    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
-    parser.add_argument("--samples", type=int, default=20)
-    parser.add_argument("--grid-size", type=int, default=16)
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--overwrite", action="store_true")
+from common import exact_reaction_solution, load_config, project_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Generate fake PINNsformer reaction data.")
+    parser.add_argument("--config", default=None, help="Path to config.yaml.")
     args = parser.parse_args()
 
-    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir)
-    data_dir = Path(cfg.datapipe.source.data_dir)
-    file_path = data_dir / cfg.datapipe.source.file_name
-    generate_darcy_file(file_path, args.samples, args.grid_size, args.seed, args.overwrite)
-    return 0
+    cfg = load_config(args.config)
+    data_cfg = cfg["data"]
+    x = np.linspace(*data_cfg["x_range"], int(data_cfg["x_num"]), dtype=np.float32)
+    t = np.linspace(*data_cfg["t_range"], int(data_cfg["t_num"]), dtype=np.float32)
+    x_mesh, t_mesh = np.meshgrid(x, t)
+    u = exact_reaction_solution(x_mesh, t_mesh, cfg).astype(np.float32)
+
+    output_path = project_path(cfg["paths"]["fake_data"])
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(output_path, x=x, t=t, u=u)
+
+    print(f"fake data saved to {output_path}")
+    print(f"x={x.shape}, t={t.shape}, u={u.shape}")
 
 
 if __name__ == "__main__":
-    raise SystemExit(main())
+    main()

scripts/inference.py

@@ -1,83 +1,85 @@
 from __future__ import annotations
 
 import argparse
-from pathlib import Path
 
 import numpy as np
 import torch
 
+from onescience.utils.pinnsformer_util import get_data, make_time_sequence
+
 from common import (
-    DEFAULT_CONFIG,
-    build_datapipe,
     build_model,
-    cleanup_distributed,
-    get_attr,
-    initialize_distributed,
+    exact_reaction_solution,
     load_config,
-    load_model_state,
-    predict_batch,
-    prepare_config,
+    project_path,
+    relative_errors,
+    select_device,
 )
 
 
-@torch.no_grad()
-def main() -> int:
-    parser = argparse.ArgumentParser(description="Run PDENNEval FNO inference.")
-    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
-    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
-    parser.add_argument("--checkpoint", default=None, help="Checkpoint path")
-    parser.add_argument("--output-dir", default=None, help="Directory for npz predictions")
-    parser.add_argument("--max-batches", type=int, default=None)
-    parser.add_argument("--force-local-datapipe", action="store_true")
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run PINNsformer inference for the 1D reaction equation.")
+    parser.add_argument("--config", default=None, help="Path to config.yaml.")
+    parser.add_argument("--checkpoint", default=None, help="Path to a trained checkpoint.")
+    parser.add_argument("--device", default=None, help="Override runtime.device.")
     args = parser.parse_args()
 
-    dist = initialize_distributed()
-    device = dist.device
-    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir, checkpoint=args.checkpoint)
-    if args.output_dir:
-        cfg.inference.output_dir = str(Path(args.output_dir).expanduser().resolve())
-
-    checkpoint = Path(args.checkpoint or cfg.inference.checkpoint)
-    if not checkpoint.is_file():
-        raise FileNotFoundError(f"checkpoint not found: {checkpoint}")
+    cfg = load_config(args.config)
+    device = select_device(args.device or cfg["runtime"]["device"])
+    checkpoint = project_path(args.checkpoint or cfg["training"]["checkpoint"])
+    if not checkpoint.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint}. Run scripts/train.py first.")
 
-    datapipe = build_datapipe(
-        cfg,
-        distributed=False,
-        force_local=args.force_local_datapipe,
-    )
-    val_loader, _ = datapipe.val_dataloader()
-    model = build_model(datapipe.spatial_dim, cfg).to(device)
-    model.load_state_dict(load_model_state(checkpoint, device))
+    model = build_model(cfg).to(device)
+    try:
+        state = torch.load(checkpoint, map_location=device, weights_only=True)
+    except TypeError:
+        state = torch.load(checkpoint, map_location=device)
+    model.load_state_dict(state.get("model_state_dict", state))
     model.eval()
 
-    output_dir = Path(args.output_dir or cfg.inference.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    max_batches = args.max_batches or int(get_attr(cfg.inference, "max_batches", 1))
+    data_cfg = cfg["data"]
+    res, _, _, _, _ = get_data(
+        data_cfg["x_range"],
+        data_cfg["t_range"],
+        int(data_cfg["x_num"]),
+        int(data_cfg["t_num"]),
+    )
+    seq = make_time_sequence(
+        res,
+        num_step=int(data_cfg["sequence"]["num_step"]),
+        step=float(data_cfg["sequence"]["step"]),
+    )
+    res_tensor = torch.tensor(seq, dtype=torch.float32, requires_grad=False, device=device)
+    x_test, t_test = res_tensor[:, :, 0:1], res_tensor[:, :, 1:2]
 
-    written = 0
-    for batch_idx, (x, y, grid) in enumerate(val_loader):
-        if batch_idx >= max_batches:
-            break
-        x = x.to(device)
-        y = y.to(device)
-        grid = grid.to(device)
-        pred, target = predict_batch(model, x, y, grid, cfg)
-        path = output_dir / f"batch_{batch_idx:04d}.npz"
-        np.savez_compressed(
-            path,
-            prediction=pred.detach().cpu().numpy(),
-            target=target.detach().cpu().numpy(),
+    with torch.no_grad():
+        pred = model(x_test, t_test)[:, 0:1].detach().cpu().numpy().reshape(
+            int(data_cfg["t_num"]),
+            int(data_cfg["x_num"]),
         )
-        print(f"wrote {path}")
-        written += 1
 
-    if written == 0:
-        raise RuntimeError("no validation batches were available for inference")
+    x = res[:, 0].reshape(int(data_cfg["t_num"]), int(data_cfg["x_num"]))
+    t = res[:, 1].reshape(int(data_cfg["t_num"]), int(data_cfg["x_num"]))
+    target = exact_reaction_solution(x, t, cfg).astype(np.float32)
+    metrics = relative_errors(pred, target)
 
-    cleanup_distributed()
-    return 0
+    output_path = project_path(cfg["paths"]["prediction"])
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        output_path,
+        x=x.astype(np.float32),
+        t=t.astype(np.float32),
+        prediction=pred.astype(np.float32),
+        target=target,
+        absolute_error=np.abs(pred - target).astype(np.float32),
+        relative_l1=metrics["relative_l1"],
+        relative_l2=metrics["relative_l2"],
+    )
+    print(f"prediction saved to {output_path}")
+    print(f"relative L1 error: {metrics['relative_l1']:.6f}")
+    print(f"relative L2 error: {metrics['relative_l2']:.6f}")
 
 
 if __name__ == "__main__":
-    raise SystemExit(main())
+    main()

scripts/result.py

@@ -1,50 +1,71 @@
 from __future__ import annotations
 
 import argparse
-import json
-from pathlib import Path
 
 import numpy as np
 
-from common import DEFAULT_CONFIG, load_config, prepare_config
+from common import load_config, project_path, relative_errors, write_json
 
 
-def main() -> int:
-    parser = argparse.ArgumentParser(description="Summarize PDENNEval FNO inference outputs.")
-    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
-    parser.add_argument("--prediction-dir", default=None)
-    parser.add_argument("--output-dir", default=None)
+def plot_result(prediction: np.ndarray, target: np.ndarray, cfg: dict, figure_path) -> bool:
+    try:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+    except Exception as exc:
+        print(f"matplotlib unavailable, skip figure: {exc}")
+        return False
+
+    figure_path.parent.mkdir(parents=True, exist_ok=True)
+    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
+    x_min, x_max = cfg["data"]["x_range"]
+    t_min, t_max = cfg["data"]["t_range"]
+    items = [
+        (prediction, "Predicted u(x,t)"),
+        (target, "Exact u(x,t)"),
+        (np.abs(prediction - target), "Absolute Error"),
+    ]
+    for ax, (data, title) in zip(axes, items):
+        image = ax.imshow(data, extent=[x_min, x_max, t_max, t_min], aspect="auto")
+        ax.set_xlabel("x")
+        ax.set_ylabel("t")
+        ax.set_title(title)
+        fig.colorbar(image, ax=ax)
+
+    plt.tight_layout()
+    plt.savefig(figure_path)
+    plt.close(fig)
+    return True
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate and plot PINNsformer inference output.")
+    parser.add_argument("--config", default=None, help="Path to config.yaml.")
+    parser.add_argument("--prediction", default=None, help="Path to prediction .npz.")
     args = parser.parse_args()
 
-    cfg = prepare_config(load_config(args.config))
-    prediction_dir = Path(args.prediction_dir or cfg.result.prediction_dir).expanduser().resolve()
-    output_dir = Path(args.output_dir or cfg.result.output_dir).expanduser().resolve()
-    files = sorted(prediction_dir.glob("*.npz"))
-    if not files:
-        raise FileNotFoundError(f"no prediction files found in {prediction_dir}")
-
-    mse_values = []
-    mae_values = []
-    for path in files:
-        data = np.load(path)
-        pred = data["prediction"]
-        target = data["target"]
-        diff = pred - target
-        mse_values.append(float(np.mean(diff ** 2)))
-        mae_values.append(float(np.mean(np.abs(diff))))
-
-    metrics = {
-        "num_files": len(files),
-        "mse": float(np.mean(mse_values)),
-        "mae": float(np.mean(mae_values)),
-    }
-    output_dir.mkdir(parents=True, exist_ok=True)
-    metrics_path = output_dir / "metrics.json"
-    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
-    print(json.dumps(metrics, indent=2))
-    print(f"wrote {metrics_path}")
-    return 0
+    cfg = load_config(args.config)
+    prediction_path = project_path(args.prediction or cfg["paths"]["prediction"])
+    if not prediction_path.exists():
+        raise FileNotFoundError(f"Prediction file not found: {prediction_path}. Run scripts/inference.py first.")
+
+    data = np.load(prediction_path)
+    prediction = data["prediction"]
+    target = data["target"]
+    metrics = relative_errors(prediction, target)
+
+    metrics_path = project_path(cfg["paths"]["metrics"])
+    write_json(metrics_path, metrics)
+    figure_path = project_path(cfg["paths"]["figure"])
+    plotted = plot_result(prediction, target, cfg, figure_path)
+
+    print(f"metrics saved to {metrics_path}")
+    if plotted:
+        print(f"figure saved to {figure_path}")
+    print(f"relative L1 error: {metrics['relative_l1']:.6f}")
+    print(f"relative L2 error: {metrics['relative_l2']:.6f}")
 
 
 if __name__ == "__main__":
-    raise SystemExit(main())
+    main()

scripts/train.py

@@ -6,138 +6,167 @@ from pathlib import Path
 import numpy as np
 import torch
 import torch.nn as nn
-from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import Adam, LBFGS
+
+from onescience.utils.pinnsformer_util import get_data, get_n_params, make_time_sequence
 
 from common import (
-    DEFAULT_CONFIG,
-    build_datapipe,
     build_model,
-    cleanup_distributed,
-    get_attr,
-    initialize_distributed,
+    ensure_runtime_dirs,
+    initial_condition,
     load_config,
-    load_model_state,
-    predict_batch,
-    prepare_config,
+    project_path,
+    seed_everything,
+    select_device,
 )
 
 
-def train_one_epoch(model, loader, optimizer, loss_fn, device, cfg) -> float:
-    model.train()
-    losses = []
-    for x, y, grid in loader:
-        x = x.to(device)
-        y = y.to(device)
-        grid = grid.to(device)
-        pred, target = predict_batch(model, x, y, grid, cfg)
-        if pred.numel() == 0:
-            continue
-        loss = loss_fn(pred.reshape(pred.size(0), -1), target.reshape(target.size(0), -1))
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        losses.append(float(loss.item()))
-    return float(np.mean(losses)) if losses else 0.0
-
-
-@torch.no_grad()
-def validate(model, loader, loss_fn, device, cfg) -> float:
-    model.eval()
-    losses = []
-    for x, y, grid in loader:
-        x = x.to(device)
-        y = y.to(device)
-        grid = grid.to(device)
-        pred, target = predict_batch(model, x, y, grid, cfg)
-        if pred.numel() == 0:
-            continue
-        loss = loss_fn(pred.reshape(pred.size(0), -1), target.reshape(target.size(0), -1))
-        losses.append(float(loss.item()))
-    return float(np.mean(losses)) if losses else 0.0
-
-
-def main() -> int:
-    parser = argparse.ArgumentParser(description="Train the standard PDENNEval FNO package.")
-    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
-    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
-    parser.add_argument("--output-dir", default=None, help="Override training.output_dir")
-    parser.add_argument("--force-local-datapipe", action="store_true")
-    args = parser.parse_args()
+def init_weights(module: nn.Module) -> None:
+    if isinstance(module, nn.Linear):
+        torch.nn.init.xavier_uniform_(module.weight)
+        module.bias.data.fill_(0.01)
 
-    dist = initialize_distributed()
-    device = dist.device
-    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir, output_dir=args.output_dir)
 
-    seed = int(get_attr(cfg.training, "seed", 0))
-    torch.manual_seed(seed + dist.rank)
-    np.random.seed(seed + dist.rank)
+def tensorize(array: np.ndarray, device: torch.device) -> torch.Tensor:
+    return torch.tensor(array, dtype=torch.float32, requires_grad=True, device=device)
 
-    output_dir = Path(cfg.training.output_dir)
-    if dist.rank == 0:
-        output_dir.mkdir(parents=True, exist_ok=True)
-        print(f"Output directory: {output_dir}")
-        print(f"Device: {device}, world_size={dist.world_size}")
 
-    datapipe = build_datapipe(
-        cfg,
-        distributed=(dist.world_size > 1),
-        force_local=args.force_local_datapipe,
-    )
-    train_loader, train_sampler = datapipe.train_dataloader()
-    val_loader, _ = datapipe.val_dataloader()
-
-    model = build_model(datapipe.spatial_dim, cfg).to(device)
-    if dist.world_size > 1:
-        model = DDP(model, device_ids=[dist.local_rank], output_device=dist.local_rank)
-
-    optimizer_cfg = cfg.training.optimizer
-    scheduler_cfg = cfg.training.scheduler
-    optimizer = getattr(torch.optim, optimizer_cfg.name)(
-        model.parameters(),
-        lr=float(optimizer_cfg.lr),
-        weight_decay=float(optimizer_cfg.weight_decay),
+def prepare_tensors(cfg: dict, device: torch.device, args: argparse.Namespace):
+    data_cfg = cfg["data"]
+    x_num = int(args.x_num or data_cfg["x_num"])
+    t_num = int(args.t_num or data_cfg["t_num"])
+    num_step = int(args.num_step or data_cfg["sequence"]["num_step"])
+    step = float(data_cfg["sequence"]["step"])
+
+    res, b_left, b_right, b_upper, b_lower = get_data(
+        data_cfg["x_range"],
+        data_cfg["t_range"],
+        x_num,
+        t_num,
     )
-    scheduler = getattr(torch.optim.lr_scheduler, scheduler_cfg.name)(
-        optimizer,
-        step_size=int(scheduler_cfg.step_size),
-        gamma=float(scheduler_cfg.gamma),
+
+    tensors = []
+    for values in (res, b_left, b_right, b_upper, b_lower):
+        tensors.append(tensorize(make_time_sequence(values, num_step=num_step, step=step), device))
+
+    return tuple(tensors)
+
+
+def loss_components(model: nn.Module, tensors: tuple[torch.Tensor, ...], cfg: dict):
+    res, b_left, b_right, b_upper, b_lower = tensors
+    x_res, t_res = res[:, :, 0:1], res[:, :, 1:2]
+    x_left, t_left = b_left[:, :, 0:1], b_left[:, :, 1:2]
+    x_upper, t_upper = b_upper[:, :, 0:1], b_upper[:, :, 1:2]
+    x_lower, t_lower = b_lower[:, :, 0:1], b_lower[:, :, 1:2]
+
+    pred_res = model(x_res, t_res)
+    pred_left = model(x_left, t_left)
+    pred_upper = model(x_upper, t_upper)
+    pred_lower = model(x_lower, t_lower)
+
+    u_t = torch.autograd.grad(
+        pred_res,
+        t_res,
+        grad_outputs=torch.ones_like(pred_res),
+        retain_graph=True,
+        create_graph=True,
+    )[0]
+
+    rate = float(cfg["equation"]["reaction_rate"])
+    target_ic = initial_condition(x_left[:, 0, :], cfg)
+    loss_res = torch.mean((u_t - rate * pred_res * (1 - pred_res)) ** 2)
+    loss_bc = torch.mean((pred_upper - pred_lower) ** 2)
+    loss_ic = torch.mean((pred_left[:, 0, :] - target_ic) ** 2)
+    loss = loss_res + loss_bc + loss_ic
+    return loss, (loss_res, loss_bc, loss_ic)
+
+
+def build_optimizer(model: nn.Module, cfg: dict):
+    opt_cfg = cfg["training"]["optimizer"]
+    name = opt_cfg["name"].lower()
+    if name == "adam":
+        return Adam(model.parameters(), lr=float(opt_cfg.get("lr", 1e-3)))
+    if name == "lbfgs":
+        return LBFGS(
+            model.parameters(),
+            lr=float(opt_cfg.get("lr", 1.0)),
+            max_iter=int(opt_cfg.get("max_iter", 20)),
+            line_search_fn=opt_cfg.get("line_search_fn", "strong_wolfe"),
+        )
+    raise ValueError(f"Unsupported optimizer: {opt_cfg['name']}")
+
+
+def save_checkpoint(path: Path, model: nn.Module, cfg: dict, loss_history: list[list[float]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model_state_dict": model.state_dict(),
+            "config": cfg,
+            "loss_history": loss_history,
+        },
+        path,
     )
-    loss_fn = nn.MSELoss()
 
-    if bool(get_attr(cfg.training, "continue_training", False)) and cfg.training.model_path:
-        model_to_load = model.module if hasattr(model, "module") else model
-        model_to_load.load_state_dict(load_model_state(Path(cfg.training.model_path), device))
 
-    best_val = float("inf")
-    epochs = int(cfg.training.epochs)
-    save_period = max(1, int(cfg.training.save_period))
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train PINNsformer on the 1D reaction equation.")
+    parser.add_argument("--config", default=None, help="Path to config.yaml.")
+    parser.add_argument("--epochs", type=int, default=None, help="Override training epochs.")
+    parser.add_argument("--x-num", type=int, default=None, help="Override x grid count.")
+    parser.add_argument("--t-num", type=int, default=None, help="Override t grid count.")
+    parser.add_argument("--num-step", type=int, default=None, help="Override pseudo-sequence length.")
+    parser.add_argument("--device", default=None, help="Override runtime.device.")
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    ensure_runtime_dirs(cfg)
+    seed_everything(int(cfg["runtime"]["seed"]))
+    device = select_device(args.device or cfg["runtime"]["device"])
+
+    tensors = prepare_tensors(cfg, device, args)
+    model = build_model(cfg).to(device)
+    model.apply(init_weights)
+    optimizer = build_optimizer(model, cfg)
+    epochs = int(args.epochs or cfg["training"]["epochs"])
+
+    print(model)
+    print(f"parameters: {get_n_params(model)}")
+    print(f"device: {device}")
 
+    loss_history: list[list[float]] = []
     for epoch in range(epochs):
-        if train_sampler is not None:
-            train_sampler.set_epoch(epoch)
-        train_loss = train_one_epoch(model, train_loader, optimizer, loss_fn, device, cfg)
-        scheduler.step()
-        val_loss = validate(model, val_loader, loss_fn, device, cfg)
-
-        if dist.rank == 0:
-            print(f"epoch={epoch} train_loss={train_loss:.6e} val_loss={val_loss:.6e}")
-            model_to_save = model.module if hasattr(model, "module") else model
-            state = {
-                "epoch": epoch,
-                "model_state_dict": model_to_save.state_dict(),
-                "optimizer_state_dict": optimizer.state_dict(),
-                "loss": val_loss,
-            }
-            torch.save(state, output_dir / "latest_model.pt")
-            if (epoch + 1) % save_period == 0:
-                torch.save(state, output_dir / f"model_epoch_{epoch}.pt")
-            if val_loss <= best_val:
-                best_val = val_loss
-                torch.save(state, output_dir / "best_model.pt")
-
-    cleanup_distributed()
-    return 0
+        latest: dict[str, float] = {}
+
+        def closure():
+            loss, parts = loss_components(model, tensors, cfg)
+            optimizer.zero_grad()
+            loss.backward()
+            latest["loss"] = float(loss.detach().cpu())
+            latest["loss_res"] = float(parts[0].detach().cpu())
+            latest["loss_bc"] = float(parts[1].detach().cpu())
+            latest["loss_ic"] = float(parts[2].detach().cpu())
+            return loss
+
+        if isinstance(optimizer, LBFGS):
+            optimizer.step(closure)
+        else:
+            closure()
+            optimizer.step()
+
+        loss_history.append([latest["loss_res"], latest["loss_bc"], latest["loss_ic"], latest["loss"]])
+        print(
+            f"epoch {epoch + 1}/{epochs} "
+            f"loss={latest['loss']:.6f} "
+            f"res={latest['loss_res']:.6f} "
+            f"bc={latest['loss_bc']:.6f} "
+            f"ic={latest['loss_ic']:.6f}"
+        )
+
+    checkpoint = project_path(cfg["training"]["checkpoint"])
+    save_checkpoint(checkpoint, model, cfg, loss_history)
+    np.save(project_path(cfg["paths"]["loss"]), np.asarray(loss_history, dtype=np.float32))
+    print(f"checkpoint saved to {checkpoint}")
 
 
 if __name__ == "__main__":
-    raise SystemExit(main())
+    main()