Adding Support for Custom Variables

And moving from temporary TestParticle to original Particle class

Author: erikvansebille

parcels/particle.py

@@ -1,10 +1,7 @@
-from operator import attrgetter
 from typing import Literal
 
 import numpy as np
 
-from parcels.tools.statuscodes import StatusCode
-
 __all__ = ["InteractionParticle", "Particle", "Variable"]
 
 
@@ -110,44 +107,24 @@ class Particle:
     Additional Variables can be added via the :Class Variable: objects
     """
 
-    lon = Variable("lon", dtype=np.float32)
-    lon_nextloop = Variable("lon_nextloop", dtype=np.float32, to_write=False)
-    lat = Variable("lat", dtype=np.float32)
-    lat_nextloop = Variable("lat_nextloop", dtype=np.float32, to_write=False)
-    depth = Variable("depth", dtype=np.float32)
-    depth_nextloop = Variable("depth_nextloop", dtype=np.float32, to_write=False)
-    time = Variable("time", dtype=np.float64)
-    time_nextloop = Variable("time_nextloop", dtype=np.float64, to_write=False)
-    id = Variable("id", dtype=np.int64, to_write="once")
-    obs_written = Variable("obs_written", dtype=np.int32, initial=0, to_write=False)
-    dt = Variable("dt", dtype=np.float64, to_write=False)
-    state = Variable("state", dtype=np.int32, initial=StatusCode.Evaluate, to_write=False)
-
-    lastID = 0  # class-level variable keeping track of last Particle ID used
-
-    def __init__(self, lon, lat, pid, fieldset=None, ngrids=None, depth=0.0, time=0.0, cptr=None):
-        # Enforce default values through Variable descriptor
-        type(self).lon.initial = lon
-        type(self).lon_nextloop.initial = lon
-        type(self).lat.initial = lat
-        type(self).lat_nextloop.initial = lat
-        type(self).depth.initial = depth
-        type(self).depth_nextloop.initial = depth
-        type(self).time.initial = time
-        type(self).time_nextloop.initial = time
-        type(self).id.initial = pid
-        type(self).lastID = max(type(self).lastID, pid)
-        type(self).obs_written.initial = 0
-        type(self).dt.initial = None
-
-        ptype = self.getPType()
-        # Explicit initialisation of all particle variables
-        for v in ptype.variables:
-            if isinstance(v.initial, attrgetter):
-                initial = v.initial(self)
-            else:
-                initial = v.initial
-            setattr(self, v.name, v.dtype(initial))
+    def __init__(self, data, index=None):
+        self._data = data
+        self._index = index
+
+    def __getattr__(self, name):
+        if name in ["_data", "_index"]:
+            return object.__getattribute__(self, name)
+        _data = object.__getattribute__(self, "_data")
+        if name in _data:
+            return _data[name].values[self._index]
+        else:
+            return False
+
+    def __setattr__(self, name, value):
+        if name in ["_data", "_index"]:
+            object.__setattr__(self, name, value)
+        else:
+            self._data[name][self._index] = value
 
     def __repr__(self):
         time_string = "not_yet_set" if self.time is None or np.isnan(self.time) else f"{self.time:f}"

parcels/particleset.py

@@ -1,6 +1,7 @@
 import sys
 import warnings
 from collections.abc import Iterable
+from operator import attrgetter
 
 import numpy as np
 import xarray as xr
@@ -13,7 +14,7 @@
 from parcels.basegrid import GridType
 from parcels.interaction.interactionkernel import InteractionKernel
 from parcels.kernel import Kernel
-from parcels.particle import Particle
+from parcels.particle import Particle, Variable
 from parcels.particlefile import ParticleFile
 from parcels.tools.converters import convert_to_flat_array
 from parcels.tools.loggers import logger
@@ -23,28 +24,6 @@
 __all__ = ["ParticleSet"]
 
 
-class TestParticle:
-    # Temporary class to allow for testing of ParticleSet without needing to change v3-Particle class. TODO update the Particle class
-    def __init__(self, data, index=None):
-        self._data = data
-        self._index = index
-
-    def __getattr__(self, name):
-        if name in ["_data", "_index"]:
-            return object.__getattribute__(self, name)
-        _data = object.__getattribute__(self, "_data")
-        if name in _data:
-            return _data[name].values[self._index]
-        else:
-            return False
-
-    def __setattr__(self, name, value):
-        if name in ["_data", "_index"]:
-            object.__setattr__(self, name, value)
-        else:
-            self._data[name][self._index] = value
-
-
 class ParticleSet:
     """Class for storing particle and executing kernel over them.
 
@@ -179,6 +158,20 @@ def __init__(
                 "ptype": self._pclass.getPType(),  # TODO check why both pclass and ptype needed
             },
         )
+        # add extra fields from the custom Particle class
+        for v in self.pclass.__dict__.values():
+            if isinstance(v, Variable):
+                if isinstance(v.initial, attrgetter):
+                    initial = v.initial(self).values
+                else:
+                    initial = v.initial * np.ones(len(pid_orig), dtype=v.dtype)
+                self._data[v.name] = (["trajectory"], initial)
+
+        # update initial values provided on ParticleSet creation
+        for kwvar, kwval in kwargs.items():
+            if not hasattr(pclass, kwvar):
+                raise RuntimeError(f"Particle class does not have Variable {kwvar}")
+            self._data[kwvar][:] = kwval
 
         self._kernel = None
 
@@ -216,7 +209,7 @@ def __getattr__(self, name):
 
     def __getitem__(self, index):
         """Get a single particle by index."""
-        return TestParticle(self._data, index=index)
+        return Particle(self._data, index=index)
 
     @staticmethod
     def lonlatdepth_dtype_from_field_interp_method(field):

tests/v4/test_particleset.py

@@ -1,5 +1,6 @@
 from contextlib import nullcontext as does_not_raise
 from datetime import datetime, timedelta
+from operator import attrgetter
 
 import numpy as np
 import pytest
@@ -66,6 +67,39 @@ def test_create_empty_pset(fieldset):
     assert pset.size == 0
 
 
+def test_pset_custominit_on_pset(fieldset):
+    MyParticle = Particle.add_variable("sample_var")
+
+    pset = ParticleSet(fieldset, lon=0, lat=0, pclass=MyParticle, sample_var=5)
+
+    pset.execute(DoNothing, dt=np.timedelta64(1, "s"), runtime=np.timedelta64(21, "s"))
+    assert np.allclose([p.sample_var for p in pset], 5.0)
+
+
+def test_pset_custominit_on_pset_attrgetter(fieldset):
+    MyParticle = Particle.add_variable("sample_var", initial=attrgetter("lon"))
+
+    pset = ParticleSet(fieldset, lon=3, lat=0, pclass=MyParticle)
+
+    pset.execute(DoNothing, dt=np.timedelta64(1, "s"), runtime=np.timedelta64(21, "s"))
+    assert np.allclose([p.sample_var for p in pset], 3.0)
+
+
+@pytest.mark.parametrize("pset_override", [True, False])
+def test_pset_custominit_on_pclass(fieldset, pset_override):
+    MyParticle = Particle.add_variable("sample_var", initial=4)
+
+    if pset_override:
+        pset = ParticleSet(fieldset, lon=0, lat=0, pclass=MyParticle, sample_var=5)
+    else:
+        pset = ParticleSet(fieldset, lon=0, lat=0, pclass=MyParticle)
+
+    pset.execute(DoNothing, dt=np.timedelta64(1, "s"), runtime=np.timedelta64(21, "s"))
+
+    check_val = 5.0 if pset_override else 4.0
+    assert np.allclose([p.sample_var for p in pset], check_val)
+
+
 @pytest.mark.parametrize(
     "time, expectation",
     [