Coverage for C:\src\imod-python\imod\mf6\uzf.py: 93%

103 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import numpy as np 

2import xarray as xr 

3 

4from imod.logging import init_log_decorator 

5from imod.mf6.boundary_condition import AdvancedBoundaryCondition, BoundaryCondition 

6from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA 

7from imod.schemata import ( 

8 AllInsideNoDataSchema, 

9 AllNoDataSchema, 

10 AllValueSchema, 

11 CoordsSchema, 

12 DimsSchema, 

13 DTypeSchema, 

14 IdentityNoDataSchema, 

15 IndexesSchema, 

16 OtherCoordsSchema, 

17) 

18 

19 

20class UnsaturatedZoneFlow(AdvancedBoundaryCondition): 

21 """ 

22 Unsaturated Zone Flow (UZF) package. 

23 

24 TODO: Support timeseries file? Observations? Water Mover? 

25 

26 Parameters 

27 ---------- 

28 surface_depression_depth: array of floats (xr.DataArray) 

29 is the surface depression depth of the UZF cell. 

30 kv_sat: array of floats (xr.DataArray) 

31 is the vertical saturated hydraulic conductivity of the UZF cell. 

32 NOTE: the UZF package determines the location of inactive cells where kv_sat is np.nan 

33 theta_res: array of floats (xr.DataArray) 

34 is the residual (irreducible) water content of the UZF cell. 

35 theta_sat: array of floats (xr.DataArray) 

36 is the saturated water content of the UZF cell. 

37 theta_init: array of floats (xr.DataArray) 

38 is the initial water content of the UZF cell. 

39 epsilon: array of floats (xr.DataArray) 

40 is the epsilon exponent of the UZF cell. 

41 infiltration_rate: array of floats (xr.DataArray) 

42 defines the applied infiltration rate of the UZF cell (LT -1). 

43 et_pot: array of floats (xr.DataArray, optional) 

44 defines the potential evapotranspiration rate of the UZF cell and specified 

45 GWF cell. Evapotranspiration is first removed from the unsaturated zone and any remaining 

46 potential evapotranspiration is applied to the saturated zone. If IVERTCON is greater than zero 

47 then residual potential evapotranspiration not satisfied in the UZF cell is applied to the underlying 

48 UZF and GWF cells. 

49 extinction_depth: array of floats (xr.DataArray, optional) 

50 defines the evapotranspiration extinction depth of the UZF cell. If 

51 IVERTCON is greater than zero and EXTDP extends below the GWF cell bottom then remaining 

52 potential evapotranspiration is applied to the underlying UZF and GWF cells. EXTDP is always 

53 specified, but is only used if SIMULATE ET is specified in the OPTIONS block. 

54 extinction_theta: array of floats (xr.DataArray, optional) 

55 defines the evapotranspiration extinction water content of the UZF 

56 cell. If specified, ET in the unsaturated zone will be simulated either as a function of the 

57 specified PET rate while the water content (THETA) is greater than the ET extinction water content 

58 air_entry_potential: array of floats (xr.DataArray, optional) 

59 defines the air entry potential (head) of the UZF cell. If specified, ET will be 

60 simulated using a capillary pressure based formulation. 

61 Capillary pressure is calculated using the Brooks-Corey retention function ("air_entry") 

62 root_potential: array of floats (xr.DataArray, optional) 

63 defines the root potential (head) of the UZF cell. If specified, ET will be 

64 simulated using a capillary pressure based formulation. 

65 Capillary pressure is calculated using the Brooks-Corey retention function ("air_entry" 

66 root_activity: array of floats (xr.DataArray, optional) 

67 defines the root activity function of the UZF cell. ROOTACT is 

68 the length of roots in a given volume of soil divided by that volume. Values range from 0 to about 3 

69 cm-2, depending on the plant community and its stage of development. If specified, ET will be 

70 simulated using a capillary pressure based formulation. 

71 Capillary pressure is calculated using the Brooks-Corey retention function ("air_entry" 

72 groundwater_ET_function: ({"linear", "square"}, optional) 

73 keyword specifying that groundwater evapotranspiration will be simulated using either 

74 the original ET formulation of MODFLOW-2005 ("linear"). Or by assuming a constant ET 

75 rate for groundwater levels between land surface (TOP) and land surface minus the ET extinction 

76 depth (TOP-EXTDP) ("square"). In the latter case, groundwater ET is smoothly reduced 

77 from the PET rate to zero over a nominal interval at TOP-EXTDP. 

78 simulate_seepage: ({True, False}, optional) 

79 keyword specifying that groundwater discharge (GWSEEP) to land surface will be 

80 simulated. Groundwater discharge is nonzero when groundwater head is greater than land surface. 

81 print_input: ({True, False}, optional) 

82 keyword to indicate that the list of UZF information will be written to the listing file 

83 immediately after it is read. 

84 Default is False. 

85 print_flows: ({True, False}, optional) 

86 keyword to indicate that the list of UZF flow rates will be printed to the listing file for 

87 every stress period time step in which "BUDGET PRINT" is specified in Output Control. If there is 

88 no Output Control option and "PRINT FLOWS" is specified, then flow rates are printed for the last 

89 time step of each stress period. 

90 Default is False. 

91 save_flows: ({True, False}, optional) 

92 keyword to indicate that UZF flow terms will be written to the file specified with "BUDGET 

93 FILEOUT" in Output Control. 

94 Default is False. 

95 observations: [Not yet supported.] 

96 Default is None. 

97 water_mover: [Not yet supported.] 

98 Default is None. 

99 timeseries: [Not yet supported.] 

100 Default is None. 

101 TODO: We could allow the user to either use xarray DataArrays to specify BCS or 

102 use a pd.DataFrame and use the MF6 timeseries files to read input. The latter could 

103 save memory for laterally large-scale models, through efficient use of the UZF cell identifiers. 

104 validate: {True, False} 

105 Flag to indicate whether the package should be validated upon 

106 initialization. This raises a ValidationError if package input is 

107 provided in the wrong manner. Defaults to True. 

108 """ 

