shade-econ / sequence-jacobian

A unified framework to solve and analyze heterogeneous-agent macro models.
MIT License
253 stars 149 forks source link

Exogenous Markov chain for the wealth distribution? #22

Open raphaelhuleux opened 6 months ago

raphaelhuleux commented 6 months ago

Hi,

Is it possible to add an exogenous Markov chain for the wealth state variables? I am trying to add a Blanchard-Yaari type of death probability. There is a xi probability every period that households die and restart with 0 wealth.

I have tried to add a Pi_death Markov chain to the household block:

def death_probability(xi, nA):
    Pi_death = np.zeros((nA, nA))
    Pi_death[:,0] = xi
    Pi_death += np.eye(Pi_death.shape[0]) * (1-xi)

    return Pi_death

with the heterogenous block defined as

@het(exogenous=['Pi_death', 'Pi_e'], policy='a', backward='Va', backward_init=household_init) 
def household(Va_p, a_grid, e_grid, rh, w, beta, eis, xi):

    uc_nextgrid = (1-xi) * beta * Va_p
    c_nextgrid = uc_nextgrid ** (-eis)
    coh = (1 + rh) * a_grid[np.newaxis, :] + w * e_grid[:, np.newaxis]
    a = interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid)
    misc.setmin(a, a_grid[0])
    c = coh - a
    Va = (1 + rh) * c ** (-1 / eis)
    return Va, a, c

but I get the following error

"ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 500 is different from 7)"

mrognlie commented 6 months ago

Hi, if you use exogenous=['Pi_death', 'Pi_e'] to have two independent exogenous Markov chains, then the toolkit will expect the arrays Va, Va_p, a, and c to each have three dimensions, with the first two corresponding to the Pi_death and Pi_e Markov chains. You'd need to rewrite the household hetblock and household_init to be consistent with this.

An alternative approach is to use the Kronecker product on your own to consolidate Pi_death and Pi_e into a single one-dimensional Markov chain Pi (and then take a Kronecker product of np.ones(2) and e_grid to get the corresponding grid of endowments over this chain, etc), and then work with that. This approach is often a bit simpler because it doesn't require as many changes to the code, although it requires a bit more manual work with the Markov chain. It is also computationally a bit less efficient than using two separate Markov chains, but I doubt in this case that would be a major issue.

One note about the specific application: I think it makes more sense to have Pi_death be a iid shock. Then, when this shock is received, you can set incoming assets to zero. (Also, if you want to have annuitization a la Blanchard, then I guess you'd want to inflate the gross return on assets by $1/(1-\xi)$. Without this, one needs to take a stand on what happens to the assets at death, e.g. whether they are redistributed somehow.)

Let me know if you have any questions about either of these approaches and if one works for you!

raphaelhuleux commented 6 months ago

Thanks a lot for your help! I wrote a minimal working example, keeping the two exogenous processes and having "Pi_death" an iid shock. One last question: in this case, should I discount the future marginal utility by (1-beta) (1-xi) in the household function (i.e., replace `uc_nextgrid = beta Va_pbyuc_nextgrid = (1-xi)beta Va_p`) or this will be taken into account directly by the package when applying the Pi_death matrix?

import copy
import numpy as np
import matplotlib.pyplot as plt

from sequence_jacobian import het, simple, create_model              # functions
from sequence_jacobian import interpolate, grids, misc, estimation   # modules

def household_init(a_grid, e_grid, rh, w, eis):    
    coh = (1 + rh) * a_grid[np.newaxis, np.newaxis, :] + w * e_grid[:, np.newaxis, np.newaxis] + np.array([0,0])[np.newaxis, :, np.newaxis]
    Va = (1 + rh) * (0.1 * coh) ** (-1 / eis)
    return Va

@het(exogenous=['Pi_e','Pi_death'], policy='a', backward='Va', backward_init=household_init)
def household(Va_p, death_grid, a_grid, e_grid, rh, w, beta, eis, xi):

    uc_nextgrid = beta * Va_p
    c_nextgrid = uc_nextgrid ** (-eis)
    coh = (1 + rh) * death_grid[np.newaxis,:,np.newaxis] * a_grid[np.newaxis, np.newaxis, :] + w * e_grid[:, np.newaxis, np.newaxis]

    a = np.empty_like(coh)

    for i in range(e_grid.size):
        for j in range(death_grid.size):        
            a[i,j,:] = np.interp(coh[i,j,:], c_nextgrid[i,j,:] + a_grid, a_grid)

    a[a<0] = a_grid[0]
    c = coh - a
    Va = (1 + rh) * c ** (-1 / eis)

    return Va, a, c

def make_grid(rho_e, sd_e, nE, amin, amax, nA):
    e_grid, _, Pi_e = grids.markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE)
    a_grid = grids.agrid(amin=amin, amax=amax, n=nA)
    death_grid = np.array([1,0])
    return e_grid, Pi_e, a_grid, death_grid

def death_probability(xi):

    Pi_death = np.array([[1-xi, xi],
                   [1-xi, xi]])

    return Pi_death

household_ext = household.add_hetinputs([make_grid, death_probability])

@simple
def firm(K, L, Z, alpha, delta):
    r = alpha * Z * (K(-1) / L) ** (alpha-1) - delta
    w = (1 - alpha) * Z * (K(-1) / L) ** alpha
    Y = Z * K(-1) ** alpha * L ** (1 - alpha)
    return r, w, Y

@simple 
def death_insurance(r, xi):
    rh = (1+r) * 1/ (1-xi) - 1
    return rh 

@simple
def mkt_clearing(K, A, Y, C, delta, xi):
    asset_mkt = A - K
    goods_mkt = Y - C - delta * K 

    print('K = ', float(K))
    print('A = ', float(A))
    return asset_mkt, goods_mkt

ks = create_model([household_ext, firm, mkt_clearing, death_insurance], name="Krusell-Smith")
print(ks.inputs)

calibration = {'eis': 1, 'delta': 0.025, 'alpha': 0.11, 'rho_e': 0.966, 'sd_e': 0.5, 'L': 1.0,
               'nE': 7, 'nA': 500, 'amin': 0, 'amax': 200, 'xi':0.02}
unknowns_ss = {'beta': 0.98, 'Z': 0.85, 'K': 3.}
targets_ss = {'r': 0.01, 'Y': 1., 'asset_mkt': 0.}
ss = ks.solve_steady_state(calibration, unknowns_ss, targets_ss, solver='hybr')

calibration = ss
unknowns_ss = {'K': 3}
targets_ss = {'asset_mkt': 0.}

print('We can check that the Va when dead is equal, for any level of wealth, to Va when alive and with 0 wealth')

print('Va_p being alive and zero wealth = ', ss.internals['household']['Va'][3,0,0])
print('Va_p being dead = ', ss.internals['household']['Va'][3,1,:])
bbardoczy commented 4 months ago

Hello! uc_nextgrid = beta * Va_p should still be fine. Va_p is the expected partial value function tomorrow. The expectation is taken automatically with respect to every exogenous state. So Pi_death in particular is applied.