{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Projection implementation of stochastic growth model adapted from Fabrice Collard's Matlab code, http://fabcol.free.fr/\n",
    "# tested in Julia 0.6.1.1\n",
    "# this code is part of chapter 5, \"Advanced Numerical Techniques\" from the book: \"Introduction to Quantitative Macroeconomics using Julia\"\n",
    "# Academic Press - Elsevier\n",
    "# for comments, email at: petre(dot)caraiani(at)gmail(dot)com\n",
    "\n",
    "# The Optimal Growth Model\n",
    "# Galerkin method (AR(1) case)                 \n",
    "\n",
    "global nbk, kmin, ksup, XK, kt;\n",
    "global nba, amin, asup, XA, at;\n",
    "global nstate, nodea, nodek, wmat, wgrid;\n",
    "\n",
    "# Parameters                        \n",
    "nbk   = 4;    # Degree of polynomials (capital)\n",
    "nba   = 2;    # Degree of polynomials (technology shock)\n",
    "ncoef = (nbk+1); # of coefficients\n",
    "nodek = 20;   # of Nodes\n",
    "nodea = 10;# of Nodes\n",
    "nstate= 12;\n",
    "                                                           \n",
    "# Structural Parameters                    \n",
    "delta  = 0.1;\n",
    "beta   = 0.95;\n",
    "alpha  = 0.3;\n",
    "sigma  = 1.5;\n",
    "ysk    =(1-beta*(1-delta))/(alpha*beta);\n",
    "ksy    = 1/ysk;\n",
    "ys     = ksy^(alpha/(1-alpha));\n",
    "ks     = ys^(1/alpha);\n",
    "is     = delta*ks;\n",
    "cs     = ys-is;\n",
    "ab     = 0;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "function hernodes(nstate)\n",
    "\n",
    "TOL    = sqrt(2.2204e-16);\n",
    "MAXIT  = 30;\n",
    "PIM4   = pi^(-1/4);\n",
    "\n",
    "n  = nstate;\n",
    "m  = (n+1)/2;\n",
    "m  =  convert(Int64, floor(m));\n",
    "x  = zeros(n,1);\n",
    "w  = zeros(n,1);\n",
    "\n",
    "z  = 0;\n",
    "z1 = 0;\n",
    "pp = 0;\n",
    "for i= 1:m\n",
    "\n",
    "# Initialize the first four roots\n",
    "if i == 1; \n",
    "   z= sqrt(2*n+1) - 1.85575 * (2*n+1)^(-1/6);\n",
    "elseif i== 2;\n",
    "   z= z - 1.14 * (n^.426)/z;\n",
    "elseif i== 3;\n",
    "   z= 1.86 * z - .86 * x[1];\n",
    "elseif i== 4;\n",
    "   z= 1.91 * z - .91 * x[2];\n",
    "else;\n",
    "   z= 2 * z - x[i-2];\n",
    "end;\n",
    "\n",
    "for its= 1:MAXIT\n",
    "    p1 = PIM4;\n",
    "    p2 = 0;\n",
    " for j  = 1:n\n",
    "    p3 = p2;\n",
    "    p2 = p1;\n",
    "    p1 = z*sqrt(2/j)*p2 - sqrt((j-1)/j)*p3;\n",
    " end;\n",
    "\n",
    "    pp = p2 * sqrt(2*n);\n",
    "    z1 = z;\n",
    "    z  = z1 - p1/pp;\n",
    "if abs(z-z1) < TOL; break; end;\n",
    "end;\n",
    "\n",
    "x[i]    = z;\n",
    "x[n+1-i]= -z;\n",
    "w[i]= 2/(pp*pp);\n",
    "w[n+1-i]= w[i];\n",
    "end;       \n",
    "        \n",
    "return x,w\n",
    "    end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Technology shock process\n",
    "rho    = 0.8;\n",
    "se     = 0.2;\n",
    "ma     = 0;\n",
    "wgrid,wmat=hernodes(nstate);\n",
    "wgrid=wgrid*sqrt(2)*se;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#define functions rcheb transfo itransfo\n",
    "\n",
    "function rcheb(nn);\n",
    "mod=nn-floor(nn/2)*2;\n",
    "n1=floor(nn/2);\n",
    "k=collect(1:n1).';\n",
    "r1=cos.((2*k-1)*pi/(2*nn));\n",
    "r1=[r1 -r1];\n",
    "if mod==1;\n",
    "  r1=[r1 0];\n",
    "end;\n",
    "r1=vec(r1);\n",
    "rr=real(sort!(r1));\n",
    "return rr\n",
    "end;\n",
    "\n",
    "function transfo(x,xmin,xmax);\n",
    "  z=(2*(x-xmin)/(xmax-xmin))-1;  \n",
    "return z\n",
    "end;\n",
    "\n",
    "\n",
    "function itransfo(x,xmin,xmax);\n",
    "z=0.5*(x+1)*(xmax-xmin)+xmin;\n",
    "return z\n",
    "end;\n",
    "\n",
    "function cheb(xx,nn);\n",
    "cc=real(cos.(kron(acos.(complex(xx)),nn)))       \n",
    "return cc\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grid for the income\n",
    "amin=(ma+wgrid[nstate]);\n",
    "asup=(ma+wgrid[1]);\n",
    "ra=rcheb(nodea);           #roots\n",
    "at=itransfo(ra,amin,asup); #grid\n",
    "vnba = collect(0:nba).';\n",
    "XA=cheb(ra,vnba);       # Polynomials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grid for the capital stock\n",
    "kmin  = log(1.2);\n",
    "ksup  = log(6);\n",
    "rk    = rcheb(nodek);               #roots\n",
    "kt    = exp.(itransfo(rk,kmin,ksup)) #grid\n",
    "vrbk = collect(0:nbk).';\n",
    "XK    = cheb(rk,vrbk);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initial conditions\n",
    "a0=[ \n",
    "  -0.23759592487257;\n",
    "   0.60814488103911;\n",
    "   0.03677400318790;\n",
    "   0.69025680170443;\n",
    "  -0.21654209984197;\n",
    "   0.00551243342828;\n",
    "   0.03499834613714;\n",
    "  -0.00341171507904;\n",
    "  -0.00449139656933;\n",
    "   0.00285737302122;\n",
    "  -0.00002348542016;\n",
    "  -0.00011606672164;\n",
    "];\n",
    "param=[alpha beta delta sigma ma rho]';"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function makepoly\n",
    "function makepoly(XA,XW);\n",
    "nba   = size(XA,2);\n",
    "nbw   = size(XW,2);\n",
    "nmax  = max(nba,nbw);\n",
    "XX    = Float64[]  ;  \n",
    "    \n",
    "for i=1:nbw\n",
    "   for j=1:nba\n",
    "      if i+j-2<=(nmax-1);\n",
    "         XX=[XX; kron(XW[:,i],XA[:,j])];\n",
    "      end\n",
    "   end\n",
    "end\n",
    "return XX\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "f  =function(theta);\n",
    "RHS=Float64[];\n",
    "LHS=Float64[];\n",
    "XX =Float64[];\n",
    "XX0=Float64[];    \n",
    "ct = 0.0;    \n",
    "c1 = 0.0;    \n",
    "for i   = 1:nodek\n",
    "    for j=1:nodea\n",
    "        XX0  = makepoly(XA[j,:]',XK[i,:]');\n",
    "        ct   = exp.(dot(XX0,theta));\n",
    "        XX   = [XX;XX0];\n",
    "        k1   = exp.(at[j])*kt[i].^alpha+(1-delta)*kt[i]-ct;\n",
    "        rk1  = transfo(log.(k1),kmin,ksup);\n",
    "        #if abs(rk1)>1;disp('problème k(t+1)');end\n",
    "        vrbk = collect(0:nbk).';    \n",
    "        xk1  = cheb(complex(rk1),vrbk);      \n",
    "        a1   = rho*at[j]+(1-rho)*ma+wgrid;\n",
    "        ra1  = transfo(a1,amin,asup);           # grid\n",
    "        #if abs(ra1)>1;disp('problème a(t+1)');end\n",
    "        XA1  = cheb(complex(ra1),vnba);              # Polynomials\n",
    "        XX1  = makepoly(XA1,xk1);\n",
    "        XX1  = reshape(XX1,size(XA1,1), size(XA1,2)*size(xk1,2)-3)    \n",
    "        c1   = exp.(XX1*theta);\n",
    "        aux  = wmat'*(beta*(alpha*exp.(a1)*k1.^(alpha-1)+1-delta).*c1.^(-sigma))/sqrt(pi);\n",
    "        RHS  = [RHS;log.(aux)];\n",
    "        LHS  = [LHS;log.(ct.^(-sigma))];\n",
    "    end;      \n",
    "end;\n",
    "XX  = reshape(XX,size(XX0,1), nodek*nodea)    \n",
    "resid     = XX*(LHS-RHS);\n",
    "return vec(resid);\n",
    "end;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mCloning cache of NLsolve from https://github.com/JuliaNLSolvers/NLsolve.jl.git\n",
      "\u001b[39m\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mInstalling NLsolve v0.12.1\n",
      "\u001b[39m\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mBuilding SpecialFunctions\n",
      "\u001b[39m\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mPackage database updated\n",
      "\u001b[39m\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mMETADATA is out-of-date — you may not have the latest version of NLsolve\n",
      "\u001b[39m\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mUse `Pkg.update()` to get the latest versions of your packages\n",
      "\u001b[39m\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mPrecompiling module NLsolve.\n",
      "\u001b[39m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Display final rule:[0.11913, 0.494021, 0.0398671, 0.361982, -0.104605, -0.00147829, 0.00829982, 0.00134686, -0.000856568, 6.58987e-5, 0.000271861, -8.72393e-7]\n"
     ]
    }
   ],
   "source": [
    "#Pkg.add(\"NLsolve\");\n",
    "using NLsolve;\n",
    "sol = nlsolve(not_in_place(f),a0);\n",
    "println(\"Display final rule:\",sol.zero)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#final\n",
    "lt  = length(sol.zero);\n",
    "nb  = 100;\n",
    "kt  = collect(kmin:(ksup-kmin)/(nb-1):ksup);\n",
    "rk  = transfo(kt,kmin,ksup);\n",
    "rk  = vec(rk);\n",
    "vnbk= collect(0:nbk).';\n",
    "XT  = cheb(rk,vnbk);\n",
    "kt  = exp.(kt);\n",
    "c=[];k1=[];ii=[];\n",
    "XX0=[];\n",
    "\n",
    "for i=1:nodea \n",
    "XX0 = makepoly(XA[i,:]',XT)\n",
    "    XX0 = reshape(XX0,size(XT,1), size(XA,2)*size(XT,2)-3) \n",
    "    c  = [c; exp.(XX0*sol.zero)];\n",
    "end; \n",
    "c = reshape(c,100, nodea)\n",
    "for i=1:nodea \n",
    "    k1 = [k1; exp.(at[i])*kt.^alpha+(1-delta)*kt.-c[:,i]];\n",
    "    ii = [ii; exp.(at[i])*kt.^alpha-c[:,i]];\n",
    "end;\n",
    "k1 = reshape(k1,100, nodea);\n",
    "ii = reshape(ii,100, nodea);"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Julia 0.6.1",
   "language": "julia",
   "name": "julia-0.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}