Repository: HDembinski/numba-stats (language: python)

Issue:
Add ArgusBG (generalised ARGUS) function
Hi,
scipy provides only the [ARGUS function](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.argus.html#scipy.stats.argus).

In ROOT is instead some function implemented called ARGUSBG, see [here](https://root.cern.ch/doc/master/classRooArgusBG.html). If I see it correctly this relates to what Wikipedia calls a generalised [ARGUS function](https://en.wikipedia.org/wiki/ARGUS_distribution).

It would be nice if this could be added to numba-stats.


Discussion / hints:
@jpwgnr Would you like to try a PR for that?

Looking at the Wikipedia, you need the gamma function and the upper incomplete gamma function. The former is available as `math.lgamma` (log of gamma function) and latter is available as `numba_stats._special.gammaincc`.

The implementations contain differences of gamma functions that can grow very large. I recommend to implement the `logpdf` and compute those differences with `numpy.logaddexp(lgamma(...), -lgamma(...))`.

Fix patch:
diff --git a/src/numba_stats/_special.py b/src/numba_stats/_special.py
index 8087e52..802e202 100644
--- a/src/numba_stats/_special.py
+++ b/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))
diff --git a/src/numba_stats/argus.py b/src/numba_stats/argus.py
new file mode 100644
index 0000000..4f23bbf
--- /dev/null
+++ b/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())


Test patch:
diff --git a/tests/test_argus.py b/tests/test_argus.py
new file mode 100644
index 0000000..23ed870
--- /dev/null
+++ b/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)