/* approximates the policy function for consumption with both a and k
                        as state variables, using an orthogonal collocation scheme. 
			a is a markov process */

/* initialize program */

new;
library optmum,pgraph,nlsys;
graphset;
nlset;
_pdate = "";
_psymsiz =  3;
_paxht   = .2;
_pnumht  = .15;
_ptitlht = .3; 
_paxes   = { 1, 1 }; 

/* set dimension of model  */

nvars = 5;
nshocks = 1;
eulers = 1;						@ indicates equation number(s) of euler equations @
neulers = rows(eulers);					@ counts # of euler equations in the model  @

@ establish parameters @

alp = 0.33;
bet = 0.96;
del = 0.1;
phi = 2;
rho = 0.8;
sige = 0.0067;
p = alp|bet|del|phi|rho|sige;

const = 0.5;
pmat = markovchain(const,rho);

let x1[2,1] = 1 0;
let x2[2,1] = 0 1;
x=x1~x2; 
prob = pmat*x;
stda = sqrt((1/(1-rho^2))*(sige^2));
ah = 1+stda;
al = 1-stda;

/* load log-linear approximation routines */

@ Note: p, xbar and tme must be global before loglin.prc can be compiled in stogrow.src.  @

xbar = zeros(nvars,1);
tme = 0;

#include stogrow.src

/* given drawing pdraw, obtain non-linear model solution */

pdraw = p;

ss = steady(pdraw);
kss = ss[4];
css = ss[2];
ass = ss[5];

/* establish ranges for states */

omegaa = 0.033;
asteps = 2;
astep = (ah-al)/(asteps-1);
arange = seqa(al,astep,asteps);
anorm = (arange-ass)/(omegaa*ass);
na = rows(arange);

omegak = 0.25;
kl = kss-omegak*kss;
kh = kss+omegak*kss;
ksteps = 10001;
kstep = (kh-kl)/(ksteps-1);
krange = seqa(kl,kstep,ksteps);
knorm = (krange-kss)/(omegak*kss);
nk = rows(krange);

/* locate zeros of polynomials */

orda = 2; @ ord = order of polynomial @
zersa = 1|2;

ordk = 4;
zersk = zerproc(ksteps,ordk);

/* obtain linear solution to the model */

{a,b,f,gno,q,h} = modelsol(p);

sigk = f[2,4];
siga = f[2,5];

startval = zeros(orda*ordk,1);
startval[1] = css;
startval[2] = omegak*sigk*css;
startval[ordk+1] = omegaa*siga*css;
startval[ordk+2] = 0.5*omegak*sigk*css*omegaa*siga*css;

{ gamopt,fopt,gopt,retcode } = nlprt(nlsys(&feval,startval));

omegak = 0.25;
kl = kss-omegak*kss;
kh = kss+omegak*kss;
ksteps = 101;
kstep = (kh-kl)/(ksteps-1);
krange = seqa(kl,kstep,ksteps);
knorm = (krange-kss)/(omegak*kss);
nk = rows(krange);

optcs = zeros(ksteps,asteps);
i=1; do while i<=ksteps;
    j=1; do while j<=asteps;
        optcs[i,j] = ceeofak(anorm[j],knorm[i],gamopt);
    j=j+1; endo;
i=i+1; endo;

/* calculate slopes of policy rules */

dcdkl = (optcs[2:rows(optcs),1]-optcs[1:rows(optcs)-1,1])/kstep;
dcdkh = (optcs[2:rows(optcs),2]-optcs[1:rows(optcs)-1,2])/kstep;

/* graph results */

begwind;
window(2,1,0);

_pmsgctl = { 3.53 1.14 .2 0 1 1 0 };
_pmsgstr = "(C[*], K[*])";
_psym = { 3.53 1.163352 7 7 1 1 0};
xlabel("K");
ylabel("C");
xtics(2.8,4.2,.16,0);
ytics(1.04,1.28,.04,0);


xy(krange,optcs);
nextwind;

_pmsgctl = { 3.54 .14 .2 0 1 1 0 };
_pmsgstr = "(C[*], K[*])";
_psym = { 3.53 .139326 7 7 1 1 0};
ylabel("dC/dK");
_plegctl = 1;
_plegstr = "Low State\000High State";?;
ytics(.12,.16,.01,0);

xy(krange[1:rows(dcdkl)],dcdkl~dcdkh);
endwind;





/* collect procedures here */

proc steady(p);
    /* calculates the steady state of the model */
    local alp,bet,del,kss,yss,iss,css;
    alp = p[1];
    bet = p[2];
    del = p[3];

    kss = (alp/(1/bet-1+del))^(1/(1-alp));
    yss = kss^alp;
    iss = del*kss;
    css = yss - iss;

    retp(yss|css|iss|kss|1);
endp;

