diff --git a/use_cases/fusion-fm/README.md b/use_cases/fusion-fm/README.md new file mode 100644 index 0000000..858a18a --- /dev/null +++ b/use_cases/fusion-fm/README.md @@ -0,0 +1,50 @@ +# Fusion Foundation Model — XGC AI Training Use Case + +**Domain:** Plasma physics — gyrokinetic turbulence simulation +**Simulation code:** XGC (X-point Gyrokinetic Code) +**Data format:** ADIOS2 BP5 (one directory per simulation run) +**Goal:** Produce GNN-ready npz files and a PyTorch Dataset for training plasma surrogate models + +## Data + +Three example simulation cases live in `example_xgc_data/`: + +| Case | Machine | Nodes | nphi | Steps | Has f3d | +|------|---------|-------|------|-------|---------| +| `n560fr_ITER_PFPO_W_Ne` | ITER | 1,277,797 | 32 | 99 | yes (every 2) | +| `n613fr_KSTART_30306_q4_rmp_turbulence` | KSTART | 45,291 | 32 | 99 | yes (every 10) | +| `ti316_NSTX_small_dt_from_ti313` | NSTX | 86,513 | 16 | 999 | no | + +Each case directory contains: +- `xgc.mesh.bp` — static 2D poloidal mesh (R/Z coordinates, triangle connectivity, flux surfaces) +- `xgc.3d.NNNNN.bp` — per-timestep field snapshots (`dpot`, `eden`, `iden`, shape `[nphi, n_nodes]` or `[n_nodes, nphi]`) +- `xgc.f3d.NNNNN.bp` — fluid moments (`e_den`, `e_T_para/perp`, `e_u_para`, `i_*` equivalents) where available + +## Skill + +See [`skills/xgc-ai-training/SKILL.md`](skills/xgc-ai-training/SKILL.md) for the full agent pipeline. + +## Quick Start (manual) + +```bash +# 1. Check case structure +python skills/xgc-ai-training/scripts/check_xgc_structure.py example_xgc_data/n613fr_KSTART_30306_q4_rmp_turbulence + +# 2. Summarize physics content +python skills/xgc-ai-training/scripts/xgc_summarize.py example_xgc_data/n613fr_KSTART_30306_q4_rmp_turbulence + +# 3. Preprocess to npz +python skills/xgc-ai-training/scripts/xgc_preprocess.py \ + example_xgc_data/n613fr_KSTART_30306_q4_rmp_turbulence \ + /tmp/kstart_npz + +# 4. Validate output +python skills/xgc-ai-training/scripts/check_xgc_preprocessed.py /tmp/kstart_npz + +# 5. Use the dataset in training +python skills/xgc-ai-training/scripts/xgc_dataset.py /tmp/kstart_npz +``` + +## Variable Reference + +See [`skills/xgc-ai-training/references/xgc_fields.md`](skills/xgc-ai-training/references/xgc_fields.md). diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/SKILL.md b/use_cases/fusion-fm/skills/xgc-ai-training/SKILL.md new file mode 100644 index 0000000..96275a6 --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/SKILL.md @@ -0,0 +1,190 @@ +--- +name: xgc-ai-training +description: > + Convert XGC plasma turbulence simulation data (ADIOS2 BP5 format) into + GNN-ready npz files and a PyTorch Dataset for AI/surrogate model training. + Use when the user asks to preprocess XGC data, create training datasets from + XGC simulations, or prepare fusion simulation data for machine learning. +--- + +# XGC → AI Training Pipeline + +Converts raw XGC simulation directories into graph-structured training data +for MeshGraphNets-style surrogate models of plasma turbulence. + +**Variable reference:** [`references/xgc_fields.md`](references/xgc_fields.md) + +--- + +## Pipeline Checklist + +Copy and track progress: + +``` +- [ ] 1. Check XGC case structure (pre-flight) +- [ ] 2. Summarize physics content (audit) +- [ ] 3. Preprocess BP files → npz +- [ ] 4. Validate preprocessed output (post-flight) +- [ ] 5. Confirm dataset class is usable +``` + +--- + +## Stage 1 — Check XGC Structure (check) + +Run **before** preprocessing. Verifies mesh file, counts 3d/f3d steps, detects +nphi and axis order, flags inconsistencies. + +```bash +python scripts/check_xgc_structure.py [--output audit/step1_pre.json] +``` + +| Parameter | Required | Description | +|-----------|----------|-------------| +| `case_dir` | yes | XGC simulation directory | +| `--output` | no | Save JSON report to file (audit trail) | + +**Expect:** `"status": "ok"` and all checks passed before proceeding. +**Watch for:** missing `xgc.mesh.bp`, axis order (`first` vs `last`), f3d alignment. + +--- + +## Stage 2 — Physics Summary (operation → audit) + +Reads mesh and field files; emits structured JSON with node counts, field +names, physical parameters, and time range. + +```bash +python scripts/xgc_summarize.py [--output audit/xgc_summary.json] +``` + +| Parameter | Required | Description | +|-----------|----------|-------------| +| `case_dir` | yes | XGC simulation directory | +| `--output` | no | Save JSON report to file | + +**Key output fields:** `mesh.n_nodes`, `mesh.n_edges`, `fields_3d.phi_field_names`, +`fields_f3d.phi_field_names`, `physical_params.sml_dt`. + +--- + +## Stage 3 — Preprocess BP → npz (operation) + +Reads all `xgc.3d.*.bp` (and `xgc.f3d.*.bp` where available). Normalizes +axis order to `[nphi, n_nodes]` float32. Writes `mesh.npz`, `step_NNNNN.npz`, +and `meta.json` under `out_dir`. + +```bash +python scripts/xgc_preprocess.py \ + [--fields_3d dpot,eden,iden] \ + [--fields_f3d e_den,e_T_para,e_T_perp,e_u_para,i_den,i_T_para,i_T_perp,i_u_para] \ + [--no_f3d] \ + [--steps 10,20,30] \ + [--overwrite] \ + [--output audit/step3_op.json] +``` + +| Parameter | Required | Default | Description | +|-----------|----------|---------|-------------| +| `case_dir` | yes | — | XGC simulation directory | +| `out_dir` | yes | — | Output directory (created if absent) | +| `--fields_3d` | no | `dpot,eden,iden` | Fields to read from xgc.3d files | +| `--fields_f3d` | no | all 8 fluid moments | Fields to read from xgc.f3d files | +| `--no_f3d` | no | false | Skip f3d files entirely | +| `--steps` | no | all | Comma-separated step indices to process | +| `--overwrite` | no | false | Overwrite existing npz files | +| `--output` | no | — | Save JSON status to file | + +**Output layout:** +``` +out_dir/ + mesh.npz rz[N,2], conn[C,3], psi[N], region[N], nextnode[N] + meta.json n_nodes, nphi, steps, field_names, field_availability + step_00002.npz dpot[nphi,N], eden[nphi,N], ... + step_00004.npz ... +``` + +**Typical run times (KSTART, 45K nodes, all 99 steps):** ~5 min on login node. + +--- + +## Stage 4 — Validate Preprocessed Output (check) + +Run **after** preprocessing. Spot-checks 5 step files for shape, dtype, field +completeness, and time monotonicity. + +```bash +python scripts/check_xgc_preprocessed.py [--output audit/step3_post.json] +``` + +| Parameter | Required | Description | +|-----------|----------|-------------| +| `out_dir` | yes | Directory produced by `xgc_preprocess.py` | +| `--output` | no | Save JSON report to file | + +**Expect:** `"status": "ok"`, all shapes `[nphi, n_nodes]`, dtype `float32`. + +--- + +## Stage 5 — Dataset Class (library) + +`scripts/xgc_dataset.py` provides `XGCGraphDataset`, a PyTorch Dataset +subclassing `BaseCFDGraphDataset` from the MATEY project. + +On first use it builds a cached `topology.pt` (edge_index, edge_attr, static +node context) then reads npz files at getitem time — efficient for large meshes. + +```python +from xgc_dataset import XGCGraphDataset, build_datasets + +# Single case +ds = XGCGraphDataset( + "/data/kstart_npz", + field_names=["dpot", "e_den", "e_T_para", "e_T_perp"], + n_steps=2, + leadtime_max=5, + split="train", + train_val_test=[0.7, 0.15, 0.15], +) +sample, bcs = ds[0] +# sample.x → [N, n_steps, 7+F] +# sample.y → [N, F] +# sample.pos → [N, 2] (R, Z) + +# Multi-case +splits = build_datasets( + ["/data/kstart_npz", "/data/nstx_npz"], + field_names=["dpot"], + n_steps=1, + leadtime_max=1, +) +train_ds = splits["train"] +``` + +**Smoke test:** +```bash +python scripts/xgc_dataset.py --n_steps 1 --leadtime_max 1 +``` + +--- + +## Decisions the Agent Must Confirm With the User + +1. **Which fields to include** — at minimum `dpot`; add fluid moments if f3d + files are present and downstream model needs them. +2. **Which steps to process** — all (default) or a subset for a quick test. +3. **Output directory** — where the npz files will be written. +4. **`n_steps` and `leadtime_max`** — depend on the downstream GNN architecture. + +--- + +## Notes + +- ITER (1.28M nodes) is large; preprocessing all 99 steps takes ~30 min. + Use `--steps 10,20,30` for a fast sanity check first. +- NSTX has axis order `[n_nodes, nphi]`; `xgc_preprocess.py` normalizes + automatically (verified by `check_xgc_preprocessed.py`). +- f3d files in KSTART appear every 10 steps; 3d every 2 steps. Steps without + f3d will only have `dpot`, `eden`, `iden`. +- `meta.json` records `field_availability` per field so the dataset class can + automatically filter to steps where all requested fields are present. diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/references/xgc_fields.md b/use_cases/fusion-fm/skills/xgc-ai-training/references/xgc_fields.md new file mode 100644 index 0000000..61796b9 --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/references/xgc_fields.md @@ -0,0 +1,72 @@ +# XGC Field Reference + +## Mesh (xgc.mesh.bp) — static, one file per case + +| Variable | Shape | Description | +|----------|-------|-------------| +| `rz` | [N, 2] | (R, Z) node coordinates in meters | +| `nd_connect_list` | [C, 3] | Triangle connectivity (0-based node indices) | +| `psi` | [N] | Poloidal magnetic flux at each node (Wb/rad) | +| `region` | [N] | Node region code (see table below) | +| `nextnode` | [N] | Next node along the magnetic field line | +| `n_n` | scalar | Total node count | +| `n_t` | scalar | Total triangle count | +| `nsurf` | scalar | Number of flux surfaces | + +**Region codes** + +| Code | Meaning | +|------|---------| +| 1 | Core (inside separatrix) | +| 2 | Edge / separatrix | +| 3 | Scrape-off layer (SOL) | +| 100 | Wall / boundary nodes | + +## Fields (xgc.3d.NNNNN.bp) — one file per timestep + +| Variable | Shape | Description | +|----------|-------|-------------| +| `dpot` | [nphi, N] or [N, nphi] | Perturbed electrostatic potential (V) | +| `eden` | [nphi, N] or [N, nphi] | Electron density perturbation | +| `iden` | [nphi, N] or [N, nphi] | Ion density perturbation | +| `pot0` | [N] | Flux-surface-averaged potential (V) | +| `nphi` | scalar | Number of toroidal planes | +| `time` | scalar | Physical time (s) | + +## Fluid moments (xgc.f3d.NNNNN.bp) — subset of timesteps + +| Variable | Shape | Description | +|----------|-------|-------------| +| `e_den` | [nphi, N] | Electron density | +| `e_T_para` | [nphi, N] | Electron parallel temperature (eV) | +| `e_T_perp` | [nphi, N] | Electron perpendicular temperature (eV) | +| `e_u_para` | [nphi, N] | Electron parallel flow velocity | +| `i_den` | [nphi, N] | Ion density | +| `i_T_para` | [nphi, N] | Ion parallel temperature (eV) | +| `i_T_perp` | [nphi, N] | Ion perpendicular temperature (eV) | +| `i_u_para` | [nphi, N] | Ion parallel flow velocity | +| `dpot` | [nphi, N] | Perturbed potential (same as xgc.3d) | + +## Axis order + +XGC field arrays appear in two layouts depending on the simulation: + +| Case | Layout | Notes | +|------|--------|-------| +| ITER, KSTART | `[nphi, n_nodes]` | phi-first (most common) | +| NSTX | `[n_nodes, nphi]` | phi-last (transposed) | + +`xgc_preprocess.py` normalizes all outputs to `[nphi, n_nodes]` float32. + +## Dataset-level conventions + +Each preprocessed case produces: +- `mesh.npz` — rz, conn, psi, region, nextnode +- `step_NNNNN.npz` — all fields at that step, shape `[nphi, n_nodes]` +- `meta.json` — n_nodes, nphi, step list, field_availability per field + +Node feature vector assembled by `xgc_dataset.py`: +``` +x = [R, Z, psi, region_oh(4), field_1, ..., field_F] shape [N, 7+F] +``` +Target `y = x[:, 7:]` (field values only at target timestep). diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/scripts/check_xgc_preprocessed.py b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/check_xgc_preprocessed.py new file mode 100644 index 0000000..515dae1 --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/check_xgc_preprocessed.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC + +"""Validate preprocessed XGC npz output before dataset use. + +Reads meta.json, spot-checks a sample of step_NNNNN.npz files: + - All declared fields present + - Shapes match [nphi, n_nodes] + - Dtype is float32 + - Time values are monotonically increasing + - Field coverage matches meta.json field_availability + +Output: JSON report to stdout (and optionally --output file). +Exit 0 if all checks pass; exit 1 if any fail. +""" + +import argparse +import json +import os +import sys + +try: + import numpy as np +except ImportError: + sys.exit('{"status": "error", "message": "numpy not available"}') + +_SPOT_CHECK_COUNT = 5 + + +def check(out_dir, output_path=None): + out_dir = os.path.abspath(out_dir) + checks = [] + warnings = [] + + def ok(name, detail=""): + checks.append({"check": name, "passed": True, "detail": str(detail)}) + + def fail(name, detail=""): + checks.append({"check": name, "passed": False, "detail": str(detail)}) + + # 1. meta.json exists and is valid + meta_path = os.path.join(out_dir, "meta.json") + if not os.path.exists(meta_path): + report = {"status": "error", "message": f"meta.json not found in {out_dir}"} + print(json.dumps(report, indent=2)) + return False + try: + with open(meta_path) as fp: + meta = json.load(fp) + ok("meta_json_valid", f"n_nodes={meta['n_nodes']}, nphi={meta['nphi']}, n_steps={meta['n_steps']}") + except Exception as e: + fail("meta_json_valid", str(e)) + report = {"status": "fail", "out_dir": out_dir, "checks": checks, "warnings": warnings} + print(json.dumps(report, indent=2)) + return False + + n_nodes = meta["n_nodes"] + nphi = meta["nphi"] + steps = meta["steps"] + fields = meta["field_names"] + + # 2. mesh.npz exists + mesh_path = os.path.join(out_dir, "mesh.npz") + if os.path.exists(mesh_path): + d = np.load(mesh_path) + rz_ok = "rz" in d and d["rz"].shape == (n_nodes, 2) + conn_ok = "conn" in d + ok("mesh_npz_valid", f"rz={d['rz'].shape if 'rz' in d else 'missing'}, conn={'present' if conn_ok else 'missing'}") + if not rz_ok: + fail("mesh_rz_shape", f"expected ({n_nodes}, 2), got {d['rz'].shape if 'rz' in d else 'missing'}") + else: + fail("mesh_npz_valid", "mesh.npz not found") + + # 3. Step file count matches meta + step_files = [os.path.join(out_dir, f"step_{s:05d}.npz") for s in steps] + present = [f for f in step_files if os.path.exists(f)] + if len(present) == len(steps): + ok("step_file_count", f"{len(present)}/{len(steps)} files present") + else: + fail("step_file_count", f"only {len(present)}/{len(steps)} step files found") + + # 4. Spot-check a sample of step files + indices = list(range(0, len(steps), max(1, len(steps) // _SPOT_CHECK_COUNT)))[:_SPOT_CHECK_COUNT] + checked_steps, shape_errors, dtype_errors, field_errors = [], [], [], [] + time_vals = [] + + for i in indices: + s = steps[i] + path = os.path.join(out_dir, f"step_{s:05d}.npz") + if not os.path.exists(path): + continue + try: + d = np.load(path) + except Exception as e: + warnings.append(f"step_{s:05d}.npz load error: {e}") + continue + + checked_steps.append(s) + + # Field presence + missing = [fn for fn in fields if fn not in d.files] + if missing: + field_errors.append(f"step {s}: missing {missing}") + + # Shape + dtype + for fn in fields: + if fn not in d: + continue + arr = d[fn] + if arr.shape != (nphi, n_nodes): + shape_errors.append(f"step {s} {fn}: {arr.shape} != ({nphi},{n_nodes})") + if arr.dtype != np.float32: + dtype_errors.append(f"step {s} {fn}: dtype={arr.dtype}") + + if "time" in d: + time_vals.append((s, float(d["time"]))) + + if field_errors: + fail("spot_check_fields", "; ".join(field_errors[:3])) + else: + ok("spot_check_fields", f"all fields present in {len(checked_steps)} sampled steps") + + if shape_errors: + fail("spot_check_shapes", "; ".join(shape_errors[:3])) + else: + ok("spot_check_shapes", f"all shapes [nphi={nphi}, n_nodes={n_nodes}] correct") + + if dtype_errors: + fail("spot_check_dtypes", "; ".join(dtype_errors[:3])) + else: + ok("spot_check_dtypes", "all fields are float32") + + # 5. Time monotonicity + if len(time_vals) > 1: + times = [t for _, t in sorted(time_vals)] + monotone = all(times[i] < times[i + 1] for i in range(len(times) - 1)) + if monotone: + ok("time_monotone", f"range [{times[0]:.4e}, {times[-1]:.4e}] s") + else: + fail("time_monotone", "time values are not monotonically increasing") + + # 6. field_availability coverage + fa = meta.get("field_availability", {}) + coverage_issues = [] + for fn in fields: + if fn in fa and len(fa[fn]) == 0: + coverage_issues.append(fn) + if coverage_issues: + warnings.append(f"fields with zero coverage in field_availability: {coverage_issues}") + else: + ok("field_availability_coverage", f"{len(fields)} fields all have coverage entries") + + passed = all(c["passed"] for c in checks) + report = { + "status": "ok" if passed else "fail", + "out_dir": out_dir, + "checks": checks, + "warnings": warnings, + "summary": { + "n_steps_checked": len(steps), + "n_steps_spot_checked": len(checked_steps), + "fields": fields, + "n_nodes": n_nodes, + "nphi": nphi, + }, + } + + out = json.dumps(report, indent=2) + print(out) + if output_path: + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + with open(output_path, "w") as fp: + fp.write(out) + + return passed + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("out_dir", help="Preprocessed output directory (contains meta.json)") + parser.add_argument("--output", default=None, + help="Write JSON report to this file (in addition to stdout)") + args = parser.parse_args() + + passed = check(args.out_dir, args.output) + sys.exit(0 if passed else 1) + + +if __name__ == "__main__": + main() diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/scripts/check_xgc_structure.py b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/check_xgc_structure.py new file mode 100644 index 0000000..52d65b3 --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/check_xgc_structure.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC + +"""Check XGC simulation directory structure before preprocessing. + +Verifies mesh file presence, counts 3d/f3d timestep files, detects +nphi and axis order, and flags any structural inconsistencies. + +Output: JSON report to stdout (and optionally --output file). +Exit 0 if all checks pass; exit 1 if any check fails. +""" + +import argparse +import glob +import json +import os +import re +import sys + +try: + import adios2 +except ImportError: + sys.exit('{"status": "error", "message": "adios2 not available"}') + + +def _parse_step(filename): + m = re.search(r'\.(\d+)\.bp$', filename) + return int(m.group(1)) if m else -1 + + +def _bp_size_mb(path): + total = 0 + for root, _, files in os.walk(path): + for fn in files: + try: + total += os.path.getsize(os.path.join(root, fn)) + except OSError: + pass + return total / 1024 ** 2 + + +def _get_nphi_and_axis(bp_path): + try: + with adios2.FileReader(bp_path) as f: + avail = f.available_variables() + nphi = int(f.read("nphi")) if "nphi" in avail else None + for name in ("dpot", "eden", "iden"): + if name in avail: + shape_str = avail[name]["Shape"] + dims = [int(x.strip()) for x in shape_str.split(",")] + if nphi and len(dims) == 2: + axis = "first" if dims[0] == nphi else "last" + return nphi, axis + return nphi, None + except Exception as e: + return None, str(e) + + +def check(case_dir, output_path=None): + case_dir = os.path.abspath(case_dir) + checks = [] + warnings = [] + + def ok(name, detail=""): + checks.append({"check": name, "passed": True, "detail": detail}) + + def fail(name, detail=""): + checks.append({"check": name, "passed": False, "detail": detail}) + + # 1. Directory exists + if not os.path.isdir(case_dir): + report = {"status": "error", "message": f"Not a directory: {case_dir}"} + print(json.dumps(report, indent=2)) + return False + ok("directory_exists", case_dir) + + # 2. Mesh file + mesh_path = os.path.join(case_dir, "xgc.mesh.bp") + if os.path.exists(mesh_path): + ok("mesh_file_present", f"{_bp_size_mb(mesh_path):.1f} MB") + else: + fail("mesh_file_present", "xgc.mesh.bp not found") + + # 3. Count 3d files + files_3d = sorted(glob.glob(os.path.join(case_dir, "xgc.3d.*.bp"))) + steps_3d = [_parse_step(f) for f in files_3d] + if files_3d: + stride = (steps_3d[1] - steps_3d[0]) if len(steps_3d) > 1 else None + ok("3d_files_found", f"{len(files_3d)} files, steps {steps_3d[0]}–{steps_3d[-1]}, stride {stride}") + else: + fail("3d_files_found", "no xgc.3d.*.bp files found") + + # 4. Count f3d files (optional) + files_f3d = sorted(glob.glob(os.path.join(case_dir, "xgc.f3d.*.bp"))) + steps_f3d = [_parse_step(f) for f in files_f3d] + if files_f3d: + ok("f3d_files_found", f"{len(files_f3d)} files, steps {steps_f3d[0]}–{steps_f3d[-1]}") + # Check f3d steps are a subset of 3d steps + orphans = [s for s in steps_f3d if s not in set(steps_3d)] + if orphans: + warnings.append(f"f3d steps not in 3d step list: {orphans[:5]}") + else: + ok("f3d_files_found", "none (3d-only dataset)") + + # 5. Probe first 3d file for nphi and axis order + if files_3d: + nphi, axis = _get_nphi_and_axis(files_3d[0]) + if nphi: + detail = f"nphi={nphi}, axis_order={axis} ({'[nphi,N]' if axis == 'first' else '[N,nphi]' if axis == 'last' else 'unknown'})" + ok("nphi_detected", detail) + else: + fail("nphi_detected", f"could not read nphi: {axis}") + + # 6. Mesh node count consistency + if os.path.exists(mesh_path) and files_3d: + try: + with adios2.FileReader(mesh_path) as f: + n_n = int(f.read("n_n")) + with adios2.FileReader(files_3d[0]) as f: + avail = f.available_variables() + nnode = int(f.read("nnode")) if "nnode" in avail else None + if nnode is not None and nnode != n_n: + fail("node_count_consistent", f"mesh n_n={n_n} != 3d nnode={nnode}") + else: + ok("node_count_consistent", f"n_nodes={n_n}") + except Exception as e: + warnings.append(f"node count check error: {e}") + + passed = all(c["passed"] for c in checks) + report = { + "status": "ok" if passed else "fail", + "case_dir": case_dir, + "checks": checks, + "warnings": warnings, + "summary": { + "n_3d_steps": len(files_3d), + "n_f3d_steps": len(files_f3d), + "nphi": nphi if files_3d else None, + "axis_order": axis if files_3d else None, + }, + } + + out = json.dumps(report, indent=2) + print(out) + if output_path: + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + with open(output_path, "w") as fp: + fp.write(out) + + return passed + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("case_dir", help="XGC simulation directory to check") + parser.add_argument("--output", default=None, + help="Write JSON report to this file (in addition to stdout)") + args = parser.parse_args() + + passed = check(args.case_dir, args.output) + sys.exit(0 if passed else 1) + + +if __name__ == "__main__": + main() diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_dataset.py b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_dataset.py new file mode 100644 index 0000000..83e3a2a --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_dataset.py @@ -0,0 +1,346 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC + +"""XGC Graph Dataset for AI training. + +Reads preprocessed npz files (from preprocess_xgc.py) and provides +PyG graph-structured training samples. + +Each sample: (Data, bcs) where Data has: + x : [N, n_steps, 7+F] – input sequence (pos + psi + region_oh + fields) + y : [N, F] – target field values + pos : [N, 2] – (R, Z) node positions + edge_index : [2, E] + edge_attr : [E, 3] – (dR, dZ, |d|) + leadtime : [1, 1] float32 + phi : int – toroidal plane index + step0 : int – first input step + target_step: int – target step + +Training samples are indexed by (phi_plane, start_timestep) pairs. +Splits are made by phi-plane group (all steps from a plane stay together). +""" + +import json +import os +import random +from dataclasses import dataclass +from typing import List, Optional + +import numpy as np +import torch +import torch.nn.functional as F +from torch_geometric.data import Data + +from graph_datasets import BaseCFDGraphDataset, SampleId + +# XGC region codes → class indices 0-3 +_REGION_MAP = {1: 0, 2: 1, 3: 2, 100: 3} +_NUM_REGION_TYPES = 4 +_TOPOLOGY_FILE = "topology.pt" + + +class XGCGraphDataset(BaseCFDGraphDataset): + """ + PyG Dataset for XGC plasma turbulence simulations. + + Workflow: + 1. Run preprocess_xgc.py to produce npz files. + 2. Instantiate XGCGraphDataset – on first use it builds a cached + topology.pt (edge_index, edge_attr, static node context) and + writes an index.json under path/{split}/processed/. + 3. Each __getitem__ call reads the relevant step_NNNNN.npz files, + slices the requested phi plane, and assembles a PyG Data graph. + + Args: + path : root directory with mesh.npz, meta.json, step_*.npz + field_names : fields to use as node features; must be keys in npz + (default: all fields in meta.json) + phi_planes : list of toroidal plane indices to include + (default: all planes 0 … nphi-1) + require_all_fields : if True (default), drop steps missing any field + **kwargs : forwarded to BaseCFDGraphDataset + (split, n_steps, leadtime_max, train_val_test, …) + """ + + # ── class-level helpers (static) ──────────────────────────────────────── + + @staticmethod + def _specifics(): + """Placeholder; real values come from _specifics(self) instance method.""" + return ["dpot"], "xgc", None, _NUM_REGION_TYPES + + # ── construction ──────────────────────────────────────────────────────── + + def __init__( + self, + path: str, + field_names: Optional[List[str]] = None, + phi_planes: Optional[List[int]] = None, + require_all_fields: bool = True, + **kwargs, + ): + meta_path = os.path.join(path, "meta.json") + if not os.path.exists(meta_path): + raise FileNotFoundError( + f"meta.json not found in {path}. Run preprocess_xgc.py first." + ) + with open(meta_path) as fp: + self._meta = json.load(fp) + + self._nphi = self._meta["nphi"] + self._n_nodes = self._meta["n_nodes"] + + # Resolve field names + avail_fields = self._meta["field_names"] + if field_names is None: + self._field_names_data = avail_fields + else: + missing = [f for f in field_names if f not in avail_fields] + if missing: + raise ValueError(f"Requested fields not in meta.json: {missing}") + self._field_names_data = list(field_names) + + # Resolve step list (drop steps missing any requested field) + field_avail = self._meta.get("field_availability", {}) + step_list = self._meta["steps"] + if require_all_fields and field_avail: + valid = set(step_list) + for fn in self._field_names_data: + if fn in field_avail: + valid &= set(field_avail[fn]) + step_list = sorted(valid) + self._step_list = step_list + + # Resolve phi planes + self._phi_planes = ( + list(phi_planes) if phi_planes is not None + else list(range(self._nphi)) + ) + + # Cache for topology (lazy-loaded in __getitem__) + self._topo_cache: Optional[dict] = None + + super().__init__(path, **kwargs) + + # ── specifics (instance override) ─────────────────────────────────────── + + def _specifics(self): + return ( + self._field_names_data, + "xgc", + len(self._step_list), + _NUM_REGION_TYPES, + ) + + # ── one-time processing ────────────────────────────────────────────────── + + def _load_mesh_arrays(self): + d = np.load(os.path.join(self.path, "mesh.npz")) + rz = d["rz"] # [N, 2] + conn = d["conn"] # [C, 3] + psi = d["psi"] if "psi" in d else None + region = d["region"] if "region" in d else None + return rz, conn, psi, region + + def _map_region(self, region: np.ndarray) -> np.ndarray: + mapped = np.zeros_like(region) + for code, idx in _REGION_MAP.items(): + mapped[region == code] = idx + return mapped + + def process(self): + """Build cached topology.pt and write index.json. Called once.""" + rz, conn, psi, region = self._load_mesh_arrays() + n_nodes = self._n_nodes + + pos = torch.as_tensor(rz, dtype=torch.float32) # [N, 2] + edge_index = self.cells_to_edge_index(conn, num_nodes=n_nodes, undirected=True) + edge_attr = self.mesh_edge_attr(pos, edge_index) # [E, 3] + + psi_t = ( + torch.as_tensor(psi, dtype=torch.float32).unsqueeze(1) + if psi is not None + else torch.zeros(n_nodes, 1) + ) + + if region is not None: + reg_idx = torch.as_tensor(self._map_region(region), dtype=torch.long) + region_oh = self.one_hot_node_type(reg_idx, _NUM_REGION_TYPES) # [N, 4] + else: + region_oh = torch.zeros(n_nodes, _NUM_REGION_TYPES) + + # static_ctx: [R, Z, psi, r0, r1, r2, r3] → [N, 7] + static_ctx = torch.cat([pos, psi_t, region_oh], dim=-1) + + topo = dict( + pos=pos, + edge_index=edge_index, + edge_attr=edge_attr, + static_ctx=static_ctx, + ) + topo_path = os.path.join(self.processed_dir, _TOPOLOGY_FILE) + torch.save(topo, topo_path) + print(f" Topology saved: {topo_path}") + print(f" nodes={n_nodes} edges={edge_index.shape[1]}") + + index_obj = { + "version": 1, + "n_nodes": n_nodes, + "nphi": self._nphi, + "phi_planes": self._phi_planes, + "steps": self._step_list, + "field_names": self._field_names_data, + } + with open(self.processed_index, "w") as fp: + json.dump(index_obj, fp, indent=2) + + # ── sample discovery ───────────────────────────────────────────────────── + + def discover_samples(self) -> List[SampleId]: + """Return all (phi_plane, start_step) pairs valid as input starts.""" + n_steps = len(self._step_list) + # need n_in input steps + at least 1 future step for target + max_start = n_steps - self.nsteps_input - 1 + if max_start < 0: + return [] + + samples = [] + for ip in self._phi_planes: + group = f"phi_{ip:03d}" + for si in range(max_start + 1): + step = self._step_list[si] + samples.append( + SampleId(group=group, item=f"step_{step:05d}", path="", t=step) + ) + return samples + + # ── getitem ────────────────────────────────────────────────────────────── + + def _load_topo(self) -> dict: + if self._topo_cache is None: + topo_path = os.path.join(self.processed_dir, _TOPOLOGY_FILE) + self._topo_cache = torch.load(topo_path, weights_only=False, map_location="cpu") + return self._topo_cache + + def _load_phi_fields(self, step: int, ip: int) -> torch.Tensor: + """Read one phi plane from step_NNNNN.npz → [N, F] float32 tensor.""" + npz_path = os.path.join(self.path, f"step_{step:05d}.npz") + d = np.load(npz_path) + vecs = [] + for fn in self._field_names_data: + if fn in d: + vecs.append(d[fn][ip].astype(np.float32)) # [N] + else: + vecs.append(np.zeros(self._n_nodes, dtype=np.float32)) + return torch.as_tensor(np.stack(vecs, axis=-1), dtype=torch.float32) # [N, F] + + def __getitem__(self, index): + base_idx = self.active_indices[index] + meta = self.samples[base_idx] + + step_start = meta.t + ip = int(meta.group.split("_")[1]) # phi plane index + + si_start = self._step_list.index(step_start) + n_in = self.nsteps_input + max_lead = len(self._step_list) - si_start - n_in + leadtime = torch.randint(1, max(2, min(self.leadtime_max + 1, max_lead + 1)), (1,)) + si_target = si_start + n_in + leadtime.item() - 1 + + topo = self._load_topo() + static_ctx = topo["static_ctx"] # [N, 7] + pos = topo["pos"] + edge_index = topo["edge_index"] + edge_attr = topo["edge_attr"] + + # Build input sequence: [N, n_in, 7+F] + x_list = [] + for si in range(si_start, si_start + n_in): + s = self._step_list[si] + fields_t = self._load_phi_fields(s, ip) # [N, F] + x_list.append(torch.cat([static_ctx, fields_t], dim=-1)) # [N, 7+F] + x_seq = torch.stack(x_list, dim=0).permute(1, 0, 2) # [N, n_in, 7+F] + + # Target: field values only [N, F] + target_step = self._step_list[si_target] + y = self._load_phi_fields(target_step, ip) + + data = Data( + x=x_seq, + y=y, + pos=pos, + edge_index=edge_index, + edge_attr=edge_attr, + ) + data.step0 = step_start + data.target_step = target_step + data.phi = ip + data.leadtime = leadtime.reshape(-1, 1).to(torch.float32) + + bcs = self._get_specific_bcs() + return data, bcs + + +# ── convenience helpers ────────────────────────────────────────────────────── + +def build_datasets( + preprocessed_dirs: List[str], + field_names: Optional[List[str]] = None, + phi_planes: Optional[List[int]] = None, + n_steps: int = 1, + leadtime_max: int = 1, + train_val_test=(0.7, 0.15, 0.15), + **kwargs, +) -> dict: + """Build train/val/test datasets from multiple preprocessed case directories. + + Returns {'train': ds, 'val': ds, 'test': ds} using ConcatDataset. + """ + from torch.utils.data import ConcatDataset + + splits = {} + for split in ("train", "val", "test"): + datasets = [] + for pdir in preprocessed_dirs: + ds = XGCGraphDataset( + pdir, + field_names=field_names, + phi_planes=phi_planes, + n_steps=n_steps, + leadtime_max=leadtime_max, + split=split, + train_val_test=list(train_val_test), + **kwargs, + ) + datasets.append(ds) + splits[split] = ConcatDataset(datasets) if len(datasets) > 1 else datasets[0] + return splits + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Smoke-test XGCGraphDataset") + parser.add_argument("preprocessed_dir", + help="Directory produced by preprocess_xgc.py") + parser.add_argument("--n_steps", type=int, default=1) + parser.add_argument("--leadtime_max", type=int, default=1) + args = parser.parse_args() + + ds = XGCGraphDataset( + args.preprocessed_dir, + n_steps=args.n_steps, + leadtime_max=args.leadtime_max, + split="train", + train_val_test=[0.7, 0.15, 0.15], + ) + print(f"Dataset: {len(ds)} training samples") + sample, bcs = ds[0] + print(f" x shape : {sample.x.shape}") # [N, n_steps, 7+F] + print(f" y shape : {sample.y.shape}") # [N, F] + print(f" pos shape : {sample.pos.shape}") + print(f" edge_index : {sample.edge_index.shape}") + print(f" edge_attr : {sample.edge_attr.shape}") + print(f" phi={sample.phi} step0={sample.step0} target={sample.target_step} lead={sample.leadtime.item():.0f}") diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_preprocess.py b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_preprocess.py new file mode 100644 index 0000000..7dc6672 --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_preprocess.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC + +"""Preprocess an XGC simulation directory into npz files for AI training. + +Reads ADIOS2 BP files and writes: + out_dir/mesh.npz - static mesh (rz, conn, psi, region, nextnode) + out_dir/step_NNNNN.npz - per-timestep fields, shape [nphi, n_nodes] float32 + out_dir/meta.json - case metadata, step lists, field availability + +Output: JSON status to stdout (and optionally --output file). +Exit 0 on success; exit 1 on error. +""" + +import argparse +import glob +import json +import os +import re +import sys + +try: + import adios2 + import numpy as np +except ImportError as e: + sys.exit(f'{{"status": "error", "message": "{e}"}}') + +_FIELDS_3D = ["dpot", "eden", "iden"] +_FIELDS_F3D = ["e_den", "e_T_para", "e_T_perp", "e_u_para", + "i_den", "i_T_para", "i_T_perp", "i_u_para"] + + +def _parse_step(filename): + m = re.search(r'\.(\d+)\.bp$', filename) + return int(m.group(1)) if m else -1 + + +def _find_steps(case_dir, pattern): + files = sorted(glob.glob(os.path.join(case_dir, pattern))) + return {_parse_step(f): f for f in files if _parse_step(f) >= 0} + + +def _normalize_phi_first(arr, nphi, n_nodes): + if arr.shape == (nphi, n_nodes): + return arr.astype(np.float32) + if arr.shape == (n_nodes, nphi): + return arr.T.astype(np.float32) + raise ValueError(f"Shape {arr.shape} does not match ({nphi},{n_nodes}) or ({n_nodes},{nphi})") + + +def _get_nphi(bp_path): + with adios2.FileReader(bp_path) as f: + avail = f.available_variables() + if "nphi" in avail: + return int(f.read("nphi")) + for name in ("dpot", "eden", "iden"): + if name in avail: + dims = [int(x.strip()) for x in avail[name]["Shape"].split(",")] + return min(dims) + return None + + +def _read_mesh(case_dir): + mesh_path = os.path.join(case_dir, "xgc.mesh.bp") + if not os.path.exists(mesh_path): + raise FileNotFoundError(f"xgc.mesh.bp not found in {case_dir}") + arrays = {} + with adios2.FileReader(mesh_path) as f: + avail = f.available_variables() + arrays["n_nodes"] = np.array(int(f.read("n_n")), dtype=np.int32) + arrays["n_cells"] = np.array(int(f.read("n_t")), dtype=np.int32) + arrays["rz"] = f.read("rz").astype(np.float32) + arrays["conn"] = f.read("nd_connect_list").astype(np.int32) + for name in ("psi", "theta"): + if name in avail: + arrays[name] = f.read(name).astype(np.float32) + for name in ("region", "nextnode"): + if name in avail: + arrays[name] = f.read(name).astype(np.int32) + return arrays + + +def _read_fields(bp_path, field_names, nphi, n_nodes): + out, time_val = {}, None + with adios2.FileReader(bp_path) as f: + avail = f.available_variables() + if "time" in avail: + time_val = float(f.read("time")) + for name in field_names: + if name not in avail: + continue + try: + out[name] = _normalize_phi_first(f.read(name), nphi, n_nodes) + except ValueError as e: + print(f" WARNING: skipping {name}: {e}", file=sys.stderr) + return out, time_val + + +def preprocess(case_dir, out_dir, fields_3d=None, fields_f3d=None, + steps=None, no_f3d=False, overwrite=False): + case_dir = os.path.abspath(case_dir) + out_dir = os.path.abspath(out_dir) + os.makedirs(out_dir, exist_ok=True) + + fields_3d = fields_3d or _FIELDS_3D + fields_f3d = fields_f3d or _FIELDS_F3D + + mesh = _read_mesh(case_dir) + n_nodes = int(mesh["n_nodes"]) + + mesh_out = os.path.join(out_dir, "mesh.npz") + if overwrite or not os.path.exists(mesh_out): + np.savez_compressed(mesh_out, **mesh) + + step_3d = _find_steps(case_dir, "xgc.3d.*.bp") + step_f3d = _find_steps(case_dir, "xgc.f3d.*.bp") if not no_f3d else {} + + if not step_3d: + raise RuntimeError(f"No xgc.3d.*.bp files in {case_dir}") + + nphi = _get_nphi(next(iter(step_3d.values()))) + + if steps is not None: + steps_to_do = sorted(s for s in step_3d if s in set(steps)) + else: + steps_to_do = sorted(step_3d) + + field_availability = {} + saved_steps = [] + + for step in steps_to_do: + out_path = os.path.join(out_dir, f"step_{step:05d}.npz") + if not overwrite and os.path.exists(out_path): + d = np.load(out_path) + for k in d.files: + if k not in ("time", "step", "has_f3d"): + field_availability.setdefault(k, []).append(step) + saved_steps.append(step) + continue + + save_dict = {} + data_3d, time_val = _read_fields(step_3d[step], fields_3d, nphi, n_nodes) + save_dict.update(data_3d) + + has_f3d = step in step_f3d + if has_f3d: + data_f3d, _ = _read_fields(step_f3d[step], fields_f3d, nphi, n_nodes) + save_dict.update(data_f3d) + + if not save_dict: + continue + + save_dict["step"] = np.array(step, dtype=np.int32) + save_dict["has_f3d"] = np.array(has_f3d, dtype=bool) + if time_val is not None: + save_dict["time"] = np.array(time_val, dtype=np.float64) + + np.savez_compressed(out_path, **save_dict) + saved_steps.append(step) + for k in save_dict: + if k not in ("time", "step", "has_f3d"): + field_availability.setdefault(k, []).append(step) + + all_fields = list(field_availability.keys()) + meta = { + "case_dir": case_dir, + "case_name": os.path.basename(case_dir), + "n_nodes": n_nodes, + "nphi": nphi, + "n_cells": int(mesh["n_cells"]), + "field_names": all_fields, + "steps": saved_steps, + "n_steps": len(saved_steps), + "f3d_steps": sorted(step_f3d.keys()), + "field_availability": field_availability, + } + with open(os.path.join(out_dir, "meta.json"), "w") as fp: + json.dump(meta, fp, indent=2) + + return { + "status": "ok", + "case_dir": case_dir, + "out_dir": out_dir, + "n_nodes": n_nodes, + "nphi": nphi, + "n_steps": len(saved_steps), + "field_names": all_fields, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("case_dir", help="XGC simulation directory") + parser.add_argument("out_dir", help="Output directory for npz files") + parser.add_argument("--fields_3d", default=",".join(_FIELDS_3D), + help=f"Comma-separated 3d fields (default: {','.join(_FIELDS_3D)})") + parser.add_argument("--fields_f3d", default=",".join(_FIELDS_F3D), + help=f"Comma-separated f3d fields (default: {','.join(_FIELDS_F3D)})") + parser.add_argument("--no_f3d", action="store_true", help="Skip f3d files") + parser.add_argument("--steps", default=None, + help="Comma-separated step indices (default: all)") + parser.add_argument("--overwrite", action="store_true", help="Overwrite existing files") + parser.add_argument("--output", default=None, + help="Write JSON status to this file (in addition to stdout)") + args = parser.parse_args() + + fields_3d = [f.strip() for f in args.fields_3d.split(",") if f.strip()] + fields_f3d = [f.strip() for f in args.fields_f3d.split(",") if f.strip()] + steps = [int(s) for s in args.steps.split(",")] if args.steps else None + + try: + report = preprocess( + args.case_dir, args.out_dir, + fields_3d=fields_3d, fields_f3d=fields_f3d, + steps=steps, no_f3d=args.no_f3d, overwrite=args.overwrite, + ) + except Exception as e: + report = {"status": "error", "message": str(e)} + + out = json.dumps(report, indent=2) + print(out) + if args.output: + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + with open(args.output, "w") as fp: + fp.write(out) + + sys.exit(0 if report.get("status") == "ok" else 1) + + +if __name__ == "__main__": + main() diff --git a/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_summarize.py b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_summarize.py new file mode 100644 index 0000000..25aaa48 --- /dev/null +++ b/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_summarize.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: 2026 UT-Battelle, LLC + +"""Summarize an XGC simulation directory and emit a structured JSON report. + +Reads xgc.mesh.bp, xgc.3d.*.bp, xgc.f3d.*.bp, and units.m to produce +a machine-readable summary suitable for audit trails and data cards. + +Output: JSON to stdout (and optionally --output file). +""" + +import argparse +import glob +import json +import os +import re +import sys + +try: + import adios2 + import numpy as np +except ImportError as e: + sys.exit(f'{{"status": "error", "message": "{e}"}}') + + +def _parse_step(filename): + m = re.search(r'\.(\d+)\.bp$', filename) + return int(m.group(1)) if m else -1 + + +def _bp_size_mb(path): + total = 0 + for root, _, files in os.walk(path): + for fn in files: + try: + total += os.path.getsize(os.path.join(root, fn)) + except OSError: + pass + return total / 1024 ** 2 + + +def _read_units(path): + params = {} + for fname in ("units.m", "units.txt"): + fp = os.path.join(path, fname) + if not os.path.exists(fp): + continue + with open(fp) as f: + for line in f: + line = line.strip().rstrip(";") + if "=" in line and not line.startswith(("!", "#")): + k, _, v = line.partition("=") + try: + params[k.strip()] = float(v.strip()) + except ValueError: + pass + break + return params + + +def _summarize_mesh(path): + mesh_path = os.path.join(path, "xgc.mesh.bp") + if not os.path.exists(mesh_path): + return None + info = {"size_mb": _bp_size_mb(mesh_path)} + with adios2.FileReader(mesh_path) as f: + avail = f.available_variables() + info["n_nodes"] = int(f.read("n_n")) + info["n_cells"] = int(f.read("n_t")) + if "rz" in avail: + rz = f.read("rz") + info["R_range_m"] = [float(rz[:, 0].min()), float(rz[:, 0].max())] + info["Z_range_m"] = [float(rz[:, 1].min()), float(rz[:, 1].max())] + if "psi" in avail: + psi = f.read("psi") + info["psi_range"] = [float(psi.min()), float(psi.max())] + if "region" in avail: + region = f.read("region") + info["region_codes"] = sorted(int(x) for x in np.unique(region).tolist()) + if "nd_connect_list" in avail: + conn = f.read("nd_connect_list") + pairs = np.concatenate([conn[:, [0, 1]], conn[:, [1, 2]], conn[:, [0, 2]]], axis=0) + pairs = np.sort(pairs, axis=1) + info["n_edges"] = int(len(np.unique(pairs, axis=0))) + info["has_psi"] = "psi" in avail + info["has_region"] = "region" in avail + info["has_nextnode"] = "nextnode" in avail + info["n_flux_surfaces"] = int(f.read("nsurf")) if "nsurf" in avail else None + return info + + +def _summarize_steps(path, pattern): + files = sorted(glob.glob(os.path.join(path, pattern))) + if not files: + return None + steps = [_parse_step(f) for f in files] + + # Read nphi and variable list from first file + nphi, axis_order, array_vars = None, None, {} + time_vals = [] + with adios2.FileReader(files[0]) as f: + avail = f.available_variables() + if "nphi" in avail: + nphi = int(f.read("nphi")) + if "time" in avail: + time_vals.append(float(f.read("time"))) + for k, v in avail.items(): + if v["SingleValue"] == "false" and "," in v.get("Shape", ""): + dims = [int(x.strip()) for x in v["Shape"].split(",")] + if nphi and len(dims) == 2 and nphi in dims: + array_vars[k] = v["Shape"] + if axis_order is None: + axis_order = "first" if dims[0] == nphi else "last" + + # Collect timestamps from remaining files + for fp in files[1:]: + try: + with adios2.FileReader(fp) as f: + avail = f.available_variables() + if "time" in avail: + time_vals.append(float(f.read("time"))) + except Exception: + pass + + stride = (steps[1] - steps[0]) if len(steps) > 1 else None + dt_phys = ((max(time_vals) - min(time_vals)) / (len(time_vals) - 1) + if len(time_vals) > 1 else None) + + return { + "n_files": len(files), + "step_range": [steps[0], steps[-1]], + "step_stride": stride, + "steps": steps, + "nphi": nphi, + "axis_order": axis_order, + "time_range_s": [min(time_vals), max(time_vals)] if time_vals else None, + "dt_phys_s": dt_phys, + "total_size_mb": sum(_bp_size_mb(f) for f in files), + "phi_field_names": sorted(array_vars.keys()), + } + + +def summarize(case_dir): + case_dir = os.path.abspath(case_dir) + units = _read_units(case_dir) + + keys_of_interest = [ + "sml_dt", "sml_tran", "vth", "eq_axis_r", "eq_axis_z", + "eq_axis_b", "eq_tempi_v1", "ptl_ion_mass_au", "ptl_ion_charge_eu", + ] + physical_params = {k: units[k] for k in keys_of_interest if k in units} + + return { + "status": "ok", + "case_dir": case_dir, + "case_name": os.path.basename(case_dir), + "physical_params": physical_params, + "mesh": _summarize_mesh(case_dir), + "fields_3d": _summarize_steps(case_dir, "xgc.3d.*.bp"), + "fields_f3d": _summarize_steps(case_dir, "xgc.f3d.*.bp"), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("case_dir", help="XGC simulation directory") + parser.add_argument("--output", default=None, + help="Write JSON report to this file (in addition to stdout)") + args = parser.parse_args() + + if not os.path.isdir(args.case_dir): + sys.exit(json.dumps({"status": "error", "message": f"Not a directory: {args.case_dir}"})) + + report = summarize(args.case_dir) + out = json.dumps(report, indent=2) + print(out) + if args.output: + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + with open(args.output, "w") as fp: + fp.write(out) + + +if __name__ == "__main__": + main()