BUG: stats: Fix zipf.pmf and zipfian.pmf for int32 k (#20702) […

…wheel build] * BUG: fix zipf.pmf for integer k * ENH: increase range for zipfian.pmf for integer k --------- Co-authored-by: Tyler Reddy <[email protected]> [wheel build]
scipy · May 22, 2024 · 11e99ba · 11e99ba
1 parent 506cbeb
commit 11e99ba
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/scipy/stats/_discrete_distns.py b/scipy/stats/_discrete_distns.py
@@ -1305,8 +1305,9 @@ def _argcheck(self, a):
         return a > 1
 
     def _pmf(self, k, a):
+        k = k.astype(np.float64)
         # zipf.pmf(k, a) = 1/(zeta(a) * k**a)
-        Pk = 1.0 / special.zeta(a, 1) / k**a
+        Pk = 1.0 / special.zeta(a, 1) * k**-a
         return Pk
 
     def _munp(self, n, a):
@@ -1404,7 +1405,8 @@ def _get_support(self, a, n):
         return 1, n
 
     def _pmf(self, k, a, n):
-        return 1.0 / _gen_harmonic(n, a) / k**a
+        k = k.astype(np.float64)
+        return 1.0 / _gen_harmonic(n, a) * k**-a
 
     def _cdf(self, k, a, n):
         return _gen_harmonic(k, a) / _gen_harmonic(n, a)

diff --git a/scipy/stats/tests/test_discrete_distns.py b/scipy/stats/tests/test_discrete_distns.py
@@ -351,6 +351,14 @@ def pzip(k, a, n):
         assert_allclose(zipfian.stats(a, n, moments="mvsk"),
                         [mean, var, skew, kurtosis])
 
+    def test_pmf_integer_k(self):
+        k = np.arange(0, 1000)
+        k_int32 = k.astype(np.int32)
+        dist = zipfian(111, 22)
+        pmf = dist.pmf(k)
+        pmf_k_int32 = dist.pmf(k_int32)
+        assert_equal(pmf, pmf_k_int32)
+
 
 class TestNCH:
     np.random.seed(2)  # seeds 0 and 1 had some xl = xu; randint failed
@@ -627,3 +635,14 @@ def test_betanbinom_kurtosis(self, n, a, b, ref):
         #      return float(fourth_moment/var**2 - 3.)
         assert_allclose(betanbinom.stats(n, a, b, moments="k"),
                         ref, rtol=3e-15)
+
+
+class TestZipf:
+    def test_gh20692(self):
+        # test that int32 data for k generates same output as double
+        k = np.arange(0, 1000)
+        k_int32 = k.astype(np.int32)
+        dist = zipf(9)
+        pmf = dist.pmf(k)
+        pmf_k_int32 = dist.pmf(k_int32)
+        assert_equal(pmf, pmf_k_int32)