diff --git a/adept/_pic1d/helpers.py b/adept/_pic1d/helpers.py index e85c1ce5..fecb5deb 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,9 @@ def _inverse_cdf_supergaussian( sampling that exactly reproduces the requested moments in expectation. """ v_thermal = np.sqrt(T0 / mass) - alpha = np.sqrt(3.0 * gamma_3_over_m(supergaussian_order) / gamma_5_over_m(supergaussian_order)) + alpha = np.sqrt( + gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order) + ) # 1D width: variance = T0/mass for every m (see _vlasov1d.helpers) nv = max(8192, 8 * n_particles) v_grid = np.linspace(-vmax, vmax, nv) @@ -56,7 +58,9 @@ def _random_supergaussian( ) -> np.ndarray: """Random velocity sampling via rejection on a bounded supergaussian.""" v_thermal = np.sqrt(T0 / mass) - alpha = np.sqrt(3.0 * gamma_3_over_m(supergaussian_order) / gamma_5_over_m(supergaussian_order)) + alpha = np.sqrt( + gamma_1_over_m(supergaussian_order) / gamma_3_over_m(supergaussian_order) + ) # 1D width: variance = T0/mass for every m (see _vlasov1d.helpers) out = np.empty(n_particles) filled = 0 diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 4161a491..af61fbc7 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -127,6 +127,67 @@ class IntensityWavelengthDriverConfig(BaseModel): intensity: str wavelength: str leftgoing: bool = False + dw0: float = 0.0 + phase: float = 0.0 + + +class BroadbandIntensitiesConfig(BaseModel): + """Per-line intensity weights of a broadband (multi-color) driver. + + The weights ``w_j`` set the line *amplitudes* as ``a_j = a0 * sqrt(w_j / sum_k w_k)``, + where ``a0`` is the monochromatic amplitude for ``base_intensity``; hence + ``sum_j a_j^2 = a0^2`` (same time-averaged power as a single line at + ``base_intensity``) and the per-line intensity is ``I_j = base_intensity * w_j / sum_k w_k``. + """ + + base_intensity: str # e.g. "2.378e+14 W/cm^2"; the total (monochromatic-equivalent) intensity + init: Literal["uniform", "random"] = "uniform" + seed: int | None = None # required for init: random + range: tuple[float, float] = (0.0, 2.0) # uniform draw bounds for init: random + + @model_validator(mode="after") + def check_seed(self) -> "BroadbandIntensitiesConfig": + """``init: random`` must be reproducible, so it needs a seed.""" + if self.init == "random" and self.seed is None: + raise ValueError("intensities.init == 'random' requires intensities.seed") + return self + + +class BroadbandPhasesConfig(BaseModel): + """Per-line spectral phases of a broadband (multi-color) driver.""" + + init: Literal["uniform", "random"] = "random" + seed: int | None = None # required for init: random + range: tuple[float, float] = (0.0, 2.0 * 3.141592653589793) # uniform draw bounds, radians + base_phase: float = 0.0 # the common phase for init: uniform + + @model_validator(mode="after") + def check_seed(self) -> "BroadbandPhasesConfig": + """``init: random`` must be reproducible, so it needs a seed.""" + if self.init == "random" and self.seed is None: + raise ValueError("phases.init == 'random' requires phases.seed") + return self + + +class BroadbandConfig(BaseModel): + """Broadband (multi-color) laser driver: a comb of ``num_colors`` lines. + + Line frequencies are ``w_j = w0 * (1 + d_j)`` with ``d_j`` spaced uniformly on + ``[-delta_omega, +delta_omega]`` (``num_colors: 1`` puts the single line exactly at + ``w0``). **``delta_omega`` is the HALF-width of the comb as a fraction of ``w0``; + the full bandwidth is ``2 * delta_omega``** -- a "0.25% bandwidth" run is + ``delta_omega: 0.00125``. All lines share the carrier wavenumber ``k0`` of + ``wavelength`` (an off-center line is not launched on its own dispersion branch), + which is accurate for a localized antenna (``source_type: point`` / a narrow + spatial envelope) and not for a source extended across the box. + """ + + num_colors: int = Field(ge=1) + delta_omega: float = Field(ge=0.0) # half-width, fraction of w0 + wavelength: str + intensities: BroadbandIntensitiesConfig + phases: BroadbandPhasesConfig + leftgoing: bool = False class AKWDriverConfig(BaseModel): @@ -136,6 +197,7 @@ class AKWDriverConfig(BaseModel): k0: float | None = None w0: float | None = None dw0: float + phase: float = 0.0 @model_validator(mode="after") def check_w_or_k(self) -> "AKWDriverConfig": @@ -148,7 +210,7 @@ def check_w_or_k(self) -> "AKWDriverConfig": class EMDriverConfig(BaseModel): """One electromagnetic driver with parameters, envelope, and source geometry.""" - params: IntensityWavelengthDriverConfig | AKWDriverConfig + params: IntensityWavelengthDriverConfig | AKWDriverConfig | BroadbandConfig envelope: SpaceTimeEnvelopeConfig source_type: Literal["extended", "point"] = "extended" diff --git a/adept/_vlasov1d/helpers.py b/adept/_vlasov1d/helpers.py index 534047ac..23730130 100644 --- a/adept/_vlasov1d/helpers.py +++ b/adept/_vlasov1d/helpers.py @@ -2,19 +2,22 @@ # Copyright (c) Ergodic LLC 2023 # research@ergodic.io +import math import os from time import time +import equinox as eqx import numpy as np import xarray from diffrax import Solution from jax import numpy as jnp +from jax import tree_util as jtu from matplotlib import pyplot as plt from scipy.special import gamma from adept._vlasov1d.simulation import SubspeciesDistributionSpec, Vlasov1DSimulation from adept._vlasov1d.storage import store_f, store_fields -from adept.normalization import PlasmaNormalization +from adept.normalization import UREG, PlasmaNormalization, normalize from .. import patched_mlflow as mlflow @@ -24,6 +27,11 @@ # g_5_m = np.squeeze(gamma_da.loc[{"gamma": "5/m"}].data) +def gamma_1_over_m(m): + """Evaluate Gamma(1 / m) for super-Gaussian normalization.""" + return gamma(1.0 / m) + + def gamma_3_over_m(m): """Evaluate Gamma(3 / m) for super-Gaussian normalization.""" return gamma(3.0 / m) # np.interp(m, m_ax, g_3_m) @@ -72,11 +80,16 @@ def _initialize_supergaussian_distribution_( # Thermal velocity: v_t = sqrt(T/m) v_thermal = np.sqrt(T0 / mass) - # Alpha factor for supergaussian normalization. This fixes the moment ratio - # / = 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)) @@ -161,10 +174,192 @@ def _initialize_total_distribution_(cfg, simulation: Vlasov1DSimulation): return species_distributions +def get_akw_from_intensity_wavelength(intensity, wavelength, leftgoing, norm: PlasmaNormalization | None = None): + """getting amplitude (a), wave number (k) and angular frequency (w) + from intensity and wavelength (passed in as args to this function) defined in + 'intensity_wavelength' type of configs""" + + intensity = UREG.Quantity(intensity).to("W/m^2") + wavelength = UREG.Quantity(wavelength).to("nm") + + e = UREG.e + m_e = UREG.m_e + eps0 = UREG.epsilon_0 + c = UREG.c + + # Standard a0 = eE0/(m_e c w0) — identical to HermiteSRS1D formula + a0_std = ((e * wavelength / (m_e * math.pi)) * (intensity / (2 * eps0 * c**5)) ** 0.5).to("").magnitude + # Vlasov normalization: a0_vlasov = a0_std / β (β = v0/c) + a0 = a0_std * norm.speed_of_light_norm() + + # k0 in Debye-length units: k0_vlasov = k_phys x v0/wp0 + k0_phys = (2 * math.pi / wavelength).to("1/m") + k_sign = -1.0 if leftgoing else 1.0 + k0 = k_sign * float((k0_phys * norm.L0).to("").magnitude) + + # w0 normalized to wp0 (same normalization as Hermite) + w0_phys = (2 * math.pi * c / wavelength).to("1/s") + w0 = float((w0_phys * norm.tau).to("").magnitude) + + return a0, k0, w0 + + +def plot_driver_spectra(cfg: dict, td: str, args: dict): + """Per-line intensity and phase vs frequency offset, for each multi-line driver. + + Reads the LIVE driver objects out of `args["drivers"]` rather than re-deriving + the line set from the config's init/seed. That matters twice over: the plot cannot + drift if BroadbandDriver's construction changes, and it shows optimized line sets + from a backward pass (which never appear in the config) automatically. A + `BroadbandDriver` contributes its per-line arrays (`amplitudes`/`delta_omega`/ + `phases`); plain `EMDriver`s contribute their scalars, so a hand-built list of + mono drivers still plots. + + Per-line intensity needs no normalization constant. The driver builds amplitudes as + A_j = a0 * sqrt(w_j / sum_k w_k), so + + I_j / I_base = A_j^2 / sum_k A_k^2 + + exactly, independent of a0 and of the plasma normalization. `base_intensity` is + read from the config only to put an absolute scale on the axis. + + Panels: I_j (linear), log10 I_j (dynamic range -- the informative one once an + optimizer produces structure, since I ~ A^2 doubles the span), and phi_j. + Single-line (monochromatic) drivers are skipped. + """ + drivers = (args or {}).get("drivers") + if drivers is None: + return + out_dir = os.path.join(td, "plots", "drivers") + + for field in ("ex", "ey"): + dlist = getattr(drivers, field, None) + if not dlist: + continue + + amp_list, dw_list, phase_list = [], [], [] + for d in dlist: + if hasattr(d, "amplitudes"): # BroadbandDriver: (N,) array leaves + amp_list.extend(np.asarray(d.amplitudes, dtype=float)) + dw_list.extend(np.asarray(d.delta_omega, dtype=float) / float(d.w0)) + phase_list.extend(np.asarray(d.phases, dtype=float)) + else: # plain EMDriver: scalar leaves + amp_list.append(float(d.a0)) + dw_list.append(float(d.dw0) / float(d.w0)) + phase_list.append(float(d.phase)) + if len(amp_list) < 2: + continue # monochromatic -> no spectrum to show + + amp = np.asarray(amp_list) + dw = np.asarray(dw_list) # dw_j/w0 + phases = np.asarray(phase_list) + + power = amp**2 + frac = power / power.sum() if power.sum() > 0 else power + # absolute scale: base_intensity of the broadband driver in this field (whichever + # key it was given under -- a driver keyed '1' must not lose the axis scale) + base = None + for dcfg in (cfg.get("drivers", {}).get(field, {}) or {}).values(): + ints = ((dcfg or {}).get("params", {}) or {}).get("intensities") + if isinstance(ints, dict) and ints.get("base_intensity") is not None: + base = ints["base_intensity"] + break + + I_j, unit = frac, "" + if isinstance(base, str) and base.split(): # "2.378e+14 W/cm^2" + val, _, u = base.partition(" ") + try: + I_j, unit = frac * float(val), u.strip() + except ValueError: + pass + elif isinstance(base, (int, float)): + I_j = frac * float(base) + + order = np.argsort(dw) + dw, I_j, phases = dw[order], I_j[order], phases[order] + + # constrained_layout sizes the suptitle band to the actual text; do NOT pair + # it with tight_layout(rect=...) + suptitle(y=...), which reserve a fixed band + # and leave a gap when the title is shorter than the reservation. + fig, axes = plt.subplots(3, 1, figsize=(7.2, 8.6), sharex=True, constrained_layout=True) + bw_pct = (dw.max() - dw.min()) * 100.0 + spacing = float(np.diff(np.sort(dw)).mean()) if len(dw) > 1 else 0.0 + run_name = ((cfg.get("mlflow") or {}).get("run")) or "" + + title = f"{field} driver — broadband line spectrum" + if run_name: + title += f"\n{run_name}" + title += ( + f"\n{len(dlist)} lines | full width $\\Delta\\omega/\\omega_0$ = " + f"{bw_pct:.3g}% | spacing $\\delta\\omega/\\omega_0$ = {spacing:.3g}" + ) + if base is not None: + title += f"\n$I_{{base}}$ = {base}" + if unit: + title += f" | $I_j$ = {I_j.mean():.4g} {unit} mean per line" + fig.suptitle(title, fontsize=9.5, linespacing=1.4) + + axes[0].plot(dw, I_j, "o", ms=4, color="#1f77b4") + axes[0].set_ylabel("line intensity $I_j$" + (f" [{unit}]" if unit else " [$I_j/I_{base}$]")) + axes[0].set_ylim(bottom=0) + axes[0].annotate( + rf"$\Sigma_j I_j$ = {I_j.sum():.4g}", xy=(0.02, 0.06), xycoords="axes fraction", fontsize=8, color="0.35" + ) + + pos = I_j > 0 + if pos.any(): + axes[1].plot(dw[pos], np.log10(I_j[pos]), "o", ms=4, color="#d62728") + lo_, hi_ = np.log10(I_j[pos]).min(), np.log10(I_j[pos]).max() + if hi_ - lo_ < 0.1: + axes[1].set_ylim(lo_ - 0.5, hi_ + 0.5) + axes[1].annotate( + f"dynamic range: {10 ** (hi_ - lo_):.3g}x", + xy=(0.02, 0.88), + xycoords="axes fraction", + fontsize=8, + color="0.35", + ) + if (~pos).any(): # an optimizer can drive lines to zero; don't hide them + floor = np.log10(I_j[pos]).min() if pos.any() else 0.0 + axes[1].plot(dw[~pos], np.full(int((~pos).sum()), floor), "x", ms=6, color="0.5") + axes[1].annotate( + f"{int((~pos).sum())} line(s) at I=0 (x)", + xy=(0.02, 0.06), + xycoords="axes fraction", + fontsize=8, + color="0.4", + ) + axes[1].set_ylabel(r"$\log_{10} I_j$") + + axes[2].plot(dw, phases, "o", ms=4, color="#2ca02c") + axes[2].set_ylabel(r"phase $\phi_j$ [rad]") + axes[2].set_xlabel(r"$\delta\omega_j/\omega_0$") + axes[2].set_ylim(-0.25, 2 * np.pi + 0.25) + axes[2].set_yticks([0, np.pi / 2, np.pi, 3 * np.pi / 2, 2 * np.pi]) + axes[2].set_yticklabels(["0", r"$\pi/2$", r"$\pi$", r"$3\pi/2$", r"$2\pi$"]) + + for ax in axes: + ax.grid(alpha=0.3) + ax.axvline(0.0, color="0.6", lw=0.8, ls="--") + + os.makedirs(out_dir, exist_ok=True) + # no tight_layout here — constrained_layout (set on the figure) already sized + # the title band; calling both would re-reserve a fixed strip and reopen the gap + fig.savefig(os.path.join(out_dir, f"{field}-lines.png"), bbox_inches="tight", dpi=150) + plt.close(fig) + + def post_process(result: Solution, cfg: dict, td: str, args: dict): """Write binary output and diagnostic plots from a completed Vlasov-1D solve.""" t0 = time() + # Driver line spectra (multi-line drivers only). Guarded: a diagnostics failure + # here must never cost a completed solve its binary output and field/dist plots. + try: + plot_driver_spectra(cfg, td, args) + except Exception as exc: + print(f"[post_process] driver spectrum plot skipped: {type(exc).__name__}: {exc}", flush=True) + # Get species names for directory creation species_names = list(cfg["grid"]["species_grids"].keys()) diff --git a/adept/_vlasov1d/modules.py b/adept/_vlasov1d/modules.py index 5d279cdc..9a6f2df8 100644 --- a/adept/_vlasov1d/modules.py +++ b/adept/_vlasov1d/modules.py @@ -337,8 +337,19 @@ def init_diffeqsolve(self): def __call__(self, trainable_modules: dict, args: dict | None = None): """Run the configured Vlasov-1D solve and return the raw Diffrax result.""" - if args is None: - args = self.args + # Merge rather than replace, so a caller passing partial args cannot + # accidentally drop the "drivers" entry that the field pushers + # unconditionally read from args. + args = self.args | args if args is not None else self.args + + # The trainable-module loop (lpse2d pattern): each module may transform the + # initial state and/or inject its (possibly optimizer-updated) leaves into + # args before the solve. Callers may pass None (the pre-module-era calling + # convention) -- treat it as "no trainable modules". + state = self.state + for name, module in (trainable_modules or {}).items(): + state, args = module(state, args) + grid = self.simulation.grid solver_result = diffeqsolve( terms=self.diffeqsolve_quants["terms"], @@ -347,7 +358,7 @@ def __call__(self, trainable_modules: dict, args: dict | None = None): t1=self.time_quantities["t1"], max_steps=grid.max_steps, dt0=grid.dt, - y0=self.state, + y0=state, args=args, saveat=SaveAt(**self.diffeqsolve_quants["saveat"]), progress_meter=TqdmProgressMeter(refresh_steps=grid.max_steps // 100) diff --git a/adept/_vlasov1d/simulation.py b/adept/_vlasov1d/simulation.py index 8623408c..b3d8f4b5 100644 --- a/adept/_vlasov1d/simulation.py +++ b/adept/_vlasov1d/simulation.py @@ -1,15 +1,17 @@ """Domain objects that represent a configured Vlasov-1D simulation.""" -import math import warnings import equinox as eqx import jax +import jax.numpy as jnp import numpy as np -from jax import numpy as jnp +from jax import tree_util as jtu +from jaxtyping import Array from adept._vlasov1d.datamodel import ( AKWDriverConfig, + BroadbandConfig, EMDriverConfig, EMDriverSetConfig, IntensityWavelengthDriverConfig, @@ -31,7 +33,7 @@ SpaceTimeEnvelopeFunction, UniformFunction, ) -from adept.normalization import UREG, PlasmaNormalization, normalize +from adept.normalization import PlasmaNormalization class EMDriver(eqx.Module): @@ -40,16 +42,19 @@ class EMDriver(eqx.Module): 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": + 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) 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(): @@ -62,35 +67,144 @@ 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" + broadband_driver = BroadbandDriver(params.model_dump(), a0, k0, w0, envelope, is_point) + return broadband_driver - # 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) + case _: + raise NotImplementedError(f"Unsupported driver params type: {type(cfg.params).__name__}") - dw0 = 0.0 # ??? - is_point = cfg.source_type == "point" - return EMDriver(a0, k0, w0, dw0, envelope, is_point_source=is_point) +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 + 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 + + 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 + + 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.intensity_weights = jnp.array(int_rng.uniform(int_lo, int_hi, n_colors)) + elif self.params["intensities"]["init"] == "uniform": + self.intensity_weights = jnp.ones(n_colors) + else: + raise NotImplementedError(f"Initialization type -- {self.params['intensities']['init']} -- not implemented") + # 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; + # 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"] + else: + raise NotImplementedError(f"Initialization type -- {self.params['phases']['init']} -- not implemented") + + 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": + 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): + """ + Get the partition spec for the model + + Depends what is learned based on the driver being passed in + + Returns + ------- + filter_spec : pytree with the same structure as the model + + """ + 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.intensity_weights) + # # intensities = intensities / jnp.sum(intensities) + + # args["drivers"]["ey"] = { + # "delta_omega": self.delta_omega, + # "phases": self.phases, + # "amplitudes": self.amplitudes, + # } | self.envelope + + # return state, args class StochasticDriver(eqx.Module): @@ -150,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 @@ -158,8 +272,16 @@ def from_config( cfg: EMDriverSetConfig, norm: PlasmaNormalization | None = None, grid: Grid | None = None ) -> "EMDriverSet": """Build normalized Ex and Ey driver lists (and optional stochastic forcing) 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]) + ex_stochastic = None if cfg.ex_stochastic is not None: if grid is None: diff --git a/adept/_vlasov1d/solvers/pushers/field.py b/adept/_vlasov1d/solvers/pushers/field.py index e28673ae..6eb3dfb6 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 BroadbandDriver, EMDriver class LongitudinalElectricFieldDriver: @@ -54,40 +55,59 @@ 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 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, mask, scale, current_time): + 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 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[:, 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) + # (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 0b69f02e..116323e4 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 @@ -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_absorbing_wave.py b/tests/test_vlasov1d/test_absorbing_wave.py index 659da30d..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 @@ -49,12 +49,17 @@ 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 = {} + # 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 new file mode 100644 index 00000000..2ae4bf0f --- /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 +import jax +from jax import numpy as jnp +from pydantic import ValidationError + +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() +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=r"phases\.seed"): + _comb(4, 0.001, phases={"init": "random"}) + 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"}) + + +@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