feat(lpse2d): parity with the original LPSE (deck translator, term/solver/kinetic parity, cross-checks)

adept/_lpse2d/core/laser.py

@@ -1,3 +1,4 @@
+import numpy as np
 from jax import Array
 from jax import numpy as jnp
 
@@ -16,7 +17,24 @@ def __init__(self, cfg) -> None:
         self.Lx = self.nx * self.dx  # box length in x (um)
         self.dE0x = jnp.zeros((cfg["grid"]["nx"], cfg["grid"]["ny"]))
         self.x = cfg["grid"]["x"]
+        self.y = cfg["grid"]["y"]
+        self.dy = cfg["grid"]["dy"]
+        self.dky = 2.0 * jnp.pi / (self.ny * self.dy)
         self.background_density = cfg["grid"]["background_density"]
+        # in-plane angle of incidence from +x (drivers.E0.angle, degrees; LPSE laser.N.direction).
+        # The pump is one k-mode snapped to the grid in kx *and* ky (LPSE makeStaticField),
+        # polarised in the plane perpendicular to the snapped k (LPSE polarization 0)
+        self.angle = float(np.deg2rad(float(cfg["drivers"].get("E0", {}).get("angle", 0.0))))
+        pump = cfg["drivers"].get("E0", {}).get("derived", {})
+        self.beam_angle = np.atleast_1d(np.asarray(pump.get("beam_angle", [self.angle]), dtype=np.float64))
+        self.beam_fraction = np.atleast_1d(np.asarray(pump.get("beam_fraction", [1.0]), dtype=np.float64))
+        self.beam_phase = np.atleast_1d(np.asarray(pump.get("beam_phase", [0.0]), dtype=np.float64))
+        self.beam_delta_omega = np.atleast_1d(np.asarray(pump.get("beam_delta_omega", [0.0]), dtype=np.float64))
+        self.kap_bandwidth = float(pump.get("kap_bandwidth", 0.0) or 0.0)
+        self.kap_seed = int(pump.get("kap_seed", 0) or 0)
+        self.pulse_t = jnp.asarray(pump["pulse_t"]) if "pulse_t" in pump else None
+        self.pulse_amp = jnp.asarray(pump["pulse_amp"]) if "pulse_amp" in pump else None
+        self.multi_beam = len(self.beam_angle) > 1 or float(self.beam_angle[0]) != 0.0
 
         # Speckle state
         self.speckle_profile = None
