From 5073964ed2ecfaf0e3e7783c1dc04cd0b71df775 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 22 May 2026 12:50:27 -0700 Subject: [PATCH 01/49] Initial osiris wrapper commit --- adept/osiris/__init__.py | 12 + adept/osiris/base.py | 137 ++++++++ adept/osiris/deck.py | 340 ++++++++++++++++++ adept/osiris/io.py | 192 ++++++++++ adept/osiris/plots.py | 387 +++++++++++++++++++++ adept/osiris/post.py | 167 +++++++++ adept/osiris/runner.py | 194 +++++++++++ configs/osiris/twostream-1d-short.yaml | 16 + configs/osiris/twostream-1d-uploadall.yaml | 18 + configs/osiris/twostream-1d.yaml | 21 ++ osiris-adept-usage.md | 140 ++++++++ tests/test_osiris/test_deck_roundtrip.py | 137 ++++++++ tests/test_osiris/test_io_and_plots.py | 93 +++++ tests/test_osiris/test_runner.py | 58 +++ 14 files changed, 1912 insertions(+) create mode 100644 adept/osiris/__init__.py create mode 100644 adept/osiris/base.py create mode 100644 adept/osiris/deck.py create mode 100644 adept/osiris/io.py create mode 100644 adept/osiris/plots.py create mode 100644 adept/osiris/post.py create mode 100644 adept/osiris/runner.py create mode 100644 configs/osiris/twostream-1d-short.yaml create mode 100644 configs/osiris/twostream-1d-uploadall.yaml create mode 100644 configs/osiris/twostream-1d.yaml create mode 100644 osiris-adept-usage.md create mode 100644 tests/test_osiris/test_deck_roundtrip.py create mode 100644 tests/test_osiris/test_io_and_plots.py create mode 100644 tests/test_osiris/test_runner.py diff --git a/adept/osiris/__init__.py b/adept/osiris/__init__.py new file mode 100644 index 00000000..b17268f8 --- /dev/null +++ b/adept/osiris/__init__.py @@ -0,0 +1,12 @@ +"""adept OSIRIS solver wrapper.""" + +from __future__ import annotations + +__all__ = ["BaseOsiris"] + + +def __getattr__(name): + if name == "BaseOsiris": + from adept.osiris.base import BaseOsiris + return BaseOsiris + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/adept/osiris/base.py b/adept/osiris/base.py new file mode 100644 index 00000000..95302c7e --- /dev/null +++ b/adept/osiris/base.py @@ -0,0 +1,137 @@ +"""``BaseOsiris`` — the ADEPT module that drives OSIRIS.""" + +from __future__ import annotations + +from typing import Any + +from adept._base_ import ADEPTModule +from adept.osiris import deck as _deck +from adept.osiris import post as _post +from adept.osiris import runner as _runner + + +class BaseOsiris(ADEPTModule): + """Wraps an external OSIRIS binary as an adept solver. + + The OSIRIS native deck is the canonical simulation spec. The YAML + manifest layered on top supplies MLflow metadata, the binary path, + MPI rank count, and optional in-place deck overrides. + """ + + def __init__(self, cfg: dict) -> None: + super().__init__(cfg) + osiris_cfg = cfg.get("osiris", {}) + deck_path = osiris_cfg.get("deck") + if not deck_path: + raise ValueError( + "BaseOsiris: cfg['osiris']['deck'] is required" + ) + + self._sections = _deck.parse_deck_file(deck_path) + overrides = osiris_cfg.pop("overrides", None) or {} + if overrides: + _deck.merge_overrides(self._sections, overrides) + + # Surface the parsed (post-override) deck inside cfg so adept's + # log_params picks every parameter up as a flat MLflow param. + # The raw overrides dict is intentionally popped above to avoid + # confusing log_params/flatdict with integer keys; the applied + # values now live verbatim under cfg["deck"]. + cfg["deck"] = _deck.deck_to_flat_dict(self._sections) + + def write_units(self) -> dict: + return {} + + def get_derived_quantities(self) -> dict: + """Lift a few useful scalars out of the deck for MLflow visibility.""" + grid = dict(self._iter_first_section("grid")) + time = dict(self._iter_first_section("time")) + time_step = dict(self._iter_first_section("time_step")) + space = dict(self._iter_first_section("space")) + + derived: dict[str, Any] = {} + nx = self._first_array_value(grid, "nx_p") + xmin = self._first_array_value(space, "xmin") + xmax = self._first_array_value(space, "xmax") + dt = time_step.get("dt") + tmax = time.get("tmax") + + if nx and xmin is not None and xmax is not None: + derived["dx"] = [ + (xmax[d] - xmin[d]) / nx[d] for d in range(len(nx)) + ] + if dt is not None and nx: + derived["cfl_ratio"] = float(dt) / min(derived["dx"]) + if dt is not None and tmax is not None: + try: + derived["num_steps"] = int(float(tmax) / float(dt)) + except (TypeError, ValueError): + pass + if derived: + self.cfg.setdefault("derived", {}).update(derived) + return derived + + def get_solver_quantities(self) -> None: + return None + + def init_state_and_args(self) -> dict: + return {} + + def init_diffeqsolve(self) -> None: + return None + + def init_modules(self) -> dict: + return {} + + def __call__(self, trainable_modules: dict, args: dict) -> dict: + osiris_cfg = self.cfg.get("osiris", {}) + binary = _runner.discover_binary( + osiris_cfg.get("binary"), + dim=self._infer_dim(), + ) + mpi_ranks = int(osiris_cfg.get("mpi_ranks", 1)) + run_root = osiris_cfg.get("run_root", "./osiris_runs") + deck_text = _deck.render_deck(self._sections) + result = _runner.run_osiris( + deck_text, + binary=binary, + mpi_ranks=mpi_ranks, + run_root=run_root, + extra_mpi_args=osiris_cfg.get("extra_mpi_args"), + ) + return {"solver result": result} + + def post_process(self, run_output: dict, td: str) -> dict: + return _post.collect(run_output, self.cfg, td) + + def vg(self, trainable_modules: dict, args: dict): + raise NotImplementedError( + "OSIRIS is not differentiable inside adept" + ) + + # --- helpers ---------------------------------------------------------- + + def _iter_first_section(self, name: str) -> dict: + for sec_name, params in self._sections: + if sec_name == name: + return params + return {} + + @staticmethod + def _first_array_value(params: dict, base_key: str) -> list | None: + """Look up either ``base_key`` or ``base_key(...)`` and return as list.""" + for k, v in params.items(): + if k == base_key or k.startswith(base_key + "("): + if isinstance(v, list): + return v + return [v] + return None + + def _infer_dim(self) -> int | None: + """Best-effort: read dimensionality from ``grid.nx_p(1:D)``.""" + grid = self._iter_first_section("grid") + for k in grid: + if k.startswith("nx_p"): + v = grid[k] + return len(v) if isinstance(v, list) else 1 + return None diff --git a/adept/osiris/deck.py b/adept/osiris/deck.py new file mode 100644 index 00000000..00c10f35 --- /dev/null +++ b/adept/osiris/deck.py @@ -0,0 +1,340 @@ +"""OSIRIS namelist parser / renderer. + +OSIRIS uses a Fortran-namelist-like syntax that is NOT a standard Fortran +namelist: + + ! comment + section_name + { + key = scalar, + key(1:N) = a, b, c, ! comments allowed inline + key = "string", + key = .true., + key(1:2,1) = 0.0, 1.0, + } + +Multiple sections with the same name are common (e.g. ``species``) and the +order matters because it corresponds to species 1, 2, ... in OSIRIS. + +Canonical representation: an ordered list of ``(section_name, params)`` +pairs where ``params`` is a ``dict`` mapping a key-string (with any slice +spec preserved verbatim, e.g. ``"nx_p(1:1)"``) to a Python value (int, +float, str, bool, or list of those). + +The invariant we rely on is dict-level round-trip: + + parse(render(parse(text))) == parse(text) +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any + +Section = tuple[str, dict[str, Any]] +Sections = list[Section] + + +_TRUE_TOKENS = {".true.", ".t."} +_FALSE_TOKENS = {".false.", ".f."} + +_IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +_KEY_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*(?:\([^)]*\))?") +_INT_RE = re.compile(r"^[+-]?\d+$") +_FLOAT_RE = re.compile(r"^[+-]?(?:\d+\.\d*|\.\d+|\d+)(?:[eEdD][+-]?\d+)?$") + + +def _strip_comment(line: str) -> str: + """Drop everything from the first unquoted ``!`` to the end of line.""" + in_string = False + out: list[str] = [] + for ch in line: + if ch == '"': + in_string = not in_string + out.append(ch) + elif ch == "!" and not in_string: + break + else: + out.append(ch) + return "".join(out) + + +def _split_top_commas(s: str) -> list[str]: + """Split on commas not inside double-quotes.""" + pieces: list[str] = [] + buf: list[str] = [] + in_string = False + for ch in s: + if ch == '"': + in_string = not in_string + buf.append(ch) + elif ch == "," and not in_string: + pieces.append("".join(buf).strip()) + buf = [] + else: + buf.append(ch) + last = "".join(buf).strip() + if last: + pieces.append(last) + return pieces + + +def _parse_atom(tok: str) -> Any: + """Convert a single token to a Python value.""" + t = tok.strip() + if not t: + return None + if t.startswith('"') and t.endswith('"') and len(t) >= 2: + return t[1:-1] + low = t.lower() + if low in _TRUE_TOKENS: + return True + if low in _FALSE_TOKENS: + return False + if _INT_RE.match(t): + return int(t) + if _FLOAT_RE.match(t): + return float(t.replace("d", "e").replace("D", "e")) + return t # unrecognized — keep as-is + + +def _parse_value(rhs: str) -> Any: + pieces = _split_top_commas(rhs) + if not pieces: + return None + if len(pieces) == 1: + return _parse_atom(pieces[0]) + return [_parse_atom(p) for p in pieces] + + +def parse_deck(text: str) -> Sections: + """Parse an OSIRIS native deck (as a string) into ordered sections.""" + src = "\n".join(_strip_comment(ln) for ln in text.splitlines()) + n = len(src) + i = 0 + + def skip_ws(j: int) -> int: + while j < n and src[j] in " \t\n\r": + j += 1 + return j + + sections: Sections = [] + while True: + i = skip_ws(i) + if i >= n: + break + m = _IDENT_RE.match(src, i) + if not m: + raise ValueError( + f"Expected section name at offset {i}: {src[i:i+30]!r}" + ) + name = m.group(0) + i = m.end() + i = skip_ws(i) + if i >= n or src[i] != "{": + raise ValueError( + f"Expected '{{' after section {name!r} at offset {i}: " + f"{src[i:i+30]!r}" + ) + i += 1 + params: dict[str, Any] = {} + while True: + i = skip_ws(i) + if i >= n: + raise ValueError(f"Unterminated section {name!r}") + if src[i] == "}": + i += 1 + break + km = _KEY_RE.match(src, i) + if not km: + raise ValueError( + f"Expected key in section {name!r} at offset {i}: " + f"{src[i:i+30]!r}" + ) + key = km.group(0) + i = km.end() + i = skip_ws(i) + if i >= n or src[i] != "=": + raise ValueError( + f"Expected '=' after key {key!r} at offset {i}" + ) + i += 1 + value_start = i + in_string = False + while i < n: + ch = src[i] + if ch == '"': + in_string = not in_string + i += 1 + elif not in_string and ch == "}": + break + elif not in_string and ch == ",": + # Peek past whitespace: if the next token is an identifier + # followed by '=', this comma terminates the current + # key=value pair; otherwise it's an intra-value separator + # (e.g. ``uth(1:3) = 0.1, 0.2, 0.3``). + j = i + 1 + while j < n and src[j] in " \t\n\r": + j += 1 + peek = _KEY_RE.match(src, j) + if peek: + jj = peek.end() + while jj < n and src[jj] in " \t\n\r": + jj += 1 + if jj < n and src[jj] == "=": + break + i += 1 + else: + i += 1 + value_text = src[value_start:i].strip() + if i < n and src[i] == ",": + i += 1 + params[key] = _parse_value(value_text) + sections.append((name, params)) + return sections + + +def parse_deck_file(path: str | Path) -> Sections: + return parse_deck(Path(path).read_text()) + + +def _render_value(v: Any) -> str: + if isinstance(v, bool): + return ".true." if v else ".false." + if isinstance(v, str): + return f'"{v}"' + if isinstance(v, int): + return str(v) + if isinstance(v, float): + return repr(v) + if isinstance(v, list): + return ", ".join(_render_value(x) for x in v) + if v is None: + return "" + raise TypeError(f"Cannot render value of type {type(v).__name__}: {v!r}") + + +def render_deck(sections: Sections) -> str: + """Render canonical sections back to OSIRIS native namelist text.""" + out: list[str] = [] + for name, params in sections: + out.append(name) + if not params: + out.append("{") + out.append("}") + out.append("") + continue + out.append("{") + kw = max(len(k) for k in params) + for k, v in params.items(): + out.append(f" {k:<{kw}} = {_render_value(v)},") + out.append("}") + out.append("") + return "\n".join(out) + + +def _find_param_key(params: dict[str, Any], requested: str) -> str: + """Resolve an override key against a section's existing keys. + + Accepts either an exact key (``"nx_p(1:1)"``) or the base name + (``"nx_p"``), the latter only when it's unambiguous within the section. + Returns the matching existing key, or the requested string verbatim if + the key is brand-new. + """ + if requested in params: + return requested + candidates = [k for k in params if k.split("(", 1)[0] == requested] + if len(candidates) == 1: + return candidates[0] + if len(candidates) == 0: + return requested + raise ValueError( + f"Ambiguous override key {requested!r}; candidates: {candidates}" + ) + + +def _merge_params(params: dict[str, Any], over: dict[str, Any]) -> None: + for k, v in over.items(): + target = _find_param_key(params, k) + params[target] = v + + +def merge_overrides(sections: Sections, overrides: dict[str, Any]) -> None: + """Apply ``overrides`` to ``sections`` in place. + + ``overrides`` shape:: + + { + "grid": {"nx_p": [256]}, # apply to all occurrences + "species": {0: {"num_par_x": [512]}}, # indexed for repeated sections + } + """ + if not overrides: + return + by_name: dict[str, list[int]] = {} + for idx, (name, _) in enumerate(sections): + by_name.setdefault(name, []).append(idx) + + for sec_name, sec_over in overrides.items(): + if sec_name not in by_name: + raise KeyError( + f"Override references unknown section: {sec_name!r}" + ) + occurrences = by_name[sec_name] + if isinstance(sec_over, dict) and sec_over and all( + isinstance(k, int) for k in sec_over.keys() + ): + for idx, params_over in sec_over.items(): + if idx < 0 or idx >= len(occurrences): + raise IndexError( + f"Override section {sec_name!r}[{idx}] out of range; " + f"{len(occurrences)} occurrence(s) present" + ) + _merge_params(sections[occurrences[idx]][1], params_over) + else: + for occ in occurrences: + _merge_params(sections[occ][1], sec_over) + + +_KEY_SANITIZE_RE = re.compile(r"[^A-Za-z0-9_./:\- ]+") + + +def _sanitize_key(k: str) -> str: + """Map an OSIRIS key (which may contain ``()``, ``[]``) to an MLflow-safe + parameter name. MLflow allows alphanumerics, ``_``, ``-``, ``.``, ``:``, + ``/``, and spaces. Everything else is folded to ``_`` and trailing + runs of underscores are collapsed.""" + out = _KEY_SANITIZE_RE.sub("_", k) + while "__" in out: + out = out.replace("__", "_") + return out.strip("_") or "_" + + +def deck_to_flat_dict(sections: Sections) -> dict[str, Any]: + """Flat representation for MLflow param logging. + + Repeated sections get bracketed indices in the source, then sanitized + to MLflow-safe names: ``species[0].num_par_x(1:1)`` → + ``species_0.num_par_x_1:1``. List values get numeric suffixes. + """ + name_count: dict[str, int] = {} + for name, _ in sections: + name_count[name] = name_count.get(name, 0) + 1 + name_idx: dict[str, int] = {} + flat: dict[str, Any] = {} + for name, params in sections: + if name_count[name] > 1: + i = name_idx.get(name, 0) + prefix = _sanitize_key(f"{name}[{i}]") + name_idx[name] = i + 1 + else: + prefix = _sanitize_key(name) + for k, v in params.items(): + ks = _sanitize_key(k) + if isinstance(v, list): + for j, x in enumerate(v): + flat[f"{prefix}.{ks}.{j}"] = x + else: + flat[f"{prefix}.{ks}"] = v + return flat diff --git a/adept/osiris/io.py b/adept/osiris/io.py new file mode 100644 index 00000000..7bbfe528 --- /dev/null +++ b/adept/osiris/io.py @@ -0,0 +1,192 @@ +"""Lightweight loaders for OSIRIS HDF5 output. + +OSIRIS dump files have a tiny, well-defined structure: + +- one dataset at the root, named after the diagnostic (``e1``, ``charge``, + ``x1p1``, ...); +- an ``AXIS`` group with ``AXIS1``, ``AXIS2``, ... datasets each holding + ``[min, max]`` plus ``LONG_NAME``, ``NAME``, ``UNITS`` attrs; +- a ``SIMULATION`` group with run-wide metadata as attrs (``DT``, ``NX``, + ``XMIN``, ``XMAX``, ``NDIMS``); +- root attrs ``TIME``, ``ITER``, ``NAME``, ``LABEL``, ``UNITS``. + +OSIRIS writes datasets in Fortran-axis order, so the *first* HDF5 axis +becomes the *last* numpy axis. This loader reverses the AXIS list when +constructing coordinates so the returned ``xarray.DataArray`` has dims +in numpy order (rows first), with each dim's ``name`` matching the +OSIRIS axis label. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr + + +_ITER_RE = re.compile(r"-(\d+)\.h5$") + + +def _decode(v) -> str: + """OSIRIS stores attrs as ``S256`` bytes inside length-1 arrays.""" + if isinstance(v, np.ndarray): + v = v.flat[0] + if isinstance(v, bytes): + return v.decode("utf-8", errors="replace").rstrip() + return str(v) + + +def _axis_metadata(h5f: h5py.File) -> list[dict]: + """Return per-axis dicts in *OSIRIS* order (axis 1 first). + + Each dict has ``name``, ``long_name``, ``units``, ``min``, ``max``. + """ + axes: list[dict] = [] + if "AXIS" not in h5f: + return axes + grp = h5f["AXIS"] + for ax_name in sorted(grp.keys()): # AXIS1, AXIS2, ... + ax = grp[ax_name] + vals = ax[:] + axes.append( + { + "name": _decode(ax.attrs.get("NAME", ax_name)), + "long_name": _decode(ax.attrs.get("LONG_NAME", ax_name)), + "units": _decode(ax.attrs.get("UNITS", "")), + "min": float(vals[0]), + "max": float(vals[-1]), + } + ) + return axes + + +def _iter_from_name(path: Path) -> int: + m = _ITER_RE.search(path.name) + return int(m.group(1)) if m else -1 + + +def load_grid_h5(path: str | Path) -> xr.DataArray: + """Load one OSIRIS field/charge/current dump into an xarray DataArray. + + The returned array has dims in numpy order (e.g. ``("x2", "x1")`` for + a 2D dump, ``("x1",)`` for 1D). Use ``.transpose(...)`` if you prefer + the OSIRIS convention. + """ + path = Path(path) + with h5py.File(path, "r") as f: + # Identify the data dataset (everything that's not AXIS / SIMULATION). + data_keys = [k for k in f.keys() if k not in ("AXIS", "SIMULATION")] + if len(data_keys) != 1: + raise ValueError( + f"Expected exactly one data dataset in {path}; got {data_keys}" + ) + name = data_keys[0] + arr = f[name][...].astype("float64") + axes_osiris = _axis_metadata(f) + # Reverse to numpy order. + axes_numpy = list(reversed(axes_osiris)) + + coords = {} + dims = [] + for i, ax in enumerate(axes_numpy): + n_i = arr.shape[i] + coords[ax["name"]] = np.linspace(ax["min"], ax["max"], n_i) + dims.append(ax["name"]) + + attrs = { + "time": float(f.attrs["TIME"][0]) if "TIME" in f.attrs else float("nan"), + "iter": int(f.attrs["ITER"][0]) if "ITER" in f.attrs else _iter_from_name(path), + "long_name": _decode(f.attrs.get("LABEL", name)), + "units": _decode(f.attrs.get("UNITS", "")), + "time_units": _decode(f.attrs.get("TIME UNITS", "")), + "source": str(path), + "axis_units": {ax["name"]: ax["units"] for ax in axes_numpy}, + "axis_long_names": {ax["name"]: ax["long_name"] for ax in axes_numpy}, + } + if "SIMULATION" in f: + sim = f["SIMULATION"].attrs + for key in ("DT", "NDIMS", "XMIN", "XMAX", "NX", "PERIODIC"): + if key in sim: + val = sim[key] + if hasattr(val, "tolist"): + val = val.tolist() + attrs[f"sim.{key}"] = val + + return xr.DataArray(arr, coords=coords, dims=dims, name=name, attrs=attrs) + + +def load_phasespace_h5(path: str | Path) -> xr.DataArray: + """Phase-space dumps have the exact same on-disk structure as grid + dumps, so this is just a semantic alias.""" + return load_grid_h5(path) + + +def _sort_dumps(directory: Path) -> list[Path]: + files = [p for p in directory.iterdir() if p.is_file() and p.suffix == ".h5"] + return sorted(files, key=_iter_from_name) + + +def load_series(directory: str | Path) -> xr.DataArray: + """Stack every dump in ``directory`` into a ``(t, ...)`` DataArray. + + All files must share the same diagnostic name, the same spatial / + phase-space shape, and the same axis bounds — this is the standard + OSIRIS convention for a single diagnostic's time history. + """ + directory = Path(directory) + dumps = _sort_dumps(directory) + if not dumps: + raise FileNotFoundError(f"No .h5 dumps in {directory}") + first = load_grid_h5(dumps[0]) + + times = np.empty(len(dumps), dtype="float64") + iters = np.empty(len(dumps), dtype="int64") + data = np.empty((len(dumps), *first.shape), dtype=first.dtype) + data[0] = first.values + times[0] = first.attrs["time"] + iters[0] = first.attrs["iter"] + for i, p in enumerate(dumps[1:], start=1): + da = load_grid_h5(p) + if da.shape != first.shape: + raise ValueError( + f"Shape mismatch in series: {p} has {da.shape}, " + f"expected {first.shape}" + ) + data[i] = da.values + times[i] = da.attrs["time"] + iters[i] = da.attrs["iter"] + + coords = {"t": times, "iter": ("t", iters)} + coords.update({d: first.coords[d] for d in first.dims}) + dims = ("t", *first.dims) + attrs = dict(first.attrs) + attrs.pop("time", None) + attrs.pop("iter", None) + attrs.pop("source", None) + attrs["source_dir"] = str(directory) + return xr.DataArray( + data, coords=coords, dims=dims, name=first.name, attrs=attrs + ) + + +def list_diagnostics(run_dir: str | Path) -> dict[str, Path]: + """Map every diagnostic name to its directory under ``run_dir/MS/``. + + Discovery rule: any directory that directly contains ``.h5`` dumps is + a diagnostic. The returned key is the directory's relative path under + ``MS/`` (e.g. ``"FLD/e1"``, ``"PHA/x1p1/beam_pos"``). + """ + run_dir = Path(run_dir) + ms = run_dir / "MS" + if not ms.is_dir(): + raise FileNotFoundError(f"No MS/ directory under {run_dir}") + out: dict[str, Path] = {} + for d in ms.rglob("*"): + if not d.is_dir(): + continue + if any(c.suffix == ".h5" for c in d.iterdir()): + out[str(d.relative_to(ms))] = d + return out diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py new file mode 100644 index 00000000..ed58aa0a --- /dev/null +++ b/adept/osiris/plots.py @@ -0,0 +1,387 @@ +"""Canned matplotlib views over OSIRIS diagnostics. + +Each plotter accepts either an already-loaded ``xarray.DataArray`` or a +filesystem path / directory, and returns the ``Axes`` it drew on so the +caller can tweak labels / colorbars / save the figure. + +``save_canned_plots(run_dir, out_dir)`` ties them together and writes +PNGs for the standard set: + + out_dir/ + spacetime/.png (t, x) heatmaps of every FLD diagnostic + phasespace//.png final-step (x, p) heatmap per species + energy_vs_time.png total field energy time-trace + omega_k/.png 2-D FFT (k, ω) dispersion plot +""" + +from __future__ import annotations + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from adept.osiris import io as _io + + +# --- low-level plotters ---------------------------------------------------- + + +def plot_spacetime( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = False, + cmap: str = "RdBu_r", + title: str | None = None, +) -> plt.Axes: + """``(t, x)`` heatmap of a 1D-field time-series. + + Accepts a pre-loaded ``DataArray`` (dims must include ``t`` and one + spatial axis) or a directory of dumps that ``load_series`` will eat. + """ + da = _ensure_series(series) + if da.ndim != 2: + raise ValueError( + f"plot_spacetime expects a 2D (t, x) array; got dims {da.dims}" + ) + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + data = np.log10(np.abs(da.values) + 1e-30) if log else da.values + t = da.coords["t"].values + xname = next(d for d in da.dims if d != "t") + x = da.coords[xname].values + mesh = ax.pcolormesh(t, x, data.T, shading="auto", cmap=cmap) + plt.colorbar( + mesh, + ax=ax, + label=f"log10 |{da.name}|" if log else da.name, + ) + ax.set_xlabel(_axis_label(da, "t")) + ax.set_ylabel(_axis_label(da, xname)) + ax.set_title(title or f"{da.name} spacetime") + return ax + + +def plot_phasespace( + da: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = True, + cmap: str = "viridis", + title: str | None = None, +) -> plt.Axes: + """Heatmap of a 2D phase-space density (e.g. ``x1p1``).""" + if not isinstance(da, xr.DataArray): + da = _io.load_phasespace_h5(da) + if da.ndim != 2: + raise ValueError( + f"plot_phasespace expects 2D data; got {da.dims} {da.shape}" + ) + if ax is None: + _, ax = plt.subplots(figsize=(5, 4)) + arr = da.values + if log: + arr = np.log10(np.abs(arr) + 1e-30) + d0, d1 = da.dims + x0, x1 = da.coords[d0].values, da.coords[d1].values + mesh = ax.pcolormesh(x0, x1, arr.T, shading="auto", cmap=cmap) + plt.colorbar( + mesh, ax=ax, label=f"log10 {da.name}" if log else da.name + ) + ax.set_xlabel(_axis_label(da, d0)) + ax.set_ylabel(_axis_label(da, d1)) + t = da.attrs.get("time", float("nan")) + ax.set_title(title or f"{da.name} t = {t:.3g}") + return ax + + +def field_energy_series(run_dir: str | Path) -> xr.DataArray: + """Sum ``(|E|^2 + |B|^2) / 2`` over space at every saved step. + + Walks every component dir under ``MS/FLD/`` (``e1``, ``e2``, ..., ``b1``, + ``b2``, ``b3``), aligns by iteration count, and returns a 1D + ``DataArray`` of total field energy in code units vs time. Components + that aren't dumped (or that have a different save cadence) are + skipped. + """ + run_dir = Path(run_dir) + fld = run_dir / "MS" / "FLD" + if not fld.is_dir(): + raise FileNotFoundError(f"No MS/FLD under {run_dir}") + components = ("e1", "e2", "e3", "b1", "b2", "b3") + by_iter: dict[int, dict[str, float]] = {} + times: dict[int, float] = {} + for comp in components: + d = fld / comp + if not d.is_dir(): + continue + for h5 in sorted(d.iterdir()): + if h5.suffix != ".h5": + continue + da = _io.load_grid_h5(h5) + # cell volume from coord spacing in each dim + dvol = 1.0 + for dim in da.dims: + c = da.coords[dim].values + if c.size > 1: + dvol *= float(c[1] - c[0]) + e = 0.5 * float((da.values ** 2).sum()) * dvol + it = int(da.attrs["iter"]) + by_iter.setdefault(it, {})[comp] = e + times.setdefault(it, float(da.attrs["time"])) + if not by_iter: + raise RuntimeError(f"No field dumps under {fld}") + iters = sorted(by_iter) + totals = np.array([sum(by_iter[i].values()) for i in iters]) + t = np.array([times[i] for i in iters]) + return xr.DataArray( + totals, + coords={"t": t, "iter": ("t", np.asarray(iters))}, + dims=("t",), + name="field_energy", + attrs={"long_name": "total field energy", "units": "code"}, + ) + + +def plot_energy_vs_time( + src: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = True, + title: str | None = None, +) -> plt.Axes: + """Plot total field energy vs time. + + ``src`` may be a precomputed ``DataArray`` from ``field_energy_series`` + or a ``run_dir`` path (in which case the series is computed on the fly). + """ + if isinstance(src, xr.DataArray): + da = src + else: + da = field_energy_series(src) + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + ax.plot(da.coords["t"].values, da.values) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel(r"$\int (|E|^2 + |B|^2)/2\ d^Dx$") + ax.set_title(title or "field energy") + ax.grid(True, which="both", alpha=0.3) + return ax + + +def plot_omega_k( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = True, + cmap: str = "magma", + show_em: bool = True, + show_langmuir: bool = False, + v_th: float | None = None, + omega_p: float = 1.0, + k_max: float | None = None, + omega_max: float | None = None, + title: str | None = None, +) -> plt.Axes: + """2-D FFT of a ``(t, x)`` field series; show power in ``(k, ω)`` space. + + The OSIRIS convention has ω_p = 1, c = 1 in code units, so the + relativistic-EM dispersion ω² = ω_p² + k² and the Langmuir + Bohm–Gross dispersion ω² = ω_p² + 3 k² v_th² overlay directly when + you ask for them. + """ + da = _ensure_series(series) + if da.ndim != 2: + raise ValueError( + f"plot_omega_k expects (t, x); got dims {da.dims}" + ) + if ax is None: + _, ax = plt.subplots(figsize=(6, 5)) + + t = da.coords["t"].values + xname = next(d for d in da.dims if d != "t") + x = da.coords[xname].values + dt = float(t[1] - t[0]) + dx = float(x[1] - x[0]) + nt = t.size + nx = x.size + + # 2-D FFT then shift so DC sits in the middle. + F = np.fft.fftshift(np.fft.fft2(da.values)) + P = np.abs(F) ** 2 + if log: + P = np.log10(P + 1e-30) + + omega = np.fft.fftshift(np.fft.fftfreq(nt, d=dt)) * 2 * np.pi + k = np.fft.fftshift(np.fft.fftfreq(nx, d=dx)) * 2 * np.pi + + mesh = ax.pcolormesh(k, omega, P, shading="auto", cmap=cmap) + plt.colorbar(mesh, ax=ax, label=("log10 |F|^2" if log else "|F|^2")) + + if k_max is None: + k_max = float(np.max(np.abs(k))) + if omega_max is None: + omega_max = float(np.max(np.abs(omega))) + ax.set_xlim(-k_max, k_max) + ax.set_ylim(-omega_max, omega_max) + + # Overlay analytical dispersion lines. + k_line = np.linspace(-k_max, k_max, 401) + if show_em: + w_em = np.sqrt(omega_p ** 2 + k_line ** 2) + ax.plot(k_line, +w_em, "w--", lw=1, alpha=0.7, + label=r"EM: $\omega^2 = \omega_p^2 + k^2$") + ax.plot(k_line, -w_em, "w--", lw=1, alpha=0.7) + if show_langmuir: + if v_th is None: + raise ValueError("show_langmuir=True requires v_th=...") + w_l = np.sqrt(omega_p ** 2 + 3 * (k_line * v_th) ** 2) + ax.plot(k_line, +w_l, "c:", lw=1, alpha=0.8, + label=r"Langmuir: $\omega^2 = \omega_p^2 + 3 k^2 v_{th}^2$") + ax.plot(k_line, -w_l, "c:", lw=1, alpha=0.8) + if show_em or show_langmuir: + ax.legend(loc="upper right", fontsize=8, framealpha=0.6) + + ax.axhline(0, color="w", lw=0.4, alpha=0.4) + ax.axvline(0, color="w", lw=0.4, alpha=0.4) + ax.set_xlabel(r"$k\ [\omega_p / c]$") + ax.set_ylabel(r"$\omega\ [\omega_p]$") + ax.set_title(title or f"{da.name} omega-k spectrum") + return ax + + +# --- driver --------------------------------------------------------------- + + +def save_canned_plots( + run_dir: str | Path, + out_dir: str | Path, + *, + v_th: float | None = None, + dpi: int = 120, +) -> dict[str, Path]: + """Generate the standard set of PNGs for a finished OSIRIS run. + + Returns a mapping of plot-name → output PNG path. + """ + run_dir = Path(run_dir) + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + written: dict[str, Path] = {} + + diags = _io.list_diagnostics(run_dir) + + # Spacetime + omega-k for each field component (those live under FLD/). + fld_st = out_dir / "spacetime" + fld_st.mkdir(exist_ok=True) + omk_dir = out_dir / "omega_k" + omk_dir.mkdir(exist_ok=True) + for diag_rel, diag_path in diags.items(): + if not diag_rel.startswith("FLD/"): + continue + comp = diag_rel.split("/", 1)[1] + try: + ser = _io.load_series(diag_path) + except Exception as e: + print(f"[plots] skipping series {diag_rel}: {e}") + continue + if ser.ndim != 2: + continue # higher-dim FLD plots skipped for now + fig, ax = plt.subplots(figsize=(6, 4)) + plot_spacetime(ser, ax=ax) + p = fld_st / f"{comp}.png" + fig.savefig(p, bbox_inches="tight", dpi=dpi) + plt.close(fig) + written[f"spacetime/{comp}"] = p + + fig, ax = plt.subplots(figsize=(6, 5)) + plot_omega_k( + ser, + ax=ax, + show_langmuir=v_th is not None, + v_th=v_th, + ) + p = omk_dir / f"{comp}.png" + fig.savefig(p, bbox_inches="tight", dpi=dpi) + plt.close(fig) + written[f"omega_k/{comp}"] = p + + # Final-step phase space per species. + ph_dir = out_dir / "phasespace" + for diag_rel, diag_path in diags.items(): + if not diag_rel.startswith("PHA/"): + continue + # PHA/x1p1/beam_pos -> ps_name="x1p1", species="beam_pos" + parts = diag_rel.split("/") + if len(parts) < 3: + continue + ps_name, species = parts[1], parts[2] + last = _latest_h5(diag_path) + if last is None: + continue + sp_out = ph_dir / species + sp_out.mkdir(parents=True, exist_ok=True) + da = _io.load_phasespace_h5(last) + fig, ax = plt.subplots(figsize=(5, 4)) + plot_phasespace(da, ax=ax) + p = sp_out / f"{ps_name}.png" + fig.savefig(p, bbox_inches="tight", dpi=dpi) + plt.close(fig) + written[f"phasespace/{species}/{ps_name}"] = p + + # Energy vs time (uses any/all FLD components present). + try: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_energy_vs_time(run_dir, ax=ax) + p = out_dir / "energy_vs_time.png" + fig.savefig(p, bbox_inches="tight", dpi=dpi) + plt.close(fig) + written["energy_vs_time"] = p + except (FileNotFoundError, RuntimeError) as e: + print(f"[plots] skipping energy_vs_time: {e}") + + return written + + +# --- helpers -------------------------------------------------------------- + + +def _ensure_series(src) -> xr.DataArray: + if isinstance(src, xr.DataArray): + return src + p = Path(src) + if p.is_dir(): + return _io.load_series(p) + raise TypeError( + f"Expected an xr.DataArray or a directory path; got {type(src).__name__}" + ) + + +def _axis_label(da: xr.DataArray, dim: str) -> str: + if dim == "t": + u = da.attrs.get("time_units", r"$1/\omega_p$") + return f"t [{u}]" if u else "t" + units = da.attrs.get("axis_units", {}).get(dim, "") + long = da.attrs.get("axis_long_names", {}).get(dim, dim) + if units: + # OSIRIS axis units are TeX-ish (e.g. "c / \\omega_p"); wrap in $...$. + return rf"{long} [${units}$]" + return long + + +def _latest_h5(diag_dir: Path) -> Path | None: + best: tuple[int, Path] | None = None + for p in diag_dir.iterdir(): + if p.suffix != ".h5": + continue + m = _io._ITER_RE.search(p.name) + if not m: + continue + it = int(m.group(1)) + if best is None or it > best[0]: + best = (it, p) + return best[1] if best else None diff --git a/adept/osiris/post.py b/adept/osiris/post.py new file mode 100644 index 00000000..744ee41e --- /dev/null +++ b/adept/osiris/post.py @@ -0,0 +1,167 @@ +"""Post-processing for OSIRIS runs. + +After ``BaseOsiris.__call__`` finishes, ``collect`` walks the ``MS/`` +output tree, copies the final-step HDF5 file from each diagnostic to the +adept temp dir (so MLflow uploads it), optionally tarballs the whole +tree, and returns scalar metrics. +""" + +from __future__ import annotations + +import re +import shutil +import tarfile +from pathlib import Path +from typing import Any + + +# OSIRIS dump filenames look like ``e1-000600.h5`` or +# ``x1p1-beam_pos-000060.h5``. The iteration number is always a +# zero-padded integer right before ``.h5``. +_ITER_RE = re.compile(r"-(\d+)\.h5$") + + +def _iter_h5_files(d: Path) -> list[Path]: + if not d.is_dir(): + return [] + return [p for p in d.iterdir() if p.is_file() and p.suffix == ".h5"] + + +def _latest_h5(d: Path) -> Path | None: + files = _iter_h5_files(d) + if not files: + return None + + def keyfn(p: Path) -> int: + m = _ITER_RE.search(p.name) + return int(m.group(1)) if m else -1 + + return max(files, key=keyfn) + + +def _walk_diag_dirs(ms: Path) -> list[Path]: + """Return every directory under ``ms`` that contains ``.h5`` files.""" + out: list[Path] = [] + for p in ms.rglob("*"): + if p.is_dir() and any(c.suffix == ".h5" for c in p.iterdir()): + out.append(p) + return out + + +def _field_energy_from_dump(h5_path: Path) -> float: + """Integrate ``field^2 * cell_volume`` for a single dump file. + + OSIRIS stores one component per file (e1, e2, ..., b3) so this returns + the contribution of *that one component*. + """ + import h5py # local import keeps adept import light + + with h5py.File(h5_path, "r") as f: + # The dataset name matches the diagnostic name (e.g. "e1"). + name = h5_path.stem.rsplit("-", 1)[0] + if name not in f: + return float("nan") + arr = f[name][...] + # Reconstruct cell volume from SIMULATION attrs. + sim = f["SIMULATION"] + ndims = int(sim.attrs["NDIMS"][0]) + xmin = sim.attrs["XMIN"] + xmax = sim.attrs["XMAX"] + nx = sim.attrs["NX"] + dvol = 1.0 + for d in range(ndims): + dvol *= (float(xmax[d]) - float(xmin[d])) / int(nx[d]) + return 0.5 * float((arr.astype("float64") ** 2).sum()) * dvol + + +def _diagnostic_name(diag_dir: Path) -> str: + """Short tag for a diagnostic directory, used in metric / artifact names.""" + return diag_dir.name + + +def _total_field_energy(ms: Path) -> float: + """Sum |E|^2/2 and |B|^2/2 across components present in MS/FLD/. + + Returns NaN if no field dumps are present. + """ + fld = ms / "FLD" + if not fld.is_dir(): + return float("nan") + total = 0.0 + found = False + for comp_dir in fld.iterdir(): + if not comp_dir.is_dir(): + continue + last = _latest_h5(comp_dir) + if last is None: + continue + e = _field_energy_from_dump(last) + if e == e: # not NaN + total += e + found = True + return total if found else float("nan") + + +def _final_iter(ms: Path) -> int: + """Largest iteration number seen across all FLD dumps; -1 if none.""" + fld = ms / "FLD" + if not fld.is_dir(): + return -1 + best = -1 + for comp_dir in fld.iterdir(): + for p in _iter_h5_files(comp_dir): + m = _ITER_RE.search(p.name) + if m: + best = max(best, int(m.group(1))) + return best + + +def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: + """Post-process a finished OSIRIS run. + + Side effects: copies the final-step HDF5 file for each diagnostic into + ``td``; copies ``stdout.log`` / ``stderr.log``; copies the rendered + ``os-stdin``; optionally writes ``ms.tar.gz``. + Returns ``{"metrics": {...}}`` for adept to log to MLflow. + """ + solver = run_output["solver result"] + run_dir: Path = Path(solver["run_dir"]) + ms = run_dir / "MS" + td = Path(td) + upload_all = bool(cfg.get("output", {}).get("upload_all", False)) + whitelist = cfg.get("output", {}).get("diagnostics_to_log") or None + + metrics: dict[str, float] = { + "wall_time_s": float(solver["wall_time"]), + "exit_code": float(solver["exit_code"]), + "field_energy_final": _total_field_energy(ms), + "final_iter": float(_final_iter(ms)), + } + + # Copy final-step HDF5 per diagnostic. + if ms.is_dir(): + for diag_dir in _walk_diag_dirs(ms): + tag = _diagnostic_name(diag_dir) + if whitelist is not None and tag not in whitelist: + continue + last = _latest_h5(diag_dir) + if last is None: + continue + rel = diag_dir.relative_to(ms) + dest_dir = td / "final_step" / rel + dest_dir.mkdir(parents=True, exist_ok=True) + shutil.copy(last, dest_dir / last.name) + + # Always include the rendered deck + stdout/stderr. + for fname in ("os-stdin", "stdout.log", "stderr.log"): + src = run_dir / fname + if src.exists(): + shutil.copy(src, td / fname) + + # Optional full archive. + if upload_all and ms.is_dir(): + archive = td / "ms.tar.gz" + with tarfile.open(archive, "w:gz") as tf: + tf.add(ms, arcname="MS") + + return {"metrics": metrics} diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py new file mode 100644 index 00000000..e1c581e0 --- /dev/null +++ b/adept/osiris/runner.py @@ -0,0 +1,194 @@ +"""Subprocess driver for OSIRIS. + +OSIRIS reads its deck from a file named ``os-stdin`` in the current +working directory by default. This module sets up a per-run work +directory, writes the rendered deck there, invokes ``mpirun`` (or runs +the binary directly when ``mpi_ranks == 1``), and captures stdout/stderr +to files for later artifact upload. +""" + +from __future__ import annotations + +import datetime as _dt +import os +import shlex +import shutil +import subprocess +import threading +import time +import uuid +from pathlib import Path +from typing import Any + + +_OSIRIS_ERR_TOKENS = ("error", "aborting", "(*error*)") +# stderr noise emitted by X11 / mpirun that we should NOT treat as an error. +_OSIRIS_STDERR_NOISE = ( + "No protocol specified", + "MPI_INIT", # benign MPI banner noise on some setups +) + + +def _looks_like_osiris_error(line: str) -> bool: + low = line.strip().lower() + if not low: + return False + if any(noise.lower() in low for noise in _OSIRIS_STDERR_NOISE): + return False + return any(tok in low for tok in _OSIRIS_ERR_TOKENS) + + +def _stream_to_file_and_buffer( + stream, file_path: Path, tail: list[str], tail_max: int = 200 +) -> None: + """Tee a subprocess stream to disk and a bounded in-memory tail.""" + with file_path.open("w") as fh: + for raw in iter(stream.readline, b""): + line = raw.decode("utf-8", errors="replace") + fh.write(line) + fh.flush() + tail.append(line) + if len(tail) > tail_max: + del tail[: len(tail) - tail_max] + + +def _make_run_dir(run_root: Path) -> Path: + run_root.mkdir(parents=True, exist_ok=True) + stamp = _dt.datetime.now().strftime("%Y%m%dT%H%M%S") + name = f"{stamp}_{uuid.uuid4().hex[:8]}" + rd = run_root / name + rd.mkdir() + return rd + + +def run_osiris( + deck_text: str, + *, + binary: str | Path, + mpi_ranks: int = 1, + run_root: str | Path = "./osiris_runs", + env: dict[str, str] | None = None, + mpirun: str = "mpirun", + extra_mpi_args: list[str] | None = None, +) -> dict[str, Any]: + """Run OSIRIS and return run metadata. + + Returns a dict with keys ``run_dir`` (Path), ``exit_code`` (int), + ``wall_time`` (float, seconds), and ``cmd`` (list[str]). + + Raises ``RuntimeError`` on non-zero exit code, with the last lines of + stderr included in the message. + """ + binary = Path(binary).expanduser().resolve() + if not binary.exists(): + raise FileNotFoundError(f"OSIRIS binary not found: {binary}") + + run_dir = _make_run_dir(Path(run_root).expanduser().resolve()) + (run_dir / "os-stdin").write_text(deck_text) + + if mpi_ranks > 1: + cmd = [mpirun, "-n", str(mpi_ranks)] + if extra_mpi_args: + cmd.extend(extra_mpi_args) + cmd.append(str(binary)) + else: + cmd = [str(binary)] + + merged_env = os.environ.copy() + if env: + merged_env.update(env) + + stdout_path = run_dir / "stdout.log" + stderr_path = run_dir / "stderr.log" + stdout_tail: list[str] = [] + stderr_tail: list[str] = [] + + t0 = time.time() + proc = subprocess.Popen( + cmd, + cwd=run_dir, + env=merged_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + bufsize=0, + ) + t_out = threading.Thread( + target=_stream_to_file_and_buffer, + args=(proc.stdout, stdout_path, stdout_tail), + daemon=True, + ) + t_err = threading.Thread( + target=_stream_to_file_and_buffer, + args=(proc.stderr, stderr_path, stderr_tail), + daemon=True, + ) + t_out.start() + t_err.start() + rc = proc.wait() + t_out.join() + t_err.join() + wall_time = time.time() - t0 + + if rc != 0: + tail = "".join(stderr_tail[-50:]) or "(empty stderr)" + raise RuntimeError( + f"OSIRIS exited with status {rc}.\n" + f" cmd: {shlex.join(cmd)}\n" + f" cwd: {run_dir}\n" + f" stderr tail:\n{tail}" + ) + + # OSIRIS can exit 0 even on input-file errors: it prints something + # like 'Error reading ... / aborting...' to stderr (and '(*error*)' + # to stdout) before terminating. Detect both. + err_lines = [ln for ln in stderr_tail if _looks_like_osiris_error(ln)] + err_lines += [ln for ln in stdout_tail if "(*error*)" in ln] + if err_lines: + out_tail = "".join(stdout_tail[-20:]) + err_tail = "".join(stderr_tail[-20:]) + raise RuntimeError( + "OSIRIS reported an error despite exit-code 0:\n" + f" cmd: {shlex.join(cmd)}\n" + f" cwd: {run_dir}\n" + f" stderr tail:\n{err_tail}\n" + f" stdout tail:\n{out_tail}" + ) + + return { + "run_dir": run_dir, + "exit_code": rc, + "wall_time": wall_time, + "cmd": cmd, + } + + +def discover_binary(cfg_binary: str | None, *, dim: int | None = None) -> Path: + """Resolve the OSIRIS binary path. + + Precedence: explicit ``cfg_binary`` > ``OSIRIS_BIN_D`` env var > + ``OSIRIS_BIN`` env var. Returns an existing Path or raises. + """ + candidates: list[str] = [] + if cfg_binary: + candidates.append(cfg_binary) + if dim is not None: + env_key = f"OSIRIS_BIN_{dim}D" + if env_key in os.environ: + candidates.append(os.environ[env_key]) + if "OSIRIS_BIN" in os.environ: + candidates.append(os.environ["OSIRIS_BIN"]) + + for c in candidates: + p = Path(c).expanduser() + if p.exists(): + return p.resolve() + + raise FileNotFoundError( + "No OSIRIS binary found. Set osiris.binary in the manifest or " + "OSIRIS_BIN / OSIRIS_BIN_D in the environment. Tried: " + f"{candidates}" + ) + + +def have_mpirun() -> bool: + return shutil.which("mpirun") is not None diff --git a/configs/osiris/twostream-1d-short.yaml b/configs/osiris/twostream-1d-short.yaml new file mode 100644 index 00000000..4ccec265 --- /dev/null +++ b/configs/osiris/twostream-1d-short.yaml @@ -0,0 +1,16 @@ +solver: osiris + +mlflow: + experiment: osiris-pic1d-twostream + run: smoke-tmax-1 + +osiris: + deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d + binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e + mpi_ranks: 1 + run_root: ./osiris_runs + overrides: + time: {tmax: 1.0} + +output: + upload_all: false diff --git a/configs/osiris/twostream-1d-uploadall.yaml b/configs/osiris/twostream-1d-uploadall.yaml new file mode 100644 index 00000000..2bc0db75 --- /dev/null +++ b/configs/osiris/twostream-1d-uploadall.yaml @@ -0,0 +1,18 @@ +solver: osiris + +mlflow: + experiment: osiris-pic1d-twostream + run: smoke-uploadall + +osiris: + deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d + binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e + mpi_ranks: 1 + run_root: ./osiris_runs + overrides: + time: {tmax: 0.5} + species: + 0: {num_par_x: [128]} + +output: + upload_all: true diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml new file mode 100644 index 00000000..0f82fdcb --- /dev/null +++ b/configs/osiris/twostream-1d.yaml @@ -0,0 +1,21 @@ +solver: osiris + +mlflow: + experiment: osiris-pic1d-twostream + run: cold-equal-beams + +osiris: + deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d + binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e + mpi_ranks: 1 + run_root: ./osiris_runs + # overrides: dict-of-dicts merged into the parsed deck before render. + # Repeated sections (species, zpulse, ...) use integer indices. + # overrides: + # time: {tmax: 5.0} + # species: + # 0: {num_par_x: [1024]} + +output: + upload_all: false + # diagnostics_to_log: [e1] # null/omitted = log final-step for all diagnostics diff --git a/osiris-adept-usage.md b/osiris-adept-usage.md new file mode 100644 index 00000000..2816b243 --- /dev/null +++ b/osiris-adept-usage.md @@ -0,0 +1,140 @@ +# OSIRIS adept module — usage overview + +## File structure + +``` +adept/ (repo root, ~/Desktop/adept/adept/) +├── run.py ← existing CLI entry (unchanged) +├── adept/ +│ ├── _base_.py ← MODIFIED: dispatcher gained `osiris` branch +│ └── osiris/ ← NEW 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, optional MS/ tarball, +│ scalar metrics +├── configs/ +│ └── osiris/ ← NEW: example manifests +│ ├── twostream-1d.yaml full 30-ω_p run +│ ├── twostream-1d-short.yaml tmax=1.0 smoke +│ └── twostream-1d-uploadall.yaml tmax=0.5 + ms.tar.gz +└── tests/ + └── test_osiris/ ← NEW + ├── test_deck_roundtrip.py 15 parser tests + └── test_runner.py 5 runner tests +``` + +## End-to-end data flow + +``` + YAML manifest OSIRIS native deck + (configs/osiris/*.yaml) (any path you point at) + │ │ + ▼ ▼ + run.py --cfg … parse_deck_file() ──┐ + │ │ + ▼ ▼ + ergoExo.setup ──► BaseOsiris.__init__ ──► merge_overrides ──► render_deck + │ │ + ▼ ▼ + cfg["deck"] = flat_dict run_dir/os-stdin + │ │ + ▼ ▼ + log_params → MLflow runner.run_osiris(mpirun…) + │ + ▼ + run_dir/MS/{FLD,PHA,…} + │ + ▼ + post.collect → td/ → MLflow +``` + +## How to run + +```bash +conda activate adept +cd ~/Desktop/adept/adept +python run.py --cfg configs/osiris/twostream-1d-short # smoke +python run.py --cfg configs/osiris/twostream-1d # full +mlflow ui --backend-store-uri file://$(pwd)/mlruns # browse +``` + +## Manifest schema + +```yaml +solver: osiris # required, dispatch key + +mlflow: + experiment: osiris-pic1d-twostream # required + run: cold-equal-beams # required + +osiris: + deck: /path/to/native/deck # required: source of truth + binary: /path/to/osiris-1D.e # or OSIRIS_BIN / OSIRIS_BIN_D + mpi_ranks: 1 # 1 → direct, >1 → mpirun -n N + run_root: ./osiris_runs # parent of per-run dirs + extra_mpi_args: ["--oversubscribe"] # optional, passed to mpirun + + overrides: # optional: applied before render + time: {tmax: 50.0} # merge into the (one) time block + grid: {nx_p: [256]} # refresh an array key + species: # indexed for repeated sections: + 0: {num_par_x: [512]} # species 1: bump ppc + 1: {ufl: [-2.0, 0.0, 0.0]} # species 2: change drift + +output: + upload_all: false # true → also write ms.tar.gz + diagnostics_to_log: null # null = all; or [e1, charge, …] +``` + +Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1)`). Indexed `{0: …, 1: …}` form addresses occurrences of repeated sections (`species`, `udist`, `profile`, `spe_bound`, `diag_species`, `zpulse`, …) in source order. + +## What lands in MLflow + +| Kind | Content | +| ----------- | -------------------------------------------------------------------------------- | +| Params | every deck key, flattened — `deck.grid.nx_p_1:1`, `deck.species_0.ufl_1:3.0`, … | +| Params | the manifest itself — `solver`, `osiris.binary`, `output.upload_all`, … | +| Metrics | `wall_time_s`, `exit_code`, `field_energy_final`, `final_iter`, `run_time`, `postprocess_time` | +| Artifacts | `config.yaml`, `derived_config.yaml`, `units.yaml` (adept stock) | +| Artifacts | `os-stdin` (rendered OSIRIS deck), `stdout.log`, `stderr.log` | +| Artifacts | `final_step///-NNNNNN.h5` (one h5 per diagnostic) | +| Artifacts | `ms.tar.gz` (only when `output.upload_all: true`) | + +## Programmatic use + +```python +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///-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. + +```bash +cp my-new.deck ~/Desktop/pic/projects/whatever/ +cp configs/osiris/twostream-1d.yaml configs/osiris/my-new.yaml +$EDITOR configs/osiris/my-new.yaml # change deck path + mlflow.run +python run.py --cfg configs/osiris/my-new +``` diff --git a/tests/test_osiris/test_deck_roundtrip.py b/tests/test_osiris/test_deck_roundtrip.py new file mode 100644 index 00000000..19190f00 --- /dev/null +++ b/tests/test_osiris/test_deck_roundtrip.py @@ -0,0 +1,137 @@ +"""Round-trip tests for the OSIRIS namelist parser/renderer.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from adept.osiris import deck as osd + + +REAL_DECKS = [ + "/home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d", + "/home/phil/Desktop/pic/projects/twostream-pic/example_deck", + "/home/phil/Desktop/pic/osiris/decks/cuda/os-stdin-1d-therm", + "/home/phil/Desktop/pic/osiris/decks/cuda/os-stdin-2d-therm", +] + + +@pytest.mark.parametrize("path", REAL_DECKS) +def test_roundtrip_identity(path: str) -> None: + text = Path(path).read_text() + parsed = osd.parse_deck(text) + rendered = osd.render_deck(parsed) + reparsed = osd.parse_deck(rendered) + assert parsed == reparsed, f"Round-trip mismatch for {path}" + + +def test_parse_basic_section() -> None: + text = """ + grid + { + nx_p(1:1) = 64, + coordinates = "cartesian", + } + """ + sections = osd.parse_deck(text) + assert sections == [("grid", {"nx_p(1:1)": 64, "coordinates": "cartesian"})] + + +def test_parse_repeated_sections_preserved_order() -> None: + text = """ + species { name = "a", rqm = -1.0, } + species { name = "b", rqm = +1.0, } + """ + sections = osd.parse_deck(text) + assert [(n, p["name"]) for n, p in sections] == [ + ("species", "a"), + ("species", "b"), + ] + + +def test_parse_empty_section() -> None: + text = "current{}\n diag_current { }\n" + sections = osd.parse_deck(text) + assert sections == [("current", {}), ("diag_current", {})] + + +def test_parse_booleans_and_lists() -> None: + text = """ + udist { + uth(1:3) = 0.0707, 0.0707, 0.0707, + ufl(1:3) = -1.0, 0.0, 0.0, + } + spe_bound { if_periodic = .true., } + """ + sections = osd.parse_deck(text) + assert sections[0][1]["uth(1:3)"] == [0.0707, 0.0707, 0.0707] + assert sections[0][1]["ufl(1:3)"] == [-1.0, 0.0, 0.0] + assert sections[1][1]["if_periodic"] is True + + +def test_comment_stripping_preserves_string_with_bang() -> None: + text = 'foo { msg = "hello! world", x = 1, }' + sections = osd.parse_deck(text) + assert sections[0][1]["msg"] == "hello! world" + assert sections[0][1]["x"] == 1 + + +def test_merge_overrides_simple() -> None: + text = "time { tmin = 0.0, tmax = 30.0, }" + s = osd.parse_deck(text) + osd.merge_overrides(s, {"time": {"tmax": 50.0}}) + assert s[0][1]["tmax"] == 50.0 + assert s[0][1]["tmin"] == 0.0 + + +def test_merge_overrides_indexed_repeated_section() -> None: + text = """ + species { name = "a", num_par_x(1:1) = 256, } + species { name = "b", num_par_x(1:1) = 256, } + """ + s = osd.parse_deck(text) + osd.merge_overrides(s, {"species": {0: {"num_par_x": [512]}}}) + assert s[0][1]["num_par_x(1:1)"] == [512] + assert s[1][1]["num_par_x(1:1)"] == 256 + + +def test_merge_overrides_unknown_section_raises() -> None: + s = osd.parse_deck("grid { nx_p(1:1) = 64, }") + with pytest.raises(KeyError): + osd.merge_overrides(s, {"nonexistent": {"foo": 1}}) + + +def test_deck_to_flat_dict_indexes_repeated_sections() -> None: + text = """ + species { name = "a", num_par_x(1:1) = 256, } + species { name = "b", num_par_x(1:1) = 256, } + grid { nx_p(1:1) = 64, } + """ + s = osd.parse_deck(text) + flat = osd.deck_to_flat_dict(s) + assert flat["species_0.name"] == "a" + assert flat["species_1.name"] == "b" + assert flat["grid.nx_p_1:1"] == 64 + + +def test_deck_to_flat_dict_expands_lists() -> None: + s = osd.parse_deck("udist { uth(1:3) = 0.1, 0.2, 0.3, }") + flat = osd.deck_to_flat_dict(s) + assert flat["udist.uth_1:3.0"] == 0.1 + assert flat["udist.uth_1:3.1"] == 0.2 + assert flat["udist.uth_1:3.2"] == 0.3 + + +def test_deck_to_flat_dict_keys_are_mlflow_safe() -> None: + import re + + s = osd.parse_deck( + Path( + "/home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d" + ).read_text() + ) + flat = osd.deck_to_flat_dict(s) + allowed = re.compile(r"^[A-Za-z0-9_./:\- ]+$") + for k in flat: + assert allowed.match(k), f"Unsafe MLflow param key: {k!r}" diff --git a/tests/test_osiris/test_io_and_plots.py b/tests/test_osiris/test_io_and_plots.py new file mode 100644 index 00000000..5dd2b998 --- /dev/null +++ b/tests/test_osiris/test_io_and_plots.py @@ -0,0 +1,93 @@ +"""Tests for the io.py loaders and plots.py canned views.""" + +from __future__ import annotations + +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import numpy as np +import pytest +import xarray as xr + +from adept.osiris import io as oio +from adept.osiris import plots as oplt + + +EXISTING_RUN = Path( + "/home/phil/Desktop/pic/projects/twostream-pic/run0001" +) + + +@pytest.fixture(scope="module") +def run_dir() -> Path: + if not EXISTING_RUN.is_dir(): + pytest.skip("Existing twostream run not present") + return EXISTING_RUN + + +def test_load_grid_h5_shape_and_coords(run_dir: Path) -> None: + da = oio.load_grid_h5(run_dir / "MS" / "FLD" / "e1" / "e1-000600.h5") + assert da.dims == ("x1",) + assert da.shape == (64,) + assert "time" in da.attrs and "iter" in da.attrs + assert da.attrs["iter"] == 600 + assert da.coords["x1"].size == 64 + + +def test_load_phasespace_shape(run_dir: Path) -> None: + p = next((run_dir / "MS/PHA/x1p1/beam_pos").glob("*.h5")) + da = oio.load_phasespace_h5(p) + assert da.ndim == 2 + assert set(da.dims) <= {"x1", "p1"} + + +def test_load_series_stacks_time(run_dir: Path) -> None: + da = oio.load_series(run_dir / "MS" / "FLD" / "e1") + assert da.dims[0] == "t" + # Two-stream baseline ran 30/0.05 = 600 steps, 601 snapshots saved. + assert da.shape[0] == 601 + assert da.shape[1] == 64 + t = da.coords["t"].values + assert t[0] == 0.0 + assert np.isclose(t[-1], 30.0, atol=0.05) + + +def test_list_diagnostics(run_dir: Path) -> None: + diags = oio.list_diagnostics(run_dir) + assert "FLD/e1" in diags + assert any(k.startswith("PHA/x1p1/") for k in diags) + + +def test_field_energy_series_nontrivial(run_dir: Path) -> None: + series = oplt.field_energy_series(run_dir) + assert series.dims == ("t",) + assert series.size > 0 + assert np.all(series.values >= 0) + + +def test_omega_k_returns_axes(run_dir: Path) -> None: + import matplotlib.pyplot as plt + + ser = oio.load_series(run_dir / "MS/FLD/e1") + fig, ax = plt.subplots() + out = oplt.plot_omega_k(ser, ax=ax, show_em=True, show_langmuir=False) + assert out is ax + # Y limits should bracket some non-zero omega. + ylo, yhi = ax.get_ylim() + assert ylo < 0 < yhi + plt.close(fig) + + +def test_save_canned_plots_writes_all_expected(tmp_path: Path, run_dir: Path) -> None: + written = oplt.save_canned_plots(run_dir, tmp_path, v_th=0.0707) + # At minimum: one spacetime + omega_k for e1, energy_vs_time, both species PS. + assert "spacetime/e1" in written + assert "omega_k/e1" in written + assert "energy_vs_time" in written + assert "phasespace/beam_pos/x1p1" in written + assert "phasespace/beam_neg/x1p1" in written + for path in written.values(): + assert path.exists() and path.stat().st_size > 0 diff --git a/tests/test_osiris/test_runner.py b/tests/test_osiris/test_runner.py new file mode 100644 index 00000000..1371b116 --- /dev/null +++ b/tests/test_osiris/test_runner.py @@ -0,0 +1,58 @@ +"""Smoke tests for the OSIRIS subprocess runner.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from adept.osiris import runner + + +def test_discover_binary_explicit_wins(tmp_path: Path) -> None: + fake = tmp_path / "fake-osiris" + fake.write_text("") + out = runner.discover_binary(str(fake)) + assert out == fake.resolve() + + +def test_discover_binary_env_fallback(tmp_path: Path, monkeypatch) -> None: + fake = tmp_path / "fake-osiris-1d" + fake.write_text("") + monkeypatch.setenv("OSIRIS_BIN_1D", str(fake)) + out = runner.discover_binary(None, dim=1) + assert out == fake.resolve() + + +def test_discover_binary_missing_raises() -> None: + with pytest.raises(FileNotFoundError): + runner.discover_binary("/no/such/path/exists", dim=1) + + +def test_run_osiris_missing_binary_raises(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError): + runner.run_osiris( + "node_conf {}", + binary="/no/such/binary", + mpi_ranks=1, + run_root=tmp_path, + ) + + +@pytest.mark.skipif( + not Path("/home/phil/Desktop/pic/osiris/bin/osiris-1D.e").exists(), + reason="osiris-1D.e not built", +) +def test_run_osiris_invalid_deck_raises(tmp_path: Path) -> None: + # A deck with a recognized section but garbage inside: OSIRIS exits 0 + # but writes 'Error reading ... / aborting...' to stderr. Our runner + # turns that into a RuntimeError so it doesn't slip past silently. + with pytest.raises(RuntimeError) as excinfo: + runner.run_osiris( + "node_conf { node_number(1:1) = junk_value, }", + binary="/home/phil/Desktop/pic/osiris/bin/osiris-1D.e", + mpi_ranks=1, + run_root=tmp_path, + ) + assert "OSIRIS" in str(excinfo.value) From 509cad2c4b655c9bd91e357d4f19cae3d2ae704f Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 22 May 2026 12:55:18 -0700 Subject: [PATCH 02/49] .gitignore changes --- .gitignore | 4 ++++ osiris-adept-usage.md => docs/osiris-adept-usage.md | 0 2 files changed, 4 insertions(+) rename osiris-adept-usage.md => docs/osiris-adept-usage.md (100%) diff --git a/.gitignore b/.gitignore index d691a617..e633fc85 100644 --- a/.gitignore +++ b/.gitignore @@ -129,6 +129,10 @@ dmypy.json *mlruns/ mlflow.db +# OSIRIS subprocess scratch (per-run dirs created by adept.osiris.runner) +osiris_runs/ +scratch/ + # Pyre type checker .pyre/ diff --git a/osiris-adept-usage.md b/docs/osiris-adept-usage.md similarity index 100% rename from osiris-adept-usage.md rename to docs/osiris-adept-usage.md From 49d583f37f69523376ed0206c14d9fa9e5ae7b86 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 27 May 2026 10:19:30 -0700 Subject: [PATCH 03/49] added osiris as a solver --- adept/_base_.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/adept/_base_.py b/adept/_base_.py index 949bb5a5..0eb28c1d 100644 --- a/adept/_base_.py +++ b/adept/_base_.py @@ -326,6 +326,9 @@ def _get_adept_module_(self, cfg: dict) -> ADEPTModule: elif cfg["solver"] == "pic-1d": from adept.pic1d import BasePIC1D as this_module + elif cfg["solver"] == "osiris": + from adept.osiris import BaseOsiris as this_module + else: raise NotImplementedError("This solver approach has not been implemented yet") From 761a23d925b9eb4e97e902b03b2da9b898a512bd Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 27 May 2026 23:36:38 -0700 Subject: [PATCH 04/49] Uploaded original config.yaml instead of translated osiris deck --- adept/_base_.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/adept/_base_.py b/adept/_base_.py index 0eb28c1d..d2c7d761 100644 --- a/adept/_base_.py +++ b/adept/_base_.py @@ -337,6 +337,12 @@ def _get_adept_module_(self, cfg: dict) -> ADEPTModule: def _setup_(self, cfg: dict, td: str, adept_module: ADEPTModule = None, log: bool = True) -> dict[str, Module]: from adept.utils import log_params + # Snapshot the config as provided, before the module mutates it in + # place (e.g. the OSIRIS module injects the parsed deck). config.yaml + # is the original config; the processed cfg lands in derived_config.yaml + # and the logged params. + original_cfg = deepcopy(cfg) + if adept_module is None: self.adept_module = self._get_adept_module_(cfg) else: @@ -345,7 +351,7 @@ def _setup_(self, cfg: dict, td: str, adept_module: ADEPTModule = None, log: boo # dump raw config if log: with open(os.path.join(td, "config.yaml"), "w") as fi: - yaml.dump(self.adept_module.cfg, fi) + yaml.dump(original_cfg, fi) # dump units quants_dict = self.adept_module.write_units() # writes the units to the temporary directory From 33ca8a06ba8b053f4d1b5144deddfc6139b0dd4c Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Thu, 28 May 2026 12:36:31 -0700 Subject: [PATCH 05/49] Aggregate osiris h5 files into an xarrary; save and upload --- adept/osiris/io.py | 70 +++++++++++- adept/osiris/post.py | 52 ++------- configs/osiris/twostream-1d-short.yaml | 3 - configs/osiris/twostream-1d-uploadall.yaml | 18 --- configs/osiris/twostream-1d.yaml | 3 +- docs/osiris-adept-usage.md | 6 +- tests/test_osiris/test_post_netcdf.py | 126 +++++++++++++++++++++ 7 files changed, 209 insertions(+), 69 deletions(-) delete mode 100644 configs/osiris/twostream-1d-uploadall.yaml create mode 100644 tests/test_osiris/test_post_netcdf.py diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 7bbfe528..18298629 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -19,6 +19,7 @@ from __future__ import annotations +import json import re from pathlib import Path @@ -26,7 +27,6 @@ import numpy as np import xarray as xr - _ITER_RE = re.compile(r"-(\d+)\.h5$") @@ -190,3 +190,71 @@ def list_diagnostics(run_dir: str | Path) -> dict[str, Path]: if any(c.suffix == ".h5" for c in d.iterdir()): out[str(d.relative_to(ms))] = d return out + + +def _coerce_attr(v): + """Make an attr value writable by the netCDF backends. + + ``load_series`` stores some attrs as dicts (per-axis units/long-names), + which netCDF cannot serialize; those are handled separately by + :func:`series_to_dataset`. Here we JSON-encode any stray dict and turn + lists/tuples into arrays. + """ + if isinstance(v, dict): + return json.dumps(v) + if isinstance(v, (list, tuple)): + return np.asarray(v) + return v + + +def series_to_dataset(da: xr.DataArray) -> xr.Dataset: + """Wrap a :func:`load_series` DataArray into a netCDF-serializable Dataset. + + The dict-valued ``axis_units`` / ``axis_long_names`` attrs are lifted onto + the matching coordinate variables (the CF-idiomatic place for them), and + any remaining non-scalar attrs are coerced so ``to_netcdf`` succeeds. + """ + da = da.copy() + axis_units = da.attrs.pop("axis_units", {}) or {} + axis_long = da.attrs.pop("axis_long_names", {}) or {} + da.attrs = {k: _coerce_attr(v) for k, v in da.attrs.items() if v is not None} + ds = da.to_dataset(name=da.name) + for dim, units in axis_units.items(): + if dim in ds.coords: + ds[dim].attrs["units"] = units + for dim, long_name in axis_long.items(): + if dim in ds.coords: + ds[dim].attrs["long_name"] = long_name + return ds + + +def save_run_datasets( + run_dir: str | Path, + out_dir: str | Path, + diagnostics: list[str] | set[str] | None = None, +) -> list[Path]: + """Convert each diagnostic's full time history to a netCDF file. + + One file per diagnostic is written under ``out_dir``, mirroring the OSIRIS + ``MS/`` layout (e.g. ``out_dir/FLD/e1.nc``, + ``out_dir/PHA/x1p1/beam_pos.nc``). Each file holds the stacked ``(t, ...)`` + series — every time slice OSIRIS dumped for that diagnostic. + + ``diagnostics``, when given, whitelists which diagnostics to convert, + matched against either the relative path (``"FLD/e1"``) or the leaf name + (``"e1"``). Returns the list of written paths. + """ + out_dir = Path(out_dir) + diags = list_diagnostics(run_dir) + written: list[Path] = [] + for relpath in sorted(diags): + if diagnostics is not None and ( + relpath not in diagnostics and Path(relpath).name not in diagnostics + ): + continue + ds = series_to_dataset(load_series(diags[relpath])) + dest = out_dir / f"{relpath}.nc" + dest.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(dest, engine="h5netcdf") + written.append(dest) + return written diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 744ee41e..a9a73cac 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -1,19 +1,19 @@ """Post-processing for OSIRIS runs. After ``BaseOsiris.__call__`` finishes, ``collect`` walks the ``MS/`` -output tree, copies the final-step HDF5 file from each diagnostic to the -adept temp dir (so MLflow uploads it), optionally tarballs the whole -tree, and returns scalar metrics. +output tree, converts each diagnostic's full time history into an xarray +netCDF file in the adept temp dir (so MLflow uploads it instead of the raw +HDF5 dumps), copies the deck / stdout / stderr, and returns scalar metrics. """ from __future__ import annotations import re import shutil -import tarfile from pathlib import Path from typing import Any +from adept.osiris import io as _io # OSIRIS dump filenames look like ``e1-000600.h5`` or # ``x1p1-beam_pos-000060.h5``. The iteration number is always a @@ -39,15 +39,6 @@ def keyfn(p: Path) -> int: return max(files, key=keyfn) -def _walk_diag_dirs(ms: Path) -> list[Path]: - """Return every directory under ``ms`` that contains ``.h5`` files.""" - out: list[Path] = [] - for p in ms.rglob("*"): - if p.is_dir() and any(c.suffix == ".h5" for c in p.iterdir()): - out.append(p) - return out - - def _field_energy_from_dump(h5_path: Path) -> float: """Integrate ``field^2 * cell_volume`` for a single dump file. @@ -74,11 +65,6 @@ def _field_energy_from_dump(h5_path: Path) -> float: return 0.5 * float((arr.astype("float64") ** 2).sum()) * dvol -def _diagnostic_name(diag_dir: Path) -> str: - """Short tag for a diagnostic directory, used in metric / artifact names.""" - return diag_dir.name - - def _total_field_energy(ms: Path) -> float: """Sum |E|^2/2 and |B|^2/2 across components present in MS/FLD/. @@ -119,17 +105,17 @@ def _final_iter(ms: Path) -> int: def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: """Post-process a finished OSIRIS run. - Side effects: copies the final-step HDF5 file for each diagnostic into - ``td``; copies ``stdout.log`` / ``stderr.log``; copies the rendered - ``os-stdin``; optionally writes ``ms.tar.gz``. + Side effects: converts each diagnostic's full time history into a netCDF + file under ``td/binary/`` (mirroring the ``MS/`` layout) so MLflow gets + combined xarray datasets rather than the raw HDF5 dumps; copies the + rendered ``os-stdin`` plus ``stdout.log`` / ``stderr.log``. Returns ``{"metrics": {...}}`` for adept to log to MLflow. """ solver = run_output["solver result"] run_dir: Path = Path(solver["run_dir"]) ms = run_dir / "MS" td = Path(td) - upload_all = bool(cfg.get("output", {}).get("upload_all", False)) - whitelist = cfg.get("output", {}).get("diagnostics_to_log") or None + whitelist = (cfg.get("output") or {}).get("diagnostics_to_log") or None metrics: dict[str, float] = { "wall_time_s": float(solver["wall_time"]), @@ -138,19 +124,9 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: "final_iter": float(_final_iter(ms)), } - # Copy final-step HDF5 per diagnostic. + # Convert each diagnostic's time history to an xarray netCDF. if ms.is_dir(): - for diag_dir in _walk_diag_dirs(ms): - tag = _diagnostic_name(diag_dir) - if whitelist is not None and tag not in whitelist: - continue - last = _latest_h5(diag_dir) - if last is None: - continue - rel = diag_dir.relative_to(ms) - dest_dir = td / "final_step" / rel - dest_dir.mkdir(parents=True, exist_ok=True) - shutil.copy(last, dest_dir / last.name) + _io.save_run_datasets(run_dir, td / "binary", diagnostics=whitelist) # Always include the rendered deck + stdout/stderr. for fname in ("os-stdin", "stdout.log", "stderr.log"): @@ -158,10 +134,4 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: if src.exists(): shutil.copy(src, td / fname) - # Optional full archive. - if upload_all and ms.is_dir(): - archive = td / "ms.tar.gz" - with tarfile.open(archive, "w:gz") as tf: - tf.add(ms, arcname="MS") - return {"metrics": metrics} diff --git a/configs/osiris/twostream-1d-short.yaml b/configs/osiris/twostream-1d-short.yaml index 4ccec265..fcbad707 100644 --- a/configs/osiris/twostream-1d-short.yaml +++ b/configs/osiris/twostream-1d-short.yaml @@ -11,6 +11,3 @@ osiris: run_root: ./osiris_runs overrides: time: {tmax: 1.0} - -output: - upload_all: false diff --git a/configs/osiris/twostream-1d-uploadall.yaml b/configs/osiris/twostream-1d-uploadall.yaml deleted file mode 100644 index 2bc0db75..00000000 --- a/configs/osiris/twostream-1d-uploadall.yaml +++ /dev/null @@ -1,18 +0,0 @@ -solver: osiris - -mlflow: - experiment: osiris-pic1d-twostream - run: smoke-uploadall - -osiris: - deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d - binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e - mpi_ranks: 1 - run_root: ./osiris_runs - overrides: - time: {tmax: 0.5} - species: - 0: {num_par_x: [128]} - -output: - upload_all: true diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index 0f82fdcb..cc3ac5c3 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -17,5 +17,4 @@ osiris: # 0: {num_par_x: [1024]} output: - upload_all: false - # diagnostics_to_log: [e1] # null/omitted = log final-step for all diagnostics + # diagnostics_to_log: [e1] # null/omitted = convert every diagnostic's full time history to netCDF diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index 2816b243..eaa2e641 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -91,7 +91,6 @@ osiris: 1: {ufl: [-2.0, 0.0, 0.0]} # species 2: change drift output: - upload_all: false # true → also write ms.tar.gz diagnostics_to_log: null # null = all; or [e1, charge, …] ``` @@ -102,12 +101,11 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | Kind | Content | | ----------- | -------------------------------------------------------------------------------- | | Params | every deck key, flattened — `deck.grid.nx_p_1:1`, `deck.species_0.ufl_1:3.0`, … | -| Params | the manifest itself — `solver`, `osiris.binary`, `output.upload_all`, … | +| Params | the manifest itself — `solver`, `osiris.binary`, `output.diagnostics_to_log`, … | | Metrics | `wall_time_s`, `exit_code`, `field_energy_final`, `final_iter`, `run_time`, `postprocess_time` | | Artifacts | `config.yaml`, `derived_config.yaml`, `units.yaml` (adept stock) | | Artifacts | `os-stdin` (rendered OSIRIS deck), `stdout.log`, `stderr.log` | -| Artifacts | `final_step///-NNNNNN.h5` (one h5 per diagnostic) | -| Artifacts | `ms.tar.gz` (only when `output.upload_all: true`) | +| Artifacts | `binary//.nc` — one xarray netCDF per diagnostic, holding the full `(t, …)` time history (replaces the raw h5 dumps) | ## Programmatic use diff --git a/tests/test_osiris/test_post_netcdf.py b/tests/test_osiris/test_post_netcdf.py new file mode 100644 index 00000000..81bb7527 --- /dev/null +++ b/tests/test_osiris/test_post_netcdf.py @@ -0,0 +1,126 @@ +"""post.collect converts OSIRIS diagnostics to netCDF time-series. + +These tests synthesize a tiny OSIRIS-shaped run so they're self-contained +(no dependency on a real run on disk). +""" + +from __future__ import annotations + +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr + +from adept.osiris import io as oio +from adept.osiris import post as opost + + +def _write_dump(path: Path, name: str, data: np.ndarray, t: float, it: int, axes) -> None: + """Write one OSIRIS-style HDF5 dump. + + ``axes`` is a list of ``(name, long_name, units, min, max)`` in OSIRIS + (Fortran) order, i.e. the last entry is the fastest-varying numpy axis. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + f.attrs["NAME"] = np.array([name.encode()], dtype="S256") + f.attrs["LABEL"] = np.array([name.encode()], dtype="S256") + f.attrs["UNITS"] = np.array([b"a.u."], dtype="S256") + f.attrs["TIME UNITS"] = np.array([rb"1 / \omega_p"], dtype="S256") + f.attrs["TYPE"] = np.array([b"grid"], dtype="S4") + ax = f.create_group("AXIS") + for i, (an, ln, un, lo, hi) in enumerate(axes, start=1): + d = ax.create_dataset(f"AXIS{i}", data=np.array([lo, hi], dtype="float64")) + d.attrs["NAME"] = np.array([an.encode()], dtype="S256") + d.attrs["LONG_NAME"] = np.array([ln.encode()], dtype="S256") + d.attrs["UNITS"] = np.array([un.encode()], dtype="S256") + d.attrs["TYPE"] = np.array([b"linear"], dtype="S6") + sim = f.create_group("SIMULATION") + sim.attrs["DT"] = np.array([0.05]) + sim.attrs["NDIMS"] = np.array([1], dtype="int32") + sim.attrs["NX"] = np.array([data.shape[-1]], dtype="int32") + sim.attrs["XMIN"] = np.array([axes[-1][3]]) + sim.attrs["XMAX"] = np.array([axes[-1][4]]) + f.create_dataset(name, data=data.astype("float32")) + + +def _make_run(root: Path, n_steps: int = 4, nx: int = 8) -> Path: + """Build a run dir with an FLD/e1 field diagnostic over ``n_steps`` dumps.""" + run_dir = root / "run" + e1 = run_dir / "MS" / "FLD" / "e1" + rng = np.random.default_rng(0) + for k in range(n_steps): + it = k * 10 + _write_dump( + e1 / f"e1-{it:06d}.h5", + "e1", + rng.standard_normal(nx), + t=k * 0.5, + it=it, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + (run_dir / "os-stdin").write_text("node_conf\n{\n}\n") + (run_dir / "stdout.log").write_text("ok\n") + (run_dir / "stderr.log").write_text("") + return run_dir + + +def test_save_run_datasets_writes_full_timeseries(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=4, nx=8) + out = tmp_path / "binary" + written = oio.save_run_datasets(run_dir, out) + + assert written == [out / "FLD" / "e1.nc"] + ds = xr.open_dataset(out / "FLD" / "e1.nc", engine="h5netcdf") + try: + assert "e1" in ds.data_vars + assert ds.sizes == {"t": 4, "x1": 8} # every time slice kept + assert list(ds["t"].values) == [0.0, 0.5, 1.0, 1.5] + assert list(ds["iter"].values) == [0, 10, 20, 30] + # axis metadata moved onto the coordinate (netCDF-serializable) + assert ds["x1"].attrs["units"] == r"c / \omega_p" + finally: + ds.close() + + +def test_collect_emits_netcdf_not_h5(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=3, nx=8) + run_output = {"solver result": {"run_dir": str(run_dir), "wall_time": 1.0, "exit_code": 0}} + cfg = {"osiris": {"deck": str(run_dir / "os-stdin")}, "output": {}} + + td = tmp_path / "td" + td.mkdir() + result = opost.collect(run_output, cfg, str(td)) + + # No raw OSIRIS dumps uploaded; the diagnostic is a netCDF time-series. + assert not list(td.rglob("*.h5")) + assert (td / "binary" / "FLD" / "e1.nc").exists() + # Deck + logs still copied; metrics still produced. + assert (td / "os-stdin").exists() + assert (td / "stdout.log").exists() + assert result["metrics"]["final_iter"] == 20.0 + + +def test_collect_respects_diagnostics_whitelist(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=2, nx=8) + # Add a second diagnostic that should be filtered out. + _write_dump( + run_dir / "MS" / "FLD" / "e2" / "e2-000000.h5", + "e2", + np.zeros(8), + t=0.0, + it=0, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + run_output = {"solver result": {"run_dir": str(run_dir), "wall_time": 1.0, "exit_code": 0}} + cfg = {"osiris": {"deck": str(run_dir / "os-stdin")}, "output": {"diagnostics_to_log": ["e1"]}} + + td = tmp_path / "td" + td.mkdir() + opost.collect(run_output, cfg, str(td)) + + assert (td / "binary" / "FLD" / "e1.nc").exists() + assert not (td / "binary" / "FLD" / "e2.nc").exists() From 3c6f257ed67adda42702243772bb79217505a1b1 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Thu, 28 May 2026 15:40:36 -0700 Subject: [PATCH 06/49] Brought plots up to parity with other solvers --- adept/osiris/io.py | 103 +++++ adept/osiris/plots.py | 421 +++++++++++++++++--- adept/osiris/post.py | 25 ++ tests/test_osiris/test_diagnostics_plots.py | 166 ++++++++ 4 files changed, 655 insertions(+), 60 deletions(-) create mode 100644 tests/test_osiris/test_diagnostics_plots.py diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 18298629..5be574ea 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -258,3 +258,106 @@ def save_run_datasets( ds.to_netcdf(dest, engine="h5netcdf") written.append(dest) return written + + +def _parse_hist_table(path: Path) -> tuple[np.ndarray, np.ndarray] | None: + """Read a whitespace-delimited OSIRIS ``HIST`` table. + + Skips blank and ``!``/``#`` comment lines and any non-numeric header row. + Assumes the OSIRIS column convention ``iteration time `` and + returns ``(t, values)`` where ``t`` is shape ``(n,)`` and ``values`` is + ``(n, ncols - 2)``. Returns ``None`` if nothing numeric is found. + """ + rows: list[list[float]] = [] + with open(path) as fh: + for line in fh: + s = line.strip() + if not s or s[0] in "!#": + continue + try: + rows.append([float(tok) for tok in s.split()]) + except ValueError: + continue # header row of column names + if not rows: + return None + width = min(len(r) for r in rows) + if width < 2: + return None + arr = np.array([r[:width] for r in rows], dtype="float64") + return arr[:, 1], arr[:, 2:] + + +def _interp_onto(src_t: np.ndarray, src_v: np.ndarray, ref_t: np.ndarray) -> np.ndarray: + """Linear-interpolate ``src_v(src_t)`` onto ``ref_t`` (identity if aligned).""" + if src_t.shape == ref_t.shape and np.allclose(src_t, ref_t): + return src_v + return np.interp(ref_t, src_t, src_v) + + +def load_hist_energy(run_dir: str | Path) -> xr.Dataset | None: + """Parse OSIRIS ``HIST/`` scalar energy time-history files, if present. + + OSIRIS writes whitespace-delimited ASCII time-history tables under + ``HIST/`` when energy diagnostics are enabled in the deck. Recognized: + + - ``*fld_ene*`` — per-component field energy; value columns are summed into + a single ``field_energy`` series. + - ``*par*_ene*`` — per-species particle (kinetic) energy; each file's value + columns are summed into ``kinetic_``. + + Returns an ``xr.Dataset`` on a shared ``t`` axis containing whatever was + found — ``field_energy``, ``kinetic_``, ``kinetic_total``, and + ``total`` (= field + kinetic, with ``attrs['total_drift_frac']`` the + fractional spread of the total) when both halves are present — or ``None`` + if there is no ``HIST/`` directory or nothing parseable in it. Per-file time + axes that disagree are interpolated onto the field-energy time axis. + + Note: the column convention assumed is the documented OSIRIS + ``iteration time `` layout; validate against a real run with + energy diagnostics enabled before relying on absolute magnitudes. + """ + hist = Path(run_dir) / "HIST" + if not hist.is_dir(): + return None + + field: tuple[np.ndarray, np.ndarray] | None = None + kinetic: dict[str, tuple[np.ndarray, np.ndarray]] = {} + for f in sorted(hist.iterdir()): + if not f.is_file() or "ene" not in f.name.lower(): + continue + parsed = _parse_hist_table(f) + if parsed is None or parsed[1].size == 0: + continue + t, values = parsed + total_per_row = values.sum(axis=1) + if "fld" in f.name.lower(): + field = (t, total_per_row) + elif "par" in f.name.lower(): + kinetic[f.stem] = (t, total_per_row) + + if field is None and not kinetic: + return None + + ref_t = field[0] if field is not None else next(iter(kinetic.values()))[0] + data_vars: dict[str, tuple] = {} + if field is not None: + data_vars["field_energy"] = ("t", _interp_onto(field[0], field[1], ref_t)) + + kin_total: np.ndarray | None = None + for stem, (t, e) in kinetic.items(): + ei = _interp_onto(t, e, ref_t) + data_vars[f"kinetic_{stem}"] = ("t", ei) + kin_total = ei if kin_total is None else kin_total + ei + if kin_total is not None: + data_vars["kinetic_total"] = ("t", kin_total) + + attrs: dict[str, float] = {} + if field is not None and kin_total is not None: + total = data_vars["field_energy"][1] + kin_total + data_vars["total"] = ("t", total) + denom = float(np.max(np.abs(total))) or 1.0 + attrs["total_drift_frac"] = float((total.max() - total.min()) / denom) + + ds = xr.Dataset(data_vars, coords={"t": ref_t}, attrs=attrs) + ds["t"].attrs.update(long_name="time", units=r"1/\omega_p") + return ds diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index ed58aa0a..d2fc079c 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -8,23 +8,36 @@ PNGs for the standard set: out_dir/ - spacetime/.png (t, x) heatmaps of every FLD diagnostic - phasespace//.png final-step (x, p) heatmap per species - energy_vs_time.png total field energy time-trace - omega_k/.png 2-D FFT (k, ω) dispersion plot + spacetime/.png (t, x) heatmap of every FLD diagnostic + spacetime_log/.png log10|·| of the same + lineouts/.png value-vs-x snapshots at sampled times + moments//.png (t, x) per-species density moments + moments//_log.png log10 of the same + moments//lineouts/.png moment snapshots at sampled times + phasespace//.png final-step (x, p) heatmap per species + phasespace_evolution//.png (x, p) heatmaps at sampled times + omega_k/.png 2-D FFT (k, ω) dispersion plot + energy_vs_time.png total field energy time-trace + energy_components_vs_time.png E-field / B-field / total energy + total_energy_vs_time.png field + kinetic conservation (needs HIST/) """ from __future__ import annotations from pathlib import Path +import matplotlib + +# Canned plots are written to disk in batch / headless runs (e.g. Perlmutter +# compute nodes), so force a non-interactive backend before importing pyplot. +matplotlib.use("Agg") + import matplotlib.pyplot as plt import numpy as np import xarray as xr from adept.osiris import io as _io - # --- low-level plotters ---------------------------------------------------- @@ -56,14 +69,47 @@ def plot_spacetime( plt.colorbar( mesh, ax=ax, - label=f"log10 |{da.name}|" if log else da.name, + label=f"log10 |{_value_label(da)}|" if log else _value_label(da), ) ax.set_xlabel(_axis_label(da, "t")) ax.set_ylabel(_axis_label(da, xname)) - ax.set_title(title or f"{da.name} spacetime") + scale = "log10 " if log else "" + ax.set_title(title or f"{_long_name(da)} — {scale}spacetime (x vs t)") return ax +def plot_lineouts( + series: xr.DataArray | str | Path, + *, + n_panels: int = 8, + col_wrap: int = 4, + title: str | None = None, +) -> plt.Figure: + """Faceted value-vs-x snapshots of a ``(t, x)`` series at sampled times. + + Subsamples the time axis to roughly ``n_panels`` snapshots (matching the + ``t_skip = nt // 8`` convention used by the other adept solvers) and lays + them out in a ``col_wrap`` grid. Returns the ``Figure`` so the caller can + save / close it. + """ + da = _decorate(_ensure_series(series)) + if da.ndim != 2: + raise ValueError( + f"plot_lineouts expects a 2D (t, x) array; got dims {da.dims}" + ) + nt = da.coords["t"].size + t_skip = max(1, nt // n_panels) + sl = da.isel(t=slice(0, None, t_skip)) + xname = next(d for d in da.dims if d != "t") + g = sl.plot(x=xname, col="t", col_wrap=min(col_wrap, sl.coords["t"].size)) + g.set_xlabels(_axis_label(da, xname)) + g.set_ylabels(_value_label(da)) + g.fig.suptitle( + title or f"{_long_name(da)} — lineouts vs x at sampled times", y=1.02 + ) + return g.fig + + def plot_phasespace( da: xr.DataArray | str | Path, ax: plt.Axes | None = None, @@ -88,15 +134,58 @@ def plot_phasespace( x0, x1 = da.coords[d0].values, da.coords[d1].values mesh = ax.pcolormesh(x0, x1, arr.T, shading="auto", cmap=cmap) plt.colorbar( - mesh, ax=ax, label=f"log10 {da.name}" if log else da.name + mesh, ax=ax, label=f"log10 {_value_label(da)}" if log else _value_label(da) ) ax.set_xlabel(_axis_label(da, d0)) ax.set_ylabel(_axis_label(da, d1)) t = da.attrs.get("time", float("nan")) - ax.set_title(title or f"{da.name} t = {t:.3g}") + scale = "log10 " if log else "" + ax.set_title(title or f"{_long_name(da)} — {scale}phase space (t = {t:.3g} 1/$\\omega_p$)") return ax +def plot_phasespace_evolution( + series: xr.DataArray | str | Path, + *, + n_panels: int = 8, + col_wrap: int = 4, + log: bool = True, + cmap: str = "viridis", + title: str | None = None, +) -> plt.Figure: + """Faceted ``(x, p)`` phase-space heatmaps at sampled times. + + Subsamples a stacked ``(t, p, x)`` phase-space series (as returned by + ``io.load_series``) to ~``n_panels`` snapshots so the time evolution is + visible, rather than only the final step. Returns the ``Figure``. + """ + da = series if isinstance(series, xr.DataArray) else _io.load_series(series) + da = _decorate(da) + if da.ndim != 3: + raise ValueError( + f"plot_phasespace_evolution expects (t, p, x); got dims {da.dims}" + ) + nt = da.coords["t"].size + t_skip = max(1, nt // n_panels) + sl = da.isel(t=slice(0, None, t_skip)) + plot_da = np.log10(np.abs(sl) + 1e-30) if log else sl + # Convention: spatial axis horizontal, momentum axis vertical. + spatial = [d for d in da.dims if d != "t" and str(d).startswith("x")] + moment = [d for d in da.dims if d != "t" and str(d).startswith("p")] + facet_kw: dict = {"cmap": cmap} + if spatial and moment: + facet_kw.update(x=spatial[0], y=moment[0]) + g = plot_da.plot(col="t", col_wrap=min(col_wrap, sl.coords["t"].size), **facet_kw) + if spatial and moment: + g.set_xlabels(_axis_label(da, spatial[0])) + g.set_ylabels(_axis_label(da, moment[0])) + scale = "log10 " if log else "" + g.fig.suptitle( + title or f"{_long_name(da)} — {scale}phase space at sampled times", y=1.02 + ) + return g.fig + + def field_energy_series(run_dir: str | Path) -> xr.DataArray: """Sum ``(|E|^2 + |B|^2) / 2`` over space at every saved step. @@ -121,13 +210,7 @@ def field_energy_series(run_dir: str | Path) -> xr.DataArray: if h5.suffix != ".h5": continue da = _io.load_grid_h5(h5) - # cell volume from coord spacing in each dim - dvol = 1.0 - for dim in da.dims: - c = da.coords[dim].values - if c.size > 1: - dvol *= float(c[1] - c[0]) - e = 0.5 * float((da.values ** 2).sum()) * dvol + e = 0.5 * float((da.values ** 2).sum()) * _cell_volume(da) it = int(da.attrs["iter"]) by_iter.setdefault(it, {})[comp] = e times.setdefault(it, float(da.attrs["time"])) @@ -167,8 +250,125 @@ def plot_energy_vs_time( if log: ax.set_yscale("log") ax.set_xlabel(r"t [$1/\omega_p$]") - ax.set_ylabel(r"$\int (|E|^2 + |B|^2)/2\ d^Dx$") - ax.set_title(title or "field energy") + ax.set_ylabel(r"$\int (|E|^2 + |B|^2)/2\ d^Dx$ [code units]") + ax.set_title(title or "Total field energy vs time") + ax.grid(True, which="both", alpha=0.3) + return ax + + +def field_energy_components(run_dir: str | Path) -> xr.Dataset: + """E-field, B-field, and total field energy vs time from ``MS/FLD`` dumps. + + Like :func:`field_energy_series` but keeps the electric (``e1/e2/e3``) and + magnetic (``b1/b2/b3``) contributions separate. Returns a ``Dataset`` with + data variables ``E_energy``, ``B_energy`` and ``total_field_energy`` on a + shared ``t`` axis. Components not dumped are simply omitted from their sum. + """ + run_dir = Path(run_dir) + fld = run_dir / "MS" / "FLD" + if not fld.is_dir(): + raise FileNotFoundError(f"No MS/FLD under {run_dir}") + groups = {"E_energy": ("e1", "e2", "e3"), "B_energy": ("b1", "b2", "b3")} + by_iter: dict[int, dict[str, float]] = {} + times: dict[int, float] = {} + for group, comps in groups.items(): + for comp in comps: + d = fld / comp + if not d.is_dir(): + continue + for h5 in sorted(d.iterdir()): + if h5.suffix != ".h5": + continue + da = _io.load_grid_h5(h5) + e = 0.5 * float((da.values ** 2).sum()) * _cell_volume(da) + it = int(da.attrs["iter"]) + rec = by_iter.setdefault(it, {"E_energy": 0.0, "B_energy": 0.0}) + rec[group] += e + times.setdefault(it, float(da.attrs["time"])) + if not by_iter: + raise RuntimeError(f"No field dumps under {fld}") + iters = sorted(by_iter) + t = np.array([times[i] for i in iters]) + e_arr = np.array([by_iter[i]["E_energy"] for i in iters]) + b_arr = np.array([by_iter[i]["B_energy"] for i in iters]) + coords = {"t": t, "iter": ("t", np.asarray(iters))} + ds = xr.Dataset( + { + "E_energy": ("t", e_arr), + "B_energy": ("t", b_arr), + "total_field_energy": ("t", e_arr + b_arr), + }, + coords=coords, + ) + ds["E_energy"].attrs.update(long_name="electric field energy", units="code") + ds["B_energy"].attrs.update(long_name="magnetic field energy", units="code") + ds["total_field_energy"].attrs.update(long_name="total field energy", units="code") + ds["t"].attrs.update(long_name="time", units=r"1/\omega_p") + return ds + + +def plot_energy_components( + src: xr.Dataset | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = True, + title: str | None = None, +) -> plt.Axes: + """Overlay E-field, B-field, and total field energy vs time.""" + ds = src if isinstance(src, xr.Dataset) else field_energy_components(src) + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + t = ds["t"].values + labels = { + "E_energy": r"$\int |E|^2/2$ (electric)", + "B_energy": r"$\int |B|^2/2$ (magnetic)", + "total_field_energy": "total field", + } + for name, label in labels.items(): + if name in ds: + ax.plot(t, ds[name].values, label=label) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel("field energy [code units]") + ax.set_title(title or "Field energy components vs time") + ax.legend() + ax.grid(True, which="both", alpha=0.3) + return ax + + +def plot_energy_conservation( + energy: xr.Dataset, + ax: plt.Axes | None = None, + *, + log: bool = False, + title: str | None = None, +) -> plt.Axes: + """Plot field, kinetic, and total energy vs time as a conservation check. + + ``energy`` is the ``Dataset`` from :func:`io.load_hist_energy`; it must + contain a ``total`` variable (field + kinetic). The total-energy drift, + ``(max - min) / max(|total|)``, is annotated in the title. + """ + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + t = energy["t"].values + series_labels = { + "field_energy": "field", + "kinetic_total": "kinetic (all species)", + "total": "total (field + kinetic)", + } + for name, label in series_labels.items(): + if name in energy: + ax.plot(t, energy[name].values, label=label) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel("energy [code units]") + drift = energy.attrs.get("total_drift_frac") + base = title or "Energy conservation vs time" + ax.set_title(base if drift is None else f"{base} (total drift {drift:.2%})") + ax.legend() ax.grid(True, which="both", alpha=0.3) return ax @@ -250,7 +450,7 @@ def plot_omega_k( ax.axvline(0, color="w", lw=0.4, alpha=0.4) ax.set_xlabel(r"$k\ [\omega_p / c]$") ax.set_ylabel(r"$\omega\ [\omega_p]$") - ax.set_title(title or f"{da.name} omega-k spectrum") + ax.set_title(title or f"{_long_name(da)} — $(k, \\omega)$ power spectrum") return ax @@ -263,24 +463,30 @@ def save_canned_plots( *, v_th: float | None = None, dpi: int = 120, + n_panels: int = 8, ) -> dict[str, Path]: """Generate the standard set of PNGs for a finished OSIRIS run. - Returns a mapping of plot-name → output PNG path. + Returns a mapping of plot-name → output PNG path. Each diagnostic family is + best-effort: a failure on one diagnostic logs and is skipped rather than + aborting the rest. """ run_dir = Path(run_dir) out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) written: dict[str, Path] = {} + def _write(fig: plt.Figure, rel: str) -> Path: + p = out_dir / rel + p.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(p, bbox_inches="tight", dpi=dpi) + plt.close(fig) + return p + diags = _io.list_diagnostics(run_dir) - # Spacetime + omega-k for each field component (those live under FLD/). - fld_st = out_dir / "spacetime" - fld_st.mkdir(exist_ok=True) - omk_dir = out_dir / "omega_k" - omk_dir.mkdir(exist_ok=True) - for diag_rel, diag_path in diags.items(): + # --- Fields (FLD/, incl. currents j1-j3): spacetime, log, lineouts, ω-k --- + for diag_rel, diag_path in sorted(diags.items()): if not diag_rel.startswith("FLD/"): continue comp = diag_rel.split("/", 1)[1] @@ -290,60 +496,108 @@ def save_canned_plots( print(f"[plots] skipping series {diag_rel}: {e}") continue if ser.ndim != 2: - continue # higher-dim FLD plots skipped for now + continue # 2D-in-space field plots deferred + fig, ax = plt.subplots(figsize=(6, 4)) plot_spacetime(ser, ax=ax) - p = fld_st / f"{comp}.png" - fig.savefig(p, bbox_inches="tight", dpi=dpi) - plt.close(fig) - written[f"spacetime/{comp}"] = p + written[f"spacetime/{comp}"] = _write(fig, f"spacetime/{comp}.png") + + fig, ax = plt.subplots(figsize=(6, 4)) + plot_spacetime(ser, ax=ax, log=True) + written[f"spacetime_log/{comp}"] = _write(fig, f"spacetime_log/{comp}.png") + + written[f"lineouts/{comp}"] = _write( + plot_lineouts(ser, n_panels=n_panels), f"lineouts/{comp}.png" + ) fig, ax = plt.subplots(figsize=(6, 5)) - plot_omega_k( - ser, - ax=ax, - show_langmuir=v_th is not None, - v_th=v_th, + plot_omega_k(ser, ax=ax, show_langmuir=v_th is not None, v_th=v_th) + written[f"omega_k/{comp}"] = _write(fig, f"omega_k/{comp}.png") + + # --- Per-species moments (DENSITY//) --- + for diag_rel, diag_path in sorted(diags.items()): + if not diag_rel.startswith("DENSITY/"): + continue + parts = diag_rel.split("/") + if len(parts) < 3: + continue + species, quantity = parts[1], "/".join(parts[2:]) + try: + ser = _io.load_series(diag_path) + except Exception as e: + print(f"[plots] skipping moment {diag_rel}: {e}") + continue + if ser.ndim != 2: + continue + + fig, ax = plt.subplots(figsize=(6, 4)) + plot_spacetime(ser, ax=ax) + written[f"moments/{species}/{quantity}"] = _write( + fig, f"moments/{species}/{quantity}.png" + ) + + fig, ax = plt.subplots(figsize=(6, 4)) + plot_spacetime(ser, ax=ax, log=True) + written[f"moments/{species}/{quantity}_log"] = _write( + fig, f"moments/{species}/{quantity}_log.png" + ) + + written[f"moments/{species}/lineouts/{quantity}"] = _write( + plot_lineouts(ser, n_panels=n_panels), + f"moments/{species}/lineouts/{quantity}.png", ) - p = omk_dir / f"{comp}.png" - fig.savefig(p, bbox_inches="tight", dpi=dpi) - plt.close(fig) - written[f"omega_k/{comp}"] = p - # Final-step phase space per species. - ph_dir = out_dir / "phasespace" - for diag_rel, diag_path in diags.items(): + # --- Phase space (PHA//): final step + time evolution --- + for diag_rel, diag_path in sorted(diags.items()): if not diag_rel.startswith("PHA/"): continue - # PHA/x1p1/beam_pos -> ps_name="x1p1", species="beam_pos" parts = diag_rel.split("/") if len(parts) < 3: continue ps_name, species = parts[1], parts[2] last = _latest_h5(diag_path) - if last is None: - continue - sp_out = ph_dir / species - sp_out.mkdir(parents=True, exist_ok=True) - da = _io.load_phasespace_h5(last) - fig, ax = plt.subplots(figsize=(5, 4)) - plot_phasespace(da, ax=ax) - p = sp_out / f"{ps_name}.png" - fig.savefig(p, bbox_inches="tight", dpi=dpi) - plt.close(fig) - written[f"phasespace/{species}/{ps_name}"] = p + if last is not None: + fig, ax = plt.subplots(figsize=(5, 4)) + plot_phasespace(_io.load_phasespace_h5(last), ax=ax) + written[f"phasespace/{species}/{ps_name}"] = _write( + fig, f"phasespace/{species}/{ps_name}.png" + ) + try: + ser = _io.load_series(diag_path) + if ser.ndim == 3: + written[f"phasespace_evolution/{species}/{ps_name}"] = _write( + plot_phasespace_evolution(ser, n_panels=n_panels), + f"phasespace_evolution/{species}/{ps_name}.png", + ) + except Exception as e: + print(f"[plots] skipping phasespace evolution {diag_rel}: {e}") - # Energy vs time (uses any/all FLD components present). + # --- Energy diagnostics --- try: fig, ax = plt.subplots(figsize=(6, 4)) plot_energy_vs_time(run_dir, ax=ax) - p = out_dir / "energy_vs_time.png" - fig.savefig(p, bbox_inches="tight", dpi=dpi) - plt.close(fig) - written["energy_vs_time"] = p + written["energy_vs_time"] = _write(fig, "energy_vs_time.png") except (FileNotFoundError, RuntimeError) as e: print(f"[plots] skipping energy_vs_time: {e}") + try: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_energy_components(run_dir, ax=ax) + written["energy_components_vs_time"] = _write( + fig, "energy_components_vs_time.png" + ) + except (FileNotFoundError, RuntimeError) as e: + print(f"[plots] skipping energy_components_vs_time: {e}") + + try: + energy = _io.load_hist_energy(run_dir) + if energy is not None and "total" in energy: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_energy_conservation(energy, ax=ax) + written["total_energy_vs_time"] = _write(fig, "total_energy_vs_time.png") + except Exception as e: + print(f"[plots] skipping total_energy_vs_time: {e}") + return written @@ -373,6 +627,53 @@ def _axis_label(da: xr.DataArray, dim: str) -> str: return long +def _long_name(da: xr.DataArray) -> str: + """Human-readable label for a diagnostic: its OSIRIS LABEL, else its name.""" + return str(da.attrs.get("long_name") or da.name or "") + + +def _value_label(da: xr.DataArray) -> str: + """Quantity label with units (for colorbars / y-axes).""" + name = _long_name(da) + units = da.attrs.get("units", "") + return rf"{name} [${units}$]" if units else name + + +def _decorate(da: xr.DataArray) -> xr.DataArray: + """Lift OSIRIS per-axis units / long-names onto coordinate attrs. + + ``io.load_series`` stashes axis metadata in dict-valued DataArray attrs; + copying it onto the coordinates lets xarray's faceted plotters auto-label + each panel's axes. Returns a decorated copy (the input is left untouched). + """ + da = da.copy() + axis_units = da.attrs.get("axis_units", {}) or {} + axis_long = da.attrs.get("axis_long_names", {}) or {} + for dim in da.dims: + if dim not in da.coords: + continue + if dim in axis_long: + da.coords[dim].attrs.setdefault("long_name", axis_long[dim]) + if dim in axis_units: + da.coords[dim].attrs.setdefault("units", axis_units[dim]) + if "t" in da.coords: + da.coords["t"].attrs.setdefault("long_name", "t") + da.coords["t"].attrs.setdefault("units", da.attrs.get("time_units", r"1/\omega_p")) + return da + + +def _cell_volume(da: xr.DataArray) -> float: + """Product of uniform coordinate spacings over the spatial dims.""" + dvol = 1.0 + for dim in da.dims: + if dim == "t": + continue + c = da.coords[dim].values + if c.size > 1: + dvol *= float(c[1] - c[0]) + return dvol + + def _latest_h5(diag_dir: Path) -> Path | None: best: tuple[int, Path] | None = None for p in diag_dir.iterdir(): diff --git a/adept/osiris/post.py b/adept/osiris/post.py index a9a73cac..0d5b92a4 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -128,6 +128,31 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: if ms.is_dir(): _io.save_run_datasets(run_dir, td / "binary", diagnostics=whitelist) + # plots imports matplotlib; do it lazily to keep `import adept.osiris` light. + from adept.osiris import plots as _plots + + # E/B-field split + energy-conservation metrics (best effort). + try: + comps = _plots.field_energy_components(run_dir) + metrics["efield_energy_final"] = float(comps["E_energy"].values[-1]) + metrics["bfield_energy_final"] = float(comps["B_energy"].values[-1]) + except (FileNotFoundError, RuntimeError) as e: + print(f"[post] field-energy components unavailable: {e}") + try: + energy = _io.load_hist_energy(run_dir) + if energy is not None and "total_drift_frac" in energy.attrs: + metrics["energy_drift_frac"] = float(energy.attrs["total_drift_frac"]) + except Exception as e: + print(f"[post] HIST energy unavailable: {e}") + + # Render the standard plot set into td/plots so MLflow logs them as + # artifacts. Never let a plotting failure abort metric/artifact logging. + try: + v_th = (cfg.get("output") or {}).get("v_th") + _plots.save_canned_plots(run_dir, td / "plots", v_th=v_th) + except Exception as e: + print(f"[post] plotting failed: {e}") + # Always include the rendered deck + stdout/stderr. for fname in ("os-stdin", "stdout.log", "stderr.log"): src = run_dir / fname diff --git a/tests/test_osiris/test_diagnostics_plots.py b/tests/test_osiris/test_diagnostics_plots.py new file mode 100644 index 00000000..cb469e09 --- /dev/null +++ b/tests/test_osiris/test_diagnostics_plots.py @@ -0,0 +1,166 @@ +"""Self-contained tests for the expanded OSIRIS diagnostics / plots. + +These synthesize a tiny OSIRIS-shaped run (fields, per-species density moment, +phase space, and HIST energy histories) so they need no real run on disk, and +exercise the additions in plots.py / io.py / post.py: + + - log-scale spacetime + lineout field plots + - per-species DENSITY moment plots + - phase-space time-evolution panels + - E/B-field energy split and HIST-based energy conservation + - post.collect wiring (plots land under td/plots) +""" + +from __future__ import annotations + +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import h5py +import numpy as np + +from adept.osiris import io as oio +from adept.osiris import plots as oplt +from adept.osiris import post as opost + + +def _write_dump(path: Path, name: str, data: np.ndarray, t: float, it: int, axes) -> None: + """Write one OSIRIS-style HDF5 grid/phasespace dump. + + ``axes`` is a list of ``(name, long_name, units, min, max)`` in OSIRIS + (Fortran) order; the last entry is the fastest-varying numpy axis. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + f.attrs["NAME"] = np.array([name.encode()], dtype="S256") + f.attrs["LABEL"] = np.array([name.encode()], dtype="S256") + f.attrs["UNITS"] = np.array([b"a.u."], dtype="S256") + f.attrs["TIME UNITS"] = np.array([rb"1 / \omega_p"], dtype="S256") + ax = f.create_group("AXIS") + for i, (an, ln, un, lo, hi) in enumerate(axes, start=1): + d = ax.create_dataset(f"AXIS{i}", data=np.array([lo, hi], dtype="float64")) + d.attrs["NAME"] = np.array([an.encode()], dtype="S256") + d.attrs["LONG_NAME"] = np.array([ln.encode()], dtype="S256") + d.attrs["UNITS"] = np.array([un.encode()], dtype="S256") + sim = f.create_group("SIMULATION") + sim.attrs["DT"] = np.array([0.05]) + sim.attrs["NDIMS"] = np.array([1], dtype="int32") + sim.attrs["NX"] = np.array([data.shape[-1]], dtype="int32") + sim.attrs["XMIN"] = np.array([axes[-1][3]]) + sim.attrs["XMAX"] = np.array([axes[-1][4]]) + f.create_dataset(name, data=data.astype("float32")) + + +def _make_full_run(root: Path, n_steps: int = 6, nx: int = 16, npx: int = 12) -> Path: + """Build a run with FLD/e1, DENSITY/electron/charge, PHA, and HIST energy.""" + run_dir = root / "run" + rng = np.random.default_rng(0) + x_ax = ("x1", "x_1", r"c / \omega_p", 0.0, 10.0) + p_ax = ("p1", "p_1", r"m_e c", -2.0, 2.0) + + for k in range(n_steps): + it = k * 10 + t = k * 0.5 + _write_dump(run_dir / "MS/FLD/e1" / f"e1-{it:06d}.h5", + "e1", rng.standard_normal(nx), t=t, it=it, axes=[x_ax]) + _write_dump(run_dir / "MS/DENSITY/electron/charge" / f"charge-electron-{it:06d}.h5", + "charge", rng.standard_normal(nx), t=t, it=it, axes=[x_ax]) + _write_dump(run_dir / "MS/PHA/x1p1/electron" / f"x1p1-electron-{it:06d}.h5", + "x1p1", rng.random((npx, nx)), t=t, it=it, axes=[x_ax, p_ax]) + + # HIST energy histories: iter, time, then value columns. + hist = run_dir / "HIST" + hist.mkdir(parents=True, exist_ok=True) + times = [k * 0.5 for k in range(n_steps)] + fld = ["! iter time e1 e2 e3 b1 b2 b3"] + par = ["! iter time ene"] + for k, t in enumerate(times): + fld.append(f"{k*10} {t} {1.0 + 0.1*k} 0.0 0.0 0.0 0.0 0.0") + par.append(f"{k*10} {t} {100.0 - 0.1*k}") + (hist / "fld_ene").write_text("\n".join(fld) + "\n") + (hist / "par01_ene").write_text("\n".join(par) + "\n") + + (run_dir / "os-stdin").write_text("node_conf\n{\n}\n") + (run_dir / "stdout.log").write_text("ok\n") + (run_dir / "stderr.log").write_text("") + return run_dir + + +def test_field_energy_components_splits_e_and_b(tmp_path: Path) -> None: + run_dir = _make_full_run(tmp_path) + ds = oplt.field_energy_components(run_dir) + assert set(ds.data_vars) == {"E_energy", "B_energy", "total_field_energy"} + # Only e1 was dumped, so all energy is electric and B is identically zero. + assert np.all(ds["E_energy"].values > 0) + assert np.all(ds["B_energy"].values == 0) + np.testing.assert_allclose( + ds["total_field_energy"].values, ds["E_energy"].values + ) + + +def test_load_hist_energy_builds_conservation_total(tmp_path: Path) -> None: + run_dir = _make_full_run(tmp_path, n_steps=5) + energy = oio.load_hist_energy(run_dir) + assert energy is not None + assert "field_energy" in energy + assert "kinetic_par01_ene" in energy + assert "kinetic_total" in energy + assert "total" in energy + # total = field + kinetic, elementwise. + np.testing.assert_allclose( + energy["total"].values, + energy["field_energy"].values + energy["kinetic_total"].values, + ) + assert "total_drift_frac" in energy.attrs + assert 0.0 <= energy.attrs["total_drift_frac"] <= 1.0 + + +def test_load_hist_energy_absent_returns_none(tmp_path: Path) -> None: + # A run dir with no HIST/ yields None rather than raising. + (tmp_path / "run" / "MS").mkdir(parents=True) + assert oio.load_hist_energy(tmp_path / "run") is None + + +def test_save_canned_plots_emits_full_set(tmp_path: Path) -> None: + run_dir = _make_full_run(tmp_path) + out = tmp_path / "plots" + written = oplt.save_canned_plots(run_dir, out, v_th=0.1) + + expected = { + "spacetime/e1", + "spacetime_log/e1", + "lineouts/e1", + "omega_k/e1", + "moments/electron/charge", + "moments/electron/charge_log", + "moments/electron/lineouts/charge", + "phasespace/electron/x1p1", + "phasespace_evolution/electron/x1p1", + "energy_vs_time", + "energy_components_vs_time", + "total_energy_vs_time", # present because HIST/ supplies kinetic energy + } + assert expected <= set(written) + for path in written.values(): + assert path.exists() and path.stat().st_size > 0 + + +def test_collect_writes_plots_under_td(tmp_path: Path) -> None: + run_dir = _make_full_run(tmp_path) + run_output = {"solver result": {"run_dir": str(run_dir), "wall_time": 1.0, "exit_code": 0}} + cfg = {"osiris": {"deck": str(run_dir / "os-stdin")}, "output": {}} + + td = tmp_path / "td" + td.mkdir() + result = opost.collect(run_output, cfg, str(td)) + + # Plots were generated as artifacts and energy metrics added. + assert (td / "plots" / "spacetime" / "e1.png").exists() + assert (td / "plots" / "energy_components_vs_time.png").exists() + assert "efield_energy_final" in result["metrics"] + assert "energy_drift_frac" in result["metrics"] From 9f818e8ec37e11a7142dc2f13dca97f51ce72a6d Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 29 May 2026 17:09:13 -0700 Subject: [PATCH 07/49] Fixed issue when multiple datasets were in an h5 file --- adept/osiris/io.py | 149 +++++++++++++++++++++++++- tests/test_osiris/test_post_netcdf.py | 123 +++++++++++++++++++++ 2 files changed, 267 insertions(+), 5 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 5be574ea..304dfb11 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -124,6 +124,137 @@ def load_phasespace_h5(path: str | Path) -> xr.DataArray: return load_grid_h5(path) +def _is_raw_h5(f: h5py.File) -> bool: + """Heuristic: is this an OSIRIS RAW (particle) dump rather than a grid dump? + + RAW dumps hold several 1-D per-particle datasets and have no ``AXIS`` + group, whereas grid / phase-space dumps hold a single gridded dataset plus + an ``AXIS`` group. Treat anything with more than one data dataset, or no + ``AXIS`` group, as non-grid. + """ + data_keys = [k for k in f.keys() if k not in ("AXIS", "SIMULATION")] + return len(data_keys) != 1 or "AXIS" not in f + + +def load_raw_h5(path: str | Path) -> xr.Dataset: + """Load one OSIRIS RAW (particle) dump into an ``xarray.Dataset``. + + Unlike field / charge / phase-space *grid* dumps (one gridded dataset plus + an ``AXIS`` group), a RAW dump holds several 1-D per-particle datasets + (``x1``, ``p1``, ``p2``, ``p3``, ``ene``, ``q``, ...), a ``SIMULATION`` + attrs group, and root attrs ``TIME`` / ``ITER`` — but typically no + ``AXIS`` group, because the data is not gridded. + + Dataset names are discovered dynamically (different decks dump different + quantities); each becomes a data variable indexed by a single particle + dimension ``"pidx"``. Per-quantity ``UNITS`` / ``LONG_NAME`` attrs, when + present, ride on the matching variable; ``TIME`` / ``ITER`` ride on the + Dataset. + """ + path = Path(path) + with h5py.File(path, "r") as f: + data_keys = sorted( + k + for k in f.keys() + if k not in ("AXIS", "SIMULATION") and isinstance(f[k], h5py.Dataset) + ) + data_vars: dict[str, tuple] = {} + for name in data_keys: + dset = f[name] + arr = dset[...].astype("float64").reshape(-1) + var_attrs = {} + if "UNITS" in dset.attrs: + var_attrs["units"] = _decode(dset.attrs["UNITS"]) + if "LONG_NAME" in dset.attrs: + var_attrs["long_name"] = _decode(dset.attrs["LONG_NAME"]) + data_vars[name] = ("pidx", arr, var_attrs) + + npart = max((v[1].shape[0] for v in data_vars.values()), default=0) + + attrs = { + "time": float(f.attrs["TIME"][0]) if "TIME" in f.attrs else float("nan"), + "iter": int(f.attrs["ITER"][0]) + if "ITER" in f.attrs + else _iter_from_name(path), + "long_name": _decode(f.attrs.get("LABEL", path.stem)), + "time_units": _decode(f.attrs.get("TIME UNITS", "")), + "source": str(path), + "n_particles": int(npart), + } + if "SIMULATION" in f: + sim = f["SIMULATION"].attrs + for key in ("DT", "NDIMS", "XMIN", "XMAX", "NX", "PERIODIC"): + if key in sim: + val = sim[key] + if hasattr(val, "tolist"): + val = val.tolist() + attrs[f"sim.{key}"] = val + + return xr.Dataset(data_vars, attrs=attrs) + + +def load_raw_series(directory: str | Path) -> xr.Dataset: + """Concatenate every RAW (particle) dump in ``directory`` long-form. + + RAW dumps have a *variable* particle count per timestep (OSIRIS samples a + ``raw_fraction`` of particles each dump), so they cannot be stacked into a + rectangular ``(t, particle)`` array the way grid dumps are. Instead every + dump is concatenated along the particle dimension ``"pidx"`` with a per-row + ``t`` / ``iter`` coordinate identifying which dump each particle came from. + The union of quantities across dumps is preserved (missing quantities fill + with NaN for that dump's rows). + """ + directory = Path(directory) + dumps = _sort_dumps(directory) + if not dumps: + raise FileNotFoundError(f"No .h5 dumps in {directory}") + + per_dump: list[xr.Dataset] = [] + times: list[np.ndarray] = [] + iters: list[np.ndarray] = [] + for p in dumps: + ds = load_raw_h5(p) + n = ds.sizes.get("pidx", 0) + per_dump.append(ds) + times.append(np.full(n, ds.attrs["time"], dtype="float64")) + iters.append(np.full(n, ds.attrs["iter"], dtype="int64")) + + combined = xr.concat(per_dump, dim="pidx", data_vars="all", coords="minimal") + combined = combined.assign_coords( + t=("pidx", np.concatenate(times) if times else np.empty(0)), + iter=( + "pidx", + np.concatenate(iters) if iters else np.empty(0, dtype="int64"), + ), + ) + attrs = dict(per_dump[0].attrs) + for k in ("time", "iter", "source", "n_particles"): + attrs.pop(k, None) + attrs["source_dir"] = str(directory) + attrs["n_dumps"] = len(dumps) + combined.attrs = attrs + return combined + + +def _diag_is_raw(relpath: str, directory: str | Path) -> bool: + """Detect a RAW (particle) diagnostic. + + Primary signal: the diagnostic relpath starts with ``"RAW/"``. As a + defensive fallback, peek at the first dump and treat it as RAW when it + fails the grid heuristic (more than one data dataset, or no ``AXIS``). + """ + if relpath.startswith("RAW/") or Path(relpath).parts[0] == "RAW": + return True + dumps = _sort_dumps(Path(directory)) + if not dumps: + return False + try: + with h5py.File(dumps[0], "r") as f: + return _is_raw_h5(f) + except Exception: + return False + + def _sort_dumps(directory: Path) -> list[Path]: files = [p for p in directory.iterdir() if p.is_file() and p.suffix == ".h5"] return sorted(files, key=_iter_from_name) @@ -252,11 +383,19 @@ def save_run_datasets( relpath not in diagnostics and Path(relpath).name not in diagnostics ): continue - ds = series_to_dataset(load_series(diags[relpath])) - dest = out_dir / f"{relpath}.nc" - dest.parent.mkdir(parents=True, exist_ok=True) - ds.to_netcdf(dest, engine="h5netcdf") - written.append(dest) + try: + if _diag_is_raw(relpath, diags[relpath]): + # RAW (particle) dumps: per-particle datasets, no grid/AXIS. + ds: xr.Dataset = load_raw_series(diags[relpath]) + else: + ds = series_to_dataset(load_series(diags[relpath])) + dest = out_dir / f"{relpath}.nc" + dest.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(dest, engine="h5netcdf") + written.append(dest) + except Exception as e: # one bad diagnostic must not abort the rest + print(f"[post] skipping diagnostic {relpath}: {e}") + continue return written diff --git a/tests/test_osiris/test_post_netcdf.py b/tests/test_osiris/test_post_netcdf.py index 81bb7527..e2e489a2 100644 --- a/tests/test_osiris/test_post_netcdf.py +++ b/tests/test_osiris/test_post_netcdf.py @@ -124,3 +124,126 @@ def test_collect_respects_diagnostics_whitelist(tmp_path: Path) -> None: assert (td / "binary" / "FLD" / "e1.nc").exists() assert not (td / "binary" / "FLD" / "e2.nc").exists() + + +def _write_raw_dump( + path: Path, quantities: dict[str, np.ndarray], t: float, it: int +) -> None: + """Write one OSIRIS-style RAW (particle) HDF5 dump. + + ``quantities`` maps a per-particle quantity name (``"p1"``, ``"x1"``, ...) + to its 1-D array. RAW dumps have NO ``AXIS`` group; each quantity is a + top-level 1-D dataset, all the same length, plus TIME / ITER root attrs and + a SIMULATION group. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + f.attrs["NAME"] = np.array([b"RAW"], dtype="S256") + f.attrs["LABEL"] = np.array([b"RAW"], dtype="S256") + f.attrs["TIME UNITS"] = np.array([rb"1 / \omega_p"], dtype="S256") + f.attrs["TYPE"] = np.array([b"particles"], dtype="S16") + for qn, arr in quantities.items(): + d = f.create_dataset(qn, data=arr.astype("float32")) + d.attrs["UNITS"] = np.array([b"a.u."], dtype="S256") + d.attrs["LONG_NAME"] = np.array([qn.encode()], dtype="S256") + sim = f.create_group("SIMULATION") + sim.attrs["DT"] = np.array([0.05]) + sim.attrs["NDIMS"] = np.array([1], dtype="int32") + + +def test_load_raw_h5_returns_particle_dataset(tmp_path: Path) -> None: + p = tmp_path / "MS" / "RAW" / "species_1" / "RAW-species_1-000030.h5" + npart = 7 + quants = { + q: np.arange(npart, dtype="float64") for q in ("ene", "p1", "p2", "p3", "q", "x1") + } + _write_raw_dump(p, quants, t=1.5, it=30) + + ds = oio.load_raw_h5(p) + assert isinstance(ds, xr.Dataset) + assert set(ds.dims) == {"pidx"} + assert ds.sizes["pidx"] == npart + assert set(ds.data_vars) == set(quants) + assert ds.attrs["iter"] == 30 + assert ds.attrs["time"] == 1.5 + assert ds["p1"].attrs["units"] == "a.u." + + +def test_load_grid_h5_still_works_for_grid_dumps(tmp_path: Path) -> None: + # A normal grid dump (single dataset + AXIS) must still load as a DataArray. + p = tmp_path / "e1-000000.h5" + _write_dump( + p, + "e1", + np.arange(8, dtype="float64"), + t=0.0, + it=0, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + da = oio.load_grid_h5(p) + assert isinstance(da, xr.DataArray) + assert da.dims == ("x1",) + assert da.sizes["x1"] == 8 + + +def test_load_raw_series_handles_variable_particle_counts(tmp_path: Path) -> None: + raw = tmp_path / "MS" / "RAW" / "species_1" + # Two dumps with DIFFERENT particle counts (raw_fraction sampling). + _write_raw_dump( + raw / "RAW-species_1-000000.h5", + {q: np.zeros(5) for q in ("p1", "x1")}, + t=0.0, + it=0, + ) + _write_raw_dump( + raw / "RAW-species_1-000100.h5", + {q: np.ones(9) for q in ("p1", "x1")}, + t=5.0, + it=100, + ) + + ds = oio.load_raw_series(raw) + assert ds.sizes["pidx"] == 14 # 5 + 9 — no equal-shape assumption + assert sorted(set(ds["iter"].values.tolist())) == [0, 100] + assert (ds["t"].values[:5] == 0.0).all() + assert (ds["t"].values[5:] == 5.0).all() + + +def test_save_run_datasets_routes_raw_to_netcdf(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=2, nx=8) # FLD/e1 grid diagnostic + raw = run_dir / "MS" / "RAW" / "species_1" + _write_raw_dump( + raw / "RAW-species_1-000000.h5", + {q: np.arange(4, dtype="float64") for q in ("ene", "p1", "x1", "q")}, + t=0.0, + it=0, + ) + out = tmp_path / "binary" + oio.save_run_datasets(run_dir, out) + + # Both the grid and the RAW diagnostic produced netCDF — no crash on RAW. + assert (out / "FLD" / "e1.nc").exists() + assert (out / "RAW" / "species_1.nc").exists() + ds = xr.open_dataset(out / "RAW" / "species_1.nc", engine="h5netcdf") + try: + assert "pidx" in ds.dims + assert "p1" in ds.data_vars + finally: + ds.close() + + +def test_save_run_datasets_skips_unloadable_diagnostic(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=2, nx=8) # good FLD/e1 + # A corrupt diagnostic: a .h5 file that isn't valid HDF5. + bad = run_dir / "MS" / "FLD" / "bad" + bad.mkdir(parents=True) + (bad / "bad-000000.h5").write_bytes(b"not an hdf5 file") + + out = tmp_path / "binary" + written = oio.save_run_datasets(run_dir, out) # must not raise + + assert (out / "FLD" / "e1.nc").exists() # good diagnostic still produced + assert not (out / "FLD" / "bad.nc").exists() + assert (out / "FLD" / "e1.nc") in written From cb78418bc4cbc1c61b79a4cb0142016928709645 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 1 Jun 2026 16:28:37 -0700 Subject: [PATCH 08/49] More plots added: zoomed omega-k, current, averaged distribution functions (should probably be moved to another branch). Added density and temperature profiles. Fixed LaTeX issues --- adept/osiris/plots.py | 530 +++++++++++++++++++++- adept/osiris/post.py | 9 +- configs/osiris/twostream-1d.yaml | 3 + docs/osiris-adept-usage.md | 28 ++ tests/test_osiris/test_plots_new_views.py | 243 ++++++++++ 5 files changed, 788 insertions(+), 25 deletions(-) create mode 100644 tests/test_osiris/test_plots_new_views.py diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index d2fc079c..fa678929 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -1,4 +1,4 @@ -"""Canned matplotlib views over OSIRIS diagnostics. +r"""Canned matplotlib views over OSIRIS diagnostics. Each plotter accepts either an already-loaded ``xarray.DataArray`` or a filesystem path / directory, and returns the ``Axes`` it drew on so the @@ -11,15 +11,25 @@ spacetime/.png (t, x) heatmap of every FLD diagnostic spacetime_log/.png log10|·| of the same lineouts/.png value-vs-x snapshots at sampled times + omega_k/.png 2-D FFT (k, ω) dispersion plot + omega_k_zoom/.png same, zoomed to the ω = k line (plasma waves) + currents/spacetime.png j1/j2/j3 (J_x/J_y/J_z) side-by-side spacetime + currents/lineouts.png j1/j2/j3 profiles vs x (final + late mean) moments//.png (t, x) per-species density moments moments//_log.png log10 of the same moments//lineouts/.png moment snapshots at sampled times + profiles//density.png density profile vs x (final + late mean) + profiles//temperature.png temperature profile vs x (if thermal moments) phasespace//.png final-step (x, p) heatmap per species phasespace_evolution//.png (x, p) heatmaps at sampled times - omega_k/.png 2-D FFT (k, ω) dispersion plot + distribution_lineouts//.png f(p) averaged over the right-boundary cells + field_decomp/.png left/right-going transverse E (Riemann split) energy_vs_time.png total field energy time-trace energy_components_vs_time.png E-field / B-field / total energy total_energy_vs_time.png field + kinetic conservation (needs HIST/) + +All axis / colorbar / title labels are emitted as proper LaTeX (``$\omega$``, +``$c/\omega_p$``, …) via the ``_tex`` helper. """ from __future__ import annotations @@ -69,12 +79,12 @@ def plot_spacetime( plt.colorbar( mesh, ax=ax, - label=f"log10 |{_value_label(da)}|" if log else _value_label(da), + label=rf"$\log_{{10}}$ |{_value_label(da)}|" if log else _value_label(da), ) ax.set_xlabel(_axis_label(da, "t")) ax.set_ylabel(_axis_label(da, xname)) - scale = "log10 " if log else "" - ax.set_title(title or f"{_long_name(da)} — {scale}spacetime (x vs t)") + scale = r"$\log_{10}$ " if log else "" + ax.set_title(title or f"{_display_name(da)} — {scale}spacetime ($x$ vs $t$)") return ax @@ -105,7 +115,7 @@ def plot_lineouts( g.set_xlabels(_axis_label(da, xname)) g.set_ylabels(_value_label(da)) g.fig.suptitle( - title or f"{_long_name(da)} — lineouts vs x at sampled times", y=1.02 + title or rf"{_display_name(da)} — lineouts vs $x$ at sampled times", y=1.02 ) return g.fig @@ -134,13 +144,16 @@ def plot_phasespace( x0, x1 = da.coords[d0].values, da.coords[d1].values mesh = ax.pcolormesh(x0, x1, arr.T, shading="auto", cmap=cmap) plt.colorbar( - mesh, ax=ax, label=f"log10 {_value_label(da)}" if log else _value_label(da) + mesh, ax=ax, label=rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da) ) ax.set_xlabel(_axis_label(da, d0)) ax.set_ylabel(_axis_label(da, d1)) t = da.attrs.get("time", float("nan")) - scale = "log10 " if log else "" - ax.set_title(title or f"{_long_name(da)} — {scale}phase space (t = {t:.3g} 1/$\\omega_p$)") + scale = r"$\log_{10}$ " if log else "" + ax.set_title( + title + or rf"{_display_name(da)} — {scale}phase space ($t = {t:.3g}\ 1/\omega_p$)" + ) return ax @@ -179,9 +192,9 @@ def plot_phasespace_evolution( if spatial and moment: g.set_xlabels(_axis_label(da, spatial[0])) g.set_ylabels(_axis_label(da, moment[0])) - scale = "log10 " if log else "" + scale = r"$\log_{10}$ " if log else "" g.fig.suptitle( - title or f"{_long_name(da)} — {scale}phase space at sampled times", y=1.02 + title or rf"{_display_name(da)} — {scale}phase space at sampled times", y=1.02 ) return g.fig @@ -381,6 +394,7 @@ def plot_omega_k( cmap: str = "magma", show_em: bool = True, show_langmuir: bool = False, + show_light_line: bool = False, v_th: float | None = None, omega_p: float = 1.0, k_max: float | None = None, @@ -392,7 +406,10 @@ def plot_omega_k( The OSIRIS convention has ω_p = 1, c = 1 in code units, so the relativistic-EM dispersion ω² = ω_p² + k² and the Langmuir Bohm–Gross dispersion ω² = ω_p² + 3 k² v_th² overlay directly when - you ask for them. + you ask for them. ``show_light_line`` overlays the vacuum light line + ω = ±k (the ω = k diagonal that the EM branch asymptotes to), which is + the useful guide when ``k_max`` / ``omega_max`` are set to zoom into the + low-(k, ω) region where the plasma (Langmuir) waves live. """ da = _ensure_series(series) if da.ndim != 2: @@ -420,7 +437,9 @@ def plot_omega_k( k = np.fft.fftshift(np.fft.fftfreq(nx, d=dx)) * 2 * np.pi mesh = ax.pcolormesh(k, omega, P, shading="auto", cmap=cmap) - plt.colorbar(mesh, ax=ax, label=("log10 |F|^2" if log else "|F|^2")) + plt.colorbar( + mesh, ax=ax, label=(r"$\log_{10}\,|\tilde{F}|^2$" if log else r"$|\tilde{F}|^2$") + ) if k_max is None: k_max = float(np.max(np.abs(k))) @@ -429,8 +448,12 @@ def plot_omega_k( ax.set_xlim(-k_max, k_max) ax.set_ylim(-omega_max, omega_max) - # Overlay analytical dispersion lines. + # Overlay analytical dispersion lines (sampled across the visible k range). k_line = np.linspace(-k_max, k_max, 401) + if show_light_line: + ax.plot(k_line, +k_line, "w-", lw=0.8, alpha=0.6, + label=r"light line: $\omega = \pm k$") + ax.plot(k_line, -k_line, "w-", lw=0.8, alpha=0.6) if show_em: w_em = np.sqrt(omega_p ** 2 + k_line ** 2) ax.plot(k_line, +w_em, "w--", lw=1, alpha=0.7, @@ -443,20 +466,390 @@ def plot_omega_k( ax.plot(k_line, +w_l, "c:", lw=1, alpha=0.8, label=r"Langmuir: $\omega^2 = \omega_p^2 + 3 k^2 v_{th}^2$") ax.plot(k_line, -w_l, "c:", lw=1, alpha=0.8) - if show_em or show_langmuir: + if show_em or show_langmuir or show_light_line: ax.legend(loc="upper right", fontsize=8, framealpha=0.6) ax.axhline(0, color="w", lw=0.4, alpha=0.4) ax.axvline(0, color="w", lw=0.4, alpha=0.4) ax.set_xlabel(r"$k\ [\omega_p / c]$") ax.set_ylabel(r"$\omega\ [\omega_p]$") - ax.set_title(title or f"{_long_name(da)} — $(k, \\omega)$ power spectrum") + ax.set_title(title or rf"{_display_name(da)} — $(k, \omega)$ power spectrum") + return ax + + +def plot_distribution_lineout( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + n_cells: int = 10, + n_times: int = 6, + log: bool = False, + cmap: str = "viridis", + title: str | None = None, +) -> plt.Axes: + """Velocity distribution ``f(p)`` averaged over the rightmost ``n_cells``. + + Takes a phase-space series with a spatial (``x*``) and a momentum + (``p*``) axis — either a stacked ``(t, p, x)`` series from + :func:`io.load_series` or a single ``(p, x)`` snapshot — averages it over + the ``n_cells`` cells nearest the right-hand boundary to get ``f(p)`` + there, and overlays that lineout at ~``n_times`` sampled times (colour = + time). This is the standard "what does the distribution look like as it + leaves the box" diagnostic for driven / open-boundary runs. + """ + da = series if isinstance(series, xr.DataArray) else _io.load_series(series) + da = _decorate(da) + spatial = [d for d in da.dims if d != "t" and str(d).startswith("x")] + moment = [d for d in da.dims if d != "t" and str(d).startswith("p")] + if not spatial or not moment: + raise ValueError( + f"plot_distribution_lineout needs an x* and a p* dim; got {da.dims}" + ) + xdim, pdim = spatial[0], moment[0] + nx = da.sizes[xdim] + k = max(1, min(n_cells, nx)) + # Average over the k rightmost spatial cells -> f over the momentum axis. + f_xp = da.isel({xdim: slice(nx - k, nx)}).mean(dim=xdim) + p = f_xp.coords[pdim].values + xv = da.coords[xdim].values + x_lo, x_hi = float(xv[nx - k]), float(xv[-1]) + + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + + if "t" in f_xp.dims: + nt = f_xp.sizes["t"] + idx = np.unique(np.linspace(0, nt - 1, min(n_times, nt)).astype(int)) + colours = plt.get_cmap(cmap)(np.linspace(0.15, 0.95, len(idx))) + tvals = f_xp.coords["t"].values + for colour, i in zip(colours, idx, strict=False): + y = f_xp.isel(t=i).values + y = np.abs(y) if log else y + ax.plot(p, y, color=colour, lw=1.3, label=f"{float(tvals[i]):.2g}") + ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) + else: + y = np.abs(f_xp.values) if log else f_xp.values + ax.plot(p, y, lw=1.4) + + if log: + ax.set_yscale("log") + units = da.attrs.get("units", "") + ylabel = rf"$\langle f \rangle_x$ [{_tex(units)}]" if units else r"$\langle f \rangle_x$" + ax.set_xlabel(_axis_label(da, pdim)) + ax.set_ylabel(ylabel) + ax.set_title( + title + or rf"{_display_name(da)} — $f$ over rightmost {k} cells " + rf"($x \in [{x_lo:.3g}, {x_hi:.3g}]$)" + ) + ax.grid(True, alpha=0.3) return ax +# --- currents (j1/j2/j3) -------------------------------------------------- + + +def _current_components(run_dir: str | Path) -> dict[str, xr.DataArray]: + """Load whichever of ``FLD/j1``, ``FLD/j2``, ``FLD/j3`` were dumped.""" + diags = _io.list_diagnostics(run_dir) + out: dict[str, xr.DataArray] = {} + for comp in ("j1", "j2", "j3"): + rel = f"FLD/{comp}" + if rel in diags: + try: + ser = _io.load_series(diags[rel]) + except Exception as e: # skip a bad component + print(f"[plots] skipping current {comp}: {e}") + continue + if ser.ndim == 2: + out[comp] = ser + return out + + +def plot_currents_spacetime(run_dir: str | Path) -> plt.Figure | None: + """Side-by-side ``(t, x)`` heatmaps of the current components present. + + OSIRIS dumps current density under ``MS/FLD/j1`` … ``j3`` (the labels map + to ``J_x``, ``J_y``, ``J_z`` for a run aligned with ``x1``). This stacks + whatever is present into one figure for an at-a-glance comparison; + returns ``None`` if no current was dumped. + """ + comps = _current_components(run_dir) + if not comps: + return None + fig, axes = plt.subplots(1, len(comps), figsize=(5 * len(comps), 4), squeeze=False) + for ax, (comp, ser) in zip(axes[0], comps.items(), strict=False): + plot_spacetime(ser, ax=ax) + fig.suptitle(r"Current density components — spacetime ($x$ vs $t$)", y=1.03) + fig.tight_layout() + return fig + + +def plot_currents_lineouts(run_dir: str | Path, *, n_avg_frac: float = 0.2) -> plt.Figure | None: + """Overlay the current components vs ``x`` (final snapshot + late-time mean). + + All available ``j1/j2/j3`` are drawn on a single axis so their relative + magnitudes and spatial structure are directly comparable. Solid = final + dump, dashed = mean over the last ``n_avg_frac`` of the run. + """ + comps = _current_components(run_dir) + if not comps: + return None + fig, ax = plt.subplots(figsize=(6, 4)) + for comp, ser in comps.items(): + da = _decorate(ser) + xdim = next(d for d in da.dims if d != "t") + x = da.coords[xdim].values + nt = da.sizes["t"] + w = max(1, round(n_avg_frac * nt)) + (line,) = ax.plot(x, da.isel(t=-1).values, lw=1.4, label=_tex(_long_name(da))) + ax.plot( + x, da.isel(t=slice(nt - w, nt)).mean("t").values, + lw=1.0, ls="--", alpha=0.7, color=line.get_color(), + ) + any_ser = _decorate(next(iter(comps.values()))) + xdim = next(d for d in any_ser.dims if d != "t") + ax.set_xlabel(_axis_label(any_ser, xdim)) + ax.set_ylabel(_value_label(any_ser)) + ax.set_title(r"Current density — profile vs $x$ (solid: final, dashed: late mean)") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + return fig + + +# --- per-species density / temperature profiles --------------------------- + + +def _species_diags(run_dir: str | Path) -> dict[str, list[tuple[str, str, Path]]]: + """Group per-species moment diagnostics as ``species -> [(kind, quantity, path)]``. + + Walks ``MS/DENSITY//`` and ``MS/UDIST//`` + (the OSIRIS homes for charge / mass / energy density and for fluid / + thermal velocity moments respectively). + """ + diags = _io.list_diagnostics(run_dir) + out: dict[str, list[tuple[str, str, Path]]] = {} + for rel, path in diags.items(): + parts = rel.split("/") + if len(parts) >= 3 and parts[0] in ("DENSITY", "UDIST"): + out.setdefault(parts[1], []).append((parts[0], "/".join(parts[2:]), path)) + return out + + +def _density_series(entries: list[tuple[str, str, Path]]) -> xr.DataArray | None: + """Pick the best density-like moment for a species and load it as ``(t, x)``.""" + by_q = {q: p for kind, q, p in entries if kind == "DENSITY"} + for pref in ("charge", "n", "n01", "m", "ene"): + if pref in by_q: + ser = _io.load_series(by_q[pref]) + return ser if ser.ndim == 2 else None + # Fall back to any DENSITY entry. + for kind, q, p in entries: + if kind == "DENSITY": + ser = _io.load_series(p) + if ser.ndim == 2: + return ser + return None + + +def _temperature_series(entries: list[tuple[str, str, Path]]) -> xr.DataArray | None: + r"""Build a temperature-like ``(t, x)`` profile from thermal moments. + + Two recognised sources, in priority order: + + - thermal-velocity moments ``uth1/uth2/uth3`` (OSIRIS ``UDIST``): summed in + quadrature, ``T(x) = \sum_i u_{th,i}^2`` (units of ``m c^2``, c = 1); + - a temperature/pressure tensor diagonal ``T11/T22/T33``: summed directly. + + Returns ``None`` when neither is present (e.g. a cold run that dumps no + thermal moment), so callers can skip the temperature profile silently. + """ + by_q = {q: p for _, q, p in entries} + uth = [by_q[f"uth{i}"] for i in (1, 2, 3) if f"uth{i}" in by_q] + tens = [by_q[f"T{i}{i}"] for i in (1, 2, 3) if f"T{i}{i}" in by_q] + if uth: + comps = [_io.load_series(p) for p in uth] + comps = [c for c in comps if c.ndim == 2] + if not comps: + return None + total = sum((c ** 2 for c in comps[1:]), comps[0] ** 2) + long_name = r"T = \sum_i u_{th,i}^2" + units = r"m_e c^2" + elif tens: + comps = [_io.load_series(p) for p in tens] + comps = [c for c in comps if c.ndim == 2] + if not comps: + return None + total = sum(comps[1:], comps[0]) + long_name = r"T = \mathrm{tr}\,T_{ii}" + units = comps[0].attrs.get("units", "") + else: + return None + out = total.copy() + out.attrs = dict(comps[0].attrs) + out.attrs["long_name"] = long_name + out.attrs["units"] = units + out.name = "temperature" + return out + + +def plot_profile( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + abs_value: bool = False, + n_avg_frac: float = 0.2, + value_label: str | None = None, + title: str | None = None, +) -> plt.Axes: + """Plot a ``(t, x)`` moment vs ``x``: final snapshot plus late-time mean. + + The late-time mean (over the last ``n_avg_frac`` of the dumps) smooths + out the time-dependent fluctuations to show the established profile. + """ + da = _decorate(series if isinstance(series, xr.DataArray) else _io.load_series(series)) + if "t" not in da.dims: + raise ValueError(f"plot_profile expects a (t, x) series; got dims {da.dims}") + xdim = next(d for d in da.dims if d != "t") + x = da.coords[xdim].values + nt = da.sizes["t"] + w = max(1, round(n_avg_frac * nt)) + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + final = da.isel(t=-1).values + mean = da.isel(t=slice(nt - w, nt)).mean("t").values + if abs_value: + final, mean = np.abs(final), np.abs(mean) + tlast = float(da.coords["t"].values[-1]) + ax.plot(x, final, lw=1.4, label=f"final ($t={tlast:.3g}$)") + ax.plot(x, mean, lw=1.1, ls="--", label=f"mean of last {w} dumps") + ax.set_xlabel(_axis_label(da, xdim)) + ax.set_ylabel(value_label or _value_label(da)) + ax.set_title(title or rf"{_display_name(da)} — profile vs $x$") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + return ax + + +# --- left/right-going electric-field decomposition ------------------------ + + +def efield_lr_components(run_dir: str | Path) -> dict[str, dict[str, xr.DataArray]]: + r"""Split the transverse electric field into left/right-going parts. + + For a 1D run (propagation along ``x1``), vacuum Maxwell makes the + transverse field pairs into Riemann invariants that advect at ``c`` in a + single direction: + + - ``(E_y, B_z)`` -> right-going ``E_y^R = (e2 + b3)/2``, + left-going ``E_y^L = (e2 - b3)/2``; + - ``(E_z, B_y)`` -> right-going ``E_z^R = (e3 - b2)/2``, + left-going ``E_z^L = (e3 + b2)/2``. + + (Code units: ``c = 1``, so ``|E| = |B|`` for a pure travelling wave and + one of the two parts vanishes.) Returns ``{"e2": {"right":…, "left":…}, + "e3": {…}}`` for whichever transverse pairs were both dumped. + + Caveat: the split is *exact only in vacuum / a uniform non-dispersive + medium*. In a plasma the EM wave is dispersive (``v_φ = ω/k > c``), so + ``|E| ≠ |B|`` mode-by-mode and the decomposition is approximate — still a + useful directional diagnostic, but do not read the residual as physical + counter-propagating power without checking the dispersion. The + longitudinal ``e1`` is electrostatic and is intentionally left out. + """ + diags = _io.list_diagnostics(run_dir) + + def _load(comp: str) -> xr.DataArray | None: + rel = f"FLD/{comp}" + if rel not in diags: + return None + ser = _io.load_series(diags[rel]) + return ser if ser.ndim == 2 else None + + def _pack(values: xr.DataArray, src: xr.DataArray, name: str, long: str) -> xr.DataArray: + out = values.copy() + out.attrs = dict(src.attrs) + out.attrs["long_name"] = long + out.name = name + return out + + out: dict[str, dict[str, xr.DataArray]] = {} + e2, b3 = _load("e2"), _load("b3") + if e2 is not None and b3 is not None: + out["e2"] = { + "right": _pack((e2 + b3) / 2.0, e2, "e2_R", r"E_y^{\rightarrow}"), + "left": _pack((e2 - b3) / 2.0, e2, "e2_L", r"E_y^{\leftarrow}"), + } + e3, b2 = _load("e3"), _load("b2") + if e3 is not None and b2 is not None: + out["e3"] = { + "right": _pack((e3 - b2) / 2.0, e3, "e3_R", r"E_z^{\rightarrow}"), + "left": _pack((e3 + b2) / 2.0, e3, "e3_L", r"E_z^{\leftarrow}"), + } + return out + + +def plot_field_lr_decomposition( + run_dir: str | Path, + *, + v_th: float | None = None, + omega_k_zoom: float | None = 4.0, +) -> dict[str, plt.Figure]: + """One 2x2 figure per transverse component: right/left spacetime + ω-k. + + Top row is the right/left-going ``(t, x)`` spacetime; bottom row is each + part's ``(k, ω)`` power spectrum (zoomed, with the light line drawn) so + you can confirm the right-going part really does sit on the ``ω·k > 0`` + branches and the left-going part on ``ω·k < 0``. + """ + figs: dict[str, plt.Figure] = {} + for comp, parts in efield_lr_components(run_dir).items(): + fig, axes = plt.subplots(2, 2, figsize=(11, 8)) + for col, side in enumerate(("right", "left")): + da = parts[side] + plot_spacetime(da, ax=axes[0, col], title=f"{_tex(_long_name(da))} — spacetime") + try: + zoom = _omega_k_zoom_window(da, omega_k_zoom) + plot_omega_k( + da, ax=axes[1, col], show_light_line=True, + show_langmuir=v_th is not None, v_th=v_th, + k_max=zoom, omega_max=zoom, + title=f"{_tex(_long_name(da))} — $(k, \\omega)$", + ) + except Exception as e: # ω-k overlay is best effort + axes[1, col].set_visible(False) + print(f"[plots] ω-k for {comp} {side} skipped: {e}") + fig.suptitle( + rf"{comp}: left/right-going decomposition (vacuum Riemann split)", y=1.01 + ) + fig.tight_layout() + figs[comp] = fig + return figs + + # --- driver --------------------------------------------------------------- +def _omega_k_zoom_window(series: xr.DataArray, requested: float | None) -> float | None: + """Clamp a requested ``(k, ω)`` zoom half-width to the data's Nyquist range. + + Returns the half-width to use for both axes so the whole ``ω = k`` line is + visible inside the box, or ``None`` to fall back to the full spectrum when + the series is too small to define a window. + """ + t = series.coords["t"].values + xname = next(d for d in series.dims if d != "t") + x = series.coords[xname].values + if t.size < 2 or x.size < 2: + return None + k_ny = np.pi / float(x[1] - x[0]) + w_ny = np.pi / float(t[1] - t[0]) + cap = min(k_ny, w_ny) + if requested is None or requested <= 0: + return cap + return float(min(requested, cap)) + + def save_canned_plots( run_dir: str | Path, out_dir: str | Path, @@ -464,12 +857,19 @@ def save_canned_plots( v_th: float | None = None, dpi: int = 120, n_panels: int = 8, + dist_cells: int = 10, + omega_k_zoom: float | None = 4.0, ) -> dict[str, Path]: """Generate the standard set of PNGs for a finished OSIRIS run. Returns a mapping of plot-name → output PNG path. Each diagnostic family is best-effort: a failure on one diagnostic logs and is skipped rather than aborting the rest. + + ``dist_cells`` sets how many right-boundary cells the phase-space + distribution lineouts average over; ``omega_k_zoom`` is the ``(k, ω)`` + half-width (in ``ω_p`` units) for the zoomed dispersion plots that show the + whole ``ω = k`` line (``None`` → full spectrum). """ run_dir = Path(run_dir) out_dir = Path(out_dir) @@ -514,6 +914,29 @@ def _write(fig: plt.Figure, rel: str) -> Path: plot_omega_k(ser, ax=ax, show_langmuir=v_th is not None, v_th=v_th) written[f"omega_k/{comp}"] = _write(fig, f"omega_k/{comp}.png") + # Zoomed (k, ω) view sized so the whole ω = k line is visible — the + # window where the plasma (Langmuir) waves live. + zoom = _omega_k_zoom_window(ser, omega_k_zoom) + fig, ax = plt.subplots(figsize=(6, 5)) + plot_omega_k( + ser, ax=ax, show_light_line=True, + show_langmuir=v_th is not None, v_th=v_th, + k_max=zoom, omega_max=zoom, + title=f"{_display_name(ser)} — $(k, \\omega)$ (zoom)", + ) + written[f"omega_k_zoom/{comp}"] = _write(fig, f"omega_k_zoom/{comp}.png") + + # --- Currents (j1/j2/j3) combined views --- + try: + fig = plot_currents_spacetime(run_dir) + if fig is not None: + written["currents/spacetime"] = _write(fig, "currents/spacetime.png") + fig = plot_currents_lineouts(run_dir) + if fig is not None: + written["currents/lineouts"] = _write(fig, "currents/lineouts.png") + except Exception as e: + print(f"[plots] skipping currents: {e}") + # --- Per-species moments (DENSITY//) --- for diag_rel, diag_path in sorted(diags.items()): if not diag_rel.startswith("DENSITY/"): @@ -547,6 +970,29 @@ def _write(fig: plt.Figure, rel: str) -> Path: f"moments/{species}/lineouts/{quantity}.png", ) + # --- Per-species density + temperature profiles --- + for species, entries in sorted(_species_diags(run_dir).items()): + try: + dens = _density_series(entries) + if dens is not None: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_profile(dens, ax=ax, title=f"{species} — density profile") + written[f"profiles/{species}/density"] = _write( + fig, f"profiles/{species}/density.png" + ) + except Exception as e: + print(f"[plots] skipping density profile for {species}: {e}") + try: + temp = _temperature_series(entries) + if temp is not None: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_profile(temp, ax=ax, title=f"{species} — temperature profile") + written[f"profiles/{species}/temperature"] = _write( + fig, f"profiles/{species}/temperature.png" + ) + except Exception as e: + print(f"[plots] skipping temperature profile for {species}: {e}") + # --- Phase space (PHA//): final step + time evolution --- for diag_rel, diag_path in sorted(diags.items()): if not diag_rel.startswith("PHA/"): @@ -569,9 +1015,24 @@ def _write(fig: plt.Figure, rel: str) -> Path: plot_phasespace_evolution(ser, n_panels=n_panels), f"phasespace_evolution/{species}/{ps_name}.png", ) + # f(p) averaged over the right-boundary cells, vs time. + fig, ax = plt.subplots(figsize=(6, 4)) + plot_distribution_lineout(ser, ax=ax, n_cells=dist_cells) + written[f"distribution_lineouts/{species}/{ps_name}"] = _write( + fig, f"distribution_lineouts/{species}/{ps_name}.png" + ) except Exception as e: print(f"[plots] skipping phasespace evolution {diag_rel}: {e}") + # --- Left/right-going transverse E-field decomposition --- + try: + for comp, fig in plot_field_lr_decomposition( + run_dir, v_th=v_th, omega_k_zoom=omega_k_zoom + ).items(): + written[f"field_decomp/{comp}"] = _write(fig, f"field_decomp/{comp}.png") + except Exception as e: + print(f"[plots] skipping field decomposition: {e}") + # --- Energy diagnostics --- try: fig, ax = plt.subplots(figsize=(6, 4)) @@ -615,16 +1076,34 @@ def _ensure_series(src) -> xr.DataArray: ) +def _tex(s) -> str: + r"""Wrap an OSIRIS label/unit in math mode so its TeX actually renders. + + OSIRIS stores labels and units as bare TeX fragments (``E_1``, + ``\omega_p``, ``c / \omega_p``, ``1 / \omega_p``). Matplotlib only + typesets those inside ``$...$``, so a raw ``\omega`` would otherwise be + drawn literally. We add the delimiters when the string carries TeX + markup (a backslash, a sub/superscript, or a fraction slash) and leave + plain words ("charge", "a.u.") untouched. Idempotent: a string that is + already ``$``-delimited is returned unchanged. + """ + s = str(s) + if not s or (s.startswith("$") and s.endswith("$")): + return s + if any(c in s for c in "\\_^/"): + return f"${s}$" + return s + + def _axis_label(da: xr.DataArray, dim: str) -> str: if dim == "t": - u = da.attrs.get("time_units", r"$1/\omega_p$") - return f"t [{u}]" if u else "t" + u = da.attrs.get("time_units") or r"1/\omega_p" + return rf"$t$ [{_tex(u)}]" if u else r"$t$" units = da.attrs.get("axis_units", {}).get(dim, "") long = da.attrs.get("axis_long_names", {}).get(dim, dim) if units: - # OSIRIS axis units are TeX-ish (e.g. "c / \\omega_p"); wrap in $...$. - return rf"{long} [${units}$]" - return long + return rf"{_tex(long)} [{_tex(units)}]" + return _tex(long) def _long_name(da: xr.DataArray) -> str: @@ -632,11 +1111,16 @@ def _long_name(da: xr.DataArray) -> str: return str(da.attrs.get("long_name") or da.name or "") +def _display_name(da: xr.DataArray) -> str: + """``_long_name`` wrapped for math-mode rendering (used in titles).""" + return _tex(_long_name(da)) + + def _value_label(da: xr.DataArray) -> str: """Quantity label with units (for colorbars / y-axes).""" - name = _long_name(da) + name = _tex(_long_name(da)) units = da.attrs.get("units", "") - return rf"{name} [${units}$]" if units else name + return rf"{name} [{_tex(units)}]" if units else name def _decorate(da: xr.DataArray) -> xr.DataArray: diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 0d5b92a4..09c7964e 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -148,8 +148,13 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: # Render the standard plot set into td/plots so MLflow logs them as # artifacts. Never let a plotting failure abort metric/artifact logging. try: - v_th = (cfg.get("output") or {}).get("v_th") - _plots.save_canned_plots(run_dir, td / "plots", v_th=v_th) + output = cfg.get("output") or {} + kwargs: dict[str, Any] = {"v_th": output.get("v_th")} + if output.get("dist_cells") is not None: + kwargs["dist_cells"] = int(output["dist_cells"]) + if "omega_k_zoom" in output: # may be explicitly null to disable zoom + kwargs["omega_k_zoom"] = output["omega_k_zoom"] + _plots.save_canned_plots(run_dir, td / "plots", **kwargs) except Exception as e: print(f"[post] plotting failed: {e}") diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index cc3ac5c3..53380380 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -18,3 +18,6 @@ osiris: output: # diagnostics_to_log: [e1] # null/omitted = convert every diagnostic's full time history to netCDF + # v_th: 0.1 # overlay the Langmuir (Bohm-Gross) branch on omega-k plots + # dist_cells: 10 # right-boundary cells averaged for the phase-space f(p) lineouts + # omega_k_zoom: 4.0 # (k, omega) half-width [omega_p] for the zoomed dispersion view; null = full diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index eaa2e641..176941dd 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -92,6 +92,12 @@ osiris: output: diagnostics_to_log: null # null = all; or [e1, charge, …] + v_th: 0.1 # optional: overlays the Bohm–Gross + # Langmuir branch on ω–k plots + dist_cells: 10 # right-boundary cells averaged for + # the phase-space f(p) lineouts + omega_k_zoom: 4.0 # (k, ω) half-width [ω_p] for the + # zoomed dispersion view; null = full ``` Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1)`). Indexed `{0: …, 1: …}` form addresses occurrences of repeated sections (`species`, `udist`, `profile`, `spe_bound`, `diag_species`, `zpulse`, …) in source order. @@ -106,6 +112,28 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | Artifacts | `config.yaml`, `derived_config.yaml`, `units.yaml` (adept stock) | | Artifacts | `os-stdin` (rendered OSIRIS deck), `stdout.log`, `stderr.log` | | Artifacts | `binary//.nc` — one xarray netCDF per diagnostic, holding the full `(t, …)` time history (replaces the raw h5 dumps) | +| Artifacts | `plots/…` — canned PNGs (see below) | + +## Canned plots (`plots/` artifacts) + +`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 | +| --------------------------------------------- | ------------- | +| `spacetime/.png`, `spacetime_log/.png` | `(t, x)` heatmap of each FLD diagnostic (lin + log) | +| `lineouts/.png` | value-vs-`x` snapshots at sampled times | +| `omega_k/.png` | full 2-D FFT `(k, ω)` dispersion | +| `omega_k_zoom/.png` | zoomed `(k, ω)` showing the whole `ω = k` light line — where the plasma (Langmuir) waves live | +| `currents/spacetime.png`, `currents/lineouts.png` | combined `J_x/J_y/J_z` (`j1/j2/j3`) views | +| `moments//…` | per-species density-moment spacetime + lineouts | +| `profiles//density.png` | density profile vs `x` (final snapshot + late-time mean) | +| `profiles//temperature.png` | temperature profile vs `x`, from `uth1/2/3` or `T11/22/33` moments (omitted if neither is dumped) | +| `phasespace//.png`, `phasespace_evolution/…` | `(x, p)` phase-space heatmaps | +| `distribution_lineouts//.png` | `f(p)` averaged over the rightmost `dist_cells` cells, overlaid at sampled times | +| `field_decomp/.png` | left/right-going transverse `E` (vacuum Riemann split `(e2±b3)/2`, `(e3∓b2)/2`), spacetime + `ω–k` | +| `energy_vs_time.png`, `energy_components_vs_time.png`, `total_energy_vs_time.png` | 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 longitudinal `e1` is electrostatic and is intentionally excluded. ## Programmatic use diff --git a/tests/test_osiris/test_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py new file mode 100644 index 00000000..f28e3653 --- /dev/null +++ b/tests/test_osiris/test_plots_new_views.py @@ -0,0 +1,243 @@ +"""Tests for the OSIRIS plotting additions. + +Synthesizes tiny OSIRIS-shaped runs (no real solver needed) to exercise: + + - proper-LaTeX axis/value labels (``_tex`` wrapping) + - phase-space distribution lineouts averaged over the right-boundary cells + - the zoomed (k, ω) dispersion view with the light line + - combined J_x/J_y/J_z (j1/j2/j3) current figures + - per-species density + temperature profiles + - the left/right-going transverse E-field decomposition (incl. correctness) + - that ``save_canned_plots`` emits all of the above +""" + +from __future__ import annotations + +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import h5py +import numpy as np +import xarray as xr + +from adept.osiris import io as oio +from adept.osiris import plots as oplt + + +def _write_dump(path: Path, name: str, data: np.ndarray, t: float, it: int, axes) -> None: + """Write one OSIRIS-style HDF5 grid/phasespace dump. + + ``axes`` is a list of ``(name, long_name, units, min, max)`` in OSIRIS + (Fortran) order; the last entry is the fastest-varying numpy axis. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + f.attrs["NAME"] = np.array([name.encode()], dtype="S256") + f.attrs["LABEL"] = np.array([name.encode()], dtype="S256") + f.attrs["UNITS"] = np.array([b"a.u."], dtype="S256") + f.attrs["TIME UNITS"] = np.array([rb"1 / \omega_p"], dtype="S256") + ax = f.create_group("AXIS") + for i, (an, ln, un, lo, hi) in enumerate(axes, start=1): + d = ax.create_dataset(f"AXIS{i}", data=np.array([lo, hi], dtype="float64")) + d.attrs["NAME"] = np.array([an.encode()], dtype="S256") + d.attrs["LONG_NAME"] = np.array([ln.encode()], dtype="S256") + d.attrs["UNITS"] = np.array([un.encode()], dtype="S256") + sim = f.create_group("SIMULATION") + sim.attrs["DT"] = np.array([0.05]) + sim.attrs["NDIMS"] = np.array([1], dtype="int32") + sim.attrs["NX"] = np.array([data.shape[-1]], dtype="int32") + sim.attrs["XMIN"] = np.array([axes[-1][3]]) + sim.attrs["XMAX"] = np.array([axes[-1][4]]) + f.create_dataset(name, data=data.astype("float32")) + + +X_AX = ("x1", "x_1", r"c / \omega_p", 0.0, 10.0) +P_AX = ("p1", "p_1", r"m_e c", -2.0, 2.0) + + +def _make_rich_run(root: Path, n_steps: int = 6, nx: int = 16, npx: int = 12) -> Path: + """Run with EM fields, currents, density+thermal moments, and phase space.""" + run_dir = root / "run" + rng = np.random.default_rng(0) + x = np.linspace(X_AX[3], X_AX[4], nx) + + for k in range(n_steps): + it, t = k * 10, k * 0.5 + # Right-going transverse pairs: e2 = b3, e3 = -b2 (so left part ~ 0). + wave = np.sin(2 * np.pi * (x - t) / 10.0) + _write_dump(run_dir / "MS/FLD/e2" / f"e2-{it:06d}.h5", "e2", wave, t, it, [X_AX]) + _write_dump(run_dir / "MS/FLD/b3" / f"b3-{it:06d}.h5", "b3", wave, t, it, [X_AX]) + _write_dump(run_dir / "MS/FLD/e3" / f"e3-{it:06d}.h5", "e3", wave, t, it, [X_AX]) + _write_dump(run_dir / "MS/FLD/b2" / f"b2-{it:06d}.h5", "b2", -wave, t, it, [X_AX]) + # Currents j1/j2/j3. + for j in ("j1", "j2", "j3"): + _write_dump(run_dir / f"MS/FLD/{j}" / f"{j}-{it:06d}.h5", j, + rng.standard_normal(nx), t, it, [X_AX]) + # Density + thermal-velocity moments for a species. + _write_dump(run_dir / "MS/DENSITY/electron/charge" / f"charge-electron-{it:06d}.h5", + "charge", -np.abs(rng.standard_normal(nx)) - 1.0, t, it, [X_AX]) + for u in ("uth1", "uth2", "uth3"): + _write_dump(run_dir / f"MS/UDIST/electron/{u}" / f"{u}-electron-{it:06d}.h5", + u, 0.1 + 0.01 * rng.standard_normal(nx), t, it, [X_AX]) + # Phase space (p, x). + _write_dump(run_dir / "MS/PHA/x1p1/electron" / f"x1p1-electron-{it:06d}.h5", + "x1p1", rng.random((npx, nx)), t, it, [X_AX, P_AX]) + return run_dir + + +# --- LaTeX labels --------------------------------------------------------- + + +def test_tex_wraps_only_tex_fragments() -> None: + assert oplt._tex(r"\omega_p") == r"$\omega_p$" + assert oplt._tex("x_1") == "$x_1$" + assert oplt._tex(r"c / \omega_p") == r"$c / \omega_p$" + assert oplt._tex("charge") == "charge" # plain word, left alone + assert oplt._tex("a.u.") == "a.u." + assert oplt._tex(r"$E_1$") == r"$E_1$" # idempotent + + +def test_axis_label_is_math_mode(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path, n_steps=3) + da = oio.load_series(run_dir / "MS/FLD/e2") + # Time-axis units wrapped in $...$ instead of rendered literally. + assert oplt._axis_label(da, "t") == r"$t$ [$1 / \omega_p$]" + # Spatial axis: both long-name and units in math mode. + assert oplt._axis_label(da, "x1") == r"$x_1$ [$c / \omega_p$]" + + +# --- distribution lineouts ------------------------------------------------ + + +def test_distribution_lineout_averages_right_cells() -> None: + p = np.linspace(-2, 2, 11) + x = np.linspace(0, 10, 16) + fp = np.exp(-(p**2)) + data = np.tile(fp[:, None], (1, x.size)) # f independent of x + da = xr.DataArray( + data, coords={"p1": p, "x1": x}, dims=("p1", "x1"), name="x1p1", + attrs={"axis_units": {"p1": r"m_e c", "x1": r"c / \omega_p"}, + "axis_long_names": {"p1": "p_1", "x1": "x_1"}, "units": "a.u."}, + ) + ax = oplt.plot_distribution_lineout(da, n_cells=10) + y = ax.lines[0].get_ydata() + np.testing.assert_allclose(y, fp, atol=1e-12) + + +def test_distribution_lineout_time_series(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + ser = oio.load_series(run_dir / "MS/PHA/x1p1/electron") + ax = oplt.plot_distribution_lineout(ser, n_cells=5, n_times=4) + assert len(ax.lines) == 4 # one curve per sampled time + + +# --- zoomed omega-k ------------------------------------------------------- + + +def test_omega_k_zoom_window_clamps_to_nyquist(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + ser = oio.load_series(run_dir / "MS/FLD/e2") + big = oplt._omega_k_zoom_window(ser, requested=1e6) # huge -> clamp to Nyquist + small = oplt._omega_k_zoom_window(ser, requested=2.0) + assert small == 2.0 + assert 0 < big < 1e6 + + +def test_omega_k_light_line_drawn(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + ser = oio.load_series(run_dir / "MS/FLD/e2") + import matplotlib.pyplot as plt + + _, ax = plt.subplots() + oplt.plot_omega_k(ser, ax=ax, show_em=False, show_light_line=True, k_max=4, omega_max=4) + labels = [ln.get_label() for ln in ax.lines] + assert any("light line" in str(lbl) for lbl in labels) + plt.close("all") + + +# --- currents ------------------------------------------------------------- + + +def test_current_components_loaded(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + comps = oplt._current_components(run_dir) + assert set(comps) == {"j1", "j2", "j3"} + assert oplt.plot_currents_spacetime(run_dir) is not None + assert oplt.plot_currents_lineouts(run_dir) is not None + + +def test_currents_absent_returns_none(tmp_path: Path) -> None: + (tmp_path / "run" / "MS").mkdir(parents=True) + assert oplt.plot_currents_spacetime(tmp_path / "run") is None + + +# --- density / temperature profiles --------------------------------------- + + +def test_temperature_series_from_uth(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + entries = oplt._species_diags(run_dir)["electron"] + temp = oplt._temperature_series(entries) + assert temp is not None + # T = uth1^2 + uth2^2 + uth3^2, so ~ 3 * 0.1^2 = 0.03, and positive. + assert float(temp.isel(t=-1).mean()) > 0 + assert temp.attrs["long_name"].startswith("T =") + dens = oplt._density_series(entries) + assert dens is not None and dens.name == "charge" + + +def test_temperature_series_absent_returns_none(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + _write_dump(run_dir / "MS/DENSITY/ion/charge" / "charge-ion-000000.h5", + "charge", np.ones(8), 0.0, 0, [X_AX]) + entries = oplt._species_diags(run_dir)["ion"] + assert oplt._temperature_series(entries) is None + + +# --- left/right-going E decomposition ------------------------------------- + + +def test_efield_decomposition_isolates_right_going(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + parts = oplt.efield_lr_components(run_dir) + assert set(parts) == {"e2", "e3"} + e2 = oio.load_series(run_dir / "MS/FLD/e2") + # e2 == b3 is a pure right-going wave: right == e2, left == 0. + np.testing.assert_allclose(parts["e2"]["right"].values, e2.values, atol=1e-6) + np.testing.assert_allclose(parts["e2"]["left"].values, 0.0, atol=1e-6) + # e3 == -b2 is also pure right-going. + e3 = oio.load_series(run_dir / "MS/FLD/e3") + np.testing.assert_allclose(parts["e3"]["right"].values, e3.values, atol=1e-6) + np.testing.assert_allclose(parts["e3"]["left"].values, 0.0, atol=1e-6) + + +def test_field_decomposition_figures(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + figs = oplt.plot_field_lr_decomposition(run_dir, omega_k_zoom=4.0) + assert set(figs) == {"e2", "e3"} + + +# --- end-to-end driver ---------------------------------------------------- + + +def test_save_canned_plots_emits_new_views(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + written = oplt.save_canned_plots(run_dir, tmp_path / "plots", v_th=0.1) + expected = { + "omega_k_zoom/e2", + "currents/spacetime", + "currents/lineouts", + "profiles/electron/density", + "profiles/electron/temperature", + "distribution_lineouts/electron/x1p1", + "field_decomp/e2", + "field_decomp/e3", + } + assert expected <= set(written) + for path in written.values(): + assert path.exists() and path.stat().st_size > 0 From 704aba83863fb5c9130ec1970bbbe511816e042d Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 1 Jun 2026 16:49:14 -0700 Subject: [PATCH 09/49] =?UTF-8?q?Modified=20adept=20(io.py=20+=20plots.py)?= =?UTF-8?q?=20so=20the=20full=20OSIRIS=20canned-plot=20set=20regenerates?= =?UTF-8?q?=20from=20the=20saved=20NetCDF=20artifacts=20alone=20=E2=80=94?= =?UTF-8?q?=20no=20rerun,=20no=20raw=20MS/=20tree.=20Made=20list=5Fdiagnos?= =?UTF-8?q?tics/load=5Fseries/load=5Fhist=5Fenergy=20dispatch=20between=20?= =?UTF-8?q?the=20MS/=20HDF5=20tree=20and=20a=20binary/=20NetCDF=20dir,=20m?= =?UTF-8?q?ade=20field-energy=20source-agnostic,=20and=20now=20persist=20H?= =?UTF-8?q?IST/energy.nc=20in=20save=5Frun=5Fdatasets.=20Added=203=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- adept/osiris/io.py | 122 ++++++++++++++++++--- adept/osiris/plots.py | 146 +++++++++++++------------- tests/test_osiris/test_post_netcdf.py | 146 ++++++++++++++++++++++++++ 3 files changed, 329 insertions(+), 85 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 304dfb11..058276d2 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -266,8 +266,14 @@ def load_series(directory: str | Path) -> xr.DataArray: All files must share the same diagnostic name, the same spatial / phase-space shape, and the same axis bounds — this is the standard OSIRIS convention for a single diagnostic's time history. + + For regenerating plots from saved artifacts, ``directory`` may instead be a + single ``.nc`` file written by :func:`save_run_datasets`; it is then loaded + via :func:`load_series_nc`, returning the same stacked ``(t, ...)`` array. """ directory = Path(directory) + if directory.is_file() and directory.suffix == ".nc": + return load_series_nc(directory) dumps = _sort_dumps(directory) if not dumps: raise FileNotFoundError(f"No .h5 dumps in {directory}") @@ -304,23 +310,36 @@ def load_series(directory: str | Path) -> xr.DataArray: def list_diagnostics(run_dir: str | Path) -> dict[str, Path]: - """Map every diagnostic name to its directory under ``run_dir/MS/``. + """Map every diagnostic name to its data source. + + Two layouts are supported transparently: + + - a raw OSIRIS run directory with an ``MS/`` tree: any directory that + directly contains ``.h5`` dumps is a diagnostic, keyed by its relative + path under ``MS/`` (e.g. ``"FLD/e1"``, ``"PHA/x1p1/beam_pos"``); + - the ``binary/`` directory of saved NetCDFs from :func:`save_run_datasets` + (no ``MS/``): each ``.nc`` is a diagnostic, keyed by its relative path + without the suffix (same keys as above). - Discovery rule: any directory that directly contains ``.h5`` dumps is - a diagnostic. The returned key is the directory's relative path under - ``MS/`` (e.g. ``"FLD/e1"``, ``"PHA/x1p1/beam_pos"``). + The returned handles (dump dirs or ``.nc`` files) are both accepted by + :func:`load_series`, so callers regenerate plots from either source without + caring which. Raises ``FileNotFoundError`` if neither layout is found. """ run_dir = Path(run_dir) ms = run_dir / "MS" - if not ms.is_dir(): - raise FileNotFoundError(f"No MS/ directory under {run_dir}") - out: dict[str, Path] = {} - for d in ms.rglob("*"): - if not d.is_dir(): - continue - if any(c.suffix == ".h5" for c in d.iterdir()): - out[str(d.relative_to(ms))] = d - return out + if ms.is_dir(): + out: dict[str, Path] = {} + for d in ms.rglob("*"): + if not d.is_dir(): + continue + if any(c.suffix == ".h5" for c in d.iterdir()): + out[str(d.relative_to(ms))] = d + return out + # No MS/ tree — fall back to the saved-NetCDF layout. + nc = list_diagnostics_nc(run_dir) + if nc: + return nc + raise FileNotFoundError(f"No MS/ directory or saved NetCDFs under {run_dir}") def _coerce_attr(v): @@ -359,6 +378,61 @@ def series_to_dataset(da: xr.DataArray) -> xr.Dataset: return ds +def load_series_nc(path: str | Path) -> xr.DataArray: + """Load a grid/phase-space series NetCDF back into a plotter-ready DataArray. + + This is the inverse of :func:`series_to_dataset` (the form + :func:`save_run_datasets` writes to disk): it reopens a single-diagnostic + ``.nc`` file and rebuilds the ``axis_units`` / ``axis_long_names`` dict + attrs from the per-coordinate metadata, so the returned ``DataArray`` is + indistinguishable (for plotting purposes) from one produced by + :func:`load_series` off the raw HDF5 tree. The whole point is that the + canned plots can be regenerated from the saved NetCDFs alone. + + The file is read fully into memory and closed before returning. + """ + ds = xr.load_dataset(path, engine="h5netcdf") + names = list(ds.data_vars) + if not names: + raise ValueError(f"No data variable in {path}") + # series_to_dataset writes exactly one data variable (the diagnostic); + # `iter` rides along as a coordinate, not a data var. + name = names[0] + da = ds[name] + + axis_units: dict[str, str] = {} + axis_long: dict[str, str] = {} + for dim in da.dims: + if dim == "t" or dim not in ds.coords: + continue + cu = ds[dim].attrs.get("units") + cl = ds[dim].attrs.get("long_name") + if cu: + axis_units[str(dim)] = cu + if cl: + axis_long[str(dim)] = cl + da.attrs["axis_units"] = axis_units + da.attrs["axis_long_names"] = axis_long + da.attrs.setdefault("time_units", da.attrs.get("time_units", r"1/\omega_p")) + return da + + +def list_diagnostics_nc(binary_dir: str | Path) -> dict[str, Path]: + """Map every saved-NetCDF diagnostic to its ``.nc`` file. + + The inverse layout of :func:`save_run_datasets`: walks ``binary_dir`` for + ``*.nc`` files and keys each by its relative path with the suffix dropped + (e.g. ``"FLD/e1"``, ``"PHA/x1p1/beam_pos"``, ``"HIST/energy"``), mirroring + :func:`list_diagnostics` over the raw ``MS/`` tree. + """ + binary_dir = Path(binary_dir) + out: dict[str, Path] = {} + for p in sorted(binary_dir.rglob("*.nc")): + rel = str(p.relative_to(binary_dir).with_suffix("")) + out[rel] = p + return out + + def save_run_datasets( run_dir: str | Path, out_dir: str | Path, @@ -396,6 +470,20 @@ def save_run_datasets( except Exception as e: # one bad diagnostic must not abort the rest print(f"[post] skipping diagnostic {relpath}: {e}") continue + + # Persist the HIST scalar-energy history (field + per-species kinetic) so + # the energy-conservation plot can be regenerated from the saved NetCDFs + # alone — it is the one plot input that lives outside the MS/ dump tree. + try: + energy = load_hist_energy(run_dir) + if energy is not None: + dest = out_dir / "HIST" / "energy.nc" + dest.parent.mkdir(parents=True, exist_ok=True) + energy.to_netcdf(dest, engine="h5netcdf") + written.append(dest) + except Exception as e: + print(f"[post] skipping HIST energy: {e}") + return written @@ -454,7 +542,15 @@ def load_hist_energy(run_dir: str | Path) -> xr.Dataset | None: Note: the column convention assumed is the documented OSIRIS ``iteration time `` layout; validate against a real run with energy diagnostics enabled before relying on absolute magnitudes. + + When given the ``binary/`` directory of saved NetCDFs (no raw ``HIST/`` + ASCII), this instead reloads the pre-parsed ``HIST/energy.nc`` written by + :func:`save_run_datasets`, so the energy-conservation plot is reproducible + from the saved artifacts alone. """ + saved = Path(run_dir) / "HIST" / "energy.nc" + if saved.is_file(): + return xr.load_dataset(saved, engine="h5netcdf") hist = Path(run_dir) / "HIST" if not hist.is_dir(): return None diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index fa678929..e76627fb 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -202,43 +202,15 @@ def plot_phasespace_evolution( def field_energy_series(run_dir: str | Path) -> xr.DataArray: """Sum ``(|E|^2 + |B|^2) / 2`` over space at every saved step. - Walks every component dir under ``MS/FLD/`` (``e1``, ``e2``, ..., ``b1``, - ``b2``, ``b3``), aligns by iteration count, and returns a 1D - ``DataArray`` of total field energy in code units vs time. Components - that aren't dumped (or that have a different save cadence) are - skipped. + Returns a 1D ``DataArray`` of total field energy in code units vs time. + ``run_dir`` may be a raw OSIRIS run directory **or** the ``binary/`` + directory of saved NetCDFs (it is resolved through :func:`io.load_series`), + so the trace is reproducible from the saved artifacts alone. """ - run_dir = Path(run_dir) - fld = run_dir / "MS" / "FLD" - if not fld.is_dir(): - raise FileNotFoundError(f"No MS/FLD under {run_dir}") - components = ("e1", "e2", "e3", "b1", "b2", "b3") - by_iter: dict[int, dict[str, float]] = {} - times: dict[int, float] = {} - for comp in components: - d = fld / comp - if not d.is_dir(): - continue - for h5 in sorted(d.iterdir()): - if h5.suffix != ".h5": - continue - da = _io.load_grid_h5(h5) - e = 0.5 * float((da.values ** 2).sum()) * _cell_volume(da) - it = int(da.attrs["iter"]) - by_iter.setdefault(it, {})[comp] = e - times.setdefault(it, float(da.attrs["time"])) - if not by_iter: - raise RuntimeError(f"No field dumps under {fld}") - iters = sorted(by_iter) - totals = np.array([sum(by_iter[i].values()) for i in iters]) - t = np.array([times[i] for i in iters]) - return xr.DataArray( - totals, - coords={"t": t, "iter": ("t", np.asarray(iters))}, - dims=("t",), - name="field_energy", - attrs={"long_name": "total field energy", "units": "code"}, - ) + ds = field_energy_components(run_dir) + da = ds["total_field_energy"].rename("field_energy") + da.attrs.update(long_name="total field energy", units="code") + return da def plot_energy_vs_time( @@ -270,36 +242,50 @@ def plot_energy_vs_time( def field_energy_components(run_dir: str | Path) -> xr.Dataset: - """E-field, B-field, and total field energy vs time from ``MS/FLD`` dumps. + """E-field, B-field, and total field energy vs time from the FLD diagnostics. Like :func:`field_energy_series` but keeps the electric (``e1/e2/e3``) and magnetic (``b1/b2/b3``) contributions separate. Returns a ``Dataset`` with data variables ``E_energy``, ``B_energy`` and ``total_field_energy`` on a shared ``t`` axis. Components not dumped are simply omitted from their sum. + + Sources are resolved via :func:`io.list_diagnostics` / :func:`io.load_series`, + so ``run_dir`` may be a raw OSIRIS run directory or the ``binary/`` directory + of saved NetCDFs — the energy is recomputed from the stored ``(t, x)`` field + series either way. """ - run_dir = Path(run_dir) - fld = run_dir / "MS" / "FLD" - if not fld.is_dir(): - raise FileNotFoundError(f"No MS/FLD under {run_dir}") + diags = _io.list_diagnostics(run_dir) groups = {"E_energy": ("e1", "e2", "e3"), "B_energy": ("b1", "b2", "b3")} by_iter: dict[int, dict[str, float]] = {} times: dict[int, float] = {} + found = False for group, comps in groups.items(): for comp in comps: - d = fld / comp - if not d.is_dir(): + rel = f"FLD/{comp}" + if rel not in diags: continue - for h5 in sorted(d.iterdir()): - if h5.suffix != ".h5": - continue - da = _io.load_grid_h5(h5) - e = 0.5 * float((da.values ** 2).sum()) * _cell_volume(da) - it = int(da.attrs["iter"]) + try: + ser = _io.load_series(diags[rel]) + except Exception as e: # a bad component must not sink the rest + print(f"[plots] skipping field-energy component {comp}: {e}") + continue + if ser.ndim != 2: + continue + found = True + e_t = _field_energy_from_series(ser) + its = ( + np.asarray(ser.coords["iter"].values) + if "iter" in ser.coords + else np.arange(ser.sizes["t"]) + ) + ts = np.asarray(ser.coords["t"].values, dtype="float64") + for k in range(ts.size): + it = int(its[k]) rec = by_iter.setdefault(it, {"E_energy": 0.0, "B_energy": 0.0}) - rec[group] += e - times.setdefault(it, float(da.attrs["time"])) - if not by_iter: - raise RuntimeError(f"No field dumps under {fld}") + rec[group] += float(e_t[k]) + times.setdefault(it, float(ts[k])) + if not found or not by_iter: + raise RuntimeError(f"No field dumps available under {run_dir}") iters = sorted(by_iter) t = np.array([times[i] for i in iters]) e_arr = np.array([by_iter[i]["E_energy"] for i in iters]) @@ -866,6 +852,15 @@ def save_canned_plots( best-effort: a failure on one diagnostic logs and is skipped rather than aborting the rest. + ``run_dir`` may be a raw OSIRIS run directory (with an ``MS/`` HDF5 tree and + optional ``HIST/``) **or** the ``binary/`` directory of saved NetCDFs + written by :func:`io.save_run_datasets`. The data source is detected + automatically (every diagnostic is fetched through :func:`io.list_diagnostics` + / :func:`io.load_series`, which handle both layouts), so the full plot set — + including the field-energy and energy-conservation traces — can be + regenerated from the saved NetCDF artifacts alone, with no rerun and no raw + HDF5 dumps. + ``dist_cells`` sets how many right-boundary cells the phase-space distribution lineouts average over; ``omega_k_zoom`` is the ``(k, ω)`` half-width (in ``ω_p`` units) for the zoomed dispersion plots that show the @@ -1001,15 +996,23 @@ def _write(fig: plt.Figure, rel: str) -> Path: if len(parts) < 3: continue ps_name, species = parts[1], parts[2] - last = _latest_h5(diag_path) - if last is not None: - fig, ax = plt.subplots(figsize=(5, 4)) - plot_phasespace(_io.load_phasespace_h5(last), ax=ax) - written[f"phasespace/{species}/{ps_name}"] = _write( - fig, f"phasespace/{species}/{ps_name}.png" - ) try: ser = _io.load_series(diag_path) + except Exception as e: + print(f"[plots] skipping phasespace {diag_rel}: {e}") + continue + # Final-step heatmap: the last time slice of the series (so it works + # identically off raw dumps and off the saved NetCDF series). + if "t" in ser.dims and ser.sizes["t"] > 0: + final = ser.isel(t=-1) + final.attrs["time"] = float(np.asarray(final.coords["t"].values)) + if final.ndim == 2: + fig, ax = plt.subplots(figsize=(5, 4)) + plot_phasespace(final, ax=ax) + written[f"phasespace/{species}/{ps_name}"] = _write( + fig, f"phasespace/{species}/{ps_name}.png" + ) + try: if ser.ndim == 3: written[f"phasespace_evolution/{species}/{ps_name}"] = _write( plot_phasespace_evolution(ser, n_panels=n_panels), @@ -1158,15 +1161,14 @@ def _cell_volume(da: xr.DataArray) -> float: return dvol -def _latest_h5(diag_dir: Path) -> Path | None: - best: tuple[int, Path] | None = None - for p in diag_dir.iterdir(): - if p.suffix != ".h5": - continue - m = _io._ITER_RE.search(p.name) - if not m: - continue - it = int(m.group(1)) - if best is None or it > best[0]: - best = (it, p) - return best[1] if best else None +def _field_energy_from_series(ser: xr.DataArray) -> np.ndarray: + """Per-timestep field energy ``0.5 * ∫ f^2 dx`` for a ``(t, x)`` series. + + Returns a 1-D array aligned with the series' ``t`` axis. Equivalent to + summing ``0.5 * f^2 * cell_volume`` over space at each dump, but vectorized + over the whole stacked series so it works identically off the raw HDF5 + dumps and off the saved NetCDFs. + """ + spatial = [d for d in ser.dims if d != "t"] + sq = (ser.astype("float64") ** 2).sum(dim=spatial) + return 0.5 * np.asarray(sq.values, dtype="float64") * _cell_volume(ser) diff --git a/tests/test_osiris/test_post_netcdf.py b/tests/test_osiris/test_post_netcdf.py index e2e489a2..64d096c3 100644 --- a/tests/test_osiris/test_post_netcdf.py +++ b/tests/test_osiris/test_post_netcdf.py @@ -247,3 +247,149 @@ def test_save_run_datasets_skips_unloadable_diagnostic(tmp_path: Path) -> None: assert (out / "FLD" / "e1.nc").exists() # good diagnostic still produced assert not (out / "FLD" / "bad.nc").exists() assert (out / "FLD" / "e1.nc") in written + + +# --- regenerating plots from the saved NetCDFs (no rerun, no raw MS/ tree) --- + + +def _write_field(run_dir: Path, comp: str, n_steps: int, nx: int, seed: int) -> None: + """Write an ``MS/FLD/`` 1-D field diagnostic over ``n_steps`` dumps.""" + rng = np.random.default_rng(seed) + d = run_dir / "MS" / "FLD" / comp + for k in range(n_steps): + it = k * 10 + _write_dump( + d / f"{comp}-{it:06d}.h5", + comp, + rng.standard_normal(nx), + t=k * 0.5, + it=it, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + + +def _write_phasespace(run_dir: Path, n_steps: int, nx: int, npmom: int) -> None: + """Write an ``MS/PHA/p1x1/electrons`` 2-D phase-space series.""" + rng = np.random.default_rng(7) + d = run_dir / "MS" / "PHA" / "p1x1" / "electrons" + for k in range(n_steps): + it = k * 10 + # data shape (p1, x1): the AXIS list is OSIRIS order (AXIS1 = x1). + _write_dump( + d / f"p1x1-electrons-{it:06d}.h5", + "p1x1", + rng.standard_normal((npmom, nx)), + t=k * 0.5, + it=it, + axes=[ + ("x1", "x_1", r"c / \omega_p", 0.0, 10.0), + ("p1", "p_1", r"m_e c", -1.0, 1.0), + ], + ) + + +def _write_hist_energy(run_dir: Path, n_steps: int) -> None: + """Write OSIRIS-style ``HIST/`` field + kinetic energy ASCII tables.""" + hist = run_dir / "HIST" + hist.mkdir(parents=True, exist_ok=True) + fld = ["! iter time e1 e2 e3 b1 b2 b3"] + par = [] + for k in range(n_steps): + t = k * 0.5 + # field falls, kinetic rises by the same amount -> total conserved. + fld.append(f"{k * 10} {t} {1.0 - 0.1 * k} 0 0 0 0 0") + par.append(f"{k * 10} {t} {0.1 * k}") + (hist / "fld_ene").write_text("\n".join(fld) + "\n") + (hist / "par01_ene").write_text("\n".join(par) + "\n") + + +def _make_full_run(tmp_path: Path, n_steps: int = 5, nx: int = 16) -> Path: + """A run with fields, a density moment, a phase space, and HIST energy.""" + run_dir = _make_run(tmp_path, n_steps=n_steps, nx=nx) # FLD/e1 + _write_field(run_dir, "e2", n_steps, nx, seed=2) + _write_field(run_dir, "b3", n_steps, nx, seed=3) + rng = np.random.default_rng(11) + dens = run_dir / "MS" / "DENSITY" / "electrons" / "charge" + for k in range(n_steps): + it = k * 10 + _write_dump( + dens / f"charge-{it:06d}.h5", + "charge", + rng.standard_normal(nx), + t=k * 0.5, + it=it, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + _write_phasespace(run_dir, n_steps, nx, npmom=6) + _write_hist_energy(run_dir, n_steps) + return run_dir + + +def test_load_series_nc_roundtrips_grid_series(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=4, nx=8) + out = tmp_path / "binary" + oio.save_run_datasets(run_dir, out) + + from_nc = oio.load_series_nc(out / "FLD" / "e1.nc") + from_ms = oio.load_series(run_dir / "MS" / "FLD" / "e1") + + assert from_nc.dims == from_ms.dims + assert from_nc.shape == from_ms.shape + np.testing.assert_allclose(from_nc.values, from_ms.values) + np.testing.assert_allclose(from_nc["t"].values, from_ms["t"].values) + assert list(from_nc["iter"].values) == list(from_ms["iter"].values) + # axis metadata is rebuilt into the dict attrs the plotters read. + assert from_nc.attrs["axis_units"]["x1"] == from_ms.attrs["axis_units"]["x1"] + # load_series dispatches to load_series_nc when handed a .nc file. + np.testing.assert_allclose( + oio.load_series(out / "FLD" / "e1.nc").values, from_ms.values + ) + + +def test_save_run_datasets_persists_hist_energy(tmp_path: Path) -> None: + run_dir = _make_full_run(tmp_path) + out = tmp_path / "binary" + written = oio.save_run_datasets(run_dir, out) + + nc = out / "HIST" / "energy.nc" + assert nc.exists() + assert nc in written + # load_hist_energy reads the saved NetCDF when there's no raw HIST/ ASCII. + energy = oio.load_hist_energy(out) + assert energy is not None + assert "total" in energy + assert energy.attrs.get("total_drift_frac") == 0.0 # conserved by construction + + +def test_save_canned_plots_regenerates_from_netcdf(tmp_path: Path) -> None: + from adept.osiris import plots as oplots + + run_dir = _make_full_run(tmp_path) + + # Plots from the raw OSIRIS run... + out_ms = tmp_path / "plots_ms" + written_ms = oplots.save_canned_plots(run_dir, out_ms) + + # ...vs plots regenerated from the saved NetCDFs alone (no MS/ tree). + binary = tmp_path / "binary" + oio.save_run_datasets(run_dir, binary) + assert not (binary / "MS").exists() + out_nc = tmp_path / "plots_nc" + written_nc = oplots.save_canned_plots(binary, out_nc) + + # The NetCDF-only regeneration reproduces the exact same plot set. + assert set(written_nc) == set(written_ms) + # ...spanning every family, incl. the energy traces that read outside MS/. + for key in ( + "spacetime/e1", + "omega_k/e1", + "moments/electrons/charge", + "profiles/electrons/density", + "phasespace/electrons/p1x1", + "phasespace_evolution/electrons/p1x1", + "energy_vs_time", + "energy_components_vs_time", + "total_energy_vs_time", + ): + assert key in written_nc, f"missing {key}" + assert written_nc[key].exists() From 08a2422480a827b036bc851cf91c7f72eb66ec14 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 2 Jun 2026 21:56:43 -0700 Subject: [PATCH 10/49] feat(osiris): NetCDF plot regeneration + plot-set refinements - regen harness: rebuild the full canned plot set from saved NetCDFs (no rerun) - f(p) and delta-f lineouts; temperature profile from phase-space Maxwellian fits - number-density profiles (initial/final/late-mean); 2-panel equal-aspect omega-k - phase-space & spacetime: space on x-axis, cropped to box, log-contrast floor - proper-LaTeX titles (prose vs math) Co-Authored-By: Claude Opus 4.8 --- adept/osiris/plots.py | 572 ++++++++++++++++++---- adept/osiris/post.py | 7 +- adept/osiris/regen.py | 161 ++++++ configs/osiris/twostream-1d.yaml | 2 +- docs/osiris-adept-usage.md | 6 +- tests/test_osiris/test_plots_new_views.py | 7 +- 6 files changed, 639 insertions(+), 116 deletions(-) create mode 100644 adept/osiris/regen.py diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index e76627fb..b7b6bf8e 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -11,18 +11,20 @@ spacetime/.png (t, x) heatmap of every FLD diagnostic spacetime_log/.png log10|·| of the same lineouts/.png value-vs-x snapshots at sampled times - omega_k/.png 2-D FFT (k, ω) dispersion plot - omega_k_zoom/.png same, zoomed to the ω = k line (plasma waves) + omega_k/.png 2-D FFT (k, ω) dispersion: full range + + equal-aspect square window (ω = k at 45°) currents/spacetime.png j1/j2/j3 (J_x/J_y/J_z) side-by-side spacetime currents/lineouts.png j1/j2/j3 profiles vs x (final + late mean) moments//.png (t, x) per-species density moments moments//_log.png log10 of the same moments//lineouts/.png moment snapshots at sampled times - profiles//density.png density profile vs x (final + late mean) - profiles//temperature.png temperature profile vs x (if thermal moments) + profiles//density.png number-density profile vs x (initial + final + late mean) + profiles//temperature.png T(x) from a Maxwellian fit to the phase space + (initial + final + late mean); else uth / T_ii moments phasespace//.png final-step (x, p) heatmap per species phasespace_evolution//.png (x, p) heatmaps at sampled times distribution_lineouts//.png f(p) averaged over the right-boundary cells + deltaf_lineouts//.png delta-f = f - (fitted Maxwellian), same averaging field_decomp/.png left/right-going transverse E (Riemann split) energy_vs_time.png total field energy time-trace energy_components_vs_time.png E-field / B-field / total energy @@ -57,12 +59,17 @@ def plot_spacetime( *, log: bool = False, cmap: str = "RdBu_r", + space_on_x: bool = True, title: str | None = None, ) -> plt.Axes: """``(t, x)`` heatmap of a 1D-field time-series. Accepts a pre-loaded ``DataArray`` (dims must include ``t`` and one spatial axis) or a directory of dumps that ``load_series`` will eat. + + By default space is on the horizontal axis and time on the vertical + (``t`` vs ``x``); set ``space_on_x=False`` to transpose so time is on the + horizontal axis and space on the vertical (``x`` vs ``t``). """ da = _ensure_series(series) if da.ndim != 2: @@ -71,20 +78,27 @@ def plot_spacetime( ) if ax is None: _, ax = plt.subplots(figsize=(6, 4)) - data = np.log10(np.abs(da.values) + 1e-30) if log else da.values + data = np.log10(np.abs(da.values) + 1e-30) if log else da.values # (t, x) t = da.coords["t"].values xname = next(d for d in da.dims if d != "t") x = da.coords[xname].values - mesh = ax.pcolormesh(t, x, data.T, shading="auto", cmap=cmap) + if space_on_x: + mesh = ax.pcolormesh(x, t, data, shading="auto", cmap=cmap) + ax.set_xlabel(_axis_label(da, xname)) + ax.set_ylabel(_axis_label(da, "t")) + orient = r"$t$ vs $x$" + else: + mesh = ax.pcolormesh(t, x, data.T, shading="auto", cmap=cmap) + ax.set_xlabel(_axis_label(da, "t")) + ax.set_ylabel(_axis_label(da, xname)) + orient = r"$x$ vs $t$" plt.colorbar( mesh, ax=ax, label=rf"$\log_{{10}}$ |{_value_label(da)}|" if log else _value_label(da), ) - ax.set_xlabel(_axis_label(da, "t")) - ax.set_ylabel(_axis_label(da, xname)) scale = r"$\log_{10}$ " if log else "" - ax.set_title(title or f"{_display_name(da)} — {scale}spacetime ($x$ vs $t$)") + ax.set_title(title or f"{_display_name(da)} — {scale}spacetime ({orient})") return ax @@ -128,26 +142,41 @@ def plot_phasespace( cmap: str = "viridis", title: str | None = None, ) -> plt.Axes: - """Heatmap of a 2D phase-space density (e.g. ``x1p1``).""" + """Heatmap of a 2D phase-space density (e.g. ``x1p1``). + + For an x-p phase space the spatial axis is cropped to the physical box + (``sim.XMAX``; phase-space dumps may pad it out past the box) and drawn on + the horizontal axis with momentum on the vertical, matching the orientation + of :func:`plot_phasespace_evolution`. + """ if not isinstance(da, xr.DataArray): da = _io.load_phasespace_h5(da) if da.ndim != 2: raise ValueError( f"plot_phasespace expects 2D data; got {da.dims} {da.shape}" ) + da = _crop_spatial_to_box(da) if ax is None: _, ax = plt.subplots(figsize=(5, 4)) - arr = da.values + # Space on the horizontal axis, momentum on the vertical; fall back to dim + # order for momentum-momentum spaces (e.g. p1p2) with no spatial axis. + spatial = [d for d in da.dims if str(d).startswith("x")] + moment = [d for d in da.dims if str(d).startswith("p")] + xdim, ydim = (spatial[0], moment[0]) if spatial and moment else da.dims + raw = da.transpose(ydim, xdim).values + vmin = vmax = None if log: - arr = np.log10(np.abs(arr) + 1e-30) - d0, d1 = da.dims - x0, x1 = da.coords[d0].values, da.coords[d1].values - mesh = ax.pcolormesh(x0, x1, arr.T, shading="auto", cmap=cmap) + vmin, vmax = _nonzero_log_clim(raw) + plot_arr = np.log10(np.abs(raw) + 1e-30) + else: + plot_arr = raw + xc, yc = da.coords[xdim].values, da.coords[ydim].values + mesh = ax.pcolormesh(xc, yc, plot_arr, shading="auto", cmap=cmap, vmin=vmin, vmax=vmax) plt.colorbar( mesh, ax=ax, label=rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da) ) - ax.set_xlabel(_axis_label(da, d0)) - ax.set_ylabel(_axis_label(da, d1)) + ax.set_xlabel(_axis_label(da, xdim)) + ax.set_ylabel(_axis_label(da, ydim)) t = da.attrs.get("time", float("nan")) scale = r"$\log_{10}$ " if log else "" ax.set_title( @@ -178,6 +207,7 @@ def plot_phasespace_evolution( raise ValueError( f"plot_phasespace_evolution expects (t, p, x); got dims {da.dims}" ) + da = _crop_spatial_to_box(da) nt = da.coords["t"].size t_skip = max(1, nt // n_panels) sl = da.isel(t=slice(0, None, t_skip)) @@ -188,6 +218,11 @@ def plot_phasespace_evolution( facet_kw: dict = {"cmap": cmap} if spatial and moment: facet_kw.update(x=spatial[0], y=moment[0]) + if log: + # Floor the shared colour scale at the lowest non-zero value so empty + # cells (log -> -30) don't crush the contrast across the facets. + vmin, vmax = _nonzero_log_clim(sl.values) + facet_kw.update(vmin=vmin, vmax=vmax) g = plot_da.plot(col="t", col_wrap=min(col_wrap, sl.coords["t"].size), **facet_kw) if spatial and moment: g.set_xlabels(_axis_label(da, spatial[0])) @@ -385,6 +420,7 @@ def plot_omega_k( omega_p: float = 1.0, k_max: float | None = None, omega_max: float | None = None, + equal_aspect: bool = False, title: str | None = None, ) -> plt.Axes: """2-D FFT of a ``(t, x)`` field series; show power in ``(k, ω)`` space. @@ -457,12 +493,178 @@ def plot_omega_k( ax.axhline(0, color="w", lw=0.4, alpha=0.4) ax.axvline(0, color="w", lw=0.4, alpha=0.4) + if equal_aspect: + # One unit of k displays as one unit of ω, so the light line ω = k is a + # true 45° slope (lets you read off where power sits relative to it). + ax.set_aspect("equal", adjustable="box") ax.set_xlabel(r"$k\ [\omega_p / c]$") ax.set_ylabel(r"$\omega\ [\omega_p]$") ax.set_title(title or rf"{_display_name(da)} — $(k, \omega)$ power spectrum") return ax +def plot_omega_k_figure( + series: xr.DataArray | str | Path, + *, + v_th: float | None = None, + omega_k_zoom: float | None = 4.0, + cmap: str = "magma", + log: bool = True, +) -> plt.Figure: + """Stacked ``(k, ω)`` spectra: full range on top, equal-aspect below. + + The top panel is the full 2-D FFT over the data's Nyquist range. The bottom + panel shows the same power with k and ω on the same scale (equal aspect, so + the light line ``ω = k`` is a true 45° slope), limited to a square + ``±z`` window with ``z = min(omega_k_zoom, Nyquist)`` so the slope-1 line — + and any low-frequency power along it — is visible. + + With a coarse dump cadence the ω-Nyquist is small, so that window is only a + few k-cells wide and the lower panel looks blocky; that is the expected + consequence of the time-sampling, not a plotting fault. + """ + da = _ensure_series(series) + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(6, 10)) + plot_omega_k( + da, ax=ax_top, log=log, cmap=cmap, + show_langmuir=v_th is not None, v_th=v_th, + ) + z = _omega_k_zoom_window(da, omega_k_zoom) + plot_omega_k( + da, ax=ax_bot, log=log, cmap=cmap, + show_light_line=True, show_langmuir=v_th is not None, v_th=v_th, + k_max=z, omega_max=z, equal_aspect=True, + title=rf"{_display_name(da)} — $(k, \omega)$ (equal aspect)", + ) + fig.tight_layout() + return fig + + +def _sim_box_bound(da: xr.DataArray, xdim: str, *, upper: bool) -> float: + """Edge of the *physical* simulation box along ``xdim`` (``sim.XMIN/XMAX``). + + Phase-space diagnostics may be binned over a spatial range wider than the + field grid (deck ``ps_xmax`` > ``xmax``), so the coordinate axis runs past + the box. ``sim.XMIN`` / ``sim.XMAX`` (carried on every diagnostic) record + the true box edges — per spatial dimension, so ``x1`` -> entry 0, + ``x2`` -> entry 1, …. Falls back to the matching axis end when the attr is + absent, which makes callers degrade to the full axis exactly when it + already coincides with the box. + """ + xv = da.coords[xdim].values + bound = da.attrs.get("sim.XMAX" if upper else "sim.XMIN") + if bound is None: + return float(xv[-1] if upper else xv[0]) + if isinstance(bound, (list, tuple, np.ndarray)): + digits = "".join(ch for ch in str(xdim) if ch.isdigit()) + i = int(digits) - 1 if digits else 0 + bound = bound[i] if 0 <= i < len(bound) else bound[-1] + return float(bound) + + +def _sim_box_xmax(da: xr.DataArray, xdim: str) -> float: + """Right edge of the physical box along ``xdim`` (see :func:`_sim_box_bound`).""" + return _sim_box_bound(da, xdim, upper=True) + + +def _nonzero_log_clim(values) -> tuple[float | None, float | None]: + """``(vmin, vmax)`` for a ``log10`` heatmap, floored at the lowest non-zero. + + Phase-space dumps are mostly empty (zero) cells; mapping those through + ``log10`` sends them to a huge negative floor that dominates the colour + range and washes out the real distribution. Restricting ``vmin`` to the + lowest *non-zero* magnitude (and ``vmax`` to the largest) keeps the dynamic + range on the populated cells. Returns ``(None, None)`` for all-zero input. + """ + mag = np.abs(np.asarray(values)) + nz = mag > 0 + if not nz.any(): + return None, None + return float(np.log10(mag[nz].min())), float(np.log10(mag.max())) + + +def _crop_spatial_to_box(da: xr.DataArray) -> xr.DataArray: + """Trim spatial (``x*``) dims to the physical box ``[sim.XMIN, sim.XMAX]``. + + Phase-space dumps can be binned over a spatial range wider than the box + (deck ``ps_xmax`` > ``xmax``), padding the high-``x`` end of the axis with + empty cells. This keeps only the cells within the simulation domain so a + phase-space plot's spatial axis spans just the box. Returns ``da`` unchanged + when no spatial dim extends past the box (e.g. ``p1p2`` momentum spaces). + """ + for d in list(da.dims): + if not str(d).startswith("x"): + continue + xv = da.coords[d].values + if xv.size < 2: + continue + left = int(np.searchsorted(xv, _sim_box_bound(da, d, upper=False), side="left")) + right = int(np.searchsorted(xv, _sim_box_bound(da, d, upper=True), side="right")) + right = max(left + 1, min(right, xv.size)) + if left > 0 or right < xv.size: + da = da.isel({d: slice(left, right)}) + return da + + +def _maxwellian_moments(p, f) -> tuple[float, float, float] | None: + r"""Density, mean momentum, and variance of a 1-D distribution ``f(p)``. + + Returns ``(W, p0, var)`` with ``W = sum(f)`` (negative weights clipped to + zero), ``p0 =

`` and ``var = <(p - p0)^2>`` — the parameters of the + moment-matched Maxwellian, ``var`` being the kinetic temperature ``T`` for + ``m = 1`` (electrons, ``p`` in ``m_e c``). ``None`` if ``f`` has no positive + weight. + """ + f = np.clip(np.asarray(f, dtype=float), 0.0, None) + W = float(f.sum()) + if W <= 0: + return None + p = np.asarray(p, dtype=float) + p0 = float((f * p).sum() / W) + var = float((f * (p - p0) ** 2).sum() / W) + return W, p0, var + + +def _maxwellian_from_moments(p, W: float, p0: float, var: float): + """Gaussian on grid ``p`` with the same sum ``W``, mean ``p0`` and ``var``.""" + p = np.asarray(p, dtype=float) + if var <= 0: + return np.zeros_like(p) + g = np.exp(-((p - p0) ** 2) / (2.0 * var)) + s = float(g.sum()) + return g * (W / s) if s > 0 else g + + +def _boundary_distribution(da: xr.DataArray, n_cells: int): + r"""``f(p)`` averaged over the ``n_cells`` cells inside the right box edge. + + Shared by the ``f(p)`` and ``\delta f`` lineouts. Returns + ``(f_xp, da, pdim, x_lo, x_hi, k)``: the decorated input ``da``, the + sign-corrected (non-negative) distribution ``f_xp`` with dims ``(t?, p)``, + the momentum axis name ``pdim``, the averaged spatial window ``[x_lo, + x_hi]`` and the cell count ``k``. The window is anchored at the physical box + edge (``sim.XMAX``) so the empty ``ps_xmax`` padding beyond the box is + excluded. + """ + da = _decorate(da) + spatial = [d for d in da.dims if d != "t" and str(d).startswith("x")] + moment = [d for d in da.dims if d != "t" and str(d).startswith("p")] + if not spatial or not moment: + raise ValueError(f"need an x* and a p* dim; got dims {da.dims}") + xdim, pdim = spatial[0], moment[0] + nx = da.sizes[xdim] + xv = da.coords[xdim].values + right = int(np.searchsorted(xv, _sim_box_xmax(da, xdim), side="right")) - 1 + right = min(max(right, 0), nx - 1) + k = max(1, min(n_cells, right + 1)) + lo = right - k + 1 + f_xp = da.isel({xdim: slice(lo, right + 1)}).mean(dim=xdim) + # Charge density q*f is single-signed per species (negative for electrons); + # divide out the sign so f(p) is non-negative. + q_sign = np.sign(float(np.nansum(da.values))) or 1.0 + return f_xp / q_sign, da, pdim, float(xv[lo]), float(xv[right]), k + + def plot_distribution_lineout( series: xr.DataArray | str | Path, ax: plt.Axes | None = None, @@ -473,32 +675,24 @@ def plot_distribution_lineout( cmap: str = "viridis", title: str | None = None, ) -> plt.Axes: - """Velocity distribution ``f(p)`` averaged over the rightmost ``n_cells``. + """Distribution function ``f(p)`` near the right boundary. Takes a phase-space series with a spatial (``x*``) and a momentum (``p*``) axis — either a stacked ``(t, p, x)`` series from :func:`io.load_series` or a single ``(p, x)`` snapshot — averages it over - the ``n_cells`` cells nearest the right-hand boundary to get ``f(p)`` - there, and overlays that lineout at ~``n_times`` sampled times (colour = - time). This is the standard "what does the distribution look like as it - leaves the box" diagnostic for driven / open-boundary runs. + the ``n_cells`` cells just inside the right edge of the *physical box* + (``sim.XMAX``; the phase-space axis may extend past the box when the deck's + ``ps_xmax`` exceeds ``xmax``), and overlays that lineout at ~``n_times`` + sampled times (colour = time). This is the standard "what does the + distribution look like as it leaves the box" diagnostic for driven / + open-boundary runs. + + OSIRIS species phase spaces store the charge density ``q*f``; the charge + sign is divided out so the curve reads as a non-negative ``f(p)``. """ da = series if isinstance(series, xr.DataArray) else _io.load_series(series) - da = _decorate(da) - spatial = [d for d in da.dims if d != "t" and str(d).startswith("x")] - moment = [d for d in da.dims if d != "t" and str(d).startswith("p")] - if not spatial or not moment: - raise ValueError( - f"plot_distribution_lineout needs an x* and a p* dim; got {da.dims}" - ) - xdim, pdim = spatial[0], moment[0] - nx = da.sizes[xdim] - k = max(1, min(n_cells, nx)) - # Average over the k rightmost spatial cells -> f over the momentum axis. - f_xp = da.isel({xdim: slice(nx - k, nx)}).mean(dim=xdim) + f_xp, da, pdim, x_lo, x_hi, k = _boundary_distribution(da, n_cells) p = f_xp.coords[pdim].values - xv = da.coords[xdim].values - x_lo, x_hi = float(xv[nx - k]), float(xv[-1]) if ax is None: _, ax = plt.subplots(figsize=(6, 4)) @@ -525,8 +719,63 @@ def plot_distribution_lineout( ax.set_ylabel(ylabel) ax.set_title( title - or rf"{_display_name(da)} — $f$ over rightmost {k} cells " - rf"($x \in [{x_lo:.3g}, {x_hi:.3g}]$)" + or rf"{_display_name(da)} — $f(p)$ over {k} cells inside the right boundary " + rf"($x \in [{x_lo:.4g}, {x_hi:.4g}]$)" + ) + ax.grid(True, alpha=0.3) + return ax + + +def plot_deltaf_lineout( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + n_cells: int = 10, + n_times: int = 6, + cmap: str = "viridis", + title: str | None = None, +) -> plt.Axes: + r"""``\delta f = f - f_M`` near the right boundary, overlaid at sampled times. + + Averages the phase space over the ``n_cells`` cells inside the right box + edge to get ``f(p)`` (as :func:`plot_distribution_lineout`), then at each + sampled time subtracts the moment-matched Maxwellian ``f_M`` — the Gaussian + with the same density, mean momentum and temperature — and overlays the + residual. Since the species starts Maxwellian, the residual isolates the + non-thermal structure that grows away from it: beams and hot tails. + """ + da = series if isinstance(series, xr.DataArray) else _io.load_series(series) + f_xp, da, pdim, x_lo, x_hi, k = _boundary_distribution(da, n_cells) + p = f_xp.coords[pdim].values + + def _residual(fvals): + m = _maxwellian_moments(p, fvals) + return fvals if m is None else fvals - _maxwellian_from_moments(p, *m) + + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + + if "t" in f_xp.dims: + nt = f_xp.sizes["t"] + idx = np.unique(np.linspace(0, nt - 1, min(n_times, nt)).astype(int)) + colours = plt.get_cmap(cmap)(np.linspace(0.15, 0.95, len(idx))) + tvals = f_xp.coords["t"].values + for colour, i in zip(colours, idx, strict=False): + ax.plot(p, _residual(f_xp.isel(t=i).values), color=colour, lw=1.3, + label=f"{float(tvals[i]):.2g}") + ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) + else: + ax.plot(p, _residual(f_xp.values), lw=1.4) + + ax.axhline(0.0, color="0.6", lw=0.6) + units = da.attrs.get("units", "") + ylabel = r"$\langle \delta f \rangle_x$" + (f" [{_tex(units)}]" if units else "") + ax.set_xlabel(_axis_label(da, pdim)) + ax.set_ylabel(ylabel) + ax.set_title( + title + or rf"{_display_name(da)} — $\delta f = f - f_M$ over {k} cells inside the " + rf"right boundary ($x \in [{x_lo:.4g}, {x_hi:.4g}]$)" ) ax.grid(True, alpha=0.3) return ax @@ -566,7 +815,7 @@ def plot_currents_spacetime(run_dir: str | Path) -> plt.Figure | None: fig, axes = plt.subplots(1, len(comps), figsize=(5 * len(comps), 4), squeeze=False) for ax, (comp, ser) in zip(axes[0], comps.items(), strict=False): plot_spacetime(ser, ax=ax) - fig.suptitle(r"Current density components — spacetime ($x$ vs $t$)", y=1.03) + fig.suptitle(r"Current density components — spacetime ($t$ vs $x$)", y=1.03) fig.tight_layout() return fig @@ -623,14 +872,36 @@ def _species_diags(run_dir: str | Path) -> dict[str, list[tuple[str, str, Path]] def _density_series(entries: list[tuple[str, str, Path]]) -> xr.DataArray | None: - """Pick the best density-like moment for a species and load it as ``(t, x)``.""" + """Best density-like ``(t, x)`` moment for a species, as a number density. + + Prefers an explicit number-density report; otherwise falls back to the + charge density and divides out the species charge sign (OSIRIS reports + charge density ``q*n``, negative for electrons) so the result is a + non-negative number density ``n``. Returns ``None`` if no spatial density + moment is present. + """ by_q = {q: p for kind, q, p in entries if kind == "DENSITY"} - for pref in ("charge", "n", "n01", "m", "ene"): - if pref in by_q: - ser = _io.load_series(by_q[pref]) - return ser if ser.ndim == 2 else None - # Fall back to any DENSITY entry. - for kind, q, p in entries: + for pref in ("n", "n01", "n02", "charge", "m"): + if pref not in by_q: + continue + ser = _io.load_series(by_q[pref]) + if ser.ndim != 2: + continue + if pref == "charge": + q_sign = np.sign(float(np.nansum(ser.values))) or 1.0 + attrs, name = dict(ser.attrs), ser.name + ser = ser / q_sign + ser.attrs, ser.name = attrs, name + ser.attrs["long_name"] = "n" # now a number density, not charge + # Charge-density units lead with the charge ``e``; drop it so the + # label reads as a number density (the magnitudes match in code + # units where e = 1). + units = str(ser.attrs.get("units", "")) + if units.startswith("e "): + ser.attrs["units"] = units[2:].lstrip() + return ser + # Fall back to any DENSITY entry as-is. + for kind, _q, p in entries: if kind == "DENSITY": ser = _io.load_series(p) if ser.ndim == 2: @@ -679,25 +950,72 @@ def _temperature_series(entries: list[tuple[str, str, Path]]) -> xr.DataArray | return out +def _species_phasespace(diags: dict[str, Path], species: str) -> Path | None: + """Path to an x-p phase space for ``species`` (e.g. ``PHA/p1x1/``).""" + for rel, path in sorted(diags.items()): + parts = rel.split("/") + if len(parts) >= 3 and parts[0] == "PHA" and parts[-1] == species: + ps = parts[1] # e.g. "p1x1" — needs both a space and a momentum axis + if "x" in ps and "p" in ps: + return path + return None + + +def _temperature_from_phasespace(series: xr.DataArray | str | Path | None) -> xr.DataArray | None: + r"""Temperature profile ``T(t, x)`` from an x-p phase space. + + Fits a moment-matched Maxwellian in ``p`` at every ``(t, x)`` and returns + its temperature ``T = <(p -

)^2>`` (the variance of the momentum; + ``m_e c^2`` units with ``m = 1``) as a ``(t, x)`` DataArray, cropped to the + physical box. The charge sign is divided out so the weights are + non-negative. ``None`` if the input is missing or is not an x-p phase space. + """ + if series is None: + return None + ser = _decorate(series if isinstance(series, xr.DataArray) else _io.load_series(series)) + spatial = [d for d in ser.dims if d != "t" and str(d).startswith("x")] + moment = [d for d in ser.dims if d != "t" and str(d).startswith("p")] + if not spatial or not moment or "t" not in ser.dims: + return None + ser = _crop_spatial_to_box(ser) + pdim = moment[0] + pc = ser.coords[pdim] + q_sign = np.sign(float(np.nansum(ser.values))) or 1.0 + w = (ser / q_sign).clip(min=0.0) # non-negative weights f(t, p, x) + weight = w.sum(pdim) + p0 = (w * pc).sum(pdim) / weight + var = (w * (pc - p0) ** 2).sum(pdim) / weight # (t, x) = T + out = var.where(weight > 0) + out.attrs = {k: v for k, v in ser.attrs.items() if k != "units"} + out.attrs["long_name"] = "T" + out.attrs["units"] = r"m_e c^2" + out.name = "temperature" + return out + + def plot_profile( series: xr.DataArray | str | Path, ax: plt.Axes | None = None, *, abs_value: bool = False, n_avg_frac: float = 0.2, + show_initial: bool = False, value_label: str | None = None, title: str | None = None, ) -> plt.Axes: """Plot a ``(t, x)`` moment vs ``x``: final snapshot plus late-time mean. The late-time mean (over the last ``n_avg_frac`` of the dumps) smooths - out the time-dependent fluctuations to show the established profile. + out the time-dependent fluctuations to show the established profile. With + ``show_initial`` the ``t = 0`` profile is overlaid too, so the change from + the initial state is visible. """ da = _decorate(series if isinstance(series, xr.DataArray) else _io.load_series(series)) if "t" not in da.dims: raise ValueError(f"plot_profile expects a (t, x) series; got dims {da.dims}") xdim = next(d for d in da.dims if d != "t") x = da.coords[xdim].values + tvals = da.coords["t"].values nt = da.sizes["t"] w = max(1, round(n_avg_frac * nt)) if ax is None: @@ -706,9 +1024,16 @@ def plot_profile( mean = da.isel(t=slice(nt - w, nt)).mean("t").values if abs_value: final, mean = np.abs(final), np.abs(mean) - tlast = float(da.coords["t"].values[-1]) - ax.plot(x, final, lw=1.4, label=f"final ($t={tlast:.3g}$)") - ax.plot(x, mean, lw=1.1, ls="--", label=f"mean of last {w} dumps") + if show_initial: + initial = da.isel(t=0).values + if abs_value: + initial = np.abs(initial) + # Dotted and on top (high zorder) so the initial profile stays visible + # where the final / mean curves overlap it. + ax.plot(x, initial, lw=1.6, ls=":", color="k", zorder=3, + label=f"initial ($t={float(tvals[0]):.3g}$)") + ax.plot(x, final, lw=1.4, zorder=2, label=f"final ($t={float(tvals[-1]):.3g}$)") + ax.plot(x, mean, lw=1.1, ls="--", zorder=2.2, label=f"mean of last {w} dumps") ax.set_xlabel(_axis_label(da, xdim)) ax.set_ylabel(value_label or _value_label(da)) ax.set_title(title or rf"{_display_name(da)} — profile vs $x$") @@ -775,38 +1100,24 @@ def _pack(values: xr.DataArray, src: xr.DataArray, name: str, long: str) -> xr.D return out -def plot_field_lr_decomposition( - run_dir: str | Path, - *, - v_th: float | None = None, - omega_k_zoom: float | None = 4.0, -) -> dict[str, plt.Figure]: - """One 2x2 figure per transverse component: right/left spacetime + ω-k. +def plot_field_lr_decomposition(run_dir: str | Path) -> dict[str, plt.Figure]: + """One figure per transverse component: right/left-going spacetime. - Top row is the right/left-going ``(t, x)`` spacetime; bottom row is each - part's ``(k, ω)`` power spectrum (zoomed, with the light line drawn) so - you can confirm the right-going part really does sit on the ``ω·k > 0`` - branches and the left-going part on ``ω·k < 0``. + Two panels side by side — the right-going and left-going parts of the + transverse E field as spacetime heatmaps with space on the horizontal axis + and time on the vertical (``t`` vs ``x``). """ figs: dict[str, plt.Figure] = {} for comp, parts in efield_lr_components(run_dir).items(): - fig, axes = plt.subplots(2, 2, figsize=(11, 8)) - for col, side in enumerate(("right", "left")): + fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True) + for ax, side in zip(axes, ("right", "left"), strict=False): da = parts[side] - plot_spacetime(da, ax=axes[0, col], title=f"{_tex(_long_name(da))} — spacetime") - try: - zoom = _omega_k_zoom_window(da, omega_k_zoom) - plot_omega_k( - da, ax=axes[1, col], show_light_line=True, - show_langmuir=v_th is not None, v_th=v_th, - k_max=zoom, omega_max=zoom, - title=f"{_tex(_long_name(da))} — $(k, \\omega)$", - ) - except Exception as e: # ω-k overlay is best effort - axes[1, col].set_visible(False) - print(f"[plots] ω-k for {comp} {side} skipped: {e}") + plot_spacetime( + da, ax=ax, space_on_x=True, + title=f"{_display_name(da)} — spacetime", + ) fig.suptitle( - rf"{comp}: left/right-going decomposition (vacuum Riemann split)", y=1.01 + rf"{comp}: left/right-going decomposition (vacuum Riemann split)", y=1.02 ) fig.tight_layout() figs[comp] = fig @@ -816,6 +1127,24 @@ def plot_field_lr_decomposition( # --- driver --------------------------------------------------------------- +def canned_plot_kwargs(output_cfg: dict | None) -> dict: + """Translate a manifest ``output:`` block into :func:`save_canned_plots` kwargs. + + Shared by the live post-processing path (``post.collect``) and the offline + regeneration harness (``regen``) so both honour the same knobs: + ``v_th``, ``dist_cells`` and ``omega_k_zoom`` (which may be explicitly + ``null`` to disable the zoom). Keys that are absent fall through to the + ``save_canned_plots`` defaults. + """ + output_cfg = output_cfg or {} + kwargs: dict = {"v_th": output_cfg.get("v_th")} + if output_cfg.get("dist_cells") is not None: + kwargs["dist_cells"] = int(output_cfg["dist_cells"]) + if "omega_k_zoom" in output_cfg: # may be explicitly null to disable zoom + kwargs["omega_k_zoom"] = output_cfg["omega_k_zoom"] + return kwargs + + def _omega_k_zoom_window(series: xr.DataArray, requested: float | None) -> float | None: """Clamp a requested ``(k, ω)`` zoom half-width to the data's Nyquist range. @@ -863,8 +1192,9 @@ def save_canned_plots( ``dist_cells`` sets how many right-boundary cells the phase-space distribution lineouts average over; ``omega_k_zoom`` is the ``(k, ω)`` - half-width (in ``ω_p`` units) for the zoomed dispersion plots that show the - whole ``ω = k`` line (``None`` → full spectrum). + half-width (in ``ω_p`` units) for the equal-aspect lower panel of the + dispersion plots, where ``ω = k`` is drawn at 45° (clamped to the data's + Nyquist; ``None`` → full Nyquist window). """ run_dir = Path(run_dir) out_dir = Path(out_dir) @@ -905,21 +1235,11 @@ def _write(fig: plt.Figure, rel: str) -> Path: plot_lineouts(ser, n_panels=n_panels), f"lineouts/{comp}.png" ) - fig, ax = plt.subplots(figsize=(6, 5)) - plot_omega_k(ser, ax=ax, show_langmuir=v_th is not None, v_th=v_th) - written[f"omega_k/{comp}"] = _write(fig, f"omega_k/{comp}.png") - - # Zoomed (k, ω) view sized so the whole ω = k line is visible — the - # window where the plasma (Langmuir) waves live. - zoom = _omega_k_zoom_window(ser, omega_k_zoom) - fig, ax = plt.subplots(figsize=(6, 5)) - plot_omega_k( - ser, ax=ax, show_light_line=True, - show_langmuir=v_th is not None, v_th=v_th, - k_max=zoom, omega_max=zoom, - title=f"{_display_name(ser)} — $(k, \\omega)$ (zoom)", + # Full (k, ω) spectrum on top, equal-aspect square window below. + written[f"omega_k/{comp}"] = _write( + plot_omega_k_figure(ser, v_th=v_th, omega_k_zoom=omega_k_zoom), + f"omega_k/{comp}.png", ) - written[f"omega_k_zoom/{comp}"] = _write(fig, f"omega_k_zoom/{comp}.png") # --- Currents (j1/j2/j3) combined views --- try: @@ -967,21 +1287,33 @@ def _write(fig: plt.Figure, rel: str) -> Path: # --- Per-species density + temperature profiles --- for species, entries in sorted(_species_diags(run_dir).items()): + label = species.replace("_", " ") # "species_1" -> "species 1" for titles try: dens = _density_series(entries) if dens is not None: fig, ax = plt.subplots(figsize=(6, 4)) - plot_profile(dens, ax=ax, title=f"{species} — density profile") + plot_profile( + dens, ax=ax, show_initial=True, + title=f"{label} — density profile", + ) written[f"profiles/{species}/density"] = _write( fig, f"profiles/{species}/density.png" ) except Exception as e: print(f"[plots] skipping density profile for {species}: {e}") try: - temp = _temperature_series(entries) + # Temperature from a Maxwellian fit to the species phase space + # (preferred); fall back to dumped thermal moments (uth / T_ii). + ps_path = _species_phasespace(diags, species) + temp = _temperature_from_phasespace(ps_path) + if temp is None: + temp = _temperature_series(entries) if temp is not None: fig, ax = plt.subplots(figsize=(6, 4)) - plot_profile(temp, ax=ax, title=f"{species} — temperature profile") + plot_profile( + temp, ax=ax, show_initial=True, + title=f"{label} — temperature profile", + ) written[f"profiles/{species}/temperature"] = _write( fig, f"profiles/{species}/temperature.png" ) @@ -1018,20 +1350,32 @@ def _write(fig: plt.Figure, rel: str) -> Path: plot_phasespace_evolution(ser, n_panels=n_panels), f"phasespace_evolution/{species}/{ps_name}.png", ) - # f(p) averaged over the right-boundary cells, vs time. + except Exception as e: + print(f"[plots] skipping phasespace evolution {diag_rel}: {e}") + + # f(p) and delta-f averaged over the right-boundary cells, vs time — + # only for x-p phase spaces (a p-p space like p1p2 has no spatial axis). + if ser.ndim == 3 and any(str(d).startswith("x") for d in ser.dims if d != "t"): + try: fig, ax = plt.subplots(figsize=(6, 4)) plot_distribution_lineout(ser, ax=ax, n_cells=dist_cells) written[f"distribution_lineouts/{species}/{ps_name}"] = _write( fig, f"distribution_lineouts/{species}/{ps_name}.png" ) - except Exception as e: - print(f"[plots] skipping phasespace evolution {diag_rel}: {e}") + except Exception as e: + print(f"[plots] skipping distribution lineout {diag_rel}: {e}") + try: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_deltaf_lineout(ser, ax=ax, n_cells=dist_cells) + written[f"deltaf_lineouts/{species}/{ps_name}"] = _write( + fig, f"deltaf_lineouts/{species}/{ps_name}.png" + ) + except Exception as e: + print(f"[plots] skipping delta-f lineout {diag_rel}: {e}") # --- Left/right-going transverse E-field decomposition --- try: - for comp, fig in plot_field_lr_decomposition( - run_dir, v_th=v_th, omega_k_zoom=omega_k_zoom - ).items(): + for comp, fig in plot_field_lr_decomposition(run_dir).items(): written[f"field_decomp/{comp}"] = _write(fig, f"field_decomp/{comp}.png") except Exception as e: print(f"[plots] skipping field decomposition: {e}") @@ -1098,6 +1442,26 @@ def _tex(s) -> str: return s +def _label_tex(s) -> str: + r"""Render an OSIRIS *label* that mixes prose words and TeX fragments. + + Diagnostic LABELs like ``species 1 p_1x_1`` interleave plain words with math + fragments. Wrapping the whole string in ``$...$`` (as :func:`_tex` does for + units, which are entirely mathematical) would typeset "species" in math + italics with no word spacing. Instead each whitespace-separated token is + wrapped on its own: tokens carrying TeX markup (``\``, ``_``, ``^``, ``{``, + ``}``) become math, while plain words and bare numbers stay as text. For a + single all-math token (``E_1``, ``\rho``) this matches :func:`_tex`. + """ + s = str(s) + if not s or (s.startswith("$") and s.endswith("$")): + return s + return " ".join( + f"${tok}$" if tok and any(c in tok for c in "\\_^{}") else tok + for tok in s.split(" ") + ) + + def _axis_label(da: xr.DataArray, dim: str) -> str: if dim == "t": u = da.attrs.get("time_units") or r"1/\omega_p" @@ -1115,13 +1479,13 @@ def _long_name(da: xr.DataArray) -> str: def _display_name(da: xr.DataArray) -> str: - """``_long_name`` wrapped for math-mode rendering (used in titles).""" - return _tex(_long_name(da)) + """``_long_name`` rendered for titles (prose stays text, TeX gets math).""" + return _label_tex(_long_name(da)) def _value_label(da: xr.DataArray) -> str: """Quantity label with units (for colorbars / y-axes).""" - name = _tex(_long_name(da)) + name = _label_tex(_long_name(da)) units = da.attrs.get("units", "") return rf"{name} [{_tex(units)}]" if units else name diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 09c7964e..235023eb 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -148,12 +148,7 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: # Render the standard plot set into td/plots so MLflow logs them as # artifacts. Never let a plotting failure abort metric/artifact logging. try: - output = cfg.get("output") or {} - kwargs: dict[str, Any] = {"v_th": output.get("v_th")} - if output.get("dist_cells") is not None: - kwargs["dist_cells"] = int(output["dist_cells"]) - if "omega_k_zoom" in output: # may be explicitly null to disable zoom - kwargs["omega_k_zoom"] = output["omega_k_zoom"] + kwargs = _plots.canned_plot_kwargs(cfg.get("output")) _plots.save_canned_plots(run_dir, td / "plots", **kwargs) except Exception as e: print(f"[post] plotting failed: {e}") diff --git a/adept/osiris/regen.py b/adept/osiris/regen.py new file mode 100644 index 00000000..88889773 --- /dev/null +++ b/adept/osiris/regen.py @@ -0,0 +1,161 @@ +"""Offline regeneration of the OSIRIS canned plot set from saved NetCDFs. + +When a run finishes, ``post.collect`` converts every diagnostic's time history +into a per-diagnostic NetCDF under ``binary/`` (mirroring the OSIRIS ``MS/`` +layout) and renders the canned plot set. This module regenerates that *same* +plot set from the ``binary/`` NetCDFs alone — no rerun and no raw ``MS/`` HDF5 +tree — so the plotting code in :mod:`adept.osiris.plots` can be iterated on +quickly against real data. + +The heavy lifting already lives in :func:`adept.osiris.io.list_diagnostics` / +:func:`adept.osiris.io.load_series`, which transparently dispatch between an +``MS/`` HDF5 tree and a ``binary/`` NetCDF directory. Regeneration is therefore +just :func:`adept.osiris.plots.save_canned_plots` pointed at the NetCDF +directory; this module supplies the ergonomics: locating the ``binary/`` dir, +reading plot knobs from the run's ``config.yaml``, and a CLI. + +Usage:: + + python -m adept.osiris.regen [--out DIR] [options] + +Examples:: + + # Regenerate into

/plots_regen, reading output.* from config.yaml if present + python -m adept.osiris.regen scratch/scan-a0-0.003 + + # Overlay the Langmuir branch (electron uth) and tighten the omega-k zoom + python -m adept.osiris.regen scratch/scan-a0-0.003 --v-th 0.0885 --omega-k-zoom 3 +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from adept.osiris import plots as _plots + + +def find_binary_dir(path: str | Path) -> Path: + """Resolve the NetCDF diagnostics directory from a run dir or the dir itself. + + Accepts a run directory that contains a ``binary/`` subdir (the common + case — the artifact layout), or a directory that *is* the NetCDF tree + already. Returns the directory to hand to ``save_canned_plots``. + """ + path = Path(path) + if (path / "binary").is_dir(): + return path / "binary" + return path + + +def default_out_dir(src: str | Path) -> Path: + """Where regenerated plots land by default: ``/plots_regen``.""" + src = Path(src) + run_dir = src if (src / "binary").is_dir() else src.parent + return run_dir / "plots_regen" + + +def load_output_cfg(src: str | Path) -> dict: + """Best-effort read of the manifest ``output:`` block for a run. + + Looks for ``config.yaml`` at ``src`` or its parent (so it is found whether + ``src`` is the run dir or the ``binary/`` dir). Returns the ``output`` dict, + or ``{}`` when no config is present / readable / has no ``output`` block. + """ + src = Path(src) + for candidate in (src / "config.yaml", src.parent / "config.yaml"): + if candidate.is_file(): + try: + import yaml + + cfg = yaml.safe_load(candidate.read_text()) or {} + except Exception: # a malformed config must not block regeneration + return {} + return cfg.get("output") or {} + return {} + + +def regenerate( + src: str | Path, + out_dir: str | Path | None = None, + *, + use_config: bool = True, + **overrides, +) -> dict[str, Path]: + """Regenerate the canned plot set from a run's saved NetCDFs. + + ``src`` is a run directory (containing ``binary/`` and optionally + ``config.yaml``) or a ``binary/`` NetCDF directory. Plot knobs default to + the run's ``output:`` config (unless ``use_config=False``); any keyword in + ``overrides`` (``v_th``, ``dist_cells``, ``omega_k_zoom``, ``dpi``, + ``n_panels``) is applied on top verbatim — pass only the knobs you mean to + set. Writes PNGs under ``out_dir`` (default ``/plots_regen``) and + returns the ``{plot-name: path}`` map from ``save_canned_plots``. + """ + binary = find_binary_dir(src) + out_dir = Path(out_dir) if out_dir is not None else default_out_dir(src) + kwargs = _plots.canned_plot_kwargs(load_output_cfg(src) if use_config else None) + kwargs.update(overrides) + return _plots.save_canned_plots(binary, out_dir, **kwargs) + + +def _summarize(written: dict[str, Path], out_dir: Path) -> None: + """Print a compact per-family summary of what was regenerated.""" + families: dict[str, int] = {} + for name in written: + fam = name.split("/", 1)[0] + families[fam] = families.get(fam, 0) + 1 + print(f"\nRegenerated {len(written)} plots -> {out_dir}") + for fam, n in sorted(families.items()): + print(f" {fam:24s} {n}") + + +def _cli_overrides(args: argparse.Namespace) -> dict: + """Collect only the plot knobs the user explicitly set on the command line.""" + overrides: dict = {} + if args.v_th is not None: + overrides["v_th"] = args.v_th + if args.dist_cells is not None: + overrides["dist_cells"] = args.dist_cells + if args.dpi is not None: + overrides["dpi"] = args.dpi + if args.n_panels is not None: + overrides["n_panels"] = args.n_panels + if args.no_zoom: + overrides["omega_k_zoom"] = None # explicit disable + elif args.omega_k_zoom is not None: + overrides["omega_k_zoom"] = args.omega_k_zoom + return overrides + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser( + prog="python -m adept.osiris.regen", + description="Regenerate the OSIRIS canned plot set from saved NetCDF artifacts.", + ) + ap.add_argument("src", help="run dir (containing binary/) or a binary/ NetCDF dir") + ap.add_argument("-o", "--out", default=None, help="output dir (default /plots_regen)") + ap.add_argument("--no-config", action="store_true", + help="ignore the run's config.yaml output block") + ap.add_argument("--v-th", type=float, default=None, + help="electron thermal velocity for the Langmuir overlay on omega-k plots") + ap.add_argument("--dist-cells", type=int, default=None, + help="right-boundary cells averaged for the f(p) distribution lineouts") + ap.add_argument("--omega-k-zoom", type=float, default=None, + help="(k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel") + ap.add_argument("--no-zoom", action="store_true", + help="use the full Nyquist window for the lower omega-k panel (omega_k_zoom=None)") + ap.add_argument("--dpi", type=int, default=None, help="figure DPI") + ap.add_argument("--n-panels", type=int, default=None, help="panels for faceted plots") + args = ap.parse_args(argv) + + out_dir = Path(args.out) if args.out else default_out_dir(args.src) + written = regenerate( + args.src, out_dir=out_dir, use_config=not args.no_config, **_cli_overrides(args) + ) + _summarize(written, out_dir) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index 53380380..5bec1328 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -20,4 +20,4 @@ output: # diagnostics_to_log: [e1] # null/omitted = convert every diagnostic's full time history to netCDF # v_th: 0.1 # overlay the Langmuir (Bohm-Gross) branch on omega-k plots # dist_cells: 10 # right-boundary cells averaged for the phase-space f(p) lineouts - # omega_k_zoom: 4.0 # (k, omega) half-width [omega_p] for the zoomed dispersion view; null = full + # omega_k_zoom: 4.0 # (k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel; null = full diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index 176941dd..b512abfa 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -97,7 +97,8 @@ output: dist_cells: 10 # right-boundary cells averaged for # the phase-space f(p) lineouts omega_k_zoom: 4.0 # (k, ω) half-width [ω_p] for the - # zoomed dispersion view; null = full + # equal-aspect lower ω–k panel + # (clamped to Nyquist); null = full ``` Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1)`). Indexed `{0: …, 1: …}` form addresses occurrences of repeated sections (`species`, `udist`, `profile`, `spe_bound`, `diag_species`, `zpulse`, …) in source order. @@ -122,8 +123,7 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | --------------------------------------------- | ------------- | | `spacetime/.png`, `spacetime_log/.png` | `(t, x)` heatmap of each FLD diagnostic (lin + log) | | `lineouts/.png` | value-vs-`x` snapshots at sampled times | -| `omega_k/.png` | full 2-D FFT `(k, ω)` dispersion | -| `omega_k_zoom/.png` | zoomed `(k, ω)` showing the whole `ω = k` light line — where the plasma (Langmuir) waves live | +| `omega_k/.png` | 2-D FFT `(k, ω)` dispersion — full Nyquist range on top, plus an equal-aspect square window below where `ω = k` is drawn at a true 45° | | `currents/spacetime.png`, `currents/lineouts.png` | combined `J_x/J_y/J_z` (`j1/j2/j3`) views | | `moments//…` | per-species density-moment spacetime + lineouts | | `profiles//density.png` | density profile vs `x` (final snapshot + late-time mean) | diff --git a/tests/test_osiris/test_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py index f28e3653..efca9a32 100644 --- a/tests/test_osiris/test_plots_new_views.py +++ b/tests/test_osiris/test_plots_new_views.py @@ -218,8 +218,11 @@ def test_efield_decomposition_isolates_right_going(tmp_path: Path) -> None: def test_field_decomposition_figures(tmp_path: Path) -> None: run_dir = _make_rich_run(tmp_path) - figs = oplt.plot_field_lr_decomposition(run_dir, omega_k_zoom=4.0) + figs = oplt.plot_field_lr_decomposition(run_dir) assert set(figs) == {"e2", "e3"} + # Each figure is now a single row of two spacetime panels (no omega-k row). + for fig in figs.values(): + assert len(fig.axes) == 4 # 2 heatmaps + their 2 colorbars # --- end-to-end driver ---------------------------------------------------- @@ -229,7 +232,7 @@ def test_save_canned_plots_emits_new_views(tmp_path: Path) -> None: run_dir = _make_rich_run(tmp_path) written = oplt.save_canned_plots(run_dir, tmp_path / "plots", v_th=0.1) expected = { - "omega_k_zoom/e2", + "omega_k/e2", "currents/spacetime", "currents/lineouts", "profiles/electron/density", From c6bb017efe3d3000c139500dd9514f6d3bd035ed Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 3 Jun 2026 08:12:52 -0700 Subject: [PATCH 11/49] Make srun default instead of mpirun for perlmutter --- adept/osiris/base.py | 1 + adept/osiris/runner.py | 12 +++++++----- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 95302c7e..357c8e52 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -97,6 +97,7 @@ def __call__(self, trainable_modules: dict, args: dict) -> dict: binary=binary, mpi_ranks=mpi_ranks, run_root=run_root, + launcher=osiris_cfg.get("mpi_launcher", "srun"), extra_mpi_args=osiris_cfg.get("extra_mpi_args"), ) return {"solver result": result} diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index e1c581e0..af75da48 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -2,9 +2,11 @@ OSIRIS reads its deck from a file named ``os-stdin`` in the current working directory by default. This module sets up a per-run work -directory, writes the rendered deck there, invokes ``mpirun`` (or runs -the binary directly when ``mpi_ranks == 1``), and captures stdout/stderr -to files for later artifact upload. +directory, writes the rendered deck there, invokes the configured +launcher (``srun`` by default — the team runs on Perlmutter/Slurm; override +with ``mpi_launcher: mpirun`` for a local MPI) when ``mpi_ranks > 1`` — or +runs the binary directly when ``mpi_ranks == 1`` — and captures +stdout/stderr to files for later artifact upload. """ from __future__ import annotations @@ -68,7 +70,7 @@ def run_osiris( mpi_ranks: int = 1, run_root: str | Path = "./osiris_runs", env: dict[str, str] | None = None, - mpirun: str = "mpirun", + launcher: str = "srun", extra_mpi_args: list[str] | None = None, ) -> dict[str, Any]: """Run OSIRIS and return run metadata. @@ -87,7 +89,7 @@ def run_osiris( (run_dir / "os-stdin").write_text(deck_text) if mpi_ranks > 1: - cmd = [mpirun, "-n", str(mpi_ranks)] + cmd = [launcher, "-n", str(mpi_ranks)] if extra_mpi_args: cmd.extend(extra_mpi_args) cmd.append(str(binary)) From 7528dde3b879d10eba223c8de921337ad5e118b7 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Thu, 4 Jun 2026 16:30:43 -0700 Subject: [PATCH 12/49] Modified distribution function plots; added wave transmission+absorption+reflection plot (should probably be moved to osiris-lpi repo) --- adept/osiris/plots.py | 225 +++++++++++++++++++++- adept/osiris/post.py | 16 ++ adept/osiris/regen.py | 49 ++++- docs/osiris-adept-usage.md | 7 +- tests/test_osiris/test_plots_new_views.py | 43 +++++ 5 files changed, 325 insertions(+), 15 deletions(-) diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index b7b6bf8e..30a2431d 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -669,11 +669,12 @@ def plot_distribution_lineout( series: xr.DataArray | str | Path, ax: plt.Axes | None = None, *, - n_cells: int = 10, + n_cells: int = 150, n_times: int = 6, log: bool = False, cmap: str = "viridis", title: str | None = None, + legend: bool = True, ) -> plt.Axes: """Distribution function ``f(p)`` near the right boundary. @@ -706,7 +707,8 @@ def plot_distribution_lineout( y = f_xp.isel(t=i).values y = np.abs(y) if log else y ax.plot(p, y, color=colour, lw=1.3, label=f"{float(tvals[i]):.2g}") - ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) + if legend: + ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) else: y = np.abs(f_xp.values) if log else f_xp.values ax.plot(p, y, lw=1.4) @@ -730,10 +732,11 @@ def plot_deltaf_lineout( series: xr.DataArray | str | Path, ax: plt.Axes | None = None, *, - n_cells: int = 10, + n_cells: int = 150, n_times: int = 6, cmap: str = "viridis", title: str | None = None, + legend: bool = True, ) -> plt.Axes: r"""``\delta f = f - f_M`` near the right boundary, overlaid at sampled times. @@ -763,7 +766,8 @@ def _residual(fvals): for colour, i in zip(colours, idx, strict=False): ax.plot(p, _residual(f_xp.isel(t=i).values), color=colour, lw=1.3, label=f"{float(tvals[i]):.2g}") - ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) + if legend: + ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) else: ax.plot(p, _residual(f_xp.values), lw=1.4) @@ -781,6 +785,41 @@ def _residual(fvals): return ax +def plot_distribution_lineout_figure( + series: xr.DataArray | str | Path, + *, + n_cells: int = 150, + n_times: int = 6, + cmap: str = "viridis", +) -> plt.Figure: + r"""Right-boundary distribution stacked three ways: ``f``, ``|f|`` log, ``\delta f``. + + Three views of the same boundary-averaged distribution (see + :func:`plot_distribution_lineout` / :func:`plot_deltaf_lineout`): the linear + ``f(p)`` panel shows the bulk, the ``|f|`` log panel the low-amplitude tail + / non-thermal structure, and the ``\delta f = f - f_M`` panel the residual + against the moment-matched Maxwellian (beams and hot tails). All three share + the momentum axis and the time sampling; the time legend is drawn once, on + the linear panel. + """ + da = _ensure_series(series) + fig, (ax_f, ax_logf, ax_df) = plt.subplots(3, 1, figsize=(6, 12), sharex=True) + plot_distribution_lineout(da, ax=ax_f, n_cells=n_cells, n_times=n_times, cmap=cmap) + plot_distribution_lineout( + da, ax=ax_logf, n_cells=n_cells, n_times=n_times, cmap=cmap, + log=True, legend=False, + title=rf"{_display_name(da)} — $|f(p)|$ (log scale)", + ) + plot_deltaf_lineout( + da, ax=ax_df, n_cells=n_cells, n_times=n_times, cmap=cmap, legend=False, + title=rf"{_display_name(da)} — $\delta f = f - f_M$", + ) + for ax in (ax_f, ax_logf): + ax.set_xlabel("") # shared axis: momentum label stays on the bottom panel only + fig.tight_layout() + return fig + + # --- currents (j1/j2/j3) -------------------------------------------------- @@ -1124,6 +1163,157 @@ def plot_field_lr_decomposition(run_dir: str | Path) -> dict[str, plt.Figure]: return figs +# --- laser energy budget (reflected / transmitted / absorbed) ------------- + + +def laser_energy_budget( + run_dir: str | Path, + *, + a0: float, + omega0: float, + last_frac: float = 0.25, + guard_cells: int = 1, + window_cells: int = 3, +) -> dict: + r"""Reflected / transmitted / absorbed laser-energy budget at the box edges. + + Uses the vacuum Riemann split (:func:`efield_lr_components`) to form the + one-way Poynting flux of the transverse light (code units, ``c = 1``, so the + flux of each travelling component is just its ``E^2`` summed over the + available polarization pairs): + + - **reflected** = left-going flux ``I_L`` at the **left** boundary, + - **transmitted** = right-going flux ``I_R`` at the **right** boundary. + + The incident intensity is taken from the drive amplitude, + ``I0 = (a0 * omega0)**2 / 2`` (the cycle-averaged flux of the launched wave, + in the same field units), so ``R`` and ``T`` are normalization-independent + ratios. The OSIRIS antenna sits at the lower-``x`` boundary, so its current + source contaminates the edge cell; ``guard_cells`` cells are skipped at each + boundary and the flux is averaged over a thin ``window_cells`` slab just + inside. Scalars are the flux averaged over the **last ``last_frac``** of the + run (the saturated phase) divided by ``I0``: ``R``, ``T`` and + ``absorbed = 1 - R - T`` (a steady-state energy balance; valid once the box + is filled and the in-flight transient has passed). + + Returns a dict with the per-dump time series (``t`` and the ``incident`` / + ``reflected`` / ``transmitted`` flux), the scalars (``I0``, ``R``, ``T``, + ``absorbed``, the ``incident_measured`` cross-check) and the settings used. + """ + comps = efield_lr_components(run_dir) + if not comps: + raise RuntimeError("no transverse field pairs (need e2/b3 or e3/b2) for the budget") + sample = next(iter(comps.values()))["right"] + xdim = next(d for d in sample.dims if d != "t") + t = np.asarray(sample.coords["t"].values, dtype=float) + # One-way Poynting flux summed over whichever polarization pairs were dumped. + I_R = sum(d["right"] ** 2 for d in comps.values()) + I_L = sum(d["left"] ** 2 for d in comps.values()) + n = I_R.sizes[xdim] + g, w = guard_cells, window_cells + if g + w > n: + g, w = 0, min(w, n) + left, right = slice(g, g + w), slice(n - g - w, n - g) + inc_t = I_R.isel({xdim: left}).mean(dim=xdim).values # right-going @ left (incident) + refl_t = I_L.isel({xdim: left}).mean(dim=xdim).values # left-going @ left (reflected) + trans_t = I_R.isel({xdim: right}).mean(dim=xdim).values # right-going @ right (transmitted) + + i0 = (a0 * omega0) ** 2 / 2.0 + t_lo = t[0] + (1.0 - last_frac) * (t[-1] - t[0]) + win = t >= t_lo + refl = float(np.mean(refl_t[win])) + trans = float(np.mean(trans_t[win])) + inc_m = float(np.mean(inc_t[win])) + r, tr = refl / i0, trans / i0 + return { + "a0": float(a0), "omega0": float(omega0), "I0": i0, + "t": t, "incident_t": inc_t, "reflected_t": refl_t, "transmitted_t": trans_t, + "t_window": (float(t_lo), float(t[-1])), "last_frac": float(last_frac), + "reflected": refl, "transmitted": trans, "incident_measured": inc_m, + "R": r, "T": tr, "absorbed": 1.0 - r - tr, + "guard_cells": g, "window_cells": w, "pairs": list(comps), + } + + +def _moving_average(y: np.ndarray, k: int) -> np.ndarray: + """Centered moving average over ``k`` samples, edge-normalized (no droop).""" + if k <= 1 or y.size < 3: + return y + ker = np.ones(min(k, y.size)) + num = np.convolve(y, ker, mode="same") + den = np.convolve(np.ones_like(y), ker, mode="same") + return num / den + + +def plot_energy_budget_figure(budget: dict, *, smooth_dumps: int = 21) -> plt.Figure: + r"""Transmitted / reflected / absorbed laser power vs time (fractions of ``I0``). + + Stacked-area energy partition over the run, each curve smoothed over + ``smooth_dumps`` dumps to suppress the ``2 omega`` intensity ripple (the + field dump cadence undersamples the laser period). The shaded band marks the + averaging window used for the scalar ``R`` / ``T``. Early on the partition is + dominated by "absorbed", which there is really energy still **filling the + box** (a transit-time transient), not true absorption — read the absorbed + curve only once it has settled. + """ + t, i0 = budget["t"], budget["I0"] + R = _moving_average(budget["reflected_t"], smooth_dumps) / i0 + T = _moving_average(budget["transmitted_t"], smooth_dumps) / i0 + A = np.clip(1.0 - R - T, 0.0, None) + + fig, ax = plt.subplots(figsize=(7, 4)) + ax.stackplot( + t, T, R, A, + labels=["transmitted", "reflected", "absorbed / in-flight"], + colors=["#4c9be8", "#e8744c", "#c2c2c2"], alpha=0.9, + ) + ax.axvspan(*budget["t_window"], color="k", alpha=0.08, lw=0, + label=f"scalar window (last {100*budget['last_frac']:g}%)") + ax.axhline(1.0, color="0.4", lw=0.6, ls="--") + ax.set_xlim(float(t[0]), float(t[-1])) + ax.set_ylim(0.0, 1.15) + ax.set_xlabel(r"$t$ [$1/\omega_p$]") + ax.set_ylabel(r"fraction of incident $I_0$") + ax.set_title( + rf"laser energy budget (incident $I_0=(a_0\omega_0)^2/2$, " + rf"$a_0={budget['a0']:g}$, $\omega_0={budget['omega0']:g}$)" + ) + ax.legend(loc="upper left", fontsize=8, ncol=2, framealpha=0.75) + fig.tight_layout() + return fig + + +def format_laser_energy_budget(budget: dict) -> str: + """Render :func:`laser_energy_budget` output as a human-readable ``.txt``.""" + b = budget + lo, hi = b["t_window"] + def pct(x: float) -> str: + return f"{100.0 * x:7.2f}%" + return "\n".join([ + "OSIRIS laser energy budget", + "==========================", + "incident intensity I0 = (a0 * omega0)^2 / 2 [code units, cycle-averaged]", + f" a0 = {b['a0']:g}", + f" omega0 = {b['omega0']:g}", + f" I0 = {b['I0']:.6e}", + "", + f"scalars: boundary flux averaged over the last {100*b['last_frac']:g}% of the run " + f"(t in [{lo:.1f}, {hi:.1f}]), divided by I0", + f" reflected (left-going @ left boundary) R = {pct(b['R'])} " + f"(flux {b['reflected']:.4e})", + f" transmitted (right-going @ right boundary) T = {pct(b['T'])} " + f"(flux {b['transmitted']:.4e})", + f" absorbed = 1 - R - T = {pct(b['absorbed'])}", + "", + f"cross-check: measured incident (right-going @ left boundary) = " + f"{b['incident_measured']:.4e} ({b['incident_measured']/b['I0']:.3f} x I0)", + f"polarization pairs used: {', '.join(b['pairs'])}", + f"boundary sampling: {b['guard_cells']} guard cell(s) skipped (antenna at lower-x), " + f"flux averaged over a {b['window_cells']}-cell slab", + "", + ]) + + # --- driver --------------------------------------------------------------- @@ -1172,8 +1362,10 @@ def save_canned_plots( v_th: float | None = None, dpi: int = 120, n_panels: int = 8, - dist_cells: int = 10, + dist_cells: int = 150, omega_k_zoom: float | None = 4.0, + a0: float | None = None, + omega0: float | None = None, ) -> dict[str, Path]: """Generate the standard set of PNGs for a finished OSIRIS run. @@ -1181,6 +1373,11 @@ def save_canned_plots( best-effort: a failure on one diagnostic logs and is skipped rather than aborting the rest. + When ``a0`` and ``omega0`` (the drive amplitude / frequency) are given, the + laser energy budget is also produced: the ``energy_budget`` time-series plot + and a ``laser_energy_budget.txt`` summary alongside it (see + :func:`laser_energy_budget`). + ``run_dir`` may be a raw OSIRIS run directory (with an ``MS/`` HDF5 tree and optional ``HIST/``) **or** the ``binary/`` directory of saved NetCDFs written by :func:`io.save_run_datasets`. The data source is detected @@ -1357,10 +1554,9 @@ def _write(fig: plt.Figure, rel: str) -> Path: # only for x-p phase spaces (a p-p space like p1p2 has no spatial axis). if ser.ndim == 3 and any(str(d).startswith("x") for d in ser.dims if d != "t"): try: - fig, ax = plt.subplots(figsize=(6, 4)) - plot_distribution_lineout(ser, ax=ax, n_cells=dist_cells) written[f"distribution_lineouts/{species}/{ps_name}"] = _write( - fig, f"distribution_lineouts/{species}/{ps_name}.png" + plot_distribution_lineout_figure(ser, n_cells=dist_cells), + f"distribution_lineouts/{species}/{ps_name}.png", ) except Exception as e: print(f"[plots] skipping distribution lineout {diag_rel}: {e}") @@ -1406,6 +1602,19 @@ def _write(fig: plt.Figure, rel: str) -> Path: except Exception as e: print(f"[plots] skipping total_energy_vs_time: {e}") + # --- Laser energy budget (reflected/transmitted/absorbed); needs a0, ω0 --- + if a0 is not None and omega0 is not None: + try: + budget = laser_energy_budget(run_dir, a0=a0, omega0=omega0) + written["energy_budget"] = _write( + plot_energy_budget_figure(budget), "energy_budget.png" + ) + (out_dir / "laser_energy_budget.txt").write_text( + format_laser_energy_budget(budget) + ) + except Exception as e: + print(f"[plots] skipping laser energy budget: {e}") + return written diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 235023eb..e0f4260c 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -145,14 +145,30 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: except Exception as e: print(f"[post] HIST energy unavailable: {e}") + # Drive amplitude/frequency from the deck → laser energy budget. + deck = cfg.get("deck") or {} + a0, omega0 = deck.get("antenna.a0"), deck.get("antenna.omega0") + # Render the standard plot set into td/plots so MLflow logs them as # artifacts. Never let a plotting failure abort metric/artifact logging. try: kwargs = _plots.canned_plot_kwargs(cfg.get("output")) + if a0 is not None and omega0 is not None: + kwargs.update(a0=float(a0), omega0=float(omega0)) _plots.save_canned_plots(run_dir, td / "plots", **kwargs) except Exception as e: print(f"[post] plotting failed: {e}") + # Reflected/transmitted/absorbed laser-energy fractions as scalar metrics. + if a0 is not None and omega0 is not None: + try: + budget = _plots.laser_energy_budget(run_dir, a0=float(a0), omega0=float(omega0)) + metrics["laser_reflectivity"] = budget["R"] + metrics["laser_transmissivity"] = budget["T"] + metrics["laser_absorbed_frac"] = budget["absorbed"] + except Exception as e: + print(f"[post] laser energy budget unavailable: {e}") + # Always include the rendered deck + stdout/stderr. for fname in ("os-stdin", "stdout.log", "stderr.log"): src = run_dir / fname diff --git a/adept/osiris/regen.py b/adept/osiris/regen.py index 88889773..bcf684bf 100644 --- a/adept/osiris/regen.py +++ b/adept/osiris/regen.py @@ -75,6 +75,31 @@ def load_output_cfg(src: str | Path) -> dict: return {} +def load_antenna_params(src: str | Path) -> dict: + """Best-effort read of the drive amplitude / frequency (``a0``, ``omega0``). + + Reads the flat ``deck`` block of ``derived_config.yaml`` (falling back to + ``config.yaml``) at ``src`` or its parent, returning ``{"a0":…, "omega0":…}`` + for the single antenna when both are present, else ``{}``. The laser energy + budget needs these; absent them (e.g. a non-antenna deck) it is skipped. + """ + src = Path(src) + for name in ("derived_config.yaml", "config.yaml"): + for candidate in (src / name, src.parent / name): + if candidate.is_file(): + try: + import yaml + + cfg = yaml.safe_load(candidate.read_text()) or {} + except Exception: # a malformed config must not block regeneration + continue + deck = cfg.get("deck") or {} + a0, omega0 = deck.get("antenna.a0"), deck.get("antenna.omega0") + if a0 is not None and omega0 is not None: + return {"a0": float(a0), "omega0": float(omega0)} + return {} + + def regenerate( src: str | Path, out_dir: str | Path | None = None, @@ -85,16 +110,20 @@ def regenerate( """Regenerate the canned plot set from a run's saved NetCDFs. ``src`` is a run directory (containing ``binary/`` and optionally - ``config.yaml``) or a ``binary/`` NetCDF directory. Plot knobs default to - the run's ``output:`` config (unless ``use_config=False``); any keyword in + ``config.yaml`` / ``derived_config.yaml``) or a ``binary/`` NetCDF + directory. Plot knobs default to the run's ``output:`` config (unless + ``use_config=False``) and the drive ``a0`` / ``omega0`` are read from + ``derived_config.yaml`` for the laser energy budget; any keyword in ``overrides`` (``v_th``, ``dist_cells``, ``omega_k_zoom``, ``dpi``, - ``n_panels``) is applied on top verbatim — pass only the knobs you mean to - set. Writes PNGs under ``out_dir`` (default ``/plots_regen``) and - returns the ``{plot-name: path}`` map from ``save_canned_plots``. + ``n_panels``, ``a0``, ``omega0``) is applied on top verbatim — pass only the + knobs you mean to set. Writes PNGs (and ``laser_energy_budget.txt`` when + ``a0``/``omega0`` are known) under ``out_dir`` (default + ``/plots_regen``) and returns the ``{plot-name: path}`` map. """ binary = find_binary_dir(src) out_dir = Path(out_dir) if out_dir is not None else default_out_dir(src) kwargs = _plots.canned_plot_kwargs(load_output_cfg(src) if use_config else None) + kwargs.update(load_antenna_params(src)) # a0/omega0 for the laser energy budget kwargs.update(overrides) return _plots.save_canned_plots(binary, out_dir, **kwargs) @@ -125,6 +154,10 @@ def _cli_overrides(args: argparse.Namespace) -> dict: overrides["omega_k_zoom"] = None # explicit disable elif args.omega_k_zoom is not None: overrides["omega_k_zoom"] = args.omega_k_zoom + if args.a0 is not None: + overrides["a0"] = args.a0 + if args.omega0 is not None: + overrides["omega0"] = args.omega0 return overrides @@ -147,6 +180,12 @@ def main(argv: list[str] | None = None) -> int: help="use the full Nyquist window for the lower omega-k panel (omega_k_zoom=None)") ap.add_argument("--dpi", type=int, default=None, help="figure DPI") ap.add_argument("--n-panels", type=int, default=None, help="panels for faceted plots") + ap.add_argument("--a0", type=float, default=None, + help="drive amplitude for the laser energy budget " + "(overrides derived_config.yaml)") + ap.add_argument("--omega0", type=float, default=None, + help="drive frequency for the laser energy budget " + "(overrides derived_config.yaml)") args = ap.parse_args(argv) out_dir = Path(args.out) if args.out else default_out_dir(args.src) diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index b512abfa..d7a1f577 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -94,7 +94,7 @@ output: diagnostics_to_log: null # null = all; or [e1, charge, …] v_th: 0.1 # optional: overlays the Bohm–Gross # Langmuir branch on ω–k plots - dist_cells: 10 # right-boundary cells averaged for + dist_cells: 150 # right-boundary cells averaged for # the phase-space f(p) lineouts omega_k_zoom: 4.0 # (k, ω) half-width [ω_p] for the # equal-aspect lower ω–k panel @@ -129,12 +129,15 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | `profiles//density.png` | density profile vs `x` (final snapshot + late-time mean) | | `profiles//temperature.png` | temperature profile vs `x`, from `uth1/2/3` or `T11/22/33` moments (omitted if neither is dumped) | | `phasespace//.png`, `phasespace_evolution/…` | `(x, p)` phase-space heatmaps | -| `distribution_lineouts//.png` | `f(p)` averaged over the rightmost `dist_cells` cells, overlaid at sampled times | +| `distribution_lineouts//.png` | `f(p)` averaged over the rightmost `dist_cells` cells, overlaid at sampled times — stacked linear `f`, `\|f\|` log, and `δf = f - f_M` panels | | `field_decomp/.png` | left/right-going transverse `E` (vacuum Riemann split `(e2±b3)/2`, `(e3∓b2)/2`), spacetime + `ω–k` | | `energy_vs_time.png`, `energy_components_vs_time.png`, `total_energy_vs_time.png` | field / kinetic energy traces | +| `energy_budget.png` + `laser_energy_budget.txt` | reflected / transmitted / absorbed laser power vs time (stacked-area), with scalar `R`/`T`/absorbed in the `.txt` — emitted only when the drive `a0`/`omega0` are known | > **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 longitudinal `e1` is electrostatic and is intentionally excluded. +> **Laser energy budget.** Reflected = left-going Poynting flux at the left boundary, transmitted = right-going flux at the right boundary (from the same Riemann split), compared to the incident intensity `I₀ = (a0·ω0)²/2` from the drive. Scalars are the boundary flux averaged over the last 25 % of the run (the saturated phase) ÷ `I₀`; `absorbed = 1 − R − T`. The antenna sits at the lower-`x` boundary, so its source cell is skipped. The same dispersive-medium caveat as `field_decomp/` applies — the split is approximate where the boundaries sit in plasma. `R`/`T`/absorbed are also logged to MLflow as `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac`. + ## Programmatic use ```python diff --git a/tests/test_osiris/test_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py index efca9a32..89517f72 100644 --- a/tests/test_osiris/test_plots_new_views.py +++ b/tests/test_osiris/test_plots_new_views.py @@ -244,3 +244,46 @@ def test_save_canned_plots_emits_new_views(tmp_path: Path) -> None: assert expected <= set(written) for path in written.values(): assert path.exists() and path.stat().st_size > 0 + + +# --- laser energy budget -------------------------------------------------- + + +def test_laser_energy_budget_pure_right_going_has_no_reflection(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + # Fixture fields are a pure right-going wave (e2=b3, e3=-b2), so the + # left-going Riemann parts vanish identically: nothing is "reflected". + # Averaging over the whole grid (one full wavelength) and all time gives + # = <2 sin^2> ~ 1, matching I0 = (a0*omega0)^2/2 = 1 -> T ~ 1. + b = oplt.laser_energy_budget( + run_dir, a0=float(np.sqrt(2.0)), omega0=1.0, + last_frac=1.0, guard_cells=0, window_cells=16, + ) + assert b["pairs"] == ["e2", "e3"] + assert b["R"] < 1e-9 # no left-going flux at all + assert b["transmitted"] > 0.0 + # window = whole grid -> left/right slabs coincide -> equal fluxes + assert abs(b["transmitted"] - b["incident_measured"]) < 1e-9 + assert 0.7 < b["T"] < 1.3 + assert abs(b["absorbed"] - (1.0 - b["R"] - b["T"])) < 1e-12 + + +def test_save_canned_plots_emits_energy_budget_with_drive(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + out = tmp_path / "plots" + written = oplt.save_canned_plots(run_dir, out, a0=0.004, omega0=1.0) + assert "energy_budget" in written + assert written["energy_budget"].exists() and written["energy_budget"].stat().st_size > 0 + txt = out / "laser_energy_budget.txt" + assert txt.exists() + body = txt.read_text() + for key in ("reflected", "transmitted", "absorbed", "I0"): + assert key in body + + +def test_save_canned_plots_skips_energy_budget_without_drive(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + out = tmp_path / "plots" + written = oplt.save_canned_plots(run_dir, out) # no a0/omega0 -> skipped + assert "energy_budget" not in written + assert not (out / "laser_energy_budget.txt").exists() From d0090ea617eb588b440b5d9aacf37e152595080c Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Thu, 4 Jun 2026 21:10:35 -0700 Subject: [PATCH 13/49] Store single precision floats, add compression (most useful for phasespace) --- adept/osiris/io.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 058276d2..8b2109e6 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -29,6 +29,13 @@ _ITER_RE = re.compile(r"-(\d+)\.h5$") +# Precision for diagnostic *data* (field/phase-space grids and RAW particle +# quantities) in the saved NetCDF artifacts. OSIRIS writes its dumps in single +# precision, so float32 matches the native precision and halves artifact size +# vs float64. Coordinates (time, spatial axes) stay float64 for axis precision +# (e.g. the omega-k FFT relies on the float64 time axis). +_DIAG_DTYPE = "float32" + def _decode(v) -> str: """OSIRIS stores attrs as ``S256`` bytes inside length-1 arrays.""" @@ -84,7 +91,7 @@ def load_grid_h5(path: str | Path) -> xr.DataArray: f"Expected exactly one data dataset in {path}; got {data_keys}" ) name = data_keys[0] - arr = f[name][...].astype("float64") + arr = f[name][...].astype(_DIAG_DTYPE) axes_osiris = _axis_metadata(f) # Reverse to numpy order. axes_numpy = list(reversed(axes_osiris)) @@ -161,7 +168,7 @@ def load_raw_h5(path: str | Path) -> xr.Dataset: data_vars: dict[str, tuple] = {} for name in data_keys: dset = f[name] - arr = dset[...].astype("float64").reshape(-1) + arr = dset[...].astype(_DIAG_DTYPE).reshape(-1) var_attrs = {} if "UNITS" in dset.attrs: var_attrs["units"] = _decode(dset.attrs["UNITS"]) @@ -433,6 +440,19 @@ def list_diagnostics_nc(binary_dir: str | Path) -> dict[str, Path]: return out +def _compression_encoding(ds: xr.Dataset) -> dict: + """Per-variable zlib settings for ``to_netcdf``. + + Field / phase-space grids are smooth and compress several-fold; ``shuffle`` + helps the float32 byte pattern. RAW (per-particle) data is noise-like and + barely compresses, but the setting does no harm. + """ + return { + name: {"zlib": True, "complevel": 4, "shuffle": True} + for name in ds.data_vars + } + + def save_run_datasets( run_dir: str | Path, out_dir: str | Path, @@ -465,7 +485,7 @@ def save_run_datasets( ds = series_to_dataset(load_series(diags[relpath])) dest = out_dir / f"{relpath}.nc" dest.parent.mkdir(parents=True, exist_ok=True) - ds.to_netcdf(dest, engine="h5netcdf") + ds.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(ds)) written.append(dest) except Exception as e: # one bad diagnostic must not abort the rest print(f"[post] skipping diagnostic {relpath}: {e}") @@ -479,7 +499,9 @@ def save_run_datasets( if energy is not None: dest = out_dir / "HIST" / "energy.nc" dest.parent.mkdir(parents=True, exist_ok=True) - energy.to_netcdf(dest, engine="h5netcdf") + energy.to_netcdf( + dest, engine="h5netcdf", encoding=_compression_encoding(energy) + ) written.append(dest) except Exception as e: print(f"[post] skipping HIST energy: {e}") From 55f662eb075012a9fed16f835daa8a74d4659174 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 5 Jun 2026 21:07:39 -0700 Subject: [PATCH 14/49] Optionally drop t=0 raw chunk (useful for massive raws) --- adept/osiris/io.py | 15 +++++++++++++-- adept/osiris/post.py | 6 +++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 8b2109e6..85a64a56 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -200,7 +200,7 @@ def load_raw_h5(path: str | Path) -> xr.Dataset: return xr.Dataset(data_vars, attrs=attrs) -def load_raw_series(directory: str | Path) -> xr.Dataset: +def load_raw_series(directory: str | Path, *, drop_initial: bool = False) -> xr.Dataset: """Concatenate every RAW (particle) dump in ``directory`` long-form. RAW dumps have a *variable* particle count per timestep (OSIRIS samples a @@ -216,6 +216,13 @@ def load_raw_series(directory: str | Path) -> xr.Dataset: if not dumps: raise FileNotFoundError(f"No .h5 dumps in {directory}") + # Optionally drop the t=0 (initial-condition) RAW dump: OSIRIS dumps RAW + # periodically from n=0, but at full raw_fraction that IC snapshot is the + # thermal start state and just bloats the artifact. Filter by filename so the + # (large) n=0 dump is never loaded. Kept if it is the sole dump. + if drop_initial and len(dumps) > 1: + dumps = [p for p in dumps if _iter_from_name(p) != 0] or dumps + per_dump: list[xr.Dataset] = [] times: list[np.ndarray] = [] iters: list[np.ndarray] = [] @@ -457,6 +464,8 @@ def save_run_datasets( run_dir: str | Path, out_dir: str | Path, diagnostics: list[str] | set[str] | None = None, + *, + raw_drop_initial: bool = False, ) -> list[Path]: """Convert each diagnostic's full time history to a netCDF file. @@ -480,7 +489,9 @@ def save_run_datasets( try: if _diag_is_raw(relpath, diags[relpath]): # RAW (particle) dumps: per-particle datasets, no grid/AXIS. - ds: xr.Dataset = load_raw_series(diags[relpath]) + ds: xr.Dataset = load_raw_series( + diags[relpath], drop_initial=raw_drop_initial + ) else: ds = series_to_dataset(load_series(diags[relpath])) dest = out_dir / f"{relpath}.nc" diff --git a/adept/osiris/post.py b/adept/osiris/post.py index e0f4260c..623fbeb8 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -116,6 +116,7 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: ms = run_dir / "MS" td = Path(td) whitelist = (cfg.get("output") or {}).get("diagnostics_to_log") or None + raw_drop_initial = bool((cfg.get("output") or {}).get("raw_drop_initial", False)) metrics: dict[str, float] = { "wall_time_s": float(solver["wall_time"]), @@ -126,7 +127,10 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: # Convert each diagnostic's time history to an xarray netCDF. if ms.is_dir(): - _io.save_run_datasets(run_dir, td / "binary", diagnostics=whitelist) + _io.save_run_datasets( + run_dir, td / "binary", diagnostics=whitelist, + raw_drop_initial=raw_drop_initial, + ) # plots imports matplotlib; do it lazily to keep `import adept.osiris` light. from adept.osiris import plots as _plots From cbaeabee3503e4582060485f5714508ae7928a4d Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 9 Jun 2026 17:33:59 -0700 Subject: [PATCH 15/49] Changed default osiris run directory to ./checkpoints so that the sync-up.sh script doesn't delete ongoing osiris runs --- adept/osiris/base.py | 2 +- adept/osiris/runner.py | 2 +- configs/osiris/twostream-1d-short.yaml | 2 +- configs/osiris/twostream-1d.yaml | 2 +- docs/osiris-adept-usage.md | 9 ++++++++- 5 files changed, 12 insertions(+), 5 deletions(-) diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 357c8e52..172299b1 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -90,7 +90,7 @@ def __call__(self, trainable_modules: dict, args: dict) -> dict: dim=self._infer_dim(), ) mpi_ranks = int(osiris_cfg.get("mpi_ranks", 1)) - run_root = osiris_cfg.get("run_root", "./osiris_runs") + run_root = osiris_cfg.get("run_root", "./checkpoints") deck_text = _deck.render_deck(self._sections) result = _runner.run_osiris( deck_text, diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index af75da48..410bfb9a 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -68,7 +68,7 @@ def run_osiris( *, binary: str | Path, mpi_ranks: int = 1, - run_root: str | Path = "./osiris_runs", + run_root: str | Path = "./checkpoints", env: dict[str, str] | None = None, launcher: str = "srun", extra_mpi_args: list[str] | None = None, diff --git a/configs/osiris/twostream-1d-short.yaml b/configs/osiris/twostream-1d-short.yaml index fcbad707..d6d154cc 100644 --- a/configs/osiris/twostream-1d-short.yaml +++ b/configs/osiris/twostream-1d-short.yaml @@ -8,6 +8,6 @@ osiris: deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e mpi_ranks: 1 - run_root: ./osiris_runs + run_root: ./checkpoints overrides: time: {tmax: 1.0} diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index 5bec1328..80f501a0 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -8,7 +8,7 @@ osiris: deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e mpi_ranks: 1 - run_root: ./osiris_runs + run_root: ./checkpoints # overrides: dict-of-dicts merged into the parsed deck before render. # Repeated sections (species, zpulse, ...) use integer indices. # overrides: diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index d7a1f577..9fc69c47 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -80,7 +80,14 @@ osiris: deck: /path/to/native/deck # required: source of truth binary: /path/to/osiris-1D.e # or OSIRIS_BIN / OSIRIS_BIN_D mpi_ranks: 1 # 1 → direct, >1 → mpirun -n N - run_root: ./osiris_runs # parent of per-run dirs + run_root: ./checkpoints # parent of per-run dirs (default) + # NOTE: the default sits inside checkpoints/ deliberately — sync-up.sh + # rsyncs with --delete but excludes checkpoints/, so in-flight and finished + # OSIRIS outputs survive a sync. If you point run_root anywhere outside an + # excluded directory, a sync-up invocation will delete those outputs. + # Nothing deletes run dirs automatically (post-processing only copies out + # of them), so clean them up manually on occasion — though $PSCRATCH's + # purge policy will take care of stale ones eventually. extra_mpi_args: ["--oversubscribe"] # optional, passed to mpirun overrides: # optional: applied before render From 78e6afc5dbca8c19b402b185712c6ed99011b692 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 15 Jun 2026 13:01:57 -0700 Subject: [PATCH 16/49] Moved LPI-specific code and plots from adept to the osiris-lpi repo and new osiris_lpi package --- adept/osiris/plots.py | 435 +--------------------- adept/osiris/post.py | 22 +- adept/osiris/regen.py | 52 +-- configs/osiris/twostream-1d.yaml | 2 +- docs/osiris-adept-usage.md | 6 +- tests/test_osiris/test_plots_new_views.py | 70 ---- 6 files changed, 30 insertions(+), 557 deletions(-) diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index 30a2431d..1a12166b 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -23,8 +23,6 @@ (initial + final + late mean); else uth / T_ii moments phasespace//.png final-step (x, p) heatmap per species phasespace_evolution//.png (x, p) heatmaps at sampled times - distribution_lineouts//.png f(p) averaged over the right-boundary cells - deltaf_lineouts//.png delta-f = f - (fitted Maxwellian), same averaging field_decomp/.png left/right-going transverse E (Riemann split) energy_vs_time.png total field energy time-trace energy_components_vs_time.png E-field / B-field / total energy @@ -606,220 +604,6 @@ def _crop_spatial_to_box(da: xr.DataArray) -> xr.DataArray: return da -def _maxwellian_moments(p, f) -> tuple[float, float, float] | None: - r"""Density, mean momentum, and variance of a 1-D distribution ``f(p)``. - - Returns ``(W, p0, var)`` with ``W = sum(f)`` (negative weights clipped to - zero), ``p0 =

`` and ``var = <(p - p0)^2>`` — the parameters of the - moment-matched Maxwellian, ``var`` being the kinetic temperature ``T`` for - ``m = 1`` (electrons, ``p`` in ``m_e c``). ``None`` if ``f`` has no positive - weight. - """ - f = np.clip(np.asarray(f, dtype=float), 0.0, None) - W = float(f.sum()) - if W <= 0: - return None - p = np.asarray(p, dtype=float) - p0 = float((f * p).sum() / W) - var = float((f * (p - p0) ** 2).sum() / W) - return W, p0, var - - -def _maxwellian_from_moments(p, W: float, p0: float, var: float): - """Gaussian on grid ``p`` with the same sum ``W``, mean ``p0`` and ``var``.""" - p = np.asarray(p, dtype=float) - if var <= 0: - return np.zeros_like(p) - g = np.exp(-((p - p0) ** 2) / (2.0 * var)) - s = float(g.sum()) - return g * (W / s) if s > 0 else g - - -def _boundary_distribution(da: xr.DataArray, n_cells: int): - r"""``f(p)`` averaged over the ``n_cells`` cells inside the right box edge. - - Shared by the ``f(p)`` and ``\delta f`` lineouts. Returns - ``(f_xp, da, pdim, x_lo, x_hi, k)``: the decorated input ``da``, the - sign-corrected (non-negative) distribution ``f_xp`` with dims ``(t?, p)``, - the momentum axis name ``pdim``, the averaged spatial window ``[x_lo, - x_hi]`` and the cell count ``k``. The window is anchored at the physical box - edge (``sim.XMAX``) so the empty ``ps_xmax`` padding beyond the box is - excluded. - """ - da = _decorate(da) - spatial = [d for d in da.dims if d != "t" and str(d).startswith("x")] - moment = [d for d in da.dims if d != "t" and str(d).startswith("p")] - if not spatial or not moment: - raise ValueError(f"need an x* and a p* dim; got dims {da.dims}") - xdim, pdim = spatial[0], moment[0] - nx = da.sizes[xdim] - xv = da.coords[xdim].values - right = int(np.searchsorted(xv, _sim_box_xmax(da, xdim), side="right")) - 1 - right = min(max(right, 0), nx - 1) - k = max(1, min(n_cells, right + 1)) - lo = right - k + 1 - f_xp = da.isel({xdim: slice(lo, right + 1)}).mean(dim=xdim) - # Charge density q*f is single-signed per species (negative for electrons); - # divide out the sign so f(p) is non-negative. - q_sign = np.sign(float(np.nansum(da.values))) or 1.0 - return f_xp / q_sign, da, pdim, float(xv[lo]), float(xv[right]), k - - -def plot_distribution_lineout( - series: xr.DataArray | str | Path, - ax: plt.Axes | None = None, - *, - n_cells: int = 150, - n_times: int = 6, - log: bool = False, - cmap: str = "viridis", - title: str | None = None, - legend: bool = True, -) -> plt.Axes: - """Distribution function ``f(p)`` near the right boundary. - - Takes a phase-space series with a spatial (``x*``) and a momentum - (``p*``) axis — either a stacked ``(t, p, x)`` series from - :func:`io.load_series` or a single ``(p, x)`` snapshot — averages it over - the ``n_cells`` cells just inside the right edge of the *physical box* - (``sim.XMAX``; the phase-space axis may extend past the box when the deck's - ``ps_xmax`` exceeds ``xmax``), and overlays that lineout at ~``n_times`` - sampled times (colour = time). This is the standard "what does the - distribution look like as it leaves the box" diagnostic for driven / - open-boundary runs. - - OSIRIS species phase spaces store the charge density ``q*f``; the charge - sign is divided out so the curve reads as a non-negative ``f(p)``. - """ - da = series if isinstance(series, xr.DataArray) else _io.load_series(series) - f_xp, da, pdim, x_lo, x_hi, k = _boundary_distribution(da, n_cells) - p = f_xp.coords[pdim].values - - if ax is None: - _, ax = plt.subplots(figsize=(6, 4)) - - if "t" in f_xp.dims: - nt = f_xp.sizes["t"] - idx = np.unique(np.linspace(0, nt - 1, min(n_times, nt)).astype(int)) - colours = plt.get_cmap(cmap)(np.linspace(0.15, 0.95, len(idx))) - tvals = f_xp.coords["t"].values - for colour, i in zip(colours, idx, strict=False): - y = f_xp.isel(t=i).values - y = np.abs(y) if log else y - ax.plot(p, y, color=colour, lw=1.3, label=f"{float(tvals[i]):.2g}") - if legend: - ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) - else: - y = np.abs(f_xp.values) if log else f_xp.values - ax.plot(p, y, lw=1.4) - - if log: - ax.set_yscale("log") - units = da.attrs.get("units", "") - ylabel = rf"$\langle f \rangle_x$ [{_tex(units)}]" if units else r"$\langle f \rangle_x$" - ax.set_xlabel(_axis_label(da, pdim)) - ax.set_ylabel(ylabel) - ax.set_title( - title - or rf"{_display_name(da)} — $f(p)$ over {k} cells inside the right boundary " - rf"($x \in [{x_lo:.4g}, {x_hi:.4g}]$)" - ) - ax.grid(True, alpha=0.3) - return ax - - -def plot_deltaf_lineout( - series: xr.DataArray | str | Path, - ax: plt.Axes | None = None, - *, - n_cells: int = 150, - n_times: int = 6, - cmap: str = "viridis", - title: str | None = None, - legend: bool = True, -) -> plt.Axes: - r"""``\delta f = f - f_M`` near the right boundary, overlaid at sampled times. - - Averages the phase space over the ``n_cells`` cells inside the right box - edge to get ``f(p)`` (as :func:`plot_distribution_lineout`), then at each - sampled time subtracts the moment-matched Maxwellian ``f_M`` — the Gaussian - with the same density, mean momentum and temperature — and overlays the - residual. Since the species starts Maxwellian, the residual isolates the - non-thermal structure that grows away from it: beams and hot tails. - """ - da = series if isinstance(series, xr.DataArray) else _io.load_series(series) - f_xp, da, pdim, x_lo, x_hi, k = _boundary_distribution(da, n_cells) - p = f_xp.coords[pdim].values - - def _residual(fvals): - m = _maxwellian_moments(p, fvals) - return fvals if m is None else fvals - _maxwellian_from_moments(p, *m) - - if ax is None: - _, ax = plt.subplots(figsize=(6, 4)) - - if "t" in f_xp.dims: - nt = f_xp.sizes["t"] - idx = np.unique(np.linspace(0, nt - 1, min(n_times, nt)).astype(int)) - colours = plt.get_cmap(cmap)(np.linspace(0.15, 0.95, len(idx))) - tvals = f_xp.coords["t"].values - for colour, i in zip(colours, idx, strict=False): - ax.plot(p, _residual(f_xp.isel(t=i).values), color=colour, lw=1.3, - label=f"{float(tvals[i]):.2g}") - if legend: - ax.legend(fontsize=7, title=_axis_label(da, "t"), ncol=2, framealpha=0.6) - else: - ax.plot(p, _residual(f_xp.values), lw=1.4) - - ax.axhline(0.0, color="0.6", lw=0.6) - units = da.attrs.get("units", "") - ylabel = r"$\langle \delta f \rangle_x$" + (f" [{_tex(units)}]" if units else "") - ax.set_xlabel(_axis_label(da, pdim)) - ax.set_ylabel(ylabel) - ax.set_title( - title - or rf"{_display_name(da)} — $\delta f = f - f_M$ over {k} cells inside the " - rf"right boundary ($x \in [{x_lo:.4g}, {x_hi:.4g}]$)" - ) - ax.grid(True, alpha=0.3) - return ax - - -def plot_distribution_lineout_figure( - series: xr.DataArray | str | Path, - *, - n_cells: int = 150, - n_times: int = 6, - cmap: str = "viridis", -) -> plt.Figure: - r"""Right-boundary distribution stacked three ways: ``f``, ``|f|`` log, ``\delta f``. - - Three views of the same boundary-averaged distribution (see - :func:`plot_distribution_lineout` / :func:`plot_deltaf_lineout`): the linear - ``f(p)`` panel shows the bulk, the ``|f|`` log panel the low-amplitude tail - / non-thermal structure, and the ``\delta f = f - f_M`` panel the residual - against the moment-matched Maxwellian (beams and hot tails). All three share - the momentum axis and the time sampling; the time legend is drawn once, on - the linear panel. - """ - da = _ensure_series(series) - fig, (ax_f, ax_logf, ax_df) = plt.subplots(3, 1, figsize=(6, 12), sharex=True) - plot_distribution_lineout(da, ax=ax_f, n_cells=n_cells, n_times=n_times, cmap=cmap) - plot_distribution_lineout( - da, ax=ax_logf, n_cells=n_cells, n_times=n_times, cmap=cmap, - log=True, legend=False, - title=rf"{_display_name(da)} — $|f(p)|$ (log scale)", - ) - plot_deltaf_lineout( - da, ax=ax_df, n_cells=n_cells, n_times=n_times, cmap=cmap, legend=False, - title=rf"{_display_name(da)} — $\delta f = f - f_M$", - ) - for ax in (ax_f, ax_logf): - ax.set_xlabel("") # shared axis: momentum label stays on the bottom panel only - fig.tight_layout() - return fig - - # --- currents (j1/j2/j3) -------------------------------------------------- @@ -1163,157 +947,6 @@ def plot_field_lr_decomposition(run_dir: str | Path) -> dict[str, plt.Figure]: return figs -# --- laser energy budget (reflected / transmitted / absorbed) ------------- - - -def laser_energy_budget( - run_dir: str | Path, - *, - a0: float, - omega0: float, - last_frac: float = 0.25, - guard_cells: int = 1, - window_cells: int = 3, -) -> dict: - r"""Reflected / transmitted / absorbed laser-energy budget at the box edges. - - Uses the vacuum Riemann split (:func:`efield_lr_components`) to form the - one-way Poynting flux of the transverse light (code units, ``c = 1``, so the - flux of each travelling component is just its ``E^2`` summed over the - available polarization pairs): - - - **reflected** = left-going flux ``I_L`` at the **left** boundary, - - **transmitted** = right-going flux ``I_R`` at the **right** boundary. - - The incident intensity is taken from the drive amplitude, - ``I0 = (a0 * omega0)**2 / 2`` (the cycle-averaged flux of the launched wave, - in the same field units), so ``R`` and ``T`` are normalization-independent - ratios. The OSIRIS antenna sits at the lower-``x`` boundary, so its current - source contaminates the edge cell; ``guard_cells`` cells are skipped at each - boundary and the flux is averaged over a thin ``window_cells`` slab just - inside. Scalars are the flux averaged over the **last ``last_frac``** of the - run (the saturated phase) divided by ``I0``: ``R``, ``T`` and - ``absorbed = 1 - R - T`` (a steady-state energy balance; valid once the box - is filled and the in-flight transient has passed). - - Returns a dict with the per-dump time series (``t`` and the ``incident`` / - ``reflected`` / ``transmitted`` flux), the scalars (``I0``, ``R``, ``T``, - ``absorbed``, the ``incident_measured`` cross-check) and the settings used. - """ - comps = efield_lr_components(run_dir) - if not comps: - raise RuntimeError("no transverse field pairs (need e2/b3 or e3/b2) for the budget") - sample = next(iter(comps.values()))["right"] - xdim = next(d for d in sample.dims if d != "t") - t = np.asarray(sample.coords["t"].values, dtype=float) - # One-way Poynting flux summed over whichever polarization pairs were dumped. - I_R = sum(d["right"] ** 2 for d in comps.values()) - I_L = sum(d["left"] ** 2 for d in comps.values()) - n = I_R.sizes[xdim] - g, w = guard_cells, window_cells - if g + w > n: - g, w = 0, min(w, n) - left, right = slice(g, g + w), slice(n - g - w, n - g) - inc_t = I_R.isel({xdim: left}).mean(dim=xdim).values # right-going @ left (incident) - refl_t = I_L.isel({xdim: left}).mean(dim=xdim).values # left-going @ left (reflected) - trans_t = I_R.isel({xdim: right}).mean(dim=xdim).values # right-going @ right (transmitted) - - i0 = (a0 * omega0) ** 2 / 2.0 - t_lo = t[0] + (1.0 - last_frac) * (t[-1] - t[0]) - win = t >= t_lo - refl = float(np.mean(refl_t[win])) - trans = float(np.mean(trans_t[win])) - inc_m = float(np.mean(inc_t[win])) - r, tr = refl / i0, trans / i0 - return { - "a0": float(a0), "omega0": float(omega0), "I0": i0, - "t": t, "incident_t": inc_t, "reflected_t": refl_t, "transmitted_t": trans_t, - "t_window": (float(t_lo), float(t[-1])), "last_frac": float(last_frac), - "reflected": refl, "transmitted": trans, "incident_measured": inc_m, - "R": r, "T": tr, "absorbed": 1.0 - r - tr, - "guard_cells": g, "window_cells": w, "pairs": list(comps), - } - - -def _moving_average(y: np.ndarray, k: int) -> np.ndarray: - """Centered moving average over ``k`` samples, edge-normalized (no droop).""" - if k <= 1 or y.size < 3: - return y - ker = np.ones(min(k, y.size)) - num = np.convolve(y, ker, mode="same") - den = np.convolve(np.ones_like(y), ker, mode="same") - return num / den - - -def plot_energy_budget_figure(budget: dict, *, smooth_dumps: int = 21) -> plt.Figure: - r"""Transmitted / reflected / absorbed laser power vs time (fractions of ``I0``). - - Stacked-area energy partition over the run, each curve smoothed over - ``smooth_dumps`` dumps to suppress the ``2 omega`` intensity ripple (the - field dump cadence undersamples the laser period). The shaded band marks the - averaging window used for the scalar ``R`` / ``T``. Early on the partition is - dominated by "absorbed", which there is really energy still **filling the - box** (a transit-time transient), not true absorption — read the absorbed - curve only once it has settled. - """ - t, i0 = budget["t"], budget["I0"] - R = _moving_average(budget["reflected_t"], smooth_dumps) / i0 - T = _moving_average(budget["transmitted_t"], smooth_dumps) / i0 - A = np.clip(1.0 - R - T, 0.0, None) - - fig, ax = plt.subplots(figsize=(7, 4)) - ax.stackplot( - t, T, R, A, - labels=["transmitted", "reflected", "absorbed / in-flight"], - colors=["#4c9be8", "#e8744c", "#c2c2c2"], alpha=0.9, - ) - ax.axvspan(*budget["t_window"], color="k", alpha=0.08, lw=0, - label=f"scalar window (last {100*budget['last_frac']:g}%)") - ax.axhline(1.0, color="0.4", lw=0.6, ls="--") - ax.set_xlim(float(t[0]), float(t[-1])) - ax.set_ylim(0.0, 1.15) - ax.set_xlabel(r"$t$ [$1/\omega_p$]") - ax.set_ylabel(r"fraction of incident $I_0$") - ax.set_title( - rf"laser energy budget (incident $I_0=(a_0\omega_0)^2/2$, " - rf"$a_0={budget['a0']:g}$, $\omega_0={budget['omega0']:g}$)" - ) - ax.legend(loc="upper left", fontsize=8, ncol=2, framealpha=0.75) - fig.tight_layout() - return fig - - -def format_laser_energy_budget(budget: dict) -> str: - """Render :func:`laser_energy_budget` output as a human-readable ``.txt``.""" - b = budget - lo, hi = b["t_window"] - def pct(x: float) -> str: - return f"{100.0 * x:7.2f}%" - return "\n".join([ - "OSIRIS laser energy budget", - "==========================", - "incident intensity I0 = (a0 * omega0)^2 / 2 [code units, cycle-averaged]", - f" a0 = {b['a0']:g}", - f" omega0 = {b['omega0']:g}", - f" I0 = {b['I0']:.6e}", - "", - f"scalars: boundary flux averaged over the last {100*b['last_frac']:g}% of the run " - f"(t in [{lo:.1f}, {hi:.1f}]), divided by I0", - f" reflected (left-going @ left boundary) R = {pct(b['R'])} " - f"(flux {b['reflected']:.4e})", - f" transmitted (right-going @ right boundary) T = {pct(b['T'])} " - f"(flux {b['transmitted']:.4e})", - f" absorbed = 1 - R - T = {pct(b['absorbed'])}", - "", - f"cross-check: measured incident (right-going @ left boundary) = " - f"{b['incident_measured']:.4e} ({b['incident_measured']/b['I0']:.3f} x I0)", - f"polarization pairs used: {', '.join(b['pairs'])}", - f"boundary sampling: {b['guard_cells']} guard cell(s) skipped (antenna at lower-x), " - f"flux averaged over a {b['window_cells']}-cell slab", - "", - ]) - - # --- driver --------------------------------------------------------------- @@ -1322,14 +955,12 @@ def canned_plot_kwargs(output_cfg: dict | None) -> dict: Shared by the live post-processing path (``post.collect``) and the offline regeneration harness (``regen``) so both honour the same knobs: - ``v_th``, ``dist_cells`` and ``omega_k_zoom`` (which may be explicitly - ``null`` to disable the zoom). Keys that are absent fall through to the - ``save_canned_plots`` defaults. + ``v_th`` and ``omega_k_zoom`` (which may be explicitly ``null`` to disable + the zoom). Keys that are absent fall through to the ``save_canned_plots`` + defaults. """ output_cfg = output_cfg or {} kwargs: dict = {"v_th": output_cfg.get("v_th")} - if output_cfg.get("dist_cells") is not None: - kwargs["dist_cells"] = int(output_cfg["dist_cells"]) if "omega_k_zoom" in output_cfg: # may be explicitly null to disable zoom kwargs["omega_k_zoom"] = output_cfg["omega_k_zoom"] return kwargs @@ -1362,10 +993,7 @@ def save_canned_plots( v_th: float | None = None, dpi: int = 120, n_panels: int = 8, - dist_cells: int = 150, omega_k_zoom: float | None = 4.0, - a0: float | None = None, - omega0: float | None = None, ) -> dict[str, Path]: """Generate the standard set of PNGs for a finished OSIRIS run. @@ -1373,11 +1001,6 @@ def save_canned_plots( best-effort: a failure on one diagnostic logs and is skipped rather than aborting the rest. - When ``a0`` and ``omega0`` (the drive amplitude / frequency) are given, the - laser energy budget is also produced: the ``energy_budget`` time-series plot - and a ``laser_energy_budget.txt`` summary alongside it (see - :func:`laser_energy_budget`). - ``run_dir`` may be a raw OSIRIS run directory (with an ``MS/`` HDF5 tree and optional ``HIST/``) **or** the ``binary/`` directory of saved NetCDFs written by :func:`io.save_run_datasets`. The data source is detected @@ -1387,11 +1010,9 @@ def save_canned_plots( regenerated from the saved NetCDF artifacts alone, with no rerun and no raw HDF5 dumps. - ``dist_cells`` sets how many right-boundary cells the phase-space - distribution lineouts average over; ``omega_k_zoom`` is the ``(k, ω)`` - half-width (in ``ω_p`` units) for the equal-aspect lower panel of the - dispersion plots, where ``ω = k`` is drawn at 45° (clamped to the data's - Nyquist; ``None`` → full Nyquist window). + ``omega_k_zoom`` is the ``(k, ω)`` half-width (in ``ω_p`` units) for the + equal-aspect lower panel of the dispersion plots, where ``ω = k`` is drawn + at 45° (clamped to the data's Nyquist; ``None`` → full Nyquist window). """ run_dir = Path(run_dir) out_dir = Path(out_dir) @@ -1550,25 +1171,6 @@ def _write(fig: plt.Figure, rel: str) -> Path: except Exception as e: print(f"[plots] skipping phasespace evolution {diag_rel}: {e}") - # f(p) and delta-f averaged over the right-boundary cells, vs time — - # only for x-p phase spaces (a p-p space like p1p2 has no spatial axis). - if ser.ndim == 3 and any(str(d).startswith("x") for d in ser.dims if d != "t"): - try: - written[f"distribution_lineouts/{species}/{ps_name}"] = _write( - plot_distribution_lineout_figure(ser, n_cells=dist_cells), - f"distribution_lineouts/{species}/{ps_name}.png", - ) - except Exception as e: - print(f"[plots] skipping distribution lineout {diag_rel}: {e}") - try: - fig, ax = plt.subplots(figsize=(6, 4)) - plot_deltaf_lineout(ser, ax=ax, n_cells=dist_cells) - written[f"deltaf_lineouts/{species}/{ps_name}"] = _write( - fig, f"deltaf_lineouts/{species}/{ps_name}.png" - ) - except Exception as e: - print(f"[plots] skipping delta-f lineout {diag_rel}: {e}") - # --- Left/right-going transverse E-field decomposition --- try: for comp, fig in plot_field_lr_decomposition(run_dir).items(): @@ -1602,19 +1204,6 @@ def _write(fig: plt.Figure, rel: str) -> Path: except Exception as e: print(f"[plots] skipping total_energy_vs_time: {e}") - # --- Laser energy budget (reflected/transmitted/absorbed); needs a0, ω0 --- - if a0 is not None and omega0 is not None: - try: - budget = laser_energy_budget(run_dir, a0=a0, omega0=omega0) - written["energy_budget"] = _write( - plot_energy_budget_figure(budget), "energy_budget.png" - ) - (out_dir / "laser_energy_budget.txt").write_text( - format_laser_energy_budget(budget) - ) - except Exception as e: - print(f"[plots] skipping laser energy budget: {e}") - return written @@ -1745,3 +1334,15 @@ def _field_energy_from_series(ser: xr.DataArray) -> np.ndarray: spatial = [d for d in ser.dims if d != "t"] sq = (ser.astype("float64") ** 2).sum(dim=spatial) return 0.5 * np.asarray(sq.values, dtype="float64") * _cell_volume(ser) + + +# --- public helpers for downstream post-processing ------------------------ +# Project repos (e.g. osiris-lpi) build their own canned plots on these shared +# label / decoration / box helpers (and on `efield_lr_components`); export +# stable public names so they need not reach into the private (underscore) API. +decorate = _decorate +ensure_series = _ensure_series +axis_label = _axis_label +display_name = _display_name +tex = _tex +sim_box_xmax = _sim_box_xmax diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 623fbeb8..8f2d0ee9 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -149,30 +149,16 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: except Exception as e: print(f"[post] HIST energy unavailable: {e}") - # Drive amplitude/frequency from the deck → laser energy budget. - deck = cfg.get("deck") or {} - a0, omega0 = deck.get("antenna.a0"), deck.get("antenna.omega0") - - # Render the standard plot set into td/plots so MLflow logs them as - # artifacts. Never let a plotting failure abort metric/artifact logging. + # Render the standard (solver-agnostic) plot set into td/plots so MLflow + # logs them as artifacts. SRS-specific views (laser energy budget, + # distribution lineouts) are layered on in osiris-lpi's OsirisLPI subclass. + # Never let a plotting failure abort metric/artifact logging. try: kwargs = _plots.canned_plot_kwargs(cfg.get("output")) - if a0 is not None and omega0 is not None: - kwargs.update(a0=float(a0), omega0=float(omega0)) _plots.save_canned_plots(run_dir, td / "plots", **kwargs) except Exception as e: print(f"[post] plotting failed: {e}") - # Reflected/transmitted/absorbed laser-energy fractions as scalar metrics. - if a0 is not None and omega0 is not None: - try: - budget = _plots.laser_energy_budget(run_dir, a0=float(a0), omega0=float(omega0)) - metrics["laser_reflectivity"] = budget["R"] - metrics["laser_transmissivity"] = budget["T"] - metrics["laser_absorbed_frac"] = budget["absorbed"] - except Exception as e: - print(f"[post] laser energy budget unavailable: {e}") - # Always include the rendered deck + stdout/stderr. for fname in ("os-stdin", "stdout.log", "stderr.log"): src = run_dir / fname diff --git a/adept/osiris/regen.py b/adept/osiris/regen.py index bcf684bf..b3dad320 100644 --- a/adept/osiris/regen.py +++ b/adept/osiris/regen.py @@ -75,31 +75,6 @@ def load_output_cfg(src: str | Path) -> dict: return {} -def load_antenna_params(src: str | Path) -> dict: - """Best-effort read of the drive amplitude / frequency (``a0``, ``omega0``). - - Reads the flat ``deck`` block of ``derived_config.yaml`` (falling back to - ``config.yaml``) at ``src`` or its parent, returning ``{"a0":…, "omega0":…}`` - for the single antenna when both are present, else ``{}``. The laser energy - budget needs these; absent them (e.g. a non-antenna deck) it is skipped. - """ - src = Path(src) - for name in ("derived_config.yaml", "config.yaml"): - for candidate in (src / name, src.parent / name): - if candidate.is_file(): - try: - import yaml - - cfg = yaml.safe_load(candidate.read_text()) or {} - except Exception: # a malformed config must not block regeneration - continue - deck = cfg.get("deck") or {} - a0, omega0 = deck.get("antenna.a0"), deck.get("antenna.omega0") - if a0 is not None and omega0 is not None: - return {"a0": float(a0), "omega0": float(omega0)} - return {} - - def regenerate( src: str | Path, out_dir: str | Path | None = None, @@ -112,18 +87,17 @@ def regenerate( ``src`` is a run directory (containing ``binary/`` and optionally ``config.yaml`` / ``derived_config.yaml``) or a ``binary/`` NetCDF directory. Plot knobs default to the run's ``output:`` config (unless - ``use_config=False``) and the drive ``a0`` / ``omega0`` are read from - ``derived_config.yaml`` for the laser energy budget; any keyword in - ``overrides`` (``v_th``, ``dist_cells``, ``omega_k_zoom``, ``dpi``, - ``n_panels``, ``a0``, ``omega0``) is applied on top verbatim — pass only the - knobs you mean to set. Writes PNGs (and ``laser_energy_budget.txt`` when - ``a0``/``omega0`` are known) under ``out_dir`` (default + ``use_config=False``); any keyword in ``overrides`` (``v_th``, + ``omega_k_zoom``, ``dpi``, ``n_panels``) is applied on top verbatim — pass + only the knobs you mean to set. Writes PNGs under ``out_dir`` (default ``/plots_regen``) and returns the ``{plot-name: path}`` map. + + SRS-specific plots (laser energy budget, distribution lineouts) live in + osiris-lpi; regenerate those with ``python -m osiris_lpi.regen``. """ binary = find_binary_dir(src) out_dir = Path(out_dir) if out_dir is not None else default_out_dir(src) kwargs = _plots.canned_plot_kwargs(load_output_cfg(src) if use_config else None) - kwargs.update(load_antenna_params(src)) # a0/omega0 for the laser energy budget kwargs.update(overrides) return _plots.save_canned_plots(binary, out_dir, **kwargs) @@ -144,8 +118,6 @@ def _cli_overrides(args: argparse.Namespace) -> dict: overrides: dict = {} if args.v_th is not None: overrides["v_th"] = args.v_th - if args.dist_cells is not None: - overrides["dist_cells"] = args.dist_cells if args.dpi is not None: overrides["dpi"] = args.dpi if args.n_panels is not None: @@ -154,10 +126,6 @@ def _cli_overrides(args: argparse.Namespace) -> dict: overrides["omega_k_zoom"] = None # explicit disable elif args.omega_k_zoom is not None: overrides["omega_k_zoom"] = args.omega_k_zoom - if args.a0 is not None: - overrides["a0"] = args.a0 - if args.omega0 is not None: - overrides["omega0"] = args.omega0 return overrides @@ -172,20 +140,12 @@ def main(argv: list[str] | None = None) -> int: help="ignore the run's config.yaml output block") ap.add_argument("--v-th", type=float, default=None, help="electron thermal velocity for the Langmuir overlay on omega-k plots") - ap.add_argument("--dist-cells", type=int, default=None, - help="right-boundary cells averaged for the f(p) distribution lineouts") ap.add_argument("--omega-k-zoom", type=float, default=None, help="(k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel") ap.add_argument("--no-zoom", action="store_true", help="use the full Nyquist window for the lower omega-k panel (omega_k_zoom=None)") ap.add_argument("--dpi", type=int, default=None, help="figure DPI") ap.add_argument("--n-panels", type=int, default=None, help="panels for faceted plots") - ap.add_argument("--a0", type=float, default=None, - help="drive amplitude for the laser energy budget " - "(overrides derived_config.yaml)") - ap.add_argument("--omega0", type=float, default=None, - help="drive frequency for the laser energy budget " - "(overrides derived_config.yaml)") args = ap.parse_args(argv) out_dir = Path(args.out) if args.out else default_out_dir(args.src) diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index 80f501a0..62ed5c09 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -19,5 +19,5 @@ osiris: output: # diagnostics_to_log: [e1] # null/omitted = convert every diagnostic's full time history to netCDF # v_th: 0.1 # overlay the Langmuir (Bohm-Gross) branch on omega-k plots - # dist_cells: 10 # right-boundary cells averaged for the phase-space f(p) lineouts + # dist_cells: 10 # (osiris-lpi OsirisLPI only) right-boundary cells averaged for the f(p) lineouts # omega_k_zoom: 4.0 # (k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel; null = full diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index 9fc69c47..449cd3ce 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -101,8 +101,6 @@ output: diagnostics_to_log: null # null = all; or [e1, charge, …] v_th: 0.1 # optional: overlays the Bohm–Gross # Langmuir branch on ω–k plots - dist_cells: 150 # right-boundary cells averaged for - # the phase-space f(p) lineouts omega_k_zoom: 4.0 # (k, ω) half-width [ω_p] for the # equal-aspect lower ω–k panel # (clamped to Nyquist); null = full @@ -136,14 +134,12 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | `profiles//density.png` | density profile vs `x` (final snapshot + late-time mean) | | `profiles//temperature.png` | temperature profile vs `x`, from `uth1/2/3` or `T11/22/33` moments (omitted if neither is dumped) | | `phasespace//.png`, `phasespace_evolution/…` | `(x, p)` phase-space heatmaps | -| `distribution_lineouts//.png` | `f(p)` averaged over the rightmost `dist_cells` cells, overlaid at sampled times — stacked linear `f`, `\|f\|` log, and `δf = f - f_M` panels | | `field_decomp/.png` | left/right-going transverse `E` (vacuum Riemann split `(e2±b3)/2`, `(e3∓b2)/2`), spacetime + `ω–k` | | `energy_vs_time.png`, `energy_components_vs_time.png`, `total_energy_vs_time.png` | field / kinetic energy traces | -| `energy_budget.png` + `laser_energy_budget.txt` | reflected / transmitted / absorbed laser power vs time (stacked-area), with scalar `R`/`T`/absorbed in the `.txt` — emitted only when the drive `a0`/`omega0` are known | > **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 longitudinal `e1` is electrostatic and is intentionally excluded. -> **Laser energy budget.** Reflected = left-going Poynting flux at the left boundary, transmitted = right-going flux at the right boundary (from the same Riemann split), compared to the incident intensity `I₀ = (a0·ω0)²/2` from the drive. Scalars are the boundary flux averaged over the last 25 % of the run (the saturated phase) ÷ `I₀`; `absorbed = 1 − R − T`. The antenna sits at the lower-`x` boundary, so its source cell is skipped. The same dispersive-medium caveat as `field_decomp/` applies — the split is approximate where the boundaries sit in plasma. `R`/`T`/absorbed are also logged to MLflow as `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac`. +> **SRS-specific plots & metrics live in osiris-lpi.** The laser-energy budget (reflected / transmitted / absorbed over time: the `energy_budget.png` + `laser_energy_budget.txt` artifacts and the `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac` metrics) and the spatio-temporal distribution-function lineouts (`distribution_lineouts/` + `deltaf_lineouts/`, `f(p)` and `δf` averaged over the rightmost `dist_cells` cells) are **not** produced by adept's general OSIRIS wrapper. They live in the [osiris-lpi](https://github.com/ergodicio/osiris-lpi) repo — `osiris_lpi.OsirisLPI` subclasses `BaseOsiris` and adds them in `post_process`, reading `output.dist_cells` and the drive `antenna.a0`/`antenna.omega0` from the deck. Regenerate them offline from saved NetCDFs with `python -m osiris_lpi.regen`. ## Programmatic use diff --git a/tests/test_osiris/test_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py index 89517f72..b9f0719e 100644 --- a/tests/test_osiris/test_plots_new_views.py +++ b/tests/test_osiris/test_plots_new_views.py @@ -3,7 +3,6 @@ Synthesizes tiny OSIRIS-shaped runs (no real solver needed) to exercise: - proper-LaTeX axis/value labels (``_tex`` wrapping) - - phase-space distribution lineouts averaged over the right-boundary cells - the zoomed (k, ω) dispersion view with the light line - combined J_x/J_y/J_z (j1/j2/j3) current figures - per-species density + temperature profiles @@ -21,7 +20,6 @@ import h5py import numpy as np -import xarray as xr from adept.osiris import io as oio from adept.osiris import plots as oplt @@ -111,31 +109,6 @@ def test_axis_label_is_math_mode(tmp_path: Path) -> None: assert oplt._axis_label(da, "x1") == r"$x_1$ [$c / \omega_p$]" -# --- distribution lineouts ------------------------------------------------ - - -def test_distribution_lineout_averages_right_cells() -> None: - p = np.linspace(-2, 2, 11) - x = np.linspace(0, 10, 16) - fp = np.exp(-(p**2)) - data = np.tile(fp[:, None], (1, x.size)) # f independent of x - da = xr.DataArray( - data, coords={"p1": p, "x1": x}, dims=("p1", "x1"), name="x1p1", - attrs={"axis_units": {"p1": r"m_e c", "x1": r"c / \omega_p"}, - "axis_long_names": {"p1": "p_1", "x1": "x_1"}, "units": "a.u."}, - ) - ax = oplt.plot_distribution_lineout(da, n_cells=10) - y = ax.lines[0].get_ydata() - np.testing.assert_allclose(y, fp, atol=1e-12) - - -def test_distribution_lineout_time_series(tmp_path: Path) -> None: - run_dir = _make_rich_run(tmp_path) - ser = oio.load_series(run_dir / "MS/PHA/x1p1/electron") - ax = oplt.plot_distribution_lineout(ser, n_cells=5, n_times=4) - assert len(ax.lines) == 4 # one curve per sampled time - - # --- zoomed omega-k ------------------------------------------------------- @@ -237,7 +210,6 @@ def test_save_canned_plots_emits_new_views(tmp_path: Path) -> None: "currents/lineouts", "profiles/electron/density", "profiles/electron/temperature", - "distribution_lineouts/electron/x1p1", "field_decomp/e2", "field_decomp/e3", } @@ -245,45 +217,3 @@ def test_save_canned_plots_emits_new_views(tmp_path: Path) -> None: for path in written.values(): assert path.exists() and path.stat().st_size > 0 - -# --- laser energy budget -------------------------------------------------- - - -def test_laser_energy_budget_pure_right_going_has_no_reflection(tmp_path: Path) -> None: - run_dir = _make_rich_run(tmp_path) - # Fixture fields are a pure right-going wave (e2=b3, e3=-b2), so the - # left-going Riemann parts vanish identically: nothing is "reflected". - # Averaging over the whole grid (one full wavelength) and all time gives - # = <2 sin^2> ~ 1, matching I0 = (a0*omega0)^2/2 = 1 -> T ~ 1. - b = oplt.laser_energy_budget( - run_dir, a0=float(np.sqrt(2.0)), omega0=1.0, - last_frac=1.0, guard_cells=0, window_cells=16, - ) - assert b["pairs"] == ["e2", "e3"] - assert b["R"] < 1e-9 # no left-going flux at all - assert b["transmitted"] > 0.0 - # window = whole grid -> left/right slabs coincide -> equal fluxes - assert abs(b["transmitted"] - b["incident_measured"]) < 1e-9 - assert 0.7 < b["T"] < 1.3 - assert abs(b["absorbed"] - (1.0 - b["R"] - b["T"])) < 1e-12 - - -def test_save_canned_plots_emits_energy_budget_with_drive(tmp_path: Path) -> None: - run_dir = _make_rich_run(tmp_path) - out = tmp_path / "plots" - written = oplt.save_canned_plots(run_dir, out, a0=0.004, omega0=1.0) - assert "energy_budget" in written - assert written["energy_budget"].exists() and written["energy_budget"].stat().st_size > 0 - txt = out / "laser_energy_budget.txt" - assert txt.exists() - body = txt.read_text() - for key in ("reflected", "transmitted", "absorbed", "I0"): - assert key in body - - -def test_save_canned_plots_skips_energy_budget_without_drive(tmp_path: Path) -> None: - run_dir = _make_rich_run(tmp_path) - out = tmp_path / "plots" - written = oplt.save_canned_plots(run_dir, out) # no a0/omega0 -> skipped - assert "energy_budget" not in written - assert not (out / "laser_energy_budget.txt").exists() From ea256c3288a53b8364f79fa4588e86f48a6fc6da Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 15 Jun 2026 14:38:17 -0700 Subject: [PATCH 17/49] Added example decks for testing so that hardcoded paths are no longer used --- tests/test_osiris/decks/F-Tsung_2d_lpi_deck | 139 ++++++++++++++ tests/test_osiris/decks/srs-1d_lpi | 146 +++++++++++++++ tests/test_osiris/decks/srs-lpi_2node | 151 ++++++++++++++++ tests/test_osiris/decks/two-stream-1d | 191 ++++++++++++++++++++ tests/test_osiris/test_deck_roundtrip.py | 20 +- tests/test_osiris/test_io_and_plots.py | 14 +- tests/test_osiris/test_runner.py | 11 +- 7 files changed, 654 insertions(+), 18 deletions(-) create mode 100644 tests/test_osiris/decks/F-Tsung_2d_lpi_deck create mode 100644 tests/test_osiris/decks/srs-1d_lpi create mode 100644 tests/test_osiris/decks/srs-lpi_2node create mode 100644 tests/test_osiris/decks/two-stream-1d diff --git a/tests/test_osiris/decks/F-Tsung_2d_lpi_deck b/tests/test_osiris/decks/F-Tsung_2d_lpi_deck new file mode 100644 index 00000000..0790052d --- /dev/null +++ b/tests/test_osiris/decks/F-Tsung_2d_lpi_deck @@ -0,0 +1,139 @@ +simulation +{ + +} + +node_conf +{ + node_number(1:2)=30,4, + if_periodic(1:2)=.false.,.true., +} + + +grid +{ +nx_p(1:2)=6000,3600, +coordinates="cartesian", +} +time_step +{ +dt=0.1259, +ndump=50, +} +restart +{ +ndump_fac = 5000, +if_restart=.false., +} +space +{ +xmin(1:2)=0.0 ,0.0, +xmax(1:2)=1076.04,645.624, +if_move(1:2)=.false., .false., +} +time +{ +tmin=0.0d0, +tmax=50000.0d0, +} +el_mag_fld +{ +} +emf_bound +{ +type(1:2,1)="vpml","vpml", +vpml_diffuse(1:2,1) = .true., .true., +vpml_bnd_size=45, +} +diag_emf +{ +ndump_fac=10, +reports="e1","e2","e3","b1","b2","b3", +} +particles +{ +num_species=1, +} + +species +{ + num_par_max=2300000, + rqm=-1.0, + num_par_x(1:2) = 8,8, +} + +! uth=0.04423 is 1keV +! uth=0.0 +udist +{ +uth(1:3) =0.076621 , 0.076621, 0.076621, +} + +profile +{ +num_x=6, +fx(1:6,1)=0.0,0.225,0.275,0.0,0.0,0.0, +x(1:6,1)=0,0.00001,1076.03,1076.04,50000,100000, +fx(1:6,2) = 1,1,1,1,1,1, +x(1:6,2) = 0,1,2,3,4,5000, +} + +! +spe_bound +{ +type(1:2,1)="thermal", "thermal", +uth_bnd(1:3,1,1)=0.076621,0.076621,0.076621, +uth_bnd(1:3,2,1)=0.076621,0.076621,0.076621, +thermal_type(1:2,1) = "half max","half max", +} + + +diag_species +{ +ndump_fac_pha=30, +ndump_fac=30, +reports="charge", +ndump_fac_raw=100, +ps_xmin(1:2)=0.0 ,0.0, +ps_xmax(1:2)=1076.04,645.624, +ps_pmin(1:3)=-0.3,-0.3,-0.3, +ps_pmax(1:3)=1.0,1.0, 1.0, +ps_nx(1:2)=1000,500, +ps_np(1:3)=100,100,100, +if_ps_p_auto(1:3)=.true., .true., .true., +phasespaces="p1x1","p1p2","p1p3","x1p1p2", +ps_gammamin=1.0, +ps_gammamax = 2000.0, +ps_ngamma= 1024, +if_ps_gamma_auto=.true., +raw_gamma_limit=1.2, +raw_fraction = 0.1, +n_ene_bins=6, +ene_bins(1:6)=0.05,0.1,0.2,0.25,0.3,0.5, +pha_ene_bin='x1_q1', +} + + +current{} + + +smooth +{ + +} + + +diag_current +{ +} + + +antenna_array +{ +n_antenna=1, +} +antenna +{ +a0=0.0033,t_rise=5,t_flat=50000.67,t_fall=5,omega0=1.0,x0=320,rad_x=500000.0, +side=2, +} \ No newline at end of file diff --git a/tests/test_osiris/decks/srs-1d_lpi b/tests/test_osiris/decks/srs-1d_lpi new file mode 100644 index 00000000..3aa183ff --- /dev/null +++ b/tests/test_osiris/decks/srs-1d_lpi @@ -0,0 +1,146 @@ +simulation +{ + ! reference density = critical density for a 351 nm laser. + ! with omega0=1, the laser frequency equals the reference plasma frequency, + ! so n0 is n_crit(351 nm) = 1.115e21 / 0.351^2 cm^-3. + n0 = 9.05e21, +} + +node_conf +{ + node_number(1)=4, + if_periodic(1)=.false., +} + +grid +{ + nx_p=6000, + coordinates="cartesian", +} +time_step +{ + dt=0.178, + ndump=2, +} +restart +{ + ndump_fac = 0, + if_restart=.false., +} +space +{ + xmin(1)=0.0 , + xmax(1)=1076.04, + if_move(1:2)=.false., .false., +} +time +{ + tmin=0.0d0, + tmax=25000.0d0, +} +el_mag_fld +{ +} +emf_bound +{ + type(1:2,1)="lindmann", "lindmann", +} +diag_emf +{ + ndump_fac=30, + reports="e1","e2","e3","b1","b2","b3", +} +particles +{ + num_species=1, +} + + +species +{ + num_par_max=1000000, + rqm=-1.0, + num_par_x(1) = 256, + n_sort=25, +} + + +udist +{ + uth(1:3) =0.0885 , 0.0885, 0.0885, +} + +profile +{ + num_x=4, + fx(1:6,1)=0.0,0.225,0.275,0.0, + x(1:6,1)=0,0.00001,1076.03,1076.04, +} + + +spe_bound +{ + type(1:2,1)="thermal", "thermal", + uth_bnd(1:3,1,1)=0.0885,0.0885,0.0885, + uth_bnd(1:3,2,1)=0.0885,0.0885,0.0885, + thermal_type(1:2,1) = "half max","half max", +} + + +diag_species +{ + ndump_fac_pha=30, + ndump_fac=30, + reports="charge", + ndump_fac_raw=100, + ps_xmin(1)=0.0, + ps_xmax(1)=1076.04, + ps_pmin(1:3)=-0.3,-0.3,-0.3, + ps_pmax(1:3)=1.0,1.0, 1.0, + ps_nx(1)=1000, + ps_np(1:3)=100,100,100, + if_ps_p_auto(1:3)=.true., .true., .true., + phasespaces="p1x1", "p1p3", "p1p2", + ps_gammamin=1.0, + ps_gammamax = 2000.0, + ps_ngamma= 1024, + if_ps_gamma_auto=.true., + raw_gamma_limit=1.2, + raw_fraction = 0.1, + n_ene_bins=6, + ene_bins(1:6)=0.005,0.04,0.08,0.12,0.2,0.5, + ! pha_ene_bin='x1p1','x1p1p2', +} + + +current{} + + +smooth +{ + type(1)="custom", + order(1)=5, + swfj(1:3,1,1)=1,2,1, + swfj(1:3,2,1)=1,2,1, + swfj(1:3,3,1)=1,2,1, + swfj(1:3,4,1)=1,2,1, + swfj(1:3,5,1)=-5,14,-5, +} + + +diag_current +{ +} + + +antenna_array +{ + n_antenna=1, +} + +antenna +{ + a0=0.004, + t_rise=10,t_flat=50000.0,t_fall=10.0,omega0=1.0, + pol=0, +} diff --git a/tests/test_osiris/decks/srs-lpi_2node b/tests/test_osiris/decks/srs-lpi_2node new file mode 100644 index 00000000..805461cf --- /dev/null +++ b/tests/test_osiris/decks/srs-lpi_2node @@ -0,0 +1,151 @@ +! 2-node (8-rank / 8-GPU) SMOKE deck for the srs-multinode campaign. +! Purpose: validate cross-node MPI + GPU binding (srun -n 8 across 2 nodes, +! MPICH_GPU_SUPPORT_ENABLED) BEFORE the full 12-node run -- NOT a physics run. +! Derived from lpi_12node but deliberately light + short: +! node_number 48 -> 8 (6144/8 = 768 cells/domain; 2 nodes x 4 GPUs) +! num_par_x 32768 -> 256 (light: mechanics test, fast) +! num_par_max 1e7 -> 1e6 (nominal 768*256 = 196k/domain) +! tmax 25000 -> 60 (~345 steps: exercises loop, halo exchange, dumps) +! nx_p=6144 and dt=0.1738 kept identical to the full deck (same grid/CFL). +simulation +{ + n0 = 9.05e21, +} + +node_conf +{ + node_number(1)=8, + if_periodic(1)=.false., +} + +grid +{ + nx_p=6144, + coordinates="cartesian", +} +time_step +{ + dt=0.1738, + ndump=2, +} +restart +{ + ndump_fac = 0, + if_restart=.false., +} +space +{ + xmin(1)=0.0 , + xmax(1)=1076.04, + if_move(1:2)=.false., .false., +} +time +{ + tmin=0.0d0, + tmax=60.0d0, +} +el_mag_fld +{ +} +emf_bound +{ + type(1:2,1)="lindmann", "lindmann", +} +diag_emf +{ + ndump_fac=30, + reports="e1","e2","e3","b1","b2","b3", +} +particles +{ + num_species=1, +} + + +species +{ + num_par_max=1000000, + rqm=-1.0, + num_par_x(1) = 256, + n_sort=25, +} + + +udist +{ + uth(1:3) =0.0885 , 0.0885, 0.0885, +} + +profile +{ + num_x=4, + fx(1:6,1)=0.0,0.225,0.275,0.0, + x(1:6,1)=0,0.00001,1076.03,1076.04, +} + + +spe_bound +{ + type(1:2,1)="thermal", "thermal", + uth_bnd(1:3,1,1)=0.0885,0.0885,0.0885, + uth_bnd(1:3,2,1)=0.0885,0.0885,0.0885, + thermal_type(1:2,1) = "half max","half max", +} + + +diag_species +{ + ndump_fac_pha=30, + ndump_fac=30, + reports="charge", + ndump_fac_raw=0, + ps_xmin(1)=0.0, + ps_xmax(1)=1076.04, + ps_pmin(1:3)=-0.3,-0.3,-0.3, + ps_pmax(1:3)=1.0,1.0, 1.0, + ps_nx(1)=1000, + ps_np(1:3)=100,100,100, + if_ps_p_auto(1:3)=.true., .true., .true., + phasespaces="p1x1", "p1p3", "p1p2", + ps_gammamin=1.0, + ps_gammamax = 2000.0, + ps_ngamma= 1024, + if_ps_gamma_auto=.true., + raw_gamma_limit=1.2, + raw_fraction = 0.1, + n_ene_bins=6, + ene_bins(1:6)=0.005,0.04,0.08,0.12,0.2,0.5, +} + + +current{} + + +smooth +{ + type(1)="custom", + order(1)=5, + swfj(1:3,1,1)=1,2,1, + swfj(1:3,2,1)=1,2,1, + swfj(1:3,3,1)=1,2,1, + swfj(1:3,4,1)=1,2,1, + swfj(1:3,5,1)=-5,14,-5, +} + + +diag_current +{ +} + + +antenna_array +{ + n_antenna=1, +} + +antenna +{ + a0=0.004, + t_rise=10,t_flat=50000.0,t_fall=10.0,omega0=1.0, + pol=0, +} diff --git a/tests/test_osiris/decks/two-stream-1d b/tests/test_osiris/decks/two-stream-1d new file mode 100644 index 00000000..afeba0fd --- /dev/null +++ b/tests/test_osiris/decks/two-stream-1d @@ -0,0 +1,191 @@ +! Two-stream instability (1D, electrostatic) +! Translated from pic-1d YAML to OSIRIS +! +! UNIT CONVERSION: The source YAML normalizes to n_ref=1e21/cc, T_ref=1eV. +! Lengths are assumed to be in units of c/ωp; velocities in units of c. +! If velocities are instead normalized to vth(T_ref)=sqrt(kT_ref/m_e)≈0.0014c, +! scale ufl and uth by ~0.0014. +! +! SOLVER: OSIRIS has no 1D Poisson solver accessible from the deck. +! The Yee EM solver with periodic BCs is equivalent for longitudinal +! electrostatic modes in a quasineutral plasma: Ampere's law gives +! ∂E1/∂t = -J1, and charge-conserving deposition (Esirkepov) maintains +! Gauss's law throughout the run. +! +! LOADING: OSIRIS loads particles uniformly within cells (equidistant). +! The source code used random positions (loading: random); velocity +! sampling noise from the Maxwellian still seeds the instability. + +!--------the node configuration for this simulation-------- +node_conf +{ + node_number(1:1) = 1, + if_periodic(1:1) = .true., +} + +!----------spatial grid---------- +grid +{ + nx_p(1:1) = 64, + coordinates = "cartesian", +} + +!----------time step and global data dump timestep number---------- +time_step +{ + dt = 0.05, + ndump = 1, ! base dump period; diagnostics multiply this via ndump_fac +} + +!----------restart information---------- +restart +{ + ndump_fac = 0, + if_restart = .false., +} + +!----------spatial limits of the simulation---------- +space +{ + xmin(1:1) = 0.0, + xmax(1:1) = 10.26566, + if_move(1:1) = .false., +} + +!----------time limits---------- +time +{ + tmin = 0.0, + tmax = 100.0, +} + +!----------field solver---------- +! drivers: ex and ey in source YAML are empty (no external drivers) +el_mag_fld +{ + ext_fld = "none", +} + +!----------boundary conditions for EM fields---------- +emf_bound +{ +!--- type(1:2,1) = "periodic", "periodic", +} + +!----------EM field diagnostics---------- +! fields: 601 outputs over [0, 30] → every step → ndump_fac = 1 +! e1 = Ex, the electrostatic field in 1D +diag_emf +{ + ndump_fac = 1, + reports = "e1", +} + +!----------number of particle species---------- +! TSC (Triangular-Shaped Cloud) = quadratic interpolation in OSIRIS +particles +{ + interpolation = "quadratic", + num_species = 2, +} + +!----------species 1: beam_pos---------- +species +{ + name = "beam_pos", + num_par_max = 32768, ! 64 cells × 512 ppc + rqm = -1.0, ! charge/mass (electrons: q=-e, m=me → -1 in code units) + num_par_x(1:1) = 512, +} + +udist +{ + uth(1:3) = 0.00447, 0.00447, 0.00447, ! sqrt(T0/me c^2) + ufl(1:3) = 0.025, 0.0, 0.0, ! 0.025 ≈ 5.6x uth (clear two-stream separation) +} + +profile +{ + density = 0.5, ! baseline = 0.5 n_ref + fx(1:2,1) = 1.0, 1.0, + x(1:2,1) = 0.0, 10.26566, +} + +spe_bound +{ +!---- type(1:2,1) = "periodic", "periodic", +} + +! beam_pos: 11 outputs over [0, 30] → every 60 steps → ndump_fac_pha = 60 +diag_species +{ + ndump_fac_pha = 60, + reports = "charge", + + ps_xmin(1:1) = 0.0, + ps_xmax(1:1) = 10.26566, + ps_nx(1:1) = 64, + + ps_pmin(1:1) = -0.05, + ps_pmax(1:1) = 0.05, + ps_np(1:1) = 256, + if_ps_p_auto(1:3) = .false., .true., .true., + + phasespaces = "x1p1", +} + +!----------species 2: beam_neg---------- +species +{ + name = "beam_neg", + num_par_max = 32768, ! 64 cells × 512 ppc + rqm = -1.0, + num_par_x(1:1) = 512, +} + +udist +{ + uth(1:3) = 0.00447, 0.00447, 0.00447, ! sqrt(T0/me c^2) + ufl(1:3) = -0.025, 0.0, 0.0, ! -0.025 ≈ 5.6x uth (clear two-stream separation) +} + +profile +{ + density = 0.5, ! baseline = 0.5 n_ref + fx(1:2,1) = 1.0, 1.0, + x(1:2,1) = 0.0, 10.26566, +} + +spe_bound +{ +!---- type(1:2,1) = "periodic", "periodic", +} + +! beam_neg: 11 outputs over [0, 30] → every 60 steps +diag_species +{ + ndump_fac_pha = 60, + reports = "charge", + + ps_xmin(1:1) = 0.0, + ps_xmax(1:1) = 10.26566, + ps_nx(1:1) = 64, + + ps_pmin(1:1) = -0.05, + ps_pmax(1:1) = 0.05, + ps_np(1:1) = 256, + if_ps_p_auto(1:3) = .false., .true., .true., + + phasespaces = "x1p1", +} + +!----------current smoothing---------- +smooth +{ + type = "none", +} + +!----------current diagnostics---------- +diag_current +{ +} diff --git a/tests/test_osiris/test_deck_roundtrip.py b/tests/test_osiris/test_deck_roundtrip.py index 19190f00..7991c055 100644 --- a/tests/test_osiris/test_deck_roundtrip.py +++ b/tests/test_osiris/test_deck_roundtrip.py @@ -9,16 +9,20 @@ from adept.osiris import deck as osd +DECKS_DIR = Path(__file__).parent / "decks" + +# Real OSIRIS decks vendored into the repo so the round-trip is exercised in +# CI without depending on any developer's local checkout. REAL_DECKS = [ - "/home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d", - "/home/phil/Desktop/pic/projects/twostream-pic/example_deck", - "/home/phil/Desktop/pic/osiris/decks/cuda/os-stdin-1d-therm", - "/home/phil/Desktop/pic/osiris/decks/cuda/os-stdin-2d-therm", + DECKS_DIR / "two-stream-1d", + DECKS_DIR / "srs-1d_lpi", + DECKS_DIR / "srs-lpi_2node", + DECKS_DIR / "F-Tsung_2d_lpi_deck", ] -@pytest.mark.parametrize("path", REAL_DECKS) -def test_roundtrip_identity(path: str) -> None: +@pytest.mark.parametrize("path", REAL_DECKS, ids=lambda p: p.name) +def test_roundtrip_identity(path: Path) -> None: text = Path(path).read_text() parsed = osd.parse_deck(text) rendered = osd.render_deck(parsed) @@ -127,9 +131,7 @@ def test_deck_to_flat_dict_keys_are_mlflow_safe() -> None: import re s = osd.parse_deck( - Path( - "/home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d" - ).read_text() + (DECKS_DIR / "two-stream-1d").read_text() ) flat = osd.deck_to_flat_dict(s) allowed = re.compile(r"^[A-Za-z0-9_./:\- ]+$") diff --git a/tests/test_osiris/test_io_and_plots.py b/tests/test_osiris/test_io_and_plots.py index 5dd2b998..9d449ba4 100644 --- a/tests/test_osiris/test_io_and_plots.py +++ b/tests/test_osiris/test_io_and_plots.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os from pathlib import Path import matplotlib @@ -16,16 +17,17 @@ from adept.osiris import plots as oplt -EXISTING_RUN = Path( - "/home/phil/Desktop/pic/projects/twostream-pic/run0001" -) +# Point this at a completed two-stream OSIRIS run to exercise the io/plots +# loaders against real data; the tests skip cleanly when it is unset. +EXISTING_RUN_ENV = "OSIRIS_TWOSTREAM_RUN" @pytest.fixture(scope="module") def run_dir() -> Path: - if not EXISTING_RUN.is_dir(): - pytest.skip("Existing twostream run not present") - return EXISTING_RUN + raw = os.environ.get(EXISTING_RUN_ENV) + if not raw or not Path(raw).is_dir(): + pytest.skip(f"set {EXISTING_RUN_ENV} to an existing two-stream run dir") + return Path(raw) def test_load_grid_h5_shape_and_coords(run_dir: Path) -> None: diff --git a/tests/test_osiris/test_runner.py b/tests/test_osiris/test_runner.py index 1371b116..388299eb 100644 --- a/tests/test_osiris/test_runner.py +++ b/tests/test_osiris/test_runner.py @@ -10,6 +10,11 @@ from adept.osiris import runner +# Resolved from the same env vars the runner itself honors, so the live-binary +# smoke test below runs wherever OSIRIS is built and skips cleanly otherwise. +OSIRIS_BIN_1D = os.environ.get("OSIRIS_BIN_1D") or os.environ.get("OSIRIS_BIN") + + def test_discover_binary_explicit_wins(tmp_path: Path) -> None: fake = tmp_path / "fake-osiris" fake.write_text("") @@ -41,8 +46,8 @@ def test_run_osiris_missing_binary_raises(tmp_path: Path) -> None: @pytest.mark.skipif( - not Path("/home/phil/Desktop/pic/osiris/bin/osiris-1D.e").exists(), - reason="osiris-1D.e not built", + not (OSIRIS_BIN_1D and Path(OSIRIS_BIN_1D).exists()), + reason="set OSIRIS_BIN_1D (or OSIRIS_BIN) to a built osiris-1D.e to run", ) def test_run_osiris_invalid_deck_raises(tmp_path: Path) -> None: # A deck with a recognized section but garbage inside: OSIRIS exits 0 @@ -51,7 +56,7 @@ def test_run_osiris_invalid_deck_raises(tmp_path: Path) -> None: with pytest.raises(RuntimeError) as excinfo: runner.run_osiris( "node_conf { node_number(1:1) = junk_value, }", - binary="/home/phil/Desktop/pic/osiris/bin/osiris-1D.e", + binary=OSIRIS_BIN_1D, mpi_ranks=1, run_root=tmp_path, ) From 33068481a52ac84550dd0459fd89070ee2faf356 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 16 Jun 2026 16:50:13 -0700 Subject: [PATCH 18/49] Fixed plots when momentum is autoscaled by osiris --- adept/osiris/io.py | 79 ++++++++++++++++++++++--- adept/osiris/plots.py | 67 +++++++++++++++------ configs/osiris/twostream-1d.yaml | 1 - docs/osiris-adept-usage.md | 2 +- tests/test_osiris/test_post_netcdf.py | 84 +++++++++++++++++++++++++++ 5 files changed, 204 insertions(+), 29 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 85a64a56..facf8143 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -293,12 +293,25 @@ def load_series(directory: str | Path) -> xr.DataArray: raise FileNotFoundError(f"No .h5 dumps in {directory}") first = load_grid_h5(dumps[0]) - times = np.empty(len(dumps), dtype="float64") - iters = np.empty(len(dumps), dtype="int64") - data = np.empty((len(dumps), *first.shape), dtype=first.dtype) - data[0] = first.values - times[0] = first.attrs["time"] - iters[0] = first.attrs["iter"] + n = len(dumps) + times = np.empty(n, dtype="float64") + iters = np.empty(n, dtype="int64") + data = np.empty((n, *first.shape), dtype=first.dtype) + # Per-dump axis bounds (min, max) for every non-time dim. OSIRIS autoscale + # (deck ``if_ps_p_auto`` / ``if_ps_gamma_auto``) re-picks a phase space's + # momentum / gamma bounds *every dump*, so an axis can move dump-to-dump + # while its bin count stays fixed — a shape-only check misses it. + bounds = {d: np.empty((n, 2), dtype="float64") for d in first.dims} + + def _record(i: int, da: xr.DataArray) -> None: + data[i] = da.values + times[i] = da.attrs["time"] + iters[i] = da.attrs["iter"] + for d in first.dims: + cv = np.asarray(da.coords[d].values) + bounds[d][i] = (cv[0], cv[-1]) if cv.size else (np.nan, np.nan) + + _record(0, first) for i, p in enumerate(dumps[1:], start=1): da = load_grid_h5(p) if da.shape != first.shape: @@ -306,23 +319,58 @@ def load_series(directory: str | Path) -> xr.DataArray: f"Shape mismatch in series: {p} has {da.shape}, " f"expected {first.shape}" ) - data[i] = da.values - times[i] = da.attrs["time"] - iters[i] = da.attrs["iter"] + _record(i, da) coords = {"t": times, "iter": ("t", iters)} coords.update({d: first.coords[d] for d in first.dims}) + # An axis whose per-dump bounds move (beyond fp noise) is autoscaled: keep + # the first dump's axis as the nominal dimension coordinate, but carry the + # true per-dump bounds along ``t`` so consumers can reconstruct the physical + # axis for any timestep (see :func:`physical_axis`) and so the bounds + # survive the NetCDF round-trip. + autoscaled: list[str] = [] + for d in first.dims: + b = bounds[d] + if not (np.allclose(b[:, 0], b[0, 0]) and np.allclose(b[:, 1], b[0, 1])): + autoscaled.append(str(d)) + coords[f"{d}_min"] = ("t", b[:, 0]) + coords[f"{d}_max"] = ("t", b[:, 1]) dims = ("t", *first.dims) attrs = dict(first.attrs) attrs.pop("time", None) attrs.pop("iter", None) attrs.pop("source", None) attrs["source_dir"] = str(directory) + if autoscaled: + attrs["autoscaled_dims"] = autoscaled return xr.DataArray( data, coords=coords, dims=dims, name=first.name, attrs=attrs ) +def physical_axis(da: xr.DataArray, dim: str, it: int = -1) -> np.ndarray: + """Physical coordinate values for ``dim``, honoring OSIRIS autoscale. + + When ``dim`` was autoscaled (its per-dump bounds move; see + :func:`load_series`), the series carries ``{dim}_min`` / ``{dim}_max`` along + ``t`` while the nominal dimension coordinate is only the first dump's axis. + This rebuilds ``linspace(min, max, n)`` for time index ``it`` (default the + last dump). Works on a still-stacked ``(t, …)`` series (the bounds are + vectors along ``t``) and on an already time-sliced array (the bounds are + scalar coordinates). Falls back to the dimension coordinate when no per-dump + bounds are present (non-autoscaled axes, or arrays not built by + :func:`load_series`). + """ + n = int(da.sizes[dim]) + lo_c, hi_c = f"{dim}_min", f"{dim}_max" + if lo_c in da.coords and hi_c in da.coords: + lo, hi = da.coords[lo_c], da.coords[hi_c] + if "t" in getattr(lo, "dims", ()): # still a (t, …) series + lo, hi = lo.isel(t=it), hi.isel(t=it) + return np.linspace(float(lo), float(hi), n) + return np.asarray(da.coords[dim].values) + + def list_diagnostics(run_dir: str | Path) -> dict[str, Path]: """Map every diagnostic name to its data source. @@ -379,6 +427,10 @@ def series_to_dataset(da: xr.DataArray) -> xr.Dataset: any remaining non-scalar attrs are coerced so ``to_netcdf`` succeeds. """ da = da.copy() + # The autoscaled-dim list is reconstructed on load from the {dim}_min/_max + # coordinates (which ride along as coords), so it need not survive as an + # attr — and dropping it avoids serializing a string array. + da.attrs.pop("autoscaled_dims", None) axis_units = da.attrs.pop("axis_units", {}) or {} axis_long = da.attrs.pop("axis_long_names", {}) or {} da.attrs = {k: _coerce_attr(v) for k, v in da.attrs.items() if v is not None} @@ -428,6 +480,15 @@ def load_series_nc(path: str | Path) -> xr.DataArray: da.attrs["axis_units"] = axis_units da.attrs["axis_long_names"] = axis_long da.attrs.setdefault("time_units", da.attrs.get("time_units", r"1/\omega_p")) + # Rebuild the autoscaled-dim list from the per-dump bound coords that were + # written by :func:`series_to_dataset`, so a round-tripped series is + # indistinguishable (for plotting) from a fresh :func:`load_series`. + autoscaled = [ + str(d) for d in da.dims + if f"{d}_min" in ds.coords and f"{d}_max" in ds.coords + ] + if autoscaled: + da.attrs["autoscaled_dims"] = autoscaled return da diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index 1a12166b..6ee87e41 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -168,7 +168,11 @@ def plot_phasespace( plot_arr = np.log10(np.abs(raw) + 1e-30) else: plot_arr = raw - xc, yc = da.coords[xdim].values, da.coords[ydim].values + # Reconstruct each axis on its OSIRIS-autoscale bounds (the momentum / gamma + # axis is re-picked per dump; ``physical_axis`` falls back to the stored + # coordinate for fixed axes like the cropped spatial one). + xc = _io.physical_axis(da, xdim) + yc = _io.physical_axis(da, ydim) mesh = ax.pcolormesh(xc, yc, plot_arr, shading="auto", cmap=cmap, vmin=vmin, vmax=vmax) plt.colorbar( mesh, ax=ax, label=rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da) @@ -208,28 +212,54 @@ def plot_phasespace_evolution( da = _crop_spatial_to_box(da) nt = da.coords["t"].size t_skip = max(1, nt // n_panels) - sl = da.isel(t=slice(0, None, t_skip)) - plot_da = np.log10(np.abs(sl) + 1e-30) if log else sl - # Convention: spatial axis horizontal, momentum axis vertical. + idx = list(range(0, nt, t_skip)) + sl = da.isel(t=idx) + # Convention: spatial axis horizontal, momentum axis vertical (fall back to + # dim order for a momentum-momentum space with no spatial axis). spatial = [d for d in da.dims if d != "t" and str(d).startswith("x")] moment = [d for d in da.dims if d != "t" and str(d).startswith("p")] - facet_kw: dict = {"cmap": cmap} - if spatial and moment: - facet_kw.update(x=spatial[0], y=moment[0]) - if log: - # Floor the shared colour scale at the lowest non-zero value so empty - # cells (log -> -30) don't crush the contrast across the facets. - vmin, vmax = _nonzero_log_clim(sl.values) - facet_kw.update(vmin=vmin, vmax=vmax) - g = plot_da.plot(col="t", col_wrap=min(col_wrap, sl.coords["t"].size), **facet_kw) - if spatial and moment: - g.set_xlabels(_axis_label(da, spatial[0])) - g.set_ylabels(_axis_label(da, moment[0])) + xdim, ydim = (spatial[0], moment[0]) if spatial and moment else (da.dims[2], da.dims[1]) + raw = sl.transpose("t", ydim, xdim).values + plot_arr = np.log10(np.abs(raw) + 1e-30) if log else raw + # Floor the shared colour scale at the lowest non-zero value so empty cells + # (log -> -30) don't crush the contrast across the facets. + vmin, vmax = _nonzero_log_clim(sl.values) if log else (None, None) + + npan = len(idx) + ncol = max(1, min(col_wrap, npan)) + nrow = int(np.ceil(npan / ncol)) + fig, axes = plt.subplots( + nrow, ncol, figsize=(3.2 * ncol, 3.0 * nrow), squeeze=False + ) + tvals = sl.coords["t"].values + mesh = None + for k in range(npan): + ax = axes[k // ncol][k % ncol] + # Each facet is drawn on ITS OWN axis: OSIRIS autoscale re-picks the + # momentum / gamma bounds every dump, so a single shared coordinate + # (the old faceting) would mislabel every panel but the first. + xc = _io.physical_axis(sl, xdim, it=k) + yc = _io.physical_axis(sl, ydim, it=k) + mesh = ax.pcolormesh( + xc, yc, plot_arr[k], shading="auto", cmap=cmap, vmin=vmin, vmax=vmax + ) + ax.set_title(rf"$t = {float(tvals[k]):.3g}$", fontsize=9) + if k % ncol == 0: + ax.set_ylabel(_axis_label(da, ydim)) + if k // ncol == nrow - 1: + ax.set_xlabel(_axis_label(da, xdim)) + for k in range(npan, nrow * ncol): # blank any unused grid cells + axes[k // ncol][k % ncol].axis("off") + if mesh is not None: + cbar = fig.colorbar(mesh, ax=axes, fraction=0.046, pad=0.02) + cbar.set_label( + rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da) + ) scale = r"$\log_{10}$ " if log else "" - g.fig.suptitle( + fig.suptitle( title or rf"{_display_name(da)} — {scale}phase space at sampled times", y=1.02 ) - return g.fig + return fig def field_energy_series(run_dir: str | Path) -> xr.DataArray: @@ -1346,3 +1376,4 @@ def _field_energy_from_series(ser: xr.DataArray) -> np.ndarray: display_name = _display_name tex = _tex sim_box_xmax = _sim_box_xmax +sim_box_bound = _sim_box_bound diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index 62ed5c09..157880dd 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -19,5 +19,4 @@ osiris: output: # diagnostics_to_log: [e1] # null/omitted = convert every diagnostic's full time history to netCDF # v_th: 0.1 # overlay the Langmuir (Bohm-Gross) branch on omega-k plots - # dist_cells: 10 # (osiris-lpi OsirisLPI only) right-boundary cells averaged for the f(p) lineouts # omega_k_zoom: 4.0 # (k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel; null = full diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index 449cd3ce..ed29b533 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -139,7 +139,7 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 > **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 longitudinal `e1` is 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.txt` artifacts and the `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac` metrics) and the spatio-temporal distribution-function lineouts (`distribution_lineouts/` + `deltaf_lineouts/`, `f(p)` and `δf` averaged over the rightmost `dist_cells` cells) are **not** produced by adept's general OSIRIS wrapper. They live in the [osiris-lpi](https://github.com/ergodicio/osiris-lpi) repo — `osiris_lpi.OsirisLPI` subclasses `BaseOsiris` and adds them in `post_process`, reading `output.dist_cells` and the drive `antenna.a0`/`antenna.omega0` from the deck. Regenerate them offline from saved NetCDFs with `python -m osiris_lpi.regen`. +> **SRS-specific plots & metrics live in osiris-lpi.** The laser-energy budget (reflected / transmitted / absorbed over time: the `energy_budget.png` + `laser_energy_budget.txt` artifacts and the `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac` metrics) 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 / `δf` views) are **not** produced by adept's general OSIRIS wrapper. They live in the [osiris-lpi](https://github.com/ergodicio/osiris-lpi) repo — `osiris_lpi.OsirisLPI` subclasses `BaseOsiris` and adds them in `post_process`, reading the drive `antenna.a0`/`antenna.omega0` from the deck. Regenerate them offline from saved NetCDFs with `python -m osiris_lpi.regen`. ## Programmatic use diff --git a/tests/test_osiris/test_post_netcdf.py b/tests/test_osiris/test_post_netcdf.py index 64d096c3..a7770e4d 100644 --- a/tests/test_osiris/test_post_netcdf.py +++ b/tests/test_osiris/test_post_netcdf.py @@ -393,3 +393,87 @@ def test_save_canned_plots_regenerates_from_netcdf(tmp_path: Path) -> None: ): assert key in written_nc, f"missing {key}" assert written_nc[key].exists() + + +# --- OSIRIS autoscale: per-dump momentum bounds (if_ps_p_auto) ------------ + + +def _write_phasespace_autoscaled(run_dir: Path, n_steps: int, nx: int, npmom: int) -> Path: + """Phase space whose momentum bounds GROW each dump (if_ps_p_auto=.true.). + + The bin count (``npmom``) stays fixed while the AXIS min/max move — exactly + what OSIRIS autoscale produces and what a shape-only series check misses. + """ + rng = np.random.default_rng(7) + d = run_dir / "MS" / "PHA" / "p1x1" / "electrons" + for k in range(n_steps): + it = k * 10 + p_hi = 1.0 + 0.5 * k # bounds move dump-to-dump + _write_dump( + d / f"p1x1-electrons-{it:06d}.h5", + "p1x1", + rng.standard_normal((npmom, nx)), + t=k * 0.5, + it=it, + axes=[ + ("x1", "x_1", r"c / \omega_p", 0.0, 10.0), + ("p1", "p_1", r"m_e c", -p_hi, p_hi), + ], + ) + return d + + +def test_load_series_captures_autoscale_bounds(tmp_path: Path) -> None: + d = _write_phasespace_autoscaled(tmp_path / "run", n_steps=4, nx=8, npmom=6) + ser = oio.load_series(d) + + # p1 is autoscaled (bounds move); x1 is fixed (no bound coords). + assert ser.attrs.get("autoscaled_dims") == ["p1"] + assert "p1_min" in ser.coords and "p1_max" in ser.coords + assert "x1_min" not in ser.coords and "x1_max" not in ser.coords + np.testing.assert_allclose(ser["p1_max"].values, [1.0, 1.5, 2.0, 2.5]) + np.testing.assert_allclose(ser["p1_min"].values, [-1.0, -1.5, -2.0, -2.5]) + + # physical_axis rebuilds each dump's true axis; the nominal dim coordinate + # (first dump) is NOT reused for later steps. + assert oio.physical_axis(ser, "p1", it=0)[-1] == 1.0 + assert oio.physical_axis(ser, "p1", it=-1)[-1] == 2.5 + # a fixed axis falls back to the stored coordinate. + np.testing.assert_allclose(oio.physical_axis(ser, "x1"), ser["x1"].values) + + +def test_autoscale_bounds_survive_netcdf_roundtrip(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + _write_phasespace_autoscaled(run_dir, n_steps=4, nx=8, npmom=6) + out = tmp_path / "binary" + oio.save_run_datasets(run_dir, out) + + from_nc = oio.load_series_nc(out / "PHA" / "p1x1" / "electrons.nc") + assert from_nc.attrs.get("autoscaled_dims") == ["p1"] + np.testing.assert_allclose( + oio.physical_axis(from_nc, "p1", it=-1), + np.linspace(-2.5, 2.5, from_nc.sizes["p1"]), + ) + + +def test_phasespace_plots_run_on_autoscaled_series(tmp_path: Path) -> None: + import matplotlib.pyplot as plt + + from adept.osiris import plots as oplots + + d = _write_phasespace_autoscaled(tmp_path / "run", n_steps=6, nx=8, npmom=6) + ser = oio.load_series(d) + + # final-step heatmap is drawn on the LAST dump's autoscaled momentum range + # (p_hi = 1 + 0.5*5 = 3.5), not the first dump's (1.0). + final = ser.isel(t=-1) + final.attrs["time"] = float(final["t"].values) + ax = oplots.plot_phasespace(final) + ylo, yhi = ax.get_ylim() + assert yhi >= 3.0 and ylo <= -3.0 + plt.close(ax.figure) + + # faceted evolution renders one panel per sampled time without error. + fig = oplots.plot_phasespace_evolution(ser, n_panels=4) + assert len(fig.axes) >= 4 # panels (+ a shared colorbar) + plt.close(fig) From 2271d2f413a6fd606fe471d7e2c961d84898d63e Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 16 Jun 2026 16:59:57 -0700 Subject: [PATCH 19/49] Ran pre-commit --- adept/osiris/__init__.py | 1 + adept/osiris/base.py | 12 +- adept/osiris/deck.py | 37 ++--- adept/osiris/io.py | 45 ++---- adept/osiris/plots.py | 159 ++++++++------------ adept/osiris/post.py | 4 +- adept/osiris/regen.py | 27 ++-- adept/osiris/runner.py | 10 +- tests/test_osiris/decks/F-Tsung_2d_lpi_deck | 2 +- tests/test_osiris/decks/two-stream-1d | 4 +- tests/test_osiris/test_deck_roundtrip.py | 5 +- tests/test_osiris/test_diagnostics_plots.py | 31 ++-- tests/test_osiris/test_io_and_plots.py | 1 - tests/test_osiris/test_plots_new_views.py | 37 +++-- tests/test_osiris/test_post_netcdf.py | 12 +- tests/test_osiris/test_runner.py | 1 - 16 files changed, 160 insertions(+), 228 deletions(-) diff --git a/adept/osiris/__init__.py b/adept/osiris/__init__.py index b17268f8..7815b551 100644 --- a/adept/osiris/__init__.py +++ b/adept/osiris/__init__.py @@ -8,5 +8,6 @@ def __getattr__(name): if name == "BaseOsiris": from adept.osiris.base import BaseOsiris + return BaseOsiris raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 172299b1..a5d937a7 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -23,9 +23,7 @@ def __init__(self, cfg: dict) -> None: osiris_cfg = cfg.get("osiris", {}) deck_path = osiris_cfg.get("deck") if not deck_path: - raise ValueError( - "BaseOsiris: cfg['osiris']['deck'] is required" - ) + raise ValueError("BaseOsiris: cfg['osiris']['deck'] is required") self._sections = _deck.parse_deck_file(deck_path) overrides = osiris_cfg.pop("overrides", None) or {} @@ -57,9 +55,7 @@ def get_derived_quantities(self) -> dict: tmax = time.get("tmax") if nx and xmin is not None and xmax is not None: - derived["dx"] = [ - (xmax[d] - xmin[d]) / nx[d] for d in range(len(nx)) - ] + derived["dx"] = [(xmax[d] - xmin[d]) / nx[d] for d in range(len(nx))] if dt is not None and nx: derived["cfl_ratio"] = float(dt) / min(derived["dx"]) if dt is not None and tmax is not None: @@ -106,9 +102,7 @@ def post_process(self, run_output: dict, td: str) -> dict: return _post.collect(run_output, self.cfg, td) def vg(self, trainable_modules: dict, args: dict): - raise NotImplementedError( - "OSIRIS is not differentiable inside adept" - ) + raise NotImplementedError("OSIRIS is not differentiable inside adept") # --- helpers ---------------------------------------------------------- diff --git a/adept/osiris/deck.py b/adept/osiris/deck.py index 00c10f35..36a14569 100644 --- a/adept/osiris/deck.py +++ b/adept/osiris/deck.py @@ -126,17 +126,12 @@ def skip_ws(j: int) -> int: break m = _IDENT_RE.match(src, i) if not m: - raise ValueError( - f"Expected section name at offset {i}: {src[i:i+30]!r}" - ) + raise ValueError(f"Expected section name at offset {i}: {src[i : i + 30]!r}") name = m.group(0) i = m.end() i = skip_ws(i) if i >= n or src[i] != "{": - raise ValueError( - f"Expected '{{' after section {name!r} at offset {i}: " - f"{src[i:i+30]!r}" - ) + raise ValueError(f"Expected '{{' after section {name!r} at offset {i}: {src[i : i + 30]!r}") i += 1 params: dict[str, Any] = {} while True: @@ -148,17 +143,12 @@ def skip_ws(j: int) -> int: break km = _KEY_RE.match(src, i) if not km: - raise ValueError( - f"Expected key in section {name!r} at offset {i}: " - f"{src[i:i+30]!r}" - ) + raise ValueError(f"Expected key in section {name!r} at offset {i}: {src[i : i + 30]!r}") key = km.group(0) i = km.end() i = skip_ws(i) if i >= n or src[i] != "=": - raise ValueError( - f"Expected '=' after key {key!r} at offset {i}" - ) + raise ValueError(f"Expected '=' after key {key!r} at offset {i}") i += 1 value_start = i in_string = False @@ -249,9 +239,7 @@ def _find_param_key(params: dict[str, Any], requested: str) -> str: return candidates[0] if len(candidates) == 0: return requested - raise ValueError( - f"Ambiguous override key {requested!r}; candidates: {candidates}" - ) + raise ValueError(f"Ambiguous override key {requested!r}; candidates: {candidates}") def _merge_params(params: dict[str, Any], over: dict[str, Any]) -> None: @@ -266,8 +254,8 @@ def merge_overrides(sections: Sections, overrides: dict[str, Any]) -> None: ``overrides`` shape:: { - "grid": {"nx_p": [256]}, # apply to all occurrences - "species": {0: {"num_par_x": [512]}}, # indexed for repeated sections + "grid": {"nx_p": [256]}, # apply to all occurrences + "species": {0: {"num_par_x": [512]}}, # indexed for repeated sections } """ if not overrides: @@ -278,18 +266,13 @@ def merge_overrides(sections: Sections, overrides: dict[str, Any]) -> None: for sec_name, sec_over in overrides.items(): if sec_name not in by_name: - raise KeyError( - f"Override references unknown section: {sec_name!r}" - ) + raise KeyError(f"Override references unknown section: {sec_name!r}") occurrences = by_name[sec_name] - if isinstance(sec_over, dict) and sec_over and all( - isinstance(k, int) for k in sec_over.keys() - ): + if isinstance(sec_over, dict) and sec_over and all(isinstance(k, int) for k in sec_over.keys()): for idx, params_over in sec_over.items(): if idx < 0 or idx >= len(occurrences): raise IndexError( - f"Override section {sec_name!r}[{idx}] out of range; " - f"{len(occurrences)} occurrence(s) present" + f"Override section {sec_name!r}[{idx}] out of range; {len(occurrences)} occurrence(s) present" ) _merge_params(sections[occurrences[idx]][1], params_over) else: diff --git a/adept/osiris/io.py b/adept/osiris/io.py index facf8143..feb28627 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -87,9 +87,7 @@ def load_grid_h5(path: str | Path) -> xr.DataArray: # Identify the data dataset (everything that's not AXIS / SIMULATION). data_keys = [k for k in f.keys() if k not in ("AXIS", "SIMULATION")] if len(data_keys) != 1: - raise ValueError( - f"Expected exactly one data dataset in {path}; got {data_keys}" - ) + raise ValueError(f"Expected exactly one data dataset in {path}; got {data_keys}") name = data_keys[0] arr = f[name][...].astype(_DIAG_DTYPE) axes_osiris = _axis_metadata(f) @@ -160,11 +158,7 @@ def load_raw_h5(path: str | Path) -> xr.Dataset: """ path = Path(path) with h5py.File(path, "r") as f: - data_keys = sorted( - k - for k in f.keys() - if k not in ("AXIS", "SIMULATION") and isinstance(f[k], h5py.Dataset) - ) + data_keys = sorted(k for k in f.keys() if k not in ("AXIS", "SIMULATION") and isinstance(f[k], h5py.Dataset)) data_vars: dict[str, tuple] = {} for name in data_keys: dset = f[name] @@ -180,9 +174,7 @@ def load_raw_h5(path: str | Path) -> xr.Dataset: attrs = { "time": float(f.attrs["TIME"][0]) if "TIME" in f.attrs else float("nan"), - "iter": int(f.attrs["ITER"][0]) - if "ITER" in f.attrs - else _iter_from_name(path), + "iter": int(f.attrs["ITER"][0]) if "ITER" in f.attrs else _iter_from_name(path), "long_name": _decode(f.attrs.get("LABEL", path.stem)), "time_units": _decode(f.attrs.get("TIME UNITS", "")), "source": str(path), @@ -315,10 +307,7 @@ def _record(i: int, da: xr.DataArray) -> None: for i, p in enumerate(dumps[1:], start=1): da = load_grid_h5(p) if da.shape != first.shape: - raise ValueError( - f"Shape mismatch in series: {p} has {da.shape}, " - f"expected {first.shape}" - ) + raise ValueError(f"Shape mismatch in series: {p} has {da.shape}, expected {first.shape}") _record(i, da) coords = {"t": times, "iter": ("t", iters)} @@ -343,9 +332,7 @@ def _record(i: int, da: xr.DataArray) -> None: attrs["source_dir"] = str(directory) if autoscaled: attrs["autoscaled_dims"] = autoscaled - return xr.DataArray( - data, coords=coords, dims=dims, name=first.name, attrs=attrs - ) + return xr.DataArray(data, coords=coords, dims=dims, name=first.name, attrs=attrs) def physical_axis(da: xr.DataArray, dim: str, it: int = -1) -> np.ndarray: @@ -483,10 +470,7 @@ def load_series_nc(path: str | Path) -> xr.DataArray: # Rebuild the autoscaled-dim list from the per-dump bound coords that were # written by :func:`series_to_dataset`, so a round-tripped series is # indistinguishable (for plotting) from a fresh :func:`load_series`. - autoscaled = [ - str(d) for d in da.dims - if f"{d}_min" in ds.coords and f"{d}_max" in ds.coords - ] + autoscaled = [str(d) for d in da.dims if f"{d}_min" in ds.coords and f"{d}_max" in ds.coords] if autoscaled: da.attrs["autoscaled_dims"] = autoscaled return da @@ -515,10 +499,7 @@ def _compression_encoding(ds: xr.Dataset) -> dict: helps the float32 byte pattern. RAW (per-particle) data is noise-like and barely compresses, but the setting does no harm. """ - return { - name: {"zlib": True, "complevel": 4, "shuffle": True} - for name in ds.data_vars - } + return {name: {"zlib": True, "complevel": 4, "shuffle": True} for name in ds.data_vars} def save_run_datasets( @@ -543,16 +524,12 @@ def save_run_datasets( diags = list_diagnostics(run_dir) written: list[Path] = [] for relpath in sorted(diags): - if diagnostics is not None and ( - relpath not in diagnostics and Path(relpath).name not in diagnostics - ): + if diagnostics is not None and (relpath not in diagnostics and Path(relpath).name not in diagnostics): continue try: if _diag_is_raw(relpath, diags[relpath]): # RAW (particle) dumps: per-particle datasets, no grid/AXIS. - ds: xr.Dataset = load_raw_series( - diags[relpath], drop_initial=raw_drop_initial - ) + ds: xr.Dataset = load_raw_series(diags[relpath], drop_initial=raw_drop_initial) else: ds = series_to_dataset(load_series(diags[relpath])) dest = out_dir / f"{relpath}.nc" @@ -571,9 +548,7 @@ def save_run_datasets( if energy is not None: dest = out_dir / "HIST" / "energy.nc" dest.parent.mkdir(parents=True, exist_ok=True) - energy.to_netcdf( - dest, engine="h5netcdf", encoding=_compression_encoding(energy) - ) + energy.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(energy)) written.append(dest) except Exception as e: print(f"[post] skipping HIST energy: {e}") diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index 6ee87e41..af819b07 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -71,9 +71,7 @@ def plot_spacetime( """ da = _ensure_series(series) if da.ndim != 2: - raise ValueError( - f"plot_spacetime expects a 2D (t, x) array; got dims {da.dims}" - ) + raise ValueError(f"plot_spacetime expects a 2D (t, x) array; got dims {da.dims}") if ax is None: _, ax = plt.subplots(figsize=(6, 4)) data = np.log10(np.abs(da.values) + 1e-30) if log else da.values # (t, x) @@ -116,9 +114,7 @@ def plot_lineouts( """ da = _decorate(_ensure_series(series)) if da.ndim != 2: - raise ValueError( - f"plot_lineouts expects a 2D (t, x) array; got dims {da.dims}" - ) + raise ValueError(f"plot_lineouts expects a 2D (t, x) array; got dims {da.dims}") nt = da.coords["t"].size t_skip = max(1, nt // n_panels) sl = da.isel(t=slice(0, None, t_skip)) @@ -126,9 +122,7 @@ def plot_lineouts( g = sl.plot(x=xname, col="t", col_wrap=min(col_wrap, sl.coords["t"].size)) g.set_xlabels(_axis_label(da, xname)) g.set_ylabels(_value_label(da)) - g.fig.suptitle( - title or rf"{_display_name(da)} — lineouts vs $x$ at sampled times", y=1.02 - ) + g.fig.suptitle(title or rf"{_display_name(da)} — lineouts vs $x$ at sampled times", y=1.02) return g.fig @@ -150,9 +144,7 @@ def plot_phasespace( if not isinstance(da, xr.DataArray): da = _io.load_phasespace_h5(da) if da.ndim != 2: - raise ValueError( - f"plot_phasespace expects 2D data; got {da.dims} {da.shape}" - ) + raise ValueError(f"plot_phasespace expects 2D data; got {da.dims} {da.shape}") da = _crop_spatial_to_box(da) if ax is None: _, ax = plt.subplots(figsize=(5, 4)) @@ -174,17 +166,12 @@ def plot_phasespace( xc = _io.physical_axis(da, xdim) yc = _io.physical_axis(da, ydim) mesh = ax.pcolormesh(xc, yc, plot_arr, shading="auto", cmap=cmap, vmin=vmin, vmax=vmax) - plt.colorbar( - mesh, ax=ax, label=rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da) - ) + plt.colorbar(mesh, ax=ax, label=rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da)) ax.set_xlabel(_axis_label(da, xdim)) ax.set_ylabel(_axis_label(da, ydim)) t = da.attrs.get("time", float("nan")) scale = r"$\log_{10}$ " if log else "" - ax.set_title( - title - or rf"{_display_name(da)} — {scale}phase space ($t = {t:.3g}\ 1/\omega_p$)" - ) + ax.set_title(title or rf"{_display_name(da)} — {scale}phase space ($t = {t:.3g}\ 1/\omega_p$)") return ax @@ -206,9 +193,7 @@ def plot_phasespace_evolution( da = series if isinstance(series, xr.DataArray) else _io.load_series(series) da = _decorate(da) if da.ndim != 3: - raise ValueError( - f"plot_phasespace_evolution expects (t, p, x); got dims {da.dims}" - ) + raise ValueError(f"plot_phasespace_evolution expects (t, p, x); got dims {da.dims}") da = _crop_spatial_to_box(da) nt = da.coords["t"].size t_skip = max(1, nt // n_panels) @@ -228,9 +213,7 @@ def plot_phasespace_evolution( npan = len(idx) ncol = max(1, min(col_wrap, npan)) nrow = int(np.ceil(npan / ncol)) - fig, axes = plt.subplots( - nrow, ncol, figsize=(3.2 * ncol, 3.0 * nrow), squeeze=False - ) + fig, axes = plt.subplots(nrow, ncol, figsize=(3.2 * ncol, 3.0 * nrow), squeeze=False) tvals = sl.coords["t"].values mesh = None for k in range(npan): @@ -240,9 +223,7 @@ def plot_phasespace_evolution( # (the old faceting) would mislabel every panel but the first. xc = _io.physical_axis(sl, xdim, it=k) yc = _io.physical_axis(sl, ydim, it=k) - mesh = ax.pcolormesh( - xc, yc, plot_arr[k], shading="auto", cmap=cmap, vmin=vmin, vmax=vmax - ) + mesh = ax.pcolormesh(xc, yc, plot_arr[k], shading="auto", cmap=cmap, vmin=vmin, vmax=vmax) ax.set_title(rf"$t = {float(tvals[k]):.3g}$", fontsize=9) if k % ncol == 0: ax.set_ylabel(_axis_label(da, ydim)) @@ -252,13 +233,9 @@ def plot_phasespace_evolution( axes[k // ncol][k % ncol].axis("off") if mesh is not None: cbar = fig.colorbar(mesh, ax=axes, fraction=0.046, pad=0.02) - cbar.set_label( - rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da) - ) + cbar.set_label(rf"$\log_{{10}}$ {_value_label(da)}" if log else _value_label(da)) scale = r"$\log_{10}$ " if log else "" - fig.suptitle( - title or rf"{_display_name(da)} — {scale}phase space at sampled times", y=1.02 - ) + fig.suptitle(title or rf"{_display_name(da)} — {scale}phase space at sampled times", y=1.02) return fig @@ -336,11 +313,7 @@ def field_energy_components(run_dir: str | Path) -> xr.Dataset: continue found = True e_t = _field_energy_from_series(ser) - its = ( - np.asarray(ser.coords["iter"].values) - if "iter" in ser.coords - else np.arange(ser.sizes["t"]) - ) + its = np.asarray(ser.coords["iter"].values) if "iter" in ser.coords else np.arange(ser.sizes["t"]) ts = np.asarray(ser.coords["t"].values, dtype="float64") for k in range(ts.size): it = int(its[k]) @@ -463,9 +436,7 @@ def plot_omega_k( """ da = _ensure_series(series) if da.ndim != 2: - raise ValueError( - f"plot_omega_k expects (t, x); got dims {da.dims}" - ) + raise ValueError(f"plot_omega_k expects (t, x); got dims {da.dims}") if ax is None: _, ax = plt.subplots(figsize=(6, 5)) @@ -487,9 +458,7 @@ def plot_omega_k( k = np.fft.fftshift(np.fft.fftfreq(nx, d=dx)) * 2 * np.pi mesh = ax.pcolormesh(k, omega, P, shading="auto", cmap=cmap) - plt.colorbar( - mesh, ax=ax, label=(r"$\log_{10}\,|\tilde{F}|^2$" if log else r"$|\tilde{F}|^2$") - ) + plt.colorbar(mesh, ax=ax, label=(r"$\log_{10}\,|\tilde{F}|^2$" if log else r"$|\tilde{F}|^2$")) if k_max is None: k_max = float(np.max(np.abs(k))) @@ -501,20 +470,17 @@ def plot_omega_k( # Overlay analytical dispersion lines (sampled across the visible k range). k_line = np.linspace(-k_max, k_max, 401) if show_light_line: - ax.plot(k_line, +k_line, "w-", lw=0.8, alpha=0.6, - label=r"light line: $\omega = \pm k$") + ax.plot(k_line, +k_line, "w-", lw=0.8, alpha=0.6, label=r"light line: $\omega = \pm k$") ax.plot(k_line, -k_line, "w-", lw=0.8, alpha=0.6) if show_em: - w_em = np.sqrt(omega_p ** 2 + k_line ** 2) - ax.plot(k_line, +w_em, "w--", lw=1, alpha=0.7, - label=r"EM: $\omega^2 = \omega_p^2 + k^2$") + w_em = np.sqrt(omega_p**2 + k_line**2) + ax.plot(k_line, +w_em, "w--", lw=1, alpha=0.7, label=r"EM: $\omega^2 = \omega_p^2 + k^2$") ax.plot(k_line, -w_em, "w--", lw=1, alpha=0.7) if show_langmuir: if v_th is None: raise ValueError("show_langmuir=True requires v_th=...") - w_l = np.sqrt(omega_p ** 2 + 3 * (k_line * v_th) ** 2) - ax.plot(k_line, +w_l, "c:", lw=1, alpha=0.8, - label=r"Langmuir: $\omega^2 = \omega_p^2 + 3 k^2 v_{th}^2$") + w_l = np.sqrt(omega_p**2 + 3 * (k_line * v_th) ** 2) + ax.plot(k_line, +w_l, "c:", lw=1, alpha=0.8, label=r"Langmuir: $\omega^2 = \omega_p^2 + 3 k^2 v_{th}^2$") ax.plot(k_line, -w_l, "c:", lw=1, alpha=0.8) if show_em or show_langmuir or show_light_line: ax.legend(loc="upper right", fontsize=8, framealpha=0.6) @@ -554,14 +520,25 @@ def plot_omega_k_figure( da = _ensure_series(series) fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(6, 10)) plot_omega_k( - da, ax=ax_top, log=log, cmap=cmap, - show_langmuir=v_th is not None, v_th=v_th, + da, + ax=ax_top, + log=log, + cmap=cmap, + show_langmuir=v_th is not None, + v_th=v_th, ) z = _omega_k_zoom_window(da, omega_k_zoom) plot_omega_k( - da, ax=ax_bot, log=log, cmap=cmap, - show_light_line=True, show_langmuir=v_th is not None, v_th=v_th, - k_max=z, omega_max=z, equal_aspect=True, + da, + ax=ax_bot, + log=log, + cmap=cmap, + show_light_line=True, + show_langmuir=v_th is not None, + v_th=v_th, + k_max=z, + omega_max=z, + equal_aspect=True, title=rf"{_display_name(da)} — $(k, \omega)$ (equal aspect)", ) fig.tight_layout() @@ -692,8 +669,12 @@ def plot_currents_lineouts(run_dir: str | Path, *, n_avg_frac: float = 0.2) -> p w = max(1, round(n_avg_frac * nt)) (line,) = ax.plot(x, da.isel(t=-1).values, lw=1.4, label=_tex(_long_name(da))) ax.plot( - x, da.isel(t=slice(nt - w, nt)).mean("t").values, - lw=1.0, ls="--", alpha=0.7, color=line.get_color(), + x, + da.isel(t=slice(nt - w, nt)).mean("t").values, + lw=1.0, + ls="--", + alpha=0.7, + color=line.get_color(), ) any_ser = _decorate(next(iter(comps.values()))) xdim = next(d for d in any_ser.dims if d != "t") @@ -782,7 +763,7 @@ def _temperature_series(entries: list[tuple[str, str, Path]]) -> xr.DataArray | comps = [c for c in comps if c.ndim == 2] if not comps: return None - total = sum((c ** 2 for c in comps[1:]), comps[0] ** 2) + total = sum((c**2 for c in comps[1:]), comps[0] ** 2) long_name = r"T = \sum_i u_{th,i}^2" units = r"m_e c^2" elif tens: @@ -883,8 +864,7 @@ def plot_profile( initial = np.abs(initial) # Dotted and on top (high zorder) so the initial profile stays visible # where the final / mean curves overlap it. - ax.plot(x, initial, lw=1.6, ls=":", color="k", zorder=3, - label=f"initial ($t={float(tvals[0]):.3g}$)") + ax.plot(x, initial, lw=1.6, ls=":", color="k", zorder=3, label=f"initial ($t={float(tvals[0]):.3g}$)") ax.plot(x, final, lw=1.4, zorder=2, label=f"final ($t={float(tvals[-1]):.3g}$)") ax.plot(x, mean, lw=1.1, ls="--", zorder=2.2, label=f"mean of last {w} dumps") ax.set_xlabel(_axis_label(da, xdim)) @@ -966,12 +946,12 @@ def plot_field_lr_decomposition(run_dir: str | Path) -> dict[str, plt.Figure]: for ax, side in zip(axes, ("right", "left"), strict=False): da = parts[side] plot_spacetime( - da, ax=ax, space_on_x=True, + da, + ax=ax, + space_on_x=True, title=f"{_display_name(da)} — spacetime", ) - fig.suptitle( - rf"{comp}: left/right-going decomposition (vacuum Riemann split)", y=1.02 - ) + fig.suptitle(rf"{comp}: left/right-going decomposition (vacuum Riemann split)", y=1.02) fig.tight_layout() figs[comp] = fig return figs @@ -1079,9 +1059,7 @@ def _write(fig: plt.Figure, rel: str) -> Path: plot_spacetime(ser, ax=ax, log=True) written[f"spacetime_log/{comp}"] = _write(fig, f"spacetime_log/{comp}.png") - written[f"lineouts/{comp}"] = _write( - plot_lineouts(ser, n_panels=n_panels), f"lineouts/{comp}.png" - ) + written[f"lineouts/{comp}"] = _write(plot_lineouts(ser, n_panels=n_panels), f"lineouts/{comp}.png") # Full (k, ω) spectrum on top, equal-aspect square window below. written[f"omega_k/{comp}"] = _write( @@ -1118,15 +1096,11 @@ def _write(fig: plt.Figure, rel: str) -> Path: fig, ax = plt.subplots(figsize=(6, 4)) plot_spacetime(ser, ax=ax) - written[f"moments/{species}/{quantity}"] = _write( - fig, f"moments/{species}/{quantity}.png" - ) + written[f"moments/{species}/{quantity}"] = _write(fig, f"moments/{species}/{quantity}.png") fig, ax = plt.subplots(figsize=(6, 4)) plot_spacetime(ser, ax=ax, log=True) - written[f"moments/{species}/{quantity}_log"] = _write( - fig, f"moments/{species}/{quantity}_log.png" - ) + written[f"moments/{species}/{quantity}_log"] = _write(fig, f"moments/{species}/{quantity}_log.png") written[f"moments/{species}/lineouts/{quantity}"] = _write( plot_lineouts(ser, n_panels=n_panels), @@ -1141,12 +1115,12 @@ def _write(fig: plt.Figure, rel: str) -> Path: if dens is not None: fig, ax = plt.subplots(figsize=(6, 4)) plot_profile( - dens, ax=ax, show_initial=True, + dens, + ax=ax, + show_initial=True, title=f"{label} — density profile", ) - written[f"profiles/{species}/density"] = _write( - fig, f"profiles/{species}/density.png" - ) + written[f"profiles/{species}/density"] = _write(fig, f"profiles/{species}/density.png") except Exception as e: print(f"[plots] skipping density profile for {species}: {e}") try: @@ -1159,12 +1133,12 @@ def _write(fig: plt.Figure, rel: str) -> Path: if temp is not None: fig, ax = plt.subplots(figsize=(6, 4)) plot_profile( - temp, ax=ax, show_initial=True, + temp, + ax=ax, + show_initial=True, title=f"{label} — temperature profile", ) - written[f"profiles/{species}/temperature"] = _write( - fig, f"profiles/{species}/temperature.png" - ) + written[f"profiles/{species}/temperature"] = _write(fig, f"profiles/{species}/temperature.png") except Exception as e: print(f"[plots] skipping temperature profile for {species}: {e}") @@ -1189,9 +1163,7 @@ def _write(fig: plt.Figure, rel: str) -> Path: if final.ndim == 2: fig, ax = plt.subplots(figsize=(5, 4)) plot_phasespace(final, ax=ax) - written[f"phasespace/{species}/{ps_name}"] = _write( - fig, f"phasespace/{species}/{ps_name}.png" - ) + written[f"phasespace/{species}/{ps_name}"] = _write(fig, f"phasespace/{species}/{ps_name}.png") try: if ser.ndim == 3: written[f"phasespace_evolution/{species}/{ps_name}"] = _write( @@ -1219,9 +1191,7 @@ def _write(fig: plt.Figure, rel: str) -> Path: try: fig, ax = plt.subplots(figsize=(6, 4)) plot_energy_components(run_dir, ax=ax) - written["energy_components_vs_time"] = _write( - fig, "energy_components_vs_time.png" - ) + written["energy_components_vs_time"] = _write(fig, "energy_components_vs_time.png") except (FileNotFoundError, RuntimeError) as e: print(f"[plots] skipping energy_components_vs_time: {e}") @@ -1246,9 +1216,7 @@ def _ensure_series(src) -> xr.DataArray: p = Path(src) if p.is_dir(): return _io.load_series(p) - raise TypeError( - f"Expected an xr.DataArray or a directory path; got {type(src).__name__}" - ) + raise TypeError(f"Expected an xr.DataArray or a directory path; got {type(src).__name__}") def _tex(s) -> str: @@ -1284,10 +1252,7 @@ def _label_tex(s) -> str: s = str(s) if not s or (s.startswith("$") and s.endswith("$")): return s - return " ".join( - f"${tok}$" if tok and any(c in tok for c in "\\_^{}") else tok - for tok in s.split(" ") - ) + return " ".join(f"${tok}$" if tok and any(c in tok for c in "\\_^{}") else tok for tok in s.split(" ")) def _axis_label(da: xr.DataArray, dim: str) -> str: diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 8f2d0ee9..23269336 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -128,7 +128,9 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: # Convert each diagnostic's time history to an xarray netCDF. if ms.is_dir(): _io.save_run_datasets( - run_dir, td / "binary", diagnostics=whitelist, + run_dir, + td / "binary", + diagnostics=whitelist, raw_drop_initial=raw_drop_initial, ) diff --git a/adept/osiris/regen.py b/adept/osiris/regen.py index b3dad320..4ec88be5 100644 --- a/adept/osiris/regen.py +++ b/adept/osiris/regen.py @@ -136,22 +136,27 @@ def main(argv: list[str] | None = None) -> int: ) ap.add_argument("src", help="run dir (containing binary/) or a binary/ NetCDF dir") ap.add_argument("-o", "--out", default=None, help="output dir (default /plots_regen)") - ap.add_argument("--no-config", action="store_true", - help="ignore the run's config.yaml output block") - ap.add_argument("--v-th", type=float, default=None, - help="electron thermal velocity for the Langmuir overlay on omega-k plots") - ap.add_argument("--omega-k-zoom", type=float, default=None, - help="(k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel") - ap.add_argument("--no-zoom", action="store_true", - help="use the full Nyquist window for the lower omega-k panel (omega_k_zoom=None)") + ap.add_argument("--no-config", action="store_true", help="ignore the run's config.yaml output block") + ap.add_argument( + "--v-th", type=float, default=None, help="electron thermal velocity for the Langmuir overlay on omega-k plots" + ) + ap.add_argument( + "--omega-k-zoom", + type=float, + default=None, + help="(k, omega) half-width [omega_p] for the equal-aspect lower omega-k panel", + ) + ap.add_argument( + "--no-zoom", + action="store_true", + help="use the full Nyquist window for the lower omega-k panel (omega_k_zoom=None)", + ) ap.add_argument("--dpi", type=int, default=None, help="figure DPI") ap.add_argument("--n-panels", type=int, default=None, help="panels for faceted plots") args = ap.parse_args(argv) out_dir = Path(args.out) if args.out else default_out_dir(args.src) - written = regenerate( - args.src, out_dir=out_dir, use_config=not args.no_config, **_cli_overrides(args) - ) + written = regenerate(args.src, out_dir=out_dir, use_config=not args.no_config, **_cli_overrides(args)) _summarize(written, out_dir) return 0 diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index 410bfb9a..11ce118e 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -22,7 +22,6 @@ from pathlib import Path from typing import Any - _OSIRIS_ERR_TOKENS = ("error", "aborting", "(*error*)") # stderr noise emitted by X11 / mpirun that we should NOT treat as an error. _OSIRIS_STDERR_NOISE = ( @@ -40,9 +39,7 @@ def _looks_like_osiris_error(line: str) -> bool: return any(tok in low for tok in _OSIRIS_ERR_TOKENS) -def _stream_to_file_and_buffer( - stream, file_path: Path, tail: list[str], tail_max: int = 200 -) -> None: +def _stream_to_file_and_buffer(stream, file_path: Path, tail: list[str], tail_max: int = 200) -> None: """Tee a subprocess stream to disk and a bounded in-memory tail.""" with file_path.open("w") as fh: for raw in iter(stream.readline, b""): @@ -134,10 +131,7 @@ def run_osiris( if rc != 0: tail = "".join(stderr_tail[-50:]) or "(empty stderr)" raise RuntimeError( - f"OSIRIS exited with status {rc}.\n" - f" cmd: {shlex.join(cmd)}\n" - f" cwd: {run_dir}\n" - f" stderr tail:\n{tail}" + f"OSIRIS exited with status {rc}.\n cmd: {shlex.join(cmd)}\n cwd: {run_dir}\n stderr tail:\n{tail}" ) # OSIRIS can exit 0 even on input-file errors: it prints something diff --git a/tests/test_osiris/decks/F-Tsung_2d_lpi_deck b/tests/test_osiris/decks/F-Tsung_2d_lpi_deck index 0790052d..5d528722 100644 --- a/tests/test_osiris/decks/F-Tsung_2d_lpi_deck +++ b/tests/test_osiris/decks/F-Tsung_2d_lpi_deck @@ -136,4 +136,4 @@ antenna { a0=0.0033,t_rise=5,t_flat=50000.67,t_fall=5,omega0=1.0,x0=320,rad_x=500000.0, side=2, -} \ No newline at end of file +} diff --git a/tests/test_osiris/decks/two-stream-1d b/tests/test_osiris/decks/two-stream-1d index afeba0fd..dae72249 100644 --- a/tests/test_osiris/decks/two-stream-1d +++ b/tests/test_osiris/decks/two-stream-1d @@ -101,7 +101,7 @@ species udist { uth(1:3) = 0.00447, 0.00447, 0.00447, ! sqrt(T0/me c^2) - ufl(1:3) = 0.025, 0.0, 0.0, ! 0.025 ≈ 5.6x uth (clear two-stream separation) + ufl(1:3) = 0.025, 0.0, 0.0, ! 0.025 ≈ 5.6x uth (clear two-stream separation) } profile @@ -146,7 +146,7 @@ species udist { uth(1:3) = 0.00447, 0.00447, 0.00447, ! sqrt(T0/me c^2) - ufl(1:3) = -0.025, 0.0, 0.0, ! -0.025 ≈ 5.6x uth (clear two-stream separation) + ufl(1:3) = -0.025, 0.0, 0.0, ! -0.025 ≈ 5.6x uth (clear two-stream separation) } profile diff --git a/tests/test_osiris/test_deck_roundtrip.py b/tests/test_osiris/test_deck_roundtrip.py index 7991c055..da16803b 100644 --- a/tests/test_osiris/test_deck_roundtrip.py +++ b/tests/test_osiris/test_deck_roundtrip.py @@ -8,7 +8,6 @@ from adept.osiris import deck as osd - DECKS_DIR = Path(__file__).parent / "decks" # Real OSIRIS decks vendored into the repo so the round-trip is exercised in @@ -130,9 +129,7 @@ def test_deck_to_flat_dict_expands_lists() -> None: def test_deck_to_flat_dict_keys_are_mlflow_safe() -> None: import re - s = osd.parse_deck( - (DECKS_DIR / "two-stream-1d").read_text() - ) + s = osd.parse_deck((DECKS_DIR / "two-stream-1d").read_text()) flat = osd.deck_to_flat_dict(s) allowed = re.compile(r"^[A-Za-z0-9_./:\- ]+$") for k in flat: diff --git a/tests/test_osiris/test_diagnostics_plots.py b/tests/test_osiris/test_diagnostics_plots.py index cb469e09..c4084d0d 100644 --- a/tests/test_osiris/test_diagnostics_plots.py +++ b/tests/test_osiris/test_diagnostics_plots.py @@ -66,12 +66,23 @@ def _make_full_run(root: Path, n_steps: int = 6, nx: int = 16, npx: int = 12) -> for k in range(n_steps): it = k * 10 t = k * 0.5 - _write_dump(run_dir / "MS/FLD/e1" / f"e1-{it:06d}.h5", - "e1", rng.standard_normal(nx), t=t, it=it, axes=[x_ax]) - _write_dump(run_dir / "MS/DENSITY/electron/charge" / f"charge-electron-{it:06d}.h5", - "charge", rng.standard_normal(nx), t=t, it=it, axes=[x_ax]) - _write_dump(run_dir / "MS/PHA/x1p1/electron" / f"x1p1-electron-{it:06d}.h5", - "x1p1", rng.random((npx, nx)), t=t, it=it, axes=[x_ax, p_ax]) + _write_dump(run_dir / "MS/FLD/e1" / f"e1-{it:06d}.h5", "e1", rng.standard_normal(nx), t=t, it=it, axes=[x_ax]) + _write_dump( + run_dir / "MS/DENSITY/electron/charge" / f"charge-electron-{it:06d}.h5", + "charge", + rng.standard_normal(nx), + t=t, + it=it, + axes=[x_ax], + ) + _write_dump( + run_dir / "MS/PHA/x1p1/electron" / f"x1p1-electron-{it:06d}.h5", + "x1p1", + rng.random((npx, nx)), + t=t, + it=it, + axes=[x_ax, p_ax], + ) # HIST energy histories: iter, time, then value columns. hist = run_dir / "HIST" @@ -80,8 +91,8 @@ def _make_full_run(root: Path, n_steps: int = 6, nx: int = 16, npx: int = 12) -> fld = ["! iter time e1 e2 e3 b1 b2 b3"] par = ["! iter time ene"] for k, t in enumerate(times): - fld.append(f"{k*10} {t} {1.0 + 0.1*k} 0.0 0.0 0.0 0.0 0.0") - par.append(f"{k*10} {t} {100.0 - 0.1*k}") + fld.append(f"{k * 10} {t} {1.0 + 0.1 * k} 0.0 0.0 0.0 0.0 0.0") + par.append(f"{k * 10} {t} {100.0 - 0.1 * k}") (hist / "fld_ene").write_text("\n".join(fld) + "\n") (hist / "par01_ene").write_text("\n".join(par) + "\n") @@ -98,9 +109,7 @@ def test_field_energy_components_splits_e_and_b(tmp_path: Path) -> None: # Only e1 was dumped, so all energy is electric and B is identically zero. assert np.all(ds["E_energy"].values > 0) assert np.all(ds["B_energy"].values == 0) - np.testing.assert_allclose( - ds["total_field_energy"].values, ds["E_energy"].values - ) + np.testing.assert_allclose(ds["total_field_energy"].values, ds["E_energy"].values) def test_load_hist_energy_builds_conservation_total(tmp_path: Path) -> None: diff --git a/tests/test_osiris/test_io_and_plots.py b/tests/test_osiris/test_io_and_plots.py index 9d449ba4..c95461e7 100644 --- a/tests/test_osiris/test_io_and_plots.py +++ b/tests/test_osiris/test_io_and_plots.py @@ -16,7 +16,6 @@ from adept.osiris import io as oio from adept.osiris import plots as oplt - # Point this at a completed two-stream OSIRIS run to exercise the io/plots # loaders against real data; the tests skip cleanly when it is unset. EXISTING_RUN_ENV = "OSIRIS_TWOSTREAM_RUN" diff --git a/tests/test_osiris/test_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py index b9f0719e..bf25861d 100644 --- a/tests/test_osiris/test_plots_new_views.py +++ b/tests/test_osiris/test_plots_new_views.py @@ -74,17 +74,34 @@ def _make_rich_run(root: Path, n_steps: int = 6, nx: int = 16, npx: int = 12) -> _write_dump(run_dir / "MS/FLD/b2" / f"b2-{it:06d}.h5", "b2", -wave, t, it, [X_AX]) # Currents j1/j2/j3. for j in ("j1", "j2", "j3"): - _write_dump(run_dir / f"MS/FLD/{j}" / f"{j}-{it:06d}.h5", j, - rng.standard_normal(nx), t, it, [X_AX]) + _write_dump(run_dir / f"MS/FLD/{j}" / f"{j}-{it:06d}.h5", j, rng.standard_normal(nx), t, it, [X_AX]) # Density + thermal-velocity moments for a species. - _write_dump(run_dir / "MS/DENSITY/electron/charge" / f"charge-electron-{it:06d}.h5", - "charge", -np.abs(rng.standard_normal(nx)) - 1.0, t, it, [X_AX]) + _write_dump( + run_dir / "MS/DENSITY/electron/charge" / f"charge-electron-{it:06d}.h5", + "charge", + -np.abs(rng.standard_normal(nx)) - 1.0, + t, + it, + [X_AX], + ) for u in ("uth1", "uth2", "uth3"): - _write_dump(run_dir / f"MS/UDIST/electron/{u}" / f"{u}-electron-{it:06d}.h5", - u, 0.1 + 0.01 * rng.standard_normal(nx), t, it, [X_AX]) + _write_dump( + run_dir / f"MS/UDIST/electron/{u}" / f"{u}-electron-{it:06d}.h5", + u, + 0.1 + 0.01 * rng.standard_normal(nx), + t, + it, + [X_AX], + ) # Phase space (p, x). - _write_dump(run_dir / "MS/PHA/x1p1/electron" / f"x1p1-electron-{it:06d}.h5", - "x1p1", rng.random((npx, nx)), t, it, [X_AX, P_AX]) + _write_dump( + run_dir / "MS/PHA/x1p1/electron" / f"x1p1-electron-{it:06d}.h5", + "x1p1", + rng.random((npx, nx)), + t, + it, + [X_AX, P_AX], + ) return run_dir @@ -166,8 +183,7 @@ def test_temperature_series_from_uth(tmp_path: Path) -> None: def test_temperature_series_absent_returns_none(tmp_path: Path) -> None: run_dir = tmp_path / "run" - _write_dump(run_dir / "MS/DENSITY/ion/charge" / "charge-ion-000000.h5", - "charge", np.ones(8), 0.0, 0, [X_AX]) + _write_dump(run_dir / "MS/DENSITY/ion/charge" / "charge-ion-000000.h5", "charge", np.ones(8), 0.0, 0, [X_AX]) entries = oplt._species_diags(run_dir)["ion"] assert oplt._temperature_series(entries) is None @@ -216,4 +232,3 @@ def test_save_canned_plots_emits_new_views(tmp_path: Path) -> None: assert expected <= set(written) for path in written.values(): assert path.exists() and path.stat().st_size > 0 - diff --git a/tests/test_osiris/test_post_netcdf.py b/tests/test_osiris/test_post_netcdf.py index a7770e4d..dddcbdf7 100644 --- a/tests/test_osiris/test_post_netcdf.py +++ b/tests/test_osiris/test_post_netcdf.py @@ -126,9 +126,7 @@ def test_collect_respects_diagnostics_whitelist(tmp_path: Path) -> None: assert not (td / "binary" / "FLD" / "e2.nc").exists() -def _write_raw_dump( - path: Path, quantities: dict[str, np.ndarray], t: float, it: int -) -> None: +def _write_raw_dump(path: Path, quantities: dict[str, np.ndarray], t: float, it: int) -> None: """Write one OSIRIS-style RAW (particle) HDF5 dump. ``quantities`` maps a per-particle quantity name (``"p1"``, ``"x1"``, ...) @@ -156,9 +154,7 @@ def _write_raw_dump( def test_load_raw_h5_returns_particle_dataset(tmp_path: Path) -> None: p = tmp_path / "MS" / "RAW" / "species_1" / "RAW-species_1-000030.h5" npart = 7 - quants = { - q: np.arange(npart, dtype="float64") for q in ("ene", "p1", "p2", "p3", "q", "x1") - } + quants = {q: np.arange(npart, dtype="float64") for q in ("ene", "p1", "p2", "p3", "q", "x1")} _write_raw_dump(p, quants, t=1.5, it=30) ds = oio.load_raw_h5(p) @@ -341,9 +337,7 @@ def test_load_series_nc_roundtrips_grid_series(tmp_path: Path) -> None: # axis metadata is rebuilt into the dict attrs the plotters read. assert from_nc.attrs["axis_units"]["x1"] == from_ms.attrs["axis_units"]["x1"] # load_series dispatches to load_series_nc when handed a .nc file. - np.testing.assert_allclose( - oio.load_series(out / "FLD" / "e1.nc").values, from_ms.values - ) + np.testing.assert_allclose(oio.load_series(out / "FLD" / "e1.nc").values, from_ms.values) def test_save_run_datasets_persists_hist_energy(tmp_path: Path) -> None: diff --git a/tests/test_osiris/test_runner.py b/tests/test_osiris/test_runner.py index 388299eb..70f77d61 100644 --- a/tests/test_osiris/test_runner.py +++ b/tests/test_osiris/test_runner.py @@ -9,7 +9,6 @@ from adept.osiris import runner - # Resolved from the same env vars the runner itself honors, so the live-binary # smoke test below runs wherever OSIRIS is built and skips cleanly otherwise. OSIRIS_BIN_1D = os.environ.get("OSIRIS_BIN_1D") or os.environ.get("OSIRIS_BIN") From ba10baad4fb6d165297643447ccd46e479ebd110 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 16 Jun 2026 21:16:32 -0700 Subject: [PATCH 20/49] Populate OSIRIS units.yaml from the deck's reference scale BaseOsiris.write_units() previously returned {} so OSIRIS runs logged an empty units.yaml. Derive the physical reference scales (wp0, tp0, n0, v0, x0, c_light, beta, box_length, sim_duration) from the deck's simulation.n0 (density) or simulation.omega_p0 (frequency); when both are present, n0 wins, as in OSIRIS. This mirrors the canonical key set the other adept solvers emit so OSIRIS runs are comparable in MLflow. Adds skin_depth_normalization and skin_depth_normalization_from_frequency to normalization.py. OSIRIS has no single global reference temperature (species carry per-species thermal momenta), so the temperature-dependent keys (T0/nuee/logLambda_ee) are omitted. Co-Authored-By: Claude Opus 4.8 --- adept/normalization.py | 52 +++++++++++++++++++++++ adept/osiris/base.py | 49 +++++++++++++++++++++- tests/test_osiris/test_units.py | 74 +++++++++++++++++++++++++++++++++ 3 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 tests/test_osiris/test_units.py diff --git a/adept/normalization.py b/adept/normalization.py index 40b34daa..fc43b6be 100644 --- a/adept/normalization.py +++ b/adept/normalization.py @@ -116,3 +116,55 @@ def laser_normalization(laser_wavelength_str, T0_str): return PlasmaNormalization( m0=UREG.m_e, q0=UREG.e, n0=ne_crit, T0=T0, L0=one_over_k, v0=1 * UREG.c, tau=1 / omega_laser ) + + +def _osiris_normalization(wp0, n0): + """ + Core OSIRIS normalization built from an angular plasma frequency ``wp0`` and + its corresponding reference density ``n0``. + + OSIRIS normalizes time to 1/wp0, length to the collisionless skin depth + c/wp0, and velocity to the speed of light. + + Unit quantities are: + - L0 = c / wp0 (skin depth) + - v0 = c (speed of light) + - tau = 1 / wp0 (inverse plasma frequency) + + There is no reference temperature: OSIRIS has no single global temperature + (species carry their own per-species thermal momenta), so ``T0`` is left + unset and temperature-dependent quantities are not defined under this + normalization. + """ + wp0 = wp0.to("rad/s") + tau = 1 / wp0 + + v0 = 1 * UREG.c + x0 = (v0 / wp0).to("nm") + + return PlasmaNormalization(m0=UREG.m_e, q0=UREG.e, n0=n0.to("1/cc"), T0=None, L0=x0, v0=v0, tau=tau) + + +def skin_depth_normalization(n0_str): + """ + OSIRIS normalization referenced to a plasma density (``simulation.n0``). + + The reference plasma frequency is computed from the density, + wp0 = sqrt(n0 e^2 / (eps0 m_e)). See :func:`_osiris_normalization`. + """ + n0 = UREG.Quantity(n0_str) + wp0 = ((n0 * UREG.e**2.0 / (UREG.m_e * UREG.epsilon_0)) ** 0.5).to("rad/s") + return _osiris_normalization(wp0, n0) + + +def skin_depth_normalization_from_frequency(wp0_str): + """ + OSIRIS normalization referenced to a plasma frequency (``simulation.omega_p0``). + + The reference density is recovered from the frequency, + n0 = wp0^2 eps0 m_e / e^2, so the reported ``n0`` stays consistent with the + density-referenced form. See :func:`_osiris_normalization`. + """ + wp0 = UREG.Quantity(wp0_str) + n0 = (wp0**2 * UREG.epsilon_0 * UREG.m_e / UREG.e**2.0).to("1/cc") + return _osiris_normalization(wp0, n0) diff --git a/adept/osiris/base.py b/adept/osiris/base.py index a5d937a7..760b32fd 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -5,6 +5,7 @@ from typing import Any from adept._base_ import ADEPTModule +from adept.normalization import skin_depth_normalization, skin_depth_normalization_from_frequency from adept.osiris import deck as _deck from adept.osiris import post as _post from adept.osiris import runner as _runner @@ -38,7 +39,53 @@ def __init__(self, cfg: dict) -> None: cfg["deck"] = _deck.deck_to_flat_dict(self._sections) def write_units(self) -> dict: - return {} + """Derive physical reference scales from the deck's ``simulation`` section. + + OSIRIS normalizes time to ``1/wp0``, length to the skin depth ``c/wp0``, + and velocity to ``c``, where ``wp0`` is the reference plasma frequency. + That reference is set in the deck's ``simulation`` section by either the + density ``n0`` (cm^-3) or the frequency ``omega_p0`` (rad/s); when both + are present OSIRIS uses ``n0``, so we do too. The returned dict mirrors + the density-derived keys the other adept solvers log to ``units.yaml`` so + OSIRIS runs are comparable in MLflow. OSIRIS has no single global + reference temperature (species carry their own per-species thermal + momenta), so the temperature-dependent keys (``T0``/``nuee``/ + ``logLambda_ee``) are intentionally omitted. + """ + sim = self._iter_first_section("simulation") + n0 = sim.get("n0") + omega_p0 = sim.get("omega_p0") + if n0 is not None: + norm = skin_depth_normalization(f"{n0} / cc") + elif omega_p0 is not None: + norm = skin_depth_normalization_from_frequency(f"{omega_p0} rad/s") + else: + return {} + + quants: dict[str, Any] = { + "wp0": (1 / norm.tau).to("rad/s"), + "tp0": norm.tau.to("fs"), + "n0": norm.n0.to("1/cc"), + "v0": norm.v0.to("m/s"), + "x0": norm.L0.to("nm"), + "c_light": norm.speed_of_light_norm(), # == 1.0; OSIRIS normalizes v to c + "beta": 1.0 / norm.speed_of_light_norm(), + } + + space = self._iter_first_section("space") + xmin = self._first_array_value(space, "xmin") + xmax = self._first_array_value(space, "xmax") + if xmin is not None and xmax is not None: + quants["box_length"] = ((xmax[0] - xmin[0]) * norm.L0).to("micron") + + time = self._iter_first_section("time") + tmax = time.get("tmax") + if tmax is not None: + tmin = time.get("tmin", 0.0) + quants["sim_duration"] = ((float(tmax) - float(tmin)) * norm.tau).to("ps") + + self.cfg.setdefault("units", {})["derived"] = quants + return quants def get_derived_quantities(self) -> dict: """Lift a few useful scalars out of the deck for MLflow visibility.""" diff --git a/tests/test_osiris/test_units.py b/tests/test_osiris/test_units.py new file mode 100644 index 00000000..252cd40c --- /dev/null +++ b/tests/test_osiris/test_units.py @@ -0,0 +1,74 @@ +"""Tests for ``BaseOsiris.write_units`` — the OSIRIS ``units.yaml`` artifact. + +OSIRIS normalizes time to ``1/wp0``, length to the skin depth ``c/wp0``, and +velocity to ``c``, all set by the reference density ``simulation.n0`` declared +in the deck. ``write_units`` derives the same canonical scales the other adept +solvers log so OSIRIS runs are comparable in MLflow. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from adept.osiris import BaseOsiris + +DECKS_DIR = Path(__file__).parent / "decks" +SRS_DECK = DECKS_DIR / "srs-1d_lpi" # has simulation{n0=9.05e21}, xmax=1076.04, tmax=25000 + + +def test_write_units_density_derived_scales() -> None: + quants = BaseOsiris({"osiris": {"deck": str(SRS_DECK)}}).write_units() + + # wp0 and n0 depend only on the reference density and match the + # corresponding kinetic-srs/adept run with the same n0. + assert quants["n0"].to("1/cc").magnitude == pytest.approx(9.05e21, rel=1e-9) + assert quants["wp0"].to("rad/s").magnitude == pytest.approx(5.3668e15, rel=1e-3) + # OSIRIS length unit is the skin depth c/wp0, not the Debye length. + assert quants["x0"].to("nm").magnitude == pytest.approx(55.86, rel=1e-3) + # Velocity is normalized to c, so c_light and beta are both unity. + assert quants["c_light"] == pytest.approx(1.0) + assert quants["beta"] == pytest.approx(1.0) + assert quants["v0"].to("m/s").magnitude == pytest.approx(299792458.0, rel=1e-6) + # Geometry: box_length = (xmax - xmin) * skin depth, duration = tmax / wp0. + assert quants["box_length"].to("micron").magnitude == pytest.approx(60.108, rel=1e-3) + assert quants["sim_duration"].to("ps").magnitude == pytest.approx(4.6583, rel=1e-3) + + +def test_write_units_omits_temperature_keys() -> None: + quants = BaseOsiris({"osiris": {"deck": str(SRS_DECK)}}).write_units() + # OSIRIS has no single global temperature, so temperature-dependent + # quantities are not defined and must not appear. + for key in ("T0", "nuee", "logLambda_ee"): + assert key not in quants + + +def test_write_units_from_reference_frequency(tmp_path: Path) -> None: + # omega_p0 = the plasma frequency of n0 = 9.05e21 cm^-3, so the frequency + # form must reproduce the density form's scales (and recover n0). + deck = tmp_path / "omega_p0_deck" + deck.write_text( + "simulation { omega_p0 = 5.3668e15, }\n" + "space { xmin(1) = 0.0, xmax(1) = 1076.04, }\n" + "time { tmin = 0.0, tmax = 25000.0, }\n" + ) + quants = BaseOsiris({"osiris": {"deck": str(deck)}}).write_units() + assert quants["wp0"].to("rad/s").magnitude == pytest.approx(5.3668e15, rel=1e-4) + assert quants["n0"].to("1/cc").magnitude == pytest.approx(9.05e21, rel=1e-3) + assert quants["x0"].to("nm").magnitude == pytest.approx(55.86, rel=1e-3) + assert quants["box_length"].to("micron").magnitude == pytest.approx(60.108, rel=1e-3) + + +def test_write_units_density_takes_precedence_over_frequency(tmp_path: Path) -> None: + # OSIRIS uses n0 when both are present; omega_p0 here is deliberately wrong. + deck = tmp_path / "both_deck" + deck.write_text("simulation { n0 = 9.05e21, omega_p0 = 1.0e14, }\n") + quants = BaseOsiris({"osiris": {"deck": str(deck)}}).write_units() + assert quants["wp0"].to("rad/s").magnitude == pytest.approx(5.3668e15, rel=1e-3) + + +def test_write_units_returns_empty_without_reference_density(tmp_path: Path) -> None: + deck = tmp_path / "no_n0" + deck.write_text("grid { nx_p(1:1) = 64, }\n") + assert BaseOsiris({"osiris": {"deck": str(deck)}}).write_units() == {} From c572b124a0235145e14e0a67d71f2354dc5bc999 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 22 Jun 2026 11:07:42 -0700 Subject: [PATCH 21/49] Fixed hardcoded paths in the configs/twostream-1d{-short}.yaml --- configs/osiris/twostream-1d-short.yaml | 5 +++-- configs/osiris/twostream-1d.yaml | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/configs/osiris/twostream-1d-short.yaml b/configs/osiris/twostream-1d-short.yaml index d6d154cc..7fde220d 100644 --- a/configs/osiris/twostream-1d-short.yaml +++ b/configs/osiris/twostream-1d-short.yaml @@ -5,8 +5,9 @@ mlflow: run: smoke-tmax-1 osiris: - deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d - binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e + deck: tests/test_osiris/decks/two-stream-1d + # binary omitted: the runner resolves it from OSIRIS_BIN_1D, then OSIRIS_BIN. + # Set it here only to override that (e.g. binary: /path/to/osiris-1D.e). mpi_ranks: 1 run_root: ./checkpoints overrides: diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml index 157880dd..f91c557e 100644 --- a/configs/osiris/twostream-1d.yaml +++ b/configs/osiris/twostream-1d.yaml @@ -5,8 +5,9 @@ mlflow: run: cold-equal-beams osiris: - deck: /home/phil/Desktop/pic/projects/twostream-pic/two-stream-1d - binary: /home/phil/Desktop/pic/osiris/bin/osiris-1D.e + deck: tests/test_osiris/decks/two-stream-1d + # binary omitted: the runner resolves it from OSIRIS_BIN_1D, then OSIRIS_BIN. + # Set it here only to override that (e.g. binary: /path/to/osiris-1D.e). mpi_ranks: 1 run_root: ./checkpoints # overrides: dict-of-dicts merged into the parsed deck before render. From 493fbda7856fe0cb729f9cd65f5f14eb311feddd Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 22 Jun 2026 11:07:55 -0700 Subject: [PATCH 22/49] Updated osiris-adept-usage --- docs/osiris-adept-usage.md | 112 ++++++++++++++++++++++--------------- 1 file changed, 68 insertions(+), 44 deletions(-) diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index ed29b533..3ffd39e5 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -3,11 +3,11 @@ ## File structure ``` -adept/ (repo root, ~/Desktop/adept/adept/) -├── run.py ← existing CLI entry (unchanged) +adept/ (repo root) +├── run.py ← CLI entry (--cfg path, no .yaml suffix) ├── adept/ -│ ├── _base_.py ← MODIFIED: dispatcher gained `osiris` branch -│ └── osiris/ ← NEW package +│ ├── _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 @@ -18,53 +18,75 @@ adept/ (repo root, ~/Desktop/adept/adep │ ├── 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, optional MS/ tarball, -│ scalar metrics +│ ├── post.py ◀── post-run collection: final-step HDF5 copy, +│ │ NetCDF export, scalar metrics +│ ├── io.py ◀── HDF5/NetCDF readers + dataset save/load +│ ├── plots.py ◀── canned plot set (save_canned_plots) +│ └── regen.py ◀── regenerate plots offline from saved NetCDFs ├── configs/ -│ └── osiris/ ← NEW: example manifests -│ ├── twostream-1d.yaml full 30-ω_p run -│ ├── twostream-1d-short.yaml tmax=1.0 smoke -│ └── twostream-1d-uploadall.yaml tmax=0.5 + ms.tar.gz +│ └── osiris/ example manifests +│ ├── twostream-1d.yaml full run (deck tmax=100) +│ └── twostream-1d-short.yaml tmax=1.0 smoke └── tests/ - └── test_osiris/ ← NEW - ├── test_deck_roundtrip.py 15 parser tests - └── test_runner.py 5 runner 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) (any path you point at) - │ │ - ▼ ▼ - run.py --cfg … parse_deck_file() ──┐ - │ │ - ▼ ▼ - ergoExo.setup ──► BaseOsiris.__init__ ──► merge_overrides ──► render_deck - │ │ - ▼ ▼ - cfg["deck"] = flat_dict run_dir/os-stdin - │ │ - ▼ ▼ - log_params → MLflow runner.run_osiris(mpirun…) - │ - ▼ - run_dir/MS/{FLD,PHA,…} - │ - ▼ - post.collect → td/ → MLflow + 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_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: + +```bash +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): + ```bash -conda activate adept -cd ~/Desktop/adept/adept -python run.py --cfg configs/osiris/twostream-1d-short # smoke -python run.py --cfg configs/osiris/twostream-1d # full -mlflow ui --backend-store-uri file://$(pwd)/mlruns # browse +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 @@ -77,8 +99,9 @@ mlflow: run: cold-equal-beams # required osiris: - deck: /path/to/native/deck # required: source of truth - binary: /path/to/osiris-1D.e # or OSIRIS_BIN / OSIRIS_BIN_D + deck: tests/test_osiris/decks/two-stream-1d # required: source of truth (repo-relative) + # binary: /path/to/osiris-1D.e # optional: overrides the env-var default + # # (OSIRIS_BIN_D → OSIRIS_BIN) mpi_ranks: 1 # 1 → direct, >1 → mpirun -n N run_root: ./checkpoints # parent of per-run dirs (default) # NOTE: the default sits inside checkpoints/ deliberately — sync-up.sh @@ -113,7 +136,7 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | Kind | Content | | ----------- | -------------------------------------------------------------------------------- | | Params | every deck key, flattened — `deck.grid.nx_p_1:1`, `deck.species_0.ufl_1:3.0`, … | -| Params | the manifest itself — `solver`, `osiris.binary`, `output.diagnostics_to_log`, … | +| Params | the manifest itself — `solver`, `osiris.deck`, `output.diagnostics_to_log`, … | | Metrics | `wall_time_s`, `exit_code`, `field_energy_final`, `final_iter`, `run_time`, `postprocess_time` | | Artifacts | `config.yaml`, `derived_config.yaml`, `units.yaml` (adept stock) | | Artifacts | `os-stdin` (rendered OSIRIS deck), `stdout.log`, `stderr.log` | @@ -121,6 +144,7 @@ Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1 | Artifacts | `plots/…` — canned PNGs (see below) | ## 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$`, …). @@ -165,8 +189,8 @@ Edit `adept/osiris/post.py:collect`. The h5 dump for each diagnostic is at `run_ Native-deck-as-truth: just write the deck, point a manifest at it, run. No code changes. ```bash -cp my-new.deck ~/Desktop/pic/projects/whatever/ +cp my-new.deck tests/test_osiris/decks/ cp configs/osiris/twostream-1d.yaml configs/osiris/my-new.yaml $EDITOR configs/osiris/my-new.yaml # change deck path + mlflow.run -python run.py --cfg configs/osiris/my-new +uv run run.py --cfg configs/osiris/my-new ``` From 28c862552f1618c5b4ee003771eee279e0cd710a Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 24 Jun 2026 23:09:53 -0700 Subject: [PATCH 23/49] Fixed erroneous error handling when osiris through a warning for the Sentoku collision model --- adept/osiris/runner.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index 11ce118e..e3896fd9 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -34,6 +34,11 @@ def _looks_like_osiris_error(line: str) -> bool: low = line.strip().lower() if not low: return False + # OSIRIS prefixes diagnostics with (*warning*) / (*error*); a warning is + # never a failure even when its text contains the word "error" (e.g. the + # Sentoku-collisions "factor of 2 error" warning). + if "(*warning*)" in low: + return False if any(noise.lower() in low for noise in _OSIRIS_STDERR_NOISE): return False return any(tok in low for tok in _OSIRIS_ERR_TOKENS) From 7f2f7b87021f2a09e520b17e423e91891c6828f1 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 26 Jun 2026 12:38:09 -0700 Subject: [PATCH 24/49] Add adaptive box sizing from density gradient scale length to OSIRIS Given osiris.density.gradient_scale_length, BaseOsiris scales the simulation box so the deck's linear density ramp realizes that gradient scale length at the reference density (default n_c/4), mirroring adept's _lpse2d / kinetic_srs grid sizing. The transform is a single spatial scale factor applied to space.xmin/xmax, every profile.x, and every diag_species phase-space window; grid.nx_p scales too, holding dx fixed (rounded up to a multiple of node_number(1)). dt/tmax are untouched, so CFL is preserved. 1D decks only. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/base.py | 11 ++ adept/osiris/density.py | 260 +++++++++++++++++++++++++ docs/osiris-adept-usage.md | 40 ++++ tests/test_osiris/test_adaptive_box.py | 170 ++++++++++++++++ 4 files changed, 481 insertions(+) create mode 100644 adept/osiris/density.py create mode 100644 tests/test_osiris/test_adaptive_box.py diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 760b32fd..7c6fd371 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -7,6 +7,7 @@ from adept._base_ import ADEPTModule from adept.normalization import skin_depth_normalization, skin_depth_normalization_from_frequency from adept.osiris import deck as _deck +from adept.osiris import density as _density from adept.osiris import post as _post from adept.osiris import runner as _runner @@ -31,6 +32,16 @@ def __init__(self, cfg: dict) -> None: if overrides: _deck.merge_overrides(self._sections, overrides) + # Optional: scale the box from a target density gradient scale length + # (mirrors adept's _lpse2d / kinetic_srs grid sizing). Runs *after* + # overrides, so an explicit space.xmax override is superseded by the + # gradient-derived box. The computed quantities are stashed back under + # osiris.density.derived for MLflow provenance. + density_cfg = osiris_cfg.get("density") + computed = _density.apply_gradient_scale_length(self._sections, density_cfg) + if computed is not None: + density_cfg.setdefault("derived", {}).update(computed) + # Surface the parsed (post-override) deck inside cfg so adept's # log_params picks every parameter up as a flat MLflow param. # The raw overrides dict is intentionally popped above to avoid diff --git a/adept/osiris/density.py b/adept/osiris/density.py new file mode 100644 index 00000000..0a191780 --- /dev/null +++ b/adept/osiris/density.py @@ -0,0 +1,260 @@ +"""Adaptive box sizing from a density gradient scale length. + +OSIRIS decks fix the simulation box (``space.xmax``) and the density ramp +(``profile.fx`` / ``profile.x``) by hand. This module lets a manifest instead +request a target *gradient scale length* ``L`` and have the box scaled so the +linear density ramp realizes that ``L`` at a chosen reference density — mirroring +ADEPT's envelope (``adept/_lpse2d/helpers.py``) and Vlasov (``kinetic_srs``) +solvers, which size their grid the same way. + +Geometry. The deck's ``profile`` is a piecewise-linear density ramp +``n(x): nmin -> nmax`` across its interior control points. For a linear ramp the +local scale length is + + L(x) = n(x) / (dn/dx) = n(x) * ramp_span / (nmax - nmin) + +so requiring ``L(n_ref) = L_target`` at the reference density ``n_ref`` (default +the quarter-critical surface, ``n_c/4 = 0.25`` in ``n_c`` units) fixes the ramp +span: + + ramp_span = L_target / n_ref * (nmax - nmin) + +which is the same relation ADEPT uses (``Lgrid = L / 0.25 * (nmax - nmin)``). + +The transform is then a single spatial **scale factor** + + s = ramp_span_target / ramp_span_deck + +applied uniformly to every length in the deck: ``space.xmin`` / ``space.xmax``, +every ``profile.x`` array, and every ``diag_species`` phase-space window +(``ps_xmin`` / ``ps_xmax``). ``grid.nx_p`` is scaled by ``s`` as well so the cell +size ``dx`` is held fixed (rounded up to a clean multiple of the per-dimension +node count for even domain decomposition). Time (``dt``, ``tmax``) is untouched, +so the CFL ratio is preserved. + +Activation. Only when the manifest carries an ``osiris.density`` block with a +``gradient_scale_length`` (alias ``gradient scale length``). Without it, decks +run with their hand-set box, unchanged. + +1D only: raises ``NotImplementedError`` for multi-dimensional decks (the SRS LPI +use case is 1D). +""" + +from __future__ import annotations + +import math +from typing import Any + +from adept.normalization import UREG, skin_depth_normalization, skin_depth_normalization_from_frequency +from adept.osiris.deck import Sections + + +def apply_gradient_scale_length(sections: Sections, density_cfg: dict | None) -> dict | None: + """Scale the OSIRIS box so the density ramp has gradient scale length ``L``. + + ``density_cfg`` is the manifest's ``osiris.density`` block. Returns a dict of + the computed quantities (for logging) or ``None`` when the feature is + inactive (no ``density_cfg``, or no gradient scale length given). + + Mutates ``sections`` in place. Recognized ``density_cfg`` keys: + + * ``gradient_scale_length`` (alias ``gradient scale length``) — target ``L``, + a unit string (``"300um"``) or a number already in ``c/wp0`` units. + * ``min`` / ``max`` — optional ``nmin`` / ``nmax`` in ``n_c`` units; default is + to read the ramp's interior endpoints from the deck's ``profile.fx``. When + given, they are also written into the primary ``profile.fx`` endpoints. + * ``reference_density`` — density (``n_c`` units) at which ``L`` is defined; + default ``0.25`` (quarter-critical). + """ + if not density_cfg: + return None + L_spec = density_cfg.get("gradient_scale_length", density_cfg.get("gradient scale length")) + if L_spec is None: + return None + + reference_density = float(density_cfg.get("reference_density", 0.25)) + + simulation = _first_section(sections, "simulation") + space = _first_section(sections, "space") + grid = _first_section(sections, "grid") + profile = _first_section(sections, "profile") + missing = [name for name, sec in (("space", space), ("grid", grid), ("profile", profile)) if sec is None] + if missing: + raise ValueError( + f"osiris.density (adaptive box sizing) requires the deck to define section(s): {', '.join(missing)}" + ) + + # --- 1D only ---------------------------------------------------------- + nx_key = _array_key(grid, "nx_p") + if nx_key is None: + raise ValueError("osiris.density requires grid.nx_p in the deck") + nx_val = grid[nx_key] + if isinstance(nx_val, list) and len(nx_val) > 1: + raise NotImplementedError( + "osiris.density adaptive box sizing is implemented for 1D decks only " + f"(grid.{nx_key} has {len(nx_val)} dimensions)" + ) + nx_old = int(nx_val[0] if isinstance(nx_val, list) else nx_val) + + # --- ramp geometry & densities from the primary profile --------------- + x_key = _array_key(profile, "x") + fx_key = _array_key(profile, "fx") + if x_key is None or fx_key is None: + raise ValueError("osiris.density requires profile.x and profile.fx in the deck") + x_old = [float(v) for v in _as_list(profile[x_key])] + fx = [float(v) for v in _as_list(profile[fx_key])] + if len(x_old) < 3 or len(fx) < 3: + raise ValueError( + f"profile needs >= 3 control points to define a density ramp " + f"(got {len(x_old)} x-points, {len(fx)} fx-values)" + ) + + given_minmax = "min" in density_cfg or "max" in density_cfg + if given_minmax: + nmin = float(density_cfg["min"]) + nmax = float(density_cfg["max"]) + else: + nmin, nmax = fx[1], fx[-2] # interior ramp endpoints + if nmax <= nmin: + raise ValueError(f"density max ({nmax}) must exceed min ({nmin})") + + ramp_span_old = x_old[-2] - x_old[1] + if ramp_span_old <= 0: + raise ValueError(f"profile.{x_key} interior control points must be increasing (got ramp span {ramp_span_old})") + + # --- target ramp span and the resulting scale factor ------------------ + L_norm = _normalized_length(L_spec, simulation) + ramp_span_new = L_norm * (nmax - nmin) / reference_density + s = ramp_span_new / ramp_span_old + + xmax_old = _scalar(space, "xmax") + if xmax_old is None: + raise ValueError("osiris.density requires space.xmax in the deck") + xmin_old = _scalar(space, "xmin") or 0.0 + box_new = (xmax_old - xmin_old) * s + + # constant dx -> nx scales with the box; round UP (don't under-resolve) + # to a clean multiple of the per-dim node count so OSIRIS splits evenly. + node_count = _node_count(sections) + nx_new = _round_up_multiple(math.ceil(nx_old * s), node_count) + + # --- apply the rescale in place --------------------------------------- + if given_minmax: + fx[1], fx[-2] = nmin, nmax + profile[fx_key] = fx if isinstance(profile[fx_key], list) else fx[0] + + _scale_array(space, "xmin", s) + _scale_array(space, "xmax", s) + for sec_name, params in sections: + if sec_name == "profile": + _scale_array(params, "x", s) + elif sec_name == "diag_species": + _scale_array(params, "ps_xmin", s) + _scale_array(params, "ps_xmax", s) + grid[nx_key] = [nx_new] if isinstance(nx_val, list) else nx_new + + return { + "gradient_scale_length_requested": str(L_spec), + "gradient_scale_length_norm": L_norm, + "reference_density": reference_density, + "density_min": nmin, + "density_max": nmax, + "ramp_span_norm": ramp_span_new, + "box_norm": box_new, + "nx": nx_new, + "scale_factor": s, + } + + +# --- helpers -------------------------------------------------------------- + + +def _first_section(sections: Sections, name: str) -> dict | None: + for sec_name, params in sections: + if sec_name == name: + return params + return None + + +def _array_key(params: dict, base: str) -> str | None: + """Key in ``params`` whose base name (before any ``(``) equals ``base``.""" + for k in params: + if k == base or k.split("(", 1)[0] == base: + return k + return None + + +def _as_list(v: Any) -> list: + return list(v) if isinstance(v, list) else [v] + + +def _scalar(params: dict, base: str) -> float | None: + key = _array_key(params, base) + if key is None: + return None + v = params[key] + return float(v[0]) if isinstance(v, list) else float(v) + + +def _scale_array(params: dict, base: str, factor: float) -> bool: + """Multiply the value(s) at ``base`` by ``factor`` in place; preserve shape.""" + key = _array_key(params, base) + if key is None: + return False + v = params[key] + if isinstance(v, list): + params[key] = [round(float(x) * factor, 10) for x in v] + else: + params[key] = round(float(v) * factor, 10) + return True + + +def _node_count(sections: Sections) -> int: + node_conf = _first_section(sections, "node_conf") + if node_conf is None: + return 1 + key = _array_key(node_conf, "node_number") + if key is None: + return 1 + val = node_conf[key] + first = val[0] if isinstance(val, list) else val + try: + return max(1, int(first)) + except (TypeError, ValueError): + return 1 + + +def _round_up_multiple(n: int, m: int) -> int: + if m <= 1: + return int(n) + return int(math.ceil(n / m) * m) + + +def _normalized_length(L_spec: Any, simulation: dict | None) -> float: + """Target ``L`` in skin-depth (``c/wp0``) units. + + A bare number is taken to be already normalized; a unit string is converted + using the deck's reference density (``simulation.n0`` / ``omega_p0``). + """ + if isinstance(L_spec, (int, float)) and not isinstance(L_spec, bool): + return float(L_spec) + norm = _skin_depth_norm(simulation) + if norm is None: + raise ValueError( + "osiris.density.gradient_scale_length was given with physical units " + f"({L_spec!r}) but the deck's simulation section has neither n0 nor " + "omega_p0 to set the c/wp0 length scale" + ) + return float((UREG.Quantity(str(L_spec)) / norm.L0).to("").magnitude) + + +def _skin_depth_norm(simulation: dict | None): + if not simulation: + return None + n0 = simulation.get("n0") + omega_p0 = simulation.get("omega_p0") + if n0 is not None: + return skin_depth_normalization(f"{n0} / cc") + if omega_p0 is not None: + return skin_depth_normalization_from_frequency(f"{omega_p0} rad/s") + return None diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index 3ffd39e5..c3743be7 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -113,6 +113,13 @@ osiris: # purge policy will take care of stale ones eventually. extra_mpi_args: ["--oversubscribe"] # optional, passed to mpirun + density: # optional: adaptive box sizing (1D) + gradient_scale_length: 300um # target L_n; scales the box so the + # deck's density ramp realizes this L + # min: 0.225 # nmin (n_c units); default: deck profile.fx + # max: 0.275 # nmax (n_c units); default: deck profile.fx + # reference_density: 0.25 # density (n_c units) where L is defined + overrides: # optional: applied before render time: {tmax: 50.0} # merge into the (one) time block grid: {nx_p: [256]} # refresh an array key @@ -131,6 +138,39 @@ output: Override keys can use the **base name** (`nx_p`) or the **exact key** (`nx_p(1:1)`). Indexed `{0: …, 1: …}` form addresses occurrences of repeated sections (`species`, `udist`, `profile`, `spe_bound`, `diag_species`, `zpulse`, …) in source order. +## Adaptive box sizing from a density gradient scale length + +`osiris.density` (1D decks) scales the simulation box so the deck's linear density +ramp realizes a target gradient scale length `L`, mirroring how adept's `_lpse2d` +and `kinetic_srs` solvers size their grids. This is the **inverse** of holding the +box fixed and steepening the ramp: here the density range (`nmin`/`nmax`) is held +and the box length follows `L`. + +For a linear ramp `n(x): nmin → nmax`, the local scale length is +`L(x) = n(x)/(dn/dx) = n(x)·ramp_span/(nmax−nmin)`, so requiring `L(n_ref) = L` at the +reference density `n_ref` (default the quarter-critical surface, `n_c/4 = 0.25`) +fixes the ramp span — the same relation adept uses, `ramp_span = L/0.25·(nmax−nmin)`. + +The result is a single spatial scale factor `s` applied to **every** length in the +deck: `space.xmin`/`space.xmax`, all `profile.x` arrays, and all `diag_species` +phase-space windows (`ps_xmin`/`ps_xmax`). `grid.nx_p` is scaled by `s` too — holding +the cell size `dx` fixed (rounded up to a multiple of `node_conf.node_number(1)` for +even domain decomposition). Time (`dt`, `tmax`) is untouched, so the CFL ratio is +preserved. + +- Activates only when `osiris.density.gradient_scale_length` is present (decks + otherwise run with their hand-set box, unchanged). Runs **after** `overrides`, so + it supersedes any `space.xmax` override. +- `gradient_scale_length` takes a unit string (`300um`, converted via the deck's + `simulation.n0`/`omega_p0`) or a bare number already in `c/wp0` units. +- `min`/`max` default to the ramp's interior `profile.fx` endpoints; if given, they + are written into the primary `profile.fx`. +- The computed quantities (`box_norm`, `nx`, `scale_factor`, …) are logged under + `osiris.density.derived.*`. +- Multi-dimensional decks raise `NotImplementedError`. Drive positions (e.g. a + `zpulse` spatial center) are **not** rescaled — boundary antennas like the SRS + deck's `antenna_array` have no position to scale. + ## What lands in MLflow | Kind | Content | diff --git a/tests/test_osiris/test_adaptive_box.py b/tests/test_osiris/test_adaptive_box.py new file mode 100644 index 00000000..fe0cdd01 --- /dev/null +++ b/tests/test_osiris/test_adaptive_box.py @@ -0,0 +1,170 @@ +"""Tests for ``osiris.density`` adaptive box sizing. + +When a manifest carries an ``osiris.density.gradient_scale_length``, +``BaseOsiris`` scales the simulation box so the deck's linear density ramp +realizes that gradient scale length at the reference density (default the +quarter-critical surface). This mirrors adept's ``_lpse2d`` / ``kinetic_srs`` +grid sizing: ``ramp_span = L / n_ref * (nmax - nmin)``. + +The reference deck ``srs-1d_lpi`` ramps n: 0.225 -> 0.275 (n_c units) across +x in (1e-5, 1076.03) c/wp0, box xmax = 1076.04, nx_p = 6000, with n0 = 9.05e21 +(critical density for 351 nm, so the skin depth c/wp0 = 55.86 nm). +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest + +from adept.osiris import BaseOsiris +from adept.osiris import deck as osd +from adept.osiris import density as den + +DECKS_DIR = Path(__file__).parent / "decks" +SRS_DECK = DECKS_DIR / "srs-1d_lpi" +TWO_D_DECK = DECKS_DIR / "F-Tsung_2d_lpi_deck" + +DECK_DX = 1076.04 / 6000 # cell size of the reference deck (c/wp0) + + +def _value(sections, name, base): + """First value in section ``name`` whose key base name is ``base``.""" + for sec_name, params in sections: + if sec_name == name: + for k, v in params.items(): + if k == base or k.split("(", 1)[0] == base: + return v + raise KeyError(f"{name}.{base} not found") + + +def _scaled(L="300um", **extra): + sections = osd.parse_deck_file(SRS_DECK) + computed = den.apply_gradient_scale_length(sections, {"gradient_scale_length": L, **extra}) + return sections, computed + + +def test_inactive_without_density_block() -> None: + sections = osd.parse_deck_file(SRS_DECK) + before = copy.deepcopy(sections) + assert den.apply_gradient_scale_length(sections, None) is None + assert den.apply_gradient_scale_length(sections, {}) is None + # no gradient scale length -> still a no-op + assert den.apply_gradient_scale_length(sections, {"min": 0.2}) is None + assert sections == before, "deck must be untouched when the feature is inactive" + + +def test_box_scaled_to_requested_scale_length() -> None: + sections, computed = _scaled("300um") + # ramp_span = L/0.25 * (0.275-0.225); L=300um at c/wp0=55.86nm -> ~5371 c/wp0 + assert _value(sections, "space", "xmax") == pytest.approx(1074.11, rel=1e-4) + assert _value(sections, "space", "xmin") == pytest.approx(0.0) + assert computed["box_norm"] == pytest.approx(1074.11, rel=1e-4) + assert computed["density_min"] == pytest.approx(0.225) + assert computed["density_max"] == pytest.approx(0.275) + assert computed["reference_density"] == pytest.approx(0.25) + + +def test_constant_dx_nx_scales_and_divides_node_count() -> None: + sections, computed = _scaled("300um") + nx = _value(sections, "grid", "nx_p") + box = _value(sections, "space", "xmax") + # cell size is held fixed (constant dx) to within node-count rounding + assert box / nx == pytest.approx(DECK_DX, rel=5e-3) + # node_number(1) = 4 in the deck -> nx_p must split evenly across 4 nodes + assert nx % 4 == 0 + assert computed["nx"] == nx + + +def test_diagnostic_window_and_profile_edges_track_box() -> None: + sections, _ = _scaled("300um") + box = _value(sections, "space", "xmax") + # phase-space diagnostic window follows the box + assert _value(sections, "diag_species", "ps_xmax") == pytest.approx(box) + assert _value(sections, "diag_species", "ps_xmin") == pytest.approx(0.0) + # profile control points (incl. the tiny vacuum edges) scale with the box + prof_x = _value(sections, "profile", "x") + assert prof_x[0] == pytest.approx(0.0) + assert prof_x[-1] == pytest.approx(box) + # density values are untouched (no min/max override given) + assert _value(sections, "profile", "fx") == [0.0, 0.225, 0.275, 0.0] + + +def test_box_scales_linearly_with_scale_length() -> None: + _, half = _scaled("150um") + _, base = _scaled("300um") + _, dbl = _scaled("600um") + assert half["box_norm"] == pytest.approx(base["box_norm"] / 2, rel=1e-6) + assert dbl["box_norm"] == pytest.approx(base["box_norm"] * 2, rel=1e-6) + + +def test_round_trips_the_deck_geometry() -> None: + # The reference deck encodes L_n = 0.25 * 1076.03 / 0.05 c/wp0 = 300.54 um; + # requesting that L must reproduce the deck's box. + sections, computed = _scaled("300.54um") + assert _value(sections, "space", "xmax") == pytest.approx(1076.04, rel=1e-4) + assert computed["scale_factor"] == pytest.approx(1.0, rel=2e-3) + + +def test_numeric_scale_length_is_treated_as_normalized() -> None: + # A bare number is taken to be already in c/wp0 units (no n0 conversion). + _, computed = _scaled(500.0) + assert computed["gradient_scale_length_norm"] == pytest.approx(500.0) + # ramp_span = 500/0.25 * 0.05 = 100 c/wp0 + assert computed["ramp_span_norm"] == pytest.approx(100.0) + + +def test_min_max_override_widens_box_and_rewrites_fx() -> None: + sections, computed = _scaled("300um", min=0.2, max=0.3) + # (nmax-nmin) doubled from 0.05 to 0.1 -> ramp span (and box) doubles + _, base = _scaled("300um") + assert computed["box_norm"] == pytest.approx(base["box_norm"] * 2, rel=1e-6) + # the new densities are written into the ramp endpoints + assert _value(sections, "profile", "fx") == [0.0, 0.2, 0.3, 0.0] + + +def test_custom_reference_density() -> None: + # ramp_span is inversely proportional to n_ref + _, qc = _scaled("300um", reference_density=0.25) + _, half = _scaled("300um", reference_density=0.125) + assert half["box_norm"] == pytest.approx(qc["box_norm"] * 2, rel=1e-6) + + +def test_multidimensional_deck_raises() -> None: + sections = osd.parse_deck_file(TWO_D_DECK) + with pytest.raises(NotImplementedError, match="1D decks only"): + den.apply_gradient_scale_length(sections, {"gradient_scale_length": "300um"}) + + +def test_physical_length_without_reference_density_raises(tmp_path: Path) -> None: + deck = tmp_path / "no_n0" + deck.write_text( + "grid { nx_p = 100, }\n" + "space { xmin(1) = 0.0, xmax(1) = 100.0, }\n" + "profile { num_x = 4, fx(1:4,1) = 0.0, 0.2, 0.3, 0.0, x(1:4,1) = 0, 1, 99, 100, }\n" + ) + sections = osd.parse_deck_file(deck) + with pytest.raises(ValueError, match="neither n0 nor"): + den.apply_gradient_scale_length(sections, {"gradient_scale_length": "300um"}) + + +def test_baseosiris_applies_box_and_records_derived() -> None: + cfg = {"osiris": {"deck": str(SRS_DECK), "density": {"gradient_scale_length": "300um"}}} + module = BaseOsiris(cfg) + # box was scaled on the live sections that __call__ will render + assert _value(module._sections, "space", "xmax") == pytest.approx(1074.11, rel=1e-4) + # computed quantities are stashed back for MLflow provenance + derived = cfg["osiris"]["density"]["derived"] + assert derived["box_norm"] == pytest.approx(1074.11, rel=1e-4) + assert derived["nx"] == _value(module._sections, "grid", "nx_p") + # write_units sees the scaled box + quants = module.write_units() + assert quants["box_length"].to("micron").magnitude == pytest.approx(60.0, rel=1e-2) + + +def test_baseosiris_unchanged_without_density() -> None: + cfg = {"osiris": {"deck": str(SRS_DECK)}} + module = BaseOsiris(cfg) + assert _value(module._sections, "space", "xmax") == pytest.approx(1076.04) + assert "density" not in cfg["osiris"] From 1a05a8ccc4e881cd56f1aa4a9addf8e03f4d800a Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 30 Jun 2026 11:05:16 -0700 Subject: [PATCH 25/49] Trimmed memory usage for omega-k plots --- adept/osiris/plots.py | 139 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 127 insertions(+), 12 deletions(-) diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index af819b07..a7d43c83 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -48,6 +48,85 @@ from adept.osiris import io as _io +# --- render-resolution downsampling --------------------------------------- + +# A heatmap is rasterized to a PNG of a fixed pixel size, so feeding pcolormesh +# a (t, x) array with vastly more cells than on-screen pixels is pure waste: the +# QuadMesh materializes an (rows+1, cols+1) float64 coordinate mesh, a float64 +# value copy, and an RGBA buffer, then Agg rasterizes ~one quad per *cell* — all +# for detail that averages away below a single pixel. Boxcar-averaging the data +# down to roughly the rendered resolution first removes that blow-up (the +# dominant term in the post-processing OOM); the small render arrays are also +# what keep peak memory in the low tens of GB. + +_RENDER_OVERSAMPLE = 3.0 # cells kept per on-screen pixel (margin for zoom / save dpi) + + +def _render_cells( + ax: plt.Axes, + *, + oversample: float = _RENDER_OVERSAMPLE, + visible_frac: tuple[float, float] = (1.0, 1.0), + cap: int = 4000, +) -> tuple[int, int]: + """Target ``(rows, cols)`` to render ``ax`` at, from its on-screen size. + + The axes' size in pixels (inches from the figure size × its subplot fraction, + times the figure dpi) sets the floor; ``oversample`` keeps a few cells per + pixel so the PNG (often saved at a higher dpi) and any later zoom stay sharp, + rather than downsampling to the exact pixel grid with no margin. The data is + only ever decimated to this, never below the eventual pixel count. + + ``visible_frac`` is the fraction of each axis ``(rows, cols)`` that a later + ``set_ylim`` / ``set_xlim`` will actually display (e.g. the equal-aspect + omega-k panel shows only a small Nyquist window): the full array is then kept + at enough cells for that *visible* window to still reach screen resolution. + ``cap`` bounds the request so a degenerate figure can't ask for a huge grid. + """ + fig = ax.figure + dpi = float(fig.dpi) + w_in, h_in = fig.get_size_inches() + try: # shrink to this axes' share of the figure when the layout is known + pos = ax.get_position() + w_in *= float(pos.width) + h_in *= float(pos.height) + except Exception: + pass + fr_rows = max(float(visible_frac[0]), 1e-3) + fr_cols = max(float(visible_frac[1]), 1e-3) + rows = int(np.ceil(h_in * dpi * oversample / fr_rows)) + cols = int(np.ceil(w_in * dpi * oversample / fr_cols)) + return min(max(rows, 1), cap), min(max(cols, 1), cap) + + +def _boxcar_downsample( + data: np.ndarray, + row_coord: np.ndarray, + col_coord: np.ndarray, + target_rows: int, + target_cols: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Boxcar-average a 2D array (and its 1D coords) toward a target shape. + + Reduces by an integer block factor per axis (``size // target``); an array + already at or below the target along an axis is left untouched on that axis, + so low-resolution data is never up-sampled or interpolated. The matching + coordinate centers are averaged over the same blocks, and any trailing cells + that don't fill a whole block are dropped. + """ + rows, cols = data.shape + fr = max(1, rows // max(target_rows, 1)) + fc = max(1, cols // max(target_cols, 1)) + if fr == 1 and fc == 1: + return data, row_coord, col_coord + nr = (rows // fr) * fr + nc = (cols // fc) * fc + reduced = data[:nr, :nc].reshape(nr // fr, fr, nc // fc, fc).mean(axis=(1, 3)) + r = row_coord[:nr].reshape(nr // fr, fr).mean(axis=1) + c = col_coord[:nc].reshape(nc // fc, fc).mean(axis=1) + return reduced, r, c + + # --- low-level plotters ---------------------------------------------------- @@ -74,10 +153,18 @@ def plot_spacetime( raise ValueError(f"plot_spacetime expects a 2D (t, x) array; got dims {da.dims}") if ax is None: _, ax = plt.subplots(figsize=(6, 4)) - data = np.log10(np.abs(da.values) + 1e-30) if log else da.values # (t, x) t = da.coords["t"].values xname = next(d for d in da.dims if d != "t") x = da.coords[xname].values + # Boxcar-average the (t, x) array down to ~the rendered pixel grid before + # rasterizing, so a 100M+-cell series doesn't build a billion-quad QuadMesh + # for sub-pixel detail. The log view averages |field| (averaging the *signed* + # field first would cancel the oscillations the log plot is meant to show). + target_rows, target_cols = _render_cells(ax) + data = np.abs(da.values) if log else da.values # (t, x) + data, t, x = _boxcar_downsample(data, t, x, target_rows, target_cols) + if log: + data = np.log10(data + 1e-30) if space_on_x: mesh = ax.pcolormesh(x, t, data, shading="auto", cmap=cmap) ax.set_xlabel(_axis_label(da, xname)) @@ -448,22 +535,48 @@ def plot_omega_k( nt = t.size nx = x.size - # 2-D FFT then shift so DC sits in the middle. - F = np.fft.fftshift(np.fft.fft2(da.values)) - P = np.abs(F) ** 2 + # Real-input 2-D FFT: rfft over the (long) time axis keeps only omega >= 0 + # and stays in the input's complex width (complex64 for float32) instead of + # upcasting the whole grid to complex128 the way fft2 would. We only want the + # power |F|^2, so the dropped, redundant negative-omega half is rebuilt by + # Hermitian symmetry *after* downsampling. Free the complex array as soon as + # the power is formed, and do the log in place, to bound peak memory. + F = np.fft.rfft2(da.values, axes=(1, 0)) # (nt//2+1, nx): rfft over t, fft over x + P = np.abs(F).astype(np.float32, copy=False) + del F + P **= 2 # |F|^2, in place + P = np.fft.fftshift(P, axes=1) # put k = 0 in the middle of the k (x) axis if log: - P = np.log10(P + 1e-30) + P += 1e-30 + np.log10(P, out=P) - omega = np.fft.fftshift(np.fft.fftfreq(nt, d=dt)) * 2 * np.pi - k = np.fft.fftshift(np.fft.fftfreq(nx, d=dx)) * 2 * np.pi + omega = np.fft.rfftfreq(nt, d=dt) * 2 * np.pi # [0, omega_Nyquist] + k = np.fft.fftshift(np.fft.fftfreq(nx, d=dx)) * 2 * np.pi # [-k_Ny, k_Ny) + + # Resolve the view window now (full Nyquist if unset, from the full-res axes) + # so the data can be decimated to the *visible* resolution: a zoomed panel + # keeps more of the full array than a full-range one (see _render_cells). + k_ny = float(np.max(np.abs(k))) + omega_ny = float(np.max(omega)) + if k_max is None: + k_max = k_ny + if omega_max is None: + omega_max = omega_ny + frac_rows = min(1.0, omega_max / omega_ny) if omega_ny > 0 else 1.0 + frac_cols = min(1.0, k_max / k_ny) if k_ny > 0 else 1.0 + target_rows, target_cols = _render_cells(ax, visible_frac=(frac_rows, frac_cols)) + P, omega, k = _boxcar_downsample(P, omega, k, target_rows, target_cols) + + # Rebuild the omega < 0 half from |F(k, -w)| = |F(-k, w)|: the negative rows + # are the omega > 0 block (excluding DC) flipped along both omega and k. + if P.shape[0] > 1: + P = np.concatenate([P[1:][::-1, ::-1], P], axis=0) + omega = np.concatenate([-omega[1:][::-1], omega]) mesh = ax.pcolormesh(k, omega, P, shading="auto", cmap=cmap) plt.colorbar(mesh, ax=ax, label=(r"$\log_{10}\,|\tilde{F}|^2$" if log else r"$|\tilde{F}|^2$")) + del P - if k_max is None: - k_max = float(np.max(np.abs(k))) - if omega_max is None: - omega_max = float(np.max(np.abs(omega))) ax.set_xlim(-k_max, k_max) ax.set_ylim(-omega_max, omega_max) @@ -911,7 +1024,9 @@ def _load(comp: str) -> xr.DataArray | None: return ser if ser.ndim == 2 else None def _pack(values: xr.DataArray, src: xr.DataArray, name: str, long: str) -> xr.DataArray: - out = values.copy() + # ``values`` is already a fresh array from the (e2 ± b3)/2 arithmetic, so + # relabel it in place rather than holding a second full (t, x) copy. + out = values out.attrs = dict(src.attrs) out.attrs["long_name"] = long out.name = name From c9a908817dc78cab576fd3b1a48bffa2b02d61c1 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 30 Jun 2026 11:27:03 -0700 Subject: [PATCH 26/49] Stream OSIRIS H5 -> NetCDF concurrently with the run Convert each diagnostic's MS/ HDF5 dumps into binary/*.nc as the run produces them, instead of one monolithic read-all/write-once pass at job end. This overlaps the (hours-long) PIC compute, reads each dump while still warm in the node page cache rather than re-opening ~700k tiny files cold off Lustre, and bounds conversion memory to a single dump instead of the whole stacked (t, x) series. Addresses the I/O wall-time and conversion-memory issues in osiris-lpi/postproc-performance.md (the pcolormesh/fft2 plot-time OOM is independent and unaffected). New adept/osiris/stream.py: - StreamWriter: append-only (t, ...) NetCDF writer using an unlimited t dim grown one slot per dump, so the file is sized to exactly the dumps produced (no deck-derived size guess, no trailing-fill trim) and a restart resumes after the on-disk slots. Schema matches io.series_to_dataset, so load_series_nc / plots / regen are unchanged. - convert_diagnostic_streaming: Stage A, memory-bounded conversion at job end. - StreamConverter: Stage B, a best-effort daemon thread that drains completed dumps live (processing dump N only once N+1 exists, to avoid partial reads); a final sweep picks up the last dump. RAW diagnostics (variable particle count) stay on the batch concat path. Fully failure-isolated: never aborts the OSIRIS run, and the batch path remains the safety net. Wiring: run_osiris spawns/finalizes the watcher around the existing subprocess (finalize before error handling so it can't mask a failure); base/post thread osiris.stream_convert + stream_poll_s through; save_run_datasets reuses the watcher's files and stream-builds any it missed. load_series_nc now flags a dim autoscaled only when its bounds actually move (the streamer always records them). On by default (osiris.stream_convert defaults true). Tests: tests/test_osiris/test_stream.py covers Stage-A equivalence vs load_series/batch, autoscale preservation, restart-resume, the N+1 hold-back rule, the live watcher thread, an end-to-end run_osiris run, RAW-skip, and the reuse paths. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/base.py | 2 + adept/osiris/io.py | 54 ++++- adept/osiris/post.py | 7 + adept/osiris/runner.py | 43 +++- adept/osiris/stream.py | 334 +++++++++++++++++++++++++++ docs/osiris-adept-usage.md | 52 +++++ tests/test_osiris/test_stream.py | 373 +++++++++++++++++++++++++++++++ 7 files changed, 855 insertions(+), 10 deletions(-) create mode 100644 adept/osiris/stream.py create mode 100644 tests/test_osiris/test_stream.py diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 7c6fd371..250c0756 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -153,6 +153,8 @@ def __call__(self, trainable_modules: dict, args: dict) -> dict: run_root=run_root, launcher=osiris_cfg.get("mpi_launcher", "srun"), extra_mpi_args=osiris_cfg.get("extra_mpi_args"), + stream_convert=bool(osiris_cfg.get("stream_convert", True)), + stream_poll_s=float(osiris_cfg.get("stream_poll_s", 10.0)), ) return {"solver result": result} diff --git a/adept/osiris/io.py b/adept/osiris/io.py index feb28627..0a12dfec 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -21,6 +21,7 @@ import json import re +import shutil from pathlib import Path import h5py @@ -467,10 +468,21 @@ def load_series_nc(path: str | Path) -> xr.DataArray: da.attrs["axis_units"] = axis_units da.attrs["axis_long_names"] = axis_long da.attrs.setdefault("time_units", da.attrs.get("time_units", r"1/\omega_p")) - # Rebuild the autoscaled-dim list from the per-dump bound coords that were - # written by :func:`series_to_dataset`, so a round-tripped series is - # indistinguishable (for plotting) from a fresh :func:`load_series`. - autoscaled = [str(d) for d in da.dims if f"{d}_min" in ds.coords and f"{d}_max" in ds.coords] + # Rebuild the autoscaled-dim list from the per-dump bound coords, so a + # round-tripped series is indistinguishable (for plotting) from a fresh + # :func:`load_series`. A dim is autoscaled only when its bounds actually + # *move*: the batch path writes {dim}_min/_max solely for moving dims, but + # the streaming writer (:mod:`adept.osiris.stream`) records them for every + # dim, so mere presence is not enough — check that they vary. + autoscaled: list[str] = [] + for d in da.dims: + lo_c, hi_c = f"{d}_min", f"{d}_max" + if lo_c not in ds.coords or hi_c not in ds.coords: + continue + lo = np.asarray(ds[lo_c].values) + hi = np.asarray(ds[hi_c].values) + if lo.size and not (np.allclose(lo, lo[0]) and np.allclose(hi, hi[0])): + autoscaled.append(str(d)) if autoscaled: da.attrs["autoscaled_dims"] = autoscaled return da @@ -508,6 +520,8 @@ def save_run_datasets( diagnostics: list[str] | set[str] | None = None, *, raw_drop_initial: bool = False, + stream: bool = False, + streamed_dir: str | Path | None = None, ) -> list[Path]: """Convert each diagnostic's full time history to a netCDF file. @@ -519,22 +533,46 @@ def save_run_datasets( ``diagnostics``, when given, whitelists which diagnostics to convert, matched against either the relative path (``"FLD/e1"``) or the leaf name (``"e1"``). Returns the list of written paths. + + Two opt-in paths bound conversion memory to a single dump (see + :mod:`adept.osiris.stream`): + + - ``streamed_dir`` — a directory where the concurrent converter already + wrote ``.nc`` *during the run*. Such a grid diagnostic is copied + into ``out_dir`` rather than rebuilt, so it is never re-read off disk. + - ``stream`` — for any grid diagnostic not already in ``streamed_dir``, + build the NetCDF a dump at a time instead of stacking the whole series in + memory. RAW diagnostics always use the (in-memory) concat path. """ out_dir = Path(out_dir) + streamed_dir = Path(streamed_dir) if streamed_dir is not None else None diags = list_diagnostics(run_dir) written: list[Path] = [] for relpath in sorted(diags): if diagnostics is not None and (relpath not in diagnostics and Path(relpath).name not in diagnostics): continue + dest = out_dir / f"{relpath}.nc" try: - if _diag_is_raw(relpath, diags[relpath]): + is_raw = _diag_is_raw(relpath, diags[relpath]) + pre = (streamed_dir / f"{relpath}.nc") if streamed_dir is not None else None + if pre is not None and not is_raw and pre.is_file(): + # The concurrent converter already produced this during the run. + dest.parent.mkdir(parents=True, exist_ok=True) + if pre.resolve() != dest.resolve(): + shutil.copy(pre, dest) + elif is_raw: # RAW (particle) dumps: per-particle datasets, no grid/AXIS. ds: xr.Dataset = load_raw_series(diags[relpath], drop_initial=raw_drop_initial) + dest.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(ds)) + elif stream: + from adept.osiris import stream as _stream # lazy: pulls in the writer only when needed + + _stream.convert_diagnostic_streaming(diags[relpath], dest, source_dir=diags[relpath]) else: ds = series_to_dataset(load_series(diags[relpath])) - dest = out_dir / f"{relpath}.nc" - dest.parent.mkdir(parents=True, exist_ok=True) - ds.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(ds)) + dest.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(ds)) written.append(dest) except Exception as e: # one bad diagnostic must not abort the rest print(f"[post] skipping diagnostic {relpath}: {e}") diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 23269336..8fbf7b0d 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -117,6 +117,11 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: td = Path(td) whitelist = (cfg.get("output") or {}).get("diagnostics_to_log") or None raw_drop_initial = bool((cfg.get("output") or {}).get("raw_drop_initial", False)) + # When the concurrent converter ran, reuse the NetCDFs it already produced + # (and stream-build any grid diagnostic it did not reach) instead of the + # in-memory batch conversion. + stream_on = bool((cfg.get("osiris") or {}).get("stream_convert", True)) + streamed_dir = solver.get("binary_dir") if stream_on else None metrics: dict[str, float] = { "wall_time_s": float(solver["wall_time"]), @@ -132,6 +137,8 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: td / "binary", diagnostics=whitelist, raw_drop_initial=raw_drop_initial, + stream=stream_on, + streamed_dir=streamed_dir, ) # plots imports matplotlib; do it lazily to keep `import adept.osiris` light. diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index e3896fd9..1113ce94 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -74,11 +74,23 @@ def run_osiris( env: dict[str, str] | None = None, launcher: str = "srun", extra_mpi_args: list[str] | None = None, + stream_convert: bool = True, + stream_poll_s: float = 10.0, ) -> dict[str, Any]: """Run OSIRIS and return run metadata. Returns a dict with keys ``run_dir`` (Path), ``exit_code`` (int), - ``wall_time`` (float, seconds), and ``cmd`` (list[str]). + ``wall_time`` (float, seconds), and ``cmd`` (list[str]). When + ``stream_convert`` is set it also returns ``binary_dir`` (str, the directory + the concurrent converter wrote into) and ``streamed_diagnostics`` + (sorted list[str] of relpaths fully converted while OSIRIS ran). + + With ``stream_convert=True`` a best-effort background thread + (:class:`adept.osiris.stream.StreamConverter`) converts grid diagnostics to + NetCDF *concurrently* with the run, polling every ``stream_poll_s`` seconds, + so the conversion overlaps compute and dumps are read while still warm in + cache. It is fully isolated: any failure there is logged and leaves the run + (and the batch post-processing fallback) untouched. Raises ``RuntimeError`` on non-zero exit code, with the last lines of stderr included in the message. @@ -107,6 +119,18 @@ def run_osiris( stdout_tail: list[str] = [] stderr_tail: list[str] = [] + # Best-effort concurrent H5 -> NetCDF converter, spawned alongside OSIRIS. + converter = None + binary_dir = run_dir / "binary" + if stream_convert: + try: + from adept.osiris import stream as _stream + + converter = _stream.StreamConverter(run_dir, binary_dir, poll_s=stream_poll_s) + except Exception as e: # never let the converter block the run + print(f"[stream] disabled (init failed): {e}") + converter = None + t0 = time.time() proc = subprocess.Popen( cmd, @@ -128,11 +152,22 @@ def run_osiris( ) t_out.start() t_err.start() + if converter is not None: + converter.start() rc = proc.wait() t_out.join() t_err.join() wall_time = time.time() - t0 + # Finalize the converter before any error handling below: the run is over, + # so flush whatever the watcher had not yet reached and close the files. + streamed: list[str] = [] + if converter is not None: + try: + streamed = sorted(converter.finalize()) + except Exception as e: + print(f"[stream] finalize failed (batch fallback will cover it): {e}") + if rc != 0: tail = "".join(stderr_tail[-50:]) or "(empty stderr)" raise RuntimeError( @@ -155,12 +190,16 @@ def run_osiris( f" stdout tail:\n{out_tail}" ) - return { + result: dict[str, Any] = { "run_dir": run_dir, "exit_code": rc, "wall_time": wall_time, "cmd": cmd, } + if stream_convert: + result["binary_dir"] = str(binary_dir) + result["streamed_diagnostics"] = streamed + return result def discover_binary(cfg_binary: str | None, *, dim: int | None = None) -> Path: diff --git a/adept/osiris/stream.py b/adept/osiris/stream.py new file mode 100644 index 00000000..9f44cd12 --- /dev/null +++ b/adept/osiris/stream.py @@ -0,0 +1,334 @@ +"""Incremental / concurrent HDF5 -> NetCDF conversion for OSIRIS runs. + +The batch converter in :mod:`adept.osiris.io` (:func:`io.save_run_datasets`) +reads a diagnostic's *entire* ``(t, ...)`` time history into memory and writes +it in one :func:`to_netcdf`. For the SRS field runs that is two problems: it +peaks at the whole stacked series in RAM, and it re-reads ~700k tiny HDF5 files +cold off Lustre at job end (see ``osiris-lpi/postproc-performance.md``). + +This module rearranges the same per-dump primitives into a *read-one -> +write-one-slot* loop: + +- **Stage A** (:func:`convert_diagnostic_streaming`): run the slot-writer at job + end. Peak memory drops from the full series to a single dump. No concurrency, + no write-completeness race. +- **Stage B** (:class:`StreamConverter`): a best-effort background thread that + drains new dumps into the NetCDF *while OSIRIS is still running*, so the + conversion overlaps the (hours-long) compute and each dump is read while still + warm in the node page cache. A final authoritative sweep at job end fills + whatever the watcher had not yet reached. + +Both write the *same* on-disk schema as :func:`io.series_to_dataset`, so the +``binary/*.nc`` contract is unchanged and every downstream reader +(:func:`io.load_series_nc`, the plotting code) works without modification. + +Design notes +------------ +- **Unlimited ``t`` dimension, append in iteration order.** OSIRIS writes dumps + sequentially and we process them sorted by iteration, so a plain append (slot + ``i`` = the ``i``-th dump) is correct and needs no stride arithmetic. The + dimension is grown one slot at a time with ``resize_dimension`` so it ends up + sized to *exactly* the dumps produced — no pre-sized guess from the deck + (whose off-by-one is genuinely ambiguous) and no trailing fill slots to trim + on early termination. This deviates from the fixed-dim sketch in the perf + doc; it is simpler and equally correct given in-order processing. +- **Write-completeness race.** A dump is only safe to read once OSIRIS has moved + on to the next one, so the watcher processes every dump *except the current + highest-numbered one*; the final sweep (run is dead) picks up that last dump. +- **Per-dump axis bounds** (``{dim}_min`` / ``{dim}_max``) are always recorded, + so OSIRIS autoscale (``if_ps_p_auto``) survives and a restart can recover them + from the file. :func:`io.load_series_nc` only *flags* a dim as autoscaled when + those bounds actually move, so the constant bounds carried for a fixed spatial + axis are harmless. +- **Failure isolation.** Nothing here may abort the OSIRIS run: the watcher + swallows and logs every exception, and the batch path in + :func:`io.save_run_datasets` remains the safety net for anything not produced. +""" + +from __future__ import annotations + +import threading +from pathlib import Path + +import h5netcdf +import numpy as np +import xarray as xr + +from adept.osiris import io as _io + +# Target uncompressed bytes per (t-chunk x spatial) chunk. A modest t-chunk +# keeps compression/read performance close to the batch path while bounding the +# read-modify-write cost of slot writes (which is hidden behind OSIRIS compute +# anyway). ~1 MiB is a reasonable HDF5 chunk size. +_TARGET_CHUNK_BYTES = 1 << 20 + + +def _default_chunk_t(spatial_shape: tuple[int, ...]) -> int: + """Pick a ``t``-chunk so each chunk is roughly :data:`_TARGET_CHUNK_BYTES`.""" + per_row = max(1, int(np.prod(spatial_shape))) * np.dtype(_io._DIAG_DTYPE).itemsize + return int(max(1, min(512, _TARGET_CHUNK_BYTES // per_row))) + + +def _data_var_attrs(template: xr.DataArray, source_dir) -> dict: + """Attrs for the diagnostic variable, matching :func:`io.series_to_dataset`. + + Starts from a single-dump :func:`io.load_grid_h5` DataArray and reproduces + exactly the attrs the batch path leaves on the data variable: the per-dump + ``time`` / ``iter`` / ``source`` are dropped, the dict-valued axis metadata + is lifted onto the coordinate variables instead, ``source_dir`` is added, and + everything left is coerced to a netCDF-writable scalar/array. + """ + attrs = dict(template.attrs) + for k in ("time", "iter", "source", "axis_units", "axis_long_names"): + attrs.pop(k, None) + attrs["source_dir"] = str(source_dir) + return {k: _io._coerce_attr(v) for k, v in attrs.items() if v is not None} + + +class StreamWriter: + """Append-only writer for one diagnostic's ``(t, ...)`` NetCDF time series. + + Created from the diagnostic's first dump (the schema template); each + :meth:`append` writes the next time slot, growing the unlimited ``t`` + dimension by one. Reopening an existing file (``mode="a"``) resumes after the + slots already on disk, which is what makes a restart idempotent. The file is + held open for the lifetime of the writer; call :meth:`close` to finalize. + """ + + def __init__(self, dest, template: xr.DataArray, *, source_dir=None, chunk_t: int | None = None, complevel: int = 4): + self.dest = Path(dest) + self.name = str(template.name) + self.spatial_dims = [str(d) for d in template.dims] + self.spatial_shape = tuple(int(s) for s in template.shape) + self._f: h5netcdf.File | None = None + + if self.dest.exists(): + # Resume: trust the existing schema, continue after written slots. + self._f = h5netcdf.File(self.dest, "a") + self._n = int(self._f.dimensions["t"].size) + else: + self.dest.parent.mkdir(parents=True, exist_ok=True) + self._f = h5netcdf.File(self.dest, "w") + self._create_schema(template, source_dir if source_dir is not None else self.dest.parent, chunk_t, complevel) + self._n = 0 + + @property + def n_written(self) -> int: + return self._n + + def _create_schema(self, template, source_dir, chunk_t, complevel) -> None: + f = self._f + axis_units = template.attrs.get("axis_units", {}) or {} + axis_long = template.attrs.get("axis_long_names", {}) or {} + + f.dimensions["t"] = None # unlimited; grown one slot per append + for d, sz in zip(self.spatial_dims, self.spatial_shape): + f.dimensions[d] = sz + + ct = chunk_t or _default_chunk_t(self.spatial_shape) + var = f.create_variable( + self.name, + ("t", *self.spatial_dims), + dtype=np.dtype(_io._DIAG_DTYPE), + chunks=(ct, *self.spatial_shape), + compression="gzip", + compression_opts=complevel, + shuffle=True, + ) + for k, v in _data_var_attrs(template, source_dir).items(): + var.attrs[k] = v + # List the non-dimension coordinate variables so xarray promotes them to + # coordinates (not data variables) on read — the CF convention the batch + # path's `to_netcdf` also emits. + nondim = ["iter"] + [f"{d}_{b}" for d in self.spatial_dims for b in ("min", "max")] + var.attrs["coordinates"] = " ".join(nondim) + + f.create_variable("t", ("t",), dtype="f8") + f.create_variable("iter", ("t",), dtype="i8") + for d in self.spatial_dims: + cv = f.create_variable(d, (d,), dtype="f8") + cv[:] = np.asarray(template.coords[d].values, dtype="f8") + if axis_units.get(d): + cv.attrs["units"] = axis_units[d] + if axis_long.get(d): + cv.attrs["long_name"] = axis_long[d] + f.create_variable(f"{d}_min", ("t",), dtype="f8") + f.create_variable(f"{d}_max", ("t",), dtype="f8") + + def append(self, da: xr.DataArray) -> None: + """Write ``da`` (one dump) into the next time slot.""" + if da.shape != self.spatial_shape: + raise ValueError(f"shape mismatch for {self.name}: {da.shape} != {self.spatial_shape}") + f = self._f + i = self._n + f.resize_dimension("t", i + 1) + f.variables[self.name][i, ...] = np.asarray(da.values, dtype=_io._DIAG_DTYPE) + f.variables["t"][i] = float(da.attrs.get("time", np.nan)) + f.variables["iter"][i] = int(da.attrs.get("iter", -1)) + for d in self.spatial_dims: + cv = np.asarray(da.coords[d].values) + lo, hi = (float(cv[0]), float(cv[-1])) if cv.size else (np.nan, np.nan) + f.variables[f"{d}_min"][i] = lo + f.variables[f"{d}_max"][i] = hi + self._n = i + 1 + + def flush(self) -> None: + if self._f is not None: + self._f.flush() + + def close(self) -> None: + if self._f is not None: + self._f.close() + self._f = None + + +def convert_diagnostic_streaming(diag_dir, dest, *, source_dir=None) -> Path: + """Stage A: convert one grid diagnostic to NetCDF a dump at a time. + + A memory-bounded drop-in for ``series_to_dataset(load_series(diag_dir))`` + followed by ``to_netcdf`` — peak RAM is one dump, not the whole stacked + series. Rebuilds ``dest`` from scratch (any partial file is removed first). + """ + diag_dir = Path(diag_dir) + dest = Path(dest) + dumps = _io._sort_dumps(diag_dir) + if not dumps: + raise FileNotFoundError(f"No .h5 dumps in {diag_dir}") + if dest.exists(): + dest.unlink() + + template = _io.load_grid_h5(dumps[0]) + writer = StreamWriter(dest, template, source_dir=source_dir or diag_dir) + try: + writer.append(template) + for p in dumps[1:]: + writer.append(_io.load_grid_h5(p)) + finally: + writer.close() + return dest + + +class StreamConverter: + """Stage B: a background thread that converts grid diagnostics during the run. + + Spawned by :func:`adept.osiris.runner.run_osiris` alongside the OSIRIS + subprocess, it polls ``run_dir/MS`` and drains each grid diagnostic's + completed dumps into ``out_dir/.nc``. RAW (particle) diagnostics do + not fit the fixed-slot model (variable particle count per dump) and are left + to the batch path. Every operation is best-effort: any error is logged and + the run is never affected. + """ + + def __init__(self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print): + self.ms = Path(run_dir) / "MS" + self.out_dir = Path(out_dir) + self.poll_s = float(poll_s) + self._log = logger + self._writers: dict[str, StreamWriter] = {} + self._completed: set[str] = set() + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + # --- lifecycle -------------------------------------------------------- + + def start(self) -> None: + self._thread = threading.Thread(target=self._run, name="osiris-stream-convert", daemon=True) + self._thread.start() + + def _run(self) -> None: + while not self._stop.is_set(): + try: + self._scan_once(final=False) + except Exception as e: # a watcher must never crash the run + self._log(f"[stream] scan error (continuing): {e}") + self._stop.wait(self.poll_s) + + def finalize(self) -> set[str]: + """Stop the thread, run the authoritative final sweep, close all writers. + + After OSIRIS has exited every dump is safe to read (no writer can still + be racing), so the final sweep processes the last dump each diagnostic + had been holding back and closes the files. Returns the set of + diagnostic relpaths fully converted by the stream (so the caller can + skip rebuilding them in the batch pass). + """ + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=max(self.poll_s, 30.0) + 30.0) + try: + self._scan_once(final=True) + except Exception as e: + self._log(f"[stream] finalize sweep error (batch fallback covers it): {e}") + for rel, w in self._writers.items(): + try: + w.close() + except Exception: + pass + self._completed.add(rel) + return set(self._completed) + + # --- scanning --------------------------------------------------------- + + def _grid_diags(self) -> dict[str, Path]: + """Grid (non-RAW) diagnostics under ``MS/``, keyed by relpath.""" + out: dict[str, Path] = {} + if not self.ms.is_dir(): + return out + for d in self.ms.rglob("*"): + if not d.is_dir(): + continue + try: + has_h5 = any(c.suffix == ".h5" for c in d.iterdir()) + except OSError: + continue + if not has_h5: + continue + rel = str(d.relative_to(self.ms)) + if _io._diag_is_raw(rel, d): + continue + out[rel] = d + return out + + def _scan_once(self, *, final: bool) -> None: + for rel, d in self._grid_diags().items(): + if rel in self._completed: + continue + try: + self._drain(rel, d, final=final) + except Exception as e: + self._log(f"[stream] {rel}: {e} (will fall back to batch)") + # Drop the writer so a half-written file is rebuilt by the batch + # pass rather than reused. + w = self._writers.pop(rel, None) + if w is not None: + try: + w.close() + except Exception: + pass + + def _drain(self, rel: str, d: Path, *, final: bool) -> None: + dumps = _io._sort_dumps(d) + if not dumps: + return + # Write-completeness: a dump is safe only once the next one exists (OSIRIS + # has moved on). The final sweep runs after OSIRIS is dead, so the last + # dump is safe too. + safe = dumps if final else dumps[:-1] + if not safe: + return # nothing safe to write yet (only the in-progress dump exists) + + w = self._writers.get(rel) + if w is None: + template = _io.load_grid_h5(dumps[0]) + w = StreamWriter(self.out_dir / f"{rel}.nc", template, source_dir=d) + self._writers[rel] = w + if w.n_written == 0: + w.append(template) # reuse the template we just loaded as slot 0 + + for p in safe[w.n_written :]: + w.append(_io.load_grid_h5(p)) + w.flush() + + if final: + w.close() + self._completed.add(rel) diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index c3743be7..16b2bbc3 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -113,6 +113,12 @@ osiris: # purge policy will take care of stale ones eventually. extra_mpi_args: ["--oversubscribe"] # optional, passed to mpirun + stream_convert: true # optional (default true): convert MS/ + # HDF5 dumps to binary/*.nc concurrently + # with the run; set false for the old + # batch conversion at job end + stream_poll_s: 10.0 # optional: watcher poll interval (s) + density: # optional: adaptive box sizing (1D) gradient_scale_length: 300um # target L_n; scales the box so the # deck's density ramp realizes this L @@ -183,6 +189,52 @@ preserved. | Artifacts | `binary//.nc` — one xarray netCDF per diagnostic, holding the full `(t, …)` time history (replaces the raw h5 dumps) | | Artifacts | `plots/…` — canned PNGs (see below) | +## 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 `/binary/.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 `t` dimension 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's `binary/*.nc` contract 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.collect` + reuses 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/.nc` reopened 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. + ## Canned plots (`plots/` artifacts) These canned plots focus on 1D simulations. diff --git a/tests/test_osiris/test_stream.py b/tests/test_osiris/test_stream.py new file mode 100644 index 00000000..92e35316 --- /dev/null +++ b/tests/test_osiris/test_stream.py @@ -0,0 +1,373 @@ +"""Incremental / concurrent HDF5 -> NetCDF conversion (adept.osiris.stream). + +These tests synthesize a tiny OSIRIS-shaped run (no real run on disk) and check +that the streaming converter — both the at-job-end Stage A path and the +concurrent Stage B watcher — produces a NetCDF that is equivalent, for every +downstream consumer, to the in-memory batch path in ``io.save_run_datasets``. +""" + +from __future__ import annotations + +import stat +import time +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr + +from adept.osiris import io as oio +from adept.osiris import post as opost +from adept.osiris import runner as orunner +from adept.osiris import stream as ostream + + +def _write_dump(path: Path, name: str, data: np.ndarray, t: float, it: int, axes) -> None: + """Write one OSIRIS-style HDF5 grid dump. + + ``axes`` is a list of ``(name, long_name, units, min, max)`` in OSIRIS + (Fortran) order, i.e. the last entry is the fastest-varying numpy axis. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + f.attrs["NAME"] = np.array([name.encode()], dtype="S256") + f.attrs["LABEL"] = np.array([name.encode()], dtype="S256") + f.attrs["UNITS"] = np.array([b"a.u."], dtype="S256") + f.attrs["TIME UNITS"] = np.array([rb"1 / \omega_p"], dtype="S256") + ax = f.create_group("AXIS") + for i, (an, ln, un, lo, hi) in enumerate(axes, start=1): + d = ax.create_dataset(f"AXIS{i}", data=np.array([lo, hi], dtype="float64")) + d.attrs["NAME"] = np.array([an.encode()], dtype="S256") + d.attrs["LONG_NAME"] = np.array([ln.encode()], dtype="S256") + d.attrs["UNITS"] = np.array([un.encode()], dtype="S256") + sim = f.create_group("SIMULATION") + sim.attrs["DT"] = np.array([0.05]) + sim.attrs["NDIMS"] = np.array([len(axes)], dtype="int32") + sim.attrs["NX"] = np.array(list(reversed(data.shape)), dtype="int32") + sim.attrs["XMIN"] = np.array([a[3] for a in axes]) + sim.attrs["XMAX"] = np.array([a[4] for a in axes]) + f.create_dataset(name, data=data.astype("float32")) + + +def _write_field(run_dir: Path, comp: str, n_steps: int, nx: int, seed: int = 0) -> Path: + """Write an ``MS/FLD/`` 1-D field diagnostic over ``n_steps`` dumps.""" + rng = np.random.default_rng(seed) + d = run_dir / "MS" / "FLD" / comp + for k in range(n_steps): + _write_dump( + d / f"{comp}-{k * 10:06d}.h5", + comp, + rng.standard_normal(nx), + t=k * 0.5, + it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + return d + + +def _write_phasespace_autoscaled(run_dir: Path, n_steps: int, nx: int, npmom: int) -> Path: + """Phase space whose momentum bounds grow each dump (if_ps_p_auto).""" + rng = np.random.default_rng(7) + d = run_dir / "MS" / "PHA" / "p1x1" / "electrons" + for k in range(n_steps): + p_hi = 1.0 + 0.5 * k + _write_dump( + d / f"p1x1-electrons-{k * 10:06d}.h5", + "p1x1", + rng.standard_normal((npmom, nx)), + t=k * 0.5, + it=k * 10, + axes=[ + ("x1", "x_1", r"c / \omega_p", 0.0, 10.0), + ("p1", "p_1", r"m_e c", -p_hi, p_hi), + ], + ) + return d + + +def _assert_series_equiv(got: xr.DataArray, ref: xr.DataArray) -> None: + """``got`` (streamed) is interchangeable with ``ref`` (raw load_series).""" + assert got.name == ref.name + assert got.dims == ref.dims + assert got.shape == ref.shape + np.testing.assert_array_equal(got.values, ref.values) # float32 -> float32, exact + np.testing.assert_allclose(got["t"].values, ref["t"].values) + assert list(got["iter"].values) == list(ref["iter"].values) + for dim in ref.dims: + if dim == "t": + continue + np.testing.assert_allclose(oio.physical_axis(got, dim), oio.physical_axis(ref, dim)) + assert got.attrs.get("axis_units") == ref.attrs.get("axis_units") + assert got.attrs.get("autoscaled_dims") == ref.attrs.get("autoscaled_dims") + + +# --- Stage A: streaming finalize ----------------------------------------- + + +def test_streaming_grid_matches_load_series(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=5, nx=8) + + dest = tmp_path / "stream" / "e1.nc" + ostream.convert_diagnostic_streaming(diag, dest) + + streamed = oio.load_series_nc(dest) + raw = oio.load_series(diag) + _assert_series_equiv(streamed, raw) + # A fixed spatial axis is not autoscaled even though the streamer always + # records its (constant) bounds. + assert "autoscaled_dims" not in streamed.attrs + + +def test_streaming_matches_batch_file(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=6, nx=12, seed=3) + + batch = oio.save_run_datasets(run_dir, tmp_path / "batch")[0] + stream = tmp_path / "stream" / "e1.nc" + ostream.convert_diagnostic_streaming(diag, stream) + + a = oio.load_series_nc(stream) + b = oio.load_series_nc(batch) + np.testing.assert_array_equal(a.values, b.values) + np.testing.assert_allclose(a["t"].values, b["t"].values) + assert list(a["iter"].values) == list(b["iter"].values) + assert a.attrs.get("long_name") == b.attrs.get("long_name") + assert a.attrs.get("units") == b.attrs.get("units") + + +def test_streaming_preserves_autoscale_bounds(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_phasespace_autoscaled(run_dir, n_steps=4, nx=8, npmom=6) + + dest = tmp_path / "stream" / "p1x1.nc" + ostream.convert_diagnostic_streaming(diag, dest) + + streamed = oio.load_series_nc(dest) + raw = oio.load_series(diag) + _assert_series_equiv(streamed, raw) + assert streamed.attrs.get("autoscaled_dims") == ["p1"] + np.testing.assert_allclose( + oio.physical_axis(streamed, "p1", it=-1), + np.linspace(-2.5, 2.5, streamed.sizes["p1"]), + ) + + +def test_streaming_compresses_output(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=8, nx=64) + dest = tmp_path / "stream" / "e1.nc" + ostream.convert_diagnostic_streaming(diag, dest) + enc = xr.open_dataset(dest, engine="h5netcdf")["e1"].encoding + assert enc.get("zlib") or (enc.get("compression") == "gzip") + + +# --- restart / resume ----------------------------------------------------- + + +def test_writer_resumes_after_reopen(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=6, nx=8) + dumps = oio._sort_dumps(diag) + dest = tmp_path / "stream" / "e1.nc" + + # First "segment": write the first 3 dumps, then close (simulates a crash + # after a checkpoint). + template = oio.load_grid_h5(dumps[0]) + w = ostream.StreamWriter(dest, template, source_dir=diag) + w.append(template) + for p in dumps[1:3]: + w.append(oio.load_grid_h5(p)) + assert w.n_written == 3 + w.close() + + # Restart: reopen and continue from where it left off. + w2 = ostream.StreamWriter(dest, oio.load_grid_h5(dumps[0]), source_dir=diag) + assert w2.n_written == 3 # picked up the on-disk slots + for p in dumps[3:]: + w2.append(oio.load_grid_h5(p)) + w2.close() + + _assert_series_equiv(oio.load_series_nc(dest), oio.load_series(diag)) + + +# --- Stage B: concurrent watcher ----------------------------------------- + + +def test_watcher_holds_back_in_progress_dump(tmp_path: Path) -> None: + """A non-final scan must not read the highest-numbered (still-writing) dump.""" + run_dir = tmp_path / "run" + _write_field(run_dir, "e1", n_steps=5, nx=8) + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + conv._scan_once(final=False) + # 5 dumps present -> 4 are "safe", the newest is held back. + ds = xr.open_dataset(tmp_path / "binary" / "FLD" / "e1.nc", engine="h5netcdf") + try: + assert ds.sizes["t"] == 4 + finally: + ds.close() + + completed = conv.finalize() # picks up the last dump and closes + assert "FLD/e1" in completed + _assert_series_equiv( + oio.load_series_nc(tmp_path / "binary" / "FLD" / "e1.nc"), + oio.load_series(run_dir / "MS" / "FLD" / "e1"), + ) + + +def test_watcher_thread_streams_dumps_as_they_appear(tmp_path: Path) -> None: + """End-to-end Stage B: dumps written over time while the thread polls.""" + run_dir = tmp_path / "run" + diag = run_dir / "MS" / "FLD" / "e1" + rng = np.random.default_rng(0) + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.02) + conv.start() + n_steps = 6 + for k in range(n_steps): + _write_dump( + diag / f"e1-{k * 10:06d}.h5", + "e1", + rng.standard_normal(8), + t=k * 0.5, + it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + time.sleep(0.03) + + completed = conv.finalize() + assert "FLD/e1" in completed + streamed = oio.load_series_nc(tmp_path / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == n_steps + _assert_series_equiv(streamed, oio.load_series(diag)) + + +def test_watcher_skips_raw_diagnostics(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + _write_field(run_dir, "e1", n_steps=3, nx=8) + # A RAW diagnostic: no AXIS group, several per-particle datasets. + raw = run_dir / "MS" / "RAW" / "species_1" + raw.mkdir(parents=True) + with h5py.File(raw / "RAW-species_1-000000.h5", "w") as f: + f.attrs["TIME"] = np.array([0.0]) + f.attrs["ITER"] = np.array([0], dtype="int32") + for q in ("p1", "x1"): + f.create_dataset(q, data=np.arange(4, dtype="float32")) + f.create_group("SIMULATION").attrs["NDIMS"] = np.array([1], dtype="int32") + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + completed = conv.finalize() + assert "FLD/e1" in completed + assert not any("RAW" in c for c in completed) + assert not (tmp_path / "binary" / "RAW" / "species_1.nc").exists() + + +# --- integration with save_run_datasets / post.collect ------------------- + + +def test_save_run_datasets_reuses_streamed_file(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=4, nx=8) + + # Pretend the watcher already produced this during the run, and tag it so we + # can tell a reuse (copy) from a rebuild. + streamed_dir = tmp_path / "run" / "binary" + ostream.convert_diagnostic_streaming(diag, streamed_dir / "FLD" / "e1.nc") + with xr.open_dataset(streamed_dir / "FLD" / "e1.nc", engine="h5netcdf") as ds: + ds = ds.load() + ds.attrs["_marker"] = "from-watcher" + ds.to_netcdf(streamed_dir / "FLD" / "e1.nc", engine="h5netcdf") + + out = tmp_path / "td" / "binary" + written = oio.save_run_datasets(run_dir, out, stream=True, streamed_dir=streamed_dir) + assert (out / "FLD" / "e1.nc") in written + with xr.open_dataset(out / "FLD" / "e1.nc", engine="h5netcdf") as ds2: + assert ds2.attrs.get("_marker") == "from-watcher" # copied, not rebuilt + + +def test_save_run_datasets_stream_builds_missing(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=5, nx=8) + # stream=True but no streamed_dir -> build via the streaming writer. + out = tmp_path / "binary" + oio.save_run_datasets(run_dir, out, stream=True) + _assert_series_equiv(oio.load_series_nc(out / "FLD" / "e1.nc"), oio.load_series(diag)) + + +def test_run_osiris_streams_concurrently(tmp_path: Path) -> None: + """run_osiris spawns the watcher alongside the (fake) OSIRIS subprocess. + + The fake binary is a shell script that drops pre-staged dumps into + ``MS/FLD/e1`` one at a time with a pause between — exactly the cadence the + watcher polls. No Python/h5py is needed in the subprocess. + """ + n_steps = 5 + stage = tmp_path / "stage" + stage.mkdir() + rng = np.random.default_rng(1) + ref = {} + for k in range(n_steps): + fn = f"e1-{k * 10:06d}.h5" + data = rng.standard_normal(8) + ref[fn] = data + _write_dump(stage / fn, "e1", data, t=k * 0.5, it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)]) + + lines = ["#!/bin/sh", "set -e", "mkdir -p MS/FLD/e1"] + for k in range(n_steps): + fn = f"e1-{k * 10:06d}.h5" + lines.append(f"cp '{stage / fn}' MS/FLD/e1/{fn}") + lines.append("sleep 0.1") + fake = tmp_path / "fake-osiris.sh" + fake.write_text("\n".join(lines) + "\n") + fake.chmod(fake.stat().st_mode | stat.S_IEXEC | stat.S_IRUSR) + + result = orunner.run_osiris( + "node_conf\n{\n}\n", + binary=fake, + mpi_ranks=1, + run_root=tmp_path / "checkpoints", + stream_convert=True, + stream_poll_s=0.05, + ) + + assert result["exit_code"] == 0 + assert result["streamed_diagnostics"] == ["FLD/e1"] + nc = Path(result["binary_dir"]) / "FLD" / "e1.nc" + assert nc.exists() + streamed = oio.load_series_nc(nc) + assert streamed.sizes["t"] == n_steps + _assert_series_equiv(streamed, oio.load_series(Path(result["run_dir"]) / "MS" / "FLD" / "e1")) + + +def test_collect_uses_streamed_binary(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=3, nx=8) + (run_dir / "os-stdin").write_text("node_conf\n{\n}\n") + (run_dir / "stdout.log").write_text("ok\n") + (run_dir / "stderr.log").write_text("") + # Simulate runner having streamed into run_dir/binary during the run. + binary_dir = run_dir / "binary" + ostream.convert_diagnostic_streaming(diag, binary_dir / "FLD" / "e1.nc") + + run_output = { + "solver result": { + "run_dir": str(run_dir), + "wall_time": 1.0, + "exit_code": 0, + "binary_dir": str(binary_dir), + "streamed_diagnostics": ["FLD/e1"], + } + } + cfg = {"osiris": {"deck": str(run_dir / "os-stdin"), "stream_convert": True}, "output": {}} + + td = tmp_path / "td" + td.mkdir() + opost.collect(run_output, cfg, str(td)) + + assert not list(td.rglob("*.h5")) + assert (td / "binary" / "FLD" / "e1.nc").exists() + _assert_series_equiv(oio.load_series_nc(td / "binary" / "FLD" / "e1.nc"), oio.load_series(diag)) From 05c6b34558a21205d74501247b5896c6cc1521e8 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 30 Jun 2026 17:00:47 -0700 Subject: [PATCH 27/49] Add /dev/shm ramdisk staging for OSIRIS diagnostic dumps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OSIRIS has no asynchronous diagnostic I/O: every dump is a synchronous barrier in the time loop, so on Lustre the ranks stall on each collective write (~halving utilization at full dump cadence). `osiris.stage_root` works around it as poor-man's async I/O: OSIRIS writes to a node-local /dev/shm ramdisk (RAM-speed, no stall) and a background drainer mirrors each completed dump to the durable run_root and reaps it from the ramdisk, so the RAM footprint stays bounded by the poll interval rather than the whole run. - stream.py: StreamConverter gains a `persist_dir` drain mode — mirrors grid AND raw dumps to persist_dir/MS, reaps the scratch, still streams grid .nc. persist_dir=None is byte-for-byte the old in-place behavior. - runner.py: `stage_root` runs OSIRIS on the ramdisk, returns the durable dir as run_dir (post.py unchanged), syncs stragglers + reclaims the ramdisk at job end. Also ignore srun's spurious "couldn't chdir ... going to /tmp instead" launcher warning, which was false-positiving the exit-0 error guard. - base.py: wire osiris.stage_root through. - docs/osiris-adept-usage.md + tests. Validated on Perlmutter (1024 ppc, nmin0.12, tmax=600): CPU 16-rank 250.7->207.0s (-17%), GPU 1xA100 128.8->123.7s (-4%); /dev/shm held bounded while 3617 dumps drained to durable storage, full SRS post-processing intact. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/base.py | 1 + adept/osiris/runner.py | 84 +++++++++++++++++-- adept/osiris/stream.py | 135 ++++++++++++++++++++++++++----- docs/osiris-adept-usage.md | 56 +++++++++++++ tests/test_osiris/test_runner.py | 10 +++ tests/test_osiris/test_stream.py | 102 +++++++++++++++++++++++ 6 files changed, 365 insertions(+), 23 deletions(-) diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 250c0756..b762079f 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -155,6 +155,7 @@ def __call__(self, trainable_modules: dict, args: dict) -> dict: extra_mpi_args=osiris_cfg.get("extra_mpi_args"), stream_convert=bool(osiris_cfg.get("stream_convert", True)), stream_poll_s=float(osiris_cfg.get("stream_poll_s", 10.0)), + stage_root=osiris_cfg.get("stage_root"), ) return {"solver result": result} diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index 1113ce94..a7a5853b 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -27,6 +27,13 @@ _OSIRIS_STDERR_NOISE = ( "No protocol specified", "MPI_INIT", # benign MPI banner noise on some setups + # srun launcher warning, not an OSIRIS error: when a task's working directory + # (e.g. a per-node /dev/shm staging dir) is not resolvable at launch, srun + # prints "error: couldn't chdir to `...`: ... going to /tmp instead". On + # Perlmutter this fires spuriously even when the tasks do chdir correctly, so + # it must not abort an otherwise-clean run. A genuine staging failure surfaces + # downstream instead (no dumps to drain -> empty/short post-processing). + "couldn't chdir to", ) @@ -65,6 +72,26 @@ def _make_run_dir(run_root: Path) -> Path: return rd +def _sync_tree(src: Path, dst: Path) -> None: + """Copy every file under ``src`` to the same relative path under ``dst``, + skipping any that already exist. + + Used at the end of a staged run to fold whatever the background drainer did + not move (``HIST/``, ``RE/``, ``TIMINGS/``, any straggler ``MS/`` dump) from + the ephemeral scratch into the durable run directory before the scratch is + reclaimed. Existing targets (the streamed ``binary/`` NetCDFs, the dumps the + drainer already mirrored) are left untouched. + """ + for p in src.rglob("*"): + if not p.is_file(): + continue + target = dst / p.relative_to(src) + if target.exists(): + continue + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(p, target) + + def run_osiris( deck_text: str, *, @@ -76,6 +103,7 @@ def run_osiris( extra_mpi_args: list[str] | None = None, stream_convert: bool = True, stream_poll_s: float = 10.0, + stage_root: str | Path | None = None, ) -> dict[str, Any]: """Run OSIRIS and return run metadata. @@ -92,6 +120,17 @@ def run_osiris( cache. It is fully isolated: any failure there is logged and leaves the run (and the batch post-processing fallback) untouched. + **Ramdisk staging (``stage_root``).** When ``stage_root`` is set (and + ``stream_convert`` is on), OSIRIS runs with its working directory on that + fast ephemeral filesystem (e.g. ``/dev/shm``) so its *synchronous* dump + writes hit RAM instead of stalling on the parallel filesystem — effectively + the asynchronous diagnostic I/O OSIRIS itself lacks. The background drainer + mirrors each completed dump to the durable run directory under ``run_root`` + and deletes the scratch copy, keeping the ramdisk footprint bounded by the + poll interval rather than the whole run. ``run_dir`` in the returned dict is + always the durable directory, so post-processing is unchanged; the scratch + is reclaimed before returning. Single-node only (``/dev/shm`` is per-node). + Raises ``RuntimeError`` on non-zero exit code, with the last lines of stderr included in the message. """ @@ -99,8 +138,23 @@ def run_osiris( if not binary.exists(): raise FileNotFoundError(f"OSIRIS binary not found: {binary}") + # The durable run directory always lives under run_root. With staging, OSIRIS + # works on a same-named scratch dir under stage_root and the drainer mirrors + # back here; without it, OSIRIS works in run_dir directly (the original path). run_dir = _make_run_dir(Path(run_root).expanduser().resolve()) - (run_dir / "os-stdin").write_text(deck_text) + staging = bool(stage_root) and stream_convert + if stage_root and not stream_convert: + print("[stream] stage_root ignored: ramdisk staging requires stream_convert=True") + if staging: + stage_dir: Path | None = Path(stage_root).expanduser().resolve() / run_dir.name + stage_dir.mkdir(parents=True, exist_ok=True) + work_dir = stage_dir + # Keep a durable copy of the deck alongside the mirrored outputs. + (run_dir / "os-stdin").write_text(deck_text) + else: + stage_dir = None + work_dir = run_dir + (work_dir / "os-stdin").write_text(deck_text) if mpi_ranks > 1: cmd = [launcher, "-n", str(mpi_ranks)] @@ -120,13 +174,21 @@ def run_osiris( stderr_tail: list[str] = [] # Best-effort concurrent H5 -> NetCDF converter, spawned alongside OSIRIS. + # It reads OSIRIS's dumps from the working dir and writes NetCDFs into the + # durable binary/ dir; in staging mode it also mirrors+reaps the scratch + # into run_dir (persist_dir). converter = None binary_dir = run_dir / "binary" if stream_convert: try: from adept.osiris import stream as _stream - converter = _stream.StreamConverter(run_dir, binary_dir, poll_s=stream_poll_s) + converter = _stream.StreamConverter( + work_dir, + binary_dir, + poll_s=stream_poll_s, + persist_dir=run_dir if staging else None, + ) except Exception as e: # never let the converter block the run print(f"[stream] disabled (init failed): {e}") converter = None @@ -134,7 +196,7 @@ def run_osiris( t0 = time.time() proc = subprocess.Popen( cmd, - cwd=run_dir, + cwd=work_dir, env=merged_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, @@ -168,10 +230,20 @@ def run_osiris( except Exception as e: print(f"[stream] finalize failed (batch fallback will cover it): {e}") + # Staging: fold anything the drainer did not move (HIST/, RE/, straggler MS/ + # dumps) from the scratch into the durable run dir, then reclaim the ramdisk. + # Done before error handling so a failed run's partial outputs still persist. + if staging and stage_dir is not None: + try: + _sync_tree(stage_dir, run_dir) + except Exception as e: + print(f"[stream] final stage->persist sync issue (partial outputs may be lost): {e}") + shutil.rmtree(stage_dir, ignore_errors=True) + if rc != 0: tail = "".join(stderr_tail[-50:]) or "(empty stderr)" raise RuntimeError( - f"OSIRIS exited with status {rc}.\n cmd: {shlex.join(cmd)}\n cwd: {run_dir}\n stderr tail:\n{tail}" + f"OSIRIS exited with status {rc}.\n cmd: {shlex.join(cmd)}\n cwd: {work_dir}\n stderr tail:\n{tail}" ) # OSIRIS can exit 0 even on input-file errors: it prints something @@ -185,7 +257,7 @@ def run_osiris( raise RuntimeError( "OSIRIS reported an error despite exit-code 0:\n" f" cmd: {shlex.join(cmd)}\n" - f" cwd: {run_dir}\n" + f" cwd: {work_dir}\n" f" stderr tail:\n{err_tail}\n" f" stdout tail:\n{out_tail}" ) @@ -199,6 +271,8 @@ def run_osiris( if stream_convert: result["binary_dir"] = str(binary_dir) result["streamed_diagnostics"] = streamed + if staging: + result["staged"] = True return result diff --git a/adept/osiris/stream.py b/adept/osiris/stream.py index 9f44cd12..de484881 100644 --- a/adept/osiris/stream.py +++ b/adept/osiris/stream.py @@ -47,6 +47,8 @@ from __future__ import annotations +import os +import shutil import threading from pathlib import Path @@ -209,7 +211,7 @@ def convert_diagnostic_streaming(diag_dir, dest, *, source_dir=None) -> Path: class StreamConverter: - """Stage B: a background thread that converts grid diagnostics during the run. + """Stage B: a background thread that drains diagnostics during the run. Spawned by :func:`adept.osiris.runner.run_osiris` alongside the OSIRIS subprocess, it polls ``run_dir/MS`` and drains each grid diagnostic's @@ -217,15 +219,33 @@ class StreamConverter: not fit the fixed-slot model (variable particle count per dump) and are left to the batch path. Every operation is best-effort: any error is logged and the run is never affected. + + **Staging mode (``persist_dir`` set).** When OSIRIS writes its ``MS/`` dumps + to a fast ephemeral scratch (e.g. a ``/dev/shm`` ramdisk passed as + ``run_dir``), the converter also *drains* the scratch: after each completed + dump is handled it is mirrored to ``persist_dir/MS//`` and + **deleted from the scratch**, so the ramdisk high-water mark is bounded by + the poll interval rather than the whole run. Grid diagnostics are mirrored + *and* streamed to ``out_dir`` (which the caller points at durable storage); + RAW diagnostics, which cannot be slot-streamed, are mirrored as HDF5 so the + batch post-processing path finds them on ``persist_dir``. The net durable + layout under ``persist_dir`` is identical to a non-staged run, so + post-processing reads it unchanged. With ``persist_dir=None`` nothing is + mirrored or deleted — the converter only streams grid diagnostics in place, + exactly as before. """ - def __init__(self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print): + def __init__(self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print, persist_dir=None): self.ms = Path(run_dir) / "MS" self.out_dir = Path(out_dir) + self.persist_dir = Path(persist_dir) if persist_dir is not None else None + self.persist_ms = (self.persist_dir / "MS") if self.persist_dir is not None else None self.poll_s = float(poll_s) self._log = logger self._writers: dict[str, StreamWriter] = {} - self._completed: set[str] = set() + self._completed: set[str] = set() # grid relpaths fully converted to .nc + self._raw_completed: set[str] = set() # raw relpaths fully mirrored to persist + self._is_raw: dict[str, bool] = {} # rel -> raw? (cached; avoids re-peeking h5) self._stop = threading.Event() self._thread: threading.Thread | None = None @@ -248,9 +268,11 @@ def finalize(self) -> set[str]: After OSIRIS has exited every dump is safe to read (no writer can still be racing), so the final sweep processes the last dump each diagnostic - had been holding back and closes the files. Returns the set of - diagnostic relpaths fully converted by the stream (so the caller can - skip rebuilding them in the batch pass). + had been holding back (and, in staging mode, drains it off the scratch) + before closing the files. Returns the set of *grid* diagnostic relpaths + fully converted by the stream (so the caller can skip rebuilding them in + the batch pass); RAW relpaths are mirrored but not returned, since the + batch pass still builds their NetCDFs from the mirrored HDF5. """ self._stop.set() if self._thread is not None: @@ -269,8 +291,8 @@ def finalize(self) -> set[str]: # --- scanning --------------------------------------------------------- - def _grid_diags(self) -> dict[str, Path]: - """Grid (non-RAW) diagnostics under ``MS/``, keyed by relpath.""" + def _diags(self) -> dict[str, Path]: + """Every diagnostic dir under ``MS/`` (grid and RAW), keyed by relpath.""" out: dict[str, Path] = {} if not self.ms.is_dir(): return out @@ -281,20 +303,32 @@ def _grid_diags(self) -> dict[str, Path]: has_h5 = any(c.suffix == ".h5" for c in d.iterdir()) except OSError: continue - if not has_h5: - continue - rel = str(d.relative_to(self.ms)) - if _io._diag_is_raw(rel, d): - continue - out[rel] = d + if has_h5: + out[str(d.relative_to(self.ms))] = d return out + def _diag_is_raw(self, rel: str, d: Path) -> bool: + cached = self._is_raw.get(rel) + if cached is None: + cached = _io._diag_is_raw(rel, d) + self._is_raw[rel] = cached + return cached + def _scan_once(self, *, final: bool) -> None: - for rel, d in self._grid_diags().items(): - if rel in self._completed: + for rel, d in self._diags().items(): + if rel in self._completed or rel in self._raw_completed: continue try: - self._drain(rel, d, final=final) + if self._diag_is_raw(rel, d): + # RAW: nothing to do unless staging (then mirror + reap the + # HDF5 so the scratch stays bounded and the batch path finds + # the dumps on persist). + if self.persist_dir is not None: + self._drain_raw(rel, d, final=final) + elif self.persist_dir is None: + self._drain_grid_inplace(rel, d, final=final) + else: + self._drain_grid_staged(rel, d, final=final) except Exception as e: self._log(f"[stream] {rel}: {e} (will fall back to batch)") # Drop the writer so a half-written file is rebuilt by the batch @@ -306,7 +340,15 @@ def _scan_once(self, *, final: bool) -> None: except Exception: pass - def _drain(self, rel: str, d: Path, *, final: bool) -> None: + # --- draining --------------------------------------------------------- + + def _drain_grid_inplace(self, rel: str, d: Path, *, final: bool) -> None: + """Stream a grid diagnostic to NetCDF in place — no mirror, no reap. + + The original non-staged behavior: dumps accumulate in ``run_dir`` and + slot ``i`` is the ``i``-th dump, indexed positionally against the full + (never-pruned) dump list. + """ dumps = _io._sort_dumps(d) if not dumps: return @@ -332,3 +374,60 @@ def _drain(self, rel: str, d: Path, *, final: bool) -> None: if final: w.close() self._completed.add(rel) + + def _drain_grid_staged(self, rel: str, d: Path, *, final: bool) -> None: + """Stream a grid diagnostic *and* drain its dumps off the scratch. + + Because every handled dump is deleted from the scratch (:meth:`_reap`), + the directory only ever holds dumps not yet processed (plus the + in-progress last one), so every ``safe`` dump is new — no positional + bookkeeping against on-disk slots is needed. + """ + dumps = _io._sort_dumps(d) + if not dumps: + return + safe = dumps if final else dumps[:-1] + if not safe: + return + + w = self._writers.get(rel) + if w is None: + w = StreamWriter(self.out_dir / f"{rel}.nc", _io.load_grid_h5(dumps[0]), source_dir=d) + self._writers[rel] = w + for p in safe: + w.append(_io.load_grid_h5(p)) + self._reap(rel, p) + w.flush() + + if final: + w.close() + self._completed.add(rel) + + def _drain_raw(self, rel: str, d: Path, *, final: bool) -> None: + """Mirror a RAW diagnostic's completed dumps to persist + reap scratch.""" + dumps = _io._sort_dumps(d) + if not dumps: + return + safe = dumps if final else dumps[:-1] + for p in safe: + self._reap(rel, p) + if final: + self._raw_completed.add(rel) + + def _reap(self, rel: str, p: Path) -> None: + """Mirror one scratch dump to ``persist_dir/MS`` then delete the scratch + copy. Best-effort: on any failure the scratch file is left in place for + the runner's final sync to catch, so data is never lost and the watcher + never crashes.""" + if self.persist_dir is None: + return + try: + dest = self.persist_ms / rel / p.name + if not dest.exists(): + dest.parent.mkdir(parents=True, exist_ok=True) + tmp = dest.with_name(dest.name + ".part") + shutil.copyfile(p, tmp) + os.replace(tmp, dest) # atomic publish on the persist filesystem + p.unlink() + except Exception as e: + self._log(f"[stream] reap {rel}/{p.name} deferred: {e}") diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index 16b2bbc3..af3fedab 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -118,6 +118,11 @@ osiris: # with the run; set false for the old # batch conversion at job end stream_poll_s: 10.0 # optional: watcher poll interval (s) + # stage_root: /dev/shm/osiris # optional: run OSIRIS on this fast + # ephemeral filesystem (ramdisk) and + # drain dumps to run_root in the + # background — see "Ramdisk staging". + # Requires stream_convert: true. density: # optional: adaptive box sizing (1D) gradient_scale_length: 300um # target L_n; scales the box so the @@ -235,6 +240,57 @@ 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. + +```yaml +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/`; all `MS/` + dumps land in RAM. The durable run directory under `run_root` is created as + usual, and `run_dir` in the result (hence everything post-processing reads) is + always that **durable** directory — so `post.collect` is unchanged. +- The `StreamConverter` runs in **drain mode**: each completed dump is mirrored + to `run_root//MS//` 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 durable + `binary/*.nc` as before; RAW (particle) diagnostics, which can't be + slot-streamed, are mirrored as HDF5 so the batch path still finds them. +- `binary/*.nc` and `stdout.log`/`stderr.log` are 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 `tmpfs` counts 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, keep + `stream_poll_s` small and leave headroom. +- **Crash window.** Whatever is still on the ramdisk when a node dies is lost + (`tmpfs` is 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/shm` is 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. diff --git a/tests/test_osiris/test_runner.py b/tests/test_osiris/test_runner.py index 70f77d61..7373e476 100644 --- a/tests/test_osiris/test_runner.py +++ b/tests/test_osiris/test_runner.py @@ -34,6 +34,16 @@ def test_discover_binary_missing_raises() -> None: runner.discover_binary("/no/such/path/exists", dim=1) +def test_srun_chdir_warning_is_not_an_error() -> None: + # srun's spurious launcher warning under /dev/shm staging must not be flagged + # as an OSIRIS error (it aborts an otherwise-clean run); a real OSIRIS error + # and a real abort still are. + chdir = "[2026-06-30T23:35:54] error: couldn't chdir to `/dev/shm/x`: No such file or directory: going to /tmp instead" + assert runner._looks_like_osiris_error(chdir) is False + assert runner._looks_like_osiris_error("(*error*) Lindman not yet allowed with tiling") is True + assert runner._looks_like_osiris_error("Error reading global simulation parameters, aborting...") is True + + def test_run_osiris_missing_binary_raises(tmp_path: Path) -> None: with pytest.raises(FileNotFoundError): runner.run_osiris( diff --git a/tests/test_osiris/test_stream.py b/tests/test_osiris/test_stream.py index 92e35316..38be8c78 100644 --- a/tests/test_osiris/test_stream.py +++ b/tests/test_osiris/test_stream.py @@ -343,6 +343,108 @@ def test_run_osiris_streams_concurrently(tmp_path: Path) -> None: _assert_series_equiv(streamed, oio.load_series(Path(result["run_dir"]) / "MS" / "FLD" / "e1")) +def _write_raw_dump(path: Path, t: float, it: int, npart: int = 4) -> None: + """Write one OSIRIS-style RAW (particle) dump: per-particle datasets, no AXIS.""" + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + for q in ("p1", "x1"): + f.create_dataset(q, data=np.arange(npart, dtype="float32")) + f.create_group("SIMULATION").attrs["NDIMS"] = np.array([1], dtype="int32") + + +# --- Staging mode: ramdisk scratch + drain to durable persist ------------- + + +def test_staged_converter_mirrors_and_reaps(tmp_path: Path) -> None: + """With ``persist_dir`` set, every completed dump is mirrored to the durable + tree and deleted from the scratch; grid is streamed, raw is mirrored as H5.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=4, nx=8) + for k in range(3): + _write_raw_dump(stage / "MS" / "RAW" / "species_1" / f"RAW-species_1-{k * 10:06d}.h5", t=k * 0.5, it=k * 10) + + conv = ostream.StreamConverter(stage, persist / "binary", poll_s=0.01, persist_dir=persist) + completed = conv.finalize() + + # Grid streamed to the durable binary/ tree; raw is mirrored but NOT returned + # as "streamed" (the batch pass still builds its NetCDF from the H5). + assert completed == {"FLD/e1"} + assert (persist / "binary" / "FLD" / "e1.nc").exists() + + # Every dump (grid + raw) mirrored to the durable MS/ tree... + assert len(list((persist / "MS" / "FLD" / "e1").glob("*.h5"))) == 4 + assert len(list((persist / "MS" / "RAW" / "species_1").glob("*.h5"))) == 3 + # ...and reaped from the scratch (RAM freed). + assert not list((stage / "MS").rglob("*.h5")) + + # The streamed grid series is faithful to the mirrored dumps. + _assert_series_equiv( + oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc"), + oio.load_series(persist / "MS" / "FLD" / "e1"), + ) + + +def test_non_staged_converter_never_deletes(tmp_path: Path) -> None: + """Without ``persist_dir`` the converter must leave the source dumps in place + (the durable run dir *is* run_dir) — no reaping of the in-place tree.""" + run_dir = tmp_path / "run" + _write_field(run_dir, "e1", n_steps=4, nx=8) + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + conv.finalize() + assert len(list((run_dir / "MS" / "FLD" / "e1").glob("*.h5"))) == 4 + assert not (tmp_path / "MS").exists() # nothing mirrored anywhere + + +def test_run_osiris_staging_mirrors_and_frees_scratch(tmp_path: Path) -> None: + """End-to-end: OSIRIS works on the ramdisk, the drainer mirrors to the + durable run dir, and the scratch is reclaimed before returning.""" + n_steps = 5 + src = tmp_path / "src" + src.mkdir() + rng = np.random.default_rng(2) + for k in range(n_steps): + _write_dump(src / f"e1-{k * 10:06d}.h5", "e1", rng.standard_normal(8), t=k * 0.5, it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)]) + + lines = ["#!/bin/sh", "set -e", "mkdir -p MS/FLD/e1"] + for k in range(n_steps): + fn = f"e1-{k * 10:06d}.h5" + lines.append(f"cp '{src / fn}' MS/FLD/e1/{fn}") + lines.append("sleep 0.1") + fake = tmp_path / "fake-osiris.sh" + fake.write_text("\n".join(lines) + "\n") + fake.chmod(fake.stat().st_mode | stat.S_IEXEC | stat.S_IRUSR) + + ramdisk = tmp_path / "ramdisk" + result = orunner.run_osiris( + "node_conf\n{\n}\n", + binary=fake, + mpi_ranks=1, + run_root=tmp_path / "checkpoints", + stream_convert=True, + stream_poll_s=0.05, + stage_root=ramdisk, + ) + + assert result["exit_code"] == 0 + assert result.get("staged") is True + run_dir = Path(result["run_dir"]) + # Durable run dir lives under run_root, and the scratch was reclaimed. + assert run_dir.parent == (tmp_path / "checkpoints").resolve() + assert not (ramdisk.resolve() / run_dir.name).exists() + # Outputs landed durably: streamed .nc, mirrored MS dumps, deck + logs. + assert (run_dir / "binary" / "FLD" / "e1.nc").exists() + assert len(list((run_dir / "MS" / "FLD" / "e1").glob("*.h5"))) == n_steps + assert (run_dir / "os-stdin").exists() + assert (run_dir / "stdout.log").exists() + streamed = oio.load_series_nc(run_dir / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == n_steps + _assert_series_equiv(streamed, oio.load_series(run_dir / "MS" / "FLD" / "e1")) + + def test_collect_uses_streamed_binary(tmp_path: Path) -> None: run_dir = tmp_path / "run" diag = _write_field(run_dir, "e1", n_steps=3, nx=8) From 4ee1c4aeadc33136012ae3978d9ed71b9d0edc9b Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 30 Jun 2026 17:05:38 -0700 Subject: [PATCH 28/49] osiris: slab-restricted transverse-field split for boundary diagnostics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add transverse_field_boundary_slabs(), a memory-bounded companion to efield_lr_components for the boundary-light diagnostics (laser budget, direction-split spectra/spectrogram) that only sample a thin slab at each box edge. It loads each raw transverse field one at a time, slices it to the two edge slabs, and frees the whole-grid array before the next — so the full (t, x) left/right-going grid (tens of GiB per field on the long, wide SRS runs) is never materialized. The Riemann split is local, so slice-then-combine equals combine-then-slice (locked by test_boundary_slabs_match_full_split). Also export render_cells / boxcar_downsample as public names so downstream spectrograms can decimate to on-screen resolution before pcolormesh the way plot_omega_k already does. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/plots.py | 97 +++++++++++++++++++++++ tests/test_osiris/test_plots_new_views.py | 31 ++++++++ 2 files changed, 128 insertions(+) diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index a7d43c83..86584013 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -1048,6 +1048,99 @@ def _pack(values: xr.DataArray, src: xr.DataArray, name: str, long: str) -> xr.D return out +def transverse_field_boundary_slabs( + run_dir: str | Path, *, guard_cells: int = 1, window_cells: int = 3 +) -> dict | None: + r"""Left/right-going transverse E-field on just the two boundary slabs. + + The slab-restricted companion to :func:`efield_lr_components`. Every + boundary-light diagnostic (laser-energy budget, direction-split spectra and + spectrogram) only samples a thin ``window_cells`` slab ``guard_cells`` in + from each edge, yet the full split holds several whole-grid ``(t, x)`` arrays + at once — prohibitive for the long, wide SRS runs (tens of GiB per field). + This loads each raw transverse field **one at a time** and immediately slices + it to the left slab ``[g, g+w)`` and the symmetric right slab + ``[n-g-w, n-g)``, freeing the whole-grid array before the next, so only the + slab columns are ever combined. + + The vacuum Riemann split is local (per cell), so slicing first and combining + on the slab is identical to combining the whole grid and then slicing. + Returns ``None`` when no transverse pair (``e2/b3`` or ``e3/b2``) was dumped, + else:: + + {"t": (n_t,), "guard_cells": g, "window_cells": w, "pairs": [...], + "edges": {"left": {pair: {"right": (t, w), "left": (t, w)}, ...}, + "right": {pair: {"right": (t, w), "left": (t, w)}, ...}}} + + where each ``(t, w)`` array is the right/left-going part of that polarization + pair on that edge's slab (native field dtype). The consumer forms the + per-channel flux / spectrum from these (the raw, un-split field is + ``right + left``). + """ + diags = _io.list_diagnostics(run_dir) + + def _load(comp: str) -> xr.DataArray | None: + rel = f"FLD/{comp}" + if rel not in diags: + return None + ser = _io.load_series(diags[rel]) + return ser if ser.ndim == 2 else None + + # Resolve the slab columns from the first available transverse field, then + # drop it (each field is reloaded one at a time below so only one whole-grid + # array is ever resident). + probe = None + for comp in ("e2", "b3", "e3", "b2"): + probe = _load(comp) + if probe is not None: + break + if probe is None: + return None + xdim = next(d for d in probe.dims if d != "t") + t = np.asarray(probe.coords["t"].values, dtype=float) + n = int(probe.sizes[xdim]) + g, w = guard_cells, window_cells + if g + w > n: + g, w = 0, min(w, n) + left, right = slice(g, g + w), slice(n - g - w, n - g) + del probe + + def _slabs(comp: str) -> tuple[np.ndarray, np.ndarray] | None: + """Raw field ``comp`` sliced to the (left, right) edge slabs. + + ``.copy()`` so the returned slabs do not keep the whole-grid array alive + as a view base — the (t, x) field is freed when this returns. + """ + da = _load(comp) + if da is None: + return None + lo = da.isel({xdim: left}).values.copy() + hi = da.isel({xdim: right}).values.copy() + return lo, hi + + edges: dict[str, dict[str, dict[str, np.ndarray]]] = {"left": {}, "right": {}} + + def _add_pair(name: str, e: str, b: str, *, right_sign: float) -> None: + # right-going = (e + right_sign*b)/2, left-going = (e - right_sign*b)/2 + # (e2/b3: right_sign=+1; e3/b2: right_sign=-1), matching efield_lr_components. + es, bs = _slabs(e), _slabs(b) + if es is None or bs is None: + return + for edge, ei, bi in (("left", es[0], bs[0]), ("right", es[1], bs[1])): + edges[edge][name] = { + "right": (ei + right_sign * bi) / 2.0, + "left": (ei - right_sign * bi) / 2.0, + } + + _add_pair("e2", "e2", "b3", right_sign=+1.0) + _add_pair("e3", "e3", "b2", right_sign=-1.0) + + pairs = list(edges["left"].keys()) + if not pairs: + return None + return {"t": t, "guard_cells": g, "window_cells": w, "pairs": pairs, "edges": edges} + + def plot_field_lr_decomposition(run_dir: str | Path) -> dict[str, plt.Figure]: """One figure per transverse component: right/left-going spacetime. @@ -1457,3 +1550,7 @@ def _field_energy_from_series(ser: xr.DataArray) -> np.ndarray: tex = _tex sim_box_xmax = _sim_box_xmax sim_box_bound = _sim_box_bound +# Shared memory-bounded rendering helpers (the (k, ω) map decimates with these +# before pcolormesh; downstream spectrograms reuse them to do the same). +render_cells = _render_cells +boxcar_downsample = _boxcar_downsample diff --git a/tests/test_osiris/test_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py index bf25861d..ed412d1a 100644 --- a/tests/test_osiris/test_plots_new_views.py +++ b/tests/test_osiris/test_plots_new_views.py @@ -205,6 +205,37 @@ def test_efield_decomposition_isolates_right_going(tmp_path: Path) -> None: np.testing.assert_allclose(parts["e3"]["left"].values, 0.0, atol=1e-6) +def test_boundary_slabs_match_full_split(tmp_path: Path) -> None: + # The slab-restricted split must equal the full split sliced to the same + # boundary columns (the Riemann split is local, so slice-then-combine == + # combine-then-slice). This is what lets the boundary diagnostics avoid + # materializing the whole (t, x) left/right grid. + run_dir = _make_rich_run(tmp_path) + g, w = 1, 3 + full = oplt.efield_lr_components(run_dir) + sl = oplt.transverse_field_boundary_slabs(run_dir, guard_cells=g, window_cells=w) + assert sl is not None and sl["pairs"] == list(full) + assert sl["guard_cells"] == g and sl["window_cells"] == w + sample = next(iter(full.values()))["right"] + xdim = next(d for d in sample.dims if d != "t") + n = sample.sizes[xdim] + left, right = slice(g, g + w), slice(n - g - w, n - g) + for pair, parts in full.items(): + for edge, sel in (("left", left), ("right", right)): + for going in ("right", "left"): + np.testing.assert_allclose( + sl["edges"][edge][pair][going], + parts[going].isel({xdim: sel}).values, + atol=1e-6, + ) + + +def test_boundary_slabs_absent_pair_returns_none(tmp_path: Path) -> None: + run_dir = tmp_path / "run" # only a density diag, no transverse FLD pairs + _write_dump(run_dir / "MS/DENSITY/ion/charge" / "charge-ion-000000.h5", "charge", np.ones(8), 0.0, 0, [X_AX]) + assert oplt.transverse_field_boundary_slabs(run_dir) is None + + def test_field_decomposition_figures(tmp_path: Path) -> None: run_dir = _make_rich_run(tmp_path) figs = oplt.plot_field_lr_decomposition(run_dir) From a31bb6d2c510c6354c12865718dd0185ac752970 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 1 Jul 2026 14:20:50 -0700 Subject: [PATCH 29/49] Added energy plots --- adept/osiris/io.py | 10 +++-- adept/osiris/plots.py | 49 ++++++++++++++++++++- tests/test_osiris/test_diagnostics_plots.py | 1 + 3 files changed, 56 insertions(+), 4 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 0a12dfec..57adc47a 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -671,11 +671,15 @@ def load_hist_energy(run_dir: str | Path) -> xr.Dataset | None: if parsed is None or parsed[1].size == 0: continue t, values = parsed - total_per_row = values.sum(axis=1) if "fld" in f.name.lower(): - field = (t, total_per_row) + # fld_ene columns are the per-component field energies + # (B1 B2 B3 E1 E2 E3); their sum is the total EM field energy. + field = (t, values.sum(axis=1)) elif "par" in f.name.lower(): - kinetic[f.stem] = (t, total_per_row) + # par_ene columns are "Total Par." (particle COUNT) then + # "Kin. Energy" — take the energy (last) column. Summing would fold + # the ~1e6 particle count into the energy and swamp it. + kinetic[f.stem] = (t, values[:, -1]) if field is None and not kinetic: return None diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index 86584013..dffe4804 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -448,7 +448,7 @@ def plot_energy_components( } for name, label in labels.items(): if name in ds: - ax.plot(t, ds[name].values, label=label) + ax.plot(t, ds[name].values, label=label, alpha=0.5) if log: ax.set_yscale("log") ax.set_xlabel(r"t [$1/\omega_p$]") @@ -495,6 +495,44 @@ def plot_energy_conservation( return ax +def plot_energy_partition( + energy: xr.Dataset, + ax: plt.Axes | None = None, + *, + log: bool = True, + title: str | None = None, +) -> plt.Axes: + """Total energy broken down into particle (kinetic) and EM (field) parts. + + Same overlay style as :func:`plot_energy_components`, but the partition is + the physical total energy rather than the E/B field split. ``energy`` is the + ``Dataset`` from :func:`io.load_hist_energy`, which carries ``field_energy`` + (EM), ``kinetic_total`` (all-species particle), and ``total`` (their sum). + + ``total`` is drawn first (underneath) so the two components stay legible + where one of them coincides with the total envelope. + """ + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + t = energy["t"].values + labels = { + "total": "total (particle + EM)", + "kinetic_total": "particle (kinetic)", + "field_energy": "EM (field)", + } + for name, label in labels.items(): + if name in energy: + ax.plot(t, energy[name].values, label=label, alpha=0.5) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel("energy [code units]") + ax.set_title(title or "Total energy partition vs time") + ax.legend() + ax.grid(True, which="both", alpha=0.3) + return ax + + def plot_omega_k( series: xr.DataArray | str | Path, ax: plt.Axes | None = None, @@ -1412,6 +1450,15 @@ def _write(fig: plt.Figure, rel: str) -> Path: except Exception as e: print(f"[plots] skipping total_energy_vs_time: {e}") + try: + energy = _io.load_hist_energy(run_dir) + if energy is not None and "total" in energy: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_energy_partition(energy, ax=ax) + written["energy_partition_vs_time"] = _write(fig, "energy_partition_vs_time.png") + except Exception as e: + print(f"[plots] skipping energy_partition_vs_time: {e}") + return written diff --git a/tests/test_osiris/test_diagnostics_plots.py b/tests/test_osiris/test_diagnostics_plots.py index c4084d0d..c6bd78dd 100644 --- a/tests/test_osiris/test_diagnostics_plots.py +++ b/tests/test_osiris/test_diagnostics_plots.py @@ -153,6 +153,7 @@ def test_save_canned_plots_emits_full_set(tmp_path: Path) -> None: "energy_vs_time", "energy_components_vs_time", "total_energy_vs_time", # present because HIST/ supplies kinetic energy + "energy_partition_vs_time", # same HIST source: total split into particle/EM } assert expected <= set(written) for path in written.values(): From a7a42bb6b951368e5bc4495fd20176af7a498710 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 1 Jul 2026 15:29:08 -0700 Subject: [PATCH 30/49] osiris: 2D-ready field energy + EPW (longitudinal e1) energy trace Prepares the general OSIRIS diagnostics for 2D decks that dump only spatially-averaged fields plus Poynting lineouts (e.g. R-Follett): - field_energy_components resolves -savg/-tavg field variants (_resolve_fld_diag), integrates over all spatial dims so 2D (t,x2,x1) series work, and emits per-component e1/e2/e3 energies. - plot_epw_energy: the longitudinal (e1) field energy = electron plasma wave energy, added to save_canned_plots as epw_energy_vs_time.png (works for 1D and 2D). - _total_field_energy sums only real field components (_is_field_component_dir), excluding s1 Poynting lineouts / slices that also live under FLD/; the per-dump reader falls back to the sole dataset when a savg dump names it differently. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/plots.py | 130 +++++++++++++++----- adept/osiris/post.py | 26 +++- tests/test_osiris/test_diagnostics_plots.py | 49 +++++++- 3 files changed, 173 insertions(+), 32 deletions(-) diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index dffe4804..2d37cddc 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -26,6 +26,7 @@ field_decomp/.png left/right-going transverse E (Riemann split) energy_vs_time.png total field energy time-trace energy_components_vs_time.png E-field / B-field / total energy + epw_energy_vs_time.png longitudinal (e1) EPW field energy total_energy_vs_time.png field + kinetic conservation (needs HIST/) All axis / colorbar / title labels are emitted as proper LaTeX (``$\omega$``, @@ -368,67 +369,132 @@ def plot_energy_vs_time( return ax +def _resolve_fld_diag(diags: dict[str, Path], comp: str) -> str | None: + """Diagnostic key for field component ``comp``: full-grid, else -savg/-tavg. + + A 2D deck that dumps ``e1,savg`` (rather than the full grid) writes it to a + distinct ``FLD/`` diagnostic (e.g. ``FLD/e1-savg``), so ``FLD/e1`` is absent. + Fall back to a spatially/time-averaged variant of the same component — but + never a lineout / slice, which is not a whole ``(t, x…)`` field. + """ + if f"FLD/{comp}" in diags: + return f"FLD/{comp}" + for rel in sorted(diags): + parts = rel.split("/") + if len(parts) == 2 and parts[0] == "FLD": + leaf = parts[1] + if leaf.split("-", 1)[0] == comp and "line" not in leaf and "slice" not in leaf: + return rel + return None + + def field_energy_components(run_dir: str | Path) -> xr.Dataset: """E-field, B-field, and total field energy vs time from the FLD diagnostics. Like :func:`field_energy_series` but keeps the electric (``e1/e2/e3``) and magnetic (``b1/b2/b3``) contributions separate. Returns a ``Dataset`` with - data variables ``E_energy``, ``B_energy`` and ``total_field_energy`` on a - shared ``t`` axis. Components not dumped are simply omitted from their sum. + ``E_energy``, ``B_energy`` and ``total_field_energy`` on a shared ``t`` axis, + plus the per-component electric energies ``e1_energy`` / ``e2_energy`` / + ``e3_energy`` — ``e1_energy`` is the longitudinal (electrostatic) field + energy, i.e. the electron-plasma-wave (EPW) energy for a run whose density + gradient lies along ``x1``. Components not dumped are omitted from their sum. + + The spatial integral is over *all* spatial dims, so this works for a 1D + ``(t, x)`` field series and a 2D ``(t, x2, x1)`` one alike; 2D decks that dump + only spatially-averaged fields (``e1,savg``) are picked up via + :func:`_resolve_fld_diag`. Sources are resolved via :func:`io.list_diagnostics` / :func:`io.load_series`, so ``run_dir`` may be a raw OSIRIS run directory or the ``binary/`` directory - of saved NetCDFs — the energy is recomputed from the stored ``(t, x)`` field - series either way. + of saved NetCDFs — the energy is recomputed from the stored field series + either way. """ diags = _io.list_diagnostics(run_dir) - groups = {"E_energy": ("e1", "e2", "e3"), "B_energy": ("b1", "b2", "b3")} + comps = ("e1", "e2", "e3", "b1", "b2", "b3") by_iter: dict[int, dict[str, float]] = {} times: dict[int, float] = {} found = False - for group, comps in groups.items(): - for comp in comps: - rel = f"FLD/{comp}" - if rel not in diags: - continue - try: - ser = _io.load_series(diags[rel]) - except Exception as e: # a bad component must not sink the rest - print(f"[plots] skipping field-energy component {comp}: {e}") - continue - if ser.ndim != 2: - continue - found = True - e_t = _field_energy_from_series(ser) - its = np.asarray(ser.coords["iter"].values) if "iter" in ser.coords else np.arange(ser.sizes["t"]) - ts = np.asarray(ser.coords["t"].values, dtype="float64") - for k in range(ts.size): - it = int(its[k]) - rec = by_iter.setdefault(it, {"E_energy": 0.0, "B_energy": 0.0}) - rec[group] += float(e_t[k]) - times.setdefault(it, float(ts[k])) + for comp in comps: + rel = _resolve_fld_diag(diags, comp) + if rel is None: + continue + try: + ser = _io.load_series(diags[rel]) + except Exception as e: # a bad component must not sink the rest + print(f"[plots] skipping field-energy component {comp}: {e}") + continue + if ser.ndim < 2 or "t" not in ser.dims: + continue + found = True + e_t = _field_energy_from_series(ser) + its = np.asarray(ser.coords["iter"].values) if "iter" in ser.coords else np.arange(ser.sizes["t"]) + ts = np.asarray(ser.coords["t"].values, dtype="float64") + for k in range(ts.size): + it = int(its[k]) + rec = by_iter.setdefault(it, {}) + rec[comp] = rec.get(comp, 0.0) + float(e_t[k]) + times.setdefault(it, float(ts[k])) if not found or not by_iter: raise RuntimeError(f"No field dumps available under {run_dir}") iters = sorted(by_iter) t = np.array([times[i] for i in iters]) - e_arr = np.array([by_iter[i]["E_energy"] for i in iters]) - b_arr = np.array([by_iter[i]["B_energy"] for i in iters]) + + def col(c: str) -> np.ndarray: + return np.array([by_iter[i].get(c, 0.0) for i in iters]) + + e1a, e2a, e3a = col("e1"), col("e2"), col("e3") + e_arr = e1a + e2a + e3a + b_arr = col("b1") + col("b2") + col("b3") coords = {"t": t, "iter": ("t", np.asarray(iters))} ds = xr.Dataset( { "E_energy": ("t", e_arr), "B_energy": ("t", b_arr), "total_field_energy": ("t", e_arr + b_arr), + "e1_energy": ("t", e1a), + "e2_energy": ("t", e2a), + "e3_energy": ("t", e3a), }, coords=coords, ) ds["E_energy"].attrs.update(long_name="electric field energy", units="code") ds["B_energy"].attrs.update(long_name="magnetic field energy", units="code") ds["total_field_energy"].attrs.update(long_name="total field energy", units="code") + ds["e1_energy"].attrs.update(long_name="longitudinal (EPW) field energy", units="code") + ds["e2_energy"].attrs.update(long_name="transverse E2 field energy", units="code") + ds["e3_energy"].attrs.update(long_name="transverse E3 field energy", units="code") ds["t"].attrs.update(long_name="time", units=r"1/\omega_p") return ds +def plot_epw_energy( + src: xr.Dataset | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = True, + title: str | None = None, +) -> plt.Axes: + r"""Plot the longitudinal (EPW) field energy ``½ ∫ e1² dV`` vs time. + + The electron-plasma-wave energy trace: the ``e1_energy`` component of + :func:`field_energy_components` (the electrostatic field along ``x1``). ``src`` + may be a precomputed ``Dataset`` or a ``run_dir`` path. + """ + ds = src if isinstance(src, xr.Dataset) else field_energy_components(src) + if "e1_energy" not in ds: + raise RuntimeError("no e1 (longitudinal) field diagnostic for the EPW energy") + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + ax.plot(ds["t"].values, ds["e1_energy"].values) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel(r"EPW energy $\frac{1}{2}\int e_1^2\,d^Dx$ [code units]") + ax.set_title(title or "Total EPW (longitudinal field) energy vs time") + ax.grid(True, which="both", alpha=0.3) + return ax + + def plot_energy_components( src: xr.Dataset | str | Path, ax: plt.Axes | None = None, @@ -1441,6 +1507,13 @@ def _write(fig: plt.Figure, rel: str) -> Path: except (FileNotFoundError, RuntimeError) as e: print(f"[plots] skipping energy_components_vs_time: {e}") + try: # total EPW (longitudinal e1) field energy vs time + fig, ax = plt.subplots(figsize=(6, 4)) + plot_epw_energy(run_dir, ax=ax) + written["epw_energy_vs_time"] = _write(fig, "epw_energy_vs_time.png") + except (FileNotFoundError, RuntimeError) as e: + print(f"[plots] skipping epw_energy_vs_time: {e}") + try: energy = _io.load_hist_energy(run_dir) if energy is not None and "total" in energy: @@ -1594,6 +1667,7 @@ def _field_energy_from_series(ser: xr.DataArray) -> np.ndarray: ensure_series = _ensure_series axis_label = _axis_label display_name = _display_name +value_label = _value_label tex = _tex sim_box_xmax = _sim_box_xmax sim_box_bound = _sim_box_bound diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 8fbf7b0d..0535d06e 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -39,6 +39,21 @@ def keyfn(p: Path) -> int: return max(files, key=keyfn) +_FIELD_COMPONENTS = ("e1", "e2", "e3", "b1", "b2", "b3") + + +def _is_field_component_dir(name: str) -> bool: + """True for a full-grid or spatially/time-averaged field-component dump dir. + + Matches ``e1`` .. ``b3`` and their ``-savg`` / ``-tavg`` variants (2D decks + dump ``e1,savg`` rather than the full grid), but NOT Poynting-flux (``s1``) + or lineout / slice diagnostics, which also live under ``FLD/`` and would + otherwise be summed into — and inflate — the total field energy. + """ + base = name.split("-", 1)[0] + return base in _FIELD_COMPONENTS and "line" not in name and "slice" not in name + + def _field_energy_from_dump(h5_path: Path) -> float: """Integrate ``field^2 * cell_volume`` for a single dump file. @@ -48,10 +63,15 @@ def _field_energy_from_dump(h5_path: Path) -> float: import h5py # local import keeps adept import light with h5py.File(h5_path, "r") as f: - # The dataset name matches the diagnostic name (e.g. "e1"). + # The dataset name usually matches the diagnostic name (e.g. "e1"), but + # spatially/time-averaged reports may name it differently; fall back to + # the single non-AXIS/SIMULATION dataset the way load_grid_h5 does. name = h5_path.stem.rsplit("-", 1)[0] if name not in f: - return float("nan") + data_keys = [k for k in f.keys() if k not in ("AXIS", "SIMULATION")] + if len(data_keys) != 1: + return float("nan") + name = data_keys[0] arr = f[name][...] # Reconstruct cell volume from SIMULATION attrs. sim = f["SIMULATION"] @@ -76,7 +96,7 @@ def _total_field_energy(ms: Path) -> float: total = 0.0 found = False for comp_dir in fld.iterdir(): - if not comp_dir.is_dir(): + if not comp_dir.is_dir() or not _is_field_component_dir(comp_dir.name): continue last = _latest_h5(comp_dir) if last is None: diff --git a/tests/test_osiris/test_diagnostics_plots.py b/tests/test_osiris/test_diagnostics_plots.py index c6bd78dd..62358da0 100644 --- a/tests/test_osiris/test_diagnostics_plots.py +++ b/tests/test_osiris/test_diagnostics_plots.py @@ -105,11 +105,58 @@ def _make_full_run(root: Path, n_steps: int = 6, nx: int = 16, npx: int = 12) -> def test_field_energy_components_splits_e_and_b(tmp_path: Path) -> None: run_dir = _make_full_run(tmp_path) ds = oplt.field_energy_components(run_dir) - assert set(ds.data_vars) == {"E_energy", "B_energy", "total_field_energy"} + assert {"E_energy", "B_energy", "total_field_energy"} <= set(ds.data_vars) # Only e1 was dumped, so all energy is electric and B is identically zero. assert np.all(ds["E_energy"].values > 0) assert np.all(ds["B_energy"].values == 0) np.testing.assert_allclose(ds["total_field_energy"].values, ds["E_energy"].values) + # e1 is the only dump, so the longitudinal (EPW) energy equals the E energy. + assert "e1_energy" in ds + np.testing.assert_allclose(ds["e1_energy"].values, ds["E_energy"].values) + + +def test_field_energy_picks_up_savg_and_excludes_poynting(tmp_path: Path) -> None: + """A 2D-style deck dumps e1,savg (not full grid) and s1 Poynting lineouts. + + field_energy_components must resolve the ``e1-savg`` variant, and the + total-field-energy metric must NOT fold the ``s1`` lineout into the sum. + """ + run_dir = tmp_path / "run" + x_ax = ("x1", "x_1", r"c / \omega_p", 0.0, 10.0) + x2_ax = ("x2", "x_2", r"c / \omega_p", 0.0, 4.0) + rng = np.random.default_rng(3) + for k in range(4): + it, t = k * 10, k * 0.5 + _write_dump(run_dir / "MS/FLD/e1-savg" / f"e1-savg-{it:06d}.h5", + "e1", rng.standard_normal(8), t=t, it=it, axes=[x_ax]) + # an s1 Poynting-flux lineout along x2 (would inflate field energy if summed) + _write_dump(run_dir / "MS/FLD/s1-line-x2-24" / f"s1-line-x2-24-{it:06d}.h5", + "s1", 5.0 + rng.standard_normal(6), t=t, it=it, axes=[x2_ax]) + + ds = oplt.field_energy_components(run_dir) + assert np.all(ds["e1_energy"].values > 0) # savg variant resolved + total = opost._total_field_energy(run_dir / "MS") + # only e1 contributes; s1 is excluded, so the total matches the e1 energy of + # the last dump (both are 0.5 * sum(e1^2) * dx). + assert np.isfinite(total) and total > 0 + + +def test_field_energy_components_2d_spatial(tmp_path: Path) -> None: + """field_energy_components integrates over BOTH spatial dims for a 2D field.""" + run_dir = tmp_path / "run" + x_ax = ("x1", "x_1", r"c / \omega_p", 0.0, 10.0) + x2_ax = ("x2", "x_2", r"c / \omega_p", 0.0, 4.0) + rng = np.random.default_rng(5) + for k in range(3): + it, t = k * 10, k * 0.5 + data = rng.standard_normal((6, 8)) # (x2, x1) + _write_dump(run_dir / "MS/FLD/e1" / f"e1-{it:06d}.h5", + "e1", data, t=t, it=it, axes=[x_ax, x2_ax]) + ds = oplt.field_energy_components(run_dir) + assert ds["e1_energy"].sizes["t"] == 3 + assert np.all(ds["e1_energy"].values > 0) + ax = oplt.plot_epw_energy(ds) + assert ax.get_ylabel() def test_load_hist_energy_builds_conservation_total(tmp_path: Path) -> None: From bcc1fa5778962df22c6c698d6c29c85a5cdef52e Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 1 Jul 2026 17:21:48 -0700 Subject: [PATCH 31/49] Fix OSIRIS deck parser to handle single-quoted strings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parser only recognized double-quoted strings, tracking string state by toggling on `"` alone. Fortran single-quoted values (the common OSIRIS convention, e.g. ext_fld='none') fell through _parse_atom's "keep as-is" branch and were stored with their quotes, then double-wrapped on render ("'none'") — which OSIRIS rejects. Single quotes were likewise invisible to _strip_comment, _split_top_commas, and the value scanner, so a `!`, `,`, or `}` inside a single-quoted string (e.g. math_func expressions) was mis-parsed. Make all four string-state trackers handle single and double quotes symmetrically via a shared _QUOTE_CHARS and an active-quote-char variable (only the matching quote closes the string). The renderer already emits double quotes, so single-quoted input now normalizes and round-trips cleanly. Adds tests for single-quote parsing, render normalization, and literal `!`/`,` inside single-quoted strings. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/deck.py | 46 ++++++++++++++++-------- tests/test_osiris/test_deck_roundtrip.py | 26 ++++++++++++++ 2 files changed, 57 insertions(+), 15 deletions(-) diff --git a/adept/osiris/deck.py b/adept/osiris/deck.py index 36a14569..ef617fe6 100644 --- a/adept/osiris/deck.py +++ b/adept/osiris/deck.py @@ -39,6 +39,10 @@ _TRUE_TOKENS = {".true.", ".t."} _FALSE_TOKENS = {".false.", ".f."} +# OSIRIS decks quote strings with either single (Fortran convention) or double +# quotes; both are treated as simple, non-nesting delimiters. +_QUOTE_CHARS = "\"'" + _IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") _KEY_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*(?:\([^)]*\))?") _INT_RE = re.compile(r"^[+-]?\d+$") @@ -47,13 +51,17 @@ def _strip_comment(line: str) -> str: """Drop everything from the first unquoted ``!`` to the end of line.""" - in_string = False + quote = "" out: list[str] = [] for ch in line: - if ch == '"': - in_string = not in_string + if quote: + if ch == quote: + quote = "" + out.append(ch) + elif ch in _QUOTE_CHARS: + quote = ch out.append(ch) - elif ch == "!" and not in_string: + elif ch == "!": break else: out.append(ch) @@ -61,15 +69,19 @@ def _strip_comment(line: str) -> str: def _split_top_commas(s: str) -> list[str]: - """Split on commas not inside double-quotes.""" + """Split on commas not inside a quoted string.""" pieces: list[str] = [] buf: list[str] = [] - in_string = False + quote = "" for ch in s: - if ch == '"': - in_string = not in_string + if quote: + if ch == quote: + quote = "" buf.append(ch) - elif ch == "," and not in_string: + elif ch in _QUOTE_CHARS: + quote = ch + buf.append(ch) + elif ch == ",": pieces.append("".join(buf).strip()) buf = [] else: @@ -85,7 +97,7 @@ def _parse_atom(tok: str) -> Any: t = tok.strip() if not t: return None - if t.startswith('"') and t.endswith('"') and len(t) >= 2: + if len(t) >= 2 and t[0] in _QUOTE_CHARS and t[-1] == t[0]: return t[1:-1] low = t.lower() if low in _TRUE_TOKENS: @@ -151,15 +163,19 @@ def skip_ws(j: int) -> int: raise ValueError(f"Expected '=' after key {key!r} at offset {i}") i += 1 value_start = i - in_string = False + quote = "" while i < n: ch = src[i] - if ch == '"': - in_string = not in_string + if quote: + if ch == quote: + quote = "" + i += 1 + elif ch in _QUOTE_CHARS: + quote = ch i += 1 - elif not in_string and ch == "}": + elif ch == "}": break - elif not in_string and ch == ",": + elif ch == ",": # Peek past whitespace: if the next token is an identifier # followed by '=', this comma terminates the current # key=value pair; otherwise it's an intra-value separator diff --git a/tests/test_osiris/test_deck_roundtrip.py b/tests/test_osiris/test_deck_roundtrip.py index da16803b..4185c465 100644 --- a/tests/test_osiris/test_deck_roundtrip.py +++ b/tests/test_osiris/test_deck_roundtrip.py @@ -80,6 +80,32 @@ def test_comment_stripping_preserves_string_with_bang() -> None: assert sections[0][1]["x"] == 1 +def test_parse_single_quoted_strings() -> None: + # OSIRIS decks commonly use Fortran single quotes; the parser must strip + # them so the value round-trips to a bare Python string (not "'none'"). + text = "emf_bound { ext_fld = 'none', type(1:2) = 'open', 'open', }" + sections = osd.parse_deck(text) + assert sections[0][1]["ext_fld"] == "none" + assert sections[0][1]["type(1:2)"] == ["open", "open"] + + +def test_render_normalizes_single_quotes_to_double() -> None: + parsed = osd.parse_deck("emf_bound { ext_fld = 'none', }") + rendered = osd.render_deck(parsed) + assert '"none"' in rendered and "'none'" not in rendered + # And the value survives a full round-trip unchanged. + assert osd.parse_deck(rendered) == parsed + + +def test_single_quoted_string_with_bang_and_comma() -> None: + # ``!`` and ``,`` inside a single-quoted string are literal, not a comment + # or a value separator (e.g. OSIRIS math_func expressions). + text = "diag { expr = 'if(x>0,1,0)!keep', x = 1, }" + sections = osd.parse_deck(text) + assert sections[0][1]["expr"] == "if(x>0,1,0)!keep" + assert sections[0][1]["x"] == 1 + + def test_merge_overrides_simple() -> None: text = "time { tmin = 0.0, tmax = 30.0, }" s = osd.parse_deck(text) From 2388a5f9d9bb3360d7469f8acabafcc90580f555 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 1 Jul 2026 18:03:14 -0700 Subject: [PATCH 32/49] osiris/runner: salvage a crashed run that produced output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OSIRIS can exit non-zero yet have written usable data — most commonly a segfault on MPI/CUDA teardown *after* "Simulation completed" (seen on Perlmutter 2D CUDA runs), but also a mid-run death with dumps on disk. Previously run_osiris hard-raised on any non-zero exit (or OSIRIS-reported error), so __call__ aborted and post_process never ran — discarding a fully completed run's diagnostics. Now: on a non-zero exit / OSIRIS error, check whether the run produced output (MS/ dumps, binary/ NetCDFs, or HIST/) via _run_produced_output(); if so, log the error as a WARNING and fall through so the caller still consolidates binaries and generates plots from the (possibly partial) data. Only a run that produced nothing still raises. result now carries "crashed". Co-Authored-By: Claude Opus 4.8 --- adept/osiris/runner.py | 77 +++++++++++++++++++++++++------- tests/test_osiris/test_runner.py | 42 +++++++++++++++++ 2 files changed, 102 insertions(+), 17 deletions(-) diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index a7a5853b..abf45104 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -92,6 +92,26 @@ def _sync_tree(src: Path, dst: Path) -> None: shutil.copyfile(p, target) +def _run_produced_output(run_dir: Path, binary_dir: Path) -> bool: + """True if the run wrote any salvageable output (raw dumps, NetCDFs, HIST). + + Used to decide whether a non-zero exit / OSIRIS error is a hard failure or + merely a crash *after* data was already written — e.g. OSIRIS segfaults on + MPI/CUDA teardown after printing "Simulation completed", or dies partway with + dumps on disk. In the latter case downstream consolidation + plotting should + still run on what exists rather than discarding the whole run. + """ + ms = run_dir / "MS" + if ms.is_dir() and next(ms.rglob("*.h5"), None) is not None: + return True + if binary_dir.is_dir() and next(binary_dir.rglob("*.nc"), None) is not None: + return True + hist = run_dir / "HIST" + if hist.is_dir() and next(hist.rglob("*"), None) is not None: + return True + return False + + def run_osiris( deck_text: str, *, @@ -240,31 +260,54 @@ def run_osiris( print(f"[stream] final stage->persist sync issue (partial outputs may be lost): {e}") shutil.rmtree(stage_dir, ignore_errors=True) + # Classify the outcome. A run that crashed (non-zero exit) or that OSIRIS + # flagged as an error can still have produced usable output — most commonly a + # segfault on MPI/CUDA teardown *after* "Simulation completed", but also a + # mid-run death with dumps already on disk. In that case we salvage: log the + # error and fall through so the caller still consolidates binaries and plots + # what was written. Only a run that produced NOTHING is a hard failure. + failure_msg: str | None = None if rc != 0: tail = "".join(stderr_tail[-50:]) or "(empty stderr)" - raise RuntimeError( - f"OSIRIS exited with status {rc}.\n cmd: {shlex.join(cmd)}\n cwd: {work_dir}\n stderr tail:\n{tail}" + failure_msg = ( + f"OSIRIS exited with status {rc}.\n cmd: {shlex.join(cmd)}\n" + f" cwd: {work_dir}\n stderr tail:\n{tail}" ) + else: + # OSIRIS can exit 0 even on input-file errors: it prints something + # like 'Error reading ... / aborting...' to stderr (and '(*error*)' + # to stdout) before terminating. Detect both. + err_lines = [ln for ln in stderr_tail if _looks_like_osiris_error(ln)] + err_lines += [ln for ln in stdout_tail if "(*error*)" in ln] + if err_lines: + out_tail = "".join(stdout_tail[-20:]) + err_tail = "".join(stderr_tail[-20:]) + failure_msg = ( + "OSIRIS reported an error despite exit-code 0:\n" + f" cmd: {shlex.join(cmd)}\n" + f" cwd: {work_dir}\n" + f" stderr tail:\n{err_tail}\n" + f" stdout tail:\n{out_tail}" + ) - # OSIRIS can exit 0 even on input-file errors: it prints something - # like 'Error reading ... / aborting...' to stderr (and '(*error*)' - # to stdout) before terminating. Detect both. - err_lines = [ln for ln in stderr_tail if _looks_like_osiris_error(ln)] - err_lines += [ln for ln in stdout_tail if "(*error*)" in ln] - if err_lines: - out_tail = "".join(stdout_tail[-20:]) - err_tail = "".join(stderr_tail[-20:]) - raise RuntimeError( - "OSIRIS reported an error despite exit-code 0:\n" - f" cmd: {shlex.join(cmd)}\n" - f" cwd: {work_dir}\n" - f" stderr tail:\n{err_tail}\n" - f" stdout tail:\n{out_tail}" - ) + crashed = failure_msg is not None + if crashed: + if _run_produced_output(run_dir, binary_dir): + print(f"[osiris] WARNING: {failure_msg}") + print( + "[osiris] the run produced output files despite the error above — " + "consolidating binaries and generating plots from the (possibly " + "partial) data anyway. Verify results carefully." + ) + else: + raise RuntimeError( + f"{failure_msg}\n (no output files were written — nothing to salvage.)" + ) result: dict[str, Any] = { "run_dir": run_dir, "exit_code": rc, + "crashed": crashed, "wall_time": wall_time, "cmd": cmd, } diff --git a/tests/test_osiris/test_runner.py b/tests/test_osiris/test_runner.py index 7373e476..4c56debb 100644 --- a/tests/test_osiris/test_runner.py +++ b/tests/test_osiris/test_runner.py @@ -54,6 +54,48 @@ def test_run_osiris_missing_binary_raises(tmp_path: Path) -> None: ) +def _write_fake_binary(path: Path, body: str) -> Path: + path.write_text("#!/bin/bash\n" + body + "\n") + path.chmod(0o755) + return path + + +def test_run_osiris_crash_with_output_is_salvaged(tmp_path: Path) -> None: + # A binary that writes a dump then exits non-zero (e.g. OSIRIS segfaulting on + # MPI/CUDA teardown *after* "Simulation completed") must NOT raise: the run + # produced data, so the runner salvages it and lets the caller consolidate + + # plot what was written. + fake = _write_fake_binary( + tmp_path / "fake-osiris", + 'mkdir -p MS/FLD/e1 && : > MS/FLD/e1/e1-000000.h5 && exit 139', + ) + result = runner.run_osiris( + "node_conf {}", + binary=str(fake), + mpi_ranks=1, + run_root=tmp_path, + stream_convert=False, + ) + assert result["exit_code"] == 139 + assert result["crashed"] is True + assert next((result["run_dir"] / "MS").rglob("*.h5"), None) is not None + + +def test_run_osiris_crash_no_output_raises(tmp_path: Path) -> None: + # A binary that exits non-zero WITHOUT writing anything is a hard failure — + # there is nothing to salvage, so the runner still raises. + fake = _write_fake_binary(tmp_path / "fake-osiris", "exit 1") + with pytest.raises(RuntimeError) as excinfo: + runner.run_osiris( + "node_conf {}", + binary=str(fake), + mpi_ranks=1, + run_root=tmp_path, + stream_convert=False, + ) + assert "nothing to salvage" in str(excinfo.value) + + @pytest.mark.skipif( not (OSIRIS_BIN_1D and Path(OSIRIS_BIN_1D).exists()), reason="set OSIRIS_BIN_1D (or OSIRIS_BIN) to a built osiris-1D.e to run", From 34b7b356e392aa14101f6dacacf38676ff1f4fde Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Wed, 1 Jul 2026 18:21:51 -0700 Subject: [PATCH 33/49] osiris/io: key diagnostic series by per-report base, not just directory Multiple line/report diagnostics of the same field share one MS/ directory and differ only by an index in the base name (e.g. two 's1,...,line' reports -> s1-tavg-line-x2-01 / -02). _iter_from_name reads both -x2-0N-000000.h5 as iteration 0, so keying a series by directory stacked them into one broken NetCDF, losing the entrance/exit distinction the laser-transmission budget needs. io: group a dir's dumps by base (_dump_base/_series_dumps); list_diagnostics exposes each base as its own diagnostic via a synthetic

/ handle; _sort_dumps resolves that handle and selects one series (bare multi-series dir raises). load_series / _diag_is_raw / convert_diagnostic_streaming compose through _sort_dumps, so batch consolidation writes one NetCDF per series. stream: the concurrent converter skips multi-report dirs (can't stream several series as one) and leaves them to the batch pass. Single-series dirs unchanged. Co-Authored-By: Claude Opus 4.8 --- adept/osiris/io.py | 77 +++++++++++++++++++++++++-- adept/osiris/stream.py | 18 +++++-- tests/test_osiris/test_post_netcdf.py | 59 ++++++++++++++++++++ 3 files changed, 146 insertions(+), 8 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 57adc47a..5ac3352f 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -262,9 +262,65 @@ def _diag_is_raw(relpath: str, directory: str | Path) -> bool: return False -def _sort_dumps(directory: Path) -> list[Path]: - files = [p for p in directory.iterdir() if p.is_file() and p.suffix == ".h5"] - return sorted(files, key=_iter_from_name) +_BASE_RE = re.compile(r"-\d+\.h5$") + + +def _dump_base(path: Path) -> str: + """A dump's series base = its filename minus the trailing ``-.h5``. + + ``e1-savg-000123.h5`` -> ``e1-savg``. A directory usually holds one base, but + multiple line/report diagnostics of the same field share a directory and + differ only by an index in the base — e.g. two ``s1,...,line`` reports produce + ``s1-tavg-line-x2-01`` and ``s1-tavg-line-x2-02``. Those are distinct time + series and must not be stacked into one. + """ + return _BASE_RE.sub("", path.name) + + +def _series_dumps(directory: Path) -> dict[str, list[Path]]: + """Group ``directory``'s ``.h5`` dumps into time series keyed by base name.""" + groups: dict[str, list[Path]] = {} + for p in directory.iterdir(): + if p.is_file() and p.suffix == ".h5": + groups.setdefault(_dump_base(p), []).append(p) + for dumps in groups.values(): + dumps.sort(key=_iter_from_name) + return groups + + +def _resolve_series(handle: Path) -> tuple[Path, str | None]: + """Resolve a diagnostic handle to ``(directory, base)``. + + A handle from :func:`list_diagnostics` is either a real dump directory (one + series, ``base`` is ``None``) or a synthetic ``/`` path naming one + report series inside a directory that holds several. + """ + handle = Path(handle) + if handle.is_dir(): + return handle, None + return handle.parent, handle.name + + +def _sort_dumps(handle: Path, base: str | None = None) -> list[Path]: + """Sorted ``.h5`` dumps for one diagnostic series. + + ``handle`` may be a dump directory or a synthetic per-report handle + (``/``); ``base`` selects one series inside a multi-report + directory. A bare directory holding more than one series is ambiguous and + raises — load one via its per-report handle from :func:`list_diagnostics`. + """ + directory, hbase = _resolve_series(handle) + if base is None: + base = hbase + groups = _series_dumps(directory) + if base is not None: + return groups.get(base, []) + if len(groups) > 1: + raise ValueError( + f"{directory} holds multiple report series {sorted(groups)}; " + "load one via its per-report handle from list_diagnostics()." + ) + return next(iter(groups.values()), []) def load_series(directory: str | Path) -> xr.DataArray: @@ -382,8 +438,19 @@ def list_diagnostics(run_dir: str | Path) -> dict[str, Path]: for d in ms.rglob("*"): if not d.is_dir(): continue - if any(c.suffix == ".h5" for c in d.iterdir()): - out[str(d.relative_to(ms))] = d + groups = _series_dumps(d) + if not groups: + continue + rel = str(d.relative_to(ms)) + if len(groups) == 1: + out[rel] = d + else: + # A directory with several report series (e.g. two s1 line + # lineouts at different perpendicular cells) exposes each as its + # own diagnostic via a synthetic / handle, so they are + # converted to separate NetCDFs instead of merged into one. + for base in sorted(groups): + out[f"{rel}/{base}"] = d / base return out # No MS/ tree — fall back to the saved-NetCDF layout. nc = list_diagnostics_nc(run_dir) diff --git a/adept/osiris/stream.py b/adept/osiris/stream.py index de484881..16182928 100644 --- a/adept/osiris/stream.py +++ b/adept/osiris/stream.py @@ -246,6 +246,7 @@ def __init__(self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print, pers self._completed: set[str] = set() # grid relpaths fully converted to .nc self._raw_completed: set[str] = set() # raw relpaths fully mirrored to persist self._is_raw: dict[str, bool] = {} # rel -> raw? (cached; avoids re-peeking h5) + self._multibase_skipped: set[str] = set() # multi-report dirs left to the batch pass self._stop = threading.Event() self._thread: threading.Thread | None = None @@ -300,11 +301,22 @@ def _diags(self) -> dict[str, Path]: if not d.is_dir(): continue try: - has_h5 = any(c.suffix == ".h5" for c in d.iterdir()) + groups = _io._series_dumps(d) except OSError: continue - if has_h5: - out[str(d.relative_to(self.ms))] = d + if not groups: + continue + rel = str(d.relative_to(self.ms)) + if len(groups) > 1: + # A directory holding several report series (e.g. two s1 line + # lineouts) can't be streamed as one series; leave it to the + # batch pass (save_run_datasets), which writes one NetCDF per + # series via the per-report handles from list_diagnostics. + if rel not in self._multibase_skipped: + self._multibase_skipped.add(rel) + self._log(f"[stream] {rel}: {len(groups)} report series -> batch pass") + continue + out[rel] = d return out def _diag_is_raw(self, rel: str, d: Path) -> bool: diff --git a/tests/test_osiris/test_post_netcdf.py b/tests/test_osiris/test_post_netcdf.py index dddcbdf7..36b7aa5a 100644 --- a/tests/test_osiris/test_post_netcdf.py +++ b/tests/test_osiris/test_post_netcdf.py @@ -68,6 +68,65 @@ def _make_run(root: Path, n_steps: int = 4, nx: int = 8) -> Path: return run_dir +def _write_lineout(dir_: Path, base: str, n_steps: int, nx: int, fill: float) -> None: + """Write one line/report series ``-.h5`` into ``dir_``.""" + for k in range(n_steps): + it = k * 10 + _write_dump( + dir_ / f"{base}-{it:06d}.h5", + "s1", + np.full(nx, fill + k, dtype="float64"), + t=k * 0.5, + it=it, + axes=[("x2", "x_2", r"c / \omega_p", 0.0, 5.0)], + ) + + +def test_multi_report_lineouts_are_separate_series(tmp_path: Path) -> None: + # Two s1 line lineouts share one MS/FLD dir, distinguished only by the report + # index in the base name (s1-tavg-line-x2-01 / -02). Both -x2-0N-000000.h5 + # read as iteration 0, so keyed by directory they would merge; they must be + # exposed and loaded as SEPARATE series instead. + run_dir = tmp_path / "run" + d = run_dir / "MS" / "FLD" / "s1-tavg-line" + _write_lineout(d, "s1-tavg-line-x2-01", n_steps=4, nx=6, fill=100.0) + _write_lineout(d, "s1-tavg-line-x2-02", n_steps=4, nx=6, fill=200.0) + + diags = oio.list_diagnostics(run_dir) + keys = {k for k in diags if k.startswith("FLD/s1-tavg-line")} + assert keys == { + "FLD/s1-tavg-line/s1-tavg-line-x2-01", + "FLD/s1-tavg-line/s1-tavg-line-x2-02", + } + ent = oio.load_series(diags["FLD/s1-tavg-line/s1-tavg-line-x2-01"]) + ex = oio.load_series(diags["FLD/s1-tavg-line/s1-tavg-line-x2-02"]) + assert ent.sizes["t"] == 4 and ex.sizes["t"] == 4 # not merged into 8 + assert float(ent.isel(t=0).values.mean()) == 100.0 + assert float(ex.isel(t=0).values.mean()) == 200.0 + + +def test_save_run_datasets_splits_multi_report_lineouts(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + d = run_dir / "MS" / "FLD" / "s1-tavg-line" + _write_lineout(d, "s1-tavg-line-x2-01", 3, 6, 100.0) + _write_lineout(d, "s1-tavg-line-x2-02", 3, 6, 200.0) + out = tmp_path / "binary" + written = oio.save_run_datasets(run_dir, out) + names = {p.relative_to(out).as_posix() for p in written} + assert "FLD/s1-tavg-line/s1-tavg-line-x2-01.nc" in names + assert "FLD/s1-tavg-line/s1-tavg-line-x2-02.nc" in names + # round-trips separately (not merged) + assert oio.load_series(out / "FLD/s1-tavg-line/s1-tavg-line-x2-01.nc").sizes["t"] == 3 + + +def test_single_report_dir_still_keys_by_dir(tmp_path: Path) -> None: + # Regression: a normal single-series dir keys by its dir relpath, unchanged. + run_dir = _make_run(tmp_path, n_steps=3, nx=8) + diags = oio.list_diagnostics(run_dir) + assert "FLD/e1" in diags + assert oio.load_series(diags["FLD/e1"]).sizes["t"] == 3 + + def test_save_run_datasets_writes_full_timeseries(tmp_path: Path) -> None: run_dir = _make_run(tmp_path, n_steps=4, nx=8) out = tmp_path / "binary" From cb76f1d2bd1373916f4b1d62ffde112ffe405b31 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 3 Jul 2026 16:14:44 -0700 Subject: [PATCH 34/49] =?UTF-8?q?osiris:=20add=20stage=5Fdiscard=5Fh5=20?= =?UTF-8?q?=E2=80=94=20stream=20grid=20dumps=20without=20mirroring=20the?= =?UTF-8?q?=20h5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In ramdisk-staging mode the converter mirrored every completed grid dump to the durable MS/ tree before deleting the scratch copy, so a staged run still paid one inode + one copy per dump on the persist filesystem (~140k h5 files per production SRS run). With osiris.stage_discard_h5: true the grid dump is deleted after being appended to the streamed NetCDF and the binary/*.nc become the only copy; RAW dumps are still mirrored (the batch path builds their NetCDFs from the HDF5). Dumps whose stream failed are left in place for the runner's final sync, so partial failures remain recoverable. Co-Authored-By: Claude Fable 5 --- adept/osiris/base.py | 1 + adept/osiris/runner.py | 11 +++++++++++ adept/osiris/stream.py | 39 ++++++++++++++++++++++++++------------- 3 files changed, 38 insertions(+), 13 deletions(-) diff --git a/adept/osiris/base.py b/adept/osiris/base.py index b762079f..833cde83 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -156,6 +156,7 @@ def __call__(self, trainable_modules: dict, args: dict) -> dict: stream_convert=bool(osiris_cfg.get("stream_convert", True)), stream_poll_s=float(osiris_cfg.get("stream_poll_s", 10.0)), stage_root=osiris_cfg.get("stage_root"), + stage_discard_h5=bool(osiris_cfg.get("stage_discard_h5", False)), ) return {"solver result": result} diff --git a/adept/osiris/runner.py b/adept/osiris/runner.py index abf45104..8edf7fc8 100644 --- a/adept/osiris/runner.py +++ b/adept/osiris/runner.py @@ -124,6 +124,7 @@ def run_osiris( stream_convert: bool = True, stream_poll_s: float = 10.0, stage_root: str | Path | None = None, + stage_discard_h5: bool = False, ) -> dict[str, Any]: """Run OSIRIS and return run metadata. @@ -151,6 +152,15 @@ def run_osiris( always the durable directory, so post-processing is unchanged; the scratch is reclaimed before returning. Single-node only (``/dev/shm`` is per-node). + **``stage_discard_h5``** (staging only): grid dumps are *not* mirrored to + the durable ``MS/`` tree -- once appended to the streamed NetCDF the scratch + HDF5 is simply deleted. Cuts the persist-filesystem inode count by ~the + grid-dump count. RAW dumps are still mirrored (the batch path needs them), + and any dump whose stream failed is left for the final sync, so partial + failures remain recoverable. The durable ``MS/`` tree is then *incomplete + by design*: offline re-conversion of grid diagnostics is impossible, the + streamed ``binary/*.nc`` are the only copy. + Raises ``RuntimeError`` on non-zero exit code, with the last lines of stderr included in the message. """ @@ -208,6 +218,7 @@ def run_osiris( binary_dir, poll_s=stream_poll_s, persist_dir=run_dir if staging else None, + discard_grid_h5=staging and stage_discard_h5, ) except Exception as e: # never let the converter block the run print(f"[stream] disabled (init failed): {e}") diff --git a/adept/osiris/stream.py b/adept/osiris/stream.py index 16182928..fbe4cef5 100644 --- a/adept/osiris/stream.py +++ b/adept/osiris/stream.py @@ -235,11 +235,21 @@ class StreamConverter: exactly as before. """ - def __init__(self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print, persist_dir=None): + def __init__( + self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print, persist_dir=None, discard_grid_h5: bool = False + ): self.ms = Path(run_dir) / "MS" self.out_dir = Path(out_dir) self.persist_dir = Path(persist_dir) if persist_dir is not None else None self.persist_ms = (self.persist_dir / "MS") if self.persist_dir is not None else None + # Staging only: when set, grid dumps are deleted from the scratch after + # being appended to the NetCDF *without* being mirrored to persist_dir/MS + # (the .nc is the durable artifact). Saves one inode+copy per dump on the + # persist filesystem. RAW dumps are always mirrored (the batch path needs + # the HDF5). Trade-off: a grid diagnostic whose stream fails mid-run can + # only be rebuilt from the dumps not yet reaped (they are left in place + # and folded to persist by the runner's final sync). + self.discard_grid_h5 = bool(discard_grid_h5) self.poll_s = float(poll_s) self._log = logger self._writers: dict[str, StreamWriter] = {} @@ -408,7 +418,7 @@ def _drain_grid_staged(self, rel: str, d: Path, *, final: bool) -> None: self._writers[rel] = w for p in safe: w.append(_io.load_grid_h5(p)) - self._reap(rel, p) + self._reap(rel, p, mirror=not self.discard_grid_h5) w.flush() if final: @@ -426,20 +436,23 @@ def _drain_raw(self, rel: str, d: Path, *, final: bool) -> None: if final: self._raw_completed.add(rel) - def _reap(self, rel: str, p: Path) -> None: - """Mirror one scratch dump to ``persist_dir/MS`` then delete the scratch - copy. Best-effort: on any failure the scratch file is left in place for - the runner's final sync to catch, so data is never lost and the watcher - never crashes.""" + def _reap(self, rel: str, p: Path, *, mirror: bool = True) -> None: + """Delete one scratch dump, first mirroring it to ``persist_dir/MS`` + unless ``mirror`` is False (``discard_grid_h5`` mode: the dump is + already in the streamed NetCDF and the raw HDF5 is not wanted on the + persist filesystem). Best-effort: on any failure the scratch file is + left in place for the runner's final sync to catch, so data is never + lost and the watcher never crashes.""" if self.persist_dir is None: return try: - dest = self.persist_ms / rel / p.name - if not dest.exists(): - dest.parent.mkdir(parents=True, exist_ok=True) - tmp = dest.with_name(dest.name + ".part") - shutil.copyfile(p, tmp) - os.replace(tmp, dest) # atomic publish on the persist filesystem + if mirror: + dest = self.persist_ms / rel / p.name + if not dest.exists(): + dest.parent.mkdir(parents=True, exist_ok=True) + tmp = dest.with_name(dest.name + ".part") + shutil.copyfile(p, tmp) + os.replace(tmp, dest) # atomic publish on the persist filesystem p.unlink() except Exception as e: self._log(f"[stream] reap {rel}/{p.name} deferred: {e}") From 1cc539d9436e693db53350ad470e7cb92802323c Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 3 Jul 2026 16:14:44 -0700 Subject: [PATCH 35/49] osiris: teach diagnostic discovery and post metrics the discard layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit list_diagnostics only walked MS/ whenever it existed, so under stage_discard_h5 (MS/ holding only RAW) every grid diagnostic was invisible to save_run_datasets and the plotting passes — the first discard-mode run uploaded 2 of 14 binary artifacts and 2 of 53 plots. Discovery now merges run_dir/binary NetCDFs for any diagnostic with no raw dumps on disk, and the no-MS/ fallback prefers a run dir's binary/ subdir. final_iter falls back to scanning the whole MS/ tree (RAW dumps carry iteration numbers) and field_energy_final is computed from the last streamed FLD/*.nc slice when no FLD h5 exists. Regression tests for the discard reap and the discovery/batch consolidation. Co-Authored-By: Claude Fable 5 --- adept/osiris/io.py | 17 +++++++++- adept/osiris/post.py | 58 +++++++++++++++++++++++++++----- tests/test_osiris/test_stream.py | 54 +++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 10 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 5ac3352f..058bb055 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -451,8 +451,23 @@ def list_diagnostics(run_dir: str | Path) -> dict[str, Path]: # converted to separate NetCDFs instead of merged into one. for base in sorted(groups): out[f"{rel}/{base}"] = d / base + # stage_discard_h5 layout: grid dumps were streamed to run_dir/binary + # and never mirrored to MS/ (which then holds only RAW). Merge the + # streamed NetCDFs for any diagnostic with no raw dumps on disk so + # post-processing and plots see the full set; load_series accepts the + # .nc handles directly. + bdir = run_dir / "binary" + if bdir.is_dir(): + for rel, p in list_diagnostics_nc(bdir).items(): + out.setdefault(rel, p) return out - # No MS/ tree — fall back to the saved-NetCDF layout. + # No MS/ tree — fall back to the saved-NetCDF layout (a run dir's binary/ + # subdir first, else run_dir itself *being* a binary dir). + bdir = run_dir / "binary" + if bdir.is_dir(): + nc = list_diagnostics_nc(bdir) + if nc: + return nc nc = list_diagnostics_nc(run_dir) if nc: return nc diff --git a/adept/osiris/post.py b/adept/osiris/post.py index 0535d06e..cda548f6 100644 --- a/adept/osiris/post.py +++ b/adept/osiris/post.py @@ -109,19 +109,54 @@ def _total_field_energy(ms: Path) -> float: def _final_iter(ms: Path) -> int: - """Largest iteration number seen across all FLD dumps; -1 if none.""" - fld = ms / "FLD" - if not fld.is_dir(): - return -1 + """Largest iteration number seen across all FLD dumps; -1 if none. + + Under ``stage_discard_h5`` the FLD h5 never reach the durable ``MS/``; + fall back to scanning the whole tree (RAW dumps still carry iteration + numbers, at their coarser cadence).""" best = -1 - for comp_dir in fld.iterdir(): - for p in _iter_h5_files(comp_dir): + for scan_dir in (ms / "FLD", ms): + if not scan_dir.is_dir(): + continue + for p in scan_dir.rglob("*.h5"): m = _ITER_RE.search(p.name) if m: best = max(best, int(m.group(1))) + if best >= 0: + return best return best +def _total_field_energy_nc(binary_dir: Path) -> float: + """|E|^2/2 + |B|^2/2 at the last streamed slice of ``binary_dir/FLD/*.nc``. + + Fallback for the ``stage_discard_h5`` layout, where no FLD h5 exists on + the durable ``MS/``. Returns NaN when no usable component is found.""" + import xarray as xr + + fld = binary_dir / "FLD" + if not fld.is_dir(): + return float("nan") + total, found = 0.0, False + for p in sorted(fld.glob("*.nc")): + if not _is_field_component_dir(p.stem): + continue + try: + with xr.open_dataset(p) as ds: + da = ds[p.stem] if p.stem in ds else ds[next(iter(ds.data_vars))] + last = da.isel(t=-1).astype("float64") + dvol = 1.0 + for dim in last.dims: + c = last[dim].values + if len(c) > 1: + dvol *= float(c[1] - c[0]) + total += 0.5 * float((last.values**2).sum()) * dvol + found = True + except Exception: + continue + return total if found else float("nan") + + def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: """Post-process a finished OSIRIS run. @@ -149,9 +184,14 @@ def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: "field_energy_final": _total_field_energy(ms), "final_iter": float(_final_iter(ms)), } - - # Convert each diagnostic's time history to an xarray netCDF. - if ms.is_dir(): + if metrics["field_energy_final"] != metrics["field_energy_final"] and streamed_dir is not None: + # stage_discard_h5: no FLD h5 on the durable MS/ — use the streamed nc. + metrics["field_energy_final"] = _total_field_energy_nc(Path(streamed_dir)) + + # Convert each diagnostic's time history to an xarray netCDF. Under + # stage_discard_h5 with no RAW diagnostics MS/ may not exist at all while + # the streamed NetCDFs do — proceed in that case too. + if ms.is_dir() or (streamed_dir is not None and Path(streamed_dir).is_dir()): _io.save_run_datasets( run_dir, td / "binary", diff --git a/tests/test_osiris/test_stream.py b/tests/test_osiris/test_stream.py index 38be8c78..140271f7 100644 --- a/tests/test_osiris/test_stream.py +++ b/tests/test_osiris/test_stream.py @@ -387,6 +387,60 @@ def test_staged_converter_mirrors_and_reaps(tmp_path: Path) -> None: ) +def test_staged_converter_discard_grid_h5(tmp_path: Path) -> None: + """With ``discard_grid_h5`` grid dumps are streamed + deleted WITHOUT a + persist mirror (inode saver); RAW dumps are still mirrored as H5.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=4, nx=8) + for k in range(3): + _write_raw_dump(stage / "MS" / "RAW" / "species_1" / f"RAW-species_1-{k * 10:06d}.h5", t=k * 0.5, it=k * 10) + + conv = ostream.StreamConverter( + stage, persist / "binary", poll_s=0.01, persist_dir=persist, discard_grid_h5=True + ) + completed = conv.finalize() + + assert completed == {"FLD/e1"} + # Grid: streamed NetCDF is the only durable artifact -- no H5 mirror. + assert (persist / "binary" / "FLD" / "e1.nc").exists() + assert not (persist / "MS" / "FLD").exists() + # RAW: still mirrored (batch path needs the H5). + assert len(list((persist / "MS" / "RAW" / "species_1").glob("*.h5"))) == 3 + # Scratch fully reaped either way. + assert not list((stage / "MS").rglob("*.h5")) + # The streamed series carries all 4 dumps. + assert oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc").sizes["t"] == 4 + + +def test_discard_layout_discovery_and_batch(tmp_path: Path) -> None: + """End-to-end stage_discard_h5 layout: MS/ holds only RAW, grid data lives + solely in the streamed NetCDFs — discovery and the batch pass must still + see and consolidate the grid diagnostics (regression: smoke-3 uploaded + 2/14 binary files because list_diagnostics only walked MS/).""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=4, nx=8) + for k in range(3): + _write_raw_dump(stage / "MS" / "RAW" / "species_1" / f"RAW-species_1-{k * 10:06d}.h5", t=k * 0.5, it=k * 10) + conv = ostream.StreamConverter( + stage, persist / "binary", poll_s=0.01, persist_dir=persist, discard_grid_h5=True + ) + conv.finalize() + + # Discovery must see the union: RAW from MS/, grid from the streamed nc. + diags = oio.list_diagnostics(persist) + assert "RAW/species_1" in diags + assert "FLD/e1" in diags and str(diags["FLD/e1"]).endswith(".nc") + + # The batch pass must consolidate the streamed grid nc into td/binary. + out = tmp_path / "td_binary" + oio.save_run_datasets(persist, out, stream=True, streamed_dir=persist / "binary") + assert (out / "FLD" / "e1.nc").exists() + assert oio.load_series_nc(out / "FLD" / "e1.nc").sizes["t"] == 4 + assert (out / "RAW" / "species_1.nc").exists() + + def test_non_staged_converter_never_deletes(tmp_path: Path) -> None: """Without ``persist_dir`` the converter must leave the source dumps in place (the durable run dir *is* run_dir) — no reaping of the in-place tree.""" From 3c681d307571f1ad42c60202599147f95b7e0877 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Sat, 4 Jul 2026 16:21:38 -0700 Subject: [PATCH 36/49] osiris: bound post-processing memory with sliced + lazy series loads A full phase-space history decompresses to tens of GB in memory (x1gamma_q1 is ~48 GB: 11705 x 1000 x 1024; ~0.19 GB gzipped on disk, sparse ~250x), and eager-loading one per plot OOM'd node bundles running many sims at once. Load only what each plot needs: - io.load_series / load_series_nc gain t_indices= (read just those time slices) plus a series_len() helper; a new open_series() context manager yields a lazily-backed array whose isel(t=it) reads one dump on demand, for consumers that walk the whole time axis reducing per dump. - save_canned_plots' phase-space plots load the final dump + evolution panels via t_indices; plot_phasespace_evolution samples the same set in both call paths and takes its colour scale from the panels only; plot_profile gains avg_window= so a slice-selected series keeps the intended mean window. - _decorate copies shallow (a deep copy materialized a lazily-opened series just to relabel its axes). Docs: the usage guide documents the sliced/lazy loading convention for new plotters. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01MMRKMZ3bRnUqLBjVtw7QJM --- adept/osiris/io.py | 104 +++++++++++++++++++++++++++++++++++-- adept/osiris/plots.py | 67 +++++++++++++++++++++--- docs/osiris-adept-usage.md | 3 ++ 3 files changed, 162 insertions(+), 12 deletions(-) diff --git a/adept/osiris/io.py b/adept/osiris/io.py index 058bb055..99f3eb0d 100644 --- a/adept/osiris/io.py +++ b/adept/osiris/io.py @@ -22,6 +22,8 @@ import json import re import shutil +from collections.abc import Sequence +from contextlib import contextmanager from pathlib import Path import h5py @@ -323,7 +325,36 @@ def _sort_dumps(handle: Path, base: str | None = None) -> list[Path]: return next(iter(groups.values()), []) -def load_series(directory: str | Path) -> xr.DataArray: +def series_len(source: str | Path) -> int: + """Number of time slices in a diagnostic series, without loading its data. + + ``source`` is any handle accepted by :func:`load_series` (a dump directory, + a per-report handle, or a saved-series ``.nc`` file). Pairs with + ``load_series(..., t_indices=...)`` so callers can pick a few slices of a + long series — e.g. a phase-space history — without ever materializing the + whole ``(t, ...)`` array. + """ + source = Path(source) + if source.is_file() and source.suffix == ".nc": + with xr.open_dataset(source, engine="h5netcdf") as ds: + return int(ds.sizes.get("t", 0)) + return len(_sort_dumps(source)) + + +def _normalize_t_indices(t_indices: Sequence[int], n: int, what: str) -> list[int]: + """Resolve possibly-negative time indices against a series of length ``n``.""" + out = [] + for i in t_indices: + j = int(i) + n if int(i) < 0 else int(i) + if not 0 <= j < n: + raise IndexError(f"t index {i} out of range for {what} with {n} time slices") + out.append(j) + if not out: + raise ValueError(f"t_indices is empty for {what}") + return out + + +def load_series(directory: str | Path, t_indices: Sequence[int] | None = None) -> xr.DataArray: """Stack every dump in ``directory`` into a ``(t, ...)`` DataArray. All files must share the same diagnostic name, the same spatial / @@ -333,13 +364,23 @@ def load_series(directory: str | Path) -> xr.DataArray: For regenerating plots from saved artifacts, ``directory`` may instead be a single ``.nc`` file written by :func:`save_run_datasets`; it is then loaded via :func:`load_series_nc`, returning the same stacked ``(t, ...)`` array. + + ``t_indices`` selects a subset of time slices (negative indices allowed, + order preserved) and reads **only** those from disk — per-dump files for + the HDF5 layout, a partial read for the NetCDF layout. A full phase-space + history can run to tens of GB while its plots need under ten slices, so + slice-selecting callers (see the phase-space paths in ``plots``) is what + keeps concurrent post-processing inside a node's memory. """ directory = Path(directory) if directory.is_file() and directory.suffix == ".nc": - return load_series_nc(directory) + return load_series_nc(directory, t_indices=t_indices) dumps = _sort_dumps(directory) if not dumps: raise FileNotFoundError(f"No .h5 dumps in {directory}") + if t_indices is not None: + idx = _normalize_t_indices(t_indices, len(dumps), str(directory)) + dumps = [dumps[j] for j in idx] first = load_grid_h5(dumps[0]) n = len(dumps) @@ -514,7 +555,7 @@ def series_to_dataset(da: xr.DataArray) -> xr.Dataset: return ds -def load_series_nc(path: str | Path) -> xr.DataArray: +def load_series_nc(path: str | Path, t_indices: Sequence[int] | None = None) -> xr.DataArray: """Load a grid/phase-space series NetCDF back into a plotter-ready DataArray. This is the inverse of :func:`series_to_dataset` (the form @@ -525,9 +566,33 @@ def load_series_nc(path: str | Path) -> xr.DataArray: :func:`load_series` off the raw HDF5 tree. The whole point is that the canned plots can be regenerated from the saved NetCDFs alone. - The file is read fully into memory and closed before returning. + The file is read into memory and closed before returning. With + ``t_indices`` (see :func:`load_series`) the read is *partial*: the file is + opened lazily and only the selected time slices are materialized — the + streamed series are written one dump at a time, so this reads a handful of + chunks instead of the (potentially tens-of-GB) full history. Reductions on + a lazily-opened xarray would silently pull in the whole variable, which is + why the selection happens here, before any consumer touches the data. + """ + if t_indices is None: + ds = xr.load_dataset(path, engine="h5netcdf") + else: + with xr.open_dataset(path, engine="h5netcdf") as f: + idx = _normalize_t_indices(t_indices, int(f.sizes.get("t", 0)), str(path)) + ds = f.isel(t=idx).load() + return _series_da_from_dataset(ds, path) + + +def _series_da_from_dataset(ds: xr.Dataset, path: str | Path) -> xr.DataArray: + """Rebuild a plotter-ready DataArray from an opened saved-series Dataset. + + Shared by the eager :func:`load_series_nc` and the lazy :func:`open_series`: + it only reads the (small) coordinate variables to restore the + ``axis_units`` / ``axis_long_names`` / ``autoscaled_dims`` metadata, and + never touches the (possibly huge) data variable — so on a lazily-opened + ``ds`` the returned array stays lazy and materializes only the slices a + caller actually indexes. """ - ds = xr.load_dataset(path, engine="h5netcdf") names = list(ds.data_vars) if not names: raise ValueError(f"No data variable in {path}") @@ -570,6 +635,35 @@ def load_series_nc(path: str | Path) -> xr.DataArray: return da +@contextmanager +def open_series(source: str | Path): + """Yield a ``(t, …)`` series whose slices are read from disk on demand. + + The memory-bounded companion to :func:`load_series` for consumers that walk + or window the time axis and reduce **one dump at a time** (e.g. the SRS + phase-space budgets in ``osiris_lpi.srs_post``, which iterate + ``da.isel(t=it)`` and sum each slice): the full ``(t, …)`` array is never + materialized, so a 48 GB phase-space history costs only one resident dump. + Unlike ``load_series(..., t_indices=)`` this does not need the caller to + know which slices up front, so it also serves the whole-history passes (the + flux-vs-time trace) that a slice list can't express. + + For a saved-series ``.nc`` the file is opened lazily and closed on exit; the + yielded DataArray is valid **only inside the with-block**. For a raw + per-dump directory (no single lazy file) it falls back to an eager + :func:`load_series` and yields that array (nothing to close). + """ + source = Path(source) + if source.is_file() and source.suffix == ".nc": + ds = xr.open_dataset(source, engine="h5netcdf") + try: + yield _series_da_from_dataset(ds, source) + finally: + ds.close() + else: + yield load_series(source) + + def list_diagnostics_nc(binary_dir: str | Path) -> dict[str, Path]: """Map every saved-NetCDF diagnostic to its ``.nc`` file. diff --git a/adept/osiris/plots.py b/adept/osiris/plots.py index 2d37cddc..3fdf72b2 100644 --- a/adept/osiris/plots.py +++ b/adept/osiris/plots.py @@ -263,6 +263,17 @@ def plot_phasespace( return ax +def _evolution_t_indices(n: int, n_panels: int) -> list[int]: + """Time indices ``plot_phasespace_evolution`` samples, plus the final slice. + + Mirrors the panel subsampling (``range(0, n, n // n_panels)``) and always + includes ``n - 1``, so a series loaded with exactly these indices serves + both the evolution facets and the final-step heatmap. + """ + t_skip = max(1, n // n_panels) + return sorted(set(range(0, n, t_skip)) | {n - 1}) + + def plot_phasespace_evolution( series: xr.DataArray | str | Path, *, @@ -277,15 +288,25 @@ def plot_phasespace_evolution( Subsamples a stacked ``(t, p, x)`` phase-space series (as returned by ``io.load_series``) to ~``n_panels`` snapshots so the time evolution is visible, rather than only the final step. Returns the ``Figure``. + + Given a path, only the sampled snapshots are read from disk + (``io.load_series`` with ``t_indices``) — a production phase-space history + can run to tens of GB while the facets need under a dozen slices. The + colour scale is likewise computed from the sampled panels only, never the + full history. """ - da = series if isinstance(series, xr.DataArray) else _io.load_series(series) + if isinstance(series, xr.DataArray): + da = series + else: + da = _io.load_series(series, t_indices=_evolution_t_indices(_io.series_len(series), n_panels)) da = _decorate(da) if da.ndim != 3: raise ValueError(f"plot_phasespace_evolution expects (t, p, x); got dims {da.dims}") da = _crop_spatial_to_box(da) nt = da.coords["t"].size - t_skip = max(1, nt // n_panels) - idx = list(range(0, nt, t_skip)) + # Same sampling rule whether da is a full in-memory series or was already + # slice-selected off disk: stride panels plus, always, the final slice. + idx = _evolution_t_indices(nt, n_panels) sl = da.isel(t=idx) # Convention: spatial axis horizontal, momentum axis vertical (fall back to # dim order for a momentum-momentum space with no spatial axis). @@ -1050,6 +1071,7 @@ def plot_profile( *, abs_value: bool = False, n_avg_frac: float = 0.2, + avg_window: int | None = None, show_initial: bool = False, value_label: str | None = None, title: str | None = None, @@ -1060,6 +1082,11 @@ def plot_profile( out the time-dependent fluctuations to show the established profile. With ``show_initial`` the ``t = 0`` profile is overlaid too, so the change from the initial state is visible. + + ``avg_window`` overrides the fraction with an explicit number of trailing + dumps. Pass it when ``series`` was slice-selected out of a longer history + (initial + the mean window), where a fraction of the *received* length no + longer means what the caller intended. """ da = _decorate(series if isinstance(series, xr.DataArray) else _io.load_series(series)) if "t" not in da.dims: @@ -1068,7 +1095,8 @@ def plot_profile( x = da.coords[xdim].values tvals = da.coords["t"].values nt = da.sizes["t"] - w = max(1, round(n_avg_frac * nt)) + w = int(avg_window) if avg_window is not None else max(1, round(n_avg_frac * nt)) + w = max(1, min(w, nt)) if ax is None: _, ax = plt.subplots(figsize=(6, 4)) final = da.isel(t=-1).values @@ -1438,15 +1466,26 @@ def _write(fig: plt.Figure, rel: str) -> Path: try: # Temperature from a Maxwellian fit to the species phase space # (preferred); fall back to dumped thermal moments (uth / T_ii). + # The fit builds several full-size temporaries of the (t, p, x) + # series, so read only the slices plot_profile uses — t=0 plus the + # trailing mean window — instead of the whole history. ps_path = _species_phasespace(diags, species) - temp = _temperature_from_phasespace(ps_path) + temp = avg_window = None + if ps_path is not None: + n_t = _io.series_len(ps_path) + if n_t > 0: + avg_window = max(1, round(0.2 * n_t)) # plot_profile's default n_avg_frac + idx = sorted({0, *range(n_t - avg_window, n_t)}) + temp = _temperature_from_phasespace(_io.load_series(ps_path, t_indices=idx)) if temp is None: temp = _temperature_series(entries) + avg_window = None # full series: the default fraction applies if temp is not None: fig, ax = plt.subplots(figsize=(6, 4)) plot_profile( temp, ax=ax, + avg_window=avg_window, show_initial=True, title=f"{label} — temperature profile", ) @@ -1463,7 +1502,15 @@ def _write(fig: plt.Figure, rel: str) -> Path: continue ps_name, species = parts[1], parts[2] try: - ser = _io.load_series(diag_path) + # Slice-selected read: only the evolution panels + the final step + # come off disk. A production phase-space history (e.g. x1g_q1 at + # ~12k dumps) is tens of GB; the plots below use under a dozen + # slices of it, and eager-loading it here is what blew node memory + # when many runs post-processed at once. + n_t = _io.series_len(diag_path) + if n_t < 1: + continue + ser = _io.load_series(diag_path, t_indices=_evolution_t_indices(n_t, n_panels)) except Exception as e: print(f"[plots] skipping phasespace {diag_rel}: {e}") continue @@ -1617,8 +1664,14 @@ def _decorate(da: xr.DataArray) -> xr.DataArray: ``io.load_series`` stashes axis metadata in dict-valued DataArray attrs; copying it onto the coordinates lets xarray's faceted plotters auto-label each panel's axes. Returns a decorated copy (the input is left untouched). + + The copy is **shallow** (``deep=False``): xarray still gives the copy its own + attrs dicts (so the labels below don't leak onto the caller's array), but the + data buffer is shared — a deep copy would pull a lazily-opened + (:func:`io.open_series`) multi-GB phase-space history fully into memory just + to relabel its axes, which is exactly the load the lazy path exists to avoid. """ - da = da.copy() + da = da.copy(deep=False) axis_units = da.attrs.get("axis_units", {}) or {} axis_long = da.attrs.get("axis_long_names", {}) or {} for dim in da.dims: diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md index af3fedab..2db39c43 100644 --- a/docs/osiris-adept-usage.md +++ b/docs/osiris-adept-usage.md @@ -21,6 +21,7 @@ adept/ (repo root) │ ├── 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/ @@ -313,6 +314,8 @@ These canned plots focus on 1D simulations. > **SRS-specific plots & metrics live in osiris-lpi.** The laser-energy budget (reflected / transmitted / absorbed over time: the `energy_budget.png` + `laser_energy_budget.txt` artifacts and the `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac` metrics) 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 / `δf` views) are **not** produced by adept's general OSIRIS wrapper. They live in the [osiris-lpi](https://github.com/ergodicio/osiris-lpi) repo — `osiris_lpi.OsirisLPI` subclasses `BaseOsiris` and adds them in `post_process`, reading the drive `antenna.a0`/`antenna.omega0` from the deck. Regenerate them offline from saved NetCDFs with `python -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_q1` is ~48 GB in memory (11705 × 1000 × 1024), ~0.19 GB gzipped on disk (sparse, ~250×) — so `io.load_series(path)` decompresses all of it (at 16 sims/node this OOM'd production). `save_canned_plots` therefore 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 before `pcolormesh`. A new plotter over a large series must do the same — `load_series(..., t_indices=…)` for a few slices, or `with 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 stray `da.values` / `da.sum()` over `t`, or a **deep** `da.copy()`, pulls the full series into RAM (`_decorate` copies shallow for this reason). + ## Programmatic use ```python From 6b04aba8fda7f5c345edf8fe0a49bad67b49a75f Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 14 Jul 2026 17:06:37 -0700 Subject: [PATCH 37/49] osiris: add laser drive in ICF/SI units to units.yaml MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit write_units now derives w_laser (rad/s), laser_wavelength (nm), laser_a0, and laser_intensity — the peak intensity of a linearly polarized drive in W/cm^2 (eps0 c E0^2/2 with E0 = a0 m_e c w_laser/e, i.e. the ICF convention I * lam_um^2 = 1.37e18 * a0^2) — from the deck's antenna / zpulse_speckle / zpulse section and the reference wp0. Laser-less decks are unchanged. Co-Authored-By: Claude Fable 5 --- adept/osiris/base.py | 31 ++++++++++++++++++++++++++++++- tests/test_osiris/test_units.py | 23 +++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/adept/osiris/base.py b/adept/osiris/base.py index 833cde83..903031c5 100644 --- a/adept/osiris/base.py +++ b/adept/osiris/base.py @@ -4,8 +4,10 @@ from typing import Any +import math + from adept._base_ import ADEPTModule -from adept.normalization import skin_depth_normalization, skin_depth_normalization_from_frequency +from adept.normalization import UREG, skin_depth_normalization, skin_depth_normalization_from_frequency from adept.osiris import deck as _deck from adept.osiris import density as _density from adept.osiris import post as _post @@ -62,6 +64,13 @@ def write_units(self) -> dict: reference temperature (species carry their own per-species thermal momenta), so the temperature-dependent keys (``T0``/``nuee``/ ``logLambda_ee``) are intentionally omitted. + + When the deck launches a laser (an ``antenna`` / ``zpulse_speckle`` / + ``zpulse`` section with ``a0`` and ``omega0``), the physical drive + scales are added too: ``w_laser`` (rad/s), ``laser_wavelength`` (nm), + ``laser_a0`` and ``laser_intensity`` — the peak intensity of a + linearly polarized drive in W/cm^2 (the ICF convention + ``I * lambda_um^2 = 1.37e18 * a0^2``). """ sim = self._iter_first_section("simulation") n0 = sim.get("n0") @@ -95,6 +104,26 @@ def write_units(self) -> dict: tmin = time.get("tmin", 0.0) quants["sim_duration"] = ((float(tmax) - float(tmin)) * norm.tau).to("ps") + # Laser drive in physical / ICF units. The deck's laser section carries + # a0 and omega0 (in wp0 units); with wp0 fixed above these give the + # laser frequency, the vacuum wavelength lambda = 2 pi c / w_laser, and + # the peak intensity of a linearly polarized drive, + # I = eps0 c E0^2 / 2 with E0 = a0 m_e c w_laser / e — equivalently the + # usual ICF convention I[W/cm^2] * lambda[um]^2 = 1.37e18 * a0^2. + # Same section priority as the SRS postproc: antenna (1D), then + # zpulse_speckle / zpulse (2D). + for sec_name in ("antenna", "zpulse_speckle", "zpulse"): + laser = self._iter_first_section(sec_name) + if laser.get("a0") is not None and laser.get("omega0") is not None: + a0, omega0 = float(laser["a0"]), float(laser["omega0"]) + w_laser = omega0 / norm.tau + e_peak = a0 * w_laser * UREG.m_e * UREG.c / UREG.e + quants["w_laser"] = w_laser.to("rad/s") + quants["laser_wavelength"] = (2.0 * math.pi * UREG.c / w_laser).to("nm") + quants["laser_a0"] = a0 + quants["laser_intensity"] = (e_peak**2 * UREG.epsilon_0 * UREG.c / 2.0).to("W/cm^2") + break + self.cfg.setdefault("units", {})["derived"] = quants return quants diff --git a/tests/test_osiris/test_units.py b/tests/test_osiris/test_units.py index 252cd40c..229fbfe3 100644 --- a/tests/test_osiris/test_units.py +++ b/tests/test_osiris/test_units.py @@ -36,6 +36,29 @@ def test_write_units_density_derived_scales() -> None: assert quants["sim_duration"].to("ps").magnitude == pytest.approx(4.6583, rel=1e-3) +def test_write_units_laser_intensity_in_icf_units() -> None: + # SRS deck: antenna{a0=0.004, omega0=1.0} with n0 = 9.05e21 cm^-3 = + # n_crit(351 nm), so w_laser = wp0 and lambda = 351 nm. The peak intensity + # of a linearly polarized drive follows the ICF convention + # I[W/cm^2] * lambda[um]^2 = 1.37e18 * a0^2. + quants = BaseOsiris({"osiris": {"deck": str(SRS_DECK)}}).write_units() + assert quants["w_laser"].to("rad/s").magnitude == pytest.approx(5.3668e15, rel=1e-3) + assert quants["laser_wavelength"].to("nm").magnitude == pytest.approx(351.0, rel=1e-3) + assert quants["laser_a0"] == pytest.approx(0.004) + lam_um = quants["laser_wavelength"].to("um").magnitude + assert quants["laser_intensity"].to("W/cm^2").magnitude == pytest.approx( + 1.37e18 * 0.004**2 / lam_um**2, rel=1e-2 + ) + + +def test_write_units_no_laser_keys_without_laser_section(tmp_path: Path) -> None: + deck = tmp_path / "no_laser" + deck.write_text("simulation { n0 = 9.05e21, }\n") + quants = BaseOsiris({"osiris": {"deck": str(deck)}}).write_units() + for key in ("w_laser", "laser_wavelength", "laser_a0", "laser_intensity"): + assert key not in quants + + def test_write_units_omits_temperature_keys() -> None: quants = BaseOsiris({"osiris": {"deck": str(SRS_DECK)}}).write_units() # OSIRIS has no single global temperature, so temperature-dependent From 86341b9e32a0f6d7f37c2bcebd3f98af1c628246 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:09 -0700 Subject: [PATCH 38/49] fix: adopt v0 = sqrt(T0/m_e) in electron_debye_normalization The engine convention for _vlasov1d (and _vlasov2d, _pic1d) is the RMS/standard-deviation thermal speed: Maxwellian exp(-v^2/2) at T=1, L0 = lambda_De, code wavenumber k*lambda_De. normalization.py:91 was the sole outlier with v0 = sqrt(2 T0/m_e), which made every logged physical unit (v0, x0, c_light, box_length) wrong by sqrt(2), every dimensional string input (um box sizes, gradient scale lengths, laser k0) off by sqrt(2), and a 'normalizing_temperature: 2000eV' run physically a 4000 eV plasma. Dimensionless numeric-input dynamics were unaffected. Also fixes a sign error in the NRL Coulomb logarithm (log(n^0.5 / T^-1.25) == log(n^0.5 * T^+1.25)), which produced logLambda_ee = -11.8 and a negative logged nuee at 2000 eV / 1.5e21cc; the correct values are +7.22 and a positive rate. vth_norm() is deliberately untouched: its only callers are in vfp1d, which is self-consistently built on the sqrt(2T/m) convention. Adds tests/test_vlasov1d/test_units_boundary.py: dimensional input in, dimensional quantity out, checked against CODATA constants and the NRL formulary - the units-boundary test class whose absence let this bug survive every existing suite. See VLASOV1D_CONVENTIONS_AUDIT.md (F1, F2) and ADEPT_CONVENTIONS_AUDIT.md. Co-Authored-By: Claude Fable 5 --- adept/normalization.py | 14 ++-- tests/test_vlasov1d/test_units_boundary.py | 74 ++++++++++++++++++++++ 2 files changed, 83 insertions(+), 5 deletions(-) create mode 100644 tests/test_vlasov1d/test_units_boundary.py diff --git a/adept/normalization.py b/adept/normalization.py index aa9f7714..ae23bbc2 100644 --- a/adept/normalization.py +++ b/adept/normalization.py @@ -34,7 +34,7 @@ class PlasmaNormalization: def logLambda_ee(self) -> float: n0_cc = self.n0.to("1/cc").magnitude T0_eV = self.T0.to("eV").magnitude - logLambda_ee = 23.5 - jnp.log(n0_cc**0.5 / T0_eV**-1.25) + logLambda_ee = 23.5 - jnp.log(n0_cc**0.5 * T0_eV**-1.25) logLambda_ee -= (1e-5 + (jnp.log(T0_eV) - 2) ** 2.0 / 16) ** 0.5 return logLambda_ee @@ -78,9 +78,13 @@ def electron_debye_normalization(n0_str, T0_str): """ Returns the electron thermal normalization for the given density and temperature. Unit quantities are: - - Debye length - - Electron thermal velocity - - Langmuir oscillation frequency + - L0 = Debye length, sqrt(eps0 T0 / (n0 e^2)) + - v0 = electron thermal velocity sqrt(T0/m_e) (RMS / standard-deviation + convention: a T=1 Maxwellian is exp(-v^2/2) with unit variance) + - tau = 1/wp0 (inverse Langmuir oscillation frequency) + + Under this convention a code wavenumber is k*lambda_De and the Bohm-Gross + dispersion reads w^2 = 1 + 3 k^2. """ n0 = UREG.Quantity(n0_str) T0 = UREG.Quantity(T0_str) @@ -88,7 +92,7 @@ def electron_debye_normalization(n0_str, T0_str): wp0 = ((n0 * UREG.e**2.0 / (UREG.m_e * UREG.epsilon_0)) ** 0.5).to("rad/s") tau = 1 / wp0 - v0 = ((2.0 * T0 / UREG.m_e) ** 0.5).to("m/s") + v0 = ((T0 / UREG.m_e) ** 0.5).to("m/s") x0 = (v0 / wp0).to("nm") return PlasmaNormalization(m0=UREG.m_e, q0=UREG.e, n0=n0, T0=T0, L0=x0, v0=v0, tau=tau) diff --git a/tests/test_vlasov1d/test_units_boundary.py b/tests/test_vlasov1d/test_units_boundary.py new file mode 100644 index 00000000..e7af1c38 --- /dev/null +++ b/tests/test_vlasov1d/test_units_boundary.py @@ -0,0 +1,74 @@ +# Copyright (c) Ergodic LLC 2026 +# research@ergodic.io +""" +Units-boundary regression tests for the Vlasov-1D normalization layer. + +Every physics test in this suite works in dimensionless code units, which +validates the engine but not the dimensional dictionary: numeric inputs bypass +``normalize()`` entirely. These tests cross the units boundary — dimensional +input in, dimensional quantity out — and compare against *independent* physical +references (CODATA constants via scipy, and the NRL formulary), so a convention +error in ``normalization.py`` cannot hide. + +The engine convention under test: v0 = sqrt(T0/m_e) (RMS / standard-deviation +thermal speed), L0 = v0/wp0 = lambda_De, Maxwellian exp(-v^2/2) at T=1. +""" + +import numpy as np +from scipy import constants as csts + +from adept.normalization import electron_debye_normalization, normalize + +N0_STR, T0_STR = "1.5e21/cc", "2000eV" +N0_SI = 1.5e21 * 1e6 # 1/m^3 +T0_J = 2000.0 * csts.e + + +def test_debye_normalization_against_codata(): + """v0, L0, tau, and c_hat must match hand-computed CODATA values.""" + norm = electron_debye_normalization(N0_STR, T0_STR) + + v_th = np.sqrt(T0_J / csts.m_e) # sigma convention: sqrt(T/m) + wp0 = np.sqrt(N0_SI * csts.e**2 / (csts.epsilon_0 * csts.m_e)) + lambda_de = v_th / wp0 + + np.testing.assert_allclose(norm.v0.to("m/s").magnitude, v_th, rtol=1e-6) + np.testing.assert_allclose(norm.L0.to("m").magnitude, lambda_de, rtol=1e-6) + np.testing.assert_allclose(norm.tau.to("s").magnitude, 1.0 / wp0, rtol=1e-6) + np.testing.assert_allclose(norm.speed_of_light_norm(), csts.c / v_th, rtol=1e-6) + + +def test_debye_length_against_nrl_formulary(): + """lambda_De = 7.43e2 sqrt(T_eV/n_cc) cm (NRL formulary) — a truly external reference.""" + norm = electron_debye_normalization(N0_STR, T0_STR) + lambda_de_nrl_m = 7.43e2 * np.sqrt(2000.0 / 1.5e21) * 1e-2 + np.testing.assert_allclose(norm.L0.to("m").magnitude, lambda_de_nrl_m, rtol=1e-3) + + +def test_dimensional_string_round_trip(): + """String inputs must convert with L0 = lambda_De (not sqrt(2) lambda_De).""" + norm = electron_debye_normalization(N0_STR, T0_STR) + + v_th = np.sqrt(T0_J / csts.m_e) + wp0 = np.sqrt(N0_SI * csts.e**2 / (csts.epsilon_0 * csts.m_e)) + lambda_de = v_th / wp0 + + # x: 100 um -> 100e-6 / lambda_De code units + np.testing.assert_allclose(normalize("100um", norm, dim="x"), 100e-6 / lambda_de, rtol=1e-6) + # k: 1/um -> lambda_De / 1e-6 code units (k lambda_De) + np.testing.assert_allclose(normalize("1/um", norm, dim="k"), lambda_de / 1e-6, rtol=1e-6) + # t: 1 ps -> wp0 * 1e-12 + np.testing.assert_allclose(normalize("1ps", norm, dim="t"), wp0 * 1e-12, rtol=1e-6) + # numeric inputs pass through untouched + assert normalize(3.25, norm, dim="x") == 3.25 + + +def test_logged_collision_frequency_is_physical(): + """logLambda and nuee must be positive and match the NRL expression.""" + norm = electron_debye_normalization(N0_STR, T0_STR) + log_lambda = float(norm.logLambda_ee()) + # NRL: 23.5 - ln(n^1/2 T^-5/4) - [1e-5 + (ln T - 2)^2/16]^1/2 + expected = 23.5 - np.log(np.sqrt(1.5e21) * 2000.0**-1.25) - np.sqrt(1e-5 + (np.log(2000.0) - 2.0) ** 2 / 16.0) + np.testing.assert_allclose(log_lambda, expected, rtol=1e-10) + assert log_lambda > 0 + assert norm.approximate_ee_collision_frequency().to("Hz").magnitude > 0 From 2a9d687e9406a28b7ef3e904b9f43e5a29e4f0e5 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:31 -0700 Subject: [PATCH 39/49] fix(tf1d): remove private sqrt(2) duplicate and logLambda sign error in write_units _tf1d re-implements the debye normalization inline; it carried the same v0 = sqrt(2T/m) outlier (its engine, like _vlasov1d, runs in sqrt(T/m) units) and the same log(n^0.5 / T^-1.25) sign error. Diagnostics-only: the logged units were wrong; dynamics are unaffected (grid beta is written but never consumed). Co-Authored-By: Claude Fable 5 --- adept/_tf1d/modules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/adept/_tf1d/modules.py b/adept/_tf1d/modules.py index bdd53b4e..9816785d 100644 --- a/adept/_tf1d/modules.py +++ b/adept/_tf1d/modules.py @@ -59,7 +59,7 @@ def write_units(self): wp0 = np.sqrt(n0 * self.ureg.e**2.0 / (self.ureg.m_e * self.ureg.epsilon_0)).to("rad/s") tp0 = (1 / wp0).to("fs") - v0 = np.sqrt(2.0 * T0 / self.ureg.m_e).to("m/s") + v0 = np.sqrt(T0 / self.ureg.m_e).to("m/s") x0 = (v0 / wp0).to("nm") c_light = _Q(1.0 * self.ureg.c).to("m/s") / v0 beta = (v0 / self.ureg.c).to("dimensionless") @@ -72,7 +72,7 @@ def write_units(self): sim_duration = (self.cfg["grid"]["tmax"] * tp0).to("ps") # collisions - logLambda_ee = 23.5 - np.log(n0.magnitude**0.5 / T0.magnitude**-1.25) + logLambda_ee = 23.5 - np.log(n0.magnitude**0.5 * T0.magnitude**-1.25) logLambda_ee -= (1e-5 + (np.log(T0.magnitude) - 2) ** 2.0 / 16) ** 0.5 nuee = _Q(2.91e-6 * n0.magnitude * logLambda_ee / T0.magnitude**1.5, "Hz") nuee_norm = nuee / wp0 From c5563478a68953b059209141a6eb9de06ba9ff80 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:32 -0700 Subject: [PATCH 40/49] fix(vlasov1d): Dougherty vbar missing 1/n; Krook target uses species T0/mass - Dougherty.compute_vbar returned the raw first moment n*u while its callers (find_self_consistent_beta, the drag term) treat it as a mean velocity. The drag therefore centered on n*u and momentum was not conserved wherever n(x) != 1. Now returns sum(v f)/sum(f). - The Krook target Maxwellian was hard-coded to exp(-v^2/2) (T=1, m=1): any species with T0 != 1 was dragged toward T=1. It is now built with variance T0/mass from the species' bulk parameters, which are threaded through cfg.grid.species_params (new T0 entry). - New test drives the collision operator in isolation with a 50% density modulation and a finite drift: per-cell density, momentum, and energy conservation, and the relaxed mean velocity must equal u0 (the old operator drifts it toward n*u0). Audit findings F4, F5. Co-Authored-By: Claude Fable 5 --- adept/_vlasov1d/modules.py | 6 ++ .../solvers/pushers/fokker_planck.py | 10 ++- .../test_fp_momentum_conservation.py | 83 +++++++++++++++++++ 3 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 tests/test_vlasov1d/test_fp_momentum_conservation.py diff --git a/adept/_vlasov1d/modules.py b/adept/_vlasov1d/modules.py index fe2c52d3..40dcd090 100644 --- a/adept/_vlasov1d/modules.py +++ b/adept/_vlasov1d/modules.py @@ -117,6 +117,9 @@ def write_units(self) -> dict: sim_duration = (grid.tmax * norm.tau).to("ps") nu_ee = norm.approximate_ee_collision_frequency() + # e-e collision rate in code units (1/wp0). This is what the FP/Krook + # `baseline` rates should be compared against. + nu_ee_norm = (nu_ee * norm.tau).to("").magnitude beta = 1.0 / norm.speed_of_light_norm() @@ -130,6 +133,7 @@ def write_units(self) -> dict: "beta": beta, "x0": norm.L0.to("nm"), "nuee": nu_ee.to("Hz"), + "nuee_norm": nu_ee_norm, "logLambda_ee": norm.logLambda_ee(), "box_length": box_length, "box_width": box_width, @@ -235,10 +239,12 @@ def get_solver_quantities(self) -> dict: cfg_grid["species_grids"][species_name]["one_over_kvr"] = jnp.array(one_over_kvr) # Build species parameters (charge, mass, charge-to-mass ratio) + # T0 is the bulk (first-listed component) temperature, used e.g. by the Krook target cfg_grid["species_params"][species_name] = { "charge": species_cfg.charge, "mass": species_cfg.mass, "charge_to_mass": species_cfg.charge / species_cfg.mass, + "T0": self.simulation.species_distributions[species_name][0].T0, } cfg_grid["n_prof_total"] = n_prof_total diff --git a/adept/_vlasov1d/solvers/pushers/fokker_planck.py b/adept/_vlasov1d/solvers/pushers/fokker_planck.py index 23166bfd..dfb2a5b1 100644 --- a/adept/_vlasov1d/solvers/pushers/fokker_planck.py +++ b/adept/_vlasov1d/solvers/pushers/fokker_planck.py @@ -76,9 +76,9 @@ def compute_vbar(self, f: Array) -> Array: f: Distribution function, shape (..., nv) Returns: - vbar: Mean velocity , shape (...) + vbar: Mean velocity = ∫v f dv / ∫f dv, shape (...) """ - return jnp.sum(f * self.v, axis=-1) * self.dv + return jnp.sum(f * self.v, axis=-1) / jnp.sum(f, axis=-1) def compute_D(self, f: Array, beta: Array) -> Array: """ @@ -242,7 +242,11 @@ def __init__(self, cfg: Mapping[str, Any]): self.cfg = cfg v = cfg["grid"]["species_grids"]["electron"]["v"] dv = cfg["grid"]["species_grids"]["electron"]["dv"] - f_mx = np.exp(-(v[None, :] ** 2.0) / 2.0) + params = cfg["grid"].get("species_params", {}).get("electron", {}) + T0 = params.get("T0", 1.0) + mass = params.get("mass", 1.0) + # Maxwellian with variance T0/m (the species' bulk thermal width) + f_mx = np.exp(-(v[None, :] ** 2.0) / (2.0 * T0 / mass)) self.f_mx = f_mx / np.sum(f_mx, axis=1)[:, None] / dv self.dv = dv diff --git a/tests/test_vlasov1d/test_fp_momentum_conservation.py b/tests/test_vlasov1d/test_fp_momentum_conservation.py new file mode 100644 index 00000000..a50d9f66 --- /dev/null +++ b/tests/test_vlasov1d/test_fp_momentum_conservation.py @@ -0,0 +1,83 @@ +# Copyright (c) Ergodic LLC 2026 +# research@ergodic.io +""" +Momentum conservation of the Dougherty operator where n(x) != 1. + +The Dougherty drag must center on the *mean velocity* u(x) = ∫v f dv / n(x). +A drag centered on the raw first moment ∫v f dv = n*u (the historical bug) +breaks momentum conservation exactly where the density deviates from 1, e.g. +for bump-on-tail or large-amplitude density perturbations with collisions on. + +This test applies the collision operator in isolation (no advection, no +fields) to a drifting Maxwellian with a strong density modulation and checks +per-cell conservation of density, momentum, and energy. +""" + +import jax + +jax.config.update("jax_enable_x64", True) + +import numpy as np +import pytest +from jax import numpy as jnp + +from adept._vlasov1d.solvers.pushers.fokker_planck import Collisions + + +def _make_cfg(nx, nv, vmax, fp_type): + dv = 2.0 * vmax / nv + v = jnp.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) + return { + "grid": { + "species_grids": {"electron": {"v": v, "dv": dv, "nv": nv, "vmax": vmax}}, + "species_params": {"electron": {"charge": -1.0, "mass": 1.0, "charge_to_mass": -1.0, "T0": 1.0}}, + }, + "terms": { + "fokker_planck": {"is_on": True, "type": fp_type}, + "krook": {"is_on": False}, + }, + } + + +# Energy tolerance is per-scheme: Chang-Cooper preserves the discrete equilibrium +# to ~1e-6, plain central differencing only to ~1e-3 at this resolution. +@pytest.mark.parametrize("fp_type,energy_rtol", [("dougherty", 5e-3), ("chang_cooper_dougherty", 1e-5)]) +def test_dougherty_momentum_conservation_with_density_perturbation(fp_type, energy_rtol): + nx, nv, vmax = 16, 512, 6.4 + cfg = _make_cfg(nx, nv, vmax, fp_type) + v = np.array(cfg["grid"]["species_grids"]["electron"]["v"]) + dv = cfg["grid"]["species_grids"]["electron"]["dv"] + + # Strong density modulation and a finite drift: exactly the regime where + # centering the drag on n*u instead of u destroys momentum conservation. + x = np.linspace(0, 2 * np.pi, nx, endpoint=False) + n_prof = 1.0 + 0.5 * np.sin(x) + u0 = 0.5 + f = n_prof[:, None] * np.exp(-((v[None, :] - u0) ** 2) / 2.0) + f = f / (np.sum(np.exp(-((v - u0) ** 2) / 2.0)) * dv) + + f = jnp.array(f) + collisions = Collisions(cfg) + + nu = jnp.ones(nx) + dt = 0.1 + f_out = f + for _ in range(50): + f_out = collisions(nu, None, f_out, dt) + + n0_ = np.sum(np.array(f), axis=1) * dv + n1 = np.sum(np.array(f_out), axis=1) * dv + p0 = np.sum(np.array(f) * v[None, :], axis=1) * dv + p1 = np.sum(np.array(f_out) * v[None, :], axis=1) * dv + e0 = np.sum(np.array(f) * v[None, :] ** 2, axis=1) * dv + e1 = np.sum(np.array(f_out) * v[None, :] ** 2, axis=1) * dv + + # The historical n*u-centered drag gives O(50%) momentum errors here; + # the corrected operator conserves to discretization level (~1e-7). + np.testing.assert_allclose(n1, n0_, rtol=1e-10, err_msg="density not conserved") + np.testing.assert_allclose(p1, p0, rtol=1e-6, err_msg="momentum not conserved where n(x) != 1") + np.testing.assert_allclose(e1, e0, rtol=energy_rtol, err_msg="energy not conserved") + + # And the distribution must still be centered on u0, not n(x)*u0 + u_final = p1 / n1 + np.testing.assert_allclose(u_final, u0, atol=1e-6) From 1ea7244af6d918c725ae0df6ffac5dd2c84b15d4 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:32 -0700 Subject: [PATCH 41/49] fix(vlasov1d): center storage moments on u not n*u; entropy sign; energy monitor - Field moments p and q were centered on the raw first moment n*u; they are now centered on the true mean velocity. The first moment is saved under the honest name 'j' and 'v' is now j/n (previously 'v' held the flux while the same integrand was called 'mean_j' in the scalars). - The field-moment '-flogf' computed +int f log|f| dv, the exact negative of the scalar of the same name; both now agree. - Adds mean_kinetic_energy / mean_field_energy / mean_total_energy scalars (electrostatic energy monitor, with the 1/2 factors). A conservation monitor of this kind would have caught the sqrt(2) convention bug far earlier. Audit findings F6, F7, F12.2. Co-Authored-By: Claude Fable 5 --- adept/_vlasov1d/storage.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/adept/_vlasov1d/storage.py b/adept/_vlasov1d/storage.py index 1075ed48..6899e049 100644 --- a/adept/_vlasov1d/storage.py +++ b/adept/_vlasov1d/storage.py @@ -137,11 +137,12 @@ def _calc_moment_(inp, _dv=dv): f = y[species_name] species_moments = {} species_moments["n"] = _calc_moment_(f) - species_moments["v"] = _calc_moment_(f * v[None, :]) + species_moments["j"] = _calc_moment_(f * v[None, :]) + species_moments["v"] = species_moments["j"] / species_moments["n"] v_m_vbar = v[None, :] - species_moments["v"][:, None] species_moments["p"] = _calc_moment_(f * v_m_vbar**2.0) species_moments["q"] = _calc_moment_(f * v_m_vbar**3.0) - species_moments["-flogf"] = _calc_moment_(f * jnp.log(jnp.abs(f))) + species_moments["-flogf"] = _calc_moment_(-jnp.abs(f) * jnp.log(jnp.abs(f))) species_moments["f^2"] = _calc_moment_(f * f) result[species_name] = species_moments @@ -292,9 +293,11 @@ def save(t, y, args): scalars = {} # Compute scalars for each species + mean_kinetic_energy = 0.0 for species_name in species_names: v = species_grids[species_name]["v"][None, :] dv = species_grids[species_name]["dv"] + mass = cfg["grid"]["species_params"][species_name]["mass"] def _calc_mean_moment_(inp, _dv=dv): return jnp.mean(jnp.sum(inp, axis=1) * _dv) @@ -306,12 +309,19 @@ def _calc_mean_moment_(inp, _dv=dv): scalars[f"mean_q_{species_name}"] = _calc_mean_moment_(f * v**3.0) scalars[f"mean_-flogf_{species_name}"] = _calc_mean_moment_(-jnp.log(jnp.abs(f)) * jnp.abs(f)) scalars[f"mean_f2_{species_name}"] = _calc_mean_moment_(f * f) + mean_kinetic_energy += 0.5 * mass * scalars[f"mean_P_{species_name}"] # Shared field scalars (not species-specific) scalars["mean_de2"] = jnp.mean(y["de"] ** 2.0) scalars["mean_e2"] = jnp.mean(y["e"] ** 2.0) scalars["mean_pond"] = jnp.mean(-0.5 * jnp.gradient(y["a"] ** 2.0, cfg["grid"]["dx"])[1:-1]) + # Energy conservation monitor (electrostatic): x-averaged kinetic + E-field + # energy density. Transverse (EM) field energy is not included. + scalars["mean_kinetic_energy"] = mean_kinetic_energy + scalars["mean_field_energy"] = 0.5 * scalars["mean_e2"] + scalars["mean_total_energy"] = mean_kinetic_energy + 0.5 * scalars["mean_e2"] + return scalars return save From 070fb044cf994da3b1b77d562a597e5821a7bf5f Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:32 -0700 Subject: [PATCH 42/49] fix(vlasov1d): dx and semi-Lagrangian period ignored xmin dx was xmax/nx and the interpolation period was xmax; both are now the box length (xmax - xmin). Latent for all shipped configs (xmin: 0.0) but wrong spacing, kx grid, and advection wrap for any xmin != 0. Audit finding F11. Co-Authored-By: Claude Fable 5 --- adept/_vlasov1d/grid.py | 2 +- adept/_vlasov1d/solvers/pushers/vlasov.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/adept/_vlasov1d/grid.py b/adept/_vlasov1d/grid.py index 91b6d09b..b16db67e 100644 --- a/adept/_vlasov1d/grid.py +++ b/adept/_vlasov1d/grid.py @@ -52,7 +52,7 @@ def __init__( self.tmin = tmin # Compute dx - self.dx = xmax / nx + self.dx = (xmax - xmin) / nx # Override dt for EM wave stability if needed if should_override_dt_for_em_waves: diff --git a/adept/_vlasov1d/solvers/pushers/vlasov.py b/adept/_vlasov1d/solvers/pushers/vlasov.py index 200da317..4a51ec4c 100644 --- a/adept/_vlasov1d/solvers/pushers/vlasov.py +++ b/adept/_vlasov1d/solvers/pushers/vlasov.py @@ -28,7 +28,7 @@ def __init__(self, cfg, interp_e): """Create interpolation helpers for x-advection and v-advection steps.""" self.x = cfg["grid"]["x"] self.v = cfg["grid"]["v"] - self.f_interp = partial(interp2d, x=self.x, y=self.v, period=cfg["grid"]["xmax"]) + self.f_interp = partial(interp2d, x=self.x, y=self.v, period=cfg["grid"]["xmax"] - cfg["grid"]["xmin"]) self.dt = cfg["grid"]["dt"] self.dummy_x = jnp.ones_like(self.x) self.dummy_v = jnp.ones_like(self.v)[None, :] From 0b626bb2529e94a74459425be654f69d20f320ca Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:32 -0700 Subject: [PATCH 43/49] chore(vlasov1d): remove dead c_light knob; fix AmpereSolver docstring GridConfig.c_light (set only by wavepacket.yaml) was never read - beta/c_light are always derived from the normalization. The AmpereSolver class docstring claimed j = sum_s (q_s/m_s) int v f dv; the code correctly has no 1/m_s. Audit findings F12.1, F12.4. Co-Authored-By: Claude Fable 5 --- adept/_vlasov1d/datamodel.py | 1 - adept/_vlasov1d/solvers/pushers/field.py | 2 +- configs/vlasov-1d/wavepacket.yaml | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 2e96c1ca..811df767 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -70,7 +70,6 @@ class GridConfig(BaseModel): nv: int | None = None vmax: float | None = None parallel: tuple[str, ...] | bool = False - c_light: float | None = None @field_validator("parallel", mode="before") @classmethod diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index 76a999af..ff73262c 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -227,7 +227,7 @@ def __call__(self, f_dict: dict, prev_ex: jnp.ndarray, dt: jnp.float64): class AmpereSolver: """Ampere solver using current density to evolve electric field. - Solves: a single Euler step of ∂E/∂t = -j, where j = Σ_s (q_s/m_s) ∫v f_s dv + Solves: a single Euler step of ∂E/∂t = -j, where j = Σ_s q_s ∫v f_s dv is the total current density from all species. """ diff --git a/configs/vlasov-1d/wavepacket.yaml b/configs/vlasov-1d/wavepacket.yaml index 3770512f..5ceba970 100644 --- a/configs/vlasov-1d/wavepacket.yaml +++ b/configs/vlasov-1d/wavepacket.yaml @@ -35,7 +35,6 @@ drivers: width: 900.0 ey: {} grid: - c_light: 10 dt: 0.1 nv: 6144 nx: 4096 From 438d4b4050c05041a3988d8e375a452ec759c2fa Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:54 -0700 Subject: [PATCH 44/49] fix(vlasov1d): retune srs-debug-small numeric driver to corrected c_hat The ey driver pair was a backscatter-matched triad under the old c_hat = 11.30 (w_pump - w_seed = 1.16 = old EPW frequency) with placeholder k0 = 1.0 on both drivers. Retuned to a matched triad under the corrected c_hat = 15.984: pump (k0 = +0.162949, w0 = 2.79), seed (k0 = -0.086080, w0 = 1.700942), EPW at k*lambda_De = 0.249. Matching condition and conventions documented in the config. Audit finding F3. Co-Authored-By: Claude Fable 5 --- configs/vlasov-1d/srs-debug-small.yaml | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/configs/vlasov-1d/srs-debug-small.yaml b/configs/vlasov-1d/srs-debug-small.yaml index c0bfa5ec..5ac73e31 100644 --- a/configs/vlasov-1d/srs-debug-small.yaml +++ b/configs/vlasov-1d/srs-debug-small.yaml @@ -48,13 +48,18 @@ mlflow: experiment: vlasov1d-srs run: debug-small +# The ey driver pair is a backscatter-SRS matched triad in code units +# (v0 = sqrt(T0/m_e), k in 1/lambda_De, w in wp0, c_hat = c/v0 = 15.984 at 2000 eV): +# pump: w0 = 2.79, k0 = +sqrt(w0^2 - 1)/c_hat = +0.162949 (rightgoing) +# seed: w0 = 1.700942, k0 = -0.086080 (leftgoing) +# EPW: k lambda_De = 0.249029, w = 1.089058 = w_pump - w_seed drivers: ex: {} ey: '0': params: a0: 1.0 - k0: 1.0 + k0: 0.162949 w0: 2.79 dw0: 0. envelope: @@ -69,8 +74,8 @@ drivers: '1': params: a0: 1.e-6 - k0: 1.0 - w0: 1.63 + k0: -0.086080 + w0: 1.700942 dw0: 0. envelope: time: From cc7fd581141931891ac8b7f2b7ec042b4d83280c Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:54 -0700 Subject: [PATCH 45/49] docs(vlasov1d): state the normalization convention; document unit traps - config.md gains a 'Normalization convention' section: v0 = sqrt(T0/m_e) (sigma convention), L0 = lambda_De, k in 1/lambda_De, Maxwellian exp(-v^2/2) at T=1, Bohm-Gross w^2 = 1 + 3k^2. - Species v0/T0 documented as code-units-only (they bypass normalize()); drift and thermal width now share the same unit. - FP/Krook 'baseline' documented as a rate in units of wp0 (no 2pi), with the newly logged nuee_norm as the reference scale, and the O(1) caveat between the Dougherty nu and the NRL nu_ee. - Super-Gaussian alpha documented (code comment + config.md): it fixes / = 3 T0/mass for all m; the variance equals T0/mass only at m=2 (x1.24 at m=3, x1.37 at m=4). Documentation only, per review. Audit findings F8, F9, F10 (nuee_norm logging itself landed with the species_params change in an earlier commit). Co-Authored-By: Claude Fable 5 --- adept/_vlasov1d/helpers.py | 5 ++++- docs/source/solvers/vlasov1d/config.md | 30 ++++++++++++++++++++++---- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index 553c8311..5a37afcf 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -68,7 +68,10 @@ def _initialize_supergaussian_distribution_( # Thermal velocity: v_t = sqrt(T/m) v_thermal = np.sqrt(T0 / mass) - # Alpha factor for supergaussian normalization + # Alpha factor for supergaussian normalization. This fixes the moment ratio + # / = 3*T0/mass for every order m; the realized *variance* equals + # T0/mass only at m=2 (Maxwellian). For m>2 flat-tops the second-moment + # temperature exceeds T0 (x1.24 at m=3, x1.37 at m=4). See docs config.md. alpha = np.sqrt(3.0 * gamma_3_over_m(supergaussian_order) / gamma_5_over_m(supergaussian_order)) single_dist = -(np.power(np.abs((vax[None, :] - v0) / (alpha * v_thermal)), supergaussian_order)) diff --git a/docs/source/solvers/vlasov1d/config.md b/docs/source/solvers/vlasov1d/config.md index 66e878f7..f092c65e 100644 --- a/docs/source/solvers/vlasov1d/config.md +++ b/docs/source/solvers/vlasov1d/config.md @@ -48,6 +48,28 @@ units: normalizing_density: 1.5e21/cc ``` +### Normalization convention + +All Vlasov-1D quantities are normalized using the following unit system, built from +`normalizing_density` ($n_0$) and `normalizing_temperature` ($T_0$): + +| Unit | Definition | Meaning | +|------|-----------|---------| +| time | $\tau = 1/\omega_{p0}$, $\omega_{p0} = \sqrt{n_0 e^2/(\epsilon_0 m_e)}$ | inverse plasma frequency | +| velocity | $v_0 = \sqrt{T_0/m_e}$ | electron thermal speed (RMS / standard-deviation convention) | +| length | $L_0 = v_0/\omega_{p0} = \lambda_{De}$ | electron Debye length | +| wavenumber | $1/L_0$ | a code wavenumber is $k \lambda_{De}$ | + +Consequences of the $v_0 = \sqrt{T_0/m_e}$ (σ) convention: + +- A Maxwellian at code temperature $T = 1$ is $f \propto e^{-v^2/2}$ (unit variance). +- The Bohm–Gross dispersion in code units is $\omega^2 = 1 + 3 k^2$. +- The normalized speed of light is $\hat c = c/\sqrt{T_0/m_e}$ (e.g. $\hat c = 15.98$ at 2000 eV). + +Dimensional inputs (strings with units, e.g. `xmax: 100um`) are converted with these +units; plain numeric inputs are taken to already be in code units and pass through +unchanged. + ## density Species and density configuration. You can define multiple "species" by using keys prefixed with `species-`. @@ -65,9 +87,9 @@ Each species is defined with a key starting with `species-` (e.g., `species-back | `noise_seed` | int | Random seed for noise initialization | | `noise_type` | string | `"gaussian"` or `"uniform"` | | `noise_val` | float | Amplitude of noise | -| `v0` | float | Drift velocity (normalized) | -| `T0` | float | Temperature (normalized to `normalizing_temperature`) | -| `m` | float | Exponent for super-Gaussian distribution. `2.0` is Maxwellian | +| `v0` | float | Drift velocity in code units of $\sqrt{T_0/m_e}$ (thermal-σ units). Numeric only — dimensional strings are not supported here | +| `T0` | float | Temperature in units of `normalizing_temperature`. The initialized distribution has velocity variance `T0/mass`. Numeric only | +| `m` | float | Exponent for super-Gaussian distribution $f \propto \exp[-\|v/(\alpha v_{th})\|^m]$. `2.0` is Maxwellian. Note: $\alpha$ is chosen to fix the moment ratio $\langle v^4\rangle/\langle v^2\rangle = 3\,T_0/\mathrm{mass}$ for all $m$; the *variance* equals `T0/mass` only at `m: 2`. For flat-top distributions (`m > 2`) the second-moment temperature diagnostic will read higher than `T0` (e.g. ×1.24 at `m: 3`, ×1.37 at `m: 4`) | | `basis` | string | Spatial profile type (see below) | #### Basis Types @@ -531,7 +553,7 @@ Both `fokker_planck` and `krook` use profile objects to define spatiotemporal va | Field | Type | Description | |-------|------|-------------| -| `baseline` | float | Base collision frequency | +| `baseline` | float | Base collision frequency in code units of $\omega_{p0}$ (i.e. $\hat\nu = \nu[\mathrm{s}^{-1}]/\omega_{p0}$, a true rate — no $2\pi$). For reference, the NRL e–e rate in these units is logged as `nuee_norm` in `units.yaml`. Note the Dougherty/Lenard–Bernstein $\nu$ agrees with the NRL $\nu_{ee}$ only up to an O(1) factor | | `bump_or_trough` | string | `"bump"` or `"trough"` | | `center` | float | Profile center | | `rise` | float | Transition steepness | From f9dfb4c085d02d35e64bd98f7fc85f75cc04d4af Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 17 Jul 2026 12:58:55 -0700 Subject: [PATCH 46/49] test(vlasov1d): regenerate config regression fixtures The fixtures locked in the sqrt(2)-wrong units (c_light 11.302, v0 2.652e7 m/s, x0 12.14 nm, negative nuee). Regenerated under the corrected normalization and the new config surface: c_light 15.984, v0 1.876e7 m/s, x0 8.584 nm, logLambda_ee +7.22, positive nuee, new nuee_norm, species_params T0, c_light knob removed. Co-Authored-By: Claude Fable 5 --- ...okker_planck_conservation_array_config.yml | 18 ++++++++++-------- ...ker_planck_conservation_derived_config.yml | 17 +++++++++-------- .../fokker_planck_conservation_units.yml | 15 ++++++++------- ...multispecies_ion_acoustic_array_config.yml | 19 +++++++++++-------- ...ltispecies_ion_acoustic_derived_config.yml | 17 +++++++++-------- .../multispecies_ion_acoustic_units.yml | 15 ++++++++------- .../resonance_array_config.yml | 18 ++++++++++-------- .../resonance_derived_config.yml | 17 +++++++++-------- .../resonance_units.yml | 15 ++++++++------- 9 files changed, 82 insertions(+), 69 deletions(-) diff --git a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml index 4134829c..9539860e 100644 --- a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml +++ b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.0 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.1 dx: 0.785 ion_charge: @@ -16507,6 +16507,7 @@ grid: vmax: 6.4 species_params: electron: + T0: 1.0 charge: -1.0 charge_to_mass: -1.0 mass: 1.0 @@ -17723,17 +17724,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.07623654008768617 micron + beta: 0.062561188941867 + box_length: 0.05390737447020297 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.004622577634581562 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml index ebca6209..f83b3fe6 100644 --- a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml +++ b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.0 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.1 dx: 0.785 max_steps: 105 @@ -111,17 +111,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.07623654008768617 micron + beta: 0.062561188941867 + box_length: 0.05390737447020297 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.004622577634581562 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml index dca5edde..f1e7ae4b 100644 --- a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml +++ b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml @@ -1,13 +1,14 @@ T0: 2000 electron_volt -beta: 0.088474881879774 -box_length: 0.07623654008768617 micron +beta: 0.062561188941867 +box_length: 0.05390737447020297 micron box_width: inf -c_light: 11.302642950785 -logLambda_ee: -11.781233290653 +c_light: 15.984350951661 +logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter -nuee: -574949910190.1326 hertz +nuee: 352401683370.27344 hertz +nuee_norm: 0.00016128753860756 sim_duration: 0.004622577634581562 picosecond tp0: 0.45768095391896646 femtosecond -v0: 26524102.30999721 meter / second +v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second -x0: 12.139576447083783 nanometer +x0: 8.58397682646544 nanometer diff --git a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml index 927d9760..e899d7fa 100644 --- a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml +++ b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml @@ -29,7 +29,7 @@ drivers: ex: {} ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.5 dx: 9.81746875 ion_charge: @@ -28711,10 +28711,12 @@ grid: vmax: 0.005 species_params: electron: + T0: 1.0 charge: -1.0 charge_to_mass: -1.0 mass: 1.0 ion: + T0: 0.01 charge: 10.0 charge_to_mass: 0.00054466230936819 mass: 18360.0 @@ -38873,17 +38875,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 3.8137571970393944 micron + beta: 0.062561188941867 + box_length: 2.696733575825556 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 2.288633610071792 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml index e7c61532..9fe52c26 100644 --- a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml +++ b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml @@ -29,7 +29,7 @@ drivers: ex: {} ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.5 dx: 9.81746875 max_steps: 10005 @@ -120,17 +120,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 3.8137571970393944 micron + beta: 0.062561188941867 + box_length: 2.696733575825556 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 2.288633610071792 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml index b05d7d96..e018602a 100644 --- a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml +++ b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml @@ -1,13 +1,14 @@ T0: 2000 electron_volt -beta: 0.088474881879774 -box_length: 3.8137571970393944 micron +beta: 0.062561188941867 +box_length: 2.696733575825556 micron box_width: inf -c_light: 11.302642950785 -logLambda_ee: -11.781233290653 +c_light: 15.984350951661 +logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter -nuee: -574949910190.1326 hertz +nuee: 352401683370.27344 hertz +nuee_norm: 0.00016128753860756 sim_duration: 2.288633610071792 picosecond tp0: 0.45768095391896646 femtosecond -v0: 26524102.30999721 meter / second +v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second -x0: 12.139576447083783 nanometer +x0: 8.58397682646544 nanometer diff --git a/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml b/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml index 8912206a..0924eec1 100644 --- a/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml +++ b/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.1598 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.25 dx: 0.654375 ion_charge: @@ -20747,6 +20747,7 @@ grid: vmax: 6.4 species_params: electron: + T0: 1.0 charge: -1.0 charge_to_mass: -1.0 mass: 1.0 @@ -23327,17 +23328,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.25420273080193445 micron + beta: 0.062561188941867 + box_length: 0.17974847474618633 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.21980127811958367 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml b/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml index 66794372..6e5e2484 100644 --- a/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml +++ b/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.1598 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.25 dx: 0.654375 max_steps: 1925 @@ -119,17 +119,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.25420273080193445 micron + beta: 0.062561188941867 + box_length: 0.17974847474618633 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.21980127811958367 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/resonance_units.yml b/tests/test_vlasov1d/test_config_regression/resonance_units.yml index cc18e176..22e9776c 100644 --- a/tests/test_vlasov1d/test_config_regression/resonance_units.yml +++ b/tests/test_vlasov1d/test_config_regression/resonance_units.yml @@ -1,13 +1,14 @@ T0: 2000 electron_volt -beta: 0.088474881879774 -box_length: 0.25420273080193445 micron +beta: 0.062561188941867 +box_length: 0.17974847474618633 micron box_width: inf -c_light: 11.302642950785 -logLambda_ee: -11.781233290653 +c_light: 15.984350951661 +logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter -nuee: -574949910190.1326 hertz +nuee: 352401683370.27344 hertz +nuee_norm: 0.00016128753860756 sim_duration: 0.21980127811958367 picosecond tp0: 0.45768095391896646 femtosecond -v0: 26524102.30999721 meter / second +v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second -x0: 12.139576447083783 nanometer +x0: 8.58397682646544 nanometer From f5785ad28c9b08521b41a594e3221f31dcb4d84c Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 20 Jul 2026 11:44:10 -0700 Subject: [PATCH 47/49] Added audit files --- .../2026-07-16_ADEPT_CONVENTIONS_AUDIT.md | 225 ++++++++++++++++ .../2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md | 249 ++++++++++++++++++ 2 files changed, 474 insertions(+) create mode 100644 docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md create mode 100644 docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md diff --git a/docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md b/docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md new file mode 100644 index 00000000..4b1ac623 --- /dev/null +++ b/docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md @@ -0,0 +1,225 @@ +# ADEPT Codebase-Wide Physics & Conventions Audit + +**Date:** 2026-07-16 +**Scope:** every solver module and shared physics file in `adept/` — `_vlasov2d`, `_pic1d`, `_tf1d`, `_lpse2d`, `_spectrax1d`, `_hermite_poisson_1d`, `vfp1d`, `osiris`, `vlasov1d2v`, plus the shared `normalization.py`, `electrostatic.py`, `functions.py`, `driftdiffusion.py`, `utils.py`, and all tests/configs. +**Companion document:** `VLASOV1D_CONVENTIONS_AUDIT.md` (the `_vlasov1d` deep audit that established the root cause). This document extends that audit to the rest of the codebase and answers: *which solvers are immune?* +**Method:** six parallel independent module audits, each cross-checked against source; load-bearing claims re-verified by hand. No source files modified. + +--- + +## 1. Executive summary + +**The √2 bug is confined to the `electron_debye_normalization` consumers — and even there, only to the unit-conversion/reporting layer.** The dimensionless *dynamics* of every solver in the codebase are internally self-consistent. The full damage inventory of `normalization.py:91` ($v_0 = \sqrt{2T_0/m_e}$, $L_0 = \sqrt2\lambda_{De}$): + +- `_vlasov1d`, `_vlasov2d`, `_pic1d`: logged physical units √2-wrong, physical temperature 2× the label, dimensional-string inputs √2-off, EM-branch $\hat c$ √2-small (details per module below). +- `_tf1d`: **independently re-implements the identical √2** in its own `write_units` (`_tf1d/modules.py:62`) — same symptom, separate code, needs its own fix. +- Everyone else: **immune** (see the table). + +**One critical amendment to the fix prescription in the vlasov1d audit:** `vth_norm()` (`normalization.py:48–50`) must **NOT** be redefined. Grep-verified, it has exactly four callers, all in `vfp1d/`, which is built self-consistently on the $v_{th} \equiv \sqrt{2T/m}$ (most-probable-speed) convention — its initializer carries a compensating factor (`vfp1d/helpers.py:111`: $\alpha = \sqrt{3\Gamma(3/m)/(2\Gamma(5/m))} = 1$ at $m{=}2$, vs. `_vlasov1d`'s $\sqrt2$) so the realized variance is exactly $T/m$. Dropping the 2 in `vth_norm()` would do nothing for `_vlasov1d` (which never calls it) and would **silently halve VFP-1D's initialized temperature** while leaving its transport coefficients at $T_0$. The safe fix touches only `electron_debye_normalization` lines 91–92. + +### 1.1 Immunity table (the answer to "which solvers are immune") + +"Bug" = the current √2 in `normalization.py:91`. "Fix" = changing `electron_debye_normalization` to $v_0=\sqrt{T_0/m_e}$ (only; `vth_norm()` untouched). + +| Solver | Working vth convention | Affected by bug NOW? | Affected by fix? | Notes | +|---|---|---|---|---| +| `_vlasov1d` | $\sqrt{T/m}$ (σ) | **YES** — logged units, physical T (2×), dimensional strings, EM $\hat c$; fixtures lock wrong values | Fix-safe for tests; must regenerate 3 fixtures | See companion doc | +| `_vlasov2d` | $\sqrt{T/m}$ (σ) | **YES (partial)** — same chain: logged units, physical T, EM $\hat c$; ES dynamics immune | **Fix-safe** — all 3 tests self-consistent, **no fixtures** | Cleanest of the affected: already fixed several vlasov1d bugs | +| `_pic1d` | $\sqrt{T/m}$ (σ) | **Logged units only** — dynamics immune (all inputs numeric) | **Fix-safe** — no fixtures, no test depends on it | Loader ≡ vlasov1d initializer convention | +| `_tf1d` | $\sqrt{T/m}$ (σ) | Not via `normalization.py` — but its **own private copy** of the √2 (`modules.py:62`) corrupts its logged units identically | Fix does **nothing** here; `modules.py:62` needs its own `2.0*T0 → T0` | Dynamics/tests fully decoupled from both | +| `_lpse2d` | $\sqrt{T/m}$ (σ), physical units | **IMMUNE** — never touches `normalization.py` | **IMMUNE** | Own ps/µm unit system; Bohm-Gross & Landau literature-exact | +| `_spectrax1d` | AW-Hermite $\alpha=\sqrt{2T/m}$ (self-consistent) | **IMMUNE** — zero imports of `normalization.py` | **IMMUNE** | α is a raw config float; physics identical to σ-convention | +| `_hermite_poisson_1d` | AW-Hermite $\alpha=\sqrt{2T/m}$ (self-consistent) | **IMMUNE** — zero imports | **IMMUNE** | Not even registered in ergoExo dispatch | +| `vfp1d` | $\sqrt{2T/m}$ (most-probable, self-consistent) | **IMMUNE** — uses `laser_normalization` ($v_0=c$), never `electron_debye_normalization` | **IMMUNE if** `vth_norm()` is left alone; **BROKEN if** `vth_norm()` is "fixed" | The reason the fix must not touch `vth_norm()` | +| `osiris` wrapper | OSIRIS native ($u_{th}=\sqrt{T/m}/c$, in-deck) | **IMMUNE** — `skin_depth_normalization` ($T_0$=None, $v_0=c$); no eV→uth conversion exists in the wrapper | **IMMUNE** | `vth_norm()` unreachable (would TypeError on $T_0$=None) | +| `vlasov1d2v` | $\sqrt{T/m}$ (σ), $m$≡1 | **IMMUNE by orphan status** — not in ergoExo dispatch, would KeyError before running | **IMMUNE** | Legacy code; carries many un-fixed bugs (§3.9) | + +**Fully immune to both bug and fix:** `_lpse2d`, `_spectrax1d`, `_hermite_poisson_1d`, `osiris`, `vfp1d` (conditional on the fix not touching `vth_norm()`), and `vlasov1d2v` (by virtue of being unrunnable). +**Affected now:** `_vlasov1d`, `_vlasov2d`, `_pic1d` (via `normalization.py`) and `_tf1d` (via its private duplicate). +**Every solver's dimensionless dynamics are unaffected in shipped configs/tests** — the corruption is confined to logged units, physical interpretation of results, dimensional-string inputs, and the EM-branch $\hat c$. + +### 1.2 Why no test ever caught the √2 + +The bug survived every test suite in the repository, and understanding why is itself a finding — the test suites, by construction, cannot see this class of error: + +1. **Every physics test works entirely in dimensionless code units.** Landau damping, ion-acoustic, Bohm-Gross, two-stream, gyro — all set $k$, $\omega$, $T_0$, box sizes as numeric floats, which pass through `normalize()` untouched (`normalization.py:57–58`). The buggy conversion layer is simply never on the tested code path. A purely numeric run *is* a correct simulation under the σ-convention; only its physical labels are wrong, and no test reads the labels. +2. **The theory references the tests compare against live in the same code-unit world.** The kinetic roots come from `electrostatic.py` with `maxwellian_convention_factor=2` and $v_{th}=1$ — i.e., the reference is expressed in the engine's internal convention, so agreement validates internal consistency, not the dimensional dictionary. The tests pin *which* convention the engine uses (invaluable for this audit) but cannot detect that `normalization.py` speaks a different one. +3. **Where a test does touch the dimensional layer, it is self-referential.** `_vlasov2d`'s `test_em_dispersion` computes $\hat c$ from the same `electron_debye_normalization` the solver uses and places the driver at $\omega^2 = 1+\hat c^2 k^2$ with that same $\hat c$ — bug and reference move in lockstep, so it passes with the wrong $c$. +4. **The only tests that pin dimensional values pin the *wrong* ones.** The `_vlasov1d` `*_derived_config.yml` regression fixtures were generated under the buggy convention, so they actively *enforce* the √2 values (`c_light: 11.302`, `x0: 12.14 nm`, …) — a fix makes tests fail, not the bug. +5. **No solver has an absolute physical-units validation** (e.g., a Landau rate checked against a rate in Hz for a stated density and temperature, or an SRS resonance checked against a wavelength in nm), and none has a total-energy conservation diagnostic. Either would have caught the mismatch the first time a dimensional input mattered. + +The general lesson: dimensionless-physics tests validate the engine; only a test that crosses the units boundary — dimensional input in, dimensional observable out, compared against an independent physical reference — can validate the normalization layer. The repo currently has zero such tests. Adding one per normalization entry point (a "round-trip units test") is the cheapest structural guard against recurrence, and belongs in the fix PR alongside item 6 of §4. + +### 1.3 The three thermal-velocity conventions in the codebase + +All are *internally* self-consistent within their modules; the mixing hazard is at module boundaries and in shared helpers: + +| Convention | Definition | Used by | +|---|---|---| +| σ (RMS / std-dev) | $v_{th} = \sqrt{T/m}$, Maxwellian $e^{-v^2/2}$ at $T{=}1$, $L_0=\lambda_{De}$ | `_vlasov1d`, `_vlasov2d`, `_pic1d`, `_tf1d`, `_lpse2d`, `vlasov1d2v` dynamics; `electrostatic.py` (mcf=2 default); OSIRIS `uth` | +| Most-probable | $v_{th} = \sqrt{2T/m}$ | `vfp1d` (with compensating init factor); `vth_norm()`; **`electron_debye_normalization` (the outlier)**; `_tf1d/modules.py:62` (private duplicate, diagnostics-only) | +| AW-Hermite scale | $\alpha = \sqrt{2T/m}$ as basis parameter; represented Maxwellian still has variance $T/m$ | `_spectrax1d`, `_hermite_poisson_1d` (Schumer–Holloway / Parker–Dellar standard) | + +--- + +## 2. Normalization wiring map + +Grep-verified callers of each `normalization.py` entry point: + +| Entry point | Definition | Callers | +|---|---|---| +| `electron_debye_normalization` | $v_0=\sqrt{2T_0/m_e}$ ← **the bug**, $L_0=v_0/\omega_{p0}$ | `_vlasov1d`, `_vlasov2d` (`modules.py:55`), `_pic1d` (`simulation.py:78`) | +| `laser_normalization` | $v_0=c$, $L_0=c/\omega_L$, $n_0=n_{crit}$, $T_0$ not self-consistent with $v_0$ (documented) | `vfp1d` (`base.py:17`) only | +| `skin_depth_normalization(_from_frequency)` | $v_0=c$, $L_0=c/\omega_{p0}$, $T_0$=None | `osiris` only | +| `vth_norm()` | $\sqrt{2T_0/m_0}/v_0$ — hard-codes most-probable | **`vfp1d` only** (grid.py:118, base.py:93, 98, 152) | +| `speed_of_light_norm()` | $c/v_0$ | `_vlasov1d`, `_vlasov2d`, `_pic1d` (√2-tainted); `osiris` (=1, exact) | +| `normalize(s, norm, dim)` | numeric passthrough; strings via $L_0$/τ/$v_0$/$T_0$ | vlasov family, `functions.py`, `vfp1d/helpers.py` profiles | +| — (no normalization import) | | `_tf1d` (own inline pint), `_lpse2d` (own unit system), `_spectrax1d`, `_hermite_poisson_1d` (raw α floats), `vlasov1d2v` (reads cfg keys nothing populates) | + +Shared physics kernels: +- `electrostatic.py` — dispersion/Z-function utilities. Self-consistent under σ-convention: `maxwellian_convention_factor=2` (default, used by **every** caller in the repo) means $f_0\propto e^{-v^2/2v_{th}^2}$, $\xi=\omega/(\sqrt2 k v_{th})$. This is the reference that pins the working convention of `_vlasov1d`, `_vlasov2d`, `_pic1d`, `_tf1d`, `_spectrax1d`, `_hermite_poisson_1d` tests. +- `driftdiffusion.py` — Dougherty/LB kernel; measures $T=\langle(v-\bar v)^2\rangle$ (or spherical $\langle v^4\rangle/3\langle v^2\rangle$), relaxes to that width. Convention-agnostic and self-adapting; verified clean in both audits. +- `functions.py` — envelope/profile layer for the vlasov family. Numeric inputs pass through; dimensional strings go through $L_0$ (√2-tainted until the fix). One bug found (§3.10). +- `utils.py` — no physics content. + +--- + +## 3. Per-module findings + +Finding IDs continue the companion doc's F-series where the bug is a copy; new IDs are per-module. + +### 3.1 `_vlasov2d` — affected now (units/EM), fix-safe, partially cleaned-up lineage + +Convention: σ, identical to `_vlasov1d` (`helpers.py:62` `v_th=√(T0/mass)`, same super-Gaussian α; 2-D init correctly normalized, $\iint f\,d^2v = n$, equal widths in $v_x,v_y$). Landau test pins σ=1 via the shared kinetic roots. Pushers all unity-coefficient; TE-mode Maxwell curl signs verified; initial 2-D Poisson correct; magnetic rotation $\theta=-(q/m)B_z dt$ verified by the gyro test. + +Bug-copy status vs `_vlasov1d`: **F5 fixed** (vbar divides by n, `fokker_planck.py:32`), **F6 fixed** (moments centered on $u=j/n$, `storage.py:100`), **F11 fixed** (`dx=(xmax-xmin)/nx`, `grid.py:60`), F7 N/A (no entropy diagnostic). **F4 still present** (Krook target $e^{-v_x^2/2}e^{-v_y^2/2}$ hard-coded, `fokker_planck.py:127–131`; latent, off in configs). **F10 still present** (`float(cfg.T0/v0x/v0y)`, `simulation.py:158–163`). **F8 present** (same α, latent at m=2). + +√2 exposure: calls `electron_debye_normalization` (`modules.py:55`) → logged units √2-wrong (F2 analog), physical T 2× label (F1 analog), EM $\hat c = c/v_0$ √2-small (F3 analog). `test_em_dispersion` passes *because it is self-consistent, not correct* — it recomputes the same wrong $\hat c$ for the driver. **No regression fixtures exist**, so the fix requires no fixture regeneration; all three tests pass unchanged (the EM test moves in lockstep with the fix). + +New (minor): `mean_KE` has the ½, field proxies $\langle E^2\rangle,\langle B^2\rangle$ don't (no combined conserved-energy diagnostic); `Txx/Tyy` are pressures ($n\cdot$variance) despite the T-name, and `T` omits the mass factor (fine for electrons); the separable per-axis Dougherty never isotropizes $T_x \leftrightarrow T_y$ (no cross-axis coupling operator exists in this module); dead config knobs `UnitsConfig.laser_wavelength`, `Z`. + +### 3.2 `_pic1d` — dynamics immune, logged units affected, fix-safe + +Convention: σ. The particle loaders (`helpers.py:34` quiet inverse-CDF, `:58` random rejection) use `v_thermal = np.sqrt(T0/mass)` — byte-for-byte the vlasov1d initializer convention. Pinned independently by `test_bohm_gross.py:34` (kinetic dielectric with vth=1) and `test_landau_damping.py:38–41` ($\omega^2=1+3k^2$, σ=1 Landau rate). + +Verified clean: deposit/gather use the identical B-spline kernel (momentum-conserving, self-force-free); uniform loading deposits exactly $n=1$ in expectation ($w=n_0L/N$, partition of unity); Poisson $E_k=-i\rho_k/k$ unity prefactor; KDK leapfrog and Yoshida4 push coefficients exact; ponderomotive $(q/m)^2(-\tfrac12\partial_x a^2)$ matches vlasov1d's verified chain. + +√2 exposure: `write_units` (`modules.py:44–57`) logs √2-wrong `v0, x0, c_light, box_length` (a 2000 eV epw run is physically 4000 eV). Dynamics immune: every shipped config and test input is numeric; loader reads raw `T0`. **No fixtures.** Fix is strictly an improvement — nothing breaks. + +New findings: +- **P1 (confirmed):** energy diagnostics mix extensivity: `mean_KE = 0.5·m·Σw v²` is a box integral with the ½; `mean_e2 = mean(e²)` is a per-cell mean without ½ or dx. `mean_KE + mean_e2` is not conserved; no valid energy monitor exists. +- **P2 (confirmed):** quiet and random loaders truncate velocity differently for drifting species — quiet clips to $[-v_{max}, v_{max}]$ absolute (`helpers.py:38`), random clips to $[v_0-v_{max}, v_0+v_{max}]$ (`helpers.py:65`). Same config, two different realized distributions; quiet wrongly clips the high-v side of a drifting Maxwellian. Latent for shipped cold-beam decks. +- **P3 (confirmed):** `mean_KE/mean_p` are sums but named `mean_*` (naming, cf. F6-class). +- Inherited: F11 (`dx=xmax/nx` while particle wrap and placement correctly use $x_{max}-x_{min}$ — the field grid and particle box disagree for $x_{min}\ne0$), F10, F8. + +### 3.3 `_tf1d` — decoupled from `normalization.py`, but carries a private copy of the √2 + +Convention: σ. Derived from the pushers: linearizing continuity + momentum ($-u\partial_x u - \frac{1}{n}\partial_x(p/m) - \frac{q}{m}E$) + adiabatic energy ($\gamma=3$) + Poisson gives $\omega^2 = 1+3k^2$, exactly what `test_resonance.py:19` asserts; the kinetic branch uses the shared roots (mcf=2). The Poisson sign ($E_k=+i\rho_k/k$, i.e. $\partial_x E=-\rho$) and the momentum sign ($-(q/m)E$) are *both* opposite the textbook and cancel — verified stable by derivation and by the passing test. Landau closure decays the field at exactly $2\,\mathrm{Im}\,\omega$, matching `test_landau_damping.py:67`. + +**T1 (confirmed, the headline):** `_tf1d/modules.py:62` re-implements `v0 = √(2·T0/m_e)` inline in its own `write_units()` (with its own pint registry — no import of `normalization.py`). Every tf1d run logs `v0, x0, c_light, beta, box_length, sim_duration` √2-wrong and implies 2× the stated temperature. Damage confined to logged diagnostics (no dimensional-string inputs exist in tf1d; the EM `WaveSolver` branch is commented out; no fixtures). **Fixing `normalization.py` does not fix this** — `modules.py:62` needs its own `2.0*T0 → T0` in lockstep. + +Other findings: +- **T2 (confirmed):** the kinetic-γ pressure closure (`pushers.py:195` `wr_corr=(wrs²-1)/k²` + γ forced to 1) yields $\omega^2 = (w_{rs}^2-1)T_0+1$ — exact only at $T_0=1$. Convention- and T₀-locked, same class as vlasov1d's Krook (F4). Latent (all configs use $T_0=1$). +- **T3 (confirmed):** `modules.py:78` `nuee_norm = nuee/wp0` — the *correct* Hz→code conversion — is computed and then **discarded** (never logged, never used). Meanwhile `physics..trapping.nuee` is an unrelated raw ML-input float. There is no collisional friction in the tf1d equations at all. (F9-class trap.) +- **T4 (confirmed):** `docs/source/usage/tf1d.md` momentum equation has a spurious $1/n$ on the E-force and mislabels the Poisson equation. Docs-only. +- **T5 (suspected, latent):** `EnergyStepper` evolves $p$ but advects/compresses $p/m$ — ambiguous mass normalization of the pressure variable; inconsistent with the momentum equation for mobile ions ($m=1836$). All configs have ions off. +- **T9 (minor):** `resonance_search.yaml` sets `ion.landau_damping: True` with `ion.is_on: False` (inert, misleading). + +### 3.4 `_lpse2d` — fully immune; thermally clean; unrelated bugs found + +Physical-units solver (ps/µm, Gaussian-cgs-derived) with its own `write_units` (`helpers.py:71`); zero imports of `normalization.py`. Convention: $v_{te} = c\sqrt{T_e/511}$ = $\sqrt{T_e/m_e}$ (σ) used consistently in the Bohm-Gross term ($e^{-i\,1.5\,v_{te}^2 k^2/\omega_{pe}\,dt}$, coefficient 3/2 correct for σ), the Landau damping rate (verified literature-exact: $\sqrt{\pi/8}(1+\tfrac32 k^2\lambda_D^2)(k\lambda_D)^{-3}e^{-3/2-1/(2k^2\lambda_D^2)}$), the sound speed, and the TPD threshold. Numerically cross-checked: the test config's driver $\omega = 1.5k^2v_{te}^2/\omega_{p0} = 19.7 \approx 20$ as configured. TPD/SRS coupling constants contain no thermal factor (convention-independent); laser $E_0=\sqrt{8\pi I/c}$ and WKB swelling $(1-n/n_c)^{-1/4}$ verified. + +Findings (none √2-class): +- **L1 (confirmed):** logged `lambda_D = vte/w0` (`helpers.py:111`) divides by the **laser** frequency instead of $\omega_{pe}$ — factor 2 too small at $n_c/4$. Diagnostic only (dynamics compute $\omega_{pe}^2/k^2v_{te}^2$ directly). +- **L2 (confirmed, latent runtime bug):** the `E2` electrostatic-driver path calls `self.epw.driver(...)` (`core/vector_field.py:84`) but the active `SpectralEPWSolver` has no `driver` attribute (only the commented-out `SpectralPotential` does) → `AttributeError` for any config with an `E2` driver. Latent only because `test_epw_frequency` is currently disabled (`pass` body) — meaning **the module's dispersion conventions have no live regression test**. +- **L3 (confirmed):** `core/trapper.py` is dead code (never instantiated by `SplitStep`); it is a stale duplicate of `_tf1d`'s trapper. Note: it uses `electrostatic.py`, *not* `driftdiffusion.py` — the vlasov1d audit's cross-module note claiming lpse2d uses the Dougherty kernel is wrong for the current tree. +- **L4 (confirmed):** the σ-convention is undocumented in the lpse2d docs — the main *future* √2 risk here is someone "harmonizing" it onto `normalization.py`. +- **L5 (minor):** the Landau/Bohm-Gross formulas are duplicated byte-for-byte in `SpectralPotential` and `SpectralEPWSolver` (drift risk); `nu_coll`'s `/2` (amplitude vs energy rate) is correct but easy to double-count. + +### 3.5 `_spectrax1d` and `_hermite_poisson_1d` — fully immune; self-consistent AW-Hermite basis + +Both use the asymmetrically-weighted Hermite basis (Schumer–Holloway / Parker–Dellar) with scale $\alpha = \sqrt{2T/m}$ taken as a **raw config float** — zero imports of `normalization.py`. The represented Maxwellian has variance $\alpha^2/2 = T/m$: physically identical to the σ-convention solvers. Verified by derivation (streaming ladder $\alpha\sqrt{n/2}$ + force ladder $\sqrt{2n}/\alpha$ + Ampère/Poisson coupling → Langmuir $\omega=\omega_{pe}$ exactly, Bohm-Gross $\omega^2=1+3(k\lambda_{De})^2$ with $\lambda_{De}=\alpha/\sqrt2$) and by tests: both modules' Landau tests set $\alpha_e = \sqrt2\,k\lambda_D/k$ explicitly against the shared mcf=2 kinetic roots (2%/5% tolerance for hermite-poisson; a cross-module test pins the two modules' E-coupling equal to 1e-12). + +The coincidence that $\alpha=\sqrt{2T/m}$ matches the `normalization.py:91` outlier is harmless — there is no code coupling in either direction. Both integrator paths (DoPri8 and Lawson-RK4 exponential operators) share identical ladder coefficients. The historical inverted E-coupling bug in hermite_poisson (`C[n+1]` vs `C[n-1]`) is fixed and regression-locked by three tests. + +Findings (all diagnostic-label/minor): **S1** logged `lambda_D` is the *total* (ion-dominated) Debye length while `k_norm` uses electron-only and hard-codes mode 1 (`base_module.py:186,189`); **S2** the ion "temperature" diagnostic is a velocity variance missing the $m_i$ factor ($T_i/m_i$, misleading next to $T_e$; `storage.py:386–405`); **S3** shipped `landau-damping.yaml` has `Nn: 4` — far too few Hermite modes to resolve the damping it is named for (tests override to 512/32); `_spectrax1d/helpers.py` is an all-stub file not on any code path. + +### 3.6 `vfp1d` — immune to the bug; the reason the fix must not touch `vth_norm()` + +Uses `laser_normalization` (`base.py:17`): $v_0 = c$, $L_0 = c/\omega_L$, $n_0 = n_{crit}$, $T_0$ = reference eV (documented as not self-consistent with $v_0$). Velocity grid in units of $c$. Thermal convention: **most-probable, $v_{th} = \sqrt{2T/m}$**, deliberately and self-consistently: + +- `vth_norm()` supplies $\sqrt{2T_0/m}/c$ (4 call sites, all vfp1d); +- the initializer's width factor `helpers.py:111` is $\alpha = \sqrt{3\Gamma(3/m)/(2\Gamma(5/m))}$ — note the extra `/2` vs the vlasov-family α — giving $\alpha=1$ at $m{=}2$ and realized variance exactly $T/m$ (the √2 and the /2 cancel); +- the v-grid extent `grid.py:118` `vmax = 8·vth_norm()/√2` = 8 standard deviations (the /√2 converts most-probable→σ); +- the IB coefficient `base.py:81` ($0.093373\,\lambda_{\mu m}^2/T_{keV}$ per $10^{15}$ W/cm²) normalizes $v_{osc}^2$ to $2T/m$, consistent; +- the temperature diagnostic `storage.py:317–322` uses the spherical variance $T = \langle v^4\rangle/3\langle v^2\rangle$ — reads the true temperature, no √2. + +Collision operators ($\nu_{ee}$ coefficient anchored to $v_0=c$ via $r_e$, Lorentz e–i with Epperlein–Haines Z*, Rosenbluth I/J integrals) are temperature-convention-independent. **No internal √2 inconsistency exists.** + +**C1 (confirmed, cross-module, high):** because items above are keyed to `vth_norm()` while the compensating α is a separate constant, redefining `vth_norm()` → $\sqrt{T_0/m}/v_0$ (as the vlasov1d audit's §5.1 originally suggested) silently initializes VFP-1D at **half** the intended temperature while all transport coefficients stay at $T_0$; the Spitzer/Epperlein–Haines gold test (`test_kappa_eh.py`, $\kappa\propto T^{5/2}$) would likely fail, and the pure-operator tests would *not* catch it. **The fix must be confined to `electron_debye_normalization:91–92`.** (The companion doc's §5.1 has been amended accordingly.) + +Other findings: **C2** `base.py:118` stores the bound method `norm.vth_norm` instead of calling it (harmless repr-string in cfg; wrong); **C3** `storage.py:170` hard-codes $9.09\times10^{21}$ cm⁻³ per $n_c$ — that is $n_{crit}(351\,\mathrm{nm})$; wrong labeling for any other `laser_wavelength` (should derive from `norm.n0`); **C4** = F11 copy (`grid.py:71` `dx=xmax/nx`); **C6/C7 (suspected, latent — IB not in any shipped config):** the production inverse-bremsstrahlung wiring looks under-normalized — `w0_norm` is identically 1.0 under laser normalization, the Langdon-factor argument $Z^2 n_i/(\omega_0 v^3)$ carries no collision-frequency coefficient, and `vosc2_per_intensity` (thermal-speed² units) is consumed as $v_{osc}^2$ in grid ($c$) units — a $\sim(c/v_{th})^2$ discrepancy between the unit tests (where the grid unit *is* the thermal speed) and production. Recommend a Spitzer-IB validation run before enabling IB in production. + +### 3.7 `osiris` wrapper — fully immune + +Uses `skin_depth_normalization(_from_frequency)` ($v_0=c$, $T_0$=None). The wrapper drives a native OSIRIS deck: **no eV→`uth` conversion exists anywhere in the wrapper** — `uth` values pass through verbatim from the deck, so the T→v √2 trap cannot occur. Verified exact: $a_0 \to E_{peak} = a_0\omega_L m_e c/e \to I = E^2\epsilon_0 c/2$ (matches $I\lambda_{\mu m}^2 = 1.37\times10^{18}a_0^2$, test-confirmed); `c_light = beta = 1.0` identically; box lengths in true skin depths; `units.yaml` deliberately omits temperature-dependent keys (consistent with $T_0$=None). Temperature diagnostics (`plots.py:1004` $T=\sum u_{th,i}^2$ in $m_ec^2$ units; `:1059` momentum variance) carry no spurious factors. `vth_norm()` is unreachable (would TypeError on $T_0$=None). + +**C8 (low, latent):** the adaptive-box feature's `reference_density = 0.25` "quarter-critical" (`density.py:75`) assumes the deck's $n_0 = n_{crit}$ (true for LPI decks with `omega0=1`); a deck normalized to a different density would silently mis-scale the density bounds (the *length* normalization stays correct). + +### 3.8 Shared `electrostatic.py` and `functions.py` + +`electrostatic.py`: self-consistent under the σ-convention throughout. $Z$ via Faddeeva, $Z' = -2(1+\xi Z)$ correct; `maxwellian_convention_factor=2` (default, used by every caller in the repo) ⇒ $f_0\propto e^{-v^2/2v_{th}^2}$, root returned as $\xi k v_{th}\sqrt{2} = \omega$. One cosmetic blemish (**T7**): the Newton `initial_root_guess` is ω-scaled ($\sqrt{\omega_p^2+3k^2v_{th}^2}$) but the root variable is ξ-scaled — converges anyway. + +`functions.py`: numeric inputs pass through; `EnvelopeFunction` amplitudes (`baseline`, `bump_height`) are raw floats (the F9 collision-frequency trap); `SineFunction.wavenumber` correctly uses `dim="k"`. **T8 (suspected):** `LinearFunction`/`ExponentialFunction` normalize `val_at_center` — a *density* — with `dim="x"` (`functions.py:159–161, 178–180`); dimensionally wrong for any string input. + +### 3.9 `vlasov1d2v` — orphaned legacy; a museum of the known bugs plus three new ones + +Not registered in the ergoExo dispatch (`_base_.py:294–333` → NotImplementedError); reads `cfg["units"]["derived"]` keys nothing populates → cannot run. Never calls `normalization.py`. Convention: σ (init `exp(-(v_x^2+v_y^2)/2T_0)`, mass≡1; the super-Gaussian machinery is commented out so `m` is silently ignored). + +Carries live copies of: **F5** (`fokker_planck.py:93,98` — vbar *and* the diffusion coefficient computed without ÷n), **F6 + F7** (`storage.py:152–158, 296`), **F10**, **F11** (`helpers.py:211`). F4 present but dead (Krook never instantiated). + +New bugs (would matter if ever revived): +- **N1 (confirmed):** the density profile is computed for every basis and folded into the ion background (`helpers.py:287`), but the line applying it to $f$ (`helpers.py:83–85`) is **commented out** — electrons always start uniform, so non-uniform configs begin with spurious net charge and no actual density perturbation in $f$. +- **N2 (confirmed):** the explicit Dougherty substeps apply ν twice (`fokker_planck.py:134–135, 141–142`: `dfdt = nu*ddx(...)` then `f + dt*nu*dfdt`) — collisional relaxation runs at $\nu^2$. +- **N4 (confirmed):** `integrator.py:234` passes the **nu_ee** envelope args when building the **nu_ei** profile (copy-paste) — the configured nu_ei profile is silently ignored. +- **N3 (design note):** without collisions the $v_y$ dimension is dynamically inert (no $v_y$ force, streaming by $v_x$ only) — the entire $n_{vy}$ grid is dead compute except through the FP operators. +- Minor: default scalars broadcast over the $v_y$ axis only (`storage.py:292–295`). + +Recommendation: either delete `vlasov1d2v` or quarantine it with a module-level comment; in its current state it is a trap for anyone who greps for reference implementations. + +### 3.10 Cross-module systemic observations + +1. **Driver amplitude conventions differ across the Vlasov family** (N5): `_vlasov1d` Ex driver $\propto \omega a_0$; `_vlasov2d` current source $\propto \omega^2 a_0$; `vlasov1d2v` $\propto |k| a_0$. Same config key, three physical meanings. Worth unifying or documenting per-solver. +2. **No solver has a valid total-energy conservation diagnostic.** vlasov1d/vlasov2d/pic1d all mix ½-factors and extensivity between kinetic and field terms. A per-solver conserved-energy scalar is the single cheapest guard that would have caught the √2 class of bugs. +3. **The F-series bugs are lineage-correlated:** `_vlasov2d` (newest refactor) fixed F5/F6/F11; `_pic1d` inherits F10/F11 by importing `_vlasov1d`'s grid/simulation; `vlasov1d2v` (oldest) has everything. When fixing `_vlasov1d`, fix the copies in the same PR: Krook targets (`_vlasov1d`, `_vlasov2d`, dead `vlasov1d2v`), `dx`/xmin (`_vlasov1d/grid.py:55`, `vfp1d/grid.py:71`, `vlasov1d2v/helpers.py:211`), raw-float T0/v0 (all three vlasov-family simulation.py files). +4. **Hz→code-unit collision-frequency conversion** exists correctly in exactly one place (`_tf1d/modules.py:78`) and is dead there; nowhere is it live. Every solver that takes a collision frequency takes it as a raw code-unit float with undocumented units. + +--- + +## 4. Consolidated fix checklist (supersedes §5 of the vlasov1d audit where they differ) + +**The √2 fix, safely scoped:** +1. `normalization.py:91–92`: $v_0 \to \sqrt{T_0/m_e}$, $x_0 \to \lambda_{De}$ — inside `electron_debye_normalization` **only**. +2. **Do NOT change `vth_norm()`** (`normalization.py:48–50`) — it is vfp1d-private and correct for vfp1d (C1). If its name is judged misleading, rename to `most_probable_speed_norm()`; do not change the value. +3. `_tf1d/modules.py:62`: `2.0*T0 → T0` (private duplicate, T1). +4. Regenerate the three `_vlasov1d` `*_derived_config.yml` regression fixtures (only vlasov1d has fixtures). +5. Fix the convention-locked Krook targets in the same PR (`_vlasov1d/…/fokker_planck.py:245`, `_vlasov2d/…/fokker_planck.py:127–131`): build from species T0/mass. +6. Re-validate: vlasov1d Landau + ion-acoustic (must pass unchanged), vlasov2d Landau/gyro/EM-dispersion (EM moves in lockstep), pic1d all three (unchanged), vfp1d `test_kappa_eh` (must pass unchanged — the sentinel that `vth_norm()` was left alone), numeric-driver SRS configs re-tuned per vlasov1d F3. +7. Document the per-solver conventions in each solver's docs page (σ vs most-probable vs AW-α; see §1.3) — the biggest residual risk is a future "harmonization" that imports the wrong convention into a currently-clean module (especially `_lpse2d`, L4). + +**Independent bug fixes, by priority:** +- High (live, wrong physics if exercised): vfp1d IB normalization audit before production use (C6/C7); lpse2d `E2` driver path + re-enable `test_epw_frequency` (L2). +- Medium (live, wrong diagnostics): lpse2d logged `lambda_D` (L1); pic1d/vlasov2d/vlasov1d energy-diagnostic normalization + add conserved-energy scalars; spectrax ion-T mass factor (S2); vfp1d hard-coded $n_{crit}(351\,nm)$ (C3); pic1d quiet-loader drift truncation (P2). +- Low / hygiene: tf1d docs (T4) and dead `nuee_norm` (T3 — log it instead of dropping it); `functions.py` `val_at_center` dim (T8); vfp1d bound-method store (C2); spectrax `Nn:4` config (S3); delete or quarantine `vlasov1d2v` and lpse2d's dead trapper (L3); unify driver-amplitude conventions or document them (N5). + +--- + +## 5. What was verified clean (highlights) + +- All solver cores' dimensionless coefficients: vlasov1d/2d pushers and field solves, pic1d push/deposit/gather/Poisson, tf1d fluid system (including its double sign cancellation, verified by derivation), Hermite ladder algebra in both spectral modules (verified by derivation to the Langmuir and Bohm-Gross limits), lpse2d Bohm-Gross/Landau/TPD/SRS coefficients against literature, vfp1d collision coefficients and temperature diagnostics, osiris a0↔intensity conversions. +- The shared `driftdiffusion.py` Dougherty kernel (both audits): self-adapting, Buet factor-2 correctly absorbed, spherical temperature correct. +- `electrostatic.py`'s Z-function algebra and mcf=2 semantics — the single reference that consistently pins six modules' tests to the σ-convention. + +--- + +*Companion deep-dive for `_vlasov1d`: `VLASOV1D_CONVENTIONS_AUDIT.md`. Method: six parallel independent module audits (one Opus agent per slice) with all load-bearing claims re-verified against source before inclusion. No source files were modified.* diff --git a/docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md b/docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md new file mode 100644 index 00000000..81a302a5 --- /dev/null +++ b/docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md @@ -0,0 +1,249 @@ +# Vlasov-1D Physics & Conventions Audit + +**Date:** 2026-07-16 +**Scope:** `adept/_vlasov1d/` (all files), plus the shared code it depends on: `adept/normalization.py`, `adept/driftdiffusion.py`, `adept/electrostatic.py`, and the `tests/test_vlasov1d/` suite and example configs. +**Trigger:** the discovery that `normalization.py:91` uses the $v_0 = \sqrt{2T_0/m}$ (most-probable-speed) convention while `_vlasov1d/helpers.py:69` uses $v_{th} = \sqrt{T_0/m}$ (RMS / standard-deviation convention). Neither has been changed yet; this audit determines which convention the module actually runs in, inventories every equation, and lists all inconsistencies found. + +--- + +## 1. Executive summary + +**The engine's working convention is $v_0 = \sqrt{T_0/m}$.** The distribution initializer, the Fokker–Planck operator, the Krook operator, the Landau-damping and ion-acoustic tests (the module's quantitative gold standards), and every numeric example config are all mutually consistent under: code velocity in units of the thermal *standard deviation* $\sqrt{T_0/m_e}$, code length unit $L_0 = \lambda_{De}$, code wavenumber $= k\lambda_{De}$, Maxwellian $\propto e^{-v^2/2}$ at $T=1$. + +**`adept/normalization.py:91–92` is the sole outlier** (`electron_debye_normalization`: $v_0 = \sqrt{2T_0/m_e}$, $L_0 = \sqrt{2}\,\lambda_{De}$). Because numeric config inputs pass through `normalize()` untouched, this does **not** corrupt the dynamics of numeric-input runs — but it means: + +- A run with `normalizing_temperature: 2000eV` and species `T0: 1.0` is physically a **4000 eV** plasma (factor 2 in temperature). +- Every dimensional **string** input converted with $L_0$ (box sizes in µm, gradient scale lengths, laser $k_0$) is off by $\sqrt{2}$. +- Every **logged** physical unit (`v0`, `x0`, `c_light`, `box_length`) is off by $\sqrt{2}$, and the regression fixtures currently lock in those wrong values. +- The EM wave speed $\hat c = c/v_0$ fed to the wave solver is $\sqrt2$ too small relative to the engine's thermal unit (with a partial cancellation for wavelength-specified drivers; see F3). + +**Recommended fix direction:** change `normalization.py` to $v_0 = \sqrt{T_0/m_e}$ (so $L_0 = \lambda_{De}$), *not* the initializer. Fixing `helpers.py:69` instead (to $\sqrt{T_0/2m}$) would break the currently-passing Landau-damping and ion-acoustic tests and invalidate every existing config's driver $(k_0, \omega_0)$ values. Section 5 lists everything that must move in lockstep. + +The collisionless solver core (`vector_field.py`, pushers) was verified **clean**: every coefficient is exactly unity under the declared normalization and is convention-independent. The convention bug lives entirely in the dimensional-conversion layer. + +Beyond the $\sqrt2$ issue, the audit found several independent bugs and traps, listed in Section 4 (notably: FP `compute_vbar` missing $1/n$; storage central moments centered on $n u$ instead of $u$; a sign inconsistency between the two `-flogf` entropy diagnostics; the super-Gaussian $\alpha$ fixing the wrong moment for $m\ne2$; the Krook target hard-coded to $T=1$; and the collision-frequency normalization chain being entirely manual). + +--- + +## 2. The two candidate conventions + +With $\omega_{p0} = \sqrt{n_0 e^2/(\epsilon_0 m_e)}$, $\tau = 1/\omega_{p0}$ in both cases: + +| Quantity | (a) `normalization.py` as written | (b) engine's actual convention | +|---|---|---| +| Velocity unit $v_0$ | $\sqrt{2T_0/m_e}$ (most-probable speed) | $\sqrt{T_0/m_e}$ (RMS / std-dev) | +| Length unit $L_0 = v_0/\omega_{p0}$ | $\sqrt{2}\,\lambda_{De}$ | $\lambda_{De}$ | +| Maxwellian at $\hat T = 1$ | $\propto e^{-v^2}$ (variance ½) | $\propto e^{-v^2/2}$ (variance 1) | +| Code wavenumber $\hat k$ | $\sqrt2\, k\lambda_{De}$ | $k\lambda_{De}$ | +| Bohm–Gross | $\hat\omega^2 = 1 + \tfrac{3}{2}\hat k^2$ | $\hat\omega^2 = 1 + 3\hat k^2$ | +| $\hat c = c/v_0$ (2000 eV) | 11.30 | 15.98 | + +**Why the test suite never caught this:** every physics test (Landau, ion-acoustic, EM dispersion) works in dimensionless code units — numeric inputs bypass `normalize()` entirely (`normalization.py:57–58`), so the buggy conversion layer is never on the tested path, and the theory references (`electrostatic.py`, mcf=2, $v_{th}=1$) are expressed in the engine's internal convention, so agreement validates internal consistency rather than the physical-units dictionary. Worse, the only tests that *do* pin dimensional values — the `*_derived_config.yml` regression fixtures — were generated under the buggy convention and therefore lock the wrong values in. No test crosses the units boundary (dimensional input in, dimensional observable out, against an independent physical reference), which is precisely the kind of test that must accompany the fix (see `ADEPT_CONVENTIONS_AUDIT.md` §1.2 for the full analysis). + +Evidence pinning convention (b) as the working one, strongest first: + +1. **Landau-damping test** (`tests/test_vlasov1d/test_landau_damping.py:24`) calls `electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, k)` — i.e. $\omega_{pe}=1$, $v_{th}=1$ with `maxwellian_convention_factor=2` (`adept/electrostatic.py:79,115,98`), which is the kinetic dielectric for $f_0 \propto e^{-v^2/2v_{th}^2}$, $\xi = \omega/(\sqrt2 k v_{th})$. The measured $\omega(k)$ matches these roots to 2 decimals only if the simulated Maxwellian has $\sigma = 1$. (E.g. $k=0.3 \Rightarrow \omega = 1.1598$, which is exactly the `w0` in `resonance.yaml` and `epw.yaml`; under convention (a) the resonance would sit near 1.07 and the test would fail.) +2. **Ion-acoustic test** (`test_ion_acoustic_wave.py`) uses $c_s^2 = ZT_e/m_i$ and $\omega^2 = k^2c_s^2/(1+k^2\lambda_D^2)$ with $\lambda_D = 1$ in code units — asserting $L_0 = \lambda_{De}$. +3. **Initializer** (`helpers.py:69–76`), **Krook target** (`fokker_planck.py:245`), and **Dougherty stationary state** (`fokker_planck.py:51,96` + `driftdiffusion.py`) all build/assume $e^{-v^2/(2T/m)}$, variance $T/m$. +4. Every numeric example config (`epw.yaml`, `resonance.yaml`, `wavepacket.yaml`, `bump-on-tail.yaml`) uses driver $(k_0,\omega_0)$ pairs consistent only with $\hat k = k\lambda_{De}$ and $\sigma = \sqrt{T/m}$. + +--- + +## 3. Equation inventory + +### 3.1 Normalization layer (`adept/normalization.py`) + +| Location | Expression | Meaning | Convention note | +|---|---|---|---| +| :88 | $\omega_{p0} = \sqrt{n_0e^2/\epsilon_0 m_e}$, $\tau = 1/\omega_{p0}$ | time unit | standard | +| **:91** | $v_0 = \sqrt{2T_0/m_e}$ | velocity unit | **outlier — convention (a)** | +| :92 | $x_0 = v_0/\omega_{p0} = \sqrt2\lambda_{De}$ | length unit | inherits $\sqrt2$ | +| :48–50 | `vth_norm()` $= \sqrt{2T_0/m_0}/v_0 = 1$ | "thermal velocity" | convention (a); **never called** in `_vlasov1d` | +| :52–53 | `speed_of_light_norm()` $= c/v_0$ | $\hat c$ | inherits $\sqrt2$ (F3) | +| :37–38 | NRL $\log\Lambda_{ee}$ | Coulomb log | uses reference $T_0$ | +| :45 | $\nu_{ee} = 2.91\times10^{-6}\, n_{cc}\log\Lambda\, T_{eV}^{-3/2}$ Hz | NRL e–e rate | **diagnostic only**, never converted to code units (F9) | +| :56–74 | `normalize()`: x/L0, t/τ, v/v0, T/T0, k·L0 | dim→code | numeric inputs pass through untouched; `dim="temp"` and `dim="v"` branches are dead code for `_vlasov1d` | + +### 3.2 Initialization & grids (`helpers.py`, `grid.py`, `modules.py`, `simulation.py`, `datamodel.py`) + +| Location | Expression | Meaning | Convention note | +|---|---|---|---| +| helpers.py:65–66 | `dv = 2 vmax/nv`; cell-centered `vax` on $[-v_{max}+dv/2,\ v_{max}-dv/2]$ | v-grid | edges at $\pm v_{max}$; half-cell center offset | +| **helpers.py:69** | $v_{th} = \sqrt{T_0/m}$ | thermal width | **convention (b)** | +| helpers.py:72 | $\alpha = \sqrt{3\,\Gamma(3/m)/\Gamma(5/m)}$ ($=\sqrt2$ at $m{=}2$) | super-Gaussian width factor | fixes $\langle v^4\rangle/\langle v^2\rangle = 3v_{th}^2$; variance $=v_{th}^2$ **only** at $m=2$ (F8) | +| helpers.py:74–76 | $f \propto \exp[-|(v-v_d)/(\alpha v_{th})|^m]$ | init EDF; $m{=}2$: $e^{-(v-v_d)^2/(2T_0/m)}$ | variance $T_0/m$ — convention (b) | +| helpers.py:80,84 | `f /= sum(f)·dv`; `f *= n_prof` | $\int f\,dv = n(x)$ | midpoint rule, consistent with all moments | +| grid.py:55 | `dx = xmax/nx` | cell width | **ignores `xmin`** (F11) | +| grid.py:59–60 | `dt = min(dt, 0.95·dx/c)`, $c = 1/\beta$ | EM CFL | $\hat c$ inherits F1's $\sqrt2$ | +| grid.py:65–66,75 | `nt = int(tmax/dt+1)`; `t = linspace(0, dt·nt, nt)` | time axis | overshoot ≤ dt; `grid.t` spacing ≠ dt (F12) | +| grid.py:74,77 | cell-centered x; `kx = 2π·fftfreq(nx, d=dx)` | spectral grid | inherits F11's dx for `xmin≠0` | +| modules.py:112 | `box_length = (xmax−xmin)·L0` | box size in µm | uses `xmin` correctly (unlike grid.py:55); value carries F1's $\sqrt2$ | +| simulation.py:213–214 | `v0 = float(cfg.v0)`; `T0 = float(cfg.T0)` | species drift & temperature | bare code floats, bypass `normalize()` (F10) | +| simulation.py:74–76 | $a_0 = a_{0,std}\cdot(c/v_0)$ | intensity → quiver velocity in $v_0$ units | internally consistent with pusher (§3.3) | +| simulation.py:79–85 | $\hat k_0 = k_{phys}L_0$, $\hat\omega_0 = \omega_{phys}\tau$ | wavelength driver → code units | $\hat k_0$ carries F1's $\sqrt2$; $\hat\omega_0/\hat k_0 = \hat c$ ✓ | +| simulation.py:87 | `dw0 = 0.0 # ???` | frequency offset placeholder | unresolved (F12) | + +### 3.3 Solver core (`solvers/vector_field.py`, `solvers/pushers/vlasov.py`, `solvers/pushers/field.py`) — **verified clean** + +The dimensionless system implemented, with every coefficient exactly unity under the declared normalization (and independent of the (a)/(b) choice, since $v_0$ cancels): + +$$\partial_t f_s + v\,\partial_x f_s + \frac{\hat q_s}{\hat m_s}\left(E - \frac{\hat q_s}{2 \hat m_s}\partial_x a^2\right)\partial_v f_s = C[f_s]$$ +$$\partial_x E = \sum_s \hat q_s \int f_s\,dv \ (+\ \text{static ion background}), \qquad \partial_t E = -\sum_s \hat q_s \int v f_s\,dv$$ +$$\partial_t^2 a = \hat c^2\,\partial_x^2 a - n_e\, a + S$$ + +| Location | Code | Check | +|---|---|---| +| vlasov.py:210–214, 220 | exact spectral shift $f(x - v\,dt)$ | coefficient $v_0\tau/L_0 = 1$ ✓ | +| vlasov.py:83–87, 127–129 | `force = q·e + (q²/m)·pond; accel = force/m` | $\hat q/\hat m$ E-push; ponderomotive $(\hat q^2/\hat m^2)$ ✓ | +| vector_field.py:411 | `pond = −0.5·∂ₓ(a²)` | exact instantaneous $-\frac{1}{2m}\partial_x p_\perp^2$ with $p_\perp = qA$; the $c/v_0$ scaling of $a_0$ (simulation.py:76) makes $a$ the quiver velocity in $v_0$ units — **no hidden $\sqrt2$; verified exact** | +| field.py:203–224 | $E_k = -i\rho_k/k$ | Poisson prefactor $e^2n_0/(\epsilon_0 m_e\omega_{p0}^2) = 1$ ✓ | +| field.py:263–264, 280 | $E^{n+1} = E^n - dt\,j$ | Ampère prefactor 1 ✓ | +| field.py:337–343 | $\Delta E_k = -i\frac{q}{k}\int dv\, f_k(e^{-ikv\,dt}-1)$ | exact (Hamiltonian) Ampère along free streaming ✓ | +| field.py:146–153 | leapfrog wave eq., plasma term $-n_e a$ | dispersion $\hat\omega^2 = \hat n_e + \hat c^2\hat k^2$ ✓; $n_e$ is electron-only (~$m_e/m_i$ approx., F12) | +| vector_field.py:292–301, 341 | $n_e = -\,\hat q_e \int f_e\,dv$, time-centered | sign keyed to electron charge $-1$ ✓ | +| field.py:21–26 | $E_x$ driver $= (\omega_0{+}\delta\omega)\,a_0\sin(k_0x - \omega t)$ | vector-potential amplitude convention ($E = -\partial_t A$); matches docs (config.md:312) ✓ | +| field.py:67–69, 78–80 | point source $F_0 = 2\omega\hat c\,a_0$ | reproduces amplitude $a_0$ in vacuum; plasma value larger by $k_{vac}/k_{plasma}$, documented in docstring ✓ | + +Charge-sign conventions were checked across Poisson / Ampère / wave-equation and the force term: consistent, no compensating double-error. Current density correctly has **no** mass factor in code (but see docstring bug F12). + +### 3.4 Collisions (`solvers/pushers/fokker_planck.py`, `adept/driftdiffusion.py`) + +Operator: $\partial_t f = \nu\,\partial_v[(v - \bar v)f + T\,\partial_v f]$ (Lenard–Bernstein/Dougherty, Buet notation $\beta = 1/2T$, $D = 1/2\beta = T$, drift $C = 2\beta D(v-\bar v) = (v-\bar v)$). + +| Location | Code | Check | +|---|---|---| +| driftdiffusion.py:101ff (`discrete_temperature`) | $T = \sum f(v-\bar v)^2 dv \,/\, \sum f\,dv$ | full 2nd central moment, no factor 2, no $m$ — variance convention, matches init ✓ | +| driftdiffusion.py:170,180 | $\beta_{init} = 1/2T$; $f_{mx} = e^{-\beta(v-\bar v)^2}$ | stationary state $e^{-(v-\bar v)^2/2T}$ ✓ | +| driftdiffusion.py:333–334 | $C = 2\beta D\,(v_{edge} - \bar v)$ | $2\beta D = 1$ exactly ✓ | +| driftdiffusion.py (flux/Chang–Cooper, implicit solve) | conservative central & positivity-preserving fluxes; $(I - dt\,\nu L)f^{n+1} = f^n$ | numerics only; $\nu\,dt$ dimensionless ✓ | +| driftdiffusion.py:20–25 | Buet "extra factor of 2" note | correctly absorbed into $\beta = 1/2T$; **no stray factor survives** (verified) | +| **fokker_planck.py:81** | `compute_vbar` $= \sum f\,v\,dv$ | returns $n\bar u$, **not** $\bar u$ — missing $1/n$ (F5) | +| **fokker_planck.py:245** | Krook target $f_{mx} \propto e^{-v^2/2}$ | hard-coded variance 1 = ($T{=}1$, $m{=}1$, convention (b)) — convention-locked (F4) | +| fokker_planck.py:262–266 | $f \to f e^{-\nu dt} + n(x) f_{mx}(1 - e^{-\nu dt})$ | BGK; conserves $n$, not momentum/energy (by design) | + +Key property: the **Dougherty operator is self-adapting** — it measures $T = \langle(v-\bar v)^2\rangle$ from $f$ and relaxes toward exactly that width, so it preserves whatever convention the initializer used and needs no change under a convention fix. The **Krook operator does not** — its width is a compile-time constant (F4). + +### 3.5 Diagnostics & storage (`storage.py`) + +| Location | Code | Meaning | Note | +|---|---|---|---| +| :134–135 | $\int(\cdot)\,dv$ = `sum·dv` | midpoint rule | consistent with init normalization ✓ | +| :139 | `n` $= \int f\,dv$ | density | ✓ | +| :140 | `v` $= \int v f\,dv$ | **raw first moment $= n\bar u$**, labeled "v" | flux, not velocity (F6) | +| :141–142 | `p` $= \int (v - n\bar u)^2 f\,dv$ | "pressure" | centered on $n\bar u$, not $\bar u$ (F6); for $n{=}1$ Maxwellian at $T_0$: `p` $= T_0$ under convention (b) ✓, would read $2T_0$ under (a) | +| :143 | `q` $= \int (v-n\bar u)^3 f\,dv$ | heat flux | same centering issue | +| :144 | `-flogf` $= \int f\log|f|\,dv$ | "entropy" | **missing minus sign** (F7) | +| :307 | `mean_-flogf` $= \langle\int -|f|\log|f|\,dv\rangle$ | entropy | has the minus; opposite sign to :144 (F7) | +| :303–306, 308 | `mean_P` $=\langle\int v^2f\rangle$, `mean_j` $=\langle\int vf\rangle$, `mean_n`, `mean_q` $=\langle\int v^3 f\rangle$, `mean_f2` | raw moments | `mean_j` and field "v" are the same integrand under different names (F6) | +| :311–312 | `mean_de2` $=\langle de^2\rangle$, `mean_e2` $=\langle e^2\rangle$ | field energy proxies | no ½; kinetic `mean_P` also lacks ½, consistently — internal conservation OK, but no combined energy monitor exists (F12) | +| :154, :313 | `pond` $= -\tfrac12\partial_x a^2$ | ponderomotive | ½ present and correct ✓ | + +--- + +## 4. Findings + +Ordered by severity. **Status: CONFIRMED** = derivation and code verified; **SUSPECTED** = probable issue, evidence stated. + +### F1 — CONFIRMED (root cause): `normalization.py` $v_0$ is $\sqrt2$ larger than the engine's velocity unit + +`normalization.py:91` ($v_0 = \sqrt{2T_0/m_e}$) vs. the engine convention $\sqrt{T_0/m_e}$ established by `helpers.py:69`, the collision operators, the dispersion tests, and all configs (§2). Since the physical temperature of the simulated plasma is $T_{phys} = m_e\,\sigma_{code}^2\,v_0^2 = T_{0,code}\cdot(m_e v_0^2)$: + +- **Under (a) as written: $T_{phys} = 2\,T_{0,code}\,T_{0,ref}$.** A `normalizing_temperature: 2000eV`, `T0: 1.0` run is a 4000 eV plasma. +- $L_0 = \sqrt2\lambda_{De}$, so all `dim="x"`/`dim="k"` string conversions carry a spurious $\sqrt2$ (box sizes, gradient scale lengths in `datamodel.py:159–177`, laser $k_0$ at `simulation.py:81`). +- The species-config `T0` parameter is effectively in units of $2T_{0,ref}$ — half the naïve expectation. + +The dimensionless dynamics of numeric-input runs are unaffected (the solver core never sees $v_0$); the damage is to the *physical interpretation* of every run and to dimensional-string inputs. + +### F2 — CONFIRMED (propagation): logged units and regression fixtures lock in the $\sqrt2$ values + +`write_units()` (`modules.py:119–141`) logs `v0`, `x0`, `box_length`, `c_light`, `nuee` computed under convention (a). For the 2000 eV resonance case the fixtures record `c_light: 11.302`, `v0: 2.652e7 m/s`, `x0: 12.14 nm` (`tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml:125–132`, same in the fokker-planck and multispecies fixtures); the convention-(b) truth values are $c/v_{th} = 15.98$, $v_{th} = 1.876\times10^7$ m/s, $\lambda_{De} = 8.585$ nm. Additionally `nuee` is evaluated at $T_{0,ref}$ while the plasma actually simulated is at $2T_{0,ref}$, so the logged collisionality does not describe the simulated plasma. Any fix of F1 must regenerate these fixtures. + +### F3 — CONFIRMED: EM branch — $\hat c$ is $\sqrt2$ too small; partial cancellation hides it for wavelength drivers only + +$\beta = 1/\hat c$ with $\hat c = c/v_0$ (`modules.py:58,121`) feeds the wave solver and the CFL limit. Under (a), $\hat c = c/(\sqrt2 v_{th})$ — $\sqrt2$ smaller than the engine's thermal unit warrants. For **wavelength-specified** drivers the $\sqrt2$'s cancel in the product $\hat c\hat k = \frac{c}{\sqrt2 v_{th}}\cdot\sqrt2 k\lambda_{De}$ — which is why `test_em_dispersion.py` and `srs.yaml`'s EM branch behave physically. For **numeric** AKW drivers (`srs-debug-small.yaml`: `k0: 1.0, w0: 2.79` used as-is at `simulation.py:59`), $\hat k$ is *not* rescaled while $\hat c$ still carries the $\sqrt2$, so the EM dispersion $\hat\omega^2 = \hat n_e + \hat c^2\hat k^2$ is $\sqrt2$-inconsistent with the ES branch's $\hat k = k\lambda_{De}$ interpretation. Mixed-input SRS runs are therefore internally inconsistent between branches. + +### F4 — CONFIRMED: Krook target Maxwellian is hard-coded to $e^{-v^2/2}$ (convention-locked, $T{=}1$, electron grid only) + +`fokker_planck.py:245`. Three separate problems: (i) it bakes in convention (b) with variance exactly 1, so it must be changed **in lockstep** with any convention fix or it will spuriously heat/cool by 2× in temperature; (ii) even today, it ignores the species `T0` and `mass` — any species initialized at $T_0 \ne 1$ (e.g. `twostream.yaml`, $T_0 = 0.2$) is dragged toward $T = 1$; (iii) it is built on the electron grid and applied only to `"electron"` (`fokker_planck.py:171`). It is a fixed-target BGK operator: conserves density only, not momentum or energy. Currently `is_on: False` in shipped configs, so latent. + +### F5 — CONFIRMED: `Dougherty.compute_vbar` is missing the $1/n$ normalization + +`fokker_planck.py:81` returns $\int v f\,dv = n\bar u$ while its docstring claims "Mean velocity ⟨v⟩". `discrete_temperature` (`driftdiffusion.py:101ff`) *does* divide by $\int f\,dv$ — the two moments are asymmetric. Consequence: the drag centers on $n\bar u$, the operator relaxes toward $e^{-\beta(v - n\bar u)^2}$, and **momentum is not conserved where $n(x) \ne 1$**; $T_{target}$ is also biased by the wrong centering. Negligible for the shipped EPW/SRS configs ($n \approx 1 \pm 10^{-4}$), real for bump-on-tail / large density perturbations with collisions on. Fix: divide by $\int f\,dv$. + +### F6 — CONFIRMED: storage central moments centered on $n\bar u$, not $\bar u$; misleading names + +`storage.py:140–143`: the field moment `"v"` is $\int v f\,dv = n\bar u$ (a flux), and `p`, `q` are centered on it. Only exact when $n = 1$; biased for `nlepw-ic.yaml` (10% density perturbation) and `bump-on-tail.yaml`. Same integrand is named `"v"` in field moments but `"mean_j"` in scalars (`storage.py:304`) — one of the labels is wrong; `j` is the honest name (or divide by $n$ and keep "v"). Note this is the same class of bug as F5, appearing independently in two places. + +### F7 — CONFIRMED: the two `-flogf` entropy diagnostics have opposite signs + +Field moment `storage.py:144` computes $+\int f\log|f|\,dv$ (no minus, uses $f$); scalar `storage.py:307` computes $-\int|f|\log|f|\,dv$. For $f > 0$ these are exact negatives. The label `-flogf` matches the scalar; the field version is sign-flipped. + +### F8 — SUSPECTED (intent unclear): super-Gaussian $\alpha$ fixes $\langle v^4\rangle/\langle v^2\rangle$, not the variance, for $m \ne 2$ + +`helpers.py:72`: $\alpha = \sqrt{3\Gamma(3/m)/\Gamma(5/m)}$ normalizes the kurtosis ratio $\langle v^4\rangle/\langle v^2\rangle = 3v_{th}^2$ for all $m$, but the variance equals $v_{th}^2$ only at $m = 2$. Numerically the realized variance is $\{1.240, 1.371, 1.449\}\times T_0/m$ for $m = \{3,4,5\}$ — so the measured `p`/`n` moment will *not* equal the input `T0` for super-Gaussian species. May be an intentional flat-top-EDF temperature definition; if so it should be documented, because the `T0` config docstring ("Temperature") and the second-moment diagnostic disagree with it. All shipped configs use $m = 2$, where it is exact. + +### F9 — CONFIRMED (gap/trap): collision-frequency normalization chain is entirely manual + +- `approximate_ee_collision_frequency` returns **Hz** and is only *logged* (`modules.py:119,132`); nothing converts it to code units or feeds it to the FP operator. +- The FP/Krook rate magnitudes (`baseline`, `bump_height`) are taken as **raw floats** (`functions.py:97–98`), not passed through `normalize()` — users must supply $\nu$ already in code units ($1/\omega_{p0}$), which is nowhere documented (`config.md` lists `baseline` with no units). +- The correct conversion is $\hat\nu = \nu_{ee}[\mathrm{Hz}]\cdot\tau = \nu_{ee}/\omega_{p0}$ with **no** $2\pi$ (the NRL rate is a true s⁻¹ rate); `_tf1d/modules.py:78` has this pattern (`nuee_norm = nuee/wp0`), `_vlasov1d` has no analog. A user thinking in cyclic frequency is one step from a silent $2\pi$ error. +- The LB/Dougherty $\nu$ and the NRL $\nu_{ee}$ agree only up to an O(1) factor — identifying them silently is itself an approximation worth a docs note. +- $\nu$ is a prescribed space-time envelope; it does not track the local $n(x)/T(x)^{3/2}$ (limitation, not a bug). + +Suggested: log a `nuee_norm` alongside `nuee`, and document `baseline`'s units. + +### F10 — CONFIRMED: species `T0` and drift `v0` bypass the unit machinery + +`simulation.py:213–214` take `float(cfg.T0)`, `float(cfg.v0)`; `datamodel.py:20–21` type them as bare floats. The `normalize(dim="temp")` and `dim="v"` branches are dead code for `_vlasov1d`. Consequences: `T0: "500 eV"` is impossible (only the global `normalizing_temperature` is dimensional), and the drift `v0` is expressed in units of the reference $v_0$ while the thermal width is in $\sigma$ units — under convention (a) these differ by $\sqrt2$ *within the same distribution*, a genuine user trap (drift "1.0" is not "one thermal width"). + +### F11 — CONFIRMED (latent): `dx = xmax/nx` ignores `xmin` + +`grid.py:55` should be `(xmax − xmin)/nx`; the x-axis itself (`grid.py:74`) spans $[x_{min}, x_{max}]$ with the wrong spacing when $x_{min} \ne 0$, and the `kx` grid (`grid.py:77`) and the semi-Lagrangian interpolation period (`pushers/vlasov.py:31`, `period = xmax`) inherit the error. `modules.py:112` uses the correct $(x_{max}-x_{min})$, highlighting the discrepancy. All shipped configs use `xmin: 0.0`, so latent but real. + +### F12 — Minor items (confirmed, low impact) + +1. **`AmpereSolver` class docstring** (`field.py:230`) claims $j = \sum_s (q_s/m_s)\int vf_s\,dv$; the code (`field.py:263–264`) correctly omits $1/m_s$. Docstring bug only. +2. **Energy diagnostics lack ½ and a total**: `mean_P`, `mean_e2`, `mean_de2` are $2\times$ the respective energies (consistently, so `mean_P + mean_e2` is still conserved), and no diagnostic sums kinetic + field energy. A dedicated conservation monitor would have caught F1 earlier. +3. **`dw0 = 0.0 # ???`** placeholder at `simulation.py:87` for the intensity/wavelength driver. +4. **Dead config knob**: `GridConfig.c_light` (`datamodel.py:73`; set in `wavepacket.yaml`) is never read — `c_light`/`beta` are always recomputed from the normalization. +5. **`grid.t` axis spacing**: `nt = int(tmax/dt + 1)`, `tmax = dt·nt` overshoots the request by up to ~dt and `linspace(0, tmax, nt)` has spacing $\ne$ dt. Cosmetic (saves use their own axes). +6. **EM plasma term is electron-only** (`vector_field.py:292–301`, `field.py:152`): ions omitted from the transverse current ($\sim m_e/m_i$ — fine, but an asymmetry vs. Poisson/Ampère which include all species). +7. **Point-source amplitude** uses the vacuum $k = \omega/c$ (`field.py:67`); realized plasma amplitude larger by $k_{vac}/k_{plasma}$ — already documented in the class docstring. +8. **Half-cell axis convention**: `vax`/`x` hold cell *centers* while the extents name the *edges* ($\pm v_{max}$) — keep in mind when labeling axes. + +### Verified correct (checked because they looked suspicious, and passed) + +- The entire collisionless solver core: all coefficients unity, signs consistent (§3.3). +- The ponderomotive chain: the $-\tfrac12\partial_x a^2$, the $(q^2/m^2)$ factor, and the $c/v_0$ scaling of $a_0$ combine *exactly* — no cycle-average ½ missing, no hidden $\sqrt2$. +- The Buet "factor of 2" in the Dougherty operator: correctly absorbed by $\beta = 1/2T$; equilibrium width equals the measured variance exactly. +- The $E_x$ driver amplitude $E = \omega a_0$: consistent vector-potential convention, matches the docs. +- Midpoint-rule ($\mathrm{sum}\cdot dv$) integration: used uniformly in init, moments, and field solves. + +--- + +## 5. Recommended resolution and lockstep checklist + +**Adopt convention (b) globally**: in `electron_debye_normalization`, set $v_0 = \sqrt{T_0/m_e}$ (so $L_0 = \lambda_{De}$), and redefine `vth_norm()` accordingly. This leaves the initializer, all collision operators, all tests, and all numeric configs untouched and correct, and makes the logged physical units true. + +Things that must change together / be re-verified: + +1. `normalization.py:91–92` ($v_0$, $x_0$) — **and nothing else in that file. Do NOT redefine `vth_norm()` (:48–50)**: the codebase-wide audit (`ADEPT_CONVENTIONS_AUDIT.md`, finding C1) found that `vth_norm()` has exactly four callers, all in `vfp1d/`, which is self-consistently built on the $\sqrt{2T/m}$ convention (its initializer carries a compensating factor). Changing `vth_norm()` does nothing for `_vlasov1d` (never called here) and would silently halve VFP-1D's initialized temperature. +2. Regenerate the three `*_derived_config.yml` regression fixtures (they lock `c_light`, `v0`, `x0` at the $\sqrt2$ values). +3. Re-check every config that uses **dimensional string** inputs (`srs.yaml`: `xmax: 100um`, gradient scale length `200um`, laser wavelength/intensity) — their physical meaning shifts by $\sqrt2$ (they become *correct*; the previously-inferred physical parameters of past runs were what was wrong). +4. `c_light`/`beta` changes for all EM runs: re-validate `test_em_dispersion.py`, the point-source amplitude test, and numeric-driver SRS configs (`srs-debug-small.yaml` `w0`/`k0` values were presumably tuned under the old $\hat c$ — see F3). +5. Krook target (`fokker_planck.py:245`) — no change needed under (b), but fix its $T_0$/species handling anyway (F4) so it stops being convention-locked. +6. Docs: state the normalization explicitly in `docs/source/solvers/vlasov1d/` (velocity unit $=\sqrt{T_0/m_e}$, $L_0 = \lambda_{De}$, $\hat k = k\lambda_{De}$, Maxwellian $e^{-v^2/2}$ at $T{=}1$), and document `baseline` units (F9) and the drift-`v0` units (F10). +7. ~~Audit sibling solvers for the same $v_0$ dependency~~ **Done — see `ADEPT_CONVENTIONS_AUDIT.md`** (codebase-wide audit with per-solver immunity verdicts). Summary: `electron_debye_normalization` is consumed only by `_vlasov1d`, `_vlasov2d`, `_pic1d` — all σ-convention, all fix-safe, and only `_vlasov1d` has fixtures to regenerate. `vfp1d` genuinely runs in $\sqrt{2T/m}$ units but via `laser_normalization` + `vth_norm()`, so it is untouched by the scoped fix (see amended item 1). `_tf1d` has a private inline duplicate of the √2 (`_tf1d/modules.py:62`) that must be fixed separately. All other modules are immune. + +Independent bug fixes (any order, no convention coupling): F5 (`compute_vbar` $1/n$), F6 (moment centering + naming), F7 (`-flogf` sign), F11 (`dx` with `xmin`), F12.1 (docstring), F12.2 (add an energy-conservation scalar). + +### Suggested verification runs after any fix + +- Landau damping and ion-acoustic tests must still pass unchanged (they pin convention (b)). +- A Dougherty-collisions run with a strong density perturbation: check $\partial_t\int v f\,dv \approx 0$ (exercises F5). +- An SRS wavelength-driver run: confirm backscatter resonance moves to the physically correct location once `c_light` and the µm conversions change together. +- Compare the second-moment temperature diagnostic against `T0` for an $m = 4$ super-Gaussian to decide F8's intended convention. + +--- + +*Audit method: four parallel independent reviews (initialization/grids/derived quantities; solver core; collisions; diagnostics/docs/tests/configs), followed by cross-checking of all load-bearing claims against the source. No source files were modified.* From d3a47c9a7adb98aae397a328060dc1ee25dff2c6 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 20 Jul 2026 11:52:06 -0700 Subject: [PATCH 48/49] jk: removed audit files --- .../2026-07-16_ADEPT_CONVENTIONS_AUDIT.md | 225 ---------------- .../2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md | 249 ------------------ 2 files changed, 474 deletions(-) delete mode 100644 docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md delete mode 100644 docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md diff --git a/docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md b/docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md deleted file mode 100644 index 4b1ac623..00000000 --- a/docs/audits/2026-07-16_ADEPT_CONVENTIONS_AUDIT.md +++ /dev/null @@ -1,225 +0,0 @@ -# ADEPT Codebase-Wide Physics & Conventions Audit - -**Date:** 2026-07-16 -**Scope:** every solver module and shared physics file in `adept/` — `_vlasov2d`, `_pic1d`, `_tf1d`, `_lpse2d`, `_spectrax1d`, `_hermite_poisson_1d`, `vfp1d`, `osiris`, `vlasov1d2v`, plus the shared `normalization.py`, `electrostatic.py`, `functions.py`, `driftdiffusion.py`, `utils.py`, and all tests/configs. -**Companion document:** `VLASOV1D_CONVENTIONS_AUDIT.md` (the `_vlasov1d` deep audit that established the root cause). This document extends that audit to the rest of the codebase and answers: *which solvers are immune?* -**Method:** six parallel independent module audits, each cross-checked against source; load-bearing claims re-verified by hand. No source files modified. - ---- - -## 1. Executive summary - -**The √2 bug is confined to the `electron_debye_normalization` consumers — and even there, only to the unit-conversion/reporting layer.** The dimensionless *dynamics* of every solver in the codebase are internally self-consistent. The full damage inventory of `normalization.py:91` ($v_0 = \sqrt{2T_0/m_e}$, $L_0 = \sqrt2\lambda_{De}$): - -- `_vlasov1d`, `_vlasov2d`, `_pic1d`: logged physical units √2-wrong, physical temperature 2× the label, dimensional-string inputs √2-off, EM-branch $\hat c$ √2-small (details per module below). -- `_tf1d`: **independently re-implements the identical √2** in its own `write_units` (`_tf1d/modules.py:62`) — same symptom, separate code, needs its own fix. -- Everyone else: **immune** (see the table). - -**One critical amendment to the fix prescription in the vlasov1d audit:** `vth_norm()` (`normalization.py:48–50`) must **NOT** be redefined. Grep-verified, it has exactly four callers, all in `vfp1d/`, which is built self-consistently on the $v_{th} \equiv \sqrt{2T/m}$ (most-probable-speed) convention — its initializer carries a compensating factor (`vfp1d/helpers.py:111`: $\alpha = \sqrt{3\Gamma(3/m)/(2\Gamma(5/m))} = 1$ at $m{=}2$, vs. `_vlasov1d`'s $\sqrt2$) so the realized variance is exactly $T/m$. Dropping the 2 in `vth_norm()` would do nothing for `_vlasov1d` (which never calls it) and would **silently halve VFP-1D's initialized temperature** while leaving its transport coefficients at $T_0$. The safe fix touches only `electron_debye_normalization` lines 91–92. - -### 1.1 Immunity table (the answer to "which solvers are immune") - -"Bug" = the current √2 in `normalization.py:91`. "Fix" = changing `electron_debye_normalization` to $v_0=\sqrt{T_0/m_e}$ (only; `vth_norm()` untouched). - -| Solver | Working vth convention | Affected by bug NOW? | Affected by fix? | Notes | -|---|---|---|---|---| -| `_vlasov1d` | $\sqrt{T/m}$ (σ) | **YES** — logged units, physical T (2×), dimensional strings, EM $\hat c$; fixtures lock wrong values | Fix-safe for tests; must regenerate 3 fixtures | See companion doc | -| `_vlasov2d` | $\sqrt{T/m}$ (σ) | **YES (partial)** — same chain: logged units, physical T, EM $\hat c$; ES dynamics immune | **Fix-safe** — all 3 tests self-consistent, **no fixtures** | Cleanest of the affected: already fixed several vlasov1d bugs | -| `_pic1d` | $\sqrt{T/m}$ (σ) | **Logged units only** — dynamics immune (all inputs numeric) | **Fix-safe** — no fixtures, no test depends on it | Loader ≡ vlasov1d initializer convention | -| `_tf1d` | $\sqrt{T/m}$ (σ) | Not via `normalization.py` — but its **own private copy** of the √2 (`modules.py:62`) corrupts its logged units identically | Fix does **nothing** here; `modules.py:62` needs its own `2.0*T0 → T0` | Dynamics/tests fully decoupled from both | -| `_lpse2d` | $\sqrt{T/m}$ (σ), physical units | **IMMUNE** — never touches `normalization.py` | **IMMUNE** | Own ps/µm unit system; Bohm-Gross & Landau literature-exact | -| `_spectrax1d` | AW-Hermite $\alpha=\sqrt{2T/m}$ (self-consistent) | **IMMUNE** — zero imports of `normalization.py` | **IMMUNE** | α is a raw config float; physics identical to σ-convention | -| `_hermite_poisson_1d` | AW-Hermite $\alpha=\sqrt{2T/m}$ (self-consistent) | **IMMUNE** — zero imports | **IMMUNE** | Not even registered in ergoExo dispatch | -| `vfp1d` | $\sqrt{2T/m}$ (most-probable, self-consistent) | **IMMUNE** — uses `laser_normalization` ($v_0=c$), never `electron_debye_normalization` | **IMMUNE if** `vth_norm()` is left alone; **BROKEN if** `vth_norm()` is "fixed" | The reason the fix must not touch `vth_norm()` | -| `osiris` wrapper | OSIRIS native ($u_{th}=\sqrt{T/m}/c$, in-deck) | **IMMUNE** — `skin_depth_normalization` ($T_0$=None, $v_0=c$); no eV→uth conversion exists in the wrapper | **IMMUNE** | `vth_norm()` unreachable (would TypeError on $T_0$=None) | -| `vlasov1d2v` | $\sqrt{T/m}$ (σ), $m$≡1 | **IMMUNE by orphan status** — not in ergoExo dispatch, would KeyError before running | **IMMUNE** | Legacy code; carries many un-fixed bugs (§3.9) | - -**Fully immune to both bug and fix:** `_lpse2d`, `_spectrax1d`, `_hermite_poisson_1d`, `osiris`, `vfp1d` (conditional on the fix not touching `vth_norm()`), and `vlasov1d2v` (by virtue of being unrunnable). -**Affected now:** `_vlasov1d`, `_vlasov2d`, `_pic1d` (via `normalization.py`) and `_tf1d` (via its private duplicate). -**Every solver's dimensionless dynamics are unaffected in shipped configs/tests** — the corruption is confined to logged units, physical interpretation of results, dimensional-string inputs, and the EM-branch $\hat c$. - -### 1.2 Why no test ever caught the √2 - -The bug survived every test suite in the repository, and understanding why is itself a finding — the test suites, by construction, cannot see this class of error: - -1. **Every physics test works entirely in dimensionless code units.** Landau damping, ion-acoustic, Bohm-Gross, two-stream, gyro — all set $k$, $\omega$, $T_0$, box sizes as numeric floats, which pass through `normalize()` untouched (`normalization.py:57–58`). The buggy conversion layer is simply never on the tested code path. A purely numeric run *is* a correct simulation under the σ-convention; only its physical labels are wrong, and no test reads the labels. -2. **The theory references the tests compare against live in the same code-unit world.** The kinetic roots come from `electrostatic.py` with `maxwellian_convention_factor=2` and $v_{th}=1$ — i.e., the reference is expressed in the engine's internal convention, so agreement validates internal consistency, not the dimensional dictionary. The tests pin *which* convention the engine uses (invaluable for this audit) but cannot detect that `normalization.py` speaks a different one. -3. **Where a test does touch the dimensional layer, it is self-referential.** `_vlasov2d`'s `test_em_dispersion` computes $\hat c$ from the same `electron_debye_normalization` the solver uses and places the driver at $\omega^2 = 1+\hat c^2 k^2$ with that same $\hat c$ — bug and reference move in lockstep, so it passes with the wrong $c$. -4. **The only tests that pin dimensional values pin the *wrong* ones.** The `_vlasov1d` `*_derived_config.yml` regression fixtures were generated under the buggy convention, so they actively *enforce* the √2 values (`c_light: 11.302`, `x0: 12.14 nm`, …) — a fix makes tests fail, not the bug. -5. **No solver has an absolute physical-units validation** (e.g., a Landau rate checked against a rate in Hz for a stated density and temperature, or an SRS resonance checked against a wavelength in nm), and none has a total-energy conservation diagnostic. Either would have caught the mismatch the first time a dimensional input mattered. - -The general lesson: dimensionless-physics tests validate the engine; only a test that crosses the units boundary — dimensional input in, dimensional observable out, compared against an independent physical reference — can validate the normalization layer. The repo currently has zero such tests. Adding one per normalization entry point (a "round-trip units test") is the cheapest structural guard against recurrence, and belongs in the fix PR alongside item 6 of §4. - -### 1.3 The three thermal-velocity conventions in the codebase - -All are *internally* self-consistent within their modules; the mixing hazard is at module boundaries and in shared helpers: - -| Convention | Definition | Used by | -|---|---|---| -| σ (RMS / std-dev) | $v_{th} = \sqrt{T/m}$, Maxwellian $e^{-v^2/2}$ at $T{=}1$, $L_0=\lambda_{De}$ | `_vlasov1d`, `_vlasov2d`, `_pic1d`, `_tf1d`, `_lpse2d`, `vlasov1d2v` dynamics; `electrostatic.py` (mcf=2 default); OSIRIS `uth` | -| Most-probable | $v_{th} = \sqrt{2T/m}$ | `vfp1d` (with compensating init factor); `vth_norm()`; **`electron_debye_normalization` (the outlier)**; `_tf1d/modules.py:62` (private duplicate, diagnostics-only) | -| AW-Hermite scale | $\alpha = \sqrt{2T/m}$ as basis parameter; represented Maxwellian still has variance $T/m$ | `_spectrax1d`, `_hermite_poisson_1d` (Schumer–Holloway / Parker–Dellar standard) | - ---- - -## 2. Normalization wiring map - -Grep-verified callers of each `normalization.py` entry point: - -| Entry point | Definition | Callers | -|---|---|---| -| `electron_debye_normalization` | $v_0=\sqrt{2T_0/m_e}$ ← **the bug**, $L_0=v_0/\omega_{p0}$ | `_vlasov1d`, `_vlasov2d` (`modules.py:55`), `_pic1d` (`simulation.py:78`) | -| `laser_normalization` | $v_0=c$, $L_0=c/\omega_L$, $n_0=n_{crit}$, $T_0$ not self-consistent with $v_0$ (documented) | `vfp1d` (`base.py:17`) only | -| `skin_depth_normalization(_from_frequency)` | $v_0=c$, $L_0=c/\omega_{p0}$, $T_0$=None | `osiris` only | -| `vth_norm()` | $\sqrt{2T_0/m_0}/v_0$ — hard-codes most-probable | **`vfp1d` only** (grid.py:118, base.py:93, 98, 152) | -| `speed_of_light_norm()` | $c/v_0$ | `_vlasov1d`, `_vlasov2d`, `_pic1d` (√2-tainted); `osiris` (=1, exact) | -| `normalize(s, norm, dim)` | numeric passthrough; strings via $L_0$/τ/$v_0$/$T_0$ | vlasov family, `functions.py`, `vfp1d/helpers.py` profiles | -| — (no normalization import) | | `_tf1d` (own inline pint), `_lpse2d` (own unit system), `_spectrax1d`, `_hermite_poisson_1d` (raw α floats), `vlasov1d2v` (reads cfg keys nothing populates) | - -Shared physics kernels: -- `electrostatic.py` — dispersion/Z-function utilities. Self-consistent under σ-convention: `maxwellian_convention_factor=2` (default, used by **every** caller in the repo) means $f_0\propto e^{-v^2/2v_{th}^2}$, $\xi=\omega/(\sqrt2 k v_{th})$. This is the reference that pins the working convention of `_vlasov1d`, `_vlasov2d`, `_pic1d`, `_tf1d`, `_spectrax1d`, `_hermite_poisson_1d` tests. -- `driftdiffusion.py` — Dougherty/LB kernel; measures $T=\langle(v-\bar v)^2\rangle$ (or spherical $\langle v^4\rangle/3\langle v^2\rangle$), relaxes to that width. Convention-agnostic and self-adapting; verified clean in both audits. -- `functions.py` — envelope/profile layer for the vlasov family. Numeric inputs pass through; dimensional strings go through $L_0$ (√2-tainted until the fix). One bug found (§3.10). -- `utils.py` — no physics content. - ---- - -## 3. Per-module findings - -Finding IDs continue the companion doc's F-series where the bug is a copy; new IDs are per-module. - -### 3.1 `_vlasov2d` — affected now (units/EM), fix-safe, partially cleaned-up lineage - -Convention: σ, identical to `_vlasov1d` (`helpers.py:62` `v_th=√(T0/mass)`, same super-Gaussian α; 2-D init correctly normalized, $\iint f\,d^2v = n$, equal widths in $v_x,v_y$). Landau test pins σ=1 via the shared kinetic roots. Pushers all unity-coefficient; TE-mode Maxwell curl signs verified; initial 2-D Poisson correct; magnetic rotation $\theta=-(q/m)B_z dt$ verified by the gyro test. - -Bug-copy status vs `_vlasov1d`: **F5 fixed** (vbar divides by n, `fokker_planck.py:32`), **F6 fixed** (moments centered on $u=j/n$, `storage.py:100`), **F11 fixed** (`dx=(xmax-xmin)/nx`, `grid.py:60`), F7 N/A (no entropy diagnostic). **F4 still present** (Krook target $e^{-v_x^2/2}e^{-v_y^2/2}$ hard-coded, `fokker_planck.py:127–131`; latent, off in configs). **F10 still present** (`float(cfg.T0/v0x/v0y)`, `simulation.py:158–163`). **F8 present** (same α, latent at m=2). - -√2 exposure: calls `electron_debye_normalization` (`modules.py:55`) → logged units √2-wrong (F2 analog), physical T 2× label (F1 analog), EM $\hat c = c/v_0$ √2-small (F3 analog). `test_em_dispersion` passes *because it is self-consistent, not correct* — it recomputes the same wrong $\hat c$ for the driver. **No regression fixtures exist**, so the fix requires no fixture regeneration; all three tests pass unchanged (the EM test moves in lockstep with the fix). - -New (minor): `mean_KE` has the ½, field proxies $\langle E^2\rangle,\langle B^2\rangle$ don't (no combined conserved-energy diagnostic); `Txx/Tyy` are pressures ($n\cdot$variance) despite the T-name, and `T` omits the mass factor (fine for electrons); the separable per-axis Dougherty never isotropizes $T_x \leftrightarrow T_y$ (no cross-axis coupling operator exists in this module); dead config knobs `UnitsConfig.laser_wavelength`, `Z`. - -### 3.2 `_pic1d` — dynamics immune, logged units affected, fix-safe - -Convention: σ. The particle loaders (`helpers.py:34` quiet inverse-CDF, `:58` random rejection) use `v_thermal = np.sqrt(T0/mass)` — byte-for-byte the vlasov1d initializer convention. Pinned independently by `test_bohm_gross.py:34` (kinetic dielectric with vth=1) and `test_landau_damping.py:38–41` ($\omega^2=1+3k^2$, σ=1 Landau rate). - -Verified clean: deposit/gather use the identical B-spline kernel (momentum-conserving, self-force-free); uniform loading deposits exactly $n=1$ in expectation ($w=n_0L/N$, partition of unity); Poisson $E_k=-i\rho_k/k$ unity prefactor; KDK leapfrog and Yoshida4 push coefficients exact; ponderomotive $(q/m)^2(-\tfrac12\partial_x a^2)$ matches vlasov1d's verified chain. - -√2 exposure: `write_units` (`modules.py:44–57`) logs √2-wrong `v0, x0, c_light, box_length` (a 2000 eV epw run is physically 4000 eV). Dynamics immune: every shipped config and test input is numeric; loader reads raw `T0`. **No fixtures.** Fix is strictly an improvement — nothing breaks. - -New findings: -- **P1 (confirmed):** energy diagnostics mix extensivity: `mean_KE = 0.5·m·Σw v²` is a box integral with the ½; `mean_e2 = mean(e²)` is a per-cell mean without ½ or dx. `mean_KE + mean_e2` is not conserved; no valid energy monitor exists. -- **P2 (confirmed):** quiet and random loaders truncate velocity differently for drifting species — quiet clips to $[-v_{max}, v_{max}]$ absolute (`helpers.py:38`), random clips to $[v_0-v_{max}, v_0+v_{max}]$ (`helpers.py:65`). Same config, two different realized distributions; quiet wrongly clips the high-v side of a drifting Maxwellian. Latent for shipped cold-beam decks. -- **P3 (confirmed):** `mean_KE/mean_p` are sums but named `mean_*` (naming, cf. F6-class). -- Inherited: F11 (`dx=xmax/nx` while particle wrap and placement correctly use $x_{max}-x_{min}$ — the field grid and particle box disagree for $x_{min}\ne0$), F10, F8. - -### 3.3 `_tf1d` — decoupled from `normalization.py`, but carries a private copy of the √2 - -Convention: σ. Derived from the pushers: linearizing continuity + momentum ($-u\partial_x u - \frac{1}{n}\partial_x(p/m) - \frac{q}{m}E$) + adiabatic energy ($\gamma=3$) + Poisson gives $\omega^2 = 1+3k^2$, exactly what `test_resonance.py:19` asserts; the kinetic branch uses the shared roots (mcf=2). The Poisson sign ($E_k=+i\rho_k/k$, i.e. $\partial_x E=-\rho$) and the momentum sign ($-(q/m)E$) are *both* opposite the textbook and cancel — verified stable by derivation and by the passing test. Landau closure decays the field at exactly $2\,\mathrm{Im}\,\omega$, matching `test_landau_damping.py:67`. - -**T1 (confirmed, the headline):** `_tf1d/modules.py:62` re-implements `v0 = √(2·T0/m_e)` inline in its own `write_units()` (with its own pint registry — no import of `normalization.py`). Every tf1d run logs `v0, x0, c_light, beta, box_length, sim_duration` √2-wrong and implies 2× the stated temperature. Damage confined to logged diagnostics (no dimensional-string inputs exist in tf1d; the EM `WaveSolver` branch is commented out; no fixtures). **Fixing `normalization.py` does not fix this** — `modules.py:62` needs its own `2.0*T0 → T0` in lockstep. - -Other findings: -- **T2 (confirmed):** the kinetic-γ pressure closure (`pushers.py:195` `wr_corr=(wrs²-1)/k²` + γ forced to 1) yields $\omega^2 = (w_{rs}^2-1)T_0+1$ — exact only at $T_0=1$. Convention- and T₀-locked, same class as vlasov1d's Krook (F4). Latent (all configs use $T_0=1$). -- **T3 (confirmed):** `modules.py:78` `nuee_norm = nuee/wp0` — the *correct* Hz→code conversion — is computed and then **discarded** (never logged, never used). Meanwhile `physics..trapping.nuee` is an unrelated raw ML-input float. There is no collisional friction in the tf1d equations at all. (F9-class trap.) -- **T4 (confirmed):** `docs/source/usage/tf1d.md` momentum equation has a spurious $1/n$ on the E-force and mislabels the Poisson equation. Docs-only. -- **T5 (suspected, latent):** `EnergyStepper` evolves $p$ but advects/compresses $p/m$ — ambiguous mass normalization of the pressure variable; inconsistent with the momentum equation for mobile ions ($m=1836$). All configs have ions off. -- **T9 (minor):** `resonance_search.yaml` sets `ion.landau_damping: True` with `ion.is_on: False` (inert, misleading). - -### 3.4 `_lpse2d` — fully immune; thermally clean; unrelated bugs found - -Physical-units solver (ps/µm, Gaussian-cgs-derived) with its own `write_units` (`helpers.py:71`); zero imports of `normalization.py`. Convention: $v_{te} = c\sqrt{T_e/511}$ = $\sqrt{T_e/m_e}$ (σ) used consistently in the Bohm-Gross term ($e^{-i\,1.5\,v_{te}^2 k^2/\omega_{pe}\,dt}$, coefficient 3/2 correct for σ), the Landau damping rate (verified literature-exact: $\sqrt{\pi/8}(1+\tfrac32 k^2\lambda_D^2)(k\lambda_D)^{-3}e^{-3/2-1/(2k^2\lambda_D^2)}$), the sound speed, and the TPD threshold. Numerically cross-checked: the test config's driver $\omega = 1.5k^2v_{te}^2/\omega_{p0} = 19.7 \approx 20$ as configured. TPD/SRS coupling constants contain no thermal factor (convention-independent); laser $E_0=\sqrt{8\pi I/c}$ and WKB swelling $(1-n/n_c)^{-1/4}$ verified. - -Findings (none √2-class): -- **L1 (confirmed):** logged `lambda_D = vte/w0` (`helpers.py:111`) divides by the **laser** frequency instead of $\omega_{pe}$ — factor 2 too small at $n_c/4$. Diagnostic only (dynamics compute $\omega_{pe}^2/k^2v_{te}^2$ directly). -- **L2 (confirmed, latent runtime bug):** the `E2` electrostatic-driver path calls `self.epw.driver(...)` (`core/vector_field.py:84`) but the active `SpectralEPWSolver` has no `driver` attribute (only the commented-out `SpectralPotential` does) → `AttributeError` for any config with an `E2` driver. Latent only because `test_epw_frequency` is currently disabled (`pass` body) — meaning **the module's dispersion conventions have no live regression test**. -- **L3 (confirmed):** `core/trapper.py` is dead code (never instantiated by `SplitStep`); it is a stale duplicate of `_tf1d`'s trapper. Note: it uses `electrostatic.py`, *not* `driftdiffusion.py` — the vlasov1d audit's cross-module note claiming lpse2d uses the Dougherty kernel is wrong for the current tree. -- **L4 (confirmed):** the σ-convention is undocumented in the lpse2d docs — the main *future* √2 risk here is someone "harmonizing" it onto `normalization.py`. -- **L5 (minor):** the Landau/Bohm-Gross formulas are duplicated byte-for-byte in `SpectralPotential` and `SpectralEPWSolver` (drift risk); `nu_coll`'s `/2` (amplitude vs energy rate) is correct but easy to double-count. - -### 3.5 `_spectrax1d` and `_hermite_poisson_1d` — fully immune; self-consistent AW-Hermite basis - -Both use the asymmetrically-weighted Hermite basis (Schumer–Holloway / Parker–Dellar) with scale $\alpha = \sqrt{2T/m}$ taken as a **raw config float** — zero imports of `normalization.py`. The represented Maxwellian has variance $\alpha^2/2 = T/m$: physically identical to the σ-convention solvers. Verified by derivation (streaming ladder $\alpha\sqrt{n/2}$ + force ladder $\sqrt{2n}/\alpha$ + Ampère/Poisson coupling → Langmuir $\omega=\omega_{pe}$ exactly, Bohm-Gross $\omega^2=1+3(k\lambda_{De})^2$ with $\lambda_{De}=\alpha/\sqrt2$) and by tests: both modules' Landau tests set $\alpha_e = \sqrt2\,k\lambda_D/k$ explicitly against the shared mcf=2 kinetic roots (2%/5% tolerance for hermite-poisson; a cross-module test pins the two modules' E-coupling equal to 1e-12). - -The coincidence that $\alpha=\sqrt{2T/m}$ matches the `normalization.py:91` outlier is harmless — there is no code coupling in either direction. Both integrator paths (DoPri8 and Lawson-RK4 exponential operators) share identical ladder coefficients. The historical inverted E-coupling bug in hermite_poisson (`C[n+1]` vs `C[n-1]`) is fixed and regression-locked by three tests. - -Findings (all diagnostic-label/minor): **S1** logged `lambda_D` is the *total* (ion-dominated) Debye length while `k_norm` uses electron-only and hard-codes mode 1 (`base_module.py:186,189`); **S2** the ion "temperature" diagnostic is a velocity variance missing the $m_i$ factor ($T_i/m_i$, misleading next to $T_e$; `storage.py:386–405`); **S3** shipped `landau-damping.yaml` has `Nn: 4` — far too few Hermite modes to resolve the damping it is named for (tests override to 512/32); `_spectrax1d/helpers.py` is an all-stub file not on any code path. - -### 3.6 `vfp1d` — immune to the bug; the reason the fix must not touch `vth_norm()` - -Uses `laser_normalization` (`base.py:17`): $v_0 = c$, $L_0 = c/\omega_L$, $n_0 = n_{crit}$, $T_0$ = reference eV (documented as not self-consistent with $v_0$). Velocity grid in units of $c$. Thermal convention: **most-probable, $v_{th} = \sqrt{2T/m}$**, deliberately and self-consistently: - -- `vth_norm()` supplies $\sqrt{2T_0/m}/c$ (4 call sites, all vfp1d); -- the initializer's width factor `helpers.py:111` is $\alpha = \sqrt{3\Gamma(3/m)/(2\Gamma(5/m))}$ — note the extra `/2` vs the vlasov-family α — giving $\alpha=1$ at $m{=}2$ and realized variance exactly $T/m$ (the √2 and the /2 cancel); -- the v-grid extent `grid.py:118` `vmax = 8·vth_norm()/√2` = 8 standard deviations (the /√2 converts most-probable→σ); -- the IB coefficient `base.py:81` ($0.093373\,\lambda_{\mu m}^2/T_{keV}$ per $10^{15}$ W/cm²) normalizes $v_{osc}^2$ to $2T/m$, consistent; -- the temperature diagnostic `storage.py:317–322` uses the spherical variance $T = \langle v^4\rangle/3\langle v^2\rangle$ — reads the true temperature, no √2. - -Collision operators ($\nu_{ee}$ coefficient anchored to $v_0=c$ via $r_e$, Lorentz e–i with Epperlein–Haines Z*, Rosenbluth I/J integrals) are temperature-convention-independent. **No internal √2 inconsistency exists.** - -**C1 (confirmed, cross-module, high):** because items above are keyed to `vth_norm()` while the compensating α is a separate constant, redefining `vth_norm()` → $\sqrt{T_0/m}/v_0$ (as the vlasov1d audit's §5.1 originally suggested) silently initializes VFP-1D at **half** the intended temperature while all transport coefficients stay at $T_0$; the Spitzer/Epperlein–Haines gold test (`test_kappa_eh.py`, $\kappa\propto T^{5/2}$) would likely fail, and the pure-operator tests would *not* catch it. **The fix must be confined to `electron_debye_normalization:91–92`.** (The companion doc's §5.1 has been amended accordingly.) - -Other findings: **C2** `base.py:118` stores the bound method `norm.vth_norm` instead of calling it (harmless repr-string in cfg; wrong); **C3** `storage.py:170` hard-codes $9.09\times10^{21}$ cm⁻³ per $n_c$ — that is $n_{crit}(351\,\mathrm{nm})$; wrong labeling for any other `laser_wavelength` (should derive from `norm.n0`); **C4** = F11 copy (`grid.py:71` `dx=xmax/nx`); **C6/C7 (suspected, latent — IB not in any shipped config):** the production inverse-bremsstrahlung wiring looks under-normalized — `w0_norm` is identically 1.0 under laser normalization, the Langdon-factor argument $Z^2 n_i/(\omega_0 v^3)$ carries no collision-frequency coefficient, and `vosc2_per_intensity` (thermal-speed² units) is consumed as $v_{osc}^2$ in grid ($c$) units — a $\sim(c/v_{th})^2$ discrepancy between the unit tests (where the grid unit *is* the thermal speed) and production. Recommend a Spitzer-IB validation run before enabling IB in production. - -### 3.7 `osiris` wrapper — fully immune - -Uses `skin_depth_normalization(_from_frequency)` ($v_0=c$, $T_0$=None). The wrapper drives a native OSIRIS deck: **no eV→`uth` conversion exists anywhere in the wrapper** — `uth` values pass through verbatim from the deck, so the T→v √2 trap cannot occur. Verified exact: $a_0 \to E_{peak} = a_0\omega_L m_e c/e \to I = E^2\epsilon_0 c/2$ (matches $I\lambda_{\mu m}^2 = 1.37\times10^{18}a_0^2$, test-confirmed); `c_light = beta = 1.0` identically; box lengths in true skin depths; `units.yaml` deliberately omits temperature-dependent keys (consistent with $T_0$=None). Temperature diagnostics (`plots.py:1004` $T=\sum u_{th,i}^2$ in $m_ec^2$ units; `:1059` momentum variance) carry no spurious factors. `vth_norm()` is unreachable (would TypeError on $T_0$=None). - -**C8 (low, latent):** the adaptive-box feature's `reference_density = 0.25` "quarter-critical" (`density.py:75`) assumes the deck's $n_0 = n_{crit}$ (true for LPI decks with `omega0=1`); a deck normalized to a different density would silently mis-scale the density bounds (the *length* normalization stays correct). - -### 3.8 Shared `electrostatic.py` and `functions.py` - -`electrostatic.py`: self-consistent under the σ-convention throughout. $Z$ via Faddeeva, $Z' = -2(1+\xi Z)$ correct; `maxwellian_convention_factor=2` (default, used by every caller in the repo) ⇒ $f_0\propto e^{-v^2/2v_{th}^2}$, root returned as $\xi k v_{th}\sqrt{2} = \omega$. One cosmetic blemish (**T7**): the Newton `initial_root_guess` is ω-scaled ($\sqrt{\omega_p^2+3k^2v_{th}^2}$) but the root variable is ξ-scaled — converges anyway. - -`functions.py`: numeric inputs pass through; `EnvelopeFunction` amplitudes (`baseline`, `bump_height`) are raw floats (the F9 collision-frequency trap); `SineFunction.wavenumber` correctly uses `dim="k"`. **T8 (suspected):** `LinearFunction`/`ExponentialFunction` normalize `val_at_center` — a *density* — with `dim="x"` (`functions.py:159–161, 178–180`); dimensionally wrong for any string input. - -### 3.9 `vlasov1d2v` — orphaned legacy; a museum of the known bugs plus three new ones - -Not registered in the ergoExo dispatch (`_base_.py:294–333` → NotImplementedError); reads `cfg["units"]["derived"]` keys nothing populates → cannot run. Never calls `normalization.py`. Convention: σ (init `exp(-(v_x^2+v_y^2)/2T_0)`, mass≡1; the super-Gaussian machinery is commented out so `m` is silently ignored). - -Carries live copies of: **F5** (`fokker_planck.py:93,98` — vbar *and* the diffusion coefficient computed without ÷n), **F6 + F7** (`storage.py:152–158, 296`), **F10**, **F11** (`helpers.py:211`). F4 present but dead (Krook never instantiated). - -New bugs (would matter if ever revived): -- **N1 (confirmed):** the density profile is computed for every basis and folded into the ion background (`helpers.py:287`), but the line applying it to $f$ (`helpers.py:83–85`) is **commented out** — electrons always start uniform, so non-uniform configs begin with spurious net charge and no actual density perturbation in $f$. -- **N2 (confirmed):** the explicit Dougherty substeps apply ν twice (`fokker_planck.py:134–135, 141–142`: `dfdt = nu*ddx(...)` then `f + dt*nu*dfdt`) — collisional relaxation runs at $\nu^2$. -- **N4 (confirmed):** `integrator.py:234` passes the **nu_ee** envelope args when building the **nu_ei** profile (copy-paste) — the configured nu_ei profile is silently ignored. -- **N3 (design note):** without collisions the $v_y$ dimension is dynamically inert (no $v_y$ force, streaming by $v_x$ only) — the entire $n_{vy}$ grid is dead compute except through the FP operators. -- Minor: default scalars broadcast over the $v_y$ axis only (`storage.py:292–295`). - -Recommendation: either delete `vlasov1d2v` or quarantine it with a module-level comment; in its current state it is a trap for anyone who greps for reference implementations. - -### 3.10 Cross-module systemic observations - -1. **Driver amplitude conventions differ across the Vlasov family** (N5): `_vlasov1d` Ex driver $\propto \omega a_0$; `_vlasov2d` current source $\propto \omega^2 a_0$; `vlasov1d2v` $\propto |k| a_0$. Same config key, three physical meanings. Worth unifying or documenting per-solver. -2. **No solver has a valid total-energy conservation diagnostic.** vlasov1d/vlasov2d/pic1d all mix ½-factors and extensivity between kinetic and field terms. A per-solver conserved-energy scalar is the single cheapest guard that would have caught the √2 class of bugs. -3. **The F-series bugs are lineage-correlated:** `_vlasov2d` (newest refactor) fixed F5/F6/F11; `_pic1d` inherits F10/F11 by importing `_vlasov1d`'s grid/simulation; `vlasov1d2v` (oldest) has everything. When fixing `_vlasov1d`, fix the copies in the same PR: Krook targets (`_vlasov1d`, `_vlasov2d`, dead `vlasov1d2v`), `dx`/xmin (`_vlasov1d/grid.py:55`, `vfp1d/grid.py:71`, `vlasov1d2v/helpers.py:211`), raw-float T0/v0 (all three vlasov-family simulation.py files). -4. **Hz→code-unit collision-frequency conversion** exists correctly in exactly one place (`_tf1d/modules.py:78`) and is dead there; nowhere is it live. Every solver that takes a collision frequency takes it as a raw code-unit float with undocumented units. - ---- - -## 4. Consolidated fix checklist (supersedes §5 of the vlasov1d audit where they differ) - -**The √2 fix, safely scoped:** -1. `normalization.py:91–92`: $v_0 \to \sqrt{T_0/m_e}$, $x_0 \to \lambda_{De}$ — inside `electron_debye_normalization` **only**. -2. **Do NOT change `vth_norm()`** (`normalization.py:48–50`) — it is vfp1d-private and correct for vfp1d (C1). If its name is judged misleading, rename to `most_probable_speed_norm()`; do not change the value. -3. `_tf1d/modules.py:62`: `2.0*T0 → T0` (private duplicate, T1). -4. Regenerate the three `_vlasov1d` `*_derived_config.yml` regression fixtures (only vlasov1d has fixtures). -5. Fix the convention-locked Krook targets in the same PR (`_vlasov1d/…/fokker_planck.py:245`, `_vlasov2d/…/fokker_planck.py:127–131`): build from species T0/mass. -6. Re-validate: vlasov1d Landau + ion-acoustic (must pass unchanged), vlasov2d Landau/gyro/EM-dispersion (EM moves in lockstep), pic1d all three (unchanged), vfp1d `test_kappa_eh` (must pass unchanged — the sentinel that `vth_norm()` was left alone), numeric-driver SRS configs re-tuned per vlasov1d F3. -7. Document the per-solver conventions in each solver's docs page (σ vs most-probable vs AW-α; see §1.3) — the biggest residual risk is a future "harmonization" that imports the wrong convention into a currently-clean module (especially `_lpse2d`, L4). - -**Independent bug fixes, by priority:** -- High (live, wrong physics if exercised): vfp1d IB normalization audit before production use (C6/C7); lpse2d `E2` driver path + re-enable `test_epw_frequency` (L2). -- Medium (live, wrong diagnostics): lpse2d logged `lambda_D` (L1); pic1d/vlasov2d/vlasov1d energy-diagnostic normalization + add conserved-energy scalars; spectrax ion-T mass factor (S2); vfp1d hard-coded $n_{crit}(351\,nm)$ (C3); pic1d quiet-loader drift truncation (P2). -- Low / hygiene: tf1d docs (T4) and dead `nuee_norm` (T3 — log it instead of dropping it); `functions.py` `val_at_center` dim (T8); vfp1d bound-method store (C2); spectrax `Nn:4` config (S3); delete or quarantine `vlasov1d2v` and lpse2d's dead trapper (L3); unify driver-amplitude conventions or document them (N5). - ---- - -## 5. What was verified clean (highlights) - -- All solver cores' dimensionless coefficients: vlasov1d/2d pushers and field solves, pic1d push/deposit/gather/Poisson, tf1d fluid system (including its double sign cancellation, verified by derivation), Hermite ladder algebra in both spectral modules (verified by derivation to the Langmuir and Bohm-Gross limits), lpse2d Bohm-Gross/Landau/TPD/SRS coefficients against literature, vfp1d collision coefficients and temperature diagnostics, osiris a0↔intensity conversions. -- The shared `driftdiffusion.py` Dougherty kernel (both audits): self-adapting, Buet factor-2 correctly absorbed, spherical temperature correct. -- `electrostatic.py`'s Z-function algebra and mcf=2 semantics — the single reference that consistently pins six modules' tests to the σ-convention. - ---- - -*Companion deep-dive for `_vlasov1d`: `VLASOV1D_CONVENTIONS_AUDIT.md`. Method: six parallel independent module audits (one Opus agent per slice) with all load-bearing claims re-verified against source before inclusion. No source files were modified.* diff --git a/docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md b/docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md deleted file mode 100644 index 81a302a5..00000000 --- a/docs/audits/2026-07-16_VLASOV1D_CONVENTIONS_AUDIT.md +++ /dev/null @@ -1,249 +0,0 @@ -# Vlasov-1D Physics & Conventions Audit - -**Date:** 2026-07-16 -**Scope:** `adept/_vlasov1d/` (all files), plus the shared code it depends on: `adept/normalization.py`, `adept/driftdiffusion.py`, `adept/electrostatic.py`, and the `tests/test_vlasov1d/` suite and example configs. -**Trigger:** the discovery that `normalization.py:91` uses the $v_0 = \sqrt{2T_0/m}$ (most-probable-speed) convention while `_vlasov1d/helpers.py:69` uses $v_{th} = \sqrt{T_0/m}$ (RMS / standard-deviation convention). Neither has been changed yet; this audit determines which convention the module actually runs in, inventories every equation, and lists all inconsistencies found. - ---- - -## 1. Executive summary - -**The engine's working convention is $v_0 = \sqrt{T_0/m}$.** The distribution initializer, the Fokker–Planck operator, the Krook operator, the Landau-damping and ion-acoustic tests (the module's quantitative gold standards), and every numeric example config are all mutually consistent under: code velocity in units of the thermal *standard deviation* $\sqrt{T_0/m_e}$, code length unit $L_0 = \lambda_{De}$, code wavenumber $= k\lambda_{De}$, Maxwellian $\propto e^{-v^2/2}$ at $T=1$. - -**`adept/normalization.py:91–92` is the sole outlier** (`electron_debye_normalization`: $v_0 = \sqrt{2T_0/m_e}$, $L_0 = \sqrt{2}\,\lambda_{De}$). Because numeric config inputs pass through `normalize()` untouched, this does **not** corrupt the dynamics of numeric-input runs — but it means: - -- A run with `normalizing_temperature: 2000eV` and species `T0: 1.0` is physically a **4000 eV** plasma (factor 2 in temperature). -- Every dimensional **string** input converted with $L_0$ (box sizes in µm, gradient scale lengths, laser $k_0$) is off by $\sqrt{2}$. -- Every **logged** physical unit (`v0`, `x0`, `c_light`, `box_length`) is off by $\sqrt{2}$, and the regression fixtures currently lock in those wrong values. -- The EM wave speed $\hat c = c/v_0$ fed to the wave solver is $\sqrt2$ too small relative to the engine's thermal unit (with a partial cancellation for wavelength-specified drivers; see F3). - -**Recommended fix direction:** change `normalization.py` to $v_0 = \sqrt{T_0/m_e}$ (so $L_0 = \lambda_{De}$), *not* the initializer. Fixing `helpers.py:69` instead (to $\sqrt{T_0/2m}$) would break the currently-passing Landau-damping and ion-acoustic tests and invalidate every existing config's driver $(k_0, \omega_0)$ values. Section 5 lists everything that must move in lockstep. - -The collisionless solver core (`vector_field.py`, pushers) was verified **clean**: every coefficient is exactly unity under the declared normalization and is convention-independent. The convention bug lives entirely in the dimensional-conversion layer. - -Beyond the $\sqrt2$ issue, the audit found several independent bugs and traps, listed in Section 4 (notably: FP `compute_vbar` missing $1/n$; storage central moments centered on $n u$ instead of $u$; a sign inconsistency between the two `-flogf` entropy diagnostics; the super-Gaussian $\alpha$ fixing the wrong moment for $m\ne2$; the Krook target hard-coded to $T=1$; and the collision-frequency normalization chain being entirely manual). - ---- - -## 2. The two candidate conventions - -With $\omega_{p0} = \sqrt{n_0 e^2/(\epsilon_0 m_e)}$, $\tau = 1/\omega_{p0}$ in both cases: - -| Quantity | (a) `normalization.py` as written | (b) engine's actual convention | -|---|---|---| -| Velocity unit $v_0$ | $\sqrt{2T_0/m_e}$ (most-probable speed) | $\sqrt{T_0/m_e}$ (RMS / std-dev) | -| Length unit $L_0 = v_0/\omega_{p0}$ | $\sqrt{2}\,\lambda_{De}$ | $\lambda_{De}$ | -| Maxwellian at $\hat T = 1$ | $\propto e^{-v^2}$ (variance ½) | $\propto e^{-v^2/2}$ (variance 1) | -| Code wavenumber $\hat k$ | $\sqrt2\, k\lambda_{De}$ | $k\lambda_{De}$ | -| Bohm–Gross | $\hat\omega^2 = 1 + \tfrac{3}{2}\hat k^2$ | $\hat\omega^2 = 1 + 3\hat k^2$ | -| $\hat c = c/v_0$ (2000 eV) | 11.30 | 15.98 | - -**Why the test suite never caught this:** every physics test (Landau, ion-acoustic, EM dispersion) works in dimensionless code units — numeric inputs bypass `normalize()` entirely (`normalization.py:57–58`), so the buggy conversion layer is never on the tested path, and the theory references (`electrostatic.py`, mcf=2, $v_{th}=1$) are expressed in the engine's internal convention, so agreement validates internal consistency rather than the physical-units dictionary. Worse, the only tests that *do* pin dimensional values — the `*_derived_config.yml` regression fixtures — were generated under the buggy convention and therefore lock the wrong values in. No test crosses the units boundary (dimensional input in, dimensional observable out, against an independent physical reference), which is precisely the kind of test that must accompany the fix (see `ADEPT_CONVENTIONS_AUDIT.md` §1.2 for the full analysis). - -Evidence pinning convention (b) as the working one, strongest first: - -1. **Landau-damping test** (`tests/test_vlasov1d/test_landau_damping.py:24`) calls `electrostatic.get_roots_to_electrostatic_dispersion(1.0, 1.0, k)` — i.e. $\omega_{pe}=1$, $v_{th}=1$ with `maxwellian_convention_factor=2` (`adept/electrostatic.py:79,115,98`), which is the kinetic dielectric for $f_0 \propto e^{-v^2/2v_{th}^2}$, $\xi = \omega/(\sqrt2 k v_{th})$. The measured $\omega(k)$ matches these roots to 2 decimals only if the simulated Maxwellian has $\sigma = 1$. (E.g. $k=0.3 \Rightarrow \omega = 1.1598$, which is exactly the `w0` in `resonance.yaml` and `epw.yaml`; under convention (a) the resonance would sit near 1.07 and the test would fail.) -2. **Ion-acoustic test** (`test_ion_acoustic_wave.py`) uses $c_s^2 = ZT_e/m_i$ and $\omega^2 = k^2c_s^2/(1+k^2\lambda_D^2)$ with $\lambda_D = 1$ in code units — asserting $L_0 = \lambda_{De}$. -3. **Initializer** (`helpers.py:69–76`), **Krook target** (`fokker_planck.py:245`), and **Dougherty stationary state** (`fokker_planck.py:51,96` + `driftdiffusion.py`) all build/assume $e^{-v^2/(2T/m)}$, variance $T/m$. -4. Every numeric example config (`epw.yaml`, `resonance.yaml`, `wavepacket.yaml`, `bump-on-tail.yaml`) uses driver $(k_0,\omega_0)$ pairs consistent only with $\hat k = k\lambda_{De}$ and $\sigma = \sqrt{T/m}$. - ---- - -## 3. Equation inventory - -### 3.1 Normalization layer (`adept/normalization.py`) - -| Location | Expression | Meaning | Convention note | -|---|---|---|---| -| :88 | $\omega_{p0} = \sqrt{n_0e^2/\epsilon_0 m_e}$, $\tau = 1/\omega_{p0}$ | time unit | standard | -| **:91** | $v_0 = \sqrt{2T_0/m_e}$ | velocity unit | **outlier — convention (a)** | -| :92 | $x_0 = v_0/\omega_{p0} = \sqrt2\lambda_{De}$ | length unit | inherits $\sqrt2$ | -| :48–50 | `vth_norm()` $= \sqrt{2T_0/m_0}/v_0 = 1$ | "thermal velocity" | convention (a); **never called** in `_vlasov1d` | -| :52–53 | `speed_of_light_norm()` $= c/v_0$ | $\hat c$ | inherits $\sqrt2$ (F3) | -| :37–38 | NRL $\log\Lambda_{ee}$ | Coulomb log | uses reference $T_0$ | -| :45 | $\nu_{ee} = 2.91\times10^{-6}\, n_{cc}\log\Lambda\, T_{eV}^{-3/2}$ Hz | NRL e–e rate | **diagnostic only**, never converted to code units (F9) | -| :56–74 | `normalize()`: x/L0, t/τ, v/v0, T/T0, k·L0 | dim→code | numeric inputs pass through untouched; `dim="temp"` and `dim="v"` branches are dead code for `_vlasov1d` | - -### 3.2 Initialization & grids (`helpers.py`, `grid.py`, `modules.py`, `simulation.py`, `datamodel.py`) - -| Location | Expression | Meaning | Convention note | -|---|---|---|---| -| helpers.py:65–66 | `dv = 2 vmax/nv`; cell-centered `vax` on $[-v_{max}+dv/2,\ v_{max}-dv/2]$ | v-grid | edges at $\pm v_{max}$; half-cell center offset | -| **helpers.py:69** | $v_{th} = \sqrt{T_0/m}$ | thermal width | **convention (b)** | -| helpers.py:72 | $\alpha = \sqrt{3\,\Gamma(3/m)/\Gamma(5/m)}$ ($=\sqrt2$ at $m{=}2$) | super-Gaussian width factor | fixes $\langle v^4\rangle/\langle v^2\rangle = 3v_{th}^2$; variance $=v_{th}^2$ **only** at $m=2$ (F8) | -| helpers.py:74–76 | $f \propto \exp[-|(v-v_d)/(\alpha v_{th})|^m]$ | init EDF; $m{=}2$: $e^{-(v-v_d)^2/(2T_0/m)}$ | variance $T_0/m$ — convention (b) | -| helpers.py:80,84 | `f /= sum(f)·dv`; `f *= n_prof` | $\int f\,dv = n(x)$ | midpoint rule, consistent with all moments | -| grid.py:55 | `dx = xmax/nx` | cell width | **ignores `xmin`** (F11) | -| grid.py:59–60 | `dt = min(dt, 0.95·dx/c)`, $c = 1/\beta$ | EM CFL | $\hat c$ inherits F1's $\sqrt2$ | -| grid.py:65–66,75 | `nt = int(tmax/dt+1)`; `t = linspace(0, dt·nt, nt)` | time axis | overshoot ≤ dt; `grid.t` spacing ≠ dt (F12) | -| grid.py:74,77 | cell-centered x; `kx = 2π·fftfreq(nx, d=dx)` | spectral grid | inherits F11's dx for `xmin≠0` | -| modules.py:112 | `box_length = (xmax−xmin)·L0` | box size in µm | uses `xmin` correctly (unlike grid.py:55); value carries F1's $\sqrt2$ | -| simulation.py:213–214 | `v0 = float(cfg.v0)`; `T0 = float(cfg.T0)` | species drift & temperature | bare code floats, bypass `normalize()` (F10) | -| simulation.py:74–76 | $a_0 = a_{0,std}\cdot(c/v_0)$ | intensity → quiver velocity in $v_0$ units | internally consistent with pusher (§3.3) | -| simulation.py:79–85 | $\hat k_0 = k_{phys}L_0$, $\hat\omega_0 = \omega_{phys}\tau$ | wavelength driver → code units | $\hat k_0$ carries F1's $\sqrt2$; $\hat\omega_0/\hat k_0 = \hat c$ ✓ | -| simulation.py:87 | `dw0 = 0.0 # ???` | frequency offset placeholder | unresolved (F12) | - -### 3.3 Solver core (`solvers/vector_field.py`, `solvers/pushers/vlasov.py`, `solvers/pushers/field.py`) — **verified clean** - -The dimensionless system implemented, with every coefficient exactly unity under the declared normalization (and independent of the (a)/(b) choice, since $v_0$ cancels): - -$$\partial_t f_s + v\,\partial_x f_s + \frac{\hat q_s}{\hat m_s}\left(E - \frac{\hat q_s}{2 \hat m_s}\partial_x a^2\right)\partial_v f_s = C[f_s]$$ -$$\partial_x E = \sum_s \hat q_s \int f_s\,dv \ (+\ \text{static ion background}), \qquad \partial_t E = -\sum_s \hat q_s \int v f_s\,dv$$ -$$\partial_t^2 a = \hat c^2\,\partial_x^2 a - n_e\, a + S$$ - -| Location | Code | Check | -|---|---|---| -| vlasov.py:210–214, 220 | exact spectral shift $f(x - v\,dt)$ | coefficient $v_0\tau/L_0 = 1$ ✓ | -| vlasov.py:83–87, 127–129 | `force = q·e + (q²/m)·pond; accel = force/m` | $\hat q/\hat m$ E-push; ponderomotive $(\hat q^2/\hat m^2)$ ✓ | -| vector_field.py:411 | `pond = −0.5·∂ₓ(a²)` | exact instantaneous $-\frac{1}{2m}\partial_x p_\perp^2$ with $p_\perp = qA$; the $c/v_0$ scaling of $a_0$ (simulation.py:76) makes $a$ the quiver velocity in $v_0$ units — **no hidden $\sqrt2$; verified exact** | -| field.py:203–224 | $E_k = -i\rho_k/k$ | Poisson prefactor $e^2n_0/(\epsilon_0 m_e\omega_{p0}^2) = 1$ ✓ | -| field.py:263–264, 280 | $E^{n+1} = E^n - dt\,j$ | Ampère prefactor 1 ✓ | -| field.py:337–343 | $\Delta E_k = -i\frac{q}{k}\int dv\, f_k(e^{-ikv\,dt}-1)$ | exact (Hamiltonian) Ampère along free streaming ✓ | -| field.py:146–153 | leapfrog wave eq., plasma term $-n_e a$ | dispersion $\hat\omega^2 = \hat n_e + \hat c^2\hat k^2$ ✓; $n_e$ is electron-only (~$m_e/m_i$ approx., F12) | -| vector_field.py:292–301, 341 | $n_e = -\,\hat q_e \int f_e\,dv$, time-centered | sign keyed to electron charge $-1$ ✓ | -| field.py:21–26 | $E_x$ driver $= (\omega_0{+}\delta\omega)\,a_0\sin(k_0x - \omega t)$ | vector-potential amplitude convention ($E = -\partial_t A$); matches docs (config.md:312) ✓ | -| field.py:67–69, 78–80 | point source $F_0 = 2\omega\hat c\,a_0$ | reproduces amplitude $a_0$ in vacuum; plasma value larger by $k_{vac}/k_{plasma}$, documented in docstring ✓ | - -Charge-sign conventions were checked across Poisson / Ampère / wave-equation and the force term: consistent, no compensating double-error. Current density correctly has **no** mass factor in code (but see docstring bug F12). - -### 3.4 Collisions (`solvers/pushers/fokker_planck.py`, `adept/driftdiffusion.py`) - -Operator: $\partial_t f = \nu\,\partial_v[(v - \bar v)f + T\,\partial_v f]$ (Lenard–Bernstein/Dougherty, Buet notation $\beta = 1/2T$, $D = 1/2\beta = T$, drift $C = 2\beta D(v-\bar v) = (v-\bar v)$). - -| Location | Code | Check | -|---|---|---| -| driftdiffusion.py:101ff (`discrete_temperature`) | $T = \sum f(v-\bar v)^2 dv \,/\, \sum f\,dv$ | full 2nd central moment, no factor 2, no $m$ — variance convention, matches init ✓ | -| driftdiffusion.py:170,180 | $\beta_{init} = 1/2T$; $f_{mx} = e^{-\beta(v-\bar v)^2}$ | stationary state $e^{-(v-\bar v)^2/2T}$ ✓ | -| driftdiffusion.py:333–334 | $C = 2\beta D\,(v_{edge} - \bar v)$ | $2\beta D = 1$ exactly ✓ | -| driftdiffusion.py (flux/Chang–Cooper, implicit solve) | conservative central & positivity-preserving fluxes; $(I - dt\,\nu L)f^{n+1} = f^n$ | numerics only; $\nu\,dt$ dimensionless ✓ | -| driftdiffusion.py:20–25 | Buet "extra factor of 2" note | correctly absorbed into $\beta = 1/2T$; **no stray factor survives** (verified) | -| **fokker_planck.py:81** | `compute_vbar` $= \sum f\,v\,dv$ | returns $n\bar u$, **not** $\bar u$ — missing $1/n$ (F5) | -| **fokker_planck.py:245** | Krook target $f_{mx} \propto e^{-v^2/2}$ | hard-coded variance 1 = ($T{=}1$, $m{=}1$, convention (b)) — convention-locked (F4) | -| fokker_planck.py:262–266 | $f \to f e^{-\nu dt} + n(x) f_{mx}(1 - e^{-\nu dt})$ | BGK; conserves $n$, not momentum/energy (by design) | - -Key property: the **Dougherty operator is self-adapting** — it measures $T = \langle(v-\bar v)^2\rangle$ from $f$ and relaxes toward exactly that width, so it preserves whatever convention the initializer used and needs no change under a convention fix. The **Krook operator does not** — its width is a compile-time constant (F4). - -### 3.5 Diagnostics & storage (`storage.py`) - -| Location | Code | Meaning | Note | -|---|---|---|---| -| :134–135 | $\int(\cdot)\,dv$ = `sum·dv` | midpoint rule | consistent with init normalization ✓ | -| :139 | `n` $= \int f\,dv$ | density | ✓ | -| :140 | `v` $= \int v f\,dv$ | **raw first moment $= n\bar u$**, labeled "v" | flux, not velocity (F6) | -| :141–142 | `p` $= \int (v - n\bar u)^2 f\,dv$ | "pressure" | centered on $n\bar u$, not $\bar u$ (F6); for $n{=}1$ Maxwellian at $T_0$: `p` $= T_0$ under convention (b) ✓, would read $2T_0$ under (a) | -| :143 | `q` $= \int (v-n\bar u)^3 f\,dv$ | heat flux | same centering issue | -| :144 | `-flogf` $= \int f\log|f|\,dv$ | "entropy" | **missing minus sign** (F7) | -| :307 | `mean_-flogf` $= \langle\int -|f|\log|f|\,dv\rangle$ | entropy | has the minus; opposite sign to :144 (F7) | -| :303–306, 308 | `mean_P` $=\langle\int v^2f\rangle$, `mean_j` $=\langle\int vf\rangle$, `mean_n`, `mean_q` $=\langle\int v^3 f\rangle$, `mean_f2` | raw moments | `mean_j` and field "v" are the same integrand under different names (F6) | -| :311–312 | `mean_de2` $=\langle de^2\rangle$, `mean_e2` $=\langle e^2\rangle$ | field energy proxies | no ½; kinetic `mean_P` also lacks ½, consistently — internal conservation OK, but no combined energy monitor exists (F12) | -| :154, :313 | `pond` $= -\tfrac12\partial_x a^2$ | ponderomotive | ½ present and correct ✓ | - ---- - -## 4. Findings - -Ordered by severity. **Status: CONFIRMED** = derivation and code verified; **SUSPECTED** = probable issue, evidence stated. - -### F1 — CONFIRMED (root cause): `normalization.py` $v_0$ is $\sqrt2$ larger than the engine's velocity unit - -`normalization.py:91` ($v_0 = \sqrt{2T_0/m_e}$) vs. the engine convention $\sqrt{T_0/m_e}$ established by `helpers.py:69`, the collision operators, the dispersion tests, and all configs (§2). Since the physical temperature of the simulated plasma is $T_{phys} = m_e\,\sigma_{code}^2\,v_0^2 = T_{0,code}\cdot(m_e v_0^2)$: - -- **Under (a) as written: $T_{phys} = 2\,T_{0,code}\,T_{0,ref}$.** A `normalizing_temperature: 2000eV`, `T0: 1.0` run is a 4000 eV plasma. -- $L_0 = \sqrt2\lambda_{De}$, so all `dim="x"`/`dim="k"` string conversions carry a spurious $\sqrt2$ (box sizes, gradient scale lengths in `datamodel.py:159–177`, laser $k_0$ at `simulation.py:81`). -- The species-config `T0` parameter is effectively in units of $2T_{0,ref}$ — half the naïve expectation. - -The dimensionless dynamics of numeric-input runs are unaffected (the solver core never sees $v_0$); the damage is to the *physical interpretation* of every run and to dimensional-string inputs. - -### F2 — CONFIRMED (propagation): logged units and regression fixtures lock in the $\sqrt2$ values - -`write_units()` (`modules.py:119–141`) logs `v0`, `x0`, `box_length`, `c_light`, `nuee` computed under convention (a). For the 2000 eV resonance case the fixtures record `c_light: 11.302`, `v0: 2.652e7 m/s`, `x0: 12.14 nm` (`tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml:125–132`, same in the fokker-planck and multispecies fixtures); the convention-(b) truth values are $c/v_{th} = 15.98$, $v_{th} = 1.876\times10^7$ m/s, $\lambda_{De} = 8.585$ nm. Additionally `nuee` is evaluated at $T_{0,ref}$ while the plasma actually simulated is at $2T_{0,ref}$, so the logged collisionality does not describe the simulated plasma. Any fix of F1 must regenerate these fixtures. - -### F3 — CONFIRMED: EM branch — $\hat c$ is $\sqrt2$ too small; partial cancellation hides it for wavelength drivers only - -$\beta = 1/\hat c$ with $\hat c = c/v_0$ (`modules.py:58,121`) feeds the wave solver and the CFL limit. Under (a), $\hat c = c/(\sqrt2 v_{th})$ — $\sqrt2$ smaller than the engine's thermal unit warrants. For **wavelength-specified** drivers the $\sqrt2$'s cancel in the product $\hat c\hat k = \frac{c}{\sqrt2 v_{th}}\cdot\sqrt2 k\lambda_{De}$ — which is why `test_em_dispersion.py` and `srs.yaml`'s EM branch behave physically. For **numeric** AKW drivers (`srs-debug-small.yaml`: `k0: 1.0, w0: 2.79` used as-is at `simulation.py:59`), $\hat k$ is *not* rescaled while $\hat c$ still carries the $\sqrt2$, so the EM dispersion $\hat\omega^2 = \hat n_e + \hat c^2\hat k^2$ is $\sqrt2$-inconsistent with the ES branch's $\hat k = k\lambda_{De}$ interpretation. Mixed-input SRS runs are therefore internally inconsistent between branches. - -### F4 — CONFIRMED: Krook target Maxwellian is hard-coded to $e^{-v^2/2}$ (convention-locked, $T{=}1$, electron grid only) - -`fokker_planck.py:245`. Three separate problems: (i) it bakes in convention (b) with variance exactly 1, so it must be changed **in lockstep** with any convention fix or it will spuriously heat/cool by 2× in temperature; (ii) even today, it ignores the species `T0` and `mass` — any species initialized at $T_0 \ne 1$ (e.g. `twostream.yaml`, $T_0 = 0.2$) is dragged toward $T = 1$; (iii) it is built on the electron grid and applied only to `"electron"` (`fokker_planck.py:171`). It is a fixed-target BGK operator: conserves density only, not momentum or energy. Currently `is_on: False` in shipped configs, so latent. - -### F5 — CONFIRMED: `Dougherty.compute_vbar` is missing the $1/n$ normalization - -`fokker_planck.py:81` returns $\int v f\,dv = n\bar u$ while its docstring claims "Mean velocity ⟨v⟩". `discrete_temperature` (`driftdiffusion.py:101ff`) *does* divide by $\int f\,dv$ — the two moments are asymmetric. Consequence: the drag centers on $n\bar u$, the operator relaxes toward $e^{-\beta(v - n\bar u)^2}$, and **momentum is not conserved where $n(x) \ne 1$**; $T_{target}$ is also biased by the wrong centering. Negligible for the shipped EPW/SRS configs ($n \approx 1 \pm 10^{-4}$), real for bump-on-tail / large density perturbations with collisions on. Fix: divide by $\int f\,dv$. - -### F6 — CONFIRMED: storage central moments centered on $n\bar u$, not $\bar u$; misleading names - -`storage.py:140–143`: the field moment `"v"` is $\int v f\,dv = n\bar u$ (a flux), and `p`, `q` are centered on it. Only exact when $n = 1$; biased for `nlepw-ic.yaml` (10% density perturbation) and `bump-on-tail.yaml`. Same integrand is named `"v"` in field moments but `"mean_j"` in scalars (`storage.py:304`) — one of the labels is wrong; `j` is the honest name (or divide by $n$ and keep "v"). Note this is the same class of bug as F5, appearing independently in two places. - -### F7 — CONFIRMED: the two `-flogf` entropy diagnostics have opposite signs - -Field moment `storage.py:144` computes $+\int f\log|f|\,dv$ (no minus, uses $f$); scalar `storage.py:307` computes $-\int|f|\log|f|\,dv$. For $f > 0$ these are exact negatives. The label `-flogf` matches the scalar; the field version is sign-flipped. - -### F8 — SUSPECTED (intent unclear): super-Gaussian $\alpha$ fixes $\langle v^4\rangle/\langle v^2\rangle$, not the variance, for $m \ne 2$ - -`helpers.py:72`: $\alpha = \sqrt{3\Gamma(3/m)/\Gamma(5/m)}$ normalizes the kurtosis ratio $\langle v^4\rangle/\langle v^2\rangle = 3v_{th}^2$ for all $m$, but the variance equals $v_{th}^2$ only at $m = 2$. Numerically the realized variance is $\{1.240, 1.371, 1.449\}\times T_0/m$ for $m = \{3,4,5\}$ — so the measured `p`/`n` moment will *not* equal the input `T0` for super-Gaussian species. May be an intentional flat-top-EDF temperature definition; if so it should be documented, because the `T0` config docstring ("Temperature") and the second-moment diagnostic disagree with it. All shipped configs use $m = 2$, where it is exact. - -### F9 — CONFIRMED (gap/trap): collision-frequency normalization chain is entirely manual - -- `approximate_ee_collision_frequency` returns **Hz** and is only *logged* (`modules.py:119,132`); nothing converts it to code units or feeds it to the FP operator. -- The FP/Krook rate magnitudes (`baseline`, `bump_height`) are taken as **raw floats** (`functions.py:97–98`), not passed through `normalize()` — users must supply $\nu$ already in code units ($1/\omega_{p0}$), which is nowhere documented (`config.md` lists `baseline` with no units). -- The correct conversion is $\hat\nu = \nu_{ee}[\mathrm{Hz}]\cdot\tau = \nu_{ee}/\omega_{p0}$ with **no** $2\pi$ (the NRL rate is a true s⁻¹ rate); `_tf1d/modules.py:78` has this pattern (`nuee_norm = nuee/wp0`), `_vlasov1d` has no analog. A user thinking in cyclic frequency is one step from a silent $2\pi$ error. -- The LB/Dougherty $\nu$ and the NRL $\nu_{ee}$ agree only up to an O(1) factor — identifying them silently is itself an approximation worth a docs note. -- $\nu$ is a prescribed space-time envelope; it does not track the local $n(x)/T(x)^{3/2}$ (limitation, not a bug). - -Suggested: log a `nuee_norm` alongside `nuee`, and document `baseline`'s units. - -### F10 — CONFIRMED: species `T0` and drift `v0` bypass the unit machinery - -`simulation.py:213–214` take `float(cfg.T0)`, `float(cfg.v0)`; `datamodel.py:20–21` type them as bare floats. The `normalize(dim="temp")` and `dim="v"` branches are dead code for `_vlasov1d`. Consequences: `T0: "500 eV"` is impossible (only the global `normalizing_temperature` is dimensional), and the drift `v0` is expressed in units of the reference $v_0$ while the thermal width is in $\sigma$ units — under convention (a) these differ by $\sqrt2$ *within the same distribution*, a genuine user trap (drift "1.0" is not "one thermal width"). - -### F11 — CONFIRMED (latent): `dx = xmax/nx` ignores `xmin` - -`grid.py:55` should be `(xmax − xmin)/nx`; the x-axis itself (`grid.py:74`) spans $[x_{min}, x_{max}]$ with the wrong spacing when $x_{min} \ne 0$, and the `kx` grid (`grid.py:77`) and the semi-Lagrangian interpolation period (`pushers/vlasov.py:31`, `period = xmax`) inherit the error. `modules.py:112` uses the correct $(x_{max}-x_{min})$, highlighting the discrepancy. All shipped configs use `xmin: 0.0`, so latent but real. - -### F12 — Minor items (confirmed, low impact) - -1. **`AmpereSolver` class docstring** (`field.py:230`) claims $j = \sum_s (q_s/m_s)\int vf_s\,dv$; the code (`field.py:263–264`) correctly omits $1/m_s$. Docstring bug only. -2. **Energy diagnostics lack ½ and a total**: `mean_P`, `mean_e2`, `mean_de2` are $2\times$ the respective energies (consistently, so `mean_P + mean_e2` is still conserved), and no diagnostic sums kinetic + field energy. A dedicated conservation monitor would have caught F1 earlier. -3. **`dw0 = 0.0 # ???`** placeholder at `simulation.py:87` for the intensity/wavelength driver. -4. **Dead config knob**: `GridConfig.c_light` (`datamodel.py:73`; set in `wavepacket.yaml`) is never read — `c_light`/`beta` are always recomputed from the normalization. -5. **`grid.t` axis spacing**: `nt = int(tmax/dt + 1)`, `tmax = dt·nt` overshoots the request by up to ~dt and `linspace(0, tmax, nt)` has spacing $\ne$ dt. Cosmetic (saves use their own axes). -6. **EM plasma term is electron-only** (`vector_field.py:292–301`, `field.py:152`): ions omitted from the transverse current ($\sim m_e/m_i$ — fine, but an asymmetry vs. Poisson/Ampère which include all species). -7. **Point-source amplitude** uses the vacuum $k = \omega/c$ (`field.py:67`); realized plasma amplitude larger by $k_{vac}/k_{plasma}$ — already documented in the class docstring. -8. **Half-cell axis convention**: `vax`/`x` hold cell *centers* while the extents name the *edges* ($\pm v_{max}$) — keep in mind when labeling axes. - -### Verified correct (checked because they looked suspicious, and passed) - -- The entire collisionless solver core: all coefficients unity, signs consistent (§3.3). -- The ponderomotive chain: the $-\tfrac12\partial_x a^2$, the $(q^2/m^2)$ factor, and the $c/v_0$ scaling of $a_0$ combine *exactly* — no cycle-average ½ missing, no hidden $\sqrt2$. -- The Buet "factor of 2" in the Dougherty operator: correctly absorbed by $\beta = 1/2T$; equilibrium width equals the measured variance exactly. -- The $E_x$ driver amplitude $E = \omega a_0$: consistent vector-potential convention, matches the docs. -- Midpoint-rule ($\mathrm{sum}\cdot dv$) integration: used uniformly in init, moments, and field solves. - ---- - -## 5. Recommended resolution and lockstep checklist - -**Adopt convention (b) globally**: in `electron_debye_normalization`, set $v_0 = \sqrt{T_0/m_e}$ (so $L_0 = \lambda_{De}$), and redefine `vth_norm()` accordingly. This leaves the initializer, all collision operators, all tests, and all numeric configs untouched and correct, and makes the logged physical units true. - -Things that must change together / be re-verified: - -1. `normalization.py:91–92` ($v_0$, $x_0$) — **and nothing else in that file. Do NOT redefine `vth_norm()` (:48–50)**: the codebase-wide audit (`ADEPT_CONVENTIONS_AUDIT.md`, finding C1) found that `vth_norm()` has exactly four callers, all in `vfp1d/`, which is self-consistently built on the $\sqrt{2T/m}$ convention (its initializer carries a compensating factor). Changing `vth_norm()` does nothing for `_vlasov1d` (never called here) and would silently halve VFP-1D's initialized temperature. -2. Regenerate the three `*_derived_config.yml` regression fixtures (they lock `c_light`, `v0`, `x0` at the $\sqrt2$ values). -3. Re-check every config that uses **dimensional string** inputs (`srs.yaml`: `xmax: 100um`, gradient scale length `200um`, laser wavelength/intensity) — their physical meaning shifts by $\sqrt2$ (they become *correct*; the previously-inferred physical parameters of past runs were what was wrong). -4. `c_light`/`beta` changes for all EM runs: re-validate `test_em_dispersion.py`, the point-source amplitude test, and numeric-driver SRS configs (`srs-debug-small.yaml` `w0`/`k0` values were presumably tuned under the old $\hat c$ — see F3). -5. Krook target (`fokker_planck.py:245`) — no change needed under (b), but fix its $T_0$/species handling anyway (F4) so it stops being convention-locked. -6. Docs: state the normalization explicitly in `docs/source/solvers/vlasov1d/` (velocity unit $=\sqrt{T_0/m_e}$, $L_0 = \lambda_{De}$, $\hat k = k\lambda_{De}$, Maxwellian $e^{-v^2/2}$ at $T{=}1$), and document `baseline` units (F9) and the drift-`v0` units (F10). -7. ~~Audit sibling solvers for the same $v_0$ dependency~~ **Done — see `ADEPT_CONVENTIONS_AUDIT.md`** (codebase-wide audit with per-solver immunity verdicts). Summary: `electron_debye_normalization` is consumed only by `_vlasov1d`, `_vlasov2d`, `_pic1d` — all σ-convention, all fix-safe, and only `_vlasov1d` has fixtures to regenerate. `vfp1d` genuinely runs in $\sqrt{2T/m}$ units but via `laser_normalization` + `vth_norm()`, so it is untouched by the scoped fix (see amended item 1). `_tf1d` has a private inline duplicate of the √2 (`_tf1d/modules.py:62`) that must be fixed separately. All other modules are immune. - -Independent bug fixes (any order, no convention coupling): F5 (`compute_vbar` $1/n$), F6 (moment centering + naming), F7 (`-flogf` sign), F11 (`dx` with `xmin`), F12.1 (docstring), F12.2 (add an energy-conservation scalar). - -### Suggested verification runs after any fix - -- Landau damping and ion-acoustic tests must still pass unchanged (they pin convention (b)). -- A Dougherty-collisions run with a strong density perturbation: check $\partial_t\int v f\,dv \approx 0$ (exercises F5). -- An SRS wavelength-driver run: confirm backscatter resonance moves to the physically correct location once `c_light` and the µm conversions change together. -- Compare the second-moment temperature diagnostic against `T0` for an $m = 4$ super-Gaussian to decide F8's intended convention. - ---- - -*Audit method: four parallel independent reviews (initialization/grids/derived quantities; solver core; collisions; diagnostics/docs/tests/configs), followed by cross-checking of all load-bearing claims against the source. No source files were modified.* From e244265fa0d0773b64fb43ec4f7a6215a245d428 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 20 Jul 2026 12:02:53 -0700 Subject: [PATCH 49/49] osiris: keep the stream drainer ahead of OSIRIS dump production Post-mortem of job 56005549 (srs-Ln100-Te4-5x-ions): at the every-3-step field cadence the per-run drainer thread sustained only ~2/3 of the ~16 MB/s dump rate, so the /dev/shm staging backlog grew ~2.2 GiB/min, filled the node's 256 GiB of RAM in ~1.8 h, and every later OSIRIS HDF5 write failed on ENOSPC (truncated 2 KB dumps, 2 OOM kills, 26 h TIMEOUT at 16% of tmax). Details in osiris-lpi NOTES.md + dev_docs/ stream-drainer-recommendations.md. Throughput (the drainer needed ~1.5x; this buys ~20x): - StreamWriter batches appends and lands them chunk-aligned: one resize_dimension + one slice write per batch instead of a gzip read-modify-write of the whole ~1 MiB chunk per row (~40x write amplification). Measured 23x on the production field shape (1765 vs 76 appends/s; the job needed ~55/s per sim). - Default compression gzip 4 -> 1 (the drainer is compression-bound; offline postproc can recompress). - Discovery is cached: the full MS/ rglob walk, whose cost grows with the backlog, runs every rediscover_every polls instead of every poll. - No thread pool on purpose: h5py serializes all HDF5 calls behind a process-global lock, so batching is where the throughput is. Safety: - Backlog spill valve (staging mode): past spill_backlog_files/_bytes, or under floor_free_bytes on the staging fs, a diagnostic flips to mirror-only draining (plain copy to persist MS/, ~10x cheaper), so the ramdisk keeps draining no matter what; the NetCDF is caught up from the mirror at finalize (mirror consumed + pruned in discard_grid_h5 mode). OSIRIS must never see ENOSPC on a dump. - Corrupt dumps are quarantined (*.h5.bad) and the stream continues; previously one bad dump dropped the writer and the watcher re-hit the same file every poll, forever. Bookkeeping is now iteration-based (recovered from the iter coordinate on resume), not positional, so skipped dumps cannot shift slots. Observability: - One stats line per stats_every_s (streamed/spilled/backlog/ quarantined + staging free space) so a smoke can assert the steady-state backlog is flat while OSIRIS runs; checking the ramdisk after the job proves nothing. - Repeated identical errors log once per error_log_every occurrences. API and on-disk schema unchanged (runner.py untouched); all 17 existing stream tests pass unmodified, plus new coverage for quarantine and the spill valve in both staging modes. Co-Authored-By: Claude Fable 5 --- adept/osiris/stream.py | 461 ++++++++++++++++++++++++------- tests/test_osiris/test_stream.py | 74 +++++ 2 files changed, 437 insertions(+), 98 deletions(-) diff --git a/adept/osiris/stream.py b/adept/osiris/stream.py index fbe4cef5..1ad00199 100644 --- a/adept/osiris/stream.py +++ b/adept/osiris/stream.py @@ -25,13 +25,45 @@ Design notes ------------ - **Unlimited ``t`` dimension, append in iteration order.** OSIRIS writes dumps - sequentially and we process them sorted by iteration, so a plain append (slot - ``i`` = the ``i``-th dump) is correct and needs no stride arithmetic. The - dimension is grown one slot at a time with ``resize_dimension`` so it ends up - sized to *exactly* the dumps produced — no pre-sized guess from the deck + sequentially and we process them sorted by iteration, so appending in that + order is correct. The dimension is grown with ``resize_dimension`` so it ends + up sized to *exactly* the dumps produced — no pre-sized guess from the deck (whose off-by-one is genuinely ambiguous) and no trailing fill slots to trim - on early termination. This deviates from the fixed-dim sketch in the perf - doc; it is simpler and equally correct given in-order processing. + on early termination. +- **Batched, chunk-aligned writes.** :meth:`StreamWriter.append` buffers rows + in memory and lands them in chunk-sized batches (one HDF5 slice write and one + ``resize_dimension`` per batch). Appending row-by-row into a gzip chunk of + ``chunk_t`` rows costs a decompress+recompress of the whole ~1 MiB chunk *per + row* — ~40x write amplification at the every-3-step field cadence, which is + what let the drainer fall behind OSIRIS and fill the staging ramdisk in job + 56005549 (see ``osiris-lpi/dev_docs/stream-drainer-recommendations.md``). + The buffer is bounded by one chunk (~1 MiB); :meth:`StreamWriter.flush` + (called once per watcher scan) spills any partial batch, so the on-disk file + trails the consumed dumps by at most one poll interval. Threads are *not* + used to parallelize conversion: h5py serializes every HDF5 call behind a + process-global lock, so batching + cheap compression is where the throughput + is. +- **Iteration-based bookkeeping.** The writer tracks the last appended OSIRIS + iteration (persisted in the ``iter`` coordinate, so restarts recover it) and + the drain loops skip any dump at or below it. Unlike positional slot + arithmetic this survives corrupt dumps being removed from the listing. +- **Corrupt dumps are quarantined, not fatal.** A dump that cannot be read + (e.g. truncated by ENOSPC) is renamed to ``*.h5.bad`` — invisible to the + ``*.h5`` listings — and the stream continues with the next dump. Previously a + single bad dump dropped the whole writer and the watcher re-hit the same file + every poll, forever. +- **Backlog spill valve (staging mode).** The ramdisk footprint is only + "bounded by the poll interval" while conversion keeps up with production. + When a diagnostic's pending backlog exceeds ``spill_backlog_files`` / + ``spill_backlog_bytes`` (or the staging filesystem drops below + ``floor_free_bytes``), the converter stops converting it and *mirrors* its + dumps straight to the persist ``MS/`` tree instead (a plain copy is ~10x + cheaper than convert+compress), so the ramdisk keeps draining no matter what. + A spilled diagnostic stays mirror-only for the rest of the run (resuming the + stream mid-run would leave an iteration gap in the NetCDF) and is caught up + from the mirrored HDF5 during :meth:`StreamConverter.finalize`. OSIRIS must + never see ENOSPC on a dump: in 56005549 that permanently corrupted its + diagnostic output for the rest of the job. - **Write-completeness race.** A dump is only safe to read once OSIRIS has moved on to the next one, so the watcher processes every dump *except the current highest-numbered one*; the final sweep (run is dead) picks up that last dump. @@ -41,8 +73,14 @@ those bounds actually move, so the constant bounds carried for a fixed spatial axis are harmless. - **Failure isolation.** Nothing here may abort the OSIRIS run: the watcher - swallows and logs every exception, and the batch path in + swallows and logs every exception (rate-limited — repeats of the same error + log once per ``error_log_every`` occurrences), and the batch path in :func:`io.save_run_datasets` remains the safety net for anything not produced. +- **Observability.** Every ``stats_every_s`` the watcher logs one line with the + interval's streamed/spilled/backlog/quarantined counts, so a smoke can assert + the steady-state backlog is flat *while* OSIRIS runs. Checking the ramdisk + after the job proves nothing — the drainer always catches up once production + stops. """ from __future__ import annotations @@ -50,6 +88,7 @@ import os import shutil import threading +import time from pathlib import Path import h5netcdf @@ -59,9 +98,9 @@ from adept.osiris import io as _io # Target uncompressed bytes per (t-chunk x spatial) chunk. A modest t-chunk -# keeps compression/read performance close to the batch path while bounding the -# read-modify-write cost of slot writes (which is hidden behind OSIRIS compute -# anyway). ~1 MiB is a reasonable HDF5 chunk size. +# keeps compression/read performance close to the batch path; batched appends +# (see StreamWriter) amortize the write cost of a chunk over its rows. ~1 MiB +# is a reasonable HDF5 chunk size. _TARGET_CHUNK_BYTES = 1 << 20 @@ -91,48 +130,63 @@ class StreamWriter: """Append-only writer for one diagnostic's ``(t, ...)`` NetCDF time series. Created from the diagnostic's first dump (the schema template); each - :meth:`append` writes the next time slot, growing the unlimited ``t`` - dimension by one. Reopening an existing file (``mode="a"``) resumes after the - slots already on disk, which is what makes a restart idempotent. The file is + :meth:`append` buffers the next time slot, and slots are landed on disk in + chunk-sized batches (or on :meth:`flush` / :meth:`close`), growing the + unlimited ``t`` dimension once per batch. Reopening an existing file + (``mode="a"``) resumes after the slots already on disk, which is what makes + a restart idempotent; ``last_iter`` (recovered from the ``iter`` coordinate + on resume) lets callers skip dumps already in the file. ``n_written`` + counts consumed dumps — on-disk slots plus the pending batch. The file is held open for the lifetime of the writer; call :meth:`close` to finalize. + + The default ``complevel`` is 1: at production dump cadences the drainer is + compression-bound, and gzip-1 is 2-4x cheaper than the old gzip-4 for a + modest size cost (offline postproc can always recompress). """ - def __init__(self, dest, template: xr.DataArray, *, source_dir=None, chunk_t: int | None = None, complevel: int = 4): + def __init__(self, dest, template: xr.DataArray, *, source_dir=None, chunk_t: int | None = None, complevel: int = 1): self.dest = Path(dest) self.name = str(template.name) self.spatial_dims = [str(d) for d in template.dims] self.spatial_shape = tuple(int(s) for s in template.shape) self._f: h5netcdf.File | None = None + self._ct = int(chunk_t or _default_chunk_t(self.spatial_shape)) + # The pending batch: rows consumed by append() but not yet on disk. + self._rows: list[np.ndarray] = [] + self._ts: list[float] = [] + self._its: list[int] = [] + self._bounds: dict[str, list[tuple[float, float]]] = {d: [] for d in self.spatial_dims} if self.dest.exists(): # Resume: trust the existing schema, continue after written slots. self._f = h5netcdf.File(self.dest, "a") - self._n = int(self._f.dimensions["t"].size) + self._n_disk = int(self._f.dimensions["t"].size) + self.last_iter = int(self._f.variables["iter"][self._n_disk - 1]) if self._n_disk else -1 else: self.dest.parent.mkdir(parents=True, exist_ok=True) self._f = h5netcdf.File(self.dest, "w") - self._create_schema(template, source_dir if source_dir is not None else self.dest.parent, chunk_t, complevel) - self._n = 0 + self._create_schema(template, source_dir if source_dir is not None else self.dest.parent, self._ct, complevel) + self._n_disk = 0 + self.last_iter = -1 @property def n_written(self) -> int: - return self._n + return self._n_disk + len(self._rows) def _create_schema(self, template, source_dir, chunk_t, complevel) -> None: f = self._f axis_units = template.attrs.get("axis_units", {}) or {} axis_long = template.attrs.get("axis_long_names", {}) or {} - f.dimensions["t"] = None # unlimited; grown one slot per append + f.dimensions["t"] = None # unlimited; grown one batch per landing for d, sz in zip(self.spatial_dims, self.spatial_shape): f.dimensions[d] = sz - ct = chunk_t or _default_chunk_t(self.spatial_shape) var = f.create_variable( self.name, ("t", *self.spatial_dims), dtype=np.dtype(_io._DIAG_DTYPE), - chunks=(ct, *self.spatial_shape), + chunks=(chunk_t, *self.spatial_shape), compression="gzip", compression_opts=complevel, shuffle=True, @@ -158,28 +212,51 @@ def _create_schema(self, template, source_dir, chunk_t, complevel) -> None: f.create_variable(f"{d}_max", ("t",), dtype="f8") def append(self, da: xr.DataArray) -> None: - """Write ``da`` (one dump) into the next time slot.""" + """Buffer ``da`` (one dump) as the next time slot; lands in batches.""" if da.shape != self.spatial_shape: raise ValueError(f"shape mismatch for {self.name}: {da.shape} != {self.spatial_shape}") - f = self._f - i = self._n - f.resize_dimension("t", i + 1) - f.variables[self.name][i, ...] = np.asarray(da.values, dtype=_io._DIAG_DTYPE) - f.variables["t"][i] = float(da.attrs.get("time", np.nan)) - f.variables["iter"][i] = int(da.attrs.get("iter", -1)) + self._rows.append(np.asarray(da.values, dtype=_io._DIAG_DTYPE)) + self._ts.append(float(da.attrs.get("time", np.nan))) + it = int(da.attrs.get("iter", -1)) + self._its.append(it) + if it > self.last_iter: + self.last_iter = it for d in self.spatial_dims: cv = np.asarray(da.coords[d].values) - lo, hi = (float(cv[0]), float(cv[-1])) if cv.size else (np.nan, np.nan) - f.variables[f"{d}_min"][i] = lo - f.variables[f"{d}_max"][i] = hi - self._n = i + 1 + self._bounds[d].append((float(cv[0]), float(cv[-1])) if cv.size else (np.nan, np.nan)) + if len(self._rows) >= self._ct: + self._write_batch() + + def _write_batch(self) -> None: + """Land the pending rows: one resize + one slice write per variable.""" + k = len(self._rows) + if not k or self._f is None: + return + f = self._f + i = self._n_disk + f.resize_dimension("t", i + k) + f.variables[self.name][i : i + k, ...] = np.stack(self._rows) + f.variables["t"][i : i + k] = np.asarray(self._ts, dtype="f8") + f.variables["iter"][i : i + k] = np.asarray(self._its, dtype="i8") + for d in self.spatial_dims: + b = np.asarray(self._bounds[d], dtype="f8") + f.variables[f"{d}_min"][i : i + k] = b[:, 0] + f.variables[f"{d}_max"][i : i + k] = b[:, 1] + self._n_disk = i + k + self._rows.clear() + self._ts.clear() + self._its.clear() + for d in self.spatial_dims: + self._bounds[d].clear() def flush(self) -> None: if self._f is not None: + self._write_batch() self._f.flush() def close(self) -> None: if self._f is not None: + self._write_batch() self._f.close() self._f = None @@ -188,8 +265,9 @@ def convert_diagnostic_streaming(diag_dir, dest, *, source_dir=None) -> Path: """Stage A: convert one grid diagnostic to NetCDF a dump at a time. A memory-bounded drop-in for ``series_to_dataset(load_series(diag_dir))`` - followed by ``to_netcdf`` — peak RAM is one dump, not the whole stacked - series. Rebuilds ``dest`` from scratch (any partial file is removed first). + followed by ``to_netcdf`` — peak RAM is one dump plus one write batch + (~1 MiB), not the whole stacked series. Rebuilds ``dest`` from scratch (any + partial file is removed first). """ diag_dir = Path(diag_dir) dest = Path(dest) @@ -217,8 +295,8 @@ class StreamConverter: subprocess, it polls ``run_dir/MS`` and drains each grid diagnostic's completed dumps into ``out_dir/.nc``. RAW (particle) diagnostics do not fit the fixed-slot model (variable particle count per dump) and are left - to the batch path. Every operation is best-effort: any error is logged and - the run is never affected. + to the batch path. Every operation is best-effort: any error is logged + (rate-limited) and the run is never affected. **Staging mode (``persist_dir`` set).** When OSIRIS writes its ``MS/`` dumps to a fast ephemeral scratch (e.g. a ``/dev/shm`` ramdisk passed as @@ -233,10 +311,35 @@ class StreamConverter: post-processing reads it unchanged. With ``persist_dir=None`` nothing is mirrored or deleted — the converter only streams grid diagnostics in place, exactly as before. + + **Backlog spill valve (staging mode).** The bound above only holds while + conversion keeps up with OSIRIS. When a diagnostic's pending backlog + exceeds ``spill_backlog_files`` or ``spill_backlog_bytes`` — or the staging + filesystem's free space drops below ``floor_free_bytes`` — that diagnostic + switches to mirror-only draining (plain copy to ``persist_dir/MS/``, ~10x + cheaper than convert+compress), so the ramdisk keeps draining no matter + what. A spilled diagnostic stays mirror-only for the rest of the run + (resuming the stream mid-run would leave an iteration gap in the NetCDF) + and its NetCDF is caught up from the mirrored HDF5 in :meth:`finalize`; in + ``discard_grid_h5`` mode the mirrored copies are deleted again once + appended. """ def __init__( - self, run_dir, out_dir, *, poll_s: float = 10.0, logger=print, persist_dir=None, discard_grid_h5: bool = False + self, + run_dir, + out_dir, + *, + poll_s: float = 10.0, + logger=print, + persist_dir=None, + discard_grid_h5: bool = False, + spill_backlog_files: int = 20_000, + spill_backlog_bytes: int = 4 << 30, + floor_free_bytes: int = 8 << 30, + stats_every_s: float = 60.0, + rediscover_every: int = 12, + error_log_every: int = 200, ): self.ms = Path(run_dir) / "MS" self.out_dir = Path(out_dir) @@ -246,17 +349,30 @@ def __init__( # being appended to the NetCDF *without* being mirrored to persist_dir/MS # (the .nc is the durable artifact). Saves one inode+copy per dump on the # persist filesystem. RAW dumps are always mirrored (the batch path needs - # the HDF5). Trade-off: a grid diagnostic whose stream fails mid-run can - # only be rebuilt from the dumps not yet reaped (they are left in place - # and folded to persist by the runner's final sync). + # the HDF5), and so are spilled grid dumps (their bytes exist nowhere + # else until the finalize catch-up appends them). self.discard_grid_h5 = bool(discard_grid_h5) self.poll_s = float(poll_s) + self.spill_backlog_files = int(spill_backlog_files or 0) + self.spill_backlog_bytes = int(spill_backlog_bytes or 0) + self.floor_free_bytes = int(floor_free_bytes or 0) + self.stats_every_s = float(stats_every_s) + self.rediscover_every = max(1, int(rediscover_every)) + self.error_log_every = max(1, int(error_log_every)) self._log = logger self._writers: dict[str, StreamWriter] = {} self._completed: set[str] = set() # grid relpaths fully converted to .nc self._raw_completed: set[str] = set() # raw relpaths fully mirrored to persist self._is_raw: dict[str, bool] = {} # rel -> raw? (cached; avoids re-peeking h5) self._multibase_skipped: set[str] = set() # multi-report dirs left to the batch pass + self._spilled: set[str] = set() # grid rels in sticky mirror-only mode + self._bad: set[Path] = set() # unreadable dumps that could not be renamed away + self._err_counts: dict[str, int] = {} + self._known: dict[str, Path] = {} # cached diagnostic discovery + self._scan_i = 0 + self._stats = {"appended": 0, "appended_bytes": 0, "spilled": 0, "spilled_bytes": 0, "quarantined": 0} + self._stats_snapshot = dict(self._stats) + self._stats_t = time.monotonic() self._stop = threading.Event() self._thread: threading.Thread | None = None @@ -271,7 +387,7 @@ def _run(self) -> None: try: self._scan_once(final=False) except Exception as e: # a watcher must never crash the run - self._log(f"[stream] scan error (continuing): {e}") + self._log_every("scan", f"[stream] scan error (continuing): {e}") self._stop.wait(self.poll_s) def finalize(self) -> set[str]: @@ -279,11 +395,12 @@ def finalize(self) -> set[str]: After OSIRIS has exited every dump is safe to read (no writer can still be racing), so the final sweep processes the last dump each diagnostic - had been holding back (and, in staging mode, drains it off the scratch) - before closing the files. Returns the set of *grid* diagnostic relpaths - fully converted by the stream (so the caller can skip rebuilding them in - the batch pass); RAW relpaths are mirrored but not returned, since the - batch pass still builds their NetCDFs from the mirrored HDF5. + had been holding back, catches spilled diagnostics up from their persist + mirror, and (in staging mode) drains everything off the scratch before + closing the files. Returns the set of *grid* diagnostic relpaths fully + converted by the stream (so the caller can skip rebuilding them in the + batch pass); RAW relpaths are mirrored but not returned, since the batch + pass still builds their NetCDFs from the mirrored HDF5. """ self._stop.set() if self._thread is not None: @@ -300,6 +417,41 @@ def finalize(self) -> set[str]: self._completed.add(rel) return set(self._completed) + # --- logging / stats -------------------------------------------------- + + def _log_every(self, key: str, msg: str) -> None: + """Log ``msg`` on the first occurrence of ``key`` and every Nth after. + + The ENOSPC failure mode repeats the *same* error for every dump on + every poll; unthrottled, that was megabytes of log spam per hour. + """ + n = self._err_counts.get(key, 0) + 1 + self._err_counts[key] = n + if n == 1 or n % self.error_log_every == 0: + self._log(msg + (f" [seen x{n}]" if n > 1 else "")) + + def _maybe_stats(self, backlog_files: int, backlog_bytes: int, *, final: bool) -> None: + """One drain-vs-production line per ``stats_every_s`` (and at finalize).""" + now = time.monotonic() + if not final and (now - self._stats_t) < self.stats_every_s: + return + delta = {k: self._stats[k] - self._stats_snapshot[k] for k in self._stats} + if final or any(delta.values()) or backlog_files: + extra = "" + if self.persist_dir is not None: + try: + extra = f", stage free {shutil.disk_usage(self.ms).free >> 30} GiB" + except OSError: + pass + self._log( + f"[stream] stats({now - self._stats_t:.0f}s): streamed {delta['appended']} dumps " + f"({delta['appended_bytes'] >> 20} MiB), spilled {delta['spilled']} " + f"({delta['spilled_bytes'] >> 20} MiB), backlog {backlog_files} dumps " + f"({backlog_bytes >> 20} MiB), quarantined {delta['quarantined']}{extra}" + ) + self._stats_t = now + self._stats_snapshot = dict(self._stats) + # --- scanning --------------------------------------------------------- def _diags(self) -> dict[str, Path]: @@ -329,6 +481,16 @@ def _diags(self) -> dict[str, Path]: out[rel] = d return out + def _discover(self, *, final: bool) -> dict[str, Path]: + """Cached discovery. The full ``rglob`` walk touches every backlogged + file, so its cost grows exactly when the drainer is behind; new + diagnostic *directories* appear rarely, so the walk runs only every + ``rediscover_every`` polls and the known dirs are re-listed directly in + between.""" + if final or not self._known or (self._scan_i % self.rediscover_every == 1): + self._known = self._diags() + return self._known + def _diag_is_raw(self, rel: str, d: Path) -> bool: cached = self._is_raw.get(rel) if cached is None: @@ -337,7 +499,22 @@ def _diag_is_raw(self, rel: str, d: Path) -> bool: return cached def _scan_once(self, *, final: bool) -> None: - for rel, d in self._diags().items(): + self._scan_i += 1 + force_spill = False + if self.persist_dir is not None and not final and self.floor_free_bytes: + try: + if shutil.disk_usage(self.ms).free < self.floor_free_bytes: + force_spill = True + self._log_every( + "floor", + f"[stream] CRITICAL: staging filesystem under {self.floor_free_bytes >> 30} GiB free" + " -- spilling all grid dumps to persist (OSIRIS must never see ENOSPC)", + ) + except OSError: + pass + backlog_files = 0 + backlog_bytes = 0 + for rel, d in self._discover(final=final).items(): if rel in self._completed or rel in self._raw_completed: continue try: @@ -347,12 +524,12 @@ def _scan_once(self, *, final: bool) -> None: # the dumps on persist). if self.persist_dir is not None: self._drain_raw(rel, d, final=final) - elif self.persist_dir is None: - self._drain_grid_inplace(rel, d, final=final) else: - self._drain_grid_staged(rel, d, final=final) + nf, nb = self._drain_grid(rel, d, final=final, force_spill=force_spill) + backlog_files += nf + backlog_bytes += nb except Exception as e: - self._log(f"[stream] {rel}: {e} (will fall back to batch)") + self._log_every(f"drain:{rel}", f"[stream] {rel}: {e} (will fall back to batch)") # Drop the writer so a half-written file is rebuilt by the batch # pass rather than reused. w = self._writers.pop(rel, None) @@ -361,69 +538,157 @@ def _scan_once(self, *, final: bool) -> None: w.close() except Exception: pass + self._maybe_stats(backlog_files, backlog_bytes, final=final) # --- draining --------------------------------------------------------- - def _drain_grid_inplace(self, rel: str, d: Path, *, final: bool) -> None: - """Stream a grid diagnostic to NetCDF in place — no mirror, no reap. + def _drain_grid(self, rel: str, d: Path, *, final: bool, force_spill: bool) -> tuple[int, int]: + """Drain one grid diagnostic; returns its (files, bytes) pending at entry. - The original non-staged behavior: dumps accumulate in ``run_dir`` and - slot ``i`` is the ``i``-th dump, indexed positionally against the full - (never-pruned) dump list. + Non-staged: stream to NetCDF in place — no mirror, no reap. Staged: + stream + reap (mirroring unless ``discard_grid_h5``), or mirror-only + when spilling. The final sweep additionally catches a spilled + diagnostic up from its persist mirror before closing. """ - dumps = _io._sort_dumps(d) - if not dumps: - return - # Write-completeness: a dump is safe only once the next one exists (OSIRIS - # has moved on). The final sweep runs after OSIRIS is dead, so the last - # dump is safe too. + staged = self.persist_dir is not None + dumps = [p for p in _io._sort_dumps(d) if p not in self._bad] + # Write-completeness: a dump is safe only once the next one exists + # (OSIRIS has moved on). The final sweep runs after OSIRIS is dead, so + # the last dump is safe too. safe = dumps if final else dumps[:-1] - if not safe: - return # nothing safe to write yet (only the in-progress dump exists) + n_pending = len(safe) + pending_bytes = 0 + if staged: + for p in safe: + try: + pending_bytes += p.stat().st_size + except OSError: + pass + + if staged and not final: + spill = ( + force_spill + or rel in self._spilled + or (self.spill_backlog_files and n_pending > self.spill_backlog_files) + or (self.spill_backlog_bytes and pending_bytes > self.spill_backlog_bytes) + ) + if spill: + if rel not in self._spilled: + self._spilled.add(rel) + self._log( + f"[stream] {rel}: backlog {n_pending} dumps / {pending_bytes >> 20} MiB " + "over spill threshold -- mirror-only until finalize" + ) + for p in safe: + # Always mirror when spilling: these bytes exist nowhere else. + self._reap(rel, p, mirror=True) + self._stats["spilled"] += n_pending + self._stats["spilled_bytes"] += pending_bytes + return n_pending, pending_bytes + + candidates = list(safe) + if final and staged and rel in self._spilled: + # Catch up from the persist mirror first; merge by iteration so the + # append order stays monotonic. + spill_dir = self.persist_ms / rel + if spill_dir.is_dir(): + candidates = sorted( + [p for p in _io._sort_dumps(spill_dir) if p not in self._bad] + candidates, + key=_io._iter_from_name, + ) + if not candidates: + return n_pending, pending_bytes w = self._writers.get(rel) if w is None: - template = _io.load_grid_h5(dumps[0]) - w = StreamWriter(self.out_dir / f"{rel}.nc", template, source_dir=d) - self._writers[rel] = w - if w.n_written == 0: - w.append(template) # reuse the template we just loaded as slot 0 - - for p in safe[w.n_written :]: - w.append(_io.load_grid_h5(p)) + w, candidates = self._open_writer(rel, d, candidates) + if w is None: + return n_pending, pending_bytes + + for p in candidates: + if _io._iter_from_name(p) <= w.last_iter: + # Already in the file (resume / idempotency): just tidy up. + self._cleanup_consumed(rel, p) + continue + try: + da = _io.load_grid_h5(p) + except Exception as e: + self._quarantine(rel, p, e) + continue + w.append(da) + self._stats["appended"] += 1 + self._stats["appended_bytes"] += int(da.values.nbytes) + self._cleanup_consumed(rel, p) w.flush() if final: w.close() self._completed.add(rel) - - def _drain_grid_staged(self, rel: str, d: Path, *, final: bool) -> None: - """Stream a grid diagnostic *and* drain its dumps off the scratch. - - Because every handled dump is deleted from the scratch (:meth:`_reap`), - the directory only ever holds dumps not yet processed (plus the - in-progress last one), so every ``safe`` dump is new — no positional - bookkeeping against on-disk slots is needed. - """ - dumps = _io._sort_dumps(d) - if not dumps: - return - safe = dumps if final else dumps[:-1] - if not safe: + if staged and self.discard_grid_h5: + self._prune_empty(self.persist_ms / rel) + return n_pending, pending_bytes + + def _open_writer(self, rel: str, d: Path, candidates: list[Path]) -> tuple[StreamWriter | None, list[Path]]: + """Create (or reopen) the writer from the first readable candidate.""" + while candidates: + p0 = candidates[0] + try: + template = _io.load_grid_h5(p0) + except Exception as e: + self._quarantine(rel, p0, e) + candidates = candidates[1:] + continue + w = StreamWriter(self.out_dir / f"{rel}.nc", template, source_dir=d) + self._writers[rel] = w + if w.n_written == 0: + w.append(template) # reuse the template we just loaded as slot 0 + self._stats["appended"] += 1 + self._stats["appended_bytes"] += int(template.values.nbytes) + self._cleanup_consumed(rel, p0) + candidates = candidates[1:] + # On resume (n_written > 0) p0 stays; the iteration filter decides. + return w, candidates + return None, [] + + def _cleanup_consumed(self, rel: str, p: Path) -> None: + """Dispose of one dump whose contents are (now) in the NetCDF.""" + if self.persist_dir is None: + return # non-staged: the in-place tree is the durable copy + if self.persist_ms is not None and self.persist_ms in p.parents: + # A spilled dump drained from the persist mirror during finalize: + # in discard mode the .nc is the durable artifact, drop the HDF5. + if self.discard_grid_h5: + try: + p.unlink() + except OSError: + pass return + self._reap(rel, p, mirror=not self.discard_grid_h5) - w = self._writers.get(rel) - if w is None: - w = StreamWriter(self.out_dir / f"{rel}.nc", _io.load_grid_h5(dumps[0]), source_dir=d) - self._writers[rel] = w - for p in safe: - w.append(_io.load_grid_h5(p)) - self._reap(rel, p, mirror=not self.discard_grid_h5) - w.flush() + def _quarantine(self, rel: str, p: Path, err: Exception) -> None: + """Sideline an unreadable dump (``*.h5.bad``) and keep streaming. - if final: - w.close() - self._completed.add(rel) + Truncated dumps are exactly what an ENOSPC-choked run leaves behind; + one of them must not stall the whole diagnostic (previously the writer + was dropped and the watcher re-hit the same file every poll).""" + self._stats["quarantined"] += 1 + self._log_every(f"bad:{rel}", f"[stream] {rel}: quarantining unreadable dump {p.name}: {err}") + try: + p.rename(p.with_name(p.name + ".bad")) # ".bad": invisible to *.h5 listings + except OSError: + self._bad.add(p) # cannot rename: remember and skip it from now on + + def _prune_empty(self, dirpath: Path) -> None: + """Remove now-empty persist-mirror directories after a discard catch-up.""" + if self.persist_ms is None: + return + cur = dirpath + try: + while (cur == self.persist_ms or self.persist_ms in cur.parents) and cur.is_dir() and not any(cur.iterdir()): + cur.rmdir() + cur = cur.parent + except OSError: + pass def _drain_raw(self, rel: str, d: Path, *, final: bool) -> None: """Mirror a RAW diagnostic's completed dumps to persist + reap scratch.""" @@ -455,4 +720,4 @@ def _reap(self, rel: str, p: Path, *, mirror: bool = True) -> None: os.replace(tmp, dest) # atomic publish on the persist filesystem p.unlink() except Exception as e: - self._log(f"[stream] reap {rel}/{p.name} deferred: {e}") + self._log_every(f"reap:{rel}", f"[stream] reap {rel}/{p.name} deferred: {e}") diff --git a/tests/test_osiris/test_stream.py b/tests/test_osiris/test_stream.py index 140271f7..b4efb529 100644 --- a/tests/test_osiris/test_stream.py +++ b/tests/test_osiris/test_stream.py @@ -527,3 +527,77 @@ def test_collect_uses_streamed_binary(tmp_path: Path) -> None: assert not list(td.rglob("*.h5")) assert (td / "binary" / "FLD" / "e1.nc").exists() _assert_series_equiv(oio.load_series_nc(td / "binary" / "FLD" / "e1.nc"), oio.load_series(diag)) + + +# --- corrupt-dump quarantine + backlog spill valve ------------------------ + + +def test_corrupt_dump_quarantined_stream_continues(tmp_path: Path) -> None: + """One unreadable dump must not stall the diagnostic: it is renamed + ``*.h5.bad`` and the remaining dumps still stream (regression: job 56005549 + dropped the writer and re-hit the same corrupt file every poll, forever).""" + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=5, nx=8) + bad = diag / "e1-000020.h5" + bad.write_bytes(b"\x00" * 2048) # ENOSPC-style truncated file + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + completed = conv.finalize() + + assert "FLD/e1" in completed + assert not bad.exists() and (diag / "e1-000020.h5.bad").exists() + streamed = oio.load_series_nc(tmp_path / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == 4 + assert list(streamed["iter"].values) == [0, 10, 30, 40] + + +def test_spill_valve_mirrors_then_catches_up(tmp_path: Path) -> None: + """Over-threshold backlog flips a staged diagnostic to mirror-only (cheap + copy to persist, ramdisk keeps draining); finalize catches the NetCDF up + from the mirror so nothing is lost.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + diag = _write_field(stage, "e1", n_steps=6, nx=8) + + conv = ostream.StreamConverter( + stage, persist / "binary", poll_s=0.01, persist_dir=persist, spill_backlog_files=2 + ) + conv._scan_once(final=False) + # 5 safe dumps > threshold 2 -> spilled: mirrored + reaped, nothing streamed. + assert "FLD/e1" in conv._spilled + assert not (persist / "binary" / "FLD" / "e1.nc").exists() + assert len(list((persist / "MS" / "FLD" / "e1").glob("*.h5"))) == 5 + assert len(list(diag.glob("*.h5"))) == 1 # only the held-back newest dump + + completed = conv.finalize() + assert "FLD/e1" in completed + streamed = oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == 6 + _assert_series_equiv(streamed, oio.load_series(persist / "MS" / "FLD" / "e1")) + assert not list((stage / "MS").rglob("*.h5")) + + +def test_spill_valve_discard_mode_prunes_mirror(tmp_path: Path) -> None: + """In discard mode the spill mirror is a temporary holding area: after the + finalize catch-up the mirrored HDF5 (and its emptied dirs) are removed and + the streamed NetCDF is again the only grid artifact.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=6, nx=8) + + conv = ostream.StreamConverter( + stage, + persist / "binary", + poll_s=0.01, + persist_dir=persist, + discard_grid_h5=True, + spill_backlog_files=2, + ) + conv._scan_once(final=False) + assert len(list((persist / "MS" / "FLD" / "e1").glob("*.h5"))) == 5 # spilled bytes exist nowhere else + + completed = conv.finalize() + assert "FLD/e1" in completed + assert oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc").sizes["t"] == 6 + assert not (persist / "MS" / "FLD").exists() # mirror consumed + pruned + assert not list((stage / "MS").rglob("*.h5"))