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

1 from typing import Optional, Tuple

2

3 import numpy as np

4

5 from imod.logging import init_log_decorator, logger

6 from imod.mf6 import GroundwaterFlowModel

7 from imod.mf6.boundary_condition import BoundaryCondition

8 from imod.mf6.interfaces.iregridpackage import IRegridPackage

9 from imod.mf6.utilities.regrid import RegridderType

10 from imod.schemata import DTypeSchema

11

12

13 def 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

21 class 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