Better code for zarr

Author: AntSimi

src/py_eddy_tracker/observations/observation.py

@@ -397,9 +397,17 @@ def zarr_dimension(filename):
         return set(dims)
 
     @classmethod
-    def load_from_zarr(cls, filename, remove_vars=None, include_vars=None):
+    def load_file(cls, filename, **kwargs):
+        if filename.endswith('.zarr'):
+            return cls.load_from_zarr(filename, **kwargs)
+        else:
+            return cls.load_from_netcdf(filename, **kwargs)
+
+    @classmethod
+    def load_from_zarr(cls, filename, raw_data=False, remove_vars=None, include_vars=None):
         # FIXME must be investigate, in zarr no dimensions name (or could be add in attr)
         array_dim = 50
+        BLOC = 5000000
         if not isinstance(filename, str):
             filename = filename.astype(str)
         h_zarr = zarr.open(filename)
@@ -428,13 +436,14 @@ def load_from_zarr(cls, filename, remove_vars=None, include_vars=None):
                 continue
             if var_inv not in cls.ELEMENTS and var_inv not in array_variables:
                 kwargs["track_extra_variables"].append(var_inv)
-        kwargs["raw_data"] = False
+        kwargs["raw_data"] = raw_data
         kwargs["only_variables"] = None if include_vars is None else [VAR_DESCR_inv[i] for i in include_vars]
         eddies = cls(size=nb_obs, **kwargs)
         for variable in var_list:
             var_inv = VAR_DESCR_inv[variable]
             if var_inv == "type_cyc":
                 continue
+            logging.debug('%s will be loaded', variable)
             # find unit factor
             factor = 1
             input_unit = h_zarr[variable].attrs.get('unit', None)
@@ -456,15 +465,22 @@ def load_from_zarr(cls, filename, remove_vars=None, include_vars=None):
                     if factor != 1:
                         logging.info('%s will be multiply by %f to take care of units(%s->%s)',
                                      variable, factor, input_unit, output_unit)
-            if factor != 1:
-                eddies.obs[var_inv] = h_zarr[variable][:] * factor
-            else:
-                eddies.obs[var_inv] = h_zarr[variable][:]
+            nb = h_zarr[variable].shape[0]
+
+            scale_factor = VAR_DESCR[var_inv].get('scale_factor', None)
+            add_offset = VAR_DESCR[var_inv].get('add_offset', None)
+            for i in range(0, nb, BLOC):
+                sl = slice(i, i + BLOC)
+                data = h_zarr[variable][sl]
+                if factor != 1:
+                    data *= factor
+                if raw_data:
+                    if add_offset is not None:
+                        data -= add_offset
+                    if scale_factor is not None:
+                        data /= scale_factor
+                eddies.obs[var_inv][sl] = data
 
-        # for variable in var_list:
-        #     var_inv = VAR_DESCR_inv[variable]
-        #     if var_inv == "type_cyc":
-        #         eddies.sign_type = h_zarr[variable][0]
         eddies.sign_type = h_zarr.attrs.get("rotation_type", 0)
         if eddies.sign_type == 0:
             logging.debug("File come from another algorithm of identification")

src/scripts/EddySubSetter

@@ -12,23 +12,29 @@ 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('-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',
-                        help='Extract path only if all obs are selected')
-    parser.add_argument('--reject_virtual', action='store_true',
-                        help="If there are only virtual observation in selection, we don't select track")
-    parser.add_argument('-a', '--area', nargs=4, type=float,
-                        metavar=('llcrnrlon', 'llcrnrlat', 'urcrnrlon', 'urcrnrlat'),
-                        help='Coordinates of bounding to extract'
-                        )
-    parser.add_argument('--remove_var', nargs='+', type=str, help='remove all listed variable')
-    parser.add_argument('--include_var', nargs='+', type=str, help='use only listed variable, remove_var will be ignored')
-    parser.add_argument('-i', '--ids', nargs='+', type=int, help='List of tracks which will be extract')
+
+    group = parser.add_argument_group('Extraction options')
+    group.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')
+    group.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')
+    group.add_argument('-f', '--full_path', action='store_true',
+                       help='Extract path, if one obs or more are selected')
+    group.add_argument('-d', '--remove_incomplete', action='store_true',
+                       help='Extract path only if all obs are selected')
+    group.add_argument('--reject_virtual', action='store_true',
+                       help="If there are only virtual observation in selection, we don't select track")
+    group.add_argument('-a', '--area', nargs=4, type=float,
+                       metavar=('llcrnrlon', 'llcrnrlat', 'urcrnrlon', 'urcrnrlat'),
+                       help='Coordinates of bounding to extract'
+                       )
+    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.add_argument('--sort_time', action='store_true', help='sort all observation with time')
+
     parser.add_argument('-n', '--no_raw_mode', action='store_true',
                         help='Uncompress all data, could be create a memory error for huge file, but is safer for extern file of py eddy tracker')
     return parser
@@ -38,7 +44,7 @@ if __name__ == '__main__':
     args = id_parser().parse_args()
 
     # Original dataset
-    dataset = TrackEddiesObservations.load_from_netcdf(
+    dataset = TrackEddiesObservations.load_file(
         args.filename,
         raw_data=False if args.no_raw_mode else True,
         remove_vars=args.remove_var,
@@ -70,6 +76,11 @@ if __name__ == '__main__':
                                             remove_incomplete=args.remove_incomplete,
                                               reject_virtual=args.reject_virtual)
 
+    if args.sort_time:
+        logging.debug('start sorting ...')
+        dataset.obs.sort(order=['time', 'lon', 'lat'])
+        logging.debug('end sorting')
+
     # if no data, no output will be written
     if len(dataset) == 0:
         logging.warning("No data are selected, out file couldn't be create")

src/scripts/EddyTracking

@@ -184,7 +184,7 @@ if __name__ == '__main__':
     NB_OBS_MIN = int(CONFIG.get('TRACK_DURATION_MIN', 14))
     CORRESPONDANCES.prepare_merging()
 
-    logging.info('The longest tracks have %d observations', CORRESPONDANCES.nb_obs_by_tracks.max())
+    logging.info('Longer track saved have %d obs', CORRESPONDANCES.nb_obs_by_tracks.max())
     logging.info('The mean length is %d observations before filtering', CORRESPONDANCES.nb_obs_by_tracks.mean())
 
     CORRESPONDANCES.get_unused_data().write_file(path=SAVE_DIR, filename='%(path)s/%(sign_type)s_untracked.nc', zarr_flag=ZARR)