% This MATLAB software solves a neoclassical stochastic growth model using
% five alternative solution methods with precomputation and compares their 
% accuracy and cost: conventional value function iteration (VFI), endogenous 
% grid method (EGM), envelope condition method (ECM), Euler equation method 
% parameterizing the integrand function Q (EEQ) and Euler equation method 
% jointly parameterizing the integrand function Q and policy function K 
% (EEQK). For a description of these solution methods, see the article 
% "How to Solve Dynamic Stochastic Models Computing Expectations Just Once"
% by Kenneth L.Judd, Lilia Maliar, Serguei Maliar and Inna Tsener, (2016),
% Quantitative Economics, forthcoming, (henceforth, JMMT, 2016).
%
% This software is based on that of Lilia  Maliar and Serguei Maliar for 
% solving the neoclassical stochastic growth model using seven alternative 
% solution methods, as described in the paper "Envelope Condition Method 
% with an Application to Default Risk Models" by Cristina Arellano, Lilia 
% Maliar, Serguei Maliar and Viktor Tsyrennikov, Journal of Economic Dynamics 
% and Control (2016), 69, 436-459 (henceforth, AMMT, 2016). 

% First version: August 19, 2011
% This version:  October 26, 2016
% ------------------------------------------------------------------------
% This software uses the following files: 
% ------------------------------------------------------------------------
% 1. "Main_6_PREC.m"        solves the model and computes the residuals by 
%                           using six alternative solution methods with
%                           precomputation
% 2. "VF_Bellman_PREC.m"    evaluates value function on the left side of
%                           the Bellman equation  
% 3. "FOC_VFI_PREC.m"       computes new capital that satisfies the FOC
% 4. "EGM_BC.m"             computes current capital that satisfies budget 
%                           constraint given next-period capital 
% 5. "Simulation_ECM.m"     simulates a time-series solution given the 
%                           constructed value function  
% 6. "Residuals_PREC.m      computes the residuals in equilibrium conditions
%                           on a given set of points in the state space for
%                           a given numerical solution by evaluating 
%                           integrals exactly
% 7. "Polynomial_2d.m"      constructs the sets of basis functions for ordinary
%                           polynomials of the degrees from one to five, for 
%                           the model with 2 state variables
% 8. "Polynomial_deriv_2d.m"constructs the derivatives of basis functions of 
%                           complete ordinary polynomials of the degrees from 
%                           one to five, for the model with 2 state
%                           variables
% 9. "Integrals_Prec_2d.m"  precomputes integrals for ordinary polynomials 
%                           of the degrees from one to five, for the model 
%                           with 2 state variables (one endogeneous and one
%                           exogenous)
% 10. "epsi_test.mat"       contains a fixed sequence of random numbers
% 11. "csolve.m"            a solver for a nonlinear equation (works faster
%                           than "fsolve"); written by Chris Sims 
% -------------------------------------------------------------------------
% Copyright © 2016 by Lilia Maliar, Serguei Maliar and Inna Tsener. All 
% rights reserved. The code may be used, modified and redistributed under 
% the terms provided in the file "License_Agreement.txt".
% -------------------------------------------------------------------------

clc;
clear all;

Method = 5;   % Choose a solution method: "1", "2", "3", "4", "5"

% "1" - conventional value function iteration (VFI)
% "2" - endogenous grid method (EGM)
% "3" - envelope condition method (ECM)
% "4" - Euler equation method parameterizing the integrand function Q (EEQ)
% "5" - Euler equation method jointly parameterizing the integrand function 
%       Q and policy function K (EEQK)

% 1. Model's parameters
% ---------------------
gam     = 1;         % Utility-function parameter
alpha   = 0.36;      % Capital share in output
beta    = 0.99;      % Discount factor
delta   = 0.025;     % Depreciation rate 
rho     = 0.95;      % Persistence of the log of the productivity level
sigma   = 0.01;      % Standard deviation of shocks to the log of the 
                     % productivity level
                     
