Merge remote-tracking branch 'origin/benchmarking-running-bugs' into particle_time_as_float

Author: erikvansebille

parcels/_index_search.py

@@ -39,6 +39,8 @@ def _search_time_index(field: Field, time: datetime):
     """
     ti = np.argmin(field._time_float <= time) - 1
     tau = (time - field._time_float[ti]) / (field._time_float[ti + 1] - field._time_float[ti])
+    if tau < 0 or tau > 1:  # TODO only for debugging; test can go?
+        raise ValueError(f"Time {time} is out of bounds for field time data {field.data.time.data}.")
     return tau, ti
 
 

parcels/field.py

@@ -246,6 +246,19 @@ def _check_velocitysampling(self):
                 stacklevel=2,
             )
 
+    def _load_timesteps(self, time):
+        """Load the appropriate timesteps of a field."""
+        ti = np.argmin(self._time_float <= time) - 1  # TODO also implement dt < 0
+        if not hasattr(self, "data"):
+            self.data = self.data_full.isel({"time": slice(ti, ti + 2)}).load()
+        elif self.data_full.time.data[ti] == self.data.time.data[1]:
+            self.data = xr.concat([self.data[1, :], self.data_full.isel({"time": ti + 1}).load()], dim="time")
+        elif self.data_full.time.data[ti] != self.data.time.data[0]:
+            self.data = self.data_full.isel({"time": slice(ti, ti + 2)}).load()
+        assert (
+            len(self.data.time) == 2
+        ), f"Field {self.name} has not been loaded correctly. Expected 2 timesteps, but got {len(self.data.time)}."
+
     def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         """Interpolate field values in space and time.
 

parcels/fieldset.py

@@ -82,6 +82,13 @@ def time_interval(self):
             return None
         return functools.reduce(lambda x, y: x.intersection(y), time_intervals)
 
+    def _load_timesteps(self, time):
+        """Load the appropriate timesteps of all fields in the fieldset."""
+        for fldname in self.fields:
+            field = self.fields[fldname]
+            if isinstance(field, Field):
+                field._load_timesteps(time)
+
     def add_field(self, field: Field, name: str | None = None):
         """Add a :class:`parcels.field.Field` object to the FieldSet.
 

parcels/particleset.py

@@ -12,6 +12,7 @@
 from parcels._reprs import particleset_repr
 from parcels.application_kernels.advection import AdvectionRK4
 from parcels.basegrid import GridType
+from parcels.field import Field
 from parcels.interaction.interactionkernel import InteractionKernel
 from parcels.kernel import Kernel
 from parcels.particle import Particle, Variable
@@ -813,14 +814,14 @@ def execute(
 
         time = start_time
 
-        for fld in [self.fieldset.U, self.fieldset.V]:  # TODO generalise to all fields and move to better place
-            fld._time_float = (fld.data_full.time.data - fld.time_interval.left) / np.timedelta64(1, "s")
+        for fldname in self.fieldset.fields:
+            field = self.fieldset.fields[fldname]
+            if isinstance(field, Field):
+                field._time_float = (field.data_full.time.data - field.time_interval.left) / np.timedelta64(1, "s")
 
         while sign_dt * (time - end_time) < 0:
-            for fld in [self.fieldset.U, self.fieldset.V]:  # TODO generalise to all fields
-                ti = np.argmin(fld._time_float <= self._data["time_nextloop"][0]) - 1  # TODO also implement dt < 0
-                if not hasattr(fld, "data") or fld.data_full.time.data[ti] != fld.data.time.data[0]:
-                    fld.data = fld.data_full.isel({"time": slice(ti, ti + 2)}).load()
+            # Load the appropriate timesteps of the fieldset
+            self.fieldset._load_timesteps(self._data["time_nextloop"][0])
 
             if sign_dt > 0:
                 next_time = min(time + dt, end_time)

parcels/xgrid.py

@@ -457,6 +457,8 @@ def _search_1d_array(
         return 0, 0.0
     i = np.argmin(arr <= x) - 1
     bcoord = (x - arr[i]) / (arr[i + 1] - arr[i])
+    if bcoord < 0 or bcoord > 1:  # TODO only for debugging; test can go?
+        raise ValueError(f"Position {x} is out of bounds for array {arr}.")
     return i, bcoord