Greetings.

I am working on a model which features competitive search and heterogeneous agents. Both unemployed and employed workers choose assets, but unemployed workers also choose the wage submarket in which to search. To capture both cases, the individual state of a worker is (a, w, e) for assets a, wage w, and employment status e, and the policy choice is (a’, w’).

The issue is that if I formulate the problem as a discrete DP, the Q matrix is very large and time-costly to compute. In this case, I use the state-action pairs formulation of the DDP with a sparse format. Specifically, I know that np.unravel_index can be replaced with the integer division operators to make compatible with numba, but how to best make the use of a sparse matrix Q compatible with numba. Alternatively, is there some other way to accelerate the computation of Q?

def populate_Q(N, r, params, grid):

beta, s, eps, phi, psi, nu, sigma, b, k, Z= params

T, burn_in, a_min, a_max, a, h, ngrid, wgrid, nstates_id = grid

util, util_c, util_h, jf, vf, theta_invert, theta_invert_fe = gen_fun(params)

n = ngrid*nstates_id*wgrid

policy = ngrid*wgrid
w = gen_w(N, Z, h, sigma, wgrid)
f = jf(theta_invert_fe(w, h, N))
L = policy*n

Q = sparse.lil_matrix((L, n))

“Loop over states”

for s_i in range(n):

a_i, e_i, w_i = np.unravel_index(s_i, (ngrid, nstates_id, wgrid))

“Loop over possible action and wage submarket choices”

for act in range(policy):

new_a_i, new_w_i = np.unravel_index(act, (ngrid, wgrid))

#adjust new wage of employed to be old wage

new_w_i = e_i*w_i + (1-e_i)*new_w_i

for new_e_i in range(nstates_id):

“index future states”

s_i_new = np.ravel_multi_index((new_a_i, new_e_i, new_w_i),

(ngrid, nstates_id, wgrid))

l = np.ravel_multi_index((s_i, act), (n, policy))

Q[l, s_i_new] = P_mat(nstates_id, f[new_w_i], s)[e_i, new_e_i]

return Q