init

.gitignore

@@ -0,0 +1,15 @@
+# Local data and generated artifacts
+data/
+outputs/
+artifacts/
+
+# Python
+.venv/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+
+.ipynb_checkpoints/
+# OS/editor
+.DS_Store

PROJECT_PLAN.md

@@ -0,0 +1,80 @@
+# AirfRANS Scaling Frontier Project
+
+This document is a scope anchor, not a specification. It should help a human or agent understand what this repository is for without forcing a particular architecture, experiment grid, or implementation shape.
+
+## What we are doing
+
+We are exploring the **scaling frontier for AirfRANS-based CFD surrogate simulations**.
+
+The project is about understanding how far useful surrogate simulation can be pushed with open AirfRANS-style data, and what factors appear to control that frontier. The goal is not to preselect a model family, reproduce every baseline, or build a large simulation platform. The goal is to produce grounded evidence about what improves surrogate quality and where returns begin to flatten.
+
+Use careful language: this is an empirical scaling/frontier study, not a claim of universal scaling laws unless the evidence later supports that.
+
+## Core questions
+
+The work should stay centered on questions like:
+
+1. What quality frontier is reachable using open AirfRANS-based simulation data?
+2. How does performance change as available data, compute, and model capacity change?
+3. Which bottleneck is most visible at a given stage: data, compute, representation, model capacity, optimization, or evaluation target?
+4. When does additional simulation data appear more valuable than better training, representation, or model choice?
+5. What is the simplest experiment that can move our understanding of the frontier forward?
+
+These questions matter more than any specific architecture list or dataset-shape detail.
+
+## Working principles
+
+- Start from open AirfRANS data and existing public context.
+- Keep experiments small enough that results can be inspected, repeated, and compared.
+- Change one major scaling axis at a time when possible.
+- Track compute and wall-clock cost alongside quality metrics.
+- Prefer evidence that informs the next decision over exhaustive sweeps.
+- Avoid committing early to a specific model family, pipeline abstraction, or benchmark layout.
+- Treat additional generated simulations as a later decision, justified only by measured need.
+
+## Scaling axes to keep in mind
+
+The project should distinguish between several sources of scale:
+
+- number of independent simulations available;
+- number of sampled/query points used from those simulations;
+- model capacity;
+- training compute;
+- inference cost;
+- evaluation target, such as field quality versus engineering quantities.
+
+Not every experiment needs to cover every axis. The important part is to avoid confusing them when interpreting results.
+
+## What counts as progress
+
+A useful step should do at least one of the following:
+
+- establish a trustworthy baseline;
+- reveal a bottleneck;
+- compare two choices under a controlled constraint;
+- improve measurement or logging so future comparisons are reliable;
+- show that a proposed direction is not worth pursuing yet;
+- produce a plot, table, or saved artifact that clarifies the frontier.
+
+A larger run is not automatically better. A small run that changes the next decision is more valuable than a broad sweep with unclear interpretation.
+
+## Non-goals for now
+
+- Do not turn this into a general CFD framework.
+- Do not make model selection the center of the project before the measurement loop is solid.
+- Do not overfit the plan to AirfRANS implementation details that are not needed for the current decision.
+- Do not generate additional CFD simulations before the open-data frontier has been measured.
+- Do not treat a large experiment grid as inherently more credible than a focused frontier measurement.
+
+## Expected final shape
+
+The eventual artifact should explain, with measured evidence:
+
+- what frontier was explored;
+- what axes were varied;
+- what improved quality or efficiency;
+- what bottleneck seems most important;
+- whether additional generated simulation data appears justified;
+- what the next most rational experiment would be.
+
+The final output should be legible to someone evaluating the project as evidence of experimental judgment, simulation awareness, and resource-aware ML engineering.

configs/mlp_tiny.toml

@@ -0,0 +1,32 @@
+[run]
+name = "mlp_tiny"
+seed = 0
+artifact_dir = "artifacts/runs"
+
+[data]
+root = "data/processed/minimal"
+train_cases = 4
+val_cases = 1
+test_cases = 1
+points_per_case = 128
+batch_size = 128
+
+[model]
+type = "mlp"
+hidden_width = 128
+depth = 4
+activation = "gelu"
+
+[optim]
+lr = 0.001
+weight_decay = 0.0
+steps = 500
+log_interval = 100
+
+[device]
+type = "cuda"
+allow_cpu_fallback = false
+benchmark_kernels = true
+
+[loss]
+type = "normalized_mse"

pyproject.toml

@@ -0,0 +1,28 @@
+[project]
+name = "airfrans-frontier"
+version = "0.1.0"
+description = "Utilities for inspecting local AirfRANS scaling-frontier data."
+requires-python = ">=3.11"
+dependencies = [
+    "numpy>=2.4.0",
+    "torch>=2.8.0",
+]
+
+[project.scripts]
+airfrans-frontier = "airfrans_frontier.cli:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/airfrans_frontier"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "ipykernel>=6.30.0",
+    "matplotlib>=3.11.1",
+    "numpy>=2.4.0",
+    "nbclient>=0.10.2",
+    "nbformat>=5.10.4",
+]

src/airfrans_frontier/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

src/airfrans_frontier/cli.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import argparse
+import sys
+
+from airfrans_frontier.paths import DEFAULT_RAW_DATA_DIR, DEFAULT_RAW_MANIFEST_PATH, resolve_path
+from airfrans_frontier.raw.inspect import format_raw_inspection, inspect_raw_subset
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="airfrans-frontier")
+    subparsers = parser.add_subparsers(required=True)
+
+    inspect_raw = subparsers.add_parser("inspect-raw", help="inspect the local raw AirfRANS subset")
+    inspect_raw.add_argument("--data-dir", default=str(DEFAULT_RAW_DATA_DIR))
+    inspect_raw.add_argument("--manifest", default=str(DEFAULT_RAW_MANIFEST_PATH))
+    inspect_raw.add_argument("--sample-limit", type=int, default=5)
+    inspect_raw.set_defaults(command="inspect-raw")
+
+    train = subparsers.add_parser("train", help="train a configured baseline model")
+    train.add_argument("config", help="path to a training config TOML file")
+    train.set_defaults(command="train")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "inspect-raw":
+        if args.sample_limit < 0:
+            print("error: --sample-limit must be non-negative", file=sys.stderr)
+            return 1
+
+        data_dir = resolve_path(args.data_dir)
+        manifest_path = resolve_path(args.manifest)
+        try:
+            report = inspect_raw_subset(data_dir, manifest_path)
+        except (FileNotFoundError, NotADirectoryError, ValueError) as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            return 1
+
+        print(format_raw_inspection(report, sample_limit=args.sample_limit))
+        return 0 if report.matches_manifest else 1
+
+    if args.command == "train":
+        from airfrans_frontier.runtime import remove_pythonpath_entries
+
+        remove_pythonpath_entries()
+        from airfrans_frontier.training.loop import train_from_config_path
+
+        try:
+            result = train_from_config_path(resolve_path(args.config))
+        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            return 1
+
+        print(f"run_dir: {result.run_dir}")
+        print(f"final_metrics: {result.run_dir / 'final_metrics.json'}")
+        return 0
+
+    parser.error(f"unknown command: {args.command}")
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/airfrans_frontier/models/__init__.py

@@ -0,0 +1,5 @@
+"""Baseline model definitions."""
+
+from airfrans_frontier.models.mlp import PointwiseMLP
+
+__all__ = ["PointwiseMLP"]

