feat: read and write zarr methods to replace read_pickle

docs/_code/prepare_hef.py

@@ -46,7 +46,7 @@
 df = workflow.calibrate_inversion_from_consensus(gdir, volume_m3_reference=refv)
 np.testing.assert_allclose(refv, df.vol_oggm_m3, rtol=0.01)
 
-cl = gdir.read_pickle('inversion_input')[-1]
+cl = gdir.read_store('inversion_input')[-1]
 mbmod = massbalance.ConstantMassBalance(gdir, y0=1985)
 mbx = mbmod.get_annual_mb(cl['hgt']) * cfg.SEC_IN_YEAR * cfg.PARAMS['ice_density']
 fdf = pd.DataFrame(index=np.arange(len(mbx))*cl['dx'])

docs/api.rst

@@ -476,7 +476,7 @@ level 5.
                                               prepro_border=80)
     gdir = gdirs[0]  # init_glacier_directories always returns a list
 
-    fls = gdir.read_pickle('model_flowlines')
+    fls = gdir.read_store('model_flowlines')
     fls
     [fl.order for fl in fls]
 

oggm/cfg.py

@@ -300,6 +300,10 @@ def _log_param_change(self, key, value):
 _doc = "The outlines of this glacier complex sub-entities (for RGI7C only!)."
 BASENAMES['complex_sub_entities'] = ('complex_sub_entities.shp', _doc)
 
+_doc = "A zarr store containing previous pickled data."
+BASENAMES['data_store'] = ('data_store.zarr', _doc)
+
+
 
 def set_logging_config(logging_level='INFO'):
     """Set the global logger parameters.

oggm/core/centerlines.py

@@ -905,7 +905,7 @@ def compute_centerlines(gdir, heads=None):
         raise InvalidParamsError('`force_one_flowline` is deprecated')
 
     # open
-    geom = gdir.read_pickle('geometries')
+    geom = gdir.read_store('geometries')
     grids_file = gdir.get_filepath('gridded_data')
     with utils.ncDataset(grids_file) as nc:
         # Variables
@@ -974,7 +974,7 @@ def compute_centerlines(gdir, heads=None):
                             'found!'.format(gdir.rgi_id))
 
     # Write the data
-    gdir.write_pickle(cls, 'centerlines')
+    gdir.write_store(cls, 'centerlines')
 
     if is_first_call:
         # For diagnostics of filtered centerlines
@@ -1009,7 +1009,7 @@ def compute_downstream_line(gdir):
     # Look for the starting points
     try:
         # Normal OGGM flowlines
-        p = gdir.read_pickle('centerlines')[-1].tail
+        p = gdir.read_store('centerlines')[-1].tail
         head = (int(p.y), int(p.x))
     except FileNotFoundError:
         # Squeezes lines
@@ -1050,7 +1050,7 @@ def compute_downstream_line(gdir):
     if line is None:
         raise GeometryError('Downstream line not found')
 
-    cl = gdir.read_pickle('inversion_flowlines')[-1]
+    cl = gdir.read_store('inversion_flowlines')[-1]
     if cl.line is not None:
         # normal OGGM lines
         lline, dline = _line_extend(cl.line, line, cl.dx)
@@ -1060,7 +1060,7 @@ def compute_downstream_line(gdir):
         _, dline = _line_extend(shpg.LineString(), line, cl.dx)
         out = dict(full_line=None, downstream_line=dline)
 
-    gdir.write_pickle(out, 'downstream_line')
+    gdir.write_store(out, 'downstream_line')
 
 
 def _approx_parabola(x, y, y0=0):
@@ -1282,8 +1282,8 @@ def compute_downstream_bedshape(gdir):
         return
 
     # We make a flowline out of the downstream for simplicity
-    tpl = gdir.read_pickle('inversion_flowlines')[-1]
-    cl = gdir.read_pickle('downstream_line')['downstream_line']
+    tpl = gdir.read_store('inversion_flowlines')[-1]
+    cl = gdir.read_store('downstream_line')['downstream_line']
     cl = Centerline(cl, dx=tpl.dx, map_dx=gdir.grid.dx)
 
     # Topography
@@ -1316,11 +1316,11 @@ def compute_downstream_bedshape(gdir):
     assert np.all(w0s >= w0_min), 'np.all(w0s >= w0_min)'
 
     # write output
-    out = gdir.read_pickle('downstream_line')
+    out = gdir.read_store('downstream_line')
     out['bedshapes'] = bs
     out['surface_h'] = hgts
     out['w0s'] = w0s
-    gdir.write_pickle(out, 'downstream_line')
+    gdir.write_store(out, 'downstream_line')
 
 
 def _mask_to_polygon(mask, gdir=None):
@@ -1562,8 +1562,8 @@ def catchment_area(gdir):
     """
 
     # Variables
