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

102 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 10:26 +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.pkgbase import MetaSwapPackage 

24from imod.msw.timeutil import to_metaswap_timeformat 

25from imod.msw.vegetation import AnnualCropFactors 

26 

27REQUIRED_PACKAGES = ( 

28 GridData, 

29 CouplerMapping, 

30 Infiltration, 

31 LanduseOptions, 

32 MeteoGrid, 

33 EvapotranspirationMapping, 

34 PrecipitationMapping, 

35 IdfMapping, 

36 TimeOutputControl, 

37 AnnualCropFactors, 

38) 

39 

40INITIAL_CONDITIONS_PACKAGES = ( 

41 InitialConditionsEquilibrium, 

42 InitialConditionsPercolation, 

43 InitialConditionsRootzonePressureHead, 

44 InitialConditionsSavedState, 

45) 

46 

47DEFAULT_SETTINGS = dict( 

48 vegetation_mdl=1, 

49 evapotranspiration_mdl=1, 

50 saltstress_mdl=0, 

51 surfacewater_mdl=0, 

52 infilimsat_opt=0, 

53 netcdf_per=0, 

54 postmsw_opt=0, 

55 dtgw=1.0, 

56 dtsw=1.0, 

57 ipstep=2, 

58 nxlvage_dim=366, 

59 co2=404.32, 

60 fact_beta2=1.0, 

61 rcsoil=0.15, 

62 iterur1=3, 

63 iterur2=5, 

64 tdbgsm=91.0, 

65 tdedsm=270.0, 

66 clocktime=0, 

67) 

68 

69 

70class Model(collections.UserDict): 

71 def __setitem__(self, key, value): 

72 # TODO: Add packagecheck 

73 super().__setitem__(key, value) 

74 

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

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

77 self[k] = v 

78 

79 

80class MetaSwapModel(Model): 

81 """ 

82 Contains data and writes consistent model input files 

83 

84 Parameters 

85 ---------- 

86 unsaturated_database: Path-like or str 

87 Path to the MetaSWAP soil physical database folder. 

88 """ 

89 

90 _pkg_id = "model" 

91 _file_name = "para_sim.inp" 

92 

93 _template = jinja2.Template( 

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

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

96 "{%endfor%}" 

97 ) 

98 

99 def __init__(self, unsaturated_database): 

100 super().__init__() 

101 

102 self.simulation_settings = copy(DEFAULT_SETTINGS) 

103 self.simulation_settings["unsa_svat_path"] = ( 

104 self._render_unsaturated_database_path(unsaturated_database) 

105 ) 

106 

107 def _render_unsaturated_database_path(self, unsaturated_database): 

108 # Force to Path object 

109 unsaturated_database = Path(unsaturated_database) 

110 

111 # Render to string for MetaSWAP 

112 if unsaturated_database.is_absolute(): 

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

114 else: 

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

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

117 

118 def _check_required_packages(self): 

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

120 missing_packages = set(REQUIRED_PACKAGES) - pkg_types_included 

121 if len(missing_packages) > 0: 

122 raise ValueError( 

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

124 ) 

125 

126 initial_condition_set = pkg_types_included & set(INITIAL_CONDITIONS_PACKAGES) 

127 if len(initial_condition_set) < 1: 

128 raise ValueError( 

129 "Missing InitialCondition package, assign one of " 

130 f"{INITIAL_CONDITIONS_PACKAGES}" 

131 ) 

132 elif len(initial_condition_set) > 1: 

133 raise ValueError( 

134 "Multiple InitialConditions assigned, choose one of " 

135 f"{initial_condition_set}" 

136 ) 

137 

138 def _check_landuse_indices_in_lookup_options(self): 

139 grid_key = self._get_pkg_key(GridData) 

140 landuse_options_key = self._get_pkg_key(LanduseOptions) 

141 

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

143 indices_in_options = set( 

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

145 ) 

146 

147 missing_indices = indices_in_grid - indices_in_options 

148 

149 if len(missing_indices) > 0: 

150 raise ValueError( 

151 "Found the following landuse indices in GridData which " 

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

153 ) 

154 

155 def _check_vegetation_indices_in_annual_crop_factors(self): 

156 landuse_options_key = self._get_pkg_key(LanduseOptions) 

157 annual_crop_factors_key = self._get_pkg_key(AnnualCropFactors) 

158 

159 indices_in_options = set( 

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

161 ) 

162 indices_in_crop_factors = set( 

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

164 ) 

165 

166 missing_indices = indices_in_options - indices_in_crop_factors 

167 

168 if len(missing_indices) > 0: 

169 raise ValueError( 

170 "Found the following vegetation indices in LanduseOptions " 

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

172 ) 

173 

174 def _get_starttime(self): 

175 """ 

176 Loop over all packages to get the minimum time. 

177 

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

179 """ 

180 

181 starttimes = [] 

182 

183 for pkgname in self: 

184 ds = self[pkgname].dataset 

185 if "time" in ds.coords: 

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

187 

188 starttime = min(starttimes) 

189 

190 year, time_since_start_year = to_metaswap_timeformat([starttime]) 

191 

192 year = int(year.values) 

193 time_since_start_year = float(time_since_start_year.values) 

194 

195 return year, time_since_start_year 

196 

197 def _get_pkg_key(self, pkg_type: MetaSwapPackage, optional_package: bool = False): 

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

199 if isinstance(pkg, pkg_type): 

200 return pkg_key 

201 

202 if not optional_package: 

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

204 

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

206 """ 

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

208 

209 Parameters 

210 ---------- 

211 directory: Path or str 

212 directory to write model in. 

213 """ 

214 

215 # Model checks 

216 self._check_required_packages() 

217 self._check_vegetation_indices_in_annual_crop_factors() 

218 self._check_landuse_indices_in_lookup_options() 

219 

220 # Force to Path 

221 directory = Path(directory) 

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

223 

224 # Add time settings 

225 year, time_since_start_year = self._get_starttime() 

226 

227 self.simulation_settings["iybg"] = year 

228 self.simulation_settings["tdbg"] = time_since_start_year 

229 

230 # Add IdfMapping settings 

231 idf_key = self._get_pkg_key(IdfMapping) 

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

233 

234 filename = directory / self._file_name 

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

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

237 f.write(rendered) 

238 

239 # Get index and svat 

240 grid_key = self._get_pkg_key(GridData) 

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

242 

243 # write package contents 

244 for pkgname in self: 

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