Support adding particlesets

Author: erikvansebille

parcels/particleset.py

@@ -233,8 +233,10 @@ def add(self, particles):
 
         """
         if isinstance(particles, type(self)):
-            particles = particles.particledata
-        self.particledata += particles
+            particles.data["trajectory"] = (
+                particles.data["trajectory"].values + self.data["trajectory"].values.max() + 1
+            )
+        self.data = xr.concat([self.data, particles.data], dim="trajectory")
         # Adding particles invalidates the neighbor search structure.
         self._dirty_neighbor = True
         return self
@@ -263,8 +265,8 @@ def remove_indices(self, indices):
         self.data = self.data.drop_sel(trajectory=indices)
 
     def _active_particles_mask(self, time, dt):
-        active_indices = (time - self.particledata.data["time"]) / dt >= 0
-        non_err_indices = np.isin(self.particledata.data["state"], [StatusCode.Success, StatusCode.Evaluate])
+        active_indices = (time - self.data["time"]) / dt >= 0
+        non_err_indices = np.isin(self.data["state"], [StatusCode.Success, StatusCode.Evaluate])
         active_indices = np.logical_and(active_indices, non_err_indices)
         self._active_particle_idx = np.where(active_indices)[0]
         return active_indices
@@ -274,9 +276,9 @@ def _compute_neighbor_tree(self, time, dt):
 
         self._values = np.vstack(
             (
-                self.particledata.data["depth"],
-                self.particledata.data["lat"],
-                self.particledata.data["lon"],
+                self.data["depth"],
+                self.data["lat"],
+                self.data["lon"],
             )
         )
         if self._dirty_neighbor:
@@ -290,14 +292,14 @@ def _neighbors_by_index(self, particle_idx):
         neighbor_idx = self._active_particle_idx[neighbor_idx]
         mask = neighbor_idx != particle_idx
         neighbor_idx = neighbor_idx[mask]
-        if "horiz_dist" in self.particledata._ptype.variables:
-            self.particledata.data["vert_dist"][neighbor_idx] = distances[0, mask]
-            self.particledata.data["horiz_dist"][neighbor_idx] = distances[1, mask]
+        if "horiz_dist" in self.data._ptype.variables:
+            self.data["vert_dist"][neighbor_idx] = distances[0, mask]
+            self.data["horiz_dist"][neighbor_idx] = distances[1, mask]
         return True  # TODO fix for v4 ParticleDataIterator(self.particledata, subset=neighbor_idx)
 
     def _neighbors_by_coor(self, coor):
         neighbor_idx = self._neighbor_tree.find_neighbors_by_coor(coor)
-        neighbor_ids = self.particledata.data["id"][neighbor_idx]
+        neighbor_ids = self.data["id"][neighbor_idx]
         return neighbor_ids
 
     # TODO: This method is only tested in tutorial notebook. Add unit test?
@@ -725,7 +727,7 @@ def _num_error_particles(self):
         int
             Number of error particles.
         """
-        return np.sum(np.isin(self.particledata.data["state"], [StatusCode.Success, StatusCode.Evaluate], invert=True))
+        return np.sum(np.isin(self.data["state"], [StatusCode.Success, StatusCode.Evaluate], invert=True))
 
     def set_variable_write_status(self, var, write_status):
         """Method to set the write status of a Variable.

tests/v4/test_particleset.py

@@ -95,6 +95,27 @@ def test_particleset_endtime_type(fieldset, endtime, expectation):
         pset.execute(endtime=endtime, dt=np.timedelta64(10, "m"), pyfunc=DoNothing)
 
 
+def test_pset_add_explicit(fieldset):
+    npart = 11
+    lon = np.linspace(0, 1, npart)
+    lat = np.linspace(1, 0, npart)
+    pset = ParticleSet(fieldset, lon=lon[0], lat=lat[0], pclass=Particle)
+    for i in range(1, npart):
+        particle = ParticleSet(pclass=Particle, lon=lon[i], lat=lat[i], fieldset=fieldset)
+        pset.add(particle)
+    assert len(pset) == npart
+    assert np.allclose(pset.data["lon"][:, 0], lon, atol=1e-12)
+    assert np.allclose(pset.data["lat"][:, 0], lat, atol=1e-12)
+    assert np.allclose(np.diff(pset.data.trajectory), np.ones(pset.data.trajectory.size - 1), atol=1e-12)
+
+
+def test_pset_add_implicit(fieldset):
+    pset = ParticleSet(fieldset, lon=np.zeros(3), lat=np.ones(3), pclass=Particle)
+    pset += ParticleSet(fieldset, lon=np.ones(4), lat=np.zeros(4), pclass=Particle)
+    assert len(pset) == 7
+    assert np.allclose(np.diff(pset.data.trajectory), np.ones(6), atol=1e-12)
+
+
 @pytest.mark.parametrize("verbose_progress", [True, False])
 def test_uxstommelgyre_pset_execute(verbose_progress):
     ds = datasets_unstructured["stommel_gyre_delaunay"]