TST: Add testing for complex normal

Add testing for complex normal for alternative paths

Author: bashtage

randomstate/randomstate.pyx

@@ -1795,11 +1795,12 @@ cdef class RandomState:
 
         >>> s = np.random.complex_normal(size=1000)
         """
+        if method != u'zig' or method != u'bm':
+            raise ValueError("method must be either 'bm' or 'zig'")
         cdef np.ndarray ogamma, orelation, oloc, randoms, v_real, v_imag, rho
         cdef double *randoms_data
         cdef double fgamma_r, fgamma_i, frelation_r, frelation_i, frho, f_v_real , f_v_imag, \
             floc_r, floc_i, f_real, f_imag, i_r_scale, r_scale, i_scale, f_rho
-        cdef complex cloc
         cdef np.npy_intp i, j, n
         cdef np.broadcast it
 
@@ -1830,8 +1831,13 @@ cdef class RandomState:
                 raise ValueError('Im(relation) ** 2 > Re(gamma ** 2 - relation** 2)')
 
             if size is None:
-                random_gauss_zig_double_fill(&self.rng_state, 1, &f_real)
-                random_gauss_zig_double_fill(&self.rng_state, 1, &f_imag)
+                if method == u'zig':
+                    random_gauss_zig_double_fill(&self.rng_state, 1, &f_real)
+                    random_gauss_zig_double_fill(&self.rng_state, 1, &f_imag)
+                else:
+                    random_gauss_fill(&self.rng_state, 1, &f_real)
+                    random_gauss_fill(&self.rng_state, 1, &f_imag)
+
                 f_imag *= sqrt(1 - f_rho * f_rho)
                 f_imag += f_rho * f_real
                 f_real *= sqrt(0.5 * f_v_real)
@@ -1847,13 +1853,22 @@ cdef class RandomState:
             r_scale = sqrt(0.5 * f_v_real)
             i_scale = sqrt(0.5 * f_v_imag)
             j = 0
-            with self.lock: # , nogil:
-                for i in range(n):
-                    random_gauss_zig_double_fill(&self.rng_state, 1, &f_real)
-                    random_gauss_zig_double_fill(&self.rng_state, 1, &f_imag)
-                    randoms_data[j+1] = floc_i + i_scale * (f_rho * f_real + i_r_scale * f_imag)
-                    randoms_data[j] = floc_r + r_scale * f_real
-                    j += 2
+            with self.lock, nogil:
+                if method == u'zig':
+                    for i in range(n):
+                        random_gauss_zig_double_fill(&self.rng_state, 1, &f_real)
+                        random_gauss_zig_double_fill(&self.rng_state, 1, &f_imag)
+                        randoms_data[j+1] = floc_i + i_scale * (f_rho * f_real + i_r_scale * f_imag)
+                        randoms_data[j] = floc_r + r_scale * f_real
+                        j += 2
+                else:
+                    for i in range(n):
+                        random_gauss_fill(&self.rng_state, 1, &f_real)
+                        random_gauss_fill(&self.rng_state, 1, &f_imag)
+                        randoms_data[j+1] = floc_i + i_scale * (f_rho * f_real + i_r_scale * f_imag)
+                        randoms_data[j] = floc_r + r_scale * f_real
+                        j += 2
+
 
             return randoms
 
@@ -1885,9 +1900,11 @@ cdef class RandomState:
         n = np.PyArray_SIZE(randoms)
 
         it = np.PyArray_MultiIterNew5(randoms, oloc, v_real, v_imag, rho)
-        # TODO: Box-Muller
         with self.lock, nogil:
-            random_gauss_zig_double_fill(&self.rng_state, 2 * n, randoms_data)
+            if method == u'zig':
+                random_gauss_zig_double_fill(&self.rng_state, 2 * n, randoms_data)
+            else:
+                random_gauss_fill(&self.rng_state, 2 * n, randoms_data)
         with nogil:
             j = 0
             for i in range(n):

randomstate/tests/test_against_numpy.py

@@ -45,11 +45,9 @@ def compare_2_input(f1, f2, is_np=False, is_scalar=False):
         inputs = inputs[:3]
 
     for i in inputs:
-        print(i[0], i[1])
         v1 = f1(*i[0], **i[1])
         v2 = f2(*i[0], **i[1])
         assert_allclose(v1, v2)
-        print('OK!' * 20)
 
 
 def compare_3_input(f1, f2, is_np=False):

randomstate/tests/test_smoke.py

@@ -287,8 +287,60 @@ def test_chisquare(self):
         params_1(self.rs.chisquare)
 
     def test_complex_normal(self):
-        vals = self.rs.complex_normal(2.0 + 7.0j, 10.0, 5.0 - 5.0j, size=10)
-        assert_(len(vals) == 10)
+        st = self.rs.get_state()
+        vals = self.rs.complex_normal(2.0 + 7.0j, 10.0, 5.0 - 5.0j, size=10, method='zig')
+        assert_(len(vals) == 10)
+
+        self.rs.set_state(st)
+        vals2 = [self.rs.complex_normal(2.0 + 7.0j, 10.0, 5.0 - 5.0j, method='zig') for _ in range(10)]
+        np.testing.assert_allclose(vals, vals2)
+
+        self.rs.set_state(st)
+        vals3 = self.rs.complex_normal(2.0 + 7.0j * np.ones(10), 10.0 * np.ones(1), 5.0 - 5.0j, method='zig')
+        np.testing.assert_allclose(vals, vals3)
+
+        self.rs.set_state(st)
+        norms = self.rs.standard_normal(size=20, method='zig')
+        norms = np.reshape(norms, (10, 2))
+        cov = 0.5 * (-5.0)
+        v_real = 7.5
+        v_imag = 2.5
+        rho = cov / np.sqrt(v_real * v_imag)
+        imag = 7 + np.sqrt(v_imag) * (rho * norms[:, 0] + np.sqrt(1 - rho ** 2) * norms[:, 1])
+        real = 2 + np.sqrt(v_real) * norms[:, 0]
+        vals4 = [re + im * (0 + 1.0j) for re, im in zip(real, imag)]
+
+        np.testing.assert_allclose(vals4, vals)
+
+    def test_complex_normal_bm(self):
+        st = self.rs.get_state()
+        vals = self.rs.complex_normal(2.0 + 7.0j, 10.0, 5.0 - 5.0j, size=10, method='bm')
+        assert_(len(vals) == 10)
+
+        self.rs.set_state(st)
+        vals2 = [self.rs.complex_normal(2.0 + 7.0j, 10.0, 5.0 - 5.0j, method='bm') for _ in range(10)]
+        np.testing.assert_allclose(vals, vals2)
+
+        self.rs.set_state(st)
+        vals3 = self.rs.complex_normal(2.0 + 7.0j * np.ones(10), 10.0 * np.ones(1), 5.0 - 5.0j, method='bm')
+        np.testing.assert_allclose(vals, vals3)
+
+    def test_complex_normal_zero_variance(self):
+        st = self.rs.get_state()
+        c = self.rs.complex_normal(0, 1.0, 1.0)
+        assert_almost_equal(c.imag, 0.0)
+        self.rs.set_state(st)
+        n = self.rs.standard_normal()
+        np.testing.assert_allclose(c, n, atol=1e-8)
+
+        st = self.rs.get_state()
+        c = self.rs.complex_normal(0, 1.0, -1.0)
+        assert_almost_equal(c.real, 0.0)
+        self.rs.set_state(st)
+        self.rs.standard_normal()
+        n = self.rs.standard_normal()
+        assert_almost_equal(c.real, 0.0)
+        np.testing.assert_allclose(c.imag, n, atol=1e-8)
 
     def test_exponential(self):
         vals = self.rs.exponential(2.0, 10)