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/adept/_base_.py b/adept/_base_.py index 949bb5a5..d2c7d761 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") @@ -334,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: @@ -342,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 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/__init__.py b/adept/osiris/__init__.py new file mode 100644 index 00000000..7815b551 --- /dev/null +++ b/adept/osiris/__init__.py @@ -0,0 +1,13 @@ +"""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..760b32fd --- /dev/null +++ b/adept/osiris/base.py @@ -0,0 +1,179 @@ +"""``BaseOsiris`` — the ADEPT module that drives OSIRIS.""" + +from __future__ import annotations + +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 + + +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: + """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.""" + 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", "./checkpoints") + deck_text = _deck.render_deck(self._sections) + result = _runner.run_osiris( + deck_text, + 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} + + 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..36a14569 --- /dev/null +++ b/adept/osiris/deck.py @@ -0,0 +1,323 @@ +"""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}: {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}: {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; {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..feb28627 --- /dev/null +++ b/adept/osiris/io.py @@ -0,0 +1,667 @@ +"""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 json +import re +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr + +_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.""" + 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(_DIAG_DTYPE) + 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 _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(_DIAG_DTYPE).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, *, 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 + ``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}") + + # 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] = [] + 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) + + +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. + + 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}") + first = load_grid_h5(dumps[0]) + + 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: + raise ValueError(f"Shape mismatch in series: {p} has {da.shape}, expected {first.shape}") + _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. + + 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). + + 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 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): + """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() + # 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} + 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 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")) + # 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 + + +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 _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, + 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. + + 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 + 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) + else: + ds = series_to_dataset(load_series(diags[relpath])) + dest = out_dir / f"{relpath}.nc" + dest.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(ds)) + written.append(dest) + 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", encoding=_compression_encoding(energy)) + written.append(dest) + except Exception as e: + print(f"[post] skipping HIST energy: {e}") + + 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. + + 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 + + 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 new file mode 100644 index 00000000..af819b07 --- /dev/null +++ b/adept/osiris/plots.py @@ -0,0 +1,1344 @@ +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 +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) 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: 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 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 + 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 + +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 ---------------------------------------------------- + + +def plot_spacetime( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + 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: + raise ValueError(f"plot_spacetime expects a 2D (t, x) array; got dims {da.dims}") + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + data = np.log10(np.abs(da.values) + 1e-30) if log else da.values # (t, x) + t = da.coords["t"].values + xname = next(d for d in da.dims if d != "t") + x = da.coords[xname].values + 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), + ) + scale = r"$\log_{10}$ " if log else "" + ax.set_title(title or f"{_display_name(da)} — {scale}spacetime ({orient})") + 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 rf"{_display_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, + *, + log: bool = True, + cmap: str = "viridis", + title: str | None = None, +) -> plt.Axes: + """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)) + # 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: + vmin, vmax = _nonzero_log_clim(raw) + plot_arr = np.log10(np.abs(raw) + 1e-30) + else: + plot_arr = raw + # 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)) + 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$)") + 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}") + da = _crop_spatial_to_box(da) + nt = da.coords["t"].size + t_skip = max(1, nt // n_panels) + idx = list(range(0, nt, t_skip)) + 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")] + 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 "" + fig.suptitle(title or rf"{_display_name(da)} — {scale}phase space at sampled times", y=1.02) + return fig + + +def field_energy_series(run_dir: str | Path) -> xr.DataArray: + """Sum ``(|E|^2 + |B|^2) / 2`` over space at every saved step. + + 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. + """ + 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( + 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$ [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 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. + """ + 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: + rel = f"FLD/{comp}" + if rel not in diags: + continue + try: + ser = _io.load_series(diags[rel]) + except Exception as e: # a bad component must not sink the rest + print(f"[plots] skipping field-energy component {comp}: {e}") + continue + if ser.ndim != 2: + continue + found = True + e_t = _field_energy_from_series(ser) + its = np.asarray(ser.coords["iter"].values) if "iter" in ser.coords else np.arange(ser.sizes["t"]) + ts = np.asarray(ser.coords["t"].values, dtype="float64") + for k in range(ts.size): + it = int(its[k]) + rec = by_iter.setdefault(it, {"E_energy": 0.0, "B_energy": 0.0}) + rec[group] += float(e_t[k]) + times.setdefault(it, float(ts[k])) + if not found or not by_iter: + raise RuntimeError(f"No field dumps available under {run_dir}") + iters = sorted(by_iter) + t = np.array([times[i] for i in iters]) + e_arr = np.array([by_iter[i]["E_energy"] for i in iters]) + b_arr = np.array([by_iter[i]["B_energy"] for i in iters]) + 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 + + +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, + show_light_line: bool = False, + v_th: float | None = None, + 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. + + 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. ``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: + 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=(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))) + 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 (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, 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 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) + 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 + + +# --- 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 ($t$ vs $x$)", 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: + """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 ("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: + 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 _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. 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: + _, 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) + 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$") + 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) -> dict[str, plt.Figure]: + """One figure per transverse component: right/left-going spacetime. + + 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(1, 2, figsize=(12, 5), sharey=True) + for ax, side in zip(axes, ("right", "left"), strict=False): + da = parts[side] + 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.02) + fig.tight_layout() + figs[comp] = fig + return figs + + +# --- 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`` 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 "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. + + 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, + *, + v_th: float | None = None, + dpi: int = 120, + n_panels: int = 8, + 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. + + ``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. + + ``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) + 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) + + # --- 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] + 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 # 2D-in-space field plots deferred + + fig, ax = plt.subplots(figsize=(6, 4)) + plot_spacetime(ser, ax=ax) + 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") + + # 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", + ) + + # --- 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/"): + 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", + ) + + # --- 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, + 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: + # 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, + show_initial=True, + title=f"{label} — 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/"): + continue + parts = diag_rel.split("/") + if len(parts) < 3: + continue + ps_name, species = parts[1], parts[2] + 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), + f"phasespace_evolution/{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).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)) + plot_energy_vs_time(run_dir, ax=ax) + 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 + + +# --- 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 _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 _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" + 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: + return rf"{_tex(long)} [{_tex(units)}]" + return _tex(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 _display_name(da: xr.DataArray) -> str: + """``_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 = _label_tex(_long_name(da)) + units = da.attrs.get("units", "") + return rf"{name} [{_tex(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 _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) + + +# --- 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 +sim_box_bound = _sim_box_bound diff --git a/adept/osiris/post.py b/adept/osiris/post.py new file mode 100644 index 00000000..23269336 --- /dev/null +++ b/adept/osiris/post.py @@ -0,0 +1,170 @@ +"""Post-processing for OSIRIS runs. + +After ``BaseOsiris.__call__`` finishes, ``collect`` walks the ``MS/`` +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 +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 +# 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 _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 _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: 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) + 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"]), + "exit_code": float(solver["exit_code"]), + "field_energy_final": _total_field_energy(ms), + "final_iter": float(_final_iter(ms)), + } + + # Convert each diagnostic's time history to an xarray netCDF. + if ms.is_dir(): + _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 + + # 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 (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")) + _plots.save_canned_plots(run_dir, td / "plots", **kwargs) + 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 + if src.exists(): + shutil.copy(src, td / fname) + + return {"metrics": metrics} diff --git a/adept/osiris/regen.py b/adept/osiris/regen.py new file mode 100644 index 00000000..4ec88be5 --- /dev/null +++ b/adept/osiris/regen.py @@ -0,0 +1,165 @@ +"""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`` / ``derived_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``, + ``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(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.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( + "--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/adept/osiris/runner.py b/adept/osiris/runner.py new file mode 100644 index 00000000..11ce118e --- /dev/null +++ b/adept/osiris/runner.py @@ -0,0 +1,190 @@ +"""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 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 + +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 = "./checkpoints", + env: dict[str, str] | None = None, + launcher: str = "srun", + 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 = [launcher, "-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 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 + # 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..d6d154cc --- /dev/null +++ b/configs/osiris/twostream-1d-short.yaml @@ -0,0 +1,13 @@ +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: ./checkpoints + overrides: + time: {tmax: 1.0} diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml new file mode 100644 index 00000000..157880dd --- /dev/null +++ b/configs/osiris/twostream-1d.yaml @@ -0,0 +1,22 @@ +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: ./checkpoints + # 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: + # 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 + # 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 new file mode 100644 index 00000000..ed29b533 --- /dev/null +++ b/docs/osiris-adept-usage.md @@ -0,0 +1,172 @@ +# 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: ./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 + 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: + diagnostics_to_log: null # null = all; or [e1, charge, …] + v_th: 0.1 # optional: overlays the Bohm–Gross + # Langmuir branch on ω–k plots + omega_k_zoom: 4.0 # (k, ω) half-width [ω_p] for the + # 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. + +## 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.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 | `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` | 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) | +| `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 | +| `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. + +> **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 + +```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/decks/F-Tsung_2d_lpi_deck b/tests/test_osiris/decks/F-Tsung_2d_lpi_deck new file mode 100644 index 00000000..5d528722 --- /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, +} 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..dae72249 --- /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 new file mode 100644 index 00000000..da16803b --- /dev/null +++ b/tests/test_osiris/test_deck_roundtrip.py @@ -0,0 +1,136 @@ +"""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 + +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 = [ + 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, 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) + 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((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: + assert allowed.match(k), f"Unsafe MLflow param key: {k!r}" diff --git a/tests/test_osiris/test_diagnostics_plots.py b/tests/test_osiris/test_diagnostics_plots.py new file mode 100644 index 00000000..c4084d0d --- /dev/null +++ b/tests/test_osiris/test_diagnostics_plots.py @@ -0,0 +1,175 @@ +"""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"] 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..c95461e7 --- /dev/null +++ b/tests/test_osiris/test_io_and_plots.py @@ -0,0 +1,94 @@ +"""Tests for the io.py loaders and plots.py canned views.""" + +from __future__ import annotations + +import os +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 + +# 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: + 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: + 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_plots_new_views.py b/tests/test_osiris/test_plots_new_views.py new file mode 100644 index 00000000..bf25861d --- /dev/null +++ b/tests/test_osiris/test_plots_new_views.py @@ -0,0 +1,234 @@ +"""Tests for the OSIRIS plotting additions. + +Synthesizes tiny OSIRIS-shaped runs (no real solver needed) to exercise: + + - proper-LaTeX axis/value labels (``_tex`` wrapping) + - 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 + +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$]" + + +# --- 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) + 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 ---------------------------------------------------- + + +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/e2", + "currents/spacetime", + "currents/lineouts", + "profiles/electron/density", + "profiles/electron/temperature", + "field_decomp/e2", + "field_decomp/e3", + } + 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 new file mode 100644 index 00000000..dddcbdf7 --- /dev/null +++ b/tests/test_osiris/test_post_netcdf.py @@ -0,0 +1,473 @@ +"""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() + + +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 + + +# --- 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() + + +# --- 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) diff --git a/tests/test_osiris/test_runner.py b/tests/test_osiris/test_runner.py new file mode 100644 index 00000000..70f77d61 --- /dev/null +++ b/tests/test_osiris/test_runner.py @@ -0,0 +1,62 @@ +"""Smoke tests for the OSIRIS subprocess runner.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +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("") + 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 (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 + # 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=OSIRIS_BIN_1D, + mpi_ranks=1, + run_root=tmp_path, + ) + assert "OSIRIS" in str(excinfo.value) 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() == {}