add poisson.rvs (#94)

Also makes norm.rvs faster and simpler

.github/workflows/ci.yml

@@ -14,11 +14,11 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.11"]
+        python-version: ["3.8", "3.12"]
 
     steps:
-    - uses: actions/checkout@v2
-    - uses: actions/setup-python@v2
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - run: python -m pip install --upgrade pip

.github/workflows/release.yml

@@ -9,9 +9,9 @@ jobs:
   upload:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
 
-    - uses: actions/setup-python@v2
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.9'
 
@@ -22,7 +22,7 @@ jobs:
     - run: python -m pip install --force-reinstall dist/*.tar.gz
     - run: python -m pytest
 
-    - uses: pypa/gh-action-pypi-publish@master
+    - uses: pypa/gh-action-pypi-publish@release/v1
       with:
         user: __token__
         password: ${{secrets.PYPI_TOKEN}}

src/numba_stats/norm.py

@@ -67,11 +67,10 @@ def _ppf(p, loc, scale):
     return r
 
 
-@_rvs_jit(2, cache=False)
+@_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)
+    return np.random.normal(loc, scale, size)
 
 
 _generate_wrappers(globals())

src/numba_stats/poisson.py

@@ -9,7 +9,8 @@
 import numpy as np
 from ._special import gammaincc as _gammaincc
 from math import lgamma as _lgamma
-from ._util import _jit, _generate_wrappers, _prange
+from ._util import _jit, _generate_wrappers, _prange, _seed
+import numba as nb
 
 _doc_par = """
 mu : float
@@ -43,4 +44,15 @@ def _cdf(k, mu):
     return r
 
 
+@nb.njit(
+    nb.int64[:](nb.float32, nb.uint64, nb.optional(nb.uint64)),
+    cache=True,
+    inline="always",
+    error_model="numpy",
+)
+def _rvs(mu, size, random_state):
+    _seed(random_state)
+    return np.random.poisson(mu, size)
+
+
 _generate_wrappers(globals())

tests/test_poisson.py

@@ -2,6 +2,7 @@
 from numba_stats import poisson
 import scipy.stats as sc
 import pytest
+import numba as nb
 
 
 @pytest.mark.parametrize("mu", np.linspace(0, 3, 5))
@@ -18,3 +19,15 @@ def test_cdf(mu):
     got = poisson.cdf(k, mu)
     expected = sc.poisson.cdf(k, mu)
     np.testing.assert_allclose(got, expected)
+
+
+@pytest.mark.parametrize("mu", np.linspace(0, 3, 5))
+def test_rvs(mu):
+    got = poisson.rvs(mu, size=1000, random_state=1)
+
+    @nb.njit
+    def expected():
+        np.random.seed(1)
+        return np.random.poisson(mu, 1000)
+
+    np.testing.assert_equal(got, expected())