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

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.prepare.layer import get_upper_active_grid_cells

8 from imod.schemata import (

9 AllInsideNoDataSchema,

10 AllNoDataSchema,

11 AllValueSchema,

12 CoordsSchema,

13 DimsSchema,

14 DTypeSchema,

15 IdentityNoDataSchema,

16 IndexesSchema,

17 OtherCoordsSchema,

18)

19

20

21 class UnsaturatedZoneFlow(AdvancedBoundaryCondition):

22 """

23 Unsaturated Zone Flow (UZF) package.

24

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

26

27 Parameters

28 ----------

29 surface_depression_depth: array of floats (xr.DataArray)

30 is the surface depression depth of the UZF cell.

31 kv_sat: array of floats (xr.DataArray)

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

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

34 theta_res: array of floats (xr.DataArray)

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

36 theta_sat: array of floats (xr.DataArray)

37 is the saturated water content of the UZF cell.

38 theta_init: array of floats (xr.DataArray)

39 is the initial water content of the UZF cell.

40 epsilon: array of floats (xr.DataArray)

41 is the epsilon exponent of the UZF cell.

42 infiltration_rate: array of floats (xr.DataArray)

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

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

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

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

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

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

49 UZF and GWF cells.

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

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

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

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

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

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

56 defines the evapotranspiration extinction water content of the UZF

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

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

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

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

61 simulated using a capillary pressure based formulation.

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

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

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

65 simulated using a capillary pressure based formulation.

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

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

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

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

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

71 simulated using a capillary pressure based formulation.

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

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

74 keyword specifying that groundwater evapotranspiration will be simulated using either

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

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

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

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

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

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

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

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

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

84 immediately after it is read.

85 Default is False.

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

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

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

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

90 time step of each stress period.

91 Default is False.

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

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

94 FILEOUT" in Output Control.

95 Default is False.

96 budget_fileout: ({"str"}, optional)

97 path to output cbc-file for UZF budgets

98 budgetcsv_fileout: ({"str"}, optional)

99 path to output csv-file for summed budgets

100 observations: [Not yet supported.]

101 Default is None.

102 water_mover: [Not yet supported.]

103 Default is None.

104 timeseries: [Not yet supported.]

105 Default is None.

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

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

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

109 validate: {True, False}

110 Flag to indicate whether the package should be validated upon

111 initialization. This raises a ValidationError if package input is

112 provided in the wrong manner. Defaults to True.

113 """

114

115 _period_data = (

116 "infiltration_rate",

117 "et_pot",

118 "extinction_depth",

119 "extinction_theta",

120 "air_entry_potential",

121 "root_potential",

122 "root_activity",

123 )

124

125 _init_schemata = {

126 "surface_depression_depth": [

127 DTypeSchema(np.floating),

128 BOUNDARY_DIMS_SCHEMA,

129 ],

130 "kv_sat": [

131 DTypeSchema(np.floating),

132 IndexesSchema(),

133 CoordsSchema(("layer",)),

134 BOUNDARY_DIMS_SCHEMA,

135 ],

136 "theta_res": [

137 DTypeSchema(np.floating),

138 BOUNDARY_DIMS_SCHEMA,

139 ],

140 "theta_sat": [

141 DTypeSchema(np.floating),

142 BOUNDARY_DIMS_SCHEMA,

143 ],

144 "theta_init": [

145 DTypeSchema(np.floating),

146 BOUNDARY_DIMS_SCHEMA,

147 ],

148 "epsilon": [

149 DTypeSchema(np.floating),

150 BOUNDARY_DIMS_SCHEMA,

151 ],

152 "infiltration_rate": [

153 DTypeSchema(np.floating),

154 BOUNDARY_DIMS_SCHEMA,

155 ],

156 "et_pot": [

157 DTypeSchema(np.floating),

158 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var

159 ],

160 "extinction_depth": [

161 DTypeSchema(np.floating),

162 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var

163 ],

164 "extinction_theta": [

165 DTypeSchema(np.floating),

166 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var

167 ],

168 "root_potential": [

169 DTypeSchema(np.floating),

170 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var

171 ],

172 "root_activity": [

173 DTypeSchema(np.floating),

174 BOUNDARY_DIMS_SCHEMA | DimsSchema(), # optional var

175 ],

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

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

178 }

179 _write_schemata = {

180 "kv_sat": [

181 AllValueSchema(">", 0.0),

182 OtherCoordsSchema("idomain"),

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

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

185 ],

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

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

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

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

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

191 "infiltration_rate": [IdentityNoDataSchema("stress_period_active")],

192 "et_pot": [IdentityNoDataSchema("stress_period_active")],

193 "extinction_depth": [IdentityNoDataSchema("stress_period_active")],

194 "extinction_theta": [IdentityNoDataSchema("stress_period_active")],

195 "root_potential": [IdentityNoDataSchema("stress_period_active")],

196 "root_activity": [IdentityNoDataSchema("stress_period_active")],

197 }

198

199 _package_data = (

200 "surface_depression_depth",

201 "kv_sat",

202 "theta_res",

203 "theta_sat",

204 "theta_init",

205 "epsilon",

206 )

207 _pkg_id = "uzf"

208

209 _template = BoundaryCondition._initialize_template(_pkg_id)

210

211 @init_log_decorator()

212 def __init__(

213 self,

214 surface_depression_depth,

215 kv_sat,

216 theta_res,

217 theta_sat,

218 theta_init,

219 epsilon,

220 infiltration_rate,

221 et_pot=None,

222 extinction_depth=None,

223 extinction_theta=None,

224 air_entry_potential=None,

225 root_potential=None,

226 root_activity=None,

227 ntrailwaves=7, # Recommended in manual

228 nwavesets=40,

229 groundwater_ET_function=None,

230 simulate_groundwater_seepage=False,

231 print_input=False,

232 print_flows=False,

233 save_flows=False,

234 budget_fileout=None,

235 budgetcsv_fileout=None,

236 observations=None,

237 water_mover=None,

238 timeseries=None,

239 validate: bool = True,

240 ):

