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

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

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:

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
+    # <v^4>/<v^2> = 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))

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

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

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 <v>, shape (...)
+            vbar: Mean velocity <v> = ∫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
 

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, :]

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

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,17 +78,21 @@ 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)
 
     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)

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