From cab64700c2d2958cdb07b0a7a297af7077fd9822 Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Tue, 4 Aug 2026 23:36:14 +0200 Subject: [PATCH 1/8] Add test: trajectory must not depend on batch-mates' release times Fails on main. Found while writing lcs_parcels. Co-authored-by: Claude --- tests/test_particleset_execute.py | 47 +++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_particleset_execute.py b/tests/test_particleset_execute.py index 244da9643..9d51403ce 100644 --- a/tests/test_particleset_execute.py +++ b/tests/test_particleset_execute.py @@ -48,6 +48,53 @@ def zonal_flow_fieldset() -> FieldSet: return FieldSet.from_sgrid_conventions(ds, mesh="flat") +@pytest.fixture +def time_varying_zonal_flow_fieldset() -> FieldSet: + """Flat A-grid whose U varies in time only (V = 0), sampled 3-hourly. + + Because U is spatially uniform, a particle's displacement depends on its own + clock and nothing else. + """ + nt = 25 + ds = simple_UV_dataset(dims=(nt, 2, 6, 6), mesh="flat") + times = np.array([np.timedelta64(3 * i, "h") for i in range(nt)]) + ds["time"] = ("time", times, {"axis": "T"}) + u = np.cos(2 * np.pi * (times / np.timedelta64(1, "s")) / 86400.0) # 1 m/s, 1-day period + ds["U"].data[:] = u[:, None, None, None] + return FieldSet.from_sgrid_conventions(ds, mesh="flat") + + +def test_execute_trajectory_independent_of_other_particles_release_times(time_varying_zonal_flow_fieldset): + """A particle's trajectory must not depend on when its batch-mates were released. + + ``ParticleSet`` accepts a per-particle ``t``, so a set can hold particles on + different time indices. Interpolation is per-particle, hence batching a + particle with later-released mates must not change its own result. + """ + fieldset = time_varying_zonal_flow_fieldset + t0 = np.timedelta64(0, "s") + + def run(release_times): + npart = len(release_times) + pset = ParticleSet( + fieldset, + pclass=Particle, + t=np.array(release_times), + z=np.zeros(npart), + y=np.zeros(npart), + x=np.zeros(npart), + ) + pset.execute(AdvectionRK4, dt=np.timedelta64(1, "h"), endtime=np.timedelta64(48, "h")) + return pset.x[0], pset.y[0] + + alone = run([t0]) + uniform = run([t0] * 4) + staggered = run([t0] + [t0 + np.timedelta64(3, "h")] * 3) + + assert uniform == pytest.approx(alone) + assert staggered == pytest.approx(alone) + + def test_pset_execute_invalid_arguments(fieldset, fieldset_no_time_interval): with pytest.raises(RuntimeWarning, match="invalid value encountered in cast.*"): ParticleSet(fieldset, x=[0.2], y=[5.0], pclass=Particle).execute(AdvectionRK4, dt=np.timedelta64(None)) From bde7b92f2000f3f40167afb70fa7a47fac659fed Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Tue, 4 Aug 2026 23:50:21 +0200 Subject: [PATCH 2/8] Fix corner index ordering in structured interpolators The index arrays were built particle-major for T/Z but corner-major for Y/X, while the reshape is corner-major. The orderings agree only when all particles share a time and depth index, so batches with a non-uniform clock or spread over depth levels were gathered from the wrong level. Derive all four index arrays from one layout in _gather_corners, and route the A-grid, C-grid velocity and C-grid tracer gathers through it. Also fixes CGrid_Tracer's np.repeat(ti) missing its repeats argument. test_nemo_3D_curvilinear_fieldset expectations recorded the old batch behaviour; they now match what each particle gets advected on its own. Co-authored-by: Claude --- src/parcels/interpolators/_xinterpolators.py | 202 +++++++++---------- tests/test_advection.py | 6 +- tests/test_interpolation.py | 74 +++++++ 3 files changed, 175 insertions(+), 107 deletions(-) diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index 54606c4b0..01de64d8f 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -2,6 +2,7 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING, Literal import numpy as np @@ -17,6 +18,63 @@ from parcels._core.xgrid import XGrid +_CORNER_AXES: tuple[ptyping.XgridAxis, ...] = ("T", "Z", "Y", "X") + + +def _gather_corners( + data: np.ndarray | xr.DataArray, + axis_dim: dict[ptyping.ptyping.XgridAxis, str], + levels: dict[ptyping.XgridAxis, tuple[np.ndarray, ...]], + npart: int, +) -> np.ndarray: + """Gather field data at the corners bracketing each particle. + + The gather is the outer product over ``_CORNER_AXES``, each axis contributing + one or two levels, with **the particle index innermost**. That layout is + stated once here -- as ``counts`` -- and drives both the index arrays and the + shape of the result, so the two cannot drift apart. + + Parameters + ---------- + data : + Field data, with dimensions ordered ``(time, Z, Y, X)``. + axis_dim : + Maps ``"X"``/``"Y"``/``"Z"`` to dimension names; ``"T"`` is ``"time"``. + levels : + Maps an axis to the per-particle index arrays of its corner levels, e.g. + ``{"T": (ti, ti + 1), "Y": (yi, yi + 1)}``. An axis absent from + ``levels`` contributes a single level. An axis the field has no + dimension for still shapes the result, but is not indexed -- its corners + simply repeat. + npart : + Number of particles. + + Returns + ------- + np.ndarray + Of shape ``(*counts, npart)``, where ``counts[i]`` is the number of + levels gathered on ``_CORNER_AXES[i]``. + """ + dims = {axis: ("time" if axis == "T" else axis_dim.get(axis)) for axis in _CORNER_AXES} + counts = tuple(len(levels[axis]) if axis in levels else 1 for axis in _CORNER_AXES) + + selection_dict = {} + before = 1 # number of corner combinations on the axes outside this one + for i, (axis, n_levels) in enumerate(zip(_CORNER_AXES, counts, strict=True)): + if axis in levels and dims[axis] is not None and dims[axis] in data.dims: + after = math.prod(counts[i + 1 :]) + # level i of particle p belongs at flat position + # ((i_before * n_levels + i) * after + i_after) * npart + p + stacked = np.stack(np.broadcast_arrays(*levels[axis])) # (n_levels, npart) + selection_dict[dims[axis]] = xr.DataArray( + np.broadcast_to(stacked[None, :, None, :], (before, n_levels, after, npart)).reshape(-1), + dims=("points"), + ) + before *= n_levels + + return data.isel(selection_dict).data.reshape(*counts, npart) + + def _get_corner_data_Agrid( data: np.ndarray | xr.DataArray, ti: int, @@ -29,40 +87,13 @@ def _get_corner_data_Agrid( axis_dim: dict[ptyping.ptyping.XgridAxis, str], ) -> np.ndarray: """Helper function to get the corner data for a given A-grid field and position.""" - # Time coordinates: 8 points at ti, then 8 points at ti+1 - if lenT == 1: - ti = np.repeat(ti, lenZ * 4) - else: - ti_1 = np.clip(ti + 1, 0, data.shape[0] - 1) - ti = np.concatenate([np.repeat(ti, lenZ * 4), np.repeat(ti_1, lenZ * 4)]) - - # Z coordinates: 4 points at zi, 4 at zi+1, repeated for both time levels - if lenZ == 1: - zi = np.repeat(zi, lenT * 4) - else: - zi_1 = np.clip(zi + 1, 0, data.shape[1] - 1) - zi = np.tile(np.array([zi, zi, zi, zi, zi_1, zi_1, zi_1, zi_1]).flatten(), lenT) - - # Y coordinates: [yi, yi, yi+1, yi+1] for each spatial point, repeated for time/z - yi_1 = np.clip(yi + 1, 0, data.shape[2] - 1) - yi = np.tile(np.array([yi, yi, yi_1, yi_1]).flatten(), lenT * lenZ) - - # X coordinates: [xi, xi+1, xi, xi+1] for each spatial point, repeated for time/z - xi_1 = np.clip(xi + 1, 0, data.shape[3] - 1) - xi = np.tile(np.array([xi, xi_1]).flatten(), lenT * lenZ * 2) - - # Create DataArrays for indexing - selection_dict = {} - if "X" in axis_dim: - selection_dict[axis_dim["X"]] = xr.DataArray(xi, dims=("points")) - if "Y" in axis_dim: - selection_dict[axis_dim["Y"]] = xr.DataArray(yi, dims=("points")) - if "Z" in axis_dim: - selection_dict[axis_dim["Z"]] = xr.DataArray(zi, dims=("points")) - if "time" in data.dims: - selection_dict["time"] = xr.DataArray(ti, dims=("points")) - - return data.isel(selection_dict).data.reshape(lenT, lenZ, 2, 2, npart) + levels = { + "T": (ti,) if lenT == 1 else (ti, np.clip(ti + 1, 0, data.shape[0] - 1)), + "Z": (zi,) if lenZ == 1 else (zi, np.clip(zi + 1, 0, data.shape[1] - 1)), + "Y": (yi, np.clip(yi + 1, 0, data.shape[2] - 1)), + "X": (xi, np.clip(xi + 1, 0, data.shape[3] - 1)), + } + return _gather_corners(data, axis_dim, levels, npart) def _get_offsets_dictionary(grid: XGrid) -> dict[ptyping.CfAxisSpatial, Literal[1, 0]]: @@ -213,39 +244,23 @@ def interp( py[3], py[0], px[3], px[0], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 0.0), py), grid.deg2m ) - def _create_selection_dict(dims, zdir=False): - """Helper function to create DataArrays for indexing.""" - axis_dim = grid.get_axis_dim_mapping(dims) - selection_dict = { - axis_dim["X"]: xr.DataArray(xi_full, dims=("points")), - axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")), - } + npart = len(xsi) + t_levels = (ti,) if lenT == 1 else (ti, np.clip(ti + 1, 0, tdim - 1)) - # Time coordinates: 2 points at ti, then 2 points at ti+1 - if "time" in dims: - if lenT == 1: - ti_full = np.repeat(ti, 2) - else: - ti_1 = np.clip(ti + 1, 0, tdim - 1) - ti_full = np.concatenate([np.repeat(ti, 2), np.repeat(ti_1, 2)]) - selection_dict["time"] = xr.DataArray(ti_full, dims=("points")) - - if "Z" in axis_dim: - if zdir: - # Z coordinates: 1 point at zi and 1 point at zi+1 repeated for lenT time levels - zi_0 = np.clip(zi + offsets["Z"], 0, zdim - 1) - zi_1 = np.clip(zi + offsets["Z"] + 1, 0, zdim - 1) - zi_full = np.tile(np.array([zi_0, zi_1]).flatten(), lenT) - else: - # Z coordinates: 2 points at zi, repeated for lenT time levels - zi_full = np.repeat(zi, lenT * 2) - selection_dict[axis_dim["Z"]] = xr.DataArray(zi_full, dims=("points")) - - return selection_dict - - def _compute_corner_data(data, selection_dict) -> np.ndarray: - """Helper function to load and reduce corner data over time dimension if needed.""" - corner_data = data.isel(selection_dict).data.reshape(lenT, 2, len(xsi)) + def _compute_corner_data(data, y_levels, x_levels, z_levels=None) -> np.ndarray: + """Gather the two bracketing face values and reduce over time if needed. + + Exactly one of the Z/Y/X axes contributes the two corners; the others + contribute a single level. + """ + levels = { + "T": t_levels, + "Z": z_levels if z_levels is not None else (zi,), + "Y": y_levels, + "X": x_levels, + } + axis_dim = grid.get_axis_dim_mapping(data.dims) + corner_data = _gather_corners(data, axis_dim, levels, npart).reshape(lenT, 2, npart) if lenT == 2: tau_full = tau[np.newaxis, :] @@ -254,29 +269,19 @@ def _compute_corner_data(data, selection_dict) -> np.ndarray: corner_data = corner_data[0, :] return corner_data - # Compute U velocity + # Compute U velocity: the two corners are the X faces yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) - yi_full = np.tile(np.array([yi_o, yi_o]).flatten(), lenT) - xi_1 = np.clip(xi + 1, 0, xdim - 1) - xi_full = np.tile(np.array([xi, xi_1]).flatten(), lenT) - - selection_dict = _create_selection_dict(U.dims) - corner_data = _compute_corner_data(U, selection_dict) + corner_data = _compute_corner_data(U, y_levels=(yi_o,), x_levels=(xi, xi_1)) U0 = corner_data[0, :] * c4 U1 = corner_data[1, :] * c2 Uvel = (1 - xsi) * U0 + xsi * U1 - # Compute V velocity + # Compute V velocity: the two corners are the Y faces yi_1 = np.clip(yi + 1, 0, ydim - 1) - yi_full = np.tile(np.array([yi, yi_1]).flatten(), lenT) - xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) - xi_full = np.tile(np.array([xi_o, xi_o]).flatten(), lenT) - - selection_dict = _create_selection_dict(V.dims) - corner_data = _compute_corner_data(V, selection_dict) + corner_data = _compute_corner_data(V, y_levels=(yi, yi_1), x_levels=(xi_o,)) V0 = corner_data[0, :] * c1 V1 = corner_data[1, :] * c3 @@ -311,16 +316,12 @@ def _compute_corner_data(data, selection_dict) -> np.ndarray: if vectorfield.W: W = vectorfield.W.data - # Y coordinates: yi+offset for each spatial point, repeated for time + # Compute W velocity: the two corners are the Z faces yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) - yi_full = np.tile(yi_o, (lenT) * 2) - - # X coordinates: xi+offset for each spatial point, repeated for time xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) - xi_full = np.tile(xi_o, (lenT) * 2) - - selection_dict = _create_selection_dict(W.dims, zdir=True) - corner_data = _compute_corner_data(W, selection_dict) + zi_0 = np.clip(zi + offsets["Z"], 0, zdim - 1) + zi_1 = np.clip(zi + offsets["Z"] + 1, 0, zdim - 1) + corner_data = _compute_corner_data(W, y_levels=(yi_o,), x_levels=(xi_o,), z_levels=(zi_0, zi_1)) w = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta if is_dask_collection(w): @@ -364,28 +365,17 @@ def interp( xi = np.clip(xi + offsets["X"], 0, data.shape[3] - 1) lenT = 2 if np.any(tau > 0) else 1 + npart = len(xi) - if lenT == 2: - ti_1 = np.clip(ti + 1, 0, data.shape[0] - 1) - ti = np.concatenate([np.repeat(ti), np.repeat(ti_1)]) - zi = np.tile(zi, (lenT) * 2) - yi = np.tile(yi, (lenT) * 2) - xi = np.tile(xi, (lenT) * 2) - - # Create DataArrays for indexing - selection_dict = { - axis_dim["X"]: xr.DataArray(xi, dims=("points")), - axis_dim["Y"]: xr.DataArray(yi, dims=("points")), + levels = { + "T": (ti,) if lenT == 1 else (ti, np.clip(ti + 1, 0, data.shape[0] - 1)), + "Z": (zi,), + "Y": (yi,), + "X": (xi,), } - if "Z" in axis_dim: - selection_dict[axis_dim["Z"]] = xr.DataArray(zi, dims=("points")) - if "time" in field.data.dims: - selection_dict["time"] = xr.DataArray(ti, dims=("points")) - - value = data.isel(selection_dict).data.reshape(lenT, len(xi)) + value = _gather_corners(data, axis_dim, levels, npart).reshape(lenT, npart) if lenT == 2: - tau = tau[:, np.newaxis] value = value[0, :] * (1 - tau) + value[1, :] * tau else: value = value[0, :] diff --git a/tests/test_advection.py b/tests/test_advection.py index e98b26014..a9d2c6708 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -466,9 +466,13 @@ def test_nemo_3D_curvilinear_fieldset(kernel): np.testing.assert_allclose([p.z for p in pset], z_initial) elif kernel == AdvectionRK4_3D: # TODO check why decimals needs to be so low in RK4_3D (compare to v3) + # These particles sit at different depth levels, so the C-grid gather used to + # mix their zi indices; the previous expectations recorded that. Each value + # below is what the same particle gets when advected on its own. np.testing.assert_allclose( [p.z for p in pset], - [0.666162, 0.8667131, 0.92150104, 0.9605109, 0.9577529, 1.0041442, 1.0284728, 1.0033542, 1.2949713, 1.3928112], + [0.66616201, 0.86671311, 0.92108649, 0.95940739, 0.95945358, 1.00413370, 1.02847278, 1.00335419, 1.27260256, 1.38021827], + rtol=1e-6, ) # fmt:skip diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 00a84c676..6d7abf2b2 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -22,6 +22,7 @@ XNearest, XPartialslip, ) +from parcels.interpolators._xinterpolators import _get_corner_data_Agrid from parcels.kernels import AdvectionRK4_3D from tests.utils import TEST_DATA @@ -201,6 +202,79 @@ def test_interpolation_mesh_type(mesh, npart=10): assert v == 0.0 +@pytest.fixture +def corner_gather_data() -> xr.DataArray: + rng = np.random.default_rng(0) + return xr.DataArray(rng.random((6, 5, 7, 8)), dims=("time", "depth", "lat", "lon")) + + +CORNER_GATHER_AXIS_DIM = {"X": "lon", "Y": "lat", "Z": "depth"} + + +def _agrid_corners(data, ti, zi, yi, xi, lenT, lenZ, npart): # noqa: N803 + return _get_corner_data_Agrid(data, ti, zi, yi, xi, lenT, lenZ, npart, CORNER_GATHER_AXIS_DIM) + + +@pytest.mark.parametrize("lenT", [1, 2]) +@pytest.mark.parametrize("lenZ", [1, 2]) +@pytest.mark.parametrize("uniform_clock", [True, False]) +def test_corner_gather_batch_matches_single_particle(corner_gather_data, lenT, lenZ, uniform_clock): # noqa: N803 + """Gathering a batch must give each particle what it would get on its own. + + Interpolation is per-particle, so batching cannot change a result. The index + arrays and the final reshape must therefore agree on where each particle sits + in the flat gather -- they did not, and particles whose time or depth index + differed from their batch-mates were gathered from the wrong level. + """ + rng = np.random.default_rng(1) + npart = 5 + if uniform_clock: + ti, zi = np.full(npart, 2), np.full(npart, 1) + else: + ti, zi = rng.integers(0, 4, npart), rng.integers(0, 3, npart) + yi, xi = rng.integers(0, 5, npart), rng.integers(0, 6, npart) + + batch = _agrid_corners(corner_gather_data, ti, zi, yi, xi, lenT, lenZ, npart) + + for p in range(npart): + single = _agrid_corners( + corner_gather_data, ti[p : p + 1], zi[p : p + 1], yi[p : p + 1], xi[p : p + 1], lenT, lenZ, 1 + ) + np.testing.assert_array_equal(batch[..., p], single[..., 0]) + + +def test_corner_gather_axes_are_ordered_t_z_y_x(corner_gather_data): + """The returned axes must mean (T, Z, Y, X, particle), as the interpolators assume.""" + rng = np.random.default_rng(2) + npart = 4 + ti, zi = rng.integers(0, 4, npart), rng.integers(0, 3, npart) + yi, xi = rng.integers(0, 5, npart), rng.integers(0, 6, npart) + + out = _agrid_corners(corner_gather_data, ti, zi, yi, xi, 2, 2, npart) + raw = corner_gather_data.values + + assert out.shape == (2, 2, 2, 2, npart) + for p in range(npart): + for it, iz, iy, ix in np.ndindex(2, 2, 2, 2): + assert out[it, iz, iy, ix, p] == raw[ti[p] + it, zi[p] + iz, yi[p] + iy, xi[p] + ix] + + +def test_corner_gather_keeps_axes_missing_from_the_mapping(): + """An axis absent from ``axis_dim`` is not indexed, but still shapes the result.""" + rng = np.random.default_rng(3) + data = xr.DataArray(rng.random((6, 1, 7, 8)), dims=("time", "depth", "lat", "lon")) + npart = 3 + ti = np.array([0, 2, 3]) + zi = np.zeros(npart, dtype=int) + yi, xi = rng.integers(0, 5, npart), rng.integers(0, 6, npart) + + out = _get_corner_data_Agrid(data, ti, zi, yi, xi, 2, 1, npart, {"X": "lon", "Y": "lat"}) + + assert out.shape == (2, 1, 2, 2, npart) + for p in range(npart): + assert out[0, 0, 0, 0, p] == data.values[ti[p], 0, yi[p], xi[p]] + + interp_methods = { "linear": XLinear, } From 523760978d1ea322adc58fba2b9fd34d2944f4b3 Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Wed, 5 Aug 2026 00:28:47 +0200 Subject: [PATCH 3/8] No changelog style docstrings --- tests/test_interpolation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 6d7abf2b2..5a34536f8 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -223,8 +223,7 @@ def test_corner_gather_batch_matches_single_particle(corner_gather_data, lenT, l Interpolation is per-particle, so batching cannot change a result. The index arrays and the final reshape must therefore agree on where each particle sits - in the flat gather -- they did not, and particles whose time or depth index - differed from their batch-mates were gathered from the wrong level. + in the flat gather. """ rng = np.random.default_rng(1) npart = 5 From 4c3d434f2739bd5337ff78e9845e71ff7f578f2f Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Wed, 5 Aug 2026 00:47:35 +0200 Subject: [PATCH 4/8] Simplify the corner index construction in _gather_corners The per-axis index arrays were laid out with manual stride arithmetic: a running `before`, a `math.prod` of the trailing counts, and a comment spelling out the flat position of each level. Build them from the output shape instead. Each axis's levels are placed in their slot of the (T, Z, Y, X, particle) grid, broadcast over the other slots and flattened, which is exactly the inverse of the reshape that ends the function. Same index arrays, no loop-carried state. Co-authored-by: Claude --- src/parcels/interpolators/_xinterpolators.py | 22 ++++++++------------ 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index 01de64d8f..b3b655e79 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -2,7 +2,6 @@ from __future__ import annotations -import math from typing import TYPE_CHECKING, Literal import numpy as np @@ -31,7 +30,7 @@ def _gather_corners( The gather is the outer product over ``_CORNER_AXES``, each axis contributing one or two levels, with **the particle index innermost**. That layout is - stated once here -- as ``counts`` -- and drives both the index arrays and the + stated once here -- as ``shape`` -- and drives both the index arrays and the shape of the result, so the two cannot drift apart. Parameters @@ -57,22 +56,19 @@ def _gather_corners( """ dims = {axis: ("time" if axis == "T" else axis_dim.get(axis)) for axis in _CORNER_AXES} counts = tuple(len(levels[axis]) if axis in levels else 1 for axis in _CORNER_AXES) + shape = (*counts, npart) selection_dict = {} - before = 1 # number of corner combinations on the axes outside this one - for i, (axis, n_levels) in enumerate(zip(_CORNER_AXES, counts, strict=True)): + for i, axis in enumerate(_CORNER_AXES): if axis in levels and dims[axis] is not None and dims[axis] in data.dims: - after = math.prod(counts[i + 1 :]) - # level i of particle p belongs at flat position - # ((i_before * n_levels + i) * after + i_after) * npart + p stacked = np.stack(np.broadcast_arrays(*levels[axis])) # (n_levels, npart) - selection_dict[dims[axis]] = xr.DataArray( - np.broadcast_to(stacked[None, :, None, :], (before, n_levels, after, npart)).reshape(-1), - dims=("points"), - ) - before *= n_levels + # Put the level axis in slot i of the corner grid, spread it over the + # other slots, and flatten -- the inverse of the reshape below. + other_slots = tuple(j for j in range(len(_CORNER_AXES)) if j != i) + in_slot_i = np.expand_dims(stacked, other_slots) + selection_dict[dims[axis]] = xr.DataArray(np.broadcast_to(in_slot_i, shape).reshape(-1), dims="points") - return data.isel(selection_dict).data.reshape(*counts, npart) + return data.isel(selection_dict).data.reshape(shape) def _get_corner_data_Agrid( From a080ff3678e2b0851ca49636070e74d3c7353925 Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Wed, 5 Aug 2026 00:53:54 +0200 Subject: [PATCH 5/8] Match house docstring style in the corner gather helpers Drop the bold emphasis, the double-hyphen asides and the semicolons from the docstrings and comments added with the corner gather rework, and give the numpydoc parameter entries the `name : type` form used elsewhere. Co-authored-by: Claude --- src/parcels/interpolators/_xinterpolators.py | 33 ++++++++++---------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index b3b655e79..498442f91 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -28,31 +28,32 @@ def _gather_corners( ) -> np.ndarray: """Gather field data at the corners bracketing each particle. - The gather is the outer product over ``_CORNER_AXES``, each axis contributing - one or two levels, with **the particle index innermost**. That layout is - stated once here -- as ``shape`` -- and drives both the index arrays and the - shape of the result, so the two cannot drift apart. + The gather is the outer product over ``_CORNER_AXES``, with each axis + contributing one or two levels and the particle index innermost. The local + variable ``shape`` states that layout, and is used both to build the index + arrays and to reshape the result. Parameters ---------- - data : + data : np.ndarray or xr.DataArray Field data, with dimensions ordered ``(time, Z, Y, X)``. - axis_dim : - Maps ``"X"``/``"Y"``/``"Z"`` to dimension names; ``"T"`` is ``"time"``. - levels : + axis_dim : dict + Maps ``"X"``, ``"Y"`` and ``"Z"`` to dimension names. The ``"T"`` + dimension is always named ``"time"``. + levels : dict Maps an axis to the per-particle index arrays of its corner levels, e.g. ``{"T": (ti, ti + 1), "Y": (yi, yi + 1)}``. An axis absent from ``levels`` contributes a single level. An axis the field has no - dimension for still shapes the result, but is not indexed -- its corners - simply repeat. - npart : + dimension for still shapes the result, but is not indexed, so its + corners repeat. + npart : int Number of particles. Returns ------- np.ndarray - Of shape ``(*counts, npart)``, where ``counts[i]`` is the number of - levels gathered on ``_CORNER_AXES[i]``. + Array of shape ``(*counts, npart)``, where ``counts[i]`` is the number + of levels gathered on ``_CORNER_AXES[i]``. """ dims = {axis: ("time" if axis == "T" else axis_dim.get(axis)) for axis in _CORNER_AXES} counts = tuple(len(levels[axis]) if axis in levels else 1 for axis in _CORNER_AXES) @@ -63,7 +64,7 @@ def _gather_corners( if axis in levels and dims[axis] is not None and dims[axis] in data.dims: stacked = np.stack(np.broadcast_arrays(*levels[axis])) # (n_levels, npart) # Put the level axis in slot i of the corner grid, spread it over the - # other slots, and flatten -- the inverse of the reshape below. + # other slots, and flatten. This is the inverse of the reshape below. other_slots = tuple(j for j in range(len(_CORNER_AXES)) if j != i) in_slot_i = np.expand_dims(stacked, other_slots) selection_dict[dims[axis]] = xr.DataArray(np.broadcast_to(in_slot_i, shape).reshape(-1), dims="points") @@ -246,8 +247,8 @@ def interp( def _compute_corner_data(data, y_levels, x_levels, z_levels=None) -> np.ndarray: """Gather the two bracketing face values and reduce over time if needed. - Exactly one of the Z/Y/X axes contributes the two corners; the others - contribute a single level. + Exactly one of the Z, Y and X axes contributes the two corners. The + other two contribute a single level each. """ levels = { "T": t_levels, From 91a7c5e3cb90a716d57a97d084b416c997738487 Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Wed, 5 Aug 2026 15:59:02 +0200 Subject: [PATCH 6/8] Drop obsolete comment Co-authored-by: Erik van Sebille --- tests/test_advection.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_advection.py b/tests/test_advection.py index a9d2c6708..e8eb62406 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -466,9 +466,6 @@ def test_nemo_3D_curvilinear_fieldset(kernel): np.testing.assert_allclose([p.z for p in pset], z_initial) elif kernel == AdvectionRK4_3D: # TODO check why decimals needs to be so low in RK4_3D (compare to v3) - # These particles sit at different depth levels, so the C-grid gather used to - # mix their zi indices; the previous expectations recorded that. Each value - # below is what the same particle gets when advected on its own. np.testing.assert_allclose( [p.z for p in pset], [0.66616201, 0.86671311, 0.92108649, 0.95940739, 0.95945358, 1.00413370, 1.02847278, 1.00335419, 1.27260256, 1.38021827], From c8372dcd52f097fe0c8445d55804a240facbc068 Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Wed, 5 Aug 2026 15:59:21 +0200 Subject: [PATCH 7/8] Update tests/test_interpolation.py Co-authored-by: Erik van Sebille --- tests/test_interpolation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 5a34536f8..4d9c246cb 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -243,7 +243,7 @@ def test_corner_gather_batch_matches_single_particle(corner_gather_data, lenT, l def test_corner_gather_axes_are_ordered_t_z_y_x(corner_gather_data): - """The returned axes must mean (T, Z, Y, X, particle), as the interpolators assume.""" + """The returned axes must be (T, Z, Y, X, particle), as the interpolators assume.""" rng = np.random.default_rng(2) npart = 4 ti, zi = rng.integers(0, 4, npart), rng.integers(0, 3, npart) From 096192e01b3080d823e9af8ba5c90ab8a0d7fadd Mon Sep 17 00:00:00 2001 From: Willi Rath Date: Wed, 5 Aug 2026 17:51:26 +0200 Subject: [PATCH 8/8] Address review comments on the corner gather helpers Make levels total over _CORNER_AXES so a missing axis fails loudly, and reword the levels docstring. Fix the annotations on both helpers: data is a DataArray or WindowedArray, never an ndarray, and the corner axes include "T". Drop the RK4_3D TODO and the rtol this branch had added, since the default tolerance passes. Co-authored-by: Claude --- src/parcels/interpolators/_xinterpolators.py | 40 ++++++++++---------- tests/test_advection.py | 2 - 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index 498442f91..651797538 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -13,17 +13,18 @@ from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator if TYPE_CHECKING: + from parcels._core._windowed_array import WindowedArray from parcels._core.field import Field, VectorField from parcels._core.xgrid import XGrid -_CORNER_AXES: tuple[ptyping.XgridAxis, ...] = ("T", "Z", "Y", "X") +_CORNER_AXES: tuple[ptyping.XgcmAxisDirection, ...] = ("T", "Z", "Y", "X") def _gather_corners( - data: np.ndarray | xr.DataArray, - axis_dim: dict[ptyping.ptyping.XgridAxis, str], - levels: dict[ptyping.XgridAxis, tuple[np.ndarray, ...]], + data: xr.DataArray | WindowedArray, + axis_dim: dict[ptyping.XgridAxis, str], + levels: dict[ptyping.XgcmAxisDirection, tuple[np.ndarray, ...]], npart: int, ) -> np.ndarray: """Gather field data at the corners bracketing each particle. @@ -35,17 +36,18 @@ def _gather_corners( Parameters ---------- - data : np.ndarray or xr.DataArray + data : xr.DataArray or WindowedArray Field data, with dimensions ordered ``(time, Z, Y, X)``. axis_dim : dict Maps ``"X"``, ``"Y"`` and ``"Z"`` to dimension names. The ``"T"`` dimension is always named ``"time"``. levels : dict - Maps an axis to the per-particle index arrays of its corner levels, e.g. - ``{"T": (ti, ti + 1), "Y": (yi, yi + 1)}``. An axis absent from - ``levels`` contributes a single level. An axis the field has no - dimension for still shapes the result, but is not indexed, so its - corners repeat. + Maps every axis in ``_CORNER_AXES`` to the per-particle index arrays of + its corner levels, e.g. ``{"T": (ti,), "Z": (zi,), "Y": (yi, yi + 1), + "X": (xi, xi + 1)}``. An axis bracketed by a single level is given a + one-tuple. If the field has no dimension for an axis, that axis still + contributes to the output shape, but is not indexed, so its corners + repeat. npart : int Number of particles. @@ -56,12 +58,12 @@ def _gather_corners( of levels gathered on ``_CORNER_AXES[i]``. """ dims = {axis: ("time" if axis == "T" else axis_dim.get(axis)) for axis in _CORNER_AXES} - counts = tuple(len(levels[axis]) if axis in levels else 1 for axis in _CORNER_AXES) + counts = tuple(len(levels[axis]) for axis in _CORNER_AXES) shape = (*counts, npart) selection_dict = {} for i, axis in enumerate(_CORNER_AXES): - if axis in levels and dims[axis] is not None and dims[axis] in data.dims: + if dims[axis] is not None and dims[axis] in data.dims: stacked = np.stack(np.broadcast_arrays(*levels[axis])) # (n_levels, npart) # Put the level axis in slot i of the corner grid, spread it over the # other slots, and flatten. This is the inverse of the reshape below. @@ -73,18 +75,18 @@ def _gather_corners( def _get_corner_data_Agrid( - data: np.ndarray | xr.DataArray, - ti: int, - zi: int, - yi: int, - xi: int, + data: xr.DataArray | WindowedArray, + ti: np.ndarray, + zi: np.ndarray, + yi: np.ndarray, + xi: np.ndarray, lenT: int, # noqa: N803 lenZ: int, # noqa: N803 npart: int, - axis_dim: dict[ptyping.ptyping.XgridAxis, str], + axis_dim: dict[ptyping.XgridAxis, str], ) -> np.ndarray: """Helper function to get the corner data for a given A-grid field and position.""" - levels = { + levels: dict[ptyping.XgcmAxisDirection, tuple[np.ndarray, ...]] = { "T": (ti,) if lenT == 1 else (ti, np.clip(ti + 1, 0, data.shape[0] - 1)), "Z": (zi,) if lenZ == 1 else (zi, np.clip(zi + 1, 0, data.shape[1] - 1)), "Y": (yi, np.clip(yi + 1, 0, data.shape[2] - 1)), diff --git a/tests/test_advection.py b/tests/test_advection.py index e8eb62406..09c6ae817 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -465,11 +465,9 @@ def test_nemo_3D_curvilinear_fieldset(kernel): if kernel == AdvectionRK4: np.testing.assert_allclose([p.z for p in pset], z_initial) elif kernel == AdvectionRK4_3D: - # TODO check why decimals needs to be so low in RK4_3D (compare to v3) np.testing.assert_allclose( [p.z for p in pset], [0.66616201, 0.86671311, 0.92108649, 0.95940739, 0.95945358, 1.00413370, 1.02847278, 1.00335419, 1.27260256, 1.38021827], - rtol=1e-6, ) # fmt:skip