IRF utility functions (#159)

* IRFs

Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiableStateSpaceModels"
 uuid = "beacd9db-9e5e-4956-9b09-459a4b2028df"
 authors = ["Jesse Perla <jesseperla@gmail.com> and contributors"]
-version = "0.4.21"
+version = "0.4.22"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/DifferentiableStateSpaceModels.jl

@@ -22,7 +22,7 @@ export make_perturbation_model, default_model_cache_location, PerturbationModel,
        SolverCache, PerturbationSolverSettings, @include_example_module,
        @make_and_include_perturbation_model, generate_perturbation,
        generate_perturbation_derivatives!, model_H_latex, model_steady_states_iv_latex,
-       model_steady_states_latex, verify_steady_state
+       model_steady_states_latex, verify_steady_state, irf
 
 # export solve # will be replaced by SciML
 

src/sequential.jl

@@ -21,4 +21,19 @@ function DifferenceEquations.QuadraticStateSpaceProblem(sol::AbstractSecondOrder
                                       sol.C_1, sol.C_2,
                                       observables_noise,
                                       observables, noise, syms)
+end
+
+function irf(sol::AbstractFirstOrderPerturbationSolution, ϵ0, T; x0 = zeros(sol.n_x))
+    @assert(length(ϵ0) == sol.n_ϵ)
+    noise = hcat(ϵ0, zeros(sol.n_ϵ, T - 1)) # the ϵ shocks are "noise" in DifferenceEquations for SciML compatibility
+    problem = LinearStateSpaceProblem(sol, x0, (0, T); noise, observables_noise = nothing) # do not include observation noise
+    return solve(problem)
+end
+
+function irf(sol::AbstractSecondOrderPerturbationSolution, ϵ0, T; x0 = zeros(sol.n_x))
+    @assert(length(ϵ0) == sol.n_ϵ)
+    noise = hcat(ϵ0, zeros(sol.n_ϵ, T - 1)) # the ϵ shocks are "noise" in DifferenceEquations for SciML compatibility
+    problem = QuadraticStateSpaceProblem(sol, x0, (0, T); noise,
+                                         observables_noise = nothing) # do not include observation noise
+    return solve(problem)
 end

test/first_order_sequence.jl

@@ -150,4 +150,38 @@ settings = PerturbationSolverSettings(; tol_cholesky = 1e9,
 kalman_test(p_d_input, p_f, m, z, settings)
 test_rrule(Zygote.ZygoteRuleConfig(),
            p_d_input -> kalman_test_alt_prior(p_d_input, p_f, m, z, settings),
-           p_d_input; rrule_f = rrule_via_ad, check_inferred = false)
+           p_d_input; rrule_f = rrule_via_ad, check_inferred = false)
+
+@testset "irf" begin
+    m = @include_example_module(Examples.rbc_observables) # const required due to #117 bug
+    p_f = (ρ = 0.2, δ = 0.02, σ = 0.01, Ω_1 = 0.1)
+    p_d = (α = 0.5, β = 0.95)
+    sol = generate_perturbation(m, p_d, p_f)
+    ϵ0 = [1.0] # only 1 shock here
+    T = 20
+    val = irf(sol, ϵ0, T)
+    @test val.u[4] ≈ [-0.07183218632163027,
+                      -0.0004000000000000003]
+    @test val.z[4] ≈ [-0.0071511417965463755
+                      -0.07183218632163027]
+end
+
+function irf_last(p_d, ϵ0; p_f, m, T = 20)
+    sol = generate_perturbation(m, p_d, p_f)
+    val = irf(sol, ϵ0, T)
+    u = val.u
+    return u[end][1] # return last element
+end
+
+# trouble putting in testset for the above reason
+p_f = (ρ = 0.2, δ = 0.02, σ = 0.01, Ω_1 = 0.1)
+p_d = (α = 0.5, β = 0.95)
+ϵ0 = [1.0]
+m = @include_example_module(Examples.rbc_observables)
+irf_last(p_d, ϵ0; p_f, m)
+
+gradient((p_d, ϵ0) -> irf_last(p_d, ϵ0; p_f, m), p_d, ϵ0)
+
+test_rrule(Zygote.ZygoteRuleConfig(),
+           (p_d, ϵ0) -> irf_last(p_d, ϵ0; p_f, m), p_d, ϵ0;
+           rrule_f = rrule_via_ad, check_inferred = false, rtol = 1e-7)

test/second_order_sequence.jl

@@ -62,3 +62,37 @@ res = gradient((args...) -> likelihood_test_joint_second(args..., m, z), p_d, p_
 @test res[4] ≈ [-562.7501911667255; -4462.244956879686] # Regression test, not sure of the correct number
 # CRTU is preferred but FiniteDifferences is having trouble at this point
 # end
+
+# @testset "irf" begin
+m = @include_example_module(Examples.rbc_observables) # const required due to #117 bug
+p_f = (ρ = 0.2, δ = 0.02, σ = 0.01, Ω_1 = 0.1)
+p_d = (α = 0.5, β = 0.95)
+sol = generate_perturbation(m, p_d, p_f, Val(2))
+ϵ0 = [1.0] # only 1 shock here
+T = 20
+val = irf(sol, ϵ0, T)
+@test val.u[4] ≈ [-0.071737439933804
+                  -0.0004000000000000003]
+@test val.z[4] ≈ [-0.007064576215795991
+                  -0.071737439933804]
+# end
+
+function irf_last_2(p_d, ϵ0; p_f, m, T = 20)
+    sol = generate_perturbation(m, p_d, p_f, Val(2))
+    val = irf(sol, ϵ0, T)
+    u = val.u
+    return u[end][1] # return last element
+end
+
+# Gradients for 2nd order differenceeuqations are not complete
+# p_f = (ρ = 0.2, δ = 0.02, σ = 0.01, Ω_1 = 0.1)
+# p_d = (α = 0.5, β = 0.95)
+# ϵ0 = [1.0]
+# m = @include_example_module(Examples.rbc_observables)
+# irf_last_2(p_d, ϵ0; p_f, m)
+
+# gradient((p_d, ϵ0) -> irf_last_2(p_d, ϵ0; p_f, m), p_d, ϵ0)
+
+# test_rrule(Zygote.ZygoteRuleConfig(),
+#            (p_d, ϵ0) -> irf_last_2(p_d, ϵ0; p_f, m), p_d, ϵ0;
+#            rrule_f = rrule_via_ad, check_inferred = false)