% MATLAB software that solves the multi-country model using the generalized 
% stochastic-simulation algorithm (GSSA) with precomputation, as described 
% in 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 multi-country model using GSSA method, as described in the 
% paper "Numerically Stable and Accurate Stochastic Simulation Approaches 
% for Solving Dynamic Economic Models" by Kenneth L. Judd, Lilia Maliar and
% Serguei Maliar, (2011), Quantitative Economics 2/2, 173–210 (henceforth, 
% JMM, 2011).  
%
% First version: August 19, 2011.
% This version: October 26, 2016. 
% 
% ------------------------------------------------------------------------
% This software uses the following files: 
% ------------------------------------------------------------------------
% 1. "Main_PREC.m"        computes GSSA solutions to the N-country model 
%                         by parametrizing N capital policy functions and
%                         N integrand functions using precomputation
% 2. "Residuals_PREC_N.m" computes approximation errors in the optimality 
%                         conditions for a given solution on a given set of
%                         points in the state space by evaluating integrals 
%                         exactly, for the N-country model
% 3. "Productivity.m"     generates random draws of the productivity shocks  
%                         and simulates the corresponding series of the  
%                         productivity levels; borrowed from JMM (2011) 
% 4. "Ord_Polynomial_N.m" constructs the sets of basis functions for ordinary 
%                         polynomials of the degrees from one to five, for
%                         the N-country model; borrowed from JMM (2011) 
% 5. "Integrals_PREC_N.m" constructs the sets of precomputed integrals 
%                         defined by equation (50) in JMMT (2016) for ordinary 
%                         polynomials of the degrees from one to five, for 
%                         the N-country model
% 6. "GH_Quadrature.m"    constructs integration nodes and weights for the 
%                         Gauss-Hermite rules with the number of nodes in 
%                         each dimension ranging from one to ten; borrowed 
%                         from JMM (2011)                      
% 7. "Monomials_1.m"      constructs integration nodes and weights for an N-
%                         dimensional monomial (non-product) integration rule 
%                         with 2N nodes; borrowed from JMM (2011)  
% 8. "Monomials_2.m"      constructs integration nodes and weights for an N-
%                         dimensional monomial (non-product) integration rule 
%                         with 2N^2+1 nodes; borrowed from JMM (2011) 
% 9. "zT12200N30.mat"     contains the series of the productivity levels of 
%                         length 12,200 for 30 countries that are used for 
%                         computing solutions and for evaluating accuracy 
% -------------------------------------------------------------------------
% 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;

% 1. Choose the number of countries and simulation length
% -------------------------------------------------------
N     = 2;      % Number of countries, 1<=N<=30 (note that the code also 
                % works for the one-country case, N=1)
T     = 2000;   % Simulation length for the solution procedure,
                % T<=2,000     
                 
% To solve models with N>30 or T>2,000, one needs to simulate new series
% of the productivity levels by enabling the code in section 5.  

% 2. 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
vcv = sigma^2*(eye(N)+ones(N,N)); 
                  % Variance-covariance matrix of the countries' productivity 
                  % shocks in which diagonal terms are equal to 2*sigma^2   
                  % and in which off-diagonal terms are equal to sigma^2; 
                  % this vcv follows from the assumption that a country's 
                  % shock has both common-for-all-countries and country-
                  % specific components; N-by-N

                                     
% 3. The normalizing constant, A, and welfare weight, tau
% -------------------------------------------------------
A       = (1-beta+beta*delta)/alpha/beta;  % The normalizing constant in output  
tau     = 1;                               % The welfare weight of country j 

% The above normalization ensures that steady state of capital of all 
% countries is equal to one 

% 4. Initial condition
% --------------------
k(1,1:N) = 1;  % Initial condition for capital (is equal to steady state)
z(1,1:N) = 1;  % Initial condition for the productivity level (is equal to 
               % steady state)

