BUG: Fix random.choice when probability is not C contiguous (#13716)

* test, fix random.choice sum for non-contiguous

4 files changed, 70 insertions, 6 deletions

diff --git a/numpy/random/generator.pyx b/numpy/random/generator.pyx
index e84a98150..368a03e8f 100644
--- a/numpy/random/generator.pyx
+++ b/numpy/random/generator.pyx
@@ -611,7 +611,8 @@ cdef class Generator:
                 if np.issubdtype(p.dtype, np.floating):
                     atol = max(atol, np.sqrt(np.finfo(p.dtype).eps))
 
-            p = <np.ndarray>np.PyArray_FROM_OTF(p, np.NPY_DOUBLE, np.NPY_ALIGNED)
+            p = <np.ndarray>np.PyArray_FROM_OTF(
+                p, np.NPY_DOUBLE, np.NPY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
             pix = <double*>np.PyArray_DATA(p)
 
             if p.ndim != 1:
@@ -3541,7 +3542,8 @@ cdef class Generator:
 
         d = len(pvals)
         on = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
-        parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
+        parr = <np.ndarray>np.PyArray_FROM_OTF(
+            pvals, np.NPY_DOUBLE, np.NPY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
         pix = <double*>np.PyArray_DATA(parr)
         check_array_constraint(parr, 'pvals', CONS_BOUNDED_0_1)
         if kahan_sum(pix, d-1) > (1.0 + 1e-12):
@@ -3695,7 +3697,8 @@ cdef class Generator:
         cdef double acc, invacc
 
         k = len(alpha)
-        alpha_arr = <np.ndarray>np.PyArray_FROM_OTF(alpha, np.NPY_DOUBLE, np.NPY_ALIGNED)
+        alpha_arr = <np.ndarray>np.PyArray_FROM_OTF(
+            alpha, np.NPY_DOUBLE, np.NPY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
         if np.any(np.less_equal(alpha_arr, 0)):
             raise ValueError('alpha <= 0')
         alpha_data = <double*>np.PyArray_DATA(alpha_arr)
diff --git a/numpy/random/mtrand.pyx b/numpy/random/mtrand.pyx
index 4184d1ebe..19f7ab86a 100644
--- a/numpy/random/mtrand.pyx
+++ b/numpy/random/mtrand.pyx
@@ -779,7 +779,8 @@ cdef class RandomState:
                 if np.issubdtype(p.dtype, np.floating):
                     atol = max(atol, np.sqrt(np.finfo(p.dtype).eps))
 
-            p = <np.ndarray>np.PyArray_FROM_OTF(p, np.NPY_DOUBLE, np.NPY_ALIGNED)
+            p = <np.ndarray>np.PyArray_FROM_OTF(
+                p, np.NPY_DOUBLE, np.NPY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
             pix = <double*>np.PyArray_DATA(p)
 
             if p.ndim != 1:
@@ -3855,7 +3856,8 @@ cdef class RandomState:
         cdef long ni
 
         d = len(pvals)
-        parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
+        parr = <np.ndarray>np.PyArray_FROM_OTF(
+            pvals, np.NPY_DOUBLE, np.NPY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
         pix = <double*>np.PyArray_DATA(parr)
         check_array_constraint(parr, 'pvals', CONS_BOUNDED_0_1)
         if kahan_sum(pix, d-1) > (1.0 + 1e-12):
@@ -3987,7 +3989,8 @@ cdef class RandomState:
         cdef double  acc, invacc
 
         k = len(alpha)
-        alpha_arr = <np.ndarray>np.PyArray_FROM_OTF(alpha, np.NPY_DOUBLE, np.NPY_ALIGNED)
+        alpha_arr = <np.ndarray>np.PyArray_FROM_OTF(
+            alpha, np.NPY_DOUBLE, np.NPY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
         if np.any(np.less_equal(alpha_arr, 0)):
             raise ValueError('alpha <= 0')
         alpha_data = <double*>np.PyArray_DATA(alpha_arr)
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
index 64d87cd71..06c88af85 100644
--- a/numpy/random/tests/test_generator_mt19937.py
+++ b/numpy/random/tests/test_generator_mt19937.py
@@ -101,6 +101,16 @@ class TestMultinomial(object):
     def test_invalid_n(self):
         assert_raises(ValueError, random.multinomial, -1, [0.8, 0.2])
         assert_raises(ValueError, random.multinomial, [-1] * 10, [0.8, 0.2])
+    
+    def test_p_non_contiguous(self):
+        p = np.arange(15.)
+        p /= np.sum(p[1::3])
+        pvals = p[1::3]
+        random.bit_generator.seed(1432985819)
+        non_contig = random.multinomial(100, pvals=pvals)
+        random.bit_generator.seed(1432985819)
+        contig = random.multinomial(100, pvals=np.ascontiguousarray(pvals))
+        assert_array_equal(non_contig, contig)
 
 
 class TestSetState(object):
@@ -656,6 +666,15 @@ class TestRandomDist(object):
         a = np.array([42, 1, 2])
         p = [None, None, None]
         assert_raises(ValueError, random.choice, a, p=p)
+    
+    def test_choice_p_non_contiguous(self):
+        p = np.ones(10) / 5
+        p[1::2] = 3.0
+        random.bit_generator.seed(self.seed)
+        non_contig = random.choice(5, 3, p=p[::2])
+        random.bit_generator.seed(self.seed)
+        contig = random.choice(5, 3, p=np.ascontiguousarray(p[::2]))
+        assert_array_equal(non_contig, contig)
 
     def test_choice_return_type(self):
         # gh 9867
@@ -804,6 +823,16 @@ class TestRandomDist(object):
         # gh-2089
         alpha = np.array([5.4e-01, -1.0e-16])
         assert_raises(ValueError, random.dirichlet, alpha)
+    
+    def test_dirichlet_alpha_non_contiguous(self):
+        a = np.array([51.72840233779265162, -1.0, 39.74494232180943953])
+        alpha = a[::2]
+        random.bit_generator.seed(self.seed)
+        non_contig = random.dirichlet(alpha, size=(3, 2))
+        random.bit_generator.seed(self.seed)
+        contig = random.dirichlet(np.ascontiguousarray(alpha),
+                                  size=(3, 2))
+        assert_array_almost_equal(non_contig, contig)
 
     def test_exponential(self):
         random.bit_generator.seed(self.seed)
diff --git a/numpy/random/tests/test_randomstate.py b/numpy/random/tests/test_randomstate.py
index d8a07e8b2..e075a97bc 100644
--- a/numpy/random/tests/test_randomstate.py
+++ b/numpy/random/tests/test_randomstate.py
@@ -157,6 +157,16 @@ class TestMultinomial(object):
     def test_invalid_n(self):
         assert_raises(ValueError, random.multinomial, -1, [0.8, 0.2])
 
+    def test_p_non_contiguous(self):
+        p = np.arange(15.)
+        p /= np.sum(p[1::3])
+        pvals = p[1::3]
+        random.seed(1432985819)
+        non_contig = random.multinomial(100, pvals=pvals)
+        random.seed(1432985819)
+        contig = random.multinomial(100, pvals=np.ascontiguousarray(pvals))
+        assert_array_equal(non_contig, contig)
+
 
 class TestSetState(object):
     def setup(self):
@@ -608,6 +618,15 @@ class TestRandomDist(object):
         a = np.array([42, 1, 2])
         p = [None, None, None]
         assert_raises(ValueError, random.choice, a, p=p)
+    
+    def test_choice_p_non_contiguous(self):
+        p = np.ones(10) / 5
+        p[1::2] = 3.0
+        random.seed(self.seed)
+        non_contig = random.choice(5, 3, p=p[::2])
+        random.seed(self.seed)
+        contig = random.choice(5, 3, p=np.ascontiguousarray(p[::2]))
+        assert_array_equal(non_contig, contig)
 
     def test_bytes(self):
         random.seed(self.seed)
@@ -734,6 +753,16 @@ class TestRandomDist(object):
         alpha = np.array([5.4e-01, -1.0e-16])
         assert_raises(ValueError, random.dirichlet, alpha)
 
+    def test_dirichlet_alpha_non_contiguous(self):
+        a = np.array([51.72840233779265162, -1.0, 39.74494232180943953])
+        alpha = a[::2]
+        random.seed(self.seed)
+        non_contig = random.dirichlet(alpha, size=(3, 2))
+        random.seed(self.seed)
+        contig = random.dirichlet(np.ascontiguousarray(alpha),
+                                  size=(3, 2))
+        assert_array_almost_equal(non_contig, contig)
+
     def test_exponential(self):
         random.seed(self.seed)
         actual = random.exponential(1.1234, size=(3, 2))