diff --git a/parcels/field.py b/parcels/field.py index 473ecbc833..3268f11c4e 100644 --- a/parcels/field.py +++ b/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,6 +323,16 @@ 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 @@ -339,37 +340,6 @@ def __contains__(self, key: str): 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 diff --git a/tests/v4/test_field.py b/tests/v4/test_field.py index b336129dfe..a555a672f2 100644 --- a/tests/v4/test_field.py +++ b/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