Coverage for C:\src\imod-python\imod\wq\pkggroup.py: 100%
57 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 abc
2import collections
3import enum
6class PackageGroup(collections.UserDict, abc.ABC):
7 """
8 Groups for packes that support multiple systems:
9 * chd
10 * drn
11 * ghb
12 * riv
13 * wel
14 """
16 def __init__(self, **kwargs):
17 collections.UserDict.__init__(self)
18 for k, v in kwargs.items():
19 self[k] = v
20 self.reorder_keys()
22 def reorder_keys(self):
23 """
24 Order packages so that the one with with concentration is first.
25 Check whether concentration for only one system has been defined.
26 """
27 n_system_concentrations = 0
28 order = []
29 for k, v in self.items():
30 if "concentration" in v.dataset.data_vars:
31 n_system_concentrations += 1
32 order.insert(0, k)
33 else:
34 order.append(k)
35 if n_system_concentrations > 1:
36 raise ValueError(
37 f"Multiple systems with concentrations detected: {order}\n"
38 "Only one system with concentration is allowed per package kind."
39 )
40 self.first_key = order[0]
41 self.key_order = order
43 def max_n_sinkssources(self):
44 return sum(pkg._ssm_cellcount for pkg in self.values())
46 def render(self, directory, globaltimes, nlayer, nrow, ncol):
47 d = {}
48 d["n_systems"] = len(self.keys())
49 d["n_max_active"] = sum(
50 [
51 v._max_active_n(self._cellcount_varname, nlayer, nrow, ncol) # pylint:disable=no-member
52 for v in self.values()
53 ]
54 )
55 d["save_budget"] = 1 if any(v.dataset.save_budget for v in self.values()) else 0
57 content = [self._template.format(**d)] # pylint: disable=no-member
58 for i, key in enumerate(self.key_order):
59 system_index = i + 1
60 content.append(
61 self[key]._render(
62 directory=directory.joinpath(key),
63 globaltimes=globaltimes,
64 system_index=system_index,
65 nlayer=nlayer,
66 )
67 )
68 return "".join(content)
70 def render_ssm(self, directory, globaltimes, nlayer):
71 # Only render for the first system, that has concentrations defined.
72 key = self.first_key
73 return self[key]._render_ssm(
74 directory.joinpath(key), globaltimes, nlayer=nlayer
75 )
78class ConstantHeadGroup(PackageGroup):
79 _cellcount_varname = "head_start"
80 _template = (
81 "[chd]\n"
82 " mchdsys = {n_systems}\n"
83 " mxactc = {n_max_active}\n"
84 " ichdcb = {save_budget}"
85 )
88class DrainageGroup(PackageGroup):
89 _cellcount_varname = "elevation"
90 _template = (
91 "[drn]\n"
92 " mdrnsys = {n_systems}\n"
93 " mxactd = {n_max_active}\n"
94 " idrncb = {save_budget}"
95 )
98class GeneralHeadBoundaryGroup(PackageGroup):
99 _cellcount_varname = "head"
100 _template = (
101 "[ghb]\n"
102 " mghbsys = {n_systems}\n"
103 " mxactb = {n_max_active}\n"
104 " ighbcb = {save_budget}"
105 )
108class RiverGroup(PackageGroup):
109 _cellcount_varname = "stage"
110 _template = (
111 "[riv]\n"
112 " mrivsys = {n_systems}\n"
113 " mxactr = {n_max_active}\n"
114 " irivcb = {save_budget}"
115 )
118class WellGroup(PackageGroup):
119 _cellcount_varname = None
120 _template = (
121 "[wel]\n"
122 " mwelsys = {n_systems}\n"
123 " mxactw = {n_max_active}\n"
124 " iwelcb = {save_budget}"
125 )
128# dict might be easier than Enumerator...
129class PackageGroups(enum.Enum):
130 chd = ConstantHeadGroup
131 drn = DrainageGroup
132 ghb = GeneralHeadBoundaryGroup
133 riv = RiverGroup
134 wel = WellGroup