Coverage for C:\src\imod-python\imod\wq\pkggroup.py: 53%
57 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 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"] = (
56 1 if any([v.dataset.save_budget for v in self.values()]) else 0
57 )
59 content = [self._template.format(**d)] # pylint: disable=no-member
60 for i, key in enumerate(self.key_order):
61 system_index = i + 1
62 content.append(
63 self[key]._render(
64 directory=directory.joinpath(key),
65 globaltimes=globaltimes,
66 system_index=system_index,
67 nlayer=nlayer,
68 )
69 )
70 return "".join(content)
72 def render_ssm(self, directory, globaltimes, nlayer):
73 # Only render for the first system, that has concentrations defined.
74 key = self.first_key
75 return self[key]._render_ssm(
76 directory.joinpath(key), globaltimes, nlayer=nlayer
77 )
80class ConstantHeadGroup(PackageGroup):
81 _cellcount_varname = "head_start"
82 _template = (
83 "[chd]\n"
84 " mchdsys = {n_systems}\n"
85 " mxactc = {n_max_active}\n"
86 " ichdcb = {save_budget}"
87 )
90class DrainageGroup(PackageGroup):
91 _cellcount_varname = "elevation"
92 _template = (
93 "[drn]\n"
94 " mdrnsys = {n_systems}\n"
95 " mxactd = {n_max_active}\n"
96 " idrncb = {save_budget}"
97 )
100class GeneralHeadBoundaryGroup(PackageGroup):
101 _cellcount_varname = "head"
102 _template = (
103 "[ghb]\n"
104 " mghbsys = {n_systems}\n"
105 " mxactb = {n_max_active}\n"
106 " ighbcb = {save_budget}"
107 )
110class RiverGroup(PackageGroup):
111 _cellcount_varname = "stage"
112 _template = (
113 "[riv]\n"
114 " mrivsys = {n_systems}\n"
115 " mxactr = {n_max_active}\n"
116 " irivcb = {save_budget}"
117 )
120class WellGroup(PackageGroup):
121 _cellcount_varname = None
122 _template = (
123 "[wel]\n"
124 " mwelsys = {n_systems}\n"
125 " mxactw = {n_max_active}\n"
126 " iwelcb = {save_budget}"
127 )
130# dict might be easier than Enumerator...
131class PackageGroups(enum.Enum):
132 chd = ConstantHeadGroup
133 drn = DrainageGroup
134 ghb = GeneralHeadBoundaryGroup
135 riv = RiverGroup
136 wel = WellGroup