Merge pull request #26 from bashtage/mv-normal-zig

ENH: Add support fo ziggurat method for all normal

Author: bashtage

randomstate/interface.pyx

@@ -1112,7 +1112,7 @@ cdef class RandomState:
 
     def randn(self, *args, method=__normal_method):
         """
-        randn(d0, d1, ..., dn)
+        randn(d0, d1, ..., dn, method='inv')
 
         Return a sample (or samples) from the "standard normal" distribution.
 
@@ -1132,6 +1132,9 @@ cdef class RandomState:
         d0, d1, ..., dn : int, optional
             The dimensions of the returned array, should be all positive.
             If no argument is given a single Python float is returned.
+        method : str, optional
+            Either 'inv' or 'zig'. 'inv' uses the default FIXME method.  'zig' uses
+            the much faster ziggurat method of FIXME.
 
         Returns
         -------
@@ -3674,9 +3677,9 @@ cdef class RandomState:
                  0.0, '', CONS_NONE)
 
     # Multivariate distributions:
-    def multivariate_normal(self, mean, cov, size=None):
+    def multivariate_normal(self, mean, cov, size=None, method=__normal_method):
         """
-        multivariate_normal(mean, cov[, size])
+        multivariate_normal(mean, cov, size=None, method='inv')
 
         Draw random samples from a multivariate normal distribution.
 
@@ -3699,6 +3702,9 @@ cdef class RandomState:
             generated, and packed in an `m`-by-`n`-by-`k` arrangement.  Because
             each sample is `N`-dimensional, the output shape is ``(m,n,k,N)``.
             If no shape is specified, a single (`N`-D) sample is returned.
+        method : str, optional
+            Either 'inv' or 'zig'. 'inv' uses the default FIXME method.  'zig' uses
+            the much faster ziggurat method of FIXME.
 
         Returns
         -------
@@ -3795,7 +3801,7 @@ cdef class RandomState:
         # form a matrix of shape final_shape.
         final_shape = tuple(shape[:])
         final_shape += (mean.shape[0],)
-        x = self.standard_normal(final_shape).reshape(-1, mean.shape[0])
+        x = self.standard_normal(final_shape, method=method).reshape(-1, mean.shape[0])
 
         # Transform matrix of standard normals into matrix where each row
         # contains multivariate normals with the desired covariance.

randomstate/tests/test_smoke.py

@@ -304,13 +304,22 @@ def test_rand(self):
 
     def test_randn(self):
         state = self.rs.get_state()
-        print(state)
         vals = self.rs.randn(10, 10, 10)
         self.rs.set_state(state)
-        print(self.rs.get_state())
         assert_equal(vals, self.rs.standard_normal((10, 10, 10)))
         assert_equal(vals.shape, (10, 10, 10))
 
+        state = self.rs.get_state()
+        vals = self.rs.randn(10, 10, 10, method='inv')
+        self.rs.set_state(state)
+        assert_equal(vals, self.rs.standard_normal((10, 10, 10), method='inv'))
+
+        state = self.rs.get_state()
+        vals_inv = self.rs.randn(10, 10, 10, method='inv')
+        self.rs.set_state(state)
+        vals_zig = self.rs.randn(10, 10, 10, method='zig')
+        assert (vals_zig != vals_inv).any()
+
     def test_noncentral_chisquare(self):
         vals = self.rs.noncentral_chisquare(10, 2, 10)
         assert len(vals) == 10
@@ -400,6 +409,11 @@ def test_multivariate_normal(self):
         cov = [[1, 0], [0, 100]]  # diagonal covariance
         x = self.rs.multivariate_normal(mean, cov, 5000)
         assert x.shape == (5000, 2)
+        x_zig = self.rs.multivariate_normal(mean, cov, 5000, method='zig')
+        assert x.shape == (5000, 2)
+        x_inv = self.rs.multivariate_normal(mean, cov, 5000, method='inv')
+        assert x.shape == (5000, 2)
+        assert (x_zig != x_inv).any()
 
     def test_multinomial(self):
         vals = self.rs.multinomial(100, [1.0 / 3, 2.0 / 3])