Coverage for C:\src\imod-python\imod\mf6\uzf.py: 94%
108 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:25 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:25 +0200
1import numpy as np
2import xarray as xr
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.prepare.layer import get_upper_active_grid_cells
8from imod.schemata import (
9 AllInsideNoDataSchema,
10 AllNoDataSchema,
11 AllValueSchema,
12 CoordsSchema,
13 DimsSchema,
14 DTypeSchema,
15 IdentityNoDataSchema,
16 IndexesSchema,
17 OtherCoordsSchema,
18)
21class UnsaturatedZoneFlow(AdvancedBoundaryCondition):
22 """
23 Unsaturated Zone Flow (UZF) package.
25 TODO: Support timeseries file? Observations? Water Mover?
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 """
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 )
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 }
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"
209 _template = BoundaryCondition._initialize_template(_pkg_id)
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)
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)
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))
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 ]
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")
328 if extinction_theta is not None:
329 self.dataset["unsat_etwc"] = True
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 )
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 )
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 )
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
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())
365 def _determine_vertical_connection(self, uzf_number):
366 return uzf_number.shift(layer=-1, fill_value=0)
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]
374 ds = self.dataset[list(self._package_data)]
376 layer = ds["layer"].values
377 arrdict = self._ds_to_arrdict(ds)
378 recarr = super()._to_struct_array(arrdict, layer)
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)
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
394 return recarr_new
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)
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)
426 # Initialize the structured array
427 recarr = np.empty(nrow, dtype=sparse_dtype)
428 # Fill in the indices
429 recarr["iuzno"] = iuzno[notnull]
431 # Fill in the data
432 for key, arr in arrdict.items():
433 recarr[key] = arr[notnull].astype(np.float64)
435 return recarr
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)
443 return errors