Add function to subset atlas

Author: AntSimi

setup.py

@@ -16,6 +16,7 @@
     scripts=[
         'src/scripts/GridFiltering',
         'src/scripts/EddyId',
+        'src/scripts/EddySubSetter',
         'src/scripts/EddyTracking',
         'src/scripts/EddyFinalTracking',
         'src/scripts/EddyMergeCorrespondances',

src/py_eddy_tracker/observations/observation.py

@@ -174,6 +174,18 @@ def __init__(
         self.active = True
         self.sign_type = None
 
+    @property
+    def time(self):
+        return self.observations["time"]
+
+    @property
+    def tracks(self):
+        return self.observations["track"]
+
+    @property
+    def observation_number(self):
+        return self.observations["n"]
+
     @property
     def sign_legend(self):
         return "Cyclonic" if self.sign_type == -1 else "Anticyclonic"

src/py_eddy_tracker/observations/tracking.py

@@ -28,21 +28,27 @@
 
 """
 from numpy import empty, arange, where, unique, \
-    interp
+    interp, ones, bool_, zeros
 from .. import VAR_DESCR_inv
 import logging
 from datetime import datetime, timedelta
 from .observation import EddiesObservations
+from numba import njit
 
 
 class TrackEddiesObservations(EddiesObservations):
     """Class to practice Tracking on observations
     """
-    __slots__ = ()
+    __slots__ = ('__obs_by_track', '__first_index_of_track')
 
-    ELEMENTS = ['lon', 'lat', 'radius_s', 'radius_e', 'amplitude', 'speed_radius', 'time',
-                'shape_error_e', 'shape_error_s', 'nb_contour_selected',
-                'height_max_speed_contour', 'height_external_contour', 'height_inner_contour', 'cost_association']
+    ELEMENTS = ['lon', 'lat', 'radius_s', 'radius_e', 'amplitude', 'speed_radius', 'time', 'shape_error_e',
+                'shape_error_s', 'nb_contour_selected', 'height_max_speed_contour', 'height_external_contour',
+                'height_inner_contour', 'cost_association']
+
+    def __init__(self, *args, **kwargs):
+        super(TrackEddiesObservations, self).__init__(*args, **kwargs)
+        self.__first_index_of_track = None
+        self.__obs_by_track = None
 
     def filled_by_interpolation(self, mask):
         """Filled selected values by interpolation
@@ -97,7 +103,7 @@ def extract_longer_eddies(self, nb_min, nb_obs, compress_id=True):
     def elements(self):
         elements = super(TrackEddiesObservations, self).elements
         elements.extend(['track', 'n', 'virtual'])
-        return elements
+        return list(set(elements))
 
     def set_global_attr_netcdf(self, h_nc):
         """Set global attr
@@ -126,19 +132,110 @@ def extract_with_period(self, period, **kwargs):
         Returns:
             same object with selected data
         """
-        mask = (self.time > period[0]) * (self.time < period[1])
+        dataset_period = self.period
+        p_min, p_max = period
+        if p_min > 0:
+            mask = self.time >= p_min
+        elif p_min < 0:
+            mask = self.time >= (dataset_period[0] - p_min)
+        else:
+            mask = ones(self.time.shape, dtype=bool_)
+        if p_max > 0:
+            mask *= self.time <= p_max
+        elif p_max < 0:
+            mask *= self.time <= (dataset_period[1] + p_max)
         return self.__extract_with_mask(mask, **kwargs)
 
-    def __extract_with_mask(self, mask, full_path=False, remove_incomplete=False):
+    @property
+    def period(self):
+        """
+        Give time coverage
+        Returns: 2 date
+        """
+        return self.time.min(), self.time.max()
+
+    def get_mask_from_id(self, tracks):
+        mask = zeros(self.tracks.shape, dtype=bool_)
+        compute_mask_from_id(tracks, self.index_from_track, self.nb_obs_by_track, mask)
+        return mask
+
+    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)
+            logging.debug('Start computing index ...')
+            compute_index(self.tracks, self.__first_index_of_track, self.__obs_by_track)
+            logging.debug('... OK')
+
+    @property
+    def index_from_track(self):
+        self.compute_index()
+        return self.__first_index_of_track
+
+    @property
+    def nb_obs_by_track(self):
+        self.compute_index()
+        return self.__obs_by_track
+
+    def __extract_with_mask(self, mask, full_path=False, remove_incomplete=False, compress_id=False):
         """
         Extract a subset of observations
         Args:
             mask: mask to select observations
             full_path: extract full path if only one part is selected
-            remove_incomplete: delete path which are not totatly selected
+            remove_incomplete: delete path which are not fully selected
+            compress_id: resample track number to use a little range
 
         Returns:
             same object with selected observations
         """
         if full_path and remove_incomplete:
             logging.warning('Incompatible option, remove_incomplete option will be remove')
+            remove_incomplete = False
+
+        if full_path:
+            tracks = unique(self.tracks[mask])
+            mask = self.get_mask_from_id(tracks)
+        elif remove_incomplete:
+            raise Exception("")
+
+        nb_obs = mask.sum()
+        new = TrackEddiesObservations(
+            size=nb_obs,
+            track_extra_variables=self.track_extra_variables,
+            track_array_variables=self.track_array_variables,
+            array_variables=self.array_variables,
+            raw_data=self.raw_data
+        )
+        new.sign_type = self.sign_type
+        if nb_obs == 0:
+            logging.warning('Empty dataset will be created')
+        else:
+            for field in self.obs.dtype.descr:
+                logging.debug('Copy of field %s ...', field)
+                var = field[0]
+                new.obs[var] = self.obs[var][mask]
+            if compress_id:
+                list_id = unique(new.obs['track'])
+                list_id.sort()
+                id_translate = arange(list_id.max() + 1)
+                id_translate[list_id] = arange(len(list_id)) + 1
+                new.obs['track'] = id_translate[new.obs['track']]
+        return new
+
+
+@njit(cache=True)
+def compute_index(tracks, index, number):
+    previous_track = -1
+    for i, track in enumerate(tracks):
+        if track != previous_track:
+            index[track] = i
+        number[track] += 1
+        previous_track = track
+
+
+@njit(cache=True)
+def compute_mask_from_id(tracks, first_index, number_of_obs, mask):
+    for track in tracks:
+        mask[first_index[track]:first_index[track] + number_of_obs[track]] = True

src/scripts/EddySubSetter

@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Subset eddy Dataset
+"""
+from py_eddy_tracker import EddyParser
+from py_eddy_tracker.observations.tracking import TrackEddiesObservations
+from os.path import basename, dirname
+
+
+def id_parser():
+    parser = EddyParser('Eddy Identification')
+    parser.add_argument('filename')
+    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('-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',
+                        help='Extract path only if all obs are selected')
+    return parser
+
+
+if __name__ == '__main__':
+    args = id_parser().parse_args()
+
+    dataset = TrackEddiesObservations.load_from_netcdf(args.filename, raw_data=True)
+    if args.period is not None:
+        dataset = dataset.extract_with_period(args.period, full_path=args.full_path,
+                                              remove_incomplete=args.remove_incomplete)
+
+    if len(dataset) != 0:
+        dataset.write_netcdf(filename=args.filename_out)