From 7510b9e9cb57aef39f598ccbc3f6d9b72362c417 Mon Sep 17 00:00:00 2001
From: Hans Dembinski <HDembinski@users.noreply.github.com>
Date: Wed, 6 Dec 2023 16:09:38 +0100
Subject: [PATCH] add laplace distribution (#88)

Closes #82
---
 src/numba_stats/expon.py   | 34 --------------------
 src/numba_stats/laplace.py | 66 ++++++++++++++++++++++++++++++++++++++
 tests/test_laplace.py      | 31 ++++++++++++++++++
 3 files changed, 97 insertions(+), 34 deletions(-)
 create mode 100644 src/numba_stats/laplace.py
 create mode 100644 tests/test_laplace.py

diff --git a/src/numba_stats/expon.py b/src/numba_stats/expon.py
index 73d73d9..c5a20e6 100644
--- a/src/numba_stats/expon.py
+++ b/src/numba_stats/expon.py
@@ -39,45 +39,11 @@ def _logpdf(x, loc, scale):
 
 @_jit(2)
 def _pdf(x, loc, scale):
-    """
-    Return probability density.
-
-    Parameters
-    ----------
-    x: ArrayLike
-        Random variate.
-    loc : float
-        Location of the mode.
-    scale : float
-        Standard deviation.
-
-    Returns
-    -------
-    ArrayLike
-        Function evaluated at x.
-    """
     return np.exp(_logpdf(x, loc, scale))
 
 
 @_jit(2)
 def _cdf(x, loc, scale):
-    """
-    Return cumulative probability.
-
-    Parameters
-    ----------
-    x: ArrayLike
-        Random variate.
-    loc : float
-        Location of the mode.
-    scale : float
-        Standard deviation.
-
-    Returns
-    -------
-    ArrayLike
-        Function evaluated at x.
-    """
     z = _trans(x, loc, scale)
     for i in _prange(len(z)):
         z[i] = _cdf1(z[i])
diff --git a/src/numba_stats/laplace.py b/src/numba_stats/laplace.py
new file mode 100644
index 0000000..c4372de
--- /dev/null
+++ b/src/numba_stats/laplace.py
@@ -0,0 +1,66 @@
+"""
+Laplace distribution.
+
+See Also
+--------
+scipy.stats.laplace: Scipy equivalent.
+"""
+import numpy as np
+from ._util import _jit, _trans, _generate_wrappers, _prange, _rvs_jit, _seed
+
+_doc_par = """
+loc : float
+    Location of the mode.
+scale : float
+    Standard deviation.
+"""
+
+
+@_jit(-1)
+def _cdf1(z):
+    return 1.0 - 0.5 * np.exp(-z) if z > 0 else 0.5 * np.exp(z)
+
+
+@_jit(-1)
+def _ppf1(p):
+    return -np.log(2 * (1 - p)) if p > 0.5 else np.log(2 * p)
+
+
+@_jit(2)
+def _logpdf(x, loc, scale):
+    z = _trans(x, loc, scale)
+    r = np.empty_like(z)
+    for i in _prange(len(r)):
+        r[i] = np.log(0.25) - np.abs(z[i])
+    return r
+
+
+@_jit(2)
+def _pdf(x, loc, scale):
+    return np.exp(_logpdf(x, loc, scale))
+
+
+@_jit(2)
+def _cdf(x, loc, scale):
+    z = _trans(x, loc, scale)
+    for i in _prange(len(z)):
+        z[i] = _cdf1(z[i])
+    return z
+
+
+@_jit(2)
+def _ppf(p, loc, scale):
+    z = np.empty_like(p)
+    for i in _prange(len(z)):
+        z[i] = _ppf1(p[i])
+    return scale * z + loc
+
+
+@_rvs_jit(2)
+def _rvs(loc, scale, size, random_state):
+    _seed(random_state)
+    p = np.random.uniform(0, 1, size)
+    return _ppf(p, loc, scale)
+
+
+_generate_wrappers(globals())
diff --git a/tests/test_laplace.py b/tests/test_laplace.py
new file mode 100644
index 0000000..1ad60be
--- /dev/null
+++ b/tests/test_laplace.py
@@ -0,0 +1,31 @@
+import scipy.stats as sc
+import numpy as np
+from numba_stats import laplace
+
+
+def test_pdf():
+    x = np.linspace(-5, 5, 20)
+    got = laplace.pdf(x, 1, 2)
+    expected = sc.laplace.pdf(x, 1, 2)
+    np.testing.assert_allclose(got, expected)
+
+
+def test_cdf():
+    x = np.linspace(-5, 5, 20) + 3
+    got = laplace.cdf(x, 3, 2)
+    expected = sc.laplace.cdf(x, 3, 2)
+    np.testing.assert_allclose(got, expected)
+
+
+def test_ppf():
+    p = np.linspace(0, 1, 20)
+    got = laplace.ppf(p, 1, 2)
+    expected = sc.laplace.ppf(p, 1, 2)
+    np.testing.assert_allclose(got, expected)
+
+
+def test_rvs():
+    args = 1, 2
+    x = laplace.rvs(*args, size=100_000, random_state=1)
+    r = sc.kstest(x, lambda x: laplace.cdf(x, *args))
+    assert r.pvalue > 0.01