-    cls = gdir.read_pickle('centerlines')
-    geom = gdir.read_pickle('geometries')
+    cls = gdir.read_store('centerlines')
+    geom = gdir.read_store('geometries')
     glacier_pix = geom['polygon_pix']
     fpath = gdir.get_filepath('gridded_data')
     with utils.ncDataset(fpath) as nc:
@@ -1676,7 +1676,7 @@ def catchment_intersections(gdir):
         where to write the data
     """
 
-    catchment_indices = gdir.read_pickle('geometries')['catchment_indices']
+    catchment_indices = gdir.read_store('geometries')['catchment_indices']
 
     # Loop over the lines
     mask = np.zeros((gdir.grid.ny, gdir.grid.nx))
@@ -1722,7 +1722,7 @@ def initialize_flowlines(gdir):
     """
 
     # variables
-    cls = gdir.read_pickle('centerlines')
+    cls = gdir.read_store('centerlines')
 
     # Initialise the flowlines
     dx = cfg.PARAMS['flowline_dx']
@@ -1811,7 +1811,7 @@ def initialize_flowlines(gdir):
         fl.set_flows_to(fls[cls.index(cl.flows_to)])
 
     # Write the data
-    gdir.write_pickle(fls, 'inversion_flowlines')
+    gdir.write_store(fls, 'inversion_flowlines')
     gdir.add_to_diagnostics('flowline_type', 'centerlines')
     if do_filter:
         out = diag_n_bad_slopes/diag_n_pix
@@ -1831,8 +1831,8 @@ def catchment_width_geom(gdir):
     """
 
     # variables
-    flowlines = gdir.read_pickle('inversion_flowlines')
-    catchment_indices = gdir.read_pickle('geometries')['catchment_indices']
+    flowlines = gdir.read_store('inversion_flowlines')
+    catchment_indices = gdir.read_store('geometries')['catchment_indices']
 
     # Topography is to filter the unrealistic lines afterwards.
     # I take the non-smoothed topography
@@ -1928,8 +1928,8 @@ def catchment_width_geom(gdir):
         fl.geometrical_widths = wlines
         fl.is_rectangular = is_rectangular
 
-    # Overwrite pickle
-    gdir.write_pickle(flowlines, 'inversion_flowlines')
+    # Overwrite
+    gdir.write_store(flowlines, 'inversion_flowlines')
 
 
 @entity_task(log, writes=['inversion_flowlines'])
@@ -1949,8 +1949,8 @@ def catchment_width_correction(gdir):
     """
 
     # variables
-    fls = gdir.read_pickle('inversion_flowlines')
-    catchment_indices = gdir.read_pickle('geometries')['catchment_indices']
+    fls = gdir.read_store('inversion_flowlines')
+    catchment_indices = gdir.read_store('geometries')['catchment_indices']
 
     # Topography for altitude-area distribution
     # I take the non-smoothed topography and remove the borders
@@ -2077,8 +2077,8 @@ def catchment_width_correction(gdir):
     for fl in fls:
         fl.widths *= fac
 
-    # Overwrite centerlines
-    gdir.write_pickle(fls, 'inversion_flowlines')
+    # Overwrite
+    gdir.write_store(fls, 'inversion_flowlines')
 
 
 @entity_task(log, writes=['inversion_flowlines'])
