%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%        Matlab code for solving a stochastic growth model in HW # 5.		
%            Jinhui Bai (jinhui.bai@yale.edu), Feb 9, 2004
% 
% Discription of the model: log utility, and two-state Markov shocks
% Algorithm: value function iteration in discrete state dynamic programming
% Note: This code is a revision of the code written by Prof George Hall at Yale
%       Main Revision: vectorize all code to avoid "For loop", which should
%                      speed up the program.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                         Outline of the Code
% Step 1: Form return function in dynamic programming		
% Step 2: Do value function iteration to find value and policy function
% Step 3: Find the invariant distribution over (K,A) pairs and calculate
%         long-run expected values of capital stock and output
% Step 4: Simulate a path and calculate the sample average
% 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear
diary brock_mirman.out;
time_start = cputime;
disp('solving a stochastic growth model in HW # 5');
disp('');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Step 1: Form return function in dynamic programming
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  set parameter values
%
ALPHA  = 0.36;              % production parameter
BETA   = 0.9;               % subjective discount factor 
DELTA  = 1;                 % capital depreciation rate
PROB   = [ .95 .05; .1 .9];   % prob(i,j) = probability (A(t+1)=Aj | A(t) = Ai)
A_HIGH = 1.1;               % high value for technology
A_LOW  = 0.9;               % low value for technology
%
%   form capital grid
%
KSS  = (ALPHA*BETA)^(1/(1 - ALPHA));    % steady-state capital stock
MIN_K =  0.95 * KSS;                     % minimum value of the capital grid
MAX_K =  1.05 * KSS;                     % maximum value of the capital grid   
INK   =  0.1 * KSS;                      % size of capital grid increments
NK   = round((MAX_K - MIN_K)/INK + 1);    % number of grid points
kgrid = [ MIN_K:INK:MAX_K ];              % capital grid
%
%   form single period return function
%
LOWER_BOUND = -10000;                % Worst-case Senario for utility
r1 = LOWER_BOUND(ones(NK,NK));       % Initialize utility when A = high
r2 = LOWER_BOUND(ones(NK,NK));       % Initialize utility when A = low
[k_Matrix1 k_Matrix2] = meshgrid(kgrid);  % Form meshgrid to vectorize the loop
consum_High = A_HIGH * k_Matrix1.^(ALPHA) + (1 - DELTA) * k_Matrix1 - k_Matrix2;
consum_Low  = A_LOW * k_Matrix1.^(ALPHA) + (1 - DELTA) * k_Matrix1 - k_Matrix2;
isvalid_High = [consum_High > 0];
isvalid_Low = [consum_Low > 0];
r1(isvalid_High) = log(consum_High(isvalid_High));
r2(isvalid_Low) = log(consum_Low(isvalid_Low));

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Step 2: Do value function iteration to find value and policy function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
v       = ones(NK,2);
decis   = ones(NK,2);
iter    = 0;
metric1 = 10;
metric2 = 10;
[rs,cs] = size(r1);
%
%  iterate on Bellman's equation and get the decision 
%  rules and the value function at the optimum         
%
while ((metric1 ~= 0) | (metric2 > 0.01));
    expectValue1 = BETA * v * (PROB(1,:))';
    expectValue2 = BETA * v * (PROB(2,:))';
    totalValueMatrix1 = r1 + expectValue1(:,ones(1,NK));
    totalValueMatrix2 = r2 + expectValue2(:,ones(1,NK));
    [tv1,tdecis1] = max (totalValueMatrix1);
    [tv2,tdecis2] = max (totalValueMatrix2);

    tdecis = [tdecis1' tdecis2'];
    tv = [tv1' tv2'];

    metric1 = max(any(tdecis-decis));
    metric2 = max(abs((v(:)-tv(:))./v(:)));
    v = tv;
    decis = tdecis;
    iter = iter + 1;
disp([ iter metric1 metric2 ]);
end;
decis = (decis - 1) * INK + MIN_K;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Step 3: Invariant Distribution and Long Run Average Value
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%   Form transition matrix: trans is the transition matrix from state at t (row)
%                           to the state at t+1 (column) 
%   Use convergence theorem to calculate Invariant distribution 
%         (Reason: Avoid inverting matrix of large dimensions)
% 
g2 = sparse(cs,cs);
g1 = sparse(cs,cs);
for i = 1:cs
    g1(i,tdecis1(i)) = 1;
    g2(i,tdecis2(i)) = 1;
end
trans = [ PROB(1,1)*g1 PROB(1,2)*g1; PROB(2,1)*g2 PROB(2,2)*g2];
probst = (1/(2*NK))*ones(2*NK,1);
test = 1;
while test > 10^(-8);
   probst1 = trans' * probst;
   test = max(abs(probst1-probst));
   probst = probst1;
end;
%
%   vectorize the decision rule to be conformable with probst
%   calculate mean level of capital
%
kk = decis(:);
yy = decis.^ALPHA * [A_LOW 0; 0 A_HIGH];
yy = yy(:);
meanK = probst' * kk;
meanY = probst' * yy;
%
%  calculate measure over (k,A) pairs
%  lambda has same dimensions as decis
%
lambda = zeros(cs,2);
lambda(:) = probst;
%
%   calculate stationary distribution of capital 
%
probk = sum(lambda,2);     
%
%   print out results
%
disp('PARAMETER VALUES');
disp('');
disp('    alpha      beta       '); 
disp([ ALPHA BETA ]);
disp(''); 
disp('RESULTS ');
disp('');
disp('      mean of K        mean of Y');
disp([ meanK meanY]);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Step 4: Simulate life histories of the agent
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

disp('SIMULATING LIFE HISTORY');  
kmark = 1;
k = kgrid(kmark,1);        % initial level of assets
N = 10000;                   % number of periods to simulate
s0 = 1;                    % initial state 
states   = zeros(N - 1,2);
controls = zeros(N - 1,2);
[chain,state] = markov(PROB,N,s0);
for i = 1:N - 1;
    if chain(i) == 1;
       kprime = decis(kmark,1);
       invest = kprime - (1 - DELTA) * k; 
       cons   = A_HIGH * k^(ALPHA) - invest;
       y      = A_HIGH * k^(ALPHA);
       kmark = tdecis(kmark,1);
    elseif chain(i) == 2;
       kprime = decis(kmark,2);
       invest = kprime - (1 - DELTA) * k; 
       cons   = A_LOW * k ^(ALPHA) - invest;
       y      = A_LOW * k^(ALPHA);
       kmark = tdecis(kmark,2);
    else;
      disp('something is wrong with chain');
    end;
    states(i,:) = [ k chain(i) ];
    controls(i,:) = [ cons kprime ];
    output(i,:) = [y];
    k = kprime;
end;

disp('RESULTS ');
disp('');
disp('    Sample mean of K     Sample mean of Y');
disp([ mean(states(:,1)) mean(y)]);


disp('Time Spent in Computation')
disp('');
time_spent = cputime - time_start;
time_spent


    
    