feat(distributions): add logk distribution and k_dist compound sampler

Introduce logk_gen (Y = ln X where X ~ K) with analytically derived
mean/variance via the CGF, numerically stable logpdf using an asymptotic
Bessel expansion for large arguments, CDF delegation to k_dist, and a
compound-gamma rvs sampler.

Add _rvs to k_dist via the same compound-gamma algorithm and extend
TestKDistPdf with stats and rvs coverage. Add a full TestLogK suite
covering pdf normalization, change-of-variable identity, CDF consistency,
analytical moment checks, and rvs moment checks.

Module-level docstring added to distributions.py
This commit is contained in:
2026-04-27 17:19:51 -03:00
parent e780bb956e
commit c59bc55fe5
2 changed files with 441 additions and 1 deletions

View File

@@ -7,7 +7,7 @@ import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from tools.distributions import k_dist, lognakagami, loggamma_dist, lograyleigh, logrice
from tools.distributions import k_dist, logk, lognakagami, loggamma_dist, lograyleigh, logrice
X = np.linspace(0.01, 10.0, 500)
@@ -73,6 +73,48 @@ class TestKDistPdf:
pdf_ba = k_dist.pdf(x_test, mu=1.0, alpha=3.0, beta=2.0)
np.testing.assert_allclose(pdf_ab, pdf_ba, rtol=1e-6)
def test_stats_mean_equals_mu(self):
"""Analytical mean must equal mu."""
for mu in [0.5, 1.0, 3.0]:
dist_mean = float(k_dist.stats(mu=mu, alpha=2.0, beta=3.0, moments="m"))
assert pytest.approx(dist_mean, rel=1e-10) == mu
def test_stats_variance_formula_eq4(self):
"""Variance must equal mu^2*(alpha+beta+1)/(alpha*beta) (equation 4)."""
mu, alpha, beta = 2.0, 3.0, 2.0
expected_var = mu**2 * (alpha + beta + 1) / (alpha * beta)
_, dist_var, *_ = k_dist.stats(mu=mu, alpha=alpha, beta=beta, moments="mv")
assert pytest.approx(float(dist_var), rel=1e-10) == expected_var
def test_stats_variance_symmetric_in_alpha_beta(self):
"""Variance is symmetric in alpha and beta."""
mu = 2.0
_, var_ab, *_ = k_dist.stats(mu=mu, alpha=2.0, beta=3.0, moments="mv")
_, var_ba, *_ = k_dist.stats(mu=mu, alpha=3.0, beta=2.0, moments="mv")
assert pytest.approx(float(var_ab), rel=1e-10) == float(var_ba)
def test_stats_variance_numerical(self):
"""Analytical variance should match sample variance from rvs."""
mu, alpha, beta = 2.0, 2.0, 3.0
rng = np.random.default_rng(42)
samples = k_dist.rvs(mu=mu, alpha=alpha, beta=beta, size=100_000, random_state=rng)
_, dist_var, *_ = k_dist.stats(mu=mu, alpha=alpha, beta=beta, moments="mv")
assert pytest.approx(float(np.var(samples)), rel=5e-2) == float(dist_var)
def test_rvs_samples_are_positive_and_finite(self):
"""K distribution samples must be positive and finite."""
rng = np.random.default_rng(7)
samples = k_dist.rvs(mu=1.0, alpha=2.0, beta=2.0, size=500, random_state=rng)
assert np.all(samples > 0)
assert np.all(np.isfinite(samples))
def test_rvs_sample_mean_near_mu(self):
"""Sample mean of k_dist rvs should be close to mu."""
mu, alpha, beta = 2.0, 3.0, 2.0
rng = np.random.default_rng(0)
samples = k_dist.rvs(mu=mu, alpha=alpha, beta=beta, size=100_000, random_state=rng)
assert pytest.approx(float(np.mean(samples)), rel=5e-2) == mu
# ── Parametric curve plots ───────────────────────────────────────────────────
@@ -514,6 +556,127 @@ class TestLogRice:
assert pytest.approx(samples.mean(), rel=5e-2) == numerical_mean
# ── logk unit tests ───────────────────────────────────────────────────────────
Y_LOGK = np.linspace(-10.0, 10.0, 500)
class TestLogK:
def test_logpdf_is_finite_on_real_line(self):
"""logpdf must be finite for all real y."""
log_vals = logk.logpdf(Y_LOGK, mu=2.0, alpha=3.0, beta=2.0)
assert np.all(np.isfinite(log_vals))
def test_pdf_integrates_to_one(self):
"""Numerical integral of PDF over the real line should be ≈ 1."""
y_fine = np.linspace(-30, 20, 200_000)
integral = np.trapezoid(logk.pdf(y_fine, mu=2.0, alpha=3.0, beta=2.0), y_fine)
assert pytest.approx(integral, abs=1e-3) == 1.0
def test_logpdf_equals_log_pdf(self):
"""logpdf must equal log(pdf) at points where pdf does not underflow."""
y_bulk = np.linspace(-5.0, 5.0, 30)
log_via_pdf = np.log(logk.pdf(y_bulk, mu=2.0, alpha=3.0, beta=2.0))
log_direct = logk.logpdf(y_bulk, mu=2.0, alpha=3.0, beta=2.0)
np.testing.assert_allclose(log_direct, log_via_pdf, rtol=1e-6)
def test_log_transform_relation_to_k_dist(self):
"""logk.pdf(y) must equal k_dist.pdf(exp(y)) * exp(y) (change-of-variable)."""
y_test = np.linspace(-3.0, 5.0, 20)
mu, alpha, beta = 2.0, 3.0, 2.0
direct = logk.pdf(y_test, mu=mu, alpha=alpha, beta=beta)
via_k = k_dist.pdf(np.exp(y_test), mu=mu, alpha=alpha, beta=beta) * np.exp(y_test)
np.testing.assert_allclose(direct, via_k, rtol=1e-6)
def test_cdf_consistent_with_k_dist(self):
"""logk.cdf(y) must equal k_dist.cdf(exp(y))."""
y_test = np.array([-2.0, 0.0, 1.0, 3.0])
mu, alpha, beta = 2.0, 2.0, 3.0
cdf_logk = logk.cdf(y_test, mu=mu, alpha=alpha, beta=beta)
cdf_k = k_dist.cdf(np.exp(y_test), mu=mu, alpha=alpha, beta=beta)
np.testing.assert_allclose(cdf_logk, cdf_k, rtol=1e-6)
def test_cdf_is_monotone_increasing(self):
"""CDF must be strictly non-decreasing (up to floating-point noise in the saturated tail)."""
y_grid = np.linspace(-5.0, 8.0, 30)
cdf_vals = logk.cdf(y_grid, mu=2.0, alpha=3.0, beta=2.0)
assert np.all(np.diff(cdf_vals) >= -1e-10)
def test_stats_mean_analytical(self):
"""Mean must equal ln(mu) - ln(alpha) - ln(beta) + psi(alpha) + psi(beta)."""
mu, alpha, beta = 2.0, 3.0, 2.0
expected = np.log(mu) - np.log(alpha) - np.log(beta) + sc.digamma(alpha) + sc.digamma(beta)
dist_mean = float(logk.stats(mu=mu, alpha=alpha, beta=beta, moments="m"))
assert pytest.approx(dist_mean, rel=1e-10) == expected
def test_stats_variance_analytical(self):
"""Variance must equal psi_1(alpha) + psi_1(beta)."""
mu, alpha, beta = 2.0, 3.0, 2.0
expected_var = sc.polygamma(1, alpha) + sc.polygamma(1, beta)
_, dist_var, *_ = logk.stats(mu=mu, alpha=alpha, beta=beta, moments="mv")
assert pytest.approx(float(dist_var), rel=1e-10) == expected_var
def test_stats_variance_mu_independent(self):
"""Variance must not depend on mu (mu is a pure shift in log-space)."""
alpha, beta = 3.0, 2.0
_, var1, *_ = logk.stats(mu=1.0, alpha=alpha, beta=beta, moments="mv")
_, var4, *_ = logk.stats(mu=4.0, alpha=alpha, beta=beta, moments="mv")
assert pytest.approx(float(var1), rel=1e-10) == float(var4)
def test_stats_mean_shifts_by_log_mu(self):
"""Doubling mu shifts the mean by ln(2) and leaves variance unchanged."""
alpha, beta = 3.0, 2.0
mean1 = float(logk.stats(mu=1.0, alpha=alpha, beta=beta, moments="m"))
mean2 = float(logk.stats(mu=2.0, alpha=alpha, beta=beta, moments="m"))
assert pytest.approx(mean2 - mean1, rel=1e-10) == np.log(2.0)
def test_argcheck_rejects_non_positive_mu(self):
"""mu <= 0 must not produce a valid PDF value."""
val = logk.pdf(0.0, mu=-1.0, alpha=2.0, beta=2.0)
assert not (np.isfinite(val) and val > 0)
def test_argcheck_rejects_non_positive_alpha(self):
"""alpha <= 0 must not produce a valid PDF value."""
val = logk.pdf(0.0, mu=1.0, alpha=-1.0, beta=2.0)
assert not (np.isfinite(val) and val > 0)
def test_argcheck_rejects_non_positive_beta(self):
"""beta <= 0 must not produce a valid PDF value."""
val = logk.pdf(0.0, mu=1.0, alpha=2.0, beta=-1.0)
assert not (np.isfinite(val) and val > 0)
def test_rvs_samples_are_finite(self):
"""Random samples must be finite real numbers."""
rng = np.random.default_rng(42)
samples = logk.rvs(mu=2.0, alpha=3.0, beta=2.0, size=200, random_state=rng)
assert samples.shape == (200,)
assert np.all(np.isfinite(samples))
def test_rvs_sample_mean_near_analytical(self):
"""Sample mean of many RVS should be close to the analytical mean."""
mu, alpha, beta = 2.0, 3.0, 2.0
rng = np.random.default_rng(0)
samples = logk.rvs(mu=mu, alpha=alpha, beta=beta, size=100_000, random_state=rng)
expected_mean = float(logk.stats(mu=mu, alpha=alpha, beta=beta, moments="m"))
assert pytest.approx(float(samples.mean()), rel=5e-2) == expected_mean
def test_rvs_sample_variance_near_analytical(self):
"""Sample variance of many RVS should be close to the analytical variance."""
mu, alpha, beta = 2.0, 3.0, 2.0
rng = np.random.default_rng(1)
samples = logk.rvs(mu=mu, alpha=alpha, beta=beta, size=100_000, random_state=rng)
_, expected_var, *_ = logk.stats(mu=mu, alpha=alpha, beta=beta, moments="mv")
assert pytest.approx(float(np.var(samples)), rel=5e-2) == float(expected_var)
def test_symmetry_in_alpha_beta(self):
"""logk PDF is symmetric in alpha and beta."""
y_test = np.linspace(-3.0, 5.0, 20)
mu = 2.0
pdf_ab = logk.pdf(y_test, mu=mu, alpha=2.0, beta=3.0)
pdf_ba = logk.pdf(y_test, mu=mu, alpha=3.0, beta=2.0)
np.testing.assert_allclose(pdf_ab, pdf_ba, rtol=1e-6)
if __name__ == "__main__":
plot_k_dist_varying_alpha()
plot_k_dist_varying_mu()