% 2. Steady state
% ---------------
A       = (1/beta-(1-delta))/alpha;  
                     % Normalize steady state capital to one             
css     = A-delta;   % Steady state consumption 
yss     = A;         % Steady state output


% 3. Construct a tensor product grid for computing a solution
%------------------------------------------------------------
% Unidimensional grid for capital
kmin    = 0.9;          % min capital on the grid
kmax    = 1.1;          % max capital on the grid
n_gridk = 10;           % number of grid points for capital 
gridk   = linspace(kmin,kmax,n_gridk)'; 
                        % Unidimensional grid for capital of n_gridk points

% Unidimensional grid for productivity                        
zmin    = 0.9;          % min productivity on the grid
zmax    = 1.1;          % max productivity on the grid
n_gridz = 10;           % number of grid points for productivity 
gridz   = linspace(zmin,zmax,n_gridz)'; 
                        % Unidimensional grid for productivity of n_gridz
                        % points

% Two-dimensional tensor product grid 
grid = []; 
for i = 1:n_gridz;
    grid = [grid;[gridk ones(n_gridk,1)*gridz(i,1)]]; 
end

grid_EGM = grid;        % Grid for the EGM method  

n_grid   = n_gridz*n_gridk;  
                        % Number of points in the two-dimensional grid 

k0 = grid(:,1);         % Grid points for capital in the tensor product grid
z0 = grid(:,2);         % Grid points for productivity in the tensor product grid
    
% 4. Compute variance of shocks
%-------------------------------
vcv    = sigma^2;       % Variance of shocks

% 5. Constructing initial guess for value function on the grid 
% ------------------------------------------------------------
D  = 2;              % Initial guess for value function is constructed using
                     % ordinary polynomial of degree 2
                     
X0   = Polynomial_2d(grid,D);
                     % Construct the matrix of explanatory variables X0 
                     % on the grid; the columns of X0 are given by basis 
                     % functions of polynomial of degree "D"        
                     
damp     = 1; if Method==5; damp = 0.1; end                     
                     % Damping parameter for (fixed-point) iteration on 
                     % value function; a default value "1" means no damping
                     % and a smaller value in the interval (0,1) means that 
                     % value function is updated only partially from one
                     % iteration to another to enhance convergence
                     
vf_coef = zeros(6,1);% Initial guess for the coefficients of value function 

V = X0*vf_coef;      % Initialize the value function on the grid; this
                     % grid function will be used to check the convergence 

difv       = 1e+10;  % Initially, set the difference between the value 
                     % functions on two subsequent iterations to exceed 
                     % the convergence criterion

c0   =  A*z0.*k0.^alpha*(css/yss);        % Initial guess for consumption 
k1   =  k0*(1-delta)+A*z0.*k0.^alpha-c0;  % Initial guess for capital
                     % Our initial guess for consumption function is that 
                     % a constant fraction css/yss of the period output 
                     % goes to consumption and the rest goes to investment,
                     % where css/yss is calculated in the steady state
                      
while difv > 1e-6                          % Unit-free convergence criterion
    
        [V_new] = VF_Bellman_PREC(c0,k1,z0,gam,beta,vcv,rho,vf_coef,D);
                                           % Recompute value function; we 
                                           % use precomputation in the
                                           % approximation of value
                                           % function
        vf_coef_new = X0\V_new;            % New vector of coefficients 
        
        vf_coef = damp*vf_coef_new + (1-damp)*vf_coef;   
                                           % Update the coefficients
        difv = max(abs(1-V_new./V));       % Compute the difference between
                                           % new and old value functions

        V = V_new;                         % Store new value function

end

% 5. Main iterative cycle: constructing polynomial approximations of degrees 
% from two to five and from one to five for value function iterative methods 
% and Euler equation methods, respectively
% --------------------------------------------------------------------------

tic;                                    % Start counting time needed to 
                                        % compute the solution

