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

1import abc 

2import collections 

3import enum 

4 

5 

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 """ 

15 

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() 

21 

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 

42 

43 def max_n_sinkssources(self): 

44 return sum(pkg._ssm_cellcount for pkg in self.values()) 

45 

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 ) 

58 

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) 

71 

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 ) 

78 

79 

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 ) 

88 

89 

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 ) 

98 

99 

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 ) 

108 

109 

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 ) 

118 

119 

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 ) 

128 

129 

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