Merge pull request #1199 from OceanParcels/dump_to_zarr

Saving Parcels output directly in zarr format

Author: erikvansebille

environment_py3_win.yml

@@ -19,7 +19,7 @@ dependencies:
   - scipy>=0.16.0
   - tqdm
   - xarray>=0.5.1
-  - dask>=2.0
+  - dask<=2022.9.0
   - cftime>=1.3.1
   - ipykernel
   - pytest

parcels/collection/collectionaos.py

@@ -1,4 +1,3 @@
-from datetime import timedelta as delta
 from operator import attrgetter  # NOQA
 
 from ctypes import c_void_p
@@ -28,24 +27,6 @@
 __all__ = ['ParticleCollectionAOS', 'ParticleCollectionIterableAOS', 'ParticleCollectionIteratorAOS']
 
 
-def _to_write_particles(pd, time):
-    """We don't want to write a particle that is not started yet.
-    Particle will be written if particle.time is between time-dt/2 and time+dt (/2)
-    """
-    return [i for i, p in enumerate(pd) if (((time - np.abs(p.dt/2) <= p.time < time + np.abs(p.dt))
-                                             or (np.isnan(p.dt) and np.equal(time, p.time)))
-                                            and np.isfinite(p.id))]
-
-
-def _is_particle_started_yet(particle, time):
-    """We don't want to write a particle that is not started yet.
-    Particle will be written if:
-      * particle.time is equal to time argument of pfile.write()
-      * particle.time is before time (in case particle was deleted between previous export and current one)
-    """
-    return (particle.dt*particle.time <= particle.dt*time or np.isclose(particle.time, time))
-
-
 def _convert_to_flat_array(var):
     """Convert lists and single integers/floats to one-dimensional numpy arrays
 
@@ -890,64 +871,33 @@ def cstruct(self):
         cstruct = self._data_c.ctypes.data_as(c_void_p)
         return cstruct
 
-    def toDictionary(self, pfile, time, deleted_only=False):
-        """
-        Convert all Particle data from one time step to a python dictionary.
-        :param pfile: ParticleFile object requesting the conversion
-        :param time: Time at which to write ParticleSet
-        :param deleted_only: Flag to write only the deleted Particles
-        returns two dictionaries: one for all variables to be written each outputdt,
-         and one for all variables to be written once
-
-        This function depends on the specific collection in question and thus needs to be specified in specific
-        derivative classes.
+    def _to_write_particles(self, pd, time):
+        """We don't want to write a particle that is not started yet.
+        Particle will be written if particle.time is between time-dt/2 and time+dt (/2)
         """
-        data_dict = {}
-        data_dict_once = {}
-
-        time = time.total_seconds() if isinstance(time, delta) else time
+        return np.array([i for i, p in enumerate(pd) if (((time - np.abs(p.dt/2) <= p.time < time + np.abs(p.dt))
+                                                         or (np.isnan(p.dt) and np.equal(time, p.time)))
+                                                         and np.isfinite(p.id))])
 
