feat: vlasov1d broadband laser driver + 1D super-Gaussian width normalization (vlasov1d & pic1d)

adept/_pic1d/helpers.py

@@ -16,7 +16,7 @@
 from scipy.special import gamma
 
 from adept._pic1d.simulation import PIC1DSimulation
-from adept._vlasov1d.helpers import gamma_3_over_m, gamma_5_over_m
+from adept._vlasov1d.helpers import gamma_1_over_m, gamma_3_over_m
 
 from .. import patched_mlflow as mlflow
 
@@ -32,7 +32,9 @@ def _inverse_cdf_supergaussian(
     sampling that exactly reproduces the requested moments in expectation.
     """
     v_thermal = np.sqrt(T0 / mass)
-    alpha = np.sqrt(3.0 * gamma_3_over_m(supergaussian_order) / gamma_5_over_m(supergaussian_order))
+    alpha = np.sqrt(
+        gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order)
+    )  # 1D width: variance = T0/mass for every m (see _vlasov1d.helpers)
 
     nv = max(8192, 8 * n_particles)
     v_grid = np.linspace(-vmax, vmax, nv)
@@ -56,7 +58,9 @@ def _random_supergaussian(
 ) -> np.ndarray:
     """Random velocity sampling via rejection on a bounded supergaussian."""
     v_thermal = np.sqrt(T0 / mass)
-    alpha = np.sqrt(3.0 * gamma_3_over_m(supergaussian_order) / gamma_5_over_m(supergaussian_order))
+    alpha = np.sqrt(
+        gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order)
+    )  # 1D width: variance = T0/mass for every m (see _vlasov1d.helpers)
 
     out = np.empty(n_particles)
     filled = 0

adept/_vlasov1d/datamodel.py

@@ -127,6 +127,67 @@ class IntensityWavelengthDriverConfig(BaseModel):
     intensity: str
     wavelength: str
     leftgoing: bool = False
+    dw0: float = 0.0
+    phase: float = 0.0
+
+
+class BroadbandIntensitiesConfig(BaseModel):
+    """Per-line intensity weights of a broadband (multi-color) driver.
+
+    The weights ``w_j`` set the line *amplitudes* as ``a_j = a0 * sqrt(w_j / sum_k w_k)``,
+    where ``a0`` is the monochromatic amplitude for ``base_intensity``; hence
+    ``sum_j a_j^2 = a0^2`` (same time-averaged power as a single line at
+    ``base_intensity``) and the per-line intensity is ``I_j = base_intensity * w_j / sum_k w_k``.
+    """
+
+    base_intensity: str  # e.g. "2.378e+14 W/cm^2"; the total (monochromatic-equivalent) intensity
+    init: Literal["uniform", "random"] = "uniform"
+    seed: int | None = None  # required for init: random
+    range: tuple[float, float] = (0.0, 2.0)  # uniform draw bounds for init: random
+
+    @model_validator(mode="after")
+    def check_seed(self) -> "BroadbandIntensitiesConfig":
+        """``init: random`` must be reproducible, so it needs a seed."""
+        if self.init == "random" and self.seed is None:
+            raise ValueError("intensities.init == 'random' requires intensities.seed")
+        return self
+
+
+class BroadbandPhasesConfig(BaseModel):
+    """Per-line spectral phases of a broadband (multi-color) driver."""
+
+    init: Literal["uniform", "random"] = "random"
+    seed: int | None = None  # required for init: random
+    range: tuple[float, float] = (0.0, 2.0 * 3.141592653589793)  # uniform draw bounds, radians
+    base_phase: float = 0.0  # the common phase for init: uniform
+
+    @model_validator(mode="after")
+    def check_seed(self) -> "BroadbandPhasesConfig":
+        """``init: random`` must be reproducible, so it needs a seed."""
+        if self.init == "random" and self.seed is None:
+            raise ValueError("phases.init == 'random' requires phases.seed")
+        return self
+
+
+class BroadbandConfig(BaseModel):
+    """Broadband (multi-color) laser driver: a comb of ``num_colors`` lines.
+
+    Line frequencies are ``w_j = w0 * (1 + d_j)`` with ``d_j`` spaced uniformly on
+    ``[-delta_omega, +delta_omega]`` (``num_colors: 1`` puts the single line exactly at
+    ``w0``). **``delta_omega`` is the HALF-width of the comb as a fraction of ``w0``;
+    the full bandwidth is ``2 * delta_omega``** -- a "0.25% bandwidth" run is
+    ``delta_omega: 0.00125``. All lines share the carrier wavenumber ``k0`` of
+    ``wavelength`` (an off-center line is not launched on its own dispersion branch),
+    which is accurate for a localized antenna (``source_type: point`` / a narrow
+    spatial envelope) and not for a source extended across the box.
+    """
+
+    num_colors: int = Field(ge=1)
+    delta_omega: float = Field(ge=0.0)  # half-width, fraction of w0
+    wavelength: str
+    intensities: BroadbandIntensitiesConfig
+    phases: BroadbandPhasesConfig
+    leftgoing: bool = False
 
 
 class AKWDriverConfig(BaseModel):
@@ -136,6 +197,7 @@ class AKWDriverConfig(BaseModel):
     k0: float | None = None
     w0: float | None = None
     dw0: float
+    phase: float = 0.0
 
     @model_validator(mode="after")
     def check_w_or_k(self) -> "AKWDriverConfig":
@@ -148,7 +210,7 @@ def check_w_or_k(self) -> "AKWDriverConfig":
 class EMDriverConfig(BaseModel):
     """One electromagnetic driver with parameters, envelope, and source geometry."""
 
-    params: IntensityWavelengthDriverConfig | AKWDriverConfig
+    params: IntensityWavelengthDriverConfig | AKWDriverConfig | BroadbandConfig
     envelope: SpaceTimeEnvelopeConfig
     source_type: Literal["extended", "point"] = "extended"
 

adept/_vlasov1d/helpers.py

@@ -2,19 +2,22 @@
 
 #  Copyright (c) Ergodic LLC 2023
 #  research@ergodic.io
+import math
 import os
 from time import time
 
+import equinox as eqx
 import numpy as np
 import xarray
 from diffrax import Solution
 from jax import numpy as jnp
+from jax import tree_util as jtu
 from matplotlib import pyplot as plt
 from scipy.special import gamma
 
 from adept._vlasov1d.simulation import SubspeciesDistributionSpec, Vlasov1DSimulation
 from adept._vlasov1d.storage import store_f, store_fields
-from adept.normalization import PlasmaNormalization
+from adept.normalization import UREG, PlasmaNormalization, normalize
 
 from .. import patched_mlflow as mlflow
 
@@ -24,6 +27,11 @@
 # g_5_m = np.squeeze(gamma_da.loc[{"gamma": "5/m"}].data)
 
 
+def gamma_1_over_m(m):
+    """Evaluate Gamma(1 / m) for super-Gaussian normalization."""
+    return gamma(1.0 / m)
+
+
 def gamma_3_over_m(m):
     """Evaluate Gamma(3 / m) for super-Gaussian normalization."""
     return gamma(3.0 / m)  # np.interp(m, m_ax, g_3_m)
@@ -72,11 +80,16 @@ def _initialize_supergaussian_distribution_(
     # Thermal velocity: v_t = sqrt(T/m)
     v_thermal = np.sqrt(T0 / mass)
 
-    # 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))
+    # 1D super-Gaussian width normalization: alpha = sqrt(Gamma(1/m)/Gamma(3/m))
+    # makes the realized VARIANCE equal T0/mass for every order m, i.e. a species
+    # labeled T0 is at temperature T0 for any m (alpha = sqrt(2) at m=2, unchanged).
+    # This is also the convention the Krook target (variance T0/mass) and the
+    # SuperGaussianDougherty temperature relation D = beta^(-2/m)*G(3/m)/G(1/m)
+    # already use. The previous alpha = sqrt(3*G(3/m)/G(5/m)) is the 3D-isotropic
+    # (Matte/DLM) normalization -- it fixes <v^2>_3D = 3*T0/mass -- and on a 1D
+    # axis it inflates the variance by F(m) = 3*G(3/m)^2/(G(5/m)*G(1/m))
+    # (x1.24 at m=3, x1.37 at m=4). See docs config.md.
+    alpha = np.sqrt(gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order))
 
     single_dist = -(np.power(np.abs((vax[None, :] - v0) / (alpha * v_thermal)), supergaussian_order))
 
@@ -161,10 +174,192 @@ def _initialize_total_distribution_(cfg, simulation: Vlasov1DSimulation):
     return species_distributions
 
 
+def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: PlasmaNormalization | None = None):
+    """getting amplitude (a), wave number (k) and angular frequency (w)
+    from intensity and wavelength (passed in as args to this function) defined in
+    'intensity_wavelength' type of configs"""
+
+    intensity = UREG.Quantity(intensity).to("W/m^2")
+    wavelength = UREG.Quantity(wavelength).to("nm")
+
+    e = UREG.e
+    m_e = UREG.m_e
+    eps0 = UREG.epsilon_0
+    c = UREG.c
+
+    # Standard a0 = eE0/(m_e c w0) — identical to HermiteSRS1D formula
+    a0_std = ((e * wavelength / (m_e * math.pi)) * (intensity / (2 * eps0 * c**5)) ** 0.5).to("").magnitude
+    # Vlasov normalization: a0_vlasov = a0_std / β  (β = v0/c)
+    a0 = a0_std * norm.speed_of_light_norm()
+
+    # k0 in Debye-length units: k0_vlasov = k_phys x v0/wp0
+    k0_phys = (2 * math.pi / wavelength).to("1/m")
+    k_sign = -1.0 if leftgoing else 1.0
+    k0 = k_sign * float((k0_phys * norm.L0).to("").magnitude)
+
+    # w0 normalized to wp0 (same normalization as Hermite)
+    w0_phys = (2 * math.pi * c / wavelength).to("1/s")
+    w0 = float((w0_phys * norm.tau).to("").magnitude)
+
+    return a0, k0, w0
+
+
+def plot_driver_spectra(cfg: dict, td: str, args: dict):
+    """Per-line intensity and phase vs frequency offset, for each multi-line driver.
+
+    Reads the LIVE driver objects out of `args["drivers"]` rather than re-deriving
+    the line set from the config's init/seed. That matters twice over: the plot cannot
+    drift if BroadbandDriver's construction changes, and it shows optimized line sets
+    from a backward pass (which never appear in the config) automatically. A
+    `BroadbandDriver` contributes its per-line arrays (`amplitudes`/`delta_omega`/
+    `phases`); plain `EMDriver`s contribute their scalars, so a hand-built list of
+    mono drivers still plots.
+
+    Per-line intensity needs no normalization constant. The driver builds amplitudes as
+    A_j = a0 * sqrt(w_j / sum_k w_k), so
+
+        I_j / I_base = A_j^2 / sum_k A_k^2
+
+    exactly, independent of a0 and of the plasma normalization. `base_intensity` is
+    read from the config only to put an absolute scale on the axis.
+
+    Panels: I_j (linear), log10 I_j (dynamic range -- the informative one once an
+    optimizer produces structure, since I ~ A^2 doubles the span), and phi_j.
+    Single-line (monochromatic) drivers are skipped.
+    """
+    drivers = (args or {}).get("drivers")
+    if drivers is None:
+        return
+    out_dir = os.path.join(td, "plots", "drivers")
+
+    for field in ("ex", "ey"):
+        dlist = getattr(drivers, field, None)
+        if not dlist:
+            continue
+
+        amp_list, dw_list, phase_list = [], [], []
+        for d in dlist:
+            if hasattr(d, "amplitudes"):  # BroadbandDriver: (N,) array leaves
+                amp_list.extend(np.asarray(d.amplitudes, dtype=float))
+                dw_list.extend(np.asarray(d.delta_omega, dtype=float) / float(d.w0))
+                phase_list.extend(np.asarray(d.phases, dtype=float))
+            else:  # plain EMDriver: scalar leaves
+                amp_list.append(float(d.a0))
+                dw_list.append(float(d.dw0) / float(d.w0))
+                phase_list.append(float(d.phase))
+        if len(amp_list) < 2:
+            continue  # monochromatic -> no spectrum to show
+
+        amp = np.asarray(amp_list)
+        dw = np.asarray(dw_list)  # dw_j/w0
+        phases = np.asarray(phase_list)
+
+        power = amp**2
+        frac = power / power.sum() if power.sum() > 0 else power
+        # absolute scale: base_intensity of the broadband driver in this field (whichever
+        # key it was given under -- a driver keyed '1' must not lose the axis scale)
+        base = None
+        for dcfg in (cfg.get("drivers", {}).get(field, {}) or {}).values():
+            ints = ((dcfg or {}).get("params", {}) or {}).get("intensities")
+            if isinstance(ints, dict) and ints.get("base_intensity") is not None:
+                base = ints["base_intensity"]
+                break
+
+        I_j, unit = frac, ""
+        if isinstance(base, str) and base.split():  # "2.378e+14 W/cm^2"
+            val, _, u = base.partition(" ")
+            try:
+                I_j, unit = frac * float(val), u.strip()
+            except ValueError:
+                pass
+        elif isinstance(base, (int, float)):
+            I_j = frac * float(base)
+
+        order = np.argsort(dw)
+        dw, I_j, phases = dw[order], I_j[order], phases[order]
+
+        # constrained_layout sizes the suptitle band to the actual text; do NOT pair
+        # it with tight_layout(rect=...) + suptitle(y=...), which reserve a fixed band
+        # and leave a gap when the title is shorter than the reservation.
+        fig, axes = plt.subplots(3, 1, figsize=(7.2, 8.6), sharex=True, constrained_layout=True)
+        bw_pct = (dw.max() - dw.min()) * 100.0
+        spacing = float(np.diff(np.sort(dw)).mean()) if len(dw) > 1 else 0.0
+        run_name = ((cfg.get("mlflow") or {}).get("run")) or ""
+
+        title = f"{field} driver — broadband line spectrum"
+        if run_name:
+            title += f"\n{run_name}"
+        title += (
+            f"\n{len(dlist)} lines   |   full width $\\Delta\\omega/\\omega_0$ = "
+            f"{bw_pct:.3g}%   |   spacing $\\delta\\omega/\\omega_0$ = {spacing:.3g}"
+        )
+        if base is not None:
+            title += f"\n$I_{{base}}$ = {base}"
+            if unit:
+                title += f"   |   $I_j$ = {I_j.mean():.4g} {unit} mean per line"
+        fig.suptitle(title, fontsize=9.5, linespacing=1.4)
+
+        axes[0].plot(dw, I_j, "o", ms=4, color="#1f77b4")
+        axes[0].set_ylabel("line intensity $I_j$" + (f"  [{unit}]" if unit else "  [$I_j/I_{base}$]"))
+        axes[0].set_ylim(bottom=0)
+        axes[0].annotate(
+            rf"$\Sigma_j I_j$ = {I_j.sum():.4g}", xy=(0.02, 0.06), xycoords="axes fraction", fontsize=8, color="0.35"
+        )
+
+        pos = I_j > 0
+        if pos.any():
+            axes[1].plot(dw[pos], np.log10(I_j[pos]), "o", ms=4, color="#d62728")
+            lo_, hi_ = np.log10(I_j[pos]).min(), np.log10(I_j[pos]).max()
+            if hi_ - lo_ < 0.1:
+                axes[1].set_ylim(lo_ - 0.5, hi_ + 0.5)
+            axes[1].annotate(
+                f"dynamic range: {10 ** (hi_ - lo_):.3g}x",
+                xy=(0.02, 0.88),
+                xycoords="axes fraction",
+                fontsize=8,
+                color="0.35",
+            )
+        if (~pos).any():  # an optimizer can drive lines to zero; don't hide them
+            floor = np.log10(I_j[pos]).min() if pos.any() else 0.0
+            axes[1].plot(dw[~pos], np.full(int((~pos).sum()), floor), "x", ms=6, color="0.5")
+            axes[1].annotate(
+                f"{int((~pos).sum())} line(s) at I=0 (x)",
+                xy=(0.02, 0.06),
+                xycoords="axes fraction",
+                fontsize=8,
+                color="0.4",
+            )
+        axes[1].set_ylabel(r"$\log_{10} I_j$")
+
+        axes[2].plot(dw, phases, "o", ms=4, color="#2ca02c")
+        axes[2].set_ylabel(r"phase $\phi_j$ [rad]")
+        axes[2].set_xlabel(r"$\delta\omega_j/\omega_0$")
+        axes[2].set_ylim(-0.25, 2 * np.pi + 0.25)
+        axes[2].set_yticks([0, np.pi / 2, np.pi, 3 * np.pi / 2, 2 * np.pi])
+        axes[2].set_yticklabels(["0", r"$\pi/2$", r"$\pi$", r"$3\pi/2$", r"$2\pi$"])
+
+        for ax in axes:
+            ax.grid(alpha=0.3)
+            ax.axvline(0.0, color="0.6", lw=0.8, ls="--")
+
+        os.makedirs(out_dir, exist_ok=True)
+        # no tight_layout here — constrained_layout (set on the figure) already sized
+        # the title band; calling both would re-reserve a fixed strip and reopen the gap
+        fig.savefig(os.path.join(out_dir, f"{field}-lines.png"), bbox_inches="tight", dpi=150)
+        plt.close(fig)
+
+
 def post_process(result: Solution, cfg: dict, td: str, args: dict):
     """Write binary output and diagnostic plots from a completed Vlasov-1D solve."""
     t0 = time()
 
+    # Driver line spectra (multi-line drivers only). Guarded: a diagnostics failure
+    # here must never cost a completed solve its binary output and field/dist plots.
+    try:
+        plot_driver_spectra(cfg, td, args)
+    except Exception as exc:
+        print(f"[post_process] driver spectrum plot skipped: {type(exc).__name__}: {exc}", flush=True)
+
     # Get species names for directory creation
     species_names = list(cfg["grid"]["species_grids"].keys())
 

adept/_vlasov1d/modules.py

@@ -337,8 +337,19 @@ def init_diffeqsolve(self):
 
     def __call__(self, trainable_modules: dict, args: dict | None = None):
         """Run the configured Vlasov-1D solve and return the raw Diffrax result."""
-        if args is None:
-            args = self.args
+        # Merge rather than replace, so a caller passing partial args cannot
+        # accidentally drop the "drivers" entry that the field pushers
+        # unconditionally read from args.
+        args = self.args | args if args is not None else self.args
+
+        # The trainable-module loop (lpse2d pattern): each module may transform the
+        # initial state and/or inject its (possibly optimizer-updated) leaves into
+        # args before the solve. Callers may pass None (the pre-module-era calling
+        # convention) -- treat it as "no trainable modules".
+        state = self.state
+        for name, module in (trainable_modules or {}).items():
+            state, args = module(state, args)
+
         grid = self.simulation.grid
         solver_result = diffeqsolve(
             terms=self.diffeqsolve_quants["terms"],
@@ -347,7 +358,7 @@ def __call__(self, trainable_modules: dict, args: dict | None = None):
             t1=self.time_quantities["t1"],
             max_steps=grid.max_steps,
             dt0=grid.dt,
-            y0=self.state,
+            y0=state,
             args=args,
             saveat=SaveAt(**self.diffeqsolve_quants["saveat"]),
             progress_meter=TqdmProgressMeter(refresh_steps=grid.max_steps // 100)