forked from QuantEcon/QuantEcon.py
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathce_util.py
123 lines (97 loc) · 3.07 KB
/
ce_util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
"""
Utility functions used in CompEcon
Based routines found in the CompEcon toolbox by Miranda and Fackler.
References
----------
Miranda, Mario J, and Paul L Fackler. Applied Computational Economics
and Finance, MIT Press, 2002.
"""
from functools import reduce
import numpy as np
def ckron(*arrays):
"""
Repeatedly applies the np.kron function to an arbitrary number of
input arrays
Parameters
----------
*arrays : tuple/list of np.ndarray
Returns
-------
out : np.ndarray
The result of repeated kronecker products.
Notes
-----
Based of original function `ckron` in CompEcon toolbox by Miranda
and Fackler.
References
----------
Miranda, Mario J, and Paul L Fackler. Applied Computational
Economics and Finance, MIT Press, 2002.
"""
return reduce(np.kron, arrays)
def gridmake(*arrays):
"""
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.
Parameters
----------
*arrays : tuple/list of np.ndarray
Tuple/list of vectors to be expanded.
Returns
-------
out : np.ndarray
The cartesian product of combinations of the input arrays.
Notes
-----
Based of original function ``gridmake`` in CompEcon toolbox by
Miranda and Fackler
References
----------
Miranda, Mario J, and Paul L Fackler. Applied Computational Economics
and Finance, MIT Press, 2002.
"""
if all([i.ndim == 1 for i in arrays]):
d = len(arrays)
if d == 2:
out = _gridmake2(*arrays)
else:
out = _gridmake2(arrays[0], arrays[1])
for arr in arrays[2:]:
out = _gridmake2(out, arr)
return out
else:
raise NotImplementedError("Come back here")
def _gridmake2(x1, x2):
"""
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.
Parameters
----------
x1 : np.ndarray
First vector to be expanded.
x2 : np.ndarray
Second vector to be expanded.
Returns
-------
out : np.ndarray
The cartesian product of combinations of the input arrays.
Notes
-----
Based of original function ``gridmake2`` in CompEcon toolbox by
Miranda and Fackler.
References
----------
Miranda, Mario J, and Paul L Fackler. Applied Computational Economics
and Finance, MIT Press, 2002.
"""
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")