proc markovchain(const,rho);

    /* builds a markov matrix from an ar1 spec. w/const, rho */

    local  p22, p11, markovmat;

    p22 = 1 - (1-rho)*const;
    p11 = const + rho*(1-const);

    markovmat = zeros(2,2);
    markovmat[1,1] = p11;
    markovmat[1,2] = 1-p22;
    markovmat[2,1] = 1-p11;
    markovmat[2,2] = p22;

    retp(markovmat);

endp;

proc zerproc(steps,ord);
    /* locates zeros of polynomial of order ord */

    local  ntees,step,normvar,tees,j,bigind,biginddif,zers;

    ntees = ord+1; 
    /* create variable ranging from [-1,1] */
    step = 2/(steps-1);
    normvar = seqa(-1,step,steps); 
    tees = zeros(steps,ntees);
    tees[.,1] = ones(steps,1);
    tees[.,2] = normvar;
    j=3; do while j<=ntees;
        tees[.,j] = 2*normvar.*tees[.,j-1]-tees[.,j-2];
    j=j+1; endo;    
    bigind = tees[.,ntees] .> zeros(steps,1);
    biginddif = bigind[2:steps]-bigind[1:steps-1];
    zers = indexcat(abs(biginddif),1);
    retp(zers);
endp;

proc ceeofak(a,k,gam);
    /* calculates c as a function of a. spits out c, and value of polynomials over the range of k */

    local  nteesa,nteesk,teesa,teesk,j,tmat,gammat,i,cee;

    nteesa = orda;
    nteesk = ordk; 
    teesa = zeros(nteesa,1);
    teesk = zeros(nteesk,1);
    teesa[1] = 1;
    teesa[2] = a;
    teesk[1] = 1;
    teesk[2] = k;
    j=3; do while j<=nteesa;
        teesa[j] = 2*a*teesa[j-1]-teesa[j-2];
    j=j+1; endo;    
    j=3; do while j<=nteesk;
        teesk[j] = 2*k*teesk[j-1]-teesk[j-2];
    j=j+1; endo;    

    tmat = teesk*teesa'; 
    gammat = zeros(nteesk,nteesa);
    i=1; do while i<=nteesa;
        gammat[.,i] = gam[(i-1)*nteesk+1:i*nteesk]; 
    i=i+1; endo;
    
    cee = sumc(sumc(tmat.*gammat));

    retp(cee);
endp;

proc feval(gam);
    local nrowsa,nrowsk,ncees,cee,f,i,j,ap,kpl,kplnorm,kph,kphnorm,cpl,cph,
        rhsl,rhsh,
kp,kpnorm,apnorm,cp,rhs;

nrowsa = rows(zersa);
nrowsk = rows(zersk); 
ncees = nrowsa*nrowsk;
cee = zeros(ncees,1);
f = zeros(ncees,1);

i=1; do while i<=nrowsa;
    j=1; do while j<=nrowsk; i~j; anorm[zersa[i]]; 
    prob = pmat*x[.,i];
    cee[(i-1)*nrowsk+j] = ceeofak(anorm[zersa[i]],knorm[zersk[j]],gam); 
    ap = (al~ah)*prob;
    kp = arange[zersa[i]]*(krange[zersk[j]]^alp)-cee[(i-1)*nrowsk+j]+(1-del)*krange[zersk[j]]; 
    apnorm = (ap-ass)/(omegaa*ass);
    kpnorm = (kp-kss)/(omegak*kss);
    cp = ceeofak(apnorm,kpnorm,gam);

    rhs = maxc(( bet*(alp*(ap*(kp^(alp-1)))+1-del)*(cp^(-phi)) )|(ap*kp^(-alp*phi)));

    f[(i-1)*nrowsk+j] = cee[(i-1)*nrowsk+j]^-phi - rhs;

/*
    kpl = arange[zersa[1]]*(krange[zersk[j]]^alp)-cee[(i-1)*nrowsk+j]+(1-del)*krange[zersk[j]]; 
    kplnorm = (kpl-kss)/(omegak*kss);
    kph = arange[zersa[2]]*(krange[zersk[j]]^alp)-cee[(i-1)*nrowsk+j]+(1-del)*krange[zersk[j]]; 
    kphnorm = (kph-kss)/(omegak*kss);
    cpl = ceeofak(anorm[1],kplnorm,gam);
    cph = ceeofak(anorm[2],kphnorm,gam);

    rhsl = maxc(( bet*(alp*(al*kpl^(alp-1))+1-del)*(cpl^(-phi)) )|(al*kpl^(-alp*phi)));
    rhsh = maxc(( bet*(alp*(ah*kph^(alp-1))+1-del)*(cph^(-phi)) )|(ah*kph^(-alp*phi)));

    f[(i-1)*nrowsk+j] = cee[(i-1)*nrowsk+j]^-phi - (rhsh~rhsl)*prob;
*/
    j=j+1; endo;
i=i+1; endo;

    retp(f);

endp;