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

1import collections 

2import os 

3from pathlib import Path 

4from typing import Optional, Tuple, Union 

5 

6import numpy as np 

7 

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 

15 

16OUTPUT_EXT_MAPPING = { 

17 "head": "hds", 

18 "concentration": "ucn", 

19 "budget": "cbc", 

20} 

21 

22 

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 

29 

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. 

34 

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. 

53 

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. 

60 

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) 

65 

66 """ 

67 

68 _pkg_id = "oc" 

69 _keyword_map = {} 

70 _template = Package._initialize_template(_pkg_id) 

71 

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 } 

83 

84 _write_schemata = {} 

85 _regrid_method: dict[str, Tuple[RegridderType, str]] = {} 

86 

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 

103 

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.") 

106 

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) 

117 

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 ) 

133 

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 

138 

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) 

148 

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)) 

155 

156 return path 

157 

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

159 d = {} 

160 

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() 

167 

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) 

179 

180 else: 

181 setting = self.dataset[datavar].item() 

182 periods[1][key] = self._get_ocsetting(setting) 

183 

184 d["periods"] = periods 

185 

186 return self._template.render(d) 

187 

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) 

197 

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 

208 

209 @property 

210 def is_budget_output(self) -> bool: 

211 return self.dataset["save_budget"].values[()] is not None 

212 

213 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: 

214 return self._regrid_method