feat: osiris wrapper

.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/
 

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

adept/normalization.py

@@ -116,3 +116,55 @@ def laser_normalization(laser_wavelength_str, T0_str):
     return PlasmaNormalization(
         m0=UREG.m_e, q0=UREG.e, n0=ne_crit, T0=T0, L0=one_over_k, v0=1 * UREG.c, tau=1 / omega_laser
     )
+
+
+def _osiris_normalization(wp0, n0):
+    """
+    Core OSIRIS normalization built from an angular plasma frequency ``wp0`` and
+    its corresponding reference density ``n0``.
+
+    OSIRIS normalizes time to 1/wp0, length to the collisionless skin depth
+    c/wp0, and velocity to the speed of light.
+
+    Unit quantities are:
+        - L0 = c / wp0 (skin depth)
+        - v0 = c (speed of light)
+        - tau = 1 / wp0 (inverse plasma frequency)
+
+    There is no reference temperature: OSIRIS has no single global temperature
+    (species carry their own per-species thermal momenta), so ``T0`` is left
+    unset and temperature-dependent quantities are not defined under this
+    normalization.
+    """
+    wp0 = wp0.to("rad/s")
+    tau = 1 / wp0
+
+    v0 = 1 * UREG.c
+    x0 = (v0 / wp0).to("nm")
+
+    return PlasmaNormalization(m0=UREG.m_e, q0=UREG.e, n0=n0.to("1/cc"), T0=None, L0=x0, v0=v0, tau=tau)
+
+
+def skin_depth_normalization(n0_str):
+    """
+    OSIRIS normalization referenced to a plasma density (``simulation.n0``).
+
+    The reference plasma frequency is computed from the density,
+    wp0 = sqrt(n0 e^2 / (eps0 m_e)). See :func:`_osiris_normalization`.
+    """
+    n0 = UREG.Quantity(n0_str)
+    wp0 = ((n0 * UREG.e**2.0 / (UREG.m_e * UREG.epsilon_0)) ** 0.5).to("rad/s")
+    return _osiris_normalization(wp0, n0)
+
+
+def skin_depth_normalization_from_frequency(wp0_str):
+    """
+    OSIRIS normalization referenced to a plasma frequency (``simulation.omega_p0``).
+
+    The reference density is recovered from the frequency,
+    n0 = wp0^2 eps0 m_e / e^2, so the reported ``n0`` stays consistent with the
+    density-referenced form. See :func:`_osiris_normalization`.
+    """
+    wp0 = UREG.Quantity(wp0_str)
+    n0 = (wp0**2 * UREG.epsilon_0 * UREG.m_e / UREG.e**2.0).to("1/cc")
+    return _osiris_normalization(wp0, n0)

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}")

adept/osiris/base.py