109 

110 _period_data = ( 

111 "infiltration_rate", 

112 "et_pot", 

113 "extinction_depth", 

114 "extinction_theta", 

115 "air_entry_potential", 

116 "root_potential", 

117 "root_activity", 

118 ) 

119 

120 _init_schemata = { 

121 "surface_depression_depth": [ 

122 DTypeSchema(np.floating), 

123 BOUNDARY_DIMS_SCHEMA, 

124 ], 

125 "kv_sat": [ 

126 DTypeSchema(np.floating), 

127 IndexesSchema(), 

128 CoordsSchema(("layer",)), 

129 BOUNDARY_DIMS_SCHEMA, 

130 ], 

131 "theta_res": [ 

132 DTypeSchema(np.floating), 

133 BOUNDARY_DIMS_SCHEMA, 

134 ], 

135 "theta_sat": [ 

136 DTypeSchema(np.floating), 

137 BOUNDARY_DIMS_SCHEMA, 

138 ], 

139 "theta_init": [ 

140 DTypeSchema(np.floating), 

141 BOUNDARY_DIMS_SCHEMA, 

142 ], 

143 "epsilon": [ 

144 DTypeSchema(np.floating), 

145 BOUNDARY_DIMS_SCHEMA, 

146 ], 

147 "infiltration_rate": [ 

148 DTypeSchema(np.floating), 

149 BOUNDARY_DIMS_SCHEMA, 

150 ], 

151 "et_pot": [ 

152 DTypeSchema(np.floating), 

153 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var 

154 ], 

155 "extinction_depth": [ 

156 DTypeSchema(np.floating), 

157 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var 

158 ], 

159 "extinction_theta": [ 

160 DTypeSchema(np.floating), 

161 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var 

162 ], 

163 "root_potential": [ 

164 DTypeSchema(np.floating), 

165 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var 

166 ], 

167 "root_activity": [ 

168 DTypeSchema(np.floating), 

169 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var 

170 ], 

171 "print_flows": [DTypeSchema(np.bool_), DimsSchema()], 

172 "save_flows": [DTypeSchema(np.bool_), DimsSchema()], 

173 } 

174 _write_schemata = { 

175 "kv_sat": [ 

176 AllValueSchema(">", 0.0), 

177 OtherCoordsSchema("idomain"), 

178 AllNoDataSchema(), # Check for all nan, can occur while clipping 

179 AllInsideNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

180 ], 

181 "surface_depression_depth": [IdentityNoDataSchema("kv_sat")], 

182 "theta_res": [IdentityNoDataSchema("kv_sat"), AllValueSchema(">=", 0.0)], 

183 "theta_sat": [IdentityNoDataSchema("kv_sat"), AllValueSchema(">=", 0.0)], 

184 "theta_init": [IdentityNoDataSchema("kv_sat"), AllValueSchema(">=", 0.0)], 

185 "epsilon": [IdentityNoDataSchema("kv_sat")], 

186 "infiltration_rate": [IdentityNoDataSchema("kv_sat")], 

187 "et_pot": [IdentityNoDataSchema("kv_sat")], 

188 "extinction_depth": [IdentityNoDataSchema("kv_sat")], 

189 "extinction_theta": [IdentityNoDataSchema("kv_sat")], 

190 "root_potential": [IdentityNoDataSchema("kv_sat")], 

191 "root_activity": [IdentityNoDataSchema("kv_sat")], 

192 } 

