BUG: stats: adapt to np.floor type promotion removal (#21283)

lucascolley · web-flow · commit ad58b229ed4a · 2024-08-10T12:16:39.000-07:00
* BUG: stats: adapt to `np.floor` type promotion removal `rv_discrete._cdf` relied on `np.floor` promoting its integer input to `np.float64`. This is no longer the case since numpy/numpy#26766.
diff --git a/scipy/stats/_distn_infrastructure.py b/scipy/stats/_distn_infrastructure.py
@@ -3393,7 +3393,7 @@ def _cdf_single(self, k, *args):
         return np.sum(self._pmf(m, *args), axis=0)
 
     def _cdf(self, x, *args):
-        k = floor(x)
+        k = floor(x).astype(np.float64)
         return self._cdfvec(k, *args)
 
     # generic _logcdf, _sf, _logsf, _ppf, _isf, _rvs defined in rv_generic
diff --git a/scipy/stats/tests/test_discrete_basic.py b/scipy/stats/tests/test_discrete_basic.py
@@ -550,3 +550,15 @@ def test_rv_sample():
     rng = np.random.default_rng(98430143469)
     rvs0 = dist.ppf(rng.random(size=100))
     assert_allclose(rvs, rvs0)
+
+def test__pmf_float_input():
+    # gh-21272
+    # test that `rvs()` can be computed when `_pmf` requires float input
+    
+    class rv_exponential(stats.rv_discrete):
+        def _pmf(self, i):
+            return (2/3)*3**(1 - i)
+    
+    rv = rv_exponential(a=0.0, b=float('inf'))
+    rvs = rv.rvs(random_state=42)  # should not crash due to integer input to `_pmf`
+    assert_allclose(rvs, 0)