% 5. Construct the productivity levels, z 
% ---------------------------------------
% z12200 = Productivity(T,N,z(1,1:N),sigma,rho);
                               % Generate a random draw of the productivity 
                               % shocks and simulate the corresponding  
                               % series of the productivity levels of length   
                               % T periods for N countries 
% save zT12200N30 z12200;      % Save the series of the productivity levels  
                               % into a file "zT12200N30.mat" 
load zT12200N30;               % Load the previously saved series of the 
                               % productivity levels of length 12,200 for 
                               % 30 countries (the first 2,000 observations
                               % are used for finding a solution, and the 
                               % remaining 10,200 observations are used for
                               % evaluating accuracy)
z = z12200(1:T,1:N);           % Restrict the series of the productivity 
                               % levels for the solution procedure to the 
                               % given T<=2,000 and N<=30
% _________________________________________________________________________
%                               
% Compute a first-degree polynomial solution using the one-node Monte Carlo  
% integration method (this solution will be used as an initial guess for
% the other cases) 
% _________________________________________________________________________
%
tic;                  % Start counting time needed to compute the solution
                            
% 6. The GSSA parameters  
% ---------------------
kdamp     = 0.1;     % Damping parameter for (fixed-point) iteration on 
                     % the coefficients of the capital policy functions
dif_1d = 1e+10;      % Set the initial difference between the series from
                     % two iterations in the convergence criterion (condition
                     % (10) in JMM, 2011) to a very large number
                     
% To achieve convergence under N>10, one may need to modify the values of 
% the damping parameter kdamp or refine the initial guess 

% 7. Initialize the first-degree capital policy functions of N countries 
%-----------------------------------------------------------------------                          
bk_1d  = [zeros(1,N); diag(ones(1,N)*0.9);diag(ones(1,N)*0.1)]; 
% Matrix of polynomial coefficients of size (1+2N)-by-N: for each country  
% (each column), 1+2N rows correspond to a constant, N coefficients on the
% capital stocks, k(t,1),...,k(t,N), and N coefficients on the productivity 
% levels, z(t,1),...,z(t,N)

% As an initial guess, assume that a country's j capital depends only on 
% its own capital and productivity level as k(t+1,j)=0.9*k(t,j)+0.1*z(t,j); 
% (note that in the steady state, we have k(t+1,j)=0.9*k(t,j)+0.1*z(t,j)=1)

% Note that diag(ones(1,N)*q) delivers an N-by-N matrix with  diagonal 
% entries equal to q. 

% 8. Initialize the capital series
% --------------------------------
k_old = ones(T+1,N);   % Initialize the series of next-period capital of N
                       % countries; these series are used to check the
                       % convergence on the subsequent iteration (initially, 
                       % capital can take any value); (T+1)-by-N
                       
% 9. Compute the initial guees on the coefficients of the capital policy
% function
% ----------------------------------------------------------------------            
while dif_1d > 1e-6;            % 10^(-6) is a convergence parameter

% 9.1 Generate time series of capital
% -----------------------------------
for t = 1:T 
    x(t,:) = [1 k(t,:) z(t,:)];   % The basis functions of the first-degree 
                                  % polynomial at time t
    k(t+1,:) = x(t,:)*bk_1d;      % Compute next-period capital using bk_1d
end

% 9.2 Compute consumption series 
%-------------------------------
C = (A*k(1:T,:).^alpha.*z(1:T,:) - k(2:T+1,:)+k(1:T,:)*(1-delta))*ones(N,1);
% Aggregate consumption is computed by summing up individual consumption, 
% which in turn, is found from the individual budget constraints; T-by-1

c = C*ones(1,N)/N;                % Individual consumption is the same for
                                  % all countries; T-by-N 

% 9.3 Evaluate the percentage (unit-free) difference between the series  
% from the previous and current iterations
% ---------------------------------------------------------------------
dif_1d = mean(mean(abs(1-k./k_old)));
                % Compute a unit-free difference between the capital series 
                % from two iterations; see condition (10) in JMM (2011)
                   
