# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernel_info:
#     name: julia-0.5
#   kernelspec:
#     display_name: Julia 0.6.0-rc1
#     language: julia
#     name: julia-0.6
# ---

# # Solving a New Keynesian model with Julia
#
# This notebook is part of a computational appendix that accompanies the paper.
#
# > MATLAB, Python, Julia: What to Choose in Economics? 
# >> Coleman, Lyon, Maliar, and Maliar (2017)
#
# In order to run the codes in this notebook you will need to install and configure a few Julia packages. We recommend downloading [JuliaPro](https://juliacomputing.com/products/juliapro.html) and/or following the instructions on [quantecon.org](https://lectures.quantecon.org/py/getting_started.html). Once your Julia installation is up and running, there are a few additional packages you will need in order to run the code here. To do this uncomment the lines in the cell below (by deleting the `#` and space at the beginning of each line) and run the cell:

# +
# Pkg.add("Sobol")
# Pkg.add("QuantEcon")
# Pkg.add("BasisMatrices")

# + [markdown] outputExpanded=false
# ## Julia Code
#
# The Julia version of our algorithm is implemented as a few functions defined on
# a core type named `Model`. This type is itself composed of three different
# types that hold the model parameters, steady state, and grids needed to
# describe the numerical model. Before we get to the types, we need to bring in
# some dependencies:

# + outputExpanded=false
# for constructing Sobol sequences
using Sobol
# for basis matrix of complete monomials and monomial quadrature rules
using BasisMatrices: complete_polynomial, complete_polynomial!, n_complete
using QuantEcon: qnwmonomial1, qnwmonomial2
    
# Set seed so we can replicate results
srand(42);

# + [markdown] outputExpanded=false
# ## Types
#
# First we have the `Params` type, which holds all the model parameters as well
# as the paramters that drive the algorithm.

# + outputExpanded=false
immutable Params
    zlb::Bool
    γ::Float64      # Utility-function parameter
    β::Float64      # Discount factor
    ϑ::Float64      # Utility-function parameter
    ϵ::Float64      # Parameter in the Dixit-Stiglitz aggregator
    ϕ_y::Float64    # Parameter of the Taylor rule
    ϕ_π::Float64    # Parameter of the Taylor rule
    μ::Float64      # Parameter of the Taylor rule
    Θ::Float64      # Share of non-reoptimizing firms (Calvo's pricing)
    πstar::Float64  # Target (gross) inflation rate
    gbar::Float64   # Steady-state share of government spending in output

    # autocorrelation coefficients
    ρηR::Float64  # See process (28) in MM (2015)
    ρηa::Float64  # See process (22) in MM (2015)
    ρηL::Float64  # See process (16) in MM (2015)
    ρηu::Float64  # See process (15) in MM (2015)
    ρηB::Float64  # See process (17) in MM (2015)
    ρηG::Float64  # See process (26) in MM (2015)

    # standard deviations
    σηR::Float64  # See process (28) in MM (2015)
    σηa::Float64  # See process (22) in MM (2015)
    σηL::Float64  # See process (16) in MM (2015)
    σηu::Float64  # See process (15) in MM (2015)
    σηB::Float64  # See process (17) in MM (2015)
    σηG::Float64  # See process (26) in MM (2015)

    # algorithm parameters
    deg::Int       # max degree of complete monomial
    damp::Float64  # dampening parameter for coefficient update
    tol::Float64   # Tolerance for stopping iterations
    m::Int         # number of points in the grid
    grid_kind::Symbol  # type of grid (:sobol or :random)
end

