/*      PEALIB.SRC:  GAUSS procedures needed to run
**      the programs PEAPRO1.PRG and PEAPROI1.PRG
**
**      November 9 1998
**
*/

proc(1) = gridfun(psi,xgrid);

/*
**      This procedure is only needed to run "equation-solver" PEA
*/


external matrix po_k po_t alpha delta nu dfactor tmin tmax kmin kmax;
local    dimgrid ,yvec ,y, i, lnk1, lnt1, lnk2;

/*
**       3.1 Compute the expected value (Step II)
**       ----------------------------------------
**       Here we calculate the conditional expectation at each grid point.
**       First, we calculate next period's capital stock as a function of
**       the log of the current capital stock and the log of the technology
**       shock. Second, we numerically integrate the function aexp over the
**       random innovation to the technology shock. psi is an input of the
**       numerical integration since it is used to evaluate the consumption
**       function at lnk2 and lnt2.
*/

dimgrid = rows(xgrid);
yvec    = zeros(dimgrid,1);
y       = yvec;
i       = 1;
do while i <= dimgrid;
    lnk1      = scalup(xgrid[i,1],kmin,kmax);
    lnt1      = scalup(xgrid[i,2],tmin,tmax);
    lnk2      = savfun(lnk1,lnt1,psi);
    yvec[i]   = numi(&aexp,lnk2,lnt1,psi);
    i       = i + 1;
endo;
y   = exp( makepoly((po_k~po_t),(xgrid[.,1]~xgrid[.,2])) * psi );
y   = y - yvec;

retp(y);
endp;

/* **********************************************************************

** **********************************************************************
*/
/*
** x=chebpol(ord,p);
** Purpose
** Create an ord-th order Chebyshev Polynomial
** if p is a vector then the polynomial will be created based
**  on that vector
**
** ----------------------------------------------------------------------
*/
proc(1) = chebpol(ord,p);

local r,x,i;

    r   = rows(p);
    x   = ones(r,ord+1);
    x[.,2]= p;
    i = 3;
    if ord > 1;
       do while i <= (ord+1);
           x[.,i] = 2.*p.*x[.,i-1]-x[.,i-2];
           i    = i + 1;
       endo;
    endif;

retp(x);
endp;

/*
**********************************************************************

**********************************************************************
*//*
** x=chebnode(n);
** Purpose
** Create n Chebyshev nodes
**
** ----------------------------------------------------------------------
*/

proc(1) = chebnode(n);

local r,x;

r       = maxc(1|n);
n       = seqa(1,1,n);
x       = cos( (pi*(n-0.5))/r );

retp(x);
endp;
/*
** **********************************************************************

** **********************************************************************
*//*
** xmat=makepoly(ord,x);
**
**            For the (Tx2) matrix x and the (1x2) vector ord
**            this creates a [T x ord(1)*ord(2)] matrix where
**            each row contains the ord(1)*ord(2) bivariate Chebyshev
**            polynomial terms.
**
** ------------------------------------------------------------------
*/

proc(1) = makepoly(ord,x);

local   po_k,po_t,xk,xt,kj,tj,xmat,r;

po_k    = ord[1];
po_t    = ord[2];
xk      = chebpol(po_k,x[.,1]);
xt      = chebpol(po_t,x[.,2]);
r     = rows(x);
xmat      = zeros(r,(po_k+1)*(po_t+1));
kj    = 0;
do while kj <= po_k;
    tj    = 0;
    do while tj <= po_t;
        xmat[.,kj*(po_t+1)+tj+1] = xk[.,kj+1].*xt[.,tj+1];
        tj               = tj + 1;
    endo;
    kj  = kj + 1;
endo;

retp(xmat);
endp;

/*
**********************************************************************

**********************************************************************
*/

/*
** function y = scalup(x,xmin,xmax);
** Linearly rescales every element of the vector x
** from [-1,1]  to  [xmin,xmax]
** where x,xmin,xmax can be vectors as long as they are all of the same
** dimension
*/

proc(1) = scalup(x,xmin,xmax);
local r,c,a,b,xu;

r   = rows(x);
c   = cols(x);
a     = (xmin+xmax)'/2;
b     = (xmax-xmin)'/2;
xu    = a*ones(1,r) + ( b*ones(1,r) ).*x';
xu    = xu';

retp(xu);
endp;

/*
**********************************************************************

**********************************************************************
*/

/*
** function xd    = scaldown(x,xmin,xmax);
** Linearly scale a variable from [min,xmax] to [-1,1],
** where x,xmin,xmax can be vectors as long as they are all of the same
** dimension
*/

proc(1) = scaldown(x,xmin,xmax);

local r,c,a,b,xd;

r   = rows(x);
c   = cols(x);
a     = 2*ones(r,c) ./ ( xmax - xmin );
b     = ones(r,c) - 2 * xmax ./ ( xmax - xmin );

xd      = b + a .* x;
retp(xd);
endp;

/*
**********************************************************************

**********************************************************************
*//*
** function c = confun(lnk1,lnt1,psi);
** Returns the LOG of consumption
*/

proc(1) = confun(lnk1,lnt1,psi);
external matrix po_k po_t nu tmin tmax kmin kmax;
local    lnk1uns,lnt1uns,xmat,c;

lnk1uns = scaldown(lnk1,kmin,kmax);
lnt1uns = scaldown(lnt1,tmin,tmax);
xmat        = makepoly((po_k~po_t),(lnk1uns~lnt1uns));
c       = -1/nu * (xmat*psi);
retp(c);
endp;

/*
*********************************************************************

**********************************************************************
*//*
** function k2=savfun(lnk1,lnt1,psi);
** Returns the LOG of new capital
*/

proc(1) = savfun(lnk1,lnt1,psi);
external matrix po_k po_t alpha delta nu dfactor tmin tmax kmin kmax;
local c1,k1,t1,k2;

c1      = exp(confun(lnk1,lnt1,psi));
k1      = exp(lnk1);
t1      = exp(lnt1);

k2      = t1.*(k1.^alpha) + (1-delta)*k1 - c1;
if k2   < 0;
    errorlog "Consumption too large.";
    end;
endif;

k2      = ln(k2);
retp(k2);
endp;

/*
**********************************************************************

**********************************************************************
*/

/*
** function aex=aexp(epsi,lnk2,lnt1,psi);
** Compute the actual marginal utility
*/

proc(1) = aexp(epsi,lnk2,lnt1,psi);

external matrix alpha dfactor delta nu rho sigma;
local r,c,lnt2,c2,k2,t2,x,aex;

r           = rows(lnt1);
c           = cols(lnt1);
lnt2        = rho * lnt1 + epsi';
lnk2        = lnk2 .* ones(r,c);
c2          = exp(confun(lnk2,lnt2,psi));
k2          = exp(lnk2);
t2          = exp(lnt2);

x           = alpha * t2 .* ( k2 .^ ( alpha - 1 ) ) + 1 - delta;
aex         = dfactor * (c2 .^ ( -nu )) .* x;

retp(aex);
endp;

/*
**********************************************************************

**********************************************************************
*/
/*
** function yvec=numi(fun,lnk2,lnt1,psi);
** Calculate the expected value of E[fun(epsi)] using Hermite Gaussian quadratur
e,
** when epsi is a normally distributed random variable with standard deviation s
igma.
*/

proc(1) = numi(&fun,lnk2,lnt1,psi);

external matrix alpha dfactor nu delta sigma rho po_k po_t;
local tnod,nodet,valu,hi,yvec,nodes,weight;
local   fun :proc;

/* Define the Nodes and weights for integration: */

{nodes, weight} = hernodes(11);
tnod            = 11;

/* -------------------------------------------------------

** Note that Hermite Gaussian quadrature uses the weighting function
** exp{-x^2}. Here, however, we have the density
** (1./sigma*sqrt(2*pi))*exp{-y^2/2*sigma^2} and want to calculate the
** integral of fun(y)*(1./sigma*sqrt(2*pi))*exp{-y^2/2*sigma^2}
**
** Adjusting the quadrature procedure requires the following two steps.
** First a change of variables is introduced. In particular, y=sqrt(2)*sigma*x.
** This results in the weighting function (1/sigma*sqrt(2*pi))*exp{-x^2}.
** Second, we have to multiply the integrand with the Jacobian of the
** transformation. This gives the weighting function (1/sqrt(pi))*exp{-x^2}
** and calculating the integral of fun(sqrt(2)*sigma*x)*(1/sqrt(pi))*exp{-x^2}.
*/

nodet = sigma*sqrt(2)*nodes;

/* Evaluate the function at the transformed nodes */

valu  = zeros(tnod,1);

hi    = 1;
do while hi <= tnod;
/*  valu(hi)    = feval(fun,nodet(hi),lnk2,lnt1,psi); */
    valu[hi]    = fun(nodet[hi],lnk2,lnt1,psi);
    hi          = hi + 1;
endo;

/* Compute the integral  */
yvec    = 1/sqrt(pi) * weight'*valu;

retp(yvec);
endp;

/*
**********************************************************************

**********************************************************************
*/

/*
** function [x,w]   = hernodes(n);
**
** Purpose:
** Compute the weights and nodes for Hermite-Polynomials
**
** Inputs:
** n    The number of nodes to compute
**
** Outputs:
** x    The nodes
** w    the associated weights
**
** Source:
** Taken from Numerical Recipes in Fortran77, 2ed
**
** -----------------------------------------------------------------
*/

proc(2) = hernodes(n);

local   TOL,MAXIT,PIM4,m,x,w,i,its,z,p1,p2,p3,pp,j,z1;

TOL     = 10^-8;
MAXIT   = 30;
PIM4        = pi^(-1/4);