-        indices_to_write = []
-        if pfile.lasttime_written != time and \
-           (pfile.write_ondelete is False or deleted_only):
-            if self._ncount == 0:
-                logger.warning("ParticleSet is empty on writing as array at time %g" % time)
-            else:
-                if deleted_only:
-                    if type(deleted_only) not in [list, np.ndarray] and deleted_only in [True, 1]:
-                        data_states = [p.state for p in self._data]
-                        indices_to_write = np.where(np.isin(data_states, [OperationCode.Delete]))[0]
-                    elif type(deleted_only) in [list, np.ndarray] and len(deleted_only) > 0:
-                        if type(deleted_only[0]) in [int, np.int32, np.uint32]:
-                            indices_to_write = deleted_only
-                        elif isinstance(deleted_only[0], ScipyParticle):
-                            indices_to_write = [i for i, p in self._data if p in deleted_only]
-                else:
-                    indices_to_write = _to_write_particles(self._data, time)
-                if len(indices_to_write) > 0:
-                    for var in pfile.var_names:
-                        if 'id' in var:
-                            data_dict[var] = np.array([np.int64(getattr(p, var)) for p in self._data[indices_to_write]])
-                        else:
-                            data_dict[var] = np.array([getattr(p, var) for p in self._data[indices_to_write]])
-
-                pset_errs = [p for p in self._data[indices_to_write] if p.state != OperationCode.Delete and abs(time-p.time) > 1e-3 and np.isfinite(p.time)]
-                for p in pset_errs:
-                    logger.warning_once('time argument in pfile.write() is %g, but a particle has time % g.' % (time, p.time))
-
-                if len(pfile.var_names_once) > 0:
-                    # _to_write_particles(self._data, time)
-                    first_write = [p for p in self._data if _is_particle_started_yet(p, time) and (np.int64(p.id) not in pfile.written_once)]
-                    if np.any(first_write):
-                        data_dict_once['id'] = np.array([p.id for p in first_write]).astype(dtype=np.int64)
-                        for var in pfile.var_names_once:
-                            data_dict_once[var] = np.array([getattr(p, var) for p in first_write])
-                        pfile.written_once.extend(np.array(data_dict_once['id']).astype(dtype=np.int64).tolist())
-
-            if deleted_only is False:
-                pfile.lasttime_written = time
+    def getvardata(self, var, indices=None):
+        if indices is None:
+            return np.array([getattr(p, var) for p in self._data])
+        else:
+            try:
+                return np.array([getattr(p, var) for p in self._data[indices]])
+            except:  # Can occur for zero-length ParticleSets
+                return None
+
+    def setvardata(self, var, index, val):
+        if isinstance(index, (np.int64, int, np.int32)):
+            setattr(self._data[index], var, val)
+        else:
+            for i, v in zip(index, val):
+                setattr(self._data[i], var, v)
 
-        return data_dict, data_dict_once
+    def setallvardata(self, var, val):
+        for i in range(len(self._data)):
+            setattr(self._data[i], var, val)
 
     def toArray(self):
         """

parcels/collection/collections.py

@@ -898,19 +898,18 @@ def __getattr__(self, name):
         else:
             return False
 
+    def has_write_once_variables(self):
+        for var in self.ptype.variables:
+            if var.to_write == 'once':
+                return True
+        return False
+
     @abstractmethod
-    def toDictionary(self):
-        """
-        Convert all Particle data from one time step to a python dictionary.
-        :param pfile: ParticleFile object requesting the conversion
-        :param time: Time at which to write ParticleSet
-        :param deleted_only: Flag to write only the deleted Particles
-        returns two dictionaries: one for all variables to be written each outputdt,
-         and one for all variables to be written once
+    def getvardata(self, var, indices=None):
+        pass
 
-         This function depends on the specific collection in question and thus needs to be specified in specific
-         derivatives classes.
-        """
+    @abstractmethod
+    def setvardata(self, var, index, val):
         pass
 
     @abstractmethod

parcels/collection/collectionsoa.py

@@ -1,4 +1,3 @@
-from datetime import timedelta as delta
 from operator import attrgetter
 from ctypes import Structure, POINTER
 from bisect import bisect_left
@@ -26,26 +25,6 @@
                                'See http://oceanparcels.org/#parallel_install for more information')
 
 
-def _to_write_particles(pd, time):
-    """We don't want to write a particle that is not started yet.
-    Particle will be written if particle.time is between time-dt/2 and time+dt (/2)
-    """
-    return ((np.less_equal(time - np.abs(pd['dt']/2), pd['time'], where=np.isfinite(pd['time']))
-             & np.greater_equal(time + np.abs(pd['dt'] / 2), pd['time'], where=np.isfinite(pd['time']))
-             | ((np.isnan(pd['dt'])) & np.equal(time, pd['time'], where=np.isfinite(pd['time']))))
-            & (np.isfinite(pd['id']))
-            & (np.isfinite(pd['time'])))
-
-
-def _is_particle_started_yet(pd, time):
-    """We don't want to write a particle that is not started yet.
-    Particle will be written if:
-      * particle.time is equal to time argument of pfile.write()
-      * particle.time is before time (in case particle was deleted between previous export and current one)
-    """
-    return np.less_equal(pd['dt']*pd['time'], pd['dt']*time) | np.isclose(pd['time'], time)
-
-
 def _convert_to_flat_array(var):
     """Convert lists and single integers/floats to one-dimensional numpy arrays
 