src/airfrans_frontier/models/mlp.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class PointwiseMLP(nn.Module):
+    def __init__(
+        self,
+        *,
+        input_dim: int,
+        output_dim: int,
+        hidden_width: int = 128,
+        depth: int = 4,
+        activation: str = "gelu",
+    ) -> None:
+        super().__init__()
+        if input_dim <= 0:
+            raise ValueError("input_dim must be positive")
+        if output_dim <= 0:
+            raise ValueError("output_dim must be positive")
+        if hidden_width <= 0:
+            raise ValueError("hidden_width must be positive")
+        if depth <= 0:
+            raise ValueError("depth must be positive")
+
+        layers: list[nn.Module] = [nn.Linear(input_dim, hidden_width), _activation(activation)]
+        for _ in range(depth - 1):
+            layers.extend((nn.Linear(hidden_width, hidden_width), _activation(activation)))
+        layers.append(nn.Linear(hidden_width, output_dim))
+        self.network = nn.Sequential(*layers)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        return self.network(features)
+
+
+def _activation(name: str) -> nn.Module:
+    normalized = name.lower()
+    if normalized == "gelu":
+        return nn.GELU()
+    if normalized == "relu":
+        return nn.ReLU()
+    if normalized == "silu":
+        return nn.SiLU()
+    if normalized == "tanh":
+        return nn.Tanh()
+    raise ValueError(f"Unsupported activation: {name}")

src/airfrans_frontier/paths.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+
+DEFAULT_RAW_DATA_DIR = Path("data/raw/OF_dataset")
+DEFAULT_RAW_MANIFEST_PATH = Path("data/raw/OF_dataset_subset_manifest.json")
+
+
+def resolve_path(value: str | Path) -> Path:
+    return Path(value).expanduser()

src/airfrans_frontier/raw/inspect.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from airfrans_frontier.raw.manifest import RawSubsetManifest, load_raw_subset_manifest
+
+
+@dataclass(frozen=True)
+class RawInspection:
+    data_dir: Path
+    manifest_path: Path
+    simulation_count: int
+    file_count: int
+    total_bytes: int
+    manifest: RawSubsetManifest
+    missing_simulations: tuple[str, ...]
+    unexpected_simulations: tuple[str, ...]
+
+    @property
+    def matches_manifest(self) -> bool:
+        return (
+            not self.missing_simulations
+            and not self.unexpected_simulations
+            and self.simulation_count == self.manifest.extracted_simulations
+            and self.file_count == self.manifest.actual_file_count_under_target
+            and self.total_bytes == self.manifest.expected_extracted_bytes
+        )
+
+
+def inspect_raw_subset(data_dir: Path, manifest_path: Path) -> RawInspection:
+    if not data_dir.exists():
+        raise FileNotFoundError(f"Raw data directory not found: {data_dir}")
+    if not data_dir.is_dir():
+        raise NotADirectoryError(f"Raw data path is not a directory: {data_dir}")
+
+    manifest = load_raw_subset_manifest(manifest_path)
+    actual_simulation_names = tuple(sorted(path.name for path in data_dir.iterdir() if path.is_dir()))
+    actual_simulation_set = set(actual_simulation_names)
+    manifest_simulation_set = set(manifest.simulation_names)
+
+    file_paths = tuple(path for path in data_dir.rglob("*") if path.is_file())
+    file_count = len(file_paths)
+    total_bytes = sum(path.stat().st_size for path in file_paths)
+
+    return RawInspection(
+        data_dir=data_dir,
+        manifest_path=manifest_path,
+        simulation_count=len(actual_simulation_names),
+        file_count=file_count,
+        total_bytes=total_bytes,
+        manifest=manifest,
+        missing_simulations=tuple(sorted(manifest_simulation_set - actual_simulation_set)),
+        unexpected_simulations=tuple(sorted(actual_simulation_set - manifest_simulation_set)),
+    )
+
+
+def format_raw_inspection(report: RawInspection, sample_limit: int = 5) -> str:
+    status = "ok" if report.matches_manifest else "mismatch"
+    gb = report.total_bytes / 1_000_000_000
+    actual_names = tuple(
+        sorted((set(report.manifest.simulation_names) - set(report.missing_simulations)) | set(report.unexpected_simulations))
+    )
+    visible_names = actual_names[: max(sample_limit, 0)]
+    remaining = len(actual_names) - len(visible_names)
+
+    lines = [
+        "AirfRANS raw subset",
+        f"status: {status}",
+        f"data_dir: {report.data_dir}",
+        f"manifest: {report.manifest_path}",
+        f"simulations: {report.simulation_count} / {report.manifest.extracted_simulations}",
+        f"files: {report.file_count} / {report.manifest.actual_file_count_under_target}",
+        f"bytes: {report.total_bytes} / {report.manifest.expected_extracted_bytes} ({gb:.2f} GB)",
+        f"source_zip_bytes: {report.manifest.source_zip_bytes}",
+        f"sample_seed: {report.manifest.sample_seed}",
+        f"sample_method: {report.manifest.sample_method}",
+        "sample_simulations:",
+    ]
+    lines.extend(f"- {name}" for name in visible_names)
+    if remaining > 0:
+        lines.append(f"... {remaining} more")
+
+    if report.missing_simulations:
+        lines.append("missing_simulations:")
+        lines.extend(f"- {name}" for name in report.missing_simulations)
+
+    if report.unexpected_simulations:
+        lines.append("unexpected_simulations:")
+        lines.extend(f"- {name}" for name in report.unexpected_simulations)
+
+    return "\n".join(lines)

src/airfrans_frontier/raw/manifest.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True)
+class RawSubsetManifest:
+    path: Path
+    source_url: str
+    source_zip_bytes: int
+    sample_seed: int
+    sample_method: str
+    requested_simulations: int
+    extracted_simulations: int
+    expected_extracted_bytes: int
+    actual_file_count_under_target: int
+    target_root: str
+    simulation_names: tuple[str, ...]
+
+
+_REQUIRED_KEYS = (
+    "source_url",
+    "source_zip_bytes",
+    "sample_seed",
+    "sample_method",
+    "requested_simulations",
+    "extracted_simulations",
+    "expected_extracted_bytes",
+    "actual_file_count_under_target",
+    "target_root",
+    "simulations",
+)
+
+
+def load_raw_subset_manifest(path: Path) -> RawSubsetManifest:
+    if not path.exists():
+        raise FileNotFoundError(f"Raw subset manifest not found: {path}")
+
+    data: dict[str, Any] = json.loads(path.read_text())
+    for key in _REQUIRED_KEYS:
+        if key not in data:
+            raise ValueError(f"Raw subset manifest missing key: {key}")
+
+    simulations = data["simulations"]
+    if not isinstance(simulations, list):
+        raise ValueError("Raw subset manifest simulations must be a list")
+
+    simulation_names: list[str] = []
+    for simulation in simulations:
+        if not isinstance(simulation, dict) or "name" not in simulation:
+            raise ValueError("Raw subset manifest simulation entry missing key: name")
+        simulation_names.append(str(simulation["name"]))
+
+    return RawSubsetManifest(
+        path=path,
+        source_url=str(data["source_url"]),
+        source_zip_bytes=int(data["source_zip_bytes"]),
+        sample_seed=int(data["sample_seed"]),
+        sample_method=str(data["sample_method"]),
+        requested_simulations=int(data["requested_simulations"]),
+        extracted_simulations=int(data["extracted_simulations"]),
+        expected_extracted_bytes=int(data["expected_extracted_bytes"]),
+        actual_file_count_under_target=int(data["actual_file_count_under_target"]),
+        target_root=str(data["target_root"]),
+        simulation_names=tuple(sorted(simulation_names)),
+    )

src/airfrans_frontier/runtime.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+
+def remove_pythonpath_entries() -> None:
+    """Prefer uv-managed project dependencies over externally injected PYTHONPATH entries.
+
+    The training stack imports compiled dependencies such as NumPy and PyTorch. If an
+    unrelated tool injects a Python-version-specific dependency directory through
+    PYTHONPATH, those imports can resolve outside the project's uv environment. The
+    CLI should use the environment it was installed into.
+    """
+    raw_pythonpath = os.environ.get("PYTHONPATH")
+    if not raw_pythonpath:
+        return
+
+    blocked = {_normalize_path(entry) for entry in raw_pythonpath.split(os.pathsep) if entry}
+    sys.path[:] = [entry for entry in sys.path if _normalize_path(entry) not in blocked]
+
+
+def _normalize_path(entry: str) -> str:
+    return str(Path(entry or ".").expanduser().resolve())

