From ad966aad09bb5c182f9868333c599e797647f34d Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Sat, 13 Jul 2024 18:55:17 +0200 Subject: [PATCH] Estimation (#84) non linear shock decomposition fast linear Kalman and inversion filter simplify keyword for parameters macro triggering symbolic equation solves input custom plot attributes (overriding defaults) improved plots custom reverse mode AD inversion and kalman filter (1st order solutions) eliminated a lot of allocations for estimation --------- Co-authored-by: Thore Kockerols --- .github/workflows/ci.yml | 38 +- .gitignore | 4 +- Project.toml | 98 +- README.md | 8 +- benchmark/optim_solver_params.jl | 13 + docs/src/index.md | 8 +- docs/src/tutorials/calibration.md | 4 +- docs/src/tutorials/estimation.md | 50 +- docs/src/tutorials/rbc.md | 4 +- docs/src/tutorials/sw03.md | 2 +- docs/src/unfinished_docs/dsl.md | 2 +- docs/src/unfinished_docs/todo.md | 34 +- models/Smets_Wouters_2007.jl | 50 +- models/Smets_Wouters_2007_linear.jl | 4 +- src/MacroModelling.jl | 3441 +++++++++++++---- src/common_docstrings.jl | 2 +- src/get_functions.jl | 219 +- src/macros.jl | 31 +- src/plotting.jl | 257 +- src/structures.jl | 15 +- test/functionality_tests.jl | 48 +- test/runtests.jl | 84 + ...t_1st_order_inversion_filter_estimation.jl | 31 +- test/test_2nd_order_estimation.jl | 10 +- test/test_3rd_order_estimation.jl | 10 +- test/test_estimation.jl | 52 +- test/test_models.jl | 228 +- test/test_standalone_function.jl | 2 +- test/test_sw07_estimation.jl | 153 + 29 files changed, 3768 insertions(+), 1134 deletions(-) create mode 100644 test/test_sw07_estimation.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1883999f..5a8c58a0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,7 +7,7 @@ jobs: env: GKS_ENCODING: "utf8" GKSwstype: "nul" - name: ${{ matrix.test_set }} - ${{ matrix.version }} - ${{ matrix.os }} + name: ${{ matrix.test_set }} - ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} runs-on: ${{ matrix.os }} continue-on-error: ${{ matrix.version == 'nightly' || matrix.version == '^1.11.0-0' }} strategy: @@ -15,7 +15,7 @@ jobs: matrix: test_set: ["basic", "plots", "higher_order", "estimation"] version: ['1.8', '1.9', '1.10'] - os: [ubuntu-latest, macos-14, windows-latest] + os: [ubuntu-latest, macos-latest, windows-latest] arch: [x64, arm64] exclude: - version: '1.8' @@ -31,50 +31,54 @@ jobs: - arch: arm64 os: windows-latest - arch: x64 - os: macos-14 + os: macos-latest include: - os: ubuntu-latest prefix: xvfb-run # - version: '1.10' - # os: macos-14 + # os: macos-latest # arch: x64 # test_set: "solver0" # - version: '1.10' - # os: macos-14 + # os: macos-latest # arch: x64 # test_set: "solver1" # - version: '1.10' - # os: macos-14 + # os: macos-latest # arch: x64 # test_set: "solver2" # - version: '1.10' - # os: macos-14 + # os: macos-latest # arch: x64 # test_set: "solver3" - version: '1.10' - os: macos-14 - arch: x64 + os: macos-latest + arch: arm64 + test_set: "estimate_sw07" + - version: '1.10' + os: macos-latest + arch: arm64 test_set: "1st_order_inversion_estimation" - version: '1.10' - os: macos-14 - arch: x64 + os: macos-latest + arch: arm64 test_set: "2nd_order_estimation" - version: '1.10' - os: macos-14 - arch: x64 + os: macos-latest + arch: arm64 test_set: "3rd_order_estimation" # - version: '1.10' # os: macOS-latest # arch: x64 # test_set: "basic" - version: 'nightly' - os: macos-14 + os: ubuntu-latest arch: x64 test_set: "basic" allow_failure: true - version: '^1.11.0-0' - os: macos-14 - arch: x64 + os: macos-latest + arch: arm64 test_set: "basic" allow_failure: true steps: @@ -90,7 +94,7 @@ jobs: if: matrix.os != 'windows-latest' run: echo "TEST_SET=${{ matrix.test_set }}" >> $GITHUB_ENV - name: Set JULIA_NUM_THREADS for estimation tests - if: matrix.version == '1.10' && (matrix.test_set == 'estimation' || matrix.test_set == 'solver0' || matrix.test_set == 'solver1' || matrix.test_set == 'solver2' || matrix.test_set == 'solver3' || matrix.test_set == '1st_order_inversion_estimation' || matrix.test_set == '2nd_order_estimation' || matrix.test_set == '3rd_order_estimation') + if: (matrix.version == '1.10' && (matrix.test_set == 'estimation' || matrix.test_set == 'estimate_sw07' || matrix.test_set == '1st_order_inversion_estimation' || matrix.test_set == '2nd_order_estimation' || matrix.test_set == '3rd_order_estimation')) run: echo "JULIA_NUM_THREADS=auto" >> $GITHUB_ENV - uses: actions/cache@v4 env: diff --git a/.gitignore b/.gitignore index caa5fc61..b7b952f1 100644 --- a/.gitignore +++ b/.gitignore @@ -58,4 +58,6 @@ results octave* *.jls samples* -*samples.csv \ No newline at end of file +*samples.csv +skewness*.png +conditional_forecast*.png diff --git a/Project.toml b/Project.toml index 59f1b2c7..96cde631 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ version = "0.1.36" [deps] AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d" +Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" BlockTriangularForm = "adeb47b7-70bf-415a-bb24-c358563e873a" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" @@ -24,6 +25,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" MatrixEquations = "99c1a7ee-ab34-5fd5-8076-27c950a045f4" NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd" Optim = "429524aa-4258-5aef-a3af-852621145aeb" +Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" @@ -47,55 +49,57 @@ StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [compat] -AbstractDifferentiation = "^0.5, 0.6" -Aqua = "^0.8" -AxisKeys = "^0.2" -BlockTriangularForm = "^0.1" -CSV = "^0.10" -ChainRulesCore = "^1" -Combinatorics = "^1" -DataFrames = "^1" -DataStructures = "^0.18" -DocStringExtensions = "^0.8, 0.9" -DynamicPPL = "0.23, 0.24" -DynarePreprocessor_jll = "^6.2" -FiniteDifferences = "^0.12" -ForwardDiff = "^0.10" -ImplicitDifferentiation = "^0.5" -JET = "0.7, 0.8" -JSON = "^0.21" -Krylov = "^0.9" -LaTeXStrings = "^1" -LineSearches = "^7" -LinearAlgebra = "^1" -LinearOperators = "^2" -MCMCChains = "^6" -MacroTools = "^0.5" -MatrixEquations = "^2" +AbstractDifferentiation = "0.5, 0.6" +Accessors = "0.1" +Aqua = "0.8" +AxisKeys = "0.2" +BlockTriangularForm = "0.1" +CSV = "0.10" +ChainRulesCore = "1" +Combinatorics = "1" +DataFrames = "1" +DataStructures = "0.18" +DocStringExtensions = "0.8, 0.9" +DynamicPPL = "0.23, 0.24, 0.25, 0.26, 0.27, 0.28" +DynarePreprocessor_jll = "6" +FiniteDifferences = "0.12" +ForwardDiff = "0.10" +ImplicitDifferentiation = "0.5" +JET = "0.7, 0.8, 0.9" +JSON = "0.21" +Krylov = "0.9" +LaTeXStrings = "1" +LineSearches = "7" +LinearAlgebra = "1" +LinearOperators = "2" +MCMCChains = "6" +MacroTools = "0.5" +MatrixEquations = "2" NLopt = "0.6, =1.0.1" -Optim = "^1" +Optim = "1" Pigeons = "0.3, 0.4" -PrecompileTools = "^1" -PythonCall = "0.9.0 - 0.9.15" -REPL = "^1" -Random = "^1" -RecursiveFactorization = "^0.2" -Reexport = "^1" -Requires = "^1" -RuntimeGeneratedFunctions = "^0.5" -SparseArrays = "^1" -SpecialFunctions = "^2" -SpeedMapping = "^0.3" -StatsPlots = "^0.15" -Subscripts = "0.1.3" -Suppressor = "^0.2" -SymPyPythonCall = "^0.2" -Symbolics = "^5" -Test = "^1" -ThreadedSparseArrays = "^0.2.3" -Turing = "0.29, 0.30" -Unicode = "^1" -Zygote = "^0.6" +Polyester = "0.7" +PrecompileTools = "1" +PythonCall = "0.9" +REPL = "1" +Random = "1" +RecursiveFactorization = "0.2" +Reexport = "1" +Requires = "1" +RuntimeGeneratedFunctions = "0.5" +SparseArrays = "1" +SpecialFunctions = "2" +SpeedMapping = "0.3, 0.4" +StatsPlots = "0.15" +Subscripts = "0.1.3, 0.2" +Suppressor = "0.2" +SymPyPythonCall = "0.2, 0.3" +Symbolics = "5, 6" +Test = "1" +ThreadedSparseArrays = "0.2.3, 0.3" +Turing = "0.30, 0.31, 0.32, 0.33" +Unicode = "1" +Zygote = "0.6" julia = "1.8" [extras] diff --git a/README.md b/README.md index 854ec07f..b6e63891 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ These kinds of models describe the behavior of a macroeconomy and are particularly suited for counterfactual analysis (economic policy evaluation) and exploring / quantifying specific mechanisms (academic research). Due to the complexity of these models, efficient numerical tools are required, as analytical solutions are often unavailable. `MacroModelling.jl` serves as a tool for handling the complexities involved, such as forward-looking expectations, nonlinearity, and high dimensionality. -The goal of this package is to reduce coding time and speed up model development by providing functions for working with discrete-time DSGE models. The user-friendly syntax, automatic variable declaration, and effective steady state solver facilitate fast prototyping of models. Furthermore, the package allows the user to work with nonlinear model solutions (up to third order (pruned) perturbation) and estimate the model using gradient based samplers (e.g. NUTS, or HMC). Currently, `DifferentiableStateSpaceModels.jl` is the only other package providing functionality to estimate using gradient based samplers but the use is limited to models with an analytical solution of the non stochastic steady state (NSSS). Larger models tend to not have an analytical solution of the NSSS and `MacroModelling.jl` can also use gradient based sampler in this case. The target audience for the package includes central bankers, regulators, graduate students, and others working in academia with an interest in DSGE modelling. +The goal of this package is to reduce coding time and speed up model development by providing functions for working with discrete-time DSGE models. The user-friendly syntax, automatic variable declaration, and effective steady state solver facilitate fast prototyping of models. Furthermore, the package allows the user to work with nonlinear model solutions (up to third order (pruned) perturbation) and estimate the model using gradient based samplers (e.g. NUTS, or HMC). Currently, `DifferentiableStateSpaceModels.jl` is the only other package providing functionality to estimate using gradient based samplers but they use the start-of-period timing convention instead of the end-of-period timing convention used in most other packages. The target audience for the package includes central bankers, regulators, graduate students, and others working in academia with an interest in DSGE modelling. As of now the package can: @@ -103,13 +103,13 @@ The package contains the following models in the `models` folder: **Host language**|julia|MATLAB|julia|Python|julia|julia|julia|MATLAB|MATLAB|MATLAB|R|MATLAB|MATLAB| **Non stochastic steady state solver**|*symbolic* or numerical solver of independent blocks; symbolic removal of variables redundant in steady state; inclusion of calibration equations in problem|numerical solver of independent blocks or user-supplied values/functions||numerical solver of independent blocks or user-supplied values/functions|numerical solver|numerical solver or user supplied values/equations|numerical solver of independent blocks or user-supplied values/functions|numerical solver of independent blocks or user-supplied values/functions|numerical solver of independent blocks or user-supplied values/functions|user-supplied steady state file or numerical solver|numerical solver; inclusion of calibration equations in problem||| **Automatic declaration of variables and parameters**|yes||||||||||||| -**Derivatives (Automatic Differentiation) wrt parameters**|yes|||||yes - for all 1st, 2nd order perturbation solution related output *if user supplied steady state equations*||||||| +**Derivatives wrt parameters**|yes|||||yes||||||| **Perturbation solution order**|1, 2, 3|k|1|1, 2, 3|1, 2, 3|1, 2|1|1|1 to 5|1|1||1 to 5| **Pruning**|yes|yes||||yes|||yes||||| **Automatic derivation of first order conditions**|||||||||||yes|| -**Handles occasionally binding constraints**|yes|yes|yes|yes|yes||||yes|||yes|| +**Occasionally binding constraints**|yes|yes|yes|yes|yes||||yes|||yes|| **Global solution**||||yes|yes|||||||yes|| -**Estimation**|yes|yes|yes|||||yes|yes|yes|yes||| +**Estimation**|yes|yes|yes|||yes||yes|yes|yes|yes||| **Balanced growth path**||yes|yes||||yes|yes|yes|yes||||| **Model input**|macro (julia)|text file|text file|text file|text file|macro (julia)|module (julia)|text file|text file|text file|text file|text file|text file| **Timing convention**|end-of-period|end-of-period||end-of-period|start-of-period|start-of-period|end-of-period|end-of-period|end-of-period|end-of-period|end-of-period|start-of-period|start-of-period| diff --git a/benchmark/optim_solver_params.jl b/benchmark/optim_solver_params.jl index 5f51dc0b..e6828623 100644 --- a/benchmark/optim_solver_params.jl +++ b/benchmark/optim_solver_params.jl @@ -39,6 +39,7 @@ include("../models/Ireland_2004.jl") include("../models/Caldara_et_al_2012.jl") include("../models/Gali_Monacelli_2005_CITR.jl") include("../models/Gali_2015_chapter_3_nonlinear.jl") +# include("../models/Gali_2015_chapter_3_obc.jl") include("../models/Aguiar_Gopinath_2007.jl") include("../models/Ascari_Sbordone_2014.jl") # stands out include("../models/FS2000.jl") @@ -315,6 +316,18 @@ all_models = [ # all_models[1].NSSS_solver_cache pars = [2.9912988764832833, 0.8725, 0.0027, 0.028948770826150612, 8.04, 4.076413176215408, 0.06375413238034794, 0.24284340766769424, 0.5634017580097571, 0.009549630552246828, 0.6342888355132347, 0.5275522227754195, 1.0, 0.06178989216048817, 0.5234277812131813, 0.422, 0.011209254402846185, 0.5047, 0.6020757011698457, 0.7688] + + +par_inputs = MacroModelling.solver_parameters(eps(), eps(), eps(), maxiters, pars..., transformation, 0.0, 2) + +SS(Guerrieri_Iacoviello_2017) +get_solution(Guerrieri_Iacoviello_2017) +simulate(Guerrieri_Iacoviello_2017) +model = Guerrieri_Iacoviello_2017 + +model.SS_solve_func(model.parameter_values, model, false, true, [par_inputs]) + + evaluate_multi_pars_loglikelihood(pars, all_models) diff --git a/docs/src/index.md b/docs/src/index.md index 73802a33..28012470 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -6,7 +6,7 @@ These kinds of models describe the behavior of a macroeconomy and are particularly suited for counterfactual analysis (economic policy evaluation) and exploring / quantifying specific mechanisms (academic research). Due to the complexity of these models, efficient numerical tools are required, as analytical solutions are often unavailable. `MacroModelling.jl` serves as a tool for handling the complexities involved, such as forward-looking expectations, nonlinearity, and high dimensionality. -The goal of this package is to reduce coding time and speed up model development by providing functions for working with discrete-time DSGE models. The user-friendly syntax, automatic variable declaration, and effective steady state solver facilitate fast prototyping of models. Furthermore, the package allows the user to work with nonlinear model solutions (up to third order (pruned) perturbation) and estimate the model using gradient based samplers (e.g. NUTS, of HMC). Currently, `DifferentiableStateSpaceModels.jl` is the only other package providing functionality to estimate using gradient based samplers but the use is limited to models with an analytical solution of the non stochastic steady state (NSSS). Larger models tend to not have an analytical solution of the NSSS and `MacroModelling.jl` can also use gradient based sampler in this case. The target audience for the package includes central bankers, regulators, graduate students, and others working in academia with an interest in DSGE modelling. +The goal of this package is to reduce coding time and speed up model development by providing functions for working with discrete-time DSGE models. The user-friendly syntax, automatic variable declaration, and effective steady state solver facilitate fast prototyping of models. Furthermore, the package allows the user to work with nonlinear model solutions (up to third order (pruned) perturbation) and estimate the model using gradient based samplers (e.g. NUTS, or HMC). Currently, `DifferentiableStateSpaceModels.jl` is the only other package providing functionality to estimate using gradient based samplers but they use the start-of-period timing convention instead of the end-of-period timing convention used in most other packages. The target audience for the package includes central bankers, regulators, graduate students, and others working in academia with an interest in DSGE modelling. As of now the package can: @@ -56,13 +56,13 @@ The package contains the following models in the `models` folder: **Host language**|julia|MATLAB|julia|Python|julia|julia|julia|MATLAB|MATLAB|MATLAB|R|MATLAB|MATLAB| **Non stochastic steady state solver**|*symbolic* or numerical solver of independent blocks; symbolic removal of variables redundant in steady state; inclusion of calibration equations in problem|numerical solver of independent blocks or user-supplied values/functions||numerical solver of independent blocks or user-supplied values/functions|numerical solver|numerical solver or user supplied values/equations|numerical solver of independent blocks or user-supplied values/functions|numerical solver of independent blocks or user-supplied values/functions|numerical solver of independent blocks or user-supplied values/functions|user-supplied steady state file or numerical solver|numerical solver; inclusion of calibration equations in problem||| **Automatic declaration of variables and parameters**|yes||||||||||||| -**Derivatives (Automatic Differentiation) wrt parameters**|yes|||||yes - for all 1st, 2nd order perturbation solution related output *if user supplied steady state equations*||||||| +**Derivatives wrt parameters**|yes|||||yes||||||| **Perturbation solution order**|1, 2, 3|k|1|1, 2, 3|1, 2, 3|1, 2|1|1|1 to 5|1|1||1 to 5| **Pruning**|yes|yes||||yes|||yes||||| **Automatic derivation of first order conditions**|||||||||||yes|| -**Handles occasionally binding constraints**|yes|yes|yes|yes|yes||||yes|||yes|| +**Occasionally binding constraints**|yes|yes|yes|yes|yes||||yes|||yes|| **Global solution**||||yes|yes|||||||yes|| -**Estimation**|yes|yes|yes|||||yes|yes|yes|yes||| +**Estimation**|yes|yes|yes|||yes||yes|yes|yes|yes||| **Balanced growth path**||yes|yes||||yes|yes|yes|yes||||| **Model input**|macro (julia)|text file|text file|text file|text file|macro (julia)|module (julia)|text file|text file|text file|text file|text file|text file| **Timing convention**|end-of-period|end-of-period||end-of-period|start-of-period|start-of-period|end-of-period|end-of-period|end-of-period|end-of-period|end-of-period|start-of-period|start-of-period| diff --git a/docs/src/tutorials/calibration.md b/docs/src/tutorials/calibration.md index 92806926..e3fa88b2 100644 --- a/docs/src/tutorials/calibration.md +++ b/docs/src/tutorials/calibration.md @@ -65,7 +65,7 @@ using MacroModelling end ``` -First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between `begin` and `end`. Equations can contain an equality sign or the expression is assumed to equal 0. Equations cannot span multiple lines (unless you wrap the expression in brackets) and the timing of endogenous variables are expressed in the squared brackets following the variable name (e.g. `[-1]` for the past period). Exogenous variables (shocks) are followed by a keyword in squared brackets indicating them being exogenous (in this case `[x]`). Note that names can leverage julia's unicode capabilities (e.g. alpha can be written as α). +First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between `begin` and `end`. Equations can contain an equality sign or the expression is assumed to equal 0. Equations cannot span multiple lines (unless you wrap the expression in brackets) and the timing of endogenous variables are expressed in the square brackets following the variable name (e.g. `[-1]` for the past period). Exogenous variables (shocks) are followed by a keyword in square brackets indicating them being exogenous (in this case `[x]`). Note that names can leverage julia's unicode capabilities (e.g. alpha can be written as α). ## Define the parameters @@ -319,7 +319,7 @@ get_statistics(Gali_2015, sol.minimizer, algorithm = :pruned_third_order, parame The solution does not match the standard deviation of inflation very well. -Potentially the partial derivatives change a lot for small changes in parameters and even though the partial derivatives for standard deviation of inflation were large wrt `std_a` they might be small for value returned from the optimisation. We can check this with: +Potentially the partial derivatives change a lot for small changes in parameters and even though the partial derivatives for standard deviation of inflation were large wrt `std_a` they might be small for values returned from the optimisation. We can check this with: ```@repl tutorial_3 get_std(Gali_2015, parameter_derivatives = [:σ, :std_a, :α], variables = [:W_real,:Pi], algorithm = :pruned_third_order, parameters = [:α, :std_a] .=> sol.minimizer) diff --git a/docs/src/tutorials/estimation.md b/docs/src/tutorials/estimation.md index b05385db..680565dc 100644 --- a/docs/src/tutorials/estimation.md +++ b/docs/src/tutorials/estimation.md @@ -50,7 +50,7 @@ using MacroModelling end ``` -First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between `begin` and `end`. Equations can contain an equality sign or the expression is assumed to equal 0. Equations cannot span multiple lines (unless you wrap the expression in brackets) and the timing of endogenous variables are expressed in the squared brackets following the variable name (e.g. `[-1]` for the past period). Exogenous variables (shocks) are followed by a keyword in squared brackets indicating them being exogenous (in this case `[x]`). Note that names can leverage julia's unicode capabilities (e.g. alpha can be written as α). +First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between `begin` and `end`. Equations can contain an equality sign or the expression is assumed to equal 0. Equations cannot span multiple lines (unless you wrap the expression in brackets) and the timing of endogenous variables are expressed in the square brackets following the variable name (e.g. `[-1]` for the past period). Exogenous variables (shocks) are followed by a keyword in square brackets indicating them being exogenous (in this case `[x]`). Note that names can leverage julia's unicode capabilities (e.g. alpha can be written as α). ## Define the parameters @@ -99,8 +99,10 @@ data = data(observables,:) Next we define the parameter priors using the Turing package. The `@model` macro of the Turing package allows us to define the prior distributions over the parameters and combine it with the (Kalman filter) loglikelihood of the model and parameters given the data with the help of the `get_loglikelihood` function. We define the prior distributions in an array and pass it on to the `arraydist` function inside the `@model` macro from the Turing package. It is also possible to define the prior distributions inside the macro but especially for reverse mode auto differentiation the `arraydist` function is substantially faster. When defining the prior distributions we can rely n the distribution implemented in the Distributions package. Note that the `μσ` parameter allows us to hand over the moments (`μ` and `σ`) of the distribution as parameters in case of the non-normal distributions (Gamma, Beta, InverseGamma), and we can also define upper and lower bounds truncating the distribution as third and fourth arguments to the distribution functions. Last but not least, we define the loglikelihood and add it to the posterior loglikelihood with the help of the `@addlogprob!` macro. ```@repl tutorial_2 +import Zygote +import DynamicPPL import Turing -import Turing: NUTS, sample, logpdf +import Turing: NUTS, sample, logpdf, AutoZygote prior_distributions = [ Beta(0.356, 0.02, μσ = true), # alp @@ -117,7 +119,9 @@ prior_distributions = [ Turing.@model function FS2000_loglikelihood_function(data, model) parameters ~ Turing.arraydist(prior_distributions) - Turing.@addlogprob! get_loglikelihood(model, data, parameters) + if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() + Turing.@addlogprob! get_loglikelihood(model, data, parameters) + end end ``` @@ -132,7 +136,7 @@ FS2000_loglikelihood = FS2000_loglikelihood_function(data, FS2000); n_samples = 1000 -chain_NUTS = sample(FS2000_loglikelihood, NUTS(), n_samples, progress = false); +chain_NUTS = sample(FS2000_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, progress = false); ``` ### Inspect posterior @@ -141,7 +145,7 @@ In order to understand the posterior distribution and the sequence of sample we ```@repl tutorial_2; setup = :(chain_NUTS = read("../assets/chain_FS2000.jls", Chains)) using StatsPlots -StatsPlots.plot(chain_NUTS); +plot(chain_NUTS); ``` ![NUTS chain](../assets/FS2000_chain_NUTS.png) @@ -149,7 +153,8 @@ StatsPlots.plot(chain_NUTS); Next, we are plotting the posterior loglikelihood along two parameters dimensions, with the other parameters ket at the posterior mean, and add the samples to the visualisation. This visualisation allows us to understand the curvature of the posterior and puts the samples in context. ```@repl tutorial_2 -using ComponentArrays, MCMCChains, DynamicPPL +using ComponentArrays, MCMCChains +import DynamicPPL: logjoint parameter_mean = mean(chain_NUTS) @@ -201,35 +206,12 @@ p ## Find posterior mode -Other than the mean and median of the posterior distribution we can also calculate the mode. To this end we will use L-BFGS optimisation routines from the Optim package. - -First, we define the posterior loglikelihood function, similar to how we defined it for the Turing model macro. +Other than the mean and median of the posterior distribution we can also calculate the mode as follows: ```@repl tutorial_2 -function calculate_posterior_loglikelihood(parameters, prior_distribuions) - alp, bet, gam, mst, rho, psi, del, z_e_a, z_e_m = parameters - log_lik = 0 - log_lik -= get_loglikelihood(FS2000, data, parameters) - - for (dist, val) in zip(prior_distribuions, parameters) - log_lik -= logpdf(dist, val) - end - - return log_lik -end -``` - -Next, we set up the optimisation problem, parameter bounds, and use the optimizer L-BFGS. - -```@repl tutorial_2 -using Optim, LineSearches - -lbs = [0,0,-10,-10,0,0,0,0,0]; -ubs = [1,1,10,10,1,1,1,100,100]; - -sol = optimize(x -> calculate_posterior_loglikelihood(x, prior_distributions), lbs, ubs , FS2000.parameter_values, Fminbox(LBFGS(linesearch = LineSearches.BackTracking(order = 3))); autodiff = :forward) - -sol.minimum +modeFS2000 = Turing.maximum_a_posteriori(FS2000_loglikelihood, + adtype = AutoZygote(), + initial_params = FS2000.parameter_values) ``` ## Model estimates given the data and the model solution @@ -237,7 +219,7 @@ sol.minimum Having found the parameters at the posterior mode we can retrieve model estimates of the shocks which explain the data used to estimate it. This can be done with the `get_estimated_shocks` function: ```@repl tutorial_2 -get_estimated_shocks(FS2000, data, parameters = sol.minimizer) +get_estimated_shocks(FS2000, data, parameters = collect(modeFS2000.values)) ``` As the first argument we pass the model, followed by the data (in levels), and then we pass the parameters at the posterior mode. The model is solved with this parameterisation and the shocks are calculated using the Kalman smoother. diff --git a/docs/src/tutorials/rbc.md b/docs/src/tutorials/rbc.md index d8ae074c..e377995c 100644 --- a/docs/src/tutorials/rbc.md +++ b/docs/src/tutorials/rbc.md @@ -1,6 +1,6 @@ # Write your first model - simple RBC -The following tutorial will walk you through the steps of writing down a model (not explained here / taken as given) and analysing it. Prior knowledge of DSGE models and their solution in practical terms (e.g. having used a mod file with dynare) is useful in understanding this tutorial. For the purpose of this tutorial we will work with a simplified version of a real business cycle (RBC) model. The model laid out below examines capital accumulation, consumption, and random technological progress. Households maximize lifetime utility from consumption, weighing current against future consumption. Firms produce using capital and a stochastic technology factor, setting capital rental rates based on marginal productivity. The model integrates households' decisions, firms' production, and random technological shifts to understand economic growth and dynamics. +The following tutorial will walk you through the steps of writing down a model and analysing it. Prior knowledge of DSGE models and their solution in practical terms (e.g. having used a mod file with dynare) is useful in understanding this tutorial. For the purpose of this tutorial we will work with a simplified version of a real business cycle (RBC) model. The model laid out below examines capital accumulation, consumption, and random technological progress. Households maximize lifetime utility from consumption, weighing current against future consumption. Firms produce using capital and a stochastic technology factor, setting capital rental rates based on marginal productivity. The model integrates households' decisions, firms' production, and random technological shifts to understand economic growth and dynamics. ## RBC - derivation of model equations @@ -85,7 +85,7 @@ z_{t} = \rho^z z_{t-1} + \sigma^z \epsilon^z_{t} The first step is always to name the model and write down the equations. Taking the RBC model described above this would go as follows. -First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between `begin` and `end`. One equation per line and timing of endogenous variables are expressed in the squared brackets following the variable name. Exogenous variables (shocks) are followed by a keyword in squared brackets indicating them being exogenous (in this case `[x]`). Note that names can leverage julias unicode capabilities (`alpha` can be written as `α`). +First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between `begin` and `end`. One equation per line and timing of endogenous variables are expressed in the square brackets following the variable name. Exogenous variables (shocks) are followed by a keyword in square brackets indicating them being exogenous (in this case `[x]`). Note that names can leverage julias unicode capabilities (`alpha` can be written as `α`). ```@setup tutorial_1 ENV["GKSwstype"] = "100" diff --git a/docs/src/tutorials/sw03.md b/docs/src/tutorials/sw03.md index 32e22327..f54b7e37 100644 --- a/docs/src/tutorials/sw03.md +++ b/docs/src/tutorials/sw03.md @@ -70,7 +70,7 @@ using MacroModelling end ``` -First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between begin and end. Equations can contain an equality sign or the expression is assumed to equal 0. Equations cannot span multiple lines (unless you wrap the expression in brackets) and the timing of endogenous variables are expressed in the squared brackets following the variable name (e.g. `[-1]` for the past period). Exogenous variables (shocks) are followed by a keyword in squared brackets indicating them being exogenous (in this case [x]). In this example there are also variables in the non stochastic steady state denoted by `[ss]`. Note that names can leverage julia's unicode capabilities (alpha can be written as α). +First, we load the package and then use the [`@model`](@ref) macro to define our model. The first argument after [`@model`](@ref) is the model name and will be the name of the object in the global environment containing all information regarding the model. The second argument to the macro are the equations, which we write down between begin and end. Equations can contain an equality sign or the expression is assumed to equal 0. Equations cannot span multiple lines (unless you wrap the expression in brackets) and the timing of endogenous variables are expressed in the square brackets following the variable name (e.g. `[-1]` for the past period). Exogenous variables (shocks) are followed by a keyword in square brackets indicating them being exogenous (in this case [x]). In this example there are also variables in the non stochastic steady state denoted by `[ss]`. Note that names can leverage julia's unicode capabilities (alpha can be written as α). ## Define the parameters diff --git a/docs/src/unfinished_docs/dsl.md b/docs/src/unfinished_docs/dsl.md index f23827c7..996ef952 100644 --- a/docs/src/unfinished_docs/dsl.md +++ b/docs/src/unfinished_docs/dsl.md @@ -10,7 +10,7 @@ MacroModelling parses models written using a user-friendly syntax: z[0] = ρ * z[-1] + std_z * eps_z[x] end ``` -The most important rule is that variables are followed by the timing in squared brackets for endogenous variables, e.g. `Y[0]`, exogenous variables are marked by certain keywords (see below), e.g. `ϵ[x]`, and parameters need no further syntax, e.g. `α`. +The most important rule is that variables are followed by the timing in square brackets for endogenous variables, e.g. `Y[0]`, exogenous variables are marked by certain keywords (see below), e.g. `ϵ[x]`, and parameters need no further syntax, e.g. `α`. A model written with this syntax allows the parser to identify, endogenous and exogenous variables and their timing as well as parameters. diff --git a/docs/src/unfinished_docs/todo.md b/docs/src/unfinished_docs/todo.md index 77182172..80fb2d30 100644 --- a/docs/src/unfinished_docs/todo.md +++ b/docs/src/unfinished_docs/todo.md @@ -3,7 +3,38 @@ ## High priority - [ ] ss transition by entering new parameters at given periods +- [ ] add argument to plotting functions to replace names in plots (e.g input a dictionnary: Dict(:dinve => "Investment growth")) +- [ ] programmatic model writing: accept {i}[0] as definition for variable +- [ ] check out dense sparse matmul on transposed matrices +- [ ] check out DiffInterface for NSSS solver +- [ ] write plotting callback for NSSS solder +- [ ] write more tests for the plots +- [ ] juliacon followup: checkout alloccheck, infiltrator, bestie, DifferentiableInterface, DepotDelivery, Interfaces, ThreadedDenseparseMul, Optimization Ensemble, redo Kalman filter with PDMats +- [ ] use IrrationalConstants for log2pi... +- [ ] checkout this invalidation precompile trick and g dalle part on precompilation +- [ ] use sobol random numbers (gives you uniform but then use norminvcdf to get norm) to integrate out future randomness when solving with neural nets +- [ ] do proper testing of ss solver with random set of params, equal across configs +- [ ] load create parts of derivatives later and not directly after parameters block +- [ ] fix model estimate plot. data not above estimate (should be red but is blue) +- [ ] make plotting options as dynamic setting instead of default, accept kwargs +- [ ] analytical derivatives of inversion filter +- [ ] implement higher order (pruned) variance decomposition +- [ ] try slicesampler instead of pigeons +- [ ] use faster derivatives for SS solver (currently forward diff) +- [ ] speed up sensitivity by caching matrix inversion from implicit diff with LRUcache +- [ ] FastDifferentiation is faster in taking derivatives and more efficient in writing functions but does not support custom functions (e.g. normlogpdf) +- [ ] fix this inference errors for large fuctions. they are slow. fix derivatives in general. - [ ] check downgrade tests +- [ ] put write_derivatives_function and lock structure inside function +- [ ] take apart solve_matrix_equation for various cases +- [ ] try static arrays in KF +- [ ] check derivatives of erfcinv with Symbolics. seems off +- [ ] have a workspace in the model object. to be accessed for example by the riccati solver at each run (instead of initialising values at each function call) +- [ ] check why PG samples are off +- [ ] implement estimation tests for all models +- [ ] optimise vanilla loglikelihood calculation and gradient thereof (incl comp time) +- [ ] checkout dynamic perturbation for obc solution: https://www.southampton.ac.uk/~alexmen/dynamic_perturbation.pdf +- [ ] checkout schedule free ADAM for global methods: https://github.com/facebookresearch/schedule_free - [ ] figure out why PG and IS return basically the prior - [ ] allow external functions to calculate the steady state (and hand it over via SS or get_loglikelihood function) - need to use the check function for implicit derivatives and cannot use it to get him a guess from which he can use internal solver going forward - [ ] go through custom SS solver once more and try to find parameters and logic that achieves best results @@ -36,6 +67,7 @@ - [ ] make inversion filter / higher order sols suitable for HMC (forward and reverse diff!!, currently only analytical pushforward, no implicitdiff) | analytic derivatives - [ ] speed up sparse matrix calcs in implicit diff of higher order funcs - [ ] compressed higher order derivatives and sparsity of jacobian +- [ ] dont use SS_solve_func but the wrapper instead (write forwarddiff wrapper) - [ ] add user facing option to choose sylvester solver - [ ] autocorr and covariance with derivatives. return 3d array - [ ] use ID for sparse output sylvester solvers (filed issue) @@ -57,7 +89,7 @@ - [ ] use cache for gradient calc in estimation (see DifferentiableStateSpaceModels) - [ ] write functions to debug (fix_SS.jl...) - [ ] model compression (speed up 2nd moment calc (derivatives) for large models; gradient loglikelihood is very slow due to large matmuls) -> model setup as maximisation problem (gEcon) -> HANK models -- [ ] implement global solution methods +- [ ] implement global solution methods - Julien Pascal, QuantEcon - [ ] add more models - [ ] use @assert for errors and @test_throws diff --git a/models/Smets_Wouters_2007.jl b/models/Smets_Wouters_2007.jl index bf083adf..c791ab86 100644 --- a/models/Smets_Wouters_2007.jl +++ b/models/Smets_Wouters_2007.jl @@ -7,21 +7,21 @@ kp[0] = inve[0] * qs[0] * (1 - Sfunc[0]) + kp[-1] * (1 - ctou) / cgamma - pdot[0] = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + pdot[-1] * cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * pinf[ss] ^ (1 - cindp) / pinf[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + pdot[0] = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + pdot[-1] * cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) - wdot[0] = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) + wdot[-1] * cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * pinf[ss] ^ (1 - cindw) / pinf[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) + wdot[0] = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) + wdot[-1] * cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) - 1 = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * pinf[ss] ^ (1 - cindp) / pinf[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + 1 = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) - 1 = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) + cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * pinf[ss] ^ (1 - cindw) / pinf[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) + 1 = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) + cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) 1 = dp[0] * (1 + pdotl[0] * curvP) / (1 + curvP) w[0] = dw[0] * (1 + curvW * wdotl[0]) / (1 + curvW) - pdotl[0] = (1 - cprobp) * Pratio[0] / dp[0] + cprobp * dp[-1] / dp[0] * pinf[-1] ^ cindp * pinf[ss] ^ (1 - cindp) / pinf[0] * pdotl[-1] + pdotl[0] = (1 - cprobp) * Pratio[0] / dp[0] + cprobp * dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0] * pdotl[-1] - wdotl[0] = (1 - cprobw) * wnew[0] / dw[0] + cprobw * dw[-1] / dw[0] * pinf[-1] ^ cindw * pinf[ss] ^ (1 - cindw) / pinf[0] * wdotl[-1] + wdotl[0] = (1 - cprobw) * wnew[0] / dw[0] + cprobw * dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0] * wdotl[-1] xi[0] = exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) * (c[0] - c[-1] * chabb / cgamma) ^ (-csigma) @@ -39,19 +39,19 @@ wnew[0] * gamw1[0] * (1 + curvw * (1 - clandaw)) / (1 + curvW) = clandaw * gamw2[0] + gamw3[0] * curvW * (clandaw - 1) / (1 + curvW) * wnew[0] ^ (1 + clandaw * (1 + curvW) / (clandaw - 1)) - gamw1[0] = lab[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) + gamw1[1] * (pinf[ss] ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * cbetabar + gamw1[0] = lab[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) + gamw1[1] * (cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * cbetabar - gamw2[0] = (c[0] - c[-1] * chabb / cgamma) * lab[0] * sw[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ csigl + gamw2[1] * (pinf[ss] ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * cbetabar + gamw2[0] = (c[0] - c[-1] * chabb / cgamma) * lab[0] * sw[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ csigl + gamw2[1] * (cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * cbetabar - gamw3[0] = lab[0] + gamw3[1] * pinf[ss] ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1] * xi[1] / xi[0] * cgamma * cprobw * cbetabar + gamw3[0] = lab[0] + gamw3[1] * cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1] * xi[1] / xi[0] * cgamma * cprobw * cbetabar Pratio[0] * gam1[0] * (1 + curvp * (1 - cfc)) / (1 + curvP) = cfc * gam2[0] + gam3[0] * (cfc - 1) * curvP / (1 + curvP) * Pratio[0] ^ (1 + cfc * (1 + curvP) / (cfc - 1)) - gam1[0] = y[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam1[1] * xi[1] / xi[0] * cgamma * cprobp * cbetabar * (pinf[ss] ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + gam1[0] = y[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam1[1] * xi[1] / xi[0] * cgamma * cprobp * cbetabar * (cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) - gam2[0] = y[0] * mc[0] * spinf[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam2[1] * xi[1] / xi[0] * cgamma * cprobp * cbetabar * (pinf[ss] ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + gam2[0] = y[0] * mc[0] * spinf[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam2[1] * xi[1] / xi[0] * cgamma * cprobp * cbetabar * (cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) - gam3[0] = y[0] + gam3[1] * pinf[ss] ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1] * xi[1] / xi[0] * cgamma * cprobp * cbetabar + gam3[0] = y[0] + gam3[1] * cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1] * xi[1] / xi[0] * cgamma * cprobp * cbetabar qsaux[0] = qs[1] @@ -67,23 +67,23 @@ SfuncD[0] = csadjcost * (cgamma * inve[0] / inve[-1] - cgamma) - a[0] = 1 - crhoa + crhoa * a[-1] + z_ea * ea[x] + a[0] = 1 - crhoa + crhoa * a[-1] + z_ea / 100 * ea[x] - b[0] = 1 - crhob + crhob * b[-1] + z_eb * eb[x] + b[0] = 1 - crhob + crhob * b[-1] + z_eb / 100 * SCALE1_eb * eb[x] - gy[0] - cg = crhog * (gy[-1] - cg) + z_eg * eg[x] + z_ea * ea[x] * cgy + gy[0] - cg = crhog * (gy[-1] - cg) + z_eg / 100 * eg[x] + z_ea / 100 * ea[x] * cgy - qs[0] = 1 - crhoqs + crhoqs * qs[-1] + z_eqs * eqs[x] + qs[0] = 1 - crhoqs + crhoqs * qs[-1] + z_eqs / 100 * SCALE1_eqs * eqs[x] - ms[0] = 1 - crhoms + crhoms * ms[-1] + z_em * em[x] + ms[0] = 1 - crhoms + crhoms * ms[-1] + z_em / 100 * em[x] spinf[0] = 1 - crhopinf + crhopinf * spinf[-1] + epinfma[0] - cmap * epinfma[-1] - epinfma[0] = z_epinf * epinf[x] + epinfma[0] = z_epinf / 100 * SCALE1_epinf * epinf[x] sw[0] = 1 - crhow + crhow * sw[-1] + ewma[0] - cmaw * ewma[-1] - ewma[0] = z_ew * ew[x] + ewma[0] = z_ew / 100 * SCALE1_ew * ew[x] yflex[0] = cflex[0] + inveflex[0] + gy[0] * yflex[ss] + afuncflex[0] * kpflex[-1] / cgamma @@ -127,18 +127,26 @@ dinve[0] = ctrend + 100 * (inve[0] / inve[-1] - 1) - pinfobs[0] = 100 * (pinf[0] - pinf[ss]) + constepinf + pinfobs[0] = 100 * (pinf[0] - 1) robs[0] = 100 * (r[0] - 1) dwobs[0] = ctrend + 100 * (w[0] / w[-1] - 1) - labobs[0] = 100 * (lab[0] / lab[ss] - 1) + labobs[0] = constelab + 100 * (lab[0] / lab[ss] - 1) end @parameters Smets_Wouters_2007 begin + SCALE1_eb = -((1 - chabb / cgamma) / (csigma * (1 + chabb / cgamma))) ^ (-1) + + SCALE1_eqs = (cgamma ^ 2 * csadjcost) * (1 + cbeta * cgamma ^ (1 - csigma)) + + SCALE1_epinf = 1 / ((1 / (1 + cbetabar * cgamma * cindp)) * ((1 - cprobp) * (1 - cbetabar * cgamma * cprobp) / cprobp) / ((cfc - 1) * curvp + 1)) + + SCALE1_ew = 1 / ((1 - cprobw) * (1 - cbetabar * cgamma * cprobw) / ((1 + cbetabar * cgamma) * cprobw) * (1 / ((clandaw - 1) * curvw + 1))) + cgamma = 1 + ctrend / 100 # gross growth rate cbeta = 1 / (1 + constebeta / 100) # discount factor diff --git a/models/Smets_Wouters_2007_linear.jl b/models/Smets_Wouters_2007_linear.jl index f3853d77..40eb7e5f 100644 --- a/models/Smets_Wouters_2007_linear.jl +++ b/models/Smets_Wouters_2007_linear.jl @@ -72,12 +72,12 @@ dinve[0] = ctrend + inve[0] - inve[-1] - dw[0] = ctrend + w[0] - w[-1] - pinfobs[0] = constepinf + pinf[0] robs[0] = r[0] + conster + dwobs[0] = ctrend + w[0] - w[-1] + labobs[0] = lab[0] + constelab end diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index 5d9982e0..30337332 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -6,8 +6,10 @@ import DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF, TYPEDSIGNATURES, TYPEDF import ThreadedSparseArrays using PrecompileTools import SpecialFunctions: erfcinv, erfc +import SpecialFunctions import SymPyPythonCall as SPyPyC import Symbolics +import Accessors # import NaNMath # import Memoization: @memoize # import LRUCache: LRU @@ -18,11 +20,15 @@ import ForwardDiff as ℱ 𝒷 = 𝒜.ForwardDiffBackend # 𝒷 = Diffractor.DiffractorForwardBackend +import Polyester import NLopt import Optim, LineSearches # import Zygote -import SparseArrays: SparseMatrixCSC, SparseVector, AbstractSparseArray#, sparse, spzeros, droptol!, sparsevec, spdiagm, findnz#, sparse! +import SparseArrays: SparseMatrixCSC, SparseVector, AbstractSparseArray, sparse! #, sparse, spzeros, droptol!, sparsevec, spdiagm, findnz#, sparse! import LinearAlgebra as ℒ +import LinearAlgebra: mul! +# import Octavian: matmul! +# import TriangularSolve as TS # import ComponentArrays as 𝒞 import Combinatorics: combinations import BlockTriangularForm @@ -40,7 +46,7 @@ import MatrixEquations # good overview: https://cscproxy.mpi-magdeburg.mpg.de/mp # using NamedArrays # using AxisKeys -import ChainRulesCore: @ignore_derivatives, ignore_derivatives, rrule, NoTangent +import ChainRulesCore: @ignore_derivatives, ignore_derivatives, rrule, NoTangent, @thunk import RecursiveFactorization as RF using RuntimeGeneratedFunctions @@ -272,7 +278,7 @@ function reverse_transformation(transformed_expr::Expr, reverse_dict::Dict{Symbo end -function transform_obc(ex::Expr) +function transform_obc(ex::Expr; avoid_solve::Bool = false) transformed_expr, reverse_dict = transform_expression(ex) for symbs in get_symbols(transformed_expr) @@ -281,8 +287,12 @@ function transform_obc(ex::Expr) eq = eval(transformed_expr) - soll = try SPyPyC.solve(eq, eval(:minmax__P)) - catch + if avoid_solve && count_ops(Meta.parse(string(eq))) > 15 + soll = nothing + else + soll = try SPyPyC.solve(eq, eval(:minmax__P)) + catch + end end if length(soll) == 1 @@ -757,6 +767,69 @@ function mat_mult_kron(A::AbstractArray{T},B::AbstractArray{T},C::AbstractArray{ sparse(rows,cols,vals,size(A,1),n_colB*n_colC) end +function kron³(A::SparseMatrixCSC{T}, M₃::third_order_auxilliary_matrices) where T <: Real + rows, cols, vals = findnz(A) + + # Dictionary to accumulate sums of values for each coordinate + result_dict = Dict{Tuple{Int, Int}, T}() + + # Using a single iteration over non-zero elements + nvals = length(vals) + + lk = ReentrantLock() + + Polyester.@batch for i in 1:nvals + for j in 1:nvals + for k in 1:nvals + r1, c1, v1 = rows[i], cols[i], vals[i] + r2, c2, v2 = rows[j], cols[j], vals[j] + r3, c3, v3 = rows[k], cols[k], vals[k] + + sorted_cols = [c1, c2, c3] + sorted_rows = [r1, r2, r3] # a lot of time spent here + sort!(sorted_rows, rev = true) # a lot of time spent here + + if haskey(M₃.𝐈₃, sorted_cols) # && haskey(M₃.𝐈₃, sorted_rows) # a lot of time spent here + row_idx = M₃.𝐈₃[sorted_rows] + col_idx = M₃.𝐈₃[sorted_cols] + + key = (row_idx, col_idx) + + begin + lock(lk) + try + if haskey(result_dict, key) + result_dict[key] += v1 * v2 * v3 + else + result_dict[key] = v1 * v2 * v3 + end + finally + unlock(lk) + end + end + end + end + end + end + + # Extract indices and values from the dictionary + result_rows = Int[] + result_cols = Int[] + result_vals = T[] + + for (ks, valu) in result_dict + push!(result_rows, ks[1]) + push!(result_cols, ks[2]) + push!(result_vals, valu) + end + + # Create the sparse matrix from the collected indices and values + if VERSION >= v"1.10" + return sparse!(result_rows, result_cols, result_vals, size(M₃.𝐂₃, 2), size(M₃.𝐔₃, 1)) + else + return sparse(result_rows, result_cols, result_vals, size(M₃.𝐂₃, 2), size(M₃.𝐔₃, 1)) + end +end function A_mult_kron_power_3_B(A::AbstractArray{R},B::AbstractArray{T}; tol::AbstractFloat = eps()) where {R <: Real, T <: Real} n_row = size(B,1) @@ -1072,9 +1145,9 @@ function get_and_check_observables(𝓂::ℳ, data::KeyedArray{Float64})::Vector @assert length(setdiff(observables_symbols, 𝓂.var)) == 0 "The following symbols in the first axis of the conditions matrix are not part of the model: " * repr(setdiff(observables_symbols,𝓂.var)) - sort!(observables) + sort!(observables_symbols) - return observables + return observables_symbols end function bivariate_moment(moment::Vector{Int}, rho::Int)::Int @@ -1227,8 +1300,20 @@ function match_pattern(strings::Union{Set,Vector}, pattern::Regex) return filter(r -> match(pattern, string(r)) !== nothing, strings) end + +function count_ops(expr) + op_count = 0 + postwalk(x -> begin + if x isa Expr && x.head == :call + op_count += 1 + end + x + end, expr) + return op_count +end + # try: run optim only if there is a violation / capture case with small shocks and set them to zero -function parse_occasionally_binding_constraints(equations_block; max_obc_horizon::Int = 40) +function parse_occasionally_binding_constraints(equations_block; max_obc_horizon::Int = 40, avoid_solve::Bool = false) # precision_factor = 1e #factor to force the optimiser to have non-relevatn shocks at zero eqs = [] @@ -1796,10 +1881,11 @@ function levenberg_marquardt(f::Function, # return undo_transform(current_guess,transformation_level,shift), (iter, Inf, Inf, upper_bounds) end - ∇̄ = RF.lu(∇̂, check = false) + ∇̄ = RF.lu!(∇̂, check = false) if !ℒ.issuccess(∇̄) - ∇̄ = ℒ.svd(∇̂) + return undo_transform(current_guess,transformation_level), (iter, Inf, Inf, upper_bounds) + # ∇̄ = ℒ.svd(∇̂) end current_guess .-= ∇̄ \ ∇' * f̂(current_guess) @@ -1845,8 +1931,10 @@ function levenberg_marquardt(f::Function, α̂ = min(α̂, ϕ̄ * α) α = max(α̂, ϕ̂ * α) - - current_guess .= previous_guess + α * guess_update + + copy!(current_guess, previous_guess) + ℒ.axpy!(α, guess_update, current_guess) + # current_guess .= previous_guess + α * guess_update minmax!(current_guess, lower_bounds, upper_bounds) P = P̋ @@ -2007,7 +2095,7 @@ end -function remove_redundant_SS_vars!(𝓂::ℳ, Symbolics::symbolics) +function remove_redundant_SS_vars!(𝓂::ℳ, Symbolics::symbolics; avoid_solve::Bool = false) ss_equations = Symbolics.ss_equations # check variables which appear in two time periods. they might be redundant in steady state @@ -2025,11 +2113,15 @@ function remove_redundant_SS_vars!(𝓂::ℳ, Symbolics::symbolics) redundant_idx = getindex(1:length(redundant_vars), (length.(redundant_vars) .> 0) .& (length.(Symbolics.var_list) .> 1)) for i in redundant_idx - for var_to_solve_for in redundant_vars[i] - soll = try SPyPyC.solve(ss_equations[i],var_to_solve_for) - catch + for var_to_solve_for in redundant_vars[i] + if avoid_solve && count_ops(Meta.parse(string(ss_equations[i]))) > 15 + soll = nothing + else + soll = try SPyPyC.solve(ss_equations[i],var_to_solve_for) + catch + end end - + if isnothing(soll) continue end @@ -2448,7 +2540,7 @@ end -function partial_solve(eqs_to_solve, vars_to_solve, incidence_matrix_subset) +function partial_solve(eqs_to_solve, vars_to_solve, incidence_matrix_subset; avoid_solve::Bool = false) for n in length(eqs_to_solve)-1:-1:2 for eq_combo in combinations(1:length(eqs_to_solve), n) var_indices_to_select_from = findall([sum(incidence_matrix_subset[:,eq_combo],dims = 2)...] .> 0) @@ -2458,9 +2550,13 @@ function partial_solve(eqs_to_solve, vars_to_solve, incidence_matrix_subset) for var_combo in combinations(var_indices_to_select_from, n) remaining_vars_in_remaining_eqs = setdiff(var_indices_in_remaining_eqs, var_combo) # println("Solving for: ",vars_to_solve[var_combo]," in: ",eqs_to_solve[eq_combo]) - if length(remaining_vars_in_remaining_eqs) == length(eqs_to_solve) - n # not sure whether this condition needs to be there. could be because if the last remaining vars not solved for in the block is not present in the remaining block he will not be able to solve it for the same reasons he wasnt able to solve the unpartitioned block - soll = try SPyPyC.solve(eqs_to_solve[eq_combo], vars_to_solve[var_combo]) - catch + if length(remaining_vars_in_remaining_eqs) == length(eqs_to_solve) - n # not sure whether this condition needs to be there. could be because if the last remaining vars not solved for in the block is not present in the remaining block he will not be able to solve it for the same reasons he wasnt able to solve the unpartitioned block + if avoid_solve && count_ops(Meta.parse(string(eqs_to_solve[eq_combo]))) > 15 + soll = nothing + else + soll = try SPyPyC.solve(eqs_to_solve[eq_combo], vars_to_solve[var_combo]) + catch + end end if !(isnothing(soll) || length(soll) == 0) @@ -3056,11 +3152,11 @@ end function write_ss_check_function!(𝓂::ℳ) - vars_in_ss_equations = sort(collect(setdiff(reduce(union,get_symbols.(𝓂.ss_equations)),union(𝓂.parameters_in_equations)))) + # vars_in_ss_equations = sort(collect(setdiff(reduce(union,get_symbols.(𝓂.ss_equations)),union(𝓂.parameters_in_equations)))) - unknowns = union(vars_in_ss_equations, 𝓂.calibration_equations_parameters) + unknowns = union(setdiff(𝓂.vars_in_ss_equations, 𝓂.➕_vars), 𝓂.calibration_equations_parameters) - ss_equations = vcat(𝓂.ss_equations,𝓂.calibration_equations) + ss_equations = vcat(𝓂.ss_equations, 𝓂.calibration_equations) pars = [] for (i, p) in enumerate(𝓂.parameters) @@ -3068,7 +3164,7 @@ function write_ss_check_function!(𝓂::ℳ) end unknwns = [] - for (i, u) in enumerate(union(vars_in_ss_equations, 𝓂.calibration_equations_parameters)) + for (i, u) in enumerate(unknowns) push!(unknwns, :($u = unknowns[$i])) end @@ -3083,7 +3179,7 @@ function write_ss_check_function!(𝓂::ℳ) end -function solve_steady_state!(𝓂::ℳ, symbolic_SS, Symbolics::symbolics; verbose::Bool = false) +function solve_steady_state!(𝓂::ℳ, symbolic_SS, Symbolics::symbolics; verbose::Bool = false, avoid_solve::Bool = false) write_ss_check_function!(𝓂) unknowns = union(Symbolics.vars_in_ss_equations, Symbolics.calibration_equations_parameters) @@ -3162,9 +3258,13 @@ function solve_steady_state!(𝓂::ℳ, symbolic_SS, Symbolics::symbolics; verbo push!(min_max_errors,:(solution_error += abs($parsed_eq_to_solve_for))) eq_to_solve = eval(minmax_fixed_eqs) end - - soll = try SPyPyC.solve(eq_to_solve,var_to_solve_for) - catch + + if avoid_solve && count_ops(Meta.parse(string(eq_to_solve))) > 15 + soll = nothing + else + soll = try SPyPyC.solve(eq_to_solve,var_to_solve_for) + catch + end end if isnothing(soll) || isempty(soll) @@ -3223,8 +3323,12 @@ function solve_steady_state!(𝓂::ℳ, symbolic_SS, Symbolics::symbolics; verbo numerical_sol = false if symbolic_SS - soll = try SPyPyC.solve(eqs_to_solve,vars_to_solve) - catch + if avoid_solve && count_ops(Meta.parse(string(eqs_to_solve))) > 15 + soll = nothing + else + soll = try SPyPyC.solve(eqs_to_solve,vars_to_solve) + catch + end end if isnothing(soll) || length(soll) == 0 || length(intersect((union(SPyPyC.free_symbols.(soll[1])...) .|> SPyPyC.:↓),(vars_to_solve .|> SPyPyC.:↓))) > 0 @@ -3261,7 +3365,7 @@ function solve_steady_state!(𝓂::ℳ, symbolic_SS, Symbolics::symbolics; verbo write_block_solution!(𝓂, SS_solve_func, vars_to_solve, eqs_to_solve, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list) # write_domain_safe_block_solution!(𝓂, SS_solve_func, vars_to_solve, eqs_to_solve, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, unique_➕_eqs) else - solved_system = partial_solve(eqs_to_solve[pe], vars_to_solve[pv], incidence_matrix_subset[pv,pe]) + solved_system = partial_solve(eqs_to_solve[pe], vars_to_solve[pv], incidence_matrix_subset[pv,pe], avoid_solve = avoid_solve) # if !isnothing(solved_system) && !any(contains.(string.(vcat(solved_system[3],solved_system[4])), "LambertW")) && !any(contains.(string.(vcat(solved_system[3],solved_system[4])), "Heaviside")) # write_reduced_block_solution!(𝓂, SS_solve_func, solved_system, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, @@ -4187,16 +4291,18 @@ second_order_stochastic_steady_state_iterative_solution = ℐ.ImplicitFunction(s function calculate_second_order_stochastic_steady_state(parameters::Vector{M}, 𝓂::ℳ; verbose::Bool = false, pruning::Bool = false, tol::AbstractFloat = 1e-12)::Tuple{Vector{M}, Bool, Vector{M}, M, AbstractMatrix{M}, SparseMatrixCSC{M}, AbstractMatrix{M}, SparseMatrixCSC{M}} where M SS_and_pars, (solution_error, iters) = 𝓂.SS_solve_func(parameters, 𝓂, verbose, false, 𝓂.solver_parameters) + all_SS = expand_steady_state(SS_and_pars,𝓂) + if solution_error > tol || isnan(solution_error) - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) end - ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂)# |> Matrix 𝐒₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) if !solved - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) end ∇₂ = calculate_hessian(parameters, SS_and_pars, 𝓂) @@ -4204,7 +4310,7 @@ function calculate_second_order_stochastic_steady_state(parameters::Vector{M}, 𝐒₂, solved2 = calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, 𝓂.solution.perturbation.second_order_auxilliary_matrices; T = 𝓂.timings) if !solved2 - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) end 𝐒₁ = [𝐒₁[:,1:𝓂.timings.nPast_not_future_and_mixed] zeros(𝓂.timings.nVars) 𝐒₁[:,𝓂.timings.nPast_not_future_and_mixed+1:end]] @@ -4217,7 +4323,7 @@ function calculate_second_order_stochastic_steady_state(parameters::Vector{M}, tmp̄ = RF.lu(tmp, check = false) if !ℒ.issuccess(tmp̄) - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0) end SSSstates = tmp̄ \ (𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2)[𝓂.timings.past_not_future_and_mixed_idx] @@ -4228,8 +4334,6 @@ function calculate_second_order_stochastic_steady_state(parameters::Vector{M}, state, converged = second_order_stochastic_steady_state_iterative_solution([sparsevec(𝐒₁); vec(𝐒₂)]; dims = [size(𝐒₁); size(𝐒₂)], 𝓂 = 𝓂) end - all_SS = expand_steady_state(SS_and_pars,𝓂) - # all_variables = sort(union(𝓂.var,𝓂.aux,𝓂.exo_present)) # all_variables[indexin(𝓂.aux,all_variables)] = map(x -> Symbol(replace(string(x), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")), 𝓂.aux) @@ -4333,16 +4437,18 @@ end function calculate_third_order_stochastic_steady_state(parameters::Vector{M}, 𝓂::ℳ; verbose::Bool = false, pruning::Bool = false, tol::AbstractFloat = 1e-12)::Tuple{Vector{M}, Bool, Vector{M}, M, AbstractMatrix{M}, SparseMatrixCSC{M}, SparseMatrixCSC{M}, AbstractMatrix{M}, SparseMatrixCSC{M}, SparseMatrixCSC{M}} where M SS_and_pars, (solution_error, iters) = 𝓂.SS_solve_func(parameters, 𝓂, verbose, false, 𝓂.solver_parameters) + all_SS = expand_steady_state(SS_and_pars,𝓂) + if solution_error > tol || isnan(solution_error) - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) end - ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂)# |> Matrix 𝐒₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) if !solved - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) end ∇₂ = calculate_hessian(parameters, SS_and_pars, 𝓂) @@ -4350,7 +4456,7 @@ function calculate_third_order_stochastic_steady_state(parameters::Vector{M}, 𝐒₂, solved2 = calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, 𝓂.solution.perturbation.second_order_auxilliary_matrices; T = 𝓂.timings, tol = tol) if !solved2 - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) end ∇₃ = calculate_third_order_derivatives(parameters, SS_and_pars, 𝓂) @@ -4358,7 +4464,7 @@ function calculate_third_order_stochastic_steady_state(parameters::Vector{M}, 𝐒₃, solved3 = calculate_third_order_solution(∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝓂.solution.perturbation.second_order_auxilliary_matrices, 𝓂.solution.perturbation.third_order_auxilliary_matrices; T = 𝓂.timings, tol = tol) if !solved3 - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) end 𝐒₁ = [𝐒₁[:,1:𝓂.timings.nPast_not_future_and_mixed] zeros(𝓂.timings.nVars) 𝐒₁[:,𝓂.timings.nPast_not_future_and_mixed+1:end]] @@ -4371,7 +4477,7 @@ function calculate_third_order_stochastic_steady_state(parameters::Vector{M}, tmp̄ = RF.lu(tmp, check = false) if !ℒ.issuccess(tmp̄) - return SS_and_pars, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) + return all_SS, false, SS_and_pars, solution_error, zeros(0,0), spzeros(0,0), spzeros(0,0), zeros(0,0), spzeros(0,0), spzeros(0,0) end SSSstates = tmp̄ \ (𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2)[𝓂.timings.past_not_future_and_mixed_idx] @@ -4382,8 +4488,6 @@ function calculate_third_order_stochastic_steady_state(parameters::Vector{M}, state, converged = third_order_stochastic_steady_state_iterative_solution([sparsevec(𝐒₁); vec(𝐒₂); vec(𝐒₃)]; dims = [size(𝐒₁); size(𝐒₂); size(𝐒₃)], 𝓂 = 𝓂) end - all_SS = expand_steady_state(SS_and_pars,𝓂) - # all_variables = sort(union(𝓂.var,𝓂.aux,𝓂.exo_present)) # all_variables[indexin(𝓂.aux,all_variables)] = map(x -> Symbol(replace(string(x), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")), 𝓂.aux) @@ -4412,14 +4516,17 @@ function solve!(𝓂::ℳ; write_parameters_input!(𝓂, parameters, verbose = verbose) - if 𝓂.model_hessian == Function[] && algorithm ∈ [:second_order, :pruned_second_order] + if 𝓂.solution.perturbation.second_order_auxilliary_matrices.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && + algorithm ∈ [:second_order, :pruned_second_order] start_time = time() + if !silent print("Take symbolic derivatives up to second order:\t\t\t\t") end write_functions_mapping!(𝓂, 2) - if !silent println("Take symbolic derivatives up to second order:\t",round(time() - start_time, digits = 3), " seconds") end - elseif 𝓂.model_third_order_derivatives == Function[] && algorithm ∈ [:third_order, :pruned_third_order] + if !silent println(round(time() - start_time, digits = 3), " seconds") end + elseif 𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && algorithm ∈ [:third_order, :pruned_third_order] start_time = time() + if !silent print("Take symbolic derivatives up to third order:\t\t\t\t") end write_functions_mapping!(𝓂, 3) - if !silent println("Take symbolic derivatives up to third order:\t",round(time() - start_time, digits = 3), " seconds") end + if !silent println(round(time() - start_time, digits = 3), " seconds") end end if dynamics @@ -4437,7 +4544,7 @@ function solve!(𝓂::ℳ; @warn "Could not find non stochastic steady steady." end - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix S₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -4452,7 +4559,7 @@ function solve!(𝓂::ℳ; if obc write_parameters_input!(𝓂, :activeᵒᵇᶜshocks => 1, verbose = false) - ∇̂₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇̂₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix Ŝ₁, solved = calculate_first_order_solution(∇̂₁; T = 𝓂.timings) @@ -4623,7 +4730,7 @@ function solve!(𝓂::ℳ; @warn "Could not find non stochastic steady steady." end - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)#|> Matrix S₁, converged = calculate_quadratic_iteration_solution(∇₁; T = 𝓂.timings) @@ -4636,7 +4743,7 @@ function solve!(𝓂::ℳ; if obc write_parameters_input!(𝓂, :activeᵒᵇᶜshocks => 1, verbose = false) - ∇̂₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇̂₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix Ŝ₁, converged = calculate_quadratic_iteration_solution(∇₁; T = 𝓂.timings) @@ -4668,7 +4775,7 @@ function solve!(𝓂::ℳ; @warn "Could not find non stochastic steady steady." end - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix S₁ = calculate_linear_time_iteration_solution(∇₁; T = 𝓂.timings) @@ -4681,7 +4788,7 @@ function solve!(𝓂::ℳ; if obc write_parameters_input!(𝓂, :activeᵒᵇᶜshocks => 1) - ∇̂₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇̂₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix Ŝ₁, converged = calculate_linear_time_iteration_solution(∇₁; T = 𝓂.timings) @@ -4737,6 +4844,13 @@ end +function add_sparse_entries!(P, perm) + n = size(P, 1) + for i in 1:n + P[perm[i], i] += 1.0 + end +end + function create_third_order_auxilliary_matrices(T::timings, ∇₃_col_indices::Vector{Int}) # Indices and number of variables @@ -4752,7 +4866,7 @@ function create_third_order_auxilliary_matrices(T::timings, ∇₃_col_indices:: colls3 = [nₑ₋^2 * (i-1) + nₑ₋ * (k-1) + l for i in 1:nₑ₋ for k in 1:i for l in 1:k] 𝐂∇₃ = sparse(colls3, 1:length(colls3) , 1.0) - idxs = [] + idxs = Int[] for k in 1:nₑ₋ for j in 1:nₑ₋ for i in 1:nₑ₋ @@ -4769,7 +4883,7 @@ function create_third_order_auxilliary_matrices(T::timings, ∇₃_col_indices:: colls3 = [nₑ₋^2 * (i-1) + nₑ₋ * (k-1) + l for i in 1:nₑ₋ for k in 1:i for l in 1:k] 𝐂₃ = sparse(colls3, 1:length(colls3) , 1.0) - idxs = [] + idxs = Int[] for k in 1:nₑ₋ for j in 1:nₑ₋ for i in 1:nₑ₋ @@ -4781,11 +4895,31 @@ function create_third_order_auxilliary_matrices(T::timings, ∇₃_col_indices:: 𝐔₃ = 𝐂₃' * sparse(idxs,1:nₑ₋ ^ 3, 1) + # Precompute 𝐈₃ + 𝐈₃ = Dict{Vector{Int}, Int}() + idx = 1 + for i in 1:nₑ₋ + for k in 1:i + for l in 1:k + 𝐈₃[[i,k,l]] = idx + idx += 1 + end + end + end + # permutation matrices M = reshape(1:nₑ₋^3,1,nₑ₋,nₑ₋,nₑ₋) - 𝐏 = @views sparse(reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 4, 2, 3])],nₑ₋^3,nₑ₋^3) - + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 4, 3])],nₑ₋^3,nₑ₋^3) - + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 3, 4])],nₑ₋^3,nₑ₋^3)) + + 𝐏 = spzeros(nₑ₋^3, nₑ₋^3) # Preallocate the sparse matrix + + # Create the permutations directly + add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 4, 2, 3))) + add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 2, 4, 3))) + add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 2, 3, 4))) + + # 𝐏 = @views sparse(reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 4, 2, 3])],nₑ₋^3,nₑ₋^3) + # + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 4, 3])],nₑ₋^3,nₑ₋^3) + # + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 3, 4])],nₑ₋^3,nₑ₋^3)) 𝐏₁ₗ = sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(2,1,3))),:]) 𝐏₁ᵣ = sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(2,1,3)))]) @@ -4802,7 +4936,7 @@ function create_third_order_auxilliary_matrices(T::timings, ∇₃_col_indices:: ∇₃_col_indices_extended = findnz(sparse(ones(Int,length(∇₃_col_indices)),∇₃_col_indices,ones(Int,length(∇₃_col_indices)),1,size(𝐔∇₃,1)) * 𝐔∇₃)[2] - nonnull_columns = Set() + nonnull_columns = Set{Int}() for i in 1:n̄ for j in i:n̄ for k in j:n̄ @@ -4817,95 +4951,185 @@ function create_third_order_auxilliary_matrices(T::timings, ∇₃_col_indices:: 𝐒𝐏 = sparse(collect(nonnull_columns), collect(nonnull_columns), 1, n̄, n̄) - return third_order_auxilliary_matrices(𝐂₃, 𝐔₃, 𝐔∇₃, 𝐏, 𝐏₁ₗ, 𝐏₁ᵣ, 𝐏₁ₗ̂, 𝐏₂ₗ̂, 𝐏₁ₗ̄, 𝐏₂ₗ̄, 𝐏₁ᵣ̃, 𝐏₂ᵣ̃, 𝐒𝐏) + return third_order_auxilliary_matrices(𝐂₃, 𝐔₃, 𝐈₃, 𝐔∇₃, 𝐏, 𝐏₁ₗ, 𝐏₁ᵣ, 𝐏₁ₗ̂, 𝐏₂ₗ̂, 𝐏₁ₗ̄, 𝐏₂ₗ̄, 𝐏₁ᵣ̃, 𝐏₂ᵣ̃, 𝐒𝐏) end -function write_functions_mapping!(𝓂::ℳ, max_perturbation_order::Int) - future_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₁₎$"))) - present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₀₎$"))) - past_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₋₁₎$"))) - shock_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₓ₎$"))) - ss_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₛₛ₎$"))) - sort!(future_varss ,by = x->replace(string(x),r"₍₁₎$"=>"")) #sort by name without time index because otherwise eps_zᴸ⁽⁻¹⁾₍₋₁₎ comes before eps_z₍₋₁₎ - sort!(present_varss ,by = x->replace(string(x),r"₍₀₎$"=>"")) - sort!(past_varss ,by = x->replace(string(x),r"₍₋₁₎$"=>"")) - sort!(shock_varss ,by = x->replace(string(x),r"₍ₓ₎$"=>"")) - sort!(ss_varss ,by = x->replace(string(x),r"₍ₛₛ₎$"=>"")) - steady_state = [] - for (i, var) in enumerate(ss_varss) - push!(steady_state,:($var = X̄[$i])) - # ii += 1 - end +function write_sparse_derivatives_function(rows::Vector{Int},columns::Vector{Int},values::Vector{Symbolics.Num},nrows::Int,ncolumns::Int,::Val{:Symbolics}) + vals_expr = Symbolics.toexpr.(values) - ii = 1 + @RuntimeGeneratedFunction( + :(𝔛 -> sparse( + $rows, + $columns, + [$(vals_expr...)], + $nrows, + $ncolumns + ) + ) + ) +end - alll = [] - for var in future_varss - push!(alll,:($var = X[$ii])) - ii += 1 - end +function write_sparse_derivatives_function(rows::Vector{Int},columns::Vector{Int},values::Vector{Symbolics.Num},nrows::Int,ncolumns::Int,::Val{:string}) + vals_expr = Meta.parse(string(values)) - for var in present_varss - push!(alll,:($var = X[$ii])) - ii += 1 - end + vals_expr.args[1] = :Float64 - for var in past_varss - push!(alll,:($var = X[$ii])) - ii += 1 - end + @RuntimeGeneratedFunction( + :(𝔛 -> sparse( + $rows, + $columns, + $vals_expr, + $nrows, + $ncolumns + ) + ) + ) +end - for var in shock_varss - push!(alll,:($var = X[$ii])) - ii += 1 - end +function write_derivatives_function(values::Vector{Symbolics.Num}, ::Val{:string}) + vals_expr = Meta.parse(string(values)) + + @RuntimeGeneratedFunction(:(𝔛 -> $(Expr(:vect, vals_expr.args[2:end]...)))) +end + +function write_derivatives_function(values::Symbolics.Num, ::Val{:string}) + vals_expr = Meta.parse(string(values)) + + @RuntimeGeneratedFunction(:(𝔛 -> $vals_expr.args)) +end - # paras = [] - # push!(paras,:((;$(vcat(𝓂.parameters,𝓂.calibration_equations_parameters)...)) = params)) +function write_derivatives_function(values::Vector{Symbolics.Num}, position::UnitRange{Int}, ::Val{:string}) + vals_expr = Meta.parse(string(values)) - paras = [] - for (i, parss) in enumerate(vcat(𝓂.parameters,𝓂.calibration_equations_parameters)) - push!(paras,:($parss = params[$i])) - end + @RuntimeGeneratedFunction(:(𝔛 -> ($(Expr(:vect, vals_expr.args[2:end]...)), $position))) +end # needed for JET tests - # # watch out with naming of parameters in model and functions - # mod_func2 = :(function model_function_uni_redux(X::Vector, params::Vector{Number}, X̄::Vector) - # $(alll...) - # $(paras...) - # $(𝓂.calibration_equations_no_var...) - # $(steady_state...) - # [$(𝓂.dyn_equations...)] - # end) +function write_derivatives_function(values::Vector{Symbolics.Num}, position::Int, ::Val{:string}) + vals_expr = Meta.parse(string(values)) + @RuntimeGeneratedFunction(:(𝔛 -> ($(Expr(:vect, vals_expr.args[2:end]...)), $position))) +end + +function write_derivatives_function(values::Symbolics.Num, position::Int, ::Val{:string}) + vals_expr = Meta.parse(string(values)) + + @RuntimeGeneratedFunction(:(𝔛 -> ($vals_expr, $position))) +end + +function write_derivatives_function(values::Symbolics.Num, position::UnitRange{Int}, ::Val{:string}) + vals_expr = Meta.parse(string(values)) + position = position[1] + @RuntimeGeneratedFunction(:(𝔛 -> ($vals_expr, $position))) +end # needed for JET tests + +function write_derivatives_function(values::Vector{Symbolics.Num}, ::Val{:Symbolics}) + vals_expr = Symbolics.toexpr.(values) + + @RuntimeGeneratedFunction(:(𝔛 -> [$(vals_expr...)])) +end - # 𝓂.model_function = @RuntimeGeneratedFunction(mod_func2) - # 𝓂.model_function = eval(mod_func2) +function write_derivatives_function(values::Symbolics.Num, ::Val{:Symbolics}) + vals_expr = Symbolics.toexpr.(values) + + @RuntimeGeneratedFunction(:(𝔛 -> $vals_expr)) +end + +# TODO: check why this takes so much longer than previous implementation +function write_functions_mapping!(𝓂::ℳ, max_perturbation_order::Int; max_exprs_per_func::Int = 1) + future_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₁₎$"))) + present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₀₎$"))) + past_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₋₁₎$"))) + shock_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₓ₎$"))) + ss_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₛₛ₎$"))) + sort!(future_varss ,by = x->replace(string(x),r"₍₁₎$"=>"")) #sort by name without time index because otherwise eps_zᴸ⁽⁻¹⁾₍₋₁₎ comes before eps_z₍₋₁₎ + sort!(present_varss ,by = x->replace(string(x),r"₍₀₎$"=>"")) + sort!(past_varss ,by = x->replace(string(x),r"₍₋₁₎$"=>"")) + sort!(shock_varss ,by = x->replace(string(x),r"₍ₓ₎$"=>"")) + sort!(ss_varss ,by = x->replace(string(x),r"₍ₛₛ₎$"=>"")) + dyn_future_list = collect(reduce(union, 𝓂.dyn_future_list)) dyn_present_list = collect(reduce(union, 𝓂.dyn_present_list)) dyn_past_list = collect(reduce(union, 𝓂.dyn_past_list)) dyn_exo_list = collect(reduce(union,𝓂.dyn_exo_list)) + # dyn_ss_list = Symbol.(string.(collect(reduce(union,𝓂.dyn_ss_list))) .* "₍ₛₛ₎") future = map(x -> Symbol(replace(string(x), r"₍₁₎" => "")),string.(dyn_future_list)) present = map(x -> Symbol(replace(string(x), r"₍₀₎" => "")),string.(dyn_present_list)) past = map(x -> Symbol(replace(string(x), r"₍₋₁₎" => "")),string.(dyn_past_list)) exo = map(x -> Symbol(replace(string(x), r"₍ₓ₎" => "")),string.(dyn_exo_list)) + # stst = map(x -> Symbol(replace(string(x), r"₍ₛₛ₎" => "")),string.(dyn_ss_list)) vars_raw = [dyn_future_list[indexin(sort(future),future)]..., - dyn_present_list[indexin(sort(present),present)]..., - dyn_past_list[indexin(sort(past),past)]..., - dyn_exo_list[indexin(sort(exo),exo)]...] + dyn_present_list[indexin(sort(present),present)]..., + dyn_past_list[indexin(sort(past),past)]..., + dyn_exo_list[indexin(sort(exo),exo)]...] + + Symbolics.@syms norminvcdf(x) norminv(x) qnorm(x) normlogpdf(x) normpdf(x) normcdf(x) pnorm(x) dnorm(x) erfc(x) erfcinv(x) # overwrite SymPyCall names - eval(:(Symbolics.@variables $(reduce(union,get_symbols.(𝓂.dyn_equations))...))) + input_args = vcat(future_varss, + present_varss, + past_varss, + ss_varss, + 𝓂.parameters, + 𝓂.calibration_equations_parameters, + shock_varss) + + eval(:(Symbolics.@variables $(input_args...))) + + Symbolics.@variables 𝔛[1:length(input_args)] + + SS_and_pars_names_lead_lag = vcat(Symbol.(string.(sort(union(𝓂.var,𝓂.exo_past,𝓂.exo_future)))), 𝓂.calibration_equations_parameters) + + calib_eq_no_vars = reduce(union, get_symbols.(𝓂.calibration_equations_no_var), init = []) |> collect + + eval(:(Symbolics.@variables $((vcat(SS_and_pars_names_lead_lag, calib_eq_no_vars))...))) vars = eval(:(Symbolics.@variables $(vars_raw...))) eqs = Symbolics.parse_expr_to_symbolic.(𝓂.dyn_equations,(@__MODULE__,)) + final_indices = vcat(𝓂.parameters, SS_and_pars_names_lead_lag) + + input_X = Pair{Symbolics.Num, Symbolics.Num}[] + input_X_no_time = Pair{Symbolics.Num, Symbolics.Num}[] + + for (v,input) in enumerate(input_args) + push!(input_X, eval(input) => eval(𝔛[v])) + + if input ∈ shock_varss + push!(input_X_no_time, eval(𝔛[v]) => 0) + else + input_no_time = Symbol(replace(string(input), r"₍₁₎$"=>"", r"₍₀₎$"=>"" , r"₍₋₁₎$"=>"", r"₍ₛₛ₎$"=>"", r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) + + vv = indexin([input_no_time], final_indices) + + if vv[1] isa Int + push!(input_X_no_time, eval(𝔛[v]) => eval(𝔛[vv[1]])) + end + end + end + + vars_X = map(x -> Symbolics.substitute(x, input_X), vars) + + calib_eqs = Dict([(eval(calib_eq.args[1]) => eval(calib_eq.args[2])) for calib_eq in reverse(𝓂.calibration_equations_no_var)]) + + eqs_sub = Symbolics.Num[] + for subst in eqs + for _ in calib_eqs + for calib_eq in calib_eqs + subst = Symbolics.substitute(subst, calib_eq) + end + end + # subst = Symbolics.fixpoint_sub(subst, calib_eqs) + subst = Symbolics.substitute(subst, input_X) + push!(eqs_sub, subst) + end + if max_perturbation_order >= 2 nk = length(vars_raw) second_order_idxs = [nk * (i-1) + k for i in 1:nk for k in 1:i] @@ -4914,274 +5138,404 @@ function write_functions_mapping!(𝓂::ℳ, max_perturbation_order::Int) end end - first_order = [] - second_order = [] - third_order = [] + first_order = Symbolics.Num[] + second_order = Symbolics.Num[] + third_order = Symbolics.Num[] row1 = Int[] row2 = Int[] row3 = Int[] column1 = Int[] column2 = Int[] column3 = Int[] - # column3ext = Int[] - i1 = 1 - i2 = 1 - i3 = 1 - - for (c1,var1) in enumerate(vars) - for (r,eq) in enumerate(eqs) + + # Polyester.@batch for rc1 in 0:length(vars_X) * length(eqs_sub) - 1 + # for rc1 in 0:length(vars_X) * length(eqs_sub) - 1 + for (c1, var1) in enumerate(vars_X) + for (r, eq) in enumerate(eqs_sub) + # r, c1 = divrem(rc1, length(vars_X)) .+ 1 + # var1 = vars_X[c1] + # eq = eqs_sub[r] if Symbol(var1) ∈ Symbol.(Symbolics.get_variables(eq)) - deriv_first = Symbolics.derivative(eq,var1) - # if deriv_first != 0 - # deriv_expr = Meta.parse(string(deriv_first.subs(SPyPyC.PI,SPyPyC.N(SPyPyC.PI)))) - # push!(first_order, :($(postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, deriv_expr)))) - deriv_first_expr = Symbolics.toexpr(deriv_first) - # deriv_first_expr_safe = postwalk(x -> x isa Expr ? - # x.args[1] == :^ ? - # :(NaNMath.pow($(x.args[2:end]...))) : - # x : - # x, - # deriv_first_expr) - - push!(first_order, deriv_first_expr) - push!(row1,r) - push!(column1,c1) - i1 += 1 - if max_perturbation_order >= 2 - for (c2,var2) in enumerate(vars) - # if Symbol(var2) ∈ Symbol.(Symbolics.get_variables(deriv_first)) - if (((c1 - 1) * length(vars) + c2) ∈ second_order_idxs) && (Symbol(var2) ∈ Symbol.(Symbolics.get_variables(deriv_first))) - deriv_second = Symbolics.derivative(deriv_first,var2) - # if deriv_second != 0 - # deriv_expr = Meta.parse(string(deriv_second.subs(SPyPyC.PI,SPyPyC.N(SPyPyC.PI)))) - # push!(second_order, :($(postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, deriv_expr)))) - push!(second_order,Symbolics.toexpr(deriv_second)) - push!(row2,r) - # push!(column2,(c1 - 1) * length(vars) + c2) - push!(column2, Int.(indexin([(c1 - 1) * length(vars) + c2], second_order_idxs))...) - i2 += 1 - if max_perturbation_order == 3 - for (c3,var3) in enumerate(vars) - # if Symbol(var3) ∈ Symbol.(Symbolics.get_variables(deriv_second)) - # push!(column3ext,(c1 - 1) * length(vars)^2 + (c2 - 1) * length(vars) + c3) - if (((c1 - 1) * length(vars)^2 + (c2 - 1) * length(vars) + c3) ∈ third_order_idxs) && (Symbol(var3) ∈ Symbol.(Symbolics.get_variables(deriv_second))) - deriv_third = Symbolics.derivative(deriv_second,var3) - # if deriv_third != 0 - # deriv_expr = Meta.parse(string(deriv_third.subs(SPyPyC.PI,SPyPyC.N(SPyPyC.PI)))) - # push!(third_order, :($(postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, deriv_expr)))) - push!(third_order,Symbolics.toexpr(deriv_third)) - push!(row3,r) - # push!(column3,(c1 - 1) * length(vars)^2 + (c2 - 1) * length(vars) + c3) - push!(column3, Int.(indexin([(c1 - 1) * length(vars)^2 + (c2 - 1) * length(vars) + c3], third_order_idxs))...) - i3 += 1 - # end - end - # end - end + deriv_first = Symbolics.derivative(eq, var1) + + push!(first_order, deriv_first) + push!(row1, r) + # push!(row1, r...) + push!(column1, c1) + if max_perturbation_order >= 2 + for (c2, var2) in enumerate(vars_X) + if (((c1 - 1) * length(vars) + c2) ∈ second_order_idxs) && (Symbol(var2) ∈ Symbol.(Symbolics.get_variables(deriv_first))) + deriv_second = Symbolics.derivative(deriv_first, var2) + + push!(second_order, deriv_second) + push!(row2, r) + push!(column2, Int.(indexin([(c1 - 1) * length(vars) + c2], second_order_idxs))...) + if max_perturbation_order == 3 + for (c3, var3) in enumerate(vars_X) + if (((c1 - 1) * length(vars)^2 + (c2 - 1) * length(vars) + c3) ∈ third_order_idxs) && (Symbol(var3) ∈ Symbol.(Symbolics.get_variables(deriv_second))) + deriv_third = Symbolics.derivative(deriv_second,var3) + + push!(third_order, deriv_third) + push!(row3, r) + push!(column3, Int.(indexin([(c1 - 1) * length(vars)^2 + (c2 - 1) * length(vars) + c3], third_order_idxs))...) end - # end + end end end end - # end + end end end end + + if max_perturbation_order >= 1 + if 𝓂.model_jacobian[2] == Int[] + write_auxilliary_indices!(𝓂) + write_derivatives_of_ss_equations!(𝓂::ℳ, max_exprs_per_func = max_exprs_per_func) - mod_func3 = :(function model_jacobian(X::Vector, params::Vector{Real}, X̄::Vector) - $(alll...) - $(paras...) - $(𝓂.calibration_equations_no_var...) - $(steady_state...) - sparse([$(row1...)], [$(column1...)], [$(first_order...)], $(length(eqs)), $(length(vars))) - end) + # derivative of jacobian wrt SS_and_pars and parameters + eqs_static = map(x -> Symbolics.substitute(x, input_X_no_time), first_order) - 𝓂.model_jacobian = @RuntimeGeneratedFunction(mod_func3) - # 𝓂.model_jacobian = FWrap{Tuple{Vector{Float64}, Vector{Number}, Vector{Float64}}, SparseMatrixCSC{Float64}}(@RuntimeGeneratedFunction(mod_func3)) + ∂jacobian_∂SS_and_pars = Symbolics.sparsejacobian(eqs_static, eval.(𝔛[1:(length(final_indices))]), simplify = false) # |> findnz - # 𝓂.model_jacobian = eval(mod_func3) + idx_conversion = (row1 + length(eqs) * (column1 .- 1)) + cols, rows, vals = findnz(∂jacobian_∂SS_and_pars) #transposed - if max_perturbation_order >= 2 && 𝓂.model_hessian == Function[] - # if length(row2) == 0 - # out = :(spzeros($(length(eqs)), $(length(second_order_idxs)))) - # else - # out = :(sparse([$(row2...)], [$(column2...)], [$(second_order...)], $(length(eqs)), $(length(second_order_idxs)))) - # end + converted_cols = idx_conversion[cols] - # mod_func4 = :(function model_hessian(X::Vector, params::Vector{Real}, X̄::Vector) - # $(alll...) - # $(paras...) - # $(𝓂.calibration_equations_no_var...) - # $(steady_state...) - # $out - # end) - - for (l,second) in enumerate(second_order) - exx = :(function(X::Vector, params::Vector{Real}, X̄::Vector) - $(alll...) - $(paras...) - $(𝓂.calibration_equations_no_var...) - $(steady_state...) - return $second, $(row2[l]), $(column2[l]) - end) - push!(𝓂.model_hessian,@RuntimeGeneratedFunction(exx)) - end + perm_vals = sortperm(converted_cols) # sparse reorders the rows and cols and sorts by column. need to do that also for the values - 𝓂.solution.perturbation.second_order_auxilliary_matrices = create_second_order_auxilliary_matrices(𝓂.timings) + min_n_funcs = length(vals) ÷ max_exprs_per_func + 1 + funcs = Function[] - # 𝓂.model_hessian = @RuntimeGeneratedFunction(mod_func4) - # 𝓂.model_hessian = eval(mod_func4) - end + lk = ReentrantLock() - if max_perturbation_order == 3 && 𝓂.model_third_order_derivatives == Function[] - # if length(row3) == 0 - # out = :(spzeros($(length(eqs)), $(length(third_order_idxs)))) - # else - # out = :(sparse([$(row3...)], [$(column3...)], [$(third_order...)], $(length(eqs)), $(length(third_order_idxs)))) - # end + if min_n_funcs == 1 + push!(funcs, write_derivatives_function(vals[perm_vals], 1:length(vals), Val(:string))) + else + Polyester.@batch minbatch = 20 for i in 1:min(min_n_funcs, length(vals)) + indices = ((i - 1) * max_exprs_per_func + 1):(i == min_n_funcs ? length(vals) : i * max_exprs_per_func) - # mod_func5 = :(function model_hessian(X::Vector, params::Vector{Real}, X̄::Vector) - # $(alll...) - # $(paras...) - # $(𝓂.calibration_equations_no_var...) - # $(steady_state...) - # $out - # end) - - for (l,third) in enumerate(third_order) - exx = :(function(X::Vector, params::Vector{Real}, X̄::Vector) - $(alll...) - $(paras...) - $(𝓂.calibration_equations_no_var...) - $(steady_state...) - return $third, $(row3[l]), $(column3[l]) - end) - push!(𝓂.model_third_order_derivatives,@RuntimeGeneratedFunction(exx)) - end + indices = length(indices) == 1 ? indices[1] : indices - 𝓂.solution.perturbation.third_order_auxilliary_matrices = create_third_order_auxilliary_matrices(𝓂.timings, unique(column3)) + func = write_derivatives_function(vals[perm_vals][indices], indices, Val(:string)) - end + begin + lock(lk) + try + push!(funcs, func) + finally + unlock(lk) + end + end + end + end + 𝓂.model_jacobian_SS_and_pars_vars = (funcs, sparse(rows, converted_cols, zero(cols), length(final_indices), length(eqs) * length(vars))) - # write indices in auxiliary objects - dyn_var_future_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₁₎"))) - dyn_var_present_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₀₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₀₎"))) - dyn_var_past_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₋₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₋₁₎"))) - dyn_exo_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₓ₎"))) - dyn_ss_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₛₛ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₛₛ₎"))) + # first order + min_n_funcs = length(first_order) ÷ max_exprs_per_func + 1 - dyn_var_future = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_var_future_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) - dyn_var_present = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_var_present_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) - dyn_var_past = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_var_past_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) - dyn_exo = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_exo_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) - dyn_ss = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_ss_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) + funcs = Function[] - SS_and_pars_names = vcat(Symbol.(replace.(string.(sort(union(𝓂.var,𝓂.exo_past,𝓂.exo_future))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")), 𝓂.calibration_equations_parameters) + lk = ReentrantLock() + if min_n_funcs == 1 + push!(funcs, write_derivatives_function(first_order, 1:length(first_order), Val(:string))) + else + Polyester.@batch minbatch = 20 for i in 1:min(min_n_funcs, length(first_order)) + indices = ((i - 1) * max_exprs_per_func + 1):(i == min_n_funcs ? length(first_order) : i * max_exprs_per_func) - dyn_var_future_idx = indexin(dyn_var_future , SS_and_pars_names) - dyn_var_present_idx = indexin(dyn_var_present , SS_and_pars_names) - dyn_var_past_idx = indexin(dyn_var_past , SS_and_pars_names) - dyn_ss_idx = indexin(dyn_ss , SS_and_pars_names) + indices = length(indices) == 1 ? indices[1] : indices - shocks_ss = zeros(length(dyn_exo)) + func = write_derivatives_function(first_order[indices], indices, Val(:string)) - 𝓂.solution.perturbation.auxilliary_indices = auxilliary_indices(dyn_var_future_idx, dyn_var_present_idx, dyn_var_past_idx, dyn_ss_idx, shocks_ss) + begin + lock(lk) + try + push!(funcs, func) + finally + unlock(lk) + end + end + end + end + 𝓂.model_jacobian = (funcs, row1 .+ (column1 .- 1) .* length(eqs_sub), zeros(length(eqs_sub), length(vars))) + end + end + + if max_perturbation_order >= 2 + # second order + if 𝓂.solution.perturbation.second_order_auxilliary_matrices.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) + 𝓂.solution.perturbation.second_order_auxilliary_matrices = create_second_order_auxilliary_matrices(𝓂.timings) - # 𝓂.model_third_order_derivatives = @RuntimeGeneratedFunction(mod_func5) - # 𝓂.model_third_order_derivatives = eval(mod_func5) + perm_vals = sortperm(column2) # sparse reorders the rows and cols and sorts by column. need to do that also for the values + min_n_funcs = length(second_order) ÷ max_exprs_per_func + 1 - # calib_eqs = [] - # for (i, eqs) in enumerate(𝓂.solved_vals) - # varss = 𝓂.solved_vars[i] - # push!(calib_eqs,:($varss = $eqs)) - # end + funcs = Function[] + + lk = ReentrantLock() - # for varss in 𝓂.exo - # push!(calib_eqs,:($varss = 0)) - # end + if min_n_funcs == 1 + push!(funcs, write_derivatives_function(second_order[perm_vals], 1:length(second_order), Val(:string))) + else + Polyester.@batch minbatch = 20 for i in 1:min(min_n_funcs, length(second_order)) + indices = ((i - 1) * max_exprs_per_func + 1):(i == min_n_funcs ? length(second_order) : i * max_exprs_per_func) + + indices = length(indices) == 1 ? indices[1] : indices - # calib_pars = [] - # for (i, parss) in enumerate(𝓂.parameters) - # push!(calib_pars,:($parss = parameters[$i])) - # end + func = write_derivatives_function(second_order[perm_vals][indices], indices, Val(:string)) - # var_out = [] - # ii = 1 - # for var in 𝓂.var - # push!(var_out,:($var = SS[$ii])) - # ii += 1 - # end + begin + lock(lk) + try + push!(funcs, func) + finally + unlock(lk) + end + end + end + end - # par_out = [] - # for cal in 𝓂.calibration_equations_parameters - # push!(par_out,:($cal = SS[$ii])) - # ii += 1 - # end + 𝓂.model_hessian = (funcs, sparse(row2, column2, zero(column2), length(eqs_sub), size(𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂,1))) + end + end - # calib_pars = [] - # for (i, parss) in enumerate(𝓂.parameters) - # push!(calib_pars,:($parss = parameters[$i])) - # end + if max_perturbation_order == 3 + # third order + if 𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) + 𝓂.solution.perturbation.third_order_auxilliary_matrices = create_third_order_auxilliary_matrices(𝓂.timings, unique(column3)) + + perm_vals = sortperm(column3) # sparse reorders the rows and cols and sorts by column. need to do that also for the values + + min_n_funcs = length(third_order) ÷ max_exprs_per_func + 1 - # test_func = :(function test_SS(parameters::Vector{Float64}, SS::Vector{Float64}) - # $(calib_pars...) - # $(var_out...) - # $(par_out...) - # [$(𝓂.ss_equations...),$(𝓂.calibration_equations...)] - # end) + funcs = Function[] + + lk = ReentrantLock() + + if min_n_funcs == 1 + push!(funcs, write_derivatives_function(third_order[perm_vals], 1:length(third_order), Val(:string))) + else + Polyester.@batch minbatch = 20 for i in 1:min(min_n_funcs, length(third_order)) + indices = ((i - 1) * max_exprs_per_func + 1):(i == min_n_funcs ? length(third_order) : i * max_exprs_per_func) + + if length(indices) == 1 + indices = indices[1] + end + + func = write_derivatives_function(third_order[perm_vals][indices], indices, Val(:string)) + + begin + lock(lk) + try + push!(funcs, func) + finally + unlock(lk) + end + end + end + end - # 𝓂.solution.valid_steady_state_solution = @RuntimeGeneratedFunction(test_func) + 𝓂.model_third_order_derivatives = (funcs, sparse(row3, column3, zero(column3), length(eqs_sub), size(𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃,1))) + end + end - # 𝓂.solution.outdated_algorithms = Set([:linear_time_iteration, :riccati, :quadratic_iteration, :first_order, :second_order, :third_order]) return nothing end +function write_derivatives_of_ss_equations!(𝓂::ℳ; max_exprs_per_func::Int = 1) + # derivative of SS equations wrt parameters and SS_and_pars + # unknowns = union(setdiff(𝓂.vars_in_ss_equations, 𝓂.➕_vars), 𝓂.calibration_equations_parameters) + SS_and_pars = Symbol.(vcat(string.(sort(collect(setdiff(reduce(union,get_symbols.(𝓂.ss_aux_equations)),union(𝓂.parameters_in_equations,𝓂.➕_vars))))), 𝓂.calibration_equations_parameters)) -write_parameters_input!(𝓂::ℳ, parameters::Nothing; verbose::Bool = true) = return parameters -write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Pair{String,Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters[1] |> Meta.parse |> replace_indices => parameters[2]), verbose = verbose) + ss_equations = vcat(𝓂.ss_equations, 𝓂.calibration_equations) + Symbolics.@syms norminvcdf(x) norminv(x) qnorm(x) normlogpdf(x) normpdf(x) normcdf(x) pnorm(x) dnorm(x) erfc(x) erfcinv(x) + # overwrite SymPyCall names + other_pars = setdiff(union(𝓂.parameters_in_equations, 𝓂.parameters_as_function_of_parameters), 𝓂.parameters) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Float64},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Union{Symbol,String},Union{Float64,Int}},Vararg{Pair{Union{Symbol,String},Union{Float64,Int}}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{String,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{String,Float64},Vararg{Pair{String,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]) -, verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{String, Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) + if length(other_pars) > 0 + eval(:(Symbolics.@variables $(other_pars...))) + end + vars = eval(:(Symbolics.@variables $(SS_and_pars...))) -write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Pair{String,Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{Symbol,Int}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{String,Int},Vararg{Pair{String,Int}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{String, Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) + pars = eval(:(Symbolics.@variables $(𝓂.parameters...))) + input_args = vcat(𝓂.parameters, SS_and_pars) + + Symbolics.@variables 𝔛[1:length(input_args)] -write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Pair{String,Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Real},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{String,Real},Vararg{Pair{String,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{String, Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) + input_X_no_time = Pair{Symbolics.Num, Symbolics.Num}[] + + for (v,input) in enumerate(input_args) + push!(input_X_no_time, eval(input) => eval(𝔛[v])) + end + ss_eqs = Symbolics.parse_expr_to_symbolic.(ss_equations,(@__MODULE__,)) + calib_eqs = Dict([(eval(calib_eq.args[1]) => eval(calib_eq.args[2])) for calib_eq in reverse(𝓂.calibration_equations_no_var)]) -function write_parameters_input!(𝓂::ℳ, parameters::Dict{Symbol,Float64}; verbose::Bool = true) - if length(setdiff(collect(keys(parameters)),𝓂.parameters))>0 - println("Parameters not part of the model: ",setdiff(collect(keys(parameters)),𝓂.parameters)) - for kk in setdiff(collect(keys(parameters)),𝓂.parameters) - delete!(parameters,kk) + eqs = Symbolics.Num[] + for subst in ss_eqs + for _ in calib_eqs # to completely substitute all calibration equations + for calib_eq in calib_eqs + subst = Symbolics.substitute(subst, calib_eq) + end + end + # subst = Symbolics.fixpoint_sub(subst, calib_eqs) + subst = Symbolics.substitute(subst, input_X_no_time) + push!(eqs, subst) + end + + ∂SS_equations_∂parameters = Symbolics.sparsejacobian(eqs, eval.(𝔛[1:length(pars)])) |> findnz + + min_n_funcs = length(∂SS_equations_∂parameters[3]) ÷ max_exprs_per_func + 1 + + funcs = Function[] + + lk = ReentrantLock() + + if min_n_funcs == 1 + push!(funcs, write_derivatives_function(∂SS_equations_∂parameters[3], 1:length(∂SS_equations_∂parameters[3]), Val(:string))) + else + Polyester.@batch minbatch = 20 for i in 1:min(min_n_funcs, length(∂SS_equations_∂parameters[3])) + indices = ((i - 1) * max_exprs_per_func + 1):(i == min_n_funcs ? length(∂SS_equations_∂parameters[3]) : i * max_exprs_per_func) + + indices = length(indices) == 1 ? indices[1] : indices + + func = write_derivatives_function(∂SS_equations_∂parameters[3][indices], indices, Val(:string)) + + begin + lock(lk) + try + push!(funcs, func) + finally + unlock(lk) + end + end + end + end + + 𝓂.∂SS_equations_∂parameters = (funcs, sparse(∂SS_equations_∂parameters[1], ∂SS_equations_∂parameters[2], zeros(Float64,length(∂SS_equations_∂parameters[3])), length(eqs), length(pars))) + + # 𝓂.∂SS_equations_∂parameters = write_sparse_derivatives_function(∂SS_equations_∂parameters[1], + # ∂SS_equations_∂parameters[2], + # ∂SS_equations_∂parameters[3], + # length(eqs), + # length(pars), + # Val(:string)); + + ∂SS_equations_∂SS_and_pars = Symbolics.sparsejacobian(eqs, eval.(𝔛[length(pars)+1:end])) |> findnz + + min_n_funcs = length(∂SS_equations_∂SS_and_pars[3]) ÷ max_exprs_per_func + 1 + + funcs = Function[] + + lk = ReentrantLock() + + if min_n_funcs == 1 + push!(funcs, write_derivatives_function(∂SS_equations_∂SS_and_pars[3], 1:length(∂SS_equations_∂SS_and_pars[3]), Val(:string))) + else + Polyester.@batch minbatch = 20 for i in 1:min(min_n_funcs, length(∂SS_equations_∂SS_and_pars[3])) + indices = ((i - 1) * max_exprs_per_func + 1):(i == min_n_funcs ? length(∂SS_equations_∂SS_and_pars[3]) : i * max_exprs_per_func) + + indices = length(indices) == 1 ? indices[1] : indices + + func = write_derivatives_function(∂SS_equations_∂SS_and_pars[3][indices], indices, Val(:string)) + + begin + lock(lk) + try + push!(funcs, func) + finally + unlock(lk) + end + end + end + end + + 𝓂.∂SS_equations_∂SS_and_pars = (funcs, ∂SS_equations_∂SS_and_pars[1] .+ (∂SS_equations_∂SS_and_pars[2] .- 1) .* length(eqs), zeros(length(eqs), length(vars))) + + # 𝓂.∂SS_equations_∂SS_and_pars = write_sparse_derivatives_function(∂SS_equations_∂SS_and_pars[1], + # ∂SS_equations_∂SS_and_pars[2], + # ∂SS_equations_∂SS_and_pars[3], + # length(eqs), + # length(vars), + # Val(:string)); +end + +function write_auxilliary_indices!(𝓂::ℳ) + # write indices in auxiliary objects + dyn_var_future_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₁₎"))) + dyn_var_present_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₀₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₀₎"))) + dyn_var_past_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₋₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₋₁₎"))) + dyn_exo_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₓ₎"))) + dyn_ss_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₛₛ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍ₛₛ₎"))) + + dyn_var_future = Symbol.(string.(sort(collect(reduce(union,dyn_var_future_list))))) + dyn_var_present = Symbol.(string.(sort(collect(reduce(union,dyn_var_present_list))))) + dyn_var_past = Symbol.(string.(sort(collect(reduce(union,dyn_var_past_list))))) + dyn_exo = Symbol.(string.(sort(collect(reduce(union,dyn_exo_list))))) + dyn_ss = Symbol.(string.(sort(collect(reduce(union,dyn_ss_list))))) + + SS_and_pars_names = vcat(Symbol.(string.(sort(union(𝓂.var,𝓂.exo_past,𝓂.exo_future)))), 𝓂.calibration_equations_parameters) + + dyn_var_future_idx = indexin(dyn_var_future , SS_and_pars_names) + dyn_var_present_idx = indexin(dyn_var_present , SS_and_pars_names) + dyn_var_past_idx = indexin(dyn_var_past , SS_and_pars_names) + dyn_ss_idx = indexin(dyn_ss , SS_and_pars_names) + + shocks_ss = zeros(length(dyn_exo)) + + 𝓂.solution.perturbation.auxilliary_indices = auxilliary_indices(dyn_var_future_idx, dyn_var_present_idx, dyn_var_past_idx, dyn_ss_idx, shocks_ss) +end + +write_parameters_input!(𝓂::ℳ, parameters::Nothing; verbose::Bool = true) = return parameters +write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Pair{String,Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters[1] |> Meta.parse |> replace_indices => parameters[2]), verbose = verbose) + + + +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Float64},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Union{Symbol,String},Union{Float64,Int}},Vararg{Pair{Union{Symbol,String},Union{Float64,Int}}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{String,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{String,Float64},Vararg{Pair{String,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]) +, verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{String, Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) + + +write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Pair{String,Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{Symbol,Int}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{String,Int},Vararg{Pair{String,Int}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{String, Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) + + +write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Pair{String,Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Real},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{String,Real},Vararg{Pair{String,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{String, Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) + + + +function write_parameters_input!(𝓂::ℳ, parameters::Dict{Symbol,Float64}; verbose::Bool = true) + if length(setdiff(collect(keys(parameters)),𝓂.parameters))>0 + println("Parameters not part of the model: ",setdiff(collect(keys(parameters)),𝓂.parameters)) + for kk in setdiff(collect(keys(parameters)),𝓂.parameters) + delete!(parameters,kk) end end @@ -5455,11 +5809,157 @@ function mean_parameter_derivatives(parameters::ℱ.Dual{Z,S,N}, parameters_idx: end +function create_timings_for_estimation!(𝓂::ℳ, observables::Vector{Symbol}) + dyn_equations = 𝓂.dyn_equations + + vars_to_exclude = setdiff(𝓂.timings.present_only, observables) + + # Mapping variables to their equation index + variable_to_equation = Dict{Symbol, Vector{Int}}() + for var in vars_to_exclude + for (eq_idx, vars_set) in enumerate(𝓂.dyn_var_present_list) + # for var in vars_set + if var in vars_set + if haskey(variable_to_equation, var) + push!(variable_to_equation[var],eq_idx) + else + variable_to_equation[var] = [eq_idx] + end + end + end + end + + # cols_to_exclude = indexin(𝓂.timings.var, setdiff(𝓂.timings.present_only, observables)) + cols_to_exclude = indexin(setdiff(𝓂.timings.present_only, observables), 𝓂.timings.var) + + present_idx = 𝓂.timings.nFuture_not_past_and_mixed .+ (setdiff(range(1, 𝓂.timings.nVars), cols_to_exclude)) + + dyn_var_future_list = deepcopy(𝓂.dyn_var_future_list) + dyn_var_present_list = deepcopy(𝓂.dyn_var_present_list) + dyn_var_past_list = deepcopy(𝓂.dyn_var_past_list) + dyn_exo_list = deepcopy(𝓂.dyn_exo_list) + dyn_ss_list = deepcopy(𝓂.dyn_ss_list) + + rows_to_exclude = Int[] + + for vidx in values(variable_to_equation) + for v in vidx + if v ∉ rows_to_exclude + push!(rows_to_exclude, v) + + for vv in vidx + dyn_var_future_list[vv] = union(dyn_var_future_list[vv], dyn_var_future_list[v]) + dyn_var_present_list[vv] = union(dyn_var_present_list[vv], dyn_var_present_list[v]) + dyn_var_past_list[vv] = union(dyn_var_past_list[vv], dyn_var_past_list[v]) + dyn_exo_list[vv] = union(dyn_exo_list[vv], dyn_exo_list[v]) + dyn_ss_list[vv] = union(dyn_ss_list[vv], dyn_ss_list[v]) + end + + break + end + end + end -function calculate_jacobian(parameters::Vector{M}, SS_and_pars::AbstractArray{N}, 𝓂::ℳ) where {M,N} + rows_to_include = setdiff(1:𝓂.timings.nVars, rows_to_exclude) + + all_symbols = setdiff(reduce(union,collect.(get_symbols.(dyn_equations)))[rows_to_include], vars_to_exclude) + parameters_in_equations = sort(setdiff(all_symbols, match_pattern(all_symbols,r"₎$"))) + + dyn_var_future = sort(setdiff(collect(reduce(union,dyn_var_future_list[rows_to_include])), vars_to_exclude)) + dyn_var_present = sort(setdiff(collect(reduce(union,dyn_var_present_list[rows_to_include])), vars_to_exclude)) + dyn_var_past = sort(setdiff(collect(reduce(union,dyn_var_past_list[rows_to_include])), vars_to_exclude)) + dyn_var_ss = sort(setdiff(collect(reduce(union,dyn_ss_list[rows_to_include])), vars_to_exclude)) + + all_dyn_vars = union(dyn_var_future, dyn_var_present, dyn_var_past) + + @assert length(setdiff(dyn_var_ss, all_dyn_vars)) == 0 "The following variables are (and cannot be) defined only in steady state (`[ss]`): $(setdiff(dyn_var_ss, all_dyn_vars))" + + all_vars = union(all_dyn_vars, dyn_var_ss) + + present_only = sort(setdiff(dyn_var_present,union(dyn_var_past,dyn_var_future))) + future_not_past = sort(setdiff(dyn_var_future, dyn_var_past)) + past_not_future = sort(setdiff(dyn_var_past, dyn_var_future)) + mixed = sort(setdiff(dyn_var_present, union(present_only, future_not_past, past_not_future))) + future_not_past_and_mixed = sort(union(future_not_past,mixed)) + past_not_future_and_mixed = sort(union(past_not_future,mixed)) + present_but_not_only = sort(setdiff(dyn_var_present,present_only)) + mixed_in_past = sort(intersect(dyn_var_past, mixed)) + not_mixed_in_past = sort(setdiff(dyn_var_past,mixed_in_past)) + mixed_in_future = sort(intersect(dyn_var_future, mixed)) + exo = sort(collect(reduce(union,dyn_exo_list))) + var = sort(dyn_var_present) + aux_tmp = sort(filter(x->occursin(r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾",string(x)), dyn_var_present)) + aux = aux_tmp[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∉ exo, aux_tmp)] + exo_future = dyn_var_future[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_future)] + exo_present = dyn_var_present[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_present)] + exo_past = dyn_var_past[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_past)] + + nPresent_only = length(present_only) + nMixed = length(mixed) + nFuture_not_past_and_mixed = length(future_not_past_and_mixed) + nPast_not_future_and_mixed = length(past_not_future_and_mixed) + nPresent_but_not_only = length(present_but_not_only) + nVars = length(all_vars) + nExo = length(collect(exo)) + + present_only_idx = indexin(present_only,var) + present_but_not_only_idx = indexin(present_but_not_only,var) + future_not_past_and_mixed_idx = indexin(future_not_past_and_mixed,var) + past_not_future_and_mixed_idx = indexin(past_not_future_and_mixed,var) + mixed_in_future_idx = indexin(mixed_in_future,dyn_var_future) + mixed_in_past_idx = indexin(mixed_in_past,dyn_var_past) + not_mixed_in_past_idx = indexin(not_mixed_in_past,dyn_var_past) + past_not_future_idx = indexin(past_not_future,var) + + reorder = indexin(var, [present_only; past_not_future; future_not_past_and_mixed]) + dynamic_order = indexin(present_but_not_only, [past_not_future; future_not_past_and_mixed]) + + @assert length(intersect(union(var,exo),parameters_in_equations)) == 0 "Parameters and variables cannot have the same name. This is the case for: " * repr(sort([intersect(union(var,exo),parameters_in_equations)...])) + + T = timings(present_only, + future_not_past, + past_not_future, + mixed, + future_not_past_and_mixed, + past_not_future_and_mixed, + present_but_not_only, + mixed_in_past, + not_mixed_in_past, + mixed_in_future, + exo, + var, + aux, + exo_present, + + nPresent_only, + nMixed, + nFuture_not_past_and_mixed, + nPast_not_future_and_mixed, + nPresent_but_not_only, + nVars, + nExo, + + present_only_idx, + present_but_not_only_idx, + future_not_past_and_mixed_idx, + not_mixed_in_past_idx, + past_not_future_and_mixed_idx, + mixed_in_past_idx, + mixed_in_future_idx, + past_not_future_idx, + + reorder, + dynamic_order) + + push!(𝓂.estimation_helper, observables => T) +end + + + +function calculate_jacobian(parameters::Vector{M}, SS_and_pars::Vector{N}, 𝓂::ℳ)::Matrix{M} where {M,N} SS = SS_and_pars[1:end - length(𝓂.calibration_equations)] calibrated_parameters = SS_and_pars[(end - length(𝓂.calibration_equations)+1):end] - # par = ComponentVector(vcat(parameters,calibrated_parameters),Axis(vcat(𝓂.parameters,𝓂.calibration_equations_parameters))) + par = vcat(parameters,calibrated_parameters) dyn_var_future_idx = 𝓂.solution.perturbation.auxilliary_indices.dyn_var_future_idx @@ -5469,12 +5969,89 @@ function calculate_jacobian(parameters::Vector{M}, SS_and_pars::AbstractArray{N} shocks_ss = 𝓂.solution.perturbation.auxilliary_indices.shocks_ss - # return 𝒜.jacobian(𝒷(), x -> 𝓂.model_function(x, par, SS), [SS_future; SS_present; SS_past; shocks_ss])#, SS_and_pars - # return Matrix(𝓂.model_jacobian(([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]))) - return 𝓂.model_jacobian([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]) + X = [SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; SS[dyn_ss_idx]; par; shocks_ss] + + # vals = M[] + + # for f in 𝓂.model_jacobian[1] + # push!(vals, f(X)...) + # end + + vals = zeros(M, length(𝓂.model_jacobian[1])) + + # lk = ReentrantLock() + + Polyester.@batch minbatch = 200 for f in 𝓂.model_jacobian[1] + # for f in 𝓂.model_jacobian[1] + # val, idx = f(X)#::Tuple{<: Real, Int} + out = f(X)#::Tuple{Vector{<: Real}, UnitRange{Int64}} + + # begin + # lock(lk) + # try + @inbounds vals[out[2]] = out[1] + # @inbounds vals[idx] = val + # finally + # unlock(lk) + # end + # end + end + + if eltype(𝓂.model_jacobian[3]) ≠ M + Accessors.@reset 𝓂.model_jacobian[3] = convert(Matrix{M}, 𝓂.model_jacobian[3]) + end + + 𝓂.model_jacobian[3][𝓂.model_jacobian[2]] .= vals + + return 𝓂.model_jacobian[3] end +function rrule(::typeof(calculate_jacobian), parameters, SS_and_pars, 𝓂) + jacobian = calculate_jacobian(parameters, SS_and_pars, 𝓂) + + function calculate_jacobian_pullback(∂∇₁) + X = [parameters; SS_and_pars] + + # vals = Float64[] + + # for f in 𝓂.model_jacobian_SS_and_pars_vars[1] + # push!(vals, f(X)...) + # end + + vals = zeros(Float64, length(𝓂.model_jacobian_SS_and_pars_vars[1])) + + # lk = ReentrantLock() + + Polyester.@batch minbatch = 200 for f in 𝓂.model_jacobian_SS_and_pars_vars[1] + out = f(X) + + # begin + # lock(lk) + # try + @inbounds vals[out[2]] = out[1] + # finally + # unlock(lk) + # end + # end + end + + Accessors.@reset 𝓂.model_jacobian_SS_and_pars_vars[2].nzval = vals + + analytical_jac_SS_and_pars_vars = 𝓂.model_jacobian_SS_and_pars_vars[2] |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + + cols_unique = unique(findnz(analytical_jac_SS_and_pars_vars)[2]) + + v∂∇₁ = ∂∇₁[cols_unique] + + ∂parameters_and_SS_and_pars = analytical_jac_SS_and_pars_vars[:,cols_unique] * v∂∇₁ + + return NoTangent(), ∂parameters_and_SS_and_pars[1:length(parameters)], ∂parameters_and_SS_and_pars[length(parameters)+1:end], NoTangent() + end + + return jacobian, calculate_jacobian_pullback +end + function calculate_hessian(parameters::Vector{M}, SS_and_pars::Vector{N}, 𝓂::ℳ) where {M,N} SS = SS_and_pars[1:end - length(𝓂.calibration_equations)] @@ -5492,25 +6069,65 @@ function calculate_hessian(parameters::Vector{M}, SS_and_pars::Vector{N}, 𝓂:: # nk = 𝓂.timings.nPast_not_future_and_mixed + 𝓂.timings.nVars + 𝓂.timings.nFuture_not_past_and_mixed + length(𝓂.exo) # return sparse(reshape(𝒜.jacobian(𝒷(), x -> 𝒜.jacobian(𝒷(), x -> (𝓂.model_function(x, par, SS)), x), [SS_future; SS_present; SS_past; shocks_ss] ), 𝓂.timings.nVars, nk^2))#, SS_and_pars - # return 𝓂.model_hessian([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]) + # return 𝓂.model_hessian([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss; par; SS[dyn_ss_idx]]) * 𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂ + + # second_out = [f([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss; par; SS[dyn_ss_idx]]) for f in 𝓂.model_hessian] + + # vals = [i[1] for i in second_out] + # rows = [i[2] for i in second_out] + # cols = [i[3] for i in second_out] - second_out = [f([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]) for f in 𝓂.model_hessian] + X = [SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; SS[dyn_ss_idx]; par; shocks_ss] - vals = [i[1] for i in second_out] - rows = [i[2] for i in second_out] - cols = [i[3] for i in second_out] + # vals = M[] + + # for f in 𝓂.model_hessian[1] + # push!(vals, f(X)...) + # end - vals = convert(Vector{M}, vals) + vals = zeros(M, length(𝓂.model_hessian[1])) - # nk = 𝓂.timings.nPast_not_future_and_mixed + 𝓂.timings.nVars + 𝓂.timings.nFuture_not_past_and_mixed + length(𝓂.exo) - # sparse(rows, cols, vals, length(𝓂.dyn_equations), nk^2) - sparse(rows, cols, vals, length(𝓂.dyn_equations), size(𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂,1)) * 𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂ + # lk = ReentrantLock() + + Polyester.@batch minbatch = 200 for f in 𝓂.model_hessian[1] + out = f(X) + + # begin + # lock(lk) + # try + @inbounds vals[out[2]] = out[1] + # finally + # unlock(lk) + # end + # end + end + + Accessors.@reset 𝓂.model_hessian[2].nzval = vals + + return 𝓂.model_hessian[2] * 𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂ + + # vals = M[] + # rows = Int[] + # cols = Int[] + + # for f in 𝓂.model_hessian + # output = f(input) + + # push!(vals, output[1]...) + # push!(rows, output[2]...) + # push!(cols, output[3]...) + # end + + # vals = convert(Vector{M}, vals) + + # # nk = 𝓂.timings.nPast_not_future_and_mixed + 𝓂.timings.nVars + 𝓂.timings.nFuture_not_past_and_mixed + length(𝓂.exo) + # # sparse(rows, cols, vals, length(𝓂.dyn_equations), nk^2) + # sparse!(rows, cols, vals, length(𝓂.dyn_equations), size(𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂,1)) * 𝓂.solution.perturbation.second_order_auxilliary_matrices.𝐔∇₂ end function calculate_third_order_derivatives(parameters::Vector{M}, SS_and_pars::Vector{N}, 𝓂::ℳ) where {M,N} - SS = SS_and_pars[1:end - length(𝓂.calibration_equations)] calibrated_parameters = SS_and_pars[(end - length(𝓂.calibration_equations)+1):end] @@ -5524,20 +6141,61 @@ function calculate_third_order_derivatives(parameters::Vector{M}, SS_and_pars::V shocks_ss = 𝓂.solution.perturbation.auxilliary_indices.shocks_ss # return sparse(reshape(𝒜.jacobian(𝒷(), x -> 𝒜.jacobian(𝒷(), x -> 𝒜.jacobian(𝒷(), x -> 𝓂.model_function(x, par, SS), x), x), [SS_future; SS_present; SS_past; shocks_ss] ), 𝓂.timings.nVars, nk^3))#, SS_and_pars - # return 𝓂.model_third_order_derivatives([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]) + # return 𝓂.model_third_order_derivatives([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss; par; SS[dyn_ss_idx]]) * 𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃ - third_out = [f([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]) for f in 𝓂.model_third_order_derivatives] + # third_out = [f([SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; shocks_ss], par, SS[dyn_ss_idx]) for f in 𝓂.model_third_order_derivatives] - vals = [i[1] for i in third_out] - rows = [i[2] for i in third_out] - cols = [i[3] for i in third_out] + # vals = [i[1] for i in third_out] + # rows = [i[2] for i in third_out] + # cols = [i[3] for i in third_out] - vals = convert(Vector{M}, vals) + # vals = convert(Vector{M}, vals) + + X = [SS[[dyn_var_future_idx; dyn_var_present_idx; dyn_var_past_idx]]; SS[dyn_ss_idx]; par; shocks_ss] + + # vals = M[] - # nk = 𝓂.timings.nPast_not_future_and_mixed + 𝓂.timings.nVars + 𝓂.timings.nFuture_not_past_and_mixed + length(𝓂.exo) - # sparse(rows, cols, vals, length(𝓂.dyn_equations), nk^3) - sparse(rows, cols, vals, length(𝓂.dyn_equations), size(𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃,1)) * 𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃ + # for f in 𝓂.model_third_order_derivatives[1] + # push!(vals, f(X)...) + # end + + vals = zeros(M, length(𝓂.model_third_order_derivatives[1])) + + # lk = ReentrantLock() + + Polyester.@batch minbatch = 200 for f in 𝓂.model_third_order_derivatives[1] + out = f(X) + + # begin + # lock(lk) + # try + @inbounds vals[out[2]] = out[1] + # finally + # unlock(lk) + # end + # end + end + + Accessors.@reset 𝓂.model_third_order_derivatives[2].nzval = vals + + return 𝓂.model_third_order_derivatives[2] * 𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃ + + # vals = M[] + # rows = Int[] + # cols = Int[] + + # for f in 𝓂.model_third_order_derivatives + # output = f(input) + + # push!(vals, output[1]...) + # push!(rows, output[2]...) + # push!(cols, output[3]...) + # end + + # # # nk = 𝓂.timings.nPast_not_future_and_mixed + 𝓂.timings.nVars + 𝓂.timings.nFuture_not_past_and_mixed + length(𝓂.exo) + # # # sparse(rows, cols, vals, length(𝓂.dyn_equations), nk^3) + # sparse(rows, cols, vals, length(𝓂.dyn_equations), size(𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃,1)) * 𝓂.solution.perturbation.third_order_auxilliary_matrices.𝐔∇₃ end @@ -5606,8 +6264,16 @@ function calculate_quadratic_iteration_solution(∇₁::AbstractMatrix{Float64}; C = similar(A) C̄ = similar(A) + E = similar(C) + sol = @suppress begin - speedmapping(zero(A); m! = (C̄, C) -> C̄ .= A + B * C^2, tol = tol, maps_limit = 10000) + speedmapping(zero(A); m! = (C̄, C) -> begin + ℒ.mul!(E, C, C) + ℒ.mul!(C̄, B, E) + ℒ.axpy!(1, A, C̄) + end, + # C̄ .= A + B * C^2, + tol = tol, maps_limit = 10000) end C = -sol.minimizer @@ -5660,56 +6326,77 @@ function calculate_quadratic_iteration_solution_AD(∇₁::AbstractMatrix{S}; T: iter += 1 end - D = -(∇₊ * -C + ∇₀) \ ∇ₑ + C̄ = ℒ.lu(∇₊ * -C + ∇₀, check = false) + + if !ℒ.issuccess(C̄) + return -C, false + end + + D = -inv(C̄) * ∇ₑ return hcat(-C[:, T.past_not_future_and_mixed_idx], D), error <= tol end function riccati_forward(∇₁::Matrix{Float64}; T::timings, explosive::Bool = false)::Tuple{Matrix{Float64},Bool} + n₀₊ = zeros(T.nVars, T.nFuture_not_past_and_mixed) + n₀₀ = zeros(T.nVars, T.nVars) + n₀₋ = zeros(T.nVars, T.nPast_not_future_and_mixed) + n₋₋ = zeros(T.nPast_not_future_and_mixed, T.nPast_not_future_and_mixed) + nₚ₋ = zeros(T.nPresent_only, T.nPast_not_future_and_mixed) + nₜₚ = zeros(T.nVars - T.nPresent_only, T.nPast_not_future_and_mixed) + ∇₊ = @view ∇₁[:,1:T.nFuture_not_past_and_mixed] - ∇₀ = @view ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1, T.nVars)] + ∇₀ = ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1, T.nVars)] ∇₋ = @view ∇₁[:,T.nFuture_not_past_and_mixed + T.nVars .+ range(1, T.nPast_not_future_and_mixed)] - Q = ℒ.qr(collect(∇₀[:,T.present_only_idx])) + Q = ℒ.qr!(∇₀[:,T.present_only_idx]) Qinv = Q.Q' - A₊ = Qinv * ∇₊ - A₀ = Qinv * ∇₀ - A₋ = Qinv * ∇₋ + mul!(n₀₊, Qinv, ∇₊) + mul!(n₀₀, Qinv, ∇₀) + mul!(n₀₋, Qinv, ∇₋) + A₊ = n₀₊ + A₀ = n₀₀ + A₋ = n₀₋ dynIndex = T.nPresent_only+1:T.nVars - Ã₊ = @view A₊[dynIndex,:] - Ã₋ = @view A₋[dynIndex,:] - Ã₀₊ = @view A₀[dynIndex, T.future_not_past_and_mixed_idx] - Ã₀₋ = @views A₀[dynIndex, T.past_not_future_idx] * ℒ.diagm(ones(T.nPast_not_future_and_mixed))[T.not_mixed_in_past_idx,:] - - Z₊ = zeros(T.nMixed,T.nFuture_not_past_and_mixed) - I₊ = @view ℒ.diagm(ones(T.nFuture_not_past_and_mixed))[T.mixed_in_future_idx,:] + Ã₊ = A₊[dynIndex,:] + Ã₋ = A₋[dynIndex,:] + Ã₀₊ = A₀[dynIndex, T.future_not_past_and_mixed_idx] + @views mul!(nₜₚ, A₀[dynIndex, T.past_not_future_idx], ℒ.I(T.nPast_not_future_and_mixed)[T.not_mixed_in_past_idx,:]) + Ã₀₋ = nₜₚ + + Z₊ = zeros(T.nMixed, T.nFuture_not_past_and_mixed) + I₊ = ℒ.I(T.nFuture_not_past_and_mixed)[T.mixed_in_future_idx,:] Z₋ = zeros(T.nMixed,T.nPast_not_future_and_mixed) - I₋ = @view ℒ.diagm(ones(T.nPast_not_future_and_mixed))[T.mixed_in_past_idx,:] + I₋ = ℒ.diagm(ones(T.nPast_not_future_and_mixed))[T.mixed_in_past_idx,:] D = vcat(hcat(Ã₀₋, Ã₊), hcat(I₋, Z₊)) - E = vcat(hcat(-Ã₋,-Ã₀₊), hcat(Z₋, I₊)) + + ℒ.rmul!(Ã₋,-1) + ℒ.rmul!(Ã₀₊,-1) + E = vcat(hcat(Ã₋,Ã₀₊), hcat(Z₋, I₊)) + # this is the companion form and by itself the linearisation of the matrix polynomial used in the linear time iteration method. see: https://opus4.kobv.de/opus4-matheon/files/209/240.pdf schdcmp = try - ℒ.schur(D, E) + ℒ.schur!(D, E) catch return zeros(T.nVars,T.nPast_not_future_and_mixed), false end - + if explosive # returns false for NaN gen. eigenvalue which is correct here bc they are > 1 eigenselect = abs.(schdcmp.β ./ schdcmp.α) .>= 1 ℒ.ordschur!(schdcmp, eigenselect) - Z₂₁ = @view schdcmp.Z[T.nPast_not_future_and_mixed+1:end, 1:T.nPast_not_future_and_mixed] + Z₂₁ = schdcmp.Z[T.nPast_not_future_and_mixed+1:end, 1:T.nPast_not_future_and_mixed] Z₁₁ = @view schdcmp.Z[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] S₁₁ = @view schdcmp.S[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] - T₁₁ = @view schdcmp.T[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] + T₁₁ = schdcmp.T[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] Ẑ₁₁ = RF.lu(Z₁₁, check = false) @@ -5729,12 +6416,11 @@ function riccati_forward(∇₁::Matrix{Float64}; T::timings, explosive::Bool = return zeros(T.nVars,T.nPast_not_future_and_mixed), false end - Z₂₁ = @view schdcmp.Z[T.nPast_not_future_and_mixed+1:end, 1:T.nPast_not_future_and_mixed] + Z₂₁ = schdcmp.Z[T.nPast_not_future_and_mixed+1:end, 1:T.nPast_not_future_and_mixed] Z₁₁ = @view schdcmp.Z[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] S₁₁ = @view schdcmp.S[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] - T₁₁ = @view schdcmp.T[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] - + T₁₁ = schdcmp.T[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] Ẑ₁₁ = RF.lu(Z₁₁, check = false) @@ -5742,34 +6428,61 @@ function riccati_forward(∇₁::Matrix{Float64}; T::timings, explosive::Bool = return zeros(T.nVars,T.nPast_not_future_and_mixed), false end end - - Ŝ₁₁ = RF.lu(S₁₁, check = false) + + if VERSION >= v"1.9" + Ŝ₁₁ = RF.lu!(S₁₁, check = false) + else + Ŝ₁₁ = RF.lu(S₁₁, check = false) + end if !ℒ.issuccess(Ŝ₁₁) return zeros(T.nVars,T.nPast_not_future_and_mixed), false end - - D = Z₂₁ / Ẑ₁₁ - L = Z₁₁ * (Ŝ₁₁ \ T₁₁) / Ẑ₁₁ - sol = @views vcat(L[T.not_mixed_in_past_idx,:], D) + # D = Z₂₁ / Ẑ₁₁ + ℒ.rdiv!(Z₂₁, Ẑ₁₁) + D = Z₂₁ + + # L = Z₁₁ * (Ŝ₁₁ \ T₁₁) / Ẑ₁₁ + ℒ.ldiv!(Ŝ₁₁, T₁₁) + mul!(n₋₋, Z₁₁, T₁₁) + ℒ.rdiv!(n₋₋, Ẑ₁₁) + L = n₋₋ + + sol = vcat(L[T.not_mixed_in_past_idx,:], D) Ā₀ᵤ = @view A₀[1:T.nPresent_only, T.present_only_idx] A₊ᵤ = @view A₊[1:T.nPresent_only,:] - Ã₀ᵤ = @view A₀[1:T.nPresent_only, T.present_but_not_only_idx] - A₋ᵤ = @view A₋[1:T.nPresent_only,:] + Ã₀ᵤ = A₀[1:T.nPresent_only, T.present_but_not_only_idx] + A₋ᵤ = A₋[1:T.nPresent_only,:] - Ā̂₀ᵤ = RF.lu(Ā₀ᵤ, check = false) + + if VERSION >= v"1.9" + Ā̂₀ᵤ = RF.lu!(Ā₀ᵤ, check = false) + else + Ā̂₀ᵤ = RF.lu(Ā₀ᵤ, check = false) + end if !ℒ.issuccess(Ā̂₀ᵤ) - Ā̂₀ᵤ = ℒ.svd(collect(Ā₀ᵤ)) + return zeros(T.nVars,T.nPast_not_future_and_mixed), false + # Ā̂₀ᵤ = ℒ.svd(collect(Ā₀ᵤ)) end - A = @views vcat(-(Ā̂₀ᵤ \ (A₊ᵤ * D * L + Ã₀ᵤ * sol[T.dynamic_order,:] + A₋ᵤ)), sol) - + # A = vcat(-(Ā̂₀ᵤ \ (A₊ᵤ * D * L + Ã₀ᵤ * sol[T.dynamic_order,:] + A₋ᵤ)), sol) + if T.nPresent_only > 0 + mul!(A₋ᵤ, Ã₀ᵤ, sol[T.dynamic_order,:], 1, 1) + mul!(nₚ₋, A₊ᵤ, D) + mul!(A₋ᵤ, nₚ₋, L, 1, 1) + ℒ.ldiv!(Ā̂₀ᵤ, A₋ᵤ) + ℒ.rmul!(A₋ᵤ,-1) + end + A = vcat(A₋ᵤ, sol) + return A[T.reorder,:], true end + + function riccati_conditions(∇₁::AbstractMatrix{M}, sol_d::AbstractMatrix{N}, solved::Bool; T::timings, explosive::Bool = false) where {M,N} expand = @ignore_derivatives [ℒ.diagm(ones(T.nVars))[T.future_not_past_and_mixed_idx,:], ℒ.diagm(ones(T.nVars))[T.past_not_future_and_mixed_idx,:]] @@ -5800,7 +6513,6 @@ function riccati_forward(∇₁::Matrix{ℱ.Dual{Z,S,N}}; T::timings, explosive: B = 𝒜.jacobian(𝒷(), x -> riccati_conditions(x, val, solved; T = T), ∇̂₁)[1] A = 𝒜.jacobian(𝒷(), x -> riccati_conditions(∇̂₁, x, solved; T = T), val)[1] - Â = RF.lu(A, check = false) if !ℒ.issuccess(Â) @@ -5820,7 +6532,7 @@ end # @memoize LRU(maxsize=50) function calculate_jacobian_transpose(∇₁::AbstractMatrix{Float64}; T::timings, explosive::Bool = false) - 𝐒₁, solved = MacroModelling.riccati_forward(∇₁;T = T, explosive = false) + 𝐒₁, solved = MacroModelling.riccati_forward(∇₁; T = T, explosive = false) sp𝐒₁ = sparse(𝐒₁) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC sp∇₁ = sparse(∇₁) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC @@ -5845,10 +6557,10 @@ function calculate_jacobian_transpose(∇₁::AbstractMatrix{Float64}; T::timing droptol!(spd𝐒₁a, 10*eps()) - d𝐒₁a = spd𝐒₁a' |> collect + # d𝐒₁a = spd𝐒₁a' |> collect # bottleneck, reduce size, avoid conversion, subselect necessary part of matrix already here (as is done in the estimation part later) # Initialize empty spd∇₁a - spd∇₁a = spzeros(length(sp𝐒₁), length(∇₁)) + spd∇₁a = spzeros(length(sp𝐒₁), length(∇₁)) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC # Directly allocate dA, dB, dC into spd∇₁a # Note: You need to calculate the column indices where each matrix starts and ends @@ -5857,37 +6569,82 @@ function calculate_jacobian_transpose(∇₁::AbstractMatrix{Float64}; T::timing dB_cols = dA_cols[end] .+ (1 : size(𝐒₁, 1)^2) dC_cols = dB_cols[end] .+ (1 : length(sp𝐒₁)) - spd∇₁a[:,dA_cols] = ℒ.kron(expand[1] * sol_buf2 * expand[2]' , ℒ.I(size(𝐒₁, 1)))' - spd∇₁a[:,dB_cols] = ℒ.kron(sp𝐒₁, ℒ.I(size(𝐒₁, 1)))' - spd∇₁a[:,dC_cols] = ℒ.I(length(𝐒₁)) + spd∇₁a[:,dA_cols] = ℒ.kron(expand[1] * sol_buf2 * expand[2]' , -ℒ.I(size(𝐒₁, 1)))' + spd∇₁a[:,dB_cols] = ℒ.kron(sp𝐒₁, -ℒ.I(size(𝐒₁, 1)))' + spd∇₁a[:,dC_cols] = -ℒ.I(length(𝐒₁)) - d𝐒₁â = ℒ.lu(d𝐒₁a, check = false) + d𝐒₁â = ℒ.lu(spd𝐒₁a', check = false) if !ℒ.issuccess(d𝐒₁â) tmp = spd∇₁a' solved = false else - tmp = -(d𝐒₁â \ spd∇₁a)' + tmp = inv(d𝐒₁â) * spd∇₁a # bottleneck, reduce size, avoid conversion end - return 𝐒₁, solved, tmp + return 𝐒₁, solved, tmp' end -function rrule(::typeof(riccati_forward), ∇₁::AbstractMatrix{Float64}; T::timings, explosive::Bool = false) +# function rrule(::typeof(riccati_forward), ∇₁; T, explosive = false) +# # Forward pass to compute the output and intermediate values needed for the backward pass +# 𝐒₁, solved, tmp = calculate_jacobian_transpose(∇₁, T = T, explosive = explosive) + +# function calculate_riccati_pullback(∂𝐒₁) +# # Backward pass to compute the derivatives with respect to inputs +# # This would involve computing the derivatives for each operation in reverse order +# # and applying chain rule to propagate through the function +# return NoTangent(), reshape(tmp * sparsevec(∂𝐒₁[1]), size(∇₁)) # Return NoTangent() for non-Array inputs or if there's no derivative w.r.t. them +# # return NoTangent(), (reshape(-d𝐒₁a \ d∇₁a * vec(∂𝐒₁) , size(∇₁))) # Return NoTangent() for non-Array inputs or if there's no derivative w.r.t. them +# end + +# return (𝐒₁, solved), calculate_riccati_pullback +# end + + + +function rrule(::typeof(riccati_forward), ∇₁; T, explosive = false) # Forward pass to compute the output and intermediate values needed for the backward pass - 𝐒₁, solved, tmp = calculate_jacobian_transpose(∇₁, T = T, explosive = explosive) + A, solved = riccati_forward(∇₁, T = T, explosive = explosive) + + expand = @views [ℒ.diagm(ones(T.nVars))[T.future_not_past_and_mixed_idx,:], + ℒ.diagm(ones(T.nVars))[T.past_not_future_and_mixed_idx,:]] + + Â = A * expand[2] + + ∇₊ = @views ∇₁[:,1:T.nFuture_not_past_and_mixed] * expand[1] + ∇₀ = @views ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] + + ∂∇₁ = zero(∇₁) + + invtmp = inv(-Â' * ∇₊' - ∇₀') + + tmp2 = invtmp * ∇₊' + + function first_order_solution_pullback(∂A) + tmp1 = invtmp * ∂A[1] * expand[2] + + coordinates = Tuple{Vector{Int}, Vector{Int}}[] + + values = vcat(vec(tmp2), vec(Â'), vec(tmp1)) + + dimensions = Tuple{Int, Int}[] + push!(dimensions,size(tmp2)) + push!(dimensions,size(Â')) + push!(dimensions,size(tmp1)) + + ss, solved = solve_matrix_equation_forward(values, coords = coordinates, dims = dimensions, solver = :sylvester)#, tol = eps()) # potentially high matrix condition numbers. precision matters + + + ∂∇₁[:,1:T.nFuture_not_past_and_mixed] .= (ss * Â' * Â')[:,T.future_not_past_and_mixed_idx] + ∂∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] .= ss * Â' + ∂∇₁[:,T.nFuture_not_past_and_mixed + T.nVars .+ range(1,T.nPast_not_future_and_mixed)] .= ss[:,T.past_not_future_and_mixed_idx] - function calculate_riccati_pullback(Δ𝐒₁) - # Backward pass to compute the derivatives with respect to inputs - # This would involve computing the derivatives for each operation in reverse order - # and applying chain rule to propagate through the function - return NoTangent(), reshape(tmp * sparsevec(Δ𝐒₁[1]), size(∇₁)) # Return NoTangent() for non-Array inputs or if there's no derivative w.r.t. them - # return NoTangent(), (reshape(-d𝐒₁a \ d∇₁a * vec(Δ𝐒₁) , size(∇₁))) # Return NoTangent() for non-Array inputs or if there's no derivative w.r.t. them + return NoTangent(), ∂∇₁ end - return (𝐒₁, solved), calculate_riccati_pullback + return (A, solved), first_order_solution_pullback end @@ -5899,7 +6656,10 @@ riccati_AD_direct = ℐ.ImplicitFunction(riccati_forward, riccati_AD = ℐ.ImplicitFunction(riccati_forward, riccati_conditions) # doesnt converge!? -function calculate_first_order_solution(∇₁::Matrix{Float64}; T::timings, explosive::Bool = false)::Tuple{Matrix{Float64},Bool} + +function calculate_first_order_solution(∇₁::Matrix{Float64}; + T::timings, + explosive::Bool = false)::Tuple{Matrix{Float64}, Bool} # A, solved = riccati_AD_direct(∇₁; T = T, explosive = explosive) A, solved = riccati_forward(∇₁; T = T, explosive = explosive) @@ -5910,27 +6670,109 @@ function calculate_first_order_solution(∇₁::Matrix{Float64}; T::timings, exp Jm = @view(ℒ.diagm(ones(T.nVars))[T.past_not_future_and_mixed_idx,:]) ∇₊ = @views ∇₁[:,1:T.nFuture_not_past_and_mixed] * ℒ.diagm(ones(T.nVars))[T.future_not_past_and_mixed_idx,:] - ∇₀ = @view ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] - ∇ₑ = @view ∇₁[:,(T.nFuture_not_past_and_mixed + T.nVars + T.nPast_not_future_and_mixed + 1):end] - - B = -((∇₊ * A * Jm + ∇₀) \ ∇ₑ) + ∇₀ = copy(∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)]) + ∇ₑ = copy(∇₁[:,(T.nFuture_not_past_and_mixed + T.nVars + T.nPast_not_future_and_mixed + 1):end]) + + M = similar(∇₀) + mul!(M, A, Jm) + mul!(∇₀, ∇₊, M, 1, 1) + C = RF.lu!(∇₀, check = false) + # C = RF.lu!(∇₊ * A * Jm + ∇₀, check = false) + + if !ℒ.issuccess(C) + return hcat(A, zeros(size(A,1),T.nExo)), solved + end + + ℒ.ldiv!(C, ∇ₑ) + ℒ.rmul!(∇ₑ, -1) + # B = -(C \ ∇ₑ) # otherwise Zygote doesnt diff it - return hcat(A, B), solved + return hcat(A, ∇ₑ), solved end -function calculate_first_order_solution(∇₁::Matrix{ℱ.Dual{Z,S,N}}; T::timings, explosive::Bool = false)::Tuple{Matrix{ℱ.Dual{Z,S,N}},Bool} where {Z,S,N} - A, solved = riccati_AD_direct(∇₁; T = T, explosive = explosive) - # A, solved = riccati_forward(∇₁; T = T, explosive = explosive) + +function rrule(::typeof(calculate_first_order_solution), ∇₁; T, explosive = false) + # Forward pass to compute the output and intermediate values needed for the backward pass + 𝐒ᵗ, solved = riccati_forward(∇₁, T = T, explosive = explosive) if !solved - return hcat(A, zeros(size(A,1),T.nExo)), solved + return (hcat(𝐒ᵗ, zeros(size(𝐒ᵗ,1),T.nExo)), solved), x -> NoTangent(), NoTangent(), NoTangent() end - Jm = @view(ℒ.diagm(ones(S,T.nVars))[T.past_not_future_and_mixed_idx,:]) - - ∇₊ = @views ∇₁[:,1:T.nFuture_not_past_and_mixed] * ℒ.diagm(ones(S,T.nVars))[T.future_not_past_and_mixed_idx,:] - ∇₀ = @view ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] + expand = @views [ℒ.diagm(ones(T.nVars))[T.future_not_past_and_mixed_idx,:], + ℒ.diagm(ones(T.nVars))[T.past_not_future_and_mixed_idx,:]] + + ∇₊ = @views ∇₁[:,1:T.nFuture_not_past_and_mixed] * expand[1] + ∇₀ = @view ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] + ∇ₑ = @view ∇₁[:,(T.nFuture_not_past_and_mixed + T.nVars + T.nPast_not_future_and_mixed + 1):end] + + M̂ = RF.lu(∇₊ * 𝐒ᵗ * expand[2] + ∇₀, check = false) + + if !ℒ.issuccess(M̂) + return (hcat(𝐒ᵗ, zeros(size(𝐒ᵗ,1),T.nExo)), solved), x -> NoTangent(), NoTangent(), NoTangent() + end + + M = inv(M̂) + + 𝐒ᵉ = -M * ∇ₑ # otherwise Zygote doesnt diff it + + 𝐒̂ᵗ = 𝐒ᵗ * expand[2] + + ∂∇₁ = zero(∇₁) + + tmp2 = -M' * ∇₊' + + function first_order_solution_pullback(∂𝐒) + ∂𝐒ᵗ = ∂𝐒[1][:,1:T.nPast_not_future_and_mixed] + ∂𝐒ᵉ = ∂𝐒[1][:,T.nPast_not_future_and_mixed + 1:end] + + ∂∇₁[:,T.nFuture_not_past_and_mixed + T.nVars + T.nPast_not_future_and_mixed + 1:end] .= -M' * ∂𝐒ᵉ + + ∂∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] .= M' * ∂𝐒ᵉ * ∇ₑ' * M' + + ∂∇₁[:,1:T.nFuture_not_past_and_mixed] .= (M' * ∂𝐒ᵉ * ∇ₑ' * M' * expand[2]' * 𝐒ᵗ')[:,T.future_not_past_and_mixed_idx] + + ∂𝐒ᵗ .+= ∇₊' * M' * ∂𝐒ᵉ * ∇ₑ' * M' * expand[2]' + + tmp1 = -M' * ∂𝐒ᵗ * expand[2] + + coordinates = Tuple{Vector{Int}, Vector{Int}}[] + + values = vcat(vec(tmp2), vec(𝐒̂ᵗ'), vec(-tmp1)) + + dimensions = Tuple{Int, Int}[] + push!(dimensions,size(tmp2)) + push!(dimensions,size(𝐒̂ᵗ')) + push!(dimensions,size(tmp1)) + + ss, solved = solve_matrix_equation_forward(values, coords = coordinates, dims = dimensions, solver = :sylvester)#, tol = eps()) # potentially high matrix condition numbers. precision matters + if !solved + NoTangent(), NoTangent(), NoTangent() + end + + ∂∇₁[:,1:T.nFuture_not_past_and_mixed] .+= (ss * 𝐒̂ᵗ' * 𝐒̂ᵗ')[:,T.future_not_past_and_mixed_idx] + ∂∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] .+= ss * 𝐒̂ᵗ' + ∂∇₁[:,T.nFuture_not_past_and_mixed + T.nVars .+ range(1,T.nPast_not_future_and_mixed)] .+= ss[:,T.past_not_future_and_mixed_idx] + + return NoTangent(), ∂∇₁, NoTangent() + end + + return (hcat(𝐒ᵗ, 𝐒ᵉ), solved), first_order_solution_pullback +end + +function calculate_first_order_solution(∇₁::Matrix{ℱ.Dual{Z,S,N}}; T::timings, explosive::Bool = false)::Tuple{Matrix{ℱ.Dual{Z,S,N}},Bool} where {Z,S,N} + A, solved = riccati_AD_direct(∇₁; T = T, explosive = explosive) + # A, solved = riccati_forward(∇₁; T = T, explosive = explosive) + + if !solved + return hcat(A, zeros(size(A,1),T.nExo)), solved + end + + Jm = @view(ℒ.diagm(ones(S,T.nVars))[T.past_not_future_and_mixed_idx,:]) + + ∇₊ = @views ∇₁[:,1:T.nFuture_not_past_and_mixed] * ℒ.diagm(ones(S,T.nVars))[T.future_not_past_and_mixed_idx,:] + ∇₀ = @view ∇₁[:,T.nFuture_not_past_and_mixed .+ range(1,T.nVars)] ∇ₑ = @view ∇₁[:,(T.nFuture_not_past_and_mixed + T.nVars + T.nPast_not_future_and_mixed + 1):end] B = -((∇₊ * A * Jm + ∇₀) \ ∇ₑ) @@ -6126,8 +6968,9 @@ function calculate_third_order_solution(∇₁::AbstractMatrix{<: Real}, #first tmpkron = ℒ.kron(𝐒₁₋╱𝟏ₑ,M₂.𝛔) C = M₃.𝐔₃ * tmpkron + M₃.𝐔₃ * M₃.𝐏₁ₗ̄ * tmpkron * M₃.𝐏₁ᵣ̃ + M₃.𝐔₃ * M₃.𝐏₂ₗ̄ * tmpkron * M₃.𝐏₂ᵣ̃ - C += M₃.𝐔₃ * ℒ.kron(𝐒₁₋╱𝟏ₑ,ℒ.kron(𝐒₁₋╱𝟏ₑ,𝐒₁₋╱𝟏ₑ)) # no speed up here from A_mult_kron_power_3_B + C += M₃.𝐔₃ * ℒ.kron(𝐒₁₋╱𝟏ₑ,ℒ.kron(𝐒₁₋╱𝟏ₑ,𝐒₁₋╱𝟏ₑ)) # no speed up here from A_mult_kron_power_3_B; this is the bottleneck. ideally have this return reduced space directly. TODO: make kron3 faster C *= M₃.𝐂₃ + # C += kron³(𝐒₁₋╱𝟏ₑ, M₃) droptol!(C,tol) r1,c1,v1 = findnz(B) @@ -6681,7 +7524,7 @@ function calculate_covariance(parameters::Vector{<: Real}, 𝓂::ℳ; verbose::B ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂) - sol, solved = calculate_first_order_solution(Matrix(∇₁); T = 𝓂.timings) + sol, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) # covar_raw, solved_cov = calculate_covariance_AD(sol, T = 𝓂.timings, subset_indices = collect(1:𝓂.timings.nVars)) @@ -6718,7 +7561,7 @@ function calculate_mean(parameters::Vector{T}, 𝓂::ℳ; verbose::Bool = false, return SS_and_pars[1:𝓂.timings.nVars], solution_error end - ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂)# |> Matrix 𝐒₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -6789,7 +7632,13 @@ function solve_matrix_equation_forward(ABC::Vector{Float64}; if length(coords) == 1 lengthA = length(coords[1][1]) vA = ABC[1:lengthA] - A = sparse(coords[1]...,vA,dims[1]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + + if VERSION >= v"1.9" + A = sparse(coords[1]...,vA,dims[1]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + else + A = sparse(coords[1]...,vA,dims[1]...)# |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + end + C = reshape(ABC[lengthA+1:end],dims[2]...) if solver != :doubling B = A' @@ -6802,9 +7651,23 @@ function solve_matrix_equation_forward(ABC::Vector{Float64}; vB = ABC[lengthA .+ (1:lengthB)] vC = ABC[lengthA + lengthB + 1:end] - A = sparse(coords[1]...,vA,dims[1]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC - B = sparse(coords[2]...,vB,dims[2]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC - C = sparse(coords[3]...,vC,dims[3]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + if VERSION >= v"1.9" + A = sparse(coords[1]...,vA,dims[1]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + B = sparse(coords[2]...,vB,dims[2]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + C = sparse(coords[3]...,vC,dims[3]...) |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + else + A = sparse(coords[1]...,vA,dims[1]...)# |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + B = sparse(coords[2]...,vB,dims[2]...)# |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + C = sparse(coords[3]...,vC,dims[3]...)# |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + end + + elseif length(dims) == 3 + lengthA = dims[1][1] * dims[1][2] + lengthB = dims[2][1] * dims[2][2] + + A = reshape(ABC[1:lengthA], dims[1]...) + B = reshape(ABC[lengthA .+ (1:lengthB)], dims[2]...) + C = reshape(ABC[lengthA + lengthB + 1:end], dims[3]...) else lengthA = dims[1][1] * dims[1][2] A = reshape(ABC[1:lengthA],dims[1]...) @@ -6815,22 +7678,36 @@ function solve_matrix_equation_forward(ABC::Vector{Float64}; end - if solver ∈ [:gmres, :bicgstab] + if solver ∈ [:gmres, :bicgstab] + # tmp̂ = similar(C) + # tmp̄ = similar(C) + # 𝐗 = similar(C) + + # function sylvester!(sol,𝐱) + # copyto!(𝐗, 𝐱) + # mul!(tmp̄, 𝐗, B) + # mul!(tmp̂, A, tmp̄) + # ℒ.axpy!(-1, tmp̂, 𝐗) + # ℒ.rmul!(𝐗, -1) + # copyto!(sol, 𝐗) + # end + # TODO: above is slower. below is fastest function sylvester!(sol,𝐱) 𝐗 = reshape(𝐱, size(C)) - sol .= vec(A * 𝐗 * B - 𝐗) - return sol + copyto!(sol, A * 𝐗 * B - 𝐗) + # sol .= vec(A * 𝐗 * B - 𝐗) + # return sol end sylvester = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, sylvester!) if solver == :gmres - 𝐂, info = Krylov.gmres(sylvester, [vec(C);]) + 𝐂, info = Krylov.gmres(sylvester, [vec(C);])#, rtol = Float64(tol)) elseif solver == :bicgstab - 𝐂, info = Krylov.bicgstab(sylvester, [vec(C);]) + 𝐂, info = Krylov.bicgstab(sylvester, [vec(C);])#, rtol = Float64(tol)) end solved = info.solved - elseif solver == :iterative + elseif solver == :iterative # this can still be optimised iter = 1 change = 1 𝐂 = C @@ -6847,22 +7724,45 @@ function solve_matrix_equation_forward(ABC::Vector{Float64}; iter += 1 end solved = change < eps(Float32) - elseif solver == :doubling + elseif solver == :doubling # cant use higher tol because rersults get weird in some cases iter = 1 change = 1 𝐂 = -C 𝐂¹ = -C + CA = similar(A) + A² = similar(A) while change > eps(Float32) && iter < 500 - 𝐂¹ = A * 𝐂 * A' + 𝐂 - A *= A - if !(A isa DenseMatrix) + if A isa DenseMatrix + + mul!(CA, 𝐂, A') + mul!(𝐂¹, A, CA, 1, 1) + + mul!(A², A, A) + copy!(A, A²) + + if iter > 10 + ℒ.axpy!(-1, 𝐂¹, 𝐂) + change = maximum(abs, 𝐂) + end + + copy!(𝐂, 𝐂¹) + + iter += 1 + else + 𝐂¹ = A * 𝐂 * A' + 𝐂 + + A *= A + droptol!(A, eps()) + + if iter > 10 + change = maximum(abs, 𝐂¹ - 𝐂) + end + + 𝐂 = 𝐂¹ + + iter += 1 end - if iter > 10 - change = maximum(abs, 𝐂¹ - 𝐂) - end - 𝐂 = 𝐂¹ - iter += 1 end solved = change < eps(Float32) elseif solver == :sylvester @@ -6876,9 +7776,16 @@ function solve_matrix_equation_forward(ABC::Vector{Float64}; 𝐂 = MatrixEquations.lyapd(collect(A),-C) solved = isapprox(𝐂, A * 𝐂 * A' - C, rtol = eps(Float32)) elseif solver == :speedmapping - + CB = similar(A) + soll = @suppress begin - speedmapping(collect(-C); m! = (X, x) -> X .= A * x * B - C, stabilize = true) + speedmapping(collect(-C); + m! = (X, x) -> begin + mul!(CB, x, B) + mul!(X, A, CB) + ℒ.axpy!(1, C, X) + end, stabilize = false)#, tol = tol) + # speedmapping(collect(-C); m! = (X, x) -> X .= A * x * B - C, stabilize = true) end 𝐂 = soll.minimizer @@ -7452,253 +8359,1227 @@ function calculate_third_order_moments(parameters::Vector{T}, end +function find_variables_to_exclude(𝓂::ℳ, observables::Vector{Symbol}) + # reduce system + vars_to_exclude = setdiff(𝓂.timings.present_only, observables) + + # Mapping variables to their equation index + variable_to_equation = Dict{Symbol, Vector{Int}}() + for var in vars_to_exclude + for (eq_idx, vars_set) in enumerate(𝓂.dyn_var_present_list) + # for var in vars_set + if var in vars_set + if haskey(variable_to_equation, var) + push!(variable_to_equation[var],eq_idx) + else + variable_to_equation[var] = [eq_idx] + end + end + end + end + + return variable_to_equation +end -function calculate_kalman_filter_loglikelihood(𝓂::ℳ, observables::Vector{Symbol}, 𝐒₁::Matrix{S}, data_in_deviations::Matrix{S})::S where S - obs_idx = @ignore_derivatives convert(Vector{Int},indexin(observables,sort(union(𝓂.aux,𝓂.var,𝓂.exo_present)))) - calculate_kalman_filter_loglikelihood(𝓂, obs_idx, 𝐒₁, data_in_deviations) +function create_broadcaster(indices::Vector{Int}, n::Int) + broadcaster = spzeros(n, length(indices)) + for (i, vid) in enumerate(indices) + broadcaster[vid,i] = 1.0 + end + return broadcaster end -function calculate_kalman_filter_loglikelihood(𝓂::ℳ, observables::Vector{String}, 𝐒₁::Matrix{S}, data_in_deviations::Matrix{S})::S where S - obs_idx = @ignore_derivatives convert(Vector{Int},indexin(observables,sort(union(𝓂.aux,𝓂.var,𝓂.exo_present)))) - calculate_kalman_filter_loglikelihood(𝓂, obs_idx, 𝐒₁, data_in_deviations) +# Specialization for :kalman filter +function calculate_loglikelihood(::Val{:kalman}, observables, 𝐒, data_in_deviations, TT, presample_periods, initial_covariance, state, warmup_iterations) + return calculate_kalman_filter_loglikelihood(observables, 𝐒, data_in_deviations, TT, presample_periods = presample_periods, initial_covariance = initial_covariance) end -function calculate_kalman_filter_loglikelihood(𝓂::ℳ, observables_index::Vector{Int}, 𝐒₁::Matrix{S}, data_in_deviations::Matrix{S})::S where S - observables_and_states = @ignore_derivatives sort(union(𝓂.timings.past_not_future_and_mixed_idx,observables_index)) +# Specialization for :inversion filter +function calculate_loglikelihood(::Val{:inversion}, observables, 𝐒, data_in_deviations, TT, presample_periods, initial_covariance, state, warmup_iterations) + return calculate_inversion_filter_loglikelihood(state, 𝐒, data_in_deviations, observables, TT, warmup_iterations = warmup_iterations, presample_periods = presample_periods) +end - A = 𝐒₁[observables_and_states,1:𝓂.timings.nPast_not_future_and_mixed] * ℒ.diagm(ones(length(observables_and_states)))[@ignore_derivatives(indexin(𝓂.timings.past_not_future_and_mixed_idx,observables_and_states)),:] - B = 𝐒₁[observables_and_states,𝓂.timings.nPast_not_future_and_mixed+1:end] +function get_non_stochastic_steady_state(𝓂::ℳ, parameter_values::Vector{S}; verbose::Bool = false, tol::AbstractFloat = 1e-12)::Tuple{Vector{S}, Tuple{S, Int}} where S <: Real + 𝓂.SS_solve_func(parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) +end - C = ℒ.diagm(ones(length(observables_and_states)))[@ignore_derivatives(indexin(sort(observables_index),observables_and_states)),:] - 𝐁 = B * B' +function rrule(::typeof(get_non_stochastic_steady_state), 𝓂, parameter_values; verbose = false, tol::AbstractFloat = 1e-12) + SS_and_pars, (solution_error, iters) = 𝓂.SS_solve_func(parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) - # Gaussian Prior - coordinates = Tuple{Vector{Int}, Vector{Int}}[] + if solution_error > tol || isnan(solution_error) + return (SS_and_pars, (solution_error, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + SS_and_pars_names_lead_lag = vcat(Symbol.(string.(sort(union(𝓂.var,𝓂.exo_past,𝓂.exo_future)))), 𝓂.calibration_equations_parameters) + + SS_and_pars_names = vcat(Symbol.(replace.(string.(sort(union(𝓂.var,𝓂.exo_past,𝓂.exo_future))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")), 𝓂.calibration_equations_parameters) - dimensions = [size(A),size(𝐁)] + # unknowns = union(setdiff(𝓂.vars_in_ss_equations, 𝓂.➕_vars), 𝓂.calibration_equations_parameters) + unknowns = Symbol.(vcat(string.(sort(collect(setdiff(reduce(union,get_symbols.(𝓂.ss_aux_equations)),union(𝓂.parameters_in_equations,𝓂.➕_vars))))), 𝓂.calibration_equations_parameters)) + # ∂SS_equations_∂parameters = try 𝓂.∂SS_equations_∂parameters(parameter_values, SS_and_pars[indexin(unknowns, SS_and_pars_names_lead_lag)]) |> Matrix + # catch + # return (SS_and_pars, (10, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) + # end + + X = [parameter_values; SS_and_pars[indexin(unknowns, SS_and_pars_names_lead_lag)]] - values = vcat(vec(A), vec(collect(-𝐁))) + # vals = Float64[] - P, _ = solve_matrix_equation_AD(values, coords = coordinates, dims = dimensions, solver = :doubling) - # P = reshape((ℒ.I - ℒ.kron(A, A)) \ reshape(𝐁, prod(size(A)), 1), size(A)) - # P = collect(ℒ.I(length(observables_and_states)) * 10.0) + # for f in 𝓂.∂SS_equations_∂parameters[1] + # push!(vals, f(X)...) + # end - u = zeros(S, length(observables_and_states)) - # u = SS_and_pars[sort(union(𝓂.timings.past_not_future_and_mixed,observables))] |> collect - z = C * u + vals = zeros(Float64, length(𝓂.∂SS_equations_∂parameters[1])) - loglik = S(0) - for t in 1:size(data_in_deviations, 2) - v = data_in_deviations[:, t] - z + # lk = ReentrantLock() - F = C * P * C' + Polyester.@batch minbatch = 200 for f in 𝓂.∂SS_equations_∂parameters[1] + out = f(X) + + # begin + # lock(lk) + # try + @inbounds vals[out[2]] = out[1] + # finally + # unlock(lk) + # end + # end + end - F̄ = ℒ.lu(F, check = false) + Accessors.@reset 𝓂.∂SS_equations_∂parameters[2].nzval = vals + + ∂SS_equations_∂parameters = 𝓂.∂SS_equations_∂parameters[2] - if !ℒ.issuccess(F̄) - return -Inf - end + # vals = Float64[] - Fdet = ℒ.det(F̄) + # for f in 𝓂.∂SS_equations_∂SS_and_pars[1] + # push!(vals, f(X)...) + # end - # Early return if determinant is too small, indicating numerical instability. - if Fdet < eps(S) - return -Inf - end + vals = zeros(Float64, length(𝓂.∂SS_equations_∂SS_and_pars[1])) - invF = inv(F̄) + # lk = ReentrantLock() - loglik += log(Fdet) + v' * invF * v + Polyester.@batch minbatch = 200 for f in 𝓂.∂SS_equations_∂SS_and_pars[1] + out = f(X) + + # begin + # lock(lk) + # try + @inbounds vals[out[2]] = out[1] + # finally + # unlock(lk) + # end + # end + end - K = P * C' * invF + 𝓂.∂SS_equations_∂SS_and_pars[3] .*= 0 + 𝓂.∂SS_equations_∂SS_and_pars[3][𝓂.∂SS_equations_∂SS_and_pars[2]] .+= vals - P = A * (P - K * C * P) * A' + 𝐁 + ∂SS_equations_∂SS_and_pars = 𝓂.∂SS_equations_∂SS_and_pars[3] - u = A * (u + K * v) + # ∂SS_equations_∂parameters = 𝓂.∂SS_equations_∂parameters(parameter_values, SS_and_pars[indexin(unknowns, SS_and_pars_names_lead_lag)]) |> Matrix + # ∂SS_equations_∂SS_and_pars = 𝓂.∂SS_equations_∂SS_and_pars(parameter_values, SS_and_pars[indexin(unknowns, SS_and_pars_names_lead_lag)]) |> Matrix + + ∂SS_equations_∂SS_and_pars_lu = RF.lu!(∂SS_equations_∂SS_and_pars, check = false) - z = C * u + if !ℒ.issuccess(∂SS_equations_∂SS_and_pars_lu) + return (SS_and_pars, (10, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) end - return -(loglik + length(data_in_deviations) * log(2 * 3.141592653589793)) / 2 -end + JVP = -(∂SS_equations_∂SS_and_pars_lu \ ∂SS_equations_∂parameters)#[indexin(SS_and_pars_names, unknowns),:] + + jvp = zeros(length(SS_and_pars_names_lead_lag), length(𝓂.parameters)) + + for (i,v) in enumerate(SS_and_pars_names) + if v in unknowns + jvp[i,:] = JVP[indexin([v], unknowns),:] + end + end + # try block-gmres here + function get_non_stochastic_steady_state_pullback(∂SS_and_pars) + # println(∂SS_and_pars) + return NoTangent(), NoTangent(), jvp' * ∂SS_and_pars[1], NoTangent() + end -# function update_loglikelihood!(loglik::S, P::Matrix{S}, u::Vector{S}, z::Vector{S}, C::Matrix{T}, A::Matrix{S}, 𝐁::Matrix{S}, data_point::Vector{S}) where {S,T} -# v = data_point - z -# F = C * P * C' + return (SS_and_pars, (solution_error, iters)), get_non_stochastic_steady_state_pullback +end + -# F̄ = ℒ.lu(F, check = false) +function calculate_kalman_filter_loglikelihood(observables::Vector{Symbol}, + 𝐒::Union{Matrix{S},Vector{AbstractMatrix{S}}}, + data_in_deviations::Matrix{S}, + T::timings; + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical)::S where S <: Real + obs_idx = @ignore_derivatives convert(Vector{Int},indexin(observables,sort(union(T.aux,T.var,T.exo_present)))) -# if !ℒ.issuccess(F̄) -# return -Inf, P, u, z -# end + calculate_kalman_filter_loglikelihood(obs_idx, 𝐒, data_in_deviations, T, presample_periods = presample_periods, initial_covariance = initial_covariance) +end -# Fdet = ℒ.det(F̄) +function calculate_kalman_filter_loglikelihood(observables::Vector{String}, + 𝐒::Union{Matrix{S},Vector{AbstractMatrix{S}}}, + data_in_deviations::Matrix{S}, + T::timings; + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical)::S where S <: Real + obs_idx = @ignore_derivatives convert(Vector{Int},indexin(observables,sort(union(T.aux,T.var,T.exo_present)))) -# # Early return if determinant is too small, indicating numerical instability. -# if Fdet < eps(S) -# return -Inf, P, u, z -# end + calculate_kalman_filter_loglikelihood(obs_idx, 𝐒, data_in_deviations, T, presample_periods = presample_periods, initial_covariance = initial_covariance) +end -# invF = inv(F̄) -# loglik_increment = log(Fdet) + v' * invF * v -# K = P * C' * invF -# P = A * (P - K * C * P) * A' + 𝐁 -# u = A * (u + K * v) -# z = C * u +function calculate_kalman_filter_loglikelihood(observables_index::Vector{Int}, + 𝐒::Union{Matrix{S},Vector{AbstractMatrix{S}}}, + data_in_deviations::Matrix{S}, + T::timings; + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical)::S where S <: Real + observables_and_states = @ignore_derivatives sort(union(T.past_not_future_and_mixed_idx,observables_index)) -# return loglik + loglik_increment, P, u, z -# end -function calculate_inversion_filter_loglikelihood(𝓂::ℳ, state::Union{Vector{Float64},Vector{Vector{Float64}}}, state_update::Function, data_in_deviations::Matrix{Float64}, observables::Union{Vector{String}, Vector{Symbol}}, warmup_iterations::Int) - if state isa Vector{Float64} - pruning = false - else - pruning = true - end + A = 𝐒[observables_and_states,1:T.nPast_not_future_and_mixed] * ℒ.diagm(ones(S, length(observables_and_states)))[@ignore_derivatives(indexin(T.past_not_future_and_mixed_idx,observables_and_states)),:] + B = 𝐒[observables_and_states,T.nPast_not_future_and_mixed+1:end] - precision_factor = 1.0 + C = ℒ.diagm(ones(length(observables_and_states)))[@ignore_derivatives(indexin(sort(observables_index), observables_and_states)),:] - n_obs = size(data_in_deviations,2) + 𝐁 = B * B' - cond_var_idx = indexin(observables,sort(union(𝓂.aux,𝓂.var,𝓂.exo_present))) + # Gaussian Prior + coordinates = @ignore_derivatives Tuple{Vector{Int}, Vector{Int}}[] + + dimensions = @ignore_derivatives [size(A),size(𝐁)] + + values = vcat(vec(A), vec(collect(-𝐁))) - shocks² = 0.0 - logabsdets = 0.0 + P = get_initial_covariance(Val(initial_covariance), values, coordinates, dimensions) - if warmup_iterations > 0 - res = Optim.optimize(x -> minimize_distance_to_initial_data(x, data_in_deviations[:,1], state, state_update, warmup_iterations, cond_var_idx, precision_factor, pruning), - zeros(𝓂.timings.nExo * warmup_iterations), - Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), - Optim.Options(f_abstol = eps(), g_tol= 1e-30); - autodiff = :forward) + return run_kalman_iterations(A, 𝐁, C, P, data_in_deviations, presample_periods = presample_periods) +end - matched = Optim.minimum(res) < 1e-12 - if !matched # for robustness try other linesearch - res = Optim.optimize(x -> minimize_distance_to_initial_data(x, data_in_deviations[:,1], state, state_update, warmup_iterations, cond_var_idx, precision_factor, pruning), - zeros(𝓂.timings.nExo * warmup_iterations), - Optim.LBFGS(), - Optim.Options(f_abstol = eps(), g_tol= 1e-30); - autodiff = :forward) - - matched = Optim.minimum(res) < 1e-12 - end +# Specialization for :theoretical +function get_initial_covariance(::Val{:theoretical}, values::Vector{S}, coordinates, dimensions)::Matrix{S} where S <: Real + P, _ = solve_matrix_equation_AD(values, coords = coordinates, dims = dimensions, solver = :doubling) + return P +end - if !matched return -Inf end +# Specialization for :diagonal +function get_initial_covariance(::Val{:diagonal}, values::Vector{S}, coordinates, dimensions)::Matrix{S} where S <: Real + P = @ignore_derivatives collect(ℒ.I(dimensions[1][1]) * 10.0) + return P +end - x = Optim.minimizer(res) - warmup_shocks = reshape(x, 𝓂.timings.nExo, warmup_iterations) +function rrule(::typeof(get_initial_covariance), + ::Val{:theoretical}, + values, + coordinates, + dimensions) - for i in 1:warmup_iterations-1 - state = state_update(state, warmup_shocks[:,i]) - end - - res = zeros(0) + P, _ = solve_matrix_equation_forward(values, coords = coordinates, dims = dimensions, solver = :doubling) - jacc = zeros(𝓂.timings.nExo * warmup_iterations, length(observables)) + A = reshape(values[1:(dimensions[1][1] * dimensions[1][2])], dimensions[1]) - match_initial_data!(res, x, jacc, data_in_deviations[:,1], state, state_update, warmup_iterations, cond_var_idx, precision_factor), zeros(size(data_in_deviations, 1)) + # pullback + function initial_covariance_pullback(∂P) + values_pb = vcat(vec(A'), vec(-∂P)) - for i in 1:warmup_iterations - if 𝓂.timings.nExo == length(observables) - logabsdets += ℒ.logabsdet(jacc[(i - 1) * 𝓂.timings.nExo .+ (1:2),:] ./ precision_factor)[1] - else - logabsdets += sum(x -> log(abs(x)), ℒ.svdvals(jacc[(i - 1) * 𝓂.timings.nExo .+ (1:2),:] ./ precision_factor)) - end - end + ∂𝐁, _ = solve_matrix_equation_forward(values_pb, coords = coordinates, dims = dimensions, solver = :doubling) + + ∂A = ∂𝐁 * A * P' + ∂𝐁' * A * P - shocks² += sum(abs2,x) + return NoTangent(), NoTangent(), vcat(vec(∂A), vec(-∂𝐁)), NoTangent(), NoTangent() end + + return P, initial_covariance_pullback +end - for i in axes(data_in_deviations,2) - res = Optim.optimize(x -> minimize_distance_to_data(x, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor, pruning), - zeros(𝓂.timings.nExo), - Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), - Optim.Options(f_abstol = eps(), g_tol= 1e-30); - autodiff = :forward) - matched = Optim.minimum(res) < 1e-12 - if !matched # for robustness try other linesearch - res = Optim.optimize(x -> minimize_distance_to_data(x, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor, pruning), - zeros(𝓂.timings.nExo), - Optim.LBFGS(), - Optim.Options(f_abstol = eps(), g_tol= 1e-30); - autodiff = :forward) - - matched = Optim.minimum(res) < 1e-12 - end +function rrule(::typeof(get_initial_covariance), + ::Val{:diagonal}, + values, + coordinates, + dimensions) - if !matched return -Inf end + # pullback + function initial_covariance_pullback(∂P) + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + return collect(ℒ.I(dimensions[1][1]) * 10.0), initial_covariance_pullback +end - x = Optim.minimizer(res) +function run_kalman_iterations(A::Matrix{S}, 𝐁::Matrix{S}, C::Matrix{Float64}, P::Matrix{S}, data_in_deviations::Matrix{S}; presample_periods::Int = 0)::S where S <: Float64 + u = zeros(S, size(C,2)) - res = zeros(0) + z = C * u - jacc = zeros(𝓂.timings.nExo, length(observables)) + ztmp = similar(z) - match_data_sequence!(res, x, jacc, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor) + loglik = S(0.0) - if 𝓂.timings.nExo == length(observables) - logabsdets += ℒ.logabsdet(jacc ./ precision_factor)[1] - else - logabsdets += sum(x -> log(abs(x)), ℒ.svdvals(jacc ./ precision_factor)) - end + utmp = similar(u) - shocks² += sum(abs2,x) + Ctmp = similar(C) - state = state_update(state, x) - end + F = similar(C * C') - return -(logabsdets + shocks² + (𝓂.timings.nExo * (warmup_iterations + n_obs)) * log(2 * 3.141592653589793)) / 2 -end + K = similar(C') + # Ktmp = similar(C') + tmp = similar(P) + Ptmp = similar(P) -function inversion_filter(𝓂::ℳ, - data_in_deviations::AbstractArray{Float64}, - algorithm::Symbol; - warmup_iterations::Int = 0, - verbose::Bool = false, - tol::AbstractFloat = 1e-12) - - observables = collect(axiskeys(data_in_deviations,1)) + for t in 1:size(data_in_deviations, 2) + ℒ.axpby!(1, data_in_deviations[:, t], -1, z) + # v = data_in_deviations[:, t] - z - data_in_deviations = collect(data_in_deviations) - - @assert observables isa Vector{String} || observables isa Vector{Symbol} "Make sure that the data has variables names as rows. They can be either Strings or Symbols." + mul!(Ctmp, C, P) # use Octavian.jl + mul!(F, Ctmp, C') + # F = C * P * C' - sort!(observables) + luF = RF.lu!(F, check = false) ### - observables = observables isa String_input ? observables .|> Meta.parse .|> replace_indices : observables + if !ℒ.issuccess(luF) + return -Inf + end - # solve model given the parameters - if algorithm == :second_order - sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(𝓂.parameter_values, 𝓂) + Fdet = ℒ.det(luF) - if !converged - @error "No solution for these parameters." + # Early return if determinant is too small, indicating numerical instability. + if Fdet < eps(Float64) + return -Inf end - all_SS = expand_steady_state(SS_and_pars,𝓂) - - state = collect(sss) - all_SS + # invF = inv(luF) ### - state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[𝓂.timings.past_not_future_and_mixed_idx] - 1 - shock] - return 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + if t > presample_periods + ℒ.ldiv!(ztmp, luF, z) + loglik += log(Fdet) + ℒ.dot(z', ztmp) ### + # loglik += log(Fdet) + z' * invF * z### + # loglik += log(Fdet) + v' * invF * v### end - elseif algorithm == :pruned_second_order - sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(𝓂.parameter_values, 𝓂, pruning = true) - if !converged - @error "No solution for these parameters." - end + # mul!(Ktmp, P, C') + # mul!(K, Ktmp, invF) + mul!(K, P, C') + ℒ.rdiv!(K, luF) + # K = P * Ct / luF + # K = P * C' * invF + + mul!(tmp, K, C) + mul!(Ptmp, tmp, P) + ℒ.axpy!(-1, Ptmp, P) + + mul!(Ptmp, A, P) + mul!(P, Ptmp, A') + ℒ.axpy!(1, 𝐁, P) + # P = A * (P - K * C * P) * A' + 𝐁 + + mul!(u, K, z, 1, 1) + mul!(utmp, A, u) + u .= utmp + # u = A * (u + K * v) + + mul!(z, C, u) + # z = C * u + end + + return -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 +end + + + +function run_kalman_iterations(A::Matrix{S}, 𝐁::Matrix{S}, C::Matrix{Float64}, P::Matrix{S}, data_in_deviations::Matrix{S}; presample_periods::Int = 0)::S where S <: ℱ.Dual + u = zeros(S, size(C,2)) + + z = C * u + + loglik = S(0.0) + + F = similar(C * C') + + K = similar(C') + + for t in 1:size(data_in_deviations, 2) + v = data_in_deviations[:, t] - z + + F = C * P * C' + + luF = ℒ.lu(F, check = false) ### + + if !ℒ.issuccess(luF) + return -Inf + end + + Fdet = ℒ.det(luF) + + # Early return if determinant is too small, indicating numerical instability. + if Fdet < eps(Float64) + return -Inf + end + + invF = inv(luF) ### + + if t > presample_periods + loglik += log(Fdet) + ℒ.dot(v, invF, v)### + end + + K = P * C' * invF + + P = A * (P - K * C * P) * A' + 𝐁 + + u = A * (u + K * v) + + z = C * u + end + + return -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 +end + + +function rrule(::typeof(run_kalman_iterations), A, 𝐁, C, P, data_in_deviations; presample_periods = 0) + T = size(data_in_deviations, 2) + 1 + + z = zeros(size(data_in_deviations, 1)) + + ū = zeros(size(C,2)) + + P̄ = deepcopy(P) + + temp_N_N = similar(P) + + PCtmp = similar(C') + + F = similar(C * C') + + u = [similar(ū) for _ in 1:T] # used in backward pass + + P = [copy(P̄) for _ in 1:T] # used in backward pass + + CP = [zero(C) for _ in 1:T] # used in backward pass + + K = [similar(C') for _ in 1:T] # used in backward pass + + invF = [similar(F) for _ in 1:T] # used in backward pass + + v = [zeros(size(data_in_deviations, 1)) for _ in 1:T] # used in backward pass + + loglik = 0.0 + + for t in 2:T + v[t] .= data_in_deviations[:, t-1] .- z#[t-1] + + # CP[t] .= C * P̄[t-1] + mul!(CP[t], C, P̄)#[t-1]) + + # F[t] .= CP[t] * C' + mul!(F, CP[t], C') + + luF = RF.lu(F, check = false) + + if !ℒ.issuccess(luF) + return -Inf, x -> NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + Fdet = ℒ.det(luF) + + # Early return if determinant is too small, indicating numerical instability. + if Fdet < eps(Float64) + return -Inf, x -> NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # invF[t] .= inv(luF) + copy!(invF[t], inv(luF)) + + if t - 1 > presample_periods + loglik += log(Fdet) + ℒ.dot(v[t], invF[t], v[t]) + end + + # K[t] .= P̄[t-1] * C' * invF[t] + mul!(PCtmp, P̄, C') + mul!(K[t], PCtmp, invF[t]) + + # P[t] .= P̄[t-1] - K[t] * CP[t] + mul!(P[t], K[t], CP[t], -1, 0) + P[t] .+= P̄ + + # P̄[t] .= A * P[t] * A' + 𝐁 + mul!(temp_N_N, P[t], A') + mul!(P̄, A, temp_N_N) + P̄ .+= 𝐁 + + # u[t] .= K[t] * v[t] + ū[t-1] + mul!(u[t], K[t], v[t]) + u[t] .+= ū + + # ū[t] .= A * u[t] + mul!(ū, A, u[t]) + + # z[t] .= C * ū[t] + mul!(z, C, ū) + end + + llh = -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 + + # initialise derivative variables + ∂A = zero(A) + ∂F = zero(F) + ∂Faccum = zero(F) + ∂P = zero(P̄) + ∂ū = zero(ū) + ∂v = zero(v[1]) + ∂𝐁 = zero(𝐁) + ∂data_in_deviations = zero(data_in_deviations) + vtmp = zero(v[1]) + Ptmp = zero(P[1]) + + # pullback + function kalman_pullback(∂llh) + ℒ.rmul!(∂A, 0) + ℒ.rmul!(∂Faccum, 0) + ℒ.rmul!(∂P, 0) + ℒ.rmul!(∂ū, 0) + ℒ.rmul!(∂𝐁, 0) + + for t in T:-1:2 + if t > presample_periods + 1 + # ∂llh∂F + # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] + # ∂F = invF[t]' - invF[t]' * v[t] * v[t]' * invF[t]' + mul!(∂F, v[t], v[t]') + mul!(invF[1], invF[t]', ∂F) # using invF[1] as temporary storage + mul!(∂F, invF[1], invF[t]') + ℒ.axpby!(1, invF[t]', -1, ∂F) + + # ∂llh∂ū + # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] + # z[t] .= C * ū[t] + # ∂v = (invF[t]' + invF[t]) * v[t] + copy!(invF[1], invF[t]' .+ invF[t]) + # copy!(invF[1], invF[t]) # using invF[1] as temporary storage + # ℒ.axpy!(1, invF[t]', invF[1]) # using invF[1] as temporary storage + mul!(∂v, invF[1], v[t]) + # mul!(∂ū∂v, C', v[1]) + else + ℒ.rmul!(∂F, 0) + ℒ.rmul!(∂v, 0) + end + + # ∂F∂P + # F[t] .= C * P̄[t-1] * C' + # ∂P += C' * (∂F + ∂Faccum) * C + ℒ.axpy!(1, ∂Faccum, ∂F) + mul!(PCtmp, C', ∂F) + mul!(∂P, PCtmp, C, 1, 1) + + # ∂ū∂P + # K[t] .= P̄[t-1] * C' * invF[t] + # u[t] .= K[t] * v[t] + ū[t-1] + # ū[t] .= A * u[t] + # ∂P += A' * ∂ū * v[t]' * invF[t]' * C + mul!(CP[1], invF[t]', C) # using CP[1] as temporary storage + mul!(PCtmp, ∂ū , v[t]') + mul!(P[1], PCtmp , CP[1]) # using P[1] as temporary storage + mul!(∂P, A', P[1], 1, 1) + + # ∂ū∂data + # v[t] .= data_in_deviations[:, t-1] .- z + # z[t] .= C * ū[t] + # ∂data_in_deviations[:,t-1] = -C * ∂ū + mul!(u[1], A', ∂ū) + mul!(v[1], K[t]', u[1]) # using v[1] as temporary storage + ℒ.axpy!(1, ∂v, v[1]) + ∂data_in_deviations[:,t-1] .= v[1] + # mul!(∂data_in_deviations[:,t-1], C, ∂ū, -1, 0) # cannot assign to columns in matrix, must be whole matrix + + # ∂ū∂ū + # z[t] .= C * ū[t] + # v[t] .= data_in_deviations[:, t-1] .- z + # K[t] .= P̄[t-1] * C' * invF[t] + # u[t] .= K[t] * v[t] + ū[t-1] + # ū[t] .= A * u[t] + # step to next iteration + # ∂ū = A' * ∂ū - C' * K[t]' * A' * ∂ū + mul!(u[1], A', ∂ū) # using u[1] as temporary storage + mul!(v[1], K[t]', u[1]) # using v[1] as temporary storage + mul!(∂ū, C', v[1]) + mul!(u[1], C', v[1], -1, 1) + copy!(∂ū, u[1]) + + # ∂llh∂ū + # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] + # v[t] .= data_in_deviations[:, t-1] .- z + # z[t] .= C * ū[t] + # ∂ū -= ∂ū∂v + mul!(u[1], C', ∂v) # using u[1] as temporary storage + ℒ.axpy!(-1, u[1], ∂ū) + + if t > 2 + # ∂ū∂A + # ū[t] .= A * u[t] + # ∂A += ∂ū * u[t-1]' + mul!(∂A, ∂ū, u[t-1]', 1, 1) + + # ∂P̄∂A and ∂P̄∂𝐁 + # P̄[t] .= A * P[t] * A' + 𝐁 + # ∂A += ∂P * A * P[t-1]' + ∂P' * A * P[t-1] + mul!(P[1], A, P[t-1]') + mul!(Ptmp ,∂P, P[1]) + mul!(P[1], A, P[t-1]) + mul!(Ptmp ,∂P', P[1], 1, 1) + ℒ.axpy!(1, Ptmp, ∂A) + + # ∂𝐁 += ∂P + ℒ.axpy!(1, ∂P, ∂𝐁) + + # ∂P∂P + # P[t] .= P̄[t-1] - K[t] * C * P̄[t-1] + # P̄[t] .= A * P[t] * A' + 𝐁 + # step to next iteration + # ∂P = A' * ∂P * A + mul!(P[1], ∂P, A) # using P[1] as temporary storage + mul!(∂P, A', P[1]) + + # ∂P̄∂P + # K[t] .= P̄[t-1] * C' * invF[t] + # P[t] .= P̄[t-1] - K[t] * CP[t] + # ∂P -= C' * K[t-1]' * ∂P + ∂P * K[t-1] * C + mul!(PCtmp, ∂P, K[t-1]) + mul!(CP[1], K[t-1]', ∂P) # using CP[1] as temporary storage + mul!(∂P, PCtmp, C, -1, 1) + mul!(∂P, C', CP[1], -1, 1) + + # ∂ū∂F + # K[t] .= P̄[t-1] * C' * invF[t] + # u[t] .= K[t] * v[t] + ū[t-1] + # ū[t] .= A * u[t] + # ∂Faccum = -invF[t-1]' * CP[t-1] * A' * ∂ū * v[t-1]' * invF[t-1]' + mul!(u[1], A', ∂ū) # using u[1] as temporary storage + mul!(v[1], CP[t-1], u[1]) # using v[1] as temporary storage + mul!(vtmp, invF[t-1]', v[1], -1, 0) + mul!(invF[1], vtmp, v[t-1]') # using invF[1] as temporary storage + mul!(∂Faccum, invF[1], invF[t-1]') + + # ∂P∂F + # K[t] .= P̄[t-1] * C' * invF[t] + # P[t] .= P̄[t-1] - K[t] * CP[t] + # ∂Faccum -= invF[t-1]' * CP[t-1] * ∂P * CP[t-1]' * invF[t-1]' + mul!(CP[1], invF[t-1]', CP[t-1]) # using CP[1] as temporary storage + mul!(PCtmp, CP[t-1]', invF[t-1]') + mul!(K[1], ∂P, PCtmp) # using K[1] as temporary storage + mul!(∂Faccum, CP[1], K[1], -1, 1) + + end + end + + ℒ.rmul!(∂P, -∂llh/2) + ℒ.rmul!(∂A, -∂llh/2) + ℒ.rmul!(∂𝐁, -∂llh/2) + ℒ.rmul!(∂data_in_deviations, -∂llh/2) + + return NoTangent(), ∂A, ∂𝐁, NoTangent(), ∂P, ∂data_in_deviations, NoTangent() + end + + return llh, kalman_pullback +end + + + + +function check_bounds(parameter_values::Vector{S}, 𝓂::ℳ)::Bool where S <: Real + if length(𝓂.bounds) > 0 + for (k,v) in 𝓂.bounds + if k ∈ 𝓂.parameters + if min(max(parameter_values[indexin([k], 𝓂.parameters)][1], v[1]), v[2]) != parameter_values[indexin([k], 𝓂.parameters)][1] + return true + end + end + end + end + + return false +end + +function get_relevant_steady_state_and_state_update(::Val{:second_order}, parameter_values::Vector{S}, 𝓂::ℳ, tol::AbstractFloat)::Tuple{timings, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{Float64}}, Bool} where S <: Real + sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(parameter_values, 𝓂) + + all_SS = expand_steady_state(SS_and_pars,𝓂) + + state = collect(sss) - all_SS + + TT = 𝓂.timings + + return TT, SS_and_pars, [𝐒₁, 𝐒₂], [state], converged +end + + + +function get_relevant_steady_state_and_state_update(::Val{:pruned_second_order}, parameter_values::Vector{S}, 𝓂::ℳ, tol::AbstractFloat)::Tuple{timings, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{Float64}}, Bool} where S <: Real + sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(parameter_values, 𝓂, pruning = true) + + all_SS = expand_steady_state(SS_and_pars,𝓂) + + state = [zeros(𝓂.timings.nVars), collect(sss) - all_SS] + + TT = 𝓂.timings + + return TT, SS_and_pars, [𝐒₁, 𝐒₂], state, converged +end + + + +function get_relevant_steady_state_and_state_update(::Val{:third_order}, parameter_values::Vector{S}, 𝓂::ℳ, tol::AbstractFloat)::Tuple{timings, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{Float64}}, Bool} where S <: Real + sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃ = calculate_third_order_stochastic_steady_state(parameter_values, 𝓂) + + all_SS = expand_steady_state(SS_and_pars,𝓂) + + state = collect(sss) - all_SS + + TT = 𝓂.timings + + return TT, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], [state], converged +end + + + +function get_relevant_steady_state_and_state_update(::Val{:pruned_third_order}, parameter_values::Vector{S}, 𝓂::ℳ, tol::AbstractFloat)::Tuple{timings, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{Float64}}, Bool} where S <: Real + sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃ = calculate_third_order_stochastic_steady_state(parameter_values, 𝓂, pruning = true) + + all_SS = expand_steady_state(SS_and_pars,𝓂) + + state = [zeros(𝓂.timings.nVars), collect(sss) - all_SS, zeros(𝓂.timings.nVars)] + + TT = 𝓂.timings + + return TT, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged +end + + +function get_relevant_steady_state_and_state_update(::Val{:first_order}, parameter_values::Vector{S}, 𝓂::ℳ, tol::AbstractFloat)::Tuple{timings, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{Float64}}, Bool} where S <: Real + SS_and_pars, (solution_error, iters) = get_non_stochastic_steady_state(𝓂, parameter_values, tol = tol) + + state = zeros(𝓂.timings.nVars) + + TT = 𝓂.timings + + if solution_error > tol || isnan(solution_error) + return TT, SS_and_pars, zeros(S, 0, 0), [state], false + end + + ∇₁ = calculate_jacobian(parameter_values, SS_and_pars, 𝓂)# |> Matrix + + # ∇₁ = Matrix{S}(sp∇₁) + + 𝐒₁, solved = calculate_first_order_solution(∇₁; T = TT) + + return TT, SS_and_pars, 𝐒₁, [state], solved +end + + # reduce_system = false + + # if reduce_system + # variable_to_equation = @ignore_derivatives find_variables_to_exclude(𝓂, observables) + + # rows_to_exclude = Int[] + # cant_exclude = Symbol[] + + # for (ks, vidx) in variable_to_equation + # iidd = @ignore_derivatives indexin([ks] ,𝓂.timings.var)[1] + # if !isnothing(iidd) + # # if all(.!(∇₁[vidx, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd] .== 0)) + # if minimum(abs, ∇₁[vidx, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd]) / maximum(abs, ∇₁[vidx, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd]) > 1e-12 + # for v in vidx + # if v ∉ rows_to_exclude + # @ignore_derivatives push!(rows_to_exclude, v) + # # ∇₁[vidx,:] .-= ∇₁[v,:]' .* ∇₁[vidx, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd] ./ ∇₁[v, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd] + # broadcaster = @ignore_derivatives create_broadcaster(vidx, size(∇₁,1)) + # # broadcaster = spzeros(size(∇₁,1), length(vidx)) + # # for (i, vid) in enumerate(vidx) + # # broadcaster[vid,i] = 1.0 + # # end + # ∇₁ -= broadcaster * (∇₁[v,:]' .* ∇₁[vidx, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd] ./ ∇₁[v, 𝓂.timings.nFuture_not_past_and_mixed .+ iidd]) + # break + # end + # end + # else + # @ignore_derivatives push!(cant_exclude, ks) + # end + # end + # end + + # rows_to_include = @ignore_derivatives setdiff(1:𝓂.timings.nVars, rows_to_exclude) + + # cols_to_exclude = @ignore_derivatives indexin(setdiff(𝓂.timings.present_only, union(observables, cant_exclude)), 𝓂.timings.var) + + # present_idx = @ignore_derivatives 𝓂.timings.nFuture_not_past_and_mixed .+ (setdiff(range(1, 𝓂.timings.nVars), cols_to_exclude)) + + # ∇₁ = Matrix{S}(∇₁[rows_to_include, vcat(1:𝓂.timings.nFuture_not_past_and_mixed, present_idx , 𝓂.timings.nFuture_not_past_and_mixed + 𝓂.timings.nVars + 1 : size(∇₁,2))]) + + # @ignore_derivatives if !haskey(𝓂.estimation_helper, union(observables, cant_exclude)) create_timings_for_estimation!(𝓂, union(observables, cant_exclude)) end + + # TT = @ignore_derivatives 𝓂.estimation_helper[union(observables, cant_exclude)] + # else + + + + + + + + +function calculate_inversion_filter_loglikelihood(state::Vector{Vector{Float64}}, + 𝐒::Matrix{Float64}, + data_in_deviations::Matrix{Float64}, + observables::Union{Vector{String}, Vector{Symbol}}, + T::timings; + warmup_iterations::Int = 0, + presample_periods::Int = 0) + # first order + state = copy(state[1]) + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + + shocks² = 0.0 + logabsdets = 0.0 + + if warmup_iterations > 0 + if warmup_iterations >= 1 + jac = 𝐒[cond_var_idx,end-T.nExo+1:end] + if warmup_iterations >= 2 + jac = hcat(𝐒[cond_var_idx,1:T.nPast_not_future_and_mixed] * 𝐒[T.past_not_future_and_mixed_idx,end-T.nExo+1:end], jac) + if warmup_iterations >= 3 + Sᵉ = 𝐒[T.past_not_future_and_mixed_idx,1:T.nPast_not_future_and_mixed] + for e in 1:warmup_iterations-2 + jac = hcat(𝐒[cond_var_idx,1:T.nPast_not_future_and_mixed] * Sᵉ * 𝐒[T.past_not_future_and_mixed_idx,end-T.nExo+1:end], jac) + Sᵉ *= 𝐒[T.past_not_future_and_mixed_idx,1:T.nPast_not_future_and_mixed] + end + end + end + end + + jacdecomp = ℒ.svd(jac) + + x = jacdecomp \ data_in_deviations[:,1] + + warmup_shocks = reshape(x, T.nExo, warmup_iterations) + + for i in 1:warmup_iterations-1 + ℒ.mul!(state, 𝐒, vcat(state[T.past_not_future_and_mixed_idx], warmup_shocks[:,i])) + # state = state_update(state, warmup_shocks[:,i]) + end + + for i in 1:warmup_iterations + if T.nExo == length(observables) + logabsdets += ℒ.logabsdet(jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor)[1] + else + logabsdets += sum(x -> log(abs(x)), ℒ.svdvals(jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor)) + end + end + + shocks² += sum(abs2,x) + end + + y = zeros(length(cond_var_idx)) + x = zeros(T.nExo) + jac = 𝐒[cond_var_idx,end-T.nExo+1:end] + + if T.nExo == length(observables) + jacdecomp = RF.lu(jac, check = false) + if !ℒ.issuccess(jacdecomp) + return -Inf + end + logabsdets = ℒ.logabsdet(jac ./ precision_factor)[1] + invjac = inv(jacdecomp) + else + jacdecomp = ℒ.svd(jac) + + logabsdets = sum(x -> log(abs(x)), ℒ.svdvals(jac ./ precision_factor)) + invjac = inv(jacdecomp) + end + + logabsdets *= size(data_in_deviations,2) - presample_periods + + 𝐒obs = 𝐒[cond_var_idx,1:end-T.nExo] + + for i in axes(data_in_deviations,2) + @views ℒ.mul!(y, 𝐒obs, state[T.past_not_future_and_mixed_idx]) + @views ℒ.axpby!(1, data_in_deviations[:,i], -1, y) + ℒ.mul!(x, invjac, y) + + # x = invjac * (data_in_deviations[:,i] - 𝐒[cond_var_idx,1:end-T.nExo] * state[T.past_not_future_and_mixed_idx]) + + if i > presample_periods + shocks² += sum(abs2,x) + end + + ℒ.mul!(state, 𝐒, vcat(state[T.past_not_future_and_mixed_idx], x)) + # state = state_update(state, x) + end + + return -(logabsdets + shocks² + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + # return -(logabsdets + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 +end + + + +function rrule(::typeof(calculate_inversion_filter_loglikelihood), state::Vector{Vector{Float64}}, 𝐒::Matrix{Float64}, data_in_deviations::Matrix{Float64}, observables::Union{Vector{String}, Vector{Symbol}}, T::timings; warmup_iterations::Int = 0, presample_periods::Int = 0) + # first order + state = copy(state[1]) + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + obs_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + + t⁻ = T.past_not_future_and_mixed_idx + + shocks² = 0.0 + logabsdets = 0.0 + + @assert warmup_iterations == 0 "Warmup iterations not yet implemented for reverse-mode automatic differentiation." + # TODO: implement warmup iterations + + state = [copy(state) for _ in 1:size(data_in_deviations,2)+1] + + shocks² = 0.0 + logabsdets = 0.0 + + y = zeros(length(obs_idx)) + x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] + + jac = 𝐒[obs_idx,end-T.nExo+1:end] + + if T.nExo == length(observables) + logabsdets = ℒ.logabsdet(-jac' ./ precision_factor)[1] + jacdecomp = ℒ.lu(jac) + invjac = inv(jacdecomp) + else + logabsdets = sum(x -> log(abs(x)), ℒ.svdvals(-jac' ./ precision_factor)) + jacdecomp = ℒ.svd(jac) + invjac = inv(jacdecomp) + end + + logabsdets *= size(data_in_deviations,2) - presample_periods + + @views 𝐒obs = 𝐒[obs_idx,1:end-T.nExo] + + for i in axes(data_in_deviations,2) + @views ℒ.mul!(y, 𝐒obs, state[i][t⁻]) + @views ℒ.axpby!(1, data_in_deviations[:,i], -1, y) + ℒ.mul!(x[i],invjac,y) + # x = 𝐒[obs_idx,end-T.nExo+1:end] \ (data_in_deviations[:,i] - 𝐒[obs_idx,1:end-T.nExo] * state[t⁻]) + + if i > presample_periods + shocks² += sum(abs2,x[i]) + end + + ℒ.mul!(state[i+1], 𝐒, vcat(state[i][t⁻], x[i])) + # state[i+1] = 𝐒 * vcat(state[i][t⁻], x[i]) + end + + llh = -(logabsdets + shocks² + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + + if llh < -1e10 + return -Inf, x -> NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + ∂𝐒 = zero(𝐒) + + ∂𝐒ᵗ⁻ = copy(∂𝐒[t⁻,:]) + + ∂data_in_deviations = zero(data_in_deviations) + + ∂data = zeros(length(t⁻), size(data_in_deviations,2) - 1) + + ∂state = zero(state[1]) + + # precomputed matrices + M¹ = 𝐒[obs_idx, 1:end-T.nExo]' * invjac' + M² = 𝐒[t⁻,1:end-T.nExo]' - M¹ * 𝐒[t⁻,end-T.nExo+1:end]' + M³ = invjac' * 𝐒[t⁻,end-T.nExo+1:end]' + + ∂Stmp = [M¹ for _ in 1:size(data_in_deviations,2)-1] + + for t in 2:size(data_in_deviations,2)-1 + ∂Stmp[t] = M² * ∂Stmp[t-1] + end + + tmp1 = zeros(Float64, T.nExo, length(t⁻) + T.nExo) + tmp2 = zeros(Float64, length(t⁻), length(t⁻) + T.nExo) + tmp3 = zeros(Float64, length(t⁻) + T.nExo) + + ∂𝐒t⁻ = copy(tmp2) + # ∂𝐒obs_idx = copy(tmp1) + + # TODO: optimize allocations + # pullback + function inversion_pullback(∂llh) + for t in reverse(axes(data_in_deviations,2)) + ∂state[t⁻] .= M² * ∂state[t⁻] + + if t > presample_periods + ∂state[t⁻] += M¹ * x[t] + + ∂data_in_deviations[:,t] -= invjac' * x[t] + + ∂𝐒[obs_idx, :] += invjac' * x[t] * vcat(state[t][t⁻], x[t])' + + if t > 1 + ∂data[:,t:end] .= M² * ∂data[:,t:end] + + ∂data[:,t-1] += M¹ * x[t] + + ∂data_in_deviations[:,t-1] += M³ * ∂data[:,t-1:end] * ones(size(data_in_deviations,2) - t + 1) + + for tt in t-1:-1:1 + for (i,v) in enumerate(t⁻) + copyto!(tmp3::Vector{Float64}, i::Int, state[tt]::Vector{Float64}, v::Int, 1) + end + + copyto!(tmp3, length(t⁻) + 1, x[tt], 1, T.nExo) + + mul!(tmp1, x[t], tmp3') + + mul!(∂𝐒t⁻, ∂Stmp[t-tt], tmp1, 1, 1) + + end + end + end + end + + ∂𝐒[t⁻,:] += ∂𝐒t⁻ + + ∂𝐒[obs_idx, :] -= M³ * ∂𝐒t⁻ + + ∂𝐒[obs_idx,end-T.nExo+1:end] -= (size(data_in_deviations,2) - presample_periods) * invjac' / 2 + + return NoTangent(), [∂state * ∂llh], ∂𝐒 * ∂llh, ∂data_in_deviations * ∂llh, NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + return llh, inversion_pullback +end + + +function calculate_inversion_filter_loglikelihood(state::Vector{Vector{Float64}}, + 𝐒::Vector{AbstractMatrix{Float64}}, + data_in_deviations::Matrix{Float64}, + observables::Union{Vector{String}, Vector{Symbol}}, + T::timings; + warmup_iterations::Int = 0, + presample_periods::Int = 0) + if length(𝐒) == 2 && length(state) == 1 # second order + function second_order_state_update(state::Vector{U}, shock::Vector{S}) where {U <: Real,S <: Real} + # state_update = function(state::Vector{T}, shock::Vector{S}) where {T <: Real,S <: Real} + aug_state = [state[T.past_not_future_and_mixed_idx] + 1 + shock] + return 𝐒[1] * aug_state + 𝐒[2] * ℒ.kron(aug_state, aug_state) / 2 + end + + state_update = second_order_state_update + + state = state[1] + + pruning = false + elseif length(𝐒) == 2 && length(state) == 2 # pruned second order + function pruned_second_order_state_update(state::Vector{Vector{U}}, shock::Vector{S}) where {U <: Real,S <: Real} + # state_update = function(state::Vector{Vector{T}}, shock::Vector{S}) where {T <: Real,S <: Real} + aug_state₁ = [state[1][T.past_not_future_and_mixed_idx]; 1; shock] + aug_state₂ = [state[2][T.past_not_future_and_mixed_idx]; 0; zero(shock)] + + return [𝐒[1] * aug_state₁, 𝐒[1] * aug_state₂ + 𝐒[2] * ℒ.kron(aug_state₁, aug_state₁) / 2] # strictly following Andreasen et al. (2018) + end + + state_update = pruned_second_order_state_update + + pruning = true + elseif length(𝐒) == 3 && length(state) == 1 # third order + function third_order_state_update(state::Vector{U}, shock::Vector{S}) where {U <: Real,S <: Real} + # state_update = function(state::Vector{T}, shock::Vector{S}) where {T <: Real,S <: Real} + aug_state = [state[T.past_not_future_and_mixed_idx] + 1 + shock] + return 𝐒[1] * aug_state + 𝐒[2] * ℒ.kron(aug_state, aug_state) / 2 + 𝐒[3] * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 + end + + state_update = third_order_state_update + + state = state[1] + + pruning = false + elseif length(𝐒) == 3 && length(state) == 3 # pruned third order + function pruned_third_order_state_update(state::Vector{Vector{U}}, shock::Vector{S}) where {U <: Real,S <: Real} + # state_update = function(state::Vector{Vector{T}}, shock::Vector{S}) where {T <: Real,S <: Real} + aug_state₁ = [state[1][T.past_not_future_and_mixed_idx]; 1; shock] + aug_state₁̂ = [state[1][T.past_not_future_and_mixed_idx]; 0; shock] + aug_state₂ = [state[2][T.past_not_future_and_mixed_idx]; 0; zero(shock)] + aug_state₃ = [state[3][T.past_not_future_and_mixed_idx]; 0; zero(shock)] + + kron_aug_state₁ = ℒ.kron(aug_state₁, aug_state₁) + + return [𝐒[1] * aug_state₁, 𝐒[1] * aug_state₂ + 𝐒[2] * kron_aug_state₁ / 2, 𝐒[1] * aug_state₃ + 𝐒[2] * ℒ.kron(aug_state₁̂, aug_state₂) + 𝐒[3] * ℒ.kron(kron_aug_state₁,aug_state₁) / 6] + end + + state_update = pruned_third_order_state_update + + pruning = true + end + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + + shocks² = 0.0 + logabsdets = 0.0 + + if warmup_iterations > 0 + res = Optim.optimize(x -> minimize_distance_to_initial_data(x, data_in_deviations[:,1], state, state_update, warmup_iterations, cond_var_idx, precision_factor, pruning), + zeros(T.nExo * warmup_iterations), + Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), + Optim.Options(f_abstol = eps(), g_tol= 1e-30); + autodiff = :forward) + + matched = Optim.minimum(res) < 1e-12 + + if !matched # for robustness try other linesearch + res = Optim.optimize(x -> minimize_distance_to_initial_data(x, data_in_deviations[:,1], state, state_update, warmup_iterations, cond_var_idx, precision_factor, pruning), + zeros(T.nExo * warmup_iterations), + Optim.LBFGS(), + Optim.Options(f_abstol = eps(), g_tol= 1e-30); + autodiff = :forward) + + matched = Optim.minimum(res) < 1e-12 + end + + if !matched return -Inf end + + x = Optim.minimizer(res) + + warmup_shocks = reshape(x, T.nExo, warmup_iterations) + + for i in 1:warmup_iterations-1 + state = state_update(state, warmup_shocks[:,i]) + end + + res = zeros(0) + + jacc = zeros(T.nExo * warmup_iterations, length(observables)) + + match_initial_data!(res, x, jacc, data_in_deviations[:,1], state, state_update, warmup_iterations, cond_var_idx, precision_factor), zeros(size(data_in_deviations, 1)) + + for i in 1:warmup_iterations + if T.nExo == length(observables) + logabsdets += ℒ.logabsdet(jacc[(i - 1) * T.nExo+1:i*T.nExo,:] ./ precision_factor)[1] + else + logabsdets += sum(x -> log(abs(x)), ℒ.svdvals(jacc[(i - 1) * T.nExo+1:i*T.nExo,:] ./ precision_factor)) + end + end + + shocks² += sum(abs2,x) + end + + for i in axes(data_in_deviations,2) + res = Optim.optimize(x -> minimize_distance_to_data(x, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor, pruning), + zeros(T.nExo), + Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), + Optim.Options(f_abstol = eps(), g_tol= 1e-30); + autodiff = :forward) + + matched = Optim.minimum(res) < 1e-12 + + if !matched # for robustness try other linesearch + res = Optim.optimize(x -> minimize_distance_to_data(x, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor, pruning), + zeros(T.nExo), + Optim.LBFGS(), + Optim.Options(f_abstol = eps(), g_tol= 1e-30); + autodiff = :forward) + + matched = Optim.minimum(res) < 1e-12 + end + + if !matched return -Inf end + + x = Optim.minimizer(res) + + res = zeros(0) + + jacc = zeros(T.nExo, length(observables)) + + match_data_sequence!(res, x, jacc, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor) + + if i > presample_periods + # due to change of variables: jacobian determinant adjustment + if T.nExo == length(observables) + logabsdets += ℒ.logabsdet(jacc ./ precision_factor)[1] + else + logabsdets += sum(x -> log(abs(x)), ℒ.svdvals(jacc ./ precision_factor)) + end + + shocks² += sum(abs2,x) + end + + state = state_update(state, x) + end + + # See: https://pcubaborda.net/documents/CGIZ-final.pdf + return -(logabsdets + shocks² + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 +end + + + +function inversion_filter(𝓂::ℳ, + data_in_deviations::AbstractArray{Float64}, + algorithm::Symbol; + warmup_iterations::Int = 0, + verbose::Bool = false, + tol::AbstractFloat = 1e-12) + + observables = collect(axiskeys(data_in_deviations,1)) + + data_in_deviations = collect(data_in_deviations) + + @assert observables isa Vector{String} || observables isa Vector{Symbol} "Make sure that the data has variables names as rows. They can be either Strings or Symbols." + + sort!(observables) + + observables = observables isa String_input ? observables .|> Meta.parse .|> replace_indices : observables + + # solve model given the parameters + if algorithm == :second_order + sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(𝓂.parameter_values, 𝓂) + + if !converged + @error "No solution for these parameters." + end + + all_SS = expand_steady_state(SS_and_pars,𝓂) + + state = collect(sss) - all_SS + + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[𝓂.timings.past_not_future_and_mixed_idx] + 1 + shock] + return 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + end + elseif algorithm == :pruned_second_order + sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(𝓂.parameter_values, 𝓂, pruning = true) + + if !converged + @error "No solution for these parameters." + end all_SS = expand_steady_state(SS_and_pars,𝓂) @@ -7757,7 +9638,7 @@ function inversion_filter(𝓂::ℳ, state = zeros(𝓂.timings.nVars) - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix 𝐒₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -7817,6 +9698,8 @@ function inversion_filter(𝓂::ℳ, end end + initial_state = state + for i in axes(data_in_deviations,2) res = @suppress begin Optim.optimize(x -> minimize_distance_to_data(x, data_in_deviations[:,i], state, state_update, cond_var_idx, precision_factor, pruning), zeros(𝓂.timings.nExo), @@ -7846,7 +9729,197 @@ function inversion_filter(𝓂::ℳ, shocks[:,i] = x end - return states, shocks + return states, shocks, initial_state +end + + +function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:first_order}, # algo + ::Val{:kalman}; # filter + warmup_iterations::Int = 0, + smooth::Bool = true, + verbose::Bool = false) + + obs_axis = collect(axiskeys(data_in_deviations,1)) + + obs_symbols = obs_axis isa String_input ? obs_axis .|> Meta.parse .|> replace_indices : obs_axis + + filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) + + variables = filtered_and_smoothed[smooth ? 1 : 5] + standard_deviations = filtered_and_smoothed[smooth ? 2 : 6] + shocks = filtered_and_smoothed[smooth ? 3 : 7] + decomposition = filtered_and_smoothed[smooth ? 4 : 8] + + return variables, shocks, standard_deviations, decomposition +end + + +function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:first_order}, # algo + ::Val{:inversion}; # filter + warmup_iterations::Int = 0, + smooth::Bool = true, + verbose::Bool = false) + + algorithm = :first_order + + variables, shocks, initial_state = inversion_filter(𝓂, data_in_deviations, algorithm, warmup_iterations = warmup_iterations) + + state_update, pruning = parse_algorithm_to_state_update(algorithm, 𝓂, false) + + decomposition = zeros(𝓂.timings.nVars, 𝓂.timings.nExo + 2, size(data_in_deviations, 2)) + + decomposition[:,end,:] .= variables + + for i in 1:𝓂.timings.nExo + sck = zeros(𝓂.timings.nExo) + sck[i] = shocks[i, 1] + decomposition[:,i,1] = state_update(initial_state , sck) + end + + decomposition[:,end - 1,1] .= decomposition[:,end,1] - sum(decomposition[:,1:end-2,1], dims=2) + + for i in 2:size(data_in_deviations,2) + for ii in 1:𝓂.timings.nExo + sck = zeros(𝓂.timings.nExo) + sck[ii] = shocks[ii, i] + decomposition[:,ii,i] = state_update(decomposition[:,ii, i-1], sck) + end + + decomposition[:,end - 1,i] .= decomposition[:,end,i] - sum(decomposition[:,1:end-2,i], dims=2) + end + + return variables, shocks, [], decomposition +end + + +function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:second_order}, # algo + ::Val{:inversion}; # filter + warmup_iterations::Int = 0, + smooth::Bool = true, + verbose::Bool = false) + + variables, shocks, initial_state = inversion_filter(𝓂, data_in_deviations, :second_order, warmup_iterations = warmup_iterations) + + return variables, shocks, [], [] +end + + +function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:pruned_second_order}, # algo + ::Val{:inversion}; # filter + warmup_iterations::Int = 0, + smooth::Bool = true, + verbose::Bool = false) + + algorithm = :pruned_second_order + + variables, shocks, initial_state = inversion_filter(𝓂, data_in_deviations, algorithm, warmup_iterations = warmup_iterations) + + state_update, pruning = parse_algorithm_to_state_update(algorithm, 𝓂, false) + + states = [initial_state for _ in 1:𝓂.timings.nExo + 1] + + decomposition = zeros(𝓂.timings.nVars, 𝓂.timings.nExo + 3, size(data_in_deviations, 2)) + + decomposition[:, end, :] .= variables + + for i in 1:𝓂.timings.nExo + sck = zeros(𝓂.timings.nExo) + sck[i] = shocks[i, 1] + states[i] = state_update(initial_state , sck) + decomposition[:,i,1] = sum(states[i]) + end + + states[end] = state_update(initial_state, shocks[:, 1]) + + decomposition[:, end - 2, 1] = sum(states[end]) - sum(decomposition[:, 1:end - 3, 1], dims = 2) + decomposition[:, end - 1, 1] .= decomposition[:, end, 1] - sum(decomposition[:, 1:end - 2, 1], dims = 2) + + for i in 2:size(data_in_deviations, 2) + for ii in 1:𝓂.timings.nExo + sck = zeros(𝓂.timings.nExo) + sck[ii] = shocks[ii, i] + states[ii] = state_update(states[ii] , sck) + decomposition[:, ii, i] = sum(states[ii]) + end + + states[end] = state_update(states[end] , shocks[:, i]) + + decomposition[:, end - 2, i] = sum(states[end]) - sum(decomposition[:, 1:end - 3, i], dims = 2) + decomposition[:, end - 1, i] .= decomposition[:, end, i] - sum(decomposition[:, 1:end - 2, i], dims = 2) + end + + return variables, shocks, [], decomposition +end + +function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:third_order}, # algo + ::Val{:inversion}; # filter + warmup_iterations::Int = 0, + smooth::Bool = true, + verbose::Bool = false) + + variables, shocks, initial_state = inversion_filter(𝓂, data_in_deviations, :third_order, warmup_iterations = warmup_iterations) + + return variables, shocks, [], [] +end + + +function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:pruned_third_order}, # algo + ::Val{:inversion}; # filter + warmup_iterations::Int = 0, + smooth::Bool = true, + verbose::Bool = false) + + algorithm = :pruned_third_order + + variables, shocks, initial_state = inversion_filter(𝓂, data_in_deviations, algorithm, warmup_iterations = warmup_iterations) + + state_update, pruning = parse_algorithm_to_state_update(algorithm, 𝓂, false) + + states = [initial_state for _ in 1:𝓂.timings.nExo + 1] + + decomposition = zeros(𝓂.timings.nVars, 𝓂.timings.nExo + 3, size(data_in_deviations, 2)) + + decomposition[:, end, :] .= variables + + for i in 1:𝓂.timings.nExo + sck = zeros(𝓂.timings.nExo) + sck[i] = shocks[i, 1] + states[i] = state_update(initial_state , sck) + decomposition[:,i,1] = sum(states[i]) + end + + states[end] = state_update(initial_state, shocks[:, 1]) + + decomposition[:,end - 2, 1] = sum(states[end]) - sum(decomposition[:,1:end - 3, 1], dims = 2) + decomposition[:,end - 1, 1] .= decomposition[:, end, 1] - sum(decomposition[:,1:end - 2, 1], dims = 2) + + for i in 2:size(data_in_deviations, 2) + for ii in 1:𝓂.timings.nExo + sck = zeros(𝓂.timings.nExo) + sck[ii] = shocks[ii, i] + states[ii] = state_update(states[ii] , sck) + decomposition[:, ii, i] = sum(states[ii]) + end + + states[end] = state_update(states[end] , shocks[:, i]) + + decomposition[:,end - 2, i] = sum(states[end]) - sum(decomposition[:,1:end - 3, i], dims = 2) + decomposition[:,end - 1, i] .= decomposition[:, end, i] - sum(decomposition[:,1:end - 2, i], dims = 2) + end + + return variables, shocks, [], decomposition end @@ -7871,7 +9944,7 @@ function filter_and_smooth(𝓂::ℳ, @assert solution_error < tol "Could not solve non stochastic steady state." - ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂)# |> Matrix sol, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -7964,80 +10037,80 @@ function filter_and_smooth(𝓂::ℳ, end -if VERSION >= v"1.9" - @setup_workload begin - # Putting some things in `setup` can reduce the size of the - # precompile file and potentially make loading faster. - @model FS2000 precompile = true begin - dA[0] = exp(gam + z_e_a * e_a[x]) - log(m[0]) = (1 - rho) * log(mst) + rho * log(m[-1]) + z_e_m * e_m[x] - - P[0] / (c[1] * P[1] * m[0]) + bet * P[1] * (alp * exp( - alp * (gam + log(e[1]))) * k[0] ^ (alp - 1) * n[1] ^ (1 - alp) + (1 - del) * exp( - (gam + log(e[1])))) / (c[2] * P[2] * m[1])=0 - W[0] = l[0] / n[0] - - (psi / (1 - psi)) * (c[0] * P[0] / (1 - n[0])) + l[0] / n[0] = 0 - R[0] = P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ ( - alp) / W[0] - 1 / (c[0] * P[0]) - bet * P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) / (m[0] * l[0] * c[1] * P[1]) = 0 - c[0] + k[0] = exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) + (1 - del) * exp( - (gam + z_e_a * e_a[x])) * k[-1] - P[0] * c[0] = m[0] - m[0] - 1 + d[0] = l[0] - e[0] = exp(z_e_a * e_a[x]) - y[0] = k[-1] ^ alp * n[0] ^ (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) - gy_obs[0] = dA[0] * y[0] / y[-1] - gp_obs[0] = (P[0] / P[-1]) * m[-1] / dA[0] - log_gy_obs[0] = log(gy_obs[0]) - log_gp_obs[0] = log(gp_obs[0]) - end - - @parameters FS2000 silent = true precompile = true begin - alp = 0.356 - bet = 0.993 - gam = 0.0085 - mst = 1.0002 - rho = 0.129 - psi = 0.65 - del = 0.01 - z_e_a = 0.035449 - z_e_m = 0.008862 - end - - ENV["GKSwstype"] = "nul" - - @compile_workload begin - # all calls in this block will be precompiled, regardless of whether - # they belong to your package or not (on Julia 1.8 and higher) - @model RBC precompile = true begin - 1 / c[0] = (0.95 / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ)) - c[0] + k[0] = (1 - δ) * k[-1] + exp(z[0]) * k[-1]^α - z[0] = 0.2 * z[-1] + 0.01 * eps_z[x] - end - - @parameters RBC silent = true precompile = true begin - δ = 0.02 - α = 0.5 - end +# if VERSION >= v"1.9" +# @setup_workload begin +# # Putting some things in `setup` can reduce the size of the +# # precompile file and potentially make loading faster. +# @model FS2000 precompile = true begin +# dA[0] = exp(gam + z_e_a * e_a[x]) +# log(m[0]) = (1 - rho) * log(mst) + rho * log(m[-1]) + z_e_m * e_m[x] +# - P[0] / (c[1] * P[1] * m[0]) + bet * P[1] * (alp * exp( - alp * (gam + log(e[1]))) * k[0] ^ (alp - 1) * n[1] ^ (1 - alp) + (1 - del) * exp( - (gam + log(e[1])))) / (c[2] * P[2] * m[1])=0 +# W[0] = l[0] / n[0] +# - (psi / (1 - psi)) * (c[0] * P[0] / (1 - n[0])) + l[0] / n[0] = 0 +# R[0] = P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ ( - alp) / W[0] +# 1 / (c[0] * P[0]) - bet * P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) / (m[0] * l[0] * c[1] * P[1]) = 0 +# c[0] + k[0] = exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) + (1 - del) * exp( - (gam + z_e_a * e_a[x])) * k[-1] +# P[0] * c[0] = m[0] +# m[0] - 1 + d[0] = l[0] +# e[0] = exp(z_e_a * e_a[x]) +# y[0] = k[-1] ^ alp * n[0] ^ (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) +# gy_obs[0] = dA[0] * y[0] / y[-1] +# gp_obs[0] = (P[0] / P[-1]) * m[-1] / dA[0] +# log_gy_obs[0] = log(gy_obs[0]) +# log_gp_obs[0] = log(gp_obs[0]) +# end - get_SS(FS2000) - get_SS(FS2000, parameters = :alp => 0.36) - get_solution(FS2000) - get_solution(FS2000, parameters = :alp => 0.35) - get_standard_deviation(FS2000) - get_correlation(FS2000) - get_autocorrelation(FS2000) - get_variance_decomposition(FS2000) - get_conditional_variance_decomposition(FS2000) - get_irf(FS2000) +# @parameters FS2000 silent = true precompile = true begin +# alp = 0.356 +# bet = 0.993 +# gam = 0.0085 +# mst = 1.0002 +# rho = 0.129 +# psi = 0.65 +# del = 0.01 +# z_e_a = 0.035449 +# z_e_m = 0.008862 +# end + +# ENV["GKSwstype"] = "nul" + +# @compile_workload begin +# # all calls in this block will be precompiled, regardless of whether +# # they belong to your package or not (on Julia 1.8 and higher) +# @model RBC precompile = true begin +# 1 / c[0] = (0.95 / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ)) +# c[0] + k[0] = (1 - δ) * k[-1] + exp(z[0]) * k[-1]^α +# z[0] = 0.2 * z[-1] + 0.01 * eps_z[x] +# end - data = simulate(FS2000)([:c,:k],:,:simulate) - get_loglikelihood(FS2000, data, FS2000.parameter_values) - get_mean(FS2000, silent = true) - # get_SSS(FS2000, silent = true) - # get_SSS(FS2000, algorithm = :third_order, silent = true) +# @parameters RBC silent = true precompile = true begin +# δ = 0.02 +# α = 0.5 +# end - # import StatsPlots - # plot_irf(FS2000) - # plot_solution(FS2000,:k) # fix warning when there is no sensitivity and all values are the same. triggers: no strict ticks found... - # plot_conditional_variance_decomposition(FS2000) - end - end -end +# get_SS(FS2000, silent = true) +# get_SS(FS2000, parameters = :alp => 0.36, silent = true) +# get_solution(FS2000, silent = true) +# get_solution(FS2000, parameters = :alp => 0.35) +# get_standard_deviation(FS2000) +# get_correlation(FS2000) +# get_autocorrelation(FS2000) +# get_variance_decomposition(FS2000) +# get_conditional_variance_decomposition(FS2000) +# get_irf(FS2000) + +# data = simulate(FS2000)([:c,:k],:,:simulate) +# get_loglikelihood(FS2000, data, FS2000.parameter_values) +# get_mean(FS2000, silent = true) +# # get_SSS(FS2000, silent = true) +# # get_SSS(FS2000, algorithm = :third_order, silent = true) + +# # import StatsPlots +# # plot_irf(FS2000) +# # plot_solution(FS2000,:k) # fix warning when there is no sensitivity and all values are the same. triggers: no strict ticks found... +# # plot_conditional_variance_decomposition(FS2000) +# end +# end +# end end diff --git a/src/common_docstrings.jl b/src/common_docstrings.jl index e8b0e527..fb7fc231 100644 --- a/src/common_docstrings.jl +++ b/src/common_docstrings.jl @@ -17,5 +17,5 @@ const CONDITIONS = "`conditions` [Type: `Union{Matrix{Union{Nothing,Float64}}, S const SHOCK_CONDITIONS = "`shocks` [Default: `nothing`, Type: `Union{Matrix{Union{Nothing,Float64}}, SparseMatrixCSC{Float64}, KeyedArray{Union{Nothing,Float64}}, KeyedArray{Float64}, Nothing} = nothing`]: known values of shocks. This entry allows the user to include certain shock values. By entering restrictions on the shock sin this way the problem to match the conditions on endogenous variables is restricted to the remaining free shocks in the repective period. The input can have multiple formats, but for all types of entries the first dimension corresponds to the number of shocks and the second dimension to the number of periods. The shocks can be specified using a matrix of type `Matrix{Union{Nothing,Float64}}`. In this case the shocks are matrix elements of type `Float64` and all remaining (free) entries are `nothing`. You can also use a `SparseMatrixCSC{Float64}` as input. In this case only non-zero elements are taken as certain shock values. Note that you cannot condition shocks to be zero using a `SparseMatrixCSC{Float64}` as input (use other input formats to do so). Another possibility to input known shocks is by using a `KeyedArray`. You can use a `KeyedArray{Union{Nothing,Float64}}` where, similar to `Matrix{Union{Nothing,Float64}}`, all entries of type `Float64` are recognised as known shocks and all other entries have to be `nothing`. Furthermore, you can specify in the primary axis a subset of shocks (of type `Symbol` or `String`) for which you specify values and all other shocks are considered free. The same goes for the case when you use `KeyedArray{Float64}}` as input, whereas in this case the values for the specified shocks bind for all periods specified in the `KeyedArray`, because there are no `nothing` entries permitted with this type." const PARAMETER_DERIVATIVES = "`parameter_derivatives` [Default: :all]: parameters for which to calculate partial derivatives. Inputs can be a parameter name passed on as either a `Symbol` or `String` (e.g. `:alpha`, or \"alpha\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. `:all` will include all parameters." const DATA = "`data` [Type: `KeyedArray`]: data matrix with variables (`String` or `Symbol`) in rows and time in columns" -const SMOOTH = "`smooth` [Default: `true`, Type: `Bool`]: whether to return smoothed (`true`) or filtered (`false`) shocks. Only works for the Kalman filter. The inversion filter only returns filtered shocks." +const SMOOTH = "`smooth` [Default: `true`, Type: `Bool`]: whether to return smoothed (`true`) or filtered (`false`) shocks/variables. Only works for the Kalman filter. The inversion filter only returns filtered shocks/variables." const DATA_IN_LEVELS = "`data_in_levels` [Default: `true`, Type: `Bool`]: indicator whether the data is provided in levels. If `true` the input to the data argument will have the non stochastic steady state substracted." diff --git a/src/get_functions.jl b/src/get_functions.jl index 7ca0af3b..d9434662 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -1,6 +1,6 @@ """ $(SIGNATURES) -Return the shock decomposition in absolute deviations from the non stochastic steady state based on the Kalman smoother or filter (depending on the `smooth` keyword argument) using the provided data and first order solution of the model. Data is by default assumed to be in levels unless `data_in_levels` is set to `false`. +Return the shock decomposition in absolute deviations from the relevant steady state (e.g. higher order perturbation algorithms are relative to the stochastic steady state) based on the Kalman smoother or filter (depending on the `smooth` keyword argument) or inversion filter using the provided data and solution of the model. Data is by default assumed to be in levels unless `data_in_levels` is set to `false`. # Arguments - $MODEL @@ -65,13 +65,25 @@ And data, 4×2×40 Array{Float64, 3}: function get_shock_decomposition(𝓂::ℳ, data::KeyedArray{Float64}; parameters::ParameterType = nothing, + filter::Symbol = :kalman, + algorithm::Symbol = :first_order, data_in_levels::Bool = true, + warmup_iterations::Int = 0, smooth::Bool = true, verbose::Bool = false) - solve!(𝓂, parameters = parameters, verbose = verbose, dynamics = true) + pruning = false + + @assert !(algorithm ∈ [:second_order, :third_order]) "Decomposition implemented for first order, pruned second and third order. Second and third order solution decomposition is not yet implemented." + + if algorithm ∈ [:pruned_second_order, :pruned_third_order] + filter = :inversion + pruning = true + end - reference_steady_state, (solution_error, iters) = 𝓂.solution.outdated_NSSS ? 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) : (copy(𝓂.solution.non_stochastic_steady_state), (eps(), 0)) + solve!(𝓂, parameters = parameters, verbose = verbose, dynamics = true, algorithm = algorithm) + + reference_steady_state, NSSS, SSS_delta = get_relevant_steady_states(𝓂, algorithm) data = data(sort(axiskeys(data,1))) @@ -82,13 +94,13 @@ function get_shock_decomposition(𝓂::ℳ, obs_idx = parse_variables_input_to_index(obs_symbols, 𝓂.timings) if data_in_levels - data_in_deviations = data .- reference_steady_state[obs_idx] + data_in_deviations = data .- NSSS[obs_idx] else data_in_deviations = data end - filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) - + variables, shocks, standard_deviations, decomposition = filter_data_with_model(𝓂, data_in_deviations, Val(algorithm), Val(filter), warmup_iterations = warmup_iterations, smooth = smooth, verbose = verbose) + axis1 = 𝓂.timings.var if any(x -> contains(string(x), "◖"), axis1) @@ -96,17 +108,30 @@ function get_shock_decomposition(𝓂::ℳ, axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] end - axis2 = vcat(𝓂.timings.exo, :Initial_values) + if pruning + axis2 = vcat(𝓂.timings.exo, :Nonlinearities, :Initial_values) + else + axis2 = vcat(𝓂.timings.exo, :Initial_values) + end if any(x -> contains(string(x), "◖"), axis2) axis2_decomposed = decompose_name.(axis2) axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] axis2[1:length(𝓂.timings.exo)] = axis2[1:length(𝓂.timings.exo)] .* "₍ₓ₎" else - axis2 = vcat(map(x->Symbol(string(x) * "₍ₓ₎"), 𝓂.timings.exo), :Initial_values) + if pruning + axis2 = vcat(map(x->Symbol(string(x) * "₍ₓ₎"), 𝓂.timings.exo), :Nonlinearities, :Initial_values) + else + axis2 = vcat(map(x->Symbol(string(x) * "₍ₓ₎"), 𝓂.timings.exo), :Initial_values) + end + end + + if pruning + decomposition[:,1:(end - 2 - pruning),:] .+= SSS_delta + decomposition[:,end - 2,:] .-= SSS_delta * (size(decomposition,2) - 4) end - return KeyedArray(filtered_and_smoothed[smooth ? 4 : 8][:,1:end-1,:]; Variables = axis1, Shocks = axis2, Periods = 1:size(data,2)) + return KeyedArray(decomposition[:,1:end-1,:]; Variables = axis1, Shocks = axis2, Periods = 1:size(data,2)) end @@ -177,7 +202,7 @@ function get_estimated_shocks(𝓂::ℳ, solve!(𝓂, parameters = parameters, algorithm = algorithm, verbose = verbose, dynamics = true) - reference_steady_state, (solution_error, iters) = 𝓂.solution.outdated_NSSS ? 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) : (copy(𝓂.solution.non_stochastic_steady_state), (eps(), 0)) + reference_steady_state, NSSS, SSS_delta = get_relevant_steady_states(𝓂, algorithm) data = data(sort(axiskeys(data,1))) @@ -188,18 +213,13 @@ function get_estimated_shocks(𝓂::ℳ, obs_idx = parse_variables_input_to_index(obs_symbols, 𝓂.timings) if data_in_levels - data_in_deviations = data .- reference_steady_state[obs_idx] + data_in_deviations = data .- NSSS[obs_idx] else data_in_deviations = data end - if filter == :kalman - filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) - shocks = filtered_and_smoothed[smooth ? 3 : 7] - elseif filter == :inversion - states, shocks = inversion_filter(𝓂, data_in_deviations, algorithm, warmup_iterations = warmup_iterations) - end - + variables, shocks, standard_deviations, decomposition = filter_data_with_model(𝓂, data_in_deviations, Val(algorithm), Val(filter), warmup_iterations = warmup_iterations, smooth = smooth, verbose = verbose) + axis1 = 𝓂.timings.exo if any(x -> contains(string(x), "◖"), axis1) @@ -288,7 +308,7 @@ function get_estimated_variables(𝓂::ℳ, solve!(𝓂, parameters = parameters, algorithm = algorithm, verbose = verbose, dynamics = true) - reference_steady_state, (solution_error, iters) = 𝓂.solution.outdated_NSSS ? 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) : (copy(𝓂.solution.non_stochastic_steady_state), (eps(), 0)) + reference_steady_state, NSSS, SSS_delta = get_relevant_steady_states(𝓂, algorithm) data = data(sort(axiskeys(data,1))) @@ -299,17 +319,12 @@ function get_estimated_variables(𝓂::ℳ, obs_idx = parse_variables_input_to_index(obs_symbols, 𝓂.timings) if data_in_levels - data_in_deviations = data .- reference_steady_state[obs_idx] + data_in_deviations = data .- NSSS[obs_idx] else data_in_deviations = data end - if filter == :kalman - filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) - states = filtered_and_smoothed[smooth ? 1 : 5] - elseif filter == :inversion - states, shocks = inversion_filter(𝓂, data_in_deviations, algorithm, warmup_iterations = warmup_iterations) - end + variables, shocks, standard_deviations, decomposition = filter_data_with_model(𝓂, data_in_deviations, Val(algorithm), Val(filter), warmup_iterations = warmup_iterations, smooth = smooth, verbose = verbose) axis1 = 𝓂.timings.var @@ -318,7 +333,7 @@ function get_estimated_variables(𝓂::ℳ, axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] end - return KeyedArray(levels ? states .+ reference_steady_state[1:length(𝓂.var)] : states; Variables = axis1, Periods = 1:size(data,2)) + return KeyedArray(levels ? variables .+ NSSS[1:length(𝓂.var)] : variables; Variables = axis1, Periods = 1:size(data,2)) end @@ -379,9 +394,11 @@ function get_estimated_variable_standard_deviations(𝓂::ℳ, smooth::Bool = true, verbose::Bool = false) + algorithm = :first_order + solve!(𝓂, parameters = parameters, verbose = verbose, dynamics = true) - reference_steady_state, (solution_error, iters) = 𝓂.solution.outdated_NSSS ? 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) : (copy(𝓂.solution.non_stochastic_steady_state), (eps(), 0)) + reference_steady_state, NSSS, SSS_delta = get_relevant_steady_states(𝓂, algorithm) data = data(sort(axiskeys(data,1))) @@ -392,12 +409,12 @@ function get_estimated_variable_standard_deviations(𝓂::ℳ, obs_idx = parse_variables_input_to_index(obs_symbols, 𝓂.timings) if data_in_levels - data_in_deviations = data .- reference_steady_state[obs_idx] + data_in_deviations = data .- NSSS[obs_idx] else data_in_deviations = data end - filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) + variables, shocks, standard_deviations, decomposition = filter_data_with_model(𝓂, data_in_deviations, Val(:first_order), Val(:kalman), smooth = smooth, verbose = verbose) axis1 = 𝓂.timings.var @@ -406,7 +423,7 @@ function get_estimated_variable_standard_deviations(𝓂::ℳ, axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] end - return KeyedArray(filtered_and_smoothed[smooth ? 2 : 6]; Standard_deviations = axis1, Periods = 1:size(data,2)) + return KeyedArray(standard_deviations; Standard_deviations = axis1, Periods = 1:size(data,2)) end @@ -861,7 +878,7 @@ function get_irf(𝓂::ℳ, reference_steady_state, (solution_error, iters) = 𝓂.SS_solve_func(parameters, 𝓂, verbose, false, 𝓂.solver_parameters) - ∇₁ = calculate_jacobian(parameters, reference_steady_state, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(parameters, reference_steady_state, 𝓂)# |> Matrix sol_mat, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -1275,7 +1292,7 @@ function get_steady_state(𝓂::ℳ; parameter_derivatives::Union{Symbol_input,String_input} = :all, return_variables_only::Bool = false, verbose::Bool = false, - silent::Bool = true, + silent::Bool = false, tol::AbstractFloat = 1e-12) if !(algorithm == :first_order) stochastic = true end @@ -1528,9 +1545,10 @@ And data, 4×4 adjoint(::Matrix{Float64}) with eltype Float64: function get_solution(𝓂::ℳ; parameters::ParameterType = nothing, algorithm::Symbol = :first_order, + silent::Bool = false, verbose::Bool = false) - solve!(𝓂, parameters = parameters, verbose = verbose, dynamics = true, algorithm = algorithm) + solve!(𝓂, parameters = parameters, verbose = verbose, dynamics = true, silent = silent, algorithm = algorithm) if algorithm == :linear_time_iteration solution_matrix = 𝓂.solution.perturbation.linear_time_iteration.solution_matrix @@ -1668,7 +1686,7 @@ function get_solution(𝓂::ℳ, end end - ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂)# |> Matrix 𝐒₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -1796,7 +1814,7 @@ function get_conditional_variance_decomposition(𝓂::ℳ; SS_and_pars, _ = 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix 𝑺₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -1938,7 +1956,7 @@ function get_variance_decomposition(𝓂::ℳ; SS_and_pars, (solution_error, iters) = 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂) |> Matrix + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂)# |> Matrix sol, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) @@ -2263,7 +2281,7 @@ function get_moments(𝓂::ℳ; algorithm::Symbol = :first_order, dependencies_tol::AbstractFloat = 1e-12, verbose::Bool = false, - silent::Bool = true)#limit output by selecting pars and vars like for plots and irfs!? + silent::Bool = false)#limit output by selecting pars and vars like for plots and irfs!? solve!(𝓂, parameters = parameters, algorithm = algorithm, verbose = verbose, silent = silent) @@ -2847,6 +2865,8 @@ This function is differentiable (so far for the Kalman filter only) and can be u - $ALGORITHM - $FILTER - `warmup_iterations` [Default: `0`, Type: `Int`]: periods added before the first observation for which shocks are computed such that the first observation is matched. A larger value alleviates the problem that the initial value is the relevant steady state. +- `presample_periods` [Default: `0`, Type: `Int`]: periods at the beginning of the data for which the loglikelihood is discarded. +- `initial_covariance` [Default: `:theoretical`, Type: `Symbol`]: defines the method to initialise the Kalman filters covariance matrix. It can be initialised with the theoretical long run values (option `:theoretical`) or large values (10.0) along the diagonal (option `:diagonal`). - $VERBOSE # Examples @@ -2881,11 +2901,16 @@ function get_loglikelihood(𝓂::ℳ, algorithm::Symbol = :first_order, filter::Symbol = :kalman, warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, tol::AbstractFloat = 1e-12, - verbose::Bool = false)::S where S + verbose::Bool = false)::S where S <: Real # checks to avoid errors further down the line and inform the user - @assert filter ∈ [:kalman, :inversion] "Currently only the kalman filter (:kalman) for linear models and the inversion filter (:inversion) for linear and nonlinear models are supported." + @assert filter ∈ [:kalman, :inversion] "Currently only the Kalman filter (:kalman) for linear models and the inversion filter (:inversion) for linear and nonlinear models are supported." + + # checks to avoid errors further down the line and inform the user + @assert initial_covariance ∈ [:theoretical, :diagonal] "Invalid method to initialise the Kalman filters covariance matrix. Supported methods are: the theoretical long run values (option `:theoretical`) or large values (10.0) along the diagonal (option `:diagonal`)." if algorithm ∈ [:second_order,:pruned_second_order,:third_order,:pruned_third_order] filter = :inversion @@ -2895,117 +2920,28 @@ function get_loglikelihood(𝓂::ℳ, @ignore_derivatives solve!(𝓂, verbose = verbose, algorithm = algorithm) - # keep the parameters within bounds - for (k,v) in 𝓂.bounds - if k ∈ 𝓂.parameters - if @ignore_derivatives min(max(parameter_values[indexin([k], 𝓂.parameters)][1], v[1]), v[2]) != parameter_values[indexin([k], 𝓂.parameters)][1] - return -Inf - end - end - end + bounds_violated = @ignore_derivatives check_bounds(parameter_values, 𝓂) - # solve model given the parameters - if algorithm == :second_order - sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(parameter_values, 𝓂) + if bounds_violated return -Inf end - if !converged return -Inf end + NSSS_labels = @ignore_derivatives [sort(union(𝓂.exo_present, 𝓂.var))..., 𝓂.calibration_equations_parameters...] - all_SS = expand_steady_state(SS_and_pars,𝓂) + obs_indices = @ignore_derivatives convert(Vector{Int}, indexin(observables, NSSS_labels)) - state = collect(sss) - all_SS + TT, SS_and_pars, 𝐒, state, solved = get_relevant_steady_state_and_state_update(Val(algorithm), parameter_values, 𝓂, tol) - state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[𝓂.timings.past_not_future_and_mixed_idx] - 1 - shock] - return 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 - end - elseif algorithm == :pruned_second_order - sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(parameter_values, 𝓂, pruning = true) - - if !converged return -Inf end - - all_SS = expand_steady_state(SS_and_pars,𝓂) - - state = [zeros(𝓂.timings.nVars), collect(sss) - all_SS] - - state_update = function(pruned_states::Vector{Vector{T}}, shock::Vector{S}) where {T,S} - aug_state₁ = [pruned_states[1][𝓂.timings.past_not_future_and_mixed_idx]; 1; shock] - aug_state₂ = [pruned_states[2][𝓂.timings.past_not_future_and_mixed_idx]; 0; zero(shock)] - - return [𝐒₁ * aug_state₁, 𝐒₁ * aug_state₂ + 𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2] # strictly following Andreasen et al. (2018) - end - elseif algorithm == :third_order - sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃ = calculate_third_order_stochastic_steady_state(parameter_values, 𝓂) - - if !converged return -Inf end - - all_SS = expand_steady_state(SS_and_pars,𝓂) - - state = collect(sss) - all_SS - - state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[𝓂.timings.past_not_future_and_mixed_idx] - 1 - shock] - return 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 𝐒₃ * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 - end - elseif algorithm == :pruned_third_order - sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃ = calculate_third_order_stochastic_steady_state(parameter_values, 𝓂, pruning = true) - - if !converged return -Inf end - - all_SS = expand_steady_state(SS_and_pars,𝓂) - - state = [zeros(𝓂.timings.nVars), collect(sss) - all_SS, zeros(𝓂.timings.nVars)] + if !solved return -Inf end - state_update = function(pruned_states::Vector{Vector{T}}, shock::Vector{S}) where {T,S} - aug_state₁ = [pruned_states[1][𝓂.timings.past_not_future_and_mixed_idx]; 1; shock] - aug_state₁̂ = [pruned_states[1][𝓂.timings.past_not_future_and_mixed_idx]; 0; shock] - aug_state₂ = [pruned_states[2][𝓂.timings.past_not_future_and_mixed_idx]; 0; zero(shock)] - aug_state₃ = [pruned_states[3][𝓂.timings.past_not_future_and_mixed_idx]; 0; zero(shock)] - - kron_aug_state₁ = ℒ.kron(aug_state₁, aug_state₁) - - return [𝐒₁ * aug_state₁, 𝐒₁ * aug_state₂ + 𝐒₂ * kron_aug_state₁ / 2, 𝐒₁ * aug_state₃ + 𝐒₂ * ℒ.kron(aug_state₁̂, aug_state₂) + 𝐒₃ * ℒ.kron(kron_aug_state₁,aug_state₁) / 6] - end - else - SS_and_pars, (solution_error, iters) = 𝓂.SS_solve_func(parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) - - if solution_error > tol || isnan(solution_error) - return -Inf - end - - state = zeros(𝓂.timings.nVars) - - ∇₁ = calculate_jacobian(parameter_values, SS_and_pars, 𝓂) |> Matrix - - 𝐒₁, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings) - # 𝐒₁, solved = calculate_quadratic_iteration_solution_AD(∇₁; T = 𝓂.timings) - - if !solved return -Inf end - - state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[𝓂.timings.past_not_future_and_mixed_idx] - shock] - return 𝐒₁ * aug_state # you need a return statement for forwarddiff to work - end + if collect(axiskeys(data,1)) isa Vector{String} + data = @ignore_derivatives rekey(data, 1 => axiskeys(data,1) .|> Meta.parse .|> replace_indices) end - # prepare data - NSSS_labels = @ignore_derivatives [sort(union(𝓂.exo_present,𝓂.var))...,𝓂.calibration_equations_parameters...] - - obs_indices = @ignore_derivatives indexin(observables,NSSS_labels) + dt = @ignore_derivatives collect(data(observables)) - data_in_deviations = collect(data) .- SS_and_pars[obs_indices] - - if filter == :kalman - loglikelihood = calculate_kalman_filter_loglikelihood(𝓂, observables, 𝐒₁, data_in_deviations) - elseif filter == :inversion - loglikelihood = @ignore_derivatives calculate_inversion_filter_loglikelihood(𝓂, state, state_update, data_in_deviations, observables, warmup_iterations) - end + # prepare data + data_in_deviations = dt .- SS_and_pars[obs_indices] - return loglikelihood + return calculate_loglikelihood(Val(filter), observables, 𝐒, data_in_deviations, TT, presample_periods, initial_covariance, state, warmup_iterations) end @@ -3055,7 +2991,6 @@ And data, 5-element Vector{Float64}: (:Equation₃) 0.0 (:Equation₄) 0.0 (:CalibrationEquation₁) 0.0 -``` get_non_stochastic_steady_state_residuals(RBC, [1.1641597, 3.0635781, 1.2254312, 0.0, 0.18157895]) # output diff --git a/src/macros.jl b/src/macros.jl index 0592fad0..652a4753 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -15,7 +15,7 @@ Parses the model equations and assigns them to an object. # Optional arguments to be placed between `𝓂` and `ex` - `max_obc_horizon` [Default: `40`, Type: `Int`]: maximum length of anticipated shocks and corresponding unconditional forecast horizon over which the occasionally binding constraint is to be enforced. Increase this number if no solution is found to enforce the constraint. -Variables must be defined with their time subscript in squared brackets. +Variables must be defined with their time subscript in square brackets. Endogenous variables can have the following: - present: `c[0]` - non-stcohastic steady state: `c[ss]` instead of `ss` any of the following is also a valid flag for the non-stochastic steady state: `ss`, `stst`, `steady`, `steadystate`, `steady_state`, and the parser is case-insensitive (`SS` or `sTst` will work as well). @@ -28,7 +28,7 @@ Exogenous variables (shocks) can have the following: - past: `eps_z[x-1]` - future: `eps_z[x+1]` -Parameters enter the equations without squared brackets. +Parameters enter the equations without square brackets. If an equation contains a `max` or `min` operator, then the default dynamic (first order) solution of the model will enforce the occasionally binding constraint. You can choose to ignore it by setting `ignore_obc = true` in the relevant function calls. @@ -93,6 +93,8 @@ macro model(𝓂,ex...) NSSS_solver_cache = CircularBuffer{Vector{Vector{Float64}}}(500) SS_solve_func = x->x SS_check_func = x->x + ∂SS_equations_∂parameters = ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)) + ∂SS_equations_∂SS_and_pars = ([], Int[], zeros(1,1)) SS_dependencies = nothing original_equations = [] @@ -857,6 +859,8 @@ macro model(𝓂,ex...) $NSSS_solver_cache, $SS_solve_func, $SS_check_func, + $∂SS_equations_∂parameters, + $∂SS_equations_∂SS_and_pars, $SS_dependencies, $➕_vars, @@ -872,10 +876,13 @@ macro model(𝓂,ex...) $bounds, - x->x, + # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)), # model_jacobian + ([], Int[], zeros(1,1)), # model_jacobian + # x->x, # model_jacobian_parameters + ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)), # model_jacobian_SS_and_pars_vars # FWrap{Tuple{Vector{Float64}, Vector{Number}, Vector{Float64}}, SparseMatrixCSC{Float64}}(model_jacobian), - [],#x->x, - [],#x->x, + ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)),#x->x, # model_hessian + ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)),#x->x, # model_third_order_derivatives $T, @@ -916,14 +923,16 @@ macro model(𝓂,ex...) third_order_perturbation_solution(SparseMatrixCSC{Float64, Int64}(ℒ.I,0,0), [], (x,y)->nothing, nothing), auxilliary_indices(Int[],Int[],Int[],Int[],Int[]), second_order_auxilliary_matrices(SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0)), - third_order_auxilliary_matrices(SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0)) + third_order_auxilliary_matrices(SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),Dict{Vector{Int}, Int}(),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0),SparseMatrixCSC{Int, Int64}(ℒ.I,0,0)) ), Float64[], Set([:first_order]), Set(all_available_algorithms), true, false - ) + ), + + Dict{Vector{Symbol}, timings}() # estimation_helper ); end end @@ -953,6 +962,7 @@ Parameters can be defined in either of the following ways: - `silent` [Default: `false`, Type: `Bool`]: do not print any information - `symbolic` [Default: `false`, Type: `Bool`]: try to solve the non stochastic steady state symbolically and fall back to a numerical solution if not possible - `perturbation_order` [Default: `1`, Type: `Int`]: take derivatives only up to the specified order at this stage. In case you want to work with higher order perturbation later on, respective derivatives will be taken at that stage. +- `simplify` [Default: `true`, Type: `Bool`]: whether to elminiate redundant variables and simplify the non stochastic steady state (NSSS) problem. Setting this to `false` can speed up the process, but might make it harder to find the NSSS. If the model does not parse at all (at step 1 or 2), setting this option to `false` might solve it. @@ -1028,6 +1038,7 @@ macro parameters(𝓂,ex...) precompile = false perturbation_order = 1 guess = Dict{Symbol,Float64}() + simplify = true for exp in ex[1:end-1] postwalk(x -> @@ -1045,6 +1056,8 @@ macro parameters(𝓂,ex...) perturbation_order = x.args[2] : x.args[1] == :guess && (isa(eval(x.args[2]), Dict{Symbol, <:Real}) || isa(eval(x.args[2]), Dict{String, <:Real})) ? guess = x.args[2] : + x.args[1] == :simplify && x.args[2] isa Bool ? + simplify = x.args[2] : begin @warn "Invalid options. See docs: `?@parameters` for valid options." x @@ -1446,7 +1459,7 @@ macro parameters(𝓂,ex...) symbolics = create_symbols_eqs!(mod.$𝓂) - remove_redundant_SS_vars!(mod.$𝓂, symbolics) + remove_redundant_SS_vars!(mod.$𝓂, symbolics, avoid_solve = !$simplify) if !$silent println(round(time() - start_time, digits = 3), " seconds") end @@ -1455,7 +1468,7 @@ macro parameters(𝓂,ex...) if !$silent print("Set up non stochastic steady state problem:\t\t\t\t") end - solve_steady_state!(mod.$𝓂, $symbolic, symbolics, verbose = $verbose) # 2nd argument is SS_symbolic + solve_steady_state!(mod.$𝓂, $symbolic, symbolics, verbose = $verbose, avoid_solve = !$simplify) # 2nd argument is SS_symbolic mod.$𝓂.obc_violation_equations = write_obc_violation_equations(mod.$𝓂) diff --git a/src/plotting.jl b/src/plotting.jl index 63175cd0..d484deb1 100644 --- a/src/plotting.jl +++ b/src/plotting.jl @@ -1,5 +1,13 @@ import LaTeXStrings +const default_attributes = Dict(:size=>(700,500), + :plot_titlefont => 10, + :titlefont => 10, + :guidefont => 8, + :legendfontsize => 8, + :tickfontsize => 8, + :framestyle => :semi) + """ ``` gr_backend() @@ -85,6 +93,7 @@ function plot_model_estimates(𝓂::ℳ, warmup_iterations::Int = 0, variables::Union{Symbol_input,String_input} = :all_excluding_obc, shocks::Union{Symbol_input,String_input} = :all, + presample_periods::Int = 0, data_in_levels::Bool = true, shock_decomposition::Bool = false, smooth::Bool = true, @@ -94,33 +103,35 @@ function plot_model_estimates(𝓂::ℳ, save_plots_path::String = ".", plots_per_page::Int = 9, transparency::Float64 = .6, + max_elements_per_legend_row::Int = 4, + extra_legend_space::Float64 = 0.0, + plot_attributes::Dict = Dict(), verbose::Bool = false) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() + attributes = merge(default_attributes, plot_attributes) + + attributes_redux = copy(attributes) - StatsPlots.default(size=(700,500), - plot_titlefont = 10, - titlefont = 10, - guidefont = 8, - legendfontsize = 8, - tickfontsize = 8, - framestyle = :box) + delete!(attributes_redux, :framestyle) + + gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() # write_parameters_input!(𝓂, parameters, verbose = verbose) @assert filter ∈ [:kalman, :inversion] "Currently only the kalman filter (:kalman) for linear models and the inversion filter (:inversion) for linear and nonlinear models are supported." - if algorithm ∈ [:second_order,:pruned_second_order,:third_order,:pruned_third_order] - filter = :inversion - end + pruning = false - if filter == :inversion - shock_decomposition = false + @assert !(algorithm ∈ [:second_order, :third_order]) "Decomposition implemented for first order, pruned second and third order. Second and third order solution decomposition is not yet implemented." + + if algorithm ∈ [:pruned_second_order, :pruned_third_order] + filter = :inversion + pruning = true end solve!(𝓂, parameters = parameters, algorithm = algorithm, verbose = verbose, dynamics = true) - reference_steady_state, (solution_error, iters) = 𝓂.solution.outdated_NSSS ? 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, verbose, false, 𝓂.solver_parameters) : (copy(𝓂.solution.non_stochastic_steady_state), (eps(), 0)) + reference_steady_state, NSSS, SSS_delta = get_relevant_steady_states(𝓂, algorithm) data = data(sort(axiskeys(data,1))) @@ -136,27 +147,41 @@ function plot_model_estimates(𝓂::ℳ, var_idx = parse_variables_input_to_index(variables, 𝓂.timings) shock_idx = parse_shocks_input_to_index(shocks,𝓂.timings) + legend_columns = 1 + + legend_items = length(shock_idx) + 3 + pruning + + max_columns = min(legend_items, max_elements_per_legend_row) + + # Try from max_columns down to 1 to find the optimal solution + for cols in max_columns:-1:1 + if legend_items % cols == 0 || legend_items % cols <= max_elements_per_legend_row + legend_columns = cols + break + end + end + if data_in_levels - data_in_deviations = data .- reference_steady_state[obs_idx] + data_in_deviations = data .- NSSS[obs_idx] else data_in_deviations = data end - # filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) + date_axis = axiskeys(data,2) - # variables_to_plot = filtered_and_smoothed[smooth ? 1 : 5] - # shocks_to_plot = filtered_and_smoothed[smooth ? 3 : 7] - # decomposition = filtered_and_smoothed[smooth ? 4 : 8] + @assert presample_periods < size(data,2) "The number of presample periods must be less than the number of periods in the data." + periods = presample_periods+1:size(data,2) - if filter == :kalman - filtered_and_smoothed = filter_and_smooth(𝓂, data_in_deviations, obs_symbols; verbose = verbose) + date_axis = date_axis[periods] - variables_to_plot = filtered_and_smoothed[smooth ? 1 : 5] - shocks_to_plot = filtered_and_smoothed[smooth ? 3 : 7] - decomposition = filtered_and_smoothed[smooth ? 4 : 8] - elseif filter == :inversion - variables_to_plot, shocks_to_plot = inversion_filter(𝓂, data_in_deviations, algorithm, warmup_iterations = warmup_iterations) + variables_to_plot, shocks_to_plot, standard_deviations, decomposition = filter_data_with_model(𝓂, data_in_deviations, Val(algorithm), Val(filter), warmup_iterations = warmup_iterations, smooth = smooth, verbose = verbose) + + if pruning + decomposition[:,1:(end - 2 - pruning),:] .+= SSS_delta + decomposition[:,end - 2,:] .-= SSS_delta * (size(decomposition,2) - 4) + variables_to_plot .+= SSS_delta + data_in_deviations .+= SSS_delta[obs_idx] end return_plots = [] @@ -174,7 +199,7 @@ function plot_model_estimates(𝓂::ℳ, if i > length(var_idx) # Shock decomposition push!(pp,begin StatsPlots.plot() - StatsPlots.plot!(shocks_to_plot[shock_idx[i - length(var_idx)],:], + StatsPlots.plot!(date_axis, shocks_to_plot[shock_idx[i - length(var_idx)],periods], title = replace_indices_in_symbol(𝓂.timings.exo[shock_idx[i - length(var_idx)]]) * "₍ₓ₎", ylabel = shock_decomposition ? "Absolute Δ" : "Level",label = "", color = shock_decomposition ? estimate_color : :auto) @@ -191,36 +216,55 @@ function plot_model_estimates(𝓂::ℳ, push!(pp,begin StatsPlots.plot() + if shock_decomposition - StatsPlots.groupedbar!(decomposition[var_idx[i],[end-1,shock_idx...],:]', + additional_indices = pruning ? [size(decomposition,2)-1, size(decomposition,2)-2] : [size(decomposition,2)-1] + + StatsPlots.groupedbar!(date_axis, + decomposition[var_idx[i],[additional_indices..., shock_idx...],periods]', bar_position = :stack, - lw = 0, + lc = :transparent, # Line color set to transparent + lw = 0, # This removes the lines around the bars legend = :none, + # yformatter = y -> round(y + SS, digits = 1), # rm Absolute Δ in this case and fix SS additions + # xformatter = x -> string(date_axis[Int(x)]), alpha = transparency) end - StatsPlots.plot!(variables_to_plot[var_idx[i],:] .+ SS, + + StatsPlots.plot!(date_axis, + variables_to_plot[var_idx[i],periods] .+ SS, title = replace_indices_in_symbol(𝓂.timings.var[var_idx[i]]), - ylabel = shock_decomposition ? "Absolute Δ" : "Level",label = "", + ylabel = shock_decomposition ? "Absolute Δ" : "Level", + label = "", + # xformatter = x -> string(date_axis[Int(x)]), color = shock_decomposition ? estimate_color : :auto) + if var_idx[i] ∈ obs_idx - StatsPlots.plot!(data_in_deviations[indexin([var_idx[i]],obs_idx),:]' .+ SS, + StatsPlots.plot!(date_axis, + data_in_deviations[indexin([var_idx[i]],obs_idx),periods]' .+ SS, title = replace_indices_in_symbol(𝓂.timings.var[var_idx[i]]), ylabel = shock_decomposition ? "Absolute Δ" : "Level", label = "", + # xformatter = x -> string(date_axis[Int(x)]), color = shock_decomposition ? data_color : :auto) end + if can_dual_axis StatsPlots.plot!(StatsPlots.twinx(), - 100*((variables_to_plot[var_idx[i],:] .+ SS) ./ SS .- 1), + date_axis, + 100*((variables_to_plot[var_idx[i],periods] .+ SS) ./ SS .- 1), ylabel = LaTeXStrings.L"\% \Delta", label = "") + if var_idx[i] ∈ obs_idx StatsPlots.plot!(StatsPlots.twinx(), - 100*((data_in_deviations[indexin([var_idx[i]],obs_idx),:]' .+ SS) ./ SS .- 1), + date_axis, + 100*((data_in_deviations[indexin([var_idx[i]],obs_idx),periods]' .+ SS) ./ SS .- 1), ylabel = LaTeXStrings.L"\% \Delta", label = "") end end + StatsPlots.hline!(can_dual_axis ? [SS 0] : [SS], color = :black, label = "") @@ -232,19 +276,21 @@ function plot_model_estimates(𝓂::ℳ, else plot_count = 1 - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) # Legend p = StatsPlots.plot(ppp,begin StatsPlots.plot(framestyle = :none) if shock_decomposition - StatsPlots.bar!(fill(0,1,length(shock_idx)+1), - label = reshape(vcat("Initial value",string.(replace_indices_in_symbol.(𝓂.exo[shock_idx]))),1,length(shock_idx)+1), + additional_labels = pruning ? ["Initial value", "Nonlinearities"] : ["Initial value"] + + StatsPlots.bar!(fill(0, 1, length(shock_idx) + 1 + pruning), + label = reshape(vcat(additional_labels, string.(replace_indices_in_symbol.(𝓂.exo[shock_idx]))), 1, length(shock_idx) + 1 + pruning), linewidth = 0, alpha = transparency, lw = 0, legend = :inside, - legend_columns = -1) + legend_columns = legend_columns) end StatsPlots.plot!(fill(0,1,1), label = "Estimate", @@ -255,8 +301,9 @@ function plot_model_estimates(𝓂::ℳ, color = shock_decomposition ? data_color : :auto, legend = :inside) end, - layout = StatsPlots.grid(2, 1, heights=[0.99, 0.01]), - plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")") + layout = StatsPlots.grid(2, 1, heights = [1 - legend_columns * 0.01 - extra_legend_space, legend_columns * 0.01 + extra_legend_space]), + plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux...) push!(return_plots,p) @@ -274,18 +321,20 @@ function plot_model_estimates(𝓂::ℳ, end if length(pp) > 0 - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp,begin StatsPlots.plot(framestyle = :none) if shock_decomposition - StatsPlots.bar!(fill(0,1,length(shock_idx)+1), - label = reshape(vcat("Initial value",string.(replace_indices_in_symbol.(𝓂.exo[shock_idx]))),1,length(shock_idx)+1), + additional_labels = pruning ? ["Initial value", "Nonlinearities"] : ["Initial value"] + + StatsPlots.bar!(fill(0,1,length(shock_idx) + 1 + pruning), + label = reshape(vcat(additional_labels..., string.(replace_indices_in_symbol.(𝓂.exo[shock_idx]))),1,length(shock_idx) + 1 + pruning), linewidth = 0, alpha = transparency, lw = 0, legend = :inside, - legend_columns = -1) + legend_columns = legend_columns) end StatsPlots.plot!(fill(0,1,1), label = "Estimate", @@ -296,8 +345,9 @@ function plot_model_estimates(𝓂::ℳ, color = shock_decomposition ? :darkred : :auto, legend = :inside) end, - layout = StatsPlots.grid(2, 1, heights=[0.99, 0.01]), - plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")") + layout = StatsPlots.grid(2, 1, heights = [1 - legend_columns * 0.01 - extra_legend_space, legend_columns * 0.01 + extra_legend_space]), + plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux...) push!(return_plots,p) @@ -389,17 +439,16 @@ function plot_irf(𝓂::ℳ; generalised_irf::Bool = false, initial_state::Union{Vector{Vector{Float64}},Vector{Float64}} = [0.0], ignore_obc::Bool = false, + plot_attributes::Dict = Dict(), verbose::Bool = false) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() + attributes = merge(default_attributes, plot_attributes) + + attributes_redux = copy(attributes) + + delete!(attributes_redux, :framestyle) - StatsPlots.default(size=(700,500), - plot_titlefont = 10, - titlefont = 10, - guidefont = 8, - legendfontsize = 8, - tickfontsize = 8, - framestyle = :box) + gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks @@ -690,7 +739,7 @@ function plot_irf(𝓂::ℳ; shock_name = "shock_matrix" end - p = StatsPlots.plot(pp...,plot_title = "Model: "*𝓂.model_name*" " * shock_dir * shock_string *" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")") + p = StatsPlots.plot(pp..., plot_title = "Model: "*𝓂.model_name*" " * shock_dir * shock_string *" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; attributes_redux...) push!(return_plots,p) @@ -724,7 +773,7 @@ function plot_irf(𝓂::ℳ; shock_name = "shock_matrix" end - p = StatsPlots.plot(pp...,plot_title = "Model: "*𝓂.model_name*" " * shock_dir * shock_string * " (" * string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page)))*")") + p = StatsPlots.plot(pp..., plot_title = "Model: "*𝓂.model_name*" " * shock_dir * shock_string * " (" * string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page)))*")"; attributes_redux...) push!(return_plots,p) @@ -840,17 +889,18 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; save_plots_format::Symbol = :pdf, save_plots_path::String = ".", plots_per_page::Int = 9, + plot_attributes::Dict = Dict(), + max_elements_per_legend_row::Int = 4, + extra_legend_space::Float64 = 0.0, verbose::Bool = false) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() + attributes = merge(default_attributes, plot_attributes) + + attributes_redux = copy(attributes) - StatsPlots.default(size=(700,500), - plot_titlefont = 10, - titlefont = 10, - guidefont = 8, - legendfontsize = 8, - tickfontsize = 8, - framestyle = :box) + delete!(attributes_redux, :framestyle) + + gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() fevds = get_conditional_variance_decomposition(𝓂, periods = 1:periods, @@ -869,6 +919,20 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; shocks_to_plot = axiskeys(fevds)[2] + legend_columns = 1 + + legend_items = length(shocks_to_plot) + + max_columns = min(legend_items, max_elements_per_legend_row) + + # Try from max_columns down to 1 to find the optimal solution + for cols in max_columns:-1:1 + if legend_items % cols == 0 || legend_items % cols <= max_elements_per_legend_row + legend_columns = cols + break + end + end + n_subplots = length(var_idx) pp = [] pane = 1 @@ -887,16 +951,16 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; else plot_count = 1 - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp,StatsPlots.bar(fill(0,1,length(shocks_to_plot)), label = reshape(string.(replace_indices_in_symbol.(shocks_to_plot)),1,length(shocks_to_plot)), linewidth = 0 , framestyle = :none, legend = :inside, - legend_columns = -1), - layout = StatsPlots.grid(2, 1, heights=[0.99, 0.01]), - plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")") + legend_columns = legend_columns), + layout = StatsPlots.grid(2, 1, heights = [1 - legend_columns * 0.01 - extra_legend_space, legend_columns * 0.01 + extra_legend_space]), + plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; attributes_redux...) push!(return_plots,gr_back ? p : ppp) @@ -914,16 +978,17 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; end if length(pp) > 0 - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp,StatsPlots.bar(fill(0,1,length(shocks_to_plot)), label = reshape(string.(replace_indices_in_symbol.(shocks_to_plot)),1,length(shocks_to_plot)), linewidth = 0 , framestyle = :none, legend = :inside, - legend_columns = -1), - layout = StatsPlots.grid(2, 1, heights=[0.99, 0.01]), - plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")") + legend_columns = legend_columns), + layout = StatsPlots.grid(2, 1, heights = [1 - legend_columns * 0.01 - extra_legend_space, legend_columns * 0.01 + extra_legend_space]), + plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux...) push!(return_plots,gr_back ? p : ppp) @@ -1020,15 +1085,14 @@ function plot_solution(𝓂::ℳ, save_plots_format::Symbol = :pdf, save_plots_path::String = ".", plots_per_page::Int = 6, + plot_attributes::Dict = Dict(), verbose::Bool = false) - StatsPlots.default(size=(700,500), - plot_titlefont = 10, - titlefont = 10, - guidefont = 8, - legendfontsize = 8, - tickfontsize = 8, - framestyle = :box) + attributes = merge(default_attributes, plot_attributes) + + attributes_redux = copy(attributes) + + delete!(attributes_redux, :framestyle) state = state isa Symbol ? state : state |> Meta.parse |> replace_indices @@ -1226,12 +1290,13 @@ function plot_solution(𝓂::ℳ, else plot_count = 1 - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp, legend_plot, layout = StatsPlots.grid(2, 1, heights = length(algorithm) > 3 ? [0.65, 0.35] : [0.8, 0.2]), - plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" + plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux... ) push!(return_plots,p) @@ -1250,12 +1315,13 @@ function plot_solution(𝓂::ℳ, end if length(pp) > 0 - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp, legend_plot, layout = StatsPlots.grid(2, 1, heights = length(algorithm) > 3 ? [0.65, 0.35] : [0.8, 0.2]), - plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" + plot_title = "Model: "*𝓂.model_name*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux... ) push!(return_plots,p) @@ -1368,17 +1434,16 @@ function plot_conditional_forecast(𝓂::ℳ, save_plots_format::Symbol = :pdf, save_plots_path::String = ".", plots_per_page::Int = 9, + plot_attributes::Dict = Dict(), verbose::Bool = false) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() + attributes = merge(default_attributes, plot_attributes) - StatsPlots.default(size=(700,500), - plot_titlefont = 10, - titlefont = 10, - guidefont = 8, - legendfontsize = 8, - tickfontsize = 8, - framestyle = :box) + attributes_redux = copy(attributes) + + delete!(attributes_redux, :framestyle) + + gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() conditions = conditions isa KeyedArray ? axiskeys(conditions,1) isa Vector{String} ? rekey(conditions, 1 => axiskeys(conditions,1) .|> Meta.parse .|> replace_indices) : conditions : conditions @@ -1524,7 +1589,7 @@ function plot_conditional_forecast(𝓂::ℳ, shock_string = "Conditional forecast" - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp,begin StatsPlots.scatter(fill(0,1,1), @@ -1548,7 +1613,8 @@ function plot_conditional_forecast(𝓂::ℳ, legend = :inside) end, layout = StatsPlots.grid(2, 1, heights=[0.99, 0.01]), - plot_title = "Model: "*𝓂.model_name*" " * shock_string * " ("*string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page)))*")") + plot_title = "Model: "*𝓂.model_name*" " * shock_string * " ("*string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux...) push!(return_plots,p) @@ -1569,7 +1635,7 @@ function plot_conditional_forecast(𝓂::ℳ, shock_string = "Conditional forecast" - ppp = StatsPlots.plot(pp...) + ppp = StatsPlots.plot(pp...; attributes...) p = StatsPlots.plot(ppp,begin StatsPlots.scatter(fill(0,1,1), @@ -1593,7 +1659,8 @@ function plot_conditional_forecast(𝓂::ℳ, legend = :inside) end, layout = StatsPlots.grid(2, 1, heights=[0.99, 0.01]), - plot_title = "Model: "*𝓂.model_name*" " * shock_string * " (" * string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page)))*")") + plot_title = "Model: "*𝓂.model_name*" " * shock_string * " (" * string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page)))*")"; + attributes_redux...) push!(return_plots,p) diff --git a/src/structures.jl b/src/structures.jl index 3aa5c53b..75a00145 100644 --- a/src/structures.jl +++ b/src/structures.jl @@ -166,7 +166,8 @@ end struct third_order_auxilliary_matrices 𝐂₃::SparseMatrixCSC{Int} 𝐔₃::SparseMatrixCSC{Int} - + 𝐈₃::Dict{Vector{Int}, Int} + 𝐔∇₃::SparseMatrixCSC{Int} 𝐏::SparseMatrixCSC{Int} @@ -358,6 +359,8 @@ mutable struct ℳ NSSS_solver_cache::CircularBuffer{Vector{Vector{Float64}}} SS_solve_func::Function SS_check_func::Function + ∂SS_equations_∂parameters::Tuple{Vector{Function}, SparseMatrixCSC{<: Real}} + ∂SS_equations_∂SS_and_pars::Tuple{Vector{Function}, Vector{Int}, Matrix{<: Real}} # nonlinear_solution_helper SS_dependencies::Any @@ -379,10 +382,13 @@ mutable struct ℳ bounds::Dict{Symbol,Tuple{Float64,Float64}} - model_jacobian::Function + # model_jacobian::Tuple{Vector{Function}, SparseMatrixCSC{Float64}} + model_jacobian::Tuple{Vector{Function}, Vector{Int}, Matrix{<: Real}} + # model_jacobian_parameters::Function + model_jacobian_SS_and_pars_vars::Tuple{Vector{Function}, SparseMatrixCSC{<: Real}} # model_jacobian::FWrap{Tuple{Vector{Float64}, Vector{Number}, Vector{Float64}}, SparseMatrixCSC{Float64}}#{typeof(model_jacobian)} - model_hessian::Vector{Function} - model_third_order_derivatives::Vector{Function} + model_hessian::Tuple{Vector{Function}, SparseMatrixCSC{<: Real}} + model_third_order_derivatives::Tuple{Vector{Function}, SparseMatrixCSC{<: Real}} timings::timings @@ -396,4 +402,5 @@ mutable struct ℳ solution::solution # symbolics::symbolics + estimation_helper::Dict{Vector{Symbol}, timings} end diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index 18b1082c..0cf1374b 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -224,18 +224,22 @@ function functionality_test(m; algorithm = :first_order, plots = true, verbose = new_cond_var_decomp3 = get_conditional_variance_decomposition(m, verbose = true, parameters = (string.(m.parameters[1:2]) .=> m.parameter_values[1:2] / 1.0001)) new_cond_var_decomp3 = fevd(m, verbose = true, parameters = (string.(m.parameters[1:2]) .=> m.parameter_values[1:2] * 1.0002)) old_cond_var_decomp = get_conditional_variance_decomposition(m, verbose = true, parameters = old_par_vals) + end - + if !(algorithm ∈ [:second_order, :third_order]) # Test filtering and smoothing sol = get_solution(m) if length(m.exo) > 1 - var_idxs = findall(vec(sum(sol[end-length(m.exo)+1:end,:] .!= 0,dims = 1)) .> 0)[1:2] + n_shocks_influence_var = vec(sum(abs.(sol[end-length(m.exo)+1:end,:]) .> eps(),dims = 1)) + var_idxs = findall(n_shocks_influence_var .== maximum(n_shocks_influence_var))[1:2] else var_idxs = [1] end - simulation = simulate(m) + Random.seed!(123) + + simulation = simulate(m, algorithm = algorithm) data_in_levels = simulation(axiskeys(simulation,1) isa Vector{String} ? MacroModelling.replace_indices_in_symbol.(m.var[var_idxs]) : m.var[var_idxs],:,:simulate) data = data_in_levels .- m.solution.non_stochastic_steady_state[var_idxs] @@ -248,22 +252,21 @@ function functionality_test(m; algorithm = :first_order, plots = true, verbose = estim_stds2 = get_estimated_variable_standard_deviations(m, data_in_levels, smooth = false, verbose = true) @test isapprox(estim_stds1,estim_stds2) - estim_stds1 = get_estimated_variable_standard_deviations(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = (m.parameters[1:2] .=> m.parameter_values[1:2] * 1.0001)) - estim_stds1 = get_estimated_variable_standard_deviations(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = (string.(m.parameters[1:2]) .=> m.parameter_values[1:2] * 1.0001)) - estim_stds2 = get_estimated_variable_standard_deviations(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = old_par_vals) - - - estim_decomp1 = get_shock_decomposition(m, data, data_in_levels = false, verbose = true) - estim_decomp2 = get_shock_decomposition(m, data_in_levels, verbose = true) + estim_decomp1 = get_shock_decomposition(m, data, algorithm = algorithm, data_in_levels = false, verbose = true) + estim_decomp2 = get_shock_decomposition(m, data_in_levels, algorithm = algorithm, verbose = true) @test isapprox(estim_decomp1,estim_decomp2) - estim_decomp1 = get_shock_decomposition(m, data, data_in_levels = false, smooth = false, verbose = true) - estim_decomp2 = get_shock_decomposition(m, data_in_levels, smooth = false, verbose = true) + estim_decomp1 = get_shock_decomposition(m, data, algorithm = algorithm, data_in_levels = false, smooth = false, verbose = true) + estim_decomp2 = get_shock_decomposition(m, data_in_levels, algorithm = algorithm, smooth = false, verbose = true) @test isapprox(estim_decomp1,estim_decomp2) - estim_decomp1 = get_shock_decomposition(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = (m.parameters[1:2] .=> m.parameter_values[1:2] * 1.0001)) - estim_decomp1 = get_shock_decomposition(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = (string.(m.parameters[1:2]) .=> m.parameter_values[1:2] * 1.0001)) - estim_decomp2 = get_shock_decomposition(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = old_par_vals) + estim_decomp1 = get_shock_decomposition(m, data, algorithm = algorithm, data_in_levels = false, smooth = false, verbose = true, parameters = (m.parameters[1:2] .=> m.parameter_values[1:2] * 1.0001)) + estim_decomp1 = get_shock_decomposition(m, data, algorithm = algorithm, data_in_levels = false, smooth = false, verbose = true, parameters = (string.(m.parameters[1:2]) .=> m.parameter_values[1:2] * 1.0001)) + estim_decomp2 = get_shock_decomposition(m, data, algorithm = algorithm, data_in_levels = false, smooth = false, verbose = true, parameters = old_par_vals) + + estim_stds1 = get_estimated_variable_standard_deviations(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = (m.parameters[1:2] .=> m.parameter_values[1:2] * 1.0001)) + estim_stds1 = get_estimated_variable_standard_deviations(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = (string.(m.parameters[1:2]) .=> m.parameter_values[1:2] * 1.0001)) + estim_stds2 = get_estimated_variable_standard_deviations(m, data, data_in_levels = false, smooth = false, verbose = true, parameters = old_par_vals) end if algorithm ∈ [:second_order, :pruned_second_order, :third_order, :pruned_third_order] @@ -277,7 +280,9 @@ function functionality_test(m; algorithm = :first_order, plots = true, verbose = varnames = axiskeys(new_sub_irfs_all,1) shocknames = axiskeys(new_sub_irfs_all,3) sol = get_solution(m) - var_idxs = findall(vec(sum(sol[end-length(shocknames)+1:end,:] .!= 0,dims = 1)) .> 0)[[1,end]] + # var_idxs = findall(vec(sum(sol[end-length(shocknames)+1:end,:] .!= 0,dims = 1)) .> 0)[[1,end]] + n_shocks_influence_var = vec(sum(abs.(sol[end-length(m.exo)+1:end,:]) .> eps(),dims = 1)) + var_idxs = findall(n_shocks_influence_var .== maximum(n_shocks_influence_var))[[1,end]] conditions = Matrix{Union{Nothing, Float64}}(undef,size(new_sub_irfs_all,1),2) conditions[var_idxs[1],1] = .01 @@ -371,8 +376,11 @@ function functionality_test(m; algorithm = :first_order, plots = true, verbose = sol = get_solution(m) if length(m.exo) > 1 - var_idxs = findall(vec(sum(abs.(sol[end-length(m.exo)+1:end,:]) .> 1e-10,dims = 1)) .> 1)[1:2] - var_idxs_kalman = findall(vec(sum(sol[end-length(m.exo)+1:end,:] .!= 0,dims = 1)) .> 0)[1:2] + n_shocks_influence_var = vec(sum(abs.(sol[end-length(m.exo)+1:end,:]) .> eps(),dims = 1)) + var_idxs = findall(n_shocks_influence_var .== maximum(n_shocks_influence_var))[1:2] + var_idxs_kalman = var_idxs + # var_idxs = findall(vec(sum(abs.(sol[end-length(m.exo)+1:end,:]) .> 1e-10,dims = 1)) .> 1)[1:2] + # var_idxs_kalman = findall(vec(sum(sol[end-length(m.exo)+1:end,:] .!= 0,dims = 1)) .> 0)[1:2] else var_idxs = [1] end @@ -798,7 +806,9 @@ function functionality_test(m; algorithm = :first_order, plots = true, verbose = sol = get_solution(m) if length(m.exo) > 1 - var_idxs = findall(vec(sum(sol[end-length(m.exo)+1:end,:] .!= 0,dims = 1)) .> 0)[1:2] + # var_idxs = findall(vec(sum(sol[end-length(m.exo)+1:end,:] .!= 0,dims = 1)) .> 0)[1:2] + n_shocks_influence_var = vec(sum(abs.(sol[end-length(m.exo)+1:end,:]) .> eps(),dims = 1)) + var_idxs = findall(n_shocks_influence_var .== maximum(n_shocks_influence_var))[1:2] else var_idxs = [1] end diff --git a/test/runtests.jl b/test/runtests.jl index 4907b942..23f3c57e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,6 +2,7 @@ using Test using MacroModelling using Random using AxisKeys, SparseArrays +import Zygote, FiniteDifferences import StatsPlots, Turing # has to come before Aqua, otherwise exports are not recognised using Aqua using JET @@ -31,6 +32,10 @@ include("functionality_tests.jl") # include("optim_solver_params.jl") # end +if test_set == "estimate_sw07" + include("test_sw07_estimation.jl") +end + if test_set == "estimation" include("test_estimation.jl") end @@ -148,6 +153,7 @@ end if test_set == "plots" plots = true + Random.seed!(1) @testset verbose = true "Backus_Kehoe_Kydland_1992" begin include("../models/Backus_Kehoe_Kydland_1992.jl") @@ -187,6 +193,19 @@ if test_set == "plots" @testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables" begin include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") functionality_test(m, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) end m = nothing GC.gc() @@ -196,6 +215,19 @@ if test_set == "plots" @testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables numerical SS" begin include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve.jl") functionality_test(m, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1, max_range = 1e-4),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) end m = nothing GC.gc() @@ -203,6 +235,19 @@ if test_set == "plots" @testset verbose = true "RBC_CME with calibration equations, parameter definitions, and special functions" begin include("models/RBC_CME_calibration_equations_and_parameter_definitions_and_specfuns.jl") functionality_test(m, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) end m = nothing GC.gc() @@ -210,6 +255,19 @@ if test_set == "plots" @testset verbose = true "RBC_CME with calibration equations and parameter definitions" begin include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") functionality_test(m, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) end m = nothing GC.gc() @@ -218,6 +276,19 @@ if test_set == "plots" @testset verbose = true "RBC_CME with calibration equations" begin include("models/RBC_CME_calibration_equations.jl") functionality_test(m, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) end m = nothing GC.gc() @@ -227,6 +298,19 @@ if test_set == "plots" @testset verbose = true "RBC_CME" begin include("models/RBC_CME.jl") functionality_test(m, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) end m = nothing GC.gc() diff --git a/test/test_1st_order_inversion_filter_estimation.jl b/test/test_1st_order_inversion_filter_estimation.jl index f484e25d..6653c531 100644 --- a/test/test_1st_order_inversion_filter_estimation.jl +++ b/test/test_1st_order_inversion_filter_estimation.jl @@ -1,10 +1,11 @@ using MacroModelling import Turing import Pigeons +import Zygote import Turing: NUTS, sample, logpdf import Optim, LineSearches using Random, CSV, DataFrames, MCMCChains, AxisKeys -import DynamicPPL: logjoint +import DynamicPPL include("../models/FS2000.jl") @@ -31,16 +32,33 @@ dists = [ InverseGamma(0.008862, Inf, μσ = true) # z_e_m ] -Turing.@model function FS2000_loglikelihood_function(data, m) +Turing.@model function FS2000_loglikelihood_function(data, m, filter) all_params ~ Turing.arraydist(dists) - Turing.@addlogprob! get_loglikelihood(m, data, all_params, filter = :inversion) + if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() + Turing.@addlogprob! get_loglikelihood(m, data, all_params, filter = filter) + end end +n_samples = 1000 + +samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :inversion), NUTS(adtype = Turing.AutoZygote()), n_samples, progress = true, init_params = FS2000.parameter_values) + +println("Mean variable values (Zygote): $(mean(samps).nt.mean)") + +sample_nuts = mean(samps).nt.mean + +modeFS2000i = Turing.maximum_a_posteriori(FS2000_loglikelihood_function(data, FS2000, :inversion), + Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), + adtype = Turing.AutoZygote(), + initial_params = FS2000.parameter_values) + +println("Mode variable values: $(modeFS2000i.values); Mode loglikelihood: $(modeFS2000i.lp)") + Random.seed!(30) # generate a Pigeons log potential -FS2000_lp = Pigeons.TuringLogPotential(FS2000_loglikelihood_function(data, FS2000)) +FS2000_lp = Pigeons.TuringLogPotential(FS2000_loglikelihood_function(data, FS2000, :inversion)) # find a feasible starting point pt = Pigeons.pigeons(target = FS2000_lp, n_rounds = 0, n_chains = 1) @@ -66,14 +84,13 @@ Pigeons.initialization(::Pigeons.TuringLogPotential{typeof(FS2000_loglikelihood_ pt = @time Pigeons.pigeons(target = FS2000_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], - n_chains = 1, + n_chains = 2, n_rounds = 10, multithreaded = true) samps = MCMCChains.Chains(pt) - -println(mean(samps).nt.mean) +println("Mean variable values (Pigeons): $(mean(samps).nt.mean)") # # estimate highly nonlinear model diff --git a/test/test_2nd_order_estimation.jl b/test/test_2nd_order_estimation.jl index 14a5b436..f32625e6 100644 --- a/test/test_2nd_order_estimation.jl +++ b/test/test_2nd_order_estimation.jl @@ -4,7 +4,7 @@ import Pigeons import Turing: NUTS, sample, logpdf import Optim, LineSearches using Random, CSV, DataFrames, MCMCChains, AxisKeys -import DynamicPPL: logjoint +import DynamicPPL include("../models/FS2000.jl") @@ -35,7 +35,9 @@ dists = [ Turing.@model function FS2000_loglikelihood_function(data, m) all_params ~ Turing.arraydist(dists) - Turing.@addlogprob! get_loglikelihood(m, data, all_params, algorithm = :pruned_second_order) + if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() + Turing.@addlogprob! get_loglikelihood(m, data, all_params, algorithm = :pruned_second_order) + end end @@ -68,14 +70,14 @@ Pigeons.initialization(::Pigeons.TuringLogPotential{typeof(FS2000_loglikelihood_ pt = @time Pigeons.pigeons(target = FS2000_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], - n_chains = 2, + n_chains = 1, n_rounds = 9, multithreaded = true) samps = MCMCChains.Chains(pt) -println(mean(samps).nt.mean) +println("Mean variable values (Pigeons): $(mean(samps).nt.mean)") # # estimate highly nonlinear model diff --git a/test/test_3rd_order_estimation.jl b/test/test_3rd_order_estimation.jl index 4f0ce9cd..c3cf46de 100644 --- a/test/test_3rd_order_estimation.jl +++ b/test/test_3rd_order_estimation.jl @@ -4,7 +4,7 @@ import Pigeons import Turing: NUTS, sample, logpdf, PG, IS import Optim, LineSearches using Random, CSV, DataFrames, MCMCChains, AxisKeys -import DynamicPPL: logjoint +import DynamicPPL # estimate highly nonlinear model @@ -47,7 +47,9 @@ dists = [ Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m) all_params ~ Turing.arraydist(dists) - Turing.@addlogprob! get_loglikelihood(m, data, all_params, algorithm = :pruned_third_order) + if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() + Turing.@addlogprob! get_loglikelihood(m, data, all_params, algorithm = :pruned_third_order) + end end @@ -88,13 +90,13 @@ Pigeons.initialization(::Pigeons.TuringLogPotential{typeof(Caldara_et_al_2012_lo pt = @time Pigeons.pigeons(target = Caldara_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], n_chains = 1, - n_rounds = 6, + n_rounds = 6, # 7 fails with out of support error multithreaded = true) samps = MCMCChains.Chains(pt) -println(mean(samps).nt.mean) +println("Mean variable values (Pigeons): $(mean(samps).nt.mean)") # include("../models/FS2000.jl") diff --git a/test/test_estimation.jl b/test/test_estimation.jl index b70d44ad..16de11b6 100644 --- a/test/test_estimation.jl +++ b/test/test_estimation.jl @@ -1,9 +1,9 @@ using MacroModelling -import Turing, Pigeons -import Turing: NUTS, sample, logpdf +import Turing, Pigeons, Zygote +import Turing: NUTS, sample, logpdf, AutoZygote import Optim, LineSearches using Random, CSV, DataFrames, MCMCChains, AxisKeys -import DynamicPPL: logjoint +import DynamicPPL include("../models/FS2000.jl") @@ -35,7 +35,9 @@ dists = [ Turing.@model function FS2000_loglikelihood_function(data, m) all_params ~ Turing.arraydist(dists) - Turing.@addlogprob! get_loglikelihood(m, data, all_params) + if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() + Turing.@addlogprob! get_loglikelihood(m, data, all_params) + end end FS2000_loglikelihood = FS2000_loglikelihood_function(data, FS2000) @@ -45,9 +47,14 @@ n_samples = 1000 # using Zygote # Turing.setadbackend(:zygote) -samps = @time sample(FS2000_loglikelihood, NUTS(), n_samples, progress = true)#, init_params = sol) +samps = @time sample(FS2000_loglikelihood, NUTS(), n_samples, progress = true, initial_params = FS2000.parameter_values) + +println("Mean variable values (ForwardDiff): $(mean(samps).nt.mean)") + +samps = @time sample(FS2000_loglikelihood, NUTS(adtype = Turing.AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) + +println("Mean variable values (Zygote): $(mean(samps).nt.mean)") -println(mean(samps).nt.mean) sample_nuts = mean(samps).nt.mean @@ -84,39 +91,32 @@ pt = @time Pigeons.pigeons(target = FS2000_lp, samps = MCMCChains.Chains(pt) -println(mean(samps).nt.mean) +println("Mean variable values (Pigeons): $(mean(samps).nt.mean)") sample_pigeons = mean(samps).nt.mean -Random.seed!(30) - - -function calculate_posterior_loglikelihood(parameters, prior_distribuions) - log_lik = 0 - - for (dist, val) in zip(prior_distribuions, parameters) - log_lik -= logpdf(dist, val) - end - - log_lik -= get_loglikelihood(FS2000, data, parameters) - - return log_lik -end +modeFS2000 = Turing.maximum_a_posteriori(FS2000_loglikelihood, + # Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 2)), + Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), + # Optim.NelderMead(), + adtype = AutoZygote(), + # maxiters = 100, + # lb = [0,0,-10,-10,0,0,0,0,0], + # ub = [1,1,10,10,1,1,1,100,100], + initial_params = FS2000.parameter_values) -sol = Optim.optimize(x -> calculate_posterior_loglikelihood(x, dists), -[0,0,-10,-10,0,0,0,0,0], [1,1,10,10,1,1,1,100,100] ,FS2000.parameter_values, -Optim.Fminbox(Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3))); autodiff = :forward) +println("Mode variable values: $(modeFS2000.values); Mode loglikelihood: $(modeFS2000.lp)") @testset "Estimation results" begin - @test isapprox(sol.minimum, -1343.7491257498598, rtol = eps(Float32)) + # @test isapprox(modeFS2000.lp, 1281.669108730447, rtol = eps(Float32)) @test isapprox(sample_nuts, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) @test isapprox(sample_pigeons[1:length(sample_nuts)], [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) end -plot_model_estimates(FS2000, data, parameters = sol.minimizer) +plot_model_estimates(FS2000, data, parameters = sample_nuts) plot_shock_decomposition(FS2000, data) FS2000 = nothing diff --git a/test/test_models.jl b/test/test_models.jl index 56088e20..4bd86f07 100644 --- a/test/test_models.jl +++ b/test/test_models.jl @@ -26,6 +26,21 @@ if !test_higher_order @test isapprox(var_dec([:y,:yflex,:labobs],:ea) * 100, [10.10, 89.47, 0.53],rtol = 1e-3) @test isapprox(var_dec([:y,:r,:c],:epinf) * 100, [36.42, 29.30, 22.62],rtol = 1e-3) + model = SW07_nonlinear + + observables = [:k,:c,:y] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(SW07_nonlinear) translate_dynare_file("SW07_nonlinear.mod") include("SW07_nonlinear.jl") @@ -43,6 +58,21 @@ if !test_higher_order @test isapprox(var_dec(["K{F}","Y{H}","Z{F}"],"E{F}") * 100, [51.34, 42.44, 52.59],rtol = 1e-3) @test isapprox(var_dec(["K{F}","Y{H}","Z{F}"],"E{H}") * 100, [48.66, 57.56, 47.41],rtol = 1e-3) + model = Backus_Kehoe_Kydland_1992 + + observables = ["C{F}","C{H}"] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Backus_Kehoe_Kydland_1992) translate_dynare_file("Backus_Kehoe_Kydland_1992.mod") include("Backus_Kehoe_Kydland_1992.jl") @@ -61,6 +91,21 @@ if !test_higher_order @test isapprox(var_dec(:EA_Y,[:EA_EPSR,:EA_EPSZ,:US_EPSZ]) * 100, [50.02, 13.54, 4.35],rtol = 1e-3) @test isapprox(var_dec(:US_K,[:EA_EPSRP,:US_EPSR,:US_EPSZ]) * 100, [17.48, 26.83, 27.76],rtol = 1e-3) + model = NAWM_EAUS_2008 + + observables = [:EA_R, :US_R, :EA_PI, :US_PI, :EA_YGAP, :US_YGAP] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end])), model.parameter_values[1:10]) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-3),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end])), model.parameter_values[1:10]) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-2) + write_to_dynare_file(NAWM_EAUS_2008) translate_dynare_file("NAWM_EAUS_2008.mod") include("NAWM_EAUS_2008.jl") @@ -80,6 +125,21 @@ if !test_higher_order @test isapprox(corrr(:k,:l),0.8553,rtol = 1e-3) @test isapprox(corrr(:r,:w),-0.9898,rtol = 1e-3) + model = Baxter_King_1993 + + observables = [:k] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Baxter_King_1993) translate_dynare_file("Baxter_King_1993.mod") include("Baxter_King_1993.jl") @@ -96,6 +156,21 @@ if !test_higher_order @test isapprox(var_dec(:π̂,:) * 100, [4.51, 0.91, 87.44, 7.14],rtol = 1e-4) + model = Ireland_2004 + + observables = [:ŷ, :π̂] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Ireland_2004) translate_dynare_file("Ireland_2004.mod") include("Ireland_2004.jl") @@ -114,6 +189,21 @@ if !test_higher_order @test isapprox(var_dec(:E_BGYN,:) * 100, [3.00, 0.23, 0.54, 0.01, 0.08, 0.38, 0.14, 0.10, 83.60, 3.28, 1.11, 3.36, 0.01, 1.56, 0.17, 0.17, 0.00, 2.24, 0.01], rtol = 1e-3) + model = QUEST3_2009 + + observables = [:outputgap, :inflation, :interest] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1, max_range = 1e-5),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-5) + write_to_dynare_file(QUEST3_2009) translate_dynare_file("QUEST3_2009.mod") # fix BGADJ1 = 0.001BGADJ2; include("QUEST3_2009.jl") @@ -132,6 +222,21 @@ if !test_higher_order @test isapprox(var_dec(:B,:) * 100, [42.97, 19.31, 11.70, 0.36, 4.45, 1.41, 0.70, 0.00, 0.61, 12.54, 2.85, 2.72, 0.38],rtol = 1e-3) + model = GNSS_2010 + + observables = [:C,:Y,:D,:BE] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-5) + write_to_dynare_file(GNSS_2010) translate_dynare_file("GNSS_2010.mod") include("GNSS_2010.jl") @@ -149,6 +254,21 @@ if !test_higher_order @test isapprox(var_dec(:pih,:) * 100, [62.06, 37.94],rtol = 1e-3) @test isapprox(var_dec(:x,:) * 100, [56.22, 43.78],rtol = 1e-3) + model = Gali_Monacelli_2005_CITR + + observables = [:pi] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Gali_Monacelli_2005_CITR) translate_dynare_file("Gali_Monacelli_2005_CITR.mod") # include("Gali_Monacelli_2005_CITR.jl") # cannot do this because the model contains Gamma @@ -166,6 +286,22 @@ if !test_higher_order @test isapprox(var_dec(:w,:) * 100, [2.02, 95.18, 2.80],rtol = 1e-3) @test isapprox(var_dec(:y,:) * 100, [0.47, 99.41, 0.12],rtol = 1e-3) + model = Ascari_Sbordone_2014 + + observables = [:y,:psi] + + Random.seed!(1) + simulated_data = simulate(model) + + get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + + back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + # use forward_cdm so that parameter values stay positive. they would return NaN otherwise + fin_grad = FiniteDifferences.grad(FiniteDifferences.forward_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + write_to_dynare_file(Ascari_Sbordone_2014) translate_dynare_file("Ascari_Sbordone_2014.mod") include("Ascari_Sbordone_2014.jl") @@ -192,6 +328,21 @@ if test_higher_order @test isapprox(autocorr_3rd(:c,5),0.4784,rtol = 1e-3) end +model = SGU_2003_debt_premium + +observables = [:r] + +Random.seed!(1) +simulated_data = simulate(model) + +get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + +back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +@test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(SGU_2003_debt_premium) translate_dynare_file("SGU_2003_debt_premium.mod") include("SGU_2003_debt_premium.jl") @@ -219,6 +370,21 @@ if test_higher_order @test isapprox(autocorr_3rd(:r,5),0.2927,rtol = 1e-3) end +model = JQ_2012_RBC + +observables = [:R,:k] + +Random.seed!(1) +simulated_data = simulate(model) + +get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + +back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +@test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(JQ_2012_RBC) translate_dynare_file("JQ_2012_RBC.mod") include("JQ_2012_RBC.jl") @@ -249,6 +415,21 @@ if test_higher_order @test isapprox(autocorr_3rd(:C,5),0.9178,rtol = 1e-3) end +model = Ghironi_Melitz_2005 + +observables = [:C,:Nx] + +Random.seed!(1) +simulated_data = simulate(model) + +get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + +back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +@test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Ghironi_Melitz_2005) translate_dynare_file("Ghironi_Melitz_2005.mod") include("Ghironi_Melitz_2005.jl") @@ -277,6 +458,21 @@ if test_higher_order @test isapprox(mean_2nd(:C),0.9156,rtol = 1e-3) end +model = Gali_2015_chapter_3_nonlinear + +observables = [:N,:Y] + +Random.seed!(1) +simulated_data = simulate(model) + +get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + +back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +@test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Gali_2015_chapter_3_nonlinear) translate_dynare_file("Gali_2015_chapter_3_nonlinear.mod") include("Gali_2015_chapter_3_nonlinear.jl") @@ -299,6 +495,21 @@ if test_higher_order @test isapprox(autocorr_3rd(:c,5),0.9424,rtol = 1e-3) end +model = Caldara_et_al_2012 + +observables = [:c] + +Random.seed!(1) +simulated_data = simulate(model) + +get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + +back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +@test isapprox(back_grad[1], fin_grad[1], rtol = 1e-3) + write_to_dynare_file(Caldara_et_al_2012) translate_dynare_file("Caldara_et_al_2012.mod") include("Caldara_et_al_2012.jl") @@ -327,8 +538,23 @@ if test_higher_order @test isapprox(autocorr_3rd(:c,5),0.9299,rtol = 1e-3) end +model = Aguiar_Gopinath_2007 + +observables = [:c,:u] + +Random.seed!(1) +simulated_data = simulate(model) + +get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) + +back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) + +@test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + write_to_dynare_file(Aguiar_Gopinath_2007) translate_dynare_file("Aguiar_Gopinath_2007.mod") include("Aguiar_Gopinath_2007.jl") get_solution(Aguiar_Gopinath_2007) -Aguiar_Gopinath_2007 = nothing \ No newline at end of file +Aguiar_Gopinath_2007 = nothing diff --git a/test/test_standalone_function.jl b/test/test_standalone_function.jl index 1222deff..be99324a 100644 --- a/test/test_standalone_function.jl +++ b/test/test_standalone_function.jl @@ -64,7 +64,7 @@ get_irf(RBC_CME, algorithm = :third_order) get_irf(RBC_CME, algorithm = :pruned_third_order) get_irf(RBC_CME, algorithm = :pruned_second_order) -∇₁ = calculate_jacobian(RBC_CME.parameter_values, SS_and_pars, RBC_CME) |> Matrix +∇₁ = calculate_jacobian(RBC_CME.parameter_values, SS_and_pars, RBC_CME)# |> Matrix ∇₂ = calculate_hessian(RBC_CME.parameter_values,SS_and_pars,RBC_CME) ∇₃ = calculate_third_order_derivatives(RBC_CME.parameter_values,SS_and_pars,RBC_CME) #SS = get_steady_state(RBC_CME, derivatives = false) diff --git a/test/test_sw07_estimation.jl b/test/test_sw07_estimation.jl new file mode 100644 index 00000000..7f4032d7 --- /dev/null +++ b/test/test_sw07_estimation.jl @@ -0,0 +1,153 @@ +using MacroModelling +using Zygote +import Turing, Pigeons +import Turing: NUTS, sample, logpdf, AutoZygote +import Optim, LineSearches +using Random, CSV, DataFrames, MCMCChains, AxisKeys +import DynamicPPL + +# load data +dat = CSV.read("data/usmodel.csv", DataFrame) + +# load data +data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) + +# declare observables as written in csv file +observables_old = [:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model + +# Subsample +# subset observables in data +sample_idx = 47:230 # 1960Q1-2004Q4 + +data = data(observables_old, sample_idx) + +# declare observables as written in model +observables = [:dy, :dc, :dinve, :labobs, :pinfobs, :dwobs, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model + +data = rekey(data, :Variable => observables) + + +# Handling distributions with varying parameters using arraydist +dists = [ +InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_ea +InverseGamma(0.1, 2.0, 0.025,5.0, μσ = true), # z_eb +InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_eg +InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_eqs +InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_em +InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_epinf +InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_ew +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhoa +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhob +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhog +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhoqs +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhoms +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhopinf +Beta(0.5, 0.2, 0.001,0.9999, μσ = true), # crhow +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # cmap +Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # cmaw +Normal(4.0, 1.5, 2.0, 15.0), # csadjcost +Normal(1.50,0.375, 0.25, 3.0), # csigma +Beta(0.7, 0.1, 0.001, 0.99, μσ = true), # chabb +Beta(0.5, 0.1, 0.3, 0.95, μσ = true), # cprobw +Normal(2.0, 0.75, 0.25, 10.0), # csigl +Beta(0.5, 0.10, 0.5, 0.95, μσ = true), # cprobp +Beta(0.5, 0.15, 0.01, 0.99, μσ = true), # cindw +Beta(0.5, 0.15, 0.01, 0.99, μσ = true), # cindp +Beta(0.5, 0.15, 0.01, 0.99999, μσ = true), # czcap +Normal(1.25, 0.125, 1.0, 3.0), # cfc +Normal(1.5, 0.25, 1.0, 3.0), # crpi +Beta(0.75, 0.10, 0.5, 0.975, μσ = true), # crr +Normal(0.125, 0.05, 0.001, 0.5), # cry +Normal(0.125, 0.05, 0.001, 0.5), # crdy +Gamma(0.625, 0.1, 0.1, 2.0, μσ = true), # constepinf +Gamma(0.25, 0.1, 0.01, 2.0, μσ = true), # constebeta +Normal(0.0, 2.0, -10.0, 10.0), # constelab +Normal(0.4, 0.10, 0.1, 0.8), # ctrend +Normal(0.5, 0.25, 0.01, 2.0), # cgy +Normal(0.3, 0.05, 0.01, 1.0), # calfa +] + +Turing.@model function SW07_loglikelihood_function(data, m, observables, fixed_parameters, filter) + all_params ~ Turing.arraydist(dists) + + z_ea, z_eb, z_eg, z_eqs, z_em, z_epinf, z_ew, crhoa, crhob, crhog, crhoqs, crhoms, crhopinf, crhow, cmap, cmaw, csadjcost, csigma, chabb, cprobw, csigl, cprobp, cindw, cindp, czcap, cfc, crpi, crr, cry, crdy, constepinf, constebeta, constelab, ctrend, cgy, calfa = all_params + + ctou, clandaw, cg, curvp, curvw = fixed_parameters + + if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() + parameters_combined = [ctou, clandaw, cg, curvp, curvw, calfa, csigma, cfc, cgy, csadjcost, chabb, cprobw, csigl, cprobp, cindw, cindp, czcap, crpi, crr, cry, crdy, crhoa, crhob, crhog, crhoqs, crhoms, crhopinf, crhow, cmap, cmaw, constelab, constepinf, constebeta, ctrend, z_ea, z_eb, z_eg, z_em, z_ew, z_eqs, z_epinf] + + llh = get_loglikelihood(m, data(observables), parameters_combined, presample_periods = 4, initial_covariance = :diagonal, filter = filter) + + Turing.@addlogprob! llh + end +end + +# estimate linear model + +include("../models/Smets_Wouters_2007_linear.jl") + +fixed_parameters = Smets_Wouters_2007_linear.parameter_values[indexin([:ctou, :clandaw, :cg, :curvp, :curvw], Smets_Wouters_2007_linear.parameters)] + +SS(Smets_Wouters_2007_linear, parameters = [:crhoms => 0.01, :crhopinf => 0.01, :crhow => 0.01,:cmap => 0.01,:cmaw => 0.01]) + +SW07_loglikelihood = SW07_loglikelihood_function(data, Smets_Wouters_2007_linear, observables, fixed_parameters, :kalman) +# inversion filter delivers similar results + +# par_names = [:z_ea, :z_eb, :z_eg, :z_eqs, :z_em, :z_epinf, :z_ew, :crhoa, :crhob, :crhog, :crhoqs, :crhoms, :crhopinf, :crhow, :cmap, :cmaw, :csadjcost, :csigma, :chabb, :cprobw, :csigl, :cprobp, :cindw, :cindp, :czcap, :cfc, :crpi, :crr, :cry, :crdy, :constepinf, :constebeta, :constelab, :ctrend, :cgy, :calfa] + +# inits = [Dict(get_parameters(Smets_Wouters_2007_linear, values = true))[string(i)] for i in par_names] + +# modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, +# Optim.SimulatedAnnealing(), +# initial_params = inits) + +# modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, +# Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), +# initial_params = modeSW2007.values) + +modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, + Optim.NelderMead()) + +println("Mode variable values (linear): $(modeSW2007.values); Mode loglikelihood: $(modeSW2007.lp)") + +n_samples = 1000 + +samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = modeSW2007.values) + +println(samps) +println("Mean variable values (linear): $(mean(samps).nt.mean)") + +# estimate nonlinear model + +include("../models/Smets_Wouters_2007.jl") + +fixed_parameters = Smets_Wouters_2007.parameter_values[indexin([:ctou, :clandaw, :cg, :curvp, :curvw], Smets_Wouters_2007.parameters)] + +SS(Smets_Wouters_2007, parameters = [:crhoms => 0.01, :crhopinf => 0.01, :crhow => 0.01, :cmap => 0.01, :cmaw => 0.01])(observables) + +SW07_loglikelihood = SW07_loglikelihood_function(data, Smets_Wouters_2007, observables, fixed_parameters, :kalman) + +# par_names = [:z_ea, :z_eb, :z_eg, :z_eqs, :z_em, :z_epinf, :z_ew, :crhoa, :crhob, :crhog, :crhoqs, :crhoms, :crhopinf, :crhow, :cmap, :cmaw, :csadjcost, :csigma, :chabb, :cprobw, :csigl, :cprobp, :cindw, :cindp, :czcap, :cfc, :crpi, :crr, :cry, :crdy, :constepinf, :constebeta, :constelab, :ctrend, :cgy, :calfa] + +# inits = [Dict(get_parameters(Smets_Wouters_2007, values = true))[string(i)] for i in par_names] + +# modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, +# Optim.SimulatedAnnealing(), +# initial_params = inits) + +# modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, +# Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), +# initial_params = modeSW2007.values) + +modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, + Optim.NelderMead()) + +println("Mode variable values (linear): $(modeSW2007.values); Mode loglikelihood: $(modeSW2007.lp)") + +n_samples = 1000 + +samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = modeSW2007.values) + +println(samps) +println("Mean variable values (nonlinear): $(mean(samps).nt.mean)") \ No newline at end of file