ai-modcon/dsagt/xgc-ai-training

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`

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.

python scripts/check_xgc_structure.py <case_dir> [--output audit/step1_pre.json]

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.

python scripts/xgc_summarize.py <case_dir> [--output audit/xgc_summary.json]

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.

python scripts/xgc_preprocess.py <case_dir> <out_dir> \
    [--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]

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.

python scripts/check_xgc_preprocessed.py <out_dir> [--output audit/step3_post.json]

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. That module (graph_datasets) is not publicly available; without it, stop after Stage 4 and report the validated npz directory as the deliverable.

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.

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:

python scripts/xgc_dataset.py <out_dir> --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.

1. Which steps to process — all (default) or a subset for a quick test.
2. Output directory — where the npz files will be written.
3. `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.