@@ -2099,7 +2099,7 @@ def terminus_width_correction(gdir, new_width=None):
     """
 
     # variables
-    fls = gdir.read_pickle('inversion_flowlines')
+    fls = gdir.read_store('inversion_flowlines')
     fl = fls[-1]
     mapdx = gdir.grid.dx
 
@@ -2120,8 +2120,8 @@ def terminus_width_correction(gdir, new_width=None):
     width[:-5] = width[:-5] * cor_factor
     fl.widths = width
 
-    # Overwrite centerlines
-    gdir.write_pickle(fls, 'inversion_flowlines')
+    # Overwrite
+    gdir.write_store(fls, 'inversion_flowlines')
 
 
 def intersect_downstream_lines(gdir, candidates=None):
@@ -2152,7 +2152,7 @@ def intersect_downstream_lines(gdir, candidates=None):
     buffer = cfg.PARAMS['kbuffer']
 
     # get main glacier downstream line and CRS
-    dline = gdir.read_pickle('downstream_line')['full_line']
+    dline = gdir.read_store('downstream_line')['full_line']
     crs = gdir.grid
 
     # return list
@@ -2165,7 +2165,7 @@ def intersect_downstream_lines(gdir, candidates=None):
             continue
 
         # get tributary glacier downstream line and CRS
-        _dline = trib.read_pickle('downstream_line')['full_line']
+        _dline = trib.read_store('downstream_line')['full_line']
         _crs = trib.grid
 
         # use salem to transform the grids
@@ -2449,5 +2449,5 @@ def fixed_dx_elevation_band_flowline(gdir, bin_variables=None,
         fl.is_rectangular[-5:] = True
         fl.is_trapezoid[-5:] = False
 
-    gdir.write_pickle([fl], 'inversion_flowlines')
+    gdir.write_store([fl], 'inversion_flowlines')
     gdir.add_to_diagnostics('flowline_type', 'elevation_band')

oggm/core/dynamic_spinup.py

@@ -259,7 +259,7 @@ def run_dynamic_spinup(gdir, init_model_filesuffix=None, init_model_yr=None,
         init_model_fls = fmod.fls
 
     if init_model_fls is None:
-        fls_spinup = gdir.read_pickle('model_flowlines',
+        fls_spinup = gdir.read_store('model_flowlines',
                                       filesuffix=model_flowline_filesuffix)
     else:
         fls_spinup = copy.deepcopy(init_model_fls)
@@ -1118,7 +1118,7 @@ def dynamic_melt_f_run_with_dynamic_spinup(
         # ATTENTION: it is assumed that the flowlines in gdir have the volume
         # we want to match during calibrate_inversion_from_consensus when we
         # set_local_variables
-        fls_ref = gdir.read_pickle('model_flowlines')
+        fls_ref = gdir.read_store('model_flowlines')
         local_variables['vol_m3_ref'] = np.sum([f.volume_m3 for f in fls_ref])
 
         # we are done with preparing the local_variables for the upcoming iterations
@@ -1150,7 +1150,7 @@ def dynamic_melt_f_run_with_dynamic_spinup(
                        np.all(getattr(fl_prov, 'bed_h') ==
                               getattr(fl_orig, 'bed_h'))
                        for fl_prov, fl_orig in
-                       zip(fls_init, gdir.read_pickle('model_flowlines'))]):
+                       zip(fls_init, gdir.read_store('model_flowlines'))]):
             raise InvalidWorkflowError('If you want to perform a dynamic '
                                        'melt_f calibration including an '
                                        'inversion, it is not possible to '
@@ -1869,7 +1869,7 @@ def run_dynamic_melt_f_calibration(
         init_model_fls = fmod.fls
 
     if init_model_fls is None:
-        fls_init = gdir.read_pickle('model_flowlines')
+        fls_init = gdir.read_store('model_flowlines')
     else:
         fls_init = copy.deepcopy(init_model_fls)