Add new method of extraction

Author: AntSimi

src/py_eddy_tracker/dataset/grid.py

@@ -2,7 +2,7 @@
 """
 """
 import logging
-from numpy import concatenate, int32, empty, where, array, \
+from numpy import concatenate, empty, where, array, \
     sin, deg2rad, pi, ones, cos, ma, int8, histogram2d, arange, float_, \
     linspace, errstate, int_, interp, meshgrid, nan, ceil, sinc, isnan, \
     percentile, zeros, arctan2, arcsin, round_, nanmean, exp, mean as np_mean
@@ -1082,8 +1082,7 @@ def convolve_filter_with_dynamic_kernel(self, grid, kernel_func, lat_max=85, ext
             if debug_active and len(dt) > 0:
                 dt_mean = np_mean(dt) * (nb_lines - i)
                 print('Remain ', dt_mean, 'ETA ', t0 + dt_mean, 'current kernel size :', k_shape, 'Step : %d/%d    ' % (i, nb_lines), end="\r")
-            
-            
+
             # Half size, k_shape must be always impair
             d_lat = int((k_shape[1] - 1) / 2)
             d_lon = int((k_shape[0] - 1) / 2)
@@ -1119,7 +1118,8 @@ def convolve_filter_with_dynamic_kernel(self, grid, kernel_func, lat_max=85, ext
             out = ma.array(data_out, mask=data_out.mask)
         else:
             out = ma.array(data_out, mask=data.mask + data_out.mask)
-        print()
+        if debug_active:
+            print()
         if out.dtype != data.dtype:
             return out.astype(data.dtype)
         return out

src/py_eddy_tracker/observations/tracking.py

@@ -223,6 +223,33 @@ def extract_ids(self, tracks):
         mask = self.get_mask_from_id(array(tracks))
         return self.__extract_with_mask(mask)
 
+    def extract_first_obs_in_box(self, res):
+        data = empty(self.obs.shape, dtype=[('lon', 'f4'), ('lat', 'f4'), ('track', 'i4')])
+        data['lon'] = self.longitude - self.longitude % res
+        data['lat'] = self.latitude - self.latitude % res
+        data['track'] = self.obs["track"]
+        _, indexs = unique(data, return_index=True)
+        mask = zeros(self.obs.shape, dtype='bool')
+        mask[indexs] = True
+        return self.__extract_with_mask(mask)
+
+    def extract_in_direction(self, direction, value=0):
+        nb_obs = self.nb_obs_by_track
+        i_start = self.index_from_track
+        i_stop = i_start + nb_obs - 1
+        if direction in ('S', 'N'):
+            d_lat = self.latitude[i_stop] - self.latitude[i_start]
+            mask = d_lat < 0 if 'S' == direction else d_lat > 0
+            mask &= abs(d_lat) > value
+        else:
+            lon_start , lon_end = self.longitude[i_start], self.longitude[i_stop]
+            lon_end = (lon_end - (lon_start - 180)) % 360 + lon_start - 180
+            d_lon = lon_end - lon_start
+            mask = d_lon < 0 if 'W' == direction else d_lon > 0
+            mask &= abs(d_lon) > value
+        mask = mask.repeat(nb_obs)
+        return self.__extract_with_mask(mask)
+
     def extract_with_length(self, bounds):
         b0, b1 = bounds
         if b0 >= 0 and b1 >= 0:

src/py_eddy_tracker/poly.py

@@ -81,7 +81,7 @@ def winding_number_grid_in_poly(x_1d, y_1d, i_x0, i_x1, x_size, i_y0, xy_poly):
     # the  winding number counter
     nb_x, nb_y = len(x_1d), len(y_1d)
     wn = empty((nb_x, nb_y), dtype=numba_types.bool_)
-    for i in range(nb_x):
+    for i in prange(nb_x):
         x_pt = x_1d[i]
         for j in range(nb_y):
             y_pt = y_1d[j]

src/scripts/EddySubSetter

@@ -11,7 +11,7 @@ logger = logging.getLogger("pet")
 
 
 def id_parser():
-    parser = EddyParser('Eddy Identification')
+    parser = EddyParser('Eddy Subsetter')
     parser.add_argument('filename')
     parser.add_argument('filename_out')
 
@@ -30,11 +30,17 @@ def id_parser():
                        metavar=('llcrnrlon', 'llcrnrlat', 'urcrnrlon', 'urcrnrlat'),
                        help='Coordinates of bounding to extract'
                        )
+    group.add_argument('--direction', choices=['E', 'W', 'S', 'N'],
+                       help='Select only track which have an end point which go in this direction')
+    group.add_argument('--minimal_degrees_displacement_in_direction', type=float,
+                       help='Minimal displacement in direction specified in --directio options')
+    group.add_argument('--select_first_observation_in_box', type=float,
+                       help='Select only the first obs in each box for each tracks, value specified must be resolution')
     group.add_argument('--remove_var', nargs='+', type=str, help='remove all listed variable')
     group.add_argument('--include_var', nargs='+', type=str, help='use only listed variable, remove_var will be ignored')
     group.add_argument('-i', '--ids', nargs='+', type=int, help='List of tracks which will be extract')
 
-    group = parser.add_argument_group('Extraction options')
+    group = parser.add_argument_group('General options')
     group.add_argument('--sort_time', action='store_true', help='sort all observation with time')
 
     parser.add_argument('-n', '--no_raw_mode', action='store_true',
@@ -78,6 +84,18 @@ if __name__ == '__main__':
                                             remove_incomplete=args.remove_incomplete,
                                               reject_virtual=args.reject_virtual)
 
+    # Select only track which go in the direction specified
+    if args.direction:
+        if args.minimal_degrees_displacement_in_direction:
+            dataset = dataset.extract_in_direction(
+                args.direction,
+                value=args.minimal_degrees_displacement_in_direction)
+        else:
+            dataset = dataset.extract_in_direction(args.direction)
+
+    if args.select_first_observation_in_box:
+        dataset = dataset.extract_first_obs_in_box(res=args.select_first_observation_in_box)
+
     if args.sort_time:
         logger.debug('start sorting ...')
         dataset.obs.sort(order=['time', 'lon', 'lat'])