% 9.4 Monte Carlo realizations of the right side of the Euler equation, Y, 
% in condition (C4) in the online Appendix C for JMM(2011)
%-------------------------------------------------------------------------
for j = 1:N
   Y(1:T-1,j) = beta*c(2:T,j).^(-gam)./c(1:T-1,j).^(-gam).*(1-delta+alpha*A*k(2:T,j).^(alpha-1).*z(2:T,j)).*k(2:T,j);
end  % (T-1)-by-N

% 9.5 Compute and update the coefficients of the capital policy functions 
% -----------------------------------------------------------------------
bk_hat_1d  = x(1:T-1,:)\Y(1:T-1,:);      
                              % Compute new coefficients of the capital 
                              % policy functions
bk_1d = kdamp*bk_hat_1d + (1-kdamp)*bk_1d; 
                              % Update the coefficients of the capital  
                              % policy functions using damping 
                                     
% 9.6 Store the capital series 
%-----------------------------
k_old = k;         % The stored capital series will be used for checking 
                   % the convergence on the subsequent iteration
end;

time_GSSA_1d     = toc;     % Time needed to compute the initial guess 

% _________________________________________________________________________                              
%
% Compute polynomial solutions of the degrees from one to D_max using
% precomputation
% _________________________________________________________________________

% 10. The GSSA parameters  
% -----------------------
kdamp     = 0.1;      % Damping parameter for (fixed-point) iteration on 
                      % the coefficients of the capital policy functions
dif_GSSA_D  = 1e+10;  % Set the initial difference between the series from
                      % two iterations in the convergence criterion (condition
                      % (10) in JMM, 2011) to a very large number

% To achieve convergence under N>10, one may need to modify the values of 
% the damping parameter kdamp or refine the initial guess 

% 11. The matrix of the polynomial coefficients
% ---------------------------------------------                             
D_max  = 2;           % Maximum degree of a polynomial: the program computes
                      % polynomial solutions of the degrees from one to D_max;
                      % (D_max can be from 1 to 5) 
for D = 1:D_max       % For the polynomial degrees from one to D_max
    npol(D) = size(Ord_Polynomial_N([k(1,:) z(1,:)],D),2);      
end
                      % Construct polynomial bases and compute their number
                      % (this is needed for finding the number of the polynomial
                      % coefficients in the policy functions)
                      
BK = zeros(npol(D_max),N,D_max); % Matrix of polynomial coefficients of the 
                                 % capital policy functions for the polynomial
                                 % solutions of the degrees from one to D_max;
                                 % npol(D_max)-by-N-by-D_max
                                 
bq_mult = zeros(npol(D_max),N,D_max);
                             % Matrix of the precomputed integrals for the polynomial
                             % solutions of the degrees from one to D_max;
                             % npol(D_max)-by-N-by-D_max
 
% 12. Construct the sets of precomputed integrals
% -----------------------------------------------
                             
for D=1:D_max;
    integrals = Integrals_PREC_N(N,N,vcv,D)';
                             % Compute the values of the integrals of future 
                             % shocks defined by equation (50) in JMMT (2016) 
                             % for a country j
                             
    % Note that "Integrals_PREC_N.m" is a function that must be used together 
    % with "Ord_Polynomial_N.m". The inputs of this polynomial routine must 
    % be in the form s = [k, z], where k and z are n1- and n2-dimensional vectors, 
    % respectively. This ordering of state variables ensures that precomputed 
    % integrals match the corresponding polynomial terms in the approximating 
    % function of degree D
                             
    bq_mult(1:npol(D),1:N,D) = integrals*ones(1,N);
                             % Store the precomputed values of the integrals 
                             % of future shocks for each polynomial degree 
                             % and country
end


% Under precomputation, conditional expectation (integral) is approximated  
% by the value of the integrand function evaluated at a vector of coefficients
% adjusted by bq_mult

% 13. Compute the polynomial solutions of the degrees from one to D_max
% ---------------------------------------------------------------------

