Adding support for OperationCode.StopAllExecution

docs/examples/tutorial_kernelloop.ipynb

@@ -326,6 +326,7 @@
       "Repeat = 2\n",
       "Delete = 3\n",
       "StopExecution = 4\n",
+      "StopAllExecution = 41\n",
       "Error = 5\n",
       "ErrorInterpolation = 51\n",
       "ErrorOutOfBounds = 6\n",

parcels/compilation/codegenerator.py

@@ -932,6 +932,7 @@ def generate(self, funcname, field_args, const_args, kernel_ast, c_include):
                                     c.Block([c.Assign("particles->state[pnum]", "EVALUATE")]),
                                     c.Block([c.Assign("particles->state[pnum]", "SUCCESS")]))
                                ]))]
+        body += [c.If("particles->state[pnum] == STOPALLEXECUTION", c.Statement("return"))]
         body += [c.Statement("particles->dt[pnum] = pre_dt")]
         body += [c.If("(particles->state[pnum] == REPEAT || particles->state[pnum] == DELETE)", c.Block([c.Statement('break')]))]
 

parcels/include/index_search.h

@@ -51,7 +51,7 @@ typedef struct
 
 typedef enum
   {
-    SUCCESS=0, EVALUATE=1, REPEAT=2, DELETE=3, STOPEXECUTION=4, ERROR=5, ERRORINTERPOLATION=51, ERROROUTOFBOUNDS=6, ERRORTHROUGHSURFACE=61, ERRORTIMEEXTRAPOLATION=7
+    SUCCESS=0, EVALUATE=1, REPEAT=2, DELETE=3, STOPEXECUTION=4, STOPALLEXECUTION=41, ERROR=5, ERRORINTERPOLATION=51, ERROROUTOFBOUNDS=6, ERRORTHROUGHSURFACE=61, ERRORTIMEEXTRAPOLATION=7
   } StatusCode;
 
 typedef enum

parcels/kernel.py

@@ -544,6 +544,8 @@ def execute_python(self, pset, endtime, dt):
 
         for p in pset:
             self.evaluate_particle(p, endtime, sign_dt)
+            if p.state == StatusCode.StopAllExecution:
+                return StatusCode.StopAllExecution
 
     def execute(self, pset, endtime, dt):
         """Execute this Kernel over a ParticleSet for several timesteps."""
@@ -576,6 +578,8 @@ def execute(self, pset, endtime, dt):
             for p in error_pset:
                 if p.state == StatusCode.StopExecution:
                     return
+                if p.state == StatusCode.StopAllExecution:
+                    return StatusCode.StopAllExecution
                 if p.state == StatusCode.Repeat:
                     p.state = StatusCode.Evaluate
                 elif p.state == StatusCode.ErrorTimeExtrapolation:

parcels/particleset.py

@@ -957,7 +957,9 @@ def execute(self, pyfunc=AdvectionRK4, pyfunc_inter=None, endtime=None, runtime=
 
             # If we don't perform interaction, only execute the normal kernel efficiently.
             if self.interaction_kernel is None:
-                self.kernel.execute(self, endtime=next_time, dt=dt)
+                res = self.kernel.execute(self, endtime=next_time, dt=dt)
+                if res == StatusCode.StopAllExecution:
+                    return StatusCode.StopAllExecution
             # Interaction: interleave the interaction and non-interaction kernel for each time step.
             # E.g. Normal -> Inter -> Normal -> Inter if endtime-time == 2*dt
             else:

parcels/tools/statuscodes.py

@@ -13,6 +13,7 @@ class StatusCode:
     Repeat = 2
     Delete = 3
     StopExecution = 4
+    StopAllExecution = 41
     Error = 5
     ErrorInterpolation = 51
     ErrorOutOfBounds = 6

tests/test_kernel_execution.py

@@ -185,6 +185,23 @@ def RecoverAllErrors(particle, fieldset, time):
     assert len(pset) == 0
 
 
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+def test_execution_check_stopallexecution(fieldset, mode):
+
+    def addoneLon(particle, fieldset, time):
+        particle_dlon += 1  # noqa
+
+        if particle.lon + particle_dlon >= 10:
+            particle.state = StatusCode.StopAllExecution
+
+    pset = ParticleSet(fieldset, pclass=ptype[mode], lon=[0, 1], lat=[0, 0])
+    pset.execute(addoneLon, endtime=20., dt=1.)
+    assert pset[0].lon == 9
+    assert pset[0].time == 9
+    assert pset[1].lon == 1
+    assert pset[1].time == 0
+
+
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 def test_execution_delete_out_of_bounds(fieldset, mode, npart=10):
     def MoveRight(particle, fieldset, time):