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

1import numpy as np 

2import pandas as pd 

3import xarray as xr 

4 

5import imod 

6 

7 

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). 

12 

13 .. math:: \Delta t \leq \frac{R\theta }{\left | q_{s} \right |} 

14 

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. 

18 

19 Returns a ipf DataFrame that includes a column for the specific discharge and 

20 resulting minimum timestep. 

21 

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. 

34 

35 Returns 

36 ------- 

37 wel: pd.DataFrame containing addition qs (specific discharge) and 

38 dt (minimum timestep) columns 

39 """ 

40 

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 ) 

47 

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) 

51 

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"] 

54 

55 return wel 

56 

57 

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): 

62 

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}} 

64 

65 This function can be used to select 

66 which cells necessitate a small timestap, thereby slowing down calculations. 

67 

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: 

70 

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) 

74 

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. 

79 

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. 

95 

96 Returns 

97 ------- 

98 dt: xr.DataArray of floats 

99 dt_xyz: xr.Dataset of floats 

100 """ 

101 

102 # top_bot reselect to bdg bounds 

103 top_bot = top_bot.sel(x=right.x, y=right.y, layer=right.layer) 

104 

105 # Compute flow velocities 

106 v_x, v_y, v_z = imod.evaluate.flow_velocity(front, lower, right, top_bot, porosity) 

107 

108 if "dz" not in top_bot: 

109 top_bot["dz"] = top_bot["top"] - top_bot["bot"] 

110 

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") 

114 

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) 

119 

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) 

124 

125 # overall dt due to advection criterion (d) 

126 dt = 1.0 / (1.0 / dt_x + 1.0 / dt_y + 1.0 / dt_z) 

127 

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 

132 

133 

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) 

144 

145 return eff_volume / Q 

146 

147 

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. 

153 

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) 

165 

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 ) 

177 

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} 

187 

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} 

196 

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} 

205 

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) 

214 

215 # determine effective volume 

216 if "dz" not in top_bot: 

217 top_bot["dz"] = top_bot["top"] - top_bot["bot"] 

218 

219 if (top_bot["dz"] <= 0.0).any(): 

220 raise ValueError("Cells with thickness <= 0 present in top_bot Dataset.") 

221 

222 eff_volume = top_bot["dz"] * top_bot.dx * np.abs(top_bot.dy) * porosity 

223 

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" 

231 

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}" 

240 

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"]) 

245 

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 ) 

253 

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 ) 

260 

261 # overall dt 

262 dt_min = dt_all.min(dim="combination") 

263 

264 return dt_min, dt_all