Coverage for C:\src\imod-python\imod\msw\model.py: 94%

101 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import collections 

2from copy import copy 

3from pathlib import Path 

4from typing import Union 

5 

6import jinja2 

7import numpy as np 

8 

9from imod.msw.coupler_mapping import CouplerMapping 

10from imod.msw.grid_data import GridData 

11from imod.msw.idf_mapping import IdfMapping 

12from imod.msw.infiltration import Infiltration 

13from imod.msw.initial_conditions import ( 

14 InitialConditionsEquilibrium, 

15 InitialConditionsPercolation, 

16 InitialConditionsRootzonePressureHead, 

17 InitialConditionsSavedState, 

18) 

19from imod.msw.landuse import LanduseOptions 

20from imod.msw.meteo_grid import MeteoGrid 

21from imod.msw.meteo_mapping import EvapotranspirationMapping, PrecipitationMapping 

22from imod.msw.output_control import TimeOutputControl 

23from imod.msw.timeutil import to_metaswap_timeformat 

24from imod.msw.vegetation import AnnualCropFactors 

25 

26REQUIRED_PACKAGES = ( 

27 GridData, 

28 CouplerMapping, 

29 Infiltration, 

30 LanduseOptions, 

31 MeteoGrid, 

32 EvapotranspirationMapping, 

33 PrecipitationMapping, 

34 IdfMapping, 

35 TimeOutputControl, 

36 AnnualCropFactors, 

37) 

38 

39INITIAL_CONDITIONS_PACKAGES = ( 

40 InitialConditionsEquilibrium, 

41 InitialConditionsPercolation, 

42 InitialConditionsRootzonePressureHead, 

43 InitialConditionsSavedState, 

44) 

45 

46DEFAULT_SETTINGS = { 

47 "vegetation_mdl": 1, 

48 "evapotranspiration_mdl": 1, 

49 "saltstress_mdl": 0, 

50 "surfacewater_mdl": 0, 

51 "infilimsat_opt": 0, 

52 "netcdf_per": 0, 

53 "postmsw_opt": 0, 

54 "dtgw": 1.0, 

55 "dtsw": 1.0, 

56 "ipstep": 2, 

57 "nxlvage_dim": 366, 

58 "co2": 404.32, 

59 "fact_beta2": 1.0, 

60 "rcsoil": 0.15, 

61 "iterur1": 3, 

62 "iterur2": 5, 

63 "tdbgsm": 91.0, 

64 "tdedsm": 270.0, 

65 "clocktime": 0, 

66} 

67 

68 

69class Model(collections.UserDict): 

70 def __setitem__(self, key, value): 

71 # TODO: Add packagecheck 

72 super().__setitem__(key, value) 

73 

74 def update(self, *args, **kwargs): 

75 for k, v in dict(*args, **kwargs).items(): 

76 self[k] = v 

77 

78 

79class MetaSwapModel(Model): 

80 """ 

81 Contains data and writes consistent model input files 

82 

83 Parameters 

84 ---------- 

85 unsaturated_database: Path-like or str 

86 Path to the MetaSWAP soil physical database folder. 

87 """ 

88 

89 _pkg_id = "model" 

90 _file_name = "para_sim.inp" 

91 

92 _template = jinja2.Template( 

93 "{%for setting, value in settings.items()%}" 

94 "{{setting}} = {{value}}\n" 

95 "{%endfor%}" 

96 ) 

97 

98 def __init__(self, unsaturated_database): 

99 super().__init__() 

100 

101 self.simulation_settings = copy(DEFAULT_SETTINGS) 

102 self.simulation_settings["unsa_svat_path"] = ( 

103 self._render_unsaturated_database_path(unsaturated_database) 

104 ) 

105 

106 def _render_unsaturated_database_path(self, unsaturated_database): 

107 # Force to Path object 

108 unsaturated_database = Path(unsaturated_database) 

109 

110 # Render to string for MetaSWAP 

111 if unsaturated_database.is_absolute(): 

112 return f'"{unsaturated_database}\\"' 

113 else: 

114 # TODO: Test if this is how MetaSWAP accepts relative paths 

115 return f'"${unsaturated_database}\\"' 

116 

117 def _check_required_packages(self): 

118 pkg_types_included = {type(pkg) for pkg in self.values()} 

119 missing_packages = set(REQUIRED_PACKAGES) - pkg_types_included 

120 if len(missing_packages) > 0: 

121 raise ValueError( 

122 f"Missing the following required packages: {missing_packages}" 

123 ) 