# constructor that takes only keyword arguments and specifies
# defualt values
function Params(;zlb::Bool=true, γ= 1, β=0.99, ϑ=2.09, ϵ=4.45, ϕ_y=0.07,
                 ϕ_π=2.21, μ=0.82, Θ=0.83, πstar=1, gbar=0.23,
                 ρηR=0.0, ρηa=0.95, ρηL=0.25, ρηu=0.92, ρηB=0.0, ρηG=0.95,
                 σηR=0.0028, σηa=0.0045, σηL=0.0500,
                 σηu=0.0054, σηB=0.0010, σηG=0.0038,
                 deg=2, damp=0.1, tol=1e-7, m=200, grid_kind=:sobol)
    Params(zlb, γ, β, ϑ, ϵ, ϕ_y, ϕ_π, μ, Θ, πstar, gbar,
           ρηR, ρηa, ρηL, ρηu, ρηB, ρηG, σηR, σηa, σηL, σηu, σηB, σηG,
           deg, damp, tol, m, grid_kind)
end

# returns the covariance matrix for the 6 shocks in the model
vcov(p::Params) = diagm([p.σηR^2, p.σηa^2, p.σηL^2, p.σηu^2, p.σηB^2, p.σηG^2])

# + [markdown] outputExpanded=false
# Next we have a type called `SteadyState` that is intended to hold the
# deterministic steady state realization for each variable in the model.

# + outputExpanded=false
immutable SteadyState
    Yn::Float64
    Y::Float64
    π::Float64
    δ::Float64
    L::Float64
    C::Float64
    F::Float64
    S::Float64
    R::Float64
    w::Float64
end

function SteadyState(p::Params)
    Yn_ss = exp(p.gbar)^(p.γ/(p.ϑ+p.γ))
    Y_ss  = Yn_ss
    π_ss  = 1.0
    δ_ss  = 1.0
    L_ss  = Y_ss/δ_ss
    C_ss  = (1-p.gbar)*Y_ss
    F_ss  = C_ss^(-p.γ)*Y_ss/(1-p.β*p.Θ*π_ss^(p.ϵ-1))
    S_ss  = L_ss^p.ϑ*Y_ss/(1-p.β*p.Θ*π_ss^p.ϵ)
    R_ss  = π_ss/p.β
    w_ss  = (L_ss^p.ϑ)*(C_ss^p.γ)

    SteadyState(Yn_ss, Y_ss, π_ss, δ_ss, L_ss, C_ss, F_ss, S_ss, R_ss, w_ss)
end

# + [markdown] outputExpanded=false
# Given an instance of `Params` and `SteadyState`, we can construct the grid on
# which we will solve the model.
#
# The `Grids` type holds this grid as well as matrices used to compute
# expectations.

# + outputExpanded=false
immutable Grids
    # period t grids
    ηR::Vector{Float64}
    ηa::Vector{Float64}
    ηL::Vector{Float64}
    ηu::Vector{Float64}
    ηB::Vector{Float64}
    ηG::Vector{Float64}
    R::Vector{Float64}
    δ::Vector{Float64}

    # combined matrix and complete polynomial version of it
    X::Matrix{Float64}
    X0_G::Dict{Int,Matrix{Float64}}

    # quadrature weights and nodes
    ϵ_nodes::Matrix{Float64}
    ω_nodes::Vector{Float64}

    # period t+1 grids at all shocks
    ηR1::Matrix{Float64}
    ηa1::Matrix{Float64}
    ηL1::Matrix{Float64}
    ηu1::Matrix{Float64}
    ηB1::Matrix{Float64}
    ηG1::Matrix{Float64}
end

