Merge branch 'main' into remove_constrain_dt_to_within_time_interval

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/_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

src/parcels/_core/model.py

@@ -7,6 +7,7 @@
 import cf_xarray  # noqa: F401
 import uxarray as ux
 import xarray as xr
+from dask import is_dask_collection
 
 import parcels._sgrid as sgrid
 import parcels._typing as ptyping
@@ -102,6 +103,12 @@ def to_windowed_arrays(self, *, max_levels: int | None = None) -> Self:
             and ``max_levels`` then bounds its size.
         """
         windowed = self.__dict__.setdefault("_windowed", {})
+        for dim in ["lon", "lat"]:
+            # ensure lon and lat are loaded into memory for dask-backed datasets
+            if is_dask_collection(self.data[dim]):
+                self.data[dim].load()
+        if "depth" in self.data.dims and is_dask_collection(self.data["depth"]):
+            self.data["depth"].load()
         for name in self.scalar_field_names:
             current = windowed.get(name, self.data[name])
             windowed[name] = maybe_windowed(current, max_levels=max_levels)

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/_reprs.py

@@ -8,7 +8,10 @@
 
 import numpy as np
 import xarray as xr
+import zarr
+from dask.base import is_dask_collection
 
+from parcels._core._windowed_array import WindowedArray
 from parcels._python import isinstance_noimport
 
 if TYPE_CHECKING:
@@ -190,13 +193,15 @@ class _FieldSetDescriptionRow:
     model_id: int | None
     name: str
     interp_method_or_value: str
+    backend: str | None = None
 
     def to_dict(self) -> dict[str, str]:
         return {
             "Name": self.name,
             "Type": self.type_,
             "Grid number": str(self.model_id) if self.model_id is not None else "-",
             "Interp method / value": self.interp_method_or_value,
+            "Parcels backend": self.backend if self.backend is not None else "-",
         }
 
 
@@ -213,6 +218,22 @@ def _print_time_interval(time_interval: TimeInterval | None) -> str:
     return repr((time_interval.left, time_interval.right))
 
 
+def _field_backend(field: Field | VectorField) -> str | None:
+    if hasattr(field, "data"):
+        if isinstance(field.data, WindowedArray):
+            return "WindowedArray"
+        elif is_dask_collection(field.data.data):
+            return "Dask"
+        elif isinstance(field.data.variable._data, zarr.Array):
+            return "Zarr"
+        elif isinstance(field.data.data, np.ndarray):
+            return "NumPy"
+        else:
+            return type(field.data).__name__
+    else:
+        return None
+
+
 def fieldset_describe(fieldset: FieldSet) -> str:
     rows: list[_FieldSetDescriptionRow] = []
     models: dict[int, int] = {}  # mapping of memory ID to a human readable ID
@@ -235,6 +256,7 @@ def fieldset_describe(fieldset: FieldSet) -> str:
                 model_id=model_id,
                 name=field.name,
                 interp_method_or_value=repr(field.interp_method),
+                backend=_field_backend(field),
             )
         )
     for k, v in fieldset.context.items():
@@ -244,6 +266,7 @@ def fieldset_describe(fieldset: FieldSet) -> str:
                 model_id=None,
                 name=k,
                 interp_method_or_value=repr(v),
+                backend=None,
             )
         )
     return (

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
 
 
@@ -10,7 +12,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]))
@@ -33,52 +35,38 @@ 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.
     """
-    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 * particles.dt
-    y1 = particles.y + v1 * 0.5 * particles.dt
-    sig_dep1 = sigma + w1 * 0.5 * particles.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 * particles.dt, dep1, y1, x1, particles)
-    (u2, v2) = fieldset.UV[particles.t + 0.5 * particles.dt, sig1, y1, x1, particles]
-    w2 = fieldset.W[particles.t + 0.5 * particles.dt, sig1, y1, x1, particles]
-    w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y1, x1]
-    x2 = particles.x + u2 * 0.5 * particles.dt
-    y2 = particles.y + v2 * 0.5 * particles.dt
-    sig_dep2 = sigma + w2 * 0.5 * particles.dt
-    dep2 = sig_dep2 * fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y2, x2]
+        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 * particles.dt, dep2, y2, x2, particles)
-    (u3, v3) = fieldset.UV[particles.t + 0.5 * particles.dt, sig2, y2, x2, particles]
-    w3 = fieldset.W[particles.t + 0.5 * particles.dt, sig2, y2, x2, particles]
-    w3 *= sig_dep2 / fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y2, x2]
-    x3 = particles.x + u3 * particles.dt
-    y3 = particles.y + v3 * particles.dt
-    sig_dep3 = sigma + w3 * particles.dt
-    dep3 = sig_dep3 * fieldset.h[particles.t + particles.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 * particles.dt
+        y1 = particles.y + v1 * 0.5 * particles.dt
+        sig_dep1 = sigma + w1 * 0.5 * particles.dt
+        dep1 = sig_dep1 * fieldset.h[particles.t, np.zeros_like(particles.z), y1, x1]
 
-    sig3 = convert_z_to_sigma_croco(fieldset, particles.t + particles.dt, dep3, y3, x3, particles)
-    (u4, v4) = fieldset.UV[particles.t + particles.dt, sig3, y3, x3, particles]
-    w4 = fieldset.W[particles.t + particles.dt, sig3, y3, x3, particles]
-    w4 *= sig_dep3 / fieldset.h[particles.t + particles.dt, np.zeros_like(particles.z), y3, x3]
-    x4 = particles.x + u4 * particles.dt
-    y4 = particles.y + v4 * particles.dt
-    sig_dep4 = sigma + w4 * particles.dt
+        sig1 = convert_z_to_sigma_croco(fieldset, particles.t + 0.5 * particles.dt, dep1, y1, x1, particles)
+        (u2, v2) = fieldset.UV[particles.t + 0.5 * particles.dt, sig1, y1, x1, particles]
+        w2 = fieldset.W[particles.t + 0.5 * particles.dt, sig1, y1, x1, particles]
+        w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y1, x1]
+        x2 = particles.x + u2 * 0.5 * particles.dt
+        y2 = particles.y + v2 * 0.5 * particles.dt
+        sig_dep2 = sigma + w2 * 0.5 * particles.dt
+        dep2 = sig_dep2 * fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y2, x2]
 
-    dep4 = sig_dep4 * fieldset.h[particles.t + particles.dt, np.zeros_like(particles.z), y4, x4]
-    particles.dx += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * particles.dt
-    particles.dy += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * particles.dt
-    particles.dz += (
-        (dep1 - particles.z) * 2 + 2 * (dep2 - particles.z) * 2 + 2 * (dep3 - particles.z) + dep4 - particles.z
-    ) / 6
+        particles.dx += u2 * particles.dt
+        particles.dy += v2 * particles.dt
+        particles.dz += (dep1 - particles.z) + (dep2 - particles.z)