Added cartesian product.

Author: albop

quantecon/cartesian.py

@@ -0,0 +1,99 @@
+"""
+Filename: cartesian.py
+
+Authors: Pablo Winant
+
+Implements cartesian products and regular cartesian grids.
+"""
+
+import numpy
+from numba import njit
+
+def cartesian(nodes, order='C'):
+    '''Cartesian product of a list of arrays
+
+    Parameters:
+    -----------
+    nodes: (list of 1d-arrays)
+    order: ('C' or 'F') order in which the product is enumerated
+
+    Returns:
+    --------
+    out: (2d-array) each line corresponds to one point of the product space
+    '''
+
+    nodes = [numpy.array(e) for e in nodes]
+    shapes = [e.shape[0] for e in nodes]
+
+    n = len(nodes)
+    l = numpy.prod(shapes)
+    out = numpy.zeros((l, n))
+
+
+    if order == 'C':
+        repetitions = numpy.cumprod([1] + shapes[:-1])
+    else:
+        shapes.reverse()
+        sh = [1] + shapes[:-1]
+        repetitions = numpy.cumprod(sh)
+        repetitions = repetitions.tolist()
+        repetitions.reverse()
+
+    for i in range(n):
+        _repeat_1d(nodes[i], repetitions[i], out[:,i])
+
+    return out
+
+def mlinspace(a, b, nums, order='C'):
+    '''Constructs a regular cartesian grid
+
+    Parameters:
+    -----------
+    a: (1d-array) lower bounds in each dimension
+    b: (1d-array) upper bounds in each dimension
+    nums: (1d-array) number of nodes along each dimension
+    order: ('C' or 'F') order in which the product is enumerated
+
+    Returns:
+    --------
+    out: (2d-array) each line corresponds to one point of the product space
+    '''
+
+    a = numpy.array(a, dtype='float64')
+    b = numpy.array(a, dtype='float64')
+    nums = numpy.array(a, dtype='int64')
+    nodes = [numpy.linspace(a[i], b[i], nums[i]) for i in orders]
+
+    return cartesian(nodes, order=order)
+
+
+@njit
+def _repeat_1d(x, K, out):
+    '''Repeats each element of a vector many times and repeats the whole result many times
+
+    Parameters
+    ----------
+
+    x: (1d array)       vector to be repeated
+    K: (int)            number of times each element of x is repeated (inner iterations)
+    out: (1d array)     placeholder for the result
+
+    Returns
+    -------
+    None
+    '''
+
+    N = x.shape[0]
+    L = out.shape[0] // (K*N) # number of outer iterations
+    # K                       # number of inner iterations
+
+    # the result out should enumerate in C-order the elements
+    # of a 3-dimensional array T of dimensions (K,N,L)
+    # such that for all k,n,l, we have T[k,n,l] == x[n]
+
+    for n in range(N):
+        val = x[n]
+        for k in range(K):
+            for l in range(L):
+                ind = k*N*L + n*L + l
+                out[ind] = val

quantecon/tests/test_cartesian.py

@@ -0,0 +1,166 @@
+"""
+Author: Pablo Winant
+Filename: test_cartesian.py
+
+Tests for cartesian.py file
+
+"""
+
+from quantecon.cartesian import cartesian, _repeat_1d
+
+def test_cartesian_C_order():
+
+    from numpy import linspace
+    x = linspace(0,9,10)
+
+    prod = cartesian([x,x,x])
+
+    correct = True
+    for i in range(999):
+        n = prod[i,0]*100+prod[i,1]*10+prod[i,2]
+        correct *= (i == n)
+
+    assert(correct)
+
+def test_cartesian_F_order():
+
+    from numpy import linspace
+    x = linspace(0,9,10)
+
+    prod = cartesian([x,x,x], order='F')
+
+    correct = True
+    for i in range(999):
+        n = prod[i,2]*100+prod[i,1]*10+prod[i,0]
+        correct *= (i == n)
+
+    assert(correct)
+
+def test_performance_C():
+
+    from numpy import linspace, column_stack, repeat, tile
+    import time
+
+    N_x = 1000
+    N_y = 7777
+    x = linspace(1,N_x,N_x)
+    y = linspace(1,N_y,N_y)
+
+    cartesian([x[:10],y[:10]]) # warmup
+
+    t1 = time.time()
+    for i in range(100):
+        prod = cartesian([x,y])
+    t2 = time.time()
+    # print(prod.shape)
+
+    # compute the same produce using numpy:
+    import numpy
+
+    t3 = time.time()
+    for i in range(100):
+        prod_numpy = column_stack([
+            repeat(x,N_y),
+            tile(y,N_x)
+        ])
+    t4 = time.time()
+
+    print("Timings for 'cartesian' (C order)")
+    print("Cartesian: {}".format(t2-t1))
+    print("Numpy:     {}".format(t4-t3))
+    assert(abs(prod-prod_numpy).max()==0)
+
+def test_performance_F():
+
+    from numpy import linspace, column_stack, repeat, tile
+    import time
+
+    N_x = 1000
+    N_y = 7777
+    x = linspace(1,N_x,N_x)
+    y = linspace(1,N_y,N_y)
+
+    cartesian([x[:10],y[:10]]) # warmup
+
+    t1 = time.time()
+    for i in range(100):
+        prod = cartesian([x,y], order='F')
+    t2 = time.time()
+    # print(prod.shape)
+
+    # compute the same produce using numpy:
+    import numpy
+
+    t3 = time.time()
+    for i in range(100):
+        prod_numpy = column_stack([
+            tile(x,N_y),
+            repeat(y,N_x)
+        ])
+    t4 = time.time()
+
+    print("Timings for 'cartesian'(Fortran order)")
+    print("Cartesian: {}".format(t2-t1))
+    print("Numpy:     {}".format(t4-t3))
+    assert(abs(prod-prod_numpy).max()==0)
+
+def test_mlinsplace():
+
+    from numpy import linspace
+    from quantecon.cartesian import mlinspace
+
+    grid1 = mlinspace([-1,-1],[2,3],[30,50])
+    grid2 = cartesian([linspace(-1,2,30), linspace(-1,3,50)])
+
+def test_tile():
+
+    from numpy import linspace, tile, zeros
+    x = linspace(1,100, 100)
+
+    import time
+    t1 = time.time()
+    t_repeat = zeros(100*1000)
+    _repeat_1d(x,1,t_repeat)
+    t2 = time.time()
+
+    t3 = time.time()
+    t_numpy = tile(x, 1000)
+    t4 = time.time()
+
+    print("Timings for 'tile' operation")
+    print("Repeat_1d: {}".format(t2-t1))
+    print("Numpy:     {}".format(t4-t3))
+
+    assert( abs(t_numpy-t_repeat).max())
+
+def test_repeat():
+
+    from numpy import linspace, repeat, zeros
+    x = linspace(1,100  , 100)
+
+    import time
+    t1 = time.time()
+    t_repeat = zeros(100*1000)
+    _repeat_1d(x,1000,t_repeat)
+    t2 = time.time()
+
+    t3 = time.time()
+    t_numpy = repeat(x, 1000)
+    t4 = time.time()
+
+    print("Timings for 'repeat' operation")
+    print("Repeat_1d: {}".format(t2-t1))
+    print("Numpy:     {}".format(t4-t3))
+
+    assert( abs(t_numpy-t_repeat).max())
+
+
+
+
+if __name__ == '__main__':
+    test_cartesian_C_order()
+    test_cartesian_F_order()
+    test_performance_C()
+    test_performance_F()
+    test_tile()
+    test_repeat()