Generic command to write file zarr or netcdf

Author: AntSimi

src/py_eddy_tracker/observations/observation.py

@@ -1113,18 +1113,25 @@ def create_variable_zarr(
         except ValueError:
             logging.warning("Data is empty")
 
-    def write_netcdf(self, path="./", filename="%(path)s/%(sign_type)s.nc"):
+    def write_file(self, path="./", filename="%(path)s/%(sign_type)s.nc", zarr_flag=False):
         """Write a netcdf with eddy obs
         """
-        eddy_size = len(self.observations)
         filename = filename % dict(
             path=path,
             sign_type=self.sign_legend,
             prod_time=datetime.now().strftime("%Y%m%d"),
         )
+        if zarr_flag:
+            filename = filename.replace('.nc', '.zarr')
+        if filename.endswith('.zarr'):
+            zarr_flag = True
         logging.info("Store in %s", filename)
-        with Dataset(filename, "w", format="NETCDF4") as handler:
-            self.to_netcdf(handler)
+        if zarr_flag:
+            handler = zarr.open(filename, 'w')
+            self.to_zarr(handler)
+        else:
+            with Dataset(filename, "w", format="NETCDF4") as handler:
+                self.to_netcdf(handler)
 
     @property
     def global_attr(self):

src/scripts/EddyFinalTracking

@@ -56,7 +56,7 @@ if __name__ == '__main__':
 
     if CONFIG.eddies_untracked_model is None:
         CONFIG.eddies_untracked_model = '%(path)s/%(sign_type)s_%(prod_time)s_untracked.nc'
-    CORRESPONDANCES.get_unused_data(raw_data=True).write_netcdf(path=SAVE_DIR, filename=CONFIG.eddies_untracked_model)
+    CORRESPONDANCES.get_unused_data(raw_data=True).write_file(path=SAVE_DIR, filename=CONFIG.eddies_untracked_model)
 
     SHORT_CORRESPONDANCES = CORRESPONDANCES._copy()
     SHORT_CORRESPONDANCES.shorter_than(size_max=NB_OBS_MIN)
@@ -84,5 +84,5 @@ if __name__ == '__main__':
         CONFIG.eddies_long_model = '%(path)s/%(sign_type)s_%(prod_time)s.nc'
     if CONFIG.eddies_short_model is None:
         CONFIG.eddies_short_model = '%(path)s/%(sign_type)s_%(prod_time)s_track_too_short.nc'
-    FINAL_EDDIES.write_netcdf(filename=CONFIG.eddies_long_model, path=SAVE_DIR)
-    SHORT_TRACK.write_netcdf(filename=CONFIG.eddies_short_model, path=SAVE_DIR)
+    FINAL_EDDIES.write_file(filename=CONFIG.eddies_long_model, path=SAVE_DIR)
+    SHORT_TRACK.write_file(filename=CONFIG.eddies_short_model, path=SAVE_DIR)

src/scripts/EddySubSetter

@@ -74,4 +74,4 @@ if __name__ == '__main__':
     if len(dataset) == 0:
         logging.warning("No data are selected, out file couldn't be create")
     else:
-        dataset.write_netcdf(filename=args.filename_out)
+        dataset.write_file(filename=args.filename_out)

src/scripts/EddyTracking

@@ -98,6 +98,9 @@ def usage():
                         action='store_true',
                         help='Stop tracking after correspondance computation,'
                              ' merging can be done with EddyFinalTracking')
+    parser.add_argument('--zarr',
+                        action='store_true',
+                        help='Output will be wrote in zarr')
     parser.add_argument('--blank_period',
                         type=int,
                         default=0,
@@ -109,11 +112,11 @@ def usage():
         config = yaml_load(stream)
     if args.correspondance_in is not None and not exists(args.correspondance_in):
         args.correspondance_in = None
-    return config, args.save_correspondance_and_stop, args.correspondance_in, args.correspondance_out, args.blank_period
+    return config, args.save_correspondance_and_stop, args.correspondance_in, args.correspondance_out, args.blank_period, args.zarr
 
 
 if __name__ == '__main__':
-    CONFIG, SAVE_STOP, CORRESPONDANCES_IN, CORRESPONDANCES_OUT, BLANK_PERIOD = usage()
+    CONFIG, SAVE_STOP, CORRESPONDANCES_IN, CORRESPONDANCES_OUT, BLANK_PERIOD, ZARR = usage()
 
     # Create output directory
     SAVE_DIR = CONFIG['PATHS'].get('SAVE_DIR', None)
@@ -184,7 +187,7 @@ if __name__ == '__main__':
     logging.info('The longest tracks have %d observations', 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_netcdf(path=SAVE_DIR, filename='%(path)s/%(sign_type)s_untracked.nc')
+    CORRESPONDANCES.get_unused_data().write_file(path=SAVE_DIR, filename='%(path)s/%(sign_type)s_untracked.nc', zarr_flag=ZARR)
 
     SHORT_CORRESPONDANCES = CORRESPONDANCES._copy()
     SHORT_CORRESPONDANCES.shorter_than(size_max=NB_OBS_MIN)
@@ -210,5 +213,5 @@ if __name__ == '__main__':
     logging.info('Longer track saved have %d obs', CORRESPONDANCES.nb_obs_by_tracks.max())
     logging.info('The mean length is %d observations after filtering', CORRESPONDANCES.nb_obs_by_tracks.mean())
 
-    FINAL_EDDIES.write_netcdf(path=SAVE_DIR)
-    SHORT_TRACK.write_netcdf(filename='%(path)s/%(sign_type)s_track_too_short.nc', path=SAVE_DIR)
+    FINAL_EDDIES.write_file(path=SAVE_DIR, zarr_flag=ZARR)
+    SHORT_TRACK.write_file(filename='%(path)s/%(sign_type)s_track_too_short.nc', path=SAVE_DIR, zarr_flag=ZARR)