From d64d8961ebe3b69836bc1768651ae877e0cc86e7 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 18 Jul 2025 09:04:54 +0200 Subject: [PATCH 1/3] Using a dictionary of arrays for particleset data --- parcels/particle.py | 2 +- parcels/particleset.py | 45 ++++++++++++++++++------------------------ 2 files changed, 20 insertions(+), 27 deletions(-) diff --git a/parcels/particle.py b/parcels/particle.py index 93ea35254..912d42bc4 100644 --- a/parcels/particle.py +++ b/parcels/particle.py @@ -115,7 +115,7 @@ def __init__(self, data: xr.Dataset, index: int): self._index = index def __getattr__(self, name): - return self._data[name].values[self._index] + return self._data[name][self._index] def __setattr__(self, name, value): if name in ["_data", "_index"]: diff --git a/parcels/particleset.py b/parcels/particleset.py index 993741d42..93de67e83 100644 --- a/parcels/particleset.py +++ b/parcels/particleset.py @@ -135,28 +135,21 @@ def __init__( lon.size == kwargs[kwvar].size ), f"{kwvar} and positions (lon, lat, depth) don't have the same lengths." - self._data = xr.Dataset( - { - "lon": (["trajectory"], lon.astype(lonlatdepth_dtype)), - "lat": (["trajectory"], lat.astype(lonlatdepth_dtype)), - "depth": (["trajectory"], depth.astype(lonlatdepth_dtype)), - "time": (["trajectory"], time), - "dt": (["trajectory"], np.timedelta64(1, "ns") * np.ones(len(trajectory_ids))), - "ei": (["trajectory", "ngrid"], np.zeros((len(trajectory_ids), len(fieldset.gridset)), dtype=np.int32)), - "state": (["trajectory"], np.zeros((len(trajectory_ids)), dtype=np.int32)), - "lon_nextloop": (["trajectory"], lon.astype(lonlatdepth_dtype)), - "lat_nextloop": (["trajectory"], lat.astype(lonlatdepth_dtype)), - "depth_nextloop": (["trajectory"], depth.astype(lonlatdepth_dtype)), - "time_nextloop": (["trajectory"], time), - }, - coords={ - "trajectory": ("trajectory", trajectory_ids), - }, - attrs={ - "ngrid": len(fieldset.gridset), - "ptype": pclass.getPType(), - }, - ) + self._data = { + "lon": lon.astype(lonlatdepth_dtype), + "lat": lat.astype(lonlatdepth_dtype), + "depth": depth.astype(lonlatdepth_dtype), + "time": time, + "dt": np.timedelta64(1, "ns") * np.ones(len(trajectory_ids)), + # "ei": (["trajectory", "ngrid"], np.zeros((len(trajectory_ids), len(fieldset.gridset)), dtype=np.int32)), + "state": np.zeros((len(trajectory_ids)), dtype=np.int32), + "lon_nextloop": lon.astype(lonlatdepth_dtype), + "lat_nextloop": lat.astype(lonlatdepth_dtype), + "depth_nextloop": depth.astype(lonlatdepth_dtype), + "time_nextloop": time, + "trajectory": trajectory_ids, + } + self.ptype = pclass.getPType() # add extra fields from the custom Particle class for v in pclass.__dict__.values(): if isinstance(v, Variable): @@ -164,7 +157,7 @@ def __init__( initial = v.initial(self).values else: initial = v.initial * np.ones(len(trajectory_ids), dtype=v.dtype) - self._data[v.name] = (["trajectory"], initial) + self._data[v.name] = initial # update initial values provided on ParticleSet creation for kwvar, kwval in kwargs.items(): @@ -591,19 +584,19 @@ def Kernel(self, pyfunc): if isinstance(pyfunc, list): return Kernel.from_list( self.fieldset, - self._data.ptype, + self.ptype, pyfunc, ) return Kernel( self.fieldset, - self._data.ptype, + self.ptype, pyfunc=pyfunc, ) def InteractionKernel(self, pyfunc_inter): if pyfunc_inter is None: return None - return InteractionKernel(self.fieldset, self._data.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.""" From c99e29f526dbcd3d3dd348eea10f4ac06f655c33 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 18 Jul 2025 09:05:22 +0200 Subject: [PATCH 2/3] Only calling remove-indices when there are particles to be removed --- parcels/kernel.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/parcels/kernel.py b/parcels/kernel.py index 4b55a7634..911c890bd 100644 --- a/parcels/kernel.py +++ b/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): From 3ee1c371c1feee9b14491a1366ed4daf7cbc695d Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 18 Jul 2025 09:51:09 +0200 Subject: [PATCH 3/3] implementing particleset adding and deleting methods for dicts And fixing some unit tests --- parcels/particleset.py | 38 ++++++++++++++++++++-------- tests/v4/test_kernel.py | 2 +- tests/v4/test_particleset.py | 6 ++--- tests/v4/test_particleset_execute.py | 3 +-- 4 files changed, 32 insertions(+), 17 deletions(-) diff --git a/parcels/particleset.py b/parcels/particleset.py index 93de67e83..706952bf6 100644 --- a/parcels/particleset.py +++ b/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 diff --git a/tests/v4/test_kernel.py b/tests/v4/test_kernel.py index 1125d925c..d9e68babf 100644 --- a/tests/v4/test_kernel.py +++ b/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")) diff --git a/tests/v4/test_particleset.py b/tests/v4/test_particleset.py index 7a9dd7be6..633d66dd6 100644 --- a/tests/v4/test_particleset.py +++ b/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) diff --git a/tests/v4/test_particleset_execute.py b/tests/v4/test_particleset_execute.py index fe97546a3..022377029 100644 --- a/tests/v4/test_particleset_execute.py +++ b/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])