Moving RK2 and RK4 methods to separate integrators.py functions

docs/user_guide/examples/tutorial_dt_integrators.ipynb

@@ -521,7 +521,9 @@
    "outputs": [],
    "source": [
     "# And the second-order Runge-Kutta (RK2) method\n",
-    "print(inspect.getsource(parcels.kernels.AdvectionRK2))"
+    "print(inspect.getsource(parcels.kernels._advection.rhs_uv))\n",
+    "print(inspect.getsource(parcels.kernels.AdvectionRK2))\n",
+    "print(inspect.getsource(parcels.kernels._integration.RK2))"
    ]
   },
   {
@@ -532,7 +534,9 @@
    "outputs": [],
    "source": [
     "# And RK4:\n",
-    "print(inspect.getsource(parcels.kernels.AdvectionRK4))"
+    "print(inspect.getsource(parcels.kernels._advection.rhs_uv))\n",
+    "print(inspect.getsource(parcels.kernels.AdvectionRK4))\n",
+    "print(inspect.getsource(parcels.kernels._integration.RK4))"
    ]
   },
   {
@@ -792,6 +796,126 @@
    "cell_type": "markdown",
    "id": "36",
    "metadata": {},
+   "source": [
+    "## Writing custom Kernels with higher-order integrators\n",
+    "\n",
+    "Similarly to the built-in AdvectionRK2 Kernel, you can also write your own custom Kernels that use higher-order integrators. To do so, you can use the `RK2` and `RK4` integrators from `parcels.integrators` in your Kernel definition. \n",
+    "\n",
+    "The 'trick' is to write a right-hand-side function (`rhs`) that computes the change in position at a given time and location, and then use this function in the `RK2` or `RK4` integrator to compute the change in position over a given timestep.\n",
+    "\n",
+    "For example, let's say we want to write a custom Kernel that computes advection and a simple windage effect. We can write the right-hand-side function as follows"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def rhs(fieldset, t, z, y, x, particles):\n",
+    "    uo, vo = fieldset.UV[t, z, y, x, particles]\n",
+    "    uw, vw = fieldset.UVw[t, z, y, x, particles]\n",
+    "    return (uo + fieldset.wind_coeff * (uw - uo), vo + fieldset.wind_coeff * (vw - vo))\n",
+    "\n",
+    "\n",
+    "def WindageRK2(particles, fieldset):\n",
+    "\n",
+    "    u, v = parcels.kernels._integration.RK2(fieldset, particles, rhs)\n",
+    "\n",
+    "    particles.dx += u * particles.dt\n",
+    "    particles.dy += v * particles.dt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "38",
+   "metadata": {},
+   "source": [
+    "Now, if we also load in the wind data for the same region with the code below, we can execute this Kernel to compute the trajectories of particles that are advected by the ocean currents and affected by the wind. The solution will be second-order accurate in time, since we are using the RK2 integrator; but it is easy to swap in the RK4 integrator if you want to increase the accuracy even more."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_w = parcels.tutorial.open_dataset(\n",
+    "    \"CopernicusMarine_data_for_Argo_tutorial/data_wind\"\n",
+    ")\n",
+    "ds_w = ds_w.expand_dims({\"depth\": [0]})\n",
+    "ds_w.load()  # load the dataset into memory\n",
+    "\n",
+    "fields = {\"Uw\": ds_w[\"eastward_wind\"], \"Vw\": ds_w[\"northward_wind\"]}\n",
+    "ds_wind = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
+    "fieldset_w = parcels.FieldSet.from_sgrid_conventions(\n",
+    "    ds_wind, vector_fields={\"UVw\": (\"Uw\", \"Vw\")}\n",
+    ")\n",
+    "\n",
+    "fieldset += fieldset_w\n",
+    "fieldset.describe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fieldset.wind_coeff = 0.01  # windage coefficient (1% of the relative wind speed)\n",
+    "\n",
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset=fieldset,\n",
+    "    pclass=parcels.Particle,\n",
+    "    t=initial_release_times,\n",
+    "    z=np.ones_like(initial_release_times),\n",
+    "    y=initial_release_lats,\n",
+    "    x=initial_release_lons,\n",
+    ")\n",
+    "\n",
+    "pfile = parcels.ParticleFile(\n",
+    "    path=OUTPUT_FOLDER / \"WindRK2.parquet\",\n",
+    "    outputdt=outputdt,\n",
+    "    mode=\"w\",\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    WindageRK2,\n",
+    "    runtime=runtime,\n",
+    "    dt=np.timedelta64(1, \"h\"),\n",
+    "    output_file=pfile,\n",
+    "    verbose_progress=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "41",
+   "metadata": {},
+   "source": [
+    "As you can see, the particles move very differently when we include the windage effect!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "42",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = parcels.read_particlefile(pfile.path)\n",
+    "for traj in df.partition_by(\"particle_id\"):\n",
+    "    plt.plot(traj[\"x\"], traj[\"y\"], alpha=0.75)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "43",
+   "metadata": {},
    "source": [
     "## Summary\n",
     "If you want to test the accuracy and efficiency of your own simulation, here is a brief list of things to consider:\n",
@@ -803,7 +927,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Parcels:test (3.14.6)",
+   "display_name": "Parcels:default (3.14.7)",
    "language": "python",
    "name": "python3"
   },
@@ -817,7 +941,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.14.6"
+   "version": "3.14.7"
   }
  },
  "nbformat": 4,

src/parcels/_datasets/remote.py

@@ -49,6 +49,7 @@ def _get_data_home() -> Path:
         "data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-cur_anfc_0.083deg_P1D-m_uo-vo_31.00E-33.00E_33.00S-30.00S_0.49-2225.08m_2024-01-01-2024-02-01.nc",
         "data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-so_anfc_0.083deg_P1D-m_so_31.00E-33.00E_33.00S-30.00S_0.49-2225.08m_2024-01-01-2024-02-01.nc",
         "data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-thetao_anfc_0.083deg_P1D-m_thetao_31.00E-33.00E_33.00S-30.00S_0.49-2225.08m_2024-01-01-2024-02-01.nc",
+        "data/CopernicusMarine_data_for_Argo_tutorial/cmems_obs-wind_glo_phy_my_l4_0.125deg_PT1H_multi-vars_31.06E-32.94E_32.94S-30.06S_2024-01-01-2024-02-01.nc",
     ]
     + ["data/CopernicusMarine_data_for_stuck_particles_tutorial/cmems_mod_glo_phy_my_0.083deg_P1D-m_NL.nc"]
     # + [
@@ -224,6 +225,7 @@ class _Purpose(enum.Enum):
     # ("Peninsula_data/T", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaT.nc"), _Purpose.TUTORIAL)),
     # ("GlobCurrent_example_data/data", (_V3Dataset(_ODIE,"data/GlobCurrent_example_data/*000000-GLOBCURRENT-L4-CUReul_hs-ALT_SUM-v02.0-fv01.0.nc", pre_decode_cf_callable=patch_dataset_v4_compat), _Purpose.TUTORIAL)),
     ("CopernicusMarine_data_for_Argo_tutorial/data", (_V3Dataset(_ODIE,"data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-*.nc"), _Purpose.TUTORIAL)),
+    ("CopernicusMarine_data_for_Argo_tutorial/data_wind", (_V3Dataset(_ODIE,"data/CopernicusMarine_data_for_Argo_tutorial/cmems_obs-wind_glo_phy*.nc"), _Purpose.TUTORIAL)),
     ("Delft3D_data/Rotterdam_tiny", (_V3Dataset(_ODIE,"data/Delft3D_data/Rotterdam_tiny.nc"), _Purpose.TUTORIAL)),
     ("CopernicusMarine_data_for_stuck_particles_tutorial/data", (_V3Dataset(_ODIE,"data/CopernicusMarine_data_for_stuck_particles_tutorial/cmems_mod_glo_phy_my_0.083deg_P1D-m_NL.nc"), _Purpose.TUTORIAL)),
     # ("DecayingMovingEddy_data/U", (_V3Dataset(_ODIE,"data/DecayingMovingEddy_data/decaying_moving_eddyU.nc"), _Purpose.TUTORIAL)),

src/parcels/kernels/__init__.py

@@ -12,6 +12,7 @@
     AdvectionDiffusionM1,
     DiffusionUniformKh,
 )
+from ._integration import RK2, RK4
 from ._sigmagrids import (
     AdvectionRK2_3D_CROCO,
     SampleOmegaCroco,
@@ -35,4 +36,7 @@
     "AdvectionRK2_3D_CROCO",
     "SampleOmegaCroco",
     "convert_z_to_sigma_croco",
+    # integration
+    "RK2",
+    "RK4",
 ]

src/parcels/kernels/_advection.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from parcels._core.statuscodes import StatusCode
+from parcels.kernels._integration import RK2, RK4
 
 __all__ = [
     "AdvectionAnalytical",
@@ -17,62 +18,42 @@
 ]
 
 
+def rhs_uv(fieldset, t, z, y, x, particles):
+    return fieldset.UV[t, z, y, x, particles]
+
+
+def rhs_uvw(fieldset, t, z, y, x, particles):
+    return fieldset.UVW[t, z, y, x, particles]
+
+
 def AdvectionRK2(particles, fieldset):  # pragma: no cover
     """Advection of particles using second-order Runge-Kutta integration."""
-    (u1, v1) = fieldset.UV[particles]
-    x1 = particles.x + u1 * 0.5 * particles.dt
-    y1 = particles.y + v1 * 0.5 * particles.dt
-    (u2, v2) = fieldset.UV[particles.t + 0.5 * particles.dt, particles.z, y1, x1, particles]
-    particles.dx += u2 * particles.dt
-    particles.dy += v2 * particles.dt
+    u, v = RK2(particles, fieldset, rhs_uv)
+    particles.dx += u * particles.dt
+    particles.dy += v * particles.dt
 
 
 def AdvectionRK2_3D(particles, fieldset):  # pragma: no cover
     """Advection of particles using second-order Runge-Kutta integration including vertical velocity."""
-    (u1, v1, w1) = fieldset.UVW[particles]
-    x1 = particles.x + u1 * 0.5 * particles.dt
-    y1 = particles.y + v1 * 0.5 * particles.dt
-    z1 = particles.z + w1 * 0.5 * particles.dt
-    (u2, v2, w2) = fieldset.UVW[particles.t + 0.5 * particles.dt, z1, y1, x1, particles]
-    particles.dx += u2 * particles.dt
-    particles.dy += v2 * particles.dt
-    particles.dz += w2 * particles.dt
+    u, v, w = RK2(particles, fieldset, rhs_uvw)
+    particles.dx += u * particles.dt
+    particles.dy += v * particles.dt
+    particles.dz += w * particles.dt
 
 
 def AdvectionRK4(particles, fieldset):  # pragma: no cover
     """Advection of particles using fourth-order Runge-Kutta integration."""
-    (u1, v1) = fieldset.UV[particles]
-    x1 = particles.x + u1 * 0.5 * particles.dt
-    y1 = particles.y + v1 * 0.5 * particles.dt
-    (u2, v2) = fieldset.UV[particles.t + 0.5 * particles.dt, particles.z, y1, x1, particles]
-    x2 = particles.x + u2 * 0.5 * particles.dt
-    y2 = particles.y + v2 * 0.5 * particles.dt
-    (u3, v3) = fieldset.UV[particles.t + 0.5 * particles.dt, particles.z, y2, x2, particles]
-    x3 = particles.x + u3 * particles.dt
-    y3 = particles.y + v3 * particles.dt
-    (u4, v4) = fieldset.UV[particles.t + particles.dt, particles.z, y3, x3, particles]
-    particles.dx += (u1 + 2 * u2 + 2 * u3 + u4) / 6.0 * particles.dt
-    particles.dy += (v1 + 2 * v2 + 2 * v3 + v4) / 6.0 * particles.dt
+    u, v = RK4(particles, fieldset, rhs_uv)
+    particles.dx += u * particles.dt
+    particles.dy += v * particles.dt
 
 
 def AdvectionRK4_3D(particles, fieldset):  # pragma: no cover
     """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity."""
-    (u1, v1, w1) = fieldset.UVW[particles]
-    x1 = particles.x + u1 * 0.5 * particles.dt
-    y1 = particles.y + v1 * 0.5 * particles.dt
-    z1 = particles.z + w1 * 0.5 * particles.dt
-    (u2, v2, w2) = fieldset.UVW[particles.t + 0.5 * particles.dt, z1, y1, x1, particles]
-    x2 = particles.x + u2 * 0.5 * particles.dt
-    y2 = particles.y + v2 * 0.5 * particles.dt
-    z2 = particles.z + w2 * 0.5 * particles.dt
-    (u3, v3, w3) = fieldset.UVW[particles.t + 0.5 * particles.dt, z2, y2, x2, particles]
-    x3 = particles.x + u3 * particles.dt
-    y3 = particles.y + v3 * particles.dt
-    z3 = particles.z + w3 * particles.dt
-    (u4, v4, w4) = fieldset.UVW[particles.t + particles.dt, z3, y3, x3, particles]
-    particles.dx += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * particles.dt
-    particles.dy += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * particles.dt
-    particles.dz += (w1 + 2 * w2 + 2 * w3 + w4) / 6 * particles.dt
+    u, v, w = RK4(particles, fieldset, rhs_uvw)
+    particles.dx += u * particles.dt
+    particles.dy += v * particles.dt
+    particles.dz += w * particles.dt
 
 
 def AdvectionEE(particles, fieldset):  # pragma: no cover

src/parcels/kernels/_integration.py

@@ -0,0 +1,56 @@
+"""Collection of time integrators for use in Parcels Kernels"""
+
+
+def _validate_rhs_output(output, integrator_name):
+    """Validate that an rhs function returned a 2- or 3-component tuple (u, v) or (u, v, w)."""
+    if not isinstance(output, tuple):
+        raise TypeError(
+            f"{integrator_name}: rhs must return a tuple of (u, v) or (u, v, w), got {type(output).__name__}."
+        )
+    if len(output) not in (2, 3):
+        raise ValueError(
+            f"{integrator_name}: rhs must return a tuple of 2 (u, v) or 3 (u, v, w) components, got {len(output)}."
+        )
+
+
+def RK2(particles, fieldset, rhs):
+    z, y, x = particles.z, particles.y, particles.x
+    fields = rhs(fieldset, particles.t, z, y, x, particles)
+    _validate_rhs_output(fields, "RK2")
+    x = particles.x + fields[0] * 0.5 * particles.dt
+    y = particles.y + fields[1] * 0.5 * particles.dt
+    if len(fields) > 2:
+        z = particles.z + fields[2] * 0.5 * particles.dt
+    t = particles.t + 0.5 * particles.dt
+    return rhs(fieldset, t, z, y, x, particles)
+
+
+def RK4(particles, fieldset, rhs):
+    z, y, x = particles.z, particles.y, particles.x
+    k1 = rhs(fieldset, particles.t, z, y, x, particles)
+    _validate_rhs_output(k1, "RK4")
+    x = particles.x + k1[0] * 0.5 * particles.dt
+    y = particles.y + k1[1] * 0.5 * particles.dt
+    if len(k1) > 2:
+        z = particles.z + k1[2] * 0.5 * particles.dt
+    t = particles.t + 0.5 * particles.dt
+    k2 = rhs(fieldset, t, z, y, x, particles)
+    x = particles.x + k2[0] * 0.5 * particles.dt
+    y = particles.y + k2[1] * 0.5 * particles.dt
+    if len(k2) > 2:
+        z = particles.z + k2[2] * 0.5 * particles.dt
+    k3 = rhs(fieldset, t, z, y, x, particles)
+    x = particles.x + k3[0] * particles.dt
+    y = particles.y + k3[1] * particles.dt
+    if len(k3) > 2:
+        z = particles.z + k3[2] * particles.dt
+    t = particles.t + particles.dt
+    k4 = rhs(fieldset, t, z, y, x, particles)
+    if len(k4) == 2:
+        return ((k1[0] + 2 * k2[0] + 2 * k3[0] + k4[0]) / 6.0, (k1[1] + 2 * k2[1] + 2 * k3[1] + k4[1]) / 6.0)
+    if len(k4) == 3:
+        return (
+            (k1[0] + 2 * k2[0] + 2 * k3[0] + k4[0]) / 6.0,
+            (k1[1] + 2 * k2[1] + 2 * k3[1] + k4[1]) / 6.0,
+            (k1[2] + 2 * k2[2] + 2 * k3[2] + k4[2]) / 6.0,
+        )

tests/test_advection.py

@@ -28,6 +28,8 @@
 )
 from parcels._datasets.structured.generic import datasets_sgrid
 from parcels.kernels import (
+    RK2,
+    RK4,
     AdvectionDiffusionEM,
     AdvectionDiffusionM1,
     AdvectionEE,
@@ -232,6 +234,25 @@ def test_length1dimensions(u_value, x_slice, v_value, y_slice, w_value, z_slice)
         np.testing.assert_allclose(np.array([p.z - z0 for p in pset]), 4 * w_value, atol=1e-5)
 
 
+@pytest.mark.parametrize("npart", [1, 2, 3, 10])
+@pytest.mark.parametrize("integrator", [RK2, RK4])
+def test_advection_with_integrators_singlefield(npart, integrator):
+    """Test that the RK2 and RK4 integrators don't work with a single field (e.g. U)"""
+    ds = simple_UV_dataset(mesh="flat").rename({"U": "T"})
+    fset = parcels.FieldSet.from_sgrid_conventions(ds, mesh="flat")
+
+    pset = parcels.ParticleSet(fset, x=np.zeros(npart), y=np.arange(npart))
+
+    def rhs(fieldset, t, z, y, x, particles):
+        return fieldset.T[t, z, y, x, particles]
+
+    def SingleField(particles, fieldset):
+        integrator(particles, fieldset, rhs)
+
+    with pytest.raises(TypeError):
+        pset.execute(SingleField, runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s"))
+
+
 def test_radialrotation(npart=10):
     ds = radial_rotation_dataset()
     fieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh="flat")