Coverage for src/gwtransport/gamma.py: 75%
60 statements
« prev ^ index » next coverage.py v7.10.7, created at 2025-09-29 20:01 +0000
1"""Functions for working with gamma distributions."""
3import numpy as np
4import numpy.typing as npt
5from scipy.special import gammainc
6from scipy.stats import gamma as gamma_dist
9def parse_parameters(
10 alpha: float | None = None,
11 beta: float | None = None,
12 mean: float | None = None,
13 std: float | None = None,
14) -> tuple[float, float]:
15 """
16 Parse parameters for gamma distribution.
18 Either alpha and beta or mean and std must be provided.
20 Parameters
21 ----------
22 alpha : float, optional
23 Shape parameter of gamma distribution (must be > 0)
24 beta : float, optional
25 Scale parameter of gamma distribution (must be > 0)
26 mean : float, optional
27 Mean of the gamma distribution.
28 std : float, optional
29 Standard deviation of the gamma distribution.
31 Returns
32 -------
33 tuple
34 Shape and scale parameters of the gamma distribution.
36 Raises
37 ------
38 ValueError
39 If both alpha and beta are None or if both mean and std are None.
40 If alpha or beta are not positive.
41 """
42 if alpha is None or beta is None:
43 if mean is None or std is None:
44 msg = "Either alpha and beta or mean and std must be provided."
45 raise ValueError(msg)
47 alpha, beta = mean_std_to_alpha_beta(mean, std)
49 if alpha <= 0 or beta <= 0:
50 msg = "Alpha and beta must be positive"
51 raise ValueError(msg)
53 return alpha, beta
56def mean_std_to_alpha_beta(mean: float, std: float) -> tuple[float, float]:
57 """
58 Convert mean and standard deviation of gamma distribution to shape and scale parameters.
60 Parameters
61 ----------
62 mean : float
63 Mean of the gamma distribution.
64 std : float
65 Standard deviation of the gamma distribution.
67 Returns
68 -------
69 tuple
70 Shape and scale parameters of the gamma distribution.
71 """
72 alpha = mean**2 / std**2
73 beta = std**2 / mean
74 return alpha, beta
77def alpha_beta_to_mean_std(alpha: float, beta: float) -> tuple[float, float]:
78 """
79 Convert shape and scale parameters of gamma distribution to mean and standard deviation.
81 Parameters
82 ----------
83 alpha : float
84 Shape parameter of the gamma distribution.
85 beta : float
86 Scale parameter of the gamma distribution.
88 Returns
89 -------
90 tuple
91 Mean and standard deviation of the gamma distribution.
92 """
93 mean = alpha * beta
94 std = np.sqrt(alpha) * beta
95 return mean, std
98def bins(
99 *,
100 alpha: float | None = None,
101 beta: float | None = None,
102 mean: float | None = None,
103 std: float | None = None,
104 n_bins: int | None = None,
105 quantile_edges: np.ndarray | None = None,
106):
107 """
108 Divide gamma distribution into bins and compute various bin properties.
110 If n_bins is provided, the gamma distribution is divided into n_bins equal-mass bins.
111 If quantile_edges is provided, the gamma distribution is divided into bins defined by
112 the quantile edges. The quantile edges must be in the range [0, 1] and of size n_bins + 1.
113 The first and last quantile edges must be 0 and 1, respectively.
115 Parameters
116 ----------
117 alpha : float, optional
118 Shape parameter of gamma distribution (must be > 0)
119 beta : float, optional
120 Scale parameter of gamma distribution (must be > 0)
121 mean : float, optional
122 Mean of the gamma distribution.
123 std : float, optional
124 Standard deviation of the gamma distribution.
125 n_bins : int, optional
126 Number of bins to divide the gamma distribution (must be > 1)
127 quantile_edges : array-like, optional
128 Quantile edges for binning. Must be in the range [0, 1] and of size n_bins + 1.
129 The first and last quantile edges must be 0 and 1, respectively.
130 If provided, n_bins is ignored.
132 Returns
133 -------
134 dict of arrays with keys:
135 - lower_bound: lower bounds of bins (first one is 0)
136 - upper_bound: upper bounds of bins (last one is inf)
137 - edges: bin edges (lower_bound[0], upper_bound[0], ..., upper_bound[-1])
138 - 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
139 - probability_mass: probability mass in bins
140 """
141 alpha, beta = parse_parameters(alpha=alpha, beta=beta, mean=mean, std=std)
143 # Calculate boundaries for equal mass bins
144 if not ((n_bins is None) ^ (quantile_edges is None)):
145 msg = "Either n_bins or quantiles must be provided"
146 raise ValueError(msg)
148 if quantile_edges is not None:
149 n_bins = len(quantile_edges) - 1
150 else:
151 if n_bins is None:
152 msg = "n_bins should not be None here due to earlier validation"
153 raise ValueError(msg)
154 quantile_edges = np.linspace(0, 1, n_bins + 1) # includes 0 and 1
156 if n_bins is None:
157 msg = "n_bins should not be None here"
158 raise ValueError(msg)
159 if n_bins <= 1:
160 msg = "Number of bins must be greater than 1"
161 raise ValueError(msg)
163 bin_edges = gamma_dist.ppf(quantile_edges, alpha, scale=beta)
164 probability_mass = np.diff(quantile_edges) # probability mass for each bin
166 # Calculate expected value for each bin
167 diff_alpha_plus_1 = bin_masses(alpha + 1, beta, bin_edges)
168 expected_values = beta * alpha * diff_alpha_plus_1 / probability_mass
170 return {
171 "lower_bound": bin_edges[:-1],
172 "upper_bound": bin_edges[1:],
173 "edges": bin_edges,
174 "expected_values": expected_values,
175 "probability_mass": probability_mass,
176 }
179def bin_masses(alpha: float, beta: float, bin_edges: npt.ArrayLike) -> np.ndarray:
180 """
181 Calculate probability mass for each bin in gamma distribution.
183 Is the area under the gamma distribution PDF between the bin edges.
185 Parameters
186 ----------
187 alpha : float
188 Shape parameter of gamma distribution (must be > 0)
189 beta : float
190 Scale parameter of gamma distribution (must be > 0)
191 bin_edges : array-like
192 Bin edges. Array of increasing values of size len(bins) + 1.
193 Must be > 0.
195 Returns
196 -------
197 array
198 Probability mass for each bin
199 """
200 # Validate inputs
201 if alpha <= 0 or beta <= 0:
202 msg = "Alpha and beta must be positive"
203 raise ValueError(msg)
204 if len(bin_edges) < 2: # noqa: PLR2004
205 msg = "Bin edges must contain at least two values"
206 raise ValueError(msg)
207 if np.any(np.diff(bin_edges) < 0):
208 msg = "Bin edges must be increasing"
209 raise ValueError(msg)
210 if np.any(bin_edges < 0):
211 msg = "Bin edges must be positive"
212 raise ValueError(msg)
214 # Convert inputs to numpy arrays
215 bin_edges = np.asarray(bin_edges)
216 val = gammainc(alpha, bin_edges / beta)
217 return val[1:] - val[:-1]