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

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"] = 1 if any(v.dataset.save_budget for v in self.values()) else 0 

56 

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) 

69 

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 ) 

76 

77 

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 ) 

86 

87 

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 ) 

96 

97 

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 ) 

106 

107 

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 ) 

116 

117 

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 ) 

126 

127 

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