Repository: OceanParcels/Parcels
Pull request #1937 (2025-03-12 15:45:27)

Moving the _search_time_index function to _index_search.py
This PR moves the `Field._search_time_index` method to its own function in `_index_search.py`, where I feel it would fit better 
<!-- Feel free to remove list items that are not relevant for your changes. -->

- [x] Chose the correct base branch (`main` for v3 changes, `v4-dev` for v4 changes)



Patch:
diff --git a/parcels/_index_search.py b/parcels/_index_search.py
index a52baf56..ddfbadd1 100644
--- a/parcels/_index_search.py
+++ b/parcels/_index_search.py
@@ -14,6 +14,7 @@ from parcels.tools.statuscodes import (
     _raise_field_out_of_bound_error,
     _raise_field_out_of_bound_surface_error,
     _raise_field_sampling_error,
+    _raise_time_extrapolation_error,
 )
 
 from .grid import GridType
@@ -23,6 +24,35 @@ if TYPE_CHECKING:
     from .grid import Grid
 
 
+def _search_time_index(grid: Grid, time: float, allow_time_extrapolation=True):
+    """Find and return the index and relative coordinate in the time array associated with a given time.
+
+    Note that we normalize to either the first or the last index
+    if the sampled value is outside the time value range.
+    """
+    if not allow_time_extrapolation and (time < grid.time[0] or time > grid.time[-1]):
+        _raise_time_extrapolation_error(time, field=None)
+    time_index = grid.time <= time
+
+    if time_index.all():
+        # If given time > last known field time, use
+        # the last field frame without interpolation
+        ti = len(grid.time) - 1
+    elif np.logical_not(time_index).all():
+        # If given time < any time in the field, use
+        # the first field frame without interpolation
+        ti = 0
+    else:
+        ti = int(time_index.argmin() - 1) if time_index.any() else 0
+    if grid.tdim == 1:
+        tau = 0
+    elif ti == len(grid.time) - 1:
+        tau = 1
+    else:
+        tau = (time - grid.time[ti]) / (grid.time[ti + 1] - grid.time[ti]) if grid.time[ti] != grid.time[ti + 1] else 0
+    return tau, ti
+
+
 def search_indices_vertical_z(grid: Grid, gridindexingtype: GridIndexingType, z: float):
     if grid.depth[-1] > grid.depth[0]:
         if z < grid.depth[0]:
diff --git a/parcels/field.py b/parcels/field.py
index f9d39500..e7f122ed 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -36,12 +36,11 @@ from parcels.tools.statuscodes import (
     FieldOutOfBoundError,
     FieldOutOfBoundSurfaceError,
     FieldSamplingError,
-    TimeExtrapolationError,
     _raise_field_out_of_bound_error,
 )
 from parcels.tools.warnings import FieldSetWarning
 
-from ._index_search import _search_indices_curvilinear, _search_indices_rectilinear
+from ._index_search import _search_indices_curvilinear, _search_indices_rectilinear, _search_time_index
 from .fieldfilebuffer import (
     NetcdfFileBuffer,
 )
@@ -543,7 +542,7 @@ class Field:
         return data
 
     def _search_indices(self, time, z, y, x, particle=None, search2D=False):
-        tau, ti = self._search_time_index(time)
+        tau, ti = _search_time_index(self.grid, time, self.allow_time_extrapolation)
 
         if self.grid._gtype in [GridType.RectilinearSGrid, GridType.RectilinearZGrid]:
             (zeta, eta, xsi, zi, yi, xi) = _search_indices_rectilinear(
@@ -602,38 +601,6 @@ class Field:
             e = add_note(e, f"Error interpolating field '{self.name}'.", before=True)
             raise e
 
-    def _search_time_index(self, time):
-        """Find and return the index and relative coordinate in the time array associated with a given time.
-
-        Note that we normalize to either the first or the last index
-        if the sampled value is outside the time value range.
-        """
-        if not self.allow_time_extrapolation and (time < self.grid.time[0] or time > self.grid.time[-1]):
-            raise TimeExtrapolationError(time, field=self)
-        time_index = self.grid.time <= time
-
-        if time_index.all():
-            # If given time > last known field time, use
-            # the last field frame without interpolation
-            ti = len(self.grid.time) - 1
-        elif np.logical_not(time_index).all():
-            # If given time < any time in the field, use
-            # the first field frame without interpolation
-            ti = 0
-        else:
-            ti = time_index.argmin() - 1 if time_index.any() else 0
-        if self.grid.tdim == 1:
-            tau = 0
-        elif ti == len(self.grid.time) - 1:
-            tau = 1
-        else:
-            tau = (
-                (time - self.grid.time[ti]) / (self.grid.time[ti + 1] - self.grid.time[ti])
-                if self.grid.time[ti] != self.grid.time[ti + 1]
-                else 0
-            )
-        return tau, ti
-
     def _check_velocitysampling(self):
         if self.name in ["U", "V", "W"]:
             warnings.warn(
diff --git a/parcels/grid.py b/parcels/grid.py
index ea62f153..9332dcb9 100644
--- a/parcels/grid.py
+++ b/parcels/grid.py
@@ -58,6 +58,7 @@ class Grid:
         self._lon = lon
         self._lat = lat
         self.time = time
+        self.tdim = time.size
         self._time_origin = TimeConverter() if time_origin is None else time_origin
         assert isinstance(self.time_origin, TimeConverter), "time_origin needs to be a TimeConverter object"
         assert_valid_mesh(mesh)
@@ -182,7 +183,6 @@ class RectilinearGrid(Grid):
             assert len(time.shape) == 1, "time is not a vector"
 
         super().__init__(lon, lat, time, time_origin, mesh)
-        self.tdim = self.time.size
 
     @property
     def xdim(self):
@@ -326,7 +326,6 @@ class CurvilinearGrid(Grid):
         lon = lon.squeeze()
         lat = lat.squeeze()
         super().__init__(lon, lat, time, time_origin, mesh)
-        self.tdim = self.time.size
 
     @property
     def xdim(self):
diff --git a/parcels/tools/statuscodes.py b/parcels/tools/statuscodes.py
index 7a9799dc..501741ec 100644
--- a/parcels/tools/statuscodes.py
+++ b/parcels/tools/statuscodes.py
@@ -10,6 +10,7 @@ __all__ = [
     "_raise_field_out_of_bound_error",
     "_raise_field_out_of_bound_surface_error",
     "_raise_field_sampling_error",
+    "_raise_time_extrapolation_error",
 ]
 
 
@@ -77,6 +78,10 @@ class TimeExtrapolationError(RuntimeError):
         super().__init__(message)
 
 
+def _raise_time_extrapolation_error(time: float, field=None):
+    raise TimeExtrapolationError(time, field)
+
+
 class KernelError(RuntimeError):
     """General particle kernel error with optional custom message."""
 


Test patch:
diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py
index 26618605..78b7317a 100644
--- a/tests/test_fieldset.py
+++ b/tests/test_fieldset.py
@@ -450,9 +450,11 @@ def test_fieldset_write(tmp_zarrfile):
     fieldset.U.to_write = True
 
     def UpdateU(particle, fieldset, time):  # pragma: no cover
+        from parcels._index_search import _search_time_index
+
         tmp1, tmp2 = fieldset.UV[particle]
         _, yi, xi = fieldset.U.unravel_index(particle.ei)
-        _, ti = fieldset.U._search_time_index(time)
+        _, ti = _search_time_index(fieldset.U.grid, time)
         fieldset.U.data[ti, yi, xi] += 1
         fieldset.U.grid.time[0] = time