Coverage for src / gwtransport / gamma.py: 87%

1""" 

2Gamma Distribution Utilities for Aquifer Pore Volume Heterogeneity. 

3 

4This module provides utilities for working with gamma distributions to model heterogeneous 

5aquifer pore volumes in groundwater transport analysis. The gamma distribution offers a 

6flexible two-parameter model for representing the natural variability in flow path lengths 

7and residence times within aquifer systems. In heterogeneous aquifers, water travels through 

8multiple flow paths with different pore volumes, and the gamma distribution provides a 

9realistic representation of this heterogeneity. 

10 

11Available functions: 

12 

13- :func:`parse_parameters` - Parse and validate gamma distribution parameters from either 

14 (alpha, beta) or (mean, std). Ensures exactly one parameter pair is provided and validates 

15 positivity constraints. 

16 

17- :func:`mean_std_to_alpha_beta` - Convert physically intuitive (mean, std) parameters to 

18 gamma shape/scale parameters. Uses formulas: alpha = mean^2 / std^2 and beta = std^2 / mean. 

19 

20- :func:`alpha_beta_to_mean_std` - Convert gamma (alpha, beta) parameters back to (mean, std) 

21 for physical interpretation. Uses formulas: mean = alpha * beta and std = sqrt(alpha) * beta. 

22 

23- :func:`bins` - Primary function for transport modeling. Creates discrete probability bins from 

24 continuous gamma distribution with equal-probability bins (default) or custom quantile edges. 

25 Returns bin edges, expected values (mean pore volume within each bin), and probability masses 

26 (weight in transport calculations). 

27 

28- :func:`bin_masses` - Calculate probability mass for custom bin edges using incomplete gamma 

29 function. Lower-level function used internally by bins(). 

30 

31This file is part of gwtransport which is released under AGPL-3.0 license. 

32See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details. 

33""" 

34 

35import numpy as np 

36import numpy.typing as npt 

37from scipy.special import gammainc 

38from scipy.stats import gamma as gamma_dist 

39 

40 

41def parse_parameters( 

42 *, 

43 alpha: float | None = None, 

44 beta: float | None = None, 

45 mean: float | None = None, 

46 std: float | None = None, 

47) -> tuple[float, float]: 

48 """ 

49 Parse parameters for gamma distribution. 

50 

51 Either alpha and beta or mean and std must be provided. 

52 

53 Parameters 

54 ---------- 

55 alpha : float, optional 

56 Shape parameter of gamma distribution (must be > 0) 

57 beta : float, optional 

58 Scale parameter of gamma distribution (must be > 0) 

59 mean : float, optional 

60 Mean of the gamma distribution. 

61 std : float, optional 

62 Standard deviation of the gamma distribution. See 

63 :ref:`concept-dispersion-scales` for what std represents depending 

64 on APVD source. 

65 

66 Returns 

67 ------- 

68 alpha : float 

69 Shape parameter of gamma distribution 

70 beta : float 

71 Scale parameter of gamma distribution 

72 

73 Raises 

74 ------ 

75 ValueError 

76 If both alpha and beta are None or if both mean and std are None. 

77 If alpha or beta are not positive. 

78 """ 

79 if (alpha is None) != (beta is None): 

80 msg = "alpha and beta must both be provided or both be None." 

81 raise ValueError(msg) 

82 

83 if alpha is None or beta is None: 

84 if mean is None or std is None: 

85 msg = "Either (alpha, beta) or (mean, std) must be provided." 

86 raise ValueError(msg) 

87 

88 alpha, beta = mean_std_to_alpha_beta(mean=mean, std=std) 

89 

90 if alpha <= 0 or beta <= 0: 

91 msg = "Alpha and beta must be positive" 

92 raise ValueError(msg) 

93 

94 return alpha, beta 

95 

96 

97def mean_std_to_alpha_beta(*, mean: float, std: float) -> tuple[float, float]: 

