/
coleman_egm.py
43 lines (36 loc) · 1.27 KB
/
coleman_egm.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
import numpy as np
def coleman_egm(g, k_grid, β, u_prime, u_prime_inv, f, f_prime, shocks):
"""
The approximate Coleman operator, updated using the endogenous grid
method.
Parameters
----------
g : function
The current guess of the policy function
k_grid : array_like(float, ndim=1)
The set of *exogenous* grid points, for capital k = y - c
β : scalar
The discount factor
u_prime : function
The derivative u'(c) of the utility function
u_prime_inv : function
The inverse of u' (which exists by assumption)
f : function
The production function f(k)
f_prime : function
The derivative f'(k)
shocks : numpy array
An array of draws from the shock, for Monte Carlo integration (to
compute expectations).
"""
# Allocate memory for value of consumption on endogenous grid points
c = np.empty_like(k_grid)
# Solve for updated consumption value
for i, k in enumerate(k_grid):
vals = u_prime(g(f(k) * shocks)) * f_prime(k) * shocks
c[i] = u_prime_inv(β * np.mean(vals))
# Determine endogenous grid
y = k_grid + c # y_i = k_i + c_i
# Update policy function and return
Kg = lambda x: np.interp(x, y, c)
return Kg