From 9cfce555a0128bd2766027c115f5b6df82835dde Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 14 Jul 2026 13:09:36 +0200 Subject: [PATCH 1/6] Use RK2 for CROCO advection --- .../examples/tutorial_croco_3D.ipynb | 27 ++++++++------ src/parcels/kernels/__init__.py | 4 +-- src/parcels/kernels/_sigmagrids.py | 35 +++---------------- tests/test_sigmagrids.py | 4 +-- 4 files changed, 25 insertions(+), 45 deletions(-) diff --git a/docs/user_guide/examples/tutorial_croco_3D.ipynb b/docs/user_guide/examples/tutorial_croco_3D.ipynb index 52f4d3602..6eae15992 100644 --- a/docs/user_guide/examples/tutorial_croco_3D.ipynb +++ b/docs/user_guide/examples/tutorial_croco_3D.ipynb @@ -37,14 +37,11 @@ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import polars as pl\n", - "import xarray as xr\n", "\n", "import parcels\n", "import parcels.tutorial\n", "\n", - "ds_fields = parcels.tutorial.open_dataset(\"CROCOidealized_data/data\")\n", - "\n", - "ds_fields.load(); # Preload data to speed up access" + "ds_fields = parcels.tutorial.open_dataset(\"CROCOidealized_data/data\")" ] }, { @@ -92,7 +89,10 @@ "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n", "\n", "# Add the critical depth (`hc`) as context to the fieldset\n", - "fieldset.add_context(\"hc\", ds_fields.hc.item())" + "fieldset.add_context(\"hc\", ds_fields.hc.item())\n", + "\n", + "# Convert the FieldSet to windowed arrays for better performance\n", + "fieldset = fieldset.to_windowed_arrays()" ] }, { @@ -132,7 +132,7 @@ ")\n", "\n", "pset.execute(\n", - " [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],\n", + " [parcels.kernels.AdvectionRK2_3D_CROCO, DeleteParticle],\n", " runtime=np.timedelta64(50_000, \"s\"),\n", " dt=np.timedelta64(100, \"s\"),\n", " output_file=outputfile,\n", @@ -173,7 +173,7 @@ "ax.set_xlabel(\"X [km]\")\n", "ax.set_xlim(30, 170)\n", "ax.set_ylabel(\"Depth [m]\")\n", - "ax.set_title(\"Particles in idealized CROCO velocity field using AdvectionRK4_3D_CROCO\")\n", + "ax.set_title(\"Particles in idealized CROCO velocity field using AdvectionRK2_3D_CROCO\")\n", "plt.tight_layout()\n", "plt.show()" ] @@ -189,7 +189,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "It may be insightful to compare this 3D run with the `AdvectionRK4_3D_CROCO` kernel with a run where the vertical velocity (`W`) is set to zero. In that case, the particles will not stay on their initial depth levels but instead follow sigma-layers." + "It may be insightful to compare this 3D run with the `AdvectionRK2_3D_CROCO` kernel with a run where the vertical velocity (`W`) is set to zero. In that case, the particles will not stay on their initial depth levels but instead follow sigma-layers." ] }, { @@ -221,7 +221,7 @@ ")\n", "\n", "pset_noW.execute(\n", - " [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],\n", + " [parcels.kernels.AdvectionRK2_3D_CROCO, DeleteParticle],\n", " runtime=np.timedelta64(50_000, \"s\"),\n", " dt=np.timedelta64(100, \"s\"),\n", " output_file=outputfile,\n", @@ -273,9 +273,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "For Advection, you will need to use the `AdvectionRK4_3D_CROCO` kernel, which works slightly different from the normal 3D advection kernel because it converts the vertical velocity in sigma-units. The conversion from z to sigma is done at every time step, using the `convert_z_to_sigma_croco()` function.\n", + "For Advection, you will need to use the `AdvectionRK2_3D_CROCO` kernel, which works slightly different from the normal 3D advection kernel because it converts the vertical velocity in sigma-units. The conversion from z to sigma is done at every time step, using the `convert_z_to_sigma_croco()` function.\n", "\n", - "In particular, the following algorithm is used (note that the RK4 version is slightly more complex than this Euler-Forward version, but the idea is identical). Also note that the vertical velocity is linearly interpolated here, which gives much better results than the default C-grid interpolation.\n", + "In particular, the following algorithm is used (note that the RK2 version is slightly more complex than this Euler-Forward version, but the idea is identical). Also note that the vertical velocity is linearly interpolated here, which gives much better results than the default C-grid interpolation.\n", "\n", "```python\n", "sigma = particles.z / fieldset.h[particles.t, 0, particles.y, particles.x]\n", @@ -297,6 +297,11 @@ "z_new = sigma_new * fieldset.h[particles.t, 0, y_new, x_new, particles]\n", "```" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] } ], "metadata": { diff --git a/src/parcels/kernels/__init__.py b/src/parcels/kernels/__init__.py index f1c613086..710fb633d 100644 --- a/src/parcels/kernels/__init__.py +++ b/src/parcels/kernels/__init__.py @@ -13,7 +13,7 @@ DiffusionUniformKh, ) from ._sigmagrids import ( - AdvectionRK4_3D_CROCO, + AdvectionRK2_3D_CROCO, SampleOmegaCroco, convert_z_to_sigma_croco, ) @@ -32,7 +32,7 @@ "AdvectionDiffusionM1", "DiffusionUniformKh", # sigmagrids - "AdvectionRK4_3D_CROCO", + "AdvectionRK2_3D_CROCO", "SampleOmegaCroco", "convert_z_to_sigma_croco", ] diff --git a/src/parcels/kernels/_sigmagrids.py b/src/parcels/kernels/_sigmagrids.py index 62dead521..07b3af9e8 100644 --- a/src/parcels/kernels/_sigmagrids.py +++ b/src/parcels/kernels/_sigmagrids.py @@ -35,9 +35,8 @@ def SampleOmegaCroco(particles, fieldset): particles.omega = fieldset.omega[particles.t, sigma, particles.y, particles.x, particles] -# TODO change to RK2 (once RK4 yields same results as v3) -def AdvectionRK4_3D_CROCO(particles, fieldset): # pragma: no cover - """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity. +def AdvectionRK2_3D_CROCO(particles, fieldset): # pragma: no cover + """Advection of particles using second-order Runge-Kutta integration including vertical velocity. This kernel assumes the vertical velocity is the 'w' field from CROCO output and works on sigma-layers. It also uses linear interpolation of the W field, which gives much better results than the default C-grid interpolation. """ @@ -57,31 +56,7 @@ def AdvectionRK4_3D_CROCO(particles, fieldset): # pragma: no cover (u2, v2) = fieldset.UV[particles.t + 0.5 * dt, sig1, y1, x1, particles] w2 = fieldset.W[particles.t + 0.5 * dt, sig1, y1, x1, particles] w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y1, x1] - x2 = particles.x + u2 * 0.5 * dt - y2 = particles.y + v2 * 0.5 * dt - sig_dep2 = sigma + w2 * 0.5 * dt - dep2 = sig_dep2 * fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y2, x2] - sig2 = convert_z_to_sigma_croco(fieldset, particles.t + 0.5 * dt, dep2, y2, x2, particles) - (u3, v3) = fieldset.UV[particles.t + 0.5 * dt, sig2, y2, x2, particles] - w3 = fieldset.W[particles.t + 0.5 * dt, sig2, y2, x2, particles] - w3 *= sig_dep2 / fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y2, x2] - x3 = particles.x + u3 * dt - y3 = particles.y + v3 * dt - sig_dep3 = sigma + w3 * dt - dep3 = sig_dep3 * fieldset.h[particles.t + dt, np.zeros_like(particles.z), y3, x3] - - sig3 = convert_z_to_sigma_croco(fieldset, particles.t + dt, dep3, y3, x3, particles) - (u4, v4) = fieldset.UV[particles.t + dt, sig3, y3, x3, particles] - w4 = fieldset.W[particles.t + dt, sig3, y3, x3, particles] - w4 *= sig_dep3 / fieldset.h[particles.t + dt, np.zeros_like(particles.z), y3, x3] - x4 = particles.x + u4 * dt - y4 = particles.y + v4 * dt - sig_dep4 = sigma + w4 * dt - - dep4 = sig_dep4 * fieldset.h[particles.t + dt, np.zeros_like(particles.z), y4, x4] - particles.dx += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt - particles.dy += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt - particles.dz += ( - (dep1 - particles.z) * 2 + 2 * (dep2 - particles.z) * 2 + 2 * (dep3 - particles.z) + dep4 - particles.z - ) / 6 + particles.dx += u2 * dt + particles.dy += v2 * dt + particles.dz += w2 * dt diff --git a/tests/test_sigmagrids.py b/tests/test_sigmagrids.py index ec05b81d0..375810c82 100644 --- a/tests/test_sigmagrids.py +++ b/tests/test_sigmagrids.py @@ -3,7 +3,7 @@ import parcels import parcels.tutorial from parcels import Particle, ParticleSet, Variable -from parcels.kernels import AdvectionRK4_3D_CROCO, SampleOmegaCroco, convert_z_to_sigma_croco +from parcels.kernels import AdvectionRK2_3D_CROCO, SampleOmegaCroco, convert_z_to_sigma_croco def test_conversion_3DCROCO(): @@ -71,7 +71,7 @@ def test_advection_3DCROCO(): pset = ParticleSet(fieldset=fieldset, pclass=pclass, x=X, y=Y, z=Z) pset.execute( - [AdvectionRK4_3D_CROCO, SampleOmegaCroco], runtime=np.timedelta64(runtime, "s"), dt=np.timedelta64(100, "s") + [AdvectionRK2_3D_CROCO, SampleOmegaCroco], runtime=np.timedelta64(runtime, "s"), dt=np.timedelta64(100, "s") ) np.testing.assert_allclose(pset.z, Z.flatten(), atol=5) # TODO lower this atol np.testing.assert_allclose(pset.x, [x + runtime for x in X.flatten()], atol=1e-3) From 64f0eaba46fdaa1d422ae3de27e59748826847c9 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 14 Jul 2026 14:38:25 +0200 Subject: [PATCH 2/6] Filter warnings and fix vertical movement in CROCO --- docs/user_guide/v4-migration.md | 2 +- src/parcels/kernels/_sigmagrids.py | 53 +++++++++++++++++++----------- 2 files changed, 34 insertions(+), 21 deletions(-) diff --git a/docs/user_guide/v4-migration.md b/docs/user_guide/v4-migration.md index 323667f21..8451f89a3 100644 --- a/docs/user_guide/v4-migration.md +++ b/docs/user_guide/v4-migration.md @@ -15,7 +15,7 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat - `math` functions should be replaced with array compatible equivalents (e.g., `math.sin` -> `np.sin`). Instead of `ParcelsRandom` you should use numpy's random functions. - `particle.depth` has been changed to `particles.z` to be consistent with the [CF conventions for trajectory data](https://cfconventions.org/cf-conventions/cf-conventions.html#trajectory-data), and to make Parcels also generalizable to atmospheric contexts. - The `InteractionKernel` class has been removed. Since normal Kernels now have access to _all_ particles, particle-particle interaction can be performed within normal Kernels. -- Users need to explicitly use `convert_z_to_sigma_croco` in sampling kernels (such as the `AdvectionRK4_3D_CROCO` or `SampleOMegaCroco` kernels) when working with CROCO data, as the automatic conversion from depth to sigma grids under the hood has been removed. +- Users need to explicitly use `convert_z_to_sigma_croco` in sampling kernels (such as the `AdvectionRK2_3D_CROCO` or `SampleOMegaCroco` kernels) when working with CROCO data, as the automatic conversion from depth to sigma grids under the hood has been removed. - We added a new AdvectionRK2 Kernel. The AdvectionRK4 kernel is still available, but RK2 is now the recommended default advection scheme as it is faster while the accuracy is comparable for most applications. See also the Choosing an integration method tutorial. - Functions shouldn't be converted to Kernels before adding to a pset.execute() call. Instead, simply pass the function(s) as a list to pset.execute(). - Kernel variables `time`, `lat` and `lon` have been renamed to `t`, `y` and `x`, and `dlat` and `dlon` have been renamed to `dy` and `dx`. These changes are also reflected on the ParticleSet as well as the ParticleFile output. diff --git a/src/parcels/kernels/_sigmagrids.py b/src/parcels/kernels/_sigmagrids.py index 07b3af9e8..72066e596 100644 --- a/src/parcels/kernels/_sigmagrids.py +++ b/src/parcels/kernels/_sigmagrids.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np from parcels.kernels._advection import _constrain_dt_to_within_time_interval @@ -40,23 +42,34 @@ def AdvectionRK2_3D_CROCO(particles, fieldset): # pragma: no cover This kernel assumes the vertical velocity is the 'w' field from CROCO output and works on sigma-layers. It also uses linear interpolation of the W field, which gives much better results than the default C-grid interpolation. """ - dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.t, particles.dt) - sigma = particles.z / fieldset.h[particles.t, np.zeros_like(particles.z), particles.y, particles.x] - - sig = convert_z_to_sigma_croco(fieldset, particles.t, particles.z, particles.y, particles.x, particles) - (u1, v1) = fieldset.UV[particles.t, sig, particles.y, particles.x, particles] - w1 = fieldset.W[particles.t, sig, particles.y, particles.x, particles] - w1 *= sigma / fieldset.h[particles.t, np.zeros_like(particles.z), particles.y, particles.x] - x1 = particles.x + u1 * 0.5 * dt - y1 = particles.y + v1 * 0.5 * dt - sig_dep1 = sigma + w1 * 0.5 * dt - dep1 = sig_dep1 * fieldset.h[particles.t, np.zeros_like(particles.z), y1, x1] - - sig1 = convert_z_to_sigma_croco(fieldset, particles.t + 0.5 * dt, dep1, y1, x1, particles) - (u2, v2) = fieldset.UV[particles.t + 0.5 * dt, sig1, y1, x1, particles] - w2 = fieldset.W[particles.t + 0.5 * dt, sig1, y1, x1, particles] - w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y1, x1] - - particles.dx += u2 * dt - particles.dy += v2 * dt - particles.dz += w2 * dt + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"^Sampling of velocities should normally be done using fieldset\.UV or fieldset\.UVW object; tread carefully$", + category=RuntimeWarning, + ) # Needed because of linear sampling of W with sigma conversion + + dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.t, particles.dt) + sigma = particles.z / fieldset.h[particles.t, np.zeros_like(particles.z), particles.y, particles.x] + + sig = convert_z_to_sigma_croco(fieldset, particles.t, particles.z, particles.y, particles.x, particles) + (u1, v1) = fieldset.UV[particles.t, sig, particles.y, particles.x, particles] + w1 = fieldset.W[particles.t, sig, particles.y, particles.x, particles] + w1 *= sigma / fieldset.h[particles.t, np.zeros_like(particles.z), particles.y, particles.x] + x1 = particles.x + u1 * 0.5 * dt + y1 = particles.y + v1 * 0.5 * dt + sig_dep1 = sigma + w1 * 0.5 * dt + dep1 = sig_dep1 * fieldset.h[particles.t, np.zeros_like(particles.z), y1, x1] + + sig1 = convert_z_to_sigma_croco(fieldset, particles.t + 0.5 * dt, dep1, y1, x1, particles) + (u2, v2) = fieldset.UV[particles.t + 0.5 * dt, sig1, y1, x1, particles] + w2 = fieldset.W[particles.t + 0.5 * dt, sig1, y1, x1, particles] + w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y1, x1] + x2 = particles.x + u2 * 0.5 * dt + y2 = particles.y + v2 * 0.5 * dt + sig_dep2 = sigma + w2 * 0.5 * dt + dep2 = sig_dep2 * fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y2, x2] + + particles.dx += u2 * dt + particles.dy += v2 * dt + particles.dz += (dep1 - particles.z) + (dep2 - particles.z) From ad334e2844631cfaea95e69d3f980c48a0f55870 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 14 Jul 2026 17:21:57 +0200 Subject: [PATCH 3/6] Compute variables that have a mockT dimension --- src/parcels/_core/_windowed_array.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/parcels/_core/_windowed_array.py b/src/parcels/_core/_windowed_array.py index 77cb3cf1c..49526e687 100644 --- a/src/parcels/_core/_windowed_array.py +++ b/src/parcels/_core/_windowed_array.py @@ -108,4 +108,6 @@ def maybe_windowed(data: xr.DataArray, max_levels: int | None = None): return data if data.dims and data.dims[0] == "time" and is_dask_collection(data.data): return WindowedArray(data, max_levels=max_levels) + elif data.dims and data.dims[0] == "mockT" and is_dask_collection(data.data): + return data.compute() return data From 6034d1fca325756c753433bcc700cc23d0b9ad65 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 14 Jul 2026 17:25:46 +0200 Subject: [PATCH 4/6] Explicitly loading lon and lat --- src/parcels/_core/xgrid.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py index 91a48af5d..36a908475 100644 --- a/src/parcels/_core/xgrid.py +++ b/src/parcels/_core/xgrid.py @@ -7,6 +7,7 @@ import numpy.typing as npt import xarray as xr import xgcm +from dask import is_dask_collection import parcels._sgrid as sgrid import parcels._typing as ptyping @@ -207,6 +208,9 @@ def lon(self): _ = self.xgcm_grid.axes["X"] except KeyError: return np.zeros(1) + # ensure lon is loaded into memory for dask-backed datasets, as it is used in the search method + if is_dask_collection(self._ds["lon"].data): + self._ds["lon"].load() return self._ds["lon"].values @property @@ -221,6 +225,9 @@ def lat(self): _ = self.xgcm_grid.axes["Y"] except KeyError: return np.zeros(1) + # ensure lat is loaded into memory for dask-backed datasets, as it is used in the search method + if is_dask_collection(self._ds["lat"].data): + self._ds["lat"].load() return self._ds["lat"].values @property From 3b20295cc4799930194fcf60614b6f8d38cd1f11 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 14 Jul 2026 17:32:09 +0200 Subject: [PATCH 5/6] Also adding to_windowed_arrays() in second experiment --- docs/user_guide/examples/tutorial_croco_3D.ipynb | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/user_guide/examples/tutorial_croco_3D.ipynb b/docs/user_guide/examples/tutorial_croco_3D.ipynb index 6eae15992..3458e3df1 100644 --- a/docs/user_guide/examples/tutorial_croco_3D.ipynb +++ b/docs/user_guide/examples/tutorial_croco_3D.ipynb @@ -205,6 +205,7 @@ "fieldset_noW = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n", "fieldset_noW.W.data[:] = 0.0\n", "fieldset_noW.add_context(\"hc\", ds_fields.hc.item())\n", + "fieldset_noW = fieldset_noW.to_windowed_arrays()\n", "\n", "X, Z = np.meshgrid(\n", " [40e3, 80e3, 120e3],\n", From a3ee2e930489b8590cd073cde8e2df10e7519136 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 14 Jul 2026 17:34:18 +0200 Subject: [PATCH 6/6] Update _sigmagrids.py --- src/parcels/kernels/_sigmagrids.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parcels/kernels/_sigmagrids.py b/src/parcels/kernels/_sigmagrids.py index 72066e596..4f321b71a 100644 --- a/src/parcels/kernels/_sigmagrids.py +++ b/src/parcels/kernels/_sigmagrids.py @@ -14,7 +14,7 @@ def convert_z_to_sigma_croco(fieldset, t, z, y, x, particle): h = fieldset.h.eval(t, np.zeros_like(z), y, x, particles=particle) zeta = fieldset.zeta.eval(t, np.zeros_like(z), y, x, particles=particle) sigma_levels = fieldset.U.grid.depth - cs_w = fieldset.Cs_w.data[0, :, 0, 0].values + cs_w = fieldset.Cs_w.data.values.flatten() z0 = fieldset.hc * sigma_levels[None, :] + (h[:, None] - fieldset.hc) * cs_w[None, :] zvec = z0 + zeta[:, None] * (1.0 + (z0 / h[:, None]))