ConsIndShockSolverNumba

Over the last few weeks I’ve worked on a numba version of ConsIndShockSolver in IndShockNumbaModel.py. For this, I created two numba helper functions, solveConsIndShockCubicNumba for solving using a cubic approximation for the consumption policy function, and addvFuncNumba for calculating the value function during the solution.

solveConsIndShockCubicNumba

In the previous post, I described the process of linearly interpolating the value function using interp from interpolation.py. For the cubic version, I needed to use a similar cubic interpolator that worked with numba, but using tools from interpolation.py did not seem straightforward. However, I had created a basic numba cubic interpolator for my own purposes previously, and added these functions to HARK.numba. A more thorough overview of what these do is presented here, but below I briefly explain.

The function splrep taxes 1-d vectors x and y, where y=f(x) for some function f, and creates interpolating coefficients using a complete cubic spline method. The function returns a 1-d vector of coefficients, which we can call z. The function splevec then takes 1-d vectors x0, x, y, and z, where y=f(x), z are interpolating coefficients, and x0 contains values where we need an approximation for f, and returns an approximation y0=g(x0), where g is the cubic spline approximation. The final code looks like this:

cFuncCoeffs = splrep(sn_cFunc_x_list, sn_cFunc_y_list)
cFuncNext = splevec(mNrmNext.flatten(), sn_cFunc_x_list, sn_cFunc_y_list, cFuncCoeffs)
vPfuncNext = (cFuncNext ** -CRRA).reshape(mNrmNext.shape)

The same process is used to calculate vPPfuncNext.

addvFuncNumba

This section uses the same process above to calculate a cubic approximation of vfuncNext on current resources. The only novel part of this section is the use of @vectorize on the different kind of utility functions. Instead of importing utilities from HARK.utilities, I have created basic numba implementations of utilities by simply adding @vectorize decorators in HARK.numba. In the future, these vectorized functions could be easily adaptable to GPU processing by adding target="cuda" as a parameter. Instead of creating lambda functions with a set CRRA parameter, however, I use the full call f(c,crra) for all f in utilities, which is used to reduce overhead and prevent numba recompiling. The solver appears to work, although there will be numerical differences given different interpolators that need to be explored further.