From 103169f9792c1bc0f8f4ef2fed9bf19ef87be1f4 Mon Sep 17 00:00:00 2001 From: roeldegoede Date: Thu, 9 Nov 2023 09:45:29 +0100 Subject: [PATCH] move storage_volume intelligence to workflow --- hydromt_sfincs/sfincs.py | 80 +------------- hydromt_sfincs/workflows/__init__.py | 1 + hydromt_sfincs/workflows/storage_volume.py | 116 +++++++++++++++++++++ 3 files changed, 121 insertions(+), 76 deletions(-) create mode 100644 hydromt_sfincs/workflows/storage_volume.py diff --git a/hydromt_sfincs/sfincs.py b/hydromt_sfincs/sfincs.py index 4727a772..e8d09600 100644 --- a/hydromt_sfincs/sfincs.py +++ b/hydromt_sfincs/sfincs.py @@ -1619,82 +1619,10 @@ def setup_storage_volume( else: da_vol = xr.full_like(self.mask, 0, dtype=np.float64) - # loop over the gdf rows and rasterize each geometry - for i, row in gdf.iterrows(): - # create a gdf with only the current row - single_gdf = gpd.GeoDataFrame(gdf.loc[[i]]).reset_index(drop=True) - - single_vol = single_gdf.get("volume") - single_height = single_gdf.get("height") - - # check if volume or height is provided in the gdf or as input - if ( - single_vol is None or np.isnan(single_vol).any() - ): # volume not provided or nan - if ( - single_height is None or np.isnan(single_height).any() - ): # height not provided or nan - if volume is not None: # volume provided as input (list) - single_vol = ( - volume if not isinstance(volume, list) else volume[i] - ) - elif height is not None: # height provided as input (list) - single_height = ( - height if not isinstance(height, list) else height[i] - ) - else: # no volume or height provided - self.logger.warning( - f"No volume or height provided for storage location {i}" - ) - continue - else: # height provided in gdf - single_height = single_height[0] - else: # volume provided in gdf - single_vol = single_vol[0] - - # check if gdf has point or polyhon geometry - if single_gdf.geometry.type[0] == "Point": - # get x and y coordinate of the point in crs of the grid - x, y = single_gdf.geometry.iloc[0].coords[0] - - # check if the grid is rotated - if da_vol.raster.rotation != 0: - # rotate the point - x, y = ~da_vol.raster.transform * (x, y) - # select the grid cell nearest to the point - closest_point = da_vol.reindex(x=da_vol.x, y=da_vol.y).sel( - x=x, y=y, method="nearest" - ) - else: - # select the grid cell nearest to the point - closest_point = da_vol.sel(x=x, y=y, method="nearest") - - # add the volume to the grid cell - if single_vol is not None and not np.isnan(single_vol).any(): - da_vol.loc[ - dict(x=closest_point.x.item(), y=closest_point.y.item()) - ] += single_vol - else: - self.logger.warning( - f"No volume provided for storage location of type Point" - ) - - elif single_gdf.geometry.type[0] == "Polygon": - # rasterize the geometry - area = da_vol.raster.rasterize_geometry(single_gdf, method="area") - total_area = area.sum().values - - # calculate the volume per cell and add it to the grid - if single_vol is not None and not np.isnan(single_vol).any(): - da_vol += area / total_area * single_vol - elif ( - single_height is not None and not np.isnan(single_height).any() - ): - da_vol += area * single_height - else: - self.logger.warning( - f"No volume or height provided for storage location of type Polygon" - ) + # add storage volumes form gdf to da_vol + da_vol = workflows.add_storage_volume( + da_vol, gdf, volume=volume, height=height, logger=self.logger, + ) # set grid mname = "vol" diff --git a/hydromt_sfincs/workflows/__init__.py b/hydromt_sfincs/workflows/__init__.py index 4ea63132..f51f7eca 100644 --- a/hydromt_sfincs/workflows/__init__.py +++ b/hydromt_sfincs/workflows/__init__.py @@ -7,3 +7,4 @@ from .merge import * from .tiling import * from .curvenumber import * +from .storage_volume import * diff --git a/hydromt_sfincs/workflows/storage_volume.py b/hydromt_sfincs/workflows/storage_volume.py new file mode 100644 index 00000000..f0ddb922 --- /dev/null +++ b/hydromt_sfincs/workflows/storage_volume.py @@ -0,0 +1,116 @@ +import logging +from typing import Union, List + +import geopandas as gpd +import numpy as np +import xarray as xr + +logger = logging.getLogger(__name__) + + +def add_storage_volume( + da_vol: xr.DataArray, + gdf: gpd.GeoDataFrame, + volume: Union[float, List[float]] = None, + height: Union[float, List[float]] = None, + logger=logger, + ) -> xr.DataArray: + """Add storage volume to a grid based on a GeoDataFrame with storage locations. + + Parameters + ---------- + da_vol : xr.DataArray + DataArray with the grid to which the storage volume should be added. + gdf : gpd.GeoDataFrame + GeoDataFrame with the storage locations (polygon or point geometry file). + Optional "volume" or "height" attributes can be provided to set the storage volume. + volume : Union[float, List[float]], optional + Volume of the storage locations [m3], by default None. + height : Union[float, List[float]], optional + Height of the storage locations [m], by default None. + + Returns + ------- + xr.DataArray + DataArray with the grid including the storage volume. + + """ + + # loop over the gdf rows and rasterize each geometry + for i, row in gdf.iterrows(): + # create a gdf with only the current row + single_gdf = gpd.GeoDataFrame(gdf.loc[[i]]).reset_index(drop=True) + + single_vol = single_gdf.get("volume") + single_height = single_gdf.get("height") + + # check if volume or height is provided in the gdf or as input + if ( + single_vol is None or np.isnan(single_vol).any() + ): # volume not provided or nan + if ( + single_height is None or np.isnan(single_height).any() + ): # height not provided or nan + if volume is not None: # volume provided as input (list) + single_vol = ( + volume if not isinstance(volume, list) else volume[i] + ) + elif height is not None: # height provided as input (list) + single_height = ( + height if not isinstance(height, list) else height[i] + ) + else: # no volume or height provided + logger.warning( + f"No volume or height provided for storage location {i}" + ) + continue + else: # height provided in gdf + single_height = single_height[0] + else: # volume provided in gdf + single_vol = single_vol[0] + + # check if gdf has point or polyhon geometry + if single_gdf.geometry.type[0] == "Point": + # get x and y coordinate of the point in crs of the grid + x, y = single_gdf.geometry.iloc[0].coords[0] + + # check if the grid is rotated + if da_vol.raster.rotation != 0: + # rotate the point + x, y = ~da_vol.raster.transform * (x, y) + # select the grid cell nearest to the point + closest_point = da_vol.reindex(x=da_vol.x, y=da_vol.y).sel( + x=x, y=y, method="nearest" + ) + else: + # select the grid cell nearest to the point + closest_point = da_vol.sel(x=x, y=y, method="nearest") + + # add the volume to the grid cell + if single_vol is not None and not np.isnan(single_vol).any(): + da_vol.loc[ + dict(x=closest_point.x.item(), y=closest_point.y.item()) + ] += single_vol + else: + logger.warning( + f"No volume provided for storage location of type Point" + ) + + elif single_gdf.geometry.type[0] == "Polygon": + # rasterize the geometry + area = da_vol.raster.rasterize_geometry(single_gdf, method="area") + total_area = area.sum().values + + # calculate the volume per cell and add it to the grid + if single_vol is not None and not np.isnan(single_vol).any(): + da_vol += area / total_area * single_vol + elif ( + single_height is not None and not np.isnan(single_height).any() + ): + da_vol += area * single_height + else: + logger.warning( + f"No volume or height provided for storage location of type Polygon" + ) + + return da_vol \ No newline at end of file