From 9cecd5fccf95594302b70453f3cd957ce580b5df Mon Sep 17 00:00:00 2001 From: Manhar Date: Thu, 30 Jul 2026 12:29:34 -0700 Subject: [PATCH 01/10] vlasov broadband forward --- adept/_vlasov1d/datamodel.py | 18 ++- adept/_vlasov1d/helpers.py | 28 +++- adept/_vlasov1d/simulation.py | 178 +++++++++++++++++++---- adept/_vlasov1d/solvers/pushers/field.py | 8 +- 4 files changed, 194 insertions(+), 38 deletions(-) diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 75d92293..2d99765d 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -102,7 +102,7 @@ class SaveConfig(BaseModel): fields: dict[str, TimeSaveConfig] - +#%% class IntensityWavelengthDriverConfig(BaseModel): """Laser driver parameters specified by physical intensity and wavelength.""" @@ -110,6 +110,15 @@ class IntensityWavelengthDriverConfig(BaseModel): wavelength: str leftgoing: bool = False +class BroadbandConfig(BaseModel): + """Broadband laser driver parameters specified by intensity and wavelength configuration (dicts)""" + + num_colors: int + delta_omega: float + wavelength: str + intensities: dict + phases: dict + leftgoing: bool = False class AKWDriverConfig(BaseModel): """Laser driver parameters specified directly as amplitude, wavenumber, and frequency.""" @@ -125,12 +134,12 @@ def check_w_or_k(self) -> "AKWDriverConfig": if self.k0 is None and self.w0 is None: raise ValueError("You must specify at least one of k0 or w0.") return self - +#%% 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" @@ -169,9 +178,6 @@ class FokkerPlanckConfig(BaseModel): type: str time: EnvelopeConfig space: EnvelopeConfig - # Super-Gaussian exponent of the operator's equilibrium (only used by - # type: super_gaussian; m=2 is Maxwellian) - m: float = Field(default=2.0, ge=1.0) class KrookConfig(BaseModel): diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index 534047ac..0b9ed521 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -3,6 +3,7 @@ # Copyright (c) Ergodic LLC 2023 # research@ergodic.io import os +import math from time import time import numpy as np @@ -14,7 +15,7 @@ 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 @@ -160,6 +161,31 @@ def _initialize_total_distribution_(cfg, simulation: Vlasov1DSimulation): return species_distributions +def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: PlasmaNormalization | None = None): + # encapsulate the logic into a separate function here + 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 post_process(result: Solution, cfg: dict, td: str, args: dict): """Write binary output and diagnostic plots from a completed Vlasov-1D solve.""" diff --git a/adept/_vlasov1d/simulation.py b/adept/_vlasov1d/simulation.py index 6158f87f..291e431e 100644 --- a/adept/_vlasov1d/simulation.py +++ b/adept/_vlasov1d/simulation.py @@ -1,16 +1,21 @@ """Domain objects that represent a configured Vlasov-1D simulation.""" -import math import warnings +from typing import List import equinox as eqx -import jax +import jax +import jax.numpy as jnp +import numpy as np +from jaxtyping import Array +from jax import tree_util as jtu from adept._vlasov1d.datamodel import ( AKWDriverConfig, EMDriverConfig, EMDriverSetConfig, IntensityWavelengthDriverConfig, + BroadbandConfig, SpeciesComponentConfig, SpeciesConfig, ) @@ -28,25 +33,29 @@ SpaceTimeEnvelopeFunction, UniformFunction, ) -from adept.normalization import UREG, PlasmaNormalization, normalize - +from adept.normalization import PlasmaNormalization +#%% # the place from where config is picked up from class EMDriver(eqx.Module): """Normalized electromagnetic driver used by longitudinal or transverse sources.""" a0: float k0: float w0: float + phase: float dw0: float envelope: SpaceTimeEnvelopeFunction is_point_source: bool = False @staticmethod - def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> "EMDriver": + # TODO: figure out new class from BaseVlasov with init_modules and vg (see for a way to ingest trainabale and non-trainable modules) + def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> List["EMDriver"]: """Convert user driver configuration into normalized solver parameters.""" envelope = SpaceTimeEnvelopeFunction.from_config(cfg.envelope, norm) params = cfg.params + # local import avoids the simulation <-> helpers module cycle + from adept._vlasov1d.helpers import get_akw_from_intensity_wavelength match cfg.params: case AKWDriverConfig(): @@ -59,36 +68,141 @@ def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> k0, w0 = params.k0, params.w0 is_point = cfg.source_type == "point" - return EMDriver(params.a0, k0, w0, params.dw0, envelope, is_point_source=is_point) + return [EMDriver(params.a0, k0, w0, params.phase, params.dw0, envelope, is_point_source=is_point)] case IntensityWavelengthDriverConfig(intensity=intensity, wavelength=wavelength, leftgoing=leftgoing): - intensity = UREG.Quantity(intensity).to("W/m^2") - wavelength = UREG.Quantity(wavelength).to("nm") + a0, k0, w0 = get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm) - e = UREG.e - m_e = UREG.m_e - eps0 = UREG.epsilon_0 - c = UREG.c + dw0 = params.dw0 + is_point = cfg.source_type == "point" + return [EMDriver(a0, k0, w0, params.phase, dw0, envelope, is_point_source=is_point)] - # 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() + case BroadbandConfig(intensities=intensities, wavelength=wavelength, leftgoing=leftgoing): + a0, k0, w0 = get_akw_from_intensity_wavelength(intensities['base_intensity'], wavelength, leftgoing, 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) + is_point = cfg.source_type == "point" + # need to see if this object remaining for later is required i.e. if the other code in the class is dead + broadband_driver = BroadbandDriver(params.model_dump(), a0, k0, w0, envelope, is_point) + return broadband_driver.driver_list - # 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) +#%% - dw0 = 0.0 # ??? +class BroadbandDriver(eqx.Module): + params: dict + a0: float + k0: float + w0: float + intensities: Array + delta_omega: Array + phases: Array + #ny: int #?? + envelope: SpaceTimeEnvelopeFunction + is_point_source: bool = False + driver_list: list + + def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): + self.params = cfg + self.a0 = a0 + self.k0 = k0 + self.w0 = w0 + self.envelope = envelope + self.is_point_source = is_point + + # intensities + if self.params["intensities"]["init"] == "random": #what do we want random to exactly be? + # note: random is not required for reproducing RK Follett results + self.intensities = jnp.array(np.random.uniform(0, 2, self.params["num_colors"])) + # 1: multiplicative shift as random amplitude broadband (changed from -1 to 1 TO 0 to 2)??? + # does a0 * self.intensities + # 2: what is the proper range (right now it is 0 to 2) + elif self.params["intensities"]["init"] == "uniform": + self.intensities = jnp.ones(self.params["num_colors"]) + else: + raise NotImplementedError( + f"Initialization type -- {self.params['intensities']['init']} -- not implemented" + ) + self.intensities = self.a0 * jnp.sqrt((self.intensities / jnp.sum(self.intensities))) # sqrt normalization to have the same power spectrum + # otherwise for uniform, power be N times the expected power + + # frequency shift + self.delta_omega = jnp.linspace( + -self.params["delta_omega"], self.params["delta_omega"], self.params["num_colors"] + ) * self.w0 + + # phases (to add: chirp very later (not to worry now)) + # opt and chirp + if self.params["intensities"]["init"] == "random": #what do we want random to exactly be? + phase_rng = np.random.default_rng(seed=self.params["phases"]["seed"]) + self.phases = jnp.array(phase_rng.uniform(-1, 1, self.params["num_colors"])) + self.phases = jnp.tanh(self.phases) * jnp.pi + elif self.params["intensities"]["init"] == "uniform": + self.phases = jnp.ones(self.params["num_colors"]) * self.params["phases"]['base_phase'] + else: + raise NotImplementedError( + f"Initialization type -- {self.params['phases']['init']} -- not implemented" + ) - is_point = cfg.source_type == "point" - return EMDriver(a0, k0, w0, dw0, envelope, is_point_source=is_point) + DriverList = [] + for a, dw, phase in zip(self.intensities, self.delta_omega, self.phases): + driver_obj = EMDriver(a, self.k0, self.w0, phase, dw, self.envelope, self.is_point_source) + DriverList.append(driver_obj) + + self.driver_list = DriverList + + def scale_intensities(self, intensities): # check if needed for the vlasov opt loops (was needed for lpse2d) + # reconfigures the intensities into weight-like values + if self.params["intensities"]["activation"] == "linear": + ints = 0.5 * (jnp.tanh(intensities) + 1.0) + elif self.params["intensities"]["activation"] == "log": + ints = 3 * (jnp.tanh(intensities) + 1.0) - 3 + ints = 10**ints + elif self.params["intensities"]["activation"] == "log-3wide": + ints = -1.5 * (jnp.tanh(intensities) + 1.0) # from 0 to -3 + ints = 10**ints + else: + raise NotImplementedError( + f"Amplitude Output type -- {self.params['intensities']['activation']} -- not implemented" + ) + return ints + + def get_partition_spec(self): # Maybe rewrite this to account for changes in the way broadband driver is injested now + """ + Get the partition spec for the model + + Only intensities and phases can be learned + + Returns + ------- + filter_spec : pytree with the same structure as the model + + """ + # figure out tracing arrays here + # jit and gradient boundary (figure if that might cause problems) + filter_spec = jtu.tree_map(lambda _: False, self) + + if self.params["intensities"]["learned"]: + filter_spec = eqx.tree_at(lambda tree: tree.intensities, filter_spec, replace=True) + + if self.params["phases"]["learned"]: + filter_spec = eqx.tree_at(lambda tree: tree.phases, filter_spec, replace=True) + + return filter_spec + + def __call__(self, state: dict, args: dict) -> tuple: + # intensities = self.scale_intensities(self.intensities) + # intensities = intensities / jnp.sum(intensities) + + ''' figure out what this really does in lpse2d --> is this needed in vlasov (only passed in additional paramteres) + but how does diffeqsolve even use these additional parameters''' + args["drivers"]["ey"] = { + "delta_omega": self.delta_omega, + "phases": jnp.tanh(self.phases) * jnp.pi, + "intensities": self.intensities, + } | self.envelope + # self.envelope configured differently in lpse2d --> connects to larger 'derived' parameters + + return state, args class EMDriverSet(eqx.Module): """Container for longitudinal (Ex) and transverse (Ey) driver lists.""" @@ -99,8 +213,16 @@ class EMDriverSet(eqx.Module): @staticmethod def from_config(cfg: EMDriverSetConfig, norm: PlasmaNormalization | None = None) -> "EMDriverSet": """Build normalized Ex and Ey driver lists from configuration.""" - ex = [EMDriver.from_config(ex_cfg, norm) for ex_cfg in cfg.ex.values()] - ey = [EMDriver.from_config(ey_cfg, norm) for ey_cfg in cfg.ey.values()] + ex = [] + for ex_cfg in cfg.ex.values(): + obj = EMDriver.from_config(ex_cfg, norm) + ex.extend(obj if isinstance(obj, list) else [obj]) + + ey = [] + for ey_cfg in cfg.ey.values(): + obj = EMDriver.from_config(ey_cfg, norm) + ey.extend(obj if isinstance(obj, list) else [obj]) + return EMDriverSet(ex, ey) diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index ff73262c..7bd40a9d 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -71,18 +71,20 @@ def __init__(self, xax, drivers: list[EMDriver], c: float = 0.0): self.point_source_masks.append(None) self.point_source_scales.append(None) +#%% def _single_driver_source(self, driver: EMDriver, mask, scale, current_time): ww = driver.w0 dw = driver.dw0 + phase = driver.phase w_total = ww + dw if driver.is_point_source: time_env = driver.envelope.time_envelope(current_time) - return scale * time_env * mask * jnp.sin(w_total * current_time) + return scale * time_env * mask * jnp.sin((w_total * current_time) + phase) else: kk = driver.k0 factor = driver.envelope(self.xax, current_time) - return -factor * w_total**2 * driver.a0 * jnp.sin(kk * self.xax - w_total * current_time) - + return -factor * w_total**2 * driver.a0 * jnp.sin((kk * self.xax - w_total * current_time) + phase) +#%% def __call__(self, t, args): """Evaluate the summed transverse current source at time t.""" total = jnp.zeros_like(self.xax) From 9b2c7bd2534ce491adca040864fd13ac227c4fa7 Mon Sep 17 00:00:00 2001 From: Manhar Date: Thu, 30 Jul 2026 21:39:51 -0700 Subject: [PATCH 02/10] range support for random and working datamodel --- adept/_vlasov1d/datamodel.py | 3 +++ adept/_vlasov1d/simulation.py | 17 +++++++++-------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 2d99765d..7eadd538 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -109,6 +109,8 @@ class IntensityWavelengthDriverConfig(BaseModel): intensity: str wavelength: str leftgoing: bool = False + dw0: float = 0.0 + phase: float = 0.0 class BroadbandConfig(BaseModel): """Broadband laser driver parameters specified by intensity and wavelength configuration (dicts)""" @@ -127,6 +129,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": diff --git a/adept/_vlasov1d/simulation.py b/adept/_vlasov1d/simulation.py index 291e431e..57a99caa 100644 --- a/adept/_vlasov1d/simulation.py +++ b/adept/_vlasov1d/simulation.py @@ -111,10 +111,9 @@ def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): # intensities if self.params["intensities"]["init"] == "random": #what do we want random to exactly be? # note: random is not required for reproducing RK Follett results - self.intensities = jnp.array(np.random.uniform(0, 2, self.params["num_colors"])) - # 1: multiplicative shift as random amplitude broadband (changed from -1 to 1 TO 0 to 2)??? - # does a0 * self.intensities - # 2: what is the proper range (right now it is 0 to 2) + int_lo, int_hi = self.params["intensities"].get("range", (0.0, 2.0)) + int_rng = np.random.default_rng(seed=self.params["intensities"]["seed"]) + self.intensities = jnp.array(int_rng.uniform(int_lo, int_hi, self.params["num_colors"])) elif self.params["intensities"]["init"] == "uniform": self.intensities = jnp.ones(self.params["num_colors"]) else: @@ -131,11 +130,13 @@ def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): # phases (to add: chirp very later (not to worry now)) # opt and chirp - if self.params["intensities"]["init"] == "random": #what do we want random to exactly be? + if self.params["phases"]["init"] == "random": #what do we want random to exactly be? + # Follett 2019: spectral phases drawn uniformly over (0, 2*pi) -- the default; + # override via phases.range + phase_lo, phase_hi = self.params["phases"].get("range", (0.0, 2.0 * np.pi)) phase_rng = np.random.default_rng(seed=self.params["phases"]["seed"]) - self.phases = jnp.array(phase_rng.uniform(-1, 1, self.params["num_colors"])) - self.phases = jnp.tanh(self.phases) * jnp.pi - elif self.params["intensities"]["init"] == "uniform": + self.phases = jnp.array(phase_rng.uniform(phase_lo, phase_hi, self.params["num_colors"])) + elif self.params["phases"]["init"] == "uniform": self.phases = jnp.ones(self.params["num_colors"]) * self.params["phases"]['base_phase'] else: raise NotImplementedError( From f251232b68b2dfb93f481e26e071ebe76013c6c0 Mon Sep 17 00:00:00 2001 From: Manhar Date: Tue, 11 Aug 2026 15:04:44 -0700 Subject: [PATCH 03/10] broadband driver spectrum plot --- adept/_vlasov1d/helpers.py | 122 +++++++++++++++++++++++ adept/_vlasov1d/solvers/pushers/field.py | 3 +- 2 files changed, 123 insertions(+), 2 deletions(-) diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index 0b9ed521..43d0acee 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -187,10 +187,132 @@ def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: Pl 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 `EMDriver` 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. + + 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 or len(dlist) < 2: + continue # absent, or monochromatic -> no spectrum to show + + amp = np.asarray([float(d.a0) for d in dlist]) + w0 = float(dlist[0].w0) + dw = np.asarray([float(d.dw0) for d in dlist]) / w0 # -> dw_j/w0 + phases = np.asarray([float(d.phase) for d in dlist]) + + power = amp**2 + frac = power / power.sum() if power.sum() > 0 else power + base = ((cfg.get("drivers", {}).get(field, {}) or {}).get("0", {}) or {}) \ + .get("params", {}).get("intensities", {}) + base = base.get("base_intensity") if isinstance(base, dict) else None + + 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: # noqa: BLE001 - diagnostics must not break post-processing + 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()) diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index 370d2802..7c1fe29a 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -71,7 +71,6 @@ def __init__(self, xax, drivers: list[EMDriver], c: float = 0.0): self.point_source_masks.append(None) self.point_source_scales.append(None) -#%% def _single_driver_source(self, driver: EMDriver, mask, scale, current_time): ww = driver.w0 dw = driver.dw0 @@ -84,7 +83,7 @@ def _single_driver_source(self, driver: EMDriver, mask, scale, current_time): kk = driver.k0 factor = driver.envelope(self.xax, current_time) return -factor * w_total**2 * driver.a0 * jnp.sin((kk * self.xax - w_total * current_time) + phase) -#%% + def __call__(self, t, args): """Evaluate the summed transverse current source at time t.""" total = jnp.zeros_like(self.xax) From 71f70e57f22404626a1f256e0f2972bbad41f961 Mon Sep 17 00:00:00 2001 From: Manhar Date: Wed, 12 Aug 2026 13:50:37 -0700 Subject: [PATCH 04/10] Shift to alpha_1D --- adept/_vlasov1d/helpers.py | 20 +++++++++++++++----- docs/source/solvers/vlasov1d/config.md | 2 +- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index 43d0acee..b7135b8a 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -25,6 +25,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) @@ -73,11 +78,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 - # / = 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 _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)) diff --git a/docs/source/solvers/vlasov1d/config.md b/docs/source/solvers/vlasov1d/config.md index 0b69f02e..8e556124 100644 --- a/docs/source/solvers/vlasov1d/config.md +++ b/docs/source/solvers/vlasov1d/config.md @@ -123,7 +123,7 @@ Each species is defined with a key starting with `species-` (e.g., `species-back | `noise_val` | float | Amplitude of noise | | `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`) | +| `m` | float | Exponent for super-Gaussian distribution $f \propto \exp[-\|v/(\alpha v_{th})\|^m]$. `2.0` is Maxwellian. $\alpha = \sqrt{\Gamma(1/m)/\Gamma(3/m)}$ (the 1D normalization), so the realized *variance* equals `T0/mass` for **every** `m` — a species labeled `T0` is at temperature `T0` regardless of `m`, consistent with the Krook target and the `super_gaussian` collision operator. (Before this change $\alpha$ was the 3D-isotropic Matte/DLM value $\sqrt{3\Gamma(3/m)/\Gamma(5/m)}$, which fixes $\langle v^4\rangle/\langle v^2\rangle = 3\,T_0/\mathrm{mass}$ and inflates the 1D variance by $F(m)=3\Gamma(3/m)^2/[\Gamma(5/m)\Gamma(1/m)]$: ×1.24 at `m: 3`, ×1.37 at `m: 4`.) | | `basis` | string | Spatial profile type (see below) | #### Basis Types From fcc6125a9ecd9018f17022f8b18a71d32d91b450 Mon Sep 17 00:00:00 2001 From: Manhar Date: Wed, 12 Aug 2026 14:34:29 -0700 Subject: [PATCH 05/10] pic1d shift to alpha_1D --- adept/_pic1d/helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/adept/_pic1d/helpers.py b/adept/_pic1d/helpers.py index e85c1ce5..f5aee320 100644 --- a/adept/_pic1d/helpers.py +++ b/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,7 @@ 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 +56,7 @@ 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 From e21e5f6a92c12b42d4931d7f372a9c1036eacbc9 Mon Sep 17 00:00:00 2001 From: Manhar Date: Wed, 12 Aug 2026 15:34:59 -0700 Subject: [PATCH 06/10] added missing defauly phase argument to absorbing_wave failing test --- tests/test_vlasov1d/test_absorbing_wave.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_vlasov1d/test_absorbing_wave.py b/tests/test_vlasov1d/test_absorbing_wave.py index 659da30d..b41ec4b9 100644 --- a/tests/test_vlasov1d/test_absorbing_wave.py +++ b/tests/test_vlasov1d/test_absorbing_wave.py @@ -49,7 +49,7 @@ def test_absorbing_boundaries(): time_env = EnvelopeFunction(center=40.0, width=30.0, rise=5.0, baseline=0.0, bump_height=1.0, is_trough=False) space_env = EnvelopeFunction(center=800.0, width=50.0, rise=10.0, baseline=0.0, bump_height=1.0, is_trough=False) envelope = SpaceTimeEnvelopeFunction(time_envelope=time_env, space_envelope=space_env) - ey_driver = EMDriver(a0=1.0e-4, k0=-1.4, w0=15.82, dw0=0.0, envelope=envelope) + ey_driver = EMDriver(a0=1.0e-4, k0=-1.4, w0=15.82, phase=0.0, dw0=0.0, envelope=envelope) drivers = [ey_driver] args = {} From 39e849cf8389a50ecc25b2f8ae9031f3abbe3e8b Mon Sep 17 00:00:00 2001 From: Manhar Date: Wed, 12 Aug 2026 18:07:13 -0700 Subject: [PATCH 07/10] Small changes for passing linting checks --- adept/_pic1d/helpers.py | 8 ++- adept/_vlasov1d/datamodel.py | 10 +++- adept/_vlasov1d/helpers.py | 47 +++++++++------ adept/_vlasov1d/simulation.py | 75 ++++++++++++------------ adept/_vlasov1d/solvers/pushers/field.py | 2 +- 5 files changed, 81 insertions(+), 61 deletions(-) diff --git a/adept/_pic1d/helpers.py b/adept/_pic1d/helpers.py index f5aee320..fecb5deb 100644 --- a/adept/_pic1d/helpers.py +++ b/adept/_pic1d/helpers.py @@ -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(gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order)) # 1D width: variance = T0/mass for every m (see _vlasov1d.helpers) + 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(gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order)) # 1D width: variance = T0/mass for every m (see _vlasov1d.helpers) + 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 diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 62e82b11..8368419b 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -120,7 +120,8 @@ class SaveConfig(BaseModel): fields: dict[str, TimeSaveConfig] -#%% + +# %% class IntensityWavelengthDriverConfig(BaseModel): """Laser driver parameters specified by physical intensity and wavelength.""" @@ -130,6 +131,7 @@ class IntensityWavelengthDriverConfig(BaseModel): dw0: float = 0.0 phase: float = 0.0 + class BroadbandConfig(BaseModel): """Broadband laser driver parameters specified by intensity and wavelength configuration (dicts)""" @@ -140,6 +142,7 @@ class BroadbandConfig(BaseModel): phases: dict leftgoing: bool = False + class AKWDriverConfig(BaseModel): """Laser driver parameters specified directly as amplitude, wavenumber, and frequency.""" @@ -155,7 +158,10 @@ def check_w_or_k(self) -> "AKWDriverConfig": if self.k0 is None and self.w0 is None: raise ValueError("You must specify at least one of k0 or w0.") return self -#%% + + +# %% + class EMDriverConfig(BaseModel): """One electromagnetic driver with parameters, envelope, and source geometry.""" diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index b7135b8a..dc387839 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -2,8 +2,8 @@ # Copyright (c) Ergodic LLC 2023 # research@ergodic.io -import os import math +import os from time import time import numpy as np @@ -171,6 +171,7 @@ def _initialize_total_distribution_(cfg, simulation: Vlasov1DSimulation): return species_distributions + def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: PlasmaNormalization | None = None): # encapsulate the logic into a separate function here intensity = UREG.Quantity(intensity).to("W/m^2") @@ -235,12 +236,13 @@ def plot_driver_spectra(cfg: dict, td: str, args: dict): power = amp**2 frac = power / power.sum() if power.sum() > 0 else power - base = ((cfg.get("drivers", {}).get(field, {}) or {}).get("0", {}) or {}) \ - .get("params", {}).get("intensities", {}) + base = ( + ((cfg.get("drivers", {}).get(field, {}) or {}).get("0", {}) or {}).get("params", {}).get("intensities", {}) + ) base = base.get("base_intensity") if isinstance(base, dict) else None I_j, unit = frac, "" - if isinstance(base, str) and base.split(): # "2.378e+14 W/cm^2" + 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() @@ -255,8 +257,7 @@ def plot_driver_spectra(cfg: dict, td: str, args: dict): # 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) + 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 "" @@ -264,8 +265,10 @@ def plot_driver_spectra(cfg: dict, td: str, args: dict): 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}") + 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: @@ -275,8 +278,9 @@ def plot_driver_spectra(cfg: dict, td: str, args: dict): 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") + 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(): @@ -284,13 +288,23 @@ def plot_driver_spectra(cfg: dict, td: str, args: dict): 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") + 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].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") @@ -319,9 +333,8 @@ def post_process(result: Solution, cfg: dict, td: str, args: dict): # 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: # noqa: BLE001 - diagnostics must not break post-processing - print(f"[post_process] driver spectrum plot skipped: {type(exc).__name__}: {exc}", - flush=True) + 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()) diff --git a/adept/_vlasov1d/simulation.py b/adept/_vlasov1d/simulation.py index d69ece0b..338aaead 100644 --- a/adept/_vlasov1d/simulation.py +++ b/adept/_vlasov1d/simulation.py @@ -1,21 +1,20 @@ """Domain objects that represent a configured Vlasov-1D simulation.""" import warnings -from typing import List import equinox as eqx -import jax +import jax import jax.numpy as jnp import numpy as np -from jaxtyping import Array from jax import tree_util as jtu +from jaxtyping import Array from adept._vlasov1d.datamodel import ( AKWDriverConfig, + BroadbandConfig, EMDriverConfig, EMDriverSetConfig, IntensityWavelengthDriverConfig, - BroadbandConfig, SpeciesComponentConfig, SpeciesConfig, StochasticDriverConfig, @@ -36,6 +35,7 @@ ) from adept.normalization import PlasmaNormalization + class EMDriver(eqx.Module): """Normalized electromagnetic driver used by longitudinal or transverse sources.""" @@ -48,8 +48,7 @@ class EMDriver(eqx.Module): is_point_source: bool = False @staticmethod - # TODO: figure out new class from BaseVlasov with init_modules and vg (see for a way to ingest trainabale and non-trainable modules) - def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> List["EMDriver"]: + def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> list["EMDriver"]: """Convert user driver configuration into normalized solver parameters.""" envelope = SpaceTimeEnvelopeFunction.from_config(cfg.envelope, norm) @@ -78,25 +77,28 @@ def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> return [EMDriver(a0, k0, w0, params.phase, dw0, envelope, is_point_source=is_point)] case BroadbandConfig(intensities=intensities, wavelength=wavelength, leftgoing=leftgoing): - a0, k0, w0 = get_akw_from_intensity_wavelength(intensities['base_intensity'], wavelength, leftgoing, norm) + a0, k0, w0 = get_akw_from_intensity_wavelength( + intensities["base_intensity"], wavelength, leftgoing, norm + ) is_point = cfg.source_type == "point" # need to see if this object remaining for later is required i.e. if the other code in the class is dead broadband_driver = BroadbandDriver(params.model_dump(), a0, k0, w0, envelope, is_point) return broadband_driver.driver_list + class BroadbandDriver(eqx.Module): params: dict a0: float k0: float - w0: float + w0: float intensities: Array delta_omega: Array phases: Array envelope: SpaceTimeEnvelopeFunction is_point_source: bool = False driver_list: list - + def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): self.params = cfg self.a0 = a0 @@ -106,49 +108,43 @@ def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): self.is_point_source = is_point # intensities - if self.params["intensities"]["init"] == "random": #what do we want random to exactly be? - # note: random is not required for reproducing RK Follett results + if self.params["intensities"]["init"] == "random": int_lo, int_hi = self.params["intensities"].get("range", (0.0, 2.0)) int_rng = np.random.default_rng(seed=self.params["intensities"]["seed"]) self.intensities = jnp.array(int_rng.uniform(int_lo, int_hi, self.params["num_colors"])) elif self.params["intensities"]["init"] == "uniform": self.intensities = jnp.ones(self.params["num_colors"]) else: - raise NotImplementedError( - f"Initialization type -- {self.params['intensities']['init']} -- not implemented" - ) - self.intensities = self.a0 * jnp.sqrt((self.intensities / jnp.sum(self.intensities))) # sqrt normalization to have the same power spectrum + raise NotImplementedError(f"Initialization type -- {self.params['intensities']['init']} -- not implemented") + self.intensities = self.a0 * jnp.sqrt( + self.intensities / jnp.sum(self.intensities) + ) # sqrt normalization to have the same power spectrum # otherwise for uniform, power be N times the expected power # frequency shift - self.delta_omega = jnp.linspace( - -self.params["delta_omega"], self.params["delta_omega"], self.params["num_colors"] - ) * self.w0 - - # phases (to add: chirp very later (not to worry now)) - # opt and chirp - if self.params["phases"]["init"] == "random": #what do we want random to exactly be? - # Follett 2019: spectral phases drawn uniformly over (0, 2*pi) -- the default; - # override via phases.range + self.delta_omega = ( + jnp.linspace(-self.params["delta_omega"], self.params["delta_omega"], self.params["num_colors"]) * self.w0 + ) + + if self.params["phases"]["init"] == "random": + # Spectral phases drawn uniformly over (0, 2*pi) -- the default; + # Override via phases.range phase_lo, phase_hi = self.params["phases"].get("range", (0.0, 2.0 * np.pi)) phase_rng = np.random.default_rng(seed=self.params["phases"]["seed"]) self.phases = jnp.array(phase_rng.uniform(phase_lo, phase_hi, self.params["num_colors"])) elif self.params["phases"]["init"] == "uniform": - self.phases = jnp.ones(self.params["num_colors"]) * self.params["phases"]['base_phase'] + self.phases = jnp.ones(self.params["num_colors"]) * self.params["phases"]["base_phase"] else: - raise NotImplementedError( - f"Initialization type -- {self.params['phases']['init']} -- not implemented" - ) + raise NotImplementedError(f"Initialization type -- {self.params['phases']['init']} -- not implemented") DriverList = [] - for a, dw, phase in zip(self.intensities, self.delta_omega, self.phases): + for a, dw, phase in zip(self.intensities, self.delta_omega, self.phases, strict=True): driver_obj = EMDriver(a, self.k0, self.w0, phase, dw, self.envelope, self.is_point_source) DriverList.append(driver_obj) - + self.driver_list = DriverList - def scale_intensities(self, intensities): # check if needed for the vlasov opt loops (was needed for lpse2d) - # reconfigures the intensities into weight-like values + def scale_intensities(self, intensities): if self.params["intensities"]["activation"] == "linear": ints = 0.5 * (jnp.tanh(intensities) + 1.0) elif self.params["intensities"]["activation"] == "log": @@ -164,7 +160,7 @@ def scale_intensities(self, intensities): # check if needed for the vlasov opt l return ints - def get_partition_spec(self): # Maybe rewrite this to account for changes in the way broadband driver is injested now + def get_partition_spec(self): """ Get the partition spec for the model @@ -187,20 +183,20 @@ def get_partition_spec(self): # Maybe rewrite this to account for changes in the return filter_spec - def __call__(self, state: dict, args: dict) -> tuple: + def __call__(self, state: dict, args: dict) -> tuple: # intensities = self.scale_intensities(self.intensities) # intensities = intensities / jnp.sum(intensities) - ''' figure out what this really does in lpse2d --> is this needed in vlasov (only passed in additional paramteres) - but how does diffeqsolve even use these additional parameters''' + """figure out what this does in lpse2d --> is this needed in vlasov (only passed in additional parameters) + but how does diffeqsolve even use these additional parameters""" args["drivers"]["ey"] = { "delta_omega": self.delta_omega, "phases": jnp.tanh(self.phases) * jnp.pi, "intensities": self.intensities, - } | self.envelope - # self.envelope configured differently in lpse2d --> connects to larger 'derived' parameters + } | self.envelope + + return state, args - return state, args class StochasticDriver(eqx.Module): """Band-limited, time-correlated (Ornstein-Uhlenbeck) longitudinal field driver. @@ -254,6 +250,7 @@ def __call__(self, t: float, x: jax.Array) -> jax.Array: phase = self.k_modes[:, None] * x[None, :] return jnp.sum(ar[:, None] * jnp.cos(phase) - ai[:, None] * jnp.sin(phase), axis=0) + class EMDriverSet(eqx.Module): """Container for longitudinal (Ex) and transverse (Ey) driver lists.""" diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index 7c1fe29a..bbbac6c3 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -83,7 +83,7 @@ def _single_driver_source(self, driver: EMDriver, mask, scale, current_time): kk = driver.k0 factor = driver.envelope(self.xax, current_time) return -factor * w_total**2 * driver.a0 * jnp.sin((kk * self.xax - w_total * current_time) + phase) - + def __call__(self, t, args): """Evaluate the summed transverse current source at time t.""" total = jnp.zeros_like(self.xax) From 4c45d6dc27dd1d1b8be4f9b113248bb1bf8ffc36 Mon Sep 17 00:00:00 2001 From: Manhar Date: Fri, 21 Aug 2026 01:04:39 -0700 Subject: [PATCH 08/10] Backward-pass compatible broadband driver --- adept/_vlasov1d/datamodel.py | 65 ++++++- adept/_vlasov1d/helpers.py | 51 ++++-- adept/_vlasov1d/modules.py | 10 +- adept/_vlasov1d/simulation.py | 93 +++++----- adept/_vlasov1d/solvers/pushers/field.py | 65 ++++--- configs/vlasov-1d/srs-broadband.yaml | 138 +++++++++++++++ docs/source/solvers/vlasov1d/config.md | 70 ++++++++ tests/test_vlasov1d/test_broadband_driver.py | 173 +++++++++++++++++++ 8 files changed, 573 insertions(+), 92 deletions(-) create mode 100644 configs/vlasov-1d/srs-broadband.yaml create mode 100644 tests/test_vlasov1d/test_broadband_driver.py diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 8368419b..af61fbc7 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -121,7 +121,6 @@ class SaveConfig(BaseModel): fields: dict[str, TimeSaveConfig] -# %% class IntensityWavelengthDriverConfig(BaseModel): """Laser driver parameters specified by physical intensity and wavelength.""" @@ -132,14 +131,62 @@ class IntensityWavelengthDriverConfig(BaseModel): 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 laser driver parameters specified by intensity and wavelength configuration (dicts)""" + """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 - delta_omega: float + num_colors: int = Field(ge=1) + delta_omega: float = Field(ge=0.0) # half-width, fraction of w0 wavelength: str - intensities: dict - phases: dict + intensities: BroadbandIntensitiesConfig + phases: BroadbandPhasesConfig leftgoing: bool = False @@ -160,9 +207,6 @@ def check_w_or_k(self) -> "AKWDriverConfig": return self -# %% - - class EMDriverConfig(BaseModel): """One electromagnetic driver with parameters, envelope, and source geometry.""" @@ -242,6 +286,9 @@ class FokkerPlanckConfig(BaseModel): type: str time: EnvelopeConfig space: EnvelopeConfig + # Super-Gaussian exponent of the operator's equilibrium (only used by + # type: super_gaussian; m=2 is Maxwellian) + m: float = Field(default=2.0, ge=1.0) class KrookConfig(BaseModel): diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index dc387839..63753948 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -10,6 +10,8 @@ import xarray from diffrax import Solution from jax import numpy as jnp +from jax import tree_util as jtu +import equinox as eqx from matplotlib import pyplot as plt from scipy.special import gamma @@ -39,7 +41,6 @@ def gamma_5_over_m(m): """Evaluate Gamma(5 / m) for super-Gaussian normalization.""" return gamma(5.0 / m) # np.interp(m, m_ax, g_5_m) - def _initialize_supergaussian_distribution_( nx: int, nv: int, @@ -173,7 +174,10 @@ def _initialize_total_distribution_(cfg, simulation: Vlasov1DSimulation): def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: PlasmaNormalization | None = None): - # encapsulate the logic into a separate function here + '''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") @@ -202,10 +206,13 @@ def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: Pl 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 `EMDriver` objects out of `args["drivers"]` rather than re-deriving + 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. + 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 @@ -226,20 +233,36 @@ def plot_driver_spectra(cfg: dict, td: str, args: dict): for field in ("ex", "ey"): dlist = getattr(drivers, field, None) - if not dlist or len(dlist) < 2: - continue # absent, or monochromatic -> no spectrum to show + if not dlist: + continue - amp = np.asarray([float(d.a0) for d in dlist]) - w0 = float(dlist[0].w0) - dw = np.asarray([float(d.dw0) for d in dlist]) / w0 # -> dw_j/w0 - phases = np.asarray([float(d.phase) for d in dlist]) + 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 - base = ( - ((cfg.get("drivers", {}).get(field, {}) or {}).get("0", {}) or {}).get("params", {}).get("intensities", {}) - ) - base = base.get("base_intensity") if isinstance(base, dict) else None + # 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" diff --git a/adept/_vlasov1d/modules.py b/adept/_vlasov1d/modules.py index 5d279cdc..cc0b182b 100644 --- a/adept/_vlasov1d/modules.py +++ b/adept/_vlasov1d/modules.py @@ -337,8 +337,14 @@ 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 + + for name, module in trainable_modules.items(): + state, args = module(state, args) + grid = self.simulation.grid solver_result = diffeqsolve( terms=self.diffeqsolve_quants["terms"], diff --git a/adept/_vlasov1d/simulation.py b/adept/_vlasov1d/simulation.py index 338aaead..b3d8f4b5 100644 --- a/adept/_vlasov1d/simulation.py +++ b/adept/_vlasov1d/simulation.py @@ -48,7 +48,7 @@ class EMDriver(eqx.Module): is_point_source: bool = False @staticmethod - def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> list["EMDriver"]: + def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> "list[EMDriver] | BroadbandDriver": """Convert user driver configuration into normalized solver parameters.""" envelope = SpaceTimeEnvelopeFunction.from_config(cfg.envelope, norm) @@ -77,27 +77,44 @@ def from_config(cfg: EMDriverConfig, norm: PlasmaNormalization | None = None) -> return [EMDriver(a0, k0, w0, params.phase, dw0, envelope, is_point_source=is_point)] case BroadbandConfig(intensities=intensities, wavelength=wavelength, leftgoing=leftgoing): - a0, k0, w0 = get_akw_from_intensity_wavelength( - intensities["base_intensity"], wavelength, leftgoing, norm - ) + a0, k0, w0 = get_akw_from_intensity_wavelength(intensities.base_intensity, wavelength, leftgoing, norm) is_point = cfg.source_type == "point" - # need to see if this object remaining for later is required i.e. if the other code in the class is dead broadband_driver = BroadbandDriver(params.model_dump(), a0, k0, w0, envelope, is_point) - return broadband_driver.driver_list + return broadband_driver + + case _: + raise NotImplementedError(f"Unsupported driver params type: {type(cfg.params).__name__}") class BroadbandDriver(eqx.Module): + """Multi-color (broadband) ey driver carrying per-line parameter arrays. + + Given the monochromatic amplitude ``a0`` for ``intensities.base_intensity`` and + per-line weights ``w_j`` (uniform or seeded-random), the line amplitudes are + ``a_j = a0 * sqrt(w_j / sum_k w_k)`` so that ``sum_j a_j^2 = a0^2`` -- the comb + carries the same *time-averaged* power as the single line (a phase-locked comb + still peaks at ``a0 * sqrt(N)`` at recurrence). Line frequencies are + ``w_j = w0 * (1 + d_j)``, ``d_j`` uniform on ``[-delta_omega, +delta_omega]`` + (``delta_omega`` is the half-width; ``num_colors == 1`` sits exactly at ``w0``). + + All lines share the carrier ``k0``: the source is + ``-env * w_j^2 * a_j * sin(k0 x - w_j t + phi_j)``, so an off-center line is not + placed on its own dispersion branch. This is accurate for a localized antenna + (``source_type: point`` / narrow spatial envelope, where ``dk * L_antenna`` is + negligible) and NOT for a source extended across the box. + """ + params: dict a0: float k0: float w0: float - intensities: Array + intensity_weights: Array # raw per-line weights w_j (dimensionless) + amplitudes: Array # per-line a_j = a0 * sqrt(w_j / sum w) delta_omega: Array phases: Array envelope: SpaceTimeEnvelopeFunction is_point_source: bool = False - driver_list: list def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): self.params = cfg @@ -107,24 +124,27 @@ def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): self.envelope = envelope self.is_point_source = is_point - # intensities + n_colors = int(self.params["num_colors"]) + + # per-line intensity weights w_j if self.params["intensities"]["init"] == "random": int_lo, int_hi = self.params["intensities"].get("range", (0.0, 2.0)) int_rng = np.random.default_rng(seed=self.params["intensities"]["seed"]) - self.intensities = jnp.array(int_rng.uniform(int_lo, int_hi, self.params["num_colors"])) + self.intensity_weights = jnp.array(int_rng.uniform(int_lo, int_hi, n_colors)) elif self.params["intensities"]["init"] == "uniform": - self.intensities = jnp.ones(self.params["num_colors"]) + self.intensity_weights = jnp.ones(n_colors) else: raise NotImplementedError(f"Initialization type -- {self.params['intensities']['init']} -- not implemented") - self.intensities = self.a0 * jnp.sqrt( - self.intensities / jnp.sum(self.intensities) - ) # sqrt normalization to have the same power spectrum - # otherwise for uniform, power be N times the expected power - - # frequency shift - self.delta_omega = ( - jnp.linspace(-self.params["delta_omega"], self.params["delta_omega"], self.params["num_colors"]) * self.w0 - ) + # amplitudes: sqrt normalization so sum_j a_j^2 = a0^2 (same time-averaged power as + # the monochromatic line; otherwise a uniform comb would carry N x the power) + self.amplitudes = self.a0 * jnp.sqrt(self.intensity_weights / jnp.sum(self.intensity_weights)) + + # frequency shifts: + if n_colors == 1: + # a single line must sit exactly at w0 (this is what makes num_colors: 1 the monochromatic driver) + self.delta_omega = jnp.zeros(1) + else: + self.delta_omega = jnp.linspace(-self.params["delta_omega"], self.params["delta_omega"], n_colors) * self.w0 if self.params["phases"]["init"] == "random": # Spectral phases drawn uniformly over (0, 2*pi) -- the default; @@ -137,13 +157,6 @@ def __init__(self, cfg: dict, a0, k0, w0, envelope, is_point): else: raise NotImplementedError(f"Initialization type -- {self.params['phases']['init']} -- not implemented") - DriverList = [] - for a, dw, phase in zip(self.intensities, self.delta_omega, self.phases, strict=True): - driver_obj = EMDriver(a, self.k0, self.w0, phase, dw, self.envelope, self.is_point_source) - DriverList.append(driver_obj) - - self.driver_list = DriverList - def scale_intensities(self, intensities): if self.params["intensities"]["activation"] == "linear": ints = 0.5 * (jnp.tanh(intensities) + 1.0) @@ -164,15 +177,13 @@ def get_partition_spec(self): """ Get the partition spec for the model - Only intensities and phases can be learned + Depends what is learned based on the driver being passed in Returns ------- filter_spec : pytree with the same structure as the model """ - # figure out tracing arrays here - # jit and gradient boundary (figure if that might cause problems) filter_spec = jtu.tree_map(lambda _: False, self) if self.params["intensities"]["learned"]: @@ -183,19 +194,17 @@ def get_partition_spec(self): return filter_spec - def __call__(self, state: dict, args: dict) -> tuple: - # intensities = self.scale_intensities(self.intensities) - # intensities = intensities / jnp.sum(intensities) + # def __call__(self, state: dict, args: dict) -> tuple: + # # intensities = self.scale_intensities(self.intensity_weights) + # # intensities = intensities / jnp.sum(intensities) - """figure out what this does in lpse2d --> is this needed in vlasov (only passed in additional parameters) - but how does diffeqsolve even use these additional parameters""" - args["drivers"]["ey"] = { - "delta_omega": self.delta_omega, - "phases": jnp.tanh(self.phases) * jnp.pi, - "intensities": self.intensities, - } | self.envelope + # args["drivers"]["ey"] = { + # "delta_omega": self.delta_omega, + # "phases": self.phases, + # "amplitudes": self.amplitudes, + # } | self.envelope - return state, args + # return state, args class StochasticDriver(eqx.Module): @@ -255,7 +264,7 @@ class EMDriverSet(eqx.Module): """Container for longitudinal (Ex) and transverse (Ey) driver lists.""" ex: list[EMDriver] - ey: list[EMDriver] + ey: list[EMDriver | BroadbandDriver] ex_stochastic: StochasticDriver | None = None @staticmethod diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index bbbac6c3..a9984b8a 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -4,10 +4,11 @@ # research@ergodic.io from jax import numpy as jnp +from jaxtyping import Array from adept._base_ import get_envelope from adept._vlasov1d.grid import Grid -from adept._vlasov1d.simulation import EMDriver +from adept._vlasov1d.simulation import EMDriver, BroadbandDriver class LongitudinalElectricFieldDriver: @@ -16,7 +17,7 @@ class LongitudinalElectricFieldDriver: def __init__(self, xax, drivers: list[EMDriver]): """Store the spatial axis and longitudinal driver list.""" self.xax = xax - self.drivers = drivers + self.drivers = drivers def _single_driver_field(self, driver: EMDriver, current_time): kk = driver.k0 @@ -54,41 +55,55 @@ def __init__(self, xax, drivers: list[EMDriver], c: float = 0.0): """Precompute point-source masks and scales for transverse drivers.""" self.xax = xax self.drivers = drivers - dx = float(xax[1] - xax[0]) - - self.point_source_masks = [] - self.point_source_scales = [] - for driver in drivers: - if driver.is_point_source: - center = driver.envelope.space_envelope.center - i0 = jnp.argmin(jnp.abs(xax - center)) - mask = jnp.zeros_like(xax).at[i0].set(1.0) - w_total = driver.w0 + driver.dw0 - F0 = 2.0 * w_total * c * driver.a0 - self.point_source_masks.append(mask) - self.point_source_scales.append(F0 / dx) - else: - self.point_source_masks.append(None) - self.point_source_scales.append(None) + self.c = c + self.dx = float(xax[1] - xax[0]) - def _single_driver_source(self, driver: EMDriver, mask, scale, current_time): + def make_scales(self, driver: EMDriver, dw: Array, amplitude: Array): + if driver.is_point_source: + center = driver.envelope.space_envelope.center + i0 = jnp.argmin(jnp.abs(self.xax - center)) + mask = jnp.zeros_like(self.xax).at[i0].set(1.0) + w_total = driver.w0 + dw + F0 = 2.0 * w_total * self.c * amplitude + point_source_masks = mask + point_source_scales = F0 / self.dx + else: + point_source_masks = None + point_source_scales = None + return point_source_masks, point_source_scales + + def _single_driver_source(self, driver: EMDriver, dw: Array, amplitude: Array, phase: Array, mask: Array, scale: Array, current_time): ww = driver.w0 - dw = driver.dw0 - phase = driver.phase w_total = ww + dw if driver.is_point_source: time_env = driver.envelope.time_envelope(current_time) - return scale * time_env * mask * jnp.sin((w_total * current_time) + phase) + return scale[:, None] * jnp.sin((w_total * current_time) + phase)[:, None] * mask[None, :] * time_env else: kk = driver.k0 factor = driver.envelope(self.xax, current_time) - return -factor * w_total**2 * driver.a0 * jnp.sin((kk * self.xax - w_total * current_time) + phase) + # (N, nx): per-color rows so the caller's axis-0 sum works for broadband + return ( + -factor[None, :] + * (w_total**2 * amplitude)[:, None] + * jnp.sin(kk * self.xax[None, :] - (w_total * current_time)[:, None] + phase[:, None]) + ) def __call__(self, t, args): """Evaluate the summed transverse current source at time t.""" total = jnp.zeros_like(self.xax) - for driver, mask, scale in zip(self.drivers, self.point_source_masks, self.point_source_scales, strict=True): - total += self._single_driver_source(driver, mask, scale, t) + # Drivers are read from args unconditionally: args is the differentiable route + # (BaseVlasov1D.__call__ guarantees args["drivers"] is present). No fallback to + # self.drivers -- a silent fallback would zero the driver gradients instead of + # erroring if the args plumbing ever broke. + ey_list = args["drivers"].ey + for driver in ey_list: + amplitudes = driver.amplitudes if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.a0) + phases = driver.phases if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.phase) + delta_omega = driver.delta_omega if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.dw0) + + point_source_masks, point_source_scales = self.make_scales(driver, delta_omega, amplitudes) + driver_source_array = self._single_driver_source(driver, delta_omega, amplitudes, phases, point_source_masks, point_source_scales, t) + total += jnp.sum(driver_source_array, axis=0) return total diff --git a/configs/vlasov-1d/srs-broadband.yaml b/configs/vlasov-1d/srs-broadband.yaml new file mode 100644 index 00000000..0b912332 --- /dev/null +++ b/configs/vlasov-1d/srs-broadband.yaml @@ -0,0 +1,138 @@ +# Broadband (multi-color) SRS example: a 50-line comb, 0.5% full bandwidth +# (delta_omega 0.0025 is the HALF-width), uniform per-line intensity summing to +# base_intensity, seeded random spectral phases. Linear density ramp 0.18 -> 0.30 nc +# over 100 um (L_n(nc/4) = 208 um), 4 keV, 40 ps -- the deck shape used by the +# kinetic-srs broadband threshold campaign (Follett 2019 comparison). Lines share +# the carrier k0, which is fine for this 1 um point-source antenna. +# See docs/source/solvers/vlasov1d/config.md, "Broadband (multi-color) ey driver". + +units: + laser_wavelength: 351nm + normalizing_temperature: 4000eV + normalizing_density: 9.05e21/cc + Z: 10 + Zp: 10 + +density: + quasineutrality: true + species-electron: + v0: 0.0 + T0: 1.0 + m: 2.0 + noise_seed: 0 + noise_type: uniform + noise_val: 0.0 + # n(x) = (baseline + bump_height*tanh_env) * val_at_center * (1 + (x-center)/L_fed), + # L_fed = gsl/val_at_center = 200.0/0.24 -> dn/dx = 0.0012 nc/um + # flat-top spans 0.18 -> 0.30 nc over 100.0um; L_n(nc/4) = 208.3um (Follett box 100 um) + basis: linear + val_at_center: 0.24 + gradient_scale_length: 200.0um + center: 60.0um + width: 100.0um + rise: 1.0um + baseline: 0.004 + bump_height: 0.996 + bump_or_trough: bump + +grid: + dt: 0.005ps + tmin: 0.0ps + tmax: 40.0ps + + nx: 8192 + xmin: 0.0um + xmax: 120um + + nv: 256 + vmax: 10.0 + +save: + fields: + t: + nt: 600 + electron: + main: + t: + nt: 4 + +solver: vlasov-1d + +mlflow: + experiment: vlasov1d-srs-broadband + run: srs-broadband-example + +drivers: + ex: {} + ey: + '0': + source_type: point + params: + num_colors: 50 + delta_omega: 0.0025 + wavelength: 351nm + leftgoing: false + intensities: + base_intensity: 2.378e+14 W/cm^2 + init: uniform + phases: + init: random + seed: 1 + envelope: + time: + center: 20500.0fs + rise: 100.0fs + width: 42000.0fs + space: + center: 5.0um + rise: 0.1um + width: 1.0um + +diagnostics: + diag-vlasov-dfdt: False + diag-fp-dfdt: False + +terms: + field: poisson + edfdv: cubic-spline + time: leapfrog + hou_li_filter: + is_on: True + alpha: 36.0 + order: 36 + dimensions: [x] + fokker_planck: + is_on: True + type: Dougherty + time: + baseline: 0.00001 + bump_or_trough: bump + center: 0.0ps + rise: 1.0ps + slope: 0.0 + bump_height: 0.0 + width: 1s + space: + baseline: 1.0 + bump_or_trough: bump + center: 0.0um + rise: 1.0um + slope: 0.0 + bump_height: 0.0 + width: 1000um + krook: + is_on: True + time: + baseline: 1.0 + bump_or_trough: bump + center: 0.0ps + rise: 1.0ps + bump_height: 0.0 + width: 1s + space: + baseline: 0.0001 + bump_or_trough: trough + center: 60.0um + rise: 1.0um + bump_height: 0.2 + width: 100.0um diff --git a/docs/source/solvers/vlasov1d/config.md b/docs/source/solvers/vlasov1d/config.md index 8e556124..116323e4 100644 --- a/docs/source/solvers/vlasov1d/config.md +++ b/docs/source/solvers/vlasov1d/config.md @@ -463,6 +463,76 @@ drivers: x_width: 1.0 ``` +### Broadband (multi-color) `ey` driver + +A single `ey` driver can launch a **comb of `num_colors` lines** instead of one +monochromatic wave, by giving it the broadband `params` form (the driver entry takes +`params:` + `envelope:` + `source_type:`; the monochromatic forms use +`params: {a0, k0, w0, dw0, phase}` or `params: {intensity, wavelength, ...}`): + +```yaml +drivers: + ex: {} + ey: + '0': + source_type: point + params: + num_colors: 50 + delta_omega: 0.0025 # HALF-width, fraction of w0 -> full bandwidth 0.5% + wavelength: 351nm + leftgoing: false + intensities: + base_intensity: 2.378e+14 W/cm^2 + init: uniform # or random (then seed is required) + phases: + init: random + seed: 1 + envelope: + time: {center: 20500.0fs, rise: 100.0fs, width: 42000.0fs} + space: {center: 5.0um, rise: 0.1um, width: 1.0um} +``` + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `num_colors` | int ≥ 1 | — | Number of lines `N`. `num_colors: 1` is exactly the monochromatic `intensity`/`wavelength` driver (line at `w0`, `dw0 = 0`). | +| `delta_omega` | float ≥ 0 | — | **Half-width** of the comb as a fraction of `w0`: lines sit at `w_j = w0 (1 + d_j)` with `d_j` uniformly spaced on `[-delta_omega, +delta_omega]`, so the **full bandwidth is `2 * delta_omega`**. A "0.25% bandwidth" run is `delta_omega: 0.00125`. | +| `wavelength` | str | — | Carrier wavelength; sets `w0` and the shared `k0`. | +| `leftgoing` | bool | `false` | Flip the sign of `k0`. | +| `intensities.base_intensity` | str | — | Total intensity, e.g. `2.378e+14 W/cm^2`. It fixes the monochromatic amplitude `a0`; the comb carries the same time-averaged power (see below). | +| `intensities.init` | `uniform` / `random` | `uniform` | Per-line weights `w_j`: all equal, or drawn uniformly from `intensities.range`. | +| `intensities.seed` | int | — | RNG seed for `init: random` (required then). | +| `intensities.range` | `[lo, hi]` | `[0, 2]` | Draw bounds for `init: random`. | +| `phases.init` | `uniform` / `random` | `random` | Spectral phases: all equal to `phases.base_phase`, or drawn uniformly from `phases.range`. | +| `phases.seed` | int | — | RNG seed for `init: random` (required then). | +| `phases.range` | `[lo, hi]` | `[0, 2π]` | Draw bounds (radians) for `init: random`. | +| `phases.base_phase` | float | `0.0` | Common phase for `init: uniform`. | + +How the lines are built (`BroadbandDriver` in `simulation.py`): + +- **Amplitudes.** With `a0` the monochromatic amplitude for `base_intensity`, line `j` gets + `a_j = a0 · sqrt(w_j / Σ_k w_k)`, so `Σ_j a_j² = a0²` — the comb has the **same + time-averaged power** as one line at `base_intensity`, and the per-line intensity is + `I_j = base_intensity · w_j / Σ_k w_k` (`base_intensity / N` for a uniform comb). A + phase-locked comb still reaches `a0·√N` at its recurrence time, which is the correct + coherent-sum physics; "same intensity" here means time-averaged, not peak. +- **Frequencies.** `w_j = w0 (1 + d_j)`. The driver is a single `BroadbandDriver` module + carrying per-line arrays (`amplitudes`, `delta_omega`, `phases`, each shape `(N,)`); + the transverse source evaluates all lines vectorized. The line parameters are read + from the solver's `args` at evaluation time, so they are differentiable inputs + (gradients w.r.t. `amplitudes`/`phases` flow; see `test_broadband_driver.py`). +- **Shared `k0`.** Every line is launched with the carrier wavenumber of `wavelength`: + the source is `-env · w_j² · a_j · sin(k0 x − w_j t + φ_j)`, so an off-center line is + **not** placed on its own dispersion branch. That is accurate for a localized antenna + (`source_type: point`, or a spatial envelope a wavelength wide, where `δk · L_antenna` is + negligible) and not for a source extended across the box (`δk · L ~ 10 rad` at 0.25% over + 200 µm). +- Both `init: random` draws use `numpy.random.default_rng(seed)`, so a deck is exactly + reproducible across builds; two runs differ only through the seeds. +- With more than one line, `post_process` writes `plots/drivers/` (per-line intensity, its + log, and phase vs `Δω/ω0`) from the live driver objects. + +An example deck is `configs/vlasov-1d/srs-broadband.yaml`. + ## diagnostics Enable/disable diagnostic outputs. diff --git a/tests/test_vlasov1d/test_broadband_driver.py b/tests/test_vlasov1d/test_broadband_driver.py new file mode 100644 index 00000000..5481d4cc --- /dev/null +++ b/tests/test_vlasov1d/test_broadband_driver.py @@ -0,0 +1,173 @@ +"""Unit tests for the broadband (multi-color) ey driver. + +Architecture under test (see BroadbandDriver in simulation.py): ``EMDriver.from_config`` +returns ONE ``BroadbandDriver`` eqx.Module carrying per-line array leaves +(``amplitudes``/``delta_omega``/``phases``, each (N,)); the transverse source pusher +reads drivers from ``args["drivers"].ey`` and evaluates all lines vectorized. No full +solve is run here. + +Covered: + * a uniform N-line comb carries the same time-averaged power as the monochromatic + driver at ``base_intensity`` (``sum_j a_j^2 == a0^2``), + * ``delta_omega`` is the comb HALF-width; ``num_colors: 1`` sits exactly at ``w0`` + and reproduces the monochromatic driver's source to machine precision at the + ``TransverseCurrentSourceDriver`` level, + * seeded line sets are reproducible across independent builds, + * config validation (``init: random`` without a seed, missing ``base_intensity``), + * gradients w.r.t. ``amplitudes`` and ``phases`` flow through the args route + (regression guard: fails if a construction-time precompute or a disconnected + driver copy is ever reintroduced). +""" + +import math + +import numpy as np +import pytest +from jax import config as jax_config + +jax_config.update("jax_enable_x64", True) + +import equinox as eqx # noqa: E402 +import jax # noqa: E402 +from jax import numpy as jnp # noqa: E402 +from pydantic import ValidationError # noqa: E402 + +from adept._vlasov1d.datamodel import EMDriverConfig # noqa: E402 +from adept._vlasov1d.simulation import BroadbandDriver, EMDriver, EMDriverSet # noqa: E402 +from adept._vlasov1d.solvers.pushers.field import TransverseCurrentSourceDriver # noqa: E402 +from adept.normalization import electron_debye_normalization # noqa: E402 + +NORM = electron_debye_normalization("9.05e21/cc", "4000eV") +C_NORM = NORM.speed_of_light_norm() +BASE_INTENSITY = "2.378e+14 W/cm^2" +WAVELENGTH = "351nm" + + +def _envelope(source_type): + # micron-scale antenna at 5 um for the point source; box-wide for extended + width = "1.0um" if source_type == "point" else "50.0um" + return { + "time": {"center": "20500.0fs", "rise": "100.0fs", "width": "42000.0fs"}, + "space": {"center": "5.0um" if source_type == "point" else "25.0um", "rise": "0.1um", "width": width}, + } + + +def _mono(source_type="point"): + cfg = EMDriverConfig( + params={"intensity": BASE_INTENSITY, "wavelength": WAVELENGTH}, + envelope=_envelope(source_type), + source_type=source_type, + ) + (d,) = EMDriver.from_config(cfg, NORM) + return d + + +def _comb(num_colors, delta_omega, intensities=None, phases=None, source_type="point"): + cfg = EMDriverConfig( + params={ + "num_colors": num_colors, + "delta_omega": delta_omega, + "wavelength": WAVELENGTH, + "intensities": intensities or {"base_intensity": BASE_INTENSITY, "init": "uniform"}, + "phases": phases or {"init": "random", "seed": 1}, + }, + envelope=_envelope(source_type), + source_type=source_type, + ) + bb = EMDriver.from_config(cfg, NORM) + assert isinstance(bb, BroadbandDriver) + return bb + + +def _source(driver, t, xax): + """Evaluate the transverse current source through the args route the solver uses.""" + pusher = TransverseCurrentSourceDriver(xax, drivers=[driver], c=C_NORM) + return pusher(t, {"drivers": EMDriverSet(ex=[], ey=[driver])}) + + +XAX = jnp.linspace(0.0, 50.0 * 213.0, 1024) # ~50 um in Debye-length units, coarse + + +def test_uniform_comb_has_monochromatic_power(): + """sum_j a_j^2 == a0^2 for a uniform comb (a_j = a0 sqrt(w_j / sum w)).""" + mono = _mono() + bb = _comb(50, 0.0025) + assert bb.amplitudes.shape == (50,) + assert np.allclose(np.asarray(bb.amplitudes), float(bb.amplitudes[0])) # uniform weights + assert math.isclose(float(jnp.sum(bb.amplitudes**2)), float(mono.a0) ** 2, rel_tol=1e-12) + # the comb shares the monochromatic carrier + assert float(bb.k0) == float(mono.k0) and float(bb.w0) == float(mono.w0) + + +def test_delta_omega_is_half_width(): + """delta_omega (N,) spans [-d, +d] * w0 uniformly (full width 2d).""" + bb = _comb(3, 0.0025) + rel = np.asarray(bb.delta_omega) / float(bb.w0) + assert np.allclose(rel, [-0.0025, 0.0, 0.0025], atol=1e-15) + bb = _comb(2, 0.0025) + rel = np.asarray(bb.delta_omega) / float(bb.w0) + assert np.allclose(rel, [-0.0025, 0.0025], atol=1e-15) + + +@pytest.mark.parametrize("source_type", ["point", "extended"]) +def test_single_line_is_monochromatic(source_type): + """num_colors: 1 IS the monochromatic driver: dw = 0, same a0, and the evaluated + source matches the plain EMDriver's to machine precision at several times.""" + mono = _mono(source_type) + bb = _comb(1, 0.0025, phases={"init": "uniform", "base_phase": 0.0}, source_type=source_type) + assert bb.delta_omega.shape == (1,) and float(bb.delta_omega[0]) == 0.0 # not -delta_omega * w0 + assert float(bb.amplitudes[0]) == float(mono.a0) + for t in (0.0, 313.7, 20500.0 * 1.88, 41000.0 * 1.88): # spread across the envelope + s_mono = _source(mono, t, XAX) + s_bb = _source(bb, t, XAX) + np.testing.assert_allclose(np.asarray(s_bb), np.asarray(s_mono), rtol=1e-12, atol=0.0) + + +def test_seeded_line_sets_are_reproducible(): + """Same seeds -> identical arrays across two builds; phase seed only moves phases.""" + kw = dict( + intensities={"base_intensity": BASE_INTENSITY, "init": "random", "seed": 7}, + phases={"init": "random", "seed": 3}, + ) + a, b = _comb(20, 0.001, **kw), _comb(20, 0.001, **kw) + np.testing.assert_array_equal(np.asarray(a.amplitudes), np.asarray(b.amplitudes)) + np.testing.assert_array_equal(np.asarray(a.phases), np.asarray(b.phases)) + np.testing.assert_array_equal(np.asarray(a.delta_omega), np.asarray(b.delta_omega)) + c = _comb(20, 0.001, intensities=kw["intensities"], phases={"init": "random", "seed": 4}) + np.testing.assert_array_equal(np.asarray(a.amplitudes), np.asarray(c.amplitudes)) + assert np.any(np.asarray(a.phases) != np.asarray(c.phases)) + # random weights still carry the monochromatic power + assert math.isclose(float(jnp.sum(a.amplitudes**2)), float(_mono().a0) ** 2, rel_tol=1e-12) + + +def test_random_init_requires_seed(): + """init: random without a seed is a validation error, not a KeyError at build time.""" + with pytest.raises(ValidationError, match="phases.seed"): + _comb(4, 0.001, phases={"init": "random"}) + with pytest.raises(ValidationError, match="intensities.seed"): + _comb(4, 0.001, intensities={"base_intensity": BASE_INTENSITY, "init": "random"}) + with pytest.raises(ValidationError): # base_intensity is required + _comb(4, 0.001, intensities={"init": "uniform"}) + + +@pytest.mark.parametrize("source_type", ["point", "extended"]) +def test_gradients_flow_through_args_route(source_type): + """d(source)/d(amplitudes) and d(source)/d(phases) are finite and nonzero. + + Regression guard for the differentiable-driver design: the pusher must read the + line parameters from args["drivers"] and compute point-source scales in-trace. + A construction-time precompute or a fallback to a disconnected self.drivers copy + zeroes these gradients, failing this test.""" + bb = _comb(8, 0.01, source_type=source_type) + pusher = TransverseCurrentSourceDriver(XAX, drivers=[bb], c=C_NORM) + t = 20500.0 * 1.88 # mid-envelope so the time envelope is ~1 + + def loss(amps, phases): + d = eqx.tree_at(lambda m: (m.amplitudes, m.phases), bb, (amps, phases)) + return jnp.sum(pusher(t, {"drivers": EMDriverSet(ex=[], ey=[d])}) ** 2) + + g_amp, g_phase = jax.grad(loss, argnums=(0, 1))(bb.amplitudes, bb.phases) + for g in (g_amp, g_phase): + assert g.shape == (8,) + assert np.all(np.isfinite(np.asarray(g))) + assert np.max(np.abs(np.asarray(g))) > 0.0 From ea1a394f6325e507679bda681853346a129a552d Mon Sep 17 00:00:00 2001 From: Manhar Date: Fri, 21 Aug 2026 01:54:44 -0700 Subject: [PATCH 09/10] Initialization level changes for backward pass --- adept/_vlasov1d/modules.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/adept/_vlasov1d/modules.py b/adept/_vlasov1d/modules.py index cc0b182b..b22471e4 100644 --- a/adept/_vlasov1d/modules.py +++ b/adept/_vlasov1d/modules.py @@ -342,6 +342,10 @@ def __call__(self, trainable_modules: dict, args: dict | None = None): # 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. + state = self.state for name, module in trainable_modules.items(): state, args = module(state, args) @@ -353,7 +357,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) From 88de80e4219678815f32d4117933e9d59f2aa8ba Mon Sep 17 00:00:00 2001 From: Manhar Date: Sat, 22 Aug 2026 00:10:38 -0700 Subject: [PATCH 10/10] Updates for pre-commit and CI tests --- adept/_vlasov1d/helpers.py | 7 ++++--- adept/_vlasov1d/modules.py | 5 +++-- adept/_vlasov1d/solvers/pushers/field.py | 18 +++++++++++------- tests/test_vlasov1d/test_absorbing_wave.py | 11 ++++++++--- tests/test_vlasov1d/test_broadband_driver.py | 20 ++++++++++---------- 5 files changed, 36 insertions(+), 25 deletions(-) diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index 63753948..23730130 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -6,12 +6,12 @@ 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 -import equinox as eqx from matplotlib import pyplot as plt from scipy.special import gamma @@ -41,6 +41,7 @@ def gamma_5_over_m(m): """Evaluate Gamma(5 / m) for super-Gaussian normalization.""" return gamma(5.0 / m) # np.interp(m, m_ax, g_5_m) + def _initialize_supergaussian_distribution_( nx: int, nv: int, @@ -174,9 +175,9 @@ def _initialize_total_distribution_(cfg, simulation: Vlasov1DSimulation): def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: PlasmaNormalization | None = None): - '''getting amplitude (a), wave number (k) and angular frequency (w) + """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_wavelength' type of configs""" intensity = UREG.Quantity(intensity).to("W/m^2") wavelength = UREG.Quantity(wavelength).to("nm") diff --git a/adept/_vlasov1d/modules.py b/adept/_vlasov1d/modules.py index b22471e4..9a6f2df8 100644 --- a/adept/_vlasov1d/modules.py +++ b/adept/_vlasov1d/modules.py @@ -344,9 +344,10 @@ def __call__(self, trainable_modules: dict, args: dict | None = None): # 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. + # 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.items(): + for name, module in (trainable_modules or {}).items(): state, args = module(state, args) grid = self.simulation.grid diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index a9984b8a..6eb3dfb6 100644 --- a/adept/_vlasov1d/solvers/pushers/field.py +++ b/adept/_vlasov1d/solvers/pushers/field.py @@ -8,7 +8,7 @@ from adept._base_ import get_envelope from adept._vlasov1d.grid import Grid -from adept._vlasov1d.simulation import EMDriver, BroadbandDriver +from adept._vlasov1d.simulation import BroadbandDriver, EMDriver class LongitudinalElectricFieldDriver: @@ -17,7 +17,7 @@ class LongitudinalElectricFieldDriver: def __init__(self, xax, drivers: list[EMDriver]): """Store the spatial axis and longitudinal driver list.""" self.xax = xax - self.drivers = drivers + self.drivers = drivers def _single_driver_field(self, driver: EMDriver, current_time): kk = driver.k0 @@ -67,12 +67,14 @@ def make_scales(self, driver: EMDriver, dw: Array, amplitude: Array): F0 = 2.0 * w_total * self.c * amplitude point_source_masks = mask point_source_scales = F0 / self.dx - else: + else: point_source_masks = None - point_source_scales = None + point_source_scales = None return point_source_masks, point_source_scales - def _single_driver_source(self, driver: EMDriver, dw: Array, amplitude: Array, phase: Array, mask: Array, scale: Array, current_time): + def _single_driver_source( + self, driver: EMDriver, dw: Array, amplitude: Array, phase: Array, mask: Array, scale: Array, current_time + ): ww = driver.w0 w_total = ww + dw if driver.is_point_source: @@ -98,11 +100,13 @@ def __call__(self, t, args): ey_list = args["drivers"].ey for driver in ey_list: amplitudes = driver.amplitudes if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.a0) - phases = driver.phases if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.phase) + phases = driver.phases if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.phase) delta_omega = driver.delta_omega if isinstance(driver, BroadbandDriver) else jnp.atleast_1d(driver.dw0) point_source_masks, point_source_scales = self.make_scales(driver, delta_omega, amplitudes) - driver_source_array = self._single_driver_source(driver, delta_omega, amplitudes, phases, point_source_masks, point_source_scales, t) + driver_source_array = self._single_driver_source( + driver, delta_omega, amplitudes, phases, point_source_masks, point_source_scales, t + ) total += jnp.sum(driver_source_array, axis=0) return total diff --git a/tests/test_vlasov1d/test_absorbing_wave.py b/tests/test_vlasov1d/test_absorbing_wave.py index b41ec4b9..57b7f760 100644 --- a/tests/test_vlasov1d/test_absorbing_wave.py +++ b/tests/test_vlasov1d/test_absorbing_wave.py @@ -4,7 +4,7 @@ from jax import jit from adept._base_ import Stepper -from adept._vlasov1d.simulation import EMDriver +from adept._vlasov1d.simulation import EMDriver, EMDriverSet from adept._vlasov1d.solvers.pushers.field import TransverseCurrentSourceDriver, WaveSolver from adept.functions import EnvelopeFunction, SpaceTimeEnvelopeFunction @@ -52,9 +52,14 @@ def test_absorbing_boundaries(): ey_driver = EMDriver(a0=1.0e-4, k0=-1.4, w0=15.82, phase=0.0, dw0=0.0, envelope=envelope) drivers = [ey_driver] - args = {} + # The transverse source pusher reads drivers from args (the differentiable + # route); standalone use must supply them the same way the solver does. + args = {"drivers": EMDriverSet(ex=[], ey=drivers)} - @jit + # filter_jit (not plain jit): args carries the EMDriverSet, whose non-array + # leaves (is_point_source bool, scalar params) must stay static -- the same + # filtering diffrax applies on the solver path. + @eqx.filter_jit def _run_(y, args): return diffeqsolve( terms=ODETerm(VectorField(c_light, dx, dt, xax, drivers)), diff --git a/tests/test_vlasov1d/test_broadband_driver.py b/tests/test_vlasov1d/test_broadband_driver.py index 5481d4cc..2ae4bf0f 100644 --- a/tests/test_vlasov1d/test_broadband_driver.py +++ b/tests/test_vlasov1d/test_broadband_driver.py @@ -27,15 +27,15 @@ jax_config.update("jax_enable_x64", True) -import equinox as eqx # noqa: E402 -import jax # noqa: E402 -from jax import numpy as jnp # noqa: E402 -from pydantic import ValidationError # noqa: E402 +import equinox as eqx +import jax +from jax import numpy as jnp +from pydantic import ValidationError -from adept._vlasov1d.datamodel import EMDriverConfig # noqa: E402 -from adept._vlasov1d.simulation import BroadbandDriver, EMDriver, EMDriverSet # noqa: E402 -from adept._vlasov1d.solvers.pushers.field import TransverseCurrentSourceDriver # noqa: E402 -from adept.normalization import electron_debye_normalization # noqa: E402 +from adept._vlasov1d.datamodel import EMDriverConfig +from adept._vlasov1d.simulation import BroadbandDriver, EMDriver, EMDriverSet +from adept._vlasov1d.solvers.pushers.field import TransverseCurrentSourceDriver +from adept.normalization import electron_debye_normalization NORM = electron_debye_normalization("9.05e21/cc", "4000eV") C_NORM = NORM.speed_of_light_norm() @@ -142,9 +142,9 @@ def test_seeded_line_sets_are_reproducible(): def test_random_init_requires_seed(): """init: random without a seed is a validation error, not a KeyError at build time.""" - with pytest.raises(ValidationError, match="phases.seed"): + with pytest.raises(ValidationError, match=r"phases\.seed"): _comb(4, 0.001, phases={"init": "random"}) - with pytest.raises(ValidationError, match="intensities.seed"): + with pytest.raises(ValidationError, match=r"intensities\.seed"): _comb(4, 0.001, intensities={"base_intensity": BASE_INTENSITY, "init": "random"}) with pytest.raises(ValidationError): # base_intensity is required _comb(4, 0.001, intensities={"init": "uniform"})