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

1 import numpy as np

2 import xarray as xr

3

4 from imod.logging import init_log_decorator

5 from imod.mf6.boundary_condition import AdvancedBoundaryCondition, BoundaryCondition

6 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA

7 from imod.schemata import (

8 AllInsideNoDataSchema,

9 AllNoDataSchema,

10 AllValueSchema,

11 CoordsSchema,

12 DimsSchema,

13 DTypeSchema,

14 IdentityNoDataSchema,

15 IndexesSchema,

16 OtherCoordsSchema,

17)

18

19

20 class 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