Add modification to manage array variable

Author: adelepoulle

src/py_eddy_tracker/observations.py

@@ -67,8 +67,10 @@ class EddiesObservations(object):
         Time
     """
 
-    def __init__(self, size=0, track_extra_variables=False):
+    def __init__(self, size=0, track_extra_variables=False, track_array_variables=0, array_variables=[]):
         self.track_extra_variables = track_extra_variables
+        self.track_array_variables = track_array_variables
+        self.array_variables = array_variables
         for elt in self.elements:
             if elt not in VAR_DESCR:
                 raise Exception('Unknown element : %s' % elt)
@@ -90,9 +92,13 @@ def dtype(self):
         """
         dtype = list()
         for elt in self.elements:
-            dtype.append((elt, VAR_DESCR[elt][
+            data_type = VAR_DESCR[elt][
                 'compute_type' if 'compute_type' in VAR_DESCR[elt] else
-                'nc_type']))
+                'nc_type']
+            if elt in self.array_variables:
+                dtype.append((elt, data_type, (self.track_array_variables,)))
+            else:
+                dtype.append((elt, data_type))
         return dtype
 
     @property
@@ -108,7 +114,9 @@ def elements(self):
             'speed_radius',  # 'uavg'
             'eke',  # 'teke'
             'time']  # 'rtime'
-
+        if self.track_array_variables > 0:
+            elements += self.array_variables
+        
         if self.track_extra_variables:
             elements += ['contour_e',
                          'contour_s',
@@ -120,7 +128,9 @@ def elements(self):
     def coherence(self, other):
         """Check coherence between two dataset
         """
-        return self.track_extra_variables == other.track_extra_variables
+        test = self.track_array_variables == other.track_array_variables
+        test *= self.array_variables == other.array_variables
+        return test
 
     def merge(self, other):
         """Merge two dataset
@@ -133,9 +143,12 @@ def merge(self, other):
         eddies.sign_type = self.sign_type
         return eddies
 
+    def reset(self):
+        self.observations = np.zeros(0, dtype=self.dtype)
+
     @property
     def obs(self):
-        """returan array observations
+        """return an array observations
         """
         return self.observations
 
@@ -161,7 +174,10 @@ def insert_observations(self, other, index):
             return self
         if index < 0:
             index = self_size + index + 1
-        eddies = self.__class__(new_size, self.track_extra_variables)
+        eddies = self.__class__(new_size, self.track_extra_variables,
+            track_array_variables=self.track_array_variables,
+            array_variables=self.array_variables
+            )
         eddies.obs[:index] = self.obs[:index]
         eddies.obs[index: index + insert_size] = other.obs
         eddies.obs[index + insert_size:] = self.obs[index:]
@@ -198,9 +214,18 @@ def index(self, index):
 
     @staticmethod
     def load_from_netcdf(filename):
+        array_dim = 'NbSample'
         with Dataset(filename) as h_nc:
             nb_obs = len(h_nc.dimensions['Nobs'])
-            eddies = EddiesObservations(size=nb_obs)
+            kwargs = dict()
+            if array_dim in h_nc.dimensions:
+                kwargs['track_array_variables'] = len(h_nc.dimensions[array_dim])
+                kwargs['array_variables'] = []
+                for variable in h_nc.variables:
+                    if array_dim in h_nc.variables[variable].dimensions:
+                        kwargs['array_variables'].append(str(variable))
+
+            eddies = EddiesObservations(size=nb_obs, ** kwargs)
             for variable in h_nc.variables:
                 if variable == 'cyc':
                     continue
@@ -297,8 +322,9 @@ def filled_by_interpolation(self, mask):
         nb_obs = len(self)
         index = arange(nb_obs)
 
-        for var, _ in self.obs.dtype.descr:
-            if var in ['n', 'virtual', 'track']:
+        for field in self.obs.dtype.descr:
+            var = field[0]
+            if var in ['n', 'virtual', 'track'] or var in self.array_variables:
                 continue
             self.obs[var][mask] = interp(index[mask], index[-mask],
                                          self.obs[var][-mask])
@@ -309,9 +335,14 @@ def extract_longer_eddies(self, nb_min, nb_obs, compress_id=True):
         mask = nb_obs >= nb_min
         nb_obs_select = mask.sum()
         logging.info('Selection of %d observations', nb_obs_select)
-        eddies = TrackEddiesObservations(size=nb_obs_select)
+        eddies = TrackEddiesObservations(
+            size=nb_obs_select,
+            track_array_variables=self.track_array_variables,
+            array_variables=self.array_variables
+            )
         eddies.sign_type = self.sign_type
-        for var, _ in eddies.obs.dtype.descr:
+        for field in self.obs.dtype.descr:
+            var = field[0]
             eddies.obs[var] = self.obs[var][mask]
         if compress_id:
             list_id = unique(eddies.obs['track'])
@@ -329,22 +360,20 @@ def elements(self):
 
     @staticmethod
     def create_variable(handler_nc, kwargs_variable,
-                        attr_variable, data, scale_factor=None):
-        """Create variable
-        """
+                        attr_variable, data, scale_factor=None, add_offset=None):
         var = handler_nc.createVariable(
             zlib=True,
             complevel=1,
             **kwargs_variable)
         for attr, attr_value in attr_variable.iteritems():
             var.setncattr(attr, attr_value)
