Coverage for C:\src\imod-python\imod\prepare\topsystem\conductance.py: 95%
142 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1from enum import Enum
2from typing import Optional
4import numpy as np
6from imod.prepare.topsystem.allocation import _enforce_layered_top
7from imod.schemata import DimsSchema
8from imod.typing import GridDataArray
9from imod.typing.grid import ones_like, preserve_gridtype, zeros_like
10from imod.util.dims import enforced_dim_order
13class DISTRIBUTING_OPTION(Enum):
14 """
15 Enumerator containing settings to distribute 2D conductance grids over
16 vertical layers for the RIV, DRN or GHB package.
18 * ``by_corrected_transmissivity``: RIV. Distribute the conductance by
19 corrected transmissivities. Crosscut thicknesses are used to compute
20 transmissivities. The crosscut thicknesses is computed based on the
21 overlap of bottom_elevation over the bottom allocated layer. Same holds
22 for the stage and top allocated layer. Furthermore the method corrects
23 distribution weights for the mismatch between the midpoints of crosscut
24 areas and model layer midpoints. This is the default method in iMOD 5.6,
25 thus DISTRCOND = 0.
26 * ``equally``: RIV, DRN, GHB. Distribute conductances equally over layers.
27 This matches iMOD 5.6 DISTRCOND = 1 option.
28 * ``by_crosscut_thickness``: RIV. Distribute the conductance by crosscut
29 thicknesses. The crosscut thicknesses is computed based on the overlap of
30 bottom_elevation over the bottom allocated layer. Same holds for the stage
31 and top allocated layer. This matches iMOD 5.6 DISTRCOND = 2 option.
32 * ``by_layer_thickness``: RIV, DRN, GHB. Distribute the conductance by model
33 layer thickness. This matches iMOD 5.6 DISTRCOND = 3 option.
34 * ``by_crosscut_transmissivity``: RIV. Distribute the conductance by
35 crosscut transmissivity. Crosscut thicknesses are used to compute
36 transmissivities. The crosscut thicknesses is computed based on the
37 overlap of bottom_elevation over the bottom allocated layer. Same holds
38 for the stage and top allocated layer. This matches iMOD 5.6 DISTRCOND = 4
39 option.
40 * ``by_conductivity``: RIV, DRN, GHB. Distribute the conductance weighted by
41 model layer hydraulic conductivities. This matches iMOD 5.6 DISTRCOND = 5
42 option.
43 * ``by_layer_transmissivity``: RIV, DRN, GHB. Distribute the conductance by
44 model layer transmissivity. This has no equivalent in iMOD 5.6.
45 """
47 by_corrected_transmissivity = 0
48 equally = 1
49 by_crosscut_thickness = 2
50 by_layer_thickness = 3
51 by_crosscut_transmissivity = 4
52 by_conductivity = 5
53 by_layer_transmissivity = 9 # Not an iMOD 5.6 option
56PLANAR_GRID = (
57 DimsSchema("time", "y", "x")
58 | DimsSchema("y", "x")
59 | DimsSchema("time", "{face_dim}")
60 | DimsSchema("{face_dim}")
61)
64@enforced_dim_order
65def distribute_riv_conductance(
66 distributing_option: DISTRIBUTING_OPTION,
67 allocated: GridDataArray,
68 conductance: GridDataArray,
69 top: GridDataArray,
70 bottom: GridDataArray,
71 k: GridDataArray,
72 stage: GridDataArray,
73 bottom_elevation: GridDataArray,
74) -> GridDataArray:
75 """
76 Function to distribute 2D conductance over vertical layers for the RIV
77 package. Multiple options are available, which need to be selected in the
78 DISTRIBUTING_OPTION enumerator.
80 Parameters
81 ----------
82 distributing_option : DISTRIBUTING_OPTION
83 Distributing option available in the DISTRIBUTING_OPTION enumerator.
84 allocated: DataArray | UgridDataArray
85 3D boolean array with river cell locations. This can be made with the
86 :func:`imod.prepare.allocate_riv_cells` function.
87 conductance: DataArray | UgridDataArray
88 Planar grid with conductances that need to be distributed over layers,
89 so grid cannot contain a layer dimension. Can contain a time dimension.
90 top: DataArray | UgridDataArray
91 Model top
92 bottom: DataArray | UgridDataArray
93 Model layer bottoms
94 k: DataArray | UgridDataArray
95 Hydraulic conductivities
96 stage: DataArray | UgridDataArray
97 Planar grid with river stages, cannot contain a layer dimension. Can
98 contain a time dimension.
99 bottom_elevation: DataArray | UgridDataArray
100 Planar grid with river bottom elevations, cannot contain a layer
101 dimension. Can contain a time dimension.
103 Returns
104 -------
105 Conductances distributed over depth.
107 Examples
108 --------
109 >>> from imod.prepare import allocate_riv_cells, distribute_riv_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION
110 >>> allocated = allocate_riv_cells(
111 ALLOCATION_OPTION.stage_to_riv_bot, active, top, bottom, stage, bottom_elevation
112 )
113 >>> conductances_distributed = distribute_riv_conductance(
114 DISTRIBUTING_OPTION.by_corrected_transmissivity, allocated,
115 conductance, top, bottom, stage, bottom_elevation, k
116 )
117 """
118 PLANAR_GRID.validate(conductance)
120 match distributing_option:
121 case DISTRIBUTING_OPTION.equally:
122 weights = _distribute_weights__equally(allocated)
123 case DISTRIBUTING_OPTION.by_layer_thickness:
124 weights = _distribute_weights__by_layer_thickness(allocated, top, bottom)
125 case DISTRIBUTING_OPTION.by_layer_transmissivity:
126 weights = _distribute_weights__by_layer_transmissivity(
127 allocated, top, bottom, k
128 )
129 case DISTRIBUTING_OPTION.by_conductivity:
130 weights = _distribute_weights__by_conductivity(allocated, k)
131 case DISTRIBUTING_OPTION.by_crosscut_thickness:
132 weights = _distribute_weights__by_crosscut_thickness(
133 allocated, top, bottom, stage, bottom_elevation
134 )
135 case DISTRIBUTING_OPTION.by_crosscut_transmissivity:
136 weights = _distribute_weights__by_crosscut_transmissivity(
137 allocated, top, bottom, k, stage, bottom_elevation
138 )
139 case DISTRIBUTING_OPTION.by_corrected_transmissivity:
140 weights = _distribute_weights__by_corrected_transmissivity(
141 allocated, top, bottom, k, stage, bottom_elevation
142 )
143 case _:
144 raise ValueError(
145 "Received incompatible setting for rivers, only"
146 f"'{DISTRIBUTING_OPTION.equally.name}', "
147 f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', "
148 f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', "
149 f"'{DISTRIBUTING_OPTION.by_conductivity.name}', "
150 f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}', "
151 f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}', and "
152 f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' supported. "
153 f"Got: '{distributing_option.name}'"
154 )
155 return (weights * conductance).where(allocated)
158@enforced_dim_order
159def distribute_drn_conductance(
160 distributing_option: DISTRIBUTING_OPTION,
161 allocated: GridDataArray,
162 conductance: GridDataArray,
163 top: GridDataArray,
164 bottom: GridDataArray,
165 k: GridDataArray,
166 elevation: GridDataArray,
167) -> GridDataArray:
168 """
169 Function to distribute 2D conductance over vertical layers for the DRN
170 package. Multiple options are available, which need to be selected in the
171 DISTRIBUTING_OPTION enumerator.
173 Parameters
174 ----------
175 distributing_option : DISTRIBUTING_OPTION
176 Distributing option available in the DISTRIBUTING_OPTION enumerator.
177 allocated: DataArray | UgridDataArray
178 3D boolean array with drain cell locations. This can be made with the
179 :func:`imod.prepare.allocate_drn_cells` function.
180 conductance: DataArray | UgridDataArray
181 Planar grid with conductances that need to be distributed over layers,
182 so grid cannot contain a layer dimension. Can contain a time dimension.
183 top: DataArray | UgridDataArray
184 Model top
185 bottom: DataArray | UgridDataArray
186 Model layer bottoms
187 k: DataArray | UgridDataArray
188 Hydraulic conductivities
189 elevation: DataArray | UgridDataArray
190 Drain elevation
192 Returns
193 -------
194 Conductances distributed over depth.
196 Examples
197 --------
198 >>> from imod.prepare import allocate_drn_cells, distribute_drn_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION
199 >>> allocated = allocate_drn_cells(
200 ALLOCATION_OPTION.at_elevation, active, top, bottom, drain_elevation
201 )
202 >>> conductances_distributed = distribute_drn_conductance(
203 DISTRIBUTING_OPTION.by_layer_transmissivity, allocated,
204 conductance, top, bottom, k, drain_elevation
205 )
206 """
207 PLANAR_GRID.validate(conductance)
209 match distributing_option:
210 case DISTRIBUTING_OPTION.equally:
211 weights = _distribute_weights__equally(allocated)
212 case DISTRIBUTING_OPTION.by_layer_thickness:
213 weights = _distribute_weights__by_layer_thickness(allocated, top, bottom)
214 case DISTRIBUTING_OPTION.by_layer_transmissivity:
215 weights = _distribute_weights__by_layer_transmissivity(
216 allocated, top, bottom, k
217 )
218 case DISTRIBUTING_OPTION.by_conductivity:
219 weights = _distribute_weights__by_conductivity(allocated, k)
220 case DISTRIBUTING_OPTION.by_crosscut_thickness:
221 weights = _distribute_weights__by_crosscut_thickness(
222 allocated, top, bottom, bc_bottom=elevation
223 )
224 case DISTRIBUTING_OPTION.by_crosscut_transmissivity:
225 weights = _distribute_weights__by_crosscut_transmissivity(
226 allocated, top, bottom, k, bc_bottom=elevation
227 )
228 case DISTRIBUTING_OPTION.by_corrected_transmissivity:
229 weights = _distribute_weights__by_corrected_transmissivity(
230 allocated, top, bottom, k, bc_bottom=elevation
231 )
232 case _:
233 raise ValueError(
234 "Received incompatible setting for drains, only"
235 f"'{DISTRIBUTING_OPTION.equally.name}', "
236 f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', "
237 f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', "
238 f"'{DISTRIBUTING_OPTION.by_conductivity.name}', "
239 f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}', "
240 f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}', and "
241 f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' supported. "
242 f"Got: '{distributing_option.name}'"
243 )
244 return (weights * conductance).where(allocated)
247@enforced_dim_order
248def distribute_ghb_conductance(
249 distributing_option: DISTRIBUTING_OPTION,
250 allocated: GridDataArray,
251 conductance: GridDataArray,
252 top: GridDataArray,
253 bottom: GridDataArray,
254 k: GridDataArray,
255) -> GridDataArray:
256 PLANAR_GRID.validate(conductance)
257 """
258 Function to distribute 2D conductance over vertical layers for the GHB
259 package. Multiple options are available, which need to be selected in the
260 DISTRIBUTING_OPTION enumerator.
262 Parameters
263 ----------
264 distributing_option : DISTRIBUTING_OPTION
265 Distributing option available in the DISTRIBUTING_OPTION enumerator.
266 allocated: DataArray | UgridDataArray
267 3D boolean array with GHB cell locations. This can be made with the
268 :func:`imod.prepare.allocate_ghb_cells` function.
269 conductance: DataArray | UgridDataArray
270 Planar grid with conductances that need to be distributed over layers,
271 so grid cannot contain a layer dimension. Can contain a time dimension.
272 top: DataArray | UgridDataArray
273 Model top
274 bottom: DataArray | UgridDataArray
275 Model layer bottoms
276 k: DataArray | UgridDataArray
277 Hydraulic conductivities
279 Returns
280 -------
281 Conductances distributed over depth.
283 Examples
284 --------
285 >>> from imod.prepare import allocate_ghb_cells, distribute_drn_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION
286 >>> allocated = allocate_ghb_cells(
287 ALLOCATION_OPTION.at_elevation, active, top, bottom, ghb_head
288 )
289 >>> conductances_distributed = distribute_ghb_conductance(
290 DISTRIBUTING_OPTION.by_layer_transmissivity, allocated,
291 conductance, top, bottom, k
292 )
293 """
294 match distributing_option:
295 case DISTRIBUTING_OPTION.equally:
296 weights = _distribute_weights__equally(allocated)
297 case DISTRIBUTING_OPTION.by_layer_thickness:
298 weights = _distribute_weights__by_layer_thickness(allocated, top, bottom)
299 case DISTRIBUTING_OPTION.by_layer_transmissivity:
300 weights = _distribute_weights__by_layer_transmissivity(
301 allocated, top, bottom, k
302 )
303 case DISTRIBUTING_OPTION.by_conductivity:
304 weights = _distribute_weights__by_conductivity(allocated, k)
305 case _:
306 raise ValueError(
307 "Received incompatible setting for general head boundary, only"
308 f"'{DISTRIBUTING_OPTION.equally.name}', "
309 f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', "
310 f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', and "
311 f"'{DISTRIBUTING_OPTION.by_conductivity.name}' supported. "
312 f"Got: '{distributing_option.name}'"
313 )
314 return (weights * conductance).where(allocated)
317@preserve_gridtype
318def _compute_layer_thickness(
319 allocated: GridDataArray, top: GridDataArray, bottom: GridDataArray
320):
321 """
322 Compute 3D grid of thicknesses in allocated cells
323 """
324 top_layered = _enforce_layered_top(top, bottom)
326 thickness = top_layered - bottom
327 return thickness.where(allocated)
330@preserve_gridtype
331def _compute_crosscut_thickness(
332 allocated: GridDataArray,
333 top: GridDataArray,
334 bottom: GridDataArray,
335 bc_top: Optional[GridDataArray] = None,
336 bc_bottom: Optional[GridDataArray] = None,
337):
338 """
339 Compute 3D grid of thicknesses crosscut by boundary condition (river/drain)
340 in allocated cells. So the thickness in the upper allocated layer is bc_top
341 - bottom and the lower allocated layer is top - bc_bottom.
342 """
343 if (bc_top is None) & (bc_bottom is None):
344 raise ValueError("`bc_top` and `bc_bottom` cannot both be None.")
346 top_layered = _enforce_layered_top(top, bottom)
347 thickness = _compute_layer_thickness(allocated, top, bottom)
348 outside = zeros_like(allocated).astype(bool)
350 if bc_top is not None:
351 upper_layer_bc = (bc_top < top_layered) & (bc_top > bottom)
352 outside = outside | (bc_top < bottom)
353 thickness = thickness.where(~upper_layer_bc, thickness - (top_layered - bc_top))
355 if bc_bottom is not None:
356 lower_layer_bc = (bc_bottom < top_layered) & (bc_bottom > bottom)
357 outside = outside | (bc_bottom > top_layered)
358 thickness = thickness.where(~lower_layer_bc, thickness - (bc_bottom - bottom))
360 thickness = thickness.where(~outside, 0.0)
362 return thickness
365def _distribute_weights__by_corrected_transmissivity(
366 allocated: GridDataArray,
367 top: GridDataArray,
368 bottom: GridDataArray,
369 k: GridDataArray,
370 bc_top: Optional[GridDataArray] = None,
371 bc_bottom: Optional[GridDataArray] = None,
372):
373 """
374 Distribute conductances according to default method in iMOD 5.6, as
375 described page 690 of the iMOD 5.6 manual (but then to distribute WEL
376 rates). The method uses crosscut thicknesses to compute transmissivities.
377 Furthermore it corrects distribution weights for the mismatch between the
378 midpoints of crosscut areas and layer midpoints.
379 """
380 crosscut_thickness = _compute_crosscut_thickness(
381 allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom
382 )
383 transmissivity = crosscut_thickness * k
385 top_layered = _enforce_layered_top(top, bottom)
386 layer_thickness = _compute_layer_thickness(allocated, top, bottom)
387 midpoints = (top_layered + bottom) / 2
388 Fc = midpoints.copy()
390 if bc_top is not None:
391 PLANAR_GRID.validate(bc_top)
392 upper_layer_bc = (bc_top < top_layered) & (bc_top > bottom)
393 # Computing vertical midpoint of river crosscutting layers.
394 Fc = Fc.where(~upper_layer_bc, (bottom + bc_top) / 2)
396 if bc_bottom is not None:
397 PLANAR_GRID.validate(bc_bottom)
398 lower_layer_bc = (bc_bottom < top_layered) & (bc_bottom > bottom)
399 # Computing vertical midpoint of river crosscutting layers.
400 Fc = Fc.where(~lower_layer_bc, (top_layered + bc_bottom) / 2)
402 # Correction factor for mismatch between midpoints of crosscut layers and
403 # layer midpoints.
404 F = 1.0 - np.abs(midpoints - Fc) / (layer_thickness * 0.5)
406 transmissivity_corrected = transmissivity * F
407 return transmissivity_corrected / transmissivity_corrected.sum(dim="layer")
410def _distribute_weights__equally(allocated: GridDataArray):
411 """Compute weights over layers equally."""
412 return ones_like(allocated) / allocated.sum(dim="layer")
415def _distribute_weights__by_layer_thickness(
416 allocated: GridDataArray,
417 top: GridDataArray,
418 bottom: GridDataArray,
419):
420 """Compute weights based on layer thickness"""
421 layer_thickness = _compute_layer_thickness(allocated, top, bottom)
423 return layer_thickness / layer_thickness.sum(dim="layer")
426def _distribute_weights__by_crosscut_thickness(
427 allocated: GridDataArray,
428 top: GridDataArray,
429 bottom: GridDataArray,
430 bc_top: Optional[GridDataArray] = None,
431 bc_bottom: Optional[GridDataArray] = None,
432):
433 """Compute weights based on crosscut thickness"""
434 if bc_top is not None:
435 PLANAR_GRID.validate(bc_top)
436 if bc_bottom is not None:
437 PLANAR_GRID.validate(bc_bottom)
439 crosscut_thickness = _compute_crosscut_thickness(
440 allocated, top, bottom, bc_top, bc_bottom
441 ).where(allocated)
443 return crosscut_thickness / crosscut_thickness.sum(dim="layer")
446def _distribute_weights__by_layer_transmissivity(
447 allocated: GridDataArray,
448 top: GridDataArray,
449 bottom: GridDataArray,
450 k: GridDataArray,
451):
452 """Compute weights based on layer transmissivity"""
453 layer_thickness = _compute_layer_thickness(allocated, top, bottom)
454 transmissivity = layer_thickness * k
456 return transmissivity / transmissivity.sum(dim="layer")
459def _distribute_weights__by_crosscut_transmissivity(
460 allocated: GridDataArray,
461 top: GridDataArray,
462 bottom: GridDataArray,
463 k: GridDataArray,
464 bc_top: Optional[GridDataArray] = None,
465 bc_bottom: Optional[GridDataArray] = None,
466):
467 """Compute weights based on crosscut transmissivity"""
468 if bc_top is not None:
469 PLANAR_GRID.validate(bc_top)
470 if bc_bottom is not None:
471 PLANAR_GRID.validate(bc_bottom)
473 crosscut_thickness = _compute_crosscut_thickness(
474 allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom
475 )
476 transmissivity = crosscut_thickness * k
478 return transmissivity / transmissivity.sum(dim="layer")
481def _distribute_weights__by_conductivity(allocated: GridDataArray, k: GridDataArray):
482 """Compute weights based on hydraulic conductivity"""
483 k_allocated = allocated * k
485 return k_allocated / k_allocated.sum(dim="layer")