Repository: OceanParcels/parcels
Pull request #397 (2018-07-26 06:56:15)

Initialisation of custom Variables directly in ParticleSet.from_list()
This PR implements #389, so that now the value of custom `Variables` can be set in `ParticleSet.from_list()`, similar to how `lon` and `lat` are set there

This PR also improves  `print()` for `ParticleSets`, so that all custom `Variables` are also printed


Patch:
diff --git a/parcels/examples/example_peninsula.py b/parcels/examples/example_peninsula.py
index 87cf7b76..2e2c31c0 100644
--- a/parcels/examples/example_peninsula.py
+++ b/parcels/examples/example_peninsula.py
@@ -101,11 +101,6 @@ def pensinsula_example(fieldset, npart, mode='jit', degree=1,
         p = Variable('p', dtype=np.float32, initial=0.)
         p_start = Variable('p_start', dtype=np.float32, initial=fieldset.P)
 
-        def __repr__(self):
-            """Custom print function which overrides the built-in"""
-            return "P(%.4f, %.4f)[p=%.5f, p_start=%f]" % (self.lon, self.lat,
-                                                          self.p, self.p_start)
-
     # Initialise particles
     if fieldset.U.grid.mesh == 'flat':
         x = 3000  # 3 km offset from boundary
diff --git a/parcels/particle.py b/parcels/particle.py
index 151f883b..aff1496a 100644
--- a/parcels/particle.py
+++ b/parcels/particle.py
@@ -185,8 +185,11 @@ class ScipyParticle(_Particle):
 
     def __repr__(self):
         time_string = 'not_yet_set' if self.time is None or np.isnan(self.time) else "{:f}".format(self.time)
-        return "P[%d](lon=%f, lat=%f, depth=%f, time=%s)" % (self.id, self.lon, self.lat,
-                                                             self.depth, time_string)
+        str = "P[%d](lon=%f, lat=%f, depth=%f, " % (self.id, self.lon, self.lat, self.depth)
+        for var in vars(type(self)):
+            if type(getattr(type(self), var)) is Variable:
+                str += "%s=%f, " % (var, getattr(self, var))
+        return str + "time=%s)" % time_string
 
     def delete(self):
         self.state = ErrorCode.Delete
@@ -228,8 +231,3 @@ class JITParticle(ScipyParticle):
                 setattr(self, index, -1*np.ones((fieldset.gridset.size), dtype=np.int32))
             setattr(self, index+'p', getattr(self, index).ctypes.data_as(c_void_p))
             setattr(self, 'c'+index, getattr(self, index+'p').value)
-
-    def __repr__(self):
-        time_string = 'not_yet_set' if self.time is None or np.isnan(self.time) else "{:f}".format(self.time)
-        return "P[%d](lon=%f, lat=%f, depth=%f, time=%s)" % (self.id, self.lon, self.lat,
-                                                             self.depth, time_string)
diff --git a/parcels/particleset.py b/parcels/particleset.py
index 882039d0..b86acf25 100644
--- a/parcels/particleset.py
+++ b/parcels/particleset.py
@@ -27,9 +27,10 @@ class ParticleSet(object):
     :param depth: Optional list of initial depth values for particles. Default is 0m
     :param time: Optional list of initial time values for particles. Default is fieldset.U.grid.time[0]
     :param repeatdt: Optional interval (in seconds) on which to repeat the release of the ParticleSet
+    Other Variables can be initialised using further arguments (e.g. v=... for a Variable named 'v')
     """
 
-    def __init__(self, fieldset, pclass=JITParticle, lon=None, lat=None, depth=None, time=None, repeatdt=None):
+    def __init__(self, fieldset, pclass=JITParticle, lon=None, lat=None, depth=None, time=None, repeatdt=None, **kwargs):
         self.fieldset = fieldset
         self.fieldset.check_complete()
 
@@ -96,11 +97,16 @@ class ParticleSet(object):
 
             for i in range(size):
                 self.particles[i] = pclass(lon[i], lat[i], fieldset=fieldset, depth=depth[i], cptr=cptr(i), time=time[i])
+                # Set other Variables if provided
+                for kwvar in kwargs:
+                    if not hasattr(self.particles[i], kwvar):
+                        raise RuntimeError('Particle class does not have Variable %s' % kwvar)
+                    setattr(self.particles[i], kwvar, kwargs[kwvar][i])
         else:
             raise ValueError("Latitude and longitude required for generating ParticleSet")
 
     @classmethod
-    def from_list(cls, fieldset, pclass, lon, lat, depth=None, time=None, repeatdt=None):
+    def from_list(cls, fieldset, pclass, lon, lat, depth=None, time=None, repeatdt=None, **kwargs):
         """Initialise the ParticleSet from lists of lon and lat
 
         :param fieldset: :mod:`parcels.fieldset.FieldSet` object from which to sample velocity
@@ -111,8 +117,9 @@ class ParticleSet(object):
         :param depth: Optional list of initial depth values for particles. Default is 0m
         :param time: Optional list of start time values for particles. Default is fieldset.U.time[0]
         :param repeatdt: Optional interval (in seconds) on which to repeat the release of the ParticleSet
+        Other Variables can be initialised using further arguments (e.g. v=... for a Variable named 'v')
        """
-        return cls(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat, depth=depth, time=time, repeatdt=repeatdt)
+        return cls(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat, depth=depth, time=time, repeatdt=repeatdt, **kwargs)
 
     @classmethod
     def from_line(cls, fieldset, pclass, start, finish, size, depth=None, time=None, repeatdt=None):


Test patch:
diff --git a/tests/test_particle_sets.py b/tests/test_particle_sets.py
index 0dfb8af1..190e923a 100644
--- a/tests/test_particle_sets.py
+++ b/tests/test_particle_sets.py
@@ -37,6 +37,21 @@ def test_pset_create_line(fieldset, mode, npart=100):
     assert np.allclose([p.lat for p in pset], lat, rtol=1e-12)
 
 
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+def test_pset_create_list_with_customvariable(fieldset, mode, npart=100):
+    lon = np.linspace(0, 1, npart, dtype=np.float32)
+    lat = np.linspace(1, 0, npart, dtype=np.float32)
+
+    class MyParticle(ptype[mode]):
+        v = Variable('v')
+
+    v_vals = np.arange(npart)
+    pset = ParticleSet.from_list(fieldset, lon=lon, lat=lat, v=v_vals, pclass=MyParticle)
+    assert np.allclose([p.lon for p in pset], lon, rtol=1e-12)
+    assert np.allclose([p.lat for p in pset], lat, rtol=1e-12)
+    assert np.allclose([p.v for p in pset], v_vals, rtol=1e-12)
+
+
 @pytest.mark.parametrize('mode', ['scipy'])
 def test_pset_create_field(fieldset, mode, npart=100):
     np.random.seed(123456)