tic;             % Start counting time needed to compute the solution

for D=1:D_max
    
    % 13.1 Using the previously computed capital series, compute the initial 
    % guess on the coefficients under the  selected approximation method
    % ----------------------------------------------------------------------
    X = Ord_Polynomial_N([k(1:T,:), z(1:T,:)],D);
                           % Construct the polynomial bases on the series 
                           % of state variables from the previously computed 
                           % time-series solution
    bk_D = X(1:T-1,:)\Y(1:T-1,:);
                           % Compute the initial guess on the coefficients
                           % using the MATLAB's backslash operator
    k_old = ones(T+1,N);   % Initialize the series of next-period capital of N
                           % countries; these series are used to check the
                           % convergence on the subsequent iteration (initially, 
                           % capital can take any value); (T+1)-by-N
    dif_GSSA_D  = 1e+10;   % Convergence criterion (initially is not satisfied)
    
    % 13.2 The main iterative cycle of GSSA
    % -------------------------------------              
    while dif_GSSA_D > 1e-10; % 10^(-10) is a convergence parameter
    
    % 13.2.1 Generate time series of capital
    % -------------------------------------- 
    
    for t = 1:T
        X(t,:) = Ord_Polynomial_N([k(t,:), z(t,:)],D);
        % The basis functions of a polynomial of degree D at time t
        k(t+1,:) = X(t,:)*bk_D;                      
        % Compute next-period capital using bk_D
    end
    
    % 13.2.2 Compute consumption series of all countries 
    %---------------------------------------------------
    k0  =  k(1:T,:);                     % N current capital stocks  
    z0  =  z(1:T,:);                     % N current productivity levels  
    k1  =  k(2:T+1,:);                   % N next-period capital stocks 
    C = (A*k0.^alpha.*z0 - k1+k0*(1-delta))*ones(N,1);
    % Aggregate consumption is computed by summing up individual consumption, 
    % which in turn, is found from the individual budget constraints; T-by-1


    c = C*ones(1,N)/N;            % Individual consumption is the same for
                                  % all countries; T-by-N 
 
    % 13.2.3 Approximate the conditional expectations for t=1,...T-1 using 
    % precomputation
    %----------------------------------------------------------------------
    %
    % 13.2.3.1 Approximate N integrand functions
    % ----------------------------------------------------------------
        Q = c(1:T,:).^(-gam).*(1-delta+alpha*A*k(1:T,:).^(alpha-1).*z0(1:T,:));
                           % The integrand of each expectation function, Q,
                           % in period t depends only on current state variables
        bq_D  = X\Q;       % Compute the coefficients of N integrand functions 
                           % using the MATLAB's backslash operator
        Q = X*bq_D;        % Compute the fitted values of the integrands
 
    % 13.2.3.2 Approximate the values of conditional expectations using Q
    % ---------------------------------------------------------------------
        Y = zeros(T,N); % Allocate memory for the variable Y
   
        EQ1 = Ord_Polynomial_N([k1(1:T,:), z(1:T,:).^rho],D)*(bq_D.*bq_mult(1:npol(D),1:N,D));
                        % Compute the expectation of each integrand function
                        % Q; Note EQ1 is computed using the coefficients of
                        % the integrand function, bq_D, adjusted by the 
                        % precomputed integrals, bq_mult 
        
        Y(1:T,:) = beta*EQ1(1:T,:).*(1-delta+alpha*A*k(1:T,:).^(alpha-1).*z(1:T,:)).*k1(1:T,:)./Q(1:T,:);
                        % Compute the values of conditional expectation
   
    % 13.2.4 Evaluate the percentage (unit-free) difference between the 
    % capital series from the previous and current iterations
    % -----------------------------------------------------------------
    
    dif_GSSA_D = mean(mean(abs(1-k./k_old)))
                % Compute a unit-free difference between the capital series 
                % from two iterations; see condition (10) in JMM (2011)

    % 13.2.5 Compute and update the coefficients of the capital policy 
    % functions 
    % ----------------------------------------------------------------
    bk_hat_D = X(1:T-1,:)\Y(1:T-1,:);
                              % Compute new coefficients of the capital 
                              % policy functions 
    bk_D = kdamp*bk_hat_D + (1-kdamp)*bk_D; 
                              % Update the coefficients of the capital  
                              % policy functions using damping 
                                     
    % 13.2.6 Store the capital series 
    %--------------------------------
    k_old = k;       % The stored capital series will be used for checking 
                     % the convergence on the subsequent iteration
    end;

    % 13.2.7 The GSSA output for the polynomial solution of degree D
    % --------------------------------------------------------------
    BK(1:npol(D),1:N,D) = bk_D; % Store the coefficients of the polynomial  
                                % of degree D that approximates capital 
                                % policy functions of N countries 
    time_GSSA(D) = toc;         % Time needed to compute the polynomial  
                                % solution of degree D
end

% 14. Accuracy test of the GSSA solutions: errors on a stochastic simulation 
% --------------------------------------------------------------------------

% 14.1 Specify a set of points on which the accuracy is evaluated
%----------------------------------------------------------------
T_test = 10200;                  % Choose the simulation length for the test 
                                 % on a stochastic simulation, T_test<=10,200 

z_test = z12200(T+1:T+T_test,1:N); % Restrict the series of the productivity 
                                   % levels for the test on a stochastic 
                                   % simulation to the given T_test<=10,200  
                                   % and N<=30                             
          
k_test(1,1:N) = 1;  % Initial condition for capital (equal to steady state)

% To implement the test on a stochastic simulation with T_test>10,200, one
% needs to simulate new series of the productivity levels with larger T_test 
% by enabling the code in section 5 of this code.

% 14.2 Compute errors on a stochastic simulation for the GSSA polynomial 
% solution of degrees D=1,...,D_max
% ----------------------------------------------------------------------

for D = 1:D_max 
    
    % 14.3.1 Simulate the time series solution under the given capital-
    % policy-function coefficients, BK(:,:,D) with D=1,...,D_max 
    %------------------------------------------------------------------
    bk = BK(1:npol(D),:,D);     % The vector of coefficients of the 
                                % polynomial of degree D 
    for t = 1:T_test-1
        X_test = Ord_Polynomial_N([k_test(t,:) z_test(t,:)],D);
        % The basis functions of a polynomial of degree D at time t
        k_test(t+1,:) = X_test*bk;
        % Compute next-period capital using bk
    end    
    
    % 14.3.2 Errors across 10,000 points on a stochastic simulation
    % -------------------------------------------------------------
    discard = 200; % Discard the first 200 observations to remove the effect
                   % of the initial conditions 
                   
    [Errors_mean(1,D),Errors_max(1,D), time_test(1,D)] = Residuals_PREC_N(k_test,z_test,bk,N,D,alpha,gam,delta,beta,A,rho,vcv,discard);
 
    % Errors_mean    is the unit-free average absolute approximation error  
    %                across 4N+1 optimality conditions (in log10) 
    % Errors_max     is the unit-free maximum absolute approximation error   
    %                across 4N+1 optimality conditions (in log10) 
end

disp(' '); disp('  GH quadrature     GSSA OUTPUT:'); disp(' '); 
disp('RUNNING TIME (in seconds):'); disp('');
disp('a) for computing the solution'); 
disp('      1        2         3         4        5');disp(time_GSSA);
disp('b) for implementing the accuracy test'); 
disp('      1        2         3         4        5'); disp(time_test);
disp('APPROXIMATION ERRORS (log10):'); disp(''); 
disp('a) mean error across 4N+1 optimality conditions'); 
disp('      1        2         3         4        5');disp(Errors_mean)
disp('b) max error across 4N+1 optimality conditions'); 
disp('      1        2         3         4        5');disp(Errors_max)
