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
« 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
6import jinja2
7import numpy as np
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
26REQUIRED_PACKAGES = (
27 GridData,
28 CouplerMapping,
29 Infiltration,
30 LanduseOptions,
31 MeteoGrid,
32 EvapotranspirationMapping,
33 PrecipitationMapping,
34 IdfMapping,
35 TimeOutputControl,
36 AnnualCropFactors,
37)
39INITIAL_CONDITIONS_PACKAGES = (
40 InitialConditionsEquilibrium,
41 InitialConditionsPercolation,
42 InitialConditionsRootzonePressureHead,
43 InitialConditionsSavedState,
44)
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}
69class Model(collections.UserDict):
70 def __setitem__(self, key, value):
71 # TODO: Add packagecheck
72 super().__setitem__(key, value)
74 def update(self, *args, **kwargs):
75 for k, v in dict(*args, **kwargs).items():
76 self[k] = v
79class MetaSwapModel(Model):
80 """
81 Contains data and writes consistent model input files
83 Parameters
84 ----------
85 unsaturated_database: Path-like or str
86 Path to the MetaSWAP soil physical database folder.
87 """
89 _pkg_id = "model"
90 _file_name = "para_sim.inp"
92 _template = jinja2.Template(
93 "{%for setting, value in settings.items()%}"
94 "{{setting}} = {{value}}\n"
95 "{%endfor%}"
96 )
98 def __init__(self, unsaturated_database):
99 super().__init__()
101 self.simulation_settings = copy(DEFAULT_SETTINGS)
102 self.simulation_settings["unsa_svat_path"] = (
103 self._render_unsaturated_database_path(unsaturated_database)
104 )
106 def _render_unsaturated_database_path(self, unsaturated_database):
107 # Force to Path object
108 unsaturated_database = Path(unsaturated_database)
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}\\"'
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 )
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 )
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)
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 )
146 missing_indices = indices_in_grid - indices_in_options
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 )
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)
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 )
165 missing_indices = indices_in_options - indices_in_crop_factors
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 )
173 def _get_starttime(self):
174 """
175 Loop over all packages to get the minimum time.
177 MetaSWAP requires a starttime in its simulation settings (para_sim.inp)
178 """
180 starttimes = []
182 for pkgname in self:
183 ds = self[pkgname].dataset
184 if "time" in ds.coords:
185 starttimes.append(ds["time"].min().values)
187 starttime = min(starttimes)
189 year, time_since_start_year = to_metaswap_timeformat([starttime])
191 year = int(year.values)
192 time_since_start_year = float(time_since_start_year.values)
194 return year, time_since_start_year
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
201 if not optional_package:
202 raise KeyError(f"Could not find package of type: {pkg_type}")
204 def write(self, directory: Union[str, Path]):
205 """
206 Write packages and simulation settings (para_sim.inp).
208 Parameters
209 ----------
210 directory: Path or str
211 directory to write model in.
212 """
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()
219 # Force to Path
220 directory = Path(directory)
221 directory.mkdir(exist_ok=True, parents=True)
223 # Add time settings
224 year, time_since_start_year = self._get_starttime()
226 self.simulation_settings["iybg"] = year
227 self.simulation_settings["tdbg"] = time_since_start_year
229 # Add IdfMapping settings
230 idf_key = self._get_pkg_key(IdfMapping)
231 self.simulation_settings.update(self[idf_key].get_output_settings())
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)
238 # Get index and svat
239 grid_key = self._get_pkg_key(GridData)
240 index, svat = self[grid_key].generate_index_array()
242 # write package contents
243 for pkgname in self:
244 self[pkgname].write(directory, index, svat)