Coverage for C:\src\imod-python\imod\mf6\ssm.py: 85%

60 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1from typing import Optional, Tuple 

2 

3import numpy as np 

4 

5from imod.logging import init_log_decorator, logger 

6from imod.mf6 import GroundwaterFlowModel 

7from imod.mf6.boundary_condition import BoundaryCondition 

8from imod.mf6.interfaces.iregridpackage import IRegridPackage 

9from imod.mf6.utilities.regrid import RegridderType 

10from imod.schemata import DTypeSchema 

11 

12 

13def with_index_dim(array_like): 

14 # At least1d will also extract the values if array_like is a DataArray. 

15 arr1d = np.atleast_1d(array_like) 

16 if arr1d.ndim > 1: 

17 raise ValueError("array must be 1d") 

18 return ("index", arr1d) 

19 

20 

21class SourceSinkMixing(BoundaryCondition, IRegridPackage): 

22 """ 

23 Parameters 

24 ---------- 

25 package_names: array_like of str 

26 concentration_boundary_type: array_like of str 

27 auxiliary_variable_name: array_like of str 

28 print_flows: bool 

29 save_flows: bool 

30 validate: bool 

31 """ 

32 

33 _pkg_id = "ssm" 

34 _template = BoundaryCondition._initialize_template(_pkg_id) 

35 

36 _init_schemata = { 

37 "package_names": [DTypeSchema(np.str_)], 

38 "concentration_boundary_type": [DTypeSchema(np.str_)], 

39 "auxiliary_variable_name": [DTypeSchema(np.str_)], 

40 "print_flows": [DTypeSchema(np.bool_)], 

41 "save_flows": [DTypeSchema(np.bool_)], 

42 } 

43 

44 _write_schemata = {} 

45 

46 _regrid_method: dict[str, Tuple[RegridderType, str]] = {} 

47 

48 @init_log_decorator() 

49 def __init__( 

50 self, 

51 package_names, 

52 concentration_boundary_type, 

53 auxiliary_variable_name, 

54 print_flows: bool = False, 

55 save_flows: bool = False, 

56 validate: bool = True, 

57 ): 

58 dict_dataset = { 

59 # By sharing the index, this will raise an error if lengths do not 

60 # match. 

61 "package_names": with_index_dim(package_names), 

62 "concentration_boundary_type": with_index_dim(concentration_boundary_type), 

63 "auxiliary_variable_name": with_index_dim(auxiliary_variable_name), 

64 "print_flows": print_flows, 

65 "save_flows": save_flows, 

66 } 

67 super().__init__(dict_dataset) 

68 self._validate_init_schemata(validate) 

69 

70 def render(self, directory, pkgname, globaltimes, binary): 

71 d = { 

72 "sources": [ 

73 (a, b, c) 

74 for a, b, c in zip( 

75 self["package_names"].values, 

76 self["concentration_boundary_type"].values, 

77 self["auxiliary_variable_name"].values, 

78 ) 

79 ], 

80 } 

81 for var in ("print_flows", "save_flows"): 

82 value = self[var].values[()] 

83 if self._valid(value): 

84 d[var] = value 

85 return self._template.render(d) 

86 

87 @staticmethod 

88 def from_flow_model( 

89 model: GroundwaterFlowModel, 

90 species: str, 

91 print_flows: bool = False, 

92 save_flows: bool = False, 

93 validate: bool = True, 

94 is_split: bool = False, 

95 ): 

96 """ 

97 Derive a Source and Sink Mixing package from a Groundwater Flow model's 

98 boundary conditions (e.g. GeneralHeadBoundary). Boundary condition 

99 packages which have the ``concentration`` variable set are included. 

100 

101 Parameters 

102 ---------- 

103 model: GroundwaterFlowModel 

104 Groundwater flow model from which sources & sinks have to be 

105 inferred. 

106 species: str 

107 Name of species to create a transport model for. This name will be 

108 looked for in the ``species`` dimensions of the ``concentration`` 

109 argument. 

110 print_flows: ({True, False}, optional) 

111 Indicates that the list of general head boundary flow rates will be 

112 printed to the listing file for every stress period time step in which 

113 "BUDGET PRINT" is specified in Output Control. If there is no Output 

114 Control option and PRINT FLOWS is specified, then flow rates are printed 

115 for the last time step of each stress period. 

116 Default is False. 

117 save_flows: ({True, False}, optional) 

118 Indicates that general head boundary flow terms will be written to the 

119 file specified with "BUDGET FILEOUT" in Output Control. 

120 Default is False. 

121 validate: ({True, False}, optional) 

122 Flag to indicate whether the package should be validated upon 

123 initialization. This raises a ValidationError if package input is 

124 provided in the wrong manner. Defaults to True. 

125 is_split: ({True, False}, optional) 

126 Flag to indicate if the simulation has been split into partitions. 

127 """ 

128 if not isinstance(model, GroundwaterFlowModel): 

129 raise TypeError( 

130 "model must be a GroundwaterFlowModel, received instead: " 

131 f"{type(model).__name__}" 

132 ) 

133 names = [] 

134 boundary_types = [] 

135 aux_var_names = [] 

136 for name, package in model.items(): 

137 if isinstance(package, BoundaryCondition): 

138 ds = package.dataset 

139 # The package should contain a concentration variable, with a 

140 # species coordinate. 

141 if "concentration" not in ds.data_vars: 

142 raise ValueError(f"concentration not present in package {name}") 

143 if "species" not in ds["concentration"].coords: 

144 raise ValueError( 

145 f"No species coordinate for concentration in package {name}" 

146 ) 

147 

148 # While somewhat far-fetched, it is possible for different 

149 # species to have different mixing behavior. 

150 type_da = ds["concentration_boundary_type"] 

151 if "species" in type_da.dims: 

152 type_da = type_da.sel(species=species) 

153 

154 names.append(name) 

155 boundary_types.append(type_da.values[()]) 

156 aux_var_names.append(species) 

157 

158 if len(names) == 0: 

159 msg = "flow model does not contain boundary conditions" 

160 if is_split: 

161 logger.info(f"{msg}, returning None instead of SourceSinkMixing") 

162 return 

163 else: 

164 raise ValueError(msg) 

165 

166 return SourceSinkMixing( 

167 names, 

168 boundary_types, 

169 aux_var_names, 

170 print_flows=print_flows, 

171 save_flows=save_flows, 

172 validate=validate, 

173 ) 

174 

175 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: 

176 return self._regrid_method