From d8196eaaf4aab5bd03401dfcc6bbfb7ca6e96d9e Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Thu, 23 Jul 2026 14:10:03 -0700 Subject: [PATCH 1/4] lpse2d: add SRS (Raman backscatter) support - New raman.py light-wave module and SRS coupling terms in epw.py - Raman seed/light-wave config in datamodel and helpers - SRS example config (configs/envelope-2d/srs.yaml) and test - Document SRS options in lpse2d config docs Co-Authored-By: Claude Fable 5 --- adept/_lpse2d/core/epw.py | 55 +++++++++ adept/_lpse2d/core/laser.py | 2 +- adept/_lpse2d/core/raman.py | 163 +++++++++++++++++++++++++++ adept/_lpse2d/core/vector_field.py | 22 +++- adept/_lpse2d/datamodel.py | 22 ++++ adept/_lpse2d/helpers.py | 110 ++++++++++++++---- configs/envelope-2d/srs.yaml | 87 ++++++++++++++ docs/source/solvers/lpse2d/config.md | 38 ++++++- tests/test_lpse2d/configs/srs.yaml | 72 ++++++++++++ tests/test_lpse2d/test_srs.py | 144 +++++++++++++++++++++++ 10 files changed, 686 insertions(+), 29 deletions(-) create mode 100644 adept/_lpse2d/core/raman.py create mode 100644 configs/envelope-2d/srs.yaml create mode 100644 tests/test_lpse2d/configs/srs.yaml create mode 100644 tests/test_lpse2d/test_srs.py diff --git a/adept/_lpse2d/core/epw.py b/adept/_lpse2d/core/epw.py index df5ceb4e..54651eee 100644 --- a/adept/_lpse2d/core/epw.py +++ b/adept/_lpse2d/core/epw.py @@ -346,6 +346,22 @@ def __init__(self, cfg: dict): if self.tpd_enabled: self.tpd_prefactor = 1j * self.e / (8.0 * self.wp0 * self.me) + # SRS parameters + self.srs_enabled = cfg["terms"]["epw"]["source"].get("srs", False) + if self.srs_enabled: + self.w1 = cfg["units"]["derived"]["w1"] + self.c = cfg["units"]["derived"]["c"] + # MATLAB line 2073: srsSourceTerm = 1i * e * wp0/(4*me*w0*w1) .* (1 + dn) .* E0_dot_E1 + self.srs_prefactor = 1j * self.e * self.wp0 / (4.0 * self.me * self.w0 * self.w1) + # high-k filter for the light fields entering the source product + # (MATLAB isSuppressHighKSource, lines 637-645): only wavevectors near the + # light-wave envelope produce physically-realistic SRS + max_source_k_multiplier = 1.2 + n_min = float(np.min(np.array(self.background_density))) + max_k1_sq = max_source_k_multiplier**2 * max(1.0 - n_min * self.w0**2 / self.w1**2, 0.0) + is_outside_max_k1 = self.k_sq * (self.c / self.w1) ** 2 > max_k1_sq + self.E1_filter = jnp.where(is_outside_max_k1, 0.0, 1.0)[..., None] + # Noise parameters self.noise_enabled = cfg["terms"]["epw"]["source"]["noise"] self.noise_amplitude = 1e-10 # matches MATLAB noiseAmp (m201805_matlabLpse_v11.m:49) @@ -504,6 +520,33 @@ def calc_tpd_source(self, t: float, phi_k: Array, ey: Array, E0_y: Array) -> Arr return source + def calc_srs_source(self, E0: Array, E1: Array) -> Array: + """ + Calculate the SRS source term for the EPW potential. + + Matches MATLAB lines 2052-2078 for isSolveForPotential=true: + E0_dot_E1 = E0 . conj(E1) (E1 high-k filtered first, evaluate_E0_dot_E1 lines 2302-2354) + srsSource = 1i * e * wp0/(4*me*w0*w1) * (1 + dn) * E0_dot_E1 + srsSource -> k-space + + The pump is static/prescribed here, so only E1 is filtered (in MATLAB the E0 + filter is skipped on the static-laser path, line 2308). + + Args: + E0: Pump field (shape: nx, ny, 2) + E1: Raman field (shape: nx, ny, 2) + + Returns: + SRS source term in k-space + """ + E1_filtered = jnp.fft.ifft2(jnp.fft.fft2(E1, axes=(0, 1)) * self.E1_filter, axes=(0, 1)) + E0_dot_E1 = E0[..., 0] * jnp.conj(E1_filtered[..., 0]) + E0[..., 1] * jnp.conj(E1_filtered[..., 1]) + + # (1 + backgroundDensityPerturbation) = n / n_envelope + source = self.srs_prefactor * self.background_density / self.envelope_density * E0_dot_E1 + + return jnp.fft.fft2(source) + def get_noise(self, t: float) -> Array: """ Generate random noise for plasma waves. @@ -606,6 +649,11 @@ def __call__(self, t: float, y, args) -> Array: E0_y = E0[..., 1] # y-component of laser field tpd_source = self.calc_tpd_source(t, phi_k, ey, E0_y) + srs_source = None + if self.srs_enabled: + # MATLAB lines 2052-2078 + srs_source = self.calc_srs_source(E0, y["E1"]) + # ======================================================================== # STEP 7: Apply density gradient to E fields (in REAL space) # ======================================================================== @@ -641,4 +689,11 @@ def __call__(self, t: float, y, args) -> Array: # MATLAB line 2109: divE = divE + tpdSourceTerm * DT phi_k = phi_k + self.dt * tpd_source + # ======================================================================== + # STEP 11: Add SRS source + # ======================================================================== + if self.srs_enabled and srs_source is not None: + # MATLAB line 2113: divE = divE + srsSourceTerm * DT + phi_k = phi_k + self.dt * srs_source + return phi_k diff --git a/adept/_lpse2d/core/laser.py b/adept/_lpse2d/core/laser.py index 76349ea9..f2830b9e 100644 --- a/adept/_lpse2d/core/laser.py +++ b/adept/_lpse2d/core/laser.py @@ -23,7 +23,7 @@ def __init__(self, cfg) -> None: self.speckle_normalization = 1.0 self.y_si = None # y-coordinates in meters - speckle_profile = cfg["drivers"]["E0"].get("speckle_profile") + speckle_profile = cfg["drivers"].get("E0", {}).get("speckle_profile") if speckle_profile is not None: # Convert y-coordinates to SI units (meters) diff --git a/adept/_lpse2d/core/raman.py b/adept/_lpse2d/core/raman.py new file mode 100644 index 00000000..623f5625 --- /dev/null +++ b/adept/_lpse2d/core/raman.py @@ -0,0 +1,163 @@ +import numpy as np +from jax import Array, lax +from jax import numpy as jnp + + +class RamanLight: + """ + Evolves the Raman scattered-light envelope E1. + + This is a port of the MATLAB `raman.solver = 'fd'` branch of m201805_matlabLpse_v11.m + (evalLaserFieldUpdate, lines 1656-1704, and lightSplitStep, lines 1377-1425). + + The envelope equation (per component, 2D with cross-derivative terms) is + + dE1/dt = i c^2/(2 w1) * (transverse Laplacian) E1 + + i w1/2 * (1 - wp0^2/w1^2 * n/n_env) * E1 + - i e/(4 w0 me) * conj(laplacian phi) * E0 (SRS coupling) + + seed injection source (optional) + + where w1 = w0 - wp0 is the Raman envelope frequency and laplacian phi is computed + spectrally from the EPW potential (MATLAB line 1604). + + Time integration is the same staggered explicit scheme as MATLAB's lightSplitStep: + the real part is updated with the RHS evaluated at t, then the imaginary part with + the RHS evaluated at t + dt/2. Because this scheme is only conditionally stable + (dt < ~dx^2 w1 / c^2), the update is sub-cycled `light_substeps` times inside each + EPW step; the EPW potential is held fixed during the sub-steps, which matches + MATLAB's `lightStepsPerEpwStep` behavior. + """ + + def __init__(self, cfg: dict): + self.cfg = cfg + derived = cfg["units"]["derived"] + self.c = derived["c"] + self.w0 = derived["w0"] + self.w1 = derived["w1"] + self.wp0 = derived["wp0"] + self.e = derived["e"] + self.me = derived["me"] + self.envelope_density = cfg["units"]["envelope density"] + + self.dx = cfg["grid"]["dx"] + self.dy = cfg["grid"]["dy"] + self.dt = cfg["grid"]["dt"] # outer (EPW) step + self.n_sub = cfg["grid"]["light_substeps"] + self.dt_l = self.dt / self.n_sub # light sub-step + self.x = cfg["grid"]["x"] + self.y = cfg["grid"]["y"] + self.k_sq = cfg["grid"]["kx"][:, None] ** 2 + cfg["grid"]["ky"][None, :] ** 2 + + background_density = cfg["grid"]["background_density"] + # local detuning of the Raman envelope (MATLAB line 1668-1670) + self.linear_coeff = ( + 1j * self.w1 / 2.0 * (1.0 - self.wp0**2 / self.w1**2 * background_density / self.envelope_density) + ) + self.diffraction_coeff = 1j * self.c**2 / (2.0 * self.w1) + self.srs_coeff = -1j * self.e / (4.0 * self.w0 * self.me) + + # absorbing boundaries are applied every sub-step so that light (group velocity ~ c) + # cannot cross the absorber between damping applications + self.sub_boundary = cfg["grid"]["absorbing_boundaries"] ** (1.0 / self.n_sub) + + # seed injection (MATLAB lines 1757-1769): a two-point antisymmetric source that + # launches a leftward-propagating (-x) wave at x = xmax - offset + if "E1" in cfg["drivers"]: + seed = cfg["drivers"]["E1"]["derived"] + x_inject = cfg["grid"]["xmax"] - seed["offset"] + self.i1 = int(np.argmin(np.abs(np.array(self.x) - x_inject))) + wpe_i1 = self.w0 * np.sqrt(background_density[self.i1, 0]) + self.wpe_sq_i1 = float(wpe_i1**2) + permittivity1 = 1.0 - self.wpe_sq_i1 / self.w1**2 + if permittivity1 <= 0: + raise ValueError( + f"The Raman seed injector at x = {float(self.x[self.i1]):.2f} um sits at density " + f"{float(background_density[self.i1, 0]):.3f} nc, above the w1 critical density " + f"{(self.w1 / self.w0) ** 2:.3f} nc where the seed is evanescent. Lower density.max, " + "or move the injector with drivers.E1.offset, or remove drivers.E1 to run noise-seeded." + ) + self.source_prefactor = self.c**2 / (2.0 * self.w1) / permittivity1**0.25 / self.dx**2 + self.seed_enabled = True + else: + self.seed_enabled = False + + def _d2x(self, f: Array) -> Array: + return (jnp.roll(f, -1, axis=0) - 2.0 * f + jnp.roll(f, 1, axis=0)) / self.dx**2 + + def _d2y(self, f: Array) -> Array: + return (jnp.roll(f, -1, axis=1) - 2.0 * f + jnp.roll(f, 1, axis=1)) / self.dy**2 + + def _dxdy(self, f: Array) -> Array: + return ( + jnp.roll(f, (-1, -1), axis=(0, 1)) + - jnp.roll(f, (1, -1), axis=(0, 1)) + - jnp.roll(f, (-1, 1), axis=(0, 1)) + + jnp.roll(f, (1, 1), axis=(0, 1)) + ) / (4.0 * self.dx * self.dy) + + def calc_seed_source(self, t: float, seed_args: dict) -> tuple[Array, Array]: + """ + Amplitude and phases for the two-point seed injector (MATLAB lines 1757-1769). + + Returns the two rows to be added to the E1_y RHS at self.i1 and self.i1 + 1. + """ + dw1 = seed_args["delta_omega"] + turn_on = 1.0 - jnp.exp(-((t / seed_args["turn_on_time"]) ** 2)) + amp = self.source_prefactor * seed_args["amplitude"] * turn_on + + if seed_args["yw"] > 0: + envelope_y = jnp.exp(-((self.y / (seed_args["yw"] / 2.0)) ** 4)) + else: + envelope_y = jnp.ones_like(self.y) + + # local seed wavenumber (MATLAB line 867) + k1 = self.w1 / self.c * jnp.sqrt((1.0 + dw1) ** 2 - self.wpe_sq_i1 / self.w1**2) + + row_i1 = -1j * amp * envelope_y * jnp.exp(-1j * k1 * self.x[self.i1 + 1] - 1j * self.w1 * dw1 * t) + row_i1p1 = 1j * amp * envelope_y * jnp.exp(-1j * k1 * self.x[self.i1] - 1j * self.w1 * dw1 * t) + return row_i1, row_i1p1 + + def rhs(self, t: float, E1: Array, E0: Array, laplacian_phi: Array, seed_args: dict | None) -> Array: + e1x, e1y = E1[..., 0], E1[..., 1] + + # paraxial propagation with cross-derivative terms (MATLAB lines 1663-1671) + k_e1x = self.diffraction_coeff * (self._d2y(e1x) - self._dxdy(e1y)) + self.linear_coeff * e1x + k_e1y = self.diffraction_coeff * (self._d2x(e1y) - self._dxdy(e1x)) + self.linear_coeff * e1y + + # SRS coupling to the EPW (MATLAB lines 1684-1689, potential formulation) + k_e1x += self.srs_coeff * jnp.conj(laplacian_phi) * E0[..., 0] + k_e1y += self.srs_coeff * jnp.conj(laplacian_phi) * E0[..., 1] + + if seed_args is not None: + row_i1, row_i1p1 = self.calc_seed_source(t, seed_args) + k_e1y = k_e1y.at[self.i1, :].add(row_i1) + k_e1y = k_e1y.at[self.i1 + 1, :].add(row_i1p1) + + return jnp.stack([k_e1x, k_e1y], axis=-1) + + def __call__(self, t: float, E1: Array, E0_fn, phi_k: Array, seed_args: dict | None) -> Array: + """ + Advance E1 over one EPW step (self.n_sub staggered light sub-steps). + + :param t: time at the start of the EPW step + :param E1: Raman field, shape (nx, ny, 2), complex + :param E0_fn: callable t -> pump field of shape (nx, ny, 2) + :param phi_k: EPW potential in k-space, held fixed during the sub-steps + :param seed_args: derived driver parameters for the seed, or None + """ + seed_args = seed_args if self.seed_enabled else None + laplacian_phi = jnp.fft.ifft2(-self.k_sq * phi_k) + + def substep(i, E1): + t_i = t + i * self.dt_l + # real-part update with the RHS at t_i (MATLAB lines 1380-1397) + k1 = self.rhs(t_i, E1, E0_fn(t_i), laplacian_phi, seed_args) + E1 = E1 + self.dt_l * jnp.real(k1) + # imaginary-part update with the RHS at t_i + dt/2 (MATLAB lines 1400-1421) + k2 = self.rhs(t_i + self.dt_l / 2.0, E1, E0_fn(t_i + self.dt_l / 2.0), laplacian_phi, seed_args) + E1 = E1 + 1j * self.dt_l * jnp.imag(k2) + # absorbing boundaries (MATLAB lines 977-983) + E1 = E1 * self.sub_boundary[..., None] + return E1 + + return lax.fori_loop(0, self.n_sub, substep, E1) diff --git a/adept/_lpse2d/core/vector_field.py b/adept/_lpse2d/core/vector_field.py index 1c8e1930..6ac08a80 100644 --- a/adept/_lpse2d/core/vector_field.py +++ b/adept/_lpse2d/core/vector_field.py @@ -4,6 +4,7 @@ from adept._base_ import get_envelope from adept._lpse2d.core import epw, laser +from adept._lpse2d.core.raman import RamanLight class SplitStep: @@ -24,6 +25,8 @@ def __init__(self, cfg): # self.epw = epw.SpectralPotential(cfg) self.epw = epw.SpectralEPWSolver(cfg) self.light = laser.Light(cfg) + # the Raman scattered light is evolved iff the SRS source term is on + self.raman = RamanLight(cfg) if cfg["terms"]["epw"]["source"].get("srs", False) else None self.complex_state_vars = ["E0", "epw", "E1"] self.boundary_envelope = cfg["grid"]["absorbing_boundaries"] self.one_over_ksq = cfg["grid"]["one_over_ksq"] @@ -66,10 +69,23 @@ def get_envelope_coefficient(self, envelope_args, t): def light_split_step(self, t, y, driver_args): if "E0" in driver_args: - t_coeff = self.get_envelope_coefficient(driver_args["E0"], t) - y["E0"] = t_coeff * self.light.laser_update(t, y, driver_args["E0"]) - y["E1"] *= self.boundary_envelope[..., None] + def E0_fn(this_t): + t_coeff = self.get_envelope_coefficient(driver_args["E0"], this_t) + return t_coeff * self.light.laser_update(this_t, y, driver_args["E0"]) + + y["E0"] = E0_fn(t) + else: + E0_now = y["E0"] + + def E0_fn(this_t): + return E0_now + + if self.raman is not None: + # evolve the Raman light; absorbing boundaries are applied per light sub-step inside + y["E1"] = self.raman(t, y["E1"], E0_fn, y["epw"], driver_args.get("E1")) + else: + y["E1"] *= self.boundary_envelope[..., None] return y diff --git a/adept/_lpse2d/datamodel.py b/adept/_lpse2d/datamodel.py index d1f25072..f5c3846e 100644 --- a/adept/_lpse2d/datamodel.py +++ b/adept/_lpse2d/datamodel.py @@ -81,6 +81,23 @@ class E0DriverModel(BaseModel): speckle: SpeckleModel | None = None +class E1DriverModel(BaseModel): + """ + Raman seed driver. + + Injects a counter-propagating (-x) scattered-light wave at x = xmax - offset + with the given (vacuum) intensity. Only used when terms.epw.source.srs is on. + """ + + intensity: str # e.g. "1.0e+12W/cm^2" + delta_omega: float = 0.0 # seed frequency shift relative to w1 = w0 - wp0 (fraction of w1) + turn_on_time: str = "10fs" + # distance of the injector from the right boundary; defaults to 1.6 * boundary_width, + # which places it just inside the absorbing boundary's tanh skirt + offset: str | None = None + yw: str | None = None # super-Gaussian width of the seed in y; omit for uniform in y + + class DriversModel(BaseModel): """ Define the drivers for the simulation @@ -88,6 +105,7 @@ class DriversModel(BaseModel): """ E0: E0DriverModel + E1: E1DriverModel | None = None class GridModel(BaseModel): @@ -105,6 +123,9 @@ class GridModel(BaseModel): tmin: str ymax: str ymin: str + # number of Raman-light sub-steps per EPW step (SRS only); computed from the + # stability limit if omitted + light_substeps: int | None = None class TimeSaveModel(BaseModel): @@ -140,6 +161,7 @@ class DampingModel(BaseModel): class SourceModel(BaseModel): noise: bool tpd: bool + srs: bool = False class EPWModel(BaseModel): diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index 66623791..c22f7c32 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -255,9 +255,56 @@ def get_derived_quantities(cfg: dict) -> dict: cfg_grid["max_steps"] = cfg_grid["nt"] + 2048 + # SRS: the Raman light is advanced with an explicit conditionally-stable scheme, so it is + # sub-cycled within each EPW step. The stability bound follows MATLAB line 500 + # (dt_max_seed), generalized to 2D + if cfg["terms"]["epw"]["source"].get("srs", False): + derived = cfg["units"]["derived"] + wpe_max_sq = derived["w0"] ** 2 * max(cfg["density"].get("max", 1.0), cfg["density"].get("min", 0.0)) + dt_max = 1.0 / ( + 2.0 * derived["c"] ** 2 / (cfg_grid["dx"] ** 2 * derived["w1"]) + + np.abs(derived["w1"] ** 2 - wpe_max_sq) / (4.0 * derived["w1"]) + ) + if "light_substeps" in cfg_grid: + n_sub = int(cfg_grid["light_substeps"]) + if cfg_grid["dt"] / n_sub > dt_max: + raise ValueError( + f"grid.light_substeps = {n_sub} gives a light step of {cfg_grid['dt'] / n_sub:.2e} ps " + f"which exceeds the Raman stability limit of {dt_max:.2e} ps" + ) + else: + n_sub = int(np.ceil(cfg_grid["dt"] / (0.9 * dt_max))) + cfg_grid["light_substeps"] = n_sub + print(f"SRS is on -- the Raman light is sub-cycled {n_sub}x per EPW step (dt_light limit {dt_max:.2e} ps)") + # change driver parameters to the right units for k in cfg["drivers"].keys(): cfg["drivers"][k]["derived"] = {} + if k == "E1": + # Raman seed injector -- different parameter set than the envelope drivers + c_cgs = 2.99792458e10 + seed_intensity = _Q(cfg["drivers"][k]["intensity"]).to("W/cm^2").value + # the injector must sit clear of the absorbing boundary, whose tanh skirt + # (rise = boundary_width / 5) extends past xmax - boundary_width into the box + boundary_width = _Q(cfg["grid"]["boundary_width"]).to("um").value + min_offset = 1.6 * boundary_width + if "offset" in cfg["drivers"][k]: + offset = _Q(cfg["drivers"][k]["offset"]).to("um").value + if offset < min_offset: + print( + f"WARNING: drivers.E1.offset = {offset}um is inside the absorbing-boundary skirt " + f"(< 1.6 * boundary_width = {min_offset}um); the seed will be damped at the source" + ) + else: + offset = min_offset + cfg["drivers"][k]["derived"] = { + "amplitude": np.sqrt(8 * np.pi * seed_intensity * 1e7 / c_cgs) / cfg["units"]["derived"]["fieldScale"], + "delta_omega": float(cfg["drivers"][k].get("delta_omega", 0.0)), + "turn_on_time": _Q(cfg["drivers"][k].get("turn_on_time", "10fs")).to("ps").value, + "offset": offset, + "yw": _Q(cfg["drivers"][k]["yw"]).to("um").value if "yw" in cfg["drivers"][k] else 0.0, + } + continue cfg["drivers"][k]["derived"]["tw"] = _Q(cfg["drivers"][k]["envelope"]["tw"]).to("ps").value cfg["drivers"][k]["derived"]["tc"] = _Q(cfg["drivers"][k]["envelope"]["tc"]).to("ps").value cfg["drivers"][k]["derived"]["tr"] = _Q(cfg["drivers"][k]["envelope"]["tr"]).to("ps").value @@ -434,7 +481,7 @@ def get_density_profile(cfg: dict) -> Array: :param cfg: Dict """ if cfg["density"]["basis"] == "uniform": - nprof = np.ones((cfg["grid"]["nx"], cfg["grid"]["ny"])) + nprof = cfg["density"].get("val", 1.0) * np.ones((cfg["grid"]["nx"], cfg["grid"]["ny"])) elif cfg["density"]["basis"] == "linear": left = cfg["grid"]["xmin"] + _Q("5.0um").to("um").value @@ -486,20 +533,20 @@ def plot_fields(fields, td): t_skip = t_skip if t_skip > 1 else 1 tslice = slice(0, -1, t_skip) - dx = fields.coords["x (um)"].data[1] - fields.coords["x (um)"].data[0] - dy = fields.coords["y (um)"].data[1] - fields.coords["y (um)"].data[0] + ny = fields.coords["y (um)"].data.size for k, v in fields.items(): fld_dir = os.path.join(td, "plots", k) os.makedirs(fld_dir) - np.abs(v[tslice]).T.plot(col="t (ps)", col_wrap=4) - plt.savefig(os.path.join(fld_dir, f"{k}_x.png"), bbox_inches="tight") - plt.close() + if ny > 1: + np.abs(v[tslice]).T.plot(col="t (ps)", col_wrap=4) + plt.savefig(os.path.join(fld_dir, f"{k}_x.png"), bbox_inches="tight") + plt.close() - np.real(v[tslice]).T.plot(col="t (ps)", col_wrap=4) - plt.savefig(os.path.join(fld_dir, f"{k}_x_r.png"), bbox_inches="tight") - plt.close() + np.real(v[tslice]).T.plot(col="t (ps)", col_wrap=4) + plt.savefig(os.path.join(fld_dir, f"{k}_x_r.png"), bbox_inches="tight") + plt.close() # fig, ax = plt.subplots(1, 1, figsize=(10, 4)) # np.abs(v[:, 1, 0]).plot(ax=ax) @@ -551,14 +598,17 @@ def plot_kt(kfields, td): kx_slice = slice(ikx_min, ikx_max) ky_slice = slice(iky_min, iky_max) + n_ky = kfields.coords[r"ky ($kc\omega_0^{-1}$)"].data.size for k, v in kfields.items(): fld_dir = os.path.join(td, "plots", k) os.makedirs(fld_dir, exist_ok=True) - # np.log10(np.abs(v[tslice, kx_slice, 0])).T.plot(col="t (ps)", col_wrap=4) - # plt.savefig(os.path.join(fld_dir, f"log_{k}_kx_k0_absmax{abs_kmax}.png"), bbox_inches="tight") - # plt.close() + if n_ky == 1: + np.log10(np.abs(v[tslice, kx_slice, 0])).plot(col="t (ps)", col_wrap=4) + plt.savefig(os.path.join(fld_dir, f"log_{k}_kx_absmax{abs_kmax}.png"), bbox_inches="tight") + plt.close() + continue np.abs(v[tslice, kx_slice, ky_slice]).T.plot(col="t (ps)", col_wrap=4) plt.savefig(os.path.join(fld_dir, f"{k}_kx_ky_absmax{abs_kmax}.png"), bbox_inches="tight") @@ -601,14 +651,7 @@ def plot_series(series, td): def make_series_xarrays(cfg, this_t, state, td): - esq = state["e_sq"] - max_phi = state["max_phi"] - series_xr = xr.Dataset( - { - "e_sq": xr.DataArray(esq, coords=(("t (ps)", this_t),)), - "max_phi": xr.DataArray(max_phi, coords=(("t (ps)", this_t),)), - } - ) + series_xr = xr.Dataset({k: xr.DataArray(v, coords=(("t (ps)", this_t),)) for k, v in state.items()}) series_xr.to_netcdf(os.path.join(td, "binary", "series.xr"), engine="h5netcdf", invalid_netcdf=True) return series_xr @@ -638,8 +681,8 @@ def make_field_xarrays(cfg, this_t, state, td): xax_tuple = ("x (um)", xax) yax_tuple = ("y (um)", yax) - # check if state["epw"] is a complex64 or 128 and choose accordingly - if state["epw"].dtype == jnp.complex128: + # the state is stored as a float view of a complex array; pick the matching complex dtype + if state["epw"].dtype in (np.float64, np.complex128): _complex = np.complex128 else: _complex = np.complex64 @@ -784,6 +827,20 @@ def save_func(t, y, args): def get_default_save_func(cfg): + srs_on = cfg["terms"]["epw"]["source"].get("srs", False) + if srs_on: + # Reflectivity probe: sample |E1_y|^2 on the low-density side, just inside the absorber. + # R = sqrt(eps1) * <|E1_y|^2>_y / E0_source^2 where eps1 = 1 - wpe^2/w1^2 accounts for + # the reduced group velocity (Poynting flux ~ sqrt(eps) |E|^2) relative to the vacuum pump. + boundary_width = _Q(cfg["grid"]["boundary_width"]).to("um").value + x_probe = cfg["grid"]["xmin"] + 1.6 * boundary_width # clear of the absorber's tanh skirt + ix_probe = int(np.argmin(np.abs(np.array(cfg["grid"]["x"]) - x_probe))) + w0 = cfg["units"]["derived"]["w0"] + w1 = cfg["units"]["derived"]["w1"] + n_probe = float(np.mean(np.array(cfg["grid"]["background_density"])[ix_probe, :])) + sqrt_eps1 = np.sqrt(max(1.0 - n_probe * w0**2 / w1**2, 0.0)) + E0_source_sq = cfg["units"]["derived"]["E0_source"] ** 2 + def save_func(t, y, args): phi_k = y["epw"].view(jnp.complex128) ex = -1j * cfg["grid"]["kx"][:, None] * phi_k @@ -792,6 +849,13 @@ def save_func(t, y, args): ey = jnp.fft.ifft2(ey) e_sq = jnp.abs(ex) ** 2 + jnp.abs(ey) ** 2 - return {"e_sq": jnp.sum(e_sq * cfg["grid"]["dx"] * cfg["grid"]["dy"]), "max_phi": jnp.max(jnp.abs(phi_k))} + out = {"e_sq": jnp.sum(e_sq * cfg["grid"]["dx"] * cfg["grid"]["dy"]), "max_phi": jnp.max(jnp.abs(phi_k))} + + if srs_on: + e1 = y["E1"].view(jnp.complex128) + out["e1_sq"] = jnp.mean(jnp.sum(jnp.abs(e1) ** 2, axis=-1)) + out["reflectivity"] = sqrt_eps1 * jnp.mean(jnp.abs(e1[ix_probe, :, 1]) ** 2) / E0_source_sq + + return out return {"t": {"ax": cfg["grid"]["t"]}, "func": save_func} diff --git a/configs/envelope-2d/srs.yaml b/configs/envelope-2d/srs.yaml new file mode 100644 index 00000000..bd9162a0 --- /dev/null +++ b/configs/envelope-2d/srs.yaml @@ -0,0 +1,87 @@ +# Backward SRS driven by a static pump on a linear density ramp. +# This is the JAX translation of the lpse-matlab `srs_1D` case (m201805_matlabLpse_v11.m) +# run as a quasi-1D 2D box: the pump propagates in +x, the Raman light E1 is evolved with a +# finite-difference paraxial solver and grows from the EPW noise (and the optional E1 seed +# below), and the reflectivity is recorded at a probe on the low-density side. +density: + basis: linear + gradient scale length: 50um + max: 0.28 + min: 0.18 + noise: + max: 1.0e-09 + min: 1.0e-10 + type: uniform +drivers: + E0: + params: + phases: + seed: 42 + shape: uniform + delta_omega_max: 0.015 + num_colors: 1 + envelope: + tw: 40ps + tr: 0.1ps + tc: 20.25ps + xr: 0.2um + xw: 1000um + xc: 50um + yr: 0.2um + yw: 1000um + yc: 50um + # Optional Raman seed. Only usable when the density at the injector (x = xmax - offset) + # is below the w1 critical density (n < 0.25 nc for envelope density 0.25) -- on this + # 0.18-0.28 ramp the seed would be evanescent, so SRS grows from the EPW noise instead + # (the MATLAB srs_1D case with I1 = 0). + # E1: + # intensity: 1.0e+12W/cm^2 + # delta_omega: 0.0 + # turn_on_time: 10fs +grid: + boundary_abs_coeff: 1.0e4 + boundary_width: 3um + low_pass_filter: 0.6 + dt: 1fs + dx: 50nm + tmax: 10.0ps + tmin: 0.0ns + ymax: 3um + ymin: -3um +mlflow: + experiment: srs + run: srs-test +save: + fields: + t: + dt: 0.5ps + tmax: 10ps + tmin: 0ps + x: + dx: 80nm + y: + dy: 128nm +solver: envelope-2d +terms: + epw: + boundary: + x: absorbing + y: periodic + damping: + collisions: 0.1 + landau: true + density_gradient: true + linear: true + source: + noise: true + tpd: false + srs: true + zero_mask: true +units: + atomic number: 1 + envelope density: 0.25 + ionization state: 1 + laser intensity: 1.5e+15W/cm^2 + laser_wavelength: 351nm + reference electron temperature: 2000.0eV + reference ion temperature: 1000eV diff --git a/docs/source/solvers/lpse2d/config.md b/docs/source/solvers/lpse2d/config.md index a02235f4..ac5a6f56 100644 --- a/docs/source/solvers/lpse2d/config.md +++ b/docs/source/solvers/lpse2d/config.md @@ -117,9 +117,12 @@ Simulation grid parameters. Note: Grid values use physical units as strings. | `tmin` | string | Start time with unit | | `ymax` | string | Domain maximum y with unit | | `ymin` | string | Domain minimum y with unit | +| `light_substeps` | int | (SRS only, optional) Raman-light sub-steps per EPW step. Computed from the explicit-scheme stability limit `dt_light < 1 / (2c^2/(dx^2 w1) + \|w1^2 - wpe_max^2\|/(4 w1))` if omitted; a `ValueError` is raised if a user-supplied value violates it | Note: `nx` and `ny` are computed automatically from the grid parameters. The grid is optimized for FFT performance (sizes with small prime factors). +Note: setting `ymax`/`ymin` smaller than `dx` collapses the box to `ny = 1`, which runs the solver in a true 1D mode (useful for cheap 1D SRS simulations; TPD requires 2D). + Example: ```yaml grid: @@ -277,6 +280,34 @@ drivers: w0: 20.0 ``` +### E1 - Raman Seed Driver (Optional) + +Injects a counter-propagating (-x) scattered-light wave for seeded SRS. Only used when +`terms.epw.source.srs` is on. The injector sits at `x = xmax - offset` and drives the +`E1` field with a two-point antisymmetric source (the MATLAB LPSE injector), so the seed +propagates toward the low-density side while backscatter growth amplifies it against the pump. + +| Field | Type | Description | +|-------|------|-------------| +| `intensity` | string | Seed vacuum intensity with unit, e.g. `"1.0e+12W/cm^2"` | +| `delta_omega` | float | Seed frequency shift relative to `w1 = w0 - wp0`, as a fraction of `w1` (default 0) | +| `turn_on_time` | string | Ramp-up time of the injector (default `10fs`) | +| `offset` | string | Distance of the injector from the right boundary. Defaults to `1.6 * boundary_width`, just inside the absorbing boundary's tanh skirt; a warning is printed for smaller values because the seed would be damped at the source | +| `yw` | string | Super-Gaussian (4th order) width of the seed in y; omit for uniform in y | + +The density at the injector must be below the `w1` critical density (`n < 0.25 n_c` for +envelope density 0.25), otherwise the seed is evanescent and setup raises an error. Without +`drivers.E1`, SRS grows from the EPW noise source instead (the noise-seeded configuration). + +Example: +```yaml +drivers: + E1: + intensity: 1.0e+12W/cm^2 + delta_omega: 0.0 + turn_on_time: 10fs +``` + ## terms Physics terms configuration. @@ -319,7 +350,7 @@ Physics terms configuration. |-------|------|-------------| | `noise` | bool | Add random noise source | | `tpd` | bool | Include two-plasmon decay source | -| `srs` | bool | Include stimulated Raman scattering source (optional) | +| `srs` | bool | Include stimulated Raman scattering (optional, default false). Turning this on also evolves the Raman scattered-light field `E1` with a finite-difference paraxial solver, sub-cycled `grid.light_substeps` times per EPW step, and adds the SRS source `i e wp0/(4 me w0 w1) (n/n_env) E0 . conj(E1)` to the EPW potential. The default time series then also records `e1_sq` and `reflectivity` (Poynting-corrected `|E1_y|^2/E0^2` at a probe on the low-density side, `x = 1.6 * boundary_width`) | #### hyperviscosity (optional) @@ -474,4 +505,7 @@ See `configs/envelope-2d/tpd.yaml` - TPD simulation with linear density gradient ### SRS / Reflection -See `configs/envelope-2d/reflection.yaml` - SRS simulation with kinetic corrections. +See `configs/envelope-2d/srs.yaml` - Noise-seeded backward SRS on a linear density ramp +(the lpse-matlab `srs_1D` case), with a reflectivity time series recorded at a probe on +the low-density side. Also see `configs/envelope-2d/reflection.yaml` - SRS simulation +with kinetic corrections. diff --git a/tests/test_lpse2d/configs/srs.yaml b/tests/test_lpse2d/configs/srs.yaml new file mode 100644 index 00000000..d63b456f --- /dev/null +++ b/tests/test_lpse2d/configs/srs.yaml @@ -0,0 +1,72 @@ +# Homogeneous-plasma backward-SRS growth test box. +# Uniform density equal to the envelope density (so there is no detuning), periodic in x, +# noise-seeded. The EPW energy should grow at ~2*gamma0 (see test_srs.py). +density: + basis: uniform + val: 0.2 + noise: + max: 1.0e-09 + min: 1.0e-10 + type: uniform +drivers: + E0: + params: + phases: + seed: 42 + shape: uniform + delta_omega_max: 0.015 + num_colors: 1 + envelope: + tw: 40ps + tr: 0.02ps + tc: 20.05ps + xr: 0.2um + xw: 1000um + xc: 50um + yr: 0.2um + yw: 1000um + yc: 50um +grid: + boundary_abs_coeff: 1.0e4 + boundary_width: 3um + low_pass_filter: 0.6 + dt: 1fs + dx: 50nm + xmax: 80um + tmax: 1.2ps + tmin: 0.0ns + ymax: 0.25um + ymin: -0.25um +mlflow: + experiment: test-lpse + run: srs-test +save: + fields: + t: + dt: 0.1ps + tmax: 1.2ps + tmin: 0ps +solver: envelope-2d +terms: + epw: + boundary: + x: periodic + y: periodic + damping: + collisions: false + landau: true + density_gradient: true + linear: true + source: + noise: true + tpd: false + srs: true + zero_mask: true +units: + atomic number: 1 + envelope density: 0.2 + ionization state: 1 + laser intensity: 1.5e+15W/cm^2 + laser_wavelength: 351nm + reference electron temperature: 2000.0eV + reference ion temperature: 1000eV diff --git a/tests/test_lpse2d/test_srs.py b/tests/test_lpse2d/test_srs.py new file mode 100644 index 00000000..c2e78dfe --- /dev/null +++ b/tests/test_lpse2d/test_srs.py @@ -0,0 +1,144 @@ +"""Tests for the SRS (stimulated Raman scattering) path of the envelope-2d solver. + +These are ports of the physics exercised by the `srs_1D` case of lpse-matlab +(m201805_matlabLpse_v11.m), run as small quasi-1D 2D boxes. +""" + +import numpy as np +import pytest +import yaml + + +def _load_cfg(): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + return cfg + + +def _run(cfg): + from adept import ergoExo + + exo = ergoExo() + modules = exo.setup(cfg) + sol, ppo, mlrunid = exo(modules) + return exo, sol + + +def _predicted_gamma0(cfg_units, cfg_grid): + """ + Homogeneous backward-SRS growth rate gamma0 = k vos/4 * wpe / sqrt(w_ek * w_s) + evaluated at the phase-matched EPW wavenumber k = k0 + k1. + """ + c = cfg_units["derived"]["c"] + w0 = cfg_units["derived"]["w0"] + wp0 = cfg_units["derived"]["wp0"] + vte_sq = cfg_units["derived"]["vte_sq"] + e = cfg_units["derived"]["e"] + me = cfg_units["derived"]["me"] + E0_source = cfg_units["derived"]["E0_source"] + n = cfg_units["envelope density"] # uniform box at the envelope density + + # pump wavenumber as launched by the solver (snapped to the FFT grid) + dk = 2.0 * np.pi / (cfg_grid["nx"] * cfg_grid["dx"]) + k0 = np.sqrt(w0**2 - wp0**2) / c + k0 = np.round(k0 / dk) * dk + + # fixed point for the phase-matched EPW wavenumber + k = k0 + for _ in range(20): + w_ek = wp0 + 1.5 * k**2 * vte_sq / wp0 + w_s = w0 - w_ek + k1 = np.sqrt(w_s**2 - wp0**2) / c + k = k0 + k1 + + # local pump field including the density swelling factor + E0_local = E0_source * (1.0 - n) ** -0.25 + vos = e * E0_local / (me * w0) + + gamma0 = k * vos / 4.0 * wp0 / np.sqrt(w_ek * w_s) + return gamma0, k1 + + +@pytest.mark.parametrize("ny", ["2d", "1d"]) +def test_srs_growth_rate(ny): + """Noise-seeded homogeneous SRS: the EPW energy should grow at ~2*gamma0. + + The "1d" variant shrinks the y extent below one cell so the box collapses to ny=1, + which is the cheap true-1D configuration (the MATLAB srs_1D case). + """ + cfg = _load_cfg() + if ny == "1d": + cfg["grid"]["ymax"] = "0.02um" + cfg["grid"]["ymin"] = "-0.02um" + cfg["mlflow"]["run"] = "srs-test-1d" + exo, sol = _run(cfg) + + result = sol["solver result"] + t = np.array(result.ts["default"]) + e_sq = np.array(result.ys["default"]["e_sq"]) + + gamma0, _ = _predicted_gamma0(exo.adept_module.cfg["units"], exo.adept_module.cfg["grid"]) + + # fit the log-slope over the last stretch of the run, where exponential growth + # dominates the noise floor + fit_window = (t > t[-1] - 0.5) & (t < t[-1] - 0.02) + slope = np.polyfit(t[fit_window], np.log(e_sq[fit_window]), 1)[0] + measured_gamma = slope / 2.0 # e_sq ~ |E|^2 grows at 2*gamma + + assert e_sq[-1] > 1e4 * np.min(e_sq[t > 0.1]), "no SRS growth occurred" + np.testing.assert_allclose(measured_gamma, gamma0, rtol=0.35) + + +def test_srs_seed_propagation(): + """The Raman seed injector should launch a backward (-x) wave at the local k1.""" + cfg = _load_cfg() + + # seed-only setup: no pump, no noise -- E1 just propagates + del cfg["drivers"]["E0"] + cfg["drivers"]["E1"] = { + "intensity": "1.0e+12W/cm^2", + "delta_omega": 0.0, + "turn_on_time": "10fs", + } + cfg["grid"]["xmax"] = "20um" + cfg["grid"]["tmax"] = "0.2ps" + cfg["save"]["fields"]["t"]["tmax"] = "0.2ps" + cfg["save"]["fields"]["t"]["dt"] = "0.02ps" + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["mlflow"]["run"] = "srs-seed-test" + + exo, sol = _run(cfg) + + result = sol["solver result"] + e1_raw = np.array(result.ys["fields"]["E1"]) + e1 = e1_raw.view(np.complex64 if e1_raw.dtype == np.float32 else np.complex128) + e1y_final = e1[-1, :, 0, 1] # final time, y=first row, y-polarization + + assert np.max(np.abs(e1y_final)) > 0, "seed was not injected" + + dcfg = exo.adept_module.cfg + x = np.array(dcfg["grid"]["x"]) + derived = dcfg["units"]["derived"] + wp0, w1, c = derived["wp0"], derived["w1"], derived["c"] + n = dcfg["units"]["envelope density"] + k1 = np.sqrt(w1**2 - wp0**2) / c # seed injected at delta_omega = 0 + + # in the bulk (away from the injector near-field and the absorbers) the seed should be + # a leftward traveling wave exp(-i k1 x) + bulk = slice(np.argmin(np.abs(x - 6.0)), np.argmin(np.abs(x - 12.0))) + phase = np.unwrap(np.angle(e1y_final[bulk])) + k_measured = np.polyfit(x[bulk], phase, 1)[0] + assert k_measured < 0, f"seed propagates the wrong way: k = {k_measured:.2f} 1/um" + np.testing.assert_allclose(np.abs(k_measured), k1, rtol=0.05) + + # amplitude calibration of the injector: |E1| = E1_source * sinc(k1 dx) / eps1^(1/4) + dx = dcfg["grid"]["dx"] + eps1 = 1.0 - n * (derived["w0"] / w1) ** 2 + expected_amp = dcfg["drivers"]["E1"]["derived"]["amplitude"] * np.sin(k1 * dx) / (k1 * dx) / eps1**0.25 + np.testing.assert_allclose(np.mean(np.abs(e1y_final[bulk])), expected_amp, rtol=0.3) + + +if __name__ == "__main__": + test_srs_seed_propagation() + test_srs_growth_rate() From f962a042d84a64e750d81300fe7ff2c3a7f1ab52 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Thu, 23 Jul 2026 15:59:32 -0700 Subject: [PATCH 2/4] lpse2d: fix field save for quasi-1D (ny=1) runs The field-save interpolators assumed >=2 transverse cells: interpax.interp2d returns NaN off a single y-node, and RegularGridInterpolator fills 0 for a single-node y axis. This blanked all real-space/k-space field artifacts (fields.xr, k-fields.xr, plots//*) for ny=1 runs while the series.xr scalars stayed valid. Add an ny==1 path that interpolates in x only and keeps the single transverse row, for both the in-solve field saver and the background-density save. Co-Authored-By: Claude Fable 5 --- adept/_lpse2d/helpers.py | 51 ++++++++++++++++++++++++++++------------ 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index c22f7c32..36ddbc6f 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -714,13 +714,24 @@ def make_field_xarrays(cfg, this_t, state, td): from scipy import interpolate - density_interpolator = interpolate.RegularGridInterpolator( - (cfg["grid"]["x"], cfg["grid"]["y"]), cfg["grid"]["background_density"], bounds_error=False, fill_value=0.0 - ) - - grid_x, grid_y = np.meshgrid(xax, yax, indexing="ij") - points = np.array([grid_x.flatten(), grid_y.flatten()]).T - density_on_save_grid = density_interpolator(points).reshape((nx, ny)) + if ny == 1: + # Quasi-1D: RegularGridInterpolator cannot take a single-node y axis (and the + # save row lies off it -> fill_value=0). Interpolate the density in x only. + density_1d = np.asarray(cfg["grid"]["background_density"])[:, 0] + density_interpolator = interpolate.interp1d( + np.asarray(cfg["grid"]["x"]), density_1d, bounds_error=False, fill_value="extrapolate" + ) + density_on_save_grid = density_interpolator(np.asarray(xax)).reshape((nx, ny)) + else: + density_interpolator = interpolate.RegularGridInterpolator( + (cfg["grid"]["x"], cfg["grid"]["y"]), + cfg["grid"]["background_density"], + bounds_error=False, + fill_value=0.0, + ) + grid_x, grid_y = np.meshgrid(xax, yax, indexing="ij") + points = np.array([grid_x.flatten(), grid_y.flatten()]).T + density_on_save_grid = density_interpolator(points).reshape((nx, ny)) background_density = xr.DataArray( np.repeat(density_on_save_grid[None, ...], repeats=len(this_t), axis=0), @@ -786,14 +797,24 @@ def get_save_quantities(cfg: dict) -> dict: xq, yq = jnp.meshgrid(cfg["save"]["fields"]["x"]["ax"], cfg["save"]["fields"]["y"]["ax"], indexing="ij") - interpolator = partial( - interpax.interp2d, - xq=jnp.reshape(xq, (nx * ny), order="F"), - yq=jnp.reshape(yq, (nx * ny), order="F"), - x=cfg["grid"]["x"], - y=cfg["grid"]["y"], - method="linear", - ) + if ny == 1: + # Quasi-1D (single transverse cell): interpax.interp2d needs >=2 nodes per + # axis and returns NaN off the lone y-node, so interpolate in x only and + # keep the single y-row. save_func reshapes the flat (nx,) result to (nx, 1). + x_save = cfg["save"]["fields"]["x"]["ax"] + x_src = cfg["grid"]["x"] + + def interpolator(f): + return interpax.interp1d(x_save, x_src, jnp.reshape(f, (-1,)), method="linear") + else: + interpolator = partial( + interpax.interp2d, + xq=jnp.reshape(xq, (nx * ny), order="F"), + yq=jnp.reshape(yq, (nx * ny), order="F"), + x=cfg["grid"]["x"], + y=cfg["grid"]["y"], + method="linear", + ) def save_func(t, y, args): save_y = {} From 2516a66b9b7ee16d4c974e308cffaf17e972893f Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 28 Jul 2026 11:43:09 -0700 Subject: [PATCH 3/4] lpse2d: OSIRIS-comparable SRS diagnostics + reproducible noise seeding - New adept/_lpse2d/diagnostics.py: laser-budget window means (names and definitions match osiris_lpi/laser_budget.py), EPW growth fit copied verbatim from osiris_lpi/epw_growth.py (per-w0 rates), electron energy as cumulative EPW dissipation. - Default save now also logs epw_energy (OSIRIS units: fields in me*c*w0/e, lengths in c/w0), epw_dissipation (using the solver's own Landau + collisional rates via the new module-level landau_damping_rate), epw_boundary_loss, and discrete two-point flux probes for the laser budget (incident/transmitted/reflected/backrefl, normalized to I0). Probes sit at 2*boundary_width, clear of the absorber skirt; the legacy reflectivity probe at 1.6*bw is unchanged for back-compat. - post_process logs these scalars as MLflow metrics (previously only write/plot times) and adds laser-budget / EPW-fit / electron-energy plots. - terms.epw.source.{noise_amplitude,noise_seed} are config-driven; the resolved seed is pinned into the cfg pre-log_params so runs are exactly reproducible. Removed the dead density.noise draws that perturbed the global RNG stream. Co-Authored-By: Claude Fable 5 --- adept/_lpse2d/core/epw.py | 50 ++++--- adept/_lpse2d/diagnostics.py | 243 ++++++++++++++++++++++++++++++++++ adept/_lpse2d/helpers.py | 145 ++++++++++++++++++-- adept/_lpse2d/modules/base.py | 23 +--- 4 files changed, 416 insertions(+), 45 deletions(-) create mode 100644 adept/_lpse2d/diagnostics.py diff --git a/adept/_lpse2d/core/epw.py b/adept/_lpse2d/core/epw.py index 54651eee..1f34dbf3 100644 --- a/adept/_lpse2d/core/epw.py +++ b/adept/_lpse2d/core/epw.py @@ -6,6 +6,30 @@ from adept._lpse2d.core.driver import Driver +def landau_damping_rate(k_sq: Array, wp0: float, vte_sq: float, zero_mask: Array) -> Array: + """ + Landau damping rate for each k mode (amplitude rate, 1/ps). + + Matches MATLAB line 913: + gammaLandauEpw = sqrt(pi/8) * (1 + 3/2*k^2*vte^2/wp^2) * wp^4/(k^3*vte^3) * exp(...) + + Module-level so the solver (`SpectralEPWSolver`) and the dissipation + diagnostic (`helpers.get_default_save_func`) use the *same* rates and can + never drift apart. + """ + k_sq_safe = jnp.where(k_sq > 0, k_sq, 1.0) + + damping = ( + jnp.sqrt(np.pi / 8.0) + * (1.0 + 1.5 * k_sq * vte_sq / wp0**2) + * wp0**4 + / (k_sq_safe**1.5 * vte_sq**1.5) + * jnp.exp(-(1.5 + 0.5 * wp0**2 / (k_sq_safe * vte_sq))) + ) + + return damping * zero_mask + + class SpectralPotential: def __init__(self, cfg) -> None: self.background_density = cfg["grid"]["background_density"] @@ -362,10 +386,14 @@ def __init__(self, cfg: dict): is_outside_max_k1 = self.k_sq * (self.c / self.w1) ** 2 > max_k1_sq self.E1_filter = jnp.where(is_outside_max_k1, 0.0, 1.0)[..., None] - # Noise parameters + # Noise parameters. Amplitude default matches MATLAB noiseAmp + # (m201805_matlabLpse_v11.m:49). The seed is resolved (and written back into + # the cfg, so MLflow logs it) in helpers.get_derived_quantities; the fallback + # here only fires if that step was skipped. self.noise_enabled = cfg["terms"]["epw"]["source"]["noise"] - self.noise_amplitude = 1e-10 # matches MATLAB noiseAmp (m201805_matlabLpse_v11.m:49) - self.noise_seed = np.random.randint(2**20) + self.noise_amplitude = float(cfg["terms"]["epw"]["source"].get("noise_amplitude", 1e-10)) + cfg_seed = cfg["terms"]["epw"]["source"].get("noise_seed") + self.noise_seed = int(cfg_seed) if cfg_seed is not None else np.random.randint(2**20) # Density gradient self.density_gradient_enabled = cfg["terms"]["epw"]["density_gradient"] @@ -383,21 +411,7 @@ def calc_landau_damping_rate(self) -> Array: Returns: Landau damping rate array (shape: nx, ny) """ - # Avoid issues at k=0 - k_sq_safe = jnp.where(self.k_sq > 0, self.k_sq, 1.0) - - damping = ( - jnp.sqrt(np.pi / 8.0) - * (1.0 + 1.5 * self.k_sq * self.vte_sq / self.wp0**2) - * self.wp0**4 - / (k_sq_safe**1.5 * self.vte_sq**1.5) - * jnp.exp(-(1.5 + 0.5 * self.wp0**2 / (k_sq_safe * self.vte_sq))) - ) - - # Zero out k=0 mode - damping = damping * self.zero_mask - - return damping + return landau_damping_rate(self.k_sq, self.wp0, self.vte_sq, self.zero_mask) def phi_k_to_e_fields(self, phi_k: Array) -> tuple[Array, Array]: """ diff --git a/adept/_lpse2d/diagnostics.py b/adept/_lpse2d/diagnostics.py new file mode 100644 index 00000000..670ee641 --- /dev/null +++ b/adept/_lpse2d/diagnostics.py @@ -0,0 +1,243 @@ +"""Scalar diagnostics for lpse2d SRS runs, mirroring the OSIRIS scan2 metrics. + +The point of this module is one-to-one MLflow comparability with the OSIRIS PIC +post-processing in the ``osiris-lpi`` repo: metric *names and definitions* here +match ``osiris_lpi/laser_budget.py`` (R/T/absorbed window means, per-segment +scalars) and ``osiris_lpi/epw_growth.py`` (growth-rate fit). adept must not +import osiris-lpi, so the pure-numpy pieces are copied verbatim; if you change +the fit or windowing on either side, change it on both. + +Unit conventions: + +- Series arrive on a ``t (ps)`` axis. OSIRIS logs times and rates in ``1/w0`` + code units, so the growth fit is fed ``t * w0`` and ``epw_growth_rate`` is + per ``w0`` (with ``epw_growth_rate_per_ps`` logged alongside for humans). +- ``epw_energy`` and ``epw_dissipation`` are already converted to OSIRIS code + units by the save function (fields in ``me*c*w0/e``, lengths in ``c/w0``). +- Flux series (``incident_flux``, ``reflected_flux``, ``transmitted_flux``, + ``backrefl_flux``) are normalized to the nominal incident flux ``I0_code``. +""" + +from __future__ import annotations + +import numpy as np + +# --- EPW growth fit: copied verbatim from osiris_lpi/epw_growth.py (fit_epw_growth) --- + +ONSET_FACTOR = 2.0 # onset = first sustained W > ONSET_FACTOR*floor +START_FACTOR = 5.0 # fit window starts at first sustained W > START_FACTOR*floor +END_FRACTION = 0.1 # fit window ends when W first reaches END_FRACTION*max(W) +MEASURABLE_FACTOR = 10.0 # need max(W) >= MEASURABLE_FACTOR*floor to fit at all +MIN_FIT_POINTS = 6 +SUSTAIN = 5 # "sustained" crossing = 4 of the next 5 samples also above + + +def _sustained_crossing(W: np.ndarray, level: float) -> int | None: + """Index of the first crossing of ``level`` where >=4 of the next 5 samples stay above.""" + above = W > level + idx = np.where(above)[0] + for i in idx: + if above[i : i + SUSTAIN].sum() >= min(SUSTAIN - 1, len(W) - i): + return int(i) + return None + + +def fit_epw_growth(t: np.ndarray, W: np.ndarray) -> dict[str, float]: + """Automated exponential-rise fit of the EPW energy history; returns metric dict. + + Copied verbatim from ``osiris_lpi/epw_growth.py:fit_epw_growth`` so the window + selection is bit-identical to the scan2 backfill. See that module's docstring + for the algorithm; do not modify one copy without the other. + """ + t = np.asarray(t, dtype=float) + W = np.asarray(W, dtype=float) + good = (W > 0) & np.isfinite(W) + t, W = t[good], W[good] + out: dict[str, float] = {} + if W.size == 0: + out["epw_growth_measurable"] = 0.0 + return out + + # iterative noise floor: seed from the earliest ~50 dumps (pre-onset even for + # the fastest-growing runs), then widen the window to the detected onset so + # slow-onset runs get the robust full pre-onset median + i_seed = min(52, len(W)) + floor = float(np.median(W[2:i_seed])) if i_seed > 4 else float(np.median(W)) + for _ in range(5): + onset_i = _sustained_crossing(W, ONSET_FACTOR * floor) + if onset_i is None or onset_i <= 60: + break + new_floor = float(np.median(W[2:onset_i])) + if abs(new_floor - floor) <= 0.02 * floor: + floor = new_floor + break + floor = new_floor + + Wmax = float(W.max()) + out["epw_energy_floor"] = floor + out["epw_energy_max"] = Wmax + + measurable = Wmax >= MEASURABLE_FACTOR * floor + i0 = _sustained_crossing(W, START_FACTOR * floor) if measurable else None + i1 = -1 + if i0 is not None: + after = np.where(W[i0:] >= END_FRACTION * Wmax)[0] + i1 = i0 + int(after[0]) if after.size else len(W) - 1 + if i1 - i0 + 1 < MIN_FIT_POINTS: # near-vertical rise: widen the window + lo = _sustained_crossing(W, ONSET_FACTOR * floor) + if lo is not None: + after = np.where(W[lo:] >= 0.5 * Wmax)[0] + i0 = lo + i1 = lo + int(after[0]) if after.size else len(W) - 1 + if i0 is None or i1 - i0 + 1 < MIN_FIT_POINTS: + out["epw_growth_measurable"] = 0.0 + return out + + def _lin_fit(j0: int, j1: int) -> dict | None: + if j1 - j0 + 1 < MIN_FIT_POINTS: + return None + ts, lnW = t[j0 : j1 + 1], np.log(W[j0 : j1 + 1]) + p = np.polyfit(ts, lnW, 1) + resid = lnW - np.polyval(p, ts) + ss_tot = float(np.sum((lnW - lnW.mean()) ** 2)) + return dict( + epw_growth_rate=float(p[0] / 2.0), + epw_growth_rate_r2=1.0 - float(np.sum(resid**2)) / ss_tot if ss_tot > 0 else 0.0, + epw_growth_fit_tstart=float(ts[0]), + epw_growth_fit_tend=float(ts[-1]), + epw_growth_efolds=float((lnW[-1] - lnW[0]) / 2.0), + ) + + best = _lin_fit(i0, i1) + if best is not None and best["epw_growth_efolds"] < 1.0: + # narrow window (marginal growth): also try extending to 0.5*max(W) and + # keep whichever window the data fit better -- near-threshold bursty runs + # stay on the standard window with their honest low R^2 + after = np.where(W[i0:] >= 0.5 * Wmax)[0] + if after.size: + wide = _lin_fit(i0, i0 + int(after[0])) + if wide is not None and wide["epw_growth_rate_r2"] > best["epw_growth_rate_r2"]: + best = wide + if best is None: + out["epw_growth_measurable"] = 0.0 + return out + + out["epw_growth_measurable"] = 1.0 + out.update(best) + return out + + +# --- Time windows: same semantics as osiris_lpi/laser_budget.py:_segment_windows --- + + +def segment_windows(t: np.ndarray, last_frac: float = 0.25, n_segments: int = 4) -> list[dict]: + """Equal-time portions of the run whose final portion is the last-``last_frac`` window. + + Same semantics as ``osiris_lpi/laser_budget.py:_segment_windows``: the final + segment coincides with the headline (last-25%) scalars; the leading span is + split into ``n_segments - 1`` equal portions. With the defaults this is four + equal quarters. + """ + t = np.asarray(t, dtype=float) + n = max(int(n_segments), 1) + span = float(t[-1] - t[0]) if t.size else 0.0 + t_lo = t[0] + (1.0 - last_frac) * span if span > 0 else float(t[0]) + edges = list(np.linspace(float(t[0]), t_lo, n)) + [float(t[-1])] + out: list[dict] = [] + for i in range(n): + lo, hi = edges[i], edges[i + 1] + is_last = i == n - 1 + mask = (t >= lo) if is_last else ((t >= lo) & (t < hi)) + out.append({"index": i + 1, "n": n, "t_lo": float(lo), "t_hi": float(hi), "is_last": is_last, "mask": mask}) + return out + + +def _wmean(v: np.ndarray, mask: np.ndarray) -> float | None: + mask = np.asarray(mask, dtype=bool) + if not mask.any(): + return None + return float(np.mean(np.asarray(v)[mask])) + + +# --- Metric assembly from the default-save series --- + + +def series_metrics(series, cfg: dict) -> dict[str, float]: + """Build the OSIRIS-comparable scalar metrics from the default-save time series. + + ``series`` is the xarray Dataset built by ``make_series_xarrays`` (keys are the + save-function outputs, coord ``t (ps)``). Returns a flat metrics dict for MLflow. + """ + metrics: dict[str, float] = {} + if "epw_energy" not in series: + return metrics + + derived = cfg["units"]["derived"] + w0 = derived["w0"] # 1/ps + t_ps = np.asarray(series["t (ps)"].values, dtype=float) + t_w0 = t_ps * w0 + + # EPW growth fit on the OSIRIS-normalized energy vs t in 1/w0, so + # epw_growth_rate / epw_energy_floor / epw_energy_max are directly comparable + # to the scan2 rows. + W = np.asarray(series["epw_energy"].values, dtype=float) + fit = fit_epw_growth(t_w0, W) + metrics.update(fit) + if "epw_growth_rate" in fit: + metrics["epw_growth_rate_per_ps"] = fit["epw_growth_rate"] * w0 + + # Electron energy: cumulative EPW dissipation (Landau + collisional), i.e. the + # energy handed to electrons. epw_dissipation is OSIRIS energy per ps, so the + # trapezoid over t (ps) gives OSIRIS energy; the incident fluence is + # I0_osiris * t_w0 with I0_osiris = a0^2/2 (flux in OSIRIS units). + a0 = derived["E0_source"] * derived["e_norm"] + I0_osiris = 0.5 * a0**2 + if "epw_dissipation" in series: + dissip_ps = np.asarray(series["epw_dissipation"].values, dtype=float) + electron_energy = np.concatenate([[0.0], np.cumsum(0.5 * (dissip_ps[1:] + dissip_ps[:-1]) * np.diff(t_ps))]) + metrics["electron_energy_final"] = float(electron_energy[-1]) + if t_w0[-1] > 0: + metrics["electron_energy_frac_final"] = float(electron_energy[-1] / (I0_osiris * t_w0[-1])) + + # Laser budget. Flux series are already normalized to the nominal incident flux + # I0_code. With the pump prescribed (no depletion) only the one-way "naive" + # estimators are meaningful; with pump depletion on, the net-Poynting R/T/absorbed + # match osiris_lpi/laser_budget.py exactly (R = 1 - S_left/I0, T = S_right/I0, + # absorbed = (S_left - S_right)/I0, so R + T + absorbed == 1 identically). + pump_evolved = bool(cfg["terms"].get("light", {}).get("pump_depletion", False)) + have_flux = all(k in series for k in ("incident_flux", "reflected_flux", "transmitted_flux", "backrefl_flux")) + if have_flux: + inc = np.asarray(series["incident_flux"].values, dtype=float) + refl = np.asarray(series["reflected_flux"].values, dtype=float) + trans = np.asarray(series["transmitted_flux"].values, dtype=float) + backrefl = np.asarray(series["backrefl_flux"].values, dtype=float) + s_left = inc - refl # F0_left + F1_left, in units of I0_code + s_right = trans + backrefl + + windows = segment_windows(t_ps, last_frac=0.25, n_segments=4) + for win in windows: + suffix = "" if win["is_last"] else f"_seg{win['index']}of{win['n']}" + r_naive = _wmean(refl, win["mask"]) + t_naive = _wmean(trans, win["mask"]) + if r_naive is None: + continue + metrics[f"laser_reflectivity_naive{suffix}"] = r_naive + metrics[f"laser_transmissivity_naive{suffix}"] = t_naive + if pump_evolved: + metrics[f"laser_reflectivity{suffix}"] = 1.0 - _wmean(s_left, win["mask"]) + metrics[f"laser_transmissivity{suffix}"] = _wmean(s_right, win["mask"]) + metrics[f"laser_absorbed_frac{suffix}"] = _wmean(s_left - s_right, win["mask"]) + else: + # pump is prescribed: the one-way Raman flux is the only measurable + # reflectivity, and there is no transmission/absorption budget + metrics[f"laser_reflectivity{suffix}"] = r_naive + if "epw_dissipation" in series: + d = _wmean(np.asarray(series["epw_dissipation"].values, dtype=float), win["mask"]) + if d is not None and I0_osiris > 0: + metrics[f"laser_absorbed_frac_epw{suffix}"] = d / w0 / I0_osiris + + # injector/prescription health: the measured incident flux should be ~1 + # (skip the first 5% of the run to avoid the fill-in transient) + metrics["laser_incident_flux_ratio"] = float(np.mean(inc[max(1, len(inc) // 20) :])) + + return metrics diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index 36ddbc6f..a0cbacd3 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -165,6 +165,14 @@ def write_units(cfg: dict) -> dict: "lambda_D": ld, "nu_coll": nu_coll, "E0_source": E0_source, + # Conversions to OSIRIS code units, for one-to-one metric comparison with PIC runs: + # multiply an E-field (code units) by e_norm to get it in me*c*w0/e units, and a + # length (um) by x_norm to get it in c/w0. I0_code is the nominal incident pump + # flux in this code's units (c * |E_envelope|^2); the WKB-swelled pump satisfies + # sqrt(eps) * |E0_local|^2 = E0_source^2 so this holds at any density below nc. + "e_norm": e / (me * c * w0), + "x_norm": w0 / c, + "I0_code": c * E0_source**2, "timeScale": timeScale, "spatialScale": spatialScale, "velocityScale": velocityScale, @@ -258,6 +266,15 @@ def get_derived_quantities(cfg: dict) -> dict: # SRS: the Raman light is advanced with an explicit conditionally-stable scheme, so it is # sub-cycled within each EPW step. The stability bound follows MATLAB line 500 # (dt_max_seed), generalized to 2D + # EPW noise source: resolve the amplitude/seed here (prior to log_params) so the + # run is reproducible and the actual seed lands in MLflow. noise_seed: null (or + # absent) draws a random seed once, then pins it in the cfg. + epw_source = cfg["terms"]["epw"]["source"] + if epw_source.get("noise", False): + epw_source.setdefault("noise_amplitude", 1e-10) + if epw_source.get("noise_seed") is None: + epw_source["noise_seed"] = int(np.random.randint(2**20)) + if cfg["terms"]["epw"]["source"].get("srs", False): derived = cfg["units"]["derived"] wpe_max_sq = derived["w0"] ** 2 * max(cfg["density"].get("max", 1.0), cfg["density"].get("min", 0.0)) @@ -622,6 +639,8 @@ def plot_kt(kfields, td): def post_process(result, cfg: dict, td: str) -> tuple[xr.Dataset, xr.Dataset]: + from adept._lpse2d.diagnostics import series_metrics + os.makedirs(os.path.join(td, "binary")) metrics = {} t0 = time.time() @@ -630,8 +649,12 @@ def post_process(result, cfg: dict, td: str) -> tuple[xr.Dataset, xr.Dataset]: metrics["write_time"] = time.time() - t0 os.makedirs(os.path.join(td, "plots")) + # OSIRIS-comparable scalars (laser budget, EPW growth fit, electron energy) + metrics.update(series_metrics(series, cfg)) + t0 = time.time() plot_series(series, td) + plot_srs_diagnostics(series, metrics, cfg, td) plot_fields(fields, td) plot_kt(kfields, td) metrics["plot_time"] = time.time() - t0 @@ -639,6 +662,53 @@ def post_process(result, cfg: dict, td: str) -> tuple[xr.Dataset, xr.Dataset]: return {"k": kfields, "x": fields, "series": series, "metrics": metrics} +def plot_srs_diagnostics(series, metrics, cfg, td): + """Composite diagnostic plots mirroring the OSIRIS scan2 figures: the laser + budget channels vs time, the EPW energy with the growth-fit window shaded, and + the cumulative electron energy (integrated EPW dissipation).""" + t = np.asarray(series["t (ps)"].values, dtype=float) + + if "incident_flux" in series: + fig, ax = plt.subplots(1, 2, figsize=(9, 3.5)) + for a in ax: + for key, label in [ + ("incident_flux", "incident"), + ("reflected_flux", "reflected"), + ("transmitted_flux", "transmitted"), + ("backrefl_flux", "back-reflected"), + ]: + a.plot(t, np.asarray(series[key].values, dtype=float), label=label) + a.set_xlabel("t (ps)") + a.set_ylabel("flux / I0") + ax[1].set_yscale("log") + ax[0].legend(fontsize=8) + fig.savefig(os.path.join(td, "plots", "laser_budget_vs_t.png"), bbox_inches="tight") + plt.close(fig) + + if "epw_energy" in series: + w0 = cfg["units"]["derived"]["w0"] + fig, ax = plt.subplots(1, 1, figsize=(5, 3.5)) + ax.semilogy(t, np.asarray(series["epw_energy"].values, dtype=float)) + if metrics.get("epw_growth_measurable"): + # fit window is stored in 1/w0 (OSIRIS code time); convert back to ps + ax.axvspan(metrics["epw_growth_fit_tstart"] / w0, metrics["epw_growth_fit_tend"] / w0, alpha=0.2) + ax.set_title(f"gamma/w0 = {metrics['epw_growth_rate']:.2e}, r2 = {metrics['epw_growth_rate_r2']:.3f}") + ax.set_xlabel("t (ps)") + ax.set_ylabel("W_epw (OSIRIS units)") + fig.savefig(os.path.join(td, "plots", "epw_energy_fit.png"), bbox_inches="tight") + plt.close(fig) + + if "epw_dissipation" in series: + dissip = np.asarray(series["epw_dissipation"].values, dtype=float) + electron_energy = np.concatenate([[0.0], np.cumsum(0.5 * (dissip[1:] + dissip[:-1]) * np.diff(t))]) + fig, ax = plt.subplots(1, 1, figsize=(5, 3.5)) + ax.plot(t, electron_energy) + ax.set_xlabel("t (ps)") + ax.set_ylabel("cumulative electron energy (OSIRIS units)") + fig.savefig(os.path.join(td, "plots", "electron_energy_vs_t.png"), bbox_inches="tight") + plt.close(fig) + + def plot_series(series, td): for k in series.keys(): fig, ax = plt.subplots(1, 2, figsize=(8, 3)) @@ -848,34 +918,91 @@ def save_func(t, y, args): def get_default_save_func(cfg): + from adept._lpse2d.core.epw import landau_damping_rate + srs_on = cfg["terms"]["epw"]["source"].get("srs", False) + derived = cfg["units"]["derived"] + dt = cfg["grid"]["dt"] + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + kx, ky = cfg["grid"]["kx"], cfg["grid"]["ky"] + + # OSIRIS-normalized EPW energy: W = 1/2 * dx * sum_x _cycle with fields in + # me*c*w0/e and lengths in c/w0 (osiris_lpi/epw_growth.py convention). The complex + # envelope carries _cycle = |E|^2/2, hence the 0.25. The transverse mean makes + # the 2D case a per-unit-y version of the same quantity; for ny=1 it is identical + # to the OSIRIS 1D reduction. + epw_energy_prefactor = 0.25 * cfg["grid"]["dx"] * derived["x_norm"] * derived["e_norm"] ** 2 + + # Per-step dissipated EPW energy, evaluated with the *same* rates the solver + # applies (phi_k *= exp(-(gamma_landau + nu_coll)*dt) => energy factor exp(-2(..)dt)). + # Parseval: sum_x <.>_y |E|^2 = (1/(nx*ny^2)) * sum_k k^2 |phi_k|^2. + k_sq = np.array(kx[:, None] ** 2 + ky[None, :] ** 2) + zero_mask = np.where(k_sq > 0, 1.0, 0.0) + gamma_total = np.array(landau_damping_rate(jnp.array(k_sq), derived["wp0"], derived["vte_sq"], jnp.array(zero_mask))) + gamma_total = gamma_total + derived.get("nu_coll", 0.0) * zero_mask + energy_loss_factor = (1.0 - np.exp(-2.0 * gamma_total * dt)) / dt # 1/ps, per k mode + boundary_sq_loss = (1.0 - np.array(cfg["grid"]["absorbing_boundaries"]) ** 2) / dt # 1/ps, per cell + if srs_on: - # Reflectivity probe: sample |E1_y|^2 on the low-density side, just inside the absorber. - # R = sqrt(eps1) * <|E1_y|^2>_y / E0_source^2 where eps1 = 1 - wpe^2/w1^2 accounts for - # the reduced group velocity (Poynting flux ~ sqrt(eps) |E|^2) relative to the vacuum pump. + # Legacy reflectivity probe: sample |E1_y|^2 on the low-density side, just inside + # the absorber. R = sqrt(eps1) * <|E1_y|^2>_y / E0_source^2. Kept unchanged for + # backward comparability; it sits in the absorber's tanh skirt and reads ~10% low. boundary_width = _Q(cfg["grid"]["boundary_width"]).to("um").value - x_probe = cfg["grid"]["xmin"] + 1.6 * boundary_width # clear of the absorber's tanh skirt + x_probe = cfg["grid"]["xmin"] + 1.6 * boundary_width ix_probe = int(np.argmin(np.abs(np.array(cfg["grid"]["x"]) - x_probe))) - w0 = cfg["units"]["derived"]["w0"] - w1 = cfg["units"]["derived"]["w1"] + w0 = derived["w0"] + w1 = derived["w1"] n_probe = float(np.mean(np.array(cfg["grid"]["background_density"])[ix_probe, :])) sqrt_eps1 = np.sqrt(max(1.0 - n_probe * w0**2 / w1**2, 0.0)) - E0_source_sq = cfg["units"]["derived"]["E0_source"] ** 2 + E0_source_sq = derived["E0_source"] ** 2 + + # Flux probes for the OSIRIS-style laser budget. F_j = c^2/(w dx) * Im(E_j* E_j+1) + # is the exactly-conserved flux of the FD Schroedinger operator (equals + # c*sqrt(eps)|E|^2*sinc(k dx) for a plane wave, i.e. the grid's own dispersion is + # accounted for). Probes sit at probe_offset (default 2*boundary_width, clear of + # the absorber skirt whose transmission at 1.6*bw is only ~0.91). + if "probe_offset" in cfg["grid"]: + probe_offset = _Q(cfg["grid"]["probe_offset"]).to("um").value + else: + probe_offset = 2.0 * boundary_width + x_grid = np.array(cfg["grid"]["x"]) + ix_left = int(np.argmin(np.abs(x_grid - (cfg["grid"]["xmin"] + probe_offset)))) + ix_right = int(np.argmin(np.abs(x_grid - (cfg["grid"]["xmax"] - probe_offset)))) + flux_coeff_w0 = derived["c"] ** 2 / (derived["w0"] * cfg["grid"]["dx"]) + flux_coeff_w1 = derived["c"] ** 2 / (derived["w1"] * cfg["grid"]["dx"]) + I0_code = derived["I0_code"] + + def discrete_flux(E, ix, coeff): + # sum over polarization components, mean over y + cross = jnp.sum(jnp.conj(E[ix, :, :]) * E[ix + 1, :, :], axis=-1) + return coeff * jnp.mean(jnp.imag(cross)) def save_func(t, y, args): phi_k = y["epw"].view(jnp.complex128) - ex = -1j * cfg["grid"]["kx"][:, None] * phi_k - ey = -1j * cfg["grid"]["ky"][None, :] * phi_k + ex = -1j * kx[:, None] * phi_k + ey = -1j * ky[None, :] * phi_k ex = jnp.fft.ifft2(ex) ey = jnp.fft.ifft2(ey) e_sq = jnp.abs(ex) ** 2 + jnp.abs(ey) ** 2 out = {"e_sq": jnp.sum(e_sq * cfg["grid"]["dx"] * cfg["grid"]["dy"]), "max_phi": jnp.max(jnp.abs(phi_k))} + out["epw_energy"] = epw_energy_prefactor * jnp.sum(jnp.mean(e_sq, axis=1)) + out["epw_dissipation"] = ( + epw_energy_prefactor / (nx * ny**2) * jnp.sum(k_sq * jnp.abs(phi_k) ** 2 * energy_loss_factor) + ) + out["epw_boundary_loss"] = epw_energy_prefactor * jnp.sum(jnp.mean(e_sq * boundary_sq_loss, axis=1)) + if srs_on: e1 = y["E1"].view(jnp.complex128) + e0 = y["E0"].view(jnp.complex128) out["e1_sq"] = jnp.mean(jnp.sum(jnp.abs(e1) ** 2, axis=-1)) out["reflectivity"] = sqrt_eps1 * jnp.mean(jnp.abs(e1[ix_probe, :, 1]) ** 2) / E0_source_sq + # laser budget channels, all normalized to the nominal incident flux + out["incident_flux"] = discrete_flux(e0, ix_left, flux_coeff_w0) / I0_code + out["transmitted_flux"] = discrete_flux(e0, ix_right, flux_coeff_w0) / I0_code + out["reflected_flux"] = -discrete_flux(e1, ix_left, flux_coeff_w1) / I0_code + out["backrefl_flux"] = discrete_flux(e1, ix_right, flux_coeff_w1) / I0_code return out diff --git a/adept/_lpse2d/modules/base.py b/adept/_lpse2d/modules/base.py index b68d3834..d9bc57e0 100644 --- a/adept/_lpse2d/modules/base.py +++ b/adept/_lpse2d/modules/base.py @@ -75,24 +75,11 @@ def init_diffeqsolve(self): ) def init_state_and_args(self) -> dict: - if self.cfg["density"]["noise"]["type"] == "uniform": - random_amps = np.random.uniform( - self.cfg["density"]["noise"]["min"], - self.cfg["density"]["noise"]["max"], - (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"]), - ) - - elif self.cfg["density"]["noise"]["type"] == "normal": - loc = 0.5 * (self.cfg["density"]["noise"]["min"] + self.cfg["density"]["noise"]["max"]) - scale = 1.0 - random_amps = np.random.normal(loc, scale, (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"])) - - else: - raise NotImplementedError - - random_phases = np.random.uniform(0, 2 * np.pi, (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"])) - phi_noise = 1 * np.exp(1j * random_phases) - epw = 0 * phi_noise + # The initial EPW is identically zero; noise-seeded runs get their seeding from + # terms.epw.source.noise (a per-step source in SpectralEPWSolver). The old + # density.noise draws were dead code, but they consumed the global numpy RNG + # stream and made "identical" runs differ -- so they are gone. + epw = np.zeros((self.cfg["grid"]["nx"], self.cfg["grid"]["ny"]), dtype=np.complex128) self.cfg["grid"]["background_density"] = get_density_profile(self.cfg) E0 = np.zeros((self.cfg["grid"]["nx"], self.cfg["grid"]["ny"], 2), dtype=np.complex128) From d3f8321ba7675d9ee108b28836421befd90ec74f Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 28 Jul 2026 12:15:09 -0700 Subject: [PATCH 4/4] lpse2d: pump depletion (evolved E0) for SRS terms.light.pump_depletion: true evolves the pump with the same staggered explicit FD envelope scheme as the Raman light (new core/light.py, CoupledLight), ported from the isPumpDepletion path of lpse-matlab m201805_matlabLpse_v11.m: - pump RHS: diffraction + local detuning + depletion coupling -i e/(4 w1 me) (laplacian phi) E1 (conjugate-free, partner-frequency denominator; Manley-Rowe-consistent with the E1 and EPW couplings) - two-point boundary injector at xmin + drivers.E0.offset (default 2*bw), multi-color, MATLAB amplitude calibration - both waves advance inside one staggered real/imag update (advancing them independently would break the discrete conservation) - substep limit = min over both carriers; E0 high-k filter in the EPW SRS source on the dynamic-pump path (MATLAB skips it on the static path) - budget flux probes convert the exact discrete two-point flux to physical flux via the FD group-velocity factor sin(k_grid dx)/(k dx); metrics normalize to the measured incident flux (the injector launches sin(k0 dx)/sin(k_grid dx) ~ 0.98 of nominal amplitude at 8 cells/lambda) - default off; the prescribed-pump path is untouched (verified: the three pre-existing SRS tests pass unchanged) Tests: pump injector flux+amplitude calibration against the discrete- dispersion prediction; seeded Raman-amplifier energy-budget closure (S_left - S_right vs 2x the field-only EPW energy rates -- the kinetic sloshing half doubles the electron heating; closes to ~4%, asserted <10%); R+T+absorbed==1; depletion actually reduces transmission vs the prescribed pump; epw_energy normalization; noise-seed reproducibility. Docs: overview pump-depletion note replaced, SRS diagnostics table added, config.md new keys (terms.light, probe_offset, noise_amplitude/seed, drivers.E0.offset/turn_on_time); datamodel updated to match. Co-Authored-By: Claude Fable 5 --- adept/_lpse2d/core/epw.py | 10 ++ adept/_lpse2d/core/light.py | 230 +++++++++++++++++++++++++ adept/_lpse2d/core/vector_field.py | 20 ++- adept/_lpse2d/datamodel.py | 10 ++ adept/_lpse2d/diagnostics.py | 24 ++- adept/_lpse2d/helpers.py | 79 ++++++++- docs/source/solvers/lpse2d/config.md | 24 +-- docs/source/solvers/lpse2d/overview.md | 19 +- tests/test_lpse2d/test_srs.py | 192 +++++++++++++++++++++ 9 files changed, 581 insertions(+), 27 deletions(-) create mode 100644 adept/_lpse2d/core/light.py diff --git a/adept/_lpse2d/core/epw.py b/adept/_lpse2d/core/epw.py index 1f34dbf3..72f8eb84 100644 --- a/adept/_lpse2d/core/epw.py +++ b/adept/_lpse2d/core/epw.py @@ -385,6 +385,14 @@ def __init__(self, cfg: dict): max_k1_sq = max_source_k_multiplier**2 * max(1.0 - n_min * self.w0**2 / self.w1**2, 0.0) is_outside_max_k1 = self.k_sq * (self.c / self.w1) ** 2 > max_k1_sq self.E1_filter = jnp.where(is_outside_max_k1, 0.0, 1.0)[..., None] + # when the pump is evolved (terms.light.pump_depletion) it is filtered too, + # exactly as MATLAB's evaluate_E0_dot_E1 (lines 2302-2354) does on the + # dynamic-laser path and skips on the static path (line 2307-2308) + self.pump_depletion = cfg["terms"].get("light", {}).get("pump_depletion", False) + if self.pump_depletion: + max_k0_sq = max_source_k_multiplier**2 * max(1.0 - n_min, 0.0) + is_outside_max_k0 = self.k_sq * (self.c / self.w0) ** 2 > max_k0_sq + self.E0_filter = jnp.where(is_outside_max_k0, 0.0, 1.0)[..., None] # Noise parameters. Amplitude default matches MATLAB noiseAmp # (m201805_matlabLpse_v11.m:49). The seed is resolved (and written back into @@ -554,6 +562,8 @@ def calc_srs_source(self, E0: Array, E1: Array) -> Array: SRS source term in k-space """ E1_filtered = jnp.fft.ifft2(jnp.fft.fft2(E1, axes=(0, 1)) * self.E1_filter, axes=(0, 1)) + if self.pump_depletion: + E0 = jnp.fft.ifft2(jnp.fft.fft2(E0, axes=(0, 1)) * self.E0_filter, axes=(0, 1)) E0_dot_E1 = E0[..., 0] * jnp.conj(E1_filtered[..., 0]) + E0[..., 1] * jnp.conj(E1_filtered[..., 1]) # (1 + backgroundDensityPerturbation) = n / n_envelope diff --git a/adept/_lpse2d/core/light.py b/adept/_lpse2d/core/light.py new file mode 100644 index 00000000..034a88e8 --- /dev/null +++ b/adept/_lpse2d/core/light.py @@ -0,0 +1,230 @@ +import numpy as np +from jax import Array, lax +from jax import numpy as jnp + +from adept._base_ import get_envelope + + +class CoupledLight: + """ + Evolves the pump E0 and the Raman scattered light E1 together, with pump depletion. + + This is the `isPumpDepletion` path of m201805_matlabLpse_v11.m: the pump is no longer + prescribed analytically but advanced with the same staggered explicit scheme as the + Raman light (lightSplitStep, lines 1377-1424), sourced by a boundary injector + (lines 1707-1753) and coupled to the EPW through the conjugate-free SRS term + (lines 1611-1648): + + dE0/dt = i c^2/(2 w0) * (Laplacian terms) E0 + + i w0/2 * (1 - wp0^2/w0^2 * n/n_env) * E0 + - i e/(4 w1 me) * (laplacian phi) * E1 (pump depletion) + + boundary source injector + + dE1/dt = i c^2/(2 w1) * (Laplacian terms) E1 + + i w1/2 * (1 - wp0^2/w1^2 * n/n_env) * E1 + - i e/(4 w0 me) * conj(laplacian phi) * E0 (SRS coupling) + [+ seed injector] + + Note the coupling denominators: each wave's SRS term carries the *partner* wave's + frequency, and only E1's term conjugates the potential. Together with the EPW source + (epw.py, prefactor e*wp0/(4 me w0 w1)) these satisfy Manley-Rowe exactly: + d/dt Int(|E0|^2 + |E1|^2 + |grad phi|^2) = 0 for the coupling terms alone. + + Both fields must be advanced inside the *same* staggered update (both real parts with + the RHS at t, then both imaginary parts with the RHS at t + dt/2) -- advancing them + with two independent RamanLight-style calls would break the discrete conservation. + + The pump injector amplitude is divided by sinc(k0 dx) so the launched amplitude is + exactly E0_source * sqrt(intensity) / eps^(1/4) despite the two-point discrete + source's sinc response (the E1 seed injector intentionally keeps the MATLAB + calibration; see tests/test_lpse2d/test_srs.py::test_srs_seed_propagation). + """ + + def __init__(self, cfg: dict): + self.cfg = cfg + derived = cfg["units"]["derived"] + self.c = derived["c"] + self.w0 = derived["w0"] + self.w1 = derived["w1"] + self.wp0 = derived["wp0"] + self.e = derived["e"] + self.me = derived["me"] + self.E0_source = derived["E0_source"] + self.envelope_density = cfg["units"]["envelope density"] + + self.dx = cfg["grid"]["dx"] + self.dy = cfg["grid"]["dy"] + self.dt = cfg["grid"]["dt"] # outer (EPW) step + self.n_sub = cfg["grid"]["light_substeps"] + self.dt_l = self.dt / self.n_sub + self.x = cfg["grid"]["x"] + self.y = cfg["grid"]["y"] + self.k_sq = cfg["grid"]["kx"][:, None] ** 2 + cfg["grid"]["ky"][None, :] ** 2 + + background_density = cfg["grid"]["background_density"] + # local detuning of each envelope (MATLAB lines 1616-1626 / 1661-1671); + # with wp0^2 = w0^2 * n_env the pump coefficient reduces to i w0/2 (1 - n) + self.linear_coeff0 = ( + 1j * self.w0 / 2.0 * (1.0 - self.wp0**2 / self.w0**2 * background_density / self.envelope_density) + ) + self.linear_coeff1 = ( + 1j * self.w1 / 2.0 * (1.0 - self.wp0**2 / self.w1**2 * background_density / self.envelope_density) + ) + self.diffraction_coeff0 = 1j * self.c**2 / (2.0 * self.w0) + self.diffraction_coeff1 = 1j * self.c**2 / (2.0 * self.w1) + self.srs_coeff1 = -1j * self.e / (4.0 * self.w0 * self.me) # in dE1/dt, conj(lap phi) * E0 + self.depletion_coeff0 = -1j * self.e / (4.0 * self.w1 * self.me) # in dE0/dt, lap phi * E1 + + self.sub_boundary = cfg["grid"]["absorbing_boundaries"] ** (1.0 / self.n_sub) + + # ---- pump injector (MATLAB lines 1707-1753, mirrored to the left edge) ---- + pump = cfg["drivers"]["E0"]["derived"] + x_inject = cfg["grid"]["xmin"] + pump["offset"] + self.i0 = int(np.argmin(np.abs(np.array(self.x) - x_inject))) + n_src = float(background_density[self.i0, 0]) + permittivity0 = 1.0 - n_src + if permittivity0 <= 0: + raise ValueError( + f"The pump injector at x = {float(self.x[self.i0]):.2f} um sits at density " + f"{n_src:.3f} nc, at or above critical. Lower density.max or move drivers.E0.offset." + ) + self.n_src = n_src + self.pump_turn_on_time = pump["turn_on_time"] + self.source_prefactor0 = self.c**2 / (2.0 * self.w0) / permittivity0**0.25 / self.dx**2 + + # E1 seed injector, identical to RamanLight's (MATLAB lines 1757-1769) + if "E1" in cfg["drivers"]: + seed = cfg["drivers"]["E1"]["derived"] + x_seed = cfg["grid"]["xmax"] - seed["offset"] + self.i1 = int(np.argmin(np.abs(np.array(self.x) - x_seed))) + wpe_i1 = self.w0 * np.sqrt(background_density[self.i1, 0]) + self.wpe_sq_i1 = float(wpe_i1**2) + permittivity1 = 1.0 - self.wpe_sq_i1 / self.w1**2 + if permittivity1 <= 0: + raise ValueError( + f"The Raman seed injector at x = {float(self.x[self.i1]):.2f} um sits at density " + f"{float(background_density[self.i1, 0]):.3f} nc, above the w1 critical density " + f"{(self.w1 / self.w0) ** 2:.3f} nc where the seed is evanescent." + ) + self.source_prefactor1 = self.c**2 / (2.0 * self.w1) / permittivity1**0.25 / self.dx**2 + self.seed_enabled = True + else: + self.seed_enabled = False + + def _d2x(self, f: Array) -> Array: + return (jnp.roll(f, -1, axis=0) - 2.0 * f + jnp.roll(f, 1, axis=0)) / self.dx**2 + + def _d2y(self, f: Array) -> Array: + return (jnp.roll(f, -1, axis=1) - 2.0 * f + jnp.roll(f, 1, axis=1)) / self.dy**2 + + def _dxdy(self, f: Array) -> Array: + return ( + jnp.roll(f, (-1, -1), axis=(0, 1)) + - jnp.roll(f, (1, -1), axis=(0, 1)) + - jnp.roll(f, (-1, 1), axis=(0, 1)) + + jnp.roll(f, (1, 1), axis=(0, 1)) + ) / (4.0 * self.dx * self.dy) + + def calc_pump_source(self, t: float, pump_args: dict) -> tuple[Array, Array]: + """ + Two-point pump injector rows, summed over colors (MATLAB lines 1738-1750, + with +k0 and the left edge instead of -k1 and the right edge). + + Returns the rows added to the E0_y RHS at self.i0 and self.i0 + 1. + """ + t_env = get_envelope( + pump_args["tr"], + pump_args["tr"], + pump_args["tc"] - pump_args["tw"] / 2, + pump_args["tc"] + pump_args["tw"] / 2, + t, + ) + turn_on = 1.0 - jnp.exp(-((t / self.pump_turn_on_time) ** 2)) + + delta_omega = pump_args["delta_omega"] # (nc,) + intensities = pump_args["intensities"] # (nc, ny), fractions summing to 1 + phases = pump_args["phases"] # (nc, ny) + + # local pump wavenumber per color (MATLAB kSource0). The two-point source + # launches amplitude E_src * sin(k0 dx)/sin(k_grid dx) / eps^(1/4) -- a ~2% + # deficit at 8 cells/wavelength from the grid dispersion; the budget metrics + # normalize to the *measured* incident flux, so this bias cancels there. + k0 = self.w0 / self.c * jnp.sqrt((1.0 + delta_omega) ** 2 - self.n_src) # (nc,) + + amp = self.source_prefactor0 * self.E0_source * jnp.sqrt(intensities) * t_env * turn_on # (nc, ny) + + color_phase = jnp.exp(-1j * self.w0 * delta_omega[:, None] * t + 1j * phases) # (nc, ny) + row_i0p1 = jnp.sum(-1j * amp * jnp.exp(1j * k0[:, None] * self.x[self.i0]) * color_phase, axis=0) + row_i0 = jnp.sum(1j * amp * jnp.exp(1j * k0[:, None] * self.x[self.i0 + 1]) * color_phase, axis=0) + return row_i0, row_i0p1 + + def calc_seed_source(self, t: float, seed_args: dict) -> tuple[Array, Array]: + """Two-point Raman seed rows, identical to RamanLight.calc_seed_source.""" + dw1 = seed_args["delta_omega"] + turn_on = 1.0 - jnp.exp(-((t / seed_args["turn_on_time"]) ** 2)) + amp = self.source_prefactor1 * seed_args["amplitude"] * turn_on + + if seed_args["yw"] > 0: + envelope_y = jnp.exp(-((self.y / (seed_args["yw"] / 2.0)) ** 4)) + else: + envelope_y = jnp.ones_like(self.y) + + k1 = self.w1 / self.c * jnp.sqrt((1.0 + dw1) ** 2 - self.wpe_sq_i1 / self.w1**2) + + row_i1 = -1j * amp * envelope_y * jnp.exp(-1j * k1 * self.x[self.i1 + 1] - 1j * self.w1 * dw1 * t) + row_i1p1 = 1j * amp * envelope_y * jnp.exp(-1j * k1 * self.x[self.i1] - 1j * self.w1 * dw1 * t) + return row_i1, row_i1p1 + + def rhs( + self, t: float, E0: Array, E1: Array, laplacian_phi: Array, pump_args: dict, seed_args: dict | None + ) -> tuple[Array, Array]: + e0x, e0y = E0[..., 0], E0[..., 1] + e1x, e1y = E1[..., 0], E1[..., 1] + + # pump: propagation + detuning (MATLAB lines 1616-1626) + k_e0x = self.diffraction_coeff0 * (self._d2y(e0x) - self._dxdy(e0y)) + self.linear_coeff0 * e0x + k_e0y = self.diffraction_coeff0 * (self._d2x(e0y) - self._dxdy(e0x)) + self.linear_coeff0 * e0y + # pump depletion (MATLAB lines 1640-1646): no conjugate, w1 denominator + k_e0x += self.depletion_coeff0 * laplacian_phi * e1x + k_e0y += self.depletion_coeff0 * laplacian_phi * e1y + row_i0, row_i0p1 = self.calc_pump_source(t, pump_args) + k_e0y = k_e0y.at[self.i0, :].add(row_i0) + k_e0y = k_e0y.at[self.i0 + 1, :].add(row_i0p1) + + # Raman: propagation + detuning + SRS coupling (MATLAB lines 1661-1689) + k_e1x = self.diffraction_coeff1 * (self._d2y(e1x) - self._dxdy(e1y)) + self.linear_coeff1 * e1x + k_e1y = self.diffraction_coeff1 * (self._d2x(e1y) - self._dxdy(e1x)) + self.linear_coeff1 * e1y + k_e1x += self.srs_coeff1 * jnp.conj(laplacian_phi) * e0x + k_e1y += self.srs_coeff1 * jnp.conj(laplacian_phi) * e0y + if seed_args is not None: + row_i1, row_i1p1 = self.calc_seed_source(t, seed_args) + k_e1y = k_e1y.at[self.i1, :].add(row_i1) + k_e1y = k_e1y.at[self.i1 + 1, :].add(row_i1p1) + + return jnp.stack([k_e0x, k_e0y], axis=-1), jnp.stack([k_e1x, k_e1y], axis=-1) + + def __call__(self, t: float, E0: Array, E1: Array, phi_k: Array, pump_args: dict, seed_args: dict | None): + """ + Advance (E0, E1) over one EPW step: self.n_sub staggered light sub-steps. + + Matches MATLAB lightSplitStep: both real parts are updated with the RHS at t_i, + then both imaginary parts with the RHS at t_i + dt/2; the EPW potential is held + fixed during the sub-steps. + """ + seed_args = seed_args if self.seed_enabled else None + laplacian_phi = jnp.fft.ifft2(-self.k_sq * phi_k) + + def substep(i, fields): + E0, E1 = fields + t_i = t + i * self.dt_l + k_e0, k_e1 = self.rhs(t_i, E0, E1, laplacian_phi, pump_args, seed_args) + E0 = E0 + self.dt_l * jnp.real(k_e0) + E1 = E1 + self.dt_l * jnp.real(k_e1) + k_e0, k_e1 = self.rhs(t_i + self.dt_l / 2.0, E0, E1, laplacian_phi, pump_args, seed_args) + E0 = E0 + 1j * self.dt_l * jnp.imag(k_e0) + E1 = E1 + 1j * self.dt_l * jnp.imag(k_e1) + E0 = E0 * self.sub_boundary[..., None] + E1 = E1 * self.sub_boundary[..., None] + return (E0, E1) + + return lax.fori_loop(0, self.n_sub, substep, (E0, E1)) diff --git a/adept/_lpse2d/core/vector_field.py b/adept/_lpse2d/core/vector_field.py index 6ac08a80..5b984923 100644 --- a/adept/_lpse2d/core/vector_field.py +++ b/adept/_lpse2d/core/vector_field.py @@ -4,6 +4,7 @@ from adept._base_ import get_envelope from adept._lpse2d.core import epw, laser +from adept._lpse2d.core.light import CoupledLight from adept._lpse2d.core.raman import RamanLight @@ -25,8 +26,15 @@ def __init__(self, cfg): # self.epw = epw.SpectralPotential(cfg) self.epw = epw.SpectralEPWSolver(cfg) self.light = laser.Light(cfg) - # the Raman scattered light is evolved iff the SRS source term is on - self.raman = RamanLight(cfg) if cfg["terms"]["epw"]["source"].get("srs", False) else None + # the Raman scattered light is evolved iff the SRS source term is on; with + # terms.light.pump_depletion the pump is evolved too (one coupled solver) + self.pump_depletion = cfg["terms"].get("light", {}).get("pump_depletion", False) + srs_on = cfg["terms"]["epw"]["source"].get("srs", False) + if self.pump_depletion: + self.coupled_light = CoupledLight(cfg) + self.raman = None + else: + self.raman = RamanLight(cfg) if srs_on else None self.complex_state_vars = ["E0", "epw", "E1"] self.boundary_envelope = cfg["grid"]["absorbing_boundaries"] self.one_over_ksq = cfg["grid"]["one_over_ksq"] @@ -68,6 +76,14 @@ def get_envelope_coefficient(self, envelope_args, t): ) def light_split_step(self, t, y, driver_args): + if self.pump_depletion: + # the pump is a dynamic field sourced by its boundary injector; both light + # waves advance inside one staggered update (absorbers applied per sub-step) + y["E0"], y["E1"] = self.coupled_light( + t, y["E0"], y["E1"], y["epw"], driver_args["E0"], driver_args.get("E1") + ) + return y + if "E0" in driver_args: def E0_fn(this_t): diff --git a/adept/_lpse2d/datamodel.py b/adept/_lpse2d/datamodel.py index f5c3846e..5d9f7bca 100644 --- a/adept/_lpse2d/datamodel.py +++ b/adept/_lpse2d/datamodel.py @@ -160,6 +160,8 @@ class DampingModel(BaseModel): class SourceModel(BaseModel): noise: bool + noise_amplitude: float = 1e-10 + noise_seed: int | None = None tpd: bool srs: bool = False @@ -172,8 +174,16 @@ class EPWModel(BaseModel): source: SourceModel +class LightModel(BaseModel): + """Light-wave evolution options. pump_depletion evolves E0 with the FD envelope + solver (boundary injector + EPW coupling) instead of prescribing it analytically.""" + + pump_depletion: bool = False + + class TermsModel(BaseModel): epw: EPWModel + light: LightModel = LightModel() zero_mask: bool diff --git a/adept/_lpse2d/diagnostics.py b/adept/_lpse2d/diagnostics.py index 670ee641..a1b63a96 100644 --- a/adept/_lpse2d/diagnostics.py +++ b/adept/_lpse2d/diagnostics.py @@ -214,6 +214,18 @@ def series_metrics(series, cfg: dict) -> dict[str, float]: s_left = inc - refl # F0_left + F1_left, in units of I0_code s_right = trans + backrefl + # With an evolved pump the injector's launched amplitude carries a small grid- + # dispersion deficit (sin(k0 dx)/sin(k_grid dx), ~2% at 8 cells/wavelength), so + # the budget is normalized to the *measured* incident flux -- i.e. fractions of + # the pump that actually entered the box, exactly like OSIRIS normalizes to its + # own launched flux. With a prescribed pump the physical incident flux is + # nominal by construction, so normalize by 1. + inc_ref = 1.0 + if pump_evolved and len(inc) > 1: + inc_ref = float(np.mean(inc[len(inc) // 2 :])) + if not np.isfinite(inc_ref) or inc_ref <= 0: + inc_ref = 1.0 + windows = segment_windows(t_ps, last_frac=0.25, n_segments=4) for win in windows: suffix = "" if win["is_last"] else f"_seg{win['index']}of{win['n']}" @@ -221,16 +233,16 @@ def series_metrics(series, cfg: dict) -> dict[str, float]: t_naive = _wmean(trans, win["mask"]) if r_naive is None: continue - metrics[f"laser_reflectivity_naive{suffix}"] = r_naive - metrics[f"laser_transmissivity_naive{suffix}"] = t_naive + metrics[f"laser_reflectivity_naive{suffix}"] = r_naive / inc_ref + metrics[f"laser_transmissivity_naive{suffix}"] = t_naive / inc_ref if pump_evolved: - metrics[f"laser_reflectivity{suffix}"] = 1.0 - _wmean(s_left, win["mask"]) - metrics[f"laser_transmissivity{suffix}"] = _wmean(s_right, win["mask"]) - metrics[f"laser_absorbed_frac{suffix}"] = _wmean(s_left - s_right, win["mask"]) + metrics[f"laser_reflectivity{suffix}"] = 1.0 - _wmean(s_left, win["mask"]) / inc_ref + metrics[f"laser_transmissivity{suffix}"] = _wmean(s_right, win["mask"]) / inc_ref + metrics[f"laser_absorbed_frac{suffix}"] = _wmean(s_left - s_right, win["mask"]) / inc_ref else: # pump is prescribed: the one-way Raman flux is the only measurable # reflectivity, and there is no transmission/absorption budget - metrics[f"laser_reflectivity{suffix}"] = r_naive + metrics[f"laser_reflectivity{suffix}"] = r_naive / inc_ref if "epw_dissipation" in series: d = _wmean(np.asarray(series["epw_dissipation"].values, dtype=float), win["mask"]) if d is not None and I0_osiris > 0: diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index a0cbacd3..595bb043 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -275,6 +275,18 @@ def get_derived_quantities(cfg: dict) -> dict: if epw_source.get("noise_seed") is None: epw_source["noise_seed"] = int(np.random.randint(2**20)) + pump_depletion = cfg["terms"].get("light", {}).get("pump_depletion", False) + if pump_depletion: + if not cfg["terms"]["epw"]["source"].get("srs", False): + raise ValueError("terms.light.pump_depletion requires terms.epw.source.srs: true") + if cfg["drivers"].get("E0", {}).get("speckle", {}).get("enabled", False): + raise ValueError("terms.light.pump_depletion does not support drivers.E0.speckle yet") + if cfg["terms"]["epw"]["boundary"]["x"] != "absorbing": + raise ValueError( + "terms.light.pump_depletion requires terms.epw.boundary.x: absorbing " + "(the pump is launched by a boundary injector and must exit the box)" + ) + if cfg["terms"]["epw"]["source"].get("srs", False): derived = cfg["units"]["derived"] wpe_max_sq = derived["w0"] ** 2 * max(cfg["density"].get("max", 1.0), cfg["density"].get("min", 0.0)) @@ -282,6 +294,14 @@ def get_derived_quantities(cfg: dict) -> dict: 2.0 * derived["c"] ** 2 / (cfg_grid["dx"] ** 2 * derived["w1"]) + np.abs(derived["w1"] ** 2 - wpe_max_sq) / (4.0 * derived["w1"]) ) + if pump_depletion: + # the evolved pump has its own (looser, but not guaranteed) stability limit; + # sub-cycle to the tighter of the two carriers + dt_max_pump = 1.0 / ( + 2.0 * derived["c"] ** 2 / (cfg_grid["dx"] ** 2 * derived["w0"]) + + np.abs(derived["w0"] ** 2 - wpe_max_sq) / (4.0 * derived["w0"]) + ) + dt_max = min(dt_max, dt_max_pump) if "light_substeps" in cfg_grid: n_sub = int(cfg_grid["light_substeps"]) if cfg_grid["dt"] / n_sub > dt_max: @@ -322,6 +342,17 @@ def get_derived_quantities(cfg: dict) -> dict: "yw": _Q(cfg["drivers"][k]["yw"]).to("um").value if "yw" in cfg["drivers"][k] else 0.0, } continue + if k == "E0" and pump_depletion: + # boundary-injector parameters for the evolved pump: the injector sits at + # xmin + offset (default 2*boundary_width, clear of the absorber skirt) + boundary_width = _Q(cfg["grid"]["boundary_width"]).to("um").value + if "offset" in cfg["drivers"][k]: + cfg["drivers"][k]["derived"]["offset"] = _Q(cfg["drivers"][k]["offset"]).to("um").value + else: + cfg["drivers"][k]["derived"]["offset"] = 2.0 * boundary_width + cfg["drivers"][k]["derived"]["turn_on_time"] = ( + _Q(cfg["drivers"][k].get("turn_on_time", "10fs")).to("ps").value + ) cfg["drivers"][k]["derived"]["tw"] = _Q(cfg["drivers"][k]["envelope"]["tw"]).to("ps").value cfg["drivers"][k]["derived"]["tc"] = _Q(cfg["drivers"][k]["envelope"]["tc"]).to("ps").value cfg["drivers"][k]["derived"]["tr"] = _Q(cfg["drivers"][k]["envelope"]["tr"]).to("ps").value @@ -968,10 +999,39 @@ def get_default_save_func(cfg): x_grid = np.array(cfg["grid"]["x"]) ix_left = int(np.argmin(np.abs(x_grid - (cfg["grid"]["xmin"] + probe_offset)))) ix_right = int(np.argmin(np.abs(x_grid - (cfg["grid"]["xmax"] - probe_offset)))) + # with an evolved pump, the incident probe must sit downstream (+x) of the pump + # injector rows or it reads the near-field of the two-point source + ix_left_e0 = ix_left + if cfg["terms"].get("light", {}).get("pump_depletion", False): + pump_offset = cfg["drivers"]["E0"]["derived"]["offset"] + ix_inject = int(np.argmin(np.abs(x_grid - (cfg["grid"]["xmin"] + pump_offset)))) + ix_left_e0 = max(ix_left, ix_inject + 4) flux_coeff_w0 = derived["c"] ** 2 / (derived["w0"] * cfg["grid"]["dx"]) flux_coeff_w1 = derived["c"] ** 2 / (derived["w1"] * cfg["grid"]["dx"]) I0_code = derived["I0_code"] + def flux_correction(w, ix): + # The discrete two-point flux of the FD mode at local wavenumber k is + # |E|^2 * v_g,discrete with v_g,disc = (c^2/w) sin(k_grid dx)/dx, where + # k_grid satisfies the FD dispersion (2/dx^2)(1 - cos k_grid dx) = k^2. + # Dividing by sin(k_grid dx)/(k dx) converts it to the physical flux + # |E|^2 * c * sqrt(eps). Evanescent probes get 1 (their flux is ~0 anyway). + n_loc = float(np.mean(np.array(cfg["grid"]["background_density"])[ix, :])) + eps = 1.0 - n_loc * w0**2 / w**2 + if eps <= 0: + return 1.0 + k_dx = w / derived["c"] * np.sqrt(eps) * cfg["grid"]["dx"] + cos_kg = 1.0 - k_dx**2 / 2.0 + if cos_kg <= -1.0: + return 1.0 + sin_kg = float(np.sqrt(1.0 - cos_kg**2)) + return float(sin_kg / k_dx) + + corr_e0_left = flux_correction(w0, ix_left_e0) + corr_e0_right = flux_correction(w0, ix_right) + corr_e1_left = flux_correction(w1, ix_left) + corr_e1_right = flux_correction(w1, ix_right) + def discrete_flux(E, ix, coeff): # sum over polarization components, mean over y cross = jnp.sum(jnp.conj(E[ix, :, :]) * E[ix + 1, :, :], axis=-1) @@ -988,21 +1048,26 @@ def save_func(t, y, args): out = {"e_sq": jnp.sum(e_sq * cfg["grid"]["dx"] * cfg["grid"]["dy"]), "max_phi": jnp.max(jnp.abs(phi_k))} out["epw_energy"] = epw_energy_prefactor * jnp.sum(jnp.mean(e_sq, axis=1)) + # dissipation/boundary channels are TOTAL EPW-energy rates: the wave's kinetic + # (sloshing) energy equals the cycle-averaged electric energy that epw_energy + # counts (the OSIRIS field-only convention), so the energy actually handed to + # electrons -- and the budget sink -- is 2x the electric-part rate out["epw_dissipation"] = ( - epw_energy_prefactor / (nx * ny**2) * jnp.sum(k_sq * jnp.abs(phi_k) ** 2 * energy_loss_factor) + 2.0 * epw_energy_prefactor / (nx * ny**2) * jnp.sum(k_sq * jnp.abs(phi_k) ** 2 * energy_loss_factor) ) - out["epw_boundary_loss"] = epw_energy_prefactor * jnp.sum(jnp.mean(e_sq * boundary_sq_loss, axis=1)) + out["epw_boundary_loss"] = 2.0 * epw_energy_prefactor * jnp.sum(jnp.mean(e_sq * boundary_sq_loss, axis=1)) if srs_on: e1 = y["E1"].view(jnp.complex128) e0 = y["E0"].view(jnp.complex128) out["e1_sq"] = jnp.mean(jnp.sum(jnp.abs(e1) ** 2, axis=-1)) out["reflectivity"] = sqrt_eps1 * jnp.mean(jnp.abs(e1[ix_probe, :, 1]) ** 2) / E0_source_sq - # laser budget channels, all normalized to the nominal incident flux - out["incident_flux"] = discrete_flux(e0, ix_left, flux_coeff_w0) / I0_code - out["transmitted_flux"] = discrete_flux(e0, ix_right, flux_coeff_w0) / I0_code - out["reflected_flux"] = -discrete_flux(e1, ix_left, flux_coeff_w1) / I0_code - out["backrefl_flux"] = discrete_flux(e1, ix_right, flux_coeff_w1) / I0_code + # laser budget channels: physical fluxes (discrete fluxes converted via the + # local group-velocity factor), normalized to the nominal incident flux + out["incident_flux"] = discrete_flux(e0, ix_left_e0, flux_coeff_w0) / corr_e0_left / I0_code + out["transmitted_flux"] = discrete_flux(e0, ix_right, flux_coeff_w0) / corr_e0_right / I0_code + out["reflected_flux"] = -discrete_flux(e1, ix_left, flux_coeff_w1) / corr_e1_left / I0_code + out["backrefl_flux"] = discrete_flux(e1, ix_right, flux_coeff_w1) / corr_e1_right / I0_code return out diff --git a/docs/source/solvers/lpse2d/config.md b/docs/source/solvers/lpse2d/config.md index ac5a6f56..fcb0db47 100644 --- a/docs/source/solvers/lpse2d/config.md +++ b/docs/source/solvers/lpse2d/config.md @@ -65,15 +65,7 @@ Density profile configuration. | `gradient scale length` | string | Scale length with unit (for `linear` basis) | | `max` | float | Maximum density fraction (for `linear` basis) | | `min` | float | Minimum density fraction (for `linear` basis) | -| `noise` | object | Initial noise configuration | - -### noise - -| Field | Type | Description | -|-------|------|-------------| -| `max` | float | Maximum noise amplitude | -| `min` | float | Minimum noise amplitude | -| `type` | string | `"uniform"` or `"normal"` | +| `noise` | object | Ignored (legacy). The initial EPW is identically zero; noise-seeded runs use the per-step `terms.epw.source.noise` source instead | ### Example: Uniform Density @@ -117,7 +109,8 @@ Simulation grid parameters. Note: Grid values use physical units as strings. | `tmin` | string | Start time with unit | | `ymax` | string | Domain maximum y with unit | | `ymin` | string | Domain minimum y with unit | -| `light_substeps` | int | (SRS only, optional) Raman-light sub-steps per EPW step. Computed from the explicit-scheme stability limit `dt_light < 1 / (2c^2/(dx^2 w1) + \|w1^2 - wpe_max^2\|/(4 w1))` if omitted; a `ValueError` is raised if a user-supplied value violates it | +| `light_substeps` | int | (SRS only, optional) Raman-light sub-steps per EPW step. Computed from the explicit-scheme stability limit `dt_light < 1 / (2c^2/(dx^2 w1) + \|w1^2 - wpe_max^2\|/(4 w1))` if omitted (with `terms.light.pump_depletion` the limit is the tighter of the `w0` and `w1` carriers); a `ValueError` is raised if a user-supplied value violates it | +| `probe_offset` | string | (SRS only, optional) Distance of the laser-budget flux probes from each box edge, with unit. Default `2 * boundary_width`, which is clear of the absorber's tanh skirt (the legacy `reflectivity` probe at `1.6 * boundary_width` sits inside it and reads ~10% low) | Note: `nx` and `ny` are computed automatically from the grid parameters. The grid is optimized for FFT performance (sizes with small prime factors). @@ -214,6 +207,8 @@ The main laser pump for TPD/SRS simulations. | `delta_omega_max` | float | Maximum frequency spread (optional) | | `num_colors` | int | Number of laser colors (optional) | | `shape` | string | Amplitude shape: `"uniform"` (optional) | +| `offset` | string | (pump depletion only, optional) Distance of the pump boundary injector from `xmin`, with unit. Default `2 * boundary_width` | +| `turn_on_time` | string | (pump depletion only, optional) Gaussian turn-on time of the injector. Default `10fs` | #### envelope @@ -315,8 +310,15 @@ Physics terms configuration. | Field | Type | Description | |-------|------|-------------| | `epw` | object | Electron plasma wave configuration | +| `light` | object | (optional) Light-wave evolution configuration | | `zero_mask` | bool | Whether to zero out k=0 mode | +### light (optional) + +| Field | Type | Description | +|-------|------|-------------| +| `pump_depletion` | bool | (default `false`) Evolve the pump `E0` with the same staggered FD envelope solver as the Raman light instead of prescribing it analytically. The pump is launched by a two-point boundary injector at `x = xmin + drivers.E0.offset` and is depleted by the EPW through `-i e/(4 w1 me) (laplacian phi) E1`. Requires `terms.epw.source.srs: true` and `terms.epw.boundary.x: absorbing`; incompatible with `drivers.E0.speckle`. Enables the true net-flux `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac` metrics and lets above-threshold runs saturate | + ### epw | Field | Type | Description | @@ -349,6 +351,8 @@ Physics terms configuration. | Field | Type | Description | |-------|------|-------------| | `noise` | bool | Add random noise source | +| `noise_amplitude` | float | (optional) Amplitude of the per-step EPW noise source. Default `1.0e-10` (the MATLAB `noiseAmp`) | +| `noise_seed` | int | (optional) Seed for the EPW noise source. Default `null`, which draws a random seed once and pins it into the config before parameters are logged, so every run is exactly reproducible from its logged `noise_seed` | | `tpd` | bool | Include two-plasmon decay source | | `srs` | bool | Include stimulated Raman scattering (optional, default false). Turning this on also evolves the Raman scattered-light field `E1` with a finite-difference paraxial solver, sub-cycled `grid.light_substeps` times per EPW step, and adds the SRS source `i e wp0/(4 me w0 w1) (n/n_env) E0 . conj(E1)` to the EPW potential. The default time series then also records `e1_sq` and `reflectivity` (Poynting-corrected `|E1_y|^2/E0^2` at a probe on the low-density side, `x = 1.6 * boundary_width`) | diff --git a/docs/source/solvers/lpse2d/overview.md b/docs/source/solvers/lpse2d/overview.md index 2c5e3945..10864400 100644 --- a/docs/source/solvers/lpse2d/overview.md +++ b/docs/source/solvers/lpse2d/overview.md @@ -6,9 +6,24 @@ These equations model the evolution and interaction of the complex envelopes of ### Note on Pump Depletion -One can solve these equations with or without "pump depletion". "Pump depletion" is the effect of the plasma waves on the light waves. We do not currently have this implemented, so we have light waves that behave as external drivers for the plasma waves and we only model the plasma wave response. +One can solve these equations with or without "pump depletion". "Pump depletion" is the effect of the plasma waves on the light waves. By default the pump is prescribed analytically (an external driver for the plasma waves), which is adequate below the absolute instability threshold; above it the plasma-wave and Raman fields grow without bound because nothing depletes the pump. -This approach is adequate for modeling laser plasma instabilities below the absolute instability threshold. +For SRS, setting `terms.light.pump_depletion: true` evolves the pump `E0` with the same finite-difference envelope solver as the Raman light, sourced by a two-point boundary injector at the low-density side and coupled to the EPW through $-i e/(4 \omega_1 m_e) (\nabla^2 \phi) \mathbf{E}_1$ (the conjugate partner of the Raman coupling; the pair conserves $\int |E_0|^2 + |E_1|^2 + |\nabla\phi|^2$ together with the EPW source). This enables true transmission/absorption diagnostics and saturation of above-threshold runs. + +### SRS diagnostics + +With `terms.epw.source.srs` on, the default time series includes OSIRIS-comparable channels (fields converted to $m_e c \omega_0/e$, lengths to $c/\omega_0$, fluxes normalized to the nominal incident flux): + +| series | meaning | OSIRIS scan2 counterpart | +|---|---|---| +| `epw_energy` | $\frac{1}{4}\,dx \sum_x \langle\|E_{epw}\|^2\rangle_y$ (cycle-averaged field energy) | `W(t) = 1/2 dx sum e1^2` | +| `epw_dissipation` | total EPW energy handed to electrons per ps (Landau + collisional, the solver's own rates; includes the kinetic half) | absorbed fraction / hot-electron source | +| `epw_boundary_loss` | total EPW energy removed by the absorbing boundaries per ps | — | +| `incident_flux`, `transmitted_flux` | pump flux at probes near each edge | `incident_t`, `T_t` | +| `reflected_flux`, `backrefl_flux` | Raman flux leaving left / right | `R_t` | +| `reflectivity`, `e1_sq` | legacy probes (unchanged) | — | + +`post_process` logs scalar metrics with the same names and definitions as `osiris_lpi` (`laser_reflectivity`, `laser_transmissivity`, `laser_absorbed_frac`, per-quarter `_seg{i}of4` variants, `epw_growth_rate` per $\omega_0$ with the identical automated fit window, `electron_energy_frac_final`, and the `laser_incident_flux_ratio` health check). ### Electron Plasma Waves diff --git a/tests/test_lpse2d/test_srs.py b/tests/test_lpse2d/test_srs.py index c2e78dfe..126fbfb0 100644 --- a/tests/test_lpse2d/test_srs.py +++ b/tests/test_lpse2d/test_srs.py @@ -24,6 +24,15 @@ def _run(cfg): return exo, sol +def _run_with_ppo(cfg): + from adept import ergoExo + + exo = ergoExo() + modules = exo.setup(cfg) + sol, ppo, mlrunid = exo(modules) + return exo, sol, ppo + + def _predicted_gamma0(cfg_units, cfg_grid): """ Homogeneous backward-SRS growth rate gamma0 = k vos/4 * wpe / sqrt(w_ek * w_s) @@ -139,6 +148,189 @@ def test_srs_seed_propagation(): np.testing.assert_allclose(np.mean(np.abs(e1y_final[bulk])), expected_amp, rtol=0.3) +def _quasi_1d(cfg): + cfg["grid"]["ymax"] = "0.02um" + cfg["grid"]["ymin"] = "-0.02um" + return cfg + + +def test_pump_injector_calibration(): + """With pump depletion on and no plasma-wave activity, the evolved pump should fill + the box with the nominal flux and the nominal (swelled) amplitude.""" + cfg = _quasi_1d(_load_cfg()) + cfg["terms"]["light"] = {"pump_depletion": True} + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["grid"]["xmax"] = "20um" + cfg["grid"]["tmax"] = "0.3ps" + cfg["save"]["fields"]["t"]["tmax"] = "0.3ps" + cfg["save"]["fields"]["t"]["dt"] = "0.05ps" + cfg["mlflow"]["run"] = "srs-pump-injector-test" + + exo, sol = _run(cfg) + result = sol["solver result"] + t = np.array(result.ts["default"]) + steady = t > 0.2 # past the ~0.08 ps fill-in transit + + # the two-point source launches amplitude E_src * sin(k0 dx)/sin(k_grid dx), where + # k_grid is the FD-dispersion wavenumber -- so the measured physical flux is S^2 of + # nominal; both probes must agree with that and with each other (no spurious loss) + dcfg = exo.adept_module.cfg + derived = dcfg["units"]["derived"] + n = dcfg["units"]["envelope density"] + dx = dcfg["grid"]["dx"] + k0_dx = derived["w0"] / derived["c"] * np.sqrt(1.0 - n) * dx + kg_dx = np.arccos(1.0 - k0_dx**2 / 2.0) + S = np.sin(k0_dx) / np.sin(kg_dx) + + incident = np.array(result.ys["default"]["incident_flux"]) + transmitted = np.array(result.ys["default"]["transmitted_flux"]) + np.testing.assert_allclose(np.mean(incident[steady]), S**2, rtol=0.02) + np.testing.assert_allclose(np.mean(transmitted[steady]), np.mean(incident[steady]), rtol=0.01) + + # bulk amplitude: |E0| = E0_source * S * (1 - n)^(-1/4) + e0_raw = np.array(result.ys["fields"]["E0"]) + e0 = e0_raw.view(np.complex64 if e0_raw.dtype == np.float32 else np.complex128) + e0y_final = e0[-1, :, 0, 1] + x = np.array(dcfg["grid"]["x"]) + bulk = slice(np.argmin(np.abs(x - 8.0)), np.argmin(np.abs(x - 14.0))) + expected_amp = derived["E0_source"] * S * (1.0 - n) ** -0.25 + np.testing.assert_allclose(np.mean(np.abs(e0y_final[bulk])), expected_amp, rtol=0.03) + + +def _amplifier_cfg(): + """Seeded Raman-amplifier setup: pump + strong E1 seed, deterministic (no noise).""" + cfg = _quasi_1d(_load_cfg()) + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["drivers"]["E1"] = { + "intensity": "1.0e+14W/cm^2", + # seed on the Bohm-Gross-shifted resonance: w_seed = w0 - w_ek(k0+k1) rather + # than the envelope carrier w1 = w0 - wp0 (n = 0.2, Te = 2 keV => dw1 = -0.0335) + "delta_omega": -0.0335, + "turn_on_time": "10fs", + } + cfg["grid"]["xmax"] = "40um" + cfg["grid"]["tmax"] = "0.8ps" + cfg["save"]["fields"]["t"]["tmax"] = "0.8ps" + cfg["save"]["fields"]["t"]["dt"] = "0.1ps" + return cfg + + +def test_pump_depletion_budget_and_saturation(): + """Raman amplifier with pump depletion: the energy budget must close, and the + transmitted pump must actually deplete relative to the prescribed-pump run.""" + cfg = _amplifier_cfg() + cfg["terms"]["light"] = {"pump_depletion": True} + cfg["mlflow"]["run"] = "srs-depletion-budget-test" + exo, sol, ppo = _run_with_ppo(cfg) + ppo_metrics = ppo.get("metrics", {}) if isinstance(ppo, dict) else {} + + result = sol["solver result"] + t = np.array(result.ts["default"]) + win = t > 0.75 * t[-1] + + inc = np.array(result.ys["default"]["incident_flux"]) + refl = np.array(result.ys["default"]["reflected_flux"]) + trans = np.array(result.ys["default"]["transmitted_flux"]) + backrefl = np.array(result.ys["default"]["backrefl_flux"]) + dissip = np.array(result.ys["default"]["epw_dissipation"]) + bloss = np.array(result.ys["default"]["epw_boundary_loss"]) + W = np.array(result.ys["default"]["epw_energy"]) + + # sanity: the amplifier actually amplified and the pump actually depleted + assert np.mean(refl[win]) > 2.0 * 1e14 / 1.5e15, "seed was not amplified" + assert np.mean(trans[win]) < 0.97, "pump did not deplete" + assert np.mean(refl[win]) <= 1.05, "reflectivity exceeds the incident flux" + + # budget closure over the quasi-steady window, all in units of the incident flux: + # what goes missing between the probes (S_left - S_right) must equal EPW dissipation + # + absorber losses + EPW energy change. dissip/bloss are total-energy rates already; + # W is the OSIRIS field-only energy so the total stored EPW energy changes at 2*dW/dt. + # Convert to flux units: I0_osiris = a0^2/2, energy rates are per ps -> / w0. + derived = exo.adept_module.cfg["units"]["derived"] + w0 = derived["w0"] + a0 = derived["E0_source"] * derived["e_norm"] + I0_osiris = 0.5 * a0**2 + absorbed_measured = np.mean((inc - refl - trans - backrefl)[win]) + dWdt = np.gradient(W, t) # OSIRIS (field-only) energy per ps + absorbed_predicted = np.mean((dissip + bloss + 2.0 * dWdt)[win]) / w0 / I0_osiris + assert abs(absorbed_measured - absorbed_predicted) < 0.1 * max(np.mean(inc[win]), 1e-30), ( + f"budget does not close: measured {absorbed_measured:.3e} vs predicted {absorbed_predicted:.3e}" + ) + + # the definitional identity R + T + absorbed = 1 from the metrics + assert "laser_absorbed_frac" in ppo_metrics, "budget metrics missing from post_process" + total = ( + ppo_metrics["laser_reflectivity"] + ppo_metrics["laser_transmissivity"] + ppo_metrics["laser_absorbed_frac"] + ) + np.testing.assert_allclose(total, 1.0, atol=1e-6) + + # comparison run with the pump prescribed: transmission cannot deplete there + cfg2 = _amplifier_cfg() + cfg2["mlflow"]["run"] = "srs-prescribed-comparison-test" + exo2, sol2 = _run(cfg2) + result2 = sol2["solver result"] + trans2 = np.array(result2.ys["default"]["transmitted_flux"]) + assert np.mean(trans[win]) < np.mean(trans2[win]) - 0.02, "depletion did not reduce transmission" + + +def test_epw_energy_normalization(): + """The epw_energy save quantity equals the hand-computed OSIRIS-unit energy for a + synthetic single-mode phi_k.""" + from copy import deepcopy + + import jax.numpy as jnp + + from adept._lpse2d import helpers + + cfg = deepcopy(_quasi_1d(_load_cfg())) + helpers.write_units(cfg) + cfg = helpers.get_derived_quantities(cfg) + cfg["grid"] = helpers.get_solver_quantities(cfg) + cfg["grid"]["background_density"] = helpers.get_density_profile(cfg) + cfg = helpers.get_save_quantities(cfg) + save_func = cfg["save"]["default"]["func"] + + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + ik = 5 + amp = 1e-3 + phi_k = np.zeros((nx, ny), dtype=np.complex128) + phi_k[ik, 0] = amp + y = { + "epw": jnp.array(phi_k).view(jnp.float64), + "E0": jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64), + "E1": jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64), + } + out = save_func(0.0, y, None) + + derived = cfg["units"]["derived"] + k = cfg["grid"]["kx"][ik] + # single k-mode: |ex(x)| = k*amp/(nx*ny) everywhere; sum_x mean_y |ex|^2 = nx*(k*amp/(nx*ny))^2 + expected = ( + 0.25 * cfg["grid"]["dx"] * derived["x_norm"] * derived["e_norm"] ** 2 * nx * (k * amp / (nx * ny)) ** 2 + ) + np.testing.assert_allclose(float(out["epw_energy"]), expected, rtol=1e-10) + + +def test_noise_seed_reproducibility(): + """Two runs with the same explicit noise seed are bit-identical; a different seed is not.""" + e_sq = {} + for name, seed in [("a", 1234), ("b", 1234), ("c", 4321)]: + cfg = _quasi_1d(_load_cfg()) + cfg["grid"]["xmax"] = "20um" + cfg["grid"]["tmax"] = "0.15ps" + cfg["save"]["fields"]["t"]["tmax"] = "0.15ps" + cfg["save"]["fields"]["t"]["dt"] = "0.05ps" + cfg["terms"]["epw"]["source"]["noise_seed"] = seed + cfg["mlflow"]["run"] = f"srs-seed-repro-{name}" + _, sol = _run(cfg) + e_sq[name] = np.array(sol["solver result"].ys["default"]["e_sq"]) + + np.testing.assert_array_equal(e_sq["a"], e_sq["b"]) + assert not np.array_equal(e_sq["a"], e_sq["c"]), "different seeds gave identical noise" + + if __name__ == "__main__": test_srs_seed_propagation() test_srs_growth_rate()