function Grids(p::Params, ss::SteadyState)
    if p.grid_kind == :sobol
        # Values of exogenous state variables are in the interval +/- σ/sqrt(1-ρ^2)
        ub = [2 * p.σηR / sqrt(1 - p.ρηR^2),
              2 * p.σηa / sqrt(1 - p.ρηa^2),
              2 * p.σηL / sqrt(1 - p.ρηL^2),
              2 * p.σηu / sqrt(1 - p.ρηu^2),
              2 * p.σηB / sqrt(1 - p.ρηB^2),
              2 * p.σηG / sqrt(1 - p.ρηG^2),
              1.05,  # R
              1.0    # δ
             ]
        lb = -ub
        lb[[7, 8]] = [1.0, 0.95]  # adjust lower bound for R and δ

        # construct SobolSeq
        s = SobolSeq(length(ub), ub, lb)
        # skip(s, m)  # See note in README of Sobol.jl

        # gather points
        seq = Array{Float64}(8, p.m)
        for i in 1:p.m
            seq[:, i] = next(s)
        end
        seq = seq'  # transpose so variables are in columns

        ηR = seq[:, 1]
        ηa = seq[:, 2]
        ηL = seq[:, 3]
        ηu = seq[:, 4]
        ηB = seq[:, 5]
        ηG = seq[:, 6]
        R  = seq[:, 7]
        δ  = seq[:, 8]
    else  # assume random
        # Values of exogenous state variables are distributed uniformly
        # in the interval +/- std/sqrt(1-rho_nu^2)
        ηR = (-2*p.σηR + 4*p.σηR*rand(p.m)) / sqrt(1 - p.ρηR^2)
        ηa = (-2*p.σηa + 4*p.σηa*rand(p.m)) / sqrt(1 - p.ρηa^2)
        ηL = (-2*p.σηL + 4*p.σηL*rand(p.m)) / sqrt(1 - p.ρηL^2)
        ηu = (-2*p.σηu + 4*p.σηu*rand(p.m)) / sqrt(1 - p.ρηu^2)
        ηB = (-2*p.σηB + 4*p.σηB*rand(p.m)) / sqrt(1 - p.ρηB^2)
        ηG = (-2*p.σηG + 4*p.σηG*rand(p.m)) / sqrt(1 - p.ρηG^2)

        # Values of endogenous state variables are distributed uniformly
        # in the intervals [1 1.05] and [0.95 1], respectively
        R = 1 + 0.05*rand(p.m)
        δ = 0.95 + 0.05*rand(p.m)
    end

    X = [log.(R) log.(δ) ηR ηa ηL ηu ηB ηG]
    X0_G = Dict(
        1 => complete_polynomial(X, 1),
        p.deg => complete_polynomial(X, p.deg)
    )

    ϵ_nodes, ω_nodes = qnwmonomial1(vcov(p))

    ηR1 = p.ρηR.*ηR .+ ϵ_nodes[:, 1]'
    ηa1 = p.ρηa.*ηa .+ ϵ_nodes[:, 2]'
    ηL1 = p.ρηL.*ηL .+ ϵ_nodes[:, 3]'
    ηu1 = p.ρηu.*ηu .+ ϵ_nodes[:, 4]'
    ηB1 = p.ρηB.*ηB .+ ϵ_nodes[:, 5]'
    ηG1 = p.ρηG.*ηG .+ ϵ_nodes[:, 6]'

    Grids(ηR, ηa, ηL, ηu, ηB, ηG, R, δ, X, X0_G, ϵ_nodes, ω_nodes,
          ηR1, ηa1, ηL1, ηu1, ηB1, ηG1)
end

# + [markdown] outputExpanded=false
# Finally, we construct the `Model` type, which has an instance of `Params`,
# `SteadyState` and `Grids` as its three fields.

# + outputExpanded=false
type Model
    p::Params
    s::SteadyState
    g::Grids
end

function Model(;kwargs...)
    p = Params(;kwargs...)
    s = SteadyState(p)
    g = Grids(p, s)
    Model(p, s, g)
end

Base.show(io::IO, m::Model) = println(io, "New Keynesian model")

# + [markdown] outputExpanded=false
# Now that we have a model, we will construct one more helper function that takes
# the control variables $(S, F, C)$ and shocks $(\delta, R, \eta_G, \eta_a,
# \eta_L, \eta_R)$ and applies equilibrium conditions to produce $(\pi, \delta',
# Y, L, Y_n, R')$. We will use this function in both the solution and simulation
# routines below.

# + outputExpanded=false
function Base.step(m::Model, S, F, C, δ0, R0, ηG, ηa, ηL, ηR)
    Θ, ϵ, gbar, ϑ, γ = m.p.Θ, m.p.ϵ, m.p.gbar, m.p.ϑ, m.p.γ
    β, μ, ϕ_π, ϕ_y, πs = m.p.β, m.p.μ, m.p.ϕ_π, m.p.ϕ_y, m.s.π

    # Compute pie(t) from condition (35) in MM (2015)
    π0 = ((1-(1-Θ)*(S./F).^(1-ϵ))/Θ).^(1/(ϵ-1))

    # Compute delta(t) from condition (36) in MM (2015)
    δ1 = ((1-Θ)*((1-Θ*π0.^(ϵ-1))/(1-Θ)).^(ϵ/(ϵ-1))+Θ*π0.^ϵ./δ0).^(-1.0)

    # Compute Y(t) from condition (38) in MM (2015)
    Y0 = C./(1-gbar./exp.(ηG))

    # Compute L(t) from condition (37) in MM (2015)
    L0 = Y0./exp.(ηa)./δ1

    #  Compute Yn(t) from condition (31) in MM (2015)
    Yn0 = (exp.(ηa).^(1+ϑ).*(1-gbar./exp.(ηG)).^(-γ)./exp.(ηL)).^(1/(ϑ+γ))

    # Compute R(t) from conditions (27), (39) in MM (2015) -- Taylor rule
    R1 = πs ./ β .* (R0*β./πs).^μ.*((π0./πs).^ϕ_π .* (Y0./Yn0).^ϕ_y).^(1-μ).*exp.(ηR)

    π0, δ1, Y0, L0, Yn0, R1
end

# + [markdown] outputExpanded=false
# ### Solution routine

# + outputExpanded=false
# construct an initial guess for the solution
function initial_coefs(m::Model)
    npol = size(m.g.X0_G[1], 2)
    coefs = fill(1e-5, npol, 3)
    coefs[1, :] = [m.s.S, m.s.F, m.s.C^(-m.p.γ)]
    coefs
end

function solve(m::Model)
    # simplify notation
    n, n_nodes = size(m.g.ηR, 1), length(m.g.ω_nodes)

    ## allocate memory
    # euler equations
    e = zeros(n, 3)

    # previous iteration S, F, C
    S0_old_G = ones(n)
    F0_old_G = ones(n)
    C0_old_G = ones(n)

    # current iteration S, F, C
    S0_new_G = ones(n)
    F0_new_G = ones(n)
    C0_new_G = ones(n)

    # future S, F, C
    S1 = zeros(n, n_nodes)
    F1 = zeros(n, n_nodes)
    C1 = zeros(n, n_nodes)
    π1 = zeros(n, n_nodes)

    local coefs

    for deg in [1, m.p.deg]
        # set up matrices for this degree
        X0_G = m.g.X0_G[deg]

        # future basis matrix for S, F, C
        X1 = Array{Float64}(n, n_complete(8, deg))

        # initialize coefs
        if deg == 1
            coefs = initial_coefs(m)
        else
            # for higher order, fit the polynomial on the states e we left
            # off with in the first order solution
            coefs = X0_G \ e
        end

        err = 1.0

        # solve at this degree of complete polynomial
        while err > m.p.tol
            # Current choices (at t)
            # ------------------------------
            S0 = X0_G*coefs[:, 1]                # Compute S(t) using coefs
            F0 = X0_G*coefs[:, 2]                # Compute F(t) using coefs
            C0 = (X0_G*coefs[:, 3]).^(-1/m.p.γ)  # Compute C(t) using coefs

            π0, δ1, Y0, L0, Yn0, R1 = step(m, S0, F0, C0, m.g.δ, m.g.R, m.g.ηG,
                                           m.g.ηa, m.g.ηL, m.g.ηR)

            if m.p.zlb R1 .= max.(R1, 1.0) end

            for u in 1:n_nodes
                # Form complete polynomial of degree "Degree" (at t+1) on future state
                complete_polynomial!(
                    X1,
                    hcat(log.(R1), log.(δ1), m.g.ηR1[:, u], m.g.ηa1[:, u],
                         m.g.ηL1[:, u], m.g.ηu1[:, u], m.g.ηB1[:, u], m.g.ηG1[:, u]),
                    deg
                )

                S1[:, u] = X1*coefs[:, 1]                # Compute S(t+1)
                F1[:, u] = X1*coefs[:, 2]                # Compute F(t+1)
                C1[:, u] = (X1*coefs[:, 3]).^(-1/m.p.γ)  # Compute C(t+1)
            end

            # Compute next-period π using condition
            # (35) in MM (2015)
            π1 .= ((1-(1-m.p.Θ).*(S1./F1).^(1-m.p.ϵ))/m.p.Θ).^(1/(m.p.ϵ-1))

            # Evaluate conditional expectations in the Euler equations
            #---------------------------------------------------------
            e[:, 1] = exp.(m.g.ηu).*exp.(m.g.ηL).*L0.^m.p.ϑ.*Y0./exp.(m.g.ηa) + (m.p.β*m.p.Θ*π1.^m.p.ϵ.*S1)*m.g.ω_nodes
            e[:, 2] = exp.(m.g.ηu).*C0.^(-m.p.γ).*Y0 + (m.p.β*m.p.Θ*π1.^(m.p.ϵ-1).*F1)*m.g.ω_nodes
            e[:, 3] = m.p.β*exp.(m.g.ηB)./exp.(m.g.ηu).*R1.*((exp.(m.g.ηu1).*C1.^(-m.p.γ)./π1)*m.g.ω_nodes)

            # Variables of the current iteration
            #-----------------------------------
            copy!(S0_new_G, S0)
            copy!(F0_new_G, F0)
            copy!(C0_new_G, C0)

            # Compute and update the coefficients of the decision functions
            # -------------------------------------------------------------
            coefs_hat = X0_G\e   # Compute the new coefficients of the decision
                                 # functions using a backslash operator

            # Update the coefficients using damping
            coefs = m.p.damp*coefs_hat + (1-m.p.damp)*coefs

            # Evaluate the percentage (unit-free) difference between the values
            # on the grid from the previous and current iterations
            # -----------------------------------------------------------------
            # The convergence criterion is adjusted to the damping parameters
            err = mean(abs, 1-S0_new_G./S0_old_G) +
                  mean(abs, 1-F0_new_G./F0_old_G) +
                  mean(abs, 1-C0_new_G./C0_old_G)

            # Store the obtained values for S(t), F(t), C(t) on the grid to
            # be used on the subsequent iteration in Section 10.2.6
            #-----------------------------------------------------------------------
            copy!(S0_old_G, S0_new_G)
            copy!(F0_old_G, F0_new_G)
            copy!(C0_old_G, C0_new_G)
        end
    end

    coefs
end

# + [markdown] outputExpanded=false
# ### Simulation

# + outputExpanded=false
type Simulation
    # shocks
    ηR::Vector{Float64}
    ηa::Vector{Float64}
    ηL::Vector{Float64}
    ηu::Vector{Float64}
    ηB::Vector{Float64}
    ηG::Vector{Float64}

    # variables
    δ::Vector{Float64}
    R::Vector{Float64}
    S::Vector{Float64}
    F::Vector{Float64}
    C::Vector{Float64}
    π::Vector{Float64}
    Y::Vector{Float64}
    L::Vector{Float64}
    Yn::Vector{Float64}
    w::Vector{Float64}
end

function Simulation(m::Model, coefs::Matrix, capT::Int=10201)

    # 11. Simualating a time-series solution
    #---------------------------------------

    # Initialize the values of 6 exogenous shocks and draw innovations
    #-----------------------------------------------------------------
    ηR = zeros(capT)
    ηa = zeros(capT)
    ηL = zeros(capT)
    ηu = zeros(capT)
    ηB = zeros(capT)
    ηG = zeros(capT)

    # Generate the series for shocks
    #-------------------------------
    @inbounds for t in 1:capT-1
        ηR[t+1] = m.p.ρηR*ηR[t] + m.p.σηR*randn()
        ηa[t+1] = m.p.ρηa*ηa[t] + m.p.σηa*randn()
        ηL[t+1] = m.p.ρηL*ηL[t] + m.p.σηL*randn()
        ηu[t+1] = m.p.ρηu*ηu[t] + m.p.σηu*randn()
        ηB[t+1] = m.p.ρηB*ηB[t] + m.p.σηB*randn()
        ηG[t+1] = m.p.ρηG*ηG[t] + m.p.σηG*randn()
    end

    δ  = ones(capT+1) # Allocate memory for the time series of delta(t)
    R  = ones(capT+1) # Allocate memory for the time series of R(t)
    S  = Array{Float64}(capT)   # Allocate memory for the time series of S(t)
    F  = Array{Float64}(capT)   # Allocate memory for the time series of F(t)
    C  = Array{Float64}(capT)   # Allocate memory for the time series of C(t)
    π  = Array{Float64}(capT)   # Allocate memory for the time series of π(t)
    Y  = Array{Float64}(capT)   # Allocate memory for the time series of Y(t)
    L  = Array{Float64}(capT)   # Allocate memory for the time series of L(t)
    Yn = Array{Float64}(capT)   # Allocate memory for the time series of Yn(t)
    w  = Array{Float64}(capT)

    pol_bases = Array{Float64}(1, size(coefs, 1))
    @inbounds for t in 1:capT
        # Construct the matrix of explanatory variables "pol_bases" on the series
        # of state variables; columns of "pol_bases" are given by the basis
        # functions of the polynomial of degree 2
        complete_polynomial!(
            pol_bases,
            hcat(log(R[t]), log(δ[t]), ηR[t], ηa[t], ηL[t], ηu[t], ηB[t], ηG[t]),
            m.p.deg
        )
        S[t], F[t], MU = pol_bases*coefs
        C[t] = (MU).^(-1/m.p.γ)

        π[t], δ[t+1], Y[t], L[t], Yn[t], R[t+1] = step(m, S[t], F[t], C[t], δ[t],
                                                       R[t], ηG[t], ηa[t], ηL[t], ηR[t])

        # Compute real wage
        w[t] = exp(ηL[t])*(L[t]^m.p.ϑ)*(C[t]^m.p.γ)

        # If ZLB is imposed, set R(t)=1 if ZLB binds
        if m.p.zlb; R[t+1] = max(R[t+1],1.0); end
    end

    Simulation(ηR, ηa, ηL, ηu, ηB, ηG, δ, R, S, F, C, π, Y, L, Yn, w)
end

# + [markdown] outputExpanded=false
# ### Accuracy

# + outputExpanded=false
type Residuals
    resids::Matrix{Float64}
end

function Residuals(m::Model, coefs::Matrix, sim::Simulation; burn::Int=200)
    capT = length(sim.w)
    resids = zeros(9, capT)

    # Integration method for evaluating accuracy
    # ------------------------------------------
    # Monomial integration rule with 2N^2+1 nodes
    ϵ_nodes, ω_nodes = qnwmonomial2(vcov(m.p))
    n_nodes = length(ω_nodes)

    # Allocate for arrays needed in the loop
    basis_mat = Array{Float64}(n_nodes, 8)
    X1 = Array{Float64}(n_nodes, size(coefs, 1))

    ηR1 = Array{Float64}(n_nodes)
    ηa1 = Array{Float64}(n_nodes)
    ηL1 = Array{Float64}(n_nodes)
    ηu1 = Array{Float64}(n_nodes)
    ηB1 = Array{Float64}(n_nodes)
    ηG1 = Array{Float64}(n_nodes)

    for t = 1:capT                 # For each given point,
        # Take the corresponding value for shocks at t
        #---------------------------------------------
        ηR0 = sim.ηR[t]  # ηR(t)
        ηa0 = sim.ηa[t]  # ηa(t)
        ηL0 = sim.ηL[t]  # ηL(t)
        ηu0 = sim.ηu[t]  # ηu(t)
        ηB0 = sim.ηB[t]  # ηB(t)
        ηG0 = sim.ηG[t]  # ηG(t)

        # Exctract time t values for all other variables (and t+1 for R, δ)
        #------------------------------------------------------------------
        R0  = sim.R[t]   # R(t-1)
        δ0  = sim.δ[t]   # δ(t-1)
        R1  = sim.R[t+1] # R(t)
        δ1  = sim.δ[t+1] # δ(t)

        L0  = sim.L[t]   # L(t)
        Y0  = sim.Y[t]   # Y(t)
        Yn0 = sim.Yn[t]  # Yn(t)
        π0  = sim.π[t]   # π(t)
        S0  = sim.S[t]   # S(t)
        F0  = sim.F[t]   # F(t)
        C0  = sim.C[t]   # C(t)

        # Fill basis matrix with R1, δ1 and shocks
        #-----------------------------------------
        # Note that we do not premultiply by standard deviations as ϵ_nodes
        # already include them. All these variables are vectors of length n_nodes
        copy!(ηR1, ηR0*m.p.ρηR + ϵ_nodes[:, 1])
        copy!(ηa1, ηa0*m.p.ρηa + ϵ_nodes[:, 2])
        copy!(ηL1, ηL0*m.p.ρηL + ϵ_nodes[:, 3])
        copy!(ηu1, ηu0*m.p.ρηu + ϵ_nodes[:, 4])
        copy!(ηB1, ηB0*m.p.ρηB + ϵ_nodes[:, 5])
        copy!(ηG1, ηG0*m.p.ρηG + ϵ_nodes[:, 6])

        basis_mat[:, 1] = log(R1)
        basis_mat[:, 2] = log(δ1)
        basis_mat[:, 3] = ηR1
        basis_mat[:, 4] = ηa1
        basis_mat[:, 5] = ηL1
        basis_mat[:, 6] = ηu1
        basis_mat[:, 7] = ηB1
        basis_mat[:, 8] = ηG1

        # Future choices at t+1
        #----------------------
        # Form a complete polynomial of degree "Degree" (at t+1) on future state
        # variables; n_nodes-by-npol
        complete_polynomial!(X1, basis_mat, m.p.deg)

        # Compute S(t+1), F(t+1) and C(t+1) in all nodes using coefs
        S1 = X1*coefs[:, 1]
        F1 = X1*coefs[:, 2]
        C1 = (X1*coefs[:, 3]).^(-1/m.p.γ)

        # Compute π(t+1) using condition (35) in MM (2015)
        π1 = ((1-(1-m.p.Θ)*(S1./F1).^(1-m.p.ϵ))/m.p.Θ).^(1/(m.p.ϵ-1))

        # Compute residuals for each of the 9 equilibrium conditions
        #-----------------------------------------------------------
        resids[1, t] = 1-dot(ω_nodes,
            (exp(ηu0)*exp(ηL0)*L0^m.p.ϑ*Y0/exp(ηa0) + m.p.β*m.p.Θ*π1.^m.p.ϵ.*S1)/S0
        )
        resids[2, t] = 1 - dot(ω_nodes,
            (exp(ηu0)*C0^(-m.p.γ)*Y0 + m.p.β*m.p.Θ*π1.^(m.p.ϵ-1).*F1)/F0
        )
        resids[3, t] = 1.0 -dot(ω_nodes,
            (m.p.β*exp(ηB0)/exp(ηu0)*R1*exp.(ηu1).*C1.^(-m.p.γ)./π1)/C0^(-m.p.γ)
        )
        resids[4, t] = 1-((1-m.p.Θ*π0^(m.p.ϵ-1))/(1-m.p.Θ))^(1/(1-m.p.ϵ))*F0/S0
        resids[5, t] = 1-((1-m.p.Θ)*((1-m.p.Θ*π0^(m.p.ϵ-1))/(1-m.p.Θ))^(m.p.ϵ/(m.p.ϵ-1)) + m.p.Θ*π0^m.p.ϵ/δ0)^(-1)/δ1
        resids[6, t] = 1-exp(ηa0)*L0*δ1/Y0
        resids[7, t] = 1-(1-m.p.gbar/exp(ηG0))*Y0/C0
        resids[8, t] = 1-(exp(ηa0)^(1+m.p.ϑ)*(1-m.p.gbar/exp(ηG0))^(-m.p.γ)/exp(ηL0))^(1/(m.p.ϑ+m.p.γ))/Yn0
        resids[9, t] = 1-m.s.π/m.p.β*(R0*m.p.β/m.s.π)^m.p.μ*((π0/m.s.π)^m.p.ϕ_π * (Y0/Yn0)^m.p.ϕ_y)^(1-m.p.μ)*exp(ηR0)/R1   # Taylor rule

        # If the ZLB is imposed and R>1, the residuals in the Taylor rule (the
        # 9th equation) are zero
        if m.p.zlb && R1 <= 1; resids[9, t] = 0.0; end

    end
    # discard the first burn observations
    Residuals(resids[:, burn+1:end])
end

Base.mean(r::Residuals) = log10(mean(abs, r.resids))
Base.max(r::Residuals) = log10(maximum(abs, r.resids))
max_E(r::Residuals) = log10.(maximum(abs, r.resids, 2))[:]

# + [markdown] outputExpanded=false
# ### Running the code
#
# Now that we've done all the hard work to define the model, its solution and
# simulation, and accuracy checks, let's put things together and run the code!

# + outputExpanded=false
function main(m=Model(); io::IO=STDOUT)
    tic(); coefs = solve(m); solve_time = toq()

    # simulate the model
    tic(); sim = Simulation(m, coefs); simulation_time = toq()

    # check accuracy
    tic(); resids = Residuals(m, coefs, sim); resids_time = toq()

    err_by_eq = max_E(resids)
    l1 = log10(sum(maximum(abs, resids.resids, 2)))
    l∞ = max(resids)
    tot_time = solve_time + simulation_time + resids_time
    round3(x) = round(x, 3)
    round2(x) = round(x, 2)

    println(io, "Solver time (in seconds): $(solve_time)")
    println(io, "Simulation time (in seconds): $(simulation_time)")
    println(io, "Residuals time (in seconds): $(resids_time)")
    println(io, "total time (in seconds): $(tot_time)")
    println(io, "\nAPPROXIMATION ERRORS (log10):");
    println(io, "\ta) mean error in the model equations: $(round3(mean(resids)))");
    println(io, "\tb) sum of max error per equation: $(round3(l1))");
    println(io, "\tc) max error in the model equations: $(round3(l∞))");
    println(io, "\td) max error by equation:$(round3.(err_by_eq))");
    println(io, "tex row\n", join(round2.([l1, l∞, tot_time]), " & "))

    solve_time, simulation_time, resids_time, coefs, sim, resids
end


function build_paper_table()
    # call main once to precompile all routines
    main()

    open("output.log", "w") do f
        for params in [Dict(:πstar=>1.0, :σηL=>0.1821, :zlb=>false),
                       Dict(:πstar=>1.0, :σηL=>0.4054, :zlb=>false),
                       Dict(:πstar=>1.0, :σηL=>0.1821, :zlb=>true)]
            for grid_kind in [:sobol, :random]
                m = Model(;grid_kind=grid_kind, params...)

                println(f, "Working with params:")
                show(f, MIME"text/plain"(), params)
                println(f, "\nand grid type: $(grid_kind)")
                main(m, io=f)
                println(f, "\n"^5)
            end
        end
    end
end
# -

main();



StarcoderdataJupyterScriptsDedupFiltered