Restart on last tracking

Author: AntSimi

src/py_eddy_tracker/__init__.py

@@ -123,23 +123,23 @@ def parse_args(self, *args, **kwargs):
         ),
     segment_size=dict(
         attr_name=None,
-        nc_name=None,
+        nc_name='segment_size',
         nc_type='byte',
-        nc_dims=None,
+        nc_dims=('Nobs',),
         nc_attr=dict()
         ),
     dlon=dict(
         attr_name=None,
-        nc_name=None,
+        nc_name='dlon',
         nc_type='float64',
-        nc_dims=None,
+        nc_dims=('Nobs',),
         nc_attr=dict()
         ),
     dlat=dict(
         attr_name=None,
-        nc_name=None,
+        nc_name='dlat',
         nc_type='float64',
-        nc_dims=None,
+        nc_dims=('Nobs',),
         nc_attr=dict()
         ),
     virtual=dict(

src/py_eddy_tracker/observations.py

@@ -208,6 +208,10 @@ def __init__(self, size=0, track_extra_variables=None,
         self.active = True
         self.sign_type = None
 
+    @property
+    def shape(self):
+        return self.observations.shape
+
     def __repr__(self):
         return str(self.observations)
 
@@ -369,11 +373,26 @@ def load_from_netcdf(cls, filename):
             eddies.sign_type = h_nc.variables['cyc'][0]
         return eddies
 
+    @classmethod
+    def from_netcdf(cls, handler):
+        nb_obs = len(handler.dimensions['Nobs'])
+        kwargs = dict()
+        if hasattr(handler, 'track_array_variables'):
+            kwargs['track_array_variables'] = handler.track_array_variables
+            kwargs['array_variables'] = handler.array_variables.split(',')
+        kwargs['track_extra_variables'] = handler.track_extra_variables.split(',')
+        for variable in handler.variables:
+            var_inv = VAR_DESCR_inv[variable]
+        eddies = cls(size=nb_obs, **kwargs)
+        for variable in handler.variables:
+            eddies.obs[VAR_DESCR_inv[variable]] = handler.variables[variable][:]
+        return eddies
+
     @staticmethod
     def cost_function2(records_in, records_out, distance):
         nb_records = records_in.shape[0]
         costs = ma.empty(nb_records,dtype='f4')
-        for i_record in xrange(nb_records):
+        for i_record in range(nb_records):
             poly_in = Polygon(
                 concatenate((
                     (records_in[i_record]['contour_lon_e'],),
@@ -647,6 +666,55 @@ def tracking(self, other):
 
         return i_self, i_other
 
+    def to_netcdf(self, handler):
+        eddy_size = len(self)
+        logging.debug('Create Dimensions "Nobs" : %d', eddy_size)
+        handler.createDimension('Nobs', eddy_size)
+        handler.track_extra_variables = ','.join(self.track_extra_variables)
+        if self.track_array_variables != 0:
+            handler.createDimension('NbSample', self.track_array_variables)
+            handler.track_array_variables = self.track_array_variables
+            handler.array_variables = ','.join(self.array_variables)
+        # Iter on variables to create:
+        for field in self.observations.dtype.descr:
+            name = field[0]
+            logging.debug('Create Variable %s', VAR_DESCR[name]['nc_name'])
+            self.create_variable(
+                handler,
+                dict(varname=VAR_DESCR[name]['nc_name'],
+                     datatype=VAR_DESCR[name]['output_type'],
+                     dimensions=VAR_DESCR[name]['nc_dims']),
+                VAR_DESCR[name]['nc_attr'],
+                self.observations[name],
+                scale_factor=VAR_DESCR[name].get('scale_factor', None),
+                add_offset=VAR_DESCR[name].get('add_offset', None)
+            )
+
+    @staticmethod
+    def create_variable(handler_nc, kwargs_variable, attr_variable,
+                        data, scale_factor=None, add_offset=None):
+        var = handler_nc.createVariable(
+            zlib=True,
+            complevel=1,
+            **kwargs_variable)
+        attrs = list(attr_variable.keys())
+        attrs.sort()
+        for attr in attrs:
+            attr_value = attr_variable[attr]
+            var.setncattr(attr, attr_value)
+        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())
+        except ValueError:
+            logging.warning('Data is empty')
+
 
 class VirtualEddiesObservations(EddiesObservations):
     """Class to work with virtual obs
@@ -680,6 +748,7 @@ def move_function(cls, obs_a, obs_b, out):
     @classmethod
     def forecast_move(cls, obs_a, obs_b, out):
         """Forecast move of an eddy
+        work to do
         """
         # New dead
         for key in obs_b.dtype.fields.keys():
@@ -709,6 +778,7 @@ def forecast_move(cls, obs_a, obs_b, out):
         # Count
         out['segment_size'][:] += 1
 
+
 class TrackEddiesObservations(EddiesObservations):
     """Class to practice Tracking on observations
     """
@@ -760,37 +830,13 @@ def elements(self):
         elements.extend(['track', 'n', 'virtual'])
         return elements
 
-    @staticmethod
-    def create_variable(handler_nc, kwargs_variable, attr_variable,
-                        data, scale_factor=None, add_offset=None):
-        var = handler_nc.createVariable(
-            zlib=True,
-            complevel=1,
-            **kwargs_variable)
-        attrs = list(attr_variable.keys())
-        attrs.sort()
-        for attr in attrs:
-            attr_value = attr_variable[attr]
-            var.setncattr(attr, attr_value)
-        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())
-        except ValueError:
-            logging.warning('Data is empty')
-
     def write_netcdf(self, path='./', filename='%(path)s/%(sign_type)s.nc'):
         """Write a netcdf with eddy obs
         """
         eddy_size = len(self.observations)
         sign_type = 'Cyclonic' if self.sign_type == -1 else 'Anticyclonic'
         filename = filename % dict(path=path, sign_type=sign_type)
+        logging.info('Store in %s', filename)
         with Dataset(filename, 'w', format='NETCDF4') as h_nc:
             logging.info('Create file %s', filename)
             # Create dimensions

src/py_eddy_tracker/tracking.py

@@ -27,6 +27,8 @@
 ===========================================================================
 
 """
+from matplotlib.dates import julian2num, num2date
+
 from py_eddy_tracker.observations import EddiesObservations, \
     VirtualEddiesObservations, TrackEddiesObservations
 from numpy import bool_, array, arange, ones, setdiff1d, zeros, uint16, \
@@ -47,7 +49,7 @@ class Correspondances(list):
     # Track limit to 65535
     N_DTYPE = 'u2'
 
-    def __init__(self, datasets, virtual=0, class_method=None):
+    def __init__(self, datasets, virtual=0, class_method=None, previous_correspondance=None):
         """Initiate tracking
         """
         super(Correspondances, self).__init__()
@@ -59,6 +61,7 @@ def __init__(self, datasets, virtual=0, class_method=None):
             self.class_method = EddiesObservations
         else:
             self.class_method = class_method
+
         # To count ID
         self.current_id = 0
         # To know the number maximal of link between two state
@@ -76,10 +79,17 @@ def __init__(self, datasets, virtual=0, class_method=None):
         self.virtual = virtual > 0
         self.virtual_obs = None
         self.previous_virtual_obs = None
+
+        # Correspondance to prolongate
+        self.filename_previous_correspondance = previous_correspondance
+        self.previous_correspondance = self.load_compatible(self.filename_previous_correspondance)
+
         if self.virtual:
             # Add field to dtype to follow virtual observations
             self.correspondance_dtype += [
+                # True if it isn't a real obs
                 ('virtual', bool_),
+                # Length of virtual segment
                 ('virtual_length', self.VIRTUAL_DTYPE)]
 
         # Array to simply merged
@@ -93,6 +103,17 @@ def reset_dataset_cache(self):
         self.previous_obs = None
         self.current_obs = None
 
+    @property
+    def period(self):
+        """To rethink
+
+        Returns: period coverage by obs
+
+        """
+        date_start = num2date(julian2num(self.class_method.load_from_netcdf(self.datasets[0]).obs['time'][0] - 0.5))
+        date_stop = num2date(julian2num(self.class_method.load_from_netcdf(self.datasets[-1]).obs['time'][0] - 0.5))
+        return date_start, date_stop
+
     def swap_dataset(self, dataset):
         """ Swap to next dataset
         """
@@ -138,6 +159,8 @@ def store_correspondance(self, i_previous, i_current, nb_real_obs):
         """
         # Create array to store correspondance data
         correspondance = array(i_previous, dtype=self.correspondance_dtype)
+        if self.virtual:
+            correspondance['virtual_length'][:] = 255
         # index from current_obs
         correspondance['out'] = i_current
 
@@ -279,14 +302,33 @@ def recense_dead_id_to_extend(self):
         # Count
         self.virtual_obs['segment_size'][:] += 1
 
+    def load_state(self):
+        # If we have a previous file of correspondance, we will replay only recent part
+        if self.previous_correspondance is not None:
+            first_dataset = len(self.previous_correspondance.datasets)
+            for correspondance in self.previous_correspondance[:first_dataset]:
+                self.append(correspondance)
+            self.current_obs = self.class_method.load_from_netcdf(self.datasets[first_dataset - 2])
+            flg_virtual = self.previous_correspondance.virtual
+            with Dataset(self.filename_previous_correspondance) as general_handler:
+                self.current_id = general_handler.last_current_id
+                # Load last virtual obs
+                self.virtual_obs = VirtualEddiesObservations.from_netcdf(general_handler.groups['LastVirtualObs'])
+                # Load and last previous virtual obs to be merge with current => will be previous2_obs
+                self.current_obs = self.current_obs.merge(
+                    VirtualEddiesObservations.from_netcdf(general_handler.groups['LastPreviousVirtualObs']))
+            return first_dataset, flg_virtual
+        return 1, False
+
     def track(self):
         """Run tracking
         """
-        flg_virtual = False
         self.reset_dataset_cache()
-        self.swap_dataset(self.datasets[0])
+        first_dataset, flg_virtual = self.load_state()
+
+        self.swap_dataset(self.datasets[first_dataset - 1])
         # We begin with second file, first one is in previous
-        for i, file_name in enumerate(self.datasets[1:]):
+        for file_name in self.datasets[first_dataset:]:
             self.swap_dataset(file_name)
             logging.debug('%s match with previous state', file_name)
             logging.debug('%d obs to match', len(self.current_obs))
@@ -295,13 +337,11 @@ def track(self):
             if flg_virtual:
                 logging.debug('%d virtual obs will be add to previous',
                               len(self.virtual_obs))
-                # If you comment this the virtual fonctionnality will be
-                # disable
                 self.previous_obs = self.previous_obs.merge(self.virtual_obs)
-
             i_previous, i_current = self.previous_obs.tracking(
                 self.current_obs)
 
+            # return true if the first time (previous2obs is none)
             if self.store_correspondance(i_previous, i_current, nb_real_obs):
                 continue
 
@@ -310,9 +350,11 @@ def track(self):
             if self.virtual:
                 flg_virtual = True
 
-    def save(self, filename):
+    def save(self, filename, dict_completion=None):
         self.prepare_merging()
         nb_step = len(self.datasets) - 1
+        if isinstance(dict_completion, dict):
+            filename = filename.format(**dict_completion)
         logging.info('Create correspondance file %s', filename)
         with Dataset(filename, 'w', format='NETCDF4') as h_nc:
             # Create dimensions
@@ -337,19 +379,46 @@ def save(self, filename):
 
             for name, dtype in self.correspondance_dtype:
                 if dtype is bool_:
-                    dtype = 'byte'
+                    dtype = 'u1'
+                kwargs_cv = dict()
+                if 'u1' in dtype:
+                    kwargs_cv['fill_value'] = 255,
                 h_nc.createVariable(zlib=True,
                                     complevel=1,
                                     varname=name,
                                     datatype=dtype,
-                                    dimensions=('Nstep', 'Nlink'))
+                                    dimensions=('Nstep', 'Nlink'),
+                                    **kwargs_cv
+                                    )
 
             for i, correspondance in enumerate(self):
                 nb_elt = correspondance.shape[0]
                 var_nb_link[i] = nb_elt
                 for name, _ in self.correspondance_dtype:
                     h_nc.variables[name][i, :nb_elt] = correspondance[name]
-            h_nc.virtual = int(self.virtual)
+            h_nc.virtual_use = str(self.virtual)
+            h_nc.virtual_max_segment = self.nb_virtual
+            h_nc.last_current_id = self.current_id
+            if self.virtual_obs is not None:
+                group = h_nc.createGroup('LastVirtualObs')
+                self.virtual_obs.to_netcdf(group)
+                group = h_nc.createGroup('LastPreviousVirtualObs')
+                self.previous_virtual_obs.to_netcdf(group)
+            h_nc.module = self.class_method.__module__
+            h_nc.classname = self.class_method.__qualname__
+
+    def load_compatible(self, filename):
+        if filename is None:
+            return None
+        previous_correspondance = Correspondances.load(filename)
+        if self.nb_virtual != previous_correspondance.nb_virtual:
+            raise Exception('File of correspondance IN contains a different virtual segment size : file(%d), yaml(%d)' %
+                            (previous_correspondance.nb_virtual, self.nb_virtual))
+
+        if self.class_method != previous_correspondance.class_method:
+            raise Exception('File of correspondance IN contains a different class method: file(%s), yaml(%s)' %
+                            (previous_correspondance.class_method, self.class_method))
+        return previous_correspondance
 
     @classmethod
     def load(cls, filename):
@@ -358,7 +427,11 @@ def load(cls, filename):
             datasets = list(h_nc.variables['FileIn'][:])
             datasets.append(h_nc.variables['FileOut'][-1])
 
-            obj = cls(datasets, h_nc.virtual)
+            if hasattr(h_nc, 'module'):
+                class_method= getattr(__import__(h_nc.module, globals(), locals(), h_nc.classname), h_nc.classname)
+            else:
+                class_method= None
+            obj = cls(datasets, h_nc.virtual_max_segment, class_method=class_method)
 
             id_max = 0
             for i, nb_elt in enumerate(h_nc.variables['nb_link'][:]):
@@ -371,7 +444,8 @@ def load(cls, filename):
                 for name, _ in obj.correspondance_dtype:
                     if name == 'in':
                         continue
-                    correspondance[name] = h_nc.variables[name][i, :nb_elt]
+                    if name == 'virtual_length':
+                        correspondance[name] = 255
                     correspondance[name] = h_nc.variables[name][i, :nb_elt]
                 id_max = max(id_max, correspondance['id'].max())
                 obj.append(correspondance)