Hello,
I have a dynamic programming problem in which I need to calculate the conditional expectation of the next period value function of my problem.
The expectation is conditional on my current states and action, and is calculated integrating over the next-period states. I assume that the states follow a first order Markov process.
The problem itself is very simple but it can blow up really quick due to the combinatorial nature of the problem.
I have two actions, and four discrete states.
using Random, Distributions, BenchmarkTools
# dimensions
da = 2 # actions
dk = 5 # capital
dx = 10
dz = 15
du = 15
I generate dense transition matrices. Note that one of the transition matrices depends on the action.
Where F_k[a][i,j] = \Pr(k^\prime=j|k=i,a) and the same goes for other states.
# transition matrices
Fk = [rand(dk,dk) for a=1:da]
Fx = rand(dx,dx)
Fz = rand(dz,dz)
Fu = rand(du,du)
# normalization
Fk = [Fk[a]./= sum(Fk[a],dims=2) for a=1:da]
Fx ./= sum(Fx,dims=2)
Fz ./= sum(Fz,dims=2)
Fu ./= sum(Fu,dims=2)
I generate some random values for the value functions.
# value functions
V = rand(dk,dx,dz,du)
# expectation of V conditional on current period action and states
EV = zeros(da,dk,dx,dz,du)
Function that calculates the expectation over future states.
function expectation(a,k,x,z,u,V,Fk,Fx,Fz,Fu)
FV = 0.0
for ik=1:dk
for ix=1:dx
for iz=1:dz
for iu=1:du
FV += Fk[a][k,ik] * Fx[x,ix] * Fz[z,iz] * Fu[u,iu] * V[ik,ix,iz,iu]
end
end
end
end
return FV
end
I calculate the conditional expectations as follows.
# the expectation takes very long
@btime for ia=1:da
for ik=1:dk
for ix=1:dx
for iz=1:dz
for iu=1:du
EV[ia,ik,ix,iz,iu] += expectation(ia,ik,ix,iz,iu,V,Fk,Fx,Fz,Fu)
end
end
end
end
end
It’s clearly a very inefficient way to get these values. The memory use is high.
78.692 s (2061430725 allocations: 35.03 GiB)
I am sure that the performance tips for the for loops will help a lot, but I wonder if there is a more intelligent and elegant way to do access the data and perform the multiplications given the dimensionality of the problem.
Feel free to rearrange the problem (transition matrices, EV array, etc.) as you wish.
Thank you all.