Merge pull request #44 from kthyng/main

added default flag for dask slicing

extract_model/accessor.py

@@ -36,13 +36,24 @@ def sub_grid(self, bbox, drop=True):
 
         return em.sub_grid(self.ds, bbox)
 
-    def sub_bbox(self, bbox, drop=True):
+    def sub_bbox(self, bbox, drop=True, dask_array_chunks=True):
         """Subset DataArray in space defined by bbox.
 
         See full docs at `em.sub_bbox()`.
         """
 
-        return em.sub_bbox(self.ds, bbox, drop=drop)
+        dss = []
+        Vars = [
+            Var for Var in self.ds.data_vars if "longitude" in self.ds[Var].cf.coords
+        ]
+        for Var in Vars:
+            dss.append(
+                em.sub_bbox(
+                    self.ds[Var], bbox, drop=drop, dask_array_chunks=dask_array_chunks
+                )
+            )
+
+        return xr.merge(dss)
 
 
 @xr.register_dataarray_accessor("em")
@@ -235,10 +246,12 @@ def interp2d(
 
         return da
 
-    def sub_bbox(self, bbox, drop=True):
+    def sub_bbox(self, bbox, drop=True, dask_array_chunks=True):
         """Subset DataArray in space defined by bbox.
 
         See full docs at `em.sub_bbox()`.
         """
 
-        return em.sub_bbox(self.da, bbox, drop=drop)
+        return em.sub_bbox(
+            self.da, bbox, drop=drop, dask_array_chunks=dask_array_chunks
+        )

extract_model/utils.py

@@ -6,7 +6,7 @@
 import xarray as xr
 
 
-def sub_grid(ds, bbox):
+def sub_grid(ds, bbox, dask_array_chunks=True):
     """Subset Dataset grids.
 
     Preserves horizontal grid structure, which matters for ROMS.
@@ -24,6 +24,8 @@ def sub_grid(ds, bbox):
         xarray Dataset to select model output from.
     bbox: list
         The bounding box for subsetting is defined as [min_lon, min_lat, max_lon, max_lat]
+    dask_array_chunks: boolean, optional
+        If True, avoids creating large chunks in slicing operation. If False, accept the large chunk and silence this warning. Comes up if Slicing is producing a large chunk.
 
     Returns
     -------
@@ -69,11 +71,14 @@ def sub_grid(ds, bbox):
         if "xi_psi" in ds:
             sel_dict["xi_psi"] = xi_rho[:-1]
         # adjust dimensions of full dataset
-        ds_new = ds.sel(sel_dict)
+        import dask
+
+        with dask.config.set(**{"array.slicing.split_large_chunks": dask_array_chunks}):
+            ds_new = ds.sel(sel_dict)
 
     elif len(lon_names) == 1:
 
-        ds_new = sub_bbox(ds, bbox)
+        ds_new = sub_bbox(ds, bbox, drop=True)
 
     else:
         # raise exception
@@ -84,7 +89,7 @@ def sub_grid(ds, bbox):
     return ds_new
 
 
-def sub_bbox(da, bbox, drop=True):
+def sub_bbox(da, bbox, drop=True, dask_array_chunks=True):
     """Subset DataArray in space.
 
     Can also be used on a Dataset if there is only one horizontal grid.
@@ -98,6 +103,8 @@ def sub_bbox(da, bbox, drop=True):
     drop: bool, optional
         This is passed onto xarray's `da.where()` function. If True, coordinates outside bbox
         are dropped from the DataArray, otherwise they are kept but masked/nan'ed.
+    dask_array_chunks: boolean, optional
+        If True, avoids creating large chunks in slicing operation. If False, accept the large chunk and silence this warning. Comes up if Slicing is producing a large chunk.
 
     Notes
     -----
@@ -112,7 +119,10 @@ def sub_bbox(da, bbox, drop=True):
         & (da.cf["latitude"] < bbox[3])
     ).compute()
 
-    da_smaller = da.where(box, drop=drop)
+    import dask
+
+    with dask.config.set(**{"array.slicing.split_large_chunks": dask_array_chunks}):
+        da_smaller = da.where(box, drop=drop)
 
     return da_smaller