Remove interpolation function annotation checking and refactor (#2077)

* Move code

* Update vector interp template

* Update interp function validation to asserts

* Test interpolation function checking (and remove checking of type annotations)

Dynamically checking of type annotations like this is flaky due to how Python internally handles type annotations

* Refactor interpolation function checking

* fix test

* Add default interpolator mapping

Author: VeckoTheGecko

parcels/field.py

@@ -59,6 +59,45 @@ def _deal_with_errors(error, key, vector_type: VectorType):
         return 0
 
 
+def ZeroInterpolator(
+    field: Field,
+    ti: int,
+    position: dict[str, tuple[int, float | np.ndarray]],
+    tau: np.float32 | np.float64,
+    t: np.float32 | np.float64,
+    z: np.float32 | np.float64,
+    y: np.float32 | np.float64,
+    x: np.float32 | np.float64,
+) -> np.float32 | np.float64:
+    """Template function used for the signature check of the lateral interpolation methods."""
+    return 0.0
+
+
+_DEFAULT_INTERPOLATOR_MAPPING = {
+    XGrid: ZeroInterpolator,  # TODO v4: Update these to better defaults
+    UxGrid: ZeroInterpolator,
+}
+
+
+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.
@@ -93,44 +132,6 @@ class Field:
 
     """
 
-    @staticmethod
-    def _interp_template(
-        self,
-        ti: int,
-        position: dict[str, tuple[int, float | np.ndarray]],
-        tau: np.float32 | np.float64,
-        t: np.float32 | np.float64,
-        z: np.float32 | np.float64,
-        y: np.float32 | np.float64,
-        x: np.float32 | np.float64,
-    ) -> np.float32 | np.float64:
-        """Template function used for the signature check of the lateral interpolation methods."""
-        return 0.0
-
-    def _validate_interp_function(self, func: Callable) -> bool:
-        """Ensures that the function has the correct signature."""
-        template_sig = inspect.signature(self._interp_template)
-        func_sig = inspect.signature(func)
-
-        if len(template_sig.parameters) != len(func_sig.parameters):
-            return False
-
-        for (_name1, param1), (_name2, param2) in zip(
-            template_sig.parameters.items(), func_sig.parameters.items(), strict=False
-        ):
-            if param1.kind != param2.kind:
-                return False
-            if param1.annotation != param2.annotation:
-                return False
-
-        return_annotation = func_sig.return_annotation
-        template_return = template_sig.return_annotation
-
-        if return_annotation != template_return:
-            return False
-
-        return True
-
     def __init__(
         self,
         name: str,
@@ -178,9 +179,9 @@ def __init__(
 
         # Setting the interpolation method dynamically
         if interp_method is None:
-            self._interp_method = self._interp_template  # Default to method that returns 0 always
+            self._interp_method = _DEFAULT_INTERPOLATOR_MAPPING[type(self.grid)]
         else:
-            self._validate_interp_function(interp_method)
+            _assert_same_function_signature(interp_method, ref=ZeroInterpolator)
             self._interp_method = interp_method
 
         self.igrid = -1  # Default the grid index to -1
@@ -272,7 +273,7 @@ def interp_method(self):
 
     @interp_method.setter
     def interp_method(self, method: Callable):
-        self._validate_interp_function(method)
+        _assert_same_function_signature(method, ref=ZeroInterpolator)
         self._interp_method = method
 
     def _check_velocitysampling(self):
@@ -283,16 +284,6 @@ def _check_velocitysampling(self):
                 stacklevel=2,
             )
 
-    def __getitem__(self, key):
-        self._check_velocitysampling()
-        try:
-            if _isParticle(key):
-                return self.eval(key.time, key.depth, key.lat, key.lon, key)
-            else:
-                return self.eval(*key)
-        except tuple(AllParcelsErrorCodes.keys()) as error:
-            return _deal_with_errors(error, key, vector_type=None)
-
     def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         """Interpolate field values in space and time.
 
@@ -332,44 +323,23 @@ def _rescale_and_set_minmax(self, data):
     def __getattr__(self, key: str):
         return getattr(self.data, key)
 
+    def __getitem__(self, key):
+        self._check_velocitysampling()
+        try:
+            if _isParticle(key):
+                return self.eval(key.time, key.depth, key.lat, key.lon, key)
+            else:
+                return self.eval(*key)
+        except tuple(AllParcelsErrorCodes.keys()) as error:
+            return _deal_with_errors(error, key, vector_type=None)
+
     def __contains__(self, key: str):
         return key in self.data
 
 
 class VectorField:
     """VectorField class that holds vector field data needed to execute particles."""
 
-    @staticmethod
-    def _vector_interp_template(
-        self,
-        ti: int,
-        ei: int,
-        bcoords: np.ndarray,
-        t: np.float32 | np.float64,
-        z: np.float32 | np.float64,
-        y: np.float32 | np.float64,
-        x: np.float32 | np.float64,
-    ) -> np.float32 | np.float64:
-        """Template function used for the signature check of the lateral interpolation methods."""
-        return 0.0
-
-    def _validate_vector_interp_function(self, func: Callable):
-        """Ensures that the function has the correct signature."""
-        expected_params = ["ti", "ei", "bcoords", "t", "z", "y", "x"]
-        expected_return_types = (np.float32, np.float64)
-
-        sig = inspect.signature(func)
-        params = list(sig.parameters.keys())
-
-        # Check the parameter names and count
-        if params != expected_params:
-            raise TypeError(f"Function must have parameters {expected_params}, but got {params}")
-
-        # Check return annotation if present
-        return_annotation = sig.return_annotation
-        if return_annotation not in (inspect.Signature.empty, *expected_return_types):
-            raise TypeError(f"Function must return a float, but got {return_annotation}")
-
     def __init__(
         self, name: str, U: Field, V: Field, W: Field | None = None, vector_interp_method: Callable | None = None
     ):
@@ -395,7 +365,7 @@ def __init__(
         if vector_interp_method is None:
             self._vector_interp_method = None
         else:
-            self._validate_vector_interp_function(vector_interp_method)
+            _assert_same_function_signature(vector_interp_method, ref=ZeroInterpolator)
             self._interp_method = vector_interp_method
 
     def __repr__(self):
@@ -411,7 +381,7 @@ def vector_interp_method(self):
 
     @vector_interp_method.setter
     def vector_interp_method(self, method: Callable):
-        self._validate_vector_interp_function(method)
+        _assert_same_function_signature(method, ref=ZeroInterpolator)
         self._vector_interp_method = method
 
     # @staticmethod

tests/v4/test_field.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import numpy as np
 import pytest
 import uxarray as ux
 import xarray as xr
 
-from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, xgcm
+from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, VectorField, xgcm
 from parcels._datasets.structured.generic import T as T_structured
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
@@ -115,6 +117,34 @@ def test_vectorfield_init_different_time_intervals():
     ...
 
 
+def test_field_invalid_interpolator():
+    ds = datasets_structured["ds_2d_left"]
+    grid = XGrid(xgcm.Grid(ds))
+
+    def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid):
+        return 0.0
+
+    # Test invalid interpolator with wrong signature
+    with pytest.raises(ValueError, match=".*incorrect name.*"):
+        Field(name="test", data=ds["data_g"], grid=grid, interp_method=invalid_interpolator_wrong_signature)
+
+
+def test_vectorfield_invalid_interpolator():
+    ds = datasets_structured["ds_2d_left"]
+    grid = XGrid(xgcm.Grid(ds))
+
+    def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid):
+        return 0.0
+
+    # Create component fields
+    U = Field(name="U", data=ds["data_g"], grid=grid)
+    V = Field(name="V", data=ds["data_g"], grid=grid)
+
+    # Test invalid interpolator with wrong signature
+    with pytest.raises(ValueError, match=".*incorrect name.*"):
+        VectorField(name="UV", U=U, V=V, vector_interp_method=invalid_interpolator_wrong_signature)
+
+
 def test_field_unstructured_z_linear():
     """Tests correctness of piecewise constant and piecewise linear interpolation methods on an unstructured grid with a vertical coordinate.
     The example dataset is a FESOM2 square Delaunay grid with uniform z-coordinate. Cell centered and layer registered data are defined to be