124 

125 initial_condition_set = pkg_types_included & set(INITIAL_CONDITIONS_PACKAGES) 

126 if len(initial_condition_set) < 1: 

127 raise ValueError( 

128 "Missing InitialCondition package, assign one of " 

129 f"{INITIAL_CONDITIONS_PACKAGES}" 

130 ) 

131 elif len(initial_condition_set) > 1: 

132 raise ValueError( 

133 "Multiple InitialConditions assigned, choose one of " 

134 f"{initial_condition_set}" 

135 ) 

136 

137 def _check_landuse_indices_in_lookup_options(self): 

138 grid_key = self._get_pkg_key(GridData) 

139 landuse_options_key = self._get_pkg_key(LanduseOptions) 

140 

141 indices_in_grid = set(self[grid_key]["landuse"].values.ravel()) 

142 indices_in_options = set( 

143 self[landuse_options_key].dataset.coords["landuse_index"].values 

144 ) 

145 

146 missing_indices = indices_in_grid - indices_in_options 

147 

148 if len(missing_indices) > 0: 

149 raise ValueError( 

150 "Found the following landuse indices in GridData which " 

151 f"were not in LanduseOptions: {missing_indices}" 

152 ) 

153 

154 def _check_vegetation_indices_in_annual_crop_factors(self): 

155 landuse_options_key = self._get_pkg_key(LanduseOptions) 

156 annual_crop_factors_key = self._get_pkg_key(AnnualCropFactors) 

157 

158 indices_in_options = set( 

159 np.unique(self[landuse_options_key]["vegetation_index"]) 

160 ) 

161 indices_in_crop_factors = set( 

162 self[annual_crop_factors_key].dataset.coords["vegetation_index"].values 

163 ) 

164 

165 missing_indices = indices_in_options - indices_in_crop_factors 

166 

167 if len(missing_indices) > 0: 

168 raise ValueError( 

169 "Found the following vegetation indices in LanduseOptions " 

170 f"which were not in AnnualCropGrowth: {missing_indices}" 

171 ) 

172 

173 def _get_starttime(self): 

174 """ 

175 Loop over all packages to get the minimum time. 

176 

177 MetaSWAP requires a starttime in its simulation settings (para_sim.inp) 

178 """ 

179 

180 starttimes = [] 

181 

182 for pkgname in self: 

183 ds = self[pkgname].dataset 

184 if "time" in ds.coords: 

185 starttimes.append(ds["time"].min().values) 

186 

187 starttime = min(starttimes) 

188 

189 year, time_since_start_year = to_metaswap_timeformat([starttime]) 

190 

191 year = int(year.values) 

192 time_since_start_year = float(time_since_start_year.values) 

193 

194 return year, time_since_start_year 

195 

196 def _get_pkg_key(self, pkg_type: type, optional_package: bool = False): 

197 for pkg_key, pkg in self.items(): 

198 if isinstance(pkg, pkg_type): 

199 return pkg_key 

200 

201 if not optional_package: 

202 raise KeyError(f"Could not find package of type: {pkg_type}") 

203 

204 def write(self, directory: Union[str, Path]): 

205 """ 

206 Write packages and simulation settings (para_sim.inp). 

207 

208 Parameters 

209 ---------- 

210 directory: Path or str 

211 directory to write model in. 

212 """ 

213 

214 # Model checks 

215 self._check_required_packages() 

216 self._check_vegetation_indices_in_annual_crop_factors() 

217 self._check_landuse_indices_in_lookup_options() 

218 

219 # Force to Path 

220 directory = Path(directory) 

221 directory.mkdir(exist_ok=True, parents=True) 

222 

223 # Add time settings 

224 year, time_since_start_year = self._get_starttime() 

225 

226 self.simulation_settings["iybg"] = year 

227 self.simulation_settings["tdbg"] = time_since_start_year 

228 

229 # Add IdfMapping settings 

230 idf_key = self._get_pkg_key(IdfMapping) 

231 self.simulation_settings.update(self[idf_key].get_output_settings()) 

232 

233 filename = directory / self._file_name 

234 with open(filename, "w") as f: 

235 rendered = self._template.render(settings=self.simulation_settings) 

236 f.write(rendered) 

237 

238 # Get index and svat 

239 grid_key = self._get_pkg_key(GridData) 

240 index, svat = self[grid_key].generate_index_array() 

241 

242 # write package contents 

243 for pkgname in self: 

244 self[pkgname].write(directory, index, svat)