diff --git a/docs/user_guide/examples/tutorial_dt_integrators.ipynb b/docs/user_guide/examples/tutorial_dt_integrators.ipynb index f586df4da..b0c6ebda3 100644 --- a/docs/user_guide/examples/tutorial_dt_integrators.ipynb +++ b/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, diff --git a/src/parcels/_datasets/remote.py b/src/parcels/_datasets/remote.py index 822b96b56..be96e8148 100644 --- a/src/parcels/_datasets/remote.py +++ b/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)), diff --git a/src/parcels/kernels/__init__.py b/src/parcels/kernels/__init__.py index 710fb633d..8ea1d0686 100644 --- a/src/parcels/kernels/__init__.py +++ b/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", ] diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index d3fe15bc3..03d3aa0f0 100644 --- a/src/parcels/kernels/_advection.py +++ b/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 diff --git a/src/parcels/kernels/_integration.py b/src/parcels/kernels/_integration.py new file mode 100644 index 000000000..44d1c9492 --- /dev/null +++ b/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, + ) diff --git a/tests/test_advection.py b/tests/test_advection.py index 451e5e485..cb85c7a99 100644 --- a/tests/test_advection.py +++ b/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")