More efficient way to compute the Q matrix for discrete dynamic programming


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 = ngridnstates_idwgrid
policy = ngridwgrid
w = gen_w(N, Z, h, sigma, wgrid)
f = jf(theta_invert_fe(w, h, N))
L = policy
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

I would certainly be looking to use numba, and, once the function is jitted, I might replace for s_i in range(n) with for s_i in prange(n) and see if that makes a difference.

I’m afraid I don’t have much experience with sparse matrices or time to dig deeper.

Thanks! I will see what I can do.