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

60 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +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": list( 

73 zip( 

74 self["package_names"].values, 

75 self["concentration_boundary_type"].values, 

76 self["auxiliary_variable_name"].values, 

77 ) 

78 ), 

79 } 

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

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

82 if self._valid(value): 

83 d[var] = value 

84 return self._template.render(d) 

85 

86 @staticmethod 

87 def from_flow_model( 

88 model: GroundwaterFlowModel, 

89 species: str, 

90 print_flows: bool = False, 

91 save_flows: bool = False, 

92 validate: bool = True, 

93 is_split: bool = False, 

94 ): 

95 """ 

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

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

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

99 

100 Parameters 

101 ---------- 

102 model: GroundwaterFlowModel 

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

104 inferred. 

105 species: str 

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

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

108 argument. 

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

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

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

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

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

114 for the last time step of each stress period. 

115 Default is False. 

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

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

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

119 Default is False. 

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

121 Flag to indicate whether the package should be validated upon 

122 initialization. This raises a ValidationError if package input is 

123 provided in the wrong manner. Defaults to True. 

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

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

126 """ 

127 if not isinstance(model, GroundwaterFlowModel): 

128 raise TypeError( 

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

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

131 ) 

132 names = [] 

133 boundary_types = [] 

134 aux_var_names = [] 

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

136 if isinstance(package, BoundaryCondition): 

137 ds = package.dataset 

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

139 # species coordinate. 

140 if "concentration" not in ds.data_vars: 

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

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

143 raise ValueError( 

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

145 ) 

146 

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

148 # species to have different mixing behavior. 

149 type_da = ds["concentration_boundary_type"] 

150 if "species" in type_da.dims: 

151 type_da = type_da.sel(species=species) 

152 

153 names.append(name) 

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

155 aux_var_names.append(species) 

156 

157 if len(names) == 0: 

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

159 if is_split: 

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

161 return 

162 else: 

163 raise ValueError(msg) 

164 

165 return SourceSinkMixing( 

166 names, 

167 boundary_types, 

168 aux_var_names, 

169 print_flows=print_flows, 

170 save_flows=save_flows, 

171 validate=validate, 

172 ) 

173 

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

175 return self._regrid_method