98 """ 

99 Convert mean and standard deviation of gamma distribution to shape and scale parameters. 

100 

101 Parameters 

102 ---------- 

103 mean : float 

104 Mean of the gamma distribution. 

105 std : float 

106 Standard deviation of the gamma distribution. See 

107 :ref:`concept-dispersion-scales` for what std represents depending 

108 on APVD source. 

109 

110 Returns 

111 ------- 

112 alpha : float 

113 Shape parameter of gamma distribution 

114 beta : float 

115 Scale parameter of gamma distribution 

116 

117 See Also 

118 -------- 

119 alpha_beta_to_mean_std : Convert shape and scale parameters to mean and std 

120 parse_parameters : Parse and validate gamma distribution parameters 

121 

122 Examples 

123 -------- 

124 >>> from gwtransport.gamma import mean_std_to_alpha_beta 

125 >>> mean_pore_volume = 30000.0 # m³ 

126 >>> std_pore_volume = 8100.0 # m³ 

127 >>> alpha, beta = mean_std_to_alpha_beta(mean=mean_pore_volume, std=std_pore_volume) 

128 >>> print(f"Shape parameter (alpha): {alpha:.2f}") 

129 Shape parameter (alpha): 13.72 

130 >>> print(f"Scale parameter (beta): {beta:.2f}") 

131 Scale parameter (beta): 2187.00 

132 """ 

133 if std <= 0: 

134 msg = "std must be positive" 

135 raise ValueError(msg) 

136 if mean <= 0: 

137 msg = "mean must be positive" 

138 raise ValueError(msg) 

139 

140 alpha = mean**2 / std**2 

141 beta = std**2 / mean 

142 return alpha, beta 

143 

144 

145def alpha_beta_to_mean_std(*, alpha: float, beta: float) -> tuple[float, float]: 

146 """ 

147 Convert shape and scale parameters of gamma distribution to mean and standard deviation. 

148 

149 Parameters 

150 ---------- 

151 alpha : float 

152 Shape parameter of the gamma distribution. 

153 beta : float 

154 Scale parameter of the gamma distribution. 

155 

156 Returns 

157 ------- 

158 mean : float 

159 Mean of the gamma distribution. 

160 std : float 

161 Standard deviation of the gamma distribution. 

162 

163 See Also 

164 -------- 

165 mean_std_to_alpha_beta : Convert mean and std to shape and scale parameters 

166 parse_parameters : Parse and validate gamma distribution parameters 

167 

168 Examples 

169 -------- 

170 >>> from gwtransport.gamma import alpha_beta_to_mean_std 

171 >>> alpha = 13.72 # shape parameter 

172 >>> beta = 2187.0 # scale parameter 

173 >>> mean, std = alpha_beta_to_mean_std(alpha=alpha, beta=beta) 

174 >>> print(f"Mean pore volume: {mean:.0f} m³") # doctest: +ELLIPSIS 

175 Mean pore volume: 3000... m³ 

176 >>> print(f"Std pore volume: {std:.0f} m³") # doctest: +ELLIPSIS 

177 Std pore volume: 810... m³ 

178 """ 

179 mean = alpha * beta 

180 std = np.sqrt(alpha) * beta 

181 return mean, std 

182 

183 

184def bins( 

185 *, 

186 alpha: float | None = None, 

187 beta: float | None = None, 

188 mean: float | None = None, 

189 std: float | None = None, 

190 n_bins: int = 100, 

191 quantile_edges: np.ndarray | None = None, 

192) -> dict[str, npt.NDArray[np.floating]]: 

