diff --git a/adept/_lpse2d/helpers.py b/adept/_lpse2d/helpers.py index 8ba4b0ee..9eec11ac 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,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 +1080,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 +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", - ) + 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") - 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 + 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..1f9ea365 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,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 diff --git a/docs/source/solvers/lpse2d/overview.md b/docs/source/solvers/lpse2d/overview.md index af292ab9..ad8cefa4 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). 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_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_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}"