-
Notifications
You must be signed in to change notification settings - Fork 14
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
move storage_volume intelligence to workflow
- Loading branch information
1 parent
827bad6
commit 103169f
Showing
3 changed files
with
121 additions
and
76 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -7,3 +7,4 @@ | |
| from .merge import * | ||
| from .tiling import * | ||
| from .curvenumber import * | ||
| from .storage_volume import * | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |