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

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.timeutil import to_metaswap_timeformat

24 from imod.msw.vegetation import AnnualCropFactors

25

26 REQUIRED_PACKAGES = (

27 GridData,

28 CouplerMapping,

29 Infiltration,

30 LanduseOptions,

31 MeteoGrid,

32 EvapotranspirationMapping,

33 PrecipitationMapping,

34 IdfMapping,

35 TimeOutputControl,

36 AnnualCropFactors,

37)

38

39 INITIAL_CONDITIONS_PACKAGES = (

40 InitialConditionsEquilibrium,

41 InitialConditionsPercolation,

42 InitialConditionsRootzonePressureHead,

43 InitialConditionsSavedState,

44)

45

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

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

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