Problem with `solve_impulse_nonlinear` with permanent type

Hello,

When solving a model with two types of heterogenous households (in this case, two different levels of permanent income), it seems that the solve_impulse_nonlinear doesn't compute the proper consumption path for the two types of households.

For example, in this simple one-asset Krusell-Smith model, the variable C_low following a TFP shock is not equal to distribution * policy function. It doesn't converge back to zero and is very different from the same variable computed through solve_impulse_linear, whereas the model should be approximately linear for a small shock. I suspect it's an issue with the initial distribution used to compute the transition. I tried imposing manually ss.internals['hh_low']['Dbeg'] = ss.internals['hh_low']['D'], etc, but it didn't solve the issue.

It also seems that the aggregate variables td_nonlin['C'] and td_lin['C'] match and are accurate, so I'm unsure where the issue comes from.

I also noted that there is no issue when setting Z_low = Z_high = 1.

import numpy as np
from sequence_jacobian import simple, create_model
from sequence_jacobian.blocks.het_block import het
from sequence_jacobian import grids, interpolate, misc
import matplotlib.pyplot as plt

def make_grid_one_asset(amax, nA, nE, rho_e, sig_e):

    e_grid, pi_e, Pi = grids.markov_rouwenhorst(rho=rho_e, sigma=sig_e, N=nE)
    a_grid = grids.agrid(amax=amax, n=nA)

    Dbeg = np.zeros((nE, nA))
    Dbeg[:,0] = pi_e

    Dbeg_low = Dbeg
    Dbeg_high = Dbeg

    Pi_low = Pi
    Pi_high = Pi
    return a_grid, e_grid, Pi, Pi_low, Pi_high

def hh_init_oa(a_grid, Z, z_grid, r, sigma):
    coh = (1 + r) * a_grid[np.newaxis,:] + Z * z_grid[:,np.newaxis] 
    Va = (1 + r) * (0.1 * coh) ** (-sigma)
    return Va

@het(exogenous=['Pi'], policy='a', backward='Va', backward_init=hh_init_oa)
def hh_one_asset(Va_p, a_grid, Z, z_grid, r, beta, sigma):

    uc_nextgrid = beta * Va_p 
    c_nextgrid = uc_nextgrid ** (-1/sigma)
    coh = (1 + r) * a_grid[np.newaxis, :] + Z * z_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 + r) * c ** (-sigma)

    return Va, a, c

'''Part 1: Blocks'''

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

@simple
def mkt_clearing(A, K, C, Y, I, verbose, beta):

    asset_mkt = A - K
    goods_mkt = Y - C - I

    if verbose == True:
        print('asset_mkt = ', float(asset_mkt))
        print('A         = ', float(A))
        print('beta      = ', float(beta))
        print(' ')

    return asset_mkt, goods_mkt

@simple  
def aggregate(A_low, A_high, C_low, C_high, share_low, share_high):

    A = A_low * share_low + A_high * share_high 
    C = C_low * share_low + C_high * share_high 
    return A, C 

def income(e_grid, w):
    z_grid = w * e_grid 
    return z_grid

# Combine Blocks
hh_one_asset = hh_one_asset.add_hetinputs([income, make_grid_one_asset])

to_map = ['Z', *hh_one_asset.outputs]
hh_low = hh_one_asset.remap({k: k + '_low' for k in to_map}).rename('hh_low')
hh_high = hh_one_asset.remap({k: k + '_high' for k in to_map}).rename('hh_high')
blocks = [hh_low, hh_high, firm, mkt_clearing, aggregate]
one_asset_model = create_model(blocks, name='One-Asset HANK')

calibration = {'tfp': 1,
               'nE': 11, 'nA': 500, 'sigma':1,
               'amax': 5_000, 'rho_e': 0.91, 'sig_e': 0.9, 
               'alpha':1/3, "delta":0.09,
               'verbose': True, 'share_low':0.9, 'share_high':0.1}

calibration['K'] = ((0.05+calibration['delta']) / (calibration['alpha']))**(1/(calibration['alpha']-1))
calibration['Z_high'] = 1.8
# Normalize total endowment in permanent productivity to 1 
calibration['Z_low'] = (1 - calibration['Z_high'] * calibration['share_high']) / calibration['share_low']

# Steady state

ss =  one_asset_model.solve_steady_state(calibration, {'beta':(0.8,0.95)}, {'asset_mkt':0})
print(ss['goods_mkt']) # Check Walras law
ss['verbose']=False

# Transition
unknowns = ['K']
targets = ['asset_mkt']

T = 300
sig = 0.0005
rho = 0.15
dtfp = sig * rho ** np.arange(T)

td_nonlin = one_asset_model.solve_impulse_nonlinear(ss, unknowns, targets, {"tfp": dtfp}, internals=['hh_low', 'hh_high'])
td_lin = one_asset_model.solve_impulse_linear(ss, unknowns, targets, {"tfp": dtfp})

# This should be the same as td_nonlin['C_low'] : distribution * policy function
dC_low = np.sum(td_nonlin.internals['hh_low']['c']*(ss.internals['hh_low']['D'] + td_nonlin.internals['hh_low']['D']), axis = (1,2)) 
error = np.max(np.abs(dC_low - td_nonlin['C_low']))
print(error)

# Difference between solve_impulse_nonlinear and policy function * distribution
plt.plot(td_nonlin['C_low'], label = 'solve_impulse_nonlinear')
plt.plot(dC_low, linestyle = ':', label = 'policy function * distribution')
plt.legend()
plt.show()

# Difference beteen nonlinear and linear: should be approximately the same for a small shock
plt.plot(td_nonlin['C_low'],label =  'nonlinear')
plt.plot(td_lin['C_low'], linestyle = ':', label = 'linear')
plt.legend()
plt.show()

# Aggregate consumption match for linear and nonlinear
plt.plot(td_nonlin['C'],label =  'nonlinear')
plt.plot(td_lin['C'], linestyle = ':', label = 'linear')
plt.legend()
plt.show()

shade-econ / sequence-jacobian

Problem with `solve_impulse_nonlinear` with permanent type #23