Update CROCO tutorial (#2753)

* Use RK2 for CROCO advection

* Filter warnings and fix vertical movement in CROCO

* Compute variables that have a mockT dimension

* Explicitly loading lon and lat

* Also adding to_windowed_arrays() in second experiment

* Update _sigmagrids.py

Author: erikvansebille

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."
    ]
   },
   {
@@ -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",
@@ -221,7 +222,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 +274,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 +298,11 @@
     "z_new = sigma_new * fieldset.h[particles.t, 0, y_new, x_new, particles]\n",
     "```"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {

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.

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

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",
 ]

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
@@ -12,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]))
@@ -35,53 +37,39 @@ 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.
     """
-    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]
+    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
 
-    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]
+        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]
 
-    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]
+        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]
 
-    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
+        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]
 
-    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 += (dep1 - particles.z) + (dep2 - particles.z)

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)