241 landflag = self._determine_landflag(kv_sat)

242 iuzno = self._create_uzf_numbers(landflag)

243 ivertcon = self._determine_vertical_connection(iuzno)

244 stress_period_active = landflag.where(landflag == 1)

245

246 dict_dataset = {

247 # Package data

248 "surface_depression_depth": surface_depression_depth,

249 "kv_sat": kv_sat,

250 "theta_res": theta_res,

251 "theta_sat": theta_sat,

252 "theta_init": theta_init,

253 "epsilon": epsilon,

254 # Stress period data

255 "stress_period_active": stress_period_active,

256 "infiltration_rate": infiltration_rate,

257 "et_pot": et_pot,

258 "extinction_depth": extinction_depth,

259 "extinction_theta": extinction_theta,

260 "air_entry_potential": air_entry_potential,

261 "root_potential": root_potential,

262 "root_activity": root_activity,

263 # Dimensions

264 "ntrailwaves": ntrailwaves,

265 "nwavesets": nwavesets,

266 # Options

267 "groundwater_ET_function": groundwater_ET_function,

268 "simulate_gwseep": simulate_groundwater_seepage,

269 "print_input": print_input,

270 "print_flows": print_flows,

271 "save_flows": save_flows,

272 "budget_fileout": budget_fileout,

273 "budgetcsv_fileout": budgetcsv_fileout,

274 "observations": observations,

275 "water_mover": water_mover,

276 "timeseries": timeseries,

277 # Additonal indices for Packagedata

278 "landflag": landflag,

279 "iuzno": iuzno,

280 "ivertcon": ivertcon,

281 }

282 super().__init__(dict_dataset)

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

284 self._check_options(

285 groundwater_ET_function,

286 et_pot,

287 extinction_depth,

288 extinction_theta,

289 air_entry_potential,

290 root_potential,

291 root_activity,

292 )

293 self._validate_init_schemata(validate)

294

295 def fill_stress_perioddata(self):

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

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

298 Only an infiltration rate is required,

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

300 """

301 for var in self._period_data:

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

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

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

305 else:

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

307

308 def _check_options(

309 self,

310 groundwater_ET_function,

311 et_pot,

312 extinction_depth,

313 extinction_theta,

314 air_entry_potential,

315 root_potential,

316 root_activity,

317 ):

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

319 unsat_etae = [

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

321 ]

322

323 if all(simulate_et):

324 self.dataset["simulate_et"] = True

325 elif any(simulate_et):

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

327

328 if extinction_theta is not None:

329 self.dataset["unsat_etwc"] = True

330

331 if all(unsat_etae):

332 self.dataset["unsat_etae"] = True

333 elif any(unsat_etae):

334 raise ValueError(

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

336 )

337

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

339 raise ValueError(

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

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

342 )

343

344 if groundwater_ET_function == "linear":

345 self.dataset["linear_gwet"] = True

346 elif groundwater_ET_function == "square":

347 self.dataset["square_gwet"] = True

348 elif groundwater_ET_function is None:

349 pass

350 else:

351 raise ValueError(

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

353 )

354

355 def _create_uzf_numbers(self, landflag):

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

357 active_nodes = landflag.notnull().astype(np.int8)

358 return np.nancumsum(active_nodes).reshape(landflag.shape) * active_nodes

359

360 def _determine_landflag(self, kv_sat):

361 """returns the landflag for uzf-model. Landflag == 1 for top active UZF-nodes"""

362 land_nodes = get_upper_active_grid_cells(kv_sat).astype(np.int32)

363 return land_nodes.where(kv_sat.notnull())

364

365 def _determine_vertical_connection(self, uzf_number):

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

367

368 def _package_data_to_sparse(self):

369 notnull = self.dataset["landflag"].notnull().to_numpy()

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

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

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

373

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

375

376 layer = ds["layer"].values

377 arrdict = self._ds_to_arrdict(ds)

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

379

380 field_spec = self._get_field_spec_from_dtype(recarr)

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

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

383 field_spec = (

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

385 )

386 sparse_dtype = np.dtype(index_spec + field_spec)

387

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

389 recarr_new["iuzno"] = iuzno

390 recarr_new["landflag"] = landflag

391 recarr_new["ivertcon"] = ivertcon

392 recarr_new[field_names] = recarr

393

394 return recarr_new

395

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

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

398 d = {}

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

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

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

402 )

403 not_options = (

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

405 )

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

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

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

409 # max uzf-cells for which time period data will be supplied

410 d["nuzfcells"] = np.count_nonzero(np.isfinite(d["landflag"]))

411 return self._template.render(d)

412

413 def _to_struct_array(self, arrdict, layer):

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

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

416 # TODO add pkgcheck that period table aligns

417 # Get the number of valid values

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

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

420 nrow = notnull.sum()

421 # Define the numpy structured array dtype

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

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

424 sparse_dtype = np.dtype(index_spec + field_spec)

425

426 # Initialize the structured array

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

428 # Fill in the indices

429 recarr["iuzno"] = iuzno[notnull]

430

431 # Fill in the data

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

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

434

435 return recarr

436

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

438 # Insert additional kwargs

439 kwargs["kv_sat"] = self["kv_sat"]

440 kwargs["stress_period_active"] = self["stress_period_active"]

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

442

443 return errors