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

@@ -0,0 +1,140 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- physics-informed neural networks
+- long-horizon physical prediction
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">PINNsformer</span>
+  </strong>
+</p>
+
+# Model Overview
+
+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
+
+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 |
+| :--- | :--- |
+| 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 |
+
+
+# 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 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
+hf download --model OneScience-Sugon/PINNsformer --local-dir ./PINNsformer
+cd PINNsformer
+```
+
+### 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
+
+To use real data, download it from the link below and set `data.data_dir` in `conf/config.yaml` to the correct path.
+
+| 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 PINNsFormer model weights in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads `weight/1dreaction_pinnsformer.pt` and saves:
+
+```text
+result/prediction.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# 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 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

@@ -0,0 +1,46 @@
+project:
+  name: PINNsformer
+  task: 1d_reaction
+
+runtime:
+  seed: 0
+  device: auto
+
+model:
+  name: PINNsformer1D
+  d_out: 1
+  d_hidden: 128
+  d_model: 32
+  num_layers: 1
+  heads: 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
+
+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

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

model/__init__.py

@@ -0,0 +1,3 @@
+from .pinnsformer import PINNsformer1D, PINNsformer2D, WaveAct
+
+__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

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import json
+import random
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+
+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 = PROJECT_ROOT / path
+    with path.open("r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def project_path(path: str | Path) -> Path:
+    path = Path(path)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def 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)
+
+
+def ensure_parent(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+
+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 build_model(cfg: dict[str, Any]) -> torch.nn.Module:
+    from model import PINNsformer1D
+
+    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"]),
+    )
+
+
+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 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)
+
+
+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 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

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+
+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 = 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__":
+    main()

scripts/inference.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+import torch
+
+from onescience.utils.pinnsformer_util import get_data, make_time_sequence
+
+from common import (
+    build_model,
+    exact_reaction_solution,
+    load_config,
+    project_path,
+    relative_errors,
+    select_device,
+)
+
+
+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()
+
+    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.")
+
+    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()
+
+    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]
+
+    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"]),
+        )
+
+    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)
+
+    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__":
+    main()

scripts/result.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+
+from common import load_config, project_path, relative_errors, write_json
+
+
+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 = 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__":
+    main()

scripts/train.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import Adam, LBFGS
+
+from onescience.utils.pinnsformer_util import get_data, get_n_params, make_time_sequence
+
+from common import (
+    build_model,
+    ensure_runtime_dirs,
+    initial_condition,
+    load_config,
+    project_path,
+    seed_everything,
+    select_device,
+)
+
+
+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)
+
+
+def tensorize(array: np.ndarray, device: torch.device) -> torch.Tensor:
+    return torch.tensor(array, dtype=torch.float32, requires_grad=True, device=device)
+
+
+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,
+    )
+
+    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,
+    )
+
+
+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):
+        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__":
+    main()

weight/.gitkeep

@@ -0,0 +1 @@
+