ENH: Complete complex normal refactor

Refactor broadcasted generation

Author: bashtage

randomstate/randomstate.pyx

@@ -164,6 +164,19 @@ cdef double kahan_sum(double *darr, np.npy_intp n):
         sum = t
     return sum
 
+cdef inline void compute_complex(double *rv_r, double *rv_i, double loc_r,
+                                 double loc_i, double var_r, double var_i, double rho) nogil:
+    cdef double scale_c, scale_i, scale_r
+
+    scale_c = sqrt(1 - rho * rho)
+    scale_r = sqrt(var_r)
+    scale_i = sqrt(var_i)
+
+    rv_i[0] = loc_i + scale_i * (rho * rv_r[0]  + scale_c * rv_i[0])
+    rv_r[0] = loc_r + scale_r * rv_r[0]
+
+
+
 cdef object _ensure_string(object s):
     try:
         return ''.join(map(chr, s))
@@ -1782,11 +1795,13 @@ cdef class RandomState:
 
         >>> s = np.random.complex_normal(size=1000)
         """
-        cdef np.ndarray ogamma, orelation, oloc, randoms
+        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
-        cdef np.npy_intp i, j, n,
+            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
 
         oloc = <np.ndarray>np.PyArray_FROM_OTF(loc, np.NPY_COMPLEX128, np.NPY_ALIGNED)
         ogamma = <np.ndarray>np.PyArray_FROM_OTF(gamma, np.NPY_COMPLEX128, np.NPY_ALIGNED)
@@ -1832,7 +1847,7 @@ 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:
+            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)
@@ -1844,10 +1859,10 @@ cdef class RandomState:
 
         gpc = ogamma + orelation
         gmc = ogamma - orelation
-        v_real = 0.5 * np.real(gpc)
+        v_real = <np.ndarray>(0.5 * np.real(gpc))
         if np.any(np.less(v_real, 0)):
             raise ValueError('Re(gamma + relation) < 0')
-        v_imag = 0.5 * np.real(gmc)
+        v_imag = <np.ndarray>(0.5 * np.real(gmc))
         if np.any(np.less(v_imag, 0)):
             raise ValueError('Re(gamma - relation) < 0')
         if np.any(np.not_equal(np.imag(ogamma), 0)):
@@ -1860,23 +1875,32 @@ cdef class RandomState:
         if np.any(cov.flat[~idx] != 0) or np.any(np.abs(rho) > 1):
             raise ValueError('Im(relation) ** 2 > Re(gamma ** 2 - relation ** 2)')
 
-        if size is None:
-            size = np.broadcast(loc, gpc).shape
-        elif np.PyArray_IsAnyScalar(size):
-            size = (size,)
+        if size is not None:
+            randoms = <np.ndarray>np.empty(size, np.complex128)
         else:
-            size = tuple(size)
+            it = np.PyArray_MultiIterNew4(oloc, v_real, v_imag, rho)
+            randoms = <np.ndarray>np.empty(it.shape, np.complex128)
 
-        norms = self.standard_normal(size + (2,), method=method)
-        real = norms[...,0]
-        imag = norms[...,1]
+        randoms_data = <double *>np.PyArray_DATA(randoms)
+        n = np.PyArray_SIZE(randoms)
 
-        imag *= np.sqrt(1-rho ** 2)
-        imag += rho * real
-        real *= np.sqrt(v_real)
-        imag *= np.sqrt(v_imag)
-        
-        return oloc + real + (0+1.0j) * imag
+        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)
+        with nogil:
+            j = 0
+            for i in range(n):
+                floc_r= (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
+                floc_i= (<double*>np.PyArray_MultiIter_DATA(it, 1))[1]
+                f_v_real = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0]
+                f_v_imag = (<double*>np.PyArray_MultiIter_DATA(it, 3))[0]
+                f_rho = (<double*>np.PyArray_MultiIter_DATA(it, 4))[0]
+                compute_complex(&randoms_data[j], &randoms_data[j+1], floc_r, floc_i, f_v_real, f_v_imag, f_rho)
+                j += 2
+                np.PyArray_MultiIter_NEXT(it)
+
+        return randoms
 
     def beta(self, a, b, size=None):
         """