for D = 2:5;                            % For polynomial degrees from 2 to 5...   
        
    Degree(D) = D;                      % Degree of polynomial approximation
    intV = Integrals_PREC_2d(D,vcv)';   % Precompute integrals for constructing 
                                        % expectation functions    
    X0 = Polynomial_2d(grid,D);         % Construct the matrix of polynomial
                                        % basis functions of ordinary poly- 
                                        % nomial of degree D 
    X0der = Polynomial_deriv_2d(grid,D);% Derivative of polynomial of degree 
                                        % D
    Vder = (c0.^(-gam).*(1-delta+A*alpha*z0.*k0.^(alpha-1)));
                                        % Derivative of value function                                 
    vf_coef = X0\V;                     % Initial guess on value function
    vder_coef = X0\Vder;                % Initial guess on the  derivative 
                                        % of value function  
    K_coef  = X0\k1;                    % Initial guess on the policy function
    difk =   inf;                       % Initial criterion of convergence 
    k_old = inf(size(grid,1),1);        % Initialize capital choices on the 
                                        % grid for checking convergence 
  
    while difk > 1e-9;                  % Convergence criterion

        
        %  Method 1. Conventional value function iteration (VFI) with
        %  precomputation
        %==================================================================
        if Method==1        
        
            for i = 1:n_grid       % For each grid point ...
                k1(i) = csolve('FOC_VFI_PREC',k1(i),[],0.000001,5,k0(i),z0(i),A,alpha,gam,delta,beta,vcv,rho,vf_coef,D);                                                                     
                                    % Solve for capital that satisfies FOC;
                                    % we use precomputation in the 
                                    % approximation of derivative of value
                                    % function
            end
         
            c0 = (1-delta)*k0+A*z0.*k0.^alpha-k1; 
                                   % Find consumption from budget constraint
                                   % (5) in JMMT (2016)
            [V_new] = VF_Bellman_PREC(c0,k1,z0,gam,beta,vcv,rho,vf_coef,D);
                                   % Recompute value function using Bellman 
                                   % equation
            vf_coef_new = X0\V_new;
                                   % New coefficients for value function       
            vf_coef = damp*vf_coef_new + (1-damp)*vf_coef;   
                                   % Update the coefficients using damping
        %==================================================================
                                                             
                                   
        % Method 2. Endogenous grid method (EGM)
        %==================================================================                                   
        elseif Method==2; 
               
            k1 = grid_EGM(:,1); % Grid points for next-period capital

            X1_EGM    = Polynomial_2d([k1 z0.^rho],D); 
                               % Construct the matrix of polynomial basis
                               % functions of ordinary polynomial of 
                               % degree D 
            vf1_coef = vf_coef.*intV;       
                               % A new vector of coefficients
                               % with precomputed integrals    
            V1_EGM = X1_EGM*vf1_coef;       
                               % Compute value function using the new vector
                               % of coefficients  

            Xder1_EGM = Polynomial_deriv_2d([k1 z0.^rho],D);
                               % Construct the matrix of derivatives of basis
                               % functions of ordinary polynomial of 
                               % degree D  
            Vder1_EGM = Xder1_EGM*vf1_coef; 
                               % Compute the derivative of value function 
                               % using the new vector of coefficients
            c0 = (beta*Vder1_EGM).^(-1/gam);
                               % Find consumption from FOC (19)

            for i = 1:n_grid   % For each grid point ...
                k0(i) = csolve('EGM_BC',k0(i),[],0.000001,5,k1(i),z0(i),c0(i),A,alpha,delta);
                               % Solve for current capital that satisfies 
                               % budget constraint (5) in JMMT (2016) given 
                               % next-period capital
            end
                       
            if gam==1
                V_new = log(c0)+beta*V1_EGM; 
                                   % Bellman equation if gam=1        
            else
                V_new = (c0.^(1-gam)-1)/(1-gam)+beta*V1_EGM;
            end                    % Bellman equation otherwise
                     
            grid(:,1) = k0;        % Grid points for current capital 
            
            X0 = Polynomial_2d(grid,D);   % Construct polynomial on 
                                          % current state variables
            
            vf_coef_new = X0\V_new;
                                   % New coefficients for value function       
            vf_coef = damp*vf_coef_new + (1-damp)*vf_coef;   
                                   % Update the coefficients using damping
        %==================================================================
                                                            
                                                           
        % Method 3. Envelope condition method (ECM)
        %==================================================================                                                                     
        elseif Method==3 
       
            Vder0 = X0der*vf_coef;      % Compute the derivative of value 
                                        % function
            u0 = Vder0./(1-delta+A*alpha*z0.*k0.^(alpha-1));
                                        % Compute marginal utility from
                                        % envelope condition (20) in JMMT
                                        % (2016)
            c0 = u0.^(-1/gam);          % Find consumption from marginal 
                                        % utility  
            k1 = (1-delta)*k0+A*z0.*k0.^alpha-c0;  
                                        % Find capital from budget constraint
                                        % (5) in JMMT (2016)
            [V_new] = VF_Bellman_PREC(c0,k1,z0,gam,beta,vcv,rho,vf_coef,D);
                                        % Recompute value function; we 
                                        % use precomputation in the
                                        % approximation of value function                                  
            vf_coef_new = X0\V_new;     % New coefficients for value function       
            vf_coef = damp*vf_coef_new + (1-damp)*vf_coef;   
                                        % Update the coefficients using 
                                        % damping
        %==================================================================
  
        
        % Method 4. Euler equation method parameterizing the integrand 
        % function Q (EEQ)
        %==================================================================                                                               
        elseif Method==4  
                          
         
            Vder0 = X0*vder_coef;      % Compute the derivative of value 
                                       % function; as is argued in Section 
                                       % 2.3.3of JMMT, 2016, the derivative 
                                        % of value function coincides with 
                                        % function Q
            u0 = Vder0./(1-delta+A*alpha*z0.*k0.^(alpha-1));
                                        % Compute marginal utility from
                                        % envelope condition (20) in JMMT
                                        % (2016)
            c0 = u0.^(-1/gam);          % Find consumption from marginal utility  
            k1 = (1-delta)*k0+A*z0.*k0.^alpha-c0;  
                                        % Find capital from budget constraint
                                        % (5) in JMMT (2016)
            intV = Integrals_PREC_2d(D,vcv)'; 
                                        % Precompute integrals for 
                                        % constructing expectation functions
            vder1_coef = vder_coef.*intV; 
                                        % A new vector of coefficients
                                        % with precomputed integrals
            X1 = Polynomial_2d([k1 z0.^rho],D);
                                        % Construct the matrix of polynomial
                                        % basis functions of ordinary poly-
                                        % nomial of degree D 
                                        
            Vder1 = X1*vder1_coef;      % Compute the derivative of value  
                                        % using the new vector of coefficients

            Vder0_new = (1-delta+A*alpha*z0.*k0.^(alpha-1)).*(beta*Vder1);
                                        % Recompute the derivative of value 
                                        % function in the grid points
            vder_coef_new = X0\Vder0_new;
                                        % New coefficients for derivative of
                                        % value function       
            vder_coef = damp*vder_coef_new + (1-damp)*vder_coef;   
                                        % Update the coefficients using damping
        %==================================================================
        

        % Method 5. Euler equation method parameterizing the integrand
        % function Q and policy function K (EEQK)
        %==================================================================
        elseif Method==5  

            k1 = X0*K_coef;             % Compute next-period capital
            c0 = (1-delta)*k0+A*z0.*k0.^alpha-k1; 
                                        % Find consumption from budget 
                                        % constraint (5) in JMMT (2016)
            Vder0 = (c0.^(-gam).*(1-delta+A*alpha*z0.*k0.^(alpha-1)));
                                        % Compute the derivative of value 
                                        % function; again, the derivative
                                        % of value function coincides with
                                        % function Q
            vder_coef = X0\Vder0;       % Compute the polynomial coefficients                                       % approximation of the derivative
                                        %for polynomial approximation of the
                                        % derivative of value function on the
                                        % grid
            Vder0 = X0*vder_coef;       % Compute the fitted derivative of value 
                                        % function on the grid
            
            intV = Integrals_PREC_2d(D,vcv)'; 
                                        % Precompute integrals for 
                                        % constructing expectation functions
            vder1_coef = vder_coef.*intV; 
                                        % A new vector of coefficients
                                        % with precomputed integrals
            X1 = Polynomial_2d([k1 z0.^rho],D);
                                        % Construct the matrix of polynomial
                                        % basis functions of ordinary poly-
                                        % nomial of degree D 
            Vder1 = X1*vder1_coef;      % Compute derivative of value using 
                                        % the new vector of coefficients
            
            k1_new = (1-delta+A*alpha*z0.*k0.^(alpha-1)).*(beta*Vder1).*k1./Vder0;
                                        % Recompute capital in the grid points
            K_coef_new = X0\k1_new;
                                        % New coefficients for capital policy 
                                        % function       
            K_coef = damp*K_coef_new + (1-damp)*K_coef;   
                                        % Update the coefficients using damping
        %==================================================================
                   
        end               
        
        % Checking convergence 
        % --------------------
        if Method == 2 
            difk = max(abs(1-k0./k_old))  
                        % For EGM, we check convergence of current-period 
                        % capital
            k_old = k0;        
        else         
            difk = max(abs(1-k1./k_old))
                        % For the other methods, we check convergence of 
                        % next-period capital
            k_old = k1;
        end
    end
    
    % After the solution is computed by any method, we compute value function
    % and the derivative of value function for the constructed policy rules
    
    difv = 1e10;           % Initially, convergence criterion is not
                           % satisfied
    while difv > 1e-10     % Unit-free convergence criterion
        
        [V_new] = VF_Bellman_PREC(c0,k1,z0,gam,beta,vcv,rho,vf_coef,D);
                                        % Recompute value function; we 
                                        % use precomputation in the
                                        % approximation of value function   
        vf_coef_new = X0\V_new;         % New vector of coefficients 
        vf_coef = damp*vf_coef_new + (1-damp)*vf_coef;   
                                        % Update the coefficients
        difv = max(abs(1-V_new./V));    % Compute the difference between
                                        % new and old value functions
        V = V_new;                      % Store new value function
    end  

    V = V_new;                          % Update value function to be 
                                        % used as an initial guess for a
                                        % higher degree polynomials
       
    CPU(D) = toc;                       % Store running time
    VK(1:1+D+D*(D+1)/2,D) = vf_coef;    % Store the solution coefficients 

end
 
% Evaluate residuals in the model's equations on a simulated path
% ---------------------------------------------------------------
fprintf(1,'Method = %i:\n',Method);
fprintf(1,'ACCURACY EVALUATION AND RUNNING TIME:\n\n');
for D = 2:5 % For polynomial degrees from 2 to 5... 
    [k,z] = Simulation_ECM(A,alpha,gam,delta,sigma,rho,VK(1:1+D+D*(D+1)/2,D),D);
            % Simulate the solution under a sequence of 10,200 shocks stored
            % "epsi_test.mat" and discard the first 200 entries to eliminate
            % the effect of initial conditions
    [Mean_Residuals(D),Max_Residuals(D)] = Residuals_PREC(k,z,A,alpha,gam,delta,sigma,rho,beta,VK(1:1+D+D*(D+1)/2,D),D);
            % Compute residuals in the model's equations in each point of the 
            % simulated path and compute the mean and maximum residuals in
            % the model's equations
    fprintf(1,'Polynomial of degree = %i:\nRunning time = %.2f, Mean residuals = %.2f, Max residuals = %.2f\n\n',Degree(D),CPU(D),Mean_Residuals(D),Max_Residuals(D));
            % Display the results
end