Skip to content

Commit

Permalink
Working allocation
Browse files Browse the repository at this point in the history
  • Loading branch information
@dalmijn
dalmijn committed Dec 14, 2023
1 parent 7e4bbcc commit e7c536e
Show file tree
Hide file tree
Showing 2 changed files with 159 additions and 51 deletions.
82 changes: 59 additions & 23 deletions hydromt_wflow/wflow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2492,26 +2492,78 @@ def setup_rootzoneclim(
# update config
self.set_config("input.vertical.rootingdepth", update_toml_rootingdepth)

def setup_water_demand(
def setup_allocation(
self,
min_area: float | int = 0,
admin_bounds_fn: str = "gadm",
admin_level: int = 0,
):
"""_summary_.
_extended_summary_
Parameters
----------
min_area : float | int
_description_
admin_bounds_fn : str, optional
_description_, by default "gadm"
admin_level : int, optional
_description_, by default 0
"""
# Will be fixes but for know this is done like this
# TODO fix in the future
admin_bounds = None
if admin_bounds_fn is not None:
admin_path = self.data_catalog[admin_bounds_fn].path
admin_bounds = self.data_catalog.get_geodataframe(
admin_path, geom=self.region, layer=f"level{admin_level}"
)
# Add this identifier for usage in the workflow
admin_bounds["admin_id"] = range(len(admin_bounds))

# Create the allocation grid
alloc = workflows.demand.allocate(
ds_like=self.grid,
min_area=min_area,
admin_bounds=admin_bounds,
basins=self.geoms["basins"],
rivers=self.geoms["rivers"],
)
self.set_grid(alloc)

# Update the settings toml
self.set_config("input.vertical.waterallocation.areas", "Allocation_id")

def setup_non_irigation(
self,
non_irigation_fn: str = "pcr_globwb",
non_irigation_vars: list = ["dom", "ind", "lsk"],
non_irigation_method: str = "nearest",
population_fn: str = "worldpop_2020_constrained",
population_method: str = "sum",
admin_bounds_fn: str = "gadm",
admin_level: int = 0,
):
"""_summary_.
_extended_summary_
Parameters
----------
pcr_fn : str, optional
non_irigation_fn : str, optional
_description_, by default "pcr_globwb"
pcr_vars : list, optional
non_irigation_vars : list, optional
_description_, by default ["dom", "ind", "lsk"]
non_irigation_method : str, optional
_description_, by default "nearest"
population_fn : str, optional
_description_, by default "worldpop_2020_constrained"
population_method : str, optional
_description_, by default "sum"
Raises
------
ValueError
_description_
"""
if not all([item in ["dom", "ind", "lsk"] for item in non_irigation_vars]):
raise ValueError("")
Expand Down Expand Up @@ -2543,28 +2595,12 @@ def setup_water_demand(
self.set_grid(non_irigation)
self.set_grid(popu)

# Will be fixes but for know this is done like this
# TODO fix in the future
admin_path = self.data_catalog[admin_bounds_fn].path
admin_bounds = self.data_catalog.get_geodataframe(
admin_path, geom=self.region, layer=f"level{admin_level}"
)

# Create the allocation grid
alloc = workflows.demand.allocate(
ds_like=self.grid,
admin_bounds=admin_bounds,
basins=self.geoms["basins"],
rivers=self.geoms["rivers"],
)
self.set_grid(alloc)

# Update the settings toml
# Update the settings toml with non irigation stuff
for var in non_irigation.data_vars:
sname, suffix = var.split("_")
lname = workflows.demand.map_vars[sname]
self.set_config(
f"vertical.{lname}.demand_{suffix}",
f"input.vertical.{lname}.demand_{suffix}",
var,
)

Expand Down
128 changes: 100 additions & 28 deletions hydromt_wflow/workflows/demand.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import math

import pandas as pd
import xarray as xr
from affine import Affine
from hydromt.raster import full_from_transform
Expand Down Expand Up @@ -37,6 +38,33 @@ def transform_half_degree(
return affine, w, h


def touch_intersect(row, vector):
"""_summary_.
_extended_summary_
Parameters
----------
row : _type_
_description_
vector : _type_
_description_
Returns
-------
_type_
_description_
"""
contain = True
_t = sum(vector.touches(row.geometry))
_i = sum(vector.intersects(row.geometry))
diff = abs(_i - _t)
if diff == 0:
contain = False
row["contain"] = contain
return row


def non_irigation(
ds: xr.Dataset,
ds_like: xr.Dataset,
Expand Down Expand Up @@ -128,36 +156,57 @@ def non_irigation(

def allocate(
ds_like: xr.Dataset,
min_area: float | int,
admin_bounds: object,
basins: xr.Dataset,
rivers: xr.Dataset,
):
"""_summary_.
_extended_summary_
Parameters
----------
ds_like : xr.Dataset
_description_
min_area : float | int
_description_
admin_bounds : object
_description_
basins : xr.Dataset
_description_
rivers : xr.Dataset
_description_
"""
# Split based on admin bounds
split_basins = basins.overlay(
admin_bounds,
how="union",
)
split_basins = split_basins[~split_basins["value"].isna()]

# Remove unneccessary stuff
cols = split_basins.columns.drop(["value", "geometry", "NAME_2"]).tolist()
split_basins.drop(cols, axis=1, inplace=True)
# Use this uid to dissolve on later
split_basins["uid"] = range(len(split_basins))

# Dissolve cut pieces back
for _, row in split_basins.iterrows():
if not str(row.NAME_2).lower() == "nan":
continue
touched = split_basins[split_basins.touches(row.geometry)]
uid = touched[touched["value"] == row.value].uid.values[0]
split_basins.loc[split_basins["uid"] == row.uid, "uid"] = uid
split_basins = split_basins.dissolve("uid", sort=False, as_index=False)
sub_basins = basins.copy()
sub_basins["uid"] = range(len(sub_basins))

if admin_bounds is not None:
sub_basins = basins.overlay(
admin_bounds,
how="union",
)
sub_basins = sub_basins[~sub_basins["value"].isna()]

# Remove unneccessary stuff
cols = sub_basins.columns.drop(["value", "geometry", "admin_id"]).tolist()
sub_basins.drop(cols, axis=1, inplace=True)
# Use this uid to dissolve on later
sub_basins["uid"] = range(len(sub_basins))

# Dissolve cut pieces back
for _, row in sub_basins.iterrows():
if not str(row.admin_id).lower() == "nan":
continue
touched = sub_basins[sub_basins.touches(row.geometry)]
uid = touched[touched["value"] == row.value].uid.values[0]
sub_basins.loc[sub_basins["uid"] == row.uid, "uid"] = uid
sub_basins = sub_basins.dissolve("uid", sort=False, as_index=False)

# Set the contain flag per geom
sub_basins = sub_basins.apply(lambda row: touch_intersect(row, rivers), axis=1)
sub_basins["sqkm"] = sub_basins.geometry.to_crs(3857).area / 1000**2
_count = 0

# Create touched and not touched by rivers datasets
Expand All @@ -168,13 +217,26 @@ def allocate(

# Everything touched by river based on difference
# (is not what we want, yet)
riv_touch = split_basins.sjoin(
rivers,
)
if _count != 0:
sub_basins = sub_basins.apply(
lambda row: touch_intersect(row, rivers), axis=1
)
sub_basins["sqkm"] = sub_basins.geometry.to_crs(3857).area / 1000**2

# Set no_riv and riv (what's touched and what's not)
no_riv = split_basins[~split_basins.geometry.isin(riv_touch.geometry)]
riv = split_basins[split_basins.geometry.isin(riv_touch.geometry)]
no_riv = sub_basins[~sub_basins["contain"]]
riv = sub_basins[sub_basins["contain"]]

# Include minimal area option
min_basins = sub_basins[sub_basins["sqkm"] < min_area]
min_basins = min_basins[~min_basins.uid.isin(no_riv.uid)]
# Only concatenate if there are different small basins
# compared to basins that do not touch a river
if not min_basins.empty:
no_riv = pd.concat(
[no_riv, min_basins],
ignore_index=True,
)

_n = 0

Expand All @@ -192,7 +254,7 @@ def allocate(
uid = touched[touched["area"] == touched["area"].max()].uid.values[0]
# Set the identifier to the new value
# (i.e. the touched basin)
split_basins.loc[split_basins["uid"] == row.uid, "uid"] = uid
sub_basins.loc[sub_basins["uid"] == row.uid, "uid"] = uid
_n += 1

# Ensure a break if nothing is touched
Expand All @@ -201,6 +263,16 @@ def allocate(
if _n == 0:
break

split_basins = split_basins.dissolve("uid", sort=False, as_index=False)
sub_basins = sub_basins.dissolve("uid", sort=False, as_index=False)
_count += 1
pass

alloc = full_from_transform(
ds_like.raster.transform,
ds_like.raster.shape,
crs=ds_like.raster.crs,
lazy=True,
)
alloc = alloc.raster.rasterize(sub_basins, col_name="uid", nodata=-9999)
alloc.name = "Allocation_id"
alloc = alloc.rename(dict(zip(alloc.dims, list(ds_like.dims)[:2])))
return alloc

0 comments on commit e7c536e

Please sign in to comment.