Coverage for C:\src\imod-python\imod\flow\pkggroup.py: 98%
42 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1import abc
2import collections
3import enum
5from imod.flow.timeutil import insert_unique_package_times
6from imod.util.nested_dict import append_nested_dict, initialize_nested_dict
9class PackageGroup(collections.UserDict, abc.ABC):
10 """
11 Groups for packs that support multiple systems:
12 * chd
13 * drn
14 * ghb
15 * riv
16 * wel
17 """
19 __slots__ = ["_n_variables"]
21 def __init__(self, **kwargs):
22 collections.UserDict.__init__(self)
23 for k, v in kwargs.items():
24 self[k] = v
25 # self.reorder_keys()
27 def compose(self, directory, globaltimes, nlayer):
28 are_periodic = [pkg._is_periodic() for pkg in self.values()]
29 have_time = [pkg._hastime() for pkg in self.values()]
31 # Raise error if one system is periodic and one is not
32 all_or_not_any_periodic = all(are_periodic) or (not any(are_periodic))
33 if not all_or_not_any_periodic:
34 raise ValueError(
35 "At least one system is periodic "
36 "and at least one system is not periodic. \n"
37 "Please insert all systems as periodic or not. "
38 )
40 if (not any(are_periodic)) & (any(have_time)):
41 pkggroup_times, _ = insert_unique_package_times(self.items())
42 else:
43 # FUTURE: We could catch edge case here where different periods
44 # specified for different systems. This is uncommon practice.
45 pkggroup_times = None
47 composition = initialize_nested_dict(5)
49 for i, (key, pkg) in enumerate(self.items()):
50 system_index = i + 1
51 composed_pkg = pkg.compose(
52 directory.joinpath(key),
53 globaltimes,
54 nlayer,
55 system_index=system_index,
56 pkggroup_time=pkggroup_times,
57 )
58 append_nested_dict(composition, composed_pkg)
60 return composition
63class ConstantHeadGroup(PackageGroup):
64 _n_variables = 1
67class DrainageGroup(PackageGroup):
68 _n_variables = 2
71class GeneralHeadBoundaryGroup(PackageGroup):
72 _n_variables = 2
75class RiverGroup(PackageGroup):
76 _n_variables = 4
79class WellGroup(PackageGroup):
80 _n_variables = 1
83# dict might be easier than Enumerator...
84class PackageGroups(enum.Enum):
85 chd = ConstantHeadGroup
86 drn = DrainageGroup
87 ghb = GeneralHeadBoundaryGroup
88 riv = RiverGroup
89 wel = WellGroup