removing zarr dataset changes which were causing replication error

Author: Kit Schwarz

dscim/menu/simple_storage.py

@@ -81,11 +81,7 @@ def gmst(self):
     @property
     def gmsl(self):
         """Cached GMSL anomalies"""
-        gmsl = (
-            xr.open_dataset(self.gmsl_path, engine="zarr")
-            .gmsl.to_dataframe()
-            .reset_index()
-        )
+        gmsl = xr.open_zarr(self.gmsl_path).gmsl.to_dataframe().reset_index()
 
         return gmsl
 
@@ -120,7 +116,7 @@ def gmsl_anomalies(self):
         """This function takes coastal sector's GMSL relative to 1991-2009.
         No rebasing occurs, as coastal damages are rebased to the same period.
         """
-        df = xr.open_mfdataset(self.gmsl_fair_path, engine="zarr")
+        df = xr.open_zarr(self.gmsl_fair_path)
         df = df.chunk(df.dims)
 
         datasets = []
@@ -277,7 +273,7 @@ def __init__(self, path_econ):
     def econ_vars(self):
         """Economic variables"""
         if self.path[-3:] == "arr":
-            raw = xr.open_dataset(self.path, engine="zarr", consolidated=True)
+            raw = xr.open_zarr(self.path, consolidated=True)
         else:
             raw = xr.open_dataset(self.path)
         return raw[["gdp", "pop"]]
@@ -404,11 +400,7 @@ def adding_up_damages(self):
                 f"Adding up aggregated damages found at {mean_cc}, {mean_no_cc}. These are being loaded..."
             )
             damages = (
-                (
-                    xr.open_dataset(mean_no_cc, engine="zarr").no_cc
-                    - xr.open_dataset(mean_cc, engine="zarr").cc
-                )
-                * self.pop
+                (xr.open_zarr(mean_no_cc).no_cc - xr.open_zarr(mean_cc).cc) * self.pop
             ).sum("region")
         else:
             raise NotImplementedError(
@@ -437,4 +429,4 @@ def risk_aversion_damages(self, ce_type):
             raise NotImplementedError(
                 "Risk-aversion CEs not found. Please run CE_calculation.ipynb for `risk_aversion`."
             )
-        return self.cut(xr.open_dataset(file, engine="zarr"))
+        return self.cut(xr.open_zarr(file))

dscim/preprocessing/preprocessing.py

@@ -35,7 +35,7 @@ def ce_from_chunk(
     model = chunk.model.values
 
     gdppc = (
-        xr.open_dataset(socioec, engine="zarr", chunks=None)
+        xr.open_zarr(socioec, chunks=None)
         .sel(
             year=year, ssp=ssp, model=model, region=ce_batch_coords["region"], drop=True
         )
@@ -97,11 +97,11 @@ def reduce_damages(
     delta = params["delta"]
     outpath = f"{c['paths']['reduced_damages_library']}/{sector}"
 
-    with xr.open_dataset(damages, engine="zarr", chunks=None)[histclim] as ds:
-        with xr.open_dataset(socioec, engine="zarr", chunks=None) as gdppc:
+    with xr.open_zarr(damages, chunks=None)[histclim] as ds:
+        with xr.open_zarr(socioec, chunks=None) as gdppc:
 
             assert (
-                xr.open_dataset(damages, engine="zarr").chunks["batch"][0] == 15
+                xr.open_zarr(damages).chunks["batch"][0] == 15
             ), "'batch' dim on damages does not have chunksize of 15. Please rechunk."
 
             ce_batch_dims = [i for i in gdppc.dims] + [
@@ -112,18 +112,15 @@ def reduce_damages(
                 i for i in gdppc.region.values if i in ce_batch_coords["region"]
             ]
             ce_shapes = [len(ce_batch_coords[c]) for c in ce_batch_dims]
-            ce_chunks = [
-                xr.open_dataset(damages, engine="zarr").chunks[c][0]
-                for c in ce_batch_dims
-            ]
+            ce_chunks = [xr.open_zarr(damages).chunks[c][0] for c in ce_batch_dims]
 
     template = xr.DataArray(
         da.empty(ce_shapes, chunks=ce_chunks),
         dims=ce_batch_dims,
         coords=ce_batch_coords,
     )
 
-    other = xr.open_dataset(damages, engine="zarr")
+    other = xr.open_zarr(damages)
 
     out = other.map_blocks(
         ce_from_chunk,
@@ -212,9 +209,7 @@ def sum_AMEL(
 
         for sector in sectors:
             print(f"Opening {sector},{params[sector]['sector_path']}")
-            ds = xr.open_dataset(
-                params[sector]["sector_path"], engine="zarr", consolidated=True
-            )
+            ds = xr.open_zarr(params[sector]["sector_path"], consolidated=True)
             ds = ds[params[sector][var]].rename(var)
             ds = xr.where(np.isinf(ds), np.nan, ds)
             datasets.append(ds)
@@ -260,12 +255,12 @@ def subset_USA_reduced_damages(
 ):
 
     if recipe == "adding_up":
-        ds = xr.open_dataset(
-            f"{input_path}/{sector}/{recipe}_{reduction}.zarr", engine="zarr"
+        ds = xr.open_zarr(
+            f"{input_path}/{sector}/{recipe}_{reduction}.zarr",
         )
     elif recipe == "risk_aversion":
-        ds = xr.open_dataset(
-            f"{input_path}/{sector}/{recipe}_{reduction}_eta{eta}.zarr", engine="zarr"
+        ds = xr.open_zarr(
+            f"{input_path}/{sector}/{recipe}_{reduction}_eta{eta}.zarr",
         )
 
     subset = ds.sel(region=[i for i in ds.region.values if "USA" in i])
@@ -292,9 +287,8 @@ def subset_USA_ssp_econ(
     out_path,
 ):
 
-    zarr = xr.open_dataset(
+    zarr = xr.open_zarr(
         in_path,
-        engine="zarr",
         consolidated=True,
     )
 
@@ -340,18 +334,16 @@ def clip_damages(
     histclim = params["histclim"]
     delta = params["delta"]
 
-    with xr.open_dataset(path, engine="zarr", chunks=None)[delta] as ds:
-        with xr.open_dataset(econ_path, engine="zarr", chunks=None) as gdppc:
+    with xr.open_zarr(path, chunks=None)[delta] as ds:
+        with xr.open_zarr(econ_path, chunks=None) as gdppc:
 
             ce_batch_dims = [i for i in ds.dims]
             ce_batch_coords = {c: ds[c].values for c in ce_batch_dims}
             ce_batch_coords["region"] = [
                 i for i in ds.region.values if i in gdppc.region.values
             ]
             ce_shapes = [len(ce_batch_coords[c]) for c in ce_batch_dims]
-            ce_chunks = [
-                xr.open_dataset(path, engine="zarr").chunks[c][0] for c in ce_batch_dims
-            ]
+            ce_chunks = [xr.open_zarr(path).chunks[c][0] for c in ce_batch_dims]
             print(ce_chunks)
 
     template = xr.DataArray(
@@ -370,7 +362,7 @@ def chunk_func(
         region = damages.region.values
 
         gdppc = (
-            xr.open_dataset(econ_path, engine="zarr", chunks=None)
+            xr.open_zarr(econ_path, chunks=None)
             .sel(year=year, ssp=ssp, model=model, region=region, drop=True)
             .gdppc
         )
@@ -402,7 +394,7 @@ def chunk_func(
 
         return damages
 
-    data = xr.open_dataset(path, engine="zarr")
+    data = xr.open_zarr(path)
 
     for var in [delta, histclim]:
         out = (

dscim/tests/test_simple_storage.py

@@ -45,12 +45,8 @@ def test_adding_up_damages(stacked_damages):
         stacked_damages.adding_up_damages,
         (
             (
-                xr.open_dataset(
-                    f"{stacked_damages.ce_path}/adding_up_no_cc.zarr", engine="zarr"
-                ).no_cc
-                - xr.open_dataset(
-                    f"{stacked_damages.ce_path}/adding_up_cc.zarr", engine="zarr"
-                ).cc
+                xr.open_zarr(f"{stacked_damages.ce_path}/adding_up_no_cc.zarr").no_cc
+                - xr.open_zarr(f"{stacked_damages.ce_path}/adding_up_cc.zarr").cc
             )
             * stacked_damages.pop
         ).sum("region"),