Coverage for C:\src\imod-python\imod\mf6\buy.py: 91%

46 statements  

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

1from typing import Optional, Sequence 

2 

3import numpy as np 

4import xarray as xr 

5 

6from imod.logging import init_log_decorator 

7from imod.mf6.package import Package 

8from imod.schemata import DTypeSchema 

9 

10 

11def assign_index(arg): 

12 if isinstance(arg, xr.DataArray): 

13 arg = arg.values 

14 elif not isinstance(arg, (np.ndarray, list, tuple)): 

15 raise TypeError("should be a tuple, list, or numpy array") 

16 

17 arr = np.array(arg) 

18 if arr.ndim != 1: 

19 raise ValueError("should be 1D") 

20 

21 return xr.DataArray(arr, dims=("index",)) 

22 

23 

24class Buoyancy(Package): 

25 """ 

26 Buoyancy package. This package must be included when performing variable 

27 density simulation. 

28 

29 Note that ``reference_density`` is a single value, but 

30 ``density_concentration_slope``, ``reference_concentration`` and 

31 ``modelname`` require an entry for every active species. Please refer to the 

32 examples. 

33 

34 Parameters 

35 ---------- 

36 reference_density: float, 

37 fluid reference density used in the equation of state. 

38 density_concentration_slope: sequence of floats 

39 Slope of the (linear) density concentration line used in the density 

40 equation of state. 

41 reference_concentration: sequence of floats 

42 Reference concentration used in the density equation of state. 

43 modelname: sequence of strings, 

44 Name of the GroundwaterTransport (GWT) model used for the 

45 concentrations. 

46 species: sequence of str, 

47 Name of the species used to calculate a density value. 

48 hhformulation_rhs: bool, optional. 

49 use the variable-density hydraulic head formulation and add off-diagonal 

50 terms to the right-hand. This option will prevent the BUY Package from 

51 adding asymmetric terms to the flow matrix. Default value is ``False``. 

52 densityfile: 

53 name of the binary output file to write density information. The density 

54 file has the same format as the head file. Density values will be 

55 written to the density file whenever heads are written to the binary 

56 head file. The settings for controlling head output are contained in the 

57 Output Control option. 

58 validate: {True, False} 

59 Flag to indicate whether the package should be validated upon 

60 initialization. This raises a ValidationError if package input is 

61 provided in the wrong manner. Defaults to True. 

62 

63 Examples 

64 -------- 

65 

66 The buoyancy input for a single species called "salinity", which is 

67 simulated by a GWT model called "gwt-1" are specified as follows: 

68 

69 >>> buy = imod.mf6.Buoyance( 

70 ... reference_density=1000.0, 

71 ... density_concentration_slope=[0.025], 

72 ... reference_concentration=[0.0], 

73 ... modelname=["gwt-1"], 

74 ... species=["salinity"], 

75 ... ) 

76 

77 Multiple species can be specified by presenting multiple values with an 

78 associated species coordinate. Two species, called "c1" and "c2", simulated 

79 by the GWT models "gwt-1" and "gwt-2" are specified as: 

80 

81 >>> coords = {"species": ["c1", "c2"]} 

82 >>> buy = imod.mf6.Buoyance( 

83 ... reference_density=1000.0, 

84 ... density_concentration_slope=[0.025, 0.01], 

85 ... reference_concentration=[0.0, 0.0], 

86 ... modelname=["gwt-1", "gwt-2"], 

87 ... species=["c1", "c2"], 

88 ... ) 

89 """ 

90 

91 _pkg_id = "buy" 

92 _template = Package._initialize_template(_pkg_id) 

93 

94 _init_schemata = { 

95 "reference_density": [DTypeSchema(np.floating)], 

96 "density_concentration_slope": [DTypeSchema(np.floating)], 

97 "reference_concentration": [DTypeSchema(np.floating)], 

98 } 

99 

100 _write_schemata = {} 

101 

102 @init_log_decorator() 

103 def __init__( 

104 self, 

105 reference_density: float, 

106 density_concentration_slope: Sequence[float], 

107 reference_concentration: Sequence[float], 

108 modelname: Sequence[str], 

109 species: Sequence[str], 

110 hhformulation_rhs: bool = False, 

111 densityfile: Optional[str] = None, 

112 validate: bool = True, 

113 ): 

114 dict_dataset = { 

115 "reference_density": reference_density, 

116 # Assign a shared index: this also forces equal lenghts 

117 "density_concentration_slope": assign_index(density_concentration_slope), 

118 "reference_concentration": assign_index(reference_concentration), 

119 "modelname": assign_index(modelname), 

120 "species": assign_index(species), 

121 "hhformulation_rhs": hhformulation_rhs, 

122 "densityfile": densityfile, 

123 } 

124 super().__init__(dict_dataset) 

125 self._validate_init_schemata(validate) 

126 

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

128 ds = self.dataset 

129 packagedata = [] 

130 

131 for i, (a, b, c, d) in enumerate( 

132 zip( 

133 ds["density_concentration_slope"].values, 

134 ds["reference_concentration"].values, 

135 ds["modelname"].values, 

136 ds["species"].values, 

137 ) 

138 ): 

139 packagedata.append((i + 1, a, b, c, d)) 

140 

141 d = { 

142 "nrhospecies": self.dataset["species"].size, 

143 "packagedata": packagedata, 

144 } 

145 

146 for varname in ["hhformulation_rhs", "reference_density", "densityfile"]: 

147 value = self.dataset[varname].values[()] 

148 if self._valid(value): 

149 d[varname] = value 

150 

151 return self._template.render(d) 

152 

153 def update_transport_models(self, new_modelnames: Sequence[str]): 

154 """ 

155 The names of the transport models can change in some cases, for example 

156 when partitioning. Use this function to update the names of the 

157 transport models. 

158 """ 

159 transport_model_names = self.get_transport_model_names() 

160 if not len(transport_model_names) == len(new_modelnames): 

161 raise ValueError("the number of transport models cannot be changed.") 

162 for modelname, new_modelname in zip(transport_model_names, new_modelnames): 

163 if not modelname in new_modelname: 

164 raise ValueError( 

165 "new transport model names do not match the old ones. The new names should be equal to the old ones, with a suffix." 

166 ) 

167 self.dataset["modelname"] = assign_index(new_modelnames) 

168 

169 def get_transport_model_names(self) -> list[str]: 

170 """ 

171 Returns the names of the transport models used by this buoyancy package. 

172 """ 

173 return list(self.dataset["modelname"].values)