Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import numpy as np
import quantecon
from numba import jit
from numba import njit, prange
@njit
def cartesian_2d(x,y,out=None):
p = x.shape[0]
q = y.shape[0]
N = p*q
if out is None:
out = np.zeros((p*q,2))
for i in range(p):
for j in range(q):
ind = i*q+j
out[ind, 0] = x[i]
out[ind, 1] = y[j]
return out
@njit
def cartesian_3d(x,y,z,out=None):
p = x.shape[0]
q = y.shape[0]
r = z.shape[0]
N = p*q*r
if out is None:
out = np.zeros((p*q*r,3))
for i in range(p):
for j in range(q):
for k in range(r):
ind = i*q*r+j*r+k
out[ind, 0] = x[i]
out[ind, 1] = y[j]
out[ind, 2] = z[k]
return out
@njit
def cartesian_2d_ranges(range_1,range_2,out=None):
x_min, x_max, p = range_1
d_x = x_max-x_min
y_min, y_max, q = range_2
d_y = y_max-y_min
if out is None:
out = np.zeros((p*q,2))
for i in range(p):
for j in range(q):
ind = i*q+j
out[ind, 0] = x_min + i/(p-1)*d_x
out[ind, 1] = y_min + j/(q-1)*d_y
return out
@njit
def cartesian_3d_ranges(range_1,range_2,range_3,out=None):
x_min, x_max, p = range_1
d_x = x_max-x_min
y_min, y_max, q = range_2
d_y = y_max-y_min
z_min, z_max, r = range_3
d_z = z_max-z_min
if out is None:
out = np.zeros((p*q*r,3))
for i in range(p):
for j in range(q):
for k in range(r):
ind = i*q*r+j*r+k
out[ind, 0] = x_min + i/(p-1)*d_x
out[ind, 1] = y_min + j/(q-1)*d_y
out[ind, 2] = z_min + k/(r-1)*d_z
return out
def test_1(N=100, d=2, K=1000):
linspaces = [np.linspace(0,10,K) for i in range(d)]
for i in range(N):
p = quantecon.cartesian(linspaces)
import quantecon.ce_util
def test_2(N=100, d=2, K=1000):
linspaces = [np.linspace(0,10,K) for i in range(d)]
for i in range(N):
p = quantecon.ce_util.gridmake(*linspaces)
def test_3(N=100, d=2, K=1000):
linspaces = [np.linspace(0,10,K) for i in range(d)]
if d==2:
meth = cartesian_2d
elif d==3:
meth = cartesian_3d
out = meth(*linspaces)
for i in range(N):
p = meth(*linspaces, out=out)
def test_4(N=100, d=2, K=1000):
linspaces = [(0,10,K) for i in range(d)]
if d==2:
meth = cartesian_2d_ranges
elif d==3:
meth = cartesian_3d_ranges
out = meth(*linspaces)
for i in range(N):
p = meth(*linspaces, out=out)
import time
commands = [
"test_1(N=100, d=2, K=1000) # cartesian",
"test_2(N=100, d=2, K=1000) # gridmake",
"test_3(N=100, d=2, K=1000) # numba",
"test_4(N=100, d=2, K=1000) # numba regular",
"test_1(N=100, d=3, K=100) # cartesian",
"test_2(N=100, d=3, K=100) # gridmake",
"test_3(N=100, d=3, K=100) # numba",
"test_4(N=100, d=3, K=100) # numba regular"
]
for c in commands:
eval(c) # warmup
t1 = time.time()
eval(c)
t2 = time.time()
print("{} : {:.3f} seconds".format(c,t2-t1))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment