diff --git a/adept/_lpse2d/core/combined.py b/adept/_lpse2d/core/combined.py new file mode 100644 index 00000000..1d182eee --- /dev/null +++ b/adept/_lpse2d/core/combined.py @@ -0,0 +1,365 @@ +"""Combined TPD + SRS solver (``terms.epw.solver: combined``). + +This is the original LPSE ``lw.solver = combined`` formulation, which LPSE requires +whenever TPD and SRS are both on (users guide 1.5: the four separate envelope +equations "are not valid for simultaneous TPD and SRS"). One complex vector field +``E1``, enveloped at the reference plasma frequency ``wp0`` (= ``w0 - wp0`` at the +envelope density n_c/4, so the Raman light and the plasma wave share a carrier), +carries the Raman scattered light as its transverse part and the EPW as its +longitudinal part. The potential ``phi_k = i k . E1_k / k^2`` is *derived* from it +every step (``ZakharovSolver::makePotentialFromE_fft``) for the diagnostics, the IAW +ponderomotive drive and the HPE tracker; it is not an independent state variable. + +Per light sub-step (``LightSolver::evolveSpectral``, ``LightSolver.cpp:3600-3779``): + +1. x-space scattering potential on the whole field, ``exp(-i dt wp0/2 (n_tot/n_env - 1))`` + (the EPW density detuning, identical to the light detuning at carrier ``wp0``), and + the collisional damping ``exp(-nu_R dt (n/n_env)^2)``, ``n`` with the IAW perturbation -- + LPSE damps the combined field with its Raman absorption coefficient (``terms.light.absorption`` + at the combined carrier, ``raman.light_absorption_rates``) evaluated at the field's own + critical density, which is ``n_env``, and ignores ``lw.collisionalDampingRate`` + (``LightSolver::calculateScatteringPotential``); +2. the unified source, forward Euler (``LightSolver.cpp:4239``): + + dE1/dt = -i e/(4 me w0) e^{-i (w0 - 2 wp0) t} [ grad(E0 . E1*) + (1 - w0/wp0) E0 (div E1)* ] + + whose longitudinal part is the TPD source of ``epw.py`` (both terms) plus the SRS + source (from ``E0 . E1_T*``), and whose transverse part is the Raman-light + generation of ``raman.py`` (from ``E0 (div E1_L)*``) -- one expression for all four + couplings, with no separate-envelope approximation between them; +3. the k-space propagator with the 2x2 longitudinal/transverse projector + (``LightSolver.cpp:3684-3743``): ``exp(-i dt 3 vte^2 k^2/(2 wp0) - gamma_L dt)`` on + ``k k/k^2 E1`` (Bohm-Gross dispersion and Landau damping) and + ``exp(-i dt c^2 k^2/(2 wp0))`` on ``(I - k k/k^2) E1`` (light near its cutoff), with + modes outside the retained band zeroed; +4. the EPW noise source on the longitudinal part (``noise_model`` as in ``epw.py``), built + for the light sub-step at which it is applied (``epw.noise_step``; + ``LwSolver::addNoise_combinedSolver_spectral(dt)``); +5. the absorbing layers. + +A Raman seed (``drivers.E1``) is the spectral light solver's smooth injector on the transverse +``y`` component, added with the unified source at the combined carrier ``wp0`` (LPSE's Raman +class injects at ``W1 = wpe`` in combined mode, ``LightSolver.cpp:2404-2406``). + +With ``terms.light.pump_depletion`` the pump advances in the same sub-step with the +spectral propagator, the spectral light solver's injector (``SmoothPumpInjector``: oblique +and multiple beams, transverse profile, KAP bandwidth, pulse table) and LPSE's unified +depletion term (``LightSolver.cpp:4265``) + + dE0/dt = i e/(2 me w0) e^{+i (w0 - 2 wp0) t} E1 (div E1) + +which contains both the SRS depletion (``E1_T div E1``, ``i e/(4 me w1)`` at n_c/4) +and the TPD depletion (``E1_L div E1``); ``terms.light.tpd_projection`` applies the +transverse projection to it (LPSE's combined path does not, set it to ``false`` for +an exact match). Otherwise the pump is the prescribed field of ``laser.py``. +""" + +import numpy as np +from jax import Array, lax +from jax import numpy as jnp + +from adept._base_ import get_envelope +from adept._lpse2d.core.epw import analytic_landau_rate, noise_kick_spectrum, noise_step +from adept._lpse2d.core.raman import light_absorption_rates +from adept._lpse2d.core.spectral_light import SmoothPumpInjector, gaussian_injector_profile, smooth_seed_source + + +def longitudinal_transverse(field: Array, kx: Array, ky: Array, one_over_k_sq: Array) -> tuple[Array, Array]: + """k-space (fx_k, fy_k) split of an (nx, ny, 2) x-space field into its longitudinal and + transverse parts: returns ((lx_k, ly_k), (tx_k, ty_k)).""" + fx_k = jnp.fft.fft2(field[..., 0]) + fy_k = jnp.fft.fft2(field[..., 1]) + kdote = (kx[:, None] * fx_k + ky[None, :] * fy_k) * one_over_k_sq + lx_k, ly_k = kx[:, None] * kdote, ky[None, :] * kdote + return (lx_k, ly_k), (fx_k - lx_k, fy_k - ly_k) + + +def transverse_part(field: Array, kx: Array, ky: Array, one_over_k_sq: Array) -> Array: + """The transverse (Raman light) part of a combined field, in x-space.""" + _, (tx_k, ty_k) = longitudinal_transverse(field, kx, ky, one_over_k_sq) + return jnp.stack([jnp.fft.ifft2(tx_k), jnp.fft.ifft2(ty_k)], axis=-1) + + +def potential_from_field(field: Array, kx: Array, ky: Array, one_over_k_sq: Array, band: Array) -> Array: + """``phi_k = i k . E_k / k^2`` on the retained band (``makePotentialFromE_fft``).""" + fx_k = jnp.fft.fft2(field[..., 0]) + fy_k = jnp.fft.fft2(field[..., 1]) + return 1j * (kx[:, None] * fx_k + ky[None, :] * fy_k) * one_over_k_sq * band + + +def combined_collision_rate(cfg: dict) -> float: + """``nu_R`` (1/ps): the combined field's collisional damping rate at ``n = n_env``, the Raman light's + absorption (``terms.light.absorption``) at the combined carrier. The field is damped by + ``exp(-nu_R dt (n/n_env + dn)^2)`` in x-space; the EPW collisions play no part.""" + _, raman_rate = light_absorption_rates(cfg) + return 0.0 if raman_rate is None else float(raman_rate) + + +class CombinedSolver(SmoothPumpInjector): + """Advance the combined Raman-light + EPW field E1 (and the pump when it is evolved).""" + + def __init__(self, cfg: dict): + grid = cfg["grid"] + derived = cfg["units"]["derived"] + source_cfg = cfg["terms"]["epw"]["source"] + light_cfg = cfg["terms"].get("light", {}) + + self.tpd_enabled = bool(source_cfg.get("tpd", False)) + self.srs_enabled = bool(source_cfg.get("srs", False)) + if self.tpd_enabled != self.srs_enabled: + raise ValueError( + "terms.epw.solver: combined needs terms.epw.source.tpd and srs both on or both off " + "(LPSE: 'When lw.solver=combined, both SRS and TPD must be enabled or neither')" + ) + self.sources_on = self.tpd_enabled + self.pump_depletion = bool(light_cfg.get("pump_depletion", False)) + + self.nx, self.ny = int(grid["nx"]), int(grid["ny"]) + self.dt = grid["dt"] + self.n_sub = int(grid.get("light_substeps", 1)) + self.dt_l = self.dt / self.n_sub + self.x = grid["x"] + self.y = grid["y"] + self.dx = grid["dx"] + self.kx = jnp.asarray(grid["kx"]) + self.ky = jnp.asarray(grid["ky"]) + k_sq_np = np.asarray(grid["kx"])[:, None] ** 2 + np.asarray(grid["ky"])[None, :] ** 2 + self.k_sq = jnp.asarray(k_sq_np) + self.one_over_k_sq = jnp.asarray(np.where(k_sq_np > 0, 1.0 / np.where(k_sq_np > 0, k_sq_np, 1.0), 0.0)) + band = np.asarray(grid["low_pass_filter_grid"]) * np.where(k_sq_np > 0, 1.0, 0.0) + cap = light_cfg.get("max_wavenumber") + if cap is not None: + band = band * np.where(np.sqrt(k_sq_np) < float(cap) * derived["w0"] / derived["c"], 1.0, 0.0) + self.band = jnp.asarray(band) + + self.c = derived["c"] + self.w0 = derived["w0"] + self.wp0 = derived["wp0"] + self.w1 = derived["w1"] + self.vte_sq = derived["vte_sq"] + self.e = derived["e"] + self.me = derived["me"] + self.envelope_density = cfg["units"]["envelope density"] + self.background_density = grid["background_density"] + self.n_over_env = self.background_density / self.envelope_density + self.delta_w = self.w0 - 2.0 * self.wp0 + + # propagators over one sub-step + self.landau_enabled = bool(cfg["terms"]["epw"]["damping"].get("landau", True)) + self.hpe_enabled = bool(cfg["terms"].get("hpe", {}).get("active", False)) + self.landau_rate = analytic_landau_rate(cfg) + self.disp_L = jnp.exp(-1j * self.dt_l * 1.5 * self.vte_sq / self.wp0 * self.k_sq) * self.band + self.prop_T = jnp.exp(-1j * self.dt_l * self.c**2 / (2.0 * self.wp0) * self.k_sq) * self.band + self.detune = jnp.exp(-1j * self.dt_l * self.wp0 / 2.0 * (self.n_over_env - 1.0)) + # the Raman light's absorption at the combined carrier and the field's own critical density + # (n_env), not the EPW's collisions + self.nu_raman = combined_collision_rate(cfg) + self.collisional = jnp.exp(-self.nu_raman * self.dt_l * self.n_over_env**2) + self.boundary = grid["absorbing_boundaries"] ** (1.0 / self.n_sub) + # the pump is LPSE's laser class: its own (5e3/ps) absorber, not the EPW one -- at the EPW + # rate the injected pump reflects off both walls into a standing wave (see helpers) + self.light_boundary = grid["light_absorbing_boundaries"] ** (1.0 / self.n_sub) + + # unified source coefficient (LightSolver.cpp Sc_srs for the Raman class = (q/m)/(4 W0)) + self.source_coeff = -1j * self.e / (4.0 * self.me * self.w0) + self.rho_factor = 1.0 - self.w0 / self.wp0 + # k-filter on the pump entering the source (LPSE lw.kFilter, off by default there; + # terms.epw.source.srs_k_filter / srs_k_filter_scale, see SpectralEPWSolver) + if bool(source_cfg.get("srs_k_filter", True)): + n_min = float(np.min(np.asarray(self.background_density))) + scale = float(source_cfg.get("srs_k_filter_scale", 1.2)) + max_k0_sq = scale**2 * (self.w0 / self.c) ** 2 * max(1.0 - n_min, 0.0) + self.E0_filter = jnp.asarray(np.where(k_sq_np > max_k0_sq, 0.0, 1.0)) + else: + self.E0_filter = jnp.ones_like(self.k_sq) + + # noise on the longitudinal part, in field form: E_k += -i k kick e^{i theta} + self.noise_enabled = bool(source_cfg.get("noise", False)) + if self.noise_enabled: + import jax + + # built for the light sub-step at which add_noise applies it, not the EPW step + dt_noise, nu_noise = noise_step(cfg) + self.noise_kick = jnp.asarray(noise_kick_spectrum(cfg, dt=dt_noise, nu_coll=nu_noise)) + seed = source_cfg.get("noise_seed") + self.noise_key = jax.random.PRNGKey(int(seed) if seed is not None else np.random.randint(2**20)) + + # Raman seed (drivers.E1): the spectral light solver's smooth injector on the transverse + # (y) component of the combined field, at the combined carrier wp0 -- LPSE's Raman class + # injects at W1 = wpe in combined mode (LightSolver.cpp:2404-2406) + self.seed_enabled = "E1" in cfg["drivers"] + if self.seed_enabled: + seed = cfg["drivers"]["E1"]["derived"] + x_inject = grid["xmax"] - seed["offset"] + self.i1 = int(np.argmin(np.abs(np.asarray(self.x) - x_inject))) + self.wpe_sq_i1 = float(self.wp0**2 * self.n_over_env[self.i1, 0]) + if self.wpe_sq_i1 >= self.wp0**2: + raise ValueError( + f"The Raman seed injector at x = {float(self.x[self.i1]):.2f} um sits at density " + f"{float(self.background_density[self.i1, 0]):.3f} nc, at or above the envelope density " + f"{self.envelope_density} nc where the combined field's light is evanescent. Move the " + "injector with drivers.E1.offset, or remove drivers.E1 to run noise-seeded." + ) + k_seed = float(np.sqrt(self.wp0**2 - self.wpe_sq_i1) / self.c) + width = seed.get("injector_width", np.pi / k_seed) + self.seed_profile = jnp.asarray( + gaussian_injector_profile(np.asarray(self.x), float(self.x[self.i1]), width, self.dx) + ) + + # evolved pump (spectral) with the unified depletion term + if self.pump_depletion: + self.E0_source = derived["E0_source"] + self.linear_coeff0 = 1j * self.w0 / 2.0 * (1.0 - self.wp0**2 / self.w0**2 * self.n_over_env) + self.detune0 = jnp.exp(self.dt_l * self.linear_coeff0) + self.prop0 = jnp.exp(-1j * self.dt_l * self.c**2 / (2.0 * self.w0) * self.k_sq) * self.band + self.depletion_coeff = 1j * self.e / (2.0 * self.me * self.w0) + self.tpd_projection = bool(light_cfg.get("tpd_projection", True)) + rate0, _ = light_absorption_rates(cfg) + self.absorption_rate0 = rate0 + self.absorb0 = None if rate0 is None else jnp.exp(-rate0 * self.dt_l * self.background_density**2) + pump = cfg["drivers"]["E0"]["derived"] + x_inject = grid["xmin"] + pump["offset"] + self.i0 = int(np.argmin(np.abs(np.asarray(self.x) - x_inject))) + self.n_src = float(self.background_density[self.i0, 0]) + if self.n_src >= 1.0: + raise ValueError("The pump injector sits at or above critical density") + self.pump_turn_on_time = pump["turn_on_time"] + # the same injector as the spectral light solver (beams, profile, KAP, pulse table) + self.init_pump_injector(cfg) + + # ------------------------------------------------------------- helpers -- + + def potential(self, E1: Array) -> Array: + return potential_from_field(E1, self.kx, self.ky, self.one_over_k_sq, self.band) + + def transverse(self, E1: Array) -> Array: + return transverse_part(E1, self.kx, self.ky, self.one_over_k_sq) + + def unified_source(self, t: float, E0: Array, E1: Array) -> Array: + """``-i e/(4 me w0) e^{-i dw t} [grad(E0 . E1*) + (1 - w0/wp0) E0 (div E1)*]`` in x-space.""" + E0f = jnp.fft.ifft2(jnp.fft.fft2(E0, axes=(0, 1)) * self.E0_filter[..., None], axes=(0, 1)) + scalar_k = jnp.fft.fft2(E0f[..., 0] * jnp.conj(E1[..., 0]) + E0f[..., 1] * jnp.conj(E1[..., 1])) * self.band + grad_x = jnp.fft.ifft2(1j * self.kx[:, None] * scalar_k) + grad_y = jnp.fft.ifft2(1j * self.ky[None, :] * scalar_k) + e1x_k, e1y_k = jnp.fft.fft2(E1[..., 0]), jnp.fft.fft2(E1[..., 1]) + rho = jnp.fft.ifft2(1j * (self.kx[:, None] * e1x_k + self.ky[None, :] * e1y_k) * self.band) + term = jnp.stack([grad_x, grad_y], axis=-1) + self.rho_factor * E0f * jnp.conj(rho)[..., None] + return self.source_coeff * jnp.exp(-1j * self.delta_w * t) * term + + def unified_depletion(self, t: float, E1: Array) -> Array: + """``i e/(2 me w0) e^{+i dw t} [E1 div E1]`` (optionally transverse-projected).""" + e1x_k, e1y_k = jnp.fft.fft2(E1[..., 0]), jnp.fft.fft2(E1[..., 1]) + rho = jnp.fft.ifft2(1j * (self.kx[:, None] * e1x_k + self.ky[None, :] * e1y_k) * self.band) + f = E1 * rho[..., None] + if self.tpd_projection: + _, (tx_k, ty_k) = longitudinal_transverse(f, self.kx, self.ky, self.one_over_k_sq) + f = jnp.stack([jnp.fft.ifft2(tx_k), jnp.fft.ifft2(ty_k)], axis=-1) + return self.depletion_coeff * jnp.exp(1j * self.delta_w * t) * f + + def propagate_combined(self, E1: Array, gamma_landau) -> Array: + (lx_k, ly_k), (tx_k, ty_k) = longitudinal_transverse(E1, self.kx, self.ky, self.one_over_k_sq) + exp_l = self.disp_L * jnp.exp(-gamma_landau * self.dt_l) + fx_k = exp_l * lx_k + self.prop_T * tx_k + fy_k = exp_l * ly_k + self.prop_T * ty_k + return jnp.stack([jnp.fft.ifft2(fx_k), jnp.fft.ifft2(fy_k)], axis=-1) + + def add_noise(self, t: float, E1: Array) -> Array: + import jax + + step = jnp.round(t / self.dt_l).astype(int) + key = jax.random.fold_in(self.noise_key, step) + phases = 2.0 * np.pi * jax.random.uniform(key, (self.nx, self.ny)) + phi_kick = self.noise_kick * jnp.exp(1j * phases) + ex_k = jnp.fft.fft2(E1[..., 0]) - 1j * self.kx[:, None] * phi_kick + ey_k = jnp.fft.fft2(E1[..., 1]) - 1j * self.ky[None, :] * phi_kick + return jnp.stack([jnp.fft.ifft2(ex_k), jnp.fft.ifft2(ey_k)], axis=-1) + + def calc_seed_source(self, t: float, seed_args: dict) -> Array: + """The leftward Raman seed (``drivers.E1``) at the combined carrier ``wp0``: the (nx, ny) + source added to ``E1_y``.""" + return smooth_seed_source( + t, + seed_args, + w=self.wp0, + c=self.c, + x=self.x, + y=self.y, + x_inject=self.x[self.i1], + wpe_sq=self.wpe_sq_i1, + profile=self.seed_profile, + ) + + def propagate_pump(self, E0: Array) -> Array: + (lx_k, ly_k), (tx_k, ty_k) = longitudinal_transverse(E0, self.kx, self.ky, self.one_over_k_sq) + fx_k = lx_k + self.prop0 * tx_k + fy_k = ly_k + self.prop0 * ty_k + return jnp.stack([jnp.fft.ifft2(fx_k), jnp.fft.ifft2(fy_k)], axis=-1) + + # ---------------------------------------------------------------- step -- + + def __call__(self, t: float, y: dict, driver_args: dict, E0_fn=None) -> tuple[Array, Array, Array]: + """Advance over one EPW step. Returns (E0, E1, phi_k). + + ``E0_fn(t)`` gives the prescribed pump when it is not evolved; ``driver_args["E0"]`` + the injector parameters when it is.""" + E1 = y["E1"] + E0 = y["E0"] + iaw_density = y.get("iaw_density") + detune = self.detune + collisional = self.collisional + if iaw_density is not None: + # the IAW perturbation in units of n_env, as the detuning: absorption at the effective density + detune = detune * jnp.exp(-1j * self.dt_l * self.wp0 / 2.0 * iaw_density) + if self.nu_raman: + collisional = jnp.exp(-self.nu_raman * self.dt_l * (self.n_over_env + iaw_density) ** 2) + if self.hpe_enabled: + gamma_landau = y["gamma_L"] + elif self.landau_enabled: + gamma_landau = self.landau_rate + else: + gamma_landau = jnp.zeros_like(self.k_sq) + pump_args = driver_args.get("E0") if self.pump_depletion else None + seed_args = driver_args["E1"] if self.seed_enabled else None + if self.pump_depletion: + detune0 = self.detune0 + absorb0 = self.absorb0 + if iaw_density is not None: + detune0 = detune0 * jnp.exp(-1j * self.wp0**2 / (2.0 * self.w0) * iaw_density * self.dt_l) + if self.absorb0 is not None: + n0 = self.background_density * (1.0 + iaw_density / self.n_over_env) + absorb0 = jnp.exp(-self.absorption_rate0 * self.dt_l * n0**2) + + def substep(i, fields): + E0, E1 = fields + t_i = t + i * self.dt_l + if not self.pump_depletion: + E0 = E0_fn(t_i) + # 1. scattering potential and collisional damping on the combined field + E1 = E1 * (detune * collisional)[..., None] + # 2. the unified source (and the pump's x-space step) + if self.pump_depletion: + E0 = E0 * detune0[..., None] + if absorb0 is not None: + E0 = E0 * absorb0[..., None] + if self.sources_on: + E0 = E0 + self.dt_l * self.unified_depletion(t_i, E1) + E0 = E0 + self.dt_l * self.calc_pump_source(t_i, pump_args) + if self.sources_on: + E1 = E1 + self.dt_l * self.unified_source(t_i, E0, E1) + if seed_args is not None: + E1 = E1.at[..., 1].add(self.dt_l * self.calc_seed_source(t_i, seed_args)) + # 3. k-space propagation with the L/T projector + E1 = self.propagate_combined(E1, gamma_landau) + if self.pump_depletion: + E0 = self.propagate_pump(E0) + # 4. noise on the longitudinal part + if self.noise_enabled: + E1 = self.add_noise(t_i, E1) + # 5. absorbing layers + E1 = E1 * self.boundary[..., None] + if self.pump_depletion: + E0 = E0 * self.light_boundary[..., None] + return (E0, E1) + + E0, E1 = lax.fori_loop(0, self.n_sub, substep, (E0, E1)) + return E0, E1, self.potential(E1) diff --git a/adept/_lpse2d/core/epw.py b/adept/_lpse2d/core/epw.py index e6dce7ee..da32e01d 100644 --- a/adept/_lpse2d/core/epw.py +++ b/adept/_lpse2d/core/epw.py @@ -5,36 +5,285 @@ from adept._lpse2d.core.driver import Driver - -def landau_damping_rate(k_sq: Array, wp0: float, vte_sq: float, zero_mask: Array) -> Array: +LANDAU_FORMS = ("matlab", "lpse", "relativistic") +NOISE_MODELS = ("flat", "thermal") +TPD_FORMS = ("lpse", "matlab") + +# EPW energy ledger (terms.epw.energy_ledger): cumulative change of the EPW energy +# measure sum_k k^2 |phi_k|^2 attributed to each operation of the split step, in this +# order. LPSE asserts the same closure to 0.1 % per step (ZakharovSolver.cpp:1700-1712). +LEDGER_KEY = "epw_ledger" +LEDGER_CHANNELS = ( + "dispersion", # k-space Bohm-Gross phase (exactly 0) + "damping", # Landau + collisional (<= 0) + "dealias", # band mask after the damping (<= 0) + "noise", # random-phase source (>= 0 on average) + "detuning", # density / IAW phase in x-space (exactly 0) + "boundary", # absorbing layers (<= 0) + "reprojection", # E -> phi_k round trip with the band mask (<= 0) + "tpd", # TPD source + "srs", # SRS source + "driver", # direct EPW driver (drivers.E2) +) + + +def landau_damping_rate(k_sq: Array, wp0: float, vte_sq: float, zero_mask: Array, form: str = "matlab") -> Array: """ - Landau damping rate for each k mode (amplitude rate, 1/ps). + Non-relativistic Landau damping rate for each k mode (amplitude rate, 1/ps). + + ``form="matlab"`` is m201805_matlabLpse_v11.m line 913: + + gammaLandauEpw = sqrt(pi/8) * (1 + 3/2*k^2*vte^2/wp^2) * wp^4/(k^3*vte^3) * exp(...) + + ``form="lpse"`` is the LPSE C++ expression (ZakharovSolver.cpp landauDamping_nonRel): - Matches MATLAB line 913: - gammaLandauEpw = sqrt(pi/8) * (1 + 3/2*k^2*vte^2/wp^2) * wp^4/(k^3*vte^3) * exp(...) + gamma = sqrt(pi/8) (kde/k)^3 w_k exp(-w_k^2 / (2 k^2 vte^2)), w_k = wp sqrt(1 + 3 k^2/kde^2) - Module-level so the solver (`SpectralEPWSolver`) and the dissipation - diagnostic (`helpers.get_default_save_func`) use the *same* rates and can - never drift apart. + The exponents are identical (w_k^2/(2 k^2 vte^2) = wp^2/(2 k^2 vte^2) + 3/2 exactly); the + prefactors differ by sqrt(1 + 3 x^2) vs (1 + 1.5 x^2), x = k lambda_D -- 0.7% at x = 0.3, + 3.9% at x = 0.5. + + Module-level so the solver (`SpectralEPWSolver`), the HPE calibration 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))) - ) + if form == "lpse": + kde_sq = wp0**2 / vte_sq + w_k = wp0 * jnp.sqrt(1.0 + 3.0 * k_sq / kde_sq) + damping = ( + jnp.sqrt(np.pi / 8.0) * (kde_sq / k_sq_safe) ** 1.5 * w_k * jnp.exp(-0.5 * w_k**2 / (k_sq_safe * vte_sq)) + ) + elif form == "matlab": + 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))) + ) + else: + raise ValueError(f"landau_damping_rate: form must be 'matlab' or 'lpse', got {form!r}") return damping * zero_mask +def landau_damping_rate_relativistic( + k_sq: np.ndarray, wp0: float, vte_sq: float, c: float, ndim: int = 2 +) -> np.ndarray: + """ + Relativistic (Maxwell-Juettner) Landau damping rate, ported from LPSE + ``ZakharovSolver::landauDamping_rel_2D`` / ``landauDamping_rel_3D`` (Bessel-function + forms after Vu / Russell). ``ndim`` selects the 2-D or 3-D velocity-space integral; + modes whose phase velocity exceeds c (``arg <= 0``) are undamped, as in LPSE. + + Evaluated with numpy/scipy once at setup (the rate is a static k-array). + """ + from scipy import special + + k_sq = np.asarray(k_sq, dtype=np.float64) + vte = np.sqrt(vte_sq) + kde = wp0 / vte + k = np.sqrt(np.where(k_sq > 0.0, k_sq, 1.0)) + x = k / kde + beta_sq = vte_sq / c**2 # (Ve/C)^2 + omega = 1.0 + 1.5 * x**2 + arg = 1.0 - beta_sq * omega**2 / x**2 + ok = (k_sq > 0.0) & (arg > 0.0) + A = 1.0 / np.sqrt(np.where(ok, arg, 1.0)) + a = A / beta_sq # (C/Ve)^2 A + # e_bessel_kN in LPSE's Bessel.h are the exponentially scaled K_N: e^a K_N(a) + if ndim == 2: + gamma = ( + 0.25 + * (kde / k) ** 3 + * wp0 + * omega + * np.exp(1.0 / beta_sq - a) + * (1.0 / np.sqrt(beta_sq)) + * A**3 + * 2.0 + * (special.k0e(a) + special.k1e(a) / a) + / (1.0 + beta_sq) + ) + elif ndim == 3: + b = 1.0 / beta_sq + gamma = ( + 0.25 + * np.pi + * (kde / k) ** 3 + * wp0 + * omega + * np.exp(b - a) + * b ** (-2.5) + * (2.0 + a * (2.0 + a)) + / (special.k0e(b) + 2.0 * special.k1e(b) / b) + ) + else: + raise ValueError("landau_damping_rate_relativistic: ndim must be 2 or 3") + return np.where(ok, gamma, 0.0) + + +def analytic_landau_rate(cfg: dict) -> Array: + """ + The static Landau damping rate array (nx, ny) selected by ``terms.epw.damping``: + + - ``landau_form``: ``matlab`` (default), ``lpse`` or ``relativistic`` (LPSE 2-D Bessel form; + ``relativistic_3d`` for the 3-D velocity-space integral) + - ``landau_lower_threshold`` (1/ps): modes damped below it are treated as undamped + (LPSE ``lw.landauDamping.lowerThreshold``) + - ``landau_multiplier``: static scalar on the rate (LPSE ``LD_multiplier``) + + Returns zeros when ``terms.epw.damping.landau`` is off. + """ + derived = cfg["units"]["derived"] + damping = cfg["terms"]["epw"]["damping"] + kx = jnp.asarray(cfg["grid"]["kx"]) + ky = jnp.asarray(cfg["grid"]["ky"]) + k_sq = kx[:, None] ** 2 + ky[None, :] ** 2 + zero_mask = jnp.where(k_sq > 0, 1.0, 0.0) + if not damping.get("landau", True): + return jnp.zeros_like(k_sq) + form = str(damping.get("landau_form", "matlab")) + if form in ("relativistic", "relativistic_2d"): + rate = jnp.asarray( + landau_damping_rate_relativistic(np.asarray(k_sq), derived["wp0"], derived["vte_sq"], derived["c"], ndim=2) + ) + elif form == "relativistic_3d": + rate = jnp.asarray( + landau_damping_rate_relativistic(np.asarray(k_sq), derived["wp0"], derived["vte_sq"], derived["c"], ndim=3) + ) + elif form in ("matlab", "lpse"): + rate = landau_damping_rate(k_sq, derived["wp0"], derived["vte_sq"], zero_mask, form=form) + else: + raise ValueError( + f"terms.epw.damping.landau_form must be one of matlab, lpse, relativistic, relativistic_3d; got {form!r}" + ) + threshold = float(damping.get("landau_lower_threshold", 0.0) or 0.0) + if threshold > 0.0: + rate = jnp.where(rate > threshold, rate, 0.0) + multiplier = float(damping.get("landau_multiplier", 1.0) if damping.get("landau_multiplier") is not None else 1.0) + return rate * multiplier + + +def noise_step(cfg: dict) -> tuple[float, float]: + """ + The time step at which the active EPW solver applies the noise kick, and the collisional + rate that damps the noised field over it: the EPW step and ``terms.epw.damping.collisions`` + for the separate solver; the light sub-step ``dt / grid.light_substeps`` and the Raman light's + absorption (``terms.light.absorption`` at the combined carrier, which damps the combined field) + for the combined solver, which adds the kick once per light step (LPSE + ``LwSolver::addNoise_combinedSolver_spectral(dt)``). ``noise_kick_spectrum`` and + ``diagnostics.expected_noise_energy`` both take it, so the solver and the growth-fit noise + floor stay consistent. + """ + grid = cfg["grid"] + nu = float(cfg["units"]["derived"].get("nu_coll", 0.0)) + if cfg["terms"]["epw"].get("solver", "separate") == "combined": + from adept._lpse2d.core.raman import light_absorption_rates + + _, raman_rate = light_absorption_rates(cfg) + return grid["dt"] / int(grid.get("light_substeps", 1)), 0.0 if raman_rate is None else float(raman_rate) + return grid["dt"], nu + + +def noise_kick_spectrum(cfg: dict, dt: float | None = None, nu_coll: float | None = None) -> np.ndarray: + """ + Per-step, per-mode EPW noise kick amplitude ``D_k`` (k-space potential units), for both + noise models. ``get_noise`` multiplies it by a random phase every step; the diagnostics + (``diagnostics.expected_noise_energy``) integrate the same array, so the two never drift. + + ``noise_model: flat`` (default; the MATLAB source): ``D_k = dt * noise_amplitude`` on + every retained mode at the EPW step. A solver that kicks more often (the combined + solver, every light sub-step ``dt_l``) gets ``noise_amplitude * sqrt(dt * dt_l)``, which + keeps the injected variance per unit time -- and so the seed level -- independent of the + sub-step count. + + ``noise_model: thermal`` (LPSE ``lw.noise``, ZakharovSolver::addNoiseToPotential_fft): + a fluctuation-dissipation source balanced against the *frozen* analytic Landau + + collisional rate ``gamma_k``: + + D_k = N A / sqrt(1 + k^2 lambda_D^2) * sqrt(1 - exp(-2 gamma_k dt)) / |k| + + ``N = nx*ny`` converts an x-space amplitude to this code's unnormalized-FFT k-space + convention, and the ``1/|k|`` converts field to potential, so the steady state of + ``phi_k *= exp(-gamma_k dt); phi_k += D_k e^{i theta}`` is the Cerenkov spectrum + ``<|E_k|^2> = A^2 / (1 + k^2 lambda_D^2)`` per mode (x-space envelope amplitude + squared), independent of dt. With ``noise_calibrate: true`` the amplitude is set from + the electron temperature by equipartition -- electric energy ``kT/2`` per mode over + the box volume ``V`` (``Lz = Ly`` for the 2-D box, ``V = Lx^3`` when ny = 1, following + LPSE's ``deltaK3``): ``A_thermal = sqrt(8 pi kT / V)``, and ``noise_amplitude`` then + multiplies it (LPSE's ``isCalculated`` convention, amplitude 1 = thermal). Modes with + no damping receive no noise; a retained band with ``gamma_k <= 0`` everywhere is refused, + as in LPSE ("Cannot add LW noise without damping"). + + ``dt`` and ``nu_coll`` default to the EPW step and ``terms.epw.damping.collisions``; pass + ``noise_step(cfg)`` for the step the active solver actually uses. + """ + grid = cfg["grid"] + source = cfg["terms"]["epw"]["source"] + derived = cfg["units"]["derived"] + kx = np.asarray(grid["kx"]) + ky = np.asarray(grid["ky"]) + k_sq = kx[:, None] ** 2 + ky[None, :] ** 2 + zero_mask = np.where(k_sq > 0, 1.0, 0.0) + band = np.asarray(grid["low_pass_filter_grid"]) * zero_mask + max_k = source.get("noise_max_wavenumber") + if max_k is not None: + k0 = derived["w0"] / derived["c"] + band = band * np.where(np.sqrt(k_sq) < float(max_k) * k0, 1.0, 0.0) + + model = str(source.get("noise_model", "flat")) + amplitude = float(source.get("noise_amplitude", 1e-10)) + dt = grid["dt"] if dt is None else float(dt) + if model == "flat": + return amplitude * np.sqrt(grid["dt"] * dt) * band + if model != "thermal": + raise ValueError(f"terms.epw.source.noise_model must be 'flat' or 'thermal', got {model!r}") + + nu = derived.get("nu_coll", 0.0) if nu_coll is None else float(nu_coll) + gamma = np.asarray(analytic_landau_rate(cfg)) + nu * zero_mask + retained = band > 0.0 + if not np.any(retained): + return np.zeros_like(k_sq) + if np.all(gamma[retained] <= 0.0): + raise ValueError( + "terms.epw.source.noise_model: thermal needs EPW damping (Landau and/or collisions) on the " + "retained band -- LPSE refuses noise without damping and so does this model" + ) + if source.get("noise_calibrate", False): + # kT in this code's energy unit (massScale * spatialScale^2 / timeScale^2, cgs) + from astropy.units import Quantity as _Q + + te_kev = _Q(cfg["units"]["reference electron temperature"]).to("keV").value + energy_scale = derived["massScale"] * derived["spatialScale"] ** 2 / derived["timeScale"] ** 2 + kT = te_kev * 1.602176634e-9 / energy_scale + nx, ny = int(grid["nx"]), int(grid["ny"]) + lx = nx * grid["dx"] + ly = ny * grid["dy"] + volume = lx**3 if ny == 1 else lx * ly * ly + amplitude = amplitude * np.sqrt(8.0 * np.pi * kT / volume) + + lambda_d_sq = derived["vte_sq"] / derived["wp0"] ** 2 + n_total = float(grid["nx"] * grid["ny"]) + k_safe = np.sqrt(np.where(k_sq > 0, k_sq, 1.0)) + two_g_dt = 2.0 * gamma * dt + # LPSE switches to the small-argument form below 1e-4 to avoid a float32 cancellation; + # in float64 -expm1 is exact enough everywhere + fd_factor = np.sqrt(np.maximum(-np.expm1(-two_g_dt), 0.0)) + # LPSE carries the 1/sqrt(1 + k^2 lambda_D^2) Debye factor only on its isCalculated + # branch; noise_debye_factor: false reproduces its plain lw.noise.amplitude source + debye = np.sqrt(1.0 + k_sq * lambda_d_sq) if source.get("noise_debye_factor", True) else 1.0 + kick = n_total * amplitude / debye * fd_factor / k_safe + return np.where(retained, kick, 0.0) + + class SpectralEPWSolver: """ Spectral solver for electrostatic plasma waves in k-space. - Matches MATLAB's spectralEpwUpdate() function (lines 1966-2118). + Matches MATLAB's spectralEpwUpdate() function (lines 1966-2118), which is the same + split-step as LPSE's default ``lw.solver = spectral`` path (ZakharovSolver.cpp). State variable: phi_k (electrostatic potential in k-space) - MATLAB convention: uses fftshift, so DC is in center @@ -44,6 +293,19 @@ class SpectralEPWSolver: 1. Filter applied at exactly 2 points per timestep 2. Clear separation between operations (no combined expressions) 3. Explicit comments matching MATLAB line numbers + + LPSE-parity options (all default to the MATLAB behaviour unless stated): + + - ``terms.epw.damping.landau_form / landau_lower_threshold / landau_multiplier`` + (see ``analytic_landau_rate``) + - ``terms.epw.source.tpd_form``: ``lpse`` (default) keeps LPSE's exact coefficient + ``i e/(4 me w0)`` and the ``(w0/wp0 - 1)`` factor on the charge-density term, and uses + every pump component; ``matlab`` is the prototype's ``w0 -> 2 wp0`` form. The two are + identical at envelope density n_c/4. + - ``terms.epw.source.noise_model``: ``flat`` (default) or ``thermal`` (see + ``noise_kick_spectrum``) + - ``terms.epw.max_wavenumber``: LPSE ``lw.maxWavenumber`` (units of k0), folded into + the retained band by ``helpers.get_solver_quantities`` """ def __init__(self, cfg: dict): @@ -93,47 +355,80 @@ def __init__(self, cfg: dict): self.low_pass_filter = cfg["grid"]["low_pass_filter_grid"] # TPD parameters - self.tpd_enabled = cfg["terms"]["epw"]["source"]["tpd"] + source_cfg = cfg["terms"]["epw"]["source"] + self.tpd_enabled = source_cfg["tpd"] + self.tpd_form = str(source_cfg.get("tpd_form", "lpse")) + if self.tpd_form not in TPD_FORMS: + raise ValueError(f"terms.epw.source.tpd_form must be one of {TPD_FORMS}, got {self.tpd_form!r}") if self.tpd_enabled: - self.tpd_prefactor = 1j * self.e / (8.0 * self.wp0 * self.me) + if self.tpd_form == "lpse": + # LPSE ZakharovSolver::updatePotentialWithTpdSource_fft: + # TPD_srcFactor = i (q/m)/(4 W0) exp(-i (W0 - 2 wpe) t) + # TPD_src = TPD1 - (1 - W0/wpe) i k.TPD2 / k^2 + self.tpd_prefactor = 1j * self.e / (4.0 * self.w0 * self.me) + self.tpd_rho_factor = self.w0 / self.wp0 - 1.0 + else: + # MATLAB line 2024: w0 -> 2 wp0 inside the coefficients + self.tpd_prefactor = 1j * self.e / (8.0 * self.wp0 * self.me) + self.tpd_rho_factor = 1.0 # SRS parameters - self.srs_enabled = cfg["terms"]["epw"]["source"].get("srs", False) + self.srs_enabled = source_cfg.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 + # (MATLAB isSuppressHighKSource, lines 637-645; LPSE ``lw.kFilter``, which + # is *disabled* by default there). The cutoff ``scale * k1(n_min)`` assumes + # the scattered light sits exactly at the envelope frequency w1; in a box + # detuned from the envelope density the resonant Raman wavenumber can lie + # above it and the filter then removes the SRS mode altogether (no growth + # at all at 0.2 nc with envelope density 0.25), so translated LPSE decks + # turn it off. + self.srs_k_filter = bool(source_cfg.get("srs_k_filter", True)) + max_source_k_multiplier = float(source_cfg.get("srs_k_filter_scale", 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] + if self.srs_k_filter: + 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) + 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] + else: + self.E1_filter = jnp.ones_like(self.k_sq)[..., None] + if self.pump_depletion: + self.E0_filter = jnp.ones_like(self.k_sq)[..., 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 = float(cfg["terms"]["epw"]["source"].get("noise_amplitude", 1e-10)) - cfg_seed = cfg["terms"]["epw"]["source"].get("noise_seed") + self.noise_enabled = source_cfg["noise"] + self.noise_model = str(source_cfg.get("noise_model", "flat")) + if self.noise_model not in NOISE_MODELS: + raise ValueError(f"terms.epw.source.noise_model must be one of {NOISE_MODELS}, got {self.noise_model!r}") + self.noise_amplitude = float(source_cfg.get("noise_amplitude", 1e-10)) + cfg_seed = source_cfg.get("noise_seed") self.noise_seed = int(cfg_seed) if cfg_seed is not None else np.random.randint(2**20) # per-step keys are derived with fold_in rather than PRNGKey(step + seed): # additive seeds made nearby seeds share the same noise trajectory merely # time-shifted by a few steps, so a seed-sweep ensemble was one realization. # fold_in streams are still fully deterministic per (seed, step). self.noise_key = jax.random.PRNGKey(self.noise_seed) + if self.noise_enabled: + # per-mode kick amplitude, already including dt (flat) or the + # fluctuation-dissipation factor (thermal) and the retained-band mask + self.noise_kick = jnp.asarray(noise_kick_spectrum(cfg)) + else: + self.noise_kick = None # Density gradient self.density_gradient_enabled = cfg["terms"]["epw"]["density_gradient"] @@ -141,6 +436,8 @@ def __init__(self, cfg: dict): # Landau damping flag (previously ignored -- damping was unconditionally on) self.landau_enabled = bool(cfg["terms"]["epw"]["damping"].get("landau", True)) + # static analytic rate with the configured form / threshold / multiplier applied + self.landau_rate = analytic_landau_rate(cfg) # HPE (Follett-style particle feedback): the damping rate is read from the # state (y["gamma_L"], written by HybridParticleEvolution) instead of the # static analytic array @@ -154,15 +451,13 @@ def __init__(self, cfg: dict): def calc_landau_damping_rate(self) -> Array: """ - Calculate Landau damping rate for each k mode. - - Matches MATLAB line 913: - gammaLandauEpw = sqrt(pi/8) * (1 + 3/2*k^2*vte^2/wp^2) * wp^4/(k^3*vte^3) * exp(...) + Landau damping rate for each k mode with the configured form, threshold and + multiplier (``analytic_landau_rate``). Returns: Landau damping rate array (shape: nx, ny) """ - return landau_damping_rate(self.k_sq, self.wp0, self.vte_sq, self.zero_mask) + return self.landau_rate def phi_k_to_e_fields(self, phi_k: Array) -> tuple[Array, Array]: """ @@ -239,43 +534,42 @@ def e_fields_to_phi_k(self, ex: Array, ey: Array) -> Array: return phi_k - def calc_tpd_source(self, t: float, phi_k: Array, ey: Array, E0_y: Array) -> Array: + def calc_tpd_source(self, t: float, phi_k: Array, ex: Array, ey: Array, E0: Array) -> Array: """ - Calculate Two Plasmon Decay source term. + Calculate the Two Plasmon Decay source term. + + LPSE (ZakharovSolver.cpp, makeTpdSource / updatePotentialWithTpdSource_fft): - Matches MATLAB lines 1996-2049 for isSolveForPotential=true. + S_k = i (q/m)/(4 w0) e^{-i (w0 - 2 wp0) t} [ F(E0 . E*) - (1 - w0/wp0) i k . F(E0 rho*) / k^2 ] - TPD source has two components: - TPD1 = FFT(E0_y * conj(Ey)) - TPD2 = 1i * KY/K_sq * FFT(E0_y * conj(divE_true)) - where divE_true = IFFT(K_sq * phi_k) + with rho = div E = F^-1(k^2 phi_k) and every pump component in the dot product and + the vector term. MATLAB lines 1996-2049 are the ``tpd_form: matlab`` special case + (w0 -> 2 wp0 in the coefficients, y-polarized pump only in the prototype; here + the x component is always carried, which is exact for the prototype's E0x = 0). Args: t: Current time phi_k: Potential in k-space - ey: Electric field y-component in real space - E0_y: Laser field y-component in real space + ex, ey: Electric field components in real space + E0: Pump field in real space, shape (nx, ny, 2) Returns: TPD source term in k-space """ - # Component 1: E0 * conj(Ey) - # MATLAB line 2011-2012 - product1 = E0_y * jnp.conj(ey) - tpd1 = jnp.fft.fft2(product1) - - # Component 2: E0 * conj(divE_true) - # MATLAB line 2014-2018 - # divE_true is the actual charge density (4π times Poisson) - div_e_true = jnp.fft.ifft2(self.k_sq * phi_k) - product2 = E0_y * jnp.conj(div_e_true) - product2_k = jnp.fft.fft2(product2) - tpd2 = 1j * self.ky[None, :] * self.one_over_k_sq * product2_k - - # Combine with prefactor - # MATLAB line 2024 + e0x, e0y = E0[..., 0], E0[..., 1] + + # Component 1: F(E0 . conj(E)) (MATLAB line 2011-2012 with E0x = 0) + tpd1 = jnp.fft.fft2(e0x * jnp.conj(ex) + e0y * jnp.conj(ey)) + + # Component 2: i k . F(E0 conj(rho)) / k^2 (MATLAB line 2014-2018) + rho = jnp.fft.ifft2(self.k_sq * phi_k) + tpd2_x = jnp.fft.fft2(e0x * jnp.conj(rho)) + tpd2_y = jnp.fft.fft2(e0y * jnp.conj(rho)) + tpd2 = 1j * (self.kx[:, None] * tpd2_x + self.ky[None, :] * tpd2_y) * self.one_over_k_sq + + # Combine with prefactor (MATLAB line 2024; LPSE TPD_srcFactor) phase = jnp.exp(-1j * (self.w0 - 2.0 * self.wp0) * t) - source = self.tpd_prefactor * phase * (tpd1 + tpd2) + source = self.tpd_prefactor * phase * (tpd1 + self.tpd_rho_factor * tpd2) # Apply filter to source (MATLAB line 2032-2033) source = source * self.low_pass_filter @@ -316,7 +610,10 @@ def calc_srs_source(self, E0: Array, E1: Array) -> Array: def get_noise(self, t: float) -> Array: """ - Generate random noise for plasma waves. + Generate the random-phase noise kick for the plasma waves at this step. + + The returned array is added to phi_k directly: it already carries dt (flat model) + or the fluctuation-dissipation factor (thermal model) from ``noise_kick_spectrum``. Args: t: Current time @@ -332,20 +629,26 @@ def get_noise(self, t: float) -> Array: # Random phases phases = 2.0 * np.pi * jax.random.uniform(key, (self.nx, self.ny)) - # Uniform amplitude with random phase - noise = self.noise_amplitude * jnp.exp(1j * phases) + # per-mode amplitude with random phase; the kick spectrum is already masked to + # the retained band (MATLAB epwNoise: phi_noise(isHighWavenumberMode) = 0) and k = 0 + return self.noise_kick * jnp.exp(1j * phases) - # Suppress high-wavenumber modes (MATLAB epwNoise: phi_noise(isHighWavenumberMode) = 0) - noise = noise * self.low_pass_filter + def energy(self, phi_k: Array) -> Array: + """The ledger's energy measure sum_k k^2 |phi_k|^2 (= N sum_x |E|^2, Parseval).""" + return jnp.sum(self.k_sq * jnp.abs(phi_k) ** 2) - # Zero out k=0 - noise = noise * self.zero_mask - - return noise + def energy_x(self, ex: Array, ey: Array) -> Array: + """The same measure evaluated on the real-space fields.""" + return (self.nx * self.ny) * jnp.sum(jnp.abs(ex) ** 2 + jnp.abs(ey) ** 2) def __call__(self, t: float, y, args) -> Array: + """Advance EPW by one timestep; see ``advance``.""" + return self.advance(t, y, args)[0] + + def advance(self, t: float, y, args) -> tuple[Array, Array]: """ - Advance EPW by one timestep using spectral method. + Advance EPW by one timestep using spectral method, returning the new potential + and the energy change of each operation (``LEDGER_CHANNELS`` order). This matches MATLAB's spectralEpwUpdate() lines 1966-2118. @@ -375,6 +678,13 @@ def __call__(self, t: float, y, args) -> Array: phi_k = y["epw"] E0 = y["E0"] background_density = self.background_density + deltas = {} + w_prev = self.energy(phi_k) + + def book(name, w_now): + nonlocal w_prev + deltas[name] = w_now - w_prev + w_prev = w_now # ======================================================================== # STEP 1-2: Thermal dispersion and Landau damping @@ -382,36 +692,42 @@ def __call__(self, t: float, y, args) -> Array: # MATLAB line 1975: divE_k = divE_k .* exp(-1i*3/2*vte_sq/wp0 .* K_sq * DT) thermal_phase = jnp.exp(-1j * 1.5 * self.vte_sq / self.wp0 * self.k_sq * self.dt) phi_k = phi_k * thermal_phase + book("dispersion", self.energy(phi_k)) # MATLAB line 1981: divE = divE .* exp(-(gammaLandau + nu_coll) * DT) if self.hpe_enabled: gamma_landau = y["gamma_L"] elif self.landau_enabled: - gamma_landau = self.calc_landau_damping_rate() + gamma_landau = self.landau_rate else: gamma_landau = 0.0 damping_factor = jnp.exp(-(gamma_landau + self.nu_coll) * self.dt) phi_k = phi_k * damping_factor + book("damping", self.energy(phi_k)) # ======================================================================== # STEP 3: Apply filter ONCE after thermal + damping # ======================================================================== # MATLAB line 1976: divE_k(isHighWavenumberMode) = 0 phi_k = phi_k * self.low_pass_filter + book("dealias", self.energy(phi_k)) # ======================================================================== - # STEP 4: Add noise + # STEP 4: Add noise (after the damping sub-step, as LPSE does for the + # energetics of the fluctuation-dissipation balance) # ======================================================================== if self.noise_enabled: - # MATLAB line 1988: divE = divE + epwNoise * DT - noise = self.get_noise(t) - phi_k = phi_k + self.dt * noise + # MATLAB line 1988: divE = divE + epwNoise * DT (dt is inside noise_kick) + phi_k = phi_k + self.get_noise(t) + book("noise", self.energy(phi_k)) # ======================================================================== # STEP 5: Calculate electric fields # ======================================================================== # MATLAB line 1992: [Ex, Ey] = calculateFieldsFromDivE(...) ex, ey = self.phi_k_to_e_fields(phi_k) + # from here to the re-projection the ledger follows the x-space fields + w_prev = self.energy_x(ex, ey) # ======================================================================== # STEP 6: Calculate TPD source (in k-space, before applying density gradient) @@ -419,8 +735,7 @@ def __call__(self, t: float, y, args) -> Array: tpd_source = None if self.tpd_enabled: # MATLAB lines 1996-2035 - E0_y = E0[..., 1] # y-component of laser field - tpd_source = self.calc_tpd_source(t, phi_k, ey, E0_y) + tpd_source = self.calc_tpd_source(t, phi_k, ex, ey, E0) srs_source = None if self.srs_enabled: @@ -442,6 +757,7 @@ def __call__(self, t: float, y, args) -> Array: density_phase = jnp.exp(-1j * self.wp0 / 2.0 * density_perturbation * self.dt) ex = ex * density_phase ey = ey * density_phase + book("detuning", self.energy_x(ex, ey)) # ======================================================================== # STEP 8: Apply absorbing boundaries to E fields (in REAL space) @@ -451,6 +767,7 @@ def __call__(self, t: float, y, args) -> Array: # Ey = Ey .* exp(-DT * boundaryDampingRate) ex = ex * self.boundary_envelope ey = ey * self.boundary_envelope + book("boundary", self.energy_x(ex, ey)) # ======================================================================== # STEP 9: Convert E fields back to phi_k @@ -458,6 +775,7 @@ def __call__(self, t: float, y, args) -> Array: # MATLAB line 2103: divE = convertFieldsToDivE(Ex, Ey, ...) # This function applies filter at line 2523 phi_k = self.e_fields_to_phi_k(ex, ey) + book("reprojection", self.energy(phi_k)) # ======================================================================== # STEP 10: Add TPD source @@ -465,6 +783,7 @@ def __call__(self, t: float, y, args) -> Array: if self.tpd_enabled and tpd_source is not None: # MATLAB line 2109: divE = divE + tpdSourceTerm * DT phi_k = phi_k + self.dt * tpd_source + book("tpd", self.energy(phi_k)) # ======================================================================== # STEP 11: Add SRS source @@ -472,5 +791,7 @@ def __call__(self, t: float, y, args) -> Array: 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 + book("srs", self.energy(phi_k)) + deltas["driver"] = jnp.zeros(()) # filled by the split step (drivers.E2 acts before this call) - return phi_k + return phi_k, jnp.stack([deltas[name] for name in LEDGER_CHANNELS]) diff --git a/adept/_lpse2d/core/hpe.py b/adept/_lpse2d/core/hpe.py index 475ec2f2..32d8464e 100644 --- a/adept/_lpse2d/core/hpe.py +++ b/adept/_lpse2d/core/hpe.py @@ -24,9 +24,17 @@ from jax.scipy import special as jax_special from scipy import special -from adept._lpse2d.core.epw import landau_damping_rate +from adept._lpse2d.core.epw import analytic_landau_rate -PARTICLE_KEYS = ("x_e", "y_e", "u_e", "epw_hist", "gamma_L") +PARTICLE_KEYS = ("x_e", "y_e", "u_e", "epw_hist", "gamma_L", "hpe_wall_flux", "hpe_cone_energy", "hpe_ld_multiplier") +WALLS = ("left", "right", "bottom", "top") +DEFAULT_FLUX_BINS_KEV = (0.0, 50.0, 100.0, 1.0e9) + + +def flux_bin_edges(hpe: dict) -> np.ndarray: + """keV edges of the wall-flux instrument (LPSE hpe.metrics.flux energy.min/max per metric).""" + edges = hpe.get("flux_bins") or DEFAULT_FLUX_BINS_KEV + return np.asarray(edges, dtype=np.float64) def _hpe_cfg(cfg: dict) -> dict: @@ -34,15 +42,8 @@ def _hpe_cfg(cfg: dict) -> dict: def _analytic_rate(cfg: dict) -> np.ndarray: - derived = cfg["units"]["derived"] - kx = np.asarray(cfg["grid"]["kx"]) - ky = np.asarray(cfg["grid"]["ky"]) - k_sq = kx[:, None] ** 2 + ky[None, :] ** 2 - zero_mask = np.where(k_sq > 0.0, 1.0, 0.0) - return np.asarray( - landau_damping_rate(jnp.asarray(k_sq), derived["wp0"], derived["vte_sq"], jnp.asarray(zero_mask)), - dtype=np.float64, - ) + # the same static rate (form, threshold, multiplier) the fluid solver would apply + return np.asarray(analytic_landau_rate(cfg), dtype=np.float64) def resonance_arrays(cfg: dict) -> dict: @@ -201,6 +202,13 @@ def load_particles(cfg: dict) -> dict: "u_e": np.asarray(u, dtype=np.float64), "epw_hist": hist, "gamma_L": arrays["gamma_analytic"].copy(), + # instruments and controls (LPSE hpe.metrics.flux / power, hpe.enforceEnergyConservation): + # cumulative keV (per test particle) leaving through each wall per energy bin, the + # cumulative keV leaving inside the acceptance cone, and the energy-conservation + # multiplier on the applied Landau rate + "hpe_wall_flux": np.zeros((len(WALLS), len(flux_bin_edges(hpe)) - 1), dtype=np.float64), + "hpe_cone_energy": np.zeros((1,), dtype=np.float64), + "hpe_ld_multiplier": np.ones((1,), dtype=np.float64), } if is_2d: state["y_e"] = np.asarray(y, dtype=np.float64) @@ -261,6 +269,53 @@ def __init__(self, cfg: dict): self.mask_res = jnp.asarray(arrays["mask_res"]) self.gamma_analytic = jnp.asarray(arrays["gamma_analytic"]) self.f_tail_frac = arrays["f_tail_frac"] + # ---- LPSE HPE controls (hpe.gammaLimit.*, allowGrowth, thermalizationProbability, + # magneticField, enforceEnergyConservation, numStepsToAverageEnergyChange, metrics) + self.gamma_limit_damping = float(hpe.get("gamma_limit_damping", 1500.0)) # 1/ps + self.gamma_limit_growth = ( + float(hpe.get("gamma_limit_growth", 1500.0)) if bool(hpe.get("allow_growth", False)) else 0.0 + ) + # LPSE particle walls (ElectronTracker.cu): every crossing of a box side is counted by the + # instruments, then the particle is thermalized (re-injected from the tail) with the + # per-direction probability or else passes through periodically. The default keeps + # adept's earlier behaviour: thermalize at absorbing field boundaries, wrap at periodic ones. + p_default = [0.0 if self.periodic_x else 1.0, 0.0 if self.periodic_y else 1.0] + p_therm = list(hpe.get("thermalization_probability") or p_default) + self.p_therm_x, self.p_therm_y = float(p_therm[0]), float(p_therm[1] if len(p_therm) > 1 else p_therm[0]) + # uniform out-of-plane B (tesla): electron cyclotron frequency e B / m_e = 0.17588 rad/ps per tesla + self.omega_c = 0.175882 * float(hpe.get("magnetic_field", 0.0)) + self.energy_conservation = bool(hpe.get("energy_conservation", False)) + self.ec_steps = max(float(hpe.get("energy_conservation_steps", 1.0)), 1.0) + self.flux_edges = jnp.asarray(flux_bin_edges(hpe)) + self.n_flux_bins = int(self.flux_edges.shape[0]) - 1 + cone_angle = hpe.get("cone_angle") + self.cone_cos = float(np.cos(np.deg2rad(float(cone_angle)))) if cone_angle is not None else None + cone_dir = np.asarray(hpe.get("cone_direction", [1.0, 0.0]), dtype=np.float64) + self.cone_dir = jnp.asarray(cone_dir / max(np.linalg.norm(cone_dir), 1e-30)) + # energy scales for the energy-conservation multiplier (both in erg per cm of depth): + # particle energy = sum(gamma - 1) * (N_tail / n_p) m c^2 with N_tail the retained tail + # electrons in the box; wave energy from k^2 |phi_k|^2 through the epw_energy normalization + c_cgs, me_cgs, e_cgs = 2.99792458e10, 9.10938291e-28, 4.8032068e-10 + w0_si = derived["w0"] * 1.0e12 + lambda_um = 2.0 * np.pi * derived["c"] / derived["w0"] + nc_cgs = 1.1148e21 / lambda_um**2 + lx_cm, ly_cm = self.Lx * 1.0e-4, self.ny * self.dy * 1.0e-4 + n_tail = float(np.mean(np.asarray(grid["background_density"]))) * nc_cgs * lx_cm * ly_cm * self.f_tail_frac + self.particle_energy_scale = n_tail / self.n_p * me_cgs * c_cgs**2 + field_unit = me_cgs * c_cgs * w0_si / e_cgs # statV/cm per unit of e_norm-scaled field + epw_energy_prefactor = 0.25 * self.dx * derived["x_norm"] * derived["e_norm"] ** 2 + self.wave_energy_scale = ( + epw_energy_prefactor + / (self.nx * self.ny**2) + * 4.0 + * (derived["c"] / derived["w0"]) + * 1.0e-4 + * ly_cm + * field_unit**2 + / (8.0 * np.pi) + * 2.0 # total (electrostatic + kinetic) EPW energy + ) + self.k_sq = self.kx[:, None] ** 2 + self.ky[None, :] ** 2 self.angles = jnp.asarray(arrays["angles"]) self.directions = jnp.stack((jnp.cos(self.angles), jnp.sin(self.angles)), axis=-1) self.theta_k = jnp.asarray(arrays["theta_k"]) @@ -347,9 +402,11 @@ def _substep(self, i, carry, ex_env, t0): u_half = u + 0.5 * self.dtp * acceleration gamma_rel = jnp.sqrt(1.0 + (u_half / self.c) ** 2) x = x + self.dtp * u_half / gamma_rel - if self.periodic_x: - x = jnp.mod(x - self.xmin, self.Lx) + self.xmin - acceleration = self._accel(x, ex_env, t_i + self.dtp) + # the gather needs in-box positions; the carried x stays unwrapped so that + # _apply_boundaries still sees (and counts / thermalizes) a periodic crossing at the + # end of the step, as the 2-D push does + x_in = jnp.mod(x - self.xmin, self.Lx) + self.xmin if self.periodic_x else x + acceleration = self._accel(x_in, ex_env, t_i + self.dtp) u = u_half + 0.5 * self.dtp * acceleration return x, u, acceleration @@ -358,27 +415,65 @@ def _substep_2d(self, i, carry, e_env, t0): t_i = t0 + i * self.dtp u_half = u + 0.5 * self.dtp * acceleration gamma_rel = jnp.sqrt(1.0 + jnp.sum((u_half / self.c) ** 2, axis=-1)) + if self.omega_c != 0.0: + # uniform B_z: rotate the momentum by omega_c dt / gamma (electrons turn counter-clockwise + # for B along +z); |u| and gamma are unchanged + theta = self.omega_c * self.dtp / gamma_rel + cos_t, sin_t = jnp.cos(theta), jnp.sin(theta) + u_half = jnp.stack( + (cos_t * u_half[:, 0] - sin_t * u_half[:, 1], sin_t * u_half[:, 0] + cos_t * u_half[:, 1]), axis=-1 + ) x = x + self.dtp * u_half[:, 0] / gamma_rel y = y + self.dtp * u_half[:, 1] / gamma_rel - if self.periodic_x: - x = jnp.mod(x - self.xmin, self.Lx) + self.xmin - if self.periodic_y: - y = jnp.mod(y - self.ymin, self.Ly) + self.ymin - acceleration = self._accel_2d(x, y, e_env, t_i + self.dtp) + # the gather needs in-box positions; a crossing at a periodic side is remembered as an + # offset so that the wall routine still sees (and counts) it at the end of the step + x_in = jnp.mod(x - self.xmin, self.Lx) + self.xmin if self.periodic_x else x + y_in = jnp.mod(y - self.ymin, self.Ly) + self.ymin if self.periodic_y else y + acceleration = self._accel_2d(x_in, y_in, e_env, t_i + self.dtp) u = u_half + 0.5 * self.dtp * acceleration return x, y, u, acceleration - def _apply_boundaries(self, x: Array, u: Array, t: float) -> tuple[Array, Array]: - if self.periodic_x: - return x, u + def _wall_instruments(self, out_masks, u_vec: Array, any_out: Array) -> tuple[Array, Array]: + """Energy (keV per test particle) leaving through each wall per energy bin, and inside the + acceptance cone (LPSE hpe.metrics.flux / power), for the particles crossing this step.""" + gamma_rel = jnp.sqrt(1.0 + jnp.sum((u_vec / self.c) ** 2, axis=-1)) + ke_kev = 510.999 * (gamma_rel - 1.0) + bins = jnp.clip(jnp.searchsorted(self.flux_edges, ke_kev, side="right") - 1, 0, self.n_flux_bins - 1) + # bins are half-open [e_i, e_{i+1}); energies outside [e_0, e_last) belong to no bin (the + # clipped index would otherwise fold them into the edge bins) + in_range = (ke_kev >= self.flux_edges[0]) & (ke_kev < self.flux_edges[-1]) + wall = jnp.zeros(self.n_p, dtype=jnp.int32) + for i, mask in enumerate(out_masks): + wall = jnp.where(mask, i, wall) + flux = jnp.zeros((len(WALLS), self.n_flux_bins)).at[wall, bins].add(jnp.where(any_out & in_range, ke_kev, 0.0)) + if self.cone_cos is None: + cone = jnp.zeros((1,)) + else: + speed = jnp.sqrt(jnp.sum(u_vec**2, axis=-1)) + cos_angle = jnp.sum(u_vec * self.cone_dir[None, :], axis=-1) / jnp.where(speed > 0, speed, 1.0) + cone = jnp.sum(jnp.where(any_out & (cos_angle >= self.cone_cos), ke_kev, 0.0))[None] + return flux, cone + + def _apply_boundaries(self, x: Array, u: Array, t: float): out_left, out_right = x < self.xmin, x > self.xmax - key = jax.random.fold_in(self.wall_key, jnp.asarray(t / self.dt).astype(jnp.int32)) + any_out = out_left | out_right + key, therm_key = jax.random.split(jax.random.fold_in(self.wall_key, jnp.asarray(t / self.dt).astype(jnp.int32))) uni = jax.random.uniform(key, (self.n_p,), minval=1.0e-12, maxval=1.0) speed = jnp.minimum(jnp.sqrt(self.v_min**2 - 2.0 * self.vte**2 * jnp.log(uni)), 0.99 * self.c) u_new = speed / jnp.sqrt(1.0 - (speed / self.c) ** 2) - x = jnp.where(out_left, self.xmin, jnp.where(out_right, self.xmax, x)) - u = jnp.where(out_left, u_new, jnp.where(out_right, -u_new, u)) - return x, u + flux, cone = self._wall_instruments( + (out_left, out_right, jnp.zeros_like(out_left), jnp.zeros_like(out_left)), + jnp.stack((u, jnp.zeros_like(u)), axis=-1), + any_out, + ) + # LPSE thermalizationProbability: thermalize (re-inject from the tail) with probability p, + # otherwise pass through periodically + therm = any_out & (jax.random.uniform(therm_key, (self.n_p,)) < self.p_therm_x) + x_wrapped = jnp.mod(x - self.xmin, self.Lx) + self.xmin + x = jnp.where(therm, jnp.where(out_left, self.xmin, self.xmax), x_wrapped) + u_therm = jnp.where(out_left, u_new, -u_new) + u = jnp.where(therm, u_therm, u) + return x, u, flux, cone def _maxwell_survival(self, speed: Array) -> Array: """Survival function of the 3-D Maxwell speed law with scale ``vte``. @@ -421,15 +516,18 @@ def bisect(_, bracket): return normal_velocity, tangential_velocity def _apply_boundaries_2d(self, x, y, u, t): - if self.periodic_x and self.periodic_y: - return x, y, u out_left, out_right = x < self.xmin, x > self.xmax out_bottom, out_top = y < self.ymin, y > self.ymax - cross_x = (out_left | out_right) if not self.periodic_x else jnp.zeros_like(out_left) - cross_y = (out_bottom | out_top) if not self.periodic_y else jnp.zeros_like(out_bottom) + cross_x = out_left | out_right + cross_y = out_bottom | out_top any_out = cross_x | cross_y + flux, cone = self._wall_instruments( + (out_left & cross_x, out_right & cross_x, out_bottom & cross_y, out_top & cross_y), u, any_out + ) - step_key = jax.random.fold_in(self.wall_key, jnp.asarray(t / self.dt).astype(jnp.int32)) + step_key, therm_key = jax.random.split( + jax.random.fold_in(self.wall_key, jnp.asarray(t / self.dt).astype(jnp.int32)) + ) normal_velocity, tangential_velocity = self._sample_flux_weighted_tail(step_key) x_sign = jnp.where(out_left, 1.0, -1.0) y_sign = jnp.where(out_bottom, 1.0, -1.0) @@ -447,31 +545,58 @@ def _apply_boundaries_2d(self, x, y, u, t): gamma_rel = 1.0 / jnp.sqrt(1.0 - (speed / self.c) ** 2) u_new = gamma_rel[:, None] * velocity - if not self.periodic_x: - x = jnp.where(out_left, self.xmin, jnp.where(out_right, self.xmax, x)) - if not self.periodic_y: - y = jnp.where(out_bottom, self.ymin, jnp.where(out_top, self.ymax, y)) - u = jnp.where(any_out[:, None], u_new, u) - return x, y, u - - def push(self, x: Array, u: Array, ex_env: Array, t: float) -> tuple[Array, Array]: - """Historical quasi-1D KDK entry point.""" + # LPSE thermalizationProbability per wall direction: thermalize with probability p_x / p_y + # (clamped to the wall, re-injected inward), otherwise pass through periodically + uni = jax.random.uniform(therm_key, (self.n_p,)) + therm = (cross_x & (uni < self.p_therm_x)) | (cross_y & (uni < self.p_therm_y)) + x_wall = jnp.where(out_left, self.xmin, jnp.where(out_right, self.xmax, x)) + y_wall = jnp.where(out_bottom, self.ymin, jnp.where(out_top, self.ymax, y)) + x_wrapped = jnp.mod(x - self.xmin, self.Lx) + self.xmin + y_wrapped = jnp.mod(y - self.ymin, self.Ly) + self.ymin + x = jnp.where(therm, x_wall, x_wrapped) + y = jnp.where(therm, y_wall, y_wrapped) + u = jnp.where(therm[:, None], u_new, u) + return x, y, u, flux, cone + + def _kinetic(self, u: Array) -> Array: + """sum(gamma - 1) over the ensemble (units of m c^2 per test particle).""" + u_sq = jnp.sum((u / self.c) ** 2, axis=-1) if u.ndim == 2 else (u / self.c) ** 2 + return jnp.sum(jnp.sqrt(1.0 + u_sq) - 1.0) + + def push(self, x: Array, u: Array, ex_env: Array, t: float): + """Historical quasi-1D KDK entry point. Returns ``(x, u, wall_flux, cone_energy, d_kinetic)``.""" acceleration = self._accel(x, ex_env, t) + ke0 = self._kinetic(u) x, u, _ = jax.lax.fori_loop( 0, self.n_sub, lambda i, carry: self._substep(i, carry, ex_env, t), (x, u, acceleration) ) - return self._apply_boundaries(x, u, t) + d_kinetic = self._kinetic(u) - ke0 + return (*self._apply_boundaries(x, u, t), d_kinetic) def push_2d(self, x: Array, y: Array, u: Array, e_env: Array, t: float): - """Relativistic KDK push for one ``(x,y,p_x,p_y)`` ensemble.""" + """Relativistic KDK push for one ``(x,y,p_x,p_y)`` ensemble. + Returns ``(x, y, u, wall_flux, cone_energy, d_kinetic)``.""" acceleration = self._accel_2d(x, y, e_env, t) + ke0 = self._kinetic(u) x, y, u, _ = jax.lax.fori_loop( 0, self.n_sub, lambda i, carry: self._substep_2d(i, carry, e_env, t), (x, y, u, acceleration), ) - return self._apply_boundaries_2d(x, y, u, t) + d_kinetic = self._kinetic(u) - ke0 + return (*self._apply_boundaries_2d(x, y, u, t), d_kinetic) + + def ld_multiplier(self, previous: Array, gamma_l: Array, phi_k: Array, d_kinetic: Array) -> Array: + """LPSE ``getLD_multiplierForEnergyConservation``: scale the applied Landau rate so that the + expected EPW energy loss over this step equals the energy the particles gained, clamped to + [1/10, 10] and averaged over ``energy_conservation_steps`` steps. The expected loss uses + the linearised ``2 gamma dt`` per mode (LPSE iterates the exponential form).""" + particle_gain = d_kinetic * self.particle_energy_scale # erg per cm of depth + wave_loss = self.wave_energy_scale * jnp.sum(2.0 * gamma_l * self.dt * self.k_sq * jnp.abs(phi_k) ** 2) + ratio = jnp.where(wave_loss > 0.0, particle_gain / jnp.where(wave_loss > 0.0, wave_loss, 1.0), 1.0) + target = jnp.clip(ratio, 0.1, 10.0) + return target / self.ec_steps + previous * (1.0 - 1.0 / self.ec_steps) # ------------------------------------------------- histogram and damping -- @@ -527,8 +652,21 @@ def gather_slope(angle_index): def _gamma_raw(self, hist: Array) -> Array: return self._gamma_raw_2d(hist) if self.is_2d else self._gamma_raw_1d(hist) + def limit_resonant(self, gamma: Array) -> Array: + """The configured limits (``gamma_limit_damping`` / ``-gamma_limit_growth``) on the resonant + modes of an applied rate; the non-resonant modes keep theirs. Re-applied after the + energy-conservation multiplier, so the rate handed to the solver respects the limits (LPSE + clips only before its LD_multiplier).""" + resonant = jnp.asarray(self.mask_res) + if not self.is_2d: + resonant = jnp.broadcast_to(resonant[:, None], (self.nx, self.ny)) + return jnp.where(resonant, jnp.clip(gamma, -self.gamma_limit_growth, self.gamma_limit_damping), gamma) + def damping(self, hist: Array) -> Array: - gamma_hpe = jnp.maximum(self.calibration * self._gamma_raw(hist), 0.0) + # LPSE gammaLimit: clip(-growth, gamma, damping); growth only with allow_growth + gamma_hpe = jnp.clip( + self.calibration * self._gamma_raw(hist), -self.gamma_limit_growth, self.gamma_limit_damping + ) if self.is_2d: return jnp.where(self.mask_res, gamma_hpe, self.gamma_analytic) gamma_1d = jnp.where(self.mask_res, gamma_hpe, self.gamma_analytic[:, 0]) @@ -537,28 +675,58 @@ def damping(self, hist: Array) -> Array: # ---------------------------------------------------------------- driver -- def __call__(self, t: float, y: dict[str, Array]) -> dict[str, Array]: + instruments = (y["hpe_wall_flux"], y["hpe_cone_energy"], y["hpe_ld_multiplier"]) + + def feedback(hist, gamma_l, mult, d_kinetic): + if not self.feedback: + return gamma_l, mult + gamma_new = self.damping(hist) + if self.energy_conservation: + mult = self.ld_multiplier(mult, gamma_new, y["epw"], d_kinetic) + gamma_new = self.limit_resonant(gamma_new * mult[0]) + return gamma_new, mult + if self.is_2d: def active(operand): - x, y_position, u, hist, gamma_l = operand - x, y_position, u = self.push_2d(x, y_position, u, self.refine_e(y["epw"]), t) + x, y_position, u, hist, gamma_l, flux, cone, mult = operand + x, y_position, u, dflux, dcone, d_kinetic = self.push_2d(x, y_position, u, self.refine_e(y["epw"]), t) hist = (1.0 - self.alpha) * hist + self.alpha * self.histogram(u) - if self.feedback: - gamma_l = self.damping(hist) - return x, y_position, u, hist, gamma_l + gamma_l, mult = feedback(hist, gamma_l, mult, d_kinetic) + return x, y_position, u, hist, gamma_l, flux + dflux, cone + dcone, mult - operand = (y["x_e"], y["y_e"], y["u_e"], y["epw_hist"], y["gamma_L"]) - x, y_position, u, hist, gamma_l = jax.lax.cond(t >= self.t_start, active, lambda value: value, operand) - return {**y, "x_e": x, "y_e": y_position, "u_e": u, "epw_hist": hist, "gamma_L": gamma_l} + operand = (y["x_e"], y["y_e"], y["u_e"], y["epw_hist"], y["gamma_L"], *instruments) + x, y_position, u, hist, gamma_l, flux, cone, mult = jax.lax.cond( + t >= self.t_start, active, lambda value: value, operand + ) + return { + **y, + "x_e": x, + "y_e": y_position, + "u_e": u, + "epw_hist": hist, + "gamma_L": gamma_l, + "hpe_wall_flux": flux, + "hpe_cone_energy": cone, + "hpe_ld_multiplier": mult, + } def active_1d(operand): - x, u, hist, gamma_l = operand - x, u = self.push(x, u, self.refine_ex(y["epw"]), t) + x, u, hist, gamma_l, flux, cone, mult = operand + x, u, dflux, dcone, d_kinetic = self.push(x, u, self.refine_ex(y["epw"]), t) hist = (1.0 - self.alpha) * hist + self.alpha * self.histogram(u) - if self.feedback: - gamma_l = self.damping(hist) - return x, u, hist, gamma_l - - operand = (y["x_e"], y["u_e"], y["epw_hist"], y["gamma_L"]) - x, u, hist, gamma_l = jax.lax.cond(t >= self.t_start, active_1d, lambda value: value, operand) - return {**y, "x_e": x, "u_e": u, "epw_hist": hist, "gamma_L": gamma_l} + gamma_l, mult = feedback(hist, gamma_l, mult, d_kinetic) + return x, u, hist, gamma_l, flux + dflux, cone + dcone, mult + + operand = (y["x_e"], y["u_e"], y["epw_hist"], y["gamma_L"], *instruments) + x, u, hist, gamma_l, flux, cone, mult = jax.lax.cond(t >= self.t_start, active_1d, lambda value: value, operand) + return { + **y, + "x_e": x, + "u_e": u, + "epw_hist": hist, + "gamma_L": gamma_l, + "hpe_wall_flux": flux, + "hpe_cone_energy": cone, + "hpe_ld_multiplier": mult, + } diff --git a/adept/_lpse2d/core/iaw.py b/adept/_lpse2d/core/iaw.py index fbbe5056..c64e3eec 100644 --- a/adept/_lpse2d/core/iaw.py +++ b/adept/_lpse2d/core/iaw.py @@ -1,9 +1,7 @@ """Ion-acoustic-wave evolution for the Envelope-2D solver. -This ports the ``isEvolveIaw`` split-step path from -``m201805_matlabLpse_v11.m``. The real state variables are the fractional ion -density perturbation ``n`` (MATLAB ``Nelf``) and the ion-velocity divergence -``w`` (MATLAB ``W``): +The real state variables are the fractional ion density perturbation ``n`` +(MATLAB ``Nelf``) and the ion-velocity divergence ``w`` (MATLAB ``W``): d w / dt = -laplacian(PP) d n / dt = -w @@ -12,17 +10,99 @@ PP = cs^2 n + Z e^2/(4 me mi) [|E_epw|^2/wp0^2 + |E0|^2/w0^2 - + |E1|^2/w1^2]. + + |E1|^2/w_R^2], -The update order matches MATLAB's ``iawSplitStep``: kick ``w`` from the -potential, apply twice the requested amplitude Landau rate to ``w``, then -drift ``n``. Boundary damping acts on ``w`` and collisions act on ``n`` after -the split step, as in the original time loop. +``w_R`` the Raman light's carrier: ``w1 = w0 - wp0``, or ``wp0`` with the combined solver, whose +transverse field (the Raman light) is enveloped there. + +Two time integrators are available (``terms.iaw.solver``): + +``explicit`` (the ``isEvolveIaw`` split-step path of ``m201805_matlabLpse_v11.m``): +kick ``w`` from the potential with the FD Laplacian, apply twice the amplitude +Landau rate to ``w``, drift ``n``; boundary damping acts on ``w`` and collisions +act on ``n`` after the split step. Conditionally stable, ``omega_max dt < 2``. + +``spectral`` (the original LPSE ``iaw.solver = spectral`` path, +``ZakharovSolver.cpp`` lines 2083-2113): per k-mode the exact solution over +``dt`` of the damped oscillator + + dn/dt = -w, dw/dt = cs^2 k^2 n - 2 gamma_k w + +with ``gamma_k`` the amplitude Landau rate, i.e. with ``beta_k = sqrt(cs^2 k^2 - gamma_k^2)`` + + n' = e^{-gamma dt} [ cos(beta dt) n + sin(beta dt)/beta (gamma n - w) ] + w' = e^{-gamma dt} [ cos(beta dt) w + sin(beta dt)/beta (cs^2 k^2 n - gamma w) ] + +times the Doppler phase ``e^{-i k . V0 dt}`` of a uniform background flow ``V0``, +followed by the split-step pieces LPSE applies in the same order: collisional +damping ``w' *= e^{-2 nu dt}`` (on the velocity divergence, not on ``n``), the +ponderomotive kick ``w' += dt k^2 PP_k`` (``PP`` without its ``cs^2 n`` part, which +is inside the propagator), and the fluctuation-dissipation noise source on ``w``. +Unconditionally stable and exact for the linear acoustic part at any ``dt``, so +the IAW may also be advanced every ``terms.iaw.stride`` EPW steps as LPSE does. """ +import jax +import numpy as np from jax import Array from jax import numpy as jnp +IAW_SOLVERS = ("explicit", "spectral") +ION_LANDAU_FORMS = ("simplified", "full") + + +def _Q_kev(value) -> float: + from adept._lpse2d.helpers import _Q + + return float(_Q(value).to("keV").value) + + +def ion_landau_rate(cfg: dict, k_sq: np.ndarray) -> np.ndarray: + """Amplitude ion Landau damping rate ``gamma_k`` (1/ps) for every k mode. + + ``terms.iaw.damping.landau_form: simplified`` (default; LPSE ``isSimplified``): + ``gamma = landau * cs * |k|``. + + ``full``: the Z-generalized Krall & Trivelpiece expression (p. 390) used by LPSE's + IAW solver (``IawSolver.cpp``, modified by J. Palastro for the effective ionization + state), with ``eta = 1 + 3 Ti/(Z Te)``, ``M = (mi/me)/(eta Z)``, ``L = 1 + k^2 lambda_D^2``: + + W_i = W_r sqrt(pi/8) L^{-3/2} [ (3/(eta-1))^{3/2} exp(-(3/(eta-1))/(2L)) + sqrt(1/(eta M)) ], + W_r = cs |k| / sqrt(L), + + ``W_i`` being the damping rate of the velocity divergence, twice the amplitude rate. + """ + derived = cfg["units"]["derived"] + iaw = cfg["terms"]["iaw"] + damping = iaw["damping"] + cs = derived["cs"] + k = np.sqrt(np.asarray(k_sq)) + form = str(damping.get("landau_form", "simplified")) + if form == "simplified": + return float(damping["landau"]) * cs * k + if form != "full": + raise ValueError(f"terms.iaw.damping.landau_form must be one of {ION_LANDAU_FORMS}, got {form!r}") + + from astropy.units import Quantity as _Q + + te = _Q(cfg["units"]["reference electron temperature"]).to("keV").value + ti = _Q(cfg["units"]["reference ion temperature"]).to("keV").value + z = float(cfg["units"]["ionization state"]) + mi_over_me = derived["mi"] / derived["me"] + eta = 1.0 + 3.0 * ti / (z * te) + big_m = mi_over_me / (eta * z) + lambda_d_sq = derived["vte_sq"] / derived["wp0"] ** 2 + big_l = 1.0 + k_sq * lambda_d_sq + w_r = cs * k / np.sqrt(big_l) + three_over = 3.0 / (eta - 1.0) + w_i = ( + w_r + * np.sqrt(np.pi / 8.0) + / big_l**1.5 + * (three_over**1.5 * np.exp(-0.5 * three_over / big_l) + np.sqrt(1.0 / (eta * big_m))) + ) + return 0.5 * w_i + class IonAcousticWave: """Advance the ion-acoustic density and velocity-divergence fields.""" @@ -32,29 +112,203 @@ def __init__(self, cfg: dict): derived = cfg["units"]["derived"] iaw = cfg["terms"]["iaw"] - self.dt = grid["dt"] + self.solver = str(iaw.get("solver", "explicit")) + if self.solver not in IAW_SOLVERS: + raise ValueError(f"terms.iaw.solver must be one of {IAW_SOLVERS}, got {self.solver!r}") + self.stride = int(iaw.get("stride", 1) or 1) + if self.stride < 1: + raise ValueError("terms.iaw.stride must be a positive integer") + if self.stride > 1 and self.solver != "spectral": + raise ValueError( + "terms.iaw.stride > 1 requires terms.iaw.solver: spectral (the explicit step is not stable)" + ) + + self.dt = grid["dt"] * self.stride self.dx = grid["dx"] self.dy = grid["dy"] + self.nx = grid["nx"] self.ny = grid["ny"] self.kx = grid["kx"] self.ky = grid["ky"] self.k_sq = self.kx[:, None] ** 2 + self.ky[None, :] ** 2 self.filter = grid["low_pass_filter_grid"] * grid["zero_mask"] - self.boundary = grid["iaw_absorbing_boundaries"] + # the absorber factor is exp(-rate grid.dt) per EPW step; a strided update spans stride steps + self.boundary = grid["iaw_absorbing_boundaries"] ** self.stride self.cs = derived["cs"] self.wp0 = derived["wp0"] self.w0 = derived["w0"] self.w1 = derived["w1"] + # the Raman light's carrier: the combined solver envelopes its field at wp0 + self.combined = cfg["terms"]["epw"].get("solver", "separate") == "combined" + self.w_raman = self.wp0 if self.combined else self.w1 self.ponderomotive_prefactor = ( cfg["units"]["ionization state"] * derived["e"] ** 2 / (4.0 * derived["me"] * derived["mi"]) ) + self._init_thermal_filamentation(cfg) damping = iaw["damping"] - self.landau_rate = damping["landau"] * self.cs * jnp.sqrt(self.k_sq) - self.collisional_factor = 1.0 - damping["collisions"] * self.dt + k_sq_np = np.asarray(self.k_sq) + self.landau_rate = jnp.asarray(ion_landau_rate(cfg, k_sq_np)) + self.nu_coll = float(damping["collisions"]) + # explicit (MATLAB): collisions act on n as (1 - nu dt); spectral (LPSE): on w as e^{-2 nu dt} + self.collisional_factor = 1.0 - self.nu_coll * self.dt self.max_density_perturbation = iaw["max_density_perturbation"] + # uniform background flow (Mach number along x, y) -> Doppler phase in the propagator + flow = iaw.get("flow") + if flow is None: + self.flow = None + else: + if self.solver != "spectral": + raise ValueError("terms.iaw.flow requires terms.iaw.solver: spectral") + self.flow = np.asarray(flow, dtype=np.float64) * self.cs + if self.flow.shape != (2,): + raise ValueError("terms.iaw.flow must be [Mach_x, Mach_y]") + + if self.solver == "spectral": + gamma = np.asarray(self.landau_rate) + omega_sq = self.cs**2 * k_sq_np + # beta may be imaginary for over-damped modes (gamma > cs k): the complex + # cos/sin form below is the exact solution in both regimes + beta = np.sqrt((omega_sq - gamma**2).astype(np.complex128)) + beta_safe = np.where(np.abs(beta) > 0.0, beta, 1.0) + cos_b = np.cos(beta * self.dt) + sin_over_b = np.where(np.abs(beta) > 0.0, np.sin(beta * self.dt) / beta_safe, self.dt) + phase = -1j * (k_sq_np * 0.0) + if self.flow is not None: + kx_np, ky_np = np.asarray(self.kx)[:, None], np.asarray(self.ky)[None, :] + phase = -1j * (kx_np * self.flow[0] + ky_np * self.flow[1]) + coeff = np.exp((-gamma + phase) * self.dt) + # propagator matrix on (n_k, w_k) + self.p_nn = jnp.asarray(coeff * (cos_b + sin_over_b * gamma)) + self.p_nw = jnp.asarray(-coeff * sin_over_b) + self.p_wn = jnp.asarray(coeff * sin_over_b * omega_sq) + self.p_ww = jnp.asarray(coeff * (cos_b - sin_over_b * gamma)) + self.collisional_w_factor = float(np.exp(-2.0 * self.nu_coll * self.dt)) + + # IAW noise (LPSE IawSolver::addNoise): random-phase source on the velocity + # divergence with the fluctuation-dissipation amplitude + # A N sqrt(exp(2 dt (gamma_k + nu)) - 1) + # on the retained band; LPSE refuses it without damping + self.noise_enabled = bool(iaw.get("noise", False)) + if self.noise_enabled: + amplitude = float(iaw.get("noise_amplitude", 1.0)) + total = np.asarray(self.landau_rate) + self.nu_coll + band = np.asarray(self.filter) > 0.0 + if not np.any(total[band] > 0.0): + raise ValueError("terms.iaw.noise requires IAW damping (landau and/or collisions)") + kick = amplitude * float(self.nx * self.ny) * np.sqrt(np.expm1(2.0 * self.dt * total)) + self.noise_kick = jnp.asarray(np.where(band, kick, 0.0)) + seed = iaw.get("noise_seed") + self.noise_key = jax.random.PRNGKey(int(seed) if seed is not None else 271828) + else: + self.noise_kick = None + + def _init_thermal_filamentation(self, cfg: dict) -> None: + """LPSE ``thermalFil.{laser,raman,lw}`` (``ZakharovSolver::getThermalFilamentationSource``): + inverse-bremsstrahlung heating by the spatially varying part of each wave's intensity, + balanced by Spitzer heat conduction, drives the ion flow through the electron pressure: + + d(div v)/dt += Z Q_w / (m_i kappa'), Q_w = nu_w(n) (|E_w|^2 - <|E_w|^2>) / (8 pi) + + with ``kappa'`` the Spitzer conductivity over k_B (1/(cm s)), ``nu_w`` the wave's energy + damping rate (light: 2 nu_abs(n_c) (n/n_c)^2; EPW: 2 nu_coll n/n_env), ``<.>`` the box + average. ``nonlocal: true`` adds LPSE's k^(4/3) correction: the source's k-space content is + multiplied by ``1 + (k lambda_nl)^(4/3)`` with ``lambda_nl = 30 (k_B T_e)^2 / (4 pi e^4 + sqrt(Z+1) ln Lambda n_e)``. The LPSE form is kept; the normalization is adept's own + (derived from the heat and momentum equations, not LPSE's ZAK constants).""" + tf = cfg["terms"]["iaw"].get("thermal_filamentation") or {} + self.thermal_waves = [w for w in ("laser", "raman", "lw") if tf.get(w, False)] + if not self.thermal_waves: + return + derived = cfg["units"]["derived"] + units = cfg["units"] + from adept._lpse2d.core.raman import light_absorption_rates + + rate0, rate1 = light_absorption_rates(cfg) + z = float(units["ionization state"]) + te_kev = float(derived.get("Te", 0.0)) or float(_Q_kev(units["reference electron temperature"])) + # cgs constants + e_cgs, me_cgs, mp_cgs, kb_cgs = 4.8032068e-10, 9.10938291e-28, 1.6726219e-24, 1.380649e-16 + te_k = te_kev * 1.0e3 * 1.16045e4 + lambda_um = 2.0 * np.pi * derived["c"] / derived["w0"] + nc_cgs = 1.1148e21 / lambda_um**2 + n_cgs = np.asarray(cfg["grid"]["background_density"]) * nc_cgs + n_mean = float(np.mean(n_cgs)) + log_lambda = max( + 23.5 - np.log(np.sqrt(n_mean) * te_kev**-1.25 * 1e3**-1.25 * 1e3) + if False + else 6.68 + np.log(lambda_um * te_kev), + 2.0, + ) + g_factor = 1.0 / (1.0 + 3.3 / z) + kappa = (8.0 / np.pi) ** 1.5 * g_factor * (kb_cgs * te_k) ** 2.5 / (z * e_cgs**4 * np.sqrt(me_cgs) * log_lambda) + kappa *= float(tf.get("conductivity_multiplier", 1.0)) + mi_cgs = mp_cgs * float(units["atomic number"]) + field_scale = float(derived["fieldScale"]) # statV/cm per code field unit + # d(div v)/dt [1/s^2] = Z nu(n)[1/s] |E|^2_cgs / (8 pi m_i kappa'); -> 1/ps^2 per code |E|^2 + base = z * field_scale**2 / (8.0 * np.pi * mi_cgs * kappa) * 1.0e-24 * 1.0e12 # per (1/ps rate) + n_over_nc = np.asarray(cfg["grid"]["background_density"]) + self.thermal_coeff = {} + if "laser" in self.thermal_waves: + if rate0 is None: + raise ValueError("terms.iaw.thermal_filamentation.laser needs terms.light.absorption") + self.thermal_coeff["laser"] = jnp.asarray(base * 2.0 * rate0 * n_over_nc**2) + if "raman" in self.thermal_waves: + if rate1 is None: + raise ValueError("terms.iaw.thermal_filamentation.raman needs terms.light.absorption") + nc1 = (self.w_raman / derived["w0"]) ** 2 # the Raman light's critical density (n_env combined) + self.thermal_coeff["raman"] = jnp.asarray(base * 2.0 * rate1 * (n_over_nc / nc1) ** 2) + if "lw" in self.thermal_waves and self.combined: + # the combined solver damps the EPW (its field's longitudinal part) with the Raman absorption + # nu_R (n/n_env)^2, not the EPW collisions: the raman term's coefficient on that part + from adept._lpse2d.core.combined import combined_collision_rate + + nu_c = combined_collision_rate(cfg) + if nu_c <= 0.0: + raise ValueError( + "terms.iaw.thermal_filamentation.lw with terms.epw.solver: combined needs terms.light.absorption " + "(the combined field's damping; the EPW collisions do not act on it)" + ) + n_over_env = n_over_nc / float(units["envelope density"]) + self.thermal_coeff["lw"] = jnp.asarray(base * 2.0 * nu_c * n_over_env**2) + elif "lw" in self.thermal_waves: + nu_coll = float(derived.get("nu_coll", 0.0)) + if nu_coll <= 0.0: + raise ValueError("terms.iaw.thermal_filamentation.lw needs terms.epw.damping.collisions") + self.thermal_coeff["lw"] = jnp.asarray(base * 2.0 * nu_coll * n_over_nc / float(units["envelope density"])) + self.thermal_nonlocal = bool(tf.get("nonlocal", tf.get("nonlocal_", False))) + if self.thermal_nonlocal: + lambda_nl_cm = ( + 30.0 * (kb_cgs * te_k) ** 2 / (4.0 * np.pi * e_cgs**4 * np.sqrt(z + 1.0) * log_lambda * n_cgs) + ) + lambda_nl_um = lambda_nl_cm * 1.0e4 + k_mag = np.sqrt(self.k_sq) + # (k lambda_nl(n))^(4/3): k^(4/3) in k-space, lambda^(4/3)(x) in x-space + self.thermal_k_factor = jnp.asarray(k_mag ** (4.0 / 3.0)) + self.thermal_lambda_factor = jnp.asarray(lambda_nl_um ** (4.0 / 3.0)) + print(f"IAW thermal filamentation on ({', '.join(self.thermal_waves)}, nonlocal={self.thermal_nonlocal})") + + def thermal_filamentation_source(self, phi_k: Array, E0: Array, E1: Array) -> Array: + """The heating source on the velocity divergence (1/ps^2), x-space; zero when off.""" + if not self.thermal_waves: + return jnp.zeros((self.nx, self.ny)) + source = jnp.zeros((self.nx, self.ny)) + for wave, coeff in self.thermal_coeff.items(): + if wave == "laser": + e_sq = jnp.sum(jnp.abs(E0) ** 2, axis=-1) + elif wave == "raman": + e_sq = jnp.sum(jnp.abs(E1) ** 2, axis=-1) + else: + ex, ey = self.epw_fields(phi_k) + e_sq = jnp.abs(ex) ** 2 + jnp.abs(ey) ** 2 + source = source + coeff * (e_sq - jnp.mean(e_sq)) + if self.thermal_nonlocal: + nonlocal_part = jnp.real(jnp.fft.ifft2(self.thermal_k_factor * jnp.fft.fft2(source))) + source = source + self.thermal_lambda_factor * nonlocal_part + return source + def laplacian(self, field: Array) -> Array: """Second-order periodic finite-difference Laplacian used by MATLAB.""" lap = (jnp.roll(field, -1, axis=0) - 2.0 * field + jnp.roll(field, 1, axis=0)) / self.dx**2 @@ -68,26 +322,73 @@ def epw_fields(self, phi_k: Array) -> tuple[Array, Array]: ey = jnp.fft.ifft2(-1j * self.ky[None, :] * phi_k) return ex, ey - def ponderomotive_potential(self, phi_k: Array, E0: Array, E1: Array, density: Array) -> Array: - """Build the acoustic plus EPW/pump/Raman ponderomotive potential.""" + def ponderomotive_drive(self, phi_k: Array, E0: Array, E1: Array) -> Array: + """The EPW/pump/Raman part of the ponderomotive potential (no acoustic term).""" ex, ey = self.epw_fields(phi_k) epw_intensity = jnp.abs(ex) ** 2 + jnp.abs(ey) ** 2 pump_intensity = jnp.sum(jnp.abs(E0) ** 2, axis=-1) raman_intensity = jnp.sum(jnp.abs(E1) ** 2, axis=-1) - return self.cs**2 * density + self.ponderomotive_prefactor * ( - epw_intensity / self.wp0**2 + pump_intensity / self.w0**2 + raman_intensity / self.w1**2 + return self.ponderomotive_prefactor * ( + epw_intensity / self.wp0**2 + pump_intensity / self.w0**2 + raman_intensity / self.w_raman**2 ) - def __call__(self, y: dict[str, Array]) -> dict[str, Array]: - """Advance one MATLAB-order IAW split step and return the full state.""" + def ponderomotive_potential(self, phi_k: Array, E0: Array, E1: Array, density: Array) -> Array: + """Build the acoustic plus EPW/pump/Raman ponderomotive potential.""" + return self.cs**2 * density + self.ponderomotive_drive(phi_k, E0, E1) + + def get_noise(self, t: float) -> Array: + step = jnp.round(t / self.dt).astype(int) + key = jax.random.fold_in(self.noise_key, step) + phases = 2.0 * np.pi * jax.random.uniform(key, (self.nx, self.ny)) + return self.noise_kick * jnp.exp(1j * phases) + + def _explicit_step(self, t: float, y: dict[str, Array]) -> tuple[Array, Array]: potential = self.ponderomotive_potential(y["epw"], y["E0"], y["E1"], y["iaw_density"]) velocity_divergence = y["iaw_velocity_divergence"] - self.dt * self.laplacian(potential) + if self.thermal_waves: + velocity_divergence = velocity_divergence + self.dt * self.thermal_filamentation_source( + y["epw"], y["E0"], y["E1"] + ) velocity_k = jnp.fft.fft2(velocity_divergence) velocity_k = velocity_k * jnp.exp(-2.0 * self.landau_rate * self.dt) * self.filter + if self.noise_enabled: + velocity_k = velocity_k + self.get_noise(t) velocity_divergence = jnp.real(jnp.fft.ifft2(velocity_k)) * self.boundary density = (y["iaw_density"] - self.dt * velocity_divergence) * self.collisional_factor + return density, velocity_divergence + + def _spectral_step(self, t: float, y: dict[str, Array]) -> tuple[Array, Array]: + n_k = jnp.fft.fft2(y["iaw_density"]) + w_k = jnp.fft.fft2(y["iaw_velocity_divergence"]) + + # exact damped-oscillator propagator with the flow Doppler phase + n_new = self.p_nn * n_k + self.p_nw * w_k + w_new = self.p_wn * n_k + self.p_ww * w_k + + # split-step pieces in LPSE's order: collisions on div v, ponderomotive kick, noise + w_new = w_new * self.collisional_w_factor + drive_k = jnp.fft.fft2(self.ponderomotive_drive(y["epw"], y["E0"], y["E1"])) + w_new = w_new + self.dt * self.k_sq * drive_k + if self.thermal_waves: + w_new = w_new + self.dt * jnp.fft.fft2(self.thermal_filamentation_source(y["epw"], y["E0"], y["E1"])) + if self.noise_enabled: + w_new = w_new + self.get_noise(t) + + n_new = n_new * self.filter + w_new = w_new * self.filter + density = jnp.real(jnp.fft.ifft2(n_new)) + velocity_divergence = jnp.real(jnp.fft.ifft2(w_new)) * self.boundary + return density, velocity_divergence + + def __call__(self, y: dict[str, Array], t: float = 0.0) -> dict[str, Array]: + """Advance one IAW step (of length ``stride * grid.dt``) and return the full state.""" + if self.solver == "spectral": + density, velocity_divergence = self._spectral_step(t, y) + else: + density, velocity_divergence = self._explicit_step(t, y) + if self.max_density_perturbation is not None: density = jnp.clip( density, diff --git a/adept/_lpse2d/core/laser.py b/adept/_lpse2d/core/laser.py index f2830b9e..2acfbabb 100644 --- a/adept/_lpse2d/core/laser.py +++ b/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 diff --git a/adept/_lpse2d/core/light.py b/adept/_lpse2d/core/light.py index 243b9cf2..bc0f614a 100644 --- a/adept/_lpse2d/core/light.py +++ b/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)) diff --git a/adept/_lpse2d/core/raman.py b/adept/_lpse2d/core/raman.py index be685e29..2423e3ec 100644 --- a/adept/_lpse2d/core/raman.py +++ b/adept/_lpse2d/core/raman.py @@ -3,6 +3,52 @@ from jax import numpy as jnp +def transverse_part(field: Array, kx: Array, ky: Array, one_over_k_sq: Array) -> Array: + """The transverse (divergence-free) part of an (nx, ny, 2) x-space vector field, via the + k-space projector ``F - k (k . F) / k^2`` (LPSE ``LightSolver::getTransversePartOfSourceTerm``).""" + fx_k = jnp.fft.fft2(field[..., 0]) + fy_k = jnp.fft.fft2(field[..., 1]) + kdote = (kx[:, None] * fx_k + ky[None, :] * fy_k) * one_over_k_sq + return jnp.stack([jnp.fft.ifft2(fx_k - kx[:, None] * kdote), jnp.fft.ifft2(fy_k - ky[None, :] * kdote)], axis=-1) + + +def light_absorption_rates(cfg: dict) -> tuple[float | None, float | None]: + """Amplitude absorption rates (1/ps) at the critical density of the pump and of the + Raman light, from ``terms.light.absorption``: ``false`` (none), ``true`` (NRL plasma + formulary as coded in LPSE ``LightSolver.cpp``: ``nu = 5.11e10 Z logLambda / + (lambda_um^2 Te_keV^1.5) * 1e-12``, ``logLambda = 6.68 + ln(lambda_um Te)`` for + ``Te > 0.01 Z^2`` else ``9.13 + ln(lambda_um Te^1.5 / Z)``, evaluated with each wave's + own wavelength), or a number (rate at nc for the pump; the Raman value is scaled by + ``(lambda_1/lambda_0)^-2 = (w1/w0)^2`` as the formula's leading dependence). The combined + solver's field is the Raman light at carrier ``wp0``: its wavelength is ``lambda_0 w0/wp0``.""" + from astropy.units import Quantity as _Q + + absorption = cfg["terms"].get("light", {}).get("absorption", False) + if absorption is None or absorption is False: + return None, None + derived = cfg["units"]["derived"] + z = float(cfg["units"]["ionization state"]) + te = _Q(cfg["units"]["reference electron temperature"]).to("keV").value + lambda0 = _Q(cfg["units"]["laser_wavelength"]).to("um").value + combined = cfg["terms"]["epw"].get("solver", "separate") == "combined" + w_raman = derived["wp0"] if combined else derived["w1"] + lambda1 = lambda0 * derived["w0"] / w_raman + + def nrl(lam_um): + if te > 0.01 * z**2: + log_lambda = 6.68 + np.log(lam_um * te) + else: + log_lambda = 9.13 + np.log(lam_um * te**1.5 / z) + return 5.11e10 * z * log_lambda / (lam_um**2 * te**1.5) * 1.0e-12 + + if absorption is True: + return float(nrl(lambda0)), float(nrl(lambda1)) + rate0 = float(absorption) + if rate0 < 0.0: + raise ValueError("terms.light.absorption must be false, true or a non-negative rate in 1/ps") + return rate0, rate0 * (w_raman / derived["w0"]) ** 2 + + class RamanLight: """ Evolves the Raman scattered-light envelope E1. @@ -47,6 +93,15 @@ def __init__(self, cfg: dict): self.x = cfg["grid"]["x"] self.y = cfg["grid"]["y"] self.k_sq = cfg["grid"]["kx"][:, None] ** 2 + cfg["grid"]["ky"][None, :] ** 2 + self.kx_arr = jnp.asarray(cfg["grid"]["kx"]) + self.ky_arr = jnp.asarray(cfg["grid"]["ky"]) + self.one_over_k_sq = jnp.asarray(np.where(self.k_sq > 0, 1.0 / np.where(self.k_sq > 0, self.k_sq, 1.0), 0.0)) + # terms.light.transverse_source (LPSE takeTransversePartOfSourceTerms): the SRS + # source is projected onto its transverse part. The longitudinal part of E1 is not + # moved by any light propagator, so without the projection it accumulates the + # source and pairs with the EPW in a spurious two-wave instability (in the test_006 + # cross-check: 41/ps energy growth against LPSE's 5.8/ps) + self.transverse_source = bool(cfg["terms"].get("light", {}).get("transverse_source", True)) background_density = cfg["grid"]["background_density"] # local detuning of the Raman envelope (MATLAB line 1668-1670) @@ -58,7 +113,19 @@ def __init__(self, cfg: dict): # 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) + self.sub_boundary = cfg["grid"]["light_absorbing_boundaries"] ** (1.0 / self.n_sub) + + # collisional (inverse-bremsstrahlung) absorption, terms.light.absorption: the + # amplitude decays at nu_abs (n/nc_w)^2 per wave, nc_w its own critical density + # (LPSE calculateScatteringPotential; nu_abs from the NRL formula at nc_w with + # the wave's wavelength, or a user rate in 1/ps at nc_w) + self.n_over_env = background_density / self.envelope_density + self.n_over_nc0 = background_density # n / nc (w0) + self.n_over_nc1 = background_density * (self.w0 / self.w1) ** 2 # n / nc (w1) + self.absorption_rate0, self.absorption_rate1 = light_absorption_rates(cfg) + self.absorption_factor1 = ( + None if self.absorption_rate1 is None else jnp.exp(-self.absorption_rate1 * self.dt_l * self.n_over_nc1**2) + ) # seed injection (MATLAB lines 1757-1769): a two-point antisymmetric source that # launches a leftward-propagating (-x) wave at x = xmax - offset @@ -77,6 +144,7 @@ def __init__(self, cfg: dict): "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.k1_inject = float(np.sqrt(self.w1**2 - self.wpe_sq_i1) / self.c) self.seed_enabled = True else: self.seed_enabled = False @@ -140,8 +208,11 @@ def rhs( # SRS coupling to the EPW (MATLAB lines 1684-1689, potential formulation); # CoupledLight switches it off here when it integrates the exchange exactly if couple: - k_e1x += self.srs_coeff * jnp.conj(laplacian_phi) * E0[..., 0] - k_e1y += self.srs_coeff * jnp.conj(laplacian_phi) * E0[..., 1] + source = (self.srs_coeff * jnp.conj(laplacian_phi))[..., None] * E0 + if self.transverse_source: + source = transverse_part(source, self.kx_arr, self.ky_arr, self.one_over_k_sq) + k_e1x += source[..., 0] + k_e1y += source[..., 1] if seed_args is not None: row_i1, row_i1p1 = self.calc_seed_source(t, seed_args) @@ -186,8 +257,16 @@ def substep(i, E1): iaw_density, ) E1 = E1 + 1j * self.dt_l * jnp.imag(k2) - # absorbing boundaries (MATLAB lines 977-983) + # absorbing boundaries (MATLAB lines 977-983) and collisional absorption E1 = E1 * self.sub_boundary[..., None] + if absorb is not None: + E1 = E1 * absorb return E1 + absorb = None + if self.absorption_rate1 is not None: + n_over_nc = ( + self.n_over_nc1 if iaw_density is None else self.n_over_nc1 * (1.0 + iaw_density / self.n_over_env) + ) + absorb = jnp.exp(-self.absorption_rate1 * self.dt_l * n_over_nc**2)[..., None] return lax.fori_loop(0, self.n_sub, substep, E1) diff --git a/adept/_lpse2d/core/spectral_light.py b/adept/_lpse2d/core/spectral_light.py new file mode 100644 index 00000000..f3b27986 --- /dev/null +++ b/adept/_lpse2d/core/spectral_light.py @@ -0,0 +1,348 @@ +"""Spectral (k-space) light-wave propagators for the Envelope-2D solver. + +This is the original LPSE ``{laser|raman}.solver = spectral`` path +(``LightSolver::evolveSpectral``). Per light sub-step, in LPSE's order: + +1. injector source and the x-space split step of the scattering potential + ``E *= exp(i dt Vo)`` -- ``Vo = w/2 (1 - wpe^2/w^2 n_tot/n_env)`` plus, with + ``terms.light.absorption``, the imaginary part ``nu_abs (n/nc_w)^2`` (collisional + absorption; ``LightSolver::calculateScatteringPotential``); +2. the EPW coupling terms added with forward Euler in x-space; +3. the exact k-space propagator ``exp(-i dt c^2 k^2/(2 w))`` applied to the + transverse part of the field only (``k k/k^2`` longitudinal projector; the + longitudinal part is left unpropagated, ``LightSolver.cpp:3684-3743``), with + modes outside the retained band zeroed; +4. the absorbing layers. + +Compared with the staggered finite-difference scheme of ``raman.py``/``light.py`` +this has no CFL limit (``grid.light_substeps`` defaults to 1) and no grid +dispersion, so the group velocity at 7-8 cells per wavelength is exact and the +staggered scheme's exchange instability cannot occur. With the pump evolved the +SRS exchange between E0 and E1 is integrated with the exact local rotation of +``CoupledLight.couple`` on either side of the propagation (Strang), which is what +``terms.light.coupling: rotation`` does on the FD path. + +The injector is a smooth (Gaussian, ``injector_width`` wide, half a local +wavelength by default as LPSE's ``injectorWidth``) source ``S(x) = A v_g g(x) e^{+-i k x}``, ``int g dx = 1``, +which by ``dA/dx = S/v_g`` launches the amplitude ``A`` in the chosen direction +(to 0.2 % at half a wavelength: the paraxial correction grows with the width) and a +leakage ``~exp(-2 k^2 sigma^2)`` in the other -- the smooth analogue of LPSE's +sigmoid spectral injector (``laser.spectral.injectorWidth``). The FD two-point +injector is not usable here: its calibration relies on the FD Green's function. +Keep the injector clear of the absorbing layer (its skirt damps the source region): +the seed's default offset of 1.6 boundary widths sits inside the tanh skirt, which +the FD path tolerates by design; with this solver use ``drivers.E1.offset`` of at +least ``2 * boundary_width``. +""" + +import numpy as np +from jax import Array, lax +from jax import numpy as jnp + +from adept._base_ import get_envelope +from adept._lpse2d.core.light import CoupledLight +from adept._lpse2d.core.raman import RamanLight, transverse_part + + +def gaussian_injector_profile(x: np.ndarray, x_inject: float, width: float, dx: float) -> np.ndarray: + """Unit-integral (sum * dx = 1) Gaussian source profile centered at ``x_inject``.""" + g = np.exp(-0.5 * ((np.asarray(x) - x_inject) / width) ** 2) + return g / (np.sum(g) * dx) + + +def smooth_seed_source(t, seed_args: dict, *, w: float, c: float, x, y, x_inject, wpe_sq: float, profile) -> Array: + """Smooth injector for the leftward Raman seed at carrier ``w``: the (nx, ny) source added + to ``E1_y``, launching ``A / eps^(1/4)`` (``eps = 1 - wpe^2/w^2`` at the injector) at the + local ``k = w/c sqrt((1 + dw)^2 - wpe^2/w^2)``. ``w`` is the Raman envelope frequency + ``w1 = w0 - wp0`` of the separate solvers and ``wp0`` for the combined field (LPSE's + ``W1 = wpe`` in combined mode, ``LightSolver.cpp:2404``).""" + dw1 = seed_args["delta_omega"] + turn_on = 1.0 - jnp.exp(-((t / seed_args["turn_on_time"]) ** 2)) + if seed_args["yw"] > 0: + envelope_y = jnp.exp(-((y / (seed_args["yw"] / 2.0)) ** 4)) + else: + envelope_y = jnp.ones_like(y) + k1 = w / c * jnp.sqrt((1.0 + dw1) ** 2 - wpe_sq / w**2) + v_g = c**2 * k1 / w + eps1 = 1.0 - wpe_sq / w**2 + amplitude = seed_args["amplitude"] / eps1**0.25 * turn_on + carrier = jnp.exp(-1j * k1 * (x - x_inject) - 1j * w * dw1 * t) + return (amplitude * v_g * profile * carrier)[:, None] * envelope_y[None, :] + + +def transverse_propagate( + field: Array, + kx: Array, + ky: Array, + one_over_k_sq: Array, + propagator: Array, + band: Array | float = 1.0, + keep_longitudinal: bool = True, +) -> Array: + """Apply ``propagator`` (nx, ny) to the transverse part of ``field`` (nx, ny, 2) in k-space + and return the result in x-space. The longitudinal part is not propagated; it is kept + inside the retained light ``band`` (LPSE zeroes every component outside the band and the + anti-aliasing region, ``LightSolver.cpp:3720-3735``) or dropped altogether when + ``keep_longitudinal`` is False.""" + fx_k = jnp.fft.fft2(field[..., 0]) + fy_k = jnp.fft.fft2(field[..., 1]) + kdote = (kx[:, None] * fx_k + ky[None, :] * fy_k) * one_over_k_sq + lx_k, ly_k = kx[:, None] * kdote, ky[None, :] * kdote + tx_k, ty_k = fx_k - lx_k, fy_k - ly_k + l_factor = band if keep_longitudinal else 0.0 + fx_k = l_factor * lx_k + propagator * tx_k + fy_k = l_factor * ly_k + propagator * ty_k + return jnp.stack([jnp.fft.ifft2(fx_k), jnp.fft.ifft2(fy_k)], axis=-1) + + +class SpectralRamanLight(RamanLight): + """Raman light E1 with a prescribed pump, spectral propagator.""" + + def __init__(self, cfg: dict): + super().__init__(cfg) + grid = cfg["grid"] + self.kx_arr = jnp.asarray(grid["kx"]) + self.ky_arr = jnp.asarray(grid["ky"]) + k_sq = np.asarray(grid["kx"])[:, None] ** 2 + np.asarray(grid["ky"])[None, :] ** 2 + self.one_over_k_sq = jnp.asarray(np.where(k_sq > 0, 1.0 / np.where(k_sq > 0, k_sq, 1.0), 0.0)) + self.light_band = jnp.asarray(self._light_band(cfg, k_sq)) + self.propagator1 = jnp.exp(-1j * self.dt_l * self.c**2 / (2.0 * self.w1) * jnp.asarray(k_sq)) * self.light_band + # x-space scattering potential over one sub-step (detuning); the IAW part and the + # absorption are applied per call + self.detune1 = jnp.exp(self.dt_l * self.linear_coeff) + if self.seed_enabled: + seed = cfg["drivers"]["E1"]["derived"] + width = seed.get("injector_width", np.pi / self.k1_inject) + self.seed_profile = jnp.asarray( + gaussian_injector_profile(np.asarray(self.x), float(self.x[self.i1]), width, self.dx) + ) + + @staticmethod + def _light_band(cfg: dict, k_sq: np.ndarray) -> np.ndarray: + """Retained light band: the grid's anti-aliasing mask (without the EPW's + ``terms.epw.max_wavenumber``), optionally capped at ``terms.light.max_wavenumber * k0`` + (LPSE ``{laser|raman}.maxWavenumber``).""" + band = np.asarray(cfg["grid"]["light_low_pass_filter_grid"]) > 0.0 + cap = cfg["terms"].get("light", {}).get("max_wavenumber") + if cap is not None: + k0 = cfg["units"]["derived"]["w0"] / cfg["units"]["derived"]["c"] + band = band & (np.sqrt(k_sq) < float(cap) * k0) + return np.where(band, 1.0, 0.0) + + def calc_seed_source(self, t: float, seed_args: dict) -> Array: + """Smooth injector for the leftward seed: returns the (nx, ny) source added to E1_y.""" + return smooth_seed_source( + t, + seed_args, + w=self.w1, + c=self.c, + x=self.x, + y=self.y, + x_inject=self.x[self.i1], + wpe_sq=self.wpe_sq_i1, + profile=self.seed_profile, + ) + + def absorption_factor(self, rate, density_ratio, iaw_density): + """``exp(-nu dt_l (n/nc_w)^2)`` with the IAW perturbation included.""" + if rate is None: + return 1.0 + n_over_nc = density_ratio if iaw_density is None else density_ratio * (1.0 + iaw_density / self.n_over_env) + return jnp.exp(-rate * self.dt_l * n_over_nc**2)[..., None] + + def __call__(self, t, E1, E0_fn, phi_k, seed_args, iaw_density=None): + seed_args = seed_args if self.seed_enabled else None + laplacian_phi = jnp.fft.ifft2(-self.k_sq * phi_k) + detune = self.detune1 + if iaw_density is not None: + detune = detune * jnp.exp(-1j * self.wp0**2 / (2.0 * self.w1) * iaw_density * self.dt_l) + absorb = self.absorption_factor(self.absorption_rate1, self.n_over_nc1, iaw_density) + + def substep(i, E1): + t_i = t + i * self.dt_l + # x-space: scattering potential (detuning + absorption), then the sources (Euler) + E1 = E1 * detune[..., None] * absorb + E0 = E0_fn(t_i) + coupling = self.srs_coeff * jnp.conj(laplacian_phi)[..., None] * E0 + if self.transverse_source: + coupling = transverse_part(coupling, self.kx_arr, self.ky_arr, self.one_over_k_sq) + E1 = E1 + self.dt_l * coupling + if seed_args is not None: + E1 = E1.at[..., 1].add(self.dt_l * self.calc_seed_source(t_i, seed_args)) + # k-space: exact transverse propagation; the longitudinal part is kept only + # inside the light band + E1 = transverse_propagate( + E1, self.kx_arr, self.ky_arr, self.one_over_k_sq, self.propagator1, self.light_band + ) + return E1 * self.sub_boundary[..., None] + + return lax.fori_loop(0, self.n_sub, substep, E1) + + +class SmoothPumpInjector: + """LPSE's smooth pump injector for an evolved pump: oblique and multiple beams (snapped + transverse wavenumber, fractions, phases, frequency offsets), the transverse super-Gaussian + profile, Kubo-Anderson bandwidth and the pulse table, summed over colors. Shared by + ``SpectralCoupledLight`` and the combined solver (``combined.CombinedSolver``), which inject + the same pump; the host sets ``w0, c, x, dx, i0, n_src, E0_source, pump_turn_on_time`` + before calling ``init_pump_injector``.""" + + def init_pump_injector(self, cfg: dict): + pump = cfg["drivers"]["E0"]["derived"] + k0_inject = self.w0 / self.c * np.sqrt(1.0 - self.n_src) + # in-plane angle of incidence (drivers.E0.angle): the transverse wavenumber is snapped + # to the periodic y grid, kx follows from the local dispersion per color, and the + # field is polarised perpendicular to the snapped k (LPSE beam direction / polarization 0) + dky = 2.0 * np.pi / (cfg["grid"]["ny"] * cfg["grid"]["dy"]) + angles = np.atleast_1d(np.asarray(pump.get("beam_angle", [pump.get("angle", 0.0)]), dtype=np.float64)) + ky_beams = np.round(k0_inject * np.sin(angles) / dky) * dky if cfg["grid"]["ny"] > 1 else np.zeros_like(angles) + if np.any(np.abs(ky_beams) >= k0_inject): + raise ValueError( + f"drivers.E0 beam angles {np.rad2deg(angles)} deg cannot be launched at density {self.n_src}" + ) + self.ky_beams = jnp.asarray(ky_beams) + self.ky_pump = float(ky_beams[0]) + self.beam_fraction = jnp.asarray(np.atleast_1d(pump.get("beam_fraction", [1.0])), dtype=jnp.float64) + self.beam_phase = jnp.asarray(np.atleast_1d(pump.get("beam_phase", [0.0])), dtype=jnp.float64) + self.beam_delta_omega = jnp.asarray(np.atleast_1d(pump.get("beam_delta_omega", [0.0])), dtype=jnp.float64) + self.y_arr = jnp.asarray(cfg["grid"]["y"]) + # transverse super-Gaussian of the injected beams (LPSE laser.N.width / sgOrder / offset) + width = float(pump.get("beam_width", 0.0) or 0.0) + if width > 0.0 and cfg["grid"]["ny"] > 1: + order = float(pump.get("beam_sg_order", 2.0)) + y_rel = (np.asarray(cfg["grid"]["y"]) - float(pump.get("beam_offset", 0.0))) ** 2 / (2.0 * width**2) + self.beam_envelope_y = jnp.asarray(np.exp(-(y_rel ** (order / 2.0)))) + else: + self.beam_envelope_y = jnp.ones(cfg["grid"]["ny"]) + # Kubo-Anderson bandwidth: piecewise-constant random phase with correlation time 2 pi / (dW) + 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 + width = pump.get("injector_width", np.pi / k0_inject) + self.pump_profile = jnp.asarray( + gaussian_injector_profile(np.asarray(self.x), float(self.x[self.i0]), width, self.dx) + ) + + def calc_pump_source(self, t: float, pump_args: dict) -> Array: + """Smooth injector for the rightward (optionally oblique) pump, summed over colors: + the (nx, ny, 2) source added to E0 (the FD two-point rows are replaced).""" + 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) + phases = pump_args["phases"] # (nc, ny) + eps0 = 1.0 - self.n_src + if self.pulse_t is not None: + t_env = t_env * jnp.interp(t, self.pulse_t, self.pulse_amp) + amp = self.E0_source * jnp.sqrt(intensities) / eps0**0.25 * t_env * turn_on # (nc, ny) + amp = amp * self.beam_envelope_y[None, :] + color_phase = jnp.exp(-1j * self.w0 * delta_omega[:, None] * t + 1j * phases) # (nc, ny) + total = jnp.zeros((self.x.shape[0], self.y_arr.shape[0], 2), dtype=jnp.complex128) + for b in range(int(self.ky_beams.shape[0])): + ky_b = self.ky_beams[b] + dw_b = self.beam_delta_omega[b] + k0 = self.w0 / self.c * jnp.sqrt((1.0 + delta_omega + dw_b) ** 2 - self.n_src) # (nc,) + kx = jnp.sqrt(k0**2 - ky_b**2) # (nc,) longitudinal wavenumber of each color + v_gx = self.c**2 * kx / self.w0 # flux through the injector plane is set by the x group velocity + beam_phase = self.beam_phase[b] - self.w0 * dw_b * t + self.kap_phase(t, b) + amp_b = amp * jnp.sqrt(self.beam_fraction[b]) * jnp.exp(1j * beam_phase) + carrier = jnp.exp(1j * kx[:, None] * (self.x[None, :] - self.x[self.i0])) # (nc, nx) + source = (amp_b * v_gx[:, None] * color_phase)[:, None, :] * (carrier * self.pump_profile[None, :])[ + :, :, None + ] + source = jnp.sum(source, axis=0) * jnp.exp(1j * ky_b * self.y_arr)[None, :] # (nx, ny) + k_mag = jnp.sqrt(kx[0] ** 2 + ky_b**2) + pol = jnp.stack([-ky_b / k_mag, kx[0] / k_mag]) # perpendicular to the (first color's) k + total = total + source[..., None] * pol[None, None, :] + return total + + def kap_phase(self, t, beam: int): + """Kubo-Anderson process (LPSE bandwidth.KAP.frequency): the phase jumps to a new uniform + random value every 2 pi / (kap_bandwidth w0), with a deterministic hash of the interval.""" + if self.kap_bandwidth <= 0.0: + return 0.0 + tau = 2.0 * jnp.pi / (self.kap_bandwidth * self.w0) + index = jnp.floor(t / 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) + + +class SpectralCoupledLight(SmoothPumpInjector, CoupledLight): + """Evolved pump E0 and Raman light E1, spectral propagators, exact SRS exchange.""" + + def __init__(self, cfg: dict): + super().__init__(cfg) + grid = cfg["grid"] + self.kx_arr = jnp.asarray(grid["kx"]) + self.ky_arr = jnp.asarray(grid["ky"]) + k_sq = np.asarray(grid["kx"])[:, None] ** 2 + np.asarray(grid["ky"])[None, :] ** 2 + self.one_over_k_sq = jnp.asarray(np.where(k_sq > 0, 1.0 / np.where(k_sq > 0, k_sq, 1.0), 0.0)) + self.light_band = jnp.asarray(SpectralRamanLight._light_band(cfg, k_sq)) + self.propagator0 = jnp.exp(-1j * self.dt_l * self.c**2 / (2.0 * self.w0) * jnp.asarray(k_sq)) * self.light_band + self.propagator1 = jnp.exp(-1j * self.dt_l * self.c**2 / (2.0 * self.w1) * jnp.asarray(k_sq)) * self.light_band + self.detune0 = jnp.exp(self.dt_l * self.linear_coeff0) + self.detune1 = jnp.exp(self.dt_l * self.linear_coeff) + self.init_pump_injector(cfg) + if self.seed_enabled: + seed = cfg["drivers"]["E1"]["derived"] + width1 = seed.get("injector_width", np.pi / self.k1_inject) + self.seed_profile = jnp.asarray( + gaussian_injector_profile(np.asarray(self.x), float(self.x[self.i1]), width1, self.dx) + ) + + # the smooth injectors + calc_seed_source = SpectralRamanLight.calc_seed_source + absorption_factor = SpectralRamanLight.absorption_factor + + def __call__(self, t, E0, E1, phi_k, pump_args, seed_args, iaw_density=None): + seed_args = seed_args if self.seed_enabled else None + laplacian_phi = jnp.fft.ifft2(-self.k_sq * phi_k) + detune0, detune1 = self.detune0, self.detune1 + if iaw_density is not None: + detune0 = detune0 * jnp.exp(-1j * self.wp0**2 / (2.0 * self.w0) * iaw_density * self.dt_l) + detune1 = detune1 * jnp.exp(-1j * self.wp0**2 / (2.0 * self.w1) * iaw_density * self.dt_l) + absorb0 = self.absorption_factor(self.absorption_rate0, self.n_over_nc0, iaw_density) + absorb1 = self.absorption_factor(self.absorption_rate1, self.n_over_nc1, iaw_density) + exchange = self.srs_enabled + + def substep(i, fields): + E0, E1 = fields + t_i = t + i * self.dt_l + if exchange: + E0, E1 = self.couple(E0, E1, laplacian_phi, 0.5 * self.dt_l) + # x-space: scattering potentials, then the non-exchange sources (Euler) + E0 = E0 * detune0[..., None] * absorb0 + E1 = E1 * detune1[..., None] * absorb1 + if self.tpd_enabled: + E0 = E0 + self.dt_l * self.calc_tpd_depletion(t_i, phi_k) + E0 = E0 + self.dt_l * self.calc_pump_source(t_i, pump_args) + if seed_args is not None: + E1 = E1.at[..., 1].add(self.dt_l * self.calc_seed_source(t_i, seed_args)) + # k-space: exact transverse propagation. The exchange is an exact local rotation + # that cannot be projected term by term, so with transverse_source the fields + # themselves are kept transverse here (equivalent to projecting every source) + keep_l = not self.transverse_source + E0 = transverse_propagate( + E0, self.kx_arr, self.ky_arr, self.one_over_k_sq, self.propagator0, self.light_band, keep_l + ) + if exchange: + E1 = transverse_propagate( + E1, self.kx_arr, self.ky_arr, self.one_over_k_sq, self.propagator1, self.light_band, keep_l + ) + E0, E1 = self.couple(E0, E1, laplacian_phi, 0.5 * self.dt_l) + E0 = E0 * self.sub_boundary[..., None] + E1 = E1 * self.sub_boundary[..., None] + return (E0, E1) + + 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)) + E1 = jnp.fft.ifft2(jnp.fft.fft2(E1, axes=(0, 1)) * self.light_filter, axes=(0, 1)) + return E0, E1 diff --git a/adept/_lpse2d/core/vector_field.py b/adept/_lpse2d/core/vector_field.py index 9afaf46d..0808dccc 100644 --- a/adept/_lpse2d/core/vector_field.py +++ b/adept/_lpse2d/core/vector_field.py @@ -1,9 +1,10 @@ import numpy as np -from jax import Array +from jax import Array, lax from jax import numpy as jnp from adept._base_ import get_envelope from adept._lpse2d.core import epw, laser +from adept._lpse2d.core.epw import LEDGER_CHANNELS, LEDGER_KEY from adept._lpse2d.core.light import CoupledLight from adept._lpse2d.core.raman import RamanLight @@ -23,17 +24,47 @@ def __init__(self, cfg): self.cfg = cfg self.dt = cfg["grid"]["dt"] self.wp0 = cfg["units"]["derived"]["wp0"] - self.epw = epw.SpectralEPWSolver(cfg) self.light = laser.Light(cfg) + # terms.epw.solver: "separate" (the four envelope equations, MATLAB / LPSE spectral) + # or "combined" (LPSE lw.solver = combined: one wp0-enveloped field for Raman light + # and EPW, required by LPSE whenever TPD and SRS are both on) + self.epw_solver = str(cfg["terms"]["epw"].get("solver", "separate")) + # the separate EPW solver only on the separate path: in combined mode it is never stepped, + # and its setup (e.g. the thermal noise's EPW-damping requirement) would judge the + # combined field by the separate solver's damping + self.epw = epw.SpectralEPWSolver(cfg) if self.epw_solver == "separate" else None + if self.epw_solver == "combined": + from adept._lpse2d.core.combined import CombinedSolver + + self.combined = CombinedSolver(cfg) + elif self.epw_solver != "separate": + raise ValueError(f"terms.epw.solver must be 'separate' or 'combined', got {self.epw_solver!r}") # 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) + # terms.light.solver: the MATLAB staggered finite-difference scheme (fd) or the + # LPSE spectral propagator (spectral, spectral_light.py) + self.light_solver = str(cfg["terms"].get("light", {}).get("solver", "fd")) + if self.light_solver == "spectral": + from adept._lpse2d.core.spectral_light import SpectralCoupledLight, SpectralRamanLight + + coupled_cls, raman_cls = SpectralCoupledLight, SpectralRamanLight + elif self.light_solver == "fd": + coupled_cls, raman_cls = CoupledLight, RamanLight + else: + raise ValueError(f"terms.light.solver must be 'fd' or 'spectral', got {self.light_solver!r}") + if self.epw_solver == "combined": + # the combined solver evolves the pump and the Raman light itself; the separate light + # solvers are never stepped, and their w1-carrier setup (the seed injector's cutoff) + # would refuse combined decks the combined carrier wp0 accepts + self.coupled_light = None + self.raman = None + elif self.pump_depletion: + self.coupled_light = coupled_cls(cfg) self.raman = None else: - self.raman = RamanLight(cfg) if srs_on else None + self.raman = raman_cls(cfg) if srs_on else None if cfg["terms"].get("hpe", {}).get("active", False): from adept._lpse2d.core.hpe import HybridParticleEvolution @@ -48,6 +79,8 @@ def __init__(self, cfg): self.iaw = None # HPE particle/histogram keys are real and stay out of this list self.complex_state_vars = ["E0", "epw", "E1"] + # terms.epw.energy_ledger: accumulate the per-operation EPW energy changes in the state + self.energy_ledger = bool(cfg["terms"]["epw"].get("energy_ledger", False)) self.boundary_envelope = cfg["grid"]["absorbing_boundaries"] self.one_over_ksq = cfg["grid"]["one_over_ksq"] self.zero_mask = cfg["grid"]["zero_mask"] @@ -124,22 +157,71 @@ def E0_fn(this_t): return y + def combined_step(self, t, y, driver_args): + """One EPW step of the combined solver: pump (prescribed or evolved), the combined + Raman + EPW field, and the derived potential.""" + if self.pump_depletion: + E0_fn = None + elif "E0" in driver_args: + + 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"]) + + else: + E0_now = y["E0"] + + def E0_fn(this_t): + return E0_now + + y["E0"], y["E1"], y["epw"] = self.combined(t, y, driver_args, E0_fn) + return y + def __call__(self, t, y, args): # unpack y into complex128 new_y = self._unpack_y_(y) + if self.epw_solver == "combined": + new_y = self.combined_step(t, new_y, args["drivers"]) + if self.iaw is not None: + # the IAW ponderomotive drive sees the Raman light (transverse part) and the + # EPW (through the derived potential) separately, not the combined field + iaw_in = {**new_y, "E1": self.combined.transverse(new_y["E1"])} + if self.iaw.stride == 1: + iaw_out = self.iaw(iaw_in, t) + else: + step = jnp.round(t / self.dt).astype(int) + iaw_out = lax.cond(step % self.iaw.stride == 0, lambda yy: self.iaw(yy, t), lambda yy: yy, iaw_in) + new_y["iaw_density"] = iaw_out["iaw_density"] + new_y["iaw_velocity_divergence"] = iaw_out["iaw_velocity_divergence"] + if self.hpe is not None: + new_y = self.hpe(t, new_y) + y, new_y = self._pack_y_(y, new_y) + return new_y + # light split step new_y = self.light_split_step(t, new_y, args["drivers"]) + driver_delta = 0.0 if "E2" in args["drivers"]: + w_before = self.epw.energy(new_y["epw"]) new_y["epw"] += jnp.fft.fft2(self.dt * self.epw.driver(args["drivers"]["E2"], t)) - # epw split step - new_y["epw"] = self.epw(t, new_y, args) + driver_delta = self.epw.energy(new_y["epw"]) - w_before + # epw split step (with the per-operation energy deltas for the ledger) + new_y["epw"], deltas = self.epw.advance(t, new_y, args) + if self.energy_ledger: + deltas = deltas.at[LEDGER_CHANNELS.index("driver")].set(driver_delta) + new_y[LEDGER_KEY] = new_y[LEDGER_KEY] + deltas # ion-acoustic split step: the updated density is seen by the light and EPW # detuning terms on the next outer step, matching the MATLAB ordering if self.iaw is not None: - new_y = self.iaw(new_y) + if self.iaw.stride == 1: + new_y = self.iaw(new_y, t) + else: + # LPSE strides the IAW solver: advance by stride*dt every stride-th EPW step + step = jnp.round(t / self.dt).astype(int) + new_y = lax.cond(step % self.iaw.stride == 0, lambda yy: self.iaw(yy, t), lambda yy: yy, new_y) # particle push + Landau-damping feedback; the gamma_L written here is the # rate the EPW update applies on the next step (one-step lag) diff --git a/adept/_lpse2d/datamodel.py b/adept/_lpse2d/datamodel.py index 27da461c..867490df 100644 --- a/adept/_lpse2d/datamodel.py +++ b/adept/_lpse2d/datamodel.py @@ -31,6 +31,21 @@ class DensityModel(BaseModel): gradient_scale_length: str | None = Field(default=None, alias="gradient scale length") max: float | None = None min: float | None = None + # lpse- (linear, exp, gaussian, inverse-power, quadratic, qd, gd, file): the original + # LPSE profiles between the N_min and N_max locations (helpers._lpse_density_profile) + min_location: str | None = None + max_location: str | None = None + min_location_y: str | None = None + max_location_y: str | None = None + geometry: Literal["cartesian", "spherical"] = "cartesian" + sg_order: float = 2.0 # LPSE sgOrder (gaussian, inverse-power, gd) + central_density: float | None = None # LPSE quadratic.centralDensity + dip_depth: float | None = None # LPSE dip.depth (qd, gd) + dip_width: str | None = None # LPSE dip.width (full width) + dip_offset: str | None = None # LPSE dip.offset from origin + origin: str | None = None # x of LPSE's box centre (default: the box centre) + max_density: float = 1.25 # LPSE maxBackgroundDensity clip + file: str | None = None # lpse-file: .npy / text table / LPSE grid file noise: NoiseModel | None = None @@ -89,6 +104,21 @@ class E0DriverModel(BaseModel): num_colors: int envelope: EnvelopeModel speckle: SpeckleModel | None = None + # in-plane angle of incidence from +x in degrees (LPSE laser.N.direction): the static pump + # is the grid mode nearest to k0 (cos a, sin a); the spectral injector launches at the + # y-snapped transverse wavenumber. Not with speckle or the FD injector. + angle: float = Field(default=0.0, gt=-90.0, lt=90.0) + # only in-plane ("p", LPSE polarization 0) is representable: lpse2d carries no z component + polarization: Literal["p"] = "p" + # LPSE laser.N.* beams: [{intensity: fraction or W/cm^2 (normalised), angle: deg, phase: rad, + # delta_omega: dW/W0}]; every beam carries every color. Spectral injector and static pump only. + beams: list[dict] | None = None + beam_width: str | None = None # LPSE laser.N.width: transverse (y) standard deviation of injected beams + beam_sg_order: float = 2.0 # LPSE laser.N.sgOrder of that transverse profile (2 = Gaussian) + beam_offset: str | None = None # y position of the beam centre (LPSE laser.N.offset) + kap_bandwidth: float = Field(default=0.0, ge=0.0, lt=1.0) # LPSE bandwidth.KAP.frequency (dW/W0) + kap_seed: int = 0 + pulse_file: str | None = None # LPSE laser.pulse.file: two-column (t_ps, relative amplitude) table class E1DriverModel(BaseModel): @@ -100,12 +130,14 @@ class E1DriverModel(BaseModel): """ 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) + delta_omega: float = 0.0 # seed frequency shift, a fraction of the Raman carrier (w1 = w0 - wp0; wp0 combined) 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 + # spectral light solver: Gaussian width of the smooth injector (default one local wavelength) + injector_width: str | None = None class DriversModel(BaseModel): @@ -126,6 +158,12 @@ class GridModel(BaseModel): boundary_abs_coeff: float boundary_width: str + # absorbing-layer profile: "tanh" (the MATLAB envelope, rate boundary_abs_coeff on the + # plateau) or "exp" (LPSE absorbingBoundaries.cpp: rate = boundary_max_rate * + # (exp(lambda s/L) - 1)/(exp(lambda) - 1) over the layer, per axis combined by max) + boundary_profile: Literal["tanh", "exp"] = "tanh" + boundary_max_rate: float = 200.0 # 1/ps, LPSE lw.abc.maxDampingRate default + boundary_lambda: float = 7.0 # LPSE abc.lambda default low_pass_filter: float dealias: str = "isotropic" dt: str @@ -157,6 +195,18 @@ class SaveModel(BaseModel): t: TimeSaveModel x: XSaveModel y: YSaveModel + # write the full solver state at grid.tmax to this path (or into the run's binary/ folder + # when true); `restart.file` resumes from it (LPSE checkpoint / --restart) + checkpoint: str | bool | None = None + # Thomson-scattering probes (LPSE thomsonScattering.N): [{k: [kx, ky] in units of k0, + # bandwidth: 0.1 (k0), field: epw | iaw}] -> series thomson__re/_im/_power + thomson: list[dict] | None = None + # save.fields.poynting: true adds s0_x, s0_y, s1_x, s1_y (LPSE laser/raman.save.S0) + poynting: bool = False + + +class RestartModel(BaseModel): + file: str class BoundaryModel(BaseModel): @@ -165,16 +215,44 @@ class BoundaryModel(BaseModel): class DampingModel(BaseModel): - collisions: bool + """EPW damping. ``landau_form`` selects the static Landau rate: ``matlab`` (the + prototype's prefactor), ``lpse`` (the C++ non-relativistic form) or + ``relativistic`` / ``relativistic_3d`` (LPSE's Maxwell-Juettner Bessel forms). + ``landau_lower_threshold`` (1/ps) zeroes the rate below it (LPSE + ``lw.landauDamping.lowerThreshold``); ``landau_multiplier`` is LPSE's static + ``LD_multiplier``.""" + + collisions: bool | float landau: bool + landau_form: Literal["matlab", "lpse", "relativistic", "relativistic_2d", "relativistic_3d"] = "matlab" + landau_lower_threshold: float = 0.0 + landau_multiplier: float = 1.0 class SourceModel(BaseModel): + """EPW sources. ``noise_model: flat`` is the MATLAB per-step source ``dt * amplitude`` + with a random phase on every retained mode; ``thermal`` is LPSE's fluctuation- + dissipation source (see ``epw.noise_kick_spectrum``), optionally calibrated to the + electron temperature with ``noise_calibrate``. ``noise_max_wavenumber`` (units of k0) + is LPSE ``lw.noise.maxWavenumber``. ``tpd_form: lpse`` keeps LPSE's exact TPD + coefficient and its ``(w0/wp0 - 1)`` charge-density factor; ``matlab`` is the + prototype's ``w0 -> 2 wp0`` form (identical at envelope density 0.25).""" + noise: bool + noise_model: Literal["flat", "thermal"] = "flat" noise_amplitude: float = 1e-10 noise_seed: int | None = None + noise_calibrate: bool = False + noise_debye_factor: bool = True # thermal model: keep the 1/sqrt(1 + k^2 lambda_D^2) spectrum shape + noise_max_wavenumber: float | None = None tpd: bool + tpd_form: Literal["lpse", "matlab"] = "lpse" srs: bool = False + # high-k filter on the light entering the SRS source (LPSE lw.kFilter.enable / .scale). + # The cutoff assumes zero detuning, so it removes the resonant mode in boxes below the + # envelope density; LPSE leaves it off by default and translated decks follow suit. + srs_k_filter: bool = True + srs_k_filter_scale: float = Field(default=1.2, ge=1.0, le=10.0) class EPWModel(BaseModel): @@ -183,6 +261,16 @@ class EPWModel(BaseModel): density_gradient: bool linear: bool source: SourceModel + # LPSE lw.maxWavenumber: hard cap |k| < max_wavenumber * k0 on the retained EPW band + max_wavenumber: float | None = None + # per-operation EPW energy ledger accumulated in the state and reported in the default + # series (epw_ledger_, epw_ledger_closure); LPSE asserts the same closure + energy_ledger: bool = False + # "separate": the four envelope equations (MATLAB / LPSE spectral path); "combined": + # LPSE lw.solver = combined, one wp0-enveloped field carrying the Raman light + # (transverse part) and the EPW (longitudinal part), required by LPSE when TPD and + # SRS are both on (see core/combined.py) + solver: Literal["separate", "combined"] = "separate" class LightModel(BaseModel): @@ -198,22 +286,78 @@ class LightModel(BaseModel): # optional isotropic low-pass filter on E0/E1 once per EPW step, as a fraction of # the grid Nyquist wavenumber pi/dx (None = off) filter: float | None = None + # TPD pump depletion: transverse (divergence-free) projection of E_h div(E_h) in + # k-space (LPSE LwSolver::makeExyzDivE) and the optional LPSE lw.kFilter + # (|k| < 1.2 k0 sqrt(1 - n_min)) on the same term + tpd_projection: bool = True + tpd_k_filter: bool = False + # light propagator: "fd" (MATLAB staggered scheme, sub-cycled to its CFL limit) or + # "spectral" (LPSE exact k-space propagator with L/T projection, no CFL limit) + solver: Literal["fd", "spectral"] = "fd" + # retained light band cap |k| < max_wavenumber * k0 (spectral solver; LPSE maxWavenumber) + max_wavenumber: float | None = None + # project the SRS light sources onto their transverse part every sub-step (LPSE + # takeTransversePartOfSourceTerms; for the coupled spectral solver the fields are kept + # transverse instead). No light propagator moves the longitudinal part of E1, so + # without this it accumulates the source and pairs with the EPW in a spurious + # two-wave instability (41/ps energy growth against LPSE's 5.8/ps in test_006) + transverse_source: bool = True + # collisional (inverse-bremsstrahlung) absorption: false, true (NRL formula as in LPSE) + # or the amplitude rate at nc in 1/ps + absorption: bool | float = False + # peak amplitude damping rate (1/ps) of the light fields' absorbing layers with the exp + # profile (LPSE {laser|raman}.evolution.abc.maxDampingRate, default 5e3); None = 5e3 for the + # exp profile and the EPW absorber (boundary_abs_coeff) for the tanh profile + boundary_max_rate: float | None = None class IAWDampingModel(BaseModel): - """Ion-acoustic damping parameters from the MATLAB LPSE model.""" + """Ion-acoustic damping. ``landau_form: simplified`` is ``gamma = landau * cs * |k|`` + (LPSE ``isSimplified``); ``full`` is the Z-generalized Krall-Trivelpiece rate with its + ``k lambda_D`` and ``Z Te/Ti`` dependence (``iaw.ion_landau_rate``), in which case + ``landau`` is ignored. ``collisions`` (1/ps) damps ``n`` as ``(1 - nu dt)`` in the + explicit solver and ``div v`` as ``exp(-2 nu dt)`` in the spectral one (LPSE).""" collisions: float = 1.0e-5 # density damping rate, 1/ps landau: float = 0.1 # gamma_iaw = landau * cs * |k| + landau_form: Literal["simplified", "full"] = "simplified" + + +class ThermalFilamentationModel(BaseModel): + """LPSE thermalFil.*: inverse-bremsstrahlung heating of the spatially varying part of each + wave's intensity, balanced by Spitzer conduction, driving the ion flow through the electron + pressure (see IonAcousticWave._init_thermal_filamentation).""" + + model_config = ConfigDict(populate_by_name=True) + + laser: bool = False + raman: bool = False + lw: bool = False + nonlocal_: bool = Field(default=False, alias="nonlocal") # k^(4/3) correction + conductivity_multiplier: float = 1.0 class IAWModel(BaseModel): - """Ion-acoustic density/velocity-divergence evolution and ponderomotive drive.""" + """Ion-acoustic density/velocity-divergence evolution and ponderomotive drive. + + ``solver: explicit`` is the MATLAB kick/drift split step (stable for + ``omega_max dt < 2``); ``spectral`` is LPSE's exact per-mode damped-oscillator + propagator (unconditionally stable), which also supports a uniform background + ``flow`` ([Mach_x, Mach_y], Doppler phase), advancing the IAW only every ``stride`` + EPW steps, and the LPSE fluctuation-dissipation ``noise`` source on ``div v``.""" active: bool = False + solver: Literal["explicit", "spectral"] = "explicit" boundary: BoundaryModel | None = None # defaults to terms.epw.boundary + boundary_max_rate: float | None = None # exp absorber peak rate (1/ps); default half the EPW one damping: IAWDampingModel = IAWDampingModel() max_density_perturbation: float | None = None + flow: list[float] | None = None + stride: int = 1 + noise: bool = False + noise_amplitude: float = 1.0 + noise_seed: int | None = None + thermal_filamentation: ThermalFilamentationModel | None = None class HPEModel(BaseModel): @@ -237,6 +381,20 @@ class HPEModel(BaseModel): feedback: bool = True # False = control run: particles evolve, damping stays analytic seed: int = 42 omega_res: str = "bohm_gross" # resonance v_phi convention: "bohm_gross" or "wp0" + # ---- LPSE HPE controls and instruments (ElectronTracker.cu readParameters) + gamma_limit_damping: float = 1500.0 # hpe.gammaLimit.damping, 1/ps: applied rate <= this + gamma_limit_growth: float = 1500.0 # hpe.gammaLimit.growth, 1/ps: applied rate >= -this when allow_growth + allow_growth: bool = False # hpe.allowGrowth: negative (inverse Landau) rates allowed + # hpe.thermalizationProbability (x, y walls): every crossing is counted by the instruments, then the + # particle is thermalized with this probability or passes through periodically (LPSE particle walls + # are independent of the field boundaries). None = 1 at absorbing field boundaries, 0 at periodic + thermalization_probability: list[float] | None = None + magnetic_field: float = 0.0 # hpe.magneticField, tesla, out of plane (2-D push only) + energy_conservation: bool = False # hpe.enforceEnergyConservation: LD multiplier + energy_conservation_steps: float = 1.0 # hpe.numStepsToAverageEnergyChange + flux_bins: list[float] | None = None # keV edges of the wall-flux instrument (default 0, 50, 100, inf) + cone_angle: float | None = None # hpe.metrics.power: acceptance half-angle in degrees (None = off) + cone_direction: list[float] = [1.0, 0.0] # hpe.metrics.power direction class TermsModel(BaseModel): @@ -271,3 +429,4 @@ class ConfigModel(BaseModel): solver: str terms: TermsModel units: UnitsModel + restart: RestartModel | None = None diff --git a/adept/_lpse2d/diagnostics.py b/adept/_lpse2d/diagnostics.py index f33a6bda..06fb98b0 100644 --- a/adept/_lpse2d/diagnostics.py +++ b/adept/_lpse2d/diagnostics.py @@ -141,17 +141,18 @@ def _lin_fit(j0: int, j1: int) -> dict | None: def expected_noise_energy(cfg: dict, t_ps: float) -> float: """Expected EPW energy of the pure noise-driven field at time ``t_ps`` (OSIRIS units). - Each retained k-mode is a damped random walk: the solver adds ``dt * A * e^{i phi}`` - every step and damps by ``e^{-(gamma_k + nu) dt}``, so + Each retained k-mode is a damped random walk: the solver adds ``D_k e^{i phi}`` every + noise step ``dt`` (``epw.noise_step``: the EPW step, or the light sub-step for the combined + solver) and damps by ``e^{-(gamma_k + nu) dt}``, so - <|phi_k|^2>(t) = (dt A)^2 * (1 - e^{-2 g t}) / (1 - e^{-2 g dt}), g = gamma_k + nu + <|phi_k|^2>(t) = D_k^2 * (1 - e^{-2 g t}) / (1 - e^{-2 g dt}), g = gamma_k + nu - with the ``g -> 0`` limit ``(dt A)^2 * t/dt``. Summed to the same OSIRIS-normalized + with the ``g -> 0`` limit ``D_k^2 * t/dt``. Summed to the same OSIRIS-normalized energy as the ``epw_energy`` save quantity (Parseval). Absorbing boundaries and detuning-induced mode mixing are neglected, so on absorbing/ramped boxes this is a mild overestimate -- conservative for use as a growth-fit noise floor. """ - from adept._lpse2d.core.epw import landau_damping_rate + from adept._lpse2d.core.epw import analytic_landau_rate, noise_kick_spectrum, noise_step grid = cfg["grid"] derived = cfg["units"]["derived"] @@ -159,15 +160,14 @@ def expected_noise_energy(cfg: dict, t_ps: float) -> float: ky = np.array(grid["ky"]) k_sq = kx[:, None] ** 2 + ky[None, :] ** 2 zero_mask = np.where(k_sq > 0, 1.0, 0.0) - if cfg["terms"]["epw"]["damping"].get("landau", True): - gamma = np.array(landau_damping_rate(k_sq, derived["wp0"], derived["vte_sq"], zero_mask)) - else: - gamma = np.zeros_like(k_sq) - g = gamma + derived.get("nu_coll", 0.0) * zero_mask - - dt = grid["dt"] - amp = float(cfg["terms"]["epw"]["source"].get("noise_amplitude", 1e-10)) - mode_amp_sq = (dt * amp) ** 2 * np.array(grid["low_pass_filter_grid"]) ** 2 * zero_mask + + # the step at which the active solver kicks and damps the noise (the EPW step, or the light + # sub-step for the combined solver) and the collisional rate it damps with + dt, nu = noise_step(cfg) + g = np.asarray(analytic_landau_rate(cfg)) + nu * zero_mask + # per-step kick amplitude of the configured noise model (flat: amplitude * sqrt(dt_EPW * dt) + # on the retained band; thermal: the fluctuation-dissipation kick), squared + mode_amp_sq = np.asarray(noise_kick_spectrum(cfg, dt=dt, nu_coll=nu)) ** 2 small = 2.0 * g * dt < 1.0e-12 # undamped modes: plain random walk, variance ~ t denom = np.where(small, 1.0, 1.0 - np.exp(-2.0 * g * dt)) diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index 8ba4b0ee..606e602e 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -1,5 +1,4 @@ import os -from functools import partial import matplotlib.pyplot as plt import scienceplots @@ -239,7 +238,7 @@ def get_derived_quantities(cfg: dict) -> dict: # Default save.*.t.tmin/tmax to grid values (preserves unit strings) for save_type in cfg.get("save", {}).keys(): - if "t" in cfg["save"][save_type]: + if isinstance(cfg["save"][save_type], dict) and "t" in cfg["save"][save_type]: t_cfg = cfg["save"][save_type]["t"] t_cfg.setdefault("tmin", cfg_grid.get("tmin", "0ps")) t_cfg.setdefault("tmax", cfg_grid["tmax"]) @@ -247,7 +246,7 @@ def get_derived_quantities(cfg: dict) -> dict: # cfg_grid["xmax"] = _Q(cfg_grid["xmax"]).to("um").value # cfg_grid["xmin"] = _Q(cfg_grid["xmin"]).to("um").value - if "linear" in cfg["density"]["basis"]: + if cfg["density"]["basis"] == "linear": L = _Q(cfg["density"]["gradient scale length"]).to("um").value nmax = cfg["density"]["max"] nmin = cfg["density"]["min"] @@ -318,15 +317,23 @@ def get_derived_quantities(cfg: dict) -> dict: inv_dy_sq = 0.0 if cfg_grid["ny"] == 1 else 1.0 / cfg_grid["dy"] ** 2 omega_max = 2.0 * cfg["units"]["derived"]["cs"] * np.sqrt(1.0 / cfg_grid["dx"] ** 2 + inv_dy_sq) - if omega_max * cfg_grid["dt"] >= 2.0: - raise ValueError( - "The ion-acoustic update is unstable: omega_iaw,max * grid.dt must be < 2 " - f"(got {omega_max * cfg_grid['dt']:.3g}). Reduce grid.dt or increase grid.dx." - ) + if iaw["stride"] < 1: + raise ValueError("terms.iaw.stride must be a positive integer") + if iaw["solver"] == "explicit": + if iaw["stride"] != 1: + raise ValueError("terms.iaw.stride > 1 requires terms.iaw.solver: spectral") + if iaw["flow"] is not None: + raise ValueError("terms.iaw.flow requires terms.iaw.solver: spectral") + if omega_max * cfg_grid["dt"] >= 2.0: + raise ValueError( + "The ion-acoustic update is unstable: omega_iaw,max * grid.dt must be < 2 " + f"(got {omega_max * cfg_grid['dt']:.3g}). Reduce grid.dt or increase grid.dx, " + "or use terms.iaw.solver: spectral (unconditionally stable)." + ) cfg["terms"]["iaw"] = iaw print( - "IAWs are on -- evolving density and velocity divergence with " - f"omega_iaw,max * dt = {omega_max * cfg_grid['dt']:.3g}" + f"IAWs are on ({iaw['solver']} solver, every {iaw['stride']} EPW step(s)) -- evolving density and " + f"velocity divergence with omega_iaw,max * dt = {omega_max * cfg_grid['dt'] * iaw['stride']:.3g}" ) # HPE (Follett-style test-particle Landau damping): resolve defaults, convert @@ -363,14 +370,22 @@ def get_derived_quantities(cfg: dict) -> dict: # light-wave options: defaults and validation (coupling scheme, filter) come from # the datamodel's LightModel; unknown keys are passed through untouched + from adept._lpse2d.datamodel import LightModel + light = cfg["terms"].get("light", {}) if light: - from adept._lpse2d.datamodel import LightModel - light = {**light, **LightModel(**light).model_dump()} cfg["terms"]["light"] = light pump_depletion = light.get("pump_depletion", False) - if not pump_depletion and (light.get("coupling", "explicit") != "explicit" or light.get("filter") is not None): + light_solver = light.get("solver", "fd") + combined = cfg["terms"]["epw"].get("solver", "separate") == "combined" + if combined and not light: + cfg["terms"]["light"] = light = {**LightModel().model_dump()} + if ( + not pump_depletion + and not combined + and (light.get("coupling", "explicit") != "explicit" or light.get("filter") is not None) + ): raise ValueError( "terms.light.coupling and terms.light.filter act on the coupled (pump-depletion) light solver " "and require terms.light.pump_depletion: true" @@ -378,6 +393,28 @@ def get_derived_quantities(cfg: dict) -> dict: source_terms = cfg["terms"]["epw"]["source"] srs_on = bool(source_terms.get("srs", False)) tpd_on = bool(source_terms.get("tpd", False)) + epw_solver = cfg["terms"]["epw"].get("solver", "separate") + if epw_solver == "combined": + if srs_on != tpd_on: + raise ValueError( + "terms.epw.solver: combined needs terms.epw.source.tpd and srs both on or both off " + "(LPSE: 'When lw.solver=combined, both SRS and TPD must be enabled or neither')" + ) + if light_solver != "spectral": + raise ValueError("terms.epw.solver: combined requires terms.light.solver: spectral") + if "E2" in cfg["drivers"]: + raise ValueError("terms.epw.solver: combined does not support the direct EPW driver drivers.E2") + if cfg["terms"]["epw"].get("energy_ledger", False): + raise ValueError("terms.epw.energy_ledger is only available with terms.epw.solver: separate") + if light.get("coupling", "explicit") != "explicit" or light.get("filter") is not None: + raise ValueError("terms.light.coupling / filter are options of the separate solver's FD light path") + elif srs_on and tpd_on: + print( + "WARNING: terms.epw.source.tpd and srs are both on with terms.epw.solver: separate. The original " + "LPSE refuses this ('Must use lw.solver=combined when both SRS and TPD are enabled'): the four " + "separate envelope equations are not valid for simultaneous TPD and SRS. Use terms.epw.solver: " + "combined for the LPSE formulation." + ) if pump_depletion: if not (srs_on or tpd_on): raise ValueError("terms.light.pump_depletion requires at least one of terms.epw.source.srs/tpd") @@ -394,12 +431,13 @@ def get_derived_quantities(cfg: dict) -> dict: "(the pump is launched by a boundary injector and must exit the box)" ) - if srs_on: + if srs_on and epw_solver != "combined" and cfg["terms"]["epw"]["source"].get("srs_k_filter", True): derived = cfg["units"]["derived"] - # The SRS source filter only passes wavenumbers up to the local Raman light - # wavenumber k1(n_min). If the box's minimum density reaches the w1 critical - # density, that band is empty and E1_filter (epw.py) silently zeroes every + # The separate EPW solver's SRS source filter only passes wavenumbers up to the local + # Raman light wavenumber k1(n_min) (the combined solver filters the pump instead and its + # light is cut off at n_env, not at the w1 critical density). If the box's minimum density + # reaches the w1 critical density, that band is empty and E1_filter (epw.py) silently zeroes every # mode of the source -- the run completes with reflectivity ~0 and reads as # "below threshold". The seeded path already dies with a clear error at the # injector; fail the noise-seeded path just as loudly here. @@ -452,7 +490,25 @@ def _worst_detuning_sq(w_carrier: float) -> float: ) evolved_carriers.append("pump") dt_max = min(dt_limits) - if "light_substeps" in cfg_grid: + if light_solver == "spectral": + # the exact k-space propagator has no CFL limit, but the smooth injectors add their + # source with an Euler step: the light must not cross more than one cell per sub-step + # (LPSE's spectral decks use raman.dt ~ dx/c), or the injected amplitude is wrong + c_light = cfg["units"]["derived"]["c"] + n_needed = int(np.ceil(c_light * cfg_grid["dt"] / cfg_grid["dx"])) + if "light_substeps" in cfg_grid: + n_sub = int(cfg_grid["light_substeps"]) + if n_sub < 1: + raise ValueError("grid.light_substeps must be a positive integer") + if n_sub < n_needed: + cells = c_light * cfg_grid["dt"] / n_sub / cfg_grid["dx"] + print( + f"WARNING: grid.light_substeps = {n_sub} lets light cross {cells:.1f} cells per " + f"sub-step; the spectral injectors want <= 1 ({n_needed} sub-steps)" + ) + else: + n_sub = max(1, n_needed) + elif "light_substeps" in cfg_grid: n_sub = int(cfg_grid["light_substeps"]) if cfg_grid["dt"] / n_sub > dt_max: raise ValueError( @@ -463,7 +519,10 @@ def _worst_detuning_sq(w_carrier: float) -> float: n_sub = int(np.ceil(cfg_grid["dt"] / (0.9 * dt_max))) cfg_grid["light_substeps"] = n_sub carriers = " + ".join(evolved_carriers) - print(f"{carriers} light is sub-cycled {n_sub}x per EPW step (dt_light limit {dt_max:.2e} ps)") + print( + f"{carriers} light ({light_solver} solver) is sub-cycled {n_sub}x per EPW step " + f"(FD dt_light limit {dt_max:.2e} ps)" + ) # change driver parameters to the right units for k in cfg["drivers"].keys(): @@ -492,6 +551,8 @@ def _worst_detuning_sq(w_carrier: float) -> float: "offset": offset, "yw": _Q(cfg["drivers"][k]["yw"]).to("um").value if "yw" in cfg["drivers"][k] else 0.0, } + if cfg["drivers"][k].get("injector_width") is not None: + cfg["drivers"][k]["derived"]["injector_width"] = _Q(cfg["drivers"][k]["injector_width"]).to("um").value continue if k == "E0" and pump_depletion: # boundary-injector parameters for the evolved pump: the injector sits at @@ -504,6 +565,65 @@ def _worst_detuning_sq(w_carrier: float) -> float: cfg["drivers"][k]["derived"]["turn_on_time"] = ( _Q(cfg["drivers"][k].get("turn_on_time", "10fs")).to("ps").value ) + if cfg["drivers"][k].get("injector_width") is not None: + cfg["drivers"][k]["derived"]["injector_width"] = _Q(cfg["drivers"][k]["injector_width"]).to("um").value + if k == "E0": + angle_deg = float(cfg["drivers"][k].get("angle", 0.0)) + cfg["drivers"][k]["derived"]["angle"] = float(np.deg2rad(angle_deg)) + e0_cfg = cfg["drivers"][k] + beams = e0_cfg.get("beams") or [{"intensity": 1.0, "angle": angle_deg, "phase": 0.0, "delta_omega": 0.0}] + fraction = np.array([float(b.get("intensity", 1.0)) for b in beams], dtype=np.float64) + if np.any(fraction < 0) or fraction.sum() <= 0: + raise ValueError("drivers.E0.beams intensities must be non-negative with a positive sum") + cfg["drivers"][k]["derived"]["beam_fraction"] = fraction / fraction.sum() + cfg["drivers"][k]["derived"]["beam_angle"] = np.deg2rad( + np.array([float(b.get("angle", angle_deg)) for b in beams], dtype=np.float64) + ) + cfg["drivers"][k]["derived"]["beam_phase"] = np.array( + [float(b.get("phase", 0.0)) for b in beams], dtype=np.float64 + ) + cfg["drivers"][k]["derived"]["beam_delta_omega"] = np.array( + [float(b.get("delta_omega", 0.0)) for b in beams], dtype=np.float64 + ) + cfg["drivers"][k]["derived"]["beam_width"] = ( + _Q(e0_cfg["beam_width"]).to("um").value if e0_cfg.get("beam_width") else 0.0 + ) + cfg["drivers"][k]["derived"]["beam_sg_order"] = float(e0_cfg.get("beam_sg_order", 2.0)) + cfg["drivers"][k]["derived"]["beam_offset"] = ( + _Q(e0_cfg["beam_offset"]).to("um").value if e0_cfg.get("beam_offset") else 0.0 + ) + cfg["drivers"][k]["derived"]["kap_bandwidth"] = float(e0_cfg.get("kap_bandwidth", 0.0)) + cfg["drivers"][k]["derived"]["kap_seed"] = int(e0_cfg.get("kap_seed", 0)) + if e0_cfg.get("pulse_file"): + table = np.loadtxt(e0_cfg["pulse_file"], dtype=np.float64) + if table.ndim != 2 or table.shape[1] < 2: + raise ValueError("drivers.E0.pulse_file must be a two-column (t_ps, amplitude) table") + cfg["drivers"][k]["derived"]["pulse_t"] = table[:, 0] + cfg["drivers"][k]["derived"]["pulse_amp"] = table[:, 1] + if pump_depletion and cfg["terms"].get("light", {}).get("solver", "fd") != "spectral": + unused = [ + key + for key, on in ( + ("drivers.E0.kap_bandwidth", float(e0_cfg.get("kap_bandwidth", 0.0) or 0.0) > 0.0), + ("drivers.E0.pulse_file", bool(e0_cfg.get("pulse_file"))), + ) + if on + ] + if unused: + raise ValueError( + f"{' and '.join(unused)} with terms.light.pump_depletion need terms.light.solver: spectral " + "(the FD two-point pump injector applies neither the pulse table nor the KAP phases)" + ) + multi = len(beams) > 1 or any(float(b.get("angle", angle_deg)) != 0.0 for b in beams) + angle_deg = angle_deg if not multi else 1.0 # trips the same checks below + if angle_deg != 0.0: + if cfg["drivers"][k].get("speckle", {}).get("enabled", False): + raise ValueError("drivers.E0.angle / beams are not supported together with drivers.E0.speckle") + if pump_depletion and cfg["terms"].get("light", {}).get("solver", "fd") != "spectral": + raise ValueError( + "drivers.E0.angle / beams with terms.light.pump_depletion need terms.light.solver: spectral " + "(the FD two-point injector launches along +x only)" + ) 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 @@ -558,8 +678,12 @@ def _pump_k_support(cfg: dict) -> tuple[float, float]: ky = k0 * numerical_aperture else: ky = 0.0 + # an oblique pump (drivers.E0.angle) carries k0 sin(angle) in y + angle = np.deg2rad(float(cfg["drivers"]["E0"].get("angle", 0.0))) + kx_support = k0 * abs(np.cos(angle)) + ky = max(ky, k0 * abs(np.sin(angle))) - return float(k0), float(ky) + return float(kx_support), float(ky) def get_solver_quantities(cfg: dict) -> dict: @@ -620,28 +744,80 @@ def get_solver_quantities(cfg: dict) -> dict: boundary_width = _Q(cfg_grid["boundary_width"]).to("um").value rise = boundary_width / 5 + boundary_profile = str(cfg_grid.get("boundary_profile", "tanh")) + if boundary_profile not in ("tanh", "exp"): + raise ValueError(f"grid.boundary_profile must be 'tanh' or 'exp', got {boundary_profile!r}") + + def absorbing_rate(boundary, max_rate=None): + """Amplitude damping rate (1/ps) of the absorbing layers, shape (nx, ny). + + ``tanh``: MATLAB's envelope, ``boundary_abs_coeff * (1 - tanh-envelope)``. + ``exp``: LPSE absorbingBoundaries.cpp, ``rate = max_rate (e^{lambda s/L} - 1)/(e^lambda - 1)`` + with s the distance into the layer of width L, per axis, the two axes combined by max. + """ + if boundary_profile == "tanh": + if boundary["x"] == "absorbing": + left = cfg_grid["xmin"] + boundary_width + right = cfg_grid["xmax"] - boundary_width + envelope_x = get_envelope(rise, rise, left, right, cfg_grid["x"])[:, None] + else: + envelope_x = np.ones((cfg_grid["nx"], cfg_grid["ny"])) - def absorbing_boundary(boundary): - if boundary["x"] == "absorbing": - left = cfg_grid["xmin"] + boundary_width - right = cfg_grid["xmax"] - boundary_width - envelope_x = get_envelope(rise, rise, left, right, cfg_grid["x"])[:, None] - else: - envelope_x = np.ones((cfg_grid["nx"], cfg_grid["ny"])) + if boundary["y"] == "absorbing": + left = cfg_grid["ymin"] + boundary_width + right = cfg_grid["ymax"] - boundary_width + envelope_y = get_envelope(rise, rise, left, right, cfg_grid["y"])[None, :] + else: + envelope_y = np.ones((cfg_grid["nx"], cfg_grid["ny"])) - if boundary["y"] == "absorbing": - left = cfg_grid["ymin"] + boundary_width - right = cfg_grid["ymax"] - boundary_width - envelope_y = get_envelope(rise, rise, left, right, cfg_grid["y"])[None, :] - else: - envelope_y = np.ones((cfg_grid["nx"], cfg_grid["ny"])) + return float(cfg_grid["boundary_abs_coeff"]) * (1.0 - envelope_x * envelope_y) - return np.exp(-float(cfg_grid["boundary_abs_coeff"]) * cfg_grid["dt"] * (1.0 - envelope_x * envelope_y)) + lam = float(cfg_grid.get("boundary_lambda", 7.0)) + gamma_max = float(cfg_grid.get("boundary_max_rate", 200.0)) if max_rate is None else float(max_rate) + coeff = gamma_max / np.expm1(lam) - cfg_grid["absorbing_boundaries"] = absorbing_boundary(cfg["terms"]["epw"]["boundary"]) + def axis_rate(ax, lo, hi): + # LPSE measures the distance into the layer in whole cells from the edge cell + # (absorbingBoundaries.cpp), so the first/last cell carries the full rate and + # the layer spans boundary_width / dx cells + half = 0.5 * (ax[1] - ax[0]) if ax.size > 1 else 0.0 + s = np.maximum(lo + boundary_width + half - ax, ax - (hi - boundary_width - half)) + s = np.clip(s, 0.0, boundary_width) + return coeff * np.expm1(lam * s / boundary_width) + + rate_x = np.zeros((cfg_grid["nx"], cfg_grid["ny"])) + rate_y = np.zeros((cfg_grid["nx"], cfg_grid["ny"])) + if boundary["x"] == "absorbing": + rate_x = axis_rate(cfg_grid["x"], cfg_grid["xmin"], cfg_grid["xmax"])[:, None] * np.ones( + (1, cfg_grid["ny"]) + ) + if boundary["y"] == "absorbing": + rate_y = axis_rate(cfg_grid["y"], cfg_grid["ymin"], cfg_grid["ymax"])[None, :] * np.ones( + (cfg_grid["nx"], 1) + ) + return np.maximum(rate_x, rate_y) + + def absorbing_boundary(boundary, max_rate=None): + return np.exp(-absorbing_rate(boundary, max_rate) * cfg_grid["dt"]) + + cfg_grid["absorbing_rate"] = absorbing_rate(cfg["terms"]["epw"]["boundary"]) + cfg_grid["absorbing_boundaries"] = np.exp(-cfg_grid["absorbing_rate"] * cfg_grid["dt"]) + # the light fields get their own absorber strength (LPSE {laser|raman}.evolution.abc.maxDampingRate, + # default 5e3/ps against 200/ps for the EPW): light crosses a 3 um layer in 0.01 ps, so at the + # EPW rate the exp profile removes only ~25 % per crossing and the pump builds a coherent + # standing wave between the walls (1.45x the launched amplitude on the LPSE test_010 deck). + # The tanh profile keeps boundary_abs_coeff for both (no change for existing configs). + light_max_rate = cfg["terms"].get("light", {}).get("boundary_max_rate") + if light_max_rate is None and boundary_profile == "exp": + light_max_rate = 5.0e3 + cfg_grid["light_absorbing_boundaries"] = absorbing_boundary(cfg["terms"]["epw"]["boundary"], light_max_rate) iaw = cfg["terms"].get("iaw", {}) if iaw.get("active", False): - cfg_grid["iaw_absorbing_boundaries"] = absorbing_boundary(iaw["boundary"]) + # LPSE's IAW absorber default is half the EPW one (IawSolver.cpp abc.maxDampingRate = 100) + iaw_max_rate = iaw.get("boundary_max_rate") + if iaw_max_rate is None and boundary_profile == "exp": + iaw_max_rate = 0.5 * float(cfg_grid.get("boundary_max_rate", 200.0)) + cfg_grid["iaw_absorbing_boundaries"] = absorbing_boundary(iaw["boundary"], iaw_max_rate) cfg_grid["zero_mask"] = ( np.where(np.sqrt(cfg_grid["kx"][:, None] ** 2 + cfg_grid["ky"][None, :] ** 2) == 0, 0, 1) @@ -696,8 +872,31 @@ def absorbing_boundary(boundary): f"dealias='shifted-band': alias-free band is |kx| <= {kx_limit:.2f}, |ky| <= {ky_limit:.2f} 1/um " f"(Nyquist {kx_nyquist:.2f}, {ky_nyquist:.2f}); low_pass_filter caps |k| <= {cutoff:.2f}" ) + elif dealias == "rectangular": + # the original LPSE mask (grid.antiAliasing.range): the outer fraction of *each* k axis + # is zeroed, i.e. |kx| < low_pass_filter * kmax_x and |ky| < low_pass_filter * kmax_y + kx_nyquist = float(np.abs(cfg_grid["kx"]).max()) + ky_nyquist = float(np.abs(cfg_grid["ky"]).max()) + frac = float(cfg_grid["low_pass_filter"]) + band = (np.abs(cfg_grid["kx"])[:, None] < frac * kx_nyquist) & ( + (np.abs(cfg_grid["ky"])[None, :] < frac * ky_nyquist) | (cfg_grid["ny"] == 1) + ) + cfg_grid["low_pass_filter_grid"] = np.where(band, 1.0, 0.0) elif dealias != "isotropic": - raise ValueError(f"Unknown grid.dealias '{dealias}'. Choose 'isotropic' or 'shifted-band'.") + raise ValueError(f"Unknown grid.dealias '{dealias}'. Choose 'isotropic', 'shifted-band' or 'rectangular'.") + + # the separate light solvers keep the anti-aliased band without the EPW's cap (their own cap is + # terms.light.max_wavenumber; LPSE LightSolver::evolveSpectral takes lw.maxWavenumber only in + # combined mode, LightSolver.cpp:3689) + cfg_grid["light_low_pass_filter_grid"] = cfg_grid["low_pass_filter_grid"] + # LPSE lw.maxWavenumber: an additional hard cap on the retained EPW band, in units + # of the vacuum laser wavenumber k0 = w0/c + max_wavenumber = cfg["terms"]["epw"].get("max_wavenumber") + if max_wavenumber is not None: + k0_vac = cfg["units"]["derived"]["w0"] / cfg["units"]["derived"]["c"] + cfg_grid["low_pass_filter_grid"] = cfg_grid["low_pass_filter_grid"] * np.where( + k_mag < float(max_wavenumber) * k0_vac, 1.0, 0.0 + ) retained = float(np.mean(cfg_grid["low_pass_filter_grid"] > 0)) debye_length = np.sqrt(cfg["units"]["derived"]["vte_sq"]) / cfg["units"]["derived"]["wp0"] @@ -800,12 +999,129 @@ def get_density_profile(cfg: dict) -> Array: amp = cfg["density"]["amplitude"] kk = cfg["density"]["wavenumber"] nprof = baseline * (1.0 + amp * np.sin(kk * cfg["grid"]["x"])) + + elif str(cfg["density"]["basis"]).startswith("lpse-"): + nprof = _lpse_density_profile(cfg) else: raise NotImplementedError return nprof +LPSE_DENSITY_SHAPES = ("linear", "exp", "gaussian", "inverse-power", "quadratic", "qd", "gd", "file") + + +def _lpse_density_profile(cfg: dict) -> np.ndarray: + """The original LPSE density profiles (``ZakharovSolver::backgroundDensityShape``), + ``density.basis: lpse-`` with shape in ``LPSE_DENSITY_SHAPES``. + + With ``r`` the distance from the N_max location along the unit vector towards the N_min + location (``geometry: cartesian``) or the radial distance from it (``spherical``), ``dr`` + the min-max separation and ``p = sg_order`` (default 2): + + - ``linear``: ``N_max + (N_min - N_max) r/dr``, clipped to [N_min, N_max] + - ``exp``: ``N_max exp(-r/L)``, ``L = dr/ln(N_max/N_min)``, clipped to [N_min, N_max] + - ``gaussian``: ``N_max exp(-|r/s|^p)``, ``s = dr / ln(N_max/N_min)^(1/p)`` (not clipped below) + - ``inverse-power``: ``N_min (dr/|r|)^p`` for ``|r| > r_c = dr (N_min/N_max)^(1/p)``, else ``N_max`` + - ``quadratic``: ``A0 + A1 x + A2 x^2`` through ``(0, central_density)``, ``(x_min, N_min)``, + ``(x_max, N_max)`` with x measured from ``origin`` (LPSE's box centre) + - ``qd`` / ``gd``: linear plus a parabolic / super-Gaussian dip of depth ``dip_depth``, full + width ``dip_width`` at ``dip_offset`` from ``origin`` + - ``file``: a (nx, ny) or (nx,) array from ``file`` (``.npy``, a text table, or an LPSE + grid file read with ``lpse_deck.read_frames``) + + Every profile is clipped at ``max_density`` (LPSE ``maxBackgroundDensity``, 1.25 n_c here).""" + d = cfg["density"] + shape = str(d["basis"])[5:] + if shape not in LPSE_DENSITY_SHAPES: + raise ValueError(f"density.basis lpse-{shape}: shape must be one of {LPSE_DENSITY_SHAPES}") + x = np.asarray(cfg["grid"]["x"], dtype=np.float64) + y = np.asarray(cfg["grid"]["y"], dtype=np.float64) + nx, ny = len(x), cfg["grid"]["ny"] + max_density = float(d.get("max_density", 1.25)) + + if shape == "file": + path = str(d["file"]) + if path.endswith(".npy"): + arr = np.load(path) + elif path.endswith((".txt", ".csv", ".dat")): + arr = np.loadtxt(path) + else: + from adept._lpse2d.lpse_deck import read_frames + + arr = np.real(read_frames(path)[-1][1]) # read_frames already returns (nx, ny) + arr = np.asarray(arr, dtype=np.float64) + if arr.ndim == 1: + arr = np.repeat(arr[:, None], ny, axis=-1) + if arr.shape != (nx, ny): + raise ValueError(f"density.file {path} has shape {arr.shape}, the grid is {(nx, ny)}") + return np.minimum(arr, max_density) + + n_min, n_max = float(d["min"]), float(d["max"]) + p = float(d.get("sg_order", 2.0)) + loc_min = np.array([_Q(d["min_location"]).to("um").value, _Q(d.get("min_location_y", "0um")).to("um").value]) + loc_max = np.array([_Q(d["max_location"]).to("um").value, _Q(d.get("max_location_y", "0um")).to("um").value]) + X, Y = np.meshgrid(x, y, indexing="ij") + dr = float(np.linalg.norm(loc_min - loc_max)) + if dr == 0.0 and n_min != n_max and shape not in ("quadratic",): + raise ValueError("density: min_location and max_location coincide but min != max") + if str(d.get("geometry", "cartesian")) == "spherical": + r = np.hypot(X - loc_max[0], Y - loc_max[1]) + else: + direction = (loc_min - loc_max) / dr if dr > 0 else np.array([1.0, 0.0]) + r = (X - loc_max[0]) * direction[0] + (Y - loc_max[1]) * direction[1] + lo, hi = min(n_min, n_max), max(n_min, n_max) + + def _linear(): + if dr == 0.0 or n_min == n_max: + return np.full_like(X, n_max) + return np.clip(n_max + (n_min - n_max) * r / dr, lo, hi) + + # the quadratic and the dips use LPSE's box-centred x; ``origin`` is that point here + origin = _Q(d["origin"]).to("um").value if d.get("origin") is not None else 0.5 * (x[0] + x[-1]) + xc = X - origin + + if shape == "linear": + nprof = _linear() + elif shape == "exp": + if n_min == n_max: + nprof = np.full_like(X, n_max) + else: + nprof = np.clip(n_max * np.exp(-r / (dr / np.log(n_max / n_min))), lo, hi) + elif shape == "gaussian": + if n_max <= n_min: + raise ValueError("density lpse-gaussian needs max > min") + sd = dr / np.log(n_max / n_min) ** (1.0 / p) + nprof = n_max * np.exp(-(np.abs(r / sd) ** p)) + elif shape == "inverse-power": + if n_max <= n_min: + raise ValueError("density lpse-inverse-power needs max > min") + rc = dr * (n_min / n_max) ** (1.0 / p) + safe = np.where(np.abs(r) > 0, np.abs(r), 1.0) + nprof = np.where(np.abs(r) > rc, n_min * (dr / safe) ** p, n_max) + elif shape == "quadratic": + a0 = float(d["central_density"]) + x1, x2 = loc_min[0] - origin, loc_max[0] - origin + if x1 == 0.0 or x2 == 0.0 or x1 == x2: + raise ValueError("density lpse-quadratic: min/max locations must differ from each other and from origin") + a2 = (n_max - n_min * x2 / x1 + a0 * (x2 / x1 - 1.0)) / (x2**2 - x1 * x2) + a1 = (n_min - a0 - a2 * x1**2) / x1 + nprof = a0 + a1 * xc + a2 * xc**2 + elif shape in ("qd", "gd"): + dn = float(d["dip_depth"]) + wd = 0.5 * _Q(d["dip_width"]).to("um").value + xd = _Q(d.get("dip_offset", "0um")).to("um").value + nprof = _linear() + if shape == "qd": + inside = np.abs(xc - xd) <= wd + nprof = nprof + np.where(inside, dn * (((xc - xd) / wd) ** 2 - 1.0), 0.0) + else: + nprof = nprof - dn * np.exp(-(np.abs((xc - xd) / wd) ** p)) + if np.any(nprof < 0): + raise ValueError("density: the LPSE profile is negative somewhere") + return np.minimum(nprof, max_density) + + def plot_fields(fields, td): t_skip = int(fields.coords["t (ps)"].data.size // 8) t_skip = t_skip if t_skip > 1 else 1 @@ -907,6 +1223,16 @@ def post_process(result, cfg: dict, td: str) -> tuple[xr.Dataset, xr.Dataset]: t0 = time.time() kfields, fields = make_field_xarrays(cfg, result.ts["fields"], result.ys["fields"], td) series = make_series_xarrays(cfg, result.ts["default"], result.ys["default"], td) + if "checkpoint" in result.ys: + target = cfg["save"]["checkpoint"] + path = target if isinstance(target, str) else os.path.join(td, "binary", "checkpoint.npz") + np.savez( + path, + t=np.asarray(result.ts["checkpoint"][-1]), + **{k: np.asarray(v[-1]) for k, v in result.ys["checkpoint"].items()}, + ) + cfg["save"]["checkpoint_file"] = path + print(f"checkpoint written to {path}") metrics["write_time"] = time.time() - t0 os.makedirs(os.path.join(td, "plots")) @@ -1041,14 +1367,29 @@ def make_series_xarrays(cfg, this_t, state, td): def make_field_xarrays(cfg, this_t, state, td): + """Build (and write) the x-space ``fields`` and k-space ``kfields`` datasets of the field save. + + Full-grid save: ``state["epw"]`` is the solver's k-space potential phi_k; phi = ifft2(phi_k) + and ex, ey = ifft2(-i k phi_k) (the solver's own field) are derived from it here. + + Interpolated save (``save.fields.x`` / ``y``): the save function has already put every field + in x-space on the save grid (``phi``, ``ex``, ``ey`` from the full-grid spectral fields), so + both saves report the same field at the same point. The k-space fields are the DFT on the save + grid, scaled by (nx * ny) / (nx_save * ny_save) so that a mode resolved on both grids has the + coefficient it has in the full-grid save (up to the phase of the save grid's first point), on + the save grid's k axes. + """ fld_save = cfg["save"]["fields"] - if "x" in fld_save: + interpolated = "x" in fld_save + if interpolated: kx = fld_save["kx"] ky = fld_save["ky"] xax = fld_save["x"]["ax"] yax = fld_save["y"]["ax"] nx = fld_save["x"]["ax"].size ny = fld_save["y"]["ax"].size + nx_grid, ny_grid = cfg["grid"]["nx"], cfg["grid"]["ny"] + k_norm = nx_grid * ny_grid / (nx * ny) else: kx = cfg["grid"]["kx"] @@ -1066,15 +1407,28 @@ def make_field_xarrays(cfg, this_t, state, td): yax_tuple = ("y (um)", yax) # 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): + if state["phi" if interpolated else "epw"].dtype in (np.float64, np.complex128): _complex = np.complex128 else: _complex = np.complex64 - phi_k_np = np.array(state["epw"]).view(_complex) - phi_vs_t = np.fft.ifft2(np.array(state["epw"]).view(_complex), axes=(1, 2)) - ex_k_np = -1j * kx[None, :, None] * phi_k_np - ey_k_np = -1j * ky[None, None, :] * phi_k_np + if interpolated: + phi_vs_t = np.array(state["phi"]).view(_complex) + ex_vs_t = np.array(state["ex"]).view(_complex) + ey_vs_t = np.array(state["ey"]).view(_complex) + phi_k_np = np.fft.fft2(phi_vs_t, axes=(1, 2)) * k_norm + ex_k_np = np.fft.fft2(ex_vs_t, axes=(1, 2)) * k_norm + ey_k_np = np.fft.fft2(ey_vs_t, axes=(1, 2)) * k_norm + ex_np = ex_vs_t + ey_np = ey_vs_t + else: + phi_k_np = np.array(state["epw"]).view(_complex) + phi_vs_t = np.fft.ifft2(np.array(state["epw"]).view(_complex), axes=(1, 2)) + ex_k_np = -1j * kx[None, :, None] * phi_k_np + ey_k_np = -1j * ky[None, None, :] * phi_k_np + # E = -grad(phi) = ifft2(-i k phi_k), the solver's own field (SpectralEPWSolver.phi_k_to_e_fields) + ex_np = np.fft.ifft2(ex_k_np, axes=(1, 2)) + ey_np = np.fft.ifft2(ey_k_np, axes=(1, 2)) phi_k = xr.DataArray( np.fft.fftshift(phi_k_np, axes=(1, 2)), @@ -1089,8 +1443,8 @@ def make_field_xarrays(cfg, this_t, state, td): coords=(tax_tuple, (r"kx ($kc\omega_0^{-1}$)", shift_kx), (r"ky ($kc\omega_0^{-1}$)", shift_ky)), ) phi_x = xr.DataArray(phi_vs_t, coords=(tax_tuple, xax_tuple, yax_tuple)) - ex = xr.DataArray(np.fft.ifft2(ex_k_np, axes=(1, 2)) / nx / ny * 4, coords=(tax_tuple, xax_tuple, yax_tuple)) - ey = xr.DataArray(np.fft.ifft2(ey_k_np, axes=(1, 2)) / nx / ny * 4, coords=(tax_tuple, xax_tuple, yax_tuple)) + ex = xr.DataArray(ex_np, coords=(tax_tuple, xax_tuple, yax_tuple)) + ey = xr.DataArray(ey_np, coords=(tax_tuple, xax_tuple, yax_tuple)) e0x = xr.DataArray(np.array(state["E0"]).view(_complex)[..., 0], coords=(tax_tuple, xax_tuple, yax_tuple)) e0y = xr.DataArray(np.array(state["E0"]).view(_complex)[..., 1], coords=(tax_tuple, xax_tuple, yax_tuple)) e1x = xr.DataArray(np.array(state["E1"]).view(_complex)[..., 0], coords=(tax_tuple, xax_tuple, yax_tuple)) @@ -1123,6 +1477,27 @@ def make_field_xarrays(cfg, this_t, state, td): ) kfield_data = {"phi": phi_k, "ex": ex_k, "ey": ey_k} + poynting = {} + if cfg["save"]["fields"].get("poynting", False): + # LPSE {laser|raman}.save.S0: the envelope energy-flux density S_j = (c^2/w) Im(E* . d_j E) + # (v_g |E|^2 for a plane wave), in field^2 * um/ps, on the save grid + derived = cfg["units"]["derived"] + for name, arr, w in ( + ("s0", np.array(state["E0"]).view(_complex), derived["w0"]), + ("s1", np.array(state["E1"]).view(_complex), _raman_carrier(cfg)), + ): + dxs = float(xax[1] - xax[0]) if len(xax) > 1 else 1.0 + dys = float(yax[1] - yax[0]) if len(yax) > 1 else 1.0 + sx = np.zeros(arr.shape[:-1]) + sy = np.zeros(arr.shape[:-1]) + for comp in range(2): + e = arr[..., comp] + sx += np.imag(np.conj(e) * np.gradient(e, dxs, axis=1)) + if arr.shape[2] > 1: + sy += np.imag(np.conj(e) * np.gradient(e, dys, axis=2)) + factor = derived["c"] ** 2 / w + poynting[f"{name}_x"] = xr.DataArray(factor * sx, coords=(tax_tuple, xax_tuple, yax_tuple)) + poynting[f"{name}_y"] = xr.DataArray(factor * sy, coords=(tax_tuple, xax_tuple, yax_tuple)) field_data = { "phi": phi_x, "ex": ex, @@ -1132,6 +1507,7 @@ def make_field_xarrays(cfg, this_t, state, td): "e1_x": e1x, "e1_y": e1y, "background_density": background_density, + **poynting, } if "iaw_density" in state: iaw_density_np = np.asarray(state["iaw_density"]) @@ -1143,11 +1519,14 @@ def make_field_xarrays(cfg, this_t, state, td): (r"kx ($kc\omega_0^{-1}$)", shift_kx), (r"ky ($kc\omega_0^{-1}$)", shift_ky), ) - kfield_data["iaw_density"] = xr.DataArray( - np.fft.fftshift(np.fft.fft2(iaw_density_np, axes=(1, 2)), axes=(1, 2)), coords=k_coords - ) + iaw_density_k = np.fft.fft2(iaw_density_np, axes=(1, 2)) + iaw_velocity_k = np.fft.fft2(iaw_velocity_np, axes=(1, 2)) + if interpolated: + iaw_density_k = iaw_density_k * k_norm + iaw_velocity_k = iaw_velocity_k * k_norm + kfield_data["iaw_density"] = xr.DataArray(np.fft.fftshift(iaw_density_k, axes=(1, 2)), coords=k_coords) kfield_data["iaw_velocity_divergence"] = xr.DataArray( - np.fft.fftshift(np.fft.fft2(iaw_velocity_np, axes=(1, 2)), axes=(1, 2)), coords=k_coords + np.fft.fftshift(iaw_velocity_k, axes=(1, 2)), coords=k_coords ) kfields = xr.Dataset(kfield_data) @@ -1158,6 +1537,149 @@ def make_field_xarrays(cfg, this_t, state, td): return kfields, fields +def _raman_carrier(cfg: dict) -> float: + """The envelope carrier of the saved ``E1``: ``w1 = w0 - wp0`` for the separate Raman solver, + ``wp0`` for the combined one (``CombinedSolver.prop_T`` and its seed both envelope at ``wp0``). + The Raman flux ``(c^2/w) Im(E* . grad E)`` and its plane-wave corrections use it; the two agree + only at envelope density n_c/4.""" + derived = cfg["units"]["derived"] + return derived["wp0"] if cfg["terms"]["epw"].get("solver", "separate") == "combined" else derived["w1"] + + +def _raman_light_part(cfg: dict): + """``y -> y`` with ``E1`` reduced to the Raman light. The combined solver's ``E1`` also carries + the EPW as its longitudinal part (saved separately as ``epw``); the saved ``E1`` fields and + everything computed from them (the ``s1`` flux maps) are the transverse part, as in the + default series. Identity for the separate solver.""" + if cfg["terms"]["epw"].get("solver", "separate") != "combined": + return lambda y: y + from adept._lpse2d.core.combined import transverse_part + + kx = np.asarray(cfg["grid"]["kx"]) + ky = np.asarray(cfg["grid"]["ky"]) + k_sq = kx[:, None] ** 2 + ky[None, :] ** 2 + one_over_k_sq = jnp.asarray(np.where(k_sq > 0, 1.0 / np.where(k_sq > 0, k_sq, 1.0), 0.0)) + kx, ky = jnp.asarray(kx), jnp.asarray(ky) + + def apply(y): + if "E1" not in y: + return y + return {**y, "E1": transverse_part(y["E1"].view(jnp.complex128), kx, ky, one_over_k_sq).view(jnp.float64)} + + return apply + + +def _interpolated_save_axis(grid_axis: np.ndarray, lo: float, hi: float, d_requested: float, name: str): + """Cell-centred save axis of n_save = round((hi - lo) / d_requested) points spanning [lo, hi]. + + Returns the axis and its actual spacing (hi - lo) / n_save. Refuses a save grid with more + points than the simulation grid: it would be finer than the grid, and its outermost points + would fall outside the outermost cell centres, where linear interpolation is undefined. + """ + key = f"save.fields.{name}.d{name}" + n_grid = grid_axis.size + length = hi - lo + if not d_requested > 0.0: + raise ValueError(f"{key} must be positive, got {d_requested} um") + n_save = max(1, round(length / d_requested)) + if n_save > n_grid: + raise ValueError( + f"{key} = {d_requested:.4g} um is finer than the simulation grid (d{name} = {length / n_grid:.4g} um): " + f"it gives {n_save} save points along {name} for {n_grid} cells. An interpolated field save must not be " + f"finer than the grid -- use {key} >= {length / n_grid:.4g} um, or drop save.fields.x and " + "save.fields.y to save the full grid." + ) + d_save = length / n_save + axis = np.linspace(lo + d_save / 2.0, hi - d_save / 2.0, n_save) + # n_save <= n_grid puts every save point inside [grid_axis[0], grid_axis[-1]]; clip only the + # last-bit rounding (n_save == n_grid), where interpax would otherwise return NaN + return np.clip(axis, grid_axis[0], grid_axis[-1]), d_save + + +def _interpolated_field_save(cfg: dict): + """Set up the save grid of ``save.fields.x.dx`` / ``save.fields.y.dy`` and return its save function. + + Every saved field is put on the save grid by linear interpolation in x-space: + + - ``epw``: the state is the k-space potential phi_k. It is transformed to x-space with the + solver's convention (``SpectralEPWSolver.phi_k_to_e_fields``: phi = ifft2(phi_k), + E = ifft2(-i k phi_k)) on the full grid, and phi, ex and ey are interpolated and saved as + ``phi``, ``ex`` and ``ey``. The fields are differentiated spectrally on the full grid rather + than on the save grid: the save grid's spectral derivative would multiply the interpolation + error by k and assign the wrong k to any content above the save grid's Nyquist wavenumber. + - ``E0``, ``E1`` and the real x-space fields (IAW) are interpolated directly. + + The axes and k axes of the save grid go in ``cfg["save"]["fields"]``; ``make_field_xarrays`` + builds k-space fields from them. + """ + from adept._lpse2d.core.hpe import PARTICLE_KEYS + + fld = cfg["save"]["fields"] + grid = cfg["grid"] + spatial_scale = cfg["units"]["derived"]["spatialScale"] + dx_requested = _Q(fld["x"]["dx"]).to("m").value / spatial_scale * 100 + dy_requested = _Q(fld["y"]["dy"]).to("m").value / spatial_scale * 100 + x_save, dx_save = _interpolated_save_axis(grid["x"], grid["xmin"], grid["xmax"], dx_requested, "x") + y_save, dy_save = _interpolated_save_axis(grid["y"], grid["ymin"], grid["ymax"], dy_requested, "y") + nx_save, ny_save = x_save.size, y_save.size + + fld["x"]["dx"] = dx_save + fld["x"]["ax"] = x_save + fld["kx"] = np.fft.fftfreq(nx_save, d=dx_save / 2.0 / np.pi) + fld["y"]["dy"] = dy_save + fld["y"]["ax"] = y_save + fld["ky"] = np.fft.fftfreq(ny_save, d=dy_save / 2.0 / np.pi) + print( + f"interpolated field save: {nx_save}x{ny_save} points (dx = {dx_save:.4g} um, dy = {dy_save:.4g} um) " + f"from the {grid['nx']}x{grid['ny']} grid" + ) + + x_src = jnp.asarray(grid["x"]) + kx = jnp.asarray(grid["kx"])[:, None] + ky = jnp.asarray(grid["ky"])[None, :] + if grid["ny"] == 1: + # Quasi-1D (single transverse cell; ny_save is then 1 too): interpax.interp2d needs + # >= 2 nodes per axis, so interpolate in x only and keep the single y-row. + def to_save_grid(f): + return interpax.interp1d(jnp.asarray(x_save), x_src, f[:, 0], method="linear")[:, None] + else: + xq, yq = np.meshgrid(x_save, y_save, indexing="ij") + xq, yq = jnp.asarray(xq.ravel()), jnp.asarray(yq.ravel()) + y_src = jnp.asarray(grid["y"]) + + def to_save_grid(f): + values = interpax.interp2d(xq, yq, x_src, y_src, f, method="linear") + return jnp.reshape(values, (nx_save, ny_save) + f.shape[2:]) + + def as_complex(v): + return v.view(jnp.complex128 if v.dtype == jnp.float64 else jnp.complex64) + + def save_func(t, y, args): + save_y = {} + for k, v in y.items(): + if k in PARTICLE_KEYS: + # particle arrays are (Np,) and gamma_L/epw_hist live in k/v space -- + # none of them fit the spatial interpolator; the histogram and the + # damping-reduction scalars are saved through the default series + continue + if k == "epw": + phi_k = as_complex(v) + x_space = jnp.stack( + [jnp.fft.ifft2(phi_k), jnp.fft.ifft2(-1j * kx * phi_k), jnp.fft.ifft2(-1j * ky * phi_k)], axis=-1 + ) + on_save_grid = to_save_grid(x_space) + for i, name in enumerate(("phi", "ex", "ey")): + save_y[name] = on_save_grid[..., i].view(v.dtype) + elif k in ("E0", "E1"): + save_y[k] = to_save_grid(as_complex(v)).view(v.dtype) + else: + save_y[k] = to_save_grid(v) + + return save_y + + return save_func + + def get_save_quantities(cfg: dict) -> dict: """ This function updates the config with the quantities required for the diagnostics and saving routines @@ -1167,6 +1689,7 @@ def get_save_quantities(cfg: dict) -> dict: """ # cfg["save"]["func"] = {**cfg["save"]["func"], **{"callable": get_save_func(cfg)}} + raman_light_part = _raman_light_part(cfg) tmin = _Q(cfg["save"]["fields"]["t"]["tmin"]).to("s").value / cfg["units"]["derived"]["timeScale"] tmax = _Q(cfg["save"]["fields"]["t"]["tmax"]).to("s").value / cfg["units"]["derived"]["timeScale"] dt = _Q(cfg["save"]["fields"]["t"]["dt"]).to("s").value / cfg["units"]["derived"]["timeScale"] @@ -1175,82 +1698,33 @@ def get_save_quantities(cfg: dict) -> dict: cfg["save"]["fields"]["t"]["dt"] = dt cfg["save"]["fields"]["t"]["ax"] = jnp.linspace(tmin, tmax, nt) - if "x" in cfg["save"]["fields"]: - xmin = cfg["grid"]["xmin"] - xmax = cfg["grid"]["xmax"] - dx = _Q(cfg["save"]["fields"]["x"]["dx"]).to("m").value / cfg["units"]["derived"]["spatialScale"] * 100 - nx = cfg["grid"]["nx"] - # nx = int((xmax - xmin) / dx) - cfg["save"]["fields"]["x"]["dx"] = dx - cfg["save"]["fields"]["x"]["ax"] = jnp.linspace(xmin + dx / 2.0, xmax - dx / 2.0, nx) - cfg["save"]["fields"]["kx"] = np.fft.fftfreq(nx, d=dx / 2.0 / np.pi) - - if "y" in cfg["save"]["fields"]: - ymin = cfg["grid"]["ymin"] - ymax = cfg["grid"]["ymax"] - dy = _Q(cfg["save"]["fields"]["y"]["dy"]).to("m").value / cfg["units"]["derived"]["spatialScale"] * 100 - ny = cfg["grid"]["ny"] - cfg["save"]["fields"]["y"]["dy"] = dy - cfg["save"]["fields"]["y"]["ax"] = jnp.linspace(ymin + dy / 2.0, ymax - dy / 2.0, ny) - cfg["save"]["fields"]["ky"] = np.fft.fftfreq(ny, d=dy / 2.0 / np.pi) - else: - raise NotImplementedError("Must specify y in save") - - xq, yq = jnp.meshgrid(cfg["save"]["fields"]["x"]["ax"], cfg["save"]["fields"]["y"]["ax"], indexing="ij") + if ("x" in cfg["save"]["fields"]) != ("y" in cfg["save"]["fields"]): + raise NotImplementedError("Must specify both save.fields.x and save.fields.y for an interpolated field save") - 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", + if "x" in cfg["save"]["fields"]: + interpolate = _interpolated_field_save(cfg) + n_save_x = cfg["save"]["fields"]["x"]["ax"].size + if cfg["save"]["fields"].get("poynting", False) and n_save_x < 2: + # the flux maps differentiate along x on the save grid; refuse here rather than after the solve + raise ValueError( + f"save.fields.poynting needs at least two saved x points, but save.fields.x.dx = " + f"{cfg['save']['fields']['x']['dx']:.4g} um gives {n_save_x}: use a smaller save.fields.x.dx, " + "or drop save.fields.x / y to save the full grid" ) def save_func(t, y, args): - from adept._lpse2d.core.hpe import PARTICLE_KEYS - - save_y = {} - for k, v in y.items(): - if k in PARTICLE_KEYS: - # particle arrays are (Np,) and gamma_L/epw_hist live in k/v space -- - # none of them fit the spatial interpolator; the histogram and the - # damping-reduction scalars are saved through the default series - continue - if k in ["E0", "E1"]: - cmplx_fld = v.view(jnp.complex128) - save_y[k] = jnp.concatenate( - [ - jnp.reshape(interpolator(f=cmplx_fld[..., ivec]), (nx, ny), order="F")[..., None] - for ivec in range(2) - ], - axis=-1, - ).view(jnp.float64) - - elif k == "epw": - cmplx_fld = v.view(jnp.complex128) - save_y[k] = jnp.reshape(interpolator(f=cmplx_fld), (nx, ny), order="F").view(jnp.float64) - else: - save_y[k] = jnp.reshape(interpolator(f=v), (nx, ny), order="F") - - return save_y + # the combined solver's E1 is saved as its transverse (Raman-light) part, as in the + # full-grid save; the energy ledger is a series quantity + y = raman_light_part(y) + return interpolate(t, {k: v for k, v in y.items() if k != "epw_ledger"}, args) else: def save_func(t, y, args): from adept._lpse2d.core.hpe import PARTICLE_KEYS - return {k: v for k, v in y.items() if k not in PARTICLE_KEYS} + y = raman_light_part(y) + return {k: v for k, v in y.items() if k not in PARTICLE_KEYS and k != "epw_ledger"} cfg["save"]["fields"]["func"] = save_func @@ -1260,7 +1734,7 @@ def save_func(t, y, args): def get_default_save_func(cfg): - from adept._lpse2d.core.epw import landau_damping_rate + from adept._lpse2d.core.epw import analytic_landau_rate srs_on = cfg["terms"]["epw"]["source"].get("srs", False) pump_evolved = cfg["terms"].get("light", {}).get("pump_depletion", False) @@ -1269,6 +1743,16 @@ def get_default_save_func(cfg): dt = cfg["grid"]["dt"] nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] kx, ky = cfg["grid"]["kx"], cfg["grid"]["ky"] + thomson_probes = list(cfg["save"].get("thomson") or []) + thomson_masks = [] + k0_vac = derived["w0"] / derived["c"] + for probe in thomson_probes: + if probe.get("field", "epw") == "iaw" and not iaw_on: + raise ValueError("save.thomson probe on the IAW density needs terms.iaw.active") + kxp, kyp = (float(v) * k0_vac for v in probe["k"]) + band = float(probe.get("bandwidth", 0.1)) * k0_vac + dist = np.sqrt((np.asarray(kx)[:, None] - kxp) ** 2 + (np.asarray(ky)[None, :] - kyp) ** 2) + thomson_masks.append(jnp.asarray(dist < band)) # 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 @@ -1282,13 +1766,8 @@ def get_default_save_func(cfg): # 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) - if cfg["terms"]["epw"]["damping"].get("landau", True): - gamma_total = np.array( - landau_damping_rate(jnp.array(k_sq), derived["wp0"], derived["vte_sq"], jnp.array(zero_mask)) - ) - else: - gamma_total = np.zeros_like(k_sq) - gamma_total = gamma_total + derived.get("nu_coll", 0.0) * zero_mask + # the same static rate (form, threshold, multiplier) the solver applies + gamma_total = np.array(analytic_landau_rate(cfg)) + 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 @@ -1301,9 +1780,22 @@ def get_default_save_func(cfg): # low-density end, which the previously hard-coded 2 understated by ~28%. energy_total_factor = 2.0 * (2.0 - np.array(cfg["grid"]["background_density"]) / cfg["units"]["envelope density"]) + # terms.epw.solver: combined -- the EPW is the combined field's longitudinal part, damped by the + # Landau rate in k-space and, with the whole field, by the Raman absorption nu_R (n/n_env + dn)^2 + # in x-space (CombinedSolver); the EPW collisions play no part + combined_epw = cfg["terms"]["epw"].get("solver", "separate") == "combined" + if combined_epw: + from adept._lpse2d.core.combined import combined_collision_rate + + nu_combined = combined_collision_rate(cfg) + n_over_env = jnp.asarray(np.array(cfg["grid"]["background_density"]) / cfg["units"]["envelope density"]) + energy_loss_factor = (1.0 - np.exp(-2.0 * np.array(analytic_landau_rate(cfg)) * dt)) / dt + # HPE: the damping is dynamic (y["gamma_L"]), so the dissipation diagnostic must # use the state's rate; also emit the hot-electron scalars and the histogram hpe_on = cfg["terms"].get("hpe", {}).get("active", False) + # with terms.hpe.allow_growth the applied rate can be negative (inverse Landau damping) + signed_loss = bool(hpe_on and cfg["terms"]["hpe"].get("allow_growth", False)) if hpe_on: from adept._lpse2d.core.hpe import resonance_arrays @@ -1311,7 +1803,7 @@ def get_default_save_func(cfg): n_p = int(cfg["terms"]["hpe"]["n_particles"]) w_tail = hpe_arrays["f_tail_frac"] / n_p # fraction of all electrons per particle c_light = derived["c"] - nu_coll_arr = derived.get("nu_coll", 0.0) * zero_mask + nu_coll_arr = 0.0 if combined_epw else derived.get("nu_coll", 0.0) * zero_mask gamma_an = np.array(hpe_arrays["gamma_analytic"]) mask_res = np.array(hpe_arrays["mask_res"]) if mask_res.ndim == 1: @@ -1328,7 +1820,7 @@ def get_default_save_func(cfg): x_probe = cfg["grid"]["xmin"] + 1.6 * boundary_width ix_probe = int(np.argmin(np.abs(np.array(cfg["grid"]["x"]) - x_probe))) w0 = derived["w0"] - w1 = derived["w1"] + w1 = _raman_carrier(cfg) 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)) if srs_on else 0.0 E0_source_sq = derived["E0_source"] ** 2 @@ -1353,20 +1845,25 @@ def get_default_save_func(cfg): 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"]) + flux_coeff_w1 = derived["c"] ** 2 / (w1 * cfg["grid"]["dx"]) I0_code = derived["I0_code"] + spectral_light = cfg["terms"].get("light", {}).get("solver", "fd") == "spectral" + 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). + # The spectral solver has no grid dispersion: k_grid = k. 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"] + if spectral_light: + return float(np.sin(k_dx) / k_dx) cos_kg = 1.0 - k_dx**2 / 2.0 if cos_kg <= -1.0: return 1.0 @@ -1383,6 +1880,15 @@ def discrete_flux(E, ix, coeff): cross = jnp.sum(jnp.conj(E[ix, :, :]) * E[ix + 1, :, :], axis=-1) return coeff * jnp.mean(jnp.imag(cross)) + ledger_on = bool(cfg["terms"]["epw"].get("energy_ledger", False)) + if ledger_on: + from adept._lpse2d.core.epw import LEDGER_CHANNELS, LEDGER_KEY + + # sum_k k^2 |phi_k|^2 -> the epw_energy normalization (Parseval, see above) + ledger_prefactor = epw_energy_prefactor / (nx * ny**2) + + raman_light_part = _raman_light_part(cfg) + def save_func(t, y, args): phi_k = y["epw"].view(jnp.complex128) ex = -1j * kx[:, None] * phi_k @@ -1390,10 +1896,19 @@ def save_func(t, y, args): ex = jnp.fft.ifft2(ex) ey = jnp.fft.ifft2(ey) e_sq = jnp.abs(ex) ** 2 + jnp.abs(ey) ** 2 + # the Raman-light diagnostics see the transverse part of the combined field only + y = raman_light_part(y) 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)) + if ledger_on: + # cumulative energy attributed to each split-step operation, in epw_energy units; + # epw_energy(t) - epw_energy(0) - sum of the channels closes to round-off + ledger = y[LEDGER_KEY] * ledger_prefactor + for i, name in enumerate(LEDGER_CHANNELS): + out[f"epw_ledger_{name}"] = ledger[i] + out["epw_ledger_closure"] = out["epw_energy"] - jnp.sum(ledger) # dissipation/boundary channels are TOTAL EPW-energy rates: epw_energy counts # only the electric part (the OSIRIS field-only convention), so the energy # actually handed to electrons -- and the budget sink -- carries the local @@ -1404,14 +1919,28 @@ def save_func(t, y, args): loss_factor = (1.0 - jnp.exp(-2.0 * gamma_dyn * dt)) / dt else: loss_factor = energy_loss_factor + # the per-k loss rate does not commute with the x-dependent energy factor, so # build a local electric-energy loss density from the sqrt-weighted fields -- # its box integral equals the k-space total exactly (Parseval) and reduces to # the previous 2x k-space sum on a uniform n = n_env box - sqrt_loss = jnp.sqrt(loss_factor) - ex_loss = jnp.fft.ifft2(-1j * kx[:, None] * phi_k * sqrt_loss) - ey_loss = jnp.fft.ifft2(-1j * ky[None, :] * phi_k * sqrt_loss) - loss_density = jnp.abs(ex_loss) ** 2 + jnp.abs(ey_loss) ** 2 + def weighted_density(weight): + ex_loss = jnp.fft.ifft2(-1j * kx[:, None] * phi_k * weight) + ey_loss = jnp.fft.ifft2(-1j * ky[None, :] * phi_k * weight) + return jnp.abs(ex_loss) ** 2 + jnp.abs(ey_loss) ** 2 + + loss_density = weighted_density(jnp.sqrt(jnp.maximum(loss_factor, 0.0))) + if signed_loss: + # terms.hpe.allow_growth: a negative applied rate hands energy from the electrons to + # the EPW; those modes enter with a minus sign, so the box integral stays + # sum_k L_k |E_k|^2 and epw_dissipation is the signed net transfer + loss_density = loss_density - weighted_density(jnp.sqrt(jnp.maximum(-loss_factor, 0.0))) + if combined_epw and nu_combined: + # the x-space collisional damping of the combined field acts on the EPW energy density + # locally (the light's share is the Raman absorption, not EPW dissipation) + dn = y.get("iaw_density") + n_eff = n_over_env if dn is None else n_over_env + dn + loss_density = loss_density + (1.0 - jnp.exp(-2.0 * nu_combined * n_eff**2 * dt)) / dt * e_sq out["epw_dissipation"] = epw_energy_prefactor * jnp.sum(jnp.mean(energy_total_factor * loss_density, axis=1)) out["epw_boundary_loss"] = epw_energy_prefactor * jnp.sum( jnp.mean(energy_total_factor * e_sq * boundary_sq_loss, axis=1) @@ -1422,6 +1951,17 @@ def save_func(t, y, args): out["iaw_density_sq"] = jnp.mean(iaw_density**2) out["iaw_density_abs_max"] = jnp.max(jnp.abs(iaw_density)) + # Thomson-scattering probes (LPSE thomsonScattering.N: wavevector.lw / .iaw, bandwidth): + # the complex amplitude and the power of the EPW potential (or the IAW density) in the + # k-window |k - k_probe| < bandwidth k0, as a time series for later spectral analysis + for i_probe, mask in enumerate(thomson_masks): + field = "iaw" if thomson_probes[i_probe].get("field", "epw") == "iaw" else "epw" + spectrum = jnp.fft.fft2(y["iaw_density"]) if field == "iaw" else phi_k + amplitude = jnp.sum(jnp.where(mask, spectrum, 0.0)) + out[f"thomson_{i_probe}_re"] = jnp.real(amplitude) + out[f"thomson_{i_probe}_im"] = jnp.imag(amplitude) + out[f"thomson_{i_probe}_power"] = jnp.sum(jnp.where(mask, jnp.abs(spectrum) ** 2, 0.0)) + if hpe_on: u = y["u_e"] if u.ndim == 2: @@ -1433,6 +1973,16 @@ def save_func(t, y, args): out["fhot_100keV"] = w_tail * jnp.sum(ke_kev > 100.0) out["hpe_mean_energy_keV"] = jnp.mean(ke_kev) out["hpe_hist"] = y["epw_hist"] + # LPSE-style instruments: cumulative energy (keV, per real electron via w_tail) out of + # each wall per energy bin, inside the acceptance cone, and the LD multiplier + from adept._lpse2d.core.hpe import WALLS + + wall_flux = y["hpe_wall_flux"] + for i_wall, wall in enumerate(WALLS): + for i_bin in range(wall_flux.shape[1]): + out[f"hpe_wall_energy_{wall}_bin{i_bin}"] = w_tail * wall_flux[i_wall, i_bin] + out["hpe_cone_energy"] = w_tail * y["hpe_cone_energy"][0] + out["hpe_ld_multiplier"] = y["hpe_ld_multiplier"][0] if have_ratio_band: ratio = y["gamma_L"] / gamma_an_safe # inflation-o-meters: worst-case reduction across the resonant band diff --git a/adept/_lpse2d/lpse_deck.py b/adept/_lpse2d/lpse_deck.py new file mode 100644 index 00000000..f0bcae23 --- /dev/null +++ b/adept/_lpse2d/lpse_deck.py @@ -0,0 +1,657 @@ +"""Translate original-LPSE decks (``lpse.parms``) to envelope-2d configs and read LPSE outputs. + +LPSE (LLE, C++) runs in "ZAK" units chosen so that the EPW dispersion coefficient +``3 v_te^2/(2 wp)`` is exactly 1 (``ParameterManager::createPhysicalParameters``); adept +runs in cgs scaled to um / ps / g. ``ZakUnits`` ports the conversion factors so deck +values (and output fields) can be compared quantitatively. The translation covers the +2-D physics of the parity plan (spectral EPW / combined solver, static or spectral light, +spectral IAWs, HPE on/off); keys outside it are reported, not silently dropped. + +Conventions that differ between the two codes and are handled here: + +- LPSE grid: ``grid.nodes`` nodes spanning ``grid.sizes`` microns with both end nodes + included, so ``dx = L/(N - 1)`` and the FFT period is ``N dx``; the box is centred on 0. + adept: ``nx`` cells of ``dx`` on ``[0, nx dx)``. Locations are shifted by ``+L/2``. +- LPSE anti-aliasing (``grid.antiAliasing.range``) zeroes the outer fraction of *each* + k axis (``dealias: rectangular``); the EPW time step is + ``min(lw.spectral.dt, lw.maxLightStepsPerStep * light dt)``. +- LPSE noise (``lw.noise.isCalculated = false``) is a flat-in-potential-kick source + without the Debye factor (``noise_model: thermal, noise_debye_factor: false``); the + amplitude converts with ``A_adept = A_lpse * zakUnitsPerMicron * potential_zak_to_adept`` + (both codes use an unnormalized forward FFT). +- ``relativity`` is enabled in LPSE by default, so the Landau rate is its relativistic + 2-D form. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import numpy as np + +# constants as in LPSE Constants.h / adept helpers +C_CGS = 2.99792458e10 +ME_CGS = 9.10938291e-28 +MP_CGS = 1.6726219e-24 +QE_CGS = 4.8032068e-10 +QE_SI = 1.602176634e-19 +KB_CGS = 1.380649e-16 +KELVIN_PER_EV = 11604.5 +FIELD_SCALE_ADEPT = 1.0e14 # cgs statvolt/cm per adept field unit (massScale^1/2 spatialScale^-1/2 / timeScale) +POTENTIAL_SCALE_ADEPT = 1.0e10 # cgs statvolt per adept potential unit (field unit * um) + + +class ZakUnits: + """LPSE's ZAK unit system for a plasma (Te, Ti in keV, Z, Mi/Me), envelope density + ``n_env`` (n0/nc) and laser wavelength (um). Port of ``ParameterManager.cpp:1136-1225``.""" + + def __init__(self, te_kev: float, ti_kev: float, z: float, mi_over_me: float, n_env: float, wavelength_um: float): + self.eta = 1.0 + 3.0 * ti_kev / (z * te_kev) + self.big_m = mi_over_me / (self.eta * z) # v_e^2 / c_s^2 + lambda_cm = wavelength_um * 1.0e-4 + te_kelvin = 1000.0 * te_kev * KELVIN_PER_EV + self.laser_frequency = 2.0 * np.pi * C_CGS / lambda_cm # rad/s + self.nc = self.laser_frequency**2 * ME_CGS / (4.0 * np.pi * QE_CGS**2) # 1/cm^3 + self.n0 = n_env * self.nc + self.plasma_frequency = np.sqrt(4.0 * np.pi * self.n0 * QE_CGS**2 / ME_CGS) # rad/s + self.debye_length_cm = np.sqrt(KB_CGS * te_kelvin / (4.0 * np.pi * self.n0 * QE_CGS**2)) + self.n0_zak = 0.75 * self.big_m + self.nc_zak = self.n0_zak / n_env + self.kde = 1.5 * np.sqrt(self.big_m) + self.wpe = 1.5 * self.big_m + self.zak_per_ps = 1.0e-12 * self.plasma_frequency / self.wpe + self.zak_per_um = 1.0e-4 / (self.kde * self.debye_length_cm) + self.um_per_zak = 1.0 / self.zak_per_um + self.ps_per_zak = 1.0 / self.zak_per_ps + self.ve = np.sqrt(self.big_m) # electron thermal velocity, ZAK + self.speed_of_light = C_CGS * 1.0e4 * self.zak_per_um / (1.0e12 * self.zak_per_ps) + self.io_scale_factor = ( + (8.0 / 3.0) * (self.eta / self.big_m) * self.n0 * (QE_SI * 1000.0 * te_kev) * C_CGS + ) # W/cm^2 + self.field_zak_to_cgs = np.sqrt(1.0e7 * self.io_scale_factor * 8.0 * np.pi / C_CGS) + self.potential_zak_to_cgs = self.field_zak_to_cgs * self.um_per_zak * 1.0e-4 + self.field_zak_to_adept = self.field_zak_to_cgs / FIELD_SCALE_ADEPT + self.potential_zak_to_adept = self.potential_zak_to_cgs / POTENTIAL_SCALE_ADEPT + + def summary(self) -> dict: + return { + "critical density (1/cm^3)": self.nc, + "envelope density (1/cm^3)": self.n0, + "plasma frequency (rad/s)": self.plasma_frequency, + "Debye length (cm)": self.debye_length_cm, + "zak per ps": self.zak_per_ps, + "zak per um": self.zak_per_um, + "Io scale factor (W/cm^2)": self.io_scale_factor, + } + + +# ------------------------------------------------------------------ decks -- + + +def parse_parms(path: str | Path) -> dict[str, str]: + """Parse an LPSE ``lpse.parms`` deck (``key = value;`` lines, ``#include "file"``, + ``#`` comments) into a flat dict of raw string values, includes resolved relative + to the deck's directory.""" + path = Path(path) + parms: dict[str, str] = {} + for raw in path.read_text().splitlines(): + line = raw.strip() + if not line: + continue + inc = re.match(r'#include\s+"([^"]+)"', line) + if inc: + parms.update(parse_parms(path.parent / inc.group(1))) + continue + if line.startswith("#") or line.startswith("//"): + continue + line = line.split("//")[0].strip() + for statement in line.split(";"): + statement = statement.strip() + if not statement or "=" not in statement: + continue + key, value = statement.split("=", 1) + parms[key.strip()] = value.strip().strip('"') + return parms + + +def _floats(value: str) -> list[float]: + return [float(v) for v in value.replace(",", " ").split()] + + +def _bool(value: str | None, default: bool = False) -> bool: + if value is None: + return default + return value.strip().lower() in ("true", "t", "1", "yes") + + +def translate_parms( + parms: dict[str, str], *, experiment: str = "lpse-parity", run: str = "translated-deck" +) -> tuple[dict, dict]: + """Build an envelope-2d config from parsed LPSE deck values. + + Returns ``(cfg, report)``; ``report["unsupported"]`` lists deck settings that have no + adept counterpart (they are also printed), ``report["units"]`` the ZakUnits summary. + """ + g = parms.get + report: dict = {"unsupported": [], "notes": []} + + te = float(g("physical.Te", "2")) + ti = float(g("physical.Ti", "1")) + z = float(g("physical.Z", "1")) + mi_over_me = float(g("physical.MiOverMe", "1836")) + n_env = float(g("lw.envelopeDensity", "0.25")) + wavelength_um = float(g("laser.wavelength", "0.351")) + zak = ZakUnits(te, ti, z, mi_over_me, n_env, wavelength_um) + report["units"] = zak.summary() + + # ---- grid: LPSE nodes span the size with both end nodes, d = L / (N - 1) per axis. LPSE runs + # square cells: h = min(d) over the deck's dimensions, a mismatch sqrt(sum (h - d)^2) above + # 0.1 h is refused, and the box becomes h (N - 1) per axis (ParameterManager.cpp:1258-1281) + sizes = _floats(g("grid.sizes", "0 0")) + nodes = [int(v) for v in _floats(g("grid.nodes", "0 0"))] + nx = nodes[0] + ny = nodes[1] if len(nodes) > 1 else 1 + spacings = [sizes[i] / (nodes[i] - 1) for i in range(min(len(sizes), len(nodes))) if nodes[i] > 1] + dx = min(spacings) + if np.sqrt(sum((dx - d) ** 2 for d in spacings)) > 0.1 * dx: + raise ValueError( + f"grid.sizes {sizes} / grid.nodes {nodes}: cell sizes {spacings} differ by more than LPSE's 10 % " + "square-cell tolerance (LPSE refuses the deck); adjust grid.nodes" + ) + if any(abs(d / dx - 1.0) > 1e-9 for d in spacings): + report["notes"].append( + f"grid: LPSE runs square cells h = min(size / (nodes - 1)) = {dx:.6g} um and rescales the box to " + f"h (nodes - 1) = {[round(dx * (n - 1), 6) for n in nodes if n > 1]} um (deck grid.sizes {sizes})" + ) + lx = dx * (nx - 1) + xmax = nx * dx + if ny > 1: + half_y = 0.5 * ny * dx + else: + half_y = 0.4 * dx # collapses to ny = 1 in adept + if len(sizes) > 2 and sizes[2] > 0 and len(nodes) > 2 and nodes[2] > 1: + report["unsupported"].append("3-D grid (grid.sizes/nodes z): adept envelope-2d is 2-D") + + # ---- time stepping + lw_dt = float(g("lw.spectral.dt", g("simulation.dt", "0.002"))) + light_dt = None + for key in ("raman.dt", "laser.dt"): + if key in parms: + light_dt = float(parms[key]) + break + max_light_steps = int(float(g("lw.maxLightStepsPerStep", "10"))) + # LPSE laser.enable (default false, ParameterManager.cpp:267-275) switches the E0 light off + # entirely -- every E0 entry point returns early (LightSolver.cpp:386, 1933) -- so a disabled + # laser is no pump and no depletion, whatever beams the deck still lists + laser_on = _bool(g("laser.enable")) + laser_evolves = laser_on and g("laser.solver", "static").lower() != "static" + raman_on = _bool(g("raman.enable")) + if light_dt is not None and (laser_evolves or raman_on): + dt = min(lw_dt, max_light_steps * light_dt) + light_substeps = max(1, round(dt / light_dt)) + else: + dt = lw_dt + light_substeps = None + tmax = float(g("simulation.time.end", "1")) + sample_period = float(g("simulation.samplePeriod", str(tmax))) + + # ---- anti-aliasing: outer fraction of each axis zeroed + # LPSE: no anti-aliasing unless grid.antiAliasing.range is given (ParameterManager.cpp:237-240) + aa_range = float(g("grid.antiAliasing.range", "0")) + + # ---- density profile + shape = g("densityProfile.shape", "linear").lower() + n_min = float(g("densityProfile.NminOverNc", str(n_env))) + n_max = float(g("densityProfile.NmaxOverNc", str(n_env))) + x_min_loc = _floats(g("densityProfile.NminLocation", f"{-lx / 2}"))[0] + lx / 2.0 + x_max_loc = _floats(g("densityProfile.NmaxLocation", f"{lx / 2}"))[0] + lx / 2.0 + shape_map = { + "linear": "linear", + "exp": "exp", + "exponential": "exp", + "gaussian": "gaussian", + "inversesquare": "inverse-power", + "inversepower": "inverse-power", + "quadratic": "quadratic", + "qd": "qd", + "gd": "gd", + "file": "file", + } + if shape not in shape_map: + report["unsupported"].append(f"densityProfile.shape = {shape} (translated as linear)") + adept_shape = shape_map.get(shape, "linear") + y_min_loc = _floats(g("densityProfile.NminLocation", "0 0")) + y_max_loc = _floats(g("densityProfile.NmaxLocation", "0 0")) + density = { + "basis": f"lpse-{adept_shape}", + "min": n_min, + "max": n_max, + "min_location": f"{x_min_loc}um", + "max_location": f"{x_max_loc}um", + "min_location_y": f"{y_min_loc[1] if len(y_min_loc) > 1 else 0.0}um", + "max_location_y": f"{y_max_loc[1] if len(y_max_loc) > 1 else 0.0}um", + "geometry": "spherical" if g("densityProfile.geometry", "cartesian").lower() == "spherical" else "cartesian", + "origin": f"{lx / 2.0}um", + "max_density": min(float(g("densityProfile.maxBackgroundDensity", "1.25")), 1.25), + } + sg = g("densityProfile.sgOrder", g("densityProfile.sgPower", g("densityProfile.power"))) + if shape == "inversesquare": + density["sg_order"] = 2.0 + elif sg is not None: + density["sg_order"] = float(sg) + if adept_shape == "quadratic": + density["central_density"] = float(g("densityProfile.quadratic.centralDensity", "0")) + if adept_shape in ("qd", "gd"): + density["dip_depth"] = float(g("densityProfile.dip.depth", "0")) + density["dip_width"] = f"{float(g('densityProfile.dip.width', '0'))}um" + density["dip_offset"] = f"{float(g('densityProfile.dip.offset', '0'))}um" + if adept_shape == "file": + density["file"] = g("densityProfile.loadFilename", "") + report["notes"].append("densityProfile.shape = file: density.file must point at the LPSE grid file") + + # ---- EPW terms + labc_x = float(g("lw.Labc.min.x", g("lw.Labc", "0"))) + labc_y = float(g("lw.Labc.min.y", g("lw.Labc", "0"))) + boundary = {"x": "absorbing" if labc_x > 0 else "periodic", "y": "absorbing" if labc_y > 0 else "periodic"} + boundary_width = max(labc_x, labc_y, dx) + if labc_x > 0 and labc_y > 0 and labc_x != labc_y: + report["notes"].append("different x/y absorber widths: adept uses one boundary_width (the larger)") + landau_on = _bool(g("lw.landauDamping.enable"), True) + relativity = _bool(g("physical.isRelativistic", g("isRelativistic")), True) + if _bool(g("lw.collisionalDampingRate.isCalculated")): + collisions: bool | float = True + else: + collisions = float(g("lw.collisionalDampingRate", "0")) + tpd_on = _bool(g("lw.TPD.enable")) + srs_on = _bool(g("lw.SRS.enable")) + noise_on = _bool(g("lw.noise.enable")) + noise_amp_lpse = float(g("lw.noise.amplitude", "1" if _bool(g("lw.noise.isCalculated")) else "0")) + noise_calculated = _bool(g("lw.noise.isCalculated")) + if noise_calculated: + report["notes"].append( + "lw.noise.isCalculated: LPSE's absolute constant is not reproduced; adept's equipartition " + "calibration (noise_calibrate) is used with the deck amplitude as multiplier" + ) + source = { + "noise": noise_on, + "noise_model": "thermal", + "noise_debye_factor": noise_calculated, + "noise_calibrate": noise_calculated, + "noise_amplitude": float( + noise_amp_lpse if noise_calculated else noise_amp_lpse * zak.zak_per_um * zak.potential_zak_to_adept + ), + "noise_seed": int(float(g("seed", "1"))), + "tpd": tpd_on, + "tpd_form": "lpse", + "srs": srs_on, + # LPSE lw.kFilter is off unless enabled; adept's default (on) would zero the SRS + # source at densities detuned from the envelope density + "srs_k_filter": _bool(g("lw.kFilter.enable")), + "srs_k_filter_scale": float(g("lw.kFilter.scale", "1.2")), + } + epw_solver = "combined" if g("lw.solver", "spectral").lower() == "combined" else "separate" + if g("lw.solver", "spectral").lower() == "fd": + report["unsupported"].append("lw.solver = fd (LPSE itself disables it); translated as spectral") + epw = { + "boundary": boundary, + "damping": { + "collisions": collisions, + "landau": landau_on, + "landau_form": "relativistic" if relativity else "lpse", + "landau_lower_threshold": float(g("lw.landauDamping.lowerThreshold", "0")), + }, + "density_gradient": True, + "linear": True, + "source": source, + "solver": epw_solver, + } + if "lw.maxWavenumber" in parms: + epw["max_wavenumber"] = float(parms["lw.maxWavenumber"]) + + # ---- light + raman_solver = g("raman.solver", "spectral").lower() + laser_solver = g("laser.solver", "static").lower() + light: dict = {"pump_depletion": laser_evolves} + if epw_solver == "combined" or raman_solver == "spectral" or laser_solver == "spectral": + light["solver"] = "spectral" + else: + light["solver"] = "fd" + if laser_evolves and laser_solver != raman_solver and raman_on: + report["notes"].append( + f"laser.solver = {laser_solver} but raman.solver = {raman_solver}: " + f"one adept light solver ({light['solver']})" + ) + absorption = float(g("raman.evolution.absorption", g("laser.evolution.absorption", "0"))) + if absorption > 0: + light["absorption"] = absorption + elif absorption < 0: + light["absorption"] = True + if "raman.spectral.maxWavenumber" in parms: + light["max_wavenumber"] = float(parms["raman.spectral.maxWavenumber"]) + light["boundary_max_rate"] = float( + g("raman.evolution.abc.maxDampingRate", g("laser.evolution.abc.maxDampingRate", "5000")) + ) + # LPSE default is false; the reference decks set it true + light["transverse_source"] = _bool( + g("raman.takeTransversePartOfSourceTerms", g("laser.takeTransversePartOfSourceTerms", "false")) + ) + if laser_on and (_bool(g("laser.pumpDepletion.TPD.enable")) or _bool(g("laser.pumpDepletion.SRS.enable"))): + light["pump_depletion"] = True + if _bool(g("lw.kFilter.enable")): + light["tpd_k_filter"] = True + + # ---- laser beams (static: intensities summed; direction along +x expected). A disabled laser + # keeps a beam-free E0 driver at zero intensity, the form the shipped laser-off decks take. + n_beams = int(float(g("laser.nBeams", "1"))) if laser_on else 0 + if not laser_on: + listed = sum(float(v) for k, v in parms.items() if re.fullmatch(r"laser\.\d+\.intensity", k)) + report["notes"].append( + "laser.enable = false: no pump (zero intensity, no depletion)" + + (f"; the listed beam intensities ({listed:.6g} W/cm^2) are not applied" if listed else "") + ) + intensity = 0.0 + angle_deg = 0.0 + beams = [] + for b in range(1, n_beams + 1): + i_b = float(g(f"laser.{b}.intensity", "0")) + intensity += i_b + direction = _floats(g(f"laser.{b}.direction", "1 0 0")) + dz = direction[2] if len(direction) > 2 else 0.0 + if abs(dz) > 1e-6 * np.linalg.norm(direction) or direction[0] <= 0: + report["unsupported"].append(f"laser.{b}.direction {direction}: adept pump is in-plane and rightward") + beam_angle = float(np.degrees(np.arctan2(direction[1], direction[0]))) + if b == 1: + angle_deg = beam_angle + if float(g(f"laser.{b}.polarization", "0")) != 0.0: + report["unsupported"].append(f"laser.{b}.polarization != 0: adept pump is in-plane (p) polarized") + # LPSE reads laser.N.phase in degrees (LightSolver.cpp:1716-1717, 1877) and frequencyShift as + # dW/W0; a non-zero wavelengthShift (Angstrom) overrides it as -1e-4 dLambda / lambda_um + wavelength_shift = float(g(f"laser.{b}.wavelengthShift", "0")) + if wavelength_shift != 0.0: + delta_omega_b = -(1.0e-4 * wavelength_shift) / wavelength_um + else: + delta_omega_b = float(g(f"laser.{b}.frequencyShift", "0")) + beams.append( + { + "intensity": i_b, + "angle": beam_angle, + "phase": float(np.deg2rad(float(g(f"laser.{b}.phase", "0")))), + "delta_omega": delta_omega_b, + } + ) + beam_extras = {} + if laser_on and float(g("laser.1.width", "0")) > 0: + beam_extras["beam_width"] = f"{float(g('laser.1.width'))}um" + beam_extras["beam_sg_order"] = float(g("laser.1.sgOrder", g("laser.1.sgPower", "2"))) + off = _floats(g("laser.1.offset", "0 0")) + beam_extras["beam_offset"] = f"{off[1] if len(off) > 1 else 0.0}um" + kap = float(g("laser.bandwidth.KAP.frequency", "0")) if laser_on else 0.0 + if kap > 0: + beam_extras["kap_bandwidth"] = kap + if laser_on and g("laser.pulse.file"): + beam_extras["pulse_file"] = g("laser.pulse.file") + report["notes"].append("laser.pulse.file: adept expects a two-column (t_ps, amplitude) text table") + drivers = { + "E0": { + "shape": "uniform", + "num_colors": 1, + "delta_omega_max": 0.0, + "params": {"phases": {"seed": 42}}, + "angle": angle_deg, + # one beam keeps its list only when it carries a phase or frequency shift (the static, + # evolved and combined pumps all honor beams[0]); otherwise the deck stays beam-free + **({"beams": beams} if n_beams > 1 or any(b["phase"] or b["delta_omega"] for b in beams) else {}), + **beam_extras, + "envelope": { + "tw": f"{10 * tmax}ps", + "tr": "0.01ps", + "tc": f"{5 * tmax}ps", + "xr": "0.2um", + "xw": "1000um", + "xc": "50um", + "yr": "0.2um", + "yw": "1000um", + "yc": "0um", + }, + } + } + if laser_evolves: + labc_light = float(g("laser.evolution.Labc.min.x", g("laser.evolution.Labc", str(labc_x)))) + drivers["E0"]["offset"] = f"{2.0 * max(labc_light, dx)}um" + if raman_on and int(float(g("raman.nBeams", "0"))) > 0: + drivers["E1"] = {"intensity": f"{float(g('raman.1.intensity', '0'))}W/cm^2"} + + # ---- IAW + iaw = None + if _bool(g("iaw.enable")): + iaw = { + "active": True, + "solver": "spectral" if g("iaw.solver", "spectral").lower() == "spectral" else "explicit", + "damping": { + "landau": float(g("iaw.landauDampingRate", g("iaw.dampingRate", "0"))), + "collisions": float(g("iaw.collisionalDampingRate", "0")), + "landau_form": "simplified" if _bool(g("iaw.fd.damping.isSimplified"), True) else "full", + }, + "max_density_perturbation": float(g("iaw.amplitudeClamp", "0.9")), + "noise": _bool(g("iaw.noise.enable")), + } + if g("iaw.solver", "spectral").lower() != "spectral": + report["notes"].append("iaw.solver = fd translated as spectral (adept has no PPM IAW solver)") + if "iaw.spectral.dt" in parms: + # LPSE (Lpse.cpp:1264-1291): the IAW step is its own dt capped by maxLightStepsPerStep + # light steps (when light evolves) and by maxLwStepsPerStep (default 2) EPW steps + iaw_step = float(parms["iaw.spectral.dt"]) + if light_dt is not None and (laser_evolves or raman_on): + iaw_step = min(iaw_step, int(float(g("iaw.maxLightStepsPerStep", "10"))) * light_dt) + iaw_step = min(iaw_step, int(float(g("iaw.maxLwStepsPerStep", "2"))) * dt) + stride = max(1, int(np.floor(iaw_step * 1.0001 / dt))) + if stride > 1: + iaw["stride"] = stride + tf = {w: _bool(g(f"thermalFil.{w}.enable")) for w in ("laser", "raman", "lw")} + if tf["laser"] and not laser_on: + raise ValueError("thermalFil.laser.enable needs laser.enable = true (LPSE rejects it too)") + if any(tf.values()): + tf["nonlocal"] = _bool(g("thermalFil.isNonlocal")) + tf["conductivity_multiplier"] = float(g("thermalFil.conductivityMultiplier", "1")) + iaw["thermal_filamentation"] = tf + report["notes"].append( + "thermalFil: LPSE's source form with adept's own normalization (heat + momentum equations)" + ) + if "fluid.velocity" in parms: + vel = _floats(parms["fluid.velocity"]) + iaw["flow"] = [vel[0], vel[1] if len(vel) > 1 else 0.0] + + # ---- HPE + hpe = None + if _bool(g("hpe.enable")): + hpe = {"active": True} + for lpse_key, adept_key in ( + ("hpe.numParticles", "n_particles"), + ("hpe.nParticles", "n_particles"), + ("hpe.nElectrons", "n_particles"), + ): + if lpse_key in parms: + hpe[adept_key] = int(float(parms[lpse_key])) + if "hpe.velocityGrid" in parms: + hpe["nv"] = int(float(parms["hpe.velocityGrid"])) + if "hpe.startEvolutionAt" in parms: + hpe["t_start"] = f"{float(parms['hpe.startEvolutionAt'])}ps" + if "hpe.VminOverVminPhase" in parms and float(parms["hpe.VminOverVminPhase"]) == 0.0: + report["notes"].append( + "hpe.VminOverVminPhase = 0 (LPSE tracks the whole Maxwellian): adept keeps its tail cutoff " + "(terms.hpe.v_min, 2.5 vte); set v_min: 0.05 to reproduce LPSE's wall-flux bins below ~6 keV " + "in field-free runs -- the kinetic feedback is calibrated for a tail" + ) + if "hpe.thermalizationProbability" in parms: + tp = _floats(parms["hpe.thermalizationProbability"]) + hpe["thermalization_probability"] = [tp[0], tp[1] if len(tp) > 1 else tp[0]] + if "hpe.magneticField" in parms: + b = _floats(parms["hpe.magneticField"]) + hpe["magnetic_field"] = b[2] if len(b) > 2 else 0.0 + if any(abs(v) > 0 for v in b[:2]): + report["unsupported"].append("hpe.magneticField in-plane components (only B_z acts on the 2-D push)") + if "hpe.gammaLimit.damping" in parms: + hpe["gamma_limit_damping"] = float(parms["hpe.gammaLimit.damping"]) + if "hpe.gammaLimit.growth" in parms: + hpe["gamma_limit_growth"] = float(parms["hpe.gammaLimit.growth"]) + if _bool(g("hpe.allowGrowth")): + hpe["allow_growth"] = True + if _bool(g("hpe.enforceEnergyConservation")): + hpe["energy_conservation"] = True + hpe["energy_conservation_steps"] = float(g("hpe.numStepsToAverageEnergyChange", "1")) + n_flux = int(float(g("hpe.metrics.nFluxMetrics", "0"))) + if n_flux > 0: + # LPSE hpe.metrics.fluxMetric.N.energy.{min,max} in keV (the decks bin a 2 keV plasma at + # 0.01, 0.5, 1, 2, 4, 6, 8, 20); the adept instrument bins on one sorted edge list (keV) + edges = set() + for i in range(1, n_flux + 1): + for end, default in (("min", "0"), ("max", "1e9")): + key = f"hpe.metrics.fluxMetric.{i}.energy.{end}" + key = key if key in parms else f"hpe.metrics.flux.{i}.energy.{end}" + edges.add(float(g(key, default))) + hpe["flux_bins"] = sorted(edges) + if int(float(g("hpe.metrics.nPowerMetrics", "0"))) > 0: + prefix = "hpe.metrics.powerMetric.1" + if f"{prefix}.angle" not in parms: + prefix = "hpe.metrics.power.1" + hpe["cone_angle"] = float(g(f"{prefix}.angle", "30")) + direction = _floats(g(f"{prefix}.direction", "1 0 0")) + hpe["cone_direction"] = [direction[0], direction[1] if len(direction) > 1 else 0.0] + report["notes"].append( + "hpe: dt/dtFields/stepsPerLandauUpdate/vdf.*/blend.gamma are LPSE tracker internals without an adept " + "equivalent (adept sub-cycles at substep_courant and blends with v_blend_buffer)" + ) + + for key in parms: + if key.startswith("qle") and _bool(parms[key]): + report["unsupported"].append(f"{key} (quasilinear module)") + + cfg = { + "solver": "envelope-2d", + "units": { + "atomic number": mi_over_me / 1836.15, + "envelope density": n_env, + "ionization state": round(z), + "laser intensity": f"{intensity:.6g}W/cm^2", + "laser_wavelength": f"{wavelength_um}um", + "reference electron temperature": f"{te}keV", + "reference ion temperature": f"{ti}keV", + }, + "density": density, + "grid": { + "boundary_abs_coeff": 200.0, + "boundary_width": f"{boundary_width}um", + "boundary_profile": "exp", + "boundary_max_rate": 200.0, + "low_pass_filter": 1.0 - aa_range, + "dealias": "rectangular", + "dt": f"{dt}ps", + "dx": f"{dx}um", + "xmax": f"{xmax}um", + "tmax": f"{tmax}ps", + "tmin": "0ps", + "ymax": f"{half_y}um", + "ymin": f"{-half_y}um", + }, + "save": {"fields": {"t": {"dt": f"{sample_period}ps", "tmax": f"{tmax}ps", "tmin": "0ps"}}}, + "mlflow": {"experiment": experiment, "run": run}, + "drivers": drivers, + "terms": {"epw": epw, "light": light, "zero_mask": True}, + } + if light_substeps is not None: + cfg["grid"]["light_substeps"] = light_substeps + if iaw is not None: + cfg["terms"]["iaw"] = iaw + if hpe is not None: + cfg["terms"]["hpe"] = hpe + for item in report["unsupported"]: + print(f"lpse_deck: not translated -- {item}") + return cfg, report + + +# ---------------------------------------------------------------- outputs -- + + +def read_metrics(path: str | Path) -> dict[str, np.ndarray]: + """Read an LPSE ``lpse.metrics`` table into ``{column name: array}``.""" + names = [] + rows = [] + for line in Path(path).read_text().splitlines(): + m = re.match(r"#\s+\((\d+)\)\s+(\S+)", line) + if m: + names.append(m.group(2)) + continue + if line.startswith("#") or not line.strip(): + continue + rows.append([float(v) for v in line.split()]) + data = np.asarray(rows) + return {name: data[:, i] for i, name in enumerate(names)} + + +def read_frames(path: str | Path) -> list[tuple[dict, np.ndarray]]: + """Read an LPSE binary field file: a list of ``(header, array)`` per frame, the array + complex ``(Nx, Ny)`` (single precision in the file; ``FileType=real`` gives a real + array). Header values are parsed to float where possible.""" + blob = Path(path).read_bytes() + frames = [] + begin = b"# BeginHeaderSegment;" + data_marker = b"# BeginDataSegment;\n" + pos = 0 + while True: + i = blob.find(begin, pos) + if i < 0: + break + j = blob.find(data_marker, i) + header_text = blob[i + len(begin) : j].decode(errors="replace") + header: dict = {} + for pair in header_text.replace("# EndHeaderSegment;", "").replace(";", "").split(","): + if "=" not in pair: + continue + k, v = pair.split("=", 1) + k, v = k.strip(), v.strip() + try: + header[k] = float(v) if len(v.split()) == 1 else [float(x) for x in v.split()] + except ValueError: + header[k] = v + nx, ny, nz = int(header["Nx"]), int(header.get("Ny", 1)), int(header.get("Nz", 1)) + is_complex = str(header.get("FileType", "complex")).lower().startswith("c") + count = nx * ny * nz * (2 if is_complex else 1) + start = j + len(data_marker) + raw = np.frombuffer(blob[start : start + 4 * count], dtype=" dict: def init_diffeqsolve(self): self.cfg = get_save_quantities(self.cfg) + t0 = float(getattr(self, "restart_t0", 0.0)) self.time_quantities = { - "t0": 0.0, + "t0": t0, "t1": self.cfg["grid"]["tmax"], "max_steps": self.cfg["grid"]["max_steps"], - "save_t0": 0.0, + "save_t0": t0, "save_t1": self.cfg["grid"]["tmax"], "save_nt": self.cfg["grid"]["tmax"], } - self.diffeqsolve_quants = dict( - terms=ODETerm(SplitStep(self.cfg)), - solver=Stepper(), - saveat=dict( - subs={ - k: SubSaveAt(ts=subsave["t"]["ax"], fn=subsave["func"]) for k, subsave in self.cfg["save"].items() - } - ), - ) + subs = { + k: SubSaveAt(ts=subsave["t"]["ax"], fn=subsave["func"]) + for k, subsave in self.cfg["save"].items() + if isinstance(subsave, dict) and "func" in subsave + } + if self.cfg["save"].get("checkpoint"): + # the full state at the final time (save.checkpoint: true | path), written by + # post_process as an .npz that `restart.file` accepts + subs["checkpoint"] = SubSaveAt(ts=np.array([self.cfg["grid"]["tmax"]]), fn=lambda t, y, args: y) + self.diffeqsolve_quants = dict(terms=ODETerm(SplitStep(self.cfg)), solver=Stepper(), saveat=dict(subs=subs)) def init_state_and_args(self) -> dict: # The initial EPW is identically zero; noise-seeded runs get their seeding from @@ -86,6 +89,11 @@ def init_state_and_args(self) -> dict: E1 = np.zeros((self.cfg["grid"]["nx"], self.cfg["grid"]["ny"], 2), dtype=np.complex128) state = {"epw": epw, "E0": E0, "E1": E1} + if self.cfg["terms"]["epw"].get("energy_ledger", False): + from adept._lpse2d.core.epw import LEDGER_CHANNELS, LEDGER_KEY + + state[LEDGER_KEY] = np.zeros(len(LEDGER_CHANNELS), dtype=np.float64) + if self.cfg["terms"].get("iaw", {}).get("active", False): iaw_shape = (self.cfg["grid"]["nx"], self.cfg["grid"]["ny"]) state["iaw_density"] = np.zeros(iaw_shape, dtype=np.float64) @@ -99,6 +107,38 @@ def init_state_and_args(self) -> dict: state = state | load_particles(self.cfg) self.state = {k: v.view(dtype=np.float64) for k, v in state.items()} + # ---- restart (LPSE --restart): replace the freshly built state by a checkpoint and + # continue from its time; the per-step noise / wall keys are folded in from the + # time index, so a resumed run reproduces the unbroken one to round-off + self.restart_t0 = 0.0 + restart = self.cfg.get("restart") + if restart and restart.get("file"): + loaded = np.load(restart["file"], allow_pickle=False) + self.restart_t0 = float(loaded["t"]) + # the saved time carries the solver's float precision; snap it to the step grid so the + # resumed step times coincide with the unbroken run's + dt = float(self.cfg["grid"]["dt"]) + n_steps = round(self.restart_t0 / dt) + if abs(self.restart_t0 - n_steps * dt) < 1.0e-5 * dt: + self.restart_t0 = n_steps * dt + if self.restart_t0 >= self.cfg["grid"]["tmax"]: + raise ValueError(f"restart time {self.restart_t0} ps is not before grid.tmax") + missing = [k for k in self.state if k not in loaded.files] + if missing: + raise ValueError(f"checkpoint {restart['file']} lacks state entries {missing}") + for k, v in self.state.items(): + if loaded[k].shape != v.shape: + raise ValueError(f"checkpoint entry {k} has shape {loaded[k].shape}, expected {v.shape}") + self.state = {k: np.asarray(loaded[k], dtype=np.float64) for k in self.state} + for sub in self.cfg["save"].values(): + if isinstance(sub, dict) and isinstance(sub.get("t"), dict) and "tmin" in sub["t"]: + if _Q(sub["t"]["tmin"]).to("ps").value < self.restart_t0: + sub["t"]["tmin"] = f"{self.restart_t0:.9g}ps" + # the default series samples the grid's time axis: keep only the resumed interval + t_axis = np.asarray(self.cfg["grid"]["t"]) + kept = t_axis[t_axis >= self.restart_t0 - 1.0e-12] + self.cfg["grid"]["t"] = kept if kept.size else np.asarray([self.cfg["grid"]["tmax"]]) + print(f"restarting from {restart['file']} at t = {self.restart_t0:.6g} ps") self.args = {"drivers": {k: v["derived"] for k, v in self.cfg["drivers"].items()}} @filter_jit diff --git a/adept/_lpse2d/threshold.py b/adept/_lpse2d/threshold.py new file mode 100644 index 00000000..6afc81d0 --- /dev/null +++ b/adept/_lpse2d/threshold.py @@ -0,0 +1,114 @@ +"""Absolute-threshold bisection driver for the envelope-2d solver (LPSE ``AbsoluteThreshold``). + +Runs the same configuration at a sequence of pump intensities and bisects on a growth +criterion evaluated from the run's post-processed metrics. Each run is an ordinary +``ergoExo`` run (its own MLflow run); the bisection itself is recorded as the parent run's +parameters and metrics. + +Usage:: + + from adept._lpse2d.threshold import find_threshold + + result = find_threshold(cfg, intensity_lo="1e14W/cm^2", intensity_hi="1e15W/cm^2", n_iter=6) + result["threshold"] # W/cm^2, midpoint of the final bracket + +The default criterion is LPSE's: the run is "unstable" when the fitted EPW energy growth +rate exceeds ``growth_min`` (1/ps) and the fit is measurable; ``criterion`` can be any +callable of the metrics dict returning a bool. +""" + +from __future__ import annotations + +from collections.abc import Callable +from copy import deepcopy + +import numpy as np + +from adept._lpse2d.helpers import _Q + + +def _intensity_w_cm2(value) -> float: + return float(_Q(value).to("W/cm^2").value) if isinstance(value, str) else float(value) + + +def default_criterion(metrics: dict, growth_min: float = 0.0) -> bool: + """Unstable when the EPW energy grows: a measurable fit with a positive rate above ``growth_min``.""" + measurable = float(metrics.get("epw_growth_measurable", 0.0)) > 0.5 + rate = float(metrics.get("epw_growth_rate_per_ps", 0.0)) + return measurable and rate > growth_min + + +def run_at_intensity(cfg: dict, intensity_w_cm2: float, run_name: str | None = None) -> dict: + """One ergoExo run of ``cfg`` with ``units['laser intensity']`` replaced; returns the metrics.""" + from adept import ergoExo + + run_cfg = deepcopy(cfg) + run_cfg["units"]["laser intensity"] = f"{intensity_w_cm2:.6g}W/cm^2" + if run_name is not None: + run_cfg.setdefault("mlflow", {})["run"] = run_name + exo = ergoExo() + modules = exo.setup(run_cfg) + _, ppo, _ = exo(modules) + return dict(ppo["metrics"]) + + +def find_threshold( + cfg: dict, + intensity_lo, + intensity_hi, + n_iter: int = 6, + criterion: Callable[[dict], bool] | None = None, + growth_min: float = 0.0, + runner: Callable[[dict, float, str | None], dict] = run_at_intensity, + log_mlflow: bool = True, +) -> dict: + """Bisect the pump intensity between ``intensity_lo`` (expected stable) and ``intensity_hi`` + (expected unstable) with ``n_iter`` midpoint runs. The bracket endpoints are run first and + must straddle the criterion, mirroring LPSE's ``AbsoluteThreshold`` bracket check. + + Returns ``{"threshold", "bracket", "history": [(intensity, unstable, metrics), ...]}``.""" + crit = criterion or (lambda m: default_criterion(m, growth_min)) + lo, hi = _intensity_w_cm2(intensity_lo), _intensity_w_cm2(intensity_hi) + if not 0.0 < lo < hi: + raise ValueError( + "the bracket needs 0 < intensity_lo < intensity_hi: the geometric midpoint sqrt(lo hi) of a zero " + "lower bound never moves; start from a small positive stable intensity" + ) + base_name = cfg.get("mlflow", {}).get("run", "threshold") + history = [] + + def evaluate(intensity: float, tag: str) -> bool: + metrics = runner(cfg, intensity, f"{base_name}-{tag}") + unstable = bool(crit(metrics)) + history.append((intensity, unstable, metrics)) + return unstable + + if evaluate(lo, "lo"): + raise ValueError(f"the lower bracket {lo:.3e} W/cm^2 is already unstable; lower intensity_lo") + if not evaluate(hi, "hi"): + raise ValueError(f"the upper bracket {hi:.3e} W/cm^2 is still stable; raise intensity_hi") + for i in range(n_iter): + mid = float(np.sqrt(lo * hi)) # geometric midpoint: thresholds scale multiplicatively + if evaluate(mid, f"iter{i}"): + hi = mid + else: + lo = mid + result = {"threshold": float(np.sqrt(lo * hi)), "bracket": (lo, hi), "history": history} + if log_mlflow: + try: + import mlflow + + with mlflow.start_run(run_name=f"{base_name}-bisection"): + mlflow.log_params( + { + "intensity_lo": _intensity_w_cm2(intensity_lo), + "intensity_hi": _intensity_w_cm2(intensity_hi), + "n_iter": n_iter, + } + ) + mlflow.log_metrics({"threshold_W_cm2": result["threshold"], "bracket_lo": lo, "bracket_hi": hi}) + for step, (intensity, unstable, _) in enumerate(history): + mlflow.log_metrics({"intensity": intensity, "unstable": float(unstable)}, step=step) + except Exception as exc: # logging must never break the bisection + print(f"threshold: MLflow logging skipped ({exc})") + return result diff --git a/docs/source/solvers/lpse2d/config.md b/docs/source/solvers/lpse2d/config.md index 3d503275..2a30204f 100644 --- a/docs/source/solvers/lpse2d/config.md +++ b/docs/source/solvers/lpse2d/config.md @@ -66,6 +66,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) | +| `lpse-` bases | | The original LPSE profiles (`ZakharovSolver::backgroundDensityShape`), shape in `linear`, `exp`, `gaussian`, `inverse-power`, `quadratic`, `qd`, `gd`, `file`: `min`/`max` at `min_location`/`max_location` (x; `min_location_y`/`max_location_y` optional, `geometry: cartesian|spherical`), `sg_order` (LPSE `sgOrder`, default 2) for `gaussian`/`inverse-power`/`gd`, `central_density` for `quadratic`, `dip_depth`/`dip_width`/`dip_offset` for the `qd` (parabolic) and `gd` (super-Gaussian) dips on a linear ramp, `origin` (x of LPSE's box centre for the quadratic and the dips, default the box centre), `max_density` clip (default 1.25), `file` (`.npy`, text table, or LPSE grid file) for `lpse-file`. See `helpers._lpse_density_profile` for the formulas | | `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 @@ -102,8 +103,11 @@ Simulation grid parameters. Note: Grid values use physical units as strings. | Field | Type | Description | |-------|------|-------------| -| `boundary_abs_coeff` | float | Absorbing boundary coefficient | +| `boundary_abs_coeff` | float | Absorbing boundary coefficient (`tanh` profile: the amplitude damping rate, 1/ps, on the layer plateau) | | `boundary_width` | string | Width of absorbing boundary layer with unit | +| `boundary_profile` | string | (optional, default `tanh`) Shape of the absorbing layers. `tanh` is the MATLAB envelope (rise `boundary_width / 5`). `exp` is the original LPSE profile (`absorbingBoundaries.cpp`): rate `boundary_max_rate * (e^{lambda s/L} - 1) / (e^lambda - 1)` with `s` the distance into the layer measured from the edge cell (so the edge cell carries the full rate and the layer spans `boundary_width / dx` cells); with both axes absorbing the two rates combine by `max`, as in LPSE | +| `boundary_max_rate` | float | (`exp` profile) Peak amplitude damping rate in 1/ps (default `200`, LPSE `lw.abc.maxDampingRate`). The IAW absorber uses half of it unless `terms.iaw.boundary_max_rate` is set (LPSE `IawSolver` default 100) | +| `boundary_lambda` | float | (`exp` profile) Exponential steepness (default `7`, LPSE `abc.lambda`) | | `low_pass_filter` | float | Low-pass filter cutoff as fraction of kmax (0-1) | | `dealias` | string | Shape of the anti-aliasing mask: `isotropic` (default) or `shifted-band` | | `dt` | string | Timestep with unit | @@ -195,8 +199,8 @@ save: | Field | Type | Description | |-------|------|-------------| | `t` | object | Temporal save configuration | -| `x` | object | Optional spatial subsampling in x | -| `y` | object | Optional spatial subsampling in y | +| `x` | object | Optional interpolated (coarser) save grid in x; requires `y` | +| `y` | object | Optional interpolated (coarser) save grid in y; requires `x` | #### t (temporal) @@ -210,13 +214,29 @@ save: | Field | Type | Description | |-------|------|-------------| -| `dx` | string | Spatial resolution for saved data, with unit | +| `dx` | string | Target spacing of the saved field grid in x, with unit; must not be finer than `grid.dx` | #### y (optional) | Field | Type | Description | |-------|------|-------------| -| `dy` | string | Spatial resolution for saved data, with unit | +| `dy` | string | Target spacing of the saved field grid in y, with unit; must not be finer than the grid's `dy` (= `grid.dx`) | + +Without `x`/`y`, `fields.xr` and `k-fields.xr` are saved on the full simulation grid. With them, the +fields are saved on a coarser grid: + +- The save grid has `n_save = round(L / dx)` cell-centred points per axis (`L` is the box length, + `xmax - xmin` or `ymax - ymin`), so the actual spacing is `L / n_save` (printed at setup). A spacing that + would give more points than the simulation grid raises a `ValueError` at setup. In a quasi-1D box + (`ny = 1`) the y save grid is the single row. +- Every field is linearly interpolated onto the save grid in real space: the EPW potential and its fields + `phi`, `ex`, `ey` are computed on the full grid (spectrally, as the solver does) and then interpolated, as + are `e0_*`, `e1_*` and the IAW fields. This is point sampling, not a filtered downsample: content above + the save grid's Nyquist wavenumber `pi / (L / n_save)` aliases. +- The saved x-space values are the full-grid save's values linearly interpolated to the save points. +- `k-fields.xr` is the DFT on the save grid, on its own k axes (`n_save` points, the same spacing `2 pi / L` as the + full grid), scaled by `(nx * ny) / (nx_save * ny_save)` so that a resolved mode has the same coefficient as in + the full-grid save (its phase is referenced to the first save point instead of the first grid point). ## mlflow @@ -247,9 +267,16 @@ The main laser pump for TPD/SRS simulations. | `envelope` | object | Spatiotemporal envelope | | `delta_omega_max` | float | Maximum frequency spread (optional) | | `num_colors` | int | Number of laser colors (optional) | +| `angle` | float | (default `0`) in-plane angle of incidence from +x in degrees (LPSE `laser.N.direction`). The static pump is the single grid mode nearest to `k0 (cos a, sin a)` (both components snapped, as LPSE's `makeStaticField`), polarized perpendicular to the snapped k; the spectral pump injector (shared by the spectral light solver and `terms.epw.solver: combined`) launches at the y-snapped transverse wavenumber with the x group velocity setting the flux. Not with `speckle` or the FD injector | +| `polarization` | string | (default `p`) only in-plane polarization is representable (LPSE `polarization = 0`); lpse2d carries no out-of-plane field component | +| `beams` | list | LPSE `laser.N.*` beamlets for the static pump and the spectral injector (also the evolved pump of the combined solver): `[{intensity, angle, phase, delta_omega}]` (intensities are normalised to fractions; every beam carries every color; `angle` in degrees, `phase` in radians, `delta_omega` as a fraction of w0). A single beam's `phase` and `delta_omega` add to every color's on every pump path (static, FD and spectral evolved, combined). The deck translator converts LPSE's `laser.N.phase` from degrees and takes `laser.N.wavelengthShift` (Angstrom) over `frequencyShift` when it is non-zero, as LPSE does | +| `beam_width` | string | transverse (y) standard deviation of the injected beams (LPSE `laser.N.width`); with `beam_sg_order` (default 2, Gaussian) and `beam_offset` (y of the beam centre) it sets the super-Gaussian `exp(-((y-y0)^2/2 sigma^2)^(p/2))` of the spectral injector | +| `kap_bandwidth` | float | (default `0`) Kubo-Anderson phase bandwidth as a fraction of w0 (LPSE `laser.bandwidth.KAP.frequency`): the pump phase jumps to a new uniform value every `2 pi / (kap_bandwidth w0)`, deterministically from `kap_seed` and the beam index. Static pump and spectral injector; with `terms.light.pump_depletion` the FD injector does not apply it and setup refuses it | +| `pulse_file` | string | LPSE `laser.pulse.file`: a two-column text table `(t_ps, relative amplitude)` interpolated in time and multiplied onto the pump amplitude. Static pump and spectral injector; with `terms.light.pump_depletion` the FD injector does not apply it and setup refuses it | | `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` | +| `injector_width` | string | (pump depletion with `terms.light.solver: spectral`, optional) Gaussian width of the smooth pump injector, with unit. Default one local wavelength `2 pi / k0(n_inject)` | #### envelope @@ -330,11 +357,17 @@ propagates toward the low-density side while backscatter growth amplifies it aga | `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 | +| `injector_width` | string | (`terms.light.solver: spectral`, optional) Gaussian width of the smooth seed injector, with unit. Default one local wavelength `2 pi / k1(n_inject)` | 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). +With `terms.epw.solver: combined` the seed is the spectral solver's smooth injector added to the +transverse (`y`) component of the combined field at its carrier `wp0` (LPSE's Raman class has +`W1 = wpe` in combined mode), so `delta_omega` is a fraction of `wp0` and the injector must sit +below the envelope density (`n < n_env`). + Example: ```yaml drivers: @@ -361,8 +394,15 @@ Physics terms configuration. | Field | Type | Description | |-------|------|-------------| | `pump_depletion` | bool | (default `false`) Evolve the pump `E0` with the staggered FD envelope solver instead of prescribing it analytically. The pump is launched at `x = xmin + drivers.E0.offset`; active SRS and TPD terms each add their reciprocal pump coupling, and both act when both instabilities are enabled. Requires at least one of `terms.epw.source.srs`/`tpd` and `terms.epw.boundary.x: absorbing`; incompatible with `drivers.E0.speckle`. Enables true net-flux `laser_reflectivity` / `laser_transmissivity` / `laser_absorbed_frac` metrics and above-threshold saturation | -| `coupling` | str | (default `explicit`; `pump_depletion` with `terms.epw.source.srs` only) How the SRS exchange between `E0` and `E1` is integrated inside each light sub-step. `explicit` keeps the MATLAB staggered real/imaginary update, which treats the part of the exchange proportional to `Im(laplacian phi)` with an explicit Euler step: the light pair grows by `1 + sin^2(arg laplacian phi) (Omega dt_l)^2 / 2` per sub-step, `Omega = e \|laplacian phi\| / (4 me sqrt(w0 w1))`, for any `dt_l`. That is invisible at small EPW amplitude but manufactures light (and, through the SRS/TPD sources, EPW) energy once the pump is depleted and `Omega dt_l` reaches ~0.005, ending in a one-step NaN. `rotation` Strang-splits each sub-step as [exact exchange rotation over `dt_l/2`] [staggered propagation with the exchange off] [rotation over `dt_l/2`]; the rotation `exp(tau M) = cos(Omega tau) I + sin(Omega tau)/Omega M` conserves the light action `w1 \|E0\|^2 + w0 \|E1\|^2` pointwise and is stable for any `dt_l`. The TPD pump term and the IAW detuning stay in the staggered update under both settings. See `tests/test_lpse2d/test_light_coupling.py` | +| `coupling` | str | (default `explicit`; `pump_depletion` with `terms.epw.source.srs` only) How the SRS exchange between `E0` and `E1` is integrated inside each light sub-step. `explicit` keeps the MATLAB staggered real/imaginary update, which treats the part of the exchange proportional to `Im(laplacian phi)` with an explicit Euler step: the light pair grows by `1 + sin^2(arg laplacian phi) (Omega dt_l)^2 / 2` per sub-step, `Omega = e \|laplacian phi\| / (4 me sqrt(w0 w1))`, for any `dt_l`. That is invisible at small EPW amplitude but manufactures light (and, through the SRS/TPD sources, EPW) energy once the pump is depleted and `Omega dt_l` reaches ~0.005, ending in a one-step NaN. `rotation` Strang-splits each sub-step as [exact exchange rotation over `dt_l/2`] [staggered propagation with the exchange off] [rotation over `dt_l/2`]; the rotation `exp(tau M) = cos(Omega tau) I + sin(Omega tau)/Omega M` conserves the light action `w1 \|E0\|^2 + w0 \|E1\|^2` pointwise and is stable for any `dt_l`. The TPD pump term and the IAW detuning stay in the staggered update under both settings. The rotation is pointwise and cannot apply the transverse projection of `transverse_source`, so with the `fd` solver and SRS on `rotation` requires `transverse_source: false` set explicitly (the default `true` is refused; the longitudinal part of `E1` then stays in the exchange) -- use `explicit` or `solver: spectral` to keep the projection. See `tests/test_lpse2d/test_light_coupling.py` | | `filter` | float | (default off; `pump_depletion` only) Isotropic low-pass filter applied to both light fields once per EPW step, keeping `\|k\| <= filter * pi/dx`. The physical light content lies below ~1.2 k0; grid-scale light modes have FD group velocity `c^2 sin(k dx)/(w dx) -> 0`. Diagnostic/numerical-hygiene option | +| `tpd_projection` | bool | (default `true`; `pump_depletion` with `terms.epw.source.tpd`) Take the transverse (divergence-free) part of `E_h div(E_h)` in k-space before it acts on the pump, `F_T = F - k (k . F)/k^2` mode by mode (LPSE `LwSolver::makeExyzDivE`), so the TPD pump term never injects a longitudinal component into the light field. Both pump components receive the term | +| `tpd_k_filter` | bool | (default `false`) LPSE `lw.kFilter`: restrict the TPD pump term to `\|k\| < 1.2 k0 sqrt(1 - n_min)` | +| `transverse_source` | bool | (default `true`) project the SRS light sources onto their transverse (divergence-free) part every light sub-step (LPSE `raman.takeTransversePartOfSourceTerms`; with the coupled spectral solver the fields themselves are kept transverse; the `fd` solver's `coupling: rotation` cannot apply it and refuses `true`, see `coupling`). No light propagator moves the longitudinal part of E1, so without it that part accumulates the source and pairs with the EPW in a spurious two-wave instability. The spectral solvers also drop every component outside the retained light band, as LPSE does | +| `boundary_max_rate` | float | peak amplitude damping rate (1/ps) of the light fields' absorbing layers with `grid.boundary_profile: exp` (LPSE `{laser|raman}.evolution.abc.maxDampingRate`, default `5e3`; the EPW layer uses `grid.boundary_max_rate`, 200). Light crosses a 3 um layer in 0.01 ps, so at the EPW rate the exp layer reflects ~75 % of the amplitude and an injected pump builds a standing wave between the walls. With the `tanh` profile both use `boundary_abs_coeff` | +| `solver` | string | (default `fd`) Light propagator for the evolved fields (`E1`, and `E0` with `pump_depletion`). `fd`: the MATLAB staggered real/imaginary finite-difference scheme, sub-cycled to its CFL limit. `spectral`: the original LPSE `{laser\|raman}.solver = spectral` path -- per sub-step the x-space scattering potential `exp(i dt Vo)` (detuning + absorption), the coupling sources with forward Euler, and the exact k-space propagator `exp(-i dt c^2 k^2/(2 w))` on the transverse part of the field (the longitudinal part is left unpropagated, as in LPSE), modes outside the retained band zeroed. No CFL limit (`grid.light_substeps` defaults to 1), no grid dispersion, and with `pump_depletion` the SRS exchange is the exact local rotation of `coupling: rotation` Strang-split around the propagation. The injectors become smooth Gaussian sources (`drivers.E0/E1.injector_width`, one local wavelength by default) that launch exactly the requested amplitude with negligible leakage in the wrong direction | +| `max_wavenumber` | float | (optional; `spectral` only) Cap the retained light band at `max_wavenumber * k0` (LPSE `{laser\|raman}.maxWavenumber`). The light band is the anti-aliased band without `terms.epw.max_wavenumber`; only the combined solver takes the smaller of the two caps (LPSE `LightSolver.cpp:3689`) | +| `absorption` | bool or float | (default `false`) Collisional (inverse-bremsstrahlung) absorption of the evolved light: the amplitude decays at `nu (n/nc_w)^2` per wave, `nc_w` its own critical density and `n` including the IAW perturbation (LPSE `calculateScatteringPotential`). `true` takes `nu` from the NRL formula as coded in LPSE, `5.11e10 Z logLambda / (lambda_um^2 Te_keV^1.5) * 1e-12` 1/ps with each wave's own wavelength (`logLambda = 6.68 + ln(lambda_um Te)` for `Te > 0.01 Z^2`, else `9.13 + ln(lambda_um Te^1.5/Z)`); a number is the pump rate at `nc` in 1/ps, the Raman rate scaled by `(w1/w0)^2`. With `terms.epw.solver: combined` the Raman light's carrier is `wp0`: its wavelength is `lambda_0 w0/wp0` and a numeric rate scales by `(wp0/w0)^2`. That rate damps the combined field at `(n/n_env)^2` (`n` with the IAW perturbation) and calibrates its noise; `terms.epw.damping.collisions` does not act on the combined field, as in LPSE. Works with both light solvers (applied per sub-step). LPSE's `resonanceAbsorption` is not implemented | ### epw @@ -375,6 +415,9 @@ Physics terms configuration. | `source` | object | Source terms | | `hyperviscosity` | object | Optional hyperviscosity for numerical stability | | `kinetic real part` | bool | Include kinetic correction to real frequency | +| `max_wavenumber` | float | (optional) LPSE `lw.maxWavenumber`: hard cap `\|k\| < max_wavenumber * k0` (vacuum laser wavenumber) on the retained EPW band, applied on top of `grid.low_pass_filter` / `grid.dealias`. It does not act on the separate light solvers' band (`terms.light.max_wavenumber`); the combined field takes the smaller of the two | +| `solver` | string | (default `separate`) `separate`: the four envelope equations of the MATLAB prototype / LPSE's default spectral path (EPW potential, pump, Raman light each with its own carrier). `combined`: the original LPSE `lw.solver = combined` formulation, which LPSE *requires* whenever TPD and SRS are both on (its users guide: the separate equations "are not valid for simultaneous TPD and SRS", nor "for SRS coupled to IAWs near n_c/4"). One field `E1` enveloped at `wp0` carries the Raman light as its transverse part and the EPW as its longitudinal part; `epw` (the potential `phi_k = i k . E1_k/k^2`) is derived from it every step for the diagnostics, IAWs and HPE. Per light sub-step: the density detuning and collisional damping on the whole field, LPSE's unified source `-i e/(4 me w0) e^{-i(w0 - 2wp0)t} [grad(E0 . E1*) + (1 - w0/wp0) E0 (div E1)*]` (whose longitudinal part is the TPD + SRS source and whose transverse part is the Raman-light generation), the 2x2 longitudinal/transverse propagator (Bohm-Gross dispersion + Landau damping on `k k/k^2 E1`, `exp(-i dt c^2 k^2/(2 wp0))` on the rest), the EPW noise on the longitudinal part, the absorbers; with `terms.light.pump_depletion` the pump advances in the same sub-step, injected by the spectral light solver's injector (beams, angles, transverse profile, KAP bandwidth, pulse table), with the unified depletion `i e/(2 me w0) e^{+i(w0 - 2wp0)t} E1 div E1` (SRS + TPD depletion in one term; `terms.light.tpd_projection` applies the transverse projection LPSE's combined path omits). Requires `terms.light.solver: spectral`, `tpd` and `srs` both on or both off, and excludes `drivers.E2`, `energy_ledger`, `terms.light.coupling/filter`. With `separate` and both instabilities on a warning quotes LPSE's refusal. The Raman-light series (`e1_sq`, reflectivity, fluxes), the saved `E1` fields and the `s1` Poynting maps are the transverse part (the EPW is saved as `epw`); the IAW ponderomotive drive sees the transverse part and the derived EPW separately, the Raman light at its carrier `wp0` (coefficient `1/wp0^2`; also the thermal-filamentation Raman term, critical density `n_env`). The separate light solvers and the separate EPW solver are not built in combined mode (the combined thermal noise is balanced against the Raman absorption, so it needs no EPW Landau damping or collisions) | +| `energy_ledger` | bool | (default `false`) Accumulate, in the state, the change of the EPW energy attributed to every operation of the split step -- `dispersion`, `damping`, `dealias`, `noise`, `detuning`, `boundary`, `reprojection`, `tpd`, `srs`, `driver` -- and report the cumulative values in the default series as `epw_ledger_` (in `epw_energy` units) together with `epw_ledger_closure = epw_energy - sum(channels)`, which stays at `epw_energy(0)` to round-off. This is the per-step energy-budget check LPSE asserts to 0.1 % (`ZakharovSolver.cpp:1700-1712`); here it is exact by construction and the channels are the diagnostic | #### boundary @@ -389,6 +432,11 @@ Physics terms configuration. |-------|------|-------------| | `collisions` | bool or float | Collisional damping. `true` computes from plasma parameters, or specify rate directly | | `landau` | bool | Include Landau damping | +| `landau_form` | string | (default `matlab`) Static Landau rate. `matlab`: the prototype's `sqrt(pi/8) (1 + 1.5 x^2) wp^4/(k^3 vte^3) exp(-3/2 - 1/(2x^2))`, `x = k lambda_D`. `lpse`: the C++ `landauDamping_nonRel`, `sqrt(pi/8) (kde/k)^3 w_k exp(-w_k^2/(2 k^2 vte^2))` with `w_k = wp sqrt(1 + 3x^2)` -- identical exponent, prefactor `sqrt(1 + 3x^2)` vs `(1 + 1.5x^2)` (0.7% at x = 0.3, 3.8% at x = 0.5). `relativistic` (= `relativistic_2d`) / `relativistic_3d`: LPSE's Maxwell-Juettner Bessel-function rates (`landauDamping_rel_2D/3D`); modes with phase velocity above `c` are undamped | +| `landau_lower_threshold` | float | (default `0`, 1/ps) LPSE `lw.landauDamping.lowerThreshold`: modes whose Landau rate is below it are treated as undamped (also in the thermal-noise balance) | +| `landau_multiplier` | float | (default `1`) Static multiplier on the Landau rate (LPSE `LD_multiplier`) | + +The same rate array (form, threshold, multiplier) is used by the EPW step, the HPE calibration, the dissipation diagnostic and the thermal noise source. #### source @@ -397,7 +445,13 @@ Physics terms configuration. | `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 the two-plasmon-decay source. With `terms.light.pump_depletion`, also include its energy-reciprocal feedback on the y-polarized pump | +| `noise_model` | string | (default `flat`) `flat`: the MATLAB source, a kick `dt * noise_amplitude` with a random phase on every retained mode each step. `thermal`: the original LPSE `lw.noise` source (`ZakharovSolver::addNoiseToPotential_fft`), a fluctuation-dissipation kick `D_k = N A / sqrt(1 + k^2 lambda_D^2) * sqrt(1 - exp(-2 gamma_k dt)) / \|k\|` balanced against the frozen analytic Landau + collisional rate `gamma_k` and added after the damping sub-step, so every mode relaxes to the Cerenkov spectrum `<\|E_k\|^2> = A^2/(1 + k^2 lambda_D^2)` (x-space envelope amplitude squared) independently of `dt`; `A = noise_amplitude`. Undamped modes receive no noise; a band with no damping at all is refused, as in LPSE. Switching an existing deck to `thermal` changes its seed level -- the srs-2d-testbed calibration was done with `flat` | +| `noise_calibrate` | bool | (default `false`; `thermal` only) Set `A` from the electron temperature by equipartition -- electric energy `kT/2` per mode over the box volume `V` (`Lz = Ly` for the 2-D box, `V = Lx^3` when `ny = 1`, following LPSE's `deltaK3`): `A = noise_amplitude * sqrt(8 pi kT / V)`, so `noise_amplitude: 1` is the thermal level (LPSE `lw.noise.isCalculated` convention). LPSE's own constant carries an unexplained extra factor in its source; this calibration is the equipartition value and has not been matched to LPSE's number | +| `noise_max_wavenumber` | float | (optional, units of `k0`) LPSE `lw.noise.maxWavenumber`: no noise above `noise_max_wavenumber * k0` | +| `tpd_form` | string | (default `lpse`) `lpse`: the C++ TPD source `i e/(4 me w0) e^{-i(w0 - 2wp0)t} [F(E0 . E*) + (w0/wp0 - 1) i k . F(E0 rho*)/k^2]` with every pump component (`ZakharovSolver::updatePotentialWithTpdSource_fft`). `matlab`: the prototype's `w0 -> 2 wp0` form, `i e/(8 me wp0)` and factor 1. The two are identical at envelope density 0.25 (`w0 = 2 wp0`); at 0.23 the LPSE coefficient is 4% smaller and the charge-density factor is 1.085 | +| `srs_k_filter` | bool | (default `true`) apply a high-k filter to the light fields entering the SRS source, cutting at `srs_k_filter_scale` times the Raman wavenumber at the minimum box density (LPSE `lw.kFilter.enable`, which is *off* by default there; MATLAB `isSuppressHighKSource`). The cutoff assumes the scattered light sits exactly at the envelope frequency, so in a box detuned from the envelope density it can remove the resonant Raman mode entirely (no SRS growth at all at 0.2 nc with envelope density 0.25); translated LPSE decks turn it off. With the separate EPW solver, setup refuses a box whose minimum density reaches the `w1` critical density `(w1/w0)^2` (the filter band is empty there); the combined solver filters the pump instead and is not checked | +| `srs_k_filter_scale` | float | (default `1.2`) multiplier on the SRS source k-filter cutoff (LPSE `lw.kFilter.scale`) | +| `tpd` | bool | Include the two-plasmon-decay source. With `terms.light.pump_depletion`, also include LPSE's pump-depletion term `i e/(2 me w0) e^{+i(w0 - 2wp0)t} [E_h div E_h]_T` on both pump components (twice the EPW-side coefficient, so the pair conserves the total wave energy at envelope density n_c/4; see `terms.light.tpd_projection`) | | `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) @@ -414,12 +468,20 @@ The IAW state follows the MATLAB LPSE split update for fractional ion-density pe | Field | Type | Description | |-------|------|-------------| | `active` | bool | Enable ion-acoustic evolution (default `false`) | +| `solver` | string | (default `explicit`) `explicit`: the MATLAB kick/drift split step (FD Laplacian, stable for `omega_iaw,max * grid.dt < 2`, validated at setup). `spectral`: the original LPSE `iaw.solver = spectral` path -- per k-mode the exact solution of the damped oscillator `dn/dt = -w, dw/dt = cs^2 k^2 n - 2 gamma_k w` (`e^{-gamma dt}` times a `cos/sin(beta dt)` rotation, `beta = sqrt(cs^2 k^2 - gamma^2)`), then collisional damping `e^{-2 nu dt}` on `w`, the ponderomotive kick `dt k^2 PP_k`, and the noise. Unconditionally stable and exact for the acoustic part at any `dt` | | `boundary` | object or null | Per-axis `x`/`y` boundary modes. Defaults to `terms.epw.boundary` | -| `damping.collisions` | float | Density damping rate in `1/ps` (default `1.0e-5`) | +| `boundary_max_rate` | float or null | (`grid.boundary_profile: exp` only) Peak absorber rate for the IAW layer, 1/ps; default half of `grid.boundary_max_rate` (LPSE `IawSolver` default 100 vs 200) | +| `damping.collisions` | float | Collisional damping rate in `1/ps` (default `1.0e-5`): on `n` as `(1 - nu dt)` in the explicit solver, on `div v` as `exp(-2 nu dt)` in the spectral one (LPSE) | | `damping.landau` | float | Dimensionless coefficient in `gamma_iaw(k) = landau * cs * |k|` (default `0.1`); the velocity-divergence equation is damped at `2 gamma_iaw` | +| `damping.landau_form` | string | (default `simplified`) `simplified`: the rate above (LPSE `isSimplified`). `full`: the Z-generalized Krall-Trivelpiece expression used by LPSE's IAW solver, `W_i = W_r sqrt(pi/8) L^{-3/2} [(3/(eta-1))^{3/2} exp(-(3/(eta-1))/(2L)) + sqrt(1/(eta M))]`, `W_r = cs |k|/sqrt(L)`, `L = 1 + k^2 lambda_D^2`, `eta = 1 + 3 Ti/(Z Te)`, `M = (mi/me)/(eta Z)`, with `gamma_iaw = W_i/2`; `damping.landau` is then ignored. The real frequency in the spectral propagator stays `cs |k|`, as in LPSE's spectral branch | | `max_density_perturbation` | float or null | Optional symmetric limiter on `|delta n_i/n_0|`; unlimited by default | +| `flow` | list or null | (`spectral` only) Uniform background flow `[Mach_x, Mach_y]` in units of `cs`: every mode acquires the Doppler phase `e^{-i k . V0 dt}` (LPSE `fluid.velocity`) | +| `stride` | int | (default `1`; `spectral` only) Advance the IAW every `stride` EPW steps with `dt_iaw = stride * grid.dt` (LPSE steps its IAW solver less often than the EPW); the IAW absorber applies its damping over the whole `dt_iaw` | +| `noise` | bool | (default `false`) LPSE `iaw.noise`: a random-phase source on `div v` every IAW step with the fluctuation-dissipation amplitude `noise_amplitude * nx * ny * sqrt(exp(2 dt (gamma_k + nu)) - 1)` on the retained band; requires IAW damping | +| `noise_amplitude` | float | (default `1.0`) Amplitude of the IAW noise (x-space `div v` units per mode) | +| `noise_seed` | int or null | Seed of the IAW noise stream | -The setup validates the explicit acoustic stability condition `omega_iaw,max * grid.dt < 2`. +The setup validates the explicit acoustic stability condition `omega_iaw,max * grid.dt < 2` for `solver: explicit`. ```yaml terms: @@ -432,6 +494,10 @@ terms: max_density_perturbation: 0.1 ``` +#### thermal_filamentation (optional, LPSE `thermalFil.*`) + +`terms.iaw.thermal_filamentation: {laser: bool, raman: bool, lw: bool, nonlocal: bool, conductivity_multiplier: 1.0}` adds the thermal-filamentation source to the ion velocity-divergence equation: inverse-bremsstrahlung heating by the spatially varying part of each enabled wave's intensity (`|E_w|^2 - <|E_w|^2>`), balanced by Spitzer heat conduction, drives the flow through the electron pressure, `d(div v)/dt += Z nu_w(n) (|E_w|^2 - <|E_w|^2>) / (8 pi m_i kappa')` with `kappa'` the Spitzer conductivity over k_B (times `conductivity_multiplier`) and `nu_w` the wave's energy damping rate (light: `2 nu_abs(n_c) (n/n_c)^2`, requires `terms.light.absorption`; EPW: `2 nu_coll n/n_env`, requires `terms.epw.damping.collisions`; with `terms.epw.solver: combined` the EPW is damped by the combined field's Raman absorption instead, `2 nu_R (n/n_env)^2`, requires `terms.light.absorption`, and the EPW collisions play no part). `nonlocal: true` adds LPSE's `k^(4/3)` correction (`1 + (k lambda_nl)^(4/3)`, `lambda_nl = 30 (k_B T_e)^2 / (4 pi e^4 sqrt(Z+1) ln Lambda n_e)`). The source form follows `ZakharovSolver::getThermalFilamentationSource`; the normalization is adept's own. + ### hpe (optional) Hybrid particle evolution, following Follett et al., *Phys. Plasmas* **24**, 102134 (2017): test electrons drawn from the Maxwellian tail are pushed relativistically in the de-enveloped electrostatic field, their spatially averaged velocity distribution is accumulated by exponential moving average, and the Landau damping rate applied by the EPW solver is recomputed from that evolving distribution every step (kinetic inflation + hot-electron generation; Im-only feedback, no nonlinear frequency shift). For `ny == 1` the tracker uses `(x, p_x)`; in a 2-D box it uses `(x, y, p_x, p_y)` and gathers both $E_x$ and $E_y$. In both cases there is one box-wide ensemble, not a particle population at each grid point. HPE requires `terms.epw.damping.landau: true`. The particle push dominates runtime, so HPE runs want a GPU. @@ -456,6 +522,16 @@ At an absorbing particle wall, outgoing particles are thermalized and reinjected | `feedback` | bool | (default `true`) `false` = control run: particles evolve but the damping stays analytic (Follett's control experiment) | | `seed` | int | RNG seed for particle loading and wall re-injection (default `42`) | | `omega_res` | string | Resonance convention for `v_phi(k)`: `"bohm_gross"` (default, matches the analytic rate) or `"wp0"` (bare carrier, as in the paper) | +| `gamma_limit_damping` | float | (default `1500`, 1/ps) upper clip on the applied kinetic rate of the resonant modes (LPSE `hpe.gammaLimit.damping`); with `energy_conservation` it is applied again after the multiplier, so the applied rate respects it (LPSE clips only before its `LD_multiplier`; the energy balance then holds only where the limit does not bind) | +| `gamma_limit_growth` | float | (default `1500`, 1/ps) lower clip `-gamma_limit_growth` on the applied rate of the resonant modes when `allow_growth` is on (LPSE `hpe.gammaLimit.growth`); after the `energy_conservation` multiplier too, as `gamma_limit_damping` | +| `allow_growth` | bool | (default `false`) allow negative (inverse-Landau) rates from an inverted tail; otherwise the rate is clipped at 0 (LPSE `hpe.allowGrowth`). The series `epw_dissipation` is then signed: modes with a negative applied rate count as energy handed from the electrons to the EPW | +| `thermalization_probability` | list | probability that a particle crossing an x / y side is thermalized and re-injected from the tail; otherwise it passes through periodically. Every crossing is counted by the wall-flux instrument, on all four sides, independently of the field boundaries (LPSE particle walls, `hpe.thermalizationProbability`). Default: 1 at absorbing field boundaries, 0 at periodic ones | +| `magnetic_field` | float | (default `0`, tesla) uniform out-of-plane B: the 2-D push rotates the momentum by the cyclotron angle each sub-step (LPSE `hpe.magneticField`, z component) | +| `energy_conservation` | bool | (default `false`) scale the applied Landau rate every step by a global multiplier in [0.1, 10] so that the expected EPW energy loss equals the particles' kinetic-energy gain (LPSE `hpe.enforceEnergyConservation`; adept linearises the per-mode loss as `2 gamma dt`). Reported as `hpe_ld_multiplier` | +| `energy_conservation_steps` | float | (default `1`) running-average length of the multiplier (LPSE `hpe.numStepsToAverageEnergyChange`) | +| `flux_bins` | list | keV edges of the wall-flux instrument (default `[0, 50, 100, 1e9]`): the series `hpe_wall_energy__bin` accumulate the energy (keV per real electron) leaving through each wall in each bin (LPSE `hpe.metrics.flux`). Bins are half-open `[e_i, e_{i+1})`; electrons below the first edge or at/above the last are not counted, so a finite window such as `[50, 100]` reports only 50-100 keV | +| `cone_angle` | float | acceptance half-angle in degrees of the cone-power instrument (`hpe_cone_energy`: cumulative energy leaving inside the cone about `cone_direction`; LPSE `hpe.metrics.power`). `null` = off | +| `cone_direction` | list | (default `[1, 0]`) axis of the acceptance cone | ```yaml terms: @@ -468,6 +544,18 @@ terms: t_start: 2ps ``` +## Diagnostics: Poynting flux and Thomson probes + +`save.fields.poynting: true` adds the light energy-flux density maps `s0_x`, `s0_y` (pump) and `s1_x`, `s1_y` (Raman) to the fields output, `S_j = (c^2/omega) Im(E* . d_j E)` in field-squared times um/ps (`v_g |E|^2` for a plane wave; LPSE `laser.save.S0` / `raman.save.S0`), with `omega` the field's envelope carrier: `w0` for `E0`; for `E1`, `w1 = w0 - wp0` with the separate EPW solver and `wp0` with `terms.epw.solver: combined` (the default series' Raman fluxes and reflectivity use the same carrier). The maps are taken on the save grid; with an interpolated save they need at least two saved x points (setup refuses fewer). `save.thomson: [{k: [kx, ky], bandwidth: 0.1, field: epw}]` adds synthetic Thomson-scattering probes to the default series (LPSE `thomsonScattering.N.wavevector.lw/.iaw` and `bandwidth`): for probe `i`, `thomson_i_re` / `thomson_i_im` are the summed complex amplitude of the EPW potential (or, with `field: iaw`, of the IAW density) over the k-window `|k - k_probe| < bandwidth k0`, and `thomson_i_power` the summed spectral power there; `k` and `bandwidth` are in units of the vacuum laser wavenumber. + +## Checkpoint and restart + +`save.checkpoint: true` (or a path) writes the complete solver state at `grid.tmax` as an `.npz` (`binary/checkpoint.npz` in the run's artifacts when `true`); `restart: {file: }` in a later configuration resumes from it: the state replaces the fresh initial condition, integration starts at the checkpoint time, and the series/fields save axes start there too unless their `tmin` is later. The per-step noise and wall keys are folded in from the time index, so a run split into two resumes reproduces the unbroken run to round-off as long as the split time is a multiple of `grid.dt` (the fixed-step solver would otherwise clip one step differently). The configuration must otherwise be the same (grid, terms, drivers); mismatched state shapes are refused. + +## Absolute-threshold bisection (`adept._lpse2d.threshold`) + +`find_threshold(cfg, intensity_lo, intensity_hi, n_iter=6)` bisects the pump intensity (geometric midpoints) between a stable and an unstable bracket, running the configuration once per point with `ergoExo` and deciding "unstable" from the run's metrics (default: a measurable EPW energy growth fit with a positive rate, `growth_min` adjustable; any callable of the metrics dict can be passed as `criterion`). The bracket needs `0 < intensity_lo < intensity_hi` (a zero lower bound never moves a geometric midpoint); its endpoints are run first and must straddle the criterion, as in LPSE's `AbsoluteThreshold`. The result holds the threshold (midpoint of the final bracket), the bracket and the per-run history; the bisection is logged as an MLflow run (`-bisection`) next to the individual runs. + ### Example: TPD Simulation ```yaml diff --git a/docs/source/solvers/lpse2d/overview.md b/docs/source/solvers/lpse2d/overview.md index af292ab9..c48c0a24 100644 --- a/docs/source/solvers/lpse2d/overview.md +++ b/docs/source/solvers/lpse2d/overview.md @@ -28,7 +28,7 @@ With `terms.epw.source.srs` on, the default time series includes OSIRIS-comparab | 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_dissipation` | total EPW energy handed to electrons per ps (Landau + collisional, the solver's own rates -- with `terms.epw.solver: combined` the Landau rate plus the combined field's Raman absorption `nu_R (n/n_env + dn)^2` on the EPW energy density, not the EPW collisions; includes the kinetic half; signed with `terms.hpe.allow_growth`, negative = net transfer from the electrons to the EPW) | 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` | @@ -101,6 +101,8 @@ The HPE particle tracker supplies self-consistent Landau-damping feedback and ho When IAWs are active, `fields.xr` and `k-fields.xr` also contain `iaw_density` and `iaw_velocity_divergence`; `series.xr` includes `iaw_density_sq` and `iaw_density_abs_max`. +The plasma-wave entries are built from the solver state, the k-space potential `phi_k` (numpy FFT convention). In a full-grid save (no `save.fields.x` / `y`), `k-fields.xr` holds `phi = phi_k`, `ex = -1j * kx * phi_k` and `ey = -1j * ky * phi_k` (fftshifted), and `fields.xr` holds their `ifft2`, so its `ex`, `ey` are $-\partial_x \phi$, $-\partial_y \phi$ of its `phi`, the same field the solver evolves and that the `series.xr` quantities `e_sq` and `epw_energy` integrate. An interpolated save does not satisfy these identities: its `fields.xr` `phi`, `ex`, `ey` are the full-grid fields interpolated to the save points, and its `k-fields.xr` is their DFT on the save grid scaled by `(nx * ny) / (nx_save * ny_save)`, so neither the `ifft2` of its `k-fields.xr` nor a derivative of its coarse `phi` reproduces its `fields.xr` (see [`save.fields`](config.md#fields)). Field saves written before this was fixed (October 2026) carry an extra factor `4 / (nx * ny)` on `fields.xr` `ex` and `ey` only. + **`plots/`**, per field `k`: | File | Contents | diff --git a/tests/test_lpse2d/test_combined.py b/tests/test_lpse2d/test_combined.py new file mode 100644 index 00000000..acdbb4bc --- /dev/null +++ b/tests/test_lpse2d/test_combined.py @@ -0,0 +1,723 @@ +"""The combined TPD + SRS solver (terms.epw.solver: combined, LPSE lw.solver = combined). + +- With the sources off, the longitudinal part of the combined field evolves exactly + as the separate spectral EPW solver and the transverse part as the wp0-enveloped + light propagator. +- The unified source's longitudinal part is the separate solver's TPD source plus SRS + source (as a potential), and its transverse part is the separate Raman coupling; + the unified pump depletion is the SRS depletion plus the TPD depletion. All of these + are algebraic identities at envelope density n_c/4, checked on random band-limited + fields. +- Config validation follows LPSE (TPD and SRS both on or both off). +""" + +from copy import deepcopy + +import numpy as np +import pytest +import yaml +from jax import numpy as jnp + + +def _finish(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _cfg(*, solver="combined", sources=True, density=0.2, pump_depletion=False, periodic=True, noise=False, **extra): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["density"] = {"basis": "uniform", "val": density} + cfg["grid"].update( + { + "boundary_width": "0.6um", + "dt": "1fs", + "dx": "0.1um", + "xmax": "6.4um", + "tmax": "10fs", + "ymax": "1.6um", + "ymin": "-1.6um", + "low_pass_filter": 0.6, + } + ) + cfg["terms"]["epw"]["boundary"] = {"x": "periodic" if periodic else "absorbing", "y": "periodic"} + cfg["terms"]["epw"]["damping"] = {"collisions": False, "landau": True} + cfg["terms"]["epw"]["density_gradient"] = True + cfg["terms"]["epw"]["source"].update({"noise": noise, "noise_seed": 1, "tpd": sources, "srs": sources}) + cfg["terms"]["epw"]["solver"] = solver + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": pump_depletion} + cfg["terms"].update(extra) + return _finish(cfg) + + +def _random_phi(cfg, seed, band=None): + rng = np.random.default_rng(seed) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + phi_k = rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny)) + mask = np.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"]) + if band is not None: + mask = mask * band + return jnp.asarray(phi_k * mask) + + +def _random_transverse(cfg, seed, band=None): + """Random divergence-free (nx, ny, 2) field on the retained band, in x-space.""" + rng = np.random.default_rng(seed) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + k_sq = kx**2 + ky**2 + mask = np.asarray(cfg["grid"]["low_pass_filter_grid"]) * np.where(k_sq > 0, 1.0, 0.0) + if band is not None: + mask = mask * band + a_k = (rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) * mask + # transverse in 2-D: E_k = a_k (-ky, kx)/|k| + k_safe = np.sqrt(np.where(k_sq > 0, k_sq, 1.0)) + ex_k, ey_k = -ky / k_safe * a_k, kx / k_safe * a_k + return jnp.asarray(np.stack([np.fft.ifft2(ex_k), np.fft.ifft2(ey_k)], axis=-1)) + + +def _field_from_phi(phi_k, cfg): + kx = jnp.asarray(cfg["grid"]["kx"])[:, None] + ky = jnp.asarray(cfg["grid"]["ky"])[None, :] + return jnp.stack([jnp.fft.ifft2(-1j * kx * phi_k), jnp.fft.ifft2(-1j * ky * phi_k)], axis=-1) + + +def _phi_from_source_field(source, cfg): + """i k . S_k / k^2 on the band: the potential source equivalent to a field source.""" + kx = jnp.asarray(cfg["grid"]["kx"])[:, None] + ky = jnp.asarray(cfg["grid"]["ky"])[None, :] + k_sq = kx**2 + ky**2 + one_over = jnp.where(k_sq > 0, 1.0 / jnp.where(k_sq > 0, k_sq, 1.0), 0.0) + band = jnp.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"]) + return 1j * (kx * jnp.fft.fft2(source[..., 0]) + ky * jnp.fft.fft2(source[..., 1])) * one_over * band + + +def test_free_evolution_matches_the_separate_solvers(): + from adept._lpse2d.core.combined import CombinedSolver + from adept._lpse2d.core.epw import SpectralEPWSolver + + cfg = _cfg(sources=False, density=0.23) + combined = CombinedSolver(cfg) + epw = SpectralEPWSolver(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + phi0 = _random_phi(cfg, 5) + e_t0 = _random_transverse(cfg, 6) + E1 = _field_from_phi(phi0, cfg) + e_t0 + E0 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + # the derived potential of the composed field is the longitudinal input + np.testing.assert_allclose(np.asarray(combined.potential(E1)), np.asarray(phi0), rtol=1e-10, atol=1e-12) + + n_steps = 7 + y = {"E0": E0, "E1": E1, "epw": phi0} + phi_sep = phi0 + for i in range(n_steps): + t = i * cfg["grid"]["dt"] + _, y["E1"], y["epw"] = combined(t, y, {}, lambda tt: E0) + phi_sep = epw(t, {"epw": phi_sep, "E0": E0, "E1": E0}, None) + np.testing.assert_allclose(np.asarray(y["epw"]), np.asarray(phi_sep), rtol=1e-9, atol=1e-12 * np.abs(phi_sep).max()) + + # transverse part: exact wp0-enveloped light propagation with the density detuning + derived = cfg["units"]["derived"] + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + k_sq = kx**2 + ky**2 + t_total = n_steps * cfg["grid"]["dt"] + phase = np.exp( + -1j * t_total * (derived["wp0"] / 2.0 * (0.23 / 0.25 - 1.0) + derived["c"] ** 2 * k_sq / (2.0 * derived["wp0"])) + ) + e_t = np.asarray(combined.transverse(y["E1"])) + for comp in range(2): + expected = np.fft.ifft2(np.fft.fft2(np.asarray(e_t0)[..., comp]) * phase) + np.testing.assert_allclose(e_t[..., comp], expected, rtol=1e-9, atol=1e-12 * np.abs(expected).max()) + + +def test_unified_source_is_the_separate_sources_at_quarter_critical(): + """At envelope density 0.25 (w0 = 2 wp0, no carrier mismatch) the unified source + -i e/(4 me w0) [grad(E0 . E1*) + (1 - w0/wp0) E0 (div E1)*] splits exactly into the + separate solver's terms: longitudinal part = TPD source (from E1_L) + SRS source (from + E1_T, whose (n/n_env) factor LPSE's combined source does not carry); transverse part = + the Raman coupling +i e/(4 me w0) rho* E0 (transverse-projected).""" + from adept._lpse2d.core.combined import CombinedSolver, longitudinal_transverse + from adept._lpse2d.core.epw import SpectralEPWSolver + from adept._lpse2d.core.spectral_light import SpectralRamanLight + + n = 0.2 + cfg = _cfg(density=n) + combined = CombinedSolver(cfg) + epw = SpectralEPWSolver(cfg) + raman = SpectralRamanLight(cfg) + assert combined.rho_factor == pytest.approx(-1.0) and epw.tpd_rho_factor == pytest.approx(1.0) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + rng = np.random.default_rng(8) + # keep the light part inside the separate SRS source's k-filter (1.2 k1) so that filter is a no-op + band_light = np.asarray(epw.E1_filter[..., 0]) + phi = _random_phi(cfg, 9) + e_t = _random_transverse(cfg, 10, band=band_light) + E1 = _field_from_phi(phi, cfg) + e_t + E0 = jnp.asarray(rng.normal(size=(nx, ny, 2)) + 1j * rng.normal(size=(nx, ny, 2))) + E0 = jnp.fft.ifft2(jnp.fft.fft2(E0, axes=(0, 1)) * combined.E0_filter[..., None], axes=(0, 1)) + t = 0.3 + + source = combined.unified_source(t, E0, E1) + (lx_k, ly_k), (tx_k, ty_k) = longitudinal_transverse(source, combined.kx, combined.ky, combined.one_over_k_sq) + + # longitudinal part as a potential source vs TPD (E_L) + SRS (E_T) / (n/n_env) + phi_source = _phi_from_source_field(source, cfg) + ex, ey = epw.phi_k_to_e_fields(phi) + # the separate SRS source is returned unmasked (the next step's band filter removes + # the out-of-band product content); the unified source is band-limited when built + band = jnp.asarray(combined.band) + expected = (epw.calc_tpd_source(t, phi, ex, ey, E0) + epw.calc_srs_source(E0, e_t) / (n / 0.25)) * band + np.testing.assert_allclose( + np.asarray(phi_source), np.asarray(expected), rtol=1e-9, atol=1e-11 * np.abs(expected).max() + ) + + # transverse part vs the Raman coupling, transverse-projected + lap_phi = jnp.fft.ifft2(-epw.k_sq * phi) + coupling = raman.srs_coeff * jnp.conj(lap_phi)[..., None] * E0 + _, (cx_k, cy_k) = longitudinal_transverse(coupling, combined.kx, combined.ky, combined.one_over_k_sq) + band = np.asarray(band) + np.testing.assert_allclose( + np.asarray(tx_k) * band, np.asarray(cx_k) * band, rtol=1e-9, atol=1e-11 * np.abs(cx_k).max() + ) + np.testing.assert_allclose( + np.asarray(ty_k) * band, np.asarray(cy_k) * band, rtol=1e-9, atol=1e-11 * np.abs(cy_k).max() + ) + + +def test_unified_depletion_is_srs_plus_tpd_depletion_at_quarter_critical(): + from adept._lpse2d.core.combined import CombinedSolver, longitudinal_transverse + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + + cfg = _cfg(density=0.2, pump_depletion=True, periodic=False) + combined = CombinedSolver(cfg) + separate = SpectralCoupledLight(cfg) + assert combined.w1 == pytest.approx(combined.w0 / 2.0) + phi = _random_phi(cfg, 11) + e_t = _random_transverse(cfg, 12) + E1 = _field_from_phi(phi, cfg) + e_t + t = 0.2 + unified = combined.unified_depletion(t, E1) + lap_phi = jnp.fft.ifft2(-combined.k_sq * phi) + srs_dep = separate.srs_depletion_coeff0 * lap_phi[..., None] * e_t + _, (sx_k, sy_k) = longitudinal_transverse(srs_dep, combined.kx, combined.ky, combined.one_over_k_sq) + srs_dep_t = jnp.stack([jnp.fft.ifft2(sx_k), jnp.fft.ifft2(sy_k)], axis=-1) + expected = srs_dep_t + separate.calc_tpd_depletion(t, phi) + np.testing.assert_allclose( + np.asarray(unified), np.asarray(expected), rtol=1e-9, atol=1e-11 * np.abs(expected).max() + ) + + +def test_combined_configuration_is_validated_like_lpse(): + with pytest.raises(ValueError, match="both on or both off"): + cfg = _cfg(sources=False) + cfg["terms"]["epw"]["source"]["srs"] = True + from adept._lpse2d.core.combined import CombinedSolver + + CombinedSolver(cfg) + with pytest.raises(ValueError, match="both on or both off"): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + raw = yaml.safe_load(fi) + raw["terms"]["epw"]["solver"] = "combined" + raw["terms"]["light"] = {"solver": "spectral"} + raw["terms"]["epw"]["source"].update({"tpd": True, "srs": False}) + _finish(raw) + with pytest.raises(ValueError, match="spectral"): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + raw = yaml.safe_load(fi) + raw["terms"]["epw"]["solver"] = "combined" + raw["terms"]["epw"]["source"].update({"tpd": True, "srs": True}) + raw["terms"]["light"] = {"solver": "fd"} + _finish(raw) + + +def test_combined_split_step_runs_with_pump_depletion_iaw_noise_and_reports_transverse_light(): + from adept._lpse2d.core.vector_field import SplitStep + from adept._lpse2d.helpers import get_default_save_func, get_save_quantities + + cfg = _cfg( + density=0.22, pump_depletion=True, periodic=False, noise=True, iaw={"active": True, "solver": "spectral"} + ) + # the full-grid field save (tpd.yaml's 0.08 um save spacing is finer than this 0.1 um grid, which + # the interpolated save refuses since PR #388) + cfg["save"]["fields"].pop("x", None) + cfg["save"]["fields"].pop("y", None) + cfg = get_save_quantities(cfg) + step = SplitStep(cfg) + assert step.epw_solver == "combined" + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + state = { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128).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), + "iaw_density": jnp.zeros((nx, ny)), + "iaw_velocity_divergence": jnp.zeros((nx, ny)), + } + pump_args = { + **cfg["drivers"]["E0"]["derived"], + "delta_omega": jnp.zeros(1), + "intensities": jnp.ones((1, ny)), + "phases": jnp.zeros((1, ny)), + } + for i in range(6): + state = step(jnp.asarray(i * cfg["grid"]["dt"]), dict(state), {"drivers": {"E0": pump_args}}) + assert all(bool(jnp.all(jnp.isfinite(v))) for v in state.values()) + E1 = state["E1"].view(jnp.complex128) + assert float(jnp.max(jnp.abs(E1))) > 0.0 # noise seeded the longitudinal part + # the derived potential equals i k . E1_k / k^2 + np.testing.assert_allclose( + np.asarray(state["epw"].view(jnp.complex128)), np.asarray(step.combined.potential(E1)), rtol=1e-12, atol=0 + ) + out = get_default_save_func(cfg)["func"](0.0, state, None) + assert "e1_sq" in out and bool(jnp.isfinite(out["e1_sq"])) + # the pump has started to fill the box + assert float(jnp.max(jnp.abs(state["E0"].view(jnp.complex128)))) > 0.0 + + +# ---- noise kick at the light sub-step (PR #372 review) --------------------------------------------- +# Tolerances fixed before the runs: the steady state is a mean over ~700 retained modes x 150 +# correlated samples, so both the n_sub-independence ratio and the agreement with the analytic +# random-walk level of diagnostics.expected_noise_energy are set at 15 %. Before the fix the kick +# was built for the EPW step and applied every sub-step: ratio 8 / 2 sub-steps ~ 4 for both models. + + +def _noise_cfg(n_sub, model, landau=True): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + # n = n_env: the combined field's damping exp(-nu_R dt (n/n_env)^2) is then the uniform rate the + # noise kick is balanced against; nu_R is the Raman absorption at the combined carrier, + # terms.light.absorption (wp0/w0)^2 = 80/ps x n_env = 20/ps (the EPW collisions do not act on it) + cfg["density"] = {"basis": "uniform", "val": 0.25} + cfg["grid"].update( + { + "boundary_width": "0.6um", + "dt": "1fs", + "dx": "0.1um", + "xmax": "6.4um", + "tmax": "10fs", + "ymax": "1.6um", + "ymin": "-1.6um", + "low_pass_filter": 0.6, + "light_substeps": n_sub, + } + ) + cfg["terms"]["epw"]["boundary"] = {"x": "periodic", "y": "periodic"} + cfg["terms"]["epw"]["damping"] = {"collisions": False, "landau": landau} + cfg["terms"]["epw"]["density_gradient"] = True + cfg["terms"]["epw"]["source"].update( + {"noise": True, "noise_model": model, "noise_amplitude": 1.0, "noise_seed": 3, "tpd": False, "srs": False} + ) + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False, "absorption": 80.0} + return _finish(cfg) + + +def _combined_noise_steady_state(cfg, n_steps=450, n_warm=300): + """Mean of sum_k k^2 |phi_k|^2 over the last n_steps - n_warm EPW steps from a quiet start.""" + import jax + + from adept._lpse2d.core.combined import CombinedSolver + + solver = CombinedSolver(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + y = {"E0": jnp.zeros((nx, ny, 2), complex), "E1": jnp.zeros((nx, ny, 2), complex)} + step = jax.jit(lambda t, y: solver(t, y, {}, E0_fn=lambda t: y["E0"])) + samples = [] + for i in range(n_steps): + E0, E1, phi_k = step(i * cfg["grid"]["dt"], y) + y = {"E0": E0, "E1": E1} + if i >= n_warm: + samples.append(float(jnp.sum(solver.k_sq * jnp.abs(phi_k) ** 2))) + return solver, np.mean(samples) + + +def test_combined_thermal_noise_needs_only_the_raman_absorption(): + """The combined field's damping is the Raman absorption, so thermal noise with the EPW Landau + damping and collisions both off and terms.light.absorption > 0 is a valid combined deck (PR #372 + review): SplitStep sets up (before the fix it also built the never-stepped separate EPW solver, + whose thermal noise refused the deck for lack of EPW damping), and the combined noise kick is + nonzero and balanced against nu_R.""" + from adept._lpse2d.core.epw import noise_step + from adept._lpse2d.core.vector_field import SplitStep + + cfg = _noise_cfg(1, "thermal", landau=False) + step = SplitStep(cfg) + assert step.epw is None + d = cfg["units"]["derived"] + assert noise_step(cfg)[1] == pytest.approx(80.0 * (d["wp0"] / d["w0"]) ** 2, rel=1e-12) + assert float(jnp.max(jnp.abs(step.combined.noise_kick))) > 0.0 + + +@pytest.mark.parametrize("model", ["thermal", "flat"]) +def test_combined_noise_steady_state_is_independent_of_the_light_substeps(model): + """The combined solver adds its noise every light sub-step, so the kick is built for that step + (LwSolver::addNoise_combinedSolver_spectral(dt)); the seed level cannot depend on n_sub.""" + levels = {} + for n_sub in (2, 8): + cfg = _noise_cfg(n_sub, model) + solver, levels[n_sub] = _combined_noise_steady_state(cfg) + assert solver.n_sub == n_sub + assert levels[8] / levels[2] == pytest.approx(1.0, rel=0.15) + + +@pytest.mark.parametrize("model", ["thermal", "flat"]) +def test_combined_noise_matches_the_growth_fit_floor(model): + """diagnostics.expected_noise_energy integrates the kick the combined solver applies (at its + light sub-step), so the late-time floor equals the simulated steady state.""" + from adept._lpse2d.core.epw import noise_kick_spectrum, noise_step + from adept._lpse2d.diagnostics import expected_noise_energy + + cfg = _noise_cfg(4, model) + dt_noise, nu = noise_step(cfg) + assert dt_noise == pytest.approx(cfg["grid"]["dt"] / 4, rel=1e-12) + assert nu == pytest.approx(20.0, rel=1e-12) + solver, level = _combined_noise_steady_state(cfg) + np.testing.assert_allclose( + np.asarray(solver.noise_kick), noise_kick_spectrum(cfg, dt=dt_noise, nu_coll=nu), rtol=1e-12, atol=0 + ) + derived = cfg["units"]["derived"] + grid = cfg["grid"] + prefactor = 0.25 * grid["dx"] * derived["x_norm"] * derived["e_norm"] ** 2 / (grid["nx"] * grid["ny"] ** 2) + assert prefactor * level / expected_noise_energy(cfg, 1.0e3) == pytest.approx(1.0, rel=0.15) + + +# ---- Raman seed on the combined path (PR #372 review) ----------------------------------------------- +# Tolerances fixed before the run, the same as the separate solver's seed test +# (test_spectral_light.test_spectral_seed_injector_launches_the_exact_amplitude_leftward): bulk +# amplitude 1 %, ripple 2 %, wavenumber 1 %. Before the fix the combined solver never read +# drivers.E1, so with the pump off and noise off E1 stayed identically zero. That test's 1 % +# rightward-leakage bound is not carried over: in this geometry (n/n_env = 0.8, region > 3.2 sigma +# right of the injector) the separate spectral solver itself leaves 1.4 % at 180 steps (combined +# 1.3 %), so the bound measures the geometry, not the wiring. The seed source is instead checked +# to equal the separate solver's (an identity at n_env = n_c/4, where w1 = wp0; rtol 1e-12). + + +def test_combined_seed_injector_launches_the_exact_amplitude_leftward(): + """drivers.E1 seeds the transverse part of the combined field at the combined carrier wp0 + (LPSE's Raman class has W1 = wpe in combined mode): the smooth injector launches + |E1_y| = A / eps^(1/4), eps = 1 - n/n_env, as a leftward wave at k = sqrt(wp0^2 - wpe^2)/c, + and the purely transverse seed leaves the derived potential at zero.""" + from adept._lpse2d.core.vector_field import SplitStep + + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["units"]["envelope density"] = 0.25 + cfg["density"] = {"basis": "uniform", "val": 0.2} + cfg["grid"].update({"xmax": "20um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": True, "srs": True}) + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False} + cfg["drivers"]["E1"] = {"intensity": "1.0e+12W/cm^2", "delta_omega": 0.0, "turn_on_time": "10fs", "offset": "7um"} + cfg = _finish(cfg) + del cfg["drivers"]["E0"] + + step = SplitStep(cfg) + assert step.epw_solver == "combined" and step.combined.seed_enabled + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + state = { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128).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), + } + args = {"drivers": {"E1": cfg["drivers"]["E1"]["derived"]}} + for i in range(180): + state = step(jnp.asarray(i * cfg["grid"]["dt"]), dict(state), args) + + e1y = np.asarray(state["E1"].view(jnp.complex128))[:, 0, 1] + x = np.asarray(cfg["grid"]["x"]) + derived = cfg["units"]["derived"] + n_over_env = 0.2 / 0.25 + eps = 1.0 - n_over_env + k_seed = derived["wp0"] * np.sqrt(eps) / derived["c"] + bulk = slice(int(np.argmin(np.abs(x - 5.0))), int(np.argmin(np.abs(x - 10.0)))) + expected = cfg["drivers"]["E1"]["derived"]["amplitude"] / eps**0.25 + np.testing.assert_allclose(np.mean(np.abs(e1y[bulk])), expected, rtol=0.01) + assert np.std(np.abs(e1y[bulk])) < 0.02 * expected + k_measured = np.polyfit(x[bulk], np.unwrap(np.angle(e1y[bulk])), 1)[0] + assert k_measured < 0 + np.testing.assert_allclose(abs(k_measured), k_seed, rtol=0.01) + # the seed is transverse: no longitudinal (EPW) part, so the derived potential stays zero + assert float(jnp.max(jnp.abs(state["epw"]))) < 1e-12 * expected + + +def test_combined_seed_source_equals_the_separate_seed_source_at_quarter_critical(): + """At n_env = n_c/4 the combined carrier wp0 equals the Raman envelope frequency w1 = w0 - wp0, + so the combined seed source is the separate spectral solver's, term for term.""" + from adept._lpse2d.core.combined import CombinedSolver + from adept._lpse2d.core.spectral_light import SpectralRamanLight + + def cfg_for(solver): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["units"]["envelope density"] = 0.25 + cfg["density"] = {"basis": "uniform", "val": 0.2} + cfg["grid"].update({"xmax": "20um", "tmax": "10fs", "ymax": "0.5um", "ymin": "-0.5um"}) + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + both = solver == "combined" + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": both, "srs": True}) + if both: + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False} + cfg["drivers"]["E1"] = { + "intensity": "1.0e+12W/cm^2", + "delta_omega": 0.003, + "turn_on_time": "10fs", + "offset": "7um", + "yw": "0.6um", + } + return _finish(cfg) + + combined = CombinedSolver(cfg_for("combined")) + cfg_s = cfg_for("separate") + separate = SpectralRamanLight(cfg_s) + assert combined.wp0 == pytest.approx(separate.w1, rel=1e-12) + assert combined.i1 == separate.i1 + seed_args = cfg_s["drivers"]["E1"]["derived"] + for t in (0.004, 0.017, 0.05): + np.testing.assert_allclose( + np.asarray(combined.calc_seed_source(t, seed_args)), + np.asarray(separate.calc_seed_source(t, seed_args)), + rtol=1e-12, + atol=0, + ) + + +def test_combined_seed_injector_refuses_an_evanescent_position(): + """The combined field's light is evanescent where n >= n_env (carrier wp0); setup says so.""" + from adept._lpse2d.core.combined import CombinedSolver + + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["grid"].update({"xmax": "20um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": True, "srs": True}) + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False} + cfg["drivers"]["E1"] = {"intensity": "1.0e+12W/cm^2", "offset": "7um"} + cfg = _finish(cfg) # uniform n = n_env = 0.2 + with pytest.raises(ValueError, match="evanescent"): + CombinedSolver(cfg) + + +# ---- the combined carrier wp0 outside the combined solver (PR #372 review round 4) ----------------- + + +def _off_quarter_cfg(solver, n_env, density, *, iaw=None, absorption=False, wavelength=None, seed=False): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["units"]["envelope density"] = n_env + if wavelength is not None: + cfg["units"]["laser_wavelength"] = wavelength + cfg["density"] = {"basis": "uniform", "val": density} + cfg["grid"].update({"xmax": "20um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + combined = solver == "combined" + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": combined, "srs": True}) + cfg["terms"]["epw"]["solver"] = solver + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False, "absorption": absorption} + if iaw is not None: + cfg["terms"]["iaw"] = iaw + if seed: + cfg["drivers"]["E1"] = { + "intensity": "1.0e+12W/cm^2", + "delta_omega": 0.0, + "turn_on_time": "10fs", + "offset": "7um", + } + return _finish(cfg) + + +def test_combined_iaw_drive_thermal_term_and_raman_absorption_use_the_combined_carrier(): + """With terms.epw.solver: combined the IAW sees the Raman light at the combined carrier wp0: its + ponderomotive coefficient is 1/wp0^2, the thermal-filamentation Raman term's critical density is + n_env and its absorption rate the NRL value at lambda_0 w0/wp0 (a numeric rate scales by + (wp0/w0)^2). At n_env = 0.2 (w1 = 1.24 wp0). Exact identities, rtol 1e-12, fixed before the + run; the separate solver keeps w1.""" + from adept._lpse2d.core.iaw import IonAcousticWave + from adept._lpse2d.core.raman import light_absorption_rates + + iaw = { + "active": True, + "solver": "spectral", + "damping": {"landau": 0.1, "collisions": 0.0}, + "thermal_filamentation": {"laser": True, "raman": True}, + } + cfg_c = _off_quarter_cfg("combined", 0.2, 0.15, iaw=iaw, absorption=True) + cfg_s = _off_quarter_cfg("separate", 0.2, 0.15, iaw=iaw, absorption=True) + d = cfg_c["units"]["derived"] + assert abs(d["w1"] / d["wp0"] - 1.0) > 0.2 + nx, ny = cfg_c["grid"]["nx"], cfg_c["grid"]["ny"] + amp = 2.0e-3 + e1 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128).at[..., 1].set(amp) + zero_phi = jnp.zeros((nx, ny), dtype=jnp.complex128) + zero_e0 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + for cfg, w_expected in ((cfg_c, d["wp0"]), (cfg_s, d["w1"])): + wave = IonAcousticWave(cfg) + assert wave.w_raman == w_expected + drive = np.asarray(wave.ponderomotive_drive(zero_phi, zero_e0, e1)) + np.testing.assert_allclose(drive, wave.ponderomotive_prefactor * amp**2 / w_expected**2, rtol=1e-12) + + # NRL Raman rate at the combined wavelength lambda_0 w0/wp0 (the pump rate of a deck at that wavelength) + rate0_c, rate1_c = light_absorption_rates(cfg_c) + lam_c = 0.351 * d["w0"] / d["wp0"] + np.testing.assert_allclose( + rate1_c, + light_absorption_rates(_off_quarter_cfg("combined", 0.2, 0.15, absorption=True, wavelength=f"{lam_c}um"))[0], + rtol=1e-12, + ) + # the thermal Raman/laser coefficient ratio: (rate1 / rate0) / nc1^2, nc1 = n_env combined + wave_c = IonAcousticWave(cfg_c) + ratio = np.asarray(wave_c.thermal_coeff["raman"]) / np.asarray(wave_c.thermal_coeff["laser"]) + np.testing.assert_allclose(ratio, (rate1_c / rate0_c) / 0.2**2, rtol=1e-12) + # a numeric pump rate scales by the carrier ratio squared + for solver, w in (("combined", d["wp0"]), ("separate", d["w1"])): + rate0, rate1 = light_absorption_rates(_off_quarter_cfg(solver, 0.2, 0.15, absorption=0.5)) + assert rate0 == 0.5 + np.testing.assert_allclose(rate1, 0.5 * (w / d["w0"]) ** 2, rtol=1e-12) + + +def test_seeded_combined_deck_above_the_separate_cutoff_sets_up_and_runs(): + """n_env = 0.3: the combined light is cut off at n_env = 0.3 nc, the separate Raman light at + (w1/w0)^2 = 0.205 nc. A seeded box at 0.25 nc is valid for the combined solver; before the fix + the separate SRS k-filter check refused the box and the never-stepped separate Raman solver + refused the seed. Through setup and SplitStep (tolerance fixed before the run): the separate + light solvers are not built and after 80 steps the seed reaches the bulk at A / eps^(1/4), + eps = 1 - n/n_env, to 5 %. The separate solver still refuses the same box.""" + from adept._lpse2d.core.vector_field import SplitStep + + cfg = _off_quarter_cfg("combined", 0.3, 0.25, seed=True) + del cfg["drivers"]["E0"] + step = SplitStep(cfg) + assert step.raman is None and step.coupled_light is None and step.combined.seed_enabled + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + state = { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128).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), + } + args = {"drivers": {"E1": cfg["drivers"]["E1"]["derived"]}} + for i in range(80): + state = step(jnp.asarray(i * cfg["grid"]["dt"]), dict(state), args) + e1y = np.asarray(state["E1"].view(jnp.complex128))[:, 0, 1] + x = np.asarray(cfg["grid"]["x"]) + x_inj = float(x[step.combined.i1]) + bulk = (x > x_inj - 4.0) & (x < x_inj - 2.0) + eps = 1.0 - 0.25 / 0.3 + np.testing.assert_allclose( + np.mean(np.abs(e1y[bulk])), cfg["drivers"]["E1"]["derived"]["amplitude"] / eps**0.25, rtol=0.05 + ) + + with pytest.raises(ValueError, match="Raman critical density"): + _off_quarter_cfg("separate", 0.3, 0.25, seed=True) + + +def test_combined_dissipation_diagnostic_uses_the_combined_damping(): + """epw_dissipation in combined mode is the combined solver's own damping of the EPW (PR #372 review): + the Landau rate in k-space plus the Raman absorption nu_R (n/n_env + dn)^2 on the EPW energy + density in x-space -- not the EPW collisions, which the combined solver ignores. Uniform box + (n = 0.15 nc, n_env = 0.2, dn = 0.05), terms.light.absorption 300/ps, EPW collisions on: with + Landau off, epw_dissipation = 2 (2 - n/n_env) L_c epw_energy, L_c = (1 - exp(-2 nu_R (n/n_env + dn)^2 + dt)) / dt; with Landau on, plus the k-space Landau term (Parseval). Exact identities, rtol 1e-12, + fixed before the run. Before the fix the first case reported the EPW collisional loss.""" + from adept._lpse2d.core.epw import analytic_landau_rate + from adept._lpse2d.helpers import get_default_save_func, get_save_quantities + + def run(landau): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["units"]["envelope density"] = 0.2 + cfg["density"] = {"basis": "uniform", "val": 0.15} + cfg["grid"].update({"xmax": "12.8um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["epw"]["boundary"] = {"x": "periodic", "y": "periodic"} + cfg["terms"]["epw"]["damping"] = {"collisions": True, "landau": landau} + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": True, "srs": True}) + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False, "absorption": 300.0} + cfg = get_save_quantities(_finish(cfg)) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + rng = np.random.default_rng(13) + phi_k = 1e-3 * (rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) + phi_k = phi_k * np.asarray(cfg["grid"]["low_pass_filter_grid"]) + zeros = jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64) + state = { + "epw": jnp.asarray(phi_k).view(jnp.float64), + "E0": zeros, + "E1": zeros, + "iaw_density": jnp.full((nx, ny), 0.05), + } + out = get_default_save_func(cfg)["func"](0.0, state, None) + return cfg, phi_k, out + + cfg, _, out = run(landau=False) + d, grid = cfg["units"]["derived"], cfg["grid"] + dt = grid["dt"] + assert d.get("nu_coll", 0.0) > 0.0 # the EPW collisions are configured, and must not count + nu_r = 300.0 * (d["wp0"] / d["w0"]) ** 2 + loss_c = (1.0 - np.exp(-2.0 * nu_r * (0.75 + 0.05) ** 2 * dt)) / dt + factor = 2.0 * (2.0 - 0.75) + assert float(out["epw_dissipation"]) == pytest.approx(factor * loss_c * float(out["epw_energy"]), rel=1e-12) + + cfg, phi_k, out = run(landau=True) + grid = cfg["grid"] + nx, ny = grid["nx"], grid["ny"] + k_sq = np.asarray(grid["kx"])[:, None] ** 2 + np.asarray(grid["ky"])[None, :] ** 2 + loss_l = (1.0 - np.exp(-2.0 * np.asarray(analytic_landau_rate(cfg)) * dt)) / dt + assert np.any(loss_l * np.abs(phi_k) > 0.0) + prefactor = 0.25 * grid["dx"] * d["x_norm"] * d["e_norm"] ** 2 + landau_part = prefactor * factor * np.sum(loss_l * k_sq * np.abs(phi_k) ** 2) / (nx * ny**2) + want = landau_part + factor * loss_c * float(out["epw_energy"]) + assert float(out["epw_dissipation"]) == pytest.approx(want, rel=1e-12) + + +def test_the_combined_field_is_damped_at_the_raman_absorption(): + """LPSE damps the combined field with the Raman light's absorption at the combined carrier and the + field's own critical density n_env, not with the EPW collisions (PR #372 review). Collisions + off, terms.light.absorption 300/ps (Raman rate at wp0: 300 (wp0/w0)^2), n = 0.15 nc, n_env = 0.2, + a periodic box: one EPW step multiplies a transverse E1 plane wave by exp(-nu_R dt (n/n_env)^2), + with a uniform iaw_density dn by exp(-nu_R dt (n/n_env + dn)^2), and noise_step returns nu_R. + Exact identities, rtol 1e-10, fixed before the run; before the fix the field kept its amplitude.""" + from adept._lpse2d.core.combined import CombinedSolver + from adept._lpse2d.core.epw import noise_step + + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["units"]["envelope density"] = 0.2 + cfg["density"] = {"basis": "uniform", "val": 0.15} + cfg["grid"].update({"xmax": "12.8um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["epw"]["boundary"] = {"x": "periodic", "y": "periodic"} + cfg["terms"]["epw"]["damping"] = {"collisions": False, "landau": False} + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": True, "srs": True}) + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False, "absorption": 300.0} + cfg = _finish(cfg) + d = cfg["units"]["derived"] + nu_r = 300.0 * (d["wp0"] / d["w0"]) ** 2 + assert noise_step(cfg)[1] == pytest.approx(nu_r, rel=1e-12) + solver = CombinedSolver(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + x = np.asarray(cfg["grid"]["x"]) + k = float(np.asarray(cfg["grid"]["kx"])[8]) + zero = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + plane = zero.at[..., 1].set(jnp.asarray(1.0e-3 * np.exp(1j * k * x)[:, None] * np.ones((1, ny)))) + epw = jnp.zeros((nx, ny), dtype=jnp.complex128) + dt = cfg["grid"]["dt"] + for dn, n_eff in ((None, 0.75), (0.05, 0.75 + 0.05)): + y = {"epw": epw, "E0": zero, "E1": plane} + if dn is not None: + y["iaw_density"] = jnp.full((nx, ny), dn) + _, E1, _ = solver(0.0, y, {}, lambda t: zero) + decay = np.abs(np.asarray(E1)[..., 1]) / 1.0e-3 + np.testing.assert_allclose(decay, np.exp(-nu_r * dt * n_eff**2), rtol=1e-10) diff --git a/tests/test_lpse2d/test_diagnostics_lpse.py b/tests/test_lpse2d/test_diagnostics_lpse.py new file mode 100644 index 00000000..b8db010c --- /dev/null +++ b/tests/test_lpse2d/test_diagnostics_lpse.py @@ -0,0 +1,212 @@ +"""LPSE-style diagnostics: Poynting-flux fields (save.fields.poynting) and Thomson probes (save.thomson).""" + +from copy import deepcopy + +import numpy as np +import yaml +from jax import numpy as jnp + + +def _finish(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_save_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + cfg = get_save_quantities(cfg) + return cfg + + +def _cfg(**save): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["density"] = {"basis": "uniform", "val": 0.2} + cfg["grid"].update({"dx": "0.1um", "xmax": "12.8um", "ymax": "3.2um", "ymin": "-3.2um", "dt": "1fs", "tmax": "2fs"}) + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": False}) + # the full-grid field save (tpd.yaml's 0.08 um save spacing is finer than this 0.1 um grid, which + # the interpolated save refuses since PR #388) + cfg["save"]["fields"].pop("x", None) + cfg["save"]["fields"].pop("y", None) + cfg["save"].update(save) + return _finish(cfg) + + +def test_poynting_of_a_plane_wave_is_group_velocity_times_intensity(tmp_path): + from adept._lpse2d.helpers import make_field_xarrays + + cfg = _cfg() + cfg["save"]["fields"]["poynting"] = True + d = cfg["units"]["derived"] + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + kx = np.asarray(cfg["grid"]["kx"]) + k0 = kx[6] # a resolved grid mode + x = np.asarray(cfg["grid"]["x"]) + amp = 2.0e-3 + e0 = np.zeros((1, nx, ny, 2), dtype=np.complex128) + e0[0, :, :, 1] = amp * np.exp(1j * k0 * x)[:, None] + state = { + "epw": np.zeros((1, nx, ny), dtype=np.complex128).view(np.float64), + "E0": e0.view(np.float64), + "E1": np.zeros((1, nx, ny, 2), dtype=np.complex128).view(np.float64), + } + (tmp_path / "binary").mkdir() + _, fields = make_field_xarrays(cfg, np.array([0.0]), state, str(tmp_path)) + v_g = d["c"] ** 2 * k0 / d["w0"] + s0x = fields["s0_x"].values[0, 4:-4, :] # away from the one-sided gradient at the edges + np.testing.assert_allclose(s0x, v_g * amp**2, rtol=2e-2) + assert np.max(np.abs(fields["s0_y"].values)) < 1e-9 * v_g * amp**2 + assert np.max(np.abs(fields["s1_x"].values)) == 0.0 + + +def test_thomson_probe_picks_the_mode_in_its_window(): + cfg = _cfg(thomson=[{"k": [0.5, 0.0], "bandwidth": 0.05}, {"k": [-0.5, 0.0], "bandwidth": 0.05}]) + save_func = cfg["save"]["default"]["func"] + d = cfg["units"]["derived"] + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + kx = np.asarray(cfg["grid"]["kx"]) + k_probe = 0.5 * d["w0"] / d["c"] + ix = int(np.argmin(np.abs(kx - k_probe))) + phi_k = jnp.zeros((nx, ny), dtype=jnp.complex128).at[ix, 0].set(3.0 + 4.0j) + y = { + "epw": phi_k, + "E0": jnp.zeros((nx, ny, 2), dtype=jnp.complex128), + "E1": jnp.zeros((nx, ny, 2), dtype=jnp.complex128), + } + out = save_func(0.0, y, {"drivers": {k: v["derived"] for k, v in cfg["drivers"].items()}}) + np.testing.assert_allclose(float(out["thomson_0_re"]), 3.0) + np.testing.assert_allclose(float(out["thomson_0_im"]), 4.0) + np.testing.assert_allclose(float(out["thomson_0_power"]), 25.0) + assert float(out["thomson_1_power"]) == 0.0 + + +def _combined_cfg(envelope_density=0.25, density=0.2): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = deepcopy(yaml.safe_load(fi)) + cfg["units"]["envelope density"] = envelope_density + cfg["density"] = {"basis": "uniform", "val": density} + cfg["grid"].update({"dx": "0.1um", "xmax": "12.8um", "ymax": "3.2um", "ymin": "-3.2um", "dt": "1fs", "tmax": "2fs"}) + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": True, "srs": True}) + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": False} + cfg["save"]["fields"]["poynting"] = True + # the full-grid save path (both save paths take E1 through the same transverse projection, + # helpers._raman_light_part, before any interpolation) + cfg["save"]["fields"].pop("x", None) + cfg["save"]["fields"].pop("y", None) + return _finish(cfg) + + +def _saved_fields(cfg, e1, tmp_path): + """Run the field save function on one state and build the xarrays from its output.""" + from adept._lpse2d.helpers import make_field_xarrays + + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + y = { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128).view(jnp.float64), + "E0": jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64), + "E1": jnp.asarray(e1).view(jnp.float64), + } + saved = cfg["save"]["fields"]["func"](0.0, y, None) + state = {k: np.asarray(v)[None] for k, v in saved.items()} + (tmp_path / "binary").mkdir(exist_ok=True) + _, fields = make_field_xarrays(cfg, np.array([0.0]), state, str(tmp_path)) + return saved, fields + + +def test_combined_raman_flux_maps_exclude_the_epw(tmp_path): + """The combined solver's E1 carries the EPW as its longitudinal part; the saved E1 fields and + the s1 flux maps are the Raman light (transverse part) only, as the default series. Tolerances + fixed before the run: a pure longitudinal mode leaves < 1e-9 of v_g |A|^2 in s1, and a mixed + field's s1 equals its transverse part's to 1e-9.""" + cfg = _combined_cfg() + d = cfg["units"]["derived"] + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + kx = np.asarray(cfg["grid"]["kx"]) + k = kx[6] + x = np.asarray(cfg["grid"]["x"]) + amp = 2.0e-3 + wave = amp * np.exp(1j * k * x)[:, None] * np.ones((1, ny)) + longitudinal = np.zeros((nx, ny, 2), dtype=np.complex128) + longitudinal[..., 0] = wave # E parallel to k: a travelling plasma wave, no light + transverse = np.zeros((nx, ny, 2), dtype=np.complex128) + transverse[..., 1] = 0.5 * wave + + scale = d["c"] ** 2 * k / d["w1"] * amp**2 + saved, fields = _saved_fields(cfg, longitudinal, tmp_path) + assert np.max(np.abs(np.asarray(saved["E1"]).view(np.complex128))) < 1e-9 * amp + assert np.max(np.abs(fields["s1_x"].values)) < 1e-9 * scale + assert np.max(np.abs(fields["s1_y"].values)) < 1e-9 * scale + + _, fields_mixed = _saved_fields(cfg, longitudinal + transverse, tmp_path) + _, fields_t = _saved_fields(cfg, transverse, tmp_path) + np.testing.assert_allclose(fields_mixed["s1_x"].values, fields_t["s1_x"].values, rtol=0, atol=1e-9 * scale) + assert np.max(np.abs(fields_t["s1_x"].values)) > 0.1 * scale + + +def test_combined_raman_flux_uses_the_combined_carrier_away_from_quarter_critical(tmp_path): + """The combined solver envelopes E1 at wp0 (CombinedSolver.prop_T), so its Raman flux is + (c^2/wp0) Im(E* . grad E); the separate carrier w1 = w0 - wp0 agrees with it only at envelope + density n_c/4 (PR #372 review). At n_env = 0.2, n = 0.15 nc the two differ by 24 %. Tolerances + fixed before the run: the s1 map of a transverse plane wave equals c^2 k |A|^2 / wp0 to 2e-2 + (the central-difference error at this k, as for s0 above); the default series' Raman fluxes and + reflectivity of the carrier-wp0 light mode equal c sqrt(eps) |A|^2 / I0 and + sqrt(eps) |A|^2 / E0_source^2, eps = 1 - n/n_env, to 1e-9.""" + cfg = _combined_cfg(envelope_density=0.2, density=0.15) + d = cfg["units"]["derived"] + assert abs(d["w1"] / d["wp0"] - 1.0) > 0.2 + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + x = np.asarray(cfg["grid"]["x"]) + amp = 2.0e-3 + + k = np.asarray(cfg["grid"]["kx"])[6] + e1 = np.zeros((nx, ny, 2), dtype=np.complex128) + e1[..., 1] = amp * np.exp(1j * k * x)[:, None] + _, fields = _saved_fields(cfg, e1, tmp_path) + s1x = fields["s1_x"].values[0, 4:-4, :] # away from the one-sided gradient at the edges + np.testing.assert_allclose(s1x, d["c"] ** 2 * k / d["wp0"] * amp**2, rtol=2e-2) + + eps = 1.0 - 0.15 / 0.2 + k_mode = d["wp0"] / d["c"] * np.sqrt(eps) + save_func = cfg["save"]["default"]["func"] + args = {"drivers": {name: v["derived"] for name, v in cfg["drivers"].items()}} + for sign, key in ((-1.0, "reflected_flux"), (1.0, "backrefl_flux")): + e1 = np.zeros((nx, ny, 2), dtype=np.complex128) + e1[..., 1] = amp * np.exp(1j * sign * k_mode * x)[:, None] + y = { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128).view(jnp.float64), + "E0": jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64), + "E1": jnp.asarray(e1).view(jnp.float64), + } + out = save_func(0.0, y, args) + np.testing.assert_allclose(float(out[key]), d["c"] * np.sqrt(eps) * amp**2 / d["I0_code"], rtol=1e-9) + np.testing.assert_allclose(float(out["reflectivity"]), np.sqrt(eps) * amp**2 / d["E0_source"] ** 2, rtol=1e-9) + + +def test_poynting_refuses_a_single_point_x_save_grid(): + """An interpolated save of one x point (save.fields.x.dx the box length) cannot take the flux maps' + d/dx: with save.fields.poynting it is refused at setup, not after the solve (PR #372 review); two or + more points are accepted.""" + import pytest + + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + base = yaml.safe_load(fi) + base["density"] = {"basis": "uniform", "val": 0.2} + base["grid"].update( + {"dx": "0.1um", "xmax": "12.8um", "ymax": "3.2um", "ymin": "-3.2um", "dt": "1fs", "tmax": "2fs"} + ) + base["terms"]["epw"]["source"].update({"noise": False, "tpd": False}) + single = deepcopy(base) + single["save"]["fields"].update({"x": {"dx": "12.8um"}, "y": {"dy": "0.2um"}, "poynting": True}) + with pytest.raises(ValueError, match="at least two saved x points"): + _finish(single) + two = deepcopy(base) + two["save"]["fields"].update({"x": {"dx": "6.4um"}, "y": {"dy": "0.2um"}, "poynting": True}) + assert _finish(two)["save"]["fields"]["x"]["ax"].size == 2 diff --git a/tests/test_lpse2d/test_energy_ledger.py b/tests/test_lpse2d/test_energy_ledger.py new file mode 100644 index 00000000..8312661f --- /dev/null +++ b/tests/test_lpse2d/test_energy_ledger.py @@ -0,0 +1,172 @@ +"""The per-operation EPW energy ledger (terms.epw.energy_ledger) closes exactly. + +LPSE asserts that its EPW energy budget closes to 0.1 % per step +(ZakharovSolver.cpp:1700-1712). Here the closure is exact to round-off by +construction, so the test checks that W(t) - W(0) equals the sum of the channels, +that the sign of each channel is what physics demands, and that the ledger is +reported in the default series in epw_energy units. +""" + +from copy import deepcopy + +import numpy as np +import pytest +import yaml +from jax import numpy as jnp + + +def _finish(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _tpd_cfg(noise=True): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update( + {"boundary_width": "0.6um", "dt": "1fs", "dx": "0.1um", "tmax": "10fs", "ymax": "1.6um", "ymin": "-1.6um"} + ) + cfg["density"] = {"basis": "linear", "gradient scale length": "30um", "max": 0.27, "min": 0.23} + cfg["terms"]["epw"]["source"].update({"noise": noise, "noise_seed": 3, "noise_amplitude": 1e-4, "tpd": True}) + cfg["terms"]["epw"]["energy_ledger"] = True + cfg["terms"]["epw"]["damping"]["collisions"] = True + return _finish(cfg) + + +def _srs_cfg(): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update({"boundary_width": "2um", "xmax": "16um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["light"] = {"pump_depletion": True, "coupling": "rotation", "transverse_source": False} + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["terms"]["epw"]["source"].update({"noise_seed": 4, "noise_amplitude": 1e-4}) + cfg["terms"]["epw"]["energy_ledger"] = True + return _finish(cfg) + + +def _state(cfg, seed): + from adept._lpse2d.core.epw import LEDGER_CHANNELS, LEDGER_KEY + + rng = np.random.default_rng(seed) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + phi = (rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) * 1e-3 + phi *= np.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"]) + E0 = 0.3 * cfg["units"]["derived"]["E0_source"] * (rng.normal(size=(nx, ny, 2)) + 1j * rng.normal(size=(nx, ny, 2))) + state = { + "epw": jnp.asarray(phi).view(jnp.float64), + "E0": jnp.asarray(E0).view(jnp.float64), + "E1": jnp.asarray(0.1 * E0).view(jnp.float64), + LEDGER_KEY: jnp.zeros(len(LEDGER_CHANNELS)), + } + return state + + +def _pump_args(cfg): + ny = cfg["grid"]["ny"] + return { + **cfg["drivers"]["E0"]["derived"], + "delta_omega": jnp.zeros(1), + "intensities": jnp.ones((1, ny)), + "phases": jnp.zeros((1, ny)), + } + + +def _run(cfg, state, n_steps): + from adept._lpse2d.core.vector_field import SplitStep + + step = SplitStep(cfg) + args = {"drivers": {"E0": _pump_args(cfg)}} + dt = cfg["grid"]["dt"] + for i in range(n_steps): + state = step(jnp.asarray(i * dt), dict(state), args) + return step, state + + +def _energy(step, state): + return float(step.epw.energy(state["epw"].view(jnp.complex128))) + + +@pytest.mark.parametrize("make_cfg", [_tpd_cfg, _srs_cfg], ids=["tpd+noise+absorbers", "srs+pump_depletion"]) +def test_ledger_closes_to_round_off(make_cfg): + from adept._lpse2d.core.epw import LEDGER_CHANNELS + + cfg = make_cfg() + state0 = _state(cfg, 1) + from adept._lpse2d.core.vector_field import SplitStep + + w0 = float(SplitStep(cfg).epw.energy(state0["epw"].view(jnp.complex128))) + step, state = _run(cfg, state0, 30) + ledger = dict(zip(LEDGER_CHANNELS, np.asarray(state["epw_ledger"]), strict=False)) + w_end = _energy(step, state) + total_flow = sum(abs(v) for v in ledger.values()) + abs(w0) + abs(w_end) + assert abs(w_end - w0 - sum(ledger.values())) < 1e-11 * total_flow, ledger + + assert ledger["dispersion"] == pytest.approx(0.0, abs=1e-12 * total_flow) + assert ledger["detuning"] == pytest.approx(0.0, abs=1e-12 * total_flow) + assert ledger["damping"] < 0.0 + assert ledger["boundary"] < 0.0 + assert ledger["dealias"] <= 0.0 and ledger["reprojection"] <= 1e-12 * total_flow + # the noise channel is 2 Re + |kick|^2 per step: only its long-run + # average is positive, its sign over 30 steps is random + assert ledger["noise"] != 0.0 + assert ledger["driver"] == 0.0 + if cfg["terms"]["epw"]["source"]["tpd"]: + assert ledger["tpd"] != 0.0 and ledger["srs"] == 0.0 + else: + assert ledger["srs"] != 0.0 and ledger["tpd"] == 0.0 + + +def test_ledger_reported_in_epw_energy_units(): + from adept._lpse2d.core.epw import LEDGER_CHANNELS + from adept._lpse2d.helpers import get_default_save_func, get_save_quantities + + cfg = _tpd_cfg() + # the full-grid field save (tpd.yaml's 0.08 um save spacing is finer than this 0.1 um grid, which + # the interpolated save refuses since PR #388) + cfg["save"]["fields"].pop("x", None) + cfg["save"]["fields"].pop("y", None) + cfg = get_save_quantities(cfg) + save = get_default_save_func(cfg)["func"] + step, state = _run(cfg, _state(cfg, 2), 10) + out = save(0.0, state, None) + for name in LEDGER_CHANNELS: + assert f"epw_ledger_{name}" in out + total = sum(float(out[f"epw_ledger_{name}"]) for name in LEDGER_CHANNELS) + w0 = float(save(0.0, _state(cfg, 2), None)["epw_energy"]) + assert float(out["epw_energy"]) - w0 == pytest.approx(total, rel=1e-9) + assert float(out["epw_ledger_closure"]) == pytest.approx(w0, rel=1e-9) + # the fields save function must not try to interpolate the ledger + fields = cfg["save"]["fields"]["func"](0.0, state, None) + assert "epw_ledger" not in fields + + +def test_ledger_off_leaves_state_and_series_unchanged(): + cfg = _tpd_cfg() + cfg["terms"]["epw"]["energy_ledger"] = False + from adept._lpse2d.core.vector_field import SplitStep + from adept._lpse2d.helpers import get_default_save_func, get_save_quantities + + state = _state(cfg, 3) + del state["epw_ledger"] + step = SplitStep(cfg) + out = step(jnp.asarray(0.0), dict(state), {"drivers": {"E0": _pump_args(cfg)}}) + assert "epw_ledger" not in out + # the full-grid field save (tpd.yaml's 0.08 um save spacing is finer than this 0.1 um grid, which + # the interpolated save refuses since PR #388) + cfg["save"]["fields"].pop("x", None) + cfg["save"]["fields"].pop("y", None) + cfg = get_save_quantities(cfg) + series = get_default_save_func(cfg)["func"](0.0, out, None) + assert not any(k.startswith("epw_ledger") for k in series) diff --git a/tests/test_lpse2d/test_field_save.py b/tests/test_lpse2d/test_field_save.py new file mode 100644 index 00000000..e06b8b29 --- /dev/null +++ b/tests/test_lpse2d/test_field_save.py @@ -0,0 +1,355 @@ +"""Tests for the interpolated field save (``save.fields.x.dx`` / ``save.fields.y.dy``). + +An interpolated save puts every saved field on a coarser, cell-centred grid of +``n_save = round(L / dx_save)`` points by linear interpolation in x-space. The full-grid +save (no ``save.fields.x`` / ``y``) is the reference: sampling a full-grid field at the +save coordinates with the same linear interpolation must reproduce the interpolated save. +""" + +import numpy as np +import pytest +import yaml +from scipy.interpolate import RegularGridInterpolator + +GRID_DX_UM = 0.05 +X_LENGTH_UM = 6.4 +Y_HALF_WIDTH_UM = {"2d": 0.8, "1d": 0.02} # "1d" collapses the box to ny = 1 +EXPECTED_FIELDS = { + "phi", + "ex", + "ey", + "e0_x", + "e0_y", + "e1_x", + "e1_y", + "background_density", + "iaw_density", + "iaw_velocity_divergence", +} +KX = r"kx ($kc\omega_0^{-1}$)" +KY = r"ky ($kc\omega_0^{-1}$)" + +# Round-off budget, in units of the saved dtype's machine epsilon, for two evaluations of the +# same linear interpolation of the same data (one inside the jitted save, one here in numpy): +# * query-coordinate rounding: |df/dx| * |x| * eps / 2 <= (pi / dx) * L * eps / 2 * max|f| +# = pi * nx / 2 * eps ~ 200 eps for nx = 128 (band-limited field, |k| <= pi / dx), +# * the x-space FFT of phi (XLA vs numpy FFT): ~ log2(nx * ny) eps ~ 12 eps, +# * the interpolation weights: a few eps. +# ~ 220 eps in total; 2000 eps leaves a 9x margin and is still > 1e10 below the O(1) defect. +ROUNDOFF_EPS = 2000.0 + + +def _box_cfg(geometry: str, save_spacing: str | None, run_name: str) -> dict: + """A small TPD+SRS+IAW box: noise-seeded EPW, static uniform pump, 6 fs.""" + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + half = Y_HALF_WIDTH_UM[geometry] + cfg["density"] = {"basis": "uniform", "val": 0.24} + cfg["grid"].update( + { + "boundary_width": "1um", + "dt": "1fs", + "dx": f"{GRID_DX_UM}um", + "xmax": f"{X_LENGTH_UM}um", + "tmax": "6fs", + "ymax": f"{half}um", + "ymin": f"{-half}um", + } + ) + cfg["save"]["fields"] = {"t": {"tmin": "0fs", "tmax": "6fs", "dt": "3fs"}} + if save_spacing is not None: + cfg["save"]["fields"]["x"] = {"dx": save_spacing} + cfg["save"]["fields"]["y"] = {"dy": save_spacing} + cfg["terms"]["epw"]["source"].update({"noise": True, "noise_seed": 5, "tpd": True, "srs": True}) + cfg["terms"]["iaw"] = { + "active": True, + "boundary": {"x": "absorbing", "y": "periodic"}, + "damping": {"collisions": 1.0e-5, "landau": 0.1}, + "max_density_perturbation": 0.1, + } + cfg["drivers"]["E0"]["params"] = {"phases": {"seed": 42}} + cfg["mlflow"] = {"experiment": "test-lpse2d-field-save", "run": run_name} + return cfg + + +def _prepared_cfg(cfg: dict) -> dict: + """Run the cfg through the setup steps up to the save functions, without MLflow.""" + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_save_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return get_save_quantities(cfg) + + +def _run(cfg: dict) -> dict: + from adept import ergoExo + + exo = ergoExo() + modules = exo.setup(cfg) + _, ppo, _ = exo(modules) + return ppo + + +def _sample_linear(values: np.ndarray, x_src, y_src, x_q, y_q) -> np.ndarray: + """Linear interpolation of (nt, nx, ny) data onto the (x_q, y_q) tensor grid.""" + if len(y_src) == 1: + sampled = [np.interp(x_q, x_src, v[:, 0].real) + 1j * np.interp(x_q, x_src, v[:, 0].imag) for v in values] + return np.asarray(sampled)[..., None] + xg, yg = np.meshgrid(x_q, y_q, indexing="ij") + points = np.stack([xg.ravel(), yg.ravel()], axis=-1) + return np.asarray([RegularGridInterpolator((x_src, y_src), v)(points).reshape(xg.shape) for v in values]) + + +def _expected_save_axis(n_save: int, lo: float, hi: float) -> np.ndarray: + return lo + (np.arange(n_save) + 0.5) * (hi - lo) / n_save + + +@pytest.mark.parametrize("geometry", ["2d", "1d"]) +def test_interpolated_save_matches_full_grid_save(geometry, monkeypatch): + """Same run, saved on the full grid and with save dx = dy = 1.6 x the grid spacing. + + Every x-space field of the interpolated save (phi, ex, ey, e0_x/y, e1_x/y, + background_density, iaw_density, iaw_velocity_divergence) must equal the full-grid + field sampled at the save coordinates by linear interpolation. Both sides are the same + linear interpolation of the same data, so the tolerance is round-off: + max|interp - full(at save points)| <= ROUNDOFF_EPS * eps(saved dtype) * max|full(at save points)| + (2000 eps; derivation at ROUNDOFF_EPS). A field that is zero in the full-grid save must + be exactly zero. + + The save grid has n_save = round(L / 0.08 um) cell-centred points: 80 in x; 20 in y in + 2D, 1 in the quasi-1D box (L_y = one 0.05 um cell). Its coordinates and the k axes of + k-fields.xr are those of that grid, to 1e-13 (a few float64 operations each side, ~1e-15); + every saved value (x- and k-space) is finite. + + The full-grid reference run skips the plots (90% of its run time; full-grid plotting is + exercised elsewhere in the suite); the interpolated run plots, so post-processing is run + on the save grid. + """ + from adept._lpse2d import helpers + + with monkeypatch.context() as patch: + for plot in ("plot_series", "plot_srs_diagnostics", "plot_fields", "plot_kt"): + patch.setattr(helpers, plot, lambda *args, **kwargs: None) + full = _run(_box_cfg(geometry, None, f"field-save-full-{geometry}")) + interp = _run(_box_cfg(geometry, "0.08um", f"field-save-interp-{geometry}")) + + xs_full = full["x"].coords["x (um)"].values + ys_full = full["x"].coords["y (um)"].values + nx, ny = xs_full.size, ys_full.size + assert nx == 128 and ny == {"2d": 32, "1d": 1}[geometry] + + # the save grid: n_save cell-centred points spanning the simulation box + n_save_x, n_save_y = 80, {"2d": 20, "1d": 1}[geometry] + xs = interp["x"].coords["x (um)"].values + ys = interp["x"].coords["y (um)"].values + ly = ny * GRID_DX_UM + np.testing.assert_allclose(xs, _expected_save_axis(n_save_x, 0.0, X_LENGTH_UM), rtol=0, atol=1e-13 * X_LENGTH_UM) + np.testing.assert_allclose(ys, _expected_save_axis(n_save_y, -ly / 2, ly / 2), rtol=0, atol=1e-13 * ly) + + assert set(interp["x"].data_vars) == EXPECTED_FIELDS + assert set(full["x"].data_vars) == EXPECTED_FIELDS + nt = full["x"].coords["t (ps)"].size + for name in EXPECTED_FIELDS: + got = np.asarray(interp["x"][name].values) + assert got.shape == (nt, n_save_x, n_save_y), name + assert np.all(np.isfinite(got)), f"{name} has non-finite values" + + ref = _sample_linear(np.asarray(full["x"][name].values), xs_full, ys_full, xs, ys) + scale = np.max(np.abs(ref)) + eps = np.finfo(got.dtype).eps + err = np.max(np.abs(got - ref)) + assert err <= ROUNDOFF_EPS * eps * scale, f"{name}: max err {err:.3e} vs max |full| {scale:.3e}" + + # every derived k-space field lives on the save grid's k axes. The k coords are stored in + # units of w0/c; read that scale off the full-grid save, whose kx axis is the grid's. + kx_full_axis = np.fft.fftshift(np.fft.fftfreq(nx, d=GRID_DX_UM / 2.0 / np.pi)) + ratio = full["k"].coords[KX].values.max() / kx_full_axis.max() + np.testing.assert_allclose(full["k"].coords[KX].values, kx_full_axis * ratio, rtol=1e-13, atol=0) + dx_save = X_LENGTH_UM / n_save_x + dy_save = ly / n_save_y + kx_expected = np.fft.fftshift(np.fft.fftfreq(n_save_x, d=dx_save / 2.0 / np.pi)) * ratio + ky_expected = np.fft.fftshift(np.fft.fftfreq(n_save_y, d=dy_save / 2.0 / np.pi)) * ratio + np.testing.assert_allclose(interp["k"].coords[KX].values, kx_expected, rtol=1e-13, atol=0) + np.testing.assert_allclose(interp["k"].coords[KY].values, ky_expected, rtol=1e-13, atol=0) + for name in interp["k"].data_vars: + got = np.asarray(interp["k"][name].values) + assert got.shape == (nt, n_save_x, n_save_y), name + assert np.all(np.isfinite(got)), f"k-space {name} has non-finite values" + + +def _plane_wave(cfg, amplitude, mx, my): + """amplitude * exp(i (kx (x - x0) + ky (y - y0))) on the simulation grid, for grid mode (mx, my).""" + x, y = np.asarray(cfg["grid"]["x"]), np.asarray(cfg["grid"]["y"]) + kx, ky = cfg["grid"]["kx"][mx], cfg["grid"]["ky"][my] + return amplitude * np.exp(1j * (kx * (x[:, None] - x[0]) + ky * (y[None, :] - y[0]))), kx, ky + + +@pytest.mark.parametrize("geometry", ["2d", "1d"]) +def test_interpolated_save_of_band_limited_state(geometry, tmp_path): + """Analytic check of the save function and make_field_xarrays on a single-mode state, + save dx = dy = 1.6 x the grid spacing. + + The EPW state is phi_k with one grid mode (mx, my) = (3, 1) in 2D, (3, 0) in quasi-1D, + phi_k[mx, my] = A nx ny, i.e. phi(x) = A exp(i(kx (x - x0) + ky (y - y0))) with the + solver's convention phi(x) = ifft2(phi_k), E = ifft2(-i k phi_k). E0, E1 and the IAW + fields are other plane waves of the grid. + + Expected, on the save grid (x_s, y_s): + * x-space: phi(x_s), ex = -i kx phi(x_s), ey = -i ky phi(x_s) -- the solver's field, as in + the full-grid save -- and E0, E1, IAW at x_s. + * k-space at the mode's bin of the save grid: phi_k = A nx ny exp(i(kx (x_s0 - x0) + ky (y_s0 - y0))) + (the full-grid coefficient; the phase is the DFT origin moving to the first save point), + ex_k = -i kx phi_k, ey_k = -i ky phi_k. + Tolerance: bilinear interpolation of exp(i q.x) errs by at most (qx^2 hx^2 + qy^2 hy^2) / 8 per + real and imaginary part (tensor-product bound, h = grid spacing), so + |err| <= sqrt(2) (qx^2 + qy^2) h^2 / 8 * max|f|; the DFT of the interpolation error is bounded by + the same relative amount. Round-off allowance: ROUNDOFF_EPS * eps on top. + """ + from adept._lpse2d.helpers import make_field_xarrays + + save_spacing = "0.08um" + cfg = _prepared_cfg(_box_cfg(geometry, save_spacing, "unused")) + grid = cfg["grid"] + nx, ny = grid["nx"], grid["ny"] + my = 1 if geometry == "2d" else 0 + h = grid["dx"] + + amp = 2.0e-3 + phi_x, kx, ky = _plane_wave(cfg, amp, 3, my) + phi_k = np.fft.fft2(phi_x) + e0y, q0x, q0y = _plane_wave(cfg, 1.0e-2, 2, my) + e1x, q1x, q1y = _plane_wave(cfg, 3.0e-4, nx - 4, 0) + iaw, qix, qiy = _plane_wave(cfg, 1.0e-3, 1, my) + zeros = np.zeros_like(e0y) + state = { + "epw": phi_k.view(np.float64), + "E0": np.stack([zeros, e0y], axis=-1).view(np.float64), + "E1": np.stack([e1x, zeros], axis=-1).view(np.float64), + "iaw_density": iaw.real.copy(), + "iaw_velocity_divergence": iaw.imag.copy(), + } + import jax.numpy as jnp + + saved = cfg["save"]["fields"]["func"](0.0, {k: jnp.asarray(v) for k, v in state.items()}, None) + saved = {k: np.asarray(v)[None] for k, v in saved.items()} # one save time + (tmp_path / "binary").mkdir() + kfields, fields = make_field_xarrays(cfg, np.array([0.0]), saved, str(tmp_path)) + + n_save_x = round(X_LENGTH_UM / float(save_spacing.removesuffix("um"))) + n_save_y = 1 if geometry == "1d" else round(ny * h / float(save_spacing.removesuffix("um"))) + xs = fields.coords["x (um)"].values + ys = fields.coords["y (um)"].values + assert xs.shape == (n_save_x,) and ys.shape == (n_save_y,) + x0, y0 = grid["x"][0], grid["y"][0] + + def at_save(amplitude, qx, qy): + return amplitude * np.exp(1j * (qx * (xs[:, None] - x0) + qy * (ys[None, :] - y0))) + + def tol(qx, qy, scale, dtype): + interp_bound = np.sqrt(2.0) * (qx**2 + qy**2) * h**2 / 8.0 + return (interp_bound + ROUNDOFF_EPS * np.finfo(dtype).eps) * scale + + phi_s = at_save(amp, kx, ky) + expected_x = { + "phi": (phi_s, kx, ky), + "ex": (-1j * kx * phi_s, kx, ky), + "ey": (-1j * ky * phi_s, kx, ky), + "e0_x": (np.zeros_like(phi_s), 0.0, 0.0), + "e0_y": (at_save(1.0e-2, q0x, q0y), q0x, q0y), + "e1_x": (at_save(3.0e-4, q1x, q1y), q1x, q1y), + "e1_y": (np.zeros_like(phi_s), 0.0, 0.0), + "iaw_density": (at_save(1.0e-3, qix, qiy).real, qix, qiy), + "iaw_velocity_divergence": (at_save(1.0e-3, qix, qiy).imag, qix, qiy), + } + for name, (expected, qx, qy) in expected_x.items(): + got = np.asarray(fields[name].values)[0] + assert got.shape == (n_save_x, n_save_y), name + assert np.all(np.isfinite(got)), name + scale = np.max(np.abs(expected)) + err = np.max(np.abs(got - expected)) + assert err <= tol(qx, qy, scale, got.dtype), f"{name}: max err {err:.3e}, max |f| {scale:.3e}" + + # k-space: the mode sits in the same bin of the save grid (both grids span the same box) + phase = np.exp(1j * (kx * (xs[0] - x0) + ky * (ys[0] - y0))) + phi_k_expected = amp * nx * ny * phase + for name, factor in (("phi", 1.0), ("ex", -1j * kx), ("ey", -1j * ky)): + got_k = np.fft.ifftshift(np.asarray(kfields[name].values)[0]) + assert got_k.shape == (n_save_x, n_save_y), name + assert np.all(np.isfinite(got_k)), name + expected = factor * phi_k_expected + scale = np.abs(amp * nx * ny * factor) if factor != 0 else 0.0 + err = np.abs(got_k[3, my] - expected) + assert err <= tol(kx, ky, scale, got_k.dtype), f"k-space {name}: err {err:.3e} of {scale:.3e}" + + +@pytest.mark.parametrize("axis", ["x", "y"]) +def test_interpolated_save_spacing_bounds(axis): + """A save spacing below the simulation grid spacing (here half of it) is refused at setup; + a save spacing equal to it is accepted and reproduces the simulation grid. + + An interpolated save finer than the grid carries no information the grid does not have, + and its first/last points fall outside the outermost cell centres, where linear + interpolation is undefined (interpax returned NaN rows there). + + At equal spacing the save axis must be the simulation axis exactly (bitwise), so the + outermost save points sit on the outermost cell centres rather than a rounding error + outside them (the other axis keeps a 0.08 um save spacing). The save of a random E0 and + IAW density is then finite everywhere and equals the same linear interpolation done here + (the identity along the equal-spacing axis) to ROUNDOFF_EPS * eps * max|f|. + + Giving only one of save.fields.x / save.fields.y is refused too (y alone used to be + silently ignored). + """ + import jax.numpy as jnp + + key = {"x": "dx", "y": "dy"}[axis] + cfg = _box_cfg("2d", "0.08um", "unused") + cfg["save"]["fields"][axis][key] = f"{GRID_DX_UM / 2}um" + with pytest.raises(ValueError, match=rf"save\.fields\.{axis}\.{key}"): + _prepared_cfg(cfg) + + cfg = _box_cfg("2d", "0.08um", "unused") + del cfg["save"]["fields"][{"x": "y", "y": "x"}[axis]] + with pytest.raises(NotImplementedError, match=r"save\.fields\.x and save\.fields\.y"): + _prepared_cfg(cfg) + + cfg = _box_cfg("2d", "0.08um", "unused") + cfg["save"]["fields"][axis][key] = f"{GRID_DX_UM}um" + cfg = _prepared_cfg(cfg) + grid, fld = cfg["grid"], cfg["save"]["fields"] + np.testing.assert_array_equal(fld[axis]["ax"], grid[axis]) + + rng = np.random.default_rng(0) + nx, ny = grid["nx"], grid["ny"] + e0 = rng.standard_normal((nx, ny, 2)) + 1j * rng.standard_normal((nx, ny, 2)) + iaw = rng.standard_normal((nx, ny)) + state = {"E0": e0.view(np.float64), "iaw_density": iaw} + saved = fld["func"](0.0, {k: jnp.asarray(v) for k, v in state.items()}, None) + other = {"x": "y", "y": "x"}[axis] + for name, value in (("E0", e0), ("iaw_density", iaw)): + got = np.asarray(saved[name]) + if name == "E0": + got = got.view(np.complex128 if got.dtype == np.float64 else np.complex64) + assert np.all(np.isfinite(got)), name + # along the equal-spacing axis the save is the identity; along the other (0.08 um) axis + # it is a linear interpolation, so compare against that + ref = _sample_linear( + np.moveaxis(value, -1, 0) if name == "E0" else value[None], + grid["x"], + grid["y"], + np.asarray(fld["x"]["ax"]), + np.asarray(fld["y"]["ax"]), + ) + ref = np.moveaxis(ref, 0, -1) if name == "E0" else ref[0] + assert got.shape == ref.shape, name + err = np.max(np.abs(got - ref)) + assert err <= ROUNDOFF_EPS * np.finfo(got.dtype).eps * np.max(np.abs(ref)), f"{name}: {err:.3e}" + assert fld[other]["ax"].size == round({"x": X_LENGTH_UM, "y": ny * GRID_DX_UM}[other] / 0.08) diff --git a/tests/test_lpse2d/test_hpe.py b/tests/test_lpse2d/test_hpe.py index 2746a24c..c9b0d145 100644 --- a/tests/test_lpse2d/test_hpe.py +++ b/tests/test_lpse2d/test_hpe.py @@ -106,7 +106,7 @@ def test_free_streaming(): u = jnp.array(rng.normal(0.0, 100.0, 1000)) ex_env = jnp.zeros(hpe.nx_f, dtype=jnp.complex128) - x_new, u_new = hpe.push(x, u, ex_env, 0.0) + x_new, u_new, *_ = hpe.push(x, u, ex_env, 0.0) gamma = np.sqrt(1.0 + (np.array(u) / hpe.c) ** 2) x_expected = np.array(x) + hpe.dt * np.array(u) / gamma @@ -129,7 +129,7 @@ def test_2d_free_streaming_and_both_field_components(): u = jnp.asarray(rng.normal(0.0, 100.0, (hpe.n_p, 2))) zero_field = jnp.zeros((hpe.nx_f, hpe.ny_f, 2), dtype=jnp.complex128) - x_new, y_new, u_new = hpe.push_2d(x, y, u, zero_field, 0.0) + x_new, y_new, u_new, *_ = hpe.push_2d(x, y, u, zero_field, 0.0) gamma = np.sqrt(1.0 + np.sum((np.asarray(u) / hpe.c) ** 2, axis=-1)) x_expected = np.mod(np.asarray(x) + hpe.dt * np.asarray(u)[:, 0] / gamma - hpe.xmin, hpe.Lx) + hpe.xmin y_expected = np.mod(np.asarray(y) + hpe.dt * np.asarray(u)[:, 1] / gamma - hpe.ymin, hpe.Ly) + hpe.ymin @@ -196,7 +196,7 @@ def test_2d_wall_reinjection_is_flux_weighted(): y[2 * quarter : 3 * quarter] = hpe.ymin - hpe.dy y[3 * quarter :] = hpe.ymax + hpe.dy u = jnp.zeros((n_particles, 2)) - x_new, y_new, u_new = hpe._apply_boundaries_2d(jnp.asarray(x), jnp.asarray(y), u, 0.0) + x_new, y_new, u_new, *_ = hpe._apply_boundaries_2d(jnp.asarray(x), jnp.asarray(y), u, 0.0) gamma_rel = np.sqrt(1.0 + np.sum((np.asarray(u_new) / hpe.c) ** 2, axis=-1)) velocity = np.asarray(u_new) / gamma_rel[:, None] np.testing.assert_allclose(x_new[:quarter], hpe.xmin, rtol=0.0, atol=0.0) @@ -248,7 +248,7 @@ def run(v0): push = jax.jit(hpe.push) for i in range(n_steps): - x, u = push(x, u, ex_env, i * dt) + x, u, *_ = push(x, u, ex_env, i * dt) us[i] = np.array(u) return us @@ -432,7 +432,12 @@ def _run_exo(cfg): @pytest.mark.slow def test_srs_smoke_with_hpe(): """End-to-end: the SRS growth box with HPE active runs, stays finite, and emits - the hot-electron series/metrics.""" + the hot-electron series/metrics. + + The final hpe_gamma_ratio_kpeak is shot-noise limited: over EPW noise seeds 1-10 it spanned + 0.0-2.39 at 20k particles (4/10 above 0.5). The criterion (ratio[-1] > 0.5), 1M particles and the + pinned noise seed 7919 were fixed before validation; on fresh seeds 7919 and 11-16 the final ratio + was 0.59-1.33 (7/7 above 0.5).""" with open("tests/test_lpse2d/configs/srs.yaml") as fi: cfg = yaml.safe_load(fi) cfg["grid"]["ymax"] = "0.02um" @@ -440,7 +445,8 @@ def test_srs_smoke_with_hpe(): cfg["grid"]["tmax"] = "0.1ps" cfg["save"]["fields"]["t"]["tmax"] = "0.1ps" cfg["save"]["fields"]["t"]["dt"] = "0.05ps" - cfg["terms"]["hpe"] = {"active": True, "n_particles": 20000, "nv": 256} + cfg["terms"]["hpe"] = {"active": True, "n_particles": 1000000, "nv": 256} + cfg["terms"]["epw"]["source"]["noise_seed"] = 7919 cfg["mlflow"]["run"] = "srs-hpe-smoke" exo, sol, ppo = _run_exo(cfg) @@ -449,10 +455,15 @@ def test_srs_smoke_with_hpe(): assert "fhot_50keV" in series assert "hpe_gamma_ratio_kpeak" in series assert "hpe_hist" in series + # hpe_gamma_ratio_kpeak is NaN by design while the EPW has no energy in the resonant band + # (1989b5e: the reduction metrics nanmin over it); the EPW starts at zero, so only t = 0 for k in series: - assert np.all(np.isfinite(series[k].values)), f"non-finite values in {k}" + values = np.asarray(series[k].values) + if k == "hpe_gamma_ratio_kpeak": + values = values[1:] + assert np.all(np.isfinite(values)), f"non-finite values in {k}" # nothing exciting happens in 0.1 ps: damping at the dominant band mode should - # still be ~analytic (the band-min is shot-noise-limited at 20k particles) + # still be ~analytic (see the docstring for its shot noise) ratio = np.asarray(series["hpe_gamma_ratio_kpeak"].values, dtype=float) assert ratio[-1] > 0.5 @@ -505,12 +516,14 @@ def test_linear_closure(): t = np.array(result.ts["default"]) e_sq = np.array(result.ys["default"]["e_sq"]) - from adept._lpse2d.core.epw import landau_damping_rate + from adept._lpse2d.core.epw import analytic_landau_rate - d = exo.adept_module.cfg["units"]["derived"] - kx = np.array(exo.adept_module.cfg["grid"]["kx"]) + # the configured analytic rate (terms.epw.damping.landau_form and its threshold / multiplier): + # the rate HPE is calibrated to and the solver applies without HPE + run = exo.adept_module.cfg + kx = np.array(run["grid"]["kx"]) m = int(np.argmin(np.abs(kx - k0))) - gamma_k = float(np.array(landau_damping_rate(jnp.array(kx[m] ** 2), d["wp0"], d["vte_sq"], jnp.array(1.0)))) + gamma_k = float(np.asarray(analytic_landau_rate(run))[m, 0]) # driver is fully off past 0.31 ps; fit ~5 e-folds of clean decay fit = (t > 0.315) & (t < 0.40) @@ -551,13 +564,10 @@ def test_oneil_flattening(): e_amp = w_b**2 / (e_over_m * k) phi_k, k = _single_mode_phi_k(cfg, k_target, e_amp) - y = { - "epw": jnp.array(phi_k), - "x_e": jnp.array(state["x_e"]), - "u_e": jnp.array(state["u_e"]), - "epw_hist": jnp.array(state["epw_hist"]), - "gamma_L": jnp.array(state["gamma_L"]), - } + # the complete loaded state (the push also carries the wall-flux, cone-energy and multiplier + # instruments), with the driven wave + y = {k: jnp.asarray(v) for k, v in state.items()} + y["epw"] = jnp.array(phi_k) step = jax.jit(hpe.__call__) dt = cfg["grid"]["dt"] @@ -576,3 +586,296 @@ def test_oneil_flattening(): assert gamma_final < 0.6 * gamma_analytic, ( f"no damping reduction: gamma {gamma_final:.3g} vs analytic {gamma_analytic:.3g}" ) + + +# ---------------------------------------------------------------------------------- +# LPSE HPE controls and instruments (plan item D.1) + + +def test_gamma_limits_and_allow_growth(): + """LPSE hpe.gammaLimit: inside the resonant band the applied rate is clipped at + [-growth (if allowed) or 0, damping]; outside it the analytic rate is kept.""" + from adept._lpse2d.core.hpe import HybridParticleEvolution + + cfg = _make_cfg({"gamma_limit_damping": 0.5}) + hpe = HybridParticleEvolution(cfg) + band = np.asarray(hpe.mask_res) + # an inverted (rising) tail on the +v side gives negative raw rates for +k modes + v = np.asarray(hpe.v_centers) + rising = np.exp(np.clip(v / hpe.vte, -50, 50) * 0.5) * (np.abs(v) > hpe.v_min) + rising = rising / (rising.sum() * hpe.dv) + gamma = np.asarray(hpe.damping(jnp.asarray(rising)))[:, 0] + assert gamma[band].max() <= 0.5 + 1e-12 + assert gamma[band].min() >= 0.0 # growth not allowed by default + cfg2 = _make_cfg({"gamma_limit_damping": 0.5, "allow_growth": True, "gamma_limit_growth": 0.25}) + hpe2 = HybridParticleEvolution(cfg2) + gamma2 = np.asarray(hpe2.damping(jnp.asarray(rising)))[:, 0] + assert gamma2[band].min() < 0.0 and gamma2[band].min() >= -0.25 - 1e-12 + + +def test_dissipation_diagnostic_is_signed_with_allow_growth(): + """With terms.hpe.allow_growth the applied rate gamma_L + nu_coll can be negative (energy from the + electrons to the EPW). The default save's epw_dissipation is then the signed sum + prefactor * factor * sum_k L_k k^2 |phi_k|^2 / (nx ny^2), L_k = (1 - exp(-2 rate dt)) / dt (Parseval; + uniform box, so the total-to-electric factor 2 (2 - n/n_env) is a constant) -- before the fix + sqrt(L < 0) made it NaN (PR #387 review). With non-negative rates the series is bitwise the same + with and without allow_growth. rtol 1e-12, fixed in advance.""" + from adept._lpse2d.core.hpe import load_particles + from adept._lpse2d.helpers import get_default_save_func, get_save_quantities + + def dissipation(allow_growth, signed_rates): + cfg = get_save_quantities(_make_cfg({"allow_growth": allow_growth})) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + rng = np.random.default_rng(3) + phi_k = 1e-3 * (rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) + phi_k = phi_k * np.asarray(cfg["grid"]["low_pass_filter_grid"]) + state = {k: jnp.asarray(v) for k, v in load_particles(cfg).items()} + rate = rng.normal(size=state["gamma_L"].shape) * 50.0 + gamma = rate if signed_rates else np.abs(rate) + zeros = jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64) + state.update( + {"gamma_L": jnp.asarray(gamma), "epw": jnp.asarray(phi_k).view(jnp.float64), "E0": zeros, "E1": zeros} + ) + out = get_default_save_func(cfg)["func"](0.0, state, None) + return cfg, phi_k, gamma, float(out["epw_dissipation"]) + + cfg, phi_k, gamma, got = dissipation(True, True) + assert np.isfinite(got) + d, grid = cfg["units"]["derived"], cfg["grid"] + nx, ny, dt = grid["nx"], grid["ny"], grid["dt"] + k_sq = np.asarray(grid["kx"])[:, None] ** 2 + np.asarray(grid["ky"])[None, :] ** 2 + rate = gamma + d.get("nu_coll", 0.0) * (k_sq > 0) + assert np.any(rate * (np.abs(phi_k) > 0) < 0.0) + loss = (1.0 - np.exp(-2.0 * rate * dt)) / dt + factor = 2.0 * (2.0 - np.asarray(grid["background_density"]) / cfg["units"]["envelope density"]) + assert np.ptp(factor) == 0.0 + prefactor = 0.25 * grid["dx"] * d["x_norm"] * d["e_norm"] ** 2 + want = prefactor * float(factor.flat[0]) * np.sum(loss * k_sq * np.abs(phi_k) ** 2) / (nx * ny**2) + assert got == pytest.approx(want, rel=1e-12) + assert np.sum(loss * k_sq * np.abs(phi_k) ** 2 * (loss < 0)) < 0.0 # the gain modes contribute + assert dissipation(True, False)[3] == dissipation(False, False)[3] + + +def test_thermalization_probability_zero_passes_through_periodically(): + """LPSE particle walls: a crossing is counted, then the particle is thermalized with the + per-direction probability or else wraps around periodically.""" + from adept._lpse2d.core.hpe import HybridParticleEvolution + + for p_therm, expect_wrap in ((0.0, True), (1.0, False)): + cfg = _make_cfg_2d( + {"n_particles": 4000, "thermalization_probability": [p_therm, p_therm]}, + {"terms.epw.boundary.x": "absorbing", "terms.epw.boundary.y": "absorbing"}, + ) + hpe = HybridParticleEvolution(cfg) + n = hpe.n_p + x = jnp.full((n,), hpe.xmax + 0.1 * hpe.dx) # all just past the right wall + y = jnp.full((n,), 0.5 * (hpe.ymin + hpe.ymax)) + u = jnp.stack((jnp.full((n,), 3.0 * hpe.vte), jnp.full((n,), 1.0 * hpe.vte)), axis=-1) + x2, y2, u2, flux, cone = hpe._apply_boundaries_2d(x, y, u, 0.0) + if expect_wrap: + np.testing.assert_allclose(np.asarray(x2), hpe.xmin + 0.1 * hpe.dx, rtol=1e-9) + np.testing.assert_allclose(np.asarray(u2), np.asarray(u)) + else: + assert float(jnp.max(x2)) <= hpe.xmax + assert np.all(np.asarray(u2[:, 0]) < 0.0) # re-injected inward with resampled speeds + assert np.std(np.asarray(u2[:, 1])) > 0.1 * hpe.vte + # the wall-flux instrument saw every particle leave through the right wall either way + flux = np.asarray(flux) + assert flux[1].sum() > 0.0 and flux[0].sum() == 0.0 and flux[2].sum() == 0.0 and flux[3].sum() == 0.0 + # periodic field boundaries default to pass-through, absorbing ones to thermalization + hpe_p = HybridParticleEvolution(_make_cfg_2d({"n_particles": 100})) + assert hpe_p.p_therm_x == 0.0 and hpe_p.p_therm_y == 0.0 + hpe_a = HybridParticleEvolution(_make_cfg_2d({"n_particles": 100}, {"terms.epw.boundary.x": "absorbing"})) + assert hpe_a.p_therm_x == 1.0 and hpe_a.p_therm_y == 0.0 + + +def test_wall_flux_bins_and_cone_energy(): + from adept._lpse2d.core.hpe import HybridParticleEvolution + + cfg = _make_cfg_2d( + {"n_particles": 1000, "flux_bins": [0.0, 20.0, 1.0e9], "cone_angle": 30.0, "cone_direction": [1.0, 0.0]}, + {"terms.epw.boundary.x": "absorbing", "terms.epw.boundary.y": "absorbing"}, + ) + hpe = HybridParticleEvolution(cfg) + n = hpe.n_p + c = hpe.c + # half the particles at 10 keV leaving right along +x (inside the cone), half at 40 keV leaving + # through the top wall along +y (outside the cone) + gamma_10, gamma_40 = 1.0 + 10.0 / 510.999, 1.0 + 40.0 / 510.999 + u10, u40 = c * np.sqrt(gamma_10**2 - 1.0), c * np.sqrt(gamma_40**2 - 1.0) + half = n // 2 + x_mid, y_mid = 0.5 * (hpe.xmin + hpe.xmax), 0.5 * (hpe.ymin + hpe.ymax) + x = jnp.concatenate((jnp.full((half,), hpe.xmax + 0.1 * hpe.dx), jnp.full((n - half,), x_mid))) + y = jnp.concatenate((jnp.full((half,), y_mid), jnp.full((n - half,), hpe.ymax + 0.1 * hpe.dy))) + ux = jnp.concatenate((jnp.full((half,), u10), jnp.zeros((n - half,)))) + uy = jnp.concatenate((jnp.zeros((half,)), jnp.full((n - half,), u40))) + _, _, _, flux, cone = hpe._apply_boundaries_2d(x, y, jnp.stack((ux, uy), axis=-1), 0.0) + flux = np.asarray(flux) + np.testing.assert_allclose(flux[1, 0], 10.0 * half, rtol=1e-6) # right wall, bin [0, 20) keV + np.testing.assert_allclose(flux[3, 1], 40.0 * (n - half), rtol=1e-6) # top wall, bin [20, inf) keV + assert flux[1, 1] == 0.0 and flux[3, 0] == 0.0 + np.testing.assert_allclose(float(cone[0]), 10.0 * half, rtol=1e-6) + + +def test_wall_flux_window_excludes_energies_outside_its_edges(): + """A finite window counts only the electrons inside it (PR #372 review): of 10 / 70 / 200 keV + electrons leaving the right wall, flux_bins [50, 100] records the 70 keV ones alone. Before the + fix the clipped bin index folded the 10 and 200 keV electrons into the single bin.""" + from adept._lpse2d.core.hpe import HybridParticleEvolution + + cfg = _make_cfg_2d( + {"n_particles": 999, "flux_bins": [50.0, 100.0]}, + {"terms.epw.boundary.x": "absorbing", "terms.epw.boundary.y": "absorbing"}, + ) + hpe = HybridParticleEvolution(cfg) + n, c = hpe.n_p, hpe.c + third = n // 3 + energies = np.concatenate((np.full(third, 10.0), np.full(third, 70.0), np.full(n - 2 * third, 200.0))) + gamma = 1.0 + energies / 510.999 + ux = jnp.asarray(c * np.sqrt(gamma**2 - 1.0)) + x = jnp.full((n,), hpe.xmax + 0.1 * hpe.dx) + y = jnp.full((n,), 0.5 * (hpe.ymin + hpe.ymax)) + _, _, _, flux, _ = hpe._apply_boundaries_2d(x, y, jnp.stack((ux, jnp.zeros((n,))), axis=-1), 0.0) + flux = np.asarray(flux) + assert flux.shape == (4, 1) + np.testing.assert_allclose(flux[1, 0], 70.0 * third, rtol=1e-6) + assert flux[0].sum() == 0.0 and flux[2].sum() == 0.0 and flux[3].sum() == 0.0 + + +def test_magnetic_field_rotates_momentum_without_changing_energy(): + from adept._lpse2d.core.hpe import HybridParticleEvolution + + b_tesla = 2000.0 + cfg = _make_cfg_2d({"n_particles": 100, "magnetic_field": b_tesla}) + hpe = HybridParticleEvolution(cfg) + n = hpe.n_p + x = jnp.full((n,), 0.5 * (hpe.xmin + hpe.xmax)) + y = jnp.full((n,), 0.5 * (hpe.ymin + hpe.ymax)) + u = jnp.stack((jnp.full((n,), 2.0 * hpe.vte), jnp.zeros((n,))), axis=-1) + zero_field = jnp.zeros((hpe.nx_f, hpe.ny_f, 2), dtype=jnp.complex128) + _, _, u2, *_ = hpe.push_2d(x, y, u, zero_field, 0.0) + u2 = np.asarray(u2) + np.testing.assert_allclose(np.hypot(u2[:, 0], u2[:, 1]), 2.0 * hpe.vte, rtol=1e-9) + gamma_rel = np.sqrt(1.0 + (2.0 * hpe.vte / hpe.c) ** 2) + theta = 0.175882 * b_tesla * hpe.dt / gamma_rel # cyclotron angle over one EPW step + np.testing.assert_allclose(np.arctan2(u2[:, 1], u2[:, 0]), theta, rtol=1e-6) + + +def test_energy_conservation_multiplier_tracks_the_particle_gain(): + from adept._lpse2d.core.hpe import HybridParticleEvolution + + cfg = _make_cfg({"energy_conservation": True, "energy_conservation_steps": 1.0}) + hpe = HybridParticleEvolution(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + phi_k = jnp.zeros((nx, ny), dtype=jnp.complex128).at[3, 0].set(1.0e-3) + gamma_l = jnp.full((nx, ny), 0.1) + loss = float(hpe.wave_energy_scale * jnp.sum(2.0 * gamma_l * hpe.dt * hpe.k_sq * jnp.abs(phi_k) ** 2)) + d_kinetic = 2.0 * loss / hpe.particle_energy_scale # particles gained twice the expected wave loss + m = hpe.ld_multiplier(jnp.ones((1,)), gamma_l, phi_k, jnp.asarray(d_kinetic)) + np.testing.assert_allclose(float(m[0]), 2.0, rtol=1e-6) + m_clamped = hpe.ld_multiplier(jnp.ones((1,)), gamma_l, phi_k, jnp.asarray(100.0 * d_kinetic)) + np.testing.assert_allclose(float(m_clamped[0]), 10.0, rtol=1e-6) + # with a running average over 4 steps the multiplier moves a quarter of the way + cfg4 = _make_cfg({"energy_conservation": True, "energy_conservation_steps": 4.0}) + m4 = HybridParticleEvolution(cfg4).ld_multiplier(jnp.ones((1,)), gamma_l, phi_k, jnp.asarray(d_kinetic)) + np.testing.assert_allclose(float(m4[0]), 1.25, rtol=1e-6) + + +def test_gamma_limits_bound_the_rate_after_the_energy_conservation_multiplier(): + """Through __call__ with gamma limits and energy conservation both on (PR #372 review): the rate + handed to the solver is the multiplier times the HPE rate, clipped again to the configured limits + on the resonant modes; the non-resonant modes keep the multiplier times their analytic rate (no + clip). LPSE clips only before its LD_multiplier (ElectronTracker.cu:2792, LwSolver.cpp:823); adept + departs from it here. The multiplier is held near 5 (a long running average from 5) so the limit + binds. rtol 1e-12.""" + from adept._lpse2d.core.hpe import HybridParticleEvolution, load_particles + + limit = 0.5 + cfg = _make_cfg({"gamma_limit_damping": limit, "energy_conservation": True, "energy_conservation_steps": 1.0e6}) + hpe = HybridParticleEvolution(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + rng = np.random.default_rng(29) + phi_k = 1e-3 * (rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) + phi_k = phi_k * np.asarray(cfg["grid"]["low_pass_filter_grid"]) + zeros = jnp.zeros((nx, ny, 2), dtype=jnp.complex128).view(jnp.float64) + y = {k: jnp.asarray(v) for k, v in load_particles(cfg).items()} + y.update({"epw": jnp.asarray(phi_k), "E0": zeros, "E1": zeros, "hpe_ld_multiplier": jnp.full((1,), 5.0)}) + out = hpe(max(hpe.t_start, 0.0) + hpe.dt, y) + rate = np.asarray(hpe.damping(out["epw_hist"])) + mult = float(np.asarray(out["hpe_ld_multiplier"])[0]) + assert mult == pytest.approx(5.0, rel=1e-4) + resonant = np.broadcast_to(np.asarray(hpe.mask_res).reshape(nx, -1), (nx, ny)) + want = np.where(resonant, np.clip(mult * rate, -hpe.gamma_limit_growth, limit), mult * rate) + gamma_l = np.asarray(out["gamma_L"]) + np.testing.assert_allclose(gamma_l, want, rtol=1e-12, atol=0.0) + assert gamma_l[resonant].max() <= limit + 1e-12 + assert np.any(mult * rate[resonant] > limit) # the limit binds after the multiplier + + +def test_translator_maps_the_hpe_controls(): + from adept._lpse2d.lpse_deck import translate_parms + + parms = { + "grid.sizes": "20 5", + "grid.nodes": "201 51", + "laser.enable": "true", + "lw.enable": "true", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.1.intensity": "1e15", + "hpe.enable": "true", + "hpe.nElectrons": "12345", + "hpe.velocityGrid": "128", + "hpe.startEvolutionAt": "0.5", + "hpe.thermalizationProbability": "1 0.1 0.1", + "hpe.magneticField": "0 0 50", + "hpe.gammaLimit.damping": "300", + "hpe.gammaLimit.growth": "20", + "hpe.allowGrowth": "true", + "hpe.enforceEnergyConservation": "true", + "hpe.numStepsToAverageEnergyChange": "5", + "hpe.metrics.nFluxMetrics": "2", + "hpe.metrics.fluxMetric.1.energy.min": "0", + "hpe.metrics.fluxMetric.1.energy.max": "50", + "hpe.metrics.fluxMetric.2.energy.min": "50", + "hpe.metrics.fluxMetric.2.energy.max": "1e9", + "hpe.metrics.nPowerMetrics": "1", + "hpe.metrics.powerMetric.1.angle": "20", + "hpe.metrics.powerMetric.1.direction": "0 1 0", + } + cfg, report = translate_parms(parms, experiment="x", run="y") + h = cfg["terms"]["hpe"] + assert h["active"] and h["n_particles"] == 12345 and h["nv"] == 128 and h["t_start"] == "0.5ps" + assert h["thermalization_probability"] == [1.0, 0.1] and h["magnetic_field"] == 50.0 + assert h["gamma_limit_damping"] == 300.0 and h["gamma_limit_growth"] == 20.0 and h["allow_growth"] + assert h["energy_conservation"] and h["energy_conservation_steps"] == 5.0 + assert h["flux_bins"] == [0.0, 50.0, 1e9] and h["cone_angle"] == 20.0 and h["cone_direction"] == [0.0, 1.0] + + +def test_quasi1d_periodic_crossings_reach_the_wall_accounting(): + """PR #372 / #378 review: the quasi-1D push keeps the unwrapped position until + _apply_boundaries, so a periodic crossing is counted by the wall instrument and the + configured thermalization probability applies (as in 2-D). Exercised through push().""" + from adept._lpse2d.core.hpe import HybridParticleEvolution + + for p_therm in (0.0, 1.0): + cfg = _make_cfg({"n_particles": 1000, "thermalization_probability": [p_therm, p_therm]}) + hpe = HybridParticleEvolution(cfg) + assert hpe.periodic_x + n = hpe.n_p + gamma_10 = 1.0 + 10.0 / 510.999 + u10 = hpe.c * np.sqrt(gamma_10**2 - 1.0) + v10 = u10 / gamma_10 + # 10 keV electrons half a step from the right wall: every one crosses during this push + x = jnp.full((n,), hpe.xmax - 0.5 * hpe.dt * v10) + u = jnp.full((n,), u10) + x2, u2, flux, cone, _ = hpe.push(x, u, jnp.zeros(hpe.nx_f, dtype=jnp.complex128), 0.0) + flux = np.asarray(flux) + np.testing.assert_allclose(flux[1].sum(), 10.0 * n, rtol=1e-6) # right wall + assert flux[0].sum() == 0.0 + if p_therm == 0.0: + np.testing.assert_allclose(np.asarray(x2), hpe.xmin + 0.5 * hpe.dt * v10, rtol=0, atol=1e-9) + np.testing.assert_allclose(np.asarray(u2), np.asarray(u), rtol=0, atol=0) + else: + np.testing.assert_allclose(np.asarray(x2), hpe.xmax, rtol=0, atol=0) + assert np.all(np.asarray(u2) < 0.0) # re-injected inward diff --git a/tests/test_lpse2d/test_iaw_spectral.py b/tests/test_lpse2d/test_iaw_spectral.py new file mode 100644 index 00000000..01d08726 --- /dev/null +++ b/tests/test_lpse2d/test_iaw_spectral.py @@ -0,0 +1,271 @@ +"""LPSE-parity tests for the ion-acoustic solver: the spectral (exact damped-oscillator) +propagator, uniform background flow, the full Krall-Trivelpiece ion Landau rate, the +fluctuation-dissipation IAW noise source and the IAW stride.""" + +import numpy as np +import pytest +import yaml +from jax import numpy as jnp + + +def _make_cfg( + *, solver="spectral", landau=0.0, collisions=0.0, dt="2fs", stride=1, flow=None, grid_extra=None, **iaw_extra +): + from adept._lpse2d.helpers import get_density_profile, get_derived_quantities, get_solver_quantities, write_units + + with open("tests/test_lpse2d/configs/epw.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg["density"] = {"basis": "uniform", "val": 0.25} + cfg["grid"].update( + { + "boundary_width": "0.2um", + "dt": dt, + "dx": "0.1um", + "xmax": "6.4um", + "tmax": "20fs", + "ymax": "0.05um", + "ymin": "-0.05um", + "light_substeps": 1, + "low_pass_filter": 1.0, + } + ) + cfg["grid"].update(grid_extra or {}) + cfg["terms"]["zero_mask"] = True + cfg["terms"]["epw"]["boundary"] = {"x": "periodic", "y": "periodic"} + cfg["terms"]["epw"]["source"] = {"noise": False, "tpd": False, "srs": False} + cfg["terms"]["epw"]["density_gradient"] = False + cfg["terms"]["iaw"] = { + "active": True, + "solver": solver, + "boundary": {"x": "periodic", "y": "periodic"}, + "damping": {"collisions": collisions, "landau": landau}, + "stride": stride, + "flow": flow, + **iaw_extra, + } + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _zero_fields(cfg): + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + return { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128), + "E0": jnp.zeros((nx, ny, 2), dtype=jnp.complex128), + "E1": jnp.zeros((nx, ny, 2), dtype=jnp.complex128), + } + + +def _single_mode_state(cfg, ik, amplitude=0.01): + x = np.asarray(cfg["grid"]["x"]) + k = float(np.asarray(cfg["grid"]["kx"])[ik]) + density = jnp.asarray(amplitude * np.cos(k * (x - x[0])))[:, None] + return k, {**_zero_fields(cfg), "iaw_density": density, "iaw_velocity_divergence": jnp.zeros_like(density)} + + +def _run(solver, state, n_steps): + for i in range(n_steps): + state = solver(state, i * solver.dt) + return state + + +def test_spectral_propagator_matches_analytic_damped_oscillator(): + """n = cos(kx), w = 0: the exact solution is n(t) = e^{-g t}[cos(b t) + g/b sin(b t)] n0, + w = -dn/dt, with g = landau cs k and b = sqrt(cs^2 k^2 - g^2). The propagator reproduces + it to round-off after many steps with cs k dt ~ 1 (far beyond the explicit limit).""" + from adept._lpse2d.core.iaw import IonAcousticWave + + cfg = _make_cfg(landau=0.2, dt="1.5ps") + solver = IonAcousticWave(cfg) + k, state = _single_mode_state(cfg, 4) + cs = solver.cs + g = 0.2 * cs * k + b = np.sqrt((cs * k) ** 2 - g**2) + assert cs * k * solver.dt > 0.5 + n_steps = 25 + out = _run(solver, state, n_steps) + t = n_steps * solver.dt + x = np.asarray(cfg["grid"]["x"]) + n_exact = 0.01 * np.exp(-g * t) * (np.cos(b * t) + g / b * np.sin(b * t)) * np.cos(k * (x - x[0])) + dn_dt = ( + 0.01 + * np.exp(-g * t) + * (-g * (np.cos(b * t) + g / b * np.sin(b * t)) + (-b * np.sin(b * t) + g * np.cos(b * t))) + * np.cos(k * (x - x[0])) + ) + np.testing.assert_allclose(np.asarray(out["iaw_density"])[:, 0], n_exact, rtol=1e-10, atol=1e-14) + np.testing.assert_allclose(np.asarray(out["iaw_velocity_divergence"])[:, 0], -dn_dt, rtol=1e-10, atol=1e-12) + + +def test_spectral_propagator_is_stable_at_any_dt(): + from adept._lpse2d.core.iaw import IonAcousticWave + + cfg = _make_cfg(landau=0.05, dt="10ps") + solver = IonAcousticWave(cfg) + k_max = float(np.abs(np.asarray(cfg["grid"]["kx"])).max()) + assert solver.cs * k_max * solver.dt > 20.0 # omega_max dt far beyond the explicit limit of 2 + _, state = _single_mode_state(cfg, 20) + out = _run(solver, state, 200) + n = np.asarray(out["iaw_density"]) + assert np.all(np.isfinite(n)) and np.abs(n).max() <= 0.01 * (1.0 + 1e-9) + + +def test_explicit_solver_refuses_unstable_dt_but_spectral_accepts_it(): + with pytest.raises(ValueError, match="unstable"): + _make_cfg(solver="explicit", dt="10ps") + cfg = _make_cfg(solver="spectral", dt="10ps") + assert cfg["terms"]["iaw"]["solver"] == "spectral" + + +def test_explicit_and_spectral_agree_for_small_dt_and_long_wavelength(): + """Undamped, undriven: the explicit kick/drift step converges to the exact propagator as + dt -> 0; the FD Laplacian differs from k^2 by (k dx)^2/12 (0.1% at k = 2 pi / 6.4 um).""" + from adept._lpse2d.core.iaw import IonAcousticWave + + cfg_e = _make_cfg(solver="explicit", dt="20fs") + cfg_s = _make_cfg(solver="spectral", dt="20fs") + explicit, spectral = IonAcousticWave(cfg_e), IonAcousticWave(cfg_s) + k, state = _single_mode_state(cfg_e, 1) + n_steps = 400 + out_e = _run(explicit, state, n_steps) + out_s = _run(spectral, state, n_steps) + assert spectral.cs * k * n_steps * spectral.dt > 1.0 # a good fraction of a period + np.testing.assert_allclose( + np.asarray(out_e["iaw_density"]), np.asarray(out_s["iaw_density"]), rtol=0, atol=3e-3 * 0.01 + ) + + +def test_flow_doppler_shifts_the_mode_phase(): + """With a uniform flow V0 along x, the exact propagator multiplies every mode by + e^{-i k V0 dt}: the +k and -k components of cos(kx) pick up opposite phases.""" + from adept._lpse2d.core.iaw import IonAcousticWave + + mach = 0.3 + cfg0 = _make_cfg() + cfgv = _make_cfg(flow=[mach, 0.0]) + still, moving = IonAcousticWave(cfg0), IonAcousticWave(cfgv) + ik = 3 + k, state = _single_mode_state(cfg0, ik) + n0 = np.fft.fft(np.asarray(still(state, 0.0)["iaw_density"])[:, 0]) + nv = np.fft.fft(np.asarray(moving(state, 0.0)["iaw_density"])[:, 0]) + v0 = mach * still.cs + expected_phase = np.exp(-1j * k * v0 * still.dt) + assert nv[ik] / n0[ik] == pytest.approx(expected_phase, rel=1e-10) + assert nv[-ik] / n0[-ik] == pytest.approx(np.conj(expected_phase), rel=1e-10) + with pytest.raises(ValueError, match="flow"): + _make_cfg(solver="explicit", flow=[0.1, 0.0]) + + +def test_full_ion_landau_rate_matches_krall_trivelpiece_formula(): + from astropy.units import Quantity as _Q + + from adept._lpse2d.core.iaw import ion_landau_rate + + cfg = _make_cfg(landau=0.1) + derived = cfg["units"]["derived"] + kx = np.asarray(cfg["grid"]["kx"]) + k_sq = kx[:, None] ** 2 + np.asarray(cfg["grid"]["ky"])[None, :] ** 2 + simplified = ion_landau_rate(cfg, k_sq) + np.testing.assert_allclose(simplified, 0.1 * derived["cs"] * np.sqrt(k_sq)) + + cfg["terms"]["iaw"]["damping"]["landau_form"] = "full" + full = ion_landau_rate(cfg, k_sq) + te = _Q(cfg["units"]["reference electron temperature"]).to("keV").value + ti = _Q(cfg["units"]["reference ion temperature"]).to("keV").value + z = cfg["units"]["ionization state"] + eta = 1.0 + 3.0 * ti / (z * te) + big_m = derived["mi"] / derived["me"] / (eta * z) + big_l = 1.0 + k_sq * derived["vte_sq"] / derived["wp0"] ** 2 + w_r = derived["cs"] * np.sqrt(k_sq) / np.sqrt(big_l) + a = 3.0 / (eta - 1.0) + w_i = w_r * np.sqrt(np.pi / 8.0) / big_l**1.5 * (a**1.5 * np.exp(-0.5 * a / big_l) + np.sqrt(1.0 / (eta * big_m))) + np.testing.assert_allclose(full, 0.5 * w_i, rtol=1e-12) + # Z = 6, Te = 2 Ti: eta = 1.25, a = 12 -> ion Landau dominated, gamma/(cs k) of order 0.1 + band = k_sq > 0 + ratio = full[band] / (derived["cs"] * np.sqrt(k_sq[band])) + assert np.all((ratio > 0.01) & (ratio < 1.0)) + + +def test_iaw_noise_kick_has_the_lpse_amplitude(): + from adept._lpse2d.core.iaw import IonAcousticWave + + amplitude = 3.0e-7 + cfg = _make_cfg(landau=0.1, collisions=0.02, noise=True, noise_amplitude=amplitude, noise_seed=9) + solver = IonAcousticWave(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + state = {**_zero_fields(cfg), "iaw_density": jnp.zeros((nx, ny)), "iaw_velocity_divergence": jnp.zeros((nx, ny))} + out = solver(state, 0.0) + w_k = np.fft.fft2(np.asarray(out["iaw_velocity_divergence"])) + gamma = np.asarray(solver.landau_rate) + 0.02 + expected = amplitude * nx * ny * np.sqrt(np.expm1(2.0 * solver.dt * gamma)) + band = np.asarray(solver.filter) > 0 + # the noise is complex in k-space; taking the real part in x-space keeps the + # Hermitian half, i.e. the +k/-k pair of each real mode (LPSE's own comment) + assert np.abs(w_k[band]).max() > 0.0 + assert np.abs(w_k[band]).max() <= expected[band].max() * (1.0 + 1e-9) + assert np.median(np.abs(w_k[band]) / expected[band]) == pytest.approx(0.7, abs=0.15) + with pytest.raises(ValueError, match="noise"): + IonAcousticWave(_make_cfg(landau=0.0, collisions=0.0, noise=True)) + + +def test_iaw_stride_advances_every_nth_epw_step(): + from adept._lpse2d.core.vector_field import SplitStep + + cfg = _make_cfg(stride=3, dt="0.5ps") + step = SplitStep(cfg) + assert step.iaw.stride == 3 and step.iaw.dt == pytest.approx(3.0 * cfg["grid"]["dt"]) + k, state = _single_mode_state(cfg, 4) + packed = {key: value.view(jnp.float64) for key, value in state.items()} + dt = cfg["grid"]["dt"] + same = step(jnp.asarray(dt), dict(packed), {"drivers": {}}) + np.testing.assert_array_equal(np.asarray(same["iaw_density"]), np.asarray(packed["iaw_density"])) + advanced = step(jnp.asarray(3.0 * dt), dict(packed), {"drivers": {}}) + direct = step.iaw(state, 3.0 * dt) + np.testing.assert_allclose(np.asarray(advanced["iaw_density"]), np.asarray(direct["iaw_density"]), rtol=1e-12) + change = np.abs(np.asarray(advanced["iaw_density"]) - np.asarray(packed["iaw_density"])).max() + assert change > 1e-4 # cs k dt_iaw ~ 1: the mode has rotated appreciably + with pytest.raises(ValueError, match="stride"): + _make_cfg(solver="explicit", stride=2) + + +def test_iaw_stride_applies_the_absorber_over_the_whole_iaw_step(): + """The IAW absorber removes exp(-rate grid.dt) per EPW step; an update of stride * grid.dt + applies it stride times (PR #372 review). Tolerances fixed before the run: the boundary factor + is the per-step factor to the stride power (rtol 1e-14), and a velocity-divergence mode run + 24 fs through an absorbing x boundary at stride 2 matches stride 1 to 1 % of its initial + amplitude (the absorber/propagator splitting at a 100/ps absorber, cs k dt << 1, is ~0.2 %); + the per-step factor applied once per strided update leaves ~20 % more at the wall. + Disclosure: the first run used this config's default tanh absorber, which peaks at 1e4/ps, not + the 100/ps the 1 % bound was derived for; at rate dt ~ 20 the splitting at the layer's inner + edge alone gave 1.4 %. The exp profile at boundary_max_rate 100/ps now sets the rate the bound + assumes; the bound is unchanged.""" + from adept._lpse2d.core.iaw import IonAcousticWave + + def cfg_for(stride): + return _make_cfg( + stride=stride, + dt="1fs", + boundary={"x": "absorbing", "y": "periodic"}, + boundary_max_rate=100.0, + grid_extra={"boundary_width": "1um", "boundary_profile": "exp"}, + ) + + per_step = np.asarray(cfg_for(1)["grid"]["iaw_absorbing_boundaries"]) + np.testing.assert_allclose(-np.log(np.min(per_step)) / cfg_for(1)["grid"]["dt"], 100.0, rtol=1e-9) + assert np.min(per_step) ** 24 < 0.5 # the absorber removes most of the wall field in 24 fs + one, two = IonAcousticWave(cfg_for(1)), IonAcousticWave(cfg_for(2)) + np.testing.assert_allclose(np.asarray(two.boundary), per_step**2, rtol=1e-14) + + cfg = cfg_for(1) + x = np.asarray(cfg["grid"]["x"]) + k = float(np.asarray(cfg["grid"]["kx"])[1]) + amplitude = 0.01 + w0 = jnp.asarray(amplitude * np.cos(k * (x - x[0])))[:, None] * jnp.ones((1, cfg["grid"]["ny"])) + state = {**_zero_fields(cfg), "iaw_density": jnp.zeros_like(w0), "iaw_velocity_divergence": w0} + w_one = np.asarray(_run(one, dict(state), 24)["iaw_velocity_divergence"]) + w_two = np.asarray(_run(two, dict(state), 12)["iaw_velocity_divergence"]) + np.testing.assert_allclose(w_two, w_one, rtol=0, atol=1e-2 * amplitude) diff --git a/tests/test_lpse2d/test_light_coupling.py b/tests/test_lpse2d/test_light_coupling.py index 0519639d..fb9057fd 100644 --- a/tests/test_lpse2d/test_light_coupling.py +++ b/tests/test_lpse2d/test_light_coupling.py @@ -39,6 +39,9 @@ def _make_cfg(coupling="rotation", light_filter=None): cfg = deepcopy(cfg) cfg["grid"].update({"boundary_width": "2um", "xmax": "16um", "tmax": "10fs"}) cfg["terms"]["light"] = {"pump_depletion": True, "coupling": coupling} + if coupling == "rotation": + # the fd rotation exchange cannot project its sources (refused with the default true) + cfg["terms"]["light"]["transverse_source"] = False if light_filter is not None: cfg["terms"]["light"]["filter"] = light_filter cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" # the evolved pump needs an exit @@ -269,11 +272,38 @@ def test_invalid_coupling_and_filter_are_rejected(): with pytest.raises(ValueError, match=r"terms\.light\.coupling"): CoupledLight(cfg) cfg["terms"]["light"]["coupling"] = "rotation" + cfg["terms"]["light"]["transverse_source"] = False cfg["terms"]["light"]["filter"] = 0.0 with pytest.raises(ValueError, match=r"terms\.light\.filter"): CoupledLight(cfg) +def test_fd_rotation_refuses_the_transverse_source_projection_it_cannot_apply(): + """The fd rotation exchange is pointwise and never applied terms.light.transverse_source (PR #372 + review): with SRS on, the default true is refused instead of silently ignored; false sets up, and + so do coupling: explicit and the spectral solver (which keeps the fields transverse instead).""" + from adept._lpse2d.core.light import CoupledLight + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + + cfg = _make_cfg(coupling="rotation") + assert cfg["terms"]["light"]["transverse_source"] is False + assert CoupledLight(cfg).coupling == "rotation" + for value in (True, None): + refused = deepcopy(cfg) + if value is None: + del refused["terms"]["light"]["transverse_source"] + else: + refused["terms"]["light"]["transverse_source"] = value + with pytest.raises(ValueError, match=r"transverse_source"): + CoupledLight(refused) + explicit = _make_cfg(coupling="explicit") + assert explicit["terms"]["light"]["transverse_source"] is True + assert CoupledLight(explicit).transverse_source + spectral = deepcopy(cfg) + spectral["terms"]["light"].update({"solver": "spectral", "transverse_source": True}) + assert SpectralCoupledLight(spectral).transverse_source + + def test_config_validation_of_light_options(): """helpers.get_solver_quantities validates terms.light through LightModel and refuses the coupled-solver options on the prescribed-pump path.""" @@ -298,4 +328,4 @@ def test_config_validation_of_light_options(): cfg["terms"]["light"] = {"pump_depletion": True, "coupling": "rotation", "filter": 0.7} cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" cfg = _finish_cfg(cfg) - assert cfg["terms"]["light"] == {"pump_depletion": True, "coupling": "rotation", "filter": 0.7} + assert {"pump_depletion": True, "coupling": "rotation", "filter": 0.7}.items() <= cfg["terms"]["light"].items() diff --git a/tests/test_lpse2d/test_lpse_deck.py b/tests/test_lpse2d/test_lpse_deck.py new file mode 100644 index 00000000..9b18dc20 --- /dev/null +++ b/tests/test_lpse2d/test_lpse_deck.py @@ -0,0 +1,271 @@ +"""Translation of original-LPSE decks and reading of LPSE outputs (lpse_deck.py). + +The ZAK-unit port is checked against the derived quantities LPSE prints in its own +parameter summary for the test_025 deck; the deck parser and translator are checked on +the shipped example decks; the output readers on a reference run when one is present. +""" + +from pathlib import Path + +import numpy as np +import pytest + +LPSE_ROOT = Path("/home/phil/Desktop/Ergodic-projects/original-lpse") +DECKS = LPSE_ROOT / "examples" / "testRuns" + + +def _finish(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def test_zak_units_match_lpse_printed_summary(): + """test_025: Z = 1, Te = 2 keV, Ti = 1 keV, Mi/Me = 1836, n0/nc = 0.2, 0.351 um. LPSE prints + critical density 9.0491e+21 1/cm^3, envelope density 1.8098e+21, plasma frequency + 2.4000e+15 rad/s, Debye length 7.8148e-07 cm (runs/test_025/lpse.out).""" + from adept._lpse2d.lpse_deck import ZakUnits + + zak = ZakUnits(2.0, 1.0, 1.0, 1836.0, 0.2, 0.351) + assert zak.nc == pytest.approx(9.0491e21, rel=2e-4) + assert zak.n0 == pytest.approx(1.8098e21, rel=2e-4) + assert zak.plasma_frequency == pytest.approx(2.4000e15, rel=2e-4) + assert zak.debye_length_cm == pytest.approx(7.8148e-7, rel=2e-4) + # the dispersion coefficient 3 ve^2 / (2 wpe) is exactly 1 in ZAK units + assert 3.0 * zak.ve**2 / (2.0 * zak.wpe) == pytest.approx(1.0) + # ZAK k unit vs the Debye wavenumber: kde = wpe / ve + assert zak.kde == pytest.approx(zak.wpe / zak.ve) + assert zak.zak_per_ps > 0 and zak.zak_per_um > 0 and np.isfinite(zak.potential_zak_to_adept) + + +@pytest.mark.skipif(not DECKS.exists(), reason="original-lpse example decks not available") +def test_parse_deck_with_include(): + from adept._lpse2d.lpse_deck import parse_parms + + parms = parse_parms(DECKS / "test_006" / "lpse.parms") + assert parms["lw.solver"] == "spectral" and parms["raman.solver"] == "spectral" + assert parms["laser.1.intensity"] == "1e+15" # from laser_include.txt + assert parms["grid.nodes"] == "360 360" and parms["densityProfile.NminLocation"] == "-10 0" + + +@pytest.mark.skipif(not DECKS.exists(), reason="original-lpse example decks not available") +def test_translate_test_006_builds_a_runnable_config(): + from adept._lpse2d.lpse_deck import parse_parms, translate_parms + + parms = parse_parms(DECKS / "test_006" / "lpse.parms") + cfg, report = translate_parms(parms, run="test_006") + assert cfg["terms"]["epw"]["source"] == { + "noise": True, + "noise_model": "thermal", + "noise_debye_factor": False, + "noise_calibrate": False, + "noise_amplitude": pytest.approx(cfg["terms"]["epw"]["source"]["noise_amplitude"]), + "noise_seed": 1, + "tpd": False, + "tpd_form": "lpse", + "srs": True, + "srs_k_filter": False, # LPSE lw.kFilter is off by default + "srs_k_filter_scale": 1.2, + } + assert cfg["terms"]["light"]["solver"] == "spectral" and cfg["terms"]["light"]["pump_depletion"] is False + assert cfg["terms"]["light"]["max_wavenumber"] == 1.0 + assert cfg["terms"]["epw"]["damping"]["landau_form"] == "relativistic" + assert cfg["terms"]["epw"]["damping"]["collisions"] == 0.1 + # EPW dt = min(lw.spectral.dt, maxLightStepsPerStep * raman.dt) = min(0.002, 4 * 0.0001) + assert cfg["grid"]["dt"] == "0.0004ps" and cfg["grid"]["light_substeps"] == 4 + assert cfg["grid"]["dealias"] == "rectangular" and cfg["grid"]["low_pass_filter"] == pytest.approx(0.666) + assert cfg["units"]["laser intensity"] == "1e+15W/cm^2" + assert not report["unsupported"] + + cfg = _finish(cfg) + grid = cfg["grid"] + assert grid["nx"] == 360 and grid["ny"] == 360 + assert grid["dx"] == pytest.approx(20.0 / 359.0) + n = np.asarray(grid["background_density"])[:, 0] + x = np.asarray(grid["x"]) + # linear from 0.2 at the N_min location (x = 0) to 0.28 at N_max (x = 20), clipped outside + inside = (x > 0.0) & (x < 20.0) + np.testing.assert_allclose(n[inside], 0.2 + 0.08 * x[inside] / 20.0, rtol=1e-12) + assert n.min() >= 0.2 - 1e-12 and n.max() <= 0.28 + 1e-12 + # rectangular anti-aliasing: outer 33.4 % of each axis zeroed + band = np.asarray(grid["low_pass_filter_grid"]) + kx = np.abs(np.asarray(grid["kx"]))[:, None] + ky = np.abs(np.asarray(grid["ky"]))[None, :] + expected = (kx < 0.666 * kx.max()) & (ky < 0.666 * ky.max()) + np.testing.assert_array_equal(band > 0, expected) + + +@pytest.mark.skipif(not DECKS.exists(), reason="original-lpse example decks not available") +def test_translate_combined_iaw_deck(): + from adept._lpse2d.lpse_deck import parse_parms, translate_parms + + cfg, report = translate_parms(parse_parms(DECKS / "test_032" / "lpse.parms"), run="test_032") + assert cfg["terms"]["epw"]["solver"] == "combined" + assert cfg["terms"]["epw"]["source"]["tpd"] and cfg["terms"]["epw"]["source"]["srs"] + assert cfg["terms"]["iaw"]["active"] and cfg["terms"]["iaw"]["solver"] == "spectral" + assert cfg["terms"]["light"]["solver"] == "spectral" + cfg = _finish(cfg) + assert cfg["terms"]["iaw"]["damping"]["landau"] > 0.0 + + +def test_translator_disabled_laser_is_no_pump(): + """LPSE laser.enable = false switches the E0 light off whatever beams the deck keeps (PR #372 + review): the translation carries no intensity, no beams and no depletion even with both + pumpDepletion flags on, and the static pump it sets up is exactly zero. Exact checks.""" + from adept._lpse2d.core.laser import Light + from adept._lpse2d.lpse_deck import translate_parms + + base = { + "grid.sizes": "20 5", + "grid.nodes": "201 51", + "lw.enable": "true", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.solver": "static", + "laser.nBeams": "1", + "laser.1.intensity": "1e15", + "laser.1.frequencyShift": "0.01", + "laser.pumpDepletion.TPD.enable": "true", + "laser.pumpDepletion.SRS.enable": "true", + } + on, _ = translate_parms({**base, "laser.enable": "true"}, run="y") + assert on["units"]["laser intensity"] == "1e+15W/cm^2" and on["terms"]["light"]["pump_depletion"] is True + + for parms in ({**base, "laser.enable": "false"}, base): # LPSE's default is false + cfg, report = translate_parms(parms, run="y") + assert cfg["units"]["laser intensity"] == "0W/cm^2" + assert "beams" not in cfg["drivers"]["E0"] + assert cfg["terms"]["light"]["pump_depletion"] is False + assert any("laser.enable = false" in note for note in report["notes"]) + + cfg = _finish(cfg) + ny = cfg["grid"]["ny"] + wave = {"delta_omega": np.array([0.0]), "intensities": np.ones((1, ny)), "phases": np.zeros((1, ny))} + assert np.max(np.abs(np.asarray(Light(cfg).laser_update(0.05, None, wave)))) == 0.0 + + with pytest.raises(ValueError, match=r"thermalFil\.laser\.enable"): + translate_parms({**base, "iaw.enable": "true", "thermalFil.laser.enable": "true"}, run="y") + + +def test_translator_runs_lpse_square_cells(): + """LPSE runs square cells h = min(size / (nodes - 1)) over the deck's axes, rescales the box to + h (nodes - 1) and refuses a mismatch above 10 % of h (ParameterManager.cpp:1258-1281; PR #372 + review). Exact checks: a deck finer in y runs at h = dy on both axes (the old translation used + dx), a deck coarser in y keeps h = dx with LPSE's rescaled y box, each with a report note, and a + 100 % mismatch raises.""" + from adept._lpse2d.lpse_deck import translate_parms + + base = { + "laser.enable": "true", + "lw.enable": "true", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.nBeams": "1", + "laser.1.intensity": "1e15", + "grid.nodes": "201 50", + } + # the translated grid (adept's own setup may still round nx to an FFT-friendly size) + for size_y, h in (("4.851", 0.099), ("5.049", 0.1)): # dy = 0.099 < dx = 0.1, dy = 0.10304 > dx + cfg, report = translate_parms({**base, "grid.sizes": f"20 {size_y}"}, run="y") + assert any("square cells" in note for note in report["notes"]) + grid = cfg["grid"] + assert float(grid["dx"].removesuffix("um")) == pytest.approx(h, rel=1e-12) + assert float(grid["xmax"].removesuffix("um")) == pytest.approx(201 * h, rel=1e-12) + assert float(grid["ymax"].removesuffix("um")) == pytest.approx(0.5 * 50 * h, rel=1e-12) + assert float(grid["ymin"].removesuffix("um")) == pytest.approx(-0.5 * 50 * h, rel=1e-12) + cfg, report = translate_parms({**base, "grid.sizes": "20 4.9"}, run="y") # dy = dx: no note + assert not any("square cells" in note for note in report["notes"]) + with pytest.raises(ValueError, match="square-cell"): + translate_parms({**base, "grid.sizes": "20 9.8"}, run="y") # dy = 2 dx + + +def test_fd_pump_injector_refuses_the_time_controls_it_does_not_apply(tmp_path): + """The FD two-point pump injector (terms.light.pump_depletion with terms.light.solver: fd) reads + neither drivers.E0.pulse_file nor kap_bandwidth; setup refuses them there instead of injecting the + unmodulated pump (PR #372 review), and the spectral injector, which applies both, accepts them. + A normal-incidence pump with a zero-amplitude pulse segment and a nonzero KAP bandwidth.""" + import yaml + + pulse = tmp_path / "pulse.txt" + pulse.write_text("0.0 1.0\n0.01 1.0\n0.011 0.0\n1.0 0.0\n") + + def cfg_for(solver, **e0): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg["grid"].update({"xmax": "12.8um", "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um", "dx": "0.1um"}) + cfg["terms"]["light"] = {"solver": solver, "pump_depletion": True} + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["terms"]["epw"]["boundary"] = {"x": "absorbing", "y": "periodic"} + cfg["drivers"]["E0"].update(e0) + return cfg + + for e0 in ( + {"pulse_file": str(pulse)}, + {"kap_bandwidth": 0.005}, + {"pulse_file": str(pulse), "kap_bandwidth": 0.005}, + ): + with pytest.raises(ValueError, match=r"terms\.light\.solver: spectral"): + _finish(cfg_for("fd", **e0)) + _finish(cfg_for("spectral", **e0)) + _finish(cfg_for("fd")) + + +def test_lpse_exp_profile_and_thermal_noise_without_debye_factor(): + import yaml + + from adept._lpse2d.core.epw import noise_kick_spectrum + + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg["density"] = {"basis": "lpse-exp", "min": 0.2, "max": 0.28, "min_location": "2um", "max_location": "8um"} + cfg["grid"].update({"xmax": "10um", "dx": "0.1um", "ymax": "0.2um", "ymin": "-0.2um", "tmax": "10fs", "dt": "1fs"}) + cfg["terms"]["epw"]["source"].update( + {"noise": True, "noise_model": "thermal", "noise_debye_factor": False, "tpd": False} + ) + cfg["terms"]["epw"]["damping"]["collisions"] = 1.0 + cfg = _finish(cfg) + n = np.asarray(cfg["grid"]["background_density"])[:, 0] + x = np.asarray(cfg["grid"]["x"]) + ln = 6.0 / np.log(0.28 / 0.2) + inside = (x > 2.0) & (x < 8.0) + np.testing.assert_allclose(n[inside], 0.28 * np.exp(-(8.0 - x[inside]) / ln), rtol=1e-12) + assert n[x < 2.0].max() == pytest.approx(0.2) and n[x > 8.0].min() == pytest.approx(0.28) + + kick = np.asarray(noise_kick_spectrum(cfg)) + cfg["terms"]["epw"]["source"]["noise_debye_factor"] = True + kick_debye = np.asarray(noise_kick_spectrum(cfg)) + derived = cfg["units"]["derived"] + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + lam_sq = derived["vte_sq"] / derived["wp0"] ** 2 + band = kick > 0 + np.testing.assert_allclose(kick[band] / kick_debye[band], np.sqrt(1.0 + (kx**2 + ky**2) * lam_sq)[band], rtol=1e-12) + + +RUN_025 = LPSE_ROOT / "runs" / "test_025" / "data" + + +@pytest.mark.skipif(not RUN_025.exists(), reason="no LPSE reference run present") +def test_read_lpse_metrics_and_frames(): + from adept._lpse2d.lpse_deck import read_frames, read_metrics + + metrics = read_metrics(RUN_025 / "lpse.metrics") + assert metrics["time"][0] == pytest.approx(0.01) and metrics["time"][-1] >= 1.0 # one sample past the end + assert "EPW_energy" in metrics and np.all(metrics["EPW_energy"] > 0) + frames = read_frames(RUN_025 / "lpse.pots") + assert len(frames) == 2 + header, pots = frames[0] + assert header["Nx"] == 720 and header["Ny"] == 90 and header["time"] == pytest.approx(0.5) + assert pots.shape == (720, 90) and np.iscomplexobj(pots) and np.all(np.isfinite(pots)) + assert frames[1][0]["time"] == pytest.approx(1.0) + # the x-space field is smooth along x on the scale of one cell but not constant + assert np.abs(pots).std() > 0 diff --git a/tests/test_lpse2d/test_lpse_density.py b/tests/test_lpse2d/test_lpse_density.py new file mode 100644 index 00000000..530ac592 --- /dev/null +++ b/tests/test_lpse2d/test_lpse_density.py @@ -0,0 +1,160 @@ +"""The original LPSE density profiles (density.basis: lpse-).""" + +from copy import deepcopy + +import numpy as np +import pytest +import yaml + + +def _cfg(density, ny=8): + from adept._lpse2d.helpers import get_density_profile, get_derived_quantities, get_solver_quantities, write_units + + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update({"dx": "0.1um", "xmax": "20um", "ymax": "0.4um", "ymin": "-0.4um", "dt": "1fs", "tmax": "2fs"}) + cfg["density"] = density + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + return cfg, np.asarray(get_density_profile(cfg)) + + +def _at(cfg, prof, x_um): + return float(prof[int(np.argmin(np.abs(np.asarray(cfg["grid"]["x"]) - x_um))), 0]) + + +BASE = {"min": 0.1, "max": 0.3, "min_location": "2um", "max_location": "18um"} + + +def test_linear_and_exp_hit_the_endpoints_and_clip(): + cfg, lin = _cfg({"basis": "lpse-linear", **BASE}) + assert abs(_at(cfg, lin, 2.0) - 0.1) < 2e-3 and abs(_at(cfg, lin, 18.0) - 0.3) < 2e-3 + assert abs(_at(cfg, lin, 10.0) - 0.2) < 2e-3 + assert lin.min() >= 0.1 and lin.max() <= 0.3 # clipped beyond the endpoints + cfg, ex = _cfg({"basis": "lpse-exp", **BASE}) + assert abs(_at(cfg, ex, 2.0) - 0.1) < 2e-3 and abs(_at(cfg, ex, 18.0) - 0.3) < 2e-3 + assert abs(_at(cfg, ex, 10.0) - np.sqrt(0.1 * 0.3)) < 2e-3 # geometric mean half way + + +@pytest.mark.parametrize("p", [2.0, 4.0]) +def test_gaussian_and_inverse_power_match_lpse_formulas(p): + cfg, g = _cfg({"basis": "lpse-gaussian", "sg_order": p, **BASE}) + assert abs(_at(cfg, g, 18.0) - 0.3) < 2e-3 and abs(_at(cfg, g, 2.0) - 0.1) < 3e-3 + sd = 16.0 / np.log(3.0) ** (1.0 / p) + assert abs(_at(cfg, g, 10.0) - 0.3 * np.exp(-((8.0 / sd) ** p))) < 3e-3 + cfg, ip = _cfg({"basis": "lpse-inverse-power", "sg_order": p, **BASE}) + assert abs(_at(cfg, ip, 2.0) - 0.1) < 3e-3 + rc = 16.0 * (0.1 / 0.3) ** (1.0 / p) + assert abs(_at(cfg, ip, 18.0 - 0.5 * rc) - 0.3) < 1e-9 # flat top inside r_c + assert abs(_at(cfg, ip, 10.0) - 0.1 * (16.0 / 8.0) ** p) < 3e-3 or _at(cfg, ip, 10.0) == 0.3 + + +def test_quadratic_passes_through_its_three_points(): + d = { + "basis": "lpse-quadratic", + "min": 0.1, + "max": 0.3, + "min_location": "5um", + "max_location": "15um", + "central_density": 0.2, + "origin": "10um", + } + cfg, q = _cfg(d) + assert abs(_at(cfg, q, 10.0) - 0.2) < 2e-3 + assert abs(_at(cfg, q, 5.0) - 0.1) < 2e-3 + assert abs(_at(cfg, q, 15.0) - 0.3) < 2e-3 + + +def test_dips_sit_on_the_linear_ramp(): + cfg, lin = _cfg({"basis": "lpse-linear", **BASE}) + dip = {"dip_depth": 0.05, "dip_width": "4um", "dip_offset": "0um", "origin": "10um"} + cfg, qd = _cfg({"basis": "lpse-qd", **BASE, **dip}) + assert abs(_at(cfg, qd, 10.0) - (_at(cfg, lin, 10.0) - 0.05)) < 2e-3 # full depth at the centre + assert abs(_at(cfg, qd, 12.0) - _at(cfg, lin, 12.0)) < 2e-3 # zero at the edge of the dip + assert abs(_at(cfg, qd, 15.0) - _at(cfg, lin, 15.0)) < 1e-9 # untouched outside + cfg, gd = _cfg({"basis": "lpse-gd", "sg_order": 2.0, **BASE, **dip}) + assert abs(_at(cfg, gd, 10.0) - (_at(cfg, lin, 10.0) - 0.05)) < 2e-3 + assert abs(_at(cfg, gd, 12.0) - (_at(cfg, lin, 12.0) - 0.05 * np.exp(-1.0))) < 2e-3 + + +def test_spherical_geometry_and_clip(): + cfg, s = _cfg({"basis": "lpse-linear", "geometry": "spherical", **BASE}) + # radial distance from the max location: symmetric about x = 18 um, so 20 um reads like 16 um + assert abs(_at(cfg, s, 20.0) - _at(cfg, s, 16.0)) < 3e-3 + cfg, c = _cfg({"basis": "lpse-linear", "min": 1.0, "max": 2.0, "min_location": "2um", "max_location": "18um"}) + assert c.max() == 1.25 + + +def test_file_profile_roundtrip(tmp_path): + cfg, lin = _cfg({"basis": "lpse-linear", **BASE}) + path = tmp_path / "n.npy" + np.save(path, lin) + cfg2, loaded = _cfg({"basis": "lpse-file", "file": str(path)}) + np.testing.assert_allclose(loaded, lin) + np.save(path, lin[:, 0]) + cfg3, loaded1d = _cfg({"basis": "lpse-file", "file": str(path)}) + np.testing.assert_allclose(loaded1d, lin) + with pytest.raises(ValueError, match="shape"): + np.save(path, lin[:-1]) + _cfg({"basis": "lpse-file", "file": str(path)}) + + +def test_translator_maps_the_shape_keys(): + from adept._lpse2d.lpse_deck import translate_parms + + base = { + "grid.sizes": "20 5", + "grid.nodes": "201 51", + "laser.enable": "true", + "lw.enable": "true", + "densityProfile.NminOverNc": "0.2", + "densityProfile.NmaxOverNc": "0.28", + "densityProfile.NminLocation": "-10 0", + "densityProfile.NmaxLocation": "10 0", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.1.intensity": "1e15", + } + dip = { + "densityProfile.shape": "gd", + "densityProfile.sgOrder": "4", + "densityProfile.dip.depth": "0.02", + "densityProfile.dip.width": "3", + "densityProfile.dip.offset": "1", + } + parms = dict(base, **dip) + cfg, report = translate_parms(parms, experiment="x", run="y") + d = cfg["density"] + assert d["basis"] == "lpse-gd" and d["sg_order"] == 4.0 and d["dip_depth"] == 0.02 + assert d["dip_width"] == "3.0um" and d["dip_offset"] == "1.0um" and d["origin"] == "10.0um" + assert d["min_location"] == "0.0um" and d["max_location"] == "20.0um" + cfg, report = translate_parms(dict(base, **{"densityProfile.shape": "inverseSquare"}), experiment="x", run="y") + assert cfg["density"]["basis"] == "lpse-inverse-power" and cfg["density"]["sg_order"] == 2.0 + + +def _write_lpse_frame(path, arr): + """One real frame in LPSE's binary field format (x fastest, single precision).""" + nx, ny = arr.shape + header = f"# BeginHeaderSegment; Nx = {nx}, Ny = {ny}, Nz = 1, FileType = real, time = 0; # EndHeaderSegment;\n" + data = np.ascontiguousarray(np.asarray(arr, dtype=" infinity; +- the thermal (fluctuation-dissipation) EPW noise reaches the Cerenkov spectrum + ``<|E_k|^2> = A^2 / (1 + k^2 lambda_D^2)`` per mode, independent of dt; +- the LPSE exponential absorbing-layer profile and ``max_wavenumber`` band cap. +""" + +from copy import deepcopy + +import numpy as np +import pytest +import yaml +from jax import numpy as jnp + + +def _finish_cfg(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _base_cfg(envelope_density=0.25, box_density=0.25, **grid): + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["units"]["envelope density"] = envelope_density + cfg["density"] = {"basis": "uniform", "val": box_density} + cfg["grid"].update( + { + "boundary_width": "0.4um", + "dt": "1fs", + "dx": "0.1um", + "xmax": "6.4um", + "tmax": "10fs", + "ymax": "1.6um", + "ymin": "-1.6um", + } + ) + cfg["grid"].update(grid) + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": True, "srs": False}) + cfg["terms"]["epw"]["density_gradient"] = False + return cfg + + +def _random_phi(cfg, seed): + rng = np.random.default_rng(seed) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + phi_k = rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny)) + phi_k *= np.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"]) + return jnp.asarray(phi_k) + + +# ------------------------------------------------------------------ TPD source -- + + +def test_tpd_forms_coincide_at_quarter_critical_and_differ_elsewhere(): + """At envelope density 0.25, w0 = 2 wp0 exactly and LPSE's coefficient/factor equal the + MATLAB prototype's; at 0.23 the LPSE coefficient i e/(4 me w0) is 4% smaller than + i e/(8 me wp0) and the charge-density factor is (w0/wp0 - 1) instead of 1.""" + from adept._lpse2d.core.epw import SpectralEPWSolver + + for n_env, same in ((0.25, True), (0.23, False)): + cfg_lpse = _finish_cfg(_base_cfg(envelope_density=n_env, box_density=n_env)) + cfg_matlab = deepcopy(cfg_lpse) + cfg_matlab["terms"]["epw"]["source"]["tpd_form"] = "matlab" + lpse = SpectralEPWSolver(cfg_lpse) + matlab = SpectralEPWSolver(cfg_matlab) + assert lpse.tpd_form == "lpse" and matlab.tpd_form == "matlab" + assert lpse.tpd_prefactor == pytest.approx(1j * lpse.e / (4.0 * lpse.w0 * lpse.me)) + assert matlab.tpd_prefactor == pytest.approx(1j * lpse.e / (8.0 * lpse.wp0 * lpse.me)) + assert lpse.tpd_rho_factor == pytest.approx(lpse.w0 / lpse.wp0 - 1.0) + assert matlab.tpd_rho_factor == 1.0 + + phi_k = _random_phi(cfg_lpse, 11) + rng = np.random.default_rng(12) + nx, ny = cfg_lpse["grid"]["nx"], cfg_lpse["grid"]["ny"] + E0 = jnp.asarray(rng.normal(size=(nx, ny, 2)) + 1j * rng.normal(size=(nx, ny, 2))) + ex, ey = lpse.phi_k_to_e_fields(phi_k) + s_lpse = np.asarray(lpse.calc_tpd_source(0.3, phi_k, ex, ey, E0)) + s_matlab = np.asarray(matlab.calc_tpd_source(0.3, phi_k, ex, ey, E0)) + if same: + np.testing.assert_allclose(s_lpse, s_matlab, rtol=1e-12, atol=1e-12 * np.abs(s_lpse).max()) + else: + assert not np.allclose(s_lpse, s_matlab, rtol=1e-3) + # prefactor ratio 2 wp0/w0 = 2 sqrt(0.23) = 0.959; rho factor 1/sqrt(0.23) - 1 = 1.085 + assert abs(lpse.tpd_prefactor / matlab.tpd_prefactor) == pytest.approx(2.0 * np.sqrt(0.23), rel=1e-12) + assert lpse.tpd_rho_factor == pytest.approx(1.0 / np.sqrt(0.23) - 1.0, rel=1e-12) + + +def test_tpd_source_uses_both_pump_components(): + """An x-polarized pump drives TPD through E0x . Ex* and kx E0x rho* exactly like the + y-polarized one drives it through the y terms (LPSE uses the full dot product).""" + from adept._lpse2d.core.epw import SpectralEPWSolver + + cfg = _finish_cfg(_base_cfg()) + epw = SpectralEPWSolver(cfg) + phi_k = _random_phi(cfg, 21) + ex, ey = epw.phi_k_to_e_fields(phi_k) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + rng = np.random.default_rng(22) + field = jnp.asarray(rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) + zeros = jnp.zeros_like(field) + E0_x = jnp.stack([field, zeros], axis=-1) + E0_y = jnp.stack([zeros, field], axis=-1) + + s_x = np.asarray(epw.calc_tpd_source(0.1, phi_k, ex, ey, E0_x)) + s_y = np.asarray(epw.calc_tpd_source(0.1, phi_k, ex, ey, E0_y)) + assert np.abs(s_x).max() > 0.0 + # hand-built reference for the x-polarized pump + rho = jnp.fft.ifft2(epw.k_sq * phi_k) + tpd1 = jnp.fft.fft2(field * jnp.conj(ex)) + tpd2 = 1j * epw.kx[:, None] * jnp.fft.fft2(field * jnp.conj(rho)) * epw.one_over_k_sq + expected = epw.tpd_prefactor * jnp.exp(-1j * (epw.w0 - 2.0 * epw.wp0) * 0.1) * (tpd1 + epw.tpd_rho_factor * tpd2) + expected = np.asarray(expected * epw.low_pass_filter * epw.zero_mask) + np.testing.assert_allclose(s_x, expected, rtol=1e-11, atol=1e-11 * np.abs(expected).max()) + # linear in the pump: the sum of the two polarizations is the source of the combined pump + s_xy = np.asarray(epw.calc_tpd_source(0.1, phi_k, ex, ey, E0_x + E0_y)) + np.testing.assert_allclose(s_xy, s_x + s_y, rtol=1e-11, atol=1e-11 * np.abs(s_xy).max()) + + +# ------------------------------------------------------------ TPD depletion -- + + +def test_tpd_depletion_transverse_projection_is_divergence_free(): + """[E div E]_T has zero divergence mode by mode, the projection is idempotent, and + switching it off returns the raw E div E for both components.""" + from adept._lpse2d.core.light import CoupledLight + + cfg = _base_cfg() + cfg["terms"]["light"] = {"pump_depletion": True} + cfg = _finish_cfg(cfg) + light = CoupledLight(cfg) + assert light.tpd_projection + phi_k = _random_phi(cfg, 31) + + f_t = np.asarray(light.tpd_depletion_vector(phi_k)) + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + fx_k, fy_k = np.fft.fft2(f_t[..., 0]), np.fft.fft2(f_t[..., 1]) + div_k = kx * fx_k + ky * fy_k + assert np.abs(div_k).max() < 1e-10 * (np.abs(kx * fx_k) + np.abs(ky * fy_k)).max() + + # idempotent: projecting the projected field changes nothing + longitudinal = kx * fx_k + ky * fy_k + k_sq = kx**2 + ky**2 + k_sq_safe = np.where(k_sq > 0, k_sq, 1.0) + fx_again = fx_k - np.where(k_sq > 0, kx / k_sq_safe, 0.0) * longitudinal + np.testing.assert_allclose(fx_again, fx_k, rtol=0, atol=1e-10 * np.abs(fx_k).max()) + + cfg_raw = deepcopy(cfg) + cfg_raw["terms"]["light"]["tpd_projection"] = False + raw = CoupledLight(cfg_raw) + f_raw = np.asarray(raw.tpd_depletion_vector(phi_k)) + ex = np.asarray(jnp.fft.ifft2(-1j * jnp.asarray(cfg["grid"]["kx"])[:, None] * phi_k)) + ey = np.asarray(jnp.fft.ifft2(-1j * jnp.asarray(cfg["grid"]["ky"])[None, :] * phi_k)) + rho = np.asarray(jnp.fft.ifft2(light.k_sq * phi_k)) + np.testing.assert_allclose(f_raw[..., 0], ex * rho, rtol=1e-11, atol=1e-11 * np.abs(ex * rho).max()) + np.testing.assert_allclose(f_raw[..., 1], ey * rho, rtol=1e-11, atol=1e-11 * np.abs(ey * rho).max()) + # the projection only removes the longitudinal part: transverse content is unchanged + fx_raw_k, fy_raw_k = np.fft.fft2(f_raw[..., 0]), np.fft.fft2(f_raw[..., 1]) + transverse_raw = ky * fx_raw_k - kx * fy_raw_k + transverse_t = ky * fx_k - kx * fy_k + np.testing.assert_allclose(transverse_t, transverse_raw, rtol=0, atol=1e-10 * np.abs(transverse_raw).max()) + + +def test_tpd_depletion_k_filter_masks_light_band(): + from adept._lpse2d.core.light import CoupledLight + + cfg = _base_cfg() + cfg["terms"]["light"] = {"pump_depletion": True, "tpd_k_filter": True} + cfg = _finish_cfg(cfg) + light = CoupledLight(cfg) + mask = np.asarray(light.tpd_k_mask) + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + k_max_sq = (1.2 * light.w0 / light.c) ** 2 * (1.0 - 0.25) + np.testing.assert_array_equal(mask, np.where((kx**2 + ky**2 > 0) & (kx**2 + ky**2 < k_max_sq), 1.0, 0.0)) + f_t = np.asarray(light.tpd_depletion_vector(_random_phi(cfg, 41))) + for comp in range(2): + outside = np.fft.fft2(f_t[..., comp]) * (1.0 - mask) + assert np.abs(outside).max() < 1e-10 * np.abs(np.fft.fft2(f_t[..., comp])).max() + + +# ---------------------------------------------------------------- Landau -- + + +def test_lpse_landau_form_prefactor_ratio(): + """LPSE's sqrt(1 + 3x^2) vs MATLAB's (1 + 1.5x^2), x = k lambda_D, with identical + exponents: 0.7% at x = 0.3 and 3.9% at x = 0.5 (hand check in the vault gap audit).""" + from adept._lpse2d.core.epw import landau_damping_rate + + wp0, vte_sq = 2.0, 0.01 + lam_d = np.sqrt(vte_sq) / wp0 + for x, expected_pct in ((0.3, 0.7), (0.5, 3.8)): + k_sq = jnp.asarray([[(x / lam_d) ** 2]]) + g_matlab = float(landau_damping_rate(k_sq, wp0, vte_sq, jnp.ones((1, 1)), form="matlab")[0, 0]) + g_lpse = float(landau_damping_rate(k_sq, wp0, vte_sq, jnp.ones((1, 1)), form="lpse")[0, 0]) + ratio = np.sqrt(1.0 + 3.0 * x**2) / (1.0 + 1.5 * x**2) + assert g_lpse / g_matlab == pytest.approx(ratio, rel=1e-12) + assert abs(1.0 - ratio) * 100.0 == pytest.approx(expected_pct, abs=0.06) # 0.70%, 3.79% + + +def test_landau_threshold_multiplier_and_form_flow_through_the_solver(): + from adept._lpse2d.core.epw import SpectralEPWSolver, analytic_landau_rate, landau_damping_rate + + cfg = _finish_cfg(_base_cfg()) + base = np.asarray(analytic_landau_rate(cfg)) + k_sq = jnp.asarray(cfg["grid"]["kx"])[:, None] ** 2 + jnp.asarray(cfg["grid"]["ky"])[None, :] ** 2 + derived = cfg["units"]["derived"] + ref = np.asarray(landau_damping_rate(k_sq, derived["wp0"], derived["vte_sq"], jnp.where(k_sq > 0, 1.0, 0.0))) + np.testing.assert_allclose(base, ref, rtol=1e-12) + + threshold = float(np.median(base[base > 0])) + cfg2 = deepcopy(cfg) + cfg2["terms"]["epw"]["damping"].update( + {"landau_form": "lpse", "landau_lower_threshold": threshold, "landau_multiplier": 0.5} + ) + rate = np.asarray(analytic_landau_rate(cfg2)) + ref2 = np.asarray( + landau_damping_rate(k_sq, derived["wp0"], derived["vte_sq"], jnp.where(k_sq > 0, 1.0, 0.0), form="lpse") + ) + expected = np.where(ref2 > threshold, ref2, 0.0) * 0.5 + np.testing.assert_allclose(rate, expected, rtol=1e-12) + assert np.any(rate == 0.0) and np.any(rate > 0.0) + # the solver applies exactly this array + solver = SpectralEPWSolver(cfg2) + np.testing.assert_allclose(np.asarray(solver.calc_landau_damping_rate()), expected, rtol=1e-12) + + +def test_relativistic_landau_rate_reduces_to_nonrelativistic(): + """For c -> infinity (Ve/C -> 0) the Maxwell-Juettner 2-D/3-D rates go to + sqrt(pi/8) (kde/k)^3 wp omega exp(-omega^2/(2 x^2)) with the *expanded* Bohm-Gross + omega = 1 + 1.5 x^2 that LPSE's relativistic forms carry (K_nu(a) ~ sqrt(pi/2a) e^{-a}); + that differs from the classical sqrt(1 + 3x^2) form by exp(-1.125 x^2) in the exponent. + At 2 keV the relativistic rates differ from that limit by a few percent in the band.""" + from adept._lpse2d.core.epw import landau_damping_rate_relativistic + + wp0, vte_sq = 2.0, 0.01 + lam_d = np.sqrt(vte_sq) / wp0 + x = np.linspace(0.25, 0.5, 6) + k_sq = (x / lam_d) ** 2 + omega = 1.0 + 1.5 * x**2 + classical = np.sqrt(np.pi / 8.0) / x**3 * wp0 * omega * np.exp(-(omega**2) / (2.0 * x**2)) + c_large = np.sqrt(vte_sq) * 1.0e3 + for ndim in (2, 3): + rel = landau_damping_rate_relativistic(k_sq, wp0, vte_sq, c_large, ndim=ndim) + np.testing.assert_allclose(rel, classical, rtol=3e-3) + c_2kev = np.sqrt(vte_sq) / np.sqrt(2.0 / 511.0) + rel = landau_damping_rate_relativistic(k_sq, wp0, vte_sq, c_2kev, ndim=2) + ratio = rel / classical + # relativistic mass reduces the damping of fast modes (v_phi ~ 0.25 c at x = 0.25: + # -38%) and the correction fades toward x = 0.5 (-5%), monotonically + assert np.all((ratio > 0.5) & (ratio < 1.0)) and np.all(np.diff(ratio) > 0.0) + # modes whose phase velocity exceeds c are undamped + rel_slow = landau_damping_rate_relativistic(np.asarray([(0.02 / lam_d) ** 2]), wp0, vte_sq, c_2kev, ndim=2) + assert rel_slow[0] == 0.0 + + +# ----------------------------------------------------------------- noise -- + + +def _noise_cfg(model, **source): + with open("tests/test_lpse2d/configs/epw.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["density"] = {"basis": "uniform", "val": 0.25} + cfg["grid"].update( + { + "boundary_width": "0.2um", + "dt": "1fs", + "dx": "0.05um", + "xmax": "3.2um", + "tmax": "10fs", + "ymax": "0.8um", + "ymin": "-0.8um", + } + ) + cfg["grid"]["low_pass_filter"] = 1.0 + cfg["terms"]["zero_mask"] = True + cfg["terms"]["epw"]["boundary"] = {"x": "periodic", "y": "periodic"} + cfg["terms"]["epw"]["damping"] = {"collisions": False, "landau": True} + cfg["terms"]["epw"]["source"] = { + "noise": True, + "noise_model": model, + "noise_seed": 5, + "tpd": False, + "srs": False, + **source, + } + cfg["terms"]["epw"]["density_gradient"] = False + return _finish_cfg(cfg) + + +def test_flat_noise_model_is_unchanged(): + """noise_model: flat keeps the MATLAB kick dt * amplitude on every retained mode.""" + from adept._lpse2d.core.epw import SpectralEPWSolver, noise_kick_spectrum + + cfg = _noise_cfg("flat", noise_amplitude=3e-9) + kick = np.asarray(noise_kick_spectrum(cfg)) + band = np.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"]) + np.testing.assert_allclose(kick, cfg["grid"]["dt"] * 3e-9 * band, rtol=1e-14) + solver = SpectralEPWSolver(cfg) + noise = np.asarray(solver.get_noise(jnp.asarray(4.0 * solver.dt))) + np.testing.assert_allclose(np.abs(noise), kick, rtol=1e-12) + + +def test_thermal_noise_requires_damping(): + from adept._lpse2d.core.epw import noise_kick_spectrum + + cfg = _noise_cfg("thermal") + cfg["terms"]["epw"]["damping"]["landau"] = False + with pytest.raises(ValueError, match="thermal"): + noise_kick_spectrum(cfg) + + +@pytest.mark.parametrize("dt_fs", [1.0, 4.0]) +def test_thermal_noise_equilibrates_to_cerenkov_spectrum(dt_fs): + """Sources off, uniform periodic box: every mode is phi *= e^{-gamma dt}; phi += D e^{i theta}, + whose stationary |E_k|^2 is A^2 / (1 + k^2 lambda_D^2) (x-space amplitude squared) for + the thermal kick -- for any dt.""" + import jax + + from adept._lpse2d.core.epw import SpectralEPWSolver + + amplitude = 2.0e-3 + cfg = _noise_cfg("thermal", noise_amplitude=amplitude) + cfg["grid"]["dt"] = dt_fs * 1e-3 + cfg["grid"]["absorbing_boundaries"] = np.ones_like(np.asarray(cfg["grid"]["absorbing_boundaries"])) + solver = SpectralEPWSolver(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + derived = cfg["units"]["derived"] + gamma = np.asarray(solver.landau_rate) + k_sq = np.asarray(solver.k_sq) + lam_sq = derived["vte_sq"] / derived["wp0"] ** 2 + + # compare only modes that relax within the run and are in the retained band + n_steps, n_burn = 6000, 2000 + relaxed = (2.0 * gamma * solver.dt > 0.01) & (np.asarray(solver.noise_kick) > 0.0) + assert relaxed.sum() > 100 + + E0 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + + def step(carry, i): + phi_k, acc = carry + t = i * solver.dt + phi_k = solver(t, {"epw": phi_k, "E0": E0, "E1": E0}, None) + acc = acc + jnp.where(i >= n_burn, k_sq * jnp.abs(phi_k) ** 2, 0.0) + return (phi_k, acc), None + + phi0 = jnp.zeros((nx, ny), dtype=jnp.complex128) + (phi_k, acc), _ = jax.lax.scan(jax.jit(step), (phi0, jnp.zeros((nx, ny))), jnp.arange(n_steps)) + measured = np.asarray(acc) / (n_steps - n_burn) / (nx * ny) ** 2 + expected = amplitude**2 / (1.0 + k_sq * lam_sq) + ratio = measured[relaxed] / expected[relaxed] + # per-mode samples are exponentially distributed with ~ (n/(1/(2 gamma dt))) effective + # draws; the band mean is far tighter than any single mode + assert ratio.mean() == pytest.approx(1.0, abs=0.05), ratio.mean() + assert np.median(ratio) == pytest.approx(1.0, abs=0.1) + + +def test_thermal_noise_calibration_scales_with_temperature_and_volume(): + from adept._lpse2d.core.epw import noise_kick_spectrum + + cfg = _noise_cfg("thermal", noise_amplitude=1.0, noise_calibrate=True) + kick = np.asarray(noise_kick_spectrum(cfg)) + cfg_uncal = _noise_cfg("thermal", noise_amplitude=1.0) + kick_uncal = np.asarray(noise_kick_spectrum(cfg_uncal)) + derived = cfg["units"]["derived"] + from astropy.units import Quantity as _Q + + te_kev = _Q(cfg["units"]["reference electron temperature"]).to("keV").value + energy_scale = derived["massScale"] * derived["spatialScale"] ** 2 / derived["timeScale"] ** 2 + kT = te_kev * 1.602176634e-9 / energy_scale + lx = cfg["grid"]["nx"] * cfg["grid"]["dx"] + ly = cfg["grid"]["ny"] * cfg["grid"]["dy"] + a_thermal = np.sqrt(8.0 * np.pi * kT / (lx * ly * ly)) + band = kick_uncal > 0 + np.testing.assert_allclose(kick[band] / kick_uncal[band], a_thermal, rtol=1e-12) + assert 0.0 < a_thermal < 1.0 + + +def test_noise_max_wavenumber_masks_the_kick(): + from adept._lpse2d.core.epw import noise_kick_spectrum + + cfg = _noise_cfg("flat", noise_max_wavenumber=2.0) + kick = np.asarray(noise_kick_spectrum(cfg)) + derived = cfg["units"]["derived"] + k0 = derived["w0"] / derived["c"] + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + outside = np.sqrt(kx**2 + ky**2) >= 2.0 * k0 + assert outside.any() and (~outside).sum() > 10 + assert np.all(kick[outside] == 0.0) and np.all(kick[~outside & (kx**2 + ky**2 > 0)] > 0.0) + + +# --------------------------------------------------- absorbers and band caps -- + + +def test_exp_absorbing_profile_matches_lpse_formula(): + cfg = _base_cfg(boundary_profile="exp", boundary_max_rate=150.0, boundary_lambda=7.0) + cfg["terms"]["epw"]["boundary"] = {"x": "absorbing", "y": "periodic"} + cfg = _finish_cfg(cfg) + grid = cfg["grid"] + rate = np.asarray(grid["absorbing_rate"]) + x = np.asarray(grid["x"]) + L = 0.4 + half = 0.5 * grid["dx"] + s = np.clip(np.maximum(grid["xmin"] + L + half - x, x - (grid["xmax"] - L - half)), 0.0, L) + expected = 150.0 * np.expm1(7.0 * s / L) / np.expm1(7.0) + np.testing.assert_allclose(rate[:, 0], expected, rtol=1e-12, atol=1e-12) + assert rate[0, 0] == pytest.approx(150.0) and rate[-1, 0] == pytest.approx(150.0) # edge cells: full rate + assert (rate[:, 0] > 0).sum() == 2 * round(L / grid["dx"]) # L / dx cells per side, as LPSE + assert np.all(rate[:, 0] == rate[:, -1]) + np.testing.assert_allclose(np.asarray(grid["absorbing_boundaries"]), np.exp(-rate * grid["dt"]), rtol=1e-14) + # interior is exactly undamped + interior = (x > grid["xmin"] + L) & (x < grid["xmax"] - L) + assert np.all(rate[interior, 0] == 0.0) + + # both axes absorbing: combined by max (LPSE), never by sum + cfg2 = _base_cfg(boundary_profile="exp", boundary_max_rate=150.0) + cfg2["terms"]["epw"]["boundary"] = {"x": "absorbing", "y": "absorbing"} + cfg2 = _finish_cfg(cfg2) + assert np.asarray(cfg2["grid"]["absorbing_rate"]).max() <= 150.0 * (1.0 + 1e-9) + + +def test_tanh_profile_is_the_default_and_unchanged(): + cfg = _base_cfg() + cfg["terms"]["epw"]["boundary"] = {"x": "absorbing", "y": "periodic"} + cfg = _finish_cfg(cfg) + from adept._base_ import get_envelope + + grid = cfg["grid"] + L = 0.4 + env = get_envelope(L / 5, L / 5, grid["xmin"] + L, grid["xmax"] - L, grid["x"])[:, None] + expected = np.exp(-1.0e4 * grid["dt"] * (1.0 - env * np.ones((1, grid["ny"])))) + np.testing.assert_allclose(np.asarray(grid["absorbing_boundaries"]), expected, rtol=1e-13) + + +def test_max_wavenumber_caps_the_retained_band(): + cfg = _base_cfg() + cfg["terms"]["epw"]["max_wavenumber"] = 1.5 + cfg = _finish_cfg(cfg) + derived = cfg["units"]["derived"] + k0 = derived["w0"] / derived["c"] + kx = np.asarray(cfg["grid"]["kx"])[:, None] + ky = np.asarray(cfg["grid"]["ky"])[None, :] + k = np.sqrt(kx**2 + ky**2) + band = np.asarray(cfg["grid"]["low_pass_filter_grid"]) + assert np.all(band[k >= 1.5 * k0] == 0.0) + assert np.any(band[k < 1.5 * k0] > 0.0) + + +def test_srs_k_filter_is_optional_and_masks_the_detuned_resonance(): + """The SRS source k-filter cuts at ``scale * k1(n_min)`` for zero detuning. At 0.2 nc + with envelope density 0.25 the resonant back-scattered Raman wavenumber lies above + that cutoff, so the default filter removes the SRS mode; ``srs_k_filter: false`` + (LPSE's ``lw.kFilter`` default) passes everything and a larger scale passes it too.""" + from adept._lpse2d.core.epw import SpectralEPWSolver + + def cfg_with(**source): + cfg = _base_cfg(envelope_density=0.25, box_density=0.2) + cfg["terms"]["epw"]["source"].update({"tpd": False, "srs": True, **source}) + return _finish_cfg(cfg) + + cfg = cfg_with(srs_k_filter=True) + d = cfg["units"]["derived"] + wp = d["w0"] * np.sqrt(0.2) + k_s = np.sqrt((d["w0"] - wp) ** 2 - wp**2) / d["c"] # resonant Raman wavenumber at 0.2 nc + kx = np.asarray(cfg["grid"]["kx"]) + ix = int(np.argmin(np.abs(kx - k_s))) + assert abs(kx[ix] - k_s) < 0.1 * k_s + + default = SpectralEPWSolver(cfg) + assert float(default.E1_filter[0, 0, 0]) == 1.0 + assert float(default.E1_filter[ix, 0, 0]) == 0.0, "resonant mode should be masked by the zero-detuning cutoff" + + off = SpectralEPWSolver(cfg_with(srs_k_filter=False)) + assert bool(jnp.all(off.E1_filter == 1.0)) + + wide = SpectralEPWSolver(cfg_with(srs_k_filter=True, srs_k_filter_scale=1.5)) + assert float(wide.E1_filter[ix, 0, 0]) == 1.0 + + +def test_light_absorber_has_its_own_peak_rate_with_the_exp_profile(): + """LPSE damps the light fields at 5e3/ps in the exp layers (200/ps for the EPW); at the EPW + rate a 3 um layer reflects most of a light wave and an injected pump builds a standing wave.""" + cfg = _base_cfg() + cfg["grid"].update({"boundary_profile": "exp", "boundary_max_rate": 200.0, "boundary_width": "3um", "xmax": "12um"}) + cfg = _finish_cfg(cfg) + dt = cfg["grid"]["dt"] + epw_rate = -np.log(np.asarray(cfg["grid"]["absorbing_boundaries"])) / dt + light_rate = -np.log(np.asarray(cfg["grid"]["light_absorbing_boundaries"])) / dt + np.testing.assert_allclose(epw_rate.max(), 200.0, rtol=1e-6) + np.testing.assert_allclose(light_rate.max(), 5.0e3, rtol=1e-6) + np.testing.assert_allclose(light_rate, epw_rate * 25.0, rtol=1e-6, atol=1e-9) + # the transit-integrated attenuation of light through the layer: strong with 5e3, weak with 200 + c = cfg["units"]["derived"]["c"] + x = np.asarray(cfg["grid"]["x"]) + layer = x < cfg["grid"]["xmin"] + 3.0 + assert np.exp(-np.trapezoid(light_rate[layer, 0], x[layer]) / c) < 0.01 + assert np.exp(-np.trapezoid(epw_rate[layer, 0], x[layer]) / c) > 0.5 + cfg2 = _base_cfg() + cfg2["grid"].update( + {"boundary_profile": "exp", "boundary_max_rate": 200.0, "boundary_width": "3um", "xmax": "12um"} + ) + cfg2["terms"]["light"] = {"boundary_max_rate": 1.0e3} + cfg2 = _finish_cfg(cfg2) + light_rate2 = -np.log(np.asarray(cfg2["grid"]["light_absorbing_boundaries"])) / dt + np.testing.assert_allclose(light_rate2.max(), 1.0e3, rtol=1e-6) diff --git a/tests/test_lpse2d/test_oblique_pump.py b/tests/test_lpse2d/test_oblique_pump.py new file mode 100644 index 00000000..b2e681d6 --- /dev/null +++ b/tests/test_lpse2d/test_oblique_pump.py @@ -0,0 +1,412 @@ +"""Oblique incidence of the pump (drivers.E0.angle): static pump and spectral injector.""" + +from copy import deepcopy + +import numpy as np +import pytest +import yaml +from jax import numpy as jnp + + +def _finish(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _cfg(angle, *, pump_depletion=False): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update({"xmax": "12.8um", "tmax": "10fs", "ymax": "6.4um", "ymin": "-6.4um", "dx": "0.1um"}) + cfg["terms"]["light"] = {"solver": "spectral", "pump_depletion": pump_depletion} + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["terms"]["epw"]["boundary"] = {"x": "absorbing" if pump_depletion else "periodic", "y": "periodic"} + cfg["drivers"]["E0"]["angle"] = angle + return _finish(cfg) + + +def _dominant_mode(field, cfg): + kx = np.asarray(cfg["grid"]["kx"]) + ky = np.asarray(cfg["grid"]["ky"]) + power = np.abs(np.fft.fft2(field[..., 0])) ** 2 + np.abs(np.fft.fft2(field[..., 1])) ** 2 + i, j = np.unravel_index(np.argmax(power), power.shape) + ex = np.fft.fft2(field[..., 0])[i, j] + ey = np.fft.fft2(field[..., 1])[i, j] + return kx[i], ky[j], ex, ey, power[i, j] / power.sum() + + +@pytest.mark.parametrize("angle", [0.0, 20.0]) +def test_static_pump_is_a_transverse_grid_mode_at_the_requested_angle(angle): + from adept._lpse2d.core.laser import Light + + cfg = _cfg(angle) + ny = cfg["grid"]["ny"] + light_wave = { + "delta_omega": jnp.array([0.0]), + "intensities": jnp.ones((1, ny)), + "phases": jnp.zeros((1, ny)), + } + E0 = np.asarray(Light(cfg).laser_update(0.0, None, light_wave)) + d = cfg["units"]["derived"] + k0 = d["w0"] / d["c"] * np.sqrt(1.0 - d["wp0"] ** 2 / d["w0"] ** 2) + kx, ky, ex, ey, fraction = _dominant_mode(E0, cfg) + assert fraction > 0.999 # one grid mode + dk = 2 * np.pi / (cfg["grid"]["nx"] * cfg["grid"]["dx"]) + dky = 2 * np.pi / (ny * cfg["grid"]["dy"]) + assert abs(kx - k0 * np.cos(np.deg2rad(angle))) <= 0.5 * dk + 1e-9 + assert abs(ky - k0 * np.sin(np.deg2rad(angle))) <= 0.5 * dky + 1e-9 + # transverse: k . E = 0 + assert abs(kx * ex + ky * ey) < 1e-9 * np.hypot(kx, ky) * np.hypot(abs(ex), abs(ey)) + # the same amplitude as at normal incidence (swelling at the uniform box density) + expected = d["E0_source"] * (1.0 - float(cfg["density"]["val"])) ** -0.25 + np.testing.assert_allclose(np.abs(E0[..., 0] + 1j * 0) ** 2 + np.abs(E0[..., 1]) ** 2, expected**2, rtol=1e-6) + + +def test_spectral_injector_launches_the_oblique_pump(): + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + from adept._lpse2d.modules.driver import UniformDriver + + cfg = _cfg(20.0, pump_depletion=True) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + light = SpectralCoupledLight(cfg) + _, args = UniformDriver(cfg)({}, {"drivers": {}}) + pump_args = args["drivers"]["E0"] + E0 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + E1 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + phi_k = jnp.zeros((nx, ny), dtype=jnp.complex128) + t, dt = 0.0, cfg["grid"]["dt"] + for _ in range(int(0.05 / dt)): # 50 fs: the front crosses ~15 um at c + E0, E1 = light(t, E0, E1, phi_k, pump_args, None) + t += dt + E0 = np.asarray(E0) + d = cfg["units"]["derived"] + k0 = d["w0"] / d["c"] * np.sqrt(1.0 - float(cfg["density"]["val"])) + kx, ky, ex, ey, _ = _dominant_mode(E0, cfg) + # a launched wave with a front and absorbers is broad in kx, but lives in one ky row + power_ky = np.sum(np.abs(np.fft.fft2(E0[..., 0])) ** 2 + np.abs(np.fft.fft2(E0[..., 1])) ** 2, axis=0) + assert power_ky[np.argmin(np.abs(np.asarray(cfg["grid"]["ky"]) - ky))] > 0.95 * power_ky.sum() + dky = 2 * np.pi / (ny * cfg["grid"]["dy"]) + assert abs(ky - k0 * np.sin(np.deg2rad(20.0))) <= 0.5 * dky + 1e-9 + assert abs(kx - np.sqrt(k0**2 - ky**2)) < 0.1 * k0 + # the absorber (an x-space envelope applied after the projection) mixes a little + assert abs(kx * ex + ky * ey) < 1e-2 * np.hypot(kx, ky) * np.hypot(abs(ex), abs(ey)) + assert np.max(np.abs(E0)) > 0.0 + + +def test_fd_injector_refuses_an_oblique_pump(): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg["terms"]["light"] = {"solver": "fd", "pump_depletion": True} + cfg["terms"]["epw"]["boundary"]["x"] = "absorbing" + cfg["drivers"]["E0"]["angle"] = 10.0 + with pytest.raises(ValueError, match="spectral"): + _finish(cfg) + + +# the beam lines of LPSE's test_010 deck (examples/testRuns/test_010/laser_include.txt) and the +# main-deck keys the translator needs, checked in so the test runs without an LPSE checkout +TEST_010_DECK = """grid.sizes = 20 10; +grid.nodes = 360 180; +laser.enable = true; +laser.solver = spectral; +lw.enable = true; +lw.spectral.dt = 0.002; +simulation.time.end = 3; +#include "laser_include.txt" +""" +TEST_010_LASER_INCLUDE = """laser.nBeams= 1; +laser.wavelength= 0.351; +laser.1.intensity = 1.5e+15; +laser.1.phase = 300.26; +laser.1.polarization = 0; +laser.1.frequencyShift = 0; +laser.1.group = 0; +laser.1.direction = 0.95783 0.28735 0; +laser.1.evolution.source = min.x; +laser.1.evolution.offset = 0 0 0; +laser.1.evolution.width = 0; +laser.1.evolution.sgOrder = 0; +""" + + +def test_translator_maps_the_test_010_beam_direction(tmp_path): + from adept._lpse2d.lpse_deck import parse_parms, translate_parms + + (tmp_path / "laser_include.txt").write_text(TEST_010_LASER_INCLUDE) + deck = tmp_path / "lpse.parms" + deck.write_text(TEST_010_DECK) + cfg, report = translate_parms(parse_parms(deck), experiment="x", run="test_010") + assert abs(cfg["drivers"]["E0"]["angle"] - np.degrees(np.arctan2(0.28735, 0.95783))) < 1e-3 + assert not any("direction" in u for u in report["unsupported"]) + + +def test_two_beams_launch_both_transverse_wavenumbers(): + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + from adept._lpse2d.modules.driver import UniformDriver + + cfg = _cfg(0.0, pump_depletion=True) + cfg = _finish(_beams_cfg(cfg, [{"intensity": 1.0, "angle": 20.0}, {"intensity": 1.0, "angle": -20.0}])) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + light = SpectralCoupledLight(cfg) + _, args = UniformDriver(cfg)({}, {"drivers": {}}) + pa = args["drivers"]["E0"] + E0 = jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + E1 = jnp.zeros_like(E0) + phi_k = jnp.zeros((nx, ny), dtype=jnp.complex128) + t, dt = 0.0, cfg["grid"]["dt"] + for _ in range(int(0.05 / dt)): + E0, E1 = light(t, E0, E1, phi_k, pa, None) + t += dt + E0 = np.asarray(E0) + power_ky = np.sum(np.abs(np.fft.fft2(E0[..., 0])) ** 2 + np.abs(np.fft.fft2(E0[..., 1])) ** 2, axis=0) + ky = np.asarray(cfg["grid"]["ky"]) + k0 = cfg["units"]["derived"]["w0"] / cfg["units"]["derived"]["c"] * np.sqrt(1.0 - float(cfg["density"]["val"])) + jp = int(np.argmin(np.abs(ky - k0 * np.sin(np.deg2rad(20.0))))) + jm = int(np.argmin(np.abs(ky + k0 * np.sin(np.deg2rad(20.0))))) + assert power_ky[jp] + power_ky[jm] > 0.95 * power_ky.sum() + np.testing.assert_allclose(power_ky[jp], power_ky[jm], rtol=0.05) + + +def _beams_cfg(cfg_finished, beams, **extra): + """Rebuild the raw config with beams/extras (the finished one already carries derived arrays).""" + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + raw = yaml.safe_load(fi) + raw = deepcopy(raw) + raw["grid"].update({"xmax": "12.8um", "tmax": "10fs", "ymax": "6.4um", "ymin": "-6.4um", "dx": "0.1um"}) + raw["terms"]["light"] = {"solver": "spectral", "pump_depletion": True} + raw["terms"]["epw"]["source"]["noise"] = False + raw["terms"]["epw"]["boundary"] = {"x": "absorbing", "y": "periodic"} + raw["drivers"]["E0"]["beams"] = beams + raw["drivers"]["E0"].update(extra) + return raw + + +def test_transverse_profile_kap_and_pulse_file(tmp_path): + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + from adept._lpse2d.modules.driver import UniformDriver + + pulse = tmp_path / "pulse.txt" + np.savetxt(pulse, np.array([[0.0, 0.25], [1.0, 0.25]])) + raw = _beams_cfg( + None, + [{"intensity": 1.0, "angle": 0.0}], + beam_width="2um", + beam_sg_order=4.0, + beam_offset="1um", + kap_bandwidth=0.02, + pulse_file=str(pulse), + ) + cfg = _finish(raw) + light = SpectralCoupledLight(cfg) + y = np.asarray(cfg["grid"]["y"]) + expected = np.exp(-(((y - 1.0) ** 2 / (2.0 * 2.0**2)) ** 2.0)) + np.testing.assert_allclose(np.asarray(light.beam_envelope_y), expected, rtol=1e-12) + _, args = UniformDriver(cfg)({}, {"drivers": {}}) + pa = args["drivers"]["E0"] + # the pulse table scales the source amplitude by 0.25 + s = np.asarray(light.calc_pump_source(0.05, pa)) + raw2 = _beams_cfg(None, [{"intensity": 1.0, "angle": 0.0}], beam_width="2um", beam_sg_order=4.0, beam_offset="1um") + light2 = SpectralCoupledLight(_finish(raw2)) + s2 = np.asarray(light2.calc_pump_source(0.05, pa)) + np.testing.assert_allclose(np.abs(s), 0.25 * np.abs(s2), rtol=1e-9) + # KAP: the phase is piecewise constant with correlation time 2 pi / (kap_bandwidth w0) + w0 = cfg["units"]["derived"]["w0"] + tau = 2.0 * np.pi / (0.02 * w0) + phases = [float(light.kap_phase(t, 0)) for t in np.arange(0.0, 20.0 * tau, tau / 4.0)] + jumps = np.sum(np.abs(np.diff(phases)) > 1e-9) + assert 15 <= jumps <= 21 # one jump per correlation time + assert float(light2.kap_phase(0.3, 0)) == 0.0 + + +def test_translator_maps_multi_beam_decks(): + from adept._lpse2d.lpse_deck import translate_parms + + parms = { + "grid.sizes": "20 5", + "grid.nodes": "201 51", + "laser.enable": "true", + "lw.enable": "true", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.nBeams": "2", + "laser.1.intensity": "1e15", + "laser.1.direction": "0.94 0.34 0", + "laser.1.phase": "0.5", + "laser.2.intensity": "3e15", + "laser.2.direction": "0.94 -0.34 0", + "laser.2.frequencyShift": "0.01", + "laser.1.width": "3", + "laser.1.sgOrder": "4", + "laser.1.offset": "0 1 0", + "laser.bandwidth.KAP.frequency": "0.005", + } + cfg, report = translate_parms(parms, experiment="x", run="y") + e0 = cfg["drivers"]["E0"] + assert len(e0["beams"]) == 2 and e0["beams"][0]["intensity"] == 1e15 and e0["beams"][1]["delta_omega"] == 0.01 + assert abs(e0["beams"][1]["angle"] + e0["beams"][0]["angle"]) < 1e-9 + assert e0["beams"][0]["phase"] == pytest.approx(np.deg2rad(0.5), rel=1e-12) # LPSE phases are degrees + assert e0["beam_width"] == "3.0um" and e0["beam_sg_order"] == 4.0 and e0["beam_offset"] == "1.0um" + assert e0["kap_bandwidth"] == 0.005 + assert not any("direction differs" in u for u in report["unsupported"]) + + +# ---- one normal-incidence beam keeps its frequency shift and phase (PR #372 review) ----------------- +# Before the fix the static pump's normal-incidence branch read only the color-level delta_omega / +# phases, the FD evolved pump ignored the beam list, and the translator dropped beams for one beam. +# A beam's shift and phase add to every color's, so a detuned single beam must equal the same shift +# and phase given at the color level -- an identity, checked to rtol 1e-12. + +SINGLE_BEAM = [{"intensity": 1.0, "angle": 0.0, "delta_omega": 0.01, "phase": 0.3}] + + +def _light_wave(ny, delta_omega=0.0, phase=0.0): + return { + "delta_omega": jnp.array([delta_omega]), + "intensities": jnp.ones((1, ny)), + "phases": jnp.full((1, ny), phase), + } + + +def test_static_pump_honors_a_single_beam_frequency_shift_and_phase(): + from adept._lpse2d.core.laser import Light + + cfg_plain = _cfg(0.0) + raw = _beams_cfg(cfg_plain, SINGLE_BEAM) + raw["terms"]["light"]["pump_depletion"] = False + raw["terms"]["epw"]["boundary"]["x"] = "periodic" + cfg_beam = _finish(raw) + ny = cfg_plain["grid"]["ny"] + beam, plain = Light(cfg_beam), Light(cfg_plain) + assert not beam.multi_beam # still the normal-incidence (speckle-capable) branch + for t in (0.0, 0.013, 0.05): + shifted = np.asarray(beam.laser_update(t, None, _light_wave(ny))) + reference = np.asarray(plain.laser_update(t, None, _light_wave(ny, 0.01, 0.3))) + np.testing.assert_allclose(shifted, reference, rtol=1e-12, atol=0) + unshifted = np.asarray(plain.laser_update(0.05, None, _light_wave(ny))) + assert np.max(np.abs(shifted - unshifted)) > 1e-3 * np.max(np.abs(unshifted)) + + +def test_fd_evolved_pump_honors_a_single_beam_frequency_shift_and_phase(): + from adept._lpse2d.core.light import CoupledLight + + def cfg_for(beams): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + raw = deepcopy(yaml.safe_load(fi)) + raw["grid"].update({"xmax": "12.8um", "tmax": "10fs", "ymax": "0.05um", "ymin": "-0.05um"}) + raw["terms"]["light"] = {"solver": "fd", "pump_depletion": True} + raw["terms"]["epw"]["source"]["noise"] = False + raw["terms"]["epw"]["boundary"]["x"] = "absorbing" + if beams is not None: + raw["drivers"]["E0"]["beams"] = beams + return _finish(raw) + + cfg_beam, cfg_plain = cfg_for(SINGLE_BEAM), cfg_for(None) + beam, plain = CoupledLight(cfg_beam), CoupledLight(cfg_plain) + ny = cfg_plain["grid"]["ny"] + pump = cfg_plain["drivers"]["E0"]["derived"] + for t in (0.013, 0.05): + rows_beam = beam.calc_pump_source(t, {**pump, **_light_wave(ny)}) + rows_plain = plain.calc_pump_source(t, {**pump, **_light_wave(ny, 0.01, 0.3)}) + for a, b in zip(rows_beam, rows_plain, strict=True): + np.testing.assert_allclose(np.asarray(a), np.asarray(b), rtol=1e-12, atol=0) + + +def test_translator_keeps_a_single_beam_frequency_shift_and_phase(): + from adept._lpse2d.lpse_deck import translate_parms + + base = { + "grid.sizes": "20 5", + "grid.nodes": "201 51", + "laser.enable": "true", + "lw.enable": "true", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.nBeams": "1", + "laser.1.intensity": "1e15", + } + cfg, _ = translate_parms({**base, "laser.1.frequencyShift": "0.01", "laser.1.phase": "90"}, run="y") + (b,) = cfg["drivers"]["E0"]["beams"] + assert b["delta_omega"] == 0.01 and b["phase"] == pytest.approx(np.pi / 2, rel=1e-12) + # LPSE's wavelengthShift (Angstrom) overrides frequencyShift: dW/W0 = -1e-4 dLambda / lambda_um + cfg, _ = translate_parms({**base, "laser.1.frequencyShift": "0.01", "laser.1.wavelengthShift": "10"}, run="y") + assert cfg["drivers"]["E0"]["beams"][0]["delta_omega"] == pytest.approx(-1.0e-3 / 0.351, rel=1e-12) + # no shift and no phase: the deck stays beam-free, as before + cfg, _ = translate_parms(base, run="y") + assert "beams" not in cfg["drivers"]["E0"] + + +# ---- the combined solver injects the same pump (PR #372 review) ----------------------------------- +# The combined solver (needed for simultaneous TPD + SRS) used to build only a normal-incidence +# scalar pump and silently ignore beams, profile, KAP and the pulse table. It now shares +# SmoothPumpInjector with the spectral light solver. Thresholds as in the spectral tests above. + + +def _combined_beams_cfg(beams, **extra): + raw = _beams_cfg(None, beams, **extra) + raw["terms"]["epw"]["solver"] = "combined" + # both sources on (the combined solver's rule); with no noise and E1 = 0 the unified source + # is zero, so the pump propagates alone + raw["terms"]["epw"]["source"].update({"tpd": True, "srs": True, "noise": False}) + return _finish(raw) + + +def test_combined_solver_pump_source_is_the_spectral_injector(tmp_path): + from adept._lpse2d.core.combined import CombinedSolver + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + from adept._lpse2d.modules.driver import UniformDriver + + pulse = tmp_path / "pulse.txt" + np.savetxt(pulse, np.array([[0.0, 0.25], [0.05, 1.0], [1.0, 1.0]])) + cfg = _combined_beams_cfg( + [{"intensity": 1.0, "angle": 20.0, "phase": 0.3}, {"intensity": 0.5, "angle": -15.0}], + beam_width="3um", + beam_sg_order=4.0, + kap_bandwidth=0.02, + pulse_file=str(pulse), + ) + combined = CombinedSolver(cfg) + spectral = SpectralCoupledLight(cfg) + _, args = UniformDriver(cfg)({}, {"drivers": {}}) + pa = args["drivers"]["E0"] + for t in (0.004, 0.03, 0.2): + np.testing.assert_array_equal( + np.asarray(combined.calc_pump_source(t, pa)), np.asarray(spectral.calc_pump_source(t, pa)) + ) + # the oblique beams put power in both components (polarization perpendicular to k) + s = np.asarray(combined.calc_pump_source(0.2, pa)) + assert np.abs(s[..., 0]).max() > 0.1 * np.abs(s[..., 1]).max() + + +def test_combined_solver_launches_both_oblique_beams(): + from adept._lpse2d.core.combined import CombinedSolver + from adept._lpse2d.modules.driver import UniformDriver + + cfg = _combined_beams_cfg([{"intensity": 1.0, "angle": 20.0}, {"intensity": 1.0, "angle": -20.0}]) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + solver = CombinedSolver(cfg) + _, args = UniformDriver(cfg)({}, {"drivers": {}}) + driver_args = {"E0": args["drivers"]["E0"]} + y = {"E0": jnp.zeros((nx, ny, 2), dtype=jnp.complex128), "E1": jnp.zeros((nx, ny, 2), dtype=jnp.complex128)} + t, dt = 0.0, cfg["grid"]["dt"] + for _ in range(int(0.05 / dt)): + E0, E1, _ = solver(t, y, driver_args) + y = {"E0": E0, "E1": E1} + t += dt + E0 = np.asarray(y["E0"]) + power_ky = np.sum(np.abs(np.fft.fft2(E0[..., 0])) ** 2 + np.abs(np.fft.fft2(E0[..., 1])) ** 2, axis=0) + ky = np.asarray(cfg["grid"]["ky"]) + k0 = cfg["units"]["derived"]["w0"] / cfg["units"]["derived"]["c"] * np.sqrt(1.0 - float(cfg["density"]["val"])) + jp = int(np.argmin(np.abs(ky - k0 * np.sin(np.deg2rad(20.0))))) + jm = int(np.argmin(np.abs(ky + k0 * np.sin(np.deg2rad(20.0))))) + assert power_ky[jp] + power_ky[jm] > 0.95 * power_ky.sum() + np.testing.assert_allclose(power_ky[jp], power_ky[jm], rtol=0.05) diff --git a/tests/test_lpse2d/test_restart.py b/tests/test_lpse2d/test_restart.py new file mode 100644 index 00000000..33d507d4 --- /dev/null +++ b/tests/test_lpse2d/test_restart.py @@ -0,0 +1,55 @@ +"""Checkpoint / restart (LPSE --restart): a run split in two reproduces the unbroken run.""" + +from copy import deepcopy + +import numpy as np +import yaml + + +def _cfg(tmax_fs, checkpoint=None, restart=None, run="restart-test"): + with open("tests/test_lpse2d/configs/epw.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update( + {"tmax": f"{tmax_fs}fs", "dt": "1fs", "xmax": "8um", "dx": "0.1um", "ymax": "0.2um", "ymin": "-0.2um"} + ) + cfg["save"]["fields"]["t"].update({"tmin": "0fs", "tmax": f"{tmax_fs}fs", "dt": "5fs"}) + cfg["terms"]["epw"]["source"]["noise"] = True + cfg["terms"]["epw"]["source"]["noise_seed"] = 7 + cfg["mlflow"]["run"] = run + if checkpoint is not None: + cfg["save"]["checkpoint"] = checkpoint + if restart is not None: + cfg["restart"] = {"file": restart} + return cfg + + +def _run(cfg): + from adept import ergoExo + + exo = ergoExo() + modules = exo.setup(cfg) + _, ppo, _ = exo(modules) + return ppo + + +def test_split_run_matches_unbroken_run(tmp_path): + ckpt = str(tmp_path / "ckpt.npz") + _run(_cfg(10, checkpoint=ckpt, run="restart-first")) + saved = np.load(ckpt) + t_ckpt = float(saved["t"]) + assert 0.0099 < t_ckpt < 0.0121 and "epw" in saved.files # the solver's end time (tmax rounded to steps) + + resumed = _run(_cfg(20, restart=ckpt, run="restart-second")) + unbroken = _run(_cfg(20, run="restart-unbroken")) + + t_res = resumed["series"]["t (ps)"].values + assert t_res[0] >= t_ckpt - 1e-12 # the series starts at the restart time + for name in ("phi", "ex"): + a = resumed["x"][name].isel({"t (ps)": -1}).values + b = unbroken["x"][name].isel({"t (ps)": -1}).values + np.testing.assert_allclose(a, b, rtol=1e-9, atol=1e-12 * np.max(np.abs(b))) + e_res = float(resumed["series"]["epw_energy"].values[-1]) + e_unb = float(unbroken["series"]["epw_energy"].values[-1]) + assert e_unb > 0.0 + np.testing.assert_allclose(e_res, e_unb, rtol=1e-9) diff --git a/tests/test_lpse2d/test_saved_epw_fields.py b/tests/test_lpse2d/test_saved_epw_fields.py new file mode 100644 index 00000000..ee7d34fc --- /dev/null +++ b/tests/test_lpse2d/test_saved_epw_fields.py @@ -0,0 +1,216 @@ +"""Normalization of the EPW fields in the full-grid field save (``fields.xr`` / ``k-fields.xr``). + +The solver state is the k-space potential phi_k (numpy FFT convention, DC at [0, 0]). The +solver's own field is E = ifft2(-i k phi_k) (``SpectralEPWSolver.phi_k_to_e_fields``), and the +``series.xr`` quantities ``e_sq`` / ``epw_energy`` integrate that E. ``fields.xr`` stores +phi = ifft2(phi_k), so its ``ex`` / ``ey`` must be E = -grad(phi) of the saved phi with no further +normalization, and ``k-fields.xr`` stores phi_k and -i k phi_k. + +Tolerances (fixed before the tests were first run). Every comparison is between two evaluations of +the same quantity in the saved dtype (complex128 here) that differ only by round-off: + +* An FFT of size N has a relative L2 error of order log2(N) eps (pocketfft and XLA alike); here + N = nx ny = 4096 (2D) or 128 (quasi-1D), log2(N) <= 12. A forward + inverse pair (the spectral + derivative of the saved phi) doubles that: ~24 eps. +* The analytic single mode exp(2 pi i (mx i / nx + my j / ny)) costs ~2 pi mx eps <= 20 eps per + element for the phase rounding, plus the inverse FFT. +* The integrated e_sq: |E|^2 doubles the relative error of E, the sums add ~log2(N) eps. + +All of these are below ~60 eps; ROUNDOFF_EPS = 200 eps of the reference's scale (max |.| for the +analytic mode, L2 norm for the random field, the value itself for e_sq) is the bound. The defect this +guards against is a factor 4 / (nx ny) on the saved ex / ey: a relative error of 1 - 4 / (nx ny) +(0.999 in 2D, 0.97 in quasi-1D), about 1e13 times the bound. +""" + +import jax.numpy as jnp +import numpy as np +import pytest +import xarray as xr +import yaml + +GRID_DX_UM = 0.05 +X_LENGTH_UM = 6.4 # nx = 128 +Y_HALF_WIDTH_UM = {"2d": 0.8, "1d": 0.02} # "1d" collapses the box to ny = 1 +EXPECTED_SHAPE = {"2d": (128, 32), "1d": (128, 1)} +MODE = {"2d": (3, 1), "1d": (3, 0)} # (mx, my) of the analytic single-mode slice +AMPLITUDE = 2.0e-3 # |phi| of the single mode, in the solver's units +RANDOM_SEED = 20261020 # any seed: the assertions are identities that hold for every phi_k +ROUNDOFF_EPS = 200.0 +KX = r"kx ($kc\omega_0^{-1}$)" +KY = r"ky ($kc\omega_0^{-1}$)" + + +def _prepared_cfg(geometry: str) -> dict: + """A small TPD box taken through the setup steps up to the save functions (no MLflow, no run).""" + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_save_quantities, + get_solver_quantities, + write_units, + ) + + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + half = Y_HALF_WIDTH_UM[geometry] + cfg["density"] = {"basis": "uniform", "val": 0.24} + cfg["grid"].update( + { + "boundary_width": "1um", + "dt": "1fs", + "dx": f"{GRID_DX_UM}um", + "xmax": f"{X_LENGTH_UM}um", + "tmax": "6fs", + "ymax": f"{half}um", + "ymin": f"{-half}um", + } + ) + # full-grid save: no save.fields.x / save.fields.y + cfg["save"]["fields"] = {"t": {"tmin": "0fs", "tmax": "6fs", "dt": "3fs"}} + cfg["drivers"]["E0"]["params"] = {"phases": {"seed": 42}} + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return get_save_quantities(cfg) + + +def _phi_k_slices(cfg: dict, geometry: str) -> np.ndarray: + """(2, nx, ny) phi_k: slice 0 a single grid mode, slice 1 a random band-limited field.""" + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + mx, my = MODE[geometry] + phi_k = np.zeros((2, nx, ny), dtype=np.complex128) + # phi_k[mx, my] = A nx ny <=> phi(x_i, y_j) = A exp(2 pi i (mx i / nx + my j / ny)) + phi_k[0, mx, my] = AMPLITUDE * nx * ny + + rng = np.random.default_rng(RANDOM_SEED) + mode_x = np.abs(np.fft.fftfreq(nx, d=1.0 / nx)) + mode_y = np.abs(np.fft.fftfreq(ny, d=1.0 / ny)) + band = (mode_x[:, None] <= nx // 4) & (mode_y[None, :] <= max(ny // 4, 0)) + band[0, 0] = False # the solver zeroes k = 0 + coeffs = rng.standard_normal((nx, ny)) + 1j * rng.standard_normal((nx, ny)) + phi_k[1] = np.where(band, coeffs, 0.0) * AMPLITUDE * np.sqrt(nx * ny) + return phi_k + + +def _saved_fields(geometry: str, td) -> tuple[dict, np.ndarray, xr.Dataset, xr.Dataset]: + """Run the full-grid save path (save function + make_field_xarrays) on a two-slice state and + read fields.xr / k-fields.xr back from disk.""" + from adept._lpse2d.helpers import make_field_xarrays + + cfg = _prepared_cfg(geometry) + assert "x" not in cfg["save"]["fields"] + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + assert (nx, ny) == EXPECTED_SHAPE[geometry] + + phi_k = _phi_k_slices(cfg, geometry) + zeros = np.zeros((nx, ny, 2), dtype=np.complex128) + saved = [] + for phi_k_t in phi_k: + state = {"epw": jnp.asarray(phi_k_t).view(jnp.float64), "E0": jnp.asarray(zeros).view(jnp.float64)} + state["E1"] = state["E0"] + saved.append(cfg["save"]["fields"]["func"](0.0, state, None)) + stacked = {k: np.stack([np.asarray(s[k]) for s in saved]) for k in saved[0]} + + (td / "binary").mkdir() + make_field_xarrays(cfg, np.array([0.0, 0.003]), stacked, str(td)) + fields = xr.load_dataset(td / "binary" / "fields.xr", engine="h5netcdf") + kfields = xr.load_dataset(td / "binary" / "k-fields.xr", engine="h5netcdf") + return cfg, phi_k, fields, kfields + + +def _tol(dtype) -> float: + return ROUNDOFF_EPS * np.finfo(dtype).eps + + +def _rel_l2(got: np.ndarray, ref: np.ndarray) -> float: + return float(np.linalg.norm(got - ref) / np.linalg.norm(ref)) + + +@pytest.mark.parametrize("geometry", ["2d", "1d"]) +def test_saved_epw_fields_single_mode(geometry, tmp_path): + """Single grid mode phi_k[mx, my] = A nx ny, i.e. phi = A exp(i (kx (x - x0) + ky (y - y0))). + + fields.xr: phi = phi(x), ex = -i kx phi(x), ey = -i ky phi(x) (ey = 0 in quasi-1D); + k-fields.xr (fftshifted): phi_k, -i kx phi_k, -i ky phi_k. + Bound: max |got - expected| <= ROUNDOFF_EPS eps |k| A in x-space (|k| = sqrt(kx^2 + ky^2)), + <= ROUNDOFF_EPS eps |k| A nx ny in k-space. + """ + cfg, _, fields, kfields = _saved_fields(geometry, tmp_path) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + mx, my = MODE[geometry] + kx, ky = cfg["grid"]["kx"][mx], cfg["grid"]["ky"][my] + k_abs = np.hypot(kx, ky) + np.testing.assert_allclose(kx, 2.0 * np.pi * mx / (cfg["grid"]["xmax"] - cfg["grid"]["xmin"]), rtol=1e-14) + + i, j = np.arange(nx)[:, None], np.arange(ny)[None, :] + phi_x = AMPLITUDE * np.exp(2j * np.pi * (mx * i / nx + my * j / ny)) + expected_x = {"phi": phi_x, "ex": -1j * kx * phi_x, "ey": -1j * ky * phi_x} + for name, expected in expected_x.items(): + got = fields[name].values[0] + assert got.dtype == np.complex128, name + scale = AMPLITUDE * (1.0 if name == "phi" else k_abs) + err = np.max(np.abs(got - expected)) + assert err <= _tol(got.dtype) * scale, f"fields.xr {name}: max err {err:.3e}, scale {scale:.3e}" + + expected_k = {"phi": 1.0, "ex": -1j * kx, "ey": -1j * ky} + for name, factor in expected_k.items(): + got_k = np.fft.ifftshift(kfields[name].values[0]) + expected = np.zeros((nx, ny), dtype=np.complex128) + expected[mx, my] = factor * AMPLITUDE * nx * ny + scale = AMPLITUDE * nx * ny * (1.0 if name == "phi" else k_abs) + err = np.max(np.abs(got_k - expected)) + assert err <= _tol(got_k.dtype) * scale, f"k-fields.xr {name}: max err {err:.3e}, scale {scale:.3e}" + + +@pytest.mark.parametrize("geometry", ["2d", "1d"]) +def test_saved_epw_fields_are_the_solver_field(geometry, tmp_path): + """Both time slices (single mode and random band-limited phi_k): + + * saved ex, ey = -d(phi)/dx, -d(phi)/dy of the saved phi, differentiated spectrally on the grid; + * saved ex, ey = SpectralEPWSolver.phi_k_to_e_fields(phi_k), the field the solver evolves; + * saved ex, ey = ifft2 of the k-fields.xr ex, ey (the two datasets agree); + * dx dy sum(|ex|^2 + |ey|^2) of the saved fields = the e_sq series of the default save. + Bound: relative L2 error (relative error for e_sq) <= ROUNDOFF_EPS eps. + """ + from adept._lpse2d.core.epw import SpectralEPWSolver + + cfg, phi_k, fields, kfields = _saved_fields(geometry, tmp_path) + kx = cfg["grid"]["kx"][:, None] + ky = cfg["grid"]["ky"][None, :] + solver = SpectralEPWSolver(cfg) + series_func = cfg["save"]["default"]["func"] + + for it in range(phi_k.shape[0]): + ex = fields["ex"].values[it] + ey = fields["ey"].values[it] + tol = _tol(ex.dtype) + + phi_saved_k = np.fft.fft2(fields["phi"].values[it]) + grad_ex = np.fft.ifft2(-1j * kx * phi_saved_k) + grad_ey = np.fft.ifft2(-1j * ky * phi_saved_k) + solver_ex, solver_ey = (np.asarray(e) for e in solver.phi_k_to_e_fields(jnp.asarray(phi_k[it]))) + k_ex = np.fft.ifft2(np.fft.ifftshift(kfields["ex"].values[it])) + k_ey = np.fft.ifft2(np.fft.ifftshift(kfields["ey"].values[it])) + + references = {"-grad(saved phi)": (grad_ex, grad_ey), "solver E": (solver_ex, solver_ey)} + references["ifft2(k-fields)"] = (k_ex, k_ey) + for label, (ref_x, ref_y) in references.items(): + err_x = _rel_l2(ex, ref_x) + assert err_x <= tol, f"slice {it}: ex vs {label}: relative L2 err {err_x:.3e}" + if np.any(ref_y != 0): + err_y = _rel_l2(ey, ref_y) + assert err_y <= tol, f"slice {it}: ey vs {label}: relative L2 err {err_y:.3e}" + else: # quasi-1D: ky = 0 everywhere + assert np.all(ey == 0), f"slice {it}: ey vs {label}: expected exactly 0" + + state = { + "epw": jnp.asarray(phi_k[it]).view(jnp.float64), + "E0": jnp.zeros((*phi_k.shape[1:], 2), dtype=jnp.complex128).view(jnp.float64), + } + state["E1"] = state["E0"] + e_sq_series = float(series_func(0.0, state, None)["e_sq"]) + e_sq_saved = float(np.sum(np.abs(ex) ** 2 + np.abs(ey) ** 2) * cfg["grid"]["dx"] * cfg["grid"]["dy"]) + err = abs(e_sq_saved - e_sq_series) / e_sq_series + assert err <= tol, f"slice {it}: e_sq from fields.xr {e_sq_saved:.6e} vs series {e_sq_series:.6e}" diff --git a/tests/test_lpse2d/test_spectral_light.py b/tests/test_lpse2d/test_spectral_light.py new file mode 100644 index 00000000..c88cb349 --- /dev/null +++ b/tests/test_lpse2d/test_spectral_light.py @@ -0,0 +1,334 @@ +"""LPSE-parity tests for the spectral light propagator (terms.light.solver: spectral) +and the collisional light absorption (terms.light.absorption).""" + +from copy import deepcopy + +import numpy as np +import pytest +import yaml +from jax import numpy as jnp + + +def _finish(cfg): + from adept._lpse2d.helpers import ( + get_density_profile, + get_derived_quantities, + get_solver_quantities, + write_units, + ) + + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _srs_cfg( + solver="spectral", + *, + seed=False, + pump_depletion=False, + absorption=False, + xmax="20um", + periodic=False, + epw_max_wavenumber=None, + light_max_wavenumber=None, +): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update({"xmax": xmax, "tmax": "10fs", "ymax": "0.02um", "ymin": "-0.02um"}) + cfg["terms"]["light"] = {"solver": solver, "pump_depletion": pump_depletion, "absorption": absorption} + if pump_depletion: + cfg["terms"]["light"]["coupling"] = "rotation" + if epw_max_wavenumber is not None: + cfg["terms"]["epw"]["max_wavenumber"] = epw_max_wavenumber + if light_max_wavenumber is not None: + cfg["terms"]["light"]["max_wavenumber"] = light_max_wavenumber + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["terms"]["epw"]["boundary"]["x"] = "periodic" if periodic else "absorbing" + if seed: + # the seed's default offset (1.6 boundary widths) sits inside the absorber skirt by + # design of the FD path; the smooth injector must be clear of it + cfg["drivers"]["E1"] = { + "intensity": "1.0e+12W/cm^2", + "delta_omega": 0.0, + "turn_on_time": "10fs", + "offset": "7um", + } + return _finish(cfg) + + +def _pump_args(cfg): + ny = cfg["grid"]["ny"] + return { + **cfg["drivers"]["E0"]["derived"], + "delta_omega": jnp.zeros(1), + "intensities": jnp.ones((1, ny)), + "phases": jnp.zeros((1, ny)), + } + + +def _plane_wave(cfg, w, sign, amplitude=1.0): + """E_y = A exp(sign i k x) with k the grid-snapped local light wavenumber for carrier w.""" + derived = cfg["units"]["derived"] + n = float(np.asarray(cfg["grid"]["background_density"])[0, 0]) + k = np.sqrt(w**2 - n * derived["w0"] ** 2) / derived["c"] + dk = 2.0 * np.pi / (cfg["grid"]["nx"] * cfg["grid"]["dx"]) + k = np.round(k / dk) * dk + x = np.asarray(cfg["grid"]["x"]) + field = np.zeros((cfg["grid"]["nx"], cfg["grid"]["ny"], 2), dtype=np.complex128) + field[..., 1] = amplitude * np.exp(sign * 1j * k * x)[:, None] + return jnp.asarray(field), k + + +def test_spectral_propagator_has_no_grid_dispersion(): + """A resonant plane wave e^{-i k1 x} advances by exactly the analytic phase + dt [w1/2 (1 - n w0^2/w1^2) - c^2 k^2/(2 w1)] per step under the spectral solver; the FD + solver's grid dispersion gives a different phase at 7 cells per wavelength.""" + from adept._lpse2d.core.raman import RamanLight + from adept._lpse2d.core.spectral_light import SpectralRamanLight + + cfg_s = _srs_cfg("spectral", periodic=True, xmax="8um") + cfg_f = _srs_cfg("fd", periodic=True, xmax="8um") + spectral, fd = SpectralRamanLight(cfg_s), RamanLight(cfg_f) + derived = cfg_s["units"]["derived"] + E1, k = _plane_wave(cfg_s, derived["w1"], -1) + n_steps = 20 + E0_fn = lambda t: jnp.zeros_like(E1) + phi_k = jnp.zeros(E1.shape[:2], dtype=jnp.complex128) + out_s, out_f = E1, E1 + for i in range(n_steps): + out_s = spectral(i * spectral.dt, out_s, E0_fn, phi_k, None) + out_f = fd(i * fd.dt, out_f, E0_fn, phi_k, None) + n = float(np.asarray(cfg_s["grid"]["background_density"])[0, 0]) + w1, c = derived["w1"], derived["c"] + phase = n_steps * spectral.dt * (w1 / 2.0 * (1.0 - n * derived["w0"] ** 2 / w1**2) - c**2 * k**2 / (2.0 * w1)) + expected = np.asarray(E1) * np.exp(1j * phase) + np.testing.assert_allclose(np.asarray(out_s), expected, rtol=1e-10, atol=1e-12) + assert not np.allclose(np.asarray(out_f), expected, rtol=1e-4, atol=1e-6) + # the spectral solver has no CFL limit: it sub-cycles only for its injectors (one cell per + # sub-step), fewer steps than the FD scheme's stability bound + assert 1 <= cfg_s["grid"]["light_substeps"] < cfg_f["grid"]["light_substeps"] + + +def _fill_box(cfg, n_steps, pump_depletion): + from adept._lpse2d.core.vector_field import SplitStep + + step = SplitStep(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + state = { + "epw": jnp.zeros((nx, ny), dtype=jnp.complex128).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), + } + args = {"drivers": {"E0": _pump_args(cfg)} if pump_depletion else {}} + if "E1" in cfg["drivers"]: + args["drivers"]["E1"] = cfg["drivers"]["E1"]["derived"] + dt = cfg["grid"]["dt"] + for i in range(n_steps): + state = step(jnp.asarray(i * dt), dict(state), args) + return step, state + + +def test_spectral_seed_injector_launches_the_exact_amplitude_leftward(): + """The smooth injector launches |E1| = E1_source / eps1^(1/4) with no sinc deficit, as a + leftward wave at the local k1 (the FD two-point injector keeps its MATLAB calibration; + the width default is half a local wavelength, LPSE's injectorWidth).""" + cfg = _srs_cfg("spectral", seed=True) + del cfg["drivers"]["E0"] + step, state = _fill_box(cfg, 180, pump_depletion=False) + e1y = np.asarray(state["E1"].view(jnp.complex128))[:, 0, 1] + x = np.asarray(cfg["grid"]["x"]) + derived = cfg["units"]["derived"] + n = cfg["units"]["envelope density"] + eps1 = 1.0 - n * (derived["w0"] / derived["w1"]) ** 2 + k1 = np.sqrt(derived["w1"] ** 2 - derived["wp0"] ** 2) / derived["c"] + bulk = slice(int(np.argmin(np.abs(x - 5.0))), int(np.argmin(np.abs(x - 10.0)))) + expected = cfg["drivers"]["E1"]["derived"]["amplitude"] / eps1**0.25 + np.testing.assert_allclose(np.mean(np.abs(e1y[bulk])), expected, rtol=0.01) + assert np.std(np.abs(e1y[bulk])) < 0.02 * expected # a clean traveling wave, no standing-wave ripple + k_measured = np.polyfit(x[bulk], np.unwrap(np.angle(e1y[bulk])), 1)[0] + assert k_measured < 0 + np.testing.assert_allclose(abs(k_measured), k1, rtol=0.01) + # negligible leakage to the right of the injector (rightward wave) + # beyond the injector's own Gaussian tail (3 sigma ~ 1.6 um) only leakage remains + right = x > x[step.raman.i1] + 2.5 + assert np.abs(e1y[right]).max() < 1e-2 * expected + + +def test_spectral_pump_injector_launches_the_exact_amplitude_rightward(): + cfg = _srs_cfg("spectral", pump_depletion=True) + step, state = _fill_box(cfg, 180, pump_depletion=True) + e0y = np.asarray(state["E0"].view(jnp.complex128))[:, 0, 1] + x = np.asarray(cfg["grid"]["x"]) + derived = cfg["units"]["derived"] + n = cfg["units"]["envelope density"] + bulk = slice(int(np.argmin(np.abs(x - 8.0))), int(np.argmin(np.abs(x - 14.0)))) + expected = derived["E0_source"] * (1.0 - n) ** -0.25 + np.testing.assert_allclose(np.mean(np.abs(e0y[bulk])), expected, rtol=0.01) + k0 = np.sqrt(derived["w0"] ** 2 - derived["wp0"] ** 2) / derived["c"] + k_measured = np.polyfit(x[bulk], np.unwrap(np.angle(e0y[bulk])), 1)[0] + np.testing.assert_allclose(k_measured, k0, rtol=0.01) + left = x < x[step.coupled_light.i0] - 1.0 + assert np.abs(e0y[left]).max() < 1e-2 * expected + # no plasma wave: the Raman field stays identically zero + assert float(jnp.max(jnp.abs(state["E1"]))) == 0.0 + + +@pytest.mark.parametrize("solver", ["spectral", "fd"]) +def test_collisional_absorption_decays_the_light(solver): + """With terms.light.absorption = nu (1/ps at nc), |E1| decays as exp(-nu (n/nc1)^2 t).""" + from adept._lpse2d.core.raman import RamanLight + from adept._lpse2d.core.spectral_light import SpectralRamanLight + + nu = 40.0 + cfg = _srs_cfg(solver, periodic=True, xmax="8um", absorption=nu) + light = (SpectralRamanLight if solver == "spectral" else RamanLight)(cfg) + derived = cfg["units"]["derived"] + assert light.absorption_rate0 == nu + assert light.absorption_rate1 == pytest.approx(nu * (derived["w1"] / derived["w0"]) ** 2) + E1, _ = _plane_wave(cfg, derived["w1"], -1) + E0_fn = lambda t: jnp.zeros_like(E1) + phi_k = jnp.zeros(E1.shape[:2], dtype=jnp.complex128) + out = E1 + n_steps = 10 + for i in range(n_steps): + out = light(i * light.dt, out, E0_fn, phi_k, None) + n = cfg["units"]["envelope density"] + n_over_nc1 = n * (derived["w0"] / derived["w1"]) ** 2 + expected_decay = np.exp(-light.absorption_rate1 * n_over_nc1**2 * n_steps * light.dt) + ratio = np.abs(np.asarray(out)[..., 1]).mean() / np.abs(np.asarray(E1)[..., 1]).mean() + assert ratio == pytest.approx(expected_decay, rel=1e-3 if solver == "spectral" else 2e-2) + assert 0.3 < expected_decay < 0.95 + + +def test_nrl_absorption_rate(): + from adept._lpse2d.core.raman import light_absorption_rates + + cfg = _srs_cfg("spectral", periodic=True, xmax="8um", absorption=True) + rate0, rate1 = light_absorption_rates(cfg) + te, z, lam = 2.0, 1.0, 0.351 + log_lambda = 6.68 + np.log(lam * te) # Te > 0.01 Z^2 + expected0 = 5.11e10 * z * log_lambda / (lam**2 * te**1.5) * 1e-12 + assert rate0 == pytest.approx(expected0, rel=1e-12) + derived = cfg["units"]["derived"] + lam1 = lam * derived["w0"] / derived["w1"] + expected1 = 5.11e10 * z * (6.68 + np.log(lam1 * te)) / (lam1**2 * te**1.5) * 1e-12 + assert rate1 == pytest.approx(expected1, rel=1e-12) + assert 0.1 < rate0 < 100.0 + assert light_absorption_rates(_srs_cfg("spectral", periodic=True, xmax="8um")) == (None, None) + + +def test_spectral_coupled_solver_conserves_light_action_in_the_exchange(): + """With the propagation trivial (phi = 0 apart from the exchange test) the Strang-split + rotation is action conserving; here we check the coupled spectral step runs, keeps + finite fields, and that with SRS the exchange transfers energy between E0 and E1.""" + from adept._lpse2d.core.spectral_light import SpectralCoupledLight + + cfg = _srs_cfg("spectral", pump_depletion=True, xmax="8um") + light = SpectralCoupledLight(cfg) + rng = np.random.default_rng(3) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + E0, _ = _plane_wave(cfg, cfg["units"]["derived"]["w0"], +1, amplitude=cfg["units"]["derived"]["E0_source"]) + E1 = jnp.zeros_like(E0) + phi = (rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) * 1e-3 + phi = jnp.asarray(phi * np.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"])) + E0n, E1n = light(0.0, E0, E1, phi, _pump_args(cfg), None) + assert bool(jnp.all(jnp.isfinite(E0n))) and bool(jnp.all(jnp.isfinite(E1n))) + assert float(jnp.max(jnp.abs(E1n))) > 0.0 + # without a plasma wave the Raman field is not seeded and the pump only propagates + E0z, E1z = light(0.0, E0, E1, jnp.zeros_like(phi), _pump_args(cfg), None) + assert float(jnp.max(jnp.abs(E1z))) == 0.0 + assert bool(jnp.all(jnp.isfinite(E0z))) + + +def test_light_solver_option_is_validated(): + import pydantic + + with pytest.raises(pydantic.ValidationError): + _srs_cfg("implicit", periodic=True, xmax="8um") + + +def test_light_band_is_independent_of_the_epw_cutoff(): + """terms.epw.max_wavenumber caps only the EPW band (LPSE lw.maxWavenumber); the separate spectral + light solvers keep the anti-aliased band and their own terms.light.max_wavenumber (LPSE + {laser|raman}.maxWavenumber; LightSolver::evolveSpectral takes the EPW cap only in combined mode, + LightSolver.cpp:3689). PR #387 review: before the fix an EPW cap also zeroed the light modes above + it. 0/1 masks, compared exactly.""" + from adept._lpse2d.core.spectral_light import SpectralCoupledLight, SpectralRamanLight + + for cls, pump_depletion in ((SpectralRamanLight, False), (SpectralCoupledLight, True)): + base = _srs_cfg(pump_depletion=pump_depletion) + capped = _srs_cfg(pump_depletion=pump_depletion, epw_max_wavenumber=0.5) + epw_base = np.asarray(base["grid"]["low_pass_filter_grid"]) + epw_capped = np.asarray(capped["grid"]["low_pass_filter_grid"]) + assert epw_capped.sum() < epw_base.sum() # the EPW band did shrink + light_base = np.asarray(cls(base).light_band) + np.testing.assert_array_equal(np.asarray(cls(capped).light_band), light_base) + assert light_base.sum() > epw_capped.sum() + # the light's own cap still acts + light_capped = np.asarray(cls(_srs_cfg(pump_depletion=pump_depletion, light_max_wavenumber=0.5)).light_band) + np.testing.assert_array_equal(light_capped > 0, (light_base > 0) & (epw_capped > 0)) + + +def _srs_cfg_2d(solver="spectral", transverse_source=True): + with open("tests/test_lpse2d/configs/srs.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["grid"].update({"xmax": "6.4um", "tmax": "10fs", "ymax": "3.2um", "ymin": "-3.2um", "dx": "0.1um"}) + cfg["terms"]["light"] = {"solver": solver, "transverse_source": transverse_source} + cfg["terms"]["epw"]["source"]["noise"] = False + cfg["terms"]["epw"]["boundary"] = {"x": "periodic", "y": "periodic"} + return _finish(cfg) + + +def test_longitudinal_light_is_kept_only_inside_the_band(): + """LPSE zeroes every component of the light field outside the retained band; adept's + propagator does the same to the longitudinal part it does not move.""" + from adept._lpse2d.core.spectral_light import SpectralRamanLight, transverse_propagate + + cfg = _srs_cfg_2d() + solver = SpectralRamanLight(cfg) + ky = np.asarray(cfg["grid"]["ky"]) + y = np.asarray(cfg["grid"]["y"]) + band = np.asarray(solver.light_band) + j_in = int(np.flatnonzero(band[0] > 0)[1]) + j_out = int(np.flatnonzero(band[0] == 0)[0]) + ones = jnp.ones_like(solver.propagator1) + for j, kept in ((j_in, True), (j_out, False)): + wave = np.exp(1j * ky[j] * y)[None, :] * np.ones((len(cfg["grid"]["x"]), 1)) + field = jnp.stack([jnp.zeros_like(wave), jnp.asarray(wave)], axis=-1) # E || y, k || y: longitudinal + out = transverse_propagate(field, solver.kx_arr, solver.ky_arr, solver.one_over_k_sq, ones, solver.light_band) + if kept: + np.testing.assert_allclose(np.asarray(out), np.asarray(field), atol=1e-12) + else: + assert float(jnp.max(jnp.abs(out))) < 1e-12 + dropped = transverse_propagate( + field, solver.kx_arr, solver.ky_arr, solver.one_over_k_sq, ones, solver.light_band, keep_longitudinal=False + ) + assert float(jnp.max(jnp.abs(dropped))) < 1e-12 + + +@pytest.mark.parametrize("solver", ["spectral", "fd"]) +def test_transverse_source_removes_the_longitudinal_srs_source(solver): + """An EPW with k along y under a y-polarised pump drives a purely longitudinal E1 source; + with transverse_source (LPSE takeTransversePartOfSourceTerms) no E1 is produced.""" + from adept._lpse2d.core.raman import RamanLight + from adept._lpse2d.core.spectral_light import SpectralRamanLight + + cls = SpectralRamanLight if solver == "spectral" else RamanLight + out = {} + for transverse_source in (False, True): + cfg = _srs_cfg_2d(solver, transverse_source) + light = cls(cfg) + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + band = np.asarray(cfg["grid"]["low_pass_filter_grid"]) > 0 + j = int(np.flatnonzero(band[0])[2]) + phi_k = jnp.zeros((nx, ny), dtype=jnp.complex128).at[0, j].set(1.0e-6 * nx * ny) + E0 = jnp.stack([jnp.zeros((nx, ny)), 1.0e-3 * jnp.ones((nx, ny))], axis=-1).astype(jnp.complex128) + E1 = light(0.0, jnp.zeros((nx, ny, 2), dtype=jnp.complex128), lambda t, E0=E0: E0, phi_k, None) + out[transverse_source] = float(jnp.max(jnp.abs(E1))) + assert out[False] > 0.0 + assert out[True] < 1e-10 * out[False] diff --git a/tests/test_lpse2d/test_thermal_filamentation.py b/tests/test_lpse2d/test_thermal_filamentation.py new file mode 100644 index 00000000..f1348ce0 --- /dev/null +++ b/tests/test_lpse2d/test_thermal_filamentation.py @@ -0,0 +1,130 @@ +"""Thermal-filamentation source on the IAW velocity divergence (LPSE thermalFil.*).""" + +from copy import deepcopy + +import numpy as np +import yaml +from jax import numpy as jnp + + +def _cfg(tf, absorption=True, collisions=1.0, epw_solver="separate"): + from adept._lpse2d.helpers import get_density_profile, get_derived_quantities, get_solver_quantities, write_units + + with open("tests/test_lpse2d/configs/tpd.yaml") as fi: + cfg = yaml.safe_load(fi) + cfg = deepcopy(cfg) + cfg["density"] = {"basis": "uniform", "val": 0.2} + cfg["grid"].update({"dx": "0.1um", "xmax": "12.8um", "ymax": "3.2um", "ymin": "-3.2um", "dt": "1fs", "tmax": "2fs"}) + cfg["terms"]["epw"]["source"].update({"noise": False, "tpd": False}) + cfg["terms"]["epw"]["damping"]["collisions"] = collisions + cfg["terms"]["light"] = {"absorption": absorption} + if epw_solver == "combined": + cfg["terms"]["epw"]["solver"] = "combined" + cfg["terms"]["epw"]["source"]["srs"] = False + cfg["terms"]["light"]["solver"] = "spectral" + cfg["terms"]["iaw"] = { + "active": True, + "solver": "spectral", + "damping": {"landau": 0.1, "collisions": 0.0}, + "thermal_filamentation": tf, + } + write_units(cfg) + cfg = get_derived_quantities(cfg) + cfg["grid"] = get_solver_quantities(cfg) + cfg["grid"]["background_density"] = get_density_profile(cfg) + return cfg + + +def _fields(cfg, modulation_k=None): + nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] + x = np.asarray(cfg["grid"]["x"]) + e0 = np.zeros((nx, ny, 2), dtype=np.complex128) + amp = 1.0e-3 if modulation_k is None else 1.0e-3 * np.sqrt(1.0 + 0.5 * np.cos(modulation_k * x))[:, None] + e0[..., 1] = amp + return jnp.zeros((nx, ny), dtype=jnp.complex128), jnp.asarray(e0), jnp.zeros((nx, ny, 2), dtype=jnp.complex128) + + +def test_uniform_intensity_gives_no_source_and_modulation_the_derived_coefficient(): + from adept._lpse2d.core.iaw import IonAcousticWave + + cfg = _cfg({"laser": True}) + iaw = IonAcousticWave(cfg) + phi_k, e0, e1 = _fields(cfg) + uniform_max = float(jnp.max(jnp.abs(iaw.thermal_filamentation_source(phi_k, e0, e1)))) + kx = np.asarray(cfg["grid"]["kx"]) + km = kx[4] + phi_k, e0, e1 = _fields(cfg, modulation_k=km) + source = np.asarray(iaw.thermal_filamentation_source(phi_k, e0, e1)) + assert uniform_max < 1e-9 * np.max(np.abs(source)) # only the box-mean subtraction round-off + e_sq = np.sum(np.abs(np.asarray(e0)) ** 2, axis=-1) + expected = np.asarray(iaw.thermal_coeff["laser"]) * (e_sq - e_sq.mean()) + np.testing.assert_allclose(source, expected, rtol=1e-12, atol=1e-30) + # heating where the intensity is high drives outflow: positive divergence there + assert source[np.argmax(e_sq[:, 0]), 0] > 0.0 + # the coefficient: Z * 2 nu_abs (n/nc)^2 * fieldScale^2 / (8 pi m_i kappa') in 1/ps^2 per code |E|^2 + coeff = float(np.asarray(iaw.thermal_coeff["laser"])[0, 0]) + assert coeff > 0.0 and np.isfinite(coeff) + + +def test_nonlocal_correction_scales_as_k_to_the_four_thirds(): + from adept._lpse2d.core.iaw import IonAcousticWave + + local = IonAcousticWave(_cfg({"laser": True})) + nonlocal_ = IonAcousticWave(_cfg({"laser": True, "nonlocal": True})) + cfg = _cfg({"laser": True}) + kx = np.asarray(cfg["grid"]["kx"]) + ratios = [] + for i in (2, 4): + phi_k, e0, e1 = _fields(cfg, modulation_k=kx[i]) + s_loc = np.fft.fft2(np.asarray(local.thermal_filamentation_source(phi_k, e0, e1))) + s_nl = np.fft.fft2(np.asarray(nonlocal_.thermal_filamentation_source(phi_k, e0, e1))) + ratios.append(np.abs(s_nl[i, 0]) / np.abs(s_loc[i, 0]) - 1.0) + # the excess over the local source grows as k^(4/3): ratio of excesses = 2^(4/3) + np.testing.assert_allclose(ratios[1] / ratios[0], 2.0 ** (4.0 / 3.0), rtol=1e-6) + assert ratios[0] > 0.0 + + +def test_lw_heating_in_combined_mode_uses_the_combined_damping(): + """With terms.epw.solver: combined the EPW is the combined field's longitudinal part, damped by the + Raman absorption nu_R (n/n_env)^2, not the EPW collisions (PR #372 review, the same class as the + dissipation diagnostic): the lw heating coefficient equals the raman one on the same field (both + 2 nu_R (n/n_env)^2, rtol 1e-12), does not change with terms.epw.damping.collisions, and needs + terms.light.absorption.""" + import pytest + + from adept._lpse2d.core.iaw import IonAcousticWave + + a = IonAcousticWave(_cfg({"lw": True, "raman": True}, collisions=1.0, epw_solver="combined")) + b = IonAcousticWave(_cfg({"lw": True, "raman": True}, collisions=5.0, epw_solver="combined")) + np.testing.assert_allclose(np.asarray(a.thermal_coeff["lw"]), np.asarray(a.thermal_coeff["raman"]), rtol=1e-12) + np.testing.assert_array_equal(np.asarray(a.thermal_coeff["lw"]), np.asarray(b.thermal_coeff["lw"])) + with pytest.raises(ValueError, match="absorption"): + IonAcousticWave(_cfg({"lw": True}, absorption=False, collisions=1.0, epw_solver="combined")) + + +def test_requirements_and_translator(): + import pytest + + from adept._lpse2d.core.iaw import IonAcousticWave + from adept._lpse2d.lpse_deck import translate_parms + + with pytest.raises(ValueError, match="absorption"): + IonAcousticWave(_cfg({"laser": True}, absorption=False)) + with pytest.raises(ValueError, match="collisions"): + IonAcousticWave(_cfg({"lw": True}, collisions=0.0)) + parms = { + "grid.sizes": "20 5", + "grid.nodes": "201 51", + "laser.enable": "true", + "lw.enable": "true", + "iaw.enable": "true", + "lw.spectral.dt": "0.005", + "simulation.time.end": "1", + "laser.1.intensity": "1e15", + "thermalFil.laser.enable": "true", + "thermalFil.isNonlocal": "true", + "thermalFil.conductivityMultiplier": "2", + } + cfg, report = translate_parms(parms, experiment="x", run="y") + tf = cfg["terms"]["iaw"]["thermal_filamentation"] + assert tf["laser"] and not tf["raman"] and tf["nonlocal"] and tf["conductivity_multiplier"] == 2.0 diff --git a/tests/test_lpse2d/test_threshold.py b/tests/test_lpse2d/test_threshold.py new file mode 100644 index 00000000..87212c09 --- /dev/null +++ b/tests/test_lpse2d/test_threshold.py @@ -0,0 +1,60 @@ +"""Absolute-threshold bisection driver (LPSE AbsoluteThreshold), with a fake runner.""" + +import numpy as np +import pytest + + +def test_bisection_converges_on_a_known_threshold(): + from adept._lpse2d.threshold import find_threshold + + true_threshold = 3.0e14 + + def fake_runner(cfg, intensity, run_name): + rate = 2.0 * (intensity / true_threshold - 1.0) + return {"epw_growth_measurable": 1.0, "epw_growth_rate_per_ps": rate} + + cfg = {"units": {"laser intensity": "1e15W/cm^2"}, "mlflow": {"run": "t"}} + result = find_threshold(cfg, "1e14W/cm^2", "1e15W/cm^2", n_iter=10, runner=fake_runner, log_mlflow=False) + lo, hi = result["bracket"] + assert lo < true_threshold < hi + assert abs(result["threshold"] - true_threshold) / true_threshold < 0.01 + assert len(result["history"]) == 12 # two bracket runs + 10 bisections + assert all(h[1] == (h[0] > true_threshold) for h in result["history"]) + + +def test_bad_brackets_are_rejected(): + from adept._lpse2d.threshold import find_threshold + + def always(cfg, intensity, name): + return {"epw_growth_measurable": 1.0, "epw_growth_rate_per_ps": 5.0} + + def never(cfg, intensity, name): + return {"epw_growth_measurable": 0.0, "epw_growth_rate_per_ps": 0.0} + + cfg = {"units": {"laser intensity": "1e15W/cm^2"}} + with pytest.raises(ValueError, match="already unstable"): + find_threshold(cfg, 1e14, 1e15, n_iter=1, runner=always, log_mlflow=False) + with pytest.raises(ValueError, match="still stable"): + find_threshold(cfg, 1e14, 1e15, n_iter=1, runner=never, log_mlflow=False) + assert np.isfinite(1.0) + + +def test_a_zero_lower_bound_is_rejected_before_any_run(): + """The geometric midpoint sqrt(lo hi) of a zero lower bound never moves, so the bracket would + never narrow and the reported threshold would be 0 (PR #372 review): the bracket needs + 0 < intensity_lo < intensity_hi, checked before a single run; a positive bracket still bisects.""" + from adept._lpse2d.threshold import find_threshold + + calls = [] + + def runner(cfg, intensity, name): + calls.append(intensity) + return {"epw_growth_measurable": 1.0, "epw_growth_rate_per_ps": 2.0 * (intensity / 3.0e14 - 1.0)} + + cfg = {"units": {"laser intensity": "1e15W/cm^2"}, "mlflow": {"run": "t"}} + for lo in (0.0, "0W/cm^2", -1.0e14): + with pytest.raises(ValueError, match="0 < intensity_lo"): + find_threshold(cfg, lo, 1e15, n_iter=3, runner=runner, log_mlflow=False) + assert calls == [] + result = find_threshold(cfg, 1e13, 1e15, n_iter=12, runner=runner, log_mlflow=False) + assert abs(result["threshold"] - 3.0e14) / 3.0e14 < 0.01 diff --git a/tests/test_lpse2d/test_tpd_depletion.py b/tests/test_lpse2d/test_tpd_depletion.py index fd073502..f3f8ffab 100644 --- a/tests/test_lpse2d/test_tpd_depletion.py +++ b/tests/test_lpse2d/test_tpd_depletion.py @@ -3,6 +3,7 @@ from copy import deepcopy import numpy as np +import pytest import yaml from jax import numpy as jnp @@ -154,15 +155,19 @@ def test_tpd_accepts_box_averaged_2d_particle_feedback(): def test_tpd_reciprocal_coupling_conserves_wave_energy(): - """The new pump term is the exact discrete reciprocal of the existing TPD term. + """The pump term is LPSE's (Follett Eqs 53/55): twice the EPW-side TPD coefficient. - For coupling terms alone, d/dt [|E0|^2 + (wp0/w0)|E_epw|^2] = 0. - This locks the coefficient, complex-conjugation convention, and carrier phase. + For the coupling terms alone, d/dt [|E0|^2 + (2 wp0/w0)|E_epw|^2] = 0, i.e. the total + wave energy at envelope density n_c/4 (two plasmons at wp0 per pump photon at w0). + This locks the coefficient, complex-conjugation convention, and carrier phase. The + transverse projection is switched off here because the random test pump is not + divergence-free; it is checked separately in test_lpse_parity.py. """ from adept._lpse2d.core.epw import SpectralEPWSolver from adept._lpse2d.core.light import CoupledLight cfg = _make_cfg() + cfg["terms"]["light"]["tpd_projection"] = False pump = CoupledLight(cfg) epw = SpectralEPWSolver(cfg) nx, ny = cfg["grid"]["nx"], cfg["grid"]["ny"] @@ -173,16 +178,21 @@ def test_tpd_reciprocal_coupling_conserves_wave_energy(): # state in the same band makes the filtered source operator self-adjoint. phi_k *= np.asarray(cfg["grid"]["low_pass_filter_grid"] * cfg["grid"]["zero_mask"]) phi_k = jnp.asarray(phi_k) - E0_y = jnp.asarray(rng.normal(size=(nx, ny)) + 1j * rng.normal(size=(nx, ny))) + E0 = jnp.asarray(rng.normal(size=(nx, ny, 2)) + 1j * rng.normal(size=(nx, ny, 2))) t = 0.037 - _, ey = epw.phi_k_to_e_fields(phi_k) - tpd_epw_rhs = epw.calc_tpd_source(t, phi_k, ey, E0_y) + ex, ey = epw.phi_k_to_e_fields(phi_k) + tpd_epw_rhs = epw.calc_tpd_source(t, phi_k, ex, ey, E0) tpd_pump_rhs = pump.calc_tpd_depletion(t, phi_k) + assert tpd_pump_rhs.shape == (nx, ny, 2) - pump_energy_rate = 2.0 * jnp.real(jnp.mean(jnp.conj(E0_y) * tpd_pump_rhs)) + pump_energy_rate = 2.0 * jnp.real(jnp.mean(jnp.sum(jnp.conj(E0) * tpd_pump_rhs, axis=-1))) epw_energy_rate = 2.0 * jnp.real(jnp.vdot(phi_k * epw.k_sq, tpd_epw_rhs)) / (nx * ny) ** 2 - total_rate = pump_energy_rate + epw.wp0 / epw.w0 * epw_energy_rate + weight = 2.0 * epw.wp0 / epw.w0 + total_rate = pump_energy_rate + weight * epw_energy_rate - scale = max(abs(float(pump_energy_rate)), abs(float(epw.wp0 / epw.w0 * epw_energy_rate))) + scale = max(abs(float(pump_energy_rate)), abs(float(weight * epw_energy_rate))) assert abs(float(total_rate)) < 1.0e-7 * scale + # the pump coefficient is exactly twice the EPW source coefficient + assert pump.tpd_depletion_coeff0 == pytest.approx(2.0 * epw.tpd_prefactor * epw.w0 / epw.w0) + assert pump.tpd_depletion_coeff0 == pytest.approx(1j * epw.e / (2.0 * epw.w0 * epw.me))