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

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

2 files changed, 58 insertions, 0 deletions

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))