%%%%% Discrete-state value function iteration: An example

clear all

tic

%%%%% Parameter values

alpha = 0.33; 
beta  = 0.50;
delta = 1.00;
sigma = 2.00;

%%%%% Solve for steady state

kbar  = ((alpha*beta)/(1-beta+delta*beta))^(1/(1-alpha));

%%%%% Set up discrete state space

NK    = 1000;
Kgrid = linspace(0,5*kbar,NK)';

%%%%% Set up return matrix

R  = zeros(NK,NK);
C  = zeros(NK,NK);

for i = 1:NK,
    ki = Kgrid(i);
    for h = 1:NK,
        kh = Kgrid(h);
        
        c  = ki^alpha+(1-delta)*ki-kh;
        C(i,h) = max(0,c);                %% Imposes consumption non-negative
    end
end

R  = (1/(1-sigma))*(C.^(1-sigma));        %% Return matrix 

%%%%% Begin value function iteration

iter  = 0;
error = 100;

V0    = zeros(NK,1);

while error>10^-6 & iter<300

TV0   = max((R+beta*ones(NK,1)*V0')');    %% Bellman operation
TV0   = TV0';

error = max(abs(V0-TV0))                  %% Check if converged

V0    = TV0;                              %% If not, update 

iter  = iter+1

c0 = (Kgrid.^alpha)+(1-delta)*Kgrid-0;    %% Value if don't accumulate any capital
U0 = (1/(1-sigma))*(c0.^(1-sigma));       %% Use this to check calculations

end

V    = TV0;
time = toc/60;

if error<10^6 & iter<300
display('Congratulations, you found a fixed point!')
display(['Time taken = ' ,num2str(time),'   minutes'])
end

figure(1)
plot(Kgrid,V)
legend('V(k)',2)
title('Value function')
xlabel('Capital stock')
ylabel('Lifetime utility')

%%%%% Compute policy function
gs = zeros(NK,1);
[vs,is] = max((R+beta*ones(NK,1)*V')');

for i = 1:NK,

gs(i)  = Kgrid(is(i));

end

figure(2)
plot(Kgrid,Kgrid,'k--',Kgrid,gs,'b')
legend('45-degree','k_{t+1}=g(k_{t})',2)
title('Policy function')
xlabel('Capital stock')
ylabel('Optimal policy')