Coverage for C:\src\imod-python\imod\mf6\oc.py: 97%
91 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1import collections
2import os
3from pathlib import Path
4from typing import Optional, Tuple, Union
6import numpy as np
8from imod.logging import init_log_decorator
9from imod.mf6.interfaces.iregridpackage import IRegridPackage
10from imod.mf6.package import Package
11from imod.mf6.utilities.dataset import is_dataarray_none
12from imod.mf6.utilities.regrid import RegridderType
13from imod.mf6.write_context import WriteContext
14from imod.schemata import DTypeSchema
16OUTPUT_EXT_MAPPING = {
17 "head": "hds",
18 "concentration": "ucn",
19 "budget": "cbc",
20}
23class OutputControl(Package, IRegridPackage):
24 """
25 The Output Control Option determines how and when heads, budgets and/or
26 concentrations are printed to the listing file and/or written to a separate
27 binary output file.
28 https://water.usgs.gov/water-resources/software/MODFLOW-6/mf6io_6.4.2.pdf#page=53
30 Currently the settings "first", "last", "all", and "frequency"
31 are supported, the "steps" setting is not supported, because of
32 its ragged nature. Furthermore, only one setting per stress period
33 can be specified in imod-python.
35 Parameters
36 ----------
37 save_head : {string, integer}, or xr.DataArray of {string, integer}, optional
38 String or integer indicating output control for head file (.hds)
39 If string, should be one of ["first", "last", "all"].
40 If integer, interpreted as frequency.
41 save_budget : {string, integer}, or xr.DataArray of {string, integer}, optional
42 String or integer indicating output control for cell budgets (.cbc)
43 If string, should be one of ["first", "last", "all"].
44 If integer, interpreted as frequency.
45 save_concentration : {string, integer}, or xr.DataArray of {string, integer}, optional
46 String or integer indicating output control for concentration file (.ucn)
47 If string, should be one of ["first", "last", "all"].
48 If integer, interpreted as frequency.
49 validate: {True, False}
50 Flag to indicate whether the package should be validated upon
51 initialization. This raises a ValidationError if package input is
52 provided in the wrong manner. Defaults to True.
54 Examples
55 --------
56 To specify a mix of both 'frequency' and 'first' setting,
57 we need to specify an array with both integers and strings.
58 For this we need to create a numpy object array first,
59 otherwise xarray converts all to strings automatically.
61 >>> time = [np.datetime64("2000-01-01"), np.datetime64("2000-01-02")]
62 >>> data = np.array(["last", 5], dtype="object")
63 >>> save_head = xr.DataArray(data, coords={"time": time}, dims=("time"))
64 >>> oc = imod.mf6.OutputControl(save_head=save_head, save_budget=None, save_concentration=None)
66 """
68 _pkg_id = "oc"
69 _keyword_map = {}
70 _template = Package._initialize_template(_pkg_id)
72 _init_schemata = {
73 "save_head": [
74 DTypeSchema(np.integer) | DTypeSchema(str) | DTypeSchema(object),
75 ],
76 "save_budget": [
77 DTypeSchema(np.integer) | DTypeSchema(str) | DTypeSchema(object),
78 ],
79 "save_concentration": [
80 DTypeSchema(np.integer) | DTypeSchema(str) | DTypeSchema(object),
81 ],
82 }
84 _write_schemata = {}
85 _regrid_method: dict[str, Tuple[RegridderType, str]] = {}
87 @init_log_decorator()
88 def __init__(
89 self,
90 save_head=None,
91 save_budget=None,
92 save_concentration=None,
93 head_file=None,
94 budget_file=None,
95 concentration_file=None,
96 validate: bool = True,
97 ):
98 save_concentration = (
99 None if is_dataarray_none(save_concentration) else save_concentration
100 )
101 save_head = None if is_dataarray_none(save_head) else save_head
102 save_budget = None if is_dataarray_none(save_budget) else save_budget
104 if save_head is not None and save_concentration is not None:
105 raise ValueError("save_head and save_concentration cannot both be defined.")
107 dict_dataset = {
108 "save_head": save_head,
109 "save_concentration": save_concentration,
110 "save_budget": save_budget,
111 "head_file": head_file,
112 "budget_file": budget_file,
113 "concentration_file": concentration_file,
114 }
115 super().__init__(dict_dataset)
116 self._validate_init_schemata(validate)
118 def _get_ocsetting(self, setting):
119 """Get oc setting based on its type. If integers return f'frequency {setting}', if"""
120 if isinstance(setting, (int, np.integer)) and not isinstance(setting, bool):
121 return f"frequency {setting}"
122 elif isinstance(setting, str):
123 if setting.lower() in ["first", "last", "all"]:
124 return setting.lower()
125 else:
126 raise ValueError(
127 f"Output Control received wrong string. String should be one of ['first', 'last', 'all'], instead got {setting}"
128 )
129 else:
130 raise TypeError(
131 f"Output Control setting should be either integer or string in ['first', 'last', 'all'], instead got {setting}"
132 )
134 def _get_output_filepath(self, directory: Path, output_variable: str) -> Path:
135 varname = f"{output_variable}_file"
136 ext = OUTPUT_EXT_MAPPING[output_variable]
137 modelname = directory.stem
139 filepath = self.dataset[varname].values[()]
140 if filepath is None:
141 filepath = directory / f"{modelname}.{ext}"
142 else:
143 if not isinstance(filepath, str | Path):
144 raise ValueError(
145 f"{varname} should be of type str or Path. However it is of type {type(filepath)}"
146 )
147 filepath = Path(filepath)
149 if filepath.is_absolute():
150 path = filepath
151 else:
152 # Get path relative to the simulation name file.
153 sim_directory = directory.parent
154 path = Path(os.path.relpath(filepath, sim_directory))
156 return path
158 def render(self, directory, pkgname, globaltimes, binary):
159 d = {}
161 for output_variable in OUTPUT_EXT_MAPPING.keys():
162 save = self.dataset[f"save_{output_variable}"].values[()]
163 if save is not None:
164 varname = f"{output_variable}_file"
165 output_path = self._get_output_filepath(directory, output_variable)
166 d[varname] = output_path.as_posix()
168 periods = collections.defaultdict(dict)
169 for datavar in ("save_head", "save_concentration", "save_budget"):
170 if self.dataset[datavar].values[()] is None:
171 continue
172 key = datavar.replace("_", " ")
173 if "time" in self.dataset[datavar].coords:
174 package_times = self.dataset[datavar].coords["time"].values
175 starts = np.searchsorted(globaltimes, package_times) + 1
176 for i, s in enumerate(starts):
177 setting = self.dataset[datavar].isel(time=i).item()
178 periods[s][key] = self._get_ocsetting(setting)
180 else:
181 setting = self.dataset[datavar].item()
182 periods[1][key] = self._get_ocsetting(setting)
184 d["periods"] = periods
186 return self._template.render(d)
188 def write(
189 self,
190 pkgname: str,
191 globaltimes: Union[list[np.datetime64], np.ndarray],
192 write_context: WriteContext,
193 ):
194 # We need to overload the write here to ensure the output directory is
195 # created in advance for MODFLOW6.
196 super().write(pkgname, globaltimes, write_context)
198 for datavar in ("head_file", "concentration_file", "budget_file"):
199 path = self.dataset[datavar].values[()]
200 if path is not None:
201 if not isinstance(path, str):
202 raise ValueError(
203 f"{path} should be of type str. However it is of type {type(path)}"
204 )
205 filepath = Path(path)
206 filepath.parent.mkdir(parents=True, exist_ok=True)
207 return
209 @property
210 def is_budget_output(self) -> bool:
211 return self.dataset["save_budget"].values[()] is not None
213 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]:
214 return self._regrid_method