OSIRIS adept module — usage overview
File structure
adept/ (repo root)
├── run.py ← CLI entry (--cfg path, no .yaml suffix)
├── adept/
│ ├── _base_.py ← dispatcher: `osiris` solver branch
│ └── osiris/ OSIRIS wrapper package
│ ├── __init__.py lazy export of BaseOsiris
│ ├── base.py ◀── BaseOsiris(ADEPTModule): __init__ parses deck,
│ │ __call__ runs OSIRIS, post_process delegates
│ ├── deck.py ◀── namelist parser/renderer/merger
│ │ parse_deck(text), parse_deck_file(path),
│ │ render_deck(sections), merge_overrides(...),
│ │ deck_to_flat_dict(...) for MLflow params
│ ├── runner.py ◀── subprocess driver
│ │ run_osiris(deck_text, binary=…, mpi_ranks=…),
│ │ discover_binary(...), OSIRIS-error detection
│ ├── post.py ◀── post-run collection: final-step HDF5 copy,
│ │ NetCDF export, scalar metrics
│ ├── io.py ◀── HDF5/NetCDF readers + dataset save/load
│ │ (load_series ±t_indices, lazy open_series, series_len)
│ ├── plots.py ◀── canned plot set (save_canned_plots)
│ └── regen.py ◀── regenerate plots offline from saved NetCDFs
├── configs/
│ └── osiris/ example manifests
│ ├── twostream-1d.yaml full run (deck tmax=100)
│ └── twostream-1d-short.yaml tmax=1.0 smoke
└── tests/
└── test_osiris/
├── decks/two-stream-1d in-repo example deck (manifests point here)
├── test_deck_roundtrip.py namelist parser round-trip
├── test_runner.py subprocess runner + discover_binary
├── test_units.py units.yaml derivation
├── test_post_netcdf.py post → NetCDF export
├── test_io_and_plots.py io readers + plotting
├── test_diagnostics_plots.py per-diagnostic plots
└── test_plots_new_views.py field-decomp / phase-space views
End-to-end data flow
YAML manifest OSIRIS native deck
(configs/osiris/*.yaml) (tests/test_osiris/decks/… or any path)
│ │
└──────────► run.py --cfg ◄──────────┘
│
── setup phase ────────┼──────────────────────────────────────────────
ergoExo.setup ──► BaseOsiris.__init__:
parse_deck_file()
merge_overrides() (mutates sections in place)
cfg["deck"] = deck_to_flat_dict(merged sections)
│
▼
log_params ──► MLflow ← logs the POST-override deck
│
── run phase ──────────┼──────────────────────────────────────────────
ergoExo(modules) ─► BaseOsiris.__call__:
render_deck(merged sections) ──► run_dir/os-stdin
runner.run_osiris(mpirun…) ──► run_dir/MS/{FLD,PHA,…}
│
▼
post.collect ──► td/ ──► MLflow
The deck logged to MLflow is the one after merge_overrides is applied — the
same merged sections that render_deck later writes to os-stdin — so the logged
params always match what OSIRIS actually ran.
How to run
First, point the runner at your built OSIRIS binary. The runner resolves it in
this order: osiris.binary in the manifest → OSIRIS_BIN_<dim>D env var (e.g.
OSIRIS_BIN_1D) → OSIRIS_BIN env var. The example manifests omit osiris.binary,
so set the env var once per shell:
export OSIRIS_BIN_1D=/path/to/osiris-1D.e # per-dim, preferred
# export OSIRIS_BIN=/path/to/osiris.e # or a single default for all dims
Then run from the repo root (run.py appends .yaml to --cfg, so omit the suffix):
uv run run.py --cfg configs/osiris/twostream-1d-short # smoke
uv run run.py --cfg configs/osiris/twostream-1d # full
uv run mlflow ui --backend-store-uri file://$(pwd)/mlruns # browse
Manifest schema
The full manifest schema — the osiris: section (deck path, binary resolution, MPI, staging,
overrides deck-patching semantics, density adaptive box sizing) and the output: section — is
documented in the Configuration Reference. Minimal example:
solver: osiris
mlflow:
experiment: osiris-pic1d-twostream
run: cold-equal-beams
osiris:
deck: tests/test_osiris/decks/two-stream-1d
mpi_ranks: 1
output:
diagnostics_to_log: null
What lands in MLflow
Kind |
Content |
|---|---|
Params |
every deck key, flattened — |
Params |
the manifest itself — |
Metrics |
|
Artifacts |
|
Artifacts |
|
Artifacts |
|
Artifacts |
|
Concurrent H5 → NetCDF conversion (osiris.stream_convert)
The binary/*.nc series are converted during the run by default
(osiris.stream_convert: true). The alternative — set it false — is the old
behavior: build them after the run, where post.collect reads each
diagnostic’s whole (t, …) history into memory and writes it in one shot. For
runs that dump a field every few steps (hundreds of thousands of tiny HDF5
files) that end-of-run pass is slow (a cold Lustre re-read of every file) and
memory-hungry (the whole stacked series at once), which is why streaming is the
default.
With streaming on, a best-effort background thread
(adept/osiris/stream.py::StreamConverter), spawned by run_osiris alongside
the OSIRIS subprocess, polls MS/ every stream_poll_s seconds and appends
each completed dump into <run_dir>/binary/<diag>.nc as it lands. Effects:
I/O overlaps compute — the conversion latency is hidden behind the (hours-long) PIC run, and each dump is read while still warm in the page cache instead of re-opened cold at the end.
Bounded memory — both the watcher and the at-job-end fallback build the NetCDF a dump at a time (
load_grid_h5→ one(t, …)slot), so peak conversion RAM is a single dump rather than the full series.
Mechanics and guarantees:
Each NetCDF uses an unlimited
tdimension grown one slot per dump, so it ends up sized to exactly the dumps produced — no pre-sized guess, no trailing fill to trim on early termination. Compression (zlib + shuffle) and chunking are preserved, so the files match the batch path’sbinary/*.nccontract and every downstream reader (load_series_nc, the plotting code,regen) is unchanged.A dump is only read once the next dump exists (OSIRIS has moved on), avoiding partial-file reads; the final sweep at job end picks up the last one.
Failure-isolated: any error in the converter is logged and never touches the OSIRIS run. The standard batch conversion is the safety net —
post.collectreuses whatever the watcher finished and stream-builds (still memory-bounded) any grid diagnostic it did not reach. RAW (particle) diagnostics, whose per-dump particle count varies, always use the batch concat path.Restart-safe: a
binary/<diag>.ncreopened after a checkpoint restart resumes after the slots already on disk.
This addresses the I/O wall-time and conversion-memory problems only; the
heavy pcolormesh/fft2 plotting cost at plot time is independent and
handled by the plotting changes, not here. See
osiris-lpi/postproc-performance.md for the full analysis.
Ramdisk staging (osiris.stage_root)
OSIRIS has no asynchronous diagnostic I/O: every dump is a synchronous
barrier in the time loop, so on a parallel filesystem (Lustre) the ranks stall
on each collective write, and at full dump cadence that stall can halve compute
utilization. stage_root works around it without trimming cadence: point it at
a fast node-local ephemeral filesystem (a /dev/shm ramdisk on NERSC GPU
nodes — those nodes have no local disk, so RAM-backed tmpfs is the only
node-local option) and OSIRIS’s synchronous writes hit RAM (microseconds)
instead of Lustre. The slow durable write is moved off the critical path onto
the background drainer — effectively the async I/O OSIRIS itself lacks.
osiris:
run_root: /pscratch/.../checkpoints # durable (Lustre)
stage_root: /dev/shm/osiris # fast scratch (ramdisk); requires stream_convert
stream_convert: true
Mechanics:
OSIRIS runs with its working directory on
stage_root/<run-name>; allMS/dumps land in RAM. The durable run directory underrun_rootis created as usual, andrun_dirin the result (hence everything post-processing reads) is always that durable directory — sopost.collectis unchanged.The
StreamConverterruns in drain mode: each completed dump is mirrored torun_root/<run-name>/MS/<diag>/and then deleted from the ramdisk, so the RAM high-water mark is bounded by the poll interval × production rate, not the whole run. Grid diagnostics are additionally streamed to the durablebinary/*.ncas before; RAW (particle) diagnostics, which can’t be slot-streamed, are mirrored as HDF5 so the batch path still finds them.binary/*.ncandstdout.log/stderr.logare written straight to the durable dir (by the driver, off OSIRIS’s critical path), so they survive even if the run is killed. At job end the drainer’s final sweep flushes the held-back last dump of each diagnostic; the runner then folds any stragglers (HIST/,RE/, …) from the scratch into the durable dir and reclaims the ramdisk.
Caveats:
RAM budget. The
tmpfscounts against node memory alongside the simulation. Bounded by the drain keeping up with production (Lustre bandwidth ≫ dump rate, so it normally does); for very long full-cadence runs, keepstream_poll_ssmall and leave headroom.Crash window. Whatever is still on the ramdisk when a node dies is lost (
tmpfsis volatile) — logs and already-mirrored dumps are safe on Lustre; only the few in-flight dumps are at risk. If you enable OSIRIS checkpoint/restart, keep those files off the ramdisk.Single-node only —
/dev/shmis per-node. Fine for the 1-GPU / single-node SRS runs; multi-node would need per-node drains and non-shared-file dumps.
Canned plots (plots/ artifacts)
These canned plots focus on 1D simulations.
post.collect renders a standard plot set via adept/osiris/plots.py::save_canned_plots. All labels are emitted as proper LaTeX ($\omega$, $c/\omega_p$, …).
Path |
What it shows |
|---|---|
|
|
|
value-vs- |
|
2-D FFT |
|
combined |
|
per-species density-moment spacetime + lineouts |
|
density profile vs |
|
temperature profile vs |
|
|
|
left/right-going transverse |
|
field / kinetic energy traces |
Note on
field_decomp/. The left/right split is exact only in vacuum or a uniform non-dispersive medium (|E| = |B|for a pure travelling wave). In a plasma the EM wave is dispersive, so the split is approximate — useful for direction, but cross-check the dispersion before reading the residual as physical counter-propagating power. The longitudinale1is electrostatic and is intentionally excluded.
SRS-specific plots & metrics live in osiris-lpi. The laser-energy budget (reflected / transmitted / absorbed over time: the
energy_budget.png+laser_energy_budget.txtartifacts and thelaser_reflectivity/laser_transmissivity/laser_absorbed_fracmetrics) and the distribution-function lineouts (distribution_lineouts/:f(p)averaged over the four domain quarters and the whole box, for the last dump and the last-1/8 average, in linear / log /δfviews) are not produced by adept’s general OSIRIS wrapper. They live in the osiris-lpi repo —osiris_lpi.OsirisLPIsubclassesBaseOsirisand adds them inpost_process, reading the driveantenna.a0/antenna.omega0from the deck. Regenerate them offline from saved NetCDFs withpython -m osiris_lpi.regen.
Memory: phase-space plots read slices, not the whole history. A saved
(t, …)series can be tens of GB uncompressed —x1gamma_q1is ~48 GB in memory (11705 × 1000 × 1024), ~0.19 GB gzipped on disk (sparse, ~250×) — soio.load_series(path)decompresses all of it (at 16 sims/node this OOM’d production).save_canned_plotstherefore loads only what each plot needs:io.load_series(path, t_indices=…)for the final dump + the evolution panels (_evolution_t_indices), and it decimates field heatmaps to render resolution beforepcolormesh. A new plotter over a large series must do the same —load_series(..., t_indices=…)for a few slices, orwith io.open_series(path) as da:(lazy;da.isel(t=it)reads one dump on demand) when it walks the whole time axis reducing one dump at a time. A strayda.values/da.sum()overt, or a deepda.copy(), pulls the full series into RAM (_decoratecopies shallow for this reason).
Programmatic use
from adept import ergoExo
import yaml
cfg = yaml.safe_load(open("configs/osiris/twostream-1d-short.yaml"))
cfg["osiris"]["overrides"] = {"time": {"tmax": 2.0}}
exo = ergoExo()
modules = exo.setup(cfg) # parse + log params
solver_result, post, run_id = exo(modules) # run + log metrics/artifacts
print(run_id, post["metrics"])
Adding new metrics
Edit adept/osiris/post.py:collect. The h5 dump for each diagnostic is at run_dir/MS/<TYPE>/<diag_name>/<diag_name>-NNNNNN.h5. Helpers _latest_h5, _field_energy_from_dump, and _walk_diag_dirs are already factored out. Anything you add to the returned metrics dict shows up in MLflow automatically.
Adding a new test problem
Native-deck-as-truth: just write the deck, point a manifest at it, run. No code changes — see Adding a new test problem in the Configuration Reference.