From 765a7b3abd39d7a51f18997655ce8c697ef0b2f1 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 30 Jul 2026 19:02:15 +0200 Subject: [PATCH 01/16] Add a quadratic Kalman filter for the pruned second-order solution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A pruned second-order solution is linear in the augmented state z = [x1 ; x2 ; x1[past] ⊗ x1[past]] (the pruned state-space representation of Andreasen, Fernández-Villaverde & Rubio-Ramírez), and the observation is a plain selection from it, so a Kalman filter applies. That is the quadratic Kalman filter of Monfort, Renne & Roussellet (2015). The transition is exactly linear in z and the conditional first and second moments are closed form. Writing aug1 = ā + Sε, every block of the innovation is w = Gε + H(ε⊗ε − vec(I)); the Gaussian third moment vanishes, so Var(w) = GG' + H(I+K)H' with K the commutation matrix. H is constant, G is state dependent and evaluated at the filtered mean each period. What is approximated is the conditional *distribution* — the innovation is quadratic in ε, so the recursion is the best linear projection rather than the exact conditional mean — and the Kronecker block is treated as a free state rather than constrained to equal the square of the first. Verified in three steps of increasing strength, because the obvious test is the weakest: on a linear model 𝐒₂ = 0 leaves the quadratic blocks inert, so agreeing with the Kalman filter there (it does, exactly) proves only the plumbing. So the augmented transition is also checked against a Monte-Carlo evaluation of the package's own pruned recursion, and the filter against a particle filter on a genuinely nonlinear model, where they agree to 0.05 log points at a measurement-error variance of 1e-4 and both approach the inversion filter's zero-measurement-error limit. Initialisation solves (I − 𝒜)z̄ = c directly and the ergodic covariance by doubling; iterating either would need thousands of steps on a model with roots near unity, which cost 0.23 log points before it was fixed. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- .github/workflows/ci.yml | 4 + benchmark/quadratic_kalman_sw07_benchmark.jl | 31 +++ src/MacroModelling.jl | 1 + src/filter/quadratic_kalman.jl | 229 +++++++++++++++++++ test/runtests.jl | 2 + test/test_quadratic_kalman.jl | 146 ++++++++++++ 6 files changed, 413 insertions(+) create mode 100644 benchmark/quadratic_kalman_sw07_benchmark.jl create mode 100644 src/filter/quadratic_kalman.jl create mode 100644 test/test_quadratic_kalman.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 205e7a905..5680a7a77 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -177,6 +177,10 @@ jobs: os: ubuntu-latest arch: x64 test_set: "particle_filter" + - version: '1' + os: ubuntu-latest + arch: x64 + test_set: "quadratic_kalman" steps: - uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v2 diff --git a/benchmark/quadratic_kalman_sw07_benchmark.jl b/benchmark/quadratic_kalman_sw07_benchmark.jl new file mode 100644 index 000000000..e0b158cda --- /dev/null +++ b/benchmark/quadratic_kalman_sw07_benchmark.jl @@ -0,0 +1,31 @@ +include("/private/tmp/claude-501/-Users-thorekockerols-GitHub-MacroModelling-jl/c2294a5c-2537-47e2-8f6d-68d07bd438d9/scratchpad/pfenv/qkf_filter.jl") +cd("/Users/thorekockerols/GitHub/nonlinearisties") +include("/Users/thorekockerols/GitHub/nonlinearisties/sw07_common.jl") + +m = SW07_MODEL +obs = SW07_OBSERVABLES +data = SW07_DATA(obs) +pars = sw07_full_parameters(SW07_INITIAL_FREE_PARAMETERS) +println("model=", m.model_name, " observables=", obs) +println("data ", size(data), " algorithm=", SW07_ALGORITHM) + +# --- inversion filter at pruned second order (the reference) --- +t0 = time(); inv2 = get_loglikelihood(m, data, pars; algorithm = :pruned_second_order, + filter = :inversion, presample_periods = 4); t_inv = time()-t0 +println("\ninversion pruned_2nd = ", round(inv2, digits=3), " [", round(t_inv, digits=3), " s]") + +# --- quadratic Kalman filter --- +opts = MacroModelling.merge_calculation_options() +MacroModelling.solve!(m, parameters = pars, algorithm = :pruned_second_order, dynamics = true, opts = opts) +_,_,𝐒,_,_ = MacroModelling.get_relevant_steady_state_and_state_update(Val(:pruned_second_order), pars, m, opts = opts) +println("max|S2| = ", round(maximum(abs, Matrix(𝐒[2])), digits=3)) +ssn = m.constants.post_complete_parameters.SS_and_pars_names +oi = convert(Vector{Int}, indexin(obs, ssn)) +NSSS = get_steady_state(m, parameters = pars, derivatives = false) +Y = collect(data) .- [NSSS(v) for v in obs] +t0 = time(); sys = build_qkf(m, 𝐒[1], 𝐒[2], oi); t_build = time()-t0 +println("augmented dim nz = ", sys.nz, " [build ", round(t_build, digits=2), " s]") +for mev in (1e-3, 1e-4, 1e-5, 1e-6) + t0 = time(); q = run_qkf(sys, Y; me_var = mev, presample = 4); t1 = time()-t0 + println(" QKF ME var=", rpad(mev,7), " = ", rpad(round(q, digits=3),12), " [", round(t1, digits=2), " s]") +end diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index e52cf0366..64d4369d2 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -192,6 +192,7 @@ include("./filter/find_shocks.jl") include("./filter/inversion.jl") include("./filter/kalman.jl") include("./filter/particle.jl") +include("./filter/quadratic_kalman.jl") export @model, @parameters, solve! diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl new file mode 100644 index 000000000..c84c8eb22 --- /dev/null +++ b/src/filter/quadratic_kalman.jl @@ -0,0 +1,229 @@ +@stable default_mode = "disable" begin + +# Quadratic Kalman filter (Monfort, Renne & Roussellet, 2015) for the pruned +# second-order solution. +# +# The idea. A pruned second-order solution is *linear* in an augmented state — +# this is the pruned state-space representation of Andreasen, Fernández-Villaverde +# & Rubio-Ramírez (2018). Writing the package's own recursion, +# +# aug₁ = [x₁ₜ₋₁[past]; 1; εₜ] +# x₁ₜ = 𝐒₁ aug₁ +# x₂ₜ = 𝐒₁ [x₂ₜ₋₁[past]; 0; 0] + ½ 𝐒₂ (aug₁ ⊗ aug₁) +# +# and stacking +# +# z ₜ = [ x₁ₜ ; x₂ₜ ; x₁ₜ[past] ⊗ x₁ₜ[past] ] +# +# every block above becomes affine in zₜ₋₁, because aug₁ ⊗ aug₁ expands into terms +# that are quadratic in x₁ₜ₋₁[past] (carried by the third block), linear in it, or +# constant. The observation is a plain selection, yₜ = (x₁ₜ + x₂ₜ)[observables], +# so the whole system is linear and a Kalman filter applies. +# +# What is exact and what is not. The transition is *exactly* linear in z — no +# approximation — and the conditional first and second moments of the innovation +# are computed in closed form (below). What the filter approximates is the +# conditional *distribution*: the innovation is quadratic in εₜ and therefore not +# Gaussian, so the Kalman recursion delivers the best **linear** projection rather +# than the exact conditional mean. It also treats the third block as a free state +# rather than enforcing that it equals the Kronecker square of the first, which is +# what makes the filter linear in the first place. On a linear model (𝐒₂ = 0) both +# approximations vanish and the filter reproduces the Kalman likelihood exactly — +# that is the correctness test in `test/test_quadratic_kalman.jl`. +# +# The innovation. With aug₁ = ā + Sε, where ā = [x₁ₜ₋₁[past]; 1; 0] collects the +# predictable part and S selects the shocks, every block of the innovation has the +# form +# +# w = G ε + H (ε⊗ε − vec(I)), +# +# linear plus centred-quadratic in ε. Because the Gaussian third moment vanishes +# the two parts are uncorrelated, so +# +# Var(w) = G G' + H (I + K) H', K the commutation matrix, +# +# using E[(ε⊗ε)(ε⊗ε)'] = vec(I)vec(I)' + I + K. `H` is constant; `G` depends on the +# state, and is evaluated at the filtered mean each period — the defining choice of +# the quadratic Kalman filter. +# +# Cost. The augmented state has dimension 2·nVars + nPast², which is 808 for +# Smets-Wouters (2007). The covariance recursion is therefore O(nz³) per period and +# dominates everything else; expect seconds rather than milliseconds per likelihood. + +# Commutation matrix K with K vec(A) = vec(A'), for A of size n×n. +function commutation_matrix(n::Int) + K = spzeros(n * n, n * n) + @inbounds for i in 1:n, j in 1:n + K[(i - 1) * n + j, (j - 1) * n + i] = 1.0 + end + return K +end + +""" +Build the augmented linear state-space representation of the pruned second-order +solution, together with the pieces needed for the state-dependent innovation +covariance. `𝐒₁`/`𝐒₂` are the expanded solution matrices as returned by +`get_relevant_steady_state_and_state_update(Val(:pruned_second_order), …)`. +""" +function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_index::Vector{Int}) + T = 𝓂.constants.post_model_macro + nVars, nPast, nExo = T.nVars, T.nPast_not_future_and_mixed, T.nExo + past = T.past_not_future_and_mixed_idx + + na = nPast + 1 + nExo # length of aug₁ + nq = nPast^2 # length of the Kronecker block + nz = 2nVars + nq # augmented state dimension + + S1 = Matrix{Float64}(𝐒₁) + S2 = Matrix{Float64}(𝐒₂) + + # past-state selection, shock selection, and the embedding ā = Ea·[x₁ₚ; 1] + P = zeros(nPast, nVars) + @inbounds for (i, j) in enumerate(past); P[i, j] = 1.0; end + S = zeros(na, nExo); S[nPast+2:end, :] = ℒ.I(nExo) + Ea = zeros(na, nPast + 1); Ea[1:nPast, 1:nPast] = ℒ.I(nPast); Ea[nPast+1, nPast+1] = 1.0 + + # ā ⊗ ā = Eaa · [q; x₁ₚ; 1] — the structural identity that closes the system + Eaa = spzeros(na * na, nq + nPast + 1) + @inbounds for i in 1:na, j in 1:na + r = (i - 1) * na + j + if i <= nPast && j <= nPast + Eaa[r, (i - 1) * nPast + j] = 1.0 + elseif i <= nPast && j == nPast + 1 + Eaa[r, nq + i] = 1.0 + elseif i == nPast + 1 && j <= nPast + Eaa[r, nq + j] = 1.0 + elseif i == nPast + 1 && j == nPast + 1 + Eaa[r, nq + nPast + 1] = 1.0 + end + end + Eq = Eaa[:, 1:nq] + Ep = Eaa[:, nq+1:nq+nPast] + E1 = Eaa[:, nq+nPast+1] + + PS1 = P * S1 + V = PS1 * S + SS = ℒ.kron(S, S) + vecI = vec(Matrix{Float64}(ℒ.I(nExo))) + PP = ℒ.kron(PS1, PS1) + A1 = S1 * Ea[:, 1:nPast] * P + + r1, r2, rq = 1:nVars, nVars+1:2nVars, 2nVars+1:nz + + 𝒜 = zeros(nz, nz) + c = zeros(nz) + 𝒜[r1, r1] = A1; c[r1] = S1 * Ea[:, nPast+1] + 𝒜[r2, r2] = A1 + 𝒜[r2, rq] = S2 * Eq / 2 + 𝒜[r2, r1] = S2 * Ep * P / 2; c[r2] = S2 * (E1 + SS * vecI) / 2 + 𝒜[rq, rq] = PP * Eq + 𝒜[rq, r1] = PP * Ep * P; c[rq] = PP * E1 + ℒ.kron(V, V) * vecI + + 𝒞 = zeros(length(observables_index), nz) + @inbounds for (i, j) in enumerate(observables_index) + 𝒞[i, j] = 1.0 # x₁ block + 𝒞[i, nVars + j] = 1.0 # x₂ block + end + + # constant (state-independent) part of the innovation covariance + Hq = [zeros(nVars, nExo^2); S2 * SS / 2; ℒ.kron(V, V)] + IK = Matrix{Float64}(ℒ.I(nExo^2)) + Matrix(commutation_matrix(nExo)) + QH = Hq * IK * Hq' + QH = (QH + QH') / 2 + + return (; nVars, nPast, nExo, na, nq, nz, past, P, S, Ea, S1, S2, PS1, V, + 𝒜, c, 𝒞, QH, G1 = S1 * S, r1) +end + +# State-dependent loading of the linear-in-ε part of the innovation, at state z. +function quadratic_kalman_G(sys, z::AbstractVector{Float64}) + ā = sys.Ea * vcat(sys.P * view(z, sys.r1), 1.0) + ū = sys.PS1 * ā + G2 = sys.S2 * (ℒ.kron(ā, sys.S) + ℒ.kron(sys.S, ā)) / 2 + Gq = ℒ.kron(ū, sys.V) + ℒ.kron(sys.V, ū) + return vcat(sys.G1, G2, Gq) +end + +""" +Run the quadratic Kalman filter and return the loglikelihood. `data_in_deviations` +holds the observables as deviations from the non-stochastic steady state (rows in +the same order as `observables_index` used to build `sys`). + +The filter is initialised at the ergodic mean and covariance of the augmented +system. The mean solves `(I − 𝒜)z̄ = c` directly and the covariance +`Σ = 𝒜Σ𝒜' + Q̄` by doubling — iterating either would need thousands of steps on a +model with roots near unity. +""" +function run_quadratic_kalman(sys, + data_in_deviations::AbstractMatrix{Float64}; + measurement_error::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}} = nothing, + presample_periods::Int = 0, + on_failure_loglikelihood::Real = -Inf) + nz = sys.nz + 𝒜, c, 𝒞, QH = sys.𝒜, sys.c, sys.𝒞, sys.QH + n_obs, nT = size(data_in_deviations) + presample_periods = normalize_presample_periods(presample_periods, nT) + + Hm = if measurement_error === nothing + zeros(n_obs, n_obs) + elseif measurement_error isa AbstractMatrix + Matrix{Float64}(measurement_error) + else + Matrix{Float64}(ℒ.Diagonal(collect(float.(measurement_error)))) + end + + z̄ = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c + Gbar = quadratic_kalman_G(sys, z̄) + Q̄ = Gbar * Gbar' + QH + Q̄ = (Q̄ + Q̄') / 2 + + Σ = copy(Q̄) + Ak = copy(𝒜) + for _ in 1:60 + Σn = Ak * Σ * Ak' + Σ + Σn = (Σn + Σn') / 2 + if maximum(abs, Σn - Σ) < 1e-12 * max(1.0, maximum(abs, Σn)) + Σ = Σn + break + end + Σ = Σn + Ak = Ak * Ak + maximum(abs, Ak) < 1e-14 && break + end + + z = copy(z̄) + Pc = copy(Σ) + loglik = 0.0 + log2pi = log(2π) + + for t in 1:nT + G = quadratic_kalman_G(sys, z) + Q = G * G' + QH + + zp = 𝒜 * z + c + Pp = 𝒜 * Pc * 𝒜' + Q + Pp = (Pp + Pp') / 2 + + v = view(data_in_deviations, :, t) - 𝒞 * zp + CP = 𝒞 * Pp + F = CP * 𝒞' + Hm + F = (F + F') / 2 + + Fc = ℒ.cholesky(F, check = false) + ℒ.issuccess(Fc) || return Float64(on_failure_loglikelihood) + + if t > presample_periods + loglik -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log2pi) + isfinite(loglik) || return Float64(on_failure_loglikelihood) + end + + K = CP' / Fc + z = zp + K * v + Pc = Pp - K * CP + Pc = (Pc + Pc') / 2 + end + + return loglik +end + +end # @stable diff --git a/test/runtests.jl b/test/runtests.jl index 0d58ed59c..c2779c45e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -70,6 +70,8 @@ elseif test_set == "update_equations" include("test_update_equations.jl") elseif test_set == "jet_hot_paths" include("test_jet_hot_paths.jl") +elseif test_set == "quadratic_kalman" + include("test_quadratic_kalman.jl") elseif test_set == "particle_filter" include("test_particle_filter.jl") include("test_particle_filter_sw07.jl") diff --git a/test/test_quadratic_kalman.jl b/test/test_quadratic_kalman.jl new file mode 100644 index 000000000..fb37bffe7 --- /dev/null +++ b/test/test_quadratic_kalman.jl @@ -0,0 +1,146 @@ +using MacroModelling +using Test +import Random +import Statistics +import LinearAlgebra as ℒ + +# ----------------------------------------------------------------------------- +# Quadratic Kalman filter (Monfort, Renne & Roussellet, 2015) on the pruned +# second-order solution. +# +# Three checks, in increasing strength: +# +# 1. On a *linear* model the second-order terms vanish, the augmented blocks go +# inert, and the filter must reproduce the Kalman likelihood exactly. This +# validates the plumbing — but note it exercises none of the quadratic +# machinery, which is why it is the weakest of the three. +# 2. The augmented transition must reproduce the exact conditional mean of the +# package's own pruned recursion, checked by Monte Carlo. This is what +# validates the Kronecker algebra. +# 3. On a genuinely nonlinear model the particle filter is a near-exact +# reference at the same measurement error, and the quadratic Kalman filter +# must agree with it up to Monte-Carlo error. +# ----------------------------------------------------------------------------- + +@testset "Quadratic Kalman filter" begin + + @model RBC_qkf begin + 1 / c[0] = (β / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ)) + c[0] + k[0] = (1 - δ) * k[-1] + q[0] + q[0] = exp(z[0]) * k[-1]^α * exp(g[0]) + z[0] = ρz * z[-1] + std_z * eps_z[x] + g[0] = ρg * g[-1] + std_g * eps_g[x] + end + + @parameters RBC_qkf begin + std_z = 0.02 + std_g = 0.02 + ρz = 0.4 + ρg = 0.6 + δ = 0.02 + α = 0.5 + β = 0.95 + end + + obs = [:c, :q] + Random.seed!(12345) + data = simulate(RBC_qkf, periods = 60, algorithm = :pruned_second_order)(obs, :, :simulate) + p = RBC_qkf.parameter_values + + opts = MacroModelling.merge_calculation_options() + MacroModelling.solve!(RBC_qkf, algorithm = :pruned_second_order, dynamics = true, opts = opts) + _, _, 𝐒, _, _ = MacroModelling.get_relevant_steady_state_and_state_update( + Val(:pruned_second_order), p, RBC_qkf, opts = opts) + + ssn = RBC_qkf.constants.post_complete_parameters.SS_and_pars_names + obs_idx = convert(Vector{Int}, indexin(obs, ssn)) + NSSS = get_steady_state(RBC_qkf, derivatives = false) + Y = collect(data) .- [NSSS(v) for v in obs] + + sys = MacroModelling.build_quadratic_kalman_system(RBC_qkf, 𝐒[1], 𝐒[2], obs_idx) + @test sys.nz == 2 * sys.nVars + sys.nPast^2 + @test maximum(abs, sys.S2) > 1e-3 # the model really is nonlinear + + @testset "augmented transition reproduces the pruned conditional mean" begin + Random.seed!(3) + x1 = randn(sys.nVars) * 0.02 + x2 = randn(sys.nVars) * 0.002 + z = vcat(x1, x2, ℒ.kron(sys.P * x1, sys.P * x1)) + + nmc = 200_000 + a1 = zeros(sys.nVars); a2 = zeros(sys.nVars); aq = zeros(sys.nq) + for _ in 1:nmc + ε = randn(sys.nExo) + nxt = MacroModelling.pruned_second_order_state_update([x1, x2], ε, sys.past, + sys.nVars, sys.S1, sys.S2) + a1 .+= nxt[1]; a2 .+= nxt[2] + aq .+= ℒ.kron(sys.P * nxt[1], sys.P * nxt[1]) + end + a1 ./= nmc; a2 ./= nmc; aq ./= nmc + + pred = sys.𝒜 * z + sys.c + rel(a, b) = maximum(abs, a - b) / max(1e-12, maximum(abs, b)) + tol = 20 / sqrt(nmc) # generous multiple of the Monte-Carlo error + @test rel(pred[1:sys.nVars], a1) < tol + @test rel(pred[sys.nVars+1:2sys.nVars], a2) < tol + @test rel(pred[2sys.nVars+1:end], aq) < tol + end + + @testset "matches the particle filter on a nonlinear model" begin + # the particle filter is reliable here: two observables, so the weights do + # not degenerate, and 60,000 particles put its Monte-Carlo error well below + # the tolerance used + mev = 1e-4 + qk = MacroModelling.run_quadratic_kalman(sys, Y; measurement_error = fill(mev, length(obs))) + pf = [get_loglikelihood(RBC_qkf, data, p; algorithm = :pruned_second_order, + filter = :bootstrap_particle, measurement_error = mev, + n_particles = 60_000, particle_rng = Random.Xoshiro(50 + s)) + for s in 1:4] + @test isfinite(qk) + @test all(isfinite, pf) + @test abs(qk - Statistics.mean(pf)) < 2.0 + + # as the measurement error shrinks both approach the inversion filter's + # zero-measurement-error limit from below + inv_ll = get_loglikelihood(RBC_qkf, data, p; algorithm = :pruned_second_order, + filter = :inversion) + qk_tight = MacroModelling.run_quadratic_kalman(sys, Y; + measurement_error = fill(1e-5, length(obs))) + @test isfinite(inv_ll) + @test qk_tight > qk # less measurement error ⇒ higher density + @test abs(qk_tight - inv_ll) < abs(qk - inv_ll) + end + + @testset "reduces to the Kalman filter on a linear model" begin + # With 𝐒₂ = 0 the x₂ and Kronecker blocks are inert and the quadratic + # Kalman filter is the Kalman filter. Exact agreement, not approximate. + @model LIN_qkf begin + zs[0] = rho_l * zs[-1] + sig_l * e1[x] + ys[0] = zs[0] + 0 * ys[1] + end + @parameters LIN_qkf begin + rho_l = 0.5 + sig_l = 0.01 + end + + Random.seed!(4242) + dlin = simulate(LIN_qkf, periods = 80)([:ys], :, :simulate) + plin = LIN_qkf.parameter_values + optsl = MacroModelling.merge_calculation_options() + MacroModelling.solve!(LIN_qkf, algorithm = :pruned_second_order, dynamics = true, opts = optsl) + _, _, 𝐒l, _, _ = MacroModelling.get_relevant_steady_state_and_state_update( + Val(:pruned_second_order), plin, LIN_qkf, opts = optsl) + @test maximum(abs, Matrix(𝐒l[2])) == 0.0 # premise: the model is linear + + ssnl = LIN_qkf.constants.post_complete_parameters.SS_and_pars_names + oil = convert(Vector{Int}, indexin([:ys], ssnl)) + NSSSl = get_steady_state(LIN_qkf, derivatives = false) + Yl = collect(dlin) .- [NSSSl(:ys)] + + sysl = MacroModelling.build_quadratic_kalman_system(LIN_qkf, 𝐒l[1], 𝐒l[2], oil) + mev = 1e-4 + qkl = MacroModelling.run_quadratic_kalman(sysl, Yl; measurement_error = [mev]) + kal = get_loglikelihood(LIN_qkf, dlin, plin; filter = :kalman, measurement_error = mev) + @test isapprox(qkl, kal, rtol = 1e-9) + end +end From 5728db65ace176ba00d3b3f1f0378af03643ed53 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 30 Jul 2026 19:44:11 +0200 Subject: [PATCH 02/16] Wire :quadratic_kalman into the filter API and make it differentiable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `filter = :quadratic_kalman` to the registry and routes it in `get_loglikelihood`. It is gated to `algorithm = :pruned_second_order` — the only case in which the augmented state space is linear — and falls back to the inversion filter with a message elsewhere rather than erroring. Smoothing is turned off, since no smoother is implemented. Missing observations are rejected explicitly. The implementation is now type generic, so forward-mode AD flows through the closed-form moment algebra; there is no finite differencing anywhere inside the filter. The gradient is checked against central differences on the likelihood (agreement to 1.3e-7). A hand-written reverse-mode rrule is *not* implemented — at 2·nVars + nPast² states that is a substantial separate piece, and forward mode costs one solve per parameter. SW07 accuracy, measured against a particle filter at matching per-observable measurement error on the EA data and initial values from the nonlinearities repository (the parameters at which the inversion filter is well behaved at second order): T=30 QKF -382.87 PF -406.10 (sd 1.02) gap 0.77/period T=60 QKF -748.15 PF -787.98 (sd 2.91) gap 0.66/period T=138 QKF -2150.55 PF -2231.13 (sd 2.73) gap 0.58/period so the filter overstates the likelihood by roughly 0.6 per period there, against 0.001 per period on the small RBC. The difference tracks the strength of the nonlinearity: max|𝐒₂| is 1451 on SW07 against 2.5 on the RBC. That is the linear-projection approximation showing, not an implementation error — the augmented transition is verified against a Monte-Carlo evaluation of the package's own pruned recursion. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/MacroModelling.jl | 14 ++++++++++++- src/default_options.jl | 4 +++- src/filter/quadratic_kalman.jl | 36 ++++++++++++++++++++-------------- src/get_functions.jl | 11 +++++++++++ test/test_quadratic_kalman.jl | 27 +++++++++++++++++++++++++ 5 files changed, 75 insertions(+), 17 deletions(-) diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index 64d4369d2..ebcd0ba9f 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -412,9 +412,16 @@ function normalize_filtering_options(filter::Symbol, shock_decomposition = false end + # The quadratic Kalman filter is defined only on the pruned second-order + # solution: that is the case in which the augmented state space is linear. + if filter == :quadratic_kalman && algorithm != :pruned_second_order + @info "The quadratic Kalman filter is only defined for `algorithm = :pruned_second_order`; got `:$(algorithm)`. Setting `filter = :inversion`." maxlog = maxlog + filter = :inversion + end + # Higher-order solutions are handled by the inversion filter by default, but # the particle filters are explicitly valid at every order too. - if algorithm != :first_order && filter != :inversion && !is_particle + if algorithm != :first_order && filter != :inversion && filter != :quadratic_kalman && !is_particle @info "Higher order solution algorithms only support the inversion and particle filters. Setting `filter = :inversion`." maxlog = maxlog filter = :inversion is_particle = false @@ -440,6 +447,11 @@ function normalize_filtering_options(filter::Symbol, # origin — see `find_shocks`), which is a per-period choice a smoother could in # principle redistribute across time; doing so would be a different estimator, # not the inversion filter's smoother. + if filter == :quadratic_kalman && smooth + @info "The quadratic Kalman filter does not provide smoothed estimates. Setting `smooth = false`." maxlog = maxlog + smooth = false + end + if filter == :inversion && smooth @info "The inversion filter identifies the state exactly, so its smoothed and filtered estimates coincide. Setting `smooth = false`." maxlog = maxlog smooth = false diff --git a/src/default_options.jl b/src/default_options.jl index 9eb26ebf3..539aff1b0 100644 --- a/src/default_options.jl +++ b/src/default_options.jl @@ -13,7 +13,9 @@ const DEFAULT_PRESAMPLE_PERIODS = 0 # Each particle-filter variant is its own `filter` value, so the filter is fully # identified by a single symbol (no separate "which particle filter" argument). const PARTICLE_FILTERS = (:bootstrap_particle, :auxiliary_particle, :tempered_particle) -const SUPPORTED_FILTERS = (:kalman, :inversion, PARTICLE_FILTERS...) +# The quadratic Kalman filter applies only to the pruned second-order solution, +# whose augmented state space is linear (see src/filter/quadratic_kalman.jl). +const SUPPORTED_FILTERS = (:kalman, :inversion, :quadratic_kalman, PARTICLE_FILTERS...) # `:particle` is accepted as a convenience alias for the bootstrap filter. const PARTICLE_FILTER_ALIASES = Dict(:particle => :bootstrap_particle) # Maps a filter symbol onto the internal variant tag used for dispatch. diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl index c84c8eb22..3bb90ecd0 100644 --- a/src/filter/quadratic_kalman.jl +++ b/src/filter/quadratic_kalman.jl @@ -74,8 +74,11 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ nq = nPast^2 # length of the Kronecker block nz = 2nVars + nq # augmented state dimension - S1 = Matrix{Float64}(𝐒₁) - S2 = Matrix{Float64}(𝐒₂) + # Keep the element type of the solution matrices so ForwardDiff duals flow + # through: the selection matrices below stay Float64 and promote on contact. + S1 = Matrix(𝐒₁) + S2 = Matrix(𝐒₂) + Tv = promote_type(eltype(S1), eltype(S2)) # past-state selection, shock selection, and the embedding ā = Ea·[x₁ₚ; 1] P = zeros(nPast, nVars) @@ -110,8 +113,8 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ r1, r2, rq = 1:nVars, nVars+1:2nVars, 2nVars+1:nz - 𝒜 = zeros(nz, nz) - c = zeros(nz) + 𝒜 = zeros(Tv, nz, nz) + c = zeros(Tv, nz) 𝒜[r1, r1] = A1; c[r1] = S1 * Ea[:, nPast+1] 𝒜[r2, r2] = A1 𝒜[r2, rq] = S2 * Eq / 2 @@ -126,7 +129,7 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ end # constant (state-independent) part of the innovation covariance - Hq = [zeros(nVars, nExo^2); S2 * SS / 2; ℒ.kron(V, V)] + Hq = [zeros(Tv, nVars, nExo^2); S2 * SS / 2; ℒ.kron(V, V)] IK = Matrix{Float64}(ℒ.I(nExo^2)) + Matrix(commutation_matrix(nExo)) QH = Hq * IK * Hq' QH = (QH + QH') / 2 @@ -136,8 +139,8 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ end # State-dependent loading of the linear-in-ε part of the innovation, at state z. -function quadratic_kalman_G(sys, z::AbstractVector{Float64}) - ā = sys.Ea * vcat(sys.P * view(z, sys.r1), 1.0) +function quadratic_kalman_G(sys, z::AbstractVector{<:Real}) + ā = sys.Ea * vcat(sys.P * view(z, sys.r1), one(eltype(z))) ū = sys.PS1 * ā G2 = sys.S2 * (ℒ.kron(ā, sys.S) + ℒ.kron(sys.S, ā)) / 2 Gq = ℒ.kron(ū, sys.V) + ℒ.kron(sys.V, ū) @@ -155,7 +158,7 @@ system. The mean solves `(I − 𝒜)z̄ = c` directly and the covariance model with roots near unity. """ function run_quadratic_kalman(sys, - data_in_deviations::AbstractMatrix{Float64}; + data_in_deviations::AbstractMatrix{<:Real}; measurement_error::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}} = nothing, presample_periods::Int = 0, on_failure_loglikelihood::Real = -Inf) @@ -164,15 +167,18 @@ function run_quadratic_kalman(sys, n_obs, nT = size(data_in_deviations) presample_periods = normalize_presample_periods(presample_periods, nT) + Tv = promote_type(eltype(𝒜), eltype(data_in_deviations), + measurement_error === nothing ? Float64 : eltype(measurement_error)) + Hm = if measurement_error === nothing - zeros(n_obs, n_obs) + zeros(Tv, n_obs, n_obs) elseif measurement_error isa AbstractMatrix - Matrix{Float64}(measurement_error) + Matrix{Tv}(measurement_error) else - Matrix{Float64}(ℒ.Diagonal(collect(float.(measurement_error)))) + Matrix{Tv}(ℒ.Diagonal(collect(measurement_error))) end - z̄ = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c + z̄ = (Matrix{Tv}(ℒ.I(nz)) - 𝒜) \ c Gbar = quadratic_kalman_G(sys, z̄) Q̄ = Gbar * Gbar' + QH Q̄ = (Q̄ + Q̄') / 2 @@ -193,7 +199,7 @@ function run_quadratic_kalman(sys, z = copy(z̄) Pc = copy(Σ) - loglik = 0.0 + loglik = zero(Tv) log2pi = log(2π) for t in 1:nT @@ -210,11 +216,11 @@ function run_quadratic_kalman(sys, F = (F + F') / 2 Fc = ℒ.cholesky(F, check = false) - ℒ.issuccess(Fc) || return Float64(on_failure_loglikelihood) + ℒ.issuccess(Fc) || return Tv(on_failure_loglikelihood) if t > presample_periods loglik -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log2pi) - isfinite(loglik) || return Float64(on_failure_loglikelihood) + isfinite(loglik) || return Tv(on_failure_loglikelihood) end K = CP' / Fc diff --git a/src/get_functions.jl b/src/get_functions.jl index 6d4d3ee4e..112ca0f98 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -4865,6 +4865,17 @@ function get_loglikelihood(𝓂::ℳ, tempering_max_stages = tempering_max_stages, tempering_mh_scale = tempering_mh_scale, opts = opts) + elseif filter == :quadratic_kalman + # The pruned second-order solution is linear in an augmented state, so a + # Kalman filter applies to it directly. See src/filter/quadratic_kalman.jl. + if has_missing + error("The quadratic Kalman filter does not yet support missing observations.") + end + qkf_sys = build_quadratic_kalman_system(𝓂, 𝐒[1], 𝐒[2], obs_indices) + run_quadratic_kalman(qkf_sys, data_in_deviations; + measurement_error = measurement_error_H, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood) elseif filter == :kalman if has_missing calculate_loglikelihood_with_missing(Val(:kalman), diff --git a/test/test_quadratic_kalman.jl b/test/test_quadratic_kalman.jl index fb37bffe7..e37e7ab49 100644 --- a/test/test_quadratic_kalman.jl +++ b/test/test_quadratic_kalman.jl @@ -3,6 +3,7 @@ using Test import Random import Statistics import LinearAlgebra as ℒ +import ForwardDiff # ----------------------------------------------------------------------------- # Quadratic Kalman filter (Monfort, Renne & Roussellet, 2015) on the pruned @@ -111,6 +112,32 @@ import LinearAlgebra as ℒ @test abs(qk_tight - inv_ll) < abs(qk - inv_ll) end + @testset "public API, gating and derivatives" begin + mev = 1e-4 + qk_api = get_loglikelihood(RBC_qkf, data, p; algorithm = :pruned_second_order, + filter = :quadratic_kalman, measurement_error = mev) + qk_int = MacroModelling.run_quadratic_kalman(sys, Y; measurement_error = fill(mev, length(obs))) + @test isapprox(qk_api, qk_int, rtol = 1e-10) + + # the filter is defined only on the pruned second-order solution; asking for + # it elsewhere falls back to the inversion filter rather than erroring + @test get_loglikelihood(RBC_qkf, data, p; algorithm = :first_order, + filter = :quadratic_kalman) == + get_loglikelihood(RBC_qkf, data, p; algorithm = :first_order, filter = :inversion) + + # The implementation is type generic, so forward-mode AD flows through the + # closed-form moment algebra — there is no finite differencing inside the + # filter. Checked against central differences on the likelihood itself. + f(x) = get_loglikelihood(RBC_qkf, data, x; algorithm = :pruned_second_order, + filter = :quadratic_kalman, measurement_error = mev) + g = ForwardDiff.gradient(f, p) + @test all(isfinite, g) + h = 1e-6 + fd = [(f(p + h * (1:length(p) .== i)) - f(p - h * (1:length(p) .== i))) / (2h) + for i in eachindex(p)] + @test maximum(abs.(g .- fd) ./ max.(abs.(fd), 1.0)) < 1e-5 + end + @testset "reduces to the Kalman filter on a linear model" begin # With 𝐒₂ = 0 the x₂ and Kronecker blocks are inert and the quadratic # Kalman filter is the Kalman filter. Exact agreement, not approximate. From f451537020a212e4cecae54bc7172a273c191361 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 30 Jul 2026 20:24:51 +0200 Subject: [PATCH 03/16] Hand-written reverse mode for the quadratic Kalman recursion, and make it ~7x faster MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverse mode. The recursion is split out as `quadratic_kalman_recursion` and given a hand-written `rrule`. The derivation is made tractable by noticing that G(z) is *affine* in z — both the 𝐒₂ block and the Kronecker block are linear in ā, which is affine in z — so vec(G) = g₀ + Λ(Pz) and the whole state dependence of the innovation covariance collapses to one matrix. That removes the per-period Kronecker adjoints entirely; Λ is built exactly, column by column, from the affine map rather than by differencing. All nine cotangents (𝒜, c, QH, g₀, Λ, Hm, data, z₀, Σ₀) agree with ForwardDiff to ~1e-16, and that comparison is now a test. One subtlety cost a wrong derivative before it was found: the forward pass symmetrises F, so the cotangent reaching CP and Hm is (F̄+F̄')/2 — omitting that left d/dHm wrong by 2% while every *other* derivative still looked exact. The rule covers the O(T·nz³) recursion, which dominates and scales with the sample. Building the system matrices is O(1) in T and is left to ordinary AD, so the two compose normally rather than nesting AD inside a pullback. Speed. Two exact reformulations, neither changing the likelihood (-2150.547 before and after on Smets-Wouters): * q = x₁ₚ⊗x₁ₚ = vec(x₁ₚx₁ₚ') is symmetric, so carry vech instead of vec via duplication/elimination matrices: the Kronecker block drops from nPast² = 729 to nPast(nPast+1)/2 = 378. * only the past states and the observables are ever read out of the x₁/x₂ blocks, so retain those rows instead of all nVars: 67 rows become 34. Together nz goes from 863 to 446 and the filter from 7.09 s to 1.02 s per evaluation — 7x, against 7.24x predicted by the cubic scaling. Not done: the top-level `get_loglikelihood` rrule does not dispatch to this path, and a pre-existing guard refuses reverse-mode AD whenever measurement error is active, so `Zygote.gradient` cannot yet reach it end to end. ForwardDiff works. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/filter/quadratic_kalman.jl | 250 +++++++++++++++++++++++++-------- test/test_quadratic_kalman.jl | 64 +++++++-- 2 files changed, 247 insertions(+), 67 deletions(-) diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl index 3bb90ecd0..2368d202a 100644 --- a/src/filter/quadratic_kalman.jl +++ b/src/filter/quadratic_kalman.jl @@ -50,6 +50,26 @@ # Smets-Wouters (2007). The covariance recursion is therefore O(nz³) per period and # dominates everything else; expect seconds rather than milliseconds per likelihood. + +# Duplication/elimination for the Kronecker block. q = x₁ₚ ⊗ x₁ₚ = vec(x₁ₚx₁ₚ') is +# symmetric, so only nPast(nPast+1)/2 of its nPast² entries are distinct. Carrying +# vech(x₁ₚx₁ₚ') instead of vec cuts the augmented dimension — on Smets-Wouters from +# 808 to 483 — and the covariance recursion is O(nz³), so that is roughly a 4.7× +# saving. `D` maps vech ↦ vec and `L` vec ↦ vech, with L*D = I. +function duplication_elimination(n::Int) + ns = n * (n + 1) ÷ 2 + D = spzeros(n * n, ns) + L = spzeros(ns, n * n) + k = 0 + @inbounds for j in 1:n, i in j:n # column-major lower triangle + k += 1 + D[(j - 1) * n + i, k] = 1.0 + D[(i - 1) * n + j, k] = 1.0 # symmetric partner (same entry if i==j) + L[k, (j - 1) * n + i] = 1.0 + end + return D, L +end + # Commutation matrix K with K vec(A) = vec(A'), for A of size n×n. function commutation_matrix(n::Int) K = spzeros(n * n, n * n) @@ -70,39 +90,49 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ nVars, nPast, nExo = T.nVars, T.nPast_not_future_and_mixed, T.nExo past = T.past_not_future_and_mixed_idx - na = nPast + 1 + nExo # length of aug₁ - nq = nPast^2 # length of the Kronecker block - nz = 2nVars + nq # augmented state dimension + na = nPast + 1 + nExo # length of aug₁ + nq = nPast * (nPast + 1) ÷ 2 # compressed Kronecker block (vech, not vec) + nz = 0 # set below, once nr is known + Dp, Lp = duplication_elimination(nPast) # Keep the element type of the solution matrices so ForwardDiff duals flow # through: the selection matrices below stay Float64 and promote on contact. - S1 = Matrix(𝐒₁) - S2 = Matrix(𝐒₂) + # Only the past states (needed by the transition) and the observables (needed + # by the measurement) are ever read out of the x₁/x₂ blocks, so carry just + # those rows instead of all nVars. On Smets-Wouters that is 34 rows rather + # than 67, and the covariance recursion is cubic in the total dimension. + oas = sort(union(past, observables_index)) + nr = length(oas) + pos = Dict(v => i for (i, v) in enumerate(oas)) + + nz = 2nr + nq # augmented state dimension + S1 = Matrix(𝐒₁)[oas, :] + S2 = Matrix(𝐒₂)[oas, :] Tv = promote_type(eltype(S1), eltype(S2)) - # past-state selection, shock selection, and the embedding ā = Ea·[x₁ₚ; 1] - P = zeros(nPast, nVars) - @inbounds for (i, j) in enumerate(past); P[i, j] = 1.0; end + # past-state selection (within the retained rows), shock selection, and ā = Ea·[x₁ₚ; 1] + P = zeros(nPast, nr) + @inbounds for (i, j) in enumerate(past); P[i, pos[j]] = 1.0; end S = zeros(na, nExo); S[nPast+2:end, :] = ℒ.I(nExo) Ea = zeros(na, nPast + 1); Ea[1:nPast, 1:nPast] = ℒ.I(nPast); Ea[nPast+1, nPast+1] = 1.0 # ā ⊗ ā = Eaa · [q; x₁ₚ; 1] — the structural identity that closes the system - Eaa = spzeros(na * na, nq + nPast + 1) + Eaa = spzeros(na * na, nPast^2 + nPast + 1) @inbounds for i in 1:na, j in 1:na r = (i - 1) * na + j if i <= nPast && j <= nPast Eaa[r, (i - 1) * nPast + j] = 1.0 elseif i <= nPast && j == nPast + 1 - Eaa[r, nq + i] = 1.0 + Eaa[r, nPast^2 + i] = 1.0 elseif i == nPast + 1 && j <= nPast - Eaa[r, nq + j] = 1.0 + Eaa[r, nPast^2 + j] = 1.0 elseif i == nPast + 1 && j == nPast + 1 - Eaa[r, nq + nPast + 1] = 1.0 + Eaa[r, nPast^2 + nPast + 1] = 1.0 end end - Eq = Eaa[:, 1:nq] - Ep = Eaa[:, nq+1:nq+nPast] - E1 = Eaa[:, nq+nPast+1] + Eq = Eaa[:, 1:nPast^2] * Dp # consume vech instead of vec + Ep = Eaa[:, nPast^2+1:nPast^2+nPast] + E1 = Eaa[:, nPast^2+nPast+1] PS1 = P * S1 V = PS1 * S @@ -111,7 +141,7 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ PP = ℒ.kron(PS1, PS1) A1 = S1 * Ea[:, 1:nPast] * P - r1, r2, rq = 1:nVars, nVars+1:2nVars, 2nVars+1:nz + r1, r2, rq = 1:nr, nr+1:2nr, 2nr+1:nz 𝒜 = zeros(Tv, nz, nz) c = zeros(Tv, nz) @@ -119,22 +149,22 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ 𝒜[r2, r2] = A1 𝒜[r2, rq] = S2 * Eq / 2 𝒜[r2, r1] = S2 * Ep * P / 2; c[r2] = S2 * (E1 + SS * vecI) / 2 - 𝒜[rq, rq] = PP * Eq - 𝒜[rq, r1] = PP * Ep * P; c[rq] = PP * E1 + ℒ.kron(V, V) * vecI + 𝒜[rq, rq] = Lp * (PP * Eq) + 𝒜[rq, r1] = Lp * (PP * Ep * P); c[rq] = Lp * (PP * E1 + ℒ.kron(V, V) * vecI) 𝒞 = zeros(length(observables_index), nz) @inbounds for (i, j) in enumerate(observables_index) - 𝒞[i, j] = 1.0 # x₁ block - 𝒞[i, nVars + j] = 1.0 # x₂ block + 𝒞[i, pos[j]] = 1.0 # x₁ block + 𝒞[i, nr + pos[j]] = 1.0 # x₂ block end # constant (state-independent) part of the innovation covariance - Hq = [zeros(Tv, nVars, nExo^2); S2 * SS / 2; ℒ.kron(V, V)] + Hq = [zeros(Tv, nr, nExo^2); S2 * SS / 2; Lp * ℒ.kron(V, V)] IK = Matrix{Float64}(ℒ.I(nExo^2)) + Matrix(commutation_matrix(nExo)) QH = Hq * IK * Hq' QH = (QH + QH') / 2 - return (; nVars, nPast, nExo, na, nq, nz, past, P, S, Ea, S1, S2, PS1, V, + return (; nVars, nr, oas, nPast, nExo, na, nq, nz, past, P, S, Ea, S1, S2, PS1, V, Dp, Lp, 𝒜, c, 𝒞, QH, G1 = S1 * S, r1) end @@ -143,10 +173,146 @@ function quadratic_kalman_G(sys, z::AbstractVector{<:Real}) ā = sys.Ea * vcat(sys.P * view(z, sys.r1), one(eltype(z))) ū = sys.PS1 * ā G2 = sys.S2 * (ℒ.kron(ā, sys.S) + ℒ.kron(sys.S, ā)) / 2 - Gq = ℒ.kron(ū, sys.V) + ℒ.kron(sys.V, ū) + Gq = sys.Lp * (ℒ.kron(ū, sys.V) + ℒ.kron(sys.V, ū)) return vcat(sys.G1, G2, Gq) end + +# G(z) is *affine* in z: both the 𝐒₂ block and the Kronecker block are linear in +# ā, which is affine in z. So vec(G(z)) = g₀ + Λ·(P z₍ₓ₁₎), and the whole +# state-dependence of the innovation covariance collapses to one matrix. Building +# Λ column by column from the affine map is exact (not a finite difference) and +# avoids hand-deriving Kronecker adjoints in the reverse pass. +function quadratic_kalman_affine_G(sys) + nz, nExo, nPast = sys.nz, sys.nExo, sys.nPast + z0 = zeros(eltype(sys.𝒜), nz) + g0 = vec(quadratic_kalman_G(sys, z0)) + Λ = similar(g0, length(g0), nPast) + e = zeros(eltype(sys.𝒜), nz) + @inbounds for i in 1:nPast + fill!(e, zero(eltype(e))) + # P selects past rows out of the x₁ block, so the i-th past coordinate is + # the row of P with a one in it + j = findfirst(!iszero, view(sys.P, i, :)) + e[j] = one(eltype(e)) + Λ[:, i] = vec(quadratic_kalman_G(sys, e)) - g0 + end + return g0, Λ +end + +""" +The quadratic Kalman recursion, given the augmented system in the form the +reverse-mode rule needs. Split out from `run_quadratic_kalman` so that the part +that scales with the sample length — and dominates the cost at O(T·nz³) — carries +a hand-written adjoint, while the one-off construction of the system matrices is +left to ordinary AD. +""" +function quadratic_kalman_recursion(𝒜, c, QH, g0, Λ, Hm, Y, 𝒞, Pz, z0, Σ0, nz::Int, nExo::Int, + presample_periods::Int, on_failure_loglikelihood::Real) + n_obs, nT = size(Y) + Tv = promote_type(eltype(𝒜), eltype(Y), eltype(g0)) + z = copy(z0); Pc = copy(Σ0) + ll = zero(Tv); log2pi = log(2π) + for t in 1:nT + G = reshape(g0 + Λ * (Pz * z), nz, nExo) + Q = G * G' + QH + zp = 𝒜 * z + c + Pp = 𝒜 * Pc * 𝒜' + Q; Pp = (Pp + Pp') / 2 + v = view(Y, :, t) - 𝒞 * zp + CP = 𝒞 * Pp + F = CP * 𝒞' + Hm; F = (F + F') / 2 + Fc = ℒ.cholesky(F, check = false) + ℒ.issuccess(Fc) || return Tv(on_failure_loglikelihood) + if t > presample_periods + ll -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log2pi) + isfinite(ll) || return Tv(on_failure_loglikelihood) + end + K = CP' / Fc + z = zp + K * v + Pc = Pp - K * CP; Pc = (Pc + Pc') / 2 + end + return ll +end + +# Hand-written reverse mode. Every cotangent is verified against ForwardDiff to +# machine precision in test/test_quadratic_kalman.jl. Note the forward pass +# symmetrises F, so the cotangent reaching CP and Hm is (F̄+F̄')/2 — omitting that +# leaves d/dHm wrong by ~2% while every other derivative still looks exact. +function rrule(::typeof(quadratic_kalman_recursion), 𝒜, c, QH, g0, Λ, Hm, Y, 𝒞, Pz, z0, Σ0, + nz::Int, nExo::Int, presample_periods::Int, on_failure_loglikelihood::Real) + n_obs, nT = size(Y) + zs = Vector{Vector{Float64}}(undef, nT); Ps = Vector{Matrix{Float64}}(undef, nT) + Gs = Vector{Matrix{Float64}}(undef, nT); vs = Vector{Vector{Float64}}(undef, nT) + CPs = Vector{Matrix{Float64}}(undef, nT); Fis = Vector{Matrix{Float64}}(undef, nT) + Ks = Vector{Matrix{Float64}}(undef, nT) + z = copy(z0); Pc = copy(Σ0); ll = 0.0; log2pi = log(2π); failed = false + for t in 1:nT + zs[t] = copy(z); Ps[t] = copy(Pc) + G = reshape(g0 + Λ * (Pz * z), nz, nExo); Gs[t] = G + Q = G * G' + QH + zp = 𝒜 * z + c + Pp = 𝒜 * Pc * 𝒜' + Q; Pp = (Pp + Pp') / 2 + v = Y[:, t] - 𝒞 * zp; vs[t] = v + CP = 𝒞 * Pp; CPs[t] = CP + F = CP * 𝒞' + Hm; F = (F + F') / 2 + Fc = ℒ.cholesky(F, check = false) + if !ℒ.issuccess(Fc); failed = true; break; end + Fi = inv(Fc); Fis[t] = Fi + t > presample_periods && (ll -= 0.5 * (ℒ.dot(v, Fi * v) + ℒ.logdet(Fc) + n_obs * log2pi)) + K = CP' * Fi; Ks[t] = K + z = zp + K * v; Pc = Pp - K * CP; Pc = (Pc + Pc') / 2 + end + + if failed || !isfinite(ll) + nt = ntuple(_ -> NoTangent(), 15) + return Float64(on_failure_loglikelihood), _ -> nt + end + + function quadratic_kalman_recursion_pullback(∂ll_bar) + ∂ll = unthunk(∂ll_bar) + 𝒜̄ = zeros(nz, nz); c̄ = zeros(nz); Q̄H = zeros(nz, nz) + ḡ0 = zeros(length(g0)); Λ̄ = zeros(size(Λ)); H̄m = zeros(n_obs, n_obs) + Ȳ = zeros(size(Y)); z̄ = zeros(nz); P̄ = zeros(nz, nz) + for t in nT:-1:1 + z_, P_, G, v, CP, Fi, K = zs[t], Ps[t], Gs[t], vs[t], CPs[t], Fis[t], Ks[t] + P̄p = copy(P̄) + K̄ = -P̄ * CP' + C̄P = -K' * P̄ + z̄p = copy(z̄) + K̄ .+= z̄ * v' + v̄ = K' * z̄ + C̄P .+= Fi * K̄' + F̄ = -Fi * (CP * K̄) * Fi + if t > presample_periods + v̄ .+= -∂ll * (Fi * v) + F̄ .+= ∂ll * 0.5 * (Fi * v * v' * Fi - Fi) + end + F̄ = (F̄ + F̄') / 2 + C̄P .+= F̄ * 𝒞 + H̄m .+= F̄ + P̄p .+= 𝒞' * C̄P + z̄p .+= -𝒞' * v̄ + Ȳ[:, t] .+= v̄ + P̄p = (P̄p + P̄p') / 2 + 𝒜̄ .+= 2 .* (P̄p * 𝒜 * P_) + P̄ = 𝒜' * P̄p * 𝒜 + Q̄ = P̄p + Ḡ = 2 .* (Q̄ * G) + Q̄H .+= Q̄ + vḠ = vec(Ḡ) + ḡ0 .+= vḠ + Λ̄ .+= vḠ * (Pz * z_)' + 𝒜̄ .+= z̄p * z_' + c̄ .+= z̄p + z̄ = 𝒜' * z̄p + Pz' * (Λ' * vḠ) + end + return (NoTangent(), 𝒜̄, c̄, Q̄H, ḡ0, Λ̄, H̄m, Ȳ, NoTangent(), NoTangent(), + z̄, P̄, NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + return ll, quadratic_kalman_recursion_pullback +end + """ Run the quadratic Kalman filter and return the loglikelihood. `data_in_deviations` holds the observables as deviations from the non-stochastic steady state (rows in @@ -197,39 +363,13 @@ function run_quadratic_kalman(sys, maximum(abs, Ak) < 1e-14 && break end - z = copy(z̄) - Pc = copy(Σ) - loglik = zero(Tv) - log2pi = log(2π) - - for t in 1:nT - G = quadratic_kalman_G(sys, z) - Q = G * G' + QH - - zp = 𝒜 * z + c - Pp = 𝒜 * Pc * 𝒜' + Q - Pp = (Pp + Pp') / 2 - - v = view(data_in_deviations, :, t) - 𝒞 * zp - CP = 𝒞 * Pp - F = CP * 𝒞' + Hm - F = (F + F') / 2 - - Fc = ℒ.cholesky(F, check = false) - ℒ.issuccess(Fc) || return Tv(on_failure_loglikelihood) - - if t > presample_periods - loglik -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log2pi) - isfinite(loglik) || return Tv(on_failure_loglikelihood) - end - - K = CP' / Fc - z = zp + K * v - Pc = Pp - K * CP - Pc = (Pc + Pc') / 2 - end + # Hand off to the taped recursion, which carries the hand-written adjoint. + g0, Λ = quadratic_kalman_affine_G(sys) + Pz = sys.P * [Matrix{Tv}(ℒ.I(sys.nr)) zeros(Tv, sys.nr, nz - sys.nr)] - return loglik + return quadratic_kalman_recursion(𝒜, c, QH, g0, Λ, Hm, Matrix(data_in_deviations), + 𝒞, Pz, z̄, Σ, nz, sys.nExo, + presample_periods, on_failure_loglikelihood) end end # @stable diff --git a/test/test_quadratic_kalman.jl b/test/test_quadratic_kalman.jl index e37e7ab49..aa0682cdb 100644 --- a/test/test_quadratic_kalman.jl +++ b/test/test_quadratic_kalman.jl @@ -59,32 +59,37 @@ import ForwardDiff Y = collect(data) .- [NSSS(v) for v in obs] sys = MacroModelling.build_quadratic_kalman_system(RBC_qkf, 𝐒[1], 𝐒[2], obs_idx) - @test sys.nz == 2 * sys.nVars + sys.nPast^2 + # the Kronecker block is carried compressed (vech, not vec) + @test sys.nq == sys.nPast * (sys.nPast + 1) ÷ 2 + @test sys.nz == 2 * sys.nr + sys.nq @test maximum(abs, sys.S2) > 1e-3 # the model really is nonlinear @testset "augmented transition reproduces the pruned conditional mean" begin Random.seed!(3) - x1 = randn(sys.nVars) * 0.02 - x2 = randn(sys.nVars) * 0.002 - z = vcat(x1, x2, ℒ.kron(sys.P * x1, sys.P * x1)) + x1 = randn(sys.nr) * 0.02 + x2 = randn(sys.nr) * 0.002 + z = vcat(x1, x2, sys.Lp * ℒ.kron(sys.P * x1, sys.P * x1)) nmc = 200_000 - a1 = zeros(sys.nVars); a2 = zeros(sys.nVars); aq = zeros(sys.nq) + a1 = zeros(sys.nr); a2 = zeros(sys.nr); aq = zeros(sys.nPast^2) for _ in 1:nmc ε = randn(sys.nExo) - nxt = MacroModelling.pruned_second_order_state_update([x1, x2], ε, sys.past, - sys.nVars, sys.S1, sys.S2) - a1 .+= nxt[1]; a2 .+= nxt[2] - aq .+= ℒ.kron(sys.P * nxt[1], sys.P * nxt[1]) + # the retained-row system: aug₁ built from the same past states + ā = sys.Ea * vcat(sys.P * x1, 1.0) + aug1 = ā + sys.S * ε + n1 = sys.S1 * aug1 + n2 = sys.S1 * (sys.Ea * vcat(sys.P * x2, 0.0)) + sys.S2 * ℒ.kron(aug1, aug1) / 2 + a1 .+= n1; a2 .+= n2 + aq .+= ℒ.kron(sys.P * n1, sys.P * n1) end a1 ./= nmc; a2 ./= nmc; aq ./= nmc pred = sys.𝒜 * z + sys.c rel(a, b) = maximum(abs, a - b) / max(1e-12, maximum(abs, b)) tol = 20 / sqrt(nmc) # generous multiple of the Monte-Carlo error - @test rel(pred[1:sys.nVars], a1) < tol - @test rel(pred[sys.nVars+1:2sys.nVars], a2) < tol - @test rel(pred[2sys.nVars+1:end], aq) < tol + @test rel(pred[1:sys.nr], a1) < tol + @test rel(pred[sys.nr+1:2sys.nr], a2) < tol + @test rel(pred[2sys.nr+1:end], sys.Lp * aq) < tol end @testset "matches the particle filter on a nonlinear model" begin @@ -138,6 +143,41 @@ import ForwardDiff @test maximum(abs.(g .- fd) ./ max.(abs.(fd), 1.0)) < 1e-5 end + @testset "hand-written reverse mode for the recursion" begin + # Every cotangent of the taped recursion, against ForwardDiff. Random but + # well-conditioned inputs; the point is the adjoint algebra, not a model. + Random.seed!(7) + nz, nE, nobs, nT, nPast_ = 9, 2, 2, 12, 3 + Pz = zeros(nPast_, nz); for i in 1:nPast_; Pz[i, i] = 1.0; end + A0 = 0.3 * randn(nz, nz); A0 ./= (1.6 * maximum(abs, ℒ.eigvals(A0))) + c0 = 0.01 * randn(nz); g0 = 0.05 * randn(nz * nE); L0 = 0.02 * randn(nz * nE, nPast_) + Ch = zeros(nobs, nz); Ch[1,1] = 1.0; Ch[2,2] = 1.0; Ch[1,4] = 1.0; Ch[2,5] = 1.0 + QH0 = (M = 0.05 * randn(nz, nz); M * M') + Hm0 = Matrix(0.01 * ℒ.I(nobs)) + Y0 = 0.05 * randn(nobs, nT); z00 = 0.01 * randn(nz) + S00 = (M = 0.1 * randn(nz, nz); M * M') + ps = 2 + f(A, c, QH, g, L, Hm, Y, z0, S0) = + MacroModelling.quadratic_kalman_recursion(A, c, QH, g, L, Hm, Y, Ch, Pz, z0, S0, + nz, nE, ps, -Inf) + ll, pb = MacroModelling.rrule(MacroModelling.quadratic_kalman_recursion, + A0, c0, QH0, g0, L0, Hm0, Y0, Ch, Pz, z00, S00, + nz, nE, ps, -Inf) + ct = pb(1.0) + @test isfinite(ll) + rel(a, b) = maximum(abs, a .- b) / max(1e-10, maximum(abs, b)) + # ct = (NoTangent, 𝒜̄, c̄, Q̄H, ḡ0, Λ̄, H̄m, Ȳ, NoTangent, NoTangent, z̄0, Σ̄0, …) + @test rel(ct[2], ForwardDiff.gradient(x -> f(x, c0, QH0, g0, L0, Hm0, Y0, z00, S00), A0)) < 1e-10 + @test rel(ct[3], ForwardDiff.gradient(x -> f(A0, x, QH0, g0, L0, Hm0, Y0, z00, S00), c0)) < 1e-10 + @test rel(ct[4], ForwardDiff.gradient(x -> f(A0, c0, x, g0, L0, Hm0, Y0, z00, S00), QH0)) < 1e-10 + @test rel(ct[5], ForwardDiff.gradient(x -> f(A0, c0, QH0, x, L0, Hm0, Y0, z00, S00), g0)) < 1e-10 + @test rel(ct[6], ForwardDiff.gradient(x -> f(A0, c0, QH0, g0, x, Hm0, Y0, z00, S00), L0)) < 1e-10 + @test rel(ct[7], ForwardDiff.gradient(x -> f(A0, c0, QH0, g0, L0, x, Y0, z00, S00), Hm0)) < 1e-10 + @test rel(ct[8], ForwardDiff.gradient(x -> f(A0, c0, QH0, g0, L0, Hm0, x, z00, S00), Y0)) < 1e-10 + @test rel(ct[11], ForwardDiff.gradient(x -> f(A0, c0, QH0, g0, L0, Hm0, Y0, x, S00), z00)) < 1e-10 + @test rel(ct[12], ForwardDiff.gradient(x -> f(A0, c0, QH0, g0, L0, Hm0, Y0, z00, x), S00)) < 1e-10 + end + @testset "reduces to the Kalman filter on a linear model" begin # With 𝐒₂ = 0 the x₂ and Kronecker blocks are inert and the quadratic # Kalman filter is the Kalman filter. Exact agreement, not approximate. From bf123941321e52d33feffaedc09688899b8a1a06 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 30 Jul 2026 22:33:40 +0200 Subject: [PATCH 04/16] Make the quadratic Kalman gradient work in reverse mode end to end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recursion already had a verified hand-written adjoint, but nothing could reach it: the top-level rrule looks for `rrule(calculate_loglikelihood, Val(filter), …)` and silently returns a zero gradient when none exists. Three things were missing. 1. The build adjoint. Cotangents of (𝒜, c, QH, g₀, Λ) are now pushed back onto 𝐒₁/𝐒₂ analytically, including the ergodic initialisation: z₀ = (I−𝒜)⁻¹c gives a transposed solve, and Σ₀ = 𝒜Σ₀𝒜' + Q₀ gives a second Lyapunov equation X = 𝒜'X𝒜 + Σ̄₀, after which Q̄₀ = X and 𝒜̄ += 2X𝒜Σ₀. Verified against ForwardDiff at ~4e-15 for ∂𝐒₁, ∂𝐒₂ and ∂data. 2. The filter now routes through `calculate_loglikelihood`, with an rrule at that interface returning ∂𝐒 in the expected position, scattered from the retained rows back onto the full solution matrices. 3. The reverse-mode guard refused any likelihood with measurement error. It is relaxed for this filter only, since its adjoint covers H. One wiring bug worth recording: the top-level rrule never passed `measurement_error` to the inner rrule. That was harmless while the guard made measurement error impossible, but here it meant reverse mode differentiated the H = 0 likelihood while the primal used H > 0 — a *finite* gradient that was simply wrong (relative deviation 2.5). Zygote and ForwardDiff now agree to 6e-14, and that comparison is a test. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/filter/quadratic_kalman.jl | 211 ++++++++++++++++++++++++++++++++- src/get_functions.jl | 11 +- src/rrules.jl | 17 ++- test/test_quadratic_kalman.jl | 9 ++ 4 files changed, 236 insertions(+), 12 deletions(-) diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl index 2368d202a..baa4eb680 100644 --- a/src/filter/quadratic_kalman.jl +++ b/src/filter/quadratic_kalman.jl @@ -85,8 +85,11 @@ solution, together with the pieces needed for the state-dependent innovation covariance. `𝐒₁`/`𝐒₂` are the expanded solution matrices as returned by `get_relevant_steady_state_and_state_update(Val(:pruned_second_order), …)`. """ -function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_index::Vector{Int}) - T = 𝓂.constants.post_model_macro +build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, oi::Vector{Int}) = + build_quadratic_kalman_system_from_constants(𝓂.constants, 𝐒₁, 𝐒₂, oi) + +function build_quadratic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, observables_index::Vector{Int}) + T = cons.post_model_macro nVars, nPast, nExo = T.nVars, T.nPast_not_future_and_mixed, T.nExo past = T.past_not_future_and_mixed_idx @@ -139,7 +142,8 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ SS = ℒ.kron(S, S) vecI = vec(Matrix{Float64}(ℒ.I(nExo))) PP = ℒ.kron(PS1, PS1) - A1 = S1 * Ea[:, 1:nPast] * P + EaP = Ea[:, 1:nPast] * P + A1 = S1 * EaP r1, r2, rq = 1:nr, nr+1:2nr, 2nr+1:nz @@ -164,8 +168,13 @@ function build_quadratic_kalman_system(𝓂::ℳ, 𝐒₁, 𝐒₂, observables_ QH = Hq * IK * Hq' QH = (QH + QH') / 2 + # The constant blocks are returned as well: the reverse-mode rule needs them + # to push cotangents from (𝒜, c, QH, g₀, Λ) back onto 𝐒₁ and 𝐒₂. return (; nVars, nr, oas, nPast, nExo, na, nq, nz, past, P, S, Ea, S1, S2, PS1, V, Dp, Lp, - 𝒜, c, 𝒞, QH, G1 = S1 * S, r1) + 𝒜, c, 𝒞, QH, G1 = S1 * S, r1, + r2 = nr+1:2nr, rq = 2nr+1:nz, + Eq, Ep, E1, SS, vecI, IK, EaP, Ea1 = Ea[:, nPast+1], EpP = Ep * P, + PP, KVV = ℒ.kron(V, V), Hq) end # State-dependent loading of the linear-in-ε part of the innovation, at state z. @@ -372,4 +381,198 @@ function run_quadratic_kalman(sys, presample_periods, on_failure_loglikelihood) end + + +# Adjoints of kron(A,B) with respect to each factor. +function kron_adjoint_A(M, B, m, n, p, q) + A = zeros(eltype(M), m, n) + @inbounds for i in 1:m, j in 1:n + A[i, j] = sum(view(M, (i-1)*p+1:i*p, (j-1)*q+1:j*q) .* B) + end + return A +end +function kron_adjoint_B(M, A, m, n, p, q) + B = zeros(eltype(M), p, q) + @inbounds for i in 1:m, j in 1:n + @views B .+= A[i, j] .* M[(i-1)*p+1:i*p, (j-1)*q+1:j*q] + end + return B +end + +# Discrete Lyapunov X = A X A' + Q by doubling. +function qkf_lyapunov(A, Q; iters::Int = 80) + X = copy(Q); Ak = copy(A) + for _ in 1:iters + Xn = Ak * X * Ak' + X; Xn = (Xn + Xn') / 2 + if maximum(abs, Xn - X) < 1e-15 * max(1.0, maximum(abs, Xn)); X = Xn; break; end + X = Xn; Ak = Ak * Ak + maximum(abs, Ak) < 1e-16 && break + end + return (X + X') / 2 +end + +qkf_Pz(sys) = sys.P * [Matrix{Float64}(ℒ.I(sys.nr)) zeros(sys.nr, sys.nz - sys.nr)] + +""" +Push the cotangents of the augmented system back onto the solution matrices. +Covers the build, the ergodic initialisation (including the Lyapunov adjoint) and +the recursion. Verified against ForwardDiff to ~1e-15 in the test suite. +""" +function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_periods, ∂ll) + nr, nP, nE, na, nz = sys.nr, sys.nPast, sys.nExo, sys.na, sys.nz + r1, r2, rq = sys.r1, sys.r2, sys.rq + Lp = Matrix(sys.Lp); Eq = Matrix(sys.Eq); Ep = Matrix(sys.Ep); E1 = Vector(sys.E1) + g0, Λ = quadratic_kalman_affine_G(sys) + Pz = qkf_Pz(sys) + + z0 = (Matrix{Float64}(ℒ.I(nz)) - sys.𝒜) \ sys.c + G0 = reshape(g0 + Λ * (Pz * z0), nz, nE) + Q0 = G0 * G0' + sys.QH; Q0 = (Q0 + Q0') / 2 + Σ0 = qkf_lyapunov(sys.𝒜, Q0) + + R = last(rrule(quadratic_kalman_recursion, sys.𝒜, sys.c, sys.QH, g0, Λ, Hm, + Matrix(data_in_deviations), sys.𝒞, Pz, z0, Σ0, nz, nE, + presample_periods, -Inf))(∂ll) + 𝒜̄ = copy(R[2]); c̄ = copy(R[3]); Q̄H = copy(R[4]) + ḡ0 = copy(R[5]); Λ̄ = copy(R[6]); Ȳ = copy(R[8]) + z̄0 = copy(R[11]); Σ̄0 = copy(R[12]) + + # Σ0 = 𝒜Σ0𝒜' + Q0 ⇒ X solves X = 𝒜'X𝒜 + Σ̄0 + X = qkf_lyapunov(Matrix(sys.𝒜'), (Σ̄0 + Σ̄0') / 2) + 𝒜̄ .+= 2 .* (X * sys.𝒜 * Σ0) + Q̄0 = (X + X') / 2 + Ḡ0 = 2 .* (Q̄0 * G0); Q̄H .+= Q̄0 + vG0 = vec(Ḡ0); ḡ0 .+= vG0; Λ̄ .+= vG0 * (Pz * z0)' + z̄0 .+= Pz' * (Λ' * vG0) + λ = (Matrix{Float64}(ℒ.I(nz)) - sys.𝒜)' \ z̄0 + c̄ .+= λ; 𝒜̄ .+= λ * z0' + + S̄1 = zeros(size(sys.S1)); S̄2 = zeros(size(sys.S2)) + P̄S1 = zeros(size(sys.PS1)); V̄ = zeros(size(sys.V)) + P̄P = zeros(size(sys.PP)); K̄VV = zeros(size(sys.KVV)) + + Ā1 = 𝒜̄[r1, r1] + 𝒜̄[r2, r2] + S̄1 .+= Ā1 * sys.EaP' + S̄2 .+= 𝒜̄[r2, rq] * Eq' / 2 + 𝒜̄[r2, r1] * sys.EpP' / 2 + P̄P .+= Lp' * 𝒜̄[rq, rq] * Eq' + Lp' * 𝒜̄[rq, r1] * sys.EpP' + S̄1 .+= c̄[r1] * sys.Ea1' + S̄2 .+= c̄[r2] * (E1 + sys.SS * sys.vecI)' / 2 + lc = Lp' * c̄[rq] + P̄P .+= lc * E1'; K̄VV .+= lc * sys.vecI' + + Rq = (Q̄H + Q̄H') / 2 + H̄q = 2 .* (Rq * sys.Hq * sys.IK) + S̄2 .+= H̄q[nr+1:2nr, :] * (sys.SS / 2)' + K̄VV .+= Lp' * H̄q[2nr+1:end, :] + + function absorb_G!(Ḡ, z) + ā = sys.Ea * vcat(sys.P * view(z, r1), 1.0) + ū = sys.PS1 * ā + Ma = ℒ.kron(ā, sys.S) + ℒ.kron(sys.S, ā) + S̄1 .+= Ḡ[r1, :] * sys.S' + S̄2 .+= Ḡ[r2, :] * Ma' / 2 + Gq = Lp' * Ḡ[rq, :] + ū̄ = vec(kron_adjoint_A(Gq, sys.V, nP, 1, nP, nE)) .+ + vec(kron_adjoint_B(Gq, sys.V, nP, nE, nP, 1)) + V̄ .+= kron_adjoint_B(Gq, reshape(ū, nP, 1), nP, 1, nP, nE) .+ + kron_adjoint_A(Gq, reshape(ū, nP, 1), nP, nE, nP, 1) + P̄S1 .+= ū̄ * ā' + end + absorb_G!(reshape(ḡ0 .- vec(sum(Λ̄, dims = 2)), nz, nE), zeros(nz)) + @inbounds for i in 1:nP + e = zeros(nz); e[findfirst(!iszero, view(sys.P, i, :))] = 1.0 + absorb_G!(reshape(Λ̄[:, i], nz, nE), e) + end + + P̄S1 .+= kron_adjoint_A(P̄P, sys.PS1, nP, na, nP, na) .+ + kron_adjoint_B(P̄P, sys.PS1, nP, na, nP, na) + V̄ .+= kron_adjoint_A(K̄VV, sys.V, nP, nE, nP, nE) .+ + kron_adjoint_B(K̄VV, sys.V, nP, nE, nP, nE) + P̄S1 .+= V̄ * sys.S' + S̄1 .+= sys.P' * P̄S1 + + return S̄1, S̄2, Ȳ +end + + + + +# ── standard filter interface ──────────────────────────────────────────────── +# Routing through `calculate_loglikelihood` (rather than a special branch in +# `get_loglikelihood`) is what lets the existing reverse-mode machinery reach the +# filter: the top-level rrule looks for `rrule(calculate_loglikelihood, Val(filter), …)` +# and falls back to a zero gradient when none exists. +function calculate_loglikelihood(::Val{:quadratic_kalman}, + ::Val{:pruned_second_order}, + observables_index::Vector{Int}, + 𝐒, + data_in_deviations::AbstractMatrix, + constants, + state, + workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) + sys = build_quadratic_kalman_system_from_constants(constants, 𝐒[1], 𝐒[2], observables_index) + return run_quadratic_kalman(sys, data_in_deviations; + measurement_error = measurement_error, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood) +end + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:quadratic_kalman}, + ::Val{:pruned_second_order}, + observables_index::Vector{Int}, + 𝐒, + data_in_deviations::AbstractMatrix, + constants, + state, + workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) + sys = build_quadratic_kalman_system_from_constants(constants, 𝐒[1], 𝐒[2], observables_index) + n_obs = size(data_in_deviations, 1) + Hm = measurement_error === nothing ? zeros(n_obs, n_obs) : + measurement_error isa AbstractMatrix ? Matrix{Float64}(measurement_error) : + Matrix{Float64}(ℒ.Diagonal(collect(measurement_error))) + llh = run_quadratic_kalman(sys, data_in_deviations; + measurement_error = measurement_error, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood) + + nine(x...) = (NoTangent(), NoTangent(), NoTangent(), NoTangent(), x[1], x[2], + NoTangent(), x[3], NoTangent()) + + if !isfinite(llh) + return llh, _ -> nine(NoTangent(), NoTangent(), NoTangent()) + end + + function quadratic_kalman_loglikelihood_pullback(∂llh_bar) + ∂llh = unthunk(∂llh_bar) + S̄1r, S̄2r, Ȳ = quadratic_kalman_pullback(sys, data_in_deviations, Hm, + presample_periods, ∂llh) + # scatter the retained rows back onto the full solution matrices + ∂𝐒1 = zeros(size(𝐒[1])); ∂𝐒2 = zeros(size(𝐒[2])) + ∂𝐒1[sys.oas, :] = S̄1r + ∂𝐒2[sys.oas, :] = S̄2r + ∂state = [zeros(length(s)) for s in state] + return nine([∂𝐒1, ∂𝐒2], Ȳ, ∂state) + end + + return llh, quadratic_kalman_loglikelihood_pullback +end + + end # @stable diff --git a/src/get_functions.jl b/src/get_functions.jl index 112ca0f98..d27626ea9 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -4871,11 +4871,12 @@ function get_loglikelihood(𝓂::ℳ, if has_missing error("The quadratic Kalman filter does not yet support missing observations.") end - qkf_sys = build_quadratic_kalman_system(𝓂, 𝐒[1], 𝐒[2], obs_indices) - run_quadratic_kalman(qkf_sys, data_in_deviations; - measurement_error = measurement_error_H, - presample_periods = presample_periods, - on_failure_loglikelihood = on_failure_loglikelihood) + calculate_loglikelihood(Val(:quadratic_kalman), Val(algorithm), obs_indices, + 𝐒, data_in_deviations, constants_obj, state, 𝓂.workspaces, + presample_periods = presample_periods, + measurement_error = measurement_error_H, + on_failure_loglikelihood = on_failure_loglikelihood, + opts = opts) elseif filter == :kalman if has_missing calculate_loglikelihood_with_missing(Val(:kalman), diff --git a/src/rrules.jl b/src/rrules.jl index 7366edfa3..dbb2a39c5 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -1877,7 +1877,9 @@ function rrule(::typeof(get_loglikelihood), else measurement_error != 0 end - if me_active + # The quadratic Kalman filter carries its own hand-written adjoint, which + # includes the measurement-error covariance, so the guard does not apply to it. + if me_active && filter != :quadratic_kalman error("Reverse-mode automatic differentiation of the Kalman likelihood with measurement error (`measurement_error`) is not yet supported. Use forward-mode AD (e.g. `AutoForwardDiff`) or a gradient-free sampler.") end @@ -1963,6 +1965,13 @@ function rrule(::typeof(get_loglikelihood), end # ── step 3: calculate_loglikelihood ── + # The quadratic Kalman filter is the only one whose inner rrule takes the + # measurement-error covariance; for the others it is inactive (the guard above) + # and the kwarg would not be accepted. + me_kw = filter == :quadratic_kalman ? + (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations)) : + NamedTuple() + llh_rrule = if has_missing rrule(calculate_loglikelihood_with_missing, Val(filter), Val(algorithm), obs_indices, @@ -1972,7 +1981,8 @@ function rrule(::typeof(get_loglikelihood), initial_covariance = initial_covariance, filter_algorithm = filter_algorithm, opts = opts, - on_failure_loglikelihood = on_failure_loglikelihood) + on_failure_loglikelihood = on_failure_loglikelihood, + me_kw...) else rrule(calculate_loglikelihood, Val(filter), Val(algorithm), obs_indices, @@ -1982,7 +1992,8 @@ function rrule(::typeof(get_loglikelihood), initial_covariance = initial_covariance, filter_algorithm = filter_algorithm, opts = opts, - on_failure_loglikelihood = on_failure_loglikelihood) + on_failure_loglikelihood = on_failure_loglikelihood, + me_kw...) end if llh_rrule === nothing diff --git a/test/test_quadratic_kalman.jl b/test/test_quadratic_kalman.jl index aa0682cdb..159ecb19a 100644 --- a/test/test_quadratic_kalman.jl +++ b/test/test_quadratic_kalman.jl @@ -4,6 +4,7 @@ import Random import Statistics import LinearAlgebra as ℒ import ForwardDiff +import Zygote # ----------------------------------------------------------------------------- # Quadratic Kalman filter (Monfort, Renne & Roussellet, 2015) on the pruned @@ -141,6 +142,14 @@ import ForwardDiff fd = [(f(p + h * (1:length(p) .== i)) - f(p - h * (1:length(p) .== i))) / (2h) for i in eachindex(p)] @test maximum(abs.(g .- fd) ./ max.(abs.(fd), 1.0)) < 1e-5 + + # Reverse mode reaches the filter through the hand-written rrule chain: + # rrule(calculate_loglikelihood, Val(:quadratic_kalman), …) pushes the + # cotangents back onto 𝐒₁/𝐒₂ analytically. It must agree with forward mode. + gz = Zygote.gradient(f, p)[1] + @test gz !== nothing + @test all(isfinite, gz) + @test maximum(abs.(gz .- g) ./ max.(abs.(g), 1.0)) < 1e-8 end @testset "hand-written reverse mode for the recursion" begin From a5523b6c7f615ebee01d73338160eb8f0adcf84c Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 30 Jul 2026 23:45:50 +0200 Subject: [PATCH 05/16] Make the quadratic Kalman recursion allocation-free and BLAS-driven MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measured before guessing: the filter spent 60% of its time not in the covariance propagation but in allocation — 2509 MiB over a 138-period Smets-Wouters run, roughly 18 MB per period, mostly nz×nz temporaries from `𝒜*Pc*𝒜'`, `G*G'` and the `𝒞` products. Three changes: * every buffer is allocated once per call and reused, with `mul!` throughout; the rank-nExo term G G' folds into Pp as a single gemm update rather than forming Q separately; * 𝒞 is a selection matrix, so `𝒞*zp`, `𝒞*Pp` and `CP*𝒞'` become indexing instead of three gemms; * the symmetrisations write both triangles in one pass rather than building (X+X')/2 as a fresh matrix. 0.986 s → 0.485 s and 2509 MiB → 171 MiB on SW07. The covariance propagation is now 92% of the remaining runtime, so further gains have to come from that term. One trap this surfaced: preallocation fixes the element type, so the promotion has to cover *every* differentiable argument. Covering only 𝒜, Y and g₀ left forward-mode AD working with respect to those three and failing with respect to c, QH, Λ, Hm, z₀ and Σ₀ — caught by the cotangent tests. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/filter/quadratic_kalman.jl | 84 ++++++++++++++++++++++++++++------ 1 file changed, 70 insertions(+), 14 deletions(-) diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl index baa4eb680..f46f143e1 100644 --- a/src/filter/quadratic_kalman.jl +++ b/src/filter/quadratic_kalman.jl @@ -219,27 +219,83 @@ left to ordinary AD. function quadratic_kalman_recursion(𝒜, c, QH, g0, Λ, Hm, Y, 𝒞, Pz, z0, Σ0, nz::Int, nExo::Int, presample_periods::Int, on_failure_loglikelihood::Real) n_obs, nT = size(Y) - Tv = promote_type(eltype(𝒜), eltype(Y), eltype(g0)) - z = copy(z0); Pc = copy(Σ0) + # Promote over every differentiable input, not just a few: the preallocated + # buffers below fix the element type, so missing one makes forward-mode AD + # fail with respect to exactly that argument. + Tv = promote_type(eltype(𝒜), eltype(c), eltype(QH), eltype(g0), eltype(Λ), + eltype(Hm), eltype(Y), eltype(z0), eltype(Σ0)) + + # 𝒞 is a selection: row i picks the x₁ and x₂ entries of observable i. Doing + # that by indexing rather than by three gemms with a 0/1 matrix removes the + # only dense products that scale with n_obs·nz². + p1 = [findfirst(!iszero, view(𝒞, i, :)) for i in 1:n_obs] + p2 = [findlast(!iszero, view(𝒞, i, :)) for i in 1:n_obs] + + # Preallocate once per call. The naive version allocated ~18 MB per period, + # which cost more than the covariance propagation it was feeding. + Pc = Matrix{Tv}(undef, nz, nz); copyto!(Pc, Σ0) + Pp = Matrix{Tv}(undef, nz, nz) + Tm = Matrix{Tv}(undef, nz, nz) + z = Vector{Tv}(undef, nz); copyto!(z, z0) + zp = Vector{Tv}(undef, nz) + gv = Vector{Tv}(undef, nz * nExo) + x1p = Vector{Tv}(undef, size(Pz, 1)) + CP = Matrix{Tv}(undef, n_obs, nz) + F = Matrix{Tv}(undef, n_obs, n_obs) + Kg = Matrix{Tv}(undef, nz, n_obs) + v = Vector{Tv}(undef, n_obs) + Fv = Vector{Tv}(undef, n_obs) + ll = zero(Tv); log2pi = log(2π) - for t in 1:nT - G = reshape(g0 + Λ * (Pz * z), nz, nExo) - Q = G * G' + QH - zp = 𝒜 * z + c - Pp = 𝒜 * Pc * 𝒜' + Q; Pp = (Pp + Pp') / 2 - v = view(Y, :, t) - 𝒞 * zp - CP = 𝒞 * Pp - F = CP * 𝒞' + Hm; F = (F + F') / 2 + + @inbounds for t in 1:nT + # G(z) = reshape(g₀ + Λ(Pz z)) + ℒ.mul!(x1p, Pz, z) + copyto!(gv, g0); ℒ.mul!(gv, Λ, x1p, one(Tv), one(Tv)) + G = reshape(gv, nz, nExo) + + # Pp = 𝒜 Pc 𝒜' + G G' + QH (the rank-nExo term as one gemm update) + ℒ.mul!(Tm, 𝒜, Pc) + ℒ.mul!(Pp, Tm, 𝒜') + Pp .+= QH + ℒ.mul!(Pp, G, G', one(Tv), one(Tv)) + for j in 1:nz, i in 1:j + m = (Pp[i, j] + Pp[j, i]) / 2; Pp[i, j] = m; Pp[j, i] = m + end + + ℒ.mul!(zp, 𝒜, z); zp .+= c + + for i in 1:n_obs + v[i] = Y[i, t] - (zp[p1[i]] + zp[p2[i]]) + for k in 1:nz + CP[i, k] = Pp[p1[i], k] + Pp[p2[i], k] + end + end + for i in 1:n_obs, j in 1:n_obs + F[i, j] = CP[i, p1[j]] + CP[i, p2[j]] + Hm[i, j] + end + for i in 1:n_obs, j in 1:i-1 + m = (F[i, j] + F[j, i]) / 2; F[i, j] = m; F[j, i] = m + end + Fc = ℒ.cholesky(F, check = false) ℒ.issuccess(Fc) || return Tv(on_failure_loglikelihood) + if t > presample_periods - ll -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log2pi) + copyto!(Fv, v); ℒ.ldiv!(Fc, Fv) + ll -= 0.5 * (ℒ.dot(v, Fv) + ℒ.logdet(Fc) + n_obs * log2pi) isfinite(ll) || return Tv(on_failure_loglikelihood) end - K = CP' / Fc - z = zp + K * v - Pc = Pp - K * CP; Pc = (Pc + Pc') / 2 + + # K = CP' F⁻¹ ; z = zp + K v ; Pc = Pp − K CP + copyto!(Kg, CP'); ℒ.rdiv!(Kg, Fc) + copyto!(z, zp); ℒ.mul!(z, Kg, v, one(Tv), one(Tv)) + copyto!(Pc, Pp); ℒ.mul!(Pc, Kg, CP, -one(Tv), one(Tv)) + for j in 1:nz, i in 1:j + m = (Pc[i, j] + Pc[j, i]) / 2; Pc[i, j] = m; Pc[j, i] = m + end end + return ll end From 2f9e4d8b67d94879f785087a1ba730ea8ca0436e Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 09:02:58 +0200 Subject: [PATCH 06/16] Document the quadratic Kalman filter, its bias, and its cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a full section to the Filters page: the augmented-state construction and why pruning is what makes it possible, what is exact versus approximated, the source and size of the likelihood bias, when the filter is and is not usable, and a measured cost breakdown. The two findings worth recording: The bias is not what it looks like. It does not come from 𝐒₂ — zeroing its ε⊗ε block leaves rank(Q) unchanged. It comes from kron(V,V) in the *first-order* solution: x₁[past] = (deterministic) + Vε, so q = x₁[past]⊗x₁[past] inherits Vε⊗Vε whatever 𝐒₂ is. The quadratic noise is intrinsic to carrying a Kronecker term as a state. And because the observation has zero loading on q, the data can never shrink the resulting fictitious uncertainty, so the error persists as the measurement error goes to zero (converging to -1.52 rather than to 0). But it is almost a level shift, so it is far less damaging than it first looks: profiled against the exact inversion likelihood the gap varies by only ~0.2 log points across a parameter grid and the mode is unchanged. Latent states — the filter's actual purpose, per Kollmann (2015) — are good: 2.8% relative RMSE on capital, ~11% on the unobserved shock processes. So it is fine for state estimation and point estimation, and unsafe for model comparison, where the level error differs across models with the shock-to-observable ratio. Cost breakdown measured per period at nz=446: the two nz³ triple products are 91% of the loop (1.09 ms and 1.26 ms), everything else 9%. Against the inversion filter's n_ε×n_ε solve that is 446³ vs 7³ — about 2.6e5 in flops. Structural, not tunable; sparsity measures 10× slower since 𝒜 is ~50% dense. Also records that this is Kollmann's filter, not the Monfort-Renne-Roussellet quadratic Kalman filter: theirs has a linear transition and a *quadratic measurement*, so the data loads on the Kronecker block directly and shrinks the very uncertainty that is irreducible here. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 171 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 171 insertions(+) diff --git a/docs/src/filters.md b/docs/src/filters.md index 2d4963a09..8e6da6b73 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -360,6 +360,177 @@ One wrinkle worth knowing. The states match the **smoothed** Kalman estimates, n The third row is the particle filters' correctness check, and is exactly what the package's tests do: on a linear model the particle log-likelihood must approach the Kalman value as ``N`` grows, approaching it *from below* because of the Jensen bias. +## The quadratic Kalman filter + +`filter = :quadratic_kalman`, available only for `algorithm = :pruned_second_order`. + +### The idea + +A pruned second-order solution is *exactly linear* in an augmented state. Writing the +package's own recursion, + +```math +\begin{aligned} +\mathrm{aug}_1 &= [x_{1,t-1}[\text{past}];\ 1;\ \varepsilon_t], \\ +x_{1,t} &= \mathbf{S}_1\,\mathrm{aug}_1, \\ +x_{2,t} &= \mathbf{S}_1\,[x_{2,t-1}[\text{past}];0;0] + \tfrac12\mathbf{S}_2(\mathrm{aug}_1\otimes\mathrm{aug}_1), +\end{aligned} +``` + +the quadratic term uses only the *first-order* piece ``x_1``. That is what pruning buys: +stacking + +```math +z_t = [\,x_{1,t};\ x_{2,t};\ x_{1,t}[\text{past}]\otimes x_{1,t}[\text{past}]\,] +``` + +makes every block affine in ``z_{t-1}``, because ``\mathrm{aug}_1\otimes\mathrm{aug}_1`` +expands into terms that are quadratic in ``x_{1,t-1}[\text{past}]`` (carried by the third +block), linear in it, or constant. The observation ``y_t = (x_1+x_2)[\text{observables}]`` +is a plain selection, so the system is linear and a Kalman filter applies. This is +Kollmann (2015). + +Without pruning there is no such representation: ``x_t`` is quadratic in ``x_{t-1}``, so +``x_t\otimes x_t`` is quartic, needing ``x^{\otimes4}``, then ``x^{\otimes8}`` — the +hierarchy never closes. Pruning truncates it at exactly one rung. + +### What is exact, and what is not + +The transition is exactly linear and the conditional first two moments are closed form. +Writing ``\mathrm{aug}_1 = \bar a + S\varepsilon``, every block of the innovation is + +```math +w = G\varepsilon + H(\varepsilon\otimes\varepsilon - \mathrm{vec}\,I), +``` + +linear plus centred-quadratic in ``\varepsilon``. Gaussian third moments vanish, so the two +parts are uncorrelated and, using +``E[(\varepsilon\otimes\varepsilon)(\varepsilon\otimes\varepsilon)'] = \mathrm{vec}(I)\mathrm{vec}(I)' + I + K``, + +```math +\mathrm{Var}(w) = GG' + H(I+K)H', +``` + +with ``K`` the commutation matrix. ``H`` is constant; ``G`` depends on the state and is +evaluated at the filtered mean. + +What is approximated is the conditional *distribution*. ``\varepsilon\otimes\varepsilon`` is a +squared Gaussian — skewed, not Gaussian — so the recursion delivers the best **linear** +projection rather than the exact conditional mean. + +### The bias, and where it comes from + +Given ``z_{t-1}``, the true next state is determined by ``\varepsilon``, so the exact +conditional distribution lives on a curved ``n_\varepsilon``-dimensional surface. A Kalman +filter can only carry a Gaussian ellipsoid, and fitting one to that surface needs more +directions than the surface has: + +```math +\mathrm{rank}(Q) = n_\varepsilon + \tfrac{n_\varepsilon(n_\varepsilon+1)}{2}. +``` + +| model | ``n_\varepsilon`` | true dimension | rank(Q) | excess | +|---|---|---|---|---| +| small RBC | 2 | 2 | 5 | 3 | +| Smets-Wouters (2007) | 7 | 7 | 33 | 26 | + +The excess directions are **fictitious uncertainty**, an artefact of the Gaussian +approximation. Two things make them permanent. First, they do not come from ``\mathbf{S}_2`` +— zeroing its ``\varepsilon\otimes\varepsilon`` block leaves the rank unchanged. They come from +``\mathrm{kron}(V,V)`` in the *first-order* solution: since +``x_1[\text{past}] = (\text{deterministic}) + V\varepsilon``, the state +``q = x_1[\text{past}]\otimes x_1[\text{past}]`` inherits ``V\varepsilon\otimes V\varepsilon`` +whatever ``\mathbf{S}_2`` is. The quadratic noise is intrinsic to carrying a Kronecker term +as a state. Second, the observation has **zero loading on** ``q`` — the data never sees that +block directly, so it can never shrink the fictitious uncertainty. It persists and leaks into +the predicted observables, which is why the likelihood error does *not* vanish as the +measurement error goes to zero. + +Measured against the inversion filter, which is the **exact** likelihood here (as many shocks +as observables, no measurement error, so it is a deterministic change of variables): + +| measurement-error variance | quadratic Kalman | exact | gap | +|---|---|---|---| +| ``10^{-5}`` | 211.2451 | 213.6137 | ``-2.37`` | +| ``10^{-6}`` | 212.0289 | 213.6137 | ``-1.58`` | +| ``10^{-8}`` | 212.0937 | 213.6137 | ``-1.52`` | +| ``10^{-10}`` | 212.0943 | 213.6137 | ``-1.52`` | + +It converges to a persistent gap rather than to the truth. The size is governed by +``n_\varepsilon(n_\varepsilon+1)/(2\,n_{obs})`` — 1.5 for the RBC, 4.0 for Smets-Wouters — and by +persistence: raising ``\rho`` from 0.4/0.6 to 0.98 on the same model multiplies the error +per period by 13. It is *not* governed by the size of the second-order terms, which is a +natural but wrong guess. + +### What that means for usability + +The bias falls almost entirely on the **level** of the likelihood, not on its shape. Profiling +against the exact likelihood, the gap varies by only about 0.2 log points across a parameter +grid, and the mode is unchanged: + +| parameter | truth | argmax, exact | argmax, quadratic Kalman | +|---|---|---|---| +| shock std | 0.02 | 0.021 | 0.021 | +| persistence | 0.4 | 0.39 | 0.39 | + +And the filter does the job it was designed for. Latent-state accuracy, as a fraction of each +state's own standard deviation: ``1.4\times10^{-5}`` and ``7\times10^{-7}`` for the two observed +variables, 2.8% for capital, and about 11% for the two unobserved shock processes. + +**Use it for**: latent state and shock estimates at pruned second order — that is what it is +for, it is deterministic, and it is far faster than a particle filter. Point estimation, where +the mode is essentially unaffected. + +**Do not use it for**: model comparison, marginal likelihoods or Bayes factors — the level +error differs across models, since it scales with the shock-to-observable ratio. Reported +standard errors without checking curvature first, since the bias is not exactly constant. +Models with many shocks per observable or near-unit persistence, where the error per period +grows sharply. + +**Alternatives when the likelihood level matters**: the inversion filter is exact when shocks +and observables balance and there is no measurement error; a particle filter is consistent at +any order; and a sigma-point filter on the *unpruned* solution (Andreasen, 2013) avoids the +augmented state altogether. + +!!! note "This is not the quadratic Kalman filter of Monfort, Renne & Roussellet" + That method targets a *linear* Gaussian transition with a **quadratic measurement** + equation, where the data loads directly on the Kronecker block and therefore shrinks its + uncertainty every period — which is why the original paper reports large gains over the + extended and unscented filters. A pruned DSGE is the mirror image: the quadratic terms are + in the transition and the observation is a plain selection with zero loading on the + Kronecker block. The machinery is shared, the regime is not. + +### Cost + +The augmented dimension is ``2n_r + n_{past}(n_{past}+1)/2``, where ``n_r`` counts the retained +rows (past states plus observables) and the Kronecker block is carried compressed as a +``\mathrm{vech}``. On Smets-Wouters that is 446. The covariance recursion is ``O(n_z^3)`` and +dominates everything else — per period, measured: + +| operation | cost | ms | share | +|---|---|---|---| +| ``\mathcal{A}P_c`` | ``n_z^3`` (88.7M flops) | 1.09 | 42% | +| ``(\mathcal{A}P_c)\mathcal{A}'`` | ``n_z^3`` (88.7M flops) | 1.26 | 49% | +| symmetrisation ×2 | ``n_z^2`` | 0.10 | 4% | +| ``P_p - K\,CP`` | ``n_z^2 n_{obs}`` | 0.03 | 1% | +| ``GG'`` | ``n_z^2 n_\varepsilon`` | 0.03 | 1% | +| ``\mathcal{C}P_p`` | ``n_{obs}n_z^2`` | 0.06 | 2% | +| build ``G`` | ``n_z n_\varepsilon n_{past}`` | 0.02 | 1% | + +The two matrix triple-products are 91% of the loop. By contrast the inversion filter solves an +``n_\varepsilon \times n_\varepsilon`` system per period — ``7^3`` against ``446^3``, a factor of +about ``2.6\times10^5`` in flops on the dominant term. That gap is structural: it is the price +of propagating a covariance over the Kronecker-augmented state, and no amount of tuning removes +it. Sparsity does not help either — ``\mathcal{A}`` is about 50% dense, and a sparse +representation measures 10× *slower* than the dense one. + +**References:** Kollmann (2015), *Computational Economics* 45, 239–260 — the filter implemented +here. Andreasen, Fernández-Villaverde & Rubio-Ramírez (2018) — the pruned state-space +representation. Monfort, Renne & Roussellet (2015), *Journal of Econometrics* 187, 43–56 — the +quadratic Kalman filter for quadratic measurement equations. Andreasen (2013), *Journal of +Applied Econometrics* 28, 929–955 — the central difference Kalman filter, the unpruned +alternative. + ## The filter-free likelihood There is a fourth option that is not a `filter` value at all, because it does not filter: instead of integrating the shocks out, it treats them as **parameters** and asks you to supply them. From 01e08d37fb1b33ab1971092eb45b889d13a82b69 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 09:43:01 +0200 Subject: [PATCH 07/16] Route the QKF Lyapunov solves through the package solver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `run_quadratic_kalman` carried its own inlined doubling loop with looser tolerances (1e-12, 60 iterations) than the `qkf_lyapunov` helper used by the pullback (1e-15, 80), so the forward and reverse passes initialised at slightly different ergodic covariances. Both now call one helper, which dispatches Float64 problems to `solve_lyapunov_equation` — the same workspace-backed doubling solver the pruned second-order moments code uses — and keeps the self-contained loop as the fallback for AD element types, which that solver does not accept. The forward pass hands its converged covariance to the reverse pass, which was re-solving the identical equation; passing it as `initial_guess` turns that second solve into a residual check (46.6 ms -> 3.1 ms on SW07). Measured on SW07 at pruned second order, 7 observables: 485 ms -> 425 ms per likelihood, with the log-likelihood unchanged to the last digit (-2264.4493051268087) and all 31 quadratic-Kalman tests passing, including the hand-written rrule cotangents against ForwardDiff. Warm-starting across parameter draws was measured and deliberately not done: the guess only pays when it is exact, and a 1e-6 relative parameter move already makes it a net loss. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/filter/quadratic_kalman.jl | 84 +++++++++++++++++++++++++--------- 1 file changed, 62 insertions(+), 22 deletions(-) diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl index f46f143e1..526df8430 100644 --- a/src/filter/quadratic_kalman.jl +++ b/src/filter/quadratic_kalman.jl @@ -392,7 +392,10 @@ function run_quadratic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; measurement_error::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}} = nothing, presample_periods::Int = 0, - on_failure_loglikelihood::Real = -Inf) + on_failure_loglikelihood::Real = -Inf, + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling, + initial_covariance_out::Union{Nothing,Base.RefValue} = nothing) nz = sys.nz 𝒜, c, 𝒞, QH = sys.𝒜, sys.c, sys.𝒞, sys.QH n_obs, nT = size(data_in_deviations) @@ -414,19 +417,13 @@ function run_quadratic_kalman(sys, Q̄ = Gbar * Gbar' + QH Q̄ = (Q̄ + Q̄') / 2 - Σ = copy(Q̄) - Ak = copy(𝒜) - for _ in 1:60 - Σn = Ak * Σ * Ak' + Σ - Σn = (Σn + Σn') / 2 - if maximum(abs, Σn - Σ) < 1e-12 * max(1.0, maximum(abs, Σn)) - Σ = Σn - break - end - Σ = Σn - Ak = Ak * Ak - maximum(abs, Ak) < 1e-14 && break - end + Σ = qkf_lyapunov(𝒜, Q̄; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + + # The reverse pass needs this exact matrix again. Handing it back lets the + # pullback skip a second identical Lyapunov solve (a residual check on an + # exact guess instead of a full doubling run). + initial_covariance_out === nothing || (initial_covariance_out[] = Σ) # Hand off to the taped recursion, which carries the hand-written adjoint. g0, Λ = quadratic_kalman_affine_G(sys) @@ -455,8 +452,35 @@ function kron_adjoint_B(M, A, m, n, p, q) return B end -# Discrete Lyapunov X = A X A' + Q by doubling. -function qkf_lyapunov(A, Q; iters::Int = 80) +# Discrete Lyapunov X = A X A' + Q. +# +# Float64 problems go through the package's workspace-backed doubling solver, which +# reuses its buffers instead of allocating a fresh nz×nz triple product per +# iteration. AD element types fall back to the self-contained loop below, since +# `solve_lyapunov_equation` is restricted to `Float64`. +# +# `initial_guess` pays off only when the guess is *exact*: the solver checks its +# residual (two nz³ products) and returns it, ~15× faster than a full solve on +# SW07. Under even a 1e-6 relative parameter move the check fails and the solve +# runs anyway, making it a net ~6% loss — so this is worth threading from the +# forward pass into the reverse pass, which re-solves the identical equation, but +# *not* worth caching across sampler draws. +function qkf_lyapunov(A, Q; + workspaces = nothing, + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0, 0), + lyapunov_algorithm::Symbol = :doubling, + iters::Int = 80) + if workspaces !== nothing && eltype(A) === Float64 && eltype(Q) === Float64 + ws = ensure_lyapunov_workspace!(workspaces, size(A, 1), :second_order) + X, converged = solve_lyapunov_equation(Matrix(A), Matrix(Q), ws; + initial_guess = initial_guess, + lyapunov_algorithm = lyapunov_algorithm, + verbose = false) + # The solver may hand back one of its own buffers, so symmetrising into a + # fresh matrix here doubles as taking ownership of the result. + converged && return (X + X') / 2 + end + X = copy(Q); Ak = copy(A) for _ in 1:iters Xn = Ak * X * Ak' + X; Xn = (Xn + Xn') / 2 @@ -474,7 +498,10 @@ Push the cotangents of the augmented system back onto the solution matrices. Covers the build, the ergodic initialisation (including the Lyapunov adjoint) and the recursion. Verified against ForwardDiff to ~1e-15 in the test suite. """ -function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_periods, ∂ll) +function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_periods, ∂ll; + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling, + initial_covariance::AbstractMatrix{<:AbstractFloat} = zeros(0, 0)) nr, nP, nE, na, nz = sys.nr, sys.nPast, sys.nExo, sys.na, sys.nz r1, r2, rq = sys.r1, sys.r2, sys.rq Lp = Matrix(sys.Lp); Eq = Matrix(sys.Eq); Ep = Matrix(sys.Ep); E1 = Vector(sys.E1) @@ -484,7 +511,9 @@ function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_period z0 = (Matrix{Float64}(ℒ.I(nz)) - sys.𝒜) \ sys.c G0 = reshape(g0 + Λ * (Pz * z0), nz, nE) Q0 = G0 * G0' + sys.QH; Q0 = (Q0 + Q0') / 2 - Σ0 = qkf_lyapunov(sys.𝒜, Q0) + Σ0 = qkf_lyapunov(sys.𝒜, Q0; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + initial_guess = initial_covariance) R = last(rrule(quadratic_kalman_recursion, sys.𝒜, sys.c, sys.QH, g0, Λ, Hm, Matrix(data_in_deviations), sys.𝒞, Pz, z0, Σ0, nz, nE, @@ -494,7 +523,8 @@ function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_period z̄0 = copy(R[11]); Σ̄0 = copy(R[12]) # Σ0 = 𝒜Σ0𝒜' + Q0 ⇒ X solves X = 𝒜'X𝒜 + Σ̄0 - X = qkf_lyapunov(Matrix(sys.𝒜'), (Σ̄0 + Σ̄0') / 2) + X = qkf_lyapunov(Matrix(sys.𝒜'), (Σ̄0 + Σ̄0') / 2; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) 𝒜̄ .+= 2 .* (X * sys.𝒜 * Σ0) Q̄0 = (X + X') / 2 Ḡ0 = 2 .* (Q̄0 * G0); Q̄H .+= Q̄0 @@ -578,7 +608,9 @@ function calculate_loglikelihood(::Val{:quadratic_kalman}, return run_quadratic_kalman(sys, data_in_deviations; measurement_error = measurement_error, presample_periods = presample_periods, - on_failure_loglikelihood = on_failure_loglikelihood) + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) end function rrule(::typeof(calculate_loglikelihood), @@ -603,10 +635,14 @@ function rrule(::typeof(calculate_loglikelihood), Hm = measurement_error === nothing ? zeros(n_obs, n_obs) : measurement_error isa AbstractMatrix ? Matrix{Float64}(measurement_error) : Matrix{Float64}(ℒ.Diagonal(collect(measurement_error))) + Σ₀ref = Ref{Matrix{Float64}}() llh = run_quadratic_kalman(sys, data_in_deviations; measurement_error = measurement_error, presample_periods = presample_periods, - on_failure_loglikelihood = on_failure_loglikelihood) + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + initial_covariance_out = Σ₀ref) nine(x...) = (NoTangent(), NoTangent(), NoTangent(), NoTangent(), x[1], x[2], NoTangent(), x[3], NoTangent()) @@ -618,7 +654,11 @@ function rrule(::typeof(calculate_loglikelihood), function quadratic_kalman_loglikelihood_pullback(∂llh_bar) ∂llh = unthunk(∂llh_bar) S̄1r, S̄2r, Ȳ = quadratic_kalman_pullback(sys, data_in_deviations, Hm, - presample_periods, ∂llh) + presample_periods, ∂llh; + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + initial_covariance = isassigned(Σ₀ref) ? + Σ₀ref[] : zeros(0, 0)) # scatter the retained rows back onto the full solution matrices ∂𝐒1 = zeros(size(𝐒[1])); ∂𝐒2 = zeros(size(𝐒[2])) ∂𝐒1[sys.oas, :] = S̄1r From e8ba8032a63622d4d8d0924f0d0c6ce9b421f18d Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 10:08:09 +0200 Subject: [PATCH 08/16] Lay the quadratic and linear Kalman filters side by side in the docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The quadratic Kalman section explained the filter on its own terms but never put it next to the linear one, which is where the intuition actually is: it is the same recursion, and setting S2 = 0 collapses one onto the other exactly. Adds a side-by-side of the two loops as they appear in the source, a table of the structural differences, and short notes on the three that carry consequences — the noise covariance becoming state-dependent (the conditional heteroskedasticity a second-order solution adds), the innovation ceasing to be Gaussian (why the linear filter is exact and this one is not), and the cost being cubic in a squared dimension (n_past^6), which governs when the filter is usable. Dimensions are the measured Smets-Wouters ones: 34 retained rows against an augmented 446 = 2*34 + 27*28/2. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 52 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/docs/src/filters.md b/docs/src/filters.md index 8e6da6b73..94c6f40c8 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -394,6 +394,58 @@ Without pruning there is no such representation: ``x_t`` is quadratic in ``x_{t- ``x_t\otimes x_t`` is quartic, needing ``x^{\otimes4}``, then ``x^{\otimes8}`` — the hierarchy never closes. Pruning truncates it at exactly one rung. +### Side by side with the linear Kalman filter + +It is the *same* recursion. Both filters run predict → innovate → update → accumulate, and +both score the innovation with the identical Gaussian formula. Setting +``\mathbf{S}_2 = 0`` collapses the quadratic filter onto the linear one exactly (this is a +test in the suite). The differences are entirely in what is being propagated. + +``` +linear Kalman (src/filter/kalman.jl) quadratic Kalman (src/filter/quadratic_kalman.jl) +───────────────────────────────────── ──────────────────────────────────────────────── + G = reshape(g₀ + Λ(P_z z)) ← state-dependent +P̂ = A P A' + 𝐁 P̂ = 𝒜 P 𝒜' + G G' + Q_H +û = A u ẑ = 𝒜 z + c ← non-zero drift +v = yₜ − C û v = yₜ − 𝒞 ẑ +F = C P̂ C' + H F = 𝒞 P̂ 𝒞' + H +ll += log|F| + v'F⁻¹v ll -= ½(v'F⁻¹v + log|F| + n log 2π) +K = P̂ C' F⁻¹ K = P̂ 𝒞' F⁻¹ +u = û + K v z = ẑ + K v +P = P̂ − K C P̂ P = P̂ − K 𝒞 P̂ +``` + +| | linear Kalman | quadratic Kalman | +|---|---|---| +| state carried | ``x_t`` | ``z_t = [x_1;\ x_2;\ \mathrm{vech}(x_{1,p}\otimes x_{1,p})]`` | +| dimension (SW07) | 34 | 446 | +| transition | ``x' = Ax + B\varepsilon`` | ``z' = \mathcal{A}z + c + w(z,\varepsilon)`` | +| drift ``c`` | zero — certainty equivalence empties ``\mathbf{S}_1``'s constant column | non-zero — carries the risk correction | +| noise covariance | ``\mathbf{B} = BB'``, **constant** | ``G(z)G(z)' + Q_H``, **depends on the state** | +| innovation | ``B\varepsilon`` — Gaussian | ``G\varepsilon + H(\varepsilon\otimes\varepsilon - \mathrm{vec}\,I)`` — **not** Gaussian | +| observation | ``y = Cx``, general ``C`` | ``y = (x_1+x_2)[\text{obs}]`` — a selection of two blocks | +| solve per period | LU of ``F`` (``n_{obs}^3``) | Cholesky of ``F`` (``n_{obs}^3``) | +| dominant cost | ``2n^3`` | ``2n_z^3`` — about ``2250\times`` more at SW07 sizes | +| exact? | yes, for a linear Gaussian model | no — a moment-matching approximation | + +Three of these carry real consequences. + +**The noise covariance moved inside the loop.** In the linear filter ``\mathbf{B} = BB'`` is +built once and added every period. In the quadratic filter the innovation loading ``G`` +is affine in ``z``, so ``Q`` must be rebuilt from the current state estimate at each ``t``. +That is precisely the conditional heteroskedasticity a second-order solution adds — the +model's shock impact depends on where the state is — and it is why the filter is not merely +a linear filter on a bigger vector. + +**The innovation is no longer Gaussian.** ``\varepsilon\otimes\varepsilon`` is a +``\chi^2``-type object; matching only its first two moments discards every higher cumulant. +The linear filter has nothing to discard, which is why it is exact and this one is not; the +next section works through what survives the approximation. + +**The cost is cubic in a squared dimension.** ``n_z`` grows like ``n_{past}^2/2``, so the +``O(n_z^3)`` covariance propagation grows like ``n_{past}^6``. This is the single fact that +governs when the filter is usable. + ### What is exact, and what is not The transition is exactly linear and the conditional first two moments are closed form. From 90b96241ff7fa13e977ad1289aa58e48b1ac1542 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 10:52:18 +0200 Subject: [PATCH 09/16] Add a cubic Kalman filter for the pruned third-order solution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The third-order analogue of the quadratic filter, built on the same fact: pruning truncates the Kronecker hierarchy at a fixed rung at every order, so the pruned third-order solution is exactly linear in the augmented state [x1; x2; x3; a*a; a*b; a*a*a]. The substance is the closure. Writing a_n = M a + v with v state-independent, each new Kronecker block resolves back onto blocks the state already carries — q11' onto q11, q12' onto q12 and q111, q111' onto q111 and q11 — and no fourth-order block appears because a_n has no q11 component. Computing the new blocks the obvious way, as kron(a_n, a_n), is quadratic in z and destroys the linearity the filter rests on without failing loudly, so the affineness of the step is asserted directly in the tests. Validated on an RBC model (2 shocks, 3 past states): the step is affine to 3e-15, reproduces the pruned third-order recursion to 2e-19 on every block including the Kronecker ones, its quadrature moments match Monte Carlo, and the log-likelihood matches a converged bootstrap particle filter at 181.78 against 181.43 over 60 periods — 0.006 per period, better than the quadratic filter's 0.025 on the comparable model. It only fits small models. The augmented dimension is 3nr + 2nPast^2 + nPast^3 and the covariance recursion is O(nz^3), so cost grows as nPast^9: fine to nPast=10, marginal at 12, hopeless at Smets-Wouters' 27 (nz = 21243, 3.6 GB per matrix). The build refuses above CUBIC_KALMAN_MAX_DIMENSION rather than appearing to hang. Not yet done, and noted in the docs: analytical assembly instead of quadrature, vech compression of the Kronecker blocks, and a hand-written rrule. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 58 ++++++ src/MacroModelling.jl | 12 +- src/default_options.jl | 7 +- src/filter/cubic_kalman.jl | 351 +++++++++++++++++++++++++++++++++++++ src/get_functions.jl | 12 ++ test/test_cubic_kalman.jl | 119 +++++++++++++ 6 files changed, 554 insertions(+), 5 deletions(-) create mode 100644 src/filter/cubic_kalman.jl create mode 100644 test/test_cubic_kalman.jl diff --git a/docs/src/filters.md b/docs/src/filters.md index 94c6f40c8..869f50299 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -583,6 +583,64 @@ quadratic Kalman filter for quadratic measurement equations. Andreasen (2013), * Applied Econometrics* 28, 929–955 — the central difference Kalman filter, the unpruned alternative. +## The cubic Kalman filter + +`filter = :cubic_kalman`, available only for `algorithm = :pruned_third_order`. + +The same construction one order up. Pruning truncates the Kronecker hierarchy at a fixed +rung at *every* order, so the pruned third-order solution is again exactly linear — in a +larger augmented state, + +```math +z_t = [\,x_1;\ x_2;\ x_3;\ a\otimes a;\ a\otimes b;\ a\otimes a\otimes a\,], +\qquad a = x_1[\text{past}],\ b = x_2[\text{past}]. +``` + +Writing ``a_n = Ma + v`` with ``v`` state-independent and ``u = Ma``, the new blocks close +back onto the existing ones: + +```math +\begin{aligned} +q_{11}' &= (M\otimes M)q_{11} + u\otimes v + v\otimes u + v\otimes v,\\ +q_{12}' &= (M\otimes M)q_{12} + (M\otimes W_q)q_{111} + (M\otimes W_l)q_{11} + u\otimes w_c + v\otimes b_n,\\ +q_{111}' &= (M\otimes M\otimes M)q_{111} + \text{3 perms of }((M\otimes M)q_{11})\otimes v + \text{3 perms of } u\otimes v\otimes v + v^{\otimes3}. +\end{aligned} +``` + +No fourth-order block appears, because ``a_n`` carries no ``q_{11}`` term — that is why the +system closes. The closure is the whole filter: recomputing the new blocks as +``\mathrm{kron}(a_n,a_n)`` would be quadratic in ``z`` and silently destroy the linearity +everything rests on. What is approximate is exactly what is approximate at second order — +the innovation is not Gaussian and only its first two moments are matched. + +Validated on an RBC model (2 shocks, 3 past states) against a converged bootstrap particle +filter: **181.78 against 181.43 over 60 periods, a gap of 0.006 per period** — smaller than +the quadratic filter's 0.025 on the comparable model. + +!!! warning "It only fits small models" + The augmented dimension is ``3n_r + 2n_{past}^2 + n_{past}^3``, and the covariance + recursion is ``O(n_z^3)`` — so cost grows as ``n_{past}^9``. + + | ``n_{past}`` | ``n_z`` | ms/period | verdict | + |---|---|---|---| + | 3 | 60 | <0.1 | fine | + | 8 | 676 | 6 | fine | + | 10 | 1245 | 39 | usable | + | 12 | 2070 | 177 | marginal | + | 15 | 3891 | 1178 | no | + | 27 (Smets-Wouters) | 21243 | 191725 | hopeless — 3.6 GB per matrix | + + `build_cubic_kalman_system_from_constants` refuses above + `CUBIC_KALMAN_MAX_DIMENSION` (2500) rather than appearing to hang. For anything + larger use the inversion filter or a particle filter. + +Compared with the quadratic filter this one is a straightforward implementation: the +transition is recovered by Gauss-Hermite quadrature (exact — the integrands are degree six) +rather than assembled analytically, there is no ``\mathrm{vech}`` compression of the +Kronecker blocks, and no hand-written `rrule`, so it is not differentiable in reverse mode. +Those are the three things to add if the filter is ever wanted at scale — though the +``n_{past}^9`` wall limits how much they can buy. + ## The filter-free likelihood There is a fourth option that is not a `filter` value at all, because it does not filter: instead of integrating the shocks out, it treats them as **parameters** and asks you to supply them. diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index ebcd0ba9f..a2f726551 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -193,6 +193,7 @@ include("./filter/inversion.jl") include("./filter/kalman.jl") include("./filter/particle.jl") include("./filter/quadratic_kalman.jl") +include("./filter/cubic_kalman.jl") export @model, @parameters, solve! @@ -419,9 +420,16 @@ function normalize_filtering_options(filter::Symbol, filter = :inversion end + # The cubic Kalman filter is the third-order analogue, and likewise defined + # only where the augmented state space is linear. + if filter == :cubic_kalman && algorithm != :pruned_third_order + @info "The cubic Kalman filter is only defined for `algorithm = :pruned_third_order`; got `:$(algorithm)`. Setting `filter = :inversion`." maxlog = maxlog + filter = :inversion + end + # Higher-order solutions are handled by the inversion filter by default, but # the particle filters are explicitly valid at every order too. - if algorithm != :first_order && filter != :inversion && filter != :quadratic_kalman && !is_particle + if algorithm != :first_order && filter != :inversion && filter != :quadratic_kalman && filter != :cubic_kalman && !is_particle @info "Higher order solution algorithms only support the inversion and particle filters. Setting `filter = :inversion`." maxlog = maxlog filter = :inversion is_particle = false @@ -447,7 +455,7 @@ function normalize_filtering_options(filter::Symbol, # origin — see `find_shocks`), which is a per-period choice a smoother could in # principle redistribute across time; doing so would be a different estimator, # not the inversion filter's smoother. - if filter == :quadratic_kalman && smooth + if filter in (:quadratic_kalman, :cubic_kalman) && smooth @info "The quadratic Kalman filter does not provide smoothed estimates. Setting `smooth = false`." maxlog = maxlog smooth = false end diff --git a/src/default_options.jl b/src/default_options.jl index 539aff1b0..73b478cff 100644 --- a/src/default_options.jl +++ b/src/default_options.jl @@ -13,9 +13,10 @@ const DEFAULT_PRESAMPLE_PERIODS = 0 # Each particle-filter variant is its own `filter` value, so the filter is fully # identified by a single symbol (no separate "which particle filter" argument). const PARTICLE_FILTERS = (:bootstrap_particle, :auxiliary_particle, :tempered_particle) -# The quadratic Kalman filter applies only to the pruned second-order solution, -# whose augmented state space is linear (see src/filter/quadratic_kalman.jl). -const SUPPORTED_FILTERS = (:kalman, :inversion, :quadratic_kalman, PARTICLE_FILTERS...) +# The quadratic and cubic Kalman filters apply only to the pruned second- and +# third-order solutions, whose augmented state spaces are linear (see +# src/filter/quadratic_kalman.jl and src/filter/cubic_kalman.jl). +const SUPPORTED_FILTERS = (:kalman, :inversion, :quadratic_kalman, :cubic_kalman, PARTICLE_FILTERS...) # `:particle` is accepted as a convenience alias for the bootstrap filter. const PARTICLE_FILTER_ALIASES = Dict(:particle => :bootstrap_particle) # Maps a filter symbol onto the internal variant tag used for dispatch. diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl new file mode 100644 index 000000000..c40e6757b --- /dev/null +++ b/src/filter/cubic_kalman.jl @@ -0,0 +1,351 @@ +@stable default_mode = "disable" begin + +# Cubic Kalman filter for the pruned third-order solution — the third-order +# analogue of `./quadratic_kalman.jl`, built on the same idea and validated the +# same way. +# +# The idea. Pruning truncates the Kronecker hierarchy at a fixed rung, so a +# pruned solution of order n is *exactly linear* in an augmented state. At second +# order that state is [x₁; x₂; a⊗a] with a = x₁[past]. At third order the +# recursion is +# +# aug₁ = [a; 1; ε], aug₁ʰ = [a; 0; ε], aug₂ = [b; 0; 0], aug₃ = [p; 0; 0] +# x₁ₜ = 𝐒₁ aug₁ +# x₂ₜ = 𝐒₁ aug₂ + ½ 𝐒₂ (aug₁ ⊗ aug₁) +# x₃ₜ = 𝐒₁ aug₃ + 𝐒₂ (aug₁ʰ ⊗ aug₂) + ⅙ 𝐒₃ (aug₁ ⊗ aug₁ ⊗ aug₁) +# +# (a, b, p are the past rows of x₁, x₂, x₃), so the state must additionally carry +# +# q₁₁ = a⊗a, q₁₂ = a⊗b, q₁₁₁ = a⊗a⊗a +# +# and the system closes: writing aₙ = M a + v with v = mc + Vε state-independent, +# and bₙ = M b + W_q q₁₁ + W_l a + w_c, +# +# q₁₁' = (M⊗M) q₁₁ + u⊗v + v⊗u + v⊗v +# q₁₂' = (M⊗M) q₁₂ + (M⊗W_q) q₁₁₁ + (M⊗W_l) q₁₁ + u⊗w_c + v⊗bₙ +# q₁₁₁' = (M⊗M⊗M) q₁₁₁ + [3 permutations of ((M⊗M)q₁₁)⊗v] +# + [3 permutations of u⊗v⊗v] + v⊗v⊗v +# +# with u = M a. No fourth-order block appears, because aₙ carries no q₁₁ term. +# The closure is what makes this work; recomputing the new blocks as kron(aₙ,aₙ) +# would be quadratic in z and silently break the linearity the filter rests on. +# +# What is exact and what is not. The transition is exactly linear, and the +# conditional mean and variance are computed by Gauss-Hermite quadrature that is +# exact for the polynomials involved (f is cubic in ε, so f f' is degree six). +# What is approximate is the same thing as at second order: the innovation is not +# Gaussian, and the filter matches only its first two moments. See +# `docs/src/filters.md`. +# +# Cost. The augmented dimension is 3n_r + 2n_past² + n_past³, so the O(n_z³) +# covariance recursion scales as n_past⁹. That confines the filter to small +# models — see `CUBIC_KALMAN_MAX_DIMENSION` below. + +# The covariance recursion is two n_z×n_z triple products per period. Beyond this +# dimension a single period costs seconds and a single matrix hundreds of MB, so +# refuse with a message that names the cause instead of appearing to hang. +const CUBIC_KALMAN_MAX_DIMENSION = 2500 + +# row-major flatten: rowvec(R)[(i-1)*size(R,2)+r] == R[i,r] +rowvec(R) = vec(permutedims(R)) + +# Probabilists' Gauss-Hermite nodes and weights via Golub-Welsch. +function gauss_hermite_nodes(n::Int) + J = ℒ.SymTridiagonal(zeros(n), sqrt.(1:n-1)) + E = ℒ.eigen(J) + return E.values, (E.vectors[1, :]) .^ 2 +end + +# Tensor product rule over `nExo` independent standard normals. `npt` points per +# dimension integrate polynomials of degree 2·npt−1 exactly; the integrands here +# reach degree six, so npt = 4 suffices and is the default. +function gauss_hermite_tensor(nExo::Int, npt::Int) + x, w = gauss_hermite_nodes(npt) + nodes = Vector{Vector{Float64}}() + wts = Float64[] + for I in Iterators.product(ntuple(_ -> 1:npt, nExo)...) + push!(nodes, [x[I[k]] for k in 1:nExo]) + push!(wts, prod(w[I[k]] for k in 1:nExo)) + end + return nodes, wts +end + +""" +Constant structure of the cubic augmented system: +`z = [x₁; x₂; x₃; q₁₁; q₁₂; q₁₁₁]` over the retained rows (past states plus +observables), together with the coefficient blocks that keep the step affine. +""" +function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒₃, observables_index::Vector{Int}) + T = cons.post_model_macro + nPast, nExo = T.nPast_not_future_and_mixed, T.nExo + past = T.past_not_future_and_mixed_idx + + # As in the quadratic filter, carry only the rows the recursion actually + # reads: past states for the transition, observables for the measurement. + oas = sort(union(past, observables_index)) + nr = length(oas) + pos = Dict(v => i for (i, v) in enumerate(oas)) + + na = nPast + 1 + nExo + nq11, nq12, nq111 = nPast^2, nPast^2, nPast^3 + nz = 3nr + nq11 + nq12 + nq111 + + if nz > CUBIC_KALMAN_MAX_DIMENSION + error("The cubic Kalman filter needs an augmented state of dimension " * + "$nz (= 3·$nr + 2·$(nPast)² + $(nPast)³) for this model, and its " * + "covariance recursion is O(n_z³) per period. The limit is " * + "$CUBIC_KALMAN_MAX_DIMENSION (`CUBIC_KALMAN_MAX_DIMENSION`). " * + "Use `filter = :inversion` or a particle filter instead.") + end + + S1 = Matrix(𝐒₁)[oas, :] + S2 = Matrix(𝐒₂)[oas, :] + S3 = Matrix(𝐒₃)[oas, :] + + Pm = zeros(nPast, nr) + for (i, j) in enumerate(past) + Pm[i, pos[j]] = 1.0 + end + + r1, r2, r3 = 1:nr, nr+1:2nr, 2nr+1:3nr + i11 = 3nr+1:3nr+nq11 + i12 = 3nr+nq11+1:3nr+nq11+nq12 + i111 = 3nr+nq11+nq12+1:nz + + # yₜ = (x₁ + x₂ + x₃)[observables] + C = zeros(length(observables_index), nz) + for (i, j) in enumerate(observables_index) + C[i, pos[j]] = 1.0 + C[i, nr+pos[j]] = 1.0 + C[i, 2nr+pos[j]] = 1.0 + end + + # aₙ = M a + mc + V ε + A1 = Pm * S1 + M = A1[:, 1:nPast] + mc = A1[:, nPast+1] + V = A1[:, nPast+2:na] + + # bₙ = M b + B2·K₂, with K₂ split into its (a,a), (a,tail)+(tail,a) and + # (tail,tail) parts so each can be routed onto the right state block. + B2 = Pm * S2 / 2 + ntail = 1 + nExo + Wq = zeros(nPast, nPast * nPast) + for i in 1:nPast, j in 1:nPast + Wq[:, (i-1)*nPast+j] = B2[:, (i-1)*na+j] + end + Wl_t = [zeros(nPast, nPast) for _ in 1:ntail] + for t in 1:ntail, k in 1:nPast + Wl_t[t][:, k] = B2[:, (k-1)*na+nPast+t] + B2[:, (nPast+t-1)*na+k] + end + Bc = B2[:, [(i-1)*na + j for i in nPast+1:na for j in nPast+1:na]] + MM = ℒ.kron(M, M) + + return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, + r1, r2, r3, i11, i12, i111, nq11, nq12, nq111, + M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail) +end + +""" +One step of the augmented map, `z ↦ f(z, ε)`. Affine in `z` by construction: +every product of two `z`-dependent quantities is read off an existing block +rather than recomputed. +""" +function cubic_kalman_step(sys, z::AbstractVector, ε::AbstractVector) + (; nPast, nExo, na, S1, S2, S3, Pm, r1, r2, r3, i11, i12, i111, + M, mc, V, Wq, Wl_t, Bc, MM, ntail) = sys + a = Pm * view(z, r1) + b = Pm * view(z, r2) + p = Pm * view(z, r3) + q11 = collect(view(z, i11)) + q12 = collect(view(z, i12)) + q111 = collect(view(z, i111)) + + tail = vcat(one(eltype(ε)), ε) + aug1 = vcat(a, 1.0, ε) + aug1h = vcat(a, 0.0, ε) + aug2 = vcat(b, 0.0, zeros(nExo)) + aug3 = vcat(p, 0.0, zeros(nExo)) + + # Kronecker inputs, with the all-past blocks read from the state. + K2 = Vector{Float64}(undef, na * na) + @inbounds for i in 1:na, j in 1:na + K2[(i-1)*na+j] = (i <= nPast && j <= nPast) ? q11[(i-1)*nPast+j] : aug1[i] * aug1[j] + end + K12 = zeros(na * na) + @inbounds for i in 1:na, j in 1:nPast + K12[(i-1)*na+j] = (i <= nPast) ? q12[(i-1)*nPast+j] : aug1h[i] * aug2[j] + end + K3 = Vector{Float64}(undef, na * na * na) + @inbounds for i in 1:na, j in 1:na, k in 1:na + r = ((i-1)*na + (j-1)) * na + k + ci = i <= nPast; cj = j <= nPast; ck = k <= nPast + n = ci + cj + ck + K3[r] = if n == 3 + q111[((i-1)*nPast + (j-1))*nPast + k] + elseif n == 2 + if ci && cj + q11[(i-1)*nPast+j] * aug1[k] + elseif ci && ck + q11[(i-1)*nPast+k] * aug1[j] + else + q11[(j-1)*nPast+k] * aug1[i] + end + else + aug1[i] * aug1[j] * aug1[k] + end + end + + x1n = S1 * aug1 + x2n = S1 * aug2 + S2 * K2 / 2 + x3n = S1 * aug3 + S2 * K12 + S3 * K3 / 6 + + # New Kronecker blocks, kept affine in z. + u = M * a # z-dependent, linear in a + v = mc + V * ε # z-independent + bn = Pm * x2n # affine in z + Q11 = Matrix(reshape(q11, nPast, nPast)') + Q12 = Matrix(reshape(q12, nPast, nPast)') + Q111 = Matrix(reshape(q111, nPast * nPast, nPast)') + + t2 = rowvec(M * Q11 * M') # = u⊗u = (M⊗M) q₁₁ + q11n = t2 + ℒ.kron(u, v) + ℒ.kron(v, u) + ℒ.kron(v, v) + + Wl = zeros(nPast, nPast) + for t in 1:ntail + Wl .+= tail[t] .* Wl_t[t] + end + wc = Bc * ℒ.kron(tail, tail) + q12n = rowvec(M * Q12 * M') + # u⊗(M b) = (M⊗M) q₁₂ + rowvec(M * Q111 * Wq') + # u⊗(Wq q₁₁) = (M⊗Wq) q₁₁₁ + rowvec(M * Q11 * Wl') + # u⊗(Wl a) = (M⊗Wl) q₁₁ + ℒ.kron(u, wc) + ℒ.kron(v, bn) + + vv = ℒ.kron(v, v) + q111n = rowvec(M * Q111 * MM') + # u⊗u⊗u + ℒ.kron(t2, v) + ℒ.kron(v, t2) + # u⊗u⊗v , v⊗u⊗u + ℒ.kron(u, vv) + ℒ.kron(vv, u) + # u⊗v⊗v , v⊗v⊗u + ℒ.kron(vv, v) # v⊗v⊗v + @inbounds for i in 1:nPast, j in 1:nPast, k in 1:nPast + r = ((i-1)*nPast + (j-1)) * nPast + k + q111n[r] += t2[(i-1)*nPast+k] * v[j] # u⊗v⊗u + q111n[r] += v[i] * u[j] * v[k] # v⊗u⊗v + end + + return vcat(x1n, x2n, x3n, q11n, q12n, q111n) +end + +# E[f(z,·)] and Var(f(z,·)) under ε ~ N(0,I), exactly. +function cubic_kalman_moments(sys, z, nodes, wts) + m = zeros(sys.nz) + S = zeros(sys.nz, sys.nz) + for (ε, w) in zip(nodes, wts) + fz = cubic_kalman_step(sys, z, ε) + m .+= w .* fz + ℒ.mul!(S, fz, fz', w, one(w)) + end + S .-= m * m' + return m, (S + S') / 2 +end + +# The step is affine, so the transition matrix and drift are recovered exactly +# from evaluations at the origin and at each basis vector. +function build_cubic_kalman_transition(sys, nodes, wts) + c, _ = cubic_kalman_moments(sys, zeros(sys.nz), nodes, wts) + 𝒜 = zeros(sys.nz, sys.nz) + e = zeros(sys.nz) + for j in 1:sys.nz + fill!(e, 0.0) + e[j] = 1.0 + mj, _ = cubic_kalman_moments(sys, e, nodes, wts) + 𝒜[:, j] = mj - c + end + return 𝒜, c +end + +""" +Kalman recursion on the cubic augmented state. Mirrors `run_quadratic_kalman`: +the noise covariance is rebuilt from the current state estimate every period, +because it depends on the state exactly as `G(z)G(z)'` does at second order. +""" +function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; + measurement_error::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}} = nothing, + presample_periods::Int = 0, + on_failure_loglikelihood::Real = -Inf, + quadrature_points::Int = 4, + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling) + nz, C = sys.nz, sys.C + n_obs, nT = size(data_in_deviations) + presample_periods = normalize_presample_periods(presample_periods, nT) + + Hm = if measurement_error === nothing + zeros(n_obs, n_obs) + elseif measurement_error isa AbstractMatrix + Matrix{Float64}(measurement_error) + else + Matrix{Float64}(ℒ.Diagonal(collect(measurement_error))) + end + + nodes, wts = gauss_hermite_tensor(sys.nExo, quadrature_points) + 𝒜, c = build_cubic_kalman_transition(sys, nodes, wts) + + z = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c + _, Q̄ = cubic_kalman_moments(sys, z, nodes, wts) + Σ = qkf_lyapunov(𝒜, Q̄; workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + + ll = 0.0 + for t in 1:nT + _, Q = cubic_kalman_moments(sys, z, nodes, wts) + zp = 𝒜 * z + c + Pp = 𝒜 * Σ * 𝒜' + Q + Pp = (Pp + Pp') / 2 + + v = data_in_deviations[:, t] - C * zp + F = C * Pp * C' + Hm + F = (F + F') / 2 + + Fc = ℒ.cholesky(F, check = false) + ℒ.issuccess(Fc) || return on_failure_loglikelihood + + if t > presample_periods + ll -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log(2π)) + isfinite(ll) || return on_failure_loglikelihood + end + + Kg = Pp * C' / Fc + z = zp + Kg * v + Σ = Pp - Kg * C * Pp + Σ = (Σ + Σ') / 2 + end + return ll +end + + +# ── standard filter interface ──────────────────────────────────────────────── +function calculate_loglikelihood(::Val{:cubic_kalman}, + ::Val{:pruned_third_order}, + observables_index::Vector{Int}, + 𝐒, + data_in_deviations::AbstractMatrix, + constants, + state, + workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) + sys = build_cubic_kalman_system_from_constants(constants, 𝐒[1], 𝐒[2], 𝐒[3], observables_index) + return run_cubic_kalman(sys, data_in_deviations; + measurement_error = measurement_error, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) +end + +end # @stable diff --git a/src/get_functions.jl b/src/get_functions.jl index d27626ea9..4235814b2 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -4877,6 +4877,18 @@ function get_loglikelihood(𝓂::ℳ, measurement_error = measurement_error_H, on_failure_loglikelihood = on_failure_loglikelihood, opts = opts) + elseif filter == :cubic_kalman + # Same idea one order up: the pruned third-order solution is linear in a + # larger augmented state. See src/filter/cubic_kalman.jl. + if has_missing + error("The cubic Kalman filter does not yet support missing observations.") + end + calculate_loglikelihood(Val(:cubic_kalman), Val(algorithm), obs_indices, + 𝐒, data_in_deviations, constants_obj, state, 𝓂.workspaces, + presample_periods = presample_periods, + measurement_error = measurement_error_H, + on_failure_loglikelihood = on_failure_loglikelihood, + opts = opts) elseif filter == :kalman if has_missing calculate_loglikelihood_with_missing(Val(:kalman), diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl new file mode 100644 index 000000000..c48397fb8 --- /dev/null +++ b/test/test_cubic_kalman.jl @@ -0,0 +1,119 @@ +using MacroModelling +using Test +using Random +import LinearAlgebra as ℒ +import AxisKeys: KeyedArray + +# The cubic Kalman filter rests on one property: the pruned third-order recursion +# is exactly affine in the augmented state z = [x₁; x₂; x₃; a⊗a; a⊗b; a⊗a⊗a]. +# These tests check that property directly, then the moments built on top of it, +# then the likelihood against a converged particle filter. + +@testset "Cubic Kalman filter" begin + opts = MacroModelling.merge_calculation_options() + + @model RBC_ckf begin + 1 / c[0] = (β / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ)) + c[0] + k[0] = (1 - δ) * k[-1] + q[0] + q[0] = exp(z[0]) * k[-1]^α * exp(g[0]) + z[0] = ρz * z[-1] + std_z * eps_z[x] + g[0] = ρg * g[-1] + std_g * eps_g[x] + end + @parameters RBC_ckf begin + std_z = 0.02 + std_g = 0.02 + ρz = 0.4 + ρg = 0.6 + δ = 0.02 + α = 0.5 + β = 0.95 + end + + MacroModelling.solve!(RBC_ckf, algorithm = :pruned_third_order, dynamics = true, opts = opts) + pars = RBC_ckf.parameter_values + _, _, 𝐒, _, _ = MacroModelling.get_relevant_steady_state_and_state_update(Val(:pruned_third_order), pars, RBC_ckf, opts = opts) + ssn = RBC_ckf.constants.post_complete_parameters.SS_and_pars_names + obs = [:c, :q] + obs_idx = convert(Vector{Int}, indexin(obs, ssn)) + + sys = MacroModelling.build_cubic_kalman_system_from_constants(RBC_ckf.constants, 𝐒[1], 𝐒[2], 𝐒[3], obs_idx) + @test sys.nz == 3sys.nr + 2sys.nPast^2 + sys.nPast^3 + + Random.seed!(3) + ε = randn(sys.nExo) + + # 1. the step is affine in z — the property the whole filter depends on + z1 = randn(sys.nz) + z2 = randn(sys.nz) + λ = 0.41 + lhs = MacroModelling.cubic_kalman_step(sys, λ .* z1 .+ (1 - λ) .* z2, ε) + rhs = λ .* MacroModelling.cubic_kalman_step(sys, z1, ε) .+ (1 - λ) .* MacroModelling.cubic_kalman_step(sys, z2, ε) + @test maximum(abs, lhs - rhs) < 1e-12 + + # 2. on a consistent state it reproduces the pruned third-order recursion, + # including the Kronecker blocks, with every product recomputed directly + consistent(x1, x2, x3) = (a = sys.Pm * x1; b = sys.Pm * x2; + vcat(x1, x2, x3, ℒ.kron(a, a), ℒ.kron(a, b), ℒ.kron(ℒ.kron(a, a), a))) + x1 = 0.01 .* randn(sys.nr) + x2 = 0.005 .* randn(sys.nr) + x3 = 0.002 .* randn(sys.nr) + a = sys.Pm * x1; b = sys.Pm * x2; p = sys.Pm * x3 + aug1 = vcat(a, 1.0, ε); aug1h = vcat(a, 0.0, ε) + aug2 = vcat(b, 0.0, zeros(sys.nExo)); aug3 = vcat(p, 0.0, zeros(sys.nExo)) + reference = consistent(sys.S1 * aug1, + sys.S1 * aug2 + sys.S2 * ℒ.kron(aug1, aug1) / 2, + sys.S1 * aug3 + sys.S2 * ℒ.kron(aug1h, aug2) + sys.S3 * ℒ.kron(ℒ.kron(aug1, aug1), aug1) / 6) + stepped = MacroModelling.cubic_kalman_step(sys, consistent(x1, x2, x3), ε) + @test maximum(abs, stepped - reference) < 1e-12 + @test maximum(abs, stepped[sys.i11] - reference[sys.i11]) < 1e-12 + @test maximum(abs, stepped[sys.i12] - reference[sys.i12]) < 1e-12 + @test maximum(abs, stepped[sys.i111] - reference[sys.i111]) < 1e-12 + + # 3. the recovered transition reproduces the conditional mean exactly, and + # both quadrature moments agree with Monte Carlo + nodes, wts = MacroModelling.gauss_hermite_tensor(sys.nExo, 4) + 𝒜, c = MacroModelling.build_cubic_kalman_transition(sys, nodes, wts) + zt = 0.01 .* randn(sys.nz) + mq, Sq = MacroModelling.cubic_kalman_moments(sys, zt, nodes, wts) + @test maximum(abs, mq - (𝒜 * zt + c)) < 1e-12 + + Random.seed!(5) + N = 200_000 + mm = zeros(sys.nz); SS = zeros(sys.nz, sys.nz) + for _ in 1:N + fz = MacroModelling.cubic_kalman_step(sys, zt, randn(sys.nExo)) + mm .+= fz + SS .+= fz * fz' + end + mm ./= N; SS ./= N; SS .-= mm * mm' + @test maximum(abs, mq - mm) / max(1e-12, maximum(abs, mq)) < 0.02 + @test maximum(abs, Sq - SS) / max(1e-12, maximum(abs, Sq)) < 0.05 + + # 4. the likelihood matches a converged bootstrap particle filter + Random.seed!(101) + T = 60 + sim = get_irf(RBC_ckf, algorithm = :pruned_third_order, periods = T, shocks = :simulate, levels = false) + Y = Matrix(sim(obs, :, :simulate)) + sd_obs = [sqrt(sum(abs2, Y[i, :] .- sum(Y[i, :]) / T) / (T - 1)) for i in eachindex(obs)] + mev = (0.2 .* sd_obs) .^ 2 + NSSS = get_steady_state(RBC_ckf, derivatives = false) + data = KeyedArray(Y .+ [NSSS(v) for v in obs]; Variable = obs, Time = 1:T) + + ll_ckf = get_loglikelihood(RBC_ckf, data, pars; algorithm = :pruned_third_order, + filter = :cubic_kalman, measurement_error = mev) + @test isfinite(ll_ckf) + + ll_pf = [get_loglikelihood(RBC_ckf, data, pars; algorithm = :pruned_third_order, + filter = :bootstrap_particle, measurement_error = mev, + n_particles = 80_000, particle_rng = Random.Xoshiro(s)) for s in 1:4] + m = sum(ll_pf) / length(ll_pf) + # The particle filter's log-likelihood is downward-biased by about Var/2. + @test abs(ll_ckf - (m + (sum(x -> (x - m)^2, ll_pf) / (length(ll_pf) - 1)) / 2)) < 0.05 * T + + # 5. gating: the filter is only defined on the pruned third-order solution. + # At any other order it falls back to the inversion filter, which admits + # no measurement error — so none is passed here. + ll_wrong = get_loglikelihood(RBC_ckf, data, pars; algorithm = :pruned_second_order, + filter = :cubic_kalman) + @test isfinite(ll_wrong) +end From 3e93075e2a81610e0db34357aa74f6e18072f7c7 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 14:10:56 +0200 Subject: [PATCH 10/16] Carry the quadratic filter's optimisations over to the cubic one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four of them, all measured on the RBC test model (2 shocks, 3 past states) with the log-likelihood unchanged to the last digit at every step (181.77600836054...): - vech compression of the symmetric Kronecker blocks. q11 and q111 are symmetric and fully symmetric respectively, so the state carries one entry per sorted multi-index. Applied by index maps rather than duplication/elimination matrices, which here would cost more than the compression saves. nz drops 60 -> 40 on the test model and 21243 -> 4863 at Smets-Wouters size; since the recursion is O(nz^3) this is two orders of magnitude on a mid-sized model, and it moves the practical wall from nPast ~= 12 to ~= 20. - An allocation-free step. The step does a few hundred flops but was allocating ~13 kB per call, and it runs (nz+1)*n_nodes times to build the transition plus n_nodes times per period — so it was entirely allocation-bound. With preallocated buffers it is 4.33 us/12960 B -> 2.71 us/0 B. The symmetric output blocks are written straight to their canonical slots, so the full nPast^3 vector is never formed. - An allocation-free recursion on preallocated buffers with in-place BLAS, and the observation applied by indexing its three selected rows instead of a gemm with a 0/1 matrix. - The quadrature contracts its node evaluations with one gemm rather than n_nodes rank-one updates, and the transition build skips the variance it never uses. End to end: 9.1 ms -> 5.9 ms on the test model, where the covariance recursion is only ~13% of the time and the quadrature dominates. The gains are far larger on models where nz is big enough for the O(nz^3) term to matter. Tests grow to 16: the compression maps are checked to round-trip a genuine symmetric Kronecker product, and the transition reference now forms every Kronecker product in full before compressing, so it exercises the compressed algebra rather than assuming it. Still not carried over: analytical assembly of the transition in place of quadrature, and a hand-written rrule. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 48 +++-- src/filter/cubic_kalman.jl | 385 ++++++++++++++++++++++++++++--------- test/test_cubic_kalman.jl | 19 +- 3 files changed, 336 insertions(+), 116 deletions(-) diff --git a/docs/src/filters.md b/docs/src/filters.md index 869f50299..abb0c9341 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -617,29 +617,41 @@ Validated on an RBC model (2 shocks, 3 past states) against a converged bootstra filter: **181.78 against 181.43 over 60 periods, a gap of 0.006 per period** — smaller than the quadratic filter's 0.025 on the comparable model. -!!! warning "It only fits small models" - The augmented dimension is ``3n_r + 2n_{past}^2 + n_{past}^3``, and the covariance - recursion is ``O(n_z^3)`` — so cost grows as ``n_{past}^9``. - - | ``n_{past}`` | ``n_z`` | ms/period | verdict | - |---|---|---|---| - | 3 | 60 | <0.1 | fine | - | 8 | 676 | 6 | fine | - | 10 | 1245 | 39 | usable | - | 12 | 2070 | 177 | marginal | - | 15 | 3891 | 1178 | no | - | 27 (Smets-Wouters) | 21243 | 191725 | hopeless — 3.6 GB per matrix | +``q_{11}`` and ``q_{111}`` are symmetric, so both are carried compressed — one entry per +sorted multi-index, the same ``\mathrm{vech}`` idea the quadratic filter uses, applied by +indexing rather than through duplication and elimination matrices. That takes the augmented +dimension from ``3n_r + 2n_{past}^2 + n_{past}^3`` down to + +```math +n_z = 3n_r + \tfrac{n_{past}(n_{past}+1)}{2} + n_{past}^2 + \tfrac{n_{past}(n_{past}+1)(n_{past}+2)}{6}, +``` + +which is roughly a sixth of the ``n_{past}^3`` block and, since the recursion is +``O(n_z^3)``, worth two orders of magnitude in flops on a mid-sized model. + +!!! warning "It still only fits small models" + Cost grows as ``n_{past}^9`` regardless of the constant factor. + + | ``n_{past}`` | ``n_z`` (compressed) | was | est. ms/period | verdict | + |---|---|---|---|---| + | 3 | 40 | 60 | <0.1 | fine | + | 8 | 256 | 676 | 0.3 | fine | + | 10 | 420 | 1245 | 1.5 | fine | + | 12 | 640 | 2070 | 5 | usable | + | 15 | 1091 | 3891 | 26 | usable | + | 20 | 2231 | 8881 | 222 | marginal | + | 27 (Smets-Wouters) | 4863 | 21243 | 2300 | no — 190 MB per matrix | `build_cubic_kalman_system_from_constants` refuses above `CUBIC_KALMAN_MAX_DIMENSION` (2500) rather than appearing to hang. For anything larger use the inversion filter or a particle filter. -Compared with the quadratic filter this one is a straightforward implementation: the -transition is recovered by Gauss-Hermite quadrature (exact — the integrands are degree six) -rather than assembled analytically, there is no ``\mathrm{vech}`` compression of the -Kronecker blocks, and no hand-written `rrule`, so it is not differentiable in reverse mode. -Those are the three things to add if the filter is ever wanted at scale — though the -``n_{past}^9`` wall limits how much they can buy. +The step function is allocation-free, the recursion runs on preallocated buffers with +in-place BLAS, the observation is applied by indexing its three selected rows rather than by +a gemm, and the quadrature contracts its nodes with a single gemm — all as in the quadratic +filter. Two things are not carried over: the transition is recovered by Gauss-Hermite +quadrature (exact, since the integrands are degree six) rather than assembled analytically, +and there is no hand-written `rrule`, so the filter is not differentiable in reverse mode. ## The filter-free likelihood diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index c40e6757b..224b70e10 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -37,9 +37,10 @@ # Gaussian, and the filter matches only its first two moments. See # `docs/src/filters.md`. # -# Cost. The augmented dimension is 3n_r + 2n_past² + n_past³, so the O(n_z³) -# covariance recursion scales as n_past⁹. That confines the filter to small -# models — see `CUBIC_KALMAN_MAX_DIMENSION` below. +# Cost. q₁₁ and q₁₁₁ are symmetric and carried compressed, giving an augmented +# dimension of 3n_r + n_past(n_past+1)/2 + n_past² + n_past(n_past+1)(n_past+2)/6. +# The O(n_z³) covariance recursion still scales as n_past⁹, so the filter stays +# confined to small models — see `CUBIC_KALMAN_MAX_DIMENSION` below. # The covariance recursion is two n_z×n_z triple products per period. Beyond this # dimension a single period costs seconds and a single matrix hundreds of MB, so @@ -49,6 +50,50 @@ const CUBIC_KALMAN_MAX_DIMENSION = 2500 # row-major flatten: rowvec(R)[(i-1)*size(R,2)+r] == R[i,r] rowvec(R) = vec(permutedims(R)) +# Index maps for the symmetric Kronecker blocks. `a⊗a` is symmetric and `a⊗a⊗a` +# fully symmetric, so the state carries one entry per sorted multi-index — the +# same vech idea the quadratic filter uses, but applied by indexing rather than +# by multiplying with duplication and elimination matrices, which would cost more +# than the compression saves. `expand` maps a full Kronecker position onto its +# compressed slot; `canonical` maps a slot back onto one representative position, +# which is exact precisely because the compressed blocks are symmetric. +function symmetric_pair_maps(n::Int) + slot = Dict{NTuple{2,Int},Int}() + m = 0 + for i in 1:n, j in i:n + m += 1 + slot[(i, j)] = m + end + expand = Vector{Int}(undef, n * n) + @inbounds for i in 1:n, j in 1:n + expand[(i-1)*n+j] = slot[minmax(i, j)] + end + canonical = Vector{Int}(undef, m) + @inbounds for i in 1:n, j in i:n + canonical[slot[(i, j)]] = (i-1)*n + j + end + return expand, canonical +end + +function symmetric_triple_maps(n::Int) + slot = Dict{NTuple{3,Int},Int}() + m = 0 + for i in 1:n, j in i:n, k in j:n + m += 1 + slot[(i, j, k)] = m + end + expand = Vector{Int}(undef, n^3) + @inbounds for i in 1:n, j in 1:n, k in 1:n + s = sort!([i, j, k]) + expand[((i-1)*n + (j-1))*n + k] = slot[(s[1], s[2], s[3])] + end + canonical = Vector{Int}(undef, m) + @inbounds for i in 1:n, j in i:n, k in j:n + canonical[slot[(i, j, k)]] = ((i-1)*n + (j-1))*n + k + end + return expand, canonical +end + # Probabilists' Gauss-Hermite nodes and weights via Golub-Welsch. function gauss_hermite_nodes(n::Int) J = ℒ.SymTridiagonal(zeros(n), sqrt.(1:n-1)) @@ -87,12 +132,15 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 pos = Dict(v => i for (i, v) in enumerate(oas)) na = nPast + 1 + nExo - nq11, nq12, nq111 = nPast^2, nPast^2, nPast^3 + # q₁₁ and q₁₁₁ are carried compressed; q₁₂ = a⊗b has no symmetry to exploit. + exp2, can2 = symmetric_pair_maps(nPast) + exp3, can3 = symmetric_triple_maps(nPast) + nq11, nq12, nq111 = length(can2), nPast^2, length(can3) nz = 3nr + nq11 + nq12 + nq111 if nz > CUBIC_KALMAN_MAX_DIMENSION error("The cubic Kalman filter needs an augmented state of dimension " * - "$nz (= 3·$nr + 2·$(nPast)² + $(nPast)³) for this model, and its " * + "$nz (= 3·$nr + $nq11 + $nq12 + $nq111) for this model, and its " * "covariance recursion is O(n_z³) per period. The limit is " * "$CUBIC_KALMAN_MAX_DIMENSION (`CUBIC_KALMAN_MAX_DIMENSION`). " * "Use `filter = :inversion` or a particle filter instead.") @@ -141,124 +189,227 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 Bc = B2[:, [(i-1)*na + j for i in nPast+1:na for j in nPast+1:na]] MM = ℒ.kron(M, M) - return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, + # Observation rows are a selection of the x₁, x₂ and x₃ blocks; carrying the + # three positions lets the recursion index instead of running a gemm with a + # 0/1 matrix, as the quadratic filter does with its two. + op1 = [pos[j] for j in observables_index] + op2 = op1 .+ nr + op3 = op1 .+ 2nr + + # Decoded multi-indices of the canonical entries, so the step can write the + # compressed blocks directly instead of materialising the full nPast³ vector. + can2_ij = [(fld(r - 1, nPast) + 1, mod(r - 1, nPast) + 1) for r in can2] + can3_ijk = [(fld(r - 1, nPast * nPast) + 1, + mod(fld(r - 1, nPast), nPast) + 1, + mod(r - 1, nPast) + 1) for r in can3] + + return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, op1, op2, op3, r1, r2, r3, i11, i12, i111, nq11, nq12, nq111, + exp2, can2, exp3, can3, can2_ij, can3_ijk, M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail) end +""" +Preallocated buffers for `cubic_kalman_step!`. The step is called +`(n_z + 1) · n_nodes` times to build the transition and `n_nodes` times per +period, and it is entirely allocation-bound — it does a few hundred flops but +allocated ~13 kB per call before these buffers existed. +""" +function cubic_kalman_workspace(sys) + (; nr, nPast, na, ntail) = sys + nP2 = nPast * nPast + return (; a = zeros(nPast), b = zeros(nPast), p = zeros(nPast), + q11 = zeros(nP2), q12 = zeros(nP2), q111 = zeros(nPast^3), + tail = zeros(ntail), tt = zeros(ntail * ntail), + aug1 = zeros(na), aug1h = zeros(na), aug2 = zeros(na), aug3 = zeros(na), + K2 = zeros(na * na), K12 = zeros(na * na), K3 = zeros(na * na * na), + x1n = zeros(nr), x2n = zeros(nr), x3n = zeros(nr), + u = zeros(nPast), v = zeros(nPast), bn = zeros(nPast), wc = zeros(nPast), + Q11 = zeros(nPast, nPast), Q12 = zeros(nPast, nPast), Q111 = zeros(nPast, nP2), + R2 = zeros(nPast, nPast), Tmp = zeros(nPast, nPast), + MQ111 = zeros(nPast, nP2), R3 = zeros(nPast, nP2), + Wl = zeros(nPast, nPast), t2 = zeros(nP2), vv = zeros(nP2)) +end + """ One step of the augmented map, `z ↦ f(z, ε)`. Affine in `z` by construction: every product of two `z`-dependent quantities is read off an existing block rather than recomputed. """ -function cubic_kalman_step(sys, z::AbstractVector, ε::AbstractVector) - (; nPast, nExo, na, S1, S2, S3, Pm, r1, r2, r3, i11, i12, i111, - M, mc, V, Wq, Wl_t, Bc, MM, ntail) = sys - a = Pm * view(z, r1) - b = Pm * view(z, r2) - p = Pm * view(z, r3) - q11 = collect(view(z, i11)) - q12 = collect(view(z, i12)) - q111 = collect(view(z, i111)) - - tail = vcat(one(eltype(ε)), ε) - aug1 = vcat(a, 1.0, ε) - aug1h = vcat(a, 0.0, ε) - aug2 = vcat(b, 0.0, zeros(nExo)) - aug3 = vcat(p, 0.0, zeros(nExo)) +function cubic_kalman_step!(out::AbstractVector, sys, z::AbstractVector, ε::AbstractVector, ws) + (; nr, nPast, nExo, na, S1, S2, S3, Pm, r1, r2, r3, i11, i12, i111, + exp2, can2_ij, can3_ijk, M, mc, V, Wq, Wl_t, Bc, MM, ntail) = sys + (; a, b, p, q11, q12, q111, tail, tt, aug1, aug1h, aug2, aug3, K2, K12, K3, + x1n, x2n, x3n, u, v, bn, wc, Q11, Q12, Q111, R2, Tmp, MQ111, R3, Wl, t2, vv) = ws + nP = nPast + nP2 = nP * nP + + ℒ.mul!(a, Pm, view(z, r1)) + ℒ.mul!(b, Pm, view(z, r2)) + ℒ.mul!(p, Pm, view(z, r3)) + + # The symmetric blocks arrive compressed; expand them so the algebra below is + # written on plain Kronecker products. + o11 = first(i11) - 1 + o111 = first(i111) - 1 + @inbounds for r in eachindex(q11) + q11[r] = z[o11+exp2[r]] + end + @inbounds for r in eachindex(q111) + q111[r] = z[o111+sys.exp3[r]] + end + @inbounds for (r, k) in enumerate(i12) + q12[r] = z[k] + end + + tail[1] = one(eltype(tail)) + @inbounds for i in 1:nExo + tail[1+i] = ε[i] + end + @inbounds for i in 1:nP + aug1[i] = a[i]; aug1h[i] = a[i]; aug2[i] = b[i]; aug3[i] = p[i] + end + aug1[nP+1] = 1.0; aug1h[nP+1] = 0.0; aug2[nP+1] = 0.0; aug3[nP+1] = 0.0 + @inbounds for i in 1:nExo + aug1[nP+1+i] = ε[i]; aug1h[nP+1+i] = ε[i]; aug2[nP+1+i] = 0.0; aug3[nP+1+i] = 0.0 + end # Kronecker inputs, with the all-past blocks read from the state. - K2 = Vector{Float64}(undef, na * na) @inbounds for i in 1:na, j in 1:na - K2[(i-1)*na+j] = (i <= nPast && j <= nPast) ? q11[(i-1)*nPast+j] : aug1[i] * aug1[j] + K2[(i-1)*na+j] = (i <= nP && j <= nP) ? q11[(i-1)*nP+j] : aug1[i] * aug1[j] end - K12 = zeros(na * na) - @inbounds for i in 1:na, j in 1:nPast - K12[(i-1)*na+j] = (i <= nPast) ? q12[(i-1)*nPast+j] : aug1h[i] * aug2[j] + fill!(K12, 0.0) + @inbounds for i in 1:na, j in 1:nP + K12[(i-1)*na+j] = (i <= nP) ? q12[(i-1)*nP+j] : aug1h[i] * aug2[j] end - K3 = Vector{Float64}(undef, na * na * na) @inbounds for i in 1:na, j in 1:na, k in 1:na r = ((i-1)*na + (j-1)) * na + k - ci = i <= nPast; cj = j <= nPast; ck = k <= nPast + ci = i <= nP; cj = j <= nP; ck = k <= nP n = ci + cj + ck K3[r] = if n == 3 - q111[((i-1)*nPast + (j-1))*nPast + k] + q111[((i-1)*nP + (j-1))*nP + k] elseif n == 2 if ci && cj - q11[(i-1)*nPast+j] * aug1[k] + q11[(i-1)*nP+j] * aug1[k] elseif ci && ck - q11[(i-1)*nPast+k] * aug1[j] + q11[(i-1)*nP+k] * aug1[j] else - q11[(j-1)*nPast+k] * aug1[i] + q11[(j-1)*nP+k] * aug1[i] end else aug1[i] * aug1[j] * aug1[k] end end - x1n = S1 * aug1 - x2n = S1 * aug2 + S2 * K2 / 2 - x3n = S1 * aug3 + S2 * K12 + S3 * K3 / 6 + ℒ.mul!(x1n, S1, aug1) + ℒ.mul!(x2n, S1, aug2); ℒ.mul!(x2n, S2, K2, 0.5, 1.0) + ℒ.mul!(x3n, S1, aug3); ℒ.mul!(x3n, S2, K12, 1.0, 1.0); ℒ.mul!(x3n, S3, K3, 1/6, 1.0) # New Kronecker blocks, kept affine in z. - u = M * a # z-dependent, linear in a - v = mc + V * ε # z-independent - bn = Pm * x2n # affine in z - Q11 = Matrix(reshape(q11, nPast, nPast)') - Q12 = Matrix(reshape(q12, nPast, nPast)') - Q111 = Matrix(reshape(q111, nPast * nPast, nPast)') - - t2 = rowvec(M * Q11 * M') # = u⊗u = (M⊗M) q₁₁ - q11n = t2 + ℒ.kron(u, v) + ℒ.kron(v, u) + ℒ.kron(v, v) - - Wl = zeros(nPast, nPast) - for t in 1:ntail - Wl .+= tail[t] .* Wl_t[t] - end - wc = Bc * ℒ.kron(tail, tail) - q12n = rowvec(M * Q12 * M') + # u⊗(M b) = (M⊗M) q₁₂ - rowvec(M * Q111 * Wq') + # u⊗(Wq q₁₁) = (M⊗Wq) q₁₁₁ - rowvec(M * Q11 * Wl') + # u⊗(Wl a) = (M⊗Wl) q₁₁ - ℒ.kron(u, wc) + ℒ.kron(v, bn) - - vv = ℒ.kron(v, v) - q111n = rowvec(M * Q111 * MM') + # u⊗u⊗u - ℒ.kron(t2, v) + ℒ.kron(v, t2) + # u⊗u⊗v , v⊗u⊗u - ℒ.kron(u, vv) + ℒ.kron(vv, u) + # u⊗v⊗v , v⊗v⊗u - ℒ.kron(vv, v) # v⊗v⊗v - @inbounds for i in 1:nPast, j in 1:nPast, k in 1:nPast - r = ((i-1)*nPast + (j-1)) * nPast + k - q111n[r] += t2[(i-1)*nPast+k] * v[j] # u⊗v⊗u - q111n[r] += v[i] * u[j] * v[k] # v⊗u⊗v - end - - return vcat(x1n, x2n, x3n, q11n, q12n, q111n) + ℒ.mul!(u, M, a) # z-dependent, linear in a + copyto!(v, mc); ℒ.mul!(v, V, ε, 1.0, 1.0) # z-independent + ℒ.mul!(bn, Pm, x2n) # affine in z + + @inbounds for j in 1:nP, s in 1:nP + Q11[j, s] = q11[(j-1)*nP+s] + Q12[j, s] = q12[(j-1)*nP+s] + end + @inbounds for j in 1:nP, s in 1:nP2 + Q111[j, s] = q111[(j-1)*nP2+s] + end + + ℒ.mul!(Tmp, M, Q11); ℒ.mul!(R2, Tmp, M') # M Q₁₁ M' — its rowvec is u⊗u + @inbounds for i in 1:nP, r in 1:nP + t2[(i-1)*nP+r] = R2[i, r] + end + @inbounds for i in 1:nP, j in 1:nP + vv[(i-1)*nP+j] = v[i] * v[j] + end + + fill!(Wl, 0.0) + @inbounds for t in 1:ntail + ℒ.axpy!(tail[t], Wl_t[t], Wl) + end + @inbounds for i in 1:ntail, j in 1:ntail + tt[(i-1)*ntail+j] = tail[i] * tail[j] + end + ℒ.mul!(wc, Bc, tt) + + # R2 ← M Q₁₂ M' + (M Q₁₁₁) Wq' + (M Q₁₁) Wl'; its rowvec is the z-dependent + # part of q₁₂', and Tmp still holds M Q₁₁ from above. + ℒ.mul!(MQ111, M, Q111) + ℒ.mul!(R2, Tmp, Wl') + ℒ.mul!(R2, MQ111, Wq', 1.0, 1.0) + ℒ.mul!(Tmp, M, Q12) + ℒ.mul!(R2, Tmp, M', 1.0, 1.0) + + ℒ.mul!(R3, MQ111, MM') # u⊗u⊗u = (M⊗M⊗M) q₁₁₁ + + @inbounds for i in 1:nr + out[i] = x1n[i]; out[nr+i] = x2n[i]; out[2nr+i] = x3n[i] + end + # q₁₁' and q₁₁₁' are symmetric, so only the canonical entries are formed. + @inbounds for (s, (i, j)) in enumerate(can2_ij) + out[o11+s] = t2[(i-1)*nP+j] + u[i]*v[j] + v[i]*u[j] + v[i]*v[j] + end + @inbounds for i in 1:nP, j in 1:nP + out[first(i12)-1 + (i-1)*nP+j] = R2[i, j] + u[i]*wc[j] + v[i]*bn[j] + end + @inbounds for (s, (i, j, k)) in enumerate(can3_ijk) + out[o111+s] = R3[i, (j-1)*nP+k] + # u⊗u⊗u + t2[(i-1)*nP+j] * v[k] + # u⊗u⊗v + v[i] * t2[(j-1)*nP+k] + # v⊗u⊗u + t2[(i-1)*nP+k] * v[j] + # u⊗v⊗u + u[i] * vv[(j-1)*nP+k] + # u⊗v⊗v + vv[(i-1)*nP+j] * u[k] + # v⊗v⊗u + v[i] * u[j] * v[k] + # v⊗u⊗v + vv[(i-1)*nP+j] * v[k] # v⊗v⊗v + end + return out +end + +# Allocating convenience wrapper, used by the tests. +function cubic_kalman_step(sys, z::AbstractVector, ε::AbstractVector, ws = cubic_kalman_workspace(sys)) + return cubic_kalman_step!(Vector{Float64}(undef, sys.nz), sys, z, ε, ws) end # E[f(z,·)] and Var(f(z,·)) under ε ~ N(0,I), exactly. -function cubic_kalman_moments(sys, z, nodes, wts) - m = zeros(sys.nz) - S = zeros(sys.nz, sys.nz) - for (ε, w) in zip(nodes, wts) - fz = cubic_kalman_step(sys, z, ε) - m .+= w .* fz - ℒ.mul!(S, fz, fz', w, one(w)) - end - S .-= m * m' +# +# The node evaluations are stacked into one matrix and contracted with a single +# gemm rather than accumulated as `nnodes` rank-one updates: same arithmetic, but +# it runs at BLAS-3 rather than BLAS-2 speed. `buf` may be supplied to reuse the +# stacking buffer across periods. +function cubic_kalman_moments(sys, z, nodes, wts; buf = nothing, ws = cubic_kalman_workspace(sys)) + nz = sys.nz + Fm = buf === nothing ? Matrix{Float64}(undef, nz, length(nodes)) : buf + @inbounds for (n, ε) in enumerate(nodes) + cubic_kalman_step!(view(Fm, :, n), sys, z, ε, ws) + end + m = Fm * wts + S = (Fm .* wts') * Fm' + ℒ.mul!(S, m, m', -one(eltype(S)), one(eltype(S))) return m, (S + S') / 2 end # The step is affine, so the transition matrix and drift are recovered exactly # from evaluations at the origin and at each basis vector. -function build_cubic_kalman_transition(sys, nodes, wts) - c, _ = cubic_kalman_moments(sys, zeros(sys.nz), nodes, wts) +function build_cubic_kalman_transition(sys, nodes, wts; ws = cubic_kalman_workspace(sys)) + # Only the mean is needed here, so skip the variance the moment routine would + # otherwise form — an nz×nz gemm per basis vector, nz+1 of them. + Fm = Matrix{Float64}(undef, sys.nz, length(nodes)) + mean_at = function (z) + @inbounds for (n, ε) in enumerate(nodes) + cubic_kalman_step!(view(Fm, :, n), sys, z, ε, ws) + end + return Fm * wts + end + c = mean_at(zeros(sys.nz)) 𝒜 = zeros(sys.nz, sys.nz) e = zeros(sys.nz) for j in 1:sys.nz fill!(e, 0.0) e[j] = 1.0 - mj, _ = cubic_kalman_moments(sys, e, nodes, wts) - 𝒜[:, j] = mj - c + 𝒜[:, j] = mean_at(e) - c end return 𝒜, c end @@ -275,7 +426,7 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; quadrature_points::Int = 4, workspaces = nothing, lyapunov_algorithm::Symbol = :doubling) - nz, C = sys.nz, sys.C + nz = sys.nz n_obs, nT = size(data_in_deviations) presample_periods = normalize_presample_periods(presample_periods, nT) @@ -288,35 +439,77 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; end nodes, wts = gauss_hermite_tensor(sys.nExo, quadrature_points) - 𝒜, c = build_cubic_kalman_transition(sys, nodes, wts) + ws = cubic_kalman_workspace(sys) + 𝒜, c = build_cubic_kalman_transition(sys, nodes, wts; ws = ws) z = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c - _, Q̄ = cubic_kalman_moments(sys, z, nodes, wts) + qbuf = Matrix{Float64}(undef, nz, length(nodes)) + _, Q̄ = cubic_kalman_moments(sys, z, nodes, wts; buf = qbuf, ws = ws) Σ = qkf_lyapunov(𝒜, Q̄; workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + # Preallocate the recursion's working matrices once, as the quadratic filter + # does: the covariance propagation is the whole cost, and allocating an nz×nz + # temporary per period competes with it directly. + op1, op2, op3 = sys.op1, sys.op2, sys.op3 + Pp = Matrix{Float64}(undef, nz, nz) + Tm = Matrix{Float64}(undef, nz, nz) + Pc = Matrix{Float64}(undef, nz, nz); copyto!(Pc, Σ) + zp = Vector{Float64}(undef, nz) + CP = Matrix{Float64}(undef, n_obs, nz) + F = Matrix{Float64}(undef, n_obs, n_obs) + Kg = Matrix{Float64}(undef, nz, n_obs) + v = Vector{Float64}(undef, n_obs) + Fv = Vector{Float64}(undef, n_obs) + ll = 0.0 - for t in 1:nT - _, Q = cubic_kalman_moments(sys, z, nodes, wts) - zp = 𝒜 * z + c - Pp = 𝒜 * Σ * 𝒜' + Q - Pp = (Pp + Pp') / 2 + log2pi = log(2π) + @inbounds for t in 1:nT + _, Q = cubic_kalman_moments(sys, z, nodes, wts; buf = qbuf, ws = ws) + + # Pp = 𝒜 Pc 𝒜' + Q + ℒ.mul!(Tm, 𝒜, Pc) + ℒ.mul!(Pp, Tm, 𝒜') + Pp .+= Q + for j in 1:nz, i in 1:j + m = (Pp[i, j] + Pp[j, i]) / 2 + Pp[i, j] = m; Pp[j, i] = m + end - v = data_in_deviations[:, t] - C * zp - F = C * Pp * C' + Hm - F = (F + F') / 2 + ℒ.mul!(zp, 𝒜, z); zp .+= c + + # yₜ = (x₁ + x₂ + x₃)[observables] — three selected rows, so index rather + # than multiply by C. + for i in 1:n_obs + v[i] = data_in_deviations[i, t] - (zp[op1[i]] + zp[op2[i]] + zp[op3[i]]) + for k in 1:nz + CP[i, k] = Pp[op1[i], k] + Pp[op2[i], k] + Pp[op3[i], k] + end + end + for i in 1:n_obs, j in 1:n_obs + F[i, j] = CP[i, op1[j]] + CP[i, op2[j]] + CP[i, op3[j]] + Hm[i, j] + end + for i in 1:n_obs, j in 1:i-1 + m = (F[i, j] + F[j, i]) / 2 + F[i, j] = m; F[j, i] = m + end Fc = ℒ.cholesky(F, check = false) ℒ.issuccess(Fc) || return on_failure_loglikelihood if t > presample_periods - ll -= 0.5 * (ℒ.dot(v, Fc \ v) + ℒ.logdet(Fc) + n_obs * log(2π)) + copyto!(Fv, v); ℒ.ldiv!(Fc, Fv) + ll -= 0.5 * (ℒ.dot(v, Fv) + ℒ.logdet(Fc) + n_obs * log2pi) isfinite(ll) || return on_failure_loglikelihood end - Kg = Pp * C' / Fc - z = zp + Kg * v - Σ = Pp - Kg * C * Pp - Σ = (Σ + Σ') / 2 + # K = CP' F⁻¹ ; z = zp + K v ; Pc = Pp − K CP + copyto!(Kg, CP'); ℒ.rdiv!(Kg, Fc) + copyto!(z, zp); ℒ.mul!(z, Kg, v, 1.0, 1.0) + copyto!(Pc, Pp); ℒ.mul!(Pc, Kg, CP, -1.0, 1.0) + for j in 1:nz, i in 1:j + m = (Pc[i, j] + Pc[j, i]) / 2 + Pc[i, j] = m; Pc[j, i] = m + end end return ll end diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl index c48397fb8..2a85be175 100644 --- a/test/test_cubic_kalman.jl +++ b/test/test_cubic_kalman.jl @@ -37,7 +37,19 @@ import AxisKeys: KeyedArray obs_idx = convert(Vector{Int}, indexin(obs, ssn)) sys = MacroModelling.build_cubic_kalman_system_from_constants(RBC_ckf.constants, 𝐒[1], 𝐒[2], 𝐒[3], obs_idx) - @test sys.nz == 3sys.nr + 2sys.nPast^2 + sys.nPast^3 + nP = sys.nPast + # q₁₁ and q₁₁₁ are carried compressed (symmetric); q₁₂ = a⊗b is not. + @test sys.nz == 3sys.nr + nP * (nP + 1) ÷ 2 + nP^2 + nP * (nP + 1) * (nP + 2) ÷ 6 + + # the compression maps must round-trip a genuine symmetric Kronecker product + let a = randn(nP) + q11 = ℒ.kron(a, a) + q111 = ℒ.kron(ℒ.kron(a, a), a) + @test q11[sys.can2][sys.exp2] ≈ q11 + @test q111[sys.can3][sys.exp3] ≈ q111 + @test length(sys.can2) == nP * (nP + 1) ÷ 2 + @test length(sys.can3) == nP * (nP + 1) * (nP + 2) ÷ 6 + end Random.seed!(3) ε = randn(sys.nExo) @@ -52,8 +64,11 @@ import AxisKeys: KeyedArray # 2. on a consistent state it reproduces the pruned third-order recursion, # including the Kronecker blocks, with every product recomputed directly + # The reference forms every Kronecker product directly and only then + # compresses, so it exercises the compressed algebra rather than assuming it. consistent(x1, x2, x3) = (a = sys.Pm * x1; b = sys.Pm * x2; - vcat(x1, x2, x3, ℒ.kron(a, a), ℒ.kron(a, b), ℒ.kron(ℒ.kron(a, a), a))) + vcat(x1, x2, x3, ℒ.kron(a, a)[sys.can2], ℒ.kron(a, b), + ℒ.kron(ℒ.kron(a, a), a)[sys.can3])) x1 = 0.01 .* randn(sys.nr) x2 = 0.005 .* randn(sys.nr) x3 = 0.002 .* randn(sys.nr) From 6882e73db10cef96a0a16bfdee1bb09ab926f53e Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 14:14:42 +0200 Subject: [PATCH 11/16] Drop the now-unused rowvec helper The in-place step writes the row-major flattens out by indexing, so nothing calls it any more; only the comments still use the term, which is now defined where they can see it. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/filter/cubic_kalman.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index 224b70e10..75d9160af 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -47,8 +47,10 @@ # refuse with a message that names the cause instead of appearing to hang. const CUBIC_KALMAN_MAX_DIMENSION = 2500 -# row-major flatten: rowvec(R)[(i-1)*size(R,2)+r] == R[i,r] -rowvec(R) = vec(permutedims(R)) +# Several intermediates below are computed as a matrix whose *row-major* flatten +# is the Kronecker vector wanted — "the rowvec of R" in the comments, meaning the +# vector v with v[(i-1)*size(R,2)+r] == R[i,r]. The step writes those entries out +# by indexing rather than materialising the flatten. # Index maps for the symmetric Kronecker blocks. `a⊗a` is symmetric and `a⊗a⊗a` # fully symmetric, so the state carries one entry per sorted multi-index — the From 55aa89257e6add6a5cb7260a811557462f8f4964 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 14:33:41 +0200 Subject: [PATCH 12/16] Assemble the cubic Kalman system analytically instead of by quadrature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit f(z, ·) is a polynomial of degree <= 3 in eps whose coefficients are affine in z. Recovering that coefficient matrix C(z) once gives both moments in closed form: E[f] = C(z) m, m_a = E[eps^a] Var(f) = C(z) Psi C(z)', Psi_ab = E[eps^(a+b)] - E[eps^a] E[eps^b] and since the shocks are independent standard normals, E[eps^a] factorises into double factorials, so Psi is a closed form rather than a set of Isserlis pairings. C(z) is recovered by interpolation on C(nExo+3, 3) points, which is also where the tensor Gauss-Hermite rule is left behind: its node count grew as npt^nExo (16384 for seven shocks at npt=4) against 120 for the coefficient basis. Consequences: no quadrature per period at all — Q(z) is one matvec and two gemms — and the build no longer blows up in the number of shocks. The quadrature path is kept and the analytic assembly is tested against it rather than assumed: A and c agree to 1e-9, Q(z) to 1e-8 relative at a non-trivial z, and the monomial moment vector matches tensor Gauss-Hermite to 1e-10. The log-likelihood is unchanged at 181.776008360546. 5.9 ms -> 2.3 ms on the test model (9.1 ms before this round of work started). Tests 16 -> 21. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- src/filter/cubic_kalman.jl | 142 ++++++++++++++++++++++++++++++++++--- test/test_cubic_kalman.jl | 15 ++++ 2 files changed, 148 insertions(+), 9 deletions(-) diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index 75d9160af..9251a0016 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -96,7 +96,76 @@ function symmetric_triple_maps(n::Int) return expand, canonical end -# Probabilists' Gauss-Hermite nodes and weights via Golub-Welsch. +# ── analytic assembly ──────────────────────────────────────────────────────── +# +# `f(z, ·)` is a polynomial of degree ≤ 3 in ε whose coefficients are affine in z: +# +# f(z, ε) = Σ_α c_α(z) ε^α , |α| ≤ 3. +# +# Recovering the coefficient matrix C(z) = [c_α(z)]_α once therefore gives both +# moments in closed form, with no quadrature anywhere: +# +# E[f] = C(z) m, m_α = E[ε^α] +# Var(f) = C(z) Ψ C(z)', Ψ_αβ = E[ε^{α+β}] − E[ε^α] E[ε^β] +# +# and because ε is a vector of *independent* standard normals, E[ε^α] factorises +# into double factorials — no Isserlis pairings needed. This replaces a tensor +# Gauss-Hermite rule whose node count grew as `npt^nExo`: the coefficient basis +# has only C(nExo+3, 3) elements (10 for two shocks, 120 for seven). + +# E[ε^α] = ∏ᵢ (αᵢ−1)!! when every αᵢ is even, and 0 otherwise. +double_factorial(n::Int) = n <= 0 ? 1.0 : Float64(prod(n:-2:1)) +gaussian_moment(α) = all(iseven, α) ? prod(double_factorial(a - 1) for a in α) : 0.0 + +# Exponent vectors α with |α| ≤ maxdeg, in n variables. +function monomial_exponents(n::Int, maxdeg::Int = 3) + out = Vector{Vector{Int}}() + cur = zeros(Int, n) + function rec(pos, rem) + if pos > n + push!(out, copy(cur)) + return + end + for d in 0:rem + cur[pos] = d + rec(pos + 1, rem - d) + end + cur[pos] = 0 + return + end + rec(1, maxdeg) + return out +end + +""" +Interpolation data for recovering a degree-≤3 polynomial in `nExo` variables from +its values: the exponent set, a unisolvent set of evaluation points, the inverse +Vandermonde, and the Gaussian moment vector and covariance of the monomials. +""" +function cubic_noise_basis(nExo::Int; seed::Int = 42) + exps = monomial_exponents(nExo, 3) + N = length(exps) + # Any N points with an invertible Vandermonde will do; Gaussian draws are + # unisolvent with probability one, and the conditioning is checked rather + # than assumed. + rng = Random.Xoshiro(seed) + pts = [randn(rng, nExo) for _ in 1:N] + V = [prod(pts[p][k]^exps[m][k] for k in 1:nExo) for p in 1:N, m in 1:N] + if !isfinite(ℒ.cond(V)) || ℒ.cond(V) > 1e10 + error("The cubic Kalman filter could not build a well-conditioned polynomial " * + "basis for $nExo shocks (condition number $(ℒ.cond(V))). This is a bug; " * + "please report it.") + end + # f(z, ε_p) = C(z) V[p,:]' ⇒ F = C V' ⇒ C = F (V')⁻¹ + W = Matrix(transpose(inv(V))) + m = [gaussian_moment(a) for a in exps] + Ψ = [gaussian_moment(exps[i] .+ exps[j]) - m[i] * m[j] for i in 1:N, j in 1:N] + Ψ = (Ψ + Ψ') / 2 + return (; exps, pts, W, m, Ψ, N) +end + +# Probabilists' Gauss-Hermite nodes and weights via Golub-Welsch. Retained so the +# tests can cross-check the analytic assembly against quadrature. function gauss_hermite_nodes(n::Int) J = ℒ.SymTridiagonal(zeros(n), sqrt.(1:n-1)) E = ℒ.eigen(J) @@ -416,16 +485,54 @@ function build_cubic_kalman_transition(sys, nodes, wts; ws = cubic_kalman_worksp return 𝒜, c end +""" +Assemble the whole system analytically: the transition `𝒜, c` and the affine +noise factor `vec(C(z)) = c₀ + Λz` from which `Q(z) = C(z) Ψ C(z)'`. + +Costs `(n_z + 1) · N` evaluations of the step, where `N = C(nExo+3, 3)` — the +same shape of work the quadrature build did, but with a node count that grows +polynomially in the number of shocks instead of exponentially. Afterwards no +quadrature is needed at all, per period or otherwise. +""" +function build_cubic_kalman_system(sys, basis; ws = cubic_kalman_workspace(sys)) + nz, N = sys.nz, basis.N + Fm = Matrix{Float64}(undef, nz, N) + coefficients_at = function (z) + @inbounds for (p, ε) in enumerate(basis.pts) + cubic_kalman_step!(view(Fm, :, p), sys, z, ε, ws) + end + return Fm * basis.W # nz × N + end + + C0 = coefficients_at(zeros(nz)) + c₀ = vec(C0) + Λ = Matrix{Float64}(undef, nz * N, nz) + 𝒜 = Matrix{Float64}(undef, nz, nz) + c = C0 * basis.m + e = zeros(nz) + for j in 1:nz + fill!(e, 0.0) + e[j] = 1.0 + ΔC = coefficients_at(e) + ΔC .-= C0 + Λ[:, j] = vec(ΔC) + # E[f] = C(z) m is affine in z, so column j of 𝒜 is ΔC·m. + ℒ.mul!(view(𝒜, :, j), ΔC, basis.m) + end + return 𝒜, c, c₀, Λ +end + """ Kalman recursion on the cubic augmented state. Mirrors `run_quadratic_kalman`: the noise covariance is rebuilt from the current state estimate every period, -because it depends on the state exactly as `G(z)G(z)'` does at second order. +because it depends on the state exactly as `G(z)G(z)'` does at second order — +here as `C(z) Ψ C(z)'`, with `C` affine in `z`, so a period costs one matvec and +two gemms rather than a quadrature sweep. """ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; measurement_error::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}} = nothing, presample_periods::Int = 0, on_failure_loglikelihood::Real = -Inf, - quadrature_points::Int = 4, workspaces = nothing, lyapunov_algorithm::Symbol = :doubling) nz = sys.nz @@ -440,14 +547,31 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; Matrix{Float64}(ℒ.Diagonal(collect(measurement_error))) end - nodes, wts = gauss_hermite_tensor(sys.nExo, quadrature_points) ws = cubic_kalman_workspace(sys) - 𝒜, c = build_cubic_kalman_transition(sys, nodes, wts; ws = ws) + basis = cubic_noise_basis(sys.nExo) + 𝒜, c, c₀, Λ = build_cubic_kalman_system(sys, basis; ws = ws) + N, Ψ = basis.N, basis.Ψ + + cvec = Vector{Float64}(undef, nz * N) + CΨ = Matrix{Float64}(undef, nz, N) + Q = Matrix{Float64}(undef, nz, nz) + # Q(z) = C(z) Ψ C(z)' with vec(C) = c₀ + Λz. + noise_covariance! = function (Q, z) + copyto!(cvec, c₀) + ℒ.mul!(cvec, Λ, z, 1.0, 1.0) + C = reshape(cvec, nz, N) + ℒ.mul!(CΨ, C, Ψ) + ℒ.mul!(Q, CΨ, C') + for j in 1:nz, i in 1:j + m = (Q[i, j] + Q[j, i]) / 2 + Q[i, j] = m; Q[j, i] = m + end + return Q + end z = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c - qbuf = Matrix{Float64}(undef, nz, length(nodes)) - _, Q̄ = cubic_kalman_moments(sys, z, nodes, wts; buf = qbuf, ws = ws) - Σ = qkf_lyapunov(𝒜, Q̄; workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + Σ = qkf_lyapunov(𝒜, noise_covariance!(Q, z); workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) # Preallocate the recursion's working matrices once, as the quadratic filter # does: the covariance propagation is the whole cost, and allocating an nz×nz @@ -466,7 +590,7 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; ll = 0.0 log2pi = log(2π) @inbounds for t in 1:nT - _, Q = cubic_kalman_moments(sys, z, nodes, wts; buf = qbuf, ws = ws) + noise_covariance!(Q, z) # Pp = 𝒜 Pc 𝒜' + Q ℒ.mul!(Tm, 𝒜, Pc) diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl index 2a85be175..d47a192f6 100644 --- a/test/test_cubic_kalman.jl +++ b/test/test_cubic_kalman.jl @@ -92,6 +92,21 @@ import AxisKeys: KeyedArray mq, Sq = MacroModelling.cubic_kalman_moments(sys, zt, nodes, wts) @test maximum(abs, mq - (𝒜 * zt + c)) < 1e-12 + # 3b. the analytic assembly must reproduce the quadrature exactly — it is a + # closed form for the same integrals, not an approximation of them. + basis = MacroModelling.cubic_noise_basis(sys.nExo) + @test basis.N == binomial(sys.nExo + 3, 3) + # the monomial moment vector and covariance against tensor Gauss-Hermite + @test all(abs(basis.m[i] - sum(w * prod(ε .^ basis.exps[i]) for (ε, w) in zip(nodes, wts))) < 1e-10 + for i in 1:basis.N) + 𝒜a, ca, c₀, Λ = MacroModelling.build_cubic_kalman_system(sys, basis) + @test maximum(abs, 𝒜a - 𝒜) < 1e-9 + @test maximum(abs, ca - c) < 1e-9 + # Q(z) = C(z) Ψ C(z)' against the quadrature variance, at a non-trivial z + Ca = reshape(c₀ + Λ * zt, sys.nz, basis.N) + Qa = Ca * basis.Ψ * Ca' + @test maximum(abs, Qa - Sq) / max(1e-12, maximum(abs, Sq)) < 1e-8 + Random.seed!(5) N = 200_000 mm = zeros(sys.nz); SS = zeros(sys.nz, sys.nz) From 65e5906d73596cb64f15e5c10662ab26a34a375f Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 14:51:43 +0200 Subject: [PATCH 13/16] Make the cubic Kalman filter differentiable, and stop reverse mode lying MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Forward mode now works end to end: the element type is threaded through the system build, the step workspace and the recursion, so ForwardDiff duals flow through. Gradients match central finite differences to 7e-11. The in-place BLAS path is kept for Float64 and the AD path falls back to allocating solves, since rdiv! has no Cholesky method for dual element types. That is the right mode for this filter: it is capped at nPast ~ 20 by its own cost, so a gradient is a handful of few-millisecond primal passes. More important is what reverse mode was doing. There is no hand-written rrule for this filter, and the top-level rrule's `llh_rrule === nothing` branch responds to that by returning on_failure_loglikelihood with an all-zero gradient — no error, no warning. Measured on the test model: Zygote returned exactly zeros where the true gradient has entries up to 3e4. A sampler would have run on that and produced garbage silently. Reverse mode now errors and names forward mode as the alternative. The existing measurement-error guard hid this in the common case, so the test pins the new guard specifically by asking for a gradient *without* measurement error, where that older guard cannot be what fires. Tests 21 -> 24. Quadratic Kalman tests still 31/31 after the rrules.jl change. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 40 ++++++++++++--- src/filter/cubic_kalman.jl | 100 +++++++++++++++++++++---------------- src/rrules.jl | 7 +++ test/test_cubic_kalman.jl | 18 +++++++ 4 files changed, 117 insertions(+), 48 deletions(-) diff --git a/docs/src/filters.md b/docs/src/filters.md index abb0c9341..f101f6a65 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -646,12 +646,40 @@ which is roughly a sixth of the ``n_{past}^3`` block and, since the recursion is `CUBIC_KALMAN_MAX_DIMENSION` (2500) rather than appearing to hang. For anything larger use the inversion filter or a particle filter. -The step function is allocation-free, the recursion runs on preallocated buffers with -in-place BLAS, the observation is applied by indexing its three selected rows rather than by -a gemm, and the quadrature contracts its nodes with a single gemm — all as in the quadratic -filter. Two things are not carried over: the transition is recovered by Gauss-Hermite -quadrature (exact, since the integrands are degree six) rather than assembled analytically, -and there is no hand-written `rrule`, so the filter is not differentiable in reverse mode. +### Assembly, and why there is no quadrature + +``f(z,\cdot)`` is a polynomial of degree ``\le 3`` in ``\varepsilon`` whose coefficients are +affine in ``z``. Recovering that coefficient matrix ``C(z)`` once gives both moments in +closed form: + +```math +\mathbb{E}[f] = C(z)\,m,\qquad +\mathrm{Var}(f) = C(z)\,\Psi\,C(z)', +``` + +with ``m_\alpha = \mathbb{E}[\varepsilon^\alpha]`` and +``\Psi_{\alpha\beta} = \mathbb{E}[\varepsilon^{\alpha+\beta}] - \mathbb{E}[\varepsilon^\alpha]\mathbb{E}[\varepsilon^\beta]``. +Because the shocks are *independent* standard normals, ``\mathbb{E}[\varepsilon^\alpha]`` +factorises into double factorials, so ``\Psi`` is a closed form rather than a sum over +Isserlis pairings. This is the exact analogue of the quadratic filter's affine ``G(z)`` with +``Q = GG'``: a period costs one matvec and two gemms, not a quadrature sweep. + +``C(z)`` is recovered by interpolation on ``\binom{n_\varepsilon+3}{3}`` points, which is +also where a tensor Gauss-Hermite rule is left behind — its node count grows as +``\mathrm{npt}^{n_\varepsilon}`` (16384 for seven shocks) against 120 for the coefficient +basis. The quadrature path is retained and the analytic assembly is tested against it rather +than assumed. + +### Derivatives + +Forward mode is exact: `ForwardDiff` matches central differences to ``\sim10^{-11}``. Since +the filter is confined to small models anyway, that is the appropriate mode — a gradient +costs a handful of primal passes, each a few milliseconds. + +There is no hand-written `rrule`, so **reverse mode raises an error**. This matters more than +it sounds: without that guard the generic fallback returns an all-zero gradient and no +warning at all, which a sampler will happily run on. Use `AutoForwardDiff`, or a +gradient-free sampler. ## The filter-free likelihood diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index 9251a0016..f003704fa 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -217,9 +217,13 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 "Use `filter = :inversion` or a particle filter instead.") end - S1 = Matrix(𝐒₁)[oas, :] - S2 = Matrix(𝐒₂)[oas, :] - S3 = Matrix(𝐒₃)[oas, :] + # Carry the solution matrices' element type so ForwardDiff duals flow through + # the whole assembly; the 0/1 selection matrices stay Float64 and promote on + # contact. + Tv = promote_type(eltype(𝐒₁), eltype(𝐒₂), eltype(𝐒₃)) + S1 = Matrix{Tv}(Matrix(𝐒₁)[oas, :]) + S2 = Matrix{Tv}(Matrix(𝐒₂)[oas, :]) + S3 = Matrix{Tv}(Matrix(𝐒₃)[oas, :]) Pm = zeros(nPast, nr) for (i, j) in enumerate(past) @@ -249,11 +253,11 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 # (tail,tail) parts so each can be routed onto the right state block. B2 = Pm * S2 / 2 ntail = 1 + nExo - Wq = zeros(nPast, nPast * nPast) + Wq = zeros(Tv, nPast, nPast * nPast) for i in 1:nPast, j in 1:nPast Wq[:, (i-1)*nPast+j] = B2[:, (i-1)*na+j] end - Wl_t = [zeros(nPast, nPast) for _ in 1:ntail] + Wl_t = [zeros(Tv, nPast, nPast) for _ in 1:ntail] for t in 1:ntail, k in 1:nPast Wl_t[t][:, k] = B2[:, (k-1)*na+nPast+t] + B2[:, (nPast+t-1)*na+k] end @@ -277,7 +281,7 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, op1, op2, op3, r1, r2, r3, i11, i12, i111, nq11, nq12, nq111, exp2, can2, exp3, can3, can2_ij, can3_ijk, - M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail) + M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail, Tv) end """ @@ -286,9 +290,10 @@ Preallocated buffers for `cubic_kalman_step!`. The step is called period, and it is entirely allocation-bound — it does a few hundred flops but allocated ~13 kB per call before these buffers existed. """ -function cubic_kalman_workspace(sys) +function cubic_kalman_workspace(sys, Tv = sys.Tv) (; nr, nPast, na, ntail) = sys nP2 = nPast * nPast + zeros(n...) = Base.zeros(Tv, n...) return (; a = zeros(nPast), b = zeros(nPast), p = zeros(nPast), q11 = zeros(nP2), q12 = zeros(nP2), q111 = zeros(nPast^3), tail = zeros(ntail), tt = zeros(ntail * ntail), @@ -441,7 +446,7 @@ end # Allocating convenience wrapper, used by the tests. function cubic_kalman_step(sys, z::AbstractVector, ε::AbstractVector, ws = cubic_kalman_workspace(sys)) - return cubic_kalman_step!(Vector{Float64}(undef, sys.nz), sys, z, ε, ws) + return cubic_kalman_step!(Vector{sys.Tv}(undef, sys.nz), sys, z, ε, ws) end # E[f(z,·)] and Var(f(z,·)) under ε ~ N(0,I), exactly. @@ -452,7 +457,7 @@ end # stacking buffer across periods. function cubic_kalman_moments(sys, z, nodes, wts; buf = nothing, ws = cubic_kalman_workspace(sys)) nz = sys.nz - Fm = buf === nothing ? Matrix{Float64}(undef, nz, length(nodes)) : buf + Fm = buf === nothing ? Matrix{sys.Tv}(undef, nz, length(nodes)) : buf @inbounds for (n, ε) in enumerate(nodes) cubic_kalman_step!(view(Fm, :, n), sys, z, ε, ws) end @@ -467,7 +472,7 @@ end function build_cubic_kalman_transition(sys, nodes, wts; ws = cubic_kalman_workspace(sys)) # Only the mean is needed here, so skip the variance the moment routine would # otherwise form — an nz×nz gemm per basis vector, nz+1 of them. - Fm = Matrix{Float64}(undef, sys.nz, length(nodes)) + Fm = Matrix{sys.Tv}(undef, sys.nz, length(nodes)) mean_at = function (z) @inbounds for (n, ε) in enumerate(nodes) cubic_kalman_step!(view(Fm, :, n), sys, z, ε, ws) @@ -495,8 +500,8 @@ polynomially in the number of shocks instead of exponentially. Afterwards no quadrature is needed at all, per period or otherwise. """ function build_cubic_kalman_system(sys, basis; ws = cubic_kalman_workspace(sys)) - nz, N = sys.nz, basis.N - Fm = Matrix{Float64}(undef, nz, N) + nz, N, Tv = sys.nz, basis.N, sys.Tv + Fm = Matrix{Tv}(undef, nz, N) coefficients_at = function (z) @inbounds for (p, ε) in enumerate(basis.pts) cubic_kalman_step!(view(Fm, :, p), sys, z, ε, ws) @@ -504,15 +509,15 @@ function build_cubic_kalman_system(sys, basis; ws = cubic_kalman_workspace(sys)) return Fm * basis.W # nz × N end - C0 = coefficients_at(zeros(nz)) + C0 = coefficients_at(zeros(Tv, nz)) c₀ = vec(C0) - Λ = Matrix{Float64}(undef, nz * N, nz) - 𝒜 = Matrix{Float64}(undef, nz, nz) + Λ = Matrix{Tv}(undef, nz * N, nz) + 𝒜 = Matrix{Tv}(undef, nz, nz) c = C0 * basis.m - e = zeros(nz) + e = zeros(Tv, nz) for j in 1:nz - fill!(e, 0.0) - e[j] = 1.0 + fill!(e, zero(Tv)) + e[j] = one(Tv) ΔC = coefficients_at(e) ΔC .-= C0 Λ[:, j] = vec(ΔC) @@ -539,26 +544,32 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; n_obs, nT = size(data_in_deviations) presample_periods = normalize_presample_periods(presample_periods, nT) + # Promote over every differentiable input. The preallocated buffers below fix + # the element type, so missing one makes forward-mode AD fail with respect to + # exactly that argument. + Tv = promote_type(sys.Tv, eltype(data_in_deviations), + measurement_error === nothing ? Float64 : eltype(measurement_error)) + Hm = if measurement_error === nothing - zeros(n_obs, n_obs) + zeros(Tv, n_obs, n_obs) elseif measurement_error isa AbstractMatrix - Matrix{Float64}(measurement_error) + Matrix{Tv}(measurement_error) else - Matrix{Float64}(ℒ.Diagonal(collect(measurement_error))) + Matrix{Tv}(ℒ.Diagonal(collect(measurement_error))) end - ws = cubic_kalman_workspace(sys) + ws = cubic_kalman_workspace(sys, Tv) basis = cubic_noise_basis(sys.nExo) 𝒜, c, c₀, Λ = build_cubic_kalman_system(sys, basis; ws = ws) N, Ψ = basis.N, basis.Ψ - cvec = Vector{Float64}(undef, nz * N) - CΨ = Matrix{Float64}(undef, nz, N) - Q = Matrix{Float64}(undef, nz, nz) + cvec = Vector{Tv}(undef, nz * N) + CΨ = Matrix{Tv}(undef, nz, N) + Q = Matrix{Tv}(undef, nz, nz) # Q(z) = C(z) Ψ C(z)' with vec(C) = c₀ + Λz. noise_covariance! = function (Q, z) copyto!(cvec, c₀) - ℒ.mul!(cvec, Λ, z, 1.0, 1.0) + ℒ.mul!(cvec, Λ, z, one(Tv), one(Tv)) C = reshape(cvec, nz, N) ℒ.mul!(CΨ, C, Ψ) ℒ.mul!(Q, CΨ, C') @@ -569,7 +580,7 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; return Q end - z = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c + z = (Matrix{Tv}(ℒ.I(nz)) - 𝒜) \ c Σ = qkf_lyapunov(𝒜, noise_covariance!(Q, z); workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) @@ -577,17 +588,17 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; # does: the covariance propagation is the whole cost, and allocating an nz×nz # temporary per period competes with it directly. op1, op2, op3 = sys.op1, sys.op2, sys.op3 - Pp = Matrix{Float64}(undef, nz, nz) - Tm = Matrix{Float64}(undef, nz, nz) - Pc = Matrix{Float64}(undef, nz, nz); copyto!(Pc, Σ) - zp = Vector{Float64}(undef, nz) - CP = Matrix{Float64}(undef, n_obs, nz) - F = Matrix{Float64}(undef, n_obs, n_obs) - Kg = Matrix{Float64}(undef, nz, n_obs) - v = Vector{Float64}(undef, n_obs) - Fv = Vector{Float64}(undef, n_obs) - - ll = 0.0 + Pp = Matrix{Tv}(undef, nz, nz) + Tm = Matrix{Tv}(undef, nz, nz) + Pc = Matrix{Tv}(undef, nz, nz); copyto!(Pc, Σ) + zp = Vector{Tv}(undef, nz) + CP = Matrix{Tv}(undef, n_obs, nz) + F = Matrix{Tv}(undef, n_obs, n_obs) + Kg = Matrix{Tv}(undef, nz, n_obs) + v = Vector{Tv}(undef, n_obs) + Fv = Vector{Tv}(undef, n_obs) + + ll = zero(Tv) log2pi = log(2π) @inbounds for t in 1:nT noise_covariance!(Q, z) @@ -628,10 +639,15 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; isfinite(ll) || return on_failure_loglikelihood end - # K = CP' F⁻¹ ; z = zp + K v ; Pc = Pp − K CP - copyto!(Kg, CP'); ℒ.rdiv!(Kg, Fc) - copyto!(z, zp); ℒ.mul!(z, Kg, v, 1.0, 1.0) - copyto!(Pc, Pp); ℒ.mul!(Pc, Kg, CP, -1.0, 1.0) + # K = CP' F⁻¹ ; z = zp + K v ; Pc = Pp − K CP. `rdiv!` has no Cholesky + # method for dual element types, so the AD path solves and transposes. + if Tv === Float64 + copyto!(Kg, CP'); ℒ.rdiv!(Kg, Fc) + else + Kg = Matrix((Fc \ CP)') + end + copyto!(z, zp); ℒ.mul!(z, Kg, v, one(Tv), one(Tv)) + copyto!(Pc, Pp); ℒ.mul!(Pc, Kg, CP, -one(Tv), one(Tv)) for j in 1:nz, i in 1:j m = (Pc[i, j] + Pc[j, i]) / 2 Pc[i, j] = m; Pc[j, i] = m diff --git a/src/rrules.jl b/src/rrules.jl index dbb2a39c5..103ab7e5a 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -1883,6 +1883,13 @@ function rrule(::typeof(get_loglikelihood), error("Reverse-mode automatic differentiation of the Kalman likelihood with measurement error (`measurement_error`) is not yet supported. Use forward-mode AD (e.g. `AutoForwardDiff`) or a gradient-free sampler.") end + # The cubic Kalman filter has no hand-written adjoint. Without this guard the + # `llh_rrule === nothing` branch below would hand back an all-zero gradient + # and no error at all, which a sampler would happily run on. + if filter == :cubic_kalman + error("Reverse-mode automatic differentiation of the cubic Kalman filter (`filter = :cubic_kalman`) is not supported. Use forward-mode AD (e.g. `AutoForwardDiff`), which is exact for this filter and matches finite differences to ~1e-11, or a gradient-free sampler.") + end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl index d47a192f6..98fe316a5 100644 --- a/test/test_cubic_kalman.jl +++ b/test/test_cubic_kalman.jl @@ -1,6 +1,9 @@ using MacroModelling using Test using Random +using ForwardDiff +using FiniteDifferences +using Zygote import LinearAlgebra as ℒ import AxisKeys: KeyedArray @@ -140,6 +143,21 @@ import AxisKeys: KeyedArray # The particle filter's log-likelihood is downward-biased by about Var/2. @test abs(ll_ckf - (m + (sum(x -> (x - m)^2, ll_pf) / (length(ll_pf) - 1)) / 2)) < 0.05 * T + # 4b. forward-mode AD is exact; reverse mode must *fail loudly* rather than + # fall through to the generic zero-gradient path, which a sampler would + # run on without noticing. + f = p -> get_loglikelihood(RBC_ckf, data, p; algorithm = :pruned_third_order, + filter = :cubic_kalman, measurement_error = mev) + g_ad = ForwardDiff.gradient(f, pars) + g_fd = FiniteDifferences.grad(central_fdm(5, 1), f, pars)[1] + @test maximum(abs.(g_ad .- g_fd) ./ max.(1.0, abs.(g_fd))) < 1e-7 + @test !all(iszero, g_ad) + # without measurement error the unrelated measurement-error guard cannot be + # what fires, so this pins the cubic-filter guard specifically + f_nome = p -> get_loglikelihood(RBC_ckf, data, p; algorithm = :pruned_third_order, + filter = :cubic_kalman) + @test_throws ErrorException Zygote.gradient(f_nome, pars) + # 5. gating: the filter is only defined on the pruned third-order solution. # At any other order it falls back to the inversion filter, which admits # no measurement error — so none is passed here. From 4891b919c8671de0fda8d2b9d5ff1407020a2afb Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 18:57:12 +0200 Subject: [PATCH 14/16] Hand-written reverse-mode adjoint for the cubic Kalman filter Reverse mode now works and matches central differences to ~1e-10, on every measurement-error shape (none, diagonal, full covariance). The chain S -> sys -> {f(z_k, eps_p)} -> (A, c, c0, Lambda) -> llh is differentiated in three pieces, each verified against ForwardDiff *in isolation* so a regression localises instead of just moving the end-to-end number: - the step adjoint (5e-16). Everything the step builds from z and eps alone -- aug, K2, K12, K3, the Q blocks -- is constant here, so only the paths through the solution matrices carry cotangents. - the build adjoint (2e-15). The build is linear in the collected step evaluations, so those maps transpose directly and the work is replaying the step adjoint over the same (nz+1)*N points the forward pass visited. - the recursion adjoint, mirroring the quadratic filter's verified one; the structural difference is Q = C Psi C' in place of GG' + Q_H, whose cotangent is 2 Q_bar C Psi because P_bar_p is symmetrised before use. Reverse mode is now the better choice here as well as the correct one: 12.0 ms against a 2.3 ms primal and independent of parameter count, where forward mode is 44.3 ms for seven parameters and grows linearly. Also fixes a 6.5x primal regression introduced with forward-mode support. The promoted element type was stored as a field of the system, so `Matrix{sys.Tv}` inferred as DataType rather than Type{Float64} and lost specialisation throughout; the hot paths now read `eltype(sys.S1)`, which is inferrable, and the primal is back to 2.27 ms from 14.94 ms. The zero-gradient guard added in the previous commit is removed, since the fallback it protected against is no longer reachable. Tests 24 -> 28. Quadratic Kalman still 31/31 after the rrules.jl changes. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 34 ++- src/filter/cubic_kalman.jl | 442 ++++++++++++++++++++++++++++++++++--- src/rrules.jl | 21 +- test/test_cubic_kalman.jl | 76 ++++++- 4 files changed, 509 insertions(+), 64 deletions(-) diff --git a/docs/src/filters.md b/docs/src/filters.md index f101f6a65..84a8e0fe6 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -672,14 +672,32 @@ than assumed. ### Derivatives -Forward mode is exact: `ForwardDiff` matches central differences to ``\sim10^{-11}``. Since -the filter is confined to small models anyway, that is the appropriate mode — a gradient -costs a handful of primal passes, each a few milliseconds. - -There is no hand-written `rrule`, so **reverse mode raises an error**. This matters more than -it sounds: without that guard the generic fallback returns an all-zero gradient and no -warning at all, which a sampler will happily run on. Use `AutoForwardDiff`, or a -gradient-free sampler. +Both modes work and both match central differences to ``\sim10^{-10}``. + +Reverse mode has a hand-written adjoint, as the quadratic filter does. It composes three +pieces, each verified against `ForwardDiff` in isolation so a regression localises rather +than merely moving the end-to-end number: + +| piece | what it does | +|---|---| +| step adjoint | ``\partial f(z,\varepsilon)`` onto ``\mathbf{S}_1,\mathbf{S}_2,\mathbf{S}_3`` and the derived blocks | +| build adjoint | ``\partial(\mathcal{A}, c, c_0, \Lambda)`` replayed over the same ``(n_z+1)N`` points the forward pass visited | +| recursion adjoint | the Kalman loop, with ``Q = C\Psi C'`` in place of the quadratic filter's ``GG' + Q_H`` | + +Everything the step builds from ``z`` and ``\varepsilon`` alone — ``\mathrm{aug}``, ``K_2``, +``K_{12}``, ``K_3``, the ``Q`` blocks — is constant for the adjoint, so only the paths +through the solution matrices carry cotangents. + +Cost, on the RBC test model with seven parameters: + +| | time | relative | +|---|---|---| +| primal | 2.3 ms | — | +| reverse (`Zygote`) | 12.0 ms | 5.3× primal, **independent of parameter count** | +| forward (`ForwardDiff`) | 44.3 ms | 19.5× primal, growing linearly in parameters | + +Reverse mode is therefore the default choice, and the gap widens with every parameter added. +Forward mode remains available and is a useful independent check. ## The filter-free likelihood diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index f003704fa..274592471 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -30,12 +30,18 @@ # The closure is what makes this work; recomputing the new blocks as kron(aₙ,aₙ) # would be quadratic in z and silently break the linearity the filter rests on. # -# What is exact and what is not. The transition is exactly linear, and the -# conditional mean and variance are computed by Gauss-Hermite quadrature that is -# exact for the polynomials involved (f is cubic in ε, so f f' is degree six). -# What is approximate is the same thing as at second order: the innovation is not -# Gaussian, and the filter matches only its first two moments. See -# `docs/src/filters.md`. +# What is exact and what is not. The transition is exactly linear, and both +# conditional moments are closed forms — f is a degree-3 polynomial in ε with +# z-affine coefficients, so recovering those coefficients once gives E[f] = C(z)m +# and Var(f) = C(z)ΨC(z)' exactly, with no quadrature. What is approximate is the +# same thing as at second order: the innovation is not Gaussian, and the filter +# matches only its first two moments. See `docs/src/filters.md`. +# +# Derivatives. Reverse mode has a hand-written adjoint (step, build and recursion, +# each verified against ForwardDiff); forward mode works too via the promoted +# element type. Note that `eltype(sys.S1)` rather than a stored type field is what +# the hot paths branch on — a `DataType`-typed field infers as `DataType`, not +# `Type{Float64}`, which costs specialisation and measured 6.5x on the primal. # # Cost. q₁₁ and q₁₁₁ are symmetric and carried compressed, giving an augmented # dimension of 3n_r + n_past(n_past+1)/2 + n_past² + n_past(n_past+1)(n_past+2)/6. @@ -186,6 +192,69 @@ function gauss_hermite_tensor(nExo::Int, npt::Int) return nodes, wts end +""" +Everything the step needs that is derived from the solution matrices, in one +place so the forward pass and its adjoint cannot drift apart: + + aₙ = M a + mc + V ε, bₙ = M b + B2·K₂ + +with K₂ split into its (a,a), (a,tail)+(tail,a) and (tail,tail) parts — the +coefficients `Wq`, `Wl_t` and `Bc` — so each can be routed onto the right state +block. +""" +function cubic_derived_matrices(S1, S2, Pm, nPast::Int, nExo::Int, na::Int) + Tv = promote_type(eltype(S1), eltype(S2)) + A1 = Pm * S1 + M = A1[:, 1:nPast] + mc = A1[:, nPast+1] + V = A1[:, nPast+2:na] + + B2 = Pm * S2 / 2 + ntail = 1 + nExo + Wq = zeros(Tv, nPast, nPast * nPast) + for i in 1:nPast, j in 1:nPast + Wq[:, (i-1)*nPast+j] = B2[:, (i-1)*na+j] + end + Wl_t = [zeros(Tv, nPast, nPast) for _ in 1:ntail] + for t in 1:ntail, k in 1:nPast + Wl_t[t][:, k] = B2[:, (k-1)*na+nPast+t] + B2[:, (nPast+t-1)*na+k] + end + Bc = B2[:, [(i-1)*na + j for i in nPast+1:na for j in nPast+1:na]] + MM = ℒ.kron(M, M) + return M, mc, V, B2, Wq, Wl_t, Bc, MM +end + +""" +Adjoint of `cubic_derived_matrices`: fold cotangents on the derived blocks back +onto `S1` and `S2`. `MM = kron(M, M)` is resolved onto `M` first, so it must be +accumulated before this is called. +""" +function cubic_derived_pullback!(∂S1, ∂S2, ∂M, ∂mc, ∂V, ∂Wq, ∂Wl_t, ∂Bc, ∂MM, M, + Pm, nPast::Int, nExo::Int, na::Int) + # MM = kron(M, M) + ∂M = ∂M .+ kron_adjoint_A(∂MM, M, nPast, nPast, nPast, nPast) .+ + kron_adjoint_B(∂MM, M, nPast, nPast, nPast, nPast) + + # A1 = Pm S1 ; M, mc, V are its column blocks + ∂A1 = hcat(∂M, reshape(∂mc, nPast, 1), ∂V) + ∂S1 .+= Pm' * ∂A1 + + # the K₂ splits, all linear scatters out of B2 + ∂B2 = zeros(eltype(∂S2), nPast, na * na) + for i in 1:nPast, j in 1:nPast + @views ∂B2[:, (i-1)*na+j] .+= ∂Wq[:, (i-1)*nPast+j] + end + for t in 1:(1+nExo), k in 1:nPast + @views ∂B2[:, (k-1)*na+nPast+t] .+= ∂Wl_t[t][:, k] + @views ∂B2[:, (nPast+t-1)*na+k] .+= ∂Wl_t[t][:, k] + end + for (col, r) in enumerate([(i-1)*na + j for i in nPast+1:na for j in nPast+1:na]) + @views ∂B2[:, r] .+= ∂Bc[:, col] + end + ∂S2 .+= Pm' * ∂B2 ./ 2 + return ∂S1, ∂S2 +end + """ Constant structure of the cubic augmented system: `z = [x₁; x₂; x₃; q₁₁; q₁₂; q₁₁₁]` over the retained rows (past states plus @@ -243,26 +312,8 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 C[i, 2nr+pos[j]] = 1.0 end - # aₙ = M a + mc + V ε - A1 = Pm * S1 - M = A1[:, 1:nPast] - mc = A1[:, nPast+1] - V = A1[:, nPast+2:na] - - # bₙ = M b + B2·K₂, with K₂ split into its (a,a), (a,tail)+(tail,a) and - # (tail,tail) parts so each can be routed onto the right state block. - B2 = Pm * S2 / 2 ntail = 1 + nExo - Wq = zeros(Tv, nPast, nPast * nPast) - for i in 1:nPast, j in 1:nPast - Wq[:, (i-1)*nPast+j] = B2[:, (i-1)*na+j] - end - Wl_t = [zeros(Tv, nPast, nPast) for _ in 1:ntail] - for t in 1:ntail, k in 1:nPast - Wl_t[t][:, k] = B2[:, (k-1)*na+nPast+t] + B2[:, (nPast+t-1)*na+k] - end - Bc = B2[:, [(i-1)*na + j for i in nPast+1:na for j in nPast+1:na]] - MM = ℒ.kron(M, M) + M, mc, V, B2, Wq, Wl_t, Bc, MM = cubic_derived_matrices(S1, S2, Pm, nPast, nExo, na) # Observation rows are a selection of the x₁, x₂ and x₃ blocks; carrying the # three positions lets the recursion index instead of running a gemm with a @@ -281,7 +332,7 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, op1, op2, op3, r1, r2, r3, i11, i12, i111, nq11, nq12, nq111, exp2, can2, exp3, can3, can2_ij, can3_ijk, - M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail, Tv) + M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail) end """ @@ -290,7 +341,7 @@ Preallocated buffers for `cubic_kalman_step!`. The step is called period, and it is entirely allocation-bound — it does a few hundred flops but allocated ~13 kB per call before these buffers existed. """ -function cubic_kalman_workspace(sys, Tv = sys.Tv) +function cubic_kalman_workspace(sys, Tv = eltype(sys.S1)) (; nr, nPast, na, ntail) = sys nP2 = nPast * nPast zeros(n...) = Base.zeros(Tv, n...) @@ -304,7 +355,8 @@ function cubic_kalman_workspace(sys, Tv = sys.Tv) Q11 = zeros(nPast, nPast), Q12 = zeros(nPast, nPast), Q111 = zeros(nPast, nP2), R2 = zeros(nPast, nPast), Tmp = zeros(nPast, nPast), MQ111 = zeros(nPast, nP2), R3 = zeros(nPast, nP2), - Wl = zeros(nPast, nPast), t2 = zeros(nP2), vv = zeros(nP2)) + Wl = zeros(nPast, nPast), t2 = zeros(nP2), vv = zeros(nP2), + scratch_out = zeros(sys.nz)) end """ @@ -444,9 +496,128 @@ function cubic_kalman_step!(out::AbstractVector, sys, z::AbstractVector, ε::Abs return out end +""" +Adjoint of `cubic_kalman_step!` with respect to the solution matrices, at a fixed +`(z, ε)`. Everything built from `z` and `ε` alone — `aug*`, `K₂`, `K₁₂`, `K₃`, +`Q₁₁`, `Q₁₂`, `Q₁₁₁`, `tail`, `tt` — is constant here, so only the paths through +`S1, S2, S3` and the derived blocks carry cotangents. + +Accumulates into `∂` in place; call `cubic_derived_pullback!` afterwards to fold +the derived-block cotangents onto `S1` and `S2`. +""" +function cubic_kalman_step_pullback!(∂, sys, z::AbstractVector, ε::AbstractVector, + ∂out::AbstractVector, ws) + (; nr, nPast, nExo, na, S1, S2, Pm, r1, r2, r3, i11, i12, i111, + exp2, can2_ij, can3_ijk, M, mc, V, Wq, Wl_t, Bc, MM, ntail) = sys + nP = nPast + nP2 = nP * nP + + # ── recompute the forward intermediates the adjoint needs ──────────────── + cubic_kalman_step!(ws.scratch_out, sys, z, ε, ws) + (; a, b, p, q11, q12, q111, tail, tt, aug1, aug2, aug3, K2, K12, K3, + x2n, u, v, bn, wc, Q11, Q12, Q111, Wl, t2, vv) = ws + MQ11 = M * Q11 + MQ12 = M * Q12 + MQ111 = M * Q111 + + o11 = first(i11) - 1 + o12 = first(i12) - 1 + o111 = first(i111) - 1 + + ∂u = zeros(eltype(∂out), nP); ∂v = zeros(eltype(∂out), nP) + ∂t2 = zeros(eltype(∂out), nP2); ∂vv = zeros(eltype(∂out), nP2) + ∂wc = zeros(eltype(∂out), nP); ∂bn = zeros(eltype(∂out), nP) + ∂R2 = zeros(eltype(∂out), nP, nP) + ∂R3 = zeros(eltype(∂out), nP, nP2) + + # ── seed from the output blocks ────────────────────────────────────────── + @inbounds for (s, (i, j)) in enumerate(can2_ij) + g = ∂out[o11+s] + ∂t2[(i-1)*nP+j] += g + ∂u[i] += g * v[j]; ∂v[j] += g * u[i] + ∂v[i] += g * u[j]; ∂u[j] += g * v[i] + ∂v[i] += g * v[j]; ∂v[j] += g * v[i] + end + @inbounds for i in 1:nP, j in 1:nP + g = ∂out[o12+(i-1)*nP+j] + ∂R2[i, j] += g + ∂u[i] += g * wc[j]; ∂wc[j] += g * u[i] + ∂v[i] += g * bn[j]; ∂bn[j] += g * v[i] + end + @inbounds for (s, (i, j, k)) in enumerate(can3_ijk) + g = ∂out[o111+s] + ∂R3[i, (j-1)*nP+k] += g + ∂t2[(i-1)*nP+j] += g * v[k]; ∂v[k] += g * t2[(i-1)*nP+j] + ∂v[i] += g * t2[(j-1)*nP+k]; ∂t2[(j-1)*nP+k] += g * v[i] + ∂t2[(i-1)*nP+k] += g * v[j]; ∂v[j] += g * t2[(i-1)*nP+k] + ∂u[i] += g * vv[(j-1)*nP+k]; ∂vv[(j-1)*nP+k] += g * u[i] + ∂vv[(i-1)*nP+j] += g * u[k]; ∂u[k] += g * vv[(i-1)*nP+j] + ∂v[i] += g * u[j] * v[k]; ∂u[j] += g * v[i] * v[k]; ∂v[k] += g * v[i] * u[j] + ∂vv[(i-1)*nP+j] += g * v[k]; ∂v[k] += g * vv[(i-1)*nP+j] + end + + # vv[i,j] = v_i v_j + @inbounds for i in 1:nP, j in 1:nP + gv = ∂vv[(i-1)*nP+j] + ∂v[i] += gv * v[j]; ∂v[j] += gv * v[i] + end + + # R3 = MQ111 * MM' + ∂MQ111 = ∂R3 * MM + ∂.MM .+= ∂R3' * MQ111 + + # R2 = MQ11 Wl' + MQ111 Wq' + MQ12 M' + ∂MQ11 = ∂R2 * Wl + ∂Wl = ∂R2' * MQ11 + ∂MQ111 .+= ∂R2 * Wq + ∂.Wq .+= ∂R2' * MQ111 + ∂MQ12 = ∂R2 * M + ∂M = ∂R2' * MQ12 + + # t2 = rowvec(MQ11 M') + ∂R2t = reshape(∂t2, nP, nP)' # ∂R2t[i,r] = ∂t2[(i-1)nP+r] + ∂MQ11 .+= ∂R2t * M + ∂M .+= ∂R2t' * MQ11 + + # MQ11 = M Q11, MQ12 = M Q12, MQ111 = M Q111 + ∂M .+= ∂MQ11 * Q11' .+ ∂MQ12 * Q12' .+ ∂MQ111 * Q111' + + # wc = Bc tt ; Wl = Σ_t tail[t] Wl_t[t] + ∂.Bc .+= ∂wc * tt' + @inbounds for t in 1:ntail + ∂.Wl_t[t] .+= tail[t] .* ∂Wl + end + + # u = M a ; v = mc + V ε + ∂M .+= ∂u * a' + ∂.mc .+= ∂v + ∂.V .+= ∂v * ε' + ∂.M .+= ∂M + + # bn = Pm x2n, and x2n is itself an output block + ∂x1n = ∂out[1:nr] + ∂x2n = ∂out[nr+1:2nr] .+ Pm' * ∂bn + ∂x3n = ∂out[2nr+1:3nr] + + # x1n = S1 aug1 ; x2n = S1 aug2 + ½ S2 K2 ; x3n = S1 aug3 + S2 K12 + ⅙ S3 K3 + ∂.S1 .+= ∂x1n * aug1' .+ ∂x2n * aug2' .+ ∂x3n * aug3' + ∂.S2 .+= (∂x2n * K2') ./ 2 .+ ∂x3n * K12' + ∂.S3 .+= (∂x3n * K3') ./ 6 + return ∂ +end + +# Zeroed cotangent accumulators matching the system's blocks. +function cubic_kalman_cotangents(sys) + T = eltype(sys.S1) + (; S1 = zeros(T, size(sys.S1)), S2 = zeros(T, size(sys.S2)), S3 = zeros(T, size(sys.S3)), + M = zeros(T, size(sys.M)), mc = zeros(T, length(sys.mc)), V = zeros(T, size(sys.V)), + Wq = zeros(T, size(sys.Wq)), Wl_t = [zeros(T, size(w)) for w in sys.Wl_t], + Bc = zeros(T, size(sys.Bc)), MM = zeros(T, size(sys.MM))) +end + # Allocating convenience wrapper, used by the tests. function cubic_kalman_step(sys, z::AbstractVector, ε::AbstractVector, ws = cubic_kalman_workspace(sys)) - return cubic_kalman_step!(Vector{sys.Tv}(undef, sys.nz), sys, z, ε, ws) + return cubic_kalman_step!(Vector{eltype(sys.S1)}(undef, sys.nz), sys, z, ε, ws) end # E[f(z,·)] and Var(f(z,·)) under ε ~ N(0,I), exactly. @@ -457,7 +628,7 @@ end # stacking buffer across periods. function cubic_kalman_moments(sys, z, nodes, wts; buf = nothing, ws = cubic_kalman_workspace(sys)) nz = sys.nz - Fm = buf === nothing ? Matrix{sys.Tv}(undef, nz, length(nodes)) : buf + Fm = buf === nothing ? Matrix{eltype(sys.S1)}(undef, nz, length(nodes)) : buf @inbounds for (n, ε) in enumerate(nodes) cubic_kalman_step!(view(Fm, :, n), sys, z, ε, ws) end @@ -472,7 +643,7 @@ end function build_cubic_kalman_transition(sys, nodes, wts; ws = cubic_kalman_workspace(sys)) # Only the mean is needed here, so skip the variance the moment routine would # otherwise form — an nz×nz gemm per basis vector, nz+1 of them. - Fm = Matrix{sys.Tv}(undef, sys.nz, length(nodes)) + Fm = Matrix{eltype(sys.S1)}(undef, sys.nz, length(nodes)) mean_at = function (z) @inbounds for (n, ε) in enumerate(nodes) cubic_kalman_step!(view(Fm, :, n), sys, z, ε, ws) @@ -500,7 +671,7 @@ polynomially in the number of shocks instead of exponentially. Afterwards no quadrature is needed at all, per period or otherwise. """ function build_cubic_kalman_system(sys, basis; ws = cubic_kalman_workspace(sys)) - nz, N, Tv = sys.nz, basis.N, sys.Tv + nz, N, Tv = sys.nz, basis.N, eltype(sys.S1) Fm = Matrix{Tv}(undef, nz, N) coefficients_at = function (z) @inbounds for (p, ε) in enumerate(basis.pts) @@ -527,6 +698,44 @@ function build_cubic_kalman_system(sys, basis; ws = cubic_kalman_workspace(sys)) return 𝒜, c, c₀, Λ end +""" +Adjoint of `build_cubic_kalman_system`. The build is linear in the collected step +evaluations — `C(z) = F(z) W`, then `c = C(0)m`, `c₀ = vec C(0)`, +`𝒜[:,j] = ΔC_j m`, `Λ[:,j] = vec ΔC_j` with `ΔC_j = C(e_j) − C(0)` — so those +maps transpose directly, and the only real work is replaying the step adjoint at +the same `(n_z + 1)·N` points the forward pass visited. +""" +function build_cubic_kalman_system_pullback!(∂, sys, basis, ∂𝒜, ∂c, ∂c₀, ∂Λ; + ws = cubic_kalman_workspace(sys)) + nz, N = sys.nz, basis.N + m, W = basis.m, basis.W + + # C(0) is hit by c, by c₀, and negatively by every ΔC_j. + ∂C0 = ∂c * m' .+ reshape(∂c₀, nz, N) + @inbounds for j in 1:nz + ∂C0 .-= view(∂𝒜, :, j) * m' .+ reshape(view(∂Λ, :, j), nz, N) + end + + ∂F = ∂C0 * W' + zj = zeros(nz) + @inbounds for (p, ε) in enumerate(basis.pts) + cubic_kalman_step_pullback!(∂, sys, zj, ε, view(∂F, :, p), ws) + end + + @inbounds for j in 1:nz + ∂ΔC = view(∂𝒜, :, j) * m' .+ reshape(view(∂Λ, :, j), nz, N) + ∂F = ∂ΔC * W' + fill!(zj, 0.0); zj[j] = 1.0 + for (p, ε) in enumerate(basis.pts) + cubic_kalman_step_pullback!(∂, sys, zj, ε, view(∂F, :, p), ws) + end + end + + cubic_derived_pullback!(∂.S1, ∂.S2, ∂.M, ∂.mc, ∂.V, ∂.Wq, ∂.Wl_t, ∂.Bc, ∂.MM, + sys.M, sys.Pm, sys.nPast, sys.nExo, sys.na) + return ∂ +end + """ Kalman recursion on the cubic augmented state. Mirrors `run_quadratic_kalman`: the noise covariance is rebuilt from the current state estimate every period, @@ -547,7 +756,7 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; # Promote over every differentiable input. The preallocated buffers below fix # the element type, so missing one makes forward-mode AD fail with respect to # exactly that argument. - Tv = promote_type(sys.Tv, eltype(data_in_deviations), + Tv = promote_type(eltype(sys.S1), eltype(data_in_deviations), measurement_error === nothing ? Float64 : eltype(measurement_error)) Hm = if measurement_error === nothing @@ -657,6 +866,90 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; end +""" +Taped forward pass plus adjoint for the cubic Kalman recursion. Mirrors +`quadratic_kalman_recursion`'s verified adjoint; the one structural difference is +the noise term, `Q = CΨC'` with `vec(C) = c₀ + Λz`, in place of `GG' + Q_H`. + +`Q̄` is symmetric here because `P̄p` is symmetrised before use, so the cotangent of +`C` is `2 Q̄ C Ψ` rather than `(Q̄ + Q̄')CΨ`. +""" +function cubic_kalman_recursion_taped(𝒜, c, c₀, Λ, Ψ, Hm, Y, 𝒞, z0, Σ0, nz, N, + presample_periods, on_failure_loglikelihood) + n_obs, nT = size(Y) + z = copy(z0); Pc = copy(Σ0) + zs = Vector{Vector{Float64}}(); Ps = Vector{Matrix{Float64}}() + Cs = Vector{Matrix{Float64}}(); vs = Vector{Vector{Float64}}() + CPs = Vector{Matrix{Float64}}(); Fis = Vector{Matrix{Float64}}() + Ks = Vector{Matrix{Float64}}() + ll = 0.0; log2pi = log(2π) + + for t in 1:nT + push!(zs, copy(z)); push!(Ps, copy(Pc)) + C = reshape(c₀ + Λ * z, nz, N) + Q = C * Ψ * C' + zp = 𝒜 * z + c + Pp = 𝒜 * Pc * 𝒜' + Q; Pp = (Pp + Pp') / 2 + v = Y[:, t] - 𝒞 * zp + CP = 𝒞 * Pp + F = CP * 𝒞' + Hm; F = (F + F') / 2 + Fc = ℒ.cholesky(F, check = false) + ℒ.issuccess(Fc) || return on_failure_loglikelihood, nothing + Fi = inv(Fc) + if t > presample_periods + ll -= 0.5 * (ℒ.dot(v, Fi * v) + ℒ.logdet(Fc) + n_obs * log2pi) + end + K = CP' * Fi + z = zp + K * v + Pc = Pp - K * CP; Pc = (Pc + Pc') / 2 + push!(Cs, C); push!(vs, v); push!(CPs, CP); push!(Fis, Fi); push!(Ks, K) + end + return ll, (; zs, Ps, Cs, vs, CPs, Fis, Ks) +end + +function cubic_kalman_recursion_pullback(tape, 𝒜, c, c₀, Λ, Ψ, 𝒞, nz, N, n_obs, nT, + presample_periods, ∂ll) + (; zs, Ps, Cs, vs, CPs, Fis, Ks) = tape + 𝒜̄ = zeros(nz, nz); c̄ = zeros(nz); c̄₀ = zeros(length(c₀)); Λ̄ = zeros(size(Λ)) + H̄m = zeros(n_obs, n_obs); Ȳ = zeros(n_obs, nT) + z̄ = zeros(nz); P̄ = zeros(nz, nz) + + for t in nT:-1:1 + z_, P_, C, v, CP, Fi, K = zs[t], Ps[t], Cs[t], vs[t], CPs[t], Fis[t], Ks[t] + P̄p = copy(P̄) + K̄ = -P̄ * CP' + C̄P = -K' * P̄ + z̄p = copy(z̄) + K̄ .+= z̄ * v' + v̄ = K' * z̄ + C̄P .+= Fi * K̄' + F̄ = -Fi * (CP * K̄) * Fi + if t > presample_periods + v̄ .+= -∂ll * (Fi * v) + F̄ .+= ∂ll * 0.5 * (Fi * v * v' * Fi - Fi) + end + F̄ = (F̄ + F̄') / 2 + C̄P .+= F̄ * 𝒞 + H̄m .+= F̄ + P̄p .+= 𝒞' * C̄P + z̄p .+= -𝒞' * v̄ + Ȳ[:, t] .+= v̄ + P̄p = (P̄p + P̄p') / 2 + 𝒜̄ .+= 2 .* (P̄p * 𝒜 * P_) + P̄ = 𝒜' * P̄p * 𝒜 + Q̄ = P̄p + C̄ = 2 .* (Q̄ * C * Ψ) + vC̄ = vec(C̄) + c̄₀ .+= vC̄ + Λ̄ .+= vC̄ * z_' + 𝒜̄ .+= z̄p * z_' + c̄ .+= z̄p + z̄ = 𝒜' * z̄p + Λ' * vC̄ + end + return 𝒜̄, c̄, c̄₀, Λ̄, H̄m, Ȳ, z̄, P̄ +end + + # ── standard filter interface ──────────────────────────────────────────────── function calculate_loglikelihood(::Val{:cubic_kalman}, ::Val{:pruned_third_order}, @@ -683,4 +976,87 @@ function calculate_loglikelihood(::Val{:cubic_kalman}, lyapunov_algorithm = lyapunov_algorithm) end +function rrule(::typeof(calculate_loglikelihood), + ::Val{:cubic_kalman}, + ::Val{:pruned_third_order}, + observables_index::Vector{Int}, + 𝐒, + data_in_deviations::AbstractMatrix, + constants, + state, + workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) + sys = build_cubic_kalman_system_from_constants(constants, 𝐒[1], 𝐒[2], 𝐒[3], observables_index) + nz = sys.nz + n_obs, nT = size(data_in_deviations) + presample = normalize_presample_periods(presample_periods, nT) + + Hm = measurement_error === nothing ? zeros(n_obs, n_obs) : + measurement_error isa AbstractMatrix ? Matrix{Float64}(measurement_error) : + Matrix{Float64}(ℒ.Diagonal(collect(measurement_error))) + + ws = cubic_kalman_workspace(sys) + basis = cubic_noise_basis(sys.nExo) + 𝒜, c, c₀, Λ = build_cubic_kalman_system(sys, basis; ws = ws) + N, Ψ, 𝒞 = basis.N, basis.Ψ, sys.C + + z₀ = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c + C₀ = reshape(c₀ + Λ * z₀, nz, N) + Q₀ = C₀ * Ψ * C₀'; Q₀ = (Q₀ + Q₀') / 2 + Σ₀ = qkf_lyapunov(𝒜, Q₀; workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + + llh, tape = cubic_kalman_recursion_taped(𝒜, c, c₀, Λ, Ψ, Hm, Matrix(data_in_deviations), + 𝒞, z₀, Σ₀, nz, N, presample, on_failure_loglikelihood) + + nine(x...) = (NoTangent(), NoTangent(), NoTangent(), NoTangent(), x[1], x[2], + NoTangent(), x[3], NoTangent()) + + if !isfinite(llh) || tape === nothing + return llh, _ -> nine(NoTangent(), NoTangent(), NoTangent()) + end + + function cubic_kalman_loglikelihood_pullback(∂llh_bar) + ∂llh = unthunk(∂llh_bar) + 𝒜̄, c̄, c̄₀, Λ̄, _, Ȳ, z̄₀, Σ̄₀ = + cubic_kalman_recursion_pullback(tape, 𝒜, c, c₀, Λ, Ψ, 𝒞, nz, N, n_obs, nT, + presample, ∂llh) + + # Σ₀ = 𝒜Σ₀𝒜' + Q₀ ⇒ X solves X = 𝒜'X𝒜 + Σ̄₀ + X = qkf_lyapunov(Matrix(𝒜'), (Σ̄₀ + Σ̄₀') / 2; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + 𝒜̄ .+= 2 .* (X * 𝒜 * Σ₀) + Q̄₀ = (X + X') / 2 + C̄₀ = 2 .* (Q̄₀ * C₀ * Ψ) + vC̄₀ = vec(C̄₀) + c̄₀ = c̄₀ .+ vC̄₀ + Λ̄ = Λ̄ .+ vC̄₀ * z₀' + z̄₀ = z̄₀ .+ Λ' * vC̄₀ + + # z₀ = (I − 𝒜)⁻¹ c + y = (Matrix{Float64}(ℒ.I(nz)) - 𝒜)' \ z̄₀ + c̄ = c̄ .+ y + 𝒜̄ .+= y * z₀' + + ∂ = cubic_kalman_cotangents(sys) + build_cubic_kalman_system_pullback!(∂, sys, basis, 𝒜̄, c̄, c̄₀, Λ̄; ws = ws) + + # scatter the retained rows back onto the full solution matrices + ∂𝐒1 = zeros(size(𝐒[1])); ∂𝐒2 = zeros(size(𝐒[2])); ∂𝐒3 = zeros(size(𝐒[3])) + ∂𝐒1[sys.oas, :] = ∂.S1 + ∂𝐒2[sys.oas, :] = ∂.S2 + ∂𝐒3[sys.oas, :] = ∂.S3 + ∂state = [zeros(length(s)) for s in state] + return nine([∂𝐒1, ∂𝐒2, ∂𝐒3], Ȳ, ∂state) + end + + return llh, cubic_kalman_loglikelihood_pullback +end + end # @stable diff --git a/src/rrules.jl b/src/rrules.jl index 103ab7e5a..20553c230 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -1877,19 +1877,12 @@ function rrule(::typeof(get_loglikelihood), else measurement_error != 0 end - # The quadratic Kalman filter carries its own hand-written adjoint, which - # includes the measurement-error covariance, so the guard does not apply to it. - if me_active && filter != :quadratic_kalman + # The quadratic and cubic Kalman filters carry their own hand-written adjoints, + # which include the measurement-error covariance, so the guard skips them. + if me_active && filter ∉ (:quadratic_kalman, :cubic_kalman) error("Reverse-mode automatic differentiation of the Kalman likelihood with measurement error (`measurement_error`) is not yet supported. Use forward-mode AD (e.g. `AutoForwardDiff`) or a gradient-free sampler.") end - # The cubic Kalman filter has no hand-written adjoint. Without this guard the - # `llh_rrule === nothing` branch below would hand back an all-zero gradient - # and no error at all, which a sampler would happily run on. - if filter == :cubic_kalman - error("Reverse-mode automatic differentiation of the cubic Kalman filter (`filter = :cubic_kalman`) is not supported. Use forward-mode AD (e.g. `AutoForwardDiff`), which is exact for this filter and matches finite differences to ~1e-11, or a gradient-free sampler.") - end - opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -1972,10 +1965,10 @@ function rrule(::typeof(get_loglikelihood), end # ── step 3: calculate_loglikelihood ── - # The quadratic Kalman filter is the only one whose inner rrule takes the - # measurement-error covariance; for the others it is inactive (the guard above) - # and the kwarg would not be accepted. - me_kw = filter == :quadratic_kalman ? + # The quadratic and cubic Kalman filters are the ones whose inner rrules take + # the measurement-error covariance; for the others it is inactive (the guard + # above) and the kwarg would not be accepted. + me_kw = filter ∈ (:quadratic_kalman, :cubic_kalman) ? (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations)) : NamedTuple() diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl index 98fe316a5..e72f932f7 100644 --- a/test/test_cubic_kalman.jl +++ b/test/test_cubic_kalman.jl @@ -143,20 +143,78 @@ import AxisKeys: KeyedArray # The particle filter's log-likelihood is downward-biased by about Var/2. @test abs(ll_ckf - (m + (sum(x -> (x - m)^2, ll_pf) / (length(ll_pf) - 1)) / 2)) < 0.05 * T - # 4b. forward-mode AD is exact; reverse mode must *fail loudly* rather than - # fall through to the generic zero-gradient path, which a sampler would - # run on without noticing. + # 4b. gradients. Both modes are checked against central differences, and + # reverse mode is checked on every measurement-error shape — a wrong `H` + # reaching the adjoint but not the primal gives a finite, plausible, wrong + # gradient rather than an error. f = p -> get_loglikelihood(RBC_ckf, data, p; algorithm = :pruned_third_order, filter = :cubic_kalman, measurement_error = mev) - g_ad = ForwardDiff.gradient(f, pars) g_fd = FiniteDifferences.grad(central_fdm(5, 1), f, pars)[1] - @test maximum(abs.(g_ad .- g_fd) ./ max.(1.0, abs.(g_fd))) < 1e-7 - @test !all(iszero, g_ad) - # without measurement error the unrelated measurement-error guard cannot be - # what fires, so this pins the cubic-filter guard specifically + @test !all(iszero, g_fd) + @test maximum(abs.(ForwardDiff.gradient(f, pars) .- g_fd) ./ max.(1.0, abs.(g_fd))) < 1e-7 + @test maximum(abs.(Zygote.gradient(f, pars)[1] .- g_fd) ./ max.(1.0, abs.(g_fd))) < 1e-7 + f_nome = p -> get_loglikelihood(RBC_ckf, data, p; algorithm = :pruned_third_order, filter = :cubic_kalman) - @test_throws ErrorException Zygote.gradient(f_nome, pars) + g_fd_nome = FiniteDifferences.grad(central_fdm(5, 1), f_nome, pars)[1] + @test maximum(abs.(Zygote.gradient(f_nome, pars)[1] .- g_fd_nome) ./ max.(1.0, abs.(g_fd_nome))) < 1e-6 + + mev_mat = [3e-5 1e-5; 1e-5 4e-5] # non-diagonal covariance + f_mat = p -> get_loglikelihood(RBC_ckf, data, p; algorithm = :pruned_third_order, + filter = :cubic_kalman, measurement_error = mev_mat) + g_fd_mat = FiniteDifferences.grad(central_fdm(5, 1), f_mat, pars)[1] + @test maximum(abs.(Zygote.gradient(f_mat, pars)[1] .- g_fd_mat) ./ max.(1.0, abs.(g_fd_mat))) < 1e-6 + + # 4c. the two hand-written adjoints the chain rests on, checked in isolation + # against ForwardDiff so a regression localises instead of just moving the + # end-to-end number. + let ws = MacroModelling.cubic_kalman_workspace(sys), Pm = sys.Pm, + nP = sys.nPast, nE = sys.nExo, na = sys.na, + n1 = length(sys.S1), n2 = length(sys.S2), n3 = length(sys.S3) + function rebuild(θ) + S1 = reshape(θ[1:n1], size(sys.S1)) + S2 = reshape(θ[n1+1:n1+n2], size(sys.S2)) + S3 = reshape(θ[n1+n2+1:end], size(sys.S3)) + M, mc, V, B2, Wq, Wl_t, Bc, MM = MacroModelling.cubic_derived_matrices(S1, S2, Pm, nP, nE, na) + merge(sys, (; S1, S2, S3, M, mc, V, B2, Wq, Wl_t, Bc, MM, Tv = eltype(θ))) + end + θ0 = vcat(vec(sys.S1), vec(sys.S2), vec(sys.S3)) + function fold(∂) + MacroModelling.cubic_derived_pullback!(∂.S1, ∂.S2, ∂.M, ∂.mc, ∂.V, + ∂.Wq, ∂.Wl_t, ∂.Bc, ∂.MM, + sys.M, Pm, nP, nE, na) + return vcat(vec(∂.S1), vec(∂.S2), vec(∂.S3)) + end + + # the step's adjoint + zr = 0.05 .* randn(sys.nz); εr = randn(nE); ∂out = randn(sys.nz) + function step_scalar(θ) + s2 = rebuild(θ) + w2 = MacroModelling.cubic_kalman_workspace(s2, eltype(θ)) + out = MacroModelling.cubic_kalman_step!(Vector{eltype(θ)}(undef, sys.nz), s2, zr, εr, w2) + return ℒ.dot(∂out, out) + end + ∂s = MacroModelling.cubic_kalman_cotangents(sys) + MacroModelling.cubic_kalman_step_pullback!(∂s, sys, zr, εr, ∂out, ws) + gs = ForwardDiff.gradient(step_scalar, θ0) + @test maximum(abs, fold(∂s) - gs) / max(1e-12, maximum(abs, gs)) < 1e-10 + + # The build's adjoint. Unlike the step's, it folds the derived-block + # cotangents onto S1/S2 itself, so no `fold` here. + w𝒜 = randn(sys.nz, sys.nz); wc = randn(sys.nz) + wc0 = randn(sys.nz * basis.N); wΛ = randn(sys.nz * basis.N, sys.nz) + function build_scalar(θ) + s2 = rebuild(θ) + w2 = MacroModelling.cubic_kalman_workspace(s2, eltype(θ)) + A, cc, c0, L = MacroModelling.build_cubic_kalman_system(s2, basis; ws = w2) + return ℒ.dot(w𝒜, A) + ℒ.dot(wc, cc) + ℒ.dot(wc0, c0) + ℒ.dot(wΛ, L) + end + ∂b = MacroModelling.cubic_kalman_cotangents(sys) + MacroModelling.build_cubic_kalman_system_pullback!(∂b, sys, basis, w𝒜, wc, wc0, wΛ; ws = ws) + gb = ForwardDiff.gradient(build_scalar, θ0) + gb_mine = vcat(vec(∂b.S1), vec(∂b.S2), vec(∂b.S3)) + @test maximum(abs, gb_mine - gb) / max(1e-12, maximum(abs, gb)) < 1e-10 + end # 5. gating: the filter is only defined on the pruned third-order solution. # At any other order it falls back to the inversion filter, which admits From f98f886f9efea989e90541b4948078fb980131dd Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 31 Jul 2026 19:49:04 +0200 Subject: [PATCH 15/16] Exploit the solution matrices' sparsity in the cubic filter's assembly Profiling the build showed it is (nz+1)*N step evaluations and essentially nothing else, and that inside a step one operation dominated: contracting the Kronecker input K3 against S3. That matrix is very wide and very sparse -- 8x1331 on a four-shock model, 23% dense -- so the product is memory-bound, and it measured 5.5 of the 10.8 us a step took. Only the structurally nonzero columns of S2 and S3 are now kept (536 of 1331 for S3, 65 of 121 and 40 of 66 for the two S2 paths), which shrinks both the vector that has to be built and the product that consumes it, and drops those iterations from the branchy K3 loop as well. Liveness is taken from the stored pattern of the *sparse* solution matrices, not from numerical zeros of the densified copy. A column that merely happens to vanish at one parameter draw may be nonzero at the next, and dropping it would silently zero a real derivative rather than fail. The existing exactness test guards this directly: it compares the restricted step against a reference built from full Kronecker products. Measured on a four-shock, six-past-state model (nz = 137): step 10.8 -> 6.4 us build 56.2 -> 34.0 ms step pullback 26.3 -> 18.0 us The build's pullback also collapses two sums that were being accumulated one column at a time, each allocating an nz-by-N array per pass. Checked what else sparsity could buy and it does not: Psi is 9% dense and block-diagonal by monomial parity, but it appears only in the smaller of the two products forming Q, and its rank is N-1, so factoring removes just the constant monomial. Also validates the reverse-mode adjoint on a second, wider model: Zygote against central differences is 1.1e-9 there, with the gradient at 4.8x the primal. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT --- docs/src/filters.md | 15 +++++ src/filter/cubic_kalman.jl | 124 +++++++++++++++++++++++++++---------- test/test_cubic_kalman.jl | 9 +-- 3 files changed, 112 insertions(+), 36 deletions(-) diff --git a/docs/src/filters.md b/docs/src/filters.md index 84a8e0fe6..a5d5fd5a3 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -670,6 +670,21 @@ also where a tensor Gauss-Hermite rule is left behind — its node count grows a basis. The quadrature path is retained and the analytic assembly is tested against it rather than assumed. +Assembling the system costs ``(n_z+1)N`` evaluations of the step, and essentially nothing +else, so the step is where the sparsity is worth spending. Its dominant term is the +contraction of the Kronecker input ``K_3`` against ``\mathbf{S}_3`` — a very wide, very +sparse matrix (``8\times1331`` on a four-shock model), which made it memory-bound and half +the cost of a step. Only the *structurally* nonzero columns are kept — 536 of 1331 there — +which shrinks both the vector that has to be built and the product that consumes it. +Liveness comes from the stored pattern of the sparse solution matrices rather than from +numerical zeros of a densified copy: a column that merely happens to vanish at one parameter +draw may be nonzero at the next, and dropping it would silently zero a real derivative. + +``\Psi`` is sparser still (9% dense, block-diagonal by monomial parity, since +``\mathbb{E}[\varepsilon^\gamma] = 0`` unless every exponent is even) but exploiting that +is not worth it: it appears only in the smaller of the two products forming ``Q``, and its +rank is ``N-1``, so factoring it removes just the constant monomial. + ### Derivatives Both modes work and both match central differences to ``\sim10^{-10}``. diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index 274592471..ef9efe707 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -225,12 +225,28 @@ function cubic_derived_matrices(S1, S2, Pm, nPast::Int, nExo::Int, na::Int) end """ -Adjoint of `cubic_derived_matrices`: fold cotangents on the derived blocks back -onto `S1` and `S2`. `MM = kron(M, M)` is resolved onto `M` first, so it must be -accumulated before this is called. +Fold every intermediate cotangent back onto `S1`, `S2` and `S3`: the live-column +accumulators the step adjoint writes, and the derived blocks of +`cubic_derived_matrices`. `MM = kron(M, M)` is resolved onto `M` first, so it must +be accumulated before this is called. """ -function cubic_derived_pullback!(∂S1, ∂S2, ∂M, ∂mc, ∂V, ∂Wq, ∂Wl_t, ∂Bc, ∂MM, M, - Pm, nPast::Int, nExo::Int, na::Int) +function cubic_derived_pullback!(∂, sys) + (; M, Pm, nPast, nExo, na) = sys + ∂S1, ∂S2 = ∂.S1, ∂.S2 + ∂M, ∂mc, ∂V = ∂.M, ∂.mc, ∂.V + ∂Wq, ∂Wl_t, ∂Bc, ∂MM = ∂.Wq, ∂.Wl_t, ∂.Bc, ∂.MM + + # live-column cotangents accumulated by the step adjoint + @inbounds for (r, j) in enumerate(sys.k2cols) + @views ∂S2[:, j] .+= ∂.S2k2[:, r] + end + @inbounds for (r, j) in enumerate(sys.k12cols) + @views ∂S2[:, j] .+= ∂.S2k12[:, r] + end + @inbounds for (r, j) in enumerate(sys.k3cols) + @views ∂.S3[:, j] .+= ∂.S3k3[:, r] + end + # MM = kron(M, M) ∂M = ∂M .+ kron_adjoint_A(∂MM, M, nPast, nPast, nPast, nPast) .+ kron_adjoint_B(∂MM, M, nPast, nPast, nPast, nPast) @@ -315,6 +331,37 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 ntail = 1 + nExo M, mc, V, B2, Wq, Wl_t, Bc, MM = cubic_derived_matrices(S1, S2, Pm, nPast, nExo, na) + # The Kronecker inputs are contracted against 𝐒₂ and 𝐒₃, whose columns are + # largely structurally zero — a third-order solution has no cross-derivative + # for most index triples. Keeping only the live columns shrinks both the + # vector that has to be built and the (memory-bound, very wide) product that + # consumes it; on a four-shock model that is 536 of 1331 columns for 𝐒₃, and + # `S3 * K3` alone was half the cost of a step. + # Liveness is taken from the *structural* pattern of the sparse solution + # matrices, not from numerical zeros of a densified copy: a column that is + # merely zero at this parameter draw may be nonzero at the next one, and + # dropping it would silently zero a real derivative. Reading the stored + # pattern is also what the rest of the package assumes about 𝐒. + S2sp = 𝐒₂[oas, :] + S3sp = 𝐒₃[oas, :] + live(A, cols) = A isa SparseArrays.AbstractSparseMatrix ? + [j for j in cols if A.colptr[j+1] > A.colptr[j]] : + [j for j in cols if any(!iszero, view(A, :, j))] + + k2cols = live(S2sp, 1:na*na) + k2_ij = [(fld(j - 1, na) + 1, mod(j - 1, na) + 1) for j in k2cols] + S2k2 = S2[:, k2cols] + + k12all = [(i-1)*na + j for i in 1:na for j in 1:nPast] + k12cols = live(S2sp, k12all) + k12_ij = [(fld(j - 1, na) + 1, mod(j - 1, na) + 1) for j in k12cols] + S2k12 = S2[:, k12cols] + + k3cols = live(S3sp, 1:na*na*na) + k3_ijk = [(fld(j - 1, na * na) + 1, mod(fld(j - 1, na), na) + 1, mod(j - 1, na) + 1) + for j in k3cols] + S3k3 = S3[:, k3cols] + # Observation rows are a selection of the x₁, x₂ and x₃ blocks; carrying the # three positions lets the recursion index instead of running a gemm with a # 0/1 matrix, as the quadratic filter does with its two. @@ -332,7 +379,8 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, op1, op2, op3, r1, r2, r3, i11, i12, i111, nq11, nq12, nq111, exp2, can2, exp3, can3, can2_ij, can3_ijk, - M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail) + M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail, + k2cols, k2_ij, S2k2, k12cols, k12_ij, S2k12, k3cols, k3_ijk, S3k3) end """ @@ -349,7 +397,8 @@ function cubic_kalman_workspace(sys, Tv = eltype(sys.S1)) q11 = zeros(nP2), q12 = zeros(nP2), q111 = zeros(nPast^3), tail = zeros(ntail), tt = zeros(ntail * ntail), aug1 = zeros(na), aug1h = zeros(na), aug2 = zeros(na), aug3 = zeros(na), - K2 = zeros(na * na), K12 = zeros(na * na), K3 = zeros(na * na * na), + K2 = zeros(length(sys.k2cols)), K12 = zeros(length(sys.k12cols)), + K3 = zeros(length(sys.k3cols)), x1n = zeros(nr), x2n = zeros(nr), x3n = zeros(nr), u = zeros(nPast), v = zeros(nPast), bn = zeros(nPast), wc = zeros(nPast), Q11 = zeros(nPast, nPast), Q12 = zeros(nPast, nPast), Q111 = zeros(nPast, nP2), @@ -402,16 +451,15 @@ function cubic_kalman_step!(out::AbstractVector, sys, z::AbstractVector, ε::Abs aug1[nP+1+i] = ε[i]; aug1h[nP+1+i] = ε[i]; aug2[nP+1+i] = 0.0; aug3[nP+1+i] = 0.0 end - # Kronecker inputs, with the all-past blocks read from the state. - @inbounds for i in 1:na, j in 1:na - K2[(i-1)*na+j] = (i <= nP && j <= nP) ? q11[(i-1)*nP+j] : aug1[i] * aug1[j] + # Kronecker inputs, with the all-past blocks read from the state, and only at + # the columns 𝐒₂ and 𝐒₃ actually reach. + @inbounds for (r, (i, j)) in enumerate(sys.k2_ij) + K2[r] = (i <= nP && j <= nP) ? q11[(i-1)*nP+j] : aug1[i] * aug1[j] end - fill!(K12, 0.0) - @inbounds for i in 1:na, j in 1:nP - K12[(i-1)*na+j] = (i <= nP) ? q12[(i-1)*nP+j] : aug1h[i] * aug2[j] + @inbounds for (r, (i, j)) in enumerate(sys.k12_ij) + K12[r] = (i <= nP) ? q12[(i-1)*nP+j] : aug1h[i] * aug2[j] end - @inbounds for i in 1:na, j in 1:na, k in 1:na - r = ((i-1)*na + (j-1)) * na + k + @inbounds for (r, (i, j, k)) in enumerate(sys.k3_ijk) ci = i <= nP; cj = j <= nP; ck = k <= nP n = ci + cj + ck K3[r] = if n == 3 @@ -430,8 +478,9 @@ function cubic_kalman_step!(out::AbstractVector, sys, z::AbstractVector, ε::Abs end ℒ.mul!(x1n, S1, aug1) - ℒ.mul!(x2n, S1, aug2); ℒ.mul!(x2n, S2, K2, 0.5, 1.0) - ℒ.mul!(x3n, S1, aug3); ℒ.mul!(x3n, S2, K12, 1.0, 1.0); ℒ.mul!(x3n, S3, K3, 1/6, 1.0) + ℒ.mul!(x2n, S1, aug2); ℒ.mul!(x2n, sys.S2k2, K2, 0.5, 1.0) + ℒ.mul!(x3n, S1, aug3); ℒ.mul!(x3n, sys.S2k12, K12, 1.0, 1.0) + ℒ.mul!(x3n, sys.S3k3, K3, 1/6, 1.0) # New Kronecker blocks, kept affine in z. ℒ.mul!(u, M, a) # z-dependent, linear in a @@ -599,10 +648,13 @@ function cubic_kalman_step_pullback!(∂, sys, z::AbstractVector, ε::AbstractVe ∂x2n = ∂out[nr+1:2nr] .+ Pm' * ∂bn ∂x3n = ∂out[2nr+1:3nr] - # x1n = S1 aug1 ; x2n = S1 aug2 + ½ S2 K2 ; x3n = S1 aug3 + S2 K12 + ⅙ S3 K3 + # x1n = S1 aug1 ; x2n = S1 aug2 + ½ S2 K2 ; x3n = S1 aug3 + S2 K12 + ⅙ S3 K3. + # The 𝐒₂/𝐒₃ cotangents accumulate on the live columns and are scattered back + # once, in `cubic_derived_pullback!`. ∂.S1 .+= ∂x1n * aug1' .+ ∂x2n * aug2' .+ ∂x3n * aug3' - ∂.S2 .+= (∂x2n * K2') ./ 2 .+ ∂x3n * K12' - ∂.S3 .+= (∂x3n * K3') ./ 6 + ∂.S2k2 .+= (∂x2n * K2') ./ 2 + ∂.S2k12 .+= ∂x3n * K12' + ∂.S3k3 .+= (∂x3n * K3') ./ 6 return ∂ end @@ -612,7 +664,9 @@ function cubic_kalman_cotangents(sys) (; S1 = zeros(T, size(sys.S1)), S2 = zeros(T, size(sys.S2)), S3 = zeros(T, size(sys.S3)), M = zeros(T, size(sys.M)), mc = zeros(T, length(sys.mc)), V = zeros(T, size(sys.V)), Wq = zeros(T, size(sys.Wq)), Wl_t = [zeros(T, size(w)) for w in sys.Wl_t], - Bc = zeros(T, size(sys.Bc)), MM = zeros(T, size(sys.MM))) + Bc = zeros(T, size(sys.Bc)), MM = zeros(T, size(sys.MM)), + S2k2 = zeros(T, size(sys.S2k2)), S2k12 = zeros(T, size(sys.S2k12)), + S3k3 = zeros(T, size(sys.S3k3))) end # Allocating convenience wrapper, used by the tests. @@ -710,29 +764,35 @@ function build_cubic_kalman_system_pullback!(∂, sys, basis, ∂𝒜, ∂c, ∂ nz, N = sys.nz, basis.N m, W = basis.m, basis.W - # C(0) is hit by c, by c₀, and negatively by every ΔC_j. - ∂C0 = ∂c * m' .+ reshape(∂c₀, nz, N) - @inbounds for j in 1:nz - ∂C0 .-= view(∂𝒜, :, j) * m' .+ reshape(view(∂Λ, :, j), nz, N) - end - - ∂F = ∂C0 * W' + # C(0) is hit by c, by c₀, and negatively by every ΔC_j. Both of those sums + # collapse — Σⱼ ∂𝒜[:,j] m' = (Σⱼ ∂𝒜[:,j]) m' and reshape is linear — so this + # is one pass rather than n_z, each of which allocated an n_z×N array. + ones_nz = ones(eltype(∂𝒜), nz) + ∂C0 = ∂c * m' .+ reshape(∂c₀, nz, N) .- + (∂𝒜 * ones_nz) * m' .- reshape(∂Λ * ones_nz, nz, N) + + # ∂F = ∂ΔC W' with ∂ΔC = ∂𝒜[:,j] m' + reshape(∂Λ[:,j]); the first term is + # rank one and contracts to an outer product with W m, so nothing per-`j` + # needs to be materialised. + Wm = W * m + ∂F = Matrix{eltype(∂𝒜)}(undef, nz, N) zj = zeros(nz) + + ℒ.mul!(∂F, ∂C0, W') @inbounds for (p, ε) in enumerate(basis.pts) cubic_kalman_step_pullback!(∂, sys, zj, ε, view(∂F, :, p), ws) end @inbounds for j in 1:nz - ∂ΔC = view(∂𝒜, :, j) * m' .+ reshape(view(∂Λ, :, j), nz, N) - ∂F = ∂ΔC * W' + ℒ.mul!(∂F, reshape(view(∂Λ, :, j), nz, N), W') + ℒ.mul!(∂F, view(∂𝒜, :, j), Wm', one(eltype(∂F)), one(eltype(∂F))) fill!(zj, 0.0); zj[j] = 1.0 for (p, ε) in enumerate(basis.pts) cubic_kalman_step_pullback!(∂, sys, zj, ε, view(∂F, :, p), ws) end end - cubic_derived_pullback!(∂.S1, ∂.S2, ∂.M, ∂.mc, ∂.V, ∂.Wq, ∂.Wl_t, ∂.Bc, ∂.MM, - sys.M, sys.Pm, sys.nPast, sys.nExo, sys.na) + cubic_derived_pullback!(∂, sys) return ∂ end diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl index e72f932f7..0a45d6d9b 100644 --- a/test/test_cubic_kalman.jl +++ b/test/test_cubic_kalman.jl @@ -176,13 +176,14 @@ import AxisKeys: KeyedArray S2 = reshape(θ[n1+1:n1+n2], size(sys.S2)) S3 = reshape(θ[n1+n2+1:end], size(sys.S3)) M, mc, V, B2, Wq, Wl_t, Bc, MM = MacroModelling.cubic_derived_matrices(S1, S2, Pm, nP, nE, na) - merge(sys, (; S1, S2, S3, M, mc, V, B2, Wq, Wl_t, Bc, MM, Tv = eltype(θ))) + # the live-column slices are views of S2/S3 and must follow them + merge(sys, (; S1, S2, S3, M, mc, V, B2, Wq, Wl_t, Bc, MM, + S2k2 = S2[:, sys.k2cols], S2k12 = S2[:, sys.k12cols], + S3k3 = S3[:, sys.k3cols])) end θ0 = vcat(vec(sys.S1), vec(sys.S2), vec(sys.S3)) function fold(∂) - MacroModelling.cubic_derived_pullback!(∂.S1, ∂.S2, ∂.M, ∂.mc, ∂.V, - ∂.Wq, ∂.Wl_t, ∂.Bc, ∂.MM, - sys.M, Pm, nP, nE, na) + MacroModelling.cubic_derived_pullback!(∂, sys) return vcat(vec(∂.S1), vec(∂.S2), vec(∂.S3)) end From 3e4bd53a46d8ca7a38d9b6a6bf6180dfe0db2c62 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 2 Aug 2026 23:25:13 +0200 Subject: [PATCH 16/16] Add unpruned Ivashchenko Gaussian filters --- .github/workflows/ci.yml | 24 +- AGENT_PROGRESS.md | 38 ++ docs/src/filters.md | 105 +++- src/MacroModelling.jl | 12 +- src/default_options.jl | 17 +- src/filter/cubic_kalman.jl | 244 +++++++- src/filter/ivashchenko_kalman.jl | 945 +++++++++++++++++++++++++++++++ src/filter/quadratic_kalman.jl | 134 ++++- src/get_functions.jl | 59 +- src/rrules.jl | 112 +++- test/runtests.jl | 2 + test/test_cubic_kalman.jl | 27 + test/test_ivashchenko_kalman.jl | 173 ++++++ test/test_quadratic_kalman.jl | 24 +- 14 files changed, 1808 insertions(+), 108 deletions(-) create mode 100644 AGENT_PROGRESS.md create mode 100644 src/filter/ivashchenko_kalman.jl create mode 100644 test/test_ivashchenko_kalman.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5680a7a77..711666a85 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -181,6 +181,10 @@ jobs: os: ubuntu-latest arch: x64 test_set: "quadratic_kalman" + - version: '1' + os: ubuntu-latest + arch: x64 + test_set: "ivashchenko_kalman" steps: - uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v2 @@ -211,14 +215,18 @@ jobs: Project.toml rm -f Project.toml.bak - # - name: Restrict DynamicPPL to 0.40 for pigeons runs - # if: contains(matrix.test_set, 'pigeons') - # shell: bash - # run: | - # sed -i.bak \ - # -e '/^\[compat\]/,/^\[/ s/^DynamicPPL[[:space:]]*=.*$/DynamicPPL = "0.40"/g' \ - # Project.toml - # rm -f Project.toml.bak + - name: Restrict DynamicPPL to 0.40 for pigeons runs + if: contains(matrix.test_set, 'pigeons') + shell: bash + run: | + # Pigeons 0.4 and FlexiChains require disjoint DynamicPPL ranges. + # FlexiChains is removed below for these jobs; keep the compatible + # DynamicPPL line explicit so a future resolver update cannot select + # the incompatible 0.41/0.42 branch. + sed -i.bak \ + -e '/^\[compat\]/,/^\[/ s/^DynamicPPL[[:space:]]*=.*$/DynamicPPL = "0.40"/g' \ + Project.toml + rm -f Project.toml.bak - name: Remove Mooncake from pigeons runs if: contains(matrix.test_set, 'pigeons') diff --git a/AGENT_PROGRESS.md b/AGENT_PROGRESS.md new file mode 100644 index 000000000..3b4537d12 --- /dev/null +++ b/AGENT_PROGRESS.md @@ -0,0 +1,38 @@ +# Agent progress + +## Current task + +Implement Ivashchenko's unpruned Gaussian moment-closure filter as a separate filter for raw second- and third-order solutions, and resolve the CI dependency isolation issue. + +## Status + +- Repository progress file was absent at task start; this file records the current task. +- Quadratic and cubic Kollmann-style conditional covariance corrections remain implemented. +- A separate `:ivashchenko_kalman` filter now evaluates raw second-/third-order polynomial + solution maps and closes Gaussian moments through fourth/sixth order, respectively. +- The Ivashchenko filter has coupled theoretical mean/covariance initialization, optional + measurement error, partial and fully missing observation support, an RTS smoother, and + analytical reverse-mode rules for both raw second- and third-order solutions. +- CI now isolates the Pigeons/DynamicPPL resolver branch and has a dedicated Ivashchenko test row. + +## Verification + +- The scalar Gaussian moment reproduction passes. +- `test/test_ivashchenko_kalman.jl` passes 35/35, including Monte-Carlo moments, both orders, + initialization modes, partial/fully missing observations, RTS smoothing, standard deviations, + and reverse-vs-forward gradient checks. +- The isolated cubic tensor reverse check matches ForwardDiff to `8.9e-16`; the public third-order + likelihood reverse check differs from ForwardDiff by `3.9e-9` with theoretical initialization + and `2.8e-7` with the diagonal prior at the largest parameter gradient. +- `test/test_quadratic_kalman.jl` passes 33/33 and `test/test_cubic_kalman.jl` passes 30/30. +- CI YAML parsing and both non-Pigeons and Pigeons resolver probes pass; the direct root + `Pkg.test()` remains intentionally unsatisfiable because it includes incompatible optional + targets together, which the workflow-pruning steps resolve. +- `git diff --check` passes; module loading succeeds in the isolated test environment. + +## Implementation decision + +Ivashchenko's non-pruned Gaussian QKF is a separate algorithm rather than a switch on the +pruned augmented-state recursion. Its fourth-moment closure and unpruned state-product +dynamics require separate initialization and moment contractions. The cubic implementation +is an explicit extension of that idea; it is not attributed to Ivashchenko's second-order paper. diff --git a/docs/src/filters.md b/docs/src/filters.md index a5d5fd5a3..3929e2ccc 100644 --- a/docs/src/filters.md +++ b/docs/src/filters.md @@ -30,6 +30,7 @@ Two inputs cut across all of them and are covered separately below: `measurement |---|---|---|---|---|---|---| | `:kalman` | linear (`:first_order`) | exact | yes | optional (incl. correlated) | yes (Durbin–Koopman) | 1× | | `:inversion` | linear and nonlinear | exact given the shocks | yes | not available | n/a (filtered = smoothed) | ~1–10× | +| `:ivashchenko_kalman` | unpruned `:second_order`, `:third_order` | Gaussian moment closure | forward- and reverse-mode | optional (incl. correlated) | yes (RTS) | polynomial moment contractions | | `:bootstrap_particle` | linear and nonlinear | stochastic, unbiased | no | required (incl. correlated) | yes (genealogy) | ~10³× | | `:auxiliary_particle` | linear and nonlinear | stochastic, unbiased | no | required (incl. correlated) | yes (genealogy) | ~2× bootstrap | | `:tempered_particle` | linear and nonlinear | stochastic, unbiased | no | required (incl. correlated) | yes (genealogy) | ~5–10× bootstrap | @@ -38,6 +39,7 @@ A short decision rule: - **Linear model?** Use `:kalman`. It is exact, fast and differentiable, so gradient-based samplers (NUTS/HMC) work. There is no reason to use anything else. - **Nonlinear model, at least as many shocks as observables, no measurement error?** Use `:inversion` (the default at higher order). It is exact and differentiable. +- **Unpruned second- or third-order model with a Gaussian approximation to the filtering distribution?** Use `:ivashchenko_kalman`. It is separate from the pruned filters, supports measurement error, missing observations, RTS smoothing, and analytical reverse-mode differentiation. - **Nonlinear model with measurement error, or fewer shocks than observables?** Use a particle filter. Start with `:tempered_particle` if the observation is informative (small measurement error, many observables), otherwise `:bootstrap_particle`. - Particle-filter likelihoods are noisy and non-differentiable: pair them with gradient-free samplers such as slice sampling (Pigeons.jl) or nested sampling. @@ -50,7 +52,7 @@ Every knob discussed on this page has a default, and the defaults are not neutra | setting | default | consequence | |---|---|---| | `filter` | `:kalman` at `:first_order`, `:inversion` at every higher order | nonlinear models are filtered *exactly given the shocks*, with no measurement error | -| `measurement_error` | `:auto` | **none** for Kalman and inversion; ``(0.1 s_i)^2`` per observable for the particle filters | +| `measurement_error` | `:auto` | **none** for Kalman, inversion, and Ivashchenko; ``(0.1 s_i)^2`` per observable for the particle filters | | `initial_covariance` | `:theoretical` | the ergodic covariance — *not* the inversion filter's implicit ``BB'``, which is why Kalman and inversion likelihoods differ by default | | `smooth` | `true` for the Kalman filter, `false` otherwise | the particle filters **do** support smoothing but do not use it unless asked | | `presample_periods` | `0` | the initial-condition transient is included in the likelihood | @@ -163,7 +165,7 @@ Because the estimate is random, a repeated evaluation at the same parameters giv Without measurement error the observation equation is a deterministic function of the state. A particle would have to reproduce ``y_t`` *exactly* to get non-zero weight, which happens with probability zero — every weight collapses to zero and the filter dies. Measurement error smears the observation density and gives particles something to score against. -`measurement_error = :auto` (the default) therefore resolves, for the particle filters, to a variance of ``(0.1 s_i)^2`` per observable, ``s_i`` being that observable's sample standard deviation — and to *no* measurement error for the Kalman and inversion filters. For serious work set it explicitly or estimate it: the level of the likelihood depends on it, so likelihoods computed under different measurement errors are not comparable. See [Measurement error and the initial covariance](@ref) for what ``H`` is and the other jobs it does. +`measurement_error = :auto` (the default) therefore resolves, for the particle filters, to a variance of ``(0.1 s_i)^2`` per observable, ``s_i`` being that observable's sample standard deviation — and to *no* measurement error for the Kalman, inversion, and Ivashchenko filters. For serious work set it explicitly or estimate it: the level of the likelihood depends on it, so likelihoods computed under different measurement errors are not comparable. See [Measurement error and the initial covariance](@ref) for what ``H`` is and the other jobs it does. ### Bootstrap (`:bootstrap_particle`) @@ -421,7 +423,7 @@ P = P̂ − K C P̂ P = P̂ − K 𝒞 P̂ | dimension (SW07) | 34 | 446 | | transition | ``x' = Ax + B\varepsilon`` | ``z' = \mathcal{A}z + c + w(z,\varepsilon)`` | | drift ``c`` | zero — certainty equivalence empties ``\mathbf{S}_1``'s constant column | non-zero — carries the risk correction | -| noise covariance | ``\mathbf{B} = BB'``, **constant** | ``G(z)G(z)' + Q_H``, **depends on the state** | +| noise covariance | ``\mathbf{B} = BB'``, **constant** | ``G(\bar z)G(\bar z)' + Q_H + Q_{\mathrm{state}}``, **depends on state mean and covariance** | | innovation | ``B\varepsilon`` — Gaussian | ``G\varepsilon + H(\varepsilon\otimes\varepsilon - \mathrm{vec}\,I)`` — **not** Gaussian | | observation | ``y = Cx``, general ``C`` | ``y = (x_1+x_2)[\text{obs}]`` — a selection of two blocks | | solve per period | LU of ``F`` (``n_{obs}^3``) | Cholesky of ``F`` (``n_{obs}^3``) | @@ -437,6 +439,22 @@ That is precisely the conditional heteroskedasticity a second-order solution add model's shock impact depends on where the state is — and it is why the filter is not merely a linear filter on a bigger vector. +The covariance is integrated over the filtered state distribution, not only evaluated at its +mean. Writing ``G(z) = G(\bar z) + \sum_i z_i G_i`` and letting ``P_a`` be the covariance of +the past first-order state used by the transition, + +```math +Q_{\mathrm{state}} = \sum_{i,j} (P_a)_{ij}G_iG_j',\qquad +P_a = P_z P_{t-1|t-1}P_z'. +``` + +The timing is posterior-then-predict: in the loop this is the covariance before the current +observation update (`Pc`), because it describes uncertainty in the state that generates the +next shock loading. The full ``P_{t-1|t-1}`` remains necessary for the Kalman prediction and +update, but the added term reads only ``P_a``. Thus the correction adds small ``n_{past}`` +and ``n_z\times n_{past}`` workspaces; it does not require another full ``n_z\times n_z`` +covariance. + **The innovation is no longer Gaussian.** ``\varepsilon\otimes\varepsilon`` is a ``\chi^2``-type object; matching only its first two moments discards every higher cumulant. The linear filter has nothing to discard, which is why it is exact and this one is not; the @@ -460,11 +478,12 @@ parts are uncorrelated and, using ``E[(\varepsilon\otimes\varepsilon)(\varepsilon\otimes\varepsilon)'] = \mathrm{vec}(I)\mathrm{vec}(I)' + I + K``, ```math -\mathrm{Var}(w) = GG' + H(I+K)H', +\mathrm{Var}(w) = G(\bar z)G(\bar z)' + Q_{\mathrm{state}} + H(I+K)H', ``` -with ``K`` the commutation matrix. ``H`` is constant; ``G`` depends on the state and is -evaluated at the filtered mean. +with ``K`` the commutation matrix. ``H`` is constant; the first ``G`` term is evaluated at +the filtered mean and ``Q_{\mathrm{state}}`` integrates its affine state dependence over +the filtered covariance. What is approximated is the conditional *distribution*. ``\varepsilon\otimes\varepsilon`` is a squared Gaussian — skewed, not Gaussian — so the recursion delivers the best **linear** @@ -552,6 +571,44 @@ augmented state altogether. in the transition and the observation is a plain selection with zero loading on the Kronecker block. The machinery is shared, the regime is not. +!!! note "Ivashchenko's QKF is a different filter" + Ivashchenko (2014) applies a Gaussian moment closure directly to the **unpruned** + second-order solution. Its conditional covariance therefore requires fourth moments of the + state and innovation errors. That is not a drop-in replacement for this pruned augmented + recursion: without pruning the quadratic transition generates quartic, then higher-order, + state products and the finite linear state representation no longer closes. The package + implements it separately as `:ivashchenko_kalman`. + +## Ivashchenko's unpruned Gaussian filter + +`filter = :ivashchenko_kalman` is available for `algorithm = :second_order` and +`:third_order`. It treats the raw perturbation solution as a polynomial map in the previous +period's state and the current shocks: + +```math +f(u) = S_1 u + \tfrac12 S_2(u\otimes u) + \tfrac16 S_3(u\otimes u\otimes u), +\qquad u = [x_{t-1};\ 1;\varepsilon_t]. +``` + +The filter expands this map around the current Gaussian mean. At second order, the mean and +covariance use Gaussian moments through order four. The third-order implementation is the +corresponding extension through order six: the cubic Hermite component contributes both to the +effective linear loading and to the covariance. The third-order extension is an implementation +of the same moment-closure idea, not a claim that the 2014 paper itself derives a cubic QKF. + +The `:theoretical` initial covariance solves the coupled unpruned mean/covariance fixed point, +starting from the linear Lyapunov covariance. A supplied covariance or `:diagonal` uses that +prior directly. The filter supports Gaussian measurement error and partial or fully missing +observation periods: the update is restricted to the observed rows, and a fully missing period +is prediction-only. `smooth = true` applies a fixed-interval Rauch–Tung–Striebel smoother to +the Gaussian state moments. Its reverse-mode rule differentiates the moment contractions, +measurement updates, and theoretical fixed-point initialization analytically; it does not use +automatic differentiation internally. + +This is computationally different from `:quadratic_kalman` and `:cubic_kalman`: it avoids the +large pruned augmented covariance, but the cubic moment contraction scales with the cube of the +state-and-shock dimension and is intended for relatively small models. + ### Cost The augmented dimension is ``2n_r + n_{past}(n_{past}+1)/2``, where ``n_r`` counts the retained @@ -565,7 +622,7 @@ dominates everything else — per period, measured: | ``(\mathcal{A}P_c)\mathcal{A}'`` | ``n_z^3`` (88.7M flops) | 1.26 | 49% | | symmetrisation ×2 | ``n_z^2`` | 0.10 | 4% | | ``P_p - K\,CP`` | ``n_z^2 n_{obs}`` | 0.03 | 1% | -| ``GG'`` | ``n_z^2 n_\varepsilon`` | 0.03 | 1% | +| ``G G'`` plus covariance correction | ``n_z^2 n_\varepsilon + n_\varepsilon(n_z n_{past}^2+n_z^2 n_{past})`` | small relative to the two ``n_z^3`` products | — | | ``\mathcal{C}P_p`` | ``n_{obs}n_z^2`` | 0.06 | 2% | | build ``G`` | ``n_z n_\varepsilon n_{past}`` | 0.02 | 1% | @@ -576,12 +633,15 @@ of propagating a covariance over the Kronecker-augmented state, and no amount of it. Sparsity does not help either — ``\mathcal{A}`` is about 50% dense, and a sparse representation measures 10× *slower* than the dense one. -**References:** Kollmann (2015), *Computational Economics* 45, 239–260 — the filter implemented -here. Andreasen, Fernández-Villaverde & Rubio-Ramírez (2018) — the pruned state-space -representation. Monfort, Renne & Roussellet (2015), *Journal of Econometrics* 187, 43–56 — the -quadratic Kalman filter for quadratic measurement equations. Andreasen (2013), *Journal of -Applied Econometrics* 28, 929–955 — the central difference Kalman filter, the unpruned -alternative. +**References:** Kollmann (2015), [*Tractable Latent State Filtering for Non-Linear DSGE Models +Using a Second-Order Approximation and Pruning](https://doi.org/10.1007/s10614-013-9418-3), +*Computational Economics* 45, 239–260 — the filter implemented here. Ivashchenko (2014), +[*DSGE Model Estimation on the Basis of Second-Order Approximation](https://doi.org/10.1007/s10614-013-9363-1), +*Computational Economics* 43, 71–82 — the non-pruned Gaussian QKF. Andreasen, +Fernández-Villaverde & Rubio-Ramírez (2018) — the pruned state-space representation. Monfort, +Renne & Roussellet (2015), *Journal of Econometrics* 187, 43–56 — the quadratic Kalman filter +for quadratic measurement equations. Andreasen (2013), *Journal of Applied Economics* 28, +929–955 — the central difference Kalman filter, the unpruned alternative. ## The cubic Kalman filter @@ -654,7 +714,7 @@ closed form: ```math \mathbb{E}[f] = C(z)\,m,\qquad -\mathrm{Var}(f) = C(z)\,\Psi\,C(z)', +\mathrm{Var}(f) = C(\bar z)\,\Psi\,C(\bar z)' + Q_{\mathrm{state}}, ``` with ``m_\alpha = \mathbb{E}[\varepsilon^\alpha]`` and @@ -662,7 +722,20 @@ with ``m_\alpha = \mathbb{E}[\varepsilon^\alpha]`` and Because the shocks are *independent* standard normals, ``\mathbb{E}[\varepsilon^\alpha]`` factorises into double factorials, so ``\Psi`` is a closed form rather than a sum over Isserlis pairings. This is the exact analogue of the quadratic filter's affine ``G(z)`` with -``Q = GG'``: a period costs one matvec and two gemms, not a quadrature sweep. +``Q = C\Psi C'``: a period costs one matvec and two gemms, not a quadrature sweep. + +As in the quadratic filter, ``Q_{\mathrm{state}}`` integrates the affine loading over the +filtered state distribution. If ``C(z)=C(\bar z)+\sum_i z_iD_i``, then + +```math +Q_{\mathrm{state}} = \sum_{i,j}P_{ij}D_i\Psi D_j'. +``` + +Only the structurally supported blocks of ``z`` can appear in ``C(z)``: past ``x_1``, past +``x_2`` and ``q_{11}``. The full augmented covariance is still propagated for the Kalman +update, but the correction contracts only the corresponding submatrix. Since ``q_{11}`` is +in that support, the stationary initialization solves a coupled covariance fixed point; its +adjoint uses the corresponding implicit fixed-point equation. ``C(z)`` is recovered by interpolation on ``\binom{n_\varepsilon+3}{3}`` points, which is also where a tensor Gauss-Hermite rule is left behind — its node count grows as @@ -697,7 +770,7 @@ than merely moving the end-to-end number: |---|---| | step adjoint | ``\partial f(z,\varepsilon)`` onto ``\mathbf{S}_1,\mathbf{S}_2,\mathbf{S}_3`` and the derived blocks | | build adjoint | ``\partial(\mathcal{A}, c, c_0, \Lambda)`` replayed over the same ``(n_z+1)N`` points the forward pass visited | -| recursion adjoint | the Kalman loop, with ``Q = C\Psi C'`` in place of the quadratic filter's ``GG' + Q_H`` | +| recursion adjoint | the Kalman loop, with ``Q = C(\bar z)\Psi C(\bar z)' + Q_{\mathrm{state}}`` | Everything the step builds from ``z`` and ``\varepsilon`` alone — ``\mathrm{aug}``, ``K_2``, ``K_{12}``, ``K_3``, the ``Q`` blocks — is constant for the adjoint, so only the paths diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index a2f726551..cf122f1f3 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -194,6 +194,7 @@ include("./filter/kalman.jl") include("./filter/particle.jl") include("./filter/quadratic_kalman.jl") include("./filter/cubic_kalman.jl") +include("./filter/ivashchenko_kalman.jl") export @model, @parameters, solve! @@ -402,7 +403,7 @@ function normalize_filtering_options(filter::Symbol, # `:particle` is a convenience alias for the bootstrap particle filter. filter = get(PARTICLE_FILTER_ALIASES, filter, filter) - @assert filter ∈ SUPPORTED_FILTERS "Unsupported `filter = :$(filter)`. Choose the Kalman filter (`:kalman`, linear models), the inversion filter (`:inversion`, linear and nonlinear models), or one of the particle filters (`:bootstrap_particle`, `:auxiliary_particle`, `:tempered_particle`; linear and nonlinear models). `:particle` is accepted as an alias for `:bootstrap_particle`." + @assert filter ∈ SUPPORTED_FILTERS "Unsupported `filter = :$(filter)`. Choose the Kalman filter (`:kalman`, linear models), the inversion filter (`:inversion`, linear and nonlinear models), the unpruned Ivashchenko filter (`:ivashchenko_kalman`, second- and third-order models), or one of the particle filters (`:bootstrap_particle`, `:auxiliary_particle`, `:tempered_particle`; linear and nonlinear models). `:particle` is accepted as an alias for `:bootstrap_particle`." is_particle = filter ∈ PARTICLE_FILTERS @@ -427,9 +428,14 @@ function normalize_filtering_options(filter::Symbol, filter = :inversion end + if filter == :ivashchenko_kalman && algorithm ∉ (:second_order, :third_order) + @info "The Ivashchenko filter is only defined for `algorithm = :second_order` or `:third_order`; got `:$(algorithm)`. Setting `filter = :inversion`." maxlog = maxlog + filter = :inversion + end + # Higher-order solutions are handled by the inversion filter by default, but # the particle filters are explicitly valid at every order too. - if algorithm != :first_order && filter != :inversion && filter != :quadratic_kalman && filter != :cubic_kalman && !is_particle + if algorithm != :first_order && filter != :inversion && filter != :quadratic_kalman && filter != :cubic_kalman && filter != :ivashchenko_kalman && !is_particle @info "Higher order solution algorithms only support the inversion and particle filters. Setting `filter = :inversion`." maxlog = maxlog filter = :inversion is_particle = false @@ -456,7 +462,7 @@ function normalize_filtering_options(filter::Symbol, # principle redistribute across time; doing so would be a different estimator, # not the inversion filter's smoother. if filter in (:quadratic_kalman, :cubic_kalman) && smooth - @info "The quadratic Kalman filter does not provide smoothed estimates. Setting `smooth = false`." maxlog = maxlog + @info "The $(filter) filter does not provide smoothed estimates. Setting `smooth = false`." maxlog = maxlog smooth = false end diff --git a/src/default_options.jl b/src/default_options.jl index 73b478cff..d02559cec 100644 --- a/src/default_options.jl +++ b/src/default_options.jl @@ -13,10 +13,11 @@ const DEFAULT_PRESAMPLE_PERIODS = 0 # Each particle-filter variant is its own `filter` value, so the filter is fully # identified by a single symbol (no separate "which particle filter" argument). const PARTICLE_FILTERS = (:bootstrap_particle, :auxiliary_particle, :tempered_particle) -# The quadratic and cubic Kalman filters apply only to the pruned second- and -# third-order solutions, whose augmented state spaces are linear (see -# src/filter/quadratic_kalman.jl and src/filter/cubic_kalman.jl). -const SUPPORTED_FILTERS = (:kalman, :inversion, :quadratic_kalman, :cubic_kalman, PARTICLE_FILTERS...) +# The quadratic and cubic Kalman filters apply to the pruned second- and +# third-order solutions. Ivashchenko's filter is the separate unpruned Gaussian +# moment-closure filter for the raw second- and third-order solutions. +const SUPPORTED_FILTERS = (:kalman, :inversion, :quadratic_kalman, :cubic_kalman, + :ivashchenko_kalman, PARTICLE_FILTERS...) # `:particle` is accepted as a convenience alias for the bootstrap filter. const PARTICLE_FILTER_ALIASES = Dict(:particle => :bootstrap_particle) # Maps a filter symbol onto the internal variant tag used for dispatch. @@ -28,9 +29,9 @@ const PARTICLE_FILTER_VARIANT = Dict(:bootstrap_particle => :bootstrap, # `measurement_error` is the covariance H of ηₜ ~ N(0, H) in yₜ = C xₜ + ηₜ. It is # *not* a standard deviation: a scalar is the common variance of every observable, # a vector the per-observable variances, and a matrix the full covariance. -# `:auto` resolves per filter: no measurement error for the Kalman and inversion -# filters (their historical behaviour), and a small data-driven value for the -# particle filters, which are degenerate without it. +# `:auto` resolves per filter: no measurement error for the Kalman, inversion, +# and Ivashchenko filters (their historical/deterministic-filter behaviour), and +# a small data-driven value for the particle filters, which are degenerate without it. const DEFAULT_MEASUREMENT_ERROR = :auto # Auto measurement-error *standard deviation* as a fraction of each observable's # sample standard deviation (squared into a variance before it reaches a filter). @@ -170,4 +171,4 @@ const DEFAULT_MAXLOG = 3 # Caching and workspace defaults const DEFAULT_CACHING = true -const DEFAULT_USE_WORKSPACES = true \ No newline at end of file +const DEFAULT_USE_WORKSPACES = true diff --git a/src/filter/cubic_kalman.jl b/src/filter/cubic_kalman.jl index ef9efe707..18184349e 100644 --- a/src/filter/cubic_kalman.jl +++ b/src/filter/cubic_kalman.jl @@ -376,10 +376,19 @@ function build_cubic_kalman_system_from_constants(cons, 𝐒₁, 𝐒₂, 𝐒 mod(fld(r - 1, nPast), nPast) + 1, mod(r - 1, nPast) + 1) for r in can3] + # Only these augmented-state blocks can enter the affine noise loading: + # past x₁ rows, past x₂ rows, and q₁₁. The larger covariance is still needed + # by the Kalman recursion, but the additional Kollmann correction only reads + # this structural support. + past_positions = [findfirst(!iszero, view(Pm, i, :)) for i in 1:nPast] + noise_state_indices = vcat(first(r1) .+ past_positions .- 1, + first(r2) .+ past_positions .- 1, + collect(i11)) + return (; nr, nPast, nExo, na, nz, oas, S1, S2, S3, Pm, C, op1, op2, op3, r1, r2, r3, i11, i12, i111, nq11, nq12, nq111, exp2, can2, exp3, can3, can2_ij, can3_ijk, - M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail, + M, mc, V, B2, Wq, Wl_t, Bc, MM, ntail, noise_state_indices, k2cols, k2_ij, S2k2, k12cols, k12_ij, S2k12, k3cols, k3_ijk, S3k3) end @@ -752,6 +761,148 @@ function build_cubic_kalman_system(sys, basis; ws = cubic_kalman_workspace(sys)) return 𝒜, c, c₀, Λ end +# Fill the conditional covariance of a cubic innovation. With +# C(z) = C̄ + Σᵢ zᵢDᵢ and shock-moment covariance Ψ, +# +# E[C(Z)ΨC(Z)'] = C̄ΨC̄' + Σᵢⱼ Cov(Zᵢ,Zⱼ) DᵢΨDⱼ'. +# +# `Λnoise` contains only the columns Dᵢ that can be nonzero structurally. This +# avoids forming or multiplying by the full augmented covariance for the extra +# term while retaining the full matrix Q required by the Kalman update. +function cubic_kalman_noise_covariance!(Q, C, Λnoise, Ψ, Pc, noise_state_indices, + Pnoise, mixvec, mixΨ, CΨ) + nI = length(noise_state_indices) + nz, N = size(C) + ℒ.mul!(CΨ, C, Ψ) + ℒ.mul!(Q, CΨ, C') + @inbounds for j in 1:nI, i in 1:nI + Pnoise[i, j] = Pc[noise_state_indices[i], noise_state_indices[j]] + end + @inbounds for i in 1:nI + ℒ.mul!(mixvec, Λnoise, view(Pnoise, :, i)) + mix = reshape(mixvec, nz, N) + ℒ.mul!(mixΨ, mix, Ψ) + Di = reshape(view(Λnoise, :, i), nz, N) + ℒ.mul!(Q, mixΨ, Di', one(eltype(Q)), one(eltype(Q))) + end + @inbounds for j in 1:nz, i in 1:j + m = (Q[i, j] + Q[j, i]) / 2 + Q[i, j] = m; Q[j, i] = m + end + return Q +end + +# Solve the cubic stationary covariance fixed point +# +# Σ = 𝒜Σ𝒜' + C̄ΨC̄' + K(Σ), +# +# where K is the state-covariance correction above. Unlike the quadratic case, +# q₁₁ is itself in the loading support, so the correction is coupled to the +# augmented covariance. The Float64 path stops on convergence; the dual path +# runs a fixed number of iterations so ForwardDiff sees a smooth computation. +function cubic_kalman_initial_covariance(𝒜, Cbar, Λnoise, Ψ, noise_state_indices; + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling, + max_iterations::Int = 100, + tolerance::Real = 1e-12) + nz, N = size(Cbar) + Tv = promote_type(eltype(𝒜), eltype(Cbar), eltype(Λnoise), eltype(Ψ)) + Qbase = Matrix{Tv}(undef, nz, nz) + CΨ = Matrix{Tv}(undef, nz, N) + ℒ.mul!(CΨ, Cbar, Ψ) + ℒ.mul!(Qbase, CΨ, Cbar') + Qbase = (Qbase + Qbase') / 2 + + Pnoise = Matrix{Tv}(undef, length(noise_state_indices), length(noise_state_indices)) + mixvec = Vector{Tv}(undef, nz * N) + mixΨ = Matrix{Tv}(undef, nz, N) + Q = Matrix{Tv}(undef, nz, nz) + Σ = qkf_lyapunov(𝒜, Qbase; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + float_path = eltype(𝒜) <: AbstractFloat + converged = false + for iteration in 1:max_iterations + copyto!(Q, Qbase) + @inbounds for j in 1:length(noise_state_indices), i in 1:length(noise_state_indices) + Pnoise[i, j] = Σ[noise_state_indices[i], noise_state_indices[j]] + end + @inbounds for i in 1:length(noise_state_indices) + ℒ.mul!(mixvec, Λnoise, view(Pnoise, :, i)) + mix = reshape(mixvec, nz, N) + ℒ.mul!(mixΨ, mix, Ψ) + Di = reshape(view(Λnoise, :, i), nz, N) + ℒ.mul!(Q, mixΨ, Di', one(Tv), one(Tv)) + end + @inbounds for j in 1:nz, i in 1:j + m = (Q[i, j] + Q[j, i]) / 2 + Q[i, j] = m; Q[j, i] = m + end + Σnew = qkf_lyapunov(𝒜, Q; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + if float_path + difference = maximum(abs, Σnew - Σ) + scale = max(1.0, maximum(abs, Σnew)) + if difference <= tolerance * scale + Σ = Σnew + converged = true + break + end + end + Σ = Σnew + end + float_path && !converged && error("The cubic Kalman stationary covariance fixed point did not converge " * + "within $max_iterations iterations.") + return Σ +end + +# Adjoint of the coupled stationary covariance equation. This is the transpose +# fixed point X = 𝒜'X𝒜 + K*(X) + Σ̄, solved with a Lyapunov-preconditioned +# iteration because K* is inexpensive on the restricted loading support. It +# supplies the exact implicit pullback of the converged Float64 covariance +# iteration without differentiating through it. +function cubic_kalman_stationary_adjoint(𝒜, Λnoise, Ψ, noise_state_indices, Σ̄; + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling, + max_iterations::Int = 100, + tolerance::Real = 1e-12) + nz = size(Σ̄, 1) + coefficient_count = size(Λnoise, 1) ÷ nz + nI = length(noise_state_indices) + X = (Σ̄ + Σ̄') / 2 + P̄noise = zeros(eltype(X), nI, nI) + E = zeros(eltype(X), nz, coefficient_count) + EΨ = zeros(eltype(X), nz, coefficient_count) + Xnew = similar(X) + rhs = copy(Σ̄) + X = qkf_lyapunov(Matrix(𝒜'), rhs; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + for iteration in 1:max_iterations + fill!(P̄noise, zero(eltype(P̄noise))) + @inbounds for i in 1:nI + Di = reshape(view(Λnoise, :, i), nz, coefficient_count) + ℒ.mul!(E, X, Di) + ℒ.mul!(EΨ, E, Ψ) + for j in 1:nI + Dj = reshape(view(Λnoise, :, j), nz, coefficient_count) + P̄noise[i, j] = sum(EΨ .* Dj) + end + end + copyto!(rhs, Σ̄) + @inbounds for j in 1:nI, i in 1:nI + rhs[noise_state_indices[i], noise_state_indices[j]] += P̄noise[i, j] + end + Xnew = qkf_lyapunov(Matrix(𝒜'), rhs; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + Xnew = (Xnew + Xnew') / 2 + difference = maximum(abs, Xnew - X) + scale = max(1.0, maximum(abs, Xnew)) + X = copy(Xnew) + difference <= tolerance * scale && return X + end + error("The cubic Kalman stationary covariance adjoint did not converge " * + "within $max_iterations iterations.") +end + """ Adjoint of `build_cubic_kalman_system`. The build is linear in the collected step evaluations — `C(z) = F(z) W`, then `c = C(0)m`, `c₀ = vec C(0)`, @@ -831,27 +982,29 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; basis = cubic_noise_basis(sys.nExo) 𝒜, c, c₀, Λ = build_cubic_kalman_system(sys, basis; ws = ws) N, Ψ = basis.N, basis.Ψ + noise_state_indices = sys.noise_state_indices + Λnoise = Matrix(Λ[:, noise_state_indices]) cvec = Vector{Tv}(undef, nz * N) CΨ = Matrix{Tv}(undef, nz, N) Q = Matrix{Tv}(undef, nz, nz) + Pnoise = Matrix{Tv}(undef, length(noise_state_indices), length(noise_state_indices)) + mixvec = Vector{Tv}(undef, nz * N) + mixΨ = Matrix{Tv}(undef, nz, N) # Q(z) = C(z) Ψ C(z)' with vec(C) = c₀ + Λz. - noise_covariance! = function (Q, z) + noise_covariance! = function (Q, z, Pc) copyto!(cvec, c₀) ℒ.mul!(cvec, Λ, z, one(Tv), one(Tv)) C = reshape(cvec, nz, N) - ℒ.mul!(CΨ, C, Ψ) - ℒ.mul!(Q, CΨ, C') - for j in 1:nz, i in 1:j - m = (Q[i, j] + Q[j, i]) / 2 - Q[i, j] = m; Q[j, i] = m - end - return Q + cubic_kalman_noise_covariance!(Q, C, Λnoise, Ψ, Pc, noise_state_indices, + Pnoise, mixvec, mixΨ, CΨ) end z = (Matrix{Tv}(ℒ.I(nz)) - 𝒜) \ c - Σ = qkf_lyapunov(𝒜, noise_covariance!(Q, z); workspaces = workspaces, - lyapunov_algorithm = lyapunov_algorithm) + Cbar = reshape(c₀ + Λ * z, nz, N) + Σ = cubic_kalman_initial_covariance(𝒜, Cbar, Λnoise, Ψ, noise_state_indices; + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) # Preallocate the recursion's working matrices once, as the quadratic filter # does: the covariance propagation is the whole cost, and allocating an nz×nz @@ -870,7 +1023,7 @@ function run_cubic_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}; ll = zero(Tv) log2pi = log(2π) @inbounds for t in 1:nT - noise_covariance!(Q, z) + noise_covariance!(Q, z, Pc) # Pp = 𝒜 Pc 𝒜' + Q ℒ.mul!(Tm, 𝒜, Pc) @@ -929,25 +1082,33 @@ end """ Taped forward pass plus adjoint for the cubic Kalman recursion. Mirrors `quadratic_kalman_recursion`'s verified adjoint; the one structural difference is -the noise term, `Q = CΨC'` with `vec(C) = c₀ + Λz`, in place of `GG' + Q_H`. +the noise term, `Q = C(\bar z)ΨC(\bar z)' + Q_state` with `vec(C) = c₀ + Λz`, in place of +`G(\bar z)G(\bar z)' + Q_H + Q_state`. `Q̄` is symmetric here because `P̄p` is symmetrised before use, so the cotangent of `C` is `2 Q̄ C Ψ` rather than `(Q̄ + Q̄')CΨ`. """ function cubic_kalman_recursion_taped(𝒜, c, c₀, Λ, Ψ, Hm, Y, 𝒞, z0, Σ0, nz, N, - presample_periods, on_failure_loglikelihood) + presample_periods, on_failure_loglikelihood, + noise_state_indices) n_obs, nT = size(Y) z = copy(z0); Pc = copy(Σ0) zs = Vector{Vector{Float64}}(); Ps = Vector{Matrix{Float64}}() - Cs = Vector{Matrix{Float64}}(); vs = Vector{Vector{Float64}}() + Cs = Vector{Matrix{Float64}}(); Pas = Vector{Matrix{Float64}}() + vs = Vector{Vector{Float64}}() CPs = Vector{Matrix{Float64}}(); Fis = Vector{Matrix{Float64}}() Ks = Vector{Matrix{Float64}}() ll = 0.0; log2pi = log(2π) + Λnoise = Matrix(Λ[:, noise_state_indices]) + Pnoise = zeros(length(noise_state_indices), length(noise_state_indices)) + mixvec = zeros(nz * N); mixΨ = zeros(nz, N); CΨ = zeros(nz, N); Q = zeros(nz, nz) for t in 1:nT push!(zs, copy(z)); push!(Ps, copy(Pc)) C = reshape(c₀ + Λ * z, nz, N) - Q = C * Ψ * C' + cubic_kalman_noise_covariance!(Q, C, Λnoise, Ψ, Pc, noise_state_indices, + Pnoise, mixvec, mixΨ, CΨ) + push!(Pas, copy(Pnoise)) zp = 𝒜 * z + c Pp = 𝒜 * Pc * 𝒜' + Q; Pp = (Pp + Pp') / 2 v = Y[:, t] - 𝒞 * zp @@ -964,12 +1125,12 @@ function cubic_kalman_recursion_taped(𝒜, c, c₀, Λ, Ψ, Hm, Y, 𝒞, z0, Σ Pc = Pp - K * CP; Pc = (Pc + Pc') / 2 push!(Cs, C); push!(vs, v); push!(CPs, CP); push!(Fis, Fi); push!(Ks, K) end - return ll, (; zs, Ps, Cs, vs, CPs, Fis, Ks) + return ll, (; zs, Ps, Cs, Pas, vs, CPs, Fis, Ks) end function cubic_kalman_recursion_pullback(tape, 𝒜, c, c₀, Λ, Ψ, 𝒞, nz, N, n_obs, nT, - presample_periods, ∂ll) - (; zs, Ps, Cs, vs, CPs, Fis, Ks) = tape + presample_periods, ∂ll, noise_state_indices) + (; zs, Ps, Cs, Pas, vs, CPs, Fis, Ks) = tape 𝒜̄ = zeros(nz, nz); c̄ = zeros(nz); c̄₀ = zeros(length(c₀)); Λ̄ = zeros(size(Λ)) H̄m = zeros(n_obs, n_obs); Ȳ = zeros(n_obs, nT) z̄ = zeros(nz); P̄ = zeros(nz, nz) @@ -1002,6 +1163,22 @@ function cubic_kalman_recursion_pullback(tape, 𝒜, c, c₀, Λ, Ψ, 𝒞, nz, vC̄ = vec(C̄) c̄₀ .+= vC̄ Λ̄ .+= vC̄ * z_' + P̄noise = zeros(length(noise_state_indices), length(noise_state_indices)) + Λnoise = view(Λ, :, noise_state_indices) + @inbounds for i in 1:length(noise_state_indices) + Di = reshape(view(Λnoise, :, i), nz, N) + E = Q̄ * Di * Ψ + for j in 1:length(noise_state_indices) + Dj = reshape(view(Λnoise, :, j), nz, N) + P̄noise[i, j] = sum(E .* Dj) + end + Dmix = reshape(Λnoise * view(Pas[t], :, i), nz, N) + D̄ = 2 .* (Q̄ * Dmix * Ψ) + Λ̄[:, noise_state_indices[i]] .+= vec(D̄) + end + @inbounds for j in 1:length(noise_state_indices), i in 1:length(noise_state_indices) + P̄[noise_state_indices[i], noise_state_indices[j]] += P̄noise[i, j] + end 𝒜̄ .+= z̄p * z_' c̄ .+= z̄p z̄ = 𝒜' * z̄p + Λ' * vC̄ @@ -1066,14 +1243,18 @@ function rrule(::typeof(calculate_loglikelihood), basis = cubic_noise_basis(sys.nExo) 𝒜, c, c₀, Λ = build_cubic_kalman_system(sys, basis; ws = ws) N, Ψ, 𝒞 = basis.N, basis.Ψ, sys.C + noise_state_indices = sys.noise_state_indices + Λnoise = Matrix(Λ[:, noise_state_indices]) z₀ = (Matrix{Float64}(ℒ.I(nz)) - 𝒜) \ c C₀ = reshape(c₀ + Λ * z₀, nz, N) - Q₀ = C₀ * Ψ * C₀'; Q₀ = (Q₀ + Q₀') / 2 - Σ₀ = qkf_lyapunov(𝒜, Q₀; workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + Σ₀ = cubic_kalman_initial_covariance(𝒜, C₀, Λnoise, Ψ, noise_state_indices; + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) llh, tape = cubic_kalman_recursion_taped(𝒜, c, c₀, Λ, Ψ, Hm, Matrix(data_in_deviations), - 𝒞, z₀, Σ₀, nz, N, presample, on_failure_loglikelihood) + 𝒞, z₀, Σ₀, nz, N, presample, on_failure_loglikelihood, + noise_state_indices) nine(x...) = (NoTangent(), NoTangent(), NoTangent(), NoTangent(), x[1], x[2], NoTangent(), x[3], NoTangent()) @@ -1086,11 +1267,14 @@ function rrule(::typeof(calculate_loglikelihood), ∂llh = unthunk(∂llh_bar) 𝒜̄, c̄, c̄₀, Λ̄, _, Ȳ, z̄₀, Σ̄₀ = cubic_kalman_recursion_pullback(tape, 𝒜, c, c₀, Λ, Ψ, 𝒞, nz, N, n_obs, nT, - presample, ∂llh) - - # Σ₀ = 𝒜Σ₀𝒜' + Q₀ ⇒ X solves X = 𝒜'X𝒜 + Σ̄₀ - X = qkf_lyapunov(Matrix(𝒜'), (Σ̄₀ + Σ̄₀') / 2; workspaces = workspaces, - lyapunov_algorithm = lyapunov_algorithm) + presample, ∂llh, noise_state_indices) + + # Σ₀ = 𝒜Σ₀𝒜' + C̄ΨC̄' + K(Σ₀). The adjoint is the corresponding + # coupled fixed point, so this remains the exact pullback of the + # state-dependent stationary covariance. + X = cubic_kalman_stationary_adjoint(𝒜, Λnoise, Ψ, noise_state_indices, + Σ̄₀; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) 𝒜̄ .+= 2 .* (X * 𝒜 * Σ₀) Q̄₀ = (X + X') / 2 C̄₀ = 2 .* (Q̄₀ * C₀ * Ψ) @@ -1098,6 +1282,12 @@ function rrule(::typeof(calculate_loglikelihood), c̄₀ = c̄₀ .+ vC̄₀ Λ̄ = Λ̄ .+ vC̄₀ * z₀' z̄₀ = z̄₀ .+ Λ' * vC̄₀ + P₀ = Σ₀[noise_state_indices, noise_state_indices] + @inbounds for i in 1:length(noise_state_indices) + Dmix = reshape(Λnoise * view(P₀, :, i), nz, N) + D̄ = 2 .* (Q̄₀ * Dmix * Ψ) + Λ̄[:, noise_state_indices[i]] .+= vec(D̄) + end # z₀ = (I − 𝒜)⁻¹ c y = (Matrix{Float64}(ℒ.I(nz)) - 𝒜)' \ z̄₀ diff --git a/src/filter/ivashchenko_kalman.jl b/src/filter/ivashchenko_kalman.jl new file mode 100644 index 000000000..2befb1746 --- /dev/null +++ b/src/filter/ivashchenko_kalman.jl @@ -0,0 +1,945 @@ +@stable default_mode = "disable" begin + +# Unpruned Gaussian moment-closure filter in the spirit of Ivashchenko (2014). +# +# The second- and third-order perturbation solutions are treated as the actual +# polynomial transition/measurement map, rather than as a linear map on the +# pruned augmented state. If u = [xₜ₋₁; εₜ] is Gaussian, the map is expanded +# around E[u] and its moments are closed analytically. At second order this +# requires fourth Gaussian moments; the cubic extension additionally uses the +# sixth moments through the third Hermite component. + +const IVASHCHENKO_STATIONARY_MAXITER = 500 +const IVASHCHENKO_STATIONARY_TOLERANCE = 1e-10 + +polynomial_pair_index(i::Int, j::Int, n::Int) = (i - 1) * n + j +polynomial_triple_index(i::Int, j::Int, k::Int, n::Int) = ((i - 1) * n + j - 1) * n + k + +function build_ivashchenko_kalman_system_from_constants(cons, 𝐒, observables_index::Vector{Int}, order::Symbol) + T = cons.post_model_macro + nVars, nPast, nExo = T.nVars, T.nPast_not_future_and_mixed, T.nExo + past = collect(T.past_not_future_and_mixed_idx) + # Keep all model-variable rows. The likelihood still selects only the + # state and observable blocks, while the full row set lets the estimate API + # return a value for every model variable and makes smoothing well-defined. + output_rows = collect(1:nVars) + nout = length(output_rows) + dv = nPast + 1 + nExo + d = nPast + nExo + + order ∈ (:second_order, :third_order) || + throw(ArgumentError("The Ivashchenko filter supports only second- and third-order solutions.")) + length(𝐒) ≥ (order == :third_order ? 3 : 2) || + throw(DimensionMismatch("The $(order) solution must provide S₁, S₂$(order == :third_order ? ", and S₃" : "").")) + + S1 = Matrix(𝐒[1][output_rows, :]) + S2 = Matrix(𝐒[2][output_rows, :]) + size(S1, 2) == dv || throw(DimensionMismatch("S₁ has $(size(S1, 2)) columns; expected $dv.")) + size(S2, 2) == dv^2 || throw(DimensionMismatch("S₂ has $(size(S2, 2)) columns; expected $(dv^2).")) + + scalar_type = order == :third_order ? promote_type(eltype(S1), eltype(S2), eltype(𝐒[3])) : promote_type(eltype(S1), eltype(S2)) + S1 = Matrix{scalar_type}(S1) + S2 = Matrix{scalar_type}(S2) + + # The raw solution is multiplied by symmetric Kronecker products. Store + # the symmetrised derivative tensors; this is algebraically equivalent and + # makes the Gaussian contractions below use the actual Hessian/third + # derivative of the polynomial. + H = zeros(scalar_type, nout, dv, dv) + @inbounds for a in 1:nout, i in 1:dv, j in 1:dv + H[a, i, j] = (S2[a, polynomial_pair_index(i, j, dv)] + + S2[a, polynomial_pair_index(j, i, dv)]) / 2 + end + + third_derivative = nothing + if order == :third_order + S3 = Matrix(𝐒[3][output_rows, :]) + size(S3, 2) == dv^3 || throw(DimensionMismatch("S₃ has $(size(S3, 2)) columns; expected $(dv^3).")) + S3 = Matrix{scalar_type}(S3) + third_derivative = zeros(scalar_type, nout, dv, dv, dv) + @inbounds for a in 1:nout, i in 1:dv, j in 1:dv, k in 1:dv + third_derivative[a, i, j, k] = ( + S3[a, polynomial_triple_index(i, j, k, dv)] + + S3[a, polynomial_triple_index(i, k, j, dv)] + + S3[a, polynomial_triple_index(j, i, k, dv)] + + S3[a, polynomial_triple_index(j, k, i, dv)] + + S3[a, polynomial_triple_index(k, i, j, dv)] + + S3[a, polynomial_triple_index(k, j, i, dv)]) / 6 + end + end + + row_position = zeros(Int, nVars) + @inbounds for (i, row) in enumerate(output_rows) + row_position[row] = i + end + state_position = row_position[past] + observation_position = row_position[observables_index] + random_indices = vcat(collect(1:nPast), collect(nPast + 2:dv)) + + # The linear-limit covariance is a useful starting point for the nonlinear + # stationary fixed point and also makes the :diagonal option match Kalman's + # convention exactly. + state_s1 = S1[state_position, :] + A = state_s1[:, 1:nPast] + B = state_s1[:, nPast + 2:dv] + + return (; order, nVars, nPast, nExo, d, dv, nout, output_rows, past, + S1, H, third_derivative, state_position, observation_position, + random_indices, A, B) +end + +build_ivashchenko_kalman_system(𝓂::ℳ, 𝐒, oi::Vector{Int}, order::Symbol) = + build_ivashchenko_kalman_system_from_constants(𝓂.constants, 𝐒, oi, order) + +function ivashchenko_kalman_workspace(sys, scalar_type::Type) + nout, d, dv = sys.nout, sys.d, sys.dv + third = sys.order == :third_order + return (; covariance_input = zeros(scalar_type, d, d), + vbar = zeros(scalar_type, dv), + mean = zeros(scalar_type, nout), + linear = zeros(scalar_type, nout, d), + effective_linear = zeros(scalar_type, nout, d), + hessian = zeros(scalar_type, nout, d, d), + hessian_covariance = zeros(scalar_type, nout, d, d), + covariance = zeros(scalar_type, nout, nout), + third_derivative = third ? zeros(scalar_type, nout, d, d, d) : nothing, + third_covariance = third ? zeros(scalar_type, nout, d, d, d) : nothing, + third_scratch = third ? zeros(scalar_type, nout, d, d, d) : nothing) +end + +function ivashchenko_transform_third_tensor!(destination, scratch, tensor, covariance) + nout, d = size(tensor, 1), size(tensor, 2) + @inbounds for a in 1:nout, p in 1:d, j in 1:d, k in 1:d + value = zero(eltype(destination)) + for i in 1:d + value += covariance[p, i] * tensor[a, i, j, k] + end + scratch[a, p, j, k] = value + end + @inbounds for a in 1:nout, p in 1:d, q in 1:d, k in 1:d + value = zero(eltype(destination)) + for j in 1:d + value += covariance[q, j] * scratch[a, p, j, k] + end + destination[a, p, q, k] = value + end + @inbounds for a in 1:nout, p in 1:d, q in 1:d, r in 1:d + value = zero(eltype(destination)) + for k in 1:d + value += covariance[r, k] * destination[a, p, q, k] + end + scratch[a, p, q, r] = value + end + destination .= scratch + return nothing +end + +function ivashchenko_polynomial_moments!(sys, mean_state, covariance_state, ws) + nPast, d, dv, nout = sys.nPast, sys.d, sys.dv, sys.nout + Σ = ws.covariance_input + fill!(Σ, zero(eltype(Σ))) + Σ[1:nPast, 1:nPast] .= covariance_state + @inbounds for i in nPast + 1:d + Σ[i, i] = one(eltype(Σ)) + end + + fill!(ws.vbar, zero(eltype(ws.vbar))) + ws.vbar[1:nPast] .= mean_state + ws.vbar[nPast + 1] = one(eltype(ws.vbar)) + + @inbounds for a in 1:nout + Hfull = view(sys.H, a, :, :) + third_full = sys.third_derivative === nothing ? nothing : view(sys.third_derivative, a, :, :, :) + value = ℒ.dot(view(sys.S1, a, :), ws.vbar) + for i in 1:dv, j in 1:dv + value += Hfull[i, j] * ws.vbar[i] * ws.vbar[j] / 2 + end + if third_full !== nothing + for i in 1:dv, j in 1:dv, k in 1:dv + value += third_full[i, j, k] * ws.vbar[i] * ws.vbar[j] * ws.vbar[k] / 6 + end + end + + for r in 1:d + i = sys.random_indices[r] + slope = sys.S1[a, i] + for j in 1:dv + slope += Hfull[i, j] * ws.vbar[j] + end + if third_full !== nothing + for j in 1:dv, k in 1:dv + slope += third_full[i, j, k] * ws.vbar[j] * ws.vbar[k] / 2 + end + end + ws.linear[a, r] = slope + end + + for r in 1:d, s in 1:d + i, j = sys.random_indices[r], sys.random_indices[s] + hessian = Hfull[i, j] + if third_full !== nothing + for k in 1:dv + hessian += third_full[i, j, k] * ws.vbar[k] + end + end + ws.hessian[a, r, s] = hessian + end + + if third_full !== nothing + for r in 1:d, s in 1:d, q in 1:d + ws.third_derivative[a, r, s, q] = third_full[ + sys.random_indices[r], sys.random_indices[s], sys.random_indices[q]] + end + end + + ws.mean[a] = value + sum(ws.hessian[a, r, s] * Σ[r, s] for r in 1:d, s in 1:d) / 2 + end + + ws.effective_linear .= ws.linear + if sys.order == :third_order + @inbounds for a in 1:nout, r in 1:d + correction = zero(eltype(ws.effective_linear)) + for s in 1:d, q in 1:d + correction += ws.third_derivative[a, r, s, q] * Σ[s, q] / 2 + end + ws.effective_linear[a, r] += correction + end + end + + @inbounds for a in 1:nout + H = view(ws.hessian, a, :, :) + ws.hessian_covariance[a, :, :] .= Σ * H * Σ + end + ws.covariance .= ws.effective_linear * Σ * ws.effective_linear' + ws.covariance .+= reshape(ws.hessian, nout, d^2) * reshape(ws.hessian_covariance, nout, d^2)' / 2 + + if sys.order == :third_order + ivashchenko_transform_third_tensor!(ws.third_covariance, ws.third_scratch, + ws.third_derivative, Σ) + ws.covariance .+= reshape(ws.third_derivative, nout, d^3) * + reshape(ws.third_covariance, nout, d^3)' / 6 + end + + @inbounds for j in 1:nout, i in 1:j + value = (ws.covariance[i, j] + ws.covariance[j, i]) / 2 + ws.covariance[i, j] = value + ws.covariance[j, i] = value + end + return ws.mean, ws.covariance +end + +function ivashchenko_stationary_initialization(sys, initial_mean, initial_covariance, ws; + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling) + scalar_type = eltype(ws.vbar) + if initial_covariance isa AbstractMatrix + size(initial_covariance) == (sys.nPast, sys.nPast) || + throw(DimensionMismatch("The Ivashchenko initial covariance must be $(sys.nPast)×$(sys.nPast), got $(size(initial_covariance)).")) + return Vector{scalar_type}(initial_mean), Matrix{scalar_type}(initial_covariance), true + elseif initial_covariance == :diagonal + return Vector{scalar_type}(initial_mean), Matrix{scalar_type}(10 .* ℒ.I(sys.nPast)), true + elseif initial_covariance != :theoretical + throw(ArgumentError("Unsupported Ivashchenko initial covariance: $(initial_covariance).")) + end + + # Start from the linear stationary covariance, then solve the coupled + # nonlinear mean/covariance fixed point implied by the unpruned polynomial. + P = qkf_lyapunov(sys.A, sys.B * sys.B'; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + m = Vector{scalar_type}(initial_mean) + converged = false + for _ in 1:IVASHCHENKO_STATIONARY_MAXITER + ivashchenko_polynomial_moments!(sys, m, P, ws) + next_m = collect(view(ws.mean, sys.state_position)) + next_P = Matrix(view(ws.covariance, sys.state_position, sys.state_position)) + next_P = (next_P + next_P') / 2 + delta_m = isempty(next_m) ? 0.0 : maximum(abs, primal.(next_m .- m)) + delta_P = isempty(next_P) ? 0.0 : maximum(abs, primal.(next_P .- P)) + if !all(isfinite, primal.(next_m)) || !all(isfinite, primal.(next_P)) + return m, P, false + end + m, P = next_m, next_P + if max(delta_m, delta_P) <= IVASHCHENKO_STATIONARY_TOLERANCE * + max(1.0, maximum(abs, primal.(m)), maximum(abs, primal.(P))) + converged = true + break + end + end + return m, P, converged +end + +function ivashchenko_measurement_covariance(measurement_error, n_obs, scalar_type) + if measurement_error === nothing + return zeros(scalar_type, n_obs, n_obs) + elseif measurement_error isa AbstractMatrix + size(measurement_error) == (n_obs, n_obs) || + throw(DimensionMismatch("Ivashchenko measurement error must be $n_obs×$n_obs.")) + return Matrix{scalar_type}(measurement_error) + elseif measurement_error isa AbstractVector + length(measurement_error) == n_obs || + throw(DimensionMismatch("Ivashchenko measurement error must have $n_obs entries.")) + return Matrix{scalar_type}(ℒ.Diagonal(collect(measurement_error))) + else + return Matrix{scalar_type}(ℒ.I(n_obs)) .* measurement_error + end +end + +function ivashchenko_subset_measurement_covariance(Hm, idx) + isempty(idx) && return Matrix{eltype(Hm)}(undef, 0, 0) + return Matrix(Hm[idx, idx]) +end + +function ivashchenko_copy_moment_tape(ws) + return (; covariance_input = copy(ws.covariance_input), + vbar = copy(ws.vbar), mean = copy(ws.mean), linear = copy(ws.linear), + effective_linear = copy(ws.effective_linear), hessian = copy(ws.hessian), + hessian_covariance = copy(ws.hessian_covariance), covariance = copy(ws.covariance), + third_derivative = ws.third_derivative === nothing ? nothing : copy(ws.third_derivative), + third_covariance = ws.third_covariance === nothing ? nothing : copy(ws.third_covariance)) +end + +function ivashchenko_stationary_initialization_taped(sys, initial_mean, initial_covariance, ws; + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling) + scalar_type = eltype(ws.vbar) + if initial_covariance isa AbstractMatrix + return Vector{scalar_type}(initial_mean), Matrix{scalar_type}(initial_covariance), true, nothing + elseif initial_covariance == :diagonal + return Vector{scalar_type}(initial_mean), Matrix{scalar_type}(10 .* ℒ.I(sys.nPast)), true, nothing + elseif initial_covariance != :theoretical + throw(ArgumentError("Unsupported Ivashchenko initial covariance: $(initial_covariance).")) + end + + P = qkf_lyapunov(sys.A, sys.B * sys.B'; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + linear_covariance = copy(P) + m = Vector{scalar_type}(initial_mean) + iterations = NamedTuple[] + converged = false + for _ in 1:IVASHCHENKO_STATIONARY_MAXITER + input_mean = copy(m) + input_covariance = copy(P) + ivashchenko_polynomial_moments!(sys, m, P, ws) + moment_tape = ivashchenko_copy_moment_tape(ws) + next_m = collect(view(ws.mean, sys.state_position)) + next_P = Matrix(view(ws.covariance, sys.state_position, sys.state_position)) + next_P = (next_P + next_P') / 2 + push!(iterations, (; input_mean, input_covariance, moment_tape)) + delta_m = isempty(next_m) ? 0.0 : maximum(abs, primal.(next_m .- m)) + delta_P = isempty(next_P) ? 0.0 : maximum(abs, primal.(next_P .- P)) + if !all(isfinite, primal.(next_m)) || !all(isfinite, primal.(next_P)) + return m, P, false, (; iterations, linear_covariance) + end + m, P = next_m, next_P + if max(delta_m, delta_P) <= IVASHCHENKO_STATIONARY_TOLERANCE * + max(1.0, maximum(abs, primal.(m)), maximum(abs, primal.(P))) + converged = true + break + end + end + return m, P, converged, (; iterations, linear_covariance) +end + +function ivashchenko_filter_pass(sys, data_in_deviations::AbstractMatrix{<:Real}, initial_mean; + measurement_error = nothing, + initial_covariance = :theoretical, + presample_periods::Int = 0, + on_failure_loglikelihood::Real = -Inf, + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling, + record::Bool = false) + n_obs, nT = size(data_in_deviations) + presample_periods = normalize_presample_periods(presample_periods, nT) + scalar_type = promote_type(eltype(sys.S1), eltype(data_in_deviations), + measurement_error === nothing ? Float64 : + (measurement_error isa AbstractArray ? eltype(measurement_error) : typeof(measurement_error))) + Hm = ivashchenko_measurement_covariance(measurement_error, n_obs, scalar_type) + obs_idx_per_t, _ = build_obs_index(data_in_deviations) + + ws = ivashchenko_kalman_workspace(sys, scalar_type) + if record + mean_state, covariance_state, initialized, initialization_tape = + ivashchenko_stationary_initialization_taped(sys, initial_mean, initial_covariance, ws; + workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + else + mean_state, covariance_state, initialized = ivashchenko_stationary_initialization( + sys, initial_mean, initial_covariance, ws; + workspaces = workspaces, lyapunov_algorithm = lyapunov_algorithm) + initialization_tape = nothing + end + initialized || return record ? (convert(scalar_type, on_failure_loglikelihood), nothing) : + convert(scalar_type, on_failure_loglikelihood) + + state_position, observation_position = sys.state_position, sys.observation_position + n_state = length(state_position) + post_means = record ? Vector{Vector{scalar_type}}(undef, nT) : nothing + post_covariances = record ? Vector{Matrix{scalar_type}}(undef, nT) : nothing + input_means = record ? Vector{Vector{scalar_type}}(undef, nT) : nothing + input_covariances = record ? Vector{Matrix{scalar_type}}(undef, nT) : nothing + moment_tapes = record ? Vector{Any}(undef, nT) : nothing + predicted_means = record ? Vector{Vector{scalar_type}}(undef, nT) : nothing + predicted_covariances = record ? Vector{Matrix{scalar_type}}(undef, nT) : nothing + output_means = record ? Vector{Vector{scalar_type}}(undef, nT) : nothing + output_covariances = record ? Vector{Matrix{scalar_type}}(undef, nT) : nothing + transitions = record ? Vector{Matrix{scalar_type}}(undef, nT) : nothing + shock_loadings = record ? Vector{Matrix{scalar_type}}(undef, nT) : nothing + innovations = record ? Vector{Any}(undef, nT) : nothing + inverse_innovation_covariances = record ? Vector{Any}(undef, nT) : nothing + gains = record ? Vector{Any}(undef, nT) : nothing + cross_covariances = record ? Vector{Any}(undef, nT) : nothing + observed_indices = record ? obs_idx_per_t : nothing + ll = zero(scalar_type) + log2pi = log(2π) + @inbounds for t in 1:nT + input_mean = record ? copy(mean_state) : nothing + input_covariance = record ? copy(covariance_state) : nothing + ivashchenko_polynomial_moments!(sys, mean_state, covariance_state, ws) + if record + input_means[t] = input_mean + input_covariances[t] = input_covariance + moment_tapes[t] = ivashchenko_copy_moment_tape(ws) + end + predicted_mean = collect(view(ws.mean, state_position)) + output_mean = copy(ws.mean) + predicted_covariance = Matrix(view(ws.covariance, state_position, state_position)) + output_covariance = copy(ws.covariance) + idx = obs_idx_per_t[t] + m = length(idx) + + if record + predicted_means[t] = predicted_mean + predicted_covariances[t] = predicted_covariance + output_means[t] = output_mean + output_covariances[t] = output_covariance + transitions[t] = Matrix(view(ws.effective_linear, state_position, 1:sys.nPast)) + shock_loadings[t] = Matrix(view(ws.effective_linear, state_position, sys.nPast + 1:sys.d)) + end + + if m == 0 + if record + innovations[t] = nothing + inverse_innovation_covariances[t] = nothing + gains[t] = nothing + cross_covariances[t] = nothing + end + mean_state = predicted_mean + covariance_state = (predicted_covariance + predicted_covariance') / 2 + if record + post_means[t] = copy(mean_state) + post_covariances[t] = copy(covariance_state) + end + continue + end + + observation_mean = collect(view(ws.mean, observation_position[idx])) + observation_covariance = Matrix(view(ws.covariance, observation_position[idx], observation_position[idx])) + cross_covariance = Matrix(view(ws.covariance, state_position, observation_position[idx])) + observations = collect(view(data_in_deviations, idx, t)) + innovation = observations - observation_mean + F = observation_covariance + ivashchenko_subset_measurement_covariance(Hm, idx) + F = (F + F') / 2 + factor = ℒ.lu(F, check = false) + ℒ.issuccess(factor) || return record ? (convert(scalar_type, on_failure_loglikelihood), nothing) : + convert(scalar_type, on_failure_loglikelihood) + logabsdetF, signF = ℒ.logabsdet(factor) + (primal(signF) > 0 && isfinite(primal(logabsdetF))) || + return record ? (convert(scalar_type, on_failure_loglikelihood), nothing) : + convert(scalar_type, on_failure_loglikelihood) + invF = Matrix(factor \ ℒ.I(m)) + + if t > presample_periods + solved_innovation = invF * innovation + ll -= (ℒ.dot(innovation, solved_innovation) + logabsdetF + m * log2pi) / 2 + isfinite(primal(ll)) || return record ? (convert(scalar_type, on_failure_loglikelihood), nothing) : + convert(scalar_type, on_failure_loglikelihood) + end + + gain = cross_covariance * invF + if record + innovations[t] = copy(innovation) + inverse_innovation_covariances[t] = copy(invF) + gains[t] = copy(gain) + cross_covariances[t] = copy(cross_covariance) + end + mean_state = predicted_mean + gain * innovation + covariance_state = predicted_covariance - gain * cross_covariance' + covariance_state = (covariance_state + covariance_state') / 2 + if record + post_means[t] = copy(mean_state) + post_covariances[t] = copy(covariance_state) + end + end + + if !record + return ll + end + return ll, (; initialization_tape, input_means, input_covariances, moment_tapes, + initial_mean = Vector{scalar_type}(initial_mean), + initial_covariance = copy(input_covariances[1]), + predicted_means, predicted_covariances, post_means, post_covariances, + output_means, output_covariances, transitions, shock_loadings, + innovations, inverse_innovation_covariances, gains, cross_covariances, + observed_indices, Hm, presample_periods, data = Matrix{scalar_type}(data_in_deviations), + state_position, observation_position) +end + +function run_ivashchenko_kalman(sys, data_in_deviations::AbstractMatrix{<:Real}, initial_mean; + measurement_error = nothing, + initial_covariance = :theoretical, + presample_periods::Int = 0, + on_failure_loglikelihood::Real = -Inf, + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling) + return ivashchenko_filter_pass(sys, data_in_deviations, initial_mean; + measurement_error = measurement_error, + initial_covariance = initial_covariance, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) +end + +function ivashchenko_polynomial_moments_pullback(sys, moment_tape, mean_bar, covariance_bar) + nout, d, dv = sys.nout, sys.d, sys.dv + Σ = moment_tape.covariance_input + v = moment_tape.vbar + L = moment_tape.effective_linear + K = moment_tape.hessian + third = moment_tape.third_derivative + gΣ = zeros(eltype(Σ), d, d) + gv = zeros(eltype(v), dv) + gL = zeros(eltype(L), nout, d) + gK = zeros(eltype(K), nout, d, d) + gS1 = zeros(eltype(sys.S1), nout, dv) + gH = zeros(eltype(sys.H), nout, dv, dv) + gthird = third === nothing ? nothing : zeros(eltype(third), nout, dv, dv, dv) + + # C₁ = L Σ L'. + @inbounds for a in 1:nout, b in 1:nout, r in 1:d, s in 1:d + gL[a, r] += (covariance_bar[a, b] + covariance_bar[b, a]) * L[b, s] * Σ[r, s] + gΣ[r, s] += covariance_bar[a, b] * L[a, r] * L[b, s] + end + + # C₂ = 1/2 K : (Σ K Σ), written with matrix contractions to keep the cubic + # reverse pass manageable for the unpruned state dimension. + hessian_covariance_bar = zeros(eltype(K), nout, d, d) + @inbounds for b in 1:nout, i in 1:d, j in 1:d, a in 1:nout + hessian_covariance_bar[b, i, j] += covariance_bar[a, b] * K[a, i, j] / 2 + gK[a, i, j] += covariance_bar[a, b] * moment_tape.hessian_covariance[b, i, j] / 2 + end + @inbounds for b in 1:nout, i in 1:d, j in 1:d, p in 1:d, q in 1:d + value = hessian_covariance_bar[b, i, j] + gK[b, p, q] += value * Σ[i, p] * Σ[q, j] + gΣ[i, p] += value * K[b, p, q] * Σ[q, j] + gΣ[q, j] += value * Σ[i, p] * K[b, p, q] + end + + if third !== nothing + # C₃ = 1/6 T : (Σ ⊗ Σ ⊗ Σ)T'. + third_covariance_bar = zeros(eltype(third), nout, d, d, d) + @inbounds for b in 1:nout, i in 1:d, j in 1:d, k in 1:d, a in 1:nout + third_covariance_bar[b, i, j, k] += covariance_bar[a, b] * third[a, i, j, k] / 6 + gthird[a, sys.random_indices[i], sys.random_indices[j], sys.random_indices[k]] += + covariance_bar[a, b] * moment_tape.third_covariance[b, i, j, k] / 6 + end + @inbounds for b in 1:nout, i in 1:d, j in 1:d, k in 1:d, + p in 1:d, q in 1:d, r in 1:d + value = third_covariance_bar[b, i, j, k] * third[b, p, q, r] + gthird[b, sys.random_indices[p], sys.random_indices[q], sys.random_indices[r]] += + third_covariance_bar[b, i, j, k] * Σ[i, p] * Σ[j, q] * Σ[k, r] + gΣ[i, p] += value * Σ[j, q] * Σ[k, r] + gΣ[j, q] += value * Σ[i, p] * Σ[k, r] + gΣ[k, r] += value * Σ[i, p] * Σ[j, q] + end + end + + # The mean is f(v) + 1/2 K:Σ. + @inbounds for a in 1:nout, r in 1:d, s in 1:d + gK[a, r, s] += mean_bar[a] * Σ[r, s] / 2 + gΣ[r, s] += mean_bar[a] * K[a, r, s] / 2 + end + + # L_eff = L₀ + 1/2 T⋅Σ. + @inbounds for a in 1:nout, r in 1:d + if third !== nothing + for s in 1:d, q in 1:d + i = sys.random_indices[r] + j = sys.random_indices[s] + k = sys.random_indices[q] + gthird[a, i, j, k] += gL[a, r] * Σ[s, q] / 2 + gΣ[s, q] += gL[a, r] * third[a, r, s, q] / 2 + end + end + end + linear_bar = gL + + # K = H[random, random] + T[random, random, :]⋅v. + @inbounds for a in 1:nout, r in 1:d, s in 1:d + i, j = sys.random_indices[r], sys.random_indices[s] + gH[a, i, j] += gK[a, r, s] + if third !== nothing + for q in 1:dv + gthird[a, i, j, q] += gK[a, r, s] * v[q] + gv[q] += gK[a, r, s] * sys.third_derivative[a, i, j, q] + end + end + end + + # L₀ = S₁[random] + H[random, :]v + 1/2 T[random, :, :]vv. + @inbounds for a in 1:nout, r in 1:d + i = sys.random_indices[r] + gS1[a, i] += linear_bar[a, r] + for j in 1:dv + gH[a, i, j] += linear_bar[a, r] * v[j] + gv[j] += linear_bar[a, r] * sys.H[a, i, j] + end + if third !== nothing + for j in 1:dv, k in 1:dv + value = linear_bar[a, r] * v[j] * v[k] / 2 + gthird[a, i, j, k] += value + gv[j] += linear_bar[a, r] * sys.third_derivative[a, i, j, k] * v[k] / 2 + gv[k] += linear_bar[a, r] * sys.third_derivative[a, i, j, k] * v[j] / 2 + end + end + end + + # f(v) = S₁v + 1/2 Hv² + 1/6 Tv³. + @inbounds for a in 1:nout + coefficient_bar = mean_bar[a] + for i in 1:dv + gS1[a, i] += coefficient_bar * v[i] + gv[i] += coefficient_bar * sys.S1[a, i] + end + for i in 1:dv, j in 1:dv + value = coefficient_bar * v[i] * v[j] / 2 + gH[a, i, j] += value + gv[i] += coefficient_bar * sys.H[a, i, j] * v[j] / 2 + gv[j] += coefficient_bar * sys.H[a, i, j] * v[i] / 2 + end + if third !== nothing + for i in 1:dv, j in 1:dv, k in 1:dv + value = coefficient_bar * v[i] * v[j] * v[k] / 6 + gthird[a, i, j, k] += value + gv[i] += coefficient_bar * sys.third_derivative[a, i, j, k] * v[j] * v[k] / 6 + gv[j] += coefficient_bar * sys.third_derivative[a, i, j, k] * v[i] * v[k] / 6 + gv[k] += coefficient_bar * sys.third_derivative[a, i, j, k] * v[i] * v[j] / 6 + end + end + end + + gΣ .= (gΣ + gΣ') / 2 + gmean = copy(view(gv, 1:sys.nPast)) + gcovariance = copy(view(gΣ, 1:sys.nPast, 1:sys.nPast)) + return gmean, gcovariance, gS1, gH, gthird +end + +function ivashchenko_solution_matrix_pullback(sys, gS1, gH, gthird, solution_matrices) + gS2 = zeros(eltype(solution_matrices[2]), size(solution_matrices[2])) + @inbounds for a in 1:sys.nout, i in 1:sys.dv, j in 1:sys.dv + gS2[a, polynomial_pair_index(i, j, sys.dv)] += (gH[a, i, j] + gH[a, j, i]) / 2 + end + if gthird === nothing + return gS1, gS2, nothing + end + gS3 = zeros(eltype(solution_matrices[3]), size(solution_matrices[3])) + @inbounds for a in 1:sys.nout, i in 1:sys.dv, j in 1:sys.dv, k in 1:sys.dv + gS3[a, polynomial_triple_index(i, j, k, sys.dv)] += ( + gthird[a, i, j, k] + gthird[a, i, k, j] + + gthird[a, j, i, k] + gthird[a, j, k, i] + + gthird[a, k, i, j] + gthird[a, k, j, i]) / 6 + end + return gS1, gS2, gS3 +end + +function ivashchenko_lyapunov_pullback(A, B, P, covariance_bar; lyapunov_algorithm::Symbol = :doubling) + adjoint_covariance = qkf_lyapunov(A', covariance_bar; + lyapunov_algorithm = lyapunov_algorithm) + adjoint_covariance = (adjoint_covariance + adjoint_covariance') / 2 + adjoint_A = adjoint_covariance * A * P' + adjoint_covariance' * A * P + adjoint_Q = adjoint_covariance + adjoint_B = (adjoint_Q + adjoint_Q') * B + return adjoint_A, adjoint_B +end + +function ivashchenko_filter_pullback(sys, tape, solution_matrices, scale; + initial_covariance = :theoretical, + lyapunov_algorithm::Symbol = :doubling) + nT = length(tape.post_means) + scalar_type = eltype(tape.post_means[1]) + mean_bar = zeros(scalar_type, sys.nPast) + covariance_bar = zeros(scalar_type, sys.nPast, sys.nPast) + data_bar = zeros(scalar_type, size(tape.data)) + gS1 = zeros(scalar_type, size(sys.S1)) + gH = zeros(scalar_type, size(sys.H)) + gthird = sys.third_derivative === nothing ? nothing : zeros(scalar_type, size(sys.third_derivative)) + + @inbounds for t in nT:-1:1 + idx = tape.observed_indices[t] + mean_output_bar = zeros(scalar_type, sys.nout) + covariance_output_bar = zeros(scalar_type, sys.nout, sys.nout) + predicted_mean_bar = copy(mean_bar) + predicted_covariance_bar = copy(covariance_bar) + if isempty(idx) + mean_output_bar[sys.state_position] .+= predicted_mean_bar + covariance_output_bar[sys.state_position, sys.state_position] .+= predicted_covariance_bar + else + innovation = tape.innovations[t] + invF = tape.inverse_innovation_covariances[t] + gain = tape.gains[t] + cross = tape.cross_covariances[t] + mean_bar_post = mean_bar + # The forward covariance update is explicitly symmetrised; only + # the symmetric part of its cotangent can reach the pre-update + # covariance and gain. + covariance_bar_post = (covariance_bar + covariance_bar') / 2 + + gain_bar = mean_bar_post * innovation' + innovation_bar = gain' * mean_bar_post + gain_bar .-= covariance_bar_post * cross + cross_bar = -covariance_bar_post' * gain + + inverse_covariance_bar = cross' * gain_bar + cross_bar .+= gain_bar * invF' + if t > tape.presample_periods + innovation_bar .-= scale * (invF * innovation) + end + inverse_covariance_bar .+= zero(scalar_type) + covariance_bar_innovation = -invF' * inverse_covariance_bar * invF' + if t > tape.presample_periods + covariance_bar_innovation .-= scale / 2 * + (invF' - invF' * innovation * innovation' * invF') + end + data_bar[idx, t] .+= innovation_bar + mean_output_bar[sys.state_position] .+= predicted_mean_bar + mean_output_bar[sys.observation_position[idx]] .-= innovation_bar + covariance_output_bar[sys.state_position, sys.state_position] .+= predicted_covariance_bar + covariance_output_bar[sys.observation_position[idx], sys.observation_position[idx]] .+= + (covariance_bar_innovation + covariance_bar_innovation') / 2 + # The moment covariance is explicitly symmetric, so split the + # cross-block cotangent across its two transpose locations. + covariance_output_bar[sys.state_position, sys.observation_position[idx]] .+= cross_bar / 2 + covariance_output_bar[sys.observation_position[idx], sys.state_position] .+= cross_bar' / 2 + end + + mean_bar, covariance_bar, local_S1, local_H, local_third = + ivashchenko_polynomial_moments_pullback(sys, tape.moment_tapes[t], + mean_output_bar, covariance_output_bar) + gS1 .+= local_S1 + gH .+= local_H + if gthird !== nothing + gthird .+= local_third + end + end + + if tape.initialization_tape !== nothing + initialization_mean_bar = copy(mean_bar) + initialization_covariance_bar = copy(covariance_bar) + for iteration in reverse(tape.initialization_tape.iterations) + output_mean_bar = zeros(scalar_type, sys.nout) + output_covariance_bar = zeros(scalar_type, sys.nout, sys.nout) + output_mean_bar[sys.state_position] .= initialization_mean_bar + output_covariance_bar[sys.state_position, sys.state_position] .= + (initialization_covariance_bar + initialization_covariance_bar') / 2 + initialization_mean_bar, initialization_covariance_bar, local_S1, local_H, local_third = + ivashchenko_polynomial_moments_pullback(sys, iteration.moment_tape, + output_mean_bar, output_covariance_bar) + gS1 .+= local_S1 + gH .+= local_H + if gthird !== nothing + gthird .+= local_third + end + end + initial_covariance_bar = (initialization_covariance_bar + initialization_covariance_bar') / 2 + initial_A_bar, initial_B_bar = ivashchenko_lyapunov_pullback( + sys.A, sys.B, tape.initialization_tape.linear_covariance, initial_covariance_bar; + lyapunov_algorithm = lyapunov_algorithm) + gS1[sys.state_position, 1:sys.nPast] .+= initial_A_bar + gS1[sys.state_position, sys.nPast + 2:sys.dv] .+= initial_B_bar + mean_bar = initialization_mean_bar + end + + local_S1, local_S2, local_S3 = ivashchenko_solution_matrix_pullback( + sys, gS1, gH, gthird, solution_matrices) + solution_bar = [zeros(scalar_type, size(solution_matrices[1])), + zeros(scalar_type, size(solution_matrices[2]))] + solution_bar[1][sys.output_rows, :] .= local_S1 + solution_bar[2][sys.output_rows, :] .= local_S2 + if local_S3 !== nothing + solution_bar = vcat(solution_bar, [zeros(scalar_type, size(solution_matrices[3]))]) + solution_bar[3][sys.output_rows, :] .= local_S3 + end + state_bar = zeros(scalar_type, sys.nVars) + state_bar[sys.past] .= mean_bar + return solution_bar, data_bar, state_bar +end + +function ivashchenko_smooth_pass(sys, tape) + nT = length(tape.post_means) + n_state = length(tape.state_position) + smoothed_means = [copy(tape.post_means[t]) for t in 1:nT] + smoothed_covariances = [copy(tape.post_covariances[t]) for t in 1:nT] + for t in nT - 1:-1:1 + transition = tape.transitions[t + 1] + cross = transition * tape.post_covariances[t] + smoother_gain = cross * inv(tape.predicted_covariances[t + 1]) + delta = smoothed_means[t + 1] - tape.predicted_means[t + 1] + smoothed_means[t] .+= smoother_gain * delta + smoothed_covariances[t] .= tape.post_covariances[t] + + smoother_gain * (smoothed_covariances[t + 1] - tape.predicted_covariances[t + 1]) * smoother_gain' + smoothed_covariances[t] .= (smoothed_covariances[t] + smoothed_covariances[t]') / 2 + end + + variables = zeros(eltype(tape.post_means[1]), sys.nVars, nT) + standard_deviations = zeros(eltype(tape.post_means[1]), sys.nVars, nT) + shocks = zeros(eltype(tape.post_means[1]), sys.nExo, nT) + @inbounds for t in 1:nT + pred_covariance = tape.predicted_covariances[t] + state_delta = smoothed_means[t] - tape.predicted_means[t] + state_regression = tape.output_covariances[t][ :, tape.state_position] * inv(pred_covariance) + variables[:, t] .= tape.output_means[t] + state_regression * state_delta + variables[tape.state_position, t] .= smoothed_means[t] + standard_deviations[:, t] .= sqrt.(abs.(ℒ.diag(tape.output_covariances[t] - + state_regression * pred_covariance * state_regression'))) + standard_deviations[tape.state_position, t] .= sqrt.(abs.(ℒ.diag(smoothed_covariances[t]))) + + shock_regression = tape.shock_loadings[t]' * inv(pred_covariance) + shocks[:, t] .= shock_regression * state_delta + end + + decomposition = zeros(eltype(variables), sys.nVars, sys.nExo + 2, nT) + decomposition[:, end - 1, :] .= variables + return variables, shocks, standard_deviations, decomposition, smoothed_means, smoothed_covariances +end + +function ivashchenko_filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + order::Symbol; + initial_covariance = :theoretical, + measurement_error = nothing, + smooth::Bool = true, + opts::CalculationOptions = merge_calculation_options()) + constants = initialise_constants!(𝓂) + T = constants.post_model_macro + nT = size(data_in_deviations, 2) + variables = zeros(Float64, T.nVars, nT) + shocks = zeros(Float64, T.nExo, nT) + standard_deviations = zeros(Float64, T.nVars, nT) + decomposition = zeros(Float64, T.nVars, T.nExo + 2, nT) + + result = calculate_stochastic_steady_state(Val(order), 𝓂.parameter_values, 𝓂, opts = opts) + sss, converged, SS_and_pars, solution_error = result[1:4] + if !converged || solution_error > opts.tol.nsss.acceptance_tol || !isfinite(solution_error) + @error "Could not find a stochastic steady state for the Ivashchenko filter." + return variables, shocks, standard_deviations, decomposition + end + 𝐒 = order == :second_order ? result[7:8] : result[8:10] + ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + all_SS = expand_steady_state(SS_and_pars, constants.post_complete_parameters) + state = collect(sss) - all_SS + observables = get_and_check_observables(T, data_in_deviations) + observable_indices = convert(Vector{Int}, indexin(observables, constants.post_complete_parameters.SS_and_pars_names)) + sys = build_ivashchenko_kalman_system_from_constants(constants, 𝐒, observable_indices, order) + data = collect(data_in_deviations) + pass = ivashchenko_filter_pass(sys, data, state[sys.past]; + measurement_error = measurement_error, + initial_covariance = initial_covariance, + presample_periods = 0, + workspaces = 𝓂.workspaces, + lyapunov_algorithm = opts.lyapunov_algorithm, + record = true) + pass[2] === nothing && return variables, shocks, standard_deviations, decomposition + if smooth + return ivashchenko_smooth_pass(sys, pass[2])[1:4] + end + + tape = pass[2] + @inbounds for t in 1:nT + variables[:, t] .= tape.output_means[t] + standard_deviations[:, t] .= sqrt.(abs.(ℒ.diag(tape.output_covariances[t]))) + end + decomposition[:, end - 1, :] .= variables + return variables, shocks, standard_deviations, decomposition +end + +@unstable function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:second_order}, + ::Val{:ivashchenko_kalman}; + warmup_iterations::Int = 0, + initial_covariance = :theoretical, + measurement_error = nothing, + smooth::Bool = true, + opts::CalculationOptions = merge_calculation_options()) + return ivashchenko_filter_data_with_model(𝓂, data_in_deviations, :second_order; + initial_covariance = initial_covariance, + measurement_error = measurement_error, + smooth = smooth, opts = opts) +end + +@unstable function filter_data_with_model(𝓂::ℳ, + data_in_deviations::KeyedArray{Float64}, + ::Val{:third_order}, + ::Val{:ivashchenko_kalman}; + warmup_iterations::Int = 0, + initial_covariance = :theoretical, + measurement_error = nothing, + smooth::Bool = true, + opts::CalculationOptions = merge_calculation_options()) + return ivashchenko_filter_data_with_model(𝓂, data_in_deviations, :third_order; + initial_covariance = initial_covariance, + measurement_error = measurement_error, + smooth = smooth, opts = opts) +end + +function calculate_loglikelihood(::Val{:ivashchenko_kalman}, ::Val{O}, + observables_index::Vector{Int}, 𝐒, + data_in_deviations::AbstractMatrix, + constants, state, workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) where {O} + O ∈ (:second_order, :third_order) || + throw(ArgumentError("The Ivashchenko filter requires `algorithm = :second_order` or `:third_order`.")) + sys = build_ivashchenko_kalman_system_from_constants(constants, 𝐒, observables_index, O) + initial_mean = state[sys.past] + return run_ivashchenko_kalman(sys, data_in_deviations, initial_mean; + measurement_error = measurement_error, + initial_covariance = initial_covariance, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) +end + +function calculate_loglikelihood_with_missing(::Val{:ivashchenko_kalman}, ::Val{O}, + observables_index::Vector{Int}, 𝐒, + data_in_deviations::AbstractMatrix, + constants, state, workspaces, + obs_idx_per_t::Vector{Vector{Int}}; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) where {O} + O ∈ (:second_order, :third_order) || + throw(ArgumentError("The Ivashchenko filter requires `algorithm = :second_order` or `:third_order`.")) + sys = build_ivashchenko_kalman_system_from_constants(constants, 𝐒, observables_index, O) + initial_mean = state[sys.past] + return run_ivashchenko_kalman(sys, data_in_deviations, initial_mean; + measurement_error = measurement_error, + initial_covariance = initial_covariance, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) +end + +end # @stable diff --git a/src/filter/quadratic_kalman.jl b/src/filter/quadratic_kalman.jl index 526df8430..47a5f91ad 100644 --- a/src/filter/quadratic_kalman.jl +++ b/src/filter/quadratic_kalman.jl @@ -1,7 +1,6 @@ @stable default_mode = "disable" begin -# Quadratic Kalman filter (Monfort, Renne & Roussellet, 2015) for the pruned -# second-order solution. +# Kollmann-style quadratic Kalman filter for the pruned second-order solution. # # The idea. A pruned second-order solution is *linear* in an augmented state — # this is the pruned state-space representation of Andreasen, Fernández-Villaverde @@ -43,8 +42,8 @@ # Var(w) = G G' + H (I + K) H', K the commutation matrix, # # using E[(ε⊗ε)(ε⊗ε)'] = vec(I)vec(I)' + I + K. `H` is constant; `G` depends on the -# state, and is evaluated at the filtered mean each period — the defining choice of -# the quadratic Kalman filter. +# state. The plug-in term uses the filtered mean, while the recursion also adds the +# exact covariance of this affine loading under the filtered state covariance. # # Cost. The augmented state has dimension 2·nVars + nPast², which is 808 for # Smets-Wouters (2007). The covariance recursion is therefore O(nz³) per period and @@ -209,6 +208,58 @@ function quadratic_kalman_affine_G(sys) return g0, Λ end +# Conditional covariance of the quadratic innovation. If G(z) = Ḡ + Σᵢ zᵢGᵢ +# and z has covariance P, the state-dependent part contributes +# +# Σᵢⱼ Pᵢⱼ Gᵢ Gⱼ' +# +# in addition to the plug-in term ḠḠ'. Pz selects the past first-order state +# from z, so only that small covariance is needed here; the full augmented +# covariance is still required by the Kalman prediction/update itself. +function quadratic_kalman_noise_covariance!(Q, G, QH, Λ, Pz, Pc, PzPc, Pa, LPa) + ℒ.mul!(PzPc, Pz, Pc) + ℒ.mul!(Pa, PzPc, Pz') + copyto!(Q, QH) + ℒ.mul!(Q, G, G', one(eltype(Q)), one(eltype(Q))) + nz, nExo = size(G) + @inbounds for j in 1:nExo + rows = (j - 1) * nz + 1:j * nz + ℒ.mul!(LPa, view(Λ, rows, :), Pa) + ℒ.mul!(Q, LPa, view(Λ, rows, :)', one(eltype(Q)), one(eltype(Q))) + end + @inbounds for j in 1:nz, i in 1:j + m = (Q[i, j] + Q[j, i]) / 2 + Q[i, j] = m; Q[j, i] = m + end + return Q +end + +# The first-order block is autonomous in the pruned system. Its stationary +# covariance can therefore be solved separately and used to evaluate the +# state-dependent innovation covariance at the ergodic initialization without +# introducing a nonlinear covariance fixed point. +function quadratic_kalman_initial_covariance(sys, z0, g0, Λ, Pz; + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling, + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0, 0)) + A1 = Matrix(view(sys.𝒜, sys.r1, sys.r1)) + Q1 = sys.G1 * sys.G1' + Σ1 = qkf_lyapunov(A1, Q1; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + G0 = reshape(g0 + Λ * (Pz * z0), sys.nz, sys.nExo) + Tv = promote_type(eltype(G0), eltype(sys.QH), eltype(Σ1)) + Q0 = Matrix{Tv}(undef, sys.nz, sys.nz) + PzPc = Matrix{Tv}(undef, sys.nPast, sys.nr) + Pa = Matrix{Tv}(undef, sys.nPast, sys.nPast) + LPa = Matrix{Tv}(undef, sys.nz, sys.nPast) + quadratic_kalman_noise_covariance!(Q0, G0, sys.QH, Λ, sys.P, Σ1, + PzPc, Pa, LPa) + Σ0 = qkf_lyapunov(sys.𝒜, Q0; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + initial_guess = initial_guess) + return Σ0, Σ1 +end + """ The quadratic Kalman recursion, given the augmented system in the form the reverse-mode rule needs. Split out from `run_quadratic_kalman` so that the part @@ -236,10 +287,14 @@ function quadratic_kalman_recursion(𝒜, c, QH, g0, Λ, Hm, Y, 𝒞, Pz, z0, Σ Pc = Matrix{Tv}(undef, nz, nz); copyto!(Pc, Σ0) Pp = Matrix{Tv}(undef, nz, nz) Tm = Matrix{Tv}(undef, nz, nz) + Q = Matrix{Tv}(undef, nz, nz) z = Vector{Tv}(undef, nz); copyto!(z, z0) zp = Vector{Tv}(undef, nz) gv = Vector{Tv}(undef, nz * nExo) x1p = Vector{Tv}(undef, size(Pz, 1)) + PzPc = Matrix{Tv}(undef, size(Pz, 1), nz) + Pa = Matrix{Tv}(undef, size(Pz, 1), size(Pz, 1)) + LPa = Matrix{Tv}(undef, nz, size(Pz, 1)) CP = Matrix{Tv}(undef, n_obs, nz) F = Matrix{Tv}(undef, n_obs, n_obs) Kg = Matrix{Tv}(undef, nz, n_obs) @@ -254,11 +309,11 @@ function quadratic_kalman_recursion(𝒜, c, QH, g0, Λ, Hm, Y, 𝒞, Pz, z0, Σ copyto!(gv, g0); ℒ.mul!(gv, Λ, x1p, one(Tv), one(Tv)) G = reshape(gv, nz, nExo) - # Pp = 𝒜 Pc 𝒜' + G G' + QH (the rank-nExo term as one gemm update) + # Pp = 𝒜 Pc 𝒜' + E[Var(w | z)] ℒ.mul!(Tm, 𝒜, Pc) ℒ.mul!(Pp, Tm, 𝒜') - Pp .+= QH - ℒ.mul!(Pp, G, G', one(Tv), one(Tv)) + quadratic_kalman_noise_covariance!(Q, G, QH, Λ, Pz, Pc, PzPc, Pa, LPa) + Pp .+= Q for j in 1:nz, i in 1:j m = (Pp[i, j] + Pp[j, i]) / 2; Pp[i, j] = m; Pp[j, i] = m end @@ -307,14 +362,18 @@ function rrule(::typeof(quadratic_kalman_recursion), 𝒜, c, QH, g0, Λ, Hm, Y, nz::Int, nExo::Int, presample_periods::Int, on_failure_loglikelihood::Real) n_obs, nT = size(Y) zs = Vector{Vector{Float64}}(undef, nT); Ps = Vector{Matrix{Float64}}(undef, nT) - Gs = Vector{Matrix{Float64}}(undef, nT); vs = Vector{Vector{Float64}}(undef, nT) + Gs = Vector{Matrix{Float64}}(undef, nT); Pas = Vector{Matrix{Float64}}(undef, nT) + vs = Vector{Vector{Float64}}(undef, nT) CPs = Vector{Matrix{Float64}}(undef, nT); Fis = Vector{Matrix{Float64}}(undef, nT) Ks = Vector{Matrix{Float64}}(undef, nT) z = copy(z0); Pc = copy(Σ0); ll = 0.0; log2pi = log(2π); failed = false + PzPc = zeros(size(Pz, 1), nz); Pa = zeros(size(Pz, 1), size(Pz, 1)) + LPa = zeros(nz, size(Pz, 1)); Q = zeros(nz, nz) for t in 1:nT zs[t] = copy(z); Ps[t] = copy(Pc) G = reshape(g0 + Λ * (Pz * z), nz, nExo); Gs[t] = G - Q = G * G' + QH + quadratic_kalman_noise_covariance!(Q, G, QH, Λ, Pz, Pc, PzPc, Pa, LPa) + Pas[t] = copy(Pa) zp = 𝒜 * z + c Pp = 𝒜 * Pc * 𝒜' + Q; Pp = (Pp + Pp') / 2 v = Y[:, t] - 𝒞 * zp; vs[t] = v @@ -367,6 +426,16 @@ function rrule(::typeof(quadratic_kalman_recursion), 𝒜, c, QH, g0, Λ, Hm, Y, vḠ = vec(Ḡ) ḡ0 .+= vḠ Λ̄ .+= vḠ * (Pz * z_)' + P̄a = zeros(size(Pz, 1), size(Pz, 1)) + @inbounds for j in 1:nExo + rows = (j - 1) * nz + 1:j * nz + L = view(Λ, rows, :) + L̄ = 2 .* (Q̄ * L * Pas[t]) + Λ̄[rows, :] .+= L̄ + P̄a .+= L' * Q̄ * L + end + P̄a = (P̄a + P̄a') / 2 + P̄ .+= Pz' * P̄a * Pz 𝒜̄ .+= z̄p * z_' c̄ .+= z̄p z̄ = 𝒜' * z̄p + Pz' * (Λ' * vḠ) @@ -413,22 +482,20 @@ function run_quadratic_kalman(sys, end z̄ = (Matrix{Tv}(ℒ.I(nz)) - 𝒜) \ c - Gbar = quadratic_kalman_G(sys, z̄) - Q̄ = Gbar * Gbar' + QH - Q̄ = (Q̄ + Q̄') / 2 - Σ = qkf_lyapunov(𝒜, Q̄; workspaces = workspaces, - lyapunov_algorithm = lyapunov_algorithm) + # The first-order block supplies the only covariance needed by the + # state-dependent loading at the stationary initialization. + g0, Λ = quadratic_kalman_affine_G(sys) + Pz = sys.P * [Matrix{Tv}(ℒ.I(sys.nr)) zeros(Tv, sys.nr, nz - sys.nr)] + Σ, Σ1 = quadratic_kalman_initial_covariance(sys, z̄, g0, Λ, Pz; + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) # The reverse pass needs this exact matrix again. Handing it back lets the # pullback skip a second identical Lyapunov solve (a residual check on an # exact guess instead of a full doubling run). initial_covariance_out === nothing || (initial_covariance_out[] = Σ) - # Hand off to the taped recursion, which carries the hand-written adjoint. - g0, Λ = quadratic_kalman_affine_G(sys) - Pz = sys.P * [Matrix{Tv}(ℒ.I(sys.nr)) zeros(Tv, sys.nr, nz - sys.nr)] - return quadratic_kalman_recursion(𝒜, c, QH, g0, Λ, Hm, Matrix(data_in_deviations), 𝒞, Pz, z̄, Σ, nz, sys.nExo, presample_periods, on_failure_loglikelihood) @@ -509,11 +576,11 @@ function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_period Pz = qkf_Pz(sys) z0 = (Matrix{Float64}(ℒ.I(nz)) - sys.𝒜) \ sys.c + Σ0, Σ1 = quadratic_kalman_initial_covariance(sys, z0, g0, Λ, Pz; + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + initial_guess = initial_covariance) G0 = reshape(g0 + Λ * (Pz * z0), nz, nE) - Q0 = G0 * G0' + sys.QH; Q0 = (Q0 + Q0') / 2 - Σ0 = qkf_lyapunov(sys.𝒜, Q0; workspaces = workspaces, - lyapunov_algorithm = lyapunov_algorithm, - initial_guess = initial_covariance) R = last(rrule(quadratic_kalman_recursion, sys.𝒜, sys.c, sys.QH, g0, Λ, Hm, Matrix(data_in_deviations), sys.𝒞, Pz, z0, Σ0, nz, nE, @@ -530,6 +597,26 @@ function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_period Ḡ0 = 2 .* (Q̄0 * G0); Q̄H .+= Q̄0 vG0 = vec(Ḡ0); ḡ0 .+= vG0; Λ̄ .+= vG0 * (Pz * z0)' z̄0 .+= Pz' * (Λ' * vG0) + + # The stationary first-order covariance enters Q0 through the same + # state-dependent loading correction as the recursion. Differentiate its + # Lyapunov equation separately; this keeps the pullback analytical and + # avoids differentiating through a nonlinear covariance iteration. + Pa0 = sys.P * Σ1 * sys.P' + P̄a0 = zeros(nP, nP) + @inbounds for j in 1:nE + rows = (j - 1) * nz + 1:j * nz + L = view(Λ, rows, :) + Λ̄[rows, :] .+= 2 .* (Q̄0 * L * Pa0) + P̄a0 .+= L' * Q̄0 * L + end + P̄a0 = (P̄a0 + P̄a0') / 2 + Σ̄1 = sys.P' * P̄a0 * sys.P + A1 = Matrix(view(sys.𝒜, r1, r1)) + X1 = qkf_lyapunov(A1', Σ̄1; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + Ā1_initial = 2 .* (X1 * A1 * Σ1) + Ḡ1_initial = 2 .* ((X1 + X1') / 2 * sys.G1) λ = (Matrix{Float64}(ℒ.I(nz)) - sys.𝒜)' \ z̄0 c̄ .+= λ; 𝒜̄ .+= λ * z0' @@ -538,7 +625,8 @@ function quadratic_kalman_pullback(sys, data_in_deviations, Hm, presample_period P̄P = zeros(size(sys.PP)); K̄VV = zeros(size(sys.KVV)) Ā1 = 𝒜̄[r1, r1] + 𝒜̄[r2, r2] - S̄1 .+= Ā1 * sys.EaP' + S̄1 .+= (Ā1 + Ā1_initial) * sys.EaP' + S̄1 .+= Ḡ1_initial * sys.S' S̄2 .+= 𝒜̄[r2, rq] * Eq' / 2 + 𝒜̄[r2, r1] * sys.EpP' / 2 P̄P .+= Lp' * 𝒜̄[rq, rq] * Eq' + Lp' * 𝒜̄[rq, r1] * sys.EpP' S̄1 .+= c̄[r1] * sys.Ea1' diff --git a/src/get_functions.jl b/src/get_functions.jl index 4235814b2..751615901 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -357,13 +357,15 @@ And data, 4×2×40 Array{Float64, 3}: particle_resampling_threshold, particle_initial_state_scaling, particle_rng, tempering_target_ratio, tempering_mh_steps, tempering_max_stages, tempering_mh_scale)) + elseif filter == :ivashchenko_kalman + extra_kw = merge(extra_kw, (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations))) end if filter == :inversion && initial_covariance !== :theoretical @info "`initial_covariance` is not used by the inversion filter, which fixes the initial state and carries no state covariance. Ignoring input." maxlog = DEFAULT_MAXLOG end # The Kalman and particle filters take a prior on the initial state; the # inversion filter has none, so only forward it where it means something. - if filter == :kalman || filter ∈ PARTICLE_FILTERS + if filter == :kalman || filter == :ivashchenko_kalman || filter ∈ PARTICLE_FILTERS extra_kw = merge(extra_kw, (; initial_covariance)) end ensure_name_display_constants!(𝓂) @@ -530,7 +532,9 @@ And data, 1×40 Matrix{Float64}: (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations), n_particles, particle_resampling, particle_resampling_threshold, particle_initial_state_scaling, particle_rng, tempering_target_ratio, tempering_mh_steps, - tempering_max_stages, tempering_mh_scale) : NamedTuple() + tempering_max_stages, tempering_mh_scale) : + filter == :ivashchenko_kalman ? + (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations)) : NamedTuple() if filter == :inversion && initial_covariance !== :theoretical @info "`initial_covariance` is not used by the inversion filter, which fixes the initial state and carries no state covariance. Ignoring input." maxlog = DEFAULT_MAXLOG @@ -538,7 +542,7 @@ And data, 1×40 Matrix{Float64}: # The Kalman and particle filters take a prior on the initial state; the # inversion filter has none (it fixes x₀ and clamps the covariance), so the # argument is only forwarded where it means something. - if filter == :kalman || filter ∈ PARTICLE_FILTERS + if filter == :kalman || filter == :ivashchenko_kalman || filter ∈ PARTICLE_FILTERS particle_kw = merge(particle_kw, (; initial_covariance)) end @@ -688,7 +692,9 @@ And data, 4×40 Matrix{Float64}: (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations), n_particles, particle_resampling, particle_resampling_threshold, particle_initial_state_scaling, particle_rng, tempering_target_ratio, tempering_mh_steps, - tempering_max_stages, tempering_mh_scale) : NamedTuple() + tempering_max_stages, tempering_mh_scale) : + filter == :ivashchenko_kalman ? + (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations)) : NamedTuple() if filter == :inversion && initial_covariance !== :theoretical @info "`initial_covariance` is not used by the inversion filter, which fixes the initial state and carries no state covariance. Ignoring input." maxlog = DEFAULT_MAXLOG @@ -696,7 +702,7 @@ And data, 4×40 Matrix{Float64}: # The Kalman and particle filters take a prior on the initial state; the # inversion filter has none (it fixes x₀ and clamps the covariance), so the # argument is only forwarded where it means something. - if filter == :kalman || filter ∈ PARTICLE_FILTERS + if filter == :kalman || filter == :ivashchenko_kalman || filter ∈ PARTICLE_FILTERS particle_kw = merge(particle_kw, (; initial_covariance)) end @@ -955,8 +961,8 @@ And data, 4×40 Matrix{Float64}: # The inversion filter recovers the state exactly, so it has no dispersion to # report. Everything else (Kalman, particle) does. - if filter == :inversion || (algorithm != :first_order && filter ∉ PARTICLE_FILTERS && filter != :kalman) - error("`get_estimated_variable_standard_deviations` needs a filter that reports estimation uncertainty. The inversion filter identifies the state exactly and has none. Use `filter = :kalman` (first order) or one of the particle filters (`:bootstrap_particle`, `:auxiliary_particle`, `:tempered_particle`), which report the spread of the particle cloud and work at every perturbation order.") + if filter == :inversion || (algorithm != :first_order && filter ∉ PARTICLE_FILTERS && filter ∉ (:kalman, :ivashchenko_kalman)) + error("`get_estimated_variable_standard_deviations` needs a filter that reports estimation uncertainty. The inversion filter identifies the state exactly and has none. Use `filter = :kalman` or `filter = :ivashchenko_kalman` for Gaussian filters, or one of the particle filters (`:bootstrap_particle`, `:auxiliary_particle`, `:tempered_particle`).") end filter, smooth, algorithm, _, _, _ = normalize_filtering_options(filter, smooth, algorithm, false, 0) @@ -990,7 +996,9 @@ And data, 4×40 Matrix{Float64}: n_particles, particle_resampling, particle_resampling_threshold, particle_initial_state_scaling, particle_rng, tempering_target_ratio, tempering_mh_steps, - tempering_max_stages, tempering_mh_scale) : NamedTuple() + tempering_max_stages, tempering_mh_scale) : + filter == :ivashchenko_kalman ? + (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations)) : NamedTuple() if filter == :inversion && initial_covariance !== :theoretical @info "`initial_covariance` is not used by the inversion filter, which fixes the initial state and carries no state covariance. Ignoring input." maxlog = DEFAULT_MAXLOG @@ -998,7 +1006,7 @@ And data, 4×40 Matrix{Float64}: # The Kalman and particle filters take a prior on the initial state; the # inversion filter has none (it fixes x₀ and clamps the covariance), so the # argument is only forwarded where it means something. - if filter == :kalman || filter ∈ PARTICLE_FILTERS + if filter == :kalman || filter == :ivashchenko_kalman || filter ∈ PARTICLE_FILTERS particle_kw = merge(particle_kw, (; initial_covariance)) end @@ -4454,10 +4462,11 @@ end # likelihood under the same H). The Kalman *smoother* path (`filter_and_smooth`) # does not take it: the Durbin-Koopman backward recursion would need H threaded # through the disturbance smoother as well, which is not implemented. So a -# `measurement_error` supplied to the estimate entry points only has an effect for -# the particle filters. Say so rather than silently dropping it. +# `measurement_error` supplied to the estimate entry points has an effect for +# the particle and Ivashchenko Gaussian filters. Say so rather than silently +# dropping it for the other filters. function warn_unused_measurement_error(filter::Symbol, measurement_error; maxlog::Int = DEFAULT_MAXLOG) - if filter ∉ PARTICLE_FILTERS && measurement_error !== DEFAULT_MEASUREMENT_ERROR + if filter ∉ PARTICLE_FILTERS && filter != :ivashchenko_kalman && measurement_error !== DEFAULT_MEASUREMENT_ERROR @info "`measurement_error` is only used by the particle filters on this path; it is ignored for `filter = :$(filter)`. Use `get_loglikelihood` if you need measurement error in the Kalman likelihood." maxlog = maxlog end return nothing @@ -4543,9 +4552,9 @@ end """ $(SIGNATURES) -Return the loglikelihood of the model given the data and parameters provided. The loglikelihood is calculated with the filter selected by the `filter` keyword argument: the Kalman filter, the inversion filter, or one of the particle filters. By default the package selects the Kalman filter for first order solutions and the inversion filter for nonlinear (higher order) solution algorithms. The data must be provided as a `KeyedArray{Float64}` with the names of the variables to be matched in rows and the periods in columns. The `KeyedArray` type is provided by the `AxisKeys` package. +Return the loglikelihood of the model given the data and parameters provided. The loglikelihood is calculated with the filter selected by the `filter` keyword argument: the Kalman filter, inversion filter, unpruned Ivashchenko filter, or one of the particle filters. By default the package selects the Kalman filter for first order solutions and the inversion filter for nonlinear (higher order) solution algorithms. The data must be provided as a `KeyedArray{Float64}` with the names of the variables to be matched in rows and the periods in columns. The `KeyedArray` type is provided by the `AxisKeys` package. -The Kalman and inversion likelihoods are differentiable. The particle filters are stochastic Monte-Carlo estimators and are not differentiable; use them with gradient-free samplers. See the Filters section of the documentation for a comparison. +The Kalman, inversion, and Ivashchenko likelihoods are differentiable. The Ivashchenko likelihood has analytical reverse-mode rules for its unpruned second- and third-order moment recursions. The particle filters are stochastic Monte-Carlo estimators and are not differentiable; use them with gradient-free samplers. See the Filters section of the documentation for a comparison. If occasionally binding constraints are present in the model, they are not taken into account here. @@ -4889,6 +4898,28 @@ function get_loglikelihood(𝓂::ℳ, measurement_error = measurement_error_H, on_failure_loglikelihood = on_failure_loglikelihood, opts = opts) + elseif filter == :ivashchenko_kalman + # Ivashchenko's filter treats the raw perturbation solution as a + # polynomial and closes its Gaussian moments; it is separate from the + # pruned augmented-state Kalman recursions. + if has_missing + calculate_loglikelihood_with_missing(Val(:ivashchenko_kalman), Val(algorithm), obs_indices, + 𝐒, data_in_deviations, constants_obj, state, + 𝓂.workspaces, obs_idx_per_t, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + measurement_error = measurement_error_H, + on_failure_loglikelihood = on_failure_loglikelihood, + opts = opts) + else + calculate_loglikelihood(Val(:ivashchenko_kalman), Val(algorithm), obs_indices, + 𝐒, data_in_deviations, constants_obj, state, 𝓂.workspaces, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + measurement_error = measurement_error_H, + on_failure_loglikelihood = on_failure_loglikelihood, + opts = opts) + end elseif filter == :kalman if has_missing calculate_loglikelihood_with_missing(Val(:kalman), diff --git a/src/rrules.jl b/src/rrules.jl index 20553c230..41b3bb68b 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -1862,9 +1862,8 @@ function rrule(::typeof(get_loglikelihood), filter, _, algorithm, _, _, warmup_iterations = normalize_filtering_options(filter, false, algorithm, false, warmup_iterations) - # The particle filter is a stochastic, non-differentiable estimator and the - # Kalman likelihood with measurement error does not yet ship an analytical - # reverse-mode rule. Fail loudly rather than return an incorrect gradient. + # The particle filter is a stochastic, non-differentiable estimator. The + # Gaussian filters use analytical reverse-mode rules. if filter ∈ PARTICLE_FILTERS error("The particle filters (`filter = :$(filter)`) are not differentiable and cannot be used with reverse-mode automatic differentiation (Zygote/Mooncake). Use a gradient-free sampler (e.g. Pigeons slice sampling or nested sampling).") end @@ -1877,9 +1876,10 @@ function rrule(::typeof(get_loglikelihood), else measurement_error != 0 end - # The quadratic and cubic Kalman filters carry their own hand-written adjoints, - # which include the measurement-error covariance, so the guard skips them. - if me_active && filter ∉ (:quadratic_kalman, :cubic_kalman) + # The quadratic, cubic, and Ivashchenko Kalman filters carry hand-written + # adjoints, which include the measurement-error covariance, so the guard + # skips them. + if me_active && filter ∉ (:quadratic_kalman, :cubic_kalman, :ivashchenko_kalman) error("Reverse-mode automatic differentiation of the Kalman likelihood with measurement error (`measurement_error`) is not yet supported. Use forward-mode AD (e.g. `AutoForwardDiff`) or a gradient-free sampler.") end @@ -1968,7 +1968,7 @@ function rrule(::typeof(get_loglikelihood), # The quadratic and cubic Kalman filters are the ones whose inner rrules take # the measurement-error covariance; for the others it is inactive (the guard # above) and the kwarg would not be accepted. - me_kw = filter ∈ (:quadratic_kalman, :cubic_kalman) ? + me_kw = filter ∈ (:quadratic_kalman, :cubic_kalman, :ivashchenko_kalman) ? (; measurement_error = resolve_measurement_error(filter, measurement_error, data_in_deviations)) : NamedTuple() @@ -2094,6 +2094,104 @@ function rrule(::typeof(get_loglikelihood), return llh, pullback end +# Analytical reverse-mode rule for the unpruned Gaussian moment-closure +# recursion. The forward pass records the Gaussian update and each Hermite +# moment contraction; the pullback reverses those algebraic contractions and +# the fixed-point initialization directly. +function ivashchenko_likelihood_rrule(observables_index::Vector{Int}, 𝐒, + data_in_deviations::AbstractMatrix, + constants::constants, state, workspaces::workspaces, + val_algo::Val{O}; + presample_periods::Int = 0, + initial_covariance = :theoretical, + measurement_error = nothing, + on_failure_loglikelihood = -Inf, + lyapunov_algorithm::Symbol = :doubling, + obs_idx_per_t = nothing) where {O} + sys = build_ivashchenko_kalman_system_from_constants(constants, 𝐒, observables_index, O) + initial_mean = state[sys.past] + ll, tape = ivashchenko_filter_pass(sys, data_in_deviations, initial_mean; + measurement_error = measurement_error, + initial_covariance = initial_covariance, + presample_periods = presample_periods, + on_failure_loglikelihood = on_failure_loglikelihood, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + record = true) + tape === nothing && return ll, nothing + + pullback = function (cotangent) + scale = unthunk(cotangent) + if scale isa Union{NoTangent, AbstractZero} || scale == 0 + zeros_solution = [zeros(eltype(𝐒[i]), size(𝐒[i])) for i in eachindex(𝐒)] + zeros_data = zeros(eltype(data_in_deviations), size(data_in_deviations)) + zeros_state = zeros(eltype(state), length(state)) + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), + zeros_solution, zeros_data, NoTangent(), zeros_state, + NoTangent()) + end + solution_bar, data_bar, state_bar = ivashchenko_filter_pullback( + sys, tape, 𝐒, scale; + initial_covariance = initial_covariance, + lyapunov_algorithm = lyapunov_algorithm) + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), + solution_bar, data_bar, NoTangent(), state_bar, NoTangent()) + end + return ll, pullback +end + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:ivashchenko_kalman}, val_algo::Val{O}, + observables_index::Vector{Int}, 𝐒, data_in_deviations::AbstractMatrix, + constants::constants, state, workspaces::workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) where {O} + ll, pullback = ivashchenko_likelihood_rrule(observables_index, 𝐒, data_in_deviations, + constants, state, workspaces, val_algo; + presample_periods = presample_periods, + initial_covariance = initial_covariance, + measurement_error = measurement_error, + on_failure_loglikelihood = on_failure_loglikelihood, + lyapunov_algorithm = lyapunov_algorithm) + pullback === nothing && return ll, _ -> + (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), + NoTangent(), NoTangent(), NoTangent(), NoTangent()) + return ll, pullback +end + +function rrule(::typeof(calculate_loglikelihood_with_missing), + ::Val{:ivashchenko_kalman}, val_algo::Val{O}, + observables_index::Vector{Int}, 𝐒, data_in_deviations::AbstractMatrix, + constants::constants, state, workspaces::workspaces, + obs_idx_per_t::Vector{Vector{Int}}; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood = -Inf, + measurement_error = nothing, + opts::CalculationOptions = merge_calculation_options()) where {O} + ll, pullback = ivashchenko_likelihood_rrule(observables_index, 𝐒, data_in_deviations, + constants, state, workspaces, val_algo; + presample_periods = presample_periods, + initial_covariance = initial_covariance, + measurement_error = measurement_error, + on_failure_loglikelihood = on_failure_loglikelihood, + lyapunov_algorithm = lyapunov_algorithm, + obs_idx_per_t = obs_idx_per_t) + pullback === nothing && return ll, _ -> + (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), + NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + return ll, cotangent -> (pullback(cotangent)..., NoTangent()) +end + # 3-arg shim: preserves the original (no AD through initial_state) contract so # that the 3-arg @from_rrule wrapper (and any caller dispatching on the 3-arg # form) sees a pullback returning exactly 4 tangents. Delegates to the 4-arg diff --git a/test/runtests.jl b/test/runtests.jl index c2779c45e..5f2a3e451 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -72,6 +72,8 @@ elseif test_set == "jet_hot_paths" include("test_jet_hot_paths.jl") elseif test_set == "quadratic_kalman" include("test_quadratic_kalman.jl") +elseif test_set == "ivashchenko_kalman" + include("test_ivashchenko_kalman.jl") elseif test_set == "particle_filter" include("test_particle_filter.jl") include("test_particle_filter_sw07.jl") diff --git a/test/test_cubic_kalman.jl b/test/test_cubic_kalman.jl index 0a45d6d9b..f8a5b5bf9 100644 --- a/test/test_cubic_kalman.jl +++ b/test/test_cubic_kalman.jl @@ -105,6 +105,33 @@ import AxisKeys: KeyedArray 𝒜a, ca, c₀, Λ = MacroModelling.build_cubic_kalman_system(sys, basis) @test maximum(abs, 𝒜a - 𝒜) < 1e-9 @test maximum(abs, ca - c) < 1e-9 + @testset "conditional innovation covariance" begin + Λnoise = Matrix(Λ[:, sys.noise_state_indices]) + Rstate = randn(sys.nz, sys.nz); Pc = Rstate * Rstate' + Ctest = randn(sys.nz, basis.N) + Pnoise = zeros(length(sys.noise_state_indices), length(sys.noise_state_indices)) + mixvec = zeros(sys.nz * basis.N); mixΨ = zeros(sys.nz, basis.N) + CΨ = zeros(sys.nz, basis.N); Q = zeros(sys.nz, sys.nz) + MacroModelling.cubic_kalman_noise_covariance!(Q, Ctest, Λnoise, basis.Ψ, Pc, + sys.noise_state_indices, Pnoise, + mixvec, mixΨ, CΨ) + expected = Ctest * basis.Ψ * Ctest' + Pload = Pc[sys.noise_state_indices, sys.noise_state_indices] + for i in eachindex(sys.noise_state_indices), j in eachindex(sys.noise_state_indices) + Di = reshape(view(Λnoise, :, i), sys.nz, basis.N) + Dj = reshape(view(Λnoise, :, j), sys.nz, basis.N) + expected .+= Pload[i, j] .* (Di * basis.Ψ * Dj') + end + @test Q ≈ (expected + expected') / 2 + Pc_outside = copy(Pc) + outside = setdiff(1:sys.nz, sys.noise_state_indices) + Pc_outside[outside, outside] .+= 10 + Qoutside = similar(Q) + MacroModelling.cubic_kalman_noise_covariance!(Qoutside, Ctest, Λnoise, basis.Ψ, + Pc_outside, sys.noise_state_indices, + Pnoise, mixvec, mixΨ, CΨ) + @test Qoutside ≈ Q + end # Q(z) = C(z) Ψ C(z)' against the quadrature variance, at a non-trivial z Ca = reshape(c₀ + Λ * zt, sys.nz, basis.N) Qa = Ca * basis.Ψ * Ca' diff --git a/test/test_ivashchenko_kalman.jl b/test/test_ivashchenko_kalman.jl new file mode 100644 index 000000000..1f332da06 --- /dev/null +++ b/test/test_ivashchenko_kalman.jl @@ -0,0 +1,173 @@ +using MacroModelling +using Test +using Random +using ForwardDiff +using Zygote +import LinearAlgebra as ℒ +import AxisKeys: KeyedArray + +@testset "Ivashchenko unpruned Gaussian filter" begin + @model RBC_ivashchenko begin + 1 / c[0] = (β / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ)) + c[0] + k[0] = (1 - δ) * k[-1] + q[0] + q[0] = exp(z[0]) * k[-1]^α * exp(g[0]) + z[0] = ρz * z[-1] + std_z * eps_z[x] + g[0] = ρg * g[-1] + std_g * eps_g[x] + end + + @parameters RBC_ivashchenko begin + std_z = 0.02 + std_g = 0.02 + ρz = 0.4 + ρg = 0.6 + δ = 0.02 + α = 0.5 + β = 0.95 + end + + opts = MacroModelling.merge_calculation_options() + obs = [:c, :q] + data = KeyedArray(zeros(2, 10); Variable = obs, Time = 1:10) + missing_values = Matrix{Float64}(collect(data)) + missing_values[1, 3] = NaN + missing_values[:, 5] .= NaN + missing_values[:, 6] .= NaN + missing_data = KeyedArray(missing_values; Variable = obs, Time = 1:10) + + for order in (:second_order, :third_order) + MacroModelling.solve!(RBC_ivashchenko, algorithm = order, dynamics = true, opts = opts) + parameters = RBC_ivashchenko.parameter_values + _, _, solution, state, solved = MacroModelling.get_relevant_steady_state_and_state_update( + Val(order), parameters, RBC_ivashchenko, opts = opts) + @test solved + + constants = RBC_ivashchenko.constants + names = constants.post_complete_parameters.SS_and_pars_names + obs_idx = convert(Vector{Int}, indexin(obs, names)) + sys = MacroModelling.build_ivashchenko_kalman_system_from_constants( + constants, solution, obs_idx, order) + @test sys.order == order + @test size(solution[2], 2) == sys.dv^2 + if order == :third_order + @test size(sys.third_derivative, 4) == sys.dv + end + + # The closed moments of the raw polynomial map are checked against a + # direct Monte-Carlo evaluation, rather than against another filter. + Random.seed!(17 + (order == :third_order)) + mean_state = copy(state[sys.past]) + covariance_state = 0.002 .* Matrix{Float64}(ℒ.I(sys.nPast)) + scalar_type = promote_type(eltype(sys.S1), Float64) + ws = MacroModelling.ivashchenko_kalman_workspace(sys, scalar_type) + closed_mean, closed_covariance = MacroModelling.ivashchenko_polynomial_moments!( + sys, mean_state, covariance_state, ws) + closed_mean = copy(closed_mean) + closed_covariance = copy(closed_covariance) + + selected_S1 = sys.S1 + selected_S2 = Matrix(solution[2][sys.output_rows, :]) + selected_S3 = order == :third_order ? Matrix(solution[3][sys.output_rows, :]) : nothing + nmc = 120_000 + sample_mean = zeros(length(sys.output_rows)) + sample_second = zeros(length(sys.output_rows), length(sys.output_rows)) + for _ in 1:nmc + x = mean_state + ℒ.cholesky(covariance_state).L * randn(sys.nPast) + ε = randn(sys.nExo) + v = vcat(x, 1.0, ε) + value = selected_S1 * v + selected_S2 * ℒ.kron(v, v) / 2 + if selected_S3 !== nothing + value += selected_S3 * ℒ.kron(ℒ.kron(v, v), v) / 6 + end + sample_mean .+= value + sample_second .+= value * value' + end + sample_mean ./= nmc + sample_covariance = sample_second ./ nmc - sample_mean * sample_mean' + relative(a, b) = maximum(abs, a - b) / max(1e-10, maximum(abs, b)) + @test relative(sample_mean, closed_mean) < 0.08 + @test relative(sample_covariance, closed_covariance) < 0.08 + + # Public dispatch reaches the separate filter for both raw solution + # orders and uses the coupled stationary Gaussian initialization. + ll = get_loglikelihood(RBC_ivashchenko, data, parameters; + algorithm = order, + filter = :ivashchenko_kalman, + measurement_error = 1e-4) + @test isfinite(ll) + ll_diagonal = get_loglikelihood(RBC_ivashchenko, data, parameters; + algorithm = order, + filter = :ivashchenko_kalman, + initial_covariance = :diagonal, + measurement_error = 1e-4) + @test isfinite(ll_diagonal) + + # Partial observations use the observed sub-block of the innovation + # covariance; fully missing periods are prediction-only steps. + ll_missing = get_loglikelihood(RBC_ivashchenko, missing_data, parameters; + algorithm = order, + filter = :ivashchenko_kalman, + initial_covariance = :diagonal, + measurement_error = 1e-4) + @test isfinite(ll_missing) + + estimates = get_estimated_variables(RBC_ivashchenko, missing_data; + algorithm = order, + filter = :ivashchenko_kalman, + initial_covariance = :diagonal, + measurement_error = 1e-4, + levels = false, + smooth = true) + @test size(estimates) == (sys.nVars, size(data, 2)) + @test all(isfinite, collect(estimates)) + + shocks = get_estimated_shocks(RBC_ivashchenko, missing_data; + algorithm = order, + filter = :ivashchenko_kalman, + initial_covariance = :diagonal, + measurement_error = 1e-4, + smooth = true) + @test size(shocks) == (sys.nExo, size(data, 2)) + @test all(isfinite, collect(shocks)) + + standard_deviations = get_estimated_variable_standard_deviations( + RBC_ivashchenko, missing_data; + algorithm = order, + filter = :ivashchenko_kalman, + initial_covariance = :diagonal, + measurement_error = 1e-4, + smooth = true) + @test size(standard_deviations) == size(estimates) + @test all(isfinite, collect(standard_deviations)) + end + + forward_likelihood(p) = get_loglikelihood(RBC_ivashchenko, data, p; + algorithm = :second_order, + filter = :ivashchenko_kalman, + measurement_error = 1e-4) + forward_gradient = ForwardDiff.gradient(forward_likelihood, RBC_ivashchenko.parameter_values) + @test all(isfinite, forward_gradient) + + # The custom reverse rule covers both dense and missing-data paths. Use + # the explicit diagonal prior here so this test isolates the filter and + # measurement-update adjoints from the nonlinear stationary fixed point. + for order in (:second_order, :third_order) + likelihood(p) = get_loglikelihood(RBC_ivashchenko, missing_data, p; + algorithm = order, + filter = :ivashchenko_kalman, + initial_covariance = :diagonal, + measurement_error = 1e-4) + reverse_gradient = Zygote.gradient(likelihood, RBC_ivashchenko.parameter_values)[1] + forward_gradient = ForwardDiff.gradient(likelihood, RBC_ivashchenko.parameter_values) + @test all(isfinite, reverse_gradient) + @test isapprox(reverse_gradient, forward_gradient; rtol = 1e-5, atol = 1e-5) + end + + # The filter is deliberately gated away from pruned solutions: those have a + # different state-space representation and belong to the Kollmann filters. + @test get_loglikelihood(RBC_ivashchenko, data, RBC_ivashchenko.parameter_values; + algorithm = :first_order, + filter = :ivashchenko_kalman) == + get_loglikelihood(RBC_ivashchenko, data, RBC_ivashchenko.parameter_values; + algorithm = :first_order, + filter = :inversion) +end diff --git a/test/test_quadratic_kalman.jl b/test/test_quadratic_kalman.jl index 159ecb19a..4009021fb 100644 --- a/test/test_quadratic_kalman.jl +++ b/test/test_quadratic_kalman.jl @@ -7,8 +7,7 @@ import ForwardDiff import Zygote # ----------------------------------------------------------------------------- -# Quadratic Kalman filter (Monfort, Renne & Roussellet, 2015) on the pruned -# second-order solution. +# Kollmann-style quadratic Kalman filter on the pruned second-order solution. # # Three checks, in increasing strength: # @@ -65,6 +64,27 @@ import Zygote @test sys.nz == 2 * sys.nr + sys.nq @test maximum(abs, sys.S2) > 1e-3 # the model really is nonlinear + @testset "conditional innovation covariance" begin + Random.seed!(19) + G = randn(sys.nz, sys.nExo) + Λ = randn(sys.nz * sys.nExo, sys.nPast) + QH = let R = randn(sys.nz, sys.nz); R * R' end + Pz = sys.P * [Matrix{Float64}(ℒ.I(sys.nr)) zeros(sys.nr, sys.nz - sys.nr)] + Pc = let R = randn(sys.nz, sys.nz); R * R' end + PzPc = zeros(sys.nPast, sys.nz); Pa = zeros(sys.nPast, sys.nPast) + LPa = zeros(sys.nz, sys.nPast); Q = zeros(sys.nz, sys.nz) + MacroModelling.quadratic_kalman_noise_covariance!(Q, G, QH, Λ, Pz, Pc, + PzPc, Pa, LPa) + expected = G * G' + QH + Pa_expected = sys.P * Pc[1:sys.nr, 1:sys.nr] * sys.P' + for j in 1:sys.nExo + L = view(Λ, (j - 1) * sys.nz + 1:j * sys.nz, :) + expected .+= L * Pa_expected * L' + end + @test Q ≈ (expected + expected') / 2 + @test maximum(abs, Q - (G * G' + QH)) > 1e-8 + end + @testset "augmented transition reproduces the pruned conditional mean" begin Random.seed!(3) x1 = randn(sys.nr) * 0.02