Coverage for src/gwtransport/gamma.py: 79%
53 statements
« prev ^ index » next coverage.py v7.8.2, created at 2025-06-06 15:45 +0000
1"""Functions for working with gamma distributions."""
3import numpy as np
4from scipy.special import gammainc
5from scipy.stats import gamma as gamma_dist
8def parse_parameters(alpha=None, beta=None, mean=None, std=None):
9 """
10 Parse parameters for gamma distribution.
12 Either alpha and beta or mean and std must be provided.
14 Parameters
15 ----------
16 alpha : float, optional
17 Shape parameter of gamma distribution (must be > 0)
18 beta : float, optional
19 Scale parameter of gamma distribution (must be > 0)
20 mean : float, optional
21 Mean of the gamma distribution.
22 std : float, optional
23 Standard deviation of the gamma distribution.
25 Returns
26 -------
27 tuple
28 Shape and scale parameters of the gamma distribution.
30 Raises
31 ------
32 ValueError
33 If both alpha and beta are None or if both mean and std are None.
34 If alpha or beta are not positive.
35 """
36 if alpha is None or beta is None:
37 if mean is None or std is None:
38 msg = "Either alpha and beta or mean and std must be provided."
39 raise ValueError(msg)
41 alpha, beta = mean_std_to_alpha_beta(mean, std)
43 if alpha <= 0 or beta <= 0:
44 msg = "Alpha and beta must be positive"
45 raise ValueError(msg)
47 return alpha, beta
50def mean_std_to_alpha_beta(mean, std):
51 """
52 Convert mean and standard deviation of gamma distribution to shape and scale parameters.
54 Parameters
55 ----------
56 mean : float
57 Mean of the gamma distribution.
58 std : float
59 Standard deviation of the gamma distribution.
61 Returns
62 -------
63 tuple
64 Shape and scale parameters of the gamma distribution.
65 """
66 alpha = mean**2 / std**2
67 beta = std**2 / mean
68 return alpha, beta
71def alpha_beta_to_mean_std(alpha, beta):
72 """
73 Convert shape and scale parameters of gamma distribution to mean and standard deviation.
75 Parameters
76 ----------
77 alpha : float
78 Shape parameter of the gamma distribution.
79 beta : float
80 Scale parameter of the gamma distribution.
82 Returns
83 -------
84 tuple
85 Mean and standard deviation of the gamma distribution.
86 """
87 mean = alpha * beta
88 std = np.sqrt(alpha) * beta
89 return mean, std
92def bins(*, alpha=None, beta=None, mean=None, std=None, n_bins=None, quantile_edges=None):
93 """
94 Divide gamma distribution into bins and compute various bin properties.
96 If n_bins is provided, the gamma distribution is divided into n_bins equal-mass bins.
97 If quantile_edges is provided, the gamma distribution is divided into bins defined by
98 the quantile edges. The quantile edges must be in the range [0, 1] and of size n_bins + 1.
99 The first and last quantile edges must be 0 and 1, respectively.
101 Parameters
102 ----------
103 alpha : float, optional
104 Shape parameter of gamma distribution (must be > 0)
105 beta : float, optional
106 Scale parameter of gamma distribution (must be > 0)
107 mean : float, optional
108 Mean of the gamma distribution.
109 std : float, optional
110 Standard deviation of the gamma distribution.
111 n_bins : int, optional
112 Number of bins to divide the gamma distribution (must be > 1)
113 quantile_edges : array-like, optional
114 Quantile edges for binning. Must be in the range [0, 1] and of size n_bins + 1.
115 The first and last quantile edges must be 0 and 1, respectively.
116 If provided, n_bins is ignored.
118 Returns
119 -------
120 dict of arrays with keys:
121 - lower_bound: lower bounds of bins (first one is 0)
122 - upper_bound: upper bounds of bins (last one is inf)
123 - edges: bin edges (lower_bound[0], upper_bound[0], ..., upper_bound[-1])
124 - expected_value: expected values in bins
125 - probability_mass: probability mass in bins
126 """
127 alpha, beta = parse_parameters(alpha=alpha, beta=beta, mean=mean, std=std)
129 # Calculate boundaries for equal mass bins
130 if not ((n_bins is None) ^ (quantile_edges is None)):
131 msg = "Either n_bins or quantiles must be provided"
132 raise ValueError(msg)
134 if quantile_edges is not None:
135 n_bins = len(quantile_edges) - 1
136 else:
137 quantile_edges = np.linspace(0, 1, n_bins + 1) # includes 0 and 1
139 if n_bins <= 1:
140 msg = "Number of bins must be greater than 1"
141 raise ValueError(msg)
143 bin_edges = gamma_dist.ppf(quantile_edges, alpha, scale=beta)
144 probability_mass = np.diff(quantile_edges) # probability mass for each bin
146 # Calculate expected value for each bin
147 diff_alpha_plus_1 = bin_masses(alpha + 1, beta, bin_edges)
148 expected_values = beta * alpha * diff_alpha_plus_1 / probability_mass
150 return {
151 "lower_bound": bin_edges[:-1],
152 "upper_bound": bin_edges[1:],
153 "edges": bin_edges,
154 "expected_value": expected_values,
155 "probability_mass": probability_mass,
156 }
159def bin_masses(alpha, beta, bin_edges):
160 """
161 Calculate probability mass for each bin in gamma distribution.
163 Is the area under the gamma distribution PDF between the bin edges.
165 Parameters
166 ----------
167 alpha : float
168 Shape parameter of gamma distribution (must be > 0)
169 beta : float
170 Scale parameter of gamma distribution (must be > 0)
171 bin_edges : array-like
172 Bin edges. Array of increasing values of size len(bins) + 1.
173 Must be > 0.
175 Returns
176 -------
177 array
178 Probability mass for each bin
179 """
180 # Validate inputs
181 if alpha <= 0 or beta <= 0:
182 msg = "Alpha and beta must be positive"
183 raise ValueError(msg)
184 if len(bin_edges) < 2: # noqa: PLR2004
185 msg = "Bin edges must contain at least two values"
186 raise ValueError(msg)
187 if np.any(np.diff(bin_edges) < 0):
188 msg = "Bin edges must be increasing"
189 raise ValueError(msg)
190 if np.any(bin_edges < 0):
191 msg = "Bin edges must be positive"
192 raise ValueError(msg)
194 # Convert inputs to numpy arrays
195 bin_edges = np.asarray(bin_edges)
196 val = gammainc(alpha, bin_edges / beta)
197 return val[1:] - val[:-1]