From f0a9c10cad8370782bc99867f474d2826c36a93d Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 16 Jun 2026 09:21:08 +0200 Subject: [PATCH 01/12] Move AnalyticalAdvection unit test to v4 --- tests-v3/test_advection.py | 55 -------------------------------------- tests/test_advection.py | 43 +++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 55 deletions(-) diff --git a/tests-v3/test_advection.py b/tests-v3/test_advection.py index 3d8f06bac3..41ea495d96 100644 --- a/tests-v3/test_advection.py +++ b/tests-v3/test_advection.py @@ -57,58 +57,3 @@ def test_advection_2DCROCO(): pset.execute([AdvectionRK4], runtime=runtime, dt=100) assert np.allclose(pset.depth, Z.flatten(), atol=1e-3) assert np.allclose(pset.lon_nextloop, [x + runtime for x in X], atol=1e-3) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_analyticalAgrid(): - lon = np.arange(0, 15, dtype=np.float32) - lat = np.arange(0, 15, dtype=np.float32) - U = np.ones((lat.size, lon.size), dtype=np.float32) - V = np.ones((lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") - pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) - - with pytest.raises(NotImplementedError): - pset.execute(AdvectionAnalytical, runtime=1) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1927") -@pytest.mark.parametrize("u", [1, -0.2, -0.3, 0]) -@pytest.mark.parametrize("v", [1, -0.3, 0, -1]) -@pytest.mark.parametrize("w", [None, 1, -0.3, 0, -1]) -@pytest.mark.parametrize("direction", [1, -1]) -def test_uniform_analytical(u, v, w, direction, tmp_zarrfile): - lon = np.arange(0, 15, dtype=np.float32) - lat = np.arange(0, 15, dtype=np.float32) - if w is not None: - depth = np.arange(0, 40, 2, dtype=np.float32) - U = u * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) - V = v * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) - W = w * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V, "W": W}, {"lon": lon, "lat": lat, "depth": depth}, mesh="flat") - fieldset.W.interp_method = "cgrid_velocity" - else: - U = u * np.ones((lat.size, lon.size), dtype=np.float32) - V = v * np.ones((lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - - x0, y0, z0 = 6.1, 6.2, 20 - pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, depth=z0) - - outfile = pset.ParticleFile(name=tmp_zarrfile, outputdt=1, chunks=(1, 1)) - pset.execute(AdvectionAnalytical, runtime=4, dt=direction, output_file=outfile) - assert np.abs(pset.lon - x0 - pset.time * u) < 1e-6 - assert np.abs(pset.lat - y0 - pset.time * v) < 1e-6 - if w is not None: - assert np.abs(pset.depth - z0 - pset.time * w) < 1e-4 - - ds = xr.open_zarr(tmp_zarrfile) - times = (direction * ds["time"][:]).values.astype("timedelta64[s]")[0] - timeref = np.arange(1, 5).astype("timedelta64[s]") - assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) - lons = ds["lon"][:].values - assert np.allclose(lons, x0 + direction * u * np.arange(1, 5)) diff --git a/tests/test_advection.py b/tests/test_advection.py index 17511e82e2..c63261cb28 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -25,6 +25,7 @@ ) from parcels._datasets.structured.generic import datasets_sgrid from parcels.kernels import ( + AdvectionAnalytical, AdvectionDiffusionEM, AdvectionDiffusionM1, AdvectionEE, @@ -488,3 +489,45 @@ def test_mitgcm(): 28349.16658616, ] np.testing.assert_allclose(pset.lon, lon_v3, atol=10) + + +def test_analyticalAgrid(): + lon = np.arange(0, 15, dtype=np.float32) + lat = np.arange(0, 15, dtype=np.float32) + U = np.ones((lat.size, lon.size), dtype=np.float32) + V = np.ones((lat.size, lon.size), dtype=np.float32) + fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") + pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) + + with pytest.raises(NotImplementedError): + pset.execute(AdvectionAnalytical, runtime=1) + + +@pytest.mark.parametrize("u", [1, -0.2, -0.3, 0]) +@pytest.mark.parametrize("v", [1, -0.3, 0, -1]) +@pytest.mark.parametrize("w", [None, 1, -0.3, 0, -1]) +@pytest.mark.parametrize("direction", [1, -1]) +def test_uniform_analytical(u, v, w, direction, tmp_parquet): + ds = simple_UV_dataset(mesh="flat") + ds["U"].data[:] = u + ds["V"].data[:] = v + if w is not None: + ds["W"].data[:] = w + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + + x0, y0, z0 = 6.1, 6.2, 20 + pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, depth=z0) + + outfile = pset.ParticleFile(tmp_parquet, outputdt=1, chunks=(1, 1)) + pset.execute(AdvectionAnalytical, runtime=4, dt=direction, output_file=outfile) + assert np.abs(pset.lon - x0 - pset.time * u) < 1e-6 + assert np.abs(pset.lat - y0 - pset.time * v) < 1e-6 + if w is not None: + assert np.abs(pset.depth - z0 - pset.time * w) < 1e-4 + + df = pd.read_parquet(tmp_parquet) + times = (direction * df["time"]).values.astype("timedelta64[s]")[0] + timeref = np.arange(1, 5).astype("timedelta64[s]") + assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) + lons = df["lon"].values + assert np.allclose(lons, x0 + direction * u * np.arange(1, 5)) From feff261beafed12f6a841645160e6cb8b3a61b46 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 17 Jun 2026 10:12:11 +0200 Subject: [PATCH 02/12] First vectorization of AnalyticalAdvection --- src/parcels/_core/kernel.py | 14 ++- src/parcels/kernels/_advection.py | 136 ++++++++++++++++-------------- tests/test_advection.py | 49 ++++++----- 3 files changed, 105 insertions(+), 94 deletions(-) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 83259228f0..8800c859e3 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -6,7 +6,6 @@ import numpy as np -from parcels._core.basegrid import GridType from parcels._core.statuscodes import ( StatusCode, _raise_field_interpolation_error, @@ -19,7 +18,6 @@ from parcels._core.warnings import KernelWarning from parcels._python import assert_same_function_signature from parcels.kernels import ( - AdvectionAnalytical, AdvectionRK4, AdvectionRK45, ) @@ -126,12 +124,12 @@ def check_fieldsets_in_kernels(self, kernel): # TODO v4: this can go into anoth This function is to be called from the derived class when setting up the 'kernel'. """ if self.fieldset is not None: - if kernel is AdvectionAnalytical: - if self._fieldset.U.interp_method != "cgrid_velocity": - raise NotImplementedError("Analytical Advection only works with C-grids") - if self._fieldset.U.grid._gtype not in [GridType.CurvilinearZGrid, GridType.RectilinearZGrid]: - raise NotImplementedError("Analytical Advection only works with Z-grids in the vertical") - elif kernel is AdvectionRK45: + # if kernel is AdvectionAnalytical: + # if self._fieldset.U.interp_method != "cgrid_velocity": + # raise NotImplementedError("Analytical Advection only works with C-grids") + # if self._fieldset.U.grid._gtype not in [GridType.CurvilinearZGrid, GridType.RectilinearZGrid]: + # raise NotImplementedError("Analytical Advection only works with Z-grids in the vertical") + if kernel is AdvectionRK45: if "next_dt" not in [v.name for v in self.ptype.variables]: raise ValueError('ParticleClass requires a "next_dt" for AdvectionRK45 Kernel.') if not hasattr(self.fieldset, "RK45_tol"): diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index 792c8ee17f..9ab0af83ce 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -1,7 +1,5 @@ """Collection of pre-built advection kernels.""" -import math - import numpy as np from parcels._core.statuscodes import StatusCode @@ -182,24 +180,31 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover Note that the time-dependent scheme is currently implemented with 'intermediate timesteps' (default 10 per model timestep) and not yet with the full analytical time integration. """ - import numpy as np - import parcels._core.utils.interpolation as i_u + from parcels._core.field import _get_positions tol = 1e-10 - I_s = 10 # number of intermediate time steps + # I_s = 10 # number of intermediate time steps dt = particles.dt direction = 1.0 if dt > 0 else -1.0 withW = True if "W" in [f.name for f in fieldset.fields.values()] else False withTime = True if len(fieldset.U.grid.time) > 1 else False - tau, zeta, eta, xsi, ti, zi, yi, xi = fieldset.U._search_indices( - particles.z, particles.lat, particles.lon, particles=particles + + igrid = fieldset.U.igrid + + _, grid_positions = _get_positions( + fieldset.U, particles.time, particles.z, particles.lat, particles.lon, particles, particles.ei[:, igrid] ) - ds_t = dt - if withTime: - time_i = np.linspace(0, fieldset.U.grid.time[ti + 1] - fieldset.U.grid.time[ti], I_s) - ds_t = min(ds_t, time_i[np.where(particles.time - fieldset.U.grid.time[ti] < time_i)[0][0]]) + xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + zi, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] + ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] + ds_t = dt + # if withTime: + # time_i = np.linspace(0, fieldset.U.grid.time[ti + 1] - fieldset.U.grid.time[ti], I_s) + # ds_t = min(ds_t, time_i[np.where(particles.time - fieldset.U.grid.time[ti] < time_i)[0][0]]) + # print("withTime", withTime, "ds_t", ds_t, dt) if withW: if abs(xsi - 1) < tol: if fieldset.U.data[0, zi + 1, yi + 1, xi + 1] > 0: @@ -215,26 +220,26 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover zeta = 0 else: if abs(xsi - 1) < tol: - if fieldset.U.data[0, yi + 1, xi + 1] > 0: + if fieldset.U.data[0, 0, yi + 1, xi + 1] > 0: xi += 1 xsi = 0 if abs(eta - 1) < tol: - if fieldset.V.data[0, yi + 1, xi + 1] > 0: + if fieldset.V.data[0, 0, yi + 1, xi + 1] > 0: yi += 1 eta = 0 - particles.ei[:] = fieldset.U.ravel_index(zi, yi, xi) + # particles.ei[:, igrid] = fieldset.U.grid.ravel_index(zi, yi, xi) - grid = fieldset.U.grid - if grid._gtype < 2: + grid = fieldset.UV.grid + if grid.lon.ndim == 1: px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) else: px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) - if grid.mesh == "spherical": - px[0] = px[0] + 360 if px[0] < particles.lon - 225 else px[0] - px[0] = px[0] - 360 if px[0] > particles.lat + 225 else px[0] + + if grid._mesh == "spherical": + px = ((px + 180.0) % 360.0) - 180.0 px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) if withW: @@ -242,14 +247,22 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover dz = pz[1] - pz[0] else: dz = 1.0 + c1 = i_u._geodetic_distance( + py[0], py[1], px[0], px[1], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(0.0, xsi), py) + ) + c2 = i_u._geodetic_distance( + py[1], py[2], px[1], px[2], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 1.0), py) + ) + c3 = i_u._geodetic_distance( + py[2], py[3], px[2], px[3], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(1.0, xsi), py) + ) + c4 = i_u._geodetic_distance( + py[3], py[0], px[3], px[0], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 0.0), py) + ) - c1 = i_u._geodetic_distance(py[0], py[1], px[0], px[1], grid.mesh, np.dot(i_u.phi2D_lin(0.0, xsi), py)) - c2 = i_u._geodetic_distance(py[1], py[2], px[1], px[2], grid.mesh, np.dot(i_u.phi2D_lin(eta, 1.0), py)) - c3 = i_u._geodetic_distance(py[2], py[3], px[2], px[3], grid.mesh, np.dot(i_u.phi2D_lin(1.0, xsi), py)) - c4 = i_u._geodetic_distance(py[3], py[0], px[3], px[0], grid.mesh, np.dot(i_u.phi2D_lin(eta, 0.0), py)) rad = np.pi / 180.0 deg2m = 1852 * 60.0 - meshJac = (deg2m * deg2m * math.cos(rad * particles.lat)) if grid.mesh == "spherical" else 1 + meshJac = (deg2m * deg2m * np.cos(rad * particles.lat)) if grid._mesh == "spherical" else 1 dxdy = i_u._compute_jacobian_determinant(py, px, eta, xsi) * meshJac if withW: @@ -263,40 +276,39 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover V0 = V0 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, zi + 1, yi, xi + 1] * c1 * dz V1 = V1 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, zi + 1, yi + 1, xi + 1] * c3 * dz else: - U0 = direction * fieldset.U.data[ti, yi + 1, xi] * c4 * dz - U1 = direction * fieldset.U.data[ti, yi + 1, xi + 1] * c2 * dz - V0 = direction * fieldset.V.data[ti, yi, xi + 1] * c1 * dz - V1 = direction * fieldset.V.data[ti, yi + 1, xi + 1] * c3 * dz + U0 = (direction * fieldset.U.data[ti, 0, yi + 1, xi] * c4 * dz).values.flatten() + U1 = (direction * fieldset.U.data[ti, 0, yi + 1, xi + 1] * c2 * dz).values.flatten() + V0 = (direction * fieldset.V.data[ti, 0, yi, xi + 1] * c1 * dz).values.flatten() + V1 = (direction * fieldset.V.data[ti, 0, yi + 1, xi + 1] * c3 * dz).values.flatten() if withTime: - U0 = U0 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, yi + 1, xi] * c4 * dz - U1 = U1 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, yi + 1, xi + 1] * c2 * dz - V0 = V0 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, yi, xi + 1] * c1 * dz - V1 = V1 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, yi + 1, xi + 1] * c3 * dz + U0 = U0 * (1 - tau) + (tau * direction * fieldset.U.data[ti + 1, 0, yi + 1, xi] * c4 * dz).values.flatten() + U1 = ( + U1 * (1 - tau) + + (tau * direction * fieldset.U.data[ti + 1, 0, yi + 1, xi + 1] * c2 * dz).values.flatten() + ) + V0 = V0 * (1 - tau) + (tau * direction * fieldset.V.data[ti + 1, 0, yi, xi + 1] * c1 * dz).values.flatten() + V1 = ( + V1 * (1 - tau) + + (tau * direction * fieldset.V.data[ti + 1, 0, yi + 1, xi + 1] * c3 * dz).values.flatten() + ) def compute_ds(F0, F1, r, direction, tol): # noqa: N803 - up = F0 * (1 - r) + F1 * r - r_target = 1.0 if direction * up >= 0.0 else 0.0 - B = F0 - F1 - delta = -F0 - B = 0 if abs(B) < tol else B - - if abs(B) > tol: - F_r1 = r_target + delta / B - F_r0 = r + delta / B - else: - F_r0, F_r1 = None, None - - if abs(B) < tol and abs(delta) < tol: - ds = float("inf") - elif B == 0: - ds = -(r_target - r) / delta - elif F_r1 * F_r0 < tol: - ds = float("inf") - else: - ds = -1.0 / B * math.log(F_r1 / F_r0) - - if abs(ds) < tol: - ds = float("inf") + with np.errstate(divide="ignore", invalid="ignore"): + up = F0 * (1 - r) + F1 * r + r_target = np.where(direction * up >= 0.0, 1.0, 0.0) + B = F0 - F1 + delta = -F0 + B = np.where(np.abs(B) < tol, np.zeros_like(B), B) + + F_r1 = np.where(np.abs(B) > tol, r_target + delta / B, np.nan) + F_r0 = np.where(np.abs(B) > tol, r + delta / B, np.nan) + + ds = -1.0 / B * np.log(F_r1 / F_r0) + ds = np.where(F_r1 * F_r0 < tol, np.inf, ds) + ds = np.where(B == 0, -delta * direction / up, ds) + ds = np.where((np.abs(B) < tol) & (np.abs(delta) < tol), np.inf, ds) + + ds = np.where(ds < tol, np.inf, ds) return ds, B, delta ds_x, B_x, delta_x = compute_ds(U0, U1, xsi, direction, tol) @@ -307,19 +319,17 @@ def compute_ds(F0, F1, r, direction, tol): # noqa: N803 if withTime: W0 = W0 * (1 - tau) + tau * direction * fieldset.W.data[ti + 1, zi, yi + 1, xi + 1] * dxdy W1 = W1 * (1 - tau) + tau * direction * fieldset.W.data[ti + 1, zi + 1, yi + 1, xi + 1] * dxdy - ds_z, B_z, delta_z = compute_ds(W0, W1, zeta, direction, tol) + ds_z, B_z, delta_z = compute_ds(W0.values.flatten(), W1.values.flatten(), zeta, direction, tol) else: - ds_z = float("inf") + ds_z = np.inf # take the minimum travel time s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(ds_t / (dxdy * dz))) # calculate end position in time s_min def compute_rs(r, B, delta, s_min): # noqa: N803 - if abs(B) < tol: - return -delta * s_min + r - else: - return (r + delta / B) * math.exp(-B * s_min) - delta / B + with np.errstate(divide="ignore", invalid="ignore"): + return np.where(abs(B) < tol, -delta * s_min + r, (r + delta / B) * np.exp(-B * s_min) - delta / B) rs_x = compute_rs(xsi, B_x, delta_x, s_min) rs_y = compute_rs(eta, B_y, delta_y, s_min) @@ -344,6 +354,6 @@ def compute_rs(r, B, delta, s_min): # noqa: N803 particles.dz += (1.0 - rs_z) * pz[0] + rs_z * pz[1] - particles.z if particles.dt > 0: - particles.dt = max(direction * s_min * (dxdy * dz), 1e-7).astype("timedelta64[s]") + particles.dt = max(direction * s_min * (dxdy * dz), 1e-7) else: - particles.dt = min(direction * s_min * (dxdy * dz), -1e-7).astype("timedelta64[s]") + particles.dt = min(direction * s_min * (dxdy * dz), -1e-7) diff --git a/tests/test_advection.py b/tests/test_advection.py index c63261cb28..5794854950 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -491,22 +491,22 @@ def test_mitgcm(): np.testing.assert_allclose(pset.lon, lon_v3, atol=10) -def test_analyticalAgrid(): - lon = np.arange(0, 15, dtype=np.float32) - lat = np.arange(0, 15, dtype=np.float32) - U = np.ones((lat.size, lon.size), dtype=np.float32) - V = np.ones((lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") - pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) - - with pytest.raises(NotImplementedError): - pset.execute(AdvectionAnalytical, runtime=1) - - -@pytest.mark.parametrize("u", [1, -0.2, -0.3, 0]) -@pytest.mark.parametrize("v", [1, -0.3, 0, -1]) -@pytest.mark.parametrize("w", [None, 1, -0.3, 0, -1]) -@pytest.mark.parametrize("direction", [1, -1]) +# def test_analyticalAgrid(): +# lon = np.arange(0, 15, dtype=np.float32) +# lat = np.arange(0, 15, dtype=np.float32) +# U = np.ones((lat.size, lon.size), dtype=np.float32) +# V = np.ones((lat.size, lon.size), dtype=np.float32) +# fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") +# pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) + +# with pytest.raises(NotImplementedError): +# pset.execute(AdvectionAnalytical, runtime=1) + + +@pytest.mark.parametrize("u", [0.2, -0.3, 0]) +@pytest.mark.parametrize("v", [-0.3, 0, 1]) +@pytest.mark.parametrize("w", [None]) # , 1, -0.3, 0, -1]) +@pytest.mark.parametrize("direction", [1]) # , -1]) def test_uniform_analytical(u, v, w, direction, tmp_parquet): ds = simple_UV_dataset(mesh="flat") ds["U"].data[:] = u @@ -514,20 +514,23 @@ def test_uniform_analytical(u, v, w, direction, tmp_parquet): if w is not None: ds["W"].data[:] = w fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + print(fieldset.U.data) - x0, y0, z0 = 6.1, 6.2, 20 - pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, depth=z0) + x0, y0, z0 = 6.1, 6.2, 0.5 + pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, z=z0) - outfile = pset.ParticleFile(tmp_parquet, outputdt=1, chunks=(1, 1)) - pset.execute(AdvectionAnalytical, runtime=4, dt=direction, output_file=outfile) + outfile = ParticleFile(tmp_parquet, outputdt=np.timedelta64(1, "s")) + pset.execute( + AdvectionAnalytical, runtime=np.timedelta64(4, "s"), dt=np.timedelta64(direction, "s"), output_file=outfile + ) assert np.abs(pset.lon - x0 - pset.time * u) < 1e-6 assert np.abs(pset.lat - y0 - pset.time * v) < 1e-6 if w is not None: assert np.abs(pset.depth - z0 - pset.time * w) < 1e-4 df = pd.read_parquet(tmp_parquet) - times = (direction * df["time"]).values.astype("timedelta64[s]")[0] - timeref = np.arange(1, 5).astype("timedelta64[s]") + times = (direction * df["time"]).values.astype("timedelta64[s]") + timeref = np.arange(0, 5).astype("timedelta64[s]") assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) lons = df["lon"].values - assert np.allclose(lons, x0 + direction * u * np.arange(1, 5)) + assert np.allclose(lons, x0 + direction * u * np.arange(0, 5)) From 23a95abf41c107bbe27d7d6ccfe619555909f488 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 30 Jul 2026 17:15:42 +0200 Subject: [PATCH 03/12] Fixing merge errors --- src/parcels/kernels/_advection.py | 6 +++--- tests/test_advection.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index 3f8b5db4b0..4d6c7c90d1 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -173,7 +173,7 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover igrid = fieldset.U.igrid _, grid_positions = _get_positions( - fieldset.U, particles.time, particles.z, particles.lat, particles.lon, particles, particles.ei[:, igrid] + fieldset.U, particles.t, particles.z, particles.y, particles.x, particles, particles.ei[:, igrid] ) xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] @@ -218,7 +218,7 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) - if grid._mesh == "spherical": + if grid._mesh.is_spherical(): px = ((px + 180.0) % 360.0) - 180.0 px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) @@ -242,7 +242,7 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover rad = np.pi / 180.0 deg2m = 1852 * 60.0 - meshJac = (deg2m * deg2m * np.cos(rad * particles.lat)) if grid._mesh == "spherical" else 1 + meshJac = (deg2m * deg2m * np.cos(rad * particles.y)) if grid._mesh.is_spherical() else 1 dxdy = i_u._compute_jacobian_determinant(py, px, eta, xsi) * meshJac if withW: diff --git a/tests/test_advection.py b/tests/test_advection.py index cf442690f9..119b4511d1 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -519,20 +519,20 @@ def test_uniform_analytical(u, v, w, direction, tmp_parquet): print(fieldset.U.data) x0, y0, z0 = 6.1, 6.2, 0.5 - pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, z=z0) + pset = ParticleSet(fieldset, pclass=Particle, x=x0, y=y0, z=z0) outfile = ParticleFile(tmp_parquet, outputdt=np.timedelta64(1, "s")) pset.execute( AdvectionAnalytical, runtime=np.timedelta64(4, "s"), dt=np.timedelta64(direction, "s"), output_file=outfile ) - assert np.abs(pset.lon - x0 - pset.time * u) < 1e-6 - assert np.abs(pset.lat - y0 - pset.time * v) < 1e-6 + assert np.abs(pset.x - x0 - pset.t * u) < 1e-6 + assert np.abs(pset.y - y0 - pset.t * v) < 1e-6 if w is not None: - assert np.abs(pset.depth - z0 - pset.time * w) < 1e-4 + assert np.abs(pset.depth - z0 - pset.t * w) < 1e-4 df = pd.read_parquet(tmp_parquet) - times = (direction * df["time"]).values.astype("timedelta64[s]") + times = (direction * df["t"]).values.astype("timedelta64[s]") timeref = np.arange(0, 5).astype("timedelta64[s]") assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) - lons = df["lon"].values + lons = df["x"].values assert np.allclose(lons, x0 + direction * u * np.arange(0, 5)) From 4816462041063740dbc249204832ff758fc30d0c Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 30 Jul 2026 17:25:13 +0200 Subject: [PATCH 04/12] Also testing backward time integration --- tests/test_advection.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/test_advection.py b/tests/test_advection.py index 119b4511d1..72e52c74df 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -508,7 +508,7 @@ def test_mitgcm(): @pytest.mark.parametrize("u", [0.2, -0.3, 0]) @pytest.mark.parametrize("v", [-0.3, 0, 1]) @pytest.mark.parametrize("w", [None]) # , 1, -0.3, 0, -1]) -@pytest.mark.parametrize("direction", [1]) # , -1]) +@pytest.mark.parametrize("direction", [1, -1]) def test_uniform_analytical(u, v, w, direction, tmp_parquet): ds = simple_UV_dataset(mesh="flat") ds["U"].data[:] = u @@ -516,22 +516,25 @@ def test_uniform_analytical(u, v, w, direction, tmp_parquet): if w is not None: ds["W"].data[:] = w fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") - print(fieldset.U.data) x0, y0, z0 = 6.1, 6.2, 0.5 pset = ParticleSet(fieldset, pclass=Particle, x=x0, y=y0, z=z0) outfile = ParticleFile(tmp_parquet, outputdt=np.timedelta64(1, "s")) + runtime = 4 pset.execute( - AdvectionAnalytical, runtime=np.timedelta64(4, "s"), dt=np.timedelta64(direction, "s"), output_file=outfile + AdvectionAnalytical, + runtime=np.timedelta64(runtime, "s"), + dt=np.timedelta64(direction, "s"), + output_file=outfile, ) - assert np.abs(pset.x - x0 - pset.t * u) < 1e-6 - assert np.abs(pset.y - y0 - pset.t * v) < 1e-6 + assert np.abs(pset.x - x0 - runtime * u * direction) < 1e-6 + assert np.abs(pset.y - y0 - runtime * v * direction) < 1e-6 if w is not None: - assert np.abs(pset.depth - z0 - pset.t * w) < 1e-4 + assert np.abs(pset.depth - z0 - runtime * w * direction) < 1e-4 df = pd.read_parquet(tmp_parquet) - times = (direction * df["t"]).values.astype("timedelta64[s]") + times = (direction * (df["t"] - df["t"][0])).values.astype("timedelta64[s]") timeref = np.arange(0, 5).astype("timedelta64[s]") assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) lons = df["x"].values From db5f2b7fa2bf4ab08a48f5e15d725f296808aec7 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 12 Aug 2026 16:38:59 +0200 Subject: [PATCH 05/12] Separating data collection from Cgrid interpolator As that part of the code can be reused for Analytical Advection scheme --- src/parcels/interpolators/_xinterpolators.py | 248 ++++++++++--------- 1 file changed, 134 insertions(+), 114 deletions(-) diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index 54606c4b04..45d9ebec13 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -159,6 +159,135 @@ def interp( return u, v, w +def _get_cgrid_velocities( + vectorfield: VectorField, grid_positions: dict[ptyping.XgridAxis, dict[str, int | float | np.ndarray]] +): + # Helper function to get the edge velocities for a given C-grid vector field and position. + xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + zi, _ = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] + ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] + + U = vectorfield.U.data + V = vectorfield.V.data + grid = vectorfield.grid + offsets = _get_offsets_dictionary(grid) + tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3] + lenT = 2 if np.any(tau > 0) else 1 + + if grid.lon.ndim == 1: + px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) + py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) + else: + px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) + py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) + + if grid._mesh.is_spherical(): + px = ((px + 180.0) % 360.0) - 180.0 + px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) + px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) + c1 = i_u._geodetic_distance( + py[0], py[1], px[0], px[1], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(0.0, xsi), py), grid.deg2m + ) + c2 = i_u._geodetic_distance( + py[1], py[2], px[1], px[2], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 1.0), py), grid.deg2m + ) + c3 = i_u._geodetic_distance( + py[2], py[3], px[2], px[3], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(1.0, xsi), py), grid.deg2m + ) + c4 = i_u._geodetic_distance( + 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")), + } + + # 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)) + + if lenT == 2: + tau_full = tau[np.newaxis, :] + corner_data = corner_data[0, :] * (1 - tau_full) + corner_data[1, :] * tau_full + else: + corner_data = corner_data[0, :] + return corner_data + + # Compute U velocity + 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) + + U0 = corner_data[0, :] * c4 + U1 = corner_data[1, :] * c2 + + # Compute V velocity + 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) + + V0 = corner_data[0, :] * c1 + V1 = corner_data[1, :] * c3 + + if vectorfield.W: + W = vectorfield.W.data + + # Y coordinates: yi+offset for each spatial point, repeated for time + 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) + W0 = corner_data[0, :] + W1 = corner_data[1, :] + else: + W0 = np.zeros_like(U0) + W1 = np.zeros_like(U1) + + return U0, U1, V0, V1, W0, W1, px, py + + class CGrid_Velocity(VectorInterpolator): # noqa: N801 """ Interpolation kernel for velocity fields on a C-Grid. @@ -177,109 +306,13 @@ def interp( Following Delandmeter and Van Sebille (2019), velocity fields should be interpolated only in the direction of the grid cell faces. """ - xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] - yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] - zi, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] - ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] - - U = vectorfield.U.data - V = vectorfield.V.data + _, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + _, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + _, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] grid = vectorfield.grid - offsets = _get_offsets_dictionary(grid) - tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3] - lenT = 2 if np.any(tau > 0) else 1 - - if grid.lon.ndim == 1: - px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) - py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) - else: - px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) - py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) - - if grid._mesh.is_spherical(): - px = ((px + 180.0) % 360.0) - 180.0 - px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) - px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) - c1 = i_u._geodetic_distance( - py[0], py[1], px[0], px[1], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(0.0, xsi), py), grid.deg2m - ) - c2 = i_u._geodetic_distance( - py[1], py[2], px[1], px[2], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 1.0), py), grid.deg2m - ) - c3 = i_u._geodetic_distance( - py[2], py[3], px[2], px[3], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(1.0, xsi), py), grid.deg2m - ) - c4 = i_u._geodetic_distance( - 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")), - } - - # 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)) - - if lenT == 2: - tau_full = tau[np.newaxis, :] - corner_data = corner_data[0, :] * (1 - tau_full) + corner_data[1, :] * tau_full - else: - corner_data = corner_data[0, :] - return corner_data - - # Compute U velocity - 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) - U0 = corner_data[0, :] * c4 - U1 = corner_data[1, :] * c2 + U0, U1, V0, V1, W0, W1, px, py = _get_cgrid_velocities(vectorfield, grid_positions) Uvel = (1 - xsi) * U0 + xsi * U1 - - # Compute V velocity - 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) - - V0 = corner_data[0, :] * c1 - V1 = corner_data[1, :] * c3 Vvel = (1 - eta) * V0 + eta * V1 if grid._mesh.is_spherical(): @@ -309,20 +342,7 @@ def _compute_corner_data(data, selection_dict) -> np.ndarray: v /= conversion if vectorfield.W: - W = vectorfield.W.data - - # Y coordinates: yi+offset for each spatial point, repeated for time - 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) - - w = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta + w = W0 * (1 - zeta) + W1 * zeta if is_dask_collection(w): w = w.compute() else: From 4e771b687c40596c0e0b62b995bca9dc7f5ecadb Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 12 Aug 2026 16:39:28 +0200 Subject: [PATCH 06/12] Support for vertical analytical advection --- tests/test_advection.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_advection.py b/tests/test_advection.py index 72e52c74df..d1cf5e6088 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -507,14 +507,14 @@ def test_mitgcm(): @pytest.mark.parametrize("u", [0.2, -0.3, 0]) @pytest.mark.parametrize("v", [-0.3, 0, 1]) -@pytest.mark.parametrize("w", [None]) # , 1, -0.3, 0, -1]) +@pytest.mark.parametrize("w", [None, 0.05, 0, -0.05]) @pytest.mark.parametrize("direction", [1, -1]) def test_uniform_analytical(u, v, w, direction, tmp_parquet): ds = simple_UV_dataset(mesh="flat") ds["U"].data[:] = u ds["V"].data[:] = v if w is not None: - ds["W"].data[:] = w + ds["W"] = xr.full_like(ds["U"], w) fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") x0, y0, z0 = 6.1, 6.2, 0.5 @@ -531,7 +531,7 @@ def test_uniform_analytical(u, v, w, direction, tmp_parquet): assert np.abs(pset.x - x0 - runtime * u * direction) < 1e-6 assert np.abs(pset.y - y0 - runtime * v * direction) < 1e-6 if w is not None: - assert np.abs(pset.depth - z0 - runtime * w * direction) < 1e-4 + assert np.abs(pset.z - z0 - runtime * w * direction) < 1e-4 df = pd.read_parquet(tmp_parquet) times = (direction * (df["t"] - df["t"][0])).values.astype("timedelta64[s]") From e42942ce42ebd6fac026b2fc7b1e717cc04e0db1 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 12 Aug 2026 16:39:51 +0200 Subject: [PATCH 07/12] Support for W in analytical advection --- src/parcels/kernels/_advection.py | 121 ++++++------------------------ 1 file changed, 24 insertions(+), 97 deletions(-) diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index 4d6c7c90d1..ee0e6f69af 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -3,6 +3,7 @@ import numpy as np from parcels._core.statuscodes import StatusCode +from parcels.interpolators._xinterpolators import _get_cgrid_velocities __all__ = [ "AdvectionAnalytical", @@ -168,112 +169,43 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover dt = particles.dt direction = 1.0 if dt > 0 else -1.0 withW = True if "W" in [f.name for f in fieldset.fields.values()] else False - withTime = True if len(fieldset.U.grid.time) > 1 else False - igrid = fieldset.U.igrid + vectorfield = fieldset.UVW if withW else fieldset.UV + # withTime = True if len(vectorfield.grid.time) > 1 else False + igrid = vectorfield.igrid + grid = vectorfield.grid _, grid_positions = _get_positions( fieldset.U, particles.t, particles.z, particles.y, particles.x, particles, particles.ei[:, igrid] ) - xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] - yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + _, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + _, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] zi, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] - ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] + U0, U1, V0, V1, W0, W1, px, py = _get_cgrid_velocities(vectorfield, grid_positions) - ds_t = dt - # if withTime: - # time_i = np.linspace(0, fieldset.U.grid.time[ti + 1] - fieldset.U.grid.time[ti], I_s) - # ds_t = min(ds_t, time_i[np.where(particles.time - fieldset.U.grid.time[ti] < time_i)[0][0]]) - # print("withTime", withTime, "ds_t", ds_t, dt) - if withW: - if abs(xsi - 1) < tol: - if fieldset.U.data[0, zi + 1, yi + 1, xi + 1] > 0: - xi += 1 - xsi = 0 - if abs(eta - 1) < tol: - if fieldset.V.data[0, zi + 1, yi + 1, xi + 1] > 0: - yi += 1 - eta = 0 - if abs(zeta - 1) < tol: - if fieldset.W.data[0, zi + 1, yi + 1, xi + 1] > 0: - zi += 1 - zeta = 0 - else: - if abs(xsi - 1) < tol: - if fieldset.U.data[0, 0, yi + 1, xi + 1] > 0: - xi += 1 - xsi = 0 - if abs(eta - 1) < tol: - if fieldset.V.data[0, 0, yi + 1, xi + 1] > 0: - yi += 1 - eta = 0 - - # particles.ei[:, igrid] = fieldset.U.grid.ravel_index(zi, yi, xi) - - grid = fieldset.U.grid - if grid._gtype < 2: - px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) - py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) - else: - px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) - py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) - - if grid._mesh.is_spherical(): - px = ((px + 180.0) % 360.0) - 180.0 - px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) - px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) if withW: pz = np.array([grid.depth[zi], grid.depth[zi + 1]]) dz = pz[1] - pz[0] else: dz = 1.0 - c1 = i_u._geodetic_distance( - py[0], py[1], px[0], px[1], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(0.0, xsi), py) - ) - c2 = i_u._geodetic_distance( - py[1], py[2], px[1], px[2], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 1.0), py) - ) - c3 = i_u._geodetic_distance( - py[2], py[3], px[2], px[3], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(1.0, xsi), py) - ) - c4 = i_u._geodetic_distance( - py[3], py[0], px[3], px[0], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 0.0), py) - ) rad = np.pi / 180.0 deg2m = 1852 * 60.0 meshJac = (deg2m * deg2m * np.cos(rad * particles.y)) if grid._mesh.is_spherical() else 1 dxdy = i_u._compute_jacobian_determinant(py, px, eta, xsi) * meshJac + U0 *= direction * dz + U1 *= direction * dz + V0 *= direction * dz + V1 *= direction * dz if withW: - U0 = direction * fieldset.U.data[ti, zi + 1, yi + 1, xi] * c4 * dz - U1 = direction * fieldset.U.data[ti, zi + 1, yi + 1, xi + 1] * c2 * dz - V0 = direction * fieldset.V.data[ti, zi + 1, yi, xi + 1] * c1 * dz - V1 = direction * fieldset.V.data[ti, zi + 1, yi + 1, xi + 1] * c3 * dz - if withTime: - U0 = U0 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, zi + 1, yi + 1, xi] * c4 * dz - U1 = U1 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, zi + 1, yi + 1, xi + 1] * c2 * dz - V0 = V0 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, zi + 1, yi, xi + 1] * c1 * dz - V1 = V1 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, zi + 1, yi + 1, xi + 1] * c3 * dz - else: - U0 = (direction * fieldset.U.data[ti, 0, yi + 1, xi] * c4 * dz).values.flatten() - U1 = (direction * fieldset.U.data[ti, 0, yi + 1, xi + 1] * c2 * dz).values.flatten() - V0 = (direction * fieldset.V.data[ti, 0, yi, xi + 1] * c1 * dz).values.flatten() - V1 = (direction * fieldset.V.data[ti, 0, yi + 1, xi + 1] * c3 * dz).values.flatten() - if withTime: - U0 = U0 * (1 - tau) + (tau * direction * fieldset.U.data[ti + 1, 0, yi + 1, xi] * c4 * dz).values.flatten() - U1 = ( - U1 * (1 - tau) - + (tau * direction * fieldset.U.data[ti + 1, 0, yi + 1, xi + 1] * c2 * dz).values.flatten() - ) - V0 = V0 * (1 - tau) + (tau * direction * fieldset.V.data[ti + 1, 0, yi, xi + 1] * c1 * dz).values.flatten() - V1 = ( - V1 * (1 - tau) - + (tau * direction * fieldset.V.data[ti + 1, 0, yi + 1, xi + 1] * c3 * dz).values.flatten() - ) + W0 *= direction * dxdy + W1 *= direction * dxdy def compute_ds(F0, F1, r, direction, tol): # noqa: N803 with np.errstate(divide="ignore", invalid="ignore"): + print("NOW IN COMPUTE_DS") + print(F0, F1, r, direction) up = F0 * (1 - r) + F1 * r r_target = np.where(direction * up >= 0.0, 1.0, 0.0) B = F0 - F1 @@ -283,28 +215,23 @@ def compute_ds(F0, F1, r, direction, tol): # noqa: N803 F_r1 = np.where(np.abs(B) > tol, r_target + delta / B, np.nan) F_r0 = np.where(np.abs(B) > tol, r + delta / B, np.nan) - ds = -1.0 / B * np.log(F_r1 / F_r0) - ds = np.where(F_r1 * F_r0 < tol, np.inf, ds) - ds = np.where(B == 0, -delta * direction / up, ds) - ds = np.where((np.abs(B) < tol) & (np.abs(delta) < tol), np.inf, ds) + d_s = -1.0 / B * np.log(F_r1 / F_r0) + d_s = np.where(F_r1 * F_r0 < tol, np.inf, d_s) + d_s = np.where(B == 0, -delta * direction / up, d_s) + d_s = np.where((np.abs(B) < tol) & (np.abs(delta) < tol), np.inf, d_s) - ds = np.where(ds < tol, np.inf, ds) - return ds, B, delta + d_s = np.where(d_s < tol, np.inf, d_s) + return d_s, B, delta ds_x, B_x, delta_x = compute_ds(U0, U1, xsi, direction, tol) ds_y, B_y, delta_y = compute_ds(V0, V1, eta, direction, tol) if withW: - W0 = direction * fieldset.W.data[ti, zi, yi + 1, xi + 1] * dxdy - W1 = direction * fieldset.W.data[ti, zi + 1, yi + 1, xi + 1] * dxdy - if withTime: - W0 = W0 * (1 - tau) + tau * direction * fieldset.W.data[ti + 1, zi, yi + 1, xi + 1] * dxdy - W1 = W1 * (1 - tau) + tau * direction * fieldset.W.data[ti + 1, zi + 1, yi + 1, xi + 1] * dxdy - ds_z, B_z, delta_z = compute_ds(W0.values.flatten(), W1.values.flatten(), zeta, direction, tol) + ds_z, B_z, delta_z = compute_ds(W0, W1, zeta, direction, tol) else: ds_z = np.inf # take the minimum travel time - s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(ds_t / (dxdy * dz))) + s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(dt / (dxdy * dz))) # calculate end position in time s_min def compute_rs(r, B, delta, s_min): # noqa: N803 From 4cedd61be29d58cbe59fa00d96f792d64a32e893 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 12 Aug 2026 16:49:05 +0200 Subject: [PATCH 08/12] Removing old print statements --- src/parcels/kernels/_advection.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index ee0e6f69af..ee53f89779 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -204,8 +204,6 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover def compute_ds(F0, F1, r, direction, tol): # noqa: N803 with np.errstate(divide="ignore", invalid="ignore"): - print("NOW IN COMPUTE_DS") - print(F0, F1, r, direction) up = F0 * (1 - r) + F1 * r r_target = np.where(direction * up >= 0.0, 1.0, 0.0) B = F0 - F1 From fa1932b27e5f7e20c39e601aed328368babb6220 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 12 Aug 2026 17:01:31 +0200 Subject: [PATCH 09/12] Checking that AnalyticalAdvection only works on C grids --- src/parcels/kernels/_advection.py | 7 ++++++- tests/test_advection.py | 18 +++++++++--------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index ee53f89779..8e4e89ef5d 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -3,7 +3,6 @@ import numpy as np from parcels._core.statuscodes import StatusCode -from parcels.interpolators._xinterpolators import _get_cgrid_velocities __all__ = [ "AdvectionAnalytical", @@ -163,6 +162,7 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover """ import parcels._core.utils.interpolation as i_u from parcels._core.field import _get_positions + from parcels.interpolators._xinterpolators import CGrid_Velocity, _get_cgrid_velocities tol = 1e-10 # I_s = 10 # number of intermediate time steps @@ -171,6 +171,11 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover withW = True if "W" in [f.name for f in fieldset.fields.values()] else False vectorfield = fieldset.UVW if withW else fieldset.UV + if not isinstance(vectorfield.interp_method, CGrid_Velocity): + raise NotImplementedError( + "Analytical advection is only implemented for C-grid velocity fields, " + f"but the fieldset has interp_method={vectorfield.interp_method}" + ) # withTime = True if len(vectorfield.grid.time) > 1 else False igrid = vectorfield.igrid grid = vectorfield.grid diff --git a/tests/test_advection.py b/tests/test_advection.py index a851fd0eb3..e220ea3653 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -27,6 +27,7 @@ stommel_gyre_dataset, ) from parcels._datasets.structured.generic import datasets_sgrid +from parcels.interpolators import CGrid_Velocity from parcels.kernels import ( AdvectionAnalytical, AdvectionDiffusionEM, @@ -508,16 +509,13 @@ def test_mitgcm(): np.testing.assert_allclose(pset.y, lat_v3, atol=1) -# def test_analyticalAgrid(): -# lon = np.arange(0, 15, dtype=np.float32) -# lat = np.arange(0, 15, dtype=np.float32) -# U = np.ones((lat.size, lon.size), dtype=np.float32) -# V = np.ones((lat.size, lon.size), dtype=np.float32) -# fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") -# pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) +def test_analytical_throwserror_on_Agrid(): + ds = simple_UV_dataset(mesh="flat") + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + pset = ParticleSet(fieldset, x=1, y=1) -# with pytest.raises(NotImplementedError): -# pset.execute(AdvectionAnalytical, runtime=1) + with pytest.raises(NotImplementedError): + pset.execute(AdvectionAnalytical, runtime=1, dt=1) @pytest.mark.parametrize("u", [0.2, -0.3, 0]) @@ -531,6 +529,8 @@ def test_uniform_analytical(u, v, w, direction, tmp_parquet): if w is not None: ds["W"] = xr.full_like(ds["U"], w) fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + fieldset.UV.interp_method = CGrid_Velocity() + fieldset.UVW.interp_method = CGrid_Velocity() x0, y0, z0 = 6.1, 6.2, 0.5 pset = ParticleSet(fieldset, pclass=Particle, x=x0, y=y0, z=z0) From 9c232dc7d5590b8c90558f1af6b36703d86a30ed Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 12 Aug 2026 17:17:44 +0200 Subject: [PATCH 10/12] Fixing unit tetst --- tests/test_advection.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_advection.py b/tests/test_advection.py index e220ea3653..ab6428b6b3 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -528,9 +528,11 @@ def test_uniform_analytical(u, v, w, direction, tmp_parquet): ds["V"].data[:] = v if w is not None: ds["W"] = xr.full_like(ds["U"], w) + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") fieldset.UV.interp_method = CGrid_Velocity() - fieldset.UVW.interp_method = CGrid_Velocity() + if w is not None: + fieldset.UVW.interp_method = CGrid_Velocity() x0, y0, z0 = 6.1, 6.2, 0.5 pset = ParticleSet(fieldset, pclass=Particle, x=x0, y=y0, z=z0) From 0950e329444518b94ea6f1bcd2230c6831b2bb5b Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 13 Aug 2026 11:52:38 +0200 Subject: [PATCH 11/12] vectorizing AdvectionAnalytical --- src/parcels/kernels/_advection.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index 8e4e89ef5d..f981902c57 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -167,7 +167,7 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover tol = 1e-10 # I_s = 10 # number of intermediate time steps dt = particles.dt - direction = 1.0 if dt > 0 else -1.0 + direction = 1.0 if np.any(dt > 0) else -1.0 withW = True if "W" in [f.name for f in fieldset.fields.values()] else False vectorfield = fieldset.UVW if withW else fieldset.UV @@ -234,7 +234,7 @@ def compute_ds(F0, F1, r, direction, tol): # noqa: N803 ds_z = np.inf # take the minimum travel time - s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(dt / (dxdy * dz))) + s_min = np.minimum(np.minimum(np.abs(ds_x), np.abs(ds_y)), np.minimum(np.abs(ds_z), np.abs(dt / (dxdy * dz)))) # calculate end position in time s_min def compute_rs(r, B, delta, s_min): # noqa: N803 @@ -263,7 +263,7 @@ def compute_rs(r, B, delta, s_min): # noqa: N803 rs_z = compute_rs(zeta, B_z, delta_z, s_min) particles.dz += (1.0 - rs_z) * pz[0] + rs_z * pz[1] - particles.z - if particles.dt > 0: - particles.dt = max(direction * s_min * (dxdy * dz), 1e-7) + if direction > 0: + particles.dt = np.maximum(direction * s_min * (dxdy * dz), 1e-7) else: - particles.dt = min(direction * s_min * (dxdy * dz), -1e-7) + particles.dt = np.minimum(direction * s_min * (dxdy * dz), -1e-7) From fe2665cbe9c8f72fe3e01ccfcc3313d3f50b2415 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 14 Aug 2026 16:47:32 +0200 Subject: [PATCH 12/12] Make radial_rotation NEMO-grid --- src/parcels/_datasets/structured/generated.py | 49 ++++++++++++------- 1 file changed, 32 insertions(+), 17 deletions(-) diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py index 7eba61f559..abb5f6d603 100644 --- a/src/parcels/_datasets/structured/generated.py +++ b/src/parcels/_datasets/structured/generated.py @@ -39,41 +39,56 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"): ) -def radial_rotation_dataset(xdim=200, ydim=200): # Define 2D flat, square fieldset for testing purposes. +def radial_rotation_dataset(xdim=200, ydim=200, grid_type="A"): # Define 2D flat, square fieldset for testing purposes. lon = np.linspace(0, 60, xdim, dtype=np.float32) lat = np.linspace(0, 60, ydim, dtype=np.float32) x0 = 30.0 # Define the origin to be the centre of the Field. y0 = 30.0 + dx, dy = lon[-1] / xdim, lat[-1] / ydim # Define the grid spacing in x and y directions. + U = np.zeros((2, 1, ydim, xdim), dtype=np.float32) V = np.zeros((2, 1, ydim, xdim), dtype=np.float32) + R = np.zeros((2, 1, ydim, xdim), dtype=np.float32) omega = 2 * np.pi / 86400.0 # Define the rotational period as 1 day. + def calc_r_theta(ln, lt, x0, y0): + r = np.sqrt((ln - x0) ** 2 + (lt - y0) ** 2) + theta = np.arctan2((lt - y0), (ln - x0)) + return r, theta + for i in range(lon.size): for j in range(lat.size): - r = np.sqrt((lon[i] - x0) ** 2 + (lat[j] - y0) ** 2) - assert r >= 0.0 - assert r <= np.sqrt(x0**2 + y0**2) - - theta = np.arctan2((lat[j] - y0), (lon[i] - x0)) - assert abs(theta) <= np.pi + r, theta = calc_r_theta(lon[i], lat[j], x0, y0) + R[:, :, j, i] = r + if grid_type == "A": + r, theta = calc_r_theta(lon[i], lat[j], x0, y0) + U[:, :, j, i] = r * np.sin(theta) * omega + V[:, :, j, i] = -r * np.cos(theta) * omega + elif grid_type == "C": + r, theta = calc_r_theta(lon[i] - dx / 2, lat[j], x0, y0) + U[:, :, j, i] = r * np.sin(theta) * omega - U[:, :, j, i] = r * np.sin(theta) * omega - V[:, :, j, i] = -r * np.cos(theta) * omega + r, theta = calc_r_theta(lon[i], lat[j] - dy / 2, x0, y0) + V[:, :, j, i] = -r * np.cos(theta) * omega return xr.Dataset( - {"U": (["time", "depth", "YG", "XG"], U), "V": (["time", "depth", "YG", "XG"], V)}, + { + "U": (["time", "depth", "YG", "XC"], U), + "V": (["time", "depth", "YC", "XG"], V), + "R": (["time", "depth", "YC", "XC"], R), + }, coords={ "time": (["time"], [np.timedelta64(0, "s"), np.timedelta64(10, "D")], {"axis": "T"}), "depth": (["depth"], np.array([0.0]), {"axis": "Z"}), - "YC": (["YC"], np.arange(ydim) + 0.5, {"axis": "Y"}), - "YG": (["YG"], np.arange(ydim), {"axis": "Y", "c_grid_axis_shift": -0.5}), - "XC": (["XC"], np.arange(xdim) + 0.5, {"axis": "X"}), - "XG": (["XG"], np.arange(xdim), {"axis": "X", "c_grid_axis_shift": -0.5}), - "lat": (["YG"], lat, {"axis": "Y", "c_grid_axis_shift": 0.5}), - "lon": (["XG"], lon, {"axis": "X", "c_grid_axis_shift": -0.5}), + "YC": (["YC"], np.arange(ydim) - 0.5, {"axis": "Y", "c_grid_axis_shift": +0.5}), + "YG": (["YG"], np.arange(ydim), {"axis": "Y"}), + "XC": (["XC"], np.arange(xdim) - 0.5, {"axis": "X", "c_grid_axis_shift": +0.5}), + "XG": (["XG"], np.arange(xdim), {"axis": "X"}), + "lat": (["YG"], lat, {"axis": "Y"}), + "lon": (["XG"], lon, {"axis": "X"}), }, ).pipe( sgrid._attach_sgrid_metadata, @@ -84,7 +99,7 @@ def radial_rotation_dataset(xdim=200, ydim=200): # Define 2D flat, square field node_coordinates=("lon", "lat"), face_dimensions=( sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW), - sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH), + sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.LOW), ), vertical_dimensions=(sgrid.FaceNodePadding("ZC", "depth", sgrid.Padding.BOTH),), ),