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/_tf1d/modules.py b/adept/_tf1d/modules.py index bdd53b4e..9816785d 100644 --- a/adept/_tf1d/modules.py +++ b/adept/_tf1d/modules.py @@ -59,7 +59,7 @@ def write_units(self): wp0 = np.sqrt(n0 * self.ureg.e**2.0 / (self.ureg.m_e * self.ureg.epsilon_0)).to("rad/s") tp0 = (1 / wp0).to("fs") - v0 = np.sqrt(2.0 * T0 / self.ureg.m_e).to("m/s") + v0 = np.sqrt(T0 / self.ureg.m_e).to("m/s") x0 = (v0 / wp0).to("nm") c_light = _Q(1.0 * self.ureg.c).to("m/s") / v0 beta = (v0 / self.ureg.c).to("dimensionless") @@ -72,7 +72,7 @@ def write_units(self): sim_duration = (self.cfg["grid"]["tmax"] * tp0).to("ps") # collisions - logLambda_ee = 23.5 - np.log(n0.magnitude**0.5 / T0.magnitude**-1.25) + logLambda_ee = 23.5 - np.log(n0.magnitude**0.5 * T0.magnitude**-1.25) logLambda_ee -= (1e-5 + (np.log(T0.magnitude) - 2) ** 2.0 / 16) ** 0.5 nuee = _Q(2.91e-6 * n0.magnitude * logLambda_ee / T0.magnitude**1.5, "Hz") nuee_norm = nuee / wp0 diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 95f28d1b..13010629 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -71,7 +71,6 @@ class GridConfig(BaseModel): vmax: float | None = None vmin: float | None = None parallel: tuple[str, ...] | bool = False - c_light: float | None = None @field_validator("parallel", mode="before") @classmethod diff --git a/adept/_vlasov1d/grid.py b/adept/_vlasov1d/grid.py index 91b6d09b..b16db67e 100644 --- a/adept/_vlasov1d/grid.py +++ b/adept/_vlasov1d/grid.py @@ -52,7 +52,7 @@ def __init__( self.tmin = tmin # Compute dx - self.dx = xmax / nx + self.dx = (xmax - xmin) / nx # Override dt for EM wave stability if needed if should_override_dt_for_em_waves: diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index d34e774f..534047ac 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -72,7 +72,10 @@ def _initialize_supergaussian_distribution_( # Thermal velocity: v_t = sqrt(T/m) v_thermal = np.sqrt(T0 / mass) - # Alpha factor for supergaussian normalization + # Alpha factor for supergaussian normalization. This fixes the moment ratio + # / = 3*T0/mass for every order m; the realized *variance* equals + # T0/mass only at m=2 (Maxwellian). For m>2 flat-tops the second-moment + # temperature exceeds T0 (x1.24 at m=3, x1.37 at m=4). See docs config.md. alpha = np.sqrt(3.0 * gamma_3_over_m(supergaussian_order) / gamma_5_over_m(supergaussian_order)) single_dist = -(np.power(np.abs((vax[None, :] - v0) / (alpha * v_thermal)), supergaussian_order)) diff --git a/adept/_vlasov1d/modules.py b/adept/_vlasov1d/modules.py index 484c7ecd..c4840e77 100644 --- a/adept/_vlasov1d/modules.py +++ b/adept/_vlasov1d/modules.py @@ -118,6 +118,9 @@ def write_units(self) -> dict: sim_duration = (grid.tmax * norm.tau).to("ps") nu_ee = norm.approximate_ee_collision_frequency() + # e-e collision rate in code units (1/wp0). This is what the FP/Krook + # `baseline` rates should be compared against. + nu_ee_norm = (nu_ee * norm.tau).to("").magnitude beta = 1.0 / norm.speed_of_light_norm() @@ -131,6 +134,7 @@ def write_units(self) -> dict: "beta": beta, "x0": norm.L0.to("nm"), "nuee": nu_ee.to("Hz"), + "nuee_norm": nu_ee_norm, "logLambda_ee": norm.logLambda_ee(), "box_length": box_length, "box_width": box_width, @@ -238,10 +242,12 @@ def get_solver_quantities(self) -> dict: cfg_grid["species_grids"][species_name]["one_over_kvr"] = jnp.array(one_over_kvr) # Build species parameters (charge, mass, charge-to-mass ratio) + # T0 is the bulk (first-listed component) temperature, used e.g. by the Krook target cfg_grid["species_params"][species_name] = { "charge": species_cfg.charge, "mass": species_cfg.mass, "charge_to_mass": species_cfg.charge / species_cfg.mass, + "T0": self.simulation.species_distributions[species_name][0].T0, } cfg_grid["n_prof_total"] = n_prof_total diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index 76a999af..ff73262c 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -227,7 +227,7 @@ def __call__(self, f_dict: dict, prev_ex: jnp.ndarray, dt: jnp.float64): class AmpereSolver: """Ampere solver using current density to evolve electric field. - Solves: a single Euler step of ∂E/∂t = -j, where j = Σ_s (q_s/m_s) ∫v f_s dv + Solves: a single Euler step of ∂E/∂t = -j, where j = Σ_s q_s ∫v f_s dv is the total current density from all species. """ diff --git a/adept/_vlasov1d/solvers/pushers/fokker_planck.py b/adept/_vlasov1d/solvers/pushers/fokker_planck.py index 23166bfd..dfb2a5b1 100644 --- a/adept/_vlasov1d/solvers/pushers/fokker_planck.py +++ b/adept/_vlasov1d/solvers/pushers/fokker_planck.py @@ -76,9 +76,9 @@ def compute_vbar(self, f: Array) -> Array: f: Distribution function, shape (..., nv) Returns: - vbar: Mean velocity , shape (...) + vbar: Mean velocity = ∫v f dv / ∫f dv, shape (...) """ - return jnp.sum(f * self.v, axis=-1) * self.dv + return jnp.sum(f * self.v, axis=-1) / jnp.sum(f, axis=-1) def compute_D(self, f: Array, beta: Array) -> Array: """ @@ -242,7 +242,11 @@ def __init__(self, cfg: Mapping[str, Any]): self.cfg = cfg v = cfg["grid"]["species_grids"]["electron"]["v"] dv = cfg["grid"]["species_grids"]["electron"]["dv"] - f_mx = np.exp(-(v[None, :] ** 2.0) / 2.0) + params = cfg["grid"].get("species_params", {}).get("electron", {}) + T0 = params.get("T0", 1.0) + mass = params.get("mass", 1.0) + # Maxwellian with variance T0/m (the species' bulk thermal width) + f_mx = np.exp(-(v[None, :] ** 2.0) / (2.0 * T0 / mass)) self.f_mx = f_mx / np.sum(f_mx, axis=1)[:, None] / dv self.dv = dv diff --git a/adept/_vlasov1d/solvers/pushers/vlasov.py b/adept/_vlasov1d/solvers/pushers/vlasov.py index 384d6f43..828dbf7a 100644 --- a/adept/_vlasov1d/solvers/pushers/vlasov.py +++ b/adept/_vlasov1d/solvers/pushers/vlasov.py @@ -28,7 +28,7 @@ def __init__(self, cfg, interp_e): """Create interpolation helpers for x-advection and v-advection steps.""" self.x = cfg["grid"]["x"] self.v = cfg["grid"]["v"] - self.f_interp = partial(interp2d, x=self.x, y=self.v, period=cfg["grid"]["xmax"]) + self.f_interp = partial(interp2d, x=self.x, y=self.v, period=cfg["grid"]["xmax"] - cfg["grid"]["xmin"]) self.dt = cfg["grid"]["dt"] self.dummy_x = jnp.ones_like(self.x) self.dummy_v = jnp.ones_like(self.v)[None, :] diff --git a/adept/_vlasov1d/storage.py b/adept/_vlasov1d/storage.py index 1075ed48..6899e049 100644 --- a/adept/_vlasov1d/storage.py +++ b/adept/_vlasov1d/storage.py @@ -137,11 +137,12 @@ def _calc_moment_(inp, _dv=dv): f = y[species_name] species_moments = {} species_moments["n"] = _calc_moment_(f) - species_moments["v"] = _calc_moment_(f * v[None, :]) + species_moments["j"] = _calc_moment_(f * v[None, :]) + species_moments["v"] = species_moments["j"] / species_moments["n"] v_m_vbar = v[None, :] - species_moments["v"][:, None] species_moments["p"] = _calc_moment_(f * v_m_vbar**2.0) species_moments["q"] = _calc_moment_(f * v_m_vbar**3.0) - species_moments["-flogf"] = _calc_moment_(f * jnp.log(jnp.abs(f))) + species_moments["-flogf"] = _calc_moment_(-jnp.abs(f) * jnp.log(jnp.abs(f))) species_moments["f^2"] = _calc_moment_(f * f) result[species_name] = species_moments @@ -292,9 +293,11 @@ def save(t, y, args): scalars = {} # Compute scalars for each species + mean_kinetic_energy = 0.0 for species_name in species_names: v = species_grids[species_name]["v"][None, :] dv = species_grids[species_name]["dv"] + mass = cfg["grid"]["species_params"][species_name]["mass"] def _calc_mean_moment_(inp, _dv=dv): return jnp.mean(jnp.sum(inp, axis=1) * _dv) @@ -306,12 +309,19 @@ def _calc_mean_moment_(inp, _dv=dv): scalars[f"mean_q_{species_name}"] = _calc_mean_moment_(f * v**3.0) scalars[f"mean_-flogf_{species_name}"] = _calc_mean_moment_(-jnp.log(jnp.abs(f)) * jnp.abs(f)) scalars[f"mean_f2_{species_name}"] = _calc_mean_moment_(f * f) + mean_kinetic_energy += 0.5 * mass * scalars[f"mean_P_{species_name}"] # Shared field scalars (not species-specific) scalars["mean_de2"] = jnp.mean(y["de"] ** 2.0) scalars["mean_e2"] = jnp.mean(y["e"] ** 2.0) scalars["mean_pond"] = jnp.mean(-0.5 * jnp.gradient(y["a"] ** 2.0, cfg["grid"]["dx"])[1:-1]) + # Energy conservation monitor (electrostatic): x-averaged kinetic + E-field + # energy density. Transverse (EM) field energy is not included. + scalars["mean_kinetic_energy"] = mean_kinetic_energy + scalars["mean_field_energy"] = 0.5 * scalars["mean_e2"] + scalars["mean_total_energy"] = mean_kinetic_energy + 0.5 * scalars["mean_e2"] + return scalars return save diff --git a/adept/normalization.py b/adept/normalization.py index aa9f7714..993e6458 100644 --- a/adept/normalization.py +++ b/adept/normalization.py @@ -34,7 +34,7 @@ class PlasmaNormalization: def logLambda_ee(self) -> float: n0_cc = self.n0.to("1/cc").magnitude T0_eV = self.T0.to("eV").magnitude - logLambda_ee = 23.5 - jnp.log(n0_cc**0.5 / T0_eV**-1.25) + logLambda_ee = 23.5 - jnp.log(n0_cc**0.5 * T0_eV**-1.25) logLambda_ee -= (1e-5 + (jnp.log(T0_eV) - 2) ** 2.0 / 16) ** 0.5 return logLambda_ee @@ -78,9 +78,13 @@ def electron_debye_normalization(n0_str, T0_str): """ Returns the electron thermal normalization for the given density and temperature. Unit quantities are: - - Debye length - - Electron thermal velocity - - Langmuir oscillation frequency + - L0 = Debye length, sqrt(eps0 T0 / (n0 e^2)) + - v0 = electron thermal velocity sqrt(T0/m_e) (RMS / standard-deviation + convention: a T=1 Maxwellian is exp(-v^2/2) with unit variance) + - tau = 1/wp0 (inverse Langmuir oscillation frequency) + + Under this convention a code wavenumber is k*lambda_De and the Bohm-Gross + dispersion reads w^2 = 1 + 3 k^2. """ n0 = UREG.Quantity(n0_str) T0 = UREG.Quantity(T0_str) @@ -88,7 +92,7 @@ def electron_debye_normalization(n0_str, T0_str): wp0 = ((n0 * UREG.e**2.0 / (UREG.m_e * UREG.epsilon_0)) ** 0.5).to("rad/s") tau = 1 / wp0 - v0 = ((2.0 * T0 / UREG.m_e) ** 0.5).to("m/s") + v0 = ((T0 / UREG.m_e) ** 0.5).to("m/s") x0 = (v0 / wp0).to("nm") return PlasmaNormalization(m0=UREG.m_e, q0=UREG.e, n0=n0, T0=T0, L0=x0, v0=v0, tau=tau) @@ -116,3 +120,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..903031c5 --- /dev/null +++ b/adept/osiris/base.py @@ -0,0 +1,223 @@ +"""``BaseOsiris`` — the ADEPT module that drives OSIRIS.""" + +from __future__ import annotations + +from typing import Any + +import math + +from adept._base_ import ADEPTModule +from adept.normalization import UREG, skin_depth_normalization, skin_depth_normalization_from_frequency +from adept.osiris import deck as _deck +from adept.osiris import density as _density +from adept.osiris import post as _post +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) + + # Optional: scale the box from a target density gradient scale length + # (mirrors adept's _lpse2d / kinetic_srs grid sizing). Runs *after* + # overrides, so an explicit space.xmax override is superseded by the + # gradient-derived box. The computed quantities are stashed back under + # osiris.density.derived for MLflow provenance. + density_cfg = osiris_cfg.get("density") + computed = _density.apply_gradient_scale_length(self._sections, density_cfg) + if computed is not None: + density_cfg.setdefault("derived", {}).update(computed) + + # Surface the parsed (post-override) deck inside cfg so adept's + # log_params picks every parameter up as a flat MLflow param. + # The raw overrides dict is intentionally popped above to avoid + # 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. + + When the deck launches a laser (an ``antenna`` / ``zpulse_speckle`` / + ``zpulse`` section with ``a0`` and ``omega0``), the physical drive + scales are added too: ``w_laser`` (rad/s), ``laser_wavelength`` (nm), + ``laser_a0`` and ``laser_intensity`` — the peak intensity of a + linearly polarized drive in W/cm^2 (the ICF convention + ``I * lambda_um^2 = 1.37e18 * a0^2``). + """ + sim = self._iter_first_section("simulation") + n0 = sim.get("n0") + 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") + + # Laser drive in physical / ICF units. The deck's laser section carries + # a0 and omega0 (in wp0 units); with wp0 fixed above these give the + # laser frequency, the vacuum wavelength lambda = 2 pi c / w_laser, and + # the peak intensity of a linearly polarized drive, + # I = eps0 c E0^2 / 2 with E0 = a0 m_e c w_laser / e — equivalently the + # usual ICF convention I[W/cm^2] * lambda[um]^2 = 1.37e18 * a0^2. + # Same section priority as the SRS postproc: antenna (1D), then + # zpulse_speckle / zpulse (2D). + for sec_name in ("antenna", "zpulse_speckle", "zpulse"): + laser = self._iter_first_section(sec_name) + if laser.get("a0") is not None and laser.get("omega0") is not None: + a0, omega0 = float(laser["a0"]), float(laser["omega0"]) + w_laser = omega0 / norm.tau + e_peak = a0 * w_laser * UREG.m_e * UREG.c / UREG.e + quants["w_laser"] = w_laser.to("rad/s") + quants["laser_wavelength"] = (2.0 * math.pi * UREG.c / w_laser).to("nm") + quants["laser_a0"] = a0 + quants["laser_intensity"] = (e_peak**2 * UREG.epsilon_0 * UREG.c / 2.0).to("W/cm^2") + break + + self.cfg.setdefault("units", {})["derived"] = quants + return quants + + 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"), + stream_convert=bool(osiris_cfg.get("stream_convert", True)), + stream_poll_s=float(osiris_cfg.get("stream_poll_s", 10.0)), + stage_root=osiris_cfg.get("stage_root"), + stage_discard_h5=bool(osiris_cfg.get("stage_discard_h5", False)), + ) + return {"solver result": result} + + 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..ef617fe6 --- /dev/null +++ b/adept/osiris/deck.py @@ -0,0 +1,339 @@ +"""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."} + +# OSIRIS decks quote strings with either single (Fortran convention) or double +# quotes; both are treated as simple, non-nesting delimiters. +_QUOTE_CHARS = "\"'" + +_IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +_KEY_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*(?:\([^)]*\))?") +_INT_RE = re.compile(r"^[+-]?\d+$") +_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.""" + quote = "" + out: list[str] = [] + for ch in line: + if quote: + if ch == quote: + quote = "" + out.append(ch) + elif ch in _QUOTE_CHARS: + quote = ch + out.append(ch) + elif ch == "!": + break + else: + out.append(ch) + return "".join(out) + + +def _split_top_commas(s: str) -> list[str]: + """Split on commas not inside a quoted string.""" + pieces: list[str] = [] + buf: list[str] = [] + quote = "" + for ch in s: + if quote: + if ch == quote: + quote = "" + buf.append(ch) + elif ch in _QUOTE_CHARS: + quote = ch + buf.append(ch) + elif ch == ",": + 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 len(t) >= 2 and t[0] in _QUOTE_CHARS and t[-1] == t[0]: + return t[1:-1] + low = t.lower() + if low in _TRUE_TOKENS: + 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 + quote = "" + while i < n: + ch = src[i] + if quote: + if ch == quote: + quote = "" + i += 1 + elif ch in _QUOTE_CHARS: + quote = ch + i += 1 + elif ch == "}": + break + elif ch == ",": + # Peek past whitespace: if the next token is an identifier + # followed by '=', this comma terminates the current + # key=value pair; otherwise it's an intra-value separator + # (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/density.py b/adept/osiris/density.py new file mode 100644 index 00000000..0a191780 --- /dev/null +++ b/adept/osiris/density.py @@ -0,0 +1,260 @@ +"""Adaptive box sizing from a density gradient scale length. + +OSIRIS decks fix the simulation box (``space.xmax``) and the density ramp +(``profile.fx`` / ``profile.x``) by hand. This module lets a manifest instead +request a target *gradient scale length* ``L`` and have the box scaled so the +linear density ramp realizes that ``L`` at a chosen reference density — mirroring +ADEPT's envelope (``adept/_lpse2d/helpers.py``) and Vlasov (``kinetic_srs``) +solvers, which size their grid the same way. + +Geometry. The deck's ``profile`` is a piecewise-linear density ramp +``n(x): nmin -> nmax`` across its interior control points. For a linear ramp the +local scale length is + + L(x) = n(x) / (dn/dx) = n(x) * ramp_span / (nmax - nmin) + +so requiring ``L(n_ref) = L_target`` at the reference density ``n_ref`` (default +the quarter-critical surface, ``n_c/4 = 0.25`` in ``n_c`` units) fixes the ramp +span: + + ramp_span = L_target / n_ref * (nmax - nmin) + +which is the same relation ADEPT uses (``Lgrid = L / 0.25 * (nmax - nmin)``). + +The transform is then a single spatial **scale factor** + + s = ramp_span_target / ramp_span_deck + +applied uniformly to every length in the deck: ``space.xmin`` / ``space.xmax``, +every ``profile.x`` array, and every ``diag_species`` phase-space window +(``ps_xmin`` / ``ps_xmax``). ``grid.nx_p`` is scaled by ``s`` as well so the cell +size ``dx`` is held fixed (rounded up to a clean multiple of the per-dimension +node count for even domain decomposition). Time (``dt``, ``tmax``) is untouched, +so the CFL ratio is preserved. + +Activation. Only when the manifest carries an ``osiris.density`` block with a +``gradient_scale_length`` (alias ``gradient scale length``). Without it, decks +run with their hand-set box, unchanged. + +1D only: raises ``NotImplementedError`` for multi-dimensional decks (the SRS LPI +use case is 1D). +""" + +from __future__ import annotations + +import math +from typing import Any + +from adept.normalization import UREG, skin_depth_normalization, skin_depth_normalization_from_frequency +from adept.osiris.deck import Sections + + +def apply_gradient_scale_length(sections: Sections, density_cfg: dict | None) -> dict | None: + """Scale the OSIRIS box so the density ramp has gradient scale length ``L``. + + ``density_cfg`` is the manifest's ``osiris.density`` block. Returns a dict of + the computed quantities (for logging) or ``None`` when the feature is + inactive (no ``density_cfg``, or no gradient scale length given). + + Mutates ``sections`` in place. Recognized ``density_cfg`` keys: + + * ``gradient_scale_length`` (alias ``gradient scale length``) — target ``L``, + a unit string (``"300um"``) or a number already in ``c/wp0`` units. + * ``min`` / ``max`` — optional ``nmin`` / ``nmax`` in ``n_c`` units; default is + to read the ramp's interior endpoints from the deck's ``profile.fx``. When + given, they are also written into the primary ``profile.fx`` endpoints. + * ``reference_density`` — density (``n_c`` units) at which ``L`` is defined; + default ``0.25`` (quarter-critical). + """ + if not density_cfg: + return None + L_spec = density_cfg.get("gradient_scale_length", density_cfg.get("gradient scale length")) + if L_spec is None: + return None + + reference_density = float(density_cfg.get("reference_density", 0.25)) + + simulation = _first_section(sections, "simulation") + space = _first_section(sections, "space") + grid = _first_section(sections, "grid") + profile = _first_section(sections, "profile") + missing = [name for name, sec in (("space", space), ("grid", grid), ("profile", profile)) if sec is None] + if missing: + raise ValueError( + f"osiris.density (adaptive box sizing) requires the deck to define section(s): {', '.join(missing)}" + ) + + # --- 1D only ---------------------------------------------------------- + nx_key = _array_key(grid, "nx_p") + if nx_key is None: + raise ValueError("osiris.density requires grid.nx_p in the deck") + nx_val = grid[nx_key] + if isinstance(nx_val, list) and len(nx_val) > 1: + raise NotImplementedError( + "osiris.density adaptive box sizing is implemented for 1D decks only " + f"(grid.{nx_key} has {len(nx_val)} dimensions)" + ) + nx_old = int(nx_val[0] if isinstance(nx_val, list) else nx_val) + + # --- ramp geometry & densities from the primary profile --------------- + x_key = _array_key(profile, "x") + fx_key = _array_key(profile, "fx") + if x_key is None or fx_key is None: + raise ValueError("osiris.density requires profile.x and profile.fx in the deck") + x_old = [float(v) for v in _as_list(profile[x_key])] + fx = [float(v) for v in _as_list(profile[fx_key])] + if len(x_old) < 3 or len(fx) < 3: + raise ValueError( + f"profile needs >= 3 control points to define a density ramp " + f"(got {len(x_old)} x-points, {len(fx)} fx-values)" + ) + + given_minmax = "min" in density_cfg or "max" in density_cfg + if given_minmax: + nmin = float(density_cfg["min"]) + nmax = float(density_cfg["max"]) + else: + nmin, nmax = fx[1], fx[-2] # interior ramp endpoints + if nmax <= nmin: + raise ValueError(f"density max ({nmax}) must exceed min ({nmin})") + + ramp_span_old = x_old[-2] - x_old[1] + if ramp_span_old <= 0: + raise ValueError(f"profile.{x_key} interior control points must be increasing (got ramp span {ramp_span_old})") + + # --- target ramp span and the resulting scale factor ------------------ + L_norm = _normalized_length(L_spec, simulation) + ramp_span_new = L_norm * (nmax - nmin) / reference_density + s = ramp_span_new / ramp_span_old + + xmax_old = _scalar(space, "xmax") + if xmax_old is None: + raise ValueError("osiris.density requires space.xmax in the deck") + xmin_old = _scalar(space, "xmin") or 0.0 + box_new = (xmax_old - xmin_old) * s + + # constant dx -> nx scales with the box; round UP (don't under-resolve) + # to a clean multiple of the per-dim node count so OSIRIS splits evenly. + node_count = _node_count(sections) + nx_new = _round_up_multiple(math.ceil(nx_old * s), node_count) + + # --- apply the rescale in place --------------------------------------- + if given_minmax: + fx[1], fx[-2] = nmin, nmax + profile[fx_key] = fx if isinstance(profile[fx_key], list) else fx[0] + + _scale_array(space, "xmin", s) + _scale_array(space, "xmax", s) + for sec_name, params in sections: + if sec_name == "profile": + _scale_array(params, "x", s) + elif sec_name == "diag_species": + _scale_array(params, "ps_xmin", s) + _scale_array(params, "ps_xmax", s) + grid[nx_key] = [nx_new] if isinstance(nx_val, list) else nx_new + + return { + "gradient_scale_length_requested": str(L_spec), + "gradient_scale_length_norm": L_norm, + "reference_density": reference_density, + "density_min": nmin, + "density_max": nmax, + "ramp_span_norm": ramp_span_new, + "box_norm": box_new, + "nx": nx_new, + "scale_factor": s, + } + + +# --- helpers -------------------------------------------------------------- + + +def _first_section(sections: Sections, name: str) -> dict | None: + for sec_name, params in sections: + if sec_name == name: + return params + return None + + +def _array_key(params: dict, base: str) -> str | None: + """Key in ``params`` whose base name (before any ``(``) equals ``base``.""" + for k in params: + if k == base or k.split("(", 1)[0] == base: + return k + return None + + +def _as_list(v: Any) -> list: + return list(v) if isinstance(v, list) else [v] + + +def _scalar(params: dict, base: str) -> float | None: + key = _array_key(params, base) + if key is None: + return None + v = params[key] + return float(v[0]) if isinstance(v, list) else float(v) + + +def _scale_array(params: dict, base: str, factor: float) -> bool: + """Multiply the value(s) at ``base`` by ``factor`` in place; preserve shape.""" + key = _array_key(params, base) + if key is None: + return False + v = params[key] + if isinstance(v, list): + params[key] = [round(float(x) * factor, 10) for x in v] + else: + params[key] = round(float(v) * factor, 10) + return True + + +def _node_count(sections: Sections) -> int: + node_conf = _first_section(sections, "node_conf") + if node_conf is None: + return 1 + key = _array_key(node_conf, "node_number") + if key is None: + return 1 + val = node_conf[key] + first = val[0] if isinstance(val, list) else val + try: + return max(1, int(first)) + except (TypeError, ValueError): + return 1 + + +def _round_up_multiple(n: int, m: int) -> int: + if m <= 1: + return int(n) + return int(math.ceil(n / m) * m) + + +def _normalized_length(L_spec: Any, simulation: dict | None) -> float: + """Target ``L`` in skin-depth (``c/wp0``) units. + + A bare number is taken to be already normalized; a unit string is converted + using the deck's reference density (``simulation.n0`` / ``omega_p0``). + """ + if isinstance(L_spec, (int, float)) and not isinstance(L_spec, bool): + return float(L_spec) + norm = _skin_depth_norm(simulation) + if norm is None: + raise ValueError( + "osiris.density.gradient_scale_length was given with physical units " + f"({L_spec!r}) but the deck's simulation section has neither n0 nor " + "omega_p0 to set the c/wp0 length scale" + ) + return float((UREG.Quantity(str(L_spec)) / norm.L0).to("").magnitude) + + +def _skin_depth_norm(simulation: dict | None): + if not simulation: + return None + n0 = simulation.get("n0") + omega_p0 = simulation.get("omega_p0") + if n0 is not None: + return skin_depth_normalization(f"{n0} / cc") + if omega_p0 is not None: + return skin_depth_normalization_from_frequency(f"{omega_p0} rad/s") + return None diff --git a/adept/osiris/io.py b/adept/osiris/io.py new file mode 100644 index 00000000..99f3eb0d --- /dev/null +++ b/adept/osiris/io.py @@ -0,0 +1,885 @@ +"""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 +import shutil +from collections.abc import Sequence +from contextlib import contextmanager +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 + + +_BASE_RE = re.compile(r"-\d+\.h5$") + + +def _dump_base(path: Path) -> str: + """A dump's series base = its filename minus the trailing ``-.h5``. + + ``e1-savg-000123.h5`` -> ``e1-savg``. A directory usually holds one base, but + multiple line/report diagnostics of the same field share a directory and + differ only by an index in the base — e.g. two ``s1,...,line`` reports produce + ``s1-tavg-line-x2-01`` and ``s1-tavg-line-x2-02``. Those are distinct time + series and must not be stacked into one. + """ + return _BASE_RE.sub("", path.name) + + +def _series_dumps(directory: Path) -> dict[str, list[Path]]: + """Group ``directory``'s ``.h5`` dumps into time series keyed by base name.""" + groups: dict[str, list[Path]] = {} + for p in directory.iterdir(): + if p.is_file() and p.suffix == ".h5": + groups.setdefault(_dump_base(p), []).append(p) + for dumps in groups.values(): + dumps.sort(key=_iter_from_name) + return groups + + +def _resolve_series(handle: Path) -> tuple[Path, str | None]: + """Resolve a diagnostic handle to ``(directory, base)``. + + A handle from :func:`list_diagnostics` is either a real dump directory (one + series, ``base`` is ``None``) or a synthetic ``/`` path naming one + report series inside a directory that holds several. + """ + handle = Path(handle) + if handle.is_dir(): + return handle, None + return handle.parent, handle.name + + +def _sort_dumps(handle: Path, base: str | None = None) -> list[Path]: + """Sorted ``.h5`` dumps for one diagnostic series. + + ``handle`` may be a dump directory or a synthetic per-report handle + (``/``); ``base`` selects one series inside a multi-report + directory. A bare directory holding more than one series is ambiguous and + raises — load one via its per-report handle from :func:`list_diagnostics`. + """ + directory, hbase = _resolve_series(handle) + if base is None: + base = hbase + groups = _series_dumps(directory) + if base is not None: + return groups.get(base, []) + if len(groups) > 1: + raise ValueError( + f"{directory} holds multiple report series {sorted(groups)}; " + "load one via its per-report handle from list_diagnostics()." + ) + return next(iter(groups.values()), []) + + +def series_len(source: str | Path) -> int: + """Number of time slices in a diagnostic series, without loading its data. + + ``source`` is any handle accepted by :func:`load_series` (a dump directory, + a per-report handle, or a saved-series ``.nc`` file). Pairs with + ``load_series(..., t_indices=...)`` so callers can pick a few slices of a + long series — e.g. a phase-space history — without ever materializing the + whole ``(t, ...)`` array. + """ + source = Path(source) + if source.is_file() and source.suffix == ".nc": + with xr.open_dataset(source, engine="h5netcdf") as ds: + return int(ds.sizes.get("t", 0)) + return len(_sort_dumps(source)) + + +def _normalize_t_indices(t_indices: Sequence[int], n: int, what: str) -> list[int]: + """Resolve possibly-negative time indices against a series of length ``n``.""" + out = [] + for i in t_indices: + j = int(i) + n if int(i) < 0 else int(i) + if not 0 <= j < n: + raise IndexError(f"t index {i} out of range for {what} with {n} time slices") + out.append(j) + if not out: + raise ValueError(f"t_indices is empty for {what}") + return out + + +def load_series(directory: str | Path, t_indices: Sequence[int] | None = None) -> xr.DataArray: + """Stack every dump in ``directory`` into a ``(t, ...)`` DataArray. + + All files must share the same diagnostic name, the same spatial / + 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. + + ``t_indices`` selects a subset of time slices (negative indices allowed, + order preserved) and reads **only** those from disk — per-dump files for + the HDF5 layout, a partial read for the NetCDF layout. A full phase-space + history can run to tens of GB while its plots need under ten slices, so + slice-selecting callers (see the phase-space paths in ``plots``) is what + keeps concurrent post-processing inside a node's memory. + """ + directory = Path(directory) + if directory.is_file() and directory.suffix == ".nc": + return load_series_nc(directory, t_indices=t_indices) + dumps = _sort_dumps(directory) + if not dumps: + raise FileNotFoundError(f"No .h5 dumps in {directory}") + if t_indices is not None: + idx = _normalize_t_indices(t_indices, len(dumps), str(directory)) + dumps = [dumps[j] for j in idx] + first = load_grid_h5(dumps[0]) + + n = len(dumps) + 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 + groups = _series_dumps(d) + if not groups: + continue + rel = str(d.relative_to(ms)) + if len(groups) == 1: + out[rel] = d + else: + # A directory with several report series (e.g. two s1 line + # lineouts at different perpendicular cells) exposes each as its + # own diagnostic via a synthetic / handle, so they are + # converted to separate NetCDFs instead of merged into one. + for base in sorted(groups): + out[f"{rel}/{base}"] = d / base + # stage_discard_h5 layout: grid dumps were streamed to run_dir/binary + # and never mirrored to MS/ (which then holds only RAW). Merge the + # streamed NetCDFs for any diagnostic with no raw dumps on disk so + # post-processing and plots see the full set; load_series accepts the + # .nc handles directly. + bdir = run_dir / "binary" + if bdir.is_dir(): + for rel, p in list_diagnostics_nc(bdir).items(): + out.setdefault(rel, p) + return out + # No MS/ tree — fall back to the saved-NetCDF layout (a run dir's binary/ + # subdir first, else run_dir itself *being* a binary dir). + bdir = run_dir / "binary" + if bdir.is_dir(): + nc = list_diagnostics_nc(bdir) + if nc: + return nc + nc = list_diagnostics_nc(run_dir) + if nc: + return nc + 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, t_indices: Sequence[int] | None = None) -> xr.DataArray: + """Load a grid/phase-space series NetCDF back into a plotter-ready DataArray. + + This is the inverse of :func:`series_to_dataset` (the form + :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 into memory and closed before returning. With + ``t_indices`` (see :func:`load_series`) the read is *partial*: the file is + opened lazily and only the selected time slices are materialized — the + streamed series are written one dump at a time, so this reads a handful of + chunks instead of the (potentially tens-of-GB) full history. Reductions on + a lazily-opened xarray would silently pull in the whole variable, which is + why the selection happens here, before any consumer touches the data. + """ + if t_indices is None: + ds = xr.load_dataset(path, engine="h5netcdf") + else: + with xr.open_dataset(path, engine="h5netcdf") as f: + idx = _normalize_t_indices(t_indices, int(f.sizes.get("t", 0)), str(path)) + ds = f.isel(t=idx).load() + return _series_da_from_dataset(ds, path) + + +def _series_da_from_dataset(ds: xr.Dataset, path: str | Path) -> xr.DataArray: + """Rebuild a plotter-ready DataArray from an opened saved-series Dataset. + + Shared by the eager :func:`load_series_nc` and the lazy :func:`open_series`: + it only reads the (small) coordinate variables to restore the + ``axis_units`` / ``axis_long_names`` / ``autoscaled_dims`` metadata, and + never touches the (possibly huge) data variable — so on a lazily-opened + ``ds`` the returned array stays lazy and materializes only the slices a + caller actually indexes. + """ + 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, so a + # round-tripped series is indistinguishable (for plotting) from a fresh + # :func:`load_series`. A dim is autoscaled only when its bounds actually + # *move*: the batch path writes {dim}_min/_max solely for moving dims, but + # the streaming writer (:mod:`adept.osiris.stream`) records them for every + # dim, so mere presence is not enough — check that they vary. + autoscaled: list[str] = [] + for d in da.dims: + lo_c, hi_c = f"{d}_min", f"{d}_max" + if lo_c not in ds.coords or hi_c not in ds.coords: + continue + lo = np.asarray(ds[lo_c].values) + hi = np.asarray(ds[hi_c].values) + if lo.size and not (np.allclose(lo, lo[0]) and np.allclose(hi, hi[0])): + autoscaled.append(str(d)) + if autoscaled: + da.attrs["autoscaled_dims"] = autoscaled + return da + + +@contextmanager +def open_series(source: str | Path): + """Yield a ``(t, …)`` series whose slices are read from disk on demand. + + The memory-bounded companion to :func:`load_series` for consumers that walk + or window the time axis and reduce **one dump at a time** (e.g. the SRS + phase-space budgets in ``osiris_lpi.srs_post``, which iterate + ``da.isel(t=it)`` and sum each slice): the full ``(t, …)`` array is never + materialized, so a 48 GB phase-space history costs only one resident dump. + Unlike ``load_series(..., t_indices=)`` this does not need the caller to + know which slices up front, so it also serves the whole-history passes (the + flux-vs-time trace) that a slice list can't express. + + For a saved-series ``.nc`` the file is opened lazily and closed on exit; the + yielded DataArray is valid **only inside the with-block**. For a raw + per-dump directory (no single lazy file) it falls back to an eager + :func:`load_series` and yields that array (nothing to close). + """ + source = Path(source) + if source.is_file() and source.suffix == ".nc": + ds = xr.open_dataset(source, engine="h5netcdf") + try: + yield _series_da_from_dataset(ds, source) + finally: + ds.close() + else: + yield load_series(source) + + +def list_diagnostics_nc(binary_dir: str | Path) -> dict[str, Path]: + """Map every saved-NetCDF diagnostic to its ``.nc`` file. + + 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, + stream: bool = False, + streamed_dir: str | Path | None = None, +) -> list[Path]: + """Convert each diagnostic's full time history to a netCDF file. + + One file per diagnostic is written under ``out_dir``, mirroring the OSIRIS + ``MS/`` layout (e.g. ``out_dir/FLD/e1.nc``, + ``out_dir/PHA/x1p1/beam_pos.nc``). Each file holds the stacked ``(t, ...)`` + series — every time slice OSIRIS dumped for that diagnostic. + + ``diagnostics``, when given, whitelists which diagnostics to convert, + matched against either the relative path (``"FLD/e1"``) or the leaf name + (``"e1"``). Returns the list of written paths. + + Two opt-in paths bound conversion memory to a single dump (see + :mod:`adept.osiris.stream`): + + - ``streamed_dir`` — a directory where the concurrent converter already + wrote ``.nc`` *during the run*. Such a grid diagnostic is copied + into ``out_dir`` rather than rebuilt, so it is never re-read off disk. + - ``stream`` — for any grid diagnostic not already in ``streamed_dir``, + build the NetCDF a dump at a time instead of stacking the whole series in + memory. RAW diagnostics always use the (in-memory) concat path. + """ + out_dir = Path(out_dir) + streamed_dir = Path(streamed_dir) if streamed_dir is not None else None + diags = list_diagnostics(run_dir) + written: list[Path] = [] + for relpath in sorted(diags): + if diagnostics is not None and (relpath not in diagnostics and Path(relpath).name not in diagnostics): + continue + dest = out_dir / f"{relpath}.nc" + try: + is_raw = _diag_is_raw(relpath, diags[relpath]) + pre = (streamed_dir / f"{relpath}.nc") if streamed_dir is not None else None + if pre is not None and not is_raw and pre.is_file(): + # The concurrent converter already produced this during the run. + dest.parent.mkdir(parents=True, exist_ok=True) + if pre.resolve() != dest.resolve(): + shutil.copy(pre, dest) + elif is_raw: + # RAW (particle) dumps: per-particle datasets, no grid/AXIS. + ds: xr.Dataset = load_raw_series(diags[relpath], drop_initial=raw_drop_initial) + dest.parent.mkdir(parents=True, exist_ok=True) + ds.to_netcdf(dest, engine="h5netcdf", encoding=_compression_encoding(ds)) + elif stream: + from adept.osiris import stream as _stream # lazy: pulls in the writer only when needed + + _stream.convert_diagnostic_streaming(diags[relpath], dest, source_dir=diags[relpath]) + else: + ds = series_to_dataset(load_series(diags[relpath])) + dest.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 + if "fld" in f.name.lower(): + # fld_ene columns are the per-component field energies + # (B1 B2 B3 E1 E2 E3); their sum is the total EM field energy. + field = (t, values.sum(axis=1)) + elif "par" in f.name.lower(): + # par_ene columns are "Total Par." (particle COUNT) then + # "Kin. Energy" — take the energy (last) column. Summing would fold + # the ~1e6 particle count into the energy and swamp it. + kinetic[f.stem] = (t, values[:, -1]) + + if field is None and not kinetic: + return None + + 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..3fdf72b2 --- /dev/null +++ b/adept/osiris/plots.py @@ -0,0 +1,1730 @@ +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 + epw_energy_vs_time.png longitudinal (e1) EPW field energy + total_energy_vs_time.png field + kinetic conservation (needs HIST/) + +All axis / colorbar / title labels are emitted as proper LaTeX (``$\omega$``, +``$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 + +# --- render-resolution downsampling --------------------------------------- + +# A heatmap is rasterized to a PNG of a fixed pixel size, so feeding pcolormesh +# a (t, x) array with vastly more cells than on-screen pixels is pure waste: the +# QuadMesh materializes an (rows+1, cols+1) float64 coordinate mesh, a float64 +# value copy, and an RGBA buffer, then Agg rasterizes ~one quad per *cell* — all +# for detail that averages away below a single pixel. Boxcar-averaging the data +# down to roughly the rendered resolution first removes that blow-up (the +# dominant term in the post-processing OOM); the small render arrays are also +# what keep peak memory in the low tens of GB. + +_RENDER_OVERSAMPLE = 3.0 # cells kept per on-screen pixel (margin for zoom / save dpi) + + +def _render_cells( + ax: plt.Axes, + *, + oversample: float = _RENDER_OVERSAMPLE, + visible_frac: tuple[float, float] = (1.0, 1.0), + cap: int = 4000, +) -> tuple[int, int]: + """Target ``(rows, cols)`` to render ``ax`` at, from its on-screen size. + + The axes' size in pixels (inches from the figure size × its subplot fraction, + times the figure dpi) sets the floor; ``oversample`` keeps a few cells per + pixel so the PNG (often saved at a higher dpi) and any later zoom stay sharp, + rather than downsampling to the exact pixel grid with no margin. The data is + only ever decimated to this, never below the eventual pixel count. + + ``visible_frac`` is the fraction of each axis ``(rows, cols)`` that a later + ``set_ylim`` / ``set_xlim`` will actually display (e.g. the equal-aspect + omega-k panel shows only a small Nyquist window): the full array is then kept + at enough cells for that *visible* window to still reach screen resolution. + ``cap`` bounds the request so a degenerate figure can't ask for a huge grid. + """ + fig = ax.figure + dpi = float(fig.dpi) + w_in, h_in = fig.get_size_inches() + try: # shrink to this axes' share of the figure when the layout is known + pos = ax.get_position() + w_in *= float(pos.width) + h_in *= float(pos.height) + except Exception: + pass + fr_rows = max(float(visible_frac[0]), 1e-3) + fr_cols = max(float(visible_frac[1]), 1e-3) + rows = int(np.ceil(h_in * dpi * oversample / fr_rows)) + cols = int(np.ceil(w_in * dpi * oversample / fr_cols)) + return min(max(rows, 1), cap), min(max(cols, 1), cap) + + +def _boxcar_downsample( + data: np.ndarray, + row_coord: np.ndarray, + col_coord: np.ndarray, + target_rows: int, + target_cols: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Boxcar-average a 2D array (and its 1D coords) toward a target shape. + + Reduces by an integer block factor per axis (``size // target``); an array + already at or below the target along an axis is left untouched on that axis, + so low-resolution data is never up-sampled or interpolated. The matching + coordinate centers are averaged over the same blocks, and any trailing cells + that don't fill a whole block are dropped. + """ + rows, cols = data.shape + fr = max(1, rows // max(target_rows, 1)) + fc = max(1, cols // max(target_cols, 1)) + if fr == 1 and fc == 1: + return data, row_coord, col_coord + nr = (rows // fr) * fr + nc = (cols // fc) * fc + reduced = data[:nr, :nc].reshape(nr // fr, fr, nc // fc, fc).mean(axis=(1, 3)) + r = row_coord[:nr].reshape(nr // fr, fr).mean(axis=1) + c = col_coord[:nc].reshape(nc // fc, fc).mean(axis=1) + return reduced, r, c + + +# --- low-level plotters ---------------------------------------------------- + + +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)) + t = da.coords["t"].values + xname = next(d for d in da.dims if d != "t") + x = da.coords[xname].values + # Boxcar-average the (t, x) array down to ~the rendered pixel grid before + # rasterizing, so a 100M+-cell series doesn't build a billion-quad QuadMesh + # for sub-pixel detail. The log view averages |field| (averaging the *signed* + # field first would cancel the oscillations the log plot is meant to show). + target_rows, target_cols = _render_cells(ax) + data = np.abs(da.values) if log else da.values # (t, x) + data, t, x = _boxcar_downsample(data, t, x, target_rows, target_cols) + if log: + data = np.log10(data + 1e-30) + if space_on_x: + mesh = ax.pcolormesh(x, t, data, shading="auto", cmap=cmap) + ax.set_xlabel(_axis_label(da, xname)) + 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 _evolution_t_indices(n: int, n_panels: int) -> list[int]: + """Time indices ``plot_phasespace_evolution`` samples, plus the final slice. + + Mirrors the panel subsampling (``range(0, n, n // n_panels)``) and always + includes ``n - 1``, so a series loaded with exactly these indices serves + both the evolution facets and the final-step heatmap. + """ + t_skip = max(1, n // n_panels) + return sorted(set(range(0, n, t_skip)) | {n - 1}) + + +def plot_phasespace_evolution( + series: xr.DataArray | str | Path, + *, + 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``. + + Given a path, only the sampled snapshots are read from disk + (``io.load_series`` with ``t_indices``) — a production phase-space history + can run to tens of GB while the facets need under a dozen slices. The + colour scale is likewise computed from the sampled panels only, never the + full history. + """ + if isinstance(series, xr.DataArray): + da = series + else: + da = _io.load_series(series, t_indices=_evolution_t_indices(_io.series_len(series), n_panels)) + da = _decorate(da) + if da.ndim != 3: + raise ValueError(f"plot_phasespace_evolution expects (t, p, x); got dims {da.dims}") + da = _crop_spatial_to_box(da) + nt = da.coords["t"].size + # Same sampling rule whether da is a full in-memory series or was already + # slice-selected off disk: stride panels plus, always, the final slice. + idx = _evolution_t_indices(nt, n_panels) + sl = da.isel(t=idx) + # Convention: spatial axis horizontal, momentum axis vertical (fall back to + # dim order for a momentum-momentum space with no spatial axis). + 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 _resolve_fld_diag(diags: dict[str, Path], comp: str) -> str | None: + """Diagnostic key for field component ``comp``: full-grid, else -savg/-tavg. + + A 2D deck that dumps ``e1,savg`` (rather than the full grid) writes it to a + distinct ``FLD/`` diagnostic (e.g. ``FLD/e1-savg``), so ``FLD/e1`` is absent. + Fall back to a spatially/time-averaged variant of the same component — but + never a lineout / slice, which is not a whole ``(t, x…)`` field. + """ + if f"FLD/{comp}" in diags: + return f"FLD/{comp}" + for rel in sorted(diags): + parts = rel.split("/") + if len(parts) == 2 and parts[0] == "FLD": + leaf = parts[1] + if leaf.split("-", 1)[0] == comp and "line" not in leaf and "slice" not in leaf: + return rel + return None + + +def field_energy_components(run_dir: str | Path) -> xr.Dataset: + """E-field, B-field, and total field energy vs time from the FLD diagnostics. + + Like :func:`field_energy_series` but keeps the electric (``e1/e2/e3``) and + magnetic (``b1/b2/b3``) contributions separate. Returns a ``Dataset`` with + ``E_energy``, ``B_energy`` and ``total_field_energy`` on a shared ``t`` axis, + plus the per-component electric energies ``e1_energy`` / ``e2_energy`` / + ``e3_energy`` — ``e1_energy`` is the longitudinal (electrostatic) field + energy, i.e. the electron-plasma-wave (EPW) energy for a run whose density + gradient lies along ``x1``. Components not dumped are omitted from their sum. + + The spatial integral is over *all* spatial dims, so this works for a 1D + ``(t, x)`` field series and a 2D ``(t, x2, x1)`` one alike; 2D decks that dump + only spatially-averaged fields (``e1,savg``) are picked up via + :func:`_resolve_fld_diag`. + + Sources are resolved via :func:`io.list_diagnostics` / :func:`io.load_series`, + so ``run_dir`` may be a raw OSIRIS run directory or the ``binary/`` directory + of saved NetCDFs — the energy is recomputed from the stored field series + either way. + """ + diags = _io.list_diagnostics(run_dir) + comps = ("e1", "e2", "e3", "b1", "b2", "b3") + by_iter: dict[int, dict[str, float]] = {} + times: dict[int, float] = {} + found = False + for comp in comps: + rel = _resolve_fld_diag(diags, comp) + if rel is None: + continue + try: + ser = _io.load_series(diags[rel]) + except Exception as e: # a bad component must not sink the rest + print(f"[plots] skipping field-energy component {comp}: {e}") + continue + if ser.ndim < 2 or "t" not in ser.dims: + continue + found = True + e_t = _field_energy_from_series(ser) + its = np.asarray(ser.coords["iter"].values) if "iter" in ser.coords else np.arange(ser.sizes["t"]) + ts = np.asarray(ser.coords["t"].values, dtype="float64") + for k in range(ts.size): + it = int(its[k]) + rec = by_iter.setdefault(it, {}) + rec[comp] = rec.get(comp, 0.0) + float(e_t[k]) + times.setdefault(it, float(ts[k])) + if not found or not by_iter: + raise RuntimeError(f"No field dumps available under {run_dir}") + iters = sorted(by_iter) + t = np.array([times[i] for i in iters]) + + def col(c: str) -> np.ndarray: + return np.array([by_iter[i].get(c, 0.0) for i in iters]) + + e1a, e2a, e3a = col("e1"), col("e2"), col("e3") + e_arr = e1a + e2a + e3a + b_arr = col("b1") + col("b2") + col("b3") + coords = {"t": t, "iter": ("t", np.asarray(iters))} + ds = xr.Dataset( + { + "E_energy": ("t", e_arr), + "B_energy": ("t", b_arr), + "total_field_energy": ("t", e_arr + b_arr), + "e1_energy": ("t", e1a), + "e2_energy": ("t", e2a), + "e3_energy": ("t", e3a), + }, + coords=coords, + ) + ds["E_energy"].attrs.update(long_name="electric field energy", units="code") + ds["B_energy"].attrs.update(long_name="magnetic field energy", units="code") + ds["total_field_energy"].attrs.update(long_name="total field energy", units="code") + ds["e1_energy"].attrs.update(long_name="longitudinal (EPW) field energy", units="code") + ds["e2_energy"].attrs.update(long_name="transverse E2 field energy", units="code") + ds["e3_energy"].attrs.update(long_name="transverse E3 field energy", units="code") + ds["t"].attrs.update(long_name="time", units=r"1/\omega_p") + return ds + + +def plot_epw_energy( + src: xr.Dataset | str | Path, + ax: plt.Axes | None = None, + *, + log: bool = True, + title: str | None = None, +) -> plt.Axes: + r"""Plot the longitudinal (EPW) field energy ``½ ∫ e1² dV`` vs time. + + The electron-plasma-wave energy trace: the ``e1_energy`` component of + :func:`field_energy_components` (the electrostatic field along ``x1``). ``src`` + may be a precomputed ``Dataset`` or a ``run_dir`` path. + """ + ds = src if isinstance(src, xr.Dataset) else field_energy_components(src) + if "e1_energy" not in ds: + raise RuntimeError("no e1 (longitudinal) field diagnostic for the EPW energy") + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + ax.plot(ds["t"].values, ds["e1_energy"].values) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel(r"EPW energy $\frac{1}{2}\int e_1^2\,d^Dx$ [code units]") + ax.set_title(title or "Total EPW (longitudinal field) energy vs time") + ax.grid(True, which="both", alpha=0.3) + return ax + + +def plot_energy_components( + src: xr.Dataset | str | Path, + ax: plt.Axes | None = None, + *, + 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, alpha=0.5) + 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_energy_partition( + energy: xr.Dataset, + ax: plt.Axes | None = None, + *, + log: bool = True, + title: str | None = None, +) -> plt.Axes: + """Total energy broken down into particle (kinetic) and EM (field) parts. + + Same overlay style as :func:`plot_energy_components`, but the partition is + the physical total energy rather than the E/B field split. ``energy`` is the + ``Dataset`` from :func:`io.load_hist_energy`, which carries ``field_energy`` + (EM), ``kinetic_total`` (all-species particle), and ``total`` (their sum). + + ``total`` is drawn first (underneath) so the two components stay legible + where one of them coincides with the total envelope. + """ + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + t = energy["t"].values + labels = { + "total": "total (particle + EM)", + "kinetic_total": "particle (kinetic)", + "field_energy": "EM (field)", + } + for name, label in labels.items(): + if name in energy: + ax.plot(t, energy[name].values, label=label, alpha=0.5) + if log: + ax.set_yscale("log") + ax.set_xlabel(r"t [$1/\omega_p$]") + ax.set_ylabel("energy [code units]") + ax.set_title(title or "Total energy partition vs time") + ax.legend() + ax.grid(True, which="both", alpha=0.3) + return ax + + +def plot_omega_k( + series: xr.DataArray | str | Path, + ax: plt.Axes | None = None, + *, + 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 + + # Real-input 2-D FFT: rfft over the (long) time axis keeps only omega >= 0 + # and stays in the input's complex width (complex64 for float32) instead of + # upcasting the whole grid to complex128 the way fft2 would. We only want the + # power |F|^2, so the dropped, redundant negative-omega half is rebuilt by + # Hermitian symmetry *after* downsampling. Free the complex array as soon as + # the power is formed, and do the log in place, to bound peak memory. + F = np.fft.rfft2(da.values, axes=(1, 0)) # (nt//2+1, nx): rfft over t, fft over x + P = np.abs(F).astype(np.float32, copy=False) + del F + P **= 2 # |F|^2, in place + P = np.fft.fftshift(P, axes=1) # put k = 0 in the middle of the k (x) axis + if log: + P += 1e-30 + np.log10(P, out=P) + + omega = np.fft.rfftfreq(nt, d=dt) * 2 * np.pi # [0, omega_Nyquist] + k = np.fft.fftshift(np.fft.fftfreq(nx, d=dx)) * 2 * np.pi # [-k_Ny, k_Ny) + + # Resolve the view window now (full Nyquist if unset, from the full-res axes) + # so the data can be decimated to the *visible* resolution: a zoomed panel + # keeps more of the full array than a full-range one (see _render_cells). + k_ny = float(np.max(np.abs(k))) + omega_ny = float(np.max(omega)) + if k_max is None: + k_max = k_ny + if omega_max is None: + omega_max = omega_ny + frac_rows = min(1.0, omega_max / omega_ny) if omega_ny > 0 else 1.0 + frac_cols = min(1.0, k_max / k_ny) if k_ny > 0 else 1.0 + target_rows, target_cols = _render_cells(ax, visible_frac=(frac_rows, frac_cols)) + P, omega, k = _boxcar_downsample(P, omega, k, target_rows, target_cols) + + # Rebuild the omega < 0 half from |F(k, -w)| = |F(-k, w)|: the negative rows + # are the omega > 0 block (excluding DC) flipped along both omega and k. + if P.shape[0] > 1: + P = np.concatenate([P[1:][::-1, ::-1], P], axis=0) + omega = np.concatenate([-omega[1:][::-1], omega]) + + mesh = ax.pcolormesh(k, omega, P, shading="auto", cmap=cmap) + plt.colorbar(mesh, ax=ax, label=(r"$\log_{10}\,|\tilde{F}|^2$" if log else r"$|\tilde{F}|^2$")) + del P + + 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, + avg_window: int | None = None, + 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. + + ``avg_window`` overrides the fraction with an explicit number of trailing + dumps. Pass it when ``series`` was slice-selected out of a longer history + (initial + the mean window), where a fraction of the *received* length no + longer means what the caller intended. + """ + da = _decorate(series if isinstance(series, xr.DataArray) else _io.load_series(series)) + if "t" not in da.dims: + 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 = int(avg_window) if avg_window is not None else max(1, round(n_avg_frac * nt)) + w = max(1, min(w, nt)) + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + final = da.isel(t=-1).values + 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: + # ``values`` is already a fresh array from the (e2 ± b3)/2 arithmetic, so + # relabel it in place rather than holding a second full (t, x) copy. + out = values + out.attrs = dict(src.attrs) + out.attrs["long_name"] = long + out.name = name + 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 transverse_field_boundary_slabs( + run_dir: str | Path, *, guard_cells: int = 1, window_cells: int = 3 +) -> dict | None: + r"""Left/right-going transverse E-field on just the two boundary slabs. + + The slab-restricted companion to :func:`efield_lr_components`. Every + boundary-light diagnostic (laser-energy budget, direction-split spectra and + spectrogram) only samples a thin ``window_cells`` slab ``guard_cells`` in + from each edge, yet the full split holds several whole-grid ``(t, x)`` arrays + at once — prohibitive for the long, wide SRS runs (tens of GiB per field). + This loads each raw transverse field **one at a time** and immediately slices + it to the left slab ``[g, g+w)`` and the symmetric right slab + ``[n-g-w, n-g)``, freeing the whole-grid array before the next, so only the + slab columns are ever combined. + + The vacuum Riemann split is local (per cell), so slicing first and combining + on the slab is identical to combining the whole grid and then slicing. + Returns ``None`` when no transverse pair (``e2/b3`` or ``e3/b2``) was dumped, + else:: + + {"t": (n_t,), "guard_cells": g, "window_cells": w, "pairs": [...], + "edges": {"left": {pair: {"right": (t, w), "left": (t, w)}, ...}, + "right": {pair: {"right": (t, w), "left": (t, w)}, ...}}} + + where each ``(t, w)`` array is the right/left-going part of that polarization + pair on that edge's slab (native field dtype). The consumer forms the + per-channel flux / spectrum from these (the raw, un-split field is + ``right + left``). + """ + diags = _io.list_diagnostics(run_dir) + + def _load(comp: str) -> xr.DataArray | None: + rel = f"FLD/{comp}" + if rel not in diags: + return None + ser = _io.load_series(diags[rel]) + return ser if ser.ndim == 2 else None + + # Resolve the slab columns from the first available transverse field, then + # drop it (each field is reloaded one at a time below so only one whole-grid + # array is ever resident). + probe = None + for comp in ("e2", "b3", "e3", "b2"): + probe = _load(comp) + if probe is not None: + break + if probe is None: + return None + xdim = next(d for d in probe.dims if d != "t") + t = np.asarray(probe.coords["t"].values, dtype=float) + n = int(probe.sizes[xdim]) + g, w = guard_cells, window_cells + if g + w > n: + g, w = 0, min(w, n) + left, right = slice(g, g + w), slice(n - g - w, n - g) + del probe + + def _slabs(comp: str) -> tuple[np.ndarray, np.ndarray] | None: + """Raw field ``comp`` sliced to the (left, right) edge slabs. + + ``.copy()`` so the returned slabs do not keep the whole-grid array alive + as a view base — the (t, x) field is freed when this returns. + """ + da = _load(comp) + if da is None: + return None + lo = da.isel({xdim: left}).values.copy() + hi = da.isel({xdim: right}).values.copy() + return lo, hi + + edges: dict[str, dict[str, dict[str, np.ndarray]]] = {"left": {}, "right": {}} + + def _add_pair(name: str, e: str, b: str, *, right_sign: float) -> None: + # right-going = (e + right_sign*b)/2, left-going = (e - right_sign*b)/2 + # (e2/b3: right_sign=+1; e3/b2: right_sign=-1), matching efield_lr_components. + es, bs = _slabs(e), _slabs(b) + if es is None or bs is None: + return + for edge, ei, bi in (("left", es[0], bs[0]), ("right", es[1], bs[1])): + edges[edge][name] = { + "right": (ei + right_sign * bi) / 2.0, + "left": (ei - right_sign * bi) / 2.0, + } + + _add_pair("e2", "e2", "b3", right_sign=+1.0) + _add_pair("e3", "e3", "b2", right_sign=-1.0) + + pairs = list(edges["left"].keys()) + if not pairs: + return None + return {"t": t, "guard_cells": g, "window_cells": w, "pairs": pairs, "edges": edges} + + +def plot_field_lr_decomposition(run_dir: str | Path) -> dict[str, plt.Figure]: + """One figure per transverse component: right/left-going spacetime. + + 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). + # The fit builds several full-size temporaries of the (t, p, x) + # series, so read only the slices plot_profile uses — t=0 plus the + # trailing mean window — instead of the whole history. + ps_path = _species_phasespace(diags, species) + temp = avg_window = None + if ps_path is not None: + n_t = _io.series_len(ps_path) + if n_t > 0: + avg_window = max(1, round(0.2 * n_t)) # plot_profile's default n_avg_frac + idx = sorted({0, *range(n_t - avg_window, n_t)}) + temp = _temperature_from_phasespace(_io.load_series(ps_path, t_indices=idx)) + if temp is None: + temp = _temperature_series(entries) + avg_window = None # full series: the default fraction applies + if temp is not None: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_profile( + temp, + ax=ax, + avg_window=avg_window, + show_initial=True, + title=f"{label} — temperature profile", + ) + 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: + # Slice-selected read: only the evolution panels + the final step + # come off disk. A production phase-space history (e.g. x1g_q1 at + # ~12k dumps) is tens of GB; the plots below use under a dozen + # slices of it, and eager-loading it here is what blew node memory + # when many runs post-processed at once. + n_t = _io.series_len(diag_path) + if n_t < 1: + continue + ser = _io.load_series(diag_path, t_indices=_evolution_t_indices(n_t, n_panels)) + except Exception as e: + print(f"[plots] skipping phasespace {diag_rel}: {e}") + continue + # 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: # total EPW (longitudinal e1) field energy vs time + fig, ax = plt.subplots(figsize=(6, 4)) + plot_epw_energy(run_dir, ax=ax) + written["epw_energy_vs_time"] = _write(fig, "epw_energy_vs_time.png") + except (FileNotFoundError, RuntimeError) as e: + print(f"[plots] skipping epw_energy_vs_time: {e}") + + try: + energy = _io.load_hist_energy(run_dir) + if energy is not None and "total" in energy: + 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}") + + try: + energy = _io.load_hist_energy(run_dir) + if energy is not None and "total" in energy: + fig, ax = plt.subplots(figsize=(6, 4)) + plot_energy_partition(energy, ax=ax) + written["energy_partition_vs_time"] = _write(fig, "energy_partition_vs_time.png") + except Exception as e: + print(f"[plots] skipping energy_partition_vs_time: {e}") + + return written + + +# --- 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). + + The copy is **shallow** (``deep=False``): xarray still gives the copy its own + attrs dicts (so the labels below don't leak onto the caller's array), but the + data buffer is shared — a deep copy would pull a lazily-opened + (:func:`io.open_series`) multi-GB phase-space history fully into memory just + to relabel its axes, which is exactly the load the lazy path exists to avoid. + """ + da = da.copy(deep=False) + 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 +value_label = _value_label +tex = _tex +sim_box_xmax = _sim_box_xmax +sim_box_bound = _sim_box_bound +# Shared memory-bounded rendering helpers (the (k, ω) map decimates with these +# before pcolormesh; downstream spectrograms reuse them to do the same). +render_cells = _render_cells +boxcar_downsample = _boxcar_downsample diff --git a/adept/osiris/post.py b/adept/osiris/post.py new file mode 100644 index 00000000..cda548f6 --- /dev/null +++ b/adept/osiris/post.py @@ -0,0 +1,237 @@ +"""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) + + +_FIELD_COMPONENTS = ("e1", "e2", "e3", "b1", "b2", "b3") + + +def _is_field_component_dir(name: str) -> bool: + """True for a full-grid or spatially/time-averaged field-component dump dir. + + Matches ``e1`` .. ``b3`` and their ``-savg`` / ``-tavg`` variants (2D decks + dump ``e1,savg`` rather than the full grid), but NOT Poynting-flux (``s1``) + or lineout / slice diagnostics, which also live under ``FLD/`` and would + otherwise be summed into — and inflate — the total field energy. + """ + base = name.split("-", 1)[0] + return base in _FIELD_COMPONENTS and "line" not in name and "slice" not in name + + +def _field_energy_from_dump(h5_path: Path) -> float: + """Integrate ``field^2 * cell_volume`` for a single dump file. + + 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 usually matches the diagnostic name (e.g. "e1"), but + # spatially/time-averaged reports may name it differently; fall back to + # the single non-AXIS/SIMULATION dataset the way load_grid_h5 does. + name = h5_path.stem.rsplit("-", 1)[0] + if name not in f: + data_keys = [k for k in f.keys() if k not in ("AXIS", "SIMULATION")] + if len(data_keys) != 1: + return float("nan") + name = data_keys[0] + arr = f[name][...] + # Reconstruct cell volume from SIMULATION attrs. + sim = f["SIMULATION"] + 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() or not _is_field_component_dir(comp_dir.name): + 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. + + Under ``stage_discard_h5`` the FLD h5 never reach the durable ``MS/``; + fall back to scanning the whole tree (RAW dumps still carry iteration + numbers, at their coarser cadence).""" + best = -1 + for scan_dir in (ms / "FLD", ms): + if not scan_dir.is_dir(): + continue + for p in scan_dir.rglob("*.h5"): + m = _ITER_RE.search(p.name) + if m: + best = max(best, int(m.group(1))) + if best >= 0: + return best + return best + + +def _total_field_energy_nc(binary_dir: Path) -> float: + """|E|^2/2 + |B|^2/2 at the last streamed slice of ``binary_dir/FLD/*.nc``. + + Fallback for the ``stage_discard_h5`` layout, where no FLD h5 exists on + the durable ``MS/``. Returns NaN when no usable component is found.""" + import xarray as xr + + fld = binary_dir / "FLD" + if not fld.is_dir(): + return float("nan") + total, found = 0.0, False + for p in sorted(fld.glob("*.nc")): + if not _is_field_component_dir(p.stem): + continue + try: + with xr.open_dataset(p) as ds: + da = ds[p.stem] if p.stem in ds else ds[next(iter(ds.data_vars))] + last = da.isel(t=-1).astype("float64") + dvol = 1.0 + for dim in last.dims: + c = last[dim].values + if len(c) > 1: + dvol *= float(c[1] - c[0]) + total += 0.5 * float((last.values**2).sum()) * dvol + found = True + except Exception: + continue + return total if found else float("nan") + + +def collect(run_output: dict, cfg: dict, td: str) -> dict[str, Any]: + """Post-process a finished OSIRIS run. + + 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)) + # When the concurrent converter ran, reuse the NetCDFs it already produced + # (and stream-build any grid diagnostic it did not reach) instead of the + # in-memory batch conversion. + stream_on = bool((cfg.get("osiris") or {}).get("stream_convert", True)) + streamed_dir = solver.get("binary_dir") if stream_on else None + + metrics: dict[str, float] = { + "wall_time_s": float(solver["wall_time"]), + "exit_code": float(solver["exit_code"]), + "field_energy_final": _total_field_energy(ms), + "final_iter": float(_final_iter(ms)), + } + if metrics["field_energy_final"] != metrics["field_energy_final"] and streamed_dir is not None: + # stage_discard_h5: no FLD h5 on the durable MS/ — use the streamed nc. + metrics["field_energy_final"] = _total_field_energy_nc(Path(streamed_dir)) + + # Convert each diagnostic's time history to an xarray netCDF. Under + # stage_discard_h5 with no RAW diagnostics MS/ may not exist at all while + # the streamed NetCDFs do — proceed in that case too. + if ms.is_dir() or (streamed_dir is not None and Path(streamed_dir).is_dir()): + _io.save_run_datasets( + run_dir, + td / "binary", + diagnostics=whitelist, + raw_drop_initial=raw_drop_initial, + stream=stream_on, + streamed_dir=streamed_dir, + ) + + # plots imports matplotlib; do it lazily to keep `import adept.osiris` light. + 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..8edf7fc8 --- /dev/null +++ b/adept/osiris/runner.py @@ -0,0 +1,362 @@ +"""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 + # srun launcher warning, not an OSIRIS error: when a task's working directory + # (e.g. a per-node /dev/shm staging dir) is not resolvable at launch, srun + # prints "error: couldn't chdir to `...`: ... going to /tmp instead". On + # Perlmutter this fires spuriously even when the tasks do chdir correctly, so + # it must not abort an otherwise-clean run. A genuine staging failure surfaces + # downstream instead (no dumps to drain -> empty/short post-processing). + "couldn't chdir to", +) + + +def _looks_like_osiris_error(line: str) -> bool: + low = line.strip().lower() + if not low: + return False + # OSIRIS prefixes diagnostics with (*warning*) / (*error*); a warning is + # never a failure even when its text contains the word "error" (e.g. the + # Sentoku-collisions "factor of 2 error" warning). + if "(*warning*)" in low: + return False + if any(noise.lower() in low for noise in _OSIRIS_STDERR_NOISE): + return False + return any(tok in low for tok in _OSIRIS_ERR_TOKENS) + + +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 _sync_tree(src: Path, dst: Path) -> None: + """Copy every file under ``src`` to the same relative path under ``dst``, + skipping any that already exist. + + Used at the end of a staged run to fold whatever the background drainer did + not move (``HIST/``, ``RE/``, ``TIMINGS/``, any straggler ``MS/`` dump) from + the ephemeral scratch into the durable run directory before the scratch is + reclaimed. Existing targets (the streamed ``binary/`` NetCDFs, the dumps the + drainer already mirrored) are left untouched. + """ + for p in src.rglob("*"): + if not p.is_file(): + continue + target = dst / p.relative_to(src) + if target.exists(): + continue + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(p, target) + + +def _run_produced_output(run_dir: Path, binary_dir: Path) -> bool: + """True if the run wrote any salvageable output (raw dumps, NetCDFs, HIST). + + Used to decide whether a non-zero exit / OSIRIS error is a hard failure or + merely a crash *after* data was already written — e.g. OSIRIS segfaults on + MPI/CUDA teardown after printing "Simulation completed", or dies partway with + dumps on disk. In the latter case downstream consolidation + plotting should + still run on what exists rather than discarding the whole run. + """ + ms = run_dir / "MS" + if ms.is_dir() and next(ms.rglob("*.h5"), None) is not None: + return True + if binary_dir.is_dir() and next(binary_dir.rglob("*.nc"), None) is not None: + return True + hist = run_dir / "HIST" + if hist.is_dir() and next(hist.rglob("*"), None) is not None: + return True + return False + + +def run_osiris( + deck_text: str, + *, + 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, + stream_convert: bool = True, + stream_poll_s: float = 10.0, + stage_root: str | Path | None = None, + stage_discard_h5: bool = False, +) -> dict[str, Any]: + """Run OSIRIS and return run metadata. + + Returns a dict with keys ``run_dir`` (Path), ``exit_code`` (int), + ``wall_time`` (float, seconds), and ``cmd`` (list[str]). When + ``stream_convert`` is set it also returns ``binary_dir`` (str, the directory + the concurrent converter wrote into) and ``streamed_diagnostics`` + (sorted list[str] of relpaths fully converted while OSIRIS ran). + + With ``stream_convert=True`` a best-effort background thread + (:class:`adept.osiris.stream.StreamConverter`) converts grid diagnostics to + NetCDF *concurrently* with the run, polling every ``stream_poll_s`` seconds, + so the conversion overlaps compute and dumps are read while still warm in + cache. It is fully isolated: any failure there is logged and leaves the run + (and the batch post-processing fallback) untouched. + + **Ramdisk staging (``stage_root``).** When ``stage_root`` is set (and + ``stream_convert`` is on), OSIRIS runs with its working directory on that + fast ephemeral filesystem (e.g. ``/dev/shm``) so its *synchronous* dump + writes hit RAM instead of stalling on the parallel filesystem — effectively + the asynchronous diagnostic I/O OSIRIS itself lacks. The background drainer + mirrors each completed dump to the durable run directory under ``run_root`` + and deletes the scratch copy, keeping the ramdisk footprint bounded by the + poll interval rather than the whole run. ``run_dir`` in the returned dict is + always the durable directory, so post-processing is unchanged; the scratch + is reclaimed before returning. Single-node only (``/dev/shm`` is per-node). + + **``stage_discard_h5``** (staging only): grid dumps are *not* mirrored to + the durable ``MS/`` tree -- once appended to the streamed NetCDF the scratch + HDF5 is simply deleted. Cuts the persist-filesystem inode count by ~the + grid-dump count. RAW dumps are still mirrored (the batch path needs them), + and any dump whose stream failed is left for the final sync, so partial + failures remain recoverable. The durable ``MS/`` tree is then *incomplete + by design*: offline re-conversion of grid diagnostics is impossible, the + streamed ``binary/*.nc`` are the only copy. + + Raises ``RuntimeError`` on non-zero exit code, with the last lines of + stderr included in the message. + """ + binary = Path(binary).expanduser().resolve() + if not binary.exists(): + raise FileNotFoundError(f"OSIRIS binary not found: {binary}") + + # The durable run directory always lives under run_root. With staging, OSIRIS + # works on a same-named scratch dir under stage_root and the drainer mirrors + # back here; without it, OSIRIS works in run_dir directly (the original path). + run_dir = _make_run_dir(Path(run_root).expanduser().resolve()) + staging = bool(stage_root) and stream_convert + if stage_root and not stream_convert: + print("[stream] stage_root ignored: ramdisk staging requires stream_convert=True") + if staging: + stage_dir: Path | None = Path(stage_root).expanduser().resolve() / run_dir.name + stage_dir.mkdir(parents=True, exist_ok=True) + work_dir = stage_dir + # Keep a durable copy of the deck alongside the mirrored outputs. + (run_dir / "os-stdin").write_text(deck_text) + else: + stage_dir = None + work_dir = run_dir + (work_dir / "os-stdin").write_text(deck_text) + + if mpi_ranks > 1: + cmd = [launcher, "-n", str(mpi_ranks)] + 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] = [] + + # Best-effort concurrent H5 -> NetCDF converter, spawned alongside OSIRIS. + # It reads OSIRIS's dumps from the working dir and writes NetCDFs into the + # durable binary/ dir; in staging mode it also mirrors+reaps the scratch + # into run_dir (persist_dir). + converter = None + binary_dir = run_dir / "binary" + if stream_convert: + try: + from adept.osiris import stream as _stream + + converter = _stream.StreamConverter( + work_dir, + binary_dir, + poll_s=stream_poll_s, + persist_dir=run_dir if staging else None, + discard_grid_h5=staging and stage_discard_h5, + ) + except Exception as e: # never let the converter block the run + print(f"[stream] disabled (init failed): {e}") + converter = None + + t0 = time.time() + proc = subprocess.Popen( + cmd, + cwd=work_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() + if converter is not None: + converter.start() + rc = proc.wait() + t_out.join() + t_err.join() + wall_time = time.time() - t0 + + # Finalize the converter before any error handling below: the run is over, + # so flush whatever the watcher had not yet reached and close the files. + streamed: list[str] = [] + if converter is not None: + try: + streamed = sorted(converter.finalize()) + except Exception as e: + print(f"[stream] finalize failed (batch fallback will cover it): {e}") + + # Staging: fold anything the drainer did not move (HIST/, RE/, straggler MS/ + # dumps) from the scratch into the durable run dir, then reclaim the ramdisk. + # Done before error handling so a failed run's partial outputs still persist. + if staging and stage_dir is not None: + try: + _sync_tree(stage_dir, run_dir) + except Exception as e: + print(f"[stream] final stage->persist sync issue (partial outputs may be lost): {e}") + shutil.rmtree(stage_dir, ignore_errors=True) + + # Classify the outcome. A run that crashed (non-zero exit) or that OSIRIS + # flagged as an error can still have produced usable output — most commonly a + # segfault on MPI/CUDA teardown *after* "Simulation completed", but also a + # mid-run death with dumps already on disk. In that case we salvage: log the + # error and fall through so the caller still consolidates binaries and plots + # what was written. Only a run that produced NOTHING is a hard failure. + failure_msg: str | None = None + if rc != 0: + tail = "".join(stderr_tail[-50:]) or "(empty stderr)" + failure_msg = ( + f"OSIRIS exited with status {rc}.\n cmd: {shlex.join(cmd)}\n" + f" cwd: {work_dir}\n stderr tail:\n{tail}" + ) + else: + # OSIRIS can exit 0 even on input-file errors: it prints something + # like 'Error reading ... / aborting...' to stderr (and '(*error*)' + # to stdout) before terminating. Detect both. + err_lines = [ln for ln in stderr_tail if _looks_like_osiris_error(ln)] + err_lines += [ln for ln in stdout_tail if "(*error*)" in ln] + if err_lines: + out_tail = "".join(stdout_tail[-20:]) + err_tail = "".join(stderr_tail[-20:]) + failure_msg = ( + "OSIRIS reported an error despite exit-code 0:\n" + f" cmd: {shlex.join(cmd)}\n" + f" cwd: {work_dir}\n" + f" stderr tail:\n{err_tail}\n" + f" stdout tail:\n{out_tail}" + ) + + crashed = failure_msg is not None + if crashed: + if _run_produced_output(run_dir, binary_dir): + print(f"[osiris] WARNING: {failure_msg}") + print( + "[osiris] the run produced output files despite the error above — " + "consolidating binaries and generating plots from the (possibly " + "partial) data anyway. Verify results carefully." + ) + else: + raise RuntimeError( + f"{failure_msg}\n (no output files were written — nothing to salvage.)" + ) + + result: dict[str, Any] = { + "run_dir": run_dir, + "exit_code": rc, + "crashed": crashed, + "wall_time": wall_time, + "cmd": cmd, + } + if stream_convert: + result["binary_dir"] = str(binary_dir) + result["streamed_diagnostics"] = streamed + if staging: + result["staged"] = True + return result + + +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/adept/osiris/stream.py b/adept/osiris/stream.py new file mode 100644 index 00000000..1ad00199 --- /dev/null +++ b/adept/osiris/stream.py @@ -0,0 +1,723 @@ +"""Incremental / concurrent HDF5 -> NetCDF conversion for OSIRIS runs. + +The batch converter in :mod:`adept.osiris.io` (:func:`io.save_run_datasets`) +reads a diagnostic's *entire* ``(t, ...)`` time history into memory and writes +it in one :func:`to_netcdf`. For the SRS field runs that is two problems: it +peaks at the whole stacked series in RAM, and it re-reads ~700k tiny HDF5 files +cold off Lustre at job end (see ``osiris-lpi/postproc-performance.md``). + +This module rearranges the same per-dump primitives into a *read-one -> +write-one-slot* loop: + +- **Stage A** (:func:`convert_diagnostic_streaming`): run the slot-writer at job + end. Peak memory drops from the full series to a single dump. No concurrency, + no write-completeness race. +- **Stage B** (:class:`StreamConverter`): a best-effort background thread that + drains new dumps into the NetCDF *while OSIRIS is still running*, so the + conversion overlaps the (hours-long) compute and each dump is read while still + warm in the node page cache. A final authoritative sweep at job end fills + whatever the watcher had not yet reached. + +Both write the *same* on-disk schema as :func:`io.series_to_dataset`, so the +``binary/*.nc`` contract is unchanged and every downstream reader +(:func:`io.load_series_nc`, the plotting code) works without modification. + +Design notes +------------ +- **Unlimited ``t`` dimension, append in iteration order.** OSIRIS writes dumps + sequentially and we process them sorted by iteration, so appending in that + order is correct. The dimension is grown with ``resize_dimension`` so it ends + up sized to *exactly* the dumps produced — no pre-sized guess from the deck + (whose off-by-one is genuinely ambiguous) and no trailing fill slots to trim + on early termination. +- **Batched, chunk-aligned writes.** :meth:`StreamWriter.append` buffers rows + in memory and lands them in chunk-sized batches (one HDF5 slice write and one + ``resize_dimension`` per batch). Appending row-by-row into a gzip chunk of + ``chunk_t`` rows costs a decompress+recompress of the whole ~1 MiB chunk *per + row* — ~40x write amplification at the every-3-step field cadence, which is + what let the drainer fall behind OSIRIS and fill the staging ramdisk in job + 56005549 (see ``osiris-lpi/dev_docs/stream-drainer-recommendations.md``). + The buffer is bounded by one chunk (~1 MiB); :meth:`StreamWriter.flush` + (called once per watcher scan) spills any partial batch, so the on-disk file + trails the consumed dumps by at most one poll interval. Threads are *not* + used to parallelize conversion: h5py serializes every HDF5 call behind a + process-global lock, so batching + cheap compression is where the throughput + is. +- **Iteration-based bookkeeping.** The writer tracks the last appended OSIRIS + iteration (persisted in the ``iter`` coordinate, so restarts recover it) and + the drain loops skip any dump at or below it. Unlike positional slot + arithmetic this survives corrupt dumps being removed from the listing. +- **Corrupt dumps are quarantined, not fatal.** A dump that cannot be read + (e.g. truncated by ENOSPC) is renamed to ``*.h5.bad`` — invisible to the + ``*.h5`` listings — and the stream continues with the next dump. Previously a + single bad dump dropped the whole writer and the watcher re-hit the same file + every poll, forever. +- **Backlog spill valve (staging mode).** The ramdisk footprint is only + "bounded by the poll interval" while conversion keeps up with production. + When a diagnostic's pending backlog exceeds ``spill_backlog_files`` / + ``spill_backlog_bytes`` (or the staging filesystem drops below + ``floor_free_bytes``), the converter stops converting it and *mirrors* its + dumps straight to the persist ``MS/`` tree instead (a plain copy is ~10x + cheaper than convert+compress), so the ramdisk keeps draining no matter what. + A spilled diagnostic stays mirror-only for the rest of the run (resuming the + stream mid-run would leave an iteration gap in the NetCDF) and is caught up + from the mirrored HDF5 during :meth:`StreamConverter.finalize`. OSIRIS must + never see ENOSPC on a dump: in 56005549 that permanently corrupted its + diagnostic output for the rest of the job. +- **Write-completeness race.** A dump is only safe to read once OSIRIS has moved + on to the next one, so the watcher processes every dump *except the current + highest-numbered one*; the final sweep (run is dead) picks up that last dump. +- **Per-dump axis bounds** (``{dim}_min`` / ``{dim}_max``) are always recorded, + so OSIRIS autoscale (``if_ps_p_auto``) survives and a restart can recover them + from the file. :func:`io.load_series_nc` only *flags* a dim as autoscaled when + those bounds actually move, so the constant bounds carried for a fixed spatial + axis are harmless. +- **Failure isolation.** Nothing here may abort the OSIRIS run: the watcher + swallows and logs every exception (rate-limited — repeats of the same error + log once per ``error_log_every`` occurrences), and the batch path in + :func:`io.save_run_datasets` remains the safety net for anything not produced. +- **Observability.** Every ``stats_every_s`` the watcher logs one line with the + interval's streamed/spilled/backlog/quarantined counts, so a smoke can assert + the steady-state backlog is flat *while* OSIRIS runs. Checking the ramdisk + after the job proves nothing — the drainer always catches up once production + stops. +""" + +from __future__ import annotations + +import os +import shutil +import threading +import time +from pathlib import Path + +import h5netcdf +import numpy as np +import xarray as xr + +from adept.osiris import io as _io + +# Target uncompressed bytes per (t-chunk x spatial) chunk. A modest t-chunk +# keeps compression/read performance close to the batch path; batched appends +# (see StreamWriter) amortize the write cost of a chunk over its rows. ~1 MiB +# is a reasonable HDF5 chunk size. +_TARGET_CHUNK_BYTES = 1 << 20 + + +def _default_chunk_t(spatial_shape: tuple[int, ...]) -> int: + """Pick a ``t``-chunk so each chunk is roughly :data:`_TARGET_CHUNK_BYTES`.""" + per_row = max(1, int(np.prod(spatial_shape))) * np.dtype(_io._DIAG_DTYPE).itemsize + return int(max(1, min(512, _TARGET_CHUNK_BYTES // per_row))) + + +def _data_var_attrs(template: xr.DataArray, source_dir) -> dict: + """Attrs for the diagnostic variable, matching :func:`io.series_to_dataset`. + + Starts from a single-dump :func:`io.load_grid_h5` DataArray and reproduces + exactly the attrs the batch path leaves on the data variable: the per-dump + ``time`` / ``iter`` / ``source`` are dropped, the dict-valued axis metadata + is lifted onto the coordinate variables instead, ``source_dir`` is added, and + everything left is coerced to a netCDF-writable scalar/array. + """ + attrs = dict(template.attrs) + for k in ("time", "iter", "source", "axis_units", "axis_long_names"): + attrs.pop(k, None) + attrs["source_dir"] = str(source_dir) + return {k: _io._coerce_attr(v) for k, v in attrs.items() if v is not None} + + +class StreamWriter: + """Append-only writer for one diagnostic's ``(t, ...)`` NetCDF time series. + + Created from the diagnostic's first dump (the schema template); each + :meth:`append` buffers the next time slot, and slots are landed on disk in + chunk-sized batches (or on :meth:`flush` / :meth:`close`), growing the + unlimited ``t`` dimension once per batch. Reopening an existing file + (``mode="a"``) resumes after the slots already on disk, which is what makes + a restart idempotent; ``last_iter`` (recovered from the ``iter`` coordinate + on resume) lets callers skip dumps already in the file. ``n_written`` + counts consumed dumps — on-disk slots plus the pending batch. The file is + held open for the lifetime of the writer; call :meth:`close` to finalize. + + The default ``complevel`` is 1: at production dump cadences the drainer is + compression-bound, and gzip-1 is 2-4x cheaper than the old gzip-4 for a + modest size cost (offline postproc can always recompress). + """ + + def __init__(self, dest, template: xr.DataArray, *, source_dir=None, chunk_t: int | None = None, complevel: int = 1): + self.dest = Path(dest) + self.name = str(template.name) + self.spatial_dims = [str(d) for d in template.dims] + self.spatial_shape = tuple(int(s) for s in template.shape) + self._f: h5netcdf.File | None = None + self._ct = int(chunk_t or _default_chunk_t(self.spatial_shape)) + # The pending batch: rows consumed by append() but not yet on disk. + self._rows: list[np.ndarray] = [] + self._ts: list[float] = [] + self._its: list[int] = [] + self._bounds: dict[str, list[tuple[float, float]]] = {d: [] for d in self.spatial_dims} + + if self.dest.exists(): + # Resume: trust the existing schema, continue after written slots. + self._f = h5netcdf.File(self.dest, "a") + self._n_disk = int(self._f.dimensions["t"].size) + self.last_iter = int(self._f.variables["iter"][self._n_disk - 1]) if self._n_disk else -1 + else: + self.dest.parent.mkdir(parents=True, exist_ok=True) + self._f = h5netcdf.File(self.dest, "w") + self._create_schema(template, source_dir if source_dir is not None else self.dest.parent, self._ct, complevel) + self._n_disk = 0 + self.last_iter = -1 + + @property + def n_written(self) -> int: + return self._n_disk + len(self._rows) + + def _create_schema(self, template, source_dir, chunk_t, complevel) -> None: + f = self._f + axis_units = template.attrs.get("axis_units", {}) or {} + axis_long = template.attrs.get("axis_long_names", {}) or {} + + f.dimensions["t"] = None # unlimited; grown one batch per landing + for d, sz in zip(self.spatial_dims, self.spatial_shape): + f.dimensions[d] = sz + + var = f.create_variable( + self.name, + ("t", *self.spatial_dims), + dtype=np.dtype(_io._DIAG_DTYPE), + chunks=(chunk_t, *self.spatial_shape), + compression="gzip", + compression_opts=complevel, + shuffle=True, + ) + for k, v in _data_var_attrs(template, source_dir).items(): + var.attrs[k] = v + # List the non-dimension coordinate variables so xarray promotes them to + # coordinates (not data variables) on read — the CF convention the batch + # path's `to_netcdf` also emits. + nondim = ["iter"] + [f"{d}_{b}" for d in self.spatial_dims for b in ("min", "max")] + var.attrs["coordinates"] = " ".join(nondim) + + f.create_variable("t", ("t",), dtype="f8") + f.create_variable("iter", ("t",), dtype="i8") + for d in self.spatial_dims: + cv = f.create_variable(d, (d,), dtype="f8") + cv[:] = np.asarray(template.coords[d].values, dtype="f8") + if axis_units.get(d): + cv.attrs["units"] = axis_units[d] + if axis_long.get(d): + cv.attrs["long_name"] = axis_long[d] + f.create_variable(f"{d}_min", ("t",), dtype="f8") + f.create_variable(f"{d}_max", ("t",), dtype="f8") + + def append(self, da: xr.DataArray) -> None: + """Buffer ``da`` (one dump) as the next time slot; lands in batches.""" + if da.shape != self.spatial_shape: + raise ValueError(f"shape mismatch for {self.name}: {da.shape} != {self.spatial_shape}") + self._rows.append(np.asarray(da.values, dtype=_io._DIAG_DTYPE)) + self._ts.append(float(da.attrs.get("time", np.nan))) + it = int(da.attrs.get("iter", -1)) + self._its.append(it) + if it > self.last_iter: + self.last_iter = it + for d in self.spatial_dims: + cv = np.asarray(da.coords[d].values) + self._bounds[d].append((float(cv[0]), float(cv[-1])) if cv.size else (np.nan, np.nan)) + if len(self._rows) >= self._ct: + self._write_batch() + + def _write_batch(self) -> None: + """Land the pending rows: one resize + one slice write per variable.""" + k = len(self._rows) + if not k or self._f is None: + return + f = self._f + i = self._n_disk + f.resize_dimension("t", i + k) + f.variables[self.name][i : i + k, ...] = np.stack(self._rows) + f.variables["t"][i : i + k] = np.asarray(self._ts, dtype="f8") + f.variables["iter"][i : i + k] = np.asarray(self._its, dtype="i8") + for d in self.spatial_dims: + b = np.asarray(self._bounds[d], dtype="f8") + f.variables[f"{d}_min"][i : i + k] = b[:, 0] + f.variables[f"{d}_max"][i : i + k] = b[:, 1] + self._n_disk = i + k + self._rows.clear() + self._ts.clear() + self._its.clear() + for d in self.spatial_dims: + self._bounds[d].clear() + + def flush(self) -> None: + if self._f is not None: + self._write_batch() + self._f.flush() + + def close(self) -> None: + if self._f is not None: + self._write_batch() + self._f.close() + self._f = None + + +def convert_diagnostic_streaming(diag_dir, dest, *, source_dir=None) -> Path: + """Stage A: convert one grid diagnostic to NetCDF a dump at a time. + + A memory-bounded drop-in for ``series_to_dataset(load_series(diag_dir))`` + followed by ``to_netcdf`` — peak RAM is one dump plus one write batch + (~1 MiB), not the whole stacked series. Rebuilds ``dest`` from scratch (any + partial file is removed first). + """ + diag_dir = Path(diag_dir) + dest = Path(dest) + dumps = _io._sort_dumps(diag_dir) + if not dumps: + raise FileNotFoundError(f"No .h5 dumps in {diag_dir}") + if dest.exists(): + dest.unlink() + + template = _io.load_grid_h5(dumps[0]) + writer = StreamWriter(dest, template, source_dir=source_dir or diag_dir) + try: + writer.append(template) + for p in dumps[1:]: + writer.append(_io.load_grid_h5(p)) + finally: + writer.close() + return dest + + +class StreamConverter: + """Stage B: a background thread that drains diagnostics during the run. + + Spawned by :func:`adept.osiris.runner.run_osiris` alongside the OSIRIS + subprocess, it polls ``run_dir/MS`` and drains each grid diagnostic's + completed dumps into ``out_dir/.nc``. RAW (particle) diagnostics do + not fit the fixed-slot model (variable particle count per dump) and are left + to the batch path. Every operation is best-effort: any error is logged + (rate-limited) and the run is never affected. + + **Staging mode (``persist_dir`` set).** When OSIRIS writes its ``MS/`` dumps + to a fast ephemeral scratch (e.g. a ``/dev/shm`` ramdisk passed as + ``run_dir``), the converter also *drains* the scratch: after each completed + dump is handled it is mirrored to ``persist_dir/MS//`` and + **deleted from the scratch**, so the ramdisk high-water mark is bounded by + the poll interval rather than the whole run. Grid diagnostics are mirrored + *and* streamed to ``out_dir`` (which the caller points at durable storage); + RAW diagnostics, which cannot be slot-streamed, are mirrored as HDF5 so the + batch post-processing path finds them on ``persist_dir``. The net durable + layout under ``persist_dir`` is identical to a non-staged run, so + post-processing reads it unchanged. With ``persist_dir=None`` nothing is + mirrored or deleted — the converter only streams grid diagnostics in place, + exactly as before. + + **Backlog spill valve (staging mode).** The bound above only holds while + conversion keeps up with OSIRIS. When a diagnostic's pending backlog + exceeds ``spill_backlog_files`` or ``spill_backlog_bytes`` — or the staging + filesystem's free space drops below ``floor_free_bytes`` — that diagnostic + switches to mirror-only draining (plain copy to ``persist_dir/MS/``, ~10x + cheaper than convert+compress), so the ramdisk keeps draining no matter + what. A spilled diagnostic stays mirror-only for the rest of the run + (resuming the stream mid-run would leave an iteration gap in the NetCDF) + and its NetCDF is caught up from the mirrored HDF5 in :meth:`finalize`; in + ``discard_grid_h5`` mode the mirrored copies are deleted again once + appended. + """ + + def __init__( + self, + run_dir, + out_dir, + *, + poll_s: float = 10.0, + logger=print, + persist_dir=None, + discard_grid_h5: bool = False, + spill_backlog_files: int = 20_000, + spill_backlog_bytes: int = 4 << 30, + floor_free_bytes: int = 8 << 30, + stats_every_s: float = 60.0, + rediscover_every: int = 12, + error_log_every: int = 200, + ): + self.ms = Path(run_dir) / "MS" + self.out_dir = Path(out_dir) + self.persist_dir = Path(persist_dir) if persist_dir is not None else None + self.persist_ms = (self.persist_dir / "MS") if self.persist_dir is not None else None + # Staging only: when set, grid dumps are deleted from the scratch after + # being appended to the NetCDF *without* being mirrored to persist_dir/MS + # (the .nc is the durable artifact). Saves one inode+copy per dump on the + # persist filesystem. RAW dumps are always mirrored (the batch path needs + # the HDF5), and so are spilled grid dumps (their bytes exist nowhere + # else until the finalize catch-up appends them). + self.discard_grid_h5 = bool(discard_grid_h5) + self.poll_s = float(poll_s) + self.spill_backlog_files = int(spill_backlog_files or 0) + self.spill_backlog_bytes = int(spill_backlog_bytes or 0) + self.floor_free_bytes = int(floor_free_bytes or 0) + self.stats_every_s = float(stats_every_s) + self.rediscover_every = max(1, int(rediscover_every)) + self.error_log_every = max(1, int(error_log_every)) + self._log = logger + self._writers: dict[str, StreamWriter] = {} + self._completed: set[str] = set() # grid relpaths fully converted to .nc + self._raw_completed: set[str] = set() # raw relpaths fully mirrored to persist + self._is_raw: dict[str, bool] = {} # rel -> raw? (cached; avoids re-peeking h5) + self._multibase_skipped: set[str] = set() # multi-report dirs left to the batch pass + self._spilled: set[str] = set() # grid rels in sticky mirror-only mode + self._bad: set[Path] = set() # unreadable dumps that could not be renamed away + self._err_counts: dict[str, int] = {} + self._known: dict[str, Path] = {} # cached diagnostic discovery + self._scan_i = 0 + self._stats = {"appended": 0, "appended_bytes": 0, "spilled": 0, "spilled_bytes": 0, "quarantined": 0} + self._stats_snapshot = dict(self._stats) + self._stats_t = time.monotonic() + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + # --- lifecycle -------------------------------------------------------- + + def start(self) -> None: + self._thread = threading.Thread(target=self._run, name="osiris-stream-convert", daemon=True) + self._thread.start() + + def _run(self) -> None: + while not self._stop.is_set(): + try: + self._scan_once(final=False) + except Exception as e: # a watcher must never crash the run + self._log_every("scan", f"[stream] scan error (continuing): {e}") + self._stop.wait(self.poll_s) + + def finalize(self) -> set[str]: + """Stop the thread, run the authoritative final sweep, close all writers. + + After OSIRIS has exited every dump is safe to read (no writer can still + be racing), so the final sweep processes the last dump each diagnostic + had been holding back, catches spilled diagnostics up from their persist + mirror, and (in staging mode) drains everything off the scratch before + closing the files. Returns the set of *grid* diagnostic relpaths fully + converted by the stream (so the caller can skip rebuilding them in the + batch pass); RAW relpaths are mirrored but not returned, since the batch + pass still builds their NetCDFs from the mirrored HDF5. + """ + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=max(self.poll_s, 30.0) + 30.0) + try: + self._scan_once(final=True) + except Exception as e: + self._log(f"[stream] finalize sweep error (batch fallback covers it): {e}") + for rel, w in self._writers.items(): + try: + w.close() + except Exception: + pass + self._completed.add(rel) + return set(self._completed) + + # --- logging / stats -------------------------------------------------- + + def _log_every(self, key: str, msg: str) -> None: + """Log ``msg`` on the first occurrence of ``key`` and every Nth after. + + The ENOSPC failure mode repeats the *same* error for every dump on + every poll; unthrottled, that was megabytes of log spam per hour. + """ + n = self._err_counts.get(key, 0) + 1 + self._err_counts[key] = n + if n == 1 or n % self.error_log_every == 0: + self._log(msg + (f" [seen x{n}]" if n > 1 else "")) + + def _maybe_stats(self, backlog_files: int, backlog_bytes: int, *, final: bool) -> None: + """One drain-vs-production line per ``stats_every_s`` (and at finalize).""" + now = time.monotonic() + if not final and (now - self._stats_t) < self.stats_every_s: + return + delta = {k: self._stats[k] - self._stats_snapshot[k] for k in self._stats} + if final or any(delta.values()) or backlog_files: + extra = "" + if self.persist_dir is not None: + try: + extra = f", stage free {shutil.disk_usage(self.ms).free >> 30} GiB" + except OSError: + pass + self._log( + f"[stream] stats({now - self._stats_t:.0f}s): streamed {delta['appended']} dumps " + f"({delta['appended_bytes'] >> 20} MiB), spilled {delta['spilled']} " + f"({delta['spilled_bytes'] >> 20} MiB), backlog {backlog_files} dumps " + f"({backlog_bytes >> 20} MiB), quarantined {delta['quarantined']}{extra}" + ) + self._stats_t = now + self._stats_snapshot = dict(self._stats) + + # --- scanning --------------------------------------------------------- + + def _diags(self) -> dict[str, Path]: + """Every diagnostic dir under ``MS/`` (grid and RAW), keyed by relpath.""" + out: dict[str, Path] = {} + if not self.ms.is_dir(): + return out + for d in self.ms.rglob("*"): + if not d.is_dir(): + continue + try: + groups = _io._series_dumps(d) + except OSError: + continue + if not groups: + continue + rel = str(d.relative_to(self.ms)) + if len(groups) > 1: + # A directory holding several report series (e.g. two s1 line + # lineouts) can't be streamed as one series; leave it to the + # batch pass (save_run_datasets), which writes one NetCDF per + # series via the per-report handles from list_diagnostics. + if rel not in self._multibase_skipped: + self._multibase_skipped.add(rel) + self._log(f"[stream] {rel}: {len(groups)} report series -> batch pass") + continue + out[rel] = d + return out + + def _discover(self, *, final: bool) -> dict[str, Path]: + """Cached discovery. The full ``rglob`` walk touches every backlogged + file, so its cost grows exactly when the drainer is behind; new + diagnostic *directories* appear rarely, so the walk runs only every + ``rediscover_every`` polls and the known dirs are re-listed directly in + between.""" + if final or not self._known or (self._scan_i % self.rediscover_every == 1): + self._known = self._diags() + return self._known + + def _diag_is_raw(self, rel: str, d: Path) -> bool: + cached = self._is_raw.get(rel) + if cached is None: + cached = _io._diag_is_raw(rel, d) + self._is_raw[rel] = cached + return cached + + def _scan_once(self, *, final: bool) -> None: + self._scan_i += 1 + force_spill = False + if self.persist_dir is not None and not final and self.floor_free_bytes: + try: + if shutil.disk_usage(self.ms).free < self.floor_free_bytes: + force_spill = True + self._log_every( + "floor", + f"[stream] CRITICAL: staging filesystem under {self.floor_free_bytes >> 30} GiB free" + " -- spilling all grid dumps to persist (OSIRIS must never see ENOSPC)", + ) + except OSError: + pass + backlog_files = 0 + backlog_bytes = 0 + for rel, d in self._discover(final=final).items(): + if rel in self._completed or rel in self._raw_completed: + continue + try: + if self._diag_is_raw(rel, d): + # RAW: nothing to do unless staging (then mirror + reap the + # HDF5 so the scratch stays bounded and the batch path finds + # the dumps on persist). + if self.persist_dir is not None: + self._drain_raw(rel, d, final=final) + else: + nf, nb = self._drain_grid(rel, d, final=final, force_spill=force_spill) + backlog_files += nf + backlog_bytes += nb + except Exception as e: + self._log_every(f"drain:{rel}", f"[stream] {rel}: {e} (will fall back to batch)") + # Drop the writer so a half-written file is rebuilt by the batch + # pass rather than reused. + w = self._writers.pop(rel, None) + if w is not None: + try: + w.close() + except Exception: + pass + self._maybe_stats(backlog_files, backlog_bytes, final=final) + + # --- draining --------------------------------------------------------- + + def _drain_grid(self, rel: str, d: Path, *, final: bool, force_spill: bool) -> tuple[int, int]: + """Drain one grid diagnostic; returns its (files, bytes) pending at entry. + + Non-staged: stream to NetCDF in place — no mirror, no reap. Staged: + stream + reap (mirroring unless ``discard_grid_h5``), or mirror-only + when spilling. The final sweep additionally catches a spilled + diagnostic up from its persist mirror before closing. + """ + staged = self.persist_dir is not None + dumps = [p for p in _io._sort_dumps(d) if p not in self._bad] + # Write-completeness: a dump is safe only once the next one exists + # (OSIRIS has moved on). The final sweep runs after OSIRIS is dead, so + # the last dump is safe too. + safe = dumps if final else dumps[:-1] + n_pending = len(safe) + pending_bytes = 0 + if staged: + for p in safe: + try: + pending_bytes += p.stat().st_size + except OSError: + pass + + if staged and not final: + spill = ( + force_spill + or rel in self._spilled + or (self.spill_backlog_files and n_pending > self.spill_backlog_files) + or (self.spill_backlog_bytes and pending_bytes > self.spill_backlog_bytes) + ) + if spill: + if rel not in self._spilled: + self._spilled.add(rel) + self._log( + f"[stream] {rel}: backlog {n_pending} dumps / {pending_bytes >> 20} MiB " + "over spill threshold -- mirror-only until finalize" + ) + for p in safe: + # Always mirror when spilling: these bytes exist nowhere else. + self._reap(rel, p, mirror=True) + self._stats["spilled"] += n_pending + self._stats["spilled_bytes"] += pending_bytes + return n_pending, pending_bytes + + candidates = list(safe) + if final and staged and rel in self._spilled: + # Catch up from the persist mirror first; merge by iteration so the + # append order stays monotonic. + spill_dir = self.persist_ms / rel + if spill_dir.is_dir(): + candidates = sorted( + [p for p in _io._sort_dumps(spill_dir) if p not in self._bad] + candidates, + key=_io._iter_from_name, + ) + if not candidates: + return n_pending, pending_bytes + + w = self._writers.get(rel) + if w is None: + w, candidates = self._open_writer(rel, d, candidates) + if w is None: + return n_pending, pending_bytes + + for p in candidates: + if _io._iter_from_name(p) <= w.last_iter: + # Already in the file (resume / idempotency): just tidy up. + self._cleanup_consumed(rel, p) + continue + try: + da = _io.load_grid_h5(p) + except Exception as e: + self._quarantine(rel, p, e) + continue + w.append(da) + self._stats["appended"] += 1 + self._stats["appended_bytes"] += int(da.values.nbytes) + self._cleanup_consumed(rel, p) + w.flush() + + if final: + w.close() + self._completed.add(rel) + if staged and self.discard_grid_h5: + self._prune_empty(self.persist_ms / rel) + return n_pending, pending_bytes + + def _open_writer(self, rel: str, d: Path, candidates: list[Path]) -> tuple[StreamWriter | None, list[Path]]: + """Create (or reopen) the writer from the first readable candidate.""" + while candidates: + p0 = candidates[0] + try: + template = _io.load_grid_h5(p0) + except Exception as e: + self._quarantine(rel, p0, e) + candidates = candidates[1:] + continue + w = StreamWriter(self.out_dir / f"{rel}.nc", template, source_dir=d) + self._writers[rel] = w + if w.n_written == 0: + w.append(template) # reuse the template we just loaded as slot 0 + self._stats["appended"] += 1 + self._stats["appended_bytes"] += int(template.values.nbytes) + self._cleanup_consumed(rel, p0) + candidates = candidates[1:] + # On resume (n_written > 0) p0 stays; the iteration filter decides. + return w, candidates + return None, [] + + def _cleanup_consumed(self, rel: str, p: Path) -> None: + """Dispose of one dump whose contents are (now) in the NetCDF.""" + if self.persist_dir is None: + return # non-staged: the in-place tree is the durable copy + if self.persist_ms is not None and self.persist_ms in p.parents: + # A spilled dump drained from the persist mirror during finalize: + # in discard mode the .nc is the durable artifact, drop the HDF5. + if self.discard_grid_h5: + try: + p.unlink() + except OSError: + pass + return + self._reap(rel, p, mirror=not self.discard_grid_h5) + + def _quarantine(self, rel: str, p: Path, err: Exception) -> None: + """Sideline an unreadable dump (``*.h5.bad``) and keep streaming. + + Truncated dumps are exactly what an ENOSPC-choked run leaves behind; + one of them must not stall the whole diagnostic (previously the writer + was dropped and the watcher re-hit the same file every poll).""" + self._stats["quarantined"] += 1 + self._log_every(f"bad:{rel}", f"[stream] {rel}: quarantining unreadable dump {p.name}: {err}") + try: + p.rename(p.with_name(p.name + ".bad")) # ".bad": invisible to *.h5 listings + except OSError: + self._bad.add(p) # cannot rename: remember and skip it from now on + + def _prune_empty(self, dirpath: Path) -> None: + """Remove now-empty persist-mirror directories after a discard catch-up.""" + if self.persist_ms is None: + return + cur = dirpath + try: + while (cur == self.persist_ms or self.persist_ms in cur.parents) and cur.is_dir() and not any(cur.iterdir()): + cur.rmdir() + cur = cur.parent + except OSError: + pass + + def _drain_raw(self, rel: str, d: Path, *, final: bool) -> None: + """Mirror a RAW diagnostic's completed dumps to persist + reap scratch.""" + dumps = _io._sort_dumps(d) + if not dumps: + return + safe = dumps if final else dumps[:-1] + for p in safe: + self._reap(rel, p) + if final: + self._raw_completed.add(rel) + + def _reap(self, rel: str, p: Path, *, mirror: bool = True) -> None: + """Delete one scratch dump, first mirroring it to ``persist_dir/MS`` + unless ``mirror`` is False (``discard_grid_h5`` mode: the dump is + already in the streamed NetCDF and the raw HDF5 is not wanted on the + persist filesystem). Best-effort: on any failure the scratch file is + left in place for the runner's final sync to catch, so data is never + lost and the watcher never crashes.""" + if self.persist_dir is None: + return + try: + if mirror: + dest = self.persist_ms / rel / p.name + if not dest.exists(): + dest.parent.mkdir(parents=True, exist_ok=True) + tmp = dest.with_name(dest.name + ".part") + shutil.copyfile(p, tmp) + os.replace(tmp, dest) # atomic publish on the persist filesystem + p.unlink() + except Exception as e: + self._log_every(f"reap:{rel}", f"[stream] reap {rel}/{p.name} deferred: {e}") diff --git a/configs/osiris/twostream-1d-short.yaml b/configs/osiris/twostream-1d-short.yaml new file mode 100644 index 00000000..7fde220d --- /dev/null +++ b/configs/osiris/twostream-1d-short.yaml @@ -0,0 +1,14 @@ +solver: osiris + +mlflow: + experiment: osiris-pic1d-twostream + run: smoke-tmax-1 + +osiris: + deck: tests/test_osiris/decks/two-stream-1d + # binary omitted: the runner resolves it from OSIRIS_BIN_1D, then OSIRIS_BIN. + # Set it here only to override that (e.g. binary: /path/to/osiris-1D.e). + mpi_ranks: 1 + run_root: ./checkpoints + overrides: + time: {tmax: 1.0} diff --git a/configs/osiris/twostream-1d.yaml b/configs/osiris/twostream-1d.yaml new file mode 100644 index 00000000..f91c557e --- /dev/null +++ b/configs/osiris/twostream-1d.yaml @@ -0,0 +1,23 @@ +solver: osiris + +mlflow: + experiment: osiris-pic1d-twostream + run: cold-equal-beams + +osiris: + deck: tests/test_osiris/decks/two-stream-1d + # binary omitted: the runner resolves it from OSIRIS_BIN_1D, then OSIRIS_BIN. + # Set it here only to override that (e.g. binary: /path/to/osiris-1D.e). + mpi_ranks: 1 + run_root: ./checkpoints + # overrides: dict-of-dicts merged into the parsed deck before render. + # 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/configs/vlasov-1d/srs-debug-small.yaml b/configs/vlasov-1d/srs-debug-small.yaml index c0bfa5ec..5ac73e31 100644 --- a/configs/vlasov-1d/srs-debug-small.yaml +++ b/configs/vlasov-1d/srs-debug-small.yaml @@ -48,13 +48,18 @@ mlflow: experiment: vlasov1d-srs run: debug-small +# The ey driver pair is a backscatter-SRS matched triad in code units +# (v0 = sqrt(T0/m_e), k in 1/lambda_De, w in wp0, c_hat = c/v0 = 15.984 at 2000 eV): +# pump: w0 = 2.79, k0 = +sqrt(w0^2 - 1)/c_hat = +0.162949 (rightgoing) +# seed: w0 = 1.700942, k0 = -0.086080 (leftgoing) +# EPW: k lambda_De = 0.249029, w = 1.089058 = w_pump - w_seed drivers: ex: {} ey: '0': params: a0: 1.0 - k0: 1.0 + k0: 0.162949 w0: 2.79 dw0: 0. envelope: @@ -69,8 +74,8 @@ drivers: '1': params: a0: 1.e-6 - k0: 1.0 - w0: 1.63 + k0: -0.086080 + w0: 1.700942 dw0: 0. envelope: time: diff --git a/configs/vlasov-1d/wavepacket.yaml b/configs/vlasov-1d/wavepacket.yaml index 3770512f..5ceba970 100644 --- a/configs/vlasov-1d/wavepacket.yaml +++ b/configs/vlasov-1d/wavepacket.yaml @@ -35,7 +35,6 @@ drivers: width: 900.0 ey: {} grid: - c_light: 10 dt: 0.1 nv: 6144 nx: 4096 diff --git a/docs/osiris-adept-usage.md b/docs/osiris-adept-usage.md new file mode 100644 index 00000000..2db39c43 --- /dev/null +++ b/docs/osiris-adept-usage.md @@ -0,0 +1,347 @@ +# OSIRIS adept module — usage overview + +## File structure + +``` +adept/ (repo root) +├── run.py ← CLI entry (--cfg path, no .yaml suffix) +├── adept/ +│ ├── _base_.py ← dispatcher: `osiris` solver branch +│ └── osiris/ OSIRIS wrapper package +│ ├── __init__.py lazy export of BaseOsiris +│ ├── base.py ◀── BaseOsiris(ADEPTModule): __init__ parses deck, +│ │ __call__ runs OSIRIS, post_process delegates +│ ├── deck.py ◀── namelist parser/renderer/merger +│ │ parse_deck(text), parse_deck_file(path), +│ │ render_deck(sections), merge_overrides(...), +│ │ deck_to_flat_dict(...) for MLflow params +│ ├── runner.py ◀── subprocess driver +│ │ run_osiris(deck_text, binary=…, mpi_ranks=…), +│ │ discover_binary(...), OSIRIS-error detection +│ ├── post.py ◀── post-run collection: final-step HDF5 copy, +│ │ NetCDF export, scalar metrics +│ ├── io.py ◀── HDF5/NetCDF readers + dataset save/load +│ │ (load_series ±t_indices, lazy open_series, series_len) +│ ├── plots.py ◀── canned plot set (save_canned_plots) +│ └── regen.py ◀── regenerate plots offline from saved NetCDFs +├── configs/ +│ └── osiris/ example manifests +│ ├── twostream-1d.yaml full run (deck tmax=100) +│ └── twostream-1d-short.yaml tmax=1.0 smoke +└── tests/ + └── test_osiris/ + ├── decks/two-stream-1d in-repo example deck (manifests point here) + ├── test_deck_roundtrip.py namelist parser round-trip + ├── test_runner.py subprocess runner + discover_binary + ├── test_units.py units.yaml derivation + ├── test_post_netcdf.py post → NetCDF export + ├── test_io_and_plots.py io readers + plotting + ├── test_diagnostics_plots.py per-diagnostic plots + └── test_plots_new_views.py field-decomp / phase-space views +``` + +## End-to-end data flow + +``` + YAML manifest OSIRIS native deck + (configs/osiris/*.yaml) (tests/test_osiris/decks/… or any path) + │ │ + └──────────► run.py --cfg ◄──────────┘ + │ + ── setup phase ────────┼────────────────────────────────────────────── + ergoExo.setup ──► BaseOsiris.__init__: + parse_deck_file() + merge_overrides() (mutates sections in place) + cfg["deck"] = deck_to_flat_dict(merged sections) + │ + ▼ + log_params ──► MLflow ← logs the POST-override deck + │ + ── run phase ──────────┼────────────────────────────────────────────── + ergoExo(modules) ─► BaseOsiris.__call__: + render_deck(merged sections) ──► run_dir/os-stdin + runner.run_osiris(mpirun…) ──► run_dir/MS/{FLD,PHA,…} + │ + ▼ + post.collect ──► td/ ──► MLflow +``` + +The deck logged to MLflow is the one **after** `merge_overrides` is applied — the +same merged sections that `render_deck` later writes to `os-stdin` — so the logged +params always match what OSIRIS actually ran. + +## How to run + +First, point the runner at your built OSIRIS binary. The runner resolves it in +this order: `osiris.binary` in the manifest → `OSIRIS_BIN_D` env var (e.g. +`OSIRIS_BIN_1D`) → `OSIRIS_BIN` env var. The example manifests omit `osiris.binary`, +so set the env var once per shell: + +```bash +export OSIRIS_BIN_1D=/path/to/osiris-1D.e # per-dim, preferred +# export OSIRIS_BIN=/path/to/osiris.e # or a single default for all dims +``` + +Then run from the repo root (`run.py` appends `.yaml` to `--cfg`, so omit the suffix): + +```bash +uv run run.py --cfg configs/osiris/twostream-1d-short # smoke +uv run run.py --cfg configs/osiris/twostream-1d # full +uv run mlflow ui --backend-store-uri file://$(pwd)/mlruns # browse +``` + +## Manifest schema + +```yaml +solver: osiris # required, dispatch key + +mlflow: + experiment: osiris-pic1d-twostream # required + run: cold-equal-beams # required + +osiris: + deck: tests/test_osiris/decks/two-stream-1d # required: source of truth (repo-relative) + # binary: /path/to/osiris-1D.e # optional: overrides the env-var default + # # (OSIRIS_BIN_D → OSIRIS_BIN) + mpi_ranks: 1 # 1 → direct, >1 → mpirun -n N + run_root: ./checkpoints # parent of per-run dirs (default) + # NOTE: the default sits inside checkpoints/ deliberately — sync-up.sh + # 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 + + stream_convert: true # optional (default true): convert MS/ + # HDF5 dumps to binary/*.nc concurrently + # with the run; set false for the old + # batch conversion at job end + stream_poll_s: 10.0 # optional: watcher poll interval (s) + # stage_root: /dev/shm/osiris # optional: run OSIRIS on this fast + # ephemeral filesystem (ramdisk) and + # drain dumps to run_root in the + # background — see "Ramdisk staging". + # Requires stream_convert: true. + + density: # optional: adaptive box sizing (1D) + gradient_scale_length: 300um # target L_n; scales the box so the + # deck's density ramp realizes this L + # min: 0.225 # nmin (n_c units); default: deck profile.fx + # max: 0.275 # nmax (n_c units); default: deck profile.fx + # reference_density: 0.25 # density (n_c units) where L is defined + + overrides: # optional: applied before render + time: {tmax: 50.0} # merge into the (one) time block + grid: {nx_p: [256]} # refresh an array key + 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. + +## Adaptive box sizing from a density gradient scale length + +`osiris.density` (1D decks) scales the simulation box so the deck's linear density +ramp realizes a target gradient scale length `L`, mirroring how adept's `_lpse2d` +and `kinetic_srs` solvers size their grids. This is the **inverse** of holding the +box fixed and steepening the ramp: here the density range (`nmin`/`nmax`) is held +and the box length follows `L`. + +For a linear ramp `n(x): nmin → nmax`, the local scale length is +`L(x) = n(x)/(dn/dx) = n(x)·ramp_span/(nmax−nmin)`, so requiring `L(n_ref) = L` at the +reference density `n_ref` (default the quarter-critical surface, `n_c/4 = 0.25`) +fixes the ramp span — the same relation adept uses, `ramp_span = L/0.25·(nmax−nmin)`. + +The result is a single spatial scale factor `s` applied to **every** length in the +deck: `space.xmin`/`space.xmax`, all `profile.x` arrays, and all `diag_species` +phase-space windows (`ps_xmin`/`ps_xmax`). `grid.nx_p` is scaled by `s` too — holding +the cell size `dx` fixed (rounded up to a multiple of `node_conf.node_number(1)` for +even domain decomposition). Time (`dt`, `tmax`) is untouched, so the CFL ratio is +preserved. + +- Activates only when `osiris.density.gradient_scale_length` is present (decks + otherwise run with their hand-set box, unchanged). Runs **after** `overrides`, so + it supersedes any `space.xmax` override. +- `gradient_scale_length` takes a unit string (`300um`, converted via the deck's + `simulation.n0`/`omega_p0`) or a bare number already in `c/wp0` units. +- `min`/`max` default to the ramp's interior `profile.fx` endpoints; if given, they + are written into the primary `profile.fx`. +- The computed quantities (`box_norm`, `nx`, `scale_factor`, …) are logged under + `osiris.density.derived.*`. +- Multi-dimensional decks raise `NotImplementedError`. Drive positions (e.g. a + `zpulse` spatial center) are **not** rescaled — boundary antennas like the SRS + deck's `antenna_array` have no position to scale. + +## What lands in MLflow + +| Kind | Content | +| ----------- | -------------------------------------------------------------------------------- | +| Params | every deck key, flattened — `deck.grid.nx_p_1:1`, `deck.species_0.ufl_1:3.0`, … | +| Params | the manifest itself — `solver`, `osiris.deck`, `output.diagnostics_to_log`, … | +| Metrics | `wall_time_s`, `exit_code`, `field_energy_final`, `final_iter`, `run_time`, `postprocess_time` | +| Artifacts | `config.yaml`, `derived_config.yaml`, `units.yaml` (adept stock) | +| Artifacts | `os-stdin` (rendered OSIRIS deck), `stdout.log`, `stderr.log` | +| Artifacts | `binary//.nc` — one xarray netCDF per diagnostic, holding the full `(t, …)` time history (replaces the raw h5 dumps) | +| Artifacts | `plots/…` — canned PNGs (see below) | + +## Concurrent H5 → NetCDF conversion (`osiris.stream_convert`) + +The `binary/*.nc` series are converted **during** the run by default +(`osiris.stream_convert: true`). The alternative — set it `false` — is the old +behavior: build them **after** the run, where `post.collect` reads each +diagnostic's whole `(t, …)` history into memory and writes it in one shot. For +runs that dump a field every few steps (hundreds of thousands of tiny HDF5 +files) that end-of-run pass is slow (a cold Lustre re-read of every file) and +memory-hungry (the whole stacked series at once), which is why streaming is the +default. + +With streaming on, a best-effort background thread +(`adept/osiris/stream.py::StreamConverter`), spawned by `run_osiris` alongside +the OSIRIS subprocess, polls `MS/` every `stream_poll_s` seconds and appends +each completed dump into `/binary/.nc` as it lands. Effects: + +- **I/O overlaps compute** — the conversion latency is hidden behind the + (hours-long) PIC run, and each dump is read while still warm in the page + cache instead of re-opened cold at the end. +- **Bounded memory** — both the watcher and the at-job-end fallback build the + NetCDF a dump at a time (`load_grid_h5` → one `(t, …)` slot), so peak + conversion RAM is a single dump rather than the full series. + +Mechanics and guarantees: + +- Each NetCDF uses an **unlimited `t` dimension grown one slot per dump**, so it + ends up sized to exactly the dumps produced — no pre-sized guess, no trailing + fill to trim on early termination. Compression (zlib + shuffle) and chunking + are preserved, so the files match the batch path's `binary/*.nc` contract and + every downstream reader (`load_series_nc`, the plotting code, `regen`) is + unchanged. +- A dump is only read once the **next** dump exists (OSIRIS has moved on), + avoiding partial-file reads; the final sweep at job end picks up the last one. +- **Failure-isolated:** any error in the converter is logged and never touches + the OSIRIS run. The standard batch conversion is the safety net — `post.collect` + reuses whatever the watcher finished and stream-builds (still memory-bounded) + any grid diagnostic it did not reach. RAW (particle) diagnostics, whose + per-dump particle count varies, always use the batch concat path. +- **Restart-safe:** a `binary/.nc` reopened after a checkpoint restart + resumes after the slots already on disk. + +This addresses the I/O wall-time and conversion-memory problems only; the +heavy `pcolormesh`/`fft2` plotting cost at *plot* time is independent and +handled by the plotting changes, not here. See +`osiris-lpi/postproc-performance.md` for the full analysis. + +## Ramdisk staging (`osiris.stage_root`) + +OSIRIS has **no asynchronous diagnostic I/O**: every dump is a synchronous +barrier in the time loop, so on a parallel filesystem (Lustre) the ranks stall +on each collective write, and at full dump cadence that stall can halve compute +utilization. `stage_root` works around it without trimming cadence: point it at +a fast **node-local ephemeral filesystem** (a `/dev/shm` ramdisk on NERSC GPU +nodes — those nodes have no local disk, so RAM-backed `tmpfs` is the only +node-local option) and OSIRIS's synchronous writes hit RAM (microseconds) +instead of Lustre. The slow durable write is moved off the critical path onto +the background drainer — effectively the async I/O OSIRIS itself lacks. + +```yaml +osiris: + run_root: /pscratch/.../checkpoints # durable (Lustre) + stage_root: /dev/shm/osiris # fast scratch (ramdisk); requires stream_convert + stream_convert: true +``` + +Mechanics: + +- OSIRIS runs with its working directory on `stage_root/`; all `MS/` + dumps land in RAM. The durable run directory under `run_root` is created as + usual, and `run_dir` in the result (hence everything post-processing reads) is + always that **durable** directory — so `post.collect` is unchanged. +- The `StreamConverter` runs in **drain mode**: each completed dump is mirrored + to `run_root//MS//` and then **deleted from the ramdisk**, so + the RAM high-water mark is bounded by the poll interval × production rate, not + the whole run. Grid diagnostics are additionally streamed to the durable + `binary/*.nc` as before; RAW (particle) diagnostics, which can't be + slot-streamed, are mirrored as HDF5 so the batch path still finds them. +- `binary/*.nc` and `stdout.log`/`stderr.log` are written straight to the + durable dir (by the driver, off OSIRIS's critical path), so they survive even + if the run is killed. At job end the drainer's final sweep flushes the + held-back last dump of each diagnostic; the runner then folds any stragglers + (`HIST/`, `RE/`, …) from the scratch into the durable dir and **reclaims the + ramdisk**. + +Caveats: + +- **RAM budget.** The `tmpfs` counts against node memory *alongside* the + simulation. Bounded by the drain keeping up with production (Lustre bandwidth + ≫ dump rate, so it normally does); for very long full-cadence runs, keep + `stream_poll_s` small and leave headroom. +- **Crash window.** Whatever is still on the ramdisk when a node dies is lost + (`tmpfs` is volatile) — logs and already-mirrored dumps are safe on Lustre; + only the few in-flight dumps are at risk. If you enable OSIRIS + checkpoint/restart, keep those files off the ramdisk. +- **Single-node only** — `/dev/shm` is per-node. Fine for the 1-GPU / single-node + SRS runs; multi-node would need per-node drains and non-shared-file dumps. + +## Canned plots (`plots/` artifacts) +These canned plots focus on 1D simulations. + +`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`. + +> **Memory: phase-space plots read slices, not the whole history.** A saved `(t, …)` series can be tens of GB uncompressed — `x1gamma_q1` is ~48 GB in memory (11705 × 1000 × 1024), ~0.19 GB gzipped on disk (sparse, ~250×) — so `io.load_series(path)` decompresses all of it (at 16 sims/node this OOM'd production). `save_canned_plots` therefore loads only what each plot needs: `io.load_series(path, t_indices=…)` for the final dump + the evolution panels (`_evolution_t_indices`), and it decimates field heatmaps to render resolution before `pcolormesh`. A new plotter over a large series must do the same — `load_series(..., t_indices=…)` for a few slices, or `with io.open_series(path) as da:` (lazy; `da.isel(t=it)` reads one dump on demand) when it walks the whole time axis reducing one dump at a time. A stray `da.values` / `da.sum()` over `t`, or a **deep** `da.copy()`, pulls the full series into RAM (`_decorate` copies shallow for this reason). + +## Programmatic use + +```python +from 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 tests/test_osiris/decks/ +cp configs/osiris/twostream-1d.yaml configs/osiris/my-new.yaml +$EDITOR configs/osiris/my-new.yaml # change deck path + mlflow.run +uv run run.py --cfg configs/osiris/my-new +``` diff --git a/docs/source/solvers/vlasov1d/config.md b/docs/source/solvers/vlasov1d/config.md index 6263565d..a7e38f63 100644 --- a/docs/source/solvers/vlasov1d/config.md +++ b/docs/source/solvers/vlasov1d/config.md @@ -48,6 +48,28 @@ units: normalizing_density: 1.5e21/cc ``` +### Normalization convention + +All Vlasov-1D quantities are normalized using the following unit system, built from +`normalizing_density` ($n_0$) and `normalizing_temperature` ($T_0$): + +| Unit | Definition | Meaning | +|------|-----------|---------| +| time | $\tau = 1/\omega_{p0}$, $\omega_{p0} = \sqrt{n_0 e^2/(\epsilon_0 m_e)}$ | inverse plasma frequency | +| velocity | $v_0 = \sqrt{T_0/m_e}$ | electron thermal speed (RMS / standard-deviation convention) | +| length | $L_0 = v_0/\omega_{p0} = \lambda_{De}$ | electron Debye length | +| wavenumber | $1/L_0$ | a code wavenumber is $k \lambda_{De}$ | + +Consequences of the $v_0 = \sqrt{T_0/m_e}$ (σ) convention: + +- A Maxwellian at code temperature $T = 1$ is $f \propto e^{-v^2/2}$ (unit variance). +- The Bohm–Gross dispersion in code units is $\omega^2 = 1 + 3 k^2$. +- The normalized speed of light is $\hat c = c/\sqrt{T_0/m_e}$ (e.g. $\hat c = 15.98$ at 2000 eV). + +Dimensional inputs (strings with units, e.g. `xmax: 100um`) are converted with these +units; plain numeric inputs are taken to already be in code units and pass through +unchanged. + ## density Species and density configuration. You can define multiple "species" by using keys prefixed with `species-`. @@ -65,9 +87,9 @@ Each species is defined with a key starting with `species-` (e.g., `species-back | `noise_seed` | int | Random seed for noise initialization | | `noise_type` | string | `"gaussian"` or `"uniform"` | | `noise_val` | float | Amplitude of noise | -| `v0` | float | Drift velocity (normalized) | -| `T0` | float | Temperature (normalized to `normalizing_temperature`) | -| `m` | float | Exponent for super-Gaussian distribution. `2.0` is Maxwellian | +| `v0` | float | Drift velocity in code units of $\sqrt{T_0/m_e}$ (thermal-σ units). Numeric only — dimensional strings are not supported here | +| `T0` | float | Temperature in units of `normalizing_temperature`. The initialized distribution has velocity variance `T0/mass`. Numeric only | +| `m` | float | Exponent for super-Gaussian distribution $f \propto \exp[-\|v/(\alpha v_{th})\|^m]$. `2.0` is Maxwellian. Note: $\alpha$ is chosen to fix the moment ratio $\langle v^4\rangle/\langle v^2\rangle = 3\,T_0/\mathrm{mass}$ for all $m$; the *variance* equals `T0/mass` only at `m: 2`. For flat-top distributions (`m > 2`) the second-moment temperature diagnostic will read higher than `T0` (e.g. ×1.24 at `m: 3`, ×1.37 at `m: 4`) | | `basis` | string | Spatial profile type (see below) | #### Basis Types @@ -535,7 +557,7 @@ Both `fokker_planck` and `krook` use profile objects to define spatiotemporal va | Field | Type | Description | |-------|------|-------------| -| `baseline` | float | Base collision frequency | +| `baseline` | float | Base collision frequency in code units of $\omega_{p0}$ (i.e. $\hat\nu = \nu[\mathrm{s}^{-1}]/\omega_{p0}$, a true rate — no $2\pi$). For reference, the NRL e–e rate in these units is logged as `nuee_norm` in `units.yaml`. Note the Dougherty/Lenard–Bernstein $\nu$ agrees with the NRL $\nu_{ee}$ only up to an O(1) factor | | `bump_or_trough` | string | `"bump"` or `"trough"` | | `center` | float | Profile center | | `rise` | float | Transition steepness | 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_adaptive_box.py b/tests/test_osiris/test_adaptive_box.py new file mode 100644 index 00000000..fe0cdd01 --- /dev/null +++ b/tests/test_osiris/test_adaptive_box.py @@ -0,0 +1,170 @@ +"""Tests for ``osiris.density`` adaptive box sizing. + +When a manifest carries an ``osiris.density.gradient_scale_length``, +``BaseOsiris`` scales the simulation box so the deck's linear density ramp +realizes that gradient scale length at the reference density (default the +quarter-critical surface). This mirrors adept's ``_lpse2d`` / ``kinetic_srs`` +grid sizing: ``ramp_span = L / n_ref * (nmax - nmin)``. + +The reference deck ``srs-1d_lpi`` ramps n: 0.225 -> 0.275 (n_c units) across +x in (1e-5, 1076.03) c/wp0, box xmax = 1076.04, nx_p = 6000, with n0 = 9.05e21 +(critical density for 351 nm, so the skin depth c/wp0 = 55.86 nm). +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest + +from adept.osiris import BaseOsiris +from adept.osiris import deck as osd +from adept.osiris import density as den + +DECKS_DIR = Path(__file__).parent / "decks" +SRS_DECK = DECKS_DIR / "srs-1d_lpi" +TWO_D_DECK = DECKS_DIR / "F-Tsung_2d_lpi_deck" + +DECK_DX = 1076.04 / 6000 # cell size of the reference deck (c/wp0) + + +def _value(sections, name, base): + """First value in section ``name`` whose key base name is ``base``.""" + for sec_name, params in sections: + if sec_name == name: + for k, v in params.items(): + if k == base or k.split("(", 1)[0] == base: + return v + raise KeyError(f"{name}.{base} not found") + + +def _scaled(L="300um", **extra): + sections = osd.parse_deck_file(SRS_DECK) + computed = den.apply_gradient_scale_length(sections, {"gradient_scale_length": L, **extra}) + return sections, computed + + +def test_inactive_without_density_block() -> None: + sections = osd.parse_deck_file(SRS_DECK) + before = copy.deepcopy(sections) + assert den.apply_gradient_scale_length(sections, None) is None + assert den.apply_gradient_scale_length(sections, {}) is None + # no gradient scale length -> still a no-op + assert den.apply_gradient_scale_length(sections, {"min": 0.2}) is None + assert sections == before, "deck must be untouched when the feature is inactive" + + +def test_box_scaled_to_requested_scale_length() -> None: + sections, computed = _scaled("300um") + # ramp_span = L/0.25 * (0.275-0.225); L=300um at c/wp0=55.86nm -> ~5371 c/wp0 + assert _value(sections, "space", "xmax") == pytest.approx(1074.11, rel=1e-4) + assert _value(sections, "space", "xmin") == pytest.approx(0.0) + assert computed["box_norm"] == pytest.approx(1074.11, rel=1e-4) + assert computed["density_min"] == pytest.approx(0.225) + assert computed["density_max"] == pytest.approx(0.275) + assert computed["reference_density"] == pytest.approx(0.25) + + +def test_constant_dx_nx_scales_and_divides_node_count() -> None: + sections, computed = _scaled("300um") + nx = _value(sections, "grid", "nx_p") + box = _value(sections, "space", "xmax") + # cell size is held fixed (constant dx) to within node-count rounding + assert box / nx == pytest.approx(DECK_DX, rel=5e-3) + # node_number(1) = 4 in the deck -> nx_p must split evenly across 4 nodes + assert nx % 4 == 0 + assert computed["nx"] == nx + + +def test_diagnostic_window_and_profile_edges_track_box() -> None: + sections, _ = _scaled("300um") + box = _value(sections, "space", "xmax") + # phase-space diagnostic window follows the box + assert _value(sections, "diag_species", "ps_xmax") == pytest.approx(box) + assert _value(sections, "diag_species", "ps_xmin") == pytest.approx(0.0) + # profile control points (incl. the tiny vacuum edges) scale with the box + prof_x = _value(sections, "profile", "x") + assert prof_x[0] == pytest.approx(0.0) + assert prof_x[-1] == pytest.approx(box) + # density values are untouched (no min/max override given) + assert _value(sections, "profile", "fx") == [0.0, 0.225, 0.275, 0.0] + + +def test_box_scales_linearly_with_scale_length() -> None: + _, half = _scaled("150um") + _, base = _scaled("300um") + _, dbl = _scaled("600um") + assert half["box_norm"] == pytest.approx(base["box_norm"] / 2, rel=1e-6) + assert dbl["box_norm"] == pytest.approx(base["box_norm"] * 2, rel=1e-6) + + +def test_round_trips_the_deck_geometry() -> None: + # The reference deck encodes L_n = 0.25 * 1076.03 / 0.05 c/wp0 = 300.54 um; + # requesting that L must reproduce the deck's box. + sections, computed = _scaled("300.54um") + assert _value(sections, "space", "xmax") == pytest.approx(1076.04, rel=1e-4) + assert computed["scale_factor"] == pytest.approx(1.0, rel=2e-3) + + +def test_numeric_scale_length_is_treated_as_normalized() -> None: + # A bare number is taken to be already in c/wp0 units (no n0 conversion). + _, computed = _scaled(500.0) + assert computed["gradient_scale_length_norm"] == pytest.approx(500.0) + # ramp_span = 500/0.25 * 0.05 = 100 c/wp0 + assert computed["ramp_span_norm"] == pytest.approx(100.0) + + +def test_min_max_override_widens_box_and_rewrites_fx() -> None: + sections, computed = _scaled("300um", min=0.2, max=0.3) + # (nmax-nmin) doubled from 0.05 to 0.1 -> ramp span (and box) doubles + _, base = _scaled("300um") + assert computed["box_norm"] == pytest.approx(base["box_norm"] * 2, rel=1e-6) + # the new densities are written into the ramp endpoints + assert _value(sections, "profile", "fx") == [0.0, 0.2, 0.3, 0.0] + + +def test_custom_reference_density() -> None: + # ramp_span is inversely proportional to n_ref + _, qc = _scaled("300um", reference_density=0.25) + _, half = _scaled("300um", reference_density=0.125) + assert half["box_norm"] == pytest.approx(qc["box_norm"] * 2, rel=1e-6) + + +def test_multidimensional_deck_raises() -> None: + sections = osd.parse_deck_file(TWO_D_DECK) + with pytest.raises(NotImplementedError, match="1D decks only"): + den.apply_gradient_scale_length(sections, {"gradient_scale_length": "300um"}) + + +def test_physical_length_without_reference_density_raises(tmp_path: Path) -> None: + deck = tmp_path / "no_n0" + deck.write_text( + "grid { nx_p = 100, }\n" + "space { xmin(1) = 0.0, xmax(1) = 100.0, }\n" + "profile { num_x = 4, fx(1:4,1) = 0.0, 0.2, 0.3, 0.0, x(1:4,1) = 0, 1, 99, 100, }\n" + ) + sections = osd.parse_deck_file(deck) + with pytest.raises(ValueError, match="neither n0 nor"): + den.apply_gradient_scale_length(sections, {"gradient_scale_length": "300um"}) + + +def test_baseosiris_applies_box_and_records_derived() -> None: + cfg = {"osiris": {"deck": str(SRS_DECK), "density": {"gradient_scale_length": "300um"}}} + module = BaseOsiris(cfg) + # box was scaled on the live sections that __call__ will render + assert _value(module._sections, "space", "xmax") == pytest.approx(1074.11, rel=1e-4) + # computed quantities are stashed back for MLflow provenance + derived = cfg["osiris"]["density"]["derived"] + assert derived["box_norm"] == pytest.approx(1074.11, rel=1e-4) + assert derived["nx"] == _value(module._sections, "grid", "nx_p") + # write_units sees the scaled box + quants = module.write_units() + assert quants["box_length"].to("micron").magnitude == pytest.approx(60.0, rel=1e-2) + + +def test_baseosiris_unchanged_without_density() -> None: + cfg = {"osiris": {"deck": str(SRS_DECK)}} + module = BaseOsiris(cfg) + assert _value(module._sections, "space", "xmax") == pytest.approx(1076.04) + assert "density" not in cfg["osiris"] diff --git a/tests/test_osiris/test_deck_roundtrip.py b/tests/test_osiris/test_deck_roundtrip.py new file mode 100644 index 00000000..4185c465 --- /dev/null +++ b/tests/test_osiris/test_deck_roundtrip.py @@ -0,0 +1,162 @@ +"""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_parse_single_quoted_strings() -> None: + # OSIRIS decks commonly use Fortran single quotes; the parser must strip + # them so the value round-trips to a bare Python string (not "'none'"). + text = "emf_bound { ext_fld = 'none', type(1:2) = 'open', 'open', }" + sections = osd.parse_deck(text) + assert sections[0][1]["ext_fld"] == "none" + assert sections[0][1]["type(1:2)"] == ["open", "open"] + + +def test_render_normalizes_single_quotes_to_double() -> None: + parsed = osd.parse_deck("emf_bound { ext_fld = 'none', }") + rendered = osd.render_deck(parsed) + assert '"none"' in rendered and "'none'" not in rendered + # And the value survives a full round-trip unchanged. + assert osd.parse_deck(rendered) == parsed + + +def test_single_quoted_string_with_bang_and_comma() -> None: + # ``!`` and ``,`` inside a single-quoted string are literal, not a comment + # or a value separator (e.g. OSIRIS math_func expressions). + text = "diag { expr = 'if(x>0,1,0)!keep', x = 1, }" + sections = osd.parse_deck(text) + assert sections[0][1]["expr"] == "if(x>0,1,0)!keep" + assert sections[0][1]["x"] == 1 + + +def test_merge_overrides_simple() -> None: + text = "time { tmin = 0.0, tmax = 30.0, }" + s = osd.parse_deck(text) + 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..62358da0 --- /dev/null +++ b/tests/test_osiris/test_diagnostics_plots.py @@ -0,0 +1,223 @@ +"""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 {"E_energy", "B_energy", "total_field_energy"} <= set(ds.data_vars) + # Only e1 was dumped, so all energy is electric and B is identically zero. + assert np.all(ds["E_energy"].values > 0) + assert np.all(ds["B_energy"].values == 0) + np.testing.assert_allclose(ds["total_field_energy"].values, ds["E_energy"].values) + # e1 is the only dump, so the longitudinal (EPW) energy equals the E energy. + assert "e1_energy" in ds + np.testing.assert_allclose(ds["e1_energy"].values, ds["E_energy"].values) + + +def test_field_energy_picks_up_savg_and_excludes_poynting(tmp_path: Path) -> None: + """A 2D-style deck dumps e1,savg (not full grid) and s1 Poynting lineouts. + + field_energy_components must resolve the ``e1-savg`` variant, and the + total-field-energy metric must NOT fold the ``s1`` lineout into the sum. + """ + run_dir = tmp_path / "run" + x_ax = ("x1", "x_1", r"c / \omega_p", 0.0, 10.0) + x2_ax = ("x2", "x_2", r"c / \omega_p", 0.0, 4.0) + rng = np.random.default_rng(3) + for k in range(4): + it, t = k * 10, k * 0.5 + _write_dump(run_dir / "MS/FLD/e1-savg" / f"e1-savg-{it:06d}.h5", + "e1", rng.standard_normal(8), t=t, it=it, axes=[x_ax]) + # an s1 Poynting-flux lineout along x2 (would inflate field energy if summed) + _write_dump(run_dir / "MS/FLD/s1-line-x2-24" / f"s1-line-x2-24-{it:06d}.h5", + "s1", 5.0 + rng.standard_normal(6), t=t, it=it, axes=[x2_ax]) + + ds = oplt.field_energy_components(run_dir) + assert np.all(ds["e1_energy"].values > 0) # savg variant resolved + total = opost._total_field_energy(run_dir / "MS") + # only e1 contributes; s1 is excluded, so the total matches the e1 energy of + # the last dump (both are 0.5 * sum(e1^2) * dx). + assert np.isfinite(total) and total > 0 + + +def test_field_energy_components_2d_spatial(tmp_path: Path) -> None: + """field_energy_components integrates over BOTH spatial dims for a 2D field.""" + run_dir = tmp_path / "run" + x_ax = ("x1", "x_1", r"c / \omega_p", 0.0, 10.0) + x2_ax = ("x2", "x_2", r"c / \omega_p", 0.0, 4.0) + rng = np.random.default_rng(5) + for k in range(3): + it, t = k * 10, k * 0.5 + data = rng.standard_normal((6, 8)) # (x2, x1) + _write_dump(run_dir / "MS/FLD/e1" / f"e1-{it:06d}.h5", + "e1", data, t=t, it=it, axes=[x_ax, x2_ax]) + ds = oplt.field_energy_components(run_dir) + assert ds["e1_energy"].sizes["t"] == 3 + assert np.all(ds["e1_energy"].values > 0) + ax = oplt.plot_epw_energy(ds) + assert ax.get_ylabel() + + +def test_load_hist_energy_builds_conservation_total(tmp_path: Path) -> None: + 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 + "energy_partition_vs_time", # same HIST source: total split into particle/EM + } + 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..ed412d1a --- /dev/null +++ b/tests/test_osiris/test_plots_new_views.py @@ -0,0 +1,265 @@ +"""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_boundary_slabs_match_full_split(tmp_path: Path) -> None: + # The slab-restricted split must equal the full split sliced to the same + # boundary columns (the Riemann split is local, so slice-then-combine == + # combine-then-slice). This is what lets the boundary diagnostics avoid + # materializing the whole (t, x) left/right grid. + run_dir = _make_rich_run(tmp_path) + g, w = 1, 3 + full = oplt.efield_lr_components(run_dir) + sl = oplt.transverse_field_boundary_slabs(run_dir, guard_cells=g, window_cells=w) + assert sl is not None and sl["pairs"] == list(full) + assert sl["guard_cells"] == g and sl["window_cells"] == w + sample = next(iter(full.values()))["right"] + xdim = next(d for d in sample.dims if d != "t") + n = sample.sizes[xdim] + left, right = slice(g, g + w), slice(n - g - w, n - g) + for pair, parts in full.items(): + for edge, sel in (("left", left), ("right", right)): + for going in ("right", "left"): + np.testing.assert_allclose( + sl["edges"][edge][pair][going], + parts[going].isel({xdim: sel}).values, + atol=1e-6, + ) + + +def test_boundary_slabs_absent_pair_returns_none(tmp_path: Path) -> None: + run_dir = tmp_path / "run" # only a density diag, no transverse FLD pairs + _write_dump(run_dir / "MS/DENSITY/ion/charge" / "charge-ion-000000.h5", "charge", np.ones(8), 0.0, 0, [X_AX]) + assert oplt.transverse_field_boundary_slabs(run_dir) is None + + +def test_field_decomposition_figures(tmp_path: Path) -> None: + run_dir = _make_rich_run(tmp_path) + figs = oplt.plot_field_lr_decomposition(run_dir) + 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..36b7aa5a --- /dev/null +++ b/tests/test_osiris/test_post_netcdf.py @@ -0,0 +1,532 @@ +"""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 _write_lineout(dir_: Path, base: str, n_steps: int, nx: int, fill: float) -> None: + """Write one line/report series ``-.h5`` into ``dir_``.""" + for k in range(n_steps): + it = k * 10 + _write_dump( + dir_ / f"{base}-{it:06d}.h5", + "s1", + np.full(nx, fill + k, dtype="float64"), + t=k * 0.5, + it=it, + axes=[("x2", "x_2", r"c / \omega_p", 0.0, 5.0)], + ) + + +def test_multi_report_lineouts_are_separate_series(tmp_path: Path) -> None: + # Two s1 line lineouts share one MS/FLD dir, distinguished only by the report + # index in the base name (s1-tavg-line-x2-01 / -02). Both -x2-0N-000000.h5 + # read as iteration 0, so keyed by directory they would merge; they must be + # exposed and loaded as SEPARATE series instead. + run_dir = tmp_path / "run" + d = run_dir / "MS" / "FLD" / "s1-tavg-line" + _write_lineout(d, "s1-tavg-line-x2-01", n_steps=4, nx=6, fill=100.0) + _write_lineout(d, "s1-tavg-line-x2-02", n_steps=4, nx=6, fill=200.0) + + diags = oio.list_diagnostics(run_dir) + keys = {k for k in diags if k.startswith("FLD/s1-tavg-line")} + assert keys == { + "FLD/s1-tavg-line/s1-tavg-line-x2-01", + "FLD/s1-tavg-line/s1-tavg-line-x2-02", + } + ent = oio.load_series(diags["FLD/s1-tavg-line/s1-tavg-line-x2-01"]) + ex = oio.load_series(diags["FLD/s1-tavg-line/s1-tavg-line-x2-02"]) + assert ent.sizes["t"] == 4 and ex.sizes["t"] == 4 # not merged into 8 + assert float(ent.isel(t=0).values.mean()) == 100.0 + assert float(ex.isel(t=0).values.mean()) == 200.0 + + +def test_save_run_datasets_splits_multi_report_lineouts(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + d = run_dir / "MS" / "FLD" / "s1-tavg-line" + _write_lineout(d, "s1-tavg-line-x2-01", 3, 6, 100.0) + _write_lineout(d, "s1-tavg-line-x2-02", 3, 6, 200.0) + out = tmp_path / "binary" + written = oio.save_run_datasets(run_dir, out) + names = {p.relative_to(out).as_posix() for p in written} + assert "FLD/s1-tavg-line/s1-tavg-line-x2-01.nc" in names + assert "FLD/s1-tavg-line/s1-tavg-line-x2-02.nc" in names + # round-trips separately (not merged) + assert oio.load_series(out / "FLD/s1-tavg-line/s1-tavg-line-x2-01.nc").sizes["t"] == 3 + + +def test_single_report_dir_still_keys_by_dir(tmp_path: Path) -> None: + # Regression: a normal single-series dir keys by its dir relpath, unchanged. + run_dir = _make_run(tmp_path, n_steps=3, nx=8) + diags = oio.list_diagnostics(run_dir) + assert "FLD/e1" in diags + assert oio.load_series(diags["FLD/e1"]).sizes["t"] == 3 + + +def test_save_run_datasets_writes_full_timeseries(tmp_path: Path) -> None: + run_dir = _make_run(tmp_path, n_steps=4, nx=8) + out = tmp_path / "binary" + 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..4c56debb --- /dev/null +++ b/tests/test_osiris/test_runner.py @@ -0,0 +1,114 @@ +"""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_srun_chdir_warning_is_not_an_error() -> None: + # srun's spurious launcher warning under /dev/shm staging must not be flagged + # as an OSIRIS error (it aborts an otherwise-clean run); a real OSIRIS error + # and a real abort still are. + chdir = "[2026-06-30T23:35:54] error: couldn't chdir to `/dev/shm/x`: No such file or directory: going to /tmp instead" + assert runner._looks_like_osiris_error(chdir) is False + assert runner._looks_like_osiris_error("(*error*) Lindman not yet allowed with tiling") is True + assert runner._looks_like_osiris_error("Error reading global simulation parameters, aborting...") is True + + +def test_run_osiris_missing_binary_raises(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError): + runner.run_osiris( + "node_conf {}", + binary="/no/such/binary", + mpi_ranks=1, + run_root=tmp_path, + ) + + +def _write_fake_binary(path: Path, body: str) -> Path: + path.write_text("#!/bin/bash\n" + body + "\n") + path.chmod(0o755) + return path + + +def test_run_osiris_crash_with_output_is_salvaged(tmp_path: Path) -> None: + # A binary that writes a dump then exits non-zero (e.g. OSIRIS segfaulting on + # MPI/CUDA teardown *after* "Simulation completed") must NOT raise: the run + # produced data, so the runner salvages it and lets the caller consolidate + + # plot what was written. + fake = _write_fake_binary( + tmp_path / "fake-osiris", + 'mkdir -p MS/FLD/e1 && : > MS/FLD/e1/e1-000000.h5 && exit 139', + ) + result = runner.run_osiris( + "node_conf {}", + binary=str(fake), + mpi_ranks=1, + run_root=tmp_path, + stream_convert=False, + ) + assert result["exit_code"] == 139 + assert result["crashed"] is True + assert next((result["run_dir"] / "MS").rglob("*.h5"), None) is not None + + +def test_run_osiris_crash_no_output_raises(tmp_path: Path) -> None: + # A binary that exits non-zero WITHOUT writing anything is a hard failure — + # there is nothing to salvage, so the runner still raises. + fake = _write_fake_binary(tmp_path / "fake-osiris", "exit 1") + with pytest.raises(RuntimeError) as excinfo: + runner.run_osiris( + "node_conf {}", + binary=str(fake), + mpi_ranks=1, + run_root=tmp_path, + stream_convert=False, + ) + assert "nothing to salvage" in str(excinfo.value) + + +@pytest.mark.skipif( + not (OSIRIS_BIN_1D and Path(OSIRIS_BIN_1D).exists()), + reason="set OSIRIS_BIN_1D (or OSIRIS_BIN) to a built osiris-1D.e to run", +) +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_stream.py b/tests/test_osiris/test_stream.py new file mode 100644 index 00000000..b4efb529 --- /dev/null +++ b/tests/test_osiris/test_stream.py @@ -0,0 +1,603 @@ +"""Incremental / concurrent HDF5 -> NetCDF conversion (adept.osiris.stream). + +These tests synthesize a tiny OSIRIS-shaped run (no real run on disk) and check +that the streaming converter — both the at-job-end Stage A path and the +concurrent Stage B watcher — produces a NetCDF that is equivalent, for every +downstream consumer, to the in-memory batch path in ``io.save_run_datasets``. +""" + +from __future__ import annotations + +import stat +import time +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr + +from adept.osiris import io as oio +from adept.osiris import post as opost +from adept.osiris import runner as orunner +from adept.osiris import stream as ostream + + +def _write_dump(path: Path, name: str, data: np.ndarray, t: float, it: int, axes) -> None: + """Write one OSIRIS-style HDF5 grid dump. + + ``axes`` is a list of ``(name, long_name, units, min, max)`` in OSIRIS + (Fortran) order, i.e. the last entry is the fastest-varying numpy axis. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + f.attrs["NAME"] = np.array([name.encode()], dtype="S256") + f.attrs["LABEL"] = np.array([name.encode()], dtype="S256") + f.attrs["UNITS"] = np.array([b"a.u."], dtype="S256") + f.attrs["TIME UNITS"] = np.array([rb"1 / \omega_p"], dtype="S256") + ax = f.create_group("AXIS") + for i, (an, ln, un, lo, hi) in enumerate(axes, start=1): + d = ax.create_dataset(f"AXIS{i}", data=np.array([lo, hi], dtype="float64")) + d.attrs["NAME"] = np.array([an.encode()], dtype="S256") + d.attrs["LONG_NAME"] = np.array([ln.encode()], dtype="S256") + d.attrs["UNITS"] = np.array([un.encode()], dtype="S256") + sim = f.create_group("SIMULATION") + sim.attrs["DT"] = np.array([0.05]) + sim.attrs["NDIMS"] = np.array([len(axes)], dtype="int32") + sim.attrs["NX"] = np.array(list(reversed(data.shape)), dtype="int32") + sim.attrs["XMIN"] = np.array([a[3] for a in axes]) + sim.attrs["XMAX"] = np.array([a[4] for a in axes]) + f.create_dataset(name, data=data.astype("float32")) + + +def _write_field(run_dir: Path, comp: str, n_steps: int, nx: int, seed: int = 0) -> Path: + """Write an ``MS/FLD/`` 1-D field diagnostic over ``n_steps`` dumps.""" + rng = np.random.default_rng(seed) + d = run_dir / "MS" / "FLD" / comp + for k in range(n_steps): + _write_dump( + d / f"{comp}-{k * 10:06d}.h5", + comp, + rng.standard_normal(nx), + t=k * 0.5, + it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + return d + + +def _write_phasespace_autoscaled(run_dir: Path, n_steps: int, nx: int, npmom: int) -> Path: + """Phase space whose momentum bounds grow each dump (if_ps_p_auto).""" + rng = np.random.default_rng(7) + d = run_dir / "MS" / "PHA" / "p1x1" / "electrons" + for k in range(n_steps): + p_hi = 1.0 + 0.5 * k + _write_dump( + d / f"p1x1-electrons-{k * 10:06d}.h5", + "p1x1", + rng.standard_normal((npmom, nx)), + t=k * 0.5, + it=k * 10, + axes=[ + ("x1", "x_1", r"c / \omega_p", 0.0, 10.0), + ("p1", "p_1", r"m_e c", -p_hi, p_hi), + ], + ) + return d + + +def _assert_series_equiv(got: xr.DataArray, ref: xr.DataArray) -> None: + """``got`` (streamed) is interchangeable with ``ref`` (raw load_series).""" + assert got.name == ref.name + assert got.dims == ref.dims + assert got.shape == ref.shape + np.testing.assert_array_equal(got.values, ref.values) # float32 -> float32, exact + np.testing.assert_allclose(got["t"].values, ref["t"].values) + assert list(got["iter"].values) == list(ref["iter"].values) + for dim in ref.dims: + if dim == "t": + continue + np.testing.assert_allclose(oio.physical_axis(got, dim), oio.physical_axis(ref, dim)) + assert got.attrs.get("axis_units") == ref.attrs.get("axis_units") + assert got.attrs.get("autoscaled_dims") == ref.attrs.get("autoscaled_dims") + + +# --- Stage A: streaming finalize ----------------------------------------- + + +def test_streaming_grid_matches_load_series(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=5, nx=8) + + dest = tmp_path / "stream" / "e1.nc" + ostream.convert_diagnostic_streaming(diag, dest) + + streamed = oio.load_series_nc(dest) + raw = oio.load_series(diag) + _assert_series_equiv(streamed, raw) + # A fixed spatial axis is not autoscaled even though the streamer always + # records its (constant) bounds. + assert "autoscaled_dims" not in streamed.attrs + + +def test_streaming_matches_batch_file(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=6, nx=12, seed=3) + + batch = oio.save_run_datasets(run_dir, tmp_path / "batch")[0] + stream = tmp_path / "stream" / "e1.nc" + ostream.convert_diagnostic_streaming(diag, stream) + + a = oio.load_series_nc(stream) + b = oio.load_series_nc(batch) + np.testing.assert_array_equal(a.values, b.values) + np.testing.assert_allclose(a["t"].values, b["t"].values) + assert list(a["iter"].values) == list(b["iter"].values) + assert a.attrs.get("long_name") == b.attrs.get("long_name") + assert a.attrs.get("units") == b.attrs.get("units") + + +def test_streaming_preserves_autoscale_bounds(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_phasespace_autoscaled(run_dir, n_steps=4, nx=8, npmom=6) + + dest = tmp_path / "stream" / "p1x1.nc" + ostream.convert_diagnostic_streaming(diag, dest) + + streamed = oio.load_series_nc(dest) + raw = oio.load_series(diag) + _assert_series_equiv(streamed, raw) + assert streamed.attrs.get("autoscaled_dims") == ["p1"] + np.testing.assert_allclose( + oio.physical_axis(streamed, "p1", it=-1), + np.linspace(-2.5, 2.5, streamed.sizes["p1"]), + ) + + +def test_streaming_compresses_output(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=8, nx=64) + dest = tmp_path / "stream" / "e1.nc" + ostream.convert_diagnostic_streaming(diag, dest) + enc = xr.open_dataset(dest, engine="h5netcdf")["e1"].encoding + assert enc.get("zlib") or (enc.get("compression") == "gzip") + + +# --- restart / resume ----------------------------------------------------- + + +def test_writer_resumes_after_reopen(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=6, nx=8) + dumps = oio._sort_dumps(diag) + dest = tmp_path / "stream" / "e1.nc" + + # First "segment": write the first 3 dumps, then close (simulates a crash + # after a checkpoint). + template = oio.load_grid_h5(dumps[0]) + w = ostream.StreamWriter(dest, template, source_dir=diag) + w.append(template) + for p in dumps[1:3]: + w.append(oio.load_grid_h5(p)) + assert w.n_written == 3 + w.close() + + # Restart: reopen and continue from where it left off. + w2 = ostream.StreamWriter(dest, oio.load_grid_h5(dumps[0]), source_dir=diag) + assert w2.n_written == 3 # picked up the on-disk slots + for p in dumps[3:]: + w2.append(oio.load_grid_h5(p)) + w2.close() + + _assert_series_equiv(oio.load_series_nc(dest), oio.load_series(diag)) + + +# --- Stage B: concurrent watcher ----------------------------------------- + + +def test_watcher_holds_back_in_progress_dump(tmp_path: Path) -> None: + """A non-final scan must not read the highest-numbered (still-writing) dump.""" + run_dir = tmp_path / "run" + _write_field(run_dir, "e1", n_steps=5, nx=8) + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + conv._scan_once(final=False) + # 5 dumps present -> 4 are "safe", the newest is held back. + ds = xr.open_dataset(tmp_path / "binary" / "FLD" / "e1.nc", engine="h5netcdf") + try: + assert ds.sizes["t"] == 4 + finally: + ds.close() + + completed = conv.finalize() # picks up the last dump and closes + assert "FLD/e1" in completed + _assert_series_equiv( + oio.load_series_nc(tmp_path / "binary" / "FLD" / "e1.nc"), + oio.load_series(run_dir / "MS" / "FLD" / "e1"), + ) + + +def test_watcher_thread_streams_dumps_as_they_appear(tmp_path: Path) -> None: + """End-to-end Stage B: dumps written over time while the thread polls.""" + run_dir = tmp_path / "run" + diag = run_dir / "MS" / "FLD" / "e1" + rng = np.random.default_rng(0) + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.02) + conv.start() + n_steps = 6 + for k in range(n_steps): + _write_dump( + diag / f"e1-{k * 10:06d}.h5", + "e1", + rng.standard_normal(8), + t=k * 0.5, + it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)], + ) + time.sleep(0.03) + + completed = conv.finalize() + assert "FLD/e1" in completed + streamed = oio.load_series_nc(tmp_path / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == n_steps + _assert_series_equiv(streamed, oio.load_series(diag)) + + +def test_watcher_skips_raw_diagnostics(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + _write_field(run_dir, "e1", n_steps=3, nx=8) + # A RAW diagnostic: no AXIS group, several per-particle datasets. + raw = run_dir / "MS" / "RAW" / "species_1" + raw.mkdir(parents=True) + with h5py.File(raw / "RAW-species_1-000000.h5", "w") as f: + f.attrs["TIME"] = np.array([0.0]) + f.attrs["ITER"] = np.array([0], dtype="int32") + for q in ("p1", "x1"): + f.create_dataset(q, data=np.arange(4, dtype="float32")) + f.create_group("SIMULATION").attrs["NDIMS"] = np.array([1], dtype="int32") + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + completed = conv.finalize() + assert "FLD/e1" in completed + assert not any("RAW" in c for c in completed) + assert not (tmp_path / "binary" / "RAW" / "species_1.nc").exists() + + +# --- integration with save_run_datasets / post.collect ------------------- + + +def test_save_run_datasets_reuses_streamed_file(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=4, nx=8) + + # Pretend the watcher already produced this during the run, and tag it so we + # can tell a reuse (copy) from a rebuild. + streamed_dir = tmp_path / "run" / "binary" + ostream.convert_diagnostic_streaming(diag, streamed_dir / "FLD" / "e1.nc") + with xr.open_dataset(streamed_dir / "FLD" / "e1.nc", engine="h5netcdf") as ds: + ds = ds.load() + ds.attrs["_marker"] = "from-watcher" + ds.to_netcdf(streamed_dir / "FLD" / "e1.nc", engine="h5netcdf") + + out = tmp_path / "td" / "binary" + written = oio.save_run_datasets(run_dir, out, stream=True, streamed_dir=streamed_dir) + assert (out / "FLD" / "e1.nc") in written + with xr.open_dataset(out / "FLD" / "e1.nc", engine="h5netcdf") as ds2: + assert ds2.attrs.get("_marker") == "from-watcher" # copied, not rebuilt + + +def test_save_run_datasets_stream_builds_missing(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=5, nx=8) + # stream=True but no streamed_dir -> build via the streaming writer. + out = tmp_path / "binary" + oio.save_run_datasets(run_dir, out, stream=True) + _assert_series_equiv(oio.load_series_nc(out / "FLD" / "e1.nc"), oio.load_series(diag)) + + +def test_run_osiris_streams_concurrently(tmp_path: Path) -> None: + """run_osiris spawns the watcher alongside the (fake) OSIRIS subprocess. + + The fake binary is a shell script that drops pre-staged dumps into + ``MS/FLD/e1`` one at a time with a pause between — exactly the cadence the + watcher polls. No Python/h5py is needed in the subprocess. + """ + n_steps = 5 + stage = tmp_path / "stage" + stage.mkdir() + rng = np.random.default_rng(1) + ref = {} + for k in range(n_steps): + fn = f"e1-{k * 10:06d}.h5" + data = rng.standard_normal(8) + ref[fn] = data + _write_dump(stage / fn, "e1", data, t=k * 0.5, it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)]) + + lines = ["#!/bin/sh", "set -e", "mkdir -p MS/FLD/e1"] + for k in range(n_steps): + fn = f"e1-{k * 10:06d}.h5" + lines.append(f"cp '{stage / fn}' MS/FLD/e1/{fn}") + lines.append("sleep 0.1") + fake = tmp_path / "fake-osiris.sh" + fake.write_text("\n".join(lines) + "\n") + fake.chmod(fake.stat().st_mode | stat.S_IEXEC | stat.S_IRUSR) + + result = orunner.run_osiris( + "node_conf\n{\n}\n", + binary=fake, + mpi_ranks=1, + run_root=tmp_path / "checkpoints", + stream_convert=True, + stream_poll_s=0.05, + ) + + assert result["exit_code"] == 0 + assert result["streamed_diagnostics"] == ["FLD/e1"] + nc = Path(result["binary_dir"]) / "FLD" / "e1.nc" + assert nc.exists() + streamed = oio.load_series_nc(nc) + assert streamed.sizes["t"] == n_steps + _assert_series_equiv(streamed, oio.load_series(Path(result["run_dir"]) / "MS" / "FLD" / "e1")) + + +def _write_raw_dump(path: Path, t: float, it: int, npart: int = 4) -> None: + """Write one OSIRIS-style RAW (particle) dump: per-particle datasets, no AXIS.""" + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + f.attrs["TIME"] = np.array([t]) + f.attrs["ITER"] = np.array([it], dtype="int32") + for q in ("p1", "x1"): + f.create_dataset(q, data=np.arange(npart, dtype="float32")) + f.create_group("SIMULATION").attrs["NDIMS"] = np.array([1], dtype="int32") + + +# --- Staging mode: ramdisk scratch + drain to durable persist ------------- + + +def test_staged_converter_mirrors_and_reaps(tmp_path: Path) -> None: + """With ``persist_dir`` set, every completed dump is mirrored to the durable + tree and deleted from the scratch; grid is streamed, raw is mirrored as H5.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=4, nx=8) + for k in range(3): + _write_raw_dump(stage / "MS" / "RAW" / "species_1" / f"RAW-species_1-{k * 10:06d}.h5", t=k * 0.5, it=k * 10) + + conv = ostream.StreamConverter(stage, persist / "binary", poll_s=0.01, persist_dir=persist) + completed = conv.finalize() + + # Grid streamed to the durable binary/ tree; raw is mirrored but NOT returned + # as "streamed" (the batch pass still builds its NetCDF from the H5). + assert completed == {"FLD/e1"} + assert (persist / "binary" / "FLD" / "e1.nc").exists() + + # Every dump (grid + raw) mirrored to the durable MS/ tree... + assert len(list((persist / "MS" / "FLD" / "e1").glob("*.h5"))) == 4 + assert len(list((persist / "MS" / "RAW" / "species_1").glob("*.h5"))) == 3 + # ...and reaped from the scratch (RAM freed). + assert not list((stage / "MS").rglob("*.h5")) + + # The streamed grid series is faithful to the mirrored dumps. + _assert_series_equiv( + oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc"), + oio.load_series(persist / "MS" / "FLD" / "e1"), + ) + + +def test_staged_converter_discard_grid_h5(tmp_path: Path) -> None: + """With ``discard_grid_h5`` grid dumps are streamed + deleted WITHOUT a + persist mirror (inode saver); RAW dumps are still mirrored as H5.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=4, nx=8) + for k in range(3): + _write_raw_dump(stage / "MS" / "RAW" / "species_1" / f"RAW-species_1-{k * 10:06d}.h5", t=k * 0.5, it=k * 10) + + conv = ostream.StreamConverter( + stage, persist / "binary", poll_s=0.01, persist_dir=persist, discard_grid_h5=True + ) + completed = conv.finalize() + + assert completed == {"FLD/e1"} + # Grid: streamed NetCDF is the only durable artifact -- no H5 mirror. + assert (persist / "binary" / "FLD" / "e1.nc").exists() + assert not (persist / "MS" / "FLD").exists() + # RAW: still mirrored (batch path needs the H5). + assert len(list((persist / "MS" / "RAW" / "species_1").glob("*.h5"))) == 3 + # Scratch fully reaped either way. + assert not list((stage / "MS").rglob("*.h5")) + # The streamed series carries all 4 dumps. + assert oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc").sizes["t"] == 4 + + +def test_discard_layout_discovery_and_batch(tmp_path: Path) -> None: + """End-to-end stage_discard_h5 layout: MS/ holds only RAW, grid data lives + solely in the streamed NetCDFs — discovery and the batch pass must still + see and consolidate the grid diagnostics (regression: smoke-3 uploaded + 2/14 binary files because list_diagnostics only walked MS/).""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=4, nx=8) + for k in range(3): + _write_raw_dump(stage / "MS" / "RAW" / "species_1" / f"RAW-species_1-{k * 10:06d}.h5", t=k * 0.5, it=k * 10) + conv = ostream.StreamConverter( + stage, persist / "binary", poll_s=0.01, persist_dir=persist, discard_grid_h5=True + ) + conv.finalize() + + # Discovery must see the union: RAW from MS/, grid from the streamed nc. + diags = oio.list_diagnostics(persist) + assert "RAW/species_1" in diags + assert "FLD/e1" in diags and str(diags["FLD/e1"]).endswith(".nc") + + # The batch pass must consolidate the streamed grid nc into td/binary. + out = tmp_path / "td_binary" + oio.save_run_datasets(persist, out, stream=True, streamed_dir=persist / "binary") + assert (out / "FLD" / "e1.nc").exists() + assert oio.load_series_nc(out / "FLD" / "e1.nc").sizes["t"] == 4 + assert (out / "RAW" / "species_1.nc").exists() + + +def test_non_staged_converter_never_deletes(tmp_path: Path) -> None: + """Without ``persist_dir`` the converter must leave the source dumps in place + (the durable run dir *is* run_dir) — no reaping of the in-place tree.""" + run_dir = tmp_path / "run" + _write_field(run_dir, "e1", n_steps=4, nx=8) + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + conv.finalize() + assert len(list((run_dir / "MS" / "FLD" / "e1").glob("*.h5"))) == 4 + assert not (tmp_path / "MS").exists() # nothing mirrored anywhere + + +def test_run_osiris_staging_mirrors_and_frees_scratch(tmp_path: Path) -> None: + """End-to-end: OSIRIS works on the ramdisk, the drainer mirrors to the + durable run dir, and the scratch is reclaimed before returning.""" + n_steps = 5 + src = tmp_path / "src" + src.mkdir() + rng = np.random.default_rng(2) + for k in range(n_steps): + _write_dump(src / f"e1-{k * 10:06d}.h5", "e1", rng.standard_normal(8), t=k * 0.5, it=k * 10, + axes=[("x1", "x_1", r"c / \omega_p", 0.0, 10.0)]) + + lines = ["#!/bin/sh", "set -e", "mkdir -p MS/FLD/e1"] + for k in range(n_steps): + fn = f"e1-{k * 10:06d}.h5" + lines.append(f"cp '{src / fn}' MS/FLD/e1/{fn}") + lines.append("sleep 0.1") + fake = tmp_path / "fake-osiris.sh" + fake.write_text("\n".join(lines) + "\n") + fake.chmod(fake.stat().st_mode | stat.S_IEXEC | stat.S_IRUSR) + + ramdisk = tmp_path / "ramdisk" + result = orunner.run_osiris( + "node_conf\n{\n}\n", + binary=fake, + mpi_ranks=1, + run_root=tmp_path / "checkpoints", + stream_convert=True, + stream_poll_s=0.05, + stage_root=ramdisk, + ) + + assert result["exit_code"] == 0 + assert result.get("staged") is True + run_dir = Path(result["run_dir"]) + # Durable run dir lives under run_root, and the scratch was reclaimed. + assert run_dir.parent == (tmp_path / "checkpoints").resolve() + assert not (ramdisk.resolve() / run_dir.name).exists() + # Outputs landed durably: streamed .nc, mirrored MS dumps, deck + logs. + assert (run_dir / "binary" / "FLD" / "e1.nc").exists() + assert len(list((run_dir / "MS" / "FLD" / "e1").glob("*.h5"))) == n_steps + assert (run_dir / "os-stdin").exists() + assert (run_dir / "stdout.log").exists() + streamed = oio.load_series_nc(run_dir / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == n_steps + _assert_series_equiv(streamed, oio.load_series(run_dir / "MS" / "FLD" / "e1")) + + +def test_collect_uses_streamed_binary(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=3, nx=8) + (run_dir / "os-stdin").write_text("node_conf\n{\n}\n") + (run_dir / "stdout.log").write_text("ok\n") + (run_dir / "stderr.log").write_text("") + # Simulate runner having streamed into run_dir/binary during the run. + binary_dir = run_dir / "binary" + ostream.convert_diagnostic_streaming(diag, binary_dir / "FLD" / "e1.nc") + + run_output = { + "solver result": { + "run_dir": str(run_dir), + "wall_time": 1.0, + "exit_code": 0, + "binary_dir": str(binary_dir), + "streamed_diagnostics": ["FLD/e1"], + } + } + cfg = {"osiris": {"deck": str(run_dir / "os-stdin"), "stream_convert": True}, "output": {}} + + td = tmp_path / "td" + td.mkdir() + opost.collect(run_output, cfg, str(td)) + + assert not list(td.rglob("*.h5")) + assert (td / "binary" / "FLD" / "e1.nc").exists() + _assert_series_equiv(oio.load_series_nc(td / "binary" / "FLD" / "e1.nc"), oio.load_series(diag)) + + +# --- corrupt-dump quarantine + backlog spill valve ------------------------ + + +def test_corrupt_dump_quarantined_stream_continues(tmp_path: Path) -> None: + """One unreadable dump must not stall the diagnostic: it is renamed + ``*.h5.bad`` and the remaining dumps still stream (regression: job 56005549 + dropped the writer and re-hit the same corrupt file every poll, forever).""" + run_dir = tmp_path / "run" + diag = _write_field(run_dir, "e1", n_steps=5, nx=8) + bad = diag / "e1-000020.h5" + bad.write_bytes(b"\x00" * 2048) # ENOSPC-style truncated file + + conv = ostream.StreamConverter(run_dir, tmp_path / "binary", poll_s=0.01) + completed = conv.finalize() + + assert "FLD/e1" in completed + assert not bad.exists() and (diag / "e1-000020.h5.bad").exists() + streamed = oio.load_series_nc(tmp_path / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == 4 + assert list(streamed["iter"].values) == [0, 10, 30, 40] + + +def test_spill_valve_mirrors_then_catches_up(tmp_path: Path) -> None: + """Over-threshold backlog flips a staged diagnostic to mirror-only (cheap + copy to persist, ramdisk keeps draining); finalize catches the NetCDF up + from the mirror so nothing is lost.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + diag = _write_field(stage, "e1", n_steps=6, nx=8) + + conv = ostream.StreamConverter( + stage, persist / "binary", poll_s=0.01, persist_dir=persist, spill_backlog_files=2 + ) + conv._scan_once(final=False) + # 5 safe dumps > threshold 2 -> spilled: mirrored + reaped, nothing streamed. + assert "FLD/e1" in conv._spilled + assert not (persist / "binary" / "FLD" / "e1.nc").exists() + assert len(list((persist / "MS" / "FLD" / "e1").glob("*.h5"))) == 5 + assert len(list(diag.glob("*.h5"))) == 1 # only the held-back newest dump + + completed = conv.finalize() + assert "FLD/e1" in completed + streamed = oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc") + assert streamed.sizes["t"] == 6 + _assert_series_equiv(streamed, oio.load_series(persist / "MS" / "FLD" / "e1")) + assert not list((stage / "MS").rglob("*.h5")) + + +def test_spill_valve_discard_mode_prunes_mirror(tmp_path: Path) -> None: + """In discard mode the spill mirror is a temporary holding area: after the + finalize catch-up the mirrored HDF5 (and its emptied dirs) are removed and + the streamed NetCDF is again the only grid artifact.""" + stage = tmp_path / "stage" + persist = tmp_path / "persist" + _write_field(stage, "e1", n_steps=6, nx=8) + + conv = ostream.StreamConverter( + stage, + persist / "binary", + poll_s=0.01, + persist_dir=persist, + discard_grid_h5=True, + spill_backlog_files=2, + ) + conv._scan_once(final=False) + assert len(list((persist / "MS" / "FLD" / "e1").glob("*.h5"))) == 5 # spilled bytes exist nowhere else + + completed = conv.finalize() + assert "FLD/e1" in completed + assert oio.load_series_nc(persist / "binary" / "FLD" / "e1.nc").sizes["t"] == 6 + assert not (persist / "MS" / "FLD").exists() # mirror consumed + pruned + assert not list((stage / "MS").rglob("*.h5")) diff --git a/tests/test_osiris/test_units.py b/tests/test_osiris/test_units.py new file mode 100644 index 00000000..229fbfe3 --- /dev/null +++ b/tests/test_osiris/test_units.py @@ -0,0 +1,97 @@ +"""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_laser_intensity_in_icf_units() -> None: + # SRS deck: antenna{a0=0.004, omega0=1.0} with n0 = 9.05e21 cm^-3 = + # n_crit(351 nm), so w_laser = wp0 and lambda = 351 nm. The peak intensity + # of a linearly polarized drive follows the ICF convention + # I[W/cm^2] * lambda[um]^2 = 1.37e18 * a0^2. + quants = BaseOsiris({"osiris": {"deck": str(SRS_DECK)}}).write_units() + assert quants["w_laser"].to("rad/s").magnitude == pytest.approx(5.3668e15, rel=1e-3) + assert quants["laser_wavelength"].to("nm").magnitude == pytest.approx(351.0, rel=1e-3) + assert quants["laser_a0"] == pytest.approx(0.004) + lam_um = quants["laser_wavelength"].to("um").magnitude + assert quants["laser_intensity"].to("W/cm^2").magnitude == pytest.approx( + 1.37e18 * 0.004**2 / lam_um**2, rel=1e-2 + ) + + +def test_write_units_no_laser_keys_without_laser_section(tmp_path: Path) -> None: + deck = tmp_path / "no_laser" + deck.write_text("simulation { n0 = 9.05e21, }\n") + quants = BaseOsiris({"osiris": {"deck": str(deck)}}).write_units() + for key in ("w_laser", "laser_wavelength", "laser_a0", "laser_intensity"): + assert key not in quants + + +def test_write_units_omits_temperature_keys() -> None: + quants = BaseOsiris({"osiris": {"deck": str(SRS_DECK)}}).write_units() + # OSIRIS has no single global temperature, so temperature-dependent + # 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() == {} diff --git a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml index 4134829c..9539860e 100644 --- a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml +++ b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_array_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.0 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.1 dx: 0.785 ion_charge: @@ -16507,6 +16507,7 @@ grid: vmax: 6.4 species_params: electron: + T0: 1.0 charge: -1.0 charge_to_mass: -1.0 mass: 1.0 @@ -17723,17 +17724,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.07623654008768617 micron + beta: 0.062561188941867 + box_length: 0.05390737447020297 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.004622577634581562 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml index ebca6209..f83b3fe6 100644 --- a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml +++ b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_derived_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.0 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.1 dx: 0.785 max_steps: 105 @@ -111,17 +111,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.07623654008768617 micron + beta: 0.062561188941867 + box_length: 0.05390737447020297 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.004622577634581562 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml index dca5edde..f1e7ae4b 100644 --- a/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml +++ b/tests/test_vlasov1d/test_config_regression/fokker_planck_conservation_units.yml @@ -1,13 +1,14 @@ T0: 2000 electron_volt -beta: 0.088474881879774 -box_length: 0.07623654008768617 micron +beta: 0.062561188941867 +box_length: 0.05390737447020297 micron box_width: inf -c_light: 11.302642950785 -logLambda_ee: -11.781233290653 +c_light: 15.984350951661 +logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter -nuee: -574949910190.1326 hertz +nuee: 352401683370.27344 hertz +nuee_norm: 0.00016128753860756 sim_duration: 0.004622577634581562 picosecond tp0: 0.45768095391896646 femtosecond -v0: 26524102.30999721 meter / second +v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second -x0: 12.139576447083783 nanometer +x0: 8.58397682646544 nanometer diff --git a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml index 927d9760..e899d7fa 100644 --- a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml +++ b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_array_config.yml @@ -29,7 +29,7 @@ drivers: ex: {} ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.5 dx: 9.81746875 ion_charge: @@ -28711,10 +28711,12 @@ grid: vmax: 0.005 species_params: electron: + T0: 1.0 charge: -1.0 charge_to_mass: -1.0 mass: 1.0 ion: + T0: 0.01 charge: 10.0 charge_to_mass: 0.00054466230936819 mass: 18360.0 @@ -38873,17 +38875,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 3.8137571970393944 micron + beta: 0.062561188941867 + box_length: 2.696733575825556 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 2.288633610071792 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml index e7c61532..9fe52c26 100644 --- a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml +++ b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_derived_config.yml @@ -29,7 +29,7 @@ drivers: ex: {} ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.5 dx: 9.81746875 max_steps: 10005 @@ -120,17 +120,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 3.8137571970393944 micron + beta: 0.062561188941867 + box_length: 2.696733575825556 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 2.288633610071792 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml index b05d7d96..e018602a 100644 --- a/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml +++ b/tests/test_vlasov1d/test_config_regression/multispecies_ion_acoustic_units.yml @@ -1,13 +1,14 @@ T0: 2000 electron_volt -beta: 0.088474881879774 -box_length: 3.8137571970393944 micron +beta: 0.062561188941867 +box_length: 2.696733575825556 micron box_width: inf -c_light: 11.302642950785 -logLambda_ee: -11.781233290653 +c_light: 15.984350951661 +logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter -nuee: -574949910190.1326 hertz +nuee: 352401683370.27344 hertz +nuee_norm: 0.00016128753860756 sim_duration: 2.288633610071792 picosecond tp0: 0.45768095391896646 femtosecond -v0: 26524102.30999721 meter / second +v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second -x0: 12.139576447083783 nanometer +x0: 8.58397682646544 nanometer diff --git a/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml b/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml index 8912206a..0924eec1 100644 --- a/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml +++ b/tests/test_vlasov1d/test_config_regression/resonance_array_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.1598 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.25 dx: 0.654375 ion_charge: @@ -20747,6 +20747,7 @@ grid: vmax: 6.4 species_params: electron: + T0: 1.0 charge: -1.0 charge_to_mass: -1.0 mass: 1.0 @@ -23327,17 +23328,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.25420273080193445 micron + beta: 0.062561188941867 + box_length: 0.17974847474618633 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.21980127811958367 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml b/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml index 66794372..6e5e2484 100644 --- a/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml +++ b/tests/test_vlasov1d/test_config_regression/resonance_derived_config.yml @@ -38,7 +38,7 @@ drivers: w0: 1.1598 ey: {} grid: - beta: 0.088474881879774 + beta: 0.062561188941867 dt: 0.25 dx: 0.654375 max_steps: 1925 @@ -119,17 +119,18 @@ terms: units: derived: T0: 2000 electron_volt - beta: 0.088474881879774 - box_length: 0.25420273080193445 micron + beta: 0.062561188941867 + box_length: 0.17974847474618633 micron box_width: inf - c_light: 11.302642950785 - logLambda_ee: -11.781233290653 + c_light: 15.984350951661 + logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter - nuee: -574949910190.1326 hertz + nuee: 352401683370.27344 hertz + nuee_norm: 0.00016128753860756 sim_duration: 0.21980127811958367 picosecond tp0: 0.45768095391896646 femtosecond - v0: 26524102.30999721 meter / second + v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second - x0: 12.139576447083783 nanometer + x0: 8.58397682646544 nanometer normalizing_density: 1.5e21/cc normalizing_temperature: 2000eV diff --git a/tests/test_vlasov1d/test_config_regression/resonance_units.yml b/tests/test_vlasov1d/test_config_regression/resonance_units.yml index cc18e176..22e9776c 100644 --- a/tests/test_vlasov1d/test_config_regression/resonance_units.yml +++ b/tests/test_vlasov1d/test_config_regression/resonance_units.yml @@ -1,13 +1,14 @@ T0: 2000 electron_volt -beta: 0.088474881879774 -box_length: 0.25420273080193445 micron +beta: 0.062561188941867 +box_length: 0.17974847474618633 micron box_width: inf -c_light: 11.302642950785 -logLambda_ee: -11.781233290653 +c_light: 15.984350951661 +logLambda_ee: 7.221022858202 n0: 1.5e+21 / cubic_centimeter -nuee: -574949910190.1326 hertz +nuee: 352401683370.27344 hertz +nuee_norm: 0.00016128753860756 sim_duration: 0.21980127811958367 picosecond tp0: 0.45768095391896646 femtosecond -v0: 26524102.30999721 meter / second +v0: 18755372.608284798 meter / second wp0: 2184928150138955.5 radian / second -x0: 12.139576447083783 nanometer +x0: 8.58397682646544 nanometer diff --git a/tests/test_vlasov1d/test_fp_momentum_conservation.py b/tests/test_vlasov1d/test_fp_momentum_conservation.py new file mode 100644 index 00000000..a50d9f66 --- /dev/null +++ b/tests/test_vlasov1d/test_fp_momentum_conservation.py @@ -0,0 +1,83 @@ +# Copyright (c) Ergodic LLC 2026 +# research@ergodic.io +""" +Momentum conservation of the Dougherty operator where n(x) != 1. + +The Dougherty drag must center on the *mean velocity* u(x) = ∫v f dv / n(x). +A drag centered on the raw first moment ∫v f dv = n*u (the historical bug) +breaks momentum conservation exactly where the density deviates from 1, e.g. +for bump-on-tail or large-amplitude density perturbations with collisions on. + +This test applies the collision operator in isolation (no advection, no +fields) to a drifting Maxwellian with a strong density modulation and checks +per-cell conservation of density, momentum, and energy. +""" + +import jax + +jax.config.update("jax_enable_x64", True) + +import numpy as np +import pytest +from jax import numpy as jnp + +from adept._vlasov1d.solvers.pushers.fokker_planck import Collisions + + +def _make_cfg(nx, nv, vmax, fp_type): + dv = 2.0 * vmax / nv + v = jnp.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) + return { + "grid": { + "species_grids": {"electron": {"v": v, "dv": dv, "nv": nv, "vmax": vmax}}, + "species_params": {"electron": {"charge": -1.0, "mass": 1.0, "charge_to_mass": -1.0, "T0": 1.0}}, + }, + "terms": { + "fokker_planck": {"is_on": True, "type": fp_type}, + "krook": {"is_on": False}, + }, + } + + +# Energy tolerance is per-scheme: Chang-Cooper preserves the discrete equilibrium +# to ~1e-6, plain central differencing only to ~1e-3 at this resolution. +@pytest.mark.parametrize("fp_type,energy_rtol", [("dougherty", 5e-3), ("chang_cooper_dougherty", 1e-5)]) +def test_dougherty_momentum_conservation_with_density_perturbation(fp_type, energy_rtol): + nx, nv, vmax = 16, 512, 6.4 + cfg = _make_cfg(nx, nv, vmax, fp_type) + v = np.array(cfg["grid"]["species_grids"]["electron"]["v"]) + dv = cfg["grid"]["species_grids"]["electron"]["dv"] + + # Strong density modulation and a finite drift: exactly the regime where + # centering the drag on n*u instead of u destroys momentum conservation. + x = np.linspace(0, 2 * np.pi, nx, endpoint=False) + n_prof = 1.0 + 0.5 * np.sin(x) + u0 = 0.5 + f = n_prof[:, None] * np.exp(-((v[None, :] - u0) ** 2) / 2.0) + f = f / (np.sum(np.exp(-((v - u0) ** 2) / 2.0)) * dv) + + f = jnp.array(f) + collisions = Collisions(cfg) + + nu = jnp.ones(nx) + dt = 0.1 + f_out = f + for _ in range(50): + f_out = collisions(nu, None, f_out, dt) + + n0_ = np.sum(np.array(f), axis=1) * dv + n1 = np.sum(np.array(f_out), axis=1) * dv + p0 = np.sum(np.array(f) * v[None, :], axis=1) * dv + p1 = np.sum(np.array(f_out) * v[None, :], axis=1) * dv + e0 = np.sum(np.array(f) * v[None, :] ** 2, axis=1) * dv + e1 = np.sum(np.array(f_out) * v[None, :] ** 2, axis=1) * dv + + # The historical n*u-centered drag gives O(50%) momentum errors here; + # the corrected operator conserves to discretization level (~1e-7). + np.testing.assert_allclose(n1, n0_, rtol=1e-10, err_msg="density not conserved") + np.testing.assert_allclose(p1, p0, rtol=1e-6, err_msg="momentum not conserved where n(x) != 1") + np.testing.assert_allclose(e1, e0, rtol=energy_rtol, err_msg="energy not conserved") + + # And the distribution must still be centered on u0, not n(x)*u0 + u_final = p1 / n1 + np.testing.assert_allclose(u_final, u0, atol=1e-6) diff --git a/tests/test_vlasov1d/test_units_boundary.py b/tests/test_vlasov1d/test_units_boundary.py new file mode 100644 index 00000000..e7af1c38 --- /dev/null +++ b/tests/test_vlasov1d/test_units_boundary.py @@ -0,0 +1,74 @@ +# Copyright (c) Ergodic LLC 2026 +# research@ergodic.io +""" +Units-boundary regression tests for the Vlasov-1D normalization layer. + +Every physics test in this suite works in dimensionless code units, which +validates the engine but not the dimensional dictionary: numeric inputs bypass +``normalize()`` entirely. These tests cross the units boundary — dimensional +input in, dimensional quantity out — and compare against *independent* physical +references (CODATA constants via scipy, and the NRL formulary), so a convention +error in ``normalization.py`` cannot hide. + +The engine convention under test: v0 = sqrt(T0/m_e) (RMS / standard-deviation +thermal speed), L0 = v0/wp0 = lambda_De, Maxwellian exp(-v^2/2) at T=1. +""" + +import numpy as np +from scipy import constants as csts + +from adept.normalization import electron_debye_normalization, normalize + +N0_STR, T0_STR = "1.5e21/cc", "2000eV" +N0_SI = 1.5e21 * 1e6 # 1/m^3 +T0_J = 2000.0 * csts.e + + +def test_debye_normalization_against_codata(): + """v0, L0, tau, and c_hat must match hand-computed CODATA values.""" + norm = electron_debye_normalization(N0_STR, T0_STR) + + v_th = np.sqrt(T0_J / csts.m_e) # sigma convention: sqrt(T/m) + wp0 = np.sqrt(N0_SI * csts.e**2 / (csts.epsilon_0 * csts.m_e)) + lambda_de = v_th / wp0 + + np.testing.assert_allclose(norm.v0.to("m/s").magnitude, v_th, rtol=1e-6) + np.testing.assert_allclose(norm.L0.to("m").magnitude, lambda_de, rtol=1e-6) + np.testing.assert_allclose(norm.tau.to("s").magnitude, 1.0 / wp0, rtol=1e-6) + np.testing.assert_allclose(norm.speed_of_light_norm(), csts.c / v_th, rtol=1e-6) + + +def test_debye_length_against_nrl_formulary(): + """lambda_De = 7.43e2 sqrt(T_eV/n_cc) cm (NRL formulary) — a truly external reference.""" + norm = electron_debye_normalization(N0_STR, T0_STR) + lambda_de_nrl_m = 7.43e2 * np.sqrt(2000.0 / 1.5e21) * 1e-2 + np.testing.assert_allclose(norm.L0.to("m").magnitude, lambda_de_nrl_m, rtol=1e-3) + + +def test_dimensional_string_round_trip(): + """String inputs must convert with L0 = lambda_De (not sqrt(2) lambda_De).""" + norm = electron_debye_normalization(N0_STR, T0_STR) + + v_th = np.sqrt(T0_J / csts.m_e) + wp0 = np.sqrt(N0_SI * csts.e**2 / (csts.epsilon_0 * csts.m_e)) + lambda_de = v_th / wp0 + + # x: 100 um -> 100e-6 / lambda_De code units + np.testing.assert_allclose(normalize("100um", norm, dim="x"), 100e-6 / lambda_de, rtol=1e-6) + # k: 1/um -> lambda_De / 1e-6 code units (k lambda_De) + np.testing.assert_allclose(normalize("1/um", norm, dim="k"), lambda_de / 1e-6, rtol=1e-6) + # t: 1 ps -> wp0 * 1e-12 + np.testing.assert_allclose(normalize("1ps", norm, dim="t"), wp0 * 1e-12, rtol=1e-6) + # numeric inputs pass through untouched + assert normalize(3.25, norm, dim="x") == 3.25 + + +def test_logged_collision_frequency_is_physical(): + """logLambda and nuee must be positive and match the NRL expression.""" + norm = electron_debye_normalization(N0_STR, T0_STR) + log_lambda = float(norm.logLambda_ee()) + # NRL: 23.5 - ln(n^1/2 T^-5/4) - [1e-5 + (ln T - 2)^2/16]^1/2 + expected = 23.5 - np.log(np.sqrt(1.5e21) * 2000.0**-1.25) - np.sqrt(1e-5 + (np.log(2000.0) - 2.0) ** 2 / 16.0) + np.testing.assert_allclose(log_lambda, expected, rtol=1e-10) + assert log_lambda > 0 + assert norm.approximate_ee_collision_frequency().to("Hz").magnitude > 0