%%%%% solving the optimal growth model by value function iteration

clearvars;
close all;
clc;

%%%%% economic parameters

alpha = 1/3;        %% capital's share in production function
beta  = 0.95;       %% time discount factor   
delta = 0.05;       %% depreciation rate
sigma = 1;          %% CRRA (=1/IES)
rho   = (1/beta)-1; %% implied rate of time preference

kstar = (alpha/(rho+delta))^(1/(1-alpha)); %% steady state
kbar  = (1/delta)^(1/(1-alpha));

%%%%% numerical parameters

max_iter = 500;      %% maximum number of iterations
tol      = 1e-7;     %% treat numbers smaller than this as zero
penalty  = 10^16;    %% for penalizing constraint violations

%%%%% setting up the grid of capital stocks

n        = 1001;     %% number of nodes for k grid
kmin     = tol;      %% effectively zero 
kmax     = kbar;     %% effective upper bound on k grid

k        = linspace(kmin, kmax, n); %% linearly spaced


%%%%% return function

c    = zeros(n,n);

for j=1:n,
    
    c(:,j) = (k.^alpha) + (1-delta)*k - k(j);
    
end

%%%%% penalize violations of feasibility constraints

violations = (c<=0);

c = c.*(c>=0) + eps;

if sigma==1,
    
    u = log(c) - penalty*violations;
    
else
    
    u = (1/(1-sigma))*(c.^(1-sigma) - 1) - penalty*violations;

end

tic;
%%%%% now solve Bellman equation by value function iteration

%%%%% initial guess

v = zeros(n,1);

%%%%% iterate on Bellman operator

for i=1:max_iter,
    
%%%%% RHS of Bellman equation

RHS    = u + beta*kron(ones(n,1),v');

%%%%% maximize over this to get Tv

[Tv,argmax] = max(RHS,[],2);

%%%%% policy that attains the maximum

g = k(argmax);

%%%%% check if converged

error = norm(Tv-v,inf);

fprintf('%4i %6.2e \n',[i, error]);

if error<tol, break, end;

%%%%% if not converged, update and try again

v = Tv;

end

time      = toc;
last_iter = i;

if error<tol && last_iter<=max_iter,
   
fprintf('\n');    
fprintf('time taken (seconds) = %3.3f \n',time);
fprintf('number of iterations = %3.0f \n',last_iter);
fprintf('error at last step   = %3.10f \n',error);
fprintf('\n');   

end

figure(1)
plot(k,v,'b-')
xlabel('capital stock, k')
ylabel('value function, v(k)')

figure(2)
plot(k,g,'r-',k,k,'k--')
xlabel('capital stock, k')
ylabel('policy function, g(k)')


