Add ARGUS distribution (#109)

Closes #104

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Hans Dembinski <[email protected]>
Co-authored-by: Hans Dembinski <[email protected]>

src/numba_stats/_special.py

@@ -30,6 +30,7 @@ def get(name, signature):
 ndtri = get("ndtri", float64(float64))
 
 # binary functions (double)
+gammainc = get("gammainc", float64(float64, float64))
 gammaincc = get("gammaincc", float64(float64, float64))
 stdtrit = get("stdtrit", float64(float64, float64))
 stdtr = get("stdtr", float64(float64, float64))

src/numba_stats/argus.py

@@ -0,0 +1,93 @@
+"""
+Generalised ARGUS distribution.
+
+The ARGUS distribution is named after the particle physics experiment ARGUS and it
+describes the reconstructed invariant mass of a decayed particle candidate
+in continuum background.
+It is motivated from experimental observation. Here we have the generalised version
+of the ARGUS distribution that can be used to describe a more peaking like distribtion.
+p = 0.5 gives the normal ARGUS distribution.
+
+https://en.wikipedia.org/wiki/ARGUS_distribution
+
+See Also
+--------
+scipy.stats.argus: Scipy equivalent.
+"""
+
+from math import lgamma as _lg
+
+import numpy as np
+
+from ._special import gammainc as _ginc
+from ._util import _generate_wrappers, _jit, _prange
+
+_doc_par = """
+x : Array-like
+    Random variate, between 0 and c.
+chi : float
+    Must be larger than 0 and represents the cutoff.
+c : float
+    Must be larger than 0 and represents the curvature.
+p : float
+    Must be larger than -1 and represents the power.
+"""
+
+
+@_jit(3, cache=False)
+def _logpdf(x, chi, c, p):
+    T = type(p)
+    one = T(1)
+    two = T(2)
+    half = T(0.5)
+    half_chi2 = half * chi * chi
+    inv_c2 = one / (c * c)
+    p1 = p + one
+    r = np.empty_like(x)
+    for i in _prange(len(x)):
+        xi = x[i]
+        if 0 <= xi and xi <= c:
+            y = one - xi * xi * inv_c2
+            r[i] = (
+                -half_chi2 * y
+                + p * (np.log(y) - np.log(two))
+                + two * (p1 * np.log(chi) - np.log(c))
+                + np.log(xi)
+                - T(_lg(p1))
+                - np.log(T(_ginc(p1, half_chi2)))
+            )
+        else:
+            r[i] = -np.inf
+    return r
+
+
+@_jit(3, cache=False)
+def _pdf(x, chi, c, p):
+    return np.exp(_logpdf(x, chi, c, p))
+
+
+@_jit(3, cache=False)
+def _cdf(x, chi, c, p):
+    T = type(p)
+    zero = T(0)
+    one = T(1)
+    half = T(0.5)
+    p1 = p + one
+    half_chi2 = half * chi * chi
+    inv_c2 = one / (c * c)
+    inv_ginc = one / T(_ginc(p1, half_chi2))
+    r = np.empty_like(x)
+    for i in _prange(len(x)):
+        xi = x[i]
+        if 0 <= xi:
+            if xi <= c:
+                y = one - xi * xi * inv_c2
+                r[i] = one - T(_ginc(p1, half_chi2 * y)) * inv_ginc
+            else:
+                r[i] = one
+        else:
+            r[i] = zero
+    return r
+
+
+_generate_wrappers(globals())

tests/test_argus.py

@@ -0,0 +1,36 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+from scipy import stats as sc
+
+from numba_stats import argus
+
+
+@pytest.mark.parametrize("chi", (0.1, 0.5, 1.0, 2.0, 3.0))
+def test_logpdf(chi):
+    c = 1
+    p = 0.5
+    x = np.linspace(-1, 2, 30)
+    got = argus.logpdf(x, chi, c, p)
+    expected = sc.argus.logpdf(x, chi)
+    assert_allclose(got, expected)
+
+
+@pytest.mark.parametrize("chi", (0.1, 0.5, 1.0, 2.0, 3.0))
+def test_pdf(chi):
+    c = 1
+    p = 0.5
+    x = np.linspace(-1, 2, 30)
+    got = argus.pdf(x, chi, c, p)
+    expected = sc.argus.pdf(x, chi)
+    assert_allclose(got, expected)
+
+
+@pytest.mark.parametrize("chi", (0.1, 0.5, 1.0, 2.0, 3.0))
+def test_cdf(chi):
+    c = 1
+    p = 0.5
+    x = np.linspace(-1, 2, 30)
+    got = argus.cdf(x, chi, c, p)
+    expected = sc.argus.cdf(x, chi)
+    assert_allclose(got, expected, atol=2e-16)