Repository: OceanParcels/parcels
Pull request #534 (2019-02-18 14:54:37)

Support for custom variable initialisation from ParticleSet.__init__() with repeatdt
It was already possible to initialise a custom variable from `ParticleSet.__init()__` using a `kwargs` (e.g. `pset = ParticleSet(...,custom_var=3)`). However, this was not possible when also using repeated releases with `repeatdt` argument. This PR fixes that



Patch:
diff --git a/parcels/particleset.py b/parcels/particleset.py
index 7a4f47d6..e380f21a 100644
--- a/parcels/particleset.py
+++ b/parcels/particleset.py
@@ -64,6 +64,9 @@ class ParticleSet(object):
             raise NotImplementedError('If fieldset.time_origin is not a date, time of a particle must be a double')
         time = [self.time_origin.reltime(t) if isinstance(t, np.datetime64) else t for t in time]
 
+        for kwvar in kwargs:
+            kwargs[kwvar] = convert_to_list(kwargs[kwvar])
+
         assert len(lon) == len(time)
 
         self.repeatdt = repeatdt.total_seconds() if isinstance(repeatdt, delta) else repeatdt
@@ -77,6 +80,7 @@ class ParticleSet(object):
             self.repeatlat = lat
             self.repeatdepth = depth
             self.repeatpclass = pclass
+            self.repeatkwargs = kwargs
 
         size = len(lon)
         self.particles = np.empty(size, dtype=pclass)
@@ -379,7 +383,7 @@ class ParticleSet(object):
             if abs(time-next_prelease) < tol:
                 pset_new = ParticleSet(fieldset=self.fieldset, time=time, lon=self.repeatlon,
                                        lat=self.repeatlat, depth=self.repeatdepth,
-                                       pclass=self.repeatpclass)
+                                       pclass=self.repeatpclass, **self.repeatkwargs)
                 for p in pset_new:
                     p.dt = dt
                 self.add(pset_new)


Test patch:
diff --git a/tests/test_particle_sets.py b/tests/test_particle_sets.py
index 4aead909..3ad1fa2e 100644
--- a/tests/test_particle_sets.py
+++ b/tests/test_particle_sets.py
@@ -126,7 +126,7 @@ def test_pset_repeated_release(fieldset, mode, npart=10):
     assert np.allclose([p.lon for p in pset], np.arange(npart, 0, -1))
 
 
-@pytest.mark.parametrize('type', ['releasedt', 'timearr'])
+@pytest.mark.parametrize('type', ['repeatdt', 'timearr'])
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 @pytest.mark.parametrize('repeatdt', range(1, 3))
 @pytest.mark.parametrize('dt', [-1, 1])
@@ -136,12 +136,12 @@ def test_pset_repeated_release_delayed_adding_deleting(type, fieldset, mode, rep
 
     class MyParticle(ptype[mode]):
         sample_var = Variable('sample_var', initial=0.)
-    if type == 'releasedt':
+    if type == 'repeatdt':
         pset = ParticleSet(fieldset, lon=[0], lat=[0], pclass=MyParticle, repeatdt=repeatdt)
     elif type == 'timearr':
         pset = ParticleSet(fieldset, lon=np.zeros(runtime), lat=np.zeros(runtime), pclass=MyParticle, time=list(range(runtime)))
     outfilepath = tmpdir.join("pfile_repeated_release")
-    pfile = pset.ParticleFile(outfilepath, outputdt=abs(dt), chunksizes=[len(pset), 1])
+    pfile = pset.ParticleFile(outfilepath, outputdt=abs(dt))
 
     def IncrLon(particle, fieldset, time):
         particle.sample_var += 1.
@@ -152,7 +152,7 @@ def test_pset_repeated_release_delayed_adding_deleting(type, fieldset, mode, rep
     ncfile = Dataset(outfilepath+".nc", 'r', 'NETCDF4')
     samplevar = ncfile.variables['sample_var'][:]
     ncfile.close()
-    if type == 'releasedt':
+    if type == 'repeatdt':
         assert samplevar.shape == (runtime // repeatdt+1, min(maxvar+1, runtime)+1)
         assert np.allclose([p.sample_var for p in pset], np.arange(maxvar, -1, -repeatdt))
     elif type == 'timearr':
@@ -173,6 +173,20 @@ def test_pset_repeatdt_check_dt(fieldset):
     assert np.allclose([p.lon for p in pset], 1)  # if p.dt is nan, it won't be executed so p.lon will be 0
 
 
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+def test_pset_repeatdt_custominit(fieldset, mode):
+    class MyParticle(ptype[mode]):
+        sample_var = Variable('sample_var')
+
+    pset = ParticleSet(fieldset, lon=0, lat=0, pclass=MyParticle, repeatdt=1, sample_var=5)
+
+    def DoNothing(particle, fieldset, time):
+        return ErrorCode.Success
+
+    pset.execute(DoNothing, dt=1, runtime=21)
+    assert np.allclose([p.sample_var for p in pset], 5.)
+
+
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 def test_pset_access(fieldset, mode, npart=100):
     lon = np.linspace(0, 1, npart, dtype=np.float32)