diff --git a/code_to_optimize/discrete_riccati.py b/code_to_optimize/discrete_riccati.py index 314e311b3..25e82daba 100644 --- a/code_to_optimize/discrete_riccati.py +++ b/code_to_optimize/discrete_riccati.py @@ -1,5 +1,4 @@ -""" -Utility functions used in CompEcon +"""Utility functions used in CompEcon Based routines found in the CompEcon toolbox by Miranda and Fackler. @@ -9,13 +8,15 @@ and Finance, MIT Press, 2002. """ + from functools import reduce + import numpy as np +from numba import njit def ckron(*arrays): - """ - Repeatedly applies the np.kron function to an arbitrary number of + """Repeatedly applies the np.kron function to an arbitrary number of input arrays Parameters @@ -42,8 +43,7 @@ def ckron(*arrays): def gridmake(*arrays): - """ - Expands one or more vectors (or matrices) into a matrix where rows span the + """Expands one or more vectors (or matrices) into a matrix where rows span the cartesian product of combinations of the input arrays. Each column of the input arrays will correspond to one column of the output matrix. @@ -78,13 +78,12 @@ def gridmake(*arrays): out = _gridmake2(out, arr) return out - else: - raise NotImplementedError("Come back here") + raise NotImplementedError("Come back here") +@njit(cache=True, fastmath=True) def _gridmake2(x1, x2): - """ - Expands two vectors (or matrices) into a matrix where rows span the + """Expands two vectors (or matrices) into a matrix where rows span the cartesian product of combinations of the input arrays. Each column of the input arrays will correspond to one column of the output matrix. @@ -113,11 +112,23 @@ def _gridmake2(x1, x2): """ if x1.ndim == 1 and x2.ndim == 1: - return np.column_stack([np.tile(x1, x2.shape[0]), - np.repeat(x2, x1.shape[0])]) - elif x1.ndim > 1 and x2.ndim == 1: - first = np.tile(x1, (x2.shape[0], 1)) - second = np.repeat(x2, x1.shape[0]) - return np.column_stack([first, second]) - else: - raise NotImplementedError("Come back here") + n1 = x1.shape[0] + n2 = x2.shape[0] + out = np.empty((n1 * n2, 2), dtype=x1.dtype) + for i in range(n2): + for j in range(n1): + out[i * n1 + j, 0] = x1[j] + out[i * n1 + j, 1] = x2[i] + return out + if x1.ndim > 1 and x2.ndim == 1: + n1 = x1.shape[0] + n2 = x2.shape[0] + ncol1 = x1.shape[1] + out = np.empty((n1 * n2, ncol1 + 1), dtype=x1.dtype) + for i in range(n2): + for j in range(n1): + for k in range(ncol1): + out[i * n1 + j, k] = x1[j, k] + out[i * n1 + j, ncol1] = x2[i] + return out + raise NotImplementedError("Come back here")