Merge pull request #2201 from OceanParcels/v4-new-kernel-api

Update Kernel signature to (particles, fieldset)

Author: erikvansebille

docs/community/v4-migration.md

@@ -11,7 +11,8 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat
 - Sharing state between kernels must be done via the particle data (as the kernels are not combined under the hood anymore).
 - `particl_dlon`, `particle_dlat` etc have been renamed to `particle.dlon` and `particle.dlat`.
 - `particle.dt` is a np.timedelta64 object; be careful when multiplying `particle.dt` with a velocity, as its value may be cast to nanoseconds.
-- The `time` argument in the Kernel signature is now standard `None` (and may be removed in the Kernel API before release of v4), so can't be used. Use `particle.time` instead.
+- The `time` argument in the Kernel signature has been removed in the Kernel API, so can't be used. Use `particle.time` instead.
+- The `particle` argument in the Kernel signature has been renamed to `particles`.
 - `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.
 
 ## FieldSet

parcels/application_kernels/advection.py

@@ -8,7 +8,7 @@
 __all__ = ["AdvectionDiffusionEM", "AdvectionDiffusionM1", "DiffusionUniformKh"]
 
 
-def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
+def AdvectionDiffusionM1(particles, fieldset):  # pragma: no cover
     """Kernel for 2D advection-diffusion, solved using the Milstein scheme at first order (M1).
 
     Assumes that fieldset has fields `Kh_zonal` and `Kh_meridional`
@@ -23,30 +23,34 @@ def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
-    dt = particle.dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
+    dt = particles.dt / np.timedelta64(1, "s")  # TODO: improve API for converting dt to seconds
     # Wiener increment with zero mean and std of sqrt(dt)
     dWx = np.random.normal(0, np.sqrt(np.fabs(dt)))
     dWy = np.random.normal(0, np.sqrt(np.fabs(dt)))
 
-    Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres, particle]
-    Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres, particle]
+    Kxp1 = fieldset.Kh_zonal[particles.time, particles.depth, particles.lat, particles.lon + fieldset.dres, particles]
+    Kxm1 = fieldset.Kh_zonal[particles.time, particles.depth, particles.lat, particles.lon - fieldset.dres, particles]
     dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres)
 
-    u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon, particle]
-    bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon, particle])
+    u, v = fieldset.UV[particles.time, particles.depth, particles.lat, particles.lon, particles]
+    bx = np.sqrt(2 * fieldset.Kh_zonal[particles.time, particles.depth, particles.lat, particles.lon, particles])
 
-    Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon, particle]
-    Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon, particle]
+    Kyp1 = fieldset.Kh_meridional[
+        particles.time, particles.depth, particles.lat + fieldset.dres, particles.lon, particles
+    ]
+    Kym1 = fieldset.Kh_meridional[
+        particles.time, particles.depth, particles.lat - fieldset.dres, particles.lon, particles
+    ]
     dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres)
 
-    by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon, particle])
+    by = np.sqrt(2 * fieldset.Kh_meridional[particles.time, particles.depth, particles.lat, particles.lon, particles])
 
     # Particle positions are updated only after evaluating all terms.
-    particle.dlon += u * dt + 0.5 * dKdx * (dWx**2 + dt) + bx * dWx
-    particle.dlat += v * dt + 0.5 * dKdy * (dWy**2 + dt) + by * dWy
+    particles.dlon += u * dt + 0.5 * dKdx * (dWx**2 + dt) + bx * dWx
+    particles.dlat += v * dt + 0.5 * dKdy * (dWy**2 + dt) + by * dWy
 
 
-def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
+def AdvectionDiffusionEM(particles, fieldset):  # pragma: no cover
     """Kernel for 2D advection-diffusion, solved using the Euler-Maruyama scheme (EM).
 
     Assumes that fieldset has fields `Kh_zonal` and `Kh_meridional`
@@ -59,31 +63,35 @@ def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
-    dt = particle.dt / np.timedelta64(1, "s")
+    dt = particles.dt / np.timedelta64(1, "s")
     # Wiener increment with zero mean and std of sqrt(dt)
     dWx = np.random.normal(0, np.sqrt(np.fabs(dt)))
     dWy = np.random.normal(0, np.sqrt(np.fabs(dt)))
 
-    u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon, particle]
+    u, v = fieldset.UV[particles.time, particles.depth, particles.lat, particles.lon, particles]
 
-    Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres, particle]
-    Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres, particle]
+    Kxp1 = fieldset.Kh_zonal[particles.time, particles.depth, particles.lat, particles.lon + fieldset.dres, particles]
+    Kxm1 = fieldset.Kh_zonal[particles.time, particles.depth, particles.lat, particles.lon - fieldset.dres, particles]
     dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres)
     ax = u + dKdx
-    bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon, particle])
-
-    Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon, particle]
-    Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon, particle]
+    bx = np.sqrt(2 * fieldset.Kh_zonal[particles.time, particles.depth, particles.lat, particles.lon, particles])
+
+    Kyp1 = fieldset.Kh_meridional[
+        particles.time, particles.depth, particles.lat + fieldset.dres, particles.lon, particles
+    ]
+    Kym1 = fieldset.Kh_meridional[
+        particles.time, particles.depth, particles.lat - fieldset.dres, particles.lon, particles
+    ]
     dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres)
     ay = v + dKdy
-    by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon, particle])
+    by = np.sqrt(2 * fieldset.Kh_meridional[particles.time, particles.depth, particles.lat, particles.lon, particles])
 
     # Particle positions are updated only after evaluating all terms.
-    particle.dlon += ax * dt + bx * dWx
-    particle.dlat += ay * dt + by * dWy
+    particles.dlon += ax * dt + bx * dWx
+    particles.dlat += ay * dt + by * dWy
 
 
-def DiffusionUniformKh(particle, fieldset, time):  # pragma: no cover
+def DiffusionUniformKh(particles, fieldset):  # pragma: no cover
     """Kernel for simple 2D diffusion where diffusivity (Kh) is assumed uniform.
 
     Assumes that fieldset has constant fields `Kh_zonal` and `Kh_meridional`.
@@ -101,15 +109,13 @@ def DiffusionUniformKh(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
-    dt = particle.dt / np.timedelta64(1, "s")
+    dt = particles.dt / np.timedelta64(1, "s")
     # Wiener increment with zero mean and std of sqrt(dt)
     dWx = np.random.normal(0, np.sqrt(np.fabs(dt)))
     dWy = np.random.normal(0, np.sqrt(np.fabs(dt)))
 
-    print(particle)
-
-    bx = np.sqrt(2 * fieldset.Kh_zonal[particle])
-    by = np.sqrt(2 * fieldset.Kh_meridional[particle])
+    bx = np.sqrt(2 * fieldset.Kh_zonal[particles])
+    by = np.sqrt(2 * fieldset.Kh_meridional[particles])
 
-    particle.dlon += bx * dWx
-    particle.dlat += by * dWy
+    particles.dlon += bx * dWx
+    particles.dlat += by * dWy

parcels/application_kernels/advectiondiffusion.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import inspect
 import warnings
 from collections.abc import Callable
 from datetime import datetime
@@ -19,6 +18,7 @@
     ZeroInterpolator_Vector,
 )
 from parcels.particle import KernelParticle
+from parcels.tools._helpers import _assert_same_function_signature
 from parcels.tools.converters import (
     UnitConverter,
     unitconverters_map,
@@ -57,25 +57,6 @@ def _deal_with_errors(error, key, vector_type: VectorType):
 }
 
 
-def _assert_same_function_signature(f: Callable, *, ref: Callable) -> None:
-    """Ensures a function `f` has the same signature as the reference function `ref`."""
-    sig_ref = inspect.signature(ref)
-    sig = inspect.signature(f)
-
-    if len(sig_ref.parameters) != len(sig.parameters):
-        raise ValueError(
-            f"Interpolation function must have {len(sig_ref.parameters)} parameters, got {len(sig.parameters)}"
-        )
-
-    for (_name1, param1), (_name2, param2) in zip(sig_ref.parameters.items(), sig.parameters.items(), strict=False):
-        if param1.kind != param2.kind:
-            raise ValueError(
-                f"Parameter '{_name2}' has incorrect parameter kind. Expected {param1.kind}, got {param2.kind}"
-            )
-        if param1.name != param2.name:
-            raise ValueError(f"Parameter '{_name2}' has incorrect name. Expected '{param1.name}', got '{param2.name}'")
-
-
 class Field:
     """The Field class that holds scalar field data.
     The `Field` object is a wrapper around a xarray.DataArray or uxarray.UxDataArray object.
@@ -157,7 +138,7 @@ def __init__(
         if interp_method is None:
             self._interp_method = _DEFAULT_INTERPOLATOR_MAPPING[type(self.grid)]
         else:
-            _assert_same_function_signature(interp_method, ref=ZeroInterpolator)
+            _assert_same_function_signature(interp_method, ref=ZeroInterpolator, context="Interpolation")
             self._interp_method = interp_method
 
         self.igrid = -1  # Default the grid index to -1
@@ -213,7 +194,7 @@ def interp_method(self):
 
     @interp_method.setter
     def interp_method(self, method: Callable):
-        _assert_same_function_signature(method, ref=ZeroInterpolator)
+        _assert_same_function_signature(method, ref=ZeroInterpolator, context="Interpolation")
         self._interp_method = method
 
     def _check_velocitysampling(self):
@@ -287,7 +268,7 @@ def __init__(
         if vector_interp_method is None:
             self._vector_interp_method = None
         else:
-            _assert_same_function_signature(vector_interp_method, ref=ZeroInterpolator_Vector)
+            _assert_same_function_signature(vector_interp_method, ref=ZeroInterpolator_Vector, context="Interpolation")
             self._vector_interp_method = vector_interp_method
 
     def __repr__(self):
@@ -303,7 +284,7 @@ def vector_interp_method(self):
 
     @vector_interp_method.setter
     def vector_interp_method(self, method: Callable):
-        _assert_same_function_signature(method, ref=ZeroInterpolator_Vector)
+        _assert_same_function_signature(method, ref=ZeroInterpolator_Vector, context="Interpolation")
         self._vector_interp_method = method
 
     def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):

parcels/field.py

@@ -13,6 +13,7 @@
     AdvectionRK45,
 )
 from parcels.basegrid import GridType
+from parcels.tools._helpers import _assert_same_function_signature
 from parcels.tools.statuscodes import (
     StatusCode,
     _raise_field_interpolation_error,
@@ -23,6 +24,7 @@
     _raise_time_extrapolation_error,
 )
 from parcels.tools.warnings import KernelWarning
+from tests.common_kernels import DoNothing
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -67,6 +69,7 @@ def __init__(
         for f in pyfuncs:
             if not isinstance(f, types.FunctionType):
                 raise TypeError(f"Argument pyfunc should be a function or list of functions. Got {type(f)}")
+            _assert_same_function_signature(f, ref=DoNothing, context="Kernel")
 
         if len(pyfuncs) == 0:
             raise ValueError("List of `pyfuncs` should have at least one function.")
@@ -112,22 +115,22 @@ def remove_deleted(self, pset):
 
     def add_positionupdate_kernels(self):
         # Adding kernels that set and update the coordinate changes
-        def Setcoords(particle, fieldset, time):  # pragma: no cover
+        def Setcoords(particles, fieldset):  # pragma: no cover
             import numpy as np  # noqa
 
-            particle.dlon = 0
-            particle.dlat = 0
-            particle.ddepth = 0
-            particle.lon = particle.lon_nextloop
-            particle.lat = particle.lat_nextloop
-            particle.depth = particle.depth_nextloop
-            particle.time = particle.time_nextloop
+            particles.dlon = 0
+            particles.dlat = 0
+            particles.ddepth = 0
+            particles.lon = particles.lon_nextloop
+            particles.lat = particles.lat_nextloop
+            particles.depth = particles.depth_nextloop
+            particles.time = particles.time_nextloop
 
-        def Updatecoords(particle, fieldset, time):  # pragma: no cover
-            particle.lon_nextloop = particle.lon + particle.dlon
-            particle.lat_nextloop = particle.lat + particle.dlat
-            particle.depth_nextloop = particle.depth + particle.ddepth
-            particle.time_nextloop = particle.time + particle.dt
+        def Updatecoords(particles, fieldset):  # pragma: no cover
+            particles.lon_nextloop = particles.lon + particles.dlon
+            particles.lat_nextloop = particles.lat + particles.dlat
+            particles.depth_nextloop = particles.depth + particles.ddepth
+            particles.time_nextloop = particles.time + particles.dt
 
         self._pyfuncs = (Setcoords + self + Updatecoords)._pyfuncs
 
@@ -255,12 +258,12 @@ def execute(self, pset, endtime, dt):
             # run kernels for all particles that need to be evaluated
             evaluate_particles = (pset.state == StatusCode.Evaluate) & (pset.dt != 0)
             for f in self._pyfuncs:
-                f(pset[evaluate_particles], self._fieldset, None)
+                f(pset[evaluate_particles], self._fieldset)
 
                 # check for particles that have to be repeated
                 repeat_particles = pset.state == StatusCode.Repeat
                 while np.any(repeat_particles):
-                    f(pset[repeat_particles], self._fieldset, None)
+                    f(pset[repeat_particles], self._fieldset)
                     repeat_particles = pset.state == StatusCode.Repeat
 
             # revert to original dt (unless in RK45 mode)

parcels/kernel.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import functools
+import inspect
 import warnings
 from collections.abc import Callable
 from datetime import timedelta
@@ -75,3 +76,24 @@ def timedelta_to_float(dt: float | timedelta | np.timedelta64) -> float:
 def should_calculate_next_ti(ti: int, tau: float, tdim: int):
     """Check if the time is beyond the last time in the field"""
     return np.greater(tau, 0) and ti < tdim - 1
+
+
+def _assert_same_function_signature(f: Callable, *, ref: Callable, context: str) -> None:
+    """Ensures a function `f` has the same signature as the reference function `ref`."""
+    sig_ref = inspect.signature(ref)
+    sig = inspect.signature(f)
+
+    if len(sig_ref.parameters) != len(sig.parameters):
+        raise ValueError(
+            f"{context} function must have {len(sig_ref.parameters)} parameters, got {len(sig.parameters)}"
+        )
+
+    for param1, param2 in zip(sig_ref.parameters.values(), sig.parameters.values(), strict=False):
+        if param1.kind != param2.kind:
+            raise ValueError(
+                f"Parameter '{param2.name}' has incorrect parameter kind. Expected {param1.kind}, got {param2.kind}"
+            )
+        if param1.name != param2.name:
+            raise ValueError(
+                f"Parameter '{param2.name}' has incorrect name. Expected '{param1.name}', got '{param2.name}'"
+            )

parcels/tools/_helpers.py

@@ -1,18 +1,21 @@
 """Shared kernels between tests."""
 
+import numpy as np
 
-def DoNothing(particle, fieldset, time):  # pragma: no cover
+from parcels.tools.statuscodes import StatusCode
+
+
+def DoNothing(particles, fieldset):  # pragma: no cover
     pass
 
 
-def DeleteParticle(particle, fieldset, time):  # pragma: no cover
-    if particle.state >= 50:  # This captures all Errors
-        particle.delete()
+def DeleteParticle(particles, fieldset):  # pragma: no cover
+    particles.state = np.where(particles.state >= 50, StatusCode.Delete, particles.state)
 
 
-def MoveEast(particle, fieldset, time):  # pragma: no cover
-    particle.dlon += 0.1
+def MoveEast(particles, fieldset):  # pragma: no cover
+    particles.dlon += 0.1
 
 
-def MoveNorth(particle, fieldset, time):  # pragma: no cover
-    particle.dlat += 0.1
+def MoveNorth(particles, fieldset):  # pragma: no cover
+    particles.dlat += 0.1