Merge branch 'v4-dev' into v4-argo-float-tutorial

Author: erikvansebille

README.md

@@ -10,6 +10,9 @@
 [![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/5353/badge)](https://bestpractices.coreinfrastructure.org/projects/5353)
 [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/OceanParcels/parcels/main?labpath=docs%2Fexamples%2Fparcels_tutorial.ipynb)
 
+> [!WARNING]
+> This branch is `v4-dev` - version 4 of Parcels which is in active development. See `main` (or the tags) to browse stable versions of Parcels.
+
 **Parcels** (**P**robably **A** **R**eally **C**omputationally **E**fficient **L**agrangian **S**imulator) is a set of Python classes and methods to create customisable particle tracking simulations using output from Ocean Circulation models. Parcels can be used to track passive and active particulates such as water, plankton, [plastic](http://www.topios.org/) and [fish](https://github.com/Jacketless/IKAMOANA).
 
 ![Arctic-SO-medusaParticles](https://github.com/OceanParcels/oceanparcels_website/blob/main/images/homepage.gif)

docs/examples/tutorial_kernelloop.ipynb

@@ -26,17 +26,7 @@
     "\n",
     "When you run a Parcels simulation (i.e. a call to `pset.execute()`), the Kernel loop is the main part of the code that is executed. This part of the code loops through all particles and executes the Kernels that are defined for each particle.\n",
     "\n",
-    "In order to make sure that the displacements of a particle in the different Kernels can be summed, all Kernels add to a _change_ in position (`particle_dlon`, `particle_dlat`, and `particle_ddepth`). This is important, because there are situations where movement kernels would otherwise not commute. Take the example of advecting particles by currents _and_ winds. If the particle would first be moved by the currents and then by the winds, the result could be different from first moving by the winds and then by the currents. Instead, by adding the changes in position, the ordering of the Kernels has no consequence on the particle displacement."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "<div class=\"alert alert-info\">\n",
-    "\n",
-    "__Note__ that the variables `particle_dlon`, `particle_dlat`, and `particle_ddepth` are defined by the wrapper function that Parcels generates for each Kernel. This is why you don't have to define these variables yourself when writing a Kernel. See [here](https://github.com/OceanParcels/parcels/blob/daa4b062ed8ae0b2be3d87367d6b45599d6f95db/parcels/kernel.py#L277-L294) for the implementation of the wrapper functions.\n",
-    "</div>"
+    "In order to make sure that the displacements of a particle in the different Kernels can be summed, all Kernels add to a _change_ in position (`particles.dlon`, `particles.dlat`, and `particles.ddepth`). This is important, because there are situations where movement kernels would otherwise not commute. Take the example of advecting particles by currents _and_ winds. If the particle would first be moved by the currents and then by the winds, the result could be different from first moving by the winds and then by the currents. Instead, by adding the changes in position, the ordering of the Kernels has no consequence on the particle displacement."
    ]
   },
   {
@@ -52,29 +42,25 @@
    "source": [
     "Below is a structured overview of the Kernel loop is implemented. Note that this is for longitude only, but the same process is applied for latitude and depth.\n",
     "\n",
-    "1. Define an extra variable `particle.lon_nextloop` for each particle, which is the longitude at the end of the Kernel loop. Inititalise it to `particle.lon`.\n",
-    "\n",
-    "2. Also define an extra variable `particle.time_nextloop` for each particle, which is the time at the end of the Kernel loop. Inititalise it to `particle.time`.\n",
+    "1. Initialise an extra Variable `particles.lon=0` and `particles.time_nextloop = particles.time`\n",
     "\n",
-    "3. Within the Kernel loop, for each particle:<br>\n",
+    "2. Within the Kernel loop, for each particle:<br>\n",
     "\n",
-    "    1. Update `particle.lon` with `particle.lon_nextloop`<br>\n",
+    "    1. Update `particles.lon += particles.dlon`<br>\n",
     "\n",
-    "    2. Update `particle.time` with `particle.time_nextloop`<br>\n",
+    "    2. Set variable `particles.dlon = 0`<br>\n",
     "\n",
-    "    3. Set local variable `particle_dlon = 0`<br>\n",
+    "    3. Update `particles.time = particles.time_nextloop`\n",
     "\n",
     "    4. For each Kernel in the list of Kernels:\n",
     "        \n",
     "        1. Execute the Kernel\n",
     "        \n",
-    "        2. Update `particle_dlon` by adding the change in longitude, if needed<br>\n",
+    "        2. Update `particles.dlon` by adding the change in longitude, if needed<br>\n",
     "\n",
-    "    5. Update `particle.lon_nextloop` with `particle.lon + particle_dlon`<br>\n",
-    "    \n",
-    "    6. Update `particle.time_nextloop` with `particle.time + particle.dt`<br>\n",
+    "    5. Update `particles.time_nextloop += particles.dt`<br>\n",
     "\n",
-    "    7. If `outputdt` is a multiple of `particle.time`, write `particle.lon` and `particle.time` to zarr output file<br>\n",
+    "    6. If `outputdt` is a multiple of `particle.time`, write `particle.lon` and `particle.time` to zarr output file<br>\n",
     "\n",
     "Besides having commutable Kernels, the main advantage of this implementation is that, when using Field Sampling with e.g. `particle.temp = fieldset.Temp[particle.time, particle.depth, particle.lat, particle.lon]`, the particle location stays the same throughout the entire Kernel loop. Additionally, this implementation ensures that the particle location is the same as the location of the sampled field in the output file."
    ]
@@ -268,12 +254,12 @@
     "### 1. Avoid updating particle locations directly in Kernels\n",
     "It is better not to update `particle.lon` directly in a Kernel, as it can interfere with the loop above. Assigning a value to `particle.lon` in a Kernel will throw a warning. \n",
     "\n",
-    "Instead, update the local variable `particle_dlon`.\n",
+    "Instead, update the local variable `particle.dlon`.\n",
     "\n",
     "### 2. Be careful with updating particle variables that do not depend on Fields.\n",
     "While assigning the interpolated value of a `Field` to a Particle goes well in the loop above, this is not necessarily so for assigning other attributes. For example, a line like `particle.age += particle.dt` is executed directly so may result in the age being `dt` at `time = 0` in the output file. \n",
     "\n",
-    "A workaround is to either initialise the age to `-dt`, or to increase the `age` only when `particle.time > 0` (using an `if` statement).\n",
+    "A workaround is to either initialise the age to `-dt`, or to increase the `age` only when `particle.time > 0` (using an `np.where` statement).\n",
     "\n",
     "\n",
     "### 3. The last time is not written to file\n",

parcels/__init__.py

@@ -2,6 +2,8 @@
 
 __version__ = version
 
+import warnings as _warnings
+
 from parcels.application_kernels import *
 from parcels.field import *
 from parcels.fieldset import *
@@ -11,3 +13,9 @@
 from parcels.particlefile import *
 from parcels.particleset import *
 from parcels.tools import *
+
+_warnings.warn(
+    "This is an alpha version of Parcels v4. The API is not stable and may change without deprecation warnings.",
+    UserWarning,
+    stacklevel=2,
+)

parcels/_index_search.py

@@ -5,18 +5,17 @@
 
 import numpy as np
 
-from parcels._typing import Mesh
-from parcels.tools.statuscodes import (
-    _raise_grid_searching_error,
-    _raise_time_extrapolation_error,
-)
+from parcels.tools.statuscodes import _raise_time_extrapolation_error
 
 if TYPE_CHECKING:
     from parcels.xgrid import XGrid
 
     from .field import Field
 
 
+GRID_SEARCH_ERROR = -3
+
+
 def _search_time_index(field: Field, time: datetime):
     """Find and return the index and relative coordinate in the time array associated with a given time.
 
@@ -40,13 +39,7 @@ def _search_time_index(field: Field, time: datetime):
     return np.atleast_1d(tau), np.atleast_1d(ti)
 
 
-def _search_indices_curvilinear_2d(
-    grid: XGrid, y: float, x: float, yi_guess: int | None = None, xi_guess: int | None = None
-):  # TODO fix typing instructions to make clear that y, x etc need to be ndarrays
-    yi, xi = yi_guess, xi_guess
-    if yi is None or xi is None:
-        yi, xi = grid.get_spatial_hash().query(y, x)
-
+def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: np.ndarray):
     xsi = eta = -1.0 * np.ones(len(x), dtype=float)
     invA = np.array(
         [
@@ -56,67 +49,60 @@ def _search_indices_curvilinear_2d(
             [1, -1, 1, -1],
         ]
     )
-    maxIterSearch = 1e6
-    it = 0
-    tol = 1.0e-10
-
-    # # ! Error handling for out of bounds
-    # TODO: Re-enable in some capacity
-    # if x < field.lonlat_minmax[0] or x > field.lonlat_minmax[1]:
-    #     if grid.lon[0, 0] < grid.lon[0, -1]:
-    #         _raise_grid_searching_error(y, x)
-    #     elif x < grid.lon[0, 0] and x > grid.lon[0, -1]:  # This prevents from crashing in [160, -160]
-    #         _raise_grid_searching_error(z, y, x)
-
-    # if y < field.lonlat_minmax[2] or y > field.lonlat_minmax[3]:
-    #     _raise_grid_searching_error(z, y, x)
-
-    while np.any(xsi < -tol) or np.any(xsi > 1 + tol) or np.any(eta < -tol) or np.any(eta > 1 + tol):
-        px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]])
-
-        py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]])
-        a = np.dot(invA, px)
-        b = np.dot(invA, py)
-
-        aa = a[3] * b[2] - a[2] * b[3]
-        bb = a[3] * b[0] - a[0] * b[3] + a[1] * b[2] - a[2] * b[1] + x * b[3] - y * a[3]
-        cc = a[1] * b[0] - a[0] * b[1] + x * b[1] - y * a[1]
-
-        det2 = bb * bb - 4 * aa * cc
-        with np.errstate(divide="ignore", invalid="ignore"):
-            det = np.where(det2 > 0, np.sqrt(det2), eta)
-
-            eta = np.where(abs(aa) < 1e-12, -cc / bb, np.where(det2 > 0, (-bb + det) / (2 * aa), eta))
-
-            xsi = np.where(
-                abs(a[1] + a[3] * eta) < 1e-12,
-                ((y - py[0]) / (py[1] - py[0]) + (y - py[3]) / (py[2] - py[3])) * 0.5,
-                (x - a[0] - a[2] * eta) / (a[1] + a[3] * eta),
-            )
-
-        xi = np.where(xsi < -tol, xi - 1, np.where(xsi > 1 + tol, xi + 1, xi))
-        yi = np.where(eta < -tol, yi - 1, np.where(eta > 1 + tol, yi + 1, yi))
-
-        (yi, xi) = _reconnect_bnd_indices(yi, xi, grid.ydim, grid.xdim, grid._mesh)
-        it += 1
-        if it > maxIterSearch:
-            print(f"Correct cell not found after {maxIterSearch} iterations")
-            _raise_grid_searching_error(0, y, x)
-    xsi = np.where(xsi < 0.0, 0.0, np.where(xsi > 1.0, 1.0, xsi))
-    eta = np.where(eta < 0.0, 0.0, np.where(eta > 1.0, 1.0, eta))
-
-    if np.any((xsi < 0) | (xsi > 1) | (eta < 0) | (eta > 1)):
-        _raise_grid_searching_error(y, x)
-    return (yi, eta, xi, xsi)
 
+    px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]])
+    py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]])
+
+    a, b = np.dot(invA, px), np.dot(invA, py)
+    aa = a[3] * b[2] - a[2] * b[3]
+    bb = a[3] * b[0] - a[0] * b[3] + a[1] * b[2] - a[2] * b[1] + x * b[3] - y * a[3]
+    cc = a[1] * b[0] - a[0] * b[1] + x * b[1] - y * a[1]
+    det2 = bb * bb - 4 * aa * cc
 
-def _reconnect_bnd_indices(yi: int, xi: int, ydim: int, xdim: int, mesh: Mesh):
-    xi = np.where(xi < 0, (xdim - 2) if mesh == "spherical" else 0, xi)
-    xi = np.where(xi > xdim - 2, 0 if mesh == "spherical" else (xdim - 2), xi)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        det = np.where(det2 > 0, np.sqrt(det2), eta)
+        eta = np.where(abs(aa) < 1e-12, -cc / bb, np.where(det2 > 0, (-bb + det) / (2 * aa), eta))
 
-    xi = np.where(yi > ydim - 2, xdim - xi if mesh == "spherical" else xi, xi)
+        xsi = np.where(
+            abs(a[1] + a[3] * eta) < 1e-12,
+            ((y - py[0]) / (py[1] - py[0]) + (y - py[3]) / (py[2] - py[3])) * 0.5,
+            (x - a[0] - a[2] * eta) / (a[1] + a[3] * eta),
+        )
 
-    yi = np.where(yi < 0, 0, yi)
-    yi = np.where(yi > ydim - 2, ydim - 2, yi)
+    is_in_cell = np.where((xsi >= 0) & (xsi <= 1) & (eta >= 0) & (eta <= 1), 1, 0)
 
-    return yi, xi
+    return is_in_cell, np.column_stack((xsi, eta))
+
+
+def _search_indices_curvilinear_2d(
+    grid: XGrid, y: np.ndarray, x: np.ndarray, yi_guess: np.ndarray | None = None, xi_guess: np.ndarray | None = None
+):
+    yi_guess = np.array(yi_guess)
+    xi_guess = np.array(xi_guess)
+    xi = np.full(len(x), GRID_SEARCH_ERROR, dtype=np.int32)
+    yi = np.full(len(y), GRID_SEARCH_ERROR, dtype=np.int32)
+    if np.any(xi_guess):
+        # If an initial guess is provided, we first perform a point in cell check for all guessed indices
+        is_in_cell, coords = curvilinear_point_in_cell(grid, y, x, yi_guess, xi_guess)
+        y_check = y[is_in_cell == 0]
+        x_check = x[is_in_cell == 0]
+        zero_indices = np.where(is_in_cell == 0)[0]
+    else:
+        # Otherwise, we need to check all points
+        y_check = y
+        x_check = x
+        coords = -1.0 * np.ones((len(y), 2), dtype=np.float32)
+        zero_indices = np.arange(len(y))
+
+    # If there are any points that were not found in the first step, we query the spatial hash for those points
+    if len(zero_indices) > 0:
+        yi_q, xi_q, coords_q = grid.get_spatial_hash().query(y_check, x_check)
+        # Only those points that were not found in the first step are updated
+        coords[zero_indices, :] = coords_q
+        yi[zero_indices] = yi_q
+        xi[zero_indices] = xi_q
+
+    xsi = coords[:, 0]
+    eta = coords[:, 1]
+
+    return (yi, eta, xi, xsi)

parcels/basegrid.py

@@ -54,9 +54,9 @@ def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int,
             - Unstructured grid: {"Z": (zi, zeta), "FACE": (fi, bcoords)}
 
             Where:
-            - index (int): The cell position of a particle along the given axis
+            - index (int): The cell position of the particles along the given axis
             - barycentric_coordinates (float or np.ndarray): The coordinates defining
-              a particle's position within the grid cell. For structured grids, this
+              the particles positions within the grid cell. For structured grids, this
               is a single coordinate per axis; for unstructured grids, this can be
               an array of coordinates for the face polygon.