cleaning up _load_timesteps

Author: erikvansebille

parcels/field.py

@@ -242,6 +242,19 @@ def _check_velocitysampling(self):
                 stacklevel=2,
             )
 
+    def _load_timesteps(self, time):
+        """Load the appropriate timesteps of a field."""
+        ti = np.argmin(self.data_full.time.data <= 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

@@ -805,11 +805,8 @@ def execute(
 
         time = start_time
         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.data_full.time.data <= 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)