@@ -63,11 +81,17 @@ def laser_update(self, t_ps: float, y: jnp.ndarray, light_wave: dict) -> tuple[j
         # (envelope-density) plasma frequency wp0 and snapped to the nearest FFT grid mode, so the
         # pump is exactly periodic on the grid. The local density swelling is applied as an
         # amplitude factor *after* the transform (not as a spatially-varying phase).
+        if self.multi_beam or self.kap_bandwidth > 0.0 or self.pulse_t is not None:
+            return self._oblique_update(t_ps, light_wave)
+        # one beam at normal incidence: its own frequency shift and phase (drivers.E0.beams[0], LPSE
+        # laser.1.frequencyShift / phase) add to every color's, as in the oblique and evolved paths
+        dw_b = float(self.beam_delta_omega[0])
+        phase_b = float(self.beam_phase[0])
         E0y_k = jnp.zeros((self.nx, self.ny), dtype=jnp.complex128)
         for i in range(len(light_wave["delta_omega"])):
-            delta_omega = light_wave["delta_omega"][i]
+            delta_omega = light_wave["delta_omega"][i] + dw_b
             intensity = light_wave["intensities"][i, :]  # (ny,)
-            phase = light_wave["phases"][i, :]  # (ny,)
+            phase = light_wave["phases"][i, :] + phase_b  # (ny,)
 
             # reference-density wavenumber, snapped to the FFT grid (MATLAB lines 1558-1559)
             k0 = self.w0 / self.c * jnp.sqrt((1.0 + delta_omega) ** 2 - self.wp0**2 / self.w0**2)
@@ -99,6 +123,48 @@ def laser_update(self, t_ps: float, y: jnp.ndarray, light_wave: dict) -> tuple[j
 
         return jnp.stack([self.dE0x, dE0y], axis=-1)
 
+    def _oblique_update(self, t_ps: float, light_wave: dict) -> jnp.ndarray:
+        """Oblique static pump: each color is the single grid mode nearest to
+        ``k0(delta_omega) (cos a, sin a)`` (LPSE ``makeStaticField`` rounds both components),
+        built in x-space, with the field along ``(-sin a', cos a')`` of the *snapped* direction
+        so that E0 is exactly transverse; the swelling factor is applied as in the normal case."""
+        E0 = jnp.zeros((self.nx, self.ny, 2), dtype=jnp.complex128)
+        xx = self.x[:, None]
+        yy = self.y[None, :]
+        pulse = jnp.interp(t_ps, self.pulse_t, self.pulse_amp) if self.pulse_t is not None else 1.0
+        for b in range(len(self.beam_angle)):
+            angle = float(self.beam_angle[b])
+            dw_b = float(self.beam_delta_omega[b])
+            beam_phase = float(self.beam_phase[b]) + self._kap_phase(t_ps, b)
+            for i in range(len(light_wave["delta_omega"])):
+                delta_omega = light_wave["delta_omega"][i] + dw_b
+                intensity = light_wave["intensities"][i, :] * float(self.beam_fraction[b])  # (ny,)
+                phase = light_wave["phases"][i, :] + beam_phase  # (ny,)
+                k0 = self.w0 / self.c * jnp.sqrt((1.0 + delta_omega) ** 2 - self.wp0**2 / self.w0**2)
+                kx = jnp.round(k0 * jnp.cos(angle) / self.dk) * self.dk
+                ky = jnp.round(k0 * jnp.sin(angle) / self.dky) * self.dky
+                k_snapped = jnp.sqrt(kx**2 + ky**2)
+                pol = jnp.stack([-ky / k_snapped, kx / k_snapped])  # unit vector perpendicular to k
+                amp = (
+                    pulse * self.E0_source * jnp.sqrt(intensity) * jnp.exp(-1j * (delta_omega * self.w0 * t_ps - phase))
+                )
+                carrier = jnp.exp(1j * (kx * xx + ky * yy))
+                E0 = E0 + (amp[None, :] * carrier)[..., None] * pol[None, None, :]
+        wpe = self.w0 * jnp.sqrt(self.background_density)
+        E0 = E0 * ((1.0 - wpe**2 / self.w0**2) ** -0.25)[..., None]
+        if self.speckle_profile is not None:
+            raise NotImplementedError("drivers.E0.speckle with a non-zero drivers.E0.angle is not supported")
+        return E0
+
+    def _kap_phase(self, t_ps, beam: int):
+        """Kubo-Anderson phase jumps (see SpectralCoupledLight.kap_phase)."""
+        if self.kap_bandwidth <= 0.0:
+            return 0.0
+        tau = 2.0 * jnp.pi / (self.kap_bandwidth * self.w0)
+        index = jnp.floor(t_ps / tau)
+        seed = 12.9898 * (index + 1.0) + 78.233 * (beam + 1.0) + 37.719 * self.kap_seed
+        return 2.0 * jnp.pi * jnp.mod(jnp.sin(seed) * 43758.5453, 1.0)
+
     def calc_ey_at_one_point(self, t: float, density: Array, light_wave: dict) -> tuple[jnp.ndarray, jnp.ndarray]:
         """
         This function is used to calculate the coherence time of the laser

adept/_lpse2d/core/light.py

@@ -3,7 +3,7 @@
 from jax import numpy as jnp
 
 from adept._base_ import get_envelope
-from adept._lpse2d.core.raman import RamanLight
+from adept._lpse2d.core.raman import RamanLight, transverse_part
 
 
 class CoupledLight(RamanLight):
@@ -41,16 +41,20 @@ class CoupledLight(RamanLight):
     paths. This class adds the pump: its coefficients, boundary injector, reciprocal
     TPD/SRS couplings, and the coupled staggered loop.
 
-    The TPD pump term is constructed directly as the discrete reciprocal of
-    ``SpectralEPWSolver.calc_tpd_source`` (rather than using the abandoned historical
-    Pump2D prototype):
+    The TPD pump term is LPSE's (LightSolver.cpp, ``Sc_tpd``; LwSolver.cpp, the
+    transverse projection; Follett's equation document Eqs 53/55):
 
-        dE0_y/dt |_TPD = i e/(4 w0 me) exp(+i (w0 - 2 wp0)t)
-                          * Ehy * div(Eh)
+        dE0/dt |_TPD = i e/(2 w0 me) exp(+i (w0 - 2 wp0)t) * [Eh div(Eh)]_T
 
-    The pair conserves ``|E0|^2 + (wp0/w0)|Eh|^2`` for coupling terms alone. Only
-    the y-polarized pump participates because the existing TPD operator is written
-    for that polarization.
+    for every pump component, with ``[.]_T`` the transverse (divergence-free) part taken
+    in k-space (``terms.light.tpd_projection``, default on) so that a longitudinal
+    component is never injected into the light field; ``terms.light.tpd_k_filter``
+    additionally restricts the term to ``|k| < 1.2 k0 sqrt(1 - n_min)`` (LPSE
+    ``lw.kFilter``). The coefficient is exactly twice the EPW-side TPD source coefficient
+    ``i e/(4 me w0)``: with that ratio the pair conserves ``|E0|^2 + (2 wp0/w0)|Eh|^2``
+    for the coupling terms alone, i.e. the total wave energy at envelope density n_c/4
+    (settled against the equation document on 2026-09-14; adept main previously carried
+    half this coefficient).
 
     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
@@ -92,6 +96,7 @@ def __init__(self, cfg: dict):
         source_cfg = cfg["terms"]["epw"]["source"]
         self.srs_enabled = bool(source_cfg.get("srs", False))
         self.tpd_enabled = bool(source_cfg.get("tpd", False))
+        self.kx = cfg["grid"]["kx"]
         self.ky = cfg["grid"]["ky"]
 
         # pump detuning/diffraction (MATLAB lines 1616-1626); with wp0^2 = w0^2 * n_env
@@ -101,13 +106,46 @@ def __init__(self, cfg: dict):
         )
         self.diffraction_coeff0 = 1j * self.c**2 / (2.0 * self.w0)
         self.srs_depletion_coeff0 = -1j * self.e / (4.0 * self.w1 * self.me)
-        self.tpd_depletion_coeff0 = 1j * self.e / (4.0 * self.w0 * self.me)
+        self.tpd_depletion_coeff0 = 1j * self.e / (2.0 * self.w0 * self.me)
 
-        # how the SRS E0 <-> E1 exchange enters the light sub-step (see the class docstring)
         light_cfg = cfg["terms"].get("light", {})
+        # transverse projection of E div(E) (LPSE LwSolver::makeExyzDivE, k-space
+        # P_T = I - k k / k^2) and the optional LPSE lw.kFilter on the same term
+        self.tpd_projection = bool(light_cfg.get("tpd_projection", True))
+        self.tpd_k_filter = bool(light_cfg.get("tpd_k_filter", False))
+        kx_arr = np.asarray(cfg["grid"]["kx"])
+        ky_arr = np.asarray(cfg["grid"]["ky"])
+        k_sq_np = kx_arr[:, None] ** 2 + ky_arr[None, :] ** 2
+        k_sq_safe = np.where(k_sq_np > 0, k_sq_np, 1.0)
+        self.kx_over_k_sq = jnp.asarray(np.where(k_sq_np > 0, kx_arr[:, None] / k_sq_safe, 0.0))
+        self.ky_over_k_sq = jnp.asarray(np.where(k_sq_np > 0, ky_arr[None, :] / k_sq_safe, 0.0))
+        if self.tpd_k_filter:
+            n_min = float(np.min(np.asarray(background_density)))
+            k0_max_sq = (1.2 * self.w0 / self.c) ** 2 * max(1.0 - n_min, 0.0)
+            self.tpd_k_mask = jnp.asarray(np.where((k_sq_np > 0) & (k_sq_np < k0_max_sq), 1.0, 0.0))
+        else:
+            self.tpd_k_mask = None
+
+        # how the SRS E0 <-> E1 exchange enters the light sub-step (see the class docstring)
         self.coupling = str(light_cfg.get("coupling", "explicit"))
         if self.coupling not in ("explicit", "rotation"):
             raise ValueError(f"terms.light.coupling must be 'explicit' or 'rotation', got {self.coupling!r}")
+        if (
+            self.coupling == "rotation"
+            and self.srs_enabled
+            and self.transverse_source
+            and str(light_cfg.get("solver", "fd")) == "fd"
+        ):
+            # the exact exchange rotates E0 and E1 pointwise; the transverse projection of its
+            # sources is a k-space operation it cannot apply (the spectral solver keeps the fields
+            # transverse instead), and the FD rotation path used to skip it silently
+            raise ValueError(
+                "terms.light.coupling: rotation with the fd light solver cannot apply "
+                "terms.light.transverse_source: true (the default): use coupling: explicit or "
+                "terms.light.solver: spectral to keep the projection, or set transverse_source: false "
+                "explicitly -- what fd rotation runs did so far, with the longitudinal E1 left in the SRS "
+                "exchange"
+            )
         # local E0 <-> E1 exchange rate per unit |laplacian phi|: Omega = sqrt(|A B|) |lap phi|
         self.omega_prefactor = self.e / (4.0 * self.me * np.sqrt(self.w0 * self.w1))
 
@@ -143,6 +181,10 @@ def __init__(self, cfg: dict):
         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
+        # the single normal-incidence beam's own frequency shift and phase (drivers.E0.beams[0];
+        # oblique or multiple beams are refused on this path by setup)
+        self.beam_delta_omega0 = float(np.atleast_1d(pump.get("beam_delta_omega", [0.0]))[0])
+        self.beam_phase0 = float(np.atleast_1d(pump.get("beam_phase", [0.0]))[0])
 
     def calc_pump_source(self, t: float, pump_args: dict) -> tuple[Array, Array]:
         """
@@ -160,9 +202,9 @@ def calc_pump_source(self, t: float, pump_args: dict) -> tuple[Array, Array]:
         )
         turn_on = 1.0 - jnp.exp(-((t / self.pump_turn_on_time) ** 2))
 
-        delta_omega = pump_args["delta_omega"]  # (nc,)
+        delta_omega = pump_args["delta_omega"] + self.beam_delta_omega0  # (nc,)
         intensities = pump_args["intensities"]  # (nc, ny), fractions summing to 1
-        phases = pump_args["phases"]  # (nc, ny)
+        phases = pump_args["phases"] + self.beam_phase0  # (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%
@@ -202,28 +244,51 @@ def pump_rhs(
         k_e0x = self.diffraction_coeff0 * (self._d2y(e0x) - self._dxdy(e0y)) + linear_coeff0 * e0x
         k_e0y = self.diffraction_coeff0 * (self._d2x(e0y) - self._dxdy(e0x)) + linear_coeff0 * e0y
         if self.srs_enabled and couple:
-            k_e0x += self.srs_depletion_coeff0 * laplacian_phi * e1x
-            k_e0y += self.srs_depletion_coeff0 * laplacian_phi * e1y
+            depletion = (self.srs_depletion_coeff0 * laplacian_phi)[..., None] * E1
+            if self.transverse_source:
+                depletion = transverse_part(depletion, self.kx_arr, self.ky_arr, self.one_over_k_sq)
+            k_e0x += depletion[..., 0]
+            k_e0y += depletion[..., 1]
         if self.tpd_enabled:
             if phi_k is None:
                 raise ValueError("phi_k is required for TPD pump depletion")
-            k_e0y += self.calc_tpd_depletion(t, phi_k)
+            tpd_dep = self.calc_tpd_depletion(t, phi_k)
+            k_e0x += tpd_dep[..., 0]
+            k_e0y += tpd_dep[..., 1]
         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)
 
         return jnp.stack([k_e0x, k_e0y], axis=-1)
 
-    def calc_tpd_depletion(self, t: float, phi_k: Array) -> Array:
-        """Return the clean reciprocal TPD source for the y-polarized pump.
+    def tpd_depletion_vector(self, phi_k: Array) -> Array:
+        """``[E_h div(E_h)]_T`` in real space, shape (nx, ny, 2), from the frozen potential.
 
         ``div_e`` uses the same potential convention as the EPW TPD source:
-        ``div_e = ifft2(k^2 phi_k)``.
+        ``div_e = ifft2(k^2 phi_k)``. The transverse projection and the optional k-filter
+        are LPSE's ``LwSolver::makeExyzDivE``: ``F_T = F - k (k . F) / k^2`` mode by mode.
         """
+        ex = jnp.fft.ifft2(-1j * self.kx[:, None] * phi_k)
         ey = jnp.fft.ifft2(-1j * self.ky[None, :] * phi_k)
         div_e = jnp.fft.ifft2(self.k_sq * phi_k)
+        fx, fy = ex * div_e, ey * div_e
+        if self.tpd_projection or self.tpd_k_mask is not None:
+            fx_k = jnp.fft.fft2(fx)
+            fy_k = jnp.fft.fft2(fy)
+            if self.tpd_projection:
+                longitudinal = self.kx[:, None] * fx_k + self.ky[None, :] * fy_k
+                fx_k = fx_k - self.kx_over_k_sq * longitudinal
+                fy_k = fy_k - self.ky_over_k_sq * longitudinal
+            if self.tpd_k_mask is not None:
+                fx_k = fx_k * self.tpd_k_mask
+                fy_k = fy_k * self.tpd_k_mask
+            fx, fy = jnp.fft.ifft2(fx_k), jnp.fft.ifft2(fy_k)
+        return jnp.stack([fx, fy], axis=-1)
+
+    def calc_tpd_depletion(self, t: float, phi_k: Array) -> Array:
+        """Return the reciprocal TPD term for both pump components, shape (nx, ny, 2)."""
         phase = jnp.exp(1j * (self.w0 - 2.0 * self.wp0) * t)
-        return self.tpd_depletion_coeff0 * phase * ey * div_e
+        return self.tpd_depletion_coeff0 * phase * self.tpd_depletion_vector(phi_k)
 
     def coupled_rhs(
         self,
@@ -320,8 +385,17 @@ def substep(i, fields):
                 E0, E1 = propagate(t_i, E0, E1)
             E0 = E0 * self.sub_boundary[..., None]
             E1 = E1 * self.sub_boundary[..., None]
+            if absorb0 is not None:
+                E0 = E0 * absorb0
+                E1 = E1 * absorb1
             return (E0, E1)
 
+        absorb0 = absorb1 = None
+        if self.absorption_rate0 is not None:
+            n0 = self.n_over_nc0 if iaw_density is None else self.n_over_nc0 * (1.0 + iaw_density / self.n_over_env)
+            n1 = self.n_over_nc1 if iaw_density is None else self.n_over_nc1 * (1.0 + iaw_density / self.n_over_env)
+            absorb0 = jnp.exp(-self.absorption_rate0 * self.dt_l * n0**2)[..., None]
+            absorb1 = jnp.exp(-self.absorption_rate1 * self.dt_l * n1**2)[..., None]
         E0, E1 = lax.fori_loop(0, self.n_sub, substep, (E0, E1))
         if self.light_filter is not None:
             E0 = jnp.fft.ifft2(jnp.fft.fft2(E0, axes=(0, 1)) * self.light_filter, axes=(0, 1))