Merge pull request #290 from OceanParcels/Write_particledata_on_delete

Write particledata on delete

Author: delandmeterp

parcels/kernel.py

@@ -206,13 +206,15 @@ def execute_python(self, pset, endtime, dt):
                 else:
                     break  # Failure - stop time loop
 
-    def execute(self, pset, endtime, dt, recovery=None):
+    def execute(self, pset, endtime, dt, recovery=None, output_file=None):
         """Execute this Kernel over a ParticleSet for several timesteps"""
 
         def remove_deleted(pset):
             """Utility to remove all particles that signalled deletion"""
             indices = [i for i, p in enumerate(pset.particles)
                        if p.state in [ErrorCode.Delete]]
+            if len(indices) > 0 and output_file is not None:
+                output_file.write(pset[indices], endtime, deleted_only=True)
             pset.remove(indices)
 
         if recovery is None:

parcels/particlefile.py

@@ -32,12 +32,18 @@ class ParticleFile(object):
                      while ParticleFile is given as an argument of ParticleSet.execute()
                      It is either a timedelta object or a positive double.
     :param type: Either 'array' for default matrix style, or 'indexed' for indexed ragged array
+    :param write_ondelete: Boolean to write particle data only when they are deleted. Default is False
     """
 
-    def __init__(self, name, particleset, outputdt=np.infty, type='array'):
+    def __init__(self, name, particleset, outputdt=np.infty, type='array', write_ondelete=False):
 
         self.type = type
         self.name = name
+        self.write_ondelete = write_ondelete
+        self.outputdt = outputdt
+        if self.write_ondelete and self.type is 'array':
+            logger.warning('ParticleFile.write_ondelete=True requires type="indexed". Setting that option')
+            self.type = 'indexed'
         self.outputdt = outputdt
         self.lasttime_written = None  # variable to check if time has been written already
         self.dataset = netCDF4.Dataset("%s.nc" % name, "w", format="NETCDF4")
@@ -122,7 +128,7 @@ def sync(self):
         """Write all buffered data to disk"""
         self.dataset.sync()
 
-    def write(self, pset, time, sync=True):
+    def write(self, pset, time, sync=True, deleted_only=False):
         """Write :class:`parcels.particleset.ParticleSet` data to file
 
         :param pset: ParticleSet object to write
@@ -132,14 +138,17 @@ def write(self, pset, time, sync=True):
         """
         if isinstance(time, delta):
             time = time.total_seconds()
-        if self.lasttime_written != time:  # only write if 'time' hasn't been written yet
+        if self.lasttime_written != time and \
+           (self.write_ondelete is False or deleted_only is True):
+
             self.lasttime_written = time
             if self.type is 'array':
                 # Check if largest particle ID is smaller than the last ID in ParticleFile.
                 # Otherwise, new particles have been added and netcdf will fail
                 if pset.size > 0:
                     if max([p.id for p in pset]) > self.id[-1]:
                         logger.error("Number of particles appears to increase. Use type='indexed' for ParticleFile")
+                        exit(-1)
 
                     # Finds the indices (inds) of the particle IDs in the ParticleFile,
                     # because particles can have been deleted
@@ -165,7 +174,7 @@ def write(self, pset, time, sync=True):
                     logger.warning("Option to_write='once' is not fully functional in indexed mode! Particle properties of such variables are not written for newly added particles.")
                 ind = np.arange(pset.size) + self.idx
                 self.id[ind] = np.array([p.id for p in pset])
-                self.time[ind] = time
+                self.time[ind] = np.array([p.time for p in pset])
                 self.lat[ind] = np.array([p.lat for p in pset])
                 self.lon[ind] = np.array([p.lon for p in pset])
                 self.z[ind] = np.array([p.depth for p in pset])

parcels/particleset.py

@@ -333,7 +333,7 @@ def execute(self, pyfunc=AdvectionRK4, endtime=None, runtime=None, dt=1.,
                 time = min(next_prelease, next_input, next_output, next_movie, endtime)
             else:
                 time = max(next_prelease, next_input, next_output, next_movie, endtime)
-            self.kernel.execute(self, endtime=time, dt=dt, recovery=recovery)
+            self.kernel.execute(self, endtime=time, dt=dt, recovery=recovery, output_file=output_file)
             if abs(time-next_prelease) < tol:
                 self.add(ParticleSet(fieldset=self.fieldset, time=time, lon=self.repeatlon,
                                      lat=self.repeatlat, depth=self.repeatdepth,

tests/test_kernel_execution.py

@@ -54,10 +54,9 @@ def test_execution_runtime(fieldset, mode, start, end, substeps, dt, npart=10):
     pset = ParticleSet(fieldset, pclass=ptype[mode], time=start,
                        lon=np.linspace(0, 1, npart, dtype=np.float32),
                        lat=np.linspace(1, 0, npart, dtype=np.float32))
-    t_step = (end - start) / substeps
+    t_step = abs(end - start) / substeps
     for _ in range(substeps):
-        pset.execute(DoNothing, runtime=abs(t_step), dt=dt)
-        start += t_step
+        pset.execute(DoNothing, runtime=t_step, dt=dt)
     assert np.allclose(np.array([p.time for p in pset]), end)
 
 

tests/test_particle_sets.py

@@ -1,4 +1,5 @@
-from parcels import FieldSet, ParticleSet, Field, ScipyParticle, JITParticle, Variable
+from parcels import (FieldSet, ParticleSet, Field, ScipyParticle, JITParticle,
+                     Variable, ErrorCode)
 import numpy as np
 import pytest
 from netCDF4 import Dataset
@@ -329,3 +330,31 @@ class MyParticle(ptype[mode]):
     assert np.all([p.v_once == 11.0 for p in pset])
     assert (V_once.shape == (npart, ))
     assert (V_once[0] == 1.)
+
+
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+def test_variable_written_ondelete(fieldset, mode, tmpdir, npart=3):
+    filepath = tmpdir.join("pfile_on_delete_written_variables")
+
+    def move_west(particle, fieldset, time, dt):
+        tmp = fieldset.U[time, particle.lon, particle.lat, particle.depth]  # to trigger out-of-bounds error
+        particle.lon -= 0.1 + tmp
+
+    def DeleteP(particle, fieldset, time, dt):
+        particle.delete()
+
+    lon = np.linspace(0.05, 0.95, npart, dtype=np.float32)
+    lat = np.linspace(0.95, 0.05, npart, dtype=np.float32)
+
+    (dt, runtime) = (0.1, 0.8)
+    lon_end = lon - runtime/dt*0.1
+    noutside = len(lon_end[lon_end < 0])
+
+    pset = ParticleSet(fieldset, pclass=ptype[mode], lon=lon, lat=lat)
+
+    outfile = pset.ParticleFile(name=filepath, write_ondelete=True, type="indexed")
+    pset.execute(move_west, runtime=runtime, dt=dt, output_file=outfile,
+                 recovery={ErrorCode.ErrorOutOfBounds: DeleteP})
+    ncfile = Dataset(filepath+".nc", 'r', 'NETCDF4')
+    lon = ncfile.variables['lon'][:]
+    assert (lon.size == noutside)