% dpexample.m

clear
close all

% Stochastic dynamic programming
% Basic example of value function iteration

% 1. Setting the grid for capital and related information

% Determine stationary capital stock
% from the condition f'(k)=rr+dep;

rr=.06;
dep=.14;
a=1;

% production function is y=ak^gam

gam=.333

kbar=((a*gam)/(rr+dep))^(1/(1-gam));


b=1/(1+rr);

% capital grid

Nk=200;

index=1:Nk;

upper=1.1;
lower=.1;

kgrid=kbar*lower+kbar*(upper-lower)*(index/Nk);

% 2. Setting shock process

% support of distribution
% very wide for graphical purposes

avec=[.75 1 1.25]';

Na=length(avec);

% probability of transitting from state i (row)
% to state j (column)

transmat=[.5 .4 .1
          .2 .6 .2
          .1 .4 .5];

% Gross output at each capital stock level
% and productivity level.

ygrid=avec*kgrid.^gam + ones(size(avec))*(1-dep)*kgrid;

% y grid is an Na by Nm matrix holding gross outputs in state.

% 3. utility function curvature
% (and associated plotting parameters)
uopt=1;

if uopt==1;
    
theta=3;
ymin=-20;
ymax=0;

else
    
theta=.9;
ymin=0;
ymax=200;

end

% Initial value function at each capital stock level.
% rows of value function show v(k) for a fixed a.

v=(ygrid.^(1-theta))/(1-theta);;

% change in value function

delv=v-0;

% Initial policy (no capital accumulation)

kpol=zeros(Na,Nk);

% Matrix for holding sequence of value functions
V=v';
it=0;
itmax=20;

% 4. Iterative scheme

% continue iteration so long as change in v is small and number o
% iterations is less than maximum

while abs(max(delv))>.0001 & it<itmax;    
    
% At each point in the k grid


vnew=[];
kpolnew=[];

for i=1:Nk;


    cons=ygrid(:,i)*ones(size(kgrid))-ones(size(avec))*kgrid;
    % This is now consumption in various a states (i=row)
    % and various future k's (column) column
    % require that consumption is positive 
    testp=cons>0;
    cons=testp.*cons+(1-testp)*(.0000001);
    utility=(cons.^(1-theta))/(1-theta);
    % Calculate maximum value and associated action
    % In matlab, [mx,pt]=max(M) takes the maximum element from 
    % each colum (mx) and identifies its location (pt). 
    % To max over various k', we therefore must transpose the matrix.
    
    utranspose=(utility+b*transmat*v)';
    [vnewi poli]=max(utranspose);
    % optimal action in each a state: an Na by 1 vector
    kpolnewi=kgrid(poli)';
    % Collect the results for the policy and value function
    vnew=[vnew vnewi'];
    kpolnew=[kpolnew kpolnewi];    
end

subplot(211),plot(kgrid/kbar,v(2,:),'r-',kgrid/kbar,vnew,'k-')
title(['Value Function Updating at Iteration ',num2str(it)])
legend('v (a=1)','vnew','Location','Southeast')
xlabel('Capital (fraction of ss)')
ylabel('Value')
axis([lower upper ymin ymax])

subplot(212),plot(kgrid/kbar,kpol(2,:)/kbar,'r-',kgrid/kbar,kpolnew/kbar,'k-')
title(['Policy Function Updating at Iteration ',num2str(it)])
legend('pol (a=1)','polnew','Location','Northwest')
xlabel('Capital (fraction of ss)')
ylabel('Kprime')

axis([lower upper lower upper])

if (it==1)
    pause
else
    pause(.1)
end

kpol=kpolnew;
v=vnew;
it=it+1;

end


