Add feature to extract track on other criterions

Author: AntSimi

src/py_eddy_tracker/observations/tracking.py

@@ -177,9 +177,12 @@ def get_mask_from_id(self, tracks):
 
     def compute_index(self):
         if self.__first_index_of_track is None:
-            s = self.tracks.max()
-            self.__first_index_of_track = -ones(s, self.tracks.dtype)
-            self.__obs_by_track = zeros(s, self.observation_number.dtype)
+            s = self.tracks.max() + 1
+            # Doesn't work => core dump with numba, maybe he wait i8 instead of u4
+            # self.__first_index_of_track = -ones(s, self.tracks.dtype)
+            # self.__obs_by_track = zeros(s, self.observation_number.dtype)
+            self.__first_index_of_track = -ones(s, 'i8')
+            self.__obs_by_track = zeros(s, 'i8')
             logging.debug('Start computing index ...')
             compute_index(self.tracks, self.__first_index_of_track, self.__obs_by_track)
             logging.debug('... OK')
@@ -198,6 +201,19 @@ def extract_ids(self, tracks):
         mask = self.get_mask_from_id(array(tracks))
         return self.__extract_with_mask(mask)
 
+    def extract_with_length(self, bounds):
+        b0, b1 = bounds
+        if b0 >= 0 and b1 >=0:
+            track_mask = (self.nb_obs_by_track >= b0) * (self.nb_obs_by_track <= b1)
+        elif b0 < 0 and b1 >= 0:
+            track_mask = self.nb_obs_by_track <= b1
+        elif b0 >= 0 and b1 < 0:
+            track_mask = self.nb_obs_by_track > b0
+        else:
+            logging.warning('No valid value for bounds')
+            raise Exception('One bounds must be positiv')
+        return self.__extract_with_mask(track_mask.repeat(self.nb_obs_by_track))
+
     def __extract_with_mask(self, mask, full_path=False, remove_incomplete=False, compress_id=False):
         """
         Extract a subset of observations

src/scripts/EddySubSetter

@@ -5,7 +5,7 @@ Subset eddy Dataset
 """
 from py_eddy_tracker import EddyParser
 from py_eddy_tracker.observations.tracking import TrackEddiesObservations
-from os.path import basename, dirname
+import logging
 
 
 def id_parser():
@@ -14,6 +14,8 @@ def id_parser():
     parser.add_argument('filename_out')
     parser.add_argument('-p', '--period', nargs=2, type=int,
                         help='Start day and end day, if it s negative value we will add to day min and add to day max, if 0 it s not use')
+    parser.add_argument('-l', '--length', nargs=2, type=int,
+                        help='Minimal and maximal quantity of observation for one track, ones bounds could be negative, it will be not use')
     parser.add_argument('-f', '--full_path', action='store_true',
                         help='Extract path, if one obs or more are selected')
     parser.add_argument('-d', '--remove_incomplete', action='store_true',
@@ -34,12 +36,20 @@ if __name__ == '__main__':
     # Original dataset
     dataset = TrackEddiesObservations.load_from_netcdf(args.filename, raw_data=False if args.no_raw_mode else True)
 
+    # Select with id
     if args.ids is not None:
         dataset = dataset.extract_ids(args.ids)
+
+    # Select with length
+    if args.length is not None:
+        dataset = dataset.extract_with_length(args.length)
+
+    # Select with a start date and end date
     if args.period is not None:
         dataset = dataset.extract_with_period(args.period, full_path=args.full_path,
                                               remove_incomplete=args.remove_incomplete)
 
+    # Select track which go through an area
     if args.area is not None:
         area = dict(llcrnrlon=args.area[0],
                     llcrnrlat=args.area[1],
@@ -49,5 +59,8 @@ if __name__ == '__main__':
         dataset = dataset.extract_with_area(area, full_path=args.full_path,
                                             remove_incomplete=args.remove_incomplete)
 
-    if len(dataset) != 0:
+    # if no data, no output will be written
+    if len(dataset) == 0:
+        logging.warning("No data are selected, out file couldn't be create")
+    else:
         dataset.write_netcdf(filename=args.filename_out)