@@ -0,0 +1,179 @@
+"""``BaseOsiris`` — the ADEPT module that drives OSIRIS."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from adept._base_ import ADEPTModule
+from adept.normalization import skin_depth_normalization, skin_depth_normalization_from_frequency
+from adept.osiris import deck as _deck
+from adept.osiris import post as _post
+from adept.osiris import runner as _runner
+
+
+class BaseOsiris(ADEPTModule):
+    """Wraps an external OSIRIS binary as an adept solver.
+
+    The OSIRIS native deck is the canonical simulation spec. The YAML
+    manifest layered on top supplies MLflow metadata, the binary path,
+    MPI rank count, and optional in-place deck overrides.
+    """
+
+    def __init__(self, cfg: dict) -> None:
+        super().__init__(cfg)
+        osiris_cfg = cfg.get("osiris", {})
+        deck_path = osiris_cfg.get("deck")
+        if not deck_path:
+            raise ValueError("BaseOsiris: cfg['osiris']['deck'] is required")
+
+        self._sections = _deck.parse_deck_file(deck_path)
+        overrides = osiris_cfg.pop("overrides", None) or {}
+        if overrides:
+            _deck.merge_overrides(self._sections, overrides)
+
+        # Surface the parsed (post-override) deck inside cfg so adept's
+        # log_params picks every parameter up as a flat MLflow param.
+        # The raw overrides dict is intentionally popped above to avoid
+        # confusing log_params/flatdict with integer keys; the applied
+        # values now live verbatim under cfg["deck"].
+        cfg["deck"] = _deck.deck_to_flat_dict(self._sections)
+
+    def write_units(self) -> dict:
+        """Derive physical reference scales from the deck's ``simulation`` section.
+
+        OSIRIS normalizes time to ``1/wp0``, length to the skin depth ``c/wp0``,
+        and velocity to ``c``, where ``wp0`` is the reference plasma frequency.
+        That reference is set in the deck's ``simulation`` section by either the
+        density ``n0`` (cm^-3) or the frequency ``omega_p0`` (rad/s); when both
+        are present OSIRIS uses ``n0``, so we do too. The returned dict mirrors
+        the density-derived keys the other adept solvers log to ``units.yaml`` so
+        OSIRIS runs are comparable in MLflow. OSIRIS has no single global
+        reference temperature (species carry their own per-species thermal
+        momenta), so the temperature-dependent keys (``T0``/``nuee``/
+        ``logLambda_ee``) are intentionally omitted.
+        """
+        sim = self._iter_first_section("simulation")
+        n0 = sim.get("n0")
+        omega_p0 = sim.get("omega_p0")
+        if n0 is not None:
+            norm = skin_depth_normalization(f"{n0} / cc")
+        elif omega_p0 is not None:
+            norm = skin_depth_normalization_from_frequency(f"{omega_p0} rad/s")
+        else:
+            return {}
+
+        quants: dict[str, Any] = {
+            "wp0": (1 / norm.tau).to("rad/s"),
+            "tp0": norm.tau.to("fs"),
+            "n0": norm.n0.to("1/cc"),
+            "v0": norm.v0.to("m/s"),
+            "x0": norm.L0.to("nm"),
+            "c_light": norm.speed_of_light_norm(),  # == 1.0; OSIRIS normalizes v to c
+            "beta": 1.0 / norm.speed_of_light_norm(),
+        }
+
+        space = self._iter_first_section("space")
+        xmin = self._first_array_value(space, "xmin")
+        xmax = self._first_array_value(space, "xmax")
+        if xmin is not None and xmax is not None:
+            quants["box_length"] = ((xmax[0] - xmin[0]) * norm.L0).to("micron")
+
+        time = self._iter_first_section("time")
+        tmax = time.get("tmax")
+        if tmax is not None:
+            tmin = time.get("tmin", 0.0)
+            quants["sim_duration"] = ((float(tmax) - float(tmin)) * norm.tau).to("ps")
+
+        self.cfg.setdefault("units", {})["derived"] = quants
+        return quants
+
+    def get_derived_quantities(self) -> dict:
+        """Lift a few useful scalars out of the deck for MLflow visibility."""
+        grid = dict(self._iter_first_section("grid"))
+        time = dict(self._iter_first_section("time"))
+        time_step = dict(self._iter_first_section("time_step"))
+        space = dict(self._iter_first_section("space"))
+
+        derived: dict[str, Any] = {}
+        nx = self._first_array_value(grid, "nx_p")
+        xmin = self._first_array_value(space, "xmin")
+        xmax = self._first_array_value(space, "xmax")
+        dt = time_step.get("dt")
+        tmax = time.get("tmax")
+
+        if nx and xmin is not None and xmax is not None:
+            derived["dx"] = [(xmax[d] - xmin[d]) / nx[d] for d in range(len(nx))]
+            if dt is not None and nx:
+                derived["cfl_ratio"] = float(dt) / min(derived["dx"])
+        if dt is not None and tmax is not None:
+            try:
+                derived["num_steps"] = int(float(tmax) / float(dt))
+            except (TypeError, ValueError):
+                pass
+        if derived:
+            self.cfg.setdefault("derived", {}).update(derived)
+        return derived
+
+    def get_solver_quantities(self) -> None:
+        return None
+
+    def init_state_and_args(self) -> dict:
+        return {}
+
+    def init_diffeqsolve(self) -> None:
+        return None
+
+    def init_modules(self) -> dict:
+        return {}
+
+    def __call__(self, trainable_modules: dict, args: dict) -> dict:
+        osiris_cfg = self.cfg.get("osiris", {})
+        binary = _runner.discover_binary(
+            osiris_cfg.get("binary"),
+            dim=self._infer_dim(),
+        )
+        mpi_ranks = int(osiris_cfg.get("mpi_ranks", 1))
+        run_root = osiris_cfg.get("run_root", "./checkpoints")
+        deck_text = _deck.render_deck(self._sections)
+        result = _runner.run_osiris(
+            deck_text,
+            binary=binary,
+            mpi_ranks=mpi_ranks,
+            run_root=run_root,
+            launcher=osiris_cfg.get("mpi_launcher", "srun"),
+            extra_mpi_args=osiris_cfg.get("extra_mpi_args"),
+        )
+        return {"solver result": result}
+
+    def post_process(self, run_output: dict, td: str) -> dict:
+        return _post.collect(run_output, self.cfg, td)
+
+    def vg(self, trainable_modules: dict, args: dict):
+        raise NotImplementedError("OSIRIS is not differentiable inside adept")
+
+    # --- helpers ----------------------------------------------------------
+
+    def _iter_first_section(self, name: str) -> dict:
+        for sec_name, params in self._sections:
+            if sec_name == name:
+                return params
+        return {}
+
+    @staticmethod
+    def _first_array_value(params: dict, base_key: str) -> list | None:
+        """Look up either ``base_key`` or ``base_key(...)`` and return as list."""
+        for k, v in params.items():
+            if k == base_key or k.startswith(base_key + "("):
+                if isinstance(v, list):
+                    return v
+                return [v]
+        return None
+
+    def _infer_dim(self) -> int | None:
+        """Best-effort: read dimensionality from ``grid.nx_p(1:D)``."""
+        grid = self._iter_first_section("grid")
+        for k in grid:
+            if k.startswith("nx_p"):
+                v = grid[k]
+                return len(v) if isinstance(v, list) else 1
+        return None