Skip to content

Adding Fusion FM use case #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ajtritt merged 1 commit into from
Jun 26, 2026
Merged

Adding Fusion FM use case #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 50 additions & 0 deletions use_cases/fusion-fm/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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).
190 changes: 190 additions & 0 deletions use_cases/fusion-fm/skills/xgc-ai-training/SKILL.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 <case_dir> [--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 <case_dir> [--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 <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]
```

| 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 <out_dir> [--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 <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.
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.
72 changes: 72 additions & 0 deletions use_cases/fusion-fm/skills/xgc-ai-training/references/xgc_fields.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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).
Loading
Loading