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
« 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
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.pkgbase import MetaSwapPackage
24from imod.msw.timeutil import to_metaswap_timeformat
25from imod.msw.vegetation import AnnualCropFactors
27REQUIRED_PACKAGES = (
28 GridData,
29 CouplerMapping,
30 Infiltration,
31 LanduseOptions,
32 MeteoGrid,
33 EvapotranspirationMapping,
34 PrecipitationMapping,
35 IdfMapping,
36 TimeOutputControl,
37 AnnualCropFactors,
38)
40INITIAL_CONDITIONS_PACKAGES = (
41 InitialConditionsEquilibrium,
42 InitialConditionsPercolation,
43 InitialConditionsRootzonePressureHead,
44 InitialConditionsSavedState,
45)
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)
70class Model(collections.UserDict):
71 def __setitem__(self, key, value):
72 # TODO: Add packagecheck
73 super().__setitem__(key, value)
75 def update(self, *args, **kwargs):
76 for k, v in dict(*args, **kwargs).items():
77 self[k] = v
80class MetaSwapModel(Model):
81 """
82 Contains data and writes consistent model input files
84 Parameters
85 ----------
86 unsaturated_database: Path-like or str
87 Path to the MetaSWAP soil physical database folder.
88 """
90 _pkg_id = "model"
91 _file_name = "para_sim.inp"
93 _template = jinja2.Template(
94 "{%for setting, value in settings.items()%}"
95 "{{setting}} = {{value}}\n"
96 "{%endfor%}"
97 )
99 def __init__(self, unsaturated_database):
100 super().__init__()
102 self.simulation_settings = copy(DEFAULT_SETTINGS)
103 self.simulation_settings["unsa_svat_path"] = (
104 self._render_unsaturated_database_path(unsaturated_database)
105 )
107 def _render_unsaturated_database_path(self, unsaturated_database):
108 # Force to Path object
109 unsaturated_database = Path(unsaturated_database)
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}\\"'
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 )
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 )
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)
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 )
147 missing_indices = indices_in_grid - indices_in_options
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 )
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)
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 )
166 missing_indices = indices_in_options - indices_in_crop_factors
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 )
174 def _get_starttime(self):
175 """
176 Loop over all packages to get the minimum time.
178 MetaSWAP requires a starttime in its simulation settings (para_sim.inp)
179 """
181 starttimes = []
183 for pkgname in self:
184 ds = self[pkgname].dataset
185 if "time" in ds.coords:
186 starttimes.append(ds["time"].min().values)
188 starttime = min(starttimes)
190 year, time_since_start_year = to_metaswap_timeformat([starttime])
192 year = int(year.values)
193 time_since_start_year = float(time_since_start_year.values)
195 return year, time_since_start_year
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
202 if not optional_package:
203 raise KeyError(f"Could not find package of type: {pkg_type}")
205 def write(self, directory: Union[str, Path]):
206 """
207 Write packages and simulation settings (para_sim.inp).
209 Parameters
210 ----------
211 directory: Path or str
212 directory to write model in.
213 """
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()
220 # Force to Path
221 directory = Path(directory)
222 directory.mkdir(exist_ok=True, parents=True)
224 # Add time settings
225 year, time_since_start_year = self._get_starttime()
227 self.simulation_settings["iybg"] = year
228 self.simulation_settings["tdbg"] = time_since_start_year
230 # Add IdfMapping settings
231 idf_key = self._get_pkg_key(IdfMapping)
232 self.simulation_settings.update(self[idf_key].get_output_settings())
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)
239 # Get index and svat
240 grid_key = self._get_pkg_key(GridData)
241 index, svat = self[grid_key].generate_index_array()
243 # write package contents
244 for pkgname in self:
245 self[pkgname].write(directory, index, svat)