m       = (n+1)/2;
x       = zeros(n,1);
w       = zeros(n,1);

i   = 1;
do while i <= m;

    /* Initialize the first four roots
    ** ------------------------------- */
    if i    == 1;
        z   = sqrt(2*n+1) - 1.85575 * (2*n+1)^(-1/6);
    elseif i== 2;
        z   = z - 1.14 * (n^.426)/z;
    elseif i== 3;
        z   = 1.86 * z - .86 * x[1];
    elseif i== 4;
        z   = 1.91 * z - .91 * x[2];
    else;
        z   = 2 * z - x[i-2];
    endif;

    its = 1;
    do while its <= MAXIT;
        p1  = PIM4;
        p2  = 0;
        j   = 1;
        do while j <= n;
            p3  = p2;
            p2  = p1;
            p1  = z*sqrt(2/j)*p2 - sqrt((j-1)/j)*p3;
            j   = j + 1;
        endo;

        pp  = p2 * sqrt(2*n);
        z1  = z;
        z   = z1 - p1/pp;
        if abs(z-z1) < TOL; break; endif;
        its = its + 1;
    endo;

    x[i]        = z;
    x[n+1-i]    = -z;
    w[i]        = 2/(pp*pp);
    w[n+1-i]    = w[i];
    i       = i + 1;
endo;

retp(x,w);
endp;

/*
**********************************************************************

**********************************************************************
*/

proc(2) = nSecant( Fname,x0,noprnt,points );
external matrix alpha dfactor nu delta rho sigma po_k po_t tmin tmax kmin kmax;
external matrix maxiter psitol;
local    n, m, itnlimt, tol, h, x, F, normF, k, e, xd, Fd, Y, S, P, xnew;
local    j, ei, Fnew, status, PP;
local    fname :proc;

/*
**       [x,F,J] = nSecant( Fname,x0 ) finds  x  such that
**        F(x) = 0 for a vector of nonlinear functions  F(x).
**        The zero is found using the (n+1)-point secant method.
**
**        Fname is a string containing the name of an m-file that
**        computes the vector F.  x0 is an initial guess at the
**        solution.  For example, secant( "my_F",x0 ) finds a
**        zero of the nonlinear function defined in the mfile
**        my_F.m.  Note the quotes around "my_F".
**
**        On successful termination, x is the first point found at
**        which the norm of F is less than sqrt(eps).  Output
**        argument is   F(x) at the final x.
**
**        The meaning of each line of printed output is as follows
**
**        Itn          The current iteration number.
**
**        Norm F       The current value of the norm of F(x).
**                     In the neighborhood of a zero of F,
**                     Norm F will be small.
**
**        This version dated 07-Oct-1995.
**        Philip E. Gill, University of California, San Diego.
**
**        This version dated 01-May-1998.
**   Modified by Christian Haefke to suit the problem
**        Christian Haefke, University of California, San Diego.
** -------------------------------------------------------------------
*/


n         = rows(x0);
m         = cols(x0);
itnlimt = maxiter;      /* Limit on evaluations of F.   */
tol       = psitol;       /* Required accuracy in norm F. */

h         = 10^-8;        /* Rule-of-thumb optimal h.       */
k         = 0;
x         = x0;
F         = Fname(x,points);
normF     = sqrt(F'F);

/* Start with the finite-difference Jacobian. */

j         = 0;
ei        = 1;
Y         = zeros(rows(F),n);
do while ei  <= n;
   e      = zeros(n,1);
   e[ei]  = 1;
   xd     = x + h*e;
   Fd     = Fname(xd,points);
   j      = j + 1;
   Y[.,j] = Fd - F;
   ei     = ei + 1;
endo;

S         = h*eye(n);
PP        = seqa(1,1,n)';
PP        = PP + ones(1,cols(PP));
PP[n]     = 1;

if noprnt == 0;
   cls;
   print "\t Itn   \t  Norm F";
   print  k~normF;
endif;

/* while  normF >= tol  &  k < itnlimt */

k   = 0;
do while k < itnlimt;
   k       = k + 1;
   p       = - S*(F/Y);
   xnew    = x + p;
   Fnew  = Fname(xnew,points);

   S[.,n]  = xnew - x;
   S       = S[.,PP];
   Y[.,n]  = Fnew - F;
   Y       = Y[.,PP];

   x       = xnew;
   F       = Fnew;

   normF   = sqrt(F'F);
   if normF < tol; break; endif;

   if noprnt == 0;
      print k~normF;
   endif;
endo;

if     normF <= tol;
   status = "Solved    ";
else;
   status = "k too big ";
endif;

if     status $== "Solved    ";
   print "Solved at iteration ";; k;
   print "Norm f = ";; normF;
elseif status $== "k too big ";
   print "Secant terminated.  Iteration limit exceeded.";
endif;

retp(x,F);
endp;

/*
**********************************************************************

**********************************************************************
*/

