Adding an iterator for xarray particleset

Author: erikvansebille

parcels/kernel.py

@@ -316,8 +316,7 @@ def execute(self, pset, endtime, dt):
             self.add_positionupdate_kernels()
             self._positionupdate_kernels_added = True
 
-        for i in pset.trajectory.values:
-            p = pset[i]
+        for p in pset:
             self.evaluate_particle(p, endtime)
             if p.state == StatusCode.StopAllExecution:
                 return StatusCode.StopAllExecution

parcels/particleset.py

@@ -174,7 +174,15 @@ def __del__(self):
         self._data = None
 
     def __iter__(self):
-        return iter(self._data)  # TODO write an iter that iterates over particles (instead of variables)
+        self._index = 0
+        return self
+
+    def __next__(self):
+        if self._index < len(self):
+            p = self._data.sel(trajectory=self._index)
+            self._index += 1
+            return p
+        raise StopIteration
 
     def __getattr__(self, name):
         """

tests/v4/test_particleset.py

@@ -116,6 +116,14 @@ def test_pset_add_implicit(fieldset):
     assert np.allclose(np.diff(pset._data.trajectory), np.ones(6), atol=1e-12)
 
 
+def test_pset_iterator(fieldset):
+    npart = 10
+    pset = ParticleSet(fieldset, lon=np.zeros(npart), lat=np.ones(npart))
+    for i, particle in enumerate(pset):
+        assert particle.trajectory == i
+    assert i == npart - 1
+
+
 @pytest.mark.parametrize("verbose_progress", [True, False])
 def test_uxstommelgyre_pset_execute(verbose_progress):
     ds = datasets_unstructured["stommel_gyre_delaunay"]