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

1 import collections

2 from copy import copy

3 from pathlib import Path

4 from typing import Union

5

6 import jinja2

7 import numpy as np

8

9 from imod.msw.coupler_mapping import CouplerMapping

10 from imod.msw.grid_data import GridData

11 from imod.msw.idf_mapping import IdfMapping

12 from imod.msw.infiltration import Infiltration

13 from imod.msw.initial_conditions import (

14 InitialConditionsEquilibrium,

15 InitialConditionsPercolation,

16 InitialConditionsRootzonePressureHead,

17 InitialConditionsSavedState,

18)

19 from imod.msw.landuse import LanduseOptions

20 from imod.msw.meteo_grid import MeteoGrid

21 from imod.msw.meteo_mapping import EvapotranspirationMapping, PrecipitationMapping

22 from imod.msw.output_control import TimeOutputControl

23 from imod.msw.pkgbase import MetaSwapPackage

24 from imod.msw.timeutil import to_metaswap_timeformat

25 from imod.msw.vegetation import AnnualCropFactors

26

27 REQUIRED_PACKAGES = (

28 GridData,

29 CouplerMapping,

30 Infiltration,

31 LanduseOptions,

32 MeteoGrid,

33 EvapotranspirationMapping,

34 PrecipitationMapping,

35 IdfMapping,

36 TimeOutputControl,

37 AnnualCropFactors,

38)

39

40 INITIAL_CONDITIONS_PACKAGES = (

41 InitialConditionsEquilibrium,

42 InitialConditionsPercolation,

43 InitialConditionsRootzonePressureHead,

44 InitialConditionsSavedState,

45)

46

47 DEFAULT_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

70 class 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

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