Merge pull request #2096 from OceanParcels/kernelloop-speedups

Kernelloop speedups

Author: erikvansebille

parcels/kernel.py

@@ -75,7 +75,8 @@ def remove_deleted(self, pset):
         # TODO v4: need to implement ParticleFile writing of deleted particles
         # if len(indices) > 0 and self.fieldset.particlefile is not None:
         #     self.fieldset.particlefile.write(pset, None, indices=indices)
-        pset.remove_indices(indices)
+        if len(indices) > 0:
+            pset.remove_indices(indices)
 
 
 class Kernel(BaseKernel):
@@ -378,25 +379,24 @@ def evaluate_particle(self, p, endtime):
         dt :
             computational integration timestep
         """
+        sign_dt = 1 if p.dt >= 0 else -1
         while p.state in [StatusCode.Evaluate, StatusCode.Repeat]:
-            pre_dt = p.dt
-
-            sign_dt = np.sign(p.dt).astype(int)
-            if sign_dt * (endtime - p.time_nextloop) <= np.timedelta64(0, "ns"):
+            if sign_dt * (endtime - p.time_nextloop) <= 0:
                 return p
 
+            pre_dt = p.dt
             # TODO implement below later again
             # try:  # Use next_dt from AdvectionRK45 if it is set
             #     if abs(endtime - p.time_nextloop) < abs(p.next_dt) - 1e-6:
             #         p.next_dt = abs(endtime - p.time_nextloop) * sign_dt
             # except AttributeError:
-            if abs(endtime - p.time_nextloop) <= abs(p.dt):
-                p.dt = abs(endtime - p.time_nextloop) * sign_dt
+            if sign_dt * (endtime - p.time_nextloop) <= p.dt:
+                p.dt = sign_dt * (endtime - p.time_nextloop)
             res = self._pyfunc(p, self._fieldset, p.time_nextloop)
 
             if res is None:
                 if p.state == StatusCode.Success:
-                    if sign_dt * (p.time - endtime) > np.timedelta64(0, "ns"):
+                    if sign_dt * (p.time - endtime) > 0:
                         p.state = StatusCode.Evaluate
             else:
                 p.state = res

parcels/particleset.py

@@ -797,9 +797,9 @@ def execute(
         time = start_time
         while sign_dt * (time - end_time) < 0:
             if sign_dt > 0:
-                next_time = min(time + dt, end_time)
+                next_time = end_time  # TODO update to min(next_output, end_time) when ParticleFile works
             else:
-                next_time = max(time + dt, end_time)
+                next_time = end_time  # TODO update to max(next_output, end_time) when ParticleFile works
             res = self._kernel.execute(self, endtime=next_time, dt=dt)
             if res == StatusCode.StopAllExecution:
                 return StatusCode.StopAllExecution
@@ -813,7 +813,7 @@ def execute(
                         next_output += outputdt
 
             if verbose_progress:
-                pbar.update(dt / np.timedelta64(1, "s"))
+                pbar.update((next_time - time) / np.timedelta64(1, "s"))
 
             time = next_time
 

tests/v4/test_particleset_execute.py

@@ -113,7 +113,7 @@ def PythonFail(particle, fieldset, time):  # pragma: no cover
         pset.execute(PythonFail, runtime=np.timedelta64(20, "s"), dt=np.timedelta64(2, "s"))
     assert len(pset) == npart
     assert pset.time[0] == fieldset.time_interval.left + np.timedelta64(10, "s")
-    assert all([time == fieldset.time_interval.left + np.timedelta64(8, "s") for time in pset.time[1:]])
+    assert all([time == fieldset.time_interval.left + np.timedelta64(0, "s") for time in pset.time[1:]])
 
 
 @pytest.mark.parametrize("verbose_progress", [True, False])