% transition.m
% This program implements the algorithm described in Appendix C of
% "Productivity Losses from Financial Frictions: Can Self-financing Undo
% Capital Misallocation?"
% It computes transition dynamics from given initial wealth
% shares w0 and aggregate capital stock K0

clear all
close all
clc
 
%%%%%%%%%%%%%%
% PARAMETERS %
%%%%%%%%%%%%%%
r = 0.05;
d = 0.05;
al = 1/3;
Corr = 0.85;
nu = -log(Corr);
sig = 0.56;
sig2 = sig^2;
Var = sig2./(2*nu);

la = 1.2;

%%%%%%%%%%%%%%%%%%%
% GRID PARAMETERS %
%%%%%%%%%%%%%%%%%%%
I= 500;
zmin = 0;
%zmax = 95-percentile of log-normal distribution:
z_stan_95 = 1.6449;
logzmax = z_stan_95*sqrt(Var);
zmax = exp(logzmax);
check = logncdf(zmax,0,sqrt(Var));
zz = linspace(zmin,zmax,I)'; %GRID
h = (zmax-zmin)/(I-1);
h2 = h^2;
ZFB = zmax^al; %FIRST-BEST TFP

T=100; %TIME LENGTH
N = 500; %NUMBER OF TIME STEPS
dt = T/(N-1); %TIME STEP

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% INITIAL WEALTH SHARES = STATIONARY DIST. OF OU WITH BARRIER %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
w0 = lognpdf(zz,0,sqrt(Var));
w0 = w0./sum(h*w0);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% START FROM THE STEADY STATE CORRESPONDING TO w0 %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for i=1:I-1
    int(i) = h*sum(w0(i+1:I));
end
ibar0=min(find(la*int<1));
f0= 1/la - int(ibar0);
E0 = la*(zz(ibar0)*f0 + h*zz(ibar0+1:I)'*w0(ibar0+1:I));
Z0= E0^al; %truncated expectation
K0 = (al*Z0/(r+d))^(1/(1-al));

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% ORNSTEIN-UHLENBECK PROCESS %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

mu = -nu.*zz.*log(zz)+sig2/2*zz;
s2 = sig2.*zz.^2;

a = nu + nu*log(zz) + sig2/2;
b = -mu+ 2*sig2.*zz;
c = 1/2*sig2.*zz.^2;

x = b*dt/(2*h) - c*dt/h2;
y = 1 - a*dt + c*2*dt/h2;
z = - b*dt/(2*h) - c*dt/h2;

B = zeros(I,I);
A = zeros(I-1,I-1);

for i=2:I-1
    B(i,i-1) = x(i);
    B(i,i) = y(i);
    B(i,i+1) = z(i);
end
A1 = B(2:I-1,2:I);
A = [A1;h*ones(1,I-1)];

clear a;

K = zeros(N,1);
KFB = zeros(N,1);
YFB = zeros(N,1);
w = zeros(I,N);

K(1) = K0;
KFB(1) = K0;
w(:,1) = w0;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% COMPUTE TRANSITION DYNAMICS %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
for n=1:N
    %find zbar
for i=1:I-1
    int(i) = h*sum(w(i+1:I,n));
end
if la==1
    ibar=1;
else
    ibar=min(find(la*int<1));
end
ibar_n(n)=ibar;
zbar(n) = zz(ibar);
f(n)= 1/la - int(ibar);
mkt(n) = la*(f(n)+h*sum(w(ibar+1:I,n)))-1;
E(n) = la*(zz(ibar)*f(n) + h*zz(ibar+1:I)'*w(ibar+1:I,n));
Z(n)= E(n)^al; %truncated expectation
Y(n) = Z(n)*K(n)^al;
W(n) = (1-al)*Z(n)*K(n)^al;
ze(n) = zbar(n)/E(n);
R(n) = al*ze(n)*Z(n)*K(n)^(al-1)-d;
pi(n) = al*((1-al)/W(n))^((1-al)/al);
gK(n) = al*Y(n)/K(n) - (r+d);
for i=1:I
sav(i,n) = la*max(zz(i)*pi(n) - R(n) - d,0)+R(n)-r;
end
Bsav = diag(sav(:,n)-gK(n));
Asav = [Bsav(2:I-1,2:I);zeros(1,I-1)];
AA = A - dt*Asav;
wn = [w(2:I-1,n);1];
wn1 = AA\wn;
w(2:I,n+1) = wn1;
w(1,n+1) = 0;
wsum(n) = sum(h*w(:,n));
K(n+1) = dt*(al*Y(n)-(r+d)*K(n))+K(n);
end