diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 205e7a905..711666a85 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -177,6 +177,14 @@ jobs: os: ubuntu-latest arch: x64 test_set: "particle_filter" + - version: '1' + 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 @@ -207,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/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/docs/src/filters.md b/docs/src/filters.md index 2d4963a09..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`) @@ -360,6 +362,431 @@ 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. + +### 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(\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``) | +| 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 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 +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. +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) = G(\bar z)G(\bar z)' + Q_{\mathrm{state}} + H(I+K)H', +``` + +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** +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. + +!!! 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 +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% | +| ``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% | + +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), [*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 + +`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. + +``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. + +### 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(\bar z)\,\Psi\,C(\bar z)' + Q_{\mathrm{state}}, +``` + +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 = 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 +``\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. + +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}``. + +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(\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 +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 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 e52cf0366..cf122f1f3 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -192,6 +192,9 @@ include("./filter/find_shocks.jl") include("./filter/inversion.jl") 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! @@ -400,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 @@ -411,9 +414,28 @@ 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 + + # 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 + + 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 && !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 @@ -439,6 +461,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 in (:quadratic_kalman, :cubic_kalman) && smooth + @info "The $(filter) 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..d02559cec 100644 --- a/src/default_options.jl +++ b/src/default_options.jl @@ -13,7 +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) -const SUPPORTED_FILTERS = (:kalman, :inversion, 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. @@ -25,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). @@ -167,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 new file mode 100644 index 000000000..18184349e --- /dev/null +++ b/src/filter/cubic_kalman.jl @@ -0,0 +1,1312 @@ +@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 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. +# 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 +# refuse with a message that names the cause instead of appearing to hang. +const CUBIC_KALMAN_MAX_DIMENSION = 2500 + +# 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 +# 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 + +# ── 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) + 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 + +""" +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 + +""" +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!(βˆ‚, 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) + + # 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 +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 + # 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 + $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.") + end + + # 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) + 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 + + 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. + 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] + + # 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, noise_state_indices, + k2cols, k2_ij, S2k2, k12cols, k12_ij, S2k12, k3cols, k3_ijk, S3k3) +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, Tv = eltype(sys.S1)) + (; 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), + aug1 = zeros(na), aug1h = zeros(na), aug2 = zeros(na), aug3 = zeros(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), + 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), + scratch_out = zeros(sys.nz)) +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!(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, 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 + @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 (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 + q111[((i-1)*nP + (j-1))*nP + k] + elseif n == 2 + if ci && cj + q11[(i-1)*nP+j] * aug1[k] + elseif ci && ck + q11[(i-1)*nP+k] * aug1[j] + else + q11[(j-1)*nP+k] * aug1[i] + end + else + aug1[i] * aug1[j] * aug1[k] + end + end + + β„’.mul!(x1n, S1, aug1) + β„’.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 + 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 + +""" +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. + # The 𝐒₂/𝐒₃ cotangents accumulate on the live columns and are scattered back + # once, in `cubic_derived_pullback!`. + βˆ‚.S1 .+= βˆ‚x1n * aug1' .+ βˆ‚x2n * aug2' .+ βˆ‚x3n * aug3' + βˆ‚.S2k2 .+= (βˆ‚x2n * K2') ./ 2 + βˆ‚.S2k12 .+= βˆ‚x3n * K12' + βˆ‚.S3k3 .+= (βˆ‚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)), + 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. +function cubic_kalman_step(sys, z::AbstractVector, Ξ΅::AbstractVector, ws = cubic_kalman_workspace(sys)) + 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. +# +# 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{eltype(sys.S1)}(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; 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{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) + 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 + π’œ[:, j] = mean_at(e) - c + end + 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, 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) + cubic_kalman_step!(view(Fm, :, p), sys, z, Ξ΅, ws) + end + return Fm * basis.W # nz Γ— N + end + + C0 = coefficients_at(zeros(Tv, nz)) + cβ‚€ = vec(C0) + Ξ› = Matrix{Tv}(undef, nz * N, nz) + π’œ = Matrix{Tv}(undef, nz, nz) + c = C0 * basis.m + e = zeros(Tv, nz) + for j in 1:nz + fill!(e, zero(Tv)) + e[j] = one(Tv) + Ξ”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 + +# 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)`, +`π’œ[:,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. 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 + β„’.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!(βˆ‚, sys) + 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, +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, + workspaces = nothing, + lyapunov_algorithm::Symbol = :doubling) + nz = sys.nz + 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(eltype(sys.S1), eltype(data_in_deviations), + measurement_error === nothing ? Float64 : eltype(measurement_error)) + + Hm = if measurement_error === nothing + zeros(Tv, n_obs, n_obs) + elseif measurement_error isa AbstractMatrix + Matrix{Tv}(measurement_error) + else + Matrix{Tv}(β„’.Diagonal(collect(measurement_error))) + end + + 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.Ξ¨ + 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, Pc) + copyto!(cvec, cβ‚€) + β„’.mul!(cvec, Ξ›, z, one(Tv), one(Tv)) + C = reshape(cvec, nz, N) + cubic_kalman_noise_covariance!(Q, C, Ξ›noise, Ξ¨, Pc, noise_state_indices, + Pnoise, mixvec, mixΞ¨, CΞ¨) + end + + z = (Matrix{Tv}(β„’.I(nz)) - π’œ) \ c + 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 + # temporary per period competes with it directly. + op1, op2, op3 = sys.op1, sys.op2, sys.op3 + 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, Pc) + + # 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 + + β„’.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 + copyto!(Fv, v); β„’.ldiv!(Fc, Fv) + ll -= 0.5 * (β„’.dot(v, Fv) + β„’.logdet(Fc) + n_obs * log2pi) + isfinite(ll) || return on_failure_loglikelihood + end + + # 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 + end + end + return ll +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(\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, + 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}}(); 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) + 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 + 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, Pas, vs, CPs, Fis, Ks) +end + +function cubic_kalman_recursion_pullback(tape, π’œ, c, cβ‚€, Ξ›, Ξ¨, π’ž, nz, N, n_obs, nT, + 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) + + 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_' + 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Μ„ + end + return π’œΜ„, cΜ„, cΜ„β‚€, Ξ›Μ„, HΜ„m, Θ², zΜ„, PΜ„ +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 + +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 + noise_state_indices = sys.noise_state_indices + Ξ›noise = Matrix(Ξ›[:, noise_state_indices]) + + zβ‚€ = (Matrix{Float64}(β„’.I(nz)) - π’œ) \ c + Cβ‚€ = reshape(cβ‚€ + Ξ› * zβ‚€, nz, N) + Ξ£β‚€ = 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, + noise_state_indices) + + 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, 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β‚€ * Ξ¨) + vCΜ„β‚€ = vec(CΜ„β‚€) + 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Μ„β‚€ + 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/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 new file mode 100644 index 000000000..47a5f91ad --- /dev/null +++ b/src/filter/quadratic_kalman.jl @@ -0,0 +1,762 @@ +@stable default_mode = "disable" begin + +# 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 +# & 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. 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 +# 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) + @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), …)`. +""" +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 + + 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. + # 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 (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, 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, nPast^2 + i] = 1.0 + elseif i == nPast + 1 && j <= nPast + Eaa[r, nPast^2 + j] = 1.0 + elseif i == nPast + 1 && j == nPast + 1 + Eaa[r, nPast^2 + nPast + 1] = 1.0 + end + end + 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 + SS = β„’.kron(S, S) + vecI = vec(Matrix{Float64}(β„’.I(nExo))) + PP = β„’.kron(PS1, PS1) + EaP = Ea[:, 1:nPast] * P + A1 = S1 * EaP + + r1, r2, rq = 1:nr, nr+1:2nr, 2nr+1: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 + π’œ[r2, r1] = S2 * Ep * P / 2; c[r2] = S2 * (E1 + SS * vecI) / 2 + π’œ[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, 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, 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 + + # 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, + 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. +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 = 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 + +# Conditional covariance of the quadratic innovation. If G(z) = GΜ„ + Ξ£α΅’ zα΅’Gα΅’ +# and z has covariance P, the state-dependent part contributes +# +# Ξ£α΅’β±Ό Pα΅’β±Ό Gα΅’ Gβ±Ό' +# +# in addition to the plug-in term GΜ„GΜ„'. 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 +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) + # 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) + 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) + v = Vector{Tv}(undef, n_obs) + Fv = Vector{Tv}(undef, n_obs) + + ll = zero(Tv); log2pi = log(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 π’œ' + E[Var(w | z)] + β„’.mul!(Tm, π’œ, Pc) + β„’.mul!(Pp, Tm, π’œ') + 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 + + β„’.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 + 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' 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 + +# 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); 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 + 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 + 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_)' + 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αΈ ) + 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 +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{<:Real}; + measurement_error::Union{Nothing,AbstractVector{<:Real},AbstractMatrix{<:Real}} = nothing, + presample_periods::Int = 0, + 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) + 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(Tv, n_obs, n_obs) + elseif measurement_error isa AbstractMatrix + Matrix{Tv}(measurement_error) + else + Matrix{Tv}(β„’.Diagonal(collect(measurement_error))) + end + + zΜ„ = (Matrix{Tv}(β„’.I(nz)) - π’œ) \ c + + # 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[] = Ξ£) + + return quadratic_kalman_recursion(π’œ, c, QH, g0, Ξ›, Hm, Matrix(data_in_deviations), + π’ž, Pz, zΜ„, Ξ£, nz, sys.nExo, + 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. +# +# 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 + 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; + 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) + g0, Ξ› = quadratic_kalman_affine_G(sys) + 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) + + 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; workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) + π’œΜ„ .+= 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) + + # 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) + AΜ„1_initial = 2 .* (X1 * A1 * Ξ£1) + GΜ„1_initial = 2 .* ((X1 + X1') / 2 * sys.G1) + Ξ» = (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 + AΜ„1_initial) * sys.EaP' + SΜ„1 .+= GΜ„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' + 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, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm) +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))) + Ξ£β‚€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, + workspaces = workspaces, + lyapunov_algorithm = lyapunov_algorithm, + initial_covariance_out = Ξ£β‚€ref) + + 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; + 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 + βˆ‚π’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 6d4d3ee4e..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. @@ -4865,6 +4874,52 @@ 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 + 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 == :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 == :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 7366edfa3..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,7 +1876,10 @@ function rrule(::typeof(get_loglikelihood), else measurement_error != 0 end - if me_active + # 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 @@ -1963,6 +1965,13 @@ function rrule(::typeof(get_loglikelihood), end # ── step 3: calculate_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, :ivashchenko_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 @@ -2083,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 0d58ed59c..5f2a3e451 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -70,6 +70,10 @@ 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 == "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 new file mode 100644 index 000000000..f8a5b5bf9 --- /dev/null +++ b/test/test_cubic_kalman.jl @@ -0,0 +1,253 @@ +using MacroModelling +using Test +using Random +using ForwardDiff +using FiniteDifferences +using Zygote +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) + 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) + + # 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 + # 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)[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) + 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 + + # 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 + @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' + @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) + 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 + + # 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_fd = FiniteDifferences.grad(central_fdm(5, 1), f, pars)[1] + @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) + 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) + # 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!(βˆ‚, sys) + 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 + # 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 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 new file mode 100644 index 000000000..4009021fb --- /dev/null +++ b/test/test_quadratic_kalman.jl @@ -0,0 +1,242 @@ +using MacroModelling +using Test +import Random +import Statistics +import LinearAlgebra as β„’ +import ForwardDiff +import Zygote + +# ----------------------------------------------------------------------------- +# Kollmann-style quadratic Kalman filter 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) + # 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 "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 + 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.nr); a2 = zeros(sys.nr); aq = zeros(sys.nPast^2) + for _ in 1:nmc + Ξ΅ = randn(sys.nExo) + # 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.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 + # 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 "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 + + # 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 + # 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. + @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