% ========================================================================
%  Code for discretization method in "An Extension of Rouwenhorst Method
%  for Non-Stationary Environments" by Giulio Fella, Giovanni Gallipoli
%  and Jutong Pan
% 
%  Purpose:
%  	   Rouwenhurst method to approximate non-stationary AR(1) process 
%          z(t) = rho*z(t-1)+ epsilon(t),   epsilion(t)~iid N(0,sigmaeps)
%  	   by a discrete Markov chain
%  INPUTS:  rho - serial correlation coefficient
%           sigma_z0 - standard deviation of initial condition z0 (can be
%           zero)
%           sigma_eps - standard deviation of white noise innovation
%           N - dimension of the state space
%           T - dimension of lifetime (convention: 1 is time when z0 is drawn)
%  OUTPUT:  Pmat -  NxNxT matrix of T (NxN) transition matrices  
%                   transition probabilities are arranged by row
%           z_vect - an NxT matrix, each column stores the Markov state 
% 	    space for period t
% !========================================================================%

function [z_vect, Pmat] = rouwen_finiteT(rho,sigma_z0,sigma_eps,N,T) 

sigma_z = zeros(1,T);
z_vect = zeros(N,T); 
p_si = zeros(1,T);
Pmat = zeros(N,N,T);

if N < 2;
    disp('The state space has to have dimension N>1. Exiting.')
    return;
end

% Step 1: Construct the evenly-spaced N-state space over
%         [-sqrt(N-1)*sigma_z(t),sqrt(N-1)*sigma_z(t)] for each period t

% 1. Compute unconditional variance at time i 
sigma_z(1) = sigma_z0;
for i = 2:T;
    sigma_z(i) = sqrt(rho^2*sigma_z(i-1)^2+sigma_eps^2);
end; 

h = 2*sqrt(N-1)*sigma_z/(N-1); % grid step
z_vect = repmat(h,N,1);
z_vect(1,:)=-sqrt(N-1) * sigma_z;
z_vect = cumsum(z_vect,1); % State space

% 2. Compute the transition matrices
q = 1/2; % z0 is white noise with zero mean
for j = 1:T 
    if j>1 
        % Compute q for t>1 
        q = (sigma_z(j)+rho*sigma_z(j-1))/(2*sigma_z(j) );   
    end

    if N == 2;
        Pmat(:,:,j) = [q, 1-q; 1-q, q];  
        continue;
    end; % of if on size 2

    P1 = [q, 1-q; 1-q, q];

    for i = 2:N-1;
    
        P2 = q *     [P1,zeros(size(P1,1),1); zeros(1,size(P1,2)),0 ] + ...
             (1-q) * [zeros(size(P1,1),1),P1; 0,zeros(1,size(P1,2)) ] + ...
             (1-q) * [zeros(1,size(P1,2)),0 ; P1,zeros(size(P1,1),1)] + ...
             q *     [0,zeros(1,size(P1,2)) ; zeros(size(P1,1),1),P1];
    
        P2(2:i,:) = 0.5*P2(2:i,:);
        
        if i==N-1;
            Pmat(:,:,j) = P2;
        else
            P1 = P2;
        end;
    end; % of for


end