src/airfrans_frontier/training/__init__.py

@@ -0,0 +1,3 @@
+"""Training utilities for AirfRANS frontier baselines."""
+
+__all__: list[str] = []

src/airfrans_frontier/training/artifacts.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Mapping
+
+import torch
+
+
+class ArtifactWriter:
+    def __init__(self, run_dir: Path) -> None:
+        self.run_dir = run_dir
+
+    @classmethod
+    def create(cls, base_dir: str | Path, run_name: str) -> ArtifactWriter:
+        base_path = Path(base_dir).expanduser()
+        base_path.mkdir(parents=True, exist_ok=True)
+        timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+        safe_name = _safe_name(run_name)
+        candidate = base_path / f"{timestamp}_{safe_name}"
+        suffix = 1
+        while candidate.exists():
+            candidate = base_path / f"{timestamp}_{safe_name}_{suffix}"
+            suffix += 1
+        candidate.mkdir(parents=True)
+        return cls(candidate)
+
+    def write_config(self, config_text: str) -> None:
+        (self.run_dir / "config.toml").write_text(config_text)
+
+    def write_split_manifest(self, split_manifest: Mapping[str, Any]) -> None:
+        self.write_json("split_manifest.json", dict(split_manifest))
+
+    def write_normalization(self, normalization: Mapping[str, Any]) -> None:
+        self.write_json("normalization.json", dict(normalization))
+
+    def append_metrics(self, metrics: Mapping[str, Any]) -> None:
+        with (self.run_dir / "metrics.jsonl").open("a") as file:
+            file.write(json.dumps(dict(metrics), sort_keys=True) + "\n")
+
+    def write_final_metrics(self, metrics: Mapping[str, Any]) -> None:
+        self.write_json("final_metrics.json", dict(metrics))
+
+    def write_json(self, name: str, data: Mapping[str, Any]) -> None:
+        (self.run_dir / name).write_text(json.dumps(dict(data), indent=2, sort_keys=True) + "\n")
+
+    def save_checkpoint(self, payload: Mapping[str, Any]) -> None:
+        torch.save(_to_cpu(payload), self.run_dir / "checkpoint.pt")
+
+
+def _safe_name(value: str) -> str:
+    safe = "".join(character if character.isalnum() or character in "-_" else "_" for character in value)
+    return safe or "run"
+
+
+def _to_cpu(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, dict):
+        return {key: _to_cpu(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_to_cpu(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_to_cpu(item) for item in value)
+    return value

src/airfrans_frontier/training/config.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True)
+class RunConfig:
+    name: str
+    seed: int
+    artifact_dir: Path
+
+
+@dataclass(frozen=True)
+class DataConfig:
+    root: Path
+    train_cases: int
+    val_cases: int
+    test_cases: int
+    points_per_case: int
+    batch_size: int
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    type: str
+    hidden_width: int
+    depth: int
+    activation: str
+
+
+@dataclass(frozen=True)
+class OptimConfig:
+    lr: float
+    weight_decay: float
+    steps: int
+    log_interval: int | None = None
+
+
+@dataclass(frozen=True)
+class DeviceConfig:
+    type: str
+    allow_cpu_fallback: bool
+    benchmark_kernels: bool
+
+
+@dataclass(frozen=True)
+class LossConfig:
+    type: str
+
+
+@dataclass(frozen=True)
+class TrainingConfig:
+    path: Path
+    config_text: str
+    run: RunConfig
+    data: DataConfig
+    model: ModelConfig
+    optim: OptimConfig
+    device: DeviceConfig
+    loss: LossConfig
+
+
+_REQUIRED_SECTIONS = ("run", "data", "model", "optim", "device", "loss")
+
+
+def load_training_config(path: str | Path) -> TrainingConfig:
+    config_path = Path(path).expanduser()
+    if not config_path.exists():
+        raise FileNotFoundError(f"Training config not found: {config_path}")
+    if not config_path.is_file():
+        raise ValueError(f"Training config is not a file: {config_path}")
+
+    text = config_path.read_text()
+    try:
+        raw = tomllib.loads(text)
+    except tomllib.TOMLDecodeError as exc:
+        raise ValueError(f"Invalid TOML config {config_path}: {exc}") from exc
+
+    for section_name in _REQUIRED_SECTIONS:
+        if section_name not in raw or not isinstance(raw[section_name], dict):
+            raise ValueError(f"Training config missing [{section_name}] section")
+
+    run_raw = raw["run"]
+    data_raw = raw["data"]
+    model_raw = raw["model"]
+    optim_raw = raw["optim"]
+    device_raw = raw["device"]
+    loss_raw = raw["loss"]
+
+    run = RunConfig(
+        name=_string(run_raw, "name"),
+        seed=_integer(run_raw, "seed", minimum=0),
+        artifact_dir=_path(run_raw, "artifact_dir"),
+    )
+    data = DataConfig(
+        root=_path(data_raw, "root"),
+        train_cases=_integer(data_raw, "train_cases", minimum=1),
+        val_cases=_integer(data_raw, "val_cases", minimum=0),
+        test_cases=_integer(data_raw, "test_cases", minimum=0),
+        points_per_case=_integer(data_raw, "points_per_case", minimum=1),
+        batch_size=_integer(data_raw, "batch_size", minimum=1),
+    )
+    model = ModelConfig(
+        type=_choice(_string(model_raw, "type"), {"mlp"}, "model.type"),
+        hidden_width=_integer(model_raw, "hidden_width", minimum=1),
+        depth=_integer(model_raw, "depth", minimum=1),
+        activation=_choice(_string(model_raw, "activation").lower(), {"gelu", "relu", "silu", "tanh"}, "model.activation"),
+    )
+    optim = OptimConfig(
+        lr=_number(optim_raw, "lr", minimum=0.0, exclusive_minimum=True),
+        weight_decay=_number(optim_raw, "weight_decay", minimum=0.0),
+        steps=_integer(optim_raw, "steps", minimum=1),
+        log_interval=_optional_integer(optim_raw, "log_interval", minimum=1),
+    )
+    device = DeviceConfig(
+        type=_choice(_string(device_raw, "type").lower(), {"cuda", "cpu", "auto"}, "device.type"),
+        allow_cpu_fallback=_boolean(device_raw, "allow_cpu_fallback"),
+        benchmark_kernels=_boolean(device_raw, "benchmark_kernels"),
+    )
+    loss = LossConfig(type=_choice(_string(loss_raw, "type"), {"normalized_mse"}, "loss.type"))
+
+    requested_cases = data.train_cases + data.val_cases + data.test_cases
+    if requested_cases <= 0:
+        raise ValueError("Training config must request at least one case")
+
+    return TrainingConfig(
+        path=config_path,
+        config_text=text,
+        run=run,
+        data=data,
+        model=model,
+        optim=optim,
+        device=device,
+        loss=loss,
+    )
+
+
+def _path(section: dict[str, Any], key: str) -> Path:
+    value = _string(section, key)
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path
+    return Path.cwd() / path
+
+
+def _string(section: dict[str, Any], key: str) -> str:
+    value = _required(section, key)
+    if not isinstance(value, str) or not value:
+        raise ValueError(f"Expected non-empty string for {key}")
+    return value
+
+
+def _integer(section: dict[str, Any], key: str, *, minimum: int | None = None) -> int:
+    value = _required(section, key)
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise ValueError(f"Expected integer for {key}")
+    if minimum is not None and value < minimum:
+        raise ValueError(f"Expected {key} >= {minimum}")
+    return value
+
+
+def _optional_integer(section: dict[str, Any], key: str, *, minimum: int | None = None) -> int | None:
+    if key not in section:
+        return None
+    return _integer(section, key, minimum=minimum)
+
+
+def _number(
+    section: dict[str, Any],
+    key: str,
+    *,
+    minimum: float | None = None,
+    exclusive_minimum: bool = False,
+) -> float:
+    value = _required(section, key)
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ValueError(f"Expected number for {key}")
+    result = float(value)
+    if minimum is not None:
+        if exclusive_minimum and result <= minimum:
+            raise ValueError(f"Expected {key} > {minimum}")
+        if not exclusive_minimum and result < minimum:
+            raise ValueError(f"Expected {key} >= {minimum}")
+    return result
+
+
+def _boolean(section: dict[str, Any], key: str) -> bool:
+    value = _required(section, key)
+    if not isinstance(value, bool):
+        raise ValueError(f"Expected boolean for {key}")
+    return value
+
+
+def _choice(value: str, allowed: set[str], key: str) -> str:
+    if value not in allowed:
+        allowed_text = ", ".join(sorted(allowed))
+        raise ValueError(f"Expected {key} to be one of: {allowed_text}")
+    return value
+
+
+def _required(section: dict[str, Any], key: str) -> Any:
+    if key not in section:
+        raise ValueError(f"Training config missing key: {key}")
+    return section[key]

src/airfrans_frontier/training/data.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+FloatArray = NDArray[np.float32]
+
+
+@dataclass(frozen=True)
+class SimulationSample:
+    case_id: str
+    features: FloatArray
+    targets: FloatArray
+    feature_names: tuple[str, ...]
+    target_names: tuple[str, ...]
+    source_path: Path
+
+    @property
+    def num_points(self) -> int:
+        return int(self.features.shape[0])
+
+    @property
+    def num_features(self) -> int:
+        return int(self.features.shape[1])
+
+    @property
+    def num_targets(self) -> int:
+        return int(self.targets.shape[1])
+
+
+@dataclass(frozen=True)
+class CaseSplit:
+    train_ids: tuple[str, ...]
+    val_ids: tuple[str, ...]
+    test_ids: tuple[str, ...]
+
+    def to_dict(self) -> dict[str, list[str]]:
+        return {
+            "train_ids": list(self.train_ids),
+            "val_ids": list(self.val_ids),
+            "test_ids": list(self.test_ids),
+        }
+
+
+@dataclass(frozen=True)
+class SplitArrays:
+    features: FloatArray
+    targets: FloatArray
+    case_ids: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class DatasetBundle:
+    train: SplitArrays
+    val: SplitArrays | None
+    test: SplitArrays | None
+    split: CaseSplit
+    feature_names: tuple[str, ...]
+    target_names: tuple[str, ...]
+
+
+def load_processed_dataset(root: str | Path) -> list[SimulationSample]:
+    root_path = Path(root).expanduser()
+    if not root_path.exists():
+        raise FileNotFoundError(f"Processed data directory not found: {root_path}")
+    if not root_path.is_dir():
+        raise NotADirectoryError(f"Processed data path is not a directory: {root_path}")
+
+    paths = sorted(root_path.glob("*.npz"))
+    if not paths:
+        raise ValueError(f"No .npz simulation files found under: {root_path}")
+
+    samples = [load_simulation_npz(path) for path in paths]
+    validate_common_schema(samples)
+    return samples
+
+
+def load_simulation_npz(path: str | Path) -> SimulationSample:
+    source_path = Path(path).expanduser()
+    if not source_path.exists():
+        raise FileNotFoundError(f"Simulation file not found: {source_path}")
+    if not source_path.is_file():
+        raise ValueError(f"Simulation path is not a file: {source_path}")
+
+    try:
+        with np.load(source_path, allow_pickle=False) as npz:
+            keys = set(npz.files)
+            if "features" not in keys or "targets" not in keys:
+                raise ValueError(f"Simulation file missing features/targets arrays: {source_path}")
+            features = _float_matrix(npz["features"], "features", source_path)
+            targets = _float_matrix(npz["targets"], "targets", source_path)
+            if features.shape[0] != targets.shape[0]:
+                raise ValueError(
+                    f"Simulation features/targets row mismatch in {source_path}: "
+                    f"{features.shape[0]} != {targets.shape[0]}"
+                )
+            feature_names = _names(npz, "feature_names", features.shape[1], "feature")
+            target_names = _names(npz, "target_names", targets.shape[1], "target")
+    except OSError as exc:
+        raise ValueError(f"Could not read simulation file {source_path}: {exc}") from exc
+
+    return SimulationSample(
+        case_id=source_path.stem,
+        features=features,
+        targets=targets,
+        feature_names=feature_names,
+        target_names=target_names,
+        source_path=source_path,
+    )
+
+
+def validate_common_schema(samples: list[SimulationSample]) -> None:
+    if not samples:
+        raise ValueError("Dataset contains no simulation samples")
+    first = samples[0]
+    seen_ids: set[str] = set()
+    for sample in samples:
+        if sample.case_id in seen_ids:
+            raise ValueError(f"Duplicate simulation case id: {sample.case_id}")
+        seen_ids.add(sample.case_id)
+        if sample.feature_names != first.feature_names:
+            raise ValueError(f"Feature schema mismatch in case {sample.case_id}")
+        if sample.target_names != first.target_names:
+            raise ValueError(f"Target schema mismatch in case {sample.case_id}")
+
+
+def create_case_split(
+    case_ids: list[str] | tuple[str, ...],
+    *,
+    train_cases: int,
+    val_cases: int,
+    test_cases: int,
+    seed: int,
+) -> CaseSplit:
+    requested = train_cases + val_cases + test_cases
+    if requested > len(case_ids):
+        raise ValueError(
+            f"Requested {requested} split cases but only {len(case_ids)} cases are available"
+        )
+
+    shuffled = list(case_ids)
+    random.Random(seed).shuffle(shuffled)
+    selected = shuffled[:requested]
+    train_end = train_cases
+    val_end = train_end + val_cases
+    return CaseSplit(
+        train_ids=tuple(selected[:train_end]),
+        val_ids=tuple(selected[train_end:val_end]),
+        test_ids=tuple(selected[val_end:]),
+    )
+
+
+def build_dataset_bundle(
+    samples: list[SimulationSample],
+    *,
+    train_cases: int,
+    val_cases: int,
+    test_cases: int,
+    points_per_case: int,
+    seed: int,
+) -> DatasetBundle:
+    validate_common_schema(samples)
+    samples_by_id = {sample.case_id: sample for sample in samples}
+    split = create_case_split(
+        tuple(samples_by_id),
+        train_cases=train_cases,
+        val_cases=val_cases,
+        test_cases=test_cases,
+        seed=seed,
+    )
+    train = build_split_arrays(samples_by_id, split.train_ids, points_per_case, seed=seed + 101)
+    val = (
+        build_split_arrays(samples_by_id, split.val_ids, points_per_case, seed=seed + 202)
+        if split.val_ids
+        else None
+    )
+    test = (
+        build_split_arrays(samples_by_id, split.test_ids, points_per_case, seed=seed + 303)
+        if split.test_ids
+        else None
+    )
+    return DatasetBundle(
+        train=train,
+        val=val,
+        test=test,
+        split=split,
+        feature_names=samples[0].feature_names,
+        target_names=samples[0].target_names,
+    )
+
+
+def build_split_arrays(
+    samples_by_id: dict[str, SimulationSample],
+    case_ids: tuple[str, ...],
+    points_per_case: int,
+    *,
+    seed: int,
+) -> SplitArrays:
+    if not case_ids:
+        raise ValueError("Cannot build split arrays for an empty case split")
+    rng = np.random.default_rng(seed)
+    feature_parts: list[FloatArray] = []
+    target_parts: list[FloatArray] = []
+    for case_id in case_ids:
+        sample = samples_by_id[case_id]
+        indices = _sample_indices(sample.num_points, points_per_case, rng)
+        feature_parts.append(np.ascontiguousarray(sample.features[indices], dtype=np.float32))
+        target_parts.append(np.ascontiguousarray(sample.targets[indices], dtype=np.float32))
+
+    return SplitArrays(
+        features=np.concatenate(feature_parts, axis=0),
+        targets=np.concatenate(target_parts, axis=0),
+        case_ids=case_ids,
+    )
+
+
+def _sample_indices(num_points: int, points_per_case: int, rng: np.random.Generator) -> NDArray[np.int64]:
+    if num_points <= 0:
+        raise ValueError("Cannot sample from an empty simulation")
+    if points_per_case >= num_points:
+        return np.arange(num_points, dtype=np.int64)
+    return np.sort(rng.choice(num_points, size=points_per_case, replace=False)).astype(np.int64)
+
+
+def _float_matrix(array: np.ndarray, name: str, path: Path) -> FloatArray:
+    if array.ndim != 2:
+        raise ValueError(f"Expected {name} to be a 2D array in {path}; got shape {array.shape}")
+    if array.shape[0] == 0 or array.shape[1] == 0:
+        raise ValueError(f"Expected non-empty {name} matrix in {path}; got shape {array.shape}")
+    result = np.asarray(array, dtype=np.float32)
+    if not np.isfinite(result).all():
+        raise ValueError(f"Expected finite values in {name} array: {path}")
+    return np.ascontiguousarray(result, dtype=np.float32)
+
+
+def _names(
+    npz: np.lib.npyio.NpzFile,
+    key: str,
+    expected_count: int,
+    prefix: str,
+) -> tuple[str, ...]:
+    if key not in npz.files:
+        return tuple(f"{prefix}_{index}" for index in range(expected_count))
+    values = np.asarray(npz[key])
+    if values.ndim != 1 or values.shape[0] != expected_count:
+        raise ValueError(
+            f"Expected {key} to be a 1D array of length {expected_count}; got shape {values.shape}"
+        )
+    return tuple(str(value) for value in values.tolist())

src/airfrans_frontier/training/loop.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import random
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from airfrans_frontier.models import PointwiseMLP
+from airfrans_frontier.training.artifacts import ArtifactWriter
+from airfrans_frontier.training.config import TrainingConfig, load_training_config
+from airfrans_frontier.training.data import DatasetBundle, build_dataset_bundle, load_processed_dataset
+from airfrans_frontier.training.metrics import count_parameters, device_metrics, overall_mse, per_channel_mse
+from airfrans_frontier.training.normalize import (
+    NormalizationStats,
+    compute_normalization_stats,
+    normalize_features,
+    normalize_targets,
+)
+
+
+@dataclass(frozen=True)
+class TrainingResult:
+    run_dir: Path
+    final_metrics: dict[str, Any]
+
+
+def train_from_config_path(path: str | Path) -> TrainingResult:
+    config = load_training_config(path)
+    return train(config)
+
+
+def train(config: TrainingConfig) -> TrainingResult:
+    _seed_all(config.run.seed)
+    device = select_device(config)
+
+    samples = load_processed_dataset(config.data.root)
+    bundle = build_dataset_bundle(
+        samples,
+        train_cases=config.data.train_cases,
+        val_cases=config.data.val_cases,
+        test_cases=config.data.test_cases,
+        points_per_case=config.data.points_per_case,
+        seed=config.run.seed,
+    )
+    stats = compute_normalization_stats(
+        bundle.train.features,
+        bundle.train.targets,
+        feature_names=bundle.feature_names,
+        target_names=bundle.target_names,
+    )
+    train_features = normalize_features(bundle.train.features, stats)
+    train_targets = normalize_targets(bundle.train.targets, stats)
+    val_features = normalize_features(bundle.val.features, stats) if bundle.val is not None else None
+    val_targets = normalize_targets(bundle.val.targets, stats) if bundle.val is not None else None
+    test_features = normalize_features(bundle.test.features, stats) if bundle.test is not None else None
+    test_targets = normalize_targets(bundle.test.targets, stats) if bundle.test is not None else None
+
+    model = PointwiseMLP(
+        input_dim=train_features.shape[1],
+        output_dim=train_targets.shape[1],
+        hidden_width=config.model.hidden_width,
+        depth=config.model.depth,
+        activation=config.model.activation,
+    ).to(device)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=config.optim.lr,
+        weight_decay=config.optim.weight_decay,
+    )
+    writer = ArtifactWriter.create(config.run.artifact_dir, config.run.name)
+    writer.write_config(config.config_text)
+    writer.write_split_manifest(bundle.split.to_dict())
+    writer.write_normalization(stats.to_dict())
+
+    started = time.perf_counter()
+    initial_train = evaluate_arrays(
+        model,
+        train_features,
+        train_targets,
+        batch_size=config.data.batch_size,
+        device=device,
+        target_names=bundle.target_names,
+    )
+    initial_val = (
+        evaluate_arrays(
+            model,
+            val_features,
+            val_targets,
+            batch_size=config.data.batch_size,
+            device=device,
+            target_names=bundle.target_names,
+        )
+        if val_features is not None and val_targets is not None
+        else None
+    )
+    writer.append_metrics(
+        _log_metrics(
+            step=0,
+            train_loss=initial_train["loss"],
+            val_loss=initial_val["loss"] if initial_val is not None else None,
+            elapsed_seconds=0.0,
+        )
+    )
+
+    log_interval = config.optim.log_interval or max(1, config.optim.steps // 10)
+    rng = np.random.default_rng(config.run.seed + 404)
+    model.train()
+    for step in range(1, config.optim.steps + 1):
+        batch_features, batch_targets = _sample_batch(
+            train_features,
+            train_targets,
+            batch_size=config.data.batch_size,
+            rng=rng,
+        )
+        features_tensor = _to_device(batch_features, device)
+        targets_tensor = _to_device(batch_targets, device)
+
+        optimizer.zero_grad(set_to_none=True)
+        predictions = model(features_tensor)
+        loss = F.mse_loss(predictions, targets_tensor)
+        loss.backward()
+        optimizer.step()
+
+        if step % log_interval == 0 or step == config.optim.steps:
+            train_eval = evaluate_arrays(
+                model,
+                train_features,
+                train_targets,
+                batch_size=config.data.batch_size,
+                device=device,
+                target_names=bundle.target_names,
+            )
+            val_eval = (
+                evaluate_arrays(
+                    model,
+                    val_features,
+                    val_targets,
+                    batch_size=config.data.batch_size,
+                    device=device,
+                    target_names=bundle.target_names,
+                )
+                if val_features is not None and val_targets is not None
+                else None
+            )
+            writer.append_metrics(
+                _log_metrics(
+                    step=step,
+                    train_loss=train_eval["loss"],
+                    val_loss=val_eval["loss"] if val_eval is not None else None,
+                    elapsed_seconds=time.perf_counter() - started,
+                )
+            )
+            model.train()
+
+    final_train = evaluate_arrays(
+        model,
+        train_features,
+        train_targets,
+        batch_size=config.data.batch_size,
+        device=device,
+        target_names=bundle.target_names,
+    )
+    final_val = (
+        evaluate_arrays(
+            model,
+            val_features,
+            val_targets,
+            batch_size=config.data.batch_size,
+            device=device,
+            target_names=bundle.target_names,
+        )
+        if val_features is not None and val_targets is not None
+        else None
+    )
+    final_test = (
+        evaluate_arrays(
+            model,
+            test_features,
+            test_targets,
+            batch_size=config.data.batch_size,
+            device=device,
+            target_names=bundle.target_names,
+        )
+        if test_features is not None and test_targets is not None
+        else None
+    )
+
+    elapsed = time.perf_counter() - started
+    final_metrics: dict[str, Any] = {
+        "initial_train_loss": initial_train["loss"],
+        "train_loss": final_train["loss"],
+        "train_mse_per_channel": final_train["per_channel_mse"],
+        "val_loss": final_val["loss"] if final_val is not None else None,
+        "val_mse_per_channel": final_val["per_channel_mse"] if final_val is not None else None,
+        "test_loss": final_test["loss"] if final_test is not None else None,
+        "test_mse_per_channel": final_test["per_channel_mse"] if final_test is not None else None,
+        "parameter_count": count_parameters(model),
+        "train_cases": len(bundle.split.train_ids),
+        "val_cases": len(bundle.split.val_ids),
+        "test_cases": len(bundle.split.test_ids),
+        "points_per_case": config.data.points_per_case,
+        "steps": config.optim.steps,
+        "elapsed_seconds": elapsed,
+        **device_metrics(device),
+    }
+    writer.write_final_metrics(final_metrics)
+    writer.save_checkpoint(
+        {
+            "step": config.optim.steps,
+            "model_type": config.model.type,
+            "input_dim": train_features.shape[1],
+            "output_dim": train_targets.shape[1],
+            "model_config": asdict(config.model),
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "normalization": stats.to_dict(),
+            "target_names": bundle.target_names,
+            "feature_names": bundle.feature_names,
+            "final_metrics": final_metrics,
+        }
+    )
+    return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics)
+
+
+def select_device(config: TrainingConfig) -> torch.device:
+    requested = config.device.type
+    if requested == "cuda":
+        if torch.cuda.is_available():
+            torch.backends.cudnn.benchmark = config.device.benchmark_kernels
+            return torch.device("cuda:0")
+        if config.device.allow_cpu_fallback:
+            return torch.device("cpu")
+        raise RuntimeError("CUDA requested by config but torch.cuda.is_available() is false")
+    if requested == "auto":
+        if torch.cuda.is_available():
+            torch.backends.cudnn.benchmark = config.device.benchmark_kernels
+            return torch.device("cuda:0")
+        return torch.device("cpu")
+    return torch.device("cpu")
+
+
+def evaluate_arrays(
+    model: torch.nn.Module,
+    features: np.ndarray,
+    targets: np.ndarray,
+    *,
+    batch_size: int,
+    device: torch.device,
+    target_names: tuple[str, ...],
+) -> dict[str, Any]:
+    model.eval()
+    target_dim = targets.shape[1]
+    squared_error_sum = torch.zeros(target_dim, dtype=torch.float64)
+    count = 0
+    with torch.no_grad():
+        for start in range(0, features.shape[0], batch_size):
+            stop = min(start + batch_size, features.shape[0])
+            batch_features = _to_device(features[start:stop], device)
+            batch_targets = _to_device(targets[start:stop], device)
+            predictions = model(batch_features)
+            errors = predictions - batch_targets
+            squared_error_sum += (errors.double().pow(2).sum(dim=0)).detach().cpu()
+            count += stop - start
+    return {
+        "loss": overall_mse(squared_error_sum, count, target_dim),
+        "per_channel_mse": per_channel_mse(squared_error_sum, count, target_names),
+    }
+
+
+def _sample_batch(
+    features: np.ndarray,
+    targets: np.ndarray,
+    *,
+    batch_size: int,
+    rng: np.random.Generator,
+) -> tuple[np.ndarray, np.ndarray]:
+    indices = rng.integers(0, features.shape[0], size=batch_size)
+    return (
+        np.ascontiguousarray(features[indices], dtype=np.float32),
+        np.ascontiguousarray(targets[indices], dtype=np.float32),
+    )
+
+
+def _to_device(values: np.ndarray, device: torch.device) -> torch.Tensor:
+    tensor = torch.from_numpy(values)
+    if device.type == "cuda":
+        tensor = tensor.pin_memory()
+    return tensor.to(device, non_blocking=device.type == "cuda")
+
+
+def _seed_all(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 _log_metrics(
+    *,
+    step: int,
+    train_loss: float,
+    val_loss: float | None,
+    elapsed_seconds: float,
+) -> dict[str, Any]:
+    return {
+        "step": step,
+        "train_loss": train_loss,
+        "val_loss": val_loss,
+        "elapsed_seconds": elapsed_seconds,
+    }

src/airfrans_frontier/training/metrics.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def count_parameters(model: torch.nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
+
+
+def per_channel_mse(
+    squared_error_sum: torch.Tensor,
+    count: int,
+    target_names: tuple[str, ...],
+) -> dict[str, float]:
+    if count <= 0:
+        raise ValueError("Metric count must be positive")
+    values = (squared_error_sum / count).detach().cpu().tolist()
+    return {name: float(value) for name, value in zip(target_names, values, strict=True)}
+
+
+def overall_mse(squared_error_sum: torch.Tensor, count: int, target_dim: int) -> float:
+    if count <= 0 or target_dim <= 0:
+        raise ValueError("Metric count and target dimension must be positive")
+    return float((squared_error_sum.sum() / (count * target_dim)).detach().cpu().item())
+
+
+def device_metrics(device: torch.device) -> dict[str, Any]:
+    if device.type != "cuda":
+        return {
+            "device": str(device),
+            "gpu_name": None,
+            "gpu_memory_total_mb": None,
+            "gpu_memory_peak_allocated_mb": None,
+        }
+
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    properties = torch.cuda.get_device_properties(index)
+    return {
+        "device": f"cuda:{index}",
+        "gpu_name": torch.cuda.get_device_name(index),
+        "gpu_memory_total_mb": int(properties.total_memory // (1024 * 1024)),
+        "gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)),
+    }

src/airfrans_frontier/training/normalize.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from numpy.typing import NDArray
+
+FloatArray = NDArray[np.float32]
+
+
+@dataclass(frozen=True)
+class NormalizationStats:
+    feature_mean: FloatArray
+    feature_std: FloatArray
+    target_mean: FloatArray
+    target_std: FloatArray
+    feature_names: tuple[str, ...]
+    target_names: tuple[str, ...]
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "feature_names": list(self.feature_names),
+            "target_names": list(self.target_names),
+            "feature_mean": self.feature_mean.tolist(),
+            "feature_std": self.feature_std.tolist(),
+            "target_mean": self.target_mean.tolist(),
+            "target_std": self.target_std.tolist(),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, object]) -> NormalizationStats:
+        return cls(
+            feature_names=tuple(str(value) for value in _list(data, "feature_names")),
+            target_names=tuple(str(value) for value in _list(data, "target_names")),
+            feature_mean=np.asarray(_list(data, "feature_mean"), dtype=np.float32),
+            feature_std=np.asarray(_list(data, "feature_std"), dtype=np.float32),
+            target_mean=np.asarray(_list(data, "target_mean"), dtype=np.float32),
+            target_std=np.asarray(_list(data, "target_std"), dtype=np.float32),
+        )
+
+
+def compute_normalization_stats(
+    features: FloatArray,
+    targets: FloatArray,
+    *,
+    feature_names: tuple[str, ...],
+    target_names: tuple[str, ...],
+    min_std: float = 1e-6,
+) -> NormalizationStats:
+    _validate_matrix(features, "features")
+    _validate_matrix(targets, "targets")
+    if features.shape[1] != len(feature_names):
+        raise ValueError("Feature name count does not match feature matrix width")
+    if targets.shape[1] != len(target_names):
+        raise ValueError("Target name count does not match target matrix width")
+
+    feature_mean = features.mean(axis=0, dtype=np.float64).astype(np.float32)
+    feature_std = features.std(axis=0, dtype=np.float64).astype(np.float32)
+    target_mean = targets.mean(axis=0, dtype=np.float64).astype(np.float32)
+    target_std = targets.std(axis=0, dtype=np.float64).astype(np.float32)
+
+    feature_std = _clamp_std(feature_std, min_std)
+    target_std = _clamp_std(target_std, min_std)
+
+    return NormalizationStats(
+        feature_mean=feature_mean,
+        feature_std=feature_std,
+        target_mean=target_mean,
+        target_std=target_std,
+        feature_names=feature_names,
+        target_names=target_names,
+    )
+
+
+def normalize_features(features: FloatArray, stats: NormalizationStats) -> FloatArray:
+    _validate_matrix(features, "features")
+    if features.shape[1] != stats.feature_mean.shape[0]:
+        raise ValueError("Feature width does not match normalization stats")
+    return np.ascontiguousarray((features - stats.feature_mean) / stats.feature_std, dtype=np.float32)
+
+
+def normalize_targets(targets: FloatArray, stats: NormalizationStats) -> FloatArray:
+    _validate_matrix(targets, "targets")
+    if targets.shape[1] != stats.target_mean.shape[0]:
+        raise ValueError("Target width does not match normalization stats")
+    return np.ascontiguousarray((targets - stats.target_mean) / stats.target_std, dtype=np.float32)
+
+
+def denormalize_targets(targets: FloatArray, stats: NormalizationStats) -> FloatArray:
+    _validate_matrix(targets, "targets")
+    if targets.shape[1] != stats.target_mean.shape[0]:
+        raise ValueError("Target width does not match normalization stats")
+    return np.ascontiguousarray(targets * stats.target_std + stats.target_mean, dtype=np.float32)
+
+
+def save_normalization_stats(stats: NormalizationStats, path: str | Path) -> None:
+    Path(path).write_text(json.dumps(stats.to_dict(), indent=2, sort_keys=True) + "\n")
+
+
+def load_normalization_stats(path: str | Path) -> NormalizationStats:
+    data = json.loads(Path(path).read_text())
+    if not isinstance(data, dict):
+        raise ValueError(f"Invalid normalization stats file: {path}")
+    return NormalizationStats.from_dict(data)
+
+
+def _clamp_std(std: FloatArray, min_std: float) -> FloatArray:
+    result = std.copy()
+    result[result < min_std] = 1.0
+    return result.astype(np.float32)
+
+
+def _validate_matrix(values: FloatArray, name: str) -> None:
+    if values.ndim != 2:
+        raise ValueError(f"Expected {name} to be 2D; got shape {values.shape}")
+    if not np.isfinite(values).all():
+        raise ValueError(f"Expected finite values in {name}")
+
+
+def _list(data: dict[str, object], key: str) -> list[object]:
+    value = data.get(key)
+    if not isinstance(value, list):
+        raise ValueError(f"Expected list for normalization key: {key}")
+    return value

tests/test_cli.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+
+def make_raw_subset(root: Path) -> tuple[Path, Path, int, int]:
+    data_dir = root / "OF_dataset"
+    sim_a_file = data_dir / "sim_a" / "field.txt"
+    sim_b_file = data_dir / "sim_b" / "nested" / "field.txt"
+    sim_a_file.parent.mkdir(parents=True)
+    sim_b_file.parent.mkdir(parents=True)
+    sim_a_file.write_text("abc")
+    sim_b_file.write_text("defgh")
+
+    expected_bytes = 8
+    expected_files = 2
+    manifest_path = root / "OF_dataset_subset_manifest.json"
+    manifest_path.write_text(
+        json.dumps(
+            {
+                "source_url": "https://example.invalid/OF_dataset.zip",
+                "source_zip_bytes": 0,
+                "sample_seed": 123,
+                "sample_method": "test sample",
+                "requested_simulations": 2,
+                "extracted_simulations": 2,
+                "expected_extracted_bytes": expected_bytes,
+                "actual_file_count_under_target": expected_files,
+                "target_root": str(data_dir),
+                "simulations": [{"name": "sim_a"}, {"name": "sim_b"}],
+            }
+        )
+    )
+    return data_dir, manifest_path, expected_bytes, expected_files
+
+
+class InspectRawCliTests(unittest.TestCase):
+    def run_cli(self, *args: str) -> subprocess.CompletedProcess[str]:
+        return subprocess.run(
+            [sys.executable, "-m", "airfrans_frontier.cli", *args],
+            text=True,
+            capture_output=True,
+            check=False,
+        )
+
+    def test_inspect_raw_success_reports_manifest_match(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            data_dir, manifest_path, _, _ = make_raw_subset(Path(tmp))
+
+            result = self.run_cli(
+                "inspect-raw",
+                "--data-dir",
+                str(data_dir),
+                "--manifest",
+                str(manifest_path),
+                "--sample-limit",
+                "2",
+            )
+
+        self.assertEqual(result.returncode, 0)
+        self.assertIn("status: ok", result.stdout)
+        self.assertIn("simulations: 2 / 2", result.stdout)
+        self.assertIn("files: 2 / 2", result.stdout)
+        self.assertIn("bytes: 8 / 8", result.stdout)
+        self.assertIn("- sim_a", result.stdout)
+        self.assertIn("- sim_b", result.stdout)
+        self.assertEqual(result.stderr, "")
+
+    def test_inspect_raw_missing_data_dir_is_clear_error(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            _, manifest_path, _, _ = make_raw_subset(Path(tmp))
+            missing_dir = Path(tmp) / "missing"
+
+            result = self.run_cli(
+                "inspect-raw",
+                "--data-dir",
+                str(missing_dir),
+                "--manifest",
+                str(manifest_path),
+            )
+
+        self.assertEqual(result.returncode, 1)
+        self.assertEqual(result.stdout, "")
+        self.assertIn("error: Raw data directory not found:", result.stderr)
+
+    def test_inspect_raw_manifest_mismatch_returns_one(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            data_dir, manifest_path, _, _ = make_raw_subset(Path(tmp))
+            unexpected_file = data_dir / "sim_c" / "field.txt"
+            unexpected_file.parent.mkdir(parents=True)
+            unexpected_file.write_text("unexpected")
+
+            result = self.run_cli(
+                "inspect-raw",
+                "--data-dir",
+                str(data_dir),
+                "--manifest",
+                str(manifest_path),
+            )
+
+        self.assertEqual(result.returncode, 1)
+        self.assertIn("status: mismatch", result.stdout)
+        self.assertIn("unexpected_simulations:", result.stdout)
+        self.assertEqual(result.stderr, "")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_mlp.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import unittest
+
+from airfrans_frontier.runtime import remove_pythonpath_entries
+
+remove_pythonpath_entries()
+
+import torch
+
+from airfrans_frontier.models import PointwiseMLP
+
+
+class PointwiseMLPTests(unittest.TestCase):
+    def test_forward_pass_returns_batch_by_target_dim(self) -> None:
+        model = PointwiseMLP(input_dim=5, output_dim=4, hidden_width=16, depth=2, activation="gelu")
+        batch = torch.randn(7, 5)
+
+        output = model(batch)
+
+        self.assertEqual(tuple(output.shape), (7, 4))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_training_config.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+
+from airfrans_frontier.training.config import load_training_config
+
+
+class TrainingConfigTests(unittest.TestCase):
+    def test_config_loader_accepts_mlp_tiny(self) -> None:
+        config = load_training_config("configs/mlp_tiny.toml")
+
+        self.assertEqual(config.run.name, "mlp_tiny")
+        self.assertEqual(config.model.type, "mlp")
+        self.assertEqual(config.loss.type, "normalized_mse")
+        self.assertEqual(config.device.type, "cuda")
+        self.assertTrue(config.data.root.is_absolute())
+
+    def test_config_loader_rejects_missing_section(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            config_path = Path(tmp) / "bad.toml"
+            config_path.write_text("[run]\nname = 'bad'\n")
+
+            with self.assertRaisesRegex(ValueError, r"missing \[data\] section"):
+                load_training_config(config_path)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_training_data.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+
+from airfrans_frontier.runtime import remove_pythonpath_entries
+
+remove_pythonpath_entries()
+
+import numpy as np
+
+from airfrans_frontier.training.data import create_case_split, load_processed_dataset, load_simulation_npz
+from airfrans_frontier.training.normalize import compute_normalization_stats, normalize_targets
+
+
+def write_case(path: Path, offset: float = 0.0) -> None:
+    features = np.array(
+        [
+            [offset + 0.0, 1.0],
+            [offset + 1.0, 2.0],
+            [offset + 2.0, 3.0],
+        ],
+        dtype=np.float32,
+    )
+    targets = np.array(
+        [
+            [offset + 10.0, -1.0],
+            [offset + 11.0, 0.0],
+            [offset + 12.0, 1.0],
+        ],
+        dtype=np.float32,
+    )
+    np.savez(
+        path,
+        features=features,
+        targets=targets,
+        feature_names=np.array(["x", "y"]),
+        target_names=np.array(["pressure", "velocity"]),
+    )
+
+
+class TrainingDataTests(unittest.TestCase):
+    def test_dataset_loader_rejects_malformed_npz(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "bad.npz"
+            np.savez(path, features=np.array([1.0, 2.0], dtype=np.float32), targets=np.ones((2, 1)))
+
+            with self.assertRaisesRegex(ValueError, "features to be a 2D array"):
+                load_simulation_npz(path)
+
+    def test_dataset_loader_loads_common_schema(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            write_case(root / "case_a.npz", offset=0.0)
+            write_case(root / "case_b.npz", offset=1.0)
+
+            samples = load_processed_dataset(root)
+
+        self.assertEqual([sample.case_id for sample in samples], ["case_a", "case_b"])
+        self.assertEqual(samples[0].feature_names, ("x", "y"))
+        self.assertEqual(samples[0].target_names, ("pressure", "velocity"))
+
+    def test_case_split_is_deterministic_and_case_level(self) -> None:
+        case_ids = [f"case_{index}" for index in range(10)]
+
+        first = create_case_split(case_ids, train_cases=6, val_cases=2, test_cases=2, seed=7)
+        second = create_case_split(case_ids, train_cases=6, val_cases=2, test_cases=2, seed=7)
+
+        self.assertEqual(first, second)
+        self.assertEqual(len(set(first.train_ids) & set(first.val_ids)), 0)
+        self.assertEqual(len(set(first.train_ids) & set(first.test_ids)), 0)
+        self.assertEqual(len(first.train_ids), 6)
+        self.assertEqual(len(first.val_ids), 2)
+        self.assertEqual(len(first.test_ids), 2)
+
+    def test_normalization_uses_train_split_only(self) -> None:
+        train_features = np.array([[0.0], [2.0]], dtype=np.float32)
+        train_targets = np.array([[10.0], [14.0]], dtype=np.float32)
+        validation_targets = np.array([[1000.0]], dtype=np.float32)
+
+        stats = compute_normalization_stats(
+            train_features,
+            train_targets,
+            feature_names=("x",),
+            target_names=("pressure",),
+        )
+        normalized_validation = normalize_targets(validation_targets, stats)
+
+        self.assertAlmostEqual(float(stats.target_mean[0]), 12.0)
+        self.assertAlmostEqual(float(stats.target_std[0]), 2.0)
+        self.assertAlmostEqual(float(normalized_validation[0, 0]), 494.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_training_loop.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from airfrans_frontier.runtime import remove_pythonpath_entries
+
+remove_pythonpath_entries()
+
+import numpy as np
+import torch
+
+from airfrans_frontier.training.config import load_training_config
+from airfrans_frontier.training.loop import select_device
+
+
+FEATURE_NAMES = np.array(["re_norm", "aoa_norm", "x", "y", "sdf"])
+TARGET_NAMES = np.array(["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"])
+
+
+def write_toy_simulator_dataset(root: Path, *, cases: int = 4, points: int = 64) -> None:
+    root.mkdir(parents=True)
+    rng = np.random.default_rng(123)
+    for case_index in range(cases):
+        re_norm = np.full(points, -0.5 + 0.25 * case_index, dtype=np.float32)
+        aoa_norm = np.full(points, -0.2 + 0.15 * case_index, dtype=np.float32)
+        x = rng.uniform(-1.0, 1.0, size=points).astype(np.float32)
+        y = rng.uniform(-1.0, 1.0, size=points).astype(np.float32)
+        sdf = (np.sqrt(x * x + y * y) - 0.5).astype(np.float32)
+        features = np.stack([re_norm, aoa_norm, x, y, sdf], axis=1).astype(np.float32)
+        targets = np.stack(
+            [
+                0.5 * x + 0.2 * y + 0.1 * aoa_norm,
+                -0.3 * x + 0.7 * sdf,
+                x * y + 0.05 * re_norm,
+                sdf**2 + 0.1 * y,
+            ],
+            axis=1,
+        ).astype(np.float32)
+        np.savez(
+            root / f"case_{case_index:02d}.npz",
+            features=features,
+            targets=targets,
+            feature_names=FEATURE_NAMES,
+            target_names=TARGET_NAMES,
+        )
+
+
+def write_training_config(
+    path: Path,
+    *,
+    data_root: Path,
+    artifact_dir: Path,
+    device_type: str,
+    allow_cpu_fallback: bool = False,
+) -> None:
+    path.write_text(
+        f"""
+[run]
+name = "test_mlp"
+seed = 0
+artifact_dir = "{artifact_dir}"
+
+[data]
+root = "{data_root}"
+train_cases = 2
+val_cases = 1
+test_cases = 1
+points_per_case = 64
+batch_size = 64
+
+[model]
+type = "mlp"
+hidden_width = 128
+depth = 4
+activation = "gelu"
+
+[optim]
+lr = 0.01
+weight_decay = 0.0
+steps = 1000
+log_interval = 250
+
+[device]
+type = "{device_type}"
+allow_cpu_fallback = {str(allow_cpu_fallback).lower()}
+benchmark_kernels = true
+
+[loss]
+type = "normalized_mse"
+""".strip()
+        + "\n"
+    )
+
+
+class TrainingLoopTests(unittest.TestCase):
+    def test_cuda_config_fails_clearly_when_cuda_unavailable(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            config_path = tmp_path / "config.toml"
+            write_training_config(
+                config_path,
+                data_root=tmp_path / "data",
+                artifact_dir=tmp_path / "artifacts",
+                device_type="cuda",
+            )
+            config = load_training_config(config_path)
+
+            with patch("torch.cuda.is_available", return_value=False):
+                with self.assertRaisesRegex(RuntimeError, "CUDA requested"):
+                    select_device(config)
+
+    def test_mlp_training_smoke_learns_tiny_deterministic_simulator(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            data_root = tmp_path / "data"
+            artifact_dir = tmp_path / "artifacts"
+            config_path = tmp_path / "config.toml"
+            write_toy_simulator_dataset(data_root, cases=4, points=64)
+            device_type = "cuda" if torch.cuda.is_available() else "cpu"
+            write_training_config(
+                config_path,
+                data_root=data_root,
+                artifact_dir=artifact_dir,
+                device_type=device_type,
+            )
+
+            result = subprocess.run(
+                [sys.executable, "-m", "airfrans_frontier.cli", "train", str(config_path)],
+                text=True,
+                capture_output=True,
+                check=False,
+            )
+            self.assertEqual(result.returncode, 0, msg=f"stdout={result.stdout}\nstderr={result.stderr}")
+            run_dir_line = next(line for line in result.stdout.splitlines() if line.startswith("run_dir: "))
+            run_dir = Path(run_dir_line.removeprefix("run_dir: "))
+            final_metrics = json.loads((run_dir / "final_metrics.json").read_text())
+
+            self.assertTrue(np.isfinite(final_metrics["train_loss"]))
+            self.assertLess(final_metrics["train_loss"], final_metrics["initial_train_loss"] * 0.1)
+            self.assertLess(final_metrics["train_loss"], 1e-2)
+            self.assertTrue((run_dir / "checkpoint.pt").exists())
+            self.assertTrue((run_dir / "metrics.jsonl").exists())
+            self.assertEqual(final_metrics["device"].startswith("cuda"), torch.cuda.is_available())
+            if torch.cuda.is_available():
+                self.assertIn("T550", final_metrics["gpu_name"])
+
+
+if __name__ == "__main__":
+    unittest.main()