193 """ 

194 Divide gamma distribution into bins and compute various bin properties. 

195 

196 If n_bins is provided, the gamma distribution is divided into n_bins equal-mass bins. 

197 If quantile_edges is provided, the gamma distribution is divided into bins defined by 

198 the quantile edges. The quantile edges must be in the range [0, 1] and of size n_bins + 1. 

199 The first and last quantile edges must be 0 and 1, respectively. 

200 

201 Parameters 

202 ---------- 

203 alpha : float, optional 

204 Shape parameter of gamma distribution (must be > 0) 

205 beta : float, optional 

206 Scale parameter of gamma distribution (must be > 0) 

207 mean : float, optional 

208 Mean of the gamma distribution. 

209 std : float, optional 

210 Standard deviation of the gamma distribution. 

211 n_bins : int, optional 

212 Number of bins to divide the gamma distribution (must be > 1). Default is 100. 

213 quantile_edges : array-like, optional 

214 Quantile edges for binning. Must be in the range [0, 1] and of size n_bins + 1. 

215 The first and last quantile edges must be 0 and 1, respectively. 

216 If provided, n_bins is ignored. 

217 

218 Returns 

219 ------- 

220 dict 

221 Dictionary with keys of type str and values of type numpy.ndarray: 

222 

223 - ``lower_bound``: lower bounds of bins (first one is 0) 

224 - ``upper_bound``: upper bounds of bins (last one is inf) 

225 - ``edges``: bin edges (lower_bound[0], upper_bound[0], ..., upper_bound[-1]) 

226 - ``expected_values``: expected values in bins. Is what you would expect to observe if you repeatedly sampled from the probability distribution, but only considered samples that fall within that particular bin 

227 - ``probability_mass``: probability mass in bins 

228 

229 See Also 

230 -------- 

231 bin_masses : Calculate probability mass for bins 

232 mean_std_to_alpha_beta : Convert mean/std to alpha/beta parameters 

233 gwtransport.advection.gamma_infiltration_to_extraction : Use bins for transport modeling 

234 :ref:`concept-gamma-distribution` : Two-parameter pore volume model 

235 :ref:`concept-dispersion-scales` : What ``std`` represents (macrodispersion vs total spreading) 

236 :ref:`assumption-gamma-distribution` : When gamma distribution is adequate 

237 

238 Examples 

239 -------- 

240 Create equal-mass bins for a gamma distribution: 

241 

242 >>> from gwtransport.gamma import bins 

243 >>> # Define gamma distribution using mean and std 

244 >>> result = bins(mean=30000.0, std=8100.0, n_bins=5) 

245 

246 Create bins with custom quantile edges: 

247 

248 >>> import numpy as np 

249 >>> quantiles = np.array([0.0, 0.25, 0.5, 0.75, 1.0]) 

250 >>> result = bins(mean=30000.0, std=8100.0, quantile_edges=quantiles) 

251 >>> print(f"Number of bins: {len(result['probability_mass'])}") 

252 Number of bins: 4 

253 """ 

254 alpha, beta = parse_parameters(alpha=alpha, beta=beta, mean=mean, std=std) 

255 

256 # Calculate boundaries for equal mass bins 

257 # If quantile_edges is provided, use it (n_bins is ignored) 

258 # Otherwise, use n_bins (which defaults to 100) 

259 if quantile_edges is not None: 

260 n_bins = len(quantile_edges) - 1 

261 else: 

262 quantile_edges = np.linspace(0, 1, n_bins + 1) # includes 0 and 1 

263 

264 if n_bins <= 1: 

265 msg = "Number of bins must be greater than 1" 

266 raise ValueError(msg) 

267 

268 bin_edges = gamma_dist.ppf(quantile_edges, alpha, scale=beta) 

269 probability_mass = np.diff(quantile_edges) # probability mass for each bin 

270 

271 # Calculate expected value for each bin 

272 diff_alpha_plus_1 = bin_masses(alpha=alpha + 1, beta=beta, bin_edges=bin_edges) 

273 expected_values = beta * alpha * diff_alpha_plus_1 / probability_mass 

274 

275 return { 

276 "lower_bound": bin_edges[:-1], 

277 "upper_bound": bin_edges[1:], 

278 "edges": bin_edges, 

279 "expected_values": expected_values, 

280 "probability_mass": probability_mass, 

281 } 

282 

283 

284def bin_masses(*, alpha: float, beta: float, bin_edges: npt.ArrayLike) -> npt.NDArray[np.floating]: 

285 """ 

286 Calculate probability mass for each bin in gamma distribution. 

287 

288 Is the area under the gamma distribution PDF between the bin edges. 

289 

290 Parameters 

291 ---------- 

292 alpha : float 

293 Shape parameter of gamma distribution (must be > 0) 

294 beta : float 

295 Scale parameter of gamma distribution (must be > 0) 

296 bin_edges : array-like 

297 Bin edges. Array of increasing values of size len(bins) + 1. 

298 Must be > 0. 

299 

300 Returns 

301 ------- 

302 numpy.ndarray 

303 Probability mass for each bin 

304 """ 

305 # Convert inputs to numpy arrays 

306 bin_edges = np.asarray(bin_edges) 

307 

308 # Validate inputs 

309 if alpha <= 0 or beta <= 0: 

310 msg = "Alpha and beta must be positive" 

311 raise ValueError(msg) 

312 if len(bin_edges) < 2: # noqa: PLR2004 

313 msg = "Bin edges must contain at least two values" 

314 raise ValueError(msg) 

315 if np.any(np.diff(bin_edges) < 0): 

316 msg = "Bin edges must be increasing" 

317 raise ValueError(msg) 

318 if np.any(bin_edges < 0): 

319 msg = "Bin edges must be positive" 

320 raise ValueError(msg) 

321 val = gammainc(alpha, bin_edges / beta) 

322 return val[1:] - val[:-1]