diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index 54606c4b0..651797538 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -13,56 +13,86 @@ 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 -def _get_corner_data_Agrid( - data: np.ndarray | xr.DataArray, - ti: int, - zi: int, - yi: int, - xi: int, - lenT: int, # noqa: N803 - lenZ: int, # noqa: N803 +_CORNER_AXES: tuple[ptyping.XgcmAxisDirection, ...] = ("T", "Z", "Y", "X") + + +def _gather_corners( + data: xr.DataArray | WindowedArray, + axis_dim: dict[ptyping.XgridAxis, str], + levels: dict[ptyping.XgcmAxisDirection, tuple[np.ndarray, ...]], npart: int, - 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) + """Gather field data at the corners bracketing each particle. + + 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 : 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 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. + + Returns + ------- + np.ndarray + 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]) for axis in _CORNER_AXES) + shape = (*counts, npart) - # 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) + selection_dict = {} + for i, axis in enumerate(_CORNER_AXES): + 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. + 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") - # 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) + return data.isel(selection_dict).data.reshape(shape) - # 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) +def _get_corner_data_Agrid( + 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.XgridAxis, str], +) -> np.ndarray: + """Helper function to get the corner data for a given A-grid field and position.""" + 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)), + "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 +243,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 and X axes contributes the two corners. The + other two contribute a single level each. + """ + 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 +268,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 +315,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 +364,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..09c6ae817 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -465,10 +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.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], ) # fmt:skip diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 00a84c676..4d9c246cb 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,78 @@ 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. + """ + 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 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) + 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, } 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))