From e107c90972f88f786901958d999b168b78d21f61 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 2 Apr 2024 18:19:17 +0200 Subject: [PATCH 01/35] Write first utility functions to allocate river cells --- imod/prepare/river.py | 214 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 214 insertions(+) create mode 100644 imod/prepare/river.py diff --git a/imod/prepare/river.py b/imod/prepare/river.py new file mode 100644 index 000000000..d214d8e8f --- /dev/null +++ b/imod/prepare/river.py @@ -0,0 +1,214 @@ +""" +This module contains all kinds of utilities to prepare rivers +""" + +from enum import Enum + +from imod.mf6.utilities.grid import create_layered_top +from imod.prepare.layer import get_upper_active_layer_number +from imod.schemata import DimsSchema +from imod.typing import GridDataArray + + +class ALLOCATION_OPTION(Enum): + between_stage_and_riv_bottom = 0 + first_active = -1 + first_active_drn_above_stage = 1 + in_riv_bottom_layer = 2 + + +PLANAR_GRID = ( + DimsSchema("time", "y", "x") + | DimsSchema("y", "x") + | DimsSchema("time", "{face_dim}") + | DimsSchema("{face_dim}") +) + + +def allocate_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +): + match allocation_option: + case ALLOCATION_OPTION.between_stage_and_riv_bottom: + return allocate_cells__between_stage_and_bottom_elevation( + top, bottom, stage, bottom_elevation + ) + case ALLOCATION_OPTION.first_active: + return allocate_cells__first_active(active, bottom, bottom_elevation) + case ALLOCATION_OPTION.first_active_drn_above_stage: + return allocate_cells__first_active_drn_above_stage( + active, top, bottom, stage + ) + case ALLOCATION_OPTION.in_riv_bottom_layer: + return allocate_cells__in_bottom_elevation_layer( + top, bottom, bottom_elevation + ) + + +def _is_layered(grid: GridDataArray): + return "layer" in grid.sizes and grid.sizes["layer"] > 1 + + +def allocate_cells__between_stage_and_bottom_elevation( + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +): + """ + Allocate cells inbetween river stage and river bottom_elevation. Compared to + iMOD5.6, this is similar to setting IDEFFLAYER=0 in the RUNFILE function. + + Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer + dimension. + + Parameters + ---------- + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + stage: GridDataArray + river stage + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + GridDataArray + River cells + """ + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + + return (stage <= top_layered) & (bottom_elevation >= bottom) + + +def allocate_cells__first_active( + active: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray +): + """ + Allocate cells inbetween first active layer and river bottom elevation. + Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE + function. + + Note that ``bottom_elevation`` is not allowed to have a layer dimension. + + Parameters + ---------- + active: GridDataArray + active model cells + bottom: GridDataArray + bottom of model layers + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + GridDataArray + River cells + """ + PLANAR_GRID.validate(bottom_elevation) + + upper_active_layer = get_upper_active_layer_number(active) + layer = active.coords["layer"] + + return (layer >= upper_active_layer) & (bottom_elevation >= bottom) + + +def allocate_cells__first_active_drn_above_stage( + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +): + """ + Allocate cells inbetween first active layer and river bottom elevation. + Cells above river stage are deactivated and returned as drn cells. Compared + to iMOD5.6, this is similar to setting IDEFFLAYER=1 in the RUNFILE function. + + Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer + dimension. + + Parameters + ---------- + active: GridDataArray + active grid cells + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + stage: GridDataArray + river stage + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + riv_cells: GridDataArray + River cells (below stage) + drn_cells: GridDataArray + Drainage cells (above stage) + """ + + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + + upper_active_layer = get_upper_active_layer_number(active) + layer = active.coords["layer"] + drn_cells = (layer >= upper_active_layer) & (stage <= top_layered) + riv_cells = (layer >= upper_active_layer) & ( + bottom_elevation >= bottom + ) != drn_cells + + return riv_cells, drn_cells + + +def allocate_cells__in_bottom_elevation_layer( + top: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray +): + """ + Allocate cells in river bottom elevation layer. Compared to iMOD5.6, this is + similar to setting IDEFFLAYER=2 in the RUNFILE function. + + Note that ``bottom_elevation`` is not allowed to have a layer dimension. + + Parameters + ---------- + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + GridDataArray + River cells + """ + + PLANAR_GRID.validate(bottom_elevation) + + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + + return (bottom_elevation < top_layered) & (bottom_elevation >= bottom) From d7819a0e903e0cc623f4f53c5abf9777065c2cee Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 13:17:31 +0200 Subject: [PATCH 02/35] Add docstring --- imod/prepare/river.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/imod/prepare/river.py b/imod/prepare/river.py index d214d8e8f..8832d3405 100644 --- a/imod/prepare/river.py +++ b/imod/prepare/river.py @@ -11,6 +11,10 @@ class ALLOCATION_OPTION(Enum): + """ + Enumerator for allocation settings. Numbers match the IDEFLAYER options in + iMOD 5.6. + """ between_stage_and_riv_bottom = 0 first_active = -1 first_active_drn_above_stage = 1 From be9d2c81fb4256e7d3997a639e770217ebf72476 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 15:57:41 +0200 Subject: [PATCH 03/35] Shorter and concise names, at extra allocation functions for specific packages --- imod/prepare/river.py | 136 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 111 insertions(+), 25 deletions(-) diff --git a/imod/prepare/river.py b/imod/prepare/river.py index 8832d3405..b51e4591d 100644 --- a/imod/prepare/river.py +++ b/imod/prepare/river.py @@ -5,7 +5,10 @@ from enum import Enum from imod.mf6.utilities.grid import create_layered_top -from imod.prepare.layer import get_upper_active_layer_number +from imod.prepare.layer import ( + get_upper_active_grid_cells, + get_upper_active_layer_number, +) from imod.schemata import DimsSchema from imod.typing import GridDataArray @@ -15,10 +18,12 @@ class ALLOCATION_OPTION(Enum): Enumerator for allocation settings. Numbers match the IDEFLAYER options in iMOD 5.6. """ - between_stage_and_riv_bottom = 0 - first_active = -1 - first_active_drn_above_stage = 1 - in_riv_bottom_layer = 2 + + stage_to_riv_bot = 0 + first_active_to_riv_bot = -1 + first_active_to_riv_bot__drn = 1 + at_elevation = 2 + at_first_active = 9 # Not an iMOD 5.6 option PLANAR_GRID = ( @@ -29,7 +34,7 @@ class ALLOCATION_OPTION(Enum): ) -def allocate_cells( +def allocate_river_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray, top: GridDataArray, @@ -38,19 +43,78 @@ def allocate_cells( bottom_elevation: GridDataArray, ): match allocation_option: - case ALLOCATION_OPTION.between_stage_and_riv_bottom: - return allocate_cells__between_stage_and_bottom_elevation( + case ALLOCATION_OPTION.stage_to_riv_bot: + return allocate_cells__stage_to_riv_bot( top, bottom, stage, bottom_elevation ) - case ALLOCATION_OPTION.first_active: - return allocate_cells__first_active(active, bottom, bottom_elevation) - case ALLOCATION_OPTION.first_active_drn_above_stage: - return allocate_cells__first_active_drn_above_stage( - active, top, bottom, stage + case ALLOCATION_OPTION.first_active_to_riv_bot: + return allocate_cells__first_active_to_riv_bot( + active, bottom, bottom_elevation ) - case ALLOCATION_OPTION.in_riv_bottom_layer: - return allocate_cells__in_bottom_elevation_layer( - top, bottom, bottom_elevation + case ALLOCATION_OPTION.first_active_to_riv_bot__drn: + return allocate_cells__first_active_to_riv_bot__drn( + active, top, bottom, stage, bottom_elevation + ) + case ALLOCATION_OPTION.at_elevation: + return allocate_cells__at_elevation(top, bottom, bottom_elevation) + case ALLOCATION_OPTION.at_first_active: + return allocate_cells__at_first_active(active) + + +def allocate_drain_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + elevation: GridDataArray, +): + match allocation_option: + case ALLOCATION_OPTION.at_elevation: + return allocate_cells__at_elevation(top, bottom, elevation) + case ALLOCATION_OPTION.at_first_active: + return allocate_cells__at_first_active(active) + case _: + raise ValueError( + "Received incompatible setting for drains, only" + f"'{ALLOCATION_OPTION.at_elevation.name}' and" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) + + +def allocate_ghb_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + head: GridDataArray, +): + match allocation_option: + case ALLOCATION_OPTION.at_elevation: + return allocate_cells__at_elevation(top, bottom, head) + case ALLOCATION_OPTION.at_first_active: + return allocate_cells__at_first_active(active) + case _: + raise ValueError( + "Received incompatible setting for drains, only" + f"'{ALLOCATION_OPTION.at_elevation.name}' and" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) + + +def allocate_rch_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, +): + match allocation_option: + case ALLOCATION_OPTION.at_first_active: + return allocate_cells__at_first_active(active) + case _: + raise ValueError( + "Received incompatible setting for drains, only" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" ) @@ -58,7 +122,7 @@ def _is_layered(grid: GridDataArray): return "layer" in grid.sizes and grid.sizes["layer"] > 1 -def allocate_cells__between_stage_and_bottom_elevation( +def allocate_cells__stage_to_riv_bot( top: GridDataArray, bottom: GridDataArray, stage: GridDataArray, @@ -98,7 +162,7 @@ def allocate_cells__between_stage_and_bottom_elevation( return (stage <= top_layered) & (bottom_elevation >= bottom) -def allocate_cells__first_active( +def allocate_cells__first_active_to_riv_bot( active: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray ): """ @@ -130,7 +194,7 @@ def allocate_cells__first_active( return (layer >= upper_active_layer) & (bottom_elevation >= bottom) -def allocate_cells__first_active_drn_above_stage( +def allocate_cells__first_active_to_riv_bot__drn( active: GridDataArray, top: GridDataArray, bottom: GridDataArray, @@ -184,8 +248,8 @@ def allocate_cells__first_active_drn_above_stage( return riv_cells, drn_cells -def allocate_cells__in_bottom_elevation_layer( - top: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray +def allocate_cells__at_elevation( + top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray ): """ Allocate cells in river bottom elevation layer. Compared to iMOD5.6, this is @@ -199,8 +263,8 @@ def allocate_cells__in_bottom_elevation_layer( top of model layers bottom: GridDataArray bottom of model layers - bottom_elevation: GridDataArray - river bottom elevation + elevation: GridDataArray + elevation. Can be river bottom, drain elevation or head of GHB. Returns ------- @@ -208,11 +272,33 @@ def allocate_cells__in_bottom_elevation_layer( River cells """ - PLANAR_GRID.validate(bottom_elevation) + PLANAR_GRID.validate(elevation) if _is_layered(top): top_layered = top else: top_layered = create_layered_top(bottom, top) - return (bottom_elevation < top_layered) & (bottom_elevation >= bottom) + return (elevation < top_layered) & (elevation >= bottom) + + +def allocate_cells__at_first_active(active: GridDataArray): + """ + Allocate cells inbetween first active layer and river bottom elevation. + Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE + function. + + Note that ``bottom_elevation`` is not allowed to have a layer dimension. + + Parameters + ---------- + active: GridDataArray + active model cells + + Returns + ------- + GridDataArray + River cells + """ + + return get_upper_active_grid_cells(active) From 93d4b9264cc16da9928d2e8957834edf8563b785 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 16:00:26 +0200 Subject: [PATCH 04/35] Make functions semi-private --- imod/prepare/river.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/imod/prepare/river.py b/imod/prepare/river.py index b51e4591d..493ded664 100644 --- a/imod/prepare/river.py +++ b/imod/prepare/river.py @@ -44,21 +44,21 @@ def allocate_river_cells( ): match allocation_option: case ALLOCATION_OPTION.stage_to_riv_bot: - return allocate_cells__stage_to_riv_bot( + return _allocate_cells__stage_to_riv_bot( top, bottom, stage, bottom_elevation ) case ALLOCATION_OPTION.first_active_to_riv_bot: - return allocate_cells__first_active_to_riv_bot( + return _allocate_cells__first_active_to_riv_bot( active, bottom, bottom_elevation ) case ALLOCATION_OPTION.first_active_to_riv_bot__drn: - return allocate_cells__first_active_to_riv_bot__drn( + return _allocate_cells__first_active_to_riv_bot__drn( active, top, bottom, stage, bottom_elevation ) case ALLOCATION_OPTION.at_elevation: - return allocate_cells__at_elevation(top, bottom, bottom_elevation) + return _allocate_cells__at_elevation(top, bottom, bottom_elevation) case ALLOCATION_OPTION.at_first_active: - return allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active) def allocate_drain_cells( @@ -70,9 +70,9 @@ def allocate_drain_cells( ): match allocation_option: case ALLOCATION_OPTION.at_elevation: - return allocate_cells__at_elevation(top, bottom, elevation) + return _allocate_cells__at_elevation(top, bottom, elevation) case ALLOCATION_OPTION.at_first_active: - return allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active) case _: raise ValueError( "Received incompatible setting for drains, only" @@ -91,9 +91,9 @@ def allocate_ghb_cells( ): match allocation_option: case ALLOCATION_OPTION.at_elevation: - return allocate_cells__at_elevation(top, bottom, head) + return _allocate_cells__at_elevation(top, bottom, head) case ALLOCATION_OPTION.at_first_active: - return allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active) case _: raise ValueError( "Received incompatible setting for drains, only" @@ -109,7 +109,7 @@ def allocate_rch_cells( ): match allocation_option: case ALLOCATION_OPTION.at_first_active: - return allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active) case _: raise ValueError( "Received incompatible setting for drains, only" @@ -122,7 +122,7 @@ def _is_layered(grid: GridDataArray): return "layer" in grid.sizes and grid.sizes["layer"] > 1 -def allocate_cells__stage_to_riv_bot( +def _allocate_cells__stage_to_riv_bot( top: GridDataArray, bottom: GridDataArray, stage: GridDataArray, @@ -162,7 +162,7 @@ def allocate_cells__stage_to_riv_bot( return (stage <= top_layered) & (bottom_elevation >= bottom) -def allocate_cells__first_active_to_riv_bot( +def _allocate_cells__first_active_to_riv_bot( active: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray ): """ @@ -194,7 +194,7 @@ def allocate_cells__first_active_to_riv_bot( return (layer >= upper_active_layer) & (bottom_elevation >= bottom) -def allocate_cells__first_active_to_riv_bot__drn( +def _allocate_cells__first_active_to_riv_bot__drn( active: GridDataArray, top: GridDataArray, bottom: GridDataArray, @@ -248,7 +248,7 @@ def allocate_cells__first_active_to_riv_bot__drn( return riv_cells, drn_cells -def allocate_cells__at_elevation( +def _allocate_cells__at_elevation( top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray ): """ @@ -282,7 +282,7 @@ def allocate_cells__at_elevation( return (elevation < top_layered) & (elevation >= bottom) -def allocate_cells__at_first_active(active: GridDataArray): +def _allocate_cells__at_first_active(active: GridDataArray): """ Allocate cells inbetween first active layer and river bottom elevation. Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE From 639092e6c17e7d42870b04bd635bd42a0d1808de Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 16:03:44 +0200 Subject: [PATCH 05/35] Move module to better name and location --- imod/prepare/{river.py => topsystem/allocation.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename imod/prepare/{river.py => topsystem/allocation.py} (100%) diff --git a/imod/prepare/river.py b/imod/prepare/topsystem/allocation.py similarity index 100% rename from imod/prepare/river.py rename to imod/prepare/topsystem/allocation.py From d980d7f3e1912216cf2ef28ea36c7f2199bfe841 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 16:28:26 +0200 Subject: [PATCH 06/35] Move functions to new module and throw deprecation warning when being used --- imod/prepare/surface_water.py | 116 ++++----------------------- imod/prepare/topsystem/resistance.py | 109 +++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 102 deletions(-) create mode 100644 imod/prepare/topsystem/resistance.py diff --git a/imod/prepare/surface_water.py b/imod/prepare/surface_water.py index 5e3ff4298..d4f3bc3a1 100644 --- a/imod/prepare/surface_water.py +++ b/imod/prepare/surface_water.py @@ -1,109 +1,21 @@ -import numpy as np +# Import for backwards compatibility +import warnings +from imod.prepare.topsystem import c_leakage as _c_leakage +from imod.prepare.topsystem import c_radial as _c_radial -def c_radial(L, kh, kv, B, D): - """ - Ernst's radial resistance term to a drain. +WARNING_MESSAGE = ( + "function has been moved from imod.prepare.surface_water to" + "imod.prepare.topsystem, please update your scripts." + "imod.prepare.surface_water is going to be removed in version 1.0" +) - Parameters - ---------- - L : xr.DataArray of floats - distance between water features - kh : xr.DataArray of floats - horizontal conductivity - kv : xr.DataArray of floats - vertical conductivity - B : xr.DataArray of floats - water feature wetted perimeter - D : xr.DataArray of floats - saturated thickness of the top system - Returns - ------- - radial_c : xr.DataArray - Ernst's radial resistance for a drain - """ - # Take anisotropy into account fully - c = ( - L - / (np.pi * np.sqrt(kh * kv)) - * np.log((4.0 * D) / (np.pi * B) * np.sqrt(kh / kv)) - ) - c = c.where(~(c < 0.0), other=0.0) - return c +def c_radial(L, kh, kv, B, D): + warnings.warn(WARNING_MESSAGE, DeprecationWarning) + return _c_radial(L, kh, kv, B, D) def c_leakage(kh, kv, D, c0, c1, B, length, dx, dy): - """ - Computes the phreatic leakage resistance. - - Parameters - ---------- - kh : xr.DataArray of floats - horizontal conductivity of phreatic aquifer - kv : xr.DataArray of floats - vertical conductivity of phreatic aquifer - c0 : xr.DataArray of floats - hydraulic bed resistance of water feature - c1 : xr.DataArray of floats - hydraulic resistance of the first aquitard - D : xr.DataArray of floats - saturated thickness of the top system - B : xr.DataArray of floats - water feature wetted perimeter - length : xr.DataArray of floats - water feature length per cell - dx : xr.DataArray of floats - cellsize in x - dy : xr.DataArray of floats - cellsize in y - - Returns - ------- - c_leakage: xr.DataArray of floats - Hydraulic resistance of water features corrected for intra-cell - hydraulic resistance and surface water interaction. - """ - - def f(x): - """ - two x times cotangens hyperbolicus of x - """ - # Avoid overflow for large x values - # 20 is safe for 32 bit floats - x = x.where(~(x > 20.0), other=20.0) - exp2x = np.exp(2.0 * x) - return x * (exp2x + 1) / (exp2x - 1.0) - - # TODO: raise ValueError when cells are far from square? - # Since method of average ditch distance will not work well - # Compute geometry of ditches within cells - cell_area = abs(dy * dx) - wetted_area = length * B - fraction_wetted = wetted_area / cell_area - # Compute average distance between ditches - L = (cell_area - wetted_area) / length - - # Compute total resistance to aquifer - c1_prime = c1 + (D / kv) - # Compute influence for the part below the ditch - labda_B = np.sqrt((kh * D * c1_prime * c0) / (c1_prime + c0)) - # ... and the field - labda_L = np.sqrt(c1_prime * kh * D) - - x_B = B / (2.0 * labda_B) - x_L = L / (2.0 * labda_L) - - # Compute feeding resistance - c_rad = c_radial(L, kv, kh, B, D) - c_L = (c0 + c1_prime) * f(x_L) + (c0 * L / B) * f(x_B) - c_B = (c1_prime + c0) / (c_L - c0 * L / B) * c_L - # total feeding resistance equals the harmonic mean of resistances below - # ditch (B) and field (L) plus the radical resistance - # Can also be regarded as area weighted sum of conductances of B and L - c_total = 1.0 / (fraction_wetted / c_B + (1.0 - fraction_wetted) / c_L) + c_rad - # Subtract aquifer and aquitard resistance from feeding resistance - c = c_total - c1_prime - # Replace areas where cells are fully covered by water - c = c.where(~(L == 0.0), other=c0) - return c + warnings.warn(WARNING_MESSAGE, DeprecationWarning) + return _c_leakage(kh, kv, D, c0, c1, B, length, dx, dy) diff --git a/imod/prepare/topsystem/resistance.py b/imod/prepare/topsystem/resistance.py new file mode 100644 index 000000000..5e3ff4298 --- /dev/null +++ b/imod/prepare/topsystem/resistance.py @@ -0,0 +1,109 @@ +import numpy as np + + +def c_radial(L, kh, kv, B, D): + """ + Ernst's radial resistance term to a drain. + + Parameters + ---------- + L : xr.DataArray of floats + distance between water features + kh : xr.DataArray of floats + horizontal conductivity + kv : xr.DataArray of floats + vertical conductivity + B : xr.DataArray of floats + water feature wetted perimeter + D : xr.DataArray of floats + saturated thickness of the top system + + Returns + ------- + radial_c : xr.DataArray + Ernst's radial resistance for a drain + """ + # Take anisotropy into account fully + c = ( + L + / (np.pi * np.sqrt(kh * kv)) + * np.log((4.0 * D) / (np.pi * B) * np.sqrt(kh / kv)) + ) + c = c.where(~(c < 0.0), other=0.0) + return c + + +def c_leakage(kh, kv, D, c0, c1, B, length, dx, dy): + """ + Computes the phreatic leakage resistance. + + Parameters + ---------- + kh : xr.DataArray of floats + horizontal conductivity of phreatic aquifer + kv : xr.DataArray of floats + vertical conductivity of phreatic aquifer + c0 : xr.DataArray of floats + hydraulic bed resistance of water feature + c1 : xr.DataArray of floats + hydraulic resistance of the first aquitard + D : xr.DataArray of floats + saturated thickness of the top system + B : xr.DataArray of floats + water feature wetted perimeter + length : xr.DataArray of floats + water feature length per cell + dx : xr.DataArray of floats + cellsize in x + dy : xr.DataArray of floats + cellsize in y + + Returns + ------- + c_leakage: xr.DataArray of floats + Hydraulic resistance of water features corrected for intra-cell + hydraulic resistance and surface water interaction. + """ + + def f(x): + """ + two x times cotangens hyperbolicus of x + """ + # Avoid overflow for large x values + # 20 is safe for 32 bit floats + x = x.where(~(x > 20.0), other=20.0) + exp2x = np.exp(2.0 * x) + return x * (exp2x + 1) / (exp2x - 1.0) + + # TODO: raise ValueError when cells are far from square? + # Since method of average ditch distance will not work well + # Compute geometry of ditches within cells + cell_area = abs(dy * dx) + wetted_area = length * B + fraction_wetted = wetted_area / cell_area + # Compute average distance between ditches + L = (cell_area - wetted_area) / length + + # Compute total resistance to aquifer + c1_prime = c1 + (D / kv) + # Compute influence for the part below the ditch + labda_B = np.sqrt((kh * D * c1_prime * c0) / (c1_prime + c0)) + # ... and the field + labda_L = np.sqrt(c1_prime * kh * D) + + x_B = B / (2.0 * labda_B) + x_L = L / (2.0 * labda_L) + + # Compute feeding resistance + c_rad = c_radial(L, kv, kh, B, D) + c_L = (c0 + c1_prime) * f(x_L) + (c0 * L / B) * f(x_B) + c_B = (c1_prime + c0) / (c_L - c0 * L / B) * c_L + # total feeding resistance equals the harmonic mean of resistances below + # ditch (B) and field (L) plus the radical resistance + # Can also be regarded as area weighted sum of conductances of B and L + c_total = 1.0 / (fraction_wetted / c_B + (1.0 - fraction_wetted) / c_L) + c_rad + # Subtract aquifer and aquitard resistance from feeding resistance + c = c_total - c1_prime + # Replace areas where cells are fully covered by water + c = c.where(~(L == 0.0), other=c0) + return c From 87c8de0f718b0e7cfa8b4442f03cef7e720a65de Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 16:29:50 +0200 Subject: [PATCH 07/35] Import at module level --- imod/prepare/topsystem/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 imod/prepare/topsystem/__init__.py diff --git a/imod/prepare/topsystem/__init__.py b/imod/prepare/topsystem/__init__.py new file mode 100644 index 000000000..6e13f2801 --- /dev/null +++ b/imod/prepare/topsystem/__init__.py @@ -0,0 +1,8 @@ +from imod.prepare.topsystem.allocation import ( + ALLOCATION_OPTION, + allocate_drain_cells, + allocate_ghb_cells, + allocate_rch_cells, + allocate_river_cells, +) +from imod.prepare.topsystem.resistance import c_leakage, c_radial From e44b7daa499d9da96d3316051d38df533c313fdf Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 16:35:27 +0200 Subject: [PATCH 08/35] Add topsystem functions to public API --- docs/api/prepare.rst | 8 ++++++++ imod/prepare/__init__.py | 9 +++++++++ 2 files changed, 17 insertions(+) diff --git a/docs/api/prepare.rst b/docs/api/prepare.rst index 236c1ee4b..870fe4655 100644 --- a/docs/api/prepare.rst +++ b/docs/api/prepare.rst @@ -31,3 +31,11 @@ Prepare model input get_lower_active_layer_number get_upper_active_grid_cells get_upper_active_layer_number + + ALLOCATION_OPTION + allocate_drain_cells + allocate_ghb_cells + allocate_rch_cells + allocate_river_cells + c_leakage + c_radial \ No newline at end of file diff --git a/imod/prepare/__init__.py b/imod/prepare/__init__.py index e648576e8..600c24917 100644 --- a/imod/prepare/__init__.py +++ b/imod/prepare/__init__.py @@ -34,5 +34,14 @@ zonal_aggregate_polygons, zonal_aggregate_raster, ) +from imod.prepare.topsystem import ( + ALLOCATION_OPTION, + allocate_drain_cells, + allocate_ghb_cells, + allocate_rch_cells, + allocate_river_cells, + c_leakage, + c_radial, +) from imod.prepare.voxelize import Voxelizer from imod.prepare.wells import assign_wells From 14ecd58ffa9780d1571ead04b4effdbbdcc9754f Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 3 Apr 2024 17:23:38 +0200 Subject: [PATCH 09/35] Add first test for allocation --- imod/tests/test_topsystem/test_allocation.py | 38 ++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 imod/tests/test_topsystem/test_allocation.py diff --git a/imod/tests/test_topsystem/test_allocation.py b/imod/tests/test_topsystem/test_allocation.py new file mode 100644 index 000000000..002a555cb --- /dev/null +++ b/imod/tests/test_topsystem/test_allocation.py @@ -0,0 +1,38 @@ +import numpy as np +from pytest_cases import parametrize_with_cases + +from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_river_cells +from imod.typing.grid import zeros_like + + +class RiverCases: + def case_structured(basic_dis): + ibound, top, bottom = basic_dis + top = top.sel(layer=1) + elevation = zeros_like(ibound.sel(layer=1)) + stage = elevation - 2.5 + bottom_elevation = elevation - 10.0 + active = ibound == 1 + return active, top, bottom, stage, bottom_elevation + + def case_unstructured(basic_unstructured_dis): + ibound, top, bottom = basic_unstructured_dis + elevation = zeros_like(ibound.sel(layer=1)) + stage = elevation - 2.5 + bottom_elevation = elevation - 10.0 + active = ibound == 1 + return active, top, bottom, stage, bottom_elevation + + +class AllocationOptionCases: + def case_stage_to_riv_bot(): + option = ALLOCATION_OPTION.stage_to_riv_bot + expected = [True, True, False] + + return option, expected + + +@parametrize_with_cases(argnames=["active,top,bottom,stage,bottom_elevation", "option,expected"], cases = [RiverCases, AllocationOptionCases]) +def test_riv_allocation(active, top, bottom, stage, bottom_elevation, option, expected): + actual = allocate_river_cells(option, active, top, bottom, stage, bottom_elevation) + np.testing.assert_equal(actual.values[:, 0, 0] == expected) \ No newline at end of file From 30902f2a77c29bf97d8e173305689fa72b0ea39b Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 10:30:36 +0200 Subject: [PATCH 10/35] Generalize utility functions and work around circular imports --- imod/mf6/utilities/grid.py | 12 +------ imod/mf6/wel.py | 2 +- imod/prepare/layer.py | 12 +++++++ imod/prepare/spatial.py | 30 ++++------------ imod/prepare/topsystem/allocation.py | 2 +- imod/typing/grid.py | 4 +-- imod/util/spatial.py | 53 ++++++++++++++++++++++++++-- 7 files changed, 74 insertions(+), 41 deletions(-) diff --git a/imod/mf6/utilities/grid.py b/imod/mf6/utilities/grid.py index 0f6718ac9..d5d5ac8be 100644 --- a/imod/mf6/utilities/grid.py +++ b/imod/mf6/utilities/grid.py @@ -7,6 +7,7 @@ import xugrid as xu import imod +from imod.prepare.layer import create_layered_top from imod.typing import GridDataArray from imod.typing.grid import zeros_like from imod.util.spatial import spatial_reference @@ -48,17 +49,6 @@ def broadcast_to_full_domain( return top, bottom -def create_layered_top(bottom: GridDataArray, top: GridDataArray) -> GridDataArray: - """ - Create a top array with layers from a single top array and a full bottom array - """ - new_top = zeros_like(bottom) - new_top[0] = top - new_top[1:] = bottom[0:-1].values - - return new_top - - def to_cell_idx(idomain: xr.DataArray) -> xr.DataArray: """ Assigns an unique index to each cell in the domain diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index 655f44069..5c3d9fc43 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -21,11 +21,11 @@ from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.package import Package from imod.mf6.utilities.dataset import remove_inactive -from imod.mf6.utilities.grid import create_layered_top from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import validation_pkg_error_message from imod.mf6.write_context import WriteContext from imod.prepare import assign_wells +from imod.prepare.layer import create_layered_top from imod.schemata import ( AnyNoDataSchema, DTypeSchema, diff --git a/imod/prepare/layer.py b/imod/prepare/layer.py index a1586c5ce..9891d6988 100644 --- a/imod/prepare/layer.py +++ b/imod/prepare/layer.py @@ -3,6 +3,7 @@ """ from imod.typing import GridDataArray +from imod.typing.grid import zeros_like def get_upper_active_layer_number(active: GridDataArray) -> GridDataArray: @@ -67,3 +68,14 @@ def get_lower_active_grid_cells(active: GridDataArray) -> GridDataArray: layer = active.coords["layer"] lower_active_layer = get_lower_active_layer_number(active) return layer == lower_active_layer + + +def create_layered_top(bottom: GridDataArray, top: GridDataArray) -> GridDataArray: + """ + Create a top array with layers from a single top array and a full bottom array + """ + new_top = zeros_like(bottom) + new_top[0] = top + new_top[1:] = bottom[0:-1].values + + return new_top diff --git a/imod/prepare/spatial.py b/imod/prepare/spatial.py index 1ea005747..49b56af97 100644 --- a/imod/prepare/spatial.py +++ b/imod/prepare/spatial.py @@ -1,5 +1,5 @@ import pathlib -from typing import Callable, Union +from typing import TYPE_CHECKING, Callable, Union import dask import numba @@ -11,6 +11,7 @@ import imod from imod.prepare import common, pcg from imod.util.imports import MissingOptionalModule +from imod.util.spatial import _polygonize # since rasterio is a big dependency that is sometimes hard to install # and not always required, we made this an optional dependency @@ -21,6 +22,8 @@ except ImportError: rasterio = MissingOptionalModule("rasterio") +if TYPE_CHECKING: + import geopandas as gpd def round_extent(extent, cellsize): """Increases the extent until all sides lie on a coordinate @@ -252,7 +255,7 @@ def rasterize(geodataframe, like, column=None, fill=np.nan, **kwargs): return xr.DataArray(raster, like.coords, like.dims) -def polygonize(da): +def polygonize(da: xr.DataArray) -> "gpd.GeoDataFrame": """ Polygonize a 2D-DataArray into a GeoDataFrame of polygons. @@ -264,28 +267,7 @@ def polygonize(da): ------- polygonized : geopandas.GeoDataFrame """ - import geopandas as gpd - import shapely.geometry as sg - - if da.dims != ("y", "x"): - raise ValueError('Dimensions must be ("y", "x")') - - values = da.values - if values.dtype == np.float64: - values = values.astype(np.float32) - - transform = imod.util.spatial.transform(da) - shapes = rasterio.features.shapes(values, transform=transform) - - geometries = [] - colvalues = [] - for geom, colval in shapes: - geometries.append(sg.Polygon(geom["coordinates"][0])) - colvalues.append(colval) - - gdf = gpd.GeoDataFrame({"value": colvalues, "geometry": geometries}) - gdf.crs = da.attrs.get("crs") - return gdf + return _polygonize(da) def _handle_dtype(dtype, nodata): diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 493ded664..1bb91ce42 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -4,8 +4,8 @@ from enum import Enum -from imod.mf6.utilities.grid import create_layered_top from imod.prepare.layer import ( + create_layered_top, get_upper_active_grid_cells, get_upper_active_layer_number, ) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index a59b0d4d1..123001f80 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -7,8 +7,8 @@ import xugrid as xu from fastcore.dispatch import typedispatch -from imod.prepare import polygonize from imod.typing import GridDataArray, GridDataset, structured +from imod.util.spatial import _polygonize @typedispatch @@ -228,7 +228,7 @@ def merge_with_dictionary( def bounding_polygon(active: xr.DataArray): """Return bounding polygon of active cells""" to_polygonize = active.where(active, other=np.nan) - polygons_gdf = polygonize(to_polygonize) + polygons_gdf = _polygonize(to_polygonize) # Filter polygons with inactive values (NaN) is_active_polygon = polygons_gdf["value"] == 1.0 return polygons_gdf.loc[is_active_polygon] diff --git a/imod/util/spatial.py b/imod/util/spatial.py index 6a00564bd..6d3ddbb6b 100644 --- a/imod/util/spatial.py +++ b/imod/util/spatial.py @@ -7,7 +7,7 @@ import collections import re -from typing import Any, Dict, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Sequence, Tuple, Union import affine import numpy as np @@ -16,7 +16,28 @@ import xugrid as xu from imod.typing import FloatArray, GridDataset, IntArray - +from imod.util.imports import MissingOptionalModule + +# since rasterio, shapely, and geopandas are a big dependencies that are +# sometimes hard to install and not always required, we made this an optional +# dependency +try: + import rasterio +except ImportError: + rasterio = MissingOptionalModule("rasterio") + +try: + import shapely +except ImportError: + shapely = MissingOptionalModule("shapely") + +if TYPE_CHECKING: + import geopandas as gpd +else: + try: + import geopandas as gpd + except ImportError: + gpd = MissingOptionalModule("geopandas") def _xycoords(bounds, cellsizes) -> Dict[str, Any]: """Based on bounds and cellsizes, construct coords with spatial information""" @@ -612,3 +633,31 @@ def is_divisor(numerator: FloatArray, denominator: float) -> bool: denominator = abs(denominator) remainder = np.abs(numerator) % denominator return (np.isclose(remainder, 0.0) | np.isclose(remainder, denominator)).all() + + +def _polygonize(da: xr.DataArray) -> gpd.GeoDataFrame: + """ + Polygonize a 2D-DataArray into a GeoDataFrame of polygons. + + Private method located in util.spatial to work around circular imports. + """ + + if da.dims != ("y", "x"): + raise ValueError('Dimensions must be ("y", "x")') + + values = da.values + if values.dtype == np.float64: + values = values.astype(np.float32) + + affine_transform = transform(da) + shapes = rasterio.features.shapes(values, transform=affine_transform) + + geometries = [] + colvalues = [] + for geom, colval in shapes: + geometries.append(shapely.geometry.Polygon(geom["coordinates"][0])) + colvalues.append(colval) + + gdf = gpd.GeoDataFrame({"value": colvalues, "geometry": geometries}) + gdf.crs = da.attrs.get("crs") + return gdf From 1ae56a71a96d219e21a97695304c979b773f20c6 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:19:26 +0200 Subject: [PATCH 11/35] Add utility functions to preserve gridtype and to enforce dim order --- imod/typing/grid.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index 123001f80..34553c48b 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -336,3 +336,41 @@ def get_grid_geometry_hash(grid: xu.UgridDataArray) -> int: @typedispatch def get_grid_geometry_hash(grid: object) -> int: raise ValueError("get_grid_geometry_hash not supported for this object.") + + +@typedispatch +def enforce_dim_order(grid: xr.DataArray) -> xr.DataArray: + """Enforce dimension order to iMOD Python standard""" + return grid.transpose("species", "time", "layer", "y", "x", missing_dims="ignore") + + +@typedispatch +def enforce_dim_order(grid: xu.UgridDataArray) -> xu.UgridDataArray: + """Enforce dimension order to iMOD Python standard""" + face_dimension = grid.ugrid.grid.face_dimension + return grid.transpose( + "species", "time", "layer", face_dimension, missing_dims="ignore" + ) + + +def preserve_gridtype(func): + """ "Decorator to preserve gridtype""" + + def decorator(*args, **kwargs): + unstructured = False + grid = None + for arg in args: + if is_unstructured(arg): + unstructured = True + grid = arg.ugrid.grid + + x = func(*args, **kwargs) + + if unstructured: + # Multiple grids returned + if isinstance(x, tuple): + return tuple(xu.UgridDataArray(i, grid) for i in x) + return xu.UgridDataArray(x, grid) + return x + + return decorator From 54337c537a65df33547a2dc5d436eb295cdca45b Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:20:12 +0200 Subject: [PATCH 12/35] Add decorator to enforce dim order --- imod/util/dims.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 imod/util/dims.py diff --git a/imod/util/dims.py b/imod/util/dims.py new file mode 100644 index 000000000..b1becd2c5 --- /dev/null +++ b/imod/util/dims.py @@ -0,0 +1,14 @@ +from imod.typing.grid import enforce_dim_order + + +def enforced_dim_order(func): + """ "Decorator to enforce dimension order after function call""" + + def decorator(*args, **kwargs): + x = func(*args, **kwargs) + # Multiple grids returned + if isinstance(x, tuple): + return tuple(enforce_dim_order(i) for i in x) + return enforce_dim_order(x) + + return decorator From 6cced53d77edee09a1615195227f05f9cea2a652 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:21:30 +0200 Subject: [PATCH 13/35] Apply decorator to preserve gridtype --- imod/prepare/layer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/imod/prepare/layer.py b/imod/prepare/layer.py index 9891d6988..148b88ddf 100644 --- a/imod/prepare/layer.py +++ b/imod/prepare/layer.py @@ -3,9 +3,10 @@ """ from imod.typing import GridDataArray -from imod.typing.grid import zeros_like +from imod.typing.grid import preserve_gridtype, zeros_like +@preserve_gridtype def get_upper_active_layer_number(active: GridDataArray) -> GridDataArray: """ Returns planar grid of integers with the layer number of the upper most @@ -23,6 +24,7 @@ def get_upper_active_layer_number(active: GridDataArray) -> GridDataArray: return layer.where(active, nodata).min("layer") +@preserve_gridtype def get_upper_active_grid_cells(active: GridDataArray) -> GridDataArray: """ Returns grid of booleans designating location of the uppermost active grid @@ -38,6 +40,7 @@ def get_upper_active_grid_cells(active: GridDataArray) -> GridDataArray: return layer == upper_active_layer +@preserve_gridtype def get_lower_active_layer_number(active: GridDataArray) -> GridDataArray: """ Returns two-dimensional grid of integers with the layer number of the lower @@ -55,6 +58,7 @@ def get_lower_active_layer_number(active: GridDataArray) -> GridDataArray: return layer.where(active, nodata).max("layer") +@preserve_gridtype def get_lower_active_grid_cells(active: GridDataArray) -> GridDataArray: """ Returns grid of booleans designating location of the lowermost active grid From e248274a5888c22d31c82ebe62470abc9696aa31 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:22:43 +0200 Subject: [PATCH 14/35] format --- imod/prepare/spatial.py | 1 + imod/util/spatial.py | 1 + 2 files changed, 2 insertions(+) diff --git a/imod/prepare/spatial.py b/imod/prepare/spatial.py index 49b56af97..74387c361 100644 --- a/imod/prepare/spatial.py +++ b/imod/prepare/spatial.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: import geopandas as gpd + def round_extent(extent, cellsize): """Increases the extent until all sides lie on a coordinate divisible by cellsize.""" diff --git a/imod/util/spatial.py b/imod/util/spatial.py index 6d3ddbb6b..eb32382e5 100644 --- a/imod/util/spatial.py +++ b/imod/util/spatial.py @@ -39,6 +39,7 @@ except ImportError: gpd = MissingOptionalModule("geopandas") + def _xycoords(bounds, cellsizes) -> Dict[str, Any]: """Based on bounds and cellsizes, construct coords with spatial information""" # unpack tuples From 1ec16d9e7b9e8876f1b55fa4b22b861623e7c83a Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:23:12 +0200 Subject: [PATCH 15/35] Add tests for allocation functions --- imod/tests/test_topsystem/test_allocation.py | 74 +++++++++++++++++--- 1 file changed, 63 insertions(+), 11 deletions(-) diff --git a/imod/tests/test_topsystem/test_allocation.py b/imod/tests/test_topsystem/test_allocation.py index 002a555cb..b97a18039 100644 --- a/imod/tests/test_topsystem/test_allocation.py +++ b/imod/tests/test_topsystem/test_allocation.py @@ -2,37 +2,89 @@ from pytest_cases import parametrize_with_cases from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_river_cells -from imod.typing.grid import zeros_like +from imod.typing import GridDataArray +from imod.typing.grid import is_unstructured, zeros_like + + +def take_first_planar_cell(grid: GridDataArray): + if is_unstructured(grid): + return grid.values[:, 0] + else: + return grid.values[:, 0, 0] class RiverCases: - def case_structured(basic_dis): + def case_structured(self, basic_dis): ibound, top, bottom = basic_dis top = top.sel(layer=1) elevation = zeros_like(ibound.sel(layer=1)) - stage = elevation - 2.5 + stage = elevation - 7.5 bottom_elevation = elevation - 10.0 active = ibound == 1 return active, top, bottom, stage, bottom_elevation - def case_unstructured(basic_unstructured_dis): + def case_unstructured(self, basic_unstructured_dis): ibound, top, bottom = basic_unstructured_dis elevation = zeros_like(ibound.sel(layer=1)) - stage = elevation - 2.5 + stage = elevation - 7.5 bottom_elevation = elevation - 10.0 active = ibound == 1 return active, top, bottom, stage, bottom_elevation class AllocationOptionCases: - def case_stage_to_riv_bot(): + def case_stage_to_riv_bot(self): option = ALLOCATION_OPTION.stage_to_riv_bot + expected = [False, True, False] + + return option, expected, None + + def case_first_active_to_riv_bot(self): + option = ALLOCATION_OPTION.first_active_to_riv_bot expected = [True, True, False] - return option, expected + return option, expected, None + + def case_first_active_to_riv_bot__drn(self): + option = ALLOCATION_OPTION.first_active_to_riv_bot__drn + expected = [False, True, False] + expected__drn = [True, False, False] + + return option, expected, expected__drn + + def case_at_elevation(self): + option = ALLOCATION_OPTION.at_elevation + expected = [False, True, False] + + return option, expected, None + + def case_at_first_active(self): + option = ALLOCATION_OPTION.at_first_active + expected = [True, False, False] + + return option, expected, None + + +@parametrize_with_cases( + argnames="active,top,bottom,stage,bottom_elevation", + cases=RiverCases, +) +@parametrize_with_cases( + argnames="option,expected_riv, expected_drn", cases=AllocationOptionCases +) +def test_riv_allocation( + active, top, bottom, stage, bottom_elevation, option, expected_riv, expected_drn +): + actual_riv_da, actual_drn_da = allocate_river_cells( + option, active, top, bottom, stage, bottom_elevation + ) + + actual_riv = take_first_planar_cell(actual_riv_da) + if actual_drn_da is None: + actual_drn = actual_drn_da + else: + actual_drn = take_first_planar_cell(actual_drn_da) -@parametrize_with_cases(argnames=["active,top,bottom,stage,bottom_elevation", "option,expected"], cases = [RiverCases, AllocationOptionCases]) -def test_riv_allocation(active, top, bottom, stage, bottom_elevation, option, expected): - actual = allocate_river_cells(option, active, top, bottom, stage, bottom_elevation) - np.testing.assert_equal(actual.values[:, 0, 0] == expected) \ No newline at end of file + np.testing.assert_equal(actual_riv, expected_riv) + np.testing.assert_equal(actual_drn, expected_drn) From c9181771d5ddddef2adda6cd6d13b741281d2a25 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:24:44 +0200 Subject: [PATCH 16/35] Fix bugs where cells were incorrectly allocated when looking at riv stage, dim order and/or grid type was not preserved. --- imod/prepare/topsystem/allocation.py | 69 +++++++++++++++++++--------- 1 file changed, 47 insertions(+), 22 deletions(-) diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 1bb91ce42..3f89b9550 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -3,6 +3,7 @@ """ from enum import Enum +from typing import Optional from imod.prepare.layer import ( create_layered_top, @@ -11,6 +12,7 @@ ) from imod.schemata import DimsSchema from imod.typing import GridDataArray +from imod.util.dims import enforced_dim_order class ALLOCATION_OPTION(Enum): @@ -41,7 +43,7 @@ def allocate_river_cells( bottom: GridDataArray, stage: GridDataArray, bottom_elevation: GridDataArray, -): +) -> tuple[GridDataArray, Optional[GridDataArray]]: match allocation_option: case ALLOCATION_OPTION.stage_to_riv_bot: return _allocate_cells__stage_to_riv_bot( @@ -49,7 +51,7 @@ def allocate_river_cells( ) case ALLOCATION_OPTION.first_active_to_riv_bot: return _allocate_cells__first_active_to_riv_bot( - active, bottom, bottom_elevation + active, top, bottom, bottom_elevation ) case ALLOCATION_OPTION.first_active_to_riv_bot__drn: return _allocate_cells__first_active_to_riv_bot__drn( @@ -67,12 +69,12 @@ def allocate_drain_cells( top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray, -): +) -> GridDataArray: match allocation_option: case ALLOCATION_OPTION.at_elevation: - return _allocate_cells__at_elevation(top, bottom, elevation) + return _allocate_cells__at_elevation(top, bottom, elevation)[0] case ALLOCATION_OPTION.at_first_active: - return _allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active)[0] case _: raise ValueError( "Received incompatible setting for drains, only" @@ -88,12 +90,12 @@ def allocate_ghb_cells( top: GridDataArray, bottom: GridDataArray, head: GridDataArray, -): +) -> GridDataArray: match allocation_option: case ALLOCATION_OPTION.at_elevation: - return _allocate_cells__at_elevation(top, bottom, head) + return _allocate_cells__at_elevation(top, bottom, head)[0] case ALLOCATION_OPTION.at_first_active: - return _allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active)[0] case _: raise ValueError( "Received incompatible setting for drains, only" @@ -106,7 +108,7 @@ def allocate_ghb_cells( def allocate_rch_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray, -): +) -> GridDataArray: match allocation_option: case ALLOCATION_OPTION.at_first_active: return _allocate_cells__at_first_active(active) @@ -122,12 +124,13 @@ def _is_layered(grid: GridDataArray): return "layer" in grid.sizes and grid.sizes["layer"] > 1 +@enforced_dim_order def _allocate_cells__stage_to_riv_bot( top: GridDataArray, bottom: GridDataArray, stage: GridDataArray, bottom_elevation: GridDataArray, -): +) -> tuple[GridDataArray, None]: """ Allocate cells inbetween river stage and river bottom_elevation. Compared to iMOD5.6, this is similar to setting IDEFFLAYER=0 in the RUNFILE function. @@ -159,12 +162,18 @@ def _allocate_cells__stage_to_riv_bot( else: top_layered = create_layered_top(bottom, top) - return (stage <= top_layered) & (bottom_elevation >= bottom) + riv_cells = (stage > bottom) & (bottom_elevation < top_layered) + + return riv_cells, None +@enforced_dim_order def _allocate_cells__first_active_to_riv_bot( - active: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray -): + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + bottom_elevation: GridDataArray, +) -> tuple[GridDataArray, None]: """ Allocate cells inbetween first active layer and river bottom elevation. Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE @@ -176,6 +185,8 @@ def _allocate_cells__first_active_to_riv_bot( ---------- active: GridDataArray active model cells + top: GridDataArray + top of model layers bottom: GridDataArray bottom of model layers bottom_elevation: GridDataArray @@ -191,16 +202,24 @@ def _allocate_cells__first_active_to_riv_bot( upper_active_layer = get_upper_active_layer_number(active) layer = active.coords["layer"] - return (layer >= upper_active_layer) & (bottom_elevation >= bottom) + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + riv_cells = (upper_active_layer <= layer) & (bottom_elevation < top_layered) + return riv_cells, None + + +@enforced_dim_order def _allocate_cells__first_active_to_riv_bot__drn( active: GridDataArray, top: GridDataArray, bottom: GridDataArray, stage: GridDataArray, bottom_elevation: GridDataArray, -): +) -> tuple[GridDataArray, GridDataArray]: """ Allocate cells inbetween first active layer and river bottom elevation. Cells above river stage are deactivated and returned as drn cells. Compared @@ -240,17 +259,18 @@ def _allocate_cells__first_active_to_riv_bot__drn( upper_active_layer = get_upper_active_layer_number(active) layer = active.coords["layer"] - drn_cells = (layer >= upper_active_layer) & (stage <= top_layered) - riv_cells = (layer >= upper_active_layer) & ( - bottom_elevation >= bottom + drn_cells = (upper_active_layer <= layer) & (bottom >= stage) + riv_cells = ( + (upper_active_layer <= layer) & (bottom_elevation < top_layered) ) != drn_cells return riv_cells, drn_cells +@enforced_dim_order def _allocate_cells__at_elevation( top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray -): +) -> tuple[GridDataArray, None]: """ Allocate cells in river bottom elevation layer. Compared to iMOD5.6, this is similar to setting IDEFFLAYER=2 in the RUNFILE function. @@ -279,10 +299,15 @@ def _allocate_cells__at_elevation( else: top_layered = create_layered_top(bottom, top) - return (elevation < top_layered) & (elevation >= bottom) + riv_cells = (elevation < top_layered) & (elevation >= bottom) + return riv_cells, None -def _allocate_cells__at_first_active(active: GridDataArray): + +@enforced_dim_order +def _allocate_cells__at_first_active( + active: GridDataArray, +) -> tuple[GridDataArray, None]: """ Allocate cells inbetween first active layer and river bottom elevation. Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE @@ -301,4 +326,4 @@ def _allocate_cells__at_first_active(active: GridDataArray): River cells """ - return get_upper_active_grid_cells(active) + return get_upper_active_grid_cells(active), None From f9caec5d2da709d3ce5be275bde05a50a74b34c8 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:32:06 +0200 Subject: [PATCH 17/35] Add docstring with explanation --- imod/typing/grid.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index 34553c48b..0866d6f75 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -354,7 +354,16 @@ def enforce_dim_order(grid: xu.UgridDataArray) -> xu.UgridDataArray: def preserve_gridtype(func): - """ "Decorator to preserve gridtype""" + """ + Decorator to preserve gridtype, this is to work around the following xugrid + behavior: + + UgridDataArray * DataArray -> UgridDataArray + DataArray * UgridDataArray -> DataArray + + with this decorator: + DataArray * UgridDataArray -> UgridDataArray + """ def decorator(*args, **kwargs): unstructured = False From 16e336941175db940596e93b46ea3f7a8ece5c56 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:33:46 +0200 Subject: [PATCH 18/35] Tidy up docstring --- imod/typing/grid.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index 0866d6f75..ab69b289c 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -358,11 +358,13 @@ def preserve_gridtype(func): Decorator to preserve gridtype, this is to work around the following xugrid behavior: - UgridDataArray * DataArray -> UgridDataArray - DataArray * UgridDataArray -> DataArray + >>> UgridDataArray() * DataArray() -> UgridDataArray + >>> DataArray() * UgridDataArray() -> DataArray with this decorator: - DataArray * UgridDataArray -> UgridDataArray + + >>> UgridDataArray() * DataArray() -> UgridDataArray + >>> DataArray() * UgridDataArray() -> UgridDataArray """ def decorator(*args, **kwargs): From 9a93004caebfea839f8048ad531641398ba1fb32 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:36:13 +0200 Subject: [PATCH 19/35] Tidy up docstring --- imod/util/dims.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/util/dims.py b/imod/util/dims.py index b1becd2c5..7db0be77e 100644 --- a/imod/util/dims.py +++ b/imod/util/dims.py @@ -2,7 +2,7 @@ def enforced_dim_order(func): - """ "Decorator to enforce dimension order after function call""" + """Decorator to enforce dimension order after function call""" def decorator(*args, **kwargs): x = func(*args, **kwargs) From 44cc5db5499521097f2c49122d0d30da3e3ecf8a Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 8 Apr 2024 18:43:37 +0200 Subject: [PATCH 20/35] format --- imod/typing/grid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index ab69b289c..1b59627ac 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -362,7 +362,7 @@ def preserve_gridtype(func): >>> DataArray() * UgridDataArray() -> DataArray with this decorator: - + >>> UgridDataArray() * DataArray() -> UgridDataArray >>> DataArray() * UgridDataArray() -> UgridDataArray """ From 8ba80fe6f5da1567d16d0f2d0ed9a11f1ef33b95 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 10:23:49 +0200 Subject: [PATCH 21/35] Add tests for allocation drn, rch, ghb cells --- imod/prepare/topsystem/allocation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 3f89b9550..7c03c633e 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -111,7 +111,7 @@ def allocate_rch_cells( ) -> GridDataArray: match allocation_option: case ALLOCATION_OPTION.at_first_active: - return _allocate_cells__at_first_active(active) + return _allocate_cells__at_first_active(active)[0] case _: raise ValueError( "Received incompatible setting for drains, only" From ea613d664144c09d726f4dfd1a7e3d2da629993d Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 10:25:19 +0200 Subject: [PATCH 22/35] Return grid instead of tuple for rch --- imod/tests/test_topsystem/test_allocation.py | 135 ++++++++++++++++++- 1 file changed, 132 insertions(+), 3 deletions(-) diff --git a/imod/tests/test_topsystem/test_allocation.py b/imod/tests/test_topsystem/test_allocation.py index b97a18039..537f8fb61 100644 --- a/imod/tests/test_topsystem/test_allocation.py +++ b/imod/tests/test_topsystem/test_allocation.py @@ -1,7 +1,13 @@ import numpy as np from pytest_cases import parametrize_with_cases -from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_river_cells +from imod.prepare.topsystem import ( + ALLOCATION_OPTION, + allocate_drain_cells, + allocate_ghb_cells, + allocate_rch_cells, + allocate_river_cells, +) from imod.typing import GridDataArray from imod.typing.grid import is_unstructured, zeros_like @@ -32,7 +38,53 @@ def case_unstructured(self, basic_unstructured_dis): return active, top, bottom, stage, bottom_elevation -class AllocationOptionCases: +class DrainCases: + def case_structured(self, basic_dis): + ibound, top, bottom = basic_dis + top = top.sel(layer=1) + elevation = zeros_like(ibound.sel(layer=1)) + drain_elevation = elevation - 7.5 + active = ibound == 1 + return active, top, bottom, drain_elevation + + def case_unstructured(self, basic_unstructured_dis): + ibound, top, bottom = basic_unstructured_dis + elevation = zeros_like(ibound.sel(layer=1)) + drain_elevation = elevation - 7.5 + active = ibound == 1 + return active, top, bottom, drain_elevation + + +class GeneralHeadBoundaryCases: + def case_structured(self, basic_dis): + ibound, top, bottom = basic_dis + top = top.sel(layer=1) + elevation = zeros_like(ibound.sel(layer=1)) + head = elevation - 7.5 + active = ibound == 1 + return active, top, bottom, head + + def case_unstructured(self, basic_unstructured_dis): + ibound, top, bottom = basic_unstructured_dis + elevation = zeros_like(ibound.sel(layer=1)) + head = elevation - 7.5 + active = ibound == 1 + return active, top, bottom, head + + +class RechargeCases: + def case_structured(self, basic_dis): + ibound, _, _ = basic_dis + active = ibound == 1 + return active + + def case_unstructured(self, basic_unstructured_dis): + ibound, _, _ = basic_unstructured_dis + active = ibound == 1 + return active + + +class AllocationOptionRiverCases: def case_stage_to_riv_bot(self): option = ALLOCATION_OPTION.stage_to_riv_bot expected = [False, True, False] @@ -65,12 +117,48 @@ def case_at_first_active(self): return option, expected, None +class AllocationOptionDrainCases: + def case_at_elevation(self): + option = ALLOCATION_OPTION.at_elevation + expected = [False, True, False] + + return option, expected + + def case_at_first_active(self): + option = ALLOCATION_OPTION.at_first_active + expected = [True, False, False] + + return option, expected + + +class AllocationOptionGeneralHeadCases: + def case_at_elevation(self): + option = ALLOCATION_OPTION.at_elevation + expected = [False, True, False] + + return option, expected + + def case_at_first_active(self): + option = ALLOCATION_OPTION.at_first_active + expected = [True, False, False] + + return option, expected + + +class AllocationOptionRechargeCases: + def case_at_first_active(self): + option = ALLOCATION_OPTION.at_first_active + expected = [True, False, False] + + return option, expected + + @parametrize_with_cases( argnames="active,top,bottom,stage,bottom_elevation", cases=RiverCases, ) @parametrize_with_cases( - argnames="option,expected_riv, expected_drn", cases=AllocationOptionCases + argnames="option,expected_riv,expected_drn", cases=AllocationOptionRiverCases ) def test_riv_allocation( active, top, bottom, stage, bottom_elevation, option, expected_riv, expected_drn @@ -88,3 +176,44 @@ def test_riv_allocation( np.testing.assert_equal(actual_riv, expected_riv) np.testing.assert_equal(actual_drn, expected_drn) + + +@parametrize_with_cases( + argnames="active,top,bottom,drn_elevation", + cases=DrainCases, +) +@parametrize_with_cases(argnames="option,expected", cases=AllocationOptionDrainCases) +def test_drn_allocation(active, top, bottom, drn_elevation, option, expected): + actual_da = allocate_drain_cells(option, active, top, bottom, drn_elevation) + + actual = take_first_planar_cell(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active,top,bottom,head", + cases=GeneralHeadBoundaryCases, +) +@parametrize_with_cases( + argnames="option,expected", cases=AllocationOptionGeneralHeadCases +) +def test_ghb_allocation(active, top, bottom, head, option, expected): + actual_da = allocate_ghb_cells(option, active, top, bottom, head) + + actual = take_first_planar_cell(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active", + cases=RechargeCases, +) +@parametrize_with_cases(argnames="option,expected", cases=AllocationOptionRechargeCases) +def test_rch_allocation(active, option, expected): + actual_da = allocate_rch_cells(option, active) + + actual = take_first_planar_cell(actual_da) + + np.testing.assert_equal(actual, expected) From e3a613cf92bf22aa1473900070c9faf1983a51b2 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 11:26:09 +0200 Subject: [PATCH 23/35] rename river to riv and drain to drn --- docs/api/prepare.rst | 4 ++-- imod/prepare/__init__.py | 4 ++-- imod/prepare/topsystem/__init__.py | 4 ++-- imod/prepare/topsystem/allocation.py | 4 ++-- imod/tests/test_topsystem/test_allocation.py | 8 ++++---- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/api/prepare.rst b/docs/api/prepare.rst index 870fe4655..407207130 100644 --- a/docs/api/prepare.rst +++ b/docs/api/prepare.rst @@ -33,9 +33,9 @@ Prepare model input get_upper_active_layer_number ALLOCATION_OPTION - allocate_drain_cells + allocate_drn_cells allocate_ghb_cells allocate_rch_cells - allocate_river_cells + allocate_riv_cells c_leakage c_radial \ No newline at end of file diff --git a/imod/prepare/__init__.py b/imod/prepare/__init__.py index 600c24917..951c849d8 100644 --- a/imod/prepare/__init__.py +++ b/imod/prepare/__init__.py @@ -36,10 +36,10 @@ ) from imod.prepare.topsystem import ( ALLOCATION_OPTION, - allocate_drain_cells, + allocate_drn_cells, allocate_ghb_cells, allocate_rch_cells, - allocate_river_cells, + allocate_riv_cells, c_leakage, c_radial, ) diff --git a/imod/prepare/topsystem/__init__.py b/imod/prepare/topsystem/__init__.py index 6e13f2801..a63b10718 100644 --- a/imod/prepare/topsystem/__init__.py +++ b/imod/prepare/topsystem/__init__.py @@ -1,8 +1,8 @@ from imod.prepare.topsystem.allocation import ( ALLOCATION_OPTION, - allocate_drain_cells, + allocate_drn_cells, allocate_ghb_cells, allocate_rch_cells, - allocate_river_cells, + allocate_riv_cells, ) from imod.prepare.topsystem.resistance import c_leakage, c_radial diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 7c03c633e..b3486fec8 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -36,7 +36,7 @@ class ALLOCATION_OPTION(Enum): ) -def allocate_river_cells( +def allocate_riv_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray, top: GridDataArray, @@ -63,7 +63,7 @@ def allocate_river_cells( return _allocate_cells__at_first_active(active) -def allocate_drain_cells( +def allocate_drn_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray, top: GridDataArray, diff --git a/imod/tests/test_topsystem/test_allocation.py b/imod/tests/test_topsystem/test_allocation.py index 537f8fb61..7713ce15c 100644 --- a/imod/tests/test_topsystem/test_allocation.py +++ b/imod/tests/test_topsystem/test_allocation.py @@ -3,10 +3,10 @@ from imod.prepare.topsystem import ( ALLOCATION_OPTION, - allocate_drain_cells, + allocate_drn_cells, allocate_ghb_cells, allocate_rch_cells, - allocate_river_cells, + allocate_riv_cells, ) from imod.typing import GridDataArray from imod.typing.grid import is_unstructured, zeros_like @@ -163,7 +163,7 @@ def case_at_first_active(self): def test_riv_allocation( active, top, bottom, stage, bottom_elevation, option, expected_riv, expected_drn ): - actual_riv_da, actual_drn_da = allocate_river_cells( + actual_riv_da, actual_drn_da = allocate_riv_cells( option, active, top, bottom, stage, bottom_elevation ) @@ -184,7 +184,7 @@ def test_riv_allocation( ) @parametrize_with_cases(argnames="option,expected", cases=AllocationOptionDrainCases) def test_drn_allocation(active, top, bottom, drn_elevation, option, expected): - actual_da = allocate_drain_cells(option, active, top, bottom, drn_elevation) + actual_da = allocate_drn_cells(option, active, top, bottom, drn_elevation) actual = take_first_planar_cell(actual_da) From 47c5ba9d756b4be9bd82c5d139b753d01eca82e9 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 11:53:42 +0200 Subject: [PATCH 24/35] Add tests for enforced_dim_order decorator --- imod/tests/test_util/test_util_dims.py | 33 ++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 imod/tests/test_util/test_util_dims.py diff --git a/imod/tests/test_util/test_util_dims.py b/imod/tests/test_util/test_util_dims.py new file mode 100644 index 000000000..7436f7763 --- /dev/null +++ b/imod/tests/test_util/test_util_dims.py @@ -0,0 +1,33 @@ +from imod.util.dims import enforced_dim_order + + +def test_enforced_dim_order_structured(basic_dis): + ibound, _, _ = basic_dis + + @enforced_dim_order + def to_be_decorated(da): + return da + + ibound_wrong_order = ibound.transpose("x", "y", "layer") + + actual = to_be_decorated(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) + + +def test_enforced_dim_order_unstructured(basic_unstructured_dis): + ibound, _, _ = basic_unstructured_dis + + @enforced_dim_order + def to_be_decorated(da): + return da + + face_dim = ibound.ugrid.grid.face_dimension + + ibound_wrong_order = ibound.transpose(face_dim, "layer") + + actual = to_be_decorated(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) \ No newline at end of file From d3ca2ecf4735b28bc931dfd439c64f1fe5b7ddf9 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 15:27:55 +0200 Subject: [PATCH 25/35] Enforce unstructured grid in preserve_gridtype --- imod/typing/grid.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index 1b59627ac..8b5908f86 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -353,6 +353,11 @@ def enforce_dim_order(grid: xu.UgridDataArray) -> xu.UgridDataArray: ) +def _enforce_unstructured(obj: GridDataArray, ugrid2d = xu.Ugrid2d) -> xu.UgridDataArray: + """Force obj to unstructured""" + return xu.UgridDataArray(xr.DataArray(obj), ugrid2d) + + def preserve_gridtype(func): """ Decorator to preserve gridtype, this is to work around the following xugrid @@ -380,8 +385,8 @@ def decorator(*args, **kwargs): if unstructured: # Multiple grids returned if isinstance(x, tuple): - return tuple(xu.UgridDataArray(i, grid) for i in x) - return xu.UgridDataArray(x, grid) + return tuple(_enforce_unstructured(i, grid) for i in x) + return _enforce_unstructured(x, grid) return x return decorator From 67ed07a39c0de4c4661c6a93437637145941a703 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 15:29:08 +0200 Subject: [PATCH 26/35] Add tests to preserve gridtype and enforce dim order --- imod/tests/test_typing/test_typing_grid.py | 64 ++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 imod/tests/test_typing/test_typing_grid.py diff --git a/imod/tests/test_typing/test_typing_grid.py b/imod/tests/test_typing/test_typing_grid.py new file mode 100644 index 000000000..1d3a5f4c9 --- /dev/null +++ b/imod/tests/test_typing/test_typing_grid.py @@ -0,0 +1,64 @@ +import xarray as xr +import xugrid as xu + +from imod.typing.grid import enforce_dim_order, preserve_gridtype + + +def test_preserve_gridtype__single_output(basic_unstructured_dis): + uda, _, da = basic_unstructured_dis + + @preserve_gridtype + def to_be_decorated(a, b): + return a * b + + result1 = to_be_decorated(uda, da) + result2 = to_be_decorated(da, uda) + + # Verify fixture provides expected type + assert isinstance(da, xr.DataArray) + assert isinstance(uda, xu.UgridDataArray) + + assert isinstance(result1, xu.UgridDataArray) + assert isinstance(result2, xu.UgridDataArray) + + +def test_preserve_gridtype__multiple_outputs(basic_unstructured_dis): + uda, _, da = basic_unstructured_dis + + @preserve_gridtype + def to_be_decorated(a, b): + return a * b, a + b + + result1a, result1b = to_be_decorated(uda, da) + result2a, result2b = to_be_decorated(da, uda) + + # Verify fixture provides expected type + assert isinstance(da, xr.DataArray) + assert isinstance(uda, xu.UgridDataArray) + + assert isinstance(result1a, xu.UgridDataArray) + assert isinstance(result1b, xu.UgridDataArray) + assert isinstance(result2a, xu.UgridDataArray) + assert isinstance(result2b, xu.UgridDataArray) + +def test_enforce_dim_order__structured(basic_dis): + ibound, _, _ = basic_dis + + ibound_wrong_order = ibound.transpose("x", "y", "layer") + + actual = enforce_dim_order(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) + +def test_enforce_dim_order__unstructured(basic_unstructured_dis): + ibound, _, _ = basic_unstructured_dis + + face_dim = ibound.ugrid.grid.face_dimension + + ibound_wrong_order = ibound.transpose(face_dim, "layer") + + actual = enforce_dim_order(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) From 0f645609f5b676fe8551df627c5c69cf773ad24c Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 15:29:24 +0200 Subject: [PATCH 27/35] Format --- imod/tests/test_typing/test_typing_grid.py | 10 ++++++---- imod/tests/test_util/test_util_dims.py | 6 +++--- imod/typing/grid.py | 2 +- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/imod/tests/test_typing/test_typing_grid.py b/imod/tests/test_typing/test_typing_grid.py index 1d3a5f4c9..10a0bb953 100644 --- a/imod/tests/test_typing/test_typing_grid.py +++ b/imod/tests/test_typing/test_typing_grid.py @@ -14,17 +14,17 @@ def to_be_decorated(a, b): result1 = to_be_decorated(uda, da) result2 = to_be_decorated(da, uda) - # Verify fixture provides expected type + # Verify fixture provides expected type assert isinstance(da, xr.DataArray) assert isinstance(uda, xu.UgridDataArray) assert isinstance(result1, xu.UgridDataArray) assert isinstance(result2, xu.UgridDataArray) - + def test_preserve_gridtype__multiple_outputs(basic_unstructured_dis): uda, _, da = basic_unstructured_dis - + @preserve_gridtype def to_be_decorated(a, b): return a * b, a + b @@ -32,7 +32,7 @@ def to_be_decorated(a, b): result1a, result1b = to_be_decorated(uda, da) result2a, result2b = to_be_decorated(da, uda) - # Verify fixture provides expected type + # Verify fixture provides expected type assert isinstance(da, xr.DataArray) assert isinstance(uda, xu.UgridDataArray) @@ -41,6 +41,7 @@ def to_be_decorated(a, b): assert isinstance(result2a, xu.UgridDataArray) assert isinstance(result2b, xu.UgridDataArray) + def test_enforce_dim_order__structured(basic_dis): ibound, _, _ = basic_dis @@ -51,6 +52,7 @@ def test_enforce_dim_order__structured(basic_dis): assert actual.dims == ibound.dims assert isinstance(actual, type(ibound)) + def test_enforce_dim_order__unstructured(basic_unstructured_dis): ibound, _, _ = basic_unstructured_dis diff --git a/imod/tests/test_util/test_util_dims.py b/imod/tests/test_util/test_util_dims.py index 7436f7763..74e91a808 100644 --- a/imod/tests/test_util/test_util_dims.py +++ b/imod/tests/test_util/test_util_dims.py @@ -7,7 +7,7 @@ def test_enforced_dim_order_structured(basic_dis): @enforced_dim_order def to_be_decorated(da): return da - + ibound_wrong_order = ibound.transpose("x", "y", "layer") actual = to_be_decorated(ibound_wrong_order) @@ -22,7 +22,7 @@ def test_enforced_dim_order_unstructured(basic_unstructured_dis): @enforced_dim_order def to_be_decorated(da): return da - + face_dim = ibound.ugrid.grid.face_dimension ibound_wrong_order = ibound.transpose(face_dim, "layer") @@ -30,4 +30,4 @@ def to_be_decorated(da): actual = to_be_decorated(ibound_wrong_order) assert actual.dims == ibound.dims - assert isinstance(actual, type(ibound)) \ No newline at end of file + assert isinstance(actual, type(ibound)) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index 8b5908f86..2a74d7b25 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -353,7 +353,7 @@ def enforce_dim_order(grid: xu.UgridDataArray) -> xu.UgridDataArray: ) -def _enforce_unstructured(obj: GridDataArray, ugrid2d = xu.Ugrid2d) -> xu.UgridDataArray: +def _enforce_unstructured(obj: GridDataArray, ugrid2d=xu.Ugrid2d) -> xu.UgridDataArray: """Force obj to unstructured""" return xu.UgridDataArray(xr.DataArray(obj), ugrid2d) From 6f0325d35f4ad3aa2415d5dd7e5f8acc32bc435d Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 15:34:07 +0200 Subject: [PATCH 28/35] Fix bug where optional geopandas throws error in type annotation --- imod/util/spatial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/util/spatial.py b/imod/util/spatial.py index eb32382e5..391e1ebf3 100644 --- a/imod/util/spatial.py +++ b/imod/util/spatial.py @@ -636,7 +636,7 @@ def is_divisor(numerator: FloatArray, denominator: float) -> bool: return (np.isclose(remainder, 0.0) | np.isclose(remainder, denominator)).all() -def _polygonize(da: xr.DataArray) -> gpd.GeoDataFrame: +def _polygonize(da: xr.DataArray) -> "gpd.GeoDataFrame": """ Polygonize a 2D-DataArray into a GeoDataFrame of polygons. From 9f65a0121ddc8eb00156bd153950b140fc43011b Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Tue, 9 Apr 2024 15:41:42 +0200 Subject: [PATCH 29/35] Update changelog --- docs/api/changelog.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index 83417922e..f4f798c35 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -9,6 +9,15 @@ The format is based on `Keep a Changelog`_, and this project adheres to [Unreleased] ------------ +Added +~~~~~ +- Added functions to allocate planar grids over layers for the topsystem in + :func:`imod.prepare.allocate_drn_cells`, + :func:`imod.prepare.allocate_ghb_cells`, + :func:`imod.prepare.allocate_rch_cells`, + :func:`imod.prepare.allocate_riv_cells`, for this multiple options can be + selected, available in :func:`imod.prepare.ALLOCATION_OPTION`. + Fixed ~~~~~ - No ``ValidationError`` thrown anymore in :class:`imod.mf6.River` when From 1071d108b8fd48859ae73809384ed2f187cdc7ec Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 10 Apr 2024 16:28:59 +0200 Subject: [PATCH 30/35] Type annotation --- imod/prepare/topsystem/resistance.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/imod/prepare/topsystem/resistance.py b/imod/prepare/topsystem/resistance.py index 5e3ff4298..c340baa51 100644 --- a/imod/prepare/topsystem/resistance.py +++ b/imod/prepare/topsystem/resistance.py @@ -1,7 +1,15 @@ import numpy as np +from imod.typing import GridDataArray -def c_radial(L, kh, kv, B, D): + +def c_radial( + L: GridDataArray, + kh: GridDataArray, + kv: GridDataArray, + B: GridDataArray, + D: GridDataArray, +) -> GridDataArray: """ Ernst's radial resistance term to a drain. @@ -33,7 +41,17 @@ def c_radial(L, kh, kv, B, D): return c -def c_leakage(kh, kv, D, c0, c1, B, length, dx, dy): +def c_leakage( + kh: GridDataArray, + kv: GridDataArray, + D: GridDataArray, + c0: GridDataArray, + c1: GridDataArray, + B: GridDataArray, + length: GridDataArray, + dx: GridDataArray, + dy: GridDataArray, +) -> GridDataArray: """ Computes the phreatic leakage resistance. From 39108cabc23347b6dd3ab6954e498c34a9371030 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 10 Apr 2024 16:30:04 +0200 Subject: [PATCH 31/35] Extend docstring create_layered_top and import in prepare module --- imod/prepare/__init__.py | 1 + imod/prepare/layer.py | 16 +++++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/imod/prepare/__init__.py b/imod/prepare/__init__.py index 951c849d8..004612578 100644 --- a/imod/prepare/__init__.py +++ b/imod/prepare/__init__.py @@ -15,6 +15,7 @@ from imod.prepare import spatial, subsoil, surface_water from imod.prepare.layer import ( + create_layered_top, get_lower_active_grid_cells, get_lower_active_layer_number, get_upper_active_grid_cells, diff --git a/imod/prepare/layer.py b/imod/prepare/layer.py index 148b88ddf..daaf073eb 100644 --- a/imod/prepare/layer.py +++ b/imod/prepare/layer.py @@ -76,7 +76,21 @@ def get_lower_active_grid_cells(active: GridDataArray) -> GridDataArray: def create_layered_top(bottom: GridDataArray, top: GridDataArray) -> GridDataArray: """ - Create a top array with layers from a single top array and a full bottom array + Create a top array with a layer dimension, from a top array with no layer + dimension and a bottom array with a layer dimension. The (output) top of + layer n is assigned the bottom of layer n-1. + + Parameters + ---------- + bottom: {DataArray, UgridDataArray} + Bottoms with layer dimension + top: {DataArray, UgridDataArray} + Top, without layer dimension + + Returns + ------- + new_top: {DataArray, UgridDataArray} + Top with layer dimension. """ new_top = zeros_like(bottom) new_top[0] = top From 1a44531b8298d1630f9b8231e02abafcbbb5b1f1 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 10 Apr 2024 16:30:37 +0200 Subject: [PATCH 32/35] Add create_layered_top to API docs --- docs/api/prepare.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/api/prepare.rst b/docs/api/prepare.rst index 407207130..28cbb2c5f 100644 --- a/docs/api/prepare.rst +++ b/docs/api/prepare.rst @@ -31,6 +31,7 @@ Prepare model input get_lower_active_layer_number get_upper_active_grid_cells get_upper_active_layer_number + create_layered_top ALLOCATION_OPTION allocate_drn_cells From c2485ed9df4dc4670874e5ffee7a9710e0090c00 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 10 Apr 2024 16:31:25 +0200 Subject: [PATCH 33/35] Add ValueError in case wrong option is applied for ALLOCATION_OPTION as well --- imod/prepare/topsystem/allocation.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index b3486fec8..1bbef7f3e 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -61,7 +61,16 @@ def allocate_riv_cells( return _allocate_cells__at_elevation(top, bottom, bottom_elevation) case ALLOCATION_OPTION.at_first_active: return _allocate_cells__at_first_active(active) - + case _: + raise ValueError( + "Received incompatible setting for rivers, only" + f"'{ALLOCATION_OPTION.stage_to_riv_bot.name}' and" + f"'{ALLOCATION_OPTION.first_active_to_riv_bot.name}' and" + f"'{ALLOCATION_OPTION.first_active_to_riv_bot__drn.name}' and" + f"'{ALLOCATION_OPTION.at_elevation.name}' and" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) def allocate_drn_cells( allocation_option: ALLOCATION_OPTION, From 5d92bb0284a33c879608cdf7d22de91639ef3703 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 10 Apr 2024 16:37:12 +0200 Subject: [PATCH 34/35] Fix typos --- imod/prepare/topsystem/allocation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 1bbef7f3e..30b019339 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -107,7 +107,7 @@ def allocate_ghb_cells( return _allocate_cells__at_first_active(active)[0] case _: raise ValueError( - "Received incompatible setting for drains, only" + "Received incompatible setting for general head boundary, only" f"'{ALLOCATION_OPTION.at_elevation.name}' and" f"'{ALLOCATION_OPTION.at_first_active.name}' supported." f"got: '{allocation_option.name}'" @@ -123,7 +123,7 @@ def allocate_rch_cells( return _allocate_cells__at_first_active(active)[0] case _: raise ValueError( - "Received incompatible setting for drains, only" + "Received incompatible setting for recharge, only" f"'{ALLOCATION_OPTION.at_first_active.name}' supported." f"got: '{allocation_option.name}'" ) From b8adfacb14a6e8d78331d3688be350dc0df21acd Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Wed, 10 Apr 2024 16:37:41 +0200 Subject: [PATCH 35/35] Format --- imod/prepare/layer.py | 2 +- imod/prepare/topsystem/allocation.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/imod/prepare/layer.py b/imod/prepare/layer.py index daaf073eb..35088d54e 100644 --- a/imod/prepare/layer.py +++ b/imod/prepare/layer.py @@ -86,7 +86,7 @@ def create_layered_top(bottom: GridDataArray, top: GridDataArray) -> GridDataArr Bottoms with layer dimension top: {DataArray, UgridDataArray} Top, without layer dimension - + Returns ------- new_top: {DataArray, UgridDataArray} diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 30b019339..792bf3547 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -72,6 +72,7 @@ def allocate_riv_cells( f"got: '{allocation_option.name}'" ) + def allocate_drn_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray,