Skip to content

Adding in the Runge-Kutta 2 advection schemes #2315

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Nov 3, 2025
Merged

Adding in the Runge-Kutta 2 advection schemes #2315

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
179d306
Implementing RK2 and RK2_3D advection schemes
michaeldenes Oct 8, 2025
202dc07
Adding the two new schemes to the list of schemes
michaeldenes Oct 8, 2025
3a95f67
Updating init to include the new kernels
michaeldenes Oct 8, 2025
4be9a2b
Adding RK2 and RK2_3D into (some of) the unit tests.
michaeldenes Oct 8, 2025
c107a5d
Updating tolerances for RK2 unit tests
erikvansebille Oct 8, 2025
a39a72c
Removing distionary from method to kernels
erikvansebille Nov 3, 2025
29b8d37
Merge branch 'v4-dev' into runge-kutta2
erikvansebille Nov 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/parcels/kernels/__init__.py
View file
Open in desktop
Copy link
Member

@erikvansebille erikvansebille Oct 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@VeckoTheGecko, do we really need to add new Kernels at three locations? (two locations here in __init__.py and also at the top of advection.py). Can there not be less duplication?

Copy link
Contributor

@VeckoTheGecko VeckoTheGecko Oct 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can make advection.py private (i.e., changing it to be _advection.py). Then we wouldn't need to define __all__ in advection.py. As it stands though, we need to have it in __all__ for advection.py and __init__.py

Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
from .advection import (
AdvectionAnalytical,
AdvectionEE,
AdvectionRK2,
AdvectionRK2_3D,
AdvectionRK4,
AdvectionRK4_3D,
AdvectionRK4_3D_CROCO,
Expand All @@ -21,6 +23,8 @@
# advection
"AdvectionAnalytical",
"AdvectionEE",
"AdvectionRK2",
"AdvectionRK2_3D",
"AdvectionRK4_3D_CROCO",
"AdvectionRK4_3D",
"AdvectionRK4",
Expand Down
25 changes: 25 additions & 0 deletions src/parcels/kernels/advection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,38 @@
__all__ = [
"AdvectionAnalytical",
"AdvectionEE",
"AdvectionRK2",
"AdvectionRK2_3D",
"AdvectionRK4",
"AdvectionRK4_3D",
"AdvectionRK4_3D_CROCO",
"AdvectionRK45",
]


def AdvectionRK2(particles, fieldset): # pragma: no cover
"""Advection of particles using second-order Runge-Kutta integration."""
dt = particles.dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds
(u1, v1) = fieldset.UV[particles]
lon1, lat1 = (particles.lon + u1 * 0.5 * dt, particles.lat + v1 * 0.5 * dt)
(u2, v2) = fieldset.UV[particles.time + 0.5 * particles.dt, particles.z, lat1, lon1, particles]
particles.dlon += u2 * dt
particles.dlat += v2 * dt


def AdvectionRK2_3D(particles, fieldset): # pragma: no cover
"""Advection of particles using second-order Runge-Kutta integration including vertical velocity."""
dt = particles.dt / np.timedelta64(1, "s")
(u1, v1, w1) = fieldset.UVW[particles]
lon1 = particles.lon + u1 * 0.5 * dt
lat1 = particles.lat + v1 * 0.5 * dt
z1 = particles.z + w1 * 0.5 * dt
(u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * particles.dt, z1, lat1, lon1, particles]
particles.dlon += u2 * dt
particles.dlat += v2 * dt
particles.dz += w2 * dt


def AdvectionRK4(particles, fieldset): # pragma: no cover
"""Advection of particles using fourth-order Runge-Kutta integration."""
dt = particles.dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds
Expand Down
91 changes: 44 additions & 47 deletions tests/test_advection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,22 +18,14 @@
AdvectionDiffusionEM,
AdvectionDiffusionM1,
AdvectionEE,
AdvectionRK2,
AdvectionRK2_3D,
AdvectionRK4,
AdvectionRK4_3D,
AdvectionRK45,
)
from tests.utils import round_and_hash_float_array

kernel = {
"EE": AdvectionEE,
"RK4": AdvectionRK4,
"RK4_3D": AdvectionRK4_3D,
"RK45": AdvectionRK45,
# "AA": AdvectionAnalytical,
"AdvDiffEM": AdvectionDiffusionEM,
"AdvDiffM1": AdvectionDiffusionM1,
}


@pytest.mark.parametrize("mesh", ["spherical", "flat"])
def test_advection_zonal(mesh, npart=10):
Expand Down Expand Up @@ -256,30 +248,32 @@ def test_radialrotation(npart=10):


@pytest.mark.parametrize(
"method, rtol",
"kernel, rtol",
[
("EE", 1e-2),
("AdvDiffEM", 1e-2),
("AdvDiffM1", 1e-2),
("RK4", 1e-5),
("RK4_3D", 1e-5),
("RK45", 1e-4),
(AdvectionEE, 1e-2),
(AdvectionDiffusionEM, 1e-2),
(AdvectionDiffusionM1, 1e-2),
(AdvectionRK2, 1e-4),
(AdvectionRK2_3D, 1e-4),
(AdvectionRK4, 1e-5),
(AdvectionRK4_3D, 1e-5),
(AdvectionRK45, 1e-4),
],
)
def test_moving_eddy(method, rtol):
def test_moving_eddy(kernel, rtol):
ds = moving_eddy_dataset()
grid = XGrid.from_dataset(ds)
U = Field("U", ds["U"], grid, interp_method=XLinear)
V = Field("V", ds["V"], grid, interp_method=XLinear)
if method == "RK4_3D":
if kernel in [AdvectionRK2_3D, AdvectionRK4_3D]:
# Using W to test 3D advection (assuming same velocity as V)
W = Field("W", ds["V"], grid, interp_method=XLinear)
UVW = VectorField("UVW", U, V, W)
fieldset = FieldSet([U, V, W, UVW])
else:
UV = VectorField("UV", U, V)
fieldset = FieldSet([U, V, UV])
if method in ["AdvDiffEM", "AdvDiffM1"]:
if kernel in [AdvectionDiffusionEM, AdvectionDiffusionM1]:
# Add zero diffusivity field for diffusion kernels
ds["Kh"] = (["time", "depth", "YG", "XG"], np.full(ds["U"].shape, 0))
fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_zonal")
Expand All @@ -289,11 +283,11 @@ def test_moving_eddy(method, rtol):
start_lon, start_lat, start_z = 12000, 12500, 12500
dt = np.timedelta64(30, "m")

if method == "RK45":
if kernel == AdvectionRK45:
fieldset.add_constant("RK45_tol", rtol)

pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, z=start_z, time=np.timedelta64(0, "s"))
pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "h"))
pset.execute(kernel, dt=dt, endtime=np.timedelta64(1, "h"))

def truth_moving(x_0, y_0, t):
t /= np.timedelta64(1, "s")
Expand All @@ -304,19 +298,20 @@ def truth_moving(x_0, y_0, t):
exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time[0])
np.testing.assert_allclose(pset.lon, exp_lon, rtol=rtol)
np.testing.assert_allclose(pset.lat, exp_lat, rtol=rtol)
if method == "RK4_3D":
if kernel == AdvectionRK4_3D:
np.testing.assert_allclose(pset.z, exp_lat, rtol=rtol)


@pytest.mark.parametrize(
"method, rtol",
"kernel, rtol",
[
("EE", 1e-1),
("RK4", 1e-5),
("RK45", 1e-4),
(AdvectionEE, 1e-1),
(AdvectionRK2, 3e-3),
(AdvectionRK4, 1e-5),
(AdvectionRK45, 1e-4),
],
)
def test_decaying_moving_eddy(method, rtol):
def test_decaying_moving_eddy(kernel, rtol):
ds = decaying_moving_eddy_dataset()
grid = XGrid.from_dataset(ds)
U = Field("U", ds["U"], grid, interp_method=XLinear)
Expand All @@ -327,12 +322,12 @@ def test_decaying_moving_eddy(method, rtol):
start_lon, start_lat = 10000, 10000
dt = np.timedelta64(60, "m")

if method == "RK45":
if kernel == AdvectionRK45:
fieldset.add_constant("RK45_tol", rtol)
fieldset.add_constant("RK45_min_dt", 10 * 60)

pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s"))
pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(23, "h"))
pset.execute(kernel, dt=dt, endtime=np.timedelta64(23, "h"))

def truth_moving(x_0, y_0, t):
t /= np.timedelta64(1, "s")
Expand All @@ -354,14 +349,15 @@ def truth_moving(x_0, y_0, t):


