Interface to control evolution of lambda (#55)

* Interface to control evolution of lambda * rewrite `sim_nr!` with `damping` callback * Same sim_solver and damping options in shockdecomp and stoch_simulate * added testset "damping" * updated docstrings --------- Co-authored-by: Boyan Bejanov <[email protected]>
bankofcanada · May 27, 2024 · 983d9ac · 983d9ac
1 parent 141035e
commit 983d9ac
Show file tree

Hide file tree

Showing 12 changed files with 453 additions and 126 deletions.
diff --git a/src/StackedTimeSolver.jl b/src/StackedTimeSolver.jl
@@ -1,7 +1,7 @@
 ##################################################################################
 # This file is part of StateSpaceEcon.jl
 # BSD 3-Clause License
-# Copyright (c) 2020-2023, Bank of Canada
+# Copyright (c) 2020-2024, Bank of Canada
 # All rights reserved.
 ##################################################################################
 
@@ -54,8 +54,8 @@ include("stackedtime/sparse.jl")
 include("stackedtime/fctypes.jl")
 include("stackedtime/misc.jl")
 include("stackedtime/solverdata.jl")
-include("stackedtime/sim_lm.jl")
 include("stackedtime/sim_nr.jl")
+include("stackedtime/sim_lm.jl")
 include("stackedtime/sim_gn.jl")
 include("stackedtime/simulate.jl")
 include("stackedtime/shockdecomp.jl")

diff --git a/src/simulate.jl b/src/simulate.jl
@@ -1,7 +1,7 @@
 ##################################################################################
 # This file is part of StateSpaceEcon.jl
 # BSD 3-Clause License
-# Copyright (c) 2020-2023, Bank of Canada
+# Copyright (c) 2020-2024, Bank of Canada
 # All rights reserved.
 ##################################################################################
 
@@ -92,16 +92,44 @@ Run a simulation for the given model, simulation plan and exogenous data.
         deviation from the steady state. Default value is `false`.
     * `anticipate::Bool` - set to `false` to instruct the solver that all shocks
         are unanticipated by the agents. Default value is `true`.
+    * `solver::Symbol` - specify the simulation solver. Available options are
+      :stackedtime and :firstorder. If not given, default is :stackedtime.
     * `verbose::Bool` - control whether or not to print progress information.
         Default value is taken from `model.options`.
     * `tol::Float64` - set the desired accuracy. Default value is taken from
         `model.options`.
     * `maxiter::Int` - algorithm fails if the desired accuracy is not reached
         within this maximum number of iterations. Default value is taken from
         `model.options`.
-    * `linesearch::Bool` - When `true` the Newton-Raphson is modified to include a
-        search along the descent direction for a sufficient decrease in f. It will
-        do this at each iteration. Default is `false`.
+The following options are specific to the `:stackedtime` solver
+    * `sim_solver` - specify the non-linear solver to use. Available options are 
+      - `:sim_nr` : (default) Newton-Raphson, with possible damping, see below.
+      - `:sim_lm` : Levenberg–Marquardt
+      - `:sim_gn` : Gauss-Newton
+    * `linesearch::Bool` - When `true` the Newton-Raphson is modified to include
+      a search along the descent direction for a sufficient decrease in f. It
+      will do this at each iteration. Default is `false`. (Superseded by the
+      `damping` option described below)
+    * `damping` - Specifies the style of damping that can be applied to the
+      Newton non-linear solver. Available options are:
+      - if not given the default behaviour is no damping, i.e. the damping
+        coefficient is set to 1.0 in each iteration.
+      - number: the damping coefficient will be set to the given number (rather than 1)
+      - vector of numbers: the damping coefficient in each iteration will be set
+        the number in the corresponding entry of the given vector. If there are
+        more Newton iterations than the length of the vector, the last entry
+        will be used until in the remaining iterations.
+      - `:linesearch` or `:armijo` : same as setting `linesearch=true`. The
+        Armijo rule is taken from "C.T.Kelly, Iterative Methods for Linear and
+        Nonlinear Equations, ch.8.1, p.137"
+      - `(:armijo, :sigma => 0.5, :alpha => 1e-4)` - override the default
+        parameters of the Armijo rule.
+      - `:br81` : (experimental) implements the damping algorithm in "Bank, R.E.,
+        Rose, D.J. Global approximate Newton methods. Numer. Math. 37, 279–295
+        (1981)."
+      - `(:br81, :rateK => 10, :delta => 0.1)` : override the default parameters
+        of the Bank & Rose (1981) algorithm.
+
 """
 function simulate end
 export simulate

diff --git a/src/stackedtime/shockdecomp.jl b/src/stackedtime/shockdecomp.jl
@@ -1,7 +1,7 @@
 ##################################################################################
 # This file is part of StateSpaceEcon.jl
 # BSD 3-Clause License
-# Copyright (c) 2020-2023, Bank of Canada
+# Copyright (c) 2020-2024, Bank of Canada
 # All rights reserved.
 ##################################################################################
 
@@ -37,11 +37,17 @@ function shockdecomp(m::Model, p::Plan, exog_data::SimData;
     verbose::Bool=getoption(m, :verbose, false),
     maxiter::Int=getoption(m, :maxiter, 20),
     tol::Float64=getoption(m, :tol, 1e-12),
-    linesearch::Bool=getoption(m, :linesearch, false),
     fctype::FinalCondition=getoption(m, :fctype, fcgiven),
-    _debug=false
+    _debug=false,
+    #= Newton-Raphson options =#
+    warn_maxiter=getoption(getoption(m, :warn, Options()), :maxiter, false),
+    linesearch::Bool=getoption(m, :linesearch, false),
+    sim_solver=:sim_nr,
+    damping=nothing
 )
 
+    sim_solve!, damping = _get_solver_damping(linesearch, sim_solver, damping)
+
     refresh_med!(m, variant)
 
     if !anticipate
@@ -74,7 +80,8 @@ function shockdecomp(m::Model, p::Plan, exog_data::SimData;
     stackedtime_R!(res_shocked, shocked, shocked, gdata)    # Run the "shocked" simulation with the given exogenous data.
     if norm(res_shocked, Inf) > tol
         assign_exog_data!(shocked, exog_data, gdata)
-        sim_nr!(shocked, gdata, maxiter, tol, verbose, linesearch)
+        converged = sim_solve!(shocked, gdata, maxiter, tol, verbose, damping)
+        check_converged(converged, warn_maxiter)
     end
 
     # We need the Jacobian matrix evaluated at the control solution.

diff --git a/src/stackedtime/sim_gn.jl b/src/stackedtime/sim_gn.jl
@@ -1,12 +1,12 @@
 ##################################################################################
 # This file is part of StateSpaceEcon.jl
 # BSD 3-Clause License
-# Copyright (c) 2020-2023, Bank of Canada
+# Copyright (c) 2020-2024, Bank of Canada
 # All rights reserved.
 ##################################################################################
 
 """
-    sim_gn!(x, sd, maxiter, tol, verbose [, linesearch])
+    sim_gn!(x, sd, maxiter, tol, verbose, damping)
 
 Solve the simulation problem using the Gauss-Newton method.
   * `x` - the array of data for the simulation. All initial, final and exogenous
@@ -16,12 +16,14 @@ Solve the simulation problem using the Gauss-Newton method.
   * `maxiter` - maximum number of iterations.
   * `tol` - desired accuracy.
   * `verbose` - whether or not to print progress information.
+  * `damping` - unused, no damping is implemented.
 """
 function sim_gn!(x::AbstractArray{Float64}, sd::StackedTimeSolverData,
-    maxiter::Int64, tol::Float64, verbose::Bool, linesearch::Bool=false
+    maxiter::Int64, tol::Float64, verbose::Bool,
+    ::Function # no damping used here
 )
 
-    @warn "Gauss-Newton method is experimental."
+    # @warn "Gauss-Newton method is experimental."
     if verbose
         @info "Simulation using Gauss-Newton method."
     end

diff --git a/src/stackedtime/sim_lm.jl b/src/stackedtime/sim_lm.jl
@@ -1,7 +1,7 @@
 ##################################################################################
 # This file is part of StateSpaceEcon.jl
 # BSD 3-Clause License
-# Copyright (c) 2020-2023, Bank of Canada
+# Copyright (c) 2020-2024, Bank of Canada
 # All rights reserved.
 ##################################################################################
 
@@ -109,7 +109,7 @@ function _sim_lm_first_step(
 end
 
 """
-    sim_lm!(x, sd, maxiter, tol, verbose [, linesearch])
+    sim_lm!(x, sd, maxiter, tol, verbose, damping)
 
 Solve the simulation problem using the Levenberg–Marquardt method.
   * `x` - the array of data for the simulation. All initial, final and exogenous
@@ -119,12 +119,15 @@ Solve the simulation problem using the Levenberg–Marquardt method.
   * `maxiter` - maximum number of iterations.
   * `tol` - desired accuracy.
   * `verbose` - whether or not to print progress information.
+  * `damping` - no damping is implemented here, however the damping is forwarded
+    to [`sim_nr!`](@ref) when the LM algorithm switches to Newton-Raphson. 
 """
 function sim_lm!(x::AbstractArray{Float64}, sd::StackedTimeSolverData,
-    maxiter::Int64, tol::Float64, verbose::Bool, linesearch::Bool=false
+    maxiter::Int64, tol::Float64, verbose::Bool,
+    damping::Function # no damping in LM method, but passed on to sim_nr! when we switch to it
 )
 
-    @warn "Levenberg–Marquardt method is experimental."
+    # @warn "Levenberg–Marquardt method is experimental."
     if verbose
         @info "Simulation using Levenberg–Marquardt method."
     end
@@ -145,7 +148,7 @@ function sim_lm!(x::AbstractArray{Float64}, sd::StackedTimeSolverData,
     nR = norm(R, Inf)
     nΔx = norm(Δx, Inf)
     if verbose
-        @info "1, || Fx || = $(nR), || Δx || = $(nΔx), lambda = $(lm_params[1]), qual = $(lm_params[3])"
+        @info "1, || Fx || = $(nR), || Δx || = $(nΔx), λ = $(lm_params[1]), qual = $(lm_params[3])"
     end
     for it = 2:maxiter
         if nR < tol
@@ -157,14 +160,14 @@ function sim_lm!(x::AbstractArray{Float64}, sd::StackedTimeSolverData,
         nR = norm(R, Inf)
         nΔx = norm(Δx, Inf)
         if verbose
-            @info "$it, || Fx || = $(nR), || Δx || = $(nΔx), lambda = $(lm_params[1]), qual = $(lm_params[3])"
+            @info "$it, || Fx || = $(nR), || Δx || = $(nΔx), λ = $(lm_params[1]), qual = $(lm_params[3])"
         end
-        if  (nΔx < 1e-5 || true ) &&  ((4nR < nR0 && nR < 1e-2) || (nR / nR0 > 0.8)) && (lm_params[1] <= 1e-4)
+        if (nΔx < 1e-5 || true) && ((4nR < nR0 && nR < 1e-2) || (nR / nR0 > 0.8)) && (lm_params[1] <= 1e-4)
             if verbose
                 @info "    --- switching to Newton-Raphson ---   "
             end
             sd.J_factorized[] = nothing
-            return sim_nr!(x, sd, maxiter, tol, verbose, linesearch)
+            return sim_nr!(x, sd, maxiter - it, tol, verbose, damping)
         end
         # nR0 = nR
     end