193 

194 _package_data = ( 

195 "surface_depression_depth", 

196 "kv_sat", 

197 "theta_res", 

198 "theta_sat", 

199 "theta_init", 

200 "epsilon", 

201 ) 

202 _pkg_id = "uzf" 

203 

204 _template = BoundaryCondition._initialize_template(_pkg_id) 

205 

206 @init_log_decorator() 

207 def __init__( 

208 self, 

209 surface_depression_depth, 

210 kv_sat, 

211 theta_res, 

212 theta_sat, 

213 theta_init, 

214 epsilon, 

215 infiltration_rate, 

216 et_pot=None, 

217 extinction_depth=None, 

218 extinction_theta=None, 

219 air_entry_potential=None, 

220 root_potential=None, 

221 root_activity=None, 

222 ntrailwaves=7, # Recommended in manual 

223 nwavesets=40, 

224 groundwater_ET_function=None, 

225 simulate_groundwater_seepage=False, 

226 print_input=False, 

227 print_flows=False, 

228 save_flows=False, 

229 observations=None, 

230 water_mover=None, 

231 timeseries=None, 

232 validate: bool = True, 

233 ): 

234 landflag = self._determine_landflag(kv_sat) 

235 iuzno = self._create_uzf_numbers(landflag) 

236 ivertcon = self._determine_vertical_connection(iuzno) 

237 

238 dict_dataset = { 

239 # Package data 

240 "surface_depression_depth": surface_depression_depth, 

241 "kv_sat": kv_sat, 

242 "theta_res": theta_res, 

243 "theta_sat": theta_sat, 

244 "theta_init": theta_init, 

245 "epsilon": epsilon, 

246 # Stress period data 

247 "infiltration_rate": infiltration_rate, 

248 "et_pot": et_pot, 

249 "extinction_depth": extinction_depth, 

250 "extinction_theta": extinction_theta, 

251 "air_entry_potential": air_entry_potential, 

252 "root_potential": root_potential, 

253 "root_activity": root_activity, 

254 # Dimensions 

255 "ntrailwaves": ntrailwaves, 

256 "nwavesets": nwavesets, 

257 # Options 

258 "groundwater_ET_function": groundwater_ET_function, 

259 "simulate_gwseep": simulate_groundwater_seepage, 

260 "print_input": print_input, 

261 "print_flows": print_flows, 

262 "save_flows": save_flows, 

263 "observations": observations, 

264 "water_mover": water_mover, 

265 "timeseries": timeseries, 

266 # Additonal indices for Packagedata 

267 "landflag": landflag, 

268 "iuzno": iuzno, 

269 "ivertcon": ivertcon, 

270 } 

271 super().__init__(dict_dataset) 

272 self.dataset["iuzno"].name = "uzf_number" 

273 self._check_options( 

274 groundwater_ET_function, 

275 et_pot, 

276 extinction_depth, 

277 extinction_theta, 

278 air_entry_potential, 

279 root_potential, 

280 root_activity, 

281 ) 

282 self._validate_init_schemata(validate) 

283 

284 def fill_stress_perioddata(self): 

285 """Modflow6 requires something to be filled in the stress perioddata, 

286 even though the data is not used in the current configuration. 

287 Only an infiltration rate is required, 

288 the rest can be filled with dummy values if not provided. 

289 """ 

290 for var in self._period_data: 

291 if self.dataset[var].size == 1: # Prevent loading large arrays in memory 

292 if self.dataset[var].values[()] is None: 

293 self.dataset[var] = xr.full_like(self["infiltration_rate"], 0.0) 

294 else: 

295 raise ValueError("{} cannot be a scalar".format(var)) 

296 

297 def _check_options( 

298 self, 

299 groundwater_ET_function, 

300 et_pot, 

301 extinction_depth, 

302 extinction_theta, 

303 air_entry_potential, 

304 root_potential, 

305 root_activity, 

306 ): 

307 simulate_et = [x is not None for x in [et_pot, extinction_depth]] 

308 unsat_etae = [ 

309 x is not None for x in [air_entry_potential, root_potential, root_activity] 

310 ] 

311 

312 if all(simulate_et): 

313 self.dataset["simulate_et"] = True 

314 elif any(simulate_et): 

315 raise ValueError("To simulate ET, set both et_pot and extinction_depth") 

316 

317 if extinction_theta is not None: 

318 self.dataset["unsat_etwc"] = True 

319 

320 if all(unsat_etae): 

321 self.dataset["unsat_etae"] = True 

322 elif any(unsat_etae): 

