@ ------------------------------ CoRRAM_5.src ---------------------------------

    28 December 2010
    Alfred Maussner

    The last piece of code of the DSGE toolbox.
    
    This file contains the Gauss routines that compute the coefficients of
    a second order approximation of a SDGE model given by the
    system of equations E(g^i[x,y,z,x',y',z'])=0 as described
    in Heer and Maussner, 2nd edition, Chapter 2.6.


    - SolveQA and related procedures

--------------------------------------------------------------------------- @


/* SolveQA: computates Hxi and Hyj 
**
**
** usage: {xcube,ycube}=SolveQA(Lxx,Lxz,Lyx,Lyz,gmat,hcube);
**
**
** input: Lxx   : n(x) by n(x) matrix, the linear part of the policy function for x'
**                                     with respect to x
**        Lxz   : n(x) by n(z) matrix, the linear part of the policy function for x'
**                                     with respect to z
**        Lyx   : n(y) by n(x) matrix, the linear part of the policy function for y,
**                                     with respect to x,
**        Lyz   : n(y) by n(z) matrix, the linear part of the policy function for y,
**                                     with respect to z.
**
**      gmat    : n(x)+n(y) by 2*(n(x)+n(y)+n(z)) matrix, the Jacobian of g
**
**    hcube     : a three dimensional array, of size n(x)+n(y) (first dimension),
**                2*(n(x)+n(y)+n(z)) (second dimsion) and 2*(n(x)+n(y)+n(z))  
**                (third dimension). Each hcube[i,.,.,] stores the Hesse matrix of
**                equation gi, i=1, 2, ..., n(x)+n(y)
**
**    Rmat      : n(z) by n(z) matrix of AR(1)-coefficients from z'=rmat*z+sigma*omat*eps
**
**    Omat      : n(z) by (z) matrix  (see line above)
**
** output: xcube : n(x) by (n(x)+n(z)+1) by (n(x)+n(z)+1) array. Each page
**                                       holds the quadratic coefficients of
**                                       the policy function for x'[i] with respect
**                                       to the vector [xbar,zbar,sigma].
**
**         ycube : n(y) by (n(x)+n(z)+1) by (n(x)+n(z)+1) array. Each page
**                                       holds the quadratic coefficients of
**                                       the policy function for y[j] with respect
**                                       to the vector [xbar,zbar,sigma].
*/

proc(2)=SolveQA(Lxx,Lxz,Lyx,Lyz,gmat,hcube,Rmat,Omat);

    local _nx, _ny, _nz, _nall, indz, inds, temp1, temp2,
          amat, bvec, hvecxj, hvecxk,hveczj, hveczk, 
          xvec1,xvec2,xvec3, xvec4, xcube, ycube,
          zi, i, j, k, l1, l2, l3, ord;

/* determine dimenions */
_nx=cols(Lxx);
_nz=cols(Lxz);
_ny=rows(Lyx);

/* Compute Hxx */
_nall=(_nx+_ny)*(_nx^2);
amat=zeros(_nall,_nall);
bvec=zeros(_nall,1);

zi=0;
for j (1,_nx,1);
    for k (1,_nx,1);
        for i (1,_nx+_ny,1);
            zi=zi+1;
            for l1 (1,_ny,1);
                amat[zi,iy1(l1,j,k)]=gmat[i,_nx+l1];
            endfor;
            for l1 (1,_nx,1);
                amat[zi,ix1(l1,j,k)]=gmat[i,_nx+_ny+_nz+l1];
            endfor;
            for l1 (1,_ny,1);
                for l2 (1,_nx,1);
                    amat[zi,ix1(l2,j,k)]=amat[zi,ix1(l2,j,k)]+gmat[i,2*_nx+_nz+_ny+l1]*Lyx[l1,l2];
                endfor;
            endfor;
            for l1 (1,_ny,1);
                for l2 (1,_nx,1);
                    for l3 (1,_nx,1);
                        amat[zi,iy1(l1,l2,l3)]=amat[zi,iy1(l1,l2,l3)]+gmat[i,2*_nx+_nz+_ny+l1]*Lxx[l2,j]*Lxx[l3,k];
                    endfor;
                endfor;
            endfor;
            hvecxj=1|Lyx[.,j]|Lxx[.,j]|(Lyx*Lxx[.,j]);
            hvecxk=1|Lyx[.,k]|Lxx[.,k]|(Lyx*Lxx[.,k]);
            indz=j|seqa(_nx+1,1,_ny)|seqa(_nx+_ny+_nz+1,1,_nx+_ny);
            inds=k|seqa(_nx+1,1,_ny)|seqa(_nx+_ny+_nz+1,1,_nx+_ny);
            temp1=arraytomat(hcube[i,indz,inds]);            
            bvec[zi]=hvecxj'*temp1*hvecxk;
        endfor;
    endfor;
endfor;    

xvec1=-bvec/amat;


/* hx_xz und hy_xz */

_nall=(_nx+_ny)*(_nx*_nz);
amat=zeros(_nall,_nall);
bvec=zeros(_nall,1);

zi=0;
for j (1,_nx,1);
    for k (1,_nz,1);
        for i (1,_nx+_ny,1);
            zi=zi+1;
            for l1 (1,_ny,1);
                amat[zi,iy2(l1,j,k)]=gmat[i,_nx+l1];
            endfor;
            for l1 (1,_nx,1);
                amat[zi,ix2(l1,j,k)]=gmat[i,_nx+_ny+_nz+l1];
            endfor;
            for l1 (1,_ny,1);
                for l2 (1,_nx,1);
                    amat[zi,ix2(l2,j,k)]=amat[zi,ix2(l2,j,k)]+gmat[i,2*_nx+_ny+_nz+l1]*Lyx[l1,l2];
                endfor;
            endfor;
            for l1 (1,_ny,1);
                for l2 (1,_nx,1);
                    for l3 (1,_nz,1);
                        amat[zi,iy2(l1,l2,l3)]=amat[zi,iy2(l1,l2,l3)]+gmat[i,2*_nx+_ny+_nz+l1]*Lxx[l2,j]*Rmat[l3,k];
                    endfor;
                    for l3 (1,_nx,1);                        
                        bvec[zi]=bvec[zi]+gmat[i,2*_nx+_ny+_nz+l1]*xvec1[iy1(l1,l2,l3)]*Lxx[l2,j]*Lxz[l3,k];
                    endfor;
                endfor;
            endfor;
            hvecxj=1|Lyx[.,j]|Lxx[.,j]|(Lyx*Lxx[.,j]);
            hveczk=Lyz[.,k]|Lxz[.,k]|(Lyx*Lxz[.,k]+Lyz*Rmat[.,k])|1|Rmat[.,k];
            indz=j|seqa(_nx+1,1,_ny)|seqa(_nx+_ny+_nz+1,1,_nx+_ny);
            inds=  seqa(_nx+1,1,_ny)|seqa(_nx+_ny+_nz+1,1,_nx+_ny)|(_nx+_ny+k)|seqa(2*_nx+2*_ny+_nz+1,1,_nz);
            temp1=arraytomat(hcube[i,indz,inds]);
@            temp2=hvecxj'*temp1*hveczk;            @
            bvec[zi]=bvec[zi]+hvecxj'*temp1*hveczk;            
        endfor;
    endfor;
endfor;                
xvec2=(-bvec)/amat;
save xvec2;
/* hx_zz und hy_zz */

_nall=(_nx+_ny)*(_nz^2);

amat=zeros(_nall,_nall);
bvec=zeros(_nall,1);

zi=0;
for j (1,_nz,1);
    for k (1,_nz,1);
        for i (1,_nx+_ny,1);
            zi=zi+1;
            for l1 (1,_ny,1);
                amat[zi,iy3(l1,j,k)]=gmat[i,_nx+l1];
            endfor;
            for l1 (1,_nx,1);
                amat[zi,ix3(l1,j,k)]=gmat[i,_nx+_ny+_nz+l1];
            endfor;
            for l1 (1,_ny,1);
                for l2 (1,_nx,1);
                    amat[zi,ix3(l2,j,k)]=amat[zi,ix3(l2,j,k)]+gmat[i,2*_nx+_ny+_nz+l1]*Lyx[l1,l2];
                endfor;
            endfor;
            for l1 (1,_ny,1);
                for l2 (1,_nz,1);
                    for l3 (1,_nz,1);
                        amat[zi,iy3(l1,l2,l3)]=amat[zi,iy3(l1,l2,l3)]+gmat[i,2*_nx+_ny+_nz+l1]*Rmat[l2,j]*Rmat[l3,k];
                    endfor;
                    for l3 (1,_nx,1);                        
                        bvec[zi]=bvec[zi]+gmat[i,2*_nx+_ny+_nz+l1]*Rmat[l2,j]*xvec2[iy2(l1,l3,l2)]*Lxz[l3,k];
                    endfor;
                endfor;
                for l2 (1,_nx,1);
                    for l3 (1,_nx,1);
                        bvec[zi]=bvec[zi]+gmat[i,2*_nx+_ny+_nz+l1]*Lxz[l2,j]*xvec1[iy1(l1,l2,l3)]*Lxz[l3,k];
                    endfor;
                    for l3 (1,_nz,1);
                        bvec[zi]=bvec[zi]+gmat[i,2*_nx+_ny+_nz+l1]*Lxz[l2,j]*xvec2[iy2(l1,l2,l3)]*Rmat[l3,k];
                    endfor;
                endfor;
            endfor;
            hveczj=Lyz[.,j]|Lxz[.,j]|(Lyx*Lxz[.,j]+Lyz*Rmat[.,j])|1|Rmat[.,j];
            hveczk=Lyz[.,k]|Lxz[.,k]|(Lyx*Lxz[.,k]+Lyz*Rmat[.,k])|1|Rmat[.,k];
            indz=seqa(_nx+1,1,_ny)|seqa(_nx+_ny+_nz+1,1,_nx+_ny)|(_nx+_ny+j)|seqa(2*(_nx+_ny)+_nz+1,1,_nz);
            inds=seqa(_nx+1,1,_ny)|seqa(_nx+_ny+_nz+1,1,_nx+_ny)|(_nx+_ny+k)|seqa(2*(_nx+_ny)+_nz+1,1,_nz);
            temp1=arraytomat(hcube[i,indz,inds]);
            ord=hveczj'*temp1*hveczk;
            bvec[zi]=bvec[zi]+hveczj'*temp1*hveczk;
        endfor;
    endfor;
endfor;                
xvec3=(-bvec)/amat;

/* Hss */
amat=zeros(_nx+_ny,_nx+_ny);
bvec=zeros(_nx+_ny,1);

for i (1,_nx+_ny,1);
    for j (1,_ny,1);
        amat[i,_nx+j]=gmat[i,_nx+j]+gmat[i, 2*_nx+_ny+_nz+j];    
    endfor;
    for j (1,_nx,1);
        amat[i,j]=gmat[i,_nx+_ny+_nz+j];
        for l1 (1,_ny,1);
            amat[i,j]=amat[i,j]+gmat[i,2*_nx+_ny+_nz+l1]*Lyx[l1,j];
        endfor;
    endfor;
    for j (1,_ny,1);
        for l1 (1,_nz,1);
            for l2 (1,_nz,1);
                bvec[i]=bvec[i]+gmat[i,2*_nx+_ny+_nz+j]*xvec3[iy3(j,l1,l2)]*Omat[l1,.]*Omat[l2,.]';
            endfor;
        endfor;
    endfor;
    indz=seqa(2*_nx+_ny+_nz+1,1,_ny+_nz);
    temp1=arraytomat(hcube[i,indz,indz]);
    temp2=zeros(_ny+_nz,_ny+_nz);
    for l1 (1,_ny,1);
        for l2(1,_ny,1);
            temp2[l1,l2]=Geta(l1,l2,Lyz,Omat);
        endfor;
        for l2(1,_nz,1);
            temp2[l1,_ny+l2]=Getab(l1,l2,Lyz,Omat);
            temp2[_ny+l2,l1]=temp2[l1,_ny+l2];
        endfor;
    endfor;
    for l1 (1,_nz,1);       
        for l2(1,_nz,1);
            temp2[_ny+l1,_ny+l2]=Omat[l1,.]*Omat[l2,.]';
        endfor;
    endfor;
    bvec[i]=bvec[i]+sumc(diag(temp1*temp2));
endfor;

xvec4=(-bvec)/amat;

/* Create xcube and ycube */
ord=_nx|(_nx+_nz+1)|(_nx+_nz+1);
xcube=arrayinit(ord,0);
ord=_ny||(_nx+_nz+1)|(_nx+_nz+1);
ycube=arrayinit(ord,0);

for i (1,_nx,1);
    xcube[i,1:_nx,1:_nx]=reshape(xvec1[ix1(i,1,1):ix1(i,_nx,_nx)],_nx,_nx);
    xcube[i,1:_nx,_nx+1:_nx+_nz]=reshape(xvec2[ix2(i,1,1):ix2(i,_nx,_nz)],_nx,_nz);
    temp1=arraytomat(xcube[i,1:_nx,_nx+1:_nx+_nz]);
    xcube[i,_nx+1:_nx+_nz,1:_nx]=temp1';
    xcube[i,_nx+1:_nx+_nz,_nx+1:_nx+_nz]=reshape(xvec3[ix3(i,1,1):ix3(i,_nz,_nz)],_nz,_nz);
    xcube[i,_nx+_nz+1,_nx+_nz+1]=xvec4[i];  
endfor;
for i (1,_ny,1);
    ycube[i,1:_nx,1:_nx]=reshape(xvec1[iy1(i,1,1):iy1(i,_nx,_nx)],_nx,_nx);
    ycube[i,1:_nx,_nx+1:_nx+_nz]=reshape(xvec2[iy2(i,1,1):iy2(i,_nx,_nz)],_nx,_nz);
    temp1=arraytomat(ycube[i,1:_nx,_nx+1:_nx+_nz]);
    ycube[i,_nx+1:_nx+_nz,1:_nx]=temp1';
    ycube[i,_nx+1:_nx+_nz,_nx+1:_nx+_nz]=reshape(xvec3[iy3(i,1,1):iy3(i,_nz,_nz)],_nz,_nz);
    ycube[i,_nx+_nz+1,_nx+_nz+1]=xvec4[_nx+i];
endfor;


retp(xcube,ycube);

endp;

/* Indexfunctions used by SolveQA: */

fn ix1(i,j,k)=(i-1)*(_nx^2) + (j-1)*_nx + k;
fn iy1(i,j,k)=(_nx^3) + ix1(i,j,k);

fn ix2(i,j,k)=(i-1)*(_nz*_nx) + (j-1)*_nz + k;
fn iy2(i,j,k)=(_nx^2)*_nz + ix2(i,j,k);

fn ix3(i,j,k)=(i-1)*(_nz^2)+(j-1)*_nz + k;
fn iy3(i,j,k)=_nx*(_nz^2) + ix3(i,j,k);

/* Two other procedures used by SolveQA to compute Hss */

proc(1)=Geta(i,j,Lyz,Omat);
    local fx, n, q, s;

    n=cols(Lyz);
    fx=0;
    for s (1,n,1);
        for q (1,n, 1);
            fx=fx+Lyz[i,q]*Lyz[j,s]*(Omat[q,.]*Omat[s,.]');
        endfor;
    endfor;
    retp(fx);
endp;

proc(1)=Getab(i,j,Lyz,Omat);
    local fx, n, s;

    fx=0;
    n=cols(Lyz);
    for s (1,n,1);
        fx=fx+Lyz[i,s]*(Omat[s,.]*Omat[j,.]');
    endfor;
    retp(fx);
endp;


