compare_products.py

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))