implementing particleset adding and deleting methods for dicts

And fixing some unit tests

Author: erikvansebille

parcels/particleset.py

@@ -149,12 +149,12 @@ def __init__(
             "time_nextloop": time,
             "trajectory": trajectory_ids,
         }
-        self.ptype = pclass.getPType()
+        self._ptype = pclass.getPType()
         # add extra fields from the custom Particle class
         for v in pclass.__dict__.values():
             if isinstance(v, Variable):
                 if isinstance(v.initial, attrgetter):
-                    initial = v.initial(self).values
+                    initial = v.initial(self)
                 else:
                     initial = v.initial * np.ones(len(trajectory_ids), dtype=v.dtype)
                 self._data[v.name] = initial
@@ -231,13 +231,28 @@ def add(self, particles):
             The current ParticleSet
 
         """
+        assert (
+            particles is not None
+        ), f"Trying to add another {type(self)} to this one, but the other one is None - invalid operation."
+        assert type(particles) is type(self)
+
+        if len(particles) == 0:
+            return
+
+        if len(self) == 0:
+            self._data = particles._data
+            return
+
         if isinstance(particles, type(self)):
             if len(self._data["trajectory"]) > 0:
-                offset = self._data["trajectory"].values.max() + 1
+                offset = self._data["trajectory"].max() + 1
             else:
                 offset = 0
-            particles._data["trajectory"] = particles._data["trajectory"].values + offset
-        self._data = xr.concat([self._data, particles._data], dim="trajectory")
+            particles._data["trajectory"] = particles._data["trajectory"] + offset
+
+        for d in self._data:
+            self._data[d] = np.concatenate((self._data[d], particles._data[d]))
+
         # Adding particles invalidates the neighbor search structure.
         self._dirty_neighbor = True
         return self
@@ -263,7 +278,8 @@ def __iadd__(self, particles):
 
     def remove_indices(self, indices):
         """Method to remove particles from the ParticleSet, based on their `indices`."""
-        self._data = self._data.drop_sel(trajectory=indices)
+        for d in self._data:
+            self._data[d] = np.delete(self._data[d], indices, axis=0)
 
     def _active_particles_mask(self, time, dt):
         active_indices = (time - self._data["time"]) / dt >= 0
@@ -584,19 +600,19 @@ def Kernel(self, pyfunc):
         if isinstance(pyfunc, list):
             return Kernel.from_list(
                 self.fieldset,
-                self.ptype,
+                self._ptype,
                 pyfunc,
             )
         return Kernel(
             self.fieldset,
-            self.ptype,
+            self._ptype,
             pyfunc=pyfunc,
         )
 
     def InteractionKernel(self, pyfunc_inter):
         if pyfunc_inter is None:
             return None
-        return InteractionKernel(self.fieldset, self.ptype, pyfunc=pyfunc_inter)
+        return InteractionKernel(self.fieldset, self._ptype, pyfunc=pyfunc_inter)
 
     def ParticleFile(self, *args, **kwargs):
         """Wrapper method to initialise a :class:`parcels.particlefile.ParticleFile` object from the ParticleSet."""
@@ -740,9 +756,9 @@ def execute(
         else:
             if not np.isnat(self._data["time_nextloop"]).any():
                 if sign_dt > 0:
-                    start_time = self._data["time_nextloop"].min().values
+                    start_time = self._data["time_nextloop"].min()
                 else:
-                    start_time = self._data["time_nextloop"].max().values
+                    start_time = self._data["time_nextloop"].max()
             else:
                 if sign_dt > 0:
                     start_time = self.fieldset.time_interval.left

tests/v4/test_kernel.py

@@ -43,7 +43,7 @@ def test_unknown_var_in_kernel(fieldset):
     def ErrorKernel(particle, fieldset, time):  # pragma: no cover
         particle.unknown_varname += 0.2
 
-    with pytest.raises(KeyError, match="No variable named 'unknown_varname'"):
+    with pytest.raises(KeyError, match="'unknown_varname'"):
         pset.execute(ErrorKernel, runtime=np.timedelta64(2, "s"))
 
 

tests/v4/test_particleset.py

@@ -198,17 +198,17 @@ def test_pset_add_explicit(fieldset):
     assert len(pset) == npart
     assert np.allclose([p.lon for p in pset], lon, atol=1e-12)
     assert np.allclose([p.lat for p in pset], lat, atol=1e-12)
-    assert np.allclose(np.diff(pset._data.trajectory), np.ones(pset._data.trajectory.size - 1), 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)
+    assert np.allclose(np.diff(pset._data["trajectory"]), np.ones(6), atol=1e-12)
 
 
-def test_pset_add_implicit(fieldset, npart=10):
+def test_pset_add_implicit_in_loop(fieldset, npart=10):
     pset = ParticleSet(fieldset, lon=[], lat=[])
     for _ in range(npart):
         pset += ParticleSet(pclass=Particle, lon=0.1, lat=0.1, fieldset=fieldset)

tests/v4/test_particleset_execute.py

@@ -112,8 +112,7 @@ def PythonFail(particle, fieldset, time):  # pragma: no cover
     with pytest.raises(RuntimeError):
         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(10, "s") for time in pset.time])
 
 
 @pytest.mark.parametrize("verbose_progress", [True, False])