323 raise ValueError( 

324 "To simulate ET with a capillary based formulation, set air_entry_potential, root_potential, and root_activity" 

325 ) 

326 

327 if all(unsat_etae) and (extinction_theta is not None): 

328 raise ValueError( 

329 """Both capillary based formulation and water content based formulation set based on provided input data. 

330 Please provide either only extinction_theta or (air_entry_potential, root_potential, and root_activity)""" 

331 ) 

332 

333 if groundwater_ET_function == "linear": 

334 self.dataset["linear_gwet"] = True 

335 elif groundwater_ET_function == "square": 

336 self.dataset["square_gwet"] = True 

337 elif groundwater_ET_function is None: 

338 pass 

339 else: 

340 raise ValueError( 

341 "Groundwater ET function should be either 'linear','square' or None" 

342 ) 

343 

344 def _create_uzf_numbers(self, landflag): 

345 """Create unique UZF ID's. Inactive cells equal 0""" 

346 return np.cumsum(np.ravel(landflag)).reshape(landflag.shape) * landflag 

347 

348 def _determine_landflag(self, kv_sat): 

349 return (np.isfinite(kv_sat)).astype(np.int32) 

350 

351 def _determine_vertical_connection(self, uzf_number): 

352 return uzf_number.shift(layer=-1, fill_value=0) 

353 

354 def _package_data_to_sparse(self): 

355 notnull = self.dataset["landflag"].values == 1 

356 iuzno = self.dataset["iuzno"].values[notnull] 

357 landflag = self.dataset["landflag"].values[notnull] 

358 ivertcon = self.dataset["ivertcon"].values[notnull] 

359 

360 ds = self.dataset[list(self._package_data)] 

361 

362 layer = ds["layer"].values 

363 arrdict = self._ds_to_arrdict(ds) 

364 recarr = super()._to_struct_array(arrdict, layer) 

365 

366 field_spec = self._get_field_spec_from_dtype(recarr) 

367 field_names = [i[0] for i in field_spec] 

368 index_spec = [("iuzno", np.int32)] + field_spec[:3] 

369 field_spec = ( 

370 [("landflag", np.int32)] + [("ivertcon", np.int32)] + field_spec[3:] 

371 ) 

372 sparse_dtype = np.dtype(index_spec + field_spec) 

373 

374 recarr_new = np.empty(recarr.shape, dtype=sparse_dtype) 

375 recarr_new["iuzno"] = iuzno 

376 recarr_new["landflag"] = landflag 

377 recarr_new["ivertcon"] = ivertcon 

378 recarr_new[field_names] = recarr 

379 

380 return recarr_new 

381 

382 def render(self, directory, pkgname, globaltimes, binary): 

383 """Render fills in the template only, doesn't write binary data""" 

384 d = {} 

385 bin_ds = self.dataset[list(self._period_data)] 

386 d["periods"] = self._period_paths( 

387 directory, pkgname, globaltimes, bin_ds, binary=False 

388 ) 

389 not_options = ( 

390 list(self._period_data) + list(self._package_data) + ["iuzno" + "ivertcon"] 

391 ) 

392 d = self._get_options(d, not_options=not_options) 

393 path = directory / pkgname / f"{self._pkg_id}-pkgdata.dat" 

394 d["packagedata"] = path.as_posix() 

395 d["nuzfcells"] = self._max_active_n() 

396 return self._template.render(d) 

397 

398 def _to_struct_array(self, arrdict, layer): 

399 """Convert from dense arrays to list based input, 

400 since the perioddata does not require cellids but iuzno, we hgave to override""" 

401 # TODO add pkgcheck that period table aligns 

402 # Get the number of valid values 

403 notnull = self.dataset["landflag"].values == 1 

404 iuzno = self.dataset["iuzno"].values 

405 nrow = notnull.sum() 

406 # Define the numpy structured array dtype 

407 index_spec = [("iuzno", np.int32)] 

408 field_spec = [(key, np.float64) for key in arrdict] 

409 sparse_dtype = np.dtype(index_spec + field_spec) 

410 

411 # Initialize the structured array 

412 recarr = np.empty(nrow, dtype=sparse_dtype) 

413 # Fill in the indices 

414 recarr["iuzno"] = iuzno[notnull] 

415 

416 # Fill in the data 

417 for key, arr in arrdict.items(): 

418 recarr[key] = arr[notnull].astype(np.float64) 

419 

420 return recarr 

421 

422 def _validate(self, schemata, **kwargs): 

423 # Insert additional kwargs 

424 kwargs["kv_sat"] = self["kv_sat"] 

425 errors = super()._validate(schemata, **kwargs) 

426 

427 return errors