From 5073964ed2ecfaf0e3e7783c1dc04cd0b71df775 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Fri, 22 May 2026 12:50:27 -0700 Subject: [PATCH 01/20] 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/20] .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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] =?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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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() == {}