1 files changed, 52 insertions, 45 deletions
diff --git a/numpy/random/examples/numba/extending.py b/numpy/random/examples/numba/extending.py
index d41c2d76f..0d240596b 100644
--- a/numpy/random/examples/numba/extending.py
+++ b/numpy/random/examples/numba/extending.py
@@ -1,14 +1,57 @@
-import datetime as dt
-
 import numpy as np
 import numba as nb
 
 from numpy.random import PCG64
+from timeit import timeit
+
+bit_gen = PCG64()
+next_d = bit_gen.cffi.next_double
+state_addr = bit_gen.cffi.state_address
+
+def normals(n, state):
+    out = np.empty(n)
+    for i in range((n + 1) // 2):
+        x1 = 2.0 * next_d(state) - 1.0
+        x2 = 2.0 * next_d(state) - 1.0
+        r2 = x1 * x1 + x2 * x2
+        while r2 >= 1.0 or r2 == 0.0:
+            x1 = 2.0 * next_d(state) - 1.0
+            x2 = 2.0 * next_d(state) - 1.0
+            r2 = x1 * x1 + x2 * x2
+        f = np.sqrt(-2.0 * np.log(r2) / r2)
+        out[2 * i] = f * x1
+        if 2 * i + 1 < n:
+            out[2 * i + 1] = f * x2
+    return out
+
+# Compile using Numba
+normalsj = nb.jit(normals, nopython=True)
+# Must use state address not state with numba
+n = 10000
+
+def numbacall():
+    return normalsj(n, state_addr)
+
+rg = np.random.Generator(PCG64())
 
-x = PCG64()
-f = x.ctypes.next_uint32
-s = x.ctypes.state
+def numpycall():
+    return rg.normal(size=n)
 
+# Check that the functions work
+r1 = numbacall()
+r2 = numpycall()
+assert r1.shape == (n,)
+assert r1.shape == r2.shape
+
+t1 = timeit(numbacall, number=1000)
+print('{:.2f} secs for {} PCG64 (Numba/PCG64) gaussian randoms'.format(t1, n))
+t2 = timeit(numpycall, number=1000)
+print('{:.2f} secs for {} PCG64 (NumPy/PCG64) gaussian randoms'.format(t2, n))
+
+# example 2
+
+next_u32 = bit_gen.ctypes.next_uint32
+ctypes_state = bit_gen.ctypes.state
 
 @nb.jit(nopython=True)
 def bounded_uint(lb, ub, state):
@@ -19,14 +62,14 @@ def bounded_uint(lb, ub, state):
     mask |= mask >> 8
     mask |= mask >> 16
 
-    val = f(state) & mask
+    val = next_u32(state) & mask
     while val > delta:
-        val = f(state) & mask
+        val = next_u32(state) & mask
 
     return lb + val
 
 
-print(bounded_uint(323, 2394691, s.value))
+print(bounded_uint(323, 2394691, ctypes_state.value))
 
 
 @nb.jit(nopython=True)
@@ -36,42 +79,6 @@ def bounded_uints(lb, ub, n, state):
         out[i] = bounded_uint(lb, ub, state)
 
 
-bounded_uints(323, 2394691, 10000000, s.value)
-
-g = x.cffi.next_double
-cffi_state = x.cffi.state
-state_addr = x.cffi.state_address
-
-
-def normals(n, state):
-    out = np.empty(n)
-    for i in range((n + 1) // 2):
-        x1 = 2.0 * g(state) - 1.0
-        x2 = 2.0 * g(state) - 1.0
-        r2 = x1 * x1 + x2 * x2
-        while r2 >= 1.0 or r2 == 0.0:
-            x1 = 2.0 * g(state) - 1.0
-            x2 = 2.0 * g(state) - 1.0
-            r2 = x1 * x1 + x2 * x2
-        f = np.sqrt(-2.0 * np.log(r2) / r2)
-        out[2 * i] = f * x1
-        if 2 * i + 1 < n:
-            out[2 * i + 1] = f * x2
-    return out
+bounded_uints(323, 2394691, 10000000, ctypes_state.value)
 
 
-print(normals(10, cffi_state).var())
-# Warm up
-normalsj = nb.jit(normals, nopython=True)
-normalsj(1, state_addr)
-
-start = dt.datetime.now()
-normalsj(1000000, state_addr)
-ms = 1000 * (dt.datetime.now() - start).total_seconds()
-print('1,000,000 Polar-transform (numba/PCG64) randoms in '
-      '{ms:0.1f}ms'.format(ms=ms))
-
-start = dt.datetime.now()
-np.random.standard_normal(1000000)
-ms = 1000 * (dt.datetime.now() - start).total_seconds()
-print('1,000,000 Polar-transform (NumPy) randoms in {ms:0.1f}ms'.format(ms=ms))