Skip to content

Commit 6f8bf93

Browse files
test(families): add test suite for gamma distribution
- Add tests covering all distribution characteristics - Test all three parametrizations and conversions between them - Test score (gradient) for all parametrizations against analytical formula - Test numerical gradient verification and edge cases - Compare characteristics against scipy.stats.gamma
1 parent 97776a2 commit 6f8bf93

1 file changed

Lines changed: 388 additions & 0 deletions

File tree

Lines changed: 388 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,388 @@
1+
"""
2+
Tests for Gamma Distribution Family
3+
4+
This module tests the functionality of the gamma distribution family,
5+
including parameterizations, characteristics, and support.
6+
"""
7+
8+
__author__ = "Elizaveta Bykova"
9+
__copyright__ = "Copyright (c) 2025 PySATL project"
10+
__license__ = "SPDX-License-Identifier: MIT"
11+
12+
13+
import numpy as np
14+
import pytest
15+
from scipy.special import gammaln
16+
from scipy.stats import gamma as gamma_dist
17+
18+
from pysatl_core.distributions.support import ContinuousSupport
19+
from pysatl_core.families.configuration import configure_families_register
20+
from pysatl_core.types import (
21+
CharacteristicName,
22+
ContinuousSupportShape1D,
23+
FamilyName,
24+
UnivariateContinuous,
25+
)
26+
27+
from .base import BaseDistributionTest
28+
29+
30+
class TestGammaFamily(BaseDistributionTest):
31+
"""Test suite for Gamma distribution family."""
32+
33+
def setup_method(self):
34+
"""Setup before each test method."""
35+
registry = configure_families_register()
36+
self.gamma_family = registry.get(FamilyName.GAMMA)
37+
self.gamma_dist_example = self.gamma_family(k=2.0, theta=3.0)
38+
39+
def test_family_properties(self):
40+
"""Test basic properties of gamma family."""
41+
assert self.gamma_family.name == FamilyName.GAMMA
42+
43+
# Check parameterizations
44+
expected_parametrizations = {"shapeScale", "shapeRate", "meanVar"}
45+
assert set(self.gamma_family.parametrization_names) == expected_parametrizations
46+
assert self.gamma_family.base_parametrization_name == "shapeScale"
47+
48+
def test_shape_scale_parametrization_creation(self):
49+
"""Test creation of distribution with shape-scale parametrization."""
50+
dist = self.gamma_family(k=2.0, theta=3.0)
51+
52+
assert dist.family_name == FamilyName.GAMMA
53+
assert dist.distribution_type == UnivariateContinuous
54+
assert dist.parameters == {"k": 2.0, "theta": 3.0}
55+
assert dist.parametrization_name == "shapeScale"
56+
57+
def test_shape_rate_parametrization_creation(self):
58+
"""Test creation of distribution with shape-rate parametrization."""
59+
dist = self.gamma_family(k=2.0, beta=0.5, parametrization_name="shapeRate")
60+
61+
assert dist.parameters == {"k": 2.0, "beta": 0.5}
62+
assert dist.parametrization_name == "shapeRate"
63+
64+
def test_mean_var_parametrization_creation(self):
65+
"""Test creation of distribution with mean-variance parametrization."""
66+
dist = self.gamma_family(m=6.0, v=18.0, parametrization_name="meanVar")
67+
68+
assert dist.parameters == {"m": 6.0, "v": 18.0}
69+
assert dist.parametrization_name == "meanVar"
70+
71+
def test_parametrization_constraints(self):
72+
"""Test parameter constraints validation."""
73+
with pytest.raises(ValueError, match="k > 0"):
74+
self.gamma_family(k=-1.0, theta=1.0)
75+
76+
with pytest.raises(ValueError, match="theta > 0"):
77+
self.gamma_family(k=1.0, theta=-1.0)
78+
79+
with pytest.raises(ValueError, match="k > 0"):
80+
self.gamma_family(k=-1.0, beta=1.0, parametrization_name="shapeRate")
81+
82+
with pytest.raises(ValueError, match="beta > 0"):
83+
self.gamma_family(k=1.0, beta=-1.0, parametrization_name="shapeRate")
84+
85+
with pytest.raises(ValueError, match="m > 0"):
86+
self.gamma_family(m=-1.0, v=1.0, parametrization_name="meanVar")
87+
88+
with pytest.raises(ValueError, match="v > 0"):
89+
self.gamma_family(m=1.0, v=-1.0, parametrization_name="meanVar")
90+
91+
def test_moments(self):
92+
"""Test moment calculations for k=2, theta=3."""
93+
mean_func = self.gamma_dist_example.query_method(CharacteristicName.MEAN)
94+
assert abs(mean_func() - 6.0) < self.CALCULATION_PRECISION
95+
96+
var_func = self.gamma_dist_example.query_method(CharacteristicName.VAR)
97+
assert abs(var_func() - 18.0) < self.CALCULATION_PRECISION
98+
99+
@pytest.mark.parametrize(
100+
"parametrization_name, params, expected_k, expected_theta",
101+
[
102+
("shapeScale", {"k": 2.0, "theta": 3.0}, 2.0, 3.0),
103+
("shapeRate", {"k": 2.0, "beta": 1 / 3}, 2.0, 3.0),
104+
("meanVar", {"m": 6.0, "v": 18.0}, 2.0, 3.0),
105+
],
106+
)
107+
def test_parametrization_conversions(
108+
self, parametrization_name, params, expected_k, expected_theta
109+
):
110+
"""Test conversions between different parameterizations."""
111+
base_params = self.gamma_family.to_base(
112+
self.gamma_family.get_parametrization(parametrization_name)(**params)
113+
)
114+
115+
assert abs(base_params.parameters["k"] - expected_k) < self.CALCULATION_PRECISION
116+
assert abs(base_params.parameters["theta"] - expected_theta) < self.CALCULATION_PRECISION
117+
118+
def test_analytical_computations_availability(self):
119+
"""Test that analytical computations are available for gamma distribution."""
120+
comp = self.gamma_family(k=1.0, theta=1.0).analytical_computations
121+
122+
expected_chars = {
123+
CharacteristicName.PDF,
124+
CharacteristicName.CDF,
125+
CharacteristicName.PPF,
126+
CharacteristicName.CF,
127+
CharacteristicName.LPDF,
128+
CharacteristicName.MEAN,
129+
CharacteristicName.VAR,
130+
}
131+
assert set(comp.keys()) == expected_chars
132+
133+
@pytest.mark.parametrize(
134+
"char_name, test_data, scipy_func, scipy_kwargs",
135+
[
136+
(
137+
CharacteristicName.PDF,
138+
[0.5, 1.0, 2.0, 3.0, 5.0],
139+
gamma_dist.pdf,
140+
{"a": 2.0, "scale": 3.0},
141+
),
142+
(
143+
CharacteristicName.CDF,
144+
[0.5, 1.0, 2.0, 3.0, 5.0],
145+
gamma_dist.cdf,
146+
{"a": 2.0, "scale": 3.0},
147+
),
148+
(
149+
CharacteristicName.PPF,
150+
[0.1, 0.25, 0.5, 0.75, 0.9],
151+
gamma_dist.ppf,
152+
{"a": 2.0, "scale": 3.0},
153+
),
154+
(
155+
CharacteristicName.LPDF,
156+
[0.5, 1.0, 2.0, 3.0, 5.0],
157+
gamma_dist.logpdf,
158+
{"a": 2.0, "scale": 3.0},
159+
),
160+
],
161+
)
162+
def test_array_input_for_characteristics(self, char_name, test_data, scipy_func, scipy_kwargs):
163+
"""Test that characteristics support array inputs."""
164+
dist = self.gamma_dist_example
165+
char_func = dist.query_method(char_name)
166+
167+
input_array = np.array(test_data)
168+
result_array = char_func(input_array)
169+
170+
assert result_array.shape == input_array.shape
171+
172+
expected_array = scipy_func(input_array, **scipy_kwargs)
173+
174+
self.assert_arrays_almost_equal(result_array, expected_array)
175+
176+
def test_characteristic_function_array_input(self):
177+
"""Test characteristic function calculation with array input."""
178+
char_func = self.gamma_dist_example.query_method(CharacteristicName.CF)
179+
t_array = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
180+
181+
cf_array = char_func(t_array)
182+
assert cf_array.shape == t_array.shape
183+
184+
k, theta = 2.0, 3.0
185+
expected = (1 - 1j * theta * t_array) ** (-k)
186+
187+
self.assert_arrays_almost_equal(cf_array.real, expected.real)
188+
self.assert_arrays_almost_equal(cf_array.imag, expected.imag)
189+
190+
def test_gamma_support(self):
191+
"""Test that gamma distribution has correct support (0, inf)."""
192+
dist = self.gamma_dist_example
193+
194+
assert dist.support is not None
195+
assert isinstance(dist.support, ContinuousSupport)
196+
197+
assert dist.support.left == 0.0
198+
assert dist.support.right == float("inf")
199+
assert dist.support.left_closed
200+
assert not dist.support.right_closed
201+
202+
assert dist.support.contains(1.0) is True
203+
assert dist.support.contains(-0.1) is False
204+
assert dist.support.contains(float("inf")) is False
205+
206+
test_points = np.array([-0.1, 0.5, 1.0, 10.0])
207+
results = dist.support.contains(test_points)
208+
np.testing.assert_array_equal(results, [False, True, True, True])
209+
210+
assert dist.support.shape == ContinuousSupportShape1D.RAY_RIGHT
211+
212+
def test_pdf_at_zero(self):
213+
"""Test that PDF at x=0 returns 0 for k > 1."""
214+
dist = self.gamma_family(k=2.0, theta=1.0)
215+
pdf_func = dist.query_method(CharacteristicName.PDF)
216+
assert pdf_func(np.array([0.0]))[0] == 0.0
217+
218+
def test_cdf_at_zero(self):
219+
"""Test that CDF at x=0 returns 0."""
220+
dist = self.gamma_family(k=2.0, theta=1.0)
221+
cdf_func = dist.query_method(CharacteristicName.CDF)
222+
assert cdf_func(np.array([0.0]))[0] == 0.0
223+
224+
def test_lpdf_at_zero(self):
225+
"""Test that LPDF at x=0 returns -inf."""
226+
dist = self.gamma_family(k=2.0, theta=1.0)
227+
lpdf_func = dist.query_method(CharacteristicName.LPDF)
228+
assert lpdf_func(np.array([0.0]))[0] == float("-inf")
229+
230+
231+
class TestGammaFamilyEdgeCases(BaseDistributionTest):
232+
"""Test edge cases and error conditions for gamma distribution."""
233+
234+
def setup_method(self):
235+
"""Setup before each test method."""
236+
registry = configure_families_register()
237+
self.gamma_family = registry.get(FamilyName.GAMMA)
238+
239+
def test_invalid_parameterization(self):
240+
"""Test error for invalid parameterization name."""
241+
with pytest.raises(KeyError):
242+
self.gamma_family.distribution(parametrization_name="invalid_name", k=1.0, theta=1.0)
243+
244+
def test_missing_parameters(self):
245+
"""Test error for missing required parameters."""
246+
with pytest.raises(TypeError):
247+
self.gamma_family.distribution(k=1.0) # Missing theta
248+
249+
def test_invalid_probability_ppf(self):
250+
"""Test PPF with invalid probability values."""
251+
dist = self.gamma_family(k=1.0, theta=1.0)
252+
ppf = dist.query_method(CharacteristicName.PPF)
253+
254+
assert ppf(0.0) == 0.0
255+
assert ppf(1.0) == float("inf")
256+
257+
with pytest.raises(ValueError):
258+
ppf(-0.1)
259+
with pytest.raises(ValueError):
260+
ppf(1.1)
261+
262+
def test_characteristic_function_at_zero(self):
263+
"""Test characteristic function at zero returns 1."""
264+
dist = self.gamma_family(k=2.0, theta=3.0)
265+
char_func = dist.query_method(CharacteristicName.CF)
266+
267+
cf_value_zero = char_func(0.0)
268+
assert abs(cf_value_zero.real - 1.0) < self.CALCULATION_PRECISION
269+
assert abs(cf_value_zero.imag) < self.CALCULATION_PRECISION
270+
271+
def test_characteristic_function_large_t(self):
272+
"""Test characteristic function with large t."""
273+
dist = self.gamma_family(k=2.0, theta=3.0)
274+
char_func = dist.query_method(CharacteristicName.CF)
275+
276+
cf_value_large = char_func(1000.0)
277+
assert np.iscomplexobj(cf_value_large)
278+
assert abs(cf_value_large) <= 1.0
279+
280+
def test_score_shape_scale(self):
281+
"""Test SCORE for shapeScale (base) parametrization against analytical formula."""
282+
k, theta = 3.0, 2.0
283+
dist = self.gamma_family(k=k, theta=theta)
284+
x = np.array([1.0, 2.0, 5.0, 10.0])
285+
grad = dist.family.score(dist.parametrization, x)
286+
287+
from scipy.special import digamma as _digamma
288+
289+
expected_k = np.log(x) - np.log(theta) - _digamma(k)
290+
expected_theta = (x - k * theta) / (theta**2)
291+
expected = np.stack([expected_k, expected_theta], axis=-1)
292+
293+
np.testing.assert_allclose(grad, expected, rtol=self.CALCULATION_PRECISION)
294+
295+
def test_score_shape_rate(self):
296+
"""Test SCORE for shapeRate parametrization via chain rule."""
297+
k, beta = 3.0, 0.5
298+
dist = self.gamma_family(parametrization_name="shapeRate", k=k, beta=beta)
299+
x = np.array([1.0, 2.0, 5.0, 10.0])
300+
grad = dist.family.score(dist.parametrization, x)
301+
302+
from scipy.special import digamma as _digamma
303+
304+
theta = 1.0 / beta
305+
base_k = np.log(x) - np.log(theta) - _digamma(k)
306+
base_theta = (x - k * theta) / (theta**2)
307+
308+
expected_k = base_k
309+
expected_beta = base_theta * (-1.0 / beta**2)
310+
expected = np.stack([expected_k, expected_beta], axis=-1)
311+
312+
np.testing.assert_allclose(grad, expected, rtol=self.CALCULATION_PRECISION)
313+
314+
def test_score_mean_var(self):
315+
"""Test SCORE for meanVar parametrization via chain rule."""
316+
m, v = 6.0, 2.0
317+
dist = self.gamma_family(parametrization_name="meanVar", m=m, v=v)
318+
x = np.array([1.0, 2.0, 5.0, 10.0])
319+
grad = dist.family.score(dist.parametrization, x)
320+
321+
from scipy.special import digamma as _digamma
322+
323+
k = m**2 / v
324+
theta = v / m
325+
326+
base_k = np.log(x) - np.log(theta) - _digamma(k)
327+
base_theta = (x - k * theta) / (theta**2)
328+
329+
dk_dm = 2 * m / v
330+
dk_dv = -(m**2) / (v**2)
331+
dtheta_dm = -v / (m**2)
332+
dtheta_dv = 1.0 / m
333+
334+
expected_m = base_k * dk_dm + base_theta * dtheta_dm
335+
expected_v = base_k * dk_dv + base_theta * dtheta_dv
336+
expected = np.stack([expected_m, expected_v], axis=-1)
337+
338+
np.testing.assert_allclose(grad, expected, rtol=self.CALCULATION_PRECISION)
339+
340+
def test_score_numerical_derivative(self):
341+
"""Compare analytical SCORE with numerical gradient for base parametrization."""
342+
k, theta = 3.0, 2.0
343+
dist = self.gamma_family(k=k, theta=theta)
344+
x_val = 4.0
345+
346+
def logpdf_k(k_val: float) -> float:
347+
return (
348+
(k_val - 1) * np.log(x_val) - x_val / theta - k_val * np.log(theta) - gammaln(k_val)
349+
)
350+
351+
def logpdf_theta(theta_val: float) -> float:
352+
return (k - 1) * np.log(x_val) - x_val / theta_val - k * np.log(theta_val) - gammaln(k)
353+
354+
eps = 1e-7
355+
grad_k_num = (logpdf_k(k + eps) - logpdf_k(k - eps)) / (2 * eps)
356+
grad_theta_num = (logpdf_theta(theta + eps) - logpdf_theta(theta - eps)) / (2 * eps)
357+
358+
grad = dist.family.score(dist.parametrization, np.array([x_val]))
359+
np.testing.assert_allclose(grad[0, 0], grad_k_num, rtol=1e-5)
360+
np.testing.assert_allclose(grad[0, 1], grad_theta_num, rtol=1e-5)
361+
362+
def test_score_x_outside_support_raises(self):
363+
"""Test SCORE raises ValueError for x <= 0."""
364+
dist = self.gamma_family(k=2.0, theta=1.0)
365+
with pytest.raises(ValueError, match="base_score is undefined"):
366+
dist.family.score(dist.parametrization, np.array([-1.0, 1.0]))
367+
with pytest.raises(ValueError, match="base_score is undefined"):
368+
dist.family.score(dist.parametrization, np.array([1.0, np.inf]))
369+
with pytest.raises(ValueError, match="base_score is undefined"):
370+
dist.family.score(dist.parametrization, np.array([np.nan, 1.0]))
371+
with pytest.raises(ValueError, match="base_score is undefined"):
372+
dist.family.score(dist.parametrization, np.array([0.0, 1.0]))
373+
374+
@pytest.mark.parametrize(
375+
"parametrization_name, params",
376+
[
377+
("shapeScale", {"k": 2.0, "theta": 1.5}),
378+
("shapeRate", {"k": 2.0, "beta": 2.0 / 3.0}),
379+
("meanVar", {"m": 3.0, "v": 1.5}),
380+
],
381+
)
382+
def test_score_shape(self, parametrization_name, params):
383+
"""Test SCORE output shape for all parametrizations."""
384+
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
385+
dist = self.gamma_family(parametrization_name=parametrization_name, **params)
386+
grad = dist.family.score(dist.parametrization, x)
387+
assert grad.shape == (len(x), 2)
388+
assert grad.dtype == float

0 commit comments

Comments
 (0)