@pytest.mark.parametrize(
"method, rtol",
"kernel, rtol",
[
("RK4", 0.1),
("RK45", 0.1),
(AdvectionRK2, 0.1),
(AdvectionRK4, 0.1),
(AdvectionRK45, 0.1),
],
)
@pytest.mark.parametrize("grid_type", ["A", "C"])
def test_stommelgyre_fieldset(method, rtol, grid_type):
def test_stommelgyre_fieldset(kernel, rtol, grid_type):
npart = 2
ds = stommel_gyre_dataset(grid_type=grid_type)
grid = XGrid.from_dataset(ds)
Expand All @@ -377,7 +373,7 @@ def test_stommelgyre_fieldset(method, rtol, grid_type):
start_lon = np.linspace(10e3, 100e3, npart)
start_lat = np.ones_like(start_lon) * 5000e3

if method == "RK45":
if kernel == AdvectionRK45:
fieldset.add_constant("RK45_tol", rtol)

SampleParticle = Particle.add_variable(
Expand All @@ -389,19 +385,20 @@ def UpdateP(particles, fieldset): # pragma: no cover
particles.p_start = np.where(particles.time == 0, particles.p, particles.p_start)

pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s"))
pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime)
pset.execute([kernel, UpdateP], dt=dt, runtime=runtime)
np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol)


@pytest.mark.parametrize(
"method, rtol",
"kernel, rtol",
[
("RK4", 5e-3),
("RK45", 1e-4),
(AdvectionRK2, 2e-2),
(AdvectionRK4, 5e-3),
(AdvectionRK45, 1e-4),
],
)
@pytest.mark.parametrize("grid_type", ["A"]) # TODO also implement C-grid once available
def test_peninsula_fieldset(method, rtol, grid_type):
def test_peninsula_fieldset(kernel, rtol, grid_type):
npart = 2
ds = peninsula_dataset(grid_type=grid_type)
grid = XGrid.from_dataset(ds)
Expand All @@ -416,7 +413,7 @@ def test_peninsula_fieldset(method, rtol, grid_type):
start_lat = np.linspace(3e3, 47e3, npart)
start_lon = 3e3 * np.ones_like(start_lat)

if method == "RK45":
if kernel == AdvectionRK45:
fieldset.add_constant("RK45_tol", rtol)

SampleParticle = Particle.add_variable(
Expand All @@ -428,7 +425,7 @@ def UpdateP(particles, fieldset): # pragma: no cover
particles.p_start = np.where(particles.time == 0, particles.p, particles.p_start)

pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s"))
pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime)
pset.execute([kernel, UpdateP], dt=dt, runtime=runtime)
np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol)


Expand Down Expand Up @@ -474,8 +471,8 @@ def periodicBC(particles, fieldset): # pragma: no cover
np.testing.assert_allclose(pset.lat, latp, atol=1e-1)


@pytest.mark.parametrize("method", ["RK4", "RK4_3D"])
def test_nemo_3D_curvilinear_fieldset(method):
@pytest.mark.parametrize("kernel", [AdvectionRK4, AdvectionRK4_3D])
def test_nemo_3D_curvilinear_fieldset(kernel):
download_dir = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data")
ufiles = download_dir.glob("*U.nc")
dsu = xr.open_mfdataset(ufiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"])
Expand Down Expand Up @@ -553,10 +550,10 @@ def test_nemo_3D_curvilinear_fieldset(method):
lats = np.linspace(52.5, 51.6, npart)
pset = parcels.ParticleSet(fieldset, lon=lons, lat=lats, z=np.ones_like(lons))

pset.execute(kernel[method], runtime=np.timedelta64(3, "D") + np.timedelta64(18, "h"), dt=np.timedelta64(6, "h"))
pset.execute(kernel, runtime=np.timedelta64(3, "D") + np.timedelta64(18, "h"), dt=np.timedelta64(6, "h"))

if method == "RK4":
if kernel == AdvectionRK4:
np.testing.assert_equal(round_and_hash_float_array([p.lon for p in pset], decimals=5), 29977383852960156017546)
elif method == "RK4_3D":
elif kernel == AdvectionRK4_3D:
# TODO check why decimals needs to be so low in RK4_3D (compare to v3)
np.testing.assert_equal(round_and_hash_float_array([p.z for p in pset], decimals=1), 29747210774230389239432)
Loading