Value Function Iteration Using OpenCL

CL_RBC.py

# Basic RBC model with full depreciation
# U(c) = log(c)
# F(K) = Z k^\alpha
# where Z is a finite-state Markov chain
#
# Original version:
# Jesus Fernandez-Villaverde
# Haverford, July 3, 2013
# https://github.com/jesusfv/Comparison-Programming-Languages-Economics/blob/master/RBC_Python.py
#
# Porting to GPU using OpenCL (at the cost of the clever optimization which uses the monotonicity
# of the policy function):w
# Jyotirmoy Bhattacharya, 2016

# 0. Initialization
import numpy as np
import numpy.linalg
import math
import time
import pyopencl as cl

t1=time.time()

def main_func():

    aalpha = 1.0/3.0     # Elasticity of output w.r.t. capital
    bbeta  = 0.95        # Discount factor

    # Productivity values
    vProductivity = np.array([0.9792, 0.9896, 1.0000, 1.0106, 1.0212],np.float64)

    # Transition matrix
    mTransition   = np.array([[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]],np.float64)

    ## 2. Steady State

    capitalSteadyState     = (aalpha*bbeta)**(1/(1-aalpha))
    outputSteadyState      = capitalSteadyState**aalpha
    consumptionSteadyState = outputSteadyState-capitalSteadyState

    print("Output = ", outputSteadyState, " Capital = ", capitalSteadyState, " Consumption = ", consumptionSteadyState)

    nGridCapital           = 16384
    nGridProductivity      = len(vProductivity)

    ## 3. Required matrices and vectors

    mValueFunction    = np.zeros((nGridProductivity,nGridCapital),dtype=np.float64)
    mValueFunctionNew = np.empty((nGridProductivity,nGridCapital),dtype=np.float64)
    mPolicyFunction   = np.empty((nGridProductivity,nGridCapital),dtype=np.float64)
    expectedValueFunction = np.empty((nGridProductivity,nGridCapital),dtype=np.float64)

    ## 4. Main iteration

    maxDifference = 10.0
    tolerance = 0.0000001
    iteration = 0

    ## 4.1 OpenCL Kernels
    ## (The double braces {{ and }} are to escape it
    ##   from the Python string formatting which we use for
    ##   filling in constants)
    solve_max_k = """
        /* The capital grid*/
        double capital_at(int k){{
            return 0.5*{capitalSteadyState}+
                    ({capitalSteadyState}*(double)k)/({nGridCapital}-1);
        }}

      /* Given a value function, compute the expectation of the 
         value function given today's state and tomorrow's capital*/
     __kernel void compute_evf(
            __global const double *valueFunction,
            __constant const double *transMat,
            __global double *evf)
    {{
            int nProductivity = get_global_id(0);
            int nCapital = get_global_id(1);
            int prodNext;
            double res=0;

            for (prodNext=0;prodNext<{nGridProductivity};prodNext++)
                res += transMat[nProductivity*{nGridProductivity}+prodNext]
                       *valueFunction[prodNext*{nGridCapital}+nCapital];
            evf[nProductivity*{nGridCapital}+nCapital]=res;
    }}

    /* Value function iteration. Given a value function find the
       optimal choice and the resulting next iteration of the value
       function */
     __kernel void solve_max(
          __constant const double *vProductivity,
          __global const double *expectedValueFunction,
          __global double *mValueFunctionNew,
          __global double *mPolicyFunction)
      {{
        int nProductivity = get_global_id(0);
        int nCapital = get_global_id(1);

        int knext;
        double maxSoFar = -1.0/0;
        int capitalChoice = 0;
        double output = vProductivity[nProductivity]*pow(capital_at(nCapital),{aalpha});

        for (knext=0; knext<{nGridCapital};knext++){{
          double consumption = output - capital_at(knext);

          double valueProvisional = (1-{bbeta})*log(consumption)
                                    + {bbeta}*expectedValueFunction[nProductivity*{nGridCapital}+knext];

          if (valueProvisional>maxSoFar){{
            maxSoFar = valueProvisional;
            capitalChoice = knext;
          }} else {{
            break;
          }}
        }}
        mValueFunctionNew[nProductivity*{nGridCapital}+nCapital] = maxSoFar;
        mPolicyFunction[nProductivity*{nGridCapital}+nCapital] = capital_at(capitalChoice);
      }}

      """.format(nGridCapital=nGridCapital,nGridProductivity=nGridProductivity,bbeta=bbeta,aalpha=aalpha,capitalSteadyState=capitalSteadyState)

    ## 4.2 Initialize OpenCL and allocate memory on the device
    ctx = cl.create_some_context(interactive=False)
    queue = cl.CommandQueue(ctx)
    mf = cl.mem_flags
    prg = cl.Program(ctx,solve_max_k).build()
    vProductivity_g = cl.Buffer(ctx,mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=vProductivity)
    mTransition_g = cl.Buffer(ctx,mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=mTransition)
    mValueFunctionNew_g = cl.Buffer(ctx, mf.READ_WRITE|mf.COPY_HOST_PTR, hostbuf=mValueFunctionNew)
    mPolicyFunction_g = cl.Buffer(ctx, mf.WRITE_ONLY, mPolicyFunction.nbytes)
    expectedValueFunction_g = cl.Buffer(ctx, mf.READ_WRITE, expectedValueFunction.nbytes)

    ## Value function iteration loop
    while(maxDifference > tolerance):
        prg.compute_evf(queue,expectedValueFunction.shape,(1,64),
                        mValueFunctionNew_g,
                        mTransition_g,
                        expectedValueFunction_g).wait()
        prg.solve_max(queue, expectedValueFunction.shape, (1,64), 
                      vProductivity_g,expectedValueFunction_g,
                      mValueFunctionNew_g,mPolicyFunction_g).wait()
        cl.enqueue_copy(queue, mValueFunctionNew, mValueFunctionNew_g)
        maxDifference = (abs(mValueFunctionNew-mValueFunction)).max()

        mValueFunction[:]    = mValueFunctionNew

        iteration += 1
        if(iteration%10 == 0 or iteration == 1):
            print(" Iteration = ", iteration, ", Sup Diff = ", maxDifference)

    cl.enqueue_copy(queue, mPolicyFunction,mPolicyFunction_g)
    return (maxDifference, iteration, mValueFunction, mPolicyFunction)

maxDiff, iter, mValueF, mPolicyFunction = main_func()

print(" Iteration = ", iter, ", Sup Duff = ", maxDiff)
print(" ")
print(" My Check = ", mPolicyFunction[2,999])
print(" ")

t2 = time.time()

print("Elapsed time = is ", t2-t1)