-
-        var[:] = data
-
-        #~ var.set_auto_maskandscale(False)
         if scale_factor is not None:
             var.scale_factor = scale_factor
-
+            if add_offset is not None:
+                var.add_offset = add_offset
+            else:
+                var.add_offset = 0
+        var[:] = data
         try:
             var.setncattr('min', var[:].min())
             var.setncattr('max', var[:].max())
@@ -362,19 +391,21 @@ def write_netcdf(self, path='./'):
             # Create dimensions
             logging.debug('Create Dimensions "Nobs" : %d', eddy_size)
             h_nc.createDimension('Nobs', eddy_size)
+            if self.track_array_variables != 0:
+                h_nc.createDimension('NbSample', self.track_array_variables)
             # Iter on variables to create:
-            for name, _ in self.observations.dtype.descr:
+            for field in self.observations.dtype.descr:
+                name = field[0]
                 logging.debug('Create Variable %s', VAR_DESCR[name]['nc_name'])
                 self.create_variable(
                     h_nc,
                     dict(varname=VAR_DESCR[name]['nc_name'],
-                         datatype=VAR_DESCR[name]['nc_type'],
+                         datatype=VAR_DESCR[name]['output_type'],
                          dimensions=VAR_DESCR[name]['nc_dims']),
                     VAR_DESCR[name]['nc_attr'],
                     self.observations[name],
-                    scale_factor=None
-                        if 'scale_factor' not in VAR_DESCR[name] else
-                        VAR_DESCR[name]['scale_factor']
+                    scale_factor=VAR_DESCR[name].get('scale_factor', None),
+                    add_offset=VAR_DESCR[name].get('add_offset', None)
                     )
 
             # Add cyclonic information

src/py_eddy_tracker/tracking.py

@@ -278,8 +278,16 @@ def prepare_merging(self):
     def merge(self):
         """Merge all the correspondance in one array with all fields
         """
+        # Start loading identification again to save in the finals tracks
+        # Load first file
+        self.swap_dataset(self.datasets[0])
+
         # Start create netcdf to agglomerate all eddy
-        eddies = TrackEddiesObservations(size=self.nb_obs)
+        eddies = TrackEddiesObservations(
+            size=self.nb_obs,
+            track_array_variables=self.current_obs.track_array_variables,
+            array_variables=self.current_obs.array_variables,
+            )
 
         # Calculate the index in each tracks, we compute in u4 and translate
         # in u2 (which are limited to 65535)
@@ -291,9 +299,6 @@ def merge(self):
         eddies['track'][:] = arange(self.current_id
                                     ).repeat(self.nb_obs_by_tracks)
 
-        # Start loading identification again to save in the finals tracks
-        # Load first file
-        self.swap_dataset(self.datasets[0])
         # Set type of eddy with first file
         eddies.sign_type = self.current_obs.sign_type
         # Fields to copy
@@ -317,7 +322,8 @@ def merge(self):
                 # Index in the previous file
                 index_in = self[i]['in'][m_first_obs]
                 # Copy all variable
-                for var, _ in fields:
+                for field in fields:
+                    var = field[0]
                     eddies[var][index_final[m_first_obs]
                         ] = self.previous_obs[var][index_in]
                 # Increment
@@ -341,7 +347,8 @@ def merge(self):
             index_current = self[i]['out']
 
             # Copy all variable
-            for var, _ in fields:
+            for field in fields:
+                var = field[0]
                 eddies[var][index_final
                     ] = self.current_obs[var][index_current]