181 further improve caching by checking only relevant coordinates for…

… certain datasets (#202)

* add option in get_ds_empty to specify dimensions

* fix error for vertex grids

---------

Co-authored-by: Ruben Caljé <[email protected]>

nlmod/dims/grid.py

@@ -1306,7 +1306,7 @@ def gdf_to_bool_da(gdf, ds):
     return da
 
 
-def gdf_to_bool_ds(gdf, ds, da_name):
+def gdf_to_bool_ds(gdf, ds, da_name, keep_coords=None):
     """convert a GeoDataFrame with polygon geometries into a model dataset with
     a data_array named 'da_name' in which each cell is 1 (True) if one or more
     geometries are (partly) in that cell.
@@ -1319,14 +1319,17 @@ def gdf_to_bool_ds(gdf, ds, da_name):
         xarray with model data
     da_name : str
         The name of the variable with boolean data in the ds_out
+    keep_coords : tuple or None, optional
+        the coordinates in ds the you want keep in your empty ds. If None all
+        coordinates are kept from original ds. The default is None.
 
     Returns
     -------
     ds_out : xr.Dataset
         Dataset with a single DataArray, this DataArray is 1 if polygon is in
         cell, 0 otherwise. Grid dimensions according to ds and mfgrid.
     """
-    ds_out = util.get_ds_empty(ds)
+    ds_out = util.get_ds_empty(ds, keep_coords=keep_coords)
     ds_out[da_name] = gdf_to_bool_da(gdf, ds)
 
     return ds_out
@@ -1579,6 +1582,8 @@ def mask_model_edge(ds, idomain):
     ds_out : xarray.Dataset
         dataset with edge mask array
     """
+    ds = ds.copy()  # avoid side effects
+
     # add constant head cells at model boundaries
     if "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
         raise NotImplementedError("model edge not yet calculated for rotated grids")

nlmod/read/ahn.py

@@ -69,7 +69,7 @@ def get_ahn(ds=None, identifier="AHN4_DTM_5m", method="average", extent=None):
     if ds is None:
         return ahn_da
 
-    ds_out = get_ds_empty(ds)
+    ds_out = get_ds_empty(ds, keep_coords=("y", "x"))
     ds_out["ahn"] = ahn_da
 
     return ds_out

nlmod/read/jarkus.py

@@ -50,7 +50,7 @@ def get_bathymetry(ds, northsea, kind="jarkus", method="average"):
     data is resampled to the modelgrid. Maybe we can speed up things by
     changing the order in which operations are executed.
     """
-    ds_out = get_ds_empty(ds)
+    ds_out = get_ds_empty(ds, keep_coords=("y", "x"))
 
     # no bathymetry if we don't have northsea
     if (northsea == 0).all():

nlmod/read/knmi.py

@@ -62,7 +62,7 @@ def get_recharge(ds, method="linear", most_common_station=False):
     start = pd.Timestamp(ds.time.attrs["start"])
     end = pd.Timestamp(ds.time.data[-1])
 
-    ds_out = util.get_ds_empty(ds)
+    ds_out = util.get_ds_empty(ds, keep_coords=("time", "y", "x"))
     ds_out.attrs["gridtype"] = ds.gridtype
 
     # get recharge data array

nlmod/read/rws.py

@@ -78,7 +78,7 @@ def get_surface_water(ds, da_basename):
         peil = xr.where(area_pol > area, row["peil"], peil)
         area = xr.where(area_pol > area, area_pol, area)
 
-    ds_out = util.get_ds_empty(ds)
+    ds_out = util.get_ds_empty(ds, keep_coords=("y", "x"))
     ds_out[f"{da_basename}_area"] = area
     ds_out[f"{da_basename}_area"].attrs["units"] = "m2"
     ds_out[f"{da_basename}_cond"] = cond
@@ -127,7 +127,7 @@ def get_northsea(ds, da_name="northsea"):
         )
     ]
 
-    ds_out = dims.gdf_to_bool_ds(swater_zee, ds, da_name)
+    ds_out = dims.gdf_to_bool_ds(swater_zee, ds, da_name, keep_coords=("y", "x"))
 
     return ds_out
 

nlmod/util.py

@@ -75,23 +75,29 @@ def get_exe_path(exe_name="mf6"):
     return exe_path
 
 
-def get_ds_empty(ds):
-    """get a copy of a model dataset with only coordinate information.
+def get_ds_empty(ds, keep_coords=None):
+    """get a copy of a dataset with only coordinate information.
 
     Parameters
     ----------
     ds : xr.Dataset
         dataset with coordinates
+    keep_coords : tuple or None, optional
+        the coordinates in ds the you want keep in your empty ds. If None all
+        coordinates are kept from original ds. The default is None.
 
     Returns
     -------
     empty_ds : xr.Dataset
-        dataset with only model coordinate information
+        dataset with only coordinate information
     """
+    if keep_coords is None:
+        keep_coords = list(ds.coords)
 
     empty_ds = xr.Dataset()
     for coord in list(ds.coords):
-        empty_ds = empty_ds.assign_coords(coords={coord: ds[coord]})
+        if coord in keep_coords:
+            empty_ds = empty_ds.assign_coords(coords={coord: ds[coord]})
 
     return empty_ds