Coverage for C:\src\imod-python\imod\evaluate\constraints.py: 88%
95 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
1import numpy as np
2import pandas as pd
3import xarray as xr
5import imod
8def stability_constraint_wel(wel, top_bot, porosity=0.3, R=1.0):
9 r"""
10 Computes sink/source stability constraint as applied in MT3D for adaptive
11 timestepping (Zheng & Wang, 1999 p54).
13 .. math:: \Delta t \leq \frac{R\theta }{\left | q_{s} \right |}
15 For the WEL package, a flux is known
16 beforehand, so we can evaluate beforehand if a flux assigned to a cell
17 will necessitate a small timestap, thereby slowing down calculations.
19 Returns a ipf DataFrame that includes a column for the specific discharge and
20 resulting minimum timestep.
22 Parameters
23 ----------
24 wel: pd.DataFrame
25 pd.DataFrame that defines a WEL package. Minimally includes
26 x, y, layer and Q column.
27 top_bot: xr.Dataset of floats, containing 'top', 'bot' and optionally
28 'dz' of layers.
29 Dimensions must be exactly ``("layer", "y", "x")``.
30 porosity: float or xr.DataArray of floats, optional (default 0.3)
31 If xr.DataArray, dimensions must be exactly ``("layer", "y", "x")``.
32 R: Retardation factor, optional (default)
33 Only when sorption is a factor.
35 Returns
36 -------
37 wel: pd.DataFrame containing addition qs (specific discharge) and
38 dt (minimum timestep) columns
39 """
41 b = imod.select.points_in_bounds(
42 top_bot, x=wel["x"], y=wel["y"], layer=wel["layer"]
43 )
44 indices = imod.select.points_indices(
45 top_bot, x=wel.loc[b, "x"], y=wel.loc[b, "y"], layer=wel.loc[b, "layer"]
46 )
48 if "dz" not in top_bot: # create if not present
49 top_bot["dz"] = top_bot["top"] - top_bot["bot"]
50 top_bot["volume"] = np.abs(top_bot["dz"] * top_bot.dx * top_bot.dy)
52 wel.loc[b, "qs"] = wel.loc[b, "Q"].abs() / top_bot["volume"].isel(indices).values
53 wel.loc[b, "dt"] = R * porosity / wel.loc[b, "qs"]
55 return wel
58def stability_constraint_advection(front, lower, right, top_bot, porosity=0.3, R=1.0):
59 r"""
60 Computes advection stability constraint as applied in MT3D for adaptive
61 timestepping (Zheng & Wang, 1999 p54):
63 .. math:: \Delta t \leq \frac{R}{\frac{\left | v_{x} \right |}{\Delta x}+\frac{\left | v_{y} \right |}{\Delta y}+\frac{\left | v_{z} \right |}{\Delta z}}
65 This function can be used to select
66 which cells necessitate a small timestap, thereby slowing down calculations.
68 Front, lower, and right arguments refer to iMOD face flow budgets, in cubic
69 meters per day. In terms of flow direction these are defined as:
71 * ``front``: positive with ``y`` (negative with row index)
72 * ``lower``: positive with ``layer`` (positive with layer index)
73 * ``right``: negative with ``x`` (negative with column index)
75 Returns the minimum timestep that is required to satisfy this constraint.
76 The resulting dt xr.DataArray is the minimum timestep over all three directions,
77 dt_xyz is an xr.Dataset containing minimum timesteps for the three directions
78 separately.
80 Parameters
81 ----------
82 front: xr.DataArray of floats, optional
83 Dimensions must be exactly ``("layer", "y", "x")``.
84 lower: xr.DataArray of floats, optional
85 Dimensions must be exactly ``("layer", "y", "x")``.
86 right: xr.DataArray of floats, optional
87 Dimensions must be exactly ``("layer", "y", "x")``.
88 top_bot: xr.Dataset of floats, containing 'top', 'bot' and optionally
89 'dz' of layers.
90 Dimensions must be exactly ``("layer", "y", "x")``.
91 porosity: float or xr.DataArray of floats, optional (default 0.3)
92 If xr.DataArray, dimensions must be exactly ``("layer", "y", "x")``.
93 R: Retardation factor, optional (default)
94 Only when sorption is a factor.
96 Returns
97 -------
98 dt: xr.DataArray of floats
99 dt_xyz: xr.Dataset of floats
100 """
102 # top_bot reselect to bdg bounds
103 top_bot = top_bot.sel(x=right.x, y=right.y, layer=right.layer)
105 # Compute flow velocities
106 v_x, v_y, v_z = imod.evaluate.flow_velocity(front, lower, right, top_bot, porosity)
108 if "dz" not in top_bot:
109 top_bot["dz"] = top_bot["top"] - top_bot["bot"]
111 # assert all dz positive - Issue #140
112 if not np.all(top_bot["dz"].values[~np.isnan(top_bot["dz"].values)] >= 0):
113 raise ValueError("All dz values should be positive")
115 # absolute velocities (m/d)
116 abs_v_x = np.abs(v_x)
117 abs_v_y = np.abs(v_y)
118 abs_v_z = np.abs(v_z)
120 # dt of constituents (d), set zero velocities to nans
121 dt_x = R / (abs_v_x.where(abs_v_x > 0) / top_bot.dx)
122 dt_y = R / (abs_v_y.where(abs_v_y > 0) / np.abs(top_bot.dy))
123 dt_z = R / (abs_v_z.where(abs_v_z > 0) / top_bot.dz)
125 # overall dt due to advection criterion (d)
126 dt = 1.0 / (1.0 / dt_x + 1.0 / dt_y + 1.0 / dt_z)
128 dt_xyz = xr.concat(
129 (dt_x, dt_y, dt_z), dim=pd.Index(["x", "y", "z"], name="direction")
130 )
131 return dt, dt_xyz
134def _calculate_intra_cell_dt(
135 source_stage, source_cond, sink_stage, sink_cond, eff_volume
136):
137 """Calculate intra-cell dt by assuming a flux from a higher source_stage to a lower sink_stage,
138 ignoring other head influences. Use limiting (lowest) conductance. eff_volume is the effective
139 volume per cell (cell volume * effective porosity)"""
140 source_cond, sink_cond = xr.align(source_cond, sink_cond, join="inner", copy=False)
141 cond = np.minimum(source_cond, sink_cond)
142 Q = cond * (source_stage - sink_stage)
143 Q = Q.where(source_stage > sink_stage)
145 return eff_volume / Q
148def intra_cell_boundary_conditions(
149 top_bot, porosity=0.3, riv=None, ghb=None, drn=None, drop_allnan=True
150):
151 """Function to pre-check boundary-conditions against one another for large intra-cell fluxes.
152 ghb and river can function as source and sink, drn only as sink.
154 Parameters
155 ----------
156 top_bot : xr.Dataset of floats
157 'top_bot' should at least contain `top` and `bot` data_vars
158 porosity : float or xr.DataArray of floats, optional
159 Effective porosity of model cells
160 riv : (dict or list of) imod.RiverPackage, optional
161 ghb : (dict or list of) imod.GeneralHeadBoundaryPackage, optional
162 drn : (dict or list of) imod.DrainagePackage, optional
163 drop_allnan : boolean, optional
164 Whether source-sink combinations without overlap should be dropped from result (default True)
166 Returns
167 -------
168 dt_min: xr.DataArray of floats
169 `dt_min` is the minimum calculated timestep over all combinations of boundary conditions
170 dt_all: xr.DataArray of floats
171 `dt_all` is the calculated timestep for all combinations of boundary conditions
172 """
173 if riv is None and ghb is None:
174 raise ValueError(
175 "At least one source boundary condition must be supplied through riv or ghb."
176 )
178 # convert all inputs to dictionaries of packages
179 if riv is None:
180 riv = {}
181 elif isinstance(riv, dict):
182 pass
183 elif isinstance(riv, (list, tuple)):
184 riv = {f"riv_{i}": single_riv for i, single_riv in enumerate(riv)}
185 else:
186 riv = {"riv": riv}
188 if ghb is None:
189 ghb = {}
190 elif isinstance(ghb, dict):
191 pass
192 elif isinstance(ghb, (list, tuple)):
193 ghb = {f"ghb_{i}": single_riv for i, single_riv in enumerate(ghb)}
194 else:
195 ghb = {"ghb": ghb}
197 if drn is None:
198 drn = {}
199 elif isinstance(drn, dict):
200 pass
201 elif isinstance(drn, (list, tuple)):
202 drn = {f"drn_{i}": single_riv for i, single_riv in enumerate(drn)}
203 else:
204 drn = {"drn": drn}
206 # get sources and sinks:
207 sources = {}
208 sources.update(ghb)
209 sources.update(riv)
210 sinks = {}
211 sinks.update(ghb)
212 sinks.update(riv)
213 sinks.update(drn)
215 # determine effective volume
216 if "dz" not in top_bot:
217 top_bot["dz"] = top_bot["top"] - top_bot["bot"]
219 if (top_bot["dz"] <= 0.0).any():
220 raise ValueError("Cells with thickness <= 0 present in top_bot Dataset.")
222 eff_volume = top_bot["dz"] * top_bot.dx * np.abs(top_bot.dy) * porosity
224 def _get_stage_name(sid):
225 if sid in riv:
226 return "stage"
227 elif sid in ghb:
228 return "head"
229 elif sid in drn:
230 return "elevation"
232 # for all possible combinations: determine dt
233 resultids = []
234 results = []
235 for sourceid, source in sources.items():
236 for sinkid, sink in sinks.items():
237 if sourceid == sinkid:
238 continue
239 comb = f"{sourceid}-{sinkid}"
241 if comb not in resultids:
242 # source in riv: only where stage > bottom elev
243 if sourceid in riv:
244 source = source.where(source["stage"] > source["bottom_elevation"])
246 dt = _calculate_intra_cell_dt(
247 source_stage=source[_get_stage_name(sourceid)],
248 source_cond=source["conductance"].load(),
249 sink_stage=sink[_get_stage_name(sinkid)],
250 sink_cond=sink["conductance"].load(),
251 eff_volume=eff_volume,
252 )
254 if not drop_allnan or not dt.isnull().all():
255 results.append(dt)
256 resultids.append(comb)
257 dt_all = xr.concat(
258 results, pd.Index(resultids, name="combination"), coords="minimal"
259 )
261 # overall dt
262 dt_min = dt_all.min(dim="combination")
264 return dt_min, dt_all