@@ -153,7 +132,7 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, p
             self._data['depth'][:] = depth
             self._data['time'][:] = time
             self._data['id'][:] = pid
-            self._data['fileid'][:] = -1
+            self._data['once_written'][:] = 0
 
             # special case for exceptions which can only be handled from scipy
             self._data['exception'] = np.empty(self.ncount, dtype=object)
@@ -815,58 +794,30 @@ def flatten_dense_data_array(vname):
         cstruct = CParticles(*cdata)
         return cstruct
 
-    def toDictionary(self, pfile, time, deleted_only=False):
+    def _to_write_particles(self, pd, time):
+        """We don't want to write a particle that is not started yet.
+        Particle will be written if particle.time is between time-dt/2 and time+dt (/2)
         """
-        Convert all Particle data from one time step to a python dictionary.
-        :param pfile: ParticleFile object requesting the conversion
-        :param time: Time at which to write ParticleSet
-        :param deleted_only: Flag to write only the deleted Particles
-        returns two dictionaries: one for all variables to be written each outputdt,
-         and one for all variables to be written once
+        return np.where((np.less_equal(time - np.abs(pd['dt'] / 2), pd['time'], where=np.isfinite(pd['time']))
+                        & np.greater_equal(time + np.abs(pd['dt'] / 2), pd['time'], where=np.isfinite(pd['time']))
+                        | ((np.isnan(pd['dt'])) & np.equal(time, pd['time'], where=np.isfinite(pd['time']))))
+                        & (np.isfinite(pd['id']))
+                        & (np.isfinite(pd['time'])))[0]
 
-        This function depends on the specific collection in question and thus needs to be specified in specific
-        derivative classes.
-        """
-
-        data_dict = {}
-        data_dict_once = {}
+    def getvardata(self, var, indices=None):
+        if indices is None:
+            return self._data[var]
+        else:
+            try:
+                return self._data[var][indices]
+            except:  # Can occur for zero-length ParticleSets
+                return None
 
-        time = time.total_seconds() if isinstance(time, delta) else time
+    def setvardata(self, var, index, val):
+        self._data[var][index] = val
 
-        indices_to_write = []
-        if pfile.lasttime_written != time and \
-           (pfile.write_ondelete is False or deleted_only is not False):
-            if self._data['id'].size == 0:
-                logger.warning("ParticleSet is empty on writing as array at time %g" % time)
-            else:
-                if deleted_only is not False:
-                    if type(deleted_only) not in [list, np.ndarray] and deleted_only in [True, 1]:
-                        indices_to_write = np.where(np.isin(self._data['state'],
-                                                            [OperationCode.Delete]))[0]
-                    elif type(deleted_only) in [list, np.ndarray]:
-                        indices_to_write = deleted_only
-                else:
-                    indices_to_write = _to_write_particles(self._data, time)
-                if np.any(indices_to_write):
-                    for var in pfile.var_names:
-                        data_dict[var] = self._data[var][indices_to_write]
-
-                pset_errs = ((self._data['state'][indices_to_write] != OperationCode.Delete) & np.greater(np.abs(time - self._data['time'][indices_to_write]), 1e-3, where=np.isfinite(self._data['time'][indices_to_write])))
-                if np.count_nonzero(pset_errs) > 0:
-                    logger.warning_once('time argument in pfile.write() is {}, but particles have time {}'.format(time, self._data['time'][pset_errs]))
-
-                if len(pfile.var_names_once) > 0:
-                    first_write = (_to_write_particles(self._data, time) & _is_particle_started_yet(self._data, time) & np.isin(self._data['id'], pfile.written_once, invert=True))
-                    if np.any(first_write):
-                        data_dict_once['id'] = np.array(self._data['id'][first_write]).astype(dtype=np.int64)
-                        for var in pfile.var_names_once:
-                            data_dict_once[var] = self._data[var][first_write]
-                        pfile.written_once.extend(np.array(self._data['id'][first_write]).astype(dtype=np.int64).tolist())
-
-            if deleted_only is False:
-                pfile.lasttime_written = time
-
-        return data_dict, data_dict_once
+    def setallvardata(self, var, val):
+        self._data[var][:] = val
 
     def toArray(self):
         """