3dplot_cpolicy

import numpy as np
from scipy import interp
from quantecon.models import ConsumerProblem
from quantecon import mc_sample_path
from quantecon import compute_fixed_point
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def compute_longrun_series(cp, T=500000, verbose=False):
    """
    Simulates a time series of length T for assets and consumption
    given optimal savings behavior. cp is the default instance of ConsumerProblem
    """
    Pi, z_vals, R = cp.Pi, cp.z_vals, cp.R
    v_init, c_init = cp.initialize()
    c = compute_fixed_point(cp.coleman_operator, c_init, verbose=verbose)
    cf = lambda a, i_z: interp(a, cp.asset_grid, c[:, i_z])
    
    a = np.zeros(T+1)
    longr_c = np.zeros(T)
    z_seq = mc_sample_path(Pi, sample_size=T)
    for t in range(T):
        i_z = z_seq[t]
        a[t+1] = R * a[t] + z_vals[i_z] - cf(a[t], i_z)
        longr_c[t] = cf(a[t], i_z)
    # Returns longrun assets, optimal consumption policy and longrun consumption
    return a, c, longr_c

cp = ConsumerProblem(r=0.03, grid_max=4)
a, c, longr_c = compute_asset_series(cp)

# Easier Task: Plot Optimal Consumption Policy
def f(X, Y):
    # 0 is the borrowing constraint, 4 is gridmax, 50 is gridsize
    asset_grid = np.linspace(0, 4, 50)
    
    for i in range(len(X[0,:])):
        for j in range(len(Y[:,0])):
            cf = lambda a, i_z: interp(a, asset_grid, c[:, i_z])
            z = cf(X,Y)
    return z

fig = plt.figure(figsize=(7, 7))
ax = fig.gca(projection='3d')

# This is range of x (set to 4 since gridmax=4)
x = np.arange(0, 4)
# This is range of y (set to 2 since there were two shocks)
y = np.arange(0, 2)

X,Y = np.meshgrid(x, y)
Z = np.vectorize(f(X,Y))

surf = ax.plot_surface(X, Y, Z, linewidth=0, antialiased=False)
ax.set_zlim('tight')
plt.show()