From 19cbe2e6fb4676f2f31ddf7d9afb9ba910720d47 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Mon, 19 Oct 2026 23:18:20 -0700 Subject: [PATCH 1/3] fix(lpse2d): put the interpolated field save on a coarser grid, in x-space With save.fields.x.dx / save.fields.y.dy set, the interpolated field save (get_save_quantities -> save_func, make_field_xarrays) was wrong in four ways (measured on ce5d3fe: 6.4 x 1.6 um TPD box, 0.05 um grid, 0.08 um save spacing, 20 fs, compared against the full-grid save of the same run sampled at the save points): 1. The save grid was not coarser. It kept the simulation's nx x ny points, spread between xmin + dx_save/2 and xmax - dx_save/2, so dx_save only set the margins. 2. The saved epw state is the k-space potential phi_k (FFT-ordered), but it went through the same x-space interpolator as the real-space fields, and make_field_xarrays then read the result as k-space on kx = fftfreq(nx, d=dx_save). The saved phi, ex and ey differed from the full-grid save by 68 %, 67 % and 73 % of their maximum, and the saved k axis was scaled by dx / dx_save (x0.625). 3. With dx_save below the grid spacing, the outermost save points fell outside the outermost cell centres, where interpax returns NaN. 4. E0, E1 and the IAW fields were interpolated correctly, but onto the wrong (non-coarsened) grid. Fix: - The save grid has n_save = round(L / dx_save) cell-centred points per axis, with spacing L / n_save (_interpolated_save_axis). A save grid with more points than the simulation grid raises a ValueError at setup. Giving only one of save.fields.x / save.fields.y raises NotImplementedError (x alone already did; y alone was silently ignored). - The save function (_interpolated_field_save) transforms phi_k to x-space with the solver's convention (phi = ifft2(phi_k), E = ifft2(-i k phi_k)) on the full grid and saves phi, ex and ey on the save grid. E0, E1 and the real fields are interpolated as before. ex and ey come from the full grid rather than from a spectral derivative on the save grid, because that derivative would scale the interpolation error by k and assign the wrong k to content above the save grid's Nyquist wavenumber. - make_field_xarrays, interpolated branch: the x-space ex and ey carry the full-grid save's 4 / (nx * ny) factor with the simulation nx, ny. The k-space fields (phi, ex, ey, IAW) are the DFT on the save grid, on its own k axes, scaled by (nx * ny) / (nx_save * ny_save) so that a resolved mode keeps its full-grid coefficient. - The full-grid save path is unchanged. fields.xr and k-fields.xr of a full-grid run are bitwise identical to ce5d3fe (x64 and float32, 2D and quasi-1D, TPD+SRS+IAW box). Behavior change: interpolated saves now have n_save points per axis instead of nx x ny, and the raw diffrax "fields" save holds phi, ex and ey (x-space, on the save grid) instead of epw. Tests (tests/test_lpse2d/test_field_save.py; all 6 fail on ce5d3fe): - test_interpolated_save_matches_full_grid_save[2d, 1d]: noise-seeded TPD+SRS+IAW box with a static pump, 6 fs, save spacing 1.6 x the grid spacing. Every x-space field equals the full-grid field linearly interpolated to the save points, within 2000 eps of its maximum (round-off budget derived in the file; measured <= 2.3 eps). Shapes are 80 x 20 (2D) / 80 x 1 (quasi-1D), save and k axes match the save grid to 1e-13, and all values are finite. - test_interpolated_save_of_band_limited_state[2d, 1d]: a single-mode state checked against analytic values: the x-space phi, ex, ey, E0, E1 and IAW fields and the mode's k-space coefficient, each within the bilinear-interpolation bound sqrt(2) (q h)^2 / 8 plus round-off. - test_interpolated_save_spacing_bounds[x, y]: half the grid spacing raises. Equal spacing reproduces the grid axis bitwise and gives finite output equal to linear interpolation within 2000 eps. Giving x or y alone raises. Docs: docs/source/solvers/lpse2d/config.md describes the save grid, the spacing bound, the interpolation and the normalizations. Co-Authored-By: Claude Opus 5.5 (1M context) --- adept/_lpse2d/helpers.py | 239 ++++++++++++------ docs/source/solvers/lpse2d/config.md | 25 +- tests/test_lpse2d/test_field_save.py | 357 +++++++++++++++++++++++++++ 3 files changed, 536 insertions(+), 85 deletions(-) create mode 100644 tests/test_lpse2d/test_field_save.py diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index 8ba4b0ee..f5134415 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 @@ -1041,14 +1040,30 @@ 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, ex and ey + are derived from it here. The x-space ex and ey carry a factor 4 / (nx * ny). + + 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). ex and + ey get the same 4 / (nx * ny) with the *simulation* nx, ny, 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 +1081,27 @@ 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 / nx_grid / ny_grid * 4 + ey_np = ey_vs_t / nx_grid / ny_grid * 4 + 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 + ex_np = np.fft.ifft2(ex_k_np, axes=(1, 2)) / nx / ny * 4 + ey_np = np.fft.ifft2(ey_k_np, axes=(1, 2)) / nx / ny * 4 phi_k = xr.DataArray( np.fft.fftshift(phi_k_np, axes=(1, 2)), @@ -1089,8 +1116,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)) @@ -1143,11 +1170,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 +1188,117 @@ def make_field_xarrays(cfg, this_t, state, td): return kfields, fields +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 @@ -1175,75 +1316,11 @@ 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 ny == 1: - # Quasi-1D (single transverse cell): interpax.interp2d needs >=2 nodes per - # axis and returns NaN off the lone y-node, so interpolate in x only and - # keep the single y-row. save_func reshapes the flat (nx,) result to (nx, 1). - x_save = cfg["save"]["fields"]["x"]["ax"] - x_src = cfg["grid"]["x"] - - def interpolator(f): - return interpax.interp1d(x_save, x_src, jnp.reshape(f, (-1,)), method="linear") - else: - interpolator = partial( - interpax.interp2d, - xq=jnp.reshape(xq, (nx * ny), order="F"), - yq=jnp.reshape(yq, (nx * ny), order="F"), - x=cfg["grid"]["x"], - y=cfg["grid"]["y"], - method="linear", - ) - - def save_func(t, y, args): - from adept._lpse2d.core.hpe import PARTICLE_KEYS + 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") - 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 + if "x" in cfg["save"]["fields"]: + save_func = _interpolated_field_save(cfg) else: diff --git a/docs/source/solvers/lpse2d/config.md b/docs/source/solvers/lpse2d/config.md index 3d503275..f826e56f 100644 --- a/docs/source/solvers/lpse2d/config.md +++ b/docs/source/solvers/lpse2d/config.md @@ -195,8 +195,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 +210,30 @@ 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, including + the `4 / (nx * ny)` factor on `ex`/`ey`, which uses the simulation `nx`, `ny`. +- `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 diff --git a/tests/test_lpse2d/test_field_save.py b/tests/test_lpse2d/test_field_save.py new file mode 100644 index 00000000..6dd86488 --- /dev/null +++ b/tests/test_lpse2d/test_field_save.py @@ -0,0 +1,357 @@ +"""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 and the 4 / (nx * ny) scaling: 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) * 4 / (nx ny), ey = -i ky phi(x_s) * 4 / (nx ny) + -- the full-grid save's normalization with the SIMULATION nx, ny, so both saves report the + same field -- 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 + + norm = 4.0 / (nx * ny) + phi_s = at_save(amp, kx, ky) + expected_x = { + "phi": (phi_s, kx, ky), + "ex": (-1j * kx * phi_s * norm, kx, ky), + "ey": (-1j * ky * phi_s * norm, 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) From 9eb7ff59414746b2d6a1cede8fb224bc9320d9b0 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 20 Oct 2026 11:55:20 -0700 Subject: [PATCH 2/3] fix(lpse2d): saved EPW ex/ey without the stray 4/(nx ny) factor make_field_xarrays wrote the x-space EPW fields of fields.xr as ex = ifft2(-i kx phi_k) / nx / ny * 4 (ey likewise), while the saved phi is ifft2(phi_k) unscaled and the solver's field is E = ifft2(-i k phi_k) (SpectralEPWSolver.phi_k_to_e_fields; the IAW ponderomotive term, the TPD pump depletion and the e_sq / epw_energy / epw_dissipation series use the same unscaled E). fields.xr ex/ey were therefore smaller than the field the solver evolves, and than -grad of the saved phi, by 4 / (nx ny): 1/1024 on a 128 x 32 grid, 1/32 on 128 x 1, ~1/14000 on 512 x 108. k-fields.xr (phi_k, -i kx phi_k, -i ky phi_k) was already unscaled, so the x- and k-space datasets disagreed by the same factor. History: the factor dates from 02e112f (PR #20), where the first EPW solver stored phi in k-space with fft_norm = nx ny / 4 (its get_eh_x multiplied by nx ny / 4; the save used the reciprocal 4 / (nx ny)). f77784f ("new units and solver based on matlab code", squashed into df1510a / PR #41) moved the state to the plain numpy FFT convention but kept the factor on ex/ey, and c641495 (PR #142) moved the state to the k-space phi_k and kept it again. The MATLAB prototype has no such factor; in the current code it has no physical meaning. Fix: drop the factor. The x-space ex/ey are ifft2(-i k phi_k), i.e. -grad(phi) of the saved phi. Both branches of make_field_xarrays lose it: the full-grid one, and the interpolated (save.fields.x / y) one of the previous commit, which had applied the same factor with the simulation nx, ny to its x-space ex/ey so the two saves agreed. k-fields.xr, phi, the light fields and the IAW fields are unchanged. Plots in plot_fields / plot_kt use autoscaled colour limits, no normalization assumed the scaled values. docs/source/solvers/lpse2d/ overview.md now states what the saved EPW entries are. Tests (tests/test_lpse2d/test_saved_epw_fields.py, 4 tests, all fail on ce5d3fe; 128 x 32 and quasi-1D 128 x 1 boxes, full-grid save, a single grid mode plus a random band-limited phi_k, round-off bound 200 eps fixed before the first run): - test_saved_epw_fields_single_mode[2d, 1d]: fields.xr phi/ex/ey and k-fields.xr phi/ex/ey equal the analytic single mode. On ce5d3fe ex is off by 5.9e-3 of a 5.9e-3 field. - test_saved_epw_fields_are_the_solver_field[2d, 1d]: saved ex/ey equal -grad(saved phi) (spectral), SpectralEPWSolver.phi_k_to_e_fields(phi_k) and ifft2 of the k-fields.xr ex/ey, and dx dy sum |E|^2 of the saved fields equals the e_sq series. On ce5d3fe the relative L2 error is 0.99902 = 1 - 4/4096 (2D) and 0.96875 = 1 - 4/128 (quasi-1D). - test_field_save.py::test_interpolated_save_of_band_limited_state now expects ex = -i kx phi, ey = -i ky phi on the save grid (factor removed from its expected values and docstring; tolerances unchanged). test_interpolated_save_matches_full_grid_save needs no change: it compares the interpolated save with the full-grid save. Cherry-picked from fix/lpse2d-saved-epw-field-scale (6b6d278, off ce5d3fe); the conflict in make_field_xarrays is resolved as above. Behavior change: fields.xr ex and ey grow by nx ny / 4 relative to every earlier save (all configs, full-grid and interpolated). phi, k-fields.xr, series.xr and the logged metrics are unaffected (none read fields.xr ex/ey). Co-Authored-By: Claude Opus 5.5 (1M context) --- adept/_lpse2d/helpers.py | 24 +-- docs/source/solvers/lpse2d/config.md | 3 +- docs/source/solvers/lpse2d/overview.md | 2 + tests/test_lpse2d/test_field_save.py | 12 +- tests/test_lpse2d/test_saved_epw_fields.py | 216 +++++++++++++++++++++ 5 files changed, 236 insertions(+), 21 deletions(-) create mode 100644 tests/test_lpse2d/test_saved_epw_fields.py diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index f5134415..9eec11ac 100644 --- a/adept/_lpse2d/helpers.py +++ b/adept/_lpse2d/helpers.py @@ -1042,16 +1042,15 @@ 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, ex and ey - are derived from it here. The x-space ex and ey carry a factor 4 / (nx * ny). + 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). ex and - ey get the same 4 / (nx * ny) with the *simulation* nx, ny, 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. + 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"] interpolated = "x" in fld_save @@ -1093,15 +1092,16 @@ def make_field_xarrays(cfg, this_t, state, td): 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 / nx_grid / ny_grid * 4 - ey_np = ey_vs_t / nx_grid / ny_grid * 4 + 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 - ex_np = np.fft.ifft2(ex_k_np, axes=(1, 2)) / nx / ny * 4 - ey_np = np.fft.ifft2(ey_k_np, axes=(1, 2)) / nx / ny * 4 + # 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)), diff --git a/docs/source/solvers/lpse2d/config.md b/docs/source/solvers/lpse2d/config.md index f826e56f..1f9ea365 100644 --- a/docs/source/solvers/lpse2d/config.md +++ b/docs/source/solvers/lpse2d/config.md @@ -229,8 +229,7 @@ fields are saved on a coarser grid: `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, including - the `4 / (nx * ny)` factor on `ex`/`ey`, which uses the simulation `nx`, `ny`. +- 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). diff --git a/docs/source/solvers/lpse2d/overview.md b/docs/source/solvers/lpse2d/overview.md index af292ab9..afe4943c 100644 --- a/docs/source/solvers/lpse2d/overview.md +++ b/docs/source/solvers/lpse2d/overview.md @@ -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): `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. 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_field_save.py b/tests/test_lpse2d/test_field_save.py index 6dd86488..e06b8b29 100644 --- a/tests/test_lpse2d/test_field_save.py +++ b/tests/test_lpse2d/test_field_save.py @@ -34,7 +34,7 @@ # * 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 and the 4 / (nx * ny) scaling: a few 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 @@ -203,9 +203,8 @@ def test_interpolated_save_of_band_limited_state(geometry, tmp_path): 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) * 4 / (nx ny), ey = -i ky phi(x_s) * 4 / (nx ny) - -- the full-grid save's normalization with the SIMULATION nx, ny, so both saves report the - same field -- and E0, E1, IAW at x_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. @@ -258,12 +257,11 @@ 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 - norm = 4.0 / (nx * ny) phi_s = at_save(amp, kx, ky) expected_x = { "phi": (phi_s, kx, ky), - "ex": (-1j * kx * phi_s * norm, kx, ky), - "ey": (-1j * ky * phi_s * norm, 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), 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}" From 984f0f63e85eacce7099d11d4cb74804cc1e2085 Mon Sep 17 00:00:00 2001 From: Phil Travis Date: Tue, 20 Oct 2026 16:35:03 -0700 Subject: [PATCH 3/3] docs(lpse2d): scope the saved-field Fourier identities to full-grid saves The overview stated that k-fields.xr holds phi_k, -i kx phi_k, -i ky phi_k and that fields.xr is their ifft2, with ex, ey = -grad(phi). That holds for a full-grid save only. An interpolated save (save.fields.x / y) stores 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); neither ifft2(k-fields.xr) nor a derivative of the coarse phi reproduces its fields.xr. The paragraph now says so and links to the save.fields reference in config.md. Review thread r4331320857 (PR #388). Docs only. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/source/solvers/lpse2d/overview.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/solvers/lpse2d/overview.md b/docs/source/solvers/lpse2d/overview.md index afe4943c..ad8cefa4 100644 --- a/docs/source/solvers/lpse2d/overview.md +++ b/docs/source/solvers/lpse2d/overview.md @@ -101,7 +101,7 @@ 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): `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. Field saves written before this was fixed (October 2026) carry an extra factor `4 / (nx * ny)` on `fields.xr` `ex` and `ey` only. +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`: