From 94b53ee87770ec34117be5aa8b4785336a34a2a8 Mon Sep 17 00:00:00 2001 From: kaylimekay Date: Sun, 2 Aug 2026 16:37:44 -0400 Subject: [PATCH 1/3] Add Fisher noncentral hypergeometric distribution --- docs/api/distributions.rst | 1 + pymc_extras/distributions/__init__.py | 2 + pymc_extras/distributions/discrete.py | 196 ++++++++++++++++++++++++++ tests/distributions/test_discrete.py | 82 ++++++++++- 4 files changed, 280 insertions(+), 1 deletion(-) diff --git a/docs/api/distributions.rst b/docs/api/distributions.rst index 5de605ca2..c5a677ac2 100644 --- a/docs/api/distributions.rst +++ b/docs/api/distributions.rst @@ -19,6 +19,7 @@ like regular PyMC distributions and can be used directly inside a model. ExtGenPareto R2D2M2CP Skellam + FisherNoncentralHypergeometric histogram_approximation Transforms diff --git a/pymc_extras/distributions/__init__.py b/pymc_extras/distributions/__init__.py index c89ca9a6b..f2477ebaf 100644 --- a/pymc_extras/distributions/__init__.py +++ b/pymc_extras/distributions/__init__.py @@ -26,6 +26,7 @@ ) from pymc_extras.distributions.discrete import ( BetaNegativeBinomial, + FisherNoncentralHypergeometric, GeneralizedPoisson, Skellam, ) @@ -40,6 +41,7 @@ "Chi", "DiscreteMarkovChain", "ExtGenPareto", + "FisherNoncentralHypergeometric", "GenExtreme", "GenPareto", "GeneralizedPoisson", diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py index fc810c415..3b558e2ca 100644 --- a/pymc_extras/distributions/discrete.py +++ b/pymc_extras/distributions/discrete.py @@ -18,7 +18,11 @@ from pymc.distributions.dist_math import betaln, check_parameters, factln, logpow from pymc.distributions.shape_utils import rv_size_is_none from pytensor import tensor as pt +from pytensor.tensor.random.basic import ScipyRandomVariable from pytensor.tensor.random.op import RandomVariable +from pytensor.utils import lazy_scipy_module + +stats = lazy_scipy_module("stats") def log1mexp(x): @@ -184,6 +188,198 @@ def logp(value, mu, lam): ) +class FisherNoncentralHypergeometricRV(ScipyRandomVariable): + name = "fisher_noncentral_hypergeometric" + signature = "(),(),(),()->()" + dtype = "int64" + _print_name = ( + "FisherNoncentralHypergeometric", + "\\operatorname{FisherNoncentralHypergeometric}", + ) + + @classmethod + def rng_fn_scipy(cls, rng, good, bad, n, odds, size=None): + return stats.nchypergeom_fisher.rvs( + good + bad, + good, + n, + odds, + size=size, + random_state=rng, + ) + + +fisher_noncentral_hypergeometric = FisherNoncentralHypergeometricRV() + + +class FisherNoncentralHypergeometric(pm.distributions.Discrete): + R""" + Fisher noncentral hypergeometric distribution. + + The Fisher noncentral hypergeometric distribution is a generalization + of the standard hypergeometric distribution and models the number of + successful draws from a finite population when the sampling odds are + biased by a noncentrality parameter. The probability mass function is + + .. math:: + + f(x \mid N, k, n, \theta) = + \frac{\binom{k}{x}\binom{N-k}{n-x}\theta^{x}} + {\sum_{j=\max(0, n-(N-k))}^{\min(n, k)} + \binom{k}{j}\binom{N-k}{n-j}\theta^{j}} + + for + :math:`x \in \{\max(0, n-(N-k)), \ldots, \min(n, k)\}`. + For more information, see [1]_. + + .. plot:: + :context: close-figs + + import matplotlib.pyplot as plt + import numpy as np + from scipy.stats import nchypergeom_fisher + + plt.style.use("arviz-darkgrid") + x = np.arange(0, 12) + params = [ + (30, 10, 8, 1.0), + (30, 10, 8, 3.0), + (30, 10, 8, 10.0), + ] + for N, k, n, odds in params: + pmf = nchypergeom_fisher.pmf(x, N, k, n, odds) + plt.plot(x, pmf, "-o", label=r"$N$ = {}, $k$ = {}, $n$ = {}, $\theta$ = {}".format(N, k, n, odds)) + plt.xlabel("x", fontsize=12) + plt.ylabel("f(x)", fontsize=12) + plt.legend(loc=1) + plt.show() + + .. list-table:: + :widths: 15 85 + + * - Support + - :math:`x \in \{\max(0, n-(N-k)), \ldots, \min(n, k)\}` + * - Mean + - :math:`\mathrm{E}[X]` + * - Variance + - :math:`\mathrm{Var}[X]` + * - Median + - No closed form; computed numerically + + Parameters + ---------- + N : tensor_like of int + Total population size (N > 0) + k : tensor_like of int + Number of successful items in the population (0 <= k <= N) + n : tensor_like of int + Number of draws from the population (0 <= n <= N) + odds : tensor_like of float + Odds ratio for selecting a successful item (odds > 0) + + References + ---------- + .. [1] "Fisher's noncentral hypergeometric distribution." + https://en.wikipedia.org/wiki/Fisher%27s_noncentral_hypergeometric_distribution + """ + + rv_op = fisher_noncentral_hypergeometric + + @classmethod + def dist(cls, N, k, n, odds, *args, **kwargs): + N = pt.as_tensor_variable(N, dtype=int) + good = pt.as_tensor_variable(k, dtype=int) + bad = N - k + n = pt.as_tensor_variable(n, dtype=int) + odds = pt.as_tensor_variable(odds) + return super().dist([good, bad, n, odds], *args, **kwargs) + + def support_point(rv, size, good, bad, n, odds): + A = odds - 1 + B = n - bad - (good + n + 2) * odds + C = (good + 1) * (n + 1) * odds + mode = pt.floor(-2 * C / (B - pt.sqrt(B**2 - 4 * A * C))) + if not rv_size_is_none(size): + mode = pt.full(size, mode) + return mode + + def _norm(good, bad, n, odds): + mode = FisherNoncentralHypergeometric.support_point(None, None, good, bad, n, odds) + mode_scale = FisherNoncentralHypergeometric._logweight(mode, good, bad, n, odds) + max_draws = pt.max(n) + 1 + arange = pt.arange(max_draws).reshape((max_draws, 1)) + log_weights = FisherNoncentralHypergeometric._logweight(arange, good, bad, n, odds) + scaled_log_sum = pt.logsumexp(log_weights - mode_scale, axis=0) + return mode_scale + scaled_log_sum + + def logp(value, good, bad, n, odds): + return FisherNoncentralHypergeometric._logweight( + value, good, bad, n, odds + ) - FisherNoncentralHypergeometric._norm(good, bad, n, odds) + + def _logweight(value, good, bad, n, odds): + fails_value = n - value + result = ( + logpow(odds, value) + - factln(value) + - factln(good - value) + - factln(fails_value) + - factln(bad - fails_value) + ) + + lower_bound = n - bad + lower = pt.switch(pt.gt(lower_bound, 0), lower_bound, 0) + upper = pt.switch(pt.lt(good, n), good, n) + res = pt.switch( + pt.lt(value, lower), + -np.inf, + pt.switch( + pt.le(value, upper), + result, + -np.inf, + ), + ) + + return check_parameters( + res, + 0 <= good, + 0 <= bad, + 0 <= n, + n <= good + bad, + odds > 0, + msg="N > 0, 0 <= k <= N, 0 <= n <= N, odds > 0", + ) + + def logcdf(value, good, bad, n, odds): + if np.ndim(value): + raise TypeError( + f"FisherNoncentralHypergeometric.logcdf expects a scalar value but received a {np.ndim(value)}-dimensional object." + ) + + res = pt.switch( + pt.lt(value, 0), + -np.inf, + pt.switch( + pt.lt(value, n), + pt.logsumexp( + FisherNoncentralHypergeometric.logp(pt.arange(value + 1), good, bad, n, odds), + axis=0, + ), + 0, + ), + ) + + return check_parameters( + res, + 0 <= good, + 0 <= bad, + 0 <= n, + n <= good + bad, + odds > 0, + msg="N > 0, 0 <= k <= N, 0 <= n <= N, odds > 0", + ) + + class BetaNegativeBinomial: R""" Beta Negative Binomial distribution. diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py index 434768ac3..429e35b76 100644 --- a/tests/distributions/test_discrete.py +++ b/tests/distributions/test_discrete.py @@ -22,21 +22,31 @@ BaseTestDistributionRandom, Domain, I, + Nat, Rplus, assert_support_point_is_expected, + check_logcdf, check_logp, + check_selfconsistency_discrete_logcdf, discrete_random_tester, + seeded_scipy_distribution_builder, ) from pytensor import config, function +from pytensor.utils import lazy_scipy_module from pymc_extras.distributions import ( BetaNegativeBinomial, + FisherNoncentralHypergeometric, GeneralizedPoisson, Skellam, ) +stats = lazy_scipy_module("stats") + + pytestmark = pytest.mark.filterwarnings( - "ignore:Numba will use object mode to run generalized_poisson_rv:UserWarning" + # "ignore:Numba will use object mode to run generalized_poisson_rv:UserWarning" + "ignore:Numba will use object mode to run:UserWarning" ) @@ -130,6 +140,76 @@ def test_moment(self, mu, lam, size, expected): assert_support_point_is_expected(model, expected) +class TestFisherNoncentralHypergeometric(BaseTestDistributionRandom): + pymc_dist = FisherNoncentralHypergeometric + pymc_dist_params = {"N": 20, "k": 12, "n": 5, "odds": 2.0} + expected_rv_op_params = { + "n": pymc_dist_params["k"], + "M": pymc_dist_params["N"], + "N": pymc_dist_params["n"], + "odds": pymc_dist_params["odds"], + } + reference_dist_params = expected_rv_op_params + reference_dist = seeded_scipy_distribution_builder("nchypergeom_fisher") + checks_to_run = [ + "check_pymc_draws_match_reference", + "check_rv_size", + ] + + def test_fisher_noncentral_hypergeometric(self): + N_domain = Domain([0, 10, 20, 30, np.inf], dtype="int64") + check_logp( + FisherNoncentralHypergeometric, + Nat, + { + "N": N_domain, + "k": Nat, + "n": Nat, + "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), + }, + lambda value, N, k, n, odds: stats.nchypergeom_fisher.logpmf(value, N, k, n, odds), + n_samples=-1, + ) + check_logcdf( + FisherNoncentralHypergeometric, + Nat, + { + "N": N_domain, + "k": Nat, + "n": Nat, + "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), + }, + lambda value, N, k, n, odds: stats.nchypergeom_fisher.logcdf(value, N, k, n, odds), + n_samples=-1, + ) + check_selfconsistency_discrete_logcdf( + FisherNoncentralHypergeometric, + Nat, + { + "N": N_domain, + "k": Nat, + "n": Nat, + "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), + }, + n_samples=-1, + ) + + @pytest.mark.parametrize( + "N, k, n, odds, size, expected", + [ + (10, 5, 3, 2.0, None, 2), + (10, 5, 3, 2.0, 5, np.full(5, 2)), + (10, [3, 5], 3, 2.0, None, [1, 2]), + (10, 5, [3, 5], 2.0, None, [2, 3]), + (10, 5, [3, 5], 2.0, (5, 2), [[2, 3]] * 5), + ], + ) + def test_moment(self, N, k, n, odds, size, expected): + with pm.Model() as model: + FisherNoncentralHypergeometric("x", N=N, k=k, n=n, odds=odds, size=size) + assert_support_point_is_expected(model, expected) + + class TestBetaNegativeBinomial: """ Wrapper class so that tests of experimental additions can be dropped into From 3c8b55c8b989a75fbbe566142e523156fcb191a6 Mon Sep 17 00:00:00 2001 From: kaylimekay Date: Sat, 8 Aug 2026 19:28:41 -0400 Subject: [PATCH 2/3] Remove n_samples argument --- tests/distributions/test_discrete.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py index 429e35b76..f9b1e8d00 100644 --- a/tests/distributions/test_discrete.py +++ b/tests/distributions/test_discrete.py @@ -168,7 +168,6 @@ def test_fisher_noncentral_hypergeometric(self): "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), }, lambda value, N, k, n, odds: stats.nchypergeom_fisher.logpmf(value, N, k, n, odds), - n_samples=-1, ) check_logcdf( FisherNoncentralHypergeometric, @@ -180,7 +179,6 @@ def test_fisher_noncentral_hypergeometric(self): "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), }, lambda value, N, k, n, odds: stats.nchypergeom_fisher.logcdf(value, N, k, n, odds), - n_samples=-1, ) check_selfconsistency_discrete_logcdf( FisherNoncentralHypergeometric, @@ -191,7 +189,6 @@ def test_fisher_noncentral_hypergeometric(self): "n": Nat, "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), }, - n_samples=-1, ) @pytest.mark.parametrize( From aa6847c2a62a600b6d22f1f5b908c5088e4a1964 Mon Sep 17 00:00:00 2001 From: kaylimekay Date: Sat, 8 Aug 2026 19:40:10 -0400 Subject: [PATCH 3/3] Remove lazy scipy --- tests/distributions/test_discrete.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py index f9b1e8d00..d451e5140 100644 --- a/tests/distributions/test_discrete.py +++ b/tests/distributions/test_discrete.py @@ -32,7 +32,6 @@ seeded_scipy_distribution_builder, ) from pytensor import config, function -from pytensor.utils import lazy_scipy_module from pymc_extras.distributions import ( BetaNegativeBinomial, @@ -41,9 +40,6 @@ Skellam, ) -stats = lazy_scipy_module("stats") - - pytestmark = pytest.mark.filterwarnings( # "ignore:Numba will use object mode to run generalized_poisson_rv:UserWarning" "ignore:Numba will use object mode to run:UserWarning" @@ -167,7 +163,9 @@ def test_fisher_noncentral_hypergeometric(self): "n": Nat, "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), }, - lambda value, N, k, n, odds: stats.nchypergeom_fisher.logpmf(value, N, k, n, odds), + lambda value, N, k, n, odds: scipy.stats.nchypergeom_fisher.logpmf( + value, N, k, n, odds + ), ) check_logcdf( FisherNoncentralHypergeometric, @@ -178,7 +176,9 @@ def test_fisher_noncentral_hypergeometric(self): "n": Nat, "odds": Domain([0.01, 0.1, 0.9, 1, 1.5, 2, np.inf]), }, - lambda value, N, k, n, odds: stats.nchypergeom_fisher.logcdf(value, N, k, n, odds), + lambda value, N, k, n, odds: scipy.stats.nchypergeom_fisher.logcdf( + value, N, k, n, odds + ), ) check_selfconsistency_discrete_logcdf( FisherNoncentralHypergeometric,