Add Fisher noncentral hypergeometric distribution

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

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",

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.

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