diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 000000000..143181918 --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,23 @@ +FROM julia:latest + +# System dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + git curl sudo ca-certificates python3 python3-pip \ + && rm -rf /var/lib/apt/lists/* + +# Non-root user (Codespaces convention) +RUN groupadd -g 1000 vscode \ + && useradd -m -u 1000 -g vscode -s /bin/bash vscode \ + && echo "vscode ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/vscode + +USER vscode +WORKDIR /home/vscode + +# Pre-populate Julia depot with MacroModelling (main branch) deps +RUN git clone --depth 1 https://github.com/thorek1/MacroModelling.jl.git /tmp/MacroModelling \ + && cd /tmp/MacroModelling \ + && julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' \ + && rm -rf /tmp/MacroModelling + +# Pre-compile LanguageServer for the Julia VS Code extension +RUN julia -e 'using Pkg; Pkg.add("LanguageServer"); using LanguageServer' diff --git a/.devcontainer/add-notice.sh b/.devcontainer/add-notice.sh deleted file mode 100644 index c292bc543..000000000 --- a/.devcontainer/add-notice.sh +++ /dev/null @@ -1,19 +0,0 @@ -# Display a notice when not running in GitHub Codespaces - -cat << 'EOF' > /usr/local/etc/vscode-dev-containers/conda-notice.txt -When using "conda" from outside of GitHub Codespaces, note the Anaconda repository -contains restrictions on commercial use that may impact certain organizations. See -https://aka.ms/vscode-remote/conda/miniconda - -EOF - -notice_script="$(cat << 'EOF' -if [ -t 1 ] && [ "${IGNORE_NOTICE}" != "true" ] && [ "${TERM_PROGRAM}" = "vscode" ] && [ "${CODESPACES}" != "true" ] && [ ! -f "$HOME/.config/vscode-dev-containers/conda-notice-already-displayed" ]; then - cat "/usr/local/etc/vscode-dev-containers/conda-notice.txt" - mkdir -p "$HOME/.config/vscode-dev-containers" - ((sleep 10s; touch "$HOME/.config/vscode-dev-containers/conda-notice-already-displayed") &) -fi -EOF -)" - -echo "${notice_script}" | tee -a /etc/bash.bashrc >> /etc/zsh/zshrc diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index e52e49449..981f9b59b 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -1,24 +1,18 @@ -// See https://github.com/julia-vscode/julia-devcontainer/blob/master/Dockerfile for image contents { - "name": "Julia (Community)", - "image": "ghcr.io/julia-vscode/julia-devcontainer:latest", - // "image": "tkockerols/julia-development:latest", - - // Configure tool-specific properties. - "customizations": { - // Configure properties specific to VS Code. - "vscode": { - // Add the IDs of extensions you want installed when the container is created. - "extensions": [ - "julialang.language-julia", - "rid9.datetime", - "mutantdino.resourcemonitor", - "bungcip.better-toml" - ] - } - }, - "onCreateCommand": "julia -e 'import Pkg; Pkg.add(\"SymPy\"); using SymPy'", - "postCreateCommand": "/julia-devcontainer-scripts/postcreate.jl", - - "remoteUser": "vscode" + "name": "MacroModelling.jl", + "build": { + "dockerfile": "Dockerfile" + }, + "customizations": { + "vscode": { + "extensions": [ + "julialang.language-julia" + ], + "settings": { + "julia.executablePath": "/usr/local/julia/bin/julia" + } + } + }, + "onCreateCommand": "julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()'", + "remoteUser": "vscode" } diff --git a/.devcontainer/noop.txt b/.devcontainer/noop.txt deleted file mode 100644 index abee19541..000000000 --- a/.devcontainer/noop.txt +++ /dev/null @@ -1,3 +0,0 @@ -This file is copied into the container along with environment.yml* from the -parent folder. This is done to prevent the Dockerfile COPY instruction from -failing if no environment.yml is found. \ No newline at end of file diff --git a/.github/STYLE_GUIDE.md b/.github/STYLE_GUIDE.md new file mode 100644 index 000000000..4d964ad0f --- /dev/null +++ b/.github/STYLE_GUIDE.md @@ -0,0 +1,859 @@ +# Code Style Guide for MacroModelling.jl + +This document describes the coding conventions and style rules used throughout the MacroModelling.jl codebase. +All new code should follow these guidelines to maintain consistency. + +--- + +## Table of Contents + +1. [Naming Conventions](#naming-conventions) +2. [Formatting and Indentation](#formatting-and-indentation) +3. [Function Signatures](#function-signatures) +4. [Type System](#type-system) +5. [Module Organisation](#module-organisation) +6. [Control Flow](#control-flow) +7. [Error Handling](#error-handling) +8. [Documentation](#documentation) +9. [Performance](#performance) +10. [Collections and Arrays](#collections-and-arrays) +11. [Strings and Symbols](#strings-and-symbols) +12. [Logging and Verbosity](#logging-and-verbosity) +13. [Caching](#caching) +14. [Macros](#macros) + +--- + +## Naming Conventions + +### Functions + +Use **snake_case** for all function names: + +```julia +calculate_first_order_solution(...) +get_shock_decomposition(...) +solve_quadratic_matrix_equation(...) +``` + +Mutating functions must end with `!` per Julia convention: + +```julia +solve!(๐“‚, ...) +fast_lu!(ws, A) +ensure_lyapunov_doubling_buffers!(ws, n) +``` + +### Variables + +Use **snake_case** for multi-word variable names: + +```julia +past_not_future_and_mixed_idx +non_stochastic_steady_state +``` + +Use **Unicode mathematical symbols** for domain-specific variables to match the underlying mathematics: + +```julia +๐“‚ # model object +โˆ‡โ‚ # Jacobian +โˆ‡โ‚‚ # Hessian +๐’โ‚ # first-order solution matrix +๐’โ‚‚ # second-order solution matrix +ฯต # epsilon / shocks +ฮฃสธโ‚ # covariance matrix +``` + +Use **Unicode subscripts and superscripts** for order indices: + +```julia +nโ‚‘ # number of exogenous variables +nโ‚‹ # number of past variables +nโ‚Š # number of future variables +iโ‚Š # future indices +iโ‚‹ # past indices +``` + +Prefix counts with `n`: + +```julia +nVars +nExo +nPresent_only +nMixed +``` + +### Types and Structs + +Use **snake_case** for workspace and internal structs: + +```julia +struct second_order_indices ... end +mutable struct qme_workspace{T} ... end +mutable struct sylvester_workspace{G,H} ... end +``` + +### Constants + +Use **SCREAMING_SNAKE_CASE** for constants: + +```julia +const DEFAULT_ALGORITHM = :first_order +const DEFAULT_VERBOSE = false +const ANALYTICAL_STEP = 1 +const NUMERICAL_STEP = 2 +``` + +Docstring template constants use a `ยฎ` suffix: + +```julia +const MODELยฎ = "..." +const ALGORITHMยฎ = "..." +const VERBOSEยฎ = "..." +``` + +### Module Aliases + +Import libraries with **Unicode letter aliases**: + +```julia +import LinearAlgebra as โ„’ +import LinearSolve as ๐’ฎ +import ForwardDiff as โ„ฑ +import DifferentiationInterface as ๐’Ÿ +``` + +### Type Aliases + +Define union types for user-facing inputs: + +```julia +const Symbol_input = Union{Symbol, Vector{Symbol}, ...} +const ParameterType = Union{Nothing, Pair{Symbol, Float64}, ...} +``` + +--- + +## Formatting and Indentation + +### Indentation + +Use **4 spaces** for indentation. Never use tabs. + +```julia +function foo(x) + if x > 0 + return x + else + return -x + end +end +``` + +### Line Length + +There is no strict line-length limit. Long lines (200+ characters) are acceptable for complex mathematical expressions and function signatures. Prefer readability over arbitrary wrapping. + +### Whitespace + +Spaces around binary operators: + +```julia +nโ‚‹ + 1 + nโ‚‘ +A * X * B + C +x == nothing +``` + +No space before `(` in function calls: + +```julia +zeros(T, n, n) +size(A, 1) +push!(vec, val) +``` + +Space after commas: + +```julia +zeros(T, n, n) +solve!(๐“‚, parameters = parameters, verbose = verbose) +``` + +### Blank Lines + +No blank lines between closely related one-liner function definitions: + +```julia +get_symbols(ex::Symbol) = [ex] +get_symbols(ex::Real) = [ex] +get_symbols(ex::Int) = [ex] +``` + +Two or more blank lines between major function definitions to visually separate sections. + +### Section Headers + +Use comment banners to delineate major sections within a file: + +```julia +# ========================================================================= +# AUXILIARY MATRICES (for perturbation solution) +# ========================================================================= +``` + +### Keyword Argument Alignment + +Align keyword arguments vertically, each on its own line, indented to the opening parenthesis: + +```julia +function get_shock_decomposition(๐“‚::โ„ณ, + data::KeyedArray{Float64}; + parameters::ParameterType = nothing, + algorithm::Symbol = DEFAULT_ALGORITHM, + verbose::Bool = DEFAULT_VERBOSE) +``` + +--- + +## Function Signatures + +### Type Annotations + +Annotate return types on public-facing functions: + +```julia +function get_equations(๐“‚::โ„ณ)::Vector{String} + ... +end +``` + +Use parametric `where` clauses to constrain type parameters: + +```julia +function solve!(A::AbstractMatrix{T}, + B::AbstractMatrix{T}) where {T <: AbstractFloat} + ... +end +``` + +### Keyword Arguments + +Separate keyword arguments with `;`. Every keyword argument should have a default value, preferably drawn from `DEFAULT_*` constants: + +```julia +function get_irf(๐“‚::โ„ณ; + parameters::ParameterType = nothing, + algorithm::Symbol = DEFAULT_ALGORITHM, + verbose::Bool = DEFAULT_VERBOSE, + tol::Tolerances = Tolerances()) +``` + +### Short Functions + +Write simple functions as one-liners: + +```julia +get_symbols(ex::Symbol) = [ex] +noop_state_update(::Float64, ::Float64) = nothing +``` + +### Multiple Dispatch + +Use `Val` dispatch for compile-time-known mode selection: + +```julia +filter_data_with_model(๐“‚, data, Val(algorithm), Val(filter), ...) +``` + +Use type dispatch for workspace variants: + +```julia +fast_lu!(A::AbstractMatrix{T}) where T = ... +fast_lu!(ws::LUWorkspace, A::AbstractMatrix{T}) where T = ... +``` + +--- + +## Type System + +### Struct Definitions + +Explicitly type all struct fields: + +```julia +mutable struct qme_workspace{T <: Real, R <: Real} + A::Matrix{T} + B::Matrix{T} + solved::Bool + n::Int +end +``` + +Use `mutable struct` for workspaces and caches that change over time. +Use `struct` for immutable configuration objects. + +### Parametric Types + +Constrain type parameters to `Real`, `AbstractFloat`, or `Number` as appropriate: + +```julia +mutable struct sylvester_workspace{G <: AbstractFloat, H <: Real} + ... +end +``` + +--- + +## Module Organisation + +### Import Order + +In the main module file, follow this order: + +1. `module` declaration +2. `import` statements with Unicode aliases +3. `using` statements (only for packages that should export into scope) +4. Inline utility function definitions +5. Type aliases +6. `include` of source files (in dependency order) +7. `export` statements (grouped by functionality) +8. AD rule includes (at the very end) +9. `end` (module close) + +### `import` vs `using` + +**Prefer `import` over `using`** to keep the namespace clean: + +```julia +# Preferred +import LinearAlgebra as โ„’ +import SparseArrays: SparseMatrixCSC, sparse!, spzeros + +# Only for packages that must export into scope +using PrecompileTools +using DispatchDoctor +``` + +### Include Order + +Include files in dependency order โ€” structures before functions that use them: + +```julia +include("default_options.jl") +include("common_docstrings.jl") +include("structures.jl") +include("solver_parameters.jl") +include("options_and_caches.jl") +include("nsss_solver.jl") +include("macros.jl") +include("get_functions.jl") +# ...subdirectories +include("./algorithms/sylvester.jl") +include("./filter/kalman.jl") +``` + +### Exports + +Provide multiple aliases for discoverability: + +```julia +export get_steady_state, get_SS, get_ss, + get_non_stochastic_steady_state, + steady_state, SS, SSS, ss, sss +``` + +--- + +## Control Flow + +### Short-Circuit Returns + +Use short-circuit for early returns: + +```julia +if !solved return zeros(T, n, n), sol, false end +``` + +### Ternary Operator + +Use ternary for simple inline conditionals: + +```julia +verbose ? println("Solving...") : nothing +filter == :kalman ? :kalman : :inversion +``` + +### Inline `if` + +Use single-line `if` for simple branches: + +```julia +if opts.verbose println("Quadratic matrix equation solution failed.") end +if solved ๐“‚.caches.qme_solution = qme_sol end +``` + +### `@assert` for Preconditions + +```julia +@assert algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] "Theoretical mean available only for..." +``` + +### `for` Loops + +Standard range iteration: + +```julia +for i in 1:n + ... +end +``` + +Reverse iteration with step: + +```julia +for n in length(eqs_to_solve)-1:-1:2 + ... +end +``` + +Destructuring with `enumerate`: + +```julia +for (i, x) in enumerate(aux_vars) + ... +end +``` + +### `do` Blocks + +Use `do` blocks with `postwalk`/`prewalk` for AST manipulation: + +```julia +postwalk(expr) do x + if x isa Expr && x.head == :(=) + found = true + end + return x +end +``` + +Use `do` blocks with `open` for file I/O: + +```julia +open(filepath, "w") do io + println(io, content) +end +``` + +### `try/catch` + +For cases where failure is expected and should be silently handled, use compact `try/catch`: + +```julia +result = try SPyPyC.solve(equation, variable) + catch + end +``` + +For user-facing errors, re-raise with context: + +```julia +try + run(pipeline(...)) +catch + error("Failed to parse the model. ...") +end +``` + +--- + +## Error Handling + +### Exceptions + +Use `throw(ArgumentError(...))` for invalid arguments: + +```julia +throw(ArgumentError("invalid argument to LU factorization, info = $info")) +``` + +### Boolean Solved Flags + +Return `(result, solved::Bool)` from solver functions rather than throwing. Callers check the flag: + +```julia +sol, solved = calculate_first_order_solution(...) +if !solved + return zeros(...), sol, false +end +``` + +### Warnings + +Use `@warn` for non-fatal issues: + +```julia +@warn "Invalid option `$(x.args[1])` ignored..." +``` + +Use `@info` with `maxlog` for informational messages that should not repeat: + +```julia +@info "Higher order solution algorithms only support the inversion filter." maxlog = maxlog +``` + +--- + +## Documentation + +### Docstrings + +Use `$(SIGNATURES)` from DocStringExtensions for auto-generated signatures. + +Structure docstrings with these sections: + +```julia +""" +$(SIGNATURES) + +Short description of the function. + +# Arguments +- `arg1`: description + +# Keyword Arguments +- `kwarg1` [default: `value`]: description +$MODELยฎ +$ALGORITHMยฎ +$VERBOSEยฎ + +# Returns +- Description of return value + +# Examples +```jldoctest +using MacroModelling + +@model RBC begin + ... +end + +@parameters RBC begin + ... +end + +get_equations(RBC) +# output +... +``` +""" +``` + +### Shared Docstring Constants + +Define reusable docstring fragments as constants with the `ยฎ` suffix and reference them with `$`: + +```julia +const MODELยฎ = """ +- `๐“‚`: the model object +""" + +# In docstring: +""" +# Arguments +\$MODELยฎ +""" +``` + +### Comments + +Use inline comments to explain non-obvious fields and logic: + +```julia +A::Matrix{T} # nร—n copy of A +solved::Bool # whether QME converged +``` + +Preserve commented-out alternative approaches for reference. + +### Writing Style + +- Avoid second-person phrasing ("you") in documentation and docstrings +- Use third person or imperative mood + +--- + +## Performance + +### `@inline` + +Apply `@inline` to hot-path utility functions: + +```julia +@inline function fast_lu!(ws, A::AbstractMatrix{T}) where T + ... +end +``` + +### `@views` + +Use `@views` to avoid array copies: + +```julia +@views sol[:, 1:T.nPast_not_future_and_mixed] +@views [๐’โ‚[iโ‚Š,:]; ...] +``` + +### Pre-allocation and Workspaces + +All major solvers use pre-allocated workspace structs. Use `ensure_*_buffers!` functions that lazily resize workspaces only when dimensions change: + +```julia +function ensure_lyapunov_doubling_buffers!(ws::lyapunov_workspace, n::Int) + if size(ws.A, 1) != n + ws.A = zeros(n, n) + # ...resize all buffers... + end +end +``` + +### Type Stability + +- Annotate return types on functions +- Use parametric `where` clauses +- Avoid untyped containers in hot paths + +### Sparse Matrices + +Use `choose_matrix_format` to decide dense vs sparse based on density thresholds. +Clean up near-zero entries with `droptol!`. + +### `@ignore_derivatives` + +Use `ChainRulesCore.@ignore_derivatives` for code that should be invisible to AD: + +```julia +@ignore_derivatives begin + # cache updates, logging, etc. +end +``` + +--- + +## Collections and Arrays + +### Broadcasting + +Prefer dot syntax for element-wise operations: + +```julia +data .- NSSS[obs_idx] +obs_axis .|> Meta.parse .|> replace_indices +solved_vals .= new_values +``` + +### Comprehensions + +Use array comprehensions for constructing new arrays: + +```julia +[replace_curly_braces_in_symbols(arg) for arg in expr.args] +``` + +Use generator expressions inside aggregation functions: + +```julia +sum(k * (k + 1) รท 2 for k in 1:n) +``` + +### Pipe Operator + +Use `|>` for chaining transformations: + +```julia +parse_variables_input_to_index(obs_symbols, ๐“‚) |> sort +collect(โˆ‚block) |> findnz +``` + +### `Ref` for Broadcasting Scalars + +Wrap non-collection arguments in `Ref` when broadcasting: + +```julia +replace_symbols.(expressions, Ref(parameter_dict)) +Symbolics.substitute.(x, Ref(back_to_array_dict)) +``` + +### `push!` and `append!` + +Use `push!` for single elements, `append!` for extending with another collection: + +```julia +push!(b.step_types, ANALYTICAL_STEP) +append!(b.write_indices, write_indices) +``` + +--- + +## Strings and Symbols + +### Interpolation + +Use `$` for string interpolation: + +```julia +"invalid argument, info = $info" +``` + +### Concatenation + +Use `*` for string concatenation (Julia convention): + +```julia +string(x.args[1]) * "โ‚โ‚“โ‚Ž" +string(x.args[1]) * "แดธโฝ" * super(string(abs(k - 1))) * "โพโ‚โ‚€โ‚Ž" +``` + +### Regex + +Use `r"..."` literals, with flags as needed: + +```julia +occursin(r"^(x|ex|exo|exogenous){1}$"i, input) +``` + +### `replace` Chains + +Chain `replace` calls for multiple substitutions: + +```julia +replace(replace(replace(str, "โ‚โ‚‹โ‚โ‚Ž" => "[-1]"), "โ‚โ‚โ‚Ž" => "[1]"), "โ‚โ‚€โ‚Ž" => "[0]") +``` + +--- + +## Logging and Verbosity + +### `verbose::Bool` + +Controls solver-internal diagnostics via `println`: + +```julia +if opts.verbose println("Quadratic matrix equation solution failed.") end +``` + +### `silent::Bool` + +Controls progress printing for user-facing operations: + +```julia +if !silent print("Set up non-stochastic steady state problem:\t\t\t\t") end +# ...computation... +if !silent println(round(time() - start_time, digits = 3), " seconds") end +``` + +### `@info` / `@warn` + +Use `@info` with `maxlog` for corrections that should not repeat endlessly: + +```julia +@info "Setting filter = :inversion for higher order solution." maxlog = maxlog +``` + +Use `@warn` for non-fatal warnings: + +```julia +@warn "Solution does not have a stochastic steady state." +``` + +--- + +## Caching + +### Pattern + +Use a dedicated `caches` sub-struct with a parallel `outdated` flags struct: + +```julia +๐“‚.caches.non_stochastic_steady_state = SS_and_pars +๐“‚.caches.outdated.non_stochastic_steady_state = solution_error > tol +``` + +### Check โ†’ Recompute โ†’ Store โ†’ Clear + +```julia +if ๐“‚.caches.outdated.second_order_solution || parameters_changed + # ...recompute... + ๐“‚.caches.second_order_stochastic_steady_state = result + ๐“‚.functions.second_order_state_update = state_updateโ‚‚ + ๐“‚.caches.outdated.second_order_solution = false +end +``` + +### Lazy Allocation + +Compute constant values lazily on first use and store in the model struct cache. Subsequent calls must read from the cache. + +--- + +## Macros + +### `@model` and `@parameters` + +User-facing macros use `begin...end` blocks: + +```julia +@model RBC 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]^ฮฑ + z[0] = ฯ * z[-1] + std_z * eps_z[x] +end +``` + +### `@stable` Wrapper + +Wrap groups of functions in `@stable default_mode = "disable" begin...end` from DispatchDoctor: + +```julia +@stable default_mode = "disable" begin + +function calculate_first_order_solution(...) + ... +end + +function calculate_second_order_solution(...) + ... +end + +end # dispatch_doctor +``` + +### AST Manipulation + +Use `postwalk`/`prewalk` from MacroTools for expression tree traversal in macro implementations: + +```julia +postwalk(expr) do x + if x isa Expr && x.head == :ref + # transform variable references + end + return x +end +``` + +--- + +## Summary of Key Principles + +1. **snake_case everywhere** โ€” functions, variables, most struct names +2. **Unicode for mathematics** โ€” match the notation from the underlying papers +3. **`import` over `using`** โ€” keep the namespace clean +4. **Explicit types** โ€” annotate struct fields, return types, and `where` clauses +5. **Pre-allocate workspaces** โ€” avoid allocations in hot loops +6. **Boolean solved flags** โ€” return `(result, solved)` rather than throwing from solvers +7. **Verbose/silent kwargs** โ€” let callers control output +8. **Shared docstring constants** โ€” avoid repeating common parameter documentation +9. **No strict line limit** โ€” readability over wrapping for mathematical code +10. **`@views`, `@inline`, `Ref`** โ€” standard Julia performance patterns diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 238784142..5feb93fe1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -171,8 +171,8 @@ jobs: Project.toml rm -f Project.toml.bak - - name: Remove JET from estimation runs - if: contains(matrix.test_set, 'estimat') == true || contains(matrix.test_set, 'pigeons') == true + - name: Remove JET from non-jet runs + if: matrix.test_set != 'jet' shell: bash run: | sed -i.bak \ diff --git a/.gitignore b/.gitignore index 5141f731c..794d64fee 100644 --- a/.gitignore +++ b/.gitignore @@ -76,4 +76,6 @@ test/data/EA_data.csv test/data/SSR_Estimates_20241130.xlsx test/data/TED---Output-Labor-and-Labor-Productivity-1950-2015.xlsx estimation_results -juliaup.sh +juliaup.sh.julia_repl/ +tasks/_repl_cmd.jl +.julia_repl diff --git a/AGENTS.md b/AGENTS.md index db0fcca89..9d6d1d06b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,340 +1,54 @@ # Agent Guide for MacroModelling.jl -This file provides guidance for AI coding agents (GitHub Copilot, Claude, etc.) when working with this repository. +This file is the concise default guide for AI coding agents (GitHub Copilot, Claude, etc.). +Read this file first. Read the companion files only when needed. -## Project Overview +## Mandatory Workflow (Always Follow) -`MacroModelling.jl` is a Julia package for developing and solving dynamic stochastic general equilibrium (DSGE) models. These models describe macroeconomic behavior and are used for counterfactual analysis, economic policy evaluation, and quantifying specific mechanisms in academic research. +1. **Read session context first:** At session start, read `AGENT_PROGRESS.md` before making changes. +2. **Start with a minimal targeted script/test:** For new features or bug fixes, first create/run a minimal script or focused test that reproduces the exact error or validates the feature's correctness before editing code. +3. **Use plan mode for non-trivial work:** If a task has 3+ steps or architecture decisions, write and maintain a clear plan. + +4. **Fix root cause when addressing errors:** Do not stop at symptom-level patches when a deeper cause can be identified and corrected. +5. **Prove changes by testing:** Never claim success without running a relevant test/check. For bug fixes and new features, accept code changes only if the initial minimal script/test passes after the implementation. If a test cannot be run, state that explicitly. +6. **Do not run the full test suite:** Use focused scripts and minimal reproductions unless a targeted test set is explicitly required. +7. **Fix issues end-to-end:** Reproduce, diagnose, implement, and verify without handing debugging back to the user. -**Key capabilities:** -- Parse models with user-friendly syntax (time indices like `[0], [-1], [1]`) -- Solve models automatically from equations and parameter values -- Calculate first, second, and third order (pruned) perturbation solutions -- Handle occasionally binding constraints -- Calculate impulse response functions, simulations, and conditional forecasts -- Estimate models using gradient-based samplers (NUTS, HMC) or inversion filters -- Differentiate solutions and moments with respect to parameters +## Core Engineering Principles -**Target audience:** Central bankers, regulators, graduate students, and researchers in DSGE modeling. +- Write all output/log files to the project folder (e.g. `tasks/`), never to `/tmp`. +- Keep changes minimal, focused, and at root cause. +- Keep code parsimonious and readable; apply Occam's razor to code changes. +- Preserve performance characteristics (type stability, allocations, threading behavior). +- Update user-facing docs/docstrings when public APIs change. +- Avoid second-person phrasing ("you") in docs/docstrings. +- Cache reusable constants lazily in model caches when appropriate. +- Avoid try-catch statements except for catching numerical issues that would otherwise error. Use explicit checks and validation instead of relying on exception handling for control flow. +- **rrule implementation:** Always derive analytical results for pullback functions. Never use AD inside a pullbackโ€”compute adjoints directly via mathematical derivation. -**Timing convention:** End-of-period (not start-of-period like some other packages). +## Task Files (Required Discipline) -## Project Structure +- Track plan/progress in `tasks/todo.md`. +- After corrections, capture reusable lessons in `tasks/lessons.md`. +- Keep `AGENT_PROGRESS.md` updated with what was done and what remains. -``` -MacroModelling.jl/ -โ”œโ”€โ”€ src/ # Main source code -โ”‚ โ”œโ”€โ”€ MacroModelling.jl # Main module, exports, type definitions -โ”‚ โ”œโ”€โ”€ macros.jl # @model and @parameters macros -โ”‚ โ”œโ”€โ”€ get_functions.jl # User-facing API (IRFs, simulations, forecasts) -โ”‚ โ”œโ”€โ”€ perturbation.jl # Perturbation solution algorithms (1st-3rd order) -โ”‚ โ”œโ”€โ”€ moments.jl # Model moment calculations -โ”‚ โ”œโ”€โ”€ structures.jl # Core data structures and types -โ”‚ โ”œโ”€โ”€ options_and_caches.jl # Solution caching and calculation options -โ”‚ โ”œโ”€โ”€ dynare.jl # Dynare file import support -โ”‚ โ”œโ”€โ”€ inspect.jl # Model inspection utilities -โ”‚ โ”œโ”€โ”€ solver_parameters.jl # Solver configuration parameters -โ”‚ โ”œโ”€โ”€ default_options.jl # Default option values -โ”‚ โ”œโ”€โ”€ common_docstrings.jl # Shared documentation strings -โ”‚ โ”œโ”€โ”€ algorithms/ # Matrix equation solvers (sylvester, lyapunov, quadratic_matrix_equation, nonlinear_solver) -โ”‚ โ”œโ”€โ”€ filter/ # Kalman and inversion filters (kalman, inversion, find_shocks) -โ”‚ โ””โ”€โ”€ custom_autodiff_rules/ # AD rules (forwarddiff, zygote) -โ”œโ”€โ”€ test/ # Test suite with multiple test sets -โ”œโ”€โ”€ models/ # Example DSGE models from literature -โ”œโ”€โ”€ docs/ # Documentation (Documenter.jl) -โ”œโ”€โ”€ benchmark/ # Benchmark scripts (BenchmarkTools) -โ””โ”€โ”€ ext/ # Package extensions (StatsPlots, Turing, Optim) -``` +## Critical Non-Negotiables -## Development Setup +1. Never claim something works without test evidence. +2. Work modularly and verify each completed module. +3. Iterate on failures independently; do not rely on user retesting loops. +4. Be explicit about unknowns; do not guess. +5. Verify before marking tasks complete. -### Julia Requirements +## On-Demand Companion Guides (Read Only If Needed) -- **Julia version:** 1.10 or higher (tested on 1.10+, lts, and pre-release versions) -- **Running Julia:** Always use `julia -t auto` to enable multi-threading - -### Package Setup - -```julia -using Pkg -Pkg.activate(".") -Pkg.instantiate() -``` - -## Revise-Based Development Workflow (REQUIRED) - -**ALWAYS use Revise.jl for interactive development.** This enables hot-reloading of code changes without restarting Julia, which is essential for efficient iteration. - -### Setup Steps - -1. **Start Julia REPL** with multi-threading enabled: - - ```bash - cd /path/to/MacroModelling.jl - julia -t auto --project=. - ``` - -2. **Load Revise FIRST**, then MacroModelling: - - ```julia - using Revise - using MacroModelling - ``` - -3. **Define a test model** for quick testing: - - ```julia - @model RBC 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]^ฮฑ - z[0] = ฯ * z[-1] + std_z * eps_z[x] - end - - @parameters RBC begin - std_z = 0.01 - ฯ = 0.2 - ฮด = 0.02 - ฮฑ = 0.5 - ฮฒ = 0.95 - end - ``` - -### Development Workflow - -1. **Keep the Julia REPL running** throughout the session - never restart between edits -2. **Edit source files** in `src/` directory -3. **Revise automatically detects changes** and recompiles only affected functions -4. **Test changes immediately** in the same REPL session -5. **Iterate rapidly** - edit, test, fix, repeat without restarting - -### Practical Example - -```julia -# Initial call (before any edits) -julia> get_equations(RBC) -4-element Vector{String}: - "1 / c[0] = (ฮฒ / c[1]) * (ฮฑ * exp(z[1]) * k[0] ^ (ฮฑ - 1) + (1 - ฮด))" - ... - -# Now edit src/inspect.jl to add a print statement: -# println("๐Ÿ” get_equations called - Revise is working!") -# Save the file - Revise detects the change automatically - -# Call again - no restart needed! -julia> get_equations(RBC) -๐Ÿ” get_equations called - Revise is working! -4-element Vector{String}: - "1 / c[0] = (ฮฒ / c[1]) * (ฮฑ * exp(z[1]) * k[0] ^ (ฮฑ - 1) + (1 - ฮด))" - ... -``` - -### Why This Matters - -- **Eliminates precompilation delays** - changes apply in seconds, not minutes -- **Preserves session state** - models, variables, and computations persist -- **Enables rapid debugging** - add/remove print statements instantly -- **Essential for this package** - MacroModelling has significant compile times - -### Important Caveats - -- **Revise must be loaded BEFORE MacroModelling** - order matters! -- **Structural changes require restart** - new types, module reorganization, or changing `__init__` functions -- **Manual refresh available** - if a change isn't detected, run `Revise.revise()` - -## Testing - -**Do NOT run the full test suite** - it takes too long. Instead: - -### Quick Feature Testing - -Write a bespoke script using the simple RBC model shown above, then test your changes: - -```julia -# Test your changes here -get_irf(RBC) -simulate(RBC) -``` - -### Test Sets (CI Only) - -Tests are organized by test sets specified via `TEST_SET` environment variable: - -- `basic`, `estimation`, `higher_order_1-3`, `plots_1-5`, `estimate_sw07`, `jet` -- Estimation tests: `1st_order_inversion_estimation`, `2nd_order_estimation`, `pruned_2nd_order_estimation`, `3rd_order_estimation`, `pruned_3rd_order_estimation` -- Pigeons estimation tests: `estimation_pigeons`, `1st_order_inversion_estimation_pigeons`, `2nd_order_estimation_pigeons`, `pruned_2nd_order_estimation_pigeons`, `3rd_order_estimation_pigeons`, `pruned_3rd_order_estimation_pigeons` - -```bash -TEST_SET=basic julia --project -e 'using Pkg; Pkg.test()' -``` - -### Test Environment Setup - -```julia -using Pkg -Pkg.activate("test") -Pkg.instantiate() -``` - -## Documentation - -Build documentation locally: - -```bash -julia --project=docs docs/make.jl -``` - -Documentation is built with Documenter.jl and deployed to GitHub Pages. - -## Benchmarking - -```julia -using BenchmarkTools -include("benchmark/benchmarks.jl") -run(SUITE) -``` - -## Model Syntax - -- **Variables** use time indices: `...[2], [1], [0], [-1], [-2]...` -- **Shocks** use `[x]`: `eps_z[x]` -- **Calibration equations** use `|` syntax in `@parameters` block -- **Custom steady state** can be provided via `steady_state_function` parameter - -## Code Style and Conventions - -### General Principles - -1. **Minimal changes:** Make the smallest possible changes to accomplish the task -2. **Testing:** Test changes with simple models rather than running the full test suite -3. **Performance:** This package emphasizes performance - be mindful of type stability and allocations -4. **Documentation:** Update docstrings when modifying public APIs - -### Writing Style - -- Avoid second-person phrasing ("you") in docs and docstrings - -### Caching Guidance - -- For constant calculations that can be computed once and reused, compute lazily on first use and store in the model struct cache; subsequent use must read from the cache - -## Key Design Considerations - -- **Performance critical** - Package competes with Dynare/RISE. Be mindful of type stability and allocations. -- **Symbolic mathematics** - Uses Symbolics.jl and SymPyPythonCall for symbolic derivatives compiled to efficient numerical code. -- **Automatic differentiation** - Supports forward and reverse-mode AD for gradients w.r.t. parameters. -- **Thread safety** - Important for estimation tasks. - -## Common Tasks - -### Adding a New Feature - -1. Write the feature in the appropriate `src/` file -2. Create a minimal test script (don't rely on full test suite) -3. Test with the simple RBC model -4. Update documentation if it's a user-facing feature - -### Fixing a Bug - -1. Identify the issue location in `src/` -2. Write a minimal reproduction case -3. Fix and verify with test script -4. Ensure existing functionality isn't broken - -### Adding a New Model - -1. Place in `models/` directory -2. Follow existing model structure -3. Include citation information -4. Test that it solves and produces IRFs - -### Common Change Points - -- **New API:** add in `src/get_functions.jl` and export from `src/MacroModelling.jl` -- **New model:** add a file under `models/` using the model macros -- **Solver changes:** look in `src/perturbation.jl` and `src/algorithms/` - -## CI/CD Pipeline - -- **CI runs on:** push (pull requests are commented out in workflow) -- **Platforms:** Ubuntu, macOS, Windows (x64 and arm64 where applicable) -- **Coverage:** Uploaded to Codecov -- **Matrix testing:** Multiple test sets run in parallel across different OS/architecture combinations - -## Core Principles - -- **Simplicity First:** Make every change as simple as possible. Impact minimal code. -- **No Laziness:** Find root causes. No temporary fixes. Senior developer standards. -- **Minimal Impact:** Changes should only touch what's necessary. - -## Workflow Orchestration - -### Plan Mode Default - -- Enter plan mode for ANY non-trivial task (3+ steps or architectural decisions) -- If something goes sideways, STOP and re-plan immediately - don't keep pushing -- Use plan mode for verification steps, not just building -- Write detailed specs upfront to reduce ambiguity - -### Subagent Strategy - -- Use subagents liberally to keep main context window clean -- Offload research, exploration, and parallel analysis to subagents -- For complex problems, throw more compute at it via subagents -- One task per subagent for focused execution - -### Demand Elegance (Balanced) - -- For non-trivial changes: pause and ask "is there a more elegant way?" -- If a fix feels hacky: "Knowing everything I know now, implement the elegant solution" -- Skip this for simple, obvious fixes - don't over-engineer -- Challenge your own work before presenting it - -### Autonomous Bug Fixing - -- When given a bug report: just fix it. Don't ask for hand-holding -- Point at logs, errors, failing tests - then resolve them -- Zero context switching required from the user -- Go fix failing CI tests without being told how - -## Task Management - -1. **Plan First:** Write plan to `tasks/todo.md` with checkable items -2. **Verify Plan:** Check in before starting implementation -3. **Track Progress:** Mark items complete as you go -4. **Explain Changes:** High-level summary at each step -5. **Document Results:** Add review section to `tasks/todo.md` -6. **Capture Lessons:** Update `tasks/lessons.md` after corrections - -### Session Progress Log - -- Always take stock of what was done and what remains, and save it in `AGENT_PROGRESS.md` -- At the start of a new session, always read `AGENT_PROGRESS.md` before making changes - -### Self-Improvement Loop - -- After ANY correction from the user: update `tasks/lessons.md` with the pattern -- Write rules for yourself that prevent the same mistake -- Ruthlessly iterate on these lessons until mistake rate drops -- Review lessons at session start for relevant project - -## CRITICAL WORKFLOW REQUIREMENTS - -**These rules are non-negotiable.** - -1. **NEVER claim something works without running a test to prove it.** After writing any code, immediately write and run a test. If you cannot test it, say so explicitly. - -2. **Work modularly.** Complete one module at a time. After each module, report what you built, show test results. - -3. **Iterate and fix errors yourself.** Do not rely on the user to report errors back to you. Run the code, observe the output, and fix problems before presenting results. - -4. **Be explicit about unknowns.** If you're uncertain about something, say so. Don't guess. - -5. **Verify before done.** Never mark a task complete without proving it works. Diff behavior between main and your changes when relevant. Ask yourself: "Would a staff engineer approve this?" +- Development setup, Revise workflow, testing, docs, benchmarking: `docs/agent-guides/development-workflow.md` +- Project overview, structure, model syntax, design context: `docs/agent-guides/project-context.md` +- Task runbook, orchestration heuristics, common change points: `docs/agent-guides/task-runbook.md` ## Additional Resources -- **Documentation:** https://thorek1.github.io/MacroModelling.jl/stable -- **Issue tracker:** GitHub Issues -- **Contributing guidelines:** See CONTRIBUTING.md -- **Code of Conduct:** See CODE_OF_CONDUCT.md +- Documentation: https://thorek1.github.io/MacroModelling.jl/stable +- Issue tracker: GitHub Issues +- Contributing guidelines: `CONTRIBUTING.md` +- Code of Conduct: `CODE_OF_CONDUCT.md` diff --git a/Project.toml b/Project.toml index 4ebce38e3..660e777c6 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" DynarePreprocessor_jll = "23afba7c-24e5-5ee2-bc2c-b42e07f0492a" +FastLapackInterface = "29a986be-02c6-4525-aec4-84b980013641" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" @@ -30,7 +31,6 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" Showoff = "992d4aef-0814-514b-bc4d-f2e9a6c4116f" @@ -59,17 +59,17 @@ Accessors = "0.1" Aqua = "0.8" AxisKeys = "0.2" BlockTriangularForm = "0.1" -CSV = "0.10" ChainRulesCore = "1" Combinatorics = "1" -DataFrames = "1" DataStructures = "0.18, 0.19" Dates = "1" +DelimitedFiles = "1" DifferentiationInterface = "0.6,0.7" DispatchDoctor = "0.4" DocStringExtensions = "0.8, 0.9" DynamicPPL = "0.35 - 0.38" DynarePreprocessor_jll = "6" +FastLapackInterface = "2" FiniteDifferences = "0.12" ForwardDiff = "0.10, 1" JET = "0.07 - 0.11" @@ -92,7 +92,6 @@ Preferences = "1" PythonCall = "0.9" REPL = "1" Random = "1" -RecursiveFactorization = "0.2" Reexport = "1" RuntimeGeneratedFunctions = "0.5" Showoff = "1" @@ -113,9 +112,8 @@ julia = "1.10" [extras] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" -DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" @@ -130,4 +128,4 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["ADTypes", "Aqua", "JET", "Dates", "CSV", "DataFrames", "DynamicPPL", "MCMCChains", "LineSearches", "Optim", "Test", "Turing", "Pigeons", "FiniteDifferences", "Zygote", "StatsPlots", "Preferences"] +test = ["ADTypes", "Aqua", "JET", "Dates", "DelimitedFiles", "DynamicPPL", "MCMCChains", "LineSearches", "Optim", "Test", "Turing", "Pigeons", "FiniteDifferences", "Zygote", "StatsPlots", "Preferences"] diff --git a/benchmark/bench.jl b/benchmark/bench.jl new file mode 100644 index 000000000..e8ac65ac7 --- /dev/null +++ b/benchmark/bench.jl @@ -0,0 +1,305 @@ +using Revise +using MacroModelling +using Zygote, ForwardDiff, FiniteDifferences +using BenchmarkTools +using LinearAlgebra + +include("../models/Smets_Wouters_2007.jl") + +model = Smets_Wouters_2007 + +params = deepcopy(model.parameter_values) +param_idx = 1 + +# MacroModelling.DEFAULT_SOLVER_PARAMETERS[7] +# MacroModelling.solver_parameters(6.8658210317889115, 3.054280631509596, 9.239560890529688, 5.0330393159601705, 4.619974181880515, 2.130665389110862, 13.395678237998878, 8.95412704048986, 16.67031860308238, 4.1686309854116175, 7.193385978766233, 6.284359482297452, 1.6025436780830082, 4.080789181245917, 11.237586964445232, 0.9812514892088027, 10.182504561803604, 2.2723756926184744, 5.580529028552923, 4.761189900509761, 1, 0.0, 2) + +popfirst!(MacroModelling.DEFAULT_SOLVER_PARAMETERS) +pushfirst!(MacroModelling.DEFAULT_SOLVER_PARAMETERS, MacroModelling.DEFAULT_SOLVER_PARAMETERS[3]); + +MacroModelling.clear_solution_caches!(model, :first_order) +get_statistics(model, params, non_stochastic_steady_state = :all, verbose = true) + +out_bench = @benchmark get_statistics(model, params, non_stochastic_steady_state = :all) setup = MacroModelling.clear_solution_caches!(model, :first_order) + +@profview for i in 1:10000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_statistics(model, params, non_stochastic_steady_state = :all) +end + +@profview_allocs for i in 1:10000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_statistics(model, params, non_stochastic_steady_state = :all) +end + + +# first order solution +MacroModelling.clear_solution_caches!(model, :first_order) +get_solution(model, params)[2] + +out_bench = @benchmark get_solution(model, params) setup = MacroModelling.clear_solution_caches!(model, :first_order) + +@profview for i in 1:5000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params) + get_solution(model, params .+ 0.001) +end + +@profview_allocs for i in 1:5000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params) + get_solution(model, params .+ 0.001) +end + + +# Gradients +# Zygote +MacroModelling.clear_solution_caches!(model, :first_order) +Zygote.gradient(x->norm(get_solution(model, x)[2]),params) + + +out_bench = @benchmark Zygote.gradient(x->norm(get_solution(model, x)),params) setup = MacroModelling.clear_solution_caches!(model, :first_order) + + +@profview for i in 1:1000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + Zygote.gradient(x->norm(get_solution(model, x)),params) +end + +@profview_allocs for i in 1:1000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + Zygote.gradient(x->norm(get_solution(model, x)),params) +end + +# ForwardDiff +MacroModelling.clear_solution_caches!(model, :first_order) +first_order_one_param = x -> begin + perturbed = convert.(eltype(x),copy(params)) + perturbed[param_idx] = x + get_solution(model, perturbed)[2] +end + +ForwardDiff.derivative(first_order_one_param, params[param_idx]) + + + +out_bench = @benchmark ForwardDiff.derivative(first_order_one_param, params[param_idx]) setup = MacroModelling.clear_solution_caches!(model, :first_order) + + +@profview for i in 1:1000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + ForwardDiff.derivative(first_order_one_param, params[param_idx]) +end + +@profview_allocs for i in 1:1000 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + ForwardDiff.derivative(first_order_one_param, params[param_idx]) +end + +# FiniteDifferences +MacroModelling.clear_solution_caches!(model, :first_order) +FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) + + +out_bench = @benchmark FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) setup = MacroModelling.clear_solution_caches!(model, :first_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) +end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) +end + + +# second order solution +MacroModelling.clear_solution_caches!(model, :first_order) +get_solution(model, params, algorithm = :second_order)[3] * model.constants.second_order.๐”โ‚‚ |> norm + +out_bench = @benchmark get_solution(model, params, algorithm = :second_order) setup = MacroModelling.clear_solution_caches!(model, :second_order) + +@profview for i in 1:500 + MacroModelling.clear_solution_caches!(model, :second_order) + get_solution(model, params) + get_solution(model, params .+ 0.001, algorithm = :second_order) +end + +@profview_allocs for i in 1:500 + MacroModelling.clear_solution_caches!(model, :second_order) + get_solution(model, params) + get_solution(model, params .+ 0.001, algorithm = :second_order) +end + + +# Gradients +# Zygote +MacroModelling.clear_solution_caches!(model, :second_order) +Zygote.gradient(x->norm(get_solution(model, x, algorithm = :second_order)[3] * model.constants.second_order.๐”โ‚‚),params)[1] + + +out_bench = @benchmark Zygote.gradient(x->norm(get_solution(model, x, algorithm = :second_order)),params) setup = MacroModelling.clear_solution_caches!(model, :second_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :second_order) + get_solution(model, params .+ 0.001) + Zygote.gradient(x->norm(get_solution(model, x, algorithm = :second_order)[3]),params)[1] +end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :second_order) + get_solution(model, params .+ 0.001) + Zygote.gradient(x->norm(get_solution(model, x, algorithm = :second_order)[3]),params)[1] +end + +# ForwardDiff +MacroModelling.clear_solution_caches!(model, :second_order) +second_order_one_param = x -> begin + perturbed = convert.(eltype(x),copy(params)) + perturbed[param_idx] = x + get_solution(model, perturbed, algorithm = :second_order)[3] * model.constants.second_order.๐”โ‚‚ +end + +ForwardDiff.derivative(second_order_one_param, params[param_idx]) + + +out_bench = @benchmark ForwardDiff.derivative(second_order_one_param, params[param_idx]) setup = MacroModelling.clear_solution_caches!(model, :second_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :second_order) + get_solution(model, params .+ 0.001) + ForwardDiff.derivative(second_order_one_param, params[param_idx]) + end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :second_order) + get_solution(model, params .+ 0.001) + ForwardDiff.derivative(second_order_one_param, params[param_idx]) +end + +# FiniteDifferences +MacroModelling.clear_solution_caches!(model, :first_order) +FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) + + +out_bench = @benchmark FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) setup = MacroModelling.clear_solution_caches!(model, :first_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) +end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) +end + + +# third order solution +include("../models/FS2000.jl") +model = FS2000 + +params = deepcopy(model.parameter_values) +param_idx = 1 + +MacroModelling.clear_solution_caches!(model, :first_order) +get_solution(model, params, algorithm = :third_order)[4] * model.constants.third_order.๐”โ‚ƒ |> norm + +out_bench = @benchmark get_solution(model, params, algorithm = :third_order) setup = MacroModelling.clear_solution_caches!(model, :third_order) + +@profview for i in 1:10 + MacroModelling.clear_solution_caches!(model, :third_order) + get_solution(model, params) + get_solution(model, params .+ 0.001, algorithm = :third_order) +end + +@profview_allocs for i in 1:10 + MacroModelling.clear_solution_caches!(model, :third_order) + get_solution(model, params) + get_solution(model, params .+ 0.001, algorithm = :third_order) +end + + +# Gradients +# Zygote +MacroModelling.clear_solution_caches!(model, :third_order) +zyg_grad = Zygote.gradient(x->norm(get_solution(model, x, algorithm = :third_order)[4] * model.constants.third_order.๐”โ‚ƒ),params)[1] + + +out_bench = @benchmark Zygote.gradient(x->norm(get_solution(model, x, algorithm = :third_order)[4]),params) setup = MacroModelling.clear_solution_caches!(model, :third_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :third_order) + get_solution(model, params .+ 0.001) + Zygote.gradient(x->norm(get_solution(model, x, algorithm = :third_order)[4]),params)[1] +end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :third_order) + get_solution(model, params .+ 0.001) + Zygote.gradient(x->norm(get_solution(model, x, algorithm = :third_order)[4]),params)[1] +end + +# FiniteDifferences +MacroModelling.clear_solution_caches!(model, :first_order) +fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1),x->norm(get_solution(model, x, algorithm = :third_order)[4] * model.constants.third_order.๐”โ‚ƒ),params)[1] + +isapprox(zyg_grad,fin_grad) +zyg_grad - fin_grad +norm(zyg_grad - fin_grad)/max(norm(zyg_grad), norm(fin_grad)) + +out_bench = @benchmark FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) setup = MacroModelling.clear_solution_caches!(model, :first_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) +end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :first_order) + get_solution(model, params .+ 0.001) + FiniteDifferences.grad(FiniteDifferences.central_fdm(2,1),x->norm(get_solution(model, x)),params) +end + +# ForwardDiff +MacroModelling.clear_solution_caches!(model, :third_order) +third_order_one_param = x -> begin + perturbed = convert.(eltype(x),copy(params)) + perturbed[param_idx] = x + norm(get_solution(model, perturbed, algorithm = :third_order)[4]) + # get_solution(model, perturbed, algorithm = :third_order)[4] * model.constants.third_order.๐”โ‚ƒ +end + +ForwardDiff.derivative(third_order_one_param, params[param_idx]) + + +out_bench = @benchmark ForwardDiff.derivative(third_order_one_param, params[param_idx]) setup = MacroModelling.clear_solution_caches!(model, :third_order) + + +@profview for i in 1:100 + MacroModelling.clear_solution_caches!(model, :third_order) + get_solution(model, params .+ 0.001) + ForwardDiff.derivative(third_order_one_param, params[param_idx]) + end + +@profview_allocs for i in 1:100 + MacroModelling.clear_solution_caches!(model, :third_order) + get_solution(model, params .+ 0.001) + ForwardDiff.derivative(third_order_one_param, params[param_idx]) +end diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 0f522c9f7..7519f67d6 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -21,7 +21,7 @@ const HAS_WORKSPACE_API = isdefined(MacroModelling, :Lyapunov_workspace) # Conditionally import workspace types only if they exist if HAS_WORKSPACE_API - import MacroModelling: Lyapunov_workspace, lyapunov_workspace, ensure_lyapunov_workspace!, ensure_qme_workspace!, ensure_sylvester_1st_order_workspace! + import MacroModelling: Lyapunov_workspace, lyapunov_workspace, ensure_lyapunov_workspace!, ensure_qme_workspace! end # Version-aware wrapper for solve_lyapunov_equation benchmarking @@ -49,7 +49,7 @@ end function first_order_solution_for_bench(โˆ‡โ‚::AbstractMatrix, ๐“‚::โ„ณ; opts = merge_calculation_options()) if HAS_WORKSPACE_API qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) + sylv_ws = ๐“‚.workspaces.sylvester_1st_order out = calculate_first_order_solution(โˆ‡โ‚, ๐“‚.constants, qme_ws, sylv_ws; opts = opts) else out = calculate_first_order_solution(โˆ‡โ‚; T = timings_for_bench(๐“‚), opts = opts) @@ -58,12 +58,7 @@ function first_order_solution_for_bench(โˆ‡โ‚::AbstractMatrix, ๐“‚::โ„ณ; opts end function calculate_jacobian_for_bench(parameters, SS_and_pars, ๐“‚::โ„ณ) - if hasmethod(calculate_jacobian, Tuple{typeof(parameters), typeof(SS_and_pars), โ„ณ}) - out = calculate_jacobian(parameters, SS_and_pars, ๐“‚) - else - out = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian) - end - return out + return calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) end diff --git a/benchmark/sw07_random_parameter_ss_benchmark.jl b/benchmark/sw07_random_parameter_ss_benchmark.jl new file mode 100644 index 000000000..4d02e404c --- /dev/null +++ b/benchmark/sw07_random_parameter_ss_benchmark.jl @@ -0,0 +1,86 @@ +using Revise +using MacroModelling +using BenchmarkTools +using DelimitedFiles +using AxisKeys + +include(joinpath(@__DIR__, "..", "models", "Smets_Wouters_2007.jl")) + +model = Smets_Wouters_2007 + +# Same SW07 data preparation used in test/test_sw07_estimation.jl +raw_data, raw_header = readdlm(joinpath(@__DIR__, "..", "test", "data", "usmodel.csv"), ',', Float64, '\n'; header = true) +variable_names = Symbol.(strip.(vec(raw_header))) +data = KeyedArray(raw_data', Variable = variable_names, Time = 1:size(raw_data, 1)) + +observables_old = [:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs] +sample_idx = 47:230 +data = data(observables_old, sample_idx) + +observables = [:dy, :dc, :dinve, :labobs, :pinfobs, :dwobs, :robs] +data = rekey(data, :Variable => observables) + +llh_data = data(observables) +known_parameters = copy(model.parameter_values) +new_parameters = known_parameters .+ 0.001 + +function clear_nsss_cache!(m) + while length(m.caches.solver_cache) > 1 + pop!(m.caches.solver_cache) + end + return nothing +end + +clear_nsss_cache!(model) + +function evaluate_llh(m, data, parameters) + return get_loglikelihood( + m, + data, + parameters; + presample_periods = 4, + initial_covariance = :diagonal, + quadratic_matrix_equation_algorithm = :doubling, + filter = :kalman, + ) +end + +function setup_known_to_new_transition!(m, data, known_params) + clear_nsss_cache!(m) + evaluate_llh(m, data, known_params) + return nothing +end + +# Warm-up compile and ensure LLHs are finite before benchmarking. +llh_known = evaluate_llh( + model, + llh_data, + known_parameters, +) +llh_new = evaluate_llh( + model, + llh_data, + new_parameters, +) +println("Warm-up known LLH: ", llh_known) +println("Warm-up new LLH: ", llh_new) + +trial = @benchmark evaluate_llh( + $model, + $llh_data, + $new_parameters, +) setup = setup_known_to_new_transition!($model, $llh_data, $known_parameters) + +@profview_allocs for _ in 1:10000 + setup_known_to_new_transition!(model, llh_data, known_parameters) + evaluate_llh(model, llh_data, new_parameters) +end + +@profview for _ in 1:1000 + setup_known_to_new_transition!(model, llh_data, known_parameters) + evaluate_llh(model, llh_data, new_parameters) +end + +println(trial) +println("Minimum time: ", minimum(trial).time, " ns") +println("Minimum memory: ", minimum(trial).memory, " bytes") diff --git a/benchmark/sw07_third_order_pullback_repl.jl b/benchmark/sw07_third_order_pullback_repl.jl new file mode 100644 index 000000000..f340fc14c --- /dev/null +++ b/benchmark/sw07_third_order_pullback_repl.jl @@ -0,0 +1,511 @@ +#= + REPL-style script to step through the third-order solution pullback + for the Smetsโ€“Wouters 2007 model. + + Objective (same as bench.jl): + f(params) = norm( S3_raw * U3 ) + Tangent wrt S3_raw: + โˆ‚f/โˆ‚S3 = (S3*U3 / norm(S3*U3)) * U3' + + This script: + 1. Builds all primal inputs (โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐‘บโ‚, ๐‘บโ‚‚) + 2. Runs the rrule forward pass โ†’ captures S3_raw + closure variables + 3. Computes the cotangent seed โˆ‚S3_raw from norm(S3 * U3) + 4. PASTES THE PULLBACK CODE INLINE so you can step through it + + Run the whole file once, then use Debugger.jl to step through the pullback. +=# + +using Revise +using MacroModelling +using LinearAlgebra +using SparseArrays +using ChainRulesCore: rrule, NoTangent + +const MM = MacroModelling +const โ„’ = LinearAlgebra + +include(joinpath(@__DIR__, "..", "models", "Smets_Wouters_2007.jl")) + +model = Smets_Wouters_2007 +parameters = copy(model.parameter_values) +opts = MM.merge_calculation_options(verbose = false) + +# ============================================================================== +# STEP 1: Build primal inputs +# ============================================================================== +MM.clear_solution_caches!(model, :third_order) + +# Warm-up (derivative caches) +_, _, _, _, solved_warmup = MM.get_solution(model, parameters, + algorithm = :third_order, verbose = false) +@assert solved_warmup +MM.clear_solution_caches!(model, :third_order) + +# Non-stochastic steady state +SS_and_pars, (solution_error, _) = MM.get_NSSS_and_parameters(model, parameters; opts = opts) +@assert solution_error <= opts.tol.NSSS_acceptance_tol + +# Jacobian โˆ‡โ‚ +โˆ‡โ‚ = MM.calculate_jacobian(parameters, SS_and_pars, + model.caches, model.functions.jacobian, model.workspaces) + +# First-order perturbation solution +๐‘บโ‚, _, solved1 = MM.calculate_first_order_solution( + โˆ‡โ‚, model.constants, model.workspaces, model.caches; + opts = opts, initial_guess = model.caches.qme_solution) +@assert solved1 + +# Hessian โˆ‡โ‚‚ (compressed) +โˆ‡โ‚‚_input = MM.calculate_hessian(parameters, SS_and_pars, + model.caches, model.functions.hessian, model.workspaces) + +# Second-order perturbation solution (compressed) +๐‘บโ‚‚_input, solved2 = MM.calculate_second_order_solution( + โˆ‡โ‚, โˆ‡โ‚‚_input, ๐‘บโ‚, model.constants, model.workspaces, model.caches; + initial_guess = model.caches.second_order_solution, opts = opts) +@assert solved2 + +# Third-order derivative tensor โˆ‡โ‚ƒ +โˆ‡โ‚ƒ = MM.calculate_third_order_derivatives( + parameters, SS_and_pars, + model.caches, model.functions.third_order_derivatives, model.workspaces) + +println("Step 1 done โ€“ primal inputs ready.") + + +# ============================================================================== +# STEP 2: rrule forward pass - captures all closure variables +# ============================================================================== + +third_out, third_pb = rrule(MM.calculate_third_order_solution, + โˆ‡โ‚, โˆ‡โ‚‚_input, โˆ‡โ‚ƒ, ๐‘บโ‚, ๐‘บโ‚‚_input, + model.constants, model.workspaces, model.caches; + initial_guess = model.caches.third_order_solution, + opts = opts) + +๐’โ‚ƒ_raw, solved3 = third_out +@assert solved3 "Third-order Sylvester solve failed." + +println("Step 2 done โ€“ S3_raw: ", size(๐’โ‚ƒ_raw), " nnz = ", nnz(sparse(๐’โ‚ƒ_raw))) + + +# ============================================================================== +# STEP 3: Compute cotangent seed from f = norm(S3_raw * U3) +# ============================================================================== + +Mโ‚ƒ = model.constants.third_order +๐”โ‚ƒ = Mโ‚ƒ.๐”โ‚ƒ + +๐’โ‚ƒ_full = ๐’โ‚ƒ_raw * ๐”โ‚ƒ +loss = โ„’.norm(๐’โ‚ƒ_full) +scale = max(loss, eps(eltype(loss))) +โˆ‚๐’โ‚ƒ_raw = (๐’โ‚ƒ_full / scale) * ๐”โ‚ƒ' + +println("Step 3 done โ€“ loss = ", loss) + + +# ============================================================================== +# STEP 4: INLINE PULLBACK CODE +# ============================================================================== +# This is the exact pullback code from rrules.jl third_order_solution_pullback. +# All variables it needs are captured from the rrule closure above. + +# Access closure variables (these are what the rrule captured) +# The closure contains: A, B, C, spinv, โˆ‡โ‚โ‚Š, โˆ‡โ‚‚t, โˆ‡โ‚ƒt, D_ab_t, tmpkron22, ck3_aux_mat, +# S2p0_sigma, mm_๐’โ‚‚_kron, Mโ‚‚, Mโ‚ƒ, T, iโ‚Š, iโ‚‹, nโ‚Š, nโ‚‹, n, nโ‚‘, nโ‚‘โ‚‹, +# โ„‚, opts, and many transposes + +# We need to rebuild some intermediates that were computed in the forward pass +# but not all are captured in the closure. Let's get what we need. + +S = eltype(โˆ‡โ‚) +โ„‚ = model.workspaces.third_order +Mโ‚‚ = model.constants.second_order +T = model.constants.post_model_macro + +# Expand compressed inputs +โˆ‡โ‚‚ = โˆ‡โ‚‚_input * Mโ‚‚.๐”โˆ‡โ‚‚ +๐’โ‚‚ = sparse(๐‘บโ‚‚_input * Mโ‚‚.๐”โ‚‚)::SparseMatrixCSC{S, Int} + +iโ‚Š = T.future_not_past_and_mixed_idx +iโ‚‹ = T.past_not_future_and_mixed_idx +nโ‚‹ = T.nPast_not_future_and_mixed +nโ‚Š = T.nFuture_not_past_and_mixed +nโ‚‘ = T.nExo +n = T.nVars +nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + +# Build S1 embedding (same as forward pass) +๐’โ‚ = โ„‚.๐’โ‚::Matrix{S} +copyto!(@view(๐’โ‚[:,1:nโ‚‹]), @view(๐‘บโ‚[:,1:nโ‚‹])) +fill!(@view(๐’โ‚[:,nโ‚‹+1]), zero(S)) +copyto!(@view(๐’โ‚[:,nโ‚‹+2:end]), @view(๐‘บโ‚[:,nโ‚‹+1:end])) + +# S1_{-1e} matrix +๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} +copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹,:]), @view(๐’โ‚[iโ‚‹,:])) +fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1:end,:]), zero(S)) +@inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1,nโ‚‹+1] = one(S) +๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = MM.choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold=1.0, min_length=10, tol=opts.tol.droptol) + +# S1 stacking matrix +โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [ + (๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] + ๐’โ‚ + โ„’.I(nโ‚‘โ‚‹)[[range(1,nโ‚‹)..., nโ‚‹+1 .+ range(1,nโ‚‘)...],:] +] + +# S1 on future rows +๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:]; zeros(S, nโ‚‹+n+nโ‚‘, nโ‚‘โ‚‹)] +๐’โ‚โ‚Šโ•ฑ๐ŸŽ = MM.choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold=1.0, min_length=10, tol=opts.tol.droptol) + +# โˆ‡โ‚โ‚ŠยทS1 + โˆ‡โ‚โ‚€ +โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * Mโ‚‚.๐ˆโ‚™โ‚‹ - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] +โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = โ„’.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) +spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) +spinv = MM.choose_matrix_format(spinv) + +โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š + +# A matrix +A = spinv * โˆ‡โ‚โ‚Š + +# B matrix +kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) +B = MM.compressed_permuted_mixed_kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”, + sparse_preallocation = โ„‚.tmp_sparse_prealloc7) +B += MM.compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.droptol, + sparse_preallocation = โ„‚.tmp_sparse_prealloc1) + +# S2 stacking matrices +โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = @views [ + (๐’โ‚‚ * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:]; zeros(S, nโ‚‘+1, nโ‚‘โ‚‹^2)])[iโ‚Š,:] + ๐’โ‚‚ + zeros(S, nโ‚‹+nโ‚‘, nโ‚‘โ‚‹^2) +] +โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = MM.choose_matrix_format( + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, density_threshold=0.0, min_length=10, tol=opts.tol.droptol) + +๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚Š,:]; zeros(S, nโ‚‹+n+nโ‚‘, nโ‚‘โ‚‹^2)] +๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = MM.choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold=1.0, min_length=10, tol=opts.tol.droptol) + +aux = Mโ‚ƒ.๐’๐ * โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ + +S1p0_kron_sigma = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›” + +tmpkron22 = MM.compressed_permuted_mixed_kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + S1p0_kron_sigma, + sparse_preallocation = โ„‚.tmp_sparse_prealloc6) + +โˆ‡โ‚โ‚Š = MM.choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold=1.0, min_length=10, tol=opts.tol.droptol) + +S2p0_sigma = ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›” + +# Build X3 (C matrix ingredients) +tmpkron2 = โ„’.kron(Mโ‚‚.๐›”, MM.choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold=0.0, tol=opts.tol.droptol)) +D_ab = (tmpkron2 + Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ) * Mโ‚ƒ.๐๐‚โ‚ƒ + +๐—โ‚ƒ = MM.mat_mult_kron(โˆ‡โ‚‚, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ), D_ab, + sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc2) + +๐—โ‚ƒ += MM.mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, Mโ‚ƒ.๐๐‚โ‚ƒ, + sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc3) + +๐—โ‚ƒ += MM.mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, collect(S2p0_sigma), Mโ‚ƒ.๐๐‚โ‚ƒ, + sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + +๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = MM.choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold=0.0, tol=opts.tol.droptol) +mm_๐’โ‚‚_kron = MM.mat_mult_kron(๐’โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, + [๐’โ‚‚[iโ‚‹,:]; zeros(S, size(๐’โ‚,2)-nโ‚‹, nโ‚‘โ‚‹^2)], sparse = true, + sparse_preallocation = โ„‚.tmp_sparse_prealloc4) +๐—โ‚ƒ += โˆ‡โ‚โ‚Š * mm_๐’โ‚‚_kron * Mโ‚ƒ.๐๐‚โ‚ƒ + +๐—โ‚ƒ += โˆ‡โ‚ƒ * tmpkron22 + +ck3_aux_mat = MM.compressed_kronยณ(aux, rowmask = Mโ‚ƒ.โˆ‡โ‚ƒ_rowmask, + tol = opts.tol.droptol, + sparse_preallocation = โ„‚.tmp_sparse_prealloc5) +๐—โ‚ƒ += โˆ‡โ‚ƒ * ck3_aux_mat + +C = spinv * ๐—โ‚ƒ + +# Solve Sylvester +๐’โ‚ƒ, solved = MM.solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, + initial_guess = zeros(S, 0, 0), + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.sylvester_tol, + acceptance_tol = opts.tol.sylvester_acceptance_tol, + verbose = opts.verbose) +@assert solved + +๐’โ‚ƒ_stable = copy(๐’โ‚ƒ) + +# Precompute transposes +๐๐‚โ‚ƒt = Mโ‚ƒ.๐๐‚โ‚ƒ' +๐›”t = Mโ‚‚.๐›”' +๐”โˆ‡โ‚‚t = Mโ‚‚.๐”โˆ‡โ‚‚' +๐”โ‚‚t = Mโ‚‚.๐”โ‚‚' + +Mโ‚ƒ๐โ‚โ‚—t = Mโ‚ƒ.๐โ‚โ‚—' +Mโ‚ƒ๐โ‚แตฃt = Mโ‚ƒ.๐โ‚แตฃ' + +โˆ‡โ‚‚t = MM.choose_matrix_format(โˆ‡โ‚‚') +โˆ‡โ‚ƒt = MM.choose_matrix_format(โˆ‡โ‚ƒ') +D_ab_t = MM.choose_matrix_format(D_ab') +tmpkron22_t = MM.choose_matrix_format(tmpkron22') +ck3_aux_mat_t = MM.choose_matrix_format(ck3_aux_mat') +๐’โ‚‚t = MM.choose_matrix_format(๐’โ‚‚', density_threshold=1.0) +โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t = MM.choose_matrix_format(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹') +โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt = MM.choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ') +S2p0_sigma_t = MM.choose_matrix_format(S2p0_sigma') +mm_๐’โ‚‚_kron_t = MM.choose_matrix_format(mm_๐’โ‚‚_kron') + +tmpkron22_ck3_aux_mat_t = MM.choose_matrix_format(tmpkron22_t + ck3_aux_mat_t) + +# Ensure pullback workspaces +MM.ensure_third_order_pullback_workspaces!(โ„‚, S, T, Mโ‚‚, Mโ‚ƒ) + +println("Step 4 done โ€“ forward pass intermediates rebuilt.") + + +# ============================================================================== +# STEP 5: INLINE PULLBACK - paste the pullback code here for stepping +# ============================================================================== +# Below is the pullback code. You can use Debugger.jl to step through it: +# using Debugger +# @enter third_order_solution_pullback(โˆ‚๐’โ‚ƒ_raw) +# +# Or copy-paste sections to run them individually. + +function third_order_solution_pullback(โˆ‚๐’โ‚ƒ) + #= + Pullback for calculate_third_order_solution. + This is pasted inline so you can step through it in the REPL. + =# + + if โ„’.norm(โˆ‚๐’โ‚ƒ) < opts.tol.sylvester_tol + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + println(" [Pullback] Step 1: adjoint Sylvester") + # --- adjoint Sylvester: Aแต€ โˆ‚C_adj Bแต€ + โˆ‚๐’โ‚ƒ = โˆ‚C_adj -------------------- + โˆ‚C_adj, slvd = MM.solve_sylvester_equation(A', B', Matrix{Float64}(โˆ‚๐’โ‚ƒ), โ„‚.sylvester_workspace, + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.sylvester_tol, + acceptance_tol = opts.tol.sylvester_acceptance_tol, + verbose = opts.verbose) + if !slvd + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + โˆ‚C_adj = MM.choose_matrix_format(โˆ‚C_adj) + println(" ||โˆ‚C_adj|| = ", โ„’.norm(Matrix(โˆ‚C_adj))) + + # --- Initialize all gradient accumulators --- + println(" [Pullback] Step 2: initialize accumulators") + โˆ‚๐—โ‚ƒ = โ„‚.โˆ‚๐—โ‚ƒ_3rd + โˆ‚A = โ„‚.โˆ‚A_3rd + โˆ‚B_from_sylv = โ„‚.โˆ‚B_sylv_3rd + โˆ‚out2 = โ„‚.โˆ‚out2_3rd + โˆ‡โ‚‚t_โˆ‚out2 = โ„‚.โˆ‡โ‚‚t_โˆ‚out2_3rd + mul_tmp = โ„‚.mul_tmp_3rd + โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = โ„‚.โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd + + โˆ‚spinv = โ„‚.โˆ‚spinv_3rd + โˆ‚โˆ‡โ‚ = โ„‚.โˆ‚โˆ‡โ‚_3rd; fill!(โˆ‚โˆ‡โ‚, zero(S)) + โˆ‚๐’โ‚โ‚ƒ = โ„‚.โˆ‚๐’โ‚_3rd; fill!(โˆ‚๐’โ‚โ‚ƒ, zero(S)) + + โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) # sparse โ€” must stay fresh + + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, zero(S)) + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, zero(S)) + โˆ‚L_c = โ„‚.โˆ‚L_c_3rd; fill!(โˆ‚L_c, zero(S)) + โˆ‚R_c = โ„‚.โˆ‚R_c_3rd; fill!(โˆ‚R_c, zero(S)) + โˆ‚L_d = โ„‚.โˆ‚L_d_3rd; fill!(โˆ‚L_d, zero(S)) + โˆ‚R_d = โ„‚.โˆ‚R_d_3rd; fill!(โˆ‚R_d, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8 = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, zero(S)) + โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, zero(S)) + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ, zero(S)) + โˆ‚S1S1_stack = โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd; fill!(โˆ‚S1S1_stack, zero(S)) + โˆ‚aux = โ„‚.โˆ‚aux_3rd; fill!(โˆ‚aux, zero(S)) + โˆ‚๐›”_discard = โ„‚.โˆ‚๐›”_discard_3rd; fill!(โˆ‚๐›”_discard, zero(S)) + โˆ‚๐›”_discard2 = โ„‚.โˆ‚tmpkron0_ฯƒ_3rd; fill!(โˆ‚๐›”_discard2, zero(S)) + + # --- gradient of A, B, C from ๐’โ‚ƒ = Aยท๐’โ‚ƒยทB + C --------------------------- + println(" [Pullback] Step 3: โˆ‚A, โˆ‚B, โˆ‚spinv, โˆ‚X3") + โ„’.mul!(โˆ‚๐—โ‚ƒ, โˆ‚C_adj, B') + โ„’.mul!(โˆ‚A, โˆ‚๐—โ‚ƒ, ๐’โ‚ƒ_stable') + โ„’.mul!(โˆ‚๐—โ‚ƒ, A', โˆ‚C_adj) + โ„’.mul!(โˆ‚B_from_sylv, ๐’โ‚ƒ_stable', โˆ‚๐—โ‚ƒ) + โˆ‚๐—โ‚ƒ = MM.choose_matrix_format(spinv' * โˆ‚C_adj, density_threshold = 1.0, min_length = 0) + โ„’.mul!(โˆ‚spinv, โˆ‚C_adj, ๐—โ‚ƒ') + โ„’.mul!(โˆ‚spinv, โˆ‚A, โˆ‡โ‚โ‚Š', 1, 1) + + # โˆ‚โˆ‡โ‚ƒ + println(" [Pullback] Step 4: โˆ‚โˆ‡โ‚ƒ") + โˆ‚โˆ‡โ‚ƒ = โˆ‚๐—โ‚ƒ * tmpkron22_ck3_aux_mat_t + + # โˆ‚โˆ‡โ‚‚ + println(" [Pullback] Step 5: โˆ‚โˆ‡โ‚‚") + โ„’.mul!(โˆ‚out2, โˆ‚๐—โ‚ƒ, ๐๐‚โ‚ƒt) + โˆ‚mid_ab = โˆ‚๐—โ‚ƒ * D_ab_t + โˆ‚โˆ‡โ‚‚ = MM.mat_mult_kron(โˆ‚mid_ab, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ'), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ')) + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ + MM.mat_mult_kron(โˆ‚out2, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt) + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ + MM.mat_mult_kron(โˆ‚out2, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, S2p0_sigma_t) + println(" ||โˆ‚โˆ‡โ‚‚|| = ", โ„’.norm(Matrix(โˆ‚โˆ‡โ‚‚))) + + # โˆ‚๐’โ‚‚ + println(" [Pullback] Step 6: โˆ‚๐’โ‚‚") + โ„’.mul!(โˆ‡โ‚‚t_โˆ‚out2, โˆ‡โ‚‚t, โˆ‚out2) + โˆ‚tmpkron1 = (โˆ‡โ‚‚t * โˆ‚mid_ab) + MM.fill_kron_adjoint!(โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, โˆ‚tmpkron1, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] + + โˆ‚kron_c = (โˆ‡โ‚‚t_โˆ‚out2) + MM.fill_kron_adjoint!(โˆ‚R_c, โˆ‚L_c, โˆ‚kron_c, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + nโ‚Š_len = length(iโ‚Š) + โˆ‚top_block = โˆ‚R_c[1:nโ‚Š_len, :] + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚top_block * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + โˆ‚๐’โ‚‚_padded = ๐’โ‚' * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block + @views โˆ‚๐’โ‚‚[iโ‚‹,:] .+= โˆ‚๐’โ‚‚_padded[1:nโ‚‹, :] + @views โˆ‚๐’โ‚‚ .+= โˆ‚R_c[nโ‚Š_len .+ (1:n), :] + + MM.fill_kron_adjoint!(โˆ‚R_d, โˆ‚L_d, โˆ‚kron_c, S2p0_sigma, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d = โˆ‚R_d * ๐›”t + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d[1:length(iโ‚Š),:] + + tmp_t8 = โˆ‡โ‚โ‚Š' * โˆ‚out2 + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ + MM.mat_mult_kron(tmp_t8, collect(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘'), collect([๐’โ‚‚[iโ‚‹,:]; zeros(S, size(๐’โ‚,2)-nโ‚‹, nโ‚‘โ‚‹^2)]')) + + โˆ‚kron_term8 = ((โˆ‡โ‚โ‚Š * ๐’โ‚‚)' * โˆ‚out2) + MM.fill_kron_adjoint!(โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, โˆ‚kron_term8, [๐’โ‚‚[iโ‚‹,:]; zeros(S, size(๐’โ‚,2)-nโ‚‹, nโ‚‘โ‚‹^2)], ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + @views โˆ‚๐’โ‚‚[iโ‚‹,:] .+= โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ[1:nโ‚‹,:] + println(" ||โˆ‚๐’โ‚‚|| = ", โ„’.norm(Matrix(โˆ‚๐’โ‚‚))) + + # โˆ‚โˆ‡โ‚ + println(" [Pullback] Step 7: โˆ‚โˆ‡โ‚") + โ„’.mul!(mul_tmp, spinv', โˆ‚spinv) + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, mul_tmp, spinv') + โ„’.rmul!(โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, -1) + + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] * ๐’โ‚[iโ‚Š,1:nโ‚‹]' + โˆ‚โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ + + โˆ‚โˆ‡โ‚โ‚Š = โ„‚.โˆ‚โˆ‡โ‚โ‚Š_3rd + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š, spinv', โˆ‚A) + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š, โˆ‚out2, mm_๐’โ‚‚_kron_t, 1, 1) + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] += โˆ‚โˆ‡โ‚โ‚Š * โ„’.I(n)[:,iโ‚Š] + println(" ||โˆ‚โˆ‡โ‚|| = ", โ„’.norm(Matrix(โˆ‚โˆ‡โ‚))) + + # โˆ‚๐‘บโ‚ + println(" [Pullback] Step 8: โˆ‚๐‘บโ‚ (most complex)") + โ„’.axpy!(1, โˆ‚L_c, โˆ‚S1S1_stack) + โ„’.axpy!(1, โˆ‚L_d, โˆ‚S1S1_stack) + + โˆ‚tmpkron22 = (โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ) + โˆ‚S1S1_from_ck = โ„‚.โˆ‚S1S1_from_ck_3rd; fill!(โˆ‚S1S1_from_ck, zero(S)) + โˆ‚S1p0_kron_sigma = โ„‚.โˆ‚S1p0_kron_sigma_3rd; fill!(โˆ‚S1p0_kron_sigma, zero(S)) + MM.compressed_permuted_mixed_kron_pullback!(โˆ‚S1S1_from_ck, + โˆ‚S1p0_kron_sigma, + โˆ‚tmpkron22, + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + S1p0_kron_sigma; + tol = opts.tol.droptol) + + โˆ‚S1p0_kron = (โˆ‚S1p0_kron_sigma * ๐›”t) + โˆ‚S1p0_left = โ„‚.โˆ‚S1p0_left_3rd; fill!(โˆ‚S1p0_left, zero(S)) + โˆ‚S1p0_right = โ„‚.โˆ‚S1p0_right_3rd; fill!(โˆ‚S1p0_right, zero(S)) + MM.fill_kron_adjoint!(โˆ‚S1p0_left, โˆ‚S1p0_right, โˆ‚S1p0_kron, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + + โ„’.axpy!(1, โˆ‚S1S1_from_ck, โˆ‚S1S1_stack) + โ„’.axpy!(1, โˆ‚S1p0_left, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + โ„’.axpy!(1, โˆ‚S1p0_right, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + + โˆ‚ck3_aux = collect(โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ) + MM.compressed_kronยณ_pullback!(โˆ‚aux, โˆ‚ck3_aux, aux) + โ„’.mul!(โˆ‚S1S1_stack, Mโ‚ƒ.๐’๐', โˆ‚aux, 1, 1) + + โ„’.axpy!(1, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + + MM.compressed_permuted_mixed_kron_pullback!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚๐›”_discard, โˆ‚B_from_sylv, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”; tol = opts.tol.droptol) + MM.compressed_kronยณ_pullback!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚B_from_sylv, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + + tmp_a = collect(MM.mat_mult_kron(collect(โˆ‡โ‚‚t_โˆ‚out2'), collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ))') + โˆ‚tmpkron2 = (tmp_a + Mโ‚ƒ๐โ‚โ‚—t * tmp_a * Mโ‚ƒ๐โ‚แตฃt) + MM.fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚๐›”_discard2, โˆ‚tmpkron2, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, collect(Mโ‚‚.๐›”)) + + โ„’.axpy!(1, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ) + + โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = (๐’โ‚‚t * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block) + MM.fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + + S2_padded = [๐’โ‚‚[iโ‚‹,:]; zeros(S, nโ‚‘ + 1, nโ‚‘โ‚‹^2)] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚top_block * S2_padded' + + nโ‚Šl = length(iโ‚Š) + โˆ‚top_S1S1 = โˆ‚S1S1_stack[1:nโ‚Šl, :] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚top_S1S1 * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ .+= ๐’โ‚' * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_S1S1 + @views โˆ‚๐’โ‚โ‚ƒ .+= โˆ‚S1S1_stack[nโ‚Šl .+ (1:n), :] + + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ[1:nโ‚Šl,:] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚‹,:] .+= โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ[1:length(iโ‚‹),:] + + โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,1:nโ‚‹] -= โˆ‡โ‚[:,1:nโ‚Š]' * โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] + + โˆ‚๐‘บโ‚ = [โˆ‚๐’โ‚โ‚ƒ[:,1:nโ‚‹] โˆ‚๐’โ‚โ‚ƒ[:,nโ‚‹+2:end]] + println(" ||โˆ‚๐‘บโ‚|| = ", โ„’.norm(Matrix(โˆ‚๐‘บโ‚))) + + # Map back to compressed space + println(" [Pullback] Step 9: compress gradients") + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ * ๐”โˆ‡โ‚‚t + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ * ๐”โ‚‚t + println(" ||โˆ‚โˆ‡โ‚‚_compressed|| = ", โ„’.norm(Matrix(โˆ‚โˆ‡โ‚‚))) + println(" ||โˆ‚๐’โ‚‚_compressed|| = ", โ„’.norm(Matrix(โˆ‚๐’โ‚‚))) + + return (NoTangent(), โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚ƒ, โˆ‚๐‘บโ‚, โˆ‚๐’โ‚‚, NoTangent(), NoTangent(), NoTangent()) +end + +println("\nStep 5 done โ€“ pullback function defined.") +println("Run: third_grads = third_order_solution_pullback(โˆ‚๐’โ‚ƒ_raw)") +println("Or with Debugger: @enter third_order_solution_pullback(โˆ‚๐’โ‚ƒ_raw)") + + +# ============================================================================== +# STEP 6: Run the inline pullback +# ============================================================================== +println("\nRunning inline pullback...") +@time third_grads = third_order_solution_pullback(โˆ‚๐’โ‚ƒ_raw) + +โˆ‚โˆ‡โ‚ = third_grads[2] +โˆ‚โˆ‡โ‚‚ = third_grads[3] +โˆ‚โˆ‡โ‚ƒ = third_grads[4] +โˆ‚๐‘บโ‚ = third_grads[5] +โˆ‚๐’โ‚‚ = third_grads[6] + +println("\nPullback complete. Gradient norms:") +println(" ||โˆ‚โˆ‡โ‚|| = ", โ„’.norm(Matrix(โˆ‚โˆ‡โ‚))) +println(" ||โˆ‚โˆ‡โ‚‚|| = ", โ„’.norm(Matrix(โˆ‚โˆ‡โ‚‚))) +println(" ||โˆ‚โˆ‡โ‚ƒ|| = ", โ„’.norm(Matrix(โˆ‚โˆ‡โ‚ƒ))) +println(" ||โˆ‚๐‘บโ‚|| = ", โ„’.norm(Matrix(โˆ‚๐‘บโ‚))) +println(" ||โˆ‚๐’โ‚‚|| = ", โ„’.norm(Matrix(โˆ‚๐’โ‚‚))) + + +# ============================================================================== +# STEP 7: Verify against rrule pullback +# ============================================================================== +println("\nVerifying against rrule pullback...") +rrule_grads = third_pb((โˆ‚๐’โ‚ƒ_raw, NoTangent())) + +labels = ("โˆ‚โˆ‡โ‚", "โˆ‚โˆ‡โ‚‚", "โˆ‚โˆ‡โ‚ƒ", "โˆ‚๐‘บโ‚", "โˆ‚๐’โ‚‚") +for (k, lab) in enumerate(labels) + manual_k = Matrix(third_grads[k+1]) + rrule_k = Matrix(rrule_grads[k+1]) + ฮ” = manual_k - rrule_k + max_abs = maximum(abs, ฮ”) + rel = โ„’.norm(ฮ”) / max(โ„’.norm(rrule_k), eps()) + println(" $lab: max|ฮ”|=$max_abs rel=$rel") +end diff --git a/benchmark/sw07_third_order_pullback_walkthrough.jl b/benchmark/sw07_third_order_pullback_walkthrough.jl new file mode 100644 index 000000000..42f803213 --- /dev/null +++ b/benchmark/sw07_third_order_pullback_walkthrough.jl @@ -0,0 +1,791 @@ +using Revise +using MacroModelling +using BenchmarkTools +using LinearAlgebra +using SparseArrays +using TimerOutputs: TimerOutput, @timeit +using ChainRulesCore: rrule, NoTangent + +const MM = MacroModelling +const LL = LinearAlgebra +const โ„’ = LinearAlgebra + +function _perm_source_to_target_from_columns(P) + n = size(P, 2) + map = zeros(Int, n) + if P isa SparseMatrixCSC + @inbounds for src in 1:n + for idx in P.colptr[src]:(P.colptr[src + 1] - 1) + if !iszero(P.nzval[idx]) + map[src] = P.rowval[idx] + break + end + end + end + else + @inbounds for src in 1:n + col = @view P[:, src] + dst = findfirst(!iszero, col) + map[src] = isnothing(dst) ? 0 : dst + end + end + return map +end + +function _accumulate_kron_A_entry!(โˆ‚A, Bฯƒ, row_idx::Int, col_idx::Int, val, + nrows::Int, n1::Int, n2::Int, m1::Int, + const_n1n2::Int, const_n1n2m1::Int) + linear_idx = (col_idx - 1) * nrows + row_idx + i = (linear_idx - 1) % n1 + 1 + k = ((linear_idx - 1) รท n1) % n2 + 1 + j = ((linear_idx - 1) รท const_n1n2) % m1 + 1 + l = ((linear_idx - 1) รท const_n1n2m1) + 1 + @inbounds โˆ‚A[k, l] += Bฯƒ[i, j] * val + return nothing +end + +include(joinpath(@__DIR__, "..", "models", "Smets_Wouters_2003.jl")) + +model = Smets_Wouters_2003 + +# include(joinpath(@__DIR__, "..", "models", "FS2000.jl")) + +# model = FS2000 + +parameters = copy(model.parameter_values) +opts = MM.merge_calculation_options(verbose = false) + +# Set to true to execute the pullback immediately. +# Keep false to step through the closure manually in REPL. +# run_pullback_now = false + +# ----------------------------------------------------------------------------- +# Step 0: Build exact inputs passed to calculate_third_order_solution +# ----------------------------------------------------------------------------- +MM.clear_solution_caches!(model, :third_order) + +# Initialize derivative/function caches for third-order path once. +_, _, _, _, solved_warmup = MM.get_solution(model, parameters, algorithm = :third_order, verbose = false) +@assert solved_warmup "Warmup third-order solve failed." +MM.clear_solution_caches!(model, :third_order) + +SS_and_pars, (solution_error, nsss_iters) = MM.get_NSSS_and_parameters(model, parameters, opts = opts) +@assert solution_error <= opts.tol.NSSS_acceptance_tol "NSSS solve did not satisfy acceptance tolerance." + +โˆ‡โ‚ = MM.calculate_jacobian(parameters, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces) + +๐’โ‚, qme_sol, solved1 = MM.calculate_first_order_solution(โˆ‡โ‚, + model.constants, + model.workspaces, + model.caches; + opts = opts, + initial_guess = model.caches.qme_solution) +@assert solved1 "First-order solution failed." + +โˆ‡โ‚‚ = MM.calculate_hessian(parameters, SS_and_pars, model.caches, model.functions.hessian, model.workspaces) + +๐’โ‚‚, solved2 = MM.calculate_second_order_solution(โˆ‡โ‚, + โˆ‡โ‚‚, + ๐’โ‚, + model.constants, + model.workspaces, + model.caches; + initial_guess = model.caches.second_order_solution, + opts = opts) +@assert solved2 "Second-order solution failed." + +โˆ‡โ‚ƒ = MM.calculate_third_order_derivatives(parameters, + SS_and_pars, + model.caches, + model.functions.third_order_derivatives, + model.workspaces) + +# ----------------------------------------------------------------------------- +# Step 1: Primal + pullback for calculate_third_order_solution +# ----------------------------------------------------------------------------- +third_out, third_pb = rrule(MM.calculate_third_order_solution, + โˆ‡โ‚, + โˆ‡โ‚‚, + โˆ‡โ‚ƒ, + ๐’โ‚, + ๐’โ‚‚, + model.constants, + model.workspaces, + model.caches; + initial_guess = model.caches.third_order_solution, + opts = opts) + +๐’โ‚ƒ_raw, solved3 = third_out +@assert solved3 "Third-order primal solve in rrule forward pass failed." + +# Objective from benchmark/bench.jl: +# norm(get_solution(model, x, algorithm = :third_order)[4] * model.constants.third_order.๐”โ‚ƒ) +๐’โ‚ƒ_full = ๐’โ‚ƒ_raw * model.constants.third_order.๐”โ‚ƒ +loss = LL.norm(๐’โ‚ƒ_full) + +# Seed cotangent for ๐’โ‚ƒ_raw from f(X) = norm(X * Uโ‚ƒ): +# โˆ‚f/โˆ‚X = (X*Uโ‚ƒ / norm(X*Uโ‚ƒ)) * Uโ‚ƒ' +scale = max(loss, eps(eltype(loss))) +โˆ‚๐’โ‚ƒ_raw_rr = (๐’โ‚ƒ_full / scale) * model.constants.third_order.๐”โ‚ƒ' + +println("third_order_solved=", solved3, + " size(๐’โ‚ƒ_raw)=", size(๐’โ‚ƒ_raw), + " nnz(๐’โ‚ƒ_raw)=", nnz(sparse(๐’โ‚ƒ_raw))) +println("loss_norm_S3_full=", loss) + +println("Ready to walk through the pullback closure.") +println("Manual call:") +println(" third_grads = third_pb((โˆ‚๐’โ‚ƒ_raw_rr, NoTangent()))") +println(" โˆ‚โˆ‡โ‚ = third_grads[2]; โˆ‚โˆ‡โ‚‚ = third_grads[3]; โˆ‚โˆ‡โ‚ƒ = third_grads[4]; โˆ‚๐’โ‚ = third_grads[5]; โˆ‚๐’โ‚‚ = third_grads[6]") + +# ----------------------------------------------------------------------------- +# Step 2: REPL-style manual chain from โˆ‚๐’โ‚ƒ_raw_rr to parameter tangents +# Mirrors pullback_3rd in rrules.jl for get_solution(..., algorithm=:third_order) +# ----------------------------------------------------------------------------- +estimation = true +nVar = length(model.constants.post_model_macro.var) + +nsss_out_rr, nsss_pb = rrule(MM.get_NSSS_and_parameters, + model, + parameters; + opts = opts, + estimation = estimation) +SS_and_pars_rr = nsss_out_rr[1] + +โˆ‡โ‚_rr, jac_pb = rrule(MM.calculate_jacobian, + parameters, + SS_and_pars_rr, + model.caches, + model.functions.jacobian, + model.workspaces) + +first_out_rr, first_pb = rrule(MM.calculate_first_order_solution, + โˆ‡โ‚_rr, + model.constants, + model.workspaces, + model.caches; + opts = opts, + initial_guess = model.caches.qme_solution) +๐’โ‚_rr = first_out_rr[1] + +โˆ‡โ‚‚_rr, hess_pb = rrule(MM.calculate_hessian, + parameters, + SS_and_pars_rr, + model.caches, + model.functions.hessian, + model.workspaces) + +second_out_rr, second_pb = rrule(MM.calculate_second_order_solution, + โˆ‡โ‚_rr, + โˆ‡โ‚‚_rr, + ๐’โ‚_rr, + model.constants, + model.workspaces, + model.caches; + initial_guess = model.caches.second_order_solution, + opts = opts) +๐’โ‚‚_raw_rr = second_out_rr[1] + +โˆ‡โ‚ƒ_rr, third_deriv_pb = rrule(MM.calculate_third_order_derivatives, + parameters, + SS_and_pars_rr, + model.caches, + model.functions.third_order_derivatives, + model.workspaces) + +# third_out_rr, third_pb_rr = rrule(MM.calculate_third_order_solution, +# โˆ‡โ‚_rr, +# โˆ‡โ‚‚_rr, +# โˆ‡โ‚ƒ_rr, +# ๐’โ‚_rr, +# ๐’โ‚‚_raw_rr, +# model.constants, +# model.workspaces, +# model.caches; +# initial_guess = model.caches.third_order_solution, +# opts = opts) +# ๐’โ‚ƒ_raw_rr = third_out_rr[1] +# @assert third_out_rr[2] "third_pb_rr forward pass failed." + +๐’โ‚ƒ_full_rr = ๐’โ‚ƒ_raw * model.constants.third_order.๐”โ‚ƒ +loss_rr = LL.norm(๐’โ‚ƒ_full_rr) +scale_rr = max(loss_rr, eps(eltype(loss_rr))) +โˆ‚๐’โ‚ƒ_raw_rr = (๐’โ‚ƒ_full_rr / scale_rr) * model.constants.third_order.๐”โ‚ƒ' + +println("manual-chain seed ready: norm(S3*U3)=", loss_rr) + +# Start here in REPL when stepping manually: +# โˆ‚๐’โ‚ƒ_raw_rr +pb_seed_rr = (โˆ‚๐’โ‚ƒ_raw_rr, NoTangent()) + +println("Pullback REPL entrypoint ready.") +println("Direct call:") +println(" third_grads_rr = third_pb_rr(pb_seed_rr)") + +# Bindings to run copied rrule body snippets directly in this script/REPL. +# These provide the same names used inside rrules.jl. +workspaces = model.workspaces +constants = model.constants +cache = model.caches +initial_guess = model.caches.third_order_solution + +S = eltype(โˆ‡โ‚_rr) +R = eltype(parameters) + +โˆ‡โ‚ = โˆ‡โ‚_rr +โˆ‡โ‚‚ = โˆ‡โ‚‚_rr +โˆ‡โ‚ƒ = โˆ‡โ‚ƒ_rr +๐‘บโ‚ = ๐’โ‚_rr +๐’โ‚‚ = ๐’โ‚‚_raw_rr + +Higher_order_workspace = MM.Higher_order_workspace +choose_matrix_format = MM.choose_matrix_format +ensure_higher_order_solution_buffers! = MM.ensure_higher_order_solution_buffers! +compressed_permuted_mixed_kron = MM.compressed_permuted_mixed_kron +compressed_kronยณ = MM.compressed_kronยณ +mat_mult_kron = MM.mat_mult_kron +fill_kron_adjoint! = MM.fill_kron_adjoint! +fill_kron_adjoint_โˆ‚A! = MM.fill_kron_adjoint_โˆ‚A! +solve_sylvester_equation = MM.solve_sylvester_equation +ensure_third_order_pullback_workspaces! = MM.ensure_third_order_pullback_workspaces! +compressed_permuted_mixed_kron_pullback! = MM.compressed_permuted_mixed_kron_pullback! +compressed_kronยณ_pullback! = MM.compressed_kronยณ_pullback! + +# ----------------------------------------------------------------------------- +# Full third_order_solution_pullback reference from +# src/custom_autodiff_rules/rrules.jl +# +# This is the full closure body so you can follow the same logic in this file +# while stepping from pb_seed_rr = (โˆ‚๐’โ‚ƒ_raw_rr, NoTangent()). +# ----------------------------------------------------------------------------- + + # --- workspace / constants --------------------------------------------------- + if !(eltype(workspaces.third_order.ลœ) == S) + workspaces.third_order = Higher_order_workspace(T = S) + end + โ„‚ = workspaces.third_order + Mโ‚‚ = constants.second_order + Mโ‚ƒ = constants.third_order + T = constants.post_model_macro + + # Expand compressed inputs to full space for internal computation + โˆ‡โ‚‚ = โˆ‡โ‚‚ * Mโ‚‚.๐”โˆ‡โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚ * Mโ‚‚.๐”โ‚‚)::SparseMatrixCSC{S, Int} + + iโ‚Š = T.future_not_past_and_mixed_idx + iโ‚‹ = T.past_not_future_and_mixed_idx + nโ‚‹ = T.nPast_not_future_and_mixed + nโ‚Š = T.nFuture_not_past_and_mixed + nโ‚‘ = T.nExo + n = T.nVars + nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + + ensure_higher_order_solution_buffers!(โ„‚, n, nโ‚‘โ‚‹) + + initial_guess_sylv = if length(initial_guess) == 0 + zeros(S, 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{S} ? initial_guess : Matrix{S}(initial_guess) + else + zeros(S, 0, 0) + end + + # --- forward pass (mirrors the primal, but stores intermediates) --------------- + + # 1st-order solution with zero-column + ๐’โ‚ = โ„‚.๐’โ‚::Matrix{S} + copyto!(@view(๐’โ‚[:,1:nโ‚‹]), @view(๐‘บโ‚[:,1:nโ‚‹])) + fill!(@view(๐’โ‚[:,nโ‚‹+1]), zero(S)) + copyto!(@view(๐’โ‚[:,nโ‚‹+2:end]), @view(๐‘บโ‚[:,nโ‚‹+1:end])) + + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} + copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹,:]), @view(๐’โ‚[iโ‚‹,:])) + fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1:end,:]), zero(S)) + @inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1,nโ‚‹+1] = one(S) + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] + ๐’โ‚ + โ„’.I(nโ‚‘โ‚‹)[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]] + + ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:]; zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)] + ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * Mโ‚‚.๐ˆโ‚™โ‚‹ - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] + + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = โ„’.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) + + if !โ„’.issuccess(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + return (โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + spinv = choose_matrix_format(spinv) + + โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š + + A = spinv * โˆ‡โ‚โ‚Š + + # --- B matrix ----------------------------------------------------------------- + kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + + B = compressed_permuted_mixed_kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”, + sparse_preallocation = โ„‚.tmp_sparse_prealloc7) + + B += compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc1) + + # --- ๐—โ‚ƒ (C-matrix ingredients) ----------------------------------------------- + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = @views [(๐’โ‚‚ * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)])[iโ‚Š,:] + ๐’โ‚‚ + zeros(nโ‚‹ + nโ‚‘, nโ‚‘โ‚‹^2)] + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, density_threshold = 0.0, min_length = 10, tol = opts.tol.droptol) + + ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚Š,:]; zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹^2)] + + aux = Mโ‚ƒ.๐’๐ * โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ + + S1p0_kron_sigma = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›” + tmpkron22 = compressed_permuted_mixed_kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + S1p0_kron_sigma, + sparse_preallocation = โ„‚.tmp_sparse_prealloc6) + + ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + + โˆ‡โ‚โ‚Š = choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + + ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚‹,:]; zeros(size(๐’โ‚)[2] - nโ‚‹, nโ‚‘โ‚‹^2)] + + # Terms (a)+(b): โˆ‡โ‚‚ * kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) * [tmpkron2 + ๐โ‚โ‚— * tmpkron2 * ๐โ‚แตฃ] * ๐๐‚โ‚ƒ + tmpkron2 = โ„’.kron(Mโ‚‚.๐›”, choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.droptol)) + D_ab = (tmpkron2 + Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ) * Mโ‚ƒ.๐๐‚โ‚ƒ + ๐—โ‚ƒ = mat_mult_kron(โˆ‡โ‚‚, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ), D_ab, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc2) + + # Term (c): โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, โŽธ๐’โ‚‚k..โŽน) * ๐๐‚โ‚ƒ + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc3) + + # Term (d): โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ*๐›”) * ๐๐‚โ‚ƒ + S2p0_sigma = ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›” + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, collect(S2p0_sigma), Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + + # Term (e): โˆ‡โ‚โ‚Š * ๐’โ‚‚ * kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) * ๐๐‚โ‚ƒ + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.droptol) + mm_๐’โ‚‚_kron = mat_mult_kron(๐’โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + ๐—โ‚ƒ += โˆ‡โ‚โ‚Š * mm_๐’โ‚‚_kron * Mโ‚ƒ.๐๐‚โ‚ƒ + + ๐—โ‚ƒ += โˆ‡โ‚ƒ * tmpkron22 + + # Compute compressed_kronยณ(aux) WITHOUT rowmask: the pullback needs โˆ‚โˆ‡โ‚ƒ at ALL + # positions (including currently-zero columns of โˆ‡โ‚ƒ) so that gradients flow + # correctly through calculate_third_order_derivatives back to parameters. + ck3_aux_mat = compressed_kronยณ(aux, rowmask = Mโ‚ƒ.โˆ‡โ‚ƒ_rowmask, tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc5) + ck3_aux = โˆ‡โ‚ƒ * ck3_aux_mat + ๐—โ‚ƒ += ck3_aux + + C = spinv * ๐—โ‚ƒ + + # --- solve Sylvester Aยท๐’โ‚ƒยทB + C = ๐’โ‚ƒ ---------------------------------------- + ๐’โ‚ƒ, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, + initial_guess = initial_guess_sylv, + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.sylvester_tol, + acceptance_tol = opts.tol.sylvester_acceptance_tol, + verbose = opts.verbose) + + ๐’โ‚ƒ = choose_matrix_format(๐’โ‚ƒ, multithreaded = false, tol = opts.tol.droptol) + ๐’โ‚ƒ_stable = copy(๐’โ‚ƒ) + + if !solved + return (๐’โ‚ƒ_stable, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # cache update (same as primal) + if ๐’โ‚ƒ_stable isa Matrix{S} && cache.third_order_solution isa Matrix{S} && size(cache.third_order_solution) == size(๐’โ‚ƒ_stable) + copyto!(cache.third_order_solution, ๐’โ‚ƒ_stable) + elseif ๐’โ‚ƒ_stable isa SparseMatrixCSC{S, Int} && cache.third_order_solution isa SparseMatrixCSC{S, Int} && + size(cache.third_order_solution) == size(๐’โ‚ƒ_stable) && + cache.third_order_solution.colptr == ๐’โ‚ƒ_stable.colptr && + cache.third_order_solution.rowval == ๐’โ‚ƒ_stable.rowval + copyto!(cache.third_order_solution.nzval, ๐’โ‚ƒ_stable.nzval) + else + cache.third_order_solution = ๐’โ‚ƒ_stable + end + + # --- precompute transposed constants for pullback ----------------------------- + # Use pre-cached transposes from constants (computed once at model compile time) + ๐๐‚โ‚ƒt = Mโ‚ƒ.๐๐‚โ‚ƒแต€ + ๐›”t = Mโ‚‚.๐›”แต€ + ๐”โˆ‡โ‚‚t = Mโ‚‚.๐”โˆ‡โ‚‚แต€ + ๐”โ‚‚t = Mโ‚‚.๐”โ‚‚แต€ + + # Use pre-cached transposes of permutation matrices (for out2 terms a,b pullback) + Mโ‚ƒ๐โ‚โ‚—t = Mโ‚ƒ.๐โ‚โ‚—แต€ + Mโ‚ƒ๐โ‚แตฃt = Mโ‚ƒ.๐โ‚แตฃแต€ + + # Materialized transposes of forward-pass intermediates + โˆ‡โ‚‚t = choose_matrix_format(โˆ‡โ‚‚') + โˆ‡โ‚ƒt = choose_matrix_format(โˆ‡โ‚ƒ') + D_ab_t = choose_matrix_format(D_ab') + tmpkron22_t = choose_matrix_format(tmpkron22') + ck3_aux_mat_t = choose_matrix_format(ck3_aux_mat') + ๐’โ‚‚t = choose_matrix_format(๐’โ‚‚', density_threshold = 1.0) + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t = choose_matrix_format(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹') + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ') + S2p0_sigma_t = choose_matrix_format(S2p0_sigma') + + mm_๐’โ‚‚_kron_t = choose_matrix_format(mm_๐’โ‚‚_kron') + + # --- ensure pullback workspace buffers --- + ensure_third_order_pullback_workspaces!(โ„‚, S, T, Mโ‚‚, Mโ‚ƒ) + + tmpkron22_ck3_aux_mat_t = choose_matrix_format(tmpkron22_t + ck3_aux_mat_t) + + + + +โˆ‚๐’โ‚ƒ_solved = pb_seed_rr + + +# @profview begin +pullback_timer = TimerOutput() +# for i in 1:10 +# function third_order_solution_pullback(โˆ‚๐’โ‚ƒ_solved) +@timeit pullback_timer "total" begin + โˆ‚๐’โ‚ƒ = โˆ‚๐’โ‚ƒ_solved[1] + + # --- adjoint Sylvester: Aแต€ โˆ‚C_adj Bแต€ + โˆ‚๐’โ‚ƒ = โˆ‚C_adj -------------------- + @timeit pullback_timer "adjoint_sylvester" begin + โˆ‚C_adj, slvd = solve_sylvester_equation(A', B', Matrix{Float64}(โˆ‚๐’โ‚ƒ), โ„‚.sylvester_workspace, + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.sylvester_tol, + acceptance_tol = opts.tol.sylvester_acceptance_tol, + verbose = opts.verbose) + + โˆ‚C_adj = choose_matrix_format(โˆ‚C_adj) + end + + # --- Initialize all gradient accumulators --- + @timeit pullback_timer "initialize_accumulators" begin + # Dense workspace temporaries (overwritten by mul! each call) + โˆ‚๐—โ‚ƒ = โ„‚.โˆ‚๐—โ‚ƒ_3rd + โˆ‚A = โ„‚.โˆ‚A_3rd + โˆ‚B_from_sylv = โ„‚.โˆ‚B_sylv_3rd + โˆ‚out2 = โ„‚.โˆ‚out2_3rd + โˆ‡โ‚‚t_โˆ‚out2 = โ„‚.โˆ‡โ‚‚t_โˆ‚out2_3rd + mul_tmp = โ„‚.mul_tmp_3rd + โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = โ„‚.โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd + + # Dense workspace accumulators (need zeroing) + โˆ‚spinv = โ„‚.โˆ‚spinv_3rd + โˆ‚โˆ‡โ‚ = โ„‚.โˆ‚โˆ‡โ‚_3rd; fill!(โˆ‚โˆ‡โ‚, zero(S)) + โˆ‚๐’โ‚โ‚ƒ = โ„‚.โˆ‚๐’โ‚_3rd; fill!(โˆ‚๐’โ‚โ‚ƒ, zero(S)) + + # Sparse-preserving gradient accumulators (reuse workspace buffers) + โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) # sparse โ€” must stay fresh + + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, zero(S)) + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, zero(S)) + โˆ‚L_c = โ„‚.โˆ‚L_c_3rd; fill!(โˆ‚L_c, zero(S)) + โˆ‚R_c = โ„‚.โˆ‚R_c_3rd; fill!(โˆ‚R_c, zero(S)) + โˆ‚L_d = โ„‚.โˆ‚L_d_3rd; fill!(โˆ‚L_d, zero(S)) + โˆ‚R_d = โ„‚.โˆ‚R_d_3rd; fill!(โˆ‚R_d, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8 = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, zero(S)) + โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, zero(S)) + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ, zero(S)) + โˆ‚S1S1_stack = โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd; fill!(โˆ‚S1S1_stack, zero(S)) + โˆ‚aux = โ„‚.โˆ‚aux_3rd; fill!(โˆ‚aux, zero(S)) + โˆ‚๐›”_discard = โ„‚.โˆ‚๐›”_discard_3rd; fill!(โˆ‚๐›”_discard, zero(S)) + end + + # --- gradient of A, B, C from ๐’โ‚ƒ = Aยท๐’โ‚ƒยทB + C --------------------------- + @timeit pullback_timer "backprop_A_B_C" begin + # โˆ‚A = โˆ‚C_adj * B' * ๐’โ‚ƒ_stable' โ€” use โˆ‚๐—โ‚ƒ as temp for intermediate + โ„’.mul!(โˆ‚๐—โ‚ƒ, โˆ‚C_adj, B') + โ„’.mul!(โˆ‚A, โˆ‚๐—โ‚ƒ, ๐’โ‚ƒ_stable') + # โˆ‚B_from_sylv = ๐’โ‚ƒ_stable' * A' * โˆ‚C_adj โ€” reuse โˆ‚๐—โ‚ƒ as temp + โ„’.mul!(โˆ‚๐—โ‚ƒ, A', โˆ‚C_adj) + โ„’.mul!(โˆ‚B_from_sylv, ๐’โ‚ƒ_stable', โˆ‚๐—โ‚ƒ) + # โˆ‚๐—โ‚ƒ = spinv' * โˆ‚C_adj + โˆ‚๐—โ‚ƒ = choose_matrix_format(spinv' * โˆ‚C_adj, density_threshold = 1.0, min_length = 0) + + # C = spinv * ๐—โ‚ƒ โ†’ โˆ‚spinv + # A = spinv * โˆ‡โ‚โ‚Š โ†’ โˆ‚spinv accumulation + โ„’.mul!(โˆ‚spinv, โˆ‚C_adj, ๐—โ‚ƒ') + โ„’.mul!(โˆ‚spinv, โˆ‚A, โˆ‡โ‚โ‚Š', 1, 1) + end + + # ===================================================================== + # โˆ‚โˆ‡โ‚ƒ (linear: โˆ‡โ‚ƒ appears in two additive terms of ๐—โ‚ƒ) + # ===================================================================== + @timeit pullback_timer "nabla3" begin + โˆ‚โˆ‡โ‚ƒ = โˆ‚๐—โ‚ƒ * tmpkron22_ck3_aux_mat_t + end + + # ===================================================================== + # โˆ‚โˆ‡โ‚‚ (โˆ‡โ‚‚ is linear in out2 โ†’ ๐—โ‚ƒ_pre โ†’ ๐—โ‚ƒ) + # ===================================================================== + @timeit pullback_timer "nabla2" begin + โ„’.mul!(โˆ‚out2, โˆ‚๐—โ‚ƒ, ๐๐‚โ‚ƒt) + + โˆ‚mid_ab = โˆ‚๐—โ‚ƒ * D_ab_t + โˆ‚โˆ‡โ‚‚ = mat_mult_kron(โˆ‚mid_ab, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ'), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ')) + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ + mat_mult_kron(โˆ‚out2, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt) + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ + mat_mult_kron(โˆ‚out2, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, S2p0_sigma_t) + end + + # ===================================================================== + # โˆ‚๐’โ‚‚ (๐’โ‚‚ enters out2 via several stacking matrices) + # ===================================================================== + @timeit pullback_timer "S2" begin + โ„’.mul!(โˆ‡โ‚‚t_โˆ‚out2, โˆ‡โ‚‚t, โˆ‚out2) + โˆ‚tmpkron1 = (โˆ‡โ‚‚t * โˆ‚mid_ab) + fill_kron_adjoint!(โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, โˆ‚tmpkron1, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] + + โˆ‚kron_c = (โˆ‡โ‚‚t_โˆ‚out2) + fill_kron_adjoint!(โˆ‚R_c, โˆ‚L_c, โˆ‚kron_c, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + + nโ‚Š_len = length(iโ‚Š) + โˆ‚top_block = โˆ‚R_c[1:nโ‚Š_len, :] + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚top_block * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + โˆ‚๐’โ‚‚_padded = ๐’โ‚' * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block + @views โˆ‚๐’โ‚‚[iโ‚‹,:] .+= โˆ‚๐’โ‚‚_padded[1:nโ‚‹, :] + @views โˆ‚๐’โ‚‚ .+= โˆ‚R_c[nโ‚Š_len .+ (1:n), :] + + fill_kron_adjoint!(โˆ‚R_d, โˆ‚L_d, โˆ‚kron_c, S2p0_sigma, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d = โˆ‚R_d * ๐›”t + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d[1:length(iโ‚Š),:] + + tmp_t8 = โˆ‡โ‚โ‚Š' * โˆ‚out2 + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ + mat_mult_kron(tmp_t8, collect(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘'), collect(๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ')) + โˆ‚kron_term8 = ((โˆ‡โ‚โ‚Š * ๐’โ‚‚)' * โˆ‚out2) + fill_kron_adjoint!(โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, โˆ‚kron_term8, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + @views โˆ‚๐’โ‚‚[iโ‚‹,:] .+= โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ[1:nโ‚‹,:] + end + + # ===================================================================== + # โˆ‚โˆ‡โ‚ + # ===================================================================== + @timeit pullback_timer "nabla1" begin + โ„’.mul!(mul_tmp, spinv', โˆ‚spinv) + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, mul_tmp, spinv') + โ„’.rmul!(โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, -1) + + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] * ๐’โ‚[iโ‚Š,1:nโ‚‹]' + โˆ‚โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ + + โˆ‚โˆ‡โ‚โ‚Š = โ„‚.โˆ‚โˆ‡โ‚โ‚Š_3rd + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š, spinv', โˆ‚A) + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š, โˆ‚out2, mm_๐’โ‚‚_kron_t, 1, 1) + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] += โˆ‚โˆ‡โ‚โ‚Š * โ„’.I(n)[:,iโ‚Š] + end + + # ===================================================================== + # โˆ‚๐‘บโ‚ + # ===================================================================== + @timeit pullback_timer "S1" begin + @timeit pullback_timer "seed_stack" begin + โ„’.axpy!(1, โˆ‚L_c, โˆ‚S1S1_stack) + โ„’.axpy!(1, โˆ‚L_d, โˆ‚S1S1_stack) + end + + @timeit pullback_timer "tmpkron22_pullback" begin + โˆ‚tmpkron22 = (โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ) + โˆ‚S1S1_from_ck = โ„‚.โˆ‚S1S1_from_ck_3rd + fill!(โˆ‚S1S1_from_ck, zero(S)) + โˆ‚S1p0_kron_sigma = โ„‚.โˆ‚S1p0_kron_sigma_3rd + fill!(โˆ‚S1p0_kron_sigma, zero(S)) + compressed_permuted_mixed_kron_pullback!(โˆ‚S1S1_from_ck, + โˆ‚S1p0_kron_sigma, + โˆ‚tmpkron22, + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + S1p0_kron_sigma; + tol = opts.tol.droptol) + end + + @timeit pullback_timer "S1p0_kron_adjoint" begin + โˆ‚S1p0_kron = (โˆ‚S1p0_kron_sigma * ๐›”t) + โˆ‚S1p0_left = โ„‚.โˆ‚S1p0_left_3rd + fill!(โˆ‚S1p0_left, zero(S)) + โˆ‚S1p0_right = โ„‚.โˆ‚S1p0_right_3rd + fill!(โˆ‚S1p0_right, zero(S)) + fill_kron_adjoint!(โˆ‚S1p0_left, โˆ‚S1p0_right, โˆ‚S1p0_kron, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + + โ„’.axpy!(1, โˆ‚S1S1_from_ck, โˆ‚S1S1_stack) + โ„’.axpy!(1, โˆ‚S1p0_left, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + โ„’.axpy!(1, โˆ‚S1p0_right, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + end + + @timeit pullback_timer "ck3_aux_pullback" begin + โˆ‚ck3_aux = collect(โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ) + compressed_kronยณ_pullback!(โˆ‚aux, โˆ‚ck3_aux, aux) + โ„’.mul!(โˆ‚S1S1_stack, Mโ‚ƒ.๐’๐', โˆ‚aux, 1, 1) + + โ„’.axpy!(1, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + end + + @timeit pullback_timer "B_pullback" begin + compressed_permuted_mixed_kron_pullback!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚๐›”_discard, โˆ‚B_from_sylv, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”; tol = opts.tol.droptol) + compressed_kronยณ_pullback!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚B_from_sylv, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + end + + @timeit pullback_timer "nabla2_cross_term" begin + @timeit pullback_timer "build_tmp_a" begin + Gt = sparse(โˆ‡โ‚‚t_โˆ‚out2') + B1 = collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + C1 = collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) + + n_rowB = size(B1, 1) + n_colB = size(B1, 2) + n_rowC = size(C1, 1) + n_colC = size(C1, 2) + nrows_tmp = n_colB * n_colC + + Bฯƒ = collect(Mโ‚‚.๐›”) + n1, m1 = size(Bฯƒ) + n2 = size(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, 1) + const_n1n2 = n1 * n2 + const_n1n2m1 = n1 * n2 * m1 + + row_map = _perm_source_to_target_from_columns(Mโ‚ƒ๐โ‚โ‚—t) + col_map = _perm_source_to_target_from_columns(Mโ‚ƒ๐โ‚แตฃt') + + Aฬ„ = zeros(S, n_rowC, n_rowB) + Aฬ„B = zeros(S, n_rowC, n_colB) + CAฬ„B = zeros(S, n_colC, n_colB) + + rv = Gt isa SparseMatrixCSC ? Gt.rowval : Gt.A.rowval + active_rows = unique(rv) + for src_col in active_rows + @views copyto!(Aฬ„, Gt[src_col, :]) + โ„’.mul!(Aฬ„B, Aฬ„, B1) + โ„’.mul!(CAฬ„B, C1', Aฬ„B) + for tmp_row in eachindex(CAฬ„B) + val = CAฬ„B[tmp_row] + abs(val) > eps(S) || continue + + _accumulate_kron_A_entry!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, Bฯƒ, tmp_row, src_col, val, + nrows_tmp, n1, n2, m1, + const_n1n2, const_n1n2m1) + + perm_row = row_map[tmp_row] + perm_col = col_map[src_col] + _accumulate_kron_A_entry!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, Bฯƒ, perm_row, perm_col, val, + nrows_tmp, n1, n2, m1, + const_n1n2, const_n1n2m1) + end + end + end + + @timeit pullback_timer "axpy_t8" begin + โ„’.axpy!(1, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ) + end + + @timeit pullback_timer "top_block_kron" begin + โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = (๐’โ‚‚t * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block) + fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + end + + @timeit pullback_timer "final_assembly" begin + S2_padded = [๐’โ‚‚[iโ‚‹,:]; zeros(S, nโ‚‘ + 1, nโ‚‘โ‚‹^2)] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚top_block * S2_padded' + + nโ‚Šl = length(iโ‚Š) + โˆ‚top_S1S1 = โˆ‚S1S1_stack[1:nโ‚Šl, :] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚top_S1S1 * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ .+= ๐’โ‚' * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_S1S1 + @views โˆ‚๐’โ‚โ‚ƒ .+= โˆ‚S1S1_stack[nโ‚Šl .+ (1:n), :] + + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ[1:nโ‚Šl,:] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚‹,:] .+= โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ[1:length(iโ‚‹),:] + โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,1:nโ‚‹] -= โˆ‡โ‚[:,1:nโ‚Š]' * โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] + + โˆ‚๐‘บโ‚ = [โˆ‚๐’โ‚โ‚ƒ[:,1:nโ‚‹] โˆ‚๐’โ‚โ‚ƒ[:,nโ‚‹+2:end]] + end + end + end + + # Map โˆ‚โˆ‡โ‚‚ and โˆ‚๐’โ‚‚ back to compressed space + @timeit pullback_timer "compress_outputs" begin + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ * ๐”โˆ‡โ‚‚t + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ * ๐”โ‚‚t + end + + manual_third_pullback_grads = (NoTangent(), โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚ƒ, โˆ‚๐‘บโ‚, โˆ‚๐’โ‚‚, NoTangent(), NoTangent(), NoTangent()) + end +# end +# end +pullback_timer + +# Actual pullback execution for calculate_third_order_solution rrule. +# This runs the real closure code from src/custom_autodiff_rules/rrules.jl. +# Start from pb_seed_rr (which contains โˆ‚๐’โ‚ƒ_raw_rr) and inspect each object below. + +third_grads_rr = third_pb_rr(pb_seed_rr) + +โˆ‚โˆ‡โ‚_from_3rd_rr = third_grads_rr[2] +โˆ‚โˆ‡โ‚‚_from_3rd_rr = third_grads_rr[3] +โˆ‚โˆ‡โ‚ƒ_from_3rd_rr = third_grads_rr[4] +โˆ‚๐’โ‚_from_3rd_rr = third_grads_rr[5] +โˆ‚๐’โ‚‚_from_3rd_rr = third_grads_rr[6] + +โˆ‚parameters_manual = zeros(eltype(parameters), length(parameters)) +โˆ‚SS_and_pars_manual = zeros(eltype(parameters), length(SS_and_pars_rr)) + +third_deriv_grads_rr = third_deriv_pb(โˆ‚โˆ‡โ‚ƒ_from_3rd_rr) +โˆ‚parameters_manual .+= third_deriv_grads_rr[2] +โˆ‚SS_and_pars_manual .+= third_deriv_grads_rr[3] + +โˆ‚๐’โ‚‚_total_rr = Matrix(โˆ‚๐’โ‚‚_from_3rd_rr) +second_grads_rr = second_pb((โˆ‚๐’โ‚‚_total_rr, NoTangent())) +โˆ‚โˆ‡โ‚_from_2nd_rr = second_grads_rr[2] +โˆ‚โˆ‡โ‚‚_from_2nd_rr = second_grads_rr[3] +โˆ‚๐’โ‚_from_2nd_rr = second_grads_rr[4] + +โˆ‚โˆ‡โ‚‚_total_rr = โˆ‚โˆ‡โ‚‚_from_3rd_rr + โˆ‚โˆ‡โ‚‚_from_2nd_rr +hess_grads_rr = hess_pb(โˆ‚โˆ‡โ‚‚_total_rr) +โˆ‚parameters_manual .+= hess_grads_rr[2] +โˆ‚SS_and_pars_manual .+= hess_grads_rr[3] + +โˆ‚๐’โ‚_total_rr = โˆ‚๐’โ‚_from_3rd_rr + โˆ‚๐’โ‚_from_2nd_rr +first_grads_rr = first_pb((โˆ‚๐’โ‚_total_rr, NoTangent(), NoTangent())) + +โˆ‚โˆ‡โ‚_total_rr = โˆ‚โˆ‡โ‚_from_3rd_rr + โˆ‚โˆ‡โ‚_from_2nd_rr + first_grads_rr[2] +jac_grads_rr = jac_pb(โˆ‚โˆ‡โ‚_total_rr) +โˆ‚parameters_manual .+= jac_grads_rr[2] +โˆ‚SS_and_pars_manual .+= jac_grads_rr[3] + +nsss_grads_rr = nsss_pb((โˆ‚SS_and_pars_manual, NoTangent())) +โˆ‚parameters_manual .+= nsss_grads_rr[3] + +println("manual_chain parameter tangent norm=", LL.norm(โˆ‚parameters_manual)) +println("\nTimerOutputs report for manual third_order_solution_pullback walkthrough:") +show(pullback_timer) +println() + +# ----------------------------------------------------------------------------- +# Step 3: Compare with real pullback of bench objective path +# bench objective path: norm(get_solution(model, x, algorithm=:third_order)[4] * Uโ‚ƒ) +# ----------------------------------------------------------------------------- +sol_out_rr, sol_pb_rr = rrule(MM.get_solution, + model, + parameters; + algorithm = :third_order, + verbose = false) + +๐’โ‚ƒ_sol_raw = sol_out_rr[4] +๐’โ‚ƒ_sol_full = ๐’โ‚ƒ_sol_raw * model.constants.third_order.๐”โ‚ƒ +loss_sol = LL.norm(๐’โ‚ƒ_sol_full) +scale_sol = max(loss_sol, eps(eltype(loss_sol))) +โˆ‚๐’โ‚ƒ_sol_raw = (๐’โ‚ƒ_sol_full / scale_sol) * model.constants.third_order.๐”โ‚ƒ' + +sol_grads_rr = sol_pb_rr((NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’โ‚ƒ_sol_raw, NoTangent())) +โˆ‚parameters_real = sol_grads_rr[3] + +ฮ”p = โˆ‚parameters_manual - โˆ‚parameters_real +max_abs_diff_params = maximum(abs, ฮ”p) +rel_diff_params = LL.norm(ฮ”p) / max(LL.norm(โˆ‚parameters_real), eps(eltype(loss_sol))) + +println("real_pullback parameter tangent norm=", LL.norm(โˆ‚parameters_real)) +println("manual_vs_real params: max_abs_diff=", max_abs_diff_params, + " rel_diff=", rel_diff_params) diff --git a/benchmark/sw07_third_order_walkthrough.jl b/benchmark/sw07_third_order_walkthrough.jl new file mode 100644 index 000000000..fd3fe1c49 --- /dev/null +++ b/benchmark/sw07_third_order_walkthrough.jl @@ -0,0 +1,264 @@ +using Revise +using MacroModelling +using BenchmarkTools +using LinearAlgebra +using SparseArrays + +const MM = MacroModelling +const LL = LinearAlgebra + +include(joinpath(@__DIR__, "..", "models", "Smets_Wouters_2007.jl")) + +model = Smets_Wouters_2007 +parameters = copy(model.parameter_values) +opts = MM.merge_calculation_options(verbose = false) + +# ----------------------------------------------------------------------------- +# Step 0: Build the exact inputs passed to calculate_third_order_solution +# ----------------------------------------------------------------------------- +MM.clear_solution_caches!(model, :third_order) + +# Initialize derivative/function caches for third-order path once. +_, _, _, _, solved_warmup = MM.get_solution(model, parameters, algorithm = :third_order, verbose = false) +@assert solved_warmup "Warmup third-order solve failed." +MM.clear_solution_caches!(model, :third_order) + +SS_and_pars, (solution_error, nsss_iters) = MM.get_NSSS_and_parameters(model, parameters, opts = opts) +@assert solution_error <= opts.tol.NSSS_acceptance_tol "NSSS solve did not satisfy acceptance tolerance." + +โˆ‡โ‚ = MM.calculate_jacobian(parameters, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces) + +๐’โ‚, qme_sol, solved1 = MM.calculate_first_order_solution(โˆ‡โ‚, + model.constants, + model.workspaces, + model.caches; + opts = opts, + initial_guess = model.caches.qme_solution) +@assert solved1 "First-order solution failed." + +โˆ‡โ‚‚ = MM.calculate_hessian(parameters, SS_and_pars, model.caches, model.functions.hessian, model.workspaces) + +๐’โ‚‚, solved2 = MM.calculate_second_order_solution(โˆ‡โ‚, + โˆ‡โ‚‚, + ๐’โ‚, + model.constants, + model.workspaces, + model.caches; + initial_guess = model.caches.second_order_solution, + opts = opts) +@assert solved2 "Second-order solution failed." + +โˆ‡โ‚ƒ = MM.calculate_third_order_derivatives(parameters, + SS_and_pars, + model.caches, + model.functions.third_order_derivatives, + model.workspaces) + +โˆ‡โ‚‚_input = copy(โˆ‡โ‚‚) +๐’โ‚‚_input = copy(๐’โ‚‚) + +# Inputs you asked for (passed to calculate_third_order_solution): +# โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, model.constants, model.workspaces, model.caches + +# ----------------------------------------------------------------------------- +# Step 1: Plain code from calculate_third_order_solution (primal) +# ----------------------------------------------------------------------------- +S = eltype(โˆ‡โ‚) +if !(eltype(model.workspaces.third_order.Sฬ‚) == S) + model.workspaces.third_order = MM.Higher_order_workspace(T = S) +end + +โ„‚ = model.workspaces.third_order +Mโ‚‚ = model.constants.second_order +Mโ‚ƒ = model.constants.third_order +T = model.constants.post_model_macro + +# Expand compressed hessian to full space +โˆ‡โ‚‚ = โˆ‡โ‚‚ * Mโ‚‚.๐”โˆ‡โ‚‚ + +# Expand compressed second-order solution to full space +๐’โ‚‚ = sparse(๐’โ‚‚ * Mโ‚‚.๐”โ‚‚) + +# Indices and dimensions +iโ‚Š = T.future_not_past_and_mixed_idx +iโ‚‹ = T.past_not_future_and_mixed_idx + +nโ‚‹ = T.nPast_not_future_and_mixed +nโ‚Š = T.nFuture_not_past_and_mixed +nโ‚‘ = T.nExo +n = T.nVars +nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + +MM.ensure_higher_order_solution_buffers!(โ„‚, n, nโ‚‘โ‚‹) + +initial_guess = model.caches.third_order_solution +initial_guess_sylv = if length(initial_guess) == 0 + zeros(S, 0, 0) +elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{S} ? initial_guess : Matrix{S}(initial_guess) +else + zeros(S, 0, 0) +end + +# 1st order solution embedding +๐’โ‚buf = โ„‚.๐’โ‚::Matrix{S} +copyto!(@view(๐’โ‚buf[:, 1:nโ‚‹]), @view(๐’โ‚[:, 1:nโ‚‹])) +fill!(@view(๐’โ‚buf[:, nโ‚‹ + 1]), zero(S)) +copyto!(@view(๐’โ‚buf[:, nโ‚‹ + 2:end]), @view(๐’โ‚[:, nโ‚‹ + 1:end])) + +๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} +copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹, :]), @view(๐’โ‚buf[iโ‚‹, :])) +fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹ + 1:end, :]), zero(S)) +@inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹ + 1, nโ‚‹ + 1] = one(S) + +๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = MM.choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + +โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [ + (๐’โ‚buf * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š, :] + ๐’โ‚buf + LL.I(nโ‚‘โ‚‹)[[range(1, nโ‚‹)..., nโ‚‹ + 1 .+ range(1, nโ‚‘)...], :] +] + +๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [ + ๐’โ‚buf[iโ‚Š, :] + zeros(S, nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹) +] +๐’โ‚โ‚Šโ•ฑ๐ŸŽ = MM.choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + +โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:, 1:nโ‚Š] * ๐’โ‚buf[iโ‚Š, 1:nโ‚‹] * LL.I(n)[iโ‚‹, :] - โˆ‡โ‚[:, range(1, n) .+ nโ‚Š] + +โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = LL.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) +if !LL.issuccess(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + error("Third-order setup failed: LU factorization of โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ was unsuccessful.") +end + +โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:, 1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š +A = โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu \ โˆ‡โ‚โ‚Š + +B = MM.compressed_permuted_mixed_kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”, sparse_preallocation = โ„‚.tmp_sparse_prealloc7) +B += MM.compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc1) + +โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = @views [ + (๐’โ‚‚ * LL.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + ๐’โ‚buf * [๐’โ‚‚[iโ‚‹, :] ; zeros(S, nโ‚‘ + 1, nโ‚‘โ‚‹^2)])[iโ‚Š, :] + ๐’โ‚‚ + zeros(S, nโ‚‹ + nโ‚‘, nโ‚‘โ‚‹^2) +] +โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = MM.choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, + density_threshold = 0.0, + min_length = 10, + tol = opts.tol.droptol) + +๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [ + ๐’โ‚‚[iโ‚Š, :] + zeros(S, nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹^2) +] + +aux = Mโ‚ƒ.๐’๐ * โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ + +๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = MM.choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) +โˆ‡โ‚โ‚Š = MM.choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + +๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚‹, :] ; zeros(S, size(๐’โ‚buf, 2) - nโ‚‹, nโ‚‘โ‚‹^2)] + +# Terms (a)+(b) +tmpkron2_sp = LL.kron(Mโ‚‚.๐›”, MM.choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.droptol)) +D_ab = (tmpkron2_sp + Mโ‚ƒ.๐โ‚โ‚— * tmpkron2_sp * Mโ‚ƒ.๐โ‚แตฃ) * Mโ‚ƒ.๐๐‚โ‚ƒ + +๐—โ‚ƒ = MM.mat_mult_kron(โˆ‡โ‚‚, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ), D_ab, + sparse = true, + sparse_preallocation = โ„‚.tmp_sparse_prealloc2) + +# Term (c) +๐—โ‚ƒ += MM.mat_mult_kron(โˆ‡โ‚‚, + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, + Mโ‚ƒ.๐๐‚โ‚ƒ, + sparse = true, + sparse_preallocation = โ„‚.tmp_sparse_prealloc3) + +# Term (d) +๐—โ‚ƒ += MM.mat_mult_kron(โˆ‡โ‚‚, + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›”), + Mโ‚ƒ.๐๐‚โ‚ƒ, + sparse = true, + sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + +# Term (e) +๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = MM.choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.droptol) +๐—โ‚ƒ += MM.mat_mult_kron(โˆ‡โ‚โ‚Š * ๐’โ‚‚, + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, + ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, + Mโ‚ƒ.๐๐‚โ‚ƒ, + sparse = true) + +# Mixed โˆ‡โ‚ƒ term +if length(โ„‚.tmpkron0) > 0 && eltype(โ„‚.tmpkron0) == S + LL.kron!(โ„‚.tmpkron0, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) +else + โ„‚.tmpkron0 = LL.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) +end + +โ„‚.tmpkron0 *= Mโ‚‚.๐›” + +tmpkron22 = MM.compressed_permuted_mixed_kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + โ„‚.tmpkron0, + sparse_preallocation = โ„‚.tmp_sparse_prealloc6) +๐—โ‚ƒ += โˆ‡โ‚ƒ * tmpkron22 + +# Cubic โˆ‡โ‚ƒ term +๐—โ‚ƒ += โˆ‡โ‚ƒ * MM.compressed_kronยณ( aux, + rowmask = Mโ‚ƒ.โˆ‡โ‚ƒ_rowmask, + tol = opts.tol.droptol, + sparse_preallocation = โ„‚.tmp_sparse_prealloc5) + +C = โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu \ ๐—โ‚ƒ + +๐’โ‚ƒ, solved3 = MM.solve_sylvester_equation(A, + B, + C, + โ„‚.sylvester_workspace, + initial_guess = initial_guess_sylv, + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.sylvester_tol, + acceptance_tol = opts.tol.sylvester_acceptance_tol, + verbose = opts.verbose) + +๐’โ‚ƒ = MM.choose_matrix_format(๐’โ‚ƒ, multithreaded = false, tol = opts.tol.droptol) + +if solved3 + if ๐’โ‚ƒ isa Matrix{S} && model.caches.third_order_solution isa Matrix{S} && size(model.caches.third_order_solution) == size(๐’โ‚ƒ) + copyto!(model.caches.third_order_solution, ๐’โ‚ƒ) + elseif ๐’โ‚ƒ isa SparseMatrixCSC{S, Int} && model.caches.third_order_solution isa SparseMatrixCSC{S, Int} && + size(model.caches.third_order_solution) == size(๐’โ‚ƒ) && + model.caches.third_order_solution.colptr == ๐’โ‚ƒ.colptr && + model.caches.third_order_solution.rowval == ๐’โ‚ƒ.rowval + copyto!(model.caches.third_order_solution.nzval, ๐’โ‚ƒ.nzval) + else + model.caches.third_order_solution = copy(๐’โ‚ƒ) + end +end + +println("third_order_solved=", solved3, " size(๐’โ‚ƒ)=", size(๐’โ‚ƒ), " nnz(๐’โ‚ƒ)=", nnz(sparse(๐’โ‚ƒ))) + +# ----------------------------------------------------------------------------- +# Step 2: Check against calculate_third_order_solution output +# ----------------------------------------------------------------------------- +๐’โ‚ƒ_ref, solved3_ref = MM.calculate_third_order_solution(โˆ‡โ‚, + โˆ‡โ‚‚_input, + โˆ‡โ‚ƒ, + ๐’โ‚, + ๐’โ‚‚_input, + model.constants, + model.workspaces, + model.caches; + initial_guess = zeros(eltype(โˆ‡โ‚), 0, 0), + opts = opts) + +ฮ” = Matrix(๐’โ‚ƒ) - Matrix(๐’โ‚ƒ_ref) +max_abs_diff = maximum(abs, ฮ”) +rel_diff = norm(ฮ”) / max(norm(Matrix(๐’โ‚ƒ_ref)), eps()) + +println("third_order_ref_solved=", solved3_ref, + " max_abs_diff=", max_abs_diff, + " rel_diff=", rel_diff) diff --git a/docs/agent-guides/development-workflow.md b/docs/agent-guides/development-workflow.md new file mode 100644 index 000000000..0564c312e --- /dev/null +++ b/docs/agent-guides/development-workflow.md @@ -0,0 +1,201 @@ +# Development Workflow (On-Demand) + +Read this file only when setup, runtime workflow, testing, docs, or benchmarking details are needed. + +## Julia Setup + +- Julia version: 1.10+ +- Run Julia with threads enabled: `julia -t auto` +- If Julia is not on PATH (Linux), check `~/.juliaup/bin/julia` + +### Environment setup + +```julia +using Pkg +Pkg.activate(".") +Pkg.instantiate() +``` + +If packages are missing, install them first (for example with `Pkg.add(...)`). + +## Revise-Based Iteration (Required for Interactive Work) + +Always use Revise for iterative development. **Never use one-shot `julia -e` commands** โ€” they discard the session and force full recompilation on every call. + +### Persistent REPL via Named Pipe (for AI Agents) + +AI agents cannot type into a REPL interactively. Use a named-pipe pattern to maintain a persistent Julia session across tool calls. + +#### 1. Start the session (once per conversation) + +Use `.julia_repl/` inside the project directory (already in `.gitignore`) instead of `/tmp/` to avoid VS Code trusted-folder approval prompts. + +```bash +# Create infrastructure (inside the project โ€” no approval needed) +mkdir -p .julia_repl +rm -f .julia_repl/pipe .julia_repl/out +mkfifo .julia_repl/pipe +touch .julia_repl/out + +# Start Julia reading from pipe (background process) +tail -f .julia_repl/pipe | julia -t auto --project=. 2>&1 | tee .julia_repl/out & +``` + +Start this with `isBackground=true` so the terminal stays alive. + +#### 2. Load packages (once) + +```bash +: > .julia_repl/out && echo 'using Revise; using MacroModelling; println("REPL_READY")' > .julia_repl/pipe +for i in {1..60}; do grep -q "REPL_READY" .julia_repl/out && break; sleep 1; done; tail -5 .julia_repl/out +``` + +The polling loop checks every second for `REPL_READY` and exits immediately when found (timeout: 60s). Package loading typically takes 10-30 seconds. + +#### 3. Execute code + +**Preferred method** โ€” write code to a file, then include it: + +```bash +# Step A: Write Julia code to a .jl file (using create_file tool โ€” no terminal command needed) +# File: tasks/_repl_cmd.jl +# IMPORTANT: End the file with println("DONE") as a sentinel marker. + +# Step B: Clear output, run it, and poll for the sentinel +: > .julia_repl/out && echo 'include("tasks/_repl_cmd.jl")' > .julia_repl/pipe +for i in {1..600}; do grep -q "DONE" .julia_repl/out && break; sleep 1; done; tail -20 .julia_repl/out +``` + +**For short one-liners**, send directly: + +```bash +: > .julia_repl/out && echo 'println(1 + 1); println("DONE")' > .julia_repl/pipe +for i in {1..120}; do grep -q "DONE" .julia_repl/out && break; sleep 1; done; tail -5 .julia_repl/out +``` + +#### 4. Sentinel-based completion detection + +Always end code with a sentinel `println` (e.g., `println("DONE")`). Use a polling loop to wait for it instead of fixed `sleep` durations: + +```bash +# Pattern: clear output, send command, poll for sentinel, read result +: > .julia_repl/out && echo '...; println("DONE")' > .julia_repl/pipe +for i in {1..TIMEOUT}; do grep -q "DONE" .julia_repl/out && break; sleep 1; done; tail -20 .julia_repl/out +``` + +Choose TIMEOUT based on expected work: +- Package loading / first compilation: `120` +- Warm cached calls: `30` +- Simple one-liners: `10` + +If the sentinel is not found within the timeout, check `.julia_repl/out` for errors. + +#### 5. Key rules + +- **Always use sentinel markers** โ€” end every code block with `println("STEP_NAME_DONE")` so the polling loop can detect completion. +- **Always clear output first** โ€” run `: > .julia_repl/out` before each command to avoid matching stale sentinels. +- **Poll, don't sleep** โ€” use `for i in {1..N}; do grep -q "SENTINEL" .julia_repl/out && break; sleep 1; done` instead of fixed `sleep` durations. This returns as soon as the task finishes. +- **The session persists** โ€” variables, models, compiled methods all survive between `echo` commands. This is the whole point. +- **Revise picks up edits** โ€” after editing `src/` files with the editor tool, the running session sees the changes automatically. +- **For test project deps**, use `--project=test` instead of `--project=.` when tests need extra packages (Zygote, Turing, etc.). +- **To reset the session**, send `exit()` to the pipe, wait for the process to end, then re-run steps 1-2: + ```bash + echo 'exit()' > .julia_repl/pipe + for i in {1..10}; do jobs -l 2>/dev/null | grep -q julia || break; sleep 1; done + rm -f .julia_repl/pipe .julia_repl/out && mkfifo .julia_repl/pipe && touch .julia_repl/out + # Then restart with tail -f ... & and reload packages + ``` + +### Human Developer REPL Setup + +1. Start one REPL and keep it running: + +```bash +cd /path/to/MacroModelling.jl +julia -t auto --project=. +``` + +2. In the REPL, load Revise before MacroModelling: + +```julia +using Revise +using MacroModelling +``` + +3. Edit source files and run code in the same session. + +### Why + +- Avoids repeated precompilation cost (minutes per call โ†’ zero) +- Preserves session/model state between edits +- Enables rapid edit-test-fix loops + +### Caveats + +- Structural changes (new type layouts, module reorganization, `__init__` changes) may require restart +- If updates are missed, run `Revise.revise()` + +## Quick Testing Strategy + +Do not run the full test suite for normal iteration. + +### Preferred approach + +- Use a bespoke script or quick reproduction with a small model +- Validate only the impacted behavior first + +Example RBC model for lightweight checks: + +```julia +@model RBC 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]^ฮฑ + z[0] = ฯ * z[-1] + std_z * eps_z[x] +end + +@parameters RBC begin + std_z = 0.01 + ฯ = 0.2 + ฮด = 0.02 + ฮฑ = 0.5 + ฮฒ = 0.95 +end + +get_irf(RBC) +simulate(RBC) +``` + +## CI Test Sets (Reference) + +Only use targeted sets when needed: + +- `basic`, `estimation`, `higher_order_1-3`, `plots_1-5`, `estimate_sw07`, `jet` +- Estimation sets: `1st_order_inversion_estimation`, `2nd_order_estimation`, `pruned_2nd_order_estimation`, `3rd_order_estimation`, `pruned_3rd_order_estimation` +- Pigeons estimation sets: `estimation_pigeons`, `1st_order_inversion_estimation_pigeons`, `2nd_order_estimation_pigeons`, `pruned_2nd_order_estimation_pigeons`, `3rd_order_estimation_pigeons`, `pruned_3rd_order_estimation_pigeons` + +```bash +TEST_SET=basic julia --project -e 'using Pkg; Pkg.test()' +``` + +Test environment setup: + +```julia +using Pkg +Pkg.activate("test") +Pkg.instantiate() +``` + +## Documentation Build + +```bash +julia --project=docs docs/make.jl +``` + +## Benchmarking + +```julia +using BenchmarkTools +include("benchmark/benchmarks.jl") +run(SUITE) +``` diff --git a/docs/agent-guides/project-context.md b/docs/agent-guides/project-context.md new file mode 100644 index 000000000..e7c0e043c --- /dev/null +++ b/docs/agent-guides/project-context.md @@ -0,0 +1,57 @@ +# Project Context (On-Demand) + +Read this file only when project background or codebase orientation is needed. + +## Overview + +`MacroModelling.jl` is a Julia package for developing and solving dynamic stochastic general equilibrium (DSGE) models. + +Key capabilities: + +- Parse models with time-indexed syntax (`[0]`, `[-1]`, `[1]`) +- Solve models automatically from equations and parameters +- Compute first-, second-, and third-order (pruned) perturbation solutions +- Handle occasionally binding constraints +- Compute IRFs, simulations, and conditional forecasts +- Estimate models using gradient-based samplers (NUTS/HMC) or inversion filters +- Differentiate solutions and moments w.r.t. parameters + +Target audience: central banks, regulators, graduate students, and researchers. + +Timing convention: end-of-period. + +## High-Level Repository Structure + +```text +MacroModelling.jl/ +โ”œโ”€โ”€ src/ # Core package code +โ”œโ”€โ”€ test/ # Test suite +โ”œโ”€โ”€ models/ # Example DSGE models +โ”œโ”€โ”€ docs/ # Documenter-based docs +โ”œโ”€โ”€ benchmark/ # Benchmark scripts +โ””โ”€โ”€ ext/ # Package extensions +``` + +Common files in `src/`: + +- `MacroModelling.jl` (module/exports/types) +- `macros.jl` (`@model`, `@parameters`) +- `get_functions.jl` (user-facing API) +- `perturbation.jl` (1st-3rd order solvers) +- `moments.jl`, `structures.jl`, `options_and_caches.jl` +- `dynare.jl`, `inspect.jl`, `solver_parameters.jl`, `default_options.jl` +- `algorithms/`, `filter/`, `custom_autodiff_rules/` + +## Model Syntax Quick Reference + +- Variables use time indices: `...[2], [1], [0], [-1], [-2]...` +- Shocks use `[x]`: `eps_z[x]` +- Calibration equations use `|` in `@parameters` +- Custom steady state can be provided via `steady_state_function` + +## Design Considerations + +- Performance is critical (type stability and allocations matter) +- Symbolic stack uses Symbolics.jl and SymPyPythonCall +- Supports forward/reverse AD for parameter gradients +- Thread safety matters for estimation workloads diff --git a/docs/agent-guides/task-runbook.md b/docs/agent-guides/task-runbook.md new file mode 100644 index 000000000..93cbd3bf8 --- /dev/null +++ b/docs/agent-guides/task-runbook.md @@ -0,0 +1,68 @@ +# Task Runbook (On-Demand) + +Read this file only for operational heuristics, orchestration style, or common task checklists. + +## Common Change Points + +- New API: update `src/get_functions.jl` and exports in `src/MacroModelling.jl` +- New model: add file under `models/` using model macros +- Solver changes: inspect `src/perturbation.jl` and `src/algorithms/` + +## Typical Task Flows + +### Add a feature + +1. Implement in the appropriate `src/` location +2. Create a minimal targeted check script +3. Validate behavior with lightweight model(s) +4. Update documentation if user-facing + +### Fix a bug + +1. Reproduce minimally +2. Locate root cause +3. Implement smallest robust fix +4. Verify with focused check + +### Add a model + +1. Add model file under `models/` +2. Follow existing model conventions +3. Include citation metadata/context +4. Verify solve + IRFs + +## Workflow Orchestration Heuristics + +### Plan mode default + +- Use plan mode for non-trivial tasks (3+ steps / architecture choices) +- Re-plan quickly if assumptions fail +- Include verification steps in plan, not only implementation + +### Subagent usage + +- Offload exploration/research for complex tasks +- Keep one focused goal per subagent + +### Elegance check (for non-trivial changes) + +- Reassess whether a cleaner root-cause solution exists before finalizing +- Avoid over-engineering for obvious/simple fixes + +### Autonomous bug-fix expectation + +- Drive issue resolution end-to-end without requiring user handholding +- Use logs/errors/tests to iterate quickly to a verified result + +## Task and Learning Files + +- Plan and execution tracking: `tasks/todo.md` +- Lessons from corrections: `tasks/lessons.md` +- Session status handoff: `AGENT_PROGRESS.md` + +## CI/CD Reference + +- CI runs on push +- Matrix includes Ubuntu/macOS/Windows (x64 and arm64 where applicable) +- Coverage uploaded to Codecov +- Test sets run in parallel by matrix configuration diff --git a/docs/generate_plots.jl b/docs/generate_plots.jl index b9c9937be..7a74d7c22 100644 --- a/docs/generate_plots.jl +++ b/docs/generate_plots.jl @@ -624,7 +624,11 @@ plot_irf(Gali_2015_chapter_3_nonlinear, shocks = :eps_a, parameters = :ฮฒ => 0.9 ### tol using MacroModelling: Tolerances -custom_tol = Tolerances(qme_acceptance_tol = 1e-12, sylvester_acceptance_tol = 1e-12) +custom_tol = Tolerances( + first_order = MacroModelling.FirstOrderTolerances(qme = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)) +) plot_irf(Gali_2015_chapter_3_nonlinear, shocks = :eps_a, tol = custom_tol, algorithm = :second_order, parameters = :ฮฒ => 0.9555,verbose = true) ### quadratic_matrix_equation_algorithm diff --git a/docs/src/plot_conditional_forecast.md b/docs/src/plot_conditional_forecast.md index cd7236970..24bbf320e 100644 --- a/docs/src/plot_conditional_forecast.md +++ b/docs/src/plot_conditional_forecast.md @@ -1439,8 +1439,11 @@ The `tol` argument (default: `Tolerances()`, type: `Tolerances`) defines various The tolerances used by the numerical solvers can be adjusted. The Tolerances object allows setting tolerances for the non-stochastic steady state solver (NSSS), Sylvester equations, Lyapunov equation, and quadratic matrix equation (QME). For example, to set tighter tolerances (this example also changes parameters to force recomputation): ```julia -custom_tol = Tolerances(qme_acceptance_tol = 1e-12, - sylvester_acceptance_tol = 1e-12) +custom_tol = Tolerances( + first_order = MacroModelling.FirstOrderTolerances(qme = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)) +) conditions_ka = KeyedArray(Matrix{Union{Nothing,Float64}}(undef,3,3), Variables = [:R, :Y, :MC], diff --git a/docs/src/plot_conditional_variance_decomposition.md b/docs/src/plot_conditional_variance_decomposition.md index 12b177df3..da3e9c8b2 100644 --- a/docs/src/plot_conditional_variance_decomposition.md +++ b/docs/src/plot_conditional_variance_decomposition.md @@ -526,8 +526,11 @@ The `tol` argument (default: `Tolerances()`, type: `Tolerances`) defines various The tolerances used by the numerical solvers can be adjusted. The Tolerances object allows setting tolerances for the non-stochastic steady state solver (NSSS), Sylvester equations, Lyapunov equation, and quadratic matrix equation (QME). For example, to set tighter tolerances (this example also changes parameters to force recomputation): ```julia -custom_tol = Tolerances(qme_acceptance_tol = 1e-12, - sylvester_acceptance_tol = 1e-12) +custom_tol = Tolerances( + first_order = MacroModelling.FirstOrderTolerances(qme = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)) +) plot_fevd(Smets_Wouters_2007_linear, tol = custom_tol, diff --git a/docs/src/plot_irf.md b/docs/src/plot_irf.md index 68d2dbbb5..8c2ec717e 100644 --- a/docs/src/plot_irf.md +++ b/docs/src/plot_irf.md @@ -1468,8 +1468,11 @@ The `tol` argument (default: `Tolerances()`, type: `Tolerances`) defines various The tolerances used by the numerical solvers can be adjusted. The Tolerances object allows setting tolerances for the non-stochastic steady state solver (NSSS), Sylvester equations, Lyapunov equation, and quadratic matrix equation (QME). For example, to set tighter tolerances (this example also changes parameters to force recomputation): ```julia -custom_tol = Tolerances(qme_acceptance_tol = 1e-12, - sylvester_acceptance_tol = 1e-12) +custom_tol = Tolerances( + first_order = MacroModelling.FirstOrderTolerances(qme = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)) +) plot_irf(Gali_2015_chapter_3_nonlinear, shocks = :eps_a, diff --git a/docs/src/plot_model_estimates.md b/docs/src/plot_model_estimates.md index f092691b8..41c581035 100644 --- a/docs/src/plot_model_estimates.md +++ b/docs/src/plot_model_estimates.md @@ -1168,8 +1168,11 @@ The `tol` argument (default: `Tolerances()`, type: `Tolerances`) defines various The tolerances used by the numerical solvers can be adjusted. The Tolerances object allows setting tolerances for the non-stochastic steady state solver (NSSS), Sylvester equations, Lyapunov equation, and quadratic matrix equation (QME). For example, to set tighter tolerances (this example also changes parameters to force recomputation): ```julia -custom_tol = Tolerances(qme_acceptance_tol = 1e-12, - sylvester_acceptance_tol = 1e-12) +custom_tol = Tolerances( + first_order = MacroModelling.FirstOrderTolerances(qme = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)) +) sim_data = simulate(Gali_2015_chapter_3_nonlinear)([:Y],:,:simulate) plot_model_estimates(Gali_2015_chapter_3_nonlinear, diff --git a/docs/src/plot_solution.md b/docs/src/plot_solution.md index 38867fc94..e773038b6 100644 --- a/docs/src/plot_solution.md +++ b/docs/src/plot_solution.md @@ -800,8 +800,11 @@ The `tol` argument (default: `Tolerances()`, type: `Tolerances`) defines various The tolerances used by the numerical solvers can be adjusted. The Tolerances object allows setting tolerances for the non-stochastic steady state solver (NSSS), Sylvester equations, Lyapunov equation, and quadratic matrix equation (QME). For example, to set tighter tolerances (this example also changes parameters to force recomputation): ```julia -custom_tol = Tolerances(qme_acceptance_tol = 1e-12, - sylvester_acceptance_tol = 1e-12) +custom_tol = Tolerances( + first_order = MacroModelling.FirstOrderTolerances(qme = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)), + third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-12)) +) plot_solution(Gali_2015_chapter_3_nonlinear, :A, tol = custom_tol, diff --git a/ext/OptimExt.jl b/ext/OptimExt.jl index 2816ad58b..d1cf049f9 100644 --- a/ext/OptimExt.jl +++ b/ext/OptimExt.jl @@ -1,6 +1,6 @@ module OptimExt -import MacroModelling: find_shocks_conditional_forecast, find_SS_solver_parameters!, Tolerances, โ„ณ, calculate_SS_solver_runtime_and_loglikelihood, solver_parameters, find_shocks_workspace +import MacroModelling: find_shocks_conditional_forecast, find_SS_solver_parameters!, Tolerances, โ„ณ, calculate_SS_solver_runtime_and_loglikelihood, solver_parameters, find_shocks_workspace, solve_nsss_wrapper import Optim # Helper function for LBFGS optimization objective @@ -134,9 +134,9 @@ function find_SS_solver_parameters!(::Val{:SAMIN}, ๐“‚::โ„ณ; par_inputs = solver_parameters(pars..., 1, 0.0, 2) - SS_and_pars, (solution_error, iters) = ๐“‚.functions.NSSS_solve(๐“‚.parameter_values, ๐“‚, tol, false, true, [par_inputs]) + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(๐“‚.parameter_values, ๐“‚, tol, false, true, [par_inputs]) - if solution_error < tol.NSSS_acceptance_tol + if solution_error < tol.nsss.acceptance_tol push!(MacroModelling.DEFAULT_SOLVER_PARAMETERS, par_inputs) return true else diff --git a/ext/StatsPlotsExt.jl b/ext/StatsPlotsExt.jl index 007efcfcf..62f09fc98 100644 --- a/ext/StatsPlotsExt.jl +++ b/ext/StatsPlotsExt.jl @@ -2,7 +2,7 @@ module StatsPlotsExt using MacroModelling -import MacroModelling: ParameterType, โ„ณ, Symbol_input, String_input, Tolerances, merge_calculation_options, MODELยฎ, DATAยฎ, PARAMETERSยฎ, ALGORITHMยฎ, FILTERยฎ, VARIABLESยฎ, SMOOTHยฎ, SHOW_PLOTSยฎ, SAVE_PLOTSยฎ, SAVE_PLOTS_NAMEยฎ, SAVE_PLOTS_FORMATยฎ, SAVE_PLOTS_PATHยฎ, PLOTS_PER_PAGEยฎ, MAX_ELEMENTS_PER_LEGENDS_ROWยฎ, EXTRA_LEGEND_SPACEยฎ, PLOT_ATTRIBUTESยฎ, QMEยฎ, SYLVESTERยฎ, LYAPUNOVยฎ, TOLERANCESยฎ, VERBOSEยฎ, DATA_IN_LEVELSยฎ, PERIODSยฎ, SHOCKSยฎ, SHOCK_SIZEยฎ, NEGATIVE_SHOCKยฎ, GENERALISED_IRFยฎ, GENERALISED_IRF_WARMUP_ITERATIONSยฎ, CONDITIONS_IN_LEVELSยฎ, GENERALISED_IRF_DRAWSยฎ, INITIAL_STATEยฎ, IGNORE_OBCยฎ, CONDITIONSยฎ, SHOCK_CONDITIONSยฎ, LEVELSยฎ, LABELยฎ, RENAME_DICTIONARYยฎ, STEADY_STATE_FUNCTIONยฎ, parse_shocks_input_to_index, parse_variables_input_to_index, replace_indices, replace_indices_special, filter_data_with_model, get_relevant_steady_states, replace_indices_in_symbol, parse_algorithm_to_state_update, girf, decompose_name, obc_objective_optim_fun, obc_constraint_optim_fun, compute_irf_responses, process_ignore_obc_flag, adjust_generalised_irf_flag, process_shocks_input, normalize_filtering_options, infer_step, SteadyStateFunctionType, normalize_superscript, apply_custom_name +import MacroModelling: ParameterType, โ„ณ, Symbol_input, String_input, Tolerances, merge_calculation_options, tol_to_dict, warn_irrelevant_tol, flatten_tol_diff, MODELยฎ, DATAยฎ, PARAMETERSยฎ, ALGORITHMยฎ, FILTERยฎ, VARIABLESยฎ, SMOOTHยฎ, SHOW_PLOTSยฎ, SAVE_PLOTSยฎ, SAVE_PLOTS_NAMEยฎ, SAVE_PLOTS_FORMATยฎ, SAVE_PLOTS_PATHยฎ, PLOTS_PER_PAGEยฎ, MAX_ELEMENTS_PER_LEGENDS_ROWยฎ, EXTRA_LEGEND_SPACEยฎ, PLOT_ATTRIBUTESยฎ, QMEยฎ, SYLVESTERยฎ, LYAPUNOVยฎ, TOLERANCESยฎ, VERBOSEยฎ, DATA_IN_LEVELSยฎ, PERIODSยฎ, SHOCKSยฎ, SHOCK_SIZEยฎ, NEGATIVE_SHOCKยฎ, GENERALISED_IRFยฎ, GENERALISED_IRF_WARMUP_ITERATIONSยฎ, CONDITIONS_IN_LEVELSยฎ, GENERALISED_IRF_DRAWSยฎ, INITIAL_STATEยฎ, IGNORE_OBCยฎ, CONDITIONSยฎ, SHOCK_CONDITIONSยฎ, LEVELSยฎ, LABELยฎ, RENAME_DICTIONARYยฎ, STEADY_STATE_FUNCTIONยฎ, parse_shocks_input_to_index, parse_variables_input_to_index, replace_indices, replace_indices_special, filter_data_with_model, get_relevant_steady_states, replace_indices_in_symbol, parse_algorithm_to_state_update, girf, decompose_name, obc_objective_optim_fun, obc_constraint_optim_fun, compute_irf_responses, process_ignore_obc_flag, adjust_generalised_irf_flag, process_shocks_input, normalize_filtering_options, infer_step, SteadyStateFunctionType, normalize_superscript, apply_custom_name import MacroModelling: DEFAULT_ALGORITHM, DEFAULT_FILTER_SELECTOR, DEFAULT_WARMUP_ITERATIONS, DEFAULT_VARIABLES_EXCLUDING_OBC, DEFAULT_SHOCK_SELECTION, DEFAULT_PRESAMPLE_PERIODS, DEFAULT_DATA_IN_LEVELS, DEFAULT_SHOCK_DECOMPOSITION_SELECTOR, DEFAULT_SMOOTH_SELECTOR, DEFAULT_LABEL, DEFAULT_SHOW_PLOTS, DEFAULT_SAVE_PLOTS, DEFAULT_SAVE_PLOTS_FORMAT, DEFAULT_SAVE_PLOTS_PATH, DEFAULT_PLOTS_PER_PAGE_SMALL, DEFAULT_TRANSPARENCY, DEFAULT_MAX_ELEMENTS_PER_LEGEND_ROW, DEFAULT_EXTRA_LEGEND_SPACE, DEFAULT_VERBOSE, DEFAULT_QME_ALGORITHM, DEFAULT_SYLVESTER_SELECTOR, DEFAULT_SYLVESTER_THRESHOLD, DEFAULT_LARGE_SYLVESTER_ALGORITHM, DEFAULT_SYLVESTER_ALGORITHM, DEFAULT_LYAPUNOV_ALGORITHM, DEFAULT_PLOT_ATTRIBUTES, DEFAULT_ARGS_AND_KWARGS_NAMES, DEFAULT_PLOTS_PER_PAGE_LARGE, DEFAULT_SHOCKS_EXCLUDING_OBC, DEFAULT_VARIABLES_EXCLUDING_AUX_AND_OBC, DEFAULT_PERIODS, DEFAULT_SHOCK_SIZE, DEFAULT_NEGATIVE_SHOCK, DEFAULT_GENERALISED_IRF, DEFAULT_GENERALISED_IRF_WARMUP, DEFAULT_GENERALISED_IRF_DRAWS, DEFAULT_INITIAL_STATE, DEFAULT_IGNORE_OBC, DEFAULT_PLOT_TYPE, DEFAULT_CONDITIONS_IN_LEVELS, DEFAULT_SIGMA_RANGE, DEFAULT_FONT_SIZE, DEFAULT_VARIABLE_SELECTION, DEFAULT_FORECAST_PERIODS import DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF, TYPEDSIGNATURES, TYPEDFIELDS import LaTeXStrings @@ -164,6 +164,7 @@ function plot_model_estimates(๐“‚::โ„ณ, sylvester_algorithmยฒ = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2], lyapunov_algorithm = lyapunov_algorithm) + warn_irrelevant_tol(tol, algorithm; needs_covariance = filter == :kalman) gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() @@ -340,18 +341,7 @@ function plot_model_estimates(๐“‚::โ„ณ, # :shock_decomposition => shock_decomposition, :smooth => smooth, - :NSSS_acceptance_tol => tol.NSSS_acceptance_tol, - :NSSS_xtol => tol.NSSS_xtol, - :NSSS_ftol => tol.NSSS_ftol, - :NSSS_rel_xtol => tol.NSSS_rel_xtol, - :qme_tol => tol.qme_tol, - :qme_acceptance_tol => tol.qme_acceptance_tol, - :sylvester_tol => tol.sylvester_tol, - :sylvester_acceptance_tol => tol.sylvester_acceptance_tol, - :lyapunov_tol => tol.lyapunov_tol, - :lyapunov_acceptance_tol => tol.lyapunov_acceptance_tol, - :droptol => tol.droptol, - :dependencies_tol => tol.dependencies_tol, + :tol => tol_to_dict(tol, algorithm; needs_covariance = filter == :kalman), :quadratic_matrix_equation_algorithm => quadratic_matrix_equation_algorithm, :sylvester_algorithm => sylvester_algorithm, @@ -818,6 +808,7 @@ function plot_model_estimates!(๐“‚::โ„ณ, sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2], lyapunov_algorithm = lyapunov_algorithm) + warn_irrelevant_tol(tol, algorithm; needs_covariance = filter == :kalman) gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() if !gr_back @@ -988,18 +979,7 @@ function plot_model_estimates!(๐“‚::โ„ณ, # :shock_decomposition => shock_decomposition, :smooth => smooth, - :NSSS_acceptance_tol => tol.NSSS_acceptance_tol, - :NSSS_xtol => tol.NSSS_xtol, - :NSSS_ftol => tol.NSSS_ftol, - :NSSS_rel_xtol => tol.NSSS_rel_xtol, - :qme_tol => tol.qme_tol, - :qme_acceptance_tol => tol.qme_acceptance_tol, - :sylvester_tol => tol.sylvester_tol, - :sylvester_acceptance_tol => tol.sylvester_acceptance_tol, - :lyapunov_tol => tol.lyapunov_tol, - :lyapunov_acceptance_tol => tol.lyapunov_acceptance_tol, - :droptol => tol.droptol, - :dependencies_tol => tol.dependencies_tol, + :tol => tol_to_dict(tol, algorithm; needs_covariance = filter == :kalman), :quadratic_matrix_equation_algorithm => quadratic_matrix_equation_algorithm, :sylvester_algorithm => sylvester_algorithm, @@ -1201,6 +1181,10 @@ function plot_model_estimates!(๐“‚::โ„ณ, push!(annotate_diff_input, DEFAULT_ARGS_AND_KWARGS_NAMES[k] => reduce(vcat, diffdict[k])) end end + + if haskey(diffdict, :tol) + append!(annotate_diff_input, flatten_tol_diff(diffdict[:tol])) + end if haskey(diffdict, :shock_names) if all(length.(diffdict[:shock_names]) .== 1) @@ -1773,6 +1757,7 @@ function plot_irf(๐“‚::โ„ณ; sylvester_algorithmยฒ = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2]) + warn_irrelevant_tol(tol, algorithm; needs_covariance = false) gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() if !gr_back @@ -1933,16 +1918,7 @@ function plot_irf(๐“‚::โ„ณ; :initial_state => initial_state_input, :ignore_obc => ignore_obc, - :NSSS_acceptance_tol => tol.NSSS_acceptance_tol, - :NSSS_xtol => tol.NSSS_xtol, - :NSSS_ftol => tol.NSSS_ftol, - :NSSS_rel_xtol => tol.NSSS_rel_xtol, - :qme_tol => tol.qme_tol, - :qme_acceptance_tol => tol.qme_acceptance_tol, - :sylvester_tol => tol.sylvester_tol, - :sylvester_acceptance_tol => tol.sylvester_acceptance_tol, - :droptol => tol.droptol, - :dependencies_tol => tol.dependencies_tol, + :tol => tol_to_dict(tol, algorithm; needs_covariance = false), :quadratic_matrix_equation_algorithm => quadratic_matrix_equation_algorithm, :sylvester_algorithm => sylvester_algorithm, @@ -2467,6 +2443,7 @@ function plot_irf!(๐“‚::โ„ณ; sylvester_algorithmยฒ = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2]) + warn_irrelevant_tol(tol, algorithm; needs_covariance = false) gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() if !gr_back @@ -2618,16 +2595,7 @@ function plot_irf!(๐“‚::โ„ณ; :initial_state => initial_state_input, :ignore_obc => ignore_obc, - :NSSS_acceptance_tol => tol.NSSS_acceptance_tol, - :NSSS_xtol => tol.NSSS_xtol, - :NSSS_ftol => tol.NSSS_ftol, - :NSSS_rel_xtol => tol.NSSS_rel_xtol, - :qme_tol => tol.qme_tol, - :qme_acceptance_tol => tol.qme_acceptance_tol, - :sylvester_tol => tol.sylvester_tol, - :sylvester_acceptance_tol => tol.sylvester_acceptance_tol, - :droptol => tol.droptol, - :dependencies_tol => tol.dependencies_tol, + :tol => tol_to_dict(tol, algorithm; needs_covariance = false), :quadratic_matrix_equation_algorithm => quadratic_matrix_equation_algorithm, :sylvester_algorithm => sylvester_algorithm, @@ -2819,6 +2787,10 @@ function plot_irf!(๐“‚::โ„ณ; end end + if haskey(diffdict, :tol) + append!(annotate_diff_input, flatten_tol_diff(diffdict[:tol])) + end + legend_plot = StatsPlots.plot(framestyle = :none, @@ -3781,6 +3753,7 @@ function plot_solution(๐“‚::โ„ณ, sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2], lyapunov_algorithm = lyapunov_algorithm) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() if !gr_back @@ -3924,6 +3897,7 @@ function plot_solution(๐“‚::โ„ณ, :ฯƒ => ฯƒ, :parameters => Dict(๐“‚.constants.post_complete_parameters.parameters .=> ๐“‚.parameter_values), :ignore_obc => ignore_obc, + :tol => tol_to_dict(tol, algorithm; needs_covariance = true), :variable_output => variable_output, :has_impact => has_impact, :vars_to_plot => vars_to_plot, @@ -4109,6 +4083,10 @@ function _plot_solution_from_container(; push!(annotate_diff_input, "Ignore OBC" => reduce(vcat, diffdict[:ignore_obc])) end + if haskey(diffdict, :tol) + append!(annotate_diff_input, flatten_tol_diff(diffdict[:tol])) + end + # Determine legend labels based on what differs # If more than one input differs (besides label), use custom labels from diffdict len_diff = length(solution_active_plot_container) @@ -4509,6 +4487,7 @@ function plot_solution!(๐“‚::โ„ณ, sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2], lyapunov_algorithm = lyapunov_algorithm) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() if !gr_back @@ -4647,6 +4626,7 @@ function plot_solution!(๐“‚::โ„ณ, :ฯƒ => ฯƒ, :parameters => Dict(๐“‚.constants.post_complete_parameters.parameters .=> ๐“‚.parameter_values), :ignore_obc => ignore_obc, + :tol => tol_to_dict(tol, algorithm; needs_covariance = true), :variable_output => variable_output, :has_impact => has_impact, :vars_to_plot => vars_to_plot, @@ -4824,6 +4804,7 @@ function plot_conditional_forecast(๐“‚::โ„ณ, sylvester_algorithm = sylvester_algorithm, tol = tol, verbose = verbose) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) periods += max(size(conditions,2), isnothing(shocks) ? 1 : size(shocks,2)) @@ -4981,16 +4962,7 @@ function plot_conditional_forecast(๐“‚::โ„ณ, :var_idx => var_idx, :algorithm => algorithm, - :NSSS_acceptance_tol => tol.NSSS_acceptance_tol, - :NSSS_xtol => tol.NSSS_xtol, - :NSSS_ftol => tol.NSSS_ftol, - :NSSS_rel_xtol => tol.NSSS_rel_xtol, - :qme_tol => tol.qme_tol, - :qme_acceptance_tol => tol.qme_acceptance_tol, - :sylvester_tol => tol.sylvester_tol, - :sylvester_acceptance_tol => tol.sylvester_acceptance_tol, - :droptol => tol.droptol, - :dependencies_tol => tol.dependencies_tol, + :tol => tol_to_dict(tol, algorithm; needs_covariance = true), :quadratic_matrix_equation_algorithm => quadratic_matrix_equation_algorithm, :sylvester_algorithm => sylvester_algorithm, @@ -5288,6 +5260,7 @@ function plot_conditional_forecast!(๐“‚::โ„ณ, tol = tol, verbose = verbose) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) periods += max(size(conditions,2), isnothing(shocks) ? 1 : size(shocks,2)) full_SS = vcat(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.aux,๐“‚.constants.post_model_macro.exo_present)),map(x->Symbol(string(x) * "โ‚โ‚“โ‚Ž"),๐“‚.constants.post_model_macro.exo)) @@ -5448,16 +5421,7 @@ function plot_conditional_forecast!(๐“‚::โ„ณ, :var_idx => var_idx, :algorithm => algorithm, - :NSSS_acceptance_tol => tol.NSSS_acceptance_tol, - :NSSS_xtol => tol.NSSS_xtol, - :NSSS_ftol => tol.NSSS_ftol, - :NSSS_rel_xtol => tol.NSSS_rel_xtol, - :qme_tol => tol.qme_tol, - :qme_acceptance_tol => tol.qme_acceptance_tol, - :sylvester_tol => tol.sylvester_tol, - :sylvester_acceptance_tol => tol.sylvester_acceptance_tol, - :droptol => tol.droptol, - :dependencies_tol => tol.dependencies_tol, + :tol => tol_to_dict(tol, algorithm; needs_covariance = true), :quadratic_matrix_equation_algorithm => quadratic_matrix_equation_algorithm, :sylvester_algorithm => sylvester_algorithm, @@ -5706,6 +5670,10 @@ function plot_conditional_forecast!(๐“‚::โ„ณ, end end + if haskey(diffdict, :tol) + append!(annotate_diff_input, flatten_tol_diff(diffdict[:tol])) + end + if haskey(diffdict, :shock_names) if all(length.(diffdict[:shock_names]) .== 1) push!(annotate_diff_input, "Shock name" => map(x->x[1], diffdict[:shock_names])) diff --git a/models/Caldara_et_al_2012.jl b/models/Caldara_et_al_2012.jl index a52fb2904..a93eae445 100644 --- a/models/Caldara_et_al_2012.jl +++ b/models/Caldara_et_al_2012.jl @@ -1,4 +1,3 @@ - @model Caldara_et_al_2012 begin V[0] = ((1 - ฮฒ) * (c[0] ^ ฮฝ * (1 - l[0]) ^ (1 - ฮฝ)) ^ (1 - 1 / ฯˆ) + ฮฒ * V[1] ^ (1 - 1 / ฯˆ)) ^ (1 / (1 - 1 / ฯˆ)) diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index d1942d6bb..a1e161204 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -34,9 +34,11 @@ import LoopVectorization: @turbo # import Polyester import NLopt # import Zygote +import SparseArrays import SparseArrays: SparseMatrixCSC, SparseVector, AbstractSparseArray, AbstractSparseMatrix, sparse!, spzeros, nnz, issparse, nonzeros #, sparse, droptol!, sparsevec, spdiagm, findnz#, sparse! import LinearAlgebra as โ„’ import LinearSolve as ๐’ฎ +import FastLapackInterface # import LinearAlgebra: mul! # import Octavian: matmul! # import TriangularSolve as TS @@ -59,8 +61,8 @@ import MatrixEquations # good overview: https://cscproxy.mpi-magdeburg.mpg.de/mp # using NamedArrays # using AxisKeys -import ChainRulesCore: @ignore_derivatives, ignore_derivatives, rrule, NoTangent, @thunk, ProjectTo, unthunk, AbstractZero -import RecursiveFactorization as RF +import ChainRulesCore: rrule, NoTangent, @thunk, ProjectTo, unthunk, AbstractZero +# import RecursiveFactorization as RF using RuntimeGeneratedFunctions RuntimeGeneratedFunctions.init(@__MODULE__) @@ -163,11 +165,13 @@ include("common_docstrings.jl") include("structures.jl") include("solver_parameters.jl") include("options_and_caches.jl") +include("nsss_solver.jl") include("macros.jl") include("get_functions.jl") include("dynare.jl") include("inspect.jl") include("moments.jl") +include("./algorithms/fast_lapack_wrappers.jl") include("perturbation.jl") include("./algorithms/sylvester.jl") @@ -203,7 +207,7 @@ export get_fevd, fevd, get_forecast_error_variance_decomposition, get_conditiona export calculate_jacobian, calculate_hessian, calculate_third_order_derivatives export calculate_first_order_solution, calculate_second_order_solution, calculate_third_order_solution #, calculate_jacobian_manual, calculate_jacobian_sparse, calculate_jacobian_threaded export get_shock_decomposition, get_model_estimates, get_estimated_shocks, get_estimated_variables, get_estimated_variable_standard_deviations, get_loglikelihood -export Tolerances +export Tolerances, SolverTolerances, NsssTolerances, AdTolerances, FirstOrderTolerances, HigherOrderTolerances export translate_mod_file, translate_dynare_file, import_model, import_dynare export write_mod_file, write_dynare_file, write_to_dynare_file, write_to_dynare, export_dynare, export_to_dynare, export_mod_file, export_model @@ -367,7 +371,6 @@ Base.show(io::IO, ๐“‚::โ„ณ) = println(io, end, # "\nยน: including auxiliary variables" # "\nVariable bounds (upper,lower,any): ",sum(๐“‚.upper_bounds .< Inf),", ",sum(๐“‚.lower_bounds .> -Inf),", ",length(๐“‚.bounds), - # "\nNon-stochastic-steady-state found: ",!๐“‚.caches.outdated_NSSS ) check_for_dynamic_variables(ex::Int) = false @@ -378,17 +381,46 @@ check_for_dynamic_variables(ex::Symbol) = occursin(r"โ‚โ‚โ‚Ž|โ‚โ‚€โ‚Ž|โ‚โ‚‹ function compare_args_and_kwargs(dicts::Vector{S}) where S <: Dict N = length(dicts) - @assert N โ‰ฅ 2 "Need at least two dictionaries to compare" + + if N โ‰ค 1 + # Single entry: nothing to compare. Return every non-skipped key so + # downstream code (e.g. diffdict[:label]) works uniformly. + # Dict values are recursed into so the result shape matches the Nโ‰ฅ2 + # case (nested Dicts with leaf vectors) expected by flatten_tol_diff. + diffs = Dict{Symbol,Any}() + if N == 1 + for k in keys(dicts[1]) + k in (:plot_data, :plot_type) && continue + v = dicts[1][k] + if v isa Dict + diffs[k] = compare_args_and_kwargs([v]) + else + diffs[k] = [v] + end + end + end + return diffs + end diffs = Dict{Symbol,Any}() - # assume all dictionaries share the same set of keys - for k in keys(dicts[1]) + # use the union of all keys so dicts with different key sets + # (e.g. tol sub-dicts that conditionally include :dependencies_tol) + # are compared correctly + all_keys = reduce(union, keys.(dicts)) + + for k in all_keys if k in [:plot_data, :plot_type] # skip keys that are not relevant for comparison continue end + # when a key is missing from some dicts, the values differ by definition + if !all(haskey(d, k) for d in dicts) + diffs[k] = [get(d, k, missing) for d in dicts] + continue + end + vals = [d[k] for d in dicts] if all(v -> v isa Dict, vals) @@ -427,6 +459,34 @@ function compare_args_and_kwargs(dicts::Vector{S}) where S <: Dict end +""" + flatten_tol_diff(diff; names = DEFAULT_ARGS_AND_KWARGS_NAMES, prefix = "") -> Vector{Pair{String,Any}} + +Recursively walk a nested tolerance diff `Dict` (as returned by +`compare_args_and_kwargs` on `tol_to_dict` outputs) and produce a flat vector +of `"human-readable path" => values` pairs suitable for plot annotations. + +Path segments are translated through `names` (defaults to +`DEFAULT_ARGS_AND_KWARGS_NAMES`). For example a diff at +`:first_order => :qme => :atol` becomes `"1st order QME atol"`. +""" +function flatten_tol_diff(diff::Dict; + names::Dict{Symbol,String} = DEFAULT_ARGS_AND_KWARGS_NAMES, + prefix::String = "") + result = Pair{String,Any}[] + for (k, v) in sort(collect(diff), by = first) + seg = get(names, k, String(k)) + label = isempty(prefix) ? seg : prefix * " " * seg + if v isa Dict + append!(result, flatten_tol_diff(v; names = names, prefix = label)) + else + push!(result, label => reduce(vcat, v)) + end + end + return result +end + + function mul_reverse_AD!( C::Matrix{S}, A::AbstractMatrix{M}, B::AbstractMatrix{N}) where {S <: Real, M <: Real, N <: Real} @@ -1004,29 +1064,71 @@ end function clear_solution_caches!(๐“‚::โ„ณ, algorithm::Symbol) - # Mark all solutions as outdated - ๐“‚.caches.outdated.non_stochastic_steady_state = true - ๐“‚.caches.outdated.jacobian = true - ๐“‚.caches.outdated.hessian = true - ๐“‚.caches.outdated.third_order_derivatives = true - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true - - while length(๐“‚.caches.solver_cache) > 1 - pop!(๐“‚.caches.solver_cache) + while length(๐“‚.caches.solver) > 1 + pop!(๐“‚.caches.solver) end + ๐“‚.caches.first_order_solution_matrix = zeros(0,0) + ๐“‚.caches.first_order_obc_solution_matrix = zeros(0,0) ๐“‚.caches.qme_solution = zeros(0,0) ๐“‚.caches.second_order_solution = spzeros(0,0) ๐“‚.caches.third_order_solution = spzeros(0,0) + ๐“‚.caches.second_order_stochastic_steady_state = Float64[] + ๐“‚.caches.pruned_second_order_stochastic_steady_state = Float64[] + ๐“‚.caches.third_order_stochastic_steady_state = Float64[] + ๐“‚.caches.pruned_third_order_stochastic_steady_state = Float64[] + + resize!(๐“‚.caches.non_stochastic_steady_state, 0) + ๐“‚.caches.valid_for.non_stochastic_steady_state = Float64[] + + ๐“‚.caches.valid_for.jacobian = Float64[] + ๐“‚.caches.valid_for.hessian = Float64[] + ๐“‚.caches.valid_for.third_order_derivatives = Float64[] + ๐“‚.caches.valid_for.first_order_solution = Float64[] + ๐“‚.caches.valid_for.first_order_obc_solution = Float64[] + ๐“‚.caches.valid_for.second_order_solution = Float64[] + ๐“‚.caches.valid_for.pruned_second_order_solution = Float64[] + ๐“‚.caches.valid_for.second_order_stochastic_steady_state = Float64[] + ๐“‚.caches.valid_for.pruned_second_order_stochastic_steady_state = Float64[] + ๐“‚.caches.valid_for.third_order_solution = Float64[] + ๐“‚.caches.valid_for.pruned_third_order_solution = Float64[] + ๐“‚.caches.valid_for.third_order_stochastic_steady_state = Float64[] + ๐“‚.caches.valid_for.pruned_third_order_stochastic_steady_state = Float64[] + return nothing end +const CACHE_VALIDITY_FIELDS = ( + :non_stochastic_steady_state, + :jacobian, + :hessian, + :third_order_derivatives, + :first_order_solution, + :first_order_obc_solution, + :second_order_solution, + :pruned_second_order_solution, + :second_order_stochastic_steady_state, + :pruned_second_order_stochastic_steady_state, + :third_order_solution, + :pruned_third_order_solution, + :third_order_stochastic_steady_state, + :pruned_third_order_stochastic_steady_state, +) + + +@inline function cache_valid_for_parameters(valid_for::Vector{Float64}, parameters::AbstractVector{<:Real})::Bool + length(valid_for) == length(parameters) || return false + @inbounds for i in eachindex(parameters) + if valid_for[i] != parameters[i] + return false + end + end + return true +end + + """ set_custom_steady_state_function!(๐“‚::โ„ณ, f::SteadyStateFunctionType) @@ -1107,35 +1209,10 @@ get_irf(RBC, steady_state_function = my_steady_state) See also: [`get_variables`](@ref), [`get_parameters`](@ref), [`get_steady_state`](@ref), [`get_irf`](@ref), [`simulate`](@ref) """ function set_custom_steady_state_function!(๐“‚::โ„ณ, f::SteadyStateFunctionType) - had_custom = !isnothing(๐“‚.functions.NSSS_custom) - - # Store the custom function - if isnothing(f) + if f === nothing ๐“‚.functions.NSSS_custom = nothing - - if had_custom - ๐“‚.caches.outdated.non_stochastic_steady_state = true - ๐“‚.caches.outdated.jacobian = true - ๐“‚.caches.outdated.hessian = true - ๐“‚.caches.outdated.third_order_derivatives = true - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true - end elseif f isa Function && f !== ๐“‚.functions.NSSS_custom - ๐“‚.functions.NSSS_custom = f - - ๐“‚.caches.outdated.non_stochastic_steady_state = true - ๐“‚.caches.outdated.jacobian = true - ๐“‚.caches.outdated.hessian = true - ๐“‚.caches.outdated.third_order_derivatives = true - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true + ๐“‚.functions.NSSS_custom = f end return nothing @@ -1174,254 +1251,6 @@ function infer_step(x_axis::AbstractVector{T}) where {T<:Dates.TimeType} return d2 - d1 end -function fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, - โˆ‚B::AbstractMatrix{R}, - โˆ‚X::AbstractSparseMatrix{R}, - A::AbstractMatrix{R}, - B::AbstractMatrix{R}) where R <: Real - @assert size(โˆ‚A) == size(A) - @assert size(โˆ‚B) == size(B) - @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" - - n1, m1 = size(B) - n2 = size(A,1) - - # Precompute constants - const_n1n2 = n1 * n2 - const_n1n2m1 = n1 * n2 * m1 - - # Access the sparse matrix internal representation - if โˆ‚X isa SparseMatrixCSC - colptr = โˆ‚X.colptr # Column pointers - rowval = โˆ‚X.rowval # Row indices of non-zeros - nzval = โˆ‚X.nzval # Non-zero values - else - colptr = โˆ‚X.A.colptr # Column pointers - rowval = โˆ‚X.A.rowval # Row indices of non-zeros - nzval = โˆ‚X.A.nzval # Non-zero values - end - - # Iterate over columns of โˆ‚X - for col in 1:size(โˆ‚X, 2) - # Iterate over the non-zeros in this column - for idx in colptr[col]:(colptr[col + 1] - 1) - row = rowval[idx] - val = nzval[idx] - - linear_idx = (col - 1) * size(โˆ‚X, 1) + row - - @inbounds begin - i = (linear_idx - 1) % n1 + 1 - k = ((linear_idx - 1) รท n1) % n2 + 1 - j = ((linear_idx - 1) รท const_n1n2) % m1 + 1 - l = ((linear_idx - 1) รท const_n1n2m1) + 1 - - # Update โˆ‚B and โˆ‚A - โˆ‚A[k,l] += B[i,j] * val - โˆ‚B[i,j] += A[k,l] * val - end - end - end -end - - -function fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, - โˆ‚B::AbstractMatrix{R}, - โˆ‚X::DenseMatrix{R}, - A::AbstractMatrix{R}, - B::AbstractMatrix{R}) where R <: Real - @assert size(โˆ‚A) == size(A) - @assert size(โˆ‚B) == size(B) - @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" - - reโˆ‚X = reshape(โˆ‚X, - size(A,1), - size(B,1), - size(A,2), - size(B,2)) - - ei = 1 - for e in eachslice(reโˆ‚X; dims = (1,3)) - @inbounds โˆ‚A[ei] += โ„’.dot(B,e) - ei += 1 - end - - ei = 1 - for e in eachslice(reโˆ‚X; dims = (2,4)) - @inbounds โˆ‚B[ei] += โ„’.dot(A,e) - ei += 1 - end -end - - - -function fill_kron_adjoint!(โˆ‚A::V, โˆ‚B::V, โˆ‚X::V, A::V, B::V) where V <: Vector{<: Real} - @assert size(โˆ‚A) == size(A) - @assert size(โˆ‚B) == size(B) - @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" - - reโˆ‚X = reshape(โˆ‚X, - length(A), - length(B)) - - ei = 1 - for e in eachslice(reโˆ‚X; dims = 1) - @inbounds โˆ‚A[ei] += โ„’.dot(B,e) - ei += 1 - end - - ei = 1 - for e in eachslice(reโˆ‚X; dims = 2) - @inbounds โˆ‚B[ei] += โ„’.dot(A,e) - ei += 1 - end -end - - -function fill_kron_adjoint_โˆ‚B!(โˆ‚X::AbstractSparseMatrix{R}, โˆ‚B::AbstractArray{S}, A::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} - @assert length(โˆ‚X) == length(โˆ‚B) * length(A) "โˆ‚X must have the same length as kron(B,A)" - - n1, m1 = size(โˆ‚B) - n2 = size(A,1) - - # Precompute constants - const_n1n2 = n1 * n2 - const_n1n2m1 = n1 * n2 * m1 - - # Access the sparse matrix internal representation - colptr = โˆ‚X.colptr # Column pointers - rowval = โˆ‚X.rowval # Row indices of non-zeros - nzval = โˆ‚X.nzval # Non-zero values - - # Iterate over columns of โˆ‚X - for col in 1:size(โˆ‚X, 2) - # Iterate over the non-zeros in this column - for idx in colptr[col]:(colptr[col + 1] - 1) - row = rowval[idx] - val = nzval[idx] - - linear_idx = (col - 1) * size(โˆ‚X, 1) + row - - @inbounds begin - i = (linear_idx - 1) % n1 + 1 - k = ((linear_idx - 1) รท n1) % n2 + 1 - j = ((linear_idx - 1) รท const_n1n2) % m1 + 1 - l = ((linear_idx - 1) รท const_n1n2m1) + 1 - - # Update โˆ‚B and โˆ‚A - โˆ‚B[i,j] += A[k,l] * val - end - end - end -end - - - -function fill_kron_adjoint_โˆ‚B!(โˆ‚X::AbstractSparseMatrix{R}, โˆ‚B::Vector{S}, A::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} - @assert length(โˆ‚X) == length(โˆ‚B) * length(A) "โˆ‚X must have the same length as kron(B,A)" - - n1 = length(โˆ‚B) - n2 = size(A,1) - # println("hello") - # Precompute constants - const_n1n2 = n1 * n2 - - # Access the sparse matrix internal representation - colptr = โˆ‚X.colptr # Column pointers - rowval = โˆ‚X.rowval # Row indices of non-zeros - nzval = โˆ‚X.nzval # Non-zero values - - # Iterate over columns of โˆ‚X - for col in 1:size(โˆ‚X, 2) - # Iterate over the non-zeros in this column - for idx in colptr[col]:(colptr[col + 1] - 1) - row = rowval[idx] - val = nzval[idx] - - linear_idx = (col - 1) * size(โˆ‚X, 1) + row - - @inbounds begin - i = (linear_idx - 1) % n1 + 1 - k = ((linear_idx - 1) รท n1) % n2 + 1 - l = ((linear_idx - 1) รท const_n1n2) + 1 - - # Update โˆ‚B and โˆ‚A - โˆ‚B[i] += A[k,l] * val - end - end - end -end - - - -function fill_kron_adjoint_โˆ‚B!(โˆ‚X::DenseMatrix{R}, โˆ‚B::Vector{S}, A::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} - @assert length(โˆ‚X) == length(โˆ‚B) * length(A) "โˆ‚X must have the same length as kron(B,A)" - - reโˆ‚X = reshape(โˆ‚X, - size(A,1), - length(โˆ‚B), - size(A,2)) - - ei = 1 - for e in eachslice(reโˆ‚X; dims = 2) - @inbounds โˆ‚B[ei] += โ„’.dot(A,e) - ei += 1 - end -end - - -function fill_kron_adjoint_โˆ‚A!(โˆ‚X::DenseMatrix{R}, โˆ‚A::Vector{S}, B::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} - @assert length(โˆ‚X) == length(โˆ‚A) * length(B) "โˆ‚X must have the same length as kron(B,A)" - - reโˆ‚X = reshape(โˆ‚X, - length(โˆ‚A), - size(B,1), - size(B,2)) - - ei = 1 - for e in eachslice(reโˆ‚X; dims = 1) - @inbounds โˆ‚A[ei] += โ„’.dot(B,e) - ei += 1 - end -end - - -function fill_kron_adjoint_โˆ‚A!(โˆ‚X::AbstractSparseMatrix{R}, โˆ‚A::AbstractMatrix{S}, B::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} - @assert length(โˆ‚X) == length(B) * length(โˆ‚A) "โˆ‚X must have the same length as kron(B,A)" - - n1, m1 = size(B) - n2 = size(โˆ‚A,1) - - # Precompute constants - const_n1n2 = n1 * n2 - const_n1n2m1 = n1 * n2 * m1 - - # Access the sparse matrix internal representation - colptr = โˆ‚X.colptr # Column pointers - rowval = โˆ‚X.rowval # Row indices of non-zeros - nzval = โˆ‚X.nzval # Non-zero values - - # Iterate over columns of โˆ‚X - for col in 1:size(โˆ‚X, 2) - # Iterate over the non-zeros in this column - for idx in colptr[col]:(colptr[col + 1] - 1) - row = rowval[idx] - val = nzval[idx] - - linear_idx = (col - 1) * size(โˆ‚X, 1) + row - - @inbounds begin - i = (linear_idx - 1) % n1 + 1 - k = ((linear_idx - 1) รท n1) % n2 + 1 - j = ((linear_idx - 1) รท const_n1n2) % m1 + 1 - l = ((linear_idx - 1) รท const_n1n2m1) + 1 - - # Update โˆ‚B and โˆ‚A - โˆ‚A[k,l] += B[i,j] * val - end - end - end -end function choose_matrix_format(A::โ„’.Diagonal{S, Vector{S}}; @@ -1448,11 +1277,21 @@ function choose_matrix_format(A::โ„’.Adjoint{S, M}; min_length::Int = 1000, tol::R = 1e-14, multithreaded::Bool = true)::Union{Matrix{S}, SparseMatrixCSC{S, Int}, ThreadedSparseArrays.ThreadedSparseMatrixCSC{S, Int, SparseMatrixCSC{S, Int}}} where {R <: AbstractFloat, S <: Real, M <: AbstractMatrix{S}} - choose_matrix_format(convert(typeof(transpose(A)),A), - density_threshold = density_threshold, - min_length = min_length, - multithreaded = multithreaded, - tol = tol) + if A.parent isa AbstractSparseMatrix || A.parent isa ThreadedSparseArrays.ThreadedSparseMatrixCSC + # Materialise sparse adjoints as SparseMatrixCSC to avoid unsupported + # ThreadedSparseMatrixCSC(::Adjoint{<:ThreadedSparseMatrixCSC}) conversion. + return choose_matrix_format(sparse(A), + density_threshold = density_threshold, + min_length = min_length, + multithreaded = multithreaded, + tol = tol) + else + return choose_matrix_format(Matrix(A), + density_threshold = density_threshold, + min_length = min_length, + multithreaded = multithreaded, + tol = tol) + end end # function choose_matrix_format(A::โ„’.Adjoint{S, <: AbstractSparseMatrix{S}}; @@ -1491,7 +1330,7 @@ function choose_matrix_format(A::DenseMatrix{S}; min_length::Int = 1000, tol::R = 1e-14, multithreaded::Bool = true)::Union{Matrix{S}, SparseMatrixCSC{S, Int}, ThreadedSparseArrays.ThreadedSparseMatrixCSC{S, Int, SparseMatrixCSC{S, Int}}} where {R <: AbstractFloat, S <: Real} - if sum(abs.(A) .> tol) / length(A) < density_threshold && length(A) > min_length + if count(x -> abs(x) > tol, A) / length(A) < density_threshold && length(A) > min_length # Use dense_to_sparse to avoid Julia 1.12 SparseArrays bug in SparseMatrixCSC(::Matrix) a = dense_to_sparse(A, tol) if multithreaded @@ -1537,14 +1376,54 @@ end function mat_mult_kron(A::AbstractSparseMatrix{R}, B::AbstractMatrix{T}, C::AbstractMatrix{T}, - D::AbstractMatrix{S}) where {R <: Real, T <: Real, S <: Real} + D::AbstractMatrix{S}; + sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{T}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{T}} = (Int[], Int[], T[], Int[], Int[], Int[], T[]), + sparse::Bool = false) where {R <: Real, T <: Real, S <: Real} n_rowB = size(B,1) n_colB = size(B,2) n_rowC = size(C,1) n_colC = size(C,2) - X = zeros(T, size(A,1), size(D,2)) + estimated_nnz = 0 + I = Vector{Int}() + J = Vector{Int}() + V = Vector{T}() + X = zeros(T, 0, 0) + reused_sparse_buffers = sparse && length(sparse_preallocation[1]) > 0 + + if sparse + nnzA = nnz(A) + nnzB = sum(abs.(B) .> eps()) + nnzC = sum(abs.(C) .> eps()) + nnzD = sum(abs.(D) .> eps()) + + p = Float64(nnzA) * Float64(nnzB) * Float64(nnzC) * Float64(nnzD) / (Float64(length(A)) * Float64(length(B)) * Float64(length(C)) * Float64(length(D))) + + if length(sparse_preallocation[1]) == 0 + estimated_nnz = Int(ceil((1 - (1 - p)^size(A,1)) * size(A,1) * size(D,2))) + + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + + I = sparse_preallocation[1] + J = sparse_preallocation[2] + V = sparse_preallocation[3] + else + estimated_nnz = length(sparse_preallocation[3]) + + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + + I = sparse_preallocation[1] + J = sparse_preallocation[2] + V = sparse_preallocation[3] + end + else + X = zeros(T, size(A,1), size(D,2)) + end # vals = T[] # rows = Int[] @@ -1554,36 +1433,77 @@ function mat_mult_kron(A::AbstractSparseMatrix{R}, Aฬ„B = zeros(T, n_rowC, n_colB) CAฬ„B = zeros(T, n_colC, n_colB) vCAฬ„B = zeros(T, n_colB * n_colC) - # vCAฬ„BD = zeros(size(D,2)) + vCAฬ„BD = zeros(T, size(D,2)) + + rv = A isa SparseMatrixCSC ? A.rowval : A.A.rowval + rowmask = falses(size(A,1)) + @inbounds for r in rv + rowmask[r] = true + end - rv = unique(A isa SparseMatrixCSC ? A.rowval : A.A.rowval) + ฮฑ = .7 + k = 0 - # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique - @inbounds for row in rv + @inbounds for row in eachindex(rowmask) + rowmask[row] || continue @views copyto!(Aฬ„, A[row, :]) โ„’.mul!(Aฬ„B, Aฬ„, B) โ„’.mul!(CAฬ„B, C', Aฬ„B) copyto!(vCAฬ„B, CAฬ„B) - @views โ„’.mul!(X[row,:], D', vCAฬ„B) + โ„’.mul!(vCAฬ„BD, D', vCAฬ„B) + + if sparse + for (i,v) in enumerate(vCAฬ„BD) + if abs(v) > eps() + k += 1 + + if k > estimated_nnz + increment = max(10000, Int(ceil((ฮฑ - 1) * estimated_nnz + (1 - ฮฑ) * size(A,1) * size(D,2)))) + estimated_nnz += min(size(A,1) * size(D,2), increment) + + resize!(I, estimated_nnz) + resize!(J, estimated_nnz) + resize!(V, estimated_nnz) + end + + I[k] = row + J[k] = i + V[k] = v + end + end + else + @views copyto!(X[row,:], vCAฬ„BD) + end end - return choose_matrix_format(X) - # โ„’.mul!(vCAฬ„BD, D', vCAฬ„B) + if sparse + resize!(I, k) + resize!(J, k) + resize!(V, k) - # for (i,v) in enumerate(vCAฬ„BD) - # if abs(v) > eps() - # push!(rows, row) - # push!(cols, i) - # push!(vals, v) - # end - # end - # end + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] - # if VERSION >= v"1.10" - # return sparse!(rows, cols, vals, size(A,1), size(D,2)) - # else - # return sparse(rows, cols, vals, size(A,1), size(D,2)) - # end + resize!(klasttouch, size(D,2)) + resize!(csrrowptr, size(A, 1) + 1) + resize!(csrcolval, length(I)) + resize!(csrnzval, length(I)) + + if length(I) >= size(D,2) + 1 + out = sparse!(I, J, V, size(A, 1), size(D,2), +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) + else + out = SparseArrays.sparse(I, J, V, size(A, 1), size(D,2)) + end + # if reused_sparse_buffers + # out = copy(out) + # end + else + out = choose_matrix_format(X) + end + + return out end @@ -1659,6 +1579,7 @@ function mat_mult_kron(A::AbstractSparseMatrix{R}, J = Vector{Int}() V = Vector{T}() X = zeros(T, 0, 0) + reused_sparse_buffers = sparse && length(sparse_preallocation[1]) > 0 if sparse nnzA = nnz(A) @@ -1681,6 +1602,8 @@ function mat_mult_kron(A::AbstractSparseMatrix{R}, estimated_nnz = length(sparse_preallocation[3]) resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) I = sparse_preallocation[1] J = sparse_preallocation[2] @@ -1743,7 +1666,14 @@ function mat_mult_kron(A::AbstractSparseMatrix{R}, resize!(csrcolval, length(I)) resize!(csrnzval, length(I)) - out = sparse!(I, J, V, size(A, 1), n_colB * n_colC, +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) + if length(I) >= n_colB * n_colC + 1 + out = sparse!(I, J, V, size(A, 1), n_colB * n_colC, +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) + else + out = SparseArrays.sparse(I, J, V, size(A, 1), n_colB * n_colC) + end + # if reused_sparse_buffers + # out = copy(out) + # end # out = sparse!(I, J, V, size(A, 1), n_colB * n_colC) else out = choose_matrix_format(X) @@ -1856,61 +1786,679 @@ function sparse_preallocated!(Sฬ‚::Matrix{T}; โ„‚::higher_order_workspace{T,F,H} end -function compressed_kronยณ(a::AbstractMatrix{T}; - rowmask::Vector{Int} = Int[], - colmask::Vector{Int} = Int[], - # timer::TimerOutput = TimerOutput(), +# Loop-based compressed permuted mixed Kronecker product. +# Computes Uโ‚ƒ * (kron(A,ฯƒ) + Pโ‚โ‚—ฬ„*kron(A,ฯƒ)*Pโ‚แตฃฬƒ + Pโ‚‚โ‚—ฬ„*kron(A,ฯƒ)*Pโ‚‚แตฃฬƒ) * Cโ‚ƒ +# directly in compressed (sorted-triple) space without forming any nยณร—nยณ intermediates. +# +# A is nrร—nc (may be rectangular), ฯƒ is nrยฒร—ncยฒ. +# Output is mrโ‚ƒร—mcโ‚ƒ sparse where mrโ‚ƒ = nr(nr+1)(nr+2)/6, mcโ‚ƒ = nc(nc+1)(nc+2)/6. +# +# The uncompressed entry at row (i,j,k) col (a,b,c) of the sum is: +# A[i,a]*ฯƒ[(j-1)*nr+k,(b-1)*nc+c] (identity) +# + A[j,b]*ฯƒ[(i-1)*nr+k,(a-1)*nc+c] (Pโ‚: swap iโ†”j rows, aโ†”b cols) +# + A[j,b]*ฯƒ[(k-1)*nr+i,(c-1)*nc+a] (Pโ‚‚: cycle (i,j,k)โ†’(j,k,i), (a,b,c)โ†’(b,c,a)) +# +# Compression: Uโ‚ƒ sums over all row permutations that sort to (iโ‚โ‰ฅjโ‚โ‰ฅkโ‚); +# Cโ‚ƒ selects the sorted column representative (ฮฑโ‰ฅฮฒโ‰ฅฮณ). +function compressed_permuted_mixed_kron(A::AbstractMatrix{T}, ฯƒ::AbstractMatrix; tol::AbstractFloat = eps(), sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{T}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{T}} = (Int[], Int[], T[], Int[], Int[], Int[], T[])) where T <: Real - # @timeit_debug timer "Compressed 3rd kronecker power" begin - - # @timeit_debug timer "Preallocation" begin - - a_is_adjoint = typeof(a) <: โ„’.Adjoint{T,Matrix{T}} - - if a_is_adjoint - aฬ‚ = copy(a') - a = sparse(a') - - rmask = colmask - colmask = rowmask - rowmask = rmask - elseif typeof(a) <: DenseMatrix{T} - aฬ‚ = copy(a) - a = sparse(a) - else - aฬ‚ = convert(Matrix, a) # Convert to dense matrix for faster access - end - # Get the number of rows and columns - n_rows, n_cols = size(a) - - # Calculate the number of unique triplet indices for rows and columns - m3_rows = n_rows * (n_rows + 1) * (n_rows + 2) รท 6 # For rows: i โ‰ค j โ‰ค k - m3_cols = n_cols * (n_cols + 1) * (n_cols + 2) รท 6 # For columns: i โ‰ค j โ‰ค k - if rowmask == Int[0] || colmask == Int[0] - if a_is_adjoint - return spzeros(T, m3_cols, m3_rows) - else - return spzeros(T, m3_rows, m3_cols) - end - end - # Initialize arrays to collect indices and values - # Estimate an upper bound for non-zero entries to preallocate arrays - lennz = nnz(a) # a isa ThreadedSparseArrays.ThreadedSparseMatrixCSC ? length(a.A.nzval) : length(a.nzval) + nr = size(A, 1) + nc = size(A, 2) + size(ฯƒ) == (nr^2, nc^2) || throw(DimensionMismatch("ฯƒ must be $(nr^2)ร—$(nc^2), got $(size(ฯƒ))")) - m3_c = length(colmask) > 0 ? length(colmask) : m3_cols - m3_r = length(rowmask) > 0 ? length(rowmask) : m3_rows + # Sparse copies for support-aware iteration. + As = A isa SparseMatrixCSC{T, Int} ? A : sparse(T.(A)) + ฯƒs = ฯƒ isa SparseMatrixCSC{T, Int} ? ฯƒ : sparse(T.(ฯƒ)) - m3_exp = (length(colmask) > 0 || length(rowmask) > 0) ? 3 : 4 + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + + ranges_A = Vector{UnitRange{Int}}(undef, nc) + ranges_ฯƒ = Vector{UnitRange{Int}}(undef, nc^2) + @inbounds for col in 1:nc + ranges_A[col] = SparseArrays.nzrange(As, col) + end + @inbounds for col in 1:(nc^2) + ranges_ฯƒ[col] = SparseArrays.nzrange(ฯƒs, col) + end + + mrโ‚ƒ = nr * (nr + 1) * (nr + 2) รท 6 + mcโ‚ƒ = nc * (nc + 1) * (nc + 2) รท 6 + # --- sparse buffer management (same pattern as compressed_kronยณ) --- if length(sparse_preallocation[1]) == 0 - estimated_nnz = floor(Int, max(m3_r * m3_c * (lennz / length(a)) ^ m3_exp, 10000)) + estimated_nnz = max(min(mrโ‚ƒ, mcโ‚ƒ), 10000) resize!(sparse_preallocation[1], estimated_nnz) resize!(sparse_preallocation[2], estimated_nnz) resize!(sparse_preallocation[3], estimated_nnz) - + else + estimated_nnz = length(sparse_preallocation[3]) + + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + end + + II = sparse_preallocation[1] + JJ = sparse_preallocation[2] + VV = sparse_preallocation[3] + + cnt = 0 # non-zero counter + + # Iterate sorted output columns first (ฮฑ โ‰ฅ ฮฒ โ‰ฅ ฮณ). For each column triple, + # only traverse non-zero supports from the relevant A and ฯƒ columns. + for ฮฑ in 1:nc + rng_Aฮฑ = ranges_A[ฮฑ] + for ฮฒ in 1:ฮฑ + rng_Aฮฒ = ranges_A[ฮฒ] + for ฮณ in 1:ฮฒ + rng_Aฮณ = ranges_A[ฮณ] + + ฯƒ_col_ฮฒฮณ = (ฮฒ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮณ = (ฮฑ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮฒ = (ฮฑ - 1) * nc + ฮฒ + + rng_ฯƒฮฒฮณ = ranges_ฯƒ[ฯƒ_col_ฮฒฮณ] + rng_ฯƒฮฑฮณ = ranges_ฯƒ[ฯƒ_col_ฮฑฮณ] + rng_ฯƒฮฑฮฒ = ranges_ฯƒ[ฯƒ_col_ฮฑฮฒ] + + has_t1 = !isempty(rng_Aฮฑ) && !isempty(rng_ฯƒฮฒฮณ) + has_t2 = !isempty(rng_Aฮฒ) && !isempty(rng_ฯƒฮฑฮณ) + has_t3 = !isempty(rng_Aฮณ) && !isempty(rng_ฯƒฮฑฮฒ) + + (has_t1 || has_t2 || has_t3) || continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + # term 1: A[p, ฮฑ] * ฯƒ[(q, r), (ฮฒ, ฮณ)] + if has_t1 + @inbounds for ia in rng_Aฮฑ + p = rv_A[ia] + a_val = nzv_A[ia] + + for is in rng_ฯƒฮฒฮณ + qr = rv_ฯƒ[is] + q = (qr - 1) รท nr + 1 + r = qr - (q - 1) * nr + + val = a_val * nzv_ฯƒ[is] + abs(val) > tol || continue + + i1 = p + j1 = q + k1 = r + + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + + cnt += 1 + if cnt > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * 0.1))) + estimated_nnz = min(mrโ‚ƒ * mcโ‚ƒ, estimated_nnz) + resize!(II, estimated_nnz) + resize!(JJ, estimated_nnz) + resize!(VV, estimated_nnz) + end + + II[cnt] = row + JJ[cnt] = col + VV[cnt] = val + end + end + end + + # term 2: A[q, ฮฒ] * ฯƒ[(p, r), (ฮฑ, ฮณ)] + if has_t2 + @inbounds for ia in rng_Aฮฒ + q = rv_A[ia] + a_val = nzv_A[ia] + + for is in rng_ฯƒฮฑฮณ + pr = rv_ฯƒ[is] + p = (pr - 1) รท nr + 1 + r = pr - (p - 1) * nr + + val = a_val * nzv_ฯƒ[is] + abs(val) > tol || continue + + i1 = p + j1 = q + k1 = r + + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + + cnt += 1 + if cnt > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * 0.1))) + estimated_nnz = min(mrโ‚ƒ * mcโ‚ƒ, estimated_nnz) + resize!(II, estimated_nnz) + resize!(JJ, estimated_nnz) + resize!(VV, estimated_nnz) + end + + II[cnt] = row + JJ[cnt] = col + VV[cnt] = val + end + end + end + + # term 3: A[r, ฮณ] * ฯƒ[(p, q), (ฮฑ, ฮฒ)] + if has_t3 + @inbounds for ia in rng_Aฮณ + r = rv_A[ia] + a_val = nzv_A[ia] + + for is in rng_ฯƒฮฑฮฒ + pq = rv_ฯƒ[is] + p = (pq - 1) รท nr + 1 + q = pq - (p - 1) * nr + + val = a_val * nzv_ฯƒ[is] + abs(val) > tol || continue + + i1 = p + j1 = q + k1 = r + + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + + cnt += 1 + if cnt > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * 0.1))) + estimated_nnz = min(mrโ‚ƒ * mcโ‚ƒ, estimated_nnz) + resize!(II, estimated_nnz) + resize!(JJ, estimated_nnz) + resize!(VV, estimated_nnz) + end + + II[cnt] = row + JJ[cnt] = col + VV[cnt] = val + end + end + end + end + end + end + + resize!(II, cnt) + resize!(JJ, cnt) + resize!(VV, cnt) + + # Assemble sparse matrix using preallocated CSR workspace + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] + + resize!(klasttouch, mcโ‚ƒ) + resize!(csrrowptr, mrโ‚ƒ + 1) + resize!(csrcolval, length(II)) + resize!(csrnzval, length(II)) + + out = if length(II) >= mcโ‚ƒ + 1 + sparse!(II, JJ, VV, mrโ‚ƒ, mcโ‚ƒ, +, klasttouch, csrrowptr, csrcolval, csrnzval, II, JJ, VV) + else + SparseArrays.sparse(II, JJ, VV, mrโ‚ƒ, mcโ‚ƒ) + end + + if tol > 0 + droptol!(out, tol) + end + + return out +end + +# Fused M * compressed_permuted_mixed_kron(A, ฯƒ) +# Computes the product without materializing the large mrโ‚ƒร—mcโ‚ƒ intermediate. +# M is m ร— mrโ‚ƒ sparse, A is nr ร— nc, ฯƒ is nrยฒ ร— ncยฒ. Output: m ร— mcโ‚ƒ sparse. +function mul_compressed_permuted_mixed_kron(M::SparseMatrixCSC, A::AbstractMatrix{T}, ฯƒ::AbstractMatrix; + tol::AbstractFloat = eps(), + sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{T}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{T}} = (Int[], Int[], T[], Int[], Int[], Int[], T[])) where T <: Real + + nr = size(A, 1) + nc = size(A, 2) + m = size(M, 1) + mrโ‚ƒ = nr * (nr + 1) * (nr + 2) รท 6 + mcโ‚ƒ = nc * (nc + 1) * (nc + 2) รท 6 + + size(ฯƒ) == (nr^2, nc^2) || throw(DimensionMismatch("ฯƒ must be $(nr^2)ร—$(nc^2), got $(size(ฯƒ))")) + size(M, 2) == mrโ‚ƒ || throw(DimensionMismatch("M must have $mrโ‚ƒ columns, got $(size(M, 2))")) + + # Sparse copies for support-aware iteration + As = A isa SparseMatrixCSC{T, Int} ? A : sparse(T.(A)) + ฯƒs = ฯƒ isa SparseMatrixCSC{T, Int} ? ฯƒ : sparse(T.(ฯƒ)) + + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + rv_M = SparseArrays.rowvals(M) + nzv_M = nonzeros(M) + + ranges_A = Vector{UnitRange{Int}}(undef, nc) + ranges_ฯƒ = Vector{UnitRange{Int}}(undef, nc^2) + @inbounds for col in 1:nc + ranges_A[col] = SparseArrays.nzrange(As, col) + end + @inbounds for col in 1:(nc^2) + ranges_ฯƒ[col] = SparseArrays.nzrange(ฯƒs, col) + end + + # Small result buffer (size m, not mrโ‚ƒ) + result_col = zeros(T, m) + + # --- sparse IJV buffer management --- + if length(sparse_preallocation[1]) == 0 + estimated_nnz = max(min(m * mcโ‚ƒ รท 4, m * mcโ‚ƒ), 10000) + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + else + estimated_nnz = length(sparse_preallocation[3]) + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + end + + II = sparse_preallocation[1] + JJ = sparse_preallocation[2] + VV = sparse_preallocation[3] + cnt = 0 + + for ฮฑ in 1:nc + rng_Aฮฑ = ranges_A[ฮฑ] + for ฮฒ in 1:ฮฑ + rng_Aฮฒ = ranges_A[ฮฒ] + for ฮณ in 1:ฮฒ + rng_Aฮณ = ranges_A[ฮณ] + + ฯƒ_col_ฮฒฮณ = (ฮฒ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮณ = (ฮฑ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮฒ = (ฮฑ - 1) * nc + ฮฒ + + rng_ฯƒฮฒฮณ = ranges_ฯƒ[ฯƒ_col_ฮฒฮณ] + rng_ฯƒฮฑฮณ = ranges_ฯƒ[ฯƒ_col_ฮฑฮณ] + rng_ฯƒฮฑฮฒ = ranges_ฯƒ[ฯƒ_col_ฮฑฮฒ] + + has_t1 = !isempty(rng_Aฮฑ) && !isempty(rng_ฯƒฮฒฮณ) + has_t2 = !isempty(rng_Aฮฒ) && !isempty(rng_ฯƒฮฑฮณ) + has_t3 = !isempty(rng_Aฮณ) && !isempty(rng_ฯƒฮฑฮฒ) + + (has_t1 || has_t2 || has_t3) || continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + fill!(result_col, zero(T)) + + # term 1: A[p, ฮฑ] * ฯƒ[(q,r), (ฮฒ,ฮณ)] โ€” scatter through M + if has_t1 + @inbounds for ia in rng_Aฮฑ + p = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฒฮณ + qr = rv_ฯƒ[is] + q = (qr - 1) รท nr + 1 + r = qr - (q - 1) * nr + val = a_val * nzv_ฯƒ[is] + abs(val) > tol || continue + i1 = p; j1 = q; k1 = r + if i1 < j1; i1, j1 = j1, i1; end + if j1 < k1; j1, k1 = k1, j1; end + if i1 < j1; i1, j1 = j1, i1; end + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + rng_M = SparseArrays.nzrange(M, row) + for p_M in rng_M + result_col[rv_M[p_M]] += nzv_M[p_M] * val + end + end + end + end + + # term 2: A[q, ฮฒ] * ฯƒ[(p,r), (ฮฑ,ฮณ)] โ€” scatter through M + if has_t2 + @inbounds for ia in rng_Aฮฒ + q = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฑฮณ + pr = rv_ฯƒ[is] + p = (pr - 1) รท nr + 1 + r = pr - (p - 1) * nr + val = a_val * nzv_ฯƒ[is] + abs(val) > tol || continue + i1 = p; j1 = q; k1 = r + if i1 < j1; i1, j1 = j1, i1; end + if j1 < k1; j1, k1 = k1, j1; end + if i1 < j1; i1, j1 = j1, i1; end + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + rng_M = SparseArrays.nzrange(M, row) + for p_M in rng_M + result_col[rv_M[p_M]] += nzv_M[p_M] * val + end + end + end + end + + # term 3: A[r, ฮณ] * ฯƒ[(p,q), (ฮฑ,ฮฒ)] โ€” scatter through M + if has_t3 + @inbounds for ia in rng_Aฮณ + r = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฑฮฒ + pq = rv_ฯƒ[is] + p = (pq - 1) รท nr + 1 + q = pq - (p - 1) * nr + val = a_val * nzv_ฯƒ[is] + abs(val) > tol || continue + i1 = p; j1 = q; k1 = r + if i1 < j1; i1, j1 = j1, i1; end + if j1 < k1; j1, k1 = k1, j1; end + if i1 < j1; i1, j1 = j1, i1; end + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + rng_M = SparseArrays.nzrange(M, row) + for p_M in rng_M + result_col[rv_M[p_M]] += nzv_M[p_M] * val + end + end + end + end + + # Extract nonzeros into IJV + @inbounds for i in 1:m + v = result_col[i] + if abs(v) > tol + cnt += 1 + if cnt > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * 0.1))) + estimated_nnz = min(m * mcโ‚ƒ, estimated_nnz) + resize!(II, estimated_nnz) + resize!(JJ, estimated_nnz) + resize!(VV, estimated_nnz) + end + II[cnt] = i + JJ[cnt] = col + VV[cnt] = v + end + end + end + end + end + + resize!(II, cnt) + resize!(JJ, cnt) + resize!(VV, cnt) + + # Sparse assembly + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] + + resize!(klasttouch, mcโ‚ƒ) + resize!(csrrowptr, m + 1) + resize!(csrcolval, length(II)) + resize!(csrnzval, length(II)) + + out = if length(II) >= mcโ‚ƒ + 1 + sparse!(II, JJ, VV, m, mcโ‚ƒ, +, klasttouch, csrrowptr, csrcolval, csrnzval, II, JJ, VV) + else + SparseArrays.sparse(II, JJ, VV, m, mcโ‚ƒ) + end + + if tol > 0 + droptol!(out, tol) + end + + return out +end + +# 2-arg overload: compressed_kron(A, ฯƒ) +# Computes ๐”โˆ‡โ‚ƒ * kron(A, ฯƒ) * ๐‚โ‚ƒ +# directly in compressed (sorted-triple) space without forming any nยณร—nยณ intermediates. +# +# A is nแตฃ ร— nแถœ (may be rectangular), ฯƒ is nแตฃยฒ ร— nแถœยฒ. +# Output is mโ‚ƒแตฃ ร— mโ‚ƒแถœ sparse where mโ‚ƒแตฃ = nแตฃ(nแตฃ+1)(nแตฃ+2)/6, mโ‚ƒแถœ = nแถœ(nแถœ+1)(nแถœ+2)/6. +# +# kron(A,ฯƒ) at row (i,j,k) col (a,b,c) equals A[i,a]*ฯƒ[(j-1)*nแตฃ+k, (b-1)*nแถœ+c]. +# ๐”โˆ‡โ‚ƒ sums all row triples that sort to (iโ‚โ‰ฅjโ‚โ‰ฅkโ‚); ๐‚โ‚ƒ selects the sorted column (ฮฑโ‰ฅฮฒโ‰ฅฮณ). +function compressed_kron(A::AbstractMatrix{TA}, + ฯƒ::AbstractMatrix{Tฯƒ}; + tol::AbstractFloat = eps(), + sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{<:Real}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{<:Real}} = (Int[], Int[], Float64[], Int[], Int[], Int[], Float64[])) where {TA <: Real, Tฯƒ <: Real} + + T = promote_type(TA, Tฯƒ) + + nแตฃ, nแถœ = size(A) + size(ฯƒ) == (nแตฃ^2, nแถœ^2) || throw(DimensionMismatch("ฯƒ must be $(nแตฃ^2)ร—$(nแถœ^2), got $(size(ฯƒ))")) + + mโ‚ƒแตฃ = nแตฃ * (nแตฃ + 1) * (nแตฃ + 2) รท 6 + mโ‚ƒแถœ = nแถœ * (nแถœ + 1) * (nแถœ + 2) รท 6 + + # Convert to sparse for CSC iteration + As = A isa SparseMatrixCSC ? A : sparse(A) + ฯƒs = ฯƒ isa SparseMatrixCSC ? ฯƒ : sparse(ฯƒ) + + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + + # --- sparse buffer management --- + spI = sparse_preallocation[1] + spJ = sparse_preallocation[2] + spV_untyped = sparse_preallocation[3] + spV = if eltype(spV_untyped) == T + spV_untyped + else + Vector{T}(undef, length(spV_untyped)) + end + + lennz_A = nnz(As) + lennz_ฯƒ = nnz(ฯƒs) + len_A = length(A) + len_ฯƒ = length(ฯƒ) + + avg_density = sqrt((lennz_A / max(len_A, 1)) * (lennz_ฯƒ / max(len_ฯƒ, 1))) + + if length(spI) == 0 + estimated_nnz = floor(Int, max(mโ‚ƒแตฃ * mโ‚ƒแถœ * avg_density ^ 3, 10000)) + resize!(spI, estimated_nnz) + resize!(spJ, estimated_nnz) + resize!(spV, estimated_nnz) + else + estimated_nnz = length(spV) + resize!(spI, estimated_nnz) + resize!(spJ, estimated_nnz) + resize!(spV, estimated_nnz) + end + + II = spI + JJ = spJ + VV = spV + + cnt = 0 + + # Iterate sorted column triples (ฮฑ โ‰ฅ ฮฒ โ‰ฅ ฮณ) where ฮฑ indexes A's columns + # and (ฮฒ, ฮณ) index ฯƒ's columns via ฯƒ_col = (ฮฒ-1)*nแถœ + ฮณ. + for ฮฑ in 1:nแถœ + rng_A = SparseArrays.nzrange(As, ฮฑ) + isempty(rng_A) && continue + + for ฮฒ in 1:ฮฑ + for ฮณ in 1:ฮฒ + ฯƒ_col = (ฮฒ - 1) * nแถœ + ฮณ + rng_ฯƒ = SparseArrays.nzrange(ฯƒs, ฯƒ_col) + isempty(rng_ฯƒ) && continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + @inbounds for pA in rng_A + i = rv_A[pA] + a_val = nzv_A[pA] + + for pฯƒ in rng_ฯƒ + s = rv_ฯƒ[pฯƒ] + ฯƒ_val = nzv_ฯƒ[pฯƒ] + + val = a_val * ฯƒ_val + abs(val) > tol || continue + + # Decompose ฯƒ row: s = (j-1)*nแตฃ + k + j = (s - 1) รท nแตฃ + 1 + k = (s - 1) % nแตฃ + 1 + + # Sort row triple (i, j, k) โ†’ (iโ‚ โ‰ฅ jโ‚ โ‰ฅ kโ‚) + iโ‚ = i; jโ‚ = j; kโ‚ = k + if iโ‚ < jโ‚; iโ‚, jโ‚ = jโ‚, iโ‚; end + if jโ‚ < kโ‚; jโ‚, kโ‚ = kโ‚, jโ‚; end + if iโ‚ < jโ‚; iโ‚, jโ‚ = jโ‚, iโ‚; end + + row = (iโ‚ - 1) * iโ‚ * (iโ‚ + 1) รท 6 + (jโ‚ - 1) * jโ‚ รท 2 + kโ‚ + + cnt += 1 + + if cnt > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * 0.1))) + estimated_nnz = min(mโ‚ƒแตฃ * mโ‚ƒแถœ, estimated_nnz) + resize!(II, estimated_nnz) + resize!(JJ, estimated_nnz) + resize!(VV, estimated_nnz) + end + + II[cnt] = row + JJ[cnt] = col + VV[cnt] = val + end + end + end + end + end + + resize!(II, cnt) + resize!(JJ, cnt) + resize!(VV, cnt) + + # Sparse assembly with preallocation buffers + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval_untyped = sparse_preallocation[7] + csrnzval = if eltype(csrnzval_untyped) == T + csrnzval_untyped + else + Vector{T}(undef, length(csrnzval_untyped)) + end + + resize!(klasttouch, mโ‚ƒแถœ) + resize!(csrrowptr, mโ‚ƒแตฃ + 1) + resize!(csrcolval, length(II)) + resize!(csrnzval, length(II)) + + out = if cnt >= mโ‚ƒแถœ + 1 + sparse!(II, JJ, VV, mโ‚ƒแตฃ, mโ‚ƒแถœ, +, klasttouch, csrrowptr, csrcolval, csrnzval, II, JJ, VV) + else + SparseArrays.sparse(II, JJ, VV, mโ‚ƒแตฃ, mโ‚ƒแถœ) + end + + if tol > 0 + droptol!(out, tol) + end + + return out +end + + +function compressed_kronยณ(a::AbstractMatrix{T}; + rowmask::Vector{Int} = Int[], + colmask::Vector{Int} = Int[], + # timer::TimerOutput = TimerOutput(), + tol::AbstractFloat = eps(), + sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{T}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{T}} = (Int[], Int[], T[], Int[], Int[], Int[], T[])) where T <: Real + # @timeit_debug timer "Compressed 3rd kronecker power" begin + + # @timeit_debug timer "Preallocation" begin + + a_is_adjoint = typeof(a) <: โ„’.Adjoint{T,Matrix{T}} + reused_sparse_buffers = length(sparse_preallocation[1]) > 0 + + if a_is_adjoint + aฬ‚ = copy(a') + a = sparse(a') + + rmask = colmask + colmask = rowmask + rowmask = rmask + elseif typeof(a) <: DenseMatrix{T} + aฬ‚ = copy(a) + a = sparse(a) + else + aฬ‚ = convert(Matrix, a) # Convert to dense matrix for faster access + end + # Get the number of rows and columns + n_rows, n_cols = size(a) + + # Calculate the number of unique triplet indices for rows and columns + m3_rows = n_rows * (n_rows + 1) * (n_rows + 2) รท 6 # For rows: i โ‰ค j โ‰ค k + m3_cols = n_cols * (n_cols + 1) * (n_cols + 2) รท 6 # For columns: i โ‰ค j โ‰ค k + + if rowmask == Int[0] || colmask == Int[0] + if a_is_adjoint + return spzeros(T, m3_cols, m3_rows) + else + return spzeros(T, m3_rows, m3_cols) + end + end + # Initialize arrays to collect indices and values + # Estimate an upper bound for non-zero entries to preallocate arrays + lennz = nnz(a) # a isa ThreadedSparseArrays.ThreadedSparseMatrixCSC ? length(a.A.nzval) : length(a.nzval) + + m3_c = length(colmask) > 0 ? length(colmask) : m3_cols + m3_r = length(rowmask) > 0 ? length(rowmask) : m3_rows + + m3_exp = (length(colmask) > 0 || length(rowmask) > 0) ? 3 : 4 + + if length(sparse_preallocation[1]) == 0 + estimated_nnz = floor(Int, max(m3_r * m3_c * (lennz / length(a)) ^ m3_exp, 10000)) + + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + I = sparse_preallocation[1] J = sparse_preallocation[2] V = sparse_preallocation[3] @@ -1918,6 +2466,8 @@ function compressed_kronยณ(a::AbstractMatrix{T}; estimated_nnz = length(sparse_preallocation[3]) resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) I = sparse_preallocation[1] J = sparse_preallocation[2] @@ -1947,6 +2497,23 @@ function compressed_kronยณ(a::AbstractMatrix{T}; # Threads.@threads for i1 in ui norowmask = length(rowmask) == 0 nocolmask = length(colmask) == 0 + rowmask_lookup = norowmask ? BitVector() : falses(m3_rows) + colmask_lookup = nocolmask ? BitVector() : falses(m3_cols) + + if !norowmask && rowmask != Int[0] + @inbounds for r in rowmask + if 1 <= r <= m3_rows + rowmask_lookup[r] = true + end + end + end + if !nocolmask && colmask != Int[0] + @inbounds for c in colmask + if 1 <= c <= m3_cols + colmask_lookup[c] = true + end + end + end for i1 in ui for j1 in ui @@ -1956,7 +2523,7 @@ function compressed_kronยณ(a::AbstractMatrix{T}; row = (i1-1) * i1 * (i1+1) รท 6 + (j1-1) * j1 รท 2 + k1 - if norowmask || row in rowmask + if norowmask || rowmask_lookup[row] for i2 in uj for j2 in uj if j2 โ‰ค i2 @@ -1965,7 +2532,7 @@ function compressed_kronยณ(a::AbstractMatrix{T}; col = (i2-1) * i2 * (i2+1) รท 6 + (j2-1) * j2 รท 2 + k2 - if nocolmask || col in colmask + if nocolmask || colmask_lookup[col] # @timeit_debug timer "Multiplication" begin @inbounds aii = aฬ‚[i1, i2] @inbounds aij = aฬ‚[i1, j2] @@ -2095,27 +2662,358 @@ function compressed_kronยณ(a::AbstractMatrix{T}; # out = sparse!(I, J, V, m3_rows, m3_cols) end + # if reused_sparse_buffers + # out = copy(out) + # end + return out end +# Fused M * compressed_kronยณ(a) +# Computes the product without materializing the large mrโ‚ƒร—mcโ‚ƒ intermediate. +# M is m ร— mrโ‚ƒ sparse, a is n_rows ร— n_cols. Output: m ร— mcโ‚ƒ sparse. +# Row-outer / col-inner with sorted bounded ranges + direct IJV scatter. +# nzrange(M, row) checked once per row triple โ€” skips ALL col iterations. +# Duplicate (I,J) entries resolved by sparse!(+). +function mul_compressed_kronยณ(M::SparseMatrixCSC, a::AbstractMatrix{T}; + tol::AbstractFloat = eps(), + sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{T}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{T}} = (Int[], Int[], T[], Int[], Int[], Int[], T[])) where T <: Real -# function kronยณ(A::AbstractSparseMatrix{T}, Mโ‚ƒ::third_order) where T <: Real -# rows, cols, vals = findnz(A) + if typeof(a) <: DenseMatrix{T} + รข = a + a_sp = sparse(a) + else + รข = convert(Matrix, a) + a_sp = a isa SparseMatrixCSC ? a : sparse(a) + end -# # Dictionary to accumulate sums of values for each coordinate -# result_dict = Dict{Tuple{Int, Int}, T}() + n_rows, n_cols = size(a_sp) + m = size(M, 1) + m3_rows = n_rows * (n_rows + 1) * (n_rows + 2) รท 6 + m3_cols = n_cols * (n_cols + 1) * (n_cols + 2) รท 6 -# # Using a single iteration over non-zero elements -# nvals = length(vals) + size(M, 2) == m3_rows || throw(DimensionMismatch("M must have $m3_rows columns, got $(size(M, 2))")) -# lk = ReentrantLock() + rv_M = SparseArrays.rowvals(M) + nzv_M = nonzeros(M) -# Polyester.@batch for i in 1:nvals -# # for i in 1:nvals -# for j in 1:nvals -# for k in 1:nvals -# r1, c1, v1 = rows[i], cols[i], vals[i] -# r2, c2, v2 = rows[j], cols[j], vals[j] + # Find unique non-zero row and column indices (sorted for bounded iteration) + rowinds, colinds, _ = findnz(a_sp) + ui = sort!(unique(rowinds)) + uj = sort!(unique(colinds)) + n_ui = length(ui) + n_uj = length(uj) + + # --- sparse IJV buffer management --- + if length(sparse_preallocation[1]) == 0 + lennz = nnz(a_sp) + estimated_nnz = floor(Int, max(m * m3_cols * (lennz / length(a)) ^ 4, 10000)) + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + else + estimated_nnz = length(sparse_preallocation[3]) + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + end + + I = sparse_preallocation[1] + J = sparse_preallocation[2] + V = sparse_preallocation[3] + k = 0 + + # Row-outer loop: row triples (i1 โ‰ฅ j1 โ‰ฅ k1) with bounded index ranges + for idx_i1 in 1:n_ui + @inbounds i1 = ui[idx_i1] + for idx_j1 in 1:idx_i1 # j1 โ‰ค i1 by construction + @inbounds j1 = ui[idx_j1] + for idx_k1 in 1:idx_j1 # k1 โ‰ค j1 by construction + @inbounds k1 = ui[idx_k1] + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + + # nzrange checked ONCE per row triple โ€” skips ALL col iterations + rng_M = SparseArrays.nzrange(M, row) + isempty(rng_M) && continue + + # Divisor depends only on row triple + if i1 == j1 + divisor = i1 == k1 ? 6 : 2 + else + divisor = (i1 โ‰  k1 && j1 โ‰  k1) ? 1 : 2 + end + + # Col-inner loop: column triples (i2 โ‰ฅ j2 โ‰ฅ k2) with bounded ranges + for idx_i2 in 1:n_uj + @inbounds i2 = uj[idx_i2] + for idx_j2 in 1:idx_i2 # j2 โ‰ค i2 by construction + @inbounds j2 = uj[idx_j2] + for idx_k2 in 1:idx_j2 # k2 โ‰ค j2 by construction + @inbounds k2 = uj[idx_k2] + + @inbounds aii = รข[i1, i2] + @inbounds aij = รข[i1, j2] + @inbounds aik = รข[i1, k2] + @inbounds aji = รข[j1, i2] + @inbounds ajj = รข[j1, j2] + @inbounds ajk = รข[j1, k2] + @inbounds aki = รข[k1, i2] + @inbounds akj = รข[k1, j2] + @inbounds akk = รข[k1, k2] + + val = aii * (ajj * akk + ajk * akj) + aij * (aji * akk + ajk * aki) + aik * (aji * akj + ajj * aki) + + if abs(val) > tol + scaled_val = val / divisor + col = (i2 - 1) * i2 * (i2 + 1) รท 6 + (j2 - 1) * j2 รท 2 + k2 + + # Direct IJV scatter through M[:, row] + for p_M in rng_M + k += 1 + if k > estimated_nnz + estimated_nnz = k + max(1000, k รท 10) + resize!(I, estimated_nnz) + resize!(J, estimated_nnz) + resize!(V, estimated_nnz) + end + I[k] = @inbounds rv_M[p_M] + J[k] = col + V[k] = @inbounds(nzv_M[p_M]) * scaled_val + end + end + end + end + end + end + end + end + + resize!(I, k) + resize!(J, k) + resize!(V, k) + + # Sparse assembly โ€” sparse!(+) resolves duplicate (I,J) entries + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] + + resize!(klasttouch, m3_cols) + resize!(csrrowptr, m + 1) + resize!(csrcolval, length(I)) + resize!(csrnzval, length(I)) + + out = if length(I) >= m3_cols + 1 + sparse!(I, J, V, m, m3_cols, +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) + else + SparseArrays.sparse(I, J, V, m, m3_cols) + end + + if tol > 0 + droptol!(out, tol) + end + + return out +end + +function compressed_kronยฒ(a::AbstractMatrix{T}; + rowmask::Vector{Int} = Int[], + colmask::Vector{Int} = Int[], + tol::AbstractFloat = eps(), + sparse_preallocation::Tuple{Vector{Int}, Vector{Int}, Vector{T}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{T}} = (Int[], Int[], T[], Int[], Int[], Int[], T[])) where T <: Real + + a_is_adjoint = typeof(a) <: โ„’.Adjoint{T,Matrix{T}} + reused_sparse_buffers = length(sparse_preallocation[1]) > 0 + + if a_is_adjoint + รข = copy(a') + a = sparse(a') + + rmask = colmask + colmask = rowmask + rowmask = rmask + elseif typeof(a) <: DenseMatrix{T} + รข = copy(a) + a = sparse(a) + else + รข = convert(Matrix, a) # Convert to dense matrix for faster access + end + + # Get the number of rows and columns + n_rows, n_cols = size(a) + + # Calculate the number of unique pair indices for rows and columns + m2_rows = n_rows * (n_rows + 1) รท 2 # For rows: i โ‰ค j + m2_cols = n_cols * (n_cols + 1) รท 2 # For columns: i โ‰ค j + + if rowmask == Int[0] || colmask == Int[0] + if a_is_adjoint + return spzeros(T, m2_cols, m2_rows) + else + return spzeros(T, m2_rows, m2_cols) + end + end + + # Initialize arrays to collect indices and values + lennz = nnz(a) + + m2_c = length(colmask) > 0 ? length(colmask) : m2_cols + m2_r = length(rowmask) > 0 ? length(rowmask) : m2_rows + + m2_exp = (length(colmask) > 0 || length(rowmask) > 0) ? 2 : 3 + + if length(sparse_preallocation[1]) == 0 + estimated_nnz = floor(Int, max(m2_r * m2_c * (lennz / length(a)) ^ m2_exp, 10000)) + + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + + I = sparse_preallocation[1] + J = sparse_preallocation[2] + V = sparse_preallocation[3] + else + estimated_nnz = length(sparse_preallocation[3]) + + resize!(sparse_preallocation[1], estimated_nnz) + resize!(sparse_preallocation[2], estimated_nnz) + resize!(sparse_preallocation[3], estimated_nnz) + + I = sparse_preallocation[1] + J = sparse_preallocation[2] + V = sparse_preallocation[3] + end + + k = 0 + + # Find unique non-zero row and column indices + rowinds, colinds, _ = findnz(a) + ui = unique(rowinds) + uj = unique(colinds) + + norowmask = length(rowmask) == 0 + nocolmask = length(colmask) == 0 + rowmask_lookup = norowmask ? BitVector() : falses(m2_rows) + colmask_lookup = nocolmask ? BitVector() : falses(m2_cols) + + if !norowmask && rowmask != Int[0] + @inbounds for r in rowmask + if 1 <= r <= m2_rows + rowmask_lookup[r] = true + end + end + end + if !nocolmask && colmask != Int[0] + @inbounds for c in colmask + if 1 <= c <= m2_cols + colmask_lookup[c] = true + end + end + end + + for i1 in ui + for j1 in ui + if j1 โ‰ค i1 + + row = (i1 - 1) * i1 รท 2 + j1 + + if norowmask || rowmask_lookup[row] + for i2 in uj + for j2 in uj + if j2 โ‰ค i2 + + col = (i2 - 1) * i2 รท 2 + j2 + + if nocolmask || colmask_lookup[col] + @inbounds aii = รข[i1, i2] + @inbounds aij = รข[i1, j2] + @inbounds aji = รข[j1, i2] + @inbounds ajj = รข[j1, j2] + + # Sum over both permutations of (i2, j2) + val = aii * ajj + aij * aji + + if abs(val) > tol + divisor = i1 == j1 ? 2 : 1 + + k += 1 + + if k > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * .1))) + estimated_nnz = min(m2_cols * m2_rows, estimated_nnz) + resize!(I, estimated_nnz) + resize!(J, estimated_nnz) + resize!(V, estimated_nnz) + end + + I[k] = row + J[k] = col + V[k] = val / divisor + end + end + end + end + end + end + end + end + end + + resize!(I, k) + resize!(J, k) + resize!(V, k) + + # Create the sparse matrix from the collected indices and values + if a_is_adjoint + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] + + resize!(klasttouch, m2_rows) + resize!(csrrowptr, m2_cols + 1) + resize!(csrcolval, length(J)) + resize!(csrnzval, length(J)) + + out = sparse!(J, I, V, m2_cols, m2_rows, +, klasttouch, csrrowptr, csrcolval, csrnzval, J, I, V) + else + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] + + resize!(klasttouch, m2_cols) + resize!(csrrowptr, m2_rows + 1) + resize!(csrcolval, length(I)) + resize!(csrnzval, length(I)) + + out = sparse!(I, J, V, m2_rows, m2_cols, +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) + end + + # if reused_sparse_buffers + # out = copy(out) + # end + + return out +end +# function kronยณ(A::AbstractSparseMatrix{T}, Mโ‚ƒ::third_order) where T <: Real +# rows, cols, vals = findnz(A) + +# # Dictionary to accumulate sums of values for each coordinate +# result_dict = Dict{Tuple{Int, Int}, T}() + +# # Using a single iteration over non-zero elements +# nvals = length(vals) + +# lk = ReentrantLock() + +# Polyester.@batch for i in 1:nvals +# # for i in 1:nvals +# for j in 1:nvals +# for k in 1:nvals +# r1, c1, v1 = rows[i], cols[i], vals[i] +# r2, c2, v2 = rows[j], cols[j], vals[j] # r3, c3, v3 = rows[k], cols[k], vals[k] # sorted_cols = [c1, c2, c3] @@ -2534,8 +3432,8 @@ end function determine_efficient_order(๐’โ‚::Matrix{<: Real}, - ๐’โ‚‚::AbstractMatrix{<: Real}, - ๐’โ‚ƒ::AbstractMatrix{<: Real}, + ๐’โ‚‚::AbstractSparseMatrix{<: Real}, + ๐’โ‚ƒ::AbstractSparseMatrix{<: Real}, constants::constants, variables::Union{Symbol_input,String_input}; covariance::Union{Symbol_input,String_input} = Symbol[], @@ -2566,68 +3464,68 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, # Precompute state indices and matrix slices state_idx_in_var = indexin(T.past_not_future_and_mixed, T.var) .|> Int ๐’โ‚_states = ๐’โ‚[state_idx_in_var, 1:nหข] - ๐’โ‚‚_states = nnz(๐’โ‚‚) > 0 ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing - ๐’โ‚ƒ_states = nnz(๐’โ‚ƒ) > 0 ? ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] : nothing + has_Sโ‚‚ = nnz(๐’โ‚‚) > 0 + has_Sโ‚ƒ = nnz(๐’โ‚ƒ) > 0 + ๐’โ‚‚_states = has_Sโ‚‚ ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing + ๐’โ‚ƒ_states = has_Sโ‚ƒ ? ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] : nothing - for obs in observables - obs_in_var_idx = indexin([obs],T.var) .|> Int - + function compute_dependencies(obs_in_var_idx::Vector{Int}) # First order dependencies dependencies_in_states = vec(sum(abs, ๐’โ‚[obs_in_var_idx,1:nหข], dims=1) .> tol) .> 0 - + # Second order dependencies from quadratic terms (s โŠ— s) - if nnz(๐’โ‚‚) > 0 + if has_Sโ‚‚ s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_var_idx, kron_s_s] - # Vectorized approach: reshape and check row/column sums s_s_matrix = reshape(vec(sum(abs, s_s_to_yโ‚‚, dims=1) .> tol), nหข, nหข) dependencies_in_states = dependencies_in_states .| vec(sum(s_s_matrix, dims=2) .> 0) .| vec(sum(s_s_matrix, dims=1) .> 0) end - + # Third order dependencies from cubic terms (s โŠ— s โŠ— s) - if nnz(๐’โ‚ƒ) > 0 + if has_Sโ‚ƒ s_s_s_to_yโ‚ƒ = ๐’โ‚ƒ[obs_in_var_idx, kron_s_s_s] - # Vectorized approach: reshape to 3D and check along dimensions s_s_s_tensor = reshape(vec(sum(abs, s_s_s_to_yโ‚ƒ, dims=1) .> tol), nหข, nหข, nหข) - dependencies_in_states = dependencies_in_states .| vec(sum(s_s_s_tensor, dims=(2,3)) .> 0) .| - vec(sum(s_s_s_tensor, dims=(1,3)) .> 0) .| + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_s_tensor, dims=(2,3)) .> 0) .| + vec(sum(s_s_s_tensor, dims=(1,3)) .> 0) .| vec(sum(s_s_s_tensor, dims=(1,2)) .> 0) end # Propagate dependencies through the system (iterative closure) - # considering first, second, and third order propagation while true prev_dependencies = dependencies_in_states - + # First order propagation new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states) .> tol) - + # Second order propagation if !isnothing(๐’โ‚‚_states) - # Generate selector vector for columns where both states are dependencies selector = vec(โ„’.kron(prev_dependencies, prev_dependencies)) if any(selector) affected = vec(sum(abs, ๐’โ‚‚_states[:, selector], dims=2) .> tol) new_deps = new_deps .| affected end end - + # Third order propagation if !isnothing(๐’โ‚ƒ_states) - # Generate selector vector for columns where all three states are dependencies selector = vec(โ„’.kron(โ„’.kron(prev_dependencies, prev_dependencies), prev_dependencies)) if any(selector) affected = vec(sum(abs, ๐’โ‚ƒ_states[:, selector], dims=2) .> tol) new_deps = new_deps .| affected end end - + if new_deps == dependencies_in_states break end dependencies_in_states = new_deps end - dependencies = T.past_not_future_and_mixed[dependencies_in_states] + return T.past_not_future_and_mixed[dependencies_in_states] + end + + for obs in observables + obs_in_var_idx = indexin([obs],T.var) .|> Int + dependencies = compute_dependencies(obs_in_var_idx) push!(orders,[obs] => sort(dependencies)) end @@ -2642,67 +3540,7 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, # Check if this variable's dependencies are already computed if isnothing(findfirst(x -> covar_var in x.first, orders)) obs_in_var_idx = indexin([covar_var], T.var) .|> Int - - # First order dependencies - dependencies_in_states = vec(sum(abs, ๐’โ‚[obs_in_var_idx,1:nหข], dims=1) .> tol) .> 0 - - # Second order dependencies from quadratic terms (s โŠ— s) - if nnz(๐’โ‚‚) > 0 - s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_var_idx, kron_s_s] - # Vectorized approach: reshape to nหขร—nหข and check column/row sums - s_s_matrix = reshape(vec(sum(abs, s_s_to_yโ‚‚, dims=1) .> tol), nหข, nหข) - dependencies_in_states = dependencies_in_states .| vec(sum(s_s_matrix, dims=2) .> 0) .| vec(sum(s_s_matrix, dims=1) .> 0) - end - - # Third order dependencies from cubic terms (s โŠ— s โŠ— s) - if nnz(๐’โ‚ƒ) > 0 - s_s_s_to_yโ‚ƒ = ๐’โ‚ƒ[obs_in_var_idx, kron_s_s_s] - # Vectorized approach: reshape to 3D and check along dimensions - s_s_s_tensor = reshape(vec(sum(abs, s_s_s_to_yโ‚ƒ, dims=1) .> tol), nหข, nหข, nหข) - dependencies_in_states = dependencies_in_states .| vec(sum(s_s_s_tensor, dims=(2,3)) .> 0) .| - vec(sum(s_s_s_tensor, dims=(1,3)) .> 0) .| - vec(sum(s_s_s_tensor, dims=(1,2)) .> 0) - end - - # Propagate dependencies through the system - # Precompute matrix slices - ๐’โ‚_states_local = ๐’โ‚[state_idx_in_var, 1:nหข] - ๐’โ‚‚_states_local = nnz(๐’โ‚‚) > 0 ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing - ๐’โ‚ƒ_states_local = nnz(๐’โ‚ƒ) > 0 ? ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] : nothing - - while true - prev_dependencies = dependencies_in_states - - # First order propagation - new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states_local) .> tol) - - # Second order propagation - if !isnothing(๐’โ‚‚_states_local) - # Generate selector vector for columns where both states are dependencies - selector = vec(โ„’.kron(prev_dependencies, prev_dependencies)) - if any(selector) - affected = vec(sum(abs, ๐’โ‚‚_states_local[:, selector], dims=2) .> tol) - new_deps = new_deps .| affected - end - end - - # Third order propagation - if !isnothing(๐’โ‚ƒ_states_local) - # Generate selector vector for columns where all three states are dependencies - selector = vec(โ„’.kron(โ„’.kron(prev_dependencies, prev_dependencies), prev_dependencies)) - if any(selector) - affected = vec(sum(abs, ๐’โ‚ƒ_states_local[:, selector], dims=2) .> tol) - new_deps = new_deps .| affected - end - end - - if new_deps == dependencies_in_states - break - end - dependencies_in_states = new_deps - end - - dependencies = T.past_not_future_and_mixed[dependencies_in_states] + dependencies = compute_dependencies(obs_in_var_idx) push!(orders,[covar_var] => sort(dependencies)) end end @@ -3004,7 +3842,7 @@ end function get_relevant_steady_states(๐“‚::โ„ณ, algorithm::Symbol; opts::CalculationOptions = merge_calculation_options())::Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}} - ms = @ignore_derivatives ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) full_NSSS = ms.full_NSSS_display relevant_SS = get_steady_state(๐“‚, algorithm = algorithm, @@ -3543,35 +4381,10 @@ function decompose_name(name::Symbol) return result end -""" - get_computational_constants(๐“‚::โ„ณ) - -Return cached second-order computational constants (BitVectors and index patterns). -""" -function get_computational_constants(๐“‚::โ„ณ) - ensure_computational_constants!(๐“‚) - return ๐“‚.constants.second_order -end +function get_possible_indices_for_name(name::Symbol, all_names::Vector{Symbol}) + indices = filter(x -> length(x) < 3 && x[1] == name, decompose_name.(all_names)) -function get_computational_constants(constants::constants) - ensure_computational_constants!(constants) - return constants.second_order -end - -""" - get_model_structure(๐“‚::โ„ณ) - -Return cached model structure information (SS_and_pars_names, all_variables, NSSS_labels). -""" -function get_model_structure(๐“‚::โ„ณ) - return ๐“‚.constants.post_complete_parameters -end - - -function get_possible_indices_for_name(name::Symbol, all_names::Vector{Symbol}) - indices = filter(x -> length(x) < 3 && x[1] == name, decompose_name.(all_names)) - - indexset = [] + indexset = [] for i in indices if length(i) > 1 @@ -3816,191 +4629,67 @@ function remove_redundant_SS_vars!(๐“‚::โ„ณ, Symbolics::symbolics; avoid_solve: end -function write_block_solution!(๐“‚, - SS_solve_func, - vars_to_solve, - eqs_to_solve, - relevant_pars_across, - NSSS_solver_cache_init_tmp, - eq_idx_in_block_to_solve, - atoms_in_equations_list, - solved_vars, - solved_vals; - cse = true, - skipzeros = true, - density_threshold::Float64 = .1, - nnz_parallel_threshold::Int = 1000000, - min_length::Int = 10000) - - # โž•_vars = Symbol[] - unique_โž•_eqs = Dict{Union{Expr,Symbol},Symbol}() - - vars_to_exclude = [vcat(Symbol.(vars_to_solve), ๐“‚.constants.post_model_macro.โž•_vars),Symbol[]] - - rewritten_eqs, ss_and_aux_equations, ss_and_aux_equations_dep, ss_and_aux_equations_error, ss_and_aux_equations_error_dep = make_equation_robust_to_domain_errors(Meta.parse.(string.(eqs_to_solve)), vars_to_exclude, ๐“‚.constants.post_parameters_macro.bounds, ๐“‚.constants.post_model_macro.โž•_vars, unique_โž•_eqs) - - - push!(solved_vars, Symbol.(vars_to_solve)) - push!(solved_vals, rewritten_eqs) - - - syms_in_eqs = Set{Symbol}() - - for i in vcat(ss_and_aux_equations_dep, ss_and_aux_equations, rewritten_eqs) - push!(syms_in_eqs, get_symbols(i)...) - end - - setdiff!(syms_in_eqs,๐“‚.constants.post_model_macro.โž•_vars) - - syms_in_eqs2 = Set{Symbol}() - - for i in ss_and_aux_equations - push!(syms_in_eqs2, get_symbols(i)...) - end - - โž•_vars_alread_in_eqs = intersect(๐“‚.constants.post_model_macro.โž•_vars,reduce(union,get_symbols.(Meta.parse.(string.(eqs_to_solve))))) - - union!(syms_in_eqs, intersect(union(โž•_vars_alread_in_eqs, syms_in_eqs2), ๐“‚.constants.post_model_macro.โž•_vars)) - - push!(atoms_in_equations_list,setdiff(syms_in_eqs, solved_vars[end])) - - # guess = Expr[] - # untransformed_guess = Expr[] - result = Expr[] - # calib_pars = Expr[] - - calib_pars_input = Symbol[] - - relevant_pars = union(intersect(reduce(union, vcat(๐“‚.constants.post_model_macro.par_list_aux_SS, ๐“‚.constants.post_parameters_macro.par_calib_list)[eq_idx_in_block_to_solve]), syms_in_eqs),intersect(syms_in_eqs, ๐“‚.constants.post_model_macro.โž•_vars)) - - union!(relevant_pars_across, relevant_pars) - - sorted_vars = sort(Symbol.(vars_to_solve)) - for (i, parss) in enumerate(sorted_vars) - # push!(guess,:($parss = guess[$i])) - # push!(untransformed_guess,:($parss = undo_transform(guess[$i],transformation_level))) - push!(result,:($parss = sol[$i])) - end - - iii = 1 - for parss in union(๐“‚.constants.post_complete_parameters.parameters, ๐“‚.constants.post_parameters_macro.parameters_as_function_of_parameters) - if :($parss) โˆˆ relevant_pars - # push!(calib_pars, :($parss = parameters_and_solved_vars[$iii])) - push!(calib_pars_input, :($parss)) - iii += 1 - end - end - - # separate out auxiliary variables (nonnegativity) - # nnaux = [] - # nnaux_linear = [] - # nnaux_error = [] - # push!(nnaux_error, :(aux_error = 0)) - # solved_vals_in_place = Expr[] - # partially_solved_block = Expr[] - - other_vrs_eliminated_by_sympy = Set{Symbol}() - - for (i,val) in enumerate(solved_vals[end]) - if eq_idx_in_block_to_solve[i] โˆˆ ๐“‚.constants.post_model_macro.ss_equations_with_aux_variables - val = vcat(๐“‚.equations.steady_state_aux, ๐“‚.equations.calibration)[eq_idx_in_block_to_solve[i]] - # push!(nnaux,:($(val.args[2]) = max(eps(),$(val.args[3])))) - push!(other_vrs_eliminated_by_sympy, val.args[2]) - # push!(nnaux_linear,:($val)) - # push!(nnaux_error, :(aux_error += min(eps(),$(val.args[3])))) - end - end - - - - solved_vals_local = Expr[] - for (i,val) in enumerate(rewritten_eqs) - push!(solved_vals_local, postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, val)) - # push!(solved_vals_in_place, :(โ„ฐ[$i] = $(postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, val)))) - end - - - # if length(nnaux) > 1 - # all_symbols = map(x->x.args[1],nnaux) #relevant symbols come first in respective equations - - # nn_symbols = map(x->intersect(all_symbols,x), get_symbols.(nnaux)) - - # inc_matrix = fill(0,length(all_symbols),length(all_symbols)) - - # for i in 1:length(all_symbols) - # for k in 1:length(nn_symbols) - # inc_matrix[i,k] = collect(all_symbols)[i] โˆˆ collect(nn_symbols)[k] - # end - # end - # QQ, P, R, nmatch, n_blocks = BlockTriangularForm.order(sparse(inc_matrix)) - - # nnaux = nnaux[QQ] - # nnaux_linear = nnaux_linear[QQ] - # end - - # other_vars = Expr[] - other_vars_input = Symbol[] - other_vrs = intersect( setdiff( union(๐“‚.constants.post_model_macro.var, ๐“‚.equations.calibration_parameters, ๐“‚.constants.post_model_macro.โž•_vars), - sort(solved_vars[end]) ), - union(syms_in_eqs, other_vrs_eliminated_by_sympy ) ) - # union(syms_in_eqs, other_vrs_eliminated_by_sympy, setdiff(reduce(union, get_symbols.(nnaux), init = []), map(x->x.args[1],nnaux)) ) ) +function write_ss_check_function!(๐“‚::โ„ณ; + cse = true, + skipzeros = true, + density_threshold::Float64 = .1, + nnz_parallel_threshold::Int = 1000000, + min_length::Int = 10000) + unknowns = union(setdiff(๐“‚.constants.post_model_macro.vars_in_ss_equations, ๐“‚.constants.post_model_macro.โž•_vars), ๐“‚.equations.calibration_parameters) - for var in other_vrs - # push!(other_vars,:($(var) = parameters_and_solved_vars[$iii])) - push!(other_vars_input,:($(var))) - iii += 1 - end + ss_equations = vcat(๐“‚.equations.steady_state, ๐“‚.equations.calibration) - parameters_and_solved_vars = vcat(calib_pars_input, other_vrs) - ng = length(sorted_vars) - np = length(parameters_and_solved_vars) - nd = length(ss_and_aux_equations_dep) - nx = iii - 1 - Symbolics.@variables ๐”Š[1:ng] ๐”“[1:np] + np = length(๐“‚.constants.post_complete_parameters.parameters) + nu = length(unknowns) + # nc = length(๐“‚.calibration_equations_no_var) + Symbolics.@variables ๐”“[1:np] ๐”˜[1:nu]# โ„ญ[1:nc] parameter_dict = Dict{Symbol, Symbol}() back_to_array_dict = Dict{Symbolics.Num, Symbolics.Num}() - aux_vars = Symbol[] - aux_expr = [] - + calib_vars = Symbol[] + calib_expr = [] - for (i,v) in enumerate(sorted_vars) - push!(parameter_dict, v => :($(Symbol("๐”Š_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”Š_$i"))), @__MODULE__) => ๐”Š[i]) - end - for (i,v) in enumerate(parameters_and_solved_vars) + for (i,v) in enumerate(๐“‚.constants.post_complete_parameters.parameters) push!(parameter_dict, v => :($(Symbol("๐”“_$i")))) push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”“_$i"))), @__MODULE__) => ๐”“[i]) end - for (i,v) in enumerate(ss_and_aux_equations_dep) - push!(aux_vars, v.args[1]) - push!(aux_expr, v.args[2]) + for (i,v) in enumerate(unknowns) + push!(parameter_dict, v => :($(Symbol("๐”˜_$i")))) + push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”˜_$i"))), @__MODULE__) => ๐”˜[i]) end - - aux_replacements = Dict{Symbol,Any}() - for (i,x) in enumerate(aux_vars) - replacement = Dict(x => aux_expr[i]) - for ii in i+1:length(aux_vars) - aux_expr[ii] = replace_symbols(aux_expr[ii], replacement) + + for (i,v) in enumerate(๐“‚.equations.calibration_no_var) + push!(calib_vars, v.args[1]) + push!(calib_expr, v.args[2]) + # push!(parameter_dict, v.args[1] => :($(Symbol("โ„ญ_$i")))) + # push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("โ„ญ_$i"))), @__MODULE__) => โ„ญ[i]) + end + + calib_replacements = Dict{Symbol, Union{Expr, Symbol, Number}}() + for (i,x) in enumerate(calib_vars) + replacement = Dict{Symbol, Union{Expr, Symbol, Number}}(x => calib_expr[i]) + for ii in i+1:length(calib_vars) + calib_expr[ii] = replace_symbols(calib_expr[ii], replacement) end - push!(aux_replacements, x => aux_expr[i]) + push!(calib_replacements, x => calib_expr[i]) end - # aux_replacements = Dict{Symbol,Any}(aux_vars .=> aux_expr) - replaced_solved_vals = solved_vals_local |> - x -> replace_symbols.(x, Ref(aux_replacements)) |> + + ss_equations_sub = ss_equations |> + x -> replace_symbols.(x, Ref(calib_replacements)) |> x -> replace_symbols.(x, Ref(parameter_dict)) |> x -> Symbolics.parse_expr_to_symbolic.(x, Ref(@__MODULE__)) |> x -> Symbolics.substitute.(x, Ref(back_to_array_dict)) - lennz = length(replaced_solved_vals) + + lennz = length(ss_equations_sub) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -4008,758 +4697,83 @@ function write_block_solution!(๐“‚, parallel = Symbolics.SerialForm() end - _, calc_block! = Symbolics.build_function(replaced_solved_vals, ๐”Š, ๐”“, + _, func_exprs = Symbolics.build_function(ss_equations_sub, ๐”“, ๐”˜, cse = cse, - skipzeros = skipzeros, - # nanmath = false, + skipzeros = skipzeros, + # nanmath = false, parallel = parallel, expression_module = @__MODULE__, expression = Val(false))::Tuple{<:Function, <:Function} - # ๐ท = zeros(Symbolics.Num, nd) - - # ฯตแตƒ = zeros(nd) - - # calc_block_aux!(๐ท, ๐”Š, ๐”“) - - ฯตหข = zeros(Symbolics.Num, ng) - - ฯต = zeros(ng) - - # calc_block!(ฯตหข, ๐”Š, ๐”“, ๐ท) - โˆ‚block_โˆ‚parameters_and_solved_vars = Symbolics.sparsejacobian(replaced_solved_vals, ๐”Š) # nฯต x nx - - lennz = nnz(โˆ‚block_โˆ‚parameters_and_solved_vars) + ๐“‚.functions.NSSS_check = func_exprs - if (lennz / length(โˆ‚block_โˆ‚parameters_and_solved_vars) > density_threshold) || (length(โˆ‚block_โˆ‚parameters_and_solved_vars) < min_length) - derivatives_mat = convert(Matrix, โˆ‚block_โˆ‚parameters_and_solved_vars) - buffer = zeros(Float64, size(โˆ‚block_โˆ‚parameters_and_solved_vars)) - else - derivatives_mat = โˆ‚block_โˆ‚parameters_and_solved_vars - buffer = similar(โˆ‚block_โˆ‚parameters_and_solved_vars, Float64) - buffer.nzval .= 1 + # Ensure check_residual buffer is sized for the NSSS_check function + nres = length(ss_equations) + cr = ๐“‚.workspaces.nsss_solver.check_residual + if length(cr) != nres + resize!(cr, nres) + fill!(cr, 0.0) end - chol_buff = buffer * buffer' - - chol_buff += โ„’.I - prob = ๐’ฎ.LinearProblem(chol_buff, ฯต, ๐’ฎ.CholeskyFactorization()) - - chol_buffer = ๐’ฎ.init(prob, ๐’ฎ.CholeskyFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + # SS_and_pars = Symbol.(vcat(string.(sort(collect(setdiff(reduce(union,get_symbols.(๐“‚.ss_aux_equations)),union(๐“‚.constants.post_model_macro.parameters_in_equations,๐“‚.constants.post_model_macro.โž•_vars))))), ๐“‚.calibration_equations_parameters)) - prob = ๐’ฎ.LinearProblem(buffer, ฯต, ๐’ฎ.LUFactorization()) + # eqs = vcat(๐“‚.ss_equations, ๐“‚.calibration_equations) - lu_buffer = ๐’ฎ.init(prob, ๐’ฎ.LUFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + # nx = length(๐“‚.parameter_values) - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”Š, ๐”“, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} + # np = length(SS_and_pars) + nฯตหข = length(ss_equations) - Symbolics.@variables ๐”Š[1:ng+nx] + # nc = length(๐“‚.calibration_equations_no_var) - ext_diff = Symbolics.Num[] - for i in 1:nx - push!(ext_diff, ๐”“[i] - ๐”Š[ng + i]) - end - replaced_solved_vals_ext = vcat(replaced_solved_vals, ext_diff) + # Symbolics.@variables ๐”›ยน[1:nx] ๐”“ยน[1:np] - _, calc_ext_block! = Symbolics.build_function(replaced_solved_vals_ext, ๐”Š, ๐”“, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} + # ฯตหข = zeros(Symbolics.Num, nฯตหข) - ฯตแต‰ = zeros(ng + nx) - - # ฯตหขแต‰ = zeros(Symbolics.Num, ng + nx) + # calib_vals = zeros(Symbolics.Num, nc) - # calc_block_aux!(๐ท, ๐”Š, ๐”“) + # ๐“‚.SS_calib_func(calib_vals, ๐”“) - # Evaluate the function symbolically - # calc_ext_block!(ฯตหขแต‰, ๐”Š, ๐”“, ๐ท) + # ๐“‚.functions.NSSS_check(ฯตหข, ๐”“, ๐”˜, calib_vals) - โˆ‚ext_block_โˆ‚parameters_and_solved_vars = Symbolics.sparsejacobian(replaced_solved_vals_ext, ๐”Š) # nฯต x nx + โˆ‚SS_equations_โˆ‚parameters = Symbolics.sparsejacobian(ss_equations_sub, ๐”“) # nฯต x nx - lennz = nnz(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) + lennz = nnz(โˆ‚SS_equations_โˆ‚parameters) - if (lennz / length(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) > density_threshold) || (length(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) < min_length) - derivatives_mat_ext = convert(Matrix, โˆ‚ext_block_โˆ‚parameters_and_solved_vars) - ext_buffer = zeros(Float64, size(โˆ‚ext_block_โˆ‚parameters_and_solved_vars)) + if (lennz / length(โˆ‚SS_equations_โˆ‚parameters) > density_threshold) || (length(โˆ‚SS_equations_โˆ‚parameters) < min_length) + derivatives_mat = convert(Matrix, โˆ‚SS_equations_โˆ‚parameters) + buffer = zeros(Float64, size(โˆ‚SS_equations_โˆ‚parameters)) else - derivatives_mat_ext = โˆ‚ext_block_โˆ‚parameters_and_solved_vars - ext_buffer = similar(โˆ‚ext_block_โˆ‚parameters_and_solved_vars, Float64) - ext_buffer.nzval .= 1 + derivatives_mat = โˆ‚SS_equations_โˆ‚parameters + buffer = similar(โˆ‚SS_equations_โˆ‚parameters, Float64) + buffer.nzval .= 0 end - ext_chol_buff = ext_buffer * ext_buffer' - - ext_chol_buff += โ„’.I - - prob = ๐’ฎ.LinearProblem(ext_chol_buff, ฯตแต‰, ๐’ฎ.CholeskyFactorization()) - - ext_chol_buffer = ๐’ฎ.init(prob, ๐’ฎ.CholeskyFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) - - prob = ๐’ฎ.LinearProblem(ext_buffer, ฯตแต‰, ๐’ฎ.LUFactorization()) - - ext_lu_buffer = ๐’ฎ.init(prob, ๐’ฎ.LUFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) - if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) else parallel = Symbolics.SerialForm() end - _, ext_func_exprs = Symbolics.build_function(derivatives_mat_ext, ๐”Š, ๐”“, + _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”“, ๐”˜, cse = cse, - skipzeros = skipzeros, - # nanmath = false, + skipzeros = skipzeros, + # nanmath = false, parallel = parallel, expression_module = @__MODULE__, expression = Val(false))::Tuple{<:Function, <:Function} - - push!(NSSS_solver_cache_init_tmp, [haskey(๐“‚.constants.post_parameters_macro.guess, v) ? ๐“‚.constants.post_parameters_macro.guess[v] : Inf for v in sorted_vars]) - push!(NSSS_solver_cache_init_tmp, [Inf]) - - # WARNING: infinite bounds are transformed to 1e12 - lbs = Float64[] - ubs = Float64[] - - limit_boundaries = 1e12 - - for i in vcat(sorted_vars, calib_pars_input, other_vars_input) - if haskey(๐“‚.constants.post_parameters_macro.bounds,i) - push!(lbs,๐“‚.constants.post_parameters_macro.bounds[i][1]) - push!(ubs,๐“‚.constants.post_parameters_macro.bounds[i][2]) - else - push!(lbs,-limit_boundaries) - push!(ubs, limit_boundaries) - end - end - - push!(SS_solve_func,ss_and_aux_equations...) - - push!(SS_solve_func,:(params_and_solved_vars = [$(calib_pars_input...), $(other_vars_input...)])) - - push!(SS_solve_func,:(lbs = [$(lbs...)])) - push!(SS_solve_func,:(ubs = [$(ubs...)])) - - # n_block = length(๐“‚.ss_solve_blocks) + 1 - n_block = length(๐“‚.NSSS.solve_blocks_in_place) + 1 - - push!(SS_solve_func,:(inits = [max.(lbs[1:length(closest_solution[$(2*(n_block-1)+1)])], min.(ubs[1:length(closest_solution[$(2*(n_block-1)+1)])], closest_solution[$(2*(n_block-1)+1)])), closest_solution[$(2*n_block)]])) - - push!(SS_solve_func,:(solution = block_solver(params_and_solved_vars, - $(n_block), - ๐“‚.NSSS.solve_blocks_in_place[$(n_block)], - # ๐“‚.ss_solve_blocks[$(n_block)], - # ๐“‚.ss_solve_blocks_no_transform[$(n_block)], - # f, - inits, - lbs, - ubs, - solver_parameters, - fail_fast_solvers_only, - cold_start, - verbose))) - - push!(SS_solve_func,:(iters += solution[2][2])) - push!(SS_solve_func,:(solution_error += solution[2][1])) - push!(SS_solve_func, :(if solution_error > tol.NSSS_acceptance_tol if verbose println("Failed after solving block with error $solution_error") end; scale = scale * .3 + solved_scale * .7; continue end)) - - if length(ss_and_aux_equations_error) > 0 - push!(SS_solve_func,:(solution_error += $(Expr(:call, :+, ss_and_aux_equations_error...)))) - push!(SS_solve_func, :(if solution_error > tol.NSSS_acceptance_tol if verbose println("Failed for aux variables with error $(solution_error)") end; scale = scale * .3 + solved_scale * .7; continue end)) - end - - push!(SS_solve_func,:(sol = solution[1])) - - push!(SS_solve_func,:($(result...))) - - push!(SS_solve_func,:(NSSS_solver_cache_tmp = [NSSS_solver_cache_tmp..., typeof(sol) == Vector{Float64} ? sol : โ„ฑ.value.(sol)])) - push!(SS_solve_func,:(NSSS_solver_cache_tmp = [NSSS_solver_cache_tmp..., typeof(params_and_solved_vars) == Vector{Float64} ? params_and_solved_vars : โ„ฑ.value.(params_and_solved_vars)])) - - # Create nonlinear solver workspaces for regular and extended problems - workspace = Nonlinear_solver_workspace(ฯต, buffer, chol_buffer, lu_buffer) - ext_workspace = Nonlinear_solver_workspace(ฯตแต‰, ext_buffer, ext_chol_buffer, ext_lu_buffer) - - push!(๐“‚.NSSS.solve_blocks_in_place, ss_solve_block( - function_and_jacobian(calc_block!::Function, func_exprs::Function, workspace), - function_and_jacobian(calc_ext_block!::Function, ext_func_exprs::Function, ext_workspace) - ) - ) - - return nothing -end - - - - -function partial_solve(eqs_to_solve::Vector{E}, vars_to_solve::Vector{T}, incidence_matrix_subset; avoid_solve::Bool = false)::Tuple{Vector{T}, Vector{T}, Vector{E}, Vector{T}} where {E, T} - for n in length(eqs_to_solve)-1:-1:2 - for eq_combo in combinations(1:length(eqs_to_solve), n) - var_indices_to_select_from = findall([sum(incidence_matrix_subset[:,eq_combo],dims = 2)...] .> 0) - - var_indices_in_remaining_eqs = findall([sum(incidence_matrix_subset[:,setdiff(1:length(eqs_to_solve),eq_combo)],dims = 2)...] .> 0) - - for var_combo in combinations(var_indices_to_select_from, n) - remaining_vars_in_remaining_eqs = setdiff(var_indices_in_remaining_eqs, var_combo) - # println("Solving for: ",vars_to_solve[var_combo]," in: ",eqs_to_solve[eq_combo]) - if length(remaining_vars_in_remaining_eqs) == length(eqs_to_solve) - n # not sure whether this condition needs to be there. could be because if the last remaining vars not solved for in the block is not present in the remaining block he will not be able to solve it for the same reasons he wasn't able to solve the unpartitioned block - if avoid_solve || count_ops(Meta.parse(string(eqs_to_solve[eq_combo]))) > 15 - soll = nothing - else - soll = solve_symbolically(eqs_to_solve[eq_combo], vars_to_solve[var_combo]) - end - - if !(isnothing(soll) || isempty(soll)) - soll_collected = collect(values(soll)) - - return (vars_to_solve[setdiff(1:length(eqs_to_solve),var_combo)], - vars_to_solve[var_combo], - eqs_to_solve[setdiff(1:length(eqs_to_solve),eq_combo)], - soll_collected) - end - end - end - end - end - - return (T[], T[], E[], T[]) -end - - - -function make_equation_robust_to_domain_errors(eqs,#::Vector{Union{Symbol,Expr}}, - vars_to_exclude::Vector{Vector{Symbol}}, - bounds::Dict{Symbol,Tuple{Float64,Float64}}, - โž•_vars::Vector{Symbol}, - unique_โž•_eqs,#::Dict{Union{Expr,Symbol},Symbol}(); - precompile::Bool = false) - ss_and_aux_equations = Expr[] - ss_and_aux_equations_dep = Expr[] - ss_and_aux_equations_error = Expr[] - ss_and_aux_equations_error_dep = Expr[] - rewritten_eqs = Union{Expr,Symbol}[] - # write down ss equations including nonnegativity auxiliary variables - # find nonegative variables, parameters, or terms - for eq in eqs - if eq isa Symbol - push!(rewritten_eqs, eq) - elseif eq isa Expr - rewritten_eq = postwalk(x -> - x isa Expr ? - # x.head == :(=) ? - # Expr(:call,:(-),x.args[1],x.args[2]) : #convert = to - - # x.head == :ref ? - # occursin(r"^(x|ex|exo|exogenous){1}"i,string(x.args[2])) ? 0 : # set shocks to zero and remove time scripts - # x : - x.head == :call ? - x.args[1] == :* ? - x.args[2] isa Int ? - x.args[3] isa Int ? - x : - Expr(:call, :*, x.args[3:end]..., x.args[2]) : # 2beta => beta * 2 - x : - x.args[1] โˆˆ [:^] ? - !(x.args[3] isa Int) ? - x.args[2] isa Symbol ? # nonnegative parameters - x.args[2] โˆˆ vars_to_exclude[1] ? - begin - bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 1e12)) : (eps(), 1e12) - x - end : - begin - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if x.args[2] in vars_to_exclude[1] - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - - :($(replacement) ^ $(x.args[3])) - end : - x.args[2] isa Float64 ? - x : - x.args[2].head == :call ? # nonnegative expressions - begin - if precompile - replacement = x.args[2] - else - replacement = simplify(x.args[2]) - end - - if !(replacement isa Int) # check if the nonnegative term is just a constant - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - end - - :($(replacement) ^ $(x.args[3])) - end : - x : - x : - x.args[2] isa Float64 ? - x : - x.args[1] โˆˆ [:log] ? - x.args[2] isa Symbol ? # nonnegative parameters - x.args[2] โˆˆ vars_to_exclude[1] ? - begin - bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 1e12)) : (eps(), 1e12) - x - end : - begin - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if x.args[2] in vars_to_exclude[1] - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x.args[2].head == :call ? # nonnegative expressions - begin - if precompile - replacement = x.args[2] - else - replacement = simplify(x.args[2]) - end - - if !(replacement isa Int) # check if the nonnegative term is just a constant - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x : - x.args[1] โˆˆ [:norminvcdf, :norminv, :qnorm] ? - x.args[2] isa Symbol ? # nonnegative parameters - x.args[2] โˆˆ vars_to_exclude[1] ? - begin - bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 1-eps())) : (eps(), 1 - eps()) - x - end : - begin - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if x.args[2] in vars_to_exclude[1] - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end + ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters = buffer + ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters = func_exprs - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1 - eps())) : (eps(), 1 - eps()) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - unique_โž•_eqs[x.args[2]] = replacement - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x.args[2].head == :call ? # nonnegative expressions - begin - if precompile - replacement = x.args[2] - else - replacement = simplify(x.args[2]) - end - if !(replacement isa Int) # check if the nonnegative term is just a constant - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end + โˆ‚SS_equations_โˆ‚SS_and_pars = Symbolics.sparsejacobian(ss_equations_sub, ๐”˜) # nฯต x nx - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1 - eps())) : (eps(), 1 - eps()) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x : - x.args[1] โˆˆ [:exp] ? - x.args[2] isa Symbol ? # have exp terms bound so they dont go to Inf - x.args[2] โˆˆ vars_to_exclude[1] ? - begin - bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], -1e12), min(bounds[x.args[2]][2], 600)) : (-1e12, 600) - x - end : - begin - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if x.args[2] in vars_to_exclude[1] - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], -1e12), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 600)) : (-1e12, 600) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x.args[2].head == :call ? # have exp terms bound so they dont go to Inf - begin - if precompile - replacement = x.args[2] - else - replacement = simplify(x.args[2]) - end - - if !(replacement isa Int) # check if the nonnegative term is just a constant - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], -1e12), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 600)) : (-1e12, 600) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x : - x.args[1] โˆˆ [:erfcinv] ? - x.args[2] isa Symbol ? # nonnegative parameters - x.args[2] โˆˆ vars_to_exclude[1] ? - begin - bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 2 - eps())) : (eps(), 2 - eps()) - x - end : - begin - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if x.args[2] in vars_to_exclude[1] - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 2 - eps())) : (eps(), 2 - eps()) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x.args[2].head == :call ? # nonnegative expressions - begin - if precompile - replacement = x.args[2] - else - replacement = simplify(x.args[2]) - end - - if !(replacement isa Int) # check if the nonnegative term is just a constant - if haskey(unique_โž•_eqs, x.args[2]) - replacement = unique_โž•_eqs[x.args[2]] - else - if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) - push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - else - push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) - push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) - end - - bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 2 - eps())) : (eps(), 2 - eps()) - push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) - replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) - - unique_โž•_eqs[x.args[2]] = replacement - end - end - - :($(Expr(:call, x.args[1], replacement))) - end : - x : - x : - x : - x, - eq) - push!(rewritten_eqs,rewritten_eq) - else - @assert typeof(eq) in [Symbol, Expr] - end - end - - vars_to_exclude_from_block = vcat(vars_to_exclude...) - - found_new_dependecy = true - - while found_new_dependecy - found_new_dependecy = false - - for ssauxdep in ss_and_aux_equations_dep - push!(vars_to_exclude_from_block, ssauxdep.args[1]) - end - - for (iii, ssaux) in enumerate(ss_and_aux_equations) - if !isempty(intersect(get_symbols(ssaux), vars_to_exclude_from_block)) - found_new_dependecy = true - push!(vars_to_exclude_from_block, ssaux.args[1]) - push!(ss_and_aux_equations_dep, ssaux) - push!(ss_and_aux_equations_error_dep, ss_and_aux_equations_error[iii]) - deleteat!(ss_and_aux_equations, iii) - deleteat!(ss_and_aux_equations_error, iii) - end - end - end - - return rewritten_eqs, ss_and_aux_equations, ss_and_aux_equations_dep, ss_and_aux_equations_error, ss_and_aux_equations_error_dep -end - - - -function replace_symbols(exprs::T, remap::Dict{Symbol,S}) where {T,S} - postwalk(node -> - if node isa Symbol && haskey(remap, node) - remap[node] - else - node - end, - exprs) -end - -function write_ss_check_function!(๐“‚::โ„ณ; - cse = true, - skipzeros = true, - density_threshold::Float64 = .1, - nnz_parallel_threshold::Int = 1000000, - min_length::Int = 10000) - unknowns = union(setdiff(๐“‚.constants.post_model_macro.vars_in_ss_equations, ๐“‚.constants.post_model_macro.โž•_vars), ๐“‚.equations.calibration_parameters) - - ss_equations = vcat(๐“‚.equations.steady_state, ๐“‚.equations.calibration) - - - - np = length(๐“‚.constants.post_complete_parameters.parameters) - nu = length(unknowns) - # nc = length(๐“‚.calibration_equations_no_var) - - Symbolics.@variables ๐”“[1:np] ๐”˜[1:nu]# โ„ญ[1:nc] - - parameter_dict = Dict{Symbol, Symbol}() - back_to_array_dict = Dict{Symbolics.Num, Symbolics.Num}() - calib_vars = Symbol[] - calib_expr = [] - - - for (i,v) in enumerate(๐“‚.constants.post_complete_parameters.parameters) - push!(parameter_dict, v => :($(Symbol("๐”“_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”“_$i"))), @__MODULE__) => ๐”“[i]) - end - - for (i,v) in enumerate(unknowns) - push!(parameter_dict, v => :($(Symbol("๐”˜_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”˜_$i"))), @__MODULE__) => ๐”˜[i]) - end - - for (i,v) in enumerate(๐“‚.equations.calibration_no_var) - push!(calib_vars, v.args[1]) - push!(calib_expr, v.args[2]) - # push!(parameter_dict, v.args[1] => :($(Symbol("โ„ญ_$i")))) - # push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("โ„ญ_$i"))), @__MODULE__) => โ„ญ[i]) - end - - calib_replacements = Dict{Symbol,Any}() - for (i,x) in enumerate(calib_vars) - replacement = Dict(x => calib_expr[i]) - for ii in i+1:length(calib_vars) - calib_expr[ii] = replace_symbols(calib_expr[ii], replacement) - end - push!(calib_replacements, x => calib_expr[i]) - end - - - ss_equations_sub = ss_equations |> - x -> replace_symbols.(x, Ref(calib_replacements)) |> - x -> replace_symbols.(x, Ref(parameter_dict)) |> - x -> Symbolics.parse_expr_to_symbolic.(x, Ref(@__MODULE__)) |> - x -> Symbolics.substitute.(x, Ref(back_to_array_dict)) - - - lennz = length(ss_equations_sub) - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_exprs = Symbolics.build_function(ss_equations_sub, ๐”“, ๐”˜, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - - ๐“‚.functions.NSSS_check = func_exprs - - - # SS_and_pars = Symbol.(vcat(string.(sort(collect(setdiff(reduce(union,get_symbols.(๐“‚.ss_aux_equations)),union(๐“‚.constants.post_model_macro.parameters_in_equations,๐“‚.constants.post_model_macro.โž•_vars))))), ๐“‚.calibration_equations_parameters)) - - # eqs = vcat(๐“‚.ss_equations, ๐“‚.calibration_equations) - - # nx = length(๐“‚.parameter_values) - - # np = length(SS_and_pars) - - nฯตหข = length(ss_equations) - - # nc = length(๐“‚.calibration_equations_no_var) - - # Symbolics.@variables ๐”›ยน[1:nx] ๐”“ยน[1:np] - - # ฯตหข = zeros(Symbolics.Num, nฯตหข) - - # calib_vals = zeros(Symbolics.Num, nc) - - # ๐“‚.SS_calib_func(calib_vals, ๐”“) - - # ๐“‚.functions.NSSS_check(ฯตหข, ๐”“, ๐”˜, calib_vals) - - โˆ‚SS_equations_โˆ‚parameters = Symbolics.sparsejacobian(ss_equations_sub, ๐”“) # nฯต x nx - - lennz = nnz(โˆ‚SS_equations_โˆ‚parameters) - - if (lennz / length(โˆ‚SS_equations_โˆ‚parameters) > density_threshold) || (length(โˆ‚SS_equations_โˆ‚parameters) < min_length) - derivatives_mat = convert(Matrix, โˆ‚SS_equations_โˆ‚parameters) - buffer = zeros(Float64, size(โˆ‚SS_equations_โˆ‚parameters)) - else - derivatives_mat = โˆ‚SS_equations_โˆ‚parameters - buffer = similar(โˆ‚SS_equations_โˆ‚parameters, Float64) - buffer.nzval .= 0 - end - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”“, ๐”˜, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - ๐“‚.caches.โˆ‚equations_โˆ‚parameters = buffer - ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters = func_exprs - - - - โˆ‚SS_equations_โˆ‚SS_and_pars = Symbolics.sparsejacobian(ss_equations_sub, ๐”˜) # nฯต x nx - - lennz = nnz(โˆ‚SS_equations_โˆ‚SS_and_pars) + lennz = nnz(โˆ‚SS_equations_โˆ‚SS_and_pars) if (lennz / length(โˆ‚SS_equations_โˆ‚SS_and_pars) > density_threshold) || (length(โˆ‚SS_equations_โˆ‚SS_and_pars) < min_length) derivatives_mat = convert(Matrix, โˆ‚SS_equations_โˆ‚SS_and_pars) @@ -4768,1142 +4782,141 @@ function write_ss_check_function!(๐“‚::โ„ณ; derivatives_mat = โˆ‚SS_equations_โˆ‚SS_and_pars buffer = similar(โˆ‚SS_equations_โˆ‚SS_and_pars, Float64) buffer.nzval .= 0 - end - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”“, ๐”˜, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - ๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars = buffer - ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars = func_exprs - - return nothing -end - - -function write_steady_state_solver_function!(๐“‚::โ„ณ, symbolic_SS, Symbolics::symbolics; verbose::Bool = false, avoid_solve::Bool = false) - unknowns = union(Symbolics.calibration_equations_parameters, Symbolics.vars_in_ss_equations) - - @assert length(unknowns) <= length(Symbolics.ss_equations) + length(Symbolics.calibration_equations) "Unable to solve steady state. More unknowns than equations." - - incidence_matrix = spzeros(Int,length(unknowns),length(unknowns)) - - eq_list = vcat(union.(setdiff.(union.(Symbolics.var_list_aux_SS, - Symbolics.ss_list_aux_SS), - Symbolics.var_redundant_list), - Symbolics.par_list_aux_SS), - union.(Symbolics.ss_calib_list, - Symbolics.par_calib_list)) - - for (i,u) in enumerate(unknowns) - for (k,e) in enumerate(eq_list) - incidence_matrix[i,k] = u โˆˆ e - end - end - - Q, P, R, nmatch, n_blocks = BlockTriangularForm.order(incidence_matrix) - Rฬ‚ = Int[] - for i in 1:n_blocks - [push!(Rฬ‚, n_blocks - i + 1) for ii in R[i]:R[i+1] - 1] - end - push!(Rฬ‚,1) - - vars = hcat(P, Rฬ‚)' - eqs = hcat(Q, Rฬ‚)' - - # @assert all(eqs[1,:] .> 0) "Could not solve system of steady state and calibration equations for: " * repr([collect(Symbol.(unknowns))[vars[1,eqs[1,:] .< 0]]...]) # repr([vcat(Symbolics.ss_equations,Symbolics.calibration_equations)[-eqs[1,eqs[1,:].<0]]...]) - @assert all(eqs[1,:] .> 0) "Could not solve system of steady state and calibration equations. Number of redundant equations: " * repr(sum(eqs[1,:] .< 0)) * ". Try defining some steady state values as parameters (e.g. r[ss] -> rฬ„). Nonstationary variables are not supported as of now." # repr([vcat(Symbolics.ss_equations,Symbolics.calibration_equations)[-eqs[1,eqs[1,:].<0]]...]) - - n = n_blocks - - ss_equations = vcat(Symbolics.ss_equations,Symbolics.calibration_equations)# .|> SPyPyC.Sym - # println(ss_equations) - - SS_solve_func = [] - - atoms_in_equations = Set{Symbol}() - atoms_in_equations_list = [] - relevant_pars_across = Symbol[] - NSSS_solver_cache_init_tmp = [] - - solved_vars = [] - solved_vals = [] - - min_max_errors = [] - - unique_โž•_eqs = Dict{Union{Expr,Symbol},Symbol}() - - while n > 0 - if length(eqs[:,eqs[2,:] .== n]) == 2 - var_to_solve_for = unknowns[vars[:,vars[2,:] .== n][1]] - - eq_to_solve = ss_equations[eqs[:,eqs[2,:] .== n][1]] - - # eliminate min/max from equations if solving for variables inside min/max. set to the variable we solve for automatically - parsed_eq_to_solve_for = eq_to_solve |> string |> Meta.parse - - minmax_fixed_eqs = postwalk(x -> - x isa Expr ? - x.head == :call ? - x.args[1] โˆˆ [:Max,:Min] ? - Symbol(var_to_solve_for) โˆˆ get_symbols(x.args[2]) ? - x.args[2] : - Symbol(var_to_solve_for) โˆˆ get_symbols(x.args[3]) ? - x.args[3] : - x : - x : - x : - x, - parsed_eq_to_solve_for) - - if parsed_eq_to_solve_for != minmax_fixed_eqs - [push!(atoms_in_equations, a) for a in setdiff(get_symbols(parsed_eq_to_solve_for), get_symbols(minmax_fixed_eqs))] - push!(min_max_errors,:(solution_error += abs($parsed_eq_to_solve_for))) - push!(SS_solve_func, :(if solution_error > tol.NSSS_acceptance_tol if verbose println("Failed for min max terms in equations with error $solution_error") end; scale = scale * .3 + solved_scale * .7; continue end)) - eq_to_solve = eval(minmax_fixed_eqs) - end - - if avoid_solve || count_ops(Meta.parse(string(eq_to_solve))) > 15 - soll = nothing - else - soll = solve_symbolically(eq_to_solve,var_to_solve_for) - end - - if isnothing(soll) || isempty(soll) - println("Failed finding solution symbolically for: ",var_to_solve_for," in: ",eq_to_solve) - - eq_idx_in_block_to_solve = eqs[:,eqs[2,:] .== n][1,:] - - write_block_solution!(๐“‚, SS_solve_func, [var_to_solve_for], [eq_to_solve], relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, solved_vars, solved_vals) - # write_domain_safe_block_solution!(๐“‚, SS_solve_func, [var_to_solve_for], [eq_to_solve], relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, unique_โž•_eqs) - elseif soll[1].is_number == true - ss_equations = [replace_symbolic(eq, var_to_solve_for, soll[1]) for eq in ss_equations] - - push!(solved_vars,Symbol(var_to_solve_for)) - push!(solved_vals,Meta.parse(string(soll[1]))) - - if (solved_vars[end] โˆˆ ๐“‚.constants.post_model_macro.โž•_vars) - push!(SS_solve_func,:($(solved_vars[end]) = max(eps(),$(solved_vals[end])))) - else - push!(SS_solve_func,:($(solved_vars[end]) = $(solved_vals[end]))) - end - - push!(atoms_in_equations_list,[]) - else - push!(solved_vars,Symbol(var_to_solve_for)) - push!(solved_vals,Meta.parse(string(soll[1]))) - - [push!(atoms_in_equations, Symbol(a)) for a in soll[1].atoms()] - push!(atoms_in_equations_list, Set(union(setdiff(get_symbols(parsed_eq_to_solve_for), get_symbols(minmax_fixed_eqs)),Symbol.(soll[1].atoms())))) - - if (solved_vars[end] โˆˆ ๐“‚.constants.post_model_macro.โž•_vars) - push!(SS_solve_func,:($(solved_vars[end]) = begin - _bounds = get($(๐“‚.constants.post_parameters_macro.bounds), $(QuoteNode(solved_vars[end])), (eps(), 1e12)) - min(max(_bounds[1], $(solved_vals[end])), _bounds[2]) - end)) - push!(SS_solve_func,:(solution_error += $(Expr(:call,:abs, Expr(:call, :-, solved_vars[end], solved_vals[end]))))) - push!(SS_solve_func, :(if solution_error > tol.NSSS_acceptance_tol if verbose println("Failed for analytical aux variables with error $solution_error") end; scale = scale * .3 + solved_scale * .7; continue end)) - - unique_โž•_eqs[solved_vals[end]] = solved_vars[end] - else - vars_to_exclude = [vcat(Symbol.(var_to_solve_for), ๐“‚.constants.post_model_macro.โž•_vars), Symbol[]] - - rewritten_eqs, ss_and_aux_equations, ss_and_aux_equations_dep, ss_and_aux_equations_error, ss_and_aux_equations_error_dep = make_equation_robust_to_domain_errors([solved_vals[end]], vars_to_exclude, ๐“‚.constants.post_parameters_macro.bounds, ๐“‚.constants.post_model_macro.โž•_vars, unique_โž•_eqs) - - if length(vcat(ss_and_aux_equations_error, ss_and_aux_equations_error_dep)) > 0 - push!(SS_solve_func,vcat(ss_and_aux_equations, ss_and_aux_equations_dep)...) - push!(SS_solve_func,:(solution_error += $(Expr(:call, :+, vcat(ss_and_aux_equations_error, ss_and_aux_equations_error_dep)...)))) - push!(SS_solve_func, :(if solution_error > tol.NSSS_acceptance_tol if verbose println("Failed for analytical variables with error $solution_error") end; scale = scale * .3 + solved_scale * .7; continue end)) - end - - push!(SS_solve_func,:($(solved_vars[end]) = $(rewritten_eqs[1]))) - end - - if haskey(๐“‚.constants.post_parameters_macro.bounds, solved_vars[end]) && solved_vars[end] โˆ‰ ๐“‚.constants.post_model_macro.โž•_vars - push!(SS_solve_func,:(solution_error += abs(min(max($(๐“‚.constants.post_parameters_macro.bounds[solved_vars[end]][1]), $(solved_vars[end])), $(๐“‚.constants.post_parameters_macro.bounds[solved_vars[end]][2])) - $(solved_vars[end])))) - push!(SS_solve_func, :(if solution_error > tol.NSSS_acceptance_tol if verbose println("Failed for bounded variables with error $solution_error") end; scale = scale * .3 + solved_scale * .7; continue end)) - end - end - else - vars_to_solve = unknowns[vars[:,vars[2,:] .== n][1,:]] - - eqs_to_solve = ss_equations[eqs[:,eqs[2,:] .== n][1,:]] - - numerical_sol = false - - if symbolic_SS - if avoid_solve || count_ops(Meta.parse(string(eqs_to_solve))) > 15 - soll = nothing - else - soll = solve_symbolically(eqs_to_solve,vars_to_solve) - end - - if isnothing(soll) || isempty(soll) || length(intersect((union(SPyPyC.free_symbols.(collect(values(soll)))...) .|> SPyPyC.:โ†“),(vars_to_solve .|> SPyPyC.:โ†“))) > 0 - if verbose println("Failed finding solution symbolically for: ",vars_to_solve," in: ",eqs_to_solve,". Solving numerically.") end - - numerical_sol = true - else - if verbose println("Solved: ",string.(eqs_to_solve)," for: ",Symbol.(vars_to_solve), " symbolically.") end - - atoms = reduce(union,map(x->x.atoms(),collect(values(soll)))) - - for a in atoms push!(atoms_in_equations, Symbol(a)) end - - for vars in vars_to_solve - push!(solved_vars,Symbol(vars)) - push!(solved_vals,Meta.parse(string(soll[vars]))) #using convert(Expr,x) leads to ugly expressions - - push!(atoms_in_equations_list, Set(Symbol.(soll[vars].atoms()))) - push!(SS_solve_func,:($(solved_vars[end]) = $(solved_vals[end]))) - end - end - end - - eq_idx_in_block_to_solve = eqs[:,eqs[2,:] .== n][1,:] - - incidence_matrix_subset = incidence_matrix[vars[:,vars[2,:] .== n][1,:], eq_idx_in_block_to_solve] - - # try symbolically and use numerical if it does not work - if numerical_sol || !symbolic_SS - pv = sortperm(vars_to_solve, by = Symbol) - pe = sortperm(eqs_to_solve, by = string) - - if length(pe) > 5 - write_block_solution!(๐“‚, SS_solve_func, vars_to_solve, eqs_to_solve, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, solved_vars, solved_vals) - # write_domain_safe_block_solution!(๐“‚, SS_solve_func, vars_to_solve, eqs_to_solve, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, unique_โž•_eqs) - else - solved_system = partial_solve(eqs_to_solve[pe], vars_to_solve[pv], incidence_matrix_subset[pv,pe], avoid_solve = avoid_solve) - - # if !isnothing(solved_system) && !any(contains.(string.(vcat(solved_system[3],solved_system[4])), "LambertW")) && !any(contains.(string.(vcat(solved_system[3],solved_system[4])), "Heaviside")) - # write_reduced_block_solution!(๐“‚, SS_solve_func, solved_system, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, - # ๐“‚.constants.post_model_macro.โž•_vars, unique_โž•_eqs) - # else - write_block_solution!(๐“‚, SS_solve_func, vars_to_solve, eqs_to_solve, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, solved_vars, solved_vals) - # write_domain_safe_block_solution!(๐“‚, SS_solve_func, vars_to_solve, eqs_to_solve, relevant_pars_across, NSSS_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, unique_โž•_eqs) - # end - end - - if !symbolic_SS && verbose - println("Solved: ",string.(eqs_to_solve)," for: ",Symbol.(vars_to_solve), " numerically.") - end - end - end - n -= 1 - end - - push!(NSSS_solver_cache_init_tmp, fill(Inf, length(๐“‚.constants.post_complete_parameters.parameters))) - push!(๐“‚.caches.solver_cache, NSSS_solver_cache_init_tmp) - - unknwns = Symbol.(unknowns) - - parameters_only_in_par_defs = Set() - # add parameters from parameter definitions - if length(๐“‚.equations.calibration_no_var) > 0 - atoms = reduce(union, get_symbols.(๐“‚.equations.calibration_no_var)) - [push!(atoms_in_equations, a) for a in atoms] - [push!(parameters_only_in_par_defs, a) for a in atoms] - end - - # ๐“‚.par = union(๐“‚.par,setdiff(parameters_only_in_par_defs,๐“‚.parameters_as_function_of_parameters)) - - parameters_in_equations = [] - - for (i, parss) in enumerate(๐“‚.constants.post_complete_parameters.parameters) - if parss โˆˆ union(atoms_in_equations, relevant_pars_across) - push!(parameters_in_equations, :($parss = parameters[$i])) - end - end - - dependencies = [] - for (i, a) in enumerate(atoms_in_equations_list) - push!(dependencies, solved_vars[i] => intersect(a, union(๐“‚.constants.post_model_macro.var, ๐“‚.constants.post_complete_parameters.parameters))) - end - - push!(dependencies, :SS_relevant_calibration_parameters => intersect(reduce(union, atoms_in_equations_list), ๐“‚.constants.post_complete_parameters.parameters)) - - ๐“‚.NSSS.dependencies = dependencies - - - - dyn_exos = [] - for dex in union(๐“‚.constants.post_model_macro.exo_past, ๐“‚.constants.post_model_macro.exo_future) - push!(dyn_exos,:($dex = 0)) - end - - push!(SS_solve_func,:($(dyn_exos...))) - - push!(SS_solve_func, min_max_errors...) - # push!(SS_solve_func,:(push!(NSSS_solver_cache_tmp, params_scaled_flt))) - - push!(SS_solve_func,:(if length(NSSS_solver_cache_tmp) == 0 NSSS_solver_cache_tmp = [copy(params_flt)] else NSSS_solver_cache_tmp = [NSSS_solver_cache_tmp..., copy(params_flt)] end)) - - - # push!(SS_solve_func,:(for pars in ๐“‚.caches.solver_cache - # latest = sqrt(sum(abs2,pars[end] - params_flt))# / max(sum(abs2,pars[end]), sum(abs,params_flt)) - # if latest <= current_best - # current_best = latest - # end - # end)) - push!(SS_solve_func,:(if (current_best > 1e-8) && (solution_error < tol.NSSS_acceptance_tol) && (scale == 1) - reverse_diff_friendly_push!(๐“‚.caches.solver_cache, NSSS_solver_cache_tmp) - end)) - # push!(SS_solve_func,:(if length(๐“‚.caches.solver_cache) > 100 popfirst!(๐“‚.caches.solver_cache) end)) - - # push!(SS_solve_func,:(SS_init_guess = ([$(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))...), $(๐“‚.calibration_equations_parameters...)]))) - - # push!(SS_solve_func,:(๐“‚.SS_init_guess = typeof(SS_init_guess) == Vector{Float64} ? SS_init_guess : โ„ฑ.value.(SS_init_guess))) - - # push!(SS_solve_func,:(return ComponentVector([$(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))...), $(๐“‚.calibration_equations_parameters...)], Axis([sort(union(๐“‚.constants.post_model_macro.exo_present,๐“‚.constants.post_model_macro.var))...,๐“‚.calibration_equations_parameters...])))) - - - # fix parameter bounds - par_bounds = [] - - for varpar in intersect(๐“‚.constants.post_complete_parameters.parameters,union(atoms_in_equations, relevant_pars_across)) - if haskey(๐“‚.constants.post_parameters_macro.bounds, varpar) - push!(par_bounds, :($varpar = min(max($varpar,$(๐“‚.constants.post_parameters_macro.bounds[varpar][1])),$(๐“‚.constants.post_parameters_macro.bounds[varpar][2])))) - end - end - - solve_exp = :(function solve_SS(initial_parameters::Vector{Real}, - ๐“‚::โ„ณ, - # fail_fast_solvers_only::Bool, - tol::Tolerances, - verbose::Bool, - cold_start::Bool, - solver_parameters::Vector{solver_parameters}) - initial_parameters = typeof(initial_parameters) == Vector{Float64} ? initial_parameters : โ„ฑ.value.(initial_parameters) - - initial_parameters_tmp = copy(initial_parameters) - - parameters = copy(initial_parameters) - params_flt = copy(initial_parameters) - - current_best = sum(abs2,๐“‚.caches.solver_cache[end][end] - initial_parameters) - closest_solution_init = ๐“‚.caches.solver_cache[end] - - for pars in ๐“‚.caches.solver_cache - copy!(initial_parameters_tmp, pars[end]) - - โ„’.axpy!(-1,initial_parameters,initial_parameters_tmp) - - latest = sum(abs2,initial_parameters_tmp) - if latest <= current_best - current_best = latest - closest_solution_init = pars - end - end - - # closest_solution = copy(closest_solution_init) - # solution_error = 1.0 - # iters = 0 - range_iters = 0 - solution_error = 1.0 - solved_scale = 0 - # range_length = [ 1, 2, 4, 8,16,32,64,128,1024] - scale = 1.0 - - NSSS_solver_cache_scale = CircularBuffer{Vector{Vector{Float64}}}(500) - push!(NSSS_solver_cache_scale, closest_solution_init) - # fail_fast_solvers_only = true - while range_iters <= (cold_start ? 1 : 500) && !(solution_error < tol.NSSS_acceptance_tol && solved_scale == 1) - range_iters += 1 - fail_fast_solvers_only = range_iters > 1 ? true : false - - if abs(solved_scale - scale) < 1e-2 - # println(NSSS_solver_cache_scale[end]) - break - end - - # println("i: $range_iters - scale: $scale - solved_scale: $solved_scale") - # println(closest_solution[end]) - # for range_ in range_length - # rangee = range(0,1,range_+1) - # for scale in rangee[2:end] - # scale = 6*scale^5 - 15*scale^4 + 10*scale^3 # smootherstep - - # if scale <= solved_scale continue end - - - current_best = sum(abs2,NSSS_solver_cache_scale[end][end] - initial_parameters) - closest_solution = NSSS_solver_cache_scale[end] - - for pars in NSSS_solver_cache_scale - copy!(initial_parameters_tmp, pars[end]) - - โ„’.axpy!(-1,initial_parameters,initial_parameters_tmp) - - latest = sum(abs2,initial_parameters_tmp) - - if latest <= current_best - current_best = latest - closest_solution = pars - end - end - - # println(closest_solution) - - if all(isfinite,closest_solution[end]) && initial_parameters != closest_solution_init[end] - parameters = scale * initial_parameters + (1 - scale) * closest_solution_init[end] - else - parameters = copy(initial_parameters) - end - params_flt = parameters - - # println(parameters) - - $(parameters_in_equations...) - $(par_bounds...) - $(๐“‚.equations.calibration_no_var...) - NSSS_solver_cache_tmp = [] - solution_error = 0.0 - iters = 0 - $(SS_solve_func...) - - if solution_error < tol.NSSS_acceptance_tol - # println("solved for $scale; $range_iters") - solved_scale = scale - if scale == 1 - # return ComponentVector([$(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))...), $(๐“‚.calibration_equations_parameters...)], Axis([sort(union(๐“‚.constants.post_model_macro.exo_present,๐“‚.constants.post_model_macro.var))...,๐“‚.calibration_equations_parameters...])), solution_error - # NSSS_solution = [$(Symbol.(replace.(string.(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))...), $(๐“‚.calibration_equations_parameters...)] - # NSSS_solution[abs.(NSSS_solution) .< 1e-12] .= 0 # doesn't work with Zygote - return [$(Symbol.(replace.(string.(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))...), $(๐“‚.equations.calibration_parameters...)], (solution_error, iters) - else - reverse_diff_friendly_push!(NSSS_solver_cache_scale, NSSS_solver_cache_tmp) - end - - if scale > .95 - scale = 1 - else - # scale = (scale + 1) / 2 - scale = scale * .4 + .6 - end - # else - # println("no sol") - # scale = (scale + solved_scale) / 2 - # println("scale $scale") - # elseif scale == 1 && range_ == range_length[end] - # return [$(Symbol.(replace.(string.(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))...), $(๐“‚.calibration_equations_parameters...)], (solution_error, iters) - end - # end - end - return zeros($(length(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future)) + length(๐“‚.equations.calibration_parameters))), (1, 0) - end) - - - ๐“‚.functions.NSSS_solve = @RuntimeGeneratedFunction(solve_exp) - # ๐“‚.functions.NSSS_solve = eval(solve_exp) - - return nothing -end - - - - -function solve_steady_state!(๐“‚::โ„ณ, - opts::CalculationOptions, - ss_solver_parameters_algorithm::Symbol, - ss_solver_parameters_maxtime::Real; - silent::Bool = false)::Tuple{Vector{Float64}, Float64, Bool} - """ - Internal function to solve and constants the steady state. - Returns: (SS_and_pars, solution_error, found_solution) - """ - start_time = time() - - if ๐“‚.constants.post_parameters_macro.precompile - return Float64[], 0.0, false - end - - if !(๐“‚.functions.NSSS_custom isa Function) - if !silent - print("Find non-stochastic steady state:\t\t\t\t\t") - end - end - - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts, cold_start = true) - - found_solution = true - - if !(๐“‚.functions.NSSS_custom isa Function) - select_fastest_SS_solver_parameters!(๐“‚, tol = opts.tol) - - if solution_error > opts.tol.NSSS_acceptance_tol - found_solution = find_SS_solver_parameters!(Val(ss_solver_parameters_algorithm), ๐“‚, tol = opts.tol, verbosity = 0, maxtime = ss_solver_parameters_maxtime, maxiter = 1000000000) - - if found_solution - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts, cold_start = true) - end - end - end - - if !(๐“‚.functions.NSSS_custom isa Function) - if !silent - println(round(time() - start_time, digits = 3), " seconds") - end - end - - if !found_solution - @warn "Could not find non-stochastic steady state. Consider setting bounds on variables or calibrated parameters in the `@parameters` section (e.g. `k > 10`)." - end - - ๐“‚.caches.non_stochastic_steady_state = SS_and_pars - ๐“‚.caches.outdated.non_stochastic_steady_state = !found_solution - - return SS_and_pars, solution_error, found_solution -end - -# Centralised helper to write symbolic derivatives and map functions -function write_symbolic_derivatives!(๐“‚::โ„ณ; perturbation_order::Int = 1, silent::Bool = false) - start_time = time() - - if !silent - if perturbation_order == 1 - print("Take symbolic derivatives up to first order:\t\t\t\t") - elseif perturbation_order == 2 - print("Take symbolic derivatives up to second order:\t\t\t\t") - elseif perturbation_order == 3 - print("Take symbolic derivatives up to third order:\t\t\t\t") - end - end - - write_auxiliary_indices!(๐“‚) - - write_functions_mapping!(๐“‚, perturbation_order) - - # Mark all solutions as outdated when derivative functions are rewritten - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true - - if !silent - println(round(time() - start_time, digits = 3), " seconds") - end - - return nothing -end - - -function write_steady_state_solver_function!(๐“‚::โ„ณ; - cse = true, - skipzeros = true, - density_threshold::Float64 = .1, - nnz_parallel_threshold::Int = 1000000, - min_length::Int = 1000, - verbose::Bool = false) - unknowns = union(๐“‚.constants.post_model_macro.vars_in_ss_equations, ๐“‚.equations.calibration_parameters) - - @assert length(unknowns) <= length(๐“‚.equations.steady_state_aux) + length(๐“‚.equations.calibration) "Unable to solve steady state. More unknowns than equations." - - incidence_matrix = spzeros(Int,length(unknowns),length(unknowns)) - - eq_list = vcat(union.(union.(๐“‚.constants.post_model_macro.var_list_aux_SS, - ๐“‚.constants.post_model_macro.ss_list_aux_SS), - ๐“‚.constants.post_model_macro.par_list_aux_SS), - union.(๐“‚.constants.post_parameters_macro.ss_calib_list, - ๐“‚.constants.post_parameters_macro.par_calib_list)) - - for (i,u) in enumerate(unknowns) - for (k,e) in enumerate(eq_list) - incidence_matrix[i,k] = u โˆˆ e - end - end - - Q, P, R, nmatch, n_blocks = BlockTriangularForm.order(incidence_matrix) - Rฬ‚ = Int[] - for i in 1:n_blocks - [push!(Rฬ‚, n_blocks - i + 1) for ii in R[i]:R[i+1] - 1] - end - push!(Rฬ‚,1) - - vars = hcat(P, Rฬ‚)' - eqs = hcat(Q, Rฬ‚)' - # @assert all(eqs[1,:] .> 0) "Could not solve system of steady state and calibration equations for: " * repr([collect(Symbol.(unknowns))[vars[1,eqs[1,:] .< 0]]...]) # repr([vcat(๐“‚.ss_equations,๐“‚.calibration_equations)[-eqs[1,eqs[1,:].<0]]...]) - @assert all(eqs[1,:] .> 0) "Could not solve system of steady state and calibration equations. Number of redundant equations: " * repr(sum(eqs[1,:] .< 0)) * ". Try defining some steady state values as parameters (e.g. r[ss] -> rฬ„). Nonstationary variables are not supported as of now." # repr([vcat(๐“‚.ss_equations,๐“‚.calibration_equations)[-eqs[1,eqs[1,:].<0]]...]) - - n = n_blocks - - ss_equations = vcat(๐“‚.equations.steady_state_aux,๐“‚.equations.calibration) - - SS_solve_func = [] - - atoms_in_equations = Set{Symbol}() - atoms_in_equations_list = [] - relevant_pars_across = [] - NSSS_solver_cache_init_tmp = [] - - solved_vars = [] - solved_vals = [] - - n_block = 1 - - while n > 0 - vars_to_solve = unknowns[vars[:,vars[2,:] .== n][1,:]] - - eqs_to_solve = ss_equations[eqs[:,eqs[2,:] .== n][1,:]] - - # try symbolically and use numerical if it does not work - if verbose - println("Solved: ",string.(eqs_to_solve)," for: ",Symbol.(vars_to_solve), " numerically.") - end - - push!(solved_vars,Symbol.(vars_to_solve)) - push!(solved_vals,Meta.parse.(string.(eqs_to_solve))) - - syms_in_eqs = Set() - - for i in eqs_to_solve - push!(syms_in_eqs, get_symbols(i)...) - end - - # println(syms_in_eqs) - push!(atoms_in_equations_list,setdiff(syms_in_eqs, solved_vars[end])) - - # calib_pars = [] - calib_pars_input = [] - relevant_pars = reduce(union,vcat(๐“‚.constants.post_model_macro.par_list_aux_SS,๐“‚.constants.post_parameters_macro.par_calib_list)[eqs[:,eqs[2,:] .== n][1,:]]) - relevant_pars_across = union(relevant_pars_across,relevant_pars) - - iii = 1 - for parss in union(๐“‚.constants.post_complete_parameters.parameters,๐“‚.constants.post_parameters_macro.parameters_as_function_of_parameters) - # valss = ๐“‚.parameter_values[i] - if :($parss) โˆˆ relevant_pars - # push!(calib_pars,:($parss = parameters_and_solved_vars[$iii])) - push!(calib_pars_input,:($parss)) - iii += 1 - end - end - - - # guess = Expr[] - # untransformed_guess = Expr[] - result = Expr[] - sorted_vars = sort(solved_vars[end]) - # sorted_vars = sort(setdiff(solved_vars[end],๐“‚.constants.post_model_macro.โž•_vars)) - for (i, parss) in enumerate(sorted_vars) - # push!(guess,:($parss = guess[$i])) - # push!(untransformed_guess,:($parss = undo_transform(guess[$i],transformation_level))) - push!(result,:($parss = sol[$i])) - end - - - # separate out auxiliary variables (nonnegativity) - nnaux = [] - # nnaux_linear = [] - # nnaux_error = [] - # push!(nnaux_error, :(aux_error = 0)) - solved_vals_local = Expr[] - # solved_vals_in_place = Expr[] - - eq_idx_in_block_to_solve = eqs[:,eqs[2,:] .== n][1,:] - - - other_vrs_eliminated_by_sympy = Set() - - for (i,val) in enumerate(solved_vals[end]) - if typeof(val) โˆˆ [Symbol,Float64,Int] - push!(solved_vals_local,val) - # push!(solved_vals_in_place, :(โ„ฐ[$i] = $val)) - else - if eq_idx_in_block_to_solve[i] โˆˆ ๐“‚.constants.post_model_macro.ss_equations_with_aux_variables - val = vcat(๐“‚.equations.steady_state_aux,๐“‚.equations.calibration)[eq_idx_in_block_to_solve[i]] - push!(nnaux,:($(val.args[2]) = max(eps(),$(val.args[3])))) - push!(other_vrs_eliminated_by_sympy, val.args[2]) - # push!(nnaux_linear,:($val)) - push!(solved_vals_local,:($val)) - # push!(solved_vals_in_place,:(โ„ฐ[$i] = $val)) - # push!(nnaux_error, :(aux_error += min(eps(),$(val.args[3])))) - else - push!(solved_vals_local,postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, val)) - # push!(solved_vals_in_place, :(โ„ฐ[$i] = $(postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, val)))) - end - end - end - - # println(other_vrs_eliminated_by_sympy) - # sort nnaux vars so that they enter in right order. avoid using a variable before it is declared - # println(nnaux) - if length(nnaux) > 1 - all_symbols = map(x->x.args[1],nnaux) #relevant symbols come first in respective equations - - nn_symbols = map(x->intersect(all_symbols,x), get_symbols.(nnaux)) - - inc_matrix = fill(0,length(all_symbols),length(all_symbols)) - - for i in 1:length(all_symbols) - for k in 1:length(nn_symbols) - inc_matrix[i,k] = collect(all_symbols)[i] โˆˆ collect(nn_symbols)[k] - end - end - - QQ, P, R, nmatch, n_blocks = BlockTriangularForm.order(sparse(inc_matrix)) - - nnaux = nnaux[QQ] - # nnaux_linear = nnaux_linear[QQ] - end - - - # other_vars = [] - other_vars_input = [] - # other_vars_inverse = [] - other_vrs = intersect( setdiff( union(๐“‚.constants.post_model_macro.var, ๐“‚.equations.calibration_parameters, ๐“‚.constants.post_model_macro.โž•_vars), - sort(solved_vars[end]) ), - union(syms_in_eqs, other_vrs_eliminated_by_sympy, setdiff(reduce(union, get_symbols.(nnaux), init = []), map(x->x.args[1],nnaux)) ) ) - - for var in other_vrs - # var_idx = findfirst(x -> x == var, union(๐“‚.constants.post_model_macro.var,๐“‚.calibration_equations_parameters)) - # push!(other_vars,:($(var) = parameters_and_solved_vars[$iii])) - push!(other_vars_input,:($(var))) - iii += 1 - # push!(other_vars_inverse,:(๐“‚.SS_init_guess[$var_idx] = $(var))) - end - - parameters_and_solved_vars = vcat(calib_pars_input, other_vrs) - - ng = length(sorted_vars) - np = length(parameters_and_solved_vars) - nd = 0 - nx = iii - 1 - - - Symbolics.@variables ๐”Š[1:ng] ๐”“[1:np] - - - parameter_dict = Dict{Symbol, Symbol}() - back_to_array_dict = Dict{Symbolics.Num, Symbolics.Num}() - # aux_vars = Symbol[] - # aux_expr = [] - - - for (i,v) in enumerate(sorted_vars) - push!(parameter_dict, v => :($(Symbol("๐”Š_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”Š_$i"))), @__MODULE__) => ๐”Š[i]) - end - - for (i,v) in enumerate(parameters_and_solved_vars) - push!(parameter_dict, v => :($(Symbol("๐”“_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”“_$i"))), @__MODULE__) => ๐”“[i]) - end - - # for (i,v) in enumerate(ss_and_aux_equations_dep) - # push!(aux_vars, v.args[1]) - # push!(aux_expr, v.args[2]) - # end - - # aux_replacements = Dict(aux_vars .=> aux_expr) - - replaced_solved_vals = solved_vals_local |> - # x -> replace_symbols.(x, Ref(aux_replacements)) |> - x -> replace_symbols.(x, Ref(parameter_dict)) |> - x -> Symbolics.parse_expr_to_symbolic.(x, Ref(@__MODULE__)) |> - x -> Symbolics.substitute.(x, Ref(back_to_array_dict)) - - lennz = length(replaced_solved_vals) - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, calc_block! = Symbolics.build_function(replaced_solved_vals, ๐”Š, ๐”“, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - # ๐ท = zeros(Symbolics.Num, nd) - - # ฯตแตƒ = zeros(nd) - - # calc_block_aux!(๐ท, ๐”Š, ๐”“) - - ฯตหข = zeros(Symbolics.Num, ng) - - ฯต = zeros(ng) - - # calc_block!(ฯตหข, ๐”Š, ๐”“, ๐ท) - - โˆ‚block_โˆ‚parameters_and_solved_vars = Symbolics.sparsejacobian(replaced_solved_vals, ๐”Š) # nฯต x nx - - lennz = nnz(โˆ‚block_โˆ‚parameters_and_solved_vars) - - if (lennz / length(โˆ‚block_โˆ‚parameters_and_solved_vars) > density_threshold) || (length(โˆ‚block_โˆ‚parameters_and_solved_vars) < min_length) - derivatives_mat = convert(Matrix, โˆ‚block_โˆ‚parameters_and_solved_vars) - buffer = zeros(Float64, size(โˆ‚block_โˆ‚parameters_and_solved_vars)) - else - derivatives_mat = โˆ‚block_โˆ‚parameters_and_solved_vars - buffer = similar(โˆ‚block_โˆ‚parameters_and_solved_vars, Float64) - buffer.nzval .= 1 - end - - chol_buff = buffer * buffer' - - chol_buff += โ„’.I - - prob = ๐’ฎ.LinearProblem(chol_buff, ฯต, ๐’ฎ.CholeskyFactorization()) + end - chol_buffer = ๐’ฎ.init(prob, ๐’ฎ.CholeskyFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end - prob = ๐’ฎ.LinearProblem(buffer, ฯต, ๐’ฎ.LUFactorization()) + _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”“, ๐”˜, + cse = cse, + skipzeros = skipzeros, + # nanmath = false, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} - lu_buffer = ๐’ฎ.init(prob, ๐’ฎ.LUFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars = buffer + ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars = func_exprs - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”Š, ๐”“, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - - Symbolics.@variables ๐”Š[1:ng+nx] - - ext_diff = Symbolics.Num[] - for i in 1:nx - push!(ext_diff, ๐”“[i] - ๐”Š[ng + i]) - end - replaced_solved_vals_ext = vcat(replaced_solved_vals, ext_diff) - - _, calc_ext_block! = Symbolics.build_function(replaced_solved_vals_ext, ๐”Š, ๐”“, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - ฯตแต‰ = zeros(ng + nx) - - # ฯตหขแต‰ = zeros(Symbolics.Num, ng + nx) - - # calc_block_aux!(๐ท, ๐”Š, ๐”“) - - # Evaluate the function symbolically - # calc_ext_block!(ฯตหขแต‰, ๐”Š, ๐”“, ๐ท) - - โˆ‚ext_block_โˆ‚parameters_and_solved_vars = Symbolics.sparsejacobian(replaced_solved_vals_ext, ๐”Š) # nฯต x nx - - lennz = nnz(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) - - if (lennz / length(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) > density_threshold) || (length(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) < min_length) - derivatives_mat_ext = convert(Matrix, โˆ‚ext_block_โˆ‚parameters_and_solved_vars) - ext_buffer = zeros(Float64, size(โˆ‚ext_block_โˆ‚parameters_and_solved_vars)) - else - derivatives_mat_ext = โˆ‚ext_block_โˆ‚parameters_and_solved_vars - ext_buffer = similar(โˆ‚ext_block_โˆ‚parameters_and_solved_vars, Float64) - ext_buffer.nzval .= 1 - end - - ext_chol_buff = ext_buffer * ext_buffer' + return nothing +end - ext_chol_buff += โ„’.I - prob = ๐’ฎ.LinearProblem(ext_chol_buff, ฯตแต‰, ๐’ฎ.CholeskyFactorization()) - ext_chol_buffer = ๐’ฎ.init(prob, ๐’ฎ.CholeskyFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) - prob = ๐’ฎ.LinearProblem(ext_buffer, ฯตแต‰, ๐’ฎ.LUFactorization()) - ext_lu_buffer = ๐’ฎ.init(prob, ๐’ฎ.LUFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, ext_func_exprs = Symbolics.build_function(derivatives_mat_ext, ๐”Š, ๐”“, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - +function calculate_first_order_obc_solution!(๐“‚::โ„ณ, constants, opts::CalculationOptions) + # Cache hit: return if valid for current parameters + if cache_valid_for_parameters(๐“‚.caches.valid_for.first_order_obc_solution, ๐“‚.parameter_values) && + !isempty(๐“‚.caches.first_order_obc_solution_matrix) + return nothing + end - push!(NSSS_solver_cache_init_tmp,fill(1.205996189998029, length(sorted_vars))) - push!(NSSS_solver_cache_init_tmp,[Inf]) + write_parameters_input!(๐“‚, :activeแต’แต‡แถœshocks => 1, verbose = false) - # WARNING: infinite bounds are transformed to 1e12 - lbs = [] - ubs = [] - - limit_boundaries = 1e12 + โˆ‡ฬ‚โ‚ = calculate_jacobian(๐“‚.parameter_values, ๐“‚.caches.non_stochastic_steady_state, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces, caching = false) - for i in vcat(sorted_vars, calib_pars_input, other_vars_input) - if haskey(๐“‚.constants.post_parameters_macro.bounds, i) - push!(lbs,๐“‚.constants.post_parameters_macro.bounds[i][1] == -Inf ? -limit_boundaries+rand() : ๐“‚.constants.post_parameters_macro.bounds[i][1]) - push!(ubs,๐“‚.constants.post_parameters_macro.bounds[i][2] == Inf ? limit_boundaries-rand() : ๐“‚.constants.post_parameters_macro.bounds[i][2]) - else - push!(lbs,-limit_boundaries+rand()) - push!(ubs,limit_boundaries+rand()) - end - end + ลœโ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡ฬ‚โ‚, + constants, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + initial_guess = ๐“‚.caches.qme_solution, + caching = false) - push!(SS_solve_func,:(params_and_solved_vars = [$(calib_pars_input...),$(other_vars_input...)])) + update_perturbation_counter!(๐“‚.counters, solved, order = 1) - push!(SS_solve_func,:(lbs = [$(lbs...)])) - push!(SS_solve_func,:(ubs = [$(ubs...)])) - - push!(SS_solve_func,:(inits = [max.(lbs[1:length(closest_solution[$(2*(n_block-1)+1)])], min.(ubs[1:length(closest_solution[$(2*(n_block-1)+1)])], closest_solution[$(2*(n_block-1)+1)])), closest_solution[$(2*n_block)]])) - - push!(SS_solve_func,:(solution = block_solver(length(params_and_solved_vars) == 0 ? [0.0] : params_and_solved_vars, - $(n_block), - ๐“‚.NSSS.solve_blocks_in_place[$(n_block)], - # ๐“‚.ss_solve_blocks[$(n_block)], - # ๐“‚.ss_solve_blocks_no_transform[$(n_block)], - # f, - inits, - lbs, - ubs, - solver_parameters, - fail_fast_solvers_only, - cold_start, - verbose))) - - # push!(SS_solve_func,:(solution = block_solver_RD(length([$(calib_pars_input...),$(other_vars_input...)]) == 0 ? [0.0] : [$(calib_pars_input...),$(other_vars_input...)])))#, - - push!(SS_solve_func,:(iters += solution[2][2])) - push!(SS_solve_func,:(solution_error += solution[2][1])) - push!(SS_solve_func,:(sol = solution[1])) + write_parameters_input!(๐“‚, :activeแต’แต‡แถœshocks => 0, verbose = false) - # push!(SS_solve_func,:(solution = block_solver_RD(length([$(calib_pars_input...),$(other_vars_input...)]) == 0 ? [0.0] : [$(calib_pars_input...),$(other_vars_input...)])))#, - - # push!(SS_solve_func,:(solution_error += sum(abs2,๐“‚.ss_solve_blocks[$(n_block)](length([$(calib_pars_input...),$(other_vars_input...)]) == 0 ? [0.0] : [$(calib_pars_input...),$(other_vars_input...)],solution)))) + # Cache write + stamp + ๐“‚.caches.first_order_obc_solution_matrix = ลœโ‚ + ๐“‚.caches.valid_for.first_order_obc_solution = Float64.(๐“‚.parameter_values) - push!(SS_solve_func,:($(result...))) - - push!(SS_solve_func,:(NSSS_solver_cache_tmp = [NSSS_solver_cache_tmp..., typeof(sol) == Vector{Float64} ? sol : โ„ฑ.value.(sol)])) - push!(SS_solve_func,:(NSSS_solver_cache_tmp = [NSSS_solver_cache_tmp..., typeof(params_and_solved_vars) == Vector{Float64} ? params_and_solved_vars : โ„ฑ.value.(params_and_solved_vars)])) - - # Create nonlinear solver workspaces for regular and extended problems - workspace = Nonlinear_solver_workspace(ฯต, buffer, chol_buffer, lu_buffer) - ext_workspace = Nonlinear_solver_workspace(ฯตแต‰, ext_buffer, ext_chol_buffer, ext_lu_buffer) - - push!(๐“‚.NSSS.solve_blocks_in_place, - ss_solve_block( - function_and_jacobian(calc_block!::Function, func_exprs::Function, workspace), - function_and_jacobian(calc_ext_block!::Function, ext_func_exprs::Function, ext_workspace) - ) - ) + return nothing +end - n_block += 1 - - n -= 1 +function solve_steady_state!(๐“‚::โ„ณ, + opts::CalculationOptions, + ss_solver_parameters_algorithm::Symbol, + ss_solver_parameters_maxtime::Real; + silent::Bool = false)::Tuple{Vector{Float64}, Float64, Bool} + """ + Internal function to solve and constants the steady state. + Returns: (SS_and_pars, solution_error, found_solution) + """ + start_time = time() + + if ๐“‚.constants.post_parameters_macro.precompile + return Float64[], 0.0, false + end + + if !(๐“‚.functions.NSSS_custom isa Function) + if !silent + print("Find non-stochastic steady state:\t\t\t\t\t") + end end - - push!(NSSS_solver_cache_init_tmp,[Inf]) - push!(NSSS_solver_cache_init_tmp,fill(Inf,length(๐“‚.constants.post_complete_parameters.parameters))) - push!(๐“‚.caches.solver_cache,NSSS_solver_cache_init_tmp) - - unknwns = Symbol.(unknowns) - - parameters_only_in_par_defs = Set() - # add parameters from parameter definitions - if length(๐“‚.equations.calibration_no_var) > 0 - atoms = reduce(union, get_symbols.(๐“‚.equations.calibration_no_var)) - [push!(atoms_in_equations, a) for a in atoms] - [push!(parameters_only_in_par_defs, a) for a in atoms] - end - # ๐“‚.par = union(๐“‚.par,setdiff(parameters_only_in_par_defs,๐“‚.parameters_as_function_of_parameters)) + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts, cold_start = true) + SS_and_pars = copy(SS_and_pars) # decouple from workspace output_buffer before select_fastest overwrites it - parameters_in_equations = [] - - for (i, parss) in enumerate(๐“‚.constants.post_complete_parameters.parameters) - if parss โˆˆ union(atoms_in_equations, relevant_pars_across) - push!(parameters_in_equations, :($parss = parameters[$i])) + found_solution = true + + if !(๐“‚.functions.NSSS_custom isa Function) + select_fastest_SS_solver_parameters!(๐“‚, tol = opts.tol) + + if solution_error > opts.tol.nsss.acceptance_tol + found_solution = find_SS_solver_parameters!(Val(ss_solver_parameters_algorithm), ๐“‚, tol = opts.tol, verbosity = 0, maxtime = ss_solver_parameters_maxtime, maxiter = 1000000000) + + if found_solution + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts, cold_start = true) + end end end - dependencies = [] - for (i, a) in enumerate(atoms_in_equations_list) - push!(dependencies, solved_vars[i] => intersect(a, union(๐“‚.constants.post_model_macro.var, ๐“‚.constants.post_complete_parameters.parameters))) + if !(๐“‚.functions.NSSS_custom isa Function) + if !silent + println(round(time() - start_time, digits = 3), " seconds") + end end - - push!(dependencies, :SS_relevant_calibration_parameters => intersect(reduce(union, atoms_in_equations_list), ๐“‚.constants.post_complete_parameters.parameters)) - - ๐“‚.NSSS.dependencies = dependencies - - dyn_exos = [] - for dex in union(๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future) - push!(dyn_exos,:($dex = 0)) + if !found_solution + @warn "Could not find non-stochastic steady state. Consider setting bounds on variables or calibrated parameters in the `@parameters` section (e.g. `k > 10`)." end - - push!(SS_solve_func,:($(dyn_exos...))) - - # push!(SS_solve_func,:(push!(NSSS_solver_cache_tmp, params_scaled_flt))) - push!(SS_solve_func,:(if length(NSSS_solver_cache_tmp) == 0 NSSS_solver_cache_tmp = [copy(params_flt)] else NSSS_solver_cache_tmp = [NSSS_solver_cache_tmp..., copy(params_flt)] end)) - push!(SS_solve_func,:(current_best = sqrt(sum(abs2,๐“‚.caches.solver_cache[end][end] - params_flt))))# / max(sum(abs2,๐“‚.caches.solver_cache[end][end]), sum(abs2,params_flt)))) - - push!(SS_solve_func,:(for pars in ๐“‚.caches.solver_cache - latest = sqrt(sum(abs2,pars[end] - params_flt))# / max(sum(abs2,pars[end]), sum(abs,params_flt)) - if latest <= current_best - current_best = latest - end - end)) + return SS_and_pars, solution_error, found_solution +end - push!(SS_solve_func,:(if (current_best > 1e-8) && (solution_error < tol.NSSS_acceptance_tol) - reverse_diff_friendly_push!(๐“‚.caches.solver_cache, NSSS_solver_cache_tmp) - # solved_scale = scale - end)) +# Centralised helper to write symbolic derivatives and map functions +function write_symbolic_derivatives!(๐“‚::โ„ณ; perturbation_order::Int = 1, silent::Bool = false) + start_time = time() - # fix parameter bounds - par_bounds = [] - - for varpar in intersect(๐“‚.constants.post_complete_parameters.parameters,union(atoms_in_equations, relevant_pars_across)) - if haskey(๐“‚.constants.post_parameters_macro.bounds, varpar) - push!(par_bounds, :($varpar = min(max($varpar,$(๐“‚.constants.post_parameters_macro.bounds[varpar][1])),$(๐“‚.constants.post_parameters_macro.bounds[varpar][2])))) + if !silent + if perturbation_order == 1 + print("Take symbolic derivatives up to first order:\t\t\t\t") + elseif perturbation_order == 2 + print("Take symbolic derivatives up to second order:\t\t\t\t") + elseif perturbation_order == 3 + print("Take symbolic derivatives up to third order:\t\t\t\t") end end - solve_exp = :(function solve_SS(initial_parameters::Vector{Real}, - ๐“‚::โ„ณ, - tol::Tolerances, - # fail_fast_solvers_only::Bool, - verbose::Bool, - cold_start::Bool, - solver_parameters::Vector{solver_parameters}) - initial_parameters = typeof(initial_parameters) == Vector{Float64} ? initial_parameters : โ„ฑ.value.(initial_parameters) - - parameters = copy(initial_parameters) - params_flt = copy(initial_parameters) - - current_best = sum(abs2,๐“‚.caches.solver_cache[end][end] - initial_parameters) - closest_solution_init = ๐“‚.caches.solver_cache[end] - - for pars in ๐“‚.caches.solver_cache - latest = sum(abs2,pars[end] - initial_parameters) - if latest <= current_best - current_best = latest - closest_solution_init = pars - end - end - - # closest_solution = closest_solution_init - # solution_error = 1.0 - # iters = 0 - range_iters = 0 - solution_error = 1.0 - solved_scale = 0 - # range_length = [ 1, 2, 4, 8,16,32,64,128,1024] - scale = 1.0 - - while range_iters <= 500 && !(solution_error < tol.NSSS_acceptance_tol && solved_scale == 1) - range_iters += 1 - fail_fast_solvers_only = range_iters > 1 ? true : false - - # for range_ in range_length - # rangee = range(0,1,range_+1) - # for scale in rangee[2:end] - # scale = 6*scale^5 - 15*scale^4 + 10*scale^3 # smootherstep - - # if scale <= solved_scale continue end - - current_best = sum(abs2,๐“‚.caches.solver_cache[end][end] - initial_parameters) - closest_solution = ๐“‚.caches.solver_cache[end] - - for pars in ๐“‚.caches.solver_cache - latest = sum(abs2,pars[end] - initial_parameters) - if latest <= current_best - current_best = latest - closest_solution = pars - end - end - - # Zero initial value if starting without guess - if !isfinite(sum(abs,closest_solution[2])) - closest_solution = copy(closest_solution) - for i in 1:2:length(closest_solution) - closest_solution[i] = zeros(length(closest_solution[i])) - end - end - - # println(closest_solution) - - if all(isfinite,closest_solution[end]) && initial_parameters != closest_solution_init[end] - parameters = scale * initial_parameters + (1 - scale) * closest_solution_init[end] - else - parameters = copy(initial_parameters) - end - params_flt = parameters - - # println(parameters) - - $(parameters_in_equations...) - $(par_bounds...) - $(๐“‚.equations.calibration_no_var...) - NSSS_solver_cache_tmp = [] - solution_error = 0.0 - iters = 0 - $(SS_solve_func...) - - if solution_error < tol.NSSS_acceptance_tol - # println("solved for $scale; $range_iters") - solved_scale = scale - if scale == 1 - # return ComponentVector([$(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))...), $(๐“‚.calibration_equations_parameters...)], Axis([sort(union(๐“‚.constants.post_model_macro.exo_present,๐“‚.constants.post_model_macro.var))...,๐“‚.calibration_equations_parameters...])), solution_error - # NSSS_solution = [$(Symbol.(replace.(string.(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))...), $(๐“‚.calibration_equations_parameters...)] - # NSSS_solution[abs.(NSSS_solution) .< 1e-12] .= 0 # doesn't work with Zygote - return [$(Symbol.(replace.(string.(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))...), $(๐“‚.equations.calibration_parameters...)], (solution_error, iters) - else - reverse_diff_friendly_push!(NSSS_solver_cache_scale, NSSS_solver_cache_tmp) - end - - if scale > .95 - scale = 1 - else - # scale = (scale + 1) / 2 - scale = scale * .4 + .6 - end - # else - # println("no sol") - # scale = (scale + solved_scale) / 2 - # println("scale $scale") - # elseif scale == 1 && range_ == range_length[end] - # return [$(Symbol.(replace.(string.(sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))...), $(๐“‚.calibration_equations_parameters...)], (solution_error, iters) - end - # end - end - return zeros($(length(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_past,๐“‚.constants.post_model_macro.exo_future)) + length(๐“‚.equations.calibration_parameters))), (1, 0) - end) + write_auxiliary_indices!(๐“‚) + + write_functions_mapping!(๐“‚, perturbation_order) - ๐“‚.functions.NSSS_solve = @RuntimeGeneratedFunction(solve_exp) - # ๐“‚.functions.NSSS_solve = eval(solve_exp) + if !silent + println(round(time() - start_time, digits = 3), " seconds") + end return nothing end -function reverse_diff_friendly_push!(x,y) - @ignore_derivatives push!(x,y) -end - function calculate_SS_solver_runtime_and_loglikelihood(pars::Vector{Float64}, ๐“‚::โ„ณ; tol::Tolerances = Tolerances())::Float64 log_lik = 0.0 log_lik -= -sum(pars[1:19]) # logpdf of a gamma dist with mean and variance 1 @@ -5914,14 +4927,14 @@ function calculate_SS_solver_runtime_and_loglikelihood(pars::Vector{Float64}, par_inputs = solver_parameters(pars..., 1, 0.0, 2) - while length(๐“‚.caches.solver_cache) > 1 - pop!(๐“‚.caches.solver_cache) + while length(๐“‚.caches.solver) > 1 + pop!(๐“‚.caches.solver) end - runtime = @elapsed outmodel = try ๐“‚.functions.NSSS_solve(๐“‚.parameter_values, ๐“‚, tol, false, true, [par_inputs]) catch end + runtime = @elapsed outmodel = try solve_nsss_wrapper(๐“‚.parameter_values, ๐“‚, tol, false, true, [par_inputs]) catch end runtime = outmodel isa Tuple{Vector{Float64}, Tuple{Float64, Int64}} ? - (outmodel[2][1] > tol.NSSS_acceptance_tol) || !isfinite(outmodel[2][1]) ? + (outmodel[2][1] > tol.nsss.acceptance_tol) || !isfinite(outmodel[2][1]) ? 10 : runtime : 10 @@ -5974,10 +4987,14 @@ function find_SS_solver_parameters!(::Val{:ESCH}, ๐“‚::โ„ณ; maxtime::Real = 120 par_inputs = solver_parameters(pars..., 1, 0.0, 2) - SS_and_pars, (solution_error, iters) = ๐“‚.functions.NSSS_solve(๐“‚.parameter_values, ๐“‚, tol, false, true, [par_inputs]) + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(๐“‚.parameter_values, ๐“‚, tol, false, true, [par_inputs]) - if solution_error < tol.NSSS_acceptance_tol + if solution_error < tol.nsss.acceptance_tol push!(DEFAULT_SOLVER_PARAMETERS, par_inputs) + ๐“‚.constants.post_complete_parameters = update_post_complete_parameters( + ๐“‚.constants.post_complete_parameters; + nsss_fastest_solver_parameter_idx = length(DEFAULT_SOLVER_PARAMETERS), + ) return true else return false @@ -5985,53 +5002,87 @@ function find_SS_solver_parameters!(::Val{:ESCH}, ๐“‚::โ„ณ; maxtime::Real = 120 end -function select_fastest_SS_solver_parameters!(๐“‚::โ„ณ; tol::Tolerances = Tolerances()) - best_param = DEFAULT_SOLVER_PARAMETERS[1] +function select_fastest_SS_solver_parameters!(๐“‚::โ„ณ; + tol::Tolerances = Tolerances(), + n_samples::Int = 100)::Nothing + @assert n_samples > 1 "n_samples must be greater than 1." + @assert n_samples รท 2 >= 1 "n_samples must be at least 2." - best_time = Inf + best_idx = 1 + best_score = Inf solved = false - solved_NSSS = ๐“‚.caches.solver_cache[end] + solved_NSSS = ๐“‚.caches.solver[end] - for p in DEFAULT_SOLVER_PARAMETERS - total_time = 0.0 + for (i_param, p) in enumerate(DEFAULT_SOLVER_PARAMETERS) + times = Vector{Float64}(undef, n_samples) + valid = true - for _ in 1:100 + for i in 1:n_samples start_time = time() - while length(๐“‚.caches.solver_cache) > 1 - pop!(๐“‚.caches.solver_cache) + while length(๐“‚.caches.solver) > 1 + pop!(๐“‚.caches.solver) end - SS_and_pars, (solution_error, iters) = ๐“‚.functions.NSSS_solve(๐“‚.parameter_values, ๐“‚, tol, false, true, [p]) + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(๐“‚.parameter_values, ๐“‚, tol, false, true, [p]) elapsed_time = time() - start_time - total_time += elapsed_time - - if solution_error > tol.NSSS_acceptance_tol - total_time = 1e7 + times[i] = elapsed_time + + if solution_error > tol.nsss.acceptance_tol + valid = false break end end - if total_time < best_time - best_time = total_time - best_param = p - end + if valid + sort!(times) + score = times[n_samples รท 2] + + if !isfinite(best_score) || score < best_score + best_score = score + best_idx = i_param + end - solved = true + solved = true + end end - while length(๐“‚.caches.solver_cache) > 1 - pop!(๐“‚.caches.solver_cache) + while length(๐“‚.caches.solver) > 1 + pop!(๐“‚.caches.solver) end - push!(๐“‚.caches.solver_cache, solved_NSSS) + push!(๐“‚.caches.solver, solved_NSSS) if solved - pushfirst!(DEFAULT_SOLVER_PARAMETERS, best_param) + ๐“‚.constants.post_complete_parameters = update_post_complete_parameters( + ๐“‚.constants.post_complete_parameters; + nsss_fastest_solver_parameter_idx = best_idx, + ) + end + + return nothing +end + +function update_init_buf!(init_buf::AbstractVector{T}, lbs, ubs, n_guess, ssv_val, sv_val, guess, use_ssv::Bool) where {T} + @inbounds for i in 1:n_guess + if use_ssv + v = clamp(ssv_val, lbs[i], ubs[i]) + init_buf[i] = ubs[i] <= one(T) ? T(0.1) : v + else + g = guess[i] + v = g < T(1e12) ? g : sv_val + init_buf[i] = clamp(v, lbs[i], ubs[i]) + end + end +end + +function update_sol_values!(sol_values::AbstractVector{T}, sol_new::AbstractVector{T}, lbs::AbstractVector{T}, ubs::AbstractVector{T}, n_guess::Int) where {T} + @inbounds for i in 1:n_guess + sol_values[i] = clamp(sol_new[i], lbs[i], ubs[i]) end end @@ -6051,15 +5102,34 @@ function solve_ss(SS_optimizer::Function, solver_params::solver_parameters, extended_problem::Bool, separate_starting_value::Union{Bool,T})::Tuple{Vector{T}, Vector{Int}, T, T} where T <: AbstractFloat - xtol = tol.NSSS_xtol - ftol = tol.NSSS_ftol - rel_xtol = tol.NSSS_rel_xtol - - if separate_starting_value isa Float64 - sol_values_init = max.(lbs[1:length(guess)], min.(ubs[1:length(guess)], fill(separate_starting_value, length(guess)))) - sol_values_init[ubs[1:length(guess)] .<= 1] .= .1 # capture cases where part of values is small + ftol = tol.nsss.ftol + n_guess = length(guess) + init_buf = SS_solve_block.ss_problem.workspace.best_previous_guess + use_ssv = separate_starting_value isa Float64 + ssv_val = use_ssv ? T(separate_starting_value) : zero(T) + sv_val = T(solver_params.starting_value) + update_init_buf!(init_buf, lbs, ubs, n_guess, ssv_val, sv_val, guess, use_ssv) + + if !extended_problem + lb_core = SS_solve_block.ss_problem.workspace.l_bounds + ub_core = SS_solve_block.ss_problem.workspace.u_bounds + copyto!(lb_core, 1, lbs, 1, n_guess) + copyto!(ub_core, 1, ubs, 1, n_guess) + end + + optimizer_init = if extended_problem + ext_init = SS_solve_block.extended_ss_problem.workspace.best_previous_guess + @inbounds begin + for i in 1:n_guess + ext_init[i] = init_buf[i] + end + for i in 1:length(closest_parameters_and_solved_vars) + ext_init[n_guess + i] = closest_parameters_and_solved_vars[i] + end + end + ext_init else - sol_values_init = max.(lbs[1:length(guess)], min.(ubs[1:length(guess)], [g < 1e12 ? g : solver_params.starting_value for g in guess])) + init_buf end sol_new_tmp, info = SS_optimizer( extended_problem ? SS_solve_block.extended_ss_problem : SS_solve_block.ss_problem, @@ -6078,48 +5148,56 @@ function solve_ss(SS_optimizer::Function, # end # sol_new_tmp, info = SS_optimizer( extended_problem ? ext_function_to_optimize : function_to_optimize, - extended_problem ? vcat(sol_values_init, closest_parameters_and_solved_vars) : sol_values_init, + optimizer_init, parameters_and_solved_vars, - extended_problem ? lbs : lbs[1:length(guess)], - extended_problem ? ubs : ubs[1:length(guess)], + extended_problem ? lbs : SS_solve_block.ss_problem.workspace.l_bounds, + extended_problem ? ubs : SS_solve_block.ss_problem.workspace.u_bounds, solver_params, tol = tol ) - sol_new = isnothing(sol_new_tmp) ? sol_new_tmp : sol_new_tmp[1:length(guess)] - sol_minimum = info[4] # isnan(sum(abs, info[4])) ? Inf : โ„’.norm(info[4]) rel_sol_minimum = info[3] - - sol_values = max.(lbs[1:length(guess)], min.(ubs[1:length(guess)], sol_new)) + + sol_values = SS_solve_block.ss_problem.workspace.best_current_guess + if isnothing(sol_new_tmp) + update_sol_values!(sol_values, init_buf, lbs, ubs, n_guess) + else + update_sol_values!(sol_values, sol_new_tmp, lbs, ubs, n_guess) + end total_iters[1] += info[1] total_iters[2] += info[2] - extended_problem_str = extended_problem ? "(extended problem) " : "" + if sol_minimum < ftol && verbose + extended_problem_str = extended_problem ? "(extended problem) " : "" - if separate_starting_value isa Bool - starting_value_str = "" - else - starting_value_str = "and starting point: $separate_starting_value" - end + if separate_starting_value isa Bool + starting_value_str = "" + else + starting_value_str = "and starting point: $separate_starting_value" + end - if all(guess .< 1e12) && separate_starting_value isa Bool - any_guess_str = "previous solution, " - elseif any(guess .< 1e12) && separate_starting_value isa Bool - any_guess_str = "provided guess, " - else - any_guess_str = "" - end + has_small_guess = false + all_small_guess = true + @inbounds for i in eachindex(guess) + is_small = guess[i] < T(1e12) + has_small_guess |= is_small + all_small_guess &= is_small + end - # max_resid = maximum(abs,ss_solve_blocks(parameters_and_solved_vars, sol_values)) + if all_small_guess && separate_starting_value isa Bool + any_guess_str = "previous solution, " + elseif has_small_guess && separate_starting_value isa Bool + any_guess_str = "provided guess, " + else + any_guess_str = "" + end - SS_solve_block.ss_problem.func(SS_solve_block.ss_problem.workspace.func_buffer, sol_values, parameters_and_solved_vars) - - max_resid = maximum(abs, SS_solve_block.ss_problem.workspace.func_buffer) + SS_solve_block.ss_problem.func(SS_solve_block.ss_problem.workspace.func_buffer, sol_values, parameters_and_solved_vars) + max_resid = maximum(abs, SS_solve_block.ss_problem.workspace.func_buffer) - if sol_minimum < ftol && verbose - println("Block: $n_block - Solved $(extended_problem_str) using ",string(SS_optimizer),", $(any_guess_str)$(starting_value_str); maximum residual = $max_resid") + println("Block: $n_block - Solved $(extended_problem_str) using ",string(SS_optimizer),", $(any_guess_str)$(starting_value_str); maximum residual = $max_resid") end return sol_values, total_iters, rel_sol_minimum, sol_minimum @@ -6136,6 +5214,7 @@ function block_solver(parameters_and_solved_vars::Vector{T}, lbs::Vector{T}, ubs::Vector{T}, parameters::Vector{solver_parameters}, + preferred_solver_parameter_idx::Int, fail_fast_solvers_only::Bool, cold_start::Bool, verbose::Bool ; @@ -6166,7 +5245,7 @@ function block_solver(parameters_and_solved_vars::Vector{T}, sol_minimum = โ„’.norm(res) if !cold_start - if !isfinite(sol_minimum) || sol_minimum > tol.NSSS_acceptance_tol + if !isfinite(sol_minimum) || sol_minimum > tol.nsss.acceptance_tol # โˆ‡ = ๐’Ÿ.jacobian(x->(ss_solve_blocks(parameters_and_solved_vars, x)), backend, guess) # โˆ‡ฬ‚ = โ„’.lu!(โˆ‡, check = false) @@ -6175,14 +5254,21 @@ function block_solver(parameters_and_solved_vars::Vector{T}, โˆ‡ = SS_solve_block.ss_problem.workspace.jac_buffer - โˆ‡ฬ‚ = โ„’.lu(โˆ‡, check = false) - - if โ„’.issuccess(โˆ‡ฬ‚) - guess_update = โˆ‡ฬ‚ \ res - - new_guess = guess - guess_update - - rel_sol_minimum = โ„’.norm(guess_update) / max(โ„’.norm(new_guess), sol_minimum) + sol_cache = SS_solve_block.ss_problem.workspace.lu_buffer + # sol_cache.A = sol_cache.alg isa ๐’ฎ.FastLUFactorization ? copy(โˆ‡) : โˆ‡ + sol_cache.A = โˆ‡ + # copy!(sol_cache.A, โˆ‡) + sol_cache.b = res + sol = ๐’ฎ.solve!(sol_cache) + + if ๐’ฎ.SciMLBase.successful_retcode(sol.retcode) || sol.retcode == ๐’ฎ.SciMLBase.ReturnCode.Default + guess_update = sol_cache.u + if has_nonfinite(guess_update) + rel_sol_minimum = 1.0 + else + new_guess = guess - guess_update + rel_sol_minimum = โ„’.norm(guess_update) / max(โ„’.norm(new_guess), sol_minimum) + end else rel_sol_minimum = 1.0 end @@ -6193,7 +5279,7 @@ function block_solver(parameters_and_solved_vars::Vector{T}, rel_sol_minimum = 1.0 end - if isfinite(sol_minimum) && sol_minimum < tol.NSSS_acceptance_tol + if isfinite(sol_minimum) && sol_minimum < tol.nsss.acceptance_tol solved_yet = true if verbose @@ -6202,18 +5288,22 @@ function block_solver(parameters_and_solved_vars::Vector{T}, end total_iters = [0,0] + n_solver_parameters = length(parameters) + @assert n_solver_parameters > 0 "At least one steady-state solver parameter set is required." SS_optimizer = levenberg_marquardt + ext_candidates = (true, false) + algo_candidates = (newton, levenberg_marquardt) if cold_start guesses = any(guess .< 1e12) ? [guess, fill(1e12, length(guess))] : [guess] # if guess were provided, loop over them, and then the starting points only - start_vals = (fail_fast_solvers_only ? [false] : Any[false, 1.206, 1.5, 0.7688, 2.0, 0.897]) - + start_vals = fail_fast_solvers_only ? (false,) : (false, T(1.206), T(1.5), T(0.7688), T(2.0), T(0.897)) for g in guesses - for p in parameters - for ext in [true, false] # try first the system where values and parameters can vary, next try the system where only values can vary + for i in 1:n_solver_parameters + p = parameters[i == 1 ? preferred_solver_parameter_idx : (i <= preferred_solver_parameter_idx ? i - 1 : i)] + for ext in ext_candidates # try first the system where values and parameters can vary, next try the system where only values can vary for s in start_vals - if !isfinite(sol_minimum) || sol_minimum > tol.NSSS_acceptance_tol# || rel_sol_minimum > rtol + if !isfinite(sol_minimum) || sol_minimum > tol.nsss.acceptance_tol# || rel_sol_minimum > rtol if solved_yet continue end sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(SS_optimizer, SS_solve_block, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, total_iters, n_block, verbose, @@ -6223,7 +5313,7 @@ function block_solver(parameters_and_solved_vars::Vector{T}, ext, s) - if isfinite(sol_minimum) && sol_minimum < tol.NSSS_acceptance_tol + if isfinite(sol_minimum) && sol_minimum < tol.nsss.acceptance_tol solved_yet = true end end @@ -6233,14 +5323,24 @@ function block_solver(parameters_and_solved_vars::Vector{T}, end else !cold_start - pars = (fail_fast_solvers_only ? [parameters[end]] : unique(parameters)) - - for p in pars #[1:3] # take unique because some parameters might appear more than once - start_vals = (fail_fast_solvers_only ? [false] : Any[false,p.starting_value, 1.206, 1.5, 0.7688, 2.0, 0.897]) - for s in start_vals #, .9, .75, 1.5, -.5, 2, .25] # try first the guess and then different starting values - # for ext in [false, true] # try first the system where only values can vary, next try the system where values and parameters can vary - for algo in [newton, levenberg_marquardt] - if !isfinite(sol_minimum) || sol_minimum > tol.NSSS_acceptance_tol # || rel_sol_minimum > rtol + start_vals = Vector{Union{Bool, T}}(undef, 7) + start_vals[1] = false + start_vals[3] = T(1.206) + start_vals[4] = T(1.5) + start_vals[5] = T(0.7688) + start_vals[6] = T(2.0) + start_vals[7] = T(0.897) + + s_candidates = fail_fast_solvers_only ? @view(start_vals[1:1]) : start_vals + n_parameter_iters = fail_fast_solvers_only ? 1 : n_solver_parameters + fail_fast_parameter_idx = n_solver_parameters == 1 ? 1 : (n_solver_parameters <= preferred_solver_parameter_idx ? n_solver_parameters - 1 : n_solver_parameters) + + for i in 1:n_parameter_iters + p = parameters[fail_fast_solvers_only ? fail_fast_parameter_idx : (i == 1 ? preferred_solver_parameter_idx : (i <= preferred_solver_parameter_idx ? i - 1 : i))] + start_vals[2] = T(p.starting_value) + for s in s_candidates + for algo in algo_candidates + if sol_minimum > tol.nsss.acceptance_tol || !isfinite(sol_minimum) # || rel_sol_minimum > rtol if solved_yet continue end # println("Block: $n_block pre GN - $ext - $sol_minimum - $rel_sol_minimum") sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(algo, SS_solve_block, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, @@ -6254,14 +5354,14 @@ function block_solver(parameters_and_solved_vars::Vector{T}, false, # ext # false) s) - if isfinite(sol_minimum) && sol_minimum < tol.NSSS_acceptance_tol # || rel_sol_minimum > rtol) + if isfinite(sol_minimum) && sol_minimum < tol.nsss.acceptance_tol # || rel_sol_minimum > rtol) solved_yet = true if verbose # println("Block: $n_block, - Solved with $algo using previous solution - $(indexin([ext],[false, true])[1])/2 - $ext - $sol_minimum - $rel_sol_minimum - $total_iters") println("Block: $n_block, - Solved with $algo using previous solution - $sol_minimum - $rel_sol_minimum - $total_iters") end - end + end end end end @@ -6314,153 +5414,230 @@ function block_solver(parameters_and_solved_vars::Vector{T}, end -function calculate_second_order_stochastic_steady_state(parameters::Vector{M}, - ๐“‚::โ„ณ; - opts::CalculationOptions = merge_calculation_options(), - pruning::Bool = false, - estimation::Bool = false) where M - # timer::TimerOutput = TimerOutput(), - # tol::AbstractFloat = 1e-12) - # @timeit_debug timer "Calculate NSSS" begin - # Initialize constants at entry point +function _prepare_stochastic_steady_state_base_terms(parameters::Vector{M}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true) where M constants = initialise_constants!(๐“‚) T = constants.post_model_macro - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts, estimation = estimation) # , timer = timer) + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts, estimation = estimation, caching = caching) - # end # timeit_debug - - if solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error) - # if verbose println("NSSS not found") end # handled within solve function - return zeros(M, T.nVars), false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) + return (false, + zeros(T.nVars), + SS_and_pars, + solution_error, + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0), + constants) end - ms = @ignore_derivatives ensure_model_structure_constants!(constants, ๐“‚.equations.calibration_parameters) + ms = ensure_model_structure_constants!(constants, ๐“‚.equations.calibration_parameters) all_SS = expand_steady_state(SS_and_pars, ms) - # @timeit_debug timer "Calculate Jacobian" begin - - โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - # end # timeit_debug + โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces, caching = caching) - # @timeit_debug timer "Calculate first order solution" begin - - qme_ws = @ignore_derivatives ensure_qme_workspace!(๐“‚) - sylv_ws = @ignore_derivatives ensure_sylvester_1st_order_workspace!(๐“‚) - ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, - constants, - qme_ws, - sylv_ws; - opts = opts, - initial_guess = ๐“‚.caches.qme_solution) - - if solved ๐“‚.caches.qme_solution = qme_sol end - - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + constants, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameters, + caching = caching) - # end # timeit_debug + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) if !solved if opts.verbose println("1st order solution not found") end - return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) + return (false, + all_SS, + SS_and_pars, + solution_error, + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0), + constants) end - # @timeit_debug timer "Calculate Hessian" begin - - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ - - # end # timeit_debug - - # @timeit_debug timer "Calculate second order solution" begin - - ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; - initial_guess = ๐“‚.caches.second_order_solution, - # timer = timer, - opts = opts) - - if eltype(๐’โ‚‚) == Float64 && solved2 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end + โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces, caching = caching) - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) + ๐’โ‚‚_raw, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + initial_guess = ๐“‚.caches.second_order_solution, + opts = opts, + parameter_values = parameters, + caching = caching) - ๐’โ‚‚ = sparse(๐’โ‚‚ * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} - - # end # timeit_debug + update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) if !solved2 if opts.verbose println("2nd order solution not found") end - return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) + return (false, + all_SS, + SS_and_pars, + solution_error, + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0), + constants) end - # @timeit_debug timer "Calculate SSS" begin + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} - ๐’โ‚ = [๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) ๐’โ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] + ๐’โ‚ = [๐’โ‚[:,1:T.nPast_not_future_and_mixed] zeros(T.nVars) ๐’โ‚[:,T.nPast_not_future_and_mixed+1:end]] - aug_stateโ‚ = sparse([zeros(๐“‚.constants.post_model_macro.nPast_not_future_and_mixed); 1; zeros(๐“‚.constants.post_model_macro.nExo)]) + aug_stateโ‚ = sparse([zeros(T.nPast_not_future_and_mixed); 1; zeros(T.nExo)]) + tmp = (T.I_nPast - ๐’โ‚[T.past_not_future_and_mixed_idx,1:T.nPast_not_future_and_mixed]) + tmpฬ„ = โ„’.lu(tmp, check = false) - I_nPast = qme_ws.I_nPast + if !โ„’.issuccess(tmpฬ„) + if opts.verbose println("SSS not found") end + return (false, + all_SS, + SS_and_pars, + solution_error, + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0,0), + spzeros(M,0,0), + zeros(M,0), + constants) + end + + SSSstates = collect(tmp \ (๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2)[T.past_not_future_and_mixed_idx]) + + return (true, + all_SS, + SS_and_pars, + solution_error, + โˆ‡โ‚, + โˆ‡โ‚‚, + ๐’โ‚, + ๐’โ‚‚_raw, + SSSstates, + constants) +end + +function calculate_stochastic_steady_state(::Val{:second_order}, + parameters::Vector{M}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true) where M + # Cache hit: return cached SSS if valid for current parameters + if caching && M === Float64 && !isempty(parameters) && + cache_valid_for_parameters(๐“‚.caches.valid_for.second_order_stochastic_steady_state, parameters) + cached_sss = ๐“‚.caches.second_order_stochastic_steady_state::Vector{M} + if !isempty(cached_sss) + T = ๐“‚.constants.post_model_macro + SS_and_pars = ๐“‚.caches.non_stochastic_steady_state::Vector{M} + โˆ‡โ‚ = Matrix(๐“‚.caches.jacobian)::Matrix{M} + โˆ‡โ‚‚ = sparse(๐“‚.caches.hessian)::SparseMatrixCSC{M, Int} + ๐’โ‚_raw = Matrix(๐“‚.caches.first_order_solution_matrix)::Matrix{M} + ๐’โ‚ = [๐’โ‚_raw[:,1:T.nPast_not_future_and_mixed] zeros(M, T.nVars) ๐’โ‚_raw[:,T.nPast_not_future_and_mixed+1:end]] + ๐’โ‚‚ = sparse(๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} + return cached_sss, true, SS_and_pars, zero(M), โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ + end + end + + common = _prepare_stochastic_steady_state_base_terms(parameters, ๐“‚, opts = opts, estimation = estimation, caching = caching) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + if caching ๐“‚.caches.second_order_stochastic_steady_state = all_SS end + return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) + end - tmp = (I_nPast - ๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed]) + # Expand compressed ๐’โ‚‚_raw to full + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} - tmpฬ„ = @ignore_derivatives โ„’.lu(tmp, check = false) + so = ๐“‚.constants.second_order + kron_sโบ_sโบ = so.kron_sโบ_sโบ + A = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] + Bฬ‚ = ๐’โ‚‚[:,kron_sโบ_sโบ] - if !โ„’.issuccess(tmpฬ„) + SSSstates, converged = solve_stochastic_steady_state_newton(Val(:second_order), ๐’โ‚, ๐’โ‚‚, collect(SSSstates), ๐“‚) + + if !converged if opts.verbose println("SSS not found") end + if caching ๐“‚.caches.second_order_stochastic_steady_state = all_SS end return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) end - SSSstates = collect(tmp \ (๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2)[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]) + state = A * SSSstates + Bฬ‚ * โ„’.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 + result = all_SS + Vector{M}(state) - if pruning - state = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * SSSstates + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2 - converged = true - else - # Get cached computational constants - so = ๐“‚.constants.second_order - s_in_sโบ = @ignore_derivatives so.s_in_sโบ - kron_sโบ_sโบ = @ignore_derivatives so.kron_sโบ_sโบ - - A = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] - Bฬ‚ = ๐’โ‚‚[:,kron_sโบ_sโบ] - - SSSstates, converged = calculate_second_order_stochastic_steady_state(Val(:newton), ๐’โ‚, ๐’โ‚‚, collect(SSSstates), ๐“‚) # , timer = timer) - - if !converged - if opts.verbose println("SSS not found") end - return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) + if caching + ๐“‚.caches.second_order_stochastic_steady_state = result + ๐“‚.caches.valid_for.second_order_stochastic_steady_state = Float64.(parameters) + end + + return result, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ +end + +function calculate_stochastic_steady_state(::Val{:pruned_second_order}, + parameters::Vector{M}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true) where M + # Cache hit: return cached pruned SSS if valid for current parameters + if caching && M === Float64 && !isempty(parameters) && + cache_valid_for_parameters(๐“‚.caches.valid_for.pruned_second_order_stochastic_steady_state, parameters) + cached_sss = ๐“‚.caches.pruned_second_order_stochastic_steady_state::Vector{M} + if !isempty(cached_sss) + T = ๐“‚.constants.post_model_macro + SS_and_pars = ๐“‚.caches.non_stochastic_steady_state::Vector{M} + โˆ‡โ‚ = Matrix(๐“‚.caches.jacobian)::Matrix{M} + โˆ‡โ‚‚ = sparse(๐“‚.caches.hessian)::SparseMatrixCSC{M, Int} + ๐’โ‚_raw = Matrix(๐“‚.caches.first_order_solution_matrix)::Matrix{M} + ๐’โ‚ = [๐’โ‚_raw[:,1:T.nPast_not_future_and_mixed] zeros(M, T.nVars) ๐’โ‚_raw[:,T.nPast_not_future_and_mixed+1:end]] + ๐’โ‚‚ = sparse(๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} + return cached_sss, true, SS_and_pars, zero(M), โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ end + end - state = A * SSSstates + Bฬ‚ * โ„’.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 - # state, converged = second_order_stochastic_steady_state_iterative_solution([sparsevec(๐’โ‚); vec(๐’โ‚‚)]; dims = [size(๐’โ‚); size(๐’โ‚‚)], ๐“‚ = ๐“‚) + common = _prepare_stochastic_steady_state_base_terms(parameters, ๐“‚, opts = opts, estimation = estimation, caching = caching) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + if caching ๐“‚.caches.pruned_second_order_stochastic_steady_state = all_SS end + return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0) end - # end # timeit_debug + # Expand compressed ๐’โ‚‚_raw to full + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} - # all_variables = sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.aux,๐“‚.constants.post_model_macro.exo_present)) + state = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * SSSstates + + ๐’โ‚‚ * โ„’.kron(sparse([zeros(๐“‚.constants.post_model_macro.nPast_not_future_and_mixed); 1; zeros(๐“‚.constants.post_model_macro.nExo)]), sparse([zeros(๐“‚.constants.post_model_macro.nPast_not_future_and_mixed); 1; zeros(๐“‚.constants.post_model_macro.nExo)])) / 2 - # all_variables[indexin(๐“‚.constants.post_model_macro.aux,all_variables)] = map(x -> Symbol(replace(string(x), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => "")), ๐“‚.constants.post_model_macro.aux) - - # NSSS_labels = [sort(union(๐“‚.constants.post_model_macro.exo_present,๐“‚.constants.post_model_macro.var))...,๐“‚.calibration_equations_parameters...] - - # all_SS = [SS_and_pars[indexin([s],NSSS_labels)...] for s in all_variables] - # we need all variables for the stochastic steady state because even leads and lags have different SSS then the non-lead-lag ones (contrary to the no stochastic steady state) and we cannot recover them otherwise + result = all_SS + Vector{M}(state) - # Ensure state is a Vector{M} for type stability - state_vec = Vector{M}(state) - - return all_SS + state_vec, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ + if caching + ๐“‚.caches.pruned_second_order_stochastic_steady_state = result + ๐“‚.caches.valid_for.pruned_second_order_stochastic_steady_state = Float64.(parameters) + end + + return result, true, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ end -function calculate_second_order_stochastic_steady_state(::Val{:newton}, - ๐’โ‚::Matrix{R}, - ๐’โ‚‚::AbstractSparseMatrix{R}, - x::Vector{R}, - ๐“‚::โ„ณ; - # timer::TimerOutput = TimerOutput(), - tol::AbstractFloat = 1e-14) where R <: AbstractFloat +function solve_stochastic_steady_state_newton(::Val{:second_order}, + ๐’โ‚::Matrix{R}, + ๐’โ‚‚::AbstractSparseMatrix{R}, + x::Vector{R}, + ๐“‚::โ„ณ; + tol::AbstractFloat = 1e-14) where R <: AbstractFloat # @timeit_debug timer "Setup matrices" begin # Get cached computational constants @@ -6469,7 +5646,7 @@ function calculate_second_order_stochastic_steady_state(::Val{:newton}, T = constants.post_model_macro s_in_sโบ = so.s_in_sโบ s_in_s = so.s_in_s - I_nPast = ๐“‚.workspaces.qme.I_nPast + I_nPast = T.I_nPast kron_sโบ_sโบ = so.kron_sโบ_sโบ @@ -6482,12 +5659,18 @@ function calculate_second_order_stochastic_steady_state(::Val{:newton}, max_iters = 100 # SSS .= ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 + # Pre-allocate augmented state vector [x; 1] + x_aug = Vector{R}(undef, length(x) + 1) + x_aug[end] = one(R) + # end # timeit_debug # @timeit_debug timer "Iterations" begin for i in 1:max_iters - โˆ‚x = (A + B * โ„’.kron(vcat(x,1), I_nPast) - I_nPast) + copyto!(x_aug, 1, x, 1, length(x)) + + โˆ‚x = (A + B * โ„’.kron(x_aug, I_nPast) - I_nPast) โˆ‚xฬ‚ = โ„’.lu!(โˆ‚x, check = false) @@ -6495,7 +5678,7 @@ function calculate_second_order_stochastic_steady_state(::Val{:newton}, return x, false end - xฬ‚ = A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + xฬ‚ = A * x + Bฬ‚ * โ„’.kron(x_aug, x_aug) / 2 ฮ”x = โˆ‚xฬ‚ \ (xฬ‚ - x) @@ -6509,91 +5692,70 @@ function calculate_second_order_stochastic_steady_state(::Val{:newton}, # end # timeit_debug - return x, isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2, x, rtol = tol) -end - - - - - -function calculate_third_order_stochastic_steady_state( parameters::Vector{M}, - ๐“‚::โ„ณ; - opts::CalculationOptions = merge_calculation_options(), - pruning::Bool = false, - estimation::Bool = false)where M <: Real - # timer::TimerOutput = TimerOutput(), - # tol::AbstractFloat = 1e-12) - # Initialize constants at entry point - constants = initialise_constants!(๐“‚) - T = constants.post_model_macro - - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts, estimation = estimation) # , timer = timer) - - if solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error) - if opts.verbose println("NSSS not found") end - return zeros(M, T.nVars), false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) - end - - ms = @ignore_derivatives ensure_model_structure_constants!(constants, ๐“‚.equations.calibration_parameters) - all_SS = expand_steady_state(SS_and_pars, ms) + copyto!(x_aug, 1, x, 1, length(x)) + return x, isapprox(A * x + Bฬ‚ * โ„’.kron(x_aug, x_aug) / 2, x, rtol = tol) +end - โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - qme_ws = @ignore_derivatives ensure_qme_workspace!(๐“‚) - sylv_ws = @ignore_derivatives ensure_sylvester_1st_order_workspace!(๐“‚) - - ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, - constants, - qme_ws, - sylv_ws; - opts = opts, - initial_guess = ๐“‚.caches.qme_solution) - - if solved ๐“‚.caches.qme_solution = qme_sol end - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) - if !solved - if opts.verbose println("1st order solution not found") end - return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) - end - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ - ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; - initial_guess = ๐“‚.caches.second_order_solution, - # timer = timer, - opts = opts) +function calculate_stochastic_steady_state(::Val{:third_order}, + parameters::Vector{M}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true) where M <: Real + # Cache hit: return cached SSS if valid for current parameters + if caching && M === Float64 && !isempty(parameters) && + cache_valid_for_parameters(๐“‚.caches.valid_for.third_order_stochastic_steady_state, parameters) + cached_sss = ๐“‚.caches.third_order_stochastic_steady_state::Vector{M} + if !isempty(cached_sss) + T = ๐“‚.constants.post_model_macro + SS_and_pars = ๐“‚.caches.non_stochastic_steady_state::Vector{M} + โˆ‡โ‚ = Matrix(๐“‚.caches.jacobian)::Matrix{M} + โˆ‡โ‚‚ = sparse(๐“‚.caches.hessian)::SparseMatrixCSC{M, Int} + โˆ‡โ‚ƒ = sparse(๐“‚.caches.third_order_derivatives)::SparseMatrixCSC{M, Int} + ๐’โ‚_raw = Matrix(๐“‚.caches.first_order_solution_matrix)::Matrix{M} + ๐’โ‚ = [๐’โ‚_raw[:,1:T.nPast_not_future_and_mixed] zeros(M, T.nVars) ๐’โ‚_raw[:,T.nPast_not_future_and_mixed+1:end]] + ๐’โ‚‚ = sparse(๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} + ๐’ฬ‚โ‚ƒ = sparse(๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ)::SparseMatrixCSC{M, Int} + return cached_sss, true, SS_and_pars, zero(M), โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’ฬ‚โ‚ƒ + end + end - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) + common = _prepare_stochastic_steady_state_base_terms(parameters, ๐“‚, opts = opts, estimation = estimation, caching = caching) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common - if !solved2 - if opts.verbose println("2nd order solution not found") end + if !ok + if caching ๐“‚.caches.third_order_stochastic_steady_state = all_SS end return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) end - - if eltype(๐’โ‚‚) == Float64 && solved2 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end - ๐’โ‚‚ = sparse(๐’โ‚‚ * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} + # Expand compressed ๐’โ‚‚_raw to full + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} - โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives) #, timer = timer)# * ๐“‚.constants.third_order.๐”โˆ‡โ‚ƒ - - ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, - ๐“‚.constants, - ๐“‚.workspaces; - initial_guess = ๐“‚.caches.third_order_solution, - # timer = timer, - opts = opts) + โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces, caching = caching) + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + ๐’โ‚_raw = [๐’โ‚[:, 1:nPast] ๐’โ‚[:, nPast+2:end]] + + ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚_raw, ๐’โ‚‚_raw, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, + parameter_values = parameters, + caching = caching) - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved3, estimation = estimation, order = 3) + update_perturbation_counter!(๐“‚.counters, solved3, estimation = estimation, order = 3) if !solved3 if opts.verbose println("3rd order solution not found") end + if caching ๐“‚.caches.third_order_stochastic_steady_state = all_SS end return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) end - if eltype(๐’โ‚ƒ) == Float64 && solved3 ๐“‚.caches.third_order_solution = ๐’โ‚ƒ end - if length(๐“‚.workspaces.third_order.Sฬ‚) == 0 || !(eltype(๐’โ‚ƒ) == eltype(๐“‚.workspaces.third_order.Sฬ‚)) ๐“‚.workspaces.third_order.Sฬ‚ = ๐’โ‚ƒ * ๐“‚.constants.third_order.๐”โ‚ƒ else @@ -6601,88 +5763,127 @@ function calculate_third_order_stochastic_steady_state( parameters::Vector{M}, end Sฬ‚ = ๐“‚.workspaces.third_order.Sฬ‚ - ๐’โ‚ƒฬ‚ = sparse_preallocated!(Sฬ‚, โ„‚ = ๐“‚.workspaces.third_order)::SparseMatrixCSC{M, Int} - - # ๐’โ‚ƒ *= ๐“‚.constants.third_order.๐”โ‚ƒ - # ๐’โ‚ƒ = sparse_preallocated!(๐’โ‚ƒ, โ„‚ = ๐“‚.workspaces.third_order) - - # ๐’โ‚ƒ = sparse(Sฬ‚) # * ๐“‚.constants.third_order.๐”โ‚ƒ) - ๐’โ‚ = [๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) ๐’โ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] - - aug_stateโ‚ = sparse([zeros(๐“‚.constants.post_model_macro.nPast_not_future_and_mixed); 1; zeros(๐“‚.constants.post_model_macro.nExo)]) - - I_nPast = qme_ws.I_nPast + so = ๐“‚.constants.second_order + kron_sโบ_sโบ = so.kron_sโบ_sโบ + kron_sโบ_sโบ_sโบ = so.kron_sโบ_sโบ_sโบ - tmp = (I_nPast - ๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx, 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed]) + A = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] + Bฬ‚ = ๐’โ‚‚[:,kron_sโบ_sโบ] + Cฬ‚ = ๐’โ‚ƒฬ‚[:,kron_sโบ_sโบ_sโบ] - tmpฬ„ = @ignore_derivatives โ„’.lu(tmp, check = false) + SSSstates, converged = solve_stochastic_steady_state_newton(Val(:third_order), ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚, collect(SSSstates), ๐“‚) - if !โ„’.issuccess(tmpฬ„) + if !converged if opts.verbose println("SSS not found") end + if caching ๐“‚.caches.third_order_stochastic_steady_state = all_SS end return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) end - SSSstates = collect(tmp \ (๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2)[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]) + state = A * SSSstates + Bฬ‚ * โ„’.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 + Cฬ‚ * โ„’.kron(vcat(SSSstates,1), โ„’.kron(vcat(SSSstates,1), vcat(SSSstates,1))) / 6 - if pruning - state = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * SSSstates + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2 - converged = true - else - # Get cached computational constants - so = ๐“‚.constants.second_order - s_in_sโบ = so.s_in_sโบ - kron_sโบ_sโบ = so.kron_sโบ_sโบ - - kron_sโบ_sโบ_sโบ = so.kron_sโบ_sโบ_sโบ - - A = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] - Bฬ‚ = ๐’โ‚‚[:,kron_sโบ_sโบ] - Cฬ‚ = ๐’โ‚ƒฬ‚[:,kron_sโบ_sโบ_sโบ] - - SSSstates, converged = calculate_third_order_stochastic_steady_state(Val(:newton), ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚, SSSstates, ๐“‚) - - if !converged - if opts.verbose println("SSS not found") end - return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) + + result = all_SS + Vector{M}(state) + + if caching + ๐“‚.caches.third_order_stochastic_steady_state = result + ๐“‚.caches.valid_for.third_order_stochastic_steady_state = Float64.(parameters) + end + + return result, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚ +end + +function calculate_stochastic_steady_state(::Val{:pruned_third_order}, + parameters::Vector{M}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true) where M <: Real + # Cache hit: return cached pruned SSS if valid for current parameters + if caching && M === Float64 && !isempty(parameters) && + cache_valid_for_parameters(๐“‚.caches.valid_for.pruned_third_order_stochastic_steady_state, parameters) + cached_sss = ๐“‚.caches.pruned_third_order_stochastic_steady_state::Vector{M} + if !isempty(cached_sss) + T = ๐“‚.constants.post_model_macro + SS_and_pars = ๐“‚.caches.non_stochastic_steady_state::Vector{M} + โˆ‡โ‚ = Matrix(๐“‚.caches.jacobian)::Matrix{M} + โˆ‡โ‚‚ = sparse(๐“‚.caches.hessian)::SparseMatrixCSC{M, Int} + โˆ‡โ‚ƒ = sparse(๐“‚.caches.third_order_derivatives)::SparseMatrixCSC{M, Int} + ๐’โ‚_raw = Matrix(๐“‚.caches.first_order_solution_matrix)::Matrix{M} + ๐’โ‚ = [๐’โ‚_raw[:,1:T.nPast_not_future_and_mixed] zeros(M, T.nVars) ๐’โ‚_raw[:,T.nPast_not_future_and_mixed+1:end]] + ๐’โ‚‚ = sparse(๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} + ๐’ฬ‚โ‚ƒ = sparse(๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ)::SparseMatrixCSC{M, Int} + return cached_sss, true, SS_and_pars, zero(M), โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’ฬ‚โ‚ƒ end + end - state = A * SSSstates + Bฬ‚ * โ„’.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 + Cฬ‚ * โ„’.kron(vcat(SSSstates,1), โ„’.kron(vcat(SSSstates,1), vcat(SSSstates,1))) / 6 - # state, converged = third_order_stochastic_steady_state_iterative_solution([sparsevec(๐’โ‚); vec(๐’โ‚‚); vec(๐’โ‚ƒ)]; dims = [size(๐’โ‚); size(๐’โ‚‚); size(๐’โ‚ƒ)], ๐“‚ = ๐“‚) - # state, converged = third_order_stochastic_steady_state_iterative_solution_forward([sparsevec(๐’โ‚); vec(๐’โ‚‚); vec(๐’โ‚ƒ)]; dims = [size(๐’โ‚); size(๐’โ‚‚); size(๐’โ‚ƒ)], ๐“‚ = ๐“‚) + common = _prepare_stochastic_steady_state_base_terms(parameters, ๐“‚, opts = opts, estimation = estimation, caching = caching) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + if caching ๐“‚.caches.pruned_third_order_stochastic_steady_state = all_SS end + return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) end - # all_variables = sort(union(๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.aux,๐“‚.constants.post_model_macro.exo_present)) + # Expand compressed ๐’โ‚‚_raw to full + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{M, Int} - # all_variables[indexin(๐“‚.constants.post_model_macro.aux,all_variables)] = map(x -> Symbol(replace(string(x), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => "")), ๐“‚.constants.post_model_macro.aux) - - # NSSS_labels = [sort(union(๐“‚.constants.post_model_macro.exo_present,๐“‚.constants.post_model_macro.var))...,๐“‚.calibration_equations_parameters...] - - # all_SS = [SS_and_pars[indexin([s],NSSS_labels)...] for s in all_variables] - # we need all variables for the stochastic steady state because even leads and lags have different SSS then the non-lead-lag ones (contrary to the no stochastic steady state) and we cannot recover them otherwise + โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces, caching = caching) + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + ๐’โ‚_raw = [๐’โ‚[:, 1:nPast] ๐’โ‚[:, nPast+2:end]] + + ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚_raw, ๐’โ‚‚_raw, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, parameter_values = parameters, caching = caching) + + update_perturbation_counter!(๐“‚.counters, solved3, estimation = estimation, order = 3) + + if !solved3 + if opts.verbose println("3rd order solution not found") end + if caching ๐“‚.caches.pruned_third_order_stochastic_steady_state = all_SS end + return all_SS, false, SS_and_pars, solution_error, zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0), zeros(M,0,0), spzeros(M,0,0), spzeros(M,0,0) + end + + if length(๐“‚.workspaces.third_order.Sฬ‚) == 0 || !(eltype(๐’โ‚ƒ) == eltype(๐“‚.workspaces.third_order.Sฬ‚)) + ๐“‚.workspaces.third_order.Sฬ‚ = ๐’โ‚ƒ * ๐“‚.constants.third_order.๐”โ‚ƒ + else + mul_reverse_AD!(๐“‚.workspaces.third_order.Sฬ‚, ๐’โ‚ƒ, ๐“‚.constants.third_order.๐”โ‚ƒ) + end + + Sฬ‚ = ๐“‚.workspaces.third_order.Sฬ‚ + ๐’โ‚ƒฬ‚ = sparse_preallocated!(Sฬ‚, โ„‚ = ๐“‚.workspaces.third_order)::SparseMatrixCSC{M, Int} - # Ensure state is a Vector{M} for type stability - state_vec = Vector{M}(state) + aug_stateโ‚ = sparse([zeros(๐“‚.constants.post_model_macro.nPast_not_future_and_mixed); 1; zeros(๐“‚.constants.post_model_macro.nExo)]) + state = ๐’โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * SSSstates + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2 + + result = all_SS + Vector{M}(state) + + if caching + ๐“‚.caches.pruned_third_order_stochastic_steady_state = result + ๐“‚.caches.valid_for.pruned_third_order_stochastic_steady_state = Float64.(parameters) + end - return all_SS + state_vec, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚ + return result, true, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚ end -function calculate_third_order_stochastic_steady_state(::Val{:newton}, - ๐’โ‚::Matrix{Float64}, - ๐’โ‚‚::AbstractSparseMatrix{Float64}, - ๐’โ‚ƒ::AbstractSparseMatrix{Float64}, - x::Vector{Float64}, - ๐“‚::โ„ณ; - # timer::TimerOutput = TimerOutput(), - tol::AbstractFloat = 1e-14) +function solve_stochastic_steady_state_newton(::Val{:third_order}, + ๐’โ‚::Matrix{Float64}, + ๐’โ‚‚::AbstractSparseMatrix{Float64}, + ๐’โ‚ƒ::AbstractSparseMatrix{Float64}, + x::Vector{Float64}, + ๐“‚::โ„ณ; + tol::AbstractFloat = 1e-14) # Get cached computational constants - so = ensure_computational_constants!(๐“‚) + so = ensure_computational_constants!(๐“‚.constants) T = ๐“‚.constants.post_model_macro s_in_sโบ = so.s_in_sโบ s_in_s = so.s_in_s - I_nPast = ๐“‚.workspaces.qme.I_nPast + I_nPast = T.I_nPast kron_sโบ_sโบ = so.kron_sโบ_sโบ @@ -6700,8 +5901,17 @@ function calculate_third_order_stochastic_steady_state(::Val{:newton}, max_iters = 100 # SSS .= ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 + + # Pre-allocate augmented state vector [x; 1] + x_aug = Vector{Float64}(undef, length(x) + 1) + x_aug[end] = 1.0 + for i in 1:max_iters - โˆ‚x = (A + B * โ„’.kron(vcat(x,1), I_nPast) + C * โ„’.kron(โ„’.kron(vcat(x,1), vcat(x,1)), I_nPast) / 2 - I_nPast) + copyto!(x_aug, 1, x, 1, length(x)) + kron_x_aug = โ„’.kron(x_aug, x_aug) + kron_x_kron = โ„’.kron(x_aug, kron_x_aug) + + โˆ‚x = (A + B * โ„’.kron(x_aug, I_nPast) + C * โ„’.kron(kron_x_aug, I_nPast) / 2 - I_nPast) โˆ‚xฬ‚ = โ„’.lu!(โˆ‚x, check = false) @@ -6709,9 +5919,9 @@ function calculate_third_order_stochastic_steady_state(::Val{:newton}, return x, false end - ฮ”x = โˆ‚xฬ‚ \ (A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + Cฬ‚ * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1))) / 6 - x) + ฮ”x = โˆ‚xฬ‚ \ (A * x + Bฬ‚ * kron_x_aug / 2 + Cฬ‚ * kron_x_kron / 6 - x) - if i > 5 && isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + Cฬ‚ * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1))) / 6, x, rtol = tol) + if i > 5 && isapprox(A * x + Bฬ‚ * kron_x_aug / 2 + Cฬ‚ * kron_x_kron / 6, x, rtol = tol) break end @@ -6719,12 +5929,25 @@ function calculate_third_order_stochastic_steady_state(::Val{:newton}, โ„’.axpy!(-1, ฮ”x, x) end - return x, isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + Cฬ‚ * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1))) / 6, x, rtol = tol) + copyto!(x_aug, 1, x, 1, length(x)) + kron_x_aug = โ„’.kron(x_aug, x_aug) + kron_x_kron = โ„’.kron(x_aug, kron_x_aug) + return x, isapprox(A * x + Bฬ‚ * kron_x_aug / 2 + Cฬ‚ * kron_x_kron / 6, x, rtol = tol) end -function set_up_steady_state_solver!(๐“‚::โ„ณ; verbose::Bool, silent::Bool, avoid_solve::Bool = false, symbolic::Bool = false) +function steady_state_symbolic_mode_flags(ss_symbolic_mode::Symbol, precompile::Bool = false) + precompile && (ss_symbolic_mode = :none) + ss_symbolic_mode == :none && return true, false + ss_symbolic_mode == :single_equation && return false, false + ss_symbolic_mode == :full && return false, true + error("Invalid ss_symbolic_mode $(ss_symbolic_mode). Expected :none, :single_equation, or :full.") +end + +function set_up_steady_state_solver!(๐“‚::โ„ณ; verbose::Bool, silent::Bool, ss_symbolic_mode::Symbol = :single_equation) + avoid_solve, symbolic_enabled = steady_state_symbolic_mode_flags(ss_symbolic_mode, ๐“‚.constants.post_parameters_macro.precompile) + if !๐“‚.constants.post_parameters_macro.precompile start_time = time() @@ -6742,7 +5965,7 @@ function set_up_steady_state_solver!(๐“‚::โ„ณ; verbose::Bool, silent::Bool, avo write_ss_check_function!(๐“‚) - write_steady_state_solver_function!(๐“‚, symbolic, symbolics, verbose = verbose, avoid_solve = avoid_solve) + write_steady_state_solver_function!(๐“‚, symbolic_enabled, symbolics, verbose = verbose, avoid_solve = avoid_solve) ๐“‚.equations.obc_violation = write_obc_violation_equations(๐“‚) @@ -6756,7 +5979,7 @@ function set_up_steady_state_solver!(๐“‚::โ„ณ; verbose::Bool, silent::Bool, avo write_ss_check_function!(๐“‚) - write_steady_state_solver_function!(๐“‚, verbose = verbose) + write_steady_state_solver_function!(๐“‚, false, nothing, verbose = verbose, avoid_solve = avoid_solve) if !silent println(round(time() - start_time, digits = 3), " seconds") end end @@ -6791,9 +6014,12 @@ function solve!(๐“‚::โ„ณ; if ๐“‚.functions.functions_written && isnothing(๐“‚.functions.NSSS_custom) && - !(๐“‚.functions.NSSS_solve isa RuntimeGeneratedFunctions.RuntimeGeneratedFunction) + ๐“‚.constants.nsss_solver.n_steps == 0 - set_up_steady_state_solver!(๐“‚, verbose = opts.verbose, silent = silent) + set_up_steady_state_solver!(๐“‚, + verbose = opts.verbose, + silent = silent, + ss_symbolic_mode = ๐“‚.constants.post_parameters_macro.ss_symbolic_mode) end if !๐“‚.functions.functions_written @@ -6801,9 +6027,16 @@ function solve!(๐“‚::โ„ณ; perturbation_order = 1 - set_up_steady_state_solver!(๐“‚, verbose = verbose, silent = silent, avoid_solve = false) + set_up_steady_state_solver!(๐“‚, + verbose = verbose, + silent = silent, + ss_symbolic_mode = ๐“‚.constants.post_parameters_macro.ss_symbolic_mode) - SS_and_pars, solution_error, found_solution = solve_steady_state!(๐“‚, opts, :ESCH, 120.0, silent = silent) + SS_and_pars, solution_error, found_solution = solve_steady_state!(๐“‚, + opts, + ๐“‚.constants.post_parameters_macro.ss_solver_parameters_algorithm, + ๐“‚.constants.post_parameters_macro.ss_solver_parameters_maxtime, + silent = silent) write_symbolic_derivatives!(๐“‚; perturbation_order = perturbation_order, silent = silent) @@ -6831,233 +6064,48 @@ function solve!(๐“‚::โ„ณ; end if dynamics - obc_not_solved = isnothing(๐“‚.functions.first_order_state_update_obc(zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nExo))) - if ((:first_order == algorithm) && (๐“‚.caches.outdated.first_order_solution || (obc && obc_not_solved))) || - ((:second_order == algorithm) && (๐“‚.caches.outdated.second_order_solution || (obc && obc_not_solved))) || - ((:pruned_second_order == algorithm) && (๐“‚.caches.outdated.pruned_second_order_solution || (obc && obc_not_solved))) || - ((:third_order == algorithm) && (๐“‚.caches.outdated.third_order_solution || (obc && obc_not_solved))) || - ((:pruned_third_order == algorithm) && (๐“‚.caches.outdated.pruned_third_order_solution || (obc && obc_not_solved))) - - # @timeit_debug timer "Solve for NSSS (if necessary)" begin - - SS_and_pars, (solution_error, iters) = ๐“‚.caches.outdated.non_stochastic_steady_state ? get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) : (๐“‚.caches.non_stochastic_steady_state, (eps(), 0)) - - # end # timeit_debug - - @assert solution_error < opts.tol.NSSS_acceptance_tol "Could not find non-stochastic steady state." - - # @timeit_debug timer "Calculate Jacobian" begin + if algorithm == :first_order + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) - โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - # end # timeit_debug + @assert solution_error < opts.tol.nsss.acceptance_tol "Could not find non-stochastic steady state." - # @timeit_debug timer "Calculate first order solution" begin + โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) - qme_ws = @ignore_derivatives ensure_qme_workspace!(๐“‚) - sylv_ws = @ignore_derivatives ensure_sylvester_1st_order_workspace!(๐“‚) - Sโ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; opts = opts, - initial_guess = ๐“‚.caches.qme_solution) - - if solved ๐“‚.caches.qme_solution = qme_sol end - - update_perturbation_counter!(๐“‚.counters, solved, order = 1) + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = ๐“‚.parameter_values) - # end # timeit_debug + update_perturbation_counter!(๐“‚.counters, solved, order = 1) @assert solved "Could not find stable first order solution." - state_updateโ‚ = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx] - shock] - return Sโ‚ * aug_state # return statement needed for forwarddiff to work - end - - if obc - write_parameters_input!(๐“‚, :activeแต’แต‡แถœshocks => 1, verbose = false) - - โˆ‡ฬ‚โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - Sฬ‚โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡ฬ‚โ‚, - constants, - qme_ws, - sylv_ws; - opts = opts, - initial_guess = ๐“‚.caches.qme_solution) - if solved ๐“‚.caches.qme_solution = qme_sol end - - update_perturbation_counter!(๐“‚.counters, solved, order = 1) - - write_parameters_input!(๐“‚, :activeแต’แต‡แถœshocks => 0, verbose = false) - - state_updateโ‚ฬ‚ = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx] - shock] - return Sฬ‚โ‚ * aug_state # you need a return statement for forwarddiff to work - end - else - state_updateโ‚ฬ‚ = (x,y)->nothing - end - - ๐“‚.caches.first_order_solution_matrix = Sโ‚ - ๐“‚.functions.first_order_state_update = state_updateโ‚ - ๐“‚.functions.first_order_state_update_obc = state_updateโ‚ฬ‚ - ๐“‚.caches.outdated.first_order_solution = false - - ๐“‚.caches.non_stochastic_steady_state = SS_and_pars - ๐“‚.caches.outdated.non_stochastic_steady_state = solution_error > opts.tol.NSSS_acceptance_tol - end - - obc_not_solved = isnothing(๐“‚.functions.second_order_state_update_obc(zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nExo))) - if ((:second_order == algorithm) && (๐“‚.caches.outdated.second_order_solution || (obc && obc_not_solved))) || - ((:third_order == algorithm) && (๐“‚.caches.outdated.third_order_solution || (obc && obc_not_solved))) - - - stochastic_steady_state, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_second_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, opts = opts) # , timer = timer) - - if !converged @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - - state_updateโ‚‚ = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx] - 1 - shock] - return ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 - end - - if obc - Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) Sฬ‚โ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] - - state_updateโ‚‚ฬ‚ = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx] - 1 - shock] - return Sฬ‚โ‚ฬ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 - end - else - state_updateโ‚‚ฬ‚ = (x,y)->nothing - end - - ๐“‚.caches.second_order_stochastic_steady_state = stochastic_steady_state - ๐“‚.functions.second_order_state_update = state_updateโ‚‚ - ๐“‚.functions.second_order_state_update_obc = state_updateโ‚‚ฬ‚ - - ๐“‚.caches.outdated.second_order_solution = false - end - - obc_not_solved = isnothing(๐“‚.functions.pruned_second_order_state_update_obc([zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nVars)], zeros(๐“‚.constants.post_model_macro.nExo))) - if ((:pruned_second_order == algorithm) && (๐“‚.caches.outdated.pruned_second_order_solution || (obc && obc_not_solved))) || - ((:pruned_third_order == algorithm) && (๐“‚.caches.outdated.pruned_third_order_solution || (obc && obc_not_solved))) - - stochastic_steady_state, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_second_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, opts = opts, pruning = true) # , timer = timer) - - if !converged @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - - state_updateโ‚‚ = function(pruned_states::Vector{Vector{T}}, shock::Vector{S}) where {T,S} - aug_stateโ‚ = [pruned_states[1][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 1; shock] - aug_stateโ‚‚ = [pruned_states[2][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] - - return [๐’โ‚ * aug_stateโ‚, ๐’โ‚ * aug_stateโ‚‚ + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2] # strictly following Andreasen et al. (2018) - end + elseif algorithm == :second_order + sss_result = calculate_stochastic_steady_state(Val(:second_order), ๐“‚.parameter_values, ๐“‚, opts = opts) + if !sss_result[2] @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - if obc - Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) Sฬ‚โ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] - - state_updateโ‚‚ฬ‚ = function(pruned_states::Vector{Vector{T}}, shock::Vector{S}) where {T,S} - aug_stateโ‚ = [pruned_states[1][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 1; shock] - aug_stateโ‚‚ = [pruned_states[2][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] - - return [Sฬ‚โ‚ฬ‚ * aug_stateโ‚, Sฬ‚โ‚ฬ‚ * aug_stateโ‚‚ + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2] # strictly following Andreasen et al. (2018) - end - else - state_updateโ‚‚ฬ‚ = (x,y)->nothing - end + elseif algorithm == :pruned_second_order + sss_result = calculate_stochastic_steady_state(Val(:pruned_second_order), ๐“‚.parameter_values, ๐“‚, opts = opts) + if !sss_result[2] @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - ๐“‚.caches.pruned_second_order_stochastic_steady_state = stochastic_steady_state - ๐“‚.functions.pruned_second_order_state_update = state_updateโ‚‚ - ๐“‚.functions.pruned_second_order_state_update_obc = state_updateโ‚‚ฬ‚ + elseif algorithm == :third_order + calculate_stochastic_steady_state(Val(:second_order), ๐“‚.parameter_values, ๐“‚, opts = opts) + sss_result = calculate_stochastic_steady_state(Val(:third_order), ๐“‚.parameter_values, ๐“‚, opts = opts) + if !sss_result[2] @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - ๐“‚.caches.outdated.pruned_second_order_solution = false + elseif algorithm == :pruned_third_order + calculate_stochastic_steady_state(Val(:pruned_second_order), ๐“‚.parameter_values, ๐“‚, opts = opts) + sss_result = calculate_stochastic_steady_state(Val(:pruned_third_order), ๐“‚.parameter_values, ๐“‚, opts = opts) + if !sss_result[2] @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end end - - obc_not_solved = isnothing(๐“‚.functions.third_order_state_update_obc(zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nExo))) - if ((:third_order == algorithm) && (๐“‚.caches.outdated.third_order_solution || (obc && obc_not_solved))) - stochastic_steady_state, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_third_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, opts = opts) - - if !converged @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - - state_updateโ‚ƒ = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx] - 1 - shock] - return ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 - end - - if obc - Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) Sฬ‚โ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] - - state_updateโ‚ƒฬ‚ = function(state::Vector{T}, shock::Vector{S}) where {T,S} - aug_state = [state[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx] - 1 - shock] - return Sฬ‚โ‚ฬ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 - end - else - state_updateโ‚ƒฬ‚ = (x,y)->nothing - end - ๐“‚.caches.third_order_stochastic_steady_state = stochastic_steady_state - ๐“‚.functions.third_order_state_update = state_updateโ‚ƒ - ๐“‚.functions.third_order_state_update_obc = state_updateโ‚ƒฬ‚ - - ๐“‚.caches.outdated.third_order_solution = false + if obc + calculate_first_order_obc_solution!(๐“‚, constants, opts) end - obc_not_solved = isnothing(๐“‚.functions.pruned_third_order_state_update_obc([zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nVars)], zeros(๐“‚.constants.post_model_macro.nExo))) - if ((:pruned_third_order == algorithm) && (๐“‚.caches.outdated.pruned_third_order_solution || (obc && obc_not_solved))) - - stochastic_steady_state, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_third_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, opts = opts, pruning = true) - - if !converged @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end - - state_updateโ‚ƒ = function(pruned_states::Vector{Vector{T}}, shock::Vector{S}) where {T,S} - aug_stateโ‚ = [pruned_states[1][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 1; shock] - aug_stateโ‚ฬ‚ = [pruned_states[1][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; shock] - aug_stateโ‚‚ = [pruned_states[2][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] - aug_stateโ‚ƒ = [pruned_states[3][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] - - kron_aug_stateโ‚ = โ„’.kron(aug_stateโ‚, aug_stateโ‚) - - return [๐’โ‚ * aug_stateโ‚, ๐’โ‚ * aug_stateโ‚‚ + ๐’โ‚‚ * kron_aug_stateโ‚ / 2, ๐’โ‚ * aug_stateโ‚ƒ + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚ฬ‚, aug_stateโ‚‚) + ๐’โ‚ƒ * โ„’.kron(kron_aug_stateโ‚,aug_stateโ‚) / 6] - end - - if obc - Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) Sฬ‚โ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] - - state_updateโ‚ƒฬ‚ = function(pruned_states::Vector{Vector{T}}, shock::Vector{S}) where {T,S} - aug_stateโ‚ = [pruned_states[1][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 1; shock] - aug_stateโ‚ฬ‚ = [pruned_states[1][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; shock] - aug_stateโ‚‚ = [pruned_states[2][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] - aug_stateโ‚ƒ = [pruned_states[3][๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] - - kron_aug_stateโ‚ = โ„’.kron(aug_stateโ‚, aug_stateโ‚) - - return [Sฬ‚โ‚ฬ‚ * aug_stateโ‚, Sฬ‚โ‚ฬ‚ * aug_stateโ‚‚ + ๐’โ‚‚ * kron_aug_stateโ‚ / 2, Sฬ‚โ‚ฬ‚ * aug_stateโ‚ƒ + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚ฬ‚, aug_stateโ‚‚) + ๐’โ‚ƒ * โ„’.kron(kron_aug_stateโ‚,aug_stateโ‚) / 6] # strictly following Andreasen et al. (2018) - end - else - state_updateโ‚ƒฬ‚ = (x,y)->nothing - end - - ๐“‚.caches.pruned_third_order_stochastic_steady_state = stochastic_steady_state - ๐“‚.functions.pruned_third_order_state_update = state_updateโ‚ƒ - ๐“‚.functions.pruned_third_order_state_update_obc = state_updateโ‚ƒฬ‚ - - ๐“‚.caches.outdated.pruned_third_order_solution = false - end end return nothing @@ -7072,6 +6120,7 @@ function create_second_order_auxiliary_matrices(constants::constants) # Indices and number of variables nโ‚‹ = T.nPast_not_future_and_mixed + n = T.nVars nโ‚‘ = T.nExo # setup compression matrices for hessian matrix @@ -7085,17 +6134,44 @@ function create_second_order_auxiliary_matrices(constants::constants) redu = sparsevec(nโ‚‘โ‚‹ - nโ‚‘ + 1:nโ‚‘โ‚‹, 1) redu_idxs = findnz(โ„’.kron(redu, redu))[1] ๐›” = @views sparse(redu_idxs[Int.(range(1,nโ‚‘^2,nโ‚‘))], fill(nโ‚‹ * (nโ‚‘โ‚‹ + 1) + 1, nโ‚‘), 1, nโ‚‘โ‚‹^2, nโ‚‘โ‚‹^2) - # setup compression matrices for transition matrix colls2 = [nโ‚‘โ‚‹ * (i-1) + k for i in 1:nโ‚‘โ‚‹ for k in 1:i] ๐‚โ‚‚ = sparse(colls2, 1:length(colls2), 1) ๐”โ‚‚ = ๐‚โ‚‚' * sparse([i <= k ? (k - 1) * nโ‚‘โ‚‹ + i : (i - 1) * nโ‚‘โ‚‹ + k for k in 1:nโ‚‘โ‚‹ for i in 1:nโ‚‘โ‚‹], 1:nโ‚‘โ‚‹^2, 1) + # Build symmetrised volatility: ๐›”_sym = ๐›” + P_swap * ๐›” * P_swap + # P_swap is the commutation matrix swapping axes 1 and 2 in nโ‚‘โ‚‹ยฒ space + swap_rows = Vector{Int}(undef, nโ‚‘โ‚‹^2) + swap_cols = Vector{Int}(undef, nโ‚‘โ‚‹^2) + @inbounds for a in 1:nโ‚‘โ‚‹, b in 1:nโ‚‘โ‚‹ + idx = (a - 1) * nโ‚‘โ‚‹ + b + swap_rows[idx] = idx + swap_cols[idx] = (b - 1) * nโ‚‘โ‚‹ + a + end + P_swap = sparse(swap_rows, swap_cols, ones(Int, nโ‚‘โ‚‹^2), nโ‚‘โ‚‹^2, nโ‚‘โ‚‹^2) + ๐›”_sym = ๐›” + P_swap * ๐›” * P_swap + so = constants.second_order so.๐›” = ๐›” + so.๐›”_sym = ๐›”_sym + so.๐›”cโ‚‚ = ๐”โ‚‚ * ๐›” * ๐‚โ‚‚ + so.๐›”๐‚โ‚‚ = ๐›” * ๐‚โ‚‚ so.๐‚โ‚‚ = ๐‚โ‚‚ so.๐”โ‚‚ = ๐”โ‚‚ so.๐”โˆ‡โ‚‚ = ๐”โˆ‡โ‚‚ + so.๐ˆโ‚™โ‚Š = sparse(1:T.nFuture_not_past_and_mixed, T.future_not_past_and_mixed_idx, 1, T.nFuture_not_past_and_mixed, n) + so.๐ˆโ‚™โ‚‹ = sparse(1:T.nPast_not_future_and_mixed, T.past_not_future_and_mixed_idx, 1, T.nPast_not_future_and_mixed, n) + so.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask = Int[] + sigma_row_lookup = falses(size(so.๐›”cโ‚‚, 1)) + @inbounds for r in so.๐›”cโ‚‚.rowval + sigma_row_lookup[r] = true + end + so.๐›”๐‚โ‚‚_nonempty_row_as_kron_colmask = findall(sigma_row_lookup) + # Pre-transposed constants for rrule pullback (computed once) + so.๐›”แต€ = sparse(๐›”') + so.๐‚โ‚‚แต€ = sparse(๐‚โ‚‚') + so.๐”โ‚‚แต€ = sparse(๐”โ‚‚') + so.๐”โˆ‡โ‚‚แต€ = sparse(๐”โˆ‡โ‚‚') return so end @@ -7217,7 +6293,9 @@ function create_third_order_auxiliary_matrices(constants::constants, โˆ‡โ‚ƒ_col_ to.๐ˆโ‚ƒ = ๐ˆโ‚ƒ to.๐‚โˆ‡โ‚ƒ = ๐‚โˆ‡โ‚ƒ to.๐”โˆ‡โ‚ƒ = ๐”โˆ‡โ‚ƒ + to.โˆ‡โ‚ƒ_rowmask = sort!(unique(โˆ‡โ‚ƒ_col_indices)) to.๐ = ๐ + to.๐๐‚โ‚ƒ = ๐ * ๐‚โ‚ƒ to.๐โ‚โ‚— = ๐โ‚โ‚— to.๐โ‚แตฃ = ๐โ‚แตฃ to.๐โ‚โ‚—ฬ‚ = ๐โ‚โ‚—ฬ‚ @@ -7227,6 +6305,16 @@ function create_third_order_auxiliary_matrices(constants::constants, โˆ‡โ‚ƒ_col_ to.๐โ‚แตฃฬƒ = ๐โ‚แตฃฬƒ to.๐โ‚‚แตฃฬƒ = ๐โ‚‚แตฃฬƒ to.๐’๐ = ๐’๐ + # Pre-transposed constants for rrule pullback (computed once) + to.๐‚โ‚ƒแต€ = sparse(๐‚โ‚ƒ') + to.๐”โ‚ƒแต€ = sparse(๐”โ‚ƒ') + to.๐๐‚โ‚ƒแต€ = sparse((to.๐๐‚โ‚ƒ)') + to.๐โ‚โ‚—แต€ = sparse(๐โ‚โ‚—') + to.๐โ‚แตฃแต€ = sparse(๐โ‚แตฃ') + to.๐โ‚โ‚—ฬ„แต€ = sparse(๐โ‚โ‚—ฬ„') + to.๐โ‚‚โ‚—ฬ„แต€ = sparse(๐โ‚‚โ‚—ฬ„') + to.๐โ‚แตฃฬƒแต€ = sparse(๐โ‚แตฃฬƒ') + to.๐โ‚‚แตฃฬƒแต€ = sparse(๐โ‚‚แตฃฬƒ') return to end @@ -7510,6 +6598,25 @@ function take_nth_order_derivatives( local X_col_idx # Column index in the final spX_order_n matrix (1 to X_ncols_n) if output_compressed + # For compressed output, only include entries where variable indices + # are in non-increasing order (v_n <= v_{n-1} <= ... <= v_1). + # This matches the compression rule used for the X-matrix. + # Unsorted tuples represent the same derivative (by symmetry of + # mixed partials) but the compressed column formula maps them to + # WRONG positions, corrupting the Jacobian. + is_compressed_P = true + for k_rule = 1:(n-1) + if var_indices_full[n-k_rule+1] > var_indices_full[n-k_rule] + is_compressed_P = false + break + end + end + + if !is_compressed_P + k_temp_P += 1 + continue + end + # Calculate the compressed column index compressed_col_idx = 0 for k_formula = 1:(n-1) @@ -7662,9 +6769,9 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; end - calib_replacements = Dict{Symbol,Any}() + calib_replacements = Dict{Symbol, Union{Expr, Symbol, Number}}() for (i,x) in enumerate(calib_vars) - replacement = Dict(x => calib_expr[i]) + replacement = Dict{Symbol, Union{Expr, Symbol, Number}}(x => calib_expr[i]) for ii in i+1:length(calib_vars) calib_expr[ii] = replace_symbols(calib_expr[ii], replacement) end @@ -7680,12 +6787,27 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; derivatives = take_nth_order_derivatives(dyn_equations, ๐”™, ๐”“, SS_mapping, nps, nxs) + function prepare_sensitivity_buffer(derivative_sensitivities) + transposed = derivative_sensitivities isa SparseMatrixCSC ? sparse(transpose(derivative_sensitivities)) : permutedims(derivative_sensitivities) + lennz = nnz(transposed) + + if (lennz / length(transposed) > density_threshold) || (length(transposed) < min_length) + return convert(Matrix, transposed), zeros(Float64, size(transposed)), lennz + end + + buffer = similar(transposed, Float64) + buffer.nzval .= 0 + return transposed, buffer, lennz + end + โˆ‡โ‚_dyn = derivatives[1][1] lennz = nnz(โˆ‡โ‚_dyn) - if (lennz / length(โˆ‡โ‚_dyn) > density_threshold) || (length(โˆ‡โ‚_dyn) < min_length) + jacobian_dense_by_heuristic = (lennz / length(โˆ‡โ‚_dyn) > density_threshold) || (length(โˆ‡โ‚_dyn) < min_length) + # Re-enable `jacobian_dense_by_heuristic` directly to restore sparse Jacobian path switching. + if jacobian_dense_by_heuristic derivatives_mat = convert(Matrix, โˆ‡โ‚_dyn) buffer = zeros(Float64, size(โˆ‡โ‚_dyn)) else @@ -7711,18 +6833,7 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; ๐“‚.caches.jacobian = buffer - โˆ‡โ‚_parameters = derivatives[1][2][:,1:nps] - - lennz = nnz(โˆ‡โ‚_parameters) - - if (lennz / length(โˆ‡โ‚_parameters) > density_threshold) || (length(โˆ‡โ‚_parameters) < min_length) - โˆ‡โ‚_parameters_mat = convert(Matrix, โˆ‡โ‚_parameters) - buffer_parameters = zeros(Float64, size(โˆ‡โ‚_parameters)) - else - โˆ‡โ‚_parameters_mat = โˆ‡โ‚_parameters - buffer_parameters = similar(โˆ‡โ‚_parameters, Float64) - buffer_parameters.nzval .= 0 - end + โˆ‡โ‚_parameters_mat, buffer_parameters, lennz = prepare_sensitivity_buffer(derivatives[1][2][:,1:nps]) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -7741,18 +6852,7 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; ๐“‚.caches.jacobian_parameters = buffer_parameters - โˆ‡โ‚_SS_and_pars = derivatives[1][2][:,nps+1:end] - - lennz = nnz(โˆ‡โ‚_SS_and_pars) - - if (lennz / length(โˆ‡โ‚_SS_and_pars) > density_threshold) || (length(โˆ‡โ‚_SS_and_pars) < min_length) - โˆ‡โ‚_SS_and_pars_mat = convert(Matrix, โˆ‡โ‚_SS_and_pars) - buffer_SS_and_pars = zeros(Float64, size(โˆ‡โ‚_SS_and_pars)) - else - โˆ‡โ‚_SS_and_pars_mat = โˆ‡โ‚_SS_and_pars - buffer_SS_and_pars = similar(โˆ‡โ‚_SS_and_pars, Float64) - buffer_SS_and_pars.nzval .= 0 - end + โˆ‡โ‚_SS_and_pars_mat, buffer_SS_and_pars, lennz = prepare_sensitivity_buffer(derivatives[1][2][:,nps+1:end]) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -7826,7 +6926,6 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; # expression_module = @__MODULE__, # expression = Val(false))::Tuple{<:Function, <:Function} - # ๐“‚.caches.โˆ‚equations_โˆ‚parameters = buffer # ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters = func_exprs @@ -7858,19 +6957,19 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; # expression_module = @__MODULE__, # expression = Val(false))::Tuple{<:Function, <:Function} - # ๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars = buffer # ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars = func_exprs # end if max_perturbation_order >= 2 # second order - derivatives = take_nth_order_derivatives(dyn_equations, ๐”™, ๐”“, SS_mapping, nps, nxs; max_perturbation_order = 2, output_compressed = false) + derivatives = take_nth_order_derivatives(dyn_equations, ๐”™, ๐”“, SS_mapping, nps, nxs; max_perturbation_order = 2, output_compressed = true) if ๐“‚.constants.second_order.๐›” == SparseMatrixCSC{Int, Int64}(โ„’.I,0,0) - ๐“‚.constants.second_order = create_second_order_auxiliary_matrices(๐“‚.constants) - โˆ‡โ‚‚_dyn = derivatives[2][1] + ๐“‚.constants.second_order = create_second_order_auxiliary_matrices(๐“‚.constants) + ๐“‚.constants.second_order.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask = findall(@view(โˆ‡โ‚‚_dyn.colptr[1:end-1]) .< @view(โˆ‡โ‚‚_dyn.colptr[2:end])) + lennz = nnz(โˆ‡โ‚‚_dyn) if (lennz / length(โˆ‡โ‚‚_dyn) > density_threshold) || (length(โˆ‡โ‚‚_dyn) < min_length) @@ -7899,18 +6998,7 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; ๐“‚.caches.hessian = buffer - โˆ‡โ‚‚_parameters = derivatives[2][2][:,1:nps] - - lennz = nnz(โˆ‡โ‚‚_parameters) - - if (lennz / length(โˆ‡โ‚‚_parameters) > density_threshold) || (length(โˆ‡โ‚‚_parameters) < min_length) - โˆ‡โ‚‚_parameters_mat = convert(Matrix, โˆ‡โ‚‚_parameters) - buffer_parameters = zeros(Float64, size(โˆ‡โ‚‚_parameters)) - else - โˆ‡โ‚‚_parameters_mat = โˆ‡โ‚‚_parameters - buffer_parameters = similar(โˆ‡โ‚‚_parameters, Float64) - buffer_parameters.nzval .= 0 - end + โˆ‡โ‚‚_parameters_mat, buffer_parameters, lennz = prepare_sensitivity_buffer(derivatives[2][2][:,1:nps]) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -7929,18 +7017,7 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; ๐“‚.caches.hessian_parameters = buffer_parameters - โˆ‡โ‚‚_SS_and_pars = derivatives[2][2][:,nps+1:end] - - lennz = nnz(โˆ‡โ‚‚_SS_and_pars) - - if (lennz / length(โˆ‡โ‚‚_SS_and_pars) > density_threshold) || (length(โˆ‡โ‚‚_SS_and_pars) < min_length) - โˆ‡โ‚‚_SS_and_pars_mat = convert(Matrix, โˆ‡โ‚‚_SS_and_pars) - buffer_SS_and_pars = zeros(Float64, size(โˆ‡โ‚‚_SS_and_pars)) - else - โˆ‡โ‚‚_SS_and_pars_mat = โˆ‡โ‚‚_SS_and_pars - buffer_SS_and_pars = similar(โˆ‡โ‚‚_SS_and_pars, Float64) - buffer_SS_and_pars.nzval .= 0 - end + โˆ‡โ‚‚_SS_and_pars_mat, buffer_SS_and_pars, lennz = prepare_sensitivity_buffer(derivatives[2][2][:,nps+1:end]) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -8000,18 +7077,7 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; ๐“‚.caches.third_order_derivatives = buffer - โˆ‡โ‚ƒ_parameters = derivatives[3][2][:,1:nps] - - lennz = nnz(โˆ‡โ‚ƒ_parameters) - - if (lennz / length(โˆ‡โ‚ƒ_parameters) > density_threshold) || (length(โˆ‡โ‚ƒ_parameters) < min_length) - โˆ‡โ‚ƒ_parameters_mat = convert(Matrix, โˆ‡โ‚ƒ_parameters) - buffer_parameters = zeros(Float64, size(โˆ‡โ‚ƒ_parameters)) - else - โˆ‡โ‚ƒ_parameters_mat = โˆ‡โ‚ƒ_parameters - buffer_parameters = similar(โˆ‡โ‚ƒ_parameters, Float64) - buffer_parameters.nzval .= 0 - end + โˆ‡โ‚ƒ_parameters_mat, buffer_parameters, lennz = prepare_sensitivity_buffer(derivatives[3][2][:,1:nps]) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -8030,18 +7096,7 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; ๐“‚.caches.third_order_derivatives_parameters = buffer_parameters - โˆ‡โ‚ƒ_SS_and_pars = derivatives[3][2][:,nps+1:end] - - lennz = nnz(โˆ‡โ‚ƒ_SS_and_pars) - - if (lennz / length(โˆ‡โ‚ƒ_SS_and_pars) > density_threshold) || (length(โˆ‡โ‚ƒ_SS_and_pars) < min_length) - โˆ‡โ‚ƒ_SS_and_pars_mat = convert(Matrix, โˆ‡โ‚ƒ_SS_and_pars) - buffer_SS_and_pars = zeros(Float64, size(โˆ‡โ‚ƒ_SS_and_pars)) - else - โˆ‡โ‚ƒ_SS_and_pars_mat = โˆ‡โ‚ƒ_SS_and_pars - buffer_SS_and_pars = similar(โˆ‡โ‚ƒ_SS_and_pars, Float64) - buffer_SS_and_pars.nzval .= 0 - end + โˆ‡โ‚ƒ_SS_and_pars_mat, buffer_SS_and_pars, lennz = prepare_sensitivity_buffer(derivatives[3][2][:,nps+1:end]) if lennz > nnz_parallel_threshold parallel = Symbolics.ShardedForm(1500,4) @@ -8064,6 +7119,11 @@ function write_functions_mapping!(๐“‚::โ„ณ, max_perturbation_order::Int; end end + # Invalidate derivative stamps since buffers were replaced with fresh (zeroed) content. + # Without this, calculate_jacobian/hessian/third_order_derivatives would return stale + # zero-filled buffers on a cache hit, causing downstream DimensionMismatch errors. + ๐“‚.caches.valid_for.jacobian = Float64[] + return nothing end @@ -8148,17 +7208,6 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::D; verbose::Bool = true) # Remove the provided missing params from the missing list remaining_missing = setdiff(p.missing_parameters, missing_params_provided) - # Mark that solution needs to be recomputed - ๐“‚.caches.outdated.non_stochastic_steady_state = true - ๐“‚.caches.outdated.jacobian = true - ๐“‚.caches.outdated.hessian = true - ๐“‚.caches.outdated.third_order_derivatives = true - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true - # If all missing parameters are now provided, print a message if !isempty(remaining_missing) @info "Remaining missing parameters: ", remaining_missing @@ -8190,10 +7239,10 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::D; verbose::Bool = true) ) ๐“‚.parameter_values = vcat(declared_values, missing_values, remaining_missing_values) - # Clear the NSSS_solver_cache since parameter order/count has changed - # It will be rebuilt when write_steady_state_solver_function! is called with correct parameter count - while length(๐“‚.caches.solver_cache) > 0 - pop!(๐“‚.caches.solver_cache) + # Clear NSSS solver cache because parameter order/count changed. + # It will be rebuilt during the next NSSS setup. + while length(๐“‚.caches.solver) > 0 + pop!(๐“‚.caches.solver) end end @@ -8231,23 +7280,10 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::D; verbose::Bool = true) if !all(๐“‚.parameter_values[ntrsct_idx] .== collect(values(parameters))) && !(p.parameters[ntrsct_idx] == [:activeแต’แต‡แถœshocks]) if verbose println("Parameter changes: ") end - ๐“‚.caches.outdated.jacobian = true - ๐“‚.caches.outdated.hessian = true - ๐“‚.caches.outdated.third_order_derivatives = true - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true end for i in 1:length(parameters) if ๐“‚.parameter_values[ntrsct_idx[i]] != collect(values(parameters))[i] - if isnothing(๐“‚.NSSS.dependencies) || (collect(keys(parameters))[i] โˆˆ ๐“‚.NSSS.dependencies[end][2] && ๐“‚.caches.outdated.non_stochastic_steady_state == false) - # if !isnothing(๐“‚.NSSS.dependencies) && collect(keys(parameters))[i] โˆˆ ๐“‚.NSSS.dependencies[end][2] && ๐“‚.caches.outdated.non_stochastic_steady_state == false - ๐“‚.caches.outdated.non_stochastic_steady_state = true - end - if verbose println("\t",p.parameters[ntrsct_idx[i]],"\tfrom ",๐“‚.parameter_values[ntrsct_idx[i]],"\tto ",collect(values(parameters))[i]) end ๐“‚.parameter_values[ntrsct_idx[i]] = collect(values(parameters))[i] @@ -8255,8 +7291,6 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::D; verbose::Bool = true) end end - if ๐“‚.caches.outdated.non_stochastic_steady_state == true && verbose println("New parameters changed the steady state.") end - return nothing end @@ -8280,8 +7314,9 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::Vector{Float64}; verbose end bounds_broken = false + parameters_dict = Dict(๐“‚.constants.post_complete_parameters.parameters .=> parameters) - for (par,val) in Dict(๐“‚.constants.post_complete_parameters.parameters .=> parameters) + for (par, val) in parameters_dict if haskey(๐“‚.constants.post_parameters_macro.bounds,par) if val > ๐“‚.constants.post_parameters_macro.bounds[par][2] @warn("Calibration is out of bounds for $par < $(๐“‚.constants.post_parameters_macro.bounds[par][2])\t parameter value: $val") @@ -8300,16 +7335,6 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::Vector{Float64}; verbose @warn("Parameters unchanged.") else if !all(parameters .== ๐“‚.parameter_values[1:length(parameters)]) - ๐“‚.caches.outdated.non_stochastic_steady_state = true - ๐“‚.caches.outdated.jacobian = true - ๐“‚.caches.outdated.hessian = true - ๐“‚.caches.outdated.third_order_derivatives = true - ๐“‚.caches.outdated.first_order_solution = true - ๐“‚.caches.outdated.second_order_solution = true - ๐“‚.caches.outdated.pruned_second_order_solution = true - ๐“‚.caches.outdated.third_order_solution = true - ๐“‚.caches.outdated.pruned_third_order_solution = true - match_idx = [] for (i, v) in enumerate(parameters) if v != ๐“‚.parameter_values[i] @@ -8331,7 +7356,9 @@ function write_parameters_input!(๐“‚::โ„ณ, parameters::Vector{Float64}; verbose end end - if ๐“‚.caches.outdated.non_stochastic_steady_state == true && verbose println("New parameters changed the steady state.") end + if ๐“‚.caches.valid_for.non_stochastic_steady_state != ๐“‚.parameter_values && verbose + println("New parameters changed the steady state.") + end return nothing end @@ -8489,7 +7516,14 @@ end function calculate_jacobian(parameters::Vector{M}, SS_and_pars::Vector{N}, caches_obj::caches, - jacobian_funcs::jacobian_functions)::Matrix{M} where {M,N} + jacobian_funcs::jacobian_functions, + workspaces::workspaces; + caching::Bool = true)::Matrix{M} where {M,N} + # Cache hit: return cached jacobian if valid for current parameters + if caching && M === Float64 && cache_valid_for_parameters(caches_obj.valid_for.jacobian, parameters) && caches_obj.jacobian isa Matrix{M} && !isempty(caches_obj.jacobian) + return caches_obj.jacobian + end + if eltype(caches_obj.jacobian) != M if caches_obj.jacobian isa SparseMatrixCSC jac_buffer = similar(caches_obj.jacobian,M) @@ -8502,6 +7536,11 @@ function calculate_jacobian(parameters::Vector{M}, end jacobian_funcs.f(jac_buffer, parameters, SS_and_pars) + + if caching && M === Float64 + caches_obj.jacobian = jac_buffer + caches_obj.valid_for.jacobian = Float64.(parameters) + end return jac_buffer end @@ -8509,7 +7548,19 @@ end function calculate_hessian(parameters::Vector{M}, SS_and_pars::Vector{N}, caches_obj::caches, - hessian_funcs::hessian_functions)::SparseMatrixCSC{M, Int} where {M,N} + hessian_funcs::hessian_functions, + workspaces::workspaces; + caching::Bool = true)::SparseMatrixCSC{M, Int} where {M,N} + # Cache hit: return cached hessian if valid for current parameters + if caching && M === Float64 && cache_valid_for_parameters(caches_obj.valid_for.hessian, parameters) && caches_obj.hessian isa SparseMatrixCSC{M, Int} && !isempty(caches_obj.hessian) + return caches_obj.hessian + end + + S = promote_type(M, N) + if eltype(workspaces.second_order.Sฬ‚) != S + workspaces.second_order = Higher_order_workspace(T = S) + end + if eltype(caches_obj.hessian) != M if caches_obj.hessian isa SparseMatrixCSC hes_buffer = similar(caches_obj.hessian,M) @@ -8522,6 +7573,11 @@ function calculate_hessian(parameters::Vector{M}, end hessian_funcs.f(hes_buffer, parameters, SS_and_pars) + + if caching && M === Float64 + caches_obj.hessian = hes_buffer + caches_obj.valid_for.hessian = Float64.(parameters) + end return hes_buffer end @@ -8530,7 +7586,19 @@ end function calculate_third_order_derivatives(parameters::Vector{M}, SS_and_pars::Vector{N}, caches_obj::caches, - third_order_derivatives_funcs::third_order_derivatives_functions)::SparseMatrixCSC{M, Int} where {M,N} + third_order_derivatives_funcs::third_order_derivatives_functions, + workspaces::workspaces; + caching::Bool = true)::SparseMatrixCSC{M, Int} where {M,N} + # Cache hit: return cached third order derivatives if valid for current parameters + if caching && M === Float64 && cache_valid_for_parameters(caches_obj.valid_for.third_order_derivatives, parameters) && caches_obj.third_order_derivatives isa SparseMatrixCSC{M, Int} && !isempty(caches_obj.third_order_derivatives) + return caches_obj.third_order_derivatives + end + + S = promote_type(M, N) + if eltype(workspaces.third_order.ลœ) != S + workspaces.third_order = Higher_order_workspace(T = S) + end + if eltype(caches_obj.third_order_derivatives) != M if caches_obj.third_order_derivatives isa SparseMatrixCSC third_buffer = similar(caches_obj.third_order_derivatives,M) @@ -8543,6 +7611,11 @@ function calculate_third_order_derivatives(parameters::Vector{M}, end third_order_derivatives_funcs.f(third_buffer, parameters, SS_and_pars) + + if caching && M === Float64 + caches_obj.third_order_derivatives = third_buffer + caches_obj.valid_for.third_order_derivatives = Float64.(parameters) + end return third_buffer end @@ -9574,49 +8647,132 @@ end end # dispatch_doctor -noop_state_update(::Float64, ::Float64) = nothing +noop_state_update(state::AbstractVector{<:Real}, ::AbstractVector{<:Real}) = state +noop_state_update(state::AbstractVector{<:AbstractVector{<:Real}}, ::AbstractVector{<:Real}) = state + +function initialize_pruned_state(state::AbstractVector{T}, n_states::Int) where T <: Real + return [Vector{T}(state), zeros(T, n_states)] +end + +function initialize_pruned_state(state::AbstractVector{T}, n_states::Int, ::Val{3}) where T <: Real + return [Vector{T}(state), zeros(T, n_states), zeros(T, n_states)] +end + +function pruned_second_order_state_update(pruned_states::AbstractVector{<:AbstractVector{T}}, shock::AbstractVector{S}, past_idx, n_states::Int, ๐’โ‚, ๐’โ‚‚) where {T <: Real, S <: Real} + aug_stateโ‚ = [pruned_states[1][past_idx]; 1; shock] + aug_stateโ‚‚ = [pruned_states[2][past_idx]; 0; zero(shock)] + return [๐’โ‚ * aug_stateโ‚, ๐’โ‚ * aug_stateโ‚‚ + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2] +end + +function pruned_second_order_state_update(state::AbstractVector{T}, shock::AbstractVector{S}, past_idx, n_states::Int, ๐’โ‚, ๐’โ‚‚) where {T <: Real, S <: Real} + return pruned_second_order_state_update(initialize_pruned_state(state, n_states), shock, past_idx, n_states, ๐’โ‚, ๐’โ‚‚) +end + +function pruned_third_order_state_update(pruned_states::AbstractVector{<:AbstractVector{T}}, shock::AbstractVector{S}, past_idx, n_states::Int, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ) where {T <: Real, S <: Real} + aug_stateโ‚ = [pruned_states[1][past_idx]; 1; shock] + aug_stateโ‚ฬ‚ = [pruned_states[1][past_idx]; 0; shock] + aug_stateโ‚‚ = [pruned_states[2][past_idx]; 0; zero(shock)] + aug_stateโ‚ƒ = [pruned_states[3][past_idx]; 0; zero(shock)] + kron_aug_stateโ‚ = โ„’.kron(aug_stateโ‚, aug_stateโ‚) + return [๐’โ‚ * aug_stateโ‚, ๐’โ‚ * aug_stateโ‚‚ + ๐’โ‚‚ * kron_aug_stateโ‚ / 2, ๐’โ‚ * aug_stateโ‚ƒ + ๐’โ‚‚ * โ„’.kron(aug_stateโ‚ฬ‚, aug_stateโ‚‚) + ๐’โ‚ƒ * โ„’.kron(kron_aug_stateโ‚,aug_stateโ‚) / 6] +end + +function pruned_third_order_state_update(state::AbstractVector{T}, shock::AbstractVector{S}, past_idx, n_states::Int, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ) where {T <: Real, S <: Real} + return pruned_third_order_state_update(initialize_pruned_state(state, n_states, Val(3)), shock, past_idx, n_states, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ) +end function parse_algorithm_to_state_update(algorithm::Symbol, ๐“‚::โ„ณ, occasionally_binding_constraints::Bool)::Tuple{Function, Bool} state_update::Function = noop_state_update pruning::Bool = algorithm โˆˆ [:pruned_second_order, :pruned_third_order] + past_idx = ๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + nVars = ๐“‚.constants.post_model_macro.nVars + if occasionally_binding_constraints + Sฬ‚โ‚ = ๐“‚.caches.first_order_obc_solution_matrix + if algorithm == :first_order - state_update = ๐“‚.functions.first_order_state_update_obc::Function - elseif :second_order == algorithm - state_update = ๐“‚.functions.second_order_state_update_obc::Function - elseif :pruned_second_order == algorithm - state_update = ๐“‚.functions.pruned_second_order_state_update_obc::Function - elseif :third_order == algorithm - state_update = ๐“‚.functions.third_order_state_update_obc::Function - elseif :pruned_third_order == algorithm - state_update = ๐“‚.functions.pruned_third_order_state_update_obc::Function + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[past_idx]; shock] + return Sฬ‚โ‚ * aug_state + end + elseif algorithm โˆˆ [:second_order, :third_order] + ๐’โ‚‚ = ๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚ + Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:nPast] zeros(nVars) Sฬ‚โ‚[:,nPast+1:end]] + + if algorithm == :second_order + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[past_idx]; 1; shock] + return Sฬ‚โ‚ฬ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + end + else # :third_order + ๐’โ‚ƒ = ๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[past_idx]; 1; shock] + return Sฬ‚โ‚ฬ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 + end + end + elseif algorithm == :pruned_second_order + ๐’โ‚‚ = ๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚ + Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:nPast] zeros(nVars) Sฬ‚โ‚[:,nPast+1:end]] + state_update = (state, shock) -> pruned_second_order_state_update(state, shock, past_idx, nVars, Sฬ‚โ‚ฬ‚, ๐’โ‚‚) + elseif algorithm == :pruned_third_order + ๐’โ‚‚ = ๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚ƒ = ๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ + Sฬ‚โ‚ฬ‚ = [Sฬ‚โ‚[:,1:nPast] zeros(nVars) Sฬ‚โ‚[:,nPast+1:end]] + state_update = (state, shock) -> pruned_third_order_state_update(state, shock, past_idx, nVars, Sฬ‚โ‚ฬ‚, ๐’โ‚‚, ๐’โ‚ƒ) end else if algorithm == :first_order - state_update = ๐“‚.functions.first_order_state_update::Function - elseif :second_order == algorithm - state_update = ๐“‚.functions.second_order_state_update::Function - elseif :pruned_second_order == algorithm - state_update = ๐“‚.functions.pruned_second_order_state_update::Function - elseif :third_order == algorithm - state_update = ๐“‚.functions.third_order_state_update::Function - elseif :pruned_third_order == algorithm - state_update = ๐“‚.functions.pruned_third_order_state_update::Function + Sโ‚ = ๐“‚.caches.first_order_solution_matrix + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[past_idx]; shock] + return Sโ‚ * aug_state + end + elseif algorithm โˆˆ [:second_order, :third_order] + Sโ‚ = ๐“‚.caches.first_order_solution_matrix + ๐’โ‚ = [Sโ‚[:,1:nPast] zeros(nVars) Sโ‚[:,nPast+1:end]] + ๐’โ‚‚ = ๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚ + + if algorithm == :second_order + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[past_idx]; 1; shock] + return ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + end + else # :third_order + ๐’โ‚ƒ = ๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ + state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} + aug_state = [state[past_idx]; 1; shock] + return ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 + end + end + elseif algorithm == :pruned_second_order + Sโ‚ = ๐“‚.caches.first_order_solution_matrix + ๐’โ‚ = [Sโ‚[:,1:nPast] zeros(nVars) Sโ‚[:,nPast+1:end]] + ๐’โ‚‚ = ๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚ + state_update = (state, shock) -> pruned_second_order_state_update(state, shock, past_idx, nVars, ๐’โ‚, ๐’โ‚‚) + elseif algorithm == :pruned_third_order + Sโ‚ = ๐“‚.caches.first_order_solution_matrix + ๐’โ‚ = [Sโ‚[:,1:nPast] zeros(nVars) Sโ‚[:,nPast+1:end]] + ๐’โ‚‚ = ๐“‚.caches.second_order_solution * ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚ƒ = ๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ + state_update = (state, shock) -> pruned_third_order_state_update(state, shock, past_idx, nVars, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ) end end return (state_update, pruning) end + @stable default_mode = "disable" begin -function get_custom_steady_state_buffer!(๐“‚::โ„ณ, expected_length::Int) - buffer = ๐“‚.workspaces.custom_steady_state_buffer +function get_custom_steady_state_workspace!(๐“‚::โ„ณ, expected_length::Int) + buffer = ๐“‚.workspaces.custom_steady_state if length(buffer) != expected_length buffer = Vector{Float64}(undef, expected_length) - ๐“‚.workspaces.custom_steady_state_buffer = buffer + ๐“‚.workspaces.custom_steady_state = buffer end return buffer @@ -9633,7 +8789,7 @@ function evaluate_custom_steady_state_function(๐“‚::โ„ณ, has_inplace = hasmethod(๐“‚.functions.NSSS_custom, Tuple{typeof(parameter_values), typeof(parameter_values)}) if has_inplace - get_custom_steady_state_buffer!(๐“‚, expected_length) + get_custom_steady_state_workspace!(๐“‚, expected_length) output = Vector{S}(undef, expected_length) try @@ -9769,11 +8925,18 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, parameter_values::Vector{S}; opts::CalculationOptions = merge_calculation_options(), cold_start::Bool = false, - estimation::Bool = false)::Tuple{Vector{S}, Tuple{S, Int}} where S <: Real + estimation::Bool = false, + caching::Bool = true)::Tuple{Vector{S}, Tuple{S, Int}} where S <: Real # timer::TimerOutput = TimerOutput(), + # @timeit_debug timer "Calculate NSSS" begin ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + # Cache hit: return cached NSSS if valid for current parameters + if caching && S === Float64 && cache_valid_for_parameters(๐“‚.caches.valid_for.non_stochastic_steady_state, parameter_values) && !isempty(๐“‚.caches.non_stochastic_steady_state) + return (copy(๐“‚.caches.non_stochastic_steady_state), (zero(S), 0))::Tuple{Vector{S}, Tuple{S, Int}} + end + # Use custom steady state function if available, otherwise use default solver if ๐“‚.functions.NSSS_custom isa Function vars_in_ss_equations = ms.vars_in_ss_equations @@ -9786,7 +8949,8 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, length(๐“‚.constants.post_complete_parameters.parameters), ) - residual = zeros(length(๐“‚.equations.steady_state) + length(๐“‚.equations.calibration)) + residual = ๐“‚.workspaces.nsss_solver.check_residual + fill!(residual, 0.0) ๐“‚.functions.NSSS_check(residual, parameter_values, SS_and_pars_tmp) @@ -9794,17 +8958,19 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, iters = 0 - # if !isfinite(solution_error) || solution_error > opts.tol.NSSS_acceptance_tol - # throw(ArgumentError("Custom steady state function failed steady state check: residual $solution_error > $(opts.tol.NSSS_acceptance_tol). Parameters: $(parameter_values). Steady state and parameters returned: $(SS_and_pars_tmp).")) + # if !isfinite(solution_error) || solution_error > opts.tol.nsss.acceptance_tol + # throw(ArgumentError("Custom steady state function failed steady state check: residual $solution_error > $(opts.tol.nsss.acceptance_tol). Parameters: $(parameter_values). Steady state and parameters returned: $(SS_and_pars_tmp).")) # end - X = @ignore_derivatives ms.custom_ss_expand_matrix + X = ms.custom_ss_expand_matrix SS_and_pars = X * SS_and_pars_tmp else - SS_and_pars, (solution_error, iters) = ๐“‚.functions.NSSS_solve(parameter_values, ๐“‚, opts.tol, opts.verbose, cold_start, DEFAULT_SOLVER_PARAMETERS) + fastest_idx = ๐“‚.constants.post_complete_parameters.nsss_fastest_solver_parameter_idx + preferred_solver_parameter_idx = fastest_idx < 1 || fastest_idx > length(DEFAULT_SOLVER_PARAMETERS) ? 1 : fastest_idx + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(parameter_values, ๐“‚, opts.tol, opts.verbose, cold_start, DEFAULT_SOLVER_PARAMETERS, preferred_solver_parameter_idx = preferred_solver_parameter_idx) end # Update counters - solved = !(solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error)) + solved = !(solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error)) update_ss_counter!(๐“‚.counters, solved, estimation = estimation) if !solved @@ -9815,11 +8981,25 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, end # end # timeit_debug + + # Cache write: store NSSS result and stamp + if caching + cache_ss = ๐“‚.caches.non_stochastic_steady_state + if length(cache_ss) != length(SS_and_pars) + resize!(cache_ss, length(SS_and_pars)) + end + copyto!(cache_ss, SS_and_pars) + if solved + ๐“‚.caches.valid_for.non_stochastic_steady_state = eltype(parameter_values) <: โ„ฑ.Dual ? Float64.(โ„ฑ.value.(parameter_values)) : Float64.(parameter_values) + else + ๐“‚.caches.valid_for.non_stochastic_steady_state = Float64[] + end + end + return SS_and_pars, (solution_error, iters) end - function check_bounds(parameter_values::Vector{S}, ๐“‚::โ„ณ)::Bool where S <: Real if !all(isfinite,parameter_values) return true end @@ -9842,9 +9022,9 @@ function get_relevant_steady_state_and_state_update(::Val{:second_order}, opts::CalculationOptions = merge_calculation_options(), estimation::Bool = false) where S <: Real # timer::TimerOutput = TimerOutput(), - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_second_order_stochastic_steady_state(parameter_values, ๐“‚, opts = opts, estimation = estimation) # timer = timer, + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_stochastic_steady_state(Val(:second_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) # timer = timer, - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol if opts.verbose println("Could not find 2nd order stochastic steady state") end return ๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚], collect(sss), converged end @@ -9865,9 +9045,9 @@ function get_relevant_steady_state_and_state_update(::Val{:pruned_second_order}, opts::CalculationOptions = merge_calculation_options(), estimation::Bool = false)::Tuple{constants, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{S}}, Bool} where S <: Real # timer::TimerOutput = TimerOutput(), - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_second_order_stochastic_steady_state(parameter_values, ๐“‚, pruning = true, opts = opts, estimation = estimation) # timer = timer, + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_stochastic_steady_state(Val(:pruned_second_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) # timer = timer, - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol if opts.verbose println("Could not find 2nd order stochastic steady state") end return ๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚], [zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nVars)], converged end @@ -9888,9 +9068,9 @@ function get_relevant_steady_state_and_state_update(::Val{:third_order}, opts::CalculationOptions = merge_calculation_options(), estimation::Bool = false)::Tuple{constants, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{S}, Bool} where S <: Real # timer::TimerOutput = TimerOutput(), - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_third_order_stochastic_steady_state(parameter_values, ๐“‚, opts = opts, estimation = estimation) # timer = timer, + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_stochastic_steady_state(Val(:third_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) # timer = timer, - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol if opts.verbose println("Could not find 3rd order stochastic steady state") end return ๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ], collect(sss), converged end @@ -9911,9 +9091,9 @@ function get_relevant_steady_state_and_state_update(::Val{:pruned_third_order}, opts::CalculationOptions = merge_calculation_options(), estimation::Bool = false)::Tuple{constants, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{S}}, Bool} where S <: Real # timer::TimerOutput = TimerOutput(), - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_third_order_stochastic_steady_state(parameter_values, ๐“‚, pruning = true, opts = opts, estimation = estimation) # timer = timer, + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_stochastic_steady_state(Val(:pruned_third_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) # timer = timer, - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol if opts.verbose println("Could not find 3rd order stochastic steady state") end return ๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ], [zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nVars), zeros(๐“‚.constants.post_model_macro.nVars)], converged end @@ -9936,31 +9116,27 @@ function get_relevant_steady_state_and_state_update(::Val{:first_order}, # Initialize constants at entry point constants_obj = initialise_constants!(๐“‚) - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameter_values, opts = opts, estimation = estimation) # timer = timer, + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameter_values, opts = opts, estimation = estimation) # timer = timer, state = zeros(๐“‚.constants.post_model_macro.nVars) - if solution_error > opts.tol.NSSS_acceptance_tol # || isnan(solution_error) if it's NaN the first condition is false anyway + if solution_error > opts.tol.nsss.acceptance_tol # || isnan(solution_error) if it's NaN the first condition is false anyway # println("NSSS not found") - return ๐“‚.constants, SS_and_pars, zeros(S, 0, 0), [state], solution_error < opts.tol.NSSS_acceptance_tol + return ๐“‚.constants, SS_and_pars, zeros(S, 0, 0), [state], solution_error < opts.tol.nsss.acceptance_tol end - โˆ‡โ‚ = calculate_jacobian(parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian) # , timer = timer)# |> Matrix + โˆ‡โ‚ = calculate_jacobian(parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) # , timer = timer)# |> Matrix - qme_ws = @ignore_derivatives ensure_qme_workspace!(๐“‚) - sylv_ws = @ignore_derivatives ensure_sylvester_1st_order_workspace!(๐“‚) - ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants_obj, - qme_ws, - sylv_ws; - # timer = timer, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, initial_guess = ๐“‚.caches.qme_solution, - opts = opts) + parameter_values = parameter_values) - if solved ๐“‚.caches.qme_solution = qme_sol end - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) if !solved # println("NSSS not found") @@ -10053,6 +9229,6 @@ include("./custom_autodiff_rules/forwarddiff.jl") # Include rrule definitions for reverse-mode AD (Zygote/ChainRulesCore) # Must be at the end of the module because rrules depend on function definitions -include("./custom_autodiff_rules/zygote.jl") +include("./custom_autodiff_rules/rrules.jl") end diff --git a/src/algorithms/fast_lapack_wrappers.jl b/src/algorithms/fast_lapack_wrappers.jl new file mode 100644 index 000000000..0d4dbfa22 --- /dev/null +++ b/src/algorithms/fast_lapack_wrappers.jl @@ -0,0 +1,161 @@ +@stable default_mode = "disable" begin + +function factorize_qr!(qr_mat::AbstractMatrix, + qr_factors::AbstractMatrix{R}, + qr_ws::FastLapackInterface.QRWs{R}; + use_fastlapack_qr::Bool = true) where {R <: AbstractFloat} + if use_fastlapack_qr && R <: Union{Float32, Float64} + copyto!(qr_factors, qr_mat) + โ„’.LAPACK.geqrf!(qr_ws, qr_factors; resize = true) + return qr_factors + else + copyto!(qr_factors, qr_mat) + return โ„’.qr!(qr_factors) + end +end + +function apply_qr_transpose_left!(dest::AbstractMatrix{R}, + src::AbstractMatrix, + Q::AbstractMatrix{R}, + qr_orm_ws, + qr_orm_dims::NTuple{3, Int}, + qr_ws; + use_fastlapack_qr::Bool = true) where {R <: AbstractFloat} + orm_dims = (size(Q, 1), size(Q, 2), size(src, 2)) + if qr_orm_dims != orm_dims + qr_orm_ws = FastLapackInterface.QROrmWs(qr_ws, 'L', 'T', Q, src) + qr_orm_dims = orm_dims + end + + copyto!(dest, src) + โ„’.LAPACK.ormqr!(qr_orm_ws, 'L', 'T', Q, dest) + return qr_orm_ws, qr_orm_dims +end + +function apply_qr_transpose_left!(dest::AbstractMatrix{R}, + src::AbstractMatrix, + Q::โ„’.QRCompactWY, + qr_orm_ws, + qr_orm_dims::NTuple{3, Int}, + qr_ws; + use_fastlapack_qr::Bool = true) where {R <: AbstractFloat} + โ„’.mul!(dest, Q.Q', src) + return qr_orm_ws, qr_orm_dims +end + +function factorize_lu!(A::AbstractMatrix{R}, + lu_ws, + lu_dims::NTuple{2, Int}; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + if use_fastlapack_lu && R <: Union{Float32, Float64} + dims = (size(A, 1), size(A, 2)) + if lu_dims != dims + lu_ws = FastLapackInterface.LUWs(A) + lu_dims = dims + end + _, _, info = โ„’.LAPACK.getrf!(lu_ws, A; resize = true) + return lu_ws, lu_dims, info == 0, nothing + else + lu = โ„’.lu!(A, check = false) + return lu_ws, lu_dims, โ„’.issuccess(lu), lu + end +end + +function solve_lu_left!(A::AbstractMatrix{R}, + B::AbstractVecOrMat{R}, + lu_ws, + lu; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + if use_fastlapack_lu && R <: Union{Float32, Float64} + โ„’.LAPACK.getrs!(lu_ws, 'N', A, B) + else + โ„’.ldiv!(lu, B) + end + return B +end + +function solve_lu_left!(A::AbstractMatrix{R}, + B::AbstractVecOrMat{R}, + lu_ws, + lu::Nothing; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + โ„’.LAPACK.getrs!(lu_ws, 'N', A, B) + return B +end + +function solve_lu_right!(A::AbstractMatrix{R}, + B::AbstractMatrix{R}, + lu_ws, + lu, + rhs_t::AbstractMatrix{R}; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + if use_fastlapack_lu && R <: Union{Float32, Float64} + rhs_t_dims = (size(B, 2), size(B, 1)) + @assert size(rhs_t) == rhs_t_dims + + copyto!(rhs_t, transpose(B)) + โ„’.LAPACK.getrs!(lu_ws, 'T', A, rhs_t) + copyto!(B, transpose(rhs_t)) + else + โ„’.rdiv!(B, lu) + end + return B +end + +function solve_lu_right!(A::AbstractMatrix{R}, + B::AbstractMatrix{R}, + lu_ws, + lu::Nothing, + rhs_t::AbstractMatrix{R}; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + rhs_t_dims = (size(B, 2), size(B, 1)) + @assert size(rhs_t) == rhs_t_dims + + copyto!(rhs_t, transpose(B)) + โ„’.LAPACK.getrs!(lu_ws, 'T', A, rhs_t) + copyto!(B, transpose(rhs_t)) + return B +end + +function factorize_generalized_schur!(D::AbstractMatrix{R}, + E::AbstractMatrix{R}, + qz_ws, + qz_dims::NTuple{2, Int}, + eigenselect::AbstractVector{Bool}; + use_fastlapack_schur::Bool = true) where {R <: AbstractFloat} + if use_fastlapack_schur && R <: Union{Float32, Float64} + dims = (size(D, 1), size(D, 2)) + if qz_dims != dims + qz_ws = FastLapackInterface.GeneralizedSchurWs(D) + qz_dims = dims + end + + try + S, T, _, _, _, Z = โ„’.LAPACK.gges!(qz_ws, 'V', 'V', D, E; + select = FastLapackInterface.ed, + criterium = 1.0, + resize = true) + return qz_ws, qz_dims, (S = S, T = T, Z = Z), true + catch + return qz_ws, qz_dims, nothing, false + end + else + schdcmp = try + โ„’.schur!(D, E) + catch + return qz_ws, qz_dims, nothing, false + end + + @. eigenselect = abs(schdcmp.ฮฒ / schdcmp.ฮฑ) < 1 + + try + โ„’.ordschur!(schdcmp, eigenselect) + catch + return qz_ws, qz_dims, nothing, false + end + + return qz_ws, qz_dims, schdcmp, true + end +end + +end # dispatch_doctor diff --git a/src/algorithms/lyapunov.jl b/src/algorithms/lyapunov.jl index 36ee6f71b..82db3c31c 100644 --- a/src/algorithms/lyapunov.jl +++ b/src/algorithms/lyapunov.jl @@ -10,14 +10,68 @@ # solves: A * X * A' + C = X @stable default_mode = "disable" begin +# Pack upper triangle of a symmetric matrix into a vech vector (in-place). +function vech!(vech_vector::AbstractVector, symmetric_matrix::AbstractMatrix) + matrix_size = size(symmetric_matrix, 1) + @inbounds for column in 1:matrix_size + offset = div(column * (column - 1), 2) + @simd for row in 1:column + vech_vector[offset + row] = symmetric_matrix[row, column] + end + end + return vech_vector +end + +# Unpack a vech vector into a full symmetric matrix (in-place). +function fill_symmetric_from_vech!(symmetric_matrix::AbstractMatrix, vech_vector::AbstractVector) + matrix_size = size(symmetric_matrix, 1) + # Fill the upper triangle + @inbounds for column in 1:matrix_size + offset = div(column * (column - 1), 2) + @simd for row in 1:column + symmetric_matrix[row, column] = vech_vector[offset + row] + end + end + # Copy the upper triangle to the lower triangle + @inbounds for column in 1:matrix_size + @simd for row in (column + 1):matrix_size + symmetric_matrix[row, column] = symmetric_matrix[column, row] + end + end + return symmetric_matrix +end + +# Approximate symmetry check (allocation-free). Returns true when +# max|C[i,j] - C[j,i]| โ‰ค rtol ยท max|C[i,j]| over all off-diagonal pairs. +function _is_approx_symmetric(C::AbstractMatrix; + rtol::Real = sqrt(eps(real(eltype(C))))) + m, n = size(C) + m == n || return false + max_asym = zero(real(eltype(C))) + max_abs = zero(real(eltype(C))) + @inbounds for j in 1:n, i in 1:(j - 1) + max_asym = max(max_asym, abs(C[i, j] - C[j, i])) + max_abs = max(max_abs, abs(C[i, j]), abs(C[j, i])) + end + return max_abs == 0 ? true : max_asym โ‰ค rtol * max_abs +end + function solve_lyapunov_equation(A::AbstractMatrix{T}, C::AbstractMatrix{T}, workspace::lyapunov_workspace; + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), lyapunov_algorithm::Symbol = :doubling, - tol::AbstractFloat = 1e-14, - acceptance_tol::AbstractFloat = 1e-12, + tol::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), verbose::Bool = false)::Union{Tuple{Matrix{T}, Bool}, Tuple{ThreadedSparseArrays.ThreadedSparseMatrixCSC{T, Int, SparseMatrixCSC{T, Int}}, Bool}} where T <: Float64 # timer::TimerOutput = TimerOutput(), + # Ownership: low-level methods below are mixed. Bartels-Stewart and sparse + # doubling paths return owned matrices, while dense doubling and Krylov + # paths can return workspace-backed buffers such as workspace.๐‚/workspace.๐—. + # This dispatcher currently returns X directly, so callers must not retain + # the result across workspace reuse unless they make their own copy. # Update workspace dimension if needed (for cases like Kalman filter where dimension differs from initial setup) n = size(A, 1) if workspace.n != n @@ -36,6 +90,29 @@ function solve_lyapunov_equation(A::AbstractMatrix{T}, # C = choose_matrix_format(C, density_threshold = 0.0) C = collect(C) # C is always dense because the output will be dense in all of these cases as we use this function to compute dense covariance matrices + + initial_guess_acceptance_tol = tol.initial_guess_acceptance_tol + acceptance_tol = tol.acceptance_tol + + if length(initial_guess) > 0 + guess = initial_guess + if size(guess) == size(C) + ensure_lyapunov_doubling_buffers!(workspace) + _tmp = workspace.๐‚A + _res = workspace.๐‚ยน + โ„’.mul!(_tmp, guess, A') + โ„’.mul!(_res, A, _tmp) + โ„’.axpy!(1, C, _res) + โ„’.axpy!(-1, guess, _res) + + denom = max(โ„’.norm(guess), โ„’.norm(C)) + reached_tol = denom == 0 ? 0.0 : โ„’.norm(_res) / denom + if reached_tol < initial_guess_acceptance_tol + if verbose println("Lyapunov equation - initial guess achieves relative tol of $reached_tol (initial guess tol: $initial_guess_acceptance_tol)") end + return choose_matrix_format(guess), true + end + end + end # end # timeit_debug # @timeit_debug timer "Solve" begin @@ -92,21 +169,24 @@ function solve_lyapunov_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat ::Val{:bartels_stewart}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::AbstractFloat = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat - # Note: workspace is unused by bartels_stewart but accepted for API consistency + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned dense matrix from MatrixEquations.lyapd. ๐‚ = try MatrixEquations.lyapd(A, C)::Matrix{T} catch return C, 0, 1.0 end - # ๐‚ยน = A * ๐‚ * A' + C - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน - ๐‚) / denom + # Allocation-free residual: use doubling workspace buffers as scratch + ensure_lyapunov_doubling_buffers!(workspace) + ๐‚A_tmp = workspace.๐‚A + ๐‚ยน_tmp = workspace.๐‚ยน + โ„’.mul!(๐‚A_tmp, ๐‚, A') + โ„’.mul!(๐‚ยน_tmp, A, ๐‚A_tmp) + โ„’.axpy!(1, C, ๐‚ยน_tmp) + โ„’.axpy!(-1, ๐‚, ๐‚ยน_tmp) - reached_tol = โ„’.norm(A * ๐‚ * A' + C - ๐‚) / โ„’.norm(๐‚) + reached_tol = โ„’.norm(๐‚ยน_tmp) / โ„’.norm(๐‚) # if reached_tol > tol # println("Lyapunov: lyapunov $reached_tol") @@ -122,7 +202,8 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, ::Val{:doubling}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::Float64 = 1e-14)::Tuple{<:AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{<:AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned sparse storage created locally in this method. # Note: workspace is unused for sparse matrices but accepted for API consistency ๐‚ = copy(C) ๐€ = copy(A) @@ -140,7 +221,7 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, if i % 2 == 0 normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -171,7 +252,8 @@ function solve_lyapunov_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat ::Val{:doubling}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::Float64 = 1e-14)::Tuple{<:AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{<:AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned sparse storage created locally in this method. # Note: workspace is unused for sparse matrices but accepted for API consistency ๐‚ = copy(C) ๐€ = copy(A) @@ -192,7 +274,7 @@ function solve_lyapunov_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat if i % 2 == 0 normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -223,7 +305,8 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, ::Val{:doubling}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned dense storage created locally in this method. # Note: workspace is unused for sparse matrices but accepted for API consistency ๐‚ = copy(C) ๐€ = copy(A) @@ -235,7 +318,7 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, iters = max_iter for i in 1:max_iter - # ๐‚ยน .= ๐€ * ๐‚ * ๐€' + ๐‚ + # Sparse A: standard matmul is efficient; Symmetric wrapper lacks optimised sparse dispatch โ„’.mul!(๐‚A, ๐‚, ๐€') โ„’.mul!(๐‚ยน, ๐€, ๐‚A, 1, 1) @@ -248,8 +331,11 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, droptol!(๐€, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + copyto!(๐‚A, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚A) + normdiff = โ„’.norm(๐‚A) + maxnorm = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) + if !isfinite(normdiff) || normdiff / maxnorm < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -260,21 +346,12 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, # ๐‚ = ๐‚ยน end - # โ„’.mul!(๐‚A, ๐‚, A') - # โ„’.mul!(๐‚ยน, A, ๐‚A) - # โ„’.axpy!(1, C, ๐‚ยน) - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # โ„’.axpy!(-1, ๐‚, ๐‚ยน) + โ„’.mul!(๐‚A, ๐‚, A') + โ„’.mul!(๐‚ยน, A, ๐‚A) + โ„’.axpy!(1, C, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚ยน) - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน) / denom - - reached_tol = โ„’.norm(A * ๐‚ * A' + C - ๐‚) / โ„’.norm(๐‚) - - # if reached_tol > tol - # println("Lyapunov: doubling $reached_tol") - # end + reached_tol = โ„’.norm(๐‚ยน) / โ„’.norm(๐‚) return ๐‚, iters, reached_tol # return info on convergence end @@ -287,7 +364,8 @@ function solve_lyapunov_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat ::Val{:doubling}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense buffer workspace.๐‚. # Ensure doubling buffers are allocated ensure_lyapunov_doubling_buffers!(workspace) @@ -307,15 +385,19 @@ function solve_lyapunov_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat iters = max_iter for i in 1:max_iter + # Always use dgemm โ€” dsymm is slower at typical DSGE sizes (n โ‰ค 400) โ„’.mul!(๐‚A, ๐‚, ๐€') โ„’.mul!(๐‚ยน, ๐€, ๐‚A, 1, 1) โ„’.mul!(๐€ยฒ, ๐€, ๐€) copyto!(๐€, ๐€ยฒ) - + if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + copyto!(๐‚A, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚A) + normdiff = โ„’.norm(๐‚A) + maxnorm = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) + if !isfinite(normdiff) || normdiff / maxnorm < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -325,23 +407,14 @@ function solve_lyapunov_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat copyto!(๐‚, ๐‚ยน) end - # โ„’.mul!(๐‚A, ๐‚, A') - # โ„’.mul!(๐‚ยน, A, ๐‚A) - # โ„’.axpy!(1, C, ๐‚ยน) - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # โ„’.axpy!(-1, ๐‚, ๐‚ยน) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน) / denom + โ„’.mul!(๐‚A, ๐‚, A') + โ„’.mul!(๐‚ยน, A, ๐‚A) + โ„’.axpy!(1, C, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚ยน) - reached_tol = โ„’.norm(A * ๐‚ * A' + C - ๐‚) / โ„’.norm(๐‚) + reached_tol = โ„’.norm(๐‚ยน) / โ„’.norm(๐‚) - # if reached_tol > tol - # println("Lyapunov: doubling $reached_tol") - # end - - return copy(๐‚), iters, reached_tol # return info on convergence + return ๐‚, iters, reached_tol # return info on convergence end @@ -352,48 +425,72 @@ function solve_lyapunov_equation(A::AbstractMatrix{T}, ::Val{:bicgstab}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat - # Ensure Krylov buffers and bicgstab solver are allocated - ensure_lyapunov_bicgstab_solver!(workspace) + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense Krylov buffer workspace.๐—. - # Use workspaces - tmpฬ„ = workspace.tmpฬ„ - ๐— = workspace.๐— - b = workspace.b - - function lyapunov!(sol,๐ฑ) - copyto!(๐—, ๐ฑ) - โ„’.mul!(tmpฬ„, ๐—, A') - โ„’.mul!(๐—, A, tmpฬ„, -1, 1) - copyto!(sol, ๐—) - end + if _is_approx_symmetric(C) + # vech-space Krylov: solve for n(n+1)/2 unique elements only + ensure_lyapunov_krylov_vech_solver!(workspace, :bicgstab) + tmpฬ„ = workspace.tmpฬ„ + ๐— = workspace.๐— + n = size(A, 1) + n_vech = n * (n + 1) รท 2 + b_vech = workspace.b_vech + + function lyapunov_vech_bicgstab!(sol, ๐ฑ) + fill_symmetric_from_vech!(๐—, ๐ฑ) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(๐—, A, tmpฬ„, -1, 1) + vech!(sol, ๐—) + end - lyapunov = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov!) + lyapunov_op = LinearOperators.LinearOperator(Float64, n_vech, n_vech, true, true, lyapunov_vech_bicgstab!) - # Use vectorized C in workspace - copyto!(b, vec(C)) - - # Use pre-allocated solver - Krylov.bicgstab!(workspace.bicgstab_workspace, lyapunov, b, rtol = tol, atol = tol) + vech!(b_vech, C) - copyto!(๐—, workspace.bicgstab_workspace.x) + Krylov.bicgstab!(workspace.bicgstab_vech, lyapunov_op, b_vech, rtol = tol.rtol, atol = tol.atol) - # โ„’.mul!(tmpฬ„, A, ๐— * A') - # โ„’.axpy!(1, C, tmpฬ„) + fill_symmetric_from_vech!(๐—, workspace.bicgstab_vech.x) - # denom = max(โ„’.norm(๐—), โ„’.norm(tmpฬ„)) + # Allocation-free residual: reuse tmpฬ„ for intermediate, ๐— is the solution + ensure_lyapunov_doubling_buffers!(workspace) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(workspace.๐‚ยน, A, tmpฬ„) + โ„’.axpy!(1, C, workspace.๐‚ยน) + โ„’.axpy!(-1, ๐—, workspace.๐‚ยน) + reached_tol = โ„’.norm(workspace.๐‚ยน) / โ„’.norm(๐—) - # โ„’.axpy!(-1, ๐—, tmpฬ„) + return ๐—, workspace.bicgstab_vech.stats.niter, reached_tol + else + # Standard full-space Krylov + ensure_lyapunov_krylov_solver!(workspace, :bicgstab) + tmpฬ„ = workspace.tmpฬ„ + ๐— = workspace.๐— + b = workspace.b + + function lyapunov_bicgstab!(sol,๐ฑ) + copyto!(๐—, ๐ฑ) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(๐—, A, tmpฬ„, -1, 1) + copyto!(sol, ๐—) + end - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(tmpฬ„) / denom + lyapunov_op = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov_bicgstab!) - reached_tol = โ„’.norm(A * ๐— * A' + C - ๐—) / โ„’.norm(๐—) + copyto!(b, vec(C)) + Krylov.bicgstab!(workspace.bicgstab, lyapunov_op, b, rtol = tol.rtol, atol = tol.atol) + copyto!(๐—, workspace.bicgstab.x) - # if reached_tol > tol - # println("Lyapunov: bicgstab $reached_tol") - # end + # Allocation-free residual + ensure_lyapunov_doubling_buffers!(workspace) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(workspace.๐‚ยน, A, tmpฬ„) + โ„’.axpy!(1, C, workspace.๐‚ยน) + โ„’.axpy!(-1, ๐—, workspace.๐‚ยน) + reached_tol = โ„’.norm(workspace.๐‚ยน) / โ„’.norm(๐—) - return copy(๐—), workspace.bicgstab_workspace.stats.niter, reached_tol + return ๐—, workspace.bicgstab.stats.niter, reached_tol + end end @@ -402,50 +499,72 @@ function solve_lyapunov_equation(A::AbstractMatrix{T}, ::Val{:gmres}, workspace::lyapunov_workspace; # timer::TimerOutput = TimerOutput(), - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat - # Ensure Krylov buffers and gmres solver are allocated - ensure_lyapunov_gmres_solver!(workspace) + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense Krylov buffer workspace.๐—. - # Use workspaces - tmpฬ„ = workspace.tmpฬ„ - ๐— = workspace.๐— - b = workspace.b - - function lyapunov!(sol,๐ฑ) - copyto!(๐—, ๐ฑ) - # ๐— = @view reshape(๐ฑ, size(๐—)) - โ„’.mul!(tmpฬ„, ๐—, A') - โ„’.mul!(๐—, A, tmpฬ„, -1, 1) - copyto!(sol, ๐—) - # sol = @view reshape(๐—, size(sol)) - end + if _is_approx_symmetric(C) + # vech-space Krylov: solve for n(n+1)/2 unique elements only + ensure_lyapunov_krylov_vech_solver!(workspace, :gmres) + tmpฬ„ = workspace.tmpฬ„ + ๐— = workspace.๐— + n = size(A, 1) + n_vech = n * (n + 1) รท 2 + b_vech = workspace.b_vech + + function lyapunov_vech_gmres!(sol, ๐ฑ) + fill_symmetric_from_vech!(๐—, ๐ฑ) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(๐—, A, tmpฬ„, -1, 1) + vech!(sol, ๐—) + end - lyapunov = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov!) + lyapunov_op = LinearOperators.LinearOperator(Float64, n_vech, n_vech, true, true, lyapunov_vech_gmres!) - # Use vectorized C in workspace - copyto!(b, vec(C)) - - # Use pre-allocated solver - Krylov.gmres!(workspace.gmres_workspace, lyapunov, b, rtol = tol, atol = tol) + vech!(b_vech, C) - copyto!(๐—, workspace.gmres_workspace.x) + Krylov.gmres!(workspace.gmres_vech, lyapunov_op, b_vech, rtol = tol.rtol, atol = tol.atol) - # โ„’.mul!(tmpฬ„, A, ๐— * A') - # โ„’.axpy!(1, C, tmpฬ„) + fill_symmetric_from_vech!(๐—, workspace.gmres_vech.x) - # denom = max(โ„’.norm(๐—), โ„’.norm(tmpฬ„)) + # Allocation-free residual + ensure_lyapunov_doubling_buffers!(workspace) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(workspace.๐‚ยน, A, tmpฬ„) + โ„’.axpy!(1, C, workspace.๐‚ยน) + โ„’.axpy!(-1, ๐—, workspace.๐‚ยน) + reached_tol = โ„’.norm(workspace.๐‚ยน) / โ„’.norm(๐—) - # โ„’.axpy!(-1, ๐—, tmpฬ„) + return ๐—, workspace.gmres_vech.stats.niter, reached_tol + else + # Standard full-space Krylov + ensure_lyapunov_krylov_solver!(workspace, :gmres) + tmpฬ„ = workspace.tmpฬ„ + ๐— = workspace.๐— + b = workspace.b + + function lyapunov_gmres!(sol,๐ฑ) + copyto!(๐—, ๐ฑ) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(๐—, A, tmpฬ„, -1, 1) + copyto!(sol, ๐—) + end - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(tmpฬ„) / denom + lyapunov_op = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov_gmres!) - reached_tol = โ„’.norm(A * ๐— * A' + C - ๐—) / โ„’.norm(๐—) + copyto!(b, vec(C)) + Krylov.gmres!(workspace.gmres, lyapunov_op, b, rtol = tol.rtol, atol = tol.atol) + copyto!(๐—, workspace.gmres.x) - # if reached_tol > tol - # println("Lyapunov: gmres $reached_tol") - # end + # Allocation-free residual + ensure_lyapunov_doubling_buffers!(workspace) + โ„’.mul!(tmpฬ„, ๐—, A') + โ„’.mul!(workspace.๐‚ยน, A, tmpฬ„) + โ„’.axpy!(1, C, workspace.๐‚ยน) + โ„’.axpy!(-1, ๐—, workspace.๐‚ยน) + reached_tol = โ„’.norm(workspace.๐‚ยน) / โ„’.norm(๐—) - return copy(๐—), workspace.gmres_workspace.stats.niter, reached_tol + return ๐—, workspace.gmres.stats.niter, reached_tol + end end diff --git a/src/algorithms/nonlinear_solver.jl b/src/algorithms/nonlinear_solver.jl index 0878f87c2..f48024f04 100644 --- a/src/algorithms/nonlinear_solver.jl +++ b/src/algorithms/nonlinear_solver.jl @@ -12,9 +12,9 @@ function levenberg_marquardt( )::Tuple{Vector{T}, Tuple{Int, Int, T, T}} where {T <: AbstractFloat} # issues with optimization: https://www.gurobi.com/documentation/8.1/refman/numerics_gurobi_guidelines.html - xtol = tol.NSSS_xtol - ftol = tol.NSSS_ftol - rel_xtol = tol.NSSS_rel_xtol + xtol = tol.nsss.xtol + ftol = tol.nsss.ftol + rel_xtol = tol.nsss.rel_xtol iterations = 250 @@ -181,10 +181,10 @@ function levenberg_marquardt( # sol_cache.A = X sol_cache.A = โˆ‡ฬ‚ sol_cache.b = guess_update - ๐’ฎ.solve!(sol_cache) + sol = ๐’ฎ.solve!(sol_cache) copy!(guess_update, sol_cache.u) - if !isfinite(sum(guess_update)) + if !(๐’ฎ.SciMLBase.successful_retcode(sol.retcode) || sol.retcode == ๐’ฎ.SciMLBase.ReturnCode.Default || isfinite(sum(guess_update))) largest_relative_step = 1.0 largest_residual = 1.0 break @@ -421,9 +421,9 @@ function newton( )::Tuple{Vector{T}, Tuple{Int, Int, T, T}} where {T <: AbstractFloat} # issues with optimization: https://www.gurobi.com/documentation/8.1/refman/numerics_gurobi_guidelines.html - xtol = tol.NSSS_xtol - ftol = tol.NSSS_ftol - rel_xtol = tol.NSSS_rel_xtol + xtol = tol.nsss.xtol + ftol = tol.nsss.ftol + rel_xtol = tol.nsss.rel_xtol iterations = 250 transformation_level = 0 # parameters.transformation_level @@ -484,25 +484,25 @@ function newton( new_residuals_norm = โ„’.norm(new_residuals) - if โˆ‡ isa SparseMatrixCSC - sol_cache.A = โˆ‡ - sol_cache.b = new_residuals - ๐’ฎ.solve!(sol_cache) - guess_update .= sol_cache.u - new_residuals .= guess_update - else - factโˆ‡ = โ„’.lu!(โˆ‡, check = false) - try - if !โ„’.issuccess(factโˆ‡) - factโˆ‡ = โ„’.qr(โˆ‡, โ„’.ColumnNorm()) - end - โ„’.ldiv!(factโˆ‡, new_residuals) - catch - rel_xtol_reached = typemax(T) - new_residuals_norm = typemax(T) - break - end + # sol_cache.A = โˆ‡ + # copy!(sol_cache.A, โˆ‡) + sol_cache.A = โˆ‡ + # sol_cache.A = sol_cache.alg isa ๐’ฎ.FastLUFactorization ? copy(โˆ‡) : โˆ‡ + sol_cache.b = new_residuals + sol = ๐’ฎ.solve!(sol_cache) + if sol.retcode != ๐’ฎ.SciMLBase.ReturnCode.Default && !๐’ฎ.SciMLBase.successful_retcode(sol.retcode) + rel_xtol_reached = typemax(T) + new_residuals_norm = typemax(T) + break end + guess_update .= sol_cache.u + if has_nonfinite(guess_update) + rel_xtol_reached = typemax(T) + new_residuals_norm = typemax(T) + break + end + # new_residuals .= guess_update + copy!(new_residuals, guess_update) guess_update_norm = โ„’.norm(new_residuals) โ„’.axpy!(-1, new_residuals, new_guess) @@ -534,29 +534,31 @@ function newton( # end # sol_cache.A = โˆ‡ + # sol_cache.b = new_residuals # ๐’ฎ.solve!(sol_cache) # copy!(guess_update, sol_cache.u) - if โˆ‡ isa SparseMatrixCSC - sol_cache.A = โˆ‡ - sol_cache.b = new_residuals - ๐’ฎ.solve!(sol_cache) - guess_update .= sol_cache.u - new_residuals .= guess_update - else - factโˆ‡ = โ„’.lu!(โˆ‡, check = false) - try - if !โ„’.issuccess(factโˆ‡) - factโˆ‡ = โ„’.qr(โˆ‡, โ„’.ColumnNorm()) - end - โ„’.ldiv!(factโˆ‡, new_residuals) - catch - rel_xtol_reached = typemax(T) - new_residuals_norm = typemax(T) - break - end + # copy!(sol_cache.A, โˆ‡) + sol_cache.A = โˆ‡ + # sol_cache.A = sol_cache.alg isa ๐’ฎ.FastLUFactorization ? copy(โˆ‡) : โˆ‡ + sol_cache.b = new_residuals + sol = ๐’ฎ.solve!(sol_cache) + if sol.retcode != ๐’ฎ.SciMLBase.ReturnCode.Default && !๐’ฎ.SciMLBase.successful_retcode(sol.retcode) + rel_xtol_reached = typemax(T) + new_residuals_norm = typemax(T) + break + end + # guess_update .= sol_cache.u + copy!(guess_update, sol_cache.u) + + if has_nonfinite(guess_update) + rel_xtol_reached = typemax(T) + new_residuals_norm = typemax(T) + break end + # new_residuals .= guess_update + copy!(new_residuals, guess_update) guess_update_norm = โ„’.norm(new_residuals) โ„’.axpy!(-1, new_residuals, new_guess) diff --git a/src/algorithms/quadratic_matrix_equation.jl b/src/algorithms/quadratic_matrix_equation.jl index da4a88c3a..9c7f341c0 100644 --- a/src/algorithms/quadratic_matrix_equation.jl +++ b/src/algorithms/quadratic_matrix_equation.jl @@ -12,43 +12,79 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, B::AbstractMatrix{R}, C::AbstractMatrix{R}, constants::constants, - workspace::qme_workspace{R,S}; + workspaces::workspaces, + cache::caches; initial_guess::AbstractMatrix{R} = zeros(0,0), - quadratic_matrix_equation_algorithm::Symbol = :schur, - tol::AbstractFloat = 1e-14, - acceptance_tol::AbstractFloat = 1e-8, - verbose::Bool = false) where {R <: Real, S <: Real} + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + use_fastlapack_schur::Bool = true, + use_fastlapack_lu::Bool = true, + tol::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-8, + acceptance_tol = 1e-8), + verbose::Bool = false, + caching::Bool = true)::Tuple{Matrix{R}, Bool} where {R <: AbstractFloat} T = constants.post_model_macro + n = T.nVars - T.nPresent_only + nPfm = T.nPast_not_future_and_mixed + + qme_ws = ensure_qme_doubling_workspace!(workspaces, n) + ensure_schur_workspace!(workspaces, + n, + T.nMixed, + nPfm, + T.nFuture_not_past_and_mixed) + + initial_guess_acceptance_tol = tol.initial_guess_acceptance_tol + acceptance_tol = tol.acceptance_tol - if length(initial_guess) > 0 + if caching && length(initial_guess) > 0 X = initial_guess + Xยฒ = qme_ws.temp3 - AXX = A * X^2 - - AXXnorm = max(โ„’.norm(AXX), โ„’.norm(C)) + # Compute residual: A*Xยฒ + B*X + C + # Xยฒ into temporary buffer + โ„’.mul!(Xยฒ, X, X) + # A*Xยฒ into AXX buffer + โ„’.mul!(qme_ws.AXX, A, Xยฒ) - โ„’.mul!(AXX, B, X, 1, 1) - - โ„’.axpy!(1, C, AXX) + AXXnorm = min(โ„’.norm(qme_ws.AXX), โ„’.norm(C)) - reached_tol = โ„’.norm(AXX) / AXXnorm + # AXX += B*X + โ„’.mul!(qme_ws.AXX, B, X, 1, 1) + # AXX += C + โ„’.axpy!(1, C, qme_ws.AXX) + + reached_tol = โ„’.norm(qme_ws.AXX) / AXXnorm - if reached_tol < (acceptance_tol * length(initial_guess) / 1e6)# 1e-12 is too large eps is too small; if the low tol is used it can be that a small change in the parameters still yields an acceptable solution but as a better tol can be reached it is actually not accurate + if reached_tol < (initial_guess_acceptance_tol * length(initial_guess) / 1e6)# 1e-12 is too large eps is too small; if the low tol is used it can be that a small change in the parameters still yields an acceptable solution but as a better tol can be reached it is actually not accurate if verbose println("Quadratic matrix equation solver previous solution has tolerance: $reached_tol") end - return initial_guess, true + _existing_sol = cache.qme_solution + if _existing_sol isa Matrix{R} && size(_existing_sol) == size(initial_guess) + copyto!(_existing_sol, initial_guess) + return _existing_sol, true + else + new_sol = Matrix{R}(initial_guess) + cache.qme_solution = new_sol + return new_sol, true + end end end sol, iterations, reached_tol = solve_quadratic_matrix_equation(A, B, C, Val(quadratic_matrix_equation_algorithm), constants, - workspace; + workspaces, + cache; initial_guess = initial_guess, + use_fastlapack_schur = use_fastlapack_schur, + use_fastlapack_lu = use_fastlapack_lu, tol = tol, # timer = timer, - verbose = verbose) + verbose = verbose, + caching = caching) if verbose println("Quadratic matrix equation solver: $quadratic_matrix_equation_algorithm - converged: $(reached_tol < acceptance_tol) in $iterations iterations to tolerance: $reached_tol") end @@ -57,22 +93,30 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, sol, iterations, reached_tol = solve_quadratic_matrix_equation(A, B, C, Val(:schur), constants, - workspace; + workspaces, + cache; initial_guess = initial_guess, + use_fastlapack_schur = use_fastlapack_schur, + use_fastlapack_lu = use_fastlapack_lu, tol = tol, # timer = timer, - verbose = verbose) + verbose = verbose, + caching = caching) if verbose println("Quadratic matrix equation solver: schur - converged: $(reached_tol < acceptance_tol) in $iterations iterations to tolerance: $reached_tol") end else quadratic_matrix_equation_algorithm โ‰  :doubling sol, iterations, reached_tol = solve_quadratic_matrix_equation(A, B, C, Val(:doubling), constants, - workspace; + workspaces, + cache; initial_guess = initial_guess, + use_fastlapack_schur = use_fastlapack_schur, + use_fastlapack_lu = use_fastlapack_lu, tol = tol, # timer = timer, - verbose = verbose) + verbose = verbose, + caching = caching) if verbose println("Quadratic matrix equation solver: doubling - converged: $(reached_tol < acceptance_tol) in $iterations iterations to tolerance: $reached_tol") end end @@ -88,131 +132,190 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, C::AbstractMatrix{R}, ::Val{:schur}, constants::constants, - workspace::qme_workspace; + workspaces::workspaces, + cache::caches; initial_guess::AbstractMatrix{R} = zeros(0,0), - tol::AbstractFloat = 1e-14, + use_fastlapack_schur::Bool = true, + use_fastlapack_lu::Bool = true, + tol::SolverTolerances = SolverTolerances(), # timer::TimerOutput = TimerOutput(), - verbose::Bool = false)::Tuple{Matrix{R}, Int64, R} where R <: AbstractFloat - # Use cached identity matrix from workspace (Diagonal{Bool} supports indexing) + verbose::Bool = false, + caching::Bool = true)::Tuple{Matrix{R}, Int64, R} where R <: AbstractFloat + T = constants.post_model_macro - # @timeit_debug timer "Prepare indice" begin - I_nPast = workspace.I_nPast - - comb = union(T.future_not_past_and_mixed_idx, T.past_not_future_idx) - sort!(comb) - - future_not_past_and_mixed_in_comb = indexin(T.future_not_past_and_mixed_idx, comb) - past_not_future_and_mixed_in_comb = indexin(T.past_not_future_and_mixed_idx, comb) - indices_past_not_future_in_comb = indexin(T.past_not_future_idx, comb) - - # end # timeit_debug - # @timeit_debug timer "Assemble matrices" begin - - Aฬƒโ‚Š = A[:,future_not_past_and_mixed_in_comb] + idx_constants = constants.post_complete_parameters - Aฬƒโ‚‹ = C[:,past_not_future_and_mixed_in_comb] + # Ensure schur workspace is properly sized + n = T.nVars - T.nPresent_only + nMixed = T.nMixed + nPfm = T.nPast_not_future_and_mixed + nFnpm = T.nFuture_not_past_and_mixed - Aฬƒโ‚€โ‚Š = B[:,future_not_past_and_mixed_in_comb] - - Aฬƒโ‚€โ‚‹ = B[:,indices_past_not_future_in_comb] * I_nPast[T.not_mixed_in_past_idx,:] - - Zโ‚Š = zeros(T.nMixed, T.nFuture_not_past_and_mixed) - Iโ‚Š = โ„’.I(T.nFuture_not_past_and_mixed)[T.mixed_in_future_idx,:] + schur_ws_local = ensure_schur_workspace!(workspaces, n, nMixed, nPfm, nFnpm) - Zโ‚‹ = zeros(T.nMixed,T.nPast_not_future_and_mixed) - Iโ‚‹ = I_nPast[T.mixed_in_past_idx,:] + # Use cached indices from constants instead of recomputing + future_not_past_and_mixed_in_comb = idx_constants.future_not_past_and_mixed_in_comb + past_not_future_and_mixed_in_comb = idx_constants.past_not_future_and_mixed_in_comb + indices_past_not_future_in_comb = idx_constants.indices_past_not_future_in_comb - D = vcat(hcat(Aฬƒโ‚€โ‚‹, Aฬƒโ‚Š), hcat(Iโ‚‹, Zโ‚Š)) + # Use views for read-only slices + รƒโ‚Š_view = @view A[:, future_not_past_and_mixed_in_comb] - โ„’.rmul!(Aฬƒโ‚‹,-1) - โ„’.rmul!(Aฬƒโ‚€โ‚Š,-1) - E = vcat(hcat(Aฬƒโ‚‹,Aฬƒโ‚€โ‚Š), hcat(Zโ‚‹, Iโ‚Š)) + # Copy C and B slices that need negation into workspace buffers + copyto!(schur_ws_local.รƒโ‚‹, @view C[:, past_not_future_and_mixed_in_comb]) + copyto!(schur_ws_local.รƒโ‚€โ‚Š, @view B[:, future_not_past_and_mixed_in_comb]) - # end # timeit_debug - # @timeit_debug timer "Schur decomposition" begin - - # this is the companion form and by itself the linearisation of the matrix polynomial used in the linear time iteration method. see: https://opus4.kobv.de/opus4-matheon/files/209/240.pdf - schdcmp = try - โ„’.schur!(D, E) - catch + # Compute รƒโ‚€โ‚‹ = B[:,indices_past_not_future_in_comb] * I_nPast[not_mixed_in_past_idx,:] + # Use cached constant matrix for I_nPast_not_mixed + โ„’.mul!(schur_ws_local.รƒโ‚€โ‚‹, @view(B[:, indices_past_not_future_in_comb]), idx_constants.I_nPast_not_mixed) + + # Use cached constant matrices for zeros and identity blocks + Zโ‚Š = idx_constants.schur_Zโ‚Š + Iโ‚Š = idx_constants.schur_Iโ‚Š + Zโ‚‹ = idx_constants.schur_Zโ‚‹ + Iโ‚‹ = idx_constants.schur_Iโ‚‹ + + # Assemble D matrix in-place: D = [[รƒโ‚€โ‚‹ รƒโ‚Š], [Iโ‚‹ Zโ‚Š]] + D = schur_ws_local.D + # Top-left block: รƒโ‚€โ‚‹ + copyto!(view(D, 1:n, 1:nPfm), schur_ws_local.รƒโ‚€โ‚‹) + # Top-right block: รƒโ‚Š + copyto!(view(D, 1:n, nPfm+1:nPfm+nFnpm), รƒโ‚Š_view) + # Bottom-left block: Iโ‚‹ + copyto!(view(D, n+1:n+nMixed, 1:nPfm), Iโ‚‹) + # Bottom-right block: Zโ‚Š + copyto!(view(D, n+1:n+nMixed, nPfm+1:nPfm+nFnpm), Zโ‚Š) + + # Negate รƒโ‚‹ and รƒโ‚€โ‚Š for E matrix + โ„’.rmul!(schur_ws_local.รƒโ‚‹, -1) + โ„’.rmul!(schur_ws_local.รƒโ‚€โ‚Š, -1) + + # Assemble E matrix in-place: E = [[รƒโ‚‹ รƒโ‚€โ‚Š], [Zโ‚‹ Iโ‚Š]] + E = schur_ws_local.E + # Top-left block: รƒโ‚‹ (already negated) + copyto!(view(E, 1:n, 1:nPfm), schur_ws_local.รƒโ‚‹) + # Top-right block: รƒโ‚€โ‚Š (already negated) + copyto!(view(E, 1:n, nPfm+1:nPfm+nFnpm), schur_ws_local.รƒโ‚€โ‚Š) + # Bottom-left block: Zโ‚‹ + copyto!(view(E, n+1:n+nMixed, 1:nPfm), Zโ‚‹) + # Bottom-right block: Iโ‚Š + copyto!(view(E, n+1:n+nMixed, nPfm+1:nPfm+nFnpm), Iโ‚Š) + + schur_ws_local.fast_qz_ws, + schur_ws_local.fast_qz_dims, + schdcmp, + schur_ok = factorize_generalized_schur!(D, + E, + schur_ws_local.fast_qz_ws, + schur_ws_local.fast_qz_dims, + schur_ws_local.eigenselect; + use_fastlapack_schur = use_fastlapack_schur) + + if !schur_ok if verbose println("Quadratic matrix equation solver: schur - converged: false") end return A, 0, 1.0 end - eigenselect = abs.(schdcmp.ฮฒ ./ schdcmp.ฮฑ) .< 1 - - # end # timeit_debug - # @timeit_debug timer "Reorder Schur decomposition" begin - - try - โ„’.ordschur!(schdcmp, eigenselect) - catch + # Extract blocks from reordered Schur form (need owned copies for lu!) + copyto!(schur_ws_local.Zโ‚โ‚, @view schdcmp.Z[1:nPfm, 1:nPfm]) + copyto!(schur_ws_local.Zโ‚‚โ‚, @view schdcmp.Z[nPfm+1:end, 1:nPfm]) + # Zโ‚โ‚ can be a view for matrix multiplication, but LU factorization needs an owned copy. + Zโ‚โ‚ = @view schdcmp.Z[1:nPfm, 1:nPfm] + + copyto!(schur_ws_local.Sโ‚โ‚, @view schdcmp.S[1:nPfm, 1:nPfm]) + copyto!(schur_ws_local.Tโ‚โ‚, @view schdcmp.T[1:nPfm, 1:nPfm]) + + schur_ws_local.fast_lu_ws_z11, + schur_ws_local.fast_lu_dims_z11, + solved_Zโ‚โ‚, + แบโ‚โ‚ = factorize_lu!(schur_ws_local.Zโ‚โ‚, + schur_ws_local.fast_lu_ws_z11, + schur_ws_local.fast_lu_dims_z11; + use_fastlapack_lu = use_fastlapack_lu) + + if !solved_Zโ‚โ‚ if verbose println("Quadratic matrix equation solver: schur - converged: false") end return A, 0, 1.0 end - # end # timeit_debug - # @timeit_debug timer "Postprocess" begin - - Zโ‚‚โ‚ = schdcmp.Z[T.nPast_not_future_and_mixed+1:end, 1:T.nPast_not_future_and_mixed] - Zโ‚โ‚ = schdcmp.Z[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] - - Sโ‚โ‚ = schdcmp.S[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] - Tโ‚โ‚ = schdcmp.T[1:T.nPast_not_future_and_mixed, 1:T.nPast_not_future_and_mixed] - - # @timeit_debug timer "Matrix inversions" begin - - Zฬ‚โ‚โ‚ = โ„’.lu(Zโ‚โ‚, check = false) + # LU factorization of Sโ‚โ‚ (mutating - overwrites workspace buffer) + schur_ws_local.fast_lu_ws_s11, + schur_ws_local.fast_lu_dims_s11, + solved_Sโ‚โ‚, + ลœโ‚โ‚ = factorize_lu!(schur_ws_local.Sโ‚โ‚, + schur_ws_local.fast_lu_ws_s11, + schur_ws_local.fast_lu_dims_s11; + use_fastlapack_lu = use_fastlapack_lu) - if !โ„’.issuccess(Zฬ‚โ‚โ‚) + if !solved_Sโ‚โ‚ if verbose println("Quadratic matrix equation solver: schur - converged: false") end return A, 0, 1.0 end - Sฬ‚โ‚โ‚ = โ„’.lu!(Sโ‚โ‚, check = false) + # Compute D = Zโ‚‚โ‚ / แบโ‚โ‚ (overwrites Zโ‚‚โ‚ buffer) + solve_lu_right!(schur_ws_local.Zโ‚โ‚, + schur_ws_local.Zโ‚‚โ‚, + schur_ws_local.fast_lu_ws_z11, + แบโ‚โ‚, + schur_ws_local.fast_lu_rhs_t_z21; + use_fastlapack_lu = use_fastlapack_lu) - if !โ„’.issuccess(Sฬ‚โ‚โ‚) - if verbose println("Quadratic matrix equation solver: schur - converged: false") end - return A, 0, 1.0 + # Compute L = Zโ‚โ‚ * (ลœโ‚โ‚ \ Tโ‚โ‚) / แบโ‚โ‚ + # First: Tโ‚โ‚ โ† ลœโ‚โ‚ \ Tโ‚โ‚ (overwrites Tโ‚โ‚ buffer) + solve_lu_left!(schur_ws_local.Sโ‚โ‚, + schur_ws_local.Tโ‚โ‚, + schur_ws_local.fast_lu_ws_s11, + ลœโ‚โ‚; + use_fastlapack_lu = use_fastlapack_lu) + # Then: Sโ‚โ‚ โ† Zโ‚โ‚ * Tโ‚โ‚ (reuse Sโ‚โ‚ buffer) + โ„’.mul!(schur_ws_local.Sโ‚โ‚, Zโ‚โ‚, schur_ws_local.Tโ‚โ‚) + # Finally: Sโ‚โ‚ โ† Sโ‚โ‚ / แบโ‚โ‚ (overwrites Sโ‚โ‚ buffer) + solve_lu_right!(schur_ws_local.Zโ‚โ‚, + schur_ws_local.Sโ‚โ‚, + schur_ws_local.fast_lu_ws_z11, + แบโ‚โ‚, + schur_ws_local.fast_lu_rhs_t_s11; + use_fastlapack_lu = use_fastlapack_lu) + + # Assemble sol = vcat(L[not_mixed_in_past_idx,:], D) in-place + sol = schur_ws_local.sol + copyto!(view(sol, 1:length(T.not_mixed_in_past_idx), :), + @view schur_ws_local.Sโ‚โ‚[T.not_mixed_in_past_idx, :]) + copyto!(view(sol, length(T.not_mixed_in_past_idx)+1:size(sol,1), :), + schur_ws_local.Zโ‚‚โ‚) + + # Final reordering: X = sol[dynamic_order,:] * Ir[past_not_future_and_mixed_in_comb,:] + # n == n_comb (= nFnpm + nPfm - nMixed) so the result is (n, n), same as doubling. + # Prefer cache-backed storage to avoid extra allocations. + X = if caching + _existing_sol = cache.qme_solution + if _existing_sol isa Matrix{R} && size(_existing_sol) == (n, n) + _existing_sol + else + cache.qme_solution = zeros(R, n, n) + end + else + zeros(R, n, n) end - # end # timeit_debug - # @timeit_debug timer "Matrix divisions" begin - - # D = Zโ‚‚โ‚ / Zฬ‚โ‚โ‚ - โ„’.rdiv!(Zโ‚‚โ‚, Zฬ‚โ‚โ‚) - D = Zโ‚‚โ‚ + โ„’.mul!(X, @view(sol[T.dynamic_order, :]), idx_constants.Ir_past_selector) - # L = Zโ‚โ‚ * (Sฬ‚โ‚โ‚ \ Tโ‚โ‚) / Zฬ‚โ‚โ‚ - โ„’.ldiv!(Sฬ‚โ‚โ‚, Tโ‚โ‚) - โ„’.mul!(Sโ‚โ‚, Zโ‚โ‚, Tโ‚โ‚) - โ„’.rdiv!(Sโ‚โ‚, Zฬ‚โ‚โ‚) - L = Sโ‚โ‚ - - sol = vcat(L[T.not_mixed_in_past_idx,:], D) - - # end # timeit_debug - # end # timeit_debug - - X = sol[T.dynamic_order,:] * โ„’.I(length(comb))[past_not_future_and_mixed_in_comb,:] - - iter = 0 - - AXX = A * X^2 + # Compute residual: A*Xยฒ + B*X + C + # Xยฒ into temp_X2 buffer + โ„’.mul!(schur_ws_local.temp_X2, X, X) + # A*Xยฒ into AXX buffer + โ„’.mul!(schur_ws_local.AXX, A, schur_ws_local.temp_X2) - AXXnorm = max(โ„’.norm(AXX), โ„’.norm(C)) + AXXnorm = min(โ„’.norm(schur_ws_local.AXX), โ„’.norm(C)) - โ„’.mul!(AXX, B, X, 1, 1) - - โ„’.axpy!(1, C, AXX) + # AXX += B*X + โ„’.mul!(schur_ws_local.AXX, B, X, 1, 1) + # AXX += C + โ„’.axpy!(1, C, schur_ws_local.AXX) - reached_tol = โ„’.norm(AXX) / AXXnorm + reached_tol = โ„’.norm(schur_ws_local.AXX) / AXXnorm - # if reached_tol > tol - # println("QME: schur $reached_tol") - # end - - return X, iter, reached_tol # schur can fail + return X, 0, reached_tol end @@ -221,13 +324,19 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, C::AbstractMatrix{R}, ::Val{:doubling}, constants::constants, - workspace::qme_workspace{R,S}; + workspaces::workspaces, + cache::caches; initial_guess::AbstractMatrix{R} = zeros(0,0), - tol::AbstractFloat = 1e-14, + use_fastlapack_schur::Bool = true, + use_fastlapack_lu::Bool = true, + tol::SolverTolerances = SolverTolerances(), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - max_iter::Int = 100)::Tuple{Matrix{R}, Int64, R} where {R <: AbstractFloat, S <: Real} + max_iter::Int = 100, + caching::Bool = true)::Tuple{Matrix{R}, Int64, R} where {R <: AbstractFloat} T = constants.post_model_macro + idx_constants = ensure_first_order_constants!(constants) + workspace = ensure_qme_doubling_workspace!(workspaces, size(A, 1)) # Johannes Huber, Alexander Meyer-Gohde, Johanna Saecker (2024). Solving Linear DSGE Models with Structure Preserving Doubling Methods. # https://www.imfs-frankfurt.de/forschung/imfs-working-papers/details.html?tx_mmpublications_publicationsdetail%5Bcontroller%5D=Publication&tx_mmpublications_publicationsdetail%5Bpublication%5D=461&cHash=f53244e0345a27419a9d40a3af98c02f # https://arxiv.org/abs/2212.09491 @@ -263,15 +372,23 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, โ„’.mul!(Bฬ„, A, initial_guess, 1, 1) - Bฬ‚ = โ„’.lu!(Bฬ„, check = false) - - if !โ„’.issuccess(Bฬ‚) + workspace.fast_lu_ws_qme_a, + workspace.fast_lu_dims_qme_a, + solved_B, + Bฬ‚ = factorize_lu!(Bฬ„, + workspace.fast_lu_ws_qme_a, + workspace.fast_lu_dims_qme_a; + use_fastlapack_lu = use_fastlapack_lu) + + if !solved_B return A, 0, 1.0 end # Compute initial values X, Y, E, F - โ„’.ldiv!(E, Bฬ‚, C) - โ„’.ldiv!(F, Bฬ‚, A) + solve_lu_left!(Bฬ„, E, workspace.fast_lu_ws_qme_a, Bฬ‚; + use_fastlapack_lu = use_fastlapack_lu) + solve_lu_left!(Bฬ„, F, workspace.fast_lu_ws_qme_a, Bฬ‚; + use_fastlapack_lu = use_fastlapack_lu) # X = -E - initial_guess (in-place) copy!(X, E) @@ -283,7 +400,7 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # end # timeit_debug # @timeit_debug timer "Prellocate" begin - II = workspace.I_n # Pre-computed identity matrix reference + II = idx_constants.I_n # Pre-computed identity matrix reference Xtol = 1.0 Ytol = 1.0 @@ -305,9 +422,15 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # end # timeit_debug # @timeit_debug timer "Invert EI" begin - fEI = โ„’.lu!(temp1, check = false) + workspace.fast_lu_ws_qme_a, + workspace.fast_lu_dims_qme_a, + solved_EI, + fEI = factorize_lu!(temp1, + workspace.fast_lu_ws_qme_a, + workspace.fast_lu_dims_qme_a; + use_fastlapack_lu = use_fastlapack_lu) - if !โ„’.issuccess(fEI) + if !solved_EI return A, iter, 1.0 end @@ -315,7 +438,9 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # @timeit_debug timer "Compute E" begin # Compute E = E * EI * E - โ„’.ldiv!(temp3, fEI, E) + copyto!(temp3, E) + solve_lu_left!(temp1, temp3, workspace.fast_lu_ws_qme_a, fEI; + use_fastlapack_lu = use_fastlapack_lu) โ„’.mul!(E_new, E, temp3) # E_new = E / fEI * E @@ -332,9 +457,15 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # end # timeit_debug # @timeit_debug timer "Invert FI" begin - fFI = โ„’.lu!(temp2, check = false) + workspace.fast_lu_ws_qme_b, + workspace.fast_lu_dims_qme_b, + solved_FI, + fFI = factorize_lu!(temp2, + workspace.fast_lu_ws_qme_b, + workspace.fast_lu_dims_qme_b; + use_fastlapack_lu = use_fastlapack_lu) - if !โ„’.issuccess(fFI) + if !solved_FI return A, iter, 1.0 end @@ -342,7 +473,9 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # @timeit_debug timer "Compute F" begin # Compute F = F * FI * F - โ„’.ldiv!(temp3, fFI, F) + copyto!(temp3, F) + solve_lu_left!(temp2, temp3, workspace.fast_lu_ws_qme_b, fFI; + use_fastlapack_lu = use_fastlapack_lu) โ„’.mul!(F_new, F, temp3) # F_new = F / fFI * F @@ -351,7 +484,8 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # Compute X_new = X + F * FI * X * E โ„’.mul!(temp3, X, E) - โ„’.ldiv!(fFI, temp3) + solve_lu_left!(temp2, temp3, workspace.fast_lu_ws_qme_b, fFI; + use_fastlapack_lu = use_fastlapack_lu) โ„’.mul!(X_new, F, temp3) # X_new = F / fFI * X * E if i > 5 || guess_provided @@ -366,7 +500,8 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # Compute Y_new = Y + E * EI * Y * F โ„’.mul!(X, Y, F) # use X as temporary storage - โ„’.ldiv!(fEI, X) + solve_lu_left!(temp1, X, workspace.fast_lu_ws_qme_a, fEI; + use_fastlapack_lu = use_fastlapack_lu) โ„’.mul!(Y_new, E, X) # Y_new = E / fEI * Y * F if i > 5 || guess_provided @@ -379,7 +514,7 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, # println("Iter: $i; xtol: $Xtol; ytol: $Ytol; rel ytol: $relYtol; rel xtol: $relXtol") # Check for convergence - if Xtol < tol # && Yreltol < tol # i % 2 == 0 && + if Xtol < tol.atol # && Yreltol < tol # i % 2 == 0 && solved = true iter = i break @@ -404,20 +539,31 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, โ„’.mul!(temp1, X_new, X_new) โ„’.mul!(AXX, A, temp1) - AXXnorm = max(โ„’.norm(AXX), โ„’.norm(C)) - + AXXnorm = min(โ„’.norm(AXX), โ„’.norm(C)) + โ„’.mul!(AXX, B, X_new, 1, 1) โ„’.axpy!(1, C, AXX) reached_tol = โ„’.norm(AXX) / AXXnorm - + # if reached_tol > tol # println("QME: doubling $reached_tol") # end - # Return a copy of X_new (to avoid returning a reference to mutable workspace) - return copy(X_new), iter, reached_tol + X_cache = if caching + _existing_sol = cache.qme_solution + if _existing_sol isa Matrix{R} && size(_existing_sol) == size(X_new) + _existing_sol + else + cache.qme_solution = zeros(R, size(X_new, 1), size(X_new, 2)) + end + else + zeros(R, size(X_new, 1), size(X_new, 2)) + end + copyto!(X_cache, X_new) + + return X_cache, iter, reached_tol end diff --git a/src/algorithms/sylvester.jl b/src/algorithms/sylvester.jl index c0d739a95..c74fe73cd 100644 --- a/src/algorithms/sylvester.jl +++ b/src/algorithms/sylvester.jl @@ -16,42 +16,80 @@ function solve_sylvester_equation(A::M, ๐•Šโ„‚::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), sylvester_algorithm::Symbol = :doubling, - acceptance_tol::AbstractFloat = 1e-10, - tol::AbstractFloat = 1e-14, + tol::SolverTolerances = SolverTolerances(), verbose::Bool = false)::Union{Tuple{Matrix{Float64}, Bool}, Tuple{SparseMatrixCSC{Float64, Int}, Bool}, Tuple{ThreadedSparseArrays.ThreadedSparseMatrixCSC{Float64, Int, SparseMatrixCSC{Float64, Int}}, Bool}} where {M <: AbstractMatrix{Float64}, N <: AbstractMatrix{Float64}, O <: AbstractMatrix{Float64}} # timer::TimerOutput = TimerOutput(), + # Ownership: low-level methods below are mixed. Some return freshly allocated + # matrices, while dense doubling and Krylov paths can return workspace-backed + # buffers (for example ๐•Šโ„‚.๐‚_dbl or ๐•Šโ„‚.๐—). This dispatcher therefore returns + # an owned copy so callers do not accidentally retain aliased workspace state. # @timeit_debug timer "Choose matrix formats" begin - - if sylvester_algorithm == :bartels_stewart - b = collect(B) - else - b = choose_matrix_format(B)# |> collect - end + # Ensure doubling buffers are allocated unconditionally so they are available + # for both the primary path and fallback retry paths below. + # Doubling buffers (๐€, ๐, ๐‚_dbl, ๐‚ยน, ๐‚B) are reused in fallback Krylov/bartels_stewart + # retry paths to avoid allocating via collect(). They are NOT used by those solvers + # (they only use krylov_workspace buffers: tmp, ๐—, ๐‚), so there is no aliasing. + # + # For dqgmres refinement fallbacks (initial_guess = x), we use ๐‚ยน instead of ๐‚_dbl + # for cc, because x may alias ๐•Šโ„‚.๐‚_dbl from a prior doubling solve. + # + # The doubling retry path still uses collect() because the doubling method modifies + # its workspace copies of A/B internally (squaring) and reads the original A/B/C + # arguments for the final residualโ€”passing workspace aliases would corrupt those reads. + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) if sylvester_algorithm โˆˆ [:bicgstab, :gmres, :dqgmres, :bartels_stewart] - a = collect(A) + a = ๐•Šโ„‚.๐€ + copyto!(a, A) - c = collect(C) + c = ๐•Šโ„‚.๐‚_dbl + copyto!(c, C) + + if sylvester_algorithm == :bartels_stewart + b = ๐•Šโ„‚.๐ + copyto!(b, B) + else + b = choose_matrix_format(B) + end else - a = choose_matrix_format(A)# |> sparse + a = choose_matrix_format(A) - c = choose_matrix_format(C)# |> sparse + b = choose_matrix_format(B) + # b = B + + c = choose_matrix_format(C) end # end # timeit_debug # @timeit_debug timer "Check if guess solves it already" begin - if length(initial_guess) > 0 - ๐‚ = a * initial_guess * b + c - initial_guess - - reached_tol = โ„’.norm(๐‚) / โ„’.norm(initial_guess) + initial_guess_acceptance_tol = tol.initial_guess_acceptance_tol + acceptance_tol = tol.acceptance_tol - if reached_tol < acceptance_tol - if verbose println("Sylvester equation - previous solution achieves relative tol of $reached_tol") end + if length(initial_guess) > 0 || length(C) > 0 + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_krylov_buffers!(๐•Šโ„‚, n, m) + + guess = length(initial_guess) > 0 ? initial_guess : c + guess_name = length(initial_guess) > 0 ? "previous solution" : "C" + + _tmp = ๐•Šโ„‚.tmp + _res = ๐•Šโ„‚.๐‚ + โ„’.mul!(_tmp, guess, b) + โ„’.mul!(_res, a, _tmp) + โ„’.axpy!(1, c, _res) + โ„’.axpy!(-1, guess, _res) + + denom = max(โ„’.norm(guess), โ„’.norm(c)) + reached_tol = denom == 0 ? 0.0 : โ„’.norm(_res) / denom - # X = choose_matrix_format(initial_guess) + if reached_tol < initial_guess_acceptance_tol + if verbose println("Sylvester equation - $guess_name achieves relative tol of $reached_tol (initial guess tol: $initial_guess_acceptance_tol)") end - return initial_guess, true + return choose_matrix_format(guess), true end end @@ -69,11 +107,14 @@ function solve_sylvester_equation(A::M, end if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && (sylvester_algorithm โ‰  :bartels_stewart) && (length(B) < 5e7) # try sylvester if previous one didn't solve it - aa = collect(A) + aa = ๐•Šโ„‚.๐€ + copyto!(aa, A) - bb = collect(B) + bb = ๐•Šโ„‚.๐ + copyto!(bb, B) - cc = collect(C) + cc = ๐•Šโ„‚.๐‚_dbl + copyto!(cc, C) x, i, reached_tol = solve_sylvester_equation(aa, bb, cc, Val(:bartels_stewart), ๐•Šโ„‚, @@ -88,9 +129,12 @@ function solve_sylvester_equation(A::M, end if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && reached_tol < sqrt(acceptance_tol) - aa = collect(A) + aa = ๐•Šโ„‚.๐€ + copyto!(aa, A) - cc = collect(C) + # Use ๐‚ยน (not ๐‚_dbl) because x may alias ๐•Šโ„‚.๐‚_dbl from a prior doubling solve + cc = ๐•Šโ„‚.๐‚ยน + copyto!(cc, C) X, i, Reached_tol = solve_sylvester_equation(aa, b, cc, Val(:dqgmres), ๐•Šโ„‚, @@ -109,9 +153,11 @@ function solve_sylvester_equation(A::M, end if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && sylvester_algorithm โ‰  :gmres - aa = collect(A) + aa = ๐•Šโ„‚.๐€ + copyto!(aa, A) - cc = collect(C) + cc = ๐•Šโ„‚.๐‚_dbl + copyto!(cc, C) x, i, reached_tol = solve_sylvester_equation(aa, b, cc, Val(:gmres), ๐•Šโ„‚, @@ -126,9 +172,12 @@ function solve_sylvester_equation(A::M, end if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && reached_tol < sqrt(acceptance_tol) - aa = collect(A) + aa = ๐•Šโ„‚.๐€ + copyto!(aa, A) - cc = collect(C) + # Use ๐‚ยน (not ๐‚_dbl) because x may alias ๐•Šโ„‚.๐‚_dbl from a prior doubling solve + cc = ๐•Šโ„‚.๐‚ยน + copyto!(cc, C) X, i, Reached_tol = solve_sylvester_equation(aa, b, cc, Val(:dqgmres), ๐•Šโ„‚, @@ -147,6 +196,8 @@ function solve_sylvester_equation(A::M, end if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && sylvester_algorithm โ‰  :doubling + # Must use collect() here: the doubling method aliases ๐•Šโ„‚.๐€/๐•Šโ„‚.๐‚_dbl internally + # (squaring A, iterating C) then reads the original A/C for the final residual. aa = collect(A) cc = collect(C) @@ -237,8 +288,9 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 + # Ownership: returns owned sparse storage created locally in this method. # guess_provided = true if length(initial_guess) == 0 @@ -268,7 +320,7 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, if i % 2 == 0 normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -286,11 +338,10 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, ๐‚ += initial_guess - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + ๐‚_res = A * ๐‚ * B + ๐‚_res += C + ๐‚_res -= ๐‚ + reached_tol = โ„’.norm(๐‚_res) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -305,25 +356,34 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 + # Ownership: returns workspace-backed dense buffer ๐•Šโ„‚.๐‚_dbl. # guess_provided = true if length(initial_guess) == 0 # guess_provided = false initial_guess = zero(C) end - ๐€ = copy(A) ๐€ยน = copy(A) ๐ = copy(B) ๐ยน = copy(B) - # ๐‚ = length(init) == 0 ? copy(C) : copy(init) - ๐‚ = A * initial_guess * B + C - initial_guess #copy(C) - # โ„’.rmul!(๐‚, -1) - ๐‚ยน = similar(๐‚) - ๐‚B = copy(C) + # Use workspace for dense C-related buffers + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) + + ๐‚ = ๐•Šโ„‚.๐‚_dbl + ๐‚ยน = ๐•Šโ„‚.๐‚ยน + ๐‚B = ๐•Šโ„‚.๐‚B + + # ๐‚ = A * initial_guess * B + C - initial_guess + โ„’.mul!(๐‚B, initial_guess, B) + โ„’.mul!(๐‚, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚) + โ„’.axpy!(-1, initial_guess, ๐‚) max_iter = 500 @@ -346,8 +406,10 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, droptol!(๐, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -357,24 +419,14 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, copy!(๐‚,๐‚ยน) end - # โ„’.mul!(๐‚B, ๐‚, ๐) - # โ„’.mul!(๐‚ยน, ๐€, ๐‚B) - # โ„’.axpy!(1, ๐‚, ๐‚ยน) - # # ๐‚ยน = ๐€ * ๐‚ * ๐ + ๐‚ - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # โ„’.axpy!(-1, ๐‚, ๐‚ยน) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน) / denom - โ„’.axpy!(1, initial_guess, ๐‚) - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚ยน, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚ยน) + + reached_tol = โ„’.norm(๐‚ยน) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -390,8 +442,9 @@ function solve_sylvester_equation( A::Matrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 + # Ownership: returns workspace-backed dense buffer ๐•Šโ„‚.๐‚_dbl. # @timeit_debug timer "Doubling solve" begin # @timeit_debug timer "Setup buffers" begin @@ -401,17 +454,26 @@ function solve_sylvester_equation( A::Matrix{T}, # guess_provided = false initial_guess = zero(C) end - - ๐€ = copy(A) - ๐€ยน = copy(A) + # Use workspace for dense matrices A and C + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) + + ๐€ = ๐•Šโ„‚.๐€ + ๐€ยน = ๐•Šโ„‚.๐€ยน + copyto!(๐€, A) + ๐ = copy(B) - # ๐ยน = similar(B) - # ๐‚ = length(init) == 0 ? copy(C) : copy(init) - ๐‚ = A * initial_guess * B + C - initial_guess #copy(C) - - # โ„’.rmul!(๐‚, -1) - ๐‚ยน = similar(๐‚) - ๐‚B = similar(C) + + ๐‚ = ๐•Šโ„‚.๐‚_dbl + ๐‚ยน = ๐•Šโ„‚.๐‚ยน + ๐‚B = ๐•Šโ„‚.๐‚B + + # ๐‚ = A * initial_guess * B + C - initial_guess + โ„’.mul!(๐‚B, initial_guess, B) + โ„’.mul!(๐‚, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚) + โ„’.axpy!(-1, initial_guess, ๐‚) max_iter = 500 @@ -446,8 +508,10 @@ function solve_sylvester_equation( A::Matrix{T}, # end # timeit_debug if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -459,28 +523,14 @@ function solve_sylvester_equation( A::Matrix{T}, # end # timeit_debug end - # @timeit_debug timer "Finalise" begin - # โ„’.mul!(๐‚B, ๐‚, ๐) - # โ„’.mul!(๐‚ยน, ๐€, ๐‚B) - # โ„’.axpy!(1, ๐‚, ๐‚ยน) - # # ๐‚ยน = ๐€ * ๐‚ * ๐ + ๐‚ - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # โ„’.axpy!(-1, ๐‚, ๐‚ยน) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน) / denom - - ๐‚ += initial_guess - - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # end # timeit_debug - # end # timeit_debug + โ„’.axpy!(1, initial_guess, ๐‚) - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚ยน, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚ยน) + + reached_tol = โ„’.norm(๐‚ยน) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -494,7 +544,7 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, initial_guess::AbstractMatrix{T} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 On Smith-type iterative algorithms for the Stein matrix equation # guess_provided = true @@ -504,15 +554,25 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, end ๐€ = copy(A) - # ๐€ยน = copy(A) - ๐ = copy(B) - ๐ยน = copy(B) - # ๐‚ = length(init) == 0 ? copy(C) : copy(init) - ๐‚ = A * initial_guess * B + C - initial_guess #copy(C) - # โ„’.rmul!(๐‚, -1) - ๐‚ยน = similar(๐‚) - ๐‚B = copy(C) + # Use workspace for dense B and C buffers + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) + + ๐ = ๐•Šโ„‚.๐ + ๐ยน = ๐•Šโ„‚.๐ยน + copyto!(๐, B) + + ๐‚ = ๐•Šโ„‚.๐‚_dbl + ๐‚ยน = ๐•Šโ„‚.๐‚ยน + ๐‚B = ๐•Šโ„‚.๐‚B + + # ๐‚ = A * initial_guess * B + C - initial_guess + โ„’.mul!(๐‚B, initial_guess, B) + โ„’.mul!(๐‚, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚) + โ„’.axpy!(-1, initial_guess, ๐‚) max_iter = 500 @@ -533,8 +593,10 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, # droptol!(๐, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -544,24 +606,14 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, copy!(๐‚,๐‚ยน) end - # โ„’.mul!(๐‚B, ๐‚, ๐) - # โ„’.mul!(๐‚ยน, ๐€, ๐‚B) - # โ„’.axpy!(1, ๐‚, ๐‚ยน) - # ๐‚ยน = ๐€ * ๐‚ * ๐ + ๐‚ - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # โ„’.axpy!(-1, ๐‚, ๐‚ยน) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน) / denom - - ๐‚ += initial_guess - - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) + โ„’.axpy!(1, initial_guess, ๐‚) - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚ยน, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚ยน) + + reached_tol = โ„’.norm(๐‚ยน) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -576,7 +628,7 @@ function solve_sylvester_equation( A::Matrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 # guess_provided = true @@ -585,16 +637,20 @@ function solve_sylvester_equation( A::Matrix{T}, initial_guess = zero(C) end - ๐€ = copy(A) - ๐€ยน = copy(A) - ๐ = copy(B) - ๐ยน = copy(B) - # ๐‚ = length(init) == 0 ? copy(C) : copy(init) - ๐‚ = A * initial_guess * B + C - initial_guess #copy(C) - - # โ„’.rmul!(๐‚, -1) + # Use workspace for dense A and B buffers + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) + + ๐€ = ๐•Šโ„‚.๐€ + ๐€ยน = ๐•Šโ„‚.๐€ยน + ๐ = ๐•Šโ„‚.๐ + ๐ยน = ๐•Šโ„‚.๐ยน + copyto!(๐€, A) + copyto!(๐, B) + + ๐‚ = A * initial_guess * B + C - initial_guess ๐‚ยน = similar(๐‚) - # ๐‚B = copy(C) max_iter = 500 @@ -617,8 +673,11 @@ function solve_sylvester_equation( A::Matrix{T}, # droptol!(๐, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + ๐‚B = ๐•Šโ„‚.๐‚B + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -629,22 +688,16 @@ function solve_sylvester_equation( A::Matrix{T}, ๐‚ = ๐‚ยน end - # โ„’.mul!(๐‚B, ๐‚, ๐) - # โ„’.mul!(๐‚ยน, ๐€, ๐‚B) - # โ„’.axpy!(1, ๐‚, ๐‚ยน) - # ๐‚ยน = ๐€ * ๐‚ * ๐ + ๐‚ - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน - ๐‚) / denom - ๐‚ += initial_guess - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + ๐‚B = ๐•Šโ„‚.๐‚B + ๐‚_tmp = ๐•Šโ„‚.๐‚_dbl + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚_tmp, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚_tmp) + โ„’.axpy!(-1, ๐‚, ๐‚_tmp) + + reached_tol = โ„’.norm(๐‚_tmp) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -659,7 +712,7 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 # guess_provided = true @@ -669,15 +722,18 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, end ๐€ = copy(A) - # ๐€ยน = copy(A) - ๐ = copy(B) - ๐ยน = copy(B) - # ๐‚ = length(init) == 0 ? copy(C) : copy(init) - ๐‚ = A * initial_guess * B + C - initial_guess #copy(C) - # โ„’.rmul!(๐‚, -1) + # Use workspace for dense B buffers + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) + + ๐ = ๐•Šโ„‚.๐ + ๐ยน = ๐•Šโ„‚.๐ยน + copyto!(๐, B) + + ๐‚ = A * initial_guess * B + C - initial_guess ๐‚ยน = similar(๐‚) - # ๐‚B = copy(C) max_iter = 500 @@ -700,8 +756,11 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, # droptol!(๐, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + ๐‚B = ๐•Šโ„‚.๐‚B + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -712,22 +771,16 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, ๐‚ = ๐‚ยน end - # โ„’.mul!(๐‚B, ๐‚, ๐) - # โ„’.mul!(๐‚ยน, ๐€, ๐‚B) - # โ„’.axpy!(1, ๐‚, ๐‚ยน) - # ๐‚ยน = ๐€ * ๐‚ * ๐ + ๐‚ - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน - ๐‚) / denom - ๐‚ += initial_guess - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + ๐‚B = ๐•Šโ„‚.๐‚B + ๐‚_tmp = ๐•Šโ„‚.๐‚_dbl + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚_tmp, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚_tmp) + โ„’.axpy!(-1, ๐‚, ๐‚_tmp) + + reached_tol = โ„’.norm(๐‚_tmp) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -741,25 +794,28 @@ function solve_sylvester_equation( A::Matrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 + # Ownership: returns owned dense storage created locally in this method. # guess_provided = true if length(initial_guess) == 0 # guess_provided = false initial_guess = zero(C) end - - ๐€ = copy(A) - ๐€ยน = copy(A) + # Use workspace for dense A buffers + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, n, m) + + ๐€ = ๐•Šโ„‚.๐€ + ๐€ยน = ๐•Šโ„‚.๐€ยน + copyto!(๐€, A) + ๐ = copy(B) - # ๐ยน = copy(B) - # ๐‚ = length(init) == 0 ? copy(C) : copy(init) - ๐‚ = A * initial_guess * B + C - initial_guess #copy(C) - - # โ„’.rmul!(๐‚, -1) + + ๐‚ = A * initial_guess * B + C - initial_guess ๐‚ยน = similar(๐‚) - # ๐‚B = copy(C) max_iter = 500 @@ -782,8 +838,11 @@ function solve_sylvester_equation( A::Matrix{T}, droptol!(๐, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + ๐‚B = ๐•Šโ„‚.๐‚B + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -794,22 +853,16 @@ function solve_sylvester_equation( A::Matrix{T}, ๐‚ = ๐‚ยน end - # โ„’.mul!(๐‚B, ๐‚, ๐) - # โ„’.mul!(๐‚ยน, ๐€, ๐‚B) - # โ„’.axpy!(1, ๐‚, ๐‚ยน) - # ๐‚ยน = ๐€ * ๐‚ * ๐ + ๐‚ - - # denom = max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) - - # reached_tol = denom == 0 ? 0.0 : โ„’.norm(๐‚ยน - ๐‚) / denom - ๐‚ += initial_guess - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + ๐‚B = ๐•Šโ„‚.๐‚B + ๐‚_tmp = ๐•Šโ„‚.๐‚_dbl + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚_tmp, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚_tmp) + โ„’.axpy!(-1, ๐‚, ๐‚_tmp) + + reached_tol = โ„’.norm(๐‚_tmp) / max(โ„’.norm(๐‚), โ„’.norm(C)) return ๐‚, iters, reached_tol # return info on convergence end @@ -823,8 +876,9 @@ function solve_sylvester_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 + # Ownership: returns workspace-backed dense buffer ๐•Šโ„‚.๐‚_dbl. # @timeit_debug timer "Setup buffers" begin # guess_provided = true @@ -886,8 +940,10 @@ function solve_sylvester_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat # droptol!(๐, eps()) if i % 2 == 0 - normdiff = โ„’.norm(๐‚ยน - ๐‚) - if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol + copyto!(๐‚B, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚B) + normdiff = โ„’.norm(๐‚B) + if !isfinite(normdiff) || normdiff / max(โ„’.norm(๐‚), โ„’.norm(๐‚ยน)) < tol.rtol # if isapprox(๐‚ยน, ๐‚, rtol = tol) iters = i break @@ -914,15 +970,14 @@ function solve_sylvester_equation( A::Union{โ„’.Adjoint{T, Matrix{T}}, DenseMat โ„’.axpy!(1, initial_guess, ๐‚) - reached_tol = โ„’.norm(A * ๐‚ * B + C - ๐‚) / max(โ„’.norm(๐‚), โ„’.norm(C)) - - # end # timeit_debug - - # if reached_tol > tol - # println("Sylvester: doubling $reached_tol") - # end + โ„’.mul!(๐‚B, ๐‚, B) + โ„’.mul!(๐‚ยน, A, ๐‚B) + โ„’.axpy!(1, C, ๐‚ยน) + โ„’.axpy!(-1, ๐‚, ๐‚ยน) + + reached_tol = โ„’.norm(๐‚ยน) / max(โ„’.norm(๐‚), โ„’.norm(C)) - return copy(๐‚), iters, reached_tol # return info on convergence + return ๐‚, iters, reached_tol # return info on convergence end @@ -934,7 +989,8 @@ function solve_sylvester_equation(A::DenseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::AbstractFloat = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned dense matrix from MatrixEquations.sylvd. # guess_provided = true if length(initial_guess) == 0 @@ -995,7 +1051,8 @@ function solve_sylvester_equation(A::DenseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense Krylov buffer ๐•Šโ„‚.๐—. # @timeit_debug timer "Preallocate matrices" begin # guess_provided = true @@ -1079,20 +1136,20 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # precond = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, preconditioner!) - if ๐•Šโ„‚.krylov_workspace.bicgstab.m == 0 - ๐•Šโ„‚.krylov_workspace.bicgstab = BicgstabWorkspace(length(C), length(C), Vector{T}) + if ๐•Šโ„‚.krylov.bicgstab.m == 0 + ๐•Šโ„‚.krylov.bicgstab = BicgstabWorkspace(length(C), length(C), Vector{T}) end # @timeit_debug timer "BICGSTAB solve" begin # if length(init) == 0 # ๐‚, info = Krylov.bicgstab(sylvester, C[idxs], rtol = tol / 10, atol = tol / 10)#, M = precond) # ๐‚, info = Krylov.bicgstab(sylvester, [vec(๐•Šโ„‚.๐‚);], - Krylov.bicgstab!( ๐•Šโ„‚.krylov_workspace.bicgstab, + Krylov.bicgstab!( ๐•Šโ„‚.krylov.bicgstab, sylvester, [vec(๐‚ยน);], # [vec(initial_guess);], itmax = min(5000,max(500,Int(round(sqrt(length(๐‚ยน)*10))))), timemax = 10.0, - rtol = tol, - atol = tol)#, M = precond) + rtol = tol.rtol, + atol = tol.atol)#, M = precond) # else # ๐‚, info = Krylov.bicgstab(sylvester, [vec(C);], [vec(init);], rtol = tol / 10) # end @@ -1101,7 +1158,7 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # @timeit_debug timer "Postprocess" begin # # @inbounds ๐•Šโ„‚.๐—[idxs] = ๐‚ - copyto!(๐—, ๐•Šโ„‚.krylov_workspace.bicgstab.x) + copyto!(๐—, ๐•Šโ„‚.krylov.bicgstab.x) # โ„’.mul!(tmpฬ„, A, ๐— * B) # โ„’.axpy!(1, C, tmpฬ„) @@ -1132,7 +1189,7 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # end # iter = info.niter - iter = ๐•Šโ„‚.krylov_workspace.bicgstab.stats.niter + iter = ๐•Šโ„‚.krylov.bicgstab.stats.niter # return ๐•Šโ„‚.๐—, iter, reached_tol return ๐—, iter, reached_tol @@ -1147,7 +1204,8 @@ function solve_sylvester_equation(A::DenseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense Krylov buffer ๐•Šโ„‚.๐—. # @timeit_debug timer "Preallocate matrices" begin # guess_provided = true @@ -1231,20 +1289,20 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # precond = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, preconditioner!) - if ๐•Šโ„‚.krylov_workspace.dqgmres.m == 0 - ๐•Šโ„‚.krylov_workspace.dqgmres = DqgmresWorkspace(length(C), length(C), Vector{T}) + if ๐•Šโ„‚.krylov.dqgmres.m == 0 + ๐•Šโ„‚.krylov.dqgmres = DqgmresWorkspace(length(C), length(C), Vector{T}) end # @timeit_debug timer "DQGMRES solve" begin # if length(init) == 0 # ๐‚, info = Krylov.dqgmres(sylvester, C[idxs], rtol = tol / 10, atol = tol / 10)#, M = precond) # ๐‚, info = Krylov.dqgmres(sylvester, [vec(๐•Šโ„‚.๐‚);], - Krylov.dqgmres!(๐•Šโ„‚.krylov_workspace.dqgmres, + Krylov.dqgmres!(๐•Šโ„‚.krylov.dqgmres, sylvester, [vec(๐‚ยน);], # [vec(initial_guess);], itmax = min(5000,max(500,Int(round(sqrt(length(๐‚ยน)*10))))), timemax = 10.0, - rtol = tol, - atol = tol)#, M = precond) + rtol = tol.rtol, + atol = tol.atol)#, M = precond) # else # ๐‚, info = Krylov.dqgmres(sylvester, [vec(C);], [vec(init);], rtol = tol / 10) # end @@ -1253,7 +1311,7 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # @timeit_debug timer "Postprocess" begin # # @inbounds ๐•Šโ„‚.๐—[idxs] = ๐‚ - copyto!(๐—, ๐•Šโ„‚.krylov_workspace.dqgmres.x) + copyto!(๐—, ๐•Šโ„‚.krylov.dqgmres.x) # โ„’.mul!(tmpฬ„, A, ๐— * B) # โ„’.axpy!(1, C, tmpฬ„) @@ -1284,7 +1342,7 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # end # iter = info.niter - iter = ๐•Šโ„‚.krylov_workspace.dqgmres.stats.niter + iter = ๐•Šโ„‚.krylov.dqgmres.stats.niter # return ๐•Šโ„‚.๐—, iter, reached_tol return ๐—, iter, reached_tol @@ -1299,7 +1357,8 @@ function solve_sylvester_equation(A::DenseMatrix{T}, initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense Krylov buffer ๐•Šโ„‚.๐—. # @timeit_debug timer "Preallocate matrices" begin # guess_provided = true @@ -1383,20 +1442,20 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # precond = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, preconditioner!) - if ๐•Šโ„‚.krylov_workspace.gmres.m == 0 - ๐•Šโ„‚.krylov_workspace.gmres = GmresWorkspace(length(C), length(C), Vector{T}) + if ๐•Šโ„‚.krylov.gmres.m == 0 + ๐•Šโ„‚.krylov.gmres = GmresWorkspace(length(C), length(C), Vector{T}) end # @timeit_debug timer "GMRES solve" begin # if length(init) == 0 # ๐‚, info = Krylov.gmres(sylvester, C[idxs], rtol = tol / 10, atol = tol / 10)#, M = precond) # ๐‚, info = Krylov.gmres(sylvester, [vec(๐•Šโ„‚.๐‚);], - Krylov.gmres!(๐•Šโ„‚.krylov_workspace.gmres, + Krylov.gmres!(๐•Šโ„‚.krylov.gmres, sylvester, [vec(๐‚ยน);], # [vec(initial_guess);], itmax = min(5000,max(500,Int(round(sqrt(length(๐‚ยน)*10))))), timemax = 10.0, - rtol = tol, - atol = tol)#, M = precond) + rtol = tol.rtol, + atol = tol.atol)#, M = precond) # else # ๐‚, info = Krylov.gmres(sylvester, [vec(C);], [vec(init);], rtol = tol / 10) # end @@ -1405,7 +1464,7 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # @timeit_debug timer "Postprocess" begin # # @inbounds ๐•Šโ„‚.๐—[idxs] = ๐‚ - copyto!(๐—, ๐•Šโ„‚.krylov_workspace.gmres.x) + copyto!(๐—, ๐•Šโ„‚.krylov.gmres.x) # โ„’.mul!(tmpฬ„, A, ๐— * B) # โ„’.axpy!(1, C, tmpฬ„) @@ -1436,7 +1495,7 @@ function solve_sylvester_equation(A::DenseMatrix{T}, # end # iter = info.niter - iter = ๐•Šโ„‚.krylov_workspace.gmres.stats.niter + iter = ๐•Šโ„‚.krylov.gmres.stats.niter # return ๐•Šโ„‚.๐—, iter, reached_tol return ๐—, iter, reached_tol diff --git a/src/custom_autodiff_rules/forwarddiff.jl b/src/custom_autodiff_rules/forwarddiff.jl index d30db7a19..3d3ec3d72 100644 --- a/src/custom_autodiff_rules/forwarddiff.jl +++ b/src/custom_autodiff_rules/forwarddiff.jl @@ -24,13 +24,12 @@ function sparse_preallocated!(ลœ::Matrix{โ„ฑ.Dual{Z,S,N}}; โ„‚::higher_order_wor sparse(Sฬ‚) end -function calculate_second_order_stochastic_steady_state(::Val{:newton}, - ๐’โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, - ๐’โ‚‚::AbstractSparseMatrix{โ„ฑ.Dual{Z,S,N}}, - x::Vector{โ„ฑ.Dual{Z,S,N}}, - ๐“‚::โ„ณ; - # timer::TimerOutput = TimerOutput(), - tol::AbstractFloat = 1e-14)::Tuple{Vector{โ„ฑ.Dual{Z,S,N}}, Bool} where {Z,S,N} +function solve_stochastic_steady_state_newton(::Val{:second_order}, + ๐’โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, + ๐’โ‚‚::AbstractSparseMatrix{โ„ฑ.Dual{Z,S,N}}, + x::Vector{โ„ฑ.Dual{Z,S,N}}, + ๐“‚::โ„ณ; + tol::AbstractFloat = 1e-14)::Tuple{Vector{โ„ฑ.Dual{Z,S,N}}, Bool} where {Z,S,N} ๐’โ‚ฬ‚ = โ„ฑ.value.(๐’โ‚) ๐’โ‚‚ฬ‚ = โ„ฑ.value.(๐’โ‚‚) @@ -43,7 +42,7 @@ function calculate_second_order_stochastic_steady_state(::Val{:newton}, T = constants.post_model_macro s_in_sโบ = so.s_in_sโบ s_in_s = so.s_in_s - I_nPast = ๐“‚.workspaces.qme.I_nPast + I_nPast = T.I_nPast kron_sโบ_sโบ = so.kron_sโบ_sโบ @@ -105,25 +104,25 @@ function calculate_second_order_stochastic_steady_state(::Val{:newton}, end, size(xฬ‚)), solved end -function calculate_third_order_stochastic_steady_state(::Val{:newton}, - ๐’โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, - ๐’โ‚‚::AbstractSparseMatrix{โ„ฑ.Dual{Z,S,N}}, - ๐’โ‚ƒ::AbstractSparseMatrix{โ„ฑ.Dual{Z,S,N}}, - x::Vector{โ„ฑ.Dual{Z,S,N}}, - ๐“‚::โ„ณ; - tol::AbstractFloat = 1e-14)::Tuple{Vector{โ„ฑ.Dual{Z,S,N}}, Bool} where {Z,S,N} +function solve_stochastic_steady_state_newton(::Val{:third_order}, + ๐’โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, + ๐’โ‚‚::AbstractSparseMatrix{โ„ฑ.Dual{Z,S,N}}, + ๐’โ‚ƒ::AbstractSparseMatrix{โ„ฑ.Dual{Z,S,N}}, + x::Vector{โ„ฑ.Dual{Z,S,N}}, + ๐“‚::โ„ณ; + tol::AbstractFloat = 1e-14)::Tuple{Vector{โ„ฑ.Dual{Z,S,N}}, Bool} where {Z,S,N} ๐’โ‚ฬ‚ = โ„ฑ.value.(๐’โ‚) ๐’โ‚‚ฬ‚ = โ„ฑ.value.(๐’โ‚‚) ๐’โ‚ƒฬ‚ = โ„ฑ.value.(๐’โ‚ƒ) xฬ‚ = โ„ฑ.value.(x) # Get cached computational constants - so = ensure_computational_constants!(๐“‚) + so = ensure_computational_constants!(๐“‚.constants) T = ๐“‚.constants.post_model_macro โ„‚ = ๐“‚.workspaces.third_order s_in_sโบ = so.s_in_sโบ s_in_s = so.s_in_s - I_nPast = ๐“‚.workspaces.qme.I_nPast + I_nPast = T.I_nPast kron_sโบ_sโบ = so.kron_sโบ_sโบ @@ -222,11 +221,13 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, parameter_values_dual::Vector{โ„ฑ.Dual{Z,S,N}}; opts::CalculationOptions = merge_calculation_options(), cold_start::Bool = false, - estimation::Bool = false)::Tuple{Vector{โ„ฑ.Dual{Z,S,N}}, Tuple{S, Int}} where {Z, S <: AbstractFloat, N} + estimation::Bool = false, + caching::Bool = true)::Tuple{Vector{โ„ฑ.Dual{Z,S,N}}, Tuple{S, Int}} where {Z, S <: AbstractFloat, N} # timer::TimerOutput = TimerOutput(), parameter_values = โ„ฑ.value.(parameter_values_dual) ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) - qme_ws = ๐“‚.workspaces.qme + T = ๐“‚.constants.post_model_macro + qme_ws = ๐“‚.workspaces.first_order if ๐“‚.functions.NSSS_custom isa Function vars_in_ss_equations = ms.vars_in_ss_equations @@ -239,7 +240,8 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, length(๐“‚.constants.post_complete_parameters.parameters), ) - residual = zeros(length(๐“‚.equations.steady_state) + length(๐“‚.equations.calibration)) + residual = ๐“‚.workspaces.nsss_solver.check_residual + fill!(residual, 0.0) ๐“‚.functions.NSSS_check(residual, parameter_values, SS_and_pars_tmp) @@ -247,13 +249,15 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, iters = 0 - # if !isfinite(solution_error) || solution_error > opts.tol.NSSS_acceptance_tol - # throw(ArgumentError("Custom steady state function failed steady state check: residual $solution_error > $(opts.tol.NSSS_acceptance_tol). Parameters: $(parameter_values). Steady state and parameters returned: $(SS_and_pars_tmp).")) + # if !isfinite(solution_error) || solution_error > opts.tol.nsss.acceptance_tol + # throw(ArgumentError("Custom steady state function failed steady state check: residual $solution_error > $(opts.tol.nsss.acceptance_tol). Parameters: $(parameter_values). Steady state and parameters returned: $(SS_and_pars_tmp).")) # end - X = @ignore_derivatives ms.custom_ss_expand_matrix + X = ms.custom_ss_expand_matrix SS_and_pars = X * SS_and_pars_tmp else - SS_and_pars, (solution_error, iters) = ๐“‚.functions.NSSS_solve(parameter_values, ๐“‚, opts.tol, opts.verbose, cold_start, DEFAULT_SOLVER_PARAMETERS) + fastest_idx = ๐“‚.constants.post_complete_parameters.nsss_fastest_solver_parameter_idx + preferred_solver_parameter_idx = fastest_idx < 1 || fastest_idx > length(DEFAULT_SOLVER_PARAMETERS) ? 1 : fastest_idx + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(parameter_values, ๐“‚, opts.tol, opts.verbose, cold_start, DEFAULT_SOLVER_PARAMETERS, preferred_solver_parameter_idx = preferred_solver_parameter_idx) end # Allocate or reuse workspace for partials @@ -264,7 +268,7 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, end โˆ‚SS_and_pars = qme_ws.โˆ‚SS_and_pars - if solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error) + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) if opts.verbose println("Failed to find NSSS") end # Update failed counter @@ -275,72 +279,145 @@ function get_NSSS_and_parameters(๐“‚::โ„ณ, # Update success counter update_ss_counter!(๐“‚.counters, true, estimation = estimation) - SS_and_pars_names = ms.SS_and_pars_names - SS_and_pars_names_lead_lag = ms.SS_and_pars_names_lead_lag - - # unknowns = union(setdiff(๐“‚.vars_in_ss_equations, ๐“‚.constants.post_model_macro.โž•_vars), ๐“‚.calibration_equations_parameters) - unknowns = Symbol.(vcat(string.(sort(collect(setdiff(reduce(union,get_symbols.(๐“‚.equations.steady_state_aux)),union(๐“‚.constants.post_model_macro.parameters_in_equations,๐“‚.constants.post_model_macro.โž•_vars))))), ๐“‚.equations.calibration_parameters)) + custom_ss_expand_matrix = ms.custom_ss_expand_matrix โˆ‚ = parameter_values C = SS_and_pars[ms.SS_and_pars_no_exo_idx] # [dyn_ss_idx]) - if eltype(๐“‚.caches.โˆ‚equations_โˆ‚parameters) != eltype(parameter_values) - if ๐“‚.caches.โˆ‚equations_โˆ‚parameters isa SparseMatrixCSC - jac_buffer = similar(๐“‚.caches.โˆ‚equations_โˆ‚parameters, eltype(parameter_values)) - jac_buffer.nzval .= 0 + if eltype(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters) != eltype(parameter_values) + if ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters isa SparseMatrixCSC + jac_cache = similar(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters, eltype(parameter_values)) + jac_cache.nzval .= 0 else - jac_buffer = zeros(eltype(parameter_values), size(๐“‚.caches.โˆ‚equations_โˆ‚parameters)) + jac_cache = zeros(eltype(parameter_values), size(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters)) end else - jac_buffer = ๐“‚.caches.โˆ‚equations_โˆ‚parameters + jac_cache = ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters + if jac_cache isa SparseMatrixCSC + jac_cache.nzval .= 0 + else + fill!(jac_cache, zero(eltype(jac_cache))) + end end - ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters(jac_buffer, โˆ‚, C) + ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters(jac_cache, โˆ‚, C) - โˆ‚SS_equations_โˆ‚parameters = jac_buffer + โˆ‚SS_equations_โˆ‚parameters = jac_cache - if eltype(๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars) != eltype(parameter_values) - if ๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars isa SparseMatrixCSC - jac_buffer = similar(๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars, eltype(SS_and_pars)) - jac_buffer.nzval .= 0 + if eltype(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars) != eltype(parameter_values) + if ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars isa SparseMatrixCSC + jac_cache = similar(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars, eltype(SS_and_pars)) + jac_cache.nzval .= 0 else - jac_buffer = zeros(eltype(SS_and_pars), size(๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars)) + jac_cache = zeros(eltype(SS_and_pars), size(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars)) end else - jac_buffer = ๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars + jac_cache = ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars + if jac_cache isa SparseMatrixCSC + jac_cache.nzval .= 0 + else + fill!(jac_cache, zero(eltype(jac_cache))) + end end - ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars(jac_buffer, โˆ‚, C) + ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars(jac_cache, โˆ‚, C) - โˆ‚SS_equations_โˆ‚SS_and_pars = jac_buffer + โˆ‚SS_equations_โˆ‚SS_and_pars = jac_cache - โˆ‚SS_equations_โˆ‚SS_and_pars_lu = RF.lu(โˆ‚SS_equations_โˆ‚SS_and_pars, check = false) + if โˆ‚SS_equations_โˆ‚SS_and_pars isa SparseMatrixCSC + rhs_n_rows = size(โˆ‚SS_equations_โˆ‚SS_and_pars, 1) + rhs_n_cols = size(โˆ‚SS_equations_โˆ‚parameters, 2) - if !โ„’.issuccess(โˆ‚SS_equations_โˆ‚SS_and_pars_lu) - if opts.verbose println("Failed to calculate implicit derivative of NSSS") end - - solution_error = S(10.0) - else - JVP = -(โˆ‚SS_equations_โˆ‚SS_and_pars_lu \ โˆ‚SS_equations_โˆ‚parameters)#[indexin(SS_and_pars_names, unknowns),:] + if length(qme_ws.nsss_sparse_rhs) != rhs_n_rows + qme_ws.nsss_sparse_rhs = zeros(eltype(SS_and_pars), rhs_n_rows) + end + + if size(qme_ws.nsss_jvp_rhs) != (rhs_n_rows, rhs_n_cols) + qme_ws.nsss_jvp_rhs = zeros(eltype(SS_and_pars), rhs_n_rows, rhs_n_cols) + end + + if size(qme_ws.nsss_sparse_lu_buffer.A) != (rhs_n_rows, rhs_n_rows) + sparse_prob = ๐’ฎ.LinearProblem(โˆ‚SS_equations_โˆ‚SS_and_pars, qme_ws.nsss_sparse_rhs) + qme_ws.nsss_sparse_lu_buffer = ๐’ฎ.init(sparse_prob, + ๐’ฎ.LUFactorization(), + verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + else + qme_ws.nsss_sparse_lu_buffer.A = โˆ‚SS_equations_โˆ‚SS_and_pars + end + + sparse_solved = true + for j in 1:rhs_n_cols + @views copyto!(qme_ws.nsss_sparse_rhs, โˆ‚SS_equations_โˆ‚parameters[:, j]) + qme_ws.nsss_sparse_lu_buffer.b = qme_ws.nsss_sparse_rhs + sparse_sol = ๐’ฎ.solve!(qme_ws.nsss_sparse_lu_buffer) - jvp = zeros(length(SS_and_pars_names_lead_lag), length(๐“‚.constants.post_complete_parameters.parameters)) - - for (i,v) in enumerate(SS_and_pars_names) - if v in unknowns - jvp[i,:] = JVP[indexin([v], unknowns),:] + if sparse_sol.retcode != ๐’ฎ.SciMLBase.ReturnCode.Default && !๐’ฎ.SciMLBase.successful_retcode(sparse_sol.retcode) + sparse_solved = false + break end + + @views copyto!(qme_ws.nsss_jvp_rhs[:, j], qme_ws.nsss_sparse_lu_buffer.u) end - for i in 1:N - parameter_values_partials = โ„ฑ.partials.(parameter_values_dual, i) + if !sparse_solved + if opts.verbose println("Failed to calculate implicit derivative of NSSS") end + solution_error = S(10.0) + else + โ„’.rmul!(qme_ws.nsss_jvp_rhs, -1) + jvp_no_exo = custom_ss_expand_matrix * qme_ws.nsss_jvp_rhs + for i in 1:N + parameter_values_partials = โ„ฑ.partials.(parameter_values_dual, i) + @view(โˆ‚SS_and_pars[:,i]) .= jvp_no_exo * parameter_values_partials + end + end + else + qme_ws.fast_lu_ws_nsss, qme_ws.fast_lu_dims_nsss, solved_nsss, nsss_lu = factorize_lu!(โˆ‚SS_equations_โˆ‚SS_and_pars, + qme_ws.fast_lu_ws_nsss, + qme_ws.fast_lu_dims_nsss) - โˆ‚SS_and_pars[:,i] = jvp * parameter_values_partials + if !solved_nsss + if opts.verbose println("Failed to calculate implicit derivative of NSSS") end + solution_error = S(10.0) + else + rhs_dense = โˆ‚SS_equations_โˆ‚parameters isa Matrix ? โˆ‚SS_equations_โˆ‚parameters : Matrix(โˆ‚SS_equations_โˆ‚parameters) + + if size(qme_ws.nsss_jvp_rhs) != size(rhs_dense) + qme_ws.nsss_jvp_rhs = zeros(eltype(rhs_dense), size(rhs_dense)) + end + copyto!(qme_ws.nsss_jvp_rhs, rhs_dense) + + solve_lu_left!(โˆ‚SS_equations_โˆ‚SS_and_pars, + qme_ws.nsss_jvp_rhs, + qme_ws.fast_lu_ws_nsss, + nsss_lu) + + โ„’.rmul!(qme_ws.nsss_jvp_rhs, -1) + jvp_no_exo = custom_ss_expand_matrix * qme_ws.nsss_jvp_rhs + for i in 1:N + parameter_values_partials = โ„ฑ.partials.(parameter_values_dual, i) + @view(โˆ‚SS_and_pars[:,i]) .= jvp_no_exo * parameter_values_partials + end end end end + # Cache write: store NSSS result and stamp (using Float64 values) + if caching + cache_ss = ๐“‚.caches.non_stochastic_steady_state + if length(cache_ss) != length(SS_and_pars) + resize!(cache_ss, length(SS_and_pars)) + end + copyto!(cache_ss, SS_and_pars) + solved = !(solution_error > opts.tol.nsss.acceptance_tol) + if solved + ๐“‚.caches.valid_for.non_stochastic_steady_state = Float64.(parameter_values) + else + ๐“‚.caches.valid_for.non_stochastic_steady_state = Float64[] + end + end + return reshape(map(SS_and_pars, eachrow(โˆ‚SS_and_pars)) do v, p โ„ฑ.Dual{Z}(v, p...) # Z is the tag end, size(SS_and_pars)), (solution_error, iters) @@ -348,43 +425,92 @@ end function calculate_first_order_solution(โˆ‡โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, constants::constants, - qme_ws::qme_workspace, - sylv_ws::sylvester_workspace; + workspaces::workspaces, + cache::caches; opts::CalculationOptions = merge_calculation_options(), - initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0))::Tuple{Matrix{โ„ฑ.Dual{Z,S,N}}, Matrix{Float64}, Bool} where {Z,S,N} - โˆ‡ฬ‚โ‚ = โ„ฑ.value.(โˆ‡โ‚) + use_fastlapack_lu::Bool = true, + initial_guess::AbstractMatrix{<:Real} = zeros(0,0), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true)::Tuple{Matrix{โ„ฑ.Dual{Z,S,N}}, Matrix{Float64}, Bool} where {Z,S,N} T = constants.post_model_macro idx_constants = ensure_first_order_constants!(constants) + qme_ws = workspaces.first_order + sylv_ws = workspaces.sylvester_1st_order + ensure_first_order_workspace_buffers!(qme_ws, T, length(idx_constants.dyn_index), length(idx_constants.comb)) + ensure_sylvester_krylov_buffers!(qme_ws.sylvester, T.nVars, T.nVars) + ensure_sylvester_doubling_buffers!(qme_ws.sylvester, T.nVars, T.nVars) + + if size(qme_ws.p_tmp) != size(โˆ‡โ‚) + qme_ws.p_tmp = zeros(S, size(โˆ‡โ‚, 1), size(โˆ‡โ‚, 2)) + end + โˆ‡ฬ‚โ‚ = qme_ws.p_tmp + @inbounds for j in axes(โˆ‡โ‚, 2), i in axes(โˆ‡โ‚, 1) + โˆ‡ฬ‚โ‚[i, j] = โ„ฑ.value(โˆ‡โ‚[i, j]) + end expand_future = idx_constants.expand_future expand_past = idx_constants.expand_past - A = โˆ‡ฬ‚โ‚[:,1:T.nFuture_not_past_and_mixed] * expand_future - B = โˆ‡ฬ‚โ‚[:,idx_constants.nabla_zero_cols] + A = qme_ws.๐€โ‚€ + B = qme_ws.โˆ‡โ‚€ + X = qme_ws.sylvester.tmp + AXB = qme_ws.sylvester.๐— + AA = qme_ws.sylvester.๐‚ + Xยฒ = qme_ws.sylvester.๐€ + dA = qme_ws.sylvester.๐€ยน + dB = qme_ws.sylvester.๐ + dC = qme_ws.sylvester.๐ยน + CC = qme_ws.sylvester.๐‚_dbl + tmp = qme_ws.sylvester.๐‚ยน + B_sylv = qme_ws.sylvester.๐‚B + + # Legacy readable path (before workspace reuse): + # โˆ‡ฬ‚โ‚ = value.(โˆ‡โ‚) + # A = โˆ‡ฬ‚โ‚[:, 1:T.nFuture_not_past_and_mixed] * expand_future + # B = โˆ‡ฬ‚โ‚[:, idx_constants.nabla_zero_cols] + # X = ๐’โ‚[:, 1:end-T.nExo] * expand_past + # AXB = A * X + B + # AA = inv(AXB) * A + # Current code computes the same objects via `mul!`/`copyto!`/LU solves in reusable buffers. + + initial_guess_value = if length(initial_guess) == 0 + zeros(eltype(โˆ‡ฬ‚โ‚), 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{eltype(โˆ‡ฬ‚โ‚)} ? initial_guess : Matrix{eltype(โˆ‡ฬ‚โ‚)}(initial_guess) + else + โ„ฑ.value.(initial_guess) + end - ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡ฬ‚โ‚, constants, qme_ws, sylv_ws; opts = opts, initial_guess = initial_guess) + ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡ฬ‚โ‚, constants, workspaces, cache; opts = opts, initial_guess = initial_guess_value, caching = caching) if !solved return โˆ‡โ‚, qme_sol, false end - X = ๐’โ‚[:,1:end-T.nExo] * expand_past - - AXB = A * X + B - - AXBfact = RF.lu(AXB, check = false) + โ„’.mul!(A, @view(โˆ‡ฬ‚โ‚[:,1:T.nFuture_not_past_and_mixed]), expand_future) + copyto!(B, @view(โˆ‡ฬ‚โ‚[:,idx_constants.nabla_zero_cols])) - if !โ„’.issuccess(AXBfact) - AXBfact = โ„’.svd(AXB) - end + โ„’.mul!(X, @view(๐’โ‚[:,1:end-T.nExo]), expand_past) - invAXB = inv(AXBfact) + copyto!(AXB, B) + โ„’.mul!(AXB, A, X, 1, 1) - AA = invAXB * A + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_AXB, AXBfact = factorize_lu!(AXB, + qme_ws.fast_lu_ws_nabla0, + qme_ws.fast_lu_dims_nabla0; + use_fastlapack_lu = use_fastlapack_lu) - Xยฒ = X * X + if !solved_AXB + return โˆ‡โ‚, qme_sol, false + end + + copyto!(AA, A) + solve_lu_left!(AXB, AA, qme_ws.fast_lu_ws_nabla0, AXBfact; + use_fastlapack_lu = use_fastlapack_lu) + + โ„’.mul!(Xยฒ, X, X) - # Allocate or reuse workspace for partials (from qme_workspace) + # Allocate or reuse workspace for partials (from first_order_workspace) if size(qme_ws.Xฬƒ_first_order) != (length(๐’โ‚[:,1:end-T.nExo]), N) qme_ws.Xฬƒ_first_order = zeros(length(๐’โ‚[:,1:end-T.nExo]), N) else @@ -392,33 +518,44 @@ function calculate_first_order_solution(โˆ‡โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, end Xฬƒ = qme_ws.Xฬƒ_first_order - # Allocate or reuse workspace for temporary p matrix (from qme_workspace) - if size(qme_ws.p_tmp) != size(โˆ‡ฬ‚โ‚) - qme_ws.p_tmp = zero(โˆ‡ฬ‚โ‚) - else - fill!(qme_ws.p_tmp, zero(eltype(qme_ws.p_tmp))) - end - p = qme_ws.p_tmp + p = โˆ‡ฬ‚โ‚ + + copyto!(B_sylv, X) + โ„’.rmul!(B_sylv, -1) - initial_guess = zero(invAXB) + initial_guess = zeros(eltype(X), size(X, 1), size(X, 2)) # https://arxiv.org/abs/2011.11430 for i in 1:N p .= โ„ฑ.partials.(โˆ‡โ‚, i) - dA = p[:,1:T.nFuture_not_past_and_mixed] * expand_future - dB = p[:,idx_constants.nabla_zero_cols] - dC = p[:,idx_constants.nabla_minus_cols] * expand_past - - CC = invAXB * (dA * Xยฒ + dC + dB * X) + โ„’.mul!(dA, @view(p[:,1:T.nFuture_not_past_and_mixed]), expand_future) + copyto!(dB, @view(p[:,idx_constants.nabla_zero_cols])) + โ„’.mul!(dC, @view(p[:,idx_constants.nabla_minus_cols]), expand_past) + + copyto!(CC, dC) + โ„’.mul!(tmp, dA, Xยฒ) + CC .+= tmp + โ„’.mul!(tmp, dB, X) + CC .+= tmp + + # Legacy readable equivalent: + # CC = inv(AXB) * (dA * Xยฒ + dC + dB * X) + # followed by Sylvester solve with (-X, -CC). + # Here, `solve_lu_left!` replaces explicit inverse multiplication, + # and `B_sylv`/sign flip encode the same Sylvester system. + + solve_lu_left!(AXB, CC, qme_ws.fast_lu_ws_nabla0, AXBfact; + use_fastlapack_lu = use_fastlapack_lu) if โ„’.norm(CC) < eps() continue end - dX, solved = solve_sylvester_equation(AA, -X, -CC, sylv_ws, + โ„’.rmul!(CC, -1) + + dX, solved = solve_sylvester_equation(AA, B_sylv, CC, sylv_ws, initial_guess = initial_guess, sylvester_algorithm = opts.sylvester_algorithmยฒ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, + tol = opts.tol.first_order.ad.sylvester, verbose = opts.verbose) # if !solved @@ -434,7 +571,7 @@ function calculate_first_order_solution(โˆ‡โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, initial_guess = dX - Xฬƒ[:,i] = vec(dX[:,T.past_not_future_and_mixed_idx]) + @views copyto!(Xฬƒ[:,i],dX[:,T.past_not_future_and_mixed_idx]) end x = reshape(map(๐’โ‚[:,1:end-T.nExo], eachrow(Xฬƒ)) do v, p @@ -449,32 +586,60 @@ function calculate_first_order_solution(โˆ‡โ‚::Matrix{โ„ฑ.Dual{Z,S,N}}, B = -((โˆ‡โ‚Š * x * Jm + โˆ‡โ‚€) \ โˆ‡โ‚‘) - return hcat(x, B), qme_sol, solved + Sโ‚ = hcat(x, B) + + Sโ‚_value = โ„ฑ.value.(Sโ‚) + Sโ‚_existing = cache.first_order_solution_matrix + if Sโ‚_existing isa Matrix{S} && size(Sโ‚_existing) == size(Sโ‚_value) + copyto!(Sโ‚_existing, Sโ‚_value) + else + cache.first_order_solution_matrix = Sโ‚_value + end + + if !isempty(parameter_values) + cache.valid_for.first_order_solution = eltype(parameter_values) <: โ„ฑ.Dual ? Float64.(โ„ฑ.value.(parameter_values)) : Float64.(parameter_values) + end + + return Sโ‚, qme_sol, solved end function solve_quadratic_matrix_equation(A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, B::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, C::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, constants::constants, - workspace::qme_workspace; + workspaces::workspaces, + cache::caches; initial_guess::AbstractMatrix{<:Real} = zeros(0,0), - tol::AbstractFloat = 1e-8, - quadratic_matrix_equation_algorithm::Symbol = :schur, - verbose::Bool = false) where {Z,S,N} + tol::AdTolerances = AdTolerances(), + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + verbose::Bool = false, + caching::Bool = true) where {Z,S,N} T = constants.post_model_macro # unpack: AoS -> SoA Aฬ‚ = โ„ฑ.value.(A) Bฬ‚ = โ„ฑ.value.(B) Cฬ‚ = โ„ฑ.value.(C) - X, solved = solve_quadratic_matrix_equation(Aฬ‚, Bฬ‚, Cฬ‚, - Val(quadratic_matrix_equation_algorithm), + initial_guess_value = if length(initial_guess) == 0 + zeros(eltype(Aฬ‚), 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{eltype(Aฬ‚)} ? initial_guess : Matrix{eltype(Aฬ‚)}(initial_guess) + else + โ„ฑ.value.(initial_guess) + end + + qme_ws = ensure_qme_doubling_workspace!(workspaces, + T.nVars - T.nPresent_only) + + X, solved = solve_quadratic_matrix_equation(Aฬ‚, Bฬ‚, Cฬ‚, constants, - workspace; - tol = tol, - initial_guess = initial_guess, - # timer = timer, - verbose = verbose) + workspaces, + cache; + tol = tol.qme, + initial_guess = initial_guess_value, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + verbose = verbose, + caching = caching) AXB = Aฬ‚ * X + Bฬ‚ @@ -490,13 +655,13 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, Xยฒ = X * X - # Allocate or reuse workspace for partials (from qme_workspace) - if size(workspace.Xฬƒ) != (length(X), N) - workspace.Xฬƒ = zeros(length(X), N) + # Allocate or reuse workspace for partials (from qme_doubling_workspace) + if size(qme_ws.Xฬƒ) != (length(X), N) + qme_ws.Xฬƒ = zeros(length(X), N) else - fill!(workspace.Xฬƒ, zero(eltype(workspace.Xฬƒ))) + fill!(qme_ws.Xฬƒ, zero(eltype(qme_ws.Xฬƒ))) end - Xฬƒ = workspace.Xฬƒ + Xฬƒ = qme_ws.Xฬƒ # https://arxiv.org/abs/2011.11430 for i in 1:N @@ -508,7 +673,9 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, if โ„’.norm(CC) < eps() continue end - dX, slvd = solve_sylvester_equation(AA, -X, -CC, workspace.sylvester_ws, sylvester_algorithm = :doubling) + dX, slvd = solve_sylvester_equation(AA, -X, -CC, qme_ws.sylvester, + sylvester_algorithm = :doubling, + tol = tol.sylvester) solved = Bool(solved) && Bool(slvd) @@ -524,21 +691,34 @@ function solve_sylvester_equation( A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, B::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, C::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, ๐•Šโ„‚::sylvester_workspace; - initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), + initial_guess::AbstractMatrix{<:Real} = zeros(0,0), sylvester_algorithm::Symbol = :doubling, - acceptance_tol::AbstractFloat = 1e-10, - tol::AbstractFloat = 1e-14, + tol::SolverTolerances = SolverTolerances(), verbose::Bool = false)::Tuple{Matrix{โ„ฑ.Dual{Z,S,N}}, Bool} where {Z,S,N} # Extract Float64 values from Dual numbers Aฬ‚ = โ„ฑ.value.(A) Bฬ‚ = โ„ฑ.value.(B) Cฬ‚ = โ„ฑ.value.(C) + initial_guess_value = if length(initial_guess) == 0 + zeros(eltype(Aฬ‚), 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{eltype(Aฬ‚)} ? initial_guess : Matrix{eltype(Aฬ‚)}(initial_guess) + else + โ„ฑ.value.(initial_guess) + end + Pฬ‚, solved = solve_sylvester_equation(Aฬ‚, Bฬ‚, Cฬ‚, ๐•Šโ„‚, sylvester_algorithm = sylvester_algorithm, tol = tol, verbose = verbose, - initial_guess = initial_guess) + initial_guess = initial_guess_value) + + if size(๐•Šโ„‚.P) != size(Pฬ‚) + ๐•Šโ„‚.P = zeros(eltype(Pฬ‚), size(Pฬ‚)...) + end + copyto!(๐•Šโ„‚.P, Pฬ‚) + Pฬ‚_stable = ๐•Šโ„‚.P # Allocate or reuse workspaces for temporary copies if size(๐•Šโ„‚.รƒ_fd) != size(ร‚) @@ -575,7 +755,7 @@ function solve_sylvester_equation( A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, Bฬƒ .= โ„ฑ.partials.(B, i) Cฬƒ .= โ„ฑ.partials.(C, i) - X = Aฬƒ * Pฬ‚ * Bฬ‚ + Aฬ‚ * Pฬ‚ * Bฬƒ + Cฬƒ + X = Aฬƒ * Pฬ‚_stable * Bฬ‚ + Aฬ‚ * Pฬ‚_stable * Bฬƒ + Cฬƒ if โ„’.norm(X) < eps() continue end @@ -589,23 +769,44 @@ function solve_sylvester_equation( A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, Pฬƒ[:,i] = vec(P) end - return reshape(map(Pฬ‚, eachrow(Pฬƒ)) do v, p + return reshape(map(Pฬ‚_stable, eachrow(Pฬƒ)) do v, p โ„ฑ.Dual{Z}(v, p...) # Z is the tag - end, size(Pฬ‚)), solved + end, size(Pฬ‚_stable)), solved end function solve_lyapunov_equation( A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, C::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, workspace::lyapunov_workspace; + initial_guess::AbstractMatrix{<:Real} = zeros(0,0), lyapunov_algorithm::Symbol = :doubling, - tol::AbstractFloat = 1e-14, - acceptance_tol::AbstractFloat = 1e-12, + tol::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), verbose::Bool = false)::Tuple{Matrix{โ„ฑ.Dual{Z,S,N}}, Bool} where {Z,S,N} # Extract Float64 values from Dual numbers ร‚ = โ„ฑ.value.(A) ฤˆ = โ„ฑ.value.(C) - Pฬ‚, solved = solve_lyapunov_equation(ร‚, ฤˆ, workspace, lyapunov_algorithm = lyapunov_algorithm, tol = tol, verbose = verbose) + initial_guess_value = if length(initial_guess) == 0 + zeros(eltype(ร‚), 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{eltype(ร‚)} ? initial_guess : Matrix{eltype(ร‚)}(initial_guess) + else + โ„ฑ.value.(initial_guess) + end + + Pฬ‚, solved = solve_lyapunov_equation(ร‚, ฤˆ, workspace, + lyapunov_algorithm = lyapunov_algorithm, + initial_guess = initial_guess_value, + tol = tol, + verbose = verbose) + + if size(workspace.P) != size(Pฬ‚) + workspace.P = zeros(eltype(Pฬ‚), size(Pฬ‚)...) + end + copyto!(workspace.P, Pฬ‚) + Pฬ‚_stable = workspace.P # Allocate or reuse workspaces for temporary copies (from lyapunov_workspace) if size(workspace.รƒ_fd) != size(ร‚) @@ -635,86 +836,135 @@ function solve_lyapunov_equation( A::AbstractMatrix{โ„ฑ.Dual{Z,S,N}}, Aฬƒ .= โ„ฑ.partials.(A, i) Cฬƒ .= โ„ฑ.partials.(C, i) - X = Aฬƒ * Pฬ‚ * Aฬ‚' + Aฬ‚ * Pฬ‚ * Aฬƒ' + Cฬƒ + X = Aฬƒ * Pฬ‚_stable * Aฬ‚' + Aฬ‚ * Pฬ‚_stable * Aฬƒ' + Cฬƒ if โ„’.norm(X) < eps() continue end - P, slvd = solve_lyapunov_equation(ร‚, X, workspace, lyapunov_algorithm = lyapunov_algorithm, tol = tol, verbose = verbose) + # X = รƒ*Pฬ‚*ร‚' + ร‚*Pฬ‚*รƒ' + Cฬƒ is symmetric when C is symmetric (Pฬ‚ is always symmetric) + P, slvd = solve_lyapunov_equation(ร‚, X, workspace, + lyapunov_algorithm = lyapunov_algorithm, + tol = tol, + verbose = verbose) solved = solved && slvd Pฬƒ[:,i] = vec(P) end - return reshape(map(Pฬ‚, eachrow(Pฬƒ)) do v, p + return reshape(map(Pฬ‚_stable, eachrow(Pฬƒ)) do v, p โ„ฑ.Dual{Z}(v, p...) # Z is the tag - end, size(Pฬ‚)), solved + end, size(Pฬ‚_stable)), solved end -function run_kalman_iterations(A::Matrix{S}, - ๐::Matrix{S}, - C::Matrix{Float64}, - P::Matrix{S}, - data_in_deviations::Matrix{S}, - ws::kalman_workspace; +function calculate_loglikelihood(::Val{:kalman}, + ::Val, + observables_index::Vector{Int}, + ๐’::Union{Matrix{โ„ฑ.Dual{Z,S,N}},Vector{AbstractMatrix{โ„ฑ.Dual{Z,S,N}}}}, + data_in_deviations::Matrix{R}, + constants::constants, + state, + workspaces::workspaces; + warmup_iterations::Int = 0, presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, on_failure_loglikelihood::U = -Inf, - # timer::TimerOutput = TimerOutput(), - verbose::Bool = false)::S where {S <: โ„ฑ.Dual, U <: AbstractFloat} - # @timeit_debug timer "Calculate Kalman filter - forward mode AD" begin - # ForwardDiff requires fresh allocations - workspace not used here - u = zeros(S, size(C,2)) + opts::CalculationOptions = merge_calculation_options())::โ„ฑ.Dual{Z,S,N} where {Z,S,N,R <: Real, U <: AbstractFloat} + + T = constants.post_model_macro + idx_constants = constants.post_complete_parameters + lyap_ws = ensure_lyapunov_workspace!(workspaces, T.nVars, :first_order) + kalman_ws = workspaces.kalman - z = C * u + observables_and_states = sort(union(T.past_not_future_and_mixed_idx, observables_index)) + observables_sorted = sort(observables_index) + I_nVars = idx_constants.diag_nVars - loglik = S(0.0) + A = @views ๐’[observables_and_states,1:T.nPast_not_future_and_mixed] * I_nVars[T.past_not_future_and_mixed_idx, observables_and_states] + B = @views ๐’[observables_and_states,T.nPast_not_future_and_mixed+1:end] - F = similar(C * C') + C = @views I_nVars[observables_sorted, observables_and_states] + ๐ = B * B' - K = similar(C') + P = get_initial_covariance(Val(initial_covariance), A, ๐, lyap_ws, opts = opts) - for t in 1:size(data_in_deviations, 2) - if !all(isfinite.(z)) - if verbose println("KF not finite at step $t") end - return on_failure_loglikelihood + if !(eltype(P) <: โ„ฑ.Dual) + dual_zero = zero(A[1]) + P_float = P + P = similar(A, size(P_float, 1), size(P_float, 2)) + @inbounds for i in eachindex(P) + P[i] = dual_zero + S(P_float[i]) end + end - v = data_in_deviations[:, t] - z + u = zeros(eltype(A), size(C, 2)) + z = C * u + loglik = zero(eltype(A)) + + # Pre-allocate Dual-typed loop buffers to avoid per-iteration allocations + DT = eltype(A) + ns = size(A, 1) # n_obs_and_states + no = size(C, 1) # n_obs + v = zeros(DT, no) + CP = zeros(DT, no, ns) + F_buf = zeros(DT, no, no) + PCt = zeros(DT, ns, no) + K = zeros(DT, ns, no) + KC = zeros(DT, ns, ns) + PmKCP = zeros(DT, ns, ns) + AP = zeros(DT, ns, ns) + Kv = zeros(DT, ns) + uKv = zeros(DT, ns) + w = zeros(DT, no) - F = C * P * C' + for t in 1:size(data_in_deviations, 2) + if !all(isfinite.(z)) + if opts.verbose println("KF not finite at step $t") end + return on_failure_loglikelihood + end - luF = โ„’.lu(F, check = false) ### + @views v .= data_in_deviations[:, t] .- z + โ„’.mul!(CP, C, P) + โ„’.mul!(F_buf, CP, C') + luF = โ„’.lu(F_buf, check = false) if !โ„’.issuccess(luF) - if verbose println("KF factorisation failed step $t") end + if opts.verbose println("KF factorisation failed step $t") end return on_failure_loglikelihood end Fdet = โ„’.det(luF) - - # Early return if determinant is too small, indicating numerical instability. if Fdet < eps(Float64) - if verbose println("KF factorisation failed step $t") end + if opts.verbose println("KF factorisation failed step $t") end return on_failure_loglikelihood end - invF = inv(luF) ### - if t > presample_periods - loglik += log(Fdet) + โ„’.dot(v, invF, v)### + โ„’.ldiv!(w, luF, v) + loglik += log(Fdet) + โ„’.dot(v, w) end - K = P * C' * invF - - P = A * (P - K * C * P) * A' + ๐ - - u = A * (u + K * v) - - z = C * u + invF = inv(luF) + โ„’.mul!(PCt, P, C') + โ„’.mul!(K, PCt, invF) + + # P = A * (P - K * C * P) * A' + ๐ + โ„’.mul!(KC, K, C) + โ„’.mul!(PmKCP, KC, P) + โ„’.axpby!(1, P, -1, PmKCP) # PmKCP = P - K*C*P + โ„’.mul!(AP, A, PmKCP) + โ„’.mul!(P, AP, A') + โ„’.axpy!(1, ๐, P) # P += ๐ + + # u = A * (u + K * v) + โ„’.mul!(Kv, K, v) + copyto!(uKv, u) + โ„’.axpy!(1, Kv, uKv) # uKv = u + K*v + โ„’.mul!(u, A, uKv) # u = A*(u + K*v) + โ„’.mul!(z, C, u) end - # end # timeit_debug - - return -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 + return -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 end diff --git a/src/custom_autodiff_rules/rrules.jl b/src/custom_autodiff_rules/rrules.jl new file mode 100644 index 000000000..df5f08072 --- /dev/null +++ b/src/custom_autodiff_rules/rrules.jl @@ -0,0 +1,11296 @@ +# Zygote/ChainRulesCore rrule definitions for reverse-mode automatic differentiation +# +# This file centralizes rrule definitions for computing gradients via reverse-mode AD. +# Each rrule specifies how to propagate gradients backward through custom functions. +# +# Strategy for each rrule: +# 1. Compute the forward pass and store necessary intermediate values +# 2. Return the result and a pullback function +# 3. The pullback computes gradients w.r.t. inputs given upstream gradients +# 4. Use implicit differentiation for iterative solvers and matrix equations +# +# Functions covered: +# - Basic operations: mul_reverse_AD!, mat_mult_kron, sparse_preallocated! +# - Steady states: get_NSSS_and_parameters, calculate_second/third_order_stochastic_steady_state +# - Derivatives: calculate_jacobian, calculate_hessian, calculate_third_order_derivatives +# - Solutions: calculate_first/second/third_order_solution +# - Matrix equations: solve_sylvester_equation, solve_lyapunov_equation +# - Filters: calculate_loglikelihood, run_kalman_iterations, find_shocks + +function rrule(::typeof(mul_reverse_AD!), + C::Matrix{S}, + A::AbstractMatrix{M}, + B::AbstractMatrix{N}) where {S <: Real, M <: Real, N <: Real} + project_A = ProjectTo(A) + project_B = ProjectTo(B) + + function times_pullback(ศณ) + ศฒ = unthunk(ศณ) + dA = @thunk(project_A(ศฒ * B')) + dB = @thunk(project_B(A' * ศฒ)) + return (NoTangent(), NoTangent(), dA, dB) + end + + return โ„’.mul!(C,A,B), times_pullback +end + +function rrule(::typeof(mat_mult_kron), + A::AbstractSparseMatrix{R}, + B::AbstractMatrix{T}, + C::AbstractMatrix{T}, + D::AbstractMatrix{S}) where {R <: Real, T <: Real, S <: Real} + Y = mat_mult_kron(A, B, C, D) + + function mat_mult_kron_pullback(ศฒ) + if ศฒ isa AbstractZero + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + ศฒdense = Matrix(ศฒ) + + n_rowB = size(B, 1) + n_colB = size(B, 2) + n_rowC = size(C, 1) + n_colC = size(C, 2) + + G = promote_type(eltype(B), eltype(C), eltype(D), Float64) + + โˆ‚B = zeros(G, size(B)) + โˆ‚C = zeros(G, size(C)) + โˆ‚D = zeros(G, size(D)) + + A_csc = A isa SparseMatrixCSC ? A : A.A + nnzA = nnz(A_csc) + nz_col = Vector{Int}(undef, nnzA) + row_to_nzinds = Dict{Int, Vector{Int}}() + + for col in 1:size(A_csc, 2) + for k in A_csc.colptr[col]:(A_csc.colptr[col + 1] - 1) + nz_col[k] = col + r = A_csc.rowval[k] + push!(get!(row_to_nzinds, r, Int[]), k) + end + end + + โˆ‚A_nz = zeros(G, nnzA) + Abar_vec = zeros(G, size(A_csc, 2)) + + for (r, ks) in row_to_nzinds + fill!(Abar_vec, zero(G)) + @inbounds for k in ks + Abar_vec[nz_col[k]] = A_csc.nzval[k] + end + + Abar = reshape(Abar_vec, n_rowC, n_rowB) + AbarB = Abar * B + CAbarB = C' * AbarB + vCAbarB = vec(CAbarB) + + g_row = collect(@view ศฒdense[r, :]) + + โˆ‚D .+= vCAbarB * g_row' + + vCAbarBฬ„ = D * g_row + CAbarBฬ„ = reshape(vCAbarBฬ„, n_colC, n_colB) + + โˆ‚C .+= AbarB * CAbarBฬ„' + + AbarBฬ„ = C * CAbarBฬ„ + โˆ‚B .+= Abar' * AbarBฬ„ + + Abarฬ„ = AbarBฬ„ * B' + vecAbarฬ„ = vec(Abarฬ„) + @inbounds for k in ks + โˆ‚A_nz[k] += vecAbarฬ„[nz_col[k]] + end + end + + โˆ‚A_csc = SparseMatrixCSC(size(A_csc, 1), size(A_csc, 2), copy(A_csc.colptr), copy(A_csc.rowval), โˆ‚A_nz) + + return NoTangent(), + ProjectTo(A)(โˆ‚A_csc), + ProjectTo(B)(โˆ‚B), + ProjectTo(C)(โˆ‚C), + ProjectTo(D)(โˆ‚D) + end + + return Y, mat_mult_kron_pullback +end + + + +function rrule(::typeof(sparse_preallocated!), ลœ::Matrix{T}; โ„‚::higher_order_workspace{T,F,H} = Higher_order_workspace()) where {T <: Real, F <: AbstractFloat, H <: Real} + project_ลœ = ProjectTo(ลœ) + + function sparse_preallocated_pullback(ฮฉฬ„) + ฮ”ฮฉ = unthunk(ฮฉฬ„) + ฮ”ลœ = project_ลœ(ฮ”ฮฉ) + return NoTangent(), ฮ”ลœ, NoTangent() + end + + return sparse_preallocated!(ลœ, โ„‚ = โ„‚), sparse_preallocated_pullback +end + +function rrule(::typeof(solve_stochastic_steady_state_newton), + ::Val{:second_order}, + ๐’โ‚::Matrix{Float64}, + ๐’โ‚‚::AbstractSparseMatrix{Float64}, + x::Vector{Float64}, + ๐“‚::โ„ณ; + # timer::TimerOutput = TimerOutput(), + tol::AbstractFloat = 1e-14) + # @timeit_debug timer "Calculate SSS - forward" begin + # @timeit_debug timer "Setup indices" begin + + # Get cached computational constants + constants = initialise_constants!(๐“‚) + so = constants.second_order + T = constants.post_model_macro + s_in_sโบ = so.s_in_sโบ + s_in_s = so.s_in_s + I_nPast = T.I_nPast + + kron_sโบ_sโบ = so.kron_sโบ_sโบ + + kron_sโบ_s = so.kron_sโบ_s + + A = ๐’โ‚[T.past_not_future_and_mixed_idx,1:T.nPast_not_future_and_mixed] + B = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_s] + Bฬ‚ = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] + + # end # timeit_debug + + # @timeit_debug timer "Iterations" begin + + max_iters = 100 + # SSS .= ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 + x_aug = Vector{Float64}(undef, length(x) + 1) + x_aug[end] = 1.0 + + for i in 1:max_iters + copyto!(x_aug, 1, x, 1, length(x)) + kron_x_aug = โ„’.kron(x_aug, x_aug) + + โˆ‚x = (A + B * โ„’.kron(x_aug, I_nPast) - I_nPast) + + โˆ‚xฬ‚ = โ„’.lu!(โˆ‚x, check = false) + + if !โ„’.issuccess(โˆ‚xฬ‚) + return x, false + end + + ฮ”x = โˆ‚xฬ‚ \ (A * x + Bฬ‚ * kron_x_aug / 2 - x) + + if i > 5 && isapprox(A * x + Bฬ‚ * kron_x_aug / 2, x, rtol = tol) + break + end + + # x += ฮ”x + โ„’.axpy!(-1, ฮ”x, x) + end + copyto!(x_aug, 1, x, 1, length(x)) + kron_x_aug = โ„’.kron(x_aug, x_aug) + solved = isapprox(A * x + Bฬ‚ * kron_x_aug / 2, x, rtol = tol) + + โˆ‚๐’โ‚ = zero(๐’โ‚) + โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) + + # end # timeit_debug + # end # timeit_debug + + function second_order_stochastic_steady_state_pullback(โˆ‚x) + # @timeit_debug timer "Calculate SSS - pullback" begin + S = -โˆ‚x[1]' / (A + B * โ„’.kron(x_aug, I_nPast) - I_nPast) + + โˆ‚๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] = S' * x' + + โˆ‚๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] = S' * kron_x_aug' / 2 + # end # timeit_debug + + return NoTangent(), NoTangent(), โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, NoTangent(), NoTangent(), NoTangent() + end + + return (x, solved), second_order_stochastic_steady_state_pullback +end + + +function rrule(::typeof(solve_stochastic_steady_state_newton), + ::Val{:third_order}, + ๐’โ‚::Matrix{Float64}, + ๐’โ‚‚::AbstractSparseMatrix{Float64}, + ๐’โ‚ƒ::AbstractSparseMatrix{Float64}, + x::Vector{Float64}, + ๐“‚::โ„ณ; + tol::AbstractFloat = 1e-14) + # Get cached computational constants + so = ensure_computational_constants!(๐“‚.constants) + T = ๐“‚.constants.post_model_macro + s_in_sโบ = so.s_in_sโบ + s_in_s = so.s_in_s + I_nPast = T.I_nPast + + kron_sโบ_sโบ = so.kron_sโบ_sโบ + + kron_sโบ_s = so.kron_sโบ_s + + kron_sโบ_sโบ_sโบ = so.kron_sโบ_sโบ_sโบ + + kron_s_sโบ_sโบ = so.kron_s_sโบ_sโบ + + A = ๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] + B = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_s] + Bฬ‚ = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] + C = ๐’โ‚ƒ[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_s_sโบ_sโบ] + ฤˆ = ๐’โ‚ƒ[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ_sโบ] + + max_iters = 100 + # SSS .= ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 + x_aug = Vector{Float64}(undef, length(x) + 1) + x_aug[end] = 1.0 + + for i in 1:max_iters + copyto!(x_aug, 1, x, 1, length(x)) + kron_x_aug = โ„’.kron(x_aug, x_aug) + kron_x_kron = โ„’.kron(x_aug, kron_x_aug) + + โˆ‚x = (A + B * โ„’.kron(x_aug, I_nPast) + C * โ„’.kron(kron_x_aug, I_nPast) / 2 - I_nPast) + + โˆ‚xฬ‚ = โ„’.lu!(โˆ‚x, check = false) + + if !โ„’.issuccess(โˆ‚xฬ‚) + return x, false + end + + ฮ”x = โˆ‚xฬ‚ \ (A * x + Bฬ‚ * kron_x_aug / 2 + ฤˆ * kron_x_kron / 6 - x) + + if i > 5 && isapprox(A * x + Bฬ‚ * kron_x_aug / 2 + ฤˆ * kron_x_kron / 6, x, rtol = tol) + break + end + + # x += ฮ”x + โ„’.axpy!(-1, ฮ”x, x) + end + + copyto!(x_aug, 1, x, 1, length(x)) + kron_x_aug = โ„’.kron(x_aug, x_aug) + kron_x_kron = โ„’.kron(x_aug, kron_x_aug) + solved = isapprox(A * x + Bฬ‚ * kron_x_aug / 2 + ฤˆ * kron_x_kron / 6, x, rtol = tol) + + โˆ‚๐’โ‚ = zero(๐’โ‚) + โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) + โˆ‚๐’โ‚ƒ = zero(๐’โ‚ƒ) + + function third_order_stochastic_steady_state_pullback(โˆ‚x) + S = -โˆ‚x[1]' / (A + B * โ„’.kron(x_aug, I_nPast) + C * โ„’.kron(kron_x_aug, I_nPast) / 2 - I_nPast) + + โˆ‚๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] = S' * x' + + โˆ‚๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] = S' * kron_x_aug' / 2 + + โˆ‚๐’โ‚ƒ[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ_sโบ] = S' * kron_x_kron' / 6 + + return NoTangent(), NoTangent(), โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, โˆ‚๐’โ‚ƒ, NoTangent(), NoTangent(), NoTangent() + end + + return (x, solved), third_order_stochastic_steady_state_pullback +end + + +function rrule(::typeof(calculate_jacobian), + parameters, + SS_and_pars, + caches_obj::caches, + jacobian_funcs::jacobian_functions, + workspaces::workspaces) + jacobian = calculate_jacobian(parameters, SS_and_pars, caches_obj, jacobian_funcs, workspaces) + โˆ‚โˆ‡โ‚_vec = ensure_first_order_cotangent_buffer!(workspaces.first_order, length(jacobian)) + + function calculate_jacobian_pullback(โˆ‚โˆ‡โ‚) + if โˆ‚โˆ‡โ‚ isa Union{NoTangent, AbstractZero} + return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() + end + + โˆ‚โˆ‡โ‚u = unthunk(โˆ‚โˆ‡โ‚) + copyto!(โˆ‚โˆ‡โ‚_vec, โˆ‚โˆ‡โ‚u) + + jacobian_funcs.f_parameters(caches_obj.jacobian_parameters, parameters, SS_and_pars) + jacobian_funcs.f_SS_and_pars(caches_obj.jacobian_SS_and_pars, parameters, SS_and_pars) + + โˆ‚parameters = caches_obj.jacobian_parameters * โˆ‚โˆ‡โ‚_vec + โˆ‚SS_and_pars = caches_obj.jacobian_SS_and_pars * โˆ‚โˆ‡โ‚_vec + return NoTangent(), โˆ‚parameters, โˆ‚SS_and_pars, NoTangent(), NoTangent(), NoTangent() + end + + return jacobian, calculate_jacobian_pullback +end + + +function rrule(::typeof(calculate_hessian), + parameters, + SS_and_pars, + caches_obj::caches, + hessian_funcs::hessian_functions, + workspaces::workspaces) + hessian = calculate_hessian(parameters, SS_and_pars, caches_obj, hessian_funcs, workspaces) + โˆ‚โˆ‡โ‚‚_vec = ensure_higher_order_cotangent_buffer!(workspaces.second_order, length(hessian)) + + function calculate_hessian_pullback(โˆ‚โˆ‡โ‚‚) + if โˆ‚โˆ‡โ‚‚ isa Union{NoTangent, AbstractZero} + return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() + end + + โˆ‚โˆ‡โ‚‚u = unthunk(โˆ‚โˆ‡โ‚‚) + copyto!(โˆ‚โˆ‡โ‚‚_vec, โˆ‚โˆ‡โ‚‚u) + + hessian_funcs.f_parameters(caches_obj.hessian_parameters, parameters, SS_and_pars) + hessian_funcs.f_SS_and_pars(caches_obj.hessian_SS_and_pars, parameters, SS_and_pars) + + โˆ‚parameters = caches_obj.hessian_parameters * โˆ‚โˆ‡โ‚‚_vec + โˆ‚SS_and_pars = caches_obj.hessian_SS_and_pars * โˆ‚โˆ‡โ‚‚_vec + + return NoTangent(), โˆ‚parameters, โˆ‚SS_and_pars, NoTangent(), NoTangent(), NoTangent() + end + + return hessian, calculate_hessian_pullback +end + + +function rrule(::typeof(calculate_third_order_derivatives), + parameters, + SS_and_pars, + caches_obj::caches, + third_order_derivatives_funcs::third_order_derivatives_functions, + workspaces::workspaces) + third_order_derivatives = calculate_third_order_derivatives(parameters, SS_and_pars, caches_obj, third_order_derivatives_funcs, workspaces) + โˆ‚โˆ‡โ‚ƒ_vec = ensure_higher_order_cotangent_buffer!(workspaces.third_order, length(third_order_derivatives)) + + function calculate_third_order_derivatives_pullback(โˆ‚โˆ‡โ‚ƒ) + if โˆ‚โˆ‡โ‚ƒ isa Union{NoTangent, AbstractZero} + return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() + end + + โˆ‚โˆ‡โ‚ƒu = unthunk(โˆ‚โˆ‡โ‚ƒ) + copyto!(โˆ‚โˆ‡โ‚ƒ_vec, โˆ‚โˆ‡โ‚ƒu) + + third_order_derivatives_funcs.f_parameters(caches_obj.third_order_derivatives_parameters, parameters, SS_and_pars) + third_order_derivatives_funcs.f_SS_and_pars(caches_obj.third_order_derivatives_SS_and_pars, parameters, SS_and_pars) + + โˆ‚parameters = caches_obj.third_order_derivatives_parameters * โˆ‚โˆ‡โ‚ƒ_vec + โˆ‚SS_and_pars = caches_obj.third_order_derivatives_SS_and_pars * โˆ‚โˆ‡โ‚ƒ_vec + + return NoTangent(), โˆ‚parameters, โˆ‚SS_and_pars, NoTangent(), NoTangent(), NoTangent() + end + + return third_order_derivatives, calculate_third_order_derivatives_pullback +end + + +function _incremental_cotangent!(ฮ”, prev_ref::Base.RefValue) + if ฮ” isa Union{NoTangent, AbstractZero} + return ฮ” + end + + ฮ”u = unthunk(ฮ”) + prev = prev_ref[] + prev_ref[] = copy(ฮ”u) + + if prev === nothing + return ฮ”u + end + + return ฮ”u .- prev +end + +function rrule(::typeof(get_NSSS_and_parameters), + ๐“‚::โ„ณ, + parameter_values::Vector{S}; + opts::CalculationOptions = merge_calculation_options(), + cold_start::Bool = false, + estimation::Bool = false) where S <: Real + # timer::TimerOutput = TimerOutput(), + # @timeit_debug timer "Calculate NSSS - forward" begin + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + + # Use custom steady state function if available, otherwise use default solver + if ๐“‚.functions.NSSS_custom isa Function + vars_in_ss_equations = ms.vars_in_ss_equations + expected_length = length(vars_in_ss_equations) + length(๐“‚.equations.calibration_parameters) + + SS_and_pars_tmp = evaluate_custom_steady_state_function( + ๐“‚, + parameter_values, + expected_length, + length(๐“‚.constants.post_complete_parameters.parameters), + ) + + residual = zeros(length(๐“‚.equations.steady_state) + length(๐“‚.equations.calibration)) + + ๐“‚.functions.NSSS_check(residual, parameter_values, SS_and_pars_tmp) + + solution_error = โ„’.norm(residual) + + iters = 0 + + # if !isfinite(solution_error) || solution_error > opts.tol.nsss.acceptance_tol + # throw(ArgumentError("Custom steady state function failed steady state check: residual $solution_error > $(opts.tol.nsss.acceptance_tol). Parameters: $(parameter_values). Steady state and parameters returned: $(SS_and_pars_tmp).")) + # end + X = ms.custom_ss_expand_matrix + SS_and_pars = X * SS_and_pars_tmp + else + fastest_idx = ๐“‚.constants.post_complete_parameters.nsss_fastest_solver_parameter_idx + preferred_solver_parameter_idx = fastest_idx < 1 || fastest_idx > length(DEFAULT_SOLVER_PARAMETERS) ? 1 : fastest_idx + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(parameter_values, ๐“‚, opts.tol, opts.verbose, cold_start, DEFAULT_SOLVER_PARAMETERS, preferred_solver_parameter_idx = preferred_solver_parameter_idx) + end + + # end # timeit_debug + + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) + # Update failed counter + update_ss_counter!(๐“‚.counters, false, estimation = estimation) + return (SS_and_pars, (solution_error, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # Update success counter + update_ss_counter!(๐“‚.counters, true, estimation = estimation) + + # @timeit_debug timer "Calculate NSSS - pullback" begin + + custom_ss_expand_matrix = ms.custom_ss_expand_matrix + + โˆ‚ = parameter_values + C = SS_and_pars[ms.SS_and_pars_no_exo_idx] # [dyn_ss_idx]) + + if eltype(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters) != eltype(parameter_values) + if ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters isa SparseMatrixCSC + jac_cache = similar(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters, eltype(parameter_values)) + jac_cache.nzval .= 0 + else + jac_cache = zeros(eltype(parameter_values), size(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters)) + end + else + jac_cache = ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚parameters + end + + if jac_cache isa SparseMatrixCSC + jac_cache.nzval .= 0 + else + fill!(jac_cache, zero(eltype(jac_cache))) + end + + ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters(jac_cache, โˆ‚, C) + + โˆ‚SS_equations_โˆ‚parameters = jac_cache + + + if eltype(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars) != eltype(SS_and_pars) + if ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars isa SparseMatrixCSC + jac_cache = similar(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars, eltype(SS_and_pars)) + jac_cache.nzval .= 0 + else + jac_cache = zeros(eltype(SS_and_pars), size(๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars)) + end + else + jac_cache = ๐“‚.caches.NSSS_โˆ‚equations_โˆ‚SS_and_pars + end + + if jac_cache isa SparseMatrixCSC + jac_cache.nzval .= 0 + else + fill!(jac_cache, zero(eltype(jac_cache))) + end + + ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars(jac_cache, โˆ‚, C) + + โˆ‚SS_equations_โˆ‚SS_and_pars = jac_cache + qme_ws = ๐“‚.workspaces.first_order + if โˆ‚SS_equations_โˆ‚SS_and_pars isa SparseMatrixCSC + rhs_n_rows = size(โˆ‚SS_equations_โˆ‚SS_and_pars, 1) + rhs_n_cols = size(โˆ‚SS_equations_โˆ‚parameters, 2) + + if length(qme_ws.nsss_sparse_rhs) != rhs_n_rows + qme_ws.nsss_sparse_rhs = zeros(eltype(SS_and_pars), rhs_n_rows) + end + + if size(qme_ws.nsss_jvp_rhs) != (rhs_n_rows, rhs_n_cols) + qme_ws.nsss_jvp_rhs = zeros(eltype(SS_and_pars), rhs_n_rows, rhs_n_cols) + end + + if size(qme_ws.nsss_sparse_lu_buffer.A) != (rhs_n_rows, rhs_n_rows) + sparse_prob = ๐’ฎ.LinearProblem(โˆ‚SS_equations_โˆ‚SS_and_pars, qme_ws.nsss_sparse_rhs) + qme_ws.nsss_sparse_lu_buffer = ๐’ฎ.init(sparse_prob, + ๐’ฎ.LUFactorization(), + verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + else + qme_ws.nsss_sparse_lu_buffer.A = โˆ‚SS_equations_โˆ‚SS_and_pars + end + + for j in 1:rhs_n_cols + @views copyto!(qme_ws.nsss_sparse_rhs, โˆ‚SS_equations_โˆ‚parameters[:, j]) + qme_ws.nsss_sparse_lu_buffer.b = qme_ws.nsss_sparse_rhs + sparse_sol = ๐’ฎ.solve!(qme_ws.nsss_sparse_lu_buffer) + + if sparse_sol.retcode != ๐’ฎ.SciMLBase.ReturnCode.Default && !๐’ฎ.SciMLBase.successful_retcode(sparse_sol.retcode) + return (SS_and_pars, (10.0, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + @views copyto!(qme_ws.nsss_jvp_rhs[:, j], qme_ws.nsss_sparse_lu_buffer.u) + end + + โ„’.rmul!(qme_ws.nsss_jvp_rhs, -1) + JVP = qme_ws.nsss_jvp_rhs + else + qme_ws.fast_lu_ws_nsss, qme_ws.fast_lu_dims_nsss, solved_nsss, nsss_lu = factorize_lu!(โˆ‚SS_equations_โˆ‚SS_and_pars, + qme_ws.fast_lu_ws_nsss, + qme_ws.fast_lu_dims_nsss) + + if !solved_nsss + return (SS_and_pars, (10.0, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + rhs_dense = โˆ‚SS_equations_โˆ‚parameters isa Matrix ? โˆ‚SS_equations_โˆ‚parameters : Matrix(โˆ‚SS_equations_โˆ‚parameters) + + if size(qme_ws.nsss_jvp_rhs) != size(rhs_dense) + qme_ws.nsss_jvp_rhs = zeros(eltype(rhs_dense), size(rhs_dense)) + end + copyto!(qme_ws.nsss_jvp_rhs, rhs_dense) + + solve_lu_left!(โˆ‚SS_equations_โˆ‚SS_and_pars, + qme_ws.nsss_jvp_rhs, + qme_ws.fast_lu_ws_nsss, + nsss_lu) + + โ„’.rmul!(qme_ws.nsss_jvp_rhs, -1) + JVP = qme_ws.nsss_jvp_rhs + end + + jvp_no_exo = custom_ss_expand_matrix * JVP + + # end # timeit_debug + # end # timeit_debug + + # try block-gmres here + function get_non_stochastic_steady_state_pullback(โˆ‚SS_and_pars) + โˆ‚SS = โˆ‚SS_and_pars[1] + if โˆ‚SS isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, size(jvp_no_exo, 2)), NoTangent() + end + return NoTangent(), NoTangent(), jvp_no_exo' * โˆ‚SS, NoTangent() + end + + + return (SS_and_pars, (solution_error, iters)), get_non_stochastic_steady_state_pullback +end + +function rrule(::typeof(get_relevant_steady_state_and_state_update), + ::Val{:first_order}, + parameter_values::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) where S <: AbstractFloat + constants_obj = initialise_constants!(๐“‚) + + nsss_out, nsss_pb = rrule(get_NSSS_and_parameters, + ๐“‚, + parameter_values; + opts = opts, + estimation = estimation) + + SS_and_pars = nsss_out[1] + solution_error = nsss_out[2][1] + + state = zeros(S, ๐“‚.constants.post_model_macro.nVars) + + if solution_error > opts.tol.nsss.acceptance_tol + y = (๐“‚.constants, SS_and_pars, zeros(S, 0, 0), [state], false) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + nsss_grads = nsss_pb((ฮ”SS_and_pars, NoTangent())) + โˆ‚parameter_values = nsss_grads[3] + + return NoTangent(), NoTangent(), โˆ‚parameter_values, NoTangent() + end + + return y, pullback + end + + โˆ‡โ‚, jac_pb = rrule(calculate_jacobian, + parameter_values, + SS_and_pars, + ๐“‚.caches, + ๐“‚.functions.jacobian, + ๐“‚.workspaces) + + first_out, first_pb = rrule(calculate_first_order_solution, + โˆ‡โ‚, + constants_obj, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameter_values) + + ๐’โ‚ = first_out[1] + solved = first_out[3] + + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + + if !solved + y = (๐“‚.constants, SS_and_pars, zeros(S, 0, 0), [state], false) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + + nsss_grads = nsss_pb((ฮ”SS_and_pars, NoTangent())) + โˆ‚parameter_values = nsss_grads[3] + + return NoTangent(), NoTangent(), โˆ‚parameter_values, NoTangent() + end + + return y, pullback + end + + y = (๐“‚.constants, SS_and_pars, ๐’โ‚, [state], true) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + ฮ”๐’โ‚ = ฮ”y[3] + + # When the caller passes NoTangent for the solution matrix cotangent + # (e.g. filter failure), skip the first-order solution pullback and + # only propagate through the steady-state. + if ฮ”๐’โ‚ isa Union{NoTangent, AbstractZero} + nsss_grads = nsss_pb((ฮ”SS_and_pars, NoTangent())) + return NoTangent(), NoTangent(), nsss_grads[3], NoTangent() + end + + first_grads = first_pb((ฮ”๐’โ‚, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚ = first_grads[2] + + jac_grads = jac_pb(โˆ‚โˆ‡โ‚) + โˆ‚parameter_values = jac_grads[2] + โˆ‚SS_and_pars_from_jac = jac_grads[3] + + nsss_grads = nsss_pb((ฮ”SS_and_pars + โˆ‚SS_and_pars_from_jac, NoTangent())) + โˆ‚parameter_values .+= nsss_grads[3] + + return NoTangent(), NoTangent(), โˆ‚parameter_values, NoTangent() + end + + return y, pullback +end + +function rrule(::typeof(_prepare_stochastic_steady_state_base_terms), + parameters::Vector{Float64}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) + constants = initialise_constants!(๐“‚) + T = constants.post_model_macro + nVars = T.nVars + nPast = T.nPast_not_future_and_mixed + nExo = T.nExo + past_idx = T.past_not_future_and_mixed_idx + + (SS_and_pars, (solution_error, iters)), nsss_pullback = + rrule(get_NSSS_and_parameters, ๐“‚, parameters, opts = opts, estimation = estimation) + + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) + common = (false, + zeros(Float64, nVars), + SS_and_pars, + solution_error, + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0), + constants) + pullback = function (ฮ”common) + return NoTangent(), zeros(Float64, length(parameters)), NoTangent() + end + return common, pullback + end + + ms = ensure_model_structure_constants!(constants, ๐“‚.equations.calibration_parameters) + all_SS = expand_steady_state(SS_and_pars, ms) + + โˆ‡โ‚, jacobian_pullback = + rrule(calculate_jacobian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) + + (๐’โ‚_raw, qme_sol, solved), first_order_pullback = + rrule(calculate_first_order_solution, โˆ‡โ‚, constants, ๐“‚.workspaces, ๐“‚.caches; + opts = opts, initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameters) + + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + + if !solved + common = (false, + all_SS, + SS_and_pars, + solution_error, + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0), + constants) + pullback = function (ฮ”common) + return NoTangent(), zeros(Float64, length(parameters)), NoTangent() + end + return common, pullback + end + + โˆ‡โ‚‚, hessian_pullback = + rrule(calculate_hessian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces) + + (๐’โ‚‚_raw, solved2), second_order_pullback = + rrule(calculate_second_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚_raw, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + initial_guess = ๐“‚.caches.second_order_solution, opts = opts, + parameter_values = parameters) + + update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) + + if !solved2 + common = (false, + all_SS, + SS_and_pars, + solution_error, + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0), + constants) + pullback = function (ฮ”common) + return NoTangent(), zeros(Float64, length(parameters)), NoTangent() + end + return common, pullback + end + + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{Float64, Int} + + ๐’โ‚ = [๐’โ‚_raw[:, 1:nPast] zeros(nVars) ๐’โ‚_raw[:, nPast+1:end]] + aug_stateโ‚ = sparse([zeros(nPast); 1; zeros(nExo)]) + kron_aug1 = โ„’.kron(aug_stateโ‚, aug_stateโ‚) + + tmp = (T.I_nPast - ๐’โ‚[past_idx, 1:nPast]) + tmpฬ„_lu = โ„’.lu(tmp, check = false) + + if !โ„’.issuccess(tmpฬ„_lu) + common = (false, + all_SS, + SS_and_pars, + solution_error, + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0,0), + spzeros(Float64,0,0), + zeros(Float64,0), + constants) + pullback = function (ฮ”common) + return NoTangent(), zeros(Float64, length(parameters)), NoTangent() + end + return common, pullback + end + + SSSstates = collect(tmpฬ„_lu \ (๐’โ‚‚ * kron_aug1 / 2)[past_idx]) + + common = (true, + all_SS, + SS_and_pars, + solution_error, + โˆ‡โ‚, + โˆ‡โ‚‚, + ๐’โ‚, + ๐’โ‚‚_raw, + SSSstates, + constants) + + pullback = function (ฮ”common) + โˆ‚all_SS = zeros(Float64, length(all_SS)) + โˆ‚SS_and_pars_direct = zeros(Float64, length(SS_and_pars)) + โˆ‚โˆ‡โ‚_direct = zeros(Float64, size(โˆ‡โ‚)) + โˆ‚โˆ‡โ‚‚_direct = zeros(Float64, size(โˆ‡โ‚‚)) + โˆ‚๐’โ‚_aug = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚‚_raw_total = zeros(Float64, size(๐’โ‚‚_raw)) + โˆ‚SSSstates = zeros(Float64, length(SSSstates)) + + if !(ฮ”common isa Union{NoTangent, AbstractZero}) + v2 = ฮ”common[2] + v3 = ฮ”common[3] + v5 = ฮ”common[5] + v6 = ฮ”common[6] + v7 = ฮ”common[7] + v8 = ฮ”common[8] + v9 = ฮ”common[9] + โˆ‚all_SS = v2 isa Union{NoTangent, AbstractZero} ? โˆ‚all_SS : v2 + โˆ‚SS_and_pars_direct = v3 isa Union{NoTangent, AbstractZero} ? โˆ‚SS_and_pars_direct : v3 + โˆ‚โˆ‡โ‚_direct = v5 isa Union{NoTangent, AbstractZero} ? โˆ‚โˆ‡โ‚_direct : v5 + โˆ‚โˆ‡โ‚‚_direct = v6 isa Union{NoTangent, AbstractZero} ? โˆ‚โˆ‡โ‚‚_direct : v6 + โˆ‚๐’โ‚_aug = v7 isa Union{NoTangent, AbstractZero} ? โˆ‚๐’โ‚_aug : v7 + โˆ‚๐’โ‚‚_raw_total = v8 isa Union{NoTangent, AbstractZero} ? โˆ‚๐’โ‚‚_raw_total : v8 + โˆ‚SSSstates = v9 isa Union{NoTangent, AbstractZero} ? โˆ‚SSSstates : v9 + end + + if !isempty(โˆ‚SSSstates) + โˆ‚rhs = tmpฬ„_lu' \ โˆ‚SSSstates + โˆ‚tmp = -(tmpฬ„_lu' \ โˆ‚SSSstates) * SSSstates' + โˆ‚๐’โ‚_aug[past_idx, 1:nPast] .-= โˆ‚tmp + โˆ‚๐’โ‚‚_from_rhs = spzeros(Float64, size(๐’โ‚‚)...) + โˆ‚๐’โ‚‚_from_rhs[past_idx, :] += โˆ‚rhs * kron_aug1' / 2 + โˆ‚๐’โ‚‚_raw_total += โˆ‚๐’โ‚‚_from_rhs * ๐”โ‚‚' + end + + X = ms.steady_state_expand_matrix + โˆ‚SS_and_pars_from_allSS = X' * โˆ‚all_SS + + โˆ‚๐’โ‚_raw = hcat(โˆ‚๐’โ‚_aug[:, 1:nPast], โˆ‚๐’โ‚_aug[:, nPast+2:end]) + + so2_tangents = second_order_pullback((โˆ‚๐’โ‚‚_raw_total, NoTangent())) + โˆ‚โˆ‡โ‚_from_so2 = so2_tangents[2] + โˆ‚โˆ‡โ‚‚_from_so2 = so2_tangents[3] + โˆ‚๐’โ‚_raw_from_so2 = so2_tangents[4] + + โˆ‚โˆ‡โ‚‚_total = โˆ‚โˆ‡โ‚‚_from_so2 + โˆ‚โˆ‡โ‚‚_direct + hess_tangents = hessian_pullback(โˆ‚โˆ‡โ‚‚_total) + โˆ‚params_from_hess = hess_tangents[2] + โˆ‚SS_and_pars_from_hess = hess_tangents[3] + + โˆ‚๐’โ‚_raw_total = โˆ‚๐’โ‚_raw + โˆ‚๐’โ‚_raw_from_so2 + fo_tangents = first_order_pullback((โˆ‚๐’โ‚_raw_total, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚_from_fo = fo_tangents[2] + + โˆ‚โˆ‡โ‚_total = โˆ‚โˆ‡โ‚_from_so2 + โˆ‚โˆ‡โ‚_from_fo + โˆ‚โˆ‡โ‚_direct + jac_tangents = jacobian_pullback(โˆ‚โˆ‡โ‚_total) + โˆ‚params_from_jac = jac_tangents[2] + โˆ‚SS_and_pars_from_jac = jac_tangents[3] + + โˆ‚SS_and_pars_total = โˆ‚SS_and_pars_from_allSS + โˆ‚SS_and_pars_from_hess + โˆ‚SS_and_pars_from_jac + โˆ‚SS_and_pars_direct + nsss_tangents = nsss_pullback((โˆ‚SS_and_pars_total, NoTangent())) + โˆ‚params_from_nsss = nsss_tangents[3] + + โˆ‚parameters = โˆ‚params_from_nsss + โˆ‚params_from_jac + โˆ‚params_from_hess + + return NoTangent(), โˆ‚parameters, NoTangent() + end + + return common, pullback +end + +function rrule(::typeof(calculate_stochastic_steady_state), + ::Val{:second_order}, + parameters::Vector{Float64}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) + common, common_pullback = rrule(_prepare_stochastic_steady_state_base_terms, + parameters, + ๐“‚; + opts = opts, + estimation = estimation) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + # Expand compressed ๐’โ‚‚_raw to full for stochastic SS computation + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{Float64, Int} + + so = ๐“‚.constants.second_order + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + kron_sโบ_sโบ = so.kron_sโบ_sโบ + A = ๐’โ‚[:,1:nPast] + Bฬ‚ = ๐’โ‚‚[:,kron_sโบ_sโบ] + + (SSSstates_final, converged), newton_pullback = + rrule(solve_stochastic_steady_state_newton, Val(:second_order), ๐’โ‚, ๐’โ‚‚, collect(SSSstates), ๐“‚) + + if !converged + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + state = A * SSSstates_final + Bฬ‚ * โ„’.kron(vcat(SSSstates_final,1), vcat(SSSstates_final,1)) / 2 + sss = all_SS + vec(state) + result = (sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚) + + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(sss)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + ฮ”โˆ‡โ‚ = zeros(Float64, size(โˆ‡โ‚)) + ฮ”โˆ‡โ‚‚ = zeros(Float64, size(โˆ‡โ‚‚)) + ฮ”๐’โ‚ = zeros(Float64, size(๐’โ‚)) + ฮ”๐’โ‚‚ = spzeros(Float64, size(๐’โ‚‚)...) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + v5 = ฮ”[5] + v6 = ฮ”[6] + v7 = ฮ”[7] + v8 = ฮ”[8] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + ฮ”โˆ‡โ‚ = v5 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚ : v5 + ฮ”โˆ‡โ‚‚ = v6 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚‚ : v6 + ฮ”๐’โ‚ = v7 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚ : v7 + ฮ”๐’โ‚‚ = v8 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚‚ : v8 + end + + โˆ‚state_vec = ฮ”sss + aug_sss = vcat(SSSstates_final, 1) + kron_aug = โ„’.kron(aug_sss, aug_sss) + + โˆ‚๐’โ‚_from_state = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚_from_state[:, 1:nPast] += โˆ‚state_vec * SSSstates_final' + + โˆ‚๐’โ‚‚_from_state = spzeros(Float64, size(๐’โ‚‚)...) + โˆ‚๐’โ‚‚_from_state[:, kron_sโบ_sโบ] += โˆ‚state_vec * kron_aug' / 2 + + โˆ‚SSSstates_from_state = A' * โˆ‚state_vec + n_aug = length(aug_sss) + I_aug = Matrix{Float64}(โ„’.I, n_aug, n_aug) + pad = vcat(Matrix{Float64}(โ„’.I, nPast, nPast), zeros(1, nPast)) + dkron_dx = โ„’.kron(I_aug, aug_sss) * pad + โ„’.kron(aug_sss, I_aug) * pad + โˆ‚SSSstates_from_state += (Bฬ‚' * โˆ‚state_vec)' * dkron_dx / 2 |> vec + + newton_tangents = newton_pullback((โˆ‚SSSstates_from_state, NoTangent())) + โˆ‚๐’โ‚_newton = newton_tangents[3] + โˆ‚๐’โ‚‚_newton = newton_tangents[4] + + # Convert full-space โˆ‚๐’โ‚‚ to compressed for common_pullback + โˆ‚๐’โ‚‚_raw_total = (โˆ‚๐’โ‚‚_from_state + โˆ‚๐’โ‚‚_newton + ฮ”๐’โ‚‚) * ๐”โ‚‚' + + common_tangents = common_pullback((NoTangent(), + ฮ”sss, + ฮ”SS_and_pars, + NoTangent(), + ฮ”โˆ‡โ‚, + ฮ”โˆ‡โ‚‚, + โˆ‚๐’โ‚_from_state + โˆ‚๐’โ‚_newton + ฮ”๐’โ‚, + โˆ‚๐’โ‚‚_raw_total, + NoTangent(), + NoTangent())) + + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + + return result, pullback +end + +function rrule(::typeof(calculate_stochastic_steady_state), + ::Val{:pruned_second_order}, + parameters::Vector{Float64}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) + common, common_pullback = rrule(_prepare_stochastic_steady_state_base_terms, + parameters, + ๐“‚; + opts = opts, + estimation = estimation) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + # Expand compressed ๐’โ‚‚_raw to full for stochastic SS computation + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{Float64, Int} + + T = ๐“‚.constants.post_model_macro + nPast = T.nPast_not_future_and_mixed + aug_stateโ‚ = sparse([zeros(nPast); 1; zeros(T.nExo)]) + kron_aug1 = โ„’.kron(aug_stateโ‚, aug_stateโ‚) + + state = ๐’โ‚[:,1:nPast] * SSSstates + ๐’โ‚‚ * kron_aug1 / 2 + sss = all_SS + vec(state) + result = (sss, true, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚) + + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(sss)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + ฮ”โˆ‡โ‚ = zeros(Float64, size(โˆ‡โ‚)) + ฮ”โˆ‡โ‚‚ = zeros(Float64, size(โˆ‡โ‚‚)) + ฮ”๐’โ‚ = zeros(Float64, size(๐’โ‚)) + ฮ”๐’โ‚‚ = spzeros(Float64, size(๐’โ‚‚)...) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + v5 = ฮ”[5] + v6 = ฮ”[6] + v7 = ฮ”[7] + v8 = ฮ”[8] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + ฮ”โˆ‡โ‚ = v5 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚ : v5 + ฮ”โˆ‡โ‚‚ = v6 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚‚ : v6 + ฮ”๐’โ‚ = v7 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚ : v7 + ฮ”๐’โ‚‚ = v8 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚‚ : v8 + end + + โˆ‚state_vec = ฮ”sss + โˆ‚๐’โ‚_from_state = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚_from_state[:, 1:nPast] += โˆ‚state_vec * SSSstates' + โˆ‚๐’โ‚‚_from_state = spzeros(Float64, size(๐’โ‚‚)...) + โˆ‚๐’โ‚‚_from_state += โˆ‚state_vec * kron_aug1' / 2 + โˆ‚SSSstates = ๐’โ‚[:,1:nPast]' * โˆ‚state_vec + + # Convert full-space โˆ‚๐’โ‚‚ to compressed for common_pullback + โˆ‚๐’โ‚‚_raw_total = (โˆ‚๐’โ‚‚_from_state + ฮ”๐’โ‚‚) * ๐”โ‚‚' + + common_tangents = common_pullback((NoTangent(), + ฮ”sss, + ฮ”SS_and_pars, + NoTangent(), + ฮ”โˆ‡โ‚, + ฮ”โˆ‡โ‚‚, + โˆ‚๐’โ‚_from_state + ฮ”๐’โ‚, + โˆ‚๐’โ‚‚_raw_total, + โˆ‚SSSstates, + NoTangent())) + + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + + return result, pullback +end + +function rrule(::typeof(calculate_stochastic_steady_state), + ::Val{:third_order}, + parameters::Vector{Float64}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) + common, common_pullback = rrule(_prepare_stochastic_steady_state_base_terms, + parameters, + ๐“‚; + opts = opts, + estimation = estimation) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{Float64, Int} + + โˆ‡โ‚ƒ, third_derivatives_pullback = + rrule(calculate_third_order_derivatives, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces) + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + ๐’โ‚_raw = [๐’โ‚[:, 1:nPast] ๐’โ‚[:, nPast+2:end]] + + (๐’โ‚ƒ, solved3), third_order_solution_pullback = + rrule(calculate_third_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚_raw, ๐’โ‚‚_raw, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, + parameter_values = parameters) + + if !solved3 + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + ๐”โ‚ƒ = ๐“‚.constants.third_order.๐”โ‚ƒ + ๐’โ‚ƒฬ‚ = sparse(๐’โ‚ƒ * ๐”โ‚ƒ) + + so = ๐“‚.constants.second_order + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + kron_sโบ_sโบ = so.kron_sโบ_sโบ + kron_sโบ_sโบ_sโบ = so.kron_sโบ_sโบ_sโบ + + A = ๐’โ‚[:,1:nPast] + Bฬ‚ = ๐’โ‚‚[:,kron_sโบ_sโบ] + Cฬ‚ = ๐’โ‚ƒฬ‚[:,kron_sโบ_sโบ_sโบ] + + (SSSstates_final, converged), newton_pullback = + rrule(solve_stochastic_steady_state_newton, Val(:third_order), ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚, collect(SSSstates), ๐“‚) + + if !converged + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + aug_sss = vcat(SSSstates_final, 1) + kron_aug = โ„’.kron(aug_sss, aug_sss) + kron_aug3 = โ„’.kron(aug_sss, kron_aug) + + state = A * SSSstates_final + Bฬ‚ * kron_aug / 2 + Cฬ‚ * kron_aug3 / 6 + sss = all_SS + vec(state) + result = (sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚) + + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(sss)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + ฮ”โˆ‡โ‚ = zeros(Float64, size(โˆ‡โ‚)) + ฮ”โˆ‡โ‚‚ = zeros(Float64, size(โˆ‡โ‚‚)) + ฮ”โˆ‡โ‚ƒ = spzeros(Float64, size(โˆ‡โ‚ƒ)...) + ฮ”๐’โ‚ = zeros(Float64, size(๐’โ‚)) + ฮ”๐’โ‚‚ = spzeros(Float64, size(๐’โ‚‚)...) + ฮ”๐’โ‚ƒฬ‚ = spzeros(Float64, size(๐’โ‚ƒฬ‚)...) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + v5 = ฮ”[5] + v6 = ฮ”[6] + v7 = ฮ”[7] + v8 = ฮ”[8] + v9 = ฮ”[9] + v10 = ฮ”[10] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + ฮ”โˆ‡โ‚ = v5 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚ : v5 + ฮ”โˆ‡โ‚‚ = v6 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚‚ : v6 + ฮ”โˆ‡โ‚ƒ = v7 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚ƒ : v7 + ฮ”๐’โ‚ = v8 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚ : v8 + ฮ”๐’โ‚‚ = v9 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚‚ : v9 + ฮ”๐’โ‚ƒฬ‚ = v10 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚ƒฬ‚ : v10 + end + + โˆ‚state_vec = ฮ”sss + + โˆ‚๐’โ‚_from_state = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚_from_state[:, 1:nPast] += โˆ‚state_vec * SSSstates_final' + + โˆ‚๐’โ‚‚_from_state = spzeros(Float64, size(๐’โ‚‚)...) + โˆ‚๐’โ‚‚_from_state[:, kron_sโบ_sโบ] += โˆ‚state_vec * kron_aug' / 2 + + โˆ‚๐’โ‚ƒฬ‚_from_state = spzeros(Float64, size(๐’โ‚ƒฬ‚)...) + โˆ‚๐’โ‚ƒฬ‚_from_state[:, kron_sโบ_sโบ_sโบ] += โˆ‚state_vec * kron_aug3' / 6 + + โˆ‚SSSstates_from_state = A' * โˆ‚state_vec + n_aug = length(aug_sss) + I_aug = Matrix{Float64}(โ„’.I, n_aug, n_aug) + pad = vcat(Matrix{Float64}(โ„’.I, nPast, nPast), zeros(1, nPast)) + dkron_dx = โ„’.kron(I_aug, aug_sss) * pad + โ„’.kron(aug_sss, I_aug) * pad + โˆ‚SSSstates_from_state += (Bฬ‚' * โˆ‚state_vec)' * dkron_dx / 2 |> vec + + dkron3_dx = โ„’.kron(pad, โ„’.kron(aug_sss, aug_sss)) + + โ„’.kron(aug_sss, โ„’.kron(pad, aug_sss)) + + โ„’.kron(aug_sss, โ„’.kron(aug_sss, pad)) + โˆ‚SSSstates_from_state += (Cฬ‚' * โˆ‚state_vec)' * dkron3_dx / 6 |> vec + + newton_tangents = newton_pullback((โˆ‚SSSstates_from_state, NoTangent())) + โˆ‚๐’โ‚_newton = newton_tangents[3] + โˆ‚๐’โ‚‚_newton = newton_tangents[4] + โˆ‚๐’โ‚ƒฬ‚_newton = newton_tangents[5] + + โˆ‚๐’โ‚ƒฬ‚_total = โˆ‚๐’โ‚ƒฬ‚_from_state + โˆ‚๐’โ‚ƒฬ‚_newton + ฮ”๐’โ‚ƒฬ‚ + โˆ‚๐’โ‚ƒ_raw = Matrix(โˆ‚๐’โ‚ƒฬ‚_total) * ๐”โ‚ƒ' + + so3_tangents = third_order_solution_pullback((โˆ‚๐’โ‚ƒ_raw, NoTangent())) + โˆ‚โˆ‡โ‚_from_so3 = so3_tangents[2] isa Union{NoTangent, AbstractZero} ? zero(โˆ‡โ‚) : so3_tangents[2] + โˆ‚โˆ‡โ‚‚_from_so3 = so3_tangents[3] isa Union{NoTangent, AbstractZero} ? zero(โˆ‡โ‚‚) : so3_tangents[3] + โˆ‚โˆ‡โ‚ƒ_from_so3 = so3_tangents[4] isa Union{NoTangent, AbstractZero} ? zero(โˆ‡โ‚ƒ) : so3_tangents[4] + โˆ‚๐’โ‚_raw_from_so3 = so3_tangents[5] isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚_raw) : so3_tangents[5] + โˆ‚๐’โ‚‚_raw_from_so3 = so3_tangents[6] isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚‚_raw) : so3_tangents[6] + + โˆ‚๐’โ‚_from_so3 = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚_from_so3[:, 1:nPast] = โˆ‚๐’โ‚_raw_from_so3[:, 1:nPast] + โˆ‚๐’โ‚_from_so3[:, nPast+2:end] = โˆ‚๐’โ‚_raw_from_so3[:, nPast+1:end] + + โˆ‚โˆ‡โ‚ƒ_total = ฮ”โˆ‡โ‚ƒ + โˆ‚โˆ‡โ‚ƒ_from_so3 + third_derivatives_tangents = third_derivatives_pullback(โˆ‚โˆ‡โ‚ƒ_total) + โˆ‚params_from_โˆ‡โ‚ƒ = third_derivatives_tangents[2] + โˆ‚SS_and_pars_from_โˆ‡โ‚ƒ = third_derivatives_tangents[3] + + # Convert full-space โˆ‚๐’โ‚‚ terms to compressed, then accumulate with compressed โˆ‚๐’โ‚‚_raw_from_so3 + โˆ‚๐’โ‚‚_raw_for_common = โˆ‚๐’โ‚‚_raw_from_so3 + (โˆ‚๐’โ‚‚_from_state + โˆ‚๐’โ‚‚_newton + ฮ”๐’โ‚‚) * ๐”โ‚‚' + + common_tangents = common_pullback((NoTangent(), + ฮ”sss, + ฮ”SS_and_pars + โˆ‚SS_and_pars_from_โˆ‡โ‚ƒ, + NoTangent(), + ฮ”โˆ‡โ‚ + โˆ‚โˆ‡โ‚_from_so3, + ฮ”โˆ‡โ‚‚ + โˆ‚โˆ‡โ‚‚_from_so3, + โˆ‚๐’โ‚_from_state + โˆ‚๐’โ‚_newton + ฮ”๐’โ‚ + โˆ‚๐’โ‚_from_so3, + โˆ‚๐’โ‚‚_raw_for_common, + NoTangent(), + NoTangent())) + + โˆ‚parameters = common_tangents[2] + โˆ‚params_from_โˆ‡โ‚ƒ + return NoTangent(), NoTangent(), โˆ‚parameters, NoTangent() + end + + return result, pullback +end + +function rrule(::typeof(calculate_stochastic_steady_state), + ::Val{:pruned_third_order}, + parameters::Vector{Float64}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) + common, common_pullback = rrule(_prepare_stochastic_steady_state_base_terms, + parameters, + ๐“‚; + opts = opts, + estimation = estimation) + ok, all_SS, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚_raw, SSSstates, _ = common + + if !ok + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{Float64, Int} + + โˆ‡โ‚ƒ, third_derivatives_pullback = + rrule(calculate_third_order_derivatives, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces) + nPast = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + ๐’โ‚_raw = [๐’โ‚[:, 1:nPast] ๐’โ‚[:, nPast+2:end]] + + (๐’โ‚ƒ, solved3), third_order_solution_pullback = + rrule(calculate_third_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚_raw, ๐’โ‚‚_raw, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, + parameter_values = parameters) + + if !solved3 + result = (all_SS, false, SS_and_pars, solution_error, + zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0), zeros(Float64,0,0), spzeros(Float64,0,0), spzeros(Float64,0,0)) + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(all_SS)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + end + common_tangents = common_pullback((NoTangent(), ฮ”sss, ฮ”SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + return NoTangent(), NoTangent(), common_tangents[2], NoTangent() + end + return result, pullback + end + + ๐”โ‚ƒ = ๐“‚.constants.third_order.๐”โ‚ƒ + ๐’โ‚ƒฬ‚ = sparse(๐’โ‚ƒ * ๐”โ‚ƒ) + + T = ๐“‚.constants.post_model_macro + nPast = T.nPast_not_future_and_mixed + aug_stateโ‚ = sparse([zeros(nPast); 1; zeros(T.nExo)]) + kron_aug1 = โ„’.kron(aug_stateโ‚, aug_stateโ‚) + + state = ๐’โ‚[:,1:nPast] * SSSstates + ๐’โ‚‚ * kron_aug1 / 2 + sss = all_SS + vec(state) + result = (sss, true, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒฬ‚) + + pullback = function (ฮ”result) + ฮ” = unthunk(ฮ”result) + ฮ”sss = zeros(Float64, length(sss)) + ฮ”SS_and_pars = zeros(Float64, length(SS_and_pars)) + ฮ”โˆ‡โ‚ = zeros(Float64, size(โˆ‡โ‚)) + ฮ”โˆ‡โ‚‚ = zeros(Float64, size(โˆ‡โ‚‚)) + ฮ”โˆ‡โ‚ƒ = spzeros(Float64, size(โˆ‡โ‚ƒ)...) + ฮ”๐’โ‚ = zeros(Float64, size(๐’โ‚)) + ฮ”๐’โ‚‚ = spzeros(Float64, size(๐’โ‚‚)...) + ฮ”๐’โ‚ƒฬ‚ = spzeros(Float64, size(๐’โ‚ƒฬ‚)...) + if !(ฮ” isa Union{NoTangent, AbstractZero}) && hasmethod(getindex, Tuple{typeof(ฮ”), Int}) + v1 = ฮ”[1] + v3 = ฮ”[3] + v5 = ฮ”[5] + v6 = ฮ”[6] + v7 = ฮ”[7] + v8 = ฮ”[8] + v9 = ฮ”[9] + v10 = ฮ”[10] + ฮ”sss = v1 isa Union{NoTangent, AbstractZero} ? ฮ”sss : v1 + ฮ”SS_and_pars = v3 isa Union{NoTangent, AbstractZero} ? ฮ”SS_and_pars : v3 + ฮ”โˆ‡โ‚ = v5 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚ : v5 + ฮ”โˆ‡โ‚‚ = v6 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚‚ : v6 + ฮ”โˆ‡โ‚ƒ = v7 isa Union{NoTangent, AbstractZero} ? ฮ”โˆ‡โ‚ƒ : v7 + ฮ”๐’โ‚ = v8 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚ : v8 + ฮ”๐’โ‚‚ = v9 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚‚ : v9 + ฮ”๐’โ‚ƒฬ‚ = v10 isa Union{NoTangent, AbstractZero} ? ฮ”๐’โ‚ƒฬ‚ : v10 + end + + โˆ‚state_vec = ฮ”sss + โˆ‚๐’โ‚_from_state = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚_from_state[:, 1:nPast] += โˆ‚state_vec * SSSstates' + โˆ‚๐’โ‚‚_from_state = spzeros(Float64, size(๐’โ‚‚)...) + โˆ‚๐’โ‚‚_from_state += โˆ‚state_vec * kron_aug1' / 2 + โˆ‚SSSstates = ๐’โ‚[:,1:nPast]' * โˆ‚state_vec + + โˆ‚๐’โ‚ƒ_raw = Matrix(ฮ”๐’โ‚ƒฬ‚) * ๐”โ‚ƒ' + so3_tangents = third_order_solution_pullback((โˆ‚๐’โ‚ƒ_raw, NoTangent())) + โˆ‚โˆ‡โ‚_from_so3 = so3_tangents[2] isa Union{NoTangent, AbstractZero} ? zero(โˆ‡โ‚) : so3_tangents[2] + โˆ‚โˆ‡โ‚‚_from_so3 = so3_tangents[3] isa Union{NoTangent, AbstractZero} ? zero(โˆ‡โ‚‚) : so3_tangents[3] + โˆ‚โˆ‡โ‚ƒ_from_so3 = so3_tangents[4] isa Union{NoTangent, AbstractZero} ? zero(โˆ‡โ‚ƒ) : so3_tangents[4] + โˆ‚๐’โ‚_raw_from_so3 = so3_tangents[5] isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚_raw) : so3_tangents[5] + โˆ‚๐’โ‚‚_raw_from_so3 = so3_tangents[6] isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚‚_raw) : so3_tangents[6] + + โˆ‚๐’โ‚_from_so3 = zeros(Float64, size(๐’โ‚)) + โˆ‚๐’โ‚_from_so3[:, 1:nPast] = โˆ‚๐’โ‚_raw_from_so3[:, 1:nPast] + โˆ‚๐’โ‚_from_so3[:, nPast+2:end] = โˆ‚๐’โ‚_raw_from_so3[:, nPast+1:end] + + โˆ‚โˆ‡โ‚ƒ_total = ฮ”โˆ‡โ‚ƒ + โˆ‚โˆ‡โ‚ƒ_from_so3 + third_derivatives_tangents = third_derivatives_pullback(โˆ‚โˆ‡โ‚ƒ_total) + โˆ‚params_from_โˆ‡โ‚ƒ = third_derivatives_tangents[2] + โˆ‚SS_and_pars_from_โˆ‡โ‚ƒ = third_derivatives_tangents[3] + + # Convert full-space โˆ‚๐’โ‚‚ terms to compressed, then accumulate with compressed โˆ‚๐’โ‚‚_raw_from_so3 + โˆ‚๐’โ‚‚_raw_for_common = โˆ‚๐’โ‚‚_raw_from_so3 + (โˆ‚๐’โ‚‚_from_state + ฮ”๐’โ‚‚) * ๐”โ‚‚' + + common_tangents = common_pullback((NoTangent(), + ฮ”sss, + ฮ”SS_and_pars + โˆ‚SS_and_pars_from_โˆ‡โ‚ƒ, + NoTangent(), + ฮ”โˆ‡โ‚ + โˆ‚โˆ‡โ‚_from_so3, + ฮ”โˆ‡โ‚‚ + โˆ‚โˆ‡โ‚‚_from_so3, + โˆ‚๐’โ‚_from_state + ฮ”๐’โ‚ + โˆ‚๐’โ‚_from_so3, + โˆ‚๐’โ‚‚_raw_for_common, + โˆ‚SSSstates, + NoTangent())) + + โˆ‚parameters = common_tangents[2] + โˆ‚params_from_โˆ‡โ‚ƒ + return NoTangent(), NoTangent(), โˆ‚parameters, NoTangent() + end + + return result, pullback +end + + +function rrule(::typeof(get_relevant_steady_state_and_state_update), + ::Val{:second_order}, + parameter_values::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) where S <: AbstractFloat + # Call inner rrule in the forward pass to capture pullback (avoids re-computing in backward) + ss_rrule = rrule(calculate_stochastic_steady_state, + Val(:second_order), + parameter_values, + ๐“‚; + opts = opts, + estimation = estimation) + + if ss_rrule === nothing + y = get_relevant_steady_state_and_state_update(Val(:second_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ss_out, ss_pb = ss_rrule + sss = ss_out[1] + converged = ss_out[2] + SS_and_pars = ss_out[3] + solution_error = ss_out[4] + ๐’โ‚ = ss_out[7] + ๐’โ‚‚ = ss_out[8] + + if !converged || solution_error > opts.tol.nsss.acceptance_tol + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚], collect(sss), converged) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + all_SS = expand_steady_state(SS_and_pars, ms) + state = collect(sss) - all_SS + + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚], state, converged) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + ฮ”๐’ = ฮ”y[3] + ฮ”state = ฮ”y[4] + + # Guard against NoTangent cotangents from filter failure + ฮ”state_val = ฮ”state isa Union{NoTangent, AbstractZero} ? zeros(S, length(state)) : ฮ”state + ฮ”๐’โ‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚)) : ฮ”๐’[1] + ฮ”๐’โ‚‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚‚)) : ฮ”๐’[2] + + ฮ”sss = ฮ”state_val + E = ms.steady_state_expand_matrix + ฮ”SS_and_pars = ฮ”SS_and_pars - E' * ฮ”state_val + + ss_grads = ss_pb((ฮ”sss, + NoTangent(), + ฮ”SS_and_pars, + NoTangent(), + NoTangent(), + NoTangent(), + ฮ”๐’โ‚, + ฮ”๐’โ‚‚)) + + return NoTangent(), NoTangent(), ss_grads[3], NoTangent() + end + + return y, pullback +end + +function rrule(::typeof(get_relevant_steady_state_and_state_update), + ::Val{:pruned_second_order}, + parameter_values::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) where S <: AbstractFloat + # Call inner rrule in the forward pass to capture pullback (avoids re-computing in backward) + ss_rrule = rrule(calculate_stochastic_steady_state, + Val(:pruned_second_order), + parameter_values, + ๐“‚; + opts = opts, + estimation = estimation) + + if ss_rrule === nothing + y = get_relevant_steady_state_and_state_update(Val(:pruned_second_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ss_out, ss_pb = ss_rrule + sss = ss_out[1] + converged = ss_out[2] + SS_and_pars = ss_out[3] + solution_error = ss_out[4] + ๐’โ‚ = ss_out[7] + ๐’โ‚‚ = ss_out[8] + nVars = ๐“‚.constants.post_model_macro.nVars + + if !converged || solution_error > opts.tol.nsss.acceptance_tol + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚], [zeros(S, nVars), zeros(S, nVars)], converged) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + all_SS = expand_steady_state(SS_and_pars, ms) + state = [zeros(S, nVars), collect(sss) - all_SS] + + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚], state, converged) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + ฮ”๐’ = ฮ”y[3] + ฮ”state = ฮ”y[4] + + E = ms.steady_state_expand_matrix + # Guard against NoTangent cotangents from filter failure + ฮ”state_val = ฮ”state isa Union{NoTangent, AbstractZero} ? [zeros(S, nVars), zeros(S, nVars)] : ฮ”state + ฮ”๐’โ‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚)) : ฮ”๐’[1] + ฮ”๐’โ‚‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚‚)) : ฮ”๐’[2] + + ฮ”sss = ฮ”state_val[2] + ฮ”SS_and_pars = ฮ”SS_and_pars - E' * ฮ”state_val[2] + + ss_grads = ss_pb((ฮ”sss, + NoTangent(), + ฮ”SS_and_pars, + NoTangent(), + NoTangent(), + NoTangent(), + ฮ”๐’โ‚, + ฮ”๐’โ‚‚)) + + return NoTangent(), NoTangent(), ss_grads[3], NoTangent() + end + + return y, pullback +end + +function rrule(::typeof(get_relevant_steady_state_and_state_update), + ::Val{:third_order}, + parameter_values::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) where S <: AbstractFloat + # Call inner rrule in the forward pass to capture pullback (avoids re-computing in backward) + ss_rrule = rrule(calculate_stochastic_steady_state, + Val(:third_order), + parameter_values, + ๐“‚; + opts = opts, + estimation = estimation) + + if ss_rrule === nothing + y = get_relevant_steady_state_and_state_update(Val(:third_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ss_out, ss_pb = ss_rrule + sss = ss_out[1] + converged = ss_out[2] + SS_and_pars = ss_out[3] + solution_error = ss_out[4] + ๐’โ‚ = ss_out[8] + ๐’โ‚‚ = ss_out[9] + ๐’โ‚ƒ = ss_out[10] + + if !converged || solution_error > opts.tol.nsss.acceptance_tol + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ], collect(sss), converged) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + all_SS = expand_steady_state(SS_and_pars, ms) + state = collect(sss) - all_SS + + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ], state, converged) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + ฮ”๐’ = ฮ”y[3] + ฮ”state = ฮ”y[4] + ฮ”SS_and_pars = ฮ”SS_and_pars isa Union{NoTangent, AbstractZero} ? zero(SS_and_pars) : ฮ”SS_and_pars + + # Guard against NoTangent cotangents from filter failure + ฮ”state_val = ฮ”state isa Union{NoTangent, AbstractZero} ? zeros(S, length(state)) : ฮ”state + ฮ”๐’โ‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚) : ฮ”๐’[1] + ฮ”๐’โ‚‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚‚) : ฮ”๐’[2] + ฮ”๐’โ‚ƒ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zero(๐’โ‚ƒ) : ฮ”๐’[3] + + ฮ”sss = ฮ”state_val + E = ms.steady_state_expand_matrix + ฮ”SS_and_pars = ฮ”SS_and_pars - E' * ฮ”state_val + + ss_grads = ss_pb((ฮ”sss, + NoTangent(), + ฮ”SS_and_pars, + NoTangent(), + NoTangent(), + NoTangent(), + NoTangent(), + ฮ”๐’โ‚, + ฮ”๐’โ‚‚, + ฮ”๐’โ‚ƒ)) + + return NoTangent(), NoTangent(), ss_grads[3], NoTangent() + end + return y, pullback +end + +function rrule(::typeof(get_relevant_steady_state_and_state_update), + ::Val{:pruned_third_order}, + parameter_values::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false) where S <: AbstractFloat + # Call inner rrule in the forward pass to capture pullback (avoids re-computing in backward) + ss_rrule = rrule(calculate_stochastic_steady_state, + Val(:pruned_third_order), + parameter_values, + ๐“‚; + opts = opts, + estimation = estimation) + + if ss_rrule === nothing + y = get_relevant_steady_state_and_state_update(Val(:pruned_third_order), parameter_values, ๐“‚, opts = opts, estimation = estimation) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ss_out, ss_pb = ss_rrule + sss = ss_out[1] + converged = ss_out[2] + SS_and_pars = ss_out[3] + solution_error = ss_out[4] + ๐’โ‚ = ss_out[8] + ๐’โ‚‚ = ss_out[9] + ๐’โ‚ƒ = ss_out[10] + nVars = ๐“‚.constants.post_model_macro.nVars + + if !converged || solution_error > opts.tol.nsss.acceptance_tol + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ], [zeros(S, nVars), zeros(S, nVars), zeros(S, nVars)], converged) + return y, _ -> (NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent()) + end + + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + all_SS = expand_steady_state(SS_and_pars, ms) + state = [zeros(S, nVars), collect(sss) - all_SS, zeros(S, nVars)] + + y = (๐“‚.constants, SS_and_pars, [๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ], state, converged) + + pullback = function (ศณ) + ฮ”y = unthunk(ศณ) + if ฮ”y isa NoTangent || ฮ”y isa AbstractZero + return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() + end + + ฮ”SS_and_pars = ฮ”y[2] + ฮ”๐’ = ฮ”y[3] + ฮ”state = ฮ”y[4] + + E = ms.steady_state_expand_matrix + # Guard against NoTangent cotangents from filter failure + ฮ”state_val = ฮ”state isa Union{NoTangent, AbstractZero} ? [zeros(S, nVars), zeros(S, nVars), zeros(S, nVars)] : ฮ”state + ฮ”๐’โ‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚)) : ฮ”๐’[1] + ฮ”๐’โ‚‚ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚‚)) : ฮ”๐’[2] + ฮ”๐’โ‚ƒ = ฮ”๐’ isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚ƒ)) : ฮ”๐’[3] + + ฮ”sss = ฮ”state_val[2] + ฮ”SS_and_pars = ฮ”SS_and_pars - E' * ฮ”state_val[2] + + ss_grads = ss_pb((ฮ”sss, + NoTangent(), + ฮ”SS_and_pars, + NoTangent(), + NoTangent(), + NoTangent(), + NoTangent(), + ฮ”๐’โ‚, + ฮ”๐’โ‚‚, + ฮ”๐’โ‚ƒ)) + + return NoTangent(), NoTangent(), ss_grads[3], NoTangent() + end + + return y, pullback +end + +function rrule(::typeof(get_loglikelihood), + ๐“‚::โ„ณ, + data::KeyedArray{Float64}, + parameter_values::Vector{S}; + steady_state_function::SteadyStateFunctionType = missing, + algorithm::Symbol = DEFAULT_ALGORITHM, + filter::Symbol = DEFAULT_FILTER_SELECTOR(algorithm), + on_failure_loglikelihood::U = -Inf, + warmup_iterations::Int = DEFAULT_WARMUP_ITERATIONS, + presample_periods::Int = DEFAULT_PRESAMPLE_PERIODS, + initial_covariance::Symbol = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + tol::Tolerances = Tolerances(), + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(๐“‚), + verbose::Bool = DEFAULT_VERBOSE) where {S <: Real, U <: AbstractFloat} + + opts = merge_calculation_options(tol = tol, verbose = verbose, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + sylvester_algorithmยฒ = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], + sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2], + lyapunov_algorithm = lyapunov_algorithm) + + estimation = true + + filter, _, algorithm, _, _, warmup_iterations = normalize_filtering_options(filter, false, algorithm, false, warmup_iterations) + + observables = get_and_check_observables(๐“‚.constants.post_model_macro, data) + + solve!(๐“‚, opts = opts, steady_state_function = steady_state_function, algorithm = algorithm) + + bounds_violated = check_bounds(parameter_values, ๐“‚) + + if bounds_violated + llh = S(on_failure_loglikelihood) + return llh, _ -> (NoTangent(), NoTangent(), NoTangent(), zeros(S, length(parameter_values))) + end + + obs_indices = convert(Vector{Int}, indexin(observables, ๐“‚.constants.post_complete_parameters.SS_and_pars_names)) + + # โ”€โ”€ step 1: get_relevant_steady_state_and_state_update โ”€โ”€ + ss_rrule = rrule(get_relevant_steady_state_and_state_update, + Val(algorithm), parameter_values, ๐“‚; + opts = opts, estimation = estimation) + + if ss_rrule === nothing + # fall back to primal-only when no rrule is available + constants_obj, SS_and_pars, ๐’, state, solved = get_relevant_steady_state_and_state_update( + Val(algorithm), parameter_values, ๐“‚, opts = opts, estimation = estimation) + ss_pb = nothing + else + (constants_obj, SS_and_pars, ๐’, state, solved), ss_pb = ss_rrule + end + + if !solved + llh = S(on_failure_loglikelihood) + return llh, _ -> (NoTangent(), NoTangent(), NoTangent(), zeros(S, length(parameter_values))) + end + + # โ”€โ”€ step 2: data_in_deviations = dt .- SS_and_pars[obs_indices] โ”€โ”€ + dt = if collect(axiskeys(data, 1)) isa Vector{String} + collect(rekey(data, 1 => axiskeys(data, 1) .|> Meta.parse .|> replace_indices)(observables)) + else + collect(data(observables)) + end + + data_in_deviations = dt .- SS_and_pars[obs_indices] + + # โ”€โ”€ step 3: calculate_loglikelihood โ”€โ”€ + llh_rrule = rrule(calculate_loglikelihood, + Val(filter), Val(algorithm), obs_indices, + ๐’, data_in_deviations, constants_obj, state, ๐“‚.workspaces; + warmup_iterations = warmup_iterations, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + filter_algorithm = filter_algorithm, + opts = opts, + on_failure_loglikelihood = on_failure_loglikelihood) + + if llh_rrule === nothing + llh = calculate_loglikelihood(Val(filter), Val(algorithm), obs_indices, + ๐’, data_in_deviations, constants_obj, state, ๐“‚.workspaces; + warmup_iterations = warmup_iterations, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + filter_algorithm = filter_algorithm, + opts = opts, + on_failure_loglikelihood = on_failure_loglikelihood) + + return llh, _ -> (NoTangent(), NoTangent(), NoTangent(), zeros(S, length(parameter_values))) + end + + llh, llh_pb = llh_rrule + + # โ”€โ”€ pullback โ”€โ”€ + pullback = function (โˆ‚llh_bar) + โˆ‚llh = unthunk(โˆ‚llh_bar) + + # backprop through calculate_loglikelihood + # returns: (_, _, _, _, โˆ‚๐’, โˆ‚data_in_deviations, _, โˆ‚state, _) + llh_grads = llh_pb(โˆ‚llh) + โˆ‚๐’ = llh_grads[5] + โˆ‚data_in_devs = llh_grads[6] + โˆ‚state = llh_grads[8] + + # When the filter forward pass fails (non-finite states, factorisation + # failure, etc.) the filter rrule returns on_failure_loglikelihood with + # an all-NoTangent pullback. The loglikelihood is then a constant, so + # the parameter gradient is exactly zero. + if โˆ‚๐’ isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), NoTangent(), zeros(S, length(parameter_values)) + end + + # backprop through data_in_deviations = dt .- SS_and_pars[obs_indices] + โˆ‚SS_and_pars = zeros(S, length(SS_and_pars)) + if !(โˆ‚data_in_devs isa Union{NoTangent, AbstractZero}) + โˆ‚SS_and_pars[obs_indices] .-= vec(sum(โˆ‚data_in_devs, dims = 2)) + end + + if ss_pb === nothing + return NoTangent(), NoTangent(), NoTangent(), zeros(S, length(parameter_values)) + end + + # backprop through get_relevant_steady_state_and_state_update + # cotangent: (ฮ”constants, ฮ”SS_and_pars, ฮ”๐’, ฮ”state, ฮ”solved) + ss_grads = ss_pb((NoTangent(), โˆ‚SS_and_pars, โˆ‚๐’, โˆ‚state, NoTangent())) + โˆ‚parameter_values = ss_grads[3] + + return NoTangent(), NoTangent(), NoTangent(), โˆ‚parameter_values + end + + return llh, pullback +end + +function rrule(::typeof(get_irf), + ๐“‚::โ„ณ, + parameters::Vector{S}; + steady_state_function::SteadyStateFunctionType = missing, + periods::Int = DEFAULT_PERIODS, + variables::Union{Symbol_input,String_input} = DEFAULT_VARIABLES_EXCLUDING_OBC, + shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = DEFAULT_SHOCK_SELECTION, + negative_shock::Bool = DEFAULT_NEGATIVE_SHOCK, + initial_state::Vector{Float64} = DEFAULT_INITIAL_STATE, + levels::Bool = false, + verbose::Bool = DEFAULT_VERBOSE, + tol::Tolerances = Tolerances(), + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM) where S <: Real + + opts = merge_calculation_options(tol = tol, verbose = verbose, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm) + + estimation = true + + constants_obj = initialise_constants!(๐“‚) + + solve!(๐“‚, + steady_state_function = steady_state_function, + opts = opts) + + shocks = ๐“‚.constants.post_model_macro.nExo == 0 ? :none : shocks + + shocks, negative_shock, _, periods, shock_idx, shock_history = process_shocks_input(shocks, negative_shock, 1.0, periods, ๐“‚) + + var_idx = parse_variables_input_to_index(variables, ๐“‚) |> sort + + nVars = ๐“‚.constants.post_model_macro.nVars + nExo = ๐“‚.constants.post_model_macro.nExo + past_idx = ๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx + nPast = length(past_idx) + nShocks = shocks == :none ? 1 : length(shock_idx) + + zero_result() = zeros(S, length(var_idx), periods, nShocks) + zero_pullback(_) = (NoTangent(), NoTangent(), zeros(S, length(parameters))) + + # โ”€โ”€ step 1: NSSS โ”€โ”€ + nsss_out, nsss_pb = rrule(get_NSSS_and_parameters, + ๐“‚, + parameters; + opts = opts, + estimation = estimation) + + reference_steady_state = nsss_out[1] + solution_error = nsss_out[2][1] + + if (solution_error > tol.nsss.acceptance_tol) || isnan(solution_error) + return zero_result(), zero_pullback + end + + # โ”€โ”€ step 2: Jacobian โ”€โ”€ + โˆ‡โ‚, jac_pb = rrule(calculate_jacobian, + parameters, + reference_steady_state, + ๐“‚.caches, + ๐“‚.functions.jacobian, + ๐“‚.workspaces) + + # โ”€โ”€ step 3: First-order solution โ”€โ”€ + first_out, first_pb = rrule(calculate_first_order_solution, + โˆ‡โ‚, + constants_obj, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameters) + + sol_mat = first_out[1] + solved = first_out[3] + + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + + if !solved + return zero_result(), zero_pullback + end + + # โ”€โ”€ step 4: Forward simulation (mutation-free, storing inputs for pullback) โ”€โ”€ + init_state = initial_state == [0.0] ? zeros(S, nVars) : initial_state - reference_steady_state[1:length(๐“‚.constants.post_model_macro.var)] + + # Pre-allocate output and input storage + Y_all = zeros(S, nVars, periods, nShocks) + # Store the input vectors [state[past_idx]; shock] for each (shock_i, t) โ€” needed for pullback + inputs_all = Array{Vector{S}}(undef, nShocks, periods) + + for (si, ii) in enumerate(shock_idx) + # Build shock history for this shock index + if shocks isa Union{Symbol_input,String_input} + shock_hist = zeros(nExo, periods) + if shocks โ‰  :none + shock_hist[ii, 1] = negative_shock ? -1.0 : 1.0 + end + else + shock_hist = shock_history + end + + # t = 1 + prev_state = init_state + input_vec = vcat(prev_state[past_idx], shock_hist[:, 1]) + y_t = sol_mat * input_vec + inputs_all[si, 1] = input_vec + Y_all[:, 1, si] = y_t + + # t = 2:periods + for t in 2:periods + input_vec = vcat(y_t[past_idx], shock_hist[:, t]) + y_t = sol_mat * input_vec + inputs_all[si, t] = input_vec + Y_all[:, t, si] = y_t + end + end + + # โ”€โ”€ step 5: Assemble output โ”€โ”€ + deviations = Y_all[var_idx, :, :] + + result = if levels + deviations .+ reference_steady_state[var_idx] + else + deviations + end + + # โ”€โ”€ step 6: Pullback โ”€โ”€ + pullback = function (โˆ‚result_bar) + โˆ‚result = unthunk(โˆ‚result_bar) + + if โˆ‚result isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, length(parameters)) + end + + # Scatter var_idx back to full nVars dimension + โˆ‚Y_all = zeros(S, nVars, periods, nShocks) + โˆ‚Y_all[var_idx, :, :] .= โˆ‚result + + # SS gradient from levels mode + โˆ‚SS_and_pars = zeros(S, length(reference_steady_state)) + if levels + โˆ‚SS_and_pars[var_idx] .+= dropdims(sum(โˆ‚result, dims = (2, 3)), dims = (2, 3)) + end + + # BPTT through the linear simulation to get โˆ‚sol_mat + โˆ‚sol_mat = zeros(S, size(sol_mat)) + + for si in 1:nShocks + # Accumulated gradient flowing backward through states + โˆ‚y_accum = zeros(S, nVars) + + for t in periods:-1:1 + # Total gradient at time t = direct gradient + propagated from t+1 + โˆ‚y_t = โˆ‚Y_all[:, t, si] .+ โˆ‚y_accum + + # โˆ‚sol_mat += โˆ‚y_t * input_t' + input_t = inputs_all[si, t] + โˆ‚sol_mat .+= โˆ‚y_t * input_t' + + # Propagate gradient to previous state through sol_mat + # input_t = [y_{t-1}[past_idx]; shock_t] + # โˆ‚input_t = sol_mat' * โˆ‚y_t + โˆ‚input_t = sol_mat' * โˆ‚y_t + + # Only the first nPast entries of โˆ‚input_t flow to โˆ‚y_{t-1}[past_idx] + โˆ‚y_accum = zeros(S, nVars) + โˆ‚y_accum[past_idx] .+= โˆ‚input_t[1:nPast] + end + + # After BPTT for this shock, โˆ‚y_accum is the gradient w.r.t. init_state. + # When init_state = initial_state - reference_steady_state[1:nVar], + # propagate gradient to reference_steady_state with negative sign. + if initial_state != [0.0] + nVar_len = length(๐“‚.constants.post_model_macro.var) + โˆ‚SS_and_pars[1:nVar_len] .-= โˆ‚y_accum[1:nVar_len] + end + end + + # โ”€โ”€ Chain backward through sub-pullbacks โ”€โ”€ + # first_pb expects cotangent tuple: (โˆ‚sol_mat, โˆ‚qme_sol, โˆ‚solved) + first_grads = first_pb((โˆ‚sol_mat, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚ = first_grads[2] + + jac_grads = jac_pb(โˆ‚โˆ‡โ‚) + โˆ‚parameters_from_jac = jac_grads[2] + โˆ‚SS_from_jac = jac_grads[3] + + โˆ‚SS_and_pars .+= โˆ‚SS_from_jac + + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + โˆ‚parameters_from_nsss = nsss_grads[3] + + โˆ‚parameters_total = โˆ‚parameters_from_jac .+ โˆ‚parameters_from_nsss + + return NoTangent(), NoTangent(), โˆ‚parameters_total + end + + return result, pullback +end + +# โ”€โ”€ calculate_covariance rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +function rrule(::typeof(calculate_covariance), + parameters::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options()) where S <: Real + + # โ”€โ”€ Non-differentiable setup โ”€โ”€ + constants_obj = initialise_constants!(๐“‚) + idx_constants = constants_obj.post_complete_parameters + T = constants_obj.post_model_macro + nPast = T.nPast_not_future_and_mixed + past_idx = T.past_not_future_and_mixed_idx + P = idx_constants.diag_nVars[past_idx, :] # (nPast, nVars) constant selection matrix + + zero_result() = (zeros(S, 0, 0), zeros(S, 0, 0), zeros(S, 0, 0), zeros(S, 0), false) + zero_pb(_) = (NoTangent(), zeros(S, length(parameters)), NoTangent()) + + # โ”€โ”€ Step 1: NSSS โ”€โ”€ + nsss_out, nsss_pb = rrule(get_NSSS_and_parameters, ๐“‚, parameters; opts = opts) + SS_and_pars = nsss_out[1] + solution_error = nsss_out[2][1] + + if solution_error > opts.tol.nsss.acceptance_tol + return (zeros(S, 0, 0), zeros(S, 0, 0), zeros(S, 0, 0), SS_and_pars, false), zero_pb + end + + # โ”€โ”€ Step 2: Jacobian โ”€โ”€ + โˆ‡โ‚, jac_pb = rrule(calculate_jacobian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) + + # โ”€โ”€ Step 3: First-order solution โ”€โ”€ + first_out, first_pb = rrule(calculate_first_order_solution, + โˆ‡โ‚, + constants_obj, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.qme_solution, + opts = opts, + parameter_values = parameters) + sol = first_out[1] + solved_first = first_out[3] + + update_perturbation_counter!(๐“‚.counters, solved_first, order = 1) + + # โ”€โ”€ Step 4: A, C, CC (mutation-free) โ”€โ”€ + A = sol[:, 1:nPast] * P + C = sol[:, nPast+1:end] + CC = C * C' + + if !solved_first + return (CC, sol, โˆ‡โ‚, SS_and_pars, solved_first), zero_pb + end + + # โ”€โ”€ Step 5: Lyapunov โ”€โ”€ + lyap_ws = ensure_lyapunov_workspace!(๐“‚.workspaces, T.nVars, :first_order) + + lyap_out, lyap_pb = rrule(solve_lyapunov_equation, A, CC, lyap_ws; + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.first_order.ad.lyapunov, + verbose = opts.verbose) + covar_raw = lyap_out[1] + solved_lyap = lyap_out[2] + + solved = solved_first && solved_lyap + + result = (covar_raw, sol, โˆ‡โ‚, SS_and_pars, solved) + + # โ”€โ”€ Pullback โ”€โ”€ + function calculate_covariance_pullback(ฮ”out) + ฮ”covar, ฮ”sol_ret, ฮ”โˆ‡โ‚_ret, ฮ”SS_ret, _ = ฮ”out + + # Materialise any InplaceableThunk / Thunk wrappers + ฮ”covar = unthunk(ฮ”covar) + ฮ”sol_ret = unthunk(ฮ”sol_ret) + ฮ”โˆ‡โ‚_ret = unthunk(ฮ”โˆ‡โ‚_ret) + ฮ”SS_ret = unthunk(ฮ”SS_ret) + + # Accumulators + โˆ‚sol_total = zeros(S, size(sol)) + โˆ‚โˆ‡โ‚_total = zeros(S, size(โˆ‡โ‚)) + โˆ‚SS_total = zeros(S, length(SS_and_pars)) + + # Direct cotangents passed through the tuple + if !(ฮ”sol_ret isa AbstractZero) + โˆ‚sol_total .+= ฮ”sol_ret + end + if !(ฮ”โˆ‡โ‚_ret isa AbstractZero) + โˆ‚โˆ‡โ‚_total .+= ฮ”โˆ‡โ‚_ret + end + if !(ฮ”SS_ret isa AbstractZero) + โˆ‚SS_total .+= ฮ”SS_ret + end + + # Backprop through Lyapunov equation + if !(ฮ”covar isa AbstractZero) + lyap_grad = lyap_pb((ฮ”covar, NoTangent())) + ฮ”A = lyap_grad[2] # โˆ‚A + ฮ”CC = lyap_grad[3] # โˆ‚CC + + # CC = C * C' โ†’ โˆ‚C = (โˆ‚CC + โˆ‚CC') * C + ฮ”C = (ฮ”CC + ฮ”CC') * C + + # A = sol[:, 1:nPast] * P โ†’ โˆ‚sol[:, 1:nPast] += โˆ‚A * P' + โˆ‚sol_total[:, 1:nPast] .+= ฮ”A * P' + + # C = sol[:, nPast+1:end] + โˆ‚sol_total[:, nPast+1:end] .+= ฮ”C + end + + # Backprop through first-order solution + first_grad = first_pb((โˆ‚sol_total, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚_total .+= first_grad[2] + + # Backprop through Jacobian + jac_grad = jac_pb(โˆ‚โˆ‡โ‚_total) + โˆ‚parameters_from_jac = jac_grad[2] + โˆ‚SS_from_jac = jac_grad[3] + โˆ‚SS_total .+= โˆ‚SS_from_jac + + # Backprop through NSSS + nsss_grad = nsss_pb((โˆ‚SS_total, NoTangent())) + โˆ‚parameters_from_nsss = nsss_grad[3] + + โˆ‚parameters_total = โˆ‚parameters_from_jac .+ โˆ‚parameters_from_nsss + + return NoTangent(), โˆ‚parameters_total, NoTangent() + end + + return result, calculate_covariance_pullback +end + + +# โ”€โ”€ Helper: VJP of kron(A, B) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# Given C = kron(A, B) and cotangent โˆ‚C, returns (โˆ‚A, โˆ‚B). +function _kron_vjp(โˆ‚C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix) + m, n = size(A) + p, q = size(B) + S = eltype(โˆ‚C) + โˆ‚A = zeros(S, m, n) + โˆ‚B = zeros(S, p, q) + @inbounds for j in 1:n + for i in 1:m + blk = @view โˆ‚C[(i-1)*p+1:i*p, (j-1)*q+1:j*q] + โˆ‚A[i,j] = โ„’.dot(blk, B) + if !iszero(A[i,j]) + โˆ‚B .+= A[i,j] .* blk + end + end + end + return โˆ‚A, โˆ‚B +end + + +# โ”€โ”€ calculate_mean rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +function rrule(::typeof(calculate_mean), + parameters::Vector{S}, + ๐“‚::โ„ณ; + algorithm = :pruned_second_order, + opts::CalculationOptions = merge_calculation_options()) where S <: Real + + @assert algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] "Theoretical mean available only for first order, pruned second and pruned third order perturbation solutions." + + # โ”€โ”€ Non-differentiable setup โ”€โ”€ + constants_obj = initialise_constants!(๐“‚) + T_pm = constants_obj.post_model_macro + nVars = T_pm.nVars + np = length(parameters) + + zero_pb(_) = (NoTangent(), zeros(S, np), NoTangent()) + + # โ”€โ”€ Step 1: NSSS โ”€โ”€ + nsss_out, nsss_pb = rrule(get_NSSS_and_parameters, ๐“‚, parameters; opts = opts) + SS_and_pars = nsss_out[1] + solution_error = nsss_out[2][1] + + # โ”€โ”€ First-order path (mean = steady state) โ”€โ”€ + if algorithm == :first_order + solved = solution_error < opts.tol.nsss.acceptance_tol + mean_of_variables = SS_and_pars[1:nVars] + + function first_order_mean_pullback(โˆ‚out) + โˆ‚mean = unthunk(โˆ‚out[1]) + if โˆ‚mean isa AbstractZero + return NoTangent(), zeros(S, np), NoTangent() + end + โˆ‚SS = zeros(S, length(SS_and_pars)) + โˆ‚SS[1:nVars] .= โˆ‚mean + nsss_grad = nsss_pb((โˆ‚SS, NoTangent())) + โˆ‚params = nsss_grad[3] isa AbstractZero ? zeros(S, np) : nsss_grad[3] + return NoTangent(), โˆ‚params, NoTangent() + end + + return (mean_of_variables, solved), first_order_mean_pullback + end + + # โ”€โ”€ Higher-order path: early exit on NSSS failure โ”€โ”€ + if solution_error > opts.tol.nsss.acceptance_tol + return (SS_and_pars[1:nVars], false), zero_pb + end + + ensure_moments_constants!(constants_obj) + so = constants_obj.second_order + + nแต‰ = T_pm.nExo + nหข = T_pm.nPast_not_future_and_mixed + iหข = T_pm.past_not_future_and_mixed_idx + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + vec_Iโ‚‘ = so.vec_Iโ‚‘ + + # โ”€โ”€ Step 2: Jacobian โ”€โ”€ + โˆ‡โ‚, jac_pb = rrule(calculate_jacobian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) + + # โ”€โ”€ Step 3: First-order solution โ”€โ”€ + first_out, first_pb = rrule(calculate_first_order_solution, + โˆ‡โ‚, + constants_obj, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.qme_solution, + opts = opts, + parameter_values = parameters) + ๐’โ‚ = first_out[1] + solved_first = first_out[3] + + update_perturbation_counter!(๐“‚.counters, solved_first, order = 1) + + if !solved_first + return (SS_and_pars[1:nVars], false), zero_pb + end + + # โ”€โ”€ Step 4: Hessian โ”€โ”€ + โˆ‡โ‚‚, hess_pb = rrule(calculate_hessian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces) + + # โ”€โ”€ Step 5: Second-order solution โ”€โ”€ + so2_out, so2_pb = rrule(calculate_second_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; opts = opts, parameter_values = parameters) + ๐’โ‚‚_raw = so2_out[1] + solved2 = so2_out[2] + + update_perturbation_counter!(๐“‚.counters, solved2, order = 2) + + if !solved2 + return (SS_and_pars[1:nVars], false), zero_pb + end + + # โ”€โ”€ Step 6: Decompress Sโ‚‚ โ”€โ”€ + ๐’โ‚‚_full = ๐’โ‚‚_raw * ๐”โ‚‚ + + # โ”€โ”€ Step 7: Slicing and mean computation โ”€โ”€ + kron_s_s = so.kron_states + kron_e_e = so.kron_e_e + kron_v_v = so.kron_v_v + + # First-order slices + s_to_yโ‚ = ๐’โ‚[:, 1:nหข] + s_to_sโ‚ = ๐’โ‚[iหข, 1:nหข] + e_to_sโ‚ = ๐’โ‚[iหข, (nหข+1):end] + + # Second-order slices (dense) + s_s_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_s_s]) + e_e_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_e_e]) + v_v_to_yโ‚‚_v = vec(๐’โ‚‚_full[:, kron_v_v]) + s_s_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_s_s]) + e_e_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_e_e]) + v_v_to_sโ‚‚_v = vec(๐’โ‚‚_full[iหข, kron_v_v]) + + # Kronecker products + sโ‚_kron_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + eโ‚_kron_eโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) |> collect + + # Block transition matrix + ล_to_ลโ‚‚ = [ s_to_sโ‚ zeros(S, nหข, nหข + nหข^2) + zeros(S, nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 + zeros(S, nหข^2, 2*nหข) sโ‚_kron_sโ‚ ] + + ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2] + + ลvโ‚‚ = vcat(zeros(S, nหข), + v_v_to_sโ‚‚_v / 2 + e_e_to_sโ‚‚ * vec_Iโ‚‘ / 2, + eโ‚_kron_eโ‚ * vec_Iโ‚‘) + + yvโ‚‚ = (v_v_to_yโ‚‚_v + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 + + # Mean solve + A_mean = collect(โ„’.I(size(ล_to_ลโ‚‚, 1))) - ล_to_ลโ‚‚ + ฮผหขโบโ‚‚ = A_mean \ ลvโ‚‚ + + mean_of_variables = SS_and_pars[1:nVars] + ล_to_yโ‚‚ * ฮผหขโบโ‚‚ + yvโ‚‚ + + slvd = solved_first && solved2 + + result = (mean_of_variables, slvd) + + # โ”€โ”€ Pullback โ”€โ”€ + function calculate_mean_pullback(โˆ‚out) + โˆ‚mean_in = unthunk(โˆ‚out[1]) + + if โˆ‚mean_in isa AbstractZero + return NoTangent(), zeros(S, np), NoTangent() + end + + # Accumulators + โˆ‚๐’โ‚_acc = zeros(S, size(๐’โ‚)) + โˆ‚S2f = zeros(S, size(๐’โ‚‚_full)) + โˆ‚SS_acc = zeros(S, length(SS_and_pars)) + + โˆ‚ฮผสธ = โˆ‚mean_in + + # โ”€โ”€ Backprop through mean_of_variables โ”€โ”€ + # mean_of_variables = SS[1:n] + ล_to_yโ‚‚ * ฮผหขโบโ‚‚ + yvโ‚‚ + โˆ‚SS_acc[1:nVars] .+= โˆ‚ฮผสธ + โˆ‚ล_to_yโ‚‚ = โˆ‚ฮผสธ * ฮผหขโบโ‚‚' + โˆ‚ฮผหขโบโ‚‚ = ล_to_yโ‚‚' * โˆ‚ฮผสธ + โˆ‚yvโ‚‚ = copy(โˆ‚ฮผสธ) + + # โ”€โ”€ Backprop through (I - ล_to_ลโ‚‚) \ ลvโ‚‚ โ”€โ”€ + ฮป = A_mean' \ โˆ‚ฮผหขโบโ‚‚ + โˆ‚ลvโ‚‚ = copy(ฮป) + โˆ‚ล_to_ลโ‚‚ = ฮป * ฮผหขโบโ‚‚' # from -(I - A): sign is + + + # โ”€โ”€ yvโ‚‚ = (v_v_to_yโ‚‚_v + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 โ”€โ”€ + โˆ‚S2f[:, kron_v_v] .+= reshape(โˆ‚yvโ‚‚ / 2, :, 1) + โˆ‚S2f[:, kron_e_e] .+= (โˆ‚yvโ‚‚ / 2) * vec_Iโ‚‘' + + # โ”€โ”€ ลvโ‚‚ = [0; v_v/2 + e_eยทv/2; eโ‚โŠ—eโ‚ยทv] โ”€โ”€ + โˆ‚ลvโ‚‚_mid = โˆ‚ลvโ‚‚[nหข+1:2nหข] + โˆ‚ลvโ‚‚_bot = โˆ‚ลvโ‚‚[2nหข+1:end] + + โˆ‚S2f[iหข, kron_v_v] .+= reshape(โˆ‚ลvโ‚‚_mid / 2, :, 1) + โˆ‚S2f[iหข, kron_e_e] .+= (โˆ‚ลvโ‚‚_mid / 2) * vec_Iโ‚‘' + โˆ‚eโ‚keโ‚ = โˆ‚ลvโ‚‚_bot * vec_Iโ‚‘' + + # โ”€โ”€ ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚/2] โ”€โ”€ + โˆ‚๐’โ‚_acc[:, 1:nหข] .+= โˆ‚ล_to_yโ‚‚[:, 1:nหข] .+ โˆ‚ล_to_yโ‚‚[:, nหข+1:2nหข] + โˆ‚S2f[:, kron_s_s] .+= โˆ‚ล_to_yโ‚‚[:, 2nหข+1:end] / 2 + + # โ”€โ”€ ล_to_ลโ‚‚ block adjoints โ”€โ”€ + โˆ‚sโ‚_from_ลล = โˆ‚ล_to_ลโ‚‚[1:nหข, 1:nหข] + โˆ‚ล_to_ลโ‚‚[nหข+1:2nหข, nหข+1:2nหข] + โˆ‚ss2_from_ลล = โˆ‚ล_to_ลโ‚‚[nหข+1:2nหข, 2nหข+1:end] / 2 + โˆ‚sโ‚ksโ‚ = โˆ‚ล_to_ลโ‚‚[2nหข+1:end, 2nหข+1:end] + + # โ”€โ”€ Kron VJPs โ”€โ”€ + โˆ‚sโ‚_L, โˆ‚sโ‚_R = _kron_vjp(โˆ‚sโ‚ksโ‚, s_to_sโ‚, s_to_sโ‚) + โˆ‚eโ‚_L, โˆ‚eโ‚_R = _kron_vjp(โˆ‚eโ‚keโ‚, e_to_sโ‚, e_to_sโ‚) + + # Aggregate into ๐’โ‚ + โˆ‚๐’โ‚_acc[iหข, 1:nหข] .+= โˆ‚sโ‚_from_ลล .+ โˆ‚sโ‚_L .+ โˆ‚sโ‚_R + โˆ‚๐’โ‚_acc[iหข, nหข+1:end] .+= โˆ‚eโ‚_L .+ โˆ‚eโ‚_R + + # Aggregate into Sโ‚‚_full + โˆ‚S2f[iหข, kron_s_s] .+= โˆ‚ss2_from_ลล + + # โ”€โ”€ Sโ‚‚_full โ†’ Sโ‚‚_raw via ๐”โ‚‚ โ”€โ”€ + โˆ‚S2_raw = โˆ‚S2f * ๐”โ‚‚' + + # โ”€โ”€ Chain through sub-rrule pullbacks (reverse order) โ”€โ”€ + # Second-order solution + so2_grad = so2_pb((โˆ‚S2_raw, NoTangent())) + โˆ‚โˆ‡โ‚_acc = so2_grad[2] isa AbstractZero ? zeros(S, size(โˆ‡โ‚)) : collect(S, so2_grad[2]) + โˆ‚โˆ‡โ‚‚_total = so2_grad[3] isa AbstractZero ? zeros(S, size(โˆ‡โ‚‚)) : so2_grad[3] + โˆ‚๐’โ‚_from_so2 = so2_grad[4] isa AbstractZero ? zeros(S, size(๐’โ‚)) : collect(S, so2_grad[4]) + โˆ‚๐’โ‚_acc .+= โˆ‚๐’โ‚_from_so2 + + # Hessian + hess_grad = hess_pb(โˆ‚โˆ‡โ‚‚_total) + โˆ‚params_hess = hess_grad[2] isa AbstractZero ? zeros(S, np) : hess_grad[2] + โˆ‚SS_from_hess = hess_grad[3] isa AbstractZero ? zeros(S, length(SS_and_pars)) : hess_grad[3] + โˆ‚SS_acc .+= โˆ‚SS_from_hess + + # First-order solution + first_grad = first_pb((โˆ‚๐’โ‚_acc, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚_from_first = first_grad[2] isa AbstractZero ? zeros(S, size(โˆ‡โ‚)) : first_grad[2] + โˆ‚โˆ‡โ‚_acc .+= โˆ‚โˆ‡โ‚_from_first + + # Jacobian + jac_grad = jac_pb(โˆ‚โˆ‡โ‚_acc) + โˆ‚params_jac = jac_grad[2] isa AbstractZero ? zeros(S, np) : jac_grad[2] + โˆ‚SS_from_jac = jac_grad[3] isa AbstractZero ? zeros(S, length(SS_and_pars)) : jac_grad[3] + โˆ‚SS_acc .+= โˆ‚SS_from_jac + + # NSSS + nsss_grad = nsss_pb((โˆ‚SS_acc, NoTangent())) + โˆ‚params_nsss = nsss_grad[3] isa AbstractZero ? zeros(S, np) : nsss_grad[3] + + โˆ‚parameters_total = โˆ‚params_hess .+ โˆ‚params_jac .+ โˆ‚params_nsss + + return NoTangent(), โˆ‚parameters_total, NoTangent() + end + + return result, calculate_mean_pullback +end + + +# โ”€โ”€ calculate_second_order_moments rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +function rrule(::typeof(calculate_second_order_moments), + parameters::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options()) where S <: Real + + # โ”€โ”€ Non-differentiable setup โ”€โ”€ + constants_obj = initialise_constants!(๐“‚) + ensure_moments_constants!(constants_obj) + so = constants_obj.second_order + T_pm = constants_obj.post_model_macro + nแต‰ = T_pm.nExo + nหข = T_pm.nPast_not_future_and_mixed + nVars = T_pm.nVars + iหข = T_pm.past_not_future_and_mixed_idx + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + vec_Iโ‚‘ = so.vec_Iโ‚‘ + + zero_10() = (zeros(S,0), zeros(S,0), zeros(S,0,0), zeros(S,0,0), + zeros(S,0), zeros(S,0,0), zeros(S,0,0), spzeros(S,0,0), spzeros(S,0,0), false) + zero_pb(_) = (NoTangent(), zeros(S, length(parameters)), NoTangent()) + + # โ”€โ”€ Step 1: Covariance โ”€โ”€ + cov_out, cov_pb = rrule(calculate_covariance, parameters, ๐“‚; opts = opts) + ฮฃสธโ‚, ๐’โ‚, โˆ‡โ‚, SS_and_pars, solved = cov_out + + if !solved + return zero_10(), zero_pb + end + + ฮฃแถปโ‚ = ฮฃสธโ‚[iหข, iหข] + + # โ”€โ”€ Step 2: Hessian โ”€โ”€ + โˆ‡โ‚‚, hess_pb = rrule(calculate_hessian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces) + + # โ”€โ”€ Step 3: Second-order solution โ”€โ”€ + so2_out, so2_pb = rrule(calculate_second_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; opts = opts, parameter_values = parameters) + ๐’โ‚‚_raw = so2_out[1] + solved2 = so2_out[2] + + update_perturbation_counter!(๐“‚.counters, solved2, order = 2) + + if !solved2 + return (zeros(S,0), zeros(S,0), ฮฃสธโ‚, zeros(S,0,0), SS_and_pars, ๐’โ‚, โˆ‡โ‚, spzeros(S,0,0), โˆ‡โ‚‚, solved2), zero_pb + end + + # โ”€โ”€ Step 4: Decompress Sโ‚‚ (mutation-free) โ”€โ”€ + ๐’โ‚‚_full = ๐’โ‚‚_raw * ๐”โ‚‚ + + # โ”€โ”€ Step 5: Slicing and mean computation โ”€โ”€ + kron_s_s = so.kron_states + kron_e_e = so.kron_e_e + kron_v_v = so.kron_v_v + + # First-order slices + s_to_yโ‚ = ๐’โ‚[:, 1:nหข] + s_to_sโ‚ = ๐’โ‚[iหข, 1:nหข] + e_to_sโ‚ = ๐’โ‚[iหข, (nหข+1):end] + + # Second-order slices (dense) + s_s_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_s_s]) + e_e_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_e_e]) + v_v_to_yโ‚‚_v = vec(๐’โ‚‚_full[:, kron_v_v]) + s_s_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_s_s]) + e_e_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_e_e]) + v_v_to_sโ‚‚_v = vec(๐’โ‚‚_full[iหข, kron_v_v]) + + # Kronecker products + sโ‚_kron_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + eโ‚_kron_eโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) |> collect + + # Block matrices + ล_to_ลโ‚‚ = [ s_to_sโ‚ zeros(S, nหข, nหข + nหข^2) + zeros(S, nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 + zeros(S, nหข^2, 2*nหข) sโ‚_kron_sโ‚ ] + + ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2] + + ลvโ‚‚ = vcat(zeros(S, nหข), + v_v_to_sโ‚‚_v / 2 + e_e_to_sโ‚‚ * vec_Iโ‚‘ / 2, + eโ‚_kron_eโ‚ * vec_Iโ‚‘) + + yvโ‚‚ = (v_v_to_yโ‚‚_v + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 + + # Mean solve + A_mean = collect(โ„’.I(size(ล_to_ลโ‚‚, 1))) - ล_to_ลโ‚‚ + ฮผหขโบโ‚‚ = A_mean \ ลvโ‚‚ + + A_ฮ” = collect(โ„’.I(nหข)) - s_to_sโ‚ + rhs_ฮ” = s_s_to_sโ‚‚ * vec(ฮฃแถปโ‚) / 2 + (v_v_to_sโ‚‚_v + e_e_to_sโ‚‚ * vec_Iโ‚‘) / 2 + ฮ”ฮผหขโ‚‚ = vec(A_ฮ” \ rhs_ฮ”) + + ฮผสธโ‚‚ = SS_and_pars[1:nVars] + ล_to_yโ‚‚ * ฮผหขโบโ‚‚ + yvโ‚‚ + + slvd = solved && solved2 + ๐’โ‚‚_sp = sparse(๐’โ‚‚_full) + + result = (ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_sp, โˆ‡โ‚‚, slvd) + + # โ”€โ”€ Pullback โ”€โ”€ + function calculate_second_order_moments_pullback(โˆ‚out) + โˆ‚ฮผสธโ‚‚_in, โˆ‚ฮ”ฮผหขโ‚‚_in, โˆ‚ฮฃสธโ‚_pass, โˆ‚ฮฃแถปโ‚_pass, โˆ‚SS_pass, + โˆ‚๐’โ‚_pass, โˆ‚โˆ‡โ‚_pass, โˆ‚๐’โ‚‚_pass, โˆ‚โˆ‡โ‚‚_pass, _ = โˆ‚out + + # Materialise any InplaceableThunk / Thunk wrappers + โˆ‚ฮผสธโ‚‚_in = unthunk(โˆ‚ฮผสธโ‚‚_in) + โˆ‚ฮ”ฮผหขโ‚‚_in = unthunk(โˆ‚ฮ”ฮผหขโ‚‚_in) + โˆ‚ฮฃสธโ‚_pass = unthunk(โˆ‚ฮฃสธโ‚_pass) + โˆ‚ฮฃแถปโ‚_pass = unthunk(โˆ‚ฮฃแถปโ‚_pass) + โˆ‚SS_pass = unthunk(โˆ‚SS_pass) + โˆ‚๐’โ‚_pass = unthunk(โˆ‚๐’โ‚_pass) + โˆ‚โˆ‡โ‚_pass = unthunk(โˆ‚โˆ‡โ‚_pass) + โˆ‚๐’โ‚‚_pass = unthunk(โˆ‚๐’โ‚‚_pass) + โˆ‚โˆ‡โ‚‚_pass = unthunk(โˆ‚โˆ‡โ‚‚_pass) + + # Accumulators + โˆ‚๐’โ‚_acc = zeros(S, size(๐’โ‚)) + โˆ‚S2f = zeros(S, size(๐’โ‚‚_full)) + โˆ‚SS_acc = zeros(S, length(SS_and_pars)) + โˆ‚โˆ‡โ‚_acc = zeros(S, size(โˆ‡โ‚)) + โˆ‚ฮฃแถปโ‚_acc = zeros(S, nหข, nหข) + + # Pass-through cotangents + if !(โˆ‚๐’โ‚_pass isa AbstractZero); โˆ‚๐’โ‚_acc .+= โˆ‚๐’โ‚_pass; end + if !(โˆ‚SS_pass isa AbstractZero); โˆ‚SS_acc .+= โˆ‚SS_pass; end + if !(โˆ‚๐’โ‚‚_pass isa AbstractZero); โˆ‚S2f .+= โˆ‚๐’โ‚‚_pass; end + if !(โˆ‚โˆ‡โ‚_pass isa AbstractZero); โˆ‚โˆ‡โ‚_acc .+= โˆ‚โˆ‡โ‚_pass; end + if !(โˆ‚ฮฃแถปโ‚_pass isa AbstractZero); โˆ‚ฮฃแถปโ‚_acc .+= โˆ‚ฮฃแถปโ‚_pass; end + + # โ”€โ”€โ”€โ”€ Backprop through ฮผสธโ‚‚ โ”€โ”€โ”€โ”€ + if !(โˆ‚ฮผสธโ‚‚_in isa AbstractZero) + โˆ‚ฮผสธโ‚‚ = โˆ‚ฮผสธโ‚‚_in + # ฮผสธโ‚‚ = SS[1:n] + ล_to_yโ‚‚ * ฮผหขโบโ‚‚ + yvโ‚‚ + โˆ‚SS_acc[1:nVars] .+= โˆ‚ฮผสธโ‚‚ + โˆ‚ล_to_yโ‚‚ = โˆ‚ฮผสธโ‚‚ * ฮผหขโบโ‚‚' + โˆ‚ฮผหขโบโ‚‚ = ล_to_yโ‚‚' * โˆ‚ฮผสธโ‚‚ + โˆ‚yvโ‚‚ = copy(โˆ‚ฮผสธโ‚‚) + + # ฮผหขโบโ‚‚ = A_mean \ ลvโ‚‚ โ†’ ฮป = A_mean' \ โˆ‚ฮผหขโบโ‚‚ + ฮป = A_mean' \ โˆ‚ฮผหขโบโ‚‚ + โˆ‚ลvโ‚‚ = copy(ฮป) + โˆ‚ล_to_ลโ‚‚ = ฮป * ฮผหขโบโ‚‚' # from (I - ล_to_ลโ‚‚) + + # โ”€โ”€ yvโ‚‚ = (v_v_to_yโ‚‚_v + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 โ”€โ”€ + โˆ‚S2f[:, kron_v_v] .+= reshape(โˆ‚yvโ‚‚ / 2, :, 1) + โˆ‚S2f[:, kron_e_e] .+= (โˆ‚yvโ‚‚ / 2) * vec_Iโ‚‘' + + # โ”€โ”€ ลvโ‚‚ = [0; v_v/2 + e_eยทv/2; eโ‚โŠ—eโ‚ยทv] โ”€โ”€ + โˆ‚ลvโ‚‚_mid = โˆ‚ลvโ‚‚[nหข+1:2nหข] + โˆ‚ลvโ‚‚_bot = โˆ‚ลvโ‚‚[2nหข+1:end] + + โˆ‚S2f[iหข, kron_v_v] .+= reshape(โˆ‚ลvโ‚‚_mid / 2, :, 1) + โˆ‚S2f[iหข, kron_e_e] .+= (โˆ‚ลvโ‚‚_mid / 2) * vec_Iโ‚‘' + โˆ‚eโ‚keโ‚ = โˆ‚ลvโ‚‚_bot * vec_Iโ‚‘' + + # โ”€โ”€ ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚/2] โ”€โ”€ + โˆ‚๐’โ‚_acc[:, 1:nหข] .+= โˆ‚ล_to_yโ‚‚[:, 1:nหข] .+ โˆ‚ล_to_yโ‚‚[:, nหข+1:2nหข] + โˆ‚S2f[:, kron_s_s] .+= โˆ‚ล_to_yโ‚‚[:, 2nหข+1:end] / 2 + + # โ”€โ”€ ล_to_ลโ‚‚ blocks โ”€โ”€ + โˆ‚sโ‚_from_ลล = โˆ‚ล_to_ลโ‚‚[1:nหข, 1:nหข] + โˆ‚ล_to_ลโ‚‚[nหข+1:2nหข, nหข+1:2nหข] + โˆ‚ss2_from_ลล = โˆ‚ล_to_ลโ‚‚[nหข+1:2nหข, 2nหข+1:end] / 2 + โˆ‚sโ‚ksโ‚ = โˆ‚ล_to_ลโ‚‚[2nหข+1:end, 2nหข+1:end] + + # โ”€โ”€ Kron VJPs โ”€โ”€ + โˆ‚sโ‚_L, โˆ‚sโ‚_R = _kron_vjp(โˆ‚sโ‚ksโ‚, s_to_sโ‚, s_to_sโ‚) + โˆ‚eโ‚_L, โˆ‚eโ‚_R = _kron_vjp(โˆ‚eโ‚keโ‚, e_to_sโ‚, e_to_sโ‚) + + # Aggregate into ๐’โ‚ + โˆ‚๐’โ‚_acc[iหข, 1:nหข] .+= โˆ‚sโ‚_from_ลล .+ โˆ‚sโ‚_L .+ โˆ‚sโ‚_R + โˆ‚๐’โ‚_acc[iหข, nหข+1:end] .+= โˆ‚eโ‚_L .+ โˆ‚eโ‚_R + + # Aggregate into Sโ‚‚_full + โˆ‚S2f[iหข, kron_s_s] .+= โˆ‚ss2_from_ลล + end + + # โ”€โ”€โ”€โ”€ Backprop through ฮ”ฮผหขโ‚‚ โ”€โ”€โ”€โ”€ + if !(โˆ‚ฮ”ฮผหขโ‚‚_in isa AbstractZero) + โˆ‚ฮ”ฮผหขโ‚‚ = โˆ‚ฮ”ฮผหขโ‚‚_in + # ฮ”ฮผหขโ‚‚ = A_ฮ” \ rhs_ฮ” + ฮป_ฮ” = A_ฮ”' \ โˆ‚ฮ”ฮผหขโ‚‚ + # โˆ‚(I - s_to_sโ‚) โ†’ โˆ‚s_to_sโ‚ + โˆ‚๐’โ‚_acc[iหข, 1:nหข] .+= ฮป_ฮ” * ฮ”ฮผหขโ‚‚' + # rhs_ฮ” = s_s_to_sโ‚‚ * vec(ฮฃแถปโ‚)/2 + (v_v_to_sโ‚‚_v + e_e_to_sโ‚‚*vec_Iโ‚‘)/2 + โˆ‚S2f[iหข, kron_s_s] .+= ฮป_ฮ” * vec(ฮฃแถปโ‚)' / 2 + โˆ‚ฮฃแถปโ‚_acc .+= reshape(s_s_to_sโ‚‚' * ฮป_ฮ” / 2, nหข, nหข) + โˆ‚S2f[iหข, kron_v_v] .+= reshape(ฮป_ฮ” / 2, :, 1) + โˆ‚S2f[iหข, kron_e_e] .+= (ฮป_ฮ” / 2) * vec_Iโ‚‘' + end + + # โ”€โ”€ ฮฃแถปโ‚ โ†’ ฮฃสธโ‚ โ”€โ”€ + โˆ‚ฮฃสธโ‚ = zeros(S, size(ฮฃสธโ‚)) + โˆ‚ฮฃสธโ‚[iหข, iหข] .= โˆ‚ฮฃแถปโ‚_acc + if !(โˆ‚ฮฃสธโ‚_pass isa AbstractZero) + โˆ‚ฮฃสธโ‚ .+= โˆ‚ฮฃสธโ‚_pass + end + + # โ”€โ”€ Sโ‚‚_full โ†’ Sโ‚‚_raw via ๐”โ‚‚ โ”€โ”€ + โˆ‚S2_raw = โˆ‚S2f * ๐”โ‚‚' + + # โ”€โ”€ Chain through sub-rrule pullbacks โ”€โ”€ + # Second-order solution + so2_grad = so2_pb((โˆ‚S2_raw, NoTangent())) + # Coerce AbstractZero returns to typed zeros + โˆ‚โˆ‡โ‚_from_so2 = so2_grad[2] isa AbstractZero ? zeros(S, size(โˆ‡โ‚)) : so2_grad[2] + โˆ‚โˆ‡โ‚‚_total = so2_grad[3] isa AbstractZero ? zeros(S, size(โˆ‡โ‚‚)) : so2_grad[3] + โˆ‚๐’โ‚_from_so2 = so2_grad[4] isa AbstractZero ? zeros(S, size(๐’โ‚)) : so2_grad[4] + โˆ‚โˆ‡โ‚_acc .+= โˆ‚โˆ‡โ‚_from_so2 + โˆ‚๐’โ‚_acc .+= โˆ‚๐’โ‚_from_so2 + + if !(โˆ‚โˆ‡โ‚‚_pass isa AbstractZero) + โˆ‚โˆ‡โ‚‚_total = โˆ‚โˆ‡โ‚‚_total .+ โˆ‚โˆ‡โ‚‚_pass + end + + # Hessian + hess_grad = hess_pb(โˆ‚โˆ‡โ‚‚_total) + โˆ‚params_hess = hess_grad[2] isa AbstractZero ? zeros(S, length(parameters)) : hess_grad[2] + โˆ‚SS_from_hess = hess_grad[3] isa AbstractZero ? zeros(S, length(SS_and_pars)) : hess_grad[3] + โˆ‚SS_acc .+= โˆ‚SS_from_hess + + # Covariance (chains through NSSS โ†’ Jacobian โ†’ 1st sol โ†’ Lyapunov) + cov_grad = cov_pb((โˆ‚ฮฃสธโ‚, โˆ‚๐’โ‚_acc, โˆ‚โˆ‡โ‚_acc, โˆ‚SS_acc, NoTangent())) + โˆ‚params_cov = cov_grad[2] isa AbstractZero ? zeros(S, length(parameters)) : cov_grad[2] + + โˆ‚parameters_total = โˆ‚params_hess .+ โˆ‚params_cov + + return NoTangent(), โˆ‚parameters_total, NoTangent() + end + + return result, calculate_second_order_moments_pullback +end + + +# โ”€โ”€ calculate_second_order_moments_with_covariance rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +function rrule(::typeof(calculate_second_order_moments_with_covariance), + parameters::Vector{S}, + ๐“‚::โ„ณ; + opts::CalculationOptions = merge_calculation_options()) where S <: Real + + # โ”€โ”€ Non-differentiable setup โ”€โ”€ + constants_obj = initialise_constants!(๐“‚) + ensure_moments_constants!(constants_obj) + so = constants_obj.second_order + T_pm = constants_obj.post_model_macro + nแต‰ = T_pm.nExo + nหข = T_pm.nPast_not_future_and_mixed + nVars = T_pm.nVars + iหข = T_pm.past_not_future_and_mixed_idx + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + vec_Iโ‚‘ = so.vec_Iโ‚‘ + I_plus_s_s = so.I_plus_s_s + e4_minus = so.e4_minus_vecIโ‚‘_outer + Iโ‚‘ = collect(S, โ„’.I(nแต‰)) + + np = length(parameters) + zero_15() = (zeros(S,0,0), zeros(S,0,0), zeros(S,0), zeros(S,0), + zeros(S,0,0), zeros(S,0,0), zeros(S,0,0), + zeros(S,0,0), zeros(S,0,0), zeros(S,0), + zeros(S,0,0), zeros(S,0,0), spzeros(S,0,0), spzeros(S,0,0), false) + zero_pb(_) = (NoTangent(), zeros(S, np), NoTangent()) + + # โ”€โ”€ Step 1: Covariance โ”€โ”€ + cov_out, cov_pb = rrule(calculate_covariance, parameters, ๐“‚; opts = opts) + ฮฃสธโ‚, ๐’โ‚, โˆ‡โ‚, SS_and_pars, solved = cov_out + + if !solved; return zero_15(), zero_pb; end + + ฮฃแถปโ‚ = ฮฃสธโ‚[iหข, iหข] + + # โ”€โ”€ Step 2: Hessian โ”€โ”€ + โˆ‡โ‚‚, hess_pb = rrule(calculate_hessian, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces) + + # โ”€โ”€ Step 3: Second-order solution โ”€โ”€ + so2_out, so2_pb = rrule(calculate_second_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; opts = opts, parameter_values = parameters) + ๐’โ‚‚_raw, solved2 = so2_out + + update_perturbation_counter!(๐“‚.counters, solved2, order = 2) + + if !solved2; return zero_15(), zero_pb; end + + # โ”€โ”€ Step 4: Decompress Sโ‚‚ โ”€โ”€ + ๐’โ‚‚_full = ๐’โ‚‚_raw * ๐”โ‚‚ + + # โ”€โ”€ Step 5: Slicing โ”€โ”€ + kron_s_s = so.kron_states + kron_e_e = so.kron_e_e + kron_v_v = so.kron_v_v + kron_s_e = so.kron_s_e + + s_to_yโ‚ = ๐’โ‚[:, 1:nหข] + e_to_yโ‚ = ๐’โ‚[:, (nหข+1):end] + s_to_sโ‚ = ๐’โ‚[iหข, 1:nหข] + e_to_sโ‚ = ๐’โ‚[iหข, (nหข+1):end] + + s_s_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_s_s]) + e_e_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_e_e]) + v_v_to_yโ‚‚_v = vec(๐’โ‚‚_full[:, kron_v_v]) + s_e_to_yโ‚‚ = Matrix(๐’โ‚‚_full[:, kron_s_e]) + + s_s_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_s_s]) + e_e_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_e_e]) + v_v_to_sโ‚‚_v = vec(๐’โ‚‚_full[iหข, kron_v_v]) + s_e_to_sโ‚‚ = Matrix(๐’โ‚‚_full[iหข, kron_s_e]) + + # Kronecker products + sโ‚_kron_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + eโ‚_kron_eโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) |> collect + sโ‚_kron_eโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) |> collect + + # โ”€โ”€ Block matrices โ”€โ”€ + ล_to_ลโ‚‚ = [ s_to_sโ‚ zeros(S, nหข, nหข + nหข^2) + zeros(S, nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 + zeros(S, nหข^2, 2*nหข) sโ‚_kron_sโ‚ ] + + รช_to_ลโ‚‚ = [ e_to_sโ‚ zeros(S, nหข, nแต‰^2 + nแต‰ * nหข) + zeros(S, nหข, nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ + zeros(S, nหข^2, nแต‰) eโ‚_kron_eโ‚ I_plus_s_s * sโ‚_kron_eโ‚ ] + + ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2] + + รช_to_yโ‚‚ = [e_to_yโ‚ e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚] + + ลvโ‚‚ = vcat(zeros(S, nหข), + v_v_to_sโ‚‚_v / 2 + e_e_to_sโ‚‚ * vec_Iโ‚‘ / 2, + eโ‚_kron_eโ‚ * vec_Iโ‚‘) + + yvโ‚‚ = (v_v_to_yโ‚‚_v + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 + + # Mean solve + A_mean = collect(โ„’.I(size(ล_to_ลโ‚‚, 1))) - ล_to_ลโ‚‚ + ฮผหขโบโ‚‚ = A_mean \ ลvโ‚‚ + + A_ฮ” = collect(โ„’.I(nหข)) - s_to_sโ‚ + rhs_ฮ” = s_s_to_sโ‚‚ * vec(ฮฃแถปโ‚) / 2 + (v_v_to_sโ‚‚_v + e_e_to_sโ‚‚ * vec_Iโ‚‘) / 2 + ฮ”ฮผหขโ‚‚ = vec(A_ฮ” \ rhs_ฮ”) + + ฮผสธโ‚‚ = SS_and_pars[1:nVars] + ล_to_yโ‚‚ * ฮผหขโบโ‚‚ + yvโ‚‚ + + # โ”€โ”€ Step 6: Pruned covariance โ”€โ”€ + kron_ฮฃแถปโ‚_Iโ‚‘ = โ„’.kron(ฮฃแถปโ‚, Iโ‚‘) + + ฮ“โ‚‚ = [ Iโ‚‘ zeros(S, nแต‰, nแต‰^2 + nแต‰ * nหข) + zeros(S, nแต‰^2, nแต‰) e4_minus zeros(S, nแต‰^2, nแต‰ * nหข) + zeros(S, nหข * nแต‰, nแต‰ + nแต‰^2) kron_ฮฃแถปโ‚_Iโ‚‘ ] + + CC = รช_to_ลโ‚‚ * ฮ“โ‚‚ * รช_to_ลโ‚‚' + + lyap_ws_2nd = ensure_lyapunov_workspace!(๐“‚.workspaces, size(ล_to_ลโ‚‚, 1), :second_order) + + lyap_out, lyap_pb = rrule(solve_lyapunov_equation, + Float64.(ล_to_ลโ‚‚), Float64.(CC), lyap_ws_2nd; + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.second_order.ad.lyapunov, + verbose = opts.verbose) + ฮฃแถปโ‚‚ = lyap_out[1] + info = lyap_out[2] + + if !info; return zero_15(), zero_pb; end + + ฮฃสธโ‚‚ = ล_to_yโ‚‚ * ฮฃแถปโ‚‚ * ล_to_yโ‚‚' + รช_to_yโ‚‚ * ฮ“โ‚‚ * รช_to_yโ‚‚' + autocorr_tmp = ล_to_ลโ‚‚ * ฮฃแถปโ‚‚ * ล_to_yโ‚‚' + รช_to_ลโ‚‚ * ฮ“โ‚‚ * รช_to_yโ‚‚' + + slvd = solved && solved2 && info + + result = (ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, ล_to_ลโ‚‚, ล_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_raw, โˆ‡โ‚‚, slvd) + + # โ”€โ”€ Pullback โ”€โ”€ + function calculate_second_order_moments_with_covariance_pullback(โˆ‚out) + โˆ‚ฮฃสธโ‚‚_in, โˆ‚ฮฃแถปโ‚‚_pass, โˆ‚ฮผสธโ‚‚_in, โˆ‚ฮ”ฮผหขโ‚‚_in, โˆ‚at_in, + โˆ‚ลลโ‚‚_pass, โˆ‚ลyโ‚‚_pass, โˆ‚ฮฃสธโ‚_pass, โˆ‚ฮฃแถปโ‚_pass, โˆ‚SS_pass, + โˆ‚๐’โ‚_pass, โˆ‚โˆ‡โ‚_pass, โˆ‚๐’โ‚‚_pass, โˆ‚โˆ‡โ‚‚_pass, _ = โˆ‚out + + # Materialise any InplaceableThunk / Thunk wrappers + โˆ‚ฮฃสธโ‚‚_in = unthunk(โˆ‚ฮฃสธโ‚‚_in) + โˆ‚ฮฃแถปโ‚‚_pass = unthunk(โˆ‚ฮฃแถปโ‚‚_pass) + โˆ‚ฮผสธโ‚‚_in = unthunk(โˆ‚ฮผสธโ‚‚_in) + โˆ‚ฮ”ฮผหขโ‚‚_in = unthunk(โˆ‚ฮ”ฮผหขโ‚‚_in) + โˆ‚at_in = unthunk(โˆ‚at_in) + โˆ‚ลลโ‚‚_pass = unthunk(โˆ‚ลลโ‚‚_pass) + โˆ‚ลyโ‚‚_pass = unthunk(โˆ‚ลyโ‚‚_pass) + โˆ‚ฮฃสธโ‚_pass = unthunk(โˆ‚ฮฃสธโ‚_pass) + โˆ‚ฮฃแถปโ‚_pass = unthunk(โˆ‚ฮฃแถปโ‚_pass) + โˆ‚SS_pass = unthunk(โˆ‚SS_pass) + โˆ‚๐’โ‚_pass = unthunk(โˆ‚๐’โ‚_pass) + โˆ‚โˆ‡โ‚_pass = unthunk(โˆ‚โˆ‡โ‚_pass) + โˆ‚๐’โ‚‚_pass = unthunk(โˆ‚๐’โ‚‚_pass) + โˆ‚โˆ‡โ‚‚_pass = unthunk(โˆ‚โˆ‡โ‚‚_pass) + + # Accumulators + โˆ‚๐’โ‚_acc = zeros(S, size(๐’โ‚)) + โˆ‚S2f = zeros(S, size(๐’โ‚‚_full)) + โˆ‚SS_acc = zeros(S, length(SS_and_pars)) + โˆ‚โˆ‡โ‚_acc = zeros(S, size(โˆ‡โ‚)) + โˆ‚ฮฃแถปโ‚_acc = zeros(S, nหข, nหข) + + โˆ‚ล_to_ลโ‚‚_acc = zeros(S, size(ล_to_ลโ‚‚)) + โˆ‚ล_to_yโ‚‚_acc = zeros(S, size(ล_to_yโ‚‚)) + โˆ‚รช_to_ลโ‚‚_acc = zeros(S, size(รช_to_ลโ‚‚)) + โˆ‚รช_to_yโ‚‚_acc = zeros(S, size(รช_to_yโ‚‚)) + โˆ‚ฮ“โ‚‚_acc = zeros(S, size(ฮ“โ‚‚)) + โˆ‚ฮฃแถปโ‚‚_acc = zeros(S, size(ฮฃแถปโ‚‚)) + + # Pass-through cotangents + if !(โˆ‚๐’โ‚_pass isa AbstractZero); โˆ‚๐’โ‚_acc .+= โˆ‚๐’โ‚_pass; end + if !(โˆ‚SS_pass isa AbstractZero); โˆ‚SS_acc .+= โˆ‚SS_pass; end + # โˆ‚๐’โ‚‚_pass is now compressed โ€” accumulate after โˆ‚S2f * ๐”โ‚‚' conversion below + if !(โˆ‚โˆ‡โ‚_pass isa AbstractZero); โˆ‚โˆ‡โ‚_acc .+= โˆ‚โˆ‡โ‚_pass; end + if !(โˆ‚ฮฃแถปโ‚_pass isa AbstractZero); โˆ‚ฮฃแถปโ‚_acc .+= โˆ‚ฮฃแถปโ‚_pass; end + if !(โˆ‚ฮฃแถปโ‚‚_pass isa AbstractZero); โˆ‚ฮฃแถปโ‚‚_acc .+= โˆ‚ฮฃแถปโ‚‚_pass; end + if !(โˆ‚ลลโ‚‚_pass isa AbstractZero); โˆ‚ล_to_ลโ‚‚_acc .+= โˆ‚ลลโ‚‚_pass; end + if !(โˆ‚ลyโ‚‚_pass isa AbstractZero); โˆ‚ล_to_yโ‚‚_acc .+= โˆ‚ลyโ‚‚_pass; end + + # โ”€โ”€โ”€โ”€ Backprop through ฮฃสธโ‚‚ โ”€โ”€โ”€โ”€ + # ฮฃสธโ‚‚ = ล_to_yโ‚‚ * ฮฃแถปโ‚‚ * ล_to_yโ‚‚' + รช_to_yโ‚‚ * ฮ“โ‚‚ * รช_to_yโ‚‚' + if !(โˆ‚ฮฃสธโ‚‚_in isa AbstractZero) + โˆ‚ฮฃสธโ‚‚_sym = โˆ‚ฮฃสธโ‚‚_in + โˆ‚ฮฃสธโ‚‚_in' + โˆ‚ล_to_yโ‚‚_acc .+= โˆ‚ฮฃสธโ‚‚_sym * ล_to_yโ‚‚ * ฮฃแถปโ‚‚ + โˆ‚ฮฃแถปโ‚‚_acc .+= ล_to_yโ‚‚' * โˆ‚ฮฃสธโ‚‚_in * ล_to_yโ‚‚ + โˆ‚รช_to_yโ‚‚_acc .+= โˆ‚ฮฃสธโ‚‚_sym * รช_to_yโ‚‚ * ฮ“โ‚‚ + โˆ‚ฮ“โ‚‚_acc .+= รช_to_yโ‚‚' * โˆ‚ฮฃสธโ‚‚_in * รช_to_yโ‚‚ + end + + # โ”€โ”€โ”€โ”€ Backprop through autocorr_tmp โ”€โ”€โ”€โ”€ + # autocorr_tmp = ล_to_ลโ‚‚ * ฮฃแถปโ‚‚ * ล_to_yโ‚‚' + รช_to_ลโ‚‚ * ฮ“โ‚‚ * รช_to_yโ‚‚' + # For C = A*X*B': โˆ‚A = โˆ‚C*B*X', โˆ‚X = A'*โˆ‚C*B, โˆ‚B = โˆ‚C'*A*X + if !(โˆ‚at_in isa AbstractZero) + โˆ‚at = โˆ‚at_in + โˆ‚ล_to_ลโ‚‚_acc .+= โˆ‚at * ล_to_yโ‚‚ * ฮฃแถปโ‚‚ + โˆ‚ฮฃแถปโ‚‚_acc .+= ล_to_ลโ‚‚' * โˆ‚at * ล_to_yโ‚‚ + โˆ‚ล_to_yโ‚‚_acc .+= โˆ‚at' * ล_to_ลโ‚‚ * ฮฃแถปโ‚‚ + โˆ‚รช_to_ลโ‚‚_acc .+= โˆ‚at * รช_to_yโ‚‚ * ฮ“โ‚‚ + โˆ‚ฮ“โ‚‚_acc .+= รช_to_ลโ‚‚' * โˆ‚at * รช_to_yโ‚‚ + โˆ‚รช_to_yโ‚‚_acc .+= โˆ‚at' * รช_to_ลโ‚‚ * ฮ“โ‚‚ + end + + # โ”€โ”€โ”€โ”€ Backprop through Lyapunov: ฮฃแถปโ‚‚ = lyap(ล_to_ลโ‚‚, CC) โ”€โ”€โ”€โ”€ + lyap_grad = lyap_pb((โˆ‚ฮฃแถปโ‚‚_acc, NoTangent())) + โˆ‚ล_to_ลโ‚‚_lyap = lyap_grad[2] isa AbstractZero ? zeros(S, size(ล_to_ลโ‚‚)) : S.(lyap_grad[2]) + โˆ‚CC = lyap_grad[3] isa AbstractZero ? zeros(S, size(CC)) : S.(lyap_grad[3]) + โˆ‚ล_to_ลโ‚‚_acc .+= โˆ‚ล_to_ลโ‚‚_lyap + + # โ”€โ”€โ”€โ”€ Backprop through CC = รช_to_ลโ‚‚ * ฮ“โ‚‚ * รช_to_ลโ‚‚' โ”€โ”€โ”€โ”€ + โˆ‚CC_sym = โˆ‚CC + โˆ‚CC' + โˆ‚รช_to_ลโ‚‚_acc .+= โˆ‚CC_sym * รช_to_ลโ‚‚ * ฮ“โ‚‚ + โˆ‚ฮ“โ‚‚_acc .+= รช_to_ลโ‚‚' * โˆ‚CC * รช_to_ลโ‚‚ + + # โ”€โ”€โ”€โ”€ Backprop through ฮ“โ‚‚ โ†’ โˆ‚ฮฃแถปโ‚ โ”€โ”€โ”€โ”€ + # Only the bottom-right block kron(ฮฃแถปโ‚, Iโ‚‘) depends on parameters + br_row = nแต‰ + nแต‰^2 + โˆ‚ฮ“โ‚‚_br = โˆ‚ฮ“โ‚‚_acc[br_row+1:end, br_row+1:end] + โˆ‚ฮฃแถปโ‚_from_ฮ“โ‚‚, _ = _kron_vjp(โˆ‚ฮ“โ‚‚_br, ฮฃแถปโ‚, Iโ‚‘) + โˆ‚ฮฃแถปโ‚_acc .+= โˆ‚ฮฃแถปโ‚_from_ฮ“โ‚‚ + + # โ”€โ”€โ”€โ”€ Backprop through ฮผสธโ‚‚ (same as base) โ”€โ”€โ”€โ”€ + if !(โˆ‚ฮผสธโ‚‚_in isa AbstractZero) + โˆ‚ฮผสธโ‚‚ = โˆ‚ฮผสธโ‚‚_in + โˆ‚SS_acc[1:nVars] .+= โˆ‚ฮผสธโ‚‚ + โˆ‚ล_to_yโ‚‚_acc .+= โˆ‚ฮผสธโ‚‚ * ฮผหขโบโ‚‚' + โˆ‚ฮผหขโบโ‚‚ = ล_to_yโ‚‚' * โˆ‚ฮผสธโ‚‚ + โˆ‚yvโ‚‚ = copy(โˆ‚ฮผสธโ‚‚) + + ฮป = A_mean' \ โˆ‚ฮผหขโบโ‚‚ + โˆ‚ลvโ‚‚ = copy(ฮป) + โˆ‚ล_to_ลโ‚‚_acc .+= ฮป * ฮผหขโบโ‚‚' + + # yvโ‚‚ + โˆ‚S2f[:, kron_v_v] .+= reshape(โˆ‚yvโ‚‚ / 2, :, 1) + โˆ‚S2f[:, kron_e_e] .+= (โˆ‚yvโ‚‚ / 2) * vec_Iโ‚‘' + + # ลvโ‚‚ + โˆ‚ลvโ‚‚_mid = โˆ‚ลvโ‚‚[nหข+1:2nหข] + โˆ‚ลvโ‚‚_bot = โˆ‚ลvโ‚‚[2nหข+1:end] + โˆ‚S2f[iหข, kron_v_v] .+= reshape(โˆ‚ลvโ‚‚_mid / 2, :, 1) + โˆ‚S2f[iหข, kron_e_e] .+= (โˆ‚ลvโ‚‚_mid / 2) * vec_Iโ‚‘' + โˆ‚eโ‚keโ‚_from_ลv = โˆ‚ลvโ‚‚_bot * vec_Iโ‚‘' + else + โˆ‚eโ‚keโ‚_from_ลv = zeros(S, size(eโ‚_kron_eโ‚)) + end + + # โ”€โ”€โ”€โ”€ Backprop through ฮ”ฮผหขโ‚‚ โ”€โ”€โ”€โ”€ + if !(โˆ‚ฮ”ฮผหขโ‚‚_in isa AbstractZero) + ฮป_ฮ” = A_ฮ”' \ โˆ‚ฮ”ฮผหขโ‚‚_in + โˆ‚๐’โ‚_acc[iหข, 1:nหข] .+= ฮป_ฮ” * ฮ”ฮผหขโ‚‚' + โˆ‚S2f[iหข, kron_s_s] .+= ฮป_ฮ” * vec(ฮฃแถปโ‚)' / 2 + โˆ‚ฮฃแถปโ‚_acc .+= reshape(s_s_to_sโ‚‚' * ฮป_ฮ” / 2, nหข, nหข) + โˆ‚S2f[iหข, kron_v_v] .+= reshape(ฮป_ฮ” / 2, :, 1) + โˆ‚S2f[iหข, kron_e_e] .+= (ฮป_ฮ” / 2) * vec_Iโ‚‘' + end + + # โ”€โ”€โ”€โ”€ Distribute block matrix grads to slice grads โ”€โ”€โ”€โ”€ + # ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚/2] + โˆ‚๐’โ‚_acc[:, 1:nหข] .+= โˆ‚ล_to_yโ‚‚_acc[:, 1:nหข] .+ โˆ‚ล_to_yโ‚‚_acc[:, nหข+1:2nหข] + โˆ‚S2f[:, kron_s_s] .+= โˆ‚ล_to_yโ‚‚_acc[:, 2nหข+1:end] / 2 + + # รช_to_yโ‚‚ = [e_to_yโ‚ e_e_to_yโ‚‚/2 s_e_to_yโ‚‚] + โˆ‚๐’โ‚_acc[:, nหข+1:end] .+= โˆ‚รช_to_yโ‚‚_acc[:, 1:nแต‰] + โˆ‚S2f[:, kron_e_e] .+= โˆ‚รช_to_yโ‚‚_acc[:, nแต‰+1:nแต‰+nแต‰^2] / 2 + โˆ‚S2f[:, kron_s_e] .+= โˆ‚รช_to_yโ‚‚_acc[:, nแต‰+nแต‰^2+1:end] + + # ล_to_ลโ‚‚ blocks + โˆ‚sโ‚_from_ลล = โˆ‚ล_to_ลโ‚‚_acc[1:nหข, 1:nหข] + โˆ‚ล_to_ลโ‚‚_acc[nหข+1:2nหข, nหข+1:2nหข] + โˆ‚ss2_from_ลล = โˆ‚ล_to_ลโ‚‚_acc[nหข+1:2nหข, 2nหข+1:end] / 2 + โˆ‚sโ‚ksโ‚_from_ลล = โˆ‚ล_to_ลโ‚‚_acc[2nหข+1:end, 2nหข+1:end] + + # รช_to_ลโ‚‚ blocks + โˆ‚๐’โ‚_acc[iหข, nหข+1:end] .+= โˆ‚รช_to_ลโ‚‚_acc[1:nหข, 1:nแต‰] # e_to_sโ‚ + โˆ‚S2f[iหข, kron_e_e] .+= โˆ‚รช_to_ลโ‚‚_acc[nหข+1:2nหข, nแต‰+1:nแต‰+nแต‰^2] / 2 # e_e_to_sโ‚‚ + โˆ‚S2f[iหข, kron_s_e] .+= โˆ‚รช_to_ลโ‚‚_acc[nหข+1:2nหข, nแต‰+nแต‰^2+1:end] # s_e_to_sโ‚‚ + โˆ‚eโ‚keโ‚_from_รช = โˆ‚รช_to_ลโ‚‚_acc[2nหข+1:end, nแต‰+1:nแต‰+nแต‰^2] + โˆ‚Ips_sโ‚keโ‚ = โˆ‚รช_to_ลโ‚‚_acc[2nหข+1:end, nแต‰+nแต‰^2+1:end] + # I_plus_s_s * sโ‚_kron_eโ‚ โ†’ โˆ‚sโ‚_kron_eโ‚ += I_plus_s_s' * โˆ‚Ips_sโ‚keโ‚ + โˆ‚sโ‚keโ‚_from_รช = I_plus_s_s' * โˆ‚Ips_sโ‚keโ‚ + + # โ”€โ”€โ”€โ”€ Kron VJPs โ”€โ”€โ”€โ”€ + โˆ‚sโ‚_L, โˆ‚sโ‚_R = _kron_vjp(โˆ‚sโ‚ksโ‚_from_ลล, s_to_sโ‚, s_to_sโ‚) + โˆ‚eโ‚keโ‚_total = โˆ‚eโ‚keโ‚_from_ลv .+ โˆ‚eโ‚keโ‚_from_รช + โˆ‚eโ‚_L, โˆ‚eโ‚_R = _kron_vjp(โˆ‚eโ‚keโ‚_total, e_to_sโ‚, e_to_sโ‚) + โˆ‚sโ‚_se_L, โˆ‚eโ‚_se_R = _kron_vjp(โˆ‚sโ‚keโ‚_from_รช, s_to_sโ‚, e_to_sโ‚) + + # Aggregate into ๐’โ‚ + โˆ‚๐’โ‚_acc[iหข, 1:nหข] .+= โˆ‚sโ‚_from_ลล .+ โˆ‚sโ‚_L .+ โˆ‚sโ‚_R .+ โˆ‚sโ‚_se_L + โˆ‚๐’โ‚_acc[iหข, nหข+1:end] .+= โˆ‚eโ‚_L .+ โˆ‚eโ‚_R .+ โˆ‚eโ‚_se_R + โˆ‚S2f[iหข, kron_s_s] .+= โˆ‚ss2_from_ลล + + # โ”€โ”€ ฮฃแถปโ‚ โ†’ ฮฃสธโ‚ โ”€โ”€ + โˆ‚ฮฃสธโ‚ = zeros(S, size(ฮฃสธโ‚)) + โˆ‚ฮฃสธโ‚[iหข, iหข] .= โˆ‚ฮฃแถปโ‚_acc + if !(โˆ‚ฮฃสธโ‚_pass isa AbstractZero); โˆ‚ฮฃสธโ‚ .+= โˆ‚ฮฃสธโ‚_pass; end + + # โ”€โ”€ Sโ‚‚_full โ†’ Sโ‚‚_raw (compressed) โ”€โ”€ + โˆ‚S2_raw = โˆ‚S2f * ๐”โ‚‚' + # Add compressed pass-through from callers (position 13 now holds compressed ๐’โ‚‚_raw) + if !(โˆ‚๐’โ‚‚_pass isa AbstractZero); โˆ‚S2_raw .+= โˆ‚๐’โ‚‚_pass; end + + # โ”€โ”€ Chain through sub-rrule pullbacks โ”€โ”€ + so2_grad = so2_pb((โˆ‚S2_raw, NoTangent())) + โˆ‚โˆ‡โ‚_from_so2 = so2_grad[2] isa AbstractZero ? zeros(S, size(โˆ‡โ‚)) : so2_grad[2] + โˆ‚โˆ‡โ‚‚_total = so2_grad[3] isa AbstractZero ? zeros(S, size(โˆ‡โ‚‚)) : so2_grad[3] + โˆ‚๐’โ‚_from_so2 = so2_grad[4] isa AbstractZero ? zeros(S, size(๐’โ‚)) : so2_grad[4] + โˆ‚โˆ‡โ‚_acc .+= โˆ‚โˆ‡โ‚_from_so2 + โˆ‚๐’โ‚_acc .+= โˆ‚๐’โ‚_from_so2 + + if !(โˆ‚โˆ‡โ‚‚_pass isa AbstractZero); โˆ‚โˆ‡โ‚‚_total = โˆ‚โˆ‡โ‚‚_total .+ โˆ‚โˆ‡โ‚‚_pass; end + + hess_grad = hess_pb(โˆ‚โˆ‡โ‚‚_total) + โˆ‚params_hess = hess_grad[2] isa AbstractZero ? zeros(S, np) : hess_grad[2] + โˆ‚SS_from_hess = hess_grad[3] isa AbstractZero ? zeros(S, length(SS_and_pars)) : hess_grad[3] + โˆ‚SS_acc .+= โˆ‚SS_from_hess + + cov_grad = cov_pb((โˆ‚ฮฃสธโ‚, โˆ‚๐’โ‚_acc, โˆ‚โˆ‡โ‚_acc, โˆ‚SS_acc, NoTangent())) + โˆ‚params_cov = cov_grad[2] isa AbstractZero ? zeros(S, np) : cov_grad[2] + + โˆ‚parameters_total = โˆ‚params_hess .+ โˆ‚params_cov + + return NoTangent(), โˆ‚parameters_total, NoTangent() + end + + return result, calculate_second_order_moments_with_covariance_pullback +end + + +# โ”€โ”€ calculate_third_order_moments rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +function rrule(::typeof(calculate_third_order_moments), + parameters::Vector{T}, + observables::Union{Symbol_input,String_input}, + ๐“‚::โ„ณ; + covariance::Union{Symbol_input,String_input} = Symbol[], + opts::CalculationOptions = merge_calculation_options()) where T <: Real + + # โ”€โ”€ Non-differentiable constants โ”€โ”€ + ensure_moments_constants!(๐“‚.constants) + so = ๐“‚.constants.second_order + to = ๐“‚.constants.third_order + T_pm = ๐“‚.constants.post_model_macro + np = length(parameters) + nแต‰ = T_pm.nExo + + zero_4() = (zeros(T,0,0), zeros(T,0), zeros(T,0), false) + zero_pb(_) = (NoTangent(), zeros(T, np), NoTangent(), NoTangent()) + + # โ”€โ”€ Step 1: Second-order moments with covariance โ”€โ”€ + som2_out, som2_pb = rrule(calculate_second_order_moments_with_covariance, parameters, ๐“‚; opts = opts) + ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp_2, ล_to_ลโ‚‚, ล_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_raw, โˆ‡โ‚‚, solved = som2_out + + if !solved; return zero_4(), zero_pb; end + + # Expand compressed ๐’โ‚‚_raw to full for moments computation + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{T, Int} + + # โ”€โ”€ Step 2: Third-order derivatives โ”€โ”€ + โˆ‡โ‚ƒ, โˆ‡โ‚ƒ_pb = rrule(calculate_third_order_derivatives, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces) + + # โ”€โ”€ Step 3: Third-order solution (pass compressed ๐’โ‚‚_raw) โ”€โ”€ + so3_out, so3_pb = rrule(calculate_third_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚_raw, + ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, + parameter_values = parameters) + ๐’โ‚ƒ, solved3 = so3_out + + update_perturbation_counter!(๐“‚.counters, solved3, order = 3) + + if !solved3; return zero_4(), zero_pb; end + + # โ”€โ”€ Step 4: Decompress Sโ‚ƒ โ”€โ”€ + ๐”โ‚ƒ = ๐“‚.constants.third_order.๐”โ‚ƒ + ๐’โ‚ƒ_full = ๐’โ‚ƒ * ๐”โ‚ƒ + + ๐’โ‚ƒ_full = sparse(๐’โ‚ƒ_full) + + # โ”€โ”€ Step 5: Determine iteration groups โ”€โ”€ + orders = determine_efficient_order(๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ_full, ๐“‚.constants, observables, + covariance = covariance, tol = opts.tol.third_order.dependencies_tol) + + kron_e_e = so.kron_e_e + kron_v_v = so.kron_v_v + kron_e_v = to.kron_e_v + e_in_sโบ = so.e_in_sโบ + v_in_sโบ = so.v_in_sโบ + vec_Iโ‚‘ = so.vec_Iโ‚‘ + e4_nแต‰ยฒ_nแต‰ยฒ = so.e4_nแต‰ยฒ_nแต‰ยฒ + e4_nแต‰_nแต‰ยณ = so.e4_nแต‰_nแต‰ยณ + e4_minus_vecIโ‚‘_outer = so.e4_minus_vecIโ‚‘_outer + e6_nแต‰ยณ_nแต‰ยณ = to.e6_nแต‰ยณ_nแต‰ยณ + + ฮฃสธโ‚ƒ = zeros(T, size(ฮฃสธโ‚‚)) + solved_lyapunov = true + + # Per-iteration storage for pullback + n_iters = length(orders) + iter_data = Vector{Any}(undef, n_iters) + + for (iter_idx, ords) in enumerate(orders) + variance_observable, dependencies_all_vars = ords + + sort!(variance_observable) + sort!(dependencies_all_vars) + + dependencies = intersect(T_pm.past_not_future_and_mixed, dependencies_all_vars) + + obs_in_y = indexin(variance_observable, T_pm.var) + + dependencies_in_states_idx = indexin(dependencies, T_pm.past_not_future_and_mixed) + + dependencies_in_var_idx = Int.(indexin(dependencies, T_pm.var)) + + nหข = length(dependencies) + + iหข = dependencies_in_var_idx + + ฮฃฬ‚แถปโ‚ = ฮฃสธโ‚[iหข, iหข] + + dependencies_extended_idx = vcat(dependencies_in_states_idx, + dependencies_in_states_idx .+ T_pm.nPast_not_future_and_mixed, + findall(โ„’.kron(T_pm.past_not_future_and_mixed .โˆˆ (intersect(T_pm.past_not_future_and_mixed,dependencies),), + T_pm.past_not_future_and_mixed .โˆˆ (intersect(T_pm.past_not_future_and_mixed,dependencies),))) .+ 2*T_pm.nPast_not_future_and_mixed) + + ฮฃฬ‚แถปโ‚‚ = ฮฃแถปโ‚‚[dependencies_extended_idx, dependencies_extended_idx] + + ฮ”ฬ‚ฮผหขโ‚‚ = ฮ”ฮผหขโ‚‚[dependencies_in_states_idx] + + s_in_sโบ = BitVector(vcat(T_pm.past_not_future_and_mixed .โˆˆ (dependencies,), zeros(Bool, nแต‰ + 1))) + + substate_indices = ensure_moments_substate_indices!(๐“‚, nหข) + I_plus_s_s = substate_indices.I_plus_s_s + e_es = substate_indices.e_es + e_ss = substate_indices.e_ss + ss_s = substate_indices.ss_s + s_s = substate_indices.s_s + Dโ‚‚หข = substate_indices.Dโ‚‚หข + Lโ‚‚หข = substate_indices.Lโ‚‚หข + Dโ‚ƒหข = substate_indices.Dโ‚ƒหข + Lโ‚ƒหข = substate_indices.Lโ‚ƒหข + nโ‚‚หข = size(Dโ‚‚หข, 2) + nโ‚ƒหข = size(Dโ‚ƒหข, 2) + + # first order slices + s_to_yโ‚ = ๐’โ‚[obs_in_y,:][:,dependencies_in_states_idx] + e_to_yโ‚ = ๐’โ‚[obs_in_y,:][:, (T_pm.nPast_not_future_and_mixed + 1):end] + + s_to_sโ‚ = ๐’โ‚[iหข, dependencies_in_states_idx] + e_to_sโ‚ = ๐’โ‚[iหข, (T_pm.nPast_not_future_and_mixed + 1):end] + + # second order slices + dep_kron = ensure_moments_dependency_kron_indices!(๐“‚, dependencies, s_in_sโบ) + kron_s_s = dep_kron.kron_s_s + kron_s_e = dep_kron.kron_s_e + + s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_y,:][:, kron_s_s] + e_e_to_yโ‚‚ = ๐’โ‚‚[obs_in_y,:][:, kron_e_e] + s_e_to_yโ‚‚ = ๐’โ‚‚[obs_in_y,:][:, kron_s_e] + + s_s_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_s_s] |> collect + e_e_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_e_e] + v_v_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_v_v] |> collect + s_e_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_s_e] + + s_to_sโ‚_by_s_to_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + e_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) + s_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) + s_to_sโ‚_by_s_to_sโ‚_c = Lโ‚‚หข * s_to_sโ‚_by_s_to_sโ‚ * Dโ‚‚หข + + # third order slices + kron_s_v = dep_kron.kron_s_v + + kron_s_s_s = โ„’.kron(kron_s_s, s_in_sโบ) + kron_s_s_e = โ„’.kron(kron_s_s, e_in_sโบ) + kron_s_e_e = โ„’.kron(kron_s_e, e_in_sโบ) + kron_e_e_e = โ„’.kron(kron_e_e, e_in_sโบ) + kron_s_v_v = โ„’.kron(kron_s_v, v_in_sโบ) + kron_e_v_v = โ„’.kron(kron_e_v, v_in_sโบ) + + s_s_s_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_s_s] + s_s_e_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_s_e] + s_e_e_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_e_e] + e_e_e_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_e_e_e] + s_v_v_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_v_v] + e_v_v_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_e_v_v] + + s_s_s_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_s_s] + s_s_e_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_s_e] + s_e_e_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_e_e] + e_e_e_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_e_e_e] + s_v_v_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_v_v] + e_v_v_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_e_v_v] + + # Set up pruned state transition sub-blocks (compressed) + N_upper = 2 * nหข + nโ‚‚หข + N_lower = nหข + nหข^2 + nโ‚ƒหข + + A_UU = [s_to_sโ‚ spzeros(nหข, nหข + nโ‚‚หข) + spzeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 * Dโ‚‚หข + spzeros(nโ‚‚หข, 2 * nหข) s_to_sโ‚_by_s_to_sโ‚_c] + + A_LU = [s_v_v_to_sโ‚ƒ / 2 spzeros(nหข, nหข + nโ‚‚หข) + โ„’.kron(s_to_sโ‚,v_v_to_sโ‚‚ / 2) spzeros(nหข^2, nหข + nโ‚‚หข) + spzeros(nโ‚ƒหข, 2 * nหข + nโ‚‚หข)] + + A_LL = [s_to_sโ‚ s_s_to_sโ‚‚ s_s_s_to_sโ‚ƒ / 6 * Dโ‚ƒหข + spzeros(nหข^2, nหข) s_to_sโ‚_by_s_to_sโ‚ โ„’.kron(s_to_sโ‚,s_s_to_sโ‚‚ / 2) * Dโ‚ƒหข + spzeros(nโ‚ƒหข, nหข + nหข^2) Lโ‚ƒหข * โ„’.kron(s_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * Dโ‚ƒหข] + + รช_to_ลโ‚ƒ = [ e_to_sโ‚ zeros(nหข,nแต‰^2 + 2*nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + zeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ zeros(nหข,nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + zeros(nโ‚‚หข,nแต‰) Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ Lโ‚‚หข * I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚ zeros(nโ‚‚หข, nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + e_v_v_to_sโ‚ƒ / 2 zeros(nหข,nแต‰^2 + nแต‰ * nหข) s_e_to_sโ‚‚ s_s_e_to_sโ‚ƒ / 2 s_e_e_to_sโ‚ƒ / 2 e_e_e_to_sโ‚ƒ / 6 + โ„’.kron(e_to_sโ‚, v_v_to_sโ‚‚ / 2) zeros(nหข^2, nแต‰^2 + nแต‰ * nหข) s_s * s_to_sโ‚_by_e_to_sโ‚ โ„’.kron(s_to_sโ‚, s_e_to_sโ‚‚) + s_s * โ„’.kron(s_s_to_sโ‚‚ / 2, e_to_sโ‚) โ„’.kron(s_to_sโ‚, e_e_to_sโ‚‚ / 2) + s_s * โ„’.kron(s_e_to_sโ‚‚, e_to_sโ‚) โ„’.kron(e_to_sโ‚, e_e_to_sโ‚‚ / 2) + zeros(nโ‚ƒหข, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_s_to_sโ‚,e_to_sโ‚) + โ„’.kron(s_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * e_ss) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_e_to_sโ‚,e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_e_to_sโ‚) * e_es + โ„’.kron(e_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) * e_es) Lโ‚ƒหข * โ„’.kron(e_to_sโ‚,e_to_sโ‚_by_e_to_sโ‚)] + + ล_to_yโ‚ƒ = [s_to_yโ‚ + s_v_v_to_yโ‚ƒ / 2 s_to_yโ‚ s_s_to_yโ‚‚ / 2 * Dโ‚‚หข s_to_yโ‚ s_s_to_yโ‚‚ s_s_s_to_yโ‚ƒ / 6 * Dโ‚ƒหข] + + รช_to_yโ‚ƒ = [e_to_yโ‚ + e_v_v_to_yโ‚ƒ / 2 e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚ s_e_to_yโ‚‚ s_s_e_to_yโ‚ƒ / 2 s_e_e_to_yโ‚ƒ / 2 e_e_e_to_yโ‚ƒ / 6] + + ฮผหขโ‚ƒฮดฮผหขโ‚ = reshape((โ„’.I(size(s_to_sโ‚_by_s_to_sโ‚, 1)) - s_to_sโ‚_by_s_to_sโ‚) \ vec( + (s_s_to_sโ‚‚ * reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, nหข + 1:2*nหข] + vec(ฮฃฬ‚แถปโ‚) * ฮ”ฬ‚ฮผหขโ‚‚'),nหข^2, nหข) + + s_s_s_to_sโ‚ƒ * reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end , 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข^3, nหข) / 6 + + s_e_e_to_sโ‚ƒ * โ„’.kron(ฮฃฬ‚แถปโ‚, vec_Iโ‚‘) / 2 + + s_v_v_to_sโ‚ƒ * ฮฃฬ‚แถปโ‚ / 2) * s_to_sโ‚' + + (s_e_to_sโ‚‚ * โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚,โ„’.I(nแต‰)) + + e_e_e_to_sโ‚ƒ * e4_nแต‰_nแต‰ยณ' / 6 + + s_s_e_to_sโ‚ƒ * โ„’.kron(vec(ฮฃฬ‚แถปโ‚), โ„’.I(nแต‰)) / 2 + + e_v_v_to_sโ‚ƒ * โ„’.I(nแต‰) / 2) * e_to_sโ‚' + ), nหข, nหข) + + ฮ“โ‚ƒ = [ โ„’.I(nแต‰) spzeros(nแต‰, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚', โ„’.I(nแต‰)) โ„’.kron(vec(ฮฃฬ‚แถปโ‚)', โ„’.I(nแต‰)) spzeros(nแต‰, nหข * nแต‰^2) e4_nแต‰_nแต‰ยณ + spzeros(nแต‰^2, nแต‰) e4_minus_vecIโ‚‘_outer spzeros(nแต‰^2, 2*nหข*nแต‰ + nหข^2*nแต‰ + nหข*nแต‰^2 + nแต‰^3) + spzeros(nหข * nแต‰, nแต‰ + nแต‰^2) โ„’.kron(ฮฃฬ‚แถปโ‚, โ„’.I(nแต‰)) spzeros(nหข * nแต‰, nหข*nแต‰ + nหข^2*nแต‰ + nหข*nแต‰^2 + nแต‰^3) + โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚,โ„’.I(nแต‰)) spzeros(nแต‰ * nหข, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,nหข + 1:2*nหข] + ฮ”ฬ‚ฮผหขโ‚‚ * ฮ”ฬ‚ฮผหขโ‚‚',โ„’.I(nแต‰)) โ„’.kron(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)',โ„’.I(nแต‰)) spzeros(nแต‰ * nหข, nหข * nแต‰^2) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚, e4_nแต‰_nแต‰ยณ) + โ„’.kron(vec(ฮฃฬ‚แถปโ‚), โ„’.I(nแต‰)) spzeros(nแต‰ * nหข^2, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, nหข + 1:2*nหข] + vec(ฮฃฬ‚แถปโ‚) * ฮ”ฬ‚ฮผหขโ‚‚', โ„’.I(nแต‰)) โ„’.kron(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', โ„’.I(nแต‰)) spzeros(nแต‰ * nหข^2, nหข * nแต‰^2) โ„’.kron(vec(ฮฃฬ‚แถปโ‚), e4_nแต‰_nแต‰ยณ) + spzeros(nหข*nแต‰^2, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข + nหข^2*nแต‰) โ„’.kron(ฮฃฬ‚แถปโ‚, e4_nแต‰ยฒ_nแต‰ยฒ) spzeros(nหข*nแต‰^2,nแต‰^3) + e4_nแต‰_nแต‰ยณ' spzeros(nแต‰^3, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚', e4_nแต‰_nแต‰ยณ') โ„’.kron(vec(ฮฃฬ‚แถปโ‚)', e4_nแต‰_nแต‰ยณ') spzeros(nแต‰^3, nหข*nแต‰^2) e6_nแต‰ยณ_nแต‰ยณ] + + + Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข) + โ„’.kron(ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) zeros(nหข*nแต‰^2, nหข + nโ‚‚หข) โ„’.kron(ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3) * Lโ‚ƒหข', vec_Iโ‚‘) + spzeros(nแต‰^3, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข)] + + droptol!(A_UU, eps()) + droptol!(A_LU, eps()) + droptol!(A_LL, eps()) + droptol!(รช_to_ลโ‚ƒ, eps()) + droptol!(Eแดธแถป, eps()) + droptol!(ฮ“โ‚ƒ, eps()) + + # โ”€โ”€ Standard Lyapunov solve (compressed) โ”€โ”€ + ล_to_ลโ‚ƒ = [A_UU spzeros(N_upper, N_lower); A_LU A_LL] + + A_cross = Matrix{Float64}(รช_to_ลโ‚ƒ * Eแดธแถป) * ล_to_ลโ‚ƒ' + C_dense = Matrix{Float64}(รช_to_ลโ‚ƒ * ฮ“โ‚ƒ * รช_to_ลโ‚ƒ') + A_cross + A_cross' + + N_total = N_upper + N_lower + lyap_ws_3rd = Lyapunov_workspace(N_total) + lyap_out, lyap_pb_iter = rrule(solve_lyapunov_equation, + ล_to_ลโ‚ƒ, C_dense, lyap_ws_3rd, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.third_order.ad.lyapunov, + verbose = opts.verbose) + ฮฃแถปโ‚ƒ = lyap_out[1] + info = lyap_out[2] + + if !info + return zero_4(), zero_pb + end + + solved_lyapunov = solved_lyapunov && info + + ฮฃสธโ‚ƒtmp = ล_to_yโ‚ƒ * ฮฃแถปโ‚ƒ * ล_to_yโ‚ƒ' + รช_to_yโ‚ƒ * ฮ“โ‚ƒ * รช_to_yโ‚ƒ' + รช_to_yโ‚ƒ * Eแดธแถป * ล_to_yโ‚ƒ' + ล_to_yโ‚ƒ * Eแดธแถป' * รช_to_yโ‚ƒ' + + for obs in variance_observable + ฮฃสธโ‚ƒ[indexin([obs], T_pm.var), indexin(variance_observable, T_pm.var)] = ฮฃสธโ‚ƒtmp[indexin([obs], variance_observable), :] + end + + # Store per-iteration data for pullback + iter_data[iter_idx] = ( + variance_observable = variance_observable, + obs_in_y = obs_in_y, + iหข = iหข, + nหข = nหข, + dependencies_in_states_idx = dependencies_in_states_idx, + dependencies_extended_idx = dependencies_extended_idx, + ฮฃฬ‚แถปโ‚ = ฮฃฬ‚แถปโ‚, + ฮฃฬ‚แถปโ‚‚ = ฮฃฬ‚แถปโ‚‚, + ฮ”ฬ‚ฮผหขโ‚‚ = ฮ”ฬ‚ฮผหขโ‚‚, + s_in_sโบ = s_in_sโบ, + s_to_yโ‚ = s_to_yโ‚, + e_to_yโ‚ = e_to_yโ‚, + s_to_sโ‚ = s_to_sโ‚, + e_to_sโ‚ = e_to_sโ‚, + kron_s_s = kron_s_s, + kron_s_e = kron_s_e, + kron_s_v = kron_s_v, + kron_s_s_s = kron_s_s_s, + kron_s_s_e = kron_s_s_e, + kron_s_e_e = kron_s_e_e, + kron_e_e_e = kron_e_e_e, + kron_s_v_v = kron_s_v_v, + kron_e_v_v = kron_e_v_v, + s_s_to_yโ‚‚ = s_s_to_yโ‚‚, + e_e_to_yโ‚‚ = e_e_to_yโ‚‚, + s_e_to_yโ‚‚ = s_e_to_yโ‚‚, + s_s_to_sโ‚‚ = s_s_to_sโ‚‚, + e_e_to_sโ‚‚ = e_e_to_sโ‚‚, + v_v_to_sโ‚‚ = v_v_to_sโ‚‚, + s_e_to_sโ‚‚ = s_e_to_sโ‚‚, + s_to_sโ‚_by_s_to_sโ‚ = s_to_sโ‚_by_s_to_sโ‚, + e_to_sโ‚_by_e_to_sโ‚ = e_to_sโ‚_by_e_to_sโ‚, + s_to_sโ‚_by_e_to_sโ‚ = s_to_sโ‚_by_e_to_sโ‚, + s_s_s_to_yโ‚ƒ = s_s_s_to_yโ‚ƒ, + s_s_e_to_yโ‚ƒ = s_s_e_to_yโ‚ƒ, + s_e_e_to_yโ‚ƒ = s_e_e_to_yโ‚ƒ, + e_e_e_to_yโ‚ƒ = e_e_e_to_yโ‚ƒ, + s_v_v_to_yโ‚ƒ = s_v_v_to_yโ‚ƒ, + e_v_v_to_yโ‚ƒ = e_v_v_to_yโ‚ƒ, + s_s_s_to_sโ‚ƒ = s_s_s_to_sโ‚ƒ, + s_s_e_to_sโ‚ƒ = s_s_e_to_sโ‚ƒ, + s_e_e_to_sโ‚ƒ = s_e_e_to_sโ‚ƒ, + e_e_e_to_sโ‚ƒ = e_e_e_to_sโ‚ƒ, + s_v_v_to_sโ‚ƒ = s_v_v_to_sโ‚ƒ, + e_v_v_to_sโ‚ƒ = e_v_v_to_sโ‚ƒ, + รช_to_ลโ‚ƒ = รช_to_ลโ‚ƒ, + ล_to_yโ‚ƒ = ล_to_yโ‚ƒ, + รช_to_yโ‚ƒ = รช_to_yโ‚ƒ, + ฮ“โ‚ƒ = ฮ“โ‚ƒ, + Eแดธแถป = Eแดธแถป, + ล_to_ลโ‚ƒ = ล_to_ลโ‚ƒ, + ฮฃแถปโ‚ƒ = ฮฃแถปโ‚ƒ, + ฮฃสธโ‚ƒtmp = ฮฃสธโ‚ƒtmp, + ฮผหขโ‚ƒฮดฮผหขโ‚ = ฮผหขโ‚ƒฮดฮผหขโ‚, + lyap_pb = lyap_pb_iter, + N_upper = N_upper, + N_lower = N_lower, + Dโ‚‚หข = Dโ‚‚หข, + Lโ‚‚หข = Lโ‚‚หข, + Dโ‚ƒหข = Dโ‚ƒหข, + Lโ‚ƒหข = Lโ‚ƒหข, + nโ‚‚หข = nโ‚‚หข, + nโ‚ƒหข = nโ‚ƒหข, + s_to_sโ‚_by_s_to_sโ‚_c = s_to_sโ‚_by_s_to_sโ‚_c, + I_plus_s_s = I_plus_s_s, + ss_s = ss_s, + s_s = s_s, + e_es = e_es, + e_ss = e_ss, + ) + end + + result = (ฮฃสธโ‚ƒ, ฮผสธโ‚‚, SS_and_pars, solved && solved3 && solved_lyapunov) + + # โ”€โ”€ Pullback โ”€โ”€ + function calculate_third_order_moments_pullback(โˆ‚out) + โˆ‚ฮฃสธโ‚ƒ_in, โˆ‚ฮผสธโ‚‚_in, โˆ‚SS_in, _ = โˆ‚out + + โˆ‚ฮฃสธโ‚ƒ_in = unthunk(โˆ‚ฮฃสธโ‚ƒ_in) + โˆ‚ฮผสธโ‚‚_in = unthunk(โˆ‚ฮผสธโ‚‚_in) + โˆ‚SS_in = unthunk(โˆ‚SS_in) + + nโ‚‹ = T_pm.nPast_not_future_and_mixed + + # Accumulators for cotangents flowing to sub-rrule inputs + โˆ‚ฮฃสธโ‚_acc = zeros(T, size(ฮฃสธโ‚)) + โˆ‚ฮฃแถปโ‚‚_acc = zeros(T, size(ฮฃแถปโ‚‚)) + โˆ‚ฮ”ฮผหขโ‚‚_acc = zeros(T, length(ฮ”ฮผหขโ‚‚)) + โˆ‚๐’โ‚_acc = zeros(T, size(๐’โ‚)) + โˆ‚S2f_acc = zeros(T, size(๐’โ‚‚)) + โˆ‚S3f_acc = zeros(T, size(๐’โ‚ƒ_full)) + โˆ‚SS_acc = zeros(T, length(SS_and_pars)) + โˆ‚โˆ‡โ‚_acc = zeros(T, size(โˆ‡โ‚)) + โˆ‚โˆ‡โ‚‚_acc = zeros(T, size(โˆ‡โ‚‚)) + โˆ‚โˆ‡โ‚ƒ_acc = zeros(T, size(โˆ‡โ‚ƒ)) + + if !(โˆ‚SS_in isa AbstractZero); โˆ‚SS_acc .+= โˆ‚SS_in; end + + # โ”€โ”€โ”€โ”€ Reverse loop over iterations โ”€โ”€โ”€โ”€ + for iter_idx in n_iters:-1:1 + d = iter_data[iter_idx] + nหข_i = d.nหข + nโ‚‚หข_i = d.nโ‚‚หข + nโ‚ƒหข_i = d.nโ‚ƒหข + + # โ”€โ”€ Gather โˆ‚ฮฃสธโ‚ƒtmp from โˆ‚ฮฃสธโ‚ƒ (reverse of scatter) โ”€โ”€ + nObs_iter = length(d.variance_observable) + โˆ‚ฮฃสธโ‚ƒtmp = zeros(T, nObs_iter, nObs_iter) + + if !(โˆ‚ฮฃสธโ‚ƒ_in isa AbstractZero) + โˆ‚ฮฃสธโ‚ƒtmp .= โˆ‚ฮฃสธโ‚ƒ_in[d.obs_in_y, indexin(d.variance_observable, T_pm.var)] + end + + if โ„’.norm(โˆ‚ฮฃสธโ‚ƒtmp) < eps(T); continue; end + + โˆ‚ฮฃสธโ‚ƒtmp_sym = โˆ‚ฮฃสธโ‚ƒtmp + โˆ‚ฮฃสธโ‚ƒtmp' + + # โ”€โ”€ ฮฃสธโ‚ƒtmp = ล_y * ฮฃแถปโ‚ƒ * ล_y' + รช_y * ฮ“โ‚ƒ * รช_y' + รช_y * Eแดธแถป * ล_y' + ล_y * Eแดธแถป' * รช_y' โ”€โ”€ + # Terms 1+2 are AXA' forms; terms 3+4 form M + M' where M = รช_y * Eแดธแถป * ล_y'. + # Effective cotangent for M+M' is G_eff = โˆ‚ + โˆ‚' = โˆ‚ฮฃสธโ‚ƒtmp_sym. + + โˆ‚ล_to_yโ‚ƒ = โˆ‚ฮฃสธโ‚ƒtmp_sym * (d.ล_to_yโ‚ƒ * d.ฮฃแถปโ‚ƒ + d.รช_to_yโ‚ƒ * Matrix(d.Eแดธแถป)) + โˆ‚รช_to_yโ‚ƒ = โˆ‚ฮฃสธโ‚ƒtmp_sym * (d.รช_to_yโ‚ƒ * d.ฮ“โ‚ƒ + d.ล_to_yโ‚ƒ * Matrix(d.Eแดธแถป')) + โˆ‚ฮฃแถปโ‚ƒ = d.ล_to_yโ‚ƒ' * โˆ‚ฮฃสธโ‚ƒtmp * d.ล_to_yโ‚ƒ + โˆ‚ฮ“โ‚ƒ_iter = d.รช_to_yโ‚ƒ' * โˆ‚ฮฃสธโ‚ƒtmp * d.รช_to_yโ‚ƒ + โˆ‚Eแดธแถป_iter = d.รช_to_yโ‚ƒ' * โˆ‚ฮฃสธโ‚ƒtmp_sym * d.ล_to_yโ‚ƒ + + # โ”€โ”€ Standard Lyapunov adjoint โ”€โ”€ + Nu = d.N_upper; Nl = d.N_lower + ru_i = 1:Nu; rl_i = (Nu+1):(Nu+Nl) + + lyap_grad = d.lyap_pb((โˆ‚ฮฃแถปโ‚ƒ, NoTangent())) + โˆ‚ล_to_ลโ‚ƒ = lyap_grad[2] isa AbstractZero ? zeros(T, size(d.ล_to_ลโ‚ƒ)) : Matrix{T}(lyap_grad[2]) + โˆ‚C_lyap = lyap_grad[3] isa AbstractZero ? zeros(T, size(d.ล_to_ลโ‚ƒ)) : Matrix{T}(lyap_grad[3]) + + # Backprop through C = รช * ฮ“โ‚ƒ * รช' + M + M' where M = รช * Eแดธแถป * ล' + โˆ‚C_sym = โˆ‚C_lyap + โˆ‚C_lyap' + รช_d = Matrix{T}(d.รช_to_ลโ‚ƒ) + ล_d = Matrix{T}(d.ล_to_ลโ‚ƒ) + EL_d = Matrix{T}(d.Eแดธแถป) + ฮ“โ‚ƒ_d = Matrix{T}(d.ฮ“โ‚ƒ) + + # Term 1: รช * ฮ“โ‚ƒ * รช' + โˆ‚ฮ“โ‚ƒ_iter .+= รช_d' * โˆ‚C_lyap * รช_d + โˆ‚รช_to_ลโ‚ƒ = โˆ‚C_sym * รช_d * ฮ“โ‚ƒ_d + + # Terms 2+3: M + M' where M = รช * Eแดธแถป * ล' + โˆ‚รช_to_ลโ‚ƒ .+= โˆ‚C_sym * ล_d * EL_d' + โˆ‚Eแดธแถป_iter .+= รช_d' * โˆ‚C_sym * ล_d + โˆ‚ล_to_ลโ‚ƒ .+= โˆ‚C_sym' * รช_d * EL_d + + # Extract โˆ‚A_UU, โˆ‚A_LU, โˆ‚A_LL from โˆ‚ล_to_ลโ‚ƒ + โˆ‚A_UU = โˆ‚ล_to_ลโ‚ƒ[ru_i, ru_i] + โˆ‚A_LU = โˆ‚ล_to_ลโ‚ƒ[rl_i, ru_i] + โˆ‚A_LL = โˆ‚ล_to_ลโ‚ƒ[rl_i, rl_i] + + + # โ”€โ”€ Disaggregate ล_to_yโ‚ƒ โ†’ โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, โˆ‚๐’โ‚ƒ โ”€โ”€ + # ล_to_yโ‚ƒ = [s_to_yโ‚+svv/2 | s_to_yโ‚ | ss_to_yโ‚‚/2 | s_to_yโ‚ | ss_to_yโ‚‚ | sss_to_yโ‚ƒ/6] + c = 0 + โˆ‚blk1 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i]; c += nหข_i + โˆ‚blk2 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i]; c += nหข_i + โˆ‚blk3 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nโ‚‚หข_i]; c += nโ‚‚หข_i # compressed + โˆ‚blk4 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i]; c += nหข_i + โˆ‚blk5 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i^2]; c += nหข_i^2 + โˆ‚blk6 = โˆ‚ล_to_yโ‚ƒ[:, c+1:end] + + โˆ‚๐’โ‚_acc[d.obs_in_y, d.dependencies_in_states_idx] .+= โˆ‚blk1 .+ โˆ‚blk2 .+ โˆ‚blk4 # โˆ‚s_to_yโ‚ + โˆ‚S2f_acc[d.obs_in_y, d.kron_s_s] .+= (โˆ‚blk3 * Matrix(d.Dโ‚‚หข)') ./ 2 .+ โˆ‚blk5 # โˆ‚s_s_to_yโ‚‚ (decompress blk3) + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_v_v] .+= โˆ‚blk1 ./ 2 # โˆ‚s_v_v_to_yโ‚ƒ + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_s_s] .+= (โˆ‚blk6 * Matrix(d.Dโ‚ƒหข)') ./ 6 # โˆ‚s_s_s_to_yโ‚ƒ (decompress blk6) + + # โ”€โ”€ Disaggregate รช_to_yโ‚ƒ โ†’ โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, โˆ‚๐’โ‚ƒ โ”€โ”€ + # รช_to_yโ‚ƒ = [e_to_yโ‚+evv/2 | ee_to_yโ‚‚/2 | se_to_yโ‚‚ | se_to_yโ‚‚ | sse_to_yโ‚ƒ/2 | see_to_yโ‚ƒ/2 | eee_to_yโ‚ƒ/6] + c = 0 + โˆ‚eblk1 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nแต‰]; c += nแต‰ + โˆ‚eblk2 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nแต‰^2]; c += nแต‰^2 + โˆ‚eblk3 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i*nแต‰]; c += nหข_i*nแต‰ + โˆ‚eblk4 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i*nแต‰]; c += nหข_i*nแต‰ + โˆ‚eblk5 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i^2*nแต‰]; c += nหข_i^2*nแต‰ + โˆ‚eblk6 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i*nแต‰^2]; c += nหข_i*nแต‰^2 + โˆ‚eblk7 = โˆ‚รช_to_yโ‚ƒ[:, c+1:end] + + โˆ‚๐’โ‚_acc[d.obs_in_y, nโ‚‹+1:end] .+= โˆ‚eblk1 # โˆ‚e_to_yโ‚ + โˆ‚S2f_acc[d.obs_in_y, kron_e_e] .+= โˆ‚eblk2 ./ 2 # โˆ‚e_e_to_yโ‚‚ + โˆ‚S2f_acc[d.obs_in_y, d.kron_s_e] .+= โˆ‚eblk3 .+ โˆ‚eblk4 # โˆ‚s_e_to_yโ‚‚ + โˆ‚S3f_acc[d.obs_in_y, d.kron_e_v_v] .+= โˆ‚eblk1 ./ 2 # โˆ‚e_v_v_to_yโ‚ƒ + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_s_e] .+= โˆ‚eblk5 ./ 2 # โˆ‚s_s_e_to_yโ‚ƒ + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_e_e] .+= โˆ‚eblk6 ./ 2 # โˆ‚s_e_e_to_yโ‚ƒ + โˆ‚S3f_acc[d.obs_in_y, d.kron_e_e_e] .+= โˆ‚eblk7 ./ 6 # โˆ‚e_e_e_to_yโ‚ƒ + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # Stage 2+3: Disaggregate block matrices โ†’ slice & data cotangents + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + n = nหข_i; ne = nแต‰ + Ine = Matrix{T}(โ„’.I(ne)) + vec_Ie_col = reshape(T.(vec_Iโ‚‘), :, 1) + + # Dense copies of frequently used slices + sโ‚ = Matrix{T}(d.s_to_sโ‚) + eโ‚ = Matrix{T}(d.e_to_sโ‚) + sโ‚ยฒ = Matrix{T}(d.s_to_sโ‚_by_s_to_sโ‚) + eโ‚ยฒ = Matrix{T}(d.e_to_sโ‚_by_e_to_sโ‚) + sโ‚eโ‚ = Matrix{T}(d.s_to_sโ‚_by_e_to_sโ‚) + ssโ‚‚ = Matrix{T}(d.s_s_to_sโ‚‚) + eeโ‚‚ = Matrix{T}(d.e_e_to_sโ‚‚) + seโ‚‚ = Matrix{T}(d.s_e_to_sโ‚‚) + vvโ‚‚ = Matrix{T}(d.v_v_to_sโ‚‚) + + # Local slice cotangent accumulators + โˆ‚sโ‚_l = zeros(T, n, n) + โˆ‚eโ‚_l = zeros(T, n, ne) + โˆ‚ssโ‚‚_l = zeros(T, n, n^2) + โˆ‚eeโ‚‚_l = zeros(T, n, ne^2) + โˆ‚seโ‚‚_l = zeros(T, n, n * ne) + โˆ‚vvโ‚‚_l = zeros(T, size(vvโ‚‚)) + โˆ‚ฮฃฬ‚แถปโ‚ = zeros(T, n, n) + โˆ‚ฮฃฬ‚แถปโ‚‚ = zeros(T, size(d.ฮฃฬ‚แถปโ‚‚)) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l = zeros(T, n) + + # Block boundary arrays + sb = cumsum([0, n, n, nโ‚‚หข_i, n, n^2, nโ‚ƒหข_i]) # ล_to_ลโ‚ƒ row/col (compressed) + eb = cumsum([0, ne, ne^2, n*ne, n*ne, n^2*ne, n*ne^2, ne^3]) # รช_to_ลโ‚ƒ cols + gb = eb # ฮ“โ‚ƒ row/col (same block sizes) + + vvh = vvโ‚‚ ./ 2; ssh = ssโ‚‚ ./ 2; eeh = eeโ‚‚ ./ 2 + + # Reusable buffers for in-place kron adjoint operations + โˆ‚sโ‚ยฒ_buf = zeros(T, n^2, n^2) + โˆ‚eโ‚ยฒ_buf = zeros(T, n^2, ne^2) + โˆ‚kron_buf = zeros(T, n^2, n * ne) + โˆ‚vvh_buf = zeros(T, size(vvh)) + โˆ‚ssh_buf = zeros(T, size(ssh)) + โˆ‚eeh_buf = zeros(T, size(eeh)) + + # โ”€โ”€ 2a: A_UU, A_LU, A_LL disaggregation โ”€โ”€ + # Block boundaries within sub-matrices + bu = cumsum([0, n, n, nโ‚‚หข_i]) # A_UU row/col blocks + bl = cumsum([0, n, n^2, nโ‚ƒหข_i]) # A_LL row/col blocks (also A_LU rows) + + # โ”€โ”€ From โˆ‚A_UU โ”€โ”€ + # (1,1) sโ‚, (2,2) sโ‚ + โˆ‚sโ‚_l .+= โˆ‚A_UU[bu[1]+1:bu[2], bu[1]+1:bu[2]] .+ + โˆ‚A_UU[bu[2]+1:bu[3], bu[2]+1:bu[3]] + # (2,3) ssโ‚‚/2 * Dโ‚‚หข โ€” decompress cols + โˆ‚ssโ‚‚_l .+= โˆ‚A_UU[bu[2]+1:bu[3], bu[3]+1:bu[4]] * Matrix(d.Dโ‚‚หข)' ./ 2 + # (3,3) Lโ‚‚หข * kron(sโ‚,sโ‚) * Dโ‚‚หข โ€” decompress then kron_vjp + โˆ‚inner33 = Matrix(d.Lโ‚‚หข)' * Matrix(โˆ‚A_UU[bu[3]+1:bu[4], bu[3]+1:bu[4]]) * Matrix(d.Dโ‚‚หข)' + fill_kron_adjoint!(โˆ‚sโ‚_l, โˆ‚sโ‚_l, โˆ‚inner33, sโ‚, sโ‚) + + # โ”€โ”€ From โˆ‚A_LU โ”€โ”€ + # (1,1) s_vvโ‚ƒ/2 + โˆ‚S3f_acc[d.iหข, d.kron_s_v_v] .+= โˆ‚A_LU[bl[1]+1:bl[2], bu[1]+1:bu[2]] ./ 2 + # (2,1) kron(sโ‚, vvโ‚‚/2) + โˆ‚vvh_buf .= 0 + fill_kron_adjoint!(โˆ‚vvh_buf, โˆ‚sโ‚_l, Matrix(โˆ‚A_LU[bl[2]+1:bl[3], bu[1]+1:bu[2]]), vvh, sโ‚) + โˆ‚vvโ‚‚_l .+= โˆ‚vvh_buf ./ 2 + + # โ”€โ”€ From โˆ‚A_LL โ”€โ”€ + # (1,1) sโ‚ + โˆ‚sโ‚_l .+= โˆ‚A_LL[bl[1]+1:bl[2], bl[1]+1:bl[2]] + # (1,2) ssโ‚‚ + โˆ‚ssโ‚‚_l .+= โˆ‚A_LL[bl[1]+1:bl[2], bl[2]+1:bl[3]] + # (1,3) sssโ‚ƒ/6 * Dโ‚ƒหข โ€” decompress cols + โˆ‚S3f_acc[d.iหข, d.kron_s_s_s] .+= โˆ‚A_LL[bl[1]+1:bl[2], bl[3]+1:bl[4]] * Matrix(d.Dโ‚ƒหข)' ./ 6 + # (2,2) kron(sโ‚,sโ‚) + fill_kron_adjoint!(โˆ‚sโ‚_l, โˆ‚sโ‚_l, Matrix(โˆ‚A_LL[bl[2]+1:bl[3], bl[2]+1:bl[3]]), sโ‚, sโ‚) + # (2,3) kron(sโ‚, ssโ‚‚/2) * Dโ‚ƒหข โ€” decompress cols then kron_vjp + โˆ‚inner56 = Matrix(โˆ‚A_LL[bl[2]+1:bl[3], bl[3]+1:bl[4]]) * Matrix(d.Dโ‚ƒหข)' + โˆ‚ssh_buf .= 0 + fill_kron_adjoint!(โˆ‚ssh_buf, โˆ‚sโ‚_l, โˆ‚inner56, ssh, sโ‚) + โˆ‚ssโ‚‚_l .+= โˆ‚ssh_buf ./ 2 + # (3,3) Lโ‚ƒหข * kron(sโ‚, kron(sโ‚,sโ‚)) * Dโ‚ƒหข โ€” decompress then kron_vjp + โˆ‚inner66 = Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚A_LL[bl[3]+1:bl[4], bl[3]+1:bl[4]]) * Matrix(d.Dโ‚ƒหข)' + โˆ‚sโ‚ยฒ_buf .= 0 + fill_kron_adjoint!(โˆ‚sโ‚ยฒ_buf, โˆ‚sโ‚_l, โˆ‚inner66, sโ‚ยฒ, sโ‚) + fill_kron_adjoint!(โˆ‚sโ‚_l, โˆ‚sโ‚_l, โˆ‚sโ‚ยฒ_buf, sโ‚, sโ‚) + + + # โ”€โ”€ 2b: รช_to_ลโ‚ƒ disaggregation โ”€โ”€ + โˆ‚รชโ‚ƒ = Matrix{T}(โˆ‚รช_to_ลโ‚ƒ) + ss_s1e1 = Matrix(d.s_s) * sโ‚eโ‚ # pre-compute + + # Row 1: (1,1) eโ‚ + โˆ‚eโ‚_l .+= โˆ‚รชโ‚ƒ[sb[1]+1:sb[2], eb[1]+1:eb[2]] + # Row 2: (2,2) eeโ‚‚/2; (2,3) seโ‚‚ + โˆ‚eeโ‚‚_l .+= โˆ‚รชโ‚ƒ[sb[2]+1:sb[3], eb[2]+1:eb[3]] ./ 2 + โˆ‚seโ‚‚_l .+= โˆ‚รชโ‚ƒ[sb[2]+1:sb[3], eb[3]+1:eb[4]] + # Row 3: (3,2) Lโ‚‚หข * kron(eโ‚,eโ‚) โ€” decompress rows + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚eโ‚_l, Matrix(d.Lโ‚‚หข)' * Matrix(โˆ‚รชโ‚ƒ[sb[3]+1:sb[4], eb[2]+1:eb[3]]), eโ‚, eโ‚) + # (3,3) Lโ‚‚หข * I_plus_s_s * kron(sโ‚,eโ‚) โ€” decompress rows + โˆ‚k33 = Matrix(d.I_plus_s_s') * Matrix(d.Lโ‚‚หข)' * Matrix(โˆ‚รชโ‚ƒ[sb[3]+1:sb[4], eb[3]+1:eb[4]]) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚_l, โˆ‚k33, eโ‚, sโ‚) + # Row 4: direct Sโ‚ƒ slices + โˆ‚S3f_acc[d.iหข, d.kron_e_v_v] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[1]+1:eb[2]] ./ 2 + โˆ‚seโ‚‚_l .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[4]+1:eb[5]] + โˆ‚S3f_acc[d.iหข, d.kron_s_s_e] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[5]+1:eb[6]] ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_s_e_e] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[6]+1:eb[7]] ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_e_e_e] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[7]+1:eb[8]] ./ 6 + # Row 5: (5,1) kron(eโ‚,vvโ‚‚/2) + โˆ‚vvh_buf .= 0 + fill_kron_adjoint!(โˆ‚vvh_buf, โˆ‚eโ‚_l, Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[1]+1:eb[2]]), vvh, eโ‚) + โˆ‚vvโ‚‚_l .+= โˆ‚vvh_buf ./ 2 + # (5,4) s_s * kron(sโ‚,eโ‚) + โˆ‚k54 = Matrix(d.s_s') * Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[4]+1:eb[5]]) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚_l, โˆ‚k54, eโ‚, sโ‚) + # (5,5) kron(sโ‚,seโ‚‚) + s_s * kron(ssโ‚‚/2, eโ‚) + โˆ‚b55 = Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[5]+1:eb[6]]) + fill_kron_adjoint!(โˆ‚seโ‚‚_l, โˆ‚sโ‚_l, โˆ‚b55, seโ‚‚, sโ‚) + โˆ‚k55b = Matrix(d.s_s') * โˆ‚b55 + โˆ‚ssh_buf .= 0 + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚ssh_buf, โˆ‚k55b, eโ‚, ssh) + โˆ‚ssโ‚‚_l .+= โˆ‚ssh_buf ./ 2 + # (5,6) kron(sโ‚,eeโ‚‚/2) + s_s * kron(seโ‚‚, eโ‚) + โˆ‚b56 = Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[6]+1:eb[7]]) + โˆ‚eeh_buf .= 0 + fill_kron_adjoint!(โˆ‚eeh_buf, โˆ‚sโ‚_l, โˆ‚b56, eeh, sโ‚) + โˆ‚eeโ‚‚_l .+= โˆ‚eeh_buf ./ 2 + โˆ‚k56b = Matrix(d.s_s') * โˆ‚b56 + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚seโ‚‚_l, โˆ‚k56b, eโ‚, seโ‚‚) + # (5,7) kron(eโ‚, eeโ‚‚/2) + โˆ‚eeh_buf .= 0 + fill_kron_adjoint!(โˆ‚eeh_buf, โˆ‚eโ‚_l, Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[7]+1:eb[8]]), eeh, eโ‚) + โˆ‚eeโ‚‚_l .+= โˆ‚eeh_buf ./ 2 + # Row 6: (6,5) Lโ‚ƒหข * (kron(sโ‚ยฒ,eโ‚) + kron(sโ‚,s_s*sโ‚eโ‚) + kron(eโ‚,sโ‚ยฒ)*e_ss) โ€” decompress rows + โˆ‚b65 = Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚รชโ‚ƒ[sb[6]+1:sb[7], eb[5]+1:eb[6]]) + โˆ‚sโ‚ยฒ_buf .= 0 # Term 1: kron(sโ‚ยฒ, eโ‚) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚ยฒ_buf, โˆ‚b65, eโ‚, sโ‚ยฒ) + fill_kron_adjoint!(โˆ‚sโ‚_l, โˆ‚sโ‚_l, โˆ‚sโ‚ยฒ_buf, sโ‚, sโ‚) + โˆ‚kron_buf .= 0 # Term 2: kron(sโ‚, ss_s1e1) + fill_kron_adjoint!(โˆ‚kron_buf, โˆ‚sโ‚_l, โˆ‚b65, ss_s1e1, sโ‚) + tmpC = Matrix(d.s_s') * โˆ‚kron_buf + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚_l, tmpC, eโ‚, sโ‚) + โˆ‚k65c = โˆ‚b65 * Matrix(d.e_ss') # Term 3: kron(eโ‚, sโ‚ยฒ) * e_ss + โˆ‚sโ‚ยฒ_buf .= 0 + fill_kron_adjoint!(โˆ‚sโ‚ยฒ_buf, โˆ‚eโ‚_l, โˆ‚k65c, sโ‚ยฒ, eโ‚) + fill_kron_adjoint!(โˆ‚sโ‚_l, โˆ‚sโ‚_l, โˆ‚sโ‚ยฒ_buf, sโ‚, sโ‚) + # (6,6) Lโ‚ƒหข * (kron(sโ‚eโ‚,eโ‚) + kron(eโ‚,sโ‚eโ‚)*e_es + kron(eโ‚,s_s*sโ‚eโ‚)*e_es) โ€” decompress rows + โˆ‚b66 = Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚รชโ‚ƒ[sb[6]+1:sb[7], eb[6]+1:eb[7]]) + โˆ‚kron_buf .= 0 # Term 1: kron(sโ‚eโ‚, eโ‚) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚kron_buf, โˆ‚b66, eโ‚, sโ‚eโ‚) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚_l, โˆ‚kron_buf, eโ‚, sโ‚) + โˆ‚pre = โˆ‚b66 * Matrix(d.e_es') # shared for Terms 2+3 + โˆ‚kron_buf .= 0 # Term 2: kron(eโ‚, sโ‚eโ‚) + fill_kron_adjoint!(โˆ‚kron_buf, โˆ‚eโ‚_l, โˆ‚pre, sโ‚eโ‚, eโ‚) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚_l, โˆ‚kron_buf, eโ‚, sโ‚) + โˆ‚kron_buf .= 0 # Term 3: kron(eโ‚, ss_s1e1) + fill_kron_adjoint!(โˆ‚kron_buf, โˆ‚eโ‚_l, โˆ‚pre, ss_s1e1, eโ‚) + tmpC = Matrix(d.s_s') * โˆ‚kron_buf + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚sโ‚_l, tmpC, eโ‚, sโ‚) + # (6,7) Lโ‚ƒหข * kron(eโ‚, eโ‚ยฒ) โ€” decompress rows + โˆ‚eโ‚ยฒ_buf .= 0 + fill_kron_adjoint!(โˆ‚eโ‚ยฒ_buf, โˆ‚eโ‚_l, Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚รชโ‚ƒ[sb[6]+1:sb[7], eb[7]+1:eb[8]]), eโ‚ยฒ, eโ‚) + fill_kron_adjoint!(โˆ‚eโ‚_l, โˆ‚eโ‚_l, โˆ‚eโ‚ยฒ_buf, eโ‚, eโ‚) + + # โ”€โ”€ 3a: ฮ“โ‚ƒ disaggregation โ†’ โˆ‚ฮฃฬ‚แถปโ‚, โˆ‚ฮฃฬ‚แถปโ‚‚, โˆ‚ฮ”ฬ‚ฮผหขโ‚‚ โ”€โ”€ + โˆ‚ฮ“ = Matrix{T}(โˆ‚ฮ“โ‚ƒ_iter) + vฮฃ = vec(d.ฮฃฬ‚แถปโ‚) + + # Row 1: (1,4) kron(ฮ”ฬ‚ฮผหขโ‚‚',Ine) + โˆ‚tmp14 = _kron_vjp(โˆ‚ฮ“[gb[1]+1:gb[2], gb[4]+1:gb[5]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, 1, :), Ine)[1] + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(โˆ‚tmp14') + # (1,5) kron(vec(ฮฃฬ‚แถปโ‚)',Ine) + โˆ‚tmp15 = _kron_vjp(โˆ‚ฮ“[gb[1]+1:gb[2], gb[5]+1:gb[6]], reshape(vฮฃ, 1, :), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(vec(โˆ‚tmp15'), n, n) + # Row 3: (3,3) kron(ฮฃฬ‚แถปโ‚,Ine) + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚ฮ“[gb[3]+1:gb[4], gb[3]+1:gb[4]], Matrix(d.ฮฃฬ‚แถปโ‚), Ine)[1] + # Row 4: (4,1) kron(ฮ”ฬ‚ฮผหขโ‚‚,Ine) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(_kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[1]+1:gb[2]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Ine)[1]) + # (4,4) kron(ฮฃฬ‚แถปโ‚‚_22 + ฮ”*ฮ”', Ine) + M44 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, n+1:2n] + d.ฮ”ฬ‚ฮผหขโ‚‚ * d.ฮ”ฬ‚ฮผหขโ‚‚' + โˆ‚M44 = _kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[4]+1:gb[5]], Matrix(M44), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[n+1:2n, n+1:2n] .+= โˆ‚M44 + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= (โˆ‚M44 + โˆ‚M44') * d.ฮ”ฬ‚ฮผหขโ‚‚ + # (4,5) kron(ฮฃฬ‚แถปโ‚‚_23 + ฮ”*vฮฃ', Ine) + M45 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] + d.ฮ”ฬ‚ฮผหขโ‚‚ * vฮฃ' + โˆ‚M45 = _kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[5]+1:gb[6]], Matrix(M45), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] .+= โˆ‚M45 + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚M45 * vฮฃ + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚M45' * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + # (4,7) kron(ฮ”ฬ‚ฮผหขโ‚‚, e4_nแต‰_nแต‰ยณ) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(_kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[7]+1:gb[8]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Matrix(e4_nแต‰_nแต‰ยณ))[1]) + # Row 5: (5,1) kron(vฮฃ, Ine) + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(_kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[1]+1:gb[2]], reshape(vฮฃ, :, 1), Ine)[1], n, n) + # (5,4) kron(ฮฃฬ‚แถปโ‚‚_32 + vฮฃ*ฮ”', Ine) + M54 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] + vฮฃ * d.ฮ”ฬ‚ฮผหขโ‚‚' + โˆ‚M54 = _kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[4]+1:gb[5]], Matrix(M54), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] .+= โˆ‚M54 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚M54 * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚M54' * vฮฃ + # (5,5) kron(ฮฃฬ‚แถปโ‚‚_33 + vฮฃ*vฮฃ', Ine) + M55 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ * vฮฃ' + โˆ‚M55 = _kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[5]+1:gb[6]], Matrix(M55), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] .+= โˆ‚M55 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape((โˆ‚M55 + โˆ‚M55') * vฮฃ, n, n) + # (5,7) kron(vฮฃ, e4_nแต‰_nแต‰ยณ) + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(_kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[7]+1:gb[8]], reshape(vฮฃ, :, 1), Matrix(e4_nแต‰_nแต‰ยณ))[1], n, n) + # Row 6: (6,6) kron(ฮฃฬ‚แถปโ‚, e4_nแต‰ยฒ_nแต‰ยฒ) + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚ฮ“[gb[6]+1:gb[7], gb[6]+1:gb[7]], Matrix(d.ฮฃฬ‚แถปโ‚), Matrix(e4_nแต‰ยฒ_nแต‰ยฒ))[1] + # Row 7: (7,4) kron(ฮ”ฬ‚ฮผหขโ‚‚', e4') + โˆ‚tmp74 = _kron_vjp(โˆ‚ฮ“[gb[7]+1:gb[8], gb[4]+1:gb[5]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, 1, :), Matrix(e4_nแต‰_nแต‰ยณ'))[1] + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(โˆ‚tmp74') + # (7,5) kron(vฮฃ', e4') + โˆ‚tmp75 = _kron_vjp(โˆ‚ฮ“[gb[7]+1:gb[8], gb[5]+1:gb[6]], reshape(vฮฃ, 1, :), Matrix(e4_nแต‰_nแต‰ยณ'))[1] + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(vec(โˆ‚tmp75'), n, n) + + # โ”€โ”€ 3b: Eแดธแถป disaggregation โ”€โ”€ + โˆ‚EL = Matrix{T}(โˆ‚Eแดธแถป_iter) + # Only row block 6 is data-dependent + โˆ‚EL6 = โˆ‚EL[gb[6]+1:gb[7], :] + # Col 1: kron(ฮฃฬ‚แถปโ‚, vec_Ie) + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚EL6[:, sb[1]+1:sb[2]], Matrix(d.ฮฃฬ‚แถปโ‚), vec_Ie_col)[1] + # Col 4: kron(ฮผหขโ‚ƒฮดฮผหขโ‚', vec_Ie) + โˆ‚ฮผ_T = _kron_vjp(โˆ‚EL6[:, sb[4]+1:sb[5]], Matrix(d.ฮผหขโ‚ƒฮดฮผหขโ‚'), vec_Ie_col)[1] + โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚ = Matrix(โˆ‚ฮผ_T') # nร—n + # Col 5: kron(Cโ‚„, vec_Ie) + inner_C4 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] + d.ฮ”ฬ‚ฮผหขโ‚‚ * vฮฃ' + ss_s_M = Matrix(d.ss_s) + C4m = reshape(ss_s_M * vec(inner_C4), n, n^2) + โˆ‚C4 = _kron_vjp(โˆ‚EL6[:, sb[5]+1:sb[6]], C4m, vec_Ie_col)[1] + โˆ‚iC4 = reshape(ss_s_M' * vec(โˆ‚C4), n, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] .+= โˆ‚iC4 + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚iC4 * vฮฃ + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚iC4' * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + # Col 6: kron(Cโ‚… * Lโ‚ƒหข', vec_Ie) โ€” compress Cโ‚… cols + inner_C5 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ * vฮฃ' + C5m = reshape(Matrix(inner_C5), n, n^3) + C5m_c = C5m * Matrix(d.Lโ‚ƒหข)' + โˆ‚C5_c = _kron_vjp(โˆ‚EL6[:, sb[6]+1:sb[7]], C5m_c, vec_Ie_col)[1] + โˆ‚C5 = โˆ‚C5_c * Matrix(d.Lโ‚ƒหข) + โˆ‚iC5 = reshape(โˆ‚C5, n^2, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] .+= โˆ‚iC5 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape((โˆ‚iC5 + โˆ‚iC5') * vฮฃ, n, n) + + # โ”€โ”€ 3c: ฮผหขโ‚ƒฮดฮผหขโ‚ adjoint โ”€โ”€ + # ฮผหขโ‚ƒฮดฮผหขโ‚ = reshape((I - sโ‚ยฒ) \ vec(RHS), n, n) + โˆ‚x_ฮผ = vec(โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚) + I_m_sโ‚ยฒ = Matrix{T}(โ„’.I(n^2)) - sโ‚ยฒ + โˆ‚b_ฮผ = I_m_sโ‚ยฒ' \ โˆ‚x_ฮผ + # โˆ‚(kron(sโ‚,sโ‚)) = โˆ‚b * vec(ฮผ)' + โˆ‚sโ‚ยฒ_from_ฮผ = โˆ‚b_ฮผ * vec(d.ฮผหขโ‚ƒฮดฮผหขโ‚)' + fill_kron_adjoint!(โˆ‚sโ‚_l, โˆ‚sโ‚_l, โˆ‚sโ‚ยฒ_from_ฮผ, sโ‚, sโ‚) + + # Decompose โˆ‚RHS: RHS = Lโ‚ * sโ‚' + Lโ‚‚ * eโ‚' + โˆ‚RHS = reshape(โˆ‚b_ฮผ, n, n) + + # Reconstruct Lโ‚ and Lโ‚‚ + inner_M1 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] + vฮฃ * d.ฮ”ฬ‚ฮผหขโ‚‚' + M1 = reshape(ss_s_M * vec(inner_M1), n^2, n) + inner_M2 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ * vฮฃ' + M2 = reshape(Matrix(inner_M2), n^3, n) + M3 = โ„’.kron(Matrix(d.ฮฃฬ‚แถปโ‚), vec_Ie_col) + + Lโ‚ = ssโ‚‚ * M1 + Matrix(d.s_s_s_to_sโ‚ƒ) * M2 / 6 + + Matrix(d.s_e_e_to_sโ‚ƒ) * M3 / 2 + Matrix(d.s_v_v_to_sโ‚ƒ) * Matrix(d.ฮฃฬ‚แถปโ‚) / 2 + + M4 = โ„’.kron(reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Ine) + M5 = Matrix(e4_nแต‰_nแต‰ยณ') + M6 = โ„’.kron(reshape(vฮฃ, :, 1), Ine) + + Lโ‚‚ = seโ‚‚ * M4 + Matrix(d.e_e_e_to_sโ‚ƒ) * M5 / 6 + + Matrix(d.s_s_e_to_sโ‚ƒ) * M6 / 2 + Matrix(d.e_v_v_to_sโ‚ƒ) * Ine / 2 + + โˆ‚Lโ‚ = โˆ‚RHS * sโ‚; โˆ‚sโ‚_l .+= โˆ‚RHS' * Lโ‚ + โˆ‚Lโ‚‚ = โˆ‚RHS * eโ‚; โˆ‚eโ‚_l .+= โˆ‚RHS' * Lโ‚‚ + + # Decompose โˆ‚Lโ‚ + โˆ‚ssโ‚‚_l .+= โˆ‚Lโ‚ * M1' + โˆ‚M1_raw = ssโ‚‚' * โˆ‚Lโ‚ + โˆ‚S3f_acc[d.iหข, d.kron_s_s_s] .+= โˆ‚Lโ‚ * M2' ./ 6 + โˆ‚M2_raw = Matrix(d.s_s_s_to_sโ‚ƒ)' * โˆ‚Lโ‚ ./ 6 + โˆ‚S3f_acc[d.iหข, d.kron_s_e_e] .+= โˆ‚Lโ‚ * M3' ./ 2 + โˆ‚M3_raw = Matrix(d.s_e_e_to_sโ‚ƒ)' * โˆ‚Lโ‚ ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_s_v_v] .+= โˆ‚Lโ‚ * Matrix(d.ฮฃฬ‚แถปโ‚)' ./ 2 + โˆ‚ฮฃฬ‚แถปโ‚ .+= Matrix(d.s_v_v_to_sโ‚ƒ)' * โˆ‚Lโ‚ ./ 2 + + # Decompose โˆ‚Lโ‚‚ + โˆ‚seโ‚‚_l .+= โˆ‚Lโ‚‚ * M4' + โˆ‚M4_raw = seโ‚‚' * โˆ‚Lโ‚‚ + โˆ‚S3f_acc[d.iหข, d.kron_e_e_e] .+= โˆ‚Lโ‚‚ * M5' ./ 6 + โˆ‚S3f_acc[d.iหข, d.kron_s_s_e] .+= โˆ‚Lโ‚‚ * M6' ./ 2 + โˆ‚M6_raw = Matrix(d.s_s_e_to_sโ‚ƒ)' * โˆ‚Lโ‚‚ ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_e_v_v] .+= โˆ‚Lโ‚‚ ./ 2 + + # Decompose โˆ‚M1 โ†’ โˆ‚ฮฃฬ‚แถปโ‚‚, โˆ‚ฮฃฬ‚แถปโ‚, โˆ‚ฮ”ฬ‚ฮผหขโ‚‚ + โˆ‚iM1 = reshape(ss_s_M' * vec(โˆ‚M1_raw), n^2, n) + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] .+= โˆ‚iM1 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚iM1 * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚iM1' * vฮฃ + # Decompose โˆ‚M2 โ†’ โˆ‚ฮฃฬ‚แถปโ‚‚, โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚iM2 = reshape(โˆ‚M2_raw, n^2, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] .+= โˆ‚iM2 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape((โˆ‚iM2 + โˆ‚iM2') * vฮฃ, n, n) + # Decompose โˆ‚M3 โ†’ โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚M3_raw, Matrix(d.ฮฃฬ‚แถปโ‚), vec_Ie_col)[1] + # Decompose โˆ‚M4 โ†’ โˆ‚ฮ”ฬ‚ฮผหขโ‚‚ + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(_kron_vjp(โˆ‚M4_raw, reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Ine)[1]) + # Decompose โˆ‚M6 โ†’ โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(_kron_vjp(โˆ‚M6_raw, reshape(vฮฃ, :, 1), Ine)[1], n, n) + + # โ”€โ”€ 4: Scatter local cotangents to global accumulators โ”€โ”€ + โˆ‚๐’โ‚_acc[d.iหข, d.dependencies_in_states_idx] .+= โˆ‚sโ‚_l + โˆ‚๐’โ‚_acc[d.iหข, nโ‚‹+1:size(โˆ‚๐’โ‚_acc, 2)] .+= โˆ‚eโ‚_l + โˆ‚S2f_acc[d.iหข, d.kron_s_s] .+= โˆ‚ssโ‚‚_l + โˆ‚S2f_acc[d.iหข, kron_e_e] .+= โˆ‚eeโ‚‚_l + โˆ‚S2f_acc[d.iหข, d.kron_s_e] .+= โˆ‚seโ‚‚_l + โˆ‚S2f_acc[d.iหข, kron_v_v] .+= โˆ‚vvโ‚‚_l + โˆ‚ฮฃสธโ‚_acc[d.iหข, d.iหข] .+= โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚ฮฃแถปโ‚‚_acc[d.dependencies_extended_idx, d.dependencies_extended_idx] .+= โˆ‚ฮฃฬ‚แถปโ‚‚ + โˆ‚ฮ”ฮผหขโ‚‚_acc[d.dependencies_in_states_idx] .+= โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l + end + + # โ”€โ”€ Sub-rrule pullback chain โ”€โ”€ + + # Sโ‚ƒ_full = Sโ‚ƒ * ๐”โ‚ƒ โ†’ โˆ‚Sโ‚ƒ = โˆ‚Sโ‚ƒ_full * ๐”โ‚ƒ' + โˆ‚๐’โ‚ƒ_compressed = โˆ‚S3f_acc * ๐”โ‚ƒ' + + # Third-order solution pullback: returns (NoTangent, โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚ƒ, โˆ‚๐‘บโ‚, โˆ‚๐’โ‚‚, NT, NT, NT) + so3_grad = so3_pb((โˆ‚๐’โ‚ƒ_compressed, NoTangent())) + if !(so3_grad[2] isa AbstractZero); โˆ‚โˆ‡โ‚_acc .+= so3_grad[2]; end + if !(so3_grad[3] isa AbstractZero); โˆ‚โˆ‡โ‚‚_acc .+= so3_grad[3]; end + if !(so3_grad[4] isa AbstractZero); โˆ‚โˆ‡โ‚ƒ_acc .+= so3_grad[4]; end + if !(so3_grad[5] isa AbstractZero); โˆ‚๐’โ‚_acc .+= so3_grad[5]; end + # so3_grad[6] is now compressed โˆ‚๐’โ‚‚_raw โ€” kept separate + + # Third-order derivatives pullback: returns (NoTangent, โˆ‚params, โˆ‚SS, NT, NT) + โˆ‡โ‚ƒ_grad = โˆ‡โ‚ƒ_pb(โˆ‚โˆ‡โ‚ƒ_acc) + โˆ‚params_โˆ‡โ‚ƒ = โˆ‡โ‚ƒ_grad[2] isa AbstractZero ? zeros(T, np) : โˆ‡โ‚ƒ_grad[2] + if !(โˆ‡โ‚ƒ_grad[3] isa AbstractZero); โˆ‚SS_acc .+= โˆ‡โ‚ƒ_grad[3]; end + + # Convert full-space โˆ‚S2f_acc to compressed and add compressed so3 gradient + โˆ‚S2_raw_acc = โˆ‚S2f_acc * ๐”โ‚‚' + if !(so3_grad[6] isa AbstractZero); โˆ‚S2_raw_acc .+= so3_grad[6]; end + + # Second-order moments pullback: cotangent tuple for 15-element output + # (ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr, ลลโ‚‚, ลyโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, slvd) + โˆ‚som2 = ( + NoTangent(), # โˆ‚ฮฃสธโ‚‚ (not used by third-order) + โˆ‚ฮฃแถปโ‚‚_acc, # โˆ‚ฮฃแถปโ‚‚ + โˆ‚ฮผสธโ‚‚_in isa AbstractZero ? NoTangent() : โˆ‚ฮผสธโ‚‚_in, # โˆ‚ฮผสธโ‚‚ + โˆ‚ฮ”ฮผหขโ‚‚_acc, # โˆ‚ฮ”ฮผหขโ‚‚ + NoTangent(), # โˆ‚autocorr (not used) + NoTangent(), # โˆ‚ล_to_ลโ‚‚ (not used) + NoTangent(), # โˆ‚ล_to_yโ‚‚ (not used) + โˆ‚ฮฃสธโ‚_acc, # โˆ‚ฮฃสธโ‚ + NoTangent(), # โˆ‚ฮฃแถปโ‚ + โˆ‚SS_acc, # โˆ‚SS_and_pars + โˆ‚๐’โ‚_acc, # โˆ‚๐’โ‚ + โˆ‚โˆ‡โ‚_acc, # โˆ‚โˆ‡โ‚ + โˆ‚S2_raw_acc, # โˆ‚๐’โ‚‚ (compressed) + โˆ‚โˆ‡โ‚‚_acc, # โˆ‚โˆ‡โ‚‚ + NoTangent(), # โˆ‚slvd + ) + + som2_grad = som2_pb(โˆ‚som2) + โˆ‚params_som2 = som2_grad[2] isa AbstractZero ? zeros(T, np) : som2_grad[2] + + โˆ‚parameters_total = โˆ‚params_som2 .+ โˆ‚params_โˆ‡โ‚ƒ + + return NoTangent(), โˆ‚parameters_total, NoTangent(), NoTangent() + end + + return result, calculate_third_order_moments_pullback +end + +# โ”€โ”€ calculate_third_order_moments_with_autocorrelation rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +function rrule(::typeof(calculate_third_order_moments_with_autocorrelation), + parameters::Vector{T}, + observables::Union{Symbol_input,String_input}, + ๐“‚::โ„ณ; + autocorrelation_periods::U = 1:5, + covariance::Union{Symbol_input,String_input} = Symbol[], + opts::CalculationOptions = merge_calculation_options()) where {U, T <: Real} + + # โ”€โ”€ Non-differentiable constants โ”€โ”€ + ensure_moments_constants!(๐“‚.constants) + so = ๐“‚.constants.second_order + to = ๐“‚.constants.third_order + T_pm = ๐“‚.constants.post_model_macro + np = length(parameters) + nแต‰ = T_pm.nExo + n_ac = length(autocorrelation_periods) + + zero_5() = (zeros(T,0,0), zeros(T,0), zeros(T,0,0), zeros(T,0), false) + zero_pb(_) = (NoTangent(), zeros(T, np), NoTangent(), NoTangent()) + + # โ”€โ”€ Step 1: Second-order moments with covariance โ”€โ”€ + som2_out, som2_pb = rrule(calculate_second_order_moments_with_covariance, parameters, ๐“‚; opts = opts) + ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp_2, ล_to_ลโ‚‚, ล_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_raw, โˆ‡โ‚‚, solved = som2_out + + if !solved; return zero_5(), zero_pb; end + + # Expand compressed ๐’โ‚‚_raw to full for moments computation + ๐”โ‚‚ = ๐“‚.constants.second_order.๐”โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐”โ‚‚)::SparseMatrixCSC{T, Int} + + # โ”€โ”€ Step 2: Third-order derivatives โ”€โ”€ + โˆ‡โ‚ƒ, โˆ‡โ‚ƒ_pb = rrule(calculate_third_order_derivatives, parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces) + + # โ”€โ”€ Step 3: Third-order solution (pass compressed ๐’โ‚‚_raw) โ”€โ”€ + so3_out, so3_pb = rrule(calculate_third_order_solution, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚_raw, + ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, + parameter_values = parameters) + ๐’โ‚ƒ, solved3 = so3_out + + update_perturbation_counter!(๐“‚.counters, solved3, order = 3) + + if !solved3; return zero_5(), zero_pb; end + + # โ”€โ”€ Step 4: Decompress Sโ‚ƒ โ”€โ”€ + ๐”โ‚ƒ = ๐“‚.constants.third_order.๐”โ‚ƒ + ๐’โ‚ƒ_full = ๐’โ‚ƒ * ๐”โ‚ƒ + + ๐’โ‚ƒ_full = sparse(๐’โ‚ƒ_full) + + # โ”€โ”€ Step 5: Determine iteration groups โ”€โ”€ + orders = determine_efficient_order(๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ_full, ๐“‚.constants, observables, + covariance = covariance, tol = opts.tol.third_order.dependencies_tol) + + kron_e_e = so.kron_e_e + kron_v_v = so.kron_v_v + kron_e_v = to.kron_e_v + e_in_sโบ = so.e_in_sโบ + v_in_sโบ = so.v_in_sโบ + vec_Iโ‚‘ = so.vec_Iโ‚‘ + e4_nแต‰ยฒ_nแต‰ยฒ = so.e4_nแต‰ยฒ_nแต‰ยฒ + e4_nแต‰_nแต‰ยณ = so.e4_nแต‰_nแต‰ยณ + e4_minus_vecIโ‚‘_outer = so.e4_minus_vecIโ‚‘_outer + e6_nแต‰ยณ_nแต‰ยณ = to.e6_nแต‰ยณ_nแต‰ยณ + + ฮฃสธโ‚ƒ = zeros(T, size(ฮฃสธโ‚‚)) + autocorr = zeros(T, size(ฮฃสธโ‚‚, 1), n_ac) + solved_lyapunov = true + + # Per-iteration storage for pullback + n_iters = length(orders) + iter_data = Vector{Any}(undef, n_iters) + + for (iter_idx, ords) in enumerate(orders) + variance_observable, dependencies_all_vars = ords + + sort!(variance_observable) + sort!(dependencies_all_vars) + + dependencies = intersect(T_pm.past_not_future_and_mixed, dependencies_all_vars) + + obs_in_y = indexin(variance_observable, T_pm.var) + + dependencies_in_states_idx = indexin(dependencies, T_pm.past_not_future_and_mixed) + + dependencies_in_var_idx = Int.(indexin(dependencies, T_pm.var)) + + nหข = length(dependencies) + + iหข = dependencies_in_var_idx + + ฮฃฬ‚แถปโ‚ = ฮฃสธโ‚[iหข, iหข] + + dependencies_extended_idx = vcat(dependencies_in_states_idx, + dependencies_in_states_idx .+ T_pm.nPast_not_future_and_mixed, + findall(โ„’.kron(T_pm.past_not_future_and_mixed .โˆˆ (intersect(T_pm.past_not_future_and_mixed,dependencies),), + T_pm.past_not_future_and_mixed .โˆˆ (intersect(T_pm.past_not_future_and_mixed,dependencies),))) .+ 2*T_pm.nPast_not_future_and_mixed) + + ฮฃฬ‚แถปโ‚‚ = ฮฃแถปโ‚‚[dependencies_extended_idx, dependencies_extended_idx] + + ฮ”ฬ‚ฮผหขโ‚‚ = ฮ”ฮผหขโ‚‚[dependencies_in_states_idx] + + s_in_sโบ = BitVector(vcat(T_pm.past_not_future_and_mixed .โˆˆ (dependencies,), zeros(Bool, nแต‰ + 1))) + + substate_indices = ensure_moments_substate_indices!(๐“‚, nหข) + I_plus_s_s = substate_indices.I_plus_s_s + e_es = substate_indices.e_es + e_ss = substate_indices.e_ss + ss_s = substate_indices.ss_s + s_s = substate_indices.s_s + + # first order slices + s_to_yโ‚ = ๐’โ‚[obs_in_y,:][:,dependencies_in_states_idx] + e_to_yโ‚ = ๐’โ‚[obs_in_y,:][:, (T_pm.nPast_not_future_and_mixed + 1):end] + + s_to_sโ‚ = ๐’โ‚[iหข, dependencies_in_states_idx] + e_to_sโ‚ = ๐’โ‚[iหข, (T_pm.nPast_not_future_and_mixed + 1):end] + + # second order slices + dep_kron = ensure_moments_dependency_kron_indices!(๐“‚, dependencies, s_in_sโบ) + kron_s_s = dep_kron.kron_s_s + kron_s_e = dep_kron.kron_s_e + + s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_y,:][:, kron_s_s] + e_e_to_yโ‚‚ = ๐’โ‚‚[obs_in_y,:][:, kron_e_e] + s_e_to_yโ‚‚ = ๐’โ‚‚[obs_in_y,:][:, kron_s_e] + + s_s_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_s_s] |> collect + e_e_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_e_e] + v_v_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_v_v] |> collect + s_e_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_s_e] + + s_to_sโ‚_by_s_to_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + e_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) + s_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) + + # third order slices + kron_s_v = dep_kron.kron_s_v + + kron_s_s_s = โ„’.kron(kron_s_s, s_in_sโบ) + kron_s_s_e = โ„’.kron(kron_s_s, e_in_sโบ) + kron_s_e_e = โ„’.kron(kron_s_e, e_in_sโบ) + kron_e_e_e = โ„’.kron(kron_e_e, e_in_sโบ) + kron_s_v_v = โ„’.kron(kron_s_v, v_in_sโบ) + kron_e_v_v = โ„’.kron(kron_e_v, v_in_sโบ) + + s_s_s_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_s_s] + s_s_e_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_s_e] + s_e_e_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_e_e] + e_e_e_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_e_e_e] + s_v_v_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_s_v_v] + e_v_v_to_yโ‚ƒ = ๐’โ‚ƒ_full[obs_in_y,:][:, kron_e_v_v] + + s_s_s_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_s_s] + s_s_e_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_s_e] + s_e_e_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_e_e] + e_e_e_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_e_e_e] + s_v_v_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_s_v_v] + e_v_v_to_sโ‚ƒ = ๐’โ‚ƒ_full[iหข, kron_e_v_v] + + # Set up compression matrices + Dโ‚‚หข = substate_indices.Dโ‚‚หข + Lโ‚‚หข = substate_indices.Lโ‚‚หข + Dโ‚ƒหข = substate_indices.Dโ‚ƒหข + Lโ‚ƒหข = substate_indices.Lโ‚ƒหข + nโ‚‚หข = size(Dโ‚‚หข, 2) + nโ‚ƒหข = size(Dโ‚ƒหข, 2) + s_to_sโ‚_by_s_to_sโ‚_c = Lโ‚‚หข * s_to_sโ‚_by_s_to_sโ‚ * Dโ‚‚หข + + # Set up pruned state transition sub-blocks (compressed) + N_upper = 2 * nหข + nโ‚‚หข + N_lower = nหข + nหข^2 + nโ‚ƒหข + + A_UU = [s_to_sโ‚ spzeros(nหข, nหข + nโ‚‚หข) + spzeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 * Dโ‚‚หข + spzeros(nโ‚‚หข, 2 * nหข) s_to_sโ‚_by_s_to_sโ‚_c] + + A_LU = [s_v_v_to_sโ‚ƒ / 2 spzeros(nหข, nหข + nโ‚‚หข) + โ„’.kron(s_to_sโ‚,v_v_to_sโ‚‚ / 2) spzeros(nหข^2, nหข + nโ‚‚หข) + spzeros(nโ‚ƒหข, 2 * nหข + nโ‚‚หข)] + + A_LL = [s_to_sโ‚ s_s_to_sโ‚‚ s_s_s_to_sโ‚ƒ / 6 * Dโ‚ƒหข + spzeros(nหข^2, nหข) s_to_sโ‚_by_s_to_sโ‚ โ„’.kron(s_to_sโ‚,s_s_to_sโ‚‚ / 2) * Dโ‚ƒหข + spzeros(nโ‚ƒหข, nหข + nหข^2) Lโ‚ƒหข * โ„’.kron(s_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * Dโ‚ƒหข] + + รช_to_ลโ‚ƒ = [ e_to_sโ‚ zeros(nหข,nแต‰^2 + 2*nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + zeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ zeros(nหข,nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + zeros(nโ‚‚หข,nแต‰) Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ Lโ‚‚หข * I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚ zeros(nโ‚‚หข, nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + e_v_v_to_sโ‚ƒ / 2 zeros(nหข,nแต‰^2 + nแต‰ * nหข) s_e_to_sโ‚‚ s_s_e_to_sโ‚ƒ / 2 s_e_e_to_sโ‚ƒ / 2 e_e_e_to_sโ‚ƒ / 6 + โ„’.kron(e_to_sโ‚, v_v_to_sโ‚‚ / 2) zeros(nหข^2, nแต‰^2 + nแต‰ * nหข) s_s * s_to_sโ‚_by_e_to_sโ‚ โ„’.kron(s_to_sโ‚, s_e_to_sโ‚‚) + s_s * โ„’.kron(s_s_to_sโ‚‚ / 2, e_to_sโ‚) โ„’.kron(s_to_sโ‚, e_e_to_sโ‚‚ / 2) + s_s * โ„’.kron(s_e_to_sโ‚‚, e_to_sโ‚) โ„’.kron(e_to_sโ‚, e_e_to_sโ‚‚ / 2) + zeros(nโ‚ƒหข, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_s_to_sโ‚,e_to_sโ‚) + โ„’.kron(s_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * e_ss) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_e_to_sโ‚,e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_e_to_sโ‚) * e_es + โ„’.kron(e_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) * e_es) Lโ‚ƒหข * โ„’.kron(e_to_sโ‚,e_to_sโ‚_by_e_to_sโ‚)] + + ล_to_yโ‚ƒ = [s_to_yโ‚ + s_v_v_to_yโ‚ƒ / 2 s_to_yโ‚ s_s_to_yโ‚‚ / 2 * Dโ‚‚หข s_to_yโ‚ s_s_to_yโ‚‚ s_s_s_to_yโ‚ƒ / 6 * Dโ‚ƒหข] + + รช_to_yโ‚ƒ = [e_to_yโ‚ + e_v_v_to_yโ‚ƒ / 2 e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚ s_e_to_yโ‚‚ s_s_e_to_yโ‚ƒ / 2 s_e_e_to_yโ‚ƒ / 2 e_e_e_to_yโ‚ƒ / 6] + + ฮผหขโ‚ƒฮดฮผหขโ‚ = reshape((โ„’.I(size(s_to_sโ‚_by_s_to_sโ‚, 1)) - s_to_sโ‚_by_s_to_sโ‚) \ vec( + (s_s_to_sโ‚‚ * reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, nหข + 1:2*nหข] + vec(ฮฃฬ‚แถปโ‚) * ฮ”ฬ‚ฮผหขโ‚‚'),nหข^2, nหข) + + s_s_s_to_sโ‚ƒ * reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end , 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข^3, nหข) / 6 + + s_e_e_to_sโ‚ƒ * โ„’.kron(ฮฃฬ‚แถปโ‚, vec_Iโ‚‘) / 2 + + s_v_v_to_sโ‚ƒ * ฮฃฬ‚แถปโ‚ / 2) * s_to_sโ‚' + + (s_e_to_sโ‚‚ * โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚,โ„’.I(nแต‰)) + + e_e_e_to_sโ‚ƒ * e4_nแต‰_nแต‰ยณ' / 6 + + s_s_e_to_sโ‚ƒ * โ„’.kron(vec(ฮฃฬ‚แถปโ‚), โ„’.I(nแต‰)) / 2 + + e_v_v_to_sโ‚ƒ * โ„’.I(nแต‰) / 2) * e_to_sโ‚' + ), nหข, nหข) + + ฮ“โ‚ƒ = [ โ„’.I(nแต‰) spzeros(nแต‰, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚', โ„’.I(nแต‰)) โ„’.kron(vec(ฮฃฬ‚แถปโ‚)', โ„’.I(nแต‰)) spzeros(nแต‰, nหข * nแต‰^2) e4_nแต‰_nแต‰ยณ + spzeros(nแต‰^2, nแต‰) e4_minus_vecIโ‚‘_outer spzeros(nแต‰^2, 2*nหข*nแต‰ + nหข^2*nแต‰ + nหข*nแต‰^2 + nแต‰^3) + spzeros(nหข * nแต‰, nแต‰ + nแต‰^2) โ„’.kron(ฮฃฬ‚แถปโ‚, โ„’.I(nแต‰)) spzeros(nหข * nแต‰, nหข*nแต‰ + nหข^2*nแต‰ + nหข*nแต‰^2 + nแต‰^3) + โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚,โ„’.I(nแต‰)) spzeros(nแต‰ * nหข, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,nหข + 1:2*nหข] + ฮ”ฬ‚ฮผหขโ‚‚ * ฮ”ฬ‚ฮผหขโ‚‚',โ„’.I(nแต‰)) โ„’.kron(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)',โ„’.I(nแต‰)) spzeros(nแต‰ * nหข, nหข * nแต‰^2) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚, e4_nแต‰_nแต‰ยณ) + โ„’.kron(vec(ฮฃฬ‚แถปโ‚), โ„’.I(nแต‰)) spzeros(nแต‰ * nหข^2, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, nหข + 1:2*nหข] + vec(ฮฃฬ‚แถปโ‚) * ฮ”ฬ‚ฮผหขโ‚‚', โ„’.I(nแต‰)) โ„’.kron(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', โ„’.I(nแต‰)) spzeros(nแต‰ * nหข^2, nหข * nแต‰^2) โ„’.kron(vec(ฮฃฬ‚แถปโ‚), e4_nแต‰_nแต‰ยณ) + spzeros(nหข*nแต‰^2, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข + nหข^2*nแต‰) โ„’.kron(ฮฃฬ‚แถปโ‚, e4_nแต‰ยฒ_nแต‰ยฒ) spzeros(nหข*nแต‰^2,nแต‰^3) + e4_nแต‰_nแต‰ยณ' spzeros(nแต‰^3, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚', e4_nแต‰_nแต‰ยณ') โ„’.kron(vec(ฮฃฬ‚แถปโ‚)', e4_nแต‰_nแต‰ยณ') spzeros(nแต‰^3, nหข*nแต‰^2) e6_nแต‰ยณ_nแต‰ยณ] + + + Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข) + โ„’.kron(ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) zeros(nหข*nแต‰^2, nหข + nโ‚‚หข) โ„’.kron(ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3) * Lโ‚ƒหข', vec_Iโ‚‘) + spzeros(nแต‰^3, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข)] + + droptol!(A_UU, eps()) + droptol!(A_LU, eps()) + droptol!(A_LL, eps()) + droptol!(รช_to_ลโ‚ƒ, eps()) + droptol!(Eแดธแถป, eps()) + droptol!(ฮ“โ‚ƒ, eps()) + + # โ”€โ”€ Standard Lyapunov solve (compressed) โ”€โ”€ + N_total = N_upper + N_lower + ล_to_ลโ‚ƒ = [A_UU spzeros(N_upper, N_lower); A_LU A_LL] + A_cross = Matrix{Float64}(รช_to_ลโ‚ƒ * Eแดธแถป) * ล_to_ลโ‚ƒ' + C_dense = Matrix{Float64}(รช_to_ลโ‚ƒ * ฮ“โ‚ƒ * รช_to_ลโ‚ƒ') + A_cross + A_cross' + + lyap_ws_3rd = Lyapunov_workspace(N_total) + lyap_out, lyap_pb_iter = rrule(solve_lyapunov_equation, + ล_to_ลโ‚ƒ, C_dense, lyap_ws_3rd, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.third_order.ad.lyapunov, + verbose = opts.verbose) + ฮฃแถปโ‚ƒ = lyap_out[1] + info = lyap_out[2] + + if !info + return zero_5(), zero_pb + end + + solved_lyapunov = solved_lyapunov && info + + ฮฃสธโ‚ƒtmp = ล_to_yโ‚ƒ * ฮฃแถปโ‚ƒ * ล_to_yโ‚ƒ' + รช_to_yโ‚ƒ * ฮ“โ‚ƒ * รช_to_yโ‚ƒ' + รช_to_yโ‚ƒ * Eแดธแถป * ล_to_yโ‚ƒ' + ล_to_yโ‚ƒ * Eแดธแถป' * รช_to_yโ‚ƒ' + + for obs in variance_observable + ฮฃสธโ‚ƒ[indexin([obs], T_pm.var), indexin(variance_observable, T_pm.var)] = ฮฃสธโ‚ƒtmp[indexin([obs], variance_observable), :] + end + + # โ”€โ”€ Autocorrelation forward pass โ”€โ”€ + Eแดธแถป_orig = Eแดธแถป # save original for pullback + + autocorr_tmp_ac = ล_to_ลโ‚ƒ * Eแดธแถป' * รช_to_yโ‚ƒ' + รช_to_ลโ‚ƒ * ฮ“โ‚ƒ * รช_to_yโ‚ƒ' + + s_to_sโ‚โฑ = Matrix{T}(โ„’.I(nหข)) + ล_to_ลโ‚ƒโฑ = Matrix{T}(โ„’.I(size(ฮฃแถปโ‚ƒ, 1))) + ฮฃแถปโ‚ƒโฑ = copy(Matrix{T}(ฮฃแถปโ‚ƒ)) + + norm_diag = max.(โ„’.diag(ฮฃสธโ‚ƒtmp), eps(Float64)) + + per_period = Vector{Any}(undef, n_ac) + Eแดธแถป_cur = Eแดธแถป_orig # tracks current Eแดธแถป for step 1 + + for (pi, i) in enumerate(autocorrelation_periods) + # Snapshot before step 1 + ฮฃแถปโ‚ƒโฑ_prev = copy(ฮฃแถปโ‚ƒโฑ) + Eแดธแถป_used = Eแดธแถป_cur # Eแดธแถป used in step 1 + + # Step 1: ฮฃแถปโ‚ƒโฑ update + ฮฃแถปโ‚ƒโฑ .= Matrix(ล_to_ลโ‚ƒ) * ฮฃแถปโ‚ƒโฑ + Matrix(รช_to_ลโ‚ƒ) * Matrix(Eแดธแถป_cur) + + # Step 2: s_to_sโ‚โฑ update (snapshot before) + s_to_sโ‚โฑ_prev = copy(s_to_sโ‚โฑ) + s_to_sโ‚โฑ = s_to_sโ‚โฑ * Matrix{T}(s_to_sโ‚) + + # Step 3: rebuild Eแดธแถป with s_to_sโ‚โฑ prefix + Eแดธแถปโฑ = [ spzeros(T, nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข) + โ„’.kron(s_to_sโ‚โฑ * ฮฃฬ‚แถปโ‚, vec_Iโ‚‘) zeros(T, nหข*nแต‰^2, nหข + nโ‚‚หข) โ„’.kron(s_to_sโ‚โฑ * ฮผหขโ‚ƒฮดฮผหขโ‚', vec_Iโ‚‘) โ„’.kron(s_to_sโ‚โฑ * reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข, 2*nหข + 1:end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(s_to_sโ‚โฑ * reshape(ฮฃฬ‚แถปโ‚‚[2*nหข + 1:end, 2*nหข + 1:end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3) * Lโ‚ƒหข', vec_Iโ‚‘) + spzeros(T, nแต‰^3, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข)] + Eแดธแถป_cur = Eแดธแถปโฑ + + # Step 4: compute autocorrelation + ล_to_ลโ‚ƒโฑ_snap = copy(ล_to_ลโ‚ƒโฑ) # snapshot before step 5 + num_mat = Matrix(ล_to_yโ‚ƒ) * ฮฃแถปโ‚ƒโฑ * Matrix(ล_to_yโ‚ƒ)' + Matrix(ล_to_yโ‚ƒ) * ล_to_ลโ‚ƒโฑ * Matrix(autocorr_tmp_ac) + Matrix(รช_to_yโ‚ƒ) * Matrix(Eแดธแถปโฑ) * Matrix(ล_to_yโ‚ƒ)' + num_diag_i = โ„’.diag(num_mat) + ac_val = num_diag_i ./ norm_diag + diag_ฮฃ = โ„’.diag(ฮฃสธโ‚ƒtmp) + zero_mask_i = diag_ฮฃ .< opts.tol.third_order.ad.lyapunov.acceptance_tol + ac_val[zero_mask_i] .= 0 + + for obs in variance_observable + autocorr[indexin([obs], T_pm.var), i] .= ac_val[indexin([obs], variance_observable)] + end + + per_period[pi] = ( + ฮฃแถปโ‚ƒโฑ_prev = ฮฃแถปโ‚ƒโฑ_prev, + Eแดธแถป_used = Eแดธแถป_used, + s_to_sโ‚โฑ = copy(s_to_sโ‚โฑ), # after step 2 + s_to_sโ‚โฑ_prev = s_to_sโ‚โฑ_prev, + Eแดธแถปโฑ = Eแดธแถปโฑ, # after step 3 + ล_to_ลโ‚ƒโฑ = ล_to_ลโ‚ƒโฑ_snap, # before step 5 + ฮฃแถปโ‚ƒโฑ = copy(ฮฃแถปโ‚ƒโฑ), # after step 1 + num_diag = num_diag_i, + zero_mask = zero_mask_i, + period_index = i, + ) + + # Step 5: ล_to_ลโ‚ƒโฑ update + ล_to_ลโ‚ƒโฑ = ล_to_ลโ‚ƒโฑ * Matrix{T}(ล_to_ลโ‚ƒ) + end + + # Store per-iteration data for pullback + iter_data[iter_idx] = ( + variance_observable = variance_observable, + obs_in_y = obs_in_y, + iหข = iหข, + nหข = nหข, + dependencies_in_states_idx = dependencies_in_states_idx, + dependencies_extended_idx = dependencies_extended_idx, + ฮฃฬ‚แถปโ‚ = ฮฃฬ‚แถปโ‚, + ฮฃฬ‚แถปโ‚‚ = ฮฃฬ‚แถปโ‚‚, + ฮ”ฬ‚ฮผหขโ‚‚ = ฮ”ฬ‚ฮผหขโ‚‚, + s_in_sโบ = s_in_sโบ, + s_to_yโ‚ = s_to_yโ‚, + e_to_yโ‚ = e_to_yโ‚, + s_to_sโ‚ = s_to_sโ‚, + e_to_sโ‚ = e_to_sโ‚, + kron_s_s = kron_s_s, + kron_s_e = kron_s_e, + kron_s_v = kron_s_v, + kron_s_s_s = kron_s_s_s, + kron_s_s_e = kron_s_s_e, + kron_s_e_e = kron_s_e_e, + kron_e_e_e = kron_e_e_e, + kron_s_v_v = kron_s_v_v, + kron_e_v_v = kron_e_v_v, + s_s_to_yโ‚‚ = s_s_to_yโ‚‚, + e_e_to_yโ‚‚ = e_e_to_yโ‚‚, + s_e_to_yโ‚‚ = s_e_to_yโ‚‚, + s_s_to_sโ‚‚ = s_s_to_sโ‚‚, + e_e_to_sโ‚‚ = e_e_to_sโ‚‚, + v_v_to_sโ‚‚ = v_v_to_sโ‚‚, + s_e_to_sโ‚‚ = s_e_to_sโ‚‚, + s_to_sโ‚_by_s_to_sโ‚ = s_to_sโ‚_by_s_to_sโ‚, + e_to_sโ‚_by_e_to_sโ‚ = e_to_sโ‚_by_e_to_sโ‚, + s_to_sโ‚_by_e_to_sโ‚ = s_to_sโ‚_by_e_to_sโ‚, + s_s_s_to_yโ‚ƒ = s_s_s_to_yโ‚ƒ, + s_s_e_to_yโ‚ƒ = s_s_e_to_yโ‚ƒ, + s_e_e_to_yโ‚ƒ = s_e_e_to_yโ‚ƒ, + e_e_e_to_yโ‚ƒ = e_e_e_to_yโ‚ƒ, + s_v_v_to_yโ‚ƒ = s_v_v_to_yโ‚ƒ, + e_v_v_to_yโ‚ƒ = e_v_v_to_yโ‚ƒ, + s_s_s_to_sโ‚ƒ = s_s_s_to_sโ‚ƒ, + s_s_e_to_sโ‚ƒ = s_s_e_to_sโ‚ƒ, + s_e_e_to_sโ‚ƒ = s_e_e_to_sโ‚ƒ, + e_e_e_to_sโ‚ƒ = e_e_e_to_sโ‚ƒ, + s_v_v_to_sโ‚ƒ = s_v_v_to_sโ‚ƒ, + e_v_v_to_sโ‚ƒ = e_v_v_to_sโ‚ƒ, + ล_to_ลโ‚ƒ = ล_to_ลโ‚ƒ, + รช_to_ลโ‚ƒ = รช_to_ลโ‚ƒ, + ล_to_yโ‚ƒ = ล_to_yโ‚ƒ, + รช_to_yโ‚ƒ = รช_to_yโ‚ƒ, + ฮ“โ‚ƒ = ฮ“โ‚ƒ, + Eแดธแถป = Eแดธแถป_orig, + N_upper = N_upper, + N_lower = N_lower, + lyap_pb = lyap_pb_iter, + Dโ‚‚หข = Dโ‚‚หข, + Lโ‚‚หข = Lโ‚‚หข, + Dโ‚ƒหข = Dโ‚ƒหข, + Lโ‚ƒหข = Lโ‚ƒหข, + nโ‚‚หข = nโ‚‚หข, + nโ‚ƒหข = nโ‚ƒหข, + s_to_sโ‚_by_s_to_sโ‚_c = s_to_sโ‚_by_s_to_sโ‚_c, + ฮฃแถปโ‚ƒ = ฮฃแถปโ‚ƒ, + ฮฃสธโ‚ƒtmp = ฮฃสธโ‚ƒtmp, + ฮผหขโ‚ƒฮดฮผหขโ‚ = ฮผหขโ‚ƒฮดฮผหขโ‚, + I_plus_s_s = I_plus_s_s, + ss_s = ss_s, + s_s = s_s, + e_es = e_es, + e_ss = e_ss, + # Autocorrelation-specific + autocorr_tmp_ac = autocorr_tmp_ac, + norm_diag = norm_diag, + per_period = per_period, + ) + end + + result = (ฮฃสธโ‚ƒ, ฮผสธโ‚‚, autocorr, SS_and_pars, solved && solved3 && solved_lyapunov) + + # โ”€โ”€ Pullback โ”€โ”€ + function calculate_third_order_moments_with_autocorrelation_pullback(โˆ‚out) + โˆ‚ฮฃสธโ‚ƒ_in, โˆ‚ฮผสธโ‚‚_in, โˆ‚autocorr_in, โˆ‚SS_in, _ = โˆ‚out + + โˆ‚ฮฃสธโ‚ƒ_in = unthunk(โˆ‚ฮฃสธโ‚ƒ_in) + โˆ‚ฮผสธโ‚‚_in = unthunk(โˆ‚ฮผสธโ‚‚_in) + โˆ‚autocorr_in = unthunk(โˆ‚autocorr_in) + โˆ‚SS_in = unthunk(โˆ‚SS_in) + + nโ‚‹ = T_pm.nPast_not_future_and_mixed + + # Accumulators for cotangents flowing to sub-rrule inputs + โˆ‚ฮฃสธโ‚_acc = zeros(T, size(ฮฃสธโ‚)) + โˆ‚ฮฃแถปโ‚‚_acc = zeros(T, size(ฮฃแถปโ‚‚)) + โˆ‚ฮ”ฮผหขโ‚‚_acc = zeros(T, length(ฮ”ฮผหขโ‚‚)) + โˆ‚๐’โ‚_acc = zeros(T, size(๐’โ‚)) + โˆ‚S2f_acc = zeros(T, size(๐’โ‚‚)) + โˆ‚S3f_acc = zeros(T, size(๐’โ‚ƒ_full)) + โˆ‚SS_acc = zeros(T, length(SS_and_pars)) + โˆ‚โˆ‡โ‚_acc = zeros(T, size(โˆ‡โ‚)) + โˆ‚โˆ‡โ‚‚_acc = zeros(T, size(โˆ‡โ‚‚)) + โˆ‚โˆ‡โ‚ƒ_acc = zeros(T, size(โˆ‡โ‚ƒ)) + + if !(โˆ‚SS_in isa AbstractZero); โˆ‚SS_acc .+= โˆ‚SS_in; end + + # โ”€โ”€โ”€โ”€ Reverse loop over iterations โ”€โ”€โ”€โ”€ + for iter_idx in n_iters:-1:1 + d = iter_data[iter_idx] + nหข_i = d.nหข + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # Stage 0: Autocorrelation reverse loop + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + nObs_iter = length(d.variance_observable) + + # Initialize cotangents that accumulate through autocorrelation loop + โˆ‚ล_to_yโ‚ƒ_ac = zeros(T, size(d.ล_to_yโ‚ƒ)) + โˆ‚รช_to_yโ‚ƒ_ac = zeros(T, size(d.รช_to_yโ‚ƒ)) + โˆ‚ฮฃแถปโ‚ƒโฑ_co = zeros(T, size(d.ฮฃแถปโ‚ƒ)) # cotangent for ฮฃแถปโ‚ƒโฑ state + โˆ‚ล_to_ลโ‚ƒ_ac = zeros(T, size(d.ล_to_ลโ‚ƒ)) + โˆ‚รช_to_ลโ‚ƒ_ac = zeros(T, size(d.รช_to_ลโ‚ƒ)) + โˆ‚Eแดธแถป_ac = zeros(T, size(d.Eแดธแถป)) # cotangent for original Eแดธแถป + โˆ‚ฮ“โ‚ƒ_ac = zeros(T, size(d.ฮ“โ‚ƒ)) + โˆ‚autocorr_tmp_co = zeros(T, size(d.autocorr_tmp_ac)) + โˆ‚sโ‚_ac = zeros(T, nหข_i, nหข_i) # cotangent for s_to_sโ‚ + โˆ‚ฮฃสธโ‚ƒtmp_ac = zeros(T, nObs_iter, nObs_iter) # cotangent from norm_diag + โˆ‚ล_to_ลโ‚ƒโฑ_co = zeros(T, size(d.ฮฃแถปโ‚ƒ)) # cotangent for ล_to_ลโ‚ƒโฑ state + โˆ‚s_to_sโ‚โฑ_co = zeros(T, nหข_i, nหข_i) # cotangent for s_to_sโ‚โฑ state + # Data cotangents from Eแดธแถปโฑ disaggregation + โˆ‚ฮฃฬ‚แถปโ‚_ac = zeros(T, nหข_i, nหข_i) + โˆ‚ฮฃฬ‚แถปโ‚‚_ac = zeros(T, size(d.ฮฃฬ‚แถปโ‚‚)) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_ac = zeros(T, nหข_i) + โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚_ac = zeros(T, nหข_i, nหข_i) + + ล_y = Matrix{T}(d.ล_to_yโ‚ƒ) + รช_y = Matrix{T}(d.รช_to_yโ‚ƒ) + ล_ล = Matrix{T}(d.ล_to_ลโ‚ƒ) + รช_ล = Matrix{T}(d.รช_to_ลโ‚ƒ) + vec_Ie_col = reshape(T.(vec_Iโ‚‘), :, 1) + ss_s_M = Matrix(d.ss_s) + vฮฃ_ac = vec(d.ฮฃฬ‚แถปโ‚) + n = nหข_i; ne = nแต‰ + sb_ac = cumsum([0, n, n, d.nโ‚‚หข, n, n^2, d.nโ‚ƒหข]) + eb_ac = cumsum([0, ne, ne^2, n*ne, n*ne, n^2*ne, n*ne^2, ne^3]) + + # Reverse loop over autocorrelation periods + for pi in n_ac:-1:1 + pp = d.per_period[pi] + + # โ”€โ”€ Step 5 reverse: ล_to_ลโ‚ƒโฑ_after = ล_to_ลโ‚ƒโฑ_before * ล_to_ลโ‚ƒ โ”€โ”€ + โˆ‚ล_to_ลโ‚ƒ_ac .+= pp.ล_to_ลโ‚ƒโฑ' * โˆ‚ล_to_ลโ‚ƒโฑ_co + โˆ‚ล_to_ลโ‚ƒโฑ_co .= โˆ‚ล_to_ลโ‚ƒโฑ_co * ล_ล' + + # โ”€โ”€ Step 4 reverse: autocorrelation output โ”€โ”€ + # Gather โˆ‚autocorr for this period + โˆ‚ac = zeros(T, nObs_iter) + if !(โˆ‚autocorr_in isa AbstractZero) + for obs in d.variance_observable + obs_local = indexin([obs], d.variance_observable) + obs_global = indexin([obs], T_pm.var) + โˆ‚ac[obs_local] .+= โˆ‚autocorr_in[obs_global, pp.period_index] + end + end + + # Apply zero mask + โˆ‚ac[pp.zero_mask] .= 0 + + if โ„’.norm(โˆ‚ac) > eps(T) + # Division adjoint: ac = num_diag / norm_diag + โˆ‚num_diag = โˆ‚ac ./ d.norm_diag + โˆ‚norm_diag_from_ac = -โˆ‚ac .* pp.num_diag ./ (d.norm_diag .^ 2) + # norm_diag = max.(diag(ฮฃสธโ‚ƒtmp), eps()) โ†’ adjoint only where diag >= eps + norm_mask = โ„’.diag(d.ฮฃสธโ‚ƒtmp) .>= eps(Float64) + โˆ‚ฮฃสธโ‚ƒtmp_ac .+= โ„’.Diagonal(โˆ‚norm_diag_from_ac .* norm_mask) + + # Numerator: N = ล_y * ฮฃแถปโ‚ƒโฑ * ล_y' + ล_y * ล_ลโ‚ƒโฑ * ac_tmp + รช_y * Eแดธแถปโฑ * ล_y' + # Adjoint of diag extraction: โˆ‚D = Diagonal(โˆ‚num_diag) + โˆ‚D = โ„’.Diagonal(โˆ‚num_diag) + + ฮฃแถปโ‚ƒโฑ_i = pp.ฮฃแถปโ‚ƒโฑ + ล_ลโ‚ƒโฑ_i = pp.ล_to_ลโ‚ƒโฑ + ELโฑ = Matrix{T}(pp.Eแดธแถปโฑ) + ac_tmp = Matrix{T}(d.autocorr_tmp_ac) + + # Term 1: diag(ล_y * ฮฃแถปโ‚ƒโฑ * ล_y') + โˆ‚ล_to_yโ‚ƒ_ac .+= โˆ‚D * ล_y * (ฮฃแถปโ‚ƒโฑ_i + ฮฃแถปโ‚ƒโฑ_i') + โˆ‚ฮฃแถปโ‚ƒโฑ_co .+= ล_y' * โˆ‚D * ล_y + + # Term 2: diag(ล_y * ล_ลโ‚ƒโฑ * ac_tmp) + โˆ‚ล_to_yโ‚ƒ_ac .+= โˆ‚D * ac_tmp' * ล_ลโ‚ƒโฑ_i' + โˆ‚ล_to_ลโ‚ƒโฑ_co .+= ล_y' * โˆ‚D * ac_tmp' + โˆ‚autocorr_tmp_co .+= ล_ลโ‚ƒโฑ_i' * ล_y' * โˆ‚D + + # Term 3: diag(รช_y * Eแดธแถปโฑ * ล_y') + โˆ‚รช_to_yโ‚ƒ_ac .+= โˆ‚D * ล_y * ELโฑ' + โˆ‚ล_to_yโ‚ƒ_ac .+= โˆ‚D * รช_y * ELโฑ + โˆ‚Eแดธแถปโฑ = รช_y' * โˆ‚D * ล_y + + # โ”€โ”€ Eแดธแถปโฑ disaggregation: only row block 6 has s_to_sโ‚โฑ prefix โ”€โ”€ + โˆ‚ELโฑ6 = โˆ‚Eแดธแถปโฑ[eb_ac[6]+1:eb_ac[7], :] + + sโ‚โฑ = pp.s_to_sโ‚โฑ # sโ‚^i (after step 2) + + # Col 1: kron(sโ‚โฑ * ฮฃฬ‚แถปโ‚, vec_Ie) + A_c1 = sโ‚โฑ * Matrix{T}(d.ฮฃฬ‚แถปโ‚) + โˆ‚A_c1 = _kron_vjp(โˆ‚ELโฑ6[:, sb_ac[1]+1:sb_ac[2]], A_c1, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_c1 * Matrix{T}(d.ฮฃฬ‚แถปโ‚)' + โˆ‚ฮฃฬ‚แถปโ‚_ac .+= sโ‚โฑ' * โˆ‚A_c1 + + # Col 4: kron(sโ‚โฑ * ฮผหขโ‚ƒฮดฮผหขโ‚', vec_Ie) + A_c4 = sโ‚โฑ * Matrix{T}(d.ฮผหขโ‚ƒฮดฮผหขโ‚') + โˆ‚A_c4 = _kron_vjp(โˆ‚ELโฑ6[:, sb_ac[4]+1:sb_ac[5]], A_c4, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_c4 * Matrix{T}(d.ฮผหขโ‚ƒฮดฮผหขโ‚) + โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚_ac .+= โˆ‚A_c4' * sโ‚โฑ + + # Col 5: kron(sโ‚โฑ * C4m, vec_Ie) + inner_C4 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] + d.ฮ”ฬ‚ฮผหขโ‚‚ * vฮฃ_ac' + C4m = reshape(ss_s_M * vec(inner_C4), n, n^2) + A_c5 = sโ‚โฑ * C4m + โˆ‚A_c5 = _kron_vjp(โˆ‚ELโฑ6[:, sb_ac[5]+1:sb_ac[6]], A_c5, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_c5 * C4m' + โˆ‚C4_i = sโ‚โฑ' * โˆ‚A_c5 + โˆ‚iC4_i = reshape(ss_s_M' * vec(โˆ‚C4_i), n, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚_ac[n+1:2n, 2n+1:end] .+= โˆ‚iC4_i + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_ac .+= โˆ‚iC4_i * vฮฃ_ac + โˆ‚ฮฃฬ‚แถปโ‚_ac .+= reshape(โˆ‚iC4_i' * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + + # Col 6: kron(sโ‚โฑ * C5m * Lโ‚ƒหข', vec_Ie) + inner_C5 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ_ac * vฮฃ_ac' + C5m = reshape(Matrix{T}(inner_C5), n, n^3) + C5m_c = C5m * Matrix(d.Lโ‚ƒหข)' + A_c6 = sโ‚โฑ * C5m_c + โˆ‚A_c6 = _kron_vjp(โˆ‚ELโฑ6[:, sb_ac[6]+1:sb_ac[7]], A_c6, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_c6 * C5m_c' + โˆ‚C5m_c_i = sโ‚โฑ' * โˆ‚A_c6 + โˆ‚C5_i = โˆ‚C5m_c_i * Matrix(d.Lโ‚ƒหข) + โˆ‚iC5_i = reshape(โˆ‚C5_i, n^2, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚_ac[2n+1:end, 2n+1:end] .+= โˆ‚iC5_i + โˆ‚ฮฃฬ‚แถปโ‚_ac .+= reshape((โˆ‚iC5_i + โˆ‚iC5_i') * vฮฃ_ac, n, n) + end # norm(โˆ‚ac) check + + # โ”€โ”€ Step 2 reverse: s_to_sโ‚โฑ_after = s_to_sโ‚โฑ_prev * s_to_sโ‚ โ”€โ”€ + sโ‚_d = Matrix{T}(d.s_to_sโ‚) + โˆ‚sโ‚_ac .+= pp.s_to_sโ‚โฑ_prev' * โˆ‚s_to_sโ‚โฑ_co + โˆ‚s_to_sโ‚โฑ_co .= โˆ‚s_to_sโ‚โฑ_co * sโ‚_d' + + # โ”€โ”€ Step 1 reverse: ฮฃแถปโ‚ƒโฑ = ล_ล * ฮฃแถปโ‚ƒโฑ_prev + รช_ล * Eแดธแถป_used โ”€โ”€ + EL_used = Matrix{T}(pp.Eแดธแถป_used) + โˆ‚ล_to_ลโ‚ƒ_ac .+= โˆ‚ฮฃแถปโ‚ƒโฑ_co * pp.ฮฃแถปโ‚ƒโฑ_prev' + โˆ‚รช_to_ลโ‚ƒ_ac .+= โˆ‚ฮฃแถปโ‚ƒโฑ_co * EL_used' + # โˆ‚Eแดธแถป_used: this flows to the previous period's Eแดธแถปโฑ or to the original Eแดธแถป + โˆ‚Eแดธแถป_used = รช_ล' * โˆ‚ฮฃแถปโ‚ƒโฑ_co + if pi == 1 + โˆ‚Eแดธแถป_ac .+= โˆ‚Eแดธแถป_used + else + # Flows to previous period's Eแดธแถปโฑ โ€” need to disaggregate + # The previous Eแดธแถปโฑ has s_to_sโ‚โฑ prefix from period pi-1 + pp_prev = d.per_period[pi-1] + sโ‚โฑ_prev = pp_prev.s_to_sโ‚โฑ + โˆ‚ELprev6 = โˆ‚Eแดธแถป_used[eb_ac[6]+1:eb_ac[7], :] + + # Col 1 + A_pc1 = sโ‚โฑ_prev * Matrix{T}(d.ฮฃฬ‚แถปโ‚) + โˆ‚A_pc1 = _kron_vjp(โˆ‚ELprev6[:, sb_ac[1]+1:sb_ac[2]], A_pc1, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_pc1 * Matrix{T}(d.ฮฃฬ‚แถปโ‚)' + โˆ‚ฮฃฬ‚แถปโ‚_ac .+= sโ‚โฑ_prev' * โˆ‚A_pc1 + + # Col 4 + A_pc4 = sโ‚โฑ_prev * Matrix{T}(d.ฮผหขโ‚ƒฮดฮผหขโ‚') + โˆ‚A_pc4 = _kron_vjp(โˆ‚ELprev6[:, sb_ac[4]+1:sb_ac[5]], A_pc4, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_pc4 * Matrix{T}(d.ฮผหขโ‚ƒฮดฮผหขโ‚) + โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚_ac .+= โˆ‚A_pc4' * sโ‚โฑ_prev + + # Col 5 + inner_C4p = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] + d.ฮ”ฬ‚ฮผหขโ‚‚ * vฮฃ_ac' + C4mp = reshape(ss_s_M * vec(inner_C4p), n, n^2) + A_pc5 = sโ‚โฑ_prev * C4mp + โˆ‚A_pc5 = _kron_vjp(โˆ‚ELprev6[:, sb_ac[5]+1:sb_ac[6]], A_pc5, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_pc5 * C4mp' + โˆ‚C4p = sโ‚โฑ_prev' * โˆ‚A_pc5 + โˆ‚iC4p = reshape(ss_s_M' * vec(โˆ‚C4p), n, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚_ac[n+1:2n, 2n+1:end] .+= โˆ‚iC4p + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_ac .+= โˆ‚iC4p * vฮฃ_ac + โˆ‚ฮฃฬ‚แถปโ‚_ac .+= reshape(โˆ‚iC4p' * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + + # Col 6 + inner_C5p = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ_ac * vฮฃ_ac' + C5mp = reshape(Matrix{T}(inner_C5p), n, n^3) + C5mp_c = C5mp * Matrix(d.Lโ‚ƒหข)' + A_pc6 = sโ‚โฑ_prev * C5mp_c + โˆ‚A_pc6 = _kron_vjp(โˆ‚ELprev6[:, sb_ac[6]+1:sb_ac[7]], A_pc6, vec_Ie_col)[1] + โˆ‚s_to_sโ‚โฑ_co .+= โˆ‚A_pc6 * C5mp_c' + โˆ‚C5m_c_p = sโ‚โฑ_prev' * โˆ‚A_pc6 + โˆ‚C5p = โˆ‚C5m_c_p * Matrix(d.Lโ‚ƒหข) + โˆ‚iC5p = reshape(โˆ‚C5p, n^2, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚_ac[2n+1:end, 2n+1:end] .+= โˆ‚iC5p + โˆ‚ฮฃฬ‚แถปโ‚_ac .+= reshape((โˆ‚iC5p + โˆ‚iC5p') * vฮฃ_ac, n, n) + + # The remaining rows (1-5 and 7) of โˆ‚Eแดธแถป_used are zero (spzeros in forward) + end + + # Propagate โˆ‚ฮฃแถปโ‚ƒโฑ to previous state + โˆ‚ฮฃแถปโ‚ƒโฑ_co .= ล_ล' * โˆ‚ฮฃแถปโ‚ƒโฑ_co + end # end autocorrelation reverse loop + + # โ”€โ”€ autocorr_tmp adjoint โ”€โ”€ + # autocorr_tmp = ล_ล * Eแดธแถป' * รช_y' + รช_ล * ฮ“โ‚ƒ * รช_y' + โˆ‚act = Matrix{T}(โˆ‚autocorr_tmp_co) + EL_orig = Matrix{T}(d.Eแดธแถป) + ฮ“โ‚ƒ_d = Matrix{T}(d.ฮ“โ‚ƒ) + + # Term 1: ล_ล * Eแดธแถป' * รช_y' + โˆ‚ล_to_ลโ‚ƒ_ac .+= โˆ‚act * รช_y * EL_orig + โˆ‚Eแดธแถป_ac .+= รช_y' * โˆ‚act' * ล_ล + โˆ‚รช_to_yโ‚ƒ_ac .+= โˆ‚act' * ล_ล * EL_orig' + + # Term 2: รช_ล * ฮ“โ‚ƒ * รช_y' + โˆ‚รช_to_ลโ‚ƒ_ac .+= โˆ‚act * รช_y * ฮ“โ‚ƒ_d' + โˆ‚ฮ“โ‚ƒ_ac .+= รช_ล' * โˆ‚act * รช_y + โˆ‚รช_to_yโ‚ƒ_ac .+= โˆ‚act' * รช_ล * ฮ“โ‚ƒ_d + + # ฮฃแถปโ‚ƒโฑ_co now holds the cotangent at the initial state (ฮฃแถปโ‚ƒโฑโ‚€ = ฮฃแถปโ‚ƒ) + # This adds to โˆ‚ฮฃแถปโ‚ƒ from the Lyapunov path + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # Stage 1: Output mapping (variance) โ€” same as existing rrule + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + # โ”€โ”€ Gather โˆ‚ฮฃสธโ‚ƒtmp from โˆ‚ฮฃสธโ‚ƒ (reverse of scatter) โ”€โ”€ + โˆ‚ฮฃสธโ‚ƒtmp = zeros(T, nObs_iter, nObs_iter) + + if !(โˆ‚ฮฃสธโ‚ƒ_in isa AbstractZero) + โˆ‚ฮฃสธโ‚ƒtmp .= โˆ‚ฮฃสธโ‚ƒ_in[d.obs_in_y, indexin(d.variance_observable, T_pm.var)] + end + + # Add autocorrelation contribution to โˆ‚ฮฃสธโ‚ƒtmp (from norm_diag) + โˆ‚ฮฃสธโ‚ƒtmp .+= โˆ‚ฮฃสธโ‚ƒtmp_ac + + if โ„’.norm(โˆ‚ฮฃสธโ‚ƒtmp) + โ„’.norm(โˆ‚ล_to_yโ‚ƒ_ac) + โ„’.norm(โˆ‚รช_to_yโ‚ƒ_ac) + โ„’.norm(โˆ‚ฮฃแถปโ‚ƒโฑ_co) + โ„’.norm(โˆ‚ล_to_ลโ‚ƒ_ac) + โ„’.norm(โˆ‚รช_to_ลโ‚ƒ_ac) + โ„’.norm(โˆ‚Eแดธแถป_ac) + โ„’.norm(โˆ‚ฮ“โ‚ƒ_ac) < eps(T); continue; end + + โˆ‚ฮฃสธโ‚ƒtmp_sym = โˆ‚ฮฃสธโ‚ƒtmp + โˆ‚ฮฃสธโ‚ƒtmp' + + # โ”€โ”€ ฮฃสธโ‚ƒtmp = ล_y * ฮฃแถปโ‚ƒ * ล_y' + รช_y * ฮ“โ‚ƒ * รช_y' + รช_y * Eแดธแถป * ล_y' + ล_y * Eแดธแถป' * รช_y' โ”€โ”€ + โˆ‚ล_to_yโ‚ƒ = โˆ‚ล_to_yโ‚ƒ_ac .+ โˆ‚ฮฃสธโ‚ƒtmp_sym * (d.ล_to_yโ‚ƒ * d.ฮฃแถปโ‚ƒ + d.รช_to_yโ‚ƒ * Matrix(d.Eแดธแถป)) + โˆ‚รช_to_yโ‚ƒ = โˆ‚รช_to_yโ‚ƒ_ac .+ โˆ‚ฮฃสธโ‚ƒtmp_sym * (d.รช_to_yโ‚ƒ * d.ฮ“โ‚ƒ + d.ล_to_yโ‚ƒ * Matrix(d.Eแดธแถป')) + โˆ‚ฮฃแถปโ‚ƒ = โˆ‚ฮฃแถปโ‚ƒโฑ_co .+ d.ล_to_yโ‚ƒ' * โˆ‚ฮฃสธโ‚ƒtmp * d.ล_to_yโ‚ƒ + โˆ‚ฮ“โ‚ƒ_iter = โˆ‚ฮ“โ‚ƒ_ac .+ d.รช_to_yโ‚ƒ' * โˆ‚ฮฃสธโ‚ƒtmp * d.รช_to_yโ‚ƒ + โˆ‚Eแดธแถป_iter = โˆ‚Eแดธแถป_ac .+ d.รช_to_yโ‚ƒ' * โˆ‚ฮฃสธโ‚ƒtmp_sym * d.ล_to_yโ‚ƒ + + # โ”€โ”€ Standard Lyapunov adjoint โ”€โ”€ + Nu = d.N_upper; Nl = d.N_lower + ru_i = 1:Nu; rl_i = (Nu+1):(Nu+Nl) + + lyap_grad = d.lyap_pb((โˆ‚ฮฃแถปโ‚ƒ, NoTangent())) + โˆ‚ล_to_ลโ‚ƒ = lyap_grad[2] isa AbstractZero ? zeros(T, size(d.ล_to_ลโ‚ƒ)) : Matrix{T}(lyap_grad[2]) + โˆ‚C_lyap = lyap_grad[3] isa AbstractZero ? zeros(T, size(d.ล_to_ลโ‚ƒ)) : Matrix{T}(lyap_grad[3]) + + # Backprop through C = รช * ฮ“โ‚ƒ * รช' + M + M' where M = รช * Eแดธแถป * ล' + โˆ‚C_sym = โˆ‚C_lyap + โˆ‚C_lyap' + รช_d = Matrix{T}(d.รช_to_ลโ‚ƒ) + ล_d = Matrix{T}(d.ล_to_ลโ‚ƒ) + EL_d = Matrix{T}(d.Eแดธแถป) + ฮ“โ‚ƒ_d = Matrix{T}(d.ฮ“โ‚ƒ) + + # Term 1: รช * ฮ“โ‚ƒ * รช' + โˆ‚ฮ“โ‚ƒ_iter .+= รช_d' * โˆ‚C_lyap * รช_d + โˆ‚รช_to_ลโ‚ƒ = โˆ‚รช_to_ลโ‚ƒ_ac .+ โˆ‚C_sym * รช_d * ฮ“โ‚ƒ_d + + # Terms 2+3: M + M' where M = รช * Eแดธแถป * ล' + โˆ‚รช_to_ลโ‚ƒ .+= โˆ‚C_sym * ล_d * EL_d' + โˆ‚Eแดธแถป_iter .+= รช_d' * โˆ‚C_sym * ล_d + โˆ‚ล_to_ลโ‚ƒ .+= โˆ‚C_sym' * รช_d * EL_d + + # Add autocorrelation contributions + โˆ‚ล_to_ลโ‚ƒ .+= โˆ‚ล_to_ลโ‚ƒ_ac + + # Extract โˆ‚A_UU, โˆ‚A_LU, โˆ‚A_LL from โˆ‚ล_to_ลโ‚ƒ + โˆ‚A_UU = โˆ‚ล_to_ลโ‚ƒ[ru_i, ru_i] + โˆ‚A_LU = โˆ‚ล_to_ลโ‚ƒ[rl_i, ru_i] + โˆ‚A_LL = โˆ‚ล_to_ลโ‚ƒ[rl_i, rl_i] + + # โ”€โ”€ Disaggregate ล_to_yโ‚ƒ โ†’ โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, โˆ‚๐’โ‚ƒ โ”€โ”€ + nโ‚‚หข_i = d.nโ‚‚หข; nโ‚ƒหข_i = d.nโ‚ƒหข + c = 0 + โˆ‚blk1 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i]; c += nหข_i + โˆ‚blk2 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i]; c += nหข_i + โˆ‚blk3 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nโ‚‚หข_i]; c += nโ‚‚หข_i # compressed + โˆ‚blk4 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i]; c += nหข_i + โˆ‚blk5 = โˆ‚ล_to_yโ‚ƒ[:, c+1:c+nหข_i^2]; c += nหข_i^2 + โˆ‚blk6 = โˆ‚ล_to_yโ‚ƒ[:, c+1:end] + + โˆ‚๐’โ‚_acc[d.obs_in_y, d.dependencies_in_states_idx] .+= โˆ‚blk1 .+ โˆ‚blk2 .+ โˆ‚blk4 + โˆ‚S2f_acc[d.obs_in_y, d.kron_s_s] .+= (โˆ‚blk3 * Matrix(d.Dโ‚‚หข)') ./ 2 .+ โˆ‚blk5 # decompress blk3 + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_v_v] .+= โˆ‚blk1 ./ 2 + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_s_s] .+= (โˆ‚blk6 * Matrix(d.Dโ‚ƒหข)') ./ 6 # decompress blk6 + + # โ”€โ”€ Disaggregate รช_to_yโ‚ƒ โ†’ โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, โˆ‚๐’โ‚ƒ โ”€โ”€ + c = 0 + โˆ‚eblk1 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nแต‰]; c += nแต‰ + โˆ‚eblk2 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nแต‰^2]; c += nแต‰^2 + โˆ‚eblk3 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i*nแต‰]; c += nหข_i*nแต‰ + โˆ‚eblk4 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i*nแต‰]; c += nหข_i*nแต‰ + โˆ‚eblk5 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i^2*nแต‰]; c += nหข_i^2*nแต‰ + โˆ‚eblk6 = โˆ‚รช_to_yโ‚ƒ[:, c+1:c+nหข_i*nแต‰^2]; c += nหข_i*nแต‰^2 + โˆ‚eblk7 = โˆ‚รช_to_yโ‚ƒ[:, c+1:end] + + โˆ‚๐’โ‚_acc[d.obs_in_y, nโ‚‹+1:end] .+= โˆ‚eblk1 + โˆ‚S2f_acc[d.obs_in_y, kron_e_e] .+= โˆ‚eblk2 ./ 2 + โˆ‚S2f_acc[d.obs_in_y, d.kron_s_e] .+= โˆ‚eblk3 .+ โˆ‚eblk4 + โˆ‚S3f_acc[d.obs_in_y, d.kron_e_v_v] .+= โˆ‚eblk1 ./ 2 + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_s_e] .+= โˆ‚eblk5 ./ 2 + โˆ‚S3f_acc[d.obs_in_y, d.kron_s_e_e] .+= โˆ‚eblk6 ./ 2 + โˆ‚S3f_acc[d.obs_in_y, d.kron_e_e_e] .+= โˆ‚eblk7 ./ 6 + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # Stage 2+3: Disaggregate block matrices โ†’ slice & data cotangents + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + Ine = Matrix{T}(โ„’.I(ne)) + + # Dense copies of frequently used slices + sโ‚ = Matrix{T}(d.s_to_sโ‚) + eโ‚ = Matrix{T}(d.e_to_sโ‚) + sโ‚ยฒ = Matrix{T}(d.s_to_sโ‚_by_s_to_sโ‚) + eโ‚ยฒ = Matrix{T}(d.e_to_sโ‚_by_e_to_sโ‚) + sโ‚eโ‚ = Matrix{T}(d.s_to_sโ‚_by_e_to_sโ‚) + ssโ‚‚ = Matrix{T}(d.s_s_to_sโ‚‚) + eeโ‚‚ = Matrix{T}(d.e_e_to_sโ‚‚) + seโ‚‚ = Matrix{T}(d.s_e_to_sโ‚‚) + vvโ‚‚ = Matrix{T}(d.v_v_to_sโ‚‚) + + # Local slice cotangent accumulators + โˆ‚sโ‚_l = โˆ‚sโ‚_ac # start with autocorrelation contribution + โˆ‚eโ‚_l = zeros(T, n, ne) + โˆ‚ssโ‚‚_l = zeros(T, n, n^2) + โˆ‚eeโ‚‚_l = zeros(T, n, ne^2) + โˆ‚seโ‚‚_l = zeros(T, n, n * ne) + โˆ‚vvโ‚‚_l = zeros(T, size(vvโ‚‚)) + โˆ‚ฮฃฬ‚แถปโ‚ = โˆ‚ฮฃฬ‚แถปโ‚_ac # start with autocorrelation contribution + โˆ‚ฮฃฬ‚แถปโ‚‚ = โˆ‚ฮฃฬ‚แถปโ‚‚_ac # start with autocorrelation contribution + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l = โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_ac # start with autocorrelation contribution + + # Block boundary arrays + sb = cumsum([0, n, n, nโ‚‚หข_i, n, n^2, nโ‚ƒหข_i]) # ล_to_ลโ‚ƒ row/col (compressed) + eb = cumsum([0, ne, ne^2, n*ne, n*ne, n^2*ne, n*ne^2, ne^3]) + gb = eb + + vvh = vvโ‚‚ ./ 2; ssh = ssโ‚‚ ./ 2; eeh = eeโ‚‚ ./ 2 + + # โ”€โ”€ 2a: A_UU, A_LU, A_LL disaggregation โ”€โ”€ + # Block boundaries within sub-matrices + bu = cumsum([0, n, n, nโ‚‚หข_i]) # A_UU row/col blocks + bl = cumsum([0, n, n^2, nโ‚ƒหข_i]) # A_LL row/col blocks (also A_LU rows) + + # โ”€โ”€ From โˆ‚A_UU โ”€โ”€ + # (1,1) sโ‚, (2,2) sโ‚ + โˆ‚sโ‚_l .+= โˆ‚A_UU[bu[1]+1:bu[2], bu[1]+1:bu[2]] .+ + โˆ‚A_UU[bu[2]+1:bu[3], bu[2]+1:bu[3]] + # (2,3) ssโ‚‚/2 * Dโ‚‚หข โ€” decompress cols + โˆ‚ssโ‚‚_l .+= โˆ‚A_UU[bu[2]+1:bu[3], bu[3]+1:bu[4]] * Matrix(d.Dโ‚‚หข)' ./ 2 + # (3,3) Lโ‚‚หข * kron(sโ‚,sโ‚) * Dโ‚‚หข โ€” decompress then kron_vjp + โˆ‚inner33 = Matrix(d.Lโ‚‚หข)' * Matrix(โˆ‚A_UU[bu[3]+1:bu[4], bu[3]+1:bu[4]]) * Matrix(d.Dโ‚‚หข)' + tmpL, tmpR = _kron_vjp(โˆ‚inner33, sโ‚, sโ‚) + โˆ‚sโ‚_l .+= tmpL .+ tmpR + + # โ”€โ”€ From โˆ‚A_LU โ”€โ”€ + # (1,1) s_vvโ‚ƒ/2 + โˆ‚S3f_acc[d.iหข, d.kron_s_v_v] .+= โˆ‚A_LU[bl[1]+1:bl[2], bu[1]+1:bu[2]] ./ 2 + # (2,1) kron(sโ‚, vvโ‚‚/2) + tmpA, tmpB = _kron_vjp(Matrix(โˆ‚A_LU[bl[2]+1:bl[3], bu[1]+1:bu[2]]), sโ‚, vvh) + โˆ‚sโ‚_l .+= tmpA; โˆ‚vvโ‚‚_l .+= tmpB ./ 2 + + # โ”€โ”€ From โˆ‚A_LL โ”€โ”€ + # (1,1) sโ‚ + โˆ‚sโ‚_l .+= โˆ‚A_LL[bl[1]+1:bl[2], bl[1]+1:bl[2]] + # (1,2) ssโ‚‚ + โˆ‚ssโ‚‚_l .+= โˆ‚A_LL[bl[1]+1:bl[2], bl[2]+1:bl[3]] + # (1,3) sssโ‚ƒ/6 * Dโ‚ƒหข โ€” decompress cols + โˆ‚S3f_acc[d.iหข, d.kron_s_s_s] .+= โˆ‚A_LL[bl[1]+1:bl[2], bl[3]+1:bl[4]] * Matrix(d.Dโ‚ƒหข)' ./ 6 + # (2,2) kron(sโ‚,sโ‚) + tmpL, tmpR = _kron_vjp(Matrix(โˆ‚A_LL[bl[2]+1:bl[3], bl[2]+1:bl[3]]), sโ‚, sโ‚) + โˆ‚sโ‚_l .+= tmpL .+ tmpR + # (2,3) kron(sโ‚, ssโ‚‚/2) * Dโ‚ƒหข โ€” decompress cols then kron_vjp + โˆ‚inner56 = Matrix(โˆ‚A_LL[bl[2]+1:bl[3], bl[3]+1:bl[4]]) * Matrix(d.Dโ‚ƒหข)' + tmpA, tmpB = _kron_vjp(โˆ‚inner56, sโ‚, ssh) + โˆ‚sโ‚_l .+= tmpA; โˆ‚ssโ‚‚_l .+= tmpB ./ 2 + # (3,3) Lโ‚ƒหข * kron(sโ‚, kron(sโ‚,sโ‚)) * Dโ‚ƒหข โ€” decompress then kron_vjp + โˆ‚inner66 = Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚A_LL[bl[3]+1:bl[4], bl[3]+1:bl[4]]) * Matrix(d.Dโ‚ƒหข)' + tmpA, tmpB = _kron_vjp(โˆ‚inner66, sโ‚, sโ‚ยฒ) + โˆ‚sโ‚_l .+= tmpA + tmpL, tmpR = _kron_vjp(tmpB, sโ‚, sโ‚) + โˆ‚sโ‚_l .+= tmpL .+ tmpR + + # โ”€โ”€ 2b: รช_to_ลโ‚ƒ disaggregation โ”€โ”€ + โˆ‚รชโ‚ƒ = Matrix{T}(โˆ‚รช_to_ลโ‚ƒ) + ss_s1e1 = Matrix(d.s_s) * sโ‚eโ‚ + + # Row 1: (1,1) eโ‚ + โˆ‚eโ‚_l .+= โˆ‚รชโ‚ƒ[sb[1]+1:sb[2], eb[1]+1:eb[2]] + # Row 2: (2,2) eeโ‚‚/2; (2,3) seโ‚‚ + โˆ‚eeโ‚‚_l .+= โˆ‚รชโ‚ƒ[sb[2]+1:sb[3], eb[2]+1:eb[3]] ./ 2 + โˆ‚seโ‚‚_l .+= โˆ‚รชโ‚ƒ[sb[2]+1:sb[3], eb[3]+1:eb[4]] + # Row 3: (3,2) Lโ‚‚หข * kron(eโ‚,eโ‚) โ€” decompress rows + tmpL, tmpR = _kron_vjp(Matrix(d.Lโ‚‚หข)' * Matrix(โˆ‚รชโ‚ƒ[sb[3]+1:sb[4], eb[2]+1:eb[3]]), eโ‚, eโ‚) + โˆ‚eโ‚_l .+= tmpL .+ tmpR + # (3,3) Lโ‚‚หข * I_plus_s_s * kron(sโ‚,eโ‚) โ€” decompress rows + โˆ‚k33 = Matrix(d.I_plus_s_s') * Matrix(d.Lโ‚‚หข)' * Matrix(โˆ‚รชโ‚ƒ[sb[3]+1:sb[4], eb[3]+1:eb[4]]) + tmpA, tmpB = _kron_vjp(โˆ‚k33, sโ‚, eโ‚) + โˆ‚sโ‚_l .+= tmpA; โˆ‚eโ‚_l .+= tmpB + # Row 4: direct Sโ‚ƒ slices + โˆ‚S3f_acc[d.iหข, d.kron_e_v_v] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[1]+1:eb[2]] ./ 2 + โˆ‚seโ‚‚_l .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[4]+1:eb[5]] + โˆ‚S3f_acc[d.iหข, d.kron_s_s_e] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[5]+1:eb[6]] ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_s_e_e] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[6]+1:eb[7]] ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_e_e_e] .+= โˆ‚รชโ‚ƒ[sb[4]+1:sb[5], eb[7]+1:eb[8]] ./ 6 + # Row 5: (5,1) kron(eโ‚,vvโ‚‚/2) + tmpA, tmpB = _kron_vjp(Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[1]+1:eb[2]]), eโ‚, vvh) + โˆ‚eโ‚_l .+= tmpA; โˆ‚vvโ‚‚_l .+= tmpB ./ 2 + # (5,4) s_s * kron(sโ‚,eโ‚) + โˆ‚k54 = Matrix(d.s_s') * Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[4]+1:eb[5]]) + tmpA, tmpB = _kron_vjp(โˆ‚k54, sโ‚, eโ‚) + โˆ‚sโ‚_l .+= tmpA; โˆ‚eโ‚_l .+= tmpB + # (5,5) kron(sโ‚,seโ‚‚) + s_s * kron(ssโ‚‚/2, eโ‚) + โˆ‚b55 = Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[5]+1:eb[6]]) + tmpA, tmpB = _kron_vjp(โˆ‚b55, sโ‚, seโ‚‚) + โˆ‚sโ‚_l .+= tmpA; โˆ‚seโ‚‚_l .+= tmpB + โˆ‚k55b = Matrix(d.s_s') * โˆ‚b55 + tmpA, tmpB = _kron_vjp(โˆ‚k55b, ssh, eโ‚) + โˆ‚ssโ‚‚_l .+= tmpA ./ 2; โˆ‚eโ‚_l .+= tmpB + # (5,6) kron(sโ‚,eeโ‚‚/2) + s_s * kron(seโ‚‚, eโ‚) + โˆ‚b56 = Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[6]+1:eb[7]]) + tmpA, tmpB = _kron_vjp(โˆ‚b56, sโ‚, eeh) + โˆ‚sโ‚_l .+= tmpA; โˆ‚eeโ‚‚_l .+= tmpB ./ 2 + โˆ‚k56b = Matrix(d.s_s') * โˆ‚b56 + tmpA, tmpB = _kron_vjp(โˆ‚k56b, seโ‚‚, eโ‚) + โˆ‚seโ‚‚_l .+= tmpA; โˆ‚eโ‚_l .+= tmpB + # (5,7) kron(eโ‚, eeโ‚‚/2) + tmpA, tmpB = _kron_vjp(Matrix(โˆ‚รชโ‚ƒ[sb[5]+1:sb[6], eb[7]+1:eb[8]]), eโ‚, eeh) + โˆ‚eโ‚_l .+= tmpA; โˆ‚eeโ‚‚_l .+= tmpB ./ 2 + # Row 6: (6,5) Lโ‚ƒหข * (kron(sโ‚ยฒ,eโ‚) + kron(sโ‚,s_s*sโ‚eโ‚) + kron(eโ‚,sโ‚ยฒ)*e_ss) โ€” decompress rows + โˆ‚b65 = Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚รชโ‚ƒ[sb[6]+1:sb[7], eb[5]+1:eb[6]]) + tmpA, tmpB = _kron_vjp(โˆ‚b65, sโ‚ยฒ, eโ‚) + โˆ‚eโ‚_l .+= tmpB + tmpL, tmpR = _kron_vjp(tmpA, sโ‚, sโ‚); โˆ‚sโ‚_l .+= tmpL .+ tmpR + tmpA, tmpB = _kron_vjp(โˆ‚b65, sโ‚, ss_s1e1) + โˆ‚sโ‚_l .+= tmpA + tmpC = Matrix(d.s_s') * tmpB + tmpL, tmpR = _kron_vjp(tmpC, sโ‚, eโ‚); โˆ‚sโ‚_l .+= tmpL; โˆ‚eโ‚_l .+= tmpR + โˆ‚k65c = โˆ‚b65 * Matrix(d.e_ss') + tmpA, tmpB = _kron_vjp(โˆ‚k65c, eโ‚, sโ‚ยฒ) + โˆ‚eโ‚_l .+= tmpA + tmpL, tmpR = _kron_vjp(tmpB, sโ‚, sโ‚); โˆ‚sโ‚_l .+= tmpL .+ tmpR + # (6,6) Lโ‚ƒหข * (kron(sโ‚eโ‚,eโ‚) + kron(eโ‚,sโ‚eโ‚)*e_es + kron(eโ‚,s_s*sโ‚eโ‚)*e_es) โ€” decompress rows + โˆ‚b66 = Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚รชโ‚ƒ[sb[6]+1:sb[7], eb[6]+1:eb[7]]) + tmpA, tmpB = _kron_vjp(โˆ‚b66, sโ‚eโ‚, eโ‚) + โˆ‚eโ‚_l .+= tmpB + tmpL, tmpR = _kron_vjp(tmpA, sโ‚, eโ‚); โˆ‚sโ‚_l .+= tmpL; โˆ‚eโ‚_l .+= tmpR + โˆ‚pre = โˆ‚b66 * Matrix(d.e_es') + tmpA, tmpB = _kron_vjp(โˆ‚pre, eโ‚, sโ‚eโ‚) + โˆ‚eโ‚_l .+= tmpA + tmpL, tmpR = _kron_vjp(tmpB, sโ‚, eโ‚); โˆ‚sโ‚_l .+= tmpL; โˆ‚eโ‚_l .+= tmpR + tmpA, tmpB = _kron_vjp(โˆ‚pre, eโ‚, ss_s1e1) + โˆ‚eโ‚_l .+= tmpA + tmpC = Matrix(d.s_s') * tmpB + tmpL, tmpR = _kron_vjp(tmpC, sโ‚, eโ‚); โˆ‚sโ‚_l .+= tmpL; โˆ‚eโ‚_l .+= tmpR + # (6,7) Lโ‚ƒหข * kron(eโ‚, eโ‚ยฒ) โ€” decompress rows + tmpA, tmpB = _kron_vjp(Matrix(d.Lโ‚ƒหข)' * Matrix(โˆ‚รชโ‚ƒ[sb[6]+1:sb[7], eb[7]+1:eb[8]]), eโ‚, eโ‚ยฒ) + โˆ‚eโ‚_l .+= tmpA + tmpL, tmpR = _kron_vjp(tmpB, eโ‚, eโ‚); โˆ‚eโ‚_l .+= tmpL .+ tmpR + + # โ”€โ”€ 3a: ฮ“โ‚ƒ disaggregation โ†’ โˆ‚ฮฃฬ‚แถปโ‚, โˆ‚ฮฃฬ‚แถปโ‚‚, โˆ‚ฮ”ฬ‚ฮผหขโ‚‚ โ”€โ”€ + โˆ‚ฮ“ = Matrix{T}(โˆ‚ฮ“โ‚ƒ_iter) + vฮฃ = vec(d.ฮฃฬ‚แถปโ‚) + + # Row 1: (1,4) kron(ฮ”ฬ‚ฮผหขโ‚‚',Ine) + โˆ‚tmp14 = _kron_vjp(โˆ‚ฮ“[gb[1]+1:gb[2], gb[4]+1:gb[5]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, 1, :), Ine)[1] + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(โˆ‚tmp14') + # (1,5) kron(vec(ฮฃฬ‚แถปโ‚)',Ine) + โˆ‚tmp15 = _kron_vjp(โˆ‚ฮ“[gb[1]+1:gb[2], gb[5]+1:gb[6]], reshape(vฮฃ, 1, :), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(vec(โˆ‚tmp15'), n, n) + # Row 3: (3,3) kron(ฮฃฬ‚แถปโ‚,Ine) + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚ฮ“[gb[3]+1:gb[4], gb[3]+1:gb[4]], Matrix(d.ฮฃฬ‚แถปโ‚), Ine)[1] + # Row 4: (4,1) kron(ฮ”ฬ‚ฮผหขโ‚‚,Ine) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(_kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[1]+1:gb[2]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Ine)[1]) + # (4,4) kron(ฮฃฬ‚แถปโ‚‚_22 + ฮ”*ฮ”', Ine) + M44 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, n+1:2n] + d.ฮ”ฬ‚ฮผหขโ‚‚ * d.ฮ”ฬ‚ฮผหขโ‚‚' + โˆ‚M44 = _kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[4]+1:gb[5]], Matrix(M44), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[n+1:2n, n+1:2n] .+= โˆ‚M44 + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= (โˆ‚M44 + โˆ‚M44') * d.ฮ”ฬ‚ฮผหขโ‚‚ + # (4,5) kron(ฮฃฬ‚แถปโ‚‚_23 + ฮ”*vฮฃ', Ine) + M45 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] + d.ฮ”ฬ‚ฮผหขโ‚‚ * vฮฃ' + โˆ‚M45 = _kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[5]+1:gb[6]], Matrix(M45), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] .+= โˆ‚M45 + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚M45 * vฮฃ + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚M45' * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + # (4,7) kron(ฮ”ฬ‚ฮผหขโ‚‚, e4_nแต‰_nแต‰ยณ) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(_kron_vjp(โˆ‚ฮ“[gb[4]+1:gb[5], gb[7]+1:gb[8]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Matrix(e4_nแต‰_nแต‰ยณ))[1]) + # Row 5: (5,1) kron(vฮฃ, Ine) + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(_kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[1]+1:gb[2]], reshape(vฮฃ, :, 1), Ine)[1], n, n) + # (5,4) kron(ฮฃฬ‚แถปโ‚‚_32 + vฮฃ*ฮ”', Ine) + M54 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] + vฮฃ * d.ฮ”ฬ‚ฮผหขโ‚‚' + โˆ‚M54 = _kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[4]+1:gb[5]], Matrix(M54), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] .+= โˆ‚M54 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚M54 * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚M54' * vฮฃ + # (5,5) kron(ฮฃฬ‚แถปโ‚‚_33 + vฮฃ*vฮฃ', Ine) + M55 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ * vฮฃ' + โˆ‚M55 = _kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[5]+1:gb[6]], Matrix(M55), Ine)[1] + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] .+= โˆ‚M55 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape((โˆ‚M55 + โˆ‚M55') * vฮฃ, n, n) + # (5,7) kron(vฮฃ, e4_nแต‰_nแต‰ยณ) + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(_kron_vjp(โˆ‚ฮ“[gb[5]+1:gb[6], gb[7]+1:gb[8]], reshape(vฮฃ, :, 1), Matrix(e4_nแต‰_nแต‰ยณ))[1], n, n) + # Row 6: (6,6) kron(ฮฃฬ‚แถปโ‚, e4_nแต‰ยฒ_nแต‰ยฒ) + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚ฮ“[gb[6]+1:gb[7], gb[6]+1:gb[7]], Matrix(d.ฮฃฬ‚แถปโ‚), Matrix(e4_nแต‰ยฒ_nแต‰ยฒ))[1] + # Row 7: (7,4) kron(ฮ”ฬ‚ฮผหขโ‚‚', e4') + โˆ‚tmp74 = _kron_vjp(โˆ‚ฮ“[gb[7]+1:gb[8], gb[4]+1:gb[5]], reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, 1, :), Matrix(e4_nแต‰_nแต‰ยณ'))[1] + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(โˆ‚tmp74') + # (7,5) kron(vฮฃ', e4') + โˆ‚tmp75 = _kron_vjp(โˆ‚ฮ“[gb[7]+1:gb[8], gb[5]+1:gb[6]], reshape(vฮฃ, 1, :), Matrix(e4_nแต‰_nแต‰ยณ'))[1] + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(vec(โˆ‚tmp75'), n, n) + + # โ”€โ”€ 3b: Eแดธแถป disaggregation โ”€โ”€ + โˆ‚EL = Matrix{T}(โˆ‚Eแดธแถป_iter) + # Only row block 6 is data-dependent + โˆ‚EL6 = โˆ‚EL[gb[6]+1:gb[7], :] + # Col 1: kron(ฮฃฬ‚แถปโ‚, vec_Ie) + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚EL6[:, sb[1]+1:sb[2]], Matrix(d.ฮฃฬ‚แถปโ‚), vec_Ie_col)[1] + # Col 4: kron(ฮผหขโ‚ƒฮดฮผหขโ‚', vec_Ie) + โˆ‚ฮผ_T = _kron_vjp(โˆ‚EL6[:, sb[4]+1:sb[5]], Matrix(d.ฮผหขโ‚ƒฮดฮผหขโ‚'), vec_Ie_col)[1] + โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚ = โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚_ac .+ Matrix(โˆ‚ฮผ_T') + # Col 5: kron(Cโ‚„, vec_Ie) + inner_C4 = d.ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] + d.ฮ”ฬ‚ฮผหขโ‚‚ * vฮฃ' + C4m = reshape(ss_s_M * vec(inner_C4), n, n^2) + โˆ‚C4 = _kron_vjp(โˆ‚EL6[:, sb[5]+1:sb[6]], C4m, vec_Ie_col)[1] + โˆ‚iC4 = reshape(ss_s_M' * vec(โˆ‚C4), n, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚[n+1:2n, 2n+1:end] .+= โˆ‚iC4 + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚iC4 * vฮฃ + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚iC4' * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + # Col 6: kron(Cโ‚… * Lโ‚ƒหข', vec_Ie) โ€” compress col 6 + inner_C5 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ * vฮฃ' + C5m = reshape(Matrix(inner_C5), n, n^3) + C5m_c = C5m * Matrix(d.Lโ‚ƒหข)' + โˆ‚C5_c = _kron_vjp(โˆ‚EL6[:, sb[6]+1:sb[7]], C5m_c, vec_Ie_col)[1] + โˆ‚C5 = โˆ‚C5_c * Matrix(d.Lโ‚ƒหข) + โˆ‚iC5 = reshape(โˆ‚C5, n^2, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] .+= โˆ‚iC5 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape((โˆ‚iC5 + โˆ‚iC5') * vฮฃ, n, n) + + # โ”€โ”€ 3c: ฮผหขโ‚ƒฮดฮผหขโ‚ adjoint โ”€โ”€ + โˆ‚x_ฮผ = vec(โˆ‚ฮผหขโ‚ƒฮดฮผหขโ‚) + I_m_sโ‚ยฒ = Matrix{T}(โ„’.I(n^2)) - sโ‚ยฒ + โˆ‚b_ฮผ = I_m_sโ‚ยฒ' \ โˆ‚x_ฮผ + โˆ‚sโ‚ยฒ_from_ฮผ = โˆ‚b_ฮผ * vec(d.ฮผหขโ‚ƒฮดฮผหขโ‚)' + tmpL, tmpR = _kron_vjp(โˆ‚sโ‚ยฒ_from_ฮผ, sโ‚, sโ‚); โˆ‚sโ‚_l .+= tmpL .+ tmpR + + โˆ‚RHS = reshape(โˆ‚b_ฮผ, n, n) + + inner_M1 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] + vฮฃ * d.ฮ”ฬ‚ฮผหขโ‚‚' + M1 = reshape(ss_s_M * vec(inner_M1), n^2, n) + inner_M2 = d.ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] + vฮฃ * vฮฃ' + M2 = reshape(Matrix(inner_M2), n^3, n) + M3 = โ„’.kron(Matrix(d.ฮฃฬ‚แถปโ‚), vec_Ie_col) + + Lโ‚ = ssโ‚‚ * M1 + Matrix(d.s_s_s_to_sโ‚ƒ) * M2 / 6 + + Matrix(d.s_e_e_to_sโ‚ƒ) * M3 / 2 + Matrix(d.s_v_v_to_sโ‚ƒ) * Matrix(d.ฮฃฬ‚แถปโ‚) / 2 + + M4 = โ„’.kron(reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Ine) + M5 = Matrix(e4_nแต‰_nแต‰ยณ') + M6 = โ„’.kron(reshape(vฮฃ, :, 1), Ine) + + Lโ‚‚ = seโ‚‚ * M4 + Matrix(d.e_e_e_to_sโ‚ƒ) * M5 / 6 + + Matrix(d.s_s_e_to_sโ‚ƒ) * M6 / 2 + Matrix(d.e_v_v_to_sโ‚ƒ) * Ine / 2 + + โˆ‚Lโ‚ = โˆ‚RHS * sโ‚; โˆ‚sโ‚_l .+= โˆ‚RHS' * Lโ‚ + โˆ‚Lโ‚‚ = โˆ‚RHS * eโ‚; โˆ‚eโ‚_l .+= โˆ‚RHS' * Lโ‚‚ + + # Decompose โˆ‚Lโ‚ + โˆ‚ssโ‚‚_l .+= โˆ‚Lโ‚ * M1' + โˆ‚M1_raw = ssโ‚‚' * โˆ‚Lโ‚ + โˆ‚S3f_acc[d.iหข, d.kron_s_s_s] .+= โˆ‚Lโ‚ * M2' ./ 6 + โˆ‚M2_raw = Matrix(d.s_s_s_to_sโ‚ƒ)' * โˆ‚Lโ‚ ./ 6 + โˆ‚S3f_acc[d.iหข, d.kron_s_e_e] .+= โˆ‚Lโ‚ * M3' ./ 2 + โˆ‚M3_raw = Matrix(d.s_e_e_to_sโ‚ƒ)' * โˆ‚Lโ‚ ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_s_v_v] .+= โˆ‚Lโ‚ * Matrix(d.ฮฃฬ‚แถปโ‚)' ./ 2 + โˆ‚ฮฃฬ‚แถปโ‚ .+= Matrix(d.s_v_v_to_sโ‚ƒ)' * โˆ‚Lโ‚ ./ 2 + + # Decompose โˆ‚Lโ‚‚ + โˆ‚seโ‚‚_l .+= โˆ‚Lโ‚‚ * M4' + โˆ‚M4_raw = seโ‚‚' * โˆ‚Lโ‚‚ + โˆ‚S3f_acc[d.iหข, d.kron_e_e_e] .+= โˆ‚Lโ‚‚ * M5' ./ 6 + โˆ‚S3f_acc[d.iหข, d.kron_s_s_e] .+= โˆ‚Lโ‚‚ * M6' ./ 2 + โˆ‚M6_raw = Matrix(d.s_s_e_to_sโ‚ƒ)' * โˆ‚Lโ‚‚ ./ 2 + โˆ‚S3f_acc[d.iหข, d.kron_e_v_v] .+= โˆ‚Lโ‚‚ ./ 2 + + # Decompose โˆ‚M1 โ†’ โˆ‚ฮฃฬ‚แถปโ‚‚, โˆ‚ฮฃฬ‚แถปโ‚, โˆ‚ฮ”ฬ‚ฮผหขโ‚‚ + โˆ‚iM1 = reshape(ss_s_M' * vec(โˆ‚M1_raw), n^2, n) + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, n+1:2n] .+= โˆ‚iM1 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(โˆ‚iM1 * d.ฮ”ฬ‚ฮผหขโ‚‚, n, n) + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= โˆ‚iM1' * vฮฃ + # Decompose โˆ‚M2 โ†’ โˆ‚ฮฃฬ‚แถปโ‚‚, โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚iM2 = reshape(โˆ‚M2_raw, n^2, n^2) + โˆ‚ฮฃฬ‚แถปโ‚‚[2n+1:end, 2n+1:end] .+= โˆ‚iM2 + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape((โˆ‚iM2 + โˆ‚iM2') * vฮฃ, n, n) + # Decompose โˆ‚M3 โ†’ โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚ฮฃฬ‚แถปโ‚ .+= _kron_vjp(โˆ‚M3_raw, Matrix(d.ฮฃฬ‚แถปโ‚), vec_Ie_col)[1] + # Decompose โˆ‚M4 โ†’ โˆ‚ฮ”ฬ‚ฮผหขโ‚‚ + โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l .+= vec(_kron_vjp(โˆ‚M4_raw, reshape(d.ฮ”ฬ‚ฮผหขโ‚‚, :, 1), Ine)[1]) + # Decompose โˆ‚M6 โ†’ โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚ฮฃฬ‚แถปโ‚ .+= reshape(_kron_vjp(โˆ‚M6_raw, reshape(vฮฃ, :, 1), Ine)[1], n, n) + + # โ”€โ”€ 4: Scatter local cotangents to global accumulators โ”€โ”€ + โˆ‚๐’โ‚_acc[d.iหข, d.dependencies_in_states_idx] .+= โˆ‚sโ‚_l + โˆ‚๐’โ‚_acc[d.iหข, nโ‚‹+1:size(โˆ‚๐’โ‚_acc, 2)] .+= โˆ‚eโ‚_l + โˆ‚S2f_acc[d.iหข, d.kron_s_s] .+= โˆ‚ssโ‚‚_l + โˆ‚S2f_acc[d.iหข, kron_e_e] .+= โˆ‚eeโ‚‚_l + โˆ‚S2f_acc[d.iหข, d.kron_s_e] .+= โˆ‚seโ‚‚_l + โˆ‚S2f_acc[d.iหข, kron_v_v] .+= โˆ‚vvโ‚‚_l + โˆ‚ฮฃสธโ‚_acc[d.iหข, d.iหข] .+= โˆ‚ฮฃฬ‚แถปโ‚ + โˆ‚ฮฃแถปโ‚‚_acc[d.dependencies_extended_idx, d.dependencies_extended_idx] .+= โˆ‚ฮฃฬ‚แถปโ‚‚ + โˆ‚ฮ”ฮผหขโ‚‚_acc[d.dependencies_in_states_idx] .+= โˆ‚ฮ”ฬ‚ฮผหขโ‚‚_l + end + + # โ”€โ”€ Sub-rrule pullback chain โ”€โ”€ + + # Sโ‚ƒ_full = Sโ‚ƒ * ๐”โ‚ƒ โ†’ โˆ‚Sโ‚ƒ = โˆ‚Sโ‚ƒ_full * ๐”โ‚ƒ' + โˆ‚๐’โ‚ƒ_compressed = โˆ‚S3f_acc * ๐”โ‚ƒ' + + # Third-order solution pullback + so3_grad = so3_pb((โˆ‚๐’โ‚ƒ_compressed, NoTangent())) + if !(so3_grad[2] isa AbstractZero); โˆ‚โˆ‡โ‚_acc .+= so3_grad[2]; end + if !(so3_grad[3] isa AbstractZero); โˆ‚โˆ‡โ‚‚_acc .+= so3_grad[3]; end + if !(so3_grad[4] isa AbstractZero); โˆ‚โˆ‡โ‚ƒ_acc .+= so3_grad[4]; end + if !(so3_grad[5] isa AbstractZero); โˆ‚๐’โ‚_acc .+= so3_grad[5]; end + # so3_grad[6] is now compressed โˆ‚๐’โ‚‚_raw โ€” kept separate + + # Third-order derivatives pullback + โˆ‡โ‚ƒ_grad = โˆ‡โ‚ƒ_pb(โˆ‚โˆ‡โ‚ƒ_acc) + โˆ‚params_โˆ‡โ‚ƒ = โˆ‡โ‚ƒ_grad[2] isa AbstractZero ? zeros(T, np) : โˆ‡โ‚ƒ_grad[2] + if !(โˆ‡โ‚ƒ_grad[3] isa AbstractZero); โˆ‚SS_acc .+= โˆ‡โ‚ƒ_grad[3]; end + + # Convert full-space โˆ‚S2f_acc to compressed and add compressed so3 gradient + โˆ‚S2_raw_acc = โˆ‚S2f_acc * ๐”โ‚‚' + if !(so3_grad[6] isa AbstractZero); โˆ‚S2_raw_acc .+= so3_grad[6]; end + + # Second-order moments pullback + โˆ‚som2 = ( + NoTangent(), # โˆ‚ฮฃสธโ‚‚ + โˆ‚ฮฃแถปโ‚‚_acc, # โˆ‚ฮฃแถปโ‚‚ + โˆ‚ฮผสธโ‚‚_in isa AbstractZero ? NoTangent() : โˆ‚ฮผสธโ‚‚_in, # โˆ‚ฮผสธโ‚‚ + โˆ‚ฮ”ฮผหขโ‚‚_acc, # โˆ‚ฮ”ฮผหขโ‚‚ + NoTangent(), # โˆ‚autocorr (not used) + NoTangent(), # โˆ‚ล_to_ลโ‚‚ (not used) + NoTangent(), # โˆ‚ล_to_yโ‚‚ (not used) + โˆ‚ฮฃสธโ‚_acc, # โˆ‚ฮฃสธโ‚ + NoTangent(), # โˆ‚ฮฃแถปโ‚ + โˆ‚SS_acc, # โˆ‚SS_and_pars + โˆ‚๐’โ‚_acc, # โˆ‚๐’โ‚ + โˆ‚โˆ‡โ‚_acc, # โˆ‚โˆ‡โ‚ + โˆ‚S2_raw_acc, # โˆ‚๐’โ‚‚ (compressed) + โˆ‚โˆ‡โ‚‚_acc, # โˆ‚โˆ‡โ‚‚ + NoTangent(), # โˆ‚slvd + ) + + som2_grad = som2_pb(โˆ‚som2) + โˆ‚params_som2 = som2_grad[2] isa AbstractZero ? zeros(T, np) : som2_grad[2] + + โˆ‚parameters_total = โˆ‚params_som2 .+ โˆ‚params_โˆ‡โ‚ƒ + + return NoTangent(), โˆ‚parameters_total, NoTangent(), NoTangent() + end + + return result, calculate_third_order_moments_with_autocorrelation_pullback +end + + +function rrule(::typeof(calculate_first_order_solution), + โˆ‡โ‚::Matrix{R}, + constants::constants, + workspaces::workspaces, + cache::caches; + opts::CalculationOptions = merge_calculation_options(), + use_fastlapack_qr::Bool = true, + use_fastlapack_lu::Bool = true, + initial_guess::AbstractMatrix{R} = zeros(0,0), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true) where {R <: AbstractFloat} + # Forward pass to compute the output and intermediate values needed for the backward pass + # @timeit_debug timer "Calculate 1st order solution" begin + # @timeit_debug timer "Preprocessing" begin + + T = constants.post_model_macro + idx_constants = ensure_first_order_constants!(constants) + + dynIndex = idx_constants.dyn_index + reverse_dynamic_order = idx_constants.reverse_dynamic_order + comb = idx_constants.comb + future_not_past_and_mixed_in_comb = idx_constants.future_not_past_and_mixed_in_comb + past_not_future_and_mixed_in_comb = idx_constants.past_not_future_and_mixed_in_comb + past_not_future_and_mixed_in_present_but_not_only = idx_constants.past_not_future_and_mixed_in_present_but_not_only + Ir = idx_constants.Ir + + qme_ws = workspaces.first_order + sylv_ws = workspaces.sylvester_1st_order + ensure_sylvester_krylov_buffers!(qme_ws.sylvester, T.nVars, T.nVars) + ensure_sylvester_doubling_buffers!(qme_ws.sylvester, T.nVars, T.nVars) + + ensure_first_order_workspace_buffers!(qme_ws, T, length(dynIndex), length(comb)) + + โˆ‡โ‚Š = @view โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] + โˆ‡โ‚€ = qme_ws.โˆ‡โ‚€ + copyto!(โˆ‡โ‚€, @view(โˆ‡โ‚[:,idx_constants.nabla_zero_cols])) + โˆ‡โ‚‹ = @view โˆ‡โ‚[:,idx_constants.nabla_minus_cols] + โˆ‡ฬ‚โ‚‘ = qme_ws.โˆ‡โ‚‘ + copyto!(โˆ‡ฬ‚โ‚‘, @view(โˆ‡โ‚[:,idx_constants.nabla_e_start:end])) + + # end # timeit_debug + # @timeit_debug timer "Invert โˆ‡โ‚€" begin + + Aโ‚Š = qme_ws.๐€โ‚Š + Aโ‚€ = qme_ws.๐€โ‚€ + Aโ‚‹ = qme_ws.๐€โ‚‹ + โˆ‡โ‚€_present = @view โˆ‡โ‚€[:, T.present_only_idx] + # Legacy readable flow mirrored from primal first-order solver: + # Q = qr!(โˆ‡โ‚€[:, T.present_only_idx]) + # Aโ‚Š = Q.Q' * โˆ‡โ‚Š; Aโ‚€ = Q.Q' * โˆ‡โ‚€; Aโ‚‹ = Q.Q' * โˆ‡โ‚‹ + # The current implementation keeps the same algebra while reusing QR workspaces. + qr_factors, qr_ws = ensure_first_order_fast_qr_workspace!(qme_ws, โˆ‡โ‚€_present) + Q = factorize_qr!(โˆ‡โ‚€_present, qr_factors, qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + + qme_ws.fast_qr_orm_ws_plus, qme_ws.fast_qr_orm_dims_plus = apply_qr_transpose_left!(Aโ‚Š, โˆ‡โ‚Š, Q, + qme_ws.fast_qr_orm_ws_plus, + qme_ws.fast_qr_orm_dims_plus, + qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + qme_ws.fast_qr_orm_ws_zero, qme_ws.fast_qr_orm_dims_zero = apply_qr_transpose_left!(Aโ‚€, โˆ‡โ‚€, Q, + qme_ws.fast_qr_orm_ws_zero, + qme_ws.fast_qr_orm_dims_zero, + qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + qme_ws.fast_qr_orm_ws_minus, qme_ws.fast_qr_orm_dims_minus = apply_qr_transpose_left!(Aโ‚‹, โˆ‡โ‚‹, Q, + qme_ws.fast_qr_orm_ws_minus, + qme_ws.fast_qr_orm_dims_minus, + qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + + # end # timeit_debug + # @timeit_debug timer "Sort matrices" begin + + Aฬƒโ‚Š = qme_ws.๐€ฬƒโ‚Š + โ„’.mul!(Aฬƒโ‚Š, @view(Aโ‚Š[dynIndex,:]), Ir[future_not_past_and_mixed_in_comb,:]) + + Aฬƒโ‚€ = qme_ws.๐€ฬƒโ‚€ + copyto!(Aฬƒโ‚€, @view(Aโ‚€[dynIndex, comb])) + + Aฬƒโ‚‹ = qme_ws.๐€ฬƒโ‚‹ + โ„’.mul!(Aฬƒโ‚‹, @view(Aโ‚‹[dynIndex,:]), Ir[past_not_future_and_mixed_in_comb,:]) + + # end # timeit_debug + # @timeit_debug timer "Quadratic matrix equation solve" begin + + sol, solved = solve_quadratic_matrix_equation(Aฬƒโ‚Š, Aฬƒโ‚€, Aฬƒโ‚‹, constants, workspaces, cache; + initial_guess = initial_guess, + quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, + tol = opts.tol.first_order.ad.qme, + verbose = opts.verbose, + caching = caching) + + if !solved + return (fill(NaN, T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # end # timeit_debug + # @timeit_debug timer "Postprocessing" begin + # @timeit_debug timer "Setup matrices" begin + + sol_compact = @view sol[reverse_dynamic_order, past_not_future_and_mixed_in_comb] + + D = @view sol_compact[end - T.nFuture_not_past_and_mixed + 1:end, :] + + L = @view sol[past_not_future_and_mixed_in_present_but_not_only, past_not_future_and_mixed_in_comb] + + Aฬ„โ‚€แตค = qme_ws.๐€ฬ„โ‚€แตค + copyto!(Aฬ„โ‚€แตค, @view(Aโ‚€[1:T.nPresent_only, T.present_only_idx])) + + Aโ‚Šแตค = qme_ws.๐€โ‚Šแตค + copyto!(Aโ‚Šแตค, @view(Aโ‚Š[1:T.nPresent_only,:])) + + Aฬƒโ‚€แตค = qme_ws.๐€ฬƒโ‚€แตค + copyto!(Aฬƒโ‚€แตค, @view(Aโ‚€[1:T.nPresent_only, T.present_but_not_only_idx])) + + Aโ‚‹แตค = qme_ws.๐€โ‚‹แตค + copyto!(Aโ‚‹แตค, @view(Aโ‚‹[1:T.nPresent_only,:])) + + # end # timeit_debug + # @timeit_debug timer "Invert Aฬ„โ‚€แตค" begin + + qme_ws.fast_lu_ws_a0u, qme_ws.fast_lu_dims_a0u, solved_Aฬ„โ‚€แตค, Aฬ„ฬ‚โ‚€แตค = factorize_lu!(Aฬ„โ‚€แตค, + qme_ws.fast_lu_ws_a0u, + qme_ws.fast_lu_dims_a0u; + use_fastlapack_lu = use_fastlapack_lu) + + if !solved_Aฬ„โ‚€แตค + return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # A = vcat(-(Aฬ„ฬ‚โ‚€แตค \ (Aโ‚Šแตค * D * L + Aฬƒโ‚€แตค * sol[T.dynamic_order,:] + Aโ‚‹แตค)), sol) + if T.nPresent_only > 0 + โ„’.mul!(Aโ‚‹แตค, Aฬƒโ‚€แตค, @view(sol[:,past_not_future_and_mixed_in_comb]), 1, 1) + nโ‚šโ‚‹ = qme_ws.๐งโ‚šโ‚‹ + โ„’.mul!(nโ‚šโ‚‹, Aโ‚Šแตค, D) + โ„’.mul!(Aโ‚‹แตค, nโ‚šโ‚‹, L, 1, 1) + solve_lu_left!(Aฬ„โ‚€แตค, Aโ‚‹แตค, qme_ws.fast_lu_ws_a0u, Aฬ„ฬ‚โ‚€แตค; + use_fastlapack_lu = use_fastlapack_lu) + โ„’.rmul!(Aโ‚‹แตค, -1) + end + + # end # timeit_debug + # end # timeit_debug + # @timeit_debug timer "Exogenous part solution" begin + + expand_future = idx_constants.expand_future + expand_past = idx_constants.expand_past + + ๐’แต— = qme_ws.๐€ + + for i in 1:T.nVars + src = T.reorder[i] + if src <= T.nPresent_only + @views copyto!(๐’แต—[i, :], Aโ‚‹แตค[src, :]) + else + src_idx = src - T.nPresent_only + @views copyto!(๐’แต—[i, :], sol_compact[src_idx, :]) + end + end + + ๐’ฬ‚แต— = qme_ws.sylvester.tmp + โ„’.mul!(๐’ฬ‚แต—, ๐’แต—, expand_past) + + โˆ‡โ‚Š = qme_ws.sylvester.๐€ + โ„’.mul!(โˆ‡โ‚Š, @view(โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed]), expand_future) + + โ„’.mul!(โˆ‡โ‚€, โˆ‡โ‚Š, ๐’ฬ‚แต—, 1, 1) + + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_โˆ‡โ‚€, C = factorize_lu!(โˆ‡โ‚€, + qme_ws.fast_lu_ws_nabla0, + qme_ws.fast_lu_dims_nabla0; + use_fastlapack_lu = use_fastlapack_lu) + + if !solved_โˆ‡โ‚€ + return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + solve_lu_left!(โˆ‡โ‚€, โˆ‡ฬ‚โ‚‘, qme_ws.fast_lu_ws_nabla0, C; + use_fastlapack_lu = use_fastlapack_lu) + โ„’.rmul!(โˆ‡ฬ‚โ‚‘, -1) + + # end # timeit_debug + # end # timeit_debug + + M = qme_ws.sylvester.๐€ยน + fill!(M, zero(R)) + @inbounds for i in axes(M, 1) + M[i, i] = one(R) + end + solve_lu_left!(โˆ‡โ‚€, M, qme_ws.fast_lu_ws_nabla0, C; + use_fastlapack_lu = use_fastlapack_lu) + + tmp2 = qme_ws.sylvester.๐ + โ„’.mul!(tmp2, M', โˆ‡โ‚Š') + โ„’.rmul!(tmp2, -1) + + โˆ‡โ‚‘ = @view โˆ‡โ‚[:,idx_constants.nabla_e_start:end] + + function first_order_solution_pullback(โˆ‚๐’) + # Guard: if the cotangent for the solution matrix is NoTangent + # (e.g. because a downstream filter failure returned all-NoTangent), + # return zero gradients immediately. + if โˆ‚๐’[1] isa Union{NoTangent, AbstractZero} + return NoTangent(), zero(โˆ‡โ‚), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + โˆ‚โˆ‡โ‚ = zero(โˆ‡โ‚) + + โˆ‚๐’แต— = โˆ‚๐’[1][:,1:T.nPast_not_future_and_mixed] + โˆ‚๐’แต‰ = โˆ‚๐’[1][:,T.nPast_not_future_and_mixed + 1:end] + + # Shared sub-expression: W = M' * โˆ‚๐’แต‰ * โˆ‡โ‚‘' * M' + # Use workspace buffers to avoid repeated intermediate allocations. + # t1 = M' * โˆ‚๐’แต‰ (nVars ร— nExo) + t1 = M' * โˆ‚๐’แต‰ # one alloc for nVarsร—nExo + + # โˆ‚โˆ‡โ‚[:,nabla_e_start:end] = -t1 + @views โˆ‚โˆ‡โ‚[:,idx_constants.nabla_e_start:end] .= .-t1 + + # t2 = t1 * โˆ‡โ‚‘' (nVars ร— nVars) โ†’ store in ๐— workspace + t2 = qme_ws.sylvester.๐— + โ„’.mul!(t2, t1, โˆ‡โ‚‘') + + # W = t2 * M' (nVars ร— nVars) โ†’ store in ๐‚_dbl workspace + W = qme_ws.sylvester.๐‚_dbl + โ„’.mul!(W, t2, M') + + @views โˆ‚โˆ‡โ‚[:,idx_constants.nabla_zero_cols] .= W + + # Wp = W * expand_past' (nVars ร— nPast) โ†’ store in view of ๐‚ยน workspace (nVarsร—nVars) + Wp = @view qme_ws.sylvester.๐‚ยน[:, 1:T.nPast_not_future_and_mixed] + โ„’.mul!(Wp, W, expand_past') + + # โˆ‚โˆ‡โ‚[:,1:nFuture] = (Wp * ๐’แต—')[:,future_idx] + # WpSt = Wp * ๐’แต—' (nVars ร— nVars) โ†’ store in ๐‚B workspace + WpSt = qme_ws.sylvester.๐‚B + โ„’.mul!(WpSt, Wp, ๐’แต—') + @views โˆ‚โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] .= WpSt[:,T.future_not_past_and_mixed_idx] + + # โˆ‚๐’แต— += โˆ‡โ‚Š' * Wp (nVars ร— nPast, โˆ‡โ‚Š is nVarsร—nVars, Wp is nVarsร—nPast) + โ„’.mul!(โˆ‚๐’แต—, โˆ‡โ‚Š', Wp, 1, 1) + + tmp1 = qme_ws.sylvester.๐‚ + # tmp1 = M' * โˆ‚๐’แต— * expand_past (nVars ร— nVars) + # t_ms = M' * โˆ‚๐’แต— (nVars ร— nPast) โ†’ reuse Wp (view of ๐‚ยน, same dims) + โ„’.mul!(Wp, M', โˆ‚๐’แต—) + โ„’.mul!(tmp1, Wp, expand_past) + โ„’.lmul!(-1, tmp1) + + ss, solved = solve_sylvester_equation(tmp2, ๐’ฬ‚แต—', tmp1, sylv_ws, + sylvester_algorithm = opts.sylvester_algorithmยฒ, + tol = opts.tol.first_order.ad.sylvester, + verbose = opts.verbose) + + if !solved + NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # ss_Sht = ss * ๐’ฬ‚แต—' (nVars ร— nVars) โ†’ reuse t2 + โ„’.mul!(t2, ss, ๐’ฬ‚แต—') + @views โˆ‚โˆ‡โ‚[:,idx_constants.nabla_zero_cols] .+= t2 + + # ss_Sht_Sht = t2 * ๐’ฬ‚แต—' (nVars ร— nVars) โ†’ reuse W + โ„’.mul!(W, t2, ๐’ฬ‚แต—') + @views โˆ‚โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] .+= W[:,T.future_not_past_and_mixed_idx] + + @views โˆ‚โˆ‡โ‚[:,idx_constants.nabla_minus_cols] .+= ss[:,T.past_not_future_and_mixed_idx] + + return NoTangent(), โˆ‚โˆ‡โ‚, NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + n_rows = size(๐’แต—, 1) + n_cols_A = size(๐’แต—, 2) + n_cols_ฯต = size(โˆ‡ฬ‚โ‚‘, 2) + total_cols = n_cols_A + n_cols_ฯต + + Sโ‚_existing = cache.first_order_solution_matrix + if Sโ‚_existing isa Matrix{R} && size(Sโ‚_existing) == (n_rows, total_cols) + copyto!(@view(Sโ‚_existing[:, 1:n_cols_A]), ๐’แต—) + copyto!(@view(Sโ‚_existing[:, n_cols_A+1:total_cols]), โˆ‡ฬ‚โ‚‘) + ๐’โ‚ = Sโ‚_existing + else + ๐’โ‚ = hcat(๐’แต—, โˆ‡ฬ‚โ‚‘) + cache.first_order_solution_matrix = ๐’โ‚ + end + + if !isempty(parameter_values) + cache.valid_for.first_order_solution = Float64.(parameter_values) + end + + return (๐’โ‚, sol, solved), first_order_solution_pullback +end + +function rrule(::typeof(calculate_second_order_solution), + โˆ‡โ‚::AbstractMatrix{S}, #first order derivatives + โˆ‡โ‚‚::SparseMatrixCSC{S}, #second order derivatives + ๐‘บโ‚::AbstractMatrix{S},#first order solution + constants::constants, + workspaces::workspaces, + cache::caches; + initial_guess::AbstractMatrix{R} = zeros(0,0), + opts::CalculationOptions = merge_calculation_options(), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true) where {S <: Real, R <: Real} + if !(eltype(workspaces.second_order.Sฬ‚) == S) + workspaces.second_order = Higher_order_workspace(T = S) + end + โ„‚ = workspaces.second_order + Mโ‚‚ = constants.second_order + T = constants.post_model_macro + + # Expand compressed hessian to full space for internal computation + โˆ‡โ‚‚_full = โˆ‡โ‚‚ * Mโ‚‚.๐”โˆ‡โ‚‚ + + # @timeit_debug timer "Second order solution - forward" begin + # inspired by Levintal + + # Indices and number of variables + iโ‚Š = T.future_not_past_and_mixed_idx; + iโ‚‹ = T.past_not_future_and_mixed_idx; + + nโ‚‹ = T.nPast_not_future_and_mixed + nโ‚Š = T.nFuture_not_past_and_mixed + nโ‚‘ = T.nExo; + n = T.nVars + nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + + ensure_higher_order_solution_buffers!(โ„‚, n, nโ‚‘โ‚‹) + + # @timeit_debug timer "Setup matrices" begin + + # 1st order solution + ๐’โ‚ = โ„‚.๐’โ‚::Matrix{S} + copyto!(@view(๐’โ‚[:,1:nโ‚‹]), @view(๐‘บโ‚[:,1:nโ‚‹])) + fill!(@view(๐’โ‚[:,nโ‚‹+1]), zero(S)) + copyto!(@view(๐’โ‚[:,nโ‚‹+2:end]), @view(๐‘บโ‚[:,nโ‚‹+1:end])) + # droptol!(๐’โ‚,tol) + + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} + copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹,:]), @view(๐’โ‚[iโ‚‹,:])) + fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1:end,:]), zero(S)) + @inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1,nโ‚‹+1] = one(S) + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0) + + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] + ๐’โ‚ + โ„’.I(nโ‚‘โ‚‹)[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]] + + ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] + zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)] + + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * Mโ‚‚.๐ˆโ‚™โ‚‹ - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] + + # end # timeit_debug + # @timeit_debug timer "Invert matrix" begin + + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = โ„’.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) + + if !โ„’.issuccess(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + if opts.verbose println("Second order solution: inversion failed") end + return (โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + spinv = choose_matrix_format(spinv) + + # end # timeit_debug + # @timeit_debug timer "Setup second order matrices" begin + # @timeit_debug timer "A" begin + + โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š + + A = spinv * โˆ‡โ‚โ‚Š + + # end # timeit_debug + # @timeit_debug timer "C" begin + + kron_compressed = compressed_kronยฒ(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + rowmask = Mโ‚‚.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask, + sparse_preallocation = โ„‚.tmp_sparse_prealloc2) + + term1 = โˆ‡โ‚‚ * kron_compressed + + kron_sigma_compressed = compressed_kronยฒ(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, + rowmask = Mโ‚‚.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask, + colmask = Mโ‚‚.๐›”๐‚โ‚‚_nonempty_row_as_kron_colmask, + sparse_preallocation = โ„‚.tmp_sparse_prealloc3) + + term2 = (โˆ‡โ‚‚ * kron_sigma_compressed) * Mโ‚‚.๐›”cโ‚‚ + + โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = term1 + term2 + + C = spinv * โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน + + # end # timeit_debug + # @timeit_debug timer "B" begin + + # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0) + + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0) + B = compressed_kronยฒ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, sparse_preallocation = โ„‚.tmp_sparse_prealloc1) + Mโ‚‚.๐›”cโ‚‚ + + # end # timeit_debug + # end # timeit_debug + # @timeit_debug timer "Solve sylvester equation" begin + + ๐’โ‚‚, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, + initial_guess = initial_guess, + sylvester_algorithm = opts.sylvester_algorithmยฒ, + tol = opts.tol.second_order.ad.sylvester, + verbose = opts.verbose) + ๐’โ‚‚_stable = copy(๐’โ‚‚) + + # end # timeit_debug + # @timeit_debug timer "Post-process" begin + + if !solved + return (๐’โ‚‚_stable, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # end # timeit_debug + + # spโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t = choose_matrix_format(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹', density_threshold = 1.0) + + # sp๐’โ‚โ‚Šโ•ฑ๐ŸŽt = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ', density_threshold = 1.0) + + ๐›”t = Mโ‚‚.๐›”แต€ + + ๐”โ‚‚t = Mโ‚‚.๐”โ‚‚แต€ + + ๐‚โ‚‚t = Mโ‚‚.๐‚โ‚‚แต€ + + Bt = choose_matrix_format(B', density_threshold = 1.0) + At = choose_matrix_format(A', density_threshold = 1.0) + ๐’โ‚‚_stable_t = choose_matrix_format(๐’โ‚‚_stable', density_threshold = 1.0) + + โˆ‡โ‚‚t = choose_matrix_format(โˆ‡โ‚‚', density_threshold = 1.0) + + # end # timeit_debug + + # Ensure pullback workspaces are properly sized + if size(โ„‚.โˆ‚โˆ‡โ‚‚) != size(โˆ‡โ‚‚) + โ„‚.โˆ‚โˆ‡โ‚‚ = zeros(S, size(โˆ‡โ‚‚)) + end + if size(โ„‚.โˆ‚โˆ‡โ‚) != size(โˆ‡โ‚) + โ„‚.โˆ‚โˆ‡โ‚ = zeros(S, size(โˆ‡โ‚)) + end + if size(โ„‚.โˆ‚๐’โ‚) != size(๐’โ‚) + โ„‚.โˆ‚๐’โ‚ = zeros(S, size(๐’โ‚)) + end + if size(โ„‚.โˆ‚spinv) != size(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€) + โ„‚.โˆ‚spinv = zeros(S, size(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€)) + end + if size(โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) != size(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = zeros(S, size(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)) + end + if size(โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ) != size(๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ = zeros(S, size(๐’โ‚โ‚Šโ•ฑ๐ŸŽ)) + end + if size(โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) != size(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = zeros(S, size(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹)) + end + + function second_order_solution_pullback(โˆ‚๐’โ‚‚_solved) + # @timeit_debug timer "Second order solution - pullback" begin + + # @timeit_debug timer "Preallocate" begin + # Use workspaces and fill with zeros instead of allocating new arrays + โˆ‚โˆ‡โ‚‚ = โ„‚.โˆ‚โˆ‡โ‚‚; fill!(โˆ‚โˆ‡โ‚‚, zero(S)) + โˆ‚โˆ‡โ‚ = โ„‚.โˆ‚โˆ‡โ‚; fill!(โˆ‚โˆ‡โ‚, zero(S)) + โˆ‚๐’โ‚ = โ„‚.โˆ‚๐’โ‚; fill!(โˆ‚๐’โ‚, zero(S)) + โˆ‚spinv = โ„‚.โˆ‚spinv; fill!(โˆ‚spinv, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, zero(S)) + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, zero(S)) + โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹; fill!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, zero(S)) + + # end # timeit_debug + + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚_solved[1] + + if size(โˆ‚๐’โ‚‚, 2) == size(๐’โ‚‚_stable, 2) + nothing + elseif size(โˆ‚๐’โ‚‚, 2) == size(Mโ‚‚.๐”โ‚‚, 2) + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ * ๐”โ‚‚t + else + throw(DimensionMismatch("second_order_solution_pullback: expected โˆ‚๐’โ‚‚ to have $(size(๐’โ‚‚_stable, 2)) (compressed) or $(size(Mโ‚‚.๐”โ‚‚, 2)) (full) columns, got $(size(โˆ‚๐’โ‚‚, 2)).")) + end + + # @timeit_debug timer "Sylvester" begin + if โ„’.norm(โˆ‚๐’โ‚‚) < opts.tol.second_order.ad.sylvester.acceptance_tol + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + โˆ‚C, solved = solve_sylvester_equation(A', B', โˆ‚๐’โ‚‚, โ„‚.sylvester_workspace, + sylvester_algorithm = opts.sylvester_algorithmยฒ, + tol = opts.tol.second_order.ad.sylvester, + verbose = opts.verbose) + + if !solved + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # end # timeit_debug + + # @timeit_debug timer "Matmul" begin + + โˆ‚C = choose_matrix_format(โˆ‚C) # Dense + + โˆ‚A = โˆ‚C * Bt * ๐’โ‚‚_stable_t + + โˆ‚B = ๐’โ‚‚_stable_t * At * โˆ‚C + + # B = (Mโ‚‚.๐”โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + Mโ‚‚.๐”โ‚‚ * Mโ‚‚.๐›”) * Mโ‚‚.๐‚โ‚‚ + โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = ๐”โ‚‚t * โˆ‚B * ๐‚โ‚‚t + + # end # timeit_debug + + # @timeit_debug timer "Kron adjoint" begin + + fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + + # end # timeit_debug + + # @timeit_debug timer "Matmul2" begin + + # A = spinv * โˆ‡โ‚โ‚Š + โˆ‚โˆ‡โ‚โ‚Š = spinv' * โˆ‚A + โˆ‚spinv += โˆ‚A * โˆ‡โ‚โ‚Š' + + # โˆ‡โ‚โ‚Š = sparse(โˆ‡โ‚[:,1:nโ‚Š] * spdiagm(ones(n))[iโ‚Š,:]) + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] += โˆ‚โˆ‡โ‚โ‚Š * โ„’.I(n)[:,iโ‚Š] + + # C = spinv * โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน + โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = spinv' * โˆ‚C + + โˆ‚spinv += โˆ‚C * โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน' + + # end # timeit_debug + + # @timeit_debug timer "Matmul3" begin + + โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = choose_matrix_format(โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน, density_threshold = 1.0) + + โˆ‚term2 = โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน * Mโ‚‚.๐›”cโ‚‚' + + โˆ‚โˆ‡โ‚‚ += โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน * kron_compressed' + โˆ‚โˆ‡โ‚‚ += โˆ‚term2 * kron_sigma_compressed' + + # end # timeit_debug + + # @timeit_debug timer "Matmul4" begin + + โˆ‚kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ = โˆ‡โ‚‚t * โˆ‚term2 + + # end # timeit_debug + + # @timeit_debug timer "Kron adjoint 2" begin + + compressed_kronยฒ_pullback!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, + tol = opts.tol.second_order.droptol, rowmask = Mโ‚‚.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask, + colmask = Mโ‚‚.๐›”๐‚โ‚‚_nonempty_row_as_kron_colmask) + + # end # timeit_debug + + โˆ‚kronโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = โˆ‡โ‚‚t * โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน + + # @timeit_debug timer "Kron adjoint 3" begin + + compressed_kronยฒ_pullback!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โˆ‚kronโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + tol = opts.tol.second_order.droptol, rowmask = Mโ‚‚.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask) + + # end # timeit_debug + + # @timeit_debug timer "Matmul5" begin + + # spinv = sparse(inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€)) + โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -spinv' * โˆ‚spinv * spinv' + + # โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.diagm(ones(n))[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] * ๐’โ‚[iโ‚Š,1:nโ‚‹]' + โˆ‚โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ + + โˆ‚๐’โ‚[iโ‚Š,1:nโ‚‹] -= โˆ‡โ‚[:,1:nโ‚Š]' * โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] + + # ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] + # zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)]; + โˆ‚๐’โ‚[iโ‚Š,:] += โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] + + ###### โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] + # โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = [โ„’.I(size(๐’โ‚,1))[iโ‚Š,:] * ๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ + # ๐’โ‚ + # spdiagm(ones(nโ‚‘โ‚‹))[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]]; + โˆ‚๐’โ‚ += โ„’.I(size(๐’โ‚,1))[:,iโ‚Š] * โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[1:length(iโ‚Š),:] * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + โˆ‚๐’โ‚ += โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[length(iโ‚Š) .+ (1:size(๐’โ‚,1)),:] + + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ += ๐’โ‚' * โ„’.I(size(๐’โ‚,1))[:,iโ‚Š] * โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[1:length(iโ‚Š),:] + + # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) spdiagm(ones(nโ‚‘ + 1))[1,:] zeros(nโ‚‘ + 1, nโ‚‘)]; + โˆ‚๐’โ‚[iโ‚‹,:] += โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:length(iโ‚‹), :] + + # ๐’โ‚ = [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]] + โˆ‚๐‘บโ‚ = [โˆ‚๐’โ‚[:,1:nโ‚‹] โˆ‚๐’โ‚[:,nโ‚‹+2:end]] + + # end # timeit_debug + + # end # timeit_debug + + return NoTangent(), โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚๐‘บโ‚, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + + if solved + if ๐’โ‚‚ isa Matrix{S} && cache.second_order_solution isa Matrix{S} && size(cache.second_order_solution) == size(๐’โ‚‚) + copyto!(cache.second_order_solution, ๐’โ‚‚) + elseif ๐’โ‚‚ isa SparseMatrixCSC{S, Int} && cache.second_order_solution isa SparseMatrixCSC{S, Int} && + size(cache.second_order_solution) == size(๐’โ‚‚) && + cache.second_order_solution.colptr == ๐’โ‚‚.colptr && + cache.second_order_solution.rowval == ๐’โ‚‚.rowval + copyto!(cache.second_order_solution.nzval, ๐’โ‚‚.nzval) + else + cache.second_order_solution = ๐’โ‚‚ + end + if !isempty(parameter_values) + cache.valid_for.second_order_solution = Float64.(parameter_values) + end + empty!(cache.valid_for.pruned_second_order_solution) + end + + # return (sparse(๐’โ‚‚ * Mโ‚‚.๐”โ‚‚), solved), second_order_solution_pullback + return (๐’โ‚‚_stable, solved), second_order_solution_pullback +end + + +# โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• +# Kron-adjoint helper kernels (fill_kron_adjoint!, mul_fill_kron_adjoint!, etc.) +# โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + +function fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + โˆ‚X::AbstractSparseMatrix{R}, + A::AbstractMatrix{TA}, + B::AbstractMatrix{TB}) where {R <: Real, TA <: Real, TB <: Real} + @assert size(โˆ‚A) == size(A) + @assert size(โˆ‚B) == size(B) + @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + n1, m1 = size(B) + n2, m2 = size(A) + + # Access the sparse matrix internal representation + if โˆ‚X isa SparseMatrixCSC + colptr = โˆ‚X.colptr # Column pointers + rowval = โˆ‚X.rowval # Row indices of non-zeros + nzval = โˆ‚X.nzval # Non-zero values + else + colptr = โˆ‚X.A.colptr # Column pointers + rowval = โˆ‚X.A.rowval # Row indices of non-zeros + nzval = โˆ‚X.A.nzval # Non-zero values + end + + # Iterate over columns of โˆ‚X + for col in 1:size(โˆ‚X, 2) + # Iterate over the non-zeros in this column + for idx in colptr[col]:(colptr[col + 1] - 1) + row = rowval[idx] + val = nzval[idx] + + @inbounds begin + i = (row - 1) รท n2 + 1 + k = (row - 1) % n2 + 1 + j = (col - 1) รท m2 + 1 + l = (col - 1) % m2 + 1 + + # Update โˆ‚B and โˆ‚A + โˆ‚A[k,l] += B[i,j] * val + โˆ‚B[i,j] += A[k,l] * val + end + end + end +end + + +function fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + โˆ‚X::DenseMatrix{R}, + A::AbstractMatrix{TA}, + B::AbstractMatrix{TB}) where {R <: Real, TA <: Real, TB <: Real} + @assert size(โˆ‚A) == size(A) + @assert size(โˆ‚B) == size(B) + @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + reโˆ‚X = reshape(โˆ‚X, + size(A,1), + size(B,1), + size(A,2), + size(B,2)) + + ei = 1 + for e in eachslice(reโˆ‚X; dims = (1,3)) + @inbounds โˆ‚A[ei] += โ„’.dot(B,e) + ei += 1 + end + + ei = 1 + for e in eachslice(reโˆ‚X; dims = (2,4)) + @inbounds โˆ‚B[ei] += โ„’.dot(A,e) + ei += 1 + end +end + + +function fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + โˆ‚X::DenseMatrix{R}, + A::SparseMatrixCSC{TA, Int}, + B::SparseMatrixCSC{TB, Int}) where {R <: Real, TA <: Real, TB <: Real} + @assert size(โˆ‚A) == size(A) + @assert size(โˆ‚B) == size(B) + @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + n1, m1 = size(B) + n2, m2 = size(A) + + A_colptr = A.colptr + A_rowval = A.rowval + A_nzval = A.nzval + + B_colptr = B.colptr + B_rowval = B.rowval + B_nzval = B.nzval + + # โˆ‚A[k,l] += ฮฃ_{i,j} B[i,j] * โˆ‚X[(i-1)n2 + k, (j-1)m2 + l] + @inbounds for l in 1:m2 + base_col_l = l + for k in 1:n2 + acc = zero(R) + for j in 1:m1 + b_start = B_colptr[j] + b_stop = B_colptr[j + 1] - 1 + col_idx = (j - 1) * m2 + base_col_l + for bidx in b_start:b_stop + i = B_rowval[bidx] + row_idx = (i - 1) * n2 + k + acc += R(B_nzval[bidx]) * โˆ‚X[row_idx, col_idx] + end + end + โˆ‚A[k, l] += acc + end + end + + # โˆ‚B[i,j] += ฮฃ_{k,l} A[k,l] * โˆ‚X[(i-1)n2 + k, (j-1)m2 + l] + @inbounds for j in 1:m1 + b_start = B_colptr[j] + b_stop = B_colptr[j + 1] - 1 + for bidx in b_start:b_stop + i = B_rowval[bidx] + row_base = (i - 1) * n2 + col_base = (j - 1) * m2 + acc = zero(R) + for l in 1:m2 + a_start = A_colptr[l] + a_stop = A_colptr[l + 1] - 1 + col_idx = col_base + l + for aidx in a_start:a_stop + k = A_rowval[aidx] + row_idx = row_base + k + acc += R(A_nzval[aidx]) * โˆ‚X[row_idx, col_idx] + end + end + โˆ‚B[i, j] += acc + end + end +end + + + +function fill_kron_adjoint!(โˆ‚A::V, โˆ‚B::V, โˆ‚X::V, A::V, B::V) where V <: Vector{<: Real} + @assert size(โˆ‚A) == size(A) + @assert size(โˆ‚B) == size(B) + @assert length(โˆ‚X) == length(B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + reโˆ‚X = reshape(โˆ‚X, + length(A), + length(B)) + + ei = 1 + for e in eachslice(reโˆ‚X; dims = 1) + @inbounds โˆ‚A[ei] += โ„’.dot(B,e) + ei += 1 + end + + ei = 1 + for e in eachslice(reโˆ‚X; dims = 2) + @inbounds โˆ‚B[ei] += โ„’.dot(A,e) + ei += 1 + end +end + + +function fill_kron_adjoint_โˆ‚B!(โˆ‚X::AbstractSparseMatrix{R}, โˆ‚B::AbstractArray{S}, A::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} + @assert length(โˆ‚X) == length(โˆ‚B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + n1, m1 = size(โˆ‚B) + n2, m2 = size(A) + + # Access the sparse matrix internal representation + colptr = โˆ‚X.colptr # Column pointers + rowval = โˆ‚X.rowval # Row indices of non-zeros + nzval = โˆ‚X.nzval # Non-zero values + + # Iterate over columns of โˆ‚X + for col in 1:size(โˆ‚X, 2) + # Iterate over the non-zeros in this column + for idx in colptr[col]:(colptr[col + 1] - 1) + row = rowval[idx] + val = nzval[idx] + + @inbounds begin + i = (row - 1) รท n2 + 1 + k = (row - 1) % n2 + 1 + j = (col - 1) รท m2 + 1 + l = (col - 1) % m2 + 1 + + # Update โˆ‚B and โˆ‚A + โˆ‚B[i,j] += A[k,l] * val + end + end + end +end + + + +function fill_kron_adjoint_โˆ‚B!(โˆ‚X::AbstractSparseMatrix{R}, โˆ‚B::Vector{S}, A::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} + @assert length(โˆ‚X) == length(โˆ‚B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + n1 = length(โˆ‚B) + n2 = size(A,1) + # println("hello") + # Precompute constants + const_n1n2 = n1 * n2 + + # Access the sparse matrix internal representation + colptr = โˆ‚X.colptr # Column pointers + rowval = โˆ‚X.rowval # Row indices of non-zeros + nzval = โˆ‚X.nzval # Non-zero values + + # Iterate over columns of โˆ‚X + for col in 1:size(โˆ‚X, 2) + # Iterate over the non-zeros in this column + for idx in colptr[col]:(colptr[col + 1] - 1) + row = rowval[idx] + val = nzval[idx] + + linear_idx = (col - 1) * size(โˆ‚X, 1) + row + + @inbounds begin + i = (linear_idx - 1) % n1 + 1 + k = ((linear_idx - 1) รท n1) % n2 + 1 + l = ((linear_idx - 1) รท const_n1n2) + 1 + + # Update โˆ‚B and โˆ‚A + โˆ‚B[i] += A[k,l] * val + end + end + end +end + + + +function fill_kron_adjoint_โˆ‚B!(โˆ‚X::DenseMatrix{R}, โˆ‚B::Vector{S}, A::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} + @assert length(โˆ‚X) == length(โˆ‚B) * length(A) "โˆ‚X must have the same length as kron(B,A)" + + reโˆ‚X = reshape(โˆ‚X, + size(A,1), + length(โˆ‚B), + size(A,2)) + + ei = 1 + for e in eachslice(reโˆ‚X; dims = 2) + @inbounds โˆ‚B[ei] += โ„’.dot(A,e) + ei += 1 + end +end + + +function fill_kron_adjoint_โˆ‚A!(โˆ‚X::DenseMatrix{R}, โˆ‚A::Vector{S}, B::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} + @assert length(โˆ‚X) == length(โˆ‚A) * length(B) "โˆ‚X must have the same length as kron(B,A)" + + reโˆ‚X = reshape(โˆ‚X, + length(โˆ‚A), + size(B,1), + size(B,2)) + + ei = 1 + for e in eachslice(reโˆ‚X; dims = 1) + @inbounds โˆ‚A[ei] += โ„’.dot(B,e) + ei += 1 + end +end + + +function fill_kron_adjoint_โˆ‚A!(โˆ‚X::AbstractSparseMatrix{R}, โˆ‚A::AbstractMatrix{S}, B::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} + @assert length(โˆ‚X) == length(B) * length(โˆ‚A) "โˆ‚X must have the same length as kron(B,A)" + + n1, m1 = size(B) + n2 = size(โˆ‚A,1) + + # Precompute constants + const_n1n2 = n1 * n2 + const_n1n2m1 = n1 * n2 * m1 + + # Access the sparse matrix internal representation + colptr = โˆ‚X.colptr # Column pointers + rowval = โˆ‚X.rowval # Row indices of non-zeros + nzval = โˆ‚X.nzval # Non-zero values + + # Iterate over columns of โˆ‚X + for col in 1:size(โˆ‚X, 2) + # Iterate over the non-zeros in this column + for idx in colptr[col]:(colptr[col + 1] - 1) + row = rowval[idx] + val = nzval[idx] + + linear_idx = (col - 1) * size(โˆ‚X, 1) + row + + @inbounds begin + i = (linear_idx - 1) % n1 + 1 + k = ((linear_idx - 1) รท n1) % n2 + 1 + j = ((linear_idx - 1) รท const_n1n2) % m1 + 1 + l = ((linear_idx - 1) รท const_n1n2m1) + 1 + + # Update โˆ‚B and โˆ‚A + โˆ‚A[k,l] += B[i,j] * val + end + end + end +end + + +# Fused operation: computes fill_kron_adjoint!(โˆ‚A, โˆ‚B, M1*M2, A, B) +# without materializing the full product M1*M2. +# +# M1*M2 has shape (n1*n2, m1*m2) where kron(B,A) has the same shape, +# B is (n1,m1) and A is (n2,m2). +# +# Processes column-blocks of M1*M2 to keep memory usage at O(n1*n2*block_size) +# instead of O(n1*n2*m1*m2). +function mul_fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::AbstractMatrix{TA}, + B::AbstractMatrix{TB}; + tol::Real = 0.0, + block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0)) where {R <: Real, TA <: Real, TB <: Real} + n2, m2 = size(A) + n1, m1 = size(B) + + @assert size(M1, 1) == n1 * n2 "M1 rows ($(size(M1,1))) must equal n1*n2 ($(n1*n2))" + @assert size(M2, 2) == m1 * m2 "M2 cols ($(size(M2,2))) must equal m1*m2 ($(m1*m2))" + @assert size(M1, 2) == size(M2, 1) "M1 cols ($(size(M1,2))) must equal M2 rows ($(size(M2,1)))" + + nrows = n1 * n2 + + # Process one j-block at a time: columns (j-1)*m2+1 : j*m2 + # Each block produces a (nrows ร— m2) matrix, reshaped to (n2, n1, m2) + if size(block, 1) == nrows && size(block, 2) >= m2 + blk = view(block, :, 1:m2) + else + blk = Matrix{R}(undef, nrows, m2) + end + + @inbounds for j in 1:m1 + col_start = (j - 1) * m2 + 1 + col_end = j * m2 + # blk = M1 * M2[:, col_start:col_end] โ€” shape (n1*n2, m2) + โ„’.mul!(blk, M1, view(M2, :, col_start:col_end)) + + # Reshape blk to (n2, n1, m2) for accumulation + re_blk = reshape(blk, n2, n1, m2) + + # โˆ‚A[:,l] += re_blk[:,i,l] * B[i,j] for all i โ†’ โˆ‚A[:,l] += ฮฃ_i B[i,j]*re_blk[:,i,l] + # = re_blk[:,:,l] * B[:,j] + for l in 1:m2 + slice_l = view(re_blk, :, :, l) # (n2, n1) + for i in 1:n1 + bij = B[i, j] + if abs(bij) > tol + for k in 1:n2 + โˆ‚A[k, l] += bij * slice_l[k, i] + end + end + end + end + + # โˆ‚B[i,j] += ฮฃ_{k,l} A[k,l] * re_blk[k,i,l] = ฮฃ_l dot(A[:,l], re_blk[:,i,l]) + for i in 1:n1 + acc = zero(R) + for l in 1:m2 + for k in 1:n2 + acc += A[k, l] * re_blk[k, i, l] + end + end + โˆ‚B[i, j] += acc + end + end +end + + +# Sparse-factor variant: when A and B are sparse, exploit nzrange for dot products +function mul_fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::SparseMatrixCSC{TA, Int}, + B::SparseMatrixCSC{TB, Int}; + tol::Real = 0.0, + block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0)) where {R <: Real, TA <: Real, TB <: Real} + n2, m2 = size(A) + n1, m1 = size(B) + + @assert size(M1, 1) == n1 * n2 + @assert size(M2, 2) == m1 * m2 + @assert size(M1, 2) == size(M2, 1) + + nrows = n1 * n2 + + if size(block, 1) == nrows && size(block, 2) >= m2 + blk = view(block, :, 1:m2) + else + blk = Matrix{R}(undef, nrows, m2) + end + + B_colptr = B.colptr + B_rowval = B.rowval + B_nzval = B.nzval + A_colptr = A.colptr + A_rowval = A.rowval + A_nzval = A.nzval + + @inbounds for j in 1:m1 + col_start = (j - 1) * m2 + 1 + col_end = j * m2 + โ„’.mul!(blk, M1, view(M2, :, col_start:col_end)) + + re_blk = reshape(blk, n2, n1, m2) + + # โˆ‚A[k,l] += B[i,j] * re_blk[k,i,l] โ€” only iterate nonzero B[i,j] + b_start = B_colptr[j] + b_stop = B_colptr[j + 1] - 1 + for l in 1:m2 + for bidx in b_start:b_stop + i = B_rowval[bidx] + bij = R(B_nzval[bidx]) + for k in 1:n2 + โˆ‚A[k, l] += bij * re_blk[k, i, l] + end + end + end + + # โˆ‚B[i,j] += ฮฃ_{k,l} A[k,l] * re_blk[k,i,l] โ€” only iterate nonzero A[k,l] + for bidx in b_start:b_stop + i = B_rowval[bidx] + acc = zero(R) + for l in 1:m2 + for aidx in A_colptr[l]:(A_colptr[l + 1] - 1) + k = A_rowval[aidx] + acc += R(A_nzval[aidx]) * re_blk[k, i, l] + end + end + โˆ‚B[i, j] += acc + end + end +end + + +# Mixed-sparsity variant: A is sparse, B is dense +function mul_fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::SparseMatrixCSC{TA, Int}, + B::AbstractMatrix{TB}; + tol::Real = 0.0, + block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0)) where {R <: Real, TA <: Real, TB <: Real} + n2, m2 = size(A) + n1, m1 = size(B) + + @assert size(M1, 1) == n1 * n2 + @assert size(M2, 2) == m1 * m2 + @assert size(M1, 2) == size(M2, 1) + + nrows = n1 * n2 + + if size(block, 1) == nrows && size(block, 2) >= m2 + blk = view(block, :, 1:m2) + else + blk = Matrix{R}(undef, nrows, m2) + end + + A_colptr = A.colptr + A_rowval = A.rowval + A_nzval = A.nzval + + @inbounds for j in 1:m1 + col_start = (j - 1) * m2 + 1 + col_end = j * m2 + โ„’.mul!(blk, M1, view(M2, :, col_start:col_end)) + + re_blk = reshape(blk, n2, n1, m2) + + # โˆ‚A[k,l] += B[i,j] * re_blk[k,i,l] โ€” B is dense, use iszero guard + for l in 1:m2 + for i in 1:n1 + bij = B[i, j] + if abs(bij) > tol + for k in 1:n2 + โˆ‚A[k, l] += bij * re_blk[k, i, l] + end + end + end + end + + # โˆ‚B[i,j] += ฮฃ_{k,l} A[k,l] * re_blk[k,i,l] โ€” A is sparse, use nzrange + for i in 1:n1 + acc = zero(R) + for l in 1:m2 + for aidx in A_colptr[l]:(A_colptr[l + 1] - 1) + k = A_rowval[aidx] + acc += R(A_nzval[aidx]) * re_blk[k, i, l] + end + end + โˆ‚B[i, j] += acc + end + end +end + + +# Mixed-sparsity variant: A is dense, B is sparse +function mul_fill_kron_adjoint!(โˆ‚A::AbstractMatrix{R}, + โˆ‚B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::AbstractMatrix{TA}, + B::SparseMatrixCSC{TB, Int}; + tol::Real = 0.0, + block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0)) where {R <: Real, TA <: Real, TB <: Real} + n2, m2 = size(A) + n1, m1 = size(B) + + @assert size(M1, 1) == n1 * n2 + @assert size(M2, 2) == m1 * m2 + @assert size(M1, 2) == size(M2, 1) + + nrows = n1 * n2 + + if size(block, 1) == nrows && size(block, 2) >= m2 + blk = view(block, :, 1:m2) + else + blk = Matrix{R}(undef, nrows, m2) + end + + B_colptr = B.colptr + B_rowval = B.rowval + B_nzval = B.nzval + + @inbounds for j in 1:m1 + col_start = (j - 1) * m2 + 1 + col_end = j * m2 + โ„’.mul!(blk, M1, view(M2, :, col_start:col_end)) + + re_blk = reshape(blk, n2, n1, m2) + + # โˆ‚A[k,l] += B[i,j] * re_blk[k,i,l] โ€” B is sparse, only iterate nonzero B[i,j] + b_start = B_colptr[j] + b_stop = B_colptr[j + 1] - 1 + for l in 1:m2 + for bidx in b_start:b_stop + i = B_rowval[bidx] + bij = R(B_nzval[bidx]) + for k in 1:n2 + โˆ‚A[k, l] += bij * re_blk[k, i, l] + end + end + end + + # โˆ‚B[i,j] += ฮฃ_{k,l} A[k,l] * re_blk[k,i,l] โ€” iterate all i (โˆ‚B is dense) + for i in 1:n1 + acc = zero(R) + for l in 1:m2 + for k in 1:n2 + akl = A[k, l] + if abs(akl) > tol + acc += akl * re_blk[k, i, l] + end + end + end + โˆ‚B[i, j] += acc + end + end +end + + +# Variant that computes fill_kron_adjoint_โˆ‚A! for both the identity and a permuted +# version of โˆ‚X in a single sparse iteration pass. +# +# Equivalent to: +# fill_kron_adjoint_โˆ‚A!(โˆ‚X, โˆ‚A, B) +# fill_kron_adjoint_โˆ‚A!(Pโ‚— * โˆ‚X * Pแตฃ, โˆ‚A, B) +# but avoids materializing the permuted matrix. +# +# perm_row and perm_col are integer vectors representing the row/column permutations +# such that (Pโ‚— * โˆ‚X * Pแตฃ)[perm_row[row], perm_col[col]] = โˆ‚X[row, col]. +# Accumulates the โˆ‚A adjoint from โˆ‚X + Pโ‚โ‚— * โˆ‚X * Pโ‚แตฃ where โˆ‚X is the cotangent +# of kron(B, A) and Pโ‚ is the (2,1,3) tensor-axis swap on the dยณ row/column space +# (d = n_A = size(โˆ‚A,1)). The permutation is baked in โ€” no external vectors needed. +# +# Requires n_B = n_Aยฒ and m_B = m_Aยฒ (i.e. B is the dยฒร—dยฒ outer factor). +function fill_kron_adjoint_โˆ‚A_with_perm!(โˆ‚X::AbstractSparseMatrix{R}, + โˆ‚A::AbstractMatrix{S}, + B::AbstractMatrix{T}) where {R <: Real, S <: Real, T <: Real} + @assert length(โˆ‚X) == length(โˆ‚A) * length(B) "โˆ‚X must have the same length as kron(B,A)" + + # Convention: kron(B, A) โ€” A is inner (fastest-varying), B is outer + # Same decomposition as fill_kron_adjoint! reshape(โˆ‚X, n_A, n_B, m_A, m_B) + n_A = size(โˆ‚A, 1) + n_B = size(B, 1) + m_A = size(โˆ‚A, 2) + + @assert n_B == n_A * n_A "fill_kron_adjoint_โˆ‚A_with_perm! requires n_B == n_Aยฒ for the (2,1,3) axis swap" + @assert size(B, 2) == m_A * m_A "fill_kron_adjoint_โˆ‚A_with_perm! requires m_B == m_Aยฒ for the (2,1,3) axis swap" + + const_nAnB = n_A * n_B + const_nAnBmA = n_A * n_B * m_A + nrows = size(โˆ‚X, 1) + + colptr = โˆ‚X.colptr + rowval = โˆ‚X.rowval + nzval = โˆ‚X.nzval + + @inbounds for col in 1:size(โˆ‚X, 2) + for idx in colptr[col]:(colptr[col + 1] - 1) + row = rowval[idx] + val = nzval[idx] + + # --- Identity contribution (linear-index decomposition) --- + L = (col - 1) * nrows + row - 1 + i_A = L % n_A + 1 + i_B = (L รท n_A) % n_B + 1 + j_A = (L รท const_nAnB) % m_A + 1 + j_B = (L รท const_nAnBmA) + 1 + โˆ‚A[i_A, j_A] += B[i_B, j_B] * val + + # --- (2,1,3) axis-swap contribution --- + # The outer index i_B (1-based) encodes two sub-axes of size n_A: + # kโ‚‚ = (i_B-1) % n_A, kโ‚ƒ = (i_B-1) รท n_A + # Swapping axis 1 (i_A) with axis 2 (kโ‚‚) gives: + i_Ap = (i_B - 1) % n_A + 1 + i_Bp = (i_A - 1) + ((i_B - 1) รท n_A) * n_A + 1 + j_Ap = (j_B - 1) % m_A + 1 + j_Bp = (j_A - 1) + ((j_B - 1) รท m_A) * m_A + 1 + โˆ‚A[i_Ap, j_Ap] += B[i_Bp, j_Bp] * val + end + end +end + + +# Fused variant of fill_kron_adjoint_โˆ‚A_with_perm! that processes M1 * M2 +# in column blocks without materializing the full product. +# +# Equivalent to: +# fill_kron_adjoint_โˆ‚A_with_perm!(sparse(M1 * M2), โˆ‚A, B) +# but avoids allocating the (n_Aยณ ร— m_Aยณ) intermediate. +# +# Requires n_B = n_Aยฒ and m_B = m_Aยฒ (same as fill_kron_adjoint_โˆ‚A_with_perm!). +function mul_fill_kron_adjoint_โˆ‚A_with_perm!(M1::AbstractMatrix, + M2::AbstractMatrix, + โˆ‚A::AbstractMatrix{S}, + B::AbstractMatrix{T}; + block::AbstractMatrix{S} = Matrix{S}(undef, size(M1, 1), 0)) where {S <: Real, T <: Real} + n_A = size(โˆ‚A, 1) + m_A = size(โˆ‚A, 2) + n_B = size(B, 1) + m_B = size(B, 2) + + @assert n_B == n_A * n_A "mul_fill_kron_adjoint_โˆ‚A_with_perm! requires n_B == n_Aยฒ" + @assert m_B == m_A * m_A "mul_fill_kron_adjoint_โˆ‚A_with_perm! requires m_B == m_Aยฒ" + @assert size(M1, 1) == n_A * n_B "M1 rows ($(size(M1,1))) must equal n_A * n_B ($(n_A * n_B))" + @assert size(M2, 2) == m_A * m_B "M2 cols ($(size(M2,2))) must equal m_A * m_B ($(m_A * m_B))" + @assert size(M1, 2) == size(M2, 1) "M1 cols ($(size(M1,2))) must equal M2 rows ($(size(M2,1)))" + + nrows = n_A * n_B # = n_Aยณ + + if size(block, 1) == nrows && size(block, 2) >= m_A + blk = view(block, :, 1:m_A) + else + blk = Matrix{S}(undef, nrows, m_A) + end + + @inbounds for j in 1:m_B # j = j_B (outer column index of B) + col_start = (j - 1) * m_A + 1 + col_end = j * m_A + โ„’.mul!(blk, M1, view(M2, :, col_start:col_end)) + + # Pre-compute the fixed permuted column index for j_B = j + # (2,1,3) axis swap: j_Ap depends only on j, not on j_A + j_Ap_fixed = (j - 1) % m_A + 1 + + for j_A in 1:m_A + # (2,1,3) axis swap: j_Bp depends on both j_A and j + j_Bp = (j_A - 1) + ((j - 1) รท m_A) * m_A + 1 + + for row in 1:nrows + val = blk[row, j_A] + + # Decompose row into (i_A, i_B) for kron(B, A) convention + i_A = (row - 1) % n_A + 1 + i_B = (row - 1) รท n_A + 1 + + # Identity contribution + โˆ‚A[i_A, j_A] += B[i_B, j] * val + + # (2,1,3) axis-swap contribution + i_Ap = (i_B - 1) % n_A + 1 + i_Bp = (i_A - 1) + ((i_B - 1) รท n_A) * n_A + 1 + โˆ‚A[i_Ap, j_Ap_fixed] += B[i_Bp, j_Bp] * val + end + end + end +end + + +# Sparse-B variant of mul_fill_kron_adjoint_โˆ‚A_with_perm! that exploits B's sparsity. +# When B is ultra-sparse (e.g. ฯƒ with ~nโ‚‘ nonzeros in nโ‚‘โ‚‹ยฒ ร— nโ‚‘โ‚‹ยฒ), +# this skips ~99.7% of work by iterating only nzrange columns. +function mul_fill_kron_adjoint_โˆ‚A_with_perm!(M1::AbstractMatrix, + M2::AbstractMatrix, + โˆ‚A::AbstractMatrix{S}, + B::SparseMatrixCSC{TB, Int}; + block::AbstractMatrix{S} = Matrix{S}(undef, size(M1, 1), 0)) where {S <: Real, TB <: Real} + n_A = size(โˆ‚A, 1) + m_A = size(โˆ‚A, 2) + n_B = size(B, 1) + m_B = size(B, 2) + + @assert n_B == n_A * n_A "mul_fill_kron_adjoint_โˆ‚A_with_perm! requires n_B == n_Aยฒ" + @assert m_B == m_A * m_A "mul_fill_kron_adjoint_โˆ‚A_with_perm! requires m_B == m_Aยฒ" + @assert size(M1, 1) == n_A * n_B "M1 rows ($(size(M1,1))) must equal n_A * n_B ($(n_A * n_B))" + @assert size(M2, 2) == m_A * m_B "M2 cols ($(size(M2,2))) must equal m_A * m_B ($(m_A * m_B))" + @assert size(M1, 2) == size(M2, 1) "M1 cols ($(size(M1,2))) must equal M2 rows ($(size(M2,1)))" + + nrows = n_A * n_B # = n_Aยณ + + B_colptr = B.colptr + B_rowval = SparseArrays.rowvals(B) + B_nzval = nonzeros(B) + + # Precompute which B columns have nonzeros for fast skip checks + has_nz = falses(m_B) + @inbounds for col in 1:m_B + has_nz[col] = B_colptr[col] < B_colptr[col + 1] + end + + if size(block, 1) == nrows && size(block, 2) >= m_A + blk = view(block, :, 1:m_A) + else + blk = Matrix{S}(undef, nrows, m_A) + end + + @inbounds for j in 1:m_B # j = j_B (outer column index of B) + # Check if this j contributes anything: + # Identity path: B[:,j] has nonzeros + # Perm path: for each j_A, B[:, j_Bp(j_A, j)] has nonzeros + need_blk = has_nz[j] + if !need_blk + j_div = (j - 1) รท m_A + for j_A in 1:m_A + j_Bp = (j_A - 1) + j_div * m_A + 1 + if has_nz[j_Bp] + need_blk = true + break + end + end + end + need_blk || continue + + col_start = (j - 1) * m_A + 1 + col_end = j * m_A + โ„’.mul!(blk, M1, view(M2, :, col_start:col_end)) + + # Pre-compute for (2,1,3) axis swap + j_Ap_fixed = (j - 1) % m_A + 1 + j_div = (j - 1) รท m_A + + # Identity contribution: iterate nonzeros of B[:, j] + for bidx in B_colptr[j]:(B_colptr[j + 1] - 1) + i_B = B_rowval[bidx] + b_val = S(B_nzval[bidx]) + # i_A = (row-1) % n_A + 1 for row = (i_B-1)*n_A + 1 : i_B*n_A + row_start = (i_B - 1) * n_A + for j_A in 1:m_A + for i_A in 1:n_A + โˆ‚A[i_A, j_A] += b_val * blk[row_start + i_A, j_A] + end + end + end + + # (2,1,3) axis-swap contribution: for each j_A, iterate nonzeros of B[:, j_Bp] + for j_A in 1:m_A + j_Bp = (j_A - 1) + j_div * m_A + 1 + for bidx in B_colptr[j_Bp]:(B_colptr[j_Bp + 1] - 1) + i_Bp = B_rowval[bidx] + b_val = S(B_nzval[bidx]) + # Reverse-map: i_Ap = (i_B-1) % n_A + 1, but here i_Bp encodes + # i_Bp = (i_A-1) + ((i_B-1) รท n_A) * n_A + 1 + # So: i_A = (i_Bp-1) % n_A + 1, block_offset = (i_Bp-1) รท n_A + i_A = (i_Bp - 1) % n_A + 1 + block_k3 = (i_Bp - 1) รท n_A # = (i_B-1) รท n_A = kโ‚ƒ - 1 + + # The identity i_Ap = (i_B-1) % n_A + 1 = kโ‚‚ + # and row = (i_B-1)*n_A + i_A where i_B = kโ‚‚ + kโ‚ƒ*n_A + 1 + # We need to iterate over all kโ‚‚ (= i_Ap's corresponding i_B values) + # For a given i_Bp, we have i_A and block_k3 fixed. + # i_Ap = kโ‚‚ + 1 ranges over 1:n_A, with i_B = kโ‚‚ + block_k3*n_A + 1 + # and row = (i_B-1)*n_A + i_A = (kโ‚‚ + block_k3*n_A)*n_A + i_A + for k2 in 0:(n_A - 1) + i_Ap = k2 + 1 + row = (k2 + block_k3 * n_A) * n_A + i_A + โˆ‚A[i_Ap, j_Ap_fixed] += b_val * blk[row, j_A] + end + end + end + end +end + +# Helper: adjoint of compressed_kron(A, ฯƒ; tol) w.r.t. A and ฯƒ. +# Forward contribution for each sorted output column triple (ฮฑโ‰ฅฮฒโ‰ฅฮณ) is: +# Y[row,col] += A[i,ฮฑ] * ฯƒ[(j-1)*nแตฃ+k, (ฮฒ-1)*nแถœ+ฮณ] +# where row is obtained by sorting (i,j,k) into iโ‚โ‰ฅjโ‚โ‰ฅkโ‚. +function compressed_kron_pullback_2arg!(โˆ‚A::AbstractMatrix{T}, + โˆ‚ฯƒ::AbstractMatrix{T}, + โˆ‚Y::AbstractMatrix{T}, + A::AbstractMatrix{TA}, + ฯƒ::AbstractMatrix{Tฯƒ}; + tol::AbstractFloat = eps()) where {T <: Real, TA <: Real, Tฯƒ <: Real} + + nแตฃ, nแถœ = size(A) + size(ฯƒ) == (nแตฃ^2, nแถœ^2) || throw(DimensionMismatch("ฯƒ must be $(nแตฃ^2)ร—$(nแถœ^2), got $(size(ฯƒ))")) + + As = A isa SparseMatrixCSC ? A : sparse(A) + ฯƒs = ฯƒ isa SparseMatrixCSC ? ฯƒ : sparse(ฯƒ) + + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + + ranges_A = Vector{UnitRange{Int}}(undef, nแถœ) + ranges_ฯƒ = Vector{UnitRange{Int}}(undef, nแถœ^2) + @inbounds for col in 1:nแถœ + ranges_A[col] = SparseArrays.nzrange(As, col) + end + @inbounds for col in 1:(nแถœ^2) + ranges_ฯƒ[col] = SparseArrays.nzrange(ฯƒs, col) + end + + @inbounds for ฮฑ in 1:nแถœ + rng_A = ranges_A[ฮฑ] + isempty(rng_A) && continue + + for ฮฒ in 1:ฮฑ + for ฮณ in 1:ฮฒ + ฯƒ_col = (ฮฒ - 1) * nแถœ + ฮณ + rng_ฯƒ = ranges_ฯƒ[ฯƒ_col] + isempty(rng_ฯƒ) && continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + for pA in rng_A + i = rv_A[pA] + a_val = nzv_A[pA] + + for pฯƒ in rng_ฯƒ + s = rv_ฯƒ[pฯƒ] + ฯƒ_val = nzv_ฯƒ[pฯƒ] + + val = a_val * ฯƒ_val + abs(val) > tol || continue + + j = (s - 1) รท nแตฃ + 1 + k = (s - 1) % nแตฃ + 1 + + iโ‚ = i; jโ‚ = j; kโ‚ = k + if iโ‚ < jโ‚; iโ‚, jโ‚ = jโ‚, iโ‚; end + if jโ‚ < kโ‚; jโ‚, kโ‚ = kโ‚, jโ‚; end + if iโ‚ < jโ‚; iโ‚, jโ‚ = jโ‚, iโ‚; end + + row = (iโ‚ - 1) * iโ‚ * (iโ‚ + 1) รท 6 + (jโ‚ - 1) * jโ‚ รท 2 + kโ‚ + g = โˆ‚Y[row, col] + abs(g) <= tol && continue + + โˆ‚A[i, ฮฑ] += g * ฯƒ_val + โˆ‚ฯƒ[s, ฯƒ_col] += g * a_val + end + end + end + end + end + + return +end + +# Helper: adjoint of compressed_permuted_mixed_kron(A, ฯƒ; tol) w.r.t. A and ฯƒ. +function compressed_permuted_mixed_kron_pullback!(โˆ‚A::AbstractMatrix{T}, + โˆ‚ฯƒ::AbstractMatrix{T}, + โˆ‚Y::AbstractMatrix{T}, + A::AbstractMatrix{TA}, + ฯƒ::AbstractMatrix{Tฯƒ}; + tol::AbstractFloat = eps()) where {T <: Real, TA <: Real, Tฯƒ <: Real} + + nr, nc = size(A) + size(ฯƒ) == (nr^2, nc^2) || throw(DimensionMismatch("ฯƒ must be $(nr^2)ร—$(nc^2), got $(size(ฯƒ))")) + + As = A isa SparseMatrixCSC ? A : sparse(A) + ฯƒs = ฯƒ isa SparseMatrixCSC ? ฯƒ : sparse(ฯƒ) + + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + + ranges_A = Vector{UnitRange{Int}}(undef, nc) + ranges_ฯƒ = Vector{UnitRange{Int}}(undef, nc^2) + @inbounds for col in 1:nc + ranges_A[col] = SparseArrays.nzrange(As, col) + end + @inbounds for col in 1:(nc^2) + ranges_ฯƒ[col] = SparseArrays.nzrange(ฯƒs, col) + end + + G = Matrix(โˆ‚Y) + + @inbounds for ฮฑ in 1:nc + rng_Aฮฑ = ranges_A[ฮฑ] + for ฮฒ in 1:ฮฑ + rng_Aฮฒ = ranges_A[ฮฒ] + for ฮณ in 1:ฮฒ + rng_Aฮณ = ranges_A[ฮณ] + + ฯƒ_col_ฮฒฮณ = (ฮฒ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮณ = (ฮฑ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮฒ = (ฮฑ - 1) * nc + ฮฒ + + rng_ฯƒฮฒฮณ = ranges_ฯƒ[ฯƒ_col_ฮฒฮณ] + rng_ฯƒฮฑฮณ = ranges_ฯƒ[ฯƒ_col_ฮฑฮณ] + rng_ฯƒฮฑฮฒ = ranges_ฯƒ[ฯƒ_col_ฮฑฮฒ] + + has_t1 = !isempty(rng_Aฮฑ) && !isempty(rng_ฯƒฮฒฮณ) + has_t2 = !isempty(rng_Aฮฒ) && !isempty(rng_ฯƒฮฑฮณ) + has_t3 = !isempty(rng_Aฮณ) && !isempty(rng_ฯƒฮฑฮฒ) + + (has_t1 || has_t2 || has_t3) || continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + if has_t1 + for ia in rng_Aฮฑ + p = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฒฮณ + qr = rv_ฯƒ[is] + q = (qr - 1) รท nr + 1 + r = qr - (q - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = G[row, col] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[p, ฮฑ] += g * ฯƒ_val + โˆ‚ฯƒ[qr, ฯƒ_col_ฮฒฮณ] += g * a_val + end + end + end + + if has_t2 + for ia in rng_Aฮฒ + q = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฑฮณ + pr = rv_ฯƒ[is] + p = (pr - 1) รท nr + 1 + r = pr - (p - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = G[row, col] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[q, ฮฒ] += g * ฯƒ_val + โˆ‚ฯƒ[pr, ฯƒ_col_ฮฑฮณ] += g * a_val + end + end + end + + if has_t3 + for ia in rng_Aฮณ + r = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฑฮฒ + pq = rv_ฯƒ[is] + p = (pq - 1) รท nr + 1 + q = pq - (p - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = G[row, col] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[r, ฮณ] += g * ฯƒ_val + โˆ‚ฯƒ[pq, ฯƒ_col_ฮฑฮฒ] += g * a_val + end + end + end + end + end + end + + return +end + + +# โˆ‚A-only variant: skips โˆ‚ฯƒ accumulation (matches fill_kron_adjoint_โˆ‚A! pattern). +# Use when the โˆ‚ฯƒ output is discarded (e.g. B-pullback for Sylvester). +function compressed_permuted_mixed_kron_pullback_โˆ‚A!(โˆ‚A::AbstractMatrix{T}, + โˆ‚Y::AbstractMatrix{T}, + A::AbstractMatrix{TA}, + ฯƒ::AbstractMatrix{Tฯƒ}; + tol::AbstractFloat = eps()) where {T <: Real, TA <: Real, Tฯƒ <: Real} + + nr, nc = size(A) + size(ฯƒ) == (nr^2, nc^2) || throw(DimensionMismatch("ฯƒ must be $(nr^2)ร—$(nc^2), got $(size(ฯƒ))")) + + As = A isa SparseMatrixCSC ? A : sparse(A) + ฯƒs = ฯƒ isa SparseMatrixCSC ? ฯƒ : sparse(ฯƒ) + + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + + ranges_A = Vector{UnitRange{Int}}(undef, nc) + ranges_ฯƒ = Vector{UnitRange{Int}}(undef, nc^2) + @inbounds for col in 1:nc + ranges_A[col] = SparseArrays.nzrange(As, col) + end + @inbounds for col in 1:(nc^2) + ranges_ฯƒ[col] = SparseArrays.nzrange(ฯƒs, col) + end + + G = Matrix(โˆ‚Y) + + @inbounds for ฮฑ in 1:nc + rng_Aฮฑ = ranges_A[ฮฑ] + for ฮฒ in 1:ฮฑ + rng_Aฮฒ = ranges_A[ฮฒ] + for ฮณ in 1:ฮฒ + rng_Aฮณ = ranges_A[ฮณ] + + ฯƒ_col_ฮฒฮณ = (ฮฒ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮณ = (ฮฑ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮฒ = (ฮฑ - 1) * nc + ฮฒ + + rng_ฯƒฮฒฮณ = ranges_ฯƒ[ฯƒ_col_ฮฒฮณ] + rng_ฯƒฮฑฮณ = ranges_ฯƒ[ฯƒ_col_ฮฑฮณ] + rng_ฯƒฮฑฮฒ = ranges_ฯƒ[ฯƒ_col_ฮฑฮฒ] + + has_t1 = !isempty(rng_Aฮฑ) && !isempty(rng_ฯƒฮฒฮณ) + has_t2 = !isempty(rng_Aฮฒ) && !isempty(rng_ฯƒฮฑฮณ) + has_t3 = !isempty(rng_Aฮณ) && !isempty(rng_ฯƒฮฑฮฒ) + + (has_t1 || has_t2 || has_t3) || continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + if has_t1 + for ia in rng_Aฮฑ + p = rv_A[ia] + for is in rng_ฯƒฮฒฮณ + qr = rv_ฯƒ[is] + q = (qr - 1) รท nr + 1 + r = qr - (q - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = G[row, col] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[p, ฮฑ] += g * ฯƒ_val + end + end + end + + if has_t2 + for ia in rng_Aฮฒ + q = rv_A[ia] + for is in rng_ฯƒฮฑฮณ + pr = rv_ฯƒ[is] + p = (pr - 1) รท nr + 1 + r = pr - (p - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = G[row, col] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[q, ฮฒ] += g * ฯƒ_val + end + end + end + + if has_t3 + for ia in rng_Aฮณ + r = rv_A[ia] + for is in rng_ฯƒฮฑฮฒ + pq = rv_ฯƒ[is] + p = (pq - 1) รท nr + 1 + q = pq - (p - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = G[row, col] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[r, ฮณ] += g * ฯƒ_val + end + end + end + end + end + end + + return +end + + +# Fused variant: computes g_col = M1 * M2[:, col] lazily per (ฮฑ,ฮฒ,ฮณ) triple +# instead of materializing the full โˆ‚Y = M1 * M2 matrix. +# Equivalent to: +# compressed_permuted_mixed_kron_pullback!(โˆ‚A, โˆ‚ฯƒ, M1 * M2, A, ฯƒ; tol) +# but avoids the n_compressedยณ ร— n_compressedยณ allocation. +function mul_compressed_permuted_mixed_kron_pullback!(โˆ‚A::AbstractMatrix{T}, + โˆ‚ฯƒ::AbstractMatrix{T}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::AbstractMatrix{TA}, + ฯƒ::AbstractMatrix{Tฯƒ}; + tol::AbstractFloat = eps()) where {T <: Real, TA <: Real, Tฯƒ <: Real} + + nr, nc = size(A) + size(ฯƒ) == (nr^2, nc^2) || throw(DimensionMismatch("ฯƒ must be $(nr^2)ร—$(nc^2), got $(size(ฯƒ))")) + + As = A isa SparseMatrixCSC ? A : sparse(A) + ฯƒs = ฯƒ isa SparseMatrixCSC ? ฯƒ : sparse(ฯƒ) + + rv_A = SparseArrays.rowvals(As) + nzv_A = nonzeros(As) + rv_ฯƒ = SparseArrays.rowvals(ฯƒs) + nzv_ฯƒ = nonzeros(ฯƒs) + + ranges_A = Vector{UnitRange{Int}}(undef, nc) + ranges_ฯƒ = Vector{UnitRange{Int}}(undef, nc^2) + @inbounds for col in 1:nc + ranges_A[col] = SparseArrays.nzrange(As, col) + end + @inbounds for col in 1:(nc^2) + ranges_ฯƒ[col] = SparseArrays.nzrange(ฯƒs, col) + end + + g_col = Vector{T}(undef, size(M1, 1)) + + @inbounds for ฮฑ in 1:nc + rng_Aฮฑ = ranges_A[ฮฑ] + for ฮฒ in 1:ฮฑ + rng_Aฮฒ = ranges_A[ฮฒ] + for ฮณ in 1:ฮฒ + rng_Aฮณ = ranges_A[ฮณ] + + ฯƒ_col_ฮฒฮณ = (ฮฒ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮณ = (ฮฑ - 1) * nc + ฮณ + ฯƒ_col_ฮฑฮฒ = (ฮฑ - 1) * nc + ฮฒ + + rng_ฯƒฮฒฮณ = ranges_ฯƒ[ฯƒ_col_ฮฒฮณ] + rng_ฯƒฮฑฮณ = ranges_ฯƒ[ฯƒ_col_ฮฑฮณ] + rng_ฯƒฮฑฮฒ = ranges_ฯƒ[ฯƒ_col_ฮฑฮฒ] + + has_t1 = !isempty(rng_Aฮฑ) && !isempty(rng_ฯƒฮฒฮณ) + has_t2 = !isempty(rng_Aฮฒ) && !isempty(rng_ฯƒฮฑฮณ) + has_t3 = !isempty(rng_Aฮณ) && !isempty(rng_ฯƒฮฑฮฒ) + + (has_t1 || has_t2 || has_t3) || continue + + col = (ฮฑ - 1) * ฮฑ * (ฮฑ + 1) รท 6 + (ฮฒ - 1) * ฮฒ รท 2 + ฮณ + + # Compute g_col = M1 * M2[:, col] lazily for this triple + โ„’.mul!(g_col, M1, view(M2, :, col)) + + if has_t1 + for ia in rng_Aฮฑ + p = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฒฮณ + qr = rv_ฯƒ[is] + q = (qr - 1) รท nr + 1 + r = qr - (q - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = g_col[row] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[p, ฮฑ] += g * ฯƒ_val + โˆ‚ฯƒ[qr, ฯƒ_col_ฮฒฮณ] += g * a_val + end + end + end + + if has_t2 + for ia in rng_Aฮฒ + q = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฑฮณ + pr = rv_ฯƒ[is] + p = (pr - 1) รท nr + 1 + r = pr - (p - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = g_col[row] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[q, ฮฒ] += g * ฯƒ_val + โˆ‚ฯƒ[pr, ฯƒ_col_ฮฑฮณ] += g * a_val + end + end + end + + if has_t3 + for ia in rng_Aฮณ + r = rv_A[ia] + a_val = nzv_A[ia] + for is in rng_ฯƒฮฑฮฒ + pq = rv_ฯƒ[is] + p = (pq - 1) รท nr + 1 + q = pq - (p - 1) * nr + + i1 = p + j1 = q + k1 = r + if i1 < j1 + i1, j1 = j1, i1 + end + if j1 < k1 + j1, k1 = k1, j1 + end + if i1 < j1 + i1, j1 = j1, i1 + end + + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = g_col[row] + abs(g) <= tol && continue + + ฯƒ_val = nzv_ฯƒ[is] + โˆ‚A[r, ฮณ] += g * ฯƒ_val + โˆ‚ฯƒ[pq, ฯƒ_col_ฮฑฮฒ] += g * a_val + end + end + end + end + end + end + + return +end + + +# Helper: adjoint of compressed_kronยฒ(X; rowmask, colmask) w.r.t. X. +# Forward value at (row(i1,j1), col(i2,j2)): (X[i1,i2]*X[j1,j2] + X[i1,j2]*X[j1,i2]) / divisor, +# where divisor = 2 if i1 == j1 else 1, and only masked rows/cols are materialized. +function compressed_kronยฒ_pullback!(โˆ‚X::AbstractMatrix{T}, + โˆ‚Y::AbstractMatrix{T}, + X::AbstractMatrix{T}; + tol::Real = 0.0, + rowmask::Vector{Int} = Int[], + colmask::Vector{Int} = Int[]) where T <: Real + Xd = X isa DenseMatrix ? X : collect(X) + n_rows, n_cols = size(Xd) + + m2_rows = n_rows * (n_rows + 1) รท 2 + m2_cols = n_cols * (n_cols + 1) รท 2 + + if rowmask == Int[0] || colmask == Int[0] + return + end + + norowmask = length(rowmask) == 0 + nocolmask = length(colmask) == 0 + rowmask_lookup = norowmask ? BitVector() : falses(m2_rows) + colmask_lookup = nocolmask ? BitVector() : falses(m2_cols) + + if !norowmask + @inbounds for r in rowmask + if 1 <= r <= m2_rows + rowmask_lookup[r] = true + end + end + end + + if !nocolmask + @inbounds for c in colmask + if 1 <= c <= m2_cols + colmask_lookup[c] = true + end + end + end + + for i1 in 1:n_rows, j1 in 1:i1 + row = (i1 - 1) * i1 รท 2 + j1 + (norowmask || rowmask_lookup[row]) || continue + divisor = i1 == j1 ? 2 : 1 + + for i2 in 1:n_cols, j2 in 1:i2 + col = (i2 - 1) * i2 รท 2 + j2 + (nocolmask || colmask_lookup[col]) || continue + + g = โˆ‚Y[row, col] + abs(g) <= tol && continue + g_d = g / divisor + + @inbounds aii = Xd[i1, i2] + @inbounds aij = Xd[i1, j2] + @inbounds aji = Xd[j1, i2] + @inbounds ajj = Xd[j1, j2] + + โˆ‚X[i1, i2] += g_d * ajj + โˆ‚X[j1, j2] += g_d * aii + โˆ‚X[i1, j2] += g_d * aji + โˆ‚X[j1, i2] += g_d * aij + end + end +end + + +# Helper: adjoint of compressed_kronยณ(X) w.r.t. X. +# Forward: out[row,col] = (aii*(ajj*akk + ajk*akj) + aij*(aji*akk + ajk*aki) + aik*(aji*akj + ajj*aki)) / divisor +# where row โ†” (i1โ‰ฅj1โ‰ฅk1) and col โ†” (i2โ‰ฅj2โ‰ฅk2) and a_pq = X[p,q]. +function compressed_kronยณ_pullback!(โˆ‚X::AbstractMatrix{T}, โˆ‚Y::AbstractMatrix{T}, X::AbstractMatrix{T}; tol::Real = 0.0) where T <: Real + Xd = X isa DenseMatrix ? X : collect(X) + n_rows, n_cols = size(Xd) + # Unlike the forward pass, the pullback must iterate over ALL row/column + # indices, not just nonzero ones. The gradient at a zero entry X[r,c] can + # be non-zero because โˆ‚(X[i]*X[j]*X[k])/โˆ‚X[i] = X[j]*X[k] which is + # generically non-zero even when X[i]=0. + # However, we can skip columns that have no stored entries in sparse โˆ‚Y. + sparse_cols = if โˆ‚Y isa SparseMatrixCSC + colmask = falses(size(โˆ‚Y, 2)) + @inbounds for col in 1:size(โˆ‚Y, 2) + colmask[col] = โˆ‚Y.colptr[col] < โˆ‚Y.colptr[col + 1] + end + colmask + else + trues(size(โˆ‚Y, 2)) + end + for i2 in 1:n_cols, j2 in 1:i2 + for k2 in 1:j2 + col = (i2 - 1) * i2 * (i2 + 1) รท 6 + (j2 - 1) * j2 รท 2 + k2 + sparse_cols[col] || continue + for i1 in 1:n_rows, j1 in 1:i1 + @inbounds for k1 in 1:j1 + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = โˆ‚Y[row, col] + abs(g) <= tol && continue + # divisor for row symmetry + if i1 == j1 + divisor = (j1 == k1) ? 6 : 2 + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end + g_d = g / divisor + aii = Xd[i1, i2]; aij = Xd[i1, j2]; aik = Xd[i1, k2] + aji = Xd[j1, i2]; ajj = Xd[j1, j2]; ajk = Xd[j1, k2] + aki = Xd[k1, i2]; akj = Xd[k1, j2]; akk = Xd[k1, k2] + โˆ‚X[i1, i2] += g_d * (ajj * akk + ajk * akj) + โˆ‚X[i1, j2] += g_d * (aji * akk + ajk * aki) + โˆ‚X[i1, k2] += g_d * (aji * akj + ajj * aki) + โˆ‚X[j1, i2] += g_d * (aij * akk + aik * akj) + โˆ‚X[j1, j2] += g_d * (aii * akk + aik * aki) + โˆ‚X[j1, k2] += g_d * (aij * aki + aii * akj) + โˆ‚X[k1, i2] += g_d * (aij * ajk + aik * ajj) + โˆ‚X[k1, j2] += g_d * (aik * aji + aii * ajk) + โˆ‚X[k1, k2] += g_d * (aii * ajj + aij * aji) + end + end + end + end +end + +# Fused variant: computes g_col = M1 * M2[:, col] lazily per (i2,j2,k2) triple +# instead of materializing the full โˆ‚Y = M1 * M2 matrix. +# Equivalent to: +# compressed_kronยณ_pullback!(โˆ‚X, M1 * M2, X) +# but avoids the n_compressedยณ ร— n_compressedยณ allocation. +function mul_compressed_kronยณ_pullback!(โˆ‚X::AbstractMatrix{T}, + M1::AbstractMatrix, + M2::AbstractMatrix, + X::AbstractMatrix{T}; + tol::Real = 0.0) where T <: Real + Xd = X isa DenseMatrix ? X : collect(X) + n_rows, n_cols = size(Xd) + + g_col = Vector{T}(undef, size(M1, 1)) + + for i2 in 1:n_cols, j2 in 1:i2 + for k2 in 1:j2 + col = (i2 - 1) * i2 * (i2 + 1) รท 6 + (j2 - 1) * j2 รท 2 + k2 + + # Compute g_col = M1 * M2[:, col] lazily for this triple + โ„’.mul!(g_col, M1, view(M2, :, col)) + + for i1 in 1:n_rows, j1 in 1:i1 + @inbounds for k1 in 1:j1 + row = (i1 - 1) * i1 * (i1 + 1) รท 6 + (j1 - 1) * j1 รท 2 + k1 + g = g_col[row] + abs(g) <= tol && continue + # divisor for row symmetry + if i1 == j1 + divisor = (j1 == k1) ? 6 : 2 + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end + g_d = g / divisor + aii = Xd[i1, i2]; aij = Xd[i1, j2]; aik = Xd[i1, k2] + aji = Xd[j1, i2]; ajj = Xd[j1, j2]; ajk = Xd[j1, k2] + aki = Xd[k1, i2]; akj = Xd[k1, j2]; akk = Xd[k1, k2] + โˆ‚X[i1, i2] += g_d * (ajj * akk + ajk * akj) + โˆ‚X[i1, j2] += g_d * (aji * akk + ajk * aki) + โˆ‚X[i1, k2] += g_d * (aji * akj + ajj * aki) + โˆ‚X[j1, i2] += g_d * (aij * akk + aik * akj) + โˆ‚X[j1, j2] += g_d * (aii * akk + aik * aki) + โˆ‚X[j1, k2] += g_d * (aij * aki + aii * akj) + โˆ‚X[k1, i2] += g_d * (aij * ajk + aik * ajj) + โˆ‚X[k1, j2] += g_d * (aik * aji + aii * ajk) + โˆ‚X[k1, k2] += g_d * (aii * ajj + aij * aji) + end + end + end + end +end + +# ===================================================================================== +# Third-order solution rrule (correctness-first, allocating version) +# ===================================================================================== + +function rrule(::typeof(calculate_third_order_solution), + โˆ‡โ‚::AbstractMatrix{S}, + โˆ‡โ‚‚::SparseMatrixCSC{S}, + โˆ‡โ‚ƒ::SparseMatrixCSC{S}, + ๐‘บโ‚::AbstractMatrix{S}, + ๐’โ‚‚::AbstractMatrix{S}, + constants::constants, + workspaces::workspaces, + cache::caches; + initial_guess::AbstractMatrix{R} = zeros(0,0), + opts::CalculationOptions = merge_calculation_options(), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true) where {S <: Real, R <: Real} + + # --- workspace / constants --------------------------------------------------- + if !(eltype(workspaces.third_order.ลœ) == S) + workspaces.third_order = Higher_order_workspace(T = S) + end + โ„‚ = workspaces.third_order + Mโ‚‚ = constants.second_order + Mโ‚ƒ = constants.third_order + T = constants.post_model_macro + + # Expand compressed inputs to full space for internal computation + โˆ‡โ‚‚ = โˆ‡โ‚‚ * Mโ‚‚.๐”โˆ‡โ‚‚ + ๐’โ‚‚ = sparse(๐’โ‚‚ * Mโ‚‚.๐”โ‚‚)::SparseMatrixCSC{S, Int} + + iโ‚Š = T.future_not_past_and_mixed_idx + iโ‚‹ = T.past_not_future_and_mixed_idx + nโ‚‹ = T.nPast_not_future_and_mixed + nโ‚Š = T.nFuture_not_past_and_mixed + nโ‚‘ = T.nExo + n = T.nVars + nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + + ensure_higher_order_solution_buffers!(โ„‚, n, nโ‚‘โ‚‹) + + initial_guess_sylv = if length(initial_guess) == 0 + zeros(S, 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{S} ? initial_guess : Matrix{S}(initial_guess) + else + zeros(S, 0, 0) + end + + # --- forward pass (mirrors the primal, but stores intermediates) --------------- + + # 1st-order solution with zero-column + ๐’โ‚ = โ„‚.๐’โ‚::Matrix{S} + copyto!(@view(๐’โ‚[:,1:nโ‚‹]), @view(๐‘บโ‚[:,1:nโ‚‹])) + fill!(@view(๐’โ‚[:,nโ‚‹+1]), zero(S)) + copyto!(@view(๐’โ‚[:,nโ‚‹+2:end]), @view(๐‘บโ‚[:,nโ‚‹+1:end])) + + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} + copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹,:]), @view(๐’โ‚[iโ‚‹,:])) + fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1:end,:]), zero(S)) + @inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1,nโ‚‹+1] = one(S) + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) + + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] + ๐’โ‚ + โ„’.I(nโ‚‘โ‚‹)[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]] + + ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:]; zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)] + ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) + + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * Mโ‚‚.๐ˆโ‚™โ‚‹ - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] + + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = โ„’.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) + + if !โ„’.issuccess(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + return (โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) + spinv = choose_matrix_format(spinv) + + โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š + + A = spinv * โˆ‡โ‚โ‚Š + + # --- B matrix ----------------------------------------------------------------- + kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + + B = compressed_permuted_mixed_kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”, + sparse_preallocation = โ„‚.tmp_sparse_prealloc7) + + B += compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.third_order.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc1) + + # --- ๐—โ‚ƒ (C-matrix ingredients) ----------------------------------------------- + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = @views [(๐’โ‚‚ * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)])[iโ‚Š,:] + ๐’โ‚‚ + zeros(nโ‚‹ + nโ‚‘, nโ‚‘โ‚‹^2)] + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, density_threshold = 0.0, min_length = 10, tol = opts.tol.third_order.droptol) + + ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚Š,:]; zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹^2)] + + aux = Mโ‚ƒ.๐’๐ * โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ + + S1p0_kron_sigma = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›” + tmpkron22 = compressed_permuted_mixed_kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + S1p0_kron_sigma, + sparse_preallocation = โ„‚.tmp_sparse_prealloc6) + + ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) + + โˆ‡โ‚โ‚Š = choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) + + ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚‹,:]; zeros(size(๐’โ‚)[2] - nโ‚‹, nโ‚‘โ‚‹^2)] + + # Terms (a)+(b): โˆ‡โ‚‚ * kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) * [tmpkron2 + ๐โ‚โ‚— * tmpkron2 * ๐โ‚แตฃ] * ๐๐‚โ‚ƒ + tmpkron2 = โ„’.kron(Mโ‚‚.๐›”, choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.third_order.droptol)) + D_ab = (tmpkron2 + Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ) * Mโ‚ƒ.๐๐‚โ‚ƒ + ๐—โ‚ƒ = mat_mult_kron(โˆ‡โ‚‚, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ), D_ab, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc2) + + # Term (c): โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, โŽธ๐’โ‚‚k..โŽน) * ๐๐‚โ‚ƒ + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc3) + + # Term (d): โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ*๐›”) * ๐๐‚โ‚ƒ + S2p0_sigma = ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›” + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, collect(S2p0_sigma), Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + + # Term (e): โˆ‡โ‚โ‚Š * ๐’โ‚‚ * kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) * ๐๐‚โ‚ƒ + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.third_order.droptol) + mm_๐’โ‚‚_kron = mat_mult_kron(๐’โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + ๐—โ‚ƒ += โˆ‡โ‚โ‚Š * mm_๐’โ‚‚_kron * Mโ‚ƒ.๐๐‚โ‚ƒ + + ๐—โ‚ƒ += โˆ‡โ‚ƒ * tmpkron22 + + # Compute compressed_kronยณ(aux) WITHOUT rowmask: the pullback needs โˆ‚โˆ‡โ‚ƒ at ALL + # positions (including currently-zero columns of โˆ‡โ‚ƒ) so that gradients flow + # correctly through calculate_third_order_derivatives back to parameters. + ck3_aux_mat = compressed_kronยณ(aux, rowmask = Mโ‚ƒ.โˆ‡โ‚ƒ_rowmask, tol = opts.tol.third_order.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc5) + ck3_aux = โˆ‡โ‚ƒ * ck3_aux_mat + ๐—โ‚ƒ += ck3_aux + + C = spinv * ๐—โ‚ƒ + + # --- solve Sylvester Aยท๐’โ‚ƒยทB + C = ๐’โ‚ƒ ---------------------------------------- + ๐’โ‚ƒ, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, + initial_guess = initial_guess_sylv, + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.third_order.ad.sylvester, + verbose = opts.verbose) + + ๐’โ‚ƒ = choose_matrix_format(๐’โ‚ƒ, multithreaded = false, tol = opts.tol.third_order.droptol) + ๐’โ‚ƒ_stable = copy(๐’โ‚ƒ) + + if !solved + return (๐’โ‚ƒ_stable, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # cache update (same as primal) + if ๐’โ‚ƒ_stable isa Matrix{S} && cache.third_order_solution isa Matrix{S} && size(cache.third_order_solution) == size(๐’โ‚ƒ_stable) + copyto!(cache.third_order_solution, ๐’โ‚ƒ_stable) + elseif ๐’โ‚ƒ_stable isa SparseMatrixCSC{S, Int} && cache.third_order_solution isa SparseMatrixCSC{S, Int} && + size(cache.third_order_solution) == size(๐’โ‚ƒ_stable) && + cache.third_order_solution.colptr == ๐’โ‚ƒ_stable.colptr && + cache.third_order_solution.rowval == ๐’โ‚ƒ_stable.rowval + copyto!(cache.third_order_solution.nzval, ๐’โ‚ƒ_stable.nzval) + else + cache.third_order_solution = ๐’โ‚ƒ_stable + end + if !isempty(parameter_values) + cache.valid_for.third_order_solution = Float64.(parameter_values) + end + empty!(cache.valid_for.pruned_third_order_solution) + + # --- precompute transposed constants for pullback ----------------------------- + # Use pre-cached transposes from constants (computed once at model compile time) + ๐๐‚โ‚ƒt = Mโ‚ƒ.๐๐‚โ‚ƒแต€ + ๐›”t = Mโ‚‚.๐›”แต€ + ๐”โˆ‡โ‚‚t = Mโ‚‚.๐”โˆ‡โ‚‚แต€ + ๐”โ‚‚t = Mโ‚‚.๐”โ‚‚แต€ + + # Materialized transposes of forward-pass intermediates + At = choose_matrix_format(A') + Bt = choose_matrix_format(B') + โˆ‡โ‚‚t = choose_matrix_format(โˆ‡โ‚‚') + โˆ‡โ‚ƒt = choose_matrix_format(โˆ‡โ‚ƒ') + D_ab_t = choose_matrix_format(D_ab') + tmpkron22_t = choose_matrix_format(tmpkron22') + ck3_aux_mat_t = choose_matrix_format(ck3_aux_mat') + ๐’โ‚‚t = choose_matrix_format(๐’โ‚‚', density_threshold = 1.0) + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t = choose_matrix_format(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹') + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ') + S2p0_sigma_t = choose_matrix_format(S2p0_sigma') + + mm_๐’โ‚‚_kron_t = choose_matrix_format(mm_๐’โ‚‚_kron') + + # Precompute (โˆ‡โ‚โ‚Š ยท ๐’โ‚‚)แต€ for term 8 fused kron adjoint + โˆ‡โ‚โ‚Š_๐’โ‚‚_t = choose_matrix_format((โˆ‡โ‚โ‚Š * ๐’โ‚‚)') + + # Precompute (โˆ‡โ‚‚ ยท kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ))แต€ for fused terms a+b pullback + nabla2_kron_S1S2_t = collect(mat_mult_kron(collect(โˆ‡โ‚‚), collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ))') + + # Sparse ฯƒ for fill_kron_adjoint_โˆ‚A_with_perm! (ultra-sparse: ~nโ‚‘ nonzeros in nโ‚‘โ‚‹ยฒ ร— nโ‚‘โ‚‹ยฒ) + ฯƒ_sparse = Mโ‚‚.๐›” isa SparseMatrixCSC ? Mโ‚‚.๐›” : sparse(Mโ‚‚.๐›”) + + # --- ensure pullback workspace buffers --- + ensure_third_order_pullback_workspaces!(โ„‚, S, T, Mโ‚‚, Mโ‚ƒ) + + tmpkron22_ck3_aux_mat_t = choose_matrix_format(tmpkron22_t + ck3_aux_mat_t) + # ========================================================================= + # PULLBACK + # ========================================================================= + function third_order_solution_pullback(โˆ‚๐’โ‚ƒ_solved) + โˆ‚๐’โ‚ƒ = choose_matrix_format(โˆ‚๐’โ‚ƒ_solved[1]) + + if โ„’.norm(โˆ‚๐’โ‚ƒ) < opts.tol.third_order.ad.sylvester.acceptance_tol + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # --- adjoint Sylvester: Aแต€ โˆ‚C_adj Bแต€ + โˆ‚๐’โ‚ƒ = โˆ‚C_adj -------------------- + โˆ‚C_adj, slvd = solve_sylvester_equation(At, Bt, โˆ‚๐’โ‚ƒ, โ„‚.sylvester_workspace, + sylvester_algorithm = opts.sylvester_algorithmยณ, + tol = opts.tol.third_order.ad.sylvester, + verbose = opts.verbose) + if !slvd + return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + โˆ‚C_adj = choose_matrix_format(โˆ‚C_adj) + + # --- Initialize all gradient accumulators --- + # Dense workspace temporaries (overwritten by mul! each call) + โˆ‚๐—โ‚ƒ = โ„‚.โˆ‚๐—โ‚ƒ_3rd + โˆ‚A = โ„‚.โˆ‚A_3rd + โˆ‚B_from_sylv = โ„‚.โˆ‚B_sylv_3rd + โˆ‚out2 = โ„‚.โˆ‚out2_3rd + mul_tmp = โ„‚.mul_tmp_3rd + โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = โ„‚.โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd + + # Dense workspace accumulators (need zeroing) + โˆ‚spinv = โ„‚.โˆ‚spinv_3rd + โˆ‚โˆ‡โ‚ = โ„‚.โˆ‚โˆ‡โ‚_3rd; fill!(โˆ‚โˆ‡โ‚, zero(S)) + โˆ‚๐’โ‚โ‚ƒ = โ„‚.โˆ‚๐’โ‚_3rd; fill!(โˆ‚๐’โ‚โ‚ƒ, zero(S)) + + # Sparse-preserving gradient accumulators (reuse workspace buffers) + โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) # sparse โ€” must stay fresh + + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, zero(S)) + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, zero(S)) + โˆ‚L_c = โ„‚.โˆ‚L_c_3rd; fill!(โˆ‚L_c, zero(S)) + โˆ‚R_c = โ„‚.โˆ‚R_c_3rd; fill!(โˆ‚R_c, zero(S)) + โˆ‚L_d = โ„‚.โˆ‚L_d_3rd; fill!(โˆ‚L_d, zero(S)) + โˆ‚R_d = โ„‚.โˆ‚R_d_3rd; fill!(โˆ‚R_d, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8 = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, zero(S)) + โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, zero(S)) + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, zero(S)) + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ, zero(S)) + โˆ‚S1S1_stack = โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd; fill!(โˆ‚S1S1_stack, zero(S)) + โˆ‚aux = โ„‚.โˆ‚aux_3rd; fill!(โˆ‚aux, zero(S)) + + # --- gradient of A, B, C from ๐’โ‚ƒ = Aยท๐’โ‚ƒยทB + C --------------------------- + # โˆ‚A = โˆ‚C_adj * B' * ๐’โ‚ƒ_stable' โ€” use โˆ‚๐—โ‚ƒ as temp for intermediate + โ„’.mul!(โˆ‚๐—โ‚ƒ, โˆ‚C_adj, Bt) + โ„’.mul!(โˆ‚A, โˆ‚๐—โ‚ƒ, ๐’โ‚ƒ_stable') + # โˆ‚B_from_sylv = ๐’โ‚ƒ_stable' * A' * โˆ‚C_adj โ€” reuse โˆ‚๐—โ‚ƒ as temp + โ„’.mul!(โˆ‚๐—โ‚ƒ, At, โˆ‚C_adj) + โ„’.mul!(โˆ‚B_from_sylv, ๐’โ‚ƒ_stable', โˆ‚๐—โ‚ƒ) + # โˆ‚B_from_sylv = sparse(๐’โ‚ƒ_stable' * โˆ‚๐—โ‚ƒ) + # โˆ‚๐—โ‚ƒ = spinv' * โˆ‚C_adj (overwrite temp with real value) + # โ„’.mul!(โˆ‚๐—โ‚ƒ, sxpinv', โˆ‚C_adj) + โˆ‚๐—โ‚ƒ = choose_matrix_format(spinv' * โˆ‚C_adj) + + # C = spinv * ๐—โ‚ƒ โ†’ โˆ‚spinv + # A = spinv * โˆ‡โ‚โ‚Š โ†’ โˆ‚spinv accumulation + โ„’.mul!(โˆ‚spinv, โˆ‚C_adj, ๐—โ‚ƒ') + โ„’.mul!(โˆ‚spinv, โˆ‚A, โˆ‡โ‚โ‚Š', 1, 1) + + # ===================================================================== + # โˆ‚โˆ‡โ‚ƒ (linear: โˆ‡โ‚ƒ appears in two additive terms of ๐—โ‚ƒ) + # ===================================================================== + # ๐—โ‚ƒ = out2 * ๐๐‚โ‚ƒ + โˆ‡โ‚ƒ * tmpkron22 + โˆ‡โ‚ƒ * ck3_aux_mat + # โˆ‡โ‚ƒ has two direct linear terms; out2 maps through ๐๐‚โ‚ƒ. + โˆ‚โˆ‡โ‚ƒ = โˆ‚๐—โ‚ƒ * tmpkron22_ck3_aux_mat_t + # ===================================================================== + # โˆ‚โˆ‡โ‚‚ (โˆ‡โ‚‚ is linear in out2 โ†’ ๐—โ‚ƒ_pre โ†’ ๐—โ‚ƒ) + # ===================================================================== + # out2 enters ๐—โ‚ƒ as: ๐—โ‚ƒ = out2 ยท ๐๐‚โ‚ƒ + ... + # โˆ‚out2 = โˆ‚๐—โ‚ƒ ยท (๐๐‚โ‚ƒ)แต€ + โ„’.mul!(โˆ‚out2, โˆ‚๐—โ‚ƒ, ๐๐‚โ‚ƒt) + + # ๐—โ‚ƒ = โˆ‡โ‚‚ * kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) * D_ab (terms a+b) + # + โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, โŽธ๐’โ‚‚k..โŽน) * ๐๐‚โ‚ƒ (term c) + # + โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽยท๐›”) * ๐๐‚โ‚ƒ (term d) + # (term e = โˆ‡โ‚โ‚Š ยท ๐’โ‚‚ ยท kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) ยท ๐๐‚โ‚ƒ does not involve โˆ‡โ‚‚.) + + # โˆ‚โˆ‡โ‚‚ via mat_mult_kron (avoids materializing cubic kron transposes) + โˆ‚mid_ab = choose_matrix_format(โˆ‚๐—โ‚ƒ * D_ab_t) # n ร— nโ‚‘โ‚‹ยณ + โˆ‚โˆ‡โ‚‚ = mat_mult_kron(โˆ‚mid_ab, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ'), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ')) # terms a+b + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ + mat_mult_kron(โˆ‚out2, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt) # term c + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ + mat_mult_kron(โˆ‚out2, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, S2p0_sigma_t) # term d + + + # ===================================================================== + # โˆ‚๐’โ‚‚ (๐’โ‚‚ enters out2 via several stacking matrices) + # ===================================================================== + # ๐’โ‚‚ does NOT affect A, B, or the โˆ‡โ‚ƒ terms โ€” only out2. + # We already have โˆ‚out2 from the ๐—โ‚ƒ = out2 * ๐๐‚โ‚ƒ adjoint. + # + # out2 terms that depend on ๐’โ‚‚: + # (a) โˆ‡โ‚‚ ยท tmpkron1 ยท tmpkron2 โ€” tmpkron1 = kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) + # (b) โˆ‡โ‚‚ ยท tmpkron1 ยท ๐โ‚โ‚— ยท tmpkron2 ยท ๐โ‚แตฃ โ€” same tmpkron1 + # (c) โˆ‡โ‚‚ ยท kron(โŽธ๐’โ‚..โŽน, โŽธ๐’โ‚‚k..โŽน) โ€” second factor depends on ๐’โ‚‚ + # (d) โˆ‡โ‚‚ ยท kron(โŽธ๐’โ‚..โŽน, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽยท๐›”) โ€” second factor depends on ๐’โ‚‚ + # (8) โˆ‡โ‚โ‚Š ยท ๐’โ‚‚ ยท kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) โ€” both ๐’โ‚‚ and ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ depend on ๐’โ‚‚ + + # --- terms (a) and (b): through kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) via D_ab --- + # โˆ‚kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) = โˆ‡โ‚‚แต€ * โˆ‚๐—โ‚ƒ * D_ab' (combines terms a+b) + โˆ‚tmpkron1 = (โˆ‡โ‚‚t * โˆ‚mid_ab) + # โˆ‚tmpkron1 = sparse(โˆ‡โ‚‚t * โˆ‚mid_ab) + + # Force only the cotangent argument onto the dense fill_kron_adjoint! path here + # and in the analogous calls below. The primal factors may stay sparse/abstract, + # but the sparse โˆ‚X overload only iterates stored cotangent entries. + # kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) pullback โ†’ โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ via fill_kron_adjoint! + fill_kron_adjoint!(โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, โˆ‚tmpkron1, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + + # ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚Š,:]; 0] โ†’ โˆ‚๐’โ‚‚[iโ‚Š,:] += โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] + + # --- term (c): through โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ --- + # Fused: โˆ‡โ‚‚แต€ * โˆ‚out2 with fill_kron_adjoint! โ€” avoids materializing โˆ‡โ‚‚t_โˆ‚out2 + mul_fill_kron_adjoint!(โˆ‚R_c, โˆ‚L_c, โˆ‡โ‚‚t, โˆ‚out2, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, tol = opts.tol.third_order.droptol) + + # โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = [ (๐’โ‚‚ยทkron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ + ๐’โ‚ยท[๐’โ‚‚[iโ‚‹,:];0])[iโ‚Š,:] ; ๐’โ‚‚ ; 0 ] + # Top block (rows 1:nโ‚Š): depends on ๐’โ‚‚ through ๐’โ‚‚ยทkron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ and ๐’โ‚ยท[๐’โ‚‚[iโ‚‹,:];0] + nโ‚Š_len = length(iโ‚Š) + โˆ‚top_block = โˆ‚R_c[1:nโ‚Š_len, :] + # From ๐’โ‚‚ยทkron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘: + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚top_block * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + # From ๐’โ‚ยท[๐’โ‚‚[iโ‚‹,:];0] โ†’ โˆ‚๐’โ‚‚[iโ‚‹,:] += ๐’โ‚' * I[:,iโ‚Š] * โˆ‚top_block + # (since [๐’โ‚‚[iโ‚‹,:];0] pads with zeros, only iโ‚‹ rows of ๐’โ‚‚ contribute) + โˆ‚๐’โ‚‚_padded = ๐’โ‚' * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block # TODO: In general check if there are more optimizations that can be carried over from the non-AD call. # nโ‚‹+1+nโ‚‘ ร— nโ‚‘โ‚‹ยฒ + @views โˆ‚๐’โ‚‚[iโ‚‹,:] .+= โˆ‚๐’โ‚‚_padded[1:nโ‚‹, :] + + # Middle block (rows nโ‚Š_len+1 : nโ‚Š_len+n): directly ๐’โ‚‚ + @views โˆ‚๐’โ‚‚ .+= โˆ‚R_c[nโ‚Š_len .+ (1:n), :] + + # Bottom block is zeros + + # --- term (d): through kron(โŽธ๐’โ‚..โŽน, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽยท๐›”) --- + # Fused: โˆ‡โ‚‚แต€ * โˆ‚out2 with fill_kron_adjoint! โ€” same pattern, different kron factors + mul_fill_kron_adjoint!(โˆ‚R_d, โˆ‚L_d, โˆ‡โ‚‚t, โˆ‚out2, S2p0_sigma, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, tol = opts.tol.third_order.droptol) + + # ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽยท๐›” โ†’ โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d = โˆ‚R_d ยท ๐›”แต€ + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d = โˆ‚R_d * ๐›”t + @views โˆ‚๐’โ‚‚[iโ‚Š,:] .+= โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_d[1:length(iโ‚Š),:] + + # --- term (8): โˆ‡โ‚โ‚Š ยท ๐’โ‚‚ ยท kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) --- + # out2_term8 = โˆ‡โ‚โ‚Š ยท ๐’โ‚‚ ยท kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) + # โˆ‚(โˆ‡โ‚โ‚Šยท๐’โ‚‚ยทK) w.r.t. ๐’โ‚‚ = โˆ‡โ‚โ‚Šแต€ ยท โˆ‚out2 ยท Kแต€ + tmp_t8 = โˆ‡โ‚โ‚Š' * โˆ‚out2 + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ + mat_mult_kron(tmp_t8, collect(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘'), collect(๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ')) + + # โˆ‚(โˆ‡โ‚โ‚Šยท๐’โ‚‚ยทkron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘,๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ)) w.r.t. ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ (through the kron) + # Fused: (โˆ‡โ‚โ‚Šยท๐’โ‚‚)แต€ ยท โˆ‚out2 with fill_kron_adjoint! in one pass + mul_fill_kron_adjoint!(โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, โˆ‡โ‚โ‚Š_๐’โ‚‚_t, โˆ‚out2, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.third_order.droptol) + + # ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚‹,:]; 0] โ†’ โˆ‚๐’โ‚‚[iโ‚‹,:] += โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ[1:nโ‚‹,:] + @views โˆ‚๐’โ‚‚[iโ‚‹,:] .+= โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ[1:nโ‚‹,:] + + # ===================================================================== + # โˆ‚โˆ‡โ‚ + # ===================================================================== + # โˆ‡โ‚ enters through: + # 1. โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -โˆ‡โ‚[:,1:nโ‚Š]ยท๐’โ‚[iโ‚Š,1:nโ‚‹]ยทI[iโ‚‹,:] - โˆ‡โ‚[:,nโ‚Š+1:nโ‚Š+n] + # โ†’ spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€) โ†’ used in A and C + # 2. โˆ‡โ‚โ‚Š = โˆ‡โ‚[:,1:nโ‚Š] ยท I(n)[iโ‚Š,:] + # โ†’ A = spinvยทโˆ‡โ‚โ‚Š and out2 += โˆ‡โ‚โ‚Š ยท mm_๐’โ‚‚_kron + + # step 1: โˆ‚ through inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€) (โˆ‚spinv already accumulated) + โ„’.mul!(mul_tmp, spinv', โˆ‚spinv) + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, mul_tmp, spinv') + โ„’.rmul!(โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, -1) + + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] * ๐’โ‚[iโ‚Š,1:nโ‚‹]' + โˆ‚โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ + + # step 2: โˆ‚ through โˆ‡โ‚โ‚Š + โˆ‚โˆ‡โ‚โ‚Š = โ„‚.โˆ‚โˆ‡โ‚โ‚Š_3rd + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š, spinv', โˆ‚A) # from A = spinv ยท โˆ‡โ‚โ‚Š + โ„’.mul!(โˆ‚โˆ‡โ‚โ‚Š, โˆ‚out2, mm_๐’โ‚‚_kron_t, 1, 1) # from out2 += โˆ‡โ‚โ‚Š ยท mm_๐’โ‚‚_kron + + โˆ‚โˆ‡โ‚[:,1:nโ‚Š] += โˆ‚โˆ‡โ‚โ‚Š * โ„’.I(n)[:,iโ‚Š] + + # ===================================================================== + # โˆ‚๐‘บโ‚ (๐‘บโ‚ enters through ๐’โ‚, affecting A,B,C,out2 via many paths) + # ===================================================================== + # --- โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ : from out2 terms c,d (kron outer factors) --- + โ„’.axpy!(1, โˆ‚L_c, โˆ‚S1S1_stack) + โ„’.axpy!(1, โˆ‚L_d, โˆ‚S1S1_stack) + + # --- โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ : from โˆ‡โ‚ƒ * compressed_kron(...) --- + # Fused: compute g_col = โˆ‡โ‚ƒแต€ * โˆ‚๐—โ‚ƒ[:, col] lazily per (ฮฑ,ฮฒ,ฮณ) triple + # instead of materializing the full โˆ‚tmpkron22 = โˆ‡โ‚ƒแต€ * โˆ‚๐—โ‚ƒ matrix. + โˆ‚S1S1_from_ck = โ„‚.โˆ‚S1S1_from_ck_3rd + fill!(โˆ‚S1S1_from_ck, zero(S)) + โˆ‚S1p0_kron_sigma = โ„‚.โˆ‚S1p0_kron_sigma_3rd + fill!(โˆ‚S1p0_kron_sigma, zero(S)) + mul_compressed_permuted_mixed_kron_pullback!(โˆ‚S1S1_from_ck, + โˆ‚S1p0_kron_sigma, + โˆ‡โ‚ƒt, โˆ‚๐—โ‚ƒ, + โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + S1p0_kron_sigma; + tol = opts.tol.third_order.droptol) + + # Sparsify โˆ‚S1p0_kron_sigma: structurally bounded by ฯƒ's support, so very sparse. + # sparse ร— sparse matmul avoids dense intermediate; downstream fill_kron_adjoint! + # uses the sparse overload that iterates only nonzero cotangent entries. + โˆ‚S1p0_kron = choose_matrix_format(sparse(โˆ‚S1p0_kron_sigma) * ๐›”t) + โˆ‚S1p0_left = โ„‚.โˆ‚S1p0_left_3rd + fill!(โˆ‚S1p0_left, zero(S)) + โˆ‚S1p0_right = โ„‚.โˆ‚S1p0_right_3rd + fill!(โˆ‚S1p0_right, zero(S)) + fill_kron_adjoint!(โˆ‚S1p0_left, โˆ‚S1p0_right, โˆ‚S1p0_kron, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + + โ„’.axpy!(1, โˆ‚S1S1_from_ck, โˆ‚S1S1_stack) + โ„’.axpy!(1, โˆ‚S1p0_left, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + โ„’.axpy!(1, โˆ‚S1p0_right, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + + # --- โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ : from compressed_kronยณ(aux) โ†’ ๐—โ‚ƒ --- + # Fused: compute g_col = โˆ‡โ‚ƒแต€ * โˆ‚๐—โ‚ƒ[:, col] lazily per (i2,j2,k2) triple + mul_compressed_kronยณ_pullback!(โˆ‚aux, โˆ‡โ‚ƒt, โˆ‚๐—โ‚ƒ, aux; tol = opts.tol.third_order.droptol) + โ„’.mul!(โˆ‚S1S1_stack, Mโ‚ƒ.๐’๐', โˆ‚aux, 1, 1) + + # --- โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ : from tmpkron1 (already computed for โˆ‚๐’โ‚‚) --- + โ„’.axpy!(1, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ) + + # --- โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ : from B via compressed_permuted_mixed_kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐›”) --- + compressed_permuted_mixed_kron_pullback_โˆ‚A!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚B_from_sylv, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”; tol = opts.tol.third_order.droptol) + + # --- โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ : from B via compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) --- + compressed_kronยณ_pullback!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚B_from_sylv, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘; tol = opts.tol.third_order.droptol) + + # --- โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ : from out2 terms a,b via tmpkron2 = kron(B=๐›”, A=๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) --- + # Fused: nabla2_kron_S1S2_t * โˆ‚out2 in blocks + identity/(2,1,3) permuted โˆ‚A + # Avoids materializing both โˆ‡โ‚‚t_โˆ‚out2 (n_โˆ‡โ‚‚ ร— n_out2_c) and tmp_a (nโ‚‘โ‚‹ยณ ร— nโ‚‘โ‚‹ยณ) + mul_fill_kron_adjoint_โˆ‚A_with_perm!(nabla2_kron_S1S2_t, โˆ‚out2, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, ฯƒ_sparse) + + # --- โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ : from term 8 kron (already computed for โˆ‚๐’โ‚‚) --- + โ„’.axpy!(1, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ) + + # --- โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ : from kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ in โŽธ๐’โ‚‚k..โŽน top block --- + # โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = sparse(๐’โ‚‚t * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block) + โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ = (๐’โ‚‚t * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_block) + fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + + # --- โˆ‚๐’โ‚ : from ๐’โ‚ยท[๐’โ‚‚[iโ‚‹,:];0] in โŽธ๐’โ‚‚k..โŽน top block --- + S2_padded = [๐’โ‚‚[iโ‚‹,:]; zeros(S, nโ‚‘ + 1, nโ‚‘โ‚‹^2)] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚top_block * S2_padded' + + # === Convert โˆ‚S1S1_stack โ†’ โˆ‚๐’โ‚ and โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ === + nโ‚Šl = length(iโ‚Š) + โˆ‚top_S1S1 = โˆ‚S1S1_stack[1:nโ‚Šl, :] + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚top_S1S1 * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ .+= ๐’โ‚' * โ„’.I(n)[:,iโ‚Š] * โˆ‚top_S1S1 + @views โˆ‚๐’โ‚โ‚ƒ .+= โˆ‚S1S1_stack[nโ‚Šl .+ (1:n), :] + + # === Convert โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚“ โ†’ โˆ‚๐’โ‚ === + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,:] .+= โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽโ‚ƒ[1:nโ‚Šl,:] + + # === Convert โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ โ†’ โˆ‚๐’โ‚ === + @views โˆ‚๐’โ‚โ‚ƒ[iโ‚‹,:] .+= โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โ‚ƒ[1:length(iโ‚‹),:] + + # === โˆ‚๐’โ‚ from โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ (spinv) === + โˆ‚๐’โ‚โ‚ƒ[iโ‚Š,1:nโ‚‹] -= โˆ‡โ‚[:,1:nโ‚Š]' * โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] + + # === ๐’โ‚ = [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]] โ†’ โˆ‚๐‘บโ‚ === + โˆ‚๐‘บโ‚ = [โˆ‚๐’โ‚โ‚ƒ[:,1:nโ‚‹] โˆ‚๐’โ‚โ‚ƒ[:,nโ‚‹+2:end]] + + # Map โˆ‚โˆ‡โ‚‚ and โˆ‚๐’โ‚‚ back to compressed space + # (adjoint of โˆ‡โ‚‚_full = โˆ‡โ‚‚_compressed * ๐”โˆ‡โ‚‚ and ๐’โ‚‚_full = ๐’โ‚‚_compressed * ๐”โ‚‚) + โˆ‚โˆ‡โ‚‚ = โˆ‚โˆ‡โ‚‚ * ๐”โˆ‡โ‚‚t + โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚ * ๐”โ‚‚t + + return (NoTangent(), โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚ƒ, โˆ‚๐‘บโ‚, โˆ‚๐’โ‚‚, NoTangent(), NoTangent(), NoTangent()) + end + + return (๐’โ‚ƒ_stable, solved), third_order_solution_pullback +end + + +function rrule(::typeof(solve_sylvester_equation), + A::M, + B::N, + C::O, + ๐•Šโ„‚::sylvester_workspace; + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), + sylvester_algorithm::Symbol = :doubling, + tol::SolverTolerances = SolverTolerances(), + # timer::TimerOutput = TimerOutput(), + verbose::Bool = false) where {M <: AbstractMatrix{Float64}, N <: AbstractMatrix{Float64}, O <: AbstractMatrix{Float64}} + + P, solved = solve_sylvester_equation(A, B, C, ๐•Šโ„‚, + sylvester_algorithm = sylvester_algorithm, + tol = tol, + verbose = verbose, + initial_guess = initial_guess) + + if size(๐•Šโ„‚.P) != size(P) + ๐•Šโ„‚.P = zeros(eltype(P), size(P)...) + end + copyto!(๐•Šโ„‚.P, P) + P_cached = ๐•Šโ„‚.P + + ensure_sylvester_doubling_buffers!(๐•Šโ„‚, size(A, 1), size(B, 1)) + + # pullback + function solve_sylvester_equation_pullback(โˆ‚P) + if โ„’.norm(โˆ‚P[1]) < tol.rtol return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() end + + โˆ‚C, slvd = solve_sylvester_equation(A', B', โˆ‚P[1], ๐•Šโ„‚, + sylvester_algorithm = sylvester_algorithm, + tol = tol, + verbose = verbose) + + solved = solved && slvd + + # โˆ‚C is nร—m, B' is mร—m, P_cached is nร—m, A is nร—n + # Intermediate products are nร—m and mร—n โ€” not nร—n or mร—m, + # so workspace buffers ๐€ (nร—n) / ๐ (mร—m) are wrong shape when n โ‰  m. + โˆ‚A = (โˆ‚C * B') * P_cached' + โˆ‚B = (P_cached' * A') * โˆ‚C + + return NoTangent(), โˆ‚A, โˆ‚B, โˆ‚C, NoTangent() + end + + return (P_cached, solved), solve_sylvester_equation_pullback +end + +function rrule(::typeof(solve_lyapunov_equation), + A::AbstractMatrix{Float64}, + C::AbstractMatrix{Float64}, + workspace::lyapunov_workspace; + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), + lyapunov_algorithm::Symbol = :doubling, + tol::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), + # timer::TimerOutput = TimerOutput(), + verbose::Bool = false) + + P, solved = solve_lyapunov_equation(A, C, workspace, + initial_guess = initial_guess, + lyapunov_algorithm = lyapunov_algorithm, + tol = tol, + verbose = verbose) + if size(workspace.P) != size(P) + workspace.P = zeros(eltype(P), size(P)...) + end + copyto!(workspace.P, P) + P_cached = workspace.P + ensure_lyapunov_doubling_buffers!(workspace) + A_dense = collect(A) + + # pullback + # https://arxiv.org/abs/2011.11430 + function solve_lyapunov_equation_pullback(โˆ‚P) + if โ„’.norm(โˆ‚P[1]) < tol.rtol return NoTangent(), NoTangent(), NoTangent(), NoTangent() end + + # Adjoint Lyapunov: โˆ‚P is generally not symmetric, so issymmetric will route to full-space + # Use dense A' directly with Val(:doubling) to force BLAS-backed dense path + # (the dispatcher's choose_matrix_format would convert back to sparse) + โˆ‚C_result, adj_iters, adj_tol = solve_lyapunov_equation(A_dense', Matrix{Float64}(โˆ‚P[1]), Val(:doubling), workspace, tol = tol) + โˆ‚C = โˆ‚C_result + slvd = adj_tol < tol.acceptance_tol + + solved = solved && slvd + + tmp_n1 = workspace.๐‚A + tmp_n2 = workspace.๐€ยฒ + โˆ‚A = zeros(eltype(A), size(A)) + + โ„’.mul!(tmp_n1, โˆ‚C, A_dense) + โ„’.mul!(โˆ‚A, tmp_n1, P_cached') + + โ„’.mul!(tmp_n2, โˆ‚C', A_dense) + โ„’.mul!(โˆ‚A, tmp_n2, P_cached, 1, 1) + + return NoTangent(), โˆ‚A, โˆ‚C, NoTangent() + end + + return (P_cached, solved), solve_lyapunov_equation_pullback +end + +function rrule(::typeof(find_shocks), + ::Val{:LagrangeNewton}, + initial_guess::Vector{Float64}, + kron_buffer::Vector{Float64}, + kron_buffer2::AbstractMatrix{Float64}, + J::โ„’.Diagonal{Bool, Vector{Bool}}, + ๐’โฑ::AbstractMatrix{Float64}, + ๐’โฑยฒแต‰::AbstractMatrix{Float64}, + shock_independent::Vector{Float64}; + max_iter::Int = 1000, + tol::Float64 = 1e-13) + + x, matched = find_shocks(Val(:LagrangeNewton), + initial_guess, + kron_buffer, + kron_buffer2, + J, + ๐’โฑ, + ๐’โฑยฒแต‰, + shock_independent, + max_iter = max_iter, + tol = tol) + + tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x)), x) + + ฮป = tmp' \ x * 2 + + fXฮปp = [reshape(2 * ๐’โฑยฒแต‰' * ฮป, size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' + -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + โ„’.kron!(kron_buffer, x, x) + + xฮป = โ„’.kron(x,ฮป) + + + โˆ‚shock_independent = similar(shock_independent) + + # โˆ‚๐’โฑ = similar(๐’โฑ) + + # โˆ‚๐’โฑยฒแต‰ = similar(๐’โฑยฒแต‰) + + function find_shocks_pullback(โˆ‚x) + โˆ‚x = vcat(โˆ‚x[1], zero(ฮป)) + + S = -fXฮปp' \ โˆ‚x + + copyto!(โˆ‚shock_independent, S[length(initial_guess)+1:end]) + + # copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:length(initial_guess)], ฮป) - โ„’.kron(x, S[length(initial_guess)+1:end])) + โˆ‚๐’โฑ = S[1:length(initial_guess)] * ฮป' - S[length(initial_guess)+1:end] * x' + + # copyto!(โˆ‚๐’โฑยฒแต‰, 2 * โ„’.kron(S[1:length(initial_guess)], xฮป) - โ„’.kron(kron_buffer, S[length(initial_guess)+1:end])) + โˆ‚๐’โฑยฒแต‰ = 2 * S[1:length(initial_guess)] * xฮป' - S[length(initial_guess)+1:end] * kron_buffer' + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’โฑ, โˆ‚๐’โฑยฒแต‰, โˆ‚shock_independent, NoTangent(), NoTangent() + end + + return (x, matched), find_shocks_pullback +end + +function rrule(::typeof(find_shocks), + ::Val{:LagrangeNewton}, + initial_guess::Vector{Float64}, + kron_buffer::Vector{Float64}, + kron_bufferยฒ::Vector{Float64}, + kron_buffer2::AbstractMatrix{Float64}, + kron_buffer3::AbstractMatrix{Float64}, + kron_buffer4::AbstractMatrix{Float64}, + J::โ„’.Diagonal{Bool, Vector{Bool}}, + ๐’โฑ::AbstractMatrix{Float64}, + ๐’โฑยฒแต‰::AbstractMatrix{Float64}, + ๐’โฑยณแต‰::AbstractMatrix{Float64}, + shock_independent::Vector{Float64}; + max_iter::Int = 1000, + tol::Float64 = 1e-13) + + x, matched = find_shocks(Val(:LagrangeNewton), + initial_guess, + kron_buffer, + kron_bufferยฒ, + kron_buffer2, + kron_buffer3, + kron_buffer4, + J, + ๐’โฑ, + ๐’โฑยฒแต‰, + ๐’โฑยณแต‰, + shock_independent, + max_iter = max_iter, + tol = tol) + + โ„’.kron!(kron_buffer, x, x) + + โ„’.kron!(kron_bufferยฒ, x, kron_buffer) + + tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x)), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(length(x)), kron_buffer) + + ฮป = tmp' \ x * 2 + + fXฮปp = [reshape((2 * ๐’โฑยฒแต‰ + 6 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(length(x)), โ„’.kron(โ„’.I(length(x)),x)))' * ฮป, size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' + -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + xฮป = โ„’.kron(x,ฮป) + + xxฮป = โ„’.kron(x,xฮป) + + function find_shocks_pullback(โˆ‚x) + โˆ‚x = vcat(โˆ‚x[1], zero(ฮป)) + + S = -fXฮปp' \ โˆ‚x + + โˆ‚shock_independent = S[length(initial_guess)+1:end] + + โˆ‚๐’โฑ = โ„’.kron(S[1:length(initial_guess)], ฮป) - โ„’.kron(x, S[length(initial_guess)+1:end]) + + โˆ‚๐’โฑยฒแต‰ = 2 * โ„’.kron(S[1:length(initial_guess)], xฮป) - โ„’.kron(kron_buffer, S[length(initial_guess)+1:end]) + + โˆ‚๐’โฑยณแต‰ = 3 * โ„’.kron(S[1:length(initial_guess)], xxฮป) - โ„’.kron(kron_bufferยฒ,S[length(initial_guess)+1:end]) + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’โฑ, โˆ‚๐’โฑยฒแต‰, โˆ‚๐’โฑยณแต‰, โˆ‚shock_independent, NoTangent(), NoTangent() + end + + return (x, matched), find_shocks_pullback +end + + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:inversion}, + ::Val{:first_order}, + observables_index::Vector{Int}, + ๐’::Matrix{Float64}, + data_in_deviations::Matrix{Float64}, + constants::constants, + state::Vector{Vector{Float64}}, + workspaces::workspaces; + # timer::TimerOutput = TimerOutput(), + warmup_iterations::Int = 0, + on_failure_loglikelihood = -Inf, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + opts::CalculationOptions = merge_calculation_options(), + filter_algorithm::Symbol = :LagrangeNewton) + T = constants.post_model_macro + ws = workspaces.inversion + # @timeit_debug timer "Inversion filter - forward" begin + + # first order + state = copy(state[1]) + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + obs_idx = observables_index + + tโป = T.past_not_future_and_mixed_idx + + shocksยฒ = 0.0 + logabsdets = 0.0 + + @assert warmup_iterations == 0 "Warmup iterations not yet implemented for reverse-mode automatic differentiation." + + state = [copy(state) for _ in 1:size(data_in_deviations,2)+1] + + shocksยฒ = 0.0 + logabsdets = 0.0 + + y = zeros(length(obs_idx)) + x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] + + jac = ๐’[obs_idx,end-T.nExo+1:end] + + if T.nExo == length(observables_index) + logabsdets = โ„’.logabsdet(jac)[1] # ./ precision_factor + + jacdecomp = โ„’.lu(jac, check = false) + + if !โ„’.issuccess(jacdecomp) + if opts.verbose println("Inversion filter failed") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + invjac = inv(jacdecomp) + else + logabsdets = sum(x -> log(abs(x)), โ„’.svdvals(jac)) #' ./ precision_factor + # jacdecomp = โ„’.svd(jac) + invjac = โ„’.pinv(jac) + end + + logabsdets *= size(data_in_deviations,2) - presample_periods + + if !isfinite(logabsdets) + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + @views ๐’obs = ๐’[obs_idx,1:end-T.nExo] + + for i in axes(data_in_deviations,2) + @views โ„’.mul!(y, ๐’obs, state[i][tโป]) + @views โ„’.axpby!(1, data_in_deviations[:,i], -1, y) + โ„’.mul!(x[i],invjac,y) + # x = ๐’[obs_idx,end-T.nExo+1:end] \ (data_in_deviations[:,i] - ๐’[obs_idx,1:end-T.nExo] * state[tโป]) + + if i > presample_periods + shocksยฒ += sum(abs2,x[i]) + if !isfinite(shocksยฒ) + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + end + + โ„’.mul!(state[i+1], ๐’, vcat(state[i][tโป], x[i])) + # state[i+1] = ๐’ * vcat(state[i][tโป], x[i]) + end + + llh = -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + + if llh < -1e12 + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + โˆ‚๐’ = zero(๐’) + + โˆ‚๐’แต—โป = copy(โˆ‚๐’[tโป,:]) + + โˆ‚data_in_deviations = zero(data_in_deviations) + + # Allocate or reuse workspaces for pullback + n_periods = size(data_in_deviations,2) - 1 + if size(ws.โˆ‚data) != (length(tโป), n_periods) + ws.โˆ‚data = zeros(length(tโป), n_periods) + else + fill!(ws.โˆ‚data, zero(eltype(ws.โˆ‚data))) + end + โˆ‚data = ws.โˆ‚data + + โˆ‚state = zero(state[1]) + + # precomputed matrices + Mยน = ๐’[obs_idx, 1:end-T.nExo]' * invjac' + Mยฒ = ๐’[tโป,1:end-T.nExo]' - Mยน * ๐’[tโป,end-T.nExo+1:end]' + Mยณ = invjac' * ๐’[tโป,end-T.nExo+1:end]' + + โˆ‚Stmp = [copy(Mยน) for _ in 1:size(data_in_deviations,2)-1] + + for t in 2:size(data_in_deviations,2)-1 + โ„’.mul!(โˆ‚Stmp[t], Mยฒ, โˆ‚Stmp[t-1]) + # โˆ‚Stmp[t] = Mยฒ * โˆ‚Stmp[t-1] + end + + # Allocate or reuse workspaces for temporary matrices + if size(ws.โˆ‚_tmp1) != (T.nExo, length(tโป) + T.nExo) + ws.โˆ‚_tmp1 = zeros(Float64, T.nExo, length(tโป) + T.nExo) + else + fill!(ws.โˆ‚_tmp1, zero(Float64)) + end + tmp1 = ws.โˆ‚_tmp1 + + if size(ws.โˆ‚_tmp2) != (length(tโป), length(tโป) + T.nExo) + ws.โˆ‚_tmp2 = zeros(Float64, length(tโป), length(tโป) + T.nExo) + else + fill!(ws.โˆ‚_tmp2, zero(Float64)) + end + tmp2 = ws.โˆ‚_tmp2 + + if size(ws.โˆ‚_tmp3) != (length(tโป) + T.nExo,) + ws.โˆ‚_tmp3 = zeros(Float64, length(tโป) + T.nExo) + else + fill!(ws.โˆ‚_tmp3, zero(Float64)) + end + tmp3 = ws.โˆ‚_tmp3 + + if size(ws.โˆ‚๐’tโป) != size(tmp2) + ws.โˆ‚๐’tโป = copy(tmp2) + else + fill!(ws.โˆ‚๐’tโป, zero(Float64)) + end + โˆ‚๐’tโป = ws.โˆ‚๐’tโป + # โˆ‚๐’obs_idx = copy(tmp1) + + # end # timeit_debug + # pullback + function inversion_pullback(โˆ‚llh) + # @timeit_debug timer "Inversion filter - pullback" begin + + for t in reverse(axes(data_in_deviations,2)) + โˆ‚state[tโป] .= Mยฒ * โˆ‚state[tโป] + + if t > presample_periods + โˆ‚state[tโป] += Mยน * x[t] + + โˆ‚data_in_deviations[:,t] -= invjac' * x[t] + + โˆ‚๐’[obs_idx, :] += invjac' * x[t] * vcat(state[t][tโป], x[t])' + + if t > 1 + โˆ‚data[:,t:end] .= Mยฒ * โˆ‚data[:,t:end] + + โˆ‚data[:,t-1] += Mยน * x[t] + + โˆ‚data_in_deviations[:,t-1] += Mยณ * โˆ‚data[:,t-1:end] * ones(size(data_in_deviations,2) - t + 1) + + for tt in t-1:-1:1 + for (i,v) in enumerate(tโป) + copyto!(tmp3::Vector{Float64}, i::Int, state[tt]::Vector{Float64}, v::Int, 1) + end + + copyto!(tmp3, length(tโป) + 1, x[tt], 1, T.nExo) + + โ„’.mul!(tmp1, x[t], tmp3') + + โ„’.mul!(โˆ‚๐’tโป, โˆ‚Stmp[t-tt], tmp1, 1, 1) + + end + end + end + end + + โˆ‚๐’[tโป,:] += โˆ‚๐’tโป + + โˆ‚๐’[obs_idx, :] -= Mยณ * โˆ‚๐’tโป + + โˆ‚๐’[obs_idx,end-T.nExo+1:end] -= (size(data_in_deviations,2) - presample_periods) * invjac' / 2 + + # end # timeit_debug + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’ * โˆ‚llh, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), [โˆ‚state * โˆ‚llh], NoTangent() + end + + return llh, inversion_pullback +end + + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:inversion}, + ::Val{:pruned_second_order}, + observables_index::Vector{Int}, + ๐’::Vector{AbstractMatrix{Float64}}, + data_in_deviations::Matrix{Float64}, + constants::constants, + state::Vector{Vector{Float64}}, + workspaces::workspaces; + # timer::TimerOutput = TimerOutput(), + on_failure_loglikelihood = -Inf, + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + opts::CalculationOptions = merge_calculation_options(), + filter_algorithm::Symbol = :LagrangeNewton)# where S <: Real + T = constants.post_model_macro + ws = workspaces.inversion + # @timeit_debug timer "Inversion filter pruned 2nd - forward" begin + # @timeit_debug timer "Preallocation" begin + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + cond_var_idx = observables_index + + shocksยฒ = 0.0 + logabsdets = 0.0 + + cc = ensure_conditional_forecast_constants!(constants) + shock_idxs = cc.shock_idxs + shockยฒ_idxs = cc.shockยฒ_idxs + shockvarยฒ_idxs = cc.shockvarยฒ_idxs + var_volยฒ_idxs = cc.var_volยฒ_idxs + varยฒ_idxs = cc.varยฒ_idxs + + ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] + ๐’โปยนแต‰ = ๐’[1][T.past_not_future_and_mixed_idx,end-T.nExo+1:end] + ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] + ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] + ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] + + ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] + ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] + ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] + ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] + ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] + + ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› + ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป + ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ + ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ + ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ + + stateโ‚ = state[1][T.past_not_future_and_mixed_idx] + stateโ‚‚ = state[2][T.past_not_future_and_mixed_idx] + + kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] + + J = โ„’.I(T.nExo) + + kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) + + kron_buffer3 = โ„’.kron(J, zeros(T.nPast_not_future_and_mixed + 1)) + + x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] + + stateยนโป = stateโ‚ + + stateยนโป_vol = vcat(stateยนโป, 1) + + stateยฒโป = stateโ‚‚ + + ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) + + ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 + + aug_stateโ‚ = [copy([stateโ‚; 1; ones(T.nExo)]) for _ in 1:size(data_in_deviations,2)] + aug_stateโ‚‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] + + tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[1])), x[1]) + + jacc = [zero(tmp) for _ in 1:size(data_in_deviations,2)] + + jacct = copy(tmp') + + ฮป = [zeros(size(tmp, 1)) for _ in 1:size(data_in_deviations,2)] + + ฮป[1] = copy(tmp' \ x[1] * 2) + + fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' + -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] + + kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) + + kronxฮป = [zero(kronxฮป_tmp) for _ in 1:size(data_in_deviations,2)] + + kronstateยนโป_vol = zeros((T.nPast_not_future_and_mixed + 1)^2) + + kronaug_stateโ‚ = zeros(length(aug_stateโ‚[1])^2) + + shock_independent = zeros(size(data_in_deviations,1)) + + init_guess = zeros(size(๐’โฑ, 2)) + + tmp = zeros(size(๐’โฑ, 2) * size(๐’โฑ, 2)) + + lI = -2 * vec(โ„’.I(size(๐’โฑ, 2))) + + # end # timeit_debug + # @timeit_debug timer "Main loop" begin + + for i in axes(data_in_deviations,2) + # stateยนโป = stateโ‚ + + # stateยนโป_vol = vcat(stateยนโป, 1) + + # stateยฒโป = stateโ‚‚ + + copyto!(stateยนโป_vol, 1, stateโ‚, 1) + + copyto!(shock_independent, data_in_deviations[:,i]) + + โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + + โ„’.mul!(shock_independent, ๐’ยนโป, stateโ‚‚, -1, 1) + + โ„’.kron!(kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) + + โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kronstateยนโป_vol, -1/2, 1) + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + โ„’.kron!(kron_buffer3, J, stateยนโป_vol) + + โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer3) + + โ„’.axpy!(1, ๐’ยนแต‰, ๐’โฑ) + + init_guess *= 0 + + # @timeit_debug timer "Find shocks" begin + x[i], matched = find_shocks(Val(filter_algorithm), + init_guess, + kronxx[i], + kron_buffer2, + J, + ๐’โฑ, + ๐’โฑยฒแต‰, + shock_independent, + # max_iter = 100 + ) + # end # timeit_debug + + if !matched + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[i])), x[i]) + โ„’.kron!(kron_buffer2, J, x[i]) + + โ„’.mul!(jacc[i], ๐’โฑยฒแต‰, kron_buffer2) + + โ„’.axpby!(1, ๐’โฑ, 2, jacc[i]) + + copy!(jacct, jacc[i]') + + jacc_fact = try + โ„’.factorize(jacct) # otherwise this fails for nshocks > nexo + catch + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + try + โ„’.ldiv!(ฮป[i], jacc_fact, x[i]) + catch + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + โ„’.rmul!(ฮป[i], 2) + + # fXฮปp[i] = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) jacc[i]' + # -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + โ„’.mul!(tmp, ๐’โฑยฒแต‰', ฮป[i]) + โ„’.axpby!(1, lI, 2, tmp) + + fXฮปp[i][1:size(๐’โฑ, 2), 1:size(๐’โฑ, 2)] = tmp + fXฮปp[i][size(๐’โฑ, 2)+1:end, 1:size(๐’โฑ, 2)] = -jacc[i] + fXฮปp[i][1:size(๐’โฑ, 2), size(๐’โฑ, 2)+1:end] = jacct + + โ„’.kron!(kronxx[i], x[i], x[i]) + + โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) + + if i > presample_periods + # due to change of variables: jacobian determinant adjustment + if T.nExo == length(observables_index) + logabsdets += โ„’.logabsdet(jacc_fact)[1] + else + logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) + end + + shocksยฒ += sum(abs2,x[i]) + + if !isfinite(logabsdets) || !isfinite(shocksยฒ) + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + end + + # aug_stateโ‚[i] = [stateโ‚; 1; x[i]] + # aug_stateโ‚‚[i] = [stateโ‚‚; 0; zero(x[1])] + copyto!(aug_stateโ‚[i], 1, stateโ‚, 1) + copyto!(aug_stateโ‚[i], length(stateโ‚) + 2, x[i], 1) + copyto!(aug_stateโ‚‚[i], 1, stateโ‚‚, 1) + + # stateโ‚, stateโ‚‚ = [๐’โปยน * aug_stateโ‚, ๐’โปยน * aug_stateโ‚‚ + ๐’โปยฒ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2] # strictly following Andreasen et al. (2018) + โ„’.mul!(stateโ‚, ๐’โปยน, aug_stateโ‚[i]) + + โ„’.mul!(stateโ‚‚, ๐’โปยน, aug_stateโ‚‚[i]) + โ„’.kron!(kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) + โ„’.mul!(stateโ‚‚, ๐’โปยฒ, kronaug_stateโ‚, 1/2, 1) + end + + # end # timeit_debug + # end # timeit_debug + + โˆ‚data_in_deviations = similar(data_in_deviations) + + โˆ‚aug_stateโ‚ = zero(aug_stateโ‚[1]) + + โˆ‚aug_stateโ‚‚ = zero(aug_stateโ‚‚[1]) + + โˆ‚kronaug_stateโ‚ = zeros(length(aug_stateโ‚[1])^2) + + โˆ‚kronIx = zero(โ„’.kron(โ„’.I(length(x[1])), x[1])) + + โˆ‚kronIstateยนโป_vol = zero(โ„’.kron(J, stateยนโป_vol)) + + โˆ‚kronstateยนโป_vol = zero(โ„’.kron(stateยนโป_vol, stateยนโป_vol)) + + โˆ‚๐’โฑ = zero(๐’โฑ) + + โˆ‚๐’โฑยฒแต‰ = zero(๐’โฑยฒแต‰) + + โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) + + โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) + + โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) + + โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) + + โˆ‚๐’โปยน = zero(๐’โปยน) + + โˆ‚๐’โปยฒ = zero(๐’โปยฒ) + + โˆ‚๐’ยนโป = zero(๐’ยนโป) + + โˆ‚stateยนโป_vol = zero(stateยนโป_vol) + + โˆ‚x = zero(x[1]) + + โˆ‚state = [zeros(T.nPast_not_future_and_mixed), zeros(T.nPast_not_future_and_mixed)] + + kronSฮป = zeros(length(cond_var_idx) * T.nExo) + + kronxS = zeros(T.nExo * length(cond_var_idx)) + + โˆ‚๐’ = [zero(๐’[1]), zeros(size(๐’[2]))] + + function inversion_filter_loglikelihood_pullback(โˆ‚llh) + # @timeit_debug timer "Inversion filter pruned 2nd - pullback" begin + # @timeit_debug timer "Preallocation" begin + + fill!(โˆ‚๐’โฑ, 0) + fill!(โˆ‚๐’โฑยฒแต‰, 0) + + fill!(โˆ‚๐’ยนแต‰, 0) + fill!(โˆ‚๐’ยฒโปแต‰, 0) + + fill!(โˆ‚๐’ยนโปแต›, 0) + fill!(โˆ‚๐’ยฒโปแต›, 0) + + fill!(โˆ‚๐’โปยน, 0) + fill!(โˆ‚๐’โปยฒ, 0) + + fill!(โˆ‚๐’ยนโป, 0) + + fill!(โˆ‚stateยนโป_vol, 0) + fill!(โˆ‚x, 0) + fill!(โˆ‚state[1], 0) + fill!(โˆ‚state[2], 0) + + fill!(kronSฮป, 0) + fill!(kronxS, 0) + + # end # timeit_debug + # @timeit_debug timer "Main loop" begin + + for i in reverse(axes(data_in_deviations,2)) + # stateโ‚, stateโ‚‚ = [๐’โปยน * aug_stateโ‚[i], ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) / 2] + # stateโ‚ = ๐’โปยน * aug_stateโ‚[i] + # โˆ‚๐’โปยน += โˆ‚state[1] * aug_stateโ‚[i]' + โ„’.mul!(โˆ‚๐’โปยน, โˆ‚state[1], aug_stateโ‚[i]', 1, 1) + + # โˆ‚aug_stateโ‚ = ๐’โปยน' * โˆ‚state[1] + โ„’.mul!(โˆ‚aug_stateโ‚, ๐’โปยน', โˆ‚state[1]) + + # stateโ‚‚ = ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) / 2 + # โˆ‚๐’โปยน += โˆ‚state[2] * aug_stateโ‚‚[i]' + โ„’.mul!(โˆ‚๐’โปยน, โˆ‚state[2], aug_stateโ‚‚[i]', 1, 1) + + # โˆ‚aug_stateโ‚‚ = ๐’โปยน' * โˆ‚state[2] + โ„’.mul!(โˆ‚aug_stateโ‚‚, ๐’โปยน', โˆ‚state[2]) + + # โˆ‚๐’โปยฒ += โˆ‚state[2] * โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i])' / 2 + โ„’.kron!(kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) + โ„’.mul!(โˆ‚๐’โปยฒ, โˆ‚state[2], kronaug_stateโ‚', 1/2, 1) + + # โˆ‚kronaug_stateโ‚ = ๐’โปยฒ' * โˆ‚state[2] / 2 + โ„’.mul!(โˆ‚kronaug_stateโ‚, ๐’โปยฒ', โˆ‚state[2]) + โ„’.rdiv!(โˆ‚kronaug_stateโ‚, 2) + + fill_kron_adjoint!(โˆ‚aug_stateโ‚, โˆ‚aug_stateโ‚, โˆ‚kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) + + if i < size(data_in_deviations,2) + โˆ‚state[1] *= 0 + โˆ‚state[2] *= 0 + end + + # aug_stateโ‚ = [stateโ‚; 1; x] + # โˆ‚state[1] += โˆ‚aug_stateโ‚[1:length(โˆ‚state[1])] + โ„’.axpy!(1, โˆ‚aug_stateโ‚[1:length(โˆ‚state[1])], โˆ‚state[1]) + + โˆ‚x = โˆ‚aug_stateโ‚[T.nPast_not_future_and_mixed+2:end] + + # aug_stateโ‚‚ = [stateโ‚‚; 0; zero(x)] + # โˆ‚state[2] += โˆ‚aug_stateโ‚‚[1:length(โˆ‚state[1])] + โ„’.axpy!(1, โˆ‚aug_stateโ‚‚[1:length(โˆ‚state[1])], โˆ‚state[2]) + + # shocksยฒ += sum(abs2,x[i]) + if i < size(data_in_deviations,2) + โˆ‚x -= copy(x[i]) + else + โˆ‚x += copy(x[i]) + end + + # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] + โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) + inv(jacc[i])' + else + โ„’.pinv(jacc[i])' + end + catch + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x[1]) + # โˆ‚kronIx = ๐’โฑยฒแต‰' * โˆ‚jacc + โ„’.mul!(โˆ‚kronIx, ๐’โฑยฒแต‰', โˆ‚jacc) + + if i < size(data_in_deviations,2) + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -J) + else + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, J) + end + + # โˆ‚๐’โฑยฒแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' + โ„’.kron!(kron_buffer2, J, x[i]) + + โ„’.mul!(โˆ‚๐’โฑยฒแต‰, โˆ‚jacc, kron_buffer2', -1, 1) + + # find_shocks + โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) + # S = vcat(โˆ‚x, zero(ฮป[i])) + + S = fXฮปp[i]' \ โˆ‚xฮป + # โ„’.ldiv!(fXฮปp[i]', S) + + if i < size(data_in_deviations,2) + S *= -1 + end + + โˆ‚shock_independent = S[T.nExo+1:end] # fine + + # โˆ‚๐’โฑ = (S[1:T.nExo] * ฮป[i]' - S[T.nExo+1:end] * x[i]') # fine + # โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine + # copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) + โ„’.kron!(kronSฮป, S[1:T.nExo], ฮป[i]) + โ„’.kron!(kronxS, x[i], S[T.nExo+1:end]) + โ„’.axpy!(-1, kronxS, kronSฮป) + copyto!(โˆ‚๐’โฑ, kronSฮป) + # โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine + โ„’.axpy!(-1/2, โˆ‚jacc, โˆ‚๐’โฑ) + + โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], โ„’.kron(x[i], ฮป[i])) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) + # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo+1:end] * kronxx[i]' + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + โˆ‚stateยนโป_vol *= 0 + # โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ + โ„’.mul!(โˆ‚kronIstateยนโป_vol, ๐’ยฒโปแต‰', โˆ‚๐’โฑ) + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, J) + + stateยนโป_vol = aug_stateโ‚[i][1:T.nPast_not_future_and_mixed+1] + + # โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ + โ„’.axpy!(1, โˆ‚๐’โฑ, โˆ‚๐’ยนแต‰) + + # โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' + โ„’.kron!(โˆ‚kronIstateยนโป_vol, J, stateยนโป_vol) + โ„’.mul!(โˆ‚๐’ยฒโปแต‰, โˆ‚๐’โฑ, โˆ‚kronIstateยนโป_vol', 1, 1) + + + # shock_independent = copy(data_in_deviations[:,i]) + โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + # โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' + โ„’.mul!(โˆ‚๐’ยนโปแต›, โˆ‚shock_independent, stateยนโป_vol', -1, 1) + + # โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent + โ„’.mul!(โˆ‚stateยนโป_vol, ๐’ยนโปแต›', โˆ‚shock_independent, -1, 1) + + # โ„’.mul!(shock_independent, ๐’ยนโป, stateยฒโป, -1, 1) + # โˆ‚๐’ยนโป -= โˆ‚shock_independent * aug_stateโ‚‚[i][1:T.nPast_not_future_and_mixed]' + โ„’.mul!(โˆ‚๐’ยนโป, โˆ‚shock_independent, aug_stateโ‚‚[i][1:T.nPast_not_future_and_mixed]', -1, 1) + + # โˆ‚state[2] -= ๐’ยนโป' * โˆ‚shock_independent + โ„’.mul!(โˆ‚state[2], ๐’ยนโป', โˆ‚shock_independent, -1, 1) + + # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + # โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 + โ„’.kron!(โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) + โ„’.mul!(โˆ‚๐’ยฒโปแต›, โˆ‚shock_independent, โˆ‚kronstateยนโป_vol', -1/2, 1) + + # โˆ‚kronstateยนโป_vol = -๐’ยฒโปแต›' * โˆ‚shock_independent / 2 + โ„’.mul!(โˆ‚kronstateยนโป_vol, ๐’ยฒโปแต›', โˆ‚shock_independent) + โ„’.rdiv!(โˆ‚kronstateยนโป_vol, -2) + + fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) + + # stateยนโป_vol = vcat(stateยนโป, 1) + # โˆ‚state[1] += โˆ‚stateยนโป_vol[1:end-1] + โ„’.axpy!(1, โˆ‚stateยนโป_vol[1:end-1], โˆ‚state[1]) + end + + # end # timeit_debug + # @timeit_debug timer "Post allocation" begin + + fill!(โˆ‚๐’[1], 0) + fill!(โˆ‚๐’[2], 0) + + โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] .+= โˆ‚๐’ยนแต‰ + โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] .+= โˆ‚๐’ยฒโปแต‰ + โ„’.rdiv!(โˆ‚๐’โฑยฒแต‰, 2) + โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] .+= โˆ‚๐’โฑยฒแต‰# / 2 + + โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] .+= โˆ‚๐’ยนโปแต› + โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] .+= โˆ‚๐’ยฒโปแต› + + โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] .+= โˆ‚๐’โปยน + โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] .+= โˆ‚๐’โปยฒ + + โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] .+= โˆ‚๐’ยนโป + + # โˆ‚๐’[1] *= โˆ‚llh + # โˆ‚๐’[2] *= โˆ‚llh + โ„’.rmul!(โˆ‚๐’[1], โˆ‚llh) + โ„’.rmul!(โˆ‚๐’[2], โˆ‚llh) + + โ„’.rmul!(โˆ‚data_in_deviations, โˆ‚llh) + + โˆ‚state[1] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[1] * โˆ‚llh + โˆ‚state[2] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[2] * โˆ‚llh + + # end # timeit_debug + # end # timeit_debug + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’, โˆ‚data_in_deviations, NoTangent(), โˆ‚state, NoTangent() + end + + # See: https://pcubaborda.net/documents/CGIZ-final.pdf + llh = -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + + return llh, inversion_filter_loglikelihood_pullback +end + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:inversion}, + ::Val{:second_order}, + observables_index::Vector{Int}, + ๐’::Vector{AbstractMatrix{Float64}}, + data_in_deviations::Matrix{Float64}, + constants::constants, + state::Vector{Float64}, + workspaces::workspaces; + # timer::TimerOutput = TimerOutput(), + on_failure_loglikelihood = -Inf, + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + opts::CalculationOptions = merge_calculation_options(), + filter_algorithm::Symbol = :LagrangeNewton)# where S <: Real + T = constants.post_model_macro + ws = workspaces.inversion + # @timeit_debug timer "Inversion filter 2nd - forward" begin + + # @timeit_debug timer "Preallocation" begin + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + cond_var_idx = observables_index + + shocksยฒ = 0.0 + logabsdets = 0.0 + + cc = ensure_conditional_forecast_constants!(constants) + shock_idxs = cc.shock_idxs + shockยฒ_idxs = cc.shockยฒ_idxs + shockvarยฒ_idxs = cc.shockvarยฒ_idxs + var_volยฒ_idxs = cc.var_volยฒ_idxs + varยฒ_idxs = cc.varยฒ_idxs + + ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] + ๐’โปยนแต‰ = ๐’[1][T.past_not_future_and_mixed_idx,end-T.nExo+1:end] + ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] + ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] + ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] + + ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] + ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] + ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] + ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] + ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] + + ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› + ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป + ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ + ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ + ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ + + kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] + + J = โ„’.I(T.nExo) + + kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) + + kron_buffer3 = โ„’.kron(J, zeros(T.nPast_not_future_and_mixed + 1)) + + x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] + + stateยนโป = state[T.past_not_future_and_mixed_idx] + + stateยนโป_vol = vcat(stateยนโป, 1) + + kronstateยนโป_voltmp = โ„’.kron(stateยนโป_vol, stateยนโป_vol) + + kronstateยนโป_vol = [kronstateยนโป_voltmp for _ in 1:size(data_in_deviations,2)] + + shock_independent = zeros(size(data_in_deviations,1)) + + ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) + + ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 + + # aug_state_tmp = [zeros(T.nPast_not_future_and_mixed); 1; zeros(T.nExo)] + + aug_state = [[zeros(T.nPast_not_future_and_mixed); 1; zeros(T.nExo)] for _ in 1:size(data_in_deviations,2)] + + kronaug_state = [zeros((T.nPast_not_future_and_mixed + 1 + T.nExo)^2) for _ in 1:size(data_in_deviations,2)] + + tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[1])), x[1]) + + jacc = [zero(tmp) for _ in 1:size(data_in_deviations,2)] + + jacct = copy(tmp') + + ฮป = [zeros(size(tmp, 1)) for _ in 1:size(data_in_deviations,2)] + + ฮป[1] = tmp' \ x[1] * 2 + + fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' + -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] + + kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) + + kronxฮป = [kronxฮป_tmp for _ in 1:size(data_in_deviations,2)] + + tmp = zeros(size(๐’โฑ, 2) * size(๐’โฑ, 2)) + + lI = -2 * vec(โ„’.I(size(๐’โฑ, 2))) + + init_guess = zeros(size(๐’โฑ, 2)) + + # end # timeit_debug + # @timeit_debug timer "Main loop" begin + + @inbounds for i in axes(data_in_deviations,2) + # aug_state[i][1:T.nPast_not_future_and_mixed] = stateยนโป + copyto!(aug_state[i], 1, stateยนโป, 1) + + stateยนโป_vol = aug_state[i][1:T.nPast_not_future_and_mixed + 1] + # copyto!(stateยนโป_vol, 1, aug_state[i], 1, T.nPast_not_future_and_mixed + 1) + + copyto!(shock_independent, data_in_deviations[:,i]) + + โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + + โ„’.kron!(kronstateยนโป_vol[i], stateยนโป_vol, stateยนโป_vol) + + โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kronstateยนโป_vol[i], -1/2, 1) + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) + โ„’.kron!(kron_buffer3, J, stateยนโป_vol) + + โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer3) + + โ„’.axpy!(1, ๐’ยนแต‰, ๐’โฑ) + + init_guess *= 0 + + # @timeit_debug timer "Find shocks" begin + x[i], matched = find_shocks(Val(filter_algorithm), + init_guess, + kronxx[i], + kron_buffer2, + J, + ๐’โฑ, + ๐’โฑยฒแต‰, + shock_independent, + # max_iter = 100 + ) + # end # timeit_debug + + if !matched + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + โ„’.kron!(kron_buffer2, J, x[i]) + + โ„’.mul!(jacc[i], ๐’โฑยฒแต‰, kron_buffer2) + + โ„’.axpby!(1, ๐’โฑ, 2, jacc[i]) + # jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[i])), x[i]) + + copy!(jacct, jacc[i]') + + jacc_fact = try + โ„’.factorize(jacct) + catch + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + try + โ„’.ldiv!(ฮป[i], jacc_fact, x[i]) + catch + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # โ„’.ldiv!(ฮป[i], jacc_fact', x[i]) + โ„’.rmul!(ฮป[i], 2) + + # fXฮปp[i] = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) jacc[i]' + # -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + โ„’.mul!(tmp, ๐’โฑยฒแต‰', ฮป[i]) + โ„’.axpby!(1, lI, 2, tmp) + + fXฮปp[i][1:size(๐’โฑ, 2), 1:size(๐’โฑ, 2)] = tmp + fXฮปp[i][size(๐’โฑ, 2)+1:end, 1:size(๐’โฑ, 2)] = -jacc[i] + fXฮปp[i][1:size(๐’โฑ, 2), size(๐’โฑ, 2)+1:end] = jacct + + โ„’.kron!(kronxx[i], x[i], x[i]) + + โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) + + if i > presample_periods + # due to change of variables: jacobian determinant adjustment + if T.nExo == length(observables_index) + logabsdets += โ„’.logabsdet(jacc_fact)[1] + else + logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) + end + + shocksยฒ += sum(abs2, x[i]) + + if !isfinite(logabsdets) || !isfinite(shocksยฒ) + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + end + + # aug_state[i] = [stateยนโป; 1; x[i]] + # aug_state[i][1:T.nPast_not_future_and_mixed] = stateยนโป + # aug_state[i][end-T.nExo+1:end] = x[i] + copyto!(aug_state[i], 1, stateยนโป, 1) + copyto!(aug_state[i], length(stateยนโป) + 2, x[i], 1) + + โ„’.kron!(kronaug_state[i], aug_state[i], aug_state[i]) + โ„’.mul!(stateยนโป, ๐’โปยน, aug_state[i]) + โ„’.mul!(stateยนโป, ๐’โปยฒ, kronaug_state[i], 1/2 ,1) + end + + # end # timeit_debug + # end # timeit_debug + + โˆ‚aug_state = zero(aug_state[1]) + + โˆ‚kronaug_state = zero(kronaug_state[1]) + + โˆ‚kronstateยนโป_vol = zero(kronstateยนโป_vol[1]) + + + โˆ‚๐’ = [zero(๐’[1]), zero(๐’[2])] + + โˆ‚data_in_deviations = similar(data_in_deviations) + + โˆ‚kronIx = zero(โ„’.kron(โ„’.I(length(x[1])), x[1])) + + โˆ‚๐’โฑ = zero(๐’โฑ) + + โˆ‚๐’โฑยฒแต‰ = zero(๐’โฑยฒแต‰) + + โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) + + โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) + + โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) + + โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) + + โˆ‚๐’โปยน = zero(๐’โปยน) + + โˆ‚๐’โปยฒ = zero(๐’โปยฒ) + + โˆ‚stateยนโป_vol = zero(stateยนโป_vol) + + โˆ‚state = zeros(T.nPast_not_future_and_mixed) + + function inversion_filter_loglikelihood_pullback(โˆ‚llh) + # @timeit_debug timer "Inversion filter 2nd - pullback" begin + + # @timeit_debug timer "Preallocation" begin + + fill!(โˆ‚๐’โฑ, 0) + fill!(โˆ‚๐’โฑยฒแต‰, 0) + + # Allocate or reuse workspaces for pullback temps + if size(ws.โˆ‚๐’โฑยฒแต‰tmp) != (T.nExo, T.nExo * length(ฮป[1])) + ws.โˆ‚๐’โฑยฒแต‰tmp = zeros(T.nExo, T.nExo * length(ฮป[1])) + else + fill!(ws.โˆ‚๐’โฑยฒแต‰tmp, zero(eltype(ws.โˆ‚๐’โฑยฒแต‰tmp))) + end + โˆ‚๐’โฑยฒแต‰tmp = ws.โˆ‚๐’โฑยฒแต‰tmp + + if size(ws.โˆ‚๐’โฑยฒแต‰tmp2) != (length(ฮป[1]), T.nExo * T.nExo) + ws.โˆ‚๐’โฑยฒแต‰tmp2 = zeros(length(ฮป[1]), T.nExo * T.nExo) + else + fill!(ws.โˆ‚๐’โฑยฒแต‰tmp2, zero(eltype(ws.โˆ‚๐’โฑยฒแต‰tmp2))) + end + โˆ‚๐’โฑยฒแต‰tmp2 = ws.โˆ‚๐’โฑยฒแต‰tmp2 + + fill!(โˆ‚๐’ยนแต‰, 0) + fill!(โˆ‚๐’ยฒโปแต‰, 0) + + fill!(โˆ‚๐’ยนโปแต›, 0) + fill!(โˆ‚๐’ยฒโปแต›, 0) + + fill!(โˆ‚๐’โปยน, 0) + fill!(โˆ‚๐’โปยฒ, 0) + + fill!(โˆ‚stateยนโป_vol, 0) + # โˆ‚x = zero(x[1]) + fill!(โˆ‚state, 0) + + โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ + + # Allocate or reuse workspaces for kron products + if length(ws.kronSฮป) != length(cond_var_idx) * T.nExo + ws.kronSฮป = zeros(length(cond_var_idx) * T.nExo) + else + fill!(ws.kronSฮป, zero(eltype(ws.kronSฮป))) + end + kronSฮป = ws.kronSฮป + + if length(ws.kronxS) != T.nExo * length(cond_var_idx) + ws.kronxS = zeros(T.nExo * length(cond_var_idx)) + else + fill!(ws.kronxS, zero(eltype(ws.kronxS))) + end + kronxS = ws.kronxS + + # end # timeit_debug + # @timeit_debug timer "Main loop" begin + + for i in reverse(axes(data_in_deviations,2)) + # stt = ๐’โปยน * aug_state + ๐’โปยฒ * โ„’.kron(aug_state, aug_state) / 2 + # โˆ‚๐’โปยน += โˆ‚state * aug_state[i]' + โ„’.mul!(โˆ‚๐’โปยน, โˆ‚state, aug_state[i]', 1, 1) + + # โˆ‚๐’โปยฒ += โˆ‚state * kronaug_state[i]' / 2 + โ„’.mul!(โˆ‚๐’โปยฒ, โˆ‚state, kronaug_state[i]', 1/2, 1) + + โ„’.mul!(โˆ‚aug_state, ๐’โปยน', โˆ‚state) + # โˆ‚aug_state = ๐’โปยน' * โˆ‚state + + โ„’.mul!(โˆ‚kronaug_state, ๐’โปยฒ', โˆ‚state) + โ„’.rdiv!(โˆ‚kronaug_state, 2) + # โˆ‚kronaug_state = ๐’โปยฒ' * โˆ‚state / 2 + + fill_kron_adjoint!(โˆ‚aug_state, โˆ‚aug_state, โˆ‚kronaug_state, aug_state[i], aug_state[i]) + + if i < size(data_in_deviations,2) + โˆ‚state *= 0 + end + + # aug_state[i] = [stt; 1; x[i]] + โˆ‚state += โˆ‚aug_state[1:length(โˆ‚state)] + + # aug_state[i] = [stt; 1; x[i]] + โˆ‚x = โˆ‚aug_state[T.nPast_not_future_and_mixed+2:end] + + # shocksยฒ += sum(abs2,x[i]) + if i < size(data_in_deviations,2) + โˆ‚x -= copy(x[i]) + else + โˆ‚x += copy(x[i]) + end + + # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] + โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) + inv(jacc[i])' + else + โ„’.pinv(jacc[i])' + end + catch + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x[1]) + โ„’.mul!(โˆ‚kronIx, ๐’โฑยฒแต‰', โˆ‚jacc) + + if i < size(data_in_deviations,2) + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -J) + else + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, J) + end + + # โˆ‚๐’โฑยฒแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' + โ„’.kron!(kron_buffer2, J, x[i]) + + โ„’.mul!(โˆ‚๐’โฑยฒแต‰, โˆ‚jacc, kron_buffer2', -1, 1) + + # find_shocks + โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) + + S = fXฮปp[i]' \ โˆ‚xฮป + + if i < size(data_in_deviations,2) + S *= -1 + end + + โˆ‚shock_independent = S[T.nExo+1:end] # fine + + # โ„’.mul!(โˆ‚๐’โฑ, ฮป[i], S[1:T.nExo]') + # โ„’.mul!(โˆ‚๐’โฑ, S[T.nExo+1:end], x[i]', -1, 1) # fine + # โ„’.axpy!(-1/2, โˆ‚jacc, โˆ‚๐’โฑ) + # โˆ‚๐’โฑ = ฮป[i] * S[1:T.nExo]' - S[T.nExo+1:end] * x[i]' # fine + + # copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) + # โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine + โ„’.kron!(kronSฮป, S[1:T.nExo], ฮป[i]) + โ„’.kron!(kronxS, x[i], S[T.nExo+1:end]) + โ„’.axpy!(-1, kronxS, kronSฮป) + copyto!(โˆ‚๐’โฑ, kronSฮป) + + โ„’.axpy!(-1/2, โˆ‚jacc, โˆ‚๐’โฑ) + + โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], kronxฮป[i]) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) + # โ„’.mul!(โˆ‚๐’โฑยฒแต‰tmp, S[1:T.nExo], kronxฮป[i]', 2, 1) + # โ„’.mul!(โˆ‚๐’โฑยฒแต‰tmp2, S[T.nExo+1:end], kronxx[i]', -1, 1) + + # โ„’.mul!(โˆ‚๐’โฑยฒแต‰, S[1:T.nExo], kronxฮป[i]', 2, 1) + # โ„’.mul!(โˆ‚๐’โฑยฒแต‰, S[T.nExo+1:end], kronxx[i]', -1, 1) + # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo+1:end] * kronxx[i]' + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + โˆ‚stateยนโป_vol *= 0 + + โ„’.mul!(โˆ‚kronIstateยนโป_vol, ๐’ยฒโปแต‰', โˆ‚๐’โฑ) + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, J) + + stateยนโป_vol = aug_state[i][1:T.nPast_not_future_and_mixed + 1] + + โ„’.axpy!(1, โˆ‚๐’โฑ, โˆ‚๐’ยนแต‰) + # โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ + + โ„’.kron!(kron_buffer3, J, stateยนโป_vol) + + โ„’.mul!(โˆ‚๐’ยฒโปแต‰, โˆ‚๐’โฑ, kron_buffer3', 1, 1) + # โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' + + # shock_independent = copy(data_in_deviations[:,i]) + โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + # โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' + โ„’.mul!(โˆ‚๐’ยนโปแต›, โˆ‚shock_independent, stateยนโป_vol', -1 ,1) + + # โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent + โ„’.mul!(โˆ‚stateยนโป_vol, ๐’ยนโปแต›', โˆ‚shock_independent, -1, 1) + + # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + โ„’.kron!(kronstateยนโป_vol[i], stateยนโป_vol, stateยนโป_vol) + โ„’.mul!(โˆ‚๐’ยฒโปแต›, โˆ‚shock_independent, kronstateยนโป_vol[i]', -1/2, 1) + # โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 + + โ„’.mul!(โˆ‚kronstateยนโป_vol, ๐’ยฒโปแต›', โˆ‚shock_independent) + โ„’.rdiv!(โˆ‚kronstateยนโป_vol, -2) + # โˆ‚kronstateยนโป_vol = ๐’ยฒโปแต›' * โˆ‚shock_independent / (-2) + + fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) + + # stateยนโป_vol = vcat(stateยนโป, 1) + โˆ‚state += โˆ‚stateยนโป_vol[1:end-1] + end + + # end # timeit_debug + # @timeit_debug timer "Post allocation" begin + + fill!(โˆ‚๐’[1], 0) + fill!(โˆ‚๐’[2], 0) + + โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] += โˆ‚๐’ยนแต‰ + โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] += โˆ‚๐’ยฒโปแต‰ + โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] += โˆ‚๐’โฑยฒแต‰ / 2 + โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] += โˆ‚๐’ยนโปแต› + โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] += โˆ‚๐’ยฒโปแต› + + โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยน + โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยฒ + + โˆ‚๐’[1] *= โˆ‚llh + โˆ‚๐’[2] *= โˆ‚llh + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state * โˆ‚llh, NoTangent() + end + + # end # timeit_debug + # end # timeit_debug + + # See: https://pcubaborda.net/documents/CGIZ-final.pdf + llh = -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + + return llh, inversion_filter_loglikelihood_pullback +end + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:inversion}, + ::Val{:pruned_third_order}, + observables_index::Vector{Int}, + ๐’::Vector{AbstractMatrix{Float64}}, + data_in_deviations::Matrix{Float64}, + constants::constants, + state::Vector{Vector{Float64}}, + workspaces::workspaces; + # timer::TimerOutput = TimerOutput(), + on_failure_loglikelihood = -Inf, + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + opts::CalculationOptions = merge_calculation_options(), + filter_algorithm::Symbol = :LagrangeNewton) + T = constants.post_model_macro + ws = workspaces.inversion + # @timeit_debug timer "Inversion filter - forward" begin + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + cond_var_idx = observables_index + + shocksยฒ = 0.0 + logabsdets = 0.0 + + cc = ensure_conditional_forecast_constants!(constants; third_order = true) + tc = constants.third_order + # pruned variant needs kron(e, s_in_s) (no vol), not the cached kron(e, s_in_sโบ) + shockvar_idxs = sparse(โ„’.kron(cc.e_in_sโบ, cc.s_in_s)).nzind + shock_idxs = cc.shock_idxs + shockยฒ_idxs = cc.shockยฒ_idxs + shockvarยฒ_idxs = cc.shockvarยฒ_idxs + var_volยฒ_idxs = cc.var_volยฒ_idxs + varยฒ_idxs = cc.varยฒ_idxs + var_volยณ_idxs = tc.var_volยณ_idxs + shock_idxs2 = tc.shock_idxs2 + shock_idxs3 = tc.shock_idxs3 + shockยณ_idxs = tc.shockยณ_idxs + shockvar1_idxs = tc.shockvar1_idxs + shockvar2_idxs = tc.shockvar2_idxs + shockvar3_idxs = tc.shockvar3_idxs + shockvarยณ2_idxs = tc.shockvarยณ2_idxs + shockvarยณ_idxs = tc.shockvarยณ_idxs + + ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] + ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] + ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] + ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] + + ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] + ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] + ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] + ๐’ยฒโปแต›แต‰ = ๐’[2][cond_var_idx,shockvar_idxs] + ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] + ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] + + ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› + ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป + ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ + ๐’ยฒโปแต›แต‰ = nnz(๐’ยฒโปแต›แต‰) / length(๐’ยฒโปแต›แต‰) > .1 ? collect(๐’ยฒโปแต›แต‰) : ๐’ยฒโปแต›แต‰ + ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ + ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ + + ๐’ยณโปแต› = ๐’[3][cond_var_idx,var_volยณ_idxs] + ๐’ยณโปแต‰ยฒ = ๐’[3][cond_var_idx,shockvarยณ2_idxs] + ๐’ยณโปแต‰ = ๐’[3][cond_var_idx,shockvarยณ_idxs] + ๐’ยณแต‰ = ๐’[3][cond_var_idx,shockยณ_idxs] + ๐’โปยณ = ๐’[3][T.past_not_future_and_mixed_idx,:] + + ๐’ยณโปแต› = nnz(๐’ยณโปแต›) / length(๐’ยณโปแต›) > .1 ? collect(๐’ยณโปแต›) : ๐’ยณโปแต› + ๐’ยณโปแต‰ = nnz(๐’ยณโปแต‰) / length(๐’ยณโปแต‰) > .1 ? collect(๐’ยณโปแต‰) : ๐’ยณโปแต‰ + ๐’ยณแต‰ = nnz(๐’ยณแต‰) / length(๐’ยณแต‰) > .1 ? collect(๐’ยณแต‰) : ๐’ยณแต‰ + ๐’โปยณ = nnz(๐’โปยณ) / length(๐’โปยณ) > .1 ? collect(๐’โปยณ) : ๐’โปยณ + + stateโ‚ = state[1][T.past_not_future_and_mixed_idx] + stateโ‚‚ = state[2][T.past_not_future_and_mixed_idx] + stateโ‚ƒ = state[3][T.past_not_future_and_mixed_idx] + + kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] + + J = โ„’.I(T.nExo) + + II = sparse(โ„’.I(T.nExo^2)) + + kronxxx = [zeros(T.nExo^3) for _ in 1:size(data_in_deviations,2)] + + kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) + + kron_buffer3 = โ„’.kron(J, zeros(T.nExo^2)) + + kron_buffer4 = โ„’.kron(โ„’.kron(J, J), zeros(T.nExo)) + + x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] + + stateยนโป = stateโ‚ + + stateยนโป_vol = vcat(stateยนโป, 1) + + stateยฒโป = stateโ‚‚#[T.past_not_future_and_mixed_idx] + + stateยณโป = stateโ‚ƒ#[T.past_not_future_and_mixed_idx] + + ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + + ๐’โฑยฒแต‰ = [zero(๐’ยฒแต‰) for _ in 1:size(data_in_deviations,2)] + + aug_stateโ‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] + aug_stateโ‚ฬ‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] + aug_stateโ‚‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] + aug_stateโ‚ƒ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] + + kron_aug_stateโ‚ = [zeros(size(๐’โปยน,2)^2) for _ in 1:size(data_in_deviations,2)] + + jacc_tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰[1] * โ„’.kron(โ„’.I(T.nExo), x[1]) + + jacc = [zero(jacc_tmp) for _ in 1:size(data_in_deviations,2)] + + ฮป = [zeros(size(jacc_tmp, 1)) for _ in 1:size(data_in_deviations,2)] + + ฮป[1] = jacc_tmp' \ x[1] * 2 + + fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰[1]' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) jacc_tmp' + -jacc_tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] + + kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) + + kronxฮป = [kronxฮป_tmp for _ in 1:size(data_in_deviations,2)] + + kronxxฮป_tmp = โ„’.kron(x[1], kronxฮป_tmp) + + kronxxฮป = [kronxxฮป_tmp for _ in 1:size(data_in_deviations,2)] + + II = sparse(โ„’.I(T.nExo^2)) + + lI = 2 * โ„’.I(size(๐’โฑ, 2)) + + ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 + + # @timeit_debug timer "Loop" begin + for i in axes(data_in_deviations,2) + stateยนโป = stateโ‚ + + stateยนโป_vol = vcat(stateยนโป, 1) + + stateยฒโป = stateโ‚‚#[T.past_not_future_and_mixed_idx] + + stateยณโป = stateโ‚ƒ#[T.past_not_future_and_mixed_idx] + + shock_independent = copy(data_in_deviations[:,i]) + + โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + + โ„’.mul!(shock_independent, ๐’ยนโป, stateยฒโป, -1, 1) + + โ„’.mul!(shock_independent, ๐’ยนโป, stateยณโป, -1, 1) + + โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + + โ„’.mul!(shock_independent, ๐’ยฒโป, โ„’.kron(stateยนโป, stateยฒโป), -1, 1) + + โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) + + ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยฒโปแต›แต‰ * โ„’.kron(โ„’.I(T.nExo), stateยฒโป) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 + + ๐’โฑยฒแต‰[i] = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 + + ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 + + init_guess = zeros(size(๐’โฑ, 2)) + + # @timeit_debug timer "Find shocks" begin + x[i], matched = find_shocks(Val(filter_algorithm), + init_guess, + kronxx[i], + kronxxx[i], + kron_buffer2, + kron_buffer3, + kron_buffer4, + J, + ๐’โฑ, + ๐’โฑยฒแต‰[i], + ๐’โฑยณแต‰, + shock_independent, + # max_iter = 100 + ) + # end # timeit_debug + + if !matched + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰[i] * โ„’.kron(โ„’.I(T.nExo), x[i]) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), kronxx[i]) + + ฮป[i] = jacc[i]' \ x[i] * 2 + # โ„’.ldiv!(ฮป[i], tmp', x[i]) + # โ„’.rmul!(ฮป[i], 2) + fXฮปp[i] = [reshape((2 * ๐’โฑยฒแต‰[i] + 6 * ๐’โฑยณแต‰ * โ„’.kron(II, x[i]))' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - lI jacc[i]' + -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + โ„’.kron!(kronxx[i], x[i], x[i]) + + โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) + + โ„’.kron!(kronxxฮป[i], x[i], kronxฮป[i]) + + โ„’.kron!(kronxxx[i], x[i], kronxx[i]) + + if i > presample_periods + # due to change of variables: jacobian determinant adjustment + if T.nExo == length(observables_index) + logabsdets += โ„’.logabsdet(jacc[i])[1] + else + logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) + end + + shocksยฒ += sum(abs2,x[i]) + + if !isfinite(logabsdets) || !isfinite(shocksยฒ) + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + end + + aug_stateโ‚[i] = [stateโ‚; 1; x[i]] + aug_stateโ‚ฬ‚[i] = [stateโ‚; 0; x[i]] + aug_stateโ‚‚[i] = [stateโ‚‚; 0; zeros(T.nExo)] + aug_stateโ‚ƒ[i] = [stateโ‚ƒ; 0; zeros(T.nExo)] + + kron_aug_stateโ‚[i] = โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) + + stateโ‚, stateโ‚‚, stateโ‚ƒ = [๐’โปยน * aug_stateโ‚[i], ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * kron_aug_stateโ‚[i] / 2, ๐’โปยน * aug_stateโ‚ƒ[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i]) + ๐’โปยณ * โ„’.kron(kron_aug_stateโ‚[i], aug_stateโ‚[i]) / 6] + end + # end # timeit_debug + + # See: https://pcubaborda.net/documents/CGIZ-final.pdf + llh = -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + + + โˆ‚๐’ = [zero(๐’[1]), zero(๐’[2]), zero(๐’[3])] + + โˆ‚data_in_deviations = similar(data_in_deviations) + + # end # timeit_debug + + โˆ‚๐’โฑ = zero(๐’โฑ) + + โˆ‚๐’ยฒแต‰ = zero(๐’ยฒแต‰) + + โˆ‚๐’โฑยณแต‰ = zero(๐’โฑยณแต‰) + + โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) + + โˆ‚๐’ยนโป = zero(๐’ยนโป) + + โˆ‚๐’ยฒโป = zero(๐’ยฒโป) + + โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) + + โˆ‚๐’ยฒโปแต›แต‰ = zero(๐’ยฒโปแต›แต‰) + + โˆ‚๐’ยณโปแต‰ = zero(๐’ยณโปแต‰) + + โˆ‚๐’ยณโปแต‰ยฒ = zero(๐’ยณโปแต‰ยฒ) + + โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) + + โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) + + โˆ‚๐’ยณโปแต› = zero(๐’ยณโปแต›) + + โˆ‚๐’โปยน = zero(๐’โปยน) + + โˆ‚๐’โปยฒ = zero(๐’โปยฒ) + + โˆ‚๐’โปยณ = zero(๐’โปยณ) + + โˆ‚aug_stateโ‚ฬ‚ = zero(aug_stateโ‚ฬ‚[1]) + + โˆ‚stateยนโป_vol = zero(stateยนโป_vol) + + โˆ‚x = zero(x[1]) + + โˆ‚kronxx = zero(kronxx[1]) + + โˆ‚kronstateยนโป_vol = zeros(length(stateยนโป_vol)^2) + + โˆ‚state = [zeros(T.nPast_not_future_and_mixed), zeros(T.nPast_not_future_and_mixed), zeros(T.nPast_not_future_and_mixed)] + + function inversion_filter_loglikelihood_pullback(โˆ‚llh) + # @timeit_debug timer "Inversion filter - pullback" begin + fill!(โˆ‚๐’โฑ, 0) + fill!(โˆ‚๐’ยฒแต‰, 0) + fill!(โˆ‚๐’โฑยณแต‰, 0) + + fill!(โˆ‚๐’ยนแต‰, 0) + fill!(โˆ‚๐’ยนโป, 0) + fill!(โˆ‚๐’ยฒโป, 0) + fill!(โˆ‚๐’ยฒโปแต‰, 0) + fill!(โˆ‚๐’ยฒโปแต›แต‰, 0) + fill!(โˆ‚๐’ยณโปแต‰, 0) + fill!(โˆ‚๐’ยณโปแต‰ยฒ, 0) + + fill!(โˆ‚๐’ยนโปแต›, 0) + fill!(โˆ‚๐’ยฒโปแต›, 0) + fill!(โˆ‚๐’ยณโปแต›, 0) + + fill!(โˆ‚๐’โปยน, 0) + fill!(โˆ‚๐’โปยฒ, 0) + fill!(โˆ‚๐’โปยณ, 0) + + fill!(โˆ‚aug_stateโ‚ฬ‚, 0) + fill!(โˆ‚stateยนโป_vol, 0) + fill!(โˆ‚x, 0) + fill!(โˆ‚kronxx, 0) + fill!(โˆ‚kronstateยนโป_vol, 0) + fill!(โˆ‚state[1], 0) + fill!(โˆ‚state[2], 0) + fill!(โˆ‚state[3], 0) + + # @timeit_debug timer "Loop" begin + for i in reverse(axes(data_in_deviations,2)) + # stateโ‚ = ๐’โปยน * aug_stateโ‚[i] + โˆ‚๐’โปยน += โˆ‚state[1] * aug_stateโ‚[i]' + + โˆ‚aug_stateโ‚ = ๐’โปยน' * โˆ‚state[1] + + # stateโ‚‚ = ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * kron_aug_stateโ‚[i] / 2 + โˆ‚๐’โปยน += โˆ‚state[2] * aug_stateโ‚‚[i]' + + โˆ‚aug_stateโ‚‚ = ๐’โปยน' * โˆ‚state[2] + + โˆ‚๐’โปยฒ += โˆ‚state[2] * kron_aug_stateโ‚[i]' / 2 + + โˆ‚kronaug_stateโ‚ = ๐’โปยฒ' * โˆ‚state[2] / 2 + + # stateโ‚ƒ = ๐’โปยน * aug_stateโ‚ƒ[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i]) + ๐’โปยณ * โ„’.kron(kron_aug_stateโ‚[i],aug_stateโ‚[i]) / 6 + โˆ‚๐’โปยน += โˆ‚state[3] * aug_stateโ‚ƒ[i]' + + โˆ‚aug_stateโ‚ƒ = ๐’โปยน' * โˆ‚state[3] + + โˆ‚๐’โปยฒ += โˆ‚state[3] * โ„’.kron(aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i])' + + โˆ‚aug_stateโ‚ฬ‚ *= 0 + + โˆ‚kronaug_stateโ‚ฬ‚โ‚‚ = ๐’โปยฒ' * โˆ‚state[3] + + fill_kron_adjoint!(โˆ‚aug_stateโ‚ฬ‚, โˆ‚aug_stateโ‚‚, โˆ‚kronaug_stateโ‚ฬ‚โ‚‚, aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i]) + + โˆ‚๐’โปยณ += โˆ‚state[3] * โ„’.kron(kron_aug_stateโ‚[i],aug_stateโ‚[i])' / 6 + + โˆ‚kronkronaug_stateโ‚ = ๐’โปยณ' * โˆ‚state[3] / 6 + + fill_kron_adjoint!(โˆ‚aug_stateโ‚, โˆ‚kronaug_stateโ‚, โˆ‚kronkronaug_stateโ‚, aug_stateโ‚[i], kron_aug_stateโ‚[i]) + + # kron_aug_stateโ‚[i] = โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) + fill_kron_adjoint!(โˆ‚aug_stateโ‚, โˆ‚aug_stateโ‚, โˆ‚kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) + + if i < size(data_in_deviations,2) + โˆ‚state[1] *= 0 + โˆ‚state[2] *= 0 + โˆ‚state[3] *= 0 + end + + # aug_stateโ‚[i] = [stateโ‚; 1; x[i]] + โˆ‚state[1] += โˆ‚aug_stateโ‚[1:length(โˆ‚state[1])] + + โˆ‚x = โˆ‚aug_stateโ‚[T.nPast_not_future_and_mixed+2:end] + + # aug_stateโ‚ฬ‚[i] = [stateโ‚; 0; x[i]] + โˆ‚state[1] += โˆ‚aug_stateโ‚ฬ‚[1:length(โˆ‚state[1])] + + โˆ‚x += โˆ‚aug_stateโ‚ฬ‚[T.nPast_not_future_and_mixed+2:end] + + # aug_stateโ‚‚[i] = [stateโ‚‚; 0; zeros(T.nExo)] + โˆ‚state[2] += โˆ‚aug_stateโ‚‚[1:length(โˆ‚state[1])] + + # aug_stateโ‚ƒ[i] = [stateโ‚ƒ; 0; zeros(T.nExo)] + โˆ‚state[3] += โˆ‚aug_stateโ‚ƒ[1:length(โˆ‚state[1])] + + # shocksยฒ += sum(abs2,x[i]) + if i < size(data_in_deviations,2) + โˆ‚x -= copy(x[i]) + else + โˆ‚x += copy(x[i]) + end + + # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] + โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) + inv(jacc[i])' + else + โ„’.pinv(jacc[i])' + end + catch + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), โ„’.kron(x, x)) + # โˆ‚๐’โฑ = -โˆ‚jacc / 2 # fine + + โˆ‚kronIx = ๐’โฑยฒแต‰[i]' * โˆ‚jacc + + if i < size(data_in_deviations,2) + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -โ„’.I(T.nExo)) + else + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, โ„’.I(T.nExo)) + end + + โˆ‚๐’โฑยฒแต‰ = -โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' + + โˆ‚kronIxx = ๐’โฑยณแต‰' * โˆ‚jacc * 3 / 2 + + โˆ‚kronxx *= 0 + + if i < size(data_in_deviations,2) + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, -โ„’.I(T.nExo)) + else + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, โ„’.I(T.nExo)) + end + + fill_kron_adjoint!(โˆ‚x, โˆ‚x, โˆ‚kronxx, x[i], x[i]) + + โˆ‚๐’โฑยณแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), kronxx[i])' * 3 / 2 + + # find_shocks + โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) + + S = fXฮปp[i]' \ โˆ‚xฮป + + if i < size(data_in_deviations,2) + S *= -1 + end + + โˆ‚shock_independent = S[T.nExo+1:end] # fine + + # โˆ‚๐’โฑ += S[1:T.nExo] * ฮป[i]' - S[T.nExo + 1:end] * x[i]' # fine + copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) + โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine + + โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], โ„’.kron(x[i], ฮป[i])) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) + # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo + 1:end] * kronxx[i]' + + โˆ‚๐’โฑยณแต‰ += reshape(3 * โ„’.kron(S[1:T.nExo], โ„’.kron(โ„’.kron(x[i], x[i]), ฮป[i])) - โ„’.kron(kronxxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยณแต‰)) + # โˆ‚๐’โฑยณแต‰ += 3 * S[1:T.nExo] * kronxxฮป[i]' - S[T.nExo + 1:end] * kronxxx[i]' + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยฒโปแต›แต‰ * โ„’.kron(โ„’.I(T.nExo), stateยฒโป) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 + โˆ‚kronstateยนโป_vol *= 0 + + stateยนโป_vol = [aug_stateโ‚[i][1:T.nPast_not_future_and_mixed];1] # define here as it is used multiple times later + stateยนโป = aug_stateโ‚[i][1:T.nPast_not_future_and_mixed] + stateยฒโป = aug_stateโ‚‚[i][1:T.nPast_not_future_and_mixed] + stateยณโป = aug_stateโ‚ƒ[i][1:T.nPast_not_future_and_mixed] + + โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ + + โˆ‚stateยนโป_vol *= 0 + + โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, โ„’.I(T.nExo)) + + โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' + + โˆ‚kronIstateยฒโป = ๐’ยฒโปแต›แต‰' * โˆ‚๐’โฑ + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยฒโป, โˆ‚state[2], โ„’.I(T.nExo)) + + โˆ‚๐’ยฒโปแต›แต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยฒโป)' + + โˆ‚kronIstateยนโป_volstateยนโป_vol = ๐’ยณโปแต‰ยฒ' * โˆ‚๐’โฑ / 2 + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_volstateยนโป_vol, โˆ‚kronstateยนโป_vol, โ„’.I(T.nExo)) + + โˆ‚๐’ยณโปแต‰ยฒ += โˆ‚๐’โฑ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol)' / 2 + + # ๐’โฑยฒแต‰[i] = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 + โˆ‚๐’ยฒแต‰ += โˆ‚๐’โฑยฒแต‰ / 2 + + โˆ‚๐’ยณโปแต‰ += โˆ‚๐’โฑยฒแต‰ * โ„’.kron(II, stateยนโป_vol)' / 2 + + โˆ‚kronIIstateยนโป_vol = ๐’ยณโปแต‰' * โˆ‚๐’โฑยฒแต‰ / 2 + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIIstateยนโป_vol, โˆ‚stateยนโป_vol, II) + + # shock_independent = copy(data_in_deviations[:,i]) + โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' + + โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยนโป, stateยฒโป, -1, 1) + โˆ‚๐’ยนโป -= โˆ‚shock_independent * stateยฒโป' + + โˆ‚state[2] -= ๐’ยนโป' * โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยนโป, stateยณโป, -1, 1) + โˆ‚๐’ยนโป -= โˆ‚shock_independent * stateยณโป' + + โˆ‚state[3] -= ๐’ยนโป' * โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 + + โˆ‚kronstateยนโป_vol -= ๐’ยฒโปแต›' * โˆ‚shock_independent / 2 + + # โ„’.mul!(shock_independent, ๐’ยฒโป, โ„’.kron(stateยนโป, stateยฒโป), -1, 1) + โˆ‚๐’ยฒโป -= โˆ‚shock_independent * โ„’.kron(stateยนโป, stateยฒโป)' + + โˆ‚kronstateยนโปยฒโป = -๐’ยฒโป' * โˆ‚shock_independent + + fill_kron_adjoint!(โˆ‚state[1], โˆ‚state[2], โˆ‚kronstateยนโปยฒโป, stateยนโป, stateยฒโป) + + # โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) + โˆ‚๐’ยณโปแต› -= โˆ‚shock_independent * โ„’.kron(โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol)' / 6 + + โˆ‚kronstateยนโป_volstateยนโป_vol = -๐’ยณโปแต›' * โˆ‚shock_independent / 6 + + fill_kron_adjoint!(โˆ‚kronstateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_volstateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol) + + fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) + + # stateยนโป_vol = vcat(stateยนโป, 1) + โˆ‚state[1] += โˆ‚stateยนโป_vol[1:end-1] + end + # end # timeit_debug + + fill!(โˆ‚๐’[1], 0) + fill!(โˆ‚๐’[2], 0) + fill!(โˆ‚๐’[3], 0) + + โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] += โˆ‚๐’ยนแต‰ + โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] += โˆ‚๐’ยนโป + โˆ‚๐’[2][cond_var_idx,varยฒ_idxs] += โˆ‚๐’ยฒโป + โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] += โˆ‚๐’ยฒโปแต‰ + โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] += โˆ‚๐’ยฒแต‰ + โˆ‚๐’[2][cond_var_idx,shockvar_idxs] += โˆ‚๐’ยฒโปแต›แต‰ + โˆ‚๐’[3][cond_var_idx,shockvarยณ2_idxs] += โˆ‚๐’ยณโปแต‰ยฒ + โˆ‚๐’[3][cond_var_idx,shockvarยณ_idxs] += โˆ‚๐’ยณโปแต‰ + โˆ‚๐’[3][cond_var_idx,shockยณ_idxs] += โˆ‚๐’โฑยณแต‰ / 6 # ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 + + โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] += โˆ‚๐’ยนโปแต› + โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] += โˆ‚๐’ยฒโปแต› + โˆ‚๐’[3][cond_var_idx,var_volยณ_idxs] += โˆ‚๐’ยณโปแต› + + โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยน + โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยฒ + โˆ‚๐’[3][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยณ + + โˆ‚๐’[1] *= โˆ‚llh + โˆ‚๐’[2] *= โˆ‚llh + โˆ‚๐’[3] *= โˆ‚llh + + โˆ‚state[1] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[1] * โˆ‚llh + โˆ‚state[2] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[2] * โˆ‚llh + โˆ‚state[3] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[3] * โˆ‚llh + + # end # timeit_debug + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), โˆ‚state, NoTangent() + end + + return llh, inversion_filter_loglikelihood_pullback +end + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:inversion}, + ::Val{:third_order}, + observables_index::Vector{Int}, + ๐’::Vector{AbstractMatrix{Float64}}, + data_in_deviations::Matrix{Float64}, + constants::constants, + state::Vector{Float64}, + workspaces::workspaces; + # timer::TimerOutput = TimerOutput(), + on_failure_loglikelihood = -Inf, + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + opts::CalculationOptions = merge_calculation_options(), + filter_algorithm::Symbol = :LagrangeNewton) + T = constants.post_model_macro + ws = workspaces.inversion + # @timeit_debug timer "Inversion filter pruned 2nd - forward" begin + # @timeit_debug timer "Preallocation" begin + + precision_factor = 1.0 + + n_obs = size(data_in_deviations,2) + + cond_var_idx = observables_index + + shocksยฒ = 0.0 + logabsdets = 0.0 + + cc = ensure_conditional_forecast_constants!(constants; third_order = true) + tc = constants.third_order + shock_idxs = cc.shock_idxs + shockยฒ_idxs = cc.shockยฒ_idxs + shockvarยฒ_idxs = cc.shockvarยฒ_idxs + var_volยฒ_idxs = cc.var_volยฒ_idxs + varยฒ_idxs = cc.varยฒ_idxs + var_volยณ_idxs = tc.var_volยณ_idxs + shock_idxs2 = tc.shock_idxs2 + shock_idxs3 = tc.shock_idxs3 + shockยณ_idxs = tc.shockยณ_idxs + shockvar1_idxs = tc.shockvar1_idxs + shockvar2_idxs = tc.shockvar2_idxs + shockvar3_idxs = tc.shockvar3_idxs + shockvarยณ2_idxs = tc.shockvarยณ2_idxs + shockvarยณ_idxs = tc.shockvarยณ_idxs + + ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] + ๐’โปยนแต‰ = ๐’[1][T.past_not_future_and_mixed_idx,end-T.nExo+1:end] + ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] + ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] + ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] + + ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] + ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] + ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] + ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] + ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] + + ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› + ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป + ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ + ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ + ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ + + ๐’ยณโปแต› = ๐’[3][cond_var_idx,var_volยณ_idxs] + ๐’ยณโปแต‰ยฒ = ๐’[3][cond_var_idx,shockvarยณ2_idxs] + ๐’ยณโปแต‰ = ๐’[3][cond_var_idx,shockvarยณ_idxs] + ๐’ยณแต‰ = ๐’[3][cond_var_idx,shockยณ_idxs] + ๐’โปยณ = ๐’[3][T.past_not_future_and_mixed_idx,:] + + ๐’ยณโปแต› = nnz(๐’ยณโปแต›) / length(๐’ยณโปแต›) > .1 ? collect(๐’ยณโปแต›) : ๐’ยณโปแต› + ๐’ยณโปแต‰ = nnz(๐’ยณโปแต‰) / length(๐’ยณโปแต‰) > .1 ? collect(๐’ยณโปแต‰) : ๐’ยณโปแต‰ + ๐’ยณแต‰ = nnz(๐’ยณแต‰) / length(๐’ยณแต‰) > .1 ? collect(๐’ยณแต‰) : ๐’ยณแต‰ + ๐’โปยณ = nnz(๐’โปยณ) / length(๐’โปยณ) > .1 ? collect(๐’โปยณ) : ๐’โปยณ + + stt = state[T.past_not_future_and_mixed_idx] + + kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] + + J = โ„’.I(T.nExo) + + kronxxx = [zeros(T.nExo^3) for _ in 1:size(data_in_deviations,2)] + + kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) + + kron_buffer3 = โ„’.kron(J, zeros(T.nExo^2)) + + kron_buffer4 = โ„’.kron(โ„’.kron(J, J), zeros(T.nExo)) + + x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] + + stateยนโป = stt + + stateยนโป_vol = vcat(stateยนโป, 1) + + ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + + ๐’โฑยฒแต‰ = [zero(๐’ยฒแต‰) for _ in 1:size(data_in_deviations,2)] + + aug_state = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] + + tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰[1] * โ„’.kron(โ„’.I(T.nExo), x[1]) + + jacc = [zero(tmp) for _ in 1:size(data_in_deviations,2)] + + ฮป = [zeros(size(tmp, 1)) for _ in 1:size(data_in_deviations,2)] + + ฮป[1] = tmp' \ x[1] * 2 + + fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰[1]' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' + -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] + + kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) + + kronxฮป = [kronxฮป_tmp for _ in 1:size(data_in_deviations,2)] + + kronxxฮป_tmp = โ„’.kron(x[1], kronxฮป_tmp) + + kronxxฮป = [kronxxฮป_tmp for _ in 1:size(data_in_deviations,2)] + + II = sparse(โ„’.I(T.nExo^2)) + + lI = 2 * โ„’.I(size(๐’โฑ, 2)) + + ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 + + # end # timeit_debug + # @timeit_debug timer "Main loop" begin + + for i in axes(data_in_deviations,2) + stateยนโป = stt + + stateยนโป_vol = vcat(stateยนโป, 1) + + shock_independent = copy(data_in_deviations[:,i]) + + โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + + โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + + โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) + + ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 + + ๐’โฑยฒแต‰[i] = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 + + init_guess = zeros(size(๐’โฑ, 2)) + + # @timeit_debug timer "Find shocks" begin + x[i], matched = find_shocks(Val(filter_algorithm), + init_guess, + kronxx[i], + kronxxx[i], + kron_buffer2, + kron_buffer3, + kron_buffer4, + J, + ๐’โฑ, + ๐’โฑยฒแต‰[i], + ๐’โฑยณแต‰, + shock_independent, + # max_iter = 100 + ) + # end # timeit_debug + + if !matched + if opts.verbose println("Inversion filter failed at step $i") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰[i] * โ„’.kron(โ„’.I(T.nExo), x[i]) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), kronxx[i]) + + ฮป[i] = jacc[i]' \ x[i] * 2 + # โ„’.ldiv!(ฮป[i], tmp', x[i]) + # โ„’.rmul!(ฮป[i], 2) + fXฮปp[i] = [reshape((2 * ๐’โฑยฒแต‰[i] + 6 * ๐’โฑยณแต‰ * โ„’.kron(II, x[i]))' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - lI jacc[i]' + -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] + + โ„’.kron!(kronxx[i], x[i], x[i]) + + โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) + + โ„’.kron!(kronxxฮป[i], x[i], kronxฮป[i]) + + โ„’.kron!(kronxxx[i], x[i], kronxx[i]) + + if i > presample_periods + # due to change of variables: jacobian determinant adjustment + if T.nExo == length(observables_index) + logabsdets += โ„’.logabsdet(jacc[i])[1] + else + logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) + end + + shocksยฒ += sum(abs2,x[i]) + + if !isfinite(logabsdets) || !isfinite(shocksยฒ) + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + end + + aug_state[i] = [stt; 1; x[i]] + + stt = ๐’โปยน * aug_state[i] + ๐’โปยฒ * โ„’.kron(aug_state[i], aug_state[i]) / 2 + ๐’โปยณ * โ„’.kron(โ„’.kron(aug_state[i],aug_state[i]),aug_state[i]) / 6 + end + + # See: https://pcubaborda.net/documents/CGIZ-final.pdf + llh = -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + + # end # timeit_debug + # end # timeit_debug + + + โˆ‚๐’ = [zero(๐’[1]), zero(๐’[2]), zero(๐’[3])] + + โˆ‚data_in_deviations = similar(data_in_deviations) + + โˆ‚๐’โฑ = zero(๐’โฑ) + + โˆ‚๐’ยฒแต‰ = zero(๐’ยฒแต‰) + + โˆ‚๐’โฑยณแต‰ = zero(๐’โฑยณแต‰) + + โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) + + โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) + + โˆ‚๐’ยณโปแต‰ = zero(๐’ยณโปแต‰) + + โˆ‚๐’ยณโปแต‰ยฒ = zero(๐’ยณโปแต‰ยฒ) + + โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) + + โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) + + โˆ‚๐’ยณโปแต› = zero(๐’ยณโปแต›) + + โˆ‚๐’โปยน = zero(๐’โปยน) + + โˆ‚๐’โปยฒ = zero(๐’โปยฒ) + + โˆ‚๐’โปยณ = zero(๐’โปยณ) + + โˆ‚stateยนโป_vol = zero(stateยนโป_vol) + + โˆ‚x = zero(x[1]) + + โˆ‚kronxx = zero(kronxx[1]) + + โˆ‚kronstateยนโป_vol = zeros(length(stateยนโป_vol)^2) + + โˆ‚state = zeros(T.nPast_not_future_and_mixed) + + function inversion_filter_loglikelihood_pullback(โˆ‚llh) + # @timeit_debug timer "Inversion filter pruned 2nd - pullback" begin + # @timeit_debug timer "Preallocation" begin + + fill!(โˆ‚๐’โฑ, 0) + fill!(โˆ‚๐’ยฒแต‰, 0) + fill!(โˆ‚๐’โฑยณแต‰, 0) + + fill!(โˆ‚๐’ยนแต‰, 0) + fill!(โˆ‚๐’ยฒโปแต‰, 0) + fill!(โˆ‚๐’ยณโปแต‰, 0) + fill!(โˆ‚๐’ยณโปแต‰ยฒ, 0) + + fill!(โˆ‚๐’ยนโปแต›, 0) + fill!(โˆ‚๐’ยฒโปแต›, 0) + fill!(โˆ‚๐’ยณโปแต›, 0) + + fill!(โˆ‚๐’โปยน, 0) + fill!(โˆ‚๐’โปยฒ, 0) + fill!(โˆ‚๐’โปยณ, 0) + + fill!(โˆ‚stateยนโป_vol, 0) + fill!(โˆ‚x, 0) + fill!(โˆ‚kronxx, 0) + fill!(โˆ‚kronstateยนโป_vol, 0) + fill!(โˆ‚state, 0) + + # end # timeit_debug + # @timeit_debug timer "Main loop" begin + + for i in reverse(axes(data_in_deviations,2)) + # stt = ๐’โปยน * aug_state[i] + ๐’โปยฒ * โ„’.kron(aug_state[i], aug_state[i]) / 2 + ๐’โปยณ * โ„’.kron(โ„’.kron(aug_state[i],aug_state[i]),aug_state[i]) / 6 + โˆ‚๐’โปยน += โˆ‚state * aug_state[i]' + + โˆ‚๐’โปยฒ += โˆ‚state * โ„’.kron(aug_state[i], aug_state[i])' / 2 + + โˆ‚๐’โปยณ += โˆ‚state * โ„’.kron(โ„’.kron(aug_state[i], aug_state[i]), aug_state[i])' / 6 + + โˆ‚aug_state = ๐’โปยน' * โˆ‚state + โˆ‚kronaug_state = ๐’โปยฒ' * โˆ‚state / 2 + โˆ‚kronkronaug_state = ๐’โปยณ' * โˆ‚state / 6 + + fill_kron_adjoint!(โˆ‚aug_state, โˆ‚kronaug_state, โˆ‚kronkronaug_state, aug_state[i], โ„’.kron(aug_state[i], aug_state[i])) + + fill_kron_adjoint!(โˆ‚aug_state, โˆ‚aug_state, โˆ‚kronaug_state, aug_state[i], aug_state[i]) + + if i < size(data_in_deviations,2) + โˆ‚state *= 0 + end + + # aug_state[i] = [stt; 1; x[i]] + โˆ‚state += โˆ‚aug_state[1:length(โˆ‚state)] + + # aug_state[i] = [stt; 1; x[i]] + โˆ‚x = โˆ‚aug_state[T.nPast_not_future_and_mixed+2:end] + + # shocksยฒ += sum(abs2,x[i]) + if i < size(data_in_deviations,2) + โˆ‚x -= copy(x[i]) + else + โˆ‚x += copy(x[i]) + end + + # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] + โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) + inv(jacc[i])' + else + โ„’.pinv(jacc[i])' + end + catch + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), โ„’.kron(x, x)) + # โˆ‚๐’โฑ = -โˆ‚jacc / 2 # fine + + โˆ‚kronIx = ๐’โฑยฒแต‰[i]' * โˆ‚jacc + + if i < size(data_in_deviations,2) + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -โ„’.I(T.nExo)) + else + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, โ„’.I(T.nExo)) + end + + โˆ‚๐’โฑยฒแต‰ = -โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' + + โˆ‚kronIxx = ๐’โฑยณแต‰' * โˆ‚jacc * 3 / 2 + + โˆ‚kronxx *= 0 + + if i < size(data_in_deviations,2) + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, -โ„’.I(T.nExo)) + else + fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, โ„’.I(T.nExo)) + end + + fill_kron_adjoint!(โˆ‚x, โˆ‚x, โˆ‚kronxx, x[i], x[i]) + + โˆ‚๐’โฑยณแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), kronxx[i])' * 3 / 2 + + # find_shocks + โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) + + S = fXฮปp[i]' \ โˆ‚xฮป + + if i < size(data_in_deviations,2) + S *= -1 + end + + โˆ‚shock_independent = S[T.nExo+1:end] # fine + + # โˆ‚๐’โฑ += S[1:T.nExo] * ฮป[i]' - S[T.nExo + 1:end] * x[i]' # fine + copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) + โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine + + โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], โ„’.kron(x[i], ฮป[i])) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) + # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo + 1:end] * kronxx[i]' + + โˆ‚๐’โฑยณแต‰ += reshape(3 * โ„’.kron(S[1:T.nExo], โ„’.kron(โ„’.kron(x[i], x[i]), ฮป[i])) - โ„’.kron(kronxxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยณแต‰)) + # โˆ‚๐’โฑยณแต‰ += 3 * S[1:T.nExo] * kronxxฮป[i]' - S[T.nExo + 1:end] * kronxxx[i]' + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 + โˆ‚kronstateยนโป_vol *= 0 + + stateยนโป_vol = [aug_state[i][1:T.nPast_not_future_and_mixed];1] # define here as it is used multiple times later + + โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ + + โˆ‚stateยนโป_vol *= 0 + + โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, โ„’.I(T.nExo)) + + โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' + + โˆ‚kronIstateยนโป_volstateยนโป_vol = ๐’ยณโปแต‰ยฒ' * โˆ‚๐’โฑ / 2 + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_volstateยนโป_vol, โˆ‚kronstateยนโป_vol, โ„’.I(T.nExo)) + + โˆ‚๐’ยณโปแต‰ยฒ += โˆ‚๐’โฑ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol)' / 2 + + + # ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 + โˆ‚๐’ยฒแต‰ += โˆ‚๐’โฑยฒแต‰ / 2 + + โˆ‚๐’ยณโปแต‰ += โˆ‚๐’โฑยฒแต‰ * โ„’.kron(II, stateยนโป_vol)' / 2 + + โˆ‚kronIIstateยนโป_vol = ๐’ยณโปแต‰' * โˆ‚๐’โฑยฒแต‰ / 2 + + fill_kron_adjoint_โˆ‚A!(โˆ‚kronIIstateยนโป_vol, โˆ‚stateยนโป_vol, II) + + # shock_independent = copy(data_in_deviations[:,i]) + โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent + + + # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) + โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' + + โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent + + # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 + + โˆ‚kronstateยนโป_vol -= ๐’ยฒโปแต›' * โˆ‚shock_independent / 2 + + # โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) + โˆ‚๐’ยณโปแต› -= โˆ‚shock_independent * โ„’.kron(โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol)' / 6 + + โˆ‚kronstateยนโป_volstateยนโป_vol = -๐’ยณโปแต›' * โˆ‚shock_independent / 6 + + fill_kron_adjoint!(โˆ‚kronstateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_volstateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol) + + fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) + + # stateยนโป_vol = vcat(stateยนโป, 1) + โˆ‚state += โˆ‚stateยนโป_vol[1:end-1] + end + + # end # timeit_debug + # @timeit_debug timer "Post allocation" begin + + fill!(โˆ‚๐’[1], 0) + fill!(โˆ‚๐’[2], 0) + fill!(โˆ‚๐’[3], 0) + + โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] += โˆ‚๐’ยนแต‰ + โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] += โˆ‚๐’ยฒโปแต‰ + โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] += โˆ‚๐’ยฒแต‰ + โˆ‚๐’[3][cond_var_idx,shockvarยณ2_idxs] += โˆ‚๐’ยณโปแต‰ยฒ + โˆ‚๐’[3][cond_var_idx,shockvarยณ_idxs] += โˆ‚๐’ยณโปแต‰ + โˆ‚๐’[3][cond_var_idx,shockยณ_idxs] += โˆ‚๐’โฑยณแต‰ / 6 # ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 + + โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] += โˆ‚๐’ยนโปแต› + โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] += โˆ‚๐’ยฒโปแต› + โˆ‚๐’[3][cond_var_idx,var_volยณ_idxs] += โˆ‚๐’ยณโปแต› + + โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยน + โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยฒ + โˆ‚๐’[3][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยณ + + โˆ‚๐’[1] *= โˆ‚llh + โˆ‚๐’[2] *= โˆ‚llh + โˆ‚๐’[3] *= โˆ‚llh + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state * โˆ‚llh, NoTangent() + end + + # end # timeit_debug + # end # timeit_debug + + return llh, inversion_filter_loglikelihood_pullback +end + +function rrule(::typeof(calculate_loglikelihood), + ::Val{:kalman}, + ::Val, + observables_index::Vector{Int}, + ๐’::AbstractMatrix{Float64}, + data_in_deviations::Matrix{Float64}, + constants::constants, + state, + workspaces::workspaces; + warmup_iterations::Int = 0, + presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, + lyapunov_algorithm::Symbol = :doubling, + on_failure_loglikelihood::U = -Inf, + opts::CalculationOptions = merge_calculation_options()) where {U <: AbstractFloat} + + T = constants.post_model_macro + idx_constants = constants.post_complete_parameters + lyap_ws = ensure_lyapunov_workspace!(workspaces, T.nVars, :first_order) + observables_and_states = sort(union(T.past_not_future_and_mixed_idx, observables_index)) + observables_sorted = sort(observables_index) + I_nVars = idx_constants.diag_nVars + + A_map = @views I_nVars[T.past_not_future_and_mixed_idx, observables_and_states] + + A = @views ๐’[observables_and_states,1:T.nPast_not_future_and_mixed] * A_map + B = @views ๐’[observables_and_states,T.nPast_not_future_and_mixed+1:end] + + C = @views I_nVars[observables_sorted, observables_and_states] + + kalman_ws = ensure_kalman_workspaces!(workspaces, size(C, 1), size(C, 2)) + ๐ = kalman_ws.๐ + โ„’.mul!(๐, B, B') + + lyap_pullback = nothing + P = if initial_covariance == :theoretical + lyap_rrule_result, lyap_pullback_local = rrule(solve_lyapunov_equation, + A, + ๐, + lyap_ws, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.first_order.ad.lyapunov, + verbose = opts.verbose) + lyap_pullback = lyap_pullback_local + lyap_rrule_result[1] + else + get_initial_covariance(Val(initial_covariance), A, ๐, lyap_ws, opts = opts) + end + + Tt = size(data_in_deviations, 2) + 1 + + z = zeros(size(data_in_deviations, 1)) + uฬ„ = zeros(size(C,2)) + Pฬ„ = deepcopy(P) + + temp_N_N = similar(P) + PCtmp = similar(P, size(P, 1), size(C, 1)) + F = similar(P, size(C, 1), size(C, 1)) + + u = [similar(uฬ„) for _ in 1:Tt] + P_seq = [copy(Pฬ„) for _ in 1:Tt] + CP = [zeros(eltype(P), size(C, 1), size(P, 2)) for _ in 1:Tt] + K = [similar(P, size(P, 1), size(C, 1)) for _ in 1:Tt] + invF = [similar(F) for _ in 1:Tt] + v = [zeros(size(data_in_deviations, 1)) for _ in 1:Tt] + + loglik = 0.0 + + for t in 2:Tt + if !all(isfinite.(z)) + if opts.verbose println("KF not finite at step $t") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + v[t] .= data_in_deviations[:, t-1] .- z + + โ„’.mul!(CP[t], C, Pฬ„) + โ„’.mul!(F, CP[t], C') + + kalman_ws.fast_lu_ws_f, kalman_ws.fast_lu_dims_f, solved_F, luF = factorize_lu!(F, + kalman_ws.fast_lu_ws_f, + kalman_ws.fast_lu_dims_f) + + if !solved_F + if opts.verbose println("KF factorisation failed step $t") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + logabsdetF = 0.0 + signF = isodd(count(i -> kalman_ws.fast_lu_ws_f.ipiv[i] != i, eachindex(kalman_ws.fast_lu_ws_f.ipiv))) ? -1.0 : 1.0 + @inbounds for i in 1:size(F, 1) + di = F[i, i] + if di == 0 + if opts.verbose println("KF factorisation failed step $t") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + logabsdetF += log(abs(di)) + signF *= sign(di) + end + + if signF <= 0 || logabsdetF < log(eps(Float64)) + if opts.verbose println("KF factorisation failed step $t") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + fill!(invF[t], 0.0) + @inbounds for i in 1:size(invF[t], 1) + invF[t][i, i] = 1.0 + end + solve_lu_left!(F, invF[t], kalman_ws.fast_lu_ws_f, luF) + + if t - 1 > presample_periods + loglik += logabsdetF + โ„’.dot(v[t], invF[t], v[t]) + end + + โ„’.mul!(PCtmp, Pฬ„, C') + copyto!(K[t], PCtmp) + solve_lu_right!(F, K[t], kalman_ws.fast_lu_ws_f, luF, kalman_ws.fast_lu_rhs_t_k) + + โ„’.mul!(P_seq[t], K[t], CP[t], -1, 0) + P_seq[t] .+= Pฬ„ + + โ„’.mul!(temp_N_N, P_seq[t], A') + โ„’.mul!(Pฬ„, A, temp_N_N) + Pฬ„ .+= ๐ + + โ„’.mul!(u[t], K[t], v[t]) + u[t] .+= uฬ„ + + โ„’.mul!(uฬ„, A, u[t]) + โ„’.mul!(z, C, uฬ„) + end + + llh = -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 + + โˆ‚F = zero(F) + โˆ‚Faccum = zero(F) + โˆ‚P = zero(Pฬ„) + โˆ‚uฬ„ = zero(uฬ„) + โˆ‚v = zero(v[1]) + โˆ‚data_in_deviations = zero(data_in_deviations) + vtmp = zero(v[1]) + Ptmp = zero(P_seq[1]) + โˆ‚A_kf = zero(A) + โˆ‚๐_kf = zero(๐) + + function calculate_loglikelihood_pullback(โˆ‚llh) + โ„’.rmul!(โˆ‚A_kf, 0) + โ„’.rmul!(โˆ‚Faccum, 0) + โ„’.rmul!(โˆ‚P, 0) + โ„’.rmul!(โˆ‚uฬ„, 0) + โ„’.rmul!(โˆ‚๐_kf, 0) + + for t in Tt:-1:2 + if t > presample_periods + 1 + โ„’.mul!(โˆ‚F, v[t], v[t]') + โ„’.mul!(invF[1], invF[t]', โˆ‚F) + โ„’.mul!(โˆ‚F, invF[1], invF[t]') + โ„’.axpby!(1, invF[t]', -1, โˆ‚F) + + copy!(invF[1], invF[t]' .+ invF[t]) + โ„’.mul!(โˆ‚v, invF[1], v[t]) + else + โ„’.rmul!(โˆ‚F, 0) + โ„’.rmul!(โˆ‚v, 0) + end + + โ„’.axpy!(1, โˆ‚Faccum, โˆ‚F) + โ„’.mul!(PCtmp, C', โˆ‚F) + โ„’.mul!(โˆ‚P, PCtmp, C, 1, 1) + + โ„’.mul!(CP[1], invF[t]', C) + โ„’.mul!(PCtmp, โˆ‚uฬ„, v[t]') + โ„’.mul!(P_seq[1], PCtmp, CP[1]) + โ„’.mul!(โˆ‚P, A', P_seq[1], 1, 1) + + โ„’.mul!(u[1], A', โˆ‚uฬ„) + โ„’.mul!(v[1], K[t]', u[1]) + โ„’.axpy!(1, โˆ‚v, v[1]) + โˆ‚data_in_deviations[:,t-1] .= v[1] + + โ„’.mul!(u[1], A', โˆ‚uฬ„) + โ„’.mul!(v[1], K[t]', u[1]) + โ„’.mul!(โˆ‚uฬ„, C', v[1]) + โ„’.mul!(u[1], C', v[1], -1, 1) + copy!(โˆ‚uฬ„, u[1]) + + โ„’.mul!(u[1], C', โˆ‚v) + โ„’.axpy!(-1, u[1], โˆ‚uฬ„) + + if t > 2 + โ„’.mul!(โˆ‚A_kf, โˆ‚uฬ„, u[t-1]', 1, 1) + + โ„’.mul!(P_seq[1], A, P_seq[t-1]') + โ„’.mul!(Ptmp, โˆ‚P, P_seq[1]) + โ„’.mul!(P_seq[1], A, P_seq[t-1]) + โ„’.mul!(Ptmp, โˆ‚P', P_seq[1], 1, 1) + โ„’.axpy!(1, Ptmp, โˆ‚A_kf) + + โ„’.axpy!(1, โˆ‚P, โˆ‚๐_kf) + + โ„’.mul!(P_seq[1], โˆ‚P, A) + โ„’.mul!(โˆ‚P, A', P_seq[1]) + + โ„’.mul!(PCtmp, โˆ‚P, K[t-1]) + โ„’.mul!(CP[1], K[t-1]', โˆ‚P) + โ„’.mul!(โˆ‚P, PCtmp, C, -1, 1) + โ„’.mul!(โˆ‚P, C', CP[1], -1, 1) + + โ„’.mul!(u[1], A', โˆ‚uฬ„) + โ„’.mul!(v[1], CP[t-1], u[1]) + โ„’.mul!(vtmp, invF[t-1]', v[1], -1, 0) + โ„’.mul!(invF[1], vtmp, v[t-1]') + โ„’.mul!(โˆ‚Faccum, invF[1], invF[t-1]') + + โ„’.mul!(CP[1], invF[t-1]', CP[t-1]) + โ„’.mul!(PCtmp, CP[t-1]', invF[t-1]') + โ„’.mul!(K[1], โˆ‚P, PCtmp) + โ„’.mul!(โˆ‚Faccum, CP[1], K[1], -1, 1) + end + end + + โ„’.rmul!(โˆ‚P, -โˆ‚llh/2) + โ„’.rmul!(โˆ‚A_kf, -โˆ‚llh/2) + โ„’.rmul!(โˆ‚๐_kf, -โˆ‚llh/2) + โ„’.rmul!(โˆ‚data_in_deviations, -โˆ‚llh/2) + + โˆ‚A = copy(โˆ‚A_kf) + โˆ‚๐ = copy(โˆ‚๐_kf) + + if !isnothing(lyap_pullback) + lyap_grads = lyap_pullback((โˆ‚P, NoTangent())) + if !(lyap_grads[2] isa AbstractZero) + โ„’.axpy!(1, lyap_grads[2], โˆ‚A) + end + if !(lyap_grads[3] isa AbstractZero) + โ„’.axpy!(1, lyap_grads[3], โˆ‚๐) + end + end + + โˆ‚B = (โˆ‚๐ + โˆ‚๐') * B + + โˆ‚๐’ = zero(๐’) + @views โˆ‚๐’[observables_and_states, 1:T.nPast_not_future_and_mixed] .+= โˆ‚A * A_map' + @views โˆ‚๐’[observables_and_states, T.nPast_not_future_and_mixed+1:end] .+= โˆ‚B + + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’, โˆ‚data_in_deviations, NoTangent(), NoTangent(), NoTangent() + end + + return llh, calculate_loglikelihood_pullback +end + + +function _get_statistics_cotangent(ฮ”ret, key::Symbol) + ฮ” = unthunk(ฮ”ret) + if ฮ” isa Union{NoTangent, AbstractZero} + return NoTangent() + end + + if ฮ” isa AbstractDict + return get(ฮ”, key, NoTangent()) + end + + if ฮ” isa NamedTuple + return get(ฮ”, key, NoTangent()) + end + + if hasproperty(ฮ”, key) + return getproperty(ฮ”, key) + end + + if hasmethod(haskey, Tuple{typeof(ฮ”), Symbol}) && haskey(ฮ”, key) + return ฮ”[key] + end + + if hasmethod(pairs, Tuple{typeof(ฮ”)}) + for (k, v) in pairs(ฮ”) + if k == key + return v + end + end + end + + if hasproperty(ฮ”, :pairs) + pairs_obj = getproperty(ฮ”, :pairs) + if pairs_obj isa AbstractDict + return get(pairs_obj, key, NoTangent()) + elseif pairs_obj isa NamedTuple + return get(pairs_obj, key, NoTangent()) + elseif hasmethod(pairs, Tuple{typeof(pairs_obj)}) + for (k, v) in pairs(pairs_obj) + if k == key + return v + end + end + end + end + + return NoTangent() +end + + +function rrule(::typeof(get_statistics), + ๐“‚::โ„ณ, + parameter_values::Vector{T}; + parameters::Union{Vector{Symbol},Vector{String}} = ๐“‚.constants.post_complete_parameters.parameters, + steady_state_function::SteadyStateFunctionType = missing, + non_stochastic_steady_state::Union{Symbol_input,String_input} = Symbol[], + mean::Union{Symbol_input,String_input} = Symbol[], + standard_deviation::Union{Symbol_input,String_input} = Symbol[], + variance::Union{Symbol_input,String_input} = Symbol[], + covariance::Union{Symbol_input,String_input, Vector{Vector{Symbol}},Vector{Tuple{Symbol,Vararg{Symbol}}},Vector{Vector{Symbol}},Tuple{Tuple{Symbol,Vararg{Symbol}},Vararg{Tuple{Symbol,Vararg{Symbol}}}}, Vector{Vector{String}},Vector{Tuple{String,Vararg{String}}},Vector{Vector{String}},Tuple{Tuple{String,Vararg{String}},Vararg{Tuple{String,Vararg{String}}}}} = Symbol[], + autocorrelation::Union{Symbol_input,String_input} = Symbol[], + autocorrelation_periods::UnitRange{Int} = DEFAULT_AUTOCORRELATION_PERIODS, + algorithm::Symbol = DEFAULT_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(๐“‚), + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + verbose::Bool = DEFAULT_VERBOSE, + tol::Tolerances = Tolerances()) where T + + opts = merge_calculation_options(tol = tol, + verbose = verbose, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + sylvester_algorithmยฒ = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], + sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? sum(k * (k + 1) รท 2 for k in 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + 1 + ๐“‚.constants.post_model_macro.nExo) > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : sylvester_algorithm[2], + lyapunov_algorithm = lyapunov_algorithm) + + @assert length(parameter_values) == length(parameters) "Vector of `parameters` must correspond to `parameter_values` in length and order. Define the parameter names in the `parameters` keyword argument." + + @assert algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] || !(!(standard_deviation == Symbol[]) || !(mean == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[])) "Statistics can only be provided for first order perturbation or second and third order pruned perturbation solutions." + + @assert !(non_stochastic_steady_state == Symbol[]) || !(standard_deviation == Symbol[]) || !(mean == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[]) "Provide variables for at least one output." + + SS_var_idx = parse_variables_input_to_index(non_stochastic_steady_state, ๐“‚) + mean_var_idx = parse_variables_input_to_index(mean, ๐“‚) + std_var_idx = parse_variables_input_to_index(standard_deviation, ๐“‚) + var_var_idx = parse_variables_input_to_index(variance, ๐“‚) + covar_var_idx = parse_variables_input_to_index(covariance, ๐“‚) + covar_groups = is_grouped_covariance_input(covariance) ? parse_covariance_groups(covariance, ๐“‚.constants) : nothing + autocorr_var_idx = parse_variables_input_to_index(autocorrelation, ๐“‚) + + other_parameter_values = ๐“‚.parameter_values[indexin(setdiff(๐“‚.constants.post_complete_parameters.parameters, parameters), ๐“‚.constants.post_complete_parameters.parameters)] + sort_idx = sortperm(vcat(indexin(setdiff(๐“‚.constants.post_complete_parameters.parameters, parameters), ๐“‚.constants.post_complete_parameters.parameters), indexin(parameters, ๐“‚.constants.post_complete_parameters.parameters))) + + all_parameters = vcat(other_parameter_values, parameter_values)[sort_idx] + n_other = length(other_parameter_values) + inv_sort = invperm(sort_idx) + + run_algorithm = algorithm + if run_algorithm == :pruned_third_order && !(!(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[])) + run_algorithm = :pruned_second_order + end + + solve!(๐“‚, + algorithm = run_algorithm, + steady_state_function = steady_state_function, + opts = opts) + + nVars = length(๐“‚.constants.post_model_macro.var) + + nsss_only = !(non_stochastic_steady_state == Symbol[]) && (standard_deviation == Symbol[]) && (variance == Symbol[]) && (covariance == Symbol[]) && (autocorrelation == Symbol[]) + + nsss_pb = nothing + cov_pb = nothing + som_pb = nothing + somc_pb = nothing + tom_pb = nothing + toma_pb = nothing + + solved = true + SS_and_pars = zeros(T, 0) + SS = zeros(T, 0) + state_ฮผ = zeros(T, 0) + + covar_dcmp = zeros(T, 0, 0) + sol = zeros(T, 0, 0) + + ฮฃแถปโ‚‚ = zeros(T, 0, 0) + ฮ”ฮผหขโ‚‚ = zeros(T, 0) + autocorr_tmp = zeros(T, 0, 0) + sฬ‚_to_sฬ‚โ‚‚ = zeros(T, 0, 0) + sฬ‚_to_yโ‚‚ = zeros(T, 0, 0) + + autocorr = zeros(T, 0, 0) + first_order_A = zeros(T, 0, 0) + first_order_P = zeros(T, 0, 0) + first_order_R_seq = Matrix{T}[] + first_order_d = zeros(T, 0) + first_order_mask = BitVector() + + second_order_P_seq = Matrix{T}[] + second_order_M_seq = Matrix{T}[] + second_order_d = zeros(T, 0) + second_order_mask = BitVector() + + st_dev = zeros(T, 0) + varrs = zeros(T, 0) + diag_covar = zeros(T, 0) + diag_gate = falses(0) + + covar_dcmp_sp = zeros(T, 0, 0) + covar_group_pairs = NTuple{4,Int}[] + + if nsss_only + prev_ฮ”nsss = Ref{Any}(nothing) + + nsss_out, nsss_pb_local = rrule(get_NSSS_and_parameters, ๐“‚, all_parameters; opts = opts) + nsss_pb = nsss_pb_local + + SS_and_pars = nsss_out[1] + solution_error = nsss_out[2][1] + SS = SS_and_pars[1:end - length(๐“‚.equations.calibration)] + + ret = Dict{Symbol,AbstractArray{T}}() + ret[:non_stochastic_steady_state] = solution_error < opts.tol.nsss.acceptance_tol ? SS[SS_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(SS_var_idx) ? 0 : length(SS_var_idx)) + + function nsss_only_pullback(ฮ”ret) + ฮ”nsss = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :non_stochastic_steady_state), prev_ฮ”nsss) + if ฮ”nsss isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(T, length(parameter_values)) + end + + โˆ‚SS = zeros(T, length(SS)) + โˆ‚SS[SS_var_idx] .+= unthunk(ฮ”nsss) + + โˆ‚SS_and_pars = zeros(T, length(SS_and_pars)) + โˆ‚SS_and_pars[1:length(SS)] .+= โˆ‚SS + + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + โˆ‚all_parameters = nsss_grads[3] isa AbstractZero ? zeros(T, length(all_parameters)) : nsss_grads[3] + + โˆ‚concat = โˆ‚all_parameters[inv_sort] + โˆ‚parameter_values = โˆ‚concat[(n_other + 1):end] + + return NoTangent(), NoTangent(), โˆ‚parameter_values + end + + return ret, nsss_only_pullback + end + + if run_algorithm == :pruned_third_order + if !(autocorrelation == Symbol[]) + second_mom_third_order = union(autocorr_var_idx, std_var_idx, var_var_idx) + toma_out, toma_pb_local = rrule(calculate_third_order_moments_with_autocorrelation, + all_parameters, + ๐“‚.constants.post_model_macro.var[second_mom_third_order], + ๐“‚; + covariance = ๐“‚.constants.post_model_macro.var[covar_var_idx], + opts = opts, + autocorrelation_periods = autocorrelation_periods) + toma_pb = toma_pb_local + + covar_dcmp = toma_out[1] + state_ฮผ = toma_out[2] + autocorr = toma_out[3] + SS_and_pars = toma_out[4] + solved = toma_out[5] + elseif !(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) + tom_out, tom_pb_local = rrule(calculate_third_order_moments, + all_parameters, + ๐“‚.constants.post_model_macro.var[union(std_var_idx, var_var_idx)], + ๐“‚; + covariance = ๐“‚.constants.post_model_macro.var[covar_var_idx], + opts = opts) + tom_pb = tom_pb_local + + covar_dcmp = tom_out[1] + state_ฮผ = tom_out[2] + SS_and_pars = tom_out[3] + solved = tom_out[4] + end + elseif run_algorithm == :pruned_second_order + if !(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[]) + somc_out, somc_pb_local = rrule(calculate_second_order_moments_with_covariance, all_parameters, ๐“‚; opts = opts) + somc_pb = somc_pb_local + + covar_dcmp = somc_out[1] + ฮฃแถปโ‚‚ = somc_out[2] + state_ฮผ = somc_out[3] + ฮ”ฮผหขโ‚‚ = somc_out[4] + autocorr_tmp = somc_out[5] + sฬ‚_to_sฬ‚โ‚‚ = somc_out[6] + sฬ‚_to_yโ‚‚ = somc_out[7] + SS_and_pars = somc_out[10] + solved = somc_out[15] + else + som_out, som_pb_local = rrule(calculate_second_order_moments, all_parameters, ๐“‚; opts = opts) + som_pb = som_pb_local + + state_ฮผ = som_out[1] + ฮ”ฮผหขโ‚‚ = som_out[2] + SS_and_pars = som_out[5] + solved = som_out[10] + end + else + cov_out, cov_pb_local = rrule(calculate_covariance, all_parameters, ๐“‚; opts = opts) + cov_pb = cov_pb_local + + covar_dcmp = cov_out[1] + sol = cov_out[2] + SS_and_pars = cov_out[4] + solved = cov_out[5] + end + + SS = SS_and_pars[1:end - length(๐“‚.equations.calibration)] + + if !(variance == Symbol[]) || !(standard_deviation == Symbol[]) + diag_covar = convert(Vector{T}, โ„’.diag(covar_dcmp)) + diag_max = max.(diag_covar, eps(Float64)) + diag_gate = diag_covar .> eps(Float64) + if !(variance == Symbol[]) + varrs = convert(Vector{T}, diag_max) + end + if !(standard_deviation == Symbol[]) + st_dev = sqrt.(abs.(convert(Vector{T}, diag_max))) + end + end + + if !(autocorrelation == Symbol[]) + if run_algorithm == :pruned_second_order + P_i = Matrix{T}(โ„’.I(size(sฬ‚_to_sฬ‚โ‚‚, 1))) + autocorr = zeros(T, size(covar_dcmp, 1), length(autocorrelation_periods)) + second_order_P_seq = [zeros(T, 0, 0) for _ in 1:maximum(autocorrelation_periods)] + second_order_M_seq = [zeros(T, 0, 0) for _ in 1:maximum(autocorrelation_periods)] + second_order_d = max.(convert(Vector{T}, โ„’.diag(covar_dcmp)), eps(Float64)) + + for i in autocorrelation_periods + second_order_P_seq[i] = copy(P_i) + M_i = sฬ‚_to_yโ‚‚ * P_i * autocorr_tmp + second_order_M_seq[i] = M_i + autocorr[:, i] .= โ„’.diag(M_i) ./ second_order_d + P_i = P_i * sฬ‚_to_sฬ‚โ‚‚ + end + + second_order_mask = โ„’.diag(covar_dcmp) .< opts.tol.second_order.lyapunov.acceptance_tol + autocorr[second_order_mask, :] .= 0 + elseif !(run_algorithm == :pruned_third_order) + first_order_P = โ„’.diagm(ones(T, ๐“‚.constants.post_model_macro.nVars))[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx, :] + first_order_A = @views sol[:, 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * first_order_P + first_order_d = max.(convert(Vector{T}, โ„’.diag(covar_dcmp)), eps(Float64)) + d_inv = 1 ./ first_order_d + + autocorr = zeros(T, size(covar_dcmp, 1), length(autocorrelation_periods)) + first_order_R_seq = [zeros(T, 0, 0) for _ in 1:maximum(autocorrelation_periods)] + + R = Matrix(covar_dcmp) + for i in 1:maximum(autocorrelation_periods) + R = first_order_A * R + first_order_R_seq[i] = copy(R) + end + + for i in autocorrelation_periods + autocorr[:, i] .= โ„’.diag(first_order_R_seq[i]) .* d_inv + end + + first_order_mask = โ„’.diag(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol + autocorr[first_order_mask, :] .= 0 + end + end + + if !(covariance == Symbol[]) + covar_dcmp_sp = โ„’.triu(covar_dcmp) + + if !isnothing(covar_groups) + for group in covar_groups + for i in group + i_pos = findfirst(==(i), covar_var_idx) + isnothing(i_pos) && continue + for j in group + j_pos = findfirst(==(j), covar_var_idx) + isnothing(j_pos) && continue + push!(covar_group_pairs, (i_pos, j_pos, i, j)) + end + end + end + end + end + + ret = Dict{Symbol,AbstractArray{T}}() + + if !(non_stochastic_steady_state == Symbol[]) + ret[:non_stochastic_steady_state] = solved ? SS[SS_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(SS_var_idx) ? 0 : length(SS_var_idx)) + end + if !(mean == Symbol[]) + if run_algorithm โˆ‰ [:pruned_second_order,:pruned_third_order] + ret[:mean] = solved ? SS[mean_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(mean_var_idx) ? 0 : length(mean_var_idx)) + else + ret[:mean] = solved ? state_ฮผ[mean_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(mean_var_idx) ? 0 : length(mean_var_idx)) + end + end + if !(standard_deviation == Symbol[]) + ret[:standard_deviation] = solved ? st_dev[std_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(std_var_idx) ? 0 : length(std_var_idx)) + end + if !(variance == Symbol[]) + ret[:variance] = solved ? varrs[var_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(var_var_idx) ? 0 : length(var_var_idx)) + end + if !(covariance == Symbol[]) + if !isnothing(covar_groups) + if solved + covar_result = zeros(T, length(covar_var_idx), length(covar_var_idx)) + for (i_pos, j_pos, i, j) in covar_group_pairs + covar_result[i_pos, j_pos] = covar_dcmp_sp[i, j] + end + ret[:covariance] = covar_result + else + ret[:covariance] = fill(Inf * sum(abs2,parameter_values), length(covar_var_idx), length(covar_var_idx)) + end + else + ret[:covariance] = solved ? covar_dcmp_sp[covar_var_idx, covar_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(covar_var_idx) ? 0 : length(covar_var_idx), isnothing(covar_var_idx) ? 0 : length(covar_var_idx)) + end + end + if !(autocorrelation == Symbol[]) + ret[:autocorrelation] = solved ? autocorr[autocorr_var_idx, :] : fill(Inf * sum(abs2,parameter_values), isnothing(autocorr_var_idx) ? 0 : length(autocorr_var_idx), isnothing(autocorrelation_periods) ? 0 : length(autocorrelation_periods)) + end + + prev_ฮ”nsss = Ref{Any}(nothing) + prev_ฮ”mean = Ref{Any}(nothing) + prev_ฮ”std = Ref{Any}(nothing) + prev_ฮ”var = Ref{Any}(nothing) + prev_ฮ”cov = Ref{Any}(nothing) + prev_ฮ”autocorr = Ref{Any}(nothing) + + function get_statistics_pullback(ฮ”ret) + if !solved + return NoTangent(), NoTangent(), zeros(T, length(parameter_values)) + end + + ฮ”nsss = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :non_stochastic_steady_state), prev_ฮ”nsss) + ฮ”mean = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :mean), prev_ฮ”mean) + ฮ”std = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :standard_deviation), prev_ฮ”std) + ฮ”var = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :variance), prev_ฮ”var) + ฮ”cov = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :covariance), prev_ฮ”cov) + ฮ”autocorr = _incremental_cotangent!(_get_statistics_cotangent(ฮ”ret, :autocorrelation), prev_ฮ”autocorr) + + โˆ‚SS_and_pars = zeros(T, length(SS_and_pars)) + โˆ‚state_ฮผ = length(state_ฮผ) == 0 ? zeros(T, 0) : zeros(T, length(state_ฮผ)) + โˆ‚covar_dcmp = size(covar_dcmp, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(covar_dcmp)) + โˆ‚sol = size(sol, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(sol)) + โˆ‚autocorr_tmp = size(autocorr_tmp, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(autocorr_tmp)) + โˆ‚ล_to_sฬ‚โ‚‚ = size(sฬ‚_to_sฬ‚โ‚‚, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(sฬ‚_to_sฬ‚โ‚‚)) + โˆ‚ล_to_yโ‚‚ = size(sฬ‚_to_yโ‚‚, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(sฬ‚_to_yโ‚‚)) + + if !(ฮ”nsss isa Union{NoTangent, AbstractZero}) + โˆ‚SS_and_pars[SS_var_idx] .+= ฮ”nsss + end + + if !(ฮ”mean isa Union{NoTangent, AbstractZero}) + if run_algorithm โˆ‰ [:pruned_second_order,:pruned_third_order] + โˆ‚SS_and_pars[mean_var_idx] .+= ฮ”mean + else + โˆ‚state_ฮผ[mean_var_idx] .+= ฮ”mean + end + end + + if !(ฮ”var isa Union{NoTangent, AbstractZero}) + โˆ‚var_full = zeros(T, length(diag_covar)) + โˆ‚var_full[var_var_idx] .+= ฮ”var + @inbounds for i in eachindex(diag_covar) + if diag_gate[i] + โˆ‚covar_dcmp[i, i] += โˆ‚var_full[i] + end + end + end + + if !(ฮ”std isa Union{NoTangent, AbstractZero}) + โˆ‚std_full = zeros(T, length(diag_covar)) + โˆ‚std_full[std_var_idx] .+= ฮ”std + @inbounds for i in eachindex(diag_covar) + if diag_gate[i] + โˆ‚covar_dcmp[i, i] += โˆ‚std_full[i] / (2 * st_dev[i]) + end + end + end + + if !(ฮ”cov isa Union{NoTangent, AbstractZero}) + โˆ‚covar_dcmp_sp = zeros(T, size(covar_dcmp)) + + if !isnothing(covar_groups) + for (i_pos, j_pos, i, j) in covar_group_pairs + โˆ‚covar_dcmp_sp[i, j] += ฮ”cov[i_pos, j_pos] + end + else + โˆ‚covar_dcmp_sp[covar_var_idx, covar_var_idx] .+= ฮ”cov + end + + โˆ‚covar_dcmp .+= โ„’.triu(โˆ‚covar_dcmp_sp) + end + + if !(ฮ”autocorr isa Union{NoTangent, AbstractZero}) && !(autocorrelation == Symbol[]) + if run_algorithm == :pruned_second_order + โˆ‚autocorr_full = zeros(T, size(covar_dcmp, 1), length(autocorrelation_periods)) + โˆ‚autocorr_full[autocorr_var_idx, :] .= ฮ”autocorr + โˆ‚autocorr_full[second_order_mask, :] .= 0 + + โˆ‚d = zeros(T, length(second_order_d)) + โˆ‚P = [zeros(T, size(second_order_P_seq[i])) for i in 1:length(second_order_P_seq)] + + for i in reverse(collect(autocorrelation_periods)) + g = view(โˆ‚autocorr_full, :, i) + M_i = second_order_M_seq[i] + P_i = second_order_P_seq[i] + + โˆ‚M_i = zeros(T, size(M_i)) + @inbounds for j in 1:size(M_i, 1) + โˆ‚M_i[j, j] += g[j] / second_order_d[j] + โˆ‚d[j] -= g[j] * M_i[j, j] / (second_order_d[j]^2) + end + + P_aut = P_i * autocorr_tmp + โˆ‚ล_to_yโ‚‚ .+= โˆ‚M_i * P_aut' + + โˆ‚Paut = sฬ‚_to_yโ‚‚' * โˆ‚M_i + โˆ‚P[i] .+= โˆ‚Paut * autocorr_tmp' + โˆ‚autocorr_tmp .+= P_i' * โˆ‚Paut + end + + if length(second_order_P_seq) >= 2 + for i in reverse(1:(length(second_order_P_seq) - 1)) + โˆ‚ล_to_sฬ‚โ‚‚ .+= second_order_P_seq[i]' * โˆ‚P[i + 1] + โˆ‚P[i] .+= โˆ‚P[i + 1] * sฬ‚_to_sฬ‚โ‚‚' + end + end + + diag_raw = convert(Vector{T}, โ„’.diag(covar_dcmp)) + @inbounds for i in eachindex(โˆ‚d) + if diag_raw[i] > eps(Float64) + โˆ‚covar_dcmp[i, i] += โˆ‚d[i] + end + end + + โˆ‚state_ฮผ .+= zero(โˆ‚state_ฮผ) + elseif run_algorithm != :pruned_third_order + โˆ‚autocorr_full = zeros(T, size(covar_dcmp, 1), length(autocorrelation_periods)) + โˆ‚autocorr_full[autocorr_var_idx, :] .= ฮ”autocorr + โˆ‚autocorr_full[first_order_mask, :] .= 0 + + d_inv = 1 ./ first_order_d + โˆ‚d = zeros(T, length(first_order_d)) + max_p = maximum(autocorrelation_periods) + โˆ‚R = [zeros(T, size(covar_dcmp)) for _ in 1:max_p] + โˆ‚A = zeros(T, size(first_order_A)) + + for i in reverse(collect(autocorrelation_periods)) + g = view(โˆ‚autocorr_full, :, i) + Ri = first_order_R_seq[i] + @inbounds for j in 1:length(g) + โˆ‚R[i][j, j] += g[j] * d_inv[j] + โˆ‚d[j] -= g[j] * Ri[j, j] / (first_order_d[j]^2) + end + end + + for i in reverse(1:max_p) + if i < max_p + โˆ‚R[i] .+= first_order_A' * โˆ‚R[i + 1] + end + R_prev = (i == 1) ? Matrix(covar_dcmp) : first_order_R_seq[i - 1] + โˆ‚A .+= โˆ‚R[i] * R_prev' + end + + if max_p >= 1 + โˆ‚covar_dcmp .+= first_order_A' * โˆ‚R[1] + end + + diag_raw = convert(Vector{T}, โ„’.diag(covar_dcmp)) + @inbounds for i in eachindex(โˆ‚d) + if diag_raw[i] > eps(Float64) + โˆ‚covar_dcmp[i, i] += โˆ‚d[i] + end + end + + โˆ‚sol[:, 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] .+= โˆ‚A * first_order_P' + end + end + + โˆ‚all_parameters = zeros(T, length(all_parameters)) + + if nsss_only + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + โˆ‚all_parameters .+= (nsss_grads[3] isa AbstractZero ? zeros(T, length(all_parameters)) : nsss_grads[3]) + elseif run_algorithm == :first_order + cov_grads = cov_pb((โˆ‚covar_dcmp, โˆ‚sol, NoTangent(), โˆ‚SS_and_pars, NoTangent())) + โˆ‚all_parameters .+= (cov_grads[2] isa AbstractZero ? zeros(T, length(all_parameters)) : cov_grads[2]) + elseif run_algorithm == :pruned_second_order + if som_pb !== nothing + som_grads = som_pb((โˆ‚state_ฮผ, NoTangent(), NoTangent(), NoTangent(), โˆ‚SS_and_pars, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent())) + โˆ‚all_parameters .+= (som_grads[2] isa AbstractZero ? zeros(T, length(all_parameters)) : som_grads[2]) + else + somc_grads = somc_pb((โˆ‚covar_dcmp, + NoTangent(), + โˆ‚state_ฮผ, + NoTangent(), + run_algorithm == :pruned_second_order && !(autocorrelation == Symbol[]) ? โˆ‚autocorr_tmp : NoTangent(), + run_algorithm == :pruned_second_order && !(autocorrelation == Symbol[]) ? โˆ‚ล_to_sฬ‚โ‚‚ : NoTangent(), + run_algorithm == :pruned_second_order && !(autocorrelation == Symbol[]) ? โˆ‚ล_to_yโ‚‚ : NoTangent(), + NoTangent(), + NoTangent(), + โˆ‚SS_and_pars, + NoTangent(), + NoTangent(), + NoTangent(), + NoTangent(), + NoTangent())) + โˆ‚all_parameters .+= (somc_grads[2] isa AbstractZero ? zeros(T, length(all_parameters)) : somc_grads[2]) + end + elseif run_algorithm == :pruned_third_order + if toma_pb !== nothing + โˆ‚autocorr_full = zeros(T, size(autocorr)) + if !(ฮ”autocorr isa Union{NoTangent, AbstractZero}) + โˆ‚autocorr_full[autocorr_var_idx, :] .= ฮ”autocorr + end + toma_grads = toma_pb((โˆ‚covar_dcmp, โˆ‚state_ฮผ, โˆ‚autocorr_full, โˆ‚SS_and_pars, NoTangent())) + โˆ‚all_parameters .+= (toma_grads[2] isa AbstractZero ? zeros(T, length(all_parameters)) : toma_grads[2]) + elseif tom_pb !== nothing + tom_grads = tom_pb((โˆ‚covar_dcmp, โˆ‚state_ฮผ, โˆ‚SS_and_pars, NoTangent())) + โˆ‚all_parameters .+= (tom_grads[2] isa AbstractZero ? zeros(T, length(all_parameters)) : tom_grads[2]) + end + end + + โˆ‚concat = โˆ‚all_parameters[inv_sort] + โˆ‚parameter_values = โˆ‚concat[(n_other + 1):end] + + return NoTangent(), NoTangent(), โˆ‚parameter_values + end + + return ret, get_statistics_pullback +end + + +# โ”€โ”€ get_solution rrule โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# Custom rrule for get_solution(๐“‚, parameters; ...) that chains existing +# sub-rrules without using AD inside the pullback. +# Supports first_order, second_order/pruned_second_order, +# and third_order/pruned_third_order algorithms. + +function rrule(::typeof(get_solution), + ๐“‚::โ„ณ, + parameters::Vector{S}; + steady_state_function::SteadyStateFunctionType = missing, + algorithm::Symbol = DEFAULT_ALGORITHM, + verbose::Bool = DEFAULT_VERBOSE, + tol::Tolerances = Tolerances(), + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(๐“‚)) where S <: Real + + opts = merge_calculation_options(tol = tol, verbose = verbose, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + sylvester_algorithmยฒ = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], + sylvester_algorithmยณ = (isa(sylvester_algorithm, Symbol) || length(sylvester_algorithm) < 2) ? :bicgstab : sylvester_algorithm[2]) + + estimation = true + + constants_obj = initialise_constants!(๐“‚) + + solve!(๐“‚, + opts = opts, + steady_state_function = steady_state_function, + algorithm = algorithm) + + nVar = length(๐“‚.constants.post_model_macro.var) + + zero_pullback(_) = (NoTangent(), NoTangent(), zeros(S, length(parameters))) + + # โ”€โ”€ Check parameter bounds โ”€โ”€ + if length(๐“‚.constants.post_parameters_macro.bounds) > 0 + for (k, v) in ๐“‚.constants.post_parameters_macro.bounds + if k โˆˆ ๐“‚.constants.post_complete_parameters.parameters + idx = indexin([k], ๐“‚.constants.post_complete_parameters.parameters)[1] + if min(max(parameters[idx], v[1]), v[2]) != parameters[idx] + return -Inf, zero_pullback + end + end + end + end + + # โ”€โ”€ Step 1: NSSS โ”€โ”€ + nsss_out, nsss_pb = rrule(get_NSSS_and_parameters, + ๐“‚, + parameters; + opts = opts, + estimation = estimation) + + SS_and_pars = nsss_out[1] + solution_error = nsss_out[2][1] + + if solution_error > tol.nsss.acceptance_tol || isnan(solution_error) + if algorithm in [:second_order, :pruned_second_order] + result = (SS_and_pars[1:nVar], zeros(nVar, 2), spzeros(nVar, 2), false) + elseif algorithm in [:third_order, :pruned_third_order] + result = (SS_and_pars[1:nVar], zeros(nVar, 2), spzeros(nVar, 2), spzeros(nVar, 2), false) + else + result = (SS_and_pars[1:nVar], zeros(nVar, 2), false) + end + return result, zero_pullback + end + + # โ”€โ”€ Step 2: Jacobian โ”€โ”€ + โˆ‡โ‚, jac_pb = rrule(calculate_jacobian, + parameters, + SS_and_pars, + ๐“‚.caches, + ๐“‚.functions.jacobian, + ๐“‚.workspaces) + + # โ”€โ”€ Step 3: First-order solution โ”€โ”€ + first_out, first_pb = rrule(calculate_first_order_solution, + โˆ‡โ‚, + constants_obj, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameters) + + ๐’โ‚ = first_out[1] + solved = first_out[3] + + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + + if !solved + if algorithm in [:second_order, :pruned_second_order] + result = (SS_and_pars[1:nVar], ๐’โ‚, spzeros(nVar, 2), false) + elseif algorithm in [:third_order, :pruned_third_order] + result = (SS_and_pars[1:nVar], ๐’โ‚, spzeros(nVar, 2), spzeros(nVar, 2), false) + else + result = (SS_and_pars[1:nVar], ๐’โ‚, false) + end + return result, zero_pullback + end + + # โ”€โ”€ Branch by algorithm โ”€โ”€ + if algorithm in [:second_order, :pruned_second_order] + # โ”€โ”€ Step 4: Hessian โ”€โ”€ + โˆ‡โ‚‚, hess_pb = rrule(calculate_hessian, + parameters, + SS_and_pars, + ๐“‚.caches, + ๐“‚.functions.hessian, + ๐“‚.workspaces) + + # โ”€โ”€ Step 5: Second-order solution โ”€โ”€ + second_out, second_pb = rrule(calculate_second_order_solution, + โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.second_order_solution, + opts = opts, + parameter_values = parameters) + + ๐’โ‚‚_raw = second_out[1] + solved2 = second_out[2] + + update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) + + # Return compressed: (NSSS, ๐’โ‚, ๐’โ‚‚, solved) + result = (SS_and_pars[1:nVar], ๐’โ‚, ๐’โ‚‚_raw, true) + + pullback_2nd = function (โˆ‚result_bar) + ฮ” = unthunk(โˆ‚result_bar) + + if ฮ” isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, length(parameters)) + end + + โˆ‚NSSS = ฮ”[1] + โˆ‚๐’โ‚_ext = ฮ”[2] + โˆ‚๐’โ‚‚_ext = ฮ”[3] + # ฮ”[4] is โˆ‚solved โ€” not differentiable + + # โ”€โ”€ Accumulate โˆ‚SS_and_pars (zero-pad to full length) โ”€โ”€ + โˆ‚SS_and_pars = zeros(S, length(SS_and_pars)) + if !(โˆ‚NSSS isa Union{NoTangent, AbstractZero}) + โˆ‚SS_and_pars[1:nVar] .+= โˆ‚NSSS + end + + โˆ‚parameters = zeros(S, length(parameters)) + + # โ”€โ”€ ๐’โ‚‚ is already in compressed space โ€” no ๐”โ‚‚ adjoint needed โ”€โ”€ + โˆ‚๐’โ‚‚_raw = if โˆ‚๐’โ‚‚_ext isa Union{NoTangent, AbstractZero} + zeros(S, size(๐’โ‚‚_raw)) + else + Matrix{S}(โˆ‚๐’โ‚‚_ext) + end + + # โ”€โ”€ second_pb: (โˆ‚๐’โ‚‚_raw, โˆ‚solved2) โ”€โ”€ + second_grads = second_pb((โˆ‚๐’โ‚‚_raw, NoTangent())) + โˆ‚โˆ‡โ‚_from_2nd = second_grads[2] + โˆ‚โˆ‡โ‚‚_from_2nd = second_grads[3] + โˆ‚๐‘บโ‚_from_2nd = second_grads[4] + + # โ”€โ”€ โˆ‡โ‚‚ is internal-only; gradient comes from second-order solution path โ”€โ”€ + โˆ‚โˆ‡โ‚‚_total = โˆ‚โˆ‡โ‚‚_from_2nd + + # โ”€โ”€ hess_pb โ”€โ”€ + hess_grads = hess_pb(โˆ‚โˆ‡โ‚‚_total) + โˆ‚parameters .+= hess_grads[2] + โˆ‚SS_and_pars .+= hess_grads[3] + + # โ”€โ”€ Accumulate โˆ‚๐’โ‚ โ”€โ”€ + โˆ‚๐’โ‚_total = if โˆ‚๐’โ‚_ext isa Union{NoTangent, AbstractZero} + โˆ‚๐‘บโ‚_from_2nd + else + โˆ‚๐’โ‚_ext + โˆ‚๐‘บโ‚_from_2nd + end + + # โ”€โ”€ first_pb โ”€โ”€ + first_grads = first_pb((โˆ‚๐’โ‚_total, NoTangent(), NoTangent())) + + โˆ‚โˆ‡โ‚_total = โˆ‚โˆ‡โ‚_from_2nd + first_grads[2] + + # โ”€โ”€ jac_pb โ”€โ”€ + jac_grads = jac_pb(โˆ‚โˆ‡โ‚_total) + โˆ‚parameters .+= jac_grads[2] + โˆ‚SS_and_pars .+= jac_grads[3] + + # โ”€โ”€ nsss_pb โ”€โ”€ + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + โˆ‚parameters .+= nsss_grads[3] + + return NoTangent(), NoTangent(), โˆ‚parameters + end + + return result, pullback_2nd + + elseif algorithm in [:third_order, :pruned_third_order] + # โ”€โ”€ Step 4: Hessian โ”€โ”€ + โˆ‡โ‚‚, hess_pb = rrule(calculate_hessian, + parameters, + SS_and_pars, + ๐“‚.caches, + ๐“‚.functions.hessian, + ๐“‚.workspaces) + + # โ”€โ”€ Step 5: Second-order solution โ”€โ”€ + second_out, second_pb = rrule(calculate_second_order_solution, + โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.second_order_solution, + opts = opts, + parameter_values = parameters) + + ๐’โ‚‚_raw = second_out[1] + solved2 = second_out[2] + + update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) + + # โ”€โ”€ Step 6: Third-order derivatives โ”€โ”€ + โˆ‡โ‚ƒ, third_deriv_pb = rrule(calculate_third_order_derivatives, + parameters, + SS_and_pars, + ๐“‚.caches, + ๐“‚.functions.third_order_derivatives, + ๐“‚.workspaces) + + # โ”€โ”€ Step 7: Third-order solution โ”€โ”€ + # calculate_third_order_solution now receives compressed ๐’โ‚‚ and compressed โˆ‡โ‚‚ + third_out, third_pb = rrule(calculate_third_order_solution, + โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, + ๐’โ‚, ๐’โ‚‚_raw, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, + parameter_values = parameters) + + ๐’โ‚ƒ_raw = third_out[1] + solved3 = third_out[2] + + update_perturbation_counter!(๐“‚.counters, solved3, estimation = estimation, order = 3) + + # Return compressed: (NSSS, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ, solved) + result = (SS_and_pars[1:nVar], ๐’โ‚, ๐’โ‚‚_raw, ๐’โ‚ƒ_raw, true) + + pullback_3rd = function (โˆ‚result_bar) + ฮ” = unthunk(โˆ‚result_bar) + + if ฮ” isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, length(parameters)) + end + + โˆ‚NSSS = ฮ”[1] + โˆ‚๐’โ‚_ext = ฮ”[2] + โˆ‚๐’โ‚‚_ext = ฮ”[3] + โˆ‚๐’โ‚ƒ_ext = ฮ”[4] + # ฮ”[5] is โˆ‚solved โ€” not differentiable + + # โ”€โ”€ Accumulate โˆ‚SS_and_pars (zero-pad to full length) โ”€โ”€ + โˆ‚SS_and_pars = zeros(S, length(SS_and_pars)) + if !(โˆ‚NSSS isa Union{NoTangent, AbstractZero}) + โˆ‚SS_and_pars[1:nVar] .+= โˆ‚NSSS + end + + โˆ‚parameters = zeros(S, length(parameters)) + + # โ”€โ”€ ๐’โ‚ƒ is already in compressed space โ€” no ๐”โ‚ƒ adjoint needed โ”€โ”€ + โˆ‚๐’โ‚ƒ_raw = if โˆ‚๐’โ‚ƒ_ext isa Union{NoTangent, AbstractZero} + zeros(S, size(๐’โ‚ƒ_raw)) + else + Matrix{S}(โˆ‚๐’โ‚ƒ_ext) + end + + # โ”€โ”€ third_pb: (โˆ‚๐’โ‚ƒ_raw, โˆ‚solved3) โ”€โ”€ + # Returns (NT, โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚ƒ, โˆ‚๐‘บโ‚, โˆ‚๐’โ‚‚, NT, NT, NT) + third_grads = third_pb((โˆ‚๐’โ‚ƒ_raw, NoTangent())) + โˆ‚โˆ‡โ‚_from_3rd = third_grads[2] + โˆ‚โˆ‡โ‚‚_from_3rd = third_grads[3] + โˆ‚โˆ‡โ‚ƒ_from_3rd = third_grads[4] + โˆ‚๐‘บโ‚_from_3rd = third_grads[5] + โˆ‚๐’โ‚‚_from_3rd = third_grads[6] # w.r.t. compressed ๐’โ‚‚ + + # โ”€โ”€ โˆ‡โ‚ƒ is internal-only; gradient comes from third-order solution path โ”€โ”€ + โˆ‚โˆ‡โ‚ƒ_total = โˆ‚โˆ‡โ‚ƒ_from_3rd + third_deriv_grads = third_deriv_pb(โˆ‚โˆ‡โ‚ƒ_total) + โˆ‚parameters .+= third_deriv_grads[2] + โˆ‚SS_and_pars .+= third_deriv_grads[3] + + # โ”€โ”€ Accumulate โˆ‚๐’โ‚‚ (compressed) from external + third-order โ”€โ”€ + โˆ‚๐’โ‚‚_total = if โˆ‚๐’โ‚‚_ext isa Union{NoTangent, AbstractZero} + โˆ‚๐’โ‚‚_from_3rd isa Union{NoTangent, AbstractZero} ? zeros(S, size(๐’โ‚‚_raw)) : Matrix{S}(โˆ‚๐’โ‚‚_from_3rd) + else + โˆ‚๐’โ‚‚_from_3rd isa Union{NoTangent, AbstractZero} ? Matrix{S}(โˆ‚๐’โ‚‚_ext) : Matrix{S}(โˆ‚๐’โ‚‚_ext) + Matrix{S}(โˆ‚๐’โ‚‚_from_3rd) + end + + # โ”€โ”€ second_pb: (โˆ‚๐’โ‚‚_raw, โˆ‚solved2) โ”€โ”€ + second_grads = second_pb((โˆ‚๐’โ‚‚_total, NoTangent())) + โˆ‚โˆ‡โ‚_from_2nd = second_grads[2] + โˆ‚โˆ‡โ‚‚_from_2nd = second_grads[3] + โˆ‚๐‘บโ‚_from_2nd = second_grads[4] + + # โ”€โ”€ hess_pb (accumulate โˆ‚โˆ‡โ‚‚ from 2nd and 3rd order paths) โ”€โ”€ + โˆ‚โˆ‡โ‚‚_total = โˆ‚โˆ‡โ‚‚_from_3rd + โˆ‚โˆ‡โ‚‚_from_2nd + hess_grads = hess_pb(โˆ‚โˆ‡โ‚‚_total) + โˆ‚parameters .+= hess_grads[2] + โˆ‚SS_and_pars .+= hess_grads[3] + + # โ”€โ”€ Accumulate โˆ‚๐’โ‚ from external + 2nd + 3rd order โ”€โ”€ + โˆ‚๐’โ‚_total = if โˆ‚๐’โ‚_ext isa Union{NoTangent, AbstractZero} + โˆ‚๐‘บโ‚_from_2nd + โˆ‚๐‘บโ‚_from_3rd + else + โˆ‚๐’โ‚_ext + โˆ‚๐‘บโ‚_from_2nd + โˆ‚๐‘บโ‚_from_3rd + end + + # โ”€โ”€ first_pb โ”€โ”€ + first_grads = first_pb((โˆ‚๐’โ‚_total, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚_total = โˆ‚โˆ‡โ‚_from_3rd + โˆ‚โˆ‡โ‚_from_2nd + first_grads[2] + + # โ”€โ”€ jac_pb โ”€โ”€ + jac_grads = jac_pb(โˆ‚โˆ‡โ‚_total) + โˆ‚parameters .+= jac_grads[2] + โˆ‚SS_and_pars .+= jac_grads[3] + + # โ”€โ”€ nsss_pb โ”€โ”€ + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + โˆ‚parameters .+= nsss_grads[3] + + return NoTangent(), NoTangent(), โˆ‚parameters + end + + return result, pullback_3rd + + else + # โ”€โ”€ First order โ”€โ”€ + result = (SS_and_pars[1:nVar], ๐’โ‚, true) + + pullback_1st = function (โˆ‚result_bar) + ฮ” = unthunk(โˆ‚result_bar) + + if ฮ” isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, length(parameters)) + end + + โˆ‚NSSS = ฮ”[1] + โˆ‚๐’โ‚_ext = ฮ”[2] + # ฮ”[3] is โˆ‚solved โ€” not differentiable + + # โ”€โ”€ Accumulate โˆ‚SS_and_pars (zero-pad to full length) โ”€โ”€ + โˆ‚SS_and_pars = zeros(S, length(SS_and_pars)) + if !(โˆ‚NSSS isa Union{NoTangent, AbstractZero}) + โˆ‚SS_and_pars[1:nVar] .+= โˆ‚NSSS + end + + # Short-circuit when solution matrix cotangent is absent + if โˆ‚๐’โ‚_ext isa Union{NoTangent, AbstractZero} + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + return NoTangent(), NoTangent(), nsss_grads[3] + end + + # โ”€โ”€ first_pb: (โˆ‚๐’โ‚, โˆ‚qme_sol, โˆ‚solved) โ”€โ”€ + # Returns (NT, โˆ‚โˆ‡โ‚, NT, NT, NT, NT) + first_grads = first_pb((โˆ‚๐’โ‚_ext, NoTangent(), NoTangent())) + โˆ‚โˆ‡โ‚ = first_grads[2] + + # โ”€โ”€ jac_pb โ”€โ”€ + # Returns (NT, โˆ‚parameters, โˆ‚SS_and_pars, NT, NT) + jac_grads = jac_pb(โˆ‚โˆ‡โ‚) + โˆ‚parameters = copy(jac_grads[2]) + โˆ‚SS_and_pars .+= jac_grads[3] + + # โ”€โ”€ nsss_pb โ”€โ”€ + # Returns (NT, NT, โˆ‚parameter_values, NT) + nsss_grads = nsss_pb((โˆ‚SS_and_pars, NoTangent())) + โˆ‚parameters .+= nsss_grads[3] + + return NoTangent(), NoTangent(), โˆ‚parameters + end + + return result, pullback_1st + end +end diff --git a/src/custom_autodiff_rules/zygote.jl b/src/custom_autodiff_rules/zygote.jl deleted file mode 100644 index b4ae9bb0e..000000000 --- a/src/custom_autodiff_rules/zygote.jl +++ /dev/null @@ -1,4115 +0,0 @@ -# Zygote/ChainRulesCore rrule definitions for reverse-mode automatic differentiation -# -# This file centralizes rrule definitions for computing gradients via reverse-mode AD. -# Each rrule specifies how to propagate gradients backward through custom functions. -# -# Strategy for each rrule: -# 1. Compute the forward pass and store necessary intermediate values -# 2. Return the result and a pullback function -# 3. The pullback computes gradients w.r.t. inputs given upstream gradients -# 4. Use implicit differentiation for iterative solvers and matrix equations -# -# Functions covered: -# - Basic operations: mul_reverse_AD!, mat_mult_kron, sparse_preallocated! -# - Steady states: get_NSSS_and_parameters, calculate_second/third_order_stochastic_steady_state -# - Derivatives: calculate_jacobian, calculate_hessian, calculate_third_order_derivatives -# - Solutions: calculate_first/second/third_order_solution -# - Matrix equations: solve_sylvester_equation, solve_lyapunov_equation -# - Filters: calculate_inversion_filter_loglikelihood, run_kalman_iterations, find_shocks - -function rrule(::typeof(mul_reverse_AD!), - C::Matrix{S}, - A::AbstractMatrix{M}, - B::AbstractMatrix{N}) where {S <: Real, M <: Real, N <: Real} - project_A = ProjectTo(A) - project_B = ProjectTo(B) - - function times_pullback(ศณ) - ศฒ = unthunk(ศณ) - dA = @thunk(project_A(ศฒ * B')) - dB = @thunk(project_B(A' * ศฒ)) - return (NoTangent(), NoTangent(), dA, dB) - end - - return โ„’.mul!(C,A,B), times_pullback -end - -function rrule(::typeof(mat_mult_kron), - A::AbstractSparseMatrix{R}, - B::AbstractMatrix{T}, - C::AbstractMatrix{T}, - D::AbstractMatrix{S}) where {R <: Real, T <: Real, S <: Real} - Y = mat_mult_kron(A, B, C, D) - - function mat_mult_kron_pullback(ศฒ) - if ศฒ isa AbstractZero - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - ศฒdense = Matrix(ศฒ) - - n_rowB = size(B, 1) - n_colB = size(B, 2) - n_rowC = size(C, 1) - n_colC = size(C, 2) - - G = promote_type(eltype(B), eltype(C), eltype(D), Float64) - - โˆ‚B = zeros(G, size(B)) - โˆ‚C = zeros(G, size(C)) - โˆ‚D = zeros(G, size(D)) - - A_csc = A isa SparseMatrixCSC ? A : A.A - nnzA = nnz(A_csc) - nz_col = Vector{Int}(undef, nnzA) - row_to_nzinds = Dict{Int, Vector{Int}}() - - for col in 1:size(A_csc, 2) - for k in A_csc.colptr[col]:(A_csc.colptr[col + 1] - 1) - nz_col[k] = col - r = A_csc.rowval[k] - push!(get!(row_to_nzinds, r, Int[]), k) - end - end - - โˆ‚A_nz = zeros(G, nnzA) - Abar_vec = zeros(G, size(A_csc, 2)) - - for (r, ks) in row_to_nzinds - fill!(Abar_vec, zero(G)) - @inbounds for k in ks - Abar_vec[nz_col[k]] = A_csc.nzval[k] - end - - Abar = reshape(Abar_vec, n_rowC, n_rowB) - AbarB = Abar * B - CAbarB = C' * AbarB - vCAbarB = vec(CAbarB) - - g_row = collect(@view ศฒdense[r, :]) - - โˆ‚D .+= vCAbarB * g_row' - - vCAbarBฬ„ = D * g_row - CAbarBฬ„ = reshape(vCAbarBฬ„, n_colC, n_colB) - - โˆ‚C .+= AbarB * CAbarBฬ„' - - AbarBฬ„ = C * CAbarBฬ„ - โˆ‚B .+= Abar' * AbarBฬ„ - - Abarฬ„ = AbarBฬ„ * B' - vecAbarฬ„ = vec(Abarฬ„) - @inbounds for k in ks - โˆ‚A_nz[k] += vecAbarฬ„[nz_col[k]] - end - end - - โˆ‚A_csc = SparseMatrixCSC(size(A_csc, 1), size(A_csc, 2), copy(A_csc.colptr), copy(A_csc.rowval), โˆ‚A_nz) - - return NoTangent(), - ProjectTo(A)(โˆ‚A_csc), - ProjectTo(B)(โˆ‚B), - ProjectTo(C)(โˆ‚C), - ProjectTo(D)(โˆ‚D) - end - - return Y, mat_mult_kron_pullback -end - - - -function rrule(::typeof(sparse_preallocated!), ลœ::Matrix{T}; โ„‚::higher_order_workspace{T,F,H} = Higher_order_workspace()) where {T <: Real, F <: AbstractFloat, H <: Real} - project_ลœ = ProjectTo(ลœ) - - function sparse_preallocated_pullback(ฮฉฬ„) - ฮ”ฮฉ = unthunk(ฮฉฬ„) - ฮ”ลœ = project_ลœ(ฮ”ฮฉ) - return NoTangent(), ฮ”ลœ, NoTangent() - end - - return sparse_preallocated!(ลœ, โ„‚ = โ„‚), sparse_preallocated_pullback -end - -function rrule(::typeof(calculate_second_order_stochastic_steady_state), - ::Val{:newton}, - ๐’โ‚::Matrix{Float64}, - ๐’โ‚‚::AbstractSparseMatrix{Float64}, - x::Vector{Float64}, - ๐“‚::โ„ณ; - # timer::TimerOutput = TimerOutput(), - tol::AbstractFloat = 1e-14) - # @timeit_debug timer "Calculate SSS - forward" begin - # @timeit_debug timer "Setup indices" begin - - # Get cached computational constants - constants = initialise_constants!(๐“‚) - so = constants.second_order - T = constants.post_model_macro - s_in_sโบ = so.s_in_sโบ - s_in_s = so.s_in_s - I_nPast = ๐“‚.workspaces.qme.I_nPast - - kron_sโบ_sโบ = so.kron_sโบ_sโบ - - kron_sโบ_s = so.kron_sโบ_s - - A = ๐’โ‚[T.past_not_future_and_mixed_idx,1:T.nPast_not_future_and_mixed] - B = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_s] - Bฬ‚ = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] - - # end # timeit_debug - - # @timeit_debug timer "Iterations" begin - - max_iters = 100 - # SSS .= ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 - for i in 1:max_iters - โˆ‚x = (A + B * โ„’.kron(vcat(x,1), I_nPast) - I_nPast) - - โˆ‚xฬ‚ = โ„’.lu!(โˆ‚x, check = false) - - if !โ„’.issuccess(โˆ‚xฬ‚) - return x, false - end - - ฮ”x = โˆ‚xฬ‚ \ (A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 - x) - - if i > 5 && isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2, x, rtol = tol) - break - end - - # x += ฮ”x - โ„’.axpy!(-1, ฮ”x, x) - end - - solved = isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2, x, rtol = tol) - - # println(x) - - โˆ‚๐’โ‚ = zero(๐’โ‚) - โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) - - # end # timeit_debug - # end # timeit_debug - - function second_order_stochastic_steady_state_pullback(โˆ‚x) - # @timeit_debug timer "Calculate SSS - pullback" begin - - S = -โˆ‚x[1]' / (A + B * โ„’.kron(vcat(x,1), I_nPast) - I_nPast) - - โˆ‚๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] = S' * x' - - โˆ‚๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] = S' * โ„’.kron(vcat(x,1), vcat(x,1))' / 2 - - # end # timeit_debug - - return NoTangent(), NoTangent(), โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, NoTangent(), NoTangent(), NoTangent() - end - - return (x, solved), second_order_stochastic_steady_state_pullback -end - - -function rrule(::typeof(calculate_third_order_stochastic_steady_state), - ::Val{:newton}, - ๐’โ‚::Matrix{Float64}, - ๐’โ‚‚::AbstractSparseMatrix{Float64}, - ๐’โ‚ƒ::AbstractSparseMatrix{Float64}, - x::Vector{Float64}, - ๐“‚::โ„ณ; - tol::AbstractFloat = 1e-14) - # Get cached computational constants - so = ensure_computational_constants!(๐“‚) - T = ๐“‚.constants.post_model_macro - s_in_sโบ = so.s_in_sโบ - s_in_s = so.s_in_s - I_nPast = ๐“‚.workspaces.qme.I_nPast - - kron_sโบ_sโบ = so.kron_sโบ_sโบ - - kron_sโบ_s = so.kron_sโบ_s - - kron_sโบ_sโบ_sโบ = so.kron_sโบ_sโบ_sโบ - - kron_s_sโบ_sโบ = so.kron_s_sโบ_sโบ - - A = ๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] - B = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_s] - Bฬ‚ = ๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] - C = ๐’โ‚ƒ[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_s_sโบ_sโบ] - ฤˆ = ๐’โ‚ƒ[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ_sโบ] - - max_iters = 100 - # SSS .= ๐’โ‚ * aug_state + ๐’โ‚‚ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โ‚ƒ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 - for i in 1:max_iters - โˆ‚x = (A + B * โ„’.kron(vcat(x,1), I_nPast) + C * โ„’.kron(โ„’.kron(vcat(x,1), vcat(x,1)), I_nPast) / 2 - I_nPast) - - โˆ‚xฬ‚ = โ„’.lu!(โˆ‚x, check = false) - - if !โ„’.issuccess(โˆ‚xฬ‚) - return x, false - end - - ฮ”x = โˆ‚xฬ‚ \ (A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + ฤˆ * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1))) / 6 - x) - - if i > 5 && isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + ฤˆ * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1))) / 6, x, rtol = tol) - break - end - - # x += ฮ”x - โ„’.axpy!(-1, ฮ”x, x) - end - - solved = isapprox(A * x + Bฬ‚ * โ„’.kron(vcat(x,1), vcat(x,1)) / 2 + ฤˆ * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1))) / 6, x, rtol = tol) - - โˆ‚๐’โ‚ = zero(๐’โ‚) - โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) - โˆ‚๐’โ‚ƒ = zero(๐’โ‚ƒ) - - function third_order_stochastic_steady_state_pullback(โˆ‚x) - S = -โˆ‚x[1]' / (A + B * โ„’.kron(vcat(x,1), I_nPast) + C * โ„’.kron(โ„’.kron(vcat(x,1), vcat(x,1)), I_nPast) / 2 - I_nPast) - - โˆ‚๐’โ‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] = S' * x' - - โˆ‚๐’โ‚‚[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ] = S' * โ„’.kron(vcat(x,1), vcat(x,1))' / 2 - - โˆ‚๐’โ‚ƒ[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,kron_sโบ_sโบ_sโบ] = S' * โ„’.kron(vcat(x,1), โ„’.kron(vcat(x,1), vcat(x,1)))' / 6 - - return NoTangent(), NoTangent(), โˆ‚๐’โ‚, โˆ‚๐’โ‚‚, โˆ‚๐’โ‚ƒ, NoTangent(), NoTangent(), NoTangent() - end - - return (x, solved), third_order_stochastic_steady_state_pullback -end - - -function rrule(::typeof(calculate_jacobian), - parameters, - SS_and_pars, - caches_obj::caches, - jacobian_funcs::jacobian_functions) - jacobian = calculate_jacobian(parameters, SS_and_pars, caches_obj, jacobian_funcs) - - function calculate_jacobian_pullback(โˆ‚โˆ‡โ‚) - jacobian_funcs.f_parameters(caches_obj.jacobian_parameters, parameters, SS_and_pars) - jacobian_funcs.f_SS_and_pars(caches_obj.jacobian_SS_and_pars, parameters, SS_and_pars) - - โˆ‚parameters = caches_obj.jacobian_parameters' * vec(โˆ‚โˆ‡โ‚) - โˆ‚SS_and_pars = caches_obj.jacobian_SS_and_pars' * vec(โˆ‚โˆ‡โ‚) - return NoTangent(), โˆ‚parameters, โˆ‚SS_and_pars, NoTangent(), NoTangent() - end - - return jacobian, calculate_jacobian_pullback -end - - -function rrule(::typeof(calculate_hessian), - parameters, - SS_and_pars, - caches_obj::caches, - hessian_funcs::hessian_functions) - hessian = calculate_hessian(parameters, SS_and_pars, caches_obj, hessian_funcs) - - function calculate_hessian_pullback(โˆ‚โˆ‡โ‚‚) - hessian_funcs.f_parameters(caches_obj.hessian_parameters, parameters, SS_and_pars) - hessian_funcs.f_SS_and_pars(caches_obj.hessian_SS_and_pars, parameters, SS_and_pars) - - โˆ‚parameters = caches_obj.hessian_parameters' * vec(โˆ‚โˆ‡โ‚‚) - โˆ‚SS_and_pars = caches_obj.hessian_SS_and_pars' * vec(โˆ‚โˆ‡โ‚‚) - - return NoTangent(), โˆ‚parameters, โˆ‚SS_and_pars, NoTangent(), NoTangent() - end - - return hessian, calculate_hessian_pullback -end - - -function rrule(::typeof(calculate_third_order_derivatives), - parameters, - SS_and_pars, - caches_obj::caches, - third_order_derivatives_funcs::third_order_derivatives_functions) - third_order_derivatives = calculate_third_order_derivatives(parameters, SS_and_pars, caches_obj, third_order_derivatives_funcs) - - function calculate_third_order_derivatives_pullback(โˆ‚โˆ‡โ‚ƒ) - third_order_derivatives_funcs.f_parameters(caches_obj.third_order_derivatives_parameters, parameters, SS_and_pars) - third_order_derivatives_funcs.f_SS_and_pars(caches_obj.third_order_derivatives_SS_and_pars, parameters, SS_and_pars) - - โˆ‚parameters = caches_obj.third_order_derivatives_parameters' * vec(โˆ‚โˆ‡โ‚ƒ) - โˆ‚SS_and_pars = caches_obj.third_order_derivatives_SS_and_pars' * vec(โˆ‚โˆ‡โ‚ƒ) - - return NoTangent(), โˆ‚parameters, โˆ‚SS_and_pars, NoTangent(), NoTangent() - end - - return third_order_derivatives, calculate_third_order_derivatives_pullback -end - -function rrule(::typeof(get_NSSS_and_parameters), - ๐“‚::โ„ณ, - parameter_values::Vector{S}; - opts::CalculationOptions = merge_calculation_options(), - cold_start::Bool = false, - estimation::Bool = false) where S <: Real - # timer::TimerOutput = TimerOutput(), - # @timeit_debug timer "Calculate NSSS - forward" begin - ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) - - # Use custom steady state function if available, otherwise use default solver - if ๐“‚.functions.NSSS_custom isa Function - vars_in_ss_equations = ms.vars_in_ss_equations - expected_length = length(vars_in_ss_equations) + length(๐“‚.equations.calibration_parameters) - - SS_and_pars_tmp = evaluate_custom_steady_state_function( - ๐“‚, - parameter_values, - expected_length, - length(๐“‚.constants.post_complete_parameters.parameters), - ) - - residual = zeros(length(๐“‚.equations.steady_state) + length(๐“‚.equations.calibration)) - - ๐“‚.functions.NSSS_check(residual, parameter_values, SS_and_pars_tmp) - - solution_error = โ„’.norm(residual) - - iters = 0 - - # if !isfinite(solution_error) || solution_error > opts.tol.NSSS_acceptance_tol - # throw(ArgumentError("Custom steady state function failed steady state check: residual $solution_error > $(opts.tol.NSSS_acceptance_tol). Parameters: $(parameter_values). Steady state and parameters returned: $(SS_and_pars_tmp).")) - # end - X = @ignore_derivatives ms.custom_ss_expand_matrix - SS_and_pars = X * SS_and_pars_tmp - else - SS_and_pars, (solution_error, iters) = ๐“‚.functions.NSSS_solve(parameter_values, ๐“‚, opts.tol, opts.verbose, cold_start, DEFAULT_SOLVER_PARAMETERS) - end - - # end # timeit_debug - - if solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error) - # Update failed counter - update_ss_counter!(๐“‚.counters, false, estimation = estimation) - return (SS_and_pars, (solution_error, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # Update success counter - update_ss_counter!(๐“‚.counters, true, estimation = estimation) - - # @timeit_debug timer "Calculate NSSS - pullback" begin - - SS_and_pars_names = ms.SS_and_pars_names - SS_and_pars_names_lead_lag = ms.SS_and_pars_names_lead_lag - - # unknowns = union(setdiff(๐“‚.vars_in_ss_equations, ๐“‚.constants.post_model_macro.โž•_vars), ๐“‚.calibration_equations_parameters) - unknowns = Symbol.(vcat(string.(sort(collect(setdiff(reduce(union,get_symbols.(๐“‚.equations.steady_state_aux)),union(๐“‚.constants.post_model_macro.parameters_in_equations,๐“‚.constants.post_model_macro.โž•_vars))))), ๐“‚.equations.calibration_parameters)) - - โˆ‚ = parameter_values - C = SS_and_pars[ms.SS_and_pars_no_exo_idx] # [dyn_ss_idx]) - - if eltype(๐“‚.caches.โˆ‚equations_โˆ‚parameters) != eltype(parameter_values) - if ๐“‚.caches.โˆ‚equations_โˆ‚parameters isa SparseMatrixCSC - jac_buffer = similar(๐“‚.caches.โˆ‚equations_โˆ‚parameters, eltype(parameter_values)) - jac_buffer.nzval .= 0 - else - jac_buffer = zeros(eltype(parameter_values), size(๐“‚.caches.โˆ‚equations_โˆ‚parameters)) - end - else - jac_buffer = ๐“‚.caches.โˆ‚equations_โˆ‚parameters - end - - ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚parameters(jac_buffer, โˆ‚, C) - - โˆ‚SS_equations_โˆ‚parameters = jac_buffer - - - if eltype(๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars) != eltype(SS_and_pars) - if ๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars isa SparseMatrixCSC - jac_buffer = similar(๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars, eltype(SS_and_pars)) - jac_buffer.nzval .= 0 - else - jac_buffer = zeros(eltype(SS_and_pars), size(๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars)) - end - else - jac_buffer = ๐“‚.caches.โˆ‚equations_โˆ‚SS_and_pars - end - - ๐“‚.functions.NSSS_โˆ‚equations_โˆ‚SS_and_pars(jac_buffer, โˆ‚, C) - - โˆ‚SS_equations_โˆ‚SS_and_pars = jac_buffer - - โˆ‚SS_equations_โˆ‚SS_and_pars_lu = RF.lu(โˆ‚SS_equations_โˆ‚SS_and_pars, check = false) - - if !โ„’.issuccess(โˆ‚SS_equations_โˆ‚SS_and_pars_lu) - return (SS_and_pars, (10.0, iters)), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - JVP = -(โˆ‚SS_equations_โˆ‚SS_and_pars_lu \ โˆ‚SS_equations_โˆ‚parameters)#[indexin(SS_and_pars_names, unknowns),:] - - jvp = zeros(length(SS_and_pars_names_lead_lag), length(๐“‚.constants.post_complete_parameters.parameters)) - - for (i,v) in enumerate(SS_and_pars_names) - if v in unknowns - jvp[i,:] = JVP[indexin([v], unknowns),:] - end - end - - # end # timeit_debug - # end # timeit_debug - - # try block-gmres here - function get_non_stochastic_steady_state_pullback(โˆ‚SS_and_pars) - # println(โˆ‚SS_and_pars) - return NoTangent(), NoTangent(), jvp' * โˆ‚SS_and_pars[1], NoTangent() - end - - - return (SS_and_pars, (solution_error, iters)), get_non_stochastic_steady_state_pullback -end - - -function rrule(::typeof(calculate_first_order_solution), - โˆ‡โ‚::Matrix{R}, - constants::constants, - qme_ws::qme_workspace{R,S}, - sylv_ws::sylvester_workspace{R,S}; - opts::CalculationOptions = merge_calculation_options(), - initial_guess::AbstractMatrix{R} = zeros(0,0)) where {R <: AbstractFloat, S <: Real} - # Forward pass to compute the output and intermediate values needed for the backward pass - # @timeit_debug timer "Calculate 1st order solution" begin - # @timeit_debug timer "Preprocessing" begin - - T = constants.post_model_macro - idx_constants = ensure_first_order_constants!(constants) - - dynIndex = idx_constants.dyn_index - reverse_dynamic_order = idx_constants.reverse_dynamic_order - comb = idx_constants.comb - future_not_past_and_mixed_in_comb = idx_constants.future_not_past_and_mixed_in_comb - past_not_future_and_mixed_in_comb = idx_constants.past_not_future_and_mixed_in_comb - Ir = idx_constants.Ir - - โˆ‡โ‚Š = โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] - โˆ‡โ‚€ = โˆ‡โ‚[:,idx_constants.nabla_zero_cols] - โˆ‡โ‚‹ = โˆ‡โ‚[:,idx_constants.nabla_minus_cols] - โˆ‡ฬ‚โ‚‘ = โˆ‡โ‚[:,idx_constants.nabla_e_start:end] - - # end # timeit_debug - # @timeit_debug timer "Invert โˆ‡โ‚€" begin - - Q = โ„’.qr!(โˆ‡โ‚€[:,T.present_only_idx]) - - Aโ‚Š = Q.Q' * โˆ‡โ‚Š - Aโ‚€ = Q.Q' * โˆ‡โ‚€ - Aโ‚‹ = Q.Q' * โˆ‡โ‚‹ - - # end # timeit_debug - # @timeit_debug timer "Sort matrices" begin - - Aฬƒโ‚Š = Aโ‚Š[dynIndex,:] * Ir[future_not_past_and_mixed_in_comb,:] - Aฬƒโ‚€ = Aโ‚€[dynIndex, comb] - Aฬƒโ‚‹ = Aโ‚‹[dynIndex,:] * Ir[past_not_future_and_mixed_in_comb,:] - - # end # timeit_debug - # @timeit_debug timer "Quadratic matrix equation solve" begin - - sol, solved = solve_quadratic_matrix_equation(Aฬƒโ‚Š, Aฬƒโ‚€, Aฬƒโ‚‹, constants, qme_ws; - initial_guess = initial_guess, - quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - tol = opts.tol.qme_tol, - acceptance_tol = opts.tol.qme_acceptance_tol, - verbose = opts.verbose) - - if !solved - return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # end # timeit_debug - # @timeit_debug timer "Postprocessing" begin - # @timeit_debug timer "Setup matrices" begin - - sol_compact = sol[reverse_dynamic_order, past_not_future_and_mixed_in_comb] - - D = sol_compact[end - T.nFuture_not_past_and_mixed + 1:end, :] - - L = sol[indexin(T.past_not_future_and_mixed_idx, T.present_but_not_only_idx), past_not_future_and_mixed_in_comb] - - Aฬ„โ‚€แตค = Aโ‚€[1:T.nPresent_only, T.present_only_idx] - Aโ‚Šแตค = Aโ‚Š[1:T.nPresent_only,:] - Aฬƒโ‚€แตค = Aโ‚€[1:T.nPresent_only, T.present_but_not_only_idx] - Aโ‚‹แตค = Aโ‚‹[1:T.nPresent_only,:] - - # end # timeit_debug - # @timeit_debug timer "Invert Aฬ„โ‚€แตค" begin - - Aฬ„ฬ‚โ‚€แตค = โ„’.lu!(Aฬ„โ‚€แตค, check = false) - - if !โ„’.issuccess(Aฬ„ฬ‚โ‚€แตค) - return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # A = vcat(-(Aฬ„ฬ‚โ‚€แตค \ (Aโ‚Šแตค * D * L + Aฬƒโ‚€แตค * sol[T.dynamic_order,:] + Aโ‚‹แตค)), sol) - if T.nPresent_only > 0 - โ„’.mul!(Aโ‚‹แตค, Aฬƒโ‚€แตค, sol[:,past_not_future_and_mixed_in_comb], 1, 1) - nโ‚šโ‚‹ = Aโ‚Šแตค * D - โ„’.mul!(Aโ‚‹แตค, nโ‚šโ‚‹, L, 1, 1) - โ„’.ldiv!(Aฬ„ฬ‚โ‚€แตค, Aโ‚‹แตค) - โ„’.rmul!(Aโ‚‹แตค, -1) - end - - # end # timeit_debug - # end # timeit_debug - # @timeit_debug timer "Exogenous part solution" begin - - expand_future = idx_constants.expand_future - expand_past = idx_constants.expand_past - - ๐’แต— = vcat(Aโ‚‹แตค, sol_compact)[T.reorder,:] - - ๐’ฬ‚แต— = ๐’แต— * expand_past - - โ„’.mul!(โˆ‡โ‚€, โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] * expand_future, ๐’ฬ‚แต—, 1, 1) - - C = โ„’.lu!(โˆ‡โ‚€, check = false) - - if !โ„’.issuccess(C) - return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - โ„’.ldiv!(C, โˆ‡ฬ‚โ‚‘) - โ„’.rmul!(โˆ‡ฬ‚โ‚‘, -1) - - # end # timeit_debug - # end # timeit_debug - - M = inv(C) - - tmp2 = -M' * (โˆ‡โ‚Š * expand_future)' - - โˆ‡โ‚Š = โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] * expand_future - โˆ‡โ‚‘ = โˆ‡โ‚[:,idx_constants.nabla_e_start:end] - - function first_order_solution_pullback(โˆ‚๐’) - โˆ‚โˆ‡โ‚ = zero(โˆ‡โ‚) - - โˆ‚๐’แต— = โˆ‚๐’[1][:,1:T.nPast_not_future_and_mixed] - โˆ‚๐’แต‰ = โˆ‚๐’[1][:,T.nPast_not_future_and_mixed + 1:end] - - โˆ‚โˆ‡โ‚[:,idx_constants.nabla_e_start:end] .= -M' * โˆ‚๐’แต‰ - - โˆ‚โˆ‡โ‚[:,idx_constants.nabla_zero_cols] .= M' * โˆ‚๐’แต‰ * โˆ‡โ‚‘' * M' - - โˆ‚โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] .= (M' * โˆ‚๐’แต‰ * โˆ‡โ‚‘' * M' * expand_past' * ๐’แต—')[:,T.future_not_past_and_mixed_idx] - - โˆ‚๐’แต— .+= โˆ‡โ‚Š' * M' * โˆ‚๐’แต‰ * โˆ‡โ‚‘' * M' * expand_past' - - tmp1 = M' * โˆ‚๐’แต— * expand_past - - ss, solved = solve_sylvester_equation(tmp2, ๐’ฬ‚แต—', -tmp1, sylv_ws, - sylvester_algorithm = opts.sylvester_algorithmยฒ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, - verbose = opts.verbose) - - if !solved - NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - โˆ‚โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] .+= (ss * ๐’ฬ‚แต—' * ๐’ฬ‚แต—')[:,T.future_not_past_and_mixed_idx] - โˆ‚โˆ‡โ‚[:,idx_constants.nabla_zero_cols] .+= ss * ๐’ฬ‚แต—' - โˆ‚โˆ‡โ‚[:,idx_constants.nabla_minus_cols] .+= ss[:,T.past_not_future_and_mixed_idx] - - return NoTangent(), โˆ‚โˆ‡โ‚, NoTangent(), NoTangent(), NoTangent() - end - - return (hcat(๐’แต—, โˆ‡ฬ‚โ‚‘), sol, solved), first_order_solution_pullback -end - -function rrule(::typeof(calculate_second_order_solution), - โˆ‡โ‚::AbstractMatrix{S}, #first order derivatives - โˆ‡โ‚‚::SparseMatrixCSC{S}, #second order derivatives - ๐‘บโ‚::AbstractMatrix{S},#first order solution - constants::constants, - workspaces::workspaces; - initial_guess::AbstractMatrix{R} = zeros(0,0), - opts::CalculationOptions = merge_calculation_options()) where {S <: Real, R <: Real} - if !(eltype(workspaces.second_order.Sฬ‚) == S) - workspaces.second_order = Higher_order_workspace(T = S) - end - โ„‚ = workspaces.second_order - Mโ‚‚ = constants.second_order - T = constants.post_model_macro - # @timeit_debug timer "Second order solution - forward" begin - # inspired by Levintal - - # Indices and number of variables - iโ‚Š = T.future_not_past_and_mixed_idx; - iโ‚‹ = T.past_not_future_and_mixed_idx; - - nโ‚‹ = T.nPast_not_future_and_mixed - nโ‚Š = T.nFuture_not_past_and_mixed - nโ‚‘ = T.nExo; - n = T.nVars - nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ - - # @timeit_debug timer "Setup matrices" begin - - # 1st order solution - ๐’โ‚ = @views [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]]# |> sparse - # droptol!(๐’โ‚,tol) - - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) โ„’.I(nโ‚‘ + 1)[1,:] zeros(nโ‚‘ + 1, nโ‚‘)] - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0) - - โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] - ๐’โ‚ - โ„’.I(nโ‚‘โ‚‹)[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]] - - ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] - zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)] - - โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.I(n)[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] - - # end # timeit_debug - # @timeit_debug timer "Invert matrix" begin - - โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = โ„’.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) - - if !โ„’.issuccess(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) - if opts.verbose println("Second order solution: inversion failed") end - return (โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) - spinv = choose_matrix_format(spinv) - - # end # timeit_debug - # @timeit_debug timer "Setup second order matrices" begin - # @timeit_debug timer "A" begin - - โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * โ„’.I(n)[iโ‚Š,:] - - A = spinv * โˆ‡โ‚โ‚Š - - # end # timeit_debug - # @timeit_debug timer "C" begin - - # โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = โˆ‡โ‚‚ * (โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›”) * Mโ‚‚.๐‚โ‚‚ - โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, Mโ‚‚.๐‚โ‚‚) + mat_mult_kron(โˆ‡โ‚‚, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, Mโ‚‚.๐›” * Mโ‚‚.๐‚โ‚‚) - - C = spinv * โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน - - # end # timeit_debug - # @timeit_debug timer "B" begin - - # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0) - - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0) - B = mat_mult_kron(Mโ‚‚.๐”โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐‚โ‚‚) + Mโ‚‚.๐”โ‚‚ * Mโ‚‚.๐›” * Mโ‚‚.๐‚โ‚‚ - - # end # timeit_debug - # end # timeit_debug - # @timeit_debug timer "Solve sylvester equation" begin - - ๐’โ‚‚, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, - initial_guess = initial_guess, - sylvester_algorithm = opts.sylvester_algorithmยฒ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, - verbose = opts.verbose) - - # end # timeit_debug - # @timeit_debug timer "Post-process" begin - - if !solved - return (๐’โ‚‚, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # end # timeit_debug - - # spโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t = choose_matrix_format(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹', density_threshold = 1.0) - - # sp๐’โ‚โ‚Šโ•ฑ๐ŸŽt = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ', density_threshold = 1.0) - - ๐›”t = choose_matrix_format(Mโ‚‚.๐›”', density_threshold = 1.0) - - ๐”โ‚‚t = choose_matrix_format(Mโ‚‚.๐”โ‚‚', density_threshold = 1.0) - - ๐‚โ‚‚t = choose_matrix_format(Mโ‚‚.๐‚โ‚‚', density_threshold = 1.0) - - โˆ‡โ‚‚t = choose_matrix_format(โˆ‡โ‚‚', density_threshold = 1.0) - - # end # timeit_debug - - # Ensure pullback workspaces are properly sized - if size(โ„‚.โˆ‚โˆ‡โ‚‚) != size(โˆ‡โ‚‚) - โ„‚.โˆ‚โˆ‡โ‚‚ = zeros(S, size(โˆ‡โ‚‚)) - end - if size(โ„‚.โˆ‚โˆ‡โ‚) != size(โˆ‡โ‚) - โ„‚.โˆ‚โˆ‡โ‚ = zeros(S, size(โˆ‡โ‚)) - end - if size(โ„‚.โˆ‚๐’โ‚) != size(๐’โ‚) - โ„‚.โˆ‚๐’โ‚ = zeros(S, size(๐’โ‚)) - end - if size(โ„‚.โˆ‚spinv) != size(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€) - โ„‚.โˆ‚spinv = zeros(S, size(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€)) - end - if size(โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) != size(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = zeros(S, size(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)) - end - if size(โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ) != size(๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ = zeros(S, size(๐’โ‚โ‚Šโ•ฑ๐ŸŽ)) - end - if size(โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) != size(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) - โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = zeros(S, size(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹)) - end - - function second_order_solution_pullback(โˆ‚๐’โ‚‚_solved) - # @timeit_debug timer "Second order solution - pullback" begin - - # @timeit_debug timer "Preallocate" begin - # Use workspaces and fill with zeros instead of allocating new arrays - โˆ‚โˆ‡โ‚‚ = โ„‚.โˆ‚โˆ‡โ‚‚; fill!(โˆ‚โˆ‡โ‚‚, zero(S)) - โˆ‚โˆ‡โ‚ = โ„‚.โˆ‚โˆ‡โ‚; fill!(โˆ‚โˆ‡โ‚, zero(S)) - โˆ‚๐’โ‚ = โ„‚.โˆ‚๐’โ‚; fill!(โˆ‚๐’โ‚, zero(S)) - โˆ‚spinv = โ„‚.โˆ‚spinv; fill!(โˆ‚spinv, zero(S)) - โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘; fill!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, zero(S)) - โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ = โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ; fill!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, zero(S)) - โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹; fill!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, zero(S)) - - # end # timeit_debug - - โˆ‚๐’โ‚‚ = โˆ‚๐’โ‚‚_solved[1] - - # โˆ‚๐’โ‚‚ *= ๐”โ‚‚t - - # @timeit_debug timer "Sylvester" begin - if โ„’.norm(โˆ‚๐’โ‚‚) < opts.tol.sylvester_tol - return (๐’โ‚‚, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - โˆ‚C, solved = solve_sylvester_equation(A', B', โˆ‚๐’โ‚‚, โ„‚.sylvester_workspace, - sylvester_algorithm = opts.sylvester_algorithmยฒ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, - verbose = opts.verbose) - - if !solved - return (๐’โ‚‚, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # end # timeit_debug - - # @timeit_debug timer "Matmul" begin - - โˆ‚C = choose_matrix_format(โˆ‚C) # Dense - - โˆ‚A = โˆ‚C * B' * ๐’โ‚‚' # Dense - - โˆ‚B = ๐’โ‚‚' * A' * โˆ‚C # Dense - - # B = (Mโ‚‚.๐”โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + Mโ‚‚.๐”โ‚‚ * Mโ‚‚.๐›”) * Mโ‚‚.๐‚โ‚‚ - โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = ๐”โ‚‚t * โˆ‚B * ๐‚โ‚‚t - - # end # timeit_debug - - # @timeit_debug timer "Kron adjoint" begin - - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - - # end # timeit_debug - - # @timeit_debug timer "Matmul2" begin - - # A = spinv * โˆ‡โ‚โ‚Š - โˆ‚โˆ‡โ‚โ‚Š = spinv' * โˆ‚A - โˆ‚spinv += โˆ‚A * โˆ‡โ‚โ‚Š' - - # โˆ‡โ‚โ‚Š = sparse(โˆ‡โ‚[:,1:nโ‚Š] * spdiagm(ones(n))[iโ‚Š,:]) - โˆ‚โˆ‡โ‚[:,1:nโ‚Š] += โˆ‚โˆ‡โ‚โ‚Š * โ„’.I(n)[:,iโ‚Š] - - # C = spinv * โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน - โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚ = spinv' * โˆ‚C * ๐‚โ‚‚t - - โˆ‚spinv += โˆ‚C * โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน' - - # end # timeit_debug - - # @timeit_debug timer "Matmul3" begin - - # โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = โˆ‡โ‚‚ * โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) * Mโ‚‚.๐‚โ‚‚ + โˆ‡โ‚‚ * โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›” * Mโ‚‚.๐‚โ‚‚ - # kronโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = choose_matrix_format(โ„’.kron(spโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, spโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t), density_threshold = 1.0) - - # ๐›”kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐›”t * โ„’.kron(sp๐’โ‚โ‚Šโ•ฑ๐ŸŽt, sp๐’โ‚โ‚Šโ•ฑ๐ŸŽt), density_threshold = 1.0) - - # โ„’.mul!(โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚, ๐›”kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ, 1, 1) - - # โ„’.mul!(โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚, kronโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, 1, 1) - - โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚ = choose_matrix_format(โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚, density_threshold = 1.0) - - โˆ‚โˆ‡โ‚‚ += mat_mult_kron(โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚ * ๐›”t, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ', ๐’โ‚โ‚Šโ•ฑ๐ŸŽ') - - โˆ‚โˆ‡โ‚‚ += mat_mult_kron(โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹', โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹') - - # end # timeit_debug - - # @timeit_debug timer "Matmul4" begin - - โˆ‚kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ = โˆ‡โ‚‚t * โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚ * ๐›”t - - # end # timeit_debug - - # @timeit_debug timer "Kron adjoint 2" begin - - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - - # end # timeit_debug - - โˆ‚kronโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = โˆ‡โ‚‚t * โˆ‚โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน๐‚โ‚‚ - - # @timeit_debug timer "Kron adjoint 3" begin - - fill_kron_adjoint!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โˆ‚kronโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) # filling dense is much faster - - # end # timeit_debug - - # @timeit_debug timer "Matmul5" begin - - # spinv = sparse(inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€)) - โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -spinv' * โˆ‚spinv * spinv' - - # โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.diagm(ones(n))[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] - โˆ‚โˆ‡โ‚[:,1:nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] * ๐’โ‚[iโ‚Š,1:nโ‚‹]' - โˆ‚โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ - - โˆ‚๐’โ‚[iโ‚Š,1:nโ‚‹] -= โˆ‡โ‚[:,1:nโ‚Š]' * โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] - - # ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] - # zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)]; - โˆ‚๐’โ‚[iโ‚Š,:] += โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] - - ###### โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] - # โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = [โ„’.I(size(๐’โ‚,1))[iโ‚Š,:] * ๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ - # ๐’โ‚ - # spdiagm(ones(nโ‚‘โ‚‹))[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]]; - โˆ‚๐’โ‚ += โ„’.I(size(๐’โ‚,1))[:,iโ‚Š] * โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[1:length(iโ‚Š),:] * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' - โˆ‚๐’โ‚ += โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[length(iโ‚Š) .+ (1:size(๐’โ‚,1)),:] - - โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ += ๐’โ‚' * โ„’.I(size(๐’โ‚,1))[:,iโ‚Š] * โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[1:length(iโ‚Š),:] - - # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) spdiagm(ones(nโ‚‘ + 1))[1,:] zeros(nโ‚‘ + 1, nโ‚‘)]; - โˆ‚๐’โ‚[iโ‚‹,:] += โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:length(iโ‚‹), :] - - # ๐’โ‚ = [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]] - โˆ‚๐‘บโ‚ = [โˆ‚๐’โ‚[:,1:nโ‚‹] โˆ‚๐’โ‚[:,nโ‚‹+2:end]] - - # end # timeit_debug - - # end # timeit_debug - - return NoTangent(), โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚๐‘บโ‚, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - - # return (sparse(๐’โ‚‚ * Mโ‚‚.๐”โ‚‚), solved), second_order_solution_pullback - return (๐’โ‚‚, solved), second_order_solution_pullback -end - -function rrule(::typeof(calculate_third_order_solution), - โˆ‡โ‚::AbstractMatrix{S}, #first order derivatives - โˆ‡โ‚‚::SparseMatrixCSC{S}, #second order derivatives - โˆ‡โ‚ƒ::SparseMatrixCSC{S}, #third order derivatives - ๐‘บโ‚::AbstractMatrix{S}, #first order solution - ๐’โ‚‚::SparseMatrixCSC{S}, #second order solution - constants::constants, - workspaces::workspaces; - initial_guess::AbstractMatrix{Float64} = zeros(0,0), - opts::CalculationOptions = merge_calculation_options()) where S <: AbstractFloat - if !(eltype(workspaces.third_order.Sฬ‚) == S) - workspaces.third_order = Higher_order_workspace(T = S) - end - โ„‚ = workspaces.third_order - Mโ‚‚ = constants.second_order - Mโ‚ƒ = constants.third_order - T = constants.post_model_macro - - # @timeit_debug timer "Third order solution - forward" begin - # inspired by Levintal - - # Indices and number of variables - iโ‚Š = T.future_not_past_and_mixed_idx; - iโ‚‹ = T.past_not_future_and_mixed_idx; - - nโ‚‹ = T.nPast_not_future_and_mixed - nโ‚Š = T.nFuture_not_past_and_mixed - nโ‚‘ = T.nExo; - n = T.nVars - nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ - - # @timeit_debug timer "Setup matrices" begin - - # 1st order solution - ๐’โ‚ = @views [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]]# |> sparse - - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) โ„’.I(nโ‚‘ + 1)[1,:] zeros(nโ‚‘ + 1, nโ‚‘)] - - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0, min_length = 10) - - โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] - ๐’โ‚ - โ„’.I(nโ‚‘โ‚‹)[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]] #|> sparse - - ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] - zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)]# |> sparse - ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10) - - โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.I(n)[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] - - # end # timeit_debug - # @timeit_debug timer "Invert matrix" begin - - โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu = โ„’.lu(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, check = false) - - if !โ„’.issuccess(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) - if opts.verbose println("Second order solution: inversion failed") end - return (โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - spinv = inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu) - spinv = choose_matrix_format(spinv) - - # end # timeit_debug - - โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * โ„’.I(n)[iโ‚Š,:] - - A = spinv * โˆ‡โ‚โ‚Š - - # tmpkron = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘,Mโ‚‚.๐›”) - tmpkron = choose_matrix_format(โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘,Mโ‚‚.๐›”), density_threshold = 1.0, tol = opts.tol.droptol) - kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘,๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - - # @timeit_debug timer "Setup B" begin - # @timeit_debug timer "Add tmpkron" begin - - B = tmpkron - - # end # timeit_debug - # @timeit_debug timer "Step 1" begin - - B += Mโ‚ƒ.๐โ‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚แตฃฬƒ - - # end # timeit_debug - # @timeit_debug timer "Step 2" begin - - B += Mโ‚ƒ.๐โ‚‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚‚แตฃฬƒ - - # end # timeit_debug - # @timeit_debug timer "Mult" begin - - B *= Mโ‚ƒ.๐‚โ‚ƒ - B = choose_matrix_format(Mโ‚ƒ.๐”โ‚ƒ * B, tol = opts.tol.droptol, multithreaded = false) - - # end # timeit_debug - # @timeit_debug timer "3rd Kronecker power" begin - - B += compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc1)#, timer = timer) - - # end # timeit_debug - # end # timeit_debug - # @timeit_debug timer "Setup C" begin - # @timeit_debug timer "Initialise smaller matrices" begin - - โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = @views [(๐’โ‚‚ * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:] ; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)])[iโ‚Š,:] - ๐’โ‚‚ - zeros(nโ‚‹ + nโ‚‘, nโ‚‘โ‚‹^2)]; - - โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, density_threshold = 0.0, min_length = 10, tol = opts.tol.droptol) - - ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚Š,:] - zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹^2)]; - - aux = Mโ‚ƒ.๐’๐ * โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ - - # end # timeit_debug - # @timeit_debug timer "โˆ‡โ‚ƒ" begin - - # tmpkron0 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - # tmpkron22 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, tmpkron0 * Mโ‚‚.๐›”) - - if length(โ„‚.tmpkron0) > 0 && eltype(โ„‚.tmpkron0) == S - โ„’.kron!(โ„‚.tmpkron0, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - else - โ„‚.tmpkron0 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - end - - if length(โ„‚.tmpkron22) > 0 && eltype(โ„‚.tmpkron22) == S - โ„’.kron!(โ„‚.tmpkron22, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) - else - โ„‚.tmpkron22 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) - end - - ๐”โˆ‡โ‚ƒ = โˆ‡โ‚ƒ * Mโ‚ƒ.๐”โˆ‡โ‚ƒ - - ๐—โ‚ƒ = ๐”โˆ‡โ‚ƒ * โ„‚.tmpkron22 + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚โ‚—ฬ‚ * โ„‚.tmpkron22 * Mโ‚ƒ.๐โ‚แตฃฬƒ + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚‚โ‚—ฬ‚ * โ„‚.tmpkron22 * Mโ‚ƒ.๐โ‚‚แตฃฬƒ - - # end # timeit_debug - # @timeit_debug timer "โˆ‡โ‚‚ & โˆ‡โ‚โ‚Š" begin - - ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) - - if length(โ„‚.tmpkron1) > 0 && eltype(โ„‚.tmpkron1) == S - โ„’.kron!(โ„‚.tmpkron1, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) - else - โ„‚.tmpkron1 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) - end - - if length(โ„‚.tmpkron2) > 0 && eltype(โ„‚.tmpkron2) == S - โ„’.kron!(โ„‚.tmpkron2, Mโ‚‚.๐›”, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - else - โ„‚.tmpkron2 = โ„’.kron(Mโ‚‚.๐›”, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - end - - โˆ‡โ‚โ‚Š = choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) - - ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚‹,:] ; zeros(size(๐’โ‚)[2] - nโ‚‹, nโ‚‘โ‚‹^2)] - - ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) - - # @timeit_debug timer "Step 1" begin - out2 = โˆ‡โ‚‚ * โ„‚.tmpkron1 * โ„‚.tmpkron2 # this help - - # end # timeit_debug - # @timeit_debug timer "Step 2" begin - - # end # timeit_debug - # @timeit_debug timer "Step 3" begin - - out2 += โˆ‡โ‚‚ * โ„‚.tmpkron1 * Mโ‚ƒ.๐โ‚โ‚— * โ„‚.tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ# |> findnz - - # end # timeit_debug - # @timeit_debug timer "Step 4" begin - - out2 += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc2)# |> findnz - - # out2 += โˆ‡โ‚‚ * โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›”)# |> findnz - ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›” = ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›” - - if length(โ„‚.tmpkron11) > 0 && eltype(โ„‚.tmpkron11) == S - โ„’.kron!(โ„‚.tmpkron11, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›”) - else - โ„‚.tmpkron11 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›”) - end - out2 += โˆ‡โ‚‚ * โ„‚.tmpkron11# |> findnz - - # end # timeit_debug - # @timeit_debug timer "Step 5" begin - - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.droptol) - if length(โ„‚.tmpkron12) > 0 && eltype(โ„‚.tmpkron12) == S - โ„’.kron!(โ„‚.tmpkron12, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) - else - โ„‚.tmpkron12 = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) - end - out2 += โˆ‡โ‚โ‚Š * ๐’โ‚‚ * โ„‚.tmpkron12 - - # end # timeit_debug - # @timeit_debug timer "Mult" begin - - ๐—โ‚ƒ += out2 * Mโ‚ƒ.๐ - - ๐—โ‚ƒ *= Mโ‚ƒ.๐‚โ‚ƒ - - # end # timeit_debug - # end # timeit_debug - # @timeit_debug timer "3rd Kronecker power aux" begin - - # ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚ƒ, collect(aux), collect(โ„’.kron(aux, aux)), Mโ‚ƒ.๐‚โ‚ƒ) # slower than direct compression - ๐—โ‚ƒ += โˆ‡โ‚ƒ * compressed_kronยณ(aux, rowmask = unique(findnz(โˆ‡โ‚ƒ)[2]), tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc3) #, timer = timer) - ๐—โ‚ƒ = choose_matrix_format(๐—โ‚ƒ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) - - # end # timeit_debug - # @timeit_debug timer "Mult 2" begin - - C = spinv * ๐—โ‚ƒ - - # end # timeit_debug - # end # timeit_debug - # @timeit_debug timer "Solve sylvester equation" begin - - ๐’โ‚ƒ, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, - initial_guess = initial_guess, - sylvester_algorithm = opts.sylvester_algorithmยณ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, - verbose = opts.verbose) - - # end # timeit_debug - # # @timeit_debug timer "Refine sylvester equation" begin - - # if !solved - # ๐’โ‚ƒ, solved = solve_sylvester_equation(A, B, C, - # sylvester_algorithm = :doubling, - # initial_guess = initial_guess, - # verbose = verbose, - # # tol = tol, - # timer = timer) - # end - - if !solved - return (๐’โ‚ƒ, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - ๐’โ‚ƒ = choose_matrix_format(๐’โ‚ƒ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) - - # # end # timeit_debug - - # @timeit_debug timer "Preallocate for pullback" begin - - # At = choose_matrix_format(A')# , density_threshold = 1.0) - - # Bt = choose_matrix_format(B')# , density_threshold = 1.0) - - ๐‚โ‚ƒt = choose_matrix_format(Mโ‚ƒ.๐‚โ‚ƒ')# , density_threshold = 1.0) - - ๐”โ‚ƒt = choose_matrix_format(Mโ‚ƒ.๐”โ‚ƒ')# , density_threshold = 1.0) - - ๐t = choose_matrix_format(Mโ‚ƒ.๐')# , density_threshold = 1.0) - - ๐โ‚แตฃt = choose_matrix_format(Mโ‚ƒ.๐โ‚แตฃ')# , density_threshold = 1.0) - - ๐โ‚โ‚—t = choose_matrix_format(Mโ‚ƒ.๐โ‚โ‚—')# , density_threshold = 1.0) - - Mโ‚ƒ๐”โˆ‡โ‚ƒt = choose_matrix_format(Mโ‚ƒ.๐”โˆ‡โ‚ƒ')# , density_threshold = 1.0) - - ๐”โˆ‡โ‚ƒt = choose_matrix_format(๐”โˆ‡โ‚ƒ')# , density_threshold = 1.0) - - Mโ‚ƒ๐โ‚‚โ‚—ฬ‚t = choose_matrix_format(Mโ‚ƒ.๐โ‚‚โ‚—ฬ‚')# , density_threshold = 1.0) - - Mโ‚ƒ๐โ‚‚แตฃฬƒt = choose_matrix_format(Mโ‚ƒ.๐โ‚‚แตฃฬƒ')# , density_threshold = 1.0) - - Mโ‚ƒ๐โ‚แตฃฬƒt = choose_matrix_format(Mโ‚ƒ.๐โ‚แตฃฬƒ')# , density_threshold = 1.0) - - Mโ‚ƒ๐โ‚โ‚—ฬ‚t = choose_matrix_format(Mโ‚ƒ.๐โ‚โ‚—ฬ‚')# , density_threshold = 1.0) - - ๐›”t = choose_matrix_format(Mโ‚‚.๐›”')# , density_threshold = 1.0) - - โˆ‡โ‚‚t = choose_matrix_format(โˆ‡โ‚‚')# , density_threshold = 1.0) - - tmpkron1t = choose_matrix_format(โ„‚.tmpkron1')# , density_threshold = 1.0) - - tmpkron2t = choose_matrix_format(โ„‚.tmpkron2')# , density_threshold = 1.0) - - tmpkron22t = choose_matrix_format(โ„‚.tmpkron22')# , density_threshold = 1.0) - - tmpkron12t = choose_matrix_format(โ„‚.tmpkron12')# , density_threshold = 1.0) - - ๐’โ‚‚t = choose_matrix_format(๐’โ‚‚', density_threshold = 1.0) # this must be sparse otherwise tests fail - - kronaux = โ„’.kron(aux, aux) - - โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t = choose_matrix_format(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹') - - โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ') - - tmpkron10t = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt) - - # end # timeit_debug - # end # timeit_debug - - # Ensure pullback workspaces are properly sized (for dense matrices only) - if size(โ„‚.โˆ‚โˆ‡โ‚_3rd) != size(โˆ‡โ‚) - โ„‚.โˆ‚โˆ‡โ‚_3rd = zeros(S, size(โˆ‡โ‚)) - end - if size(โ„‚.โˆ‚๐’โ‚_3rd) != size(๐’โ‚) - โ„‚.โˆ‚๐’โ‚_3rd = zeros(S, size(๐’โ‚)) - end - if size(โ„‚.โˆ‚spinv_3rd) != size(spinv) - โ„‚.โˆ‚spinv_3rd = zeros(S, size(spinv)) - end - - function third_order_solution_pullback(โˆ‚๐’โ‚ƒ_solved) - # Use workspaces for dense matrices, zero() for sparse - โˆ‚โˆ‡โ‚ = โ„‚.โˆ‚โˆ‡โ‚_3rd; fill!(โˆ‚โˆ‡โ‚, zero(S)) - โˆ‚โˆ‡โ‚‚ = zero(โˆ‡โ‚‚) # sparse - # โˆ‚๐”โˆ‡โ‚ƒ = zero(๐”โˆ‡โ‚ƒ) - โˆ‚โˆ‡โ‚ƒ = zero(โˆ‡โ‚ƒ) # sparse - โˆ‚๐’โ‚ = โ„‚.โˆ‚๐’โ‚_3rd; fill!(โˆ‚๐’โ‚, zero(S)) - โˆ‚๐’โ‚‚ = zero(๐’โ‚‚) # sparse - โˆ‚spinv = โ„‚.โˆ‚spinv_3rd; fill!(โˆ‚spinv, zero(S)) - โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = zero(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) # may be sparse - โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = zero(kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) # may be sparse - โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ = zero(๐’โ‚โ‚Šโ•ฑ๐ŸŽ) # may be sparse - โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = zero(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) # may be sparse - โˆ‚tmpkron = zero(tmpkron) # sparse - โˆ‚tmpkron22 = zero(โ„‚.tmpkron22) # sparse - โˆ‚kronaux = zero(kronaux) # kron product - โˆ‚aux = zero(aux) - โˆ‚tmpkron0 = zero(โ„‚.tmpkron0) # sparse - โˆ‚โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = zero(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ) # may be sparse - โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = zero(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) # may be sparse - โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›” = zero(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›”) # may be sparse - โˆ‚โˆ‡โ‚โ‚Š = zero(โˆ‡โ‚โ‚Š) # may be sparse - โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = zero(๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) # may be sparse - - # @timeit_debug timer "Third order solution - pullback" begin - - # @timeit_debug timer "Solve sylvester equation" begin - - โˆ‚๐’โ‚ƒ = โˆ‚๐’โ‚ƒ_solved[1] - - # โˆ‚๐’โ‚ƒ *= ๐”โ‚ƒt - - โˆ‚C, solved = solve_sylvester_equation(A', B', โˆ‚๐’โ‚ƒ, โ„‚.sylvester_workspace, - sylvester_algorithm = opts.sylvester_algorithmยณ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, - verbose = opts.verbose) - - if !solved - return (๐’โ‚ƒ, solved), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - โˆ‚C = choose_matrix_format(โˆ‚C, density_threshold = 1.0, min_length = 0) - - # end # timeit_debug - # @timeit_debug timer "Step 0" begin - - โˆ‚A = โˆ‚C * B' * ๐’โ‚ƒ' - - # โˆ‚B = ๐’โ‚ƒ' * A' * โˆ‚C - โˆ‚B = choose_matrix_format(๐’โ‚ƒ' * A' * โˆ‚C, density_threshold = 1.0, min_length = 0) - - # end # timeit_debug - # @timeit_debug timer "Step 1" begin - - # C = spinv * ๐—โ‚ƒ - # โˆ‚๐—โ‚ƒ = spinv' * โˆ‚C * Mโ‚ƒ.๐‚โ‚ƒ' - โˆ‚๐—โ‚ƒ = choose_matrix_format(spinv' * โˆ‚C, density_threshold = 1.0, min_length = 0) - - โˆ‚spinv += โˆ‚C * ๐—โ‚ƒ' - - # ๐—โ‚ƒ = โˆ‡โ‚ƒ * compressed_kronยณ(aux, rowmask = unique(findnz(โˆ‡โ‚ƒ)[2])) - # + (๐”โˆ‡โ‚ƒ * tmpkron22 - # + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚โ‚—ฬ‚ * tmpkron22 * Mโ‚ƒ.๐โ‚แตฃฬƒ - # + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚‚โ‚—ฬ‚ * tmpkron22 * Mโ‚ƒ.๐โ‚‚แตฃฬƒ - # + โˆ‡โ‚‚ * (tmpkron10 + tmpkron1 * tmpkron2 + tmpkron1 * Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ + โ„‚.tmpkron11) * Mโ‚ƒ.๐ - # + โˆ‡โ‚โ‚Š * ๐’โ‚‚ * โ„‚.tmpkron12 * Mโ‚ƒ.๐) * Mโ‚ƒ.๐‚โ‚ƒ - - # โˆ‡โ‚โ‚Š * ๐’โ‚‚ * โ„‚.tmpkron12 * Mโ‚ƒ.๐ * Mโ‚ƒ.๐‚โ‚ƒ - โˆ‚โˆ‡โ‚โ‚Š += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * tmpkron12t * ๐’โ‚‚t - โˆ‚๐’โ‚‚ += โˆ‡โ‚โ‚Š' * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * tmpkron12t - โˆ‚tmpkron12 = ๐’โ‚‚t * โˆ‡โ‚โ‚Š' * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t - - # โ„‚.tmpkron12 = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, โˆ‚tmpkron12, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) - - # end # timeit_debug - # @timeit_debug timer "Step 2" begin - - # โˆ‡โ‚‚ * (tmpkron10 + tmpkron1 * tmpkron2 + tmpkron1 * Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ + โ„‚.tmpkron11) * Mโ‚ƒ.๐ * Mโ‚ƒ.๐‚โ‚ƒ - #improve this - # โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * ( - # tmpkron10 - # + tmpkron1 * tmpkron2 - # + tmpkron1 * Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ - # + โ„‚.tmpkron11 - # )' - - โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * tmpkron10t - # โˆ‚โˆ‡โ‚‚ += mat_mult_kron(โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹t, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽt) - # โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * (tmpkron1 * tmpkron2)' - โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * tmpkron2t * tmpkron1t - - # โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * (tmpkron1 * Mโ‚ƒ.๐โ‚โ‚— * tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ)' - โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * Mโ‚ƒ.๐โ‚แตฃ' * tmpkron2t * Mโ‚ƒ.๐โ‚โ‚—' * tmpkron1t - - โˆ‚โˆ‡โ‚‚ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * โ„‚.tmpkron11' - - โˆ‚tmpkron10 = โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t - - # end # timeit_debug - # @timeit_debug timer "Step 3" begin - - # tmpkron10 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ) - fill_kron_adjoint!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โˆ‚โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, โˆ‚tmpkron10, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ) - - โˆ‚tmpkron11 = โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t - - โˆ‚tmpkron1 = โˆ‚tmpkron11 * tmpkron2t + โˆ‚tmpkron11 * ๐โ‚แตฃt * tmpkron2t * ๐โ‚โ‚—t - - โˆ‚tmpkron2 = tmpkron1t * โˆ‚tmpkron11 - - โˆ‚tmpkron2 += ๐โ‚โ‚—t * โˆ‚tmpkron2 * ๐โ‚แตฃt - - # โˆ‚tmpkron1 = โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * tmpkron2t + โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * ๐โ‚แตฃt * tmpkron2t * ๐โ‚โ‚—t - # #improve this - # โˆ‚tmpkron2 = tmpkron1t * โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t + ๐โ‚โ‚—t * tmpkron1t * โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t * ๐โ‚แตฃt - - # โˆ‚tmpkron11 = โˆ‡โ‚‚t * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * ๐t - - # tmpkron1 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚tmpkron1, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) - - # tmpkron2 = โ„’.kron(Mโ‚‚.๐›”, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - fill_kron_adjoint_โˆ‚B!(โˆ‚tmpkron2, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”) - - # tmpkron11 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›”) - fill_kron_adjoint!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›”, โˆ‚tmpkron11, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›”) - - โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ += โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ๐›” * ๐›”t - - # end # timeit_debug - # @timeit_debug timer "Step 4" begin - - # out = (๐”โˆ‡โ‚ƒ * tmpkron22 - # + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚โ‚—ฬ‚ * tmpkron22 * Mโ‚ƒ.๐โ‚แตฃฬƒ - # + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚‚โ‚—ฬ‚ * tmpkron22 * Mโ‚ƒ.๐โ‚‚แตฃฬƒ ) * Mโ‚ƒ.๐‚โ‚ƒ - - โˆ‚โˆ‡โ‚ƒ += โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * tmpkron22t * Mโ‚ƒ๐”โˆ‡โ‚ƒt + โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * Mโ‚ƒ๐โ‚แตฃฬƒt * tmpkron22t * Mโ‚ƒ๐โ‚โ‚—ฬ‚t * Mโ‚ƒ๐”โˆ‡โ‚ƒt + โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * Mโ‚ƒ๐โ‚‚แตฃฬƒt * tmpkron22t * Mโ‚ƒ๐โ‚‚โ‚—ฬ‚t * Mโ‚ƒ๐”โˆ‡โ‚ƒt - - โˆ‚tmpkron22 += ๐”โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt + Mโ‚ƒ๐โ‚โ‚—ฬ‚t * ๐”โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * Mโ‚ƒ๐โ‚แตฃฬƒt + Mโ‚ƒ๐โ‚‚โ‚—ฬ‚t * ๐”โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt * Mโ‚ƒ๐โ‚‚แตฃฬƒt - - # tmpkron22 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›”) - fill_kron_adjoint!(โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โˆ‚tmpkron0, โˆ‚tmpkron22, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) - - โˆ‚kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ = โˆ‚tmpkron0 * ๐›”t - - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ, โˆ‚kron๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - - # -โˆ‡โ‚ƒ * โ„’.kron(โ„’.kron(aux, aux), aux) - # โˆ‚โˆ‡โ‚ƒ += โˆ‚๐—โ‚ƒ * โ„’.kron(โ„’.kron(aux', aux'), aux') - # A_mult_kron_power_3_B!(โˆ‚โˆ‡โ‚ƒ, โˆ‚๐—โ‚ƒ, aux') # not a good idea because filling an existing matrix one by one is slow - # โˆ‚โˆ‡โ‚ƒ += A_mult_kron_power_3_B(โˆ‚๐—โ‚ƒ, aux') # this is slower somehow - - # end # timeit_debug - # @timeit_debug timer "Step 5" begin - - # this is very slow - โˆ‚โˆ‡โ‚ƒ += โˆ‚๐—โ‚ƒ * compressed_kronยณ(aux', rowmask = unique(findnz(โˆ‚๐—โ‚ƒ)[2]), sparse_preallocation = โ„‚.tmp_sparse_prealloc4) # , timer = timer) - # โˆ‚โˆ‡โ‚ƒ += โˆ‚๐—โ‚ƒ * โ„’.kron(aux', aux', aux') - - # end # timeit_debug - # @timeit_debug timer "Step 6" begin - - โˆ‚kronkronaux = ๐”โˆ‡โ‚ƒt * โˆ‚๐—โ‚ƒ * ๐‚โ‚ƒt - - fill_kron_adjoint!(โˆ‚kronaux, โˆ‚aux, โˆ‚kronkronaux, kronaux, aux) - - fill_kron_adjoint!(โˆ‚aux, โˆ‚aux, โˆ‚kronaux, aux, aux) - - # end # timeit_debug - # @timeit_debug timer "Step 7" begin - - # aux = Mโ‚ƒ.๐’๐ * โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ - โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ += Mโ‚ƒ.๐’๐' * โˆ‚aux - - # ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚‹,:] ; zeros(size(๐’โ‚)[2] - nโ‚‹, nโ‚‘โ‚‹^2)] - โˆ‚๐’โ‚‚[iโ‚‹,:] += โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ[1:length(iโ‚‹),:] - - # ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚Š,:] - # zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹^2)] - โˆ‚๐’โ‚‚[iโ‚Š,:] += โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] - - - # โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = [ - ## (๐’โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:] ; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)])[iโ‚Š,:] - ## โ„’.diagm(ones(n))[iโ‚Š,:] * (๐’โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:] ; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)]) - # โ„’.diagm(ones(n))[iโ‚Š,:] * ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ - # ๐’โ‚‚ - # zeros(nโ‚‹ + nโ‚‘, nโ‚‘โ‚‹^2) - # ]; - โˆ‚๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„’.diagm(ones(n))[iโ‚Š,:]' * โˆ‚โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ[1:length(iโ‚Š),:] - - โˆ‚๐’โ‚‚ += โˆ‚โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ[length(iโ‚Š) .+ (1:size(๐’โ‚‚,1)),:] - - โˆ‚๐’โ‚‚ += โˆ‚๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ * kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' - - โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ += ๐’โ‚‚t * โˆ‚๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ - - - # ๐’โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + ๐’โ‚ * [๐’โ‚‚[iโ‚‹,:] ; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)] - # ๐’โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) + ๐’โ‚ * ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ - โˆ‚๐’โ‚ += โˆ‚๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ * [๐’โ‚‚[iโ‚‹,:] ; zeros(nโ‚‘ + 1, nโ‚‘โ‚‹^2)]' - - # โˆ‚๐’โ‚‚[iโ‚‹,:] += spdiagm(ones(size(๐’โ‚‚,1)))[iโ‚‹,:]' * ๐’โ‚' * โˆ‚๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:length(iโ‚‹),:] - โˆ‚๐’โ‚‚โ•ฑ๐ŸŽ = ๐’โ‚' * โˆ‚๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ - โˆ‚๐’โ‚‚[iโ‚‹,:] += โˆ‚๐’โ‚‚โ•ฑ๐ŸŽ[1:length(iโ‚‹),:] - - # end # timeit_debug - # @timeit_debug timer "Step 8" begin - - ### - # B = Mโ‚ƒ.๐”โ‚ƒ * (tmpkron + Mโ‚ƒ.๐โ‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚แตฃฬƒ + Mโ‚ƒ.๐โ‚‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚‚แตฃฬƒ + โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)) * Mโ‚ƒ.๐‚โ‚ƒ - โˆ‚tmpkron += ๐”โ‚ƒt * โˆ‚B * ๐‚โ‚ƒt - โˆ‚tmpkron += Mโ‚ƒ.๐โ‚โ‚—ฬ„' * ๐”โ‚ƒt * โˆ‚B * ๐‚โ‚ƒt * Mโ‚ƒ๐โ‚แตฃฬƒt - โˆ‚tmpkron += Mโ‚ƒ.๐โ‚‚โ‚—ฬ„' * ๐”โ‚ƒt * โˆ‚B * ๐‚โ‚ƒt * Mโ‚ƒ๐โ‚‚แตฃฬƒt - - โˆ‚kronkron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = ๐”โ‚ƒt * โˆ‚B * ๐‚โ‚ƒt - - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚kronkron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - - fill_kron_adjoint!(โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, โˆ‚kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - - # tmpkron = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘,Mโ‚‚.๐›”) - fill_kron_adjoint_โˆ‚A!(โˆ‚tmpkron, โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”) - # A = spinv * โˆ‡โ‚โ‚Š - โˆ‚โˆ‡โ‚โ‚Š += spinv' * โˆ‚A - โˆ‚spinv += โˆ‚A * โˆ‡โ‚โ‚Š' - - # โˆ‡โ‚โ‚Š = sparse(โˆ‡โ‚[:,1:nโ‚Š] * spdiagm(ones(n))[iโ‚Š,:]) - โˆ‚โˆ‡โ‚[:,1:nโ‚Š] += โˆ‚โˆ‡โ‚โ‚Š * โ„’.I(n)[:,iโ‚Š] - - # spinv = sparse(inv(โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€)) - โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -spinv' * โˆ‚spinv * spinv' - - # โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.diagm(ones(n))[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] - โˆ‚โˆ‡โ‚[:,1:nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] * ๐’โ‚[iโ‚Š,1:nโ‚‹]' - โˆ‚โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] -= โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ - - โˆ‚๐’โ‚[iโ‚Š,1:nโ‚‹] -= โˆ‡โ‚[:,1:nโ‚Š]' * โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ * โ„’.I(n)[:,iโ‚‹] - - # # ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] - # # zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)]; - โˆ‚๐’โ‚[iโ‚Š,:] += โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ[1:length(iโ‚Š),:] - - # ###### โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] - # # โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = [โ„’.I(size(๐’โ‚,1))[iโ‚Š,:] * ๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ - # # ๐’โ‚ - # # spdiagm(ones(nโ‚‘โ‚‹))[[range(1,nโ‚‹)...,nโ‚‹ + 1 .+ range(1,nโ‚‘)...],:]]; - โˆ‚๐’โ‚ += โ„’.I(size(๐’โ‚,1))[:,iโ‚Š] * โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[1:length(iโ‚Š),:] * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘' - โˆ‚๐’โ‚ += โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[length(iโ‚Š) .+ (1:size(๐’โ‚,1)),:] - - โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ += ๐’โ‚' * โ„’.I(size(๐’โ‚,1))[:,iโ‚Š] * โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹[1:length(iโ‚Š),:] - - # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) spdiagm(ones(nโ‚‘ + 1))[1,:] zeros(nโ‚‘ + 1, nโ‚‘)]; - โˆ‚๐’โ‚[iโ‚‹,:] += โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:length(iโ‚‹), :] - - # ๐’โ‚ = [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]] - โˆ‚๐‘บโ‚ = [โˆ‚๐’โ‚[:,1:nโ‚‹] โˆ‚๐’โ‚[:,nโ‚‹+2:end]] - - # end # timeit_debug - # end # timeit_debug - - return NoTangent(), โˆ‚โˆ‡โ‚, โˆ‚โˆ‡โ‚‚, โˆ‚โˆ‡โ‚ƒ, โˆ‚๐‘บโ‚, โˆ‚๐’โ‚‚, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - return (๐’โ‚ƒ, solved), third_order_solution_pullback -end - -function rrule(::typeof(solve_sylvester_equation), - A::M, - B::N, - C::O, - ๐•Šโ„‚::sylvester_workspace; - initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), - sylvester_algorithm::Symbol = :doubling, - acceptance_tol::AbstractFloat = 1e-10, - tol::AbstractFloat = 1e-14, - # timer::TimerOutput = TimerOutput(), - verbose::Bool = false) where {M <: AbstractMatrix{Float64}, N <: AbstractMatrix{Float64}, O <: AbstractMatrix{Float64}} - - P, solved = solve_sylvester_equation(A, B, C, ๐•Šโ„‚, - sylvester_algorithm = sylvester_algorithm, - tol = tol, - verbose = verbose, - initial_guess = initial_guess) - - println("C norm: $(โ„’.norm(C))") - # pullback - function solve_sylvester_equation_pullback(โˆ‚P) - if โ„’.norm(โˆ‚P[1]) < tol return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() end - - โˆ‚C, slvd = solve_sylvester_equation(A', B', โˆ‚P[1], ๐•Šโ„‚, - sylvester_algorithm = sylvester_algorithm, - tol = tol, - verbose = verbose) - - solved = solved && slvd - - โˆ‚A = โˆ‚C * B' * P' - - โˆ‚B = P' * A' * โˆ‚C - - return NoTangent(), โˆ‚A, โˆ‚B, โˆ‚C, NoTangent() - end - - return (P, solved), solve_sylvester_equation_pullback -end - -function rrule(::typeof(solve_lyapunov_equation), - A::AbstractMatrix{Float64}, - C::AbstractMatrix{Float64}, - workspace::lyapunov_workspace; - lyapunov_algorithm::Symbol = :doubling, - tol::AbstractFloat = 1e-14, - acceptance_tol::AbstractFloat = 1e-12, - # timer::TimerOutput = TimerOutput(), - verbose::Bool = false) - - P, solved = solve_lyapunov_equation(A, C, workspace, lyapunov_algorithm = lyapunov_algorithm, tol = tol, verbose = verbose) - - # pullback - # https://arxiv.org/abs/2011.11430 - function solve_lyapunov_equation_pullback(โˆ‚P) - if โ„’.norm(โˆ‚P[1]) < tol return NoTangent(), NoTangent(), NoTangent(), NoTangent() end - - โˆ‚C, slvd = solve_lyapunov_equation(A', โˆ‚P[1], workspace, lyapunov_algorithm = lyapunov_algorithm, tol = tol, verbose = verbose) - - solved = solved && slvd - - โˆ‚A = โˆ‚C * A * P' + โˆ‚C' * A * P - - return NoTangent(), โˆ‚A, โˆ‚C, NoTangent() - end - - return (P, solved), solve_lyapunov_equation_pullback -end - -function rrule(::typeof(find_shocks), - ::Val{:LagrangeNewton}, - initial_guess::Vector{Float64}, - kron_buffer::Vector{Float64}, - kron_buffer2::AbstractMatrix{Float64}, - J::โ„’.Diagonal{Bool, Vector{Bool}}, - ๐’โฑ::AbstractMatrix{Float64}, - ๐’โฑยฒแต‰::AbstractMatrix{Float64}, - shock_independent::Vector{Float64}; - max_iter::Int = 1000, - tol::Float64 = 1e-13) - - x, matched = find_shocks(Val(:LagrangeNewton), - initial_guess, - kron_buffer, - kron_buffer2, - J, - ๐’โฑ, - ๐’โฑยฒแต‰, - shock_independent, - max_iter = max_iter, - tol = tol) - - tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x)), x) - - ฮป = tmp' \ x * 2 - - fXฮปp = [reshape(2 * ๐’โฑยฒแต‰' * ฮป, size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' - -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - โ„’.kron!(kron_buffer, x, x) - - xฮป = โ„’.kron(x,ฮป) - - - โˆ‚shock_independent = similar(shock_independent) - - # โˆ‚๐’โฑ = similar(๐’โฑ) - - # โˆ‚๐’โฑยฒแต‰ = similar(๐’โฑยฒแต‰) - - function find_shocks_pullback(โˆ‚x) - โˆ‚x = vcat(โˆ‚x[1], zero(ฮป)) - - S = -fXฮปp' \ โˆ‚x - - copyto!(โˆ‚shock_independent, S[length(initial_guess)+1:end]) - - # copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:length(initial_guess)], ฮป) - โ„’.kron(x, S[length(initial_guess)+1:end])) - โˆ‚๐’โฑ = S[1:length(initial_guess)] * ฮป' - S[length(initial_guess)+1:end] * x' - - # copyto!(โˆ‚๐’โฑยฒแต‰, 2 * โ„’.kron(S[1:length(initial_guess)], xฮป) - โ„’.kron(kron_buffer, S[length(initial_guess)+1:end])) - โˆ‚๐’โฑยฒแต‰ = 2 * S[1:length(initial_guess)] * xฮป' - S[length(initial_guess)+1:end] * kron_buffer' - - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’โฑ, โˆ‚๐’โฑยฒแต‰, โˆ‚shock_independent, NoTangent(), NoTangent() - end - - return (x, matched), find_shocks_pullback -end - -function rrule(::typeof(find_shocks), - ::Val{:LagrangeNewton}, - initial_guess::Vector{Float64}, - kron_buffer::Vector{Float64}, - kron_bufferยฒ::Vector{Float64}, - kron_buffer2::AbstractMatrix{Float64}, - kron_buffer3::AbstractMatrix{Float64}, - kron_buffer4::AbstractMatrix{Float64}, - J::โ„’.Diagonal{Bool, Vector{Bool}}, - ๐’โฑ::AbstractMatrix{Float64}, - ๐’โฑยฒแต‰::AbstractMatrix{Float64}, - ๐’โฑยณแต‰::AbstractMatrix{Float64}, - shock_independent::Vector{Float64}; - max_iter::Int = 1000, - tol::Float64 = 1e-13) - - x, matched = find_shocks(Val(:LagrangeNewton), - initial_guess, - kron_buffer, - kron_bufferยฒ, - kron_buffer2, - kron_buffer3, - kron_buffer4, - J, - ๐’โฑ, - ๐’โฑยฒแต‰, - ๐’โฑยณแต‰, - shock_independent, - max_iter = max_iter, - tol = tol) - - โ„’.kron!(kron_buffer, x, x) - - โ„’.kron!(kron_bufferยฒ, x, kron_buffer) - - tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x)), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(length(x)), kron_buffer) - - ฮป = tmp' \ x * 2 - - fXฮปp = [reshape((2 * ๐’โฑยฒแต‰ + 6 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(length(x)), โ„’.kron(โ„’.I(length(x)),x)))' * ฮป, size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' - -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - xฮป = โ„’.kron(x,ฮป) - - xxฮป = โ„’.kron(x,xฮป) - - function find_shocks_pullback(โˆ‚x) - โˆ‚x = vcat(โˆ‚x[1], zero(ฮป)) - - S = -fXฮปp' \ โˆ‚x - - โˆ‚shock_independent = S[length(initial_guess)+1:end] - - โˆ‚๐’โฑ = โ„’.kron(S[1:length(initial_guess)], ฮป) - โ„’.kron(x, S[length(initial_guess)+1:end]) - - โˆ‚๐’โฑยฒแต‰ = 2 * โ„’.kron(S[1:length(initial_guess)], xฮป) - โ„’.kron(kron_buffer, S[length(initial_guess)+1:end]) - - โˆ‚๐’โฑยณแต‰ = 3 * โ„’.kron(S[1:length(initial_guess)], xxฮป) - โ„’.kron(kron_bufferยฒ,S[length(initial_guess)+1:end]) - - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), โˆ‚๐’โฑ, โˆ‚๐’โฑยฒแต‰, โˆ‚๐’โฑยณแต‰, โˆ‚shock_independent, NoTangent(), NoTangent() - end - - return (x, matched), find_shocks_pullback -end - - -function rrule(::typeof(calculate_inversion_filter_loglikelihood), - ::Val{:first_order}, - state::Vector{Vector{Float64}}, - ๐’::Matrix{Float64}, - data_in_deviations::Matrix{Float64}, - observables::Union{Vector{String}, Vector{Symbol}}, - constants::constants, - ws::inversion_workspace{Float64}; - # timer::TimerOutput = TimerOutput(), - warmup_iterations::Int = 0, - on_failure_loglikelihood = -Inf, - presample_periods::Int = 0, - opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton) - T = constants.post_model_macro - # @timeit_debug timer "Inversion filter - forward" begin - - # first order - state = copy(state[1]) - - precision_factor = 1.0 - - n_obs = size(data_in_deviations,2) - - obs_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - - tโป = T.past_not_future_and_mixed_idx - - shocksยฒ = 0.0 - logabsdets = 0.0 - - @assert warmup_iterations == 0 "Warmup iterations not yet implemented for reverse-mode automatic differentiation." - - state = [copy(state) for _ in 1:size(data_in_deviations,2)+1] - - shocksยฒ = 0.0 - logabsdets = 0.0 - - y = zeros(length(obs_idx)) - x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] - - jac = ๐’[obs_idx,end-T.nExo+1:end] - - if T.nExo == length(observables) - logabsdets = โ„’.logabsdet(jac)[1] # ./ precision_factor - - jacdecomp = โ„’.lu(jac, check = false) - - if !โ„’.issuccess(jacdecomp) - if opts.verbose println("Inversion filter failed") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - invjac = inv(jacdecomp) - else - logabsdets = sum(x -> log(abs(x)), โ„’.svdvals(jac)) #' ./ precision_factor - # jacdecomp = โ„’.svd(jac) - invjac = โ„’.pinv(jac) - end - - logabsdets *= size(data_in_deviations,2) - presample_periods - - if !isfinite(logabsdets) - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - @views ๐’obs = ๐’[obs_idx,1:end-T.nExo] - - for i in axes(data_in_deviations,2) - @views โ„’.mul!(y, ๐’obs, state[i][tโป]) - @views โ„’.axpby!(1, data_in_deviations[:,i], -1, y) - โ„’.mul!(x[i],invjac,y) - # x = ๐’[obs_idx,end-T.nExo+1:end] \ (data_in_deviations[:,i] - ๐’[obs_idx,1:end-T.nExo] * state[tโป]) - - if i > presample_periods - shocksยฒ += sum(abs2,x[i]) - if !isfinite(shocksยฒ) - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - end - - โ„’.mul!(state[i+1], ๐’, vcat(state[i][tโป], x[i])) - # state[i+1] = ๐’ * vcat(state[i][tโป], x[i]) - end - - llh = -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 - - if llh < -1e12 - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - โˆ‚๐’ = zero(๐’) - - โˆ‚๐’แต—โป = copy(โˆ‚๐’[tโป,:]) - - โˆ‚data_in_deviations = zero(data_in_deviations) - - # Allocate or reuse workspaces for pullback - n_periods = size(data_in_deviations,2) - 1 - if size(ws.โˆ‚data) != (length(tโป), n_periods) - ws.โˆ‚data = zeros(length(tโป), n_periods) - else - fill!(ws.โˆ‚data, zero(eltype(ws.โˆ‚data))) - end - โˆ‚data = ws.โˆ‚data - - โˆ‚state = zero(state[1]) - - # precomputed matrices - Mยน = ๐’[obs_idx, 1:end-T.nExo]' * invjac' - Mยฒ = ๐’[tโป,1:end-T.nExo]' - Mยน * ๐’[tโป,end-T.nExo+1:end]' - Mยณ = invjac' * ๐’[tโป,end-T.nExo+1:end]' - - โˆ‚Stmp = [copy(Mยน) for _ in 1:size(data_in_deviations,2)-1] - - for t in 2:size(data_in_deviations,2)-1 - โ„’.mul!(โˆ‚Stmp[t], Mยฒ, โˆ‚Stmp[t-1]) - # โˆ‚Stmp[t] = Mยฒ * โˆ‚Stmp[t-1] - end - - # Allocate or reuse workspaces for temporary matrices - if size(ws.โˆ‚_tmp1) != (T.nExo, length(tโป) + T.nExo) - ws.โˆ‚_tmp1 = zeros(Float64, T.nExo, length(tโป) + T.nExo) - else - fill!(ws.โˆ‚_tmp1, zero(Float64)) - end - tmp1 = ws.โˆ‚_tmp1 - - if size(ws.โˆ‚_tmp2) != (length(tโป), length(tโป) + T.nExo) - ws.โˆ‚_tmp2 = zeros(Float64, length(tโป), length(tโป) + T.nExo) - else - fill!(ws.โˆ‚_tmp2, zero(Float64)) - end - tmp2 = ws.โˆ‚_tmp2 - - if size(ws.โˆ‚_tmp3) != (length(tโป) + T.nExo,) - ws.โˆ‚_tmp3 = zeros(Float64, length(tโป) + T.nExo) - else - fill!(ws.โˆ‚_tmp3, zero(Float64)) - end - tmp3 = ws.โˆ‚_tmp3 - - if size(ws.โˆ‚๐’tโป) != size(tmp2) - ws.โˆ‚๐’tโป = copy(tmp2) - else - fill!(ws.โˆ‚๐’tโป, zero(Float64)) - end - โˆ‚๐’tโป = ws.โˆ‚๐’tโป - # โˆ‚๐’obs_idx = copy(tmp1) - - # end # timeit_debug - # pullback - function inversion_pullback(โˆ‚llh) - # @timeit_debug timer "Inversion filter - pullback" begin - - for t in reverse(axes(data_in_deviations,2)) - โˆ‚state[tโป] .= Mยฒ * โˆ‚state[tโป] - - if t > presample_periods - โˆ‚state[tโป] += Mยน * x[t] - - โˆ‚data_in_deviations[:,t] -= invjac' * x[t] - - โˆ‚๐’[obs_idx, :] += invjac' * x[t] * vcat(state[t][tโป], x[t])' - - if t > 1 - โˆ‚data[:,t:end] .= Mยฒ * โˆ‚data[:,t:end] - - โˆ‚data[:,t-1] += Mยน * x[t] - - โˆ‚data_in_deviations[:,t-1] += Mยณ * โˆ‚data[:,t-1:end] * ones(size(data_in_deviations,2) - t + 1) - - for tt in t-1:-1:1 - for (i,v) in enumerate(tโป) - copyto!(tmp3::Vector{Float64}, i::Int, state[tt]::Vector{Float64}, v::Int, 1) - end - - copyto!(tmp3, length(tโป) + 1, x[tt], 1, T.nExo) - - โ„’.mul!(tmp1, x[t], tmp3') - - โ„’.mul!(โˆ‚๐’tโป, โˆ‚Stmp[t-tt], tmp1, 1, 1) - - end - end - end - end - - โˆ‚๐’[tโป,:] += โˆ‚๐’tโป - - โˆ‚๐’[obs_idx, :] -= Mยณ * โˆ‚๐’tโป - - โˆ‚๐’[obs_idx,end-T.nExo+1:end] -= (size(data_in_deviations,2) - presample_periods) * invjac' / 2 - - # end # timeit_debug - - return NoTangent(), NoTangent(), [โˆ‚state * โˆ‚llh], โˆ‚๐’ * โˆ‚llh, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - return llh, inversion_pullback -end - - -function rrule(::typeof(calculate_inversion_filter_loglikelihood), - ::Val{:pruned_second_order}, - state::Vector{Vector{Float64}}, - ๐’::Vector{AbstractMatrix{Float64}}, - data_in_deviations::Matrix{Float64}, - observables::Union{Vector{String}, Vector{Symbol}}, - constants::constants, - ws::inversion_workspace{Float64}; - # timer::TimerOutput = TimerOutput(), - on_failure_loglikelihood = -Inf, - warmup_iterations::Int = 0, - presample_periods::Int = 0, - opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)# where S <: Real - T = constants.post_model_macro - # @timeit_debug timer "Inversion filter pruned 2nd - forward" begin - # @timeit_debug timer "Preallocation" begin - - precision_factor = 1.0 - - n_obs = size(data_in_deviations,2) - - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - - shocksยฒ = 0.0 - logabsdets = 0.0 - - cc = ensure_computational_constants!(constants) - s_in_sโบ = cc.s_in_s - sv_in_sโบ = cc.s_in_sโบ - e_in_sโบ = cc.e_in_sโบ - - tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, e_in_sโบ) |> sparse - shockยฒ_idxs = tmp.nzind - - shockvarยฒ_idxs = setdiff(shock_idxs, shockยฒ_idxs) - - tmp = โ„’.kron(sv_in_sโบ, sv_in_sโบ) |> sparse - var_volยฒ_idxs = tmp.nzind - - tmp = โ„’.kron(s_in_sโบ, s_in_sโบ) |> sparse - varยฒ_idxs = tmp.nzind - - ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] - ๐’โปยนแต‰ = ๐’[1][T.past_not_future_and_mixed_idx,end-T.nExo+1:end] - ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] - ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] - ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] - - ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] - ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] - ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] - ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] - ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] - - ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› - ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป - ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ - ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ - ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ - - stateโ‚ = state[1][T.past_not_future_and_mixed_idx] - stateโ‚‚ = state[2][T.past_not_future_and_mixed_idx] - - kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] - - J = โ„’.I(T.nExo) - - kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) - - kron_buffer3 = โ„’.kron(J, zeros(T.nPast_not_future_and_mixed + 1)) - - x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] - - stateยนโป = stateโ‚ - - stateยนโป_vol = vcat(stateยนโป, 1) - - stateยฒโป = stateโ‚‚ - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) - - ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 - - aug_stateโ‚ = [copy([stateโ‚; 1; ones(T.nExo)]) for _ in 1:size(data_in_deviations,2)] - aug_stateโ‚‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] - - tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[1])), x[1]) - - jacc = [zero(tmp) for _ in 1:size(data_in_deviations,2)] - - jacct = copy(tmp') - - ฮป = [zeros(size(tmp, 1)) for _ in 1:size(data_in_deviations,2)] - - ฮป[1] = copy(tmp' \ x[1] * 2) - - fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' - -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] - - kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) - - kronxฮป = [zero(kronxฮป_tmp) for _ in 1:size(data_in_deviations,2)] - - kronstateยนโป_vol = zeros((T.nPast_not_future_and_mixed + 1)^2) - - kronaug_stateโ‚ = zeros(length(aug_stateโ‚[1])^2) - - shock_independent = zeros(size(data_in_deviations,1)) - - init_guess = zeros(size(๐’โฑ, 2)) - - tmp = zeros(size(๐’โฑ, 2) * size(๐’โฑ, 2)) - - lI = -2 * vec(โ„’.I(size(๐’โฑ, 2))) - - # end # timeit_debug - # @timeit_debug timer "Main loop" begin - - for i in axes(data_in_deviations,2) - # stateยนโป = stateโ‚ - - # stateยนโป_vol = vcat(stateยนโป, 1) - - # stateยฒโป = stateโ‚‚ - - copyto!(stateยนโป_vol, 1, stateโ‚, 1) - - copyto!(shock_independent, data_in_deviations[:,i]) - - โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - - โ„’.mul!(shock_independent, ๐’ยนโป, stateโ‚‚, -1, 1) - - โ„’.kron!(kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) - - โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kronstateยนโป_vol, -1/2, 1) - - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) - โ„’.kron!(kron_buffer3, J, stateยนโป_vol) - - โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer3) - - โ„’.axpy!(1, ๐’ยนแต‰, ๐’โฑ) - - init_guess *= 0 - - # @timeit_debug timer "Find shocks" begin - x[i], matched = find_shocks(Val(filter_algorithm), - init_guess, - kronxx[i], - kron_buffer2, - J, - ๐’โฑ, - ๐’โฑยฒแต‰, - shock_independent, - # max_iter = 100 - ) - # end # timeit_debug - - if !matched - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[i])), x[i]) - โ„’.kron!(kron_buffer2, J, x[i]) - - โ„’.mul!(jacc[i], ๐’โฑยฒแต‰, kron_buffer2) - - โ„’.axpby!(1, ๐’โฑ, 2, jacc[i]) - - copy!(jacct, jacc[i]') - - jacc_fact = try - โ„’.factorize(jacct) # otherwise this fails for nshocks > nexo - catch - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - try - โ„’.ldiv!(ฮป[i], jacc_fact, x[i]) - catch - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - โ„’.rmul!(ฮป[i], 2) - - # fXฮปp[i] = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) jacc[i]' - # -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - โ„’.mul!(tmp, ๐’โฑยฒแต‰', ฮป[i]) - โ„’.axpby!(1, lI, 2, tmp) - - fXฮปp[i][1:size(๐’โฑ, 2), 1:size(๐’โฑ, 2)] = tmp - fXฮปp[i][size(๐’โฑ, 2)+1:end, 1:size(๐’โฑ, 2)] = -jacc[i] - fXฮปp[i][1:size(๐’โฑ, 2), size(๐’โฑ, 2)+1:end] = jacct - - โ„’.kron!(kronxx[i], x[i], x[i]) - - โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) - - if i > presample_periods - # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) - logabsdets += โ„’.logabsdet(jacc_fact)[1] - else - logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) - end - - shocksยฒ += sum(abs2,x[i]) - - if !isfinite(logabsdets) || !isfinite(shocksยฒ) - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - end - - # aug_stateโ‚[i] = [stateโ‚; 1; x[i]] - # aug_stateโ‚‚[i] = [stateโ‚‚; 0; zero(x[1])] - copyto!(aug_stateโ‚[i], 1, stateโ‚, 1) - copyto!(aug_stateโ‚[i], length(stateโ‚) + 2, x[i], 1) - copyto!(aug_stateโ‚‚[i], 1, stateโ‚‚, 1) - - # stateโ‚, stateโ‚‚ = [๐’โปยน * aug_stateโ‚, ๐’โปยน * aug_stateโ‚‚ + ๐’โปยฒ * โ„’.kron(aug_stateโ‚, aug_stateโ‚) / 2] # strictly following Andreasen et al. (2018) - โ„’.mul!(stateโ‚, ๐’โปยน, aug_stateโ‚[i]) - - โ„’.mul!(stateโ‚‚, ๐’โปยน, aug_stateโ‚‚[i]) - โ„’.kron!(kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) - โ„’.mul!(stateโ‚‚, ๐’โปยฒ, kronaug_stateโ‚, 1/2, 1) - end - - # end # timeit_debug - # end # timeit_debug - - โˆ‚data_in_deviations = similar(data_in_deviations) - - โˆ‚aug_stateโ‚ = zero(aug_stateโ‚[1]) - - โˆ‚aug_stateโ‚‚ = zero(aug_stateโ‚‚[1]) - - โˆ‚kronaug_stateโ‚ = zeros(length(aug_stateโ‚[1])^2) - - โˆ‚kronIx = zero(โ„’.kron(โ„’.I(length(x[1])), x[1])) - - โˆ‚kronIstateยนโป_vol = zero(โ„’.kron(J, stateยนโป_vol)) - - โˆ‚kronstateยนโป_vol = zero(โ„’.kron(stateยนโป_vol, stateยนโป_vol)) - - function inversion_filter_loglikelihood_pullback(โˆ‚llh) - # @timeit_debug timer "Inversion filter pruned 2nd - pullback" begin - # @timeit_debug timer "Preallocation" begin - - โˆ‚๐’โฑ = zero(๐’โฑ) - โˆ‚๐’โฑยฒแต‰ = zero(๐’โฑยฒแต‰) - - โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) - โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) - - โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) - โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) - - โˆ‚๐’โปยน = zero(๐’โปยน) - โˆ‚๐’โปยฒ = zero(๐’โปยฒ) - - โˆ‚๐’ยนโป = zero(๐’ยนโป) - - โˆ‚stateยนโป_vol = zero(stateยนโป_vol) - โˆ‚x = zero(x[1]) - โˆ‚state = [zeros(T.nPast_not_future_and_mixed), zeros(T.nPast_not_future_and_mixed)] - - kronSฮป = zeros(length(cond_var_idx) * T.nExo) - kronxS = zeros(T.nExo * length(cond_var_idx)) - - # end # timeit_debug - # @timeit_debug timer "Main loop" begin - - for i in reverse(axes(data_in_deviations,2)) - # stateโ‚, stateโ‚‚ = [๐’โปยน * aug_stateโ‚[i], ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) / 2] - # stateโ‚ = ๐’โปยน * aug_stateโ‚[i] - # โˆ‚๐’โปยน += โˆ‚state[1] * aug_stateโ‚[i]' - โ„’.mul!(โˆ‚๐’โปยน, โˆ‚state[1], aug_stateโ‚[i]', 1, 1) - - # โˆ‚aug_stateโ‚ = ๐’โปยน' * โˆ‚state[1] - โ„’.mul!(โˆ‚aug_stateโ‚, ๐’โปยน', โˆ‚state[1]) - - # stateโ‚‚ = ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) / 2 - # โˆ‚๐’โปยน += โˆ‚state[2] * aug_stateโ‚‚[i]' - โ„’.mul!(โˆ‚๐’โปยน, โˆ‚state[2], aug_stateโ‚‚[i]', 1, 1) - - # โˆ‚aug_stateโ‚‚ = ๐’โปยน' * โˆ‚state[2] - โ„’.mul!(โˆ‚aug_stateโ‚‚, ๐’โปยน', โˆ‚state[2]) - - # โˆ‚๐’โปยฒ += โˆ‚state[2] * โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i])' / 2 - โ„’.kron!(kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) - โ„’.mul!(โˆ‚๐’โปยฒ, โˆ‚state[2], kronaug_stateโ‚', 1/2, 1) - - # โˆ‚kronaug_stateโ‚ = ๐’โปยฒ' * โˆ‚state[2] / 2 - โ„’.mul!(โˆ‚kronaug_stateโ‚, ๐’โปยฒ', โˆ‚state[2]) - โ„’.rdiv!(โˆ‚kronaug_stateโ‚, 2) - - fill_kron_adjoint!(โˆ‚aug_stateโ‚, โˆ‚aug_stateโ‚, โˆ‚kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) - - if i > 1 && i < size(data_in_deviations,2) - โˆ‚state[1] *= 0 - โˆ‚state[2] *= 0 - end - - # aug_stateโ‚ = [stateโ‚; 1; x] - # โˆ‚state[1] += โˆ‚aug_stateโ‚[1:length(โˆ‚state[1])] - โ„’.axpy!(1, โˆ‚aug_stateโ‚[1:length(โˆ‚state[1])], โˆ‚state[1]) - - โˆ‚x = โˆ‚aug_stateโ‚[T.nPast_not_future_and_mixed+2:end] - - # aug_stateโ‚‚ = [stateโ‚‚; 0; zero(x)] - # โˆ‚state[2] += โˆ‚aug_stateโ‚‚[1:length(โˆ‚state[1])] - โ„’.axpy!(1, โˆ‚aug_stateโ‚‚[1:length(โˆ‚state[1])], โˆ‚state[2]) - - # shocksยฒ += sum(abs2,x[i]) - if i < size(data_in_deviations,2) - โˆ‚x -= copy(x[i]) - else - โˆ‚x += copy(x[i]) - end - - # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] - โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) - inv(jacc[i])' - else - โ„’.pinv(jacc[i])' - end - catch - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x[1]) - # โˆ‚kronIx = ๐’โฑยฒแต‰' * โˆ‚jacc - โ„’.mul!(โˆ‚kronIx, ๐’โฑยฒแต‰', โˆ‚jacc) - - if i < size(data_in_deviations,2) - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -J) - else - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, J) - end - - # โˆ‚๐’โฑยฒแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' - โ„’.kron!(kron_buffer2, J, x[i]) - - โ„’.mul!(โˆ‚๐’โฑยฒแต‰, โˆ‚jacc, kron_buffer2', -1, 1) - - # find_shocks - โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) - # S = vcat(โˆ‚x, zero(ฮป[i])) - - S = fXฮปp[i]' \ โˆ‚xฮป - # โ„’.ldiv!(fXฮปp[i]', S) - - if i < size(data_in_deviations,2) - S *= -1 - end - - โˆ‚shock_independent = S[T.nExo+1:end] # fine - - # โˆ‚๐’โฑ = (S[1:T.nExo] * ฮป[i]' - S[T.nExo+1:end] * x[i]') # fine - # โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine - # copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) - โ„’.kron!(kronSฮป, S[1:T.nExo], ฮป[i]) - โ„’.kron!(kronxS, x[i], S[T.nExo+1:end]) - โ„’.axpy!(-1, kronxS, kronSฮป) - copyto!(โˆ‚๐’โฑ, kronSฮป) - # โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine - โ„’.axpy!(-1/2, โˆ‚jacc, โˆ‚๐’โฑ) - - โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], โ„’.kron(x[i], ฮป[i])) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) - # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo+1:end] * kronxx[i]' - - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) - โˆ‚stateยนโป_vol *= 0 - # โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ - โ„’.mul!(โˆ‚kronIstateยนโป_vol, ๐’ยฒโปแต‰', โˆ‚๐’โฑ) - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, J) - - stateยนโป_vol = aug_stateโ‚[i][1:T.nPast_not_future_and_mixed+1] - - # โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ - โ„’.axpy!(1, โˆ‚๐’โฑ, โˆ‚๐’ยนแต‰) - - # โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' - โ„’.kron!(โˆ‚kronIstateยนโป_vol, J, stateยนโป_vol) - โ„’.mul!(โˆ‚๐’ยฒโปแต‰, โˆ‚๐’โฑ, โˆ‚kronIstateยนโป_vol', 1, 1) - - - # shock_independent = copy(data_in_deviations[:,i]) - โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - # โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' - โ„’.mul!(โˆ‚๐’ยนโปแต›, โˆ‚shock_independent, stateยนโป_vol', -1, 1) - - # โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent - โ„’.mul!(โˆ‚stateยนโป_vol, ๐’ยนโปแต›', โˆ‚shock_independent, -1, 1) - - # โ„’.mul!(shock_independent, ๐’ยนโป, stateยฒโป, -1, 1) - # โˆ‚๐’ยนโป -= โˆ‚shock_independent * aug_stateโ‚‚[i][1:T.nPast_not_future_and_mixed]' - โ„’.mul!(โˆ‚๐’ยนโป, โˆ‚shock_independent, aug_stateโ‚‚[i][1:T.nPast_not_future_and_mixed]', -1, 1) - - # โˆ‚state[2] -= ๐’ยนโป' * โˆ‚shock_independent - โ„’.mul!(โˆ‚state[2], ๐’ยนโป', โˆ‚shock_independent, -1, 1) - - # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) - # โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 - โ„’.kron!(โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) - โ„’.mul!(โˆ‚๐’ยฒโปแต›, โˆ‚shock_independent, โˆ‚kronstateยนโป_vol', -1/2, 1) - - # โˆ‚kronstateยนโป_vol = -๐’ยฒโปแต›' * โˆ‚shock_independent / 2 - โ„’.mul!(โˆ‚kronstateยนโป_vol, ๐’ยฒโปแต›', โˆ‚shock_independent) - โ„’.rdiv!(โˆ‚kronstateยนโป_vol, -2) - - fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) - - # stateยนโป_vol = vcat(stateยนโป, 1) - # โˆ‚state[1] += โˆ‚stateยนโป_vol[1:end-1] - โ„’.axpy!(1, โˆ‚stateยนโป_vol[1:end-1], โˆ‚state[1]) - end - - # end # timeit_debug - # @timeit_debug timer "Post allocation" begin - - โˆ‚๐’ = [zero(๐’[1]), zeros(size(๐’[2]))] - - โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] .+= โˆ‚๐’ยนแต‰ - โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] .+= โˆ‚๐’ยฒโปแต‰ - โ„’.rdiv!(โˆ‚๐’โฑยฒแต‰, 2) - โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] .+= โˆ‚๐’โฑยฒแต‰# / 2 - - โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] .+= โˆ‚๐’ยนโปแต› - โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] .+= โˆ‚๐’ยฒโปแต› - - โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] .+= โˆ‚๐’โปยน - โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] .+= โˆ‚๐’โปยฒ - - โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] .+= โˆ‚๐’ยนโป - - # โˆ‚๐’[1] *= โˆ‚llh - # โˆ‚๐’[2] *= โˆ‚llh - โ„’.rmul!(โˆ‚๐’[1], โˆ‚llh) - โ„’.rmul!(โˆ‚๐’[2], โˆ‚llh) - - โ„’.rmul!(โˆ‚data_in_deviations, โˆ‚llh) - - โˆ‚state[1] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[1] * โˆ‚llh - โˆ‚state[2] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[2] * โˆ‚llh - - # end # timeit_debug - # end # timeit_debug - - return NoTangent(), NoTangent(), โˆ‚state, โˆ‚๐’, โˆ‚data_in_deviations, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # See: https://pcubaborda.net/documents/CGIZ-final.pdf - llh = -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 - - return llh, inversion_filter_loglikelihood_pullback -end - -function rrule(::typeof(calculate_inversion_filter_loglikelihood), - ::Val{:second_order}, - state::Vector{Float64}, - ๐’::Vector{AbstractMatrix{Float64}}, - data_in_deviations::Matrix{Float64}, - observables::Union{Vector{String}, Vector{Symbol}}, - constants::constants, - ws::inversion_workspace{Float64}; - # timer::TimerOutput = TimerOutput(), - on_failure_loglikelihood = -Inf, - warmup_iterations::Int = 0, - presample_periods::Int = 0, - opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)# where S <: Real - T = constants.post_model_macro - # @timeit_debug timer "Inversion filter 2nd - forward" begin - - # @timeit_debug timer "Preallocation" begin - - precision_factor = 1.0 - - n_obs = size(data_in_deviations,2) - - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - - shocksยฒ = 0.0 - logabsdets = 0.0 - - cc = ensure_computational_constants!(constants) - s_in_sโบ = cc.s_in_s - sv_in_sโบ = cc.s_in_sโบ - e_in_sโบ = cc.e_in_sโบ - - tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, e_in_sโบ) |> sparse - shockยฒ_idxs = tmp.nzind - - shockvarยฒ_idxs = setdiff(shock_idxs, shockยฒ_idxs) - - tmp = โ„’.kron(sv_in_sโบ, sv_in_sโบ) |> sparse - var_volยฒ_idxs = tmp.nzind - - tmp = โ„’.kron(s_in_sโบ, s_in_sโบ) |> sparse - varยฒ_idxs = tmp.nzind - - ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] - ๐’โปยนแต‰ = ๐’[1][T.past_not_future_and_mixed_idx,end-T.nExo+1:end] - ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] - ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] - ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] - - ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] - ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] - ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] - ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] - ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] - - ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› - ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป - ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ - ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ - ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ - - kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] - - J = โ„’.I(T.nExo) - - kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) - - kron_buffer3 = โ„’.kron(J, zeros(T.nPast_not_future_and_mixed + 1)) - - x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] - - stateยนโป = state[T.past_not_future_and_mixed_idx] - - stateยนโป_vol = vcat(stateยนโป, 1) - - kronstateยนโป_voltmp = โ„’.kron(stateยนโป_vol, stateยนโป_vol) - - kronstateยนโป_vol = [kronstateยนโป_voltmp for _ in 1:size(data_in_deviations,2)] - - shock_independent = zeros(size(data_in_deviations,1)) - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) - - ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 - - # aug_state_tmp = [zeros(T.nPast_not_future_and_mixed); 1; zeros(T.nExo)] - - aug_state = [[zeros(T.nPast_not_future_and_mixed); 1; zeros(T.nExo)] for _ in 1:size(data_in_deviations,2)] - - kronaug_state = [zeros((T.nPast_not_future_and_mixed + 1 + T.nExo)^2) for _ in 1:size(data_in_deviations,2)] - - tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[1])), x[1]) - - jacc = [zero(tmp) for _ in 1:size(data_in_deviations,2)] - - jacct = copy(tmp') - - ฮป = [zeros(size(tmp, 1)) for _ in 1:size(data_in_deviations,2)] - - ฮป[1] = tmp' \ x[1] * 2 - - fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' - -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] - - kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) - - kronxฮป = [kronxฮป_tmp for _ in 1:size(data_in_deviations,2)] - - tmp = zeros(size(๐’โฑ, 2) * size(๐’โฑ, 2)) - - lI = -2 * vec(โ„’.I(size(๐’โฑ, 2))) - - init_guess = zeros(size(๐’โฑ, 2)) - - # end # timeit_debug - # @timeit_debug timer "Main loop" begin - - @inbounds for i in axes(data_in_deviations,2) - # aug_state[i][1:T.nPast_not_future_and_mixed] = stateยนโป - copyto!(aug_state[i], 1, stateยนโป, 1) - - stateยนโป_vol = aug_state[i][1:T.nPast_not_future_and_mixed + 1] - # copyto!(stateยนโป_vol, 1, aug_state[i], 1, T.nPast_not_future_and_mixed + 1) - - copyto!(shock_independent, data_in_deviations[:,i]) - - โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - - โ„’.kron!(kronstateยนโป_vol[i], stateยนโป_vol, stateยนโป_vol) - - โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kronstateยนโป_vol[i], -1/2, 1) - - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) - โ„’.kron!(kron_buffer3, J, stateยนโป_vol) - - โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer3) - - โ„’.axpy!(1, ๐’ยนแต‰, ๐’โฑ) - - init_guess *= 0 - - # @timeit_debug timer "Find shocks" begin - x[i], matched = find_shocks(Val(filter_algorithm), - init_guess, - kronxx[i], - kron_buffer2, - J, - ๐’โฑ, - ๐’โฑยฒแต‰, - shock_independent, - # max_iter = 100 - ) - # end # timeit_debug - - if !matched - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - โ„’.kron!(kron_buffer2, J, x[i]) - - โ„’.mul!(jacc[i], ๐’โฑยฒแต‰, kron_buffer2) - - โ„’.axpby!(1, ๐’โฑ, 2, jacc[i]) - # jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(length(x[i])), x[i]) - - copy!(jacct, jacc[i]') - - jacc_fact = try - โ„’.factorize(jacct) - catch - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - try - โ„’.ldiv!(ฮป[i], jacc_fact, x[i]) - catch - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # โ„’.ldiv!(ฮป[i], jacc_fact', x[i]) - โ„’.rmul!(ฮป[i], 2) - - # fXฮปp[i] = [reshape(2 * ๐’โฑยฒแต‰' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) jacc[i]' - # -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - โ„’.mul!(tmp, ๐’โฑยฒแต‰', ฮป[i]) - โ„’.axpby!(1, lI, 2, tmp) - - fXฮปp[i][1:size(๐’โฑ, 2), 1:size(๐’โฑ, 2)] = tmp - fXฮปp[i][size(๐’โฑ, 2)+1:end, 1:size(๐’โฑ, 2)] = -jacc[i] - fXฮปp[i][1:size(๐’โฑ, 2), size(๐’โฑ, 2)+1:end] = jacct - - โ„’.kron!(kronxx[i], x[i], x[i]) - - โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) - - if i > presample_periods - # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) - logabsdets += โ„’.logabsdet(jacc_fact)[1] - else - logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) - end - - shocksยฒ += sum(abs2, x[i]) - - if !isfinite(logabsdets) || !isfinite(shocksยฒ) - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - end - - # aug_state[i] = [stateยนโป; 1; x[i]] - # aug_state[i][1:T.nPast_not_future_and_mixed] = stateยนโป - # aug_state[i][end-T.nExo+1:end] = x[i] - copyto!(aug_state[i], 1, stateยนโป, 1) - copyto!(aug_state[i], length(stateยนโป) + 2, x[i], 1) - - โ„’.kron!(kronaug_state[i], aug_state[i], aug_state[i]) - โ„’.mul!(stateยนโป, ๐’โปยน, aug_state[i]) - โ„’.mul!(stateยนโป, ๐’โปยฒ, kronaug_state[i], 1/2 ,1) - end - - # end # timeit_debug - # end # timeit_debug - - โˆ‚aug_state = zero(aug_state[1]) - - โˆ‚kronaug_state = zero(kronaug_state[1]) - - โˆ‚kronstateยนโป_vol = zero(kronstateยนโป_vol[1]) - - โˆ‚state = similar(state) - - โˆ‚๐’ = copy(๐’) - - โˆ‚data_in_deviations = similar(data_in_deviations) - - โˆ‚kronIx = zero(โ„’.kron(โ„’.I(length(x[1])), x[1])) - - function inversion_filter_loglikelihood_pullback(โˆ‚llh) - # @timeit_debug timer "Inversion filter 2nd - pullback" begin - - # @timeit_debug timer "Preallocation" begin - - โˆ‚๐’โฑ = zero(๐’โฑ) - โˆ‚๐’โฑยฒแต‰ = zero(๐’โฑยฒแต‰) - - # Allocate or reuse workspaces for pullback temps - if size(ws.โˆ‚๐’โฑยฒแต‰tmp) != (T.nExo, T.nExo * length(ฮป[1])) - ws.โˆ‚๐’โฑยฒแต‰tmp = zeros(T.nExo, T.nExo * length(ฮป[1])) - else - fill!(ws.โˆ‚๐’โฑยฒแต‰tmp, zero(eltype(ws.โˆ‚๐’โฑยฒแต‰tmp))) - end - โˆ‚๐’โฑยฒแต‰tmp = ws.โˆ‚๐’โฑยฒแต‰tmp - - if size(ws.โˆ‚๐’โฑยฒแต‰tmp2) != (length(ฮป[1]), T.nExo * T.nExo) - ws.โˆ‚๐’โฑยฒแต‰tmp2 = zeros(length(ฮป[1]), T.nExo * T.nExo) - else - fill!(ws.โˆ‚๐’โฑยฒแต‰tmp2, zero(eltype(ws.โˆ‚๐’โฑยฒแต‰tmp2))) - end - โˆ‚๐’โฑยฒแต‰tmp2 = ws.โˆ‚๐’โฑยฒแต‰tmp2 - - โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) - โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) - - โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) - โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) - - โˆ‚๐’โปยน = zero(๐’โปยน) - โˆ‚๐’โปยฒ = zero(๐’โปยฒ) - - โˆ‚stateยนโป_vol = zero(stateยนโป_vol) - # โˆ‚x = zero(x[1]) - โˆ‚state = zeros(T.nPast_not_future_and_mixed) - - โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ - - # Allocate or reuse workspaces for kron products - if length(ws.kronSฮป) != length(cond_var_idx) * T.nExo - ws.kronSฮป = zeros(length(cond_var_idx) * T.nExo) - else - fill!(ws.kronSฮป, zero(eltype(ws.kronSฮป))) - end - kronSฮป = ws.kronSฮป - - if length(ws.kronxS) != T.nExo * length(cond_var_idx) - ws.kronxS = zeros(T.nExo * length(cond_var_idx)) - else - fill!(ws.kronxS, zero(eltype(ws.kronxS))) - end - kronxS = ws.kronxS - - # end # timeit_debug - # @timeit_debug timer "Main loop" begin - - for i in reverse(axes(data_in_deviations,2)) - # stt = ๐’โปยน * aug_state + ๐’โปยฒ * โ„’.kron(aug_state, aug_state) / 2 - # โˆ‚๐’โปยน += โˆ‚state * aug_state[i]' - โ„’.mul!(โˆ‚๐’โปยน, โˆ‚state, aug_state[i]', 1, 1) - - # โˆ‚๐’โปยฒ += โˆ‚state * kronaug_state[i]' / 2 - โ„’.mul!(โˆ‚๐’โปยฒ, โˆ‚state, kronaug_state[i]', 1/2, 1) - - โ„’.mul!(โˆ‚aug_state, ๐’โปยน', โˆ‚state) - # โˆ‚aug_state = ๐’โปยน' * โˆ‚state - - โ„’.mul!(โˆ‚kronaug_state, ๐’โปยฒ', โˆ‚state) - โ„’.rdiv!(โˆ‚kronaug_state, 2) - # โˆ‚kronaug_state = ๐’โปยฒ' * โˆ‚state / 2 - - fill_kron_adjoint!(โˆ‚aug_state, โˆ‚aug_state, โˆ‚kronaug_state, aug_state[i], aug_state[i]) - - if i > 1 && i < size(data_in_deviations,2) - โˆ‚state *= 0 - end - - # aug_state[i] = [stt; 1; x[i]] - โˆ‚state += โˆ‚aug_state[1:length(โˆ‚state)] - - # aug_state[i] = [stt; 1; x[i]] - โˆ‚x = โˆ‚aug_state[T.nPast_not_future_and_mixed+2:end] - - # shocksยฒ += sum(abs2,x[i]) - if i < size(data_in_deviations,2) - โˆ‚x -= copy(x[i]) - else - โˆ‚x += copy(x[i]) - end - - # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] - โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) - inv(jacc[i])' - else - โ„’.pinv(jacc[i])' - end - catch - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x[1]) - โ„’.mul!(โˆ‚kronIx, ๐’โฑยฒแต‰', โˆ‚jacc) - - if i < size(data_in_deviations,2) - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -J) - else - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, J) - end - - # โˆ‚๐’โฑยฒแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' - โ„’.kron!(kron_buffer2, J, x[i]) - - โ„’.mul!(โˆ‚๐’โฑยฒแต‰, โˆ‚jacc, kron_buffer2', -1, 1) - - # find_shocks - โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) - - S = fXฮปp[i]' \ โˆ‚xฮป - - if i < size(data_in_deviations,2) - S *= -1 - end - - โˆ‚shock_independent = S[T.nExo+1:end] # fine - - # โ„’.mul!(โˆ‚๐’โฑ, ฮป[i], S[1:T.nExo]') - # โ„’.mul!(โˆ‚๐’โฑ, S[T.nExo+1:end], x[i]', -1, 1) # fine - # โ„’.axpy!(-1/2, โˆ‚jacc, โˆ‚๐’โฑ) - # โˆ‚๐’โฑ = ฮป[i] * S[1:T.nExo]' - S[T.nExo+1:end] * x[i]' # fine - - # copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) - # โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine - โ„’.kron!(kronSฮป, S[1:T.nExo], ฮป[i]) - โ„’.kron!(kronxS, x[i], S[T.nExo+1:end]) - โ„’.axpy!(-1, kronxS, kronSฮป) - copyto!(โˆ‚๐’โฑ, kronSฮป) - - โ„’.axpy!(-1/2, โˆ‚jacc, โˆ‚๐’โฑ) - - โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], kronxฮป[i]) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) - # โ„’.mul!(โˆ‚๐’โฑยฒแต‰tmp, S[1:T.nExo], kronxฮป[i]', 2, 1) - # โ„’.mul!(โˆ‚๐’โฑยฒแต‰tmp2, S[T.nExo+1:end], kronxx[i]', -1, 1) - - # โ„’.mul!(โˆ‚๐’โฑยฒแต‰, S[1:T.nExo], kronxฮป[i]', 2, 1) - # โ„’.mul!(โˆ‚๐’โฑยฒแต‰, S[T.nExo+1:end], kronxx[i]', -1, 1) - # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo+1:end] * kronxx[i]' - - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) - โˆ‚stateยนโป_vol *= 0 - - โ„’.mul!(โˆ‚kronIstateยนโป_vol, ๐’ยฒโปแต‰', โˆ‚๐’โฑ) - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, J) - - stateยนโป_vol = aug_state[i][1:T.nPast_not_future_and_mixed + 1] - - โ„’.axpy!(1, โˆ‚๐’โฑ, โˆ‚๐’ยนแต‰) - # โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ - - โ„’.kron!(kron_buffer3, J, stateยนโป_vol) - - โ„’.mul!(โˆ‚๐’ยฒโปแต‰, โˆ‚๐’โฑ, kron_buffer3', 1, 1) - # โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' - - # shock_independent = copy(data_in_deviations[:,i]) - โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - # โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' - โ„’.mul!(โˆ‚๐’ยนโปแต›, โˆ‚shock_independent, stateยนโป_vol', -1 ,1) - - # โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent - โ„’.mul!(โˆ‚stateยนโป_vol, ๐’ยนโปแต›', โˆ‚shock_independent, -1, 1) - - # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) - โ„’.kron!(kronstateยนโป_vol[i], stateยนโป_vol, stateยนโป_vol) - โ„’.mul!(โˆ‚๐’ยฒโปแต›, โˆ‚shock_independent, kronstateยนโป_vol[i]', -1/2, 1) - # โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 - - โ„’.mul!(โˆ‚kronstateยนโป_vol, ๐’ยฒโปแต›', โˆ‚shock_independent) - โ„’.rdiv!(โˆ‚kronstateยนโป_vol, -2) - # โˆ‚kronstateยนโป_vol = ๐’ยฒโปแต›' * โˆ‚shock_independent / (-2) - - fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) - - # stateยนโป_vol = vcat(stateยนโป, 1) - โˆ‚state += โˆ‚stateยนโป_vol[1:end-1] - end - - # end # timeit_debug - # @timeit_debug timer "Post allocation" begin - - โˆ‚๐’ = [copy(๐’[1]) * 0, copy(๐’[2]) * 0] - - โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] += โˆ‚๐’ยนแต‰ - โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] += โˆ‚๐’ยฒโปแต‰ - โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] += โˆ‚๐’โฑยฒแต‰ / 2 - โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] += โˆ‚๐’ยนโปแต› - โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] += โˆ‚๐’ยฒโปแต› - - โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยน - โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยฒ - - โˆ‚๐’[1] *= โˆ‚llh - โˆ‚๐’[2] *= โˆ‚llh - - return NoTangent(), NoTangent(), โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state * โˆ‚llh, โˆ‚๐’, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # end # timeit_debug - # end # timeit_debug - - # See: https://pcubaborda.net/documents/CGIZ-final.pdf - llh = -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 - - return llh, inversion_filter_loglikelihood_pullback -end - -function rrule(::typeof(calculate_inversion_filter_loglikelihood), - ::Val{:pruned_third_order}, - state::Vector{Vector{Float64}}, - ๐’::Vector{AbstractMatrix{Float64}}, - data_in_deviations::Matrix{Float64}, - observables::Union{Vector{String}, Vector{Symbol}}, - constants::constants, - ws::inversion_workspace{Float64}; - # timer::TimerOutput = TimerOutput(), - on_failure_loglikelihood = -Inf, - warmup_iterations::Int = 0, - presample_periods::Int = 0, - opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton) - T = constants.post_model_macro - # @timeit_debug timer "Inversion filter - forward" begin - precision_factor = 1.0 - - n_obs = size(data_in_deviations,2) - - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - - shocksยฒ = 0.0 - logabsdets = 0.0 - - cc = ensure_computational_constants!(constants) - s_in_sโบ = cc.s_in_s - sv_in_sโบ = cc.s_in_sโบ - e_in_sโบ = cc.e_in_sโบ - - tmp = โ„’.kron(e_in_sโบ, s_in_sโบ) |> sparse - shockvar_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs = tmp.nzind - - tmp = โ„’.kron(zero(e_in_sโบ) .+ 1, e_in_sโบ) |> sparse - shock_idxs2 = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, e_in_sโบ) |> sparse - shockยฒ_idxs = tmp.nzind - - shockvarยฒ_idxs = setdiff(union(shock_idxs), shockยฒ_idxs) - - tmp = โ„’.kron(sv_in_sโบ, sv_in_sโบ) |> sparse - var_volยฒ_idxs = tmp.nzind - - tmp = โ„’.kron(s_in_sโบ, s_in_sโบ) |> sparse - varยฒ_idxs = tmp.nzind - - ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] - ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] - ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] - ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] - - ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] - ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] - ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] - ๐’ยฒโปแต›แต‰ = ๐’[2][cond_var_idx,shockvar_idxs] - ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] - ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] - - ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› - ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป - ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ - ๐’ยฒโปแต›แต‰ = nnz(๐’ยฒโปแต›แต‰) / length(๐’ยฒโปแต›แต‰) > .1 ? collect(๐’ยฒโปแต›แต‰) : ๐’ยฒโปแต›แต‰ - ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ - ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ - - tmp = โ„’.kron(sv_in_sโบ, โ„’.kron(sv_in_sโบ, sv_in_sโบ)) |> sparse - var_volยณ_idxs = tmp.nzind - - tmp = โ„’.kron(โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1), zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs2 = tmp.nzind - - tmp = โ„’.kron(โ„’.kron(e_in_sโบ, e_in_sโบ), zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs3 = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, โ„’.kron(e_in_sโบ, e_in_sโบ)) |> sparse - shockยณ_idxs = tmp.nzind - - tmp = โ„’.kron(zero(e_in_sโบ) .+ 1, โ„’.kron(e_in_sโบ, e_in_sโบ)) |> sparse - shockvar1_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, โ„’.kron(zero(e_in_sโบ) .+ 1, e_in_sโบ)) |> sparse - shockvar2_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1)) |> sparse - shockvar3_idxs = tmp.nzind - - shockvarยณ2_idxs = setdiff(shock_idxs2, shockยณ_idxs, shockvar1_idxs, shockvar2_idxs, shockvar3_idxs) - - shockvarยณ_idxs = setdiff(shock_idxs3, shockยณ_idxs)#, shockvar1_idxs, shockvar2_idxs, shockvar3_idxs) - - ๐’ยณโปแต› = ๐’[3][cond_var_idx,var_volยณ_idxs] - ๐’ยณโปแต‰ยฒ = ๐’[3][cond_var_idx,shockvarยณ2_idxs] - ๐’ยณโปแต‰ = ๐’[3][cond_var_idx,shockvarยณ_idxs] - ๐’ยณแต‰ = ๐’[3][cond_var_idx,shockยณ_idxs] - ๐’โปยณ = ๐’[3][T.past_not_future_and_mixed_idx,:] - - ๐’ยณโปแต› = nnz(๐’ยณโปแต›) / length(๐’ยณโปแต›) > .1 ? collect(๐’ยณโปแต›) : ๐’ยณโปแต› - ๐’ยณโปแต‰ = nnz(๐’ยณโปแต‰) / length(๐’ยณโปแต‰) > .1 ? collect(๐’ยณโปแต‰) : ๐’ยณโปแต‰ - ๐’ยณแต‰ = nnz(๐’ยณแต‰) / length(๐’ยณแต‰) > .1 ? collect(๐’ยณแต‰) : ๐’ยณแต‰ - ๐’โปยณ = nnz(๐’โปยณ) / length(๐’โปยณ) > .1 ? collect(๐’โปยณ) : ๐’โปยณ - - stateโ‚ = state[1][T.past_not_future_and_mixed_idx] - stateโ‚‚ = state[2][T.past_not_future_and_mixed_idx] - stateโ‚ƒ = state[3][T.past_not_future_and_mixed_idx] - - kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] - - J = โ„’.I(T.nExo) - - II = sparse(โ„’.I(T.nExo^2)) - - kronxxx = [zeros(T.nExo^3) for _ in 1:size(data_in_deviations,2)] - - kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) - - kron_buffer3 = โ„’.kron(J, zeros(T.nExo^2)) - - kron_buffer4 = โ„’.kron(โ„’.kron(J, J), zeros(T.nExo)) - - x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] - - stateยนโป = stateโ‚ - - stateยนโป_vol = vcat(stateยนโป, 1) - - stateยฒโป = stateโ‚‚#[T.past_not_future_and_mixed_idx] - - stateยณโป = stateโ‚ƒ#[T.past_not_future_and_mixed_idx] - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) - - ๐’โฑยฒแต‰ = [zero(๐’ยฒแต‰) for _ in 1:size(data_in_deviations,2)] - - aug_stateโ‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] - aug_stateโ‚ฬ‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] - aug_stateโ‚‚ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] - aug_stateโ‚ƒ = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] - - kron_aug_stateโ‚ = [zeros(size(๐’โปยน,2)^2) for _ in 1:size(data_in_deviations,2)] - - jacc_tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰[1] * โ„’.kron(โ„’.I(T.nExo), x[1]) - - jacc = [zero(jacc_tmp) for _ in 1:size(data_in_deviations,2)] - - ฮป = [zeros(size(jacc_tmp, 1)) for _ in 1:size(data_in_deviations,2)] - - ฮป[1] = jacc_tmp' \ x[1] * 2 - - fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰[1]' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) jacc_tmp' - -jacc_tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] - - kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) - - kronxฮป = [kronxฮป_tmp for _ in 1:size(data_in_deviations,2)] - - kronxxฮป_tmp = โ„’.kron(x[1], kronxฮป_tmp) - - kronxxฮป = [kronxxฮป_tmp for _ in 1:size(data_in_deviations,2)] - - II = sparse(โ„’.I(T.nExo^2)) - - lI = 2 * โ„’.I(size(๐’โฑ, 2)) - - ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - - # @timeit_debug timer "Loop" begin - for i in axes(data_in_deviations,2) - stateยนโป = stateโ‚ - - stateยนโป_vol = vcat(stateยนโป, 1) - - stateยฒโป = stateโ‚‚#[T.past_not_future_and_mixed_idx] - - stateยณโป = stateโ‚ƒ#[T.past_not_future_and_mixed_idx] - - shock_independent = copy(data_in_deviations[:,i]) - - โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - - โ„’.mul!(shock_independent, ๐’ยนโป, stateยฒโป, -1, 1) - - โ„’.mul!(shock_independent, ๐’ยนโป, stateยณโป, -1, 1) - - โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) - - โ„’.mul!(shock_independent, ๐’ยฒโป, โ„’.kron(stateยนโป, stateยฒโป), -1, 1) - - โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยฒโปแต›แต‰ * โ„’.kron(โ„’.I(T.nExo), stateยฒโป) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 - - ๐’โฑยฒแต‰[i] = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 - - ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - - init_guess = zeros(size(๐’โฑ, 2)) - - # @timeit_debug timer "Find shocks" begin - x[i], matched = find_shocks(Val(filter_algorithm), - init_guess, - kronxx[i], - kronxxx[i], - kron_buffer2, - kron_buffer3, - kron_buffer4, - J, - ๐’โฑ, - ๐’โฑยฒแต‰[i], - ๐’โฑยณแต‰, - shock_independent, - # max_iter = 100 - ) - # end # timeit_debug - - if !matched - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰[i] * โ„’.kron(โ„’.I(T.nExo), x[i]) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), kronxx[i]) - - ฮป[i] = jacc[i]' \ x[i] * 2 - # โ„’.ldiv!(ฮป[i], tmp', x[i]) - # โ„’.rmul!(ฮป[i], 2) - fXฮปp[i] = [reshape((2 * ๐’โฑยฒแต‰[i] + 6 * ๐’โฑยณแต‰ * โ„’.kron(II, x[i]))' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - lI jacc[i]' - -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - โ„’.kron!(kronxx[i], x[i], x[i]) - - โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) - - โ„’.kron!(kronxxฮป[i], x[i], kronxฮป[i]) - - โ„’.kron!(kronxxx[i], x[i], kronxx[i]) - - if i > presample_periods - # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) - logabsdets += โ„’.logabsdet(jacc[i])[1] - else - logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) - end - - shocksยฒ += sum(abs2,x[i]) - - if !isfinite(logabsdets) || !isfinite(shocksยฒ) - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - end - - aug_stateโ‚[i] = [stateโ‚; 1; x[i]] - aug_stateโ‚ฬ‚[i] = [stateโ‚; 0; x[i]] - aug_stateโ‚‚[i] = [stateโ‚‚; 0; zeros(T.nExo)] - aug_stateโ‚ƒ[i] = [stateโ‚ƒ; 0; zeros(T.nExo)] - - kron_aug_stateโ‚[i] = โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) - - stateโ‚, stateโ‚‚, stateโ‚ƒ = [๐’โปยน * aug_stateโ‚[i], ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * kron_aug_stateโ‚[i] / 2, ๐’โปยน * aug_stateโ‚ƒ[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i]) + ๐’โปยณ * โ„’.kron(kron_aug_stateโ‚[i], aug_stateโ‚[i]) / 6] - end - # end # timeit_debug - - # See: https://pcubaborda.net/documents/CGIZ-final.pdf - llh = -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 - - โˆ‚state = similar(state) - - โˆ‚๐’ = copy(๐’) - - โˆ‚data_in_deviations = similar(data_in_deviations) - - # end # timeit_debug - - function inversion_filter_loglikelihood_pullback(โˆ‚llh) - # @timeit_debug timer "Inversion filter - pullback" begin - โˆ‚๐’โฑ = zero(๐’โฑ) - โˆ‚๐’ยฒแต‰ = zero(๐’ยฒแต‰) - โˆ‚๐’โฑยณแต‰ = zero(๐’โฑยณแต‰) - - โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) - โˆ‚๐’ยนโป = zero(๐’ยนโป) - โˆ‚๐’ยฒโป = zero(๐’ยฒโป) - โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) - โˆ‚๐’ยฒโปแต›แต‰ = zero(๐’ยฒโปแต›แต‰) - โˆ‚๐’ยณโปแต‰ = zero(๐’ยณโปแต‰) - โˆ‚๐’ยณโปแต‰ยฒ = zero(๐’ยณโปแต‰ยฒ) - - โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) - โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) - โˆ‚๐’ยณโปแต› = zero(๐’ยณโปแต›) - - โˆ‚๐’โปยน = zero(๐’โปยน) - โˆ‚๐’โปยฒ = zero(๐’โปยฒ) - โˆ‚๐’โปยณ = zero(๐’โปยณ) - - โˆ‚aug_stateโ‚ฬ‚ = zero(aug_stateโ‚ฬ‚[1]) - โˆ‚stateยนโป_vol = zero(stateยนโป_vol) - โˆ‚x = zero(x[1]) - โˆ‚kronxx = zero(kronxx[1]) - โˆ‚kronstateยนโป_vol = zeros(length(stateยนโป_vol)^2) - โˆ‚state = [zeros(T.nPast_not_future_and_mixed), zeros(T.nPast_not_future_and_mixed), zeros(T.nPast_not_future_and_mixed)] - - # @timeit_debug timer "Loop" begin - for i in reverse(axes(data_in_deviations,2)) - # stateโ‚ = ๐’โปยน * aug_stateโ‚[i] - โˆ‚๐’โปยน += โˆ‚state[1] * aug_stateโ‚[i]' - - โˆ‚aug_stateโ‚ = ๐’โปยน' * โˆ‚state[1] - - # stateโ‚‚ = ๐’โปยน * aug_stateโ‚‚[i] + ๐’โปยฒ * kron_aug_stateโ‚[i] / 2 - โˆ‚๐’โปยน += โˆ‚state[2] * aug_stateโ‚‚[i]' - - โˆ‚aug_stateโ‚‚ = ๐’โปยน' * โˆ‚state[2] - - โˆ‚๐’โปยฒ += โˆ‚state[2] * kron_aug_stateโ‚[i]' / 2 - - โˆ‚kronaug_stateโ‚ = ๐’โปยฒ' * โˆ‚state[2] / 2 - - # stateโ‚ƒ = ๐’โปยน * aug_stateโ‚ƒ[i] + ๐’โปยฒ * โ„’.kron(aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i]) + ๐’โปยณ * โ„’.kron(kron_aug_stateโ‚[i],aug_stateโ‚[i]) / 6 - โˆ‚๐’โปยน += โˆ‚state[3] * aug_stateโ‚ƒ[i]' - - โˆ‚aug_stateโ‚ƒ = ๐’โปยน' * โˆ‚state[3] - - โˆ‚๐’โปยฒ += โˆ‚state[3] * โ„’.kron(aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i])' - - โˆ‚aug_stateโ‚ฬ‚ *= 0 - - โˆ‚kronaug_stateโ‚ฬ‚โ‚‚ = ๐’โปยฒ' * โˆ‚state[3] - - fill_kron_adjoint!(โˆ‚aug_stateโ‚ฬ‚, โˆ‚aug_stateโ‚‚, โˆ‚kronaug_stateโ‚ฬ‚โ‚‚, aug_stateโ‚ฬ‚[i], aug_stateโ‚‚[i]) - - โˆ‚๐’โปยณ += โˆ‚state[3] * โ„’.kron(kron_aug_stateโ‚[i],aug_stateโ‚[i])' / 6 - - โˆ‚kronkronaug_stateโ‚ = ๐’โปยณ' * โˆ‚state[3] / 6 - - fill_kron_adjoint!(โˆ‚aug_stateโ‚, โˆ‚kronaug_stateโ‚, โˆ‚kronkronaug_stateโ‚, aug_stateโ‚[i], kron_aug_stateโ‚[i]) - - # kron_aug_stateโ‚[i] = โ„’.kron(aug_stateโ‚[i], aug_stateโ‚[i]) - fill_kron_adjoint!(โˆ‚aug_stateโ‚, โˆ‚aug_stateโ‚, โˆ‚kronaug_stateโ‚, aug_stateโ‚[i], aug_stateโ‚[i]) - - if i > 1 && i < size(data_in_deviations,2) - โˆ‚state[1] *= 0 - โˆ‚state[2] *= 0 - โˆ‚state[3] *= 0 - end - - # aug_stateโ‚[i] = [stateโ‚; 1; x[i]] - โˆ‚state[1] += โˆ‚aug_stateโ‚[1:length(โˆ‚state[1])] - - โˆ‚x = โˆ‚aug_stateโ‚[T.nPast_not_future_and_mixed+2:end] - - # aug_stateโ‚ฬ‚[i] = [stateโ‚; 0; x[i]] - โˆ‚state[1] += โˆ‚aug_stateโ‚ฬ‚[1:length(โˆ‚state[1])] - - โˆ‚x += โˆ‚aug_stateโ‚ฬ‚[T.nPast_not_future_and_mixed+2:end] - - # aug_stateโ‚‚[i] = [stateโ‚‚; 0; zeros(T.nExo)] - โˆ‚state[2] += โˆ‚aug_stateโ‚‚[1:length(โˆ‚state[1])] - - # aug_stateโ‚ƒ[i] = [stateโ‚ƒ; 0; zeros(T.nExo)] - โˆ‚state[3] += โˆ‚aug_stateโ‚ƒ[1:length(โˆ‚state[1])] - - # shocksยฒ += sum(abs2,x[i]) - if i < size(data_in_deviations,2) - โˆ‚x -= copy(x[i]) - else - โˆ‚x += copy(x[i]) - end - - # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] - โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) - inv(jacc[i])' - else - โ„’.pinv(jacc[i])' - end - catch - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), โ„’.kron(x, x)) - # โˆ‚๐’โฑ = -โˆ‚jacc / 2 # fine - - โˆ‚kronIx = ๐’โฑยฒแต‰[i]' * โˆ‚jacc - - if i < size(data_in_deviations,2) - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -โ„’.I(T.nExo)) - else - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, โ„’.I(T.nExo)) - end - - โˆ‚๐’โฑยฒแต‰ = -โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' - - โˆ‚kronIxx = ๐’โฑยณแต‰' * โˆ‚jacc * 3 / 2 - - โˆ‚kronxx *= 0 - - if i < size(data_in_deviations,2) - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, -โ„’.I(T.nExo)) - else - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, โ„’.I(T.nExo)) - end - - fill_kron_adjoint!(โˆ‚x, โˆ‚x, โˆ‚kronxx, x[i], x[i]) - - โˆ‚๐’โฑยณแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), kronxx[i])' * 3 / 2 - - # find_shocks - โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) - - S = fXฮปp[i]' \ โˆ‚xฮป - - if i < size(data_in_deviations,2) - S *= -1 - end - - โˆ‚shock_independent = S[T.nExo+1:end] # fine - - # โˆ‚๐’โฑ += S[1:T.nExo] * ฮป[i]' - S[T.nExo + 1:end] * x[i]' # fine - copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) - โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine - - โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], โ„’.kron(x[i], ฮป[i])) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) - # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo + 1:end] * kronxx[i]' - - โˆ‚๐’โฑยณแต‰ += reshape(3 * โ„’.kron(S[1:T.nExo], โ„’.kron(โ„’.kron(x[i], x[i]), ฮป[i])) - โ„’.kron(kronxxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยณแต‰)) - # โˆ‚๐’โฑยณแต‰ += 3 * S[1:T.nExo] * kronxxฮป[i]' - S[T.nExo + 1:end] * kronxxx[i]' - - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยฒโปแต›แต‰ * โ„’.kron(โ„’.I(T.nExo), stateยฒโป) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 - โˆ‚kronstateยนโป_vol *= 0 - - stateยนโป_vol = [aug_stateโ‚[i][1:T.nPast_not_future_and_mixed];1] # define here as it is used multiple times later - stateยนโป = aug_stateโ‚[i][1:T.nPast_not_future_and_mixed] - stateยฒโป = aug_stateโ‚‚[i][1:T.nPast_not_future_and_mixed] - stateยณโป = aug_stateโ‚ƒ[i][1:T.nPast_not_future_and_mixed] - - โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ - - โˆ‚stateยนโป_vol *= 0 - - โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, โ„’.I(T.nExo)) - - โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' - - โˆ‚kronIstateยฒโป = ๐’ยฒโปแต›แต‰' * โˆ‚๐’โฑ - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยฒโป, โˆ‚state[2], โ„’.I(T.nExo)) - - โˆ‚๐’ยฒโปแต›แต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยฒโป)' - - โˆ‚kronIstateยนโป_volstateยนโป_vol = ๐’ยณโปแต‰ยฒ' * โˆ‚๐’โฑ / 2 - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_volstateยนโป_vol, โˆ‚kronstateยนโป_vol, โ„’.I(T.nExo)) - - โˆ‚๐’ยณโปแต‰ยฒ += โˆ‚๐’โฑ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol)' / 2 - - # ๐’โฑยฒแต‰[i] = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 - โˆ‚๐’ยฒแต‰ += โˆ‚๐’โฑยฒแต‰ / 2 - - โˆ‚๐’ยณโปแต‰ += โˆ‚๐’โฑยฒแต‰ * โ„’.kron(II, stateยนโป_vol)' / 2 - - โˆ‚kronIIstateยนโป_vol = ๐’ยณโปแต‰' * โˆ‚๐’โฑยฒแต‰ / 2 - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIIstateยนโป_vol, โˆ‚stateยนโป_vol, II) - - # shock_independent = copy(data_in_deviations[:,i]) - โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' - - โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยนโป, stateยฒโป, -1, 1) - โˆ‚๐’ยนโป -= โˆ‚shock_independent * stateยฒโป' - - โˆ‚state[2] -= ๐’ยนโป' * โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยนโป, stateยณโป, -1, 1) - โˆ‚๐’ยนโป -= โˆ‚shock_independent * stateยณโป' - - โˆ‚state[3] -= ๐’ยนโป' * โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) - โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 - - โˆ‚kronstateยนโป_vol -= ๐’ยฒโปแต›' * โˆ‚shock_independent / 2 - - # โ„’.mul!(shock_independent, ๐’ยฒโป, โ„’.kron(stateยนโป, stateยฒโป), -1, 1) - โˆ‚๐’ยฒโป -= โˆ‚shock_independent * โ„’.kron(stateยนโป, stateยฒโป)' - - โˆ‚kronstateยนโปยฒโป = -๐’ยฒโป' * โˆ‚shock_independent - - fill_kron_adjoint!(โˆ‚state[1], โˆ‚state[2], โˆ‚kronstateยนโปยฒโป, stateยนโป, stateยฒโป) - - # โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) - โˆ‚๐’ยณโปแต› -= โˆ‚shock_independent * โ„’.kron(โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol)' / 6 - - โˆ‚kronstateยนโป_volstateยนโป_vol = -๐’ยณโปแต›' * โˆ‚shock_independent / 6 - - fill_kron_adjoint!(โˆ‚kronstateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_volstateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol) - - fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) - - # stateยนโป_vol = vcat(stateยนโป, 1) - โˆ‚state[1] += โˆ‚stateยนโป_vol[1:end-1] - end - # end # timeit_debug - - โˆ‚๐’ = [copy(๐’[1]) * 0, copy(๐’[2]) * 0, copy(๐’[3]) * 0] - - โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] += โˆ‚๐’ยนแต‰ - โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] += โˆ‚๐’ยนโป - โˆ‚๐’[2][cond_var_idx,varยฒ_idxs] += โˆ‚๐’ยฒโป - โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] += โˆ‚๐’ยฒโปแต‰ - โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] += โˆ‚๐’ยฒแต‰ - โˆ‚๐’[2][cond_var_idx,shockvar_idxs] += โˆ‚๐’ยฒโปแต›แต‰ - โˆ‚๐’[3][cond_var_idx,shockvarยณ2_idxs] += โˆ‚๐’ยณโปแต‰ยฒ - โˆ‚๐’[3][cond_var_idx,shockvarยณ_idxs] += โˆ‚๐’ยณโปแต‰ - โˆ‚๐’[3][cond_var_idx,shockยณ_idxs] += โˆ‚๐’โฑยณแต‰ / 6 # ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - - โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] += โˆ‚๐’ยนโปแต› - โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] += โˆ‚๐’ยฒโปแต› - โˆ‚๐’[3][cond_var_idx,var_volยณ_idxs] += โˆ‚๐’ยณโปแต› - - โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยน - โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยฒ - โˆ‚๐’[3][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยณ - - โˆ‚๐’[1] *= โˆ‚llh - โˆ‚๐’[2] *= โˆ‚llh - โˆ‚๐’[3] *= โˆ‚llh - - โˆ‚state[1] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[1] * โˆ‚llh - โˆ‚state[2] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[2] * โˆ‚llh - โˆ‚state[3] = โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state[3] * โˆ‚llh - - # end # timeit_debug - - return NoTangent(), NoTangent(), โˆ‚state, โˆ‚๐’, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - return llh, inversion_filter_loglikelihood_pullback -end - -function rrule(::typeof(calculate_inversion_filter_loglikelihood), - ::Val{:third_order}, - state::Vector{Float64}, - ๐’::Vector{AbstractMatrix{Float64}}, - data_in_deviations::Matrix{Float64}, - observables::Union{Vector{String}, Vector{Symbol}}, - constants::constants, - ws::inversion_workspace{Float64}; - # timer::TimerOutput = TimerOutput(), - on_failure_loglikelihood = -Inf, - warmup_iterations::Int = 0, - presample_periods::Int = 0, - opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton) - T = constants.post_model_macro - # @timeit_debug timer "Inversion filter pruned 2nd - forward" begin - # @timeit_debug timer "Preallocation" begin - - precision_factor = 1.0 - - n_obs = size(data_in_deviations,2) - - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - - shocksยฒ = 0.0 - logabsdets = 0.0 - - cc = ensure_computational_constants!(constants) - s_in_sโบ = cc.s_in_s - sv_in_sโบ = cc.s_in_sโบ - e_in_sโบ = cc.e_in_sโบ - - tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, e_in_sโบ) |> sparse - shockยฒ_idxs = tmp.nzind - - shockvarยฒ_idxs = setdiff(shock_idxs, shockยฒ_idxs) - - tmp = โ„’.kron(sv_in_sโบ, sv_in_sโบ) |> sparse - var_volยฒ_idxs = tmp.nzind - - tmp = โ„’.kron(s_in_sโบ, s_in_sโบ) |> sparse - varยฒ_idxs = tmp.nzind - - ๐’โปยน = ๐’[1][T.past_not_future_and_mixed_idx,:] - ๐’โปยนแต‰ = ๐’[1][T.past_not_future_and_mixed_idx,end-T.nExo+1:end] - ๐’ยนโป = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] - ๐’ยนโปแต› = ๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] - ๐’ยนแต‰ = ๐’[1][cond_var_idx,end-T.nExo+1:end] - - ๐’ยฒโปแต› = ๐’[2][cond_var_idx,var_volยฒ_idxs] - ๐’ยฒโป = ๐’[2][cond_var_idx,varยฒ_idxs] - ๐’ยฒโปแต‰ = ๐’[2][cond_var_idx,shockvarยฒ_idxs] - ๐’ยฒแต‰ = ๐’[2][cond_var_idx,shockยฒ_idxs] - ๐’โปยฒ = ๐’[2][T.past_not_future_and_mixed_idx,:] - - ๐’ยฒโปแต› = nnz(๐’ยฒโปแต›) / length(๐’ยฒโปแต›) > .1 ? collect(๐’ยฒโปแต›) : ๐’ยฒโปแต› - ๐’ยฒโป = nnz(๐’ยฒโป) / length(๐’ยฒโป) > .1 ? collect(๐’ยฒโป) : ๐’ยฒโป - ๐’ยฒโปแต‰ = nnz(๐’ยฒโปแต‰) / length(๐’ยฒโปแต‰) > .1 ? collect(๐’ยฒโปแต‰) : ๐’ยฒโปแต‰ - ๐’ยฒแต‰ = nnz(๐’ยฒแต‰) / length(๐’ยฒแต‰) > .1 ? collect(๐’ยฒแต‰) : ๐’ยฒแต‰ - ๐’โปยฒ = nnz(๐’โปยฒ) / length(๐’โปยฒ) > .1 ? collect(๐’โปยฒ) : ๐’โปยฒ - - tmp = โ„’.kron(sv_in_sโบ, โ„’.kron(sv_in_sโบ, sv_in_sโบ)) |> sparse - var_volยณ_idxs = tmp.nzind - - tmp = โ„’.kron(โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1), zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs2 = tmp.nzind - - tmp = โ„’.kron(โ„’.kron(e_in_sโบ, e_in_sโบ), zero(e_in_sโบ) .+ 1) |> sparse - shock_idxs3 = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, โ„’.kron(e_in_sโบ, e_in_sโบ)) |> sparse - shockยณ_idxs = tmp.nzind - - tmp = โ„’.kron(zero(e_in_sโบ) .+ 1, โ„’.kron(e_in_sโบ, e_in_sโบ)) |> sparse - shockvar1_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, โ„’.kron(zero(e_in_sโบ) .+ 1, e_in_sโบ)) |> sparse - shockvar2_idxs = tmp.nzind - - tmp = โ„’.kron(e_in_sโบ, โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1)) |> sparse - shockvar3_idxs = tmp.nzind - - shockvarยณ2_idxs = setdiff(shock_idxs2, shockยณ_idxs, shockvar1_idxs, shockvar2_idxs, shockvar3_idxs) - - shockvarยณ_idxs = setdiff(shock_idxs3, shockยณ_idxs)#, shockvar1_idxs, shockvar2_idxs, shockvar3_idxs) - - ๐’ยณโปแต› = ๐’[3][cond_var_idx,var_volยณ_idxs] - ๐’ยณโปแต‰ยฒ = ๐’[3][cond_var_idx,shockvarยณ2_idxs] - ๐’ยณโปแต‰ = ๐’[3][cond_var_idx,shockvarยณ_idxs] - ๐’ยณแต‰ = ๐’[3][cond_var_idx,shockยณ_idxs] - ๐’โปยณ = ๐’[3][T.past_not_future_and_mixed_idx,:] - - ๐’ยณโปแต› = nnz(๐’ยณโปแต›) / length(๐’ยณโปแต›) > .1 ? collect(๐’ยณโปแต›) : ๐’ยณโปแต› - ๐’ยณโปแต‰ = nnz(๐’ยณโปแต‰) / length(๐’ยณโปแต‰) > .1 ? collect(๐’ยณโปแต‰) : ๐’ยณโปแต‰ - ๐’ยณแต‰ = nnz(๐’ยณแต‰) / length(๐’ยณแต‰) > .1 ? collect(๐’ยณแต‰) : ๐’ยณแต‰ - ๐’โปยณ = nnz(๐’โปยณ) / length(๐’โปยณ) > .1 ? collect(๐’โปยณ) : ๐’โปยณ - - stt = state[T.past_not_future_and_mixed_idx] - - kronxx = [zeros(T.nExo^2) for _ in 1:size(data_in_deviations,2)] - - J = โ„’.I(T.nExo) - - kronxxx = [zeros(T.nExo^3) for _ in 1:size(data_in_deviations,2)] - - kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) - - kron_buffer3 = โ„’.kron(J, zeros(T.nExo^2)) - - kron_buffer4 = โ„’.kron(โ„’.kron(J, J), zeros(T.nExo)) - - x = [zeros(T.nExo) for _ in 1:size(data_in_deviations,2)] - - stateยนโป = stt - - stateยนโป_vol = vcat(stateยนโป, 1) - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) - - ๐’โฑยฒแต‰ = [zero(๐’ยฒแต‰) for _ in 1:size(data_in_deviations,2)] - - aug_state = [zeros(size(๐’โปยน,2)) for _ in 1:size(data_in_deviations,2)] - - tmp = ๐’โฑ + 2 * ๐’โฑยฒแต‰[1] * โ„’.kron(โ„’.I(T.nExo), x[1]) - - jacc = [zero(tmp) for _ in 1:size(data_in_deviations,2)] - - ฮป = [zeros(size(tmp, 1)) for _ in 1:size(data_in_deviations,2)] - - ฮป[1] = tmp' \ x[1] * 2 - - fXฮปp_tmp = [reshape(2 * ๐’โฑยฒแต‰[1]' * ฮป[1], size(๐’โฑ, 2), size(๐’โฑ, 2)) - 2 * โ„’.I(size(๐’โฑ, 2)) tmp' - -tmp zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - fXฮปp = [zero(fXฮปp_tmp) for _ in 1:size(data_in_deviations,2)] - - kronxฮป_tmp = โ„’.kron(x[1], ฮป[1]) - - kronxฮป = [kronxฮป_tmp for _ in 1:size(data_in_deviations,2)] - - kronxxฮป_tmp = โ„’.kron(x[1], kronxฮป_tmp) - - kronxxฮป = [kronxxฮป_tmp for _ in 1:size(data_in_deviations,2)] - - II = sparse(โ„’.I(T.nExo^2)) - - lI = 2 * โ„’.I(size(๐’โฑ, 2)) - - ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - - # end # timeit_debug - # @timeit_debug timer "Main loop" begin - - for i in axes(data_in_deviations,2) - stateยนโป = stt - - stateยนโป_vol = vcat(stateยนโป, 1) - - shock_independent = copy(data_in_deviations[:,i]) - - โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - - โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) - - โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 - - ๐’โฑยฒแต‰[i] = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 - - init_guess = zeros(size(๐’โฑ, 2)) - - # @timeit_debug timer "Find shocks" begin - x[i], matched = find_shocks(Val(filter_algorithm), - init_guess, - kronxx[i], - kronxxx[i], - kron_buffer2, - kron_buffer3, - kron_buffer4, - J, - ๐’โฑ, - ๐’โฑยฒแต‰[i], - ๐’โฑยณแต‰, - shock_independent, - # max_iter = 100 - ) - # end # timeit_debug - - if !matched - if opts.verbose println("Inversion filter failed at step $i") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - jacc[i] = ๐’โฑ + 2 * ๐’โฑยฒแต‰[i] * โ„’.kron(โ„’.I(T.nExo), x[i]) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), kronxx[i]) - - ฮป[i] = jacc[i]' \ x[i] * 2 - # โ„’.ldiv!(ฮป[i], tmp', x[i]) - # โ„’.rmul!(ฮป[i], 2) - fXฮปp[i] = [reshape((2 * ๐’โฑยฒแต‰[i] + 6 * ๐’โฑยณแต‰ * โ„’.kron(II, x[i]))' * ฮป[i], size(๐’โฑ, 2), size(๐’โฑ, 2)) - lI jacc[i]' - -jacc[i] zeros(size(๐’โฑ, 1),size(๐’โฑ, 1))] - - โ„’.kron!(kronxx[i], x[i], x[i]) - - โ„’.kron!(kronxฮป[i], x[i], ฮป[i]) - - โ„’.kron!(kronxxฮป[i], x[i], kronxฮป[i]) - - โ„’.kron!(kronxxx[i], x[i], kronxx[i]) - - if i > presample_periods - # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) - logabsdets += โ„’.logabsdet(jacc[i])[1] - else - logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc[i])) - end - - shocksยฒ += sum(abs2,x[i]) - - if !isfinite(logabsdets) || !isfinite(shocksยฒ) - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - end - - aug_state[i] = [stt; 1; x[i]] - - stt = ๐’โปยน * aug_state[i] + ๐’โปยฒ * โ„’.kron(aug_state[i], aug_state[i]) / 2 + ๐’โปยณ * โ„’.kron(โ„’.kron(aug_state[i],aug_state[i]),aug_state[i]) / 6 - end - - # See: https://pcubaborda.net/documents/CGIZ-final.pdf - llh = -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 - - # end # timeit_debug - # end # timeit_debug - - โˆ‚state = similar(state) - - โˆ‚๐’ = copy(๐’) - - โˆ‚data_in_deviations = similar(data_in_deviations) - - function inversion_filter_loglikelihood_pullback(โˆ‚llh) - # @timeit_debug timer "Inversion filter pruned 2nd - pullback" begin - # @timeit_debug timer "Preallocation" begin - - โˆ‚๐’โฑ = zero(๐’โฑ) - โˆ‚๐’ยฒแต‰ = zero(๐’ยฒแต‰) - โˆ‚๐’โฑยณแต‰ = zero(๐’โฑยณแต‰) - - โˆ‚๐’ยนแต‰ = zero(๐’ยนแต‰) - โˆ‚๐’ยฒโปแต‰ = zero(๐’ยฒโปแต‰) - โˆ‚๐’ยณโปแต‰ = zero(๐’ยณโปแต‰) - โˆ‚๐’ยณโปแต‰ยฒ = zero(๐’ยณโปแต‰ยฒ) - - โˆ‚๐’ยนโปแต› = zero(๐’ยนโปแต›) - โˆ‚๐’ยฒโปแต› = zero(๐’ยฒโปแต›) - โˆ‚๐’ยณโปแต› = zero(๐’ยณโปแต›) - - โˆ‚๐’โปยน = zero(๐’โปยน) - โˆ‚๐’โปยฒ = zero(๐’โปยฒ) - โˆ‚๐’โปยณ = zero(๐’โปยณ) - - โˆ‚stateยนโป_vol = zero(stateยนโป_vol) - โˆ‚x = zero(x[1]) - โˆ‚kronxx = zero(kronxx[1]) - โˆ‚kronstateยนโป_vol = zeros(length(stateยนโป_vol)^2) - โˆ‚state = zeros(T.nPast_not_future_and_mixed) - - # end # timeit_debug - # @timeit_debug timer "Main loop" begin - - for i in reverse(axes(data_in_deviations,2)) - # stt = ๐’โปยน * aug_state[i] + ๐’โปยฒ * โ„’.kron(aug_state[i], aug_state[i]) / 2 + ๐’โปยณ * โ„’.kron(โ„’.kron(aug_state[i],aug_state[i]),aug_state[i]) / 6 - โˆ‚๐’โปยน += โˆ‚state * aug_state[i]' - - โˆ‚๐’โปยฒ += โˆ‚state * โ„’.kron(aug_state[i], aug_state[i])' / 2 - - โˆ‚๐’โปยณ += โˆ‚state * โ„’.kron(โ„’.kron(aug_state[i], aug_state[i]), aug_state[i])' / 6 - - โˆ‚aug_state = ๐’โปยน' * โˆ‚state - โˆ‚kronaug_state = ๐’โปยฒ' * โˆ‚state / 2 - โˆ‚kronkronaug_state = ๐’โปยณ' * โˆ‚state / 6 - - fill_kron_adjoint!(โˆ‚aug_state, โˆ‚kronaug_state, โˆ‚kronkronaug_state, aug_state[i], โ„’.kron(aug_state[i], aug_state[i])) - - fill_kron_adjoint!(โˆ‚aug_state, โˆ‚aug_state, โˆ‚kronaug_state, aug_state[i], aug_state[i]) - - if i > 1 && i < size(data_in_deviations,2) - โˆ‚state *= 0 - end - - # aug_state[i] = [stt; 1; x[i]] - โˆ‚state += โˆ‚aug_state[1:length(โˆ‚state)] - - # aug_state[i] = [stt; 1; x[i]] - โˆ‚x = โˆ‚aug_state[T.nPast_not_future_and_mixed+2:end] - - # shocksยฒ += sum(abs2,x[i]) - if i < size(data_in_deviations,2) - โˆ‚x -= copy(x[i]) - else - โˆ‚x += copy(x[i]) - end - - # logabsdets += โ„’.logabsdet(jacc ./ precision_factor)[1] - โˆ‚jacc = try if size(jacc[i], 1) == size(jacc[i], 2) - inv(jacc[i])' - else - โ„’.pinv(jacc[i])' - end - catch - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # jacc = ๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), โ„’.kron(x, x)) - # โˆ‚๐’โฑ = -โˆ‚jacc / 2 # fine - - โˆ‚kronIx = ๐’โฑยฒแต‰[i]' * โˆ‚jacc - - if i < size(data_in_deviations,2) - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, -โ„’.I(T.nExo)) - else - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIx, โˆ‚x, โ„’.I(T.nExo)) - end - - โˆ‚๐’โฑยฒแต‰ = -โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), x[i])' - - โˆ‚kronIxx = ๐’โฑยณแต‰' * โˆ‚jacc * 3 / 2 - - โˆ‚kronxx *= 0 - - if i < size(data_in_deviations,2) - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, -โ„’.I(T.nExo)) - else - fill_kron_adjoint_โˆ‚B!(โˆ‚kronIxx, โˆ‚kronxx, โ„’.I(T.nExo)) - end - - fill_kron_adjoint!(โˆ‚x, โˆ‚x, โˆ‚kronxx, x[i], x[i]) - - โˆ‚๐’โฑยณแต‰ -= โˆ‚jacc * โ„’.kron(โ„’.I(T.nExo), kronxx[i])' * 3 / 2 - - # find_shocks - โˆ‚xฮป = vcat(โˆ‚x, zero(ฮป[i])) - - S = fXฮปp[i]' \ โˆ‚xฮป - - if i < size(data_in_deviations,2) - S *= -1 - end - - โˆ‚shock_independent = S[T.nExo+1:end] # fine - - # โˆ‚๐’โฑ += S[1:T.nExo] * ฮป[i]' - S[T.nExo + 1:end] * x[i]' # fine - copyto!(โˆ‚๐’โฑ, โ„’.kron(S[1:T.nExo], ฮป[i]) - โ„’.kron(x[i], S[T.nExo+1:end])) - โˆ‚๐’โฑ -= โˆ‚jacc / 2 # fine - - โˆ‚๐’โฑยฒแต‰ += reshape(2 * โ„’.kron(S[1:T.nExo], โ„’.kron(x[i], ฮป[i])) - โ„’.kron(kronxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยฒแต‰)) - # โˆ‚๐’โฑยฒแต‰ += 2 * S[1:T.nExo] * kronxฮป[i]' - S[T.nExo + 1:end] * kronxx[i]' - - โˆ‚๐’โฑยณแต‰ += reshape(3 * โ„’.kron(S[1:T.nExo], โ„’.kron(โ„’.kron(x[i], x[i]), ฮป[i])) - โ„’.kron(kronxxx[i], S[T.nExo+1:end]), size(โˆ‚๐’โฑยณแต‰)) - # โˆ‚๐’โฑยณแต‰ += 3 * S[1:T.nExo] * kronxxฮป[i]' - S[T.nExo + 1:end] * kronxxx[i]' - - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 - โˆ‚kronstateยนโป_vol *= 0 - - stateยนโป_vol = [aug_state[i][1:T.nPast_not_future_and_mixed];1] # define here as it is used multiple times later - - โˆ‚๐’ยนแต‰ += โˆ‚๐’โฑ - - โˆ‚stateยนโป_vol *= 0 - - โˆ‚kronIstateยนโป_vol = ๐’ยฒโปแต‰' * โˆ‚๐’โฑ - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_vol, โˆ‚stateยนโป_vol, โ„’.I(T.nExo)) - - โˆ‚๐’ยฒโปแต‰ += โˆ‚๐’โฑ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol)' - - โˆ‚kronIstateยนโป_volstateยนโป_vol = ๐’ยณโปแต‰ยฒ' * โˆ‚๐’โฑ / 2 - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIstateยนโป_volstateยนโป_vol, โˆ‚kronstateยนโป_vol, โ„’.I(T.nExo)) - - โˆ‚๐’ยณโปแต‰ยฒ += โˆ‚๐’โฑ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol)' / 2 - - - # ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 - โˆ‚๐’ยฒแต‰ += โˆ‚๐’โฑยฒแต‰ / 2 - - โˆ‚๐’ยณโปแต‰ += โˆ‚๐’โฑยฒแต‰ * โ„’.kron(II, stateยนโป_vol)' / 2 - - โˆ‚kronIIstateยนโป_vol = ๐’ยณโปแต‰' * โˆ‚๐’โฑยฒแต‰ / 2 - - fill_kron_adjoint_โˆ‚A!(โˆ‚kronIIstateยนโป_vol, โˆ‚stateยนโป_vol, II) - - # shock_independent = copy(data_in_deviations[:,i]) - โˆ‚data_in_deviations[:,i] = โˆ‚shock_independent - - - # โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - โˆ‚๐’ยนโปแต› -= โˆ‚shock_independent * stateยนโป_vol' - - โˆ‚stateยนโป_vol -= ๐’ยนโปแต›' * โˆ‚shock_independent - - # โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) - โˆ‚๐’ยฒโปแต› -= โˆ‚shock_independent * โ„’.kron(stateยนโป_vol, stateยนโป_vol)' / 2 - - โˆ‚kronstateยนโป_vol -= ๐’ยฒโปแต›' * โˆ‚shock_independent / 2 - - # โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) - โˆ‚๐’ยณโปแต› -= โˆ‚shock_independent * โ„’.kron(โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol)' / 6 - - โˆ‚kronstateยนโป_volstateยนโป_vol = -๐’ยณโปแต›' * โˆ‚shock_independent / 6 - - fill_kron_adjoint!(โˆ‚kronstateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_volstateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol), stateยนโป_vol) - - fill_kron_adjoint!(โˆ‚stateยนโป_vol, โˆ‚stateยนโป_vol, โˆ‚kronstateยนโป_vol, stateยนโป_vol, stateยนโป_vol) - - # stateยนโป_vol = vcat(stateยนโป, 1) - โˆ‚state += โˆ‚stateยนโป_vol[1:end-1] - end - - # end # timeit_debug - # @timeit_debug timer "Post allocation" begin - - โˆ‚๐’ = [copy(๐’[1]) * 0, copy(๐’[2]) * 0, copy(๐’[3]) * 0] - - โˆ‚๐’[1][cond_var_idx,end-T.nExo+1:end] += โˆ‚๐’ยนแต‰ - โˆ‚๐’[2][cond_var_idx,shockvarยฒ_idxs] += โˆ‚๐’ยฒโปแต‰ - โˆ‚๐’[2][cond_var_idx,shockยฒ_idxs] += โˆ‚๐’ยฒแต‰ - โˆ‚๐’[3][cond_var_idx,shockvarยณ2_idxs] += โˆ‚๐’ยณโปแต‰ยฒ - โˆ‚๐’[3][cond_var_idx,shockvarยณ_idxs] += โˆ‚๐’ยณโปแต‰ - โˆ‚๐’[3][cond_var_idx,shockยณ_idxs] += โˆ‚๐’โฑยณแต‰ / 6 # ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - - โˆ‚๐’[1][cond_var_idx, 1:T.nPast_not_future_and_mixed+1] += โˆ‚๐’ยนโปแต› - โˆ‚๐’[2][cond_var_idx,var_volยฒ_idxs] += โˆ‚๐’ยฒโปแต› - โˆ‚๐’[3][cond_var_idx,var_volยณ_idxs] += โˆ‚๐’ยณโปแต› - - โˆ‚๐’[1][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยน - โˆ‚๐’[2][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยฒ - โˆ‚๐’[3][T.past_not_future_and_mixed_idx,:] += โˆ‚๐’โปยณ - - โˆ‚๐’[1] *= โˆ‚llh - โˆ‚๐’[2] *= โˆ‚llh - โˆ‚๐’[3] *= โˆ‚llh - - return NoTangent(), NoTangent(), โ„’.I(T.nVars)[:,T.past_not_future_and_mixed_idx] * โˆ‚state * โˆ‚llh, โˆ‚๐’, โˆ‚data_in_deviations * โˆ‚llh, NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - # end # timeit_debug - # end # timeit_debug - - return llh, inversion_filter_loglikelihood_pullback -end - -function rrule(::typeof(run_kalman_iterations), - A, - ๐, - C, - P, - data_in_deviations, - ws::kalman_workspace; - presample_periods = 0, - on_failure_loglikelihood = -Inf, - # timer::TimerOutput = TimerOutput(), - verbose::Bool = false) - # @timeit_debug timer "Calculate Kalman filter - forward" begin - # Note: The rrule requires time-indexed arrays for the backward pass that depend on data length, - # so we cannot cache them in the workspace. Only small fixed-size buffers could potentially be cached. - T = size(data_in_deviations, 2) + 1 - - z = zeros(size(data_in_deviations, 1)) - - uฬ„ = zeros(size(C,2)) - - Pฬ„ = deepcopy(P) - - temp_N_N = similar(P) - - PCtmp = similar(C') - - F = similar(C * C') - - u = [similar(uฬ„) for _ in 1:T] # used in backward pass - - P = [copy(Pฬ„) for _ in 1:T] # used in backward pass - - CP = [zero(C) for _ in 1:T] # used in backward pass - - K = [similar(C') for _ in 1:T] # used in backward pass - - invF = [similar(F) for _ in 1:T] # used in backward pass - - v = [zeros(size(data_in_deviations, 1)) for _ in 1:T] # used in backward pass - - loglik = 0.0 - - # @timeit_debug timer "Loop" begin - - for t in 2:T - if !all(isfinite.(z)) - if verbose println("KF not finite at step $t") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - v[t] .= data_in_deviations[:, t-1] .- z#[t-1] - - # CP[t] .= C * Pฬ„[t-1] - โ„’.mul!(CP[t], C, Pฬ„)#[t-1]) - - # F[t] .= CP[t] * C' - โ„’.mul!(F, CP[t], C') - - luF = RF.lu(F, check = false) - - if !โ„’.issuccess(luF) - if verbose println("KF factorisation failed step $t") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - Fdet = โ„’.det(luF) - - # Early return if determinant is too small, indicating numerical instability. - if Fdet < eps(Float64) - if verbose println("KF factorisation failed step $t") end - return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # invF[t] .= inv(luF) - copy!(invF[t], inv(luF)) - - if t - 1 > presample_periods - loglik += log(Fdet) + โ„’.dot(v[t], invF[t], v[t]) - end - - # K[t] .= Pฬ„[t-1] * C' * invF[t] - โ„’.mul!(PCtmp, Pฬ„, C') - โ„’.mul!(K[t], PCtmp, invF[t]) - - # P[t] .= Pฬ„[t-1] - K[t] * CP[t] - โ„’.mul!(P[t], K[t], CP[t], -1, 0) - P[t] .+= Pฬ„ - - # Pฬ„[t] .= A * P[t] * A' + ๐ - โ„’.mul!(temp_N_N, P[t], A') - โ„’.mul!(Pฬ„, A, temp_N_N) - Pฬ„ .+= ๐ - - # u[t] .= K[t] * v[t] + uฬ„[t-1] - โ„’.mul!(u[t], K[t], v[t]) - u[t] .+= uฬ„ - - # uฬ„[t] .= A * u[t] - โ„’.mul!(uฬ„, A, u[t]) - - # z[t] .= C * uฬ„[t] - โ„’.mul!(z, C, uฬ„) - end - - llh = -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 - - # initialise derivative variables - โˆ‚A = zero(A) - โˆ‚F = zero(F) - โˆ‚Faccum = zero(F) - โˆ‚P = zero(Pฬ„) - โˆ‚uฬ„ = zero(uฬ„) - โˆ‚v = zero(v[1]) - โˆ‚๐ = zero(๐) - โˆ‚data_in_deviations = zero(data_in_deviations) - vtmp = zero(v[1]) - Ptmp = zero(P[1]) - - # end # timeit_debug - # end # timeit_debug - - # pullback - function kalman_pullback(โˆ‚llh) - # @timeit_debug timer "Calculate Kalman filter - reverse" begin - โ„’.rmul!(โˆ‚A, 0) - โ„’.rmul!(โˆ‚Faccum, 0) - โ„’.rmul!(โˆ‚P, 0) - โ„’.rmul!(โˆ‚uฬ„, 0) - โ„’.rmul!(โˆ‚๐, 0) - - # @timeit_debug timer "Loop" begin - for t in T:-1:2 - if t > presample_periods + 1 - # โˆ‚llhโˆ‚F - # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] - # โˆ‚F = invF[t]' - invF[t]' * v[t] * v[t]' * invF[t]' - โ„’.mul!(โˆ‚F, v[t], v[t]') - โ„’.mul!(invF[1], invF[t]', โˆ‚F) # using invF[1] as temporary storage - โ„’.mul!(โˆ‚F, invF[1], invF[t]') - โ„’.axpby!(1, invF[t]', -1, โˆ‚F) - - # โˆ‚llhโˆ‚uฬ„ - # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] - # z[t] .= C * uฬ„[t] - # โˆ‚v = (invF[t]' + invF[t]) * v[t] - copy!(invF[1], invF[t]' .+ invF[t]) - # copy!(invF[1], invF[t]) # using invF[1] as temporary storage - # โ„’.axpy!(1, invF[t]', invF[1]) # using invF[1] as temporary storage - โ„’.mul!(โˆ‚v, invF[1], v[t]) - # โ„’.mul!(โˆ‚uฬ„โˆ‚v, C', v[1]) - else - โ„’.rmul!(โˆ‚F, 0) - โ„’.rmul!(โˆ‚v, 0) - end - - # โˆ‚Fโˆ‚P - # F[t] .= C * Pฬ„[t-1] * C' - # โˆ‚P += C' * (โˆ‚F + โˆ‚Faccum) * C - โ„’.axpy!(1, โˆ‚Faccum, โˆ‚F) - โ„’.mul!(PCtmp, C', โˆ‚F) - โ„’.mul!(โˆ‚P, PCtmp, C, 1, 1) - - # โˆ‚uฬ„โˆ‚P - # K[t] .= Pฬ„[t-1] * C' * invF[t] - # u[t] .= K[t] * v[t] + uฬ„[t-1] - # uฬ„[t] .= A * u[t] - # โˆ‚P += A' * โˆ‚uฬ„ * v[t]' * invF[t]' * C - โ„’.mul!(CP[1], invF[t]', C) # using CP[1] as temporary storage - โ„’.mul!(PCtmp, โˆ‚uฬ„ , v[t]') - โ„’.mul!(P[1], PCtmp , CP[1]) # using P[1] as temporary storage - โ„’.mul!(โˆ‚P, A', P[1], 1, 1) - - # โˆ‚uฬ„โˆ‚data - # v[t] .= data_in_deviations[:, t-1] .- z - # z[t] .= C * uฬ„[t] - # โˆ‚data_in_deviations[:,t-1] = -C * โˆ‚uฬ„ - โ„’.mul!(u[1], A', โˆ‚uฬ„) - โ„’.mul!(v[1], K[t]', u[1]) # using v[1] as temporary storage - โ„’.axpy!(1, โˆ‚v, v[1]) - โˆ‚data_in_deviations[:,t-1] .= v[1] - # โ„’.mul!(โˆ‚data_in_deviations[:,t-1], C, โˆ‚uฬ„, -1, 0) # cannot assign to columns in matrix, must be whole matrix - - # โˆ‚uฬ„โˆ‚uฬ„ - # z[t] .= C * uฬ„[t] - # v[t] .= data_in_deviations[:, t-1] .- z - # K[t] .= Pฬ„[t-1] * C' * invF[t] - # u[t] .= K[t] * v[t] + uฬ„[t-1] - # uฬ„[t] .= A * u[t] - # step to next iteration - # โˆ‚uฬ„ = A' * โˆ‚uฬ„ - C' * K[t]' * A' * โˆ‚uฬ„ - โ„’.mul!(u[1], A', โˆ‚uฬ„) # using u[1] as temporary storage - โ„’.mul!(v[1], K[t]', u[1]) # using v[1] as temporary storage - โ„’.mul!(โˆ‚uฬ„, C', v[1]) - โ„’.mul!(u[1], C', v[1], -1, 1) - copy!(โˆ‚uฬ„, u[1]) - - # โˆ‚llhโˆ‚uฬ„ - # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] - # v[t] .= data_in_deviations[:, t-1] .- z - # z[t] .= C * uฬ„[t] - # โˆ‚uฬ„ -= โˆ‚uฬ„โˆ‚v - โ„’.mul!(u[1], C', โˆ‚v) # using u[1] as temporary storage - โ„’.axpy!(-1, u[1], โˆ‚uฬ„) - - if t > 2 - # โˆ‚uฬ„โˆ‚A - # uฬ„[t] .= A * u[t] - # โˆ‚A += โˆ‚uฬ„ * u[t-1]' - โ„’.mul!(โˆ‚A, โˆ‚uฬ„, u[t-1]', 1, 1) - - # โˆ‚Pฬ„โˆ‚A and โˆ‚Pฬ„โˆ‚๐ - # Pฬ„[t] .= A * P[t] * A' + ๐ - # โˆ‚A += โˆ‚P * A * P[t-1]' + โˆ‚P' * A * P[t-1] - โ„’.mul!(P[1], A, P[t-1]') - โ„’.mul!(Ptmp ,โˆ‚P, P[1]) - โ„’.mul!(P[1], A, P[t-1]) - โ„’.mul!(Ptmp ,โˆ‚P', P[1], 1, 1) - โ„’.axpy!(1, Ptmp, โˆ‚A) - - # โˆ‚๐ += โˆ‚P - โ„’.axpy!(1, โˆ‚P, โˆ‚๐) - - # โˆ‚Pโˆ‚P - # P[t] .= Pฬ„[t-1] - K[t] * C * Pฬ„[t-1] - # Pฬ„[t] .= A * P[t] * A' + ๐ - # step to next iteration - # โˆ‚P = A' * โˆ‚P * A - โ„’.mul!(P[1], โˆ‚P, A) # using P[1] as temporary storage - โ„’.mul!(โˆ‚P, A', P[1]) - - # โˆ‚Pฬ„โˆ‚P - # K[t] .= Pฬ„[t-1] * C' * invF[t] - # P[t] .= Pฬ„[t-1] - K[t] * CP[t] - # โˆ‚P -= C' * K[t-1]' * โˆ‚P + โˆ‚P * K[t-1] * C - โ„’.mul!(PCtmp, โˆ‚P, K[t-1]) - โ„’.mul!(CP[1], K[t-1]', โˆ‚P) # using CP[1] as temporary storage - โ„’.mul!(โˆ‚P, PCtmp, C, -1, 1) - โ„’.mul!(โˆ‚P, C', CP[1], -1, 1) - - # โˆ‚uฬ„โˆ‚F - # K[t] .= Pฬ„[t-1] * C' * invF[t] - # u[t] .= K[t] * v[t] + uฬ„[t-1] - # uฬ„[t] .= A * u[t] - # โˆ‚Faccum = -invF[t-1]' * CP[t-1] * A' * โˆ‚uฬ„ * v[t-1]' * invF[t-1]' - โ„’.mul!(u[1], A', โˆ‚uฬ„) # using u[1] as temporary storage - โ„’.mul!(v[1], CP[t-1], u[1]) # using v[1] as temporary storage - โ„’.mul!(vtmp, invF[t-1]', v[1], -1, 0) - โ„’.mul!(invF[1], vtmp, v[t-1]') # using invF[1] as temporary storage - โ„’.mul!(โˆ‚Faccum, invF[1], invF[t-1]') - - # โˆ‚Pโˆ‚F - # K[t] .= Pฬ„[t-1] * C' * invF[t] - # P[t] .= Pฬ„[t-1] - K[t] * CP[t] - # โˆ‚Faccum -= invF[t-1]' * CP[t-1] * โˆ‚P * CP[t-1]' * invF[t-1]' - โ„’.mul!(CP[1], invF[t-1]', CP[t-1]) # using CP[1] as temporary storage - โ„’.mul!(PCtmp, CP[t-1]', invF[t-1]') - โ„’.mul!(K[1], โˆ‚P, PCtmp) # using K[1] as temporary storage - โ„’.mul!(โˆ‚Faccum, CP[1], K[1], -1, 1) - - end - end - - โ„’.rmul!(โˆ‚P, -โˆ‚llh/2) - โ„’.rmul!(โˆ‚A, -โˆ‚llh/2) - โ„’.rmul!(โˆ‚๐, -โˆ‚llh/2) - โ„’.rmul!(โˆ‚data_in_deviations, -โˆ‚llh/2) - - # end # timeit_debug - # end # timeit_debug - - return NoTangent(), โˆ‚A, โˆ‚๐, NoTangent(), โˆ‚P, โˆ‚data_in_deviations, NoTangent(), NoTangent() - end - - return llh, kalman_pullback -end diff --git a/src/default_options.jl b/src/default_options.jl index 2b690447b..af78e4296 100644 --- a/src/default_options.jl +++ b/src/default_options.jl @@ -105,18 +105,23 @@ const DEFAULT_ARGS_AND_KWARGS_NAMES = Dict( :quadratic_matrix_equation_algorithm => "Quadratic Matrix Equation Algorithm", :sylvester_algorithm => "Sylvester Algorithm", :lyapunov_algorithm => "Lyapunov Algorithm", - :NSSS_acceptance_tol => "NSSS acceptance tol", - :NSSS_xtol => "NSSS xtol", - :NSSS_ftol => "NSSS ftol", - :NSSS_rel_xtol => "NSSS rel xtol", - :qme_tol => "QME tol", - :qme_acceptance_tol => "QME acceptance tol", - :sylvester_tol => "Sylvester tol", - :sylvester_acceptance_tol => "Sylvester acceptance tol", - :lyapunov_tol => "Lyapunov tol", - :lyapunov_acceptance_tol => "Lyapunov acceptance tol", - :droptol => "Droptol", - :dependencies_tol => "Dependencies tol", + :tol => "Tolerance", + :nsss => "NSSS", + :first_order => "1st order", + :second_order => "2nd order", + :third_order => "3rd order", + :qme => "QME", + :sylvester => "Sylvester", + :lyapunov => "Lyapunov", + :atol => "atol", + :rtol => "rtol", + :initial_guess_acceptance_tol => "init. guess acc. tol", + :acceptance_tol => "acc. tol", + :xtol => "xtol", + :ftol => "ftol", + :rel_xtol => "rel. xtol", + :droptol => "droptol", + :dependencies_tol => "dep. tol", ) # Turing distribution wrapper defaults diff --git a/src/filter/find_shocks.jl b/src/filter/find_shocks.jl index a014d02c0..dd22a57d7 100644 --- a/src/filter/find_shocks.jl +++ b/src/filter/find_shocks.jl @@ -1071,6 +1071,8 @@ function find_shocks(::Val{:LagrangeNewton}, iter = 0 @inbounds for i in 1:max_iter iter = i + # Initialize x โŠ— x for the current iterate before using kron_buffer in Jacobian terms. + โ„’.kron!(kron_buffer, x, x) โ„’.kron!(kron_buffer2, J, x) โ„’.kron!(kron_buffer3, J, kron_buffer) diff --git a/src/filter/inversion.jl b/src/filter/inversion.jl index d627735ce..6c9a44b8d 100644 --- a/src/filter/inversion.jl +++ b/src/filter/inversion.jl @@ -8,67 +8,41 @@ from the origin with gradient-based solvers (including the default LagrangeNewto returns the root whose basin contains the origin rather than guaranteeing the global minimum. """ -# Specialization for :inversion filter -function calculate_loglikelihood(::Val{:inversion}, - algorithm, observables, - ๐’, - data_in_deviations, - constants_obj::constants, - presample_periods, - initial_covariance, - state, - warmup_iterations, - filter_algorithm, - opts, - on_failure_loglikelihood, - lyap_ws::lyapunov_workspace, - inv_ws::inversion_workspace, - kalman_ws::kalman_workspace) #; - # timer::TimerOutput = TimerOutput()) - return calculate_inversion_filter_loglikelihood(Val(algorithm), - state, - ๐’, - data_in_deviations, - observables, - constants_obj, - inv_ws, - warmup_iterations = warmup_iterations, - presample_periods = presample_periods, - filter_algorithm = filter_algorithm, - # timer = timer, - opts = opts, - on_failure_loglikelihood = on_failure_loglikelihood) -end - - -function calculate_inversion_filter_loglikelihood(::Val{:first_order}, - state::Vector{Vector{R}}, +function calculate_loglikelihood(::Val{:inversion}, + ::Val{:first_order}, + observables_index::Vector{Int}, ๐’::Matrix{R}, data_in_deviations::Matrix{R}, - observables::Union{Vector{String}, Vector{Symbol}}, constants::constants, - ws::inversion_workspace{Float64}; + state, + workspaces::workspaces; # timer::TimerOutput = TimerOutput(), warmup_iterations::Int = 0, presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, on_failure_loglikelihood::U = -Inf, opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: AbstractFloat,U <: AbstractFloat} + filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: Real,U <: AbstractFloat} T = constants.post_model_macro + ws = workspaces.inversion + ensure_inversion_buffers!(ws, T.nExo, T.nPast_not_future_and_mixed; third_order = false) + ensure_inversion_estimation_buffers!(ws, T.nExo, length(observables_index)) # @timeit_debug timer "Inversion filter" begin # first order - state = copy(state[1]) + state = convert(Vector{R}, state[1]) - precision_factor = 1.0 + precision_factor = one(R) n_obs = size(data_in_deviations,2) - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + cond_var_idx = observables_index + # Use workspace buffers for observation and shock vectors + state_concat = ws.state_concat - shocksยฒ = 0.0 - logabsdets = 0.0 - jac = zeros(0,0) + shocksยฒ = zero(R) + logabsdets = zero(R) + jac = zeros(R, 0, 0) if warmup_iterations > 0 if warmup_iterations >= 1 @@ -92,12 +66,14 @@ function calculate_inversion_filter_loglikelihood(::Val{:first_order}, warmup_shocks = reshape(x, T.nExo, warmup_iterations) for i in 1:warmup_iterations-1 - โ„’.mul!(state, ๐’, vcat(state[T.past_not_future_and_mixed_idx], warmup_shocks[:,i])) + copyto!(state_concat, 1, view(state, T.past_not_future_and_mixed_idx), 1, T.nPast_not_future_and_mixed) + copyto!(state_concat, T.nPast_not_future_and_mixed + 1, view(warmup_shocks, :, i), 1, T.nExo) + โ„’.mul!(state, ๐’, state_concat) # state = state_update(state, warmup_shocks[:,i]) end for i in 1:warmup_iterations - if T.nExo == length(observables) + if T.nExo == length(observables_index) logabsdets += โ„’.logabsdet(jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor)[1] else logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor)) @@ -107,11 +83,13 @@ function calculate_inversion_filter_loglikelihood(::Val{:first_order}, shocksยฒ += sum(abs2,x) end - y = zeros(length(cond_var_idx)) - x = zeros(T.nExo) + y = ws.y_obs + x = ws.x_shocks + fill!(y, zero(R)) + fill!(x, zero(R)) jac = ๐’[cond_var_idx,end-T.nExo+1:end] - if T.nExo == length(observables) + if T.nExo == length(observables_index) jacdecomp = โ„’.lu(jac, check = false) if !โ„’.issuccess(jacdecomp) @@ -155,48 +133,55 @@ function calculate_inversion_filter_loglikelihood(::Val{:first_order}, if !isfinite(shocksยฒ) return on_failure_loglikelihood end end - โ„’.mul!(state, ๐’, vcat(state[T.past_not_future_and_mixed_idx], x)) + # Use pre-allocated state_concat instead of vcat + copyto!(state_concat, 1, view(state, T.past_not_future_and_mixed_idx), 1, T.nPast_not_future_and_mixed) + copyto!(state_concat, T.nPast_not_future_and_mixed + 1, x, 1, T.nExo) + โ„’.mul!(state, ๐’, state_concat) # state = ๐’ * vcat(state[T.past_not_future_and_mixed_idx], x) end # end # timeit_debug # end # timeit_debug - return -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + return -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 # return -(logabsdets + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 end -function calculate_inversion_filter_loglikelihood(::Val{:pruned_second_order}, - state::Vector{Vector{R}}, +function calculate_loglikelihood(::Val{:inversion}, + ::Val{:pruned_second_order}, + observables_index::Vector{Int}, ๐’::Vector{AbstractMatrix{R}}, data_in_deviations::Matrix{R}, - observables::Union{Vector{String}, Vector{Symbol}}, constants::constants, - ws::inversion_workspace{Float64}; + state, + workspaces::workspaces; # timer::TimerOutput = TimerOutput(), warmup_iterations::Int = 0, on_failure_loglikelihood::U = -Inf, presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: AbstractFloat,U <: AbstractFloat} + filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: Real,U <: AbstractFloat} T = constants.post_model_macro + ws = workspaces.inversion # @timeit_debug timer "Pruned 2nd - Inversion filter" begin # @timeit_debug timer "Preallocation" begin # Ensure workspaces are properly sized n_exo = T.nExo n_past = T.nPast_not_future_and_mixed - @ignore_derivatives ensure_inversion_buffers!(ws, n_exo, n_past; third_order = false) + ensure_inversion_buffers!(ws, n_exo, n_past; third_order = false) + ensure_inversion_estimation_buffers!(ws, n_exo, length(observables_index)) n_obs = size(data_in_deviations,2) - cond_var_idx = @ignore_derivatives indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + cond_var_idx = observables_index shocksยฒ = 0.0 logabsdets = 0.0 - cc = @ignore_derivatives ensure_computational_constants!(constants) + cc = ensure_computational_constants!(constants) s_in_sโบ = cc.s_in_s sv_in_sโบ = cc.s_in_sโบ e_in_sโบ = cc.e_in_sโบ @@ -207,7 +192,7 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_second_order}, tmp = โ„’.kron(e_in_sโบ, e_in_sโบ) |> sparse shockยฒ_idxs = tmp.nzind - shockvarยฒ_idxs = @ignore_derivatives setdiff(shock_idxs, shockยฒ_idxs) + shockvarยฒ_idxs = setdiff(shock_idxs, shockยฒ_idxs) tmp = โ„’.kron(sv_in_sโบ, sv_in_sโบ) |> sparse var_volยฒ_idxs = tmp.nzind @@ -259,15 +244,20 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_second_order}, kron_buffer3 = ws.kron_buffer_state kronstateยนโป_vol = ws.kronstate_vol - shock_independent = zeros(size(data_in_deviations,1)) + # Use workspace buffers instead of fresh allocations + shock_independent = ws.shock_independent + fill!(shock_independent, 0.0) - ๐’โฑ = copy(๐’ยนแต‰) + ๐’โฑ = ws.Si_buffer + copyto!(๐’โฑ, ๐’ยนแต‰) - jacc = copy(๐’ยนแต‰) + jacc = ws.jacc_buffer + copyto!(jacc, ๐’ยนแต‰) ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 - init_guess = zeros(size(๐’โฑ, 2)) + init_guess = ws.init_guess + fill!(init_guess, 0.0) # end # timeit_debug # @timeit_debug timer "Loop" begin @@ -378,7 +368,7 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_second_order}, if i > presample_periods # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) + if T.nExo == length(observables_index) logabsdets += โ„’.logabsdet(jacc)[1] else logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc)) @@ -409,24 +399,27 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_second_order}, # end # timeit_debug # See: https://pcubaborda.net/documents/CGIZ-final.pdf and Fair and Taylor (1983) - return -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + return -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 end -function calculate_inversion_filter_loglikelihood(::Val{:second_order}, - state::Vector{R}, +function calculate_loglikelihood(::Val{:inversion}, + ::Val{:second_order}, + observables_index::Vector{Int}, ๐’::Vector{AbstractMatrix{R}}, data_in_deviations::Matrix{R}, - observables::Union{Vector{String}, Vector{Symbol}}, constants::constants, - ws::inversion_workspace{Float64}; + state, + workspaces::workspaces; # timer::TimerOutput = TimerOutput(), on_failure_loglikelihood::U = -Inf, warmup_iterations::Int = 0, presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: AbstractFloat, U <: AbstractFloat} + filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: Real, U <: AbstractFloat} T = constants.post_model_macro + ws = workspaces.inversion # @timeit_debug timer "2nd - Inversion filter" begin # @timeit_debug timer "Preallocation" begin @@ -434,17 +427,18 @@ function calculate_inversion_filter_loglikelihood(::Val{:second_order}, n_exo = T.nExo n_past = T.nPast_not_future_and_mixed ensure_inversion_buffers!(ws, n_exo, n_past; third_order = false) + ensure_inversion_estimation_buffers!(ws, n_exo, length(observables_index)) precision_factor = 1.0 n_obs = size(data_in_deviations,2) - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + cond_var_idx = observables_index shocksยฒ = 0.0 logabsdets = 0.0 - # s_in_sโบ = get_computational_constants(๐“‚).s_in_s + # s_in_sโบ = computational_constants.s_in_s cc = ensure_computational_constants!(constants) sv_in_sโบ = cc.s_in_sโบ e_in_sโบ = cc.e_in_sโบ @@ -501,17 +495,22 @@ function calculate_inversion_filter_loglikelihood(::Val{:second_order}, kron_buffer3 = ws.kron_buffer_state - shock_independent = zeros(size(data_in_deviations,1)) + # Use workspace buffers instead of fresh allocations + shock_independent = ws.shock_independent + fill!(shock_independent, 0.0) kronstateยนโป_vol = ws.kronstate_vol - ๐’โฑ = copy(๐’ยนแต‰) + ๐’โฑ = ws.Si_buffer + copyto!(๐’โฑ, ๐’ยนแต‰) - jacc = copy(๐’ยนแต‰) + jacc = ws.jacc_buffer + copyto!(jacc, ๐’ยนแต‰) ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 - init_guess = zeros(size(๐’โฑ, 2)) + init_guess = ws.init_guess + fill!(init_guess, 0.0) # end # timeit_debug # @timeit_debug timer "Loop" begin @@ -615,7 +614,7 @@ function calculate_inversion_filter_loglikelihood(::Val{:second_order}, if i > presample_periods # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) + if T.nExo == length(observables_index) logabsdets += โ„’.logabsdet(jacc)[1] # ./ precision_factor else logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc)) # ./ precision_factor @@ -647,40 +646,44 @@ function calculate_inversion_filter_loglikelihood(::Val{:second_order}, # end # timeit_debug # See: https://pcubaborda.net/documents/CGIZ-final.pdf - return -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + return -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 end -function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, - state::Vector{Vector{R}}, +function calculate_loglikelihood(::Val{:inversion}, + ::Val{:pruned_third_order}, + observables_index::Vector{Int}, ๐’::Vector{AbstractMatrix{R}}, data_in_deviations::Matrix{R}, - observables::Union{Vector{String}, Vector{Symbol}}, constants::constants, - ws::inversion_workspace{Float64}; + state, + workspaces::workspaces; # timer::TimerOutput = TimerOutput(), on_failure_loglikelihood::U = -Inf, warmup_iterations::Int = 0, presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: AbstractFloat, U <: AbstractFloat} + filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: Real, U <: AbstractFloat} T = constants.post_model_macro + ws = workspaces.inversion # @timeit_debug timer "Inversion filter" begin # Ensure workspaces are properly sized n_exo = T.nExo n_past = T.nPast_not_future_and_mixed - @ignore_derivatives ensure_inversion_buffers!(ws, n_exo, n_past; third_order = true) + ensure_inversion_buffers!(ws, n_exo, n_past; third_order = true) + ensure_inversion_estimation_buffers!(ws, n_exo, length(observables_index); third_order = true) precision_factor = 1.0 n_obs = size(data_in_deviations,2) - cond_var_idx = @ignore_derivatives indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + cond_var_idx = observables_index shocksยฒ = 0.0 logabsdets = 0.0 - cc = @ignore_derivatives ensure_computational_constants!(constants) + cc = ensure_computational_constants!(constants) s_in_sโบ = cc.s_in_s sv_in_sโบ = cc.s_in_sโบ e_in_sโบ = cc.e_in_sโบ @@ -754,61 +757,58 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, state[2] = state[2][T.past_not_future_and_mixed_idx] state[3] = state[3][T.past_not_future_and_mixed_idx] - ๐’โฑ = copy(๐’ยนแต‰) - - jacc = copy(๐’ยนแต‰) - - kron_buffer = zeros(T.nExo^2) - - kron_bufferยฒ = zeros(T.nExo^3) - - II = โ„’.I(T.nExo^2) - + # Use workspace buffers + kron_buffer = ws.kron_buffer + kron_bufferยฒ = ws.kron_bufferยฒ J = โ„’.I(T.nExo) - - kron_buffer2 = โ„’.kron(J, zeros(T.nExo)) - - kron_buffer3 = โ„’.kron(J, kron_buffer) - - kron_buffer4 = โ„’.kron(II, zeros(T.nExo)) - + II = โ„’.I(T.nExo^2) + kron_buffer2 = ws.kron_buffer2 + kron_buffer3 = ws.kron_buffer3 + kron_buffer4 = ws.kron_buffer4 + kron_buffer_state = ws.kron_buffer_state + ๐’โฑ = ws.Si_buffer + jacc = ws.jacc_buffer + shock_independent = ws.shock_independent + init_guess = ws.init_guess + state_vol = ws.state_vol + kronstate_vol = ws.kronstate_vol + kronstate_volยณ = ws.kronstate_volยณ + stateยฒโป_vol = ws.stateยฒโป_vol + + # Pruned-third specific kron buffers (not in ws, allocated once per call) kron_buffer4sv = โ„’.kron(II, vcat(1,state[1])) - - kron_buffer2s = โ„’.kron(J, vcat(state[1], zero(R))) - - kron_buffer2sv = โ„’.kron(J, vcat(1,state[1])) - kron_buffer2ss = โ„’.kron(state[1], state[1]) - - kron_buffer2svsv = โ„’.kron(vcat(1,state[1]), vcat(1,state[1])) - - kron_buffer3svsv = โ„’.kron(kron_buffer2svsv, vcat(1,state[1])) - - kron_buffer3sv = โ„’.kron(kron_buffer2sv, vcat(1,state[1])) + kron_buffer3sv = โ„’.kron(โ„’.kron(J, vcat(1,state[1])), vcat(1,state[1])) # Use workspaces for augmented state kron operations kron_aug_stateโ‚ = ws.kronaug_state kron_kron_aug_stateโ‚ = ws.kron_kron_aug_state + aug_stateโ‚ = ws.aug_stateโ‚ + aug_stateโ‚ฬ‚ = ws.aug_stateโ‚ฬ‚ + aug_stateโ‚‚ = ws.aug_stateโ‚‚ + aug_stateโ‚ƒ = ws.aug_stateโ‚ƒ + stateยนโป = state[1] stateยฒโป = state[2]#[T.past_not_future_and_mixed_idx] stateยณโป = state[3]#[T.past_not_future_and_mixed_idx] - stateยฒโป_vol = zeros(R, length(stateยฒโป) + 1) - # @timeit_debug timer "Loop" begin ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - init_guess = zeros(size(๐’โฑ, 2)) + fill!(init_guess, zero(R)) for i in axes(data_in_deviations,2) - stateยนโป_vol = vcat(stateยนโป, 1) + # stateยนโป_vol = [stateยนโป; 1] + copyto!(state_vol, 1, stateยนโป, 1, n_past) + state_vol[end] = 1 + stateยนโป_vol = state_vol - shock_independent = copy(data_in_deviations[:,i]) + copyto!(shock_independent, view(data_in_deviations, :, i)) โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) @@ -816,33 +816,31 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, โ„’.mul!(shock_independent, ๐’ยนโป, stateยณโป, -1, 1) - โ„’.kron!(kron_buffer2svsv, stateยนโป_vol, stateยนโป_vol) + โ„’.kron!(kronstate_vol, stateยนโป_vol, stateยนโป_vol) - โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kron_buffer2svsv, -1/2, 1) + โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kronstate_vol, -1/2, 1) โ„’.kron!(kron_buffer2ss, stateยนโป, stateยฒโป) โ„’.mul!(shock_independent, ๐’ยฒโป, kron_buffer2ss, -1, 1) - โ„’.kron!(kron_buffer3svsv, kron_buffer2svsv, stateยนโป_vol) + โ„’.kron!(kronstate_volยณ, kronstate_vol, stateยนโป_vol) - โ„’.mul!(shock_independent, ๐’ยณโปแต›, kron_buffer3svsv, -1/6, 1) + โ„’.mul!(shock_independent, ๐’ยณโปแต›, kronstate_volยณ, -1/6, 1) - # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(J, stateยนโป_vol) + ๐’ยฒโปแต›แต‰ * โ„’.kron(J, stateยฒโป) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(J, stateยนโป_vol), stateยนโป_vol) / 2 + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต›แต‰ * kron(J, s2_vol) + ๐’ยฒโปแต‰ * kron(J, sv) + ๐’ยณโปแต‰ยฒ * kron(kron(J, sv), sv) / 2 copyto!(stateยฒโป_vol, 1, stateยฒโป, 1) stateยฒโป_vol[end] = 0 - โ„’.kron!(kron_buffer2s, J, stateยฒโป_vol) + โ„’.kron!(kron_buffer_state, J, stateยฒโป_vol) - โ„’.mul!(๐’โฑ, ๐’ยฒโปแต›แต‰, kron_buffer2s) + โ„’.mul!(๐’โฑ, ๐’ยฒโปแต›แต‰, kron_buffer_state) - โ„’.kron!(kron_buffer2sv, J, stateยนโป_vol) + โ„’.kron!(kron_buffer_state, J, stateยนโป_vol) - โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer2sv, 1, 1) - - โ„’.kron!(kron_buffer2sv, J, stateยนโป_vol) + โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer_state, 1, 1) - โ„’.kron!(kron_buffer3sv, kron_buffer2sv, stateยนโป_vol) + โ„’.kron!(kron_buffer3sv, kron_buffer_state, stateยนโป_vol) โ„’.mul!(๐’โฑ, ๐’ยณโปแต‰ยฒ, kron_buffer3sv, 1/2, 1) @@ -1027,7 +1025,7 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, if i > presample_periods # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) + if T.nExo == length(observables_index) logabsdets += โ„’.logabsdet(jacc)[1] else logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc)) @@ -1040,10 +1038,25 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, end end - aug_stateโ‚ = [stateยนโป; 1; x] - aug_stateโ‚ฬ‚ = [stateยนโป; 0; x] - aug_stateโ‚‚ = [stateยฒโป; 0; zero(x)] - aug_stateโ‚ƒ = [stateยณโป; 0; zero(x)] + # aug_stateโ‚ = [stateยนโป; 1; x] + copyto!(aug_stateโ‚, 1, stateยนโป, 1, n_past) + aug_stateโ‚[n_past + 1] = 1 + copyto!(aug_stateโ‚, n_past + 2, x, 1, n_exo) + + # aug_stateโ‚ฬ‚ = [stateยนโป; 0; x] + copyto!(aug_stateโ‚ฬ‚, 1, stateยนโป, 1, n_past) + aug_stateโ‚ฬ‚[n_past + 1] = 0 + copyto!(aug_stateโ‚ฬ‚, n_past + 2, x, 1, n_exo) + + # aug_stateโ‚‚ = [stateยฒโป; 0; zero(x)] + copyto!(aug_stateโ‚‚, 1, stateยฒโป, 1, n_past) + aug_stateโ‚‚[n_past + 1] = 0 + fill!(view(aug_stateโ‚‚, n_past + 2:n_past + 1 + n_exo), zero(R)) + + # aug_stateโ‚ƒ = [stateยณโป; 0; zero(x)] + copyto!(aug_stateโ‚ƒ, 1, stateยณโป, 1, n_past) + aug_stateโ‚ƒ[n_past + 1] = 0 + fill!(view(aug_stateโ‚ƒ, n_past + 2:n_past + 1 + n_exo), zero(R)) # kron_aug_stateโ‚ = โ„’.kron(aug_stateโ‚, aug_stateโ‚) โ„’.kron!(kron_aug_stateโ‚, aug_stateโ‚, aug_stateโ‚) @@ -1075,24 +1088,27 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, # end # timeit_debug # See: https://pcubaborda.net/documents/CGIZ-final.pdf - return -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + return -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 end -function calculate_inversion_filter_loglikelihood(::Val{:third_order}, - state::Vector{R}, +function calculate_loglikelihood(::Val{:inversion}, + ::Val{:third_order}, + observables_index::Vector{Int}, ๐’::Vector{AbstractMatrix{R}}, data_in_deviations::Matrix{R}, - observables::Union{Vector{String}, Vector{Symbol}}, constants::constants, - ws::inversion_workspace{Float64}; + state, + workspaces::workspaces; # timer::TimerOutput = TimerOutput(), on_failure_loglikelihood::U = -Inf, warmup_iterations::Int = 0, presample_periods::Int = 0, + initial_covariance::Symbol = :theoretical, opts::CalculationOptions = merge_calculation_options(), - filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: AbstractFloat,U <: AbstractFloat} + filter_algorithm::Symbol = :LagrangeNewton)::R where {R <: Real,U <: AbstractFloat} T = constants.post_model_macro + ws = workspaces.inversion # @timeit_debug timer "3rd - Inversion filter" begin # @timeit_debug timer "Preallocation" begin @@ -1100,12 +1116,13 @@ function calculate_inversion_filter_loglikelihood(::Val{:third_order}, n_exo = T.nExo n_past = T.nPast_not_future_and_mixed ensure_inversion_buffers!(ws, n_exo, n_past; third_order = true) + ensure_inversion_estimation_buffers!(ws, n_exo, length(observables_index); third_order = true) precision_factor = 1.0 n_obs = size(data_in_deviations,2) - cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + cond_var_idx = observables_index shocksยฒ = 0.0 logabsdets = 0.0 @@ -1202,31 +1219,48 @@ function calculate_inversion_filter_loglikelihood(::Val{:third_order}, II = sparse(โ„’.I(T.nExo^2)) + # Use workspace buffers for state/estimation temporaries + state_vol = ws.state_vol + kronstate_vol = ws.kronstate_vol + kronstate_volยณ = ws.kronstate_volยณ + kron_buffer_state = ws.kron_buffer_state + shock_independent = ws.shock_independent + init_guess = ws.init_guess + ๐’โฑ = ws.Si_buffer + jacc = ws.jacc_buffer + aug_state = ws.aug_stateโ‚ + kronaug_state = ws.kronaug_state + kron_kron_aug_state = ws.kron_kron_aug_state + ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 + # end # timeit_debug # @timeit_debug timer "Loop" begin for i in axes(data_in_deviations,2) - stateยนโป = state + # Build state_vol = [state; 1] + copyto!(state_vol, 1, state, 1, n_past) + state_vol[end] = 1 + stateยนโป_vol = state_vol - stateยนโป_vol = vcat(stateยนโป, 1) - - shock_independent = copy(data_in_deviations[:,i]) + copyto!(shock_independent, view(data_in_deviations, :, i)) โ„’.mul!(shock_independent, ๐’ยนโปแต›, stateยนโป_vol, -1, 1) - โ„’.mul!(shock_independent, ๐’ยฒโปแต›, โ„’.kron(stateยนโป_vol, stateยนโป_vol), -1/2, 1) + โ„’.kron!(kronstate_vol, stateยนโป_vol, stateยนโป_vol) + โ„’.mul!(shock_independent, ๐’ยฒโปแต›, kronstate_vol, -1/2, 1) - โ„’.mul!(shock_independent, ๐’ยณโปแต›, โ„’.kron(stateยนโป_vol, โ„’.kron(stateยนโป_vol, stateยนโป_vol)), -1/6, 1) - - ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol) + ๐’ยณโปแต‰ยฒ * โ„’.kron(โ„’.kron(โ„’.I(T.nExo), stateยนโป_vol), stateยนโป_vol) / 2 + โ„’.kron!(kronstate_volยณ, stateยนโป_vol, kronstate_vol) + โ„’.mul!(shock_independent, ๐’ยณโปแต›, kronstate_volยณ, -1/6, 1) + + # ๐’โฑ = ๐’ยนแต‰ + ๐’ยฒโปแต‰ * kron(I, sv) + ๐’ยณโปแต‰ยฒ * kron(kron(I, sv), sv) / 2 + โ„’.kron!(kron_buffer_state, J, stateยนโป_vol) + copyto!(๐’โฑ, ๐’ยนแต‰) + โ„’.mul!(๐’โฑ, ๐’ยฒโปแต‰, kron_buffer_state, 1, 1) + โ„’.mul!(๐’โฑ, ๐’ยณโปแต‰ยฒ, โ„’.kron(kron_buffer_state, stateยนโป_vol), 1/2, 1) ๐’โฑยฒแต‰ = ๐’ยฒแต‰ / 2 + ๐’ยณโปแต‰ * โ„’.kron(II, stateยนโป_vol) / 2 - ๐’โฑยณแต‰ = ๐’ยณแต‰ / 6 - - # x, jacc, matchd = find_shocks(Val(:fixed_point), state isa Vector{Float64} ? [state] : state, ๐’, data_in_deviations[:,i], observables, T) - - init_guess = zeros(size(๐’โฑ, 2)) + fill!(init_guess, zero(R)) # @timeit_debug timer "Find shocks" begin x, matched = find_shocks(Val(filter_algorithm), @@ -1374,11 +1408,18 @@ function calculate_inversion_filter_loglikelihood(::Val{:third_order}, # println("LagrangeNewton restart - $mat2: $x3, $(โ„’.norm(x3))") # # end - jacc = -(๐’โฑ + 2 * ๐’โฑยฒแต‰ * โ„’.kron(โ„’.I(T.nExo), x) + 3 * ๐’โฑยณแต‰ * โ„’.kron(โ„’.I(T.nExo), โ„’.kron(x, x))) + # jacc = -(๐’โฑ + 2 * ๐’โฑยฒแต‰ * kron(I,x) + 3 * ๐’โฑยณแต‰ * kron(I, kron(x,x))) + โ„’.kron!(kron_buffer2, J, x) + โ„’.kron!(kron_buffer, x, x) + โ„’.kron!(kron_buffer3, J, kron_buffer) + copyto!(jacc, ๐’โฑ) + โ„’.mul!(jacc, ๐’โฑยฒแต‰, kron_buffer2, 2, 1) + โ„’.mul!(jacc, ๐’โฑยณแต‰, kron_buffer3, 3, 1) + โ„’.rmul!(jacc, -1) if i > presample_periods # due to change of variables: jacobian determinant adjustment - if T.nExo == length(observables) + if T.nExo == length(observables_index) logabsdets += โ„’.logabsdet(jacc)[1] else logabsdets += sum(x -> log(abs(x)), โ„’.svdvals(jacc)) @@ -1391,20 +1432,24 @@ function calculate_inversion_filter_loglikelihood(::Val{:third_order}, end end - aug_state = [state; 1; x] - - # res = ๐’[1][cond_var_idx, :] * aug_state + ๐’[2][cond_var_idx, :] * โ„’.kron(aug_state, aug_state) / 2 + ๐’[3][cond_var_idx, :] * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 - data_in_deviations[:,i] - # println("Match with data: $res") + # aug_state = [state; 1; x] + copyto!(aug_state, 1, state, 1, n_past) + aug_state[n_past + 1] = 1 + copyto!(aug_state, n_past + 2, x, 1, n_exo) - state = ๐’โปยน * aug_state + ๐’โปยฒ * โ„’.kron(aug_state, aug_state) / 2 + ๐’โปยณ * โ„’.kron(โ„’.kron(aug_state,aug_state),aug_state) / 6 - # state = state_update(state, x) + # state = ๐’โปยน * aug_state + ๐’โปยฒ * kron(aug,aug)/2 + ๐’โปยณ * kron(kron(aug,aug),aug)/6 + โ„’.kron!(kronaug_state, aug_state, aug_state) + โ„’.kron!(kron_kron_aug_state, kronaug_state, aug_state) + โ„’.mul!(state, ๐’โปยน, aug_state) + โ„’.mul!(state, ๐’โปยฒ, kronaug_state, 1/2, 1) + โ„’.mul!(state, ๐’โปยณ, kron_kron_aug_state, 1/6, 1) end # end # timeit_debug # end # timeit_debug # See: https://pcubaborda.net/documents/CGIZ-final.pdf - return -(logabsdets + shocksยฒ + (length(observables) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 + return -(logabsdets + shocksยฒ + (length(observables_index) * (warmup_iterations + n_obs - presample_periods)) * log(2 * 3.141592653589793)) / 2 end function filter_data_with_model(๐“‚::โ„ณ, @@ -1425,7 +1470,7 @@ function filter_data_with_model(๐“‚::โ„ณ, SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) - if solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error) + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) @error "No solution for these parameters." return variables, shocks, zeros(0,0), decomposition end @@ -1434,22 +1479,18 @@ function filter_data_with_model(๐“‚::โ„ณ, initial_state = zeros(T.nVars) - โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix + โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) - ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; initial_guess = ๐“‚.caches.qme_solution, - opts = opts) + opts = opts, + parameter_values = ๐“‚.parameter_values) update_perturbation_counter!(๐“‚.counters, solved, order = 1) - if solved ๐“‚.caches.qme_solution = qme_sol end - if !solved @error "No solution for these parameters." return variables, shocks, zeros(0,0), decomposition @@ -1565,9 +1606,9 @@ function filter_data_with_model(๐“‚::โ„ณ, variables = zeros(T.nVars, size(data_in_deviations,2)) shocks = zeros(T.nExo, size(data_in_deviations,2)) - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_second_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, opts = opts) + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_stochastic_steady_state(Val(:second_order), ๐“‚.parameter_values, ๐“‚, opts = opts) - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol @error "Could not find 2nd order stochastic steady state" return variables, shocks, zeros(0,0), zeros(0,0) end @@ -1585,9 +1626,10 @@ function filter_data_with_model(๐“‚::โ„ณ, cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - # s_in_sโบ = get_computational_constants(๐“‚).s_in_s - sv_in_sโบ = get_computational_constants(๐“‚).s_in_sโบ - e_in_sโบ = get_computational_constants(๐“‚).e_in_sโบ + computational_constants = ensure_computational_constants!(๐“‚.constants) + # s_in_sโบ = computational_constants.s_in_s + sv_in_sโบ = computational_constants.s_in_sโบ + e_in_sโบ = computational_constants.e_in_sโบ tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse shock_idxs = tmp.nzind @@ -1786,9 +1828,9 @@ function filter_data_with_model(๐“‚::โ„ณ, observables = get_and_check_observables(T, data_in_deviations) - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_second_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, pruning = true, opts = opts) + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐’โ‚‚ = calculate_stochastic_steady_state(Val(:pruned_second_order), ๐“‚.parameter_values, ๐“‚, opts = opts) - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol @error "Could not find pruned 2nd order stochastic steady state" return variables, shocks, zeros(0,0), zeros(0,0) end @@ -1805,8 +1847,9 @@ function filter_data_with_model(๐“‚::โ„ณ, cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) + computational_constants = ensure_computational_constants!(๐“‚.constants) s_in_sโบ = BitVector(vcat(ones(Bool, T.nPast_not_future_and_mixed), zeros(Bool, T.nExo + 1))) - sv_in_sโบ = get_computational_constants(๐“‚).s_in_sโบ + sv_in_sโบ = computational_constants.s_in_sโบ e_in_sโบ = BitVector(vcat(zeros(Bool, T.nPast_not_future_and_mixed + 1), ones(Bool, T.nExo))) tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse @@ -2056,9 +2099,9 @@ function filter_data_with_model(๐“‚::โ„ณ, observables = get_and_check_observables(T, data_in_deviations) - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_third_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, opts = opts) # timer = timer, + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_stochastic_steady_state(Val(:third_order), ๐“‚.parameter_values, ๐“‚, opts = opts) # timer = timer, - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol @error "Could not find 3rd order stochastic steady state" return variables, shocks, zeros(0,0), zeros(0,0) end @@ -2076,9 +2119,10 @@ function filter_data_with_model(๐“‚::โ„ณ, cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - s_in_sโบ = get_computational_constants(๐“‚).s_in_s - sv_in_sโบ = get_computational_constants(๐“‚).s_in_sโบ - e_in_sโบ = get_computational_constants(๐“‚).e_in_sโบ + computational_constants = ensure_computational_constants!(๐“‚.constants) + s_in_sโบ = computational_constants.s_in_s + sv_in_sโบ = computational_constants.s_in_sโบ + e_in_sโบ = computational_constants.e_in_sโบ tmp = โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1) |> sparse shock_idxs = tmp.nzind @@ -2370,9 +2414,9 @@ function filter_data_with_model(๐“‚::โ„ณ, observables = get_and_check_observables(T, data_in_deviations) - sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_third_order_stochastic_steady_state(๐“‚.parameter_values, ๐“‚, pruning = true, opts = opts) # timer = timer, + sss, converged, SS_and_pars, solution_error, โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ = calculate_stochastic_steady_state(Val(:pruned_third_order), ๐“‚.parameter_values, ๐“‚, opts = opts) # timer = timer, - if !converged || solution_error > opts.tol.NSSS_acceptance_tol + if !converged || solution_error > opts.tol.nsss.acceptance_tol @error "Could not find pruned 3rd order stochastic steady state" return variables, shocks, zeros(0,0), zeros(0,0) end @@ -2389,9 +2433,10 @@ function filter_data_with_model(๐“‚::โ„ณ, cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - s_in_sโบ = get_computational_constants(๐“‚).s_in_s - sv_in_sโบ = get_computational_constants(๐“‚).s_in_sโบ - e_in_sโบ = get_computational_constants(๐“‚).e_in_sโบ + computational_constants = ensure_computational_constants!(๐“‚.constants) + s_in_sโบ = computational_constants.s_in_s + sv_in_sโบ = computational_constants.s_in_sโบ + e_in_sโบ = computational_constants.e_in_sโบ tmp = โ„’.kron(e_in_sโบ, s_in_sโบ) |> sparse shockvar_idxs = tmp.nzind diff --git a/src/filter/kalman.jl b/src/filter/kalman.jl index a50f1e6dd..e60d6b60c 100644 --- a/src/filter/kalman.jl +++ b/src/filter/kalman.jl @@ -1,93 +1,38 @@ @stable default_mode = "disable" begin -# Specialization for :kalman filter -function calculate_loglikelihood(::Val{:kalman}, - algorithm, - observables, - ๐’, - data_in_deviations, - constants_obj::constants, - presample_periods, - initial_covariance, - state, - warmup_iterations, - filter_algorithm, - opts, - on_failure_loglikelihood, - lyap_ws::lyapunov_workspace, - inv_ws::inversion_workspace, - kalman_ws::kalman_workspace) #; - # timer::TimerOutput = TimerOutput()) - return calculate_kalman_filter_loglikelihood(observables, - ๐’, - data_in_deviations, - constants_obj, - lyap_ws, - kalman_ws, - presample_periods = presample_periods, - initial_covariance = initial_covariance, - # timer = timer, - opts = opts, - on_failure_loglikelihood = on_failure_loglikelihood) -end - -function calculate_kalman_filter_loglikelihood(observables::Vector{Symbol}, - ๐’::Union{Matrix{S},Vector{AbstractMatrix{S}}}, - data_in_deviations::Matrix{S}, - constants::constants, - lyap_ws::lyapunov_workspace, - kalman_ws::kalman_workspace; - # timer::TimerOutput = TimerOutput(), - on_failure_loglikelihood::U = -Inf, - presample_periods::Int = 0, - initial_covariance::Symbol = :theoretical, - opts::CalculationOptions = merge_calculation_options())::S where {S <: Real, U <: AbstractFloat} - T = constants.post_model_macro - obs_idx = @ignore_derivatives convert(Vector{Int},indexin(observables,sort(union(T.aux,T.var,T.exo_present)))) - - calculate_kalman_filter_loglikelihood(obs_idx, ๐’, data_in_deviations, constants, lyap_ws, kalman_ws, presample_periods = presample_periods, initial_covariance = initial_covariance, opts = opts, on_failure_loglikelihood = on_failure_loglikelihood) - # timer = timer, -end - -function calculate_kalman_filter_loglikelihood(observables::Vector{String}, +function calculate_loglikelihood(::Val{:kalman}, + ::Val, + observables_index::Vector{Int}, ๐’::Union{Matrix{S},Vector{AbstractMatrix{S}}}, data_in_deviations::Matrix{S}, constants::constants, - lyap_ws::lyapunov_workspace, - kalman_ws::kalman_workspace; - # timer::TimerOutput = TimerOutput(), - presample_periods::Int = 0, - on_failure_loglikelihood::U = -Inf, - initial_covariance::Symbol = :theoretical, - opts::CalculationOptions = merge_calculation_options())::S where {S <: Real, U <: AbstractFloat} - T = constants.post_model_macro - obs_idx = @ignore_derivatives convert(Vector{Int},indexin(observables,sort(union(T.aux,T.var,T.exo_present)))) - - calculate_kalman_filter_loglikelihood(obs_idx, ๐’, data_in_deviations, constants, lyap_ws, kalman_ws, presample_periods = presample_periods, initial_covariance = initial_covariance, opts = opts, on_failure_loglikelihood = on_failure_loglikelihood) - # timer = timer, -end - -function calculate_kalman_filter_loglikelihood(observables_index::Vector{Int}, - ๐’::Union{Matrix{S},Vector{AbstractMatrix{S}}}, - data_in_deviations::Matrix{S}, - constants::constants, - lyap_ws::lyapunov_workspace, - kalman_ws::kalman_workspace; + state, + workspaces::workspaces; # timer::TimerOutput = TimerOutput(), + warmup_iterations::Int = 0, presample_periods::Int = 0, initial_covariance::Symbol = :theoretical, + filter_algorithm::Symbol = :LagrangeNewton, lyapunov_algorithm::Symbol = :doubling, on_failure_loglikelihood::U = -Inf, opts::CalculationOptions = merge_calculation_options())::S where {S <: Real, U <: AbstractFloat} T = constants.post_model_macro - observables_and_states = @ignore_derivatives sort(union(T.past_not_future_and_mixed_idx,observables_index)) + idx_constants = constants.post_complete_parameters + lyap_ws = ensure_lyapunov_workspace!(workspaces, T.nVars, :first_order) - A = ๐’[observables_and_states,1:T.nPast_not_future_and_mixed] * โ„’.diagm(ones(S, length(observables_and_states)))[@ignore_derivatives(indexin(T.past_not_future_and_mixed_idx,observables_and_states)),:] - B = ๐’[observables_and_states,T.nPast_not_future_and_mixed+1:end] + observables_and_states = sort(union(T.past_not_future_and_mixed_idx,observables_index)) + observables_sorted = sort(observables_index) + I_nVars = idx_constants.diag_nVars - C = โ„’.diagm(ones(length(observables_and_states)))[@ignore_derivatives(indexin(sort(observables_index), observables_and_states)),:] + A = @views ๐’[observables_and_states,1:T.nPast_not_future_and_mixed] * I_nVars[T.past_not_future_and_mixed_idx, observables_and_states] + B = @views ๐’[observables_and_states,T.nPast_not_future_and_mixed+1:end] - ๐ = B * B' + C = @views I_nVars[observables_sorted, observables_and_states] + + kalman_ws = ensure_kalman_workspaces!(workspaces, size(C, 1), size(C, 2)) + + ๐ = kalman_ws.๐ + โ„’.mul!(๐, B, B') # Gaussian Prior P = get_initial_covariance(Val(initial_covariance), A, ๐, lyap_ws, opts = opts) @@ -106,11 +51,10 @@ function get_initial_covariance(::Val{:theoretical}, # timer::TimerOutput = TimerOutput(), P, _ = solve_lyapunov_equation(A, B, lyap_ws, lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, + tol = opts.tol.first_order.lyapunov, verbose = opts.verbose) # timer = timer, - return P + return copy(P) end @@ -121,28 +65,23 @@ function get_initial_covariance(::Val{:diagonal}, lyap_ws::lyapunov_workspace; opts::CalculationOptions = merge_calculation_options())::Matrix{S} where S <: Real # timer::TimerOutput = TimerOutput(), - P = @ignore_derivatives collect(โ„’.I(size(A, 1)) * 10.0) + P = collect(โ„’.I(size(A, 1)) * 10.0) return P end function run_kalman_iterations(A::Matrix{S}, ๐::Matrix{S}, - C::Matrix{Float64}, + C::AbstractMatrix{R}, P::Matrix{S}, data_in_deviations::Matrix{S}, ws::kalman_workspace; presample_periods::Int = 0, on_failure_loglikelihood::U = -Inf, # timer::TimerOutput = TimerOutput(), - verbose::Bool = false)::S where {S <: Float64, U <: AbstractFloat} + verbose::Bool = false)::S where {S <: Float64, R <: Real, U <: AbstractFloat} # @timeit_debug timer "Calculate Kalman filter" begin - # Ensure workspaces are properly sized - n_obs = size(C, 1) - n_states = size(C, 2) - @ignore_derivatives ensure_kalman_buffers!(ws, n_obs, n_states) - # Use workspaces u = ws.u z = ws.z @@ -162,12 +101,12 @@ function run_kalman_iterations(A::Matrix{S}, # @timeit_debug timer "Loop" begin for t in 1:size(data_in_deviations, 2) - if !all(isfinite.(z)) + if any(!isfinite, z) if verbose println("KF not finite at step $t") end return on_failure_loglikelihood end - โ„’.axpby!(1, data_in_deviations[:, t], -1, z) + โ„’.axpby!(1, @view(data_in_deviations[:, t]), -1, z) # v = data_in_deviations[:, t] - z โ„’.mul!(Ctmp, C, P) # use Octavian.jl @@ -175,18 +114,30 @@ function run_kalman_iterations(A::Matrix{S}, # F = C * P * C' # @timeit_debug timer "LU factorisation" begin - luF = RF.lu!(F, check = false) ### has to be LU since F will always be symmetric and positive semi-definite but not positive definite (due to linear dependencies) + ws.fast_lu_ws_f, ws.fast_lu_dims_f, solved_F, luF = factorize_lu!(F, + ws.fast_lu_ws_f, + ws.fast_lu_dims_f) # end # timeit_debug - if !โ„’.issuccess(luF) + if !solved_F if verbose println("KF factorisation failed step $t") end return on_failure_loglikelihood end - Fdet = โ„’.det(luF) + logabsdetF = zero(S) + signF = isodd(count(i -> ws.fast_lu_ws_f.ipiv[i] != i, eachindex(ws.fast_lu_ws_f.ipiv))) ? -one(S) : one(S) + @inbounds for i in 1:size(F, 1) + di = F[i, i] + if di == 0 + if verbose println("KF factorisation failed step $t") end + return on_failure_loglikelihood + end + logabsdetF += log(abs(di)) + signF *= sign(di) + end # Early return if determinant is too small, indicating numerical instability. - if Fdet < eps(Float64) + if signF <= 0 || logabsdetF < log(eps(Float64)) if verbose println("KF factorisation failed step $t") end return on_failure_loglikelihood end @@ -195,8 +146,9 @@ function run_kalman_iterations(A::Matrix{S}, # @timeit_debug timer "LU div" begin if t > presample_periods - โ„’.ldiv!(ztmp, luF, z) - loglik += log(Fdet) + โ„’.dot(z', ztmp) ### + copyto!(ztmp, z) + solve_lu_left!(F, ztmp, ws.fast_lu_ws_f, luF) + loglik += logabsdetF + โ„’.dot(z', ztmp) ### # loglik += log(Fdet) + z' * invF * z### # loglik += log(Fdet) + v' * invF * v### end @@ -204,7 +156,7 @@ function run_kalman_iterations(A::Matrix{S}, # โ„’.mul!(Ktmp, P, C') # โ„’.mul!(K, Ktmp, invF) โ„’.mul!(K, P, C') - โ„’.rdiv!(K, luF) + solve_lu_right!(F, K, ws.fast_lu_ws_f, luF, ws.fast_lu_rhs_t_k) # K = P * Ct / luF # K = P * C' * invF @@ -284,20 +236,16 @@ function filter_and_smooth(๐“‚::โ„ณ, SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts) - @assert solution_error < opts.tol.NSSS_acceptance_tol "Could not solve non-stochastic steady state." + @assert solution_error < opts.tol.nsss.acceptance_tol "Could not solve non-stochastic steady state." - โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix + โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) - sol, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; - opts = opts) - - if solved ๐“‚.caches.qme_solution = qme_sol end + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved, order = 1) @@ -306,7 +254,7 @@ function filter_and_smooth(๐“‚::โ„ณ, B = @views sol[:,T.nPast_not_future_and_mixed+1:end] - C = @views โ„’.diagm(ones(T.nVars))[sort(indexin(observables,sort(union(๐“‚.constants.post_model_macro.aux,๐“‚.constants.post_model_macro.var,๐“‚.constants.post_model_macro.exo_present)))),:] + C = @views โ„’.diagm(ones(T.nVars))[sort(indexin(observables, sort(union(T.aux, T.var, T.exo_present)))),:] ๐ = B * B' diff --git a/src/get_functions.jl b/src/get_functions.jl index d8fc003cc..1612a04ff 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -856,7 +856,7 @@ function get_conditional_forecast(๐“‚::โ„ณ, if algorithm โˆˆ [:second_order, :third_order, :pruned_second_order, :pruned_third_order] Sโ‚ = ๐“‚.caches.first_order_solution_matrix - Sโ‚ = [Sโ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) Sโ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] + Sฬ‚โ‚ = [Sโ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] zeros(๐“‚.constants.post_model_macro.nVars) Sโ‚[:,๐“‚.constants.post_model_macro.nPast_not_future_and_mixed+1:end]] Sโ‚‚ = nothing if size(๐“‚.caches.second_order_solution, 2) > 0 @@ -868,7 +868,7 @@ function get_conditional_forecast(๐“‚::โ„ณ, Sโ‚ƒ = ๐“‚.caches.third_order_solution * ๐“‚.constants.third_order.๐”โ‚ƒ end - ensure_conditional_forecast_constants!(๐“‚; third_order = !isnothing(Sโ‚ƒ)) + ensure_conditional_forecast_constants!(๐“‚.constants; third_order = !isnothing(Sโ‚ƒ)) # Use Lagrange-Newton algorithm to find shocks x, matched = find_shocks_conditional_forecast(Val(conditional_forecast_solver), @@ -878,7 +878,7 @@ function get_conditional_forecast(๐“‚::โ„ณ, cond_var_idx, free_shock_idx, state_update, - Sโ‚, + Sฬ‚โ‚, Sโ‚‚, Sโ‚ƒ, ๐“‚.constants, @@ -920,7 +920,7 @@ function get_conditional_forecast(๐“‚::โ„ณ, cond_var_idx, free_shock_idx, state_update, - Sโ‚, + Sฬ‚โ‚, Sโ‚‚, Sโ‚ƒ, ๐“‚.constants, @@ -944,7 +944,7 @@ function get_conditional_forecast(๐“‚::โ„ณ, if length(cond_var_idx) == 1 @assert any(CC .!= 0) "Free shocks have no impact on conditioned variable in period 1." elseif length(free_shock_idx) == length(cond_var_idx) - CC = RF.lu(CC, check = false) + CC = โ„’.lu(CC, check = false) @assert โ„’.issuccess(CC) "Numerical stabiltiy issues for restrictions in period 1." end @@ -973,7 +973,7 @@ function get_conditional_forecast(๐“‚::โ„ณ, @assert any(CC .!= 0) "Free shocks have no impact on conditioned variable in period " * repr(i) * "." elseif length(free_shock_idx) == length(cond_var_idx) - CC = RF.lu(CC, check = false) + CC = โ„’.lu(CC, check = false) @assert โ„’.issuccess(CC) "Numerical stabiltiy issues for restrictions in period " * repr(i) * "." end @@ -1073,9 +1073,9 @@ function get_irf(๐“‚::โ„ณ, # Initialize constants at entry point constants = initialise_constants!(๐“‚) - @ignore_derivatives solve!(๐“‚, - steady_state_function = steady_state_function, - opts = opts) + solve!(๐“‚, + steady_state_function = steady_state_function, + opts = opts) shocks = ๐“‚.constants.post_model_macro.nExo == 0 ? :none : shocks @@ -1087,28 +1087,23 @@ function get_irf(๐“‚::โ„ณ, reference_steady_state, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts, estimation = estimation) - if (solution_error > tol.NSSS_acceptance_tol) || isnan(solution_error) + if (solution_error > tol.nsss.acceptance_tol) || isnan(solution_error) return zeros(S, length(var_idx), periods, shocks == :none ? 1 : length(shock_idx)) end - โˆ‡โ‚ = calculate_jacobian(parameters, reference_steady_state, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix + โˆ‡โ‚ = calculate_jacobian(parameters, reference_steady_state, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix - # Ensure QME workspace - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) - sol_mat, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, - constants, - qme_ws, - sylv_ws; - opts = opts, - initial_guess = ๐“‚.caches.qme_solution) + constants, + ๐“‚.workspaces, + ๐“‚.caches; + opts = opts, + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameters) - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) - if solved - ๐“‚.caches.qme_solution = qme_sol - else + if !solved return zeros(S, length(var_idx), periods, shocks == :none ? 1 : length(shock_idx)) end @@ -1501,8 +1496,8 @@ function get_steady_state(๐“‚::โ„ณ; SS, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) - if solution_error > tol.NSSS_acceptance_tol - @warn "Could not find non-stochastic steady state. Solution error: $solution_error > $(tol.NSSS_acceptance_tol)" + if solution_error > tol.nsss.acceptance_tol + @warn "Could not find non-stochastic steady state. Solution error: $solution_error > $(tol.nsss.acceptance_tol)" end if stochastic @@ -1525,9 +1520,10 @@ function get_steady_state(๐“‚::โ„ณ; end end - var_idx = indexin([vars_in_ss_equations...], [๐“‚.constants.post_model_macro.var...,๐“‚.equations.calibration_parameters...]) + ms = ensure_model_structure_constants!(๐“‚.constants, ๐“‚.equations.calibration_parameters) + var_idx = ms.ss_var_idx_in_var_and_calib - calib_idx = return_variables_only ? [] : indexin([๐“‚.equations.calibration_parameters...], [๐“‚.constants.post_model_macro.var...,๐“‚.equations.calibration_parameters...]) + calib_idx = return_variables_only ? Int[] : ms.calib_idx_in_var_and_calib if length_par * length(var_idx) > 200 && derivatives @info "Most of the time is spent calculating derivatives wrt parameters. If they are not needed, add `derivatives = false` as an argument to the function call." maxlog = DEFAULT_MAXLOG @@ -1552,57 +1548,45 @@ function get_steady_state(๐“‚::โ„ณ; if derivatives if stochastic - if algorithm == :third_order - - # dSSS = ๐’œ.jacobian(๐’ท(), x->begin - # SSS = SSS_third_order_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose) - # [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - # end, ๐“‚.parameter_values[param_idx])[1] - dSSS = ๐’Ÿ.jacobian(x -> begin SSS = calculate_third_order_stochastic_steady_state(x, ๐“‚, opts = opts) - return [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - end, backend, ๐“‚.parameter_values)[:,param_idx] - - return KeyedArray(hcat(SS[[var_idx...,calib_idx...]], dSSS); Variables_and_calibrated_parameters = axis1, Steady_state_and_โˆ‚steady_stateโˆ‚parameter = axis2) - - elseif algorithm == :pruned_third_order - - # dSSS = ๐’œ.jacobian(๐’ท(), x->begin - # SSS = SSS_third_order_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose, pruning = true) - # [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - # end, ๐“‚.parameter_values[param_idx])[1] - dSSS = ๐’Ÿ.jacobian(x-> begin SSS = calculate_third_order_stochastic_steady_state(x, ๐“‚, opts = opts, pruning = true) - return [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - end, backend, ๐“‚.parameter_values)[:,param_idx] - - return KeyedArray(hcat(SS[[var_idx...,calib_idx...]], dSSS); Variables_and_calibrated_parameters = axis1, Steady_state_and_โˆ‚steady_stateโˆ‚parameter = axis2) - - elseif algorithm == :pruned_second_order - # dSSS = ๐’œ.jacobian(๐’ท(), x->begin - # SSS = SSS_second_order_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose, pruning = true) - # [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - # end, ๐“‚.parameter_values[param_idx])[1] - dSSS = ๐’Ÿ.jacobian(x->begin SSS = calculate_second_order_stochastic_steady_state(x, ๐“‚, opts = opts, pruning = true) - return [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - end, backend, ๐“‚.parameter_values)[:,param_idx] - - return KeyedArray(hcat(SS[[var_idx...,calib_idx...]], dSSS); Variables_and_calibrated_parameters = axis1, Steady_state_and_โˆ‚steady_stateโˆ‚parameter = axis2) - - else - # dSSS = ๐’œ.jacobian(๐’ท(), x->begin - # SSS = SSS_second_order_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose) - # [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - # end, ๐“‚.parameter_values[param_idx])[1] - dSSS = ๐’Ÿ.jacobian(x->begin SSS = calculate_second_order_stochastic_steady_state(x, ๐“‚, opts = opts) - return [collect(SSS[1])[var_idx]...,collect(SSS[3])[calib_idx]...] - end, backend, ๐“‚.parameter_values)[:,param_idx] - - return KeyedArray(hcat(SS[[var_idx...,calib_idx...]], dSSS); Variables_and_calibrated_parameters = axis1, Steady_state_and_โˆ‚steady_stateโˆ‚parameter = axis2) - + n_tuple = algorithm โˆˆ (:third_order, :pruned_third_order) ? 10 : 8 + SSS_result, SSS_pb = rrule(calculate_stochastic_steady_state, Val(algorithm), ๐“‚.parameter_values, ๐“‚, opts = opts) + SSS = SSS_result[1] + n_sss = length(SSS) + n_ss = length(SSS_result[3]) + nv = length(var_idx) + nc = length(calib_idx) + n_out = nv + nc + np = length(๐“‚.parameter_values) + dSSS = zeros(n_out, np) + for j in 1:n_out + if j โ‰ค nv + โˆ‚sss = zeros(n_sss); โˆ‚sss[var_idx[j]] = 1.0 + seed = ntuple(k -> k == 1 ? โˆ‚sss : NoTangent(), n_tuple) + else + โˆ‚ss = zeros(n_ss); โˆ‚ss[calib_idx[j - nv]] = 1.0 + seed = ntuple(k -> k == 3 ? โˆ‚ss : NoTangent(), n_tuple) + end + โˆ‚p = SSS_pb(seed)[3] + if !(โˆ‚p isa AbstractZero); dSSS[j, :] .= โˆ‚p; end end + dSSS = dSSS[:, param_idx] + + SS_and_pars = SSS_result[3] + steady_state_column = vcat(SSS[var_idx], SS_and_pars[calib_idx]) + return KeyedArray(hcat(steady_state_column, dSSS); Variables_and_calibrated_parameters = axis1, Steady_state_and_โˆ‚steady_stateโˆ‚parameter = axis2) else - # dSS = ๐’œ.jacobian(๐’ท(), x->๐“‚.functions.NSSS_solve(x, ๐“‚),๐“‚.parameter_values) - # dSS = ๐’œ.jacobian(๐’ท(), x->collect(SS_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose)[1])[[var_idx...,calib_idx...]], ๐“‚.parameter_values[param_idx])[1] - dSS = ๐’Ÿ.jacobian(x->get_NSSS_and_parameters(๐“‚, x, opts = opts)[1][[var_idx...,calib_idx...]], backend, ๐“‚.parameter_values)[:,param_idx] + (nsss_result, nsss_pb) = rrule(get_NSSS_and_parameters, ๐“‚, ๐“‚.parameter_values, opts = opts) + out_idx = [var_idx..., calib_idx...] + n_ss = length(nsss_result[1]) + np = length(๐“‚.parameter_values) + n_out = length(out_idx) + dSS = zeros(n_out, np) + for j in 1:n_out + โˆ‚ss = zeros(n_ss); โˆ‚ss[out_idx[j]] = 1.0 + โˆ‚p = nsss_pb((โˆ‚ss, NoTangent()))[3] + if !(โˆ‚p isa AbstractZero); dSS[j, :] .= โˆ‚p; end + end + dSS = dSS[:, param_idx] # if length(๐“‚.calibration_equations_parameters) == 0 # return KeyedArray(hcat(collect(NSSS)[1:(end-1)],dNSSS); Variables = [sort(union(๐“‚.constants.post_model_macro.exo_present,var))...], Steady_state_and_โˆ‚steady_stateโˆ‚parameter = vcat(:Steady_state, ๐“‚.constants.post_complete_parameters.parameters)) @@ -1621,8 +1605,6 @@ function get_steady_state(๐“‚::โ„ณ; # calibrated_parameters = ComponentVector(NSSS.non_stochastic_steady_state, Axis(๐“‚.calibration_equations_parameters)), # stochastic = stochastic) - # return ๐“‚.caches.outdated_NSSS ? ๐“‚.functions.NSSS_solve(๐“‚.parameter_values, ๐“‚) : ๐“‚.caches.non_stochastic_steady_state - # return ๐“‚.functions.NSSS_solve(๐“‚) # return (var .=> ๐“‚.parameter_to_steady_state(๐“‚.parameter_values...)[1:length(var)]), (๐“‚.par .=> ๐“‚.parameter_to_steady_state(๐“‚.parameter_values...)[length(var)+1:end])[getindex(1:length(๐“‚.par),map(x->x โˆˆ collect(๐“‚.calibration_equations_parameters),๐“‚.par))] end @@ -1661,22 +1643,22 @@ sss(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true """ See [`get_steady_state`](@ref) """ -SS = get_steady_state +SS(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -steady_state = get_steady_state +steady_state(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -get_SS = get_steady_state +get_SS(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -get_ss = get_steady_state +get_ss(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) @@ -1836,7 +1818,10 @@ function get_solution(๐“‚::โ„ณ; axis1 = [:Steady_state; map(x->Symbol(string(x) * "โ‚โ‚‹โ‚โ‚Ž"),๐“‚.constants.post_model_macro.past_not_future_and_mixed); map(x->Symbol(string(x) * "โ‚โ‚“โ‚Ž"),๐“‚.constants.post_model_macro.exo)] end - return KeyedArray([๐“‚.caches.non_stochastic_steady_state[1:length(๐“‚.constants.post_model_macro.var)] solution_matrix]'; + n_vars = length(๐“‚.constants.post_model_macro.var) + nsss = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts)[1][1:n_vars] + + return KeyedArray([nsss solution_matrix]'; Steady_state__States__Shocks = axis1, Variables = axis2) end @@ -1932,16 +1917,16 @@ function get_solution(๐“‚::โ„ณ, # Initialize constants at entry point constants = initialise_constants!(๐“‚) - @ignore_derivatives solve!(๐“‚, - opts = opts, - steady_state_function = steady_state_function, - algorithm = algorithm) + solve!(๐“‚, + opts = opts, + steady_state_function = steady_state_function, + algorithm = algorithm) if length(๐“‚.constants.post_parameters_macro.bounds) > 0 for (k,v) in ๐“‚.constants.post_parameters_macro.bounds if k โˆˆ ๐“‚.constants.post_complete_parameters.parameters - if @ignore_derivatives min(max(parameters[indexin([k], ๐“‚.constants.post_complete_parameters.parameters)][1], v[1]), v[2]) != parameters[indexin([k], ๐“‚.constants.post_complete_parameters.parameters)][1] + if min(max(parameters[indexin([k], ๐“‚.constants.post_complete_parameters.parameters)][1], v[1]), v[2]) != parameters[indexin([k], ๐“‚.constants.post_complete_parameters.parameters)][1] return -Inf end end @@ -1950,7 +1935,7 @@ function get_solution(๐“‚::โ„ณ, SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts, estimation = estimation) - if solution_error > tol.NSSS_acceptance_tol || isnan(solution_error) + if solution_error > tol.nsss.acceptance_tol || isnan(solution_error) if algorithm in [:second_order, :pruned_second_order] return SS_and_pars[1:length(๐“‚.constants.post_model_macro.var)], zeros(length(๐“‚.constants.post_model_macro.var),2), spzeros(length(๐“‚.constants.post_model_macro.var),2), false elseif algorithm in [:third_order, :pruned_third_order] @@ -1960,22 +1945,17 @@ function get_solution(๐“‚::โ„ณ, end end - โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - # Ensure QME workspace - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) + โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; opts = opts, - initial_guess = ๐“‚.caches.qme_solution) + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = parameters) - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) - - if solved ๐“‚.caches.qme_solution = qme_sol end + update_perturbation_counter!(๐“‚.counters, solved, estimation = estimation, order = 1) if !solved if algorithm in [:second_order, :pruned_second_order] @@ -1988,58 +1968,35 @@ function get_solution(๐“‚::โ„ณ, end if algorithm in [:second_order, :pruned_second_order] - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ + โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces) - ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; + ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; initial_guess = ๐“‚.caches.second_order_solution, - opts = opts) + opts = opts, parameter_values = parameters) - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) - - if eltype(๐’โ‚‚) == Float64 && solved2 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end - - ๐’โ‚‚ *= ๐“‚.constants.second_order.๐”โ‚‚ - - if !(typeof(๐’โ‚‚) <: AbstractSparseMatrix) - ๐’โ‚‚ = sparse(๐’โ‚‚) # * ๐“‚.constants.second_order.๐”โ‚‚) - end + update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) return SS_and_pars[1:length(๐“‚.constants.post_model_macro.var)], ๐’โ‚, ๐’โ‚‚, true elseif algorithm in [:third_order, :pruned_third_order] - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ + โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces) - ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; + ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; initial_guess = ๐“‚.caches.second_order_solution, - opts = opts) + opts = opts, parameter_values = parameters) - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) - - if eltype(๐’โ‚‚) == Float64 && solved2 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end - - ๐’โ‚‚ *= ๐“‚.constants.second_order.๐”โ‚‚ - - if !(typeof(๐’โ‚‚) <: AbstractSparseMatrix) - ๐’โ‚‚ = sparse(๐’โ‚‚) # * ๐“‚.constants.second_order.๐”โ‚‚) - end + update_perturbation_counter!(๐“‚.counters, solved2, estimation = estimation, order = 2) - โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives)# * ๐“‚.constants.third_order.๐”โˆ‡โ‚ƒ + โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces) - ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, - ๐’โ‚, ๐’โ‚‚, - ๐“‚.constants, - ๐“‚.workspaces; - initial_guess = ๐“‚.caches.third_order_solution, - opts = opts) - - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved3, estimation = estimation, order = 3) + ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, + ๐’โ‚, ๐’โ‚‚, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, parameter_values = parameters) - if eltype(๐’โ‚ƒ) == Float64 && solved3 ๐“‚.caches.third_order_solution = ๐’โ‚ƒ end - - ๐’โ‚ƒ *= ๐“‚.constants.third_order.๐”โ‚ƒ - - if !(typeof(๐’โ‚ƒ) <: AbstractSparseMatrix) - ๐’โ‚ƒ = sparse(๐’โ‚ƒ) # * ๐“‚.constants.third_order.๐”โ‚ƒ) - end + update_perturbation_counter!(๐“‚.counters, solved3, estimation = estimation, order = 3) return SS_and_pars[1:length(๐“‚.constants.post_model_macro.var)], ๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ, true else @@ -2159,23 +2116,18 @@ function get_conditional_variance_decomposition(๐“‚::โ„ณ; SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) - โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - # Ensure QME workspace - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) + โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix ๐‘บโ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; opts = opts, - initial_guess = ๐“‚.caches.qme_solution) + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = ๐“‚.parameter_values) update_perturbation_counter!(๐“‚.counters, solved, order = 1) - if solved ๐“‚.caches.qme_solution = qme_sol end - A = @views ๐‘บโ‚[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * โ„’.diagm(ones(๐“‚.constants.post_model_macro.nVars))[indexin(๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,1:๐“‚.constants.post_model_macro.nVars),:] sort!(periods) @@ -2196,12 +2148,11 @@ function get_conditional_variance_decomposition(๐“‚::โ„ณ; end if Inf in periods # Ensure lyapunov workspace is properly sized and get it - lyap_ws = ensure_lyapunov_workspace_1st_order!(๐“‚) + lyap_ws = ensure_lyapunov_workspace!(๐“‚.workspaces, ๐“‚.constants.post_model_macro.nVars, :first_order) covar_raw, _ = solve_lyapunov_equation(A, CC, lyap_ws, lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, + tol = opts.tol.first_order.lyapunov, verbose = opts.verbose) var_container[:,i,indexin(Inf,periods)] = โ„’.diag(covar_raw) # numerically more stable @@ -2210,7 +2161,7 @@ function get_conditional_variance_decomposition(๐“‚::โ„ณ; sum_var_container = max.(sum(var_container, dims=2),eps()) - var_container[var_container .< opts.tol.lyapunov_acceptance_tol] .= 0 + var_container[var_container .< opts.tol.first_order.lyapunov.acceptance_tol] .= 0 cond_var_decomp = var_container ./ sum_var_container @@ -2328,23 +2279,18 @@ function get_variance_decomposition(๐“‚::โ„ณ; SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) - โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - # Ensure QME workspace - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) + โˆ‡โ‚ = calculate_jacobian(๐“‚.parameter_values, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix sol, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; opts = opts, - initial_guess = ๐“‚.caches.qme_solution) + initial_guess = ๐“‚.caches.qme_solution, + parameter_values = ๐“‚.parameter_values) update_perturbation_counter!(๐“‚.counters, solved, order = 1) - if solved ๐“‚.caches.qme_solution = qme_sol end - variances_by_shock = zeros(๐“‚.constants.post_model_macro.nVars, ๐“‚.constants.post_model_macro.nExo) A = @views sol[:, 1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * โ„’.diagm(ones(๐“‚.constants.post_model_macro.nVars))[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,:] @@ -2355,12 +2301,11 @@ function get_variance_decomposition(๐“‚::โ„ณ; CC = C * C' # Ensure lyapunov workspace is properly sized and get it - lyap_ws = ensure_lyapunov_workspace_1st_order!(๐“‚) + lyap_ws = ensure_lyapunov_workspace!(๐“‚.workspaces, ๐“‚.constants.post_model_macro.nVars, :first_order) covar_raw, _ = solve_lyapunov_equation(A, CC, lyap_ws, lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, + tol = opts.tol.first_order.lyapunov, verbose = opts.verbose) variances_by_shock[:,i] = โ„’.diag(covar_raw) @@ -2368,7 +2313,7 @@ function get_variance_decomposition(๐“‚::โ„ณ; sum_variances_by_shock = max.(sum(variances_by_shock, dims=2), eps()) - variances_by_shock[variances_by_shock .< opts.tol.lyapunov_acceptance_tol] .= 0 + variances_by_shock[variances_by_shock .< opts.tol.first_order.lyapunov.acceptance_tol] .= 0 var_decomp = variances_by_shock ./ sum_variances_by_shock @@ -2479,7 +2424,7 @@ function get_correlation(๐“‚::โ„ณ; @assert solved "Could not find covariance matrix." end - covar_dcmp[abs.(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol] .= 0 + covar_dcmp[abs.(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol] .= 0 std = sqrt.(max.(โ„’.diag(covar_dcmp),eps(Float64))) @@ -2592,7 +2537,7 @@ function get_autocorrelation(๐“‚::โ„ณ; opts = opts, autocorrelation_periods = autocorrelation_periods) - autocorr[โ„’.diag(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol,:] .= 0 + autocorr[โ„’.diag(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol,:] .= 0 elseif algorithm == :pruned_second_order covar_dcmp, ฮฃแถปโ‚‚, state_ฮผ, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, solved = calculate_second_order_moments_with_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) @@ -2600,14 +2545,14 @@ function get_autocorrelation(๐“‚::โ„ณ; autocorr = zeros(size(covar_dcmp,1),length(autocorrelation_periods)) - covar_dcmp[abs.(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol] .= 0 + covar_dcmp[abs.(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol] .= 0 for i in autocorrelation_periods autocorr[:,i] .= โ„’.diag(sฬ‚_to_yโ‚‚ * sฬ‚_to_sฬ‚โ‚‚โฑ * autocorr_tmp) ./ โ„’.diag(covar_dcmp) sฬ‚_to_sฬ‚โ‚‚โฑ *= sฬ‚_to_sฬ‚โ‚‚ end - autocorr[โ„’.diag(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol,:] .= 0 + autocorr[โ„’.diag(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol,:] .= 0 else covar_dcmp, sol, _, SS_and_pars, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) @@ -2617,7 +2562,7 @@ function get_autocorrelation(๐“‚::โ„ณ; autocorr = reduce(hcat,[โ„’.diag(A ^ i * covar_dcmp ./ โ„’.diag(covar_dcmp)) for i in autocorrelation_periods]) - autocorr[โ„’.diag(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol,:] .= 0 + autocorr[โ„’.diag(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol,:] .= 0 end @@ -2632,13 +2577,13 @@ end """ See [`get_autocorrelation`](@ref) """ -get_autocorr = get_autocorrelation +get_autocorr(args...; kwargs...) = get_autocorrelation(args...; kwargs...) """ See [`get_autocorrelation`](@ref) """ -autocorr = get_autocorrelation +autocorr(args...; kwargs...) = get_autocorrelation(args...; kwargs...) @@ -2785,9 +2730,9 @@ function get_moments(๐“‚::โ„ณ; length_par = length(parameter_derivatives) end - NSSS, (solution_error, iters) = ๐“‚.caches.outdated.non_stochastic_steady_state ? get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) : (copy(๐“‚.caches.non_stochastic_steady_state), (eps(), 0)) + NSSS, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values, opts = opts) - @assert solution_error < tol.NSSS_acceptance_tol "Could not find non-stochastic steady state." + @assert solution_error < tol.nsss.acceptance_tol "Could not find non-stochastic steady state." if length_par * length(NSSS) > 200 && derivatives @info "Most of the time is spent calculating derivatives wrt parameters. If they are not needed, add `derivatives = false` as an argument to the function call." maxlog = DEFAULT_MAXLOG @@ -2825,8 +2770,16 @@ function get_moments(๐“‚::โ„ณ; axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] end - # dNSSS = ๐’œ.jacobian(๐’ท(), x -> collect(SS_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose)[1]), ๐“‚.parameter_values[param_idx])[1] - dNSSS = ๐’Ÿ.jacobian(x -> get_NSSS_and_parameters(๐“‚, x, opts = opts)[1], backend, ๐“‚.parameter_values)[:,param_idx] + (nsss_d_result, nsss_d_pb) = rrule(get_NSSS_and_parameters, ๐“‚, ๐“‚.parameter_values, opts = opts) + n_ss_full = length(nsss_d_result[1]) + np = length(๐“‚.parameter_values) + dNSSS = zeros(n_ss_full, np) + for j in 1:n_ss_full + โˆ‚ss = zeros(n_ss_full); โˆ‚ss[j] = 1.0 + โˆ‚p = nsss_d_pb((โˆ‚ss, NoTangent()))[3] + if !(โˆ‚p isa AbstractZero); dNSSS[j, :] .= โˆ‚p; end + end + dNSSS = dNSSS[:, param_idx] if length(๐“‚.equations.calibration_parameters) > 0 var_idx_ext = vcat(var_idx, ๐“‚.constants.post_model_macro.nVars .+ (1:length(๐“‚.equations.calibration_parameters))) @@ -2834,7 +2787,6 @@ function get_moments(๐“‚::โ„ณ; var_idx_ext = var_idx end - # dNSSS = ๐’œ.jacobian(๐’ท(), x->๐“‚.functions.NSSS_solve(x, ๐“‚),๐“‚.parameter_values) SS = KeyedArray(hcat(collect(NSSS[var_idx_ext]),dNSSS[var_idx_ext,:]); Variables = axis1, Steady_state_and_โˆ‚steady_stateโˆ‚parameter = axis2) end @@ -2845,6 +2797,40 @@ function get_moments(๐“‚::โ„ณ; axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] end + # Hoist covariance rrule call for shared use across variance/std_dev/covariance + if variance || standard_deviation || covariance + if algorithm == :pruned_second_order + _cov_result, _cov_pb = rrule(calculate_second_order_moments_with_covariance, ๐“‚.parameter_values, ๐“‚, opts = opts) + covar_dcmp = _cov_result[1] + _n_cov_tuple = 15 + elseif algorithm == :pruned_third_order + _cov_obs = covariance ? :full_covar : variables + _cov_result, _cov_pb = rrule(calculate_third_order_moments, ๐“‚.parameter_values, _cov_obs, ๐“‚, opts = opts) + covar_dcmp = _cov_result[1] + _n_cov_tuple = 4 + else + _cov_result, _cov_pb = rrule(calculate_covariance, ๐“‚.parameter_values, ๐“‚, opts = opts) + covar_dcmp = _cov_result[1] + @assert _cov_result[5] "Could not find covariance matrix." + _n_cov_tuple = 5 + end + + # Compute variance Jacobian via VJP (shared by variance & std_dev) + if variance || standard_deviation + _np_cov = length(๐“‚.parameter_values) + _nv_cov = size(covar_dcmp, 1) + _dvariance_full = zeros(_nv_cov, _np_cov) + for j in 1:_nv_cov + if covar_dcmp[j,j] > eps(Float64) + โˆ‚ฮฃ = zeros(_nv_cov, _nv_cov); โˆ‚ฮฃ[j,j] = 1.0 + seed = ntuple(k -> k == 1 ? โˆ‚ฮฃ : NoTangent(), _n_cov_tuple) + โˆ‚p = _cov_pb(seed)[2] + if !(โˆ‚p isa AbstractZero); _dvariance_full[j,:] .= โˆ‚p; end + end + end + end + end + if variance axis2 = vcat(:Variance, ๐“‚.constants.post_complete_parameters.parameters[param_idx]) @@ -2853,29 +2839,9 @@ function get_moments(๐“‚::โ„ณ; axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] end - if algorithm == :pruned_second_order - covar_dcmp, ฮฃแถปโ‚‚, state_ฮผ, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, solved = calculate_second_order_moments_with_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - # dvariance = ๐’œ.jacobian(๐’ท(), x -> covariance_parameter_derivatives_second_order(x, param_idx, ๐“‚, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, verbose = verbose), ๐“‚.parameter_values[param_idx])[1] - dvariance = ๐’Ÿ.jacobian(x -> max.(โ„’.diag(calculate_second_order_moments_with_covariance(x, ๐“‚, opts = opts)[1]),eps(Float64)), backend, ๐“‚.parameter_values)[:,param_idx] - elseif algorithm == :pruned_third_order - covar_dcmp, state_ฮผ, _, solved = calculate_third_order_moments(๐“‚.parameter_values, variables, ๐“‚, opts = opts) - - # dvariance = ๐’œ.jacobian(๐’ท(), x -> covariance_parameter_derivatives_third_order(x, variables, param_idx, ๐“‚, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, verbose = verbose), ๐“‚.parameter_values[param_idx])[1] - dvariance = ๐’Ÿ.jacobian(x -> max.(โ„’.diag(calculate_third_order_moments(x, variables, ๐“‚, opts = opts)[1]),eps(Float64)), backend, ๐“‚.parameter_values)[:,param_idx] - else - covar_dcmp, ___, __, _, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - @assert solved "Could not find covariance matrix." - - # dvariance = ๐’œ.jacobian(๐’ท(), x -> covariance_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose, lyapunov_algorithm = lyapunov_algorithm), ๐“‚.parameter_values[param_idx])[1] - dvariance = ๐’Ÿ.jacobian(x -> max.(โ„’.diag(calculate_covariance(x, ๐“‚, opts = opts)[1]),eps(Float64)), backend, ๐“‚.parameter_values)[:,param_idx] - end + dvariance = _dvariance_full[:, param_idx] vari = convert(Vector{Real},max.(โ„’.diag(covar_dcmp),eps(Float64))) - - # dvariance = ๐’œ.jacobian(๐’ท(), x-> convert(Vector{Number},max.(โ„’.diag(calculate_covariance(x, ๐“‚)),eps(Float64))), Float64.(๐“‚.parameter_values)) - varrs = KeyedArray(hcat(vari[var_idx],dvariance[var_idx,:]); Variables = axis1, Variance_and_โˆ‚varianceโˆ‚parameter = axis2) @@ -2888,17 +2854,8 @@ function get_moments(๐“‚::โ„ณ; end standard_dev = sqrt.(convert(Vector{Real},max.(โ„’.diag(covar_dcmp),eps(Float64)))) - - if algorithm == :pruned_second_order - # dst_dev = ๐’œ.jacobian(๐’ท(), x -> sqrt.(covariance_parameter_derivatives_second_order(x, param_idx, ๐“‚, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, verbose = verbose)), ๐“‚.parameter_values[param_idx])[1] - dst_dev = ๐’Ÿ.jacobian(x -> sqrt.(max.(โ„’.diag(calculate_second_order_moments_with_covariance(x, ๐“‚, opts = opts)[1]),eps(Float64))), backend, ๐“‚.parameter_values)[:,param_idx] - elseif algorithm == :pruned_third_order - # dst_dev = ๐’œ.jacobian(๐’ท(), x -> sqrt.(covariance_parameter_derivatives_third_order(x, variables, param_idx, ๐“‚, lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose)), ๐“‚.parameter_values[param_idx])[1] - dst_dev = ๐’Ÿ.jacobian(x -> sqrt.(max.(โ„’.diag(calculate_third_order_moments(x, variables, ๐“‚, opts = opts)[1]),eps(Float64))), backend, ๐“‚.parameter_values)[:,param_idx] - else - # dst_dev = ๐’œ.jacobian(๐’ท(), x -> sqrt.(covariance_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose, lyapunov_algorithm = lyapunov_algorithm)), ๐“‚.parameter_values[param_idx])[1] - dst_dev = ๐’Ÿ.jacobian(x -> sqrt.(max.(โ„’.diag(calculate_covariance(x, ๐“‚, opts = opts)[1]),eps(Float64))), backend, ๐“‚.parameter_values)[:,param_idx] - end + # Analytical: d(sqrt(v))/d(params) = dv/d(params) / (2*sqrt(v)) + dst_dev = _dvariance_full[:, param_idx] ./ (2 .* standard_dev) st_dev = KeyedArray(hcat(standard_dev[var_idx], dst_dev[var_idx, :]); Variables = axis1, Standard_deviation_and_โˆ‚standard_deviationโˆ‚parameter = axis2) end @@ -2912,26 +2869,9 @@ function get_moments(๐“‚::โ„ณ; axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] end - if algorithm == :pruned_second_order - covar_dcmp, ฮฃแถปโ‚‚, state_ฮผ, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, solved = calculate_second_order_moments_with_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - # dst_dev = ๐’œ.jacobian(๐’ท(), x -> sqrt.(covariance_parameter_derivatives_second_order(x, param_idx, ๐“‚, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, verbose = verbose)), ๐“‚.parameter_values[param_idx])[1] - dst_dev = ๐’Ÿ.jacobian(x -> sqrt.(max.(โ„’.diag(calculate_second_order_moments_with_covariance(x, ๐“‚, opts = opts)[1]),eps(Float64))), backend, ๐“‚.parameter_values)[:,param_idx] - elseif algorithm == :pruned_third_order - covar_dcmp, state_ฮผ, _, solved = calculate_third_order_moments(๐“‚.parameter_values, variables, ๐“‚, opts = opts) - - # dst_dev = ๐’œ.jacobian(๐’ท(), x -> sqrt.(covariance_parameter_derivatives_third_order(x, variables, param_idx, ๐“‚, lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose)), ๐“‚.parameter_values[param_idx])[1] - dst_dev = ๐’Ÿ.jacobian(x -> sqrt.(max.(โ„’.diag(calculate_third_order_moments(x, variables, ๐“‚, opts = opts)[1]),eps(Float64))), backend, ๐“‚.parameter_values)[:,param_idx] - else - covar_dcmp, ___, __, _, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - @assert solved "Could not find covariance matrix." - - # dst_dev = ๐’œ.jacobian(๐’ท(), x -> sqrt.(covariance_parameter_derivatives(x, param_idx, ๐“‚, verbose = verbose, lyapunov_algorithm = lyapunov_algorithm)), ๐“‚.parameter_values[param_idx])[1] - dst_dev = ๐’Ÿ.jacobian(x -> sqrt.(max.(โ„’.diag(calculate_covariance(x, ๐“‚, opts = opts)[1]),eps(Float64))), backend, ๐“‚.parameter_values)[:,param_idx] - end - standard_dev = sqrt.(convert(Vector{Real},max.(โ„’.diag(covar_dcmp),eps(Float64)))) + # Analytical: d(sqrt(v))/d(params) = dv/d(params) / (2*sqrt(v)) + dst_dev = _dvariance_full[:, param_idx] ./ (2 .* standard_dev) st_dev = KeyedArray(hcat(standard_dev[var_idx], dst_dev[var_idx, :]); Variables = axis1, Standard_deviation_and_โˆ‚standard_deviationโˆ‚parameter = axis2) end @@ -2945,24 +2885,19 @@ function get_moments(๐“‚::โ„ณ; axis3 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis3_decomposed] end - if algorithm == :pruned_second_order - covar_dcmp, ฮฃแถปโ‚‚, state_ฮผ, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, solved = calculate_second_order_moments_with_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - # Compute covariance derivatives - dcovariance = ๐’Ÿ.jacobian(x -> vec(calculate_second_order_moments_with_covariance(x, ๐“‚, opts = opts)[1]), backend, ๐“‚.parameter_values)[:,param_idx] - elseif algorithm == :pruned_third_order - covar_dcmp, state_ฮผ, _, solved = calculate_third_order_moments(๐“‚.parameter_values, :full_covar, ๐“‚, opts = opts) - - # Compute covariance derivatives - dcovariance = ๐’Ÿ.jacobian(x -> vec(calculate_third_order_moments(x, :full_covar, ๐“‚, opts = opts)[1]), backend, ๐“‚.parameter_values)[:,param_idx] - else - covar_dcmp, ___, __, _, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - @assert solved "Could not find covariance matrix." - - # Compute covariance derivatives - dcovariance = ๐’Ÿ.jacobian(x -> vec(calculate_covariance(x, ๐“‚, opts = opts)[1]), backend, ๐“‚.parameter_values)[:,param_idx] + # Compute full covariance Jacobian via VJP from hoisted rrule + _np_cov2 = length(๐“‚.parameter_values) + _nv_cov2 = size(covar_dcmp, 1) + dcovariance = zeros(_nv_cov2 * _nv_cov2, _np_cov2) + for j in 1:(_nv_cov2 * _nv_cov2) + r = mod1(j, _nv_cov2) + c = div(j - 1, _nv_cov2) + 1 + โˆ‚ฮฃ = zeros(_nv_cov2, _nv_cov2); โˆ‚ฮฃ[r,c] = 1.0 + seed = ntuple(k -> k == 1 ? โˆ‚ฮฃ : NoTangent(), _n_cov_tuple) + โˆ‚p = _cov_pb(seed)[2] + if !(โˆ‚p isa AbstractZero); dcovariance[j,:] .= โˆ‚p; end end + dcovariance = dcovariance[:, param_idx] end if mean && algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] @@ -2973,12 +2908,20 @@ function get_moments(๐“‚::โ„ณ; axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] end - state_ฮผ, solved = calculate_mean(๐“‚.parameter_values, ๐“‚, algorithm = algorithm, opts = opts) + (mean_result, mean_pb) = rrule(calculate_mean, ๐“‚.parameter_values, ๐“‚, algorithm = algorithm, opts = opts) + state_ฮผ = mean_result[1] - @assert solved "Mean not found." - - # state_ฮผ_dev = ๐’œ.jacobian(๐’ท(), x -> mean_parameter_derivatives(x, param_idx, ๐“‚, algorithm = algorithm, verbose = verbose, sylvester_algorithm = sylvester_algorithm), ๐“‚.parameter_values[param_idx])[1] - state_ฮผ_dev = ๐’Ÿ.jacobian(x -> calculate_mean(x, ๐“‚, algorithm = algorithm, opts = opts)[1], backend, ๐“‚.parameter_values)[:,param_idx] + @assert mean_result[2] "Mean not found." + + n_mean = length(state_ฮผ) + np_mean = length(๐“‚.parameter_values) + state_ฮผ_dev = zeros(n_mean, np_mean) + for j in 1:n_mean + โˆ‚mean = zeros(n_mean); โˆ‚mean[j] = 1.0 + โˆ‚p = mean_pb((โˆ‚mean, NoTangent()))[2] + if !(โˆ‚p isa AbstractZero); state_ฮผ_dev[j,:] .= โˆ‚p; end + end + state_ฮผ_dev = state_ฮผ_dev[:, param_idx] var_means = KeyedArray(hcat(state_ฮผ[var_idx], state_ฮผ_dev[var_idx, :]); Variables = axis1, Mean_and_โˆ‚meanโˆ‚parameter = axis2) end @@ -3014,7 +2957,9 @@ function get_moments(๐“‚::โ„ณ; if mean && !(variance || standard_deviation || covariance) state_ฮผ, solved = calculate_mean(๐“‚.parameter_values, ๐“‚, algorithm = algorithm, opts = opts) - @assert solved "Mean not found." + if !solved + @warn "Mean not found." + end var_means = KeyedArray(state_ฮผ[var_idx]; Variables = axis1) end @@ -3032,18 +2977,18 @@ function get_moments(๐“‚::โ„ณ; end else covar_dcmp, ___, __, _, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - @assert solved "Could not find covariance matrix." if mean && algorithm == :first_order var_means = KeyedArray(collect(NSSS)[var_idx]; Variables = ๐“‚.constants.post_model_macro.var[var_idx]) end end - varr = convert(Vector{Real},max.(โ„’.diag(covar_dcmp),eps(Float64))) + if !solved + @warn "Could not find covariance matrix." + end + varr = convert(Vector{Real},max.(โ„’.diag(covar_dcmp),eps(Float64))) varrs = KeyedArray(varr[var_idx]; Variables = axis1) - if standard_deviation st_dev = KeyedArray(sqrt.(varr)[var_idx]; Variables = axis1) end @@ -3062,13 +3007,16 @@ function get_moments(๐“‚::โ„ณ; end else covar_dcmp, ___, __, _, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - @assert solved "Could not find covariance matrix." if mean && algorithm == :first_order var_means = KeyedArray(collect(NSSS)[var_idx]; Variables = ๐“‚.constants.post_model_macro.var[var_idx]) end end + + if !solved + @warn "Could not find covariance matrix." + end + st_dev = KeyedArray(sqrt.(convert(Vector{Real},max.(โ„’.diag(covar_dcmp),eps(Float64))))[var_idx]; Variables = axis1) end @@ -3085,12 +3033,14 @@ function get_moments(๐“‚::โ„ณ; end else covar_dcmp, ___, __, _, solved = calculate_covariance(๐“‚.parameter_values, ๐“‚, opts = opts) - - @assert solved "Could not find covariance matrix." if mean && algorithm == :first_order var_means = KeyedArray(collect(NSSS)[var_idx]; Variables = ๐“‚.constants.post_model_macro.var[var_idx]) end + + if !solved + @warn "Could not find covariance matrix." + end end end end @@ -3312,7 +3262,7 @@ Dict{Symbol, AbstractArray{Float64}} with 1 entry: :covariance => [...4x4 matrix with c-k covariances filled, y-i covariances filled, and cross-group elements set to zero...] ``` """ -function get_statistics(๐“‚, +function get_statistics(๐“‚::โ„ณ, parameter_values::Vector{T}; parameters::Union{Vector{Symbol},Vector{String}} = ๐“‚.constants.post_complete_parameters.parameters, steady_state_function::SteadyStateFunctionType = missing, @@ -3342,25 +3292,25 @@ function get_statistics(๐“‚, @assert !(non_stochastic_steady_state == Symbol[]) || !(standard_deviation == Symbol[]) || !(mean == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[]) "Provide variables for at least one output." - SS_var_idx = @ignore_derivatives parse_variables_input_to_index(non_stochastic_steady_state, ๐“‚) + SS_var_idx = parse_variables_input_to_index(non_stochastic_steady_state, ๐“‚) - mean_var_idx = @ignore_derivatives parse_variables_input_to_index(mean, ๐“‚) + mean_var_idx = parse_variables_input_to_index(mean, ๐“‚) - std_var_idx = @ignore_derivatives parse_variables_input_to_index(standard_deviation, ๐“‚) + std_var_idx = parse_variables_input_to_index(standard_deviation, ๐“‚) - var_var_idx = @ignore_derivatives parse_variables_input_to_index(variance, ๐“‚) + var_var_idx = parse_variables_input_to_index(variance, ๐“‚) - covar_var_idx = @ignore_derivatives parse_variables_input_to_index(covariance, ๐“‚) + covar_var_idx = parse_variables_input_to_index(covariance, ๐“‚) # Parse covariance groups if input is grouped format - covar_groups = @ignore_derivatives is_grouped_covariance_input(covariance) ? parse_covariance_groups(covariance, ๐“‚.constants) : nothing + covar_groups = is_grouped_covariance_input(covariance) ? parse_covariance_groups(covariance, ๐“‚.constants) : nothing - autocorr_var_idx = @ignore_derivatives parse_variables_input_to_index(autocorrelation, ๐“‚) + autocorr_var_idx = parse_variables_input_to_index(autocorrelation, ๐“‚) - other_parameter_values = @ignore_derivatives ๐“‚.parameter_values[indexin(setdiff(๐“‚.constants.post_complete_parameters.parameters, parameters), ๐“‚.constants.post_complete_parameters.parameters)] + other_parameter_values = ๐“‚.parameter_values[indexin(setdiff(๐“‚.constants.post_complete_parameters.parameters, parameters), ๐“‚.constants.post_complete_parameters.parameters)] - sort_idx = @ignore_derivatives sortperm(vcat(indexin(setdiff(๐“‚.constants.post_complete_parameters.parameters, parameters), ๐“‚.constants.post_complete_parameters.parameters), indexin(parameters, ๐“‚.constants.post_complete_parameters.parameters))) + sort_idx = sortperm(vcat(indexin(setdiff(๐“‚.constants.post_complete_parameters.parameters, parameters), ๐“‚.constants.post_complete_parameters.parameters), indexin(parameters, ๐“‚.constants.post_complete_parameters.parameters))) all_parameters = vcat(other_parameter_values, parameter_values)[sort_idx] @@ -3370,10 +3320,10 @@ function get_statistics(๐“‚, algorithm = :pruned_second_order end - @ignore_derivatives solve!(๐“‚, - algorithm = algorithm, - steady_state_function = steady_state_function, - opts = opts) + solve!(๐“‚, + algorithm = algorithm, + steady_state_function = steady_state_function, + opts = opts) if !(non_stochastic_steady_state == Symbol[]) && (standard_deviation == Symbol[]) && (variance == Symbol[]) && (covariance == Symbol[]) && (autocorrelation == Symbol[]) SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, all_parameters, opts = opts) # timer = timer, @@ -3382,7 +3332,7 @@ function get_statistics(๐“‚, ret = Dict{Symbol,AbstractArray{T}}() - ret[:non_stochastic_steady_state] = solution_error < opts.tol.NSSS_acceptance_tol ? SS[SS_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(SS_var_idx) ? 0 : length(SS_var_idx)) + ret[:non_stochastic_steady_state] = solution_error < opts.tol.nsss.acceptance_tol ? SS[SS_var_idx] : fill(Inf * sum(abs2,parameter_values), isnothing(SS_var_idx) ? 0 : length(SS_var_idx)) return ret end @@ -3435,13 +3385,13 @@ function get_statistics(๐“‚, sฬ‚_to_sฬ‚โ‚‚โฑ *= sฬ‚_to_sฬ‚โ‚‚ end - autocorr[โ„’.diag(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol,:] .= 0 + autocorr[โ„’.diag(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol,:] .= 0 elseif !(algorithm == :pruned_third_order) A = @views sol[:,1:๐“‚.constants.post_model_macro.nPast_not_future_and_mixed] * โ„’.diagm(ones(๐“‚.constants.post_model_macro.nVars))[๐“‚.constants.post_model_macro.past_not_future_and_mixed_idx,:] autocorr = reduce(hcat,[โ„’.diag(A ^ i * covar_dcmp ./ max.(โ„’.diag(covar_dcmp),eps(Float64))) for i in autocorrelation_periods]) - autocorr[โ„’.diag(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol,:] .= 0 + autocorr[โ„’.diag(covar_dcmp) .< opts.tol.first_order.lyapunov.acceptance_tol,:] .= 0 end end @@ -3612,26 +3562,26 @@ function get_loglikelihood(๐“‚::โ„ณ, # checks to avoid errors further down the line and inform the user @assert initial_covariance โˆˆ [:theoretical, :diagonal] "Invalid method to initialise the Kalman filters covariance matrix. Supported methods are: the theoretical long run values (option `:theoretical`) or large values (10.0) along the diagonal (option `:diagonal`)." - filter, _, algorithm, _, _, warmup_iterations = @ignore_derivatives normalize_filtering_options(filter, false, algorithm, false, warmup_iterations) + filter, _, algorithm, _, _, warmup_iterations = normalize_filtering_options(filter, false, algorithm, false, warmup_iterations) - observables = @ignore_derivatives get_and_check_observables(๐“‚.constants.post_model_macro, data) + observables = get_and_check_observables(๐“‚.constants.post_model_macro, data) - @ignore_derivatives solve!(๐“‚, - opts = opts, - steady_state_function = steady_state_function, - # timer = timer, - algorithm = algorithm) + solve!(๐“‚, + opts = opts, + steady_state_function = steady_state_function, + # timer = timer, + algorithm = algorithm) - bounds_violated = @ignore_derivatives check_bounds(parameter_values, ๐“‚) + bounds_violated = check_bounds(parameter_values, ๐“‚) if bounds_violated # println("Bounds violated") return on_failure_loglikelihood end - NSSS_labels = @ignore_derivatives [sort(union(๐“‚.constants.post_model_macro.exo_present, ๐“‚.constants.post_model_macro.var))..., ๐“‚.equations.calibration_parameters...] + SS_and_pars_names = ๐“‚.constants.post_complete_parameters.SS_and_pars_names - obs_indices = @ignore_derivatives convert(Vector{Int}, indexin(observables, NSSS_labels)) + obs_indices = convert(Vector{Int}, indexin(observables, SS_and_pars_names)) # @timeit_debug timer "Get relevant steady state and solution" begin @@ -3646,27 +3596,30 @@ function get_loglikelihood(๐“‚::โ„ณ, end if collect(axiskeys(data,1)) isa Vector{String} - data = @ignore_derivatives rekey(data, 1 => axiskeys(data,1) .|> Meta.parse .|> replace_indices) + data = rekey(data, 1 => axiskeys(data,1) .|> Meta.parse .|> replace_indices) end - dt = @ignore_derivatives collect(data(observables)) + dt = collect(data(observables)) # prepare data data_in_deviations = dt .- SS_and_pars[obs_indices] # @timeit_debug timer "Filter" begin - # Ensure lyapunov workspace for Kalman filter initial covariance - lyap_ws = @ignore_derivatives ensure_lyapunov_workspace_1st_order!(๐“‚) - - # Ensure inversion workspace if using inversion filter - third_order = algorithm in (:pruned_third_order, :third_order) - inv_ws = @ignore_derivatives ensure_inversion_workspace!(๐“‚; third_order = third_order) - - # Ensure kalman workspace for Kalman filter iterations - kalman_ws = @ignore_derivatives ensure_kalman_workspace!(๐“‚) - - llh = calculate_loglikelihood(Val(filter), algorithm, observables, ๐’, data_in_deviations, constants_obj, presample_periods, initial_covariance, state, warmup_iterations, filter_algorithm, opts, on_failure_loglikelihood, lyap_ws, inv_ws, kalman_ws) # timer = timer + llh = calculate_loglikelihood(Val(filter), + Val(algorithm), + obs_indices, + ๐’, + data_in_deviations, + constants_obj, + state, + ๐“‚.workspaces, + warmup_iterations = warmup_iterations, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + filter_algorithm = filter_algorithm, + opts = opts, + on_failure_loglikelihood = on_failure_loglikelihood) # timer = timer # end # timeit_debug @@ -3773,7 +3726,7 @@ function get_non_stochastic_steady_state_residuals(๐“‚::โ„ณ, combined_values[key] = value end elseif isa(values, KeyedArray) - for (key, value) in Dict(axiskeys(values, 1) .=> collect(values)) + for (key, value) in zip(axiskeys(values, 1), collect(values)) if key isa String key = replace_indices(key) end @@ -3800,4 +3753,4 @@ get_residuals = get_non_stochastic_steady_state_residuals """ See [`get_non_stochastic_steady_state_residuals`](@ref) """ -check_residuals = get_non_stochastic_steady_state_residuals +check_residuals = get_non_stochastic_steady_state_residuals \ No newline at end of file diff --git a/src/inspect.jl b/src/inspect.jl index b86337858..21c8133fd 100644 --- a/src/inspect.jl +++ b/src/inspect.jl @@ -66,7 +66,7 @@ function replace_curly_braces_in_symbols(expr) result = Expr(:curly, result, content) end - remaining = rest + remaining = something(rest, "") end return result === nothing ? expr : result @@ -123,16 +123,21 @@ end Check if `expr` contains `sym` matching `pattern` (nothing = any timing). """ function expr_contains(expr, sym::Symbol, pattern) + normalize_repr(x) = replace(string(x), "โ—–" => "{", "โ——" => "}") + sym_str = normalize_repr(sym) + pattern_str = pattern === nothing ? "" : normalize_repr(pattern) + found = Ref(false) postwalk(expr) do x if pattern === nothing # Match symbol anywhere (as ref base or standalone) - if x === sym || (x isa Expr && x.head == :ref && x.args[1] === sym) + if normalize_repr(x) == sym_str || + (x isa Expr && x.head == :ref && normalize_repr(x.args[1]) == sym_str) found[] = true end else # Match exact expression pattern - x == pattern && (found[] = true) + normalize_repr(x) == pattern_str && (found[] = true) end x end @@ -348,7 +353,7 @@ function get_dynamic_equations(๐“‚::โ„ณ; filter::Union{Symbol, String, Nothing} # Parse filter term (uses user-friendly format with [-1], [0], etc.) sym, pattern = parse_filter_term(filter) - return [expr for (expr, orig) in zip(exprs, ๐“‚.equations.dynamic) if expr_contains(orig, sym, pattern)] + return [expr for expr in exprs if expr_contains(expr, sym, pattern)] end @@ -579,6 +584,7 @@ get_calibrated_parameters(RBC) """ function get_calibrated_parameters(๐“‚::โ„ณ; values::Bool = false)::Union{Vector{Pair{String, Float64}},Vector{String}} if values + get_NSSS_and_parameters(๐“‚, ๐“‚.parameter_values) return replace.(string.(๐“‚.equations.calibration_parameters), "โ—–" => "{", "โ——" => "}") .=> ๐“‚.caches.non_stochastic_steady_state[๐“‚.constants.post_model_macro.nVars + 1:end] else return replace.(string.(๐“‚.equations.calibration_parameters), "โ—–" => "{", "โ——" => "}")# |> sort diff --git a/src/macros.jl b/src/macros.jl index 6bda65dec..e263c4ab8 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -1,5 +1,84 @@ const all_available_algorithms = [:first_order, :second_order, :pruned_second_order, :third_order, :pruned_third_order] +# Pure-Julia algebraic simplification used in precompile mode (avoids SymPy). +# Applied after converting expressions to steady-state form (time subscripts stripped), +# so that patterns like c[1]/c[0] and (1-l[1])/(1-l[0]) correctly simplify to 1. +# +# Rules applied bottom-up (postwalk order) on the SS-form expression: +# A / A โ†’ 1 +# A * A โ†’ A^2 (not needed here but handled by postwalk recursion) +# A - A โ†’ 0 +# A + 0 โ†’ A, 0 + A โ†’ A +# A - 0 โ†’ A, 0 - A โ†’ -A +# A * 0 โ†’ 0, 0 * A โ†’ 0 +# A * 1 โ†’ A, 1 * A โ†’ A +# A ^ 1 โ†’ A +# A ^ 0 โ†’ 1 +# 1 ^ A โ†’ 1 +# A / 1 โ†’ A +# log(A) - log(A) โ†’ 0 (captured by A-A after postwalk recurse) +# (X/Z) * (Y/Z) โ†’ not simplified (require matching denominators case-by-case) +# +# Also handles the equation-level simplification: +# f(A)/f(C) = f(B)/f(C) โ†’ f(A) = f(B) (cancel common factor on both sides of +# equivalently: A/C - B/C => (A-B)/C an equation written as LHS - RHS = 0) +# +# Note: the function operates on the already-SS-converted expression. All transformations +# are purely syntactic (structural equality) and safe to apply without domain knowledge. +function trivial_simplify(ex) + if !(ex isa Expr) + return ex + end + # Convert to SS form: strips time subscripts l[1],l[0] โ†’ l, shocks โ†’ 0, etc. + ss_ex = convert_to_ss_equation(ex) + # Apply algebraic rewrites bottom-up so simpler forms enable further reductions + result = postwalk(ss_ex) do x + if !(x isa Expr) || x.head != :call + return x + end + f = x.args[1] + nargs = length(x.args) - 1 # number of operands + + if f == :/ && nargs == 2 + num, den = x.args[2], x.args[3] + num == den && return 1 # A / A โ†’ 1 + den == 1 && return num # A / 1 โ†’ A + num == 0 && return 0 # 0 / A โ†’ 0 + + elseif f == :* && nargs == 2 + a, b = x.args[2], x.args[3] + a == 1 && return b # 1 * B โ†’ B + b == 1 && return a # A * 1 โ†’ A + a == 0 || b == 0 && return 0 # 0 * _ or _ * 0 โ†’ 0 + + elseif f == :^ && nargs == 2 + base, exp_ = x.args[2], x.args[3] + exp_ == 1 && return base # A^1 โ†’ A + exp_ == 0 && return 1 # A^0 โ†’ 1 + base == 1 && return 1 # 1^A โ†’ 1 + base == 0 && exp_ isa Number && exp_ > 0 && return 0 # 0^n โ†’ 0 (n>0) + + elseif f == :+ && nargs == 2 + a, b = x.args[2], x.args[3] + b == 0 && return a # A + 0 โ†’ A + a == 0 && return b # 0 + A โ†’ A + + elseif f == :- && nargs == 2 + a, b = x.args[2], x.args[3] + a == b && return 0 # A - A โ†’ 0 + b == 0 && return a # A - 0 โ†’ A + a == 0 && return :(-($b)) # 0 - A โ†’ -A + + elseif f == :- && nargs == 1 + a = x.args[2] + a == 0 && return 0 # -(0) โ†’ 0 + end + + return x + end + return result +end + """ $(SIGNATURES) @@ -86,16 +165,13 @@ macro model(๐“‚,ex...) par_calib_list = [] # NSSS struct fields - NSSS_solve_blocks_in_place = ss_solve_block[] - NSSS_solver_cache = CircularBuffer{Vector{Vector{Float64}}}(500) - NSSS_solve_func = x->x + nsss_solver_cache = CircularBuffer{Vector{Vector{Float64}}}(500) NSSS_check_func = x->x NSSS_custom_function = nothing NSSS_โˆ‚equations_โˆ‚parameters = zeros(0,0) NSSS_โˆ‚equations_โˆ‚parameters_func = x->x NSSS_โˆ‚equations_โˆ‚SS_and_pars = zeros(0,0) NSSS_โˆ‚equations_โˆ‚SS_and_pars_func = x->x - NSSS_dependencies = nothing original_equations = [] calibration_equations = [] @@ -342,7 +418,7 @@ macro model(๐“‚,ex...) x.args[2].head == :call ? # nonnegative expressions begin if precompile - replacement = x.args[2] + replacement = trivial_simplify(x.args[2]) else replacement = simplify(x.args[2]) end @@ -390,7 +466,7 @@ macro model(๐“‚,ex...) x.args[2].head == :call ? # nonnegative expressions begin if precompile - replacement = x.args[2] + replacement = trivial_simplify(x.args[2]) else replacement = simplify(x.args[2]) end @@ -434,7 +510,7 @@ macro model(๐“‚,ex...) x.args[2].head == :call ? # nonnegative expressions begin if precompile - replacement = x.args[2] + replacement = trivial_simplify(x.args[2]) else replacement = simplify(x.args[2]) end @@ -478,7 +554,7 @@ macro model(๐“‚,ex...) x.args[2].head == :call ? # nonnegative expressions begin if precompile - replacement = x.args[2] + replacement = trivial_simplify(x.args[2]) else replacement = simplify(x.args[2]) end @@ -522,7 +598,7 @@ macro model(๐“‚,ex...) x.args[2].head == :call ? # nonnegative expressions begin if precompile - replacement = x.args[2] + replacement = trivial_simplify(x.args[2]) else replacement = simplify(x.args[2]) end @@ -638,13 +714,13 @@ macro model(๐“‚,ex...) if idx โˆˆ ss_equations_with_aux_variables if precompile - ss_aux_equation = Expr(:call,:-,unblock(prs_ex).args[2],unblock(prs_ex).args[3]) + ss_aux_equation = Expr(:call,:-,unblock(prs_ex).args[2],trivial_simplify(unblock(prs_ex).args[3])) else ss_aux_equation = Expr(:call,:-,unblock(prs_ex).args[2],simplify(unblock(prs_ex).args[3])) # simplify RHS if nonnegative auxiliary variable end else if precompile - ss_aux_equation = unblock(prs_ex) + ss_aux_equation = trivial_simplify(unblock(prs_ex)) else ss_aux_equation = simplify(unblock(prs_ex)) end @@ -704,6 +780,7 @@ macro model(๐“‚,ex...) nPresent_but_not_only = length(present_but_not_only) nVars = length(all_vars) nExo = length(collect(exo)) + I_nPast = โ„’.I(nPast_not_future_and_mixed) present_only_idx = indexin(present_only,var) present_but_not_only_idx = indexin(present_but_not_only,var) @@ -782,6 +859,7 @@ macro model(๐“‚,ex...) nMixed, nFuture_not_past_and_mixed, nPast_not_future_and_mixed, + I_nPast, # nPresent_but_not_only, nVars, nExo, @@ -866,25 +944,10 @@ macro model(๐“‚,ex...) # sort(collect($parameters_in_equations)), $parameter_values, - non_stochastic_steady_state( - $NSSS_solve_blocks_in_place, - $NSSS_dependencies - ), - equations($original_equations, $dyn_equations, $ss_equations, $ss_aux_equations, Expr[], $calibration_equations, Expr[], Symbol[]), caches( - outdated_caches( - true, # non_stochastic_steady_state - true, # jacobian - true, # hessian - true, # third_order_derivatives - true, # first_order_solution - true, # second_order_solution - true, # pruned_second_order_solution - true, # third_order_solution - true, # pruned_third_order_solution - ), + valid_for_caches(), zeros(0,0), # jacobian zeros(0,0), # jacobian_parameters zeros(0,0), # jacobian_SS_and_pars @@ -895,6 +958,7 @@ macro model(๐“‚,ex...) zeros(0,0), # third_order_derivatives_parameters zeros(0,0), # third_order_derivatives_SS_and_pars zeros(0,0), # first_order_solution_matrix + zeros(0,0), # first_order_obc_solution_matrix zeros(0,0), # qme_solution Float64[], # second_order_stochastic_steady_state SparseMatrixCSC{Float64, Int64}(โ„’.I,0,0), # second_order_solution @@ -903,9 +967,9 @@ macro model(๐“‚,ex...) SparseMatrixCSC{Float64, Int64}(โ„’.I,0,0), # third_order_solution Float64[], # pruned_third_order_stochastic_steady_state Float64[], # non_stochastic_steady_state - $NSSS_solver_cache, # solver_cache - $NSSS_โˆ‚equations_โˆ‚parameters, # โˆ‚equations_โˆ‚parameters - $NSSS_โˆ‚equations_โˆ‚SS_and_pars, # โˆ‚equations_โˆ‚SS_and_pars + $nsss_solver_cache, # solver + $NSSS_โˆ‚equations_โˆ‚parameters, # NSSS_โˆ‚equations_โˆ‚parameters + $NSSS_โˆ‚equations_โˆ‚SS_and_pars, # NSSS_โˆ‚equations_โˆ‚SS_and_pars ), # (x->x, SparseMatrixCSC{Float64, Int64}(โ„’.I, 0, 0), ๐’Ÿ.prepare_jacobian(x->x, ๐’Ÿ.AutoForwardDiff(), [0]), SparseMatrixCSC{Float64, Int64}(โ„’.I, 0, 0)), # third_order_derivatives # ([], SparseMatrixCSC{Float64, Int64}(โ„’.I, 0, 0)), # model_jacobian @@ -924,24 +988,15 @@ macro model(๐“‚,ex...) $๐“ฆ, model_functions( - $NSSS_solve_func, $NSSS_check_func, $NSSS_custom_function, $NSSS_โˆ‚equations_โˆ‚parameters_func, # NSSS_โˆ‚equations_โˆ‚parameters $NSSS_โˆ‚equations_โˆ‚SS_and_pars_func, # NSSS_โˆ‚equations_โˆ‚SS_and_pars + NSSSSolverFunctions(), + nothing, # nsss_param_prep! jacobian_functions(x->x, x->x, x->x), # jacobian, jacobian_parameters, jacobian_SS_and_pars hessian_functions(x->x, x->x, x->x), # hessian, hessian_parameters, hessian_SS_and_pars third_order_derivatives_functions(x->x, x->x, x->x), # third_order_derivatives, third_order_derivatives_parameters, third_order_derivatives_SS_and_pars - (x,y)->nothing, # first_order_state_update - (x,y)->nothing, # first_order_state_update_obc - (x,y)->nothing, # second_order_state_update - (x,y)->nothing, # second_order_state_update_obc - (x,y)->nothing, # pruned_second_order_state_update - (x,y)->nothing, # pruned_second_order_state_update_obc - (x,y)->nothing, # third_order_state_update - (x,y)->nothing, # third_order_state_update_obc - (x,y)->nothing, # pruned_third_order_state_update - (x,y)->nothing, # pruned_third_order_state_update_obc x->x, # obc_violation false # functions_written ), @@ -971,13 +1026,12 @@ Parameters can be defined in either of the following ways: - expressions containing a target parameter and an equations with endogenous variables in the non-stochastic steady state, and other parameters, or numbers: `k[ss] / (4 * q[ss]) = 1.5 | ฮด` or `ฮฑ | 4 * q[ss] = ฮด * k[ss]` in this case the target parameter will be solved simultaneously with the non-stochastic steady state using the equation defined with it. # Optional arguments to be placed between `๐“‚` and `ex` -- `guess` [Type: `Dict{Symbol, <:Real}, Dict{String, <:Real}}`]: Guess for the non-stochastic steady state. The keys must be the variable (and calibrated parameters) names and the values the guesses. Missing values are filled with standard starting values. +- `guess` [Type: `Dict{Symbol, <:Real}` or `Dict{String, <:Real}`]: Guess for the non-stochastic steady state. The keys must be variable (and calibrated parameter) names and the values the guesses. Missing values are filled with standard starting values. - $STEADY_STATE_FUNCTIONยฎ - `verbose` [Default: `false`, Type: `Bool`]: print more information about how the non-stochastic steady state is solved - `silent` [Default: `false`, Type: `Bool`]: do not print any information -- `symbolic` [Default: `false`, Type: `Bool`]: try to solve the non-stochastic steady state symbolically and fall back to a numerical solution if not possible +- `ss_symbolic_mode` [Default: `:single_equation`, Type: `Symbol`]: controls symbolic steps in non-stochastic steady state (NSSS) setup. Use `:none` for numerical-only setup, `:single_equation` to allow symbolic solves only for single-equation blocks, or `:full` to allow symbolic solves for both single- and multi-equation blocks. - `perturbation_order` [Default: `1`, Type: `Int`]: take derivatives only up to the specified order at this stage. When working with higher order perturbation later on, respective derivatives will be taken at that stage. -- `simplify` [Default: `true`, Type: `Bool`]: whether to eliminate redundant variables and simplify the non-stochastic steady state (NSSS) problem. Setting this to `false` can speed up the process, but might make it harder to find the NSSS. If the model does not parse at all (at step 1 or 2), setting this option to `false` might solve it. - `ss_solver_parameters_algorithm` [Default: `:ESCH`, Type: `Symbol`]: global optimization routine used when searching for steady-state solver parameters after an initial failure; choose `:ESCH` (evolutionary) or `:SAMIN` (simulated annealing). `:SAMIN` is available only when Optim.jl is loaded. - `ss_solver_parameters_maxtime` [Default: `120.0`, Type: `Real`]: time budget in seconds for the steady-state solver parameter search when `ss_solver_parameters_algorithm` is invoked @@ -1057,12 +1111,11 @@ macro parameters(๐“‚,ex...) # parse options verbose = false silent = false - symbolic = false + ss_symbolic_mode = :single_equation precompile = false report_missing_parameters = true perturbation_order = 1 guess = Dict{Symbol,Float64}() - simplify = true steady_state_function = nothing ss_solver_parameters_algorithm = :ESCH ss_solver_parameters_maxtime = 120.0 @@ -1071,8 +1124,8 @@ macro parameters(๐“‚,ex...) postwalk(x -> x isa Expr ? x.head == :(=) ? - (x.args[1] == :symbolic && x.args[2] isa Bool) ? - symbolic = x.args[2] : + (x.args[1] == :ss_symbolic_mode && (x.args[2] isa Symbol || (x.args[2] isa QuoteNode && x.args[2].value isa Symbol))) ? + ss_symbolic_mode = x.args[2] isa QuoteNode ? x.args[2].value : x.args[2] : (x.args[1] == :verbose && x.args[2] isa Bool) ? verbose = x.args[2] : (x.args[1] == :silent && x.args[2] isa Bool) ? @@ -1087,8 +1140,6 @@ macro parameters(๐“‚,ex...) guess = x.args[2] : (x.args[1] == :ss_solver_parameters_algorithm && (x.args[2] isa Symbol || (x.args[2] isa QuoteNode && x.args[2].value isa Symbol))) ? ss_solver_parameters_algorithm = x.args[2] isa QuoteNode ? x.args[2].value : x.args[2] : - (x.args[1] == :simplify && x.args[2] isa Bool) ? - simplify = x.args[2] : (x.args[1] == :steady_state_function && x.args[2] isa Symbol) ? # allow Symbol, anonymous fn, or any callable expr steady_state_function = esc(x.args[2]) : (x.args[1] == :ss_solver_parameters_maxtime && x.args[2] isa Real) ? @@ -1101,6 +1152,8 @@ macro parameters(๐“‚,ex...) x, exp) end + + @assert ss_symbolic_mode โˆˆ [:none, :single_equation, :full] "ss_symbolic_mode must be :none, :single_equation, or :full. Got $ss_symbolic_mode." @assert ss_solver_parameters_algorithm โˆˆ [:ESCH, :SAMIN] "ss_solver_parameters_algorithm must be :ESCH or :SAMIN. Got $ss_solver_parameters_algorithm. Using default :ESCH." @@ -1518,7 +1571,9 @@ macro parameters(๐“‚,ex...) mod.$๐“‚.constants.post_parameters_macro = post_parameters_macro( calib_parameters_no_var, $precompile, - $simplify, + $(QuoteNode(ss_symbolic_mode)), + $(QuoteNode(ss_solver_parameters_algorithm)), + $ss_solver_parameters_maxtime, guess_dict, ss_calib_list, par_calib_list, @@ -1545,9 +1600,6 @@ macro parameters(๐“‚,ex...) missing_parameters = missing_params, ) mod.$๐“‚.parameter_values = all_values[defined_params_idx] - # mod.$๐“‚.caches.outdated_NSSS = true - - # Store precompile and simplify flag in model container # Set custom steady state function if provided # if !isnothing($steady_state_function) @@ -1562,7 +1614,7 @@ macro parameters(๐“‚,ex...) write_ss_check_function!(mod.$๐“‚) else if !has_missing_parameters - set_up_steady_state_solver!(mod.$๐“‚, verbose = $verbose, silent = $silent, avoid_solve = !$simplify, symbolic = $symbolic) + set_up_steady_state_solver!(mod.$๐“‚, verbose = $verbose, silent = $silent, ss_symbolic_mode = $(QuoteNode(ss_symbolic_mode))) end end @@ -1577,7 +1629,7 @@ macro parameters(๐“‚,ex...) end if has_missing_parameters && $report_missing_parameters - @warn "Model has been set up with incomplete parameter definitions. Missing parameters: $(missing_params). The non-stochastic steady state and perturbation solution cannot be computed until all parameters are defined. Provide missing parameter values via the `parameters` keyword argument in functions like `get_irf`, `get_SS`, `simulate`, etc." + @warn "Model has been set up with incomplete parameter definitions. Missing parameters: $(missing_params). The non-stochastic steady state and perturbation solution cannot be computed until all parameters are defined. Provide missing parameter values via the `parameters` keyword argument in functions like `get_irf`, `get_steady_state`, `simulate`, etc." end if !$silent && $report_missing_parameters Base.show(mod.$๐“‚) end diff --git a/src/moments.jl b/src/moments.jl index 46cf36d19..5d25e43a6 100644 --- a/src/moments.jl +++ b/src/moments.jl @@ -8,28 +8,25 @@ function calculate_covariance(parameters::Vector{R}, idx_constants = constants.post_complete_parameters T = constants.post_model_macro - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts) + _nsss_result = get_NSSS_and_parameters(๐“‚, parameters, opts = opts) + SS_and_pars = _nsss_result[1]::Vector{R} + solution_error = _nsss_result[2][1] - if solution_error > opts.tol.NSSS_acceptance_tol - return zeros(0,0), zeros(0,0), zeros(0,0), SS_and_pars, solution_error < opts.tol.NSSS_acceptance_tol + if solution_error > opts.tol.nsss.acceptance_tol + return zeros(0,0), zeros(0,0), zeros(0,0), SS_and_pars, solution_error < opts.tol.nsss.acceptance_tol end - โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian) - - # Ensure QME workspace - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) + โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces) sol, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; initial_guess = ๐“‚.caches.qme_solution, - opts = opts) - - @ignore_derivatives update_perturbation_counter!(๐“‚.counters, solved, order = 1) + opts = opts, + parameter_values = parameters) - if solved ๐“‚.caches.qme_solution = qme_sol end + update_perturbation_counter!(๐“‚.counters, solved, order = 1) # Direct constants access instead of model access A = @views sol[:, 1:T.nPast_not_future_and_mixed] * idx_constants.diag_nVars[T.past_not_future_and_mixed_idx,:] @@ -43,15 +40,16 @@ function calculate_covariance(parameters::Vector{R}, end # Ensure lyapunov workspace is properly sized and get it - lyap_ws = ensure_lyapunov_workspace_1st_order!(๐“‚) + lyap_ws = ensure_lyapunov_workspace!(๐“‚.workspaces, T.nVars, :first_order) covar_raw, solved = solve_lyapunov_equation(A, CC, lyap_ws, lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, + tol = opts.tol.first_order.lyapunov, verbose = opts.verbose) - return covar_raw, sol , โˆ‡โ‚, SS_and_pars, solved + covar_stable = covar_raw + + return covar_stable, sol , โˆ‡โ‚, SS_and_pars, solved end @@ -68,52 +66,45 @@ function calculate_mean(parameters::Vector{R}, constants = initialise_constants!(๐“‚) T = constants.post_model_macro - SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(๐“‚, parameters, opts = opts) + _nsss_result = get_NSSS_and_parameters(๐“‚, parameters, opts = opts) + SS_and_pars = _nsss_result[1]::Vector{R} + solution_error = _nsss_result[2][1] if algorithm == :first_order mean_of_variables = SS_and_pars[1:T.nVars] - solved = solution_error < opts.tol.NSSS_acceptance_tol + solved = solution_error < opts.tol.nsss.acceptance_tol else ensure_moments_constants!(constants) so = constants.second_order - โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian)# |> Matrix - - # Ensure QME workspace - qme_ws = ensure_qme_workspace!(๐“‚) - sylv_ws = ensure_sylvester_1st_order_workspace!(๐“‚) + โˆ‡โ‚ = calculate_jacobian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.jacobian, ๐“‚.workspaces)# |> Matrix ๐’โ‚, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, constants, - qme_ws, - sylv_ws; + ๐“‚.workspaces, + ๐“‚.caches; initial_guess = ๐“‚.caches.qme_solution, - opts = opts) + opts = opts, + parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved, order = 1) if !solved - mean_of_variables = SS_and_pars[1:T.nVars] + mean_of_variables = fill(R(NaN), T.nVars) else - ๐“‚.caches.qme_solution = qme_sol - - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ + โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ - ๐’โ‚‚, solved = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; - opts = opts) + ๐’โ‚‚, solved = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + opts = opts, parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved, order = 2) if !solved - mean_of_variables = SS_and_pars[1:T.nVars] + mean_of_variables = fill(R(NaN), T.nVars) else - if eltype(๐’โ‚‚) == Float64 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end - ๐’โ‚‚ *= ๐“‚.constants.second_order.๐”โ‚‚ - if !(typeof(๐’โ‚‚) <: AbstractSparseMatrix) - ๐’โ‚‚ = sparse(๐’โ‚‚) # * ๐“‚.constants.second_order.๐”โ‚‚) - end + ๐’โ‚‚ = sparse(๐’โ‚‚) # ensure stable sparse type nแต‰ = T.nExo nหข = T.nPast_not_future_and_mixed @@ -196,21 +187,17 @@ function calculate_second_order_moments(parameters::Vector{R}, eโด = so.e4 # second order - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ + โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ - ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; - opts = opts) + ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + opts = opts, parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved2, order = 2) if solved2 - if eltype(๐’โ‚‚) == Float64 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end - ๐’โ‚‚ *= ๐“‚.constants.second_order.๐”โ‚‚ - if !(typeof(๐’โ‚‚) <: AbstractSparseMatrix) - ๐’โ‚‚ = sparse(๐’โ‚‚) # * ๐“‚.constants.second_order.๐”โ‚‚) - end + ๐’โ‚‚ = sparse(๐’โ‚‚) # ensure stable sparse type kron_s_s = so.kron_states kron_e_e = so.kron_e_e @@ -236,27 +223,33 @@ function calculate_second_order_moments(parameters::Vector{R}, v_v_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_v_v] |> collect s_e_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_s_e] - s_to_sโ‚_by_s_to_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + # Compression matrices + sub_idx = ensure_moments_substate_indices!(๐“‚, nหข) + Dโ‚‚หข = sub_idx.Dโ‚‚หข + Lโ‚‚หข = sub_idx.Lโ‚‚หข + nโ‚‚หข = size(Dโ‚‚หข, 2) # nหข(nหข+1)/2 + + s_to_sโ‚_by_s_to_sโ‚ = Lโ‚‚หข * โ„’.kron(s_to_sโ‚, s_to_sโ‚) * Dโ‚‚หข e_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) s_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) - # # Set up in pruned state transition matrices - sฬ‚_to_sฬ‚โ‚‚ = [ s_to_sโ‚ zeros(nหข, nหข + nหข^2) - zeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 - zeros(nหข^2, 2*nหข) s_to_sโ‚_by_s_to_sโ‚ ] + # # Set up in pruned state transition matrices (block 3 compressed: nหขยฒ โ†’ nโ‚‚หข) + ล_to_ลโ‚‚ = [ s_to_sโ‚ zeros(nหข, nหข + nโ‚‚หข) + zeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 * Dโ‚‚หข + zeros(nโ‚‚หข, 2*nหข) s_to_sโ‚_by_s_to_sโ‚ ] - eฬ‚_to_sฬ‚โ‚‚ = [ e_to_sโ‚ zeros(nหข, nแต‰^2 + nแต‰ * nหข) + รช_to_ลโ‚‚ = [ e_to_sโ‚ zeros(nหข, nแต‰^2 + nแต‰ * nหข) zeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ - zeros(nหข^2,nแต‰) e_to_sโ‚_by_e_to_sโ‚ I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚] + zeros(nโ‚‚หข,nแต‰) Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ Lโ‚‚หข * I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚] - sฬ‚_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2] + ล_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2 * Dโ‚‚หข] - eฬ‚_to_yโ‚‚ = [e_to_yโ‚ e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚] + รช_to_yโ‚‚ = [e_to_yโ‚ e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚] vec_Iโ‚‘ = so.vec_Iโ‚‘ ลvโ‚‚ = [ zeros(nหข) vec(v_v_to_sโ‚‚) / 2 + e_e_to_sโ‚‚ / 2 * vec_Iโ‚‘ - e_to_sโ‚_by_e_to_sโ‚ * vec_Iโ‚‘] + Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ * vec_Iโ‚‘] yvโ‚‚ = (vec(v_v_to_yโ‚‚) + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 @@ -297,7 +290,7 @@ end function calculate_second_order_moments_with_covariance(parameters::Vector{R}, ๐“‚::โ„ณ; - opts::CalculationOptions = merge_calculation_options())::Tuple{Matrix{R}, Matrix{R}, Vector{R}, Vector{R}, Matrix{R}, Matrix{R}, Matrix{R}, Matrix{R}, Matrix{R}, Vector{R}, Matrix{R}, Matrix{R}, AbstractSparseMatrix{R,Int}, AbstractSparseMatrix{R,Int}, Bool} where R <: Real + opts::CalculationOptions = merge_calculation_options())::Tuple{Matrix{R}, Matrix{R}, Vector{R}, Vector{R}, Matrix{R}, Matrix{R}, Matrix{R}, Matrix{R}, Matrix{R}, Vector{R}, Matrix{R}, Matrix{R}, AbstractMatrix{R}, AbstractSparseMatrix{R,Int}, Bool} where R <: Real ฮฃสธโ‚, ๐’โ‚, โˆ‡โ‚, SS_and_pars, solved = calculate_covariance(parameters, ๐“‚, opts = opts) @@ -320,27 +313,27 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, eโด = so.e4 # second order - โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ + โˆ‡โ‚‚ = calculate_hessian(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.hessian, ๐“‚.workspaces)# * ๐“‚.constants.second_order.๐”โˆ‡โ‚‚ - ๐’โ‚‚, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces; - opts = opts) + ๐’โ‚‚_raw, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, ๐’โ‚, ๐“‚.constants, ๐“‚.workspaces, ๐“‚.caches; + opts = opts, parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved2, order = 2) if solved2 - if eltype(๐’โ‚‚) == Float64 ๐“‚.caches.second_order_solution = ๐’โ‚‚ end - - ๐’โ‚‚ *= ๐“‚.constants.second_order.๐”โ‚‚ - - if !(typeof(๐’โ‚‚) <: AbstractSparseMatrix) - ๐’โ‚‚ = sparse(๐’โ‚‚) # * ๐“‚.constants.second_order.๐”โ‚‚) - end + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{R, Int} kron_s_s = so.kron_states kron_e_e = so.kron_e_e kron_v_v = so.kron_v_v kron_s_e = so.kron_s_e + # Substate duplication/elimination matrices for symmetric Kronecker compression + sub_idx = ensure_moments_substate_indices!(๐“‚, nหข) + Dโ‚‚หข = sub_idx.Dโ‚‚หข + Lโ‚‚หข = sub_idx.Lโ‚‚หข + nโ‚‚หข = size(Dโ‚‚หข, 2) # nหข(nหข+1)/2 + # first order s_to_yโ‚ = ๐’โ‚[:, 1:nหข] e_to_yโ‚ = ๐’โ‚[:, (nหข + 1):end] @@ -360,32 +353,32 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, v_v_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_v_v] |> collect s_e_to_sโ‚‚ = ๐’โ‚‚[iหข, kron_s_e] - s_to_sโ‚_by_s_to_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect + s_to_sโ‚_by_s_to_sโ‚ = Lโ‚‚หข * โ„’.kron(s_to_sโ‚, s_to_sโ‚) * Dโ‚‚หข e_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) s_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) - # # Set up in pruned state transition matrices - sฬ‚_to_sฬ‚โ‚‚ = [ s_to_sโ‚ zeros(nหข, nหข + nหข^2) - zeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 - zeros(nหข^2, 2*nหข) s_to_sโ‚_by_s_to_sโ‚ ] + # # Set up in pruned state transition matrices (block 3 compressed: nหขยฒ โ†’ nโ‚‚หข) + sฬ‚_to_sฬ‚โ‚‚ = [ s_to_sโ‚ spzeros(nหข, nหข + nโ‚‚หข) + spzeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 * Dโ‚‚หข + spzeros(nโ‚‚หข, 2*nหข) s_to_sโ‚_by_s_to_sโ‚ ] - eฬ‚_to_sฬ‚โ‚‚ = [ e_to_sโ‚ zeros(nหข, nแต‰^2 + nแต‰ * nหข) - zeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ - zeros(nหข^2,nแต‰) e_to_sโ‚_by_e_to_sโ‚ I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚] + eฬ‚_to_sฬ‚โ‚‚ = [ e_to_sโ‚ spzeros(nหข, nแต‰^2 + nแต‰ * nหข) + spzeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ + spzeros(nโ‚‚หข,nแต‰) Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ Lโ‚‚หข * I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚] - sฬ‚_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2] + sฬ‚_to_yโ‚‚ = [s_to_yโ‚ s_to_yโ‚ s_s_to_yโ‚‚ / 2 * Dโ‚‚หข] eฬ‚_to_yโ‚‚ = [e_to_yโ‚ e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚] vec_Iโ‚‘ = so.vec_Iโ‚‘ ลvโ‚‚ = [ zeros(nหข) vec(v_v_to_sโ‚‚) / 2 + e_e_to_sโ‚‚ / 2 * vec_Iโ‚‘ - e_to_sโ‚_by_e_to_sโ‚ * vec_Iโ‚‘] + Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ * vec_Iโ‚‘] yvโ‚‚ = (vec(v_v_to_yโ‚‚) + e_e_to_yโ‚‚ * vec_Iโ‚‘) / 2 ## Mean - ฮผหขโบโ‚‚ = (โ„’.I(size(ล_to_ลโ‚‚, 1)) - ล_to_ลโ‚‚) \ ลvโ‚‚ + ฮผหขโบโ‚‚ = collect(โ„’.I(size(ล_to_ลโ‚‚, 1)) - ล_to_ลโ‚‚) \ ลvโ‚‚ ฮ”ฮผหขโ‚‚ = vec((โ„’.I(size(s_to_sโ‚, 1)) - s_to_sโ‚) \ (s_s_to_sโ‚‚ * vec(ฮฃแถปโ‚) / 2 + (v_v_to_sโ‚‚ + e_e_to_sโ‚‚ * vec_Iโ‚‘) / 2)) ฮผสธโ‚‚ = SS_and_pars[1:๐“‚.constants.post_model_macro.nVars] + ล_to_yโ‚‚ * ฮผหขโบโ‚‚ + yvโ‚‚ @@ -401,8 +394,7 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, ฮฃแถปโ‚‚, info = solve_lyapunov_equation(ล_to_ลโ‚‚, C, lyap_ws_2nd, lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, + tol = opts.tol.second_order.lyapunov, verbose = opts.verbose) if info @@ -412,9 +404,10 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, slvd = solved && solved2 && info else - ฮฃสธโ‚‚ = zeros(R,0,0) + nVars = ๐“‚.constants.post_model_macro.nVars + ฮฃสธโ‚‚ = fill(R(NaN), nVars, nVars) ฮฃแถปโ‚‚ = zeros(R,0,0) - ฮผสธโ‚‚ = zeros(R,0) + ฮผสธโ‚‚ = fill(R(NaN), nVars) ฮ”ฮผหขโ‚‚ = zeros(R,0) autocorr_tmp = zeros(R,0,0) sฬ‚_to_sฬ‚โ‚‚ = zeros(R,0,0) @@ -422,9 +415,10 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, slvd = info end else - ฮฃสธโ‚‚ = zeros(R,0,0) + nVars = ๐“‚.constants.post_model_macro.nVars + ฮฃสธโ‚‚ = fill(R(NaN), nVars, nVars) ฮฃแถปโ‚‚ = zeros(R,0,0) - ฮผสธโ‚‚ = zeros(R,0) + ฮผสธโ‚‚ = fill(R(NaN), nVars) ฮ”ฮผหขโ‚‚ = zeros(R,0) autocorr_tmp = zeros(R,0,0) sฬ‚_to_sฬ‚โ‚‚ = zeros(R,0,0) @@ -432,9 +426,10 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, slvd = solved2 end else - ฮฃสธโ‚‚ = zeros(R,0,0) + nVars = ๐“‚.constants.post_model_macro.nVars + ฮฃสธโ‚‚ = fill(R(NaN), nVars, nVars) ฮฃแถปโ‚‚ = zeros(R,0,0) - ฮผสธโ‚‚ = zeros(R,0) + ฮผสธโ‚‚ = fill(R(NaN), nVars) ฮ”ฮผหขโ‚‚ = zeros(R,0) autocorr_tmp = zeros(R,0,0) sฬ‚_to_sฬ‚โ‚‚ = zeros(R,0,0) @@ -444,43 +439,145 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, # SS_and_pars = zeros(R,0) # ๐’โ‚ = zeros(R,0,0) # โˆ‡โ‚ = zeros(R,0,0) - ๐’โ‚‚ = spzeros(R,0,0) + ๐’โ‚‚_raw = zeros(R,0,0) โˆ‡โ‚‚ = spzeros(R,0,0) slvd = solved end - return ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, slvd + return ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, ล_to_ลโ‚‚, ล_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_raw, โˆ‡โ‚‚, slvd end +# Block-triangular Lyapunov solver for third-order pruned state covariance. +# Solves the block-triangular Lyapunov equation for the third-order pruned state covariance. +# Accepts pre-sliced sub-blocks of the transition matrix [A_UU 0; A_LU A_LL] +# and RHS matrix C (only C_LU and C_LL blocks needed). +# Reuses ฮฃฬ‚แถปโ‚‚ for the upper block; Sylvester for the cross-block; Lyapunov for the lower block. +function solve_block_triangular_lyapunov(A_UU::AbstractMatrix{T}, + A_LU::AbstractMatrix{T}, + A_LL::AbstractMatrix{T}, + C_LU::AbstractMatrix{T}, + C_LL::AbstractMatrix{T}, + ฮฃแถปโ‚‚_upper::AbstractMatrix{T}, + ๐“‚_workspaces::workspaces, + opts::CalculationOptions; + nโ‚ƒหข::Int = 0) where T <: Real + N_upper = size(A_UU, 1) + N_lower = size(A_LL, 1) + + # Step 1: X_UU = ฮฃฬ‚แถปโ‚‚ (already solved) + X_UU = ฮฃแถปโ‚‚_upper + + # Step 2: X_LU via discrete Sylvester (A_LL X_LU A_UU' + RHS = X_LU) + RHS_LU = A_LU * X_UU * A_UU' + C_LU + + sylv_ws = ๐“‚_workspaces.sylvester_block + X_LU, sylv_solved = solve_sylvester_equation(A_LL, A_UU', RHS_LU, sylv_ws, + tol = opts.tol.third_order.sylvester, + verbose = opts.verbose) + + # Step 3: X_LL via Lyapunov with modified RHS + C_LL_mod = C_LL + + A_LU * X_UU * A_LU' + + A_LL * X_LU * A_LU' + + A_LU * X_LU' * A_LL' + + if nโ‚ƒหข > 0 && N_lower > nโ‚ƒหข + # A_LL has sub-block structure: decompose into Aโ‚†โ‚† (lower-right nโ‚ƒหขร—nโ‚ƒหข) and upper blocks + n_upper_LL = N_lower - nโ‚ƒหข + ru_ll = 1:n_upper_LL + rl_ll = (n_upper_LL+1):N_lower + + A_LL_UU = A_LL[ru_ll, ru_ll] + A_LL_UL = A_LL[ru_ll, rl_ll] + A_LL_LL = A_LL[rl_ll, rl_ll] + + C_mod_UU = C_LL_mod[ru_ll, ru_ll] + C_mod_UL = C_LL_mod[ru_ll, rl_ll] + C_mod_LL = C_LL_mod[rl_ll, rl_ll] + + # Step 3a: Xโ‚†โ‚† via standard Lyapunov + lyap_ws_66 = ensure_lyapunov_workspace!(๐“‚_workspaces, nโ‚ƒหข, :block) + X_66, _ = solve_lyapunov_equation(A_LL_LL, C_mod_LL, lyap_ws_66, + tol = opts.tol.third_order.lyapunov, + verbose = opts.verbose) + + # Step 3b: X_{upper,6} via Sylvester + RHS_UL6 = A_LL_UL * X_66 * A_LL_LL' + C_mod_UL + X_UL6, _ = solve_sylvester_equation(A_LL_UU, A_LL_LL', RHS_UL6, sylv_ws, + tol = opts.tol.third_order.sylvester, + verbose = opts.verbose) + + # Step 3c: X_{upper,upper} via Lyapunov + C_UU_mod2 = C_mod_UU + + A_LL_UL * X_66 * A_LL_UL' + + A_LL_UU * X_UL6 * A_LL_UL' + + A_LL_UL * X_UL6' * A_LL_UU' + + lyap_ws_inner = ensure_lyapunov_workspace!(๐“‚_workspaces, n_upper_LL, :block) + X_UU_LL, _ = solve_lyapunov_equation(A_LL_UU, C_UU_mod2, lyap_ws_inner, + tol = opts.tol.third_order.lyapunov, + verbose = opts.verbose) + + X_LL = zeros(T, N_lower, N_lower) + X_LL[ru_ll, ru_ll] = X_UU_LL + X_LL[ru_ll, rl_ll] = X_UL6 + X_LL[rl_ll, ru_ll] = X_UL6' + X_LL[rl_ll, rl_ll] = X_66 + else + # Standard Lyapunov on full lower block + lyap_ws = ensure_lyapunov_workspace!(๐“‚_workspaces, N_lower, :block) + X_LL_result, _ = solve_lyapunov_equation(A_LL, C_LL_mod, lyap_ws, + tol = opts.tol.third_order.lyapunov, + verbose = opts.verbose) + X_LL = X_LL_result + end + + # Reassemble full solution + N = N_upper + N_lower + ru = 1:N_upper + rl = (N_upper+1):N + ฮฃแถปโ‚ƒ = Matrix{T}(undef, N, N) + ฮฃแถปโ‚ƒ[ru, ru] = X_UU + ฮฃแถปโ‚ƒ[ru, rl] = X_LU' + ฮฃแถปโ‚ƒ[rl, ru] = X_LU + ฮฃแถปโ‚ƒ[rl, rl] = X_LL + + return ฮฃแถปโ‚ƒ, sylv_solved +end function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T}, observables::Union{Symbol_input,String_input}, ๐“‚::โ„ณ; autocorrelation_periods::U = 1:5, + third_order_block_lyapunov_method::Bool = false, covariance::Union{Symbol_input,String_input} = Symbol[], opts::CalculationOptions = merge_calculation_options())::Tuple{Matrix{T}, Vector{T}, Matrix{T}, Vector{T}, Bool} where {U, T <: Real} second_order_moments = calculate_second_order_moments_with_covariance(parameters, ๐“‚; opts = opts) - ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, solved = second_order_moments + ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, ล_to_ลโ‚‚, ล_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_raw, โˆ‡โ‚‚, solved = second_order_moments if !solved return zeros(T,0,0), zeros(T,0), zeros(T,0,0), zeros(T,0), false end + # Expand compressed ๐’โ‚‚_raw to full for moments computation + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{T, Int} + ensure_moments_constants!(๐“‚.constants) so = ๐“‚.constants.second_order to = ๐“‚.constants.third_order - โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives)# * ๐“‚.constants.third_order.๐”โˆ‡โ‚ƒ + โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces)# * ๐“‚.constants.third_order.๐”โˆ‡โ‚ƒ - ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, - ๐“‚.constants, - ๐“‚.workspaces; - initial_guess = ๐“‚.caches.third_order_solution, - opts = opts) + ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚_raw, + ๐“‚.constants, + ๐“‚.workspaces, + ๐“‚.caches; + initial_guess = ๐“‚.caches.third_order_solution, + opts = opts, parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved3, order = 3) @@ -488,15 +585,11 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T return zeros(T,0,0), zeros(T,0), zeros(T,0,0), zeros(T,0), false end - if eltype(๐’โ‚ƒ) == Float64 && solved3 ๐“‚.caches.third_order_solution = ๐’โ‚ƒ end - ๐’โ‚ƒ *= ๐“‚.constants.third_order.๐”โ‚ƒ - if !(typeof(๐’โ‚ƒ) <: AbstractSparseMatrix) - ๐’โ‚ƒ = sparse(๐’โ‚ƒ) # * ๐“‚.constants.third_order.๐”โ‚ƒ) - end + ๐’โ‚ƒ = sparse(๐’โ‚ƒ) # ensure stable sparse type - orders = determine_efficient_order(๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ, ๐“‚.constants, observables, covariance = covariance, tol = opts.tol.dependencies_tol) + orders = determine_efficient_order(๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ, ๐“‚.constants, observables, covariance = covariance, tol = opts.tol.third_order.dependencies_tol) nแต‰ = ๐“‚.constants.post_model_macro.nExo @@ -520,6 +613,15 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T e4_minus_vecIโ‚‘_outer = so.e4_minus_vecIโ‚‘_outer e6_nแต‰ยณ_nแต‰ยณ = to.e6_nแต‰ยณ_nแต‰ยณ + # Expand compressed ฮฃแถปโ‚‚ (block 3 is vech-compressed) back to full form for third-order indexing + nหข_full = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + sub_idx_full = ensure_moments_substate_indices!(๐“‚, nหข_full) + Dโ‚‚หข_full = sub_idx_full.Dโ‚‚หข + nโ‚‚หข_full = size(Dโ‚‚หข_full, 2) + Eโ‚‚_exp = [sparse(โ„’.I, 2*nหข_full, 2*nหข_full) spzeros(2*nหข_full, nโ‚‚หข_full) + spzeros(nหข_full^2, 2*nหข_full) Dโ‚‚หข_full] + ฮฃแถปโ‚‚ = Eโ‚‚_exp * ฮฃแถปโ‚‚ * Eโ‚‚_exp' + ฮฃสธโ‚ƒ = zeros(T, size(ฮฃสธโ‚‚)) autocorr = zeros(T, size(ฮฃสธโ‚‚,1), length(autocorrelation_periods)) @@ -564,6 +666,12 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T e_ss = substate_indices.e_ss ss_s = substate_indices.ss_s s_s = substate_indices.s_s + Dโ‚‚หข = substate_indices.Dโ‚‚หข + Lโ‚‚หข = substate_indices.Lโ‚‚หข + Dโ‚ƒหข = substate_indices.Dโ‚ƒหข + Lโ‚ƒหข = substate_indices.Lโ‚ƒหข + nโ‚‚หข = size(Dโ‚‚หข, 2) + nโ‚ƒหข = size(Dโ‚ƒหข, 2) # first order s_to_yโ‚ = ๐’โ‚[obs_in_y,:][:,dependencies_in_states_idx] @@ -589,6 +697,7 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T s_to_sโ‚_by_s_to_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect e_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) s_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) + s_to_sโ‚_by_s_to_sโ‚_c = Lโ‚‚หข * s_to_sโ‚_by_s_to_sโ‚ * Dโ‚‚หข # third order kron_s_v = dep_kron.kron_s_v @@ -607,22 +716,30 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T s_v_v_to_sโ‚ƒ = ๐’โ‚ƒ[iหข, โ„’.kron(kron_s_v, v_in_sโบ)] e_v_v_to_sโ‚ƒ = ๐’โ‚ƒ[iหข, โ„’.kron(kron_e_v, v_in_sโบ)] - # Set up pruned state transition matrices - sฬ‚_to_sฬ‚โ‚ƒ = [ s_to_sโ‚ zeros(nหข, 2*nหข + 2*nหข^2 + nหข^3) - zeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 zeros(nหข, nหข + nหข^2 + nหข^3) - zeros(nหข^2, 2 * nหข) s_to_sโ‚_by_s_to_sโ‚ zeros(nหข^2, nหข + nหข^2 + nหข^3) - s_v_v_to_sโ‚ƒ / 2 zeros(nหข, nหข + nหข^2) s_to_sโ‚ s_s_to_sโ‚‚ s_s_s_to_sโ‚ƒ / 6 - โ„’.kron(s_to_sโ‚,v_v_to_sโ‚‚ / 2) zeros(nหข^2, 2*nหข + nหข^2) s_to_sโ‚_by_s_to_sโ‚ โ„’.kron(s_to_sโ‚,s_s_to_sโ‚‚ / 2) - zeros(nหข^3, 3*nหข + 2*nหข^2) โ„’.kron(s_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚)] + # Set up pruned state transition sub-blocks + N_upper = 2 * nหข + nโ‚‚หข + N_lower = nหข + nหข^2 + nโ‚ƒหข + + A_UU = [s_to_sโ‚ spzeros(nหข, nหข + nโ‚‚หข) + spzeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 * Dโ‚‚หข + spzeros(nโ‚‚หข, 2 * nหข) s_to_sโ‚_by_s_to_sโ‚_c] - eฬ‚_to_sฬ‚โ‚ƒ = [ e_to_sโ‚ zeros(nหข,nแต‰^2 + 2*nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) - zeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ zeros(nหข,nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) - zeros(nหข^2,nแต‰) e_to_sโ‚_by_e_to_sโ‚ I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚ zeros(nหข^2, nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) - e_v_v_to_sโ‚ƒ / 2 zeros(nหข,nแต‰^2 + nแต‰ * nหข) s_e_to_sโ‚‚ s_s_e_to_sโ‚ƒ / 2 s_e_e_to_sโ‚ƒ / 2 e_e_e_to_sโ‚ƒ / 6 - โ„’.kron(e_to_sโ‚, v_v_to_sโ‚‚ / 2) zeros(nหข^2, nแต‰^2 + nแต‰ * nหข) s_s * s_to_sโ‚_by_e_to_sโ‚ โ„’.kron(s_to_sโ‚, s_e_to_sโ‚‚) + s_s * โ„’.kron(s_s_to_sโ‚‚ / 2, e_to_sโ‚) โ„’.kron(s_to_sโ‚, e_e_to_sโ‚‚ / 2) + s_s * โ„’.kron(s_e_to_sโ‚‚, e_to_sโ‚) โ„’.kron(e_to_sโ‚, e_e_to_sโ‚‚ / 2) - zeros(nหข^3, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข) โ„’.kron(s_to_sโ‚_by_s_to_sโ‚,e_to_sโ‚) + โ„’.kron(s_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * e_ss โ„’.kron(s_to_sโ‚_by_e_to_sโ‚,e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_e_to_sโ‚) * e_es + โ„’.kron(e_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) * e_es โ„’.kron(e_to_sโ‚,e_to_sโ‚_by_e_to_sโ‚)] + A_LU = [s_v_v_to_sโ‚ƒ / 2 spzeros(nหข, nหข + nโ‚‚หข) + โ„’.kron(s_to_sโ‚,v_v_to_sโ‚‚ / 2) spzeros(nหข^2, nหข + nโ‚‚หข) + spzeros(nโ‚ƒหข, 2 * nหข + nโ‚‚หข)] - sฬ‚_to_yโ‚ƒ = [s_to_yโ‚ + s_v_v_to_yโ‚ƒ / 2 s_to_yโ‚ s_s_to_yโ‚‚ / 2 s_to_yโ‚ s_s_to_yโ‚‚ s_s_s_to_yโ‚ƒ / 6] + A_LL = [s_to_sโ‚ s_s_to_sโ‚‚ s_s_s_to_sโ‚ƒ / 6 * Dโ‚ƒหข + spzeros(nหข^2, nหข) s_to_sโ‚_by_s_to_sโ‚ โ„’.kron(s_to_sโ‚,s_s_to_sโ‚‚ / 2) * Dโ‚ƒหข + spzeros(nโ‚ƒหข, nหข + nหข^2) Lโ‚ƒหข * โ„’.kron(s_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * Dโ‚ƒหข] + + eฬ‚_to_sฬ‚โ‚ƒ = [ e_to_sโ‚ spzeros(nหข,nแต‰^2 + 2*nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + spzeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ spzeros(nหข,nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + spzeros(nโ‚‚หข,nแต‰) Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ Lโ‚‚หข * I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚ spzeros(nโ‚‚หข, nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + e_v_v_to_sโ‚ƒ / 2 spzeros(nหข,nแต‰^2 + nแต‰ * nหข) s_e_to_sโ‚‚ s_s_e_to_sโ‚ƒ / 2 s_e_e_to_sโ‚ƒ / 2 e_e_e_to_sโ‚ƒ / 6 + โ„’.kron(e_to_sโ‚, v_v_to_sโ‚‚ / 2) spzeros(nหข^2, nแต‰^2 + nแต‰ * nหข) s_s * s_to_sโ‚_by_e_to_sโ‚ โ„’.kron(s_to_sโ‚, s_e_to_sโ‚‚) + s_s * โ„’.kron(s_s_to_sโ‚‚ / 2, e_to_sโ‚) โ„’.kron(s_to_sโ‚, e_e_to_sโ‚‚ / 2) + s_s * โ„’.kron(s_e_to_sโ‚‚, e_to_sโ‚) โ„’.kron(e_to_sโ‚, e_e_to_sโ‚‚ / 2) + spzeros(nโ‚ƒหข, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_s_to_sโ‚,e_to_sโ‚) + โ„’.kron(s_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * e_ss) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_e_to_sโ‚,e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_e_to_sโ‚) * e_es + โ„’.kron(e_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) * e_es) Lโ‚ƒหข * โ„’.kron(e_to_sโ‚,e_to_sโ‚_by_e_to_sโ‚)] + + sฬ‚_to_yโ‚ƒ = [s_to_yโ‚ + s_v_v_to_yโ‚ƒ / 2 s_to_yโ‚ s_s_to_yโ‚‚ / 2 * Dโ‚‚หข s_to_yโ‚ s_s_to_yโ‚‚ s_s_s_to_yโ‚ƒ / 6 * Dโ‚ƒหข] eฬ‚_to_yโ‚ƒ = [e_to_yโ‚ + e_v_v_to_yโ‚ƒ / 2 e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚ s_e_to_yโ‚‚ s_s_e_to_yโ‚ƒ / 2 s_e_e_to_yโ‚ƒ / 2 e_e_e_to_yโ‚ƒ / 6] @@ -646,29 +763,59 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T e4_nแต‰_nแต‰ยณ' spzeros(nแต‰^3, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚', e4_nแต‰_nแต‰ยณ') โ„’.kron(vec(ฮฃฬ‚แถปโ‚)', e4_nแต‰_nแต‰ยณ') spzeros(nแต‰^3, nหข*nแต‰^2) e6_nแต‰ยณ_nแต‰ยณ] - Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + 2*nหข^2 +nหข^3) - โ„’.kron(ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) zeros(nหข*nแต‰^2, nหข + nหข^2) โ„’.kron(ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3), vec_Iโ‚‘) - spzeros(nแต‰^3, 3*nหข + 2*nหข^2 +nหข^3)] + Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข) + โ„’.kron(ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) spzeros(nหข*nแต‰^2, nหข + nโ‚‚หข) โ„’.kron(ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3) * Lโ‚ƒหข', vec_Iโ‚‘) + spzeros(nแต‰^3, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข)] - droptol!(sฬ‚_to_sฬ‚โ‚ƒ, eps()) - droptol!(eฬ‚_to_sฬ‚โ‚ƒ, eps()) + droptol!(A_UU, eps()) + droptol!(A_LU, eps()) + droptol!(A_LL, eps()) + droptol!(รช_to_ลโ‚ƒ, eps()) droptol!(Eแดธแถป, eps()) droptol!(ฮ“โ‚ƒ, eps()) - - A = eฬ‚_to_sฬ‚โ‚ƒ * Eแดธแถป * sฬ‚_to_sฬ‚โ‚ƒ' - droptol!(A, eps()) - C = eฬ‚_to_sฬ‚โ‚ƒ * ฮ“โ‚ƒ * eฬ‚_to_sฬ‚โ‚ƒ' + A + A' - droptol!(C, eps()) + # Third-order Lyapunov solve + if third_order_block_lyapunov_method + # Block-triangular: reuse second-order covariance + Eโ‚‚_comp = [sparse(โ„’.I, 2*nหข, 2*nหข) spzeros(2*nหข, nหข^2) + spzeros(nโ‚‚หข, 2*nหข) Lโ‚‚หข] + ฮฃฬ‚แถปโ‚‚_compressed = Eโ‚‚_comp * ฮฃฬ‚แถปโ‚‚ * Eโ‚‚_comp' + + # Compute C sub-blocks directly (avoid building full Nร—N matrix) + รช_U = รช_to_ลโ‚ƒ[1:N_upper, :] + รช_L = รช_to_ลโ‚ƒ[(N_upper+1):end, :] + E_cU = Eแดธแถป[:, 1:N_upper] + E_cL = Eแดธแถป[:, (N_upper+1):end] + + Q = E_cU * A_LU' + E_cL * A_LL' + R = E_cU * A_UU' + C_LU = รช_L * (ฮ“โ‚ƒ * รช_U' + R) + Q' * รช_U' + C_LL = รช_L * (ฮ“โ‚ƒ * รช_L' + Q) + Q' * รช_L' + droptol!(C_LU, eps()) + droptol!(C_LL, eps()) + + ฮฃแถปโ‚ƒ, info = solve_block_triangular_lyapunov(A_UU, A_LU, A_LL, C_LU, C_LL, + ฮฃฬ‚แถปโ‚‚_compressed, + ๐“‚.workspaces, opts, + nโ‚ƒหข = nโ‚ƒหข) + + # Assemble full ล_to_ลโ‚ƒ (needed for autocorrelation) + ล_to_ลโ‚ƒ = [A_UU spzeros(N_upper, N_lower); A_LU A_LL] + else + ล_to_ลโ‚ƒ = [A_UU spzeros(N_upper, N_lower); A_LU A_LL] - # Ensure third-order lyapunov workspace and solve - lyap_ws_3rd = ensure_lyapunov_workspace!(๐“‚.workspaces, size(ล_to_ลโ‚ƒ, 1), :third_order) + A = รช_to_ลโ‚ƒ * Eแดธแถป * ล_to_ลโ‚ƒ' + droptol!(A, eps()) - ฮฃแถปโ‚ƒ, info = solve_lyapunov_equation(ล_to_ลโ‚ƒ, C, lyap_ws_3rd, - lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, - verbose = opts.verbose) + C = รช_to_ลโ‚ƒ * ฮ“โ‚ƒ * รช_to_ลโ‚ƒ' + A + A' + droptol!(C, eps()) + + lyap_ws_3rd = ensure_lyapunov_workspace!(๐“‚.workspaces, size(ล_to_ลโ‚ƒ, 1), :third_order) + ฮฃแถปโ‚ƒ, info = solve_lyapunov_equation(ล_to_ลโ‚ƒ, C, lyap_ws_3rd, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.third_order.lyapunov, + verbose = opts.verbose) + end if !info return zeros(T,0,0), zeros(T,0), zeros(T,0,0), zeros(T,0), false @@ -676,7 +823,7 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T solved_lyapunov = solved_lyapunov && info - ฮฃสธโ‚ƒtmp = sฬ‚_to_yโ‚ƒ * ฮฃแถปโ‚ƒ * sฬ‚_to_yโ‚ƒ' + eฬ‚_to_yโ‚ƒ * ฮ“โ‚ƒ * eฬ‚_to_yโ‚ƒ' + eฬ‚_to_yโ‚ƒ * Eแดธแถป * sฬ‚_to_yโ‚ƒ' + sฬ‚_to_yโ‚ƒ * Eแดธแถป' * eฬ‚_to_yโ‚ƒ' + ฮฃสธโ‚ƒtmp = ล_to_yโ‚ƒ * ฮฃแถปโ‚ƒ * ล_to_yโ‚ƒ' + รช_to_yโ‚ƒ * ฮ“โ‚ƒ * รช_to_yโ‚ƒ' + รช_to_yโ‚ƒ * Eแดธแถป * ล_to_yโ‚ƒ' + ล_to_yโ‚ƒ * Eแดธแถป' * รช_to_yโ‚ƒ' for obs in variance_observable ฮฃสธโ‚ƒ[indexin([obs], ๐“‚.constants.post_model_macro.var), indexin(variance_observable, ๐“‚.constants.post_model_macro.var)] = ฮฃสธโ‚ƒtmp[indexin([obs], variance_observable), :] @@ -696,14 +843,14 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T ฮฃแถปโ‚ƒโฑ .= sฬ‚_to_sฬ‚โ‚ƒ * ฮฃแถปโ‚ƒโฑ + eฬ‚_to_sฬ‚โ‚ƒ * Eแดธแถป s_to_sโ‚โฑ *= s_to_sโ‚ - Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + 2*nหข^2 +nหข^3) - โ„’.kron(s_to_sโ‚โฑ * ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) zeros(nหข*nแต‰^2, nหข + nหข^2) โ„’.kron(s_to_sโ‚โฑ * ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(s_to_sโ‚โฑ * reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(s_to_sโ‚โฑ * reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3), vec_Iโ‚‘) - spzeros(nแต‰^3, 3*nหข + 2*nหข^2 +nหข^3)] + Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข) + โ„’.kron(s_to_sโ‚โฑ * ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) spzeros(nหข*nแต‰^2, nหข + nโ‚‚หข) โ„’.kron(s_to_sโ‚โฑ * ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(s_to_sโ‚โฑ * reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(s_to_sโ‚โฑ * reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3) * Lโ‚ƒหข', vec_Iโ‚‘) + spzeros(nแต‰^3, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข)] for obs in variance_observable autocorr[indexin([obs], ๐“‚.constants.post_model_macro.var), i] .= โ„’.diag(sฬ‚_to_yโ‚ƒ * ฮฃแถปโ‚ƒโฑ * sฬ‚_to_yโ‚ƒ' + sฬ‚_to_yโ‚ƒ * sฬ‚_to_sฬ‚โ‚ƒโฑ * autocorr_tmp + eฬ‚_to_yโ‚ƒ * Eแดธแถป * sฬ‚_to_yโ‚ƒ')[indexin([obs], variance_observable)] ./ max.(โ„’.diag(ฮฃสธโ‚ƒtmp), eps(Float64))[indexin([obs], variance_observable)] - autocorr[indexin([obs], ๐“‚.constants.post_model_macro.var), i][โ„’.diag(ฮฃสธโ‚ƒtmp)[indexin([obs], variance_observable)] .< opts.tol.lyapunov_acceptance_tol] .= 0 + autocorr[indexin([obs], ๐“‚.constants.post_model_macro.var), i][โ„’.diag(ฮฃสธโ‚ƒtmp)[indexin([obs], variance_observable)] .< opts.tol.third_order.lyapunov.acceptance_tol] .= 0 end sฬ‚_to_sฬ‚โ‚ƒโฑ *= sฬ‚_to_sฬ‚โ‚ƒ @@ -717,42 +864,45 @@ function calculate_third_order_moments(parameters::Vector{T}, observables::Union{Symbol_input,String_input}, ๐“‚::โ„ณ; covariance::Union{Symbol_input,String_input} = Symbol[], + third_order_block_lyapunov_method::Bool = false, opts::CalculationOptions = merge_calculation_options())::Tuple{Matrix{T}, Vector{T}, Vector{T}, Bool} where T <: Real second_order_moments = calculate_second_order_moments_with_covariance(parameters, ๐“‚; opts = opts) - ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, sฬ‚_to_sฬ‚โ‚‚, sฬ‚_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚, โˆ‡โ‚‚, solved = second_order_moments + ฮฃสธโ‚‚, ฮฃแถปโ‚‚, ฮผสธโ‚‚, ฮ”ฮผหขโ‚‚, autocorr_tmp, ล_to_ลโ‚‚, ล_to_yโ‚‚, ฮฃสธโ‚, ฮฃแถปโ‚, SS_and_pars, ๐’โ‚, โˆ‡โ‚, ๐’โ‚‚_raw, โˆ‡โ‚‚, solved = second_order_moments if !solved - return zeros(T,0,0), zeros(T,0), zeros(T,0), false + nVars = ๐“‚.constants.post_model_macro.nVars + return fill(T(NaN), nVars, nVars), fill(T(NaN), nVars), fill(T(NaN), nVars), false end + # Expand compressed ๐’โ‚‚_raw to full for moments computation + ๐’โ‚‚ = sparse(๐’โ‚‚_raw * ๐“‚.constants.second_order.๐”โ‚‚)::SparseMatrixCSC{T, Int} + ensure_moments_constants!(๐“‚.constants) so = ๐“‚.constants.second_order to = ๐“‚.constants.third_order - โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives)# * ๐“‚.constants.third_order.๐”โˆ‡โ‚ƒ + โˆ‡โ‚ƒ = calculate_third_order_derivatives(parameters, SS_and_pars, ๐“‚.caches, ๐“‚.functions.third_order_derivatives, ๐“‚.workspaces)# * ๐“‚.constants.third_order.๐”โˆ‡โ‚ƒ - ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚, + ๐’โ‚ƒ, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚, ๐’โ‚‚_raw, ๐“‚.constants, - ๐“‚.workspaces; + ๐“‚.workspaces, + ๐“‚.caches; initial_guess = ๐“‚.caches.third_order_solution, - opts = opts) + opts = opts, parameter_values = parameters) update_perturbation_counter!(๐“‚.counters, solved3, order = 3) if !solved3 - return zeros(T,0,0), zeros(T,0), zeros(T,0), false + nVars = ๐“‚.constants.post_model_macro.nVars + return fill(T(NaN), nVars, nVars), fill(T(NaN), nVars), fill(T(NaN), nVars), false end - if eltype(๐’โ‚ƒ) == Float64 && solved3 ๐“‚.caches.third_order_solution = ๐’โ‚ƒ end - ๐’โ‚ƒ *= ๐“‚.constants.third_order.๐”โ‚ƒ - if !(typeof(๐’โ‚ƒ) <: AbstractSparseMatrix) - ๐’โ‚ƒ = sparse(๐’โ‚ƒ) # * ๐“‚.constants.third_order.๐”โ‚ƒ) - end + ๐’โ‚ƒ = sparse(๐’โ‚ƒ) # ensure stable sparse type - orders = determine_efficient_order(๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ, ๐“‚.constants, observables, covariance = covariance, tol = opts.tol.dependencies_tol) + orders = determine_efficient_order(๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ, ๐“‚.constants, observables, covariance = covariance, tol = opts.tol.third_order.dependencies_tol) nแต‰ = ๐“‚.constants.post_model_macro.nExo @@ -776,6 +926,15 @@ function calculate_third_order_moments(parameters::Vector{T}, e4_minus_vecIโ‚‘_outer = so.e4_minus_vecIโ‚‘_outer e6_nแต‰ยณ_nแต‰ยณ = to.e6_nแต‰ยณ_nแต‰ยณ + # Expand compressed ฮฃแถปโ‚‚ (block 3 is vech-compressed) back to full form for third-order indexing + nหข_full = ๐“‚.constants.post_model_macro.nPast_not_future_and_mixed + sub_idx_full = ensure_moments_substate_indices!(๐“‚, nหข_full) + Dโ‚‚หข_full = sub_idx_full.Dโ‚‚หข + nโ‚‚หข_full = size(Dโ‚‚หข_full, 2) + Eโ‚‚_exp = [sparse(โ„’.I, 2*nหข_full, 2*nหข_full) spzeros(2*nหข_full, nโ‚‚หข_full) + spzeros(nหข_full^2, 2*nหข_full) Dโ‚‚หข_full] + ฮฃแถปโ‚‚ = Eโ‚‚_exp * ฮฃแถปโ‚‚ * Eโ‚‚_exp' + ฮฃสธโ‚ƒ = zeros(T, size(ฮฃสธโ‚‚)) solved_lyapunov = true @@ -818,6 +977,12 @@ function calculate_third_order_moments(parameters::Vector{T}, e_ss = substate_indices.e_ss ss_s = substate_indices.ss_s s_s = substate_indices.s_s + Dโ‚‚หข = substate_indices.Dโ‚‚หข + Lโ‚‚หข = substate_indices.Lโ‚‚หข + Dโ‚ƒหข = substate_indices.Dโ‚ƒหข + Lโ‚ƒหข = substate_indices.Lโ‚ƒหข + nโ‚‚หข = size(Dโ‚‚หข, 2) + nโ‚ƒหข = size(Dโ‚ƒหข, 2) # first order s_to_yโ‚ = ๐’โ‚[obs_in_y,:][:,dependencies_in_states_idx] @@ -843,6 +1008,7 @@ function calculate_third_order_moments(parameters::Vector{T}, s_to_sโ‚_by_s_to_sโ‚ = โ„’.kron(s_to_sโ‚, s_to_sโ‚) |> collect e_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(e_to_sโ‚, e_to_sโ‚) s_to_sโ‚_by_e_to_sโ‚ = โ„’.kron(s_to_sโ‚, e_to_sโ‚) + s_to_sโ‚_by_s_to_sโ‚_c = Lโ‚‚หข * s_to_sโ‚_by_s_to_sโ‚ * Dโ‚‚หข # third order kron_s_v = dep_kron.kron_s_v @@ -861,22 +1027,30 @@ function calculate_third_order_moments(parameters::Vector{T}, s_v_v_to_sโ‚ƒ = ๐’โ‚ƒ[iหข, โ„’.kron(kron_s_v, v_in_sโบ)] e_v_v_to_sโ‚ƒ = ๐’โ‚ƒ[iหข, โ„’.kron(kron_e_v, v_in_sโบ)] - # Set up pruned state transition matrices - sฬ‚_to_sฬ‚โ‚ƒ = [ s_to_sโ‚ zeros(nหข, 2*nหข + 2*nหข^2 + nหข^3) - zeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 zeros(nหข, nหข + nหข^2 + nหข^3) - zeros(nหข^2, 2 * nหข) s_to_sโ‚_by_s_to_sโ‚ zeros(nหข^2, nหข + nหข^2 + nหข^3) - s_v_v_to_sโ‚ƒ / 2 zeros(nหข, nหข + nหข^2) s_to_sโ‚ s_s_to_sโ‚‚ s_s_s_to_sโ‚ƒ / 6 - โ„’.kron(s_to_sโ‚,v_v_to_sโ‚‚ / 2) zeros(nหข^2, 2*nหข + nหข^2) s_to_sโ‚_by_s_to_sโ‚ โ„’.kron(s_to_sโ‚,s_s_to_sโ‚‚ / 2) - zeros(nหข^3, 3*nหข + 2*nหข^2) โ„’.kron(s_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚)] + # Set up pruned state transition sub-blocks + N_upper = 2 * nหข + nโ‚‚หข + N_lower = nหข + nหข^2 + nโ‚ƒหข + + A_UU = [s_to_sโ‚ spzeros(nหข, nหข + nโ‚‚หข) + spzeros(nหข, nหข) s_to_sโ‚ s_s_to_sโ‚‚ / 2 * Dโ‚‚หข + spzeros(nโ‚‚หข, 2 * nหข) s_to_sโ‚_by_s_to_sโ‚_c] - eฬ‚_to_sฬ‚โ‚ƒ = [ e_to_sโ‚ zeros(nหข,nแต‰^2 + 2*nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) - zeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ zeros(nหข,nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) - zeros(nหข^2,nแต‰) e_to_sโ‚_by_e_to_sโ‚ I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚ zeros(nหข^2, nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) - e_v_v_to_sโ‚ƒ / 2 zeros(nหข,nแต‰^2 + nแต‰ * nหข) s_e_to_sโ‚‚ s_s_e_to_sโ‚ƒ / 2 s_e_e_to_sโ‚ƒ / 2 e_e_e_to_sโ‚ƒ / 6 - โ„’.kron(e_to_sโ‚, v_v_to_sโ‚‚ / 2) zeros(nหข^2, nแต‰^2 + nแต‰ * nหข) s_s * s_to_sโ‚_by_e_to_sโ‚ โ„’.kron(s_to_sโ‚, s_e_to_sโ‚‚) + s_s * โ„’.kron(s_s_to_sโ‚‚ / 2, e_to_sโ‚) โ„’.kron(s_to_sโ‚, e_e_to_sโ‚‚ / 2) + s_s * โ„’.kron(s_e_to_sโ‚‚, e_to_sโ‚) โ„’.kron(e_to_sโ‚, e_e_to_sโ‚‚ / 2) - zeros(nหข^3, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข) โ„’.kron(s_to_sโ‚_by_s_to_sโ‚,e_to_sโ‚) + โ„’.kron(s_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * e_ss โ„’.kron(s_to_sโ‚_by_e_to_sโ‚,e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_e_to_sโ‚) * e_es + โ„’.kron(e_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) * e_es โ„’.kron(e_to_sโ‚,e_to_sโ‚_by_e_to_sโ‚)] + A_LU = [s_v_v_to_sโ‚ƒ / 2 spzeros(nหข, nหข + nโ‚‚หข) + โ„’.kron(s_to_sโ‚,v_v_to_sโ‚‚ / 2) spzeros(nหข^2, nหข + nโ‚‚หข) + spzeros(nโ‚ƒหข, 2 * nหข + nโ‚‚หข)] - sฬ‚_to_yโ‚ƒ = [s_to_yโ‚ + s_v_v_to_yโ‚ƒ / 2 s_to_yโ‚ s_s_to_yโ‚‚ / 2 s_to_yโ‚ s_s_to_yโ‚‚ s_s_s_to_yโ‚ƒ / 6] + A_LL = [s_to_sโ‚ s_s_to_sโ‚‚ s_s_s_to_sโ‚ƒ / 6 * Dโ‚ƒหข + spzeros(nหข^2, nหข) s_to_sโ‚_by_s_to_sโ‚ โ„’.kron(s_to_sโ‚,s_s_to_sโ‚‚ / 2) * Dโ‚ƒหข + spzeros(nโ‚ƒหข, nหข + nหข^2) Lโ‚ƒหข * โ„’.kron(s_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * Dโ‚ƒหข] + + eฬ‚_to_sฬ‚โ‚ƒ = [ e_to_sโ‚ spzeros(nหข,nแต‰^2 + 2*nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + spzeros(nหข,nแต‰) e_e_to_sโ‚‚ / 2 s_e_to_sโ‚‚ spzeros(nหข,nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + spzeros(nโ‚‚หข,nแต‰) Lโ‚‚หข * e_to_sโ‚_by_e_to_sโ‚ Lโ‚‚หข * I_plus_s_s * s_to_sโ‚_by_e_to_sโ‚ spzeros(nโ‚‚หข, nแต‰ * nหข + nแต‰ * nหข^2 + nแต‰^2 * nหข + nแต‰^3) + e_v_v_to_sโ‚ƒ / 2 spzeros(nหข,nแต‰^2 + nแต‰ * nหข) s_e_to_sโ‚‚ s_s_e_to_sโ‚ƒ / 2 s_e_e_to_sโ‚ƒ / 2 e_e_e_to_sโ‚ƒ / 6 + โ„’.kron(e_to_sโ‚, v_v_to_sโ‚‚ / 2) spzeros(nหข^2, nแต‰^2 + nแต‰ * nหข) s_s * s_to_sโ‚_by_e_to_sโ‚ โ„’.kron(s_to_sโ‚, s_e_to_sโ‚‚) + s_s * โ„’.kron(s_s_to_sโ‚‚ / 2, e_to_sโ‚) โ„’.kron(s_to_sโ‚, e_e_to_sโ‚‚ / 2) + s_s * โ„’.kron(s_e_to_sโ‚‚, e_to_sโ‚) โ„’.kron(e_to_sโ‚, e_e_to_sโ‚‚ / 2) + spzeros(nโ‚ƒหข, nแต‰ + nแต‰^2 + 2*nแต‰ * nหข) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_s_to_sโ‚,e_to_sโ‚) + โ„’.kron(s_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_s_to_sโ‚) * e_ss) Lโ‚ƒหข * (โ„’.kron(s_to_sโ‚_by_e_to_sโ‚,e_to_sโ‚) + โ„’.kron(e_to_sโ‚,s_to_sโ‚_by_e_to_sโ‚) * e_es + โ„’.kron(e_to_sโ‚, s_s * s_to_sโ‚_by_e_to_sโ‚) * e_es) Lโ‚ƒหข * โ„’.kron(e_to_sโ‚,e_to_sโ‚_by_e_to_sโ‚)] + + sฬ‚_to_yโ‚ƒ = [s_to_yโ‚ + s_v_v_to_yโ‚ƒ / 2 s_to_yโ‚ s_s_to_yโ‚‚ / 2 * Dโ‚‚หข s_to_yโ‚ s_s_to_yโ‚‚ s_s_s_to_yโ‚ƒ / 6 * Dโ‚ƒหข] eฬ‚_to_yโ‚ƒ = [e_to_yโ‚ + e_v_v_to_yโ‚ƒ / 2 e_e_to_yโ‚‚ / 2 s_e_to_yโ‚‚ s_e_to_yโ‚‚ s_s_e_to_yโ‚ƒ / 2 s_e_e_to_yโ‚ƒ / 2 e_e_e_to_yโ‚ƒ / 6] @@ -900,38 +1074,65 @@ function calculate_third_order_moments(parameters::Vector{T}, e4_nแต‰_nแต‰ยณ' spzeros(nแต‰^3, nแต‰^2 + nแต‰ * nหข) โ„’.kron(ฮ”ฬ‚ฮผหขโ‚‚', e4_nแต‰_nแต‰ยณ') โ„’.kron(vec(ฮฃฬ‚แถปโ‚)', e4_nแต‰_nแต‰ยณ') spzeros(nแต‰^3, nหข*nแต‰^2) e6_nแต‰ยณ_nแต‰ยณ] - Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + 2*nหข^2 +nหข^3) - โ„’.kron(ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) zeros(nหข*nแต‰^2, nหข + nหข^2) โ„’.kron(ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3), vec_Iโ‚‘) - spzeros(nแต‰^3, 3*nหข + 2*nหข^2 +nหข^3)] + Eแดธแถป = [ spzeros(nแต‰ + nแต‰^2 + 2*nแต‰*nหข + nแต‰*nหข^2, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข) + โ„’.kron(ฮฃฬ‚แถปโ‚,vec_Iโ‚‘) spzeros(nหข*nแต‰^2, nหข + nโ‚‚หข) โ„’.kron(ฮผหขโ‚ƒฮดฮผหขโ‚',vec_Iโ‚‘) โ„’.kron(reshape(ss_s * vec(ฮฃฬ‚แถปโ‚‚[nหข + 1:2*nหข,2 * nหข + 1 : end] + ฮ”ฬ‚ฮผหขโ‚‚ * vec(ฮฃฬ‚แถปโ‚)'), nหข, nหข^2), vec_Iโ‚‘) โ„’.kron(reshape(ฮฃฬ‚แถปโ‚‚[2 * nหข + 1 : end, 2 * nหข + 1 : end] + vec(ฮฃฬ‚แถปโ‚) * vec(ฮฃฬ‚แถปโ‚)', nหข, nหข^3) * Lโ‚ƒหข', vec_Iโ‚‘) + spzeros(nแต‰^3, 3*nหข + nโ‚‚หข + nหข^2 + nโ‚ƒหข)] - droptol!(sฬ‚_to_sฬ‚โ‚ƒ, eps()) - droptol!(eฬ‚_to_sฬ‚โ‚ƒ, eps()) + droptol!(A_UU, eps()) + droptol!(A_LU, eps()) + droptol!(A_LL, eps()) + droptol!(รช_to_ลโ‚ƒ, eps()) droptol!(Eแดธแถป, eps()) droptol!(ฮ“โ‚ƒ, eps()) - - A = eฬ‚_to_sฬ‚โ‚ƒ * Eแดธแถป * sฬ‚_to_sฬ‚โ‚ƒ' - droptol!(A, eps()) - C = eฬ‚_to_sฬ‚โ‚ƒ * ฮ“โ‚ƒ * eฬ‚_to_sฬ‚โ‚ƒ' + A + A' - droptol!(C, eps()) + # Third-order Lyapunov solve + if third_order_block_lyapunov_method + # Block-triangular: reuse second-order covariance + Eโ‚‚_comp = [sparse(โ„’.I, 2*nหข, 2*nหข) spzeros(2*nหข, nหข^2) + spzeros(nโ‚‚หข, 2*nหข) Lโ‚‚หข] + ฮฃฬ‚แถปโ‚‚_compressed = Eโ‚‚_comp * ฮฃฬ‚แถปโ‚‚ * Eโ‚‚_comp' + + # Compute C sub-blocks directly (avoid building full Nร—N matrix) + รช_U = รช_to_ลโ‚ƒ[1:N_upper, :] + รช_L = รช_to_ลโ‚ƒ[(N_upper+1):end, :] + E_cU = Eแดธแถป[:, 1:N_upper] + E_cL = Eแดธแถป[:, (N_upper+1):end] + + Q = E_cU * A_LU' + E_cL * A_LL' + R = E_cU * A_UU' + C_LU = รช_L * (ฮ“โ‚ƒ * รช_U' + R) + Q' * รช_U' + C_LL = รช_L * (ฮ“โ‚ƒ * รช_L' + Q) + Q' * รช_L' + droptol!(C_LU, eps()) + droptol!(C_LL, eps()) + + ฮฃแถปโ‚ƒ, info = solve_block_triangular_lyapunov(A_UU, A_LU, A_LL, C_LU, C_LL, + ฮฃฬ‚แถปโ‚‚_compressed, + ๐“‚.workspaces, opts, + nโ‚ƒหข = nโ‚ƒหข) + else + ล_to_ลโ‚ƒ = [A_UU spzeros(N_upper, N_lower); A_LU A_LL] - # Ensure third-order lyapunov workspace and solve - lyap_ws_3rd = ensure_lyapunov_workspace!(๐“‚.workspaces, size(ล_to_ลโ‚ƒ, 1), :third_order) + A = รช_to_ลโ‚ƒ * Eแดธแถป * ล_to_ลโ‚ƒ' + droptol!(A, eps()) - ฮฃแถปโ‚ƒ, info = solve_lyapunov_equation(ล_to_ลโ‚ƒ, C, lyap_ws_3rd, - lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, - verbose = opts.verbose) + C = รช_to_ลโ‚ƒ * ฮ“โ‚ƒ * รช_to_ลโ‚ƒ' + A + A' + droptol!(C, eps()) + + lyap_ws_3rd = ensure_lyapunov_workspace!(๐“‚.workspaces, size(ล_to_ลโ‚ƒ, 1), :third_order) + ฮฃแถปโ‚ƒ, info = solve_lyapunov_equation(ล_to_ลโ‚ƒ, C, lyap_ws_3rd, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.third_order.lyapunov, + verbose = opts.verbose) + end if !info - return zeros(T,0,0), zeros(T,0), zeros(T,0), false + nVars = ๐“‚.constants.post_model_macro.nVars + return fill(T(NaN), nVars, nVars), fill(T(NaN), nVars), fill(T(NaN), nVars), false end solved_lyapunov = solved_lyapunov && info ฮฃสธโ‚ƒtmp = sฬ‚_to_yโ‚ƒ * ฮฃแถปโ‚ƒ * sฬ‚_to_yโ‚ƒ' + eฬ‚_to_yโ‚ƒ * ฮ“โ‚ƒ * eฬ‚_to_yโ‚ƒ' + eฬ‚_to_yโ‚ƒ * Eแดธแถป * sฬ‚_to_yโ‚ƒ' + sฬ‚_to_yโ‚ƒ * Eแดธแถป' * eฬ‚_to_yโ‚ƒ' - for obs in variance_observable ฮฃสธโ‚ƒ[indexin([obs], ๐“‚.constants.post_model_macro.var), indexin(variance_observable, ๐“‚.constants.post_model_macro.var)] = ฮฃสธโ‚ƒtmp[indexin([obs], variance_observable), :] end diff --git a/src/nsss_solver.jl b/src/nsss_solver.jl new file mode 100644 index 000000000..be940b744 --- /dev/null +++ b/src/nsss_solver.jl @@ -0,0 +1,2090 @@ +# Non-stochastic steady state (NSSS) solver +# +# This file contains: +# 1. Builder for accumulating step data into model sub-structs +# 2. Step execution function dispatching on step type +# 3. The solve_nsss_steps orchestrator that iterates over steps +# 4. The solve_nsss_wrapper that handles cache management and continuation method + +# ============================================================================ +# Pipeline builder +# ============================================================================ + +const EMPTY_NSSS_STEP_CACHE = Vector{Vector{Float64}}() +const NOOP_NSSS_FUNC! = (_out, _sol_vec, _params_vec) -> nothing +const NOOP_NSSS_EVAL! = (_out, _sol_vec, _params_vec) -> nothing + +""" +Mutable accumulator used during `write_steady_state_solver_function!` to collect step data. +After all steps are appended, call `build_nsss_solver!(๐“‚, builder, param_prep!)` to assign +data into the model's functions, constants, and workspaces sub-structs. +""" +mutable struct NSSSSolverBuilder + # Per-step parallel vectors (functions) + aux_funcs::Vector{Function} + error_funcs::Vector{Function} + eval_funcs::Vector{Function} + solve_blocks::Vector{Union{Nothing, ss_solve_block}} + # Per-step metadata + step_types::Vector{UInt8} + descriptions::Vector{String} + block_indices::Vector{Int} + # Flat index accumulators + write_indices::Vector{Int} + write_ranges::Vector{UnitRange{Int}} + aux_write_indices::Vector{Int} + aux_write_ranges::Vector{UnitRange{Int}} + param_gather_indices::Vector{Int} + param_gather_ranges::Vector{UnitRange{Int}} + var_gather_indices::Vector{Int} + var_gather_ranges::Vector{UnitRange{Int}} + # Flat bounds accumulators (analytical) + lower_bounds::Vector{Float64} + upper_bounds::Vector{Float64} + has_bounds::BitVector + bounds_ranges::Vector{UnitRange{Int}} + # Flat bounds accumulators (numerical) + numerical_lbs::Vector{Float64} + numerical_ubs::Vector{Float64} + numerical_bounds_ranges::Vector{UnitRange{Int}} + # Error sizes + error_sizes::Vector{Int} + aux_error_sizes::Vector{Int} + # Workspace size tracking + max_main_buffer::Int + max_aux_buffer::Int + max_error_buffer::Int + max_guess_buffer::Int +end + +function NSSSSolverBuilder() + NSSSSolverBuilder( + Function[], Function[], + Function[], Union{Nothing,ss_solve_block}[], + UInt8[], String[], Int[], + Int[], UnitRange{Int}[], + Int[], UnitRange{Int}[], + Int[], UnitRange{Int}[], + Int[], UnitRange{Int}[], + Float64[], Float64[], BitVector(), UnitRange{Int}[], + Float64[], Float64[], UnitRange{Int}[], + Int[], Int[], + 0, 0, 0, 0, + ) +end + +"""Append an analytical step to the builder.""" +function push_analytical_step!(b::NSSSSolverBuilder; + aux_func!::Function = NOOP_NSSS_FUNC!, + aux_write_indices::Vector{Int} = Int[], + error_func!::Function = NOOP_NSSS_FUNC!, + error_size::Int = 0, + eval_func!::Function, + write_indices::Vector{Int}, + lower_bounds::Vector{Float64} = Float64[], + upper_bounds::Vector{Float64} = Float64[], + has_bounds::BitVector = falses(length(write_indices)), + description::String = "") + push!(b.step_types, ANALYTICAL_STEP) + push!(b.descriptions, description) + push!(b.block_indices, 0) + + # Functions + push!(b.aux_funcs, aux_func!) + push!(b.error_funcs, error_func!) + push!(b.eval_funcs, eval_func!) + push!(b.solve_blocks, nothing) + + # Write indices + off = length(b.write_indices) + append!(b.write_indices, write_indices) + push!(b.write_ranges, (off+1):(off+length(write_indices))) + + # Aux write indices + off = length(b.aux_write_indices) + append!(b.aux_write_indices, aux_write_indices) + push!(b.aux_write_ranges, (off+1):(off+length(aux_write_indices))) + + # No param/var gather for analytical + push!(b.param_gather_ranges, 1:0) + push!(b.var_gather_ranges, 1:0) + + # Bounds (analytical) + off = length(b.lower_bounds) + append!(b.lower_bounds, lower_bounds) + append!(b.upper_bounds, upper_bounds) + append!(b.has_bounds, has_bounds) + push!(b.bounds_ranges, (off+1):(off+length(lower_bounds))) + + # No numerical bounds + push!(b.numerical_bounds_ranges, 1:0) + + # Error sizes + push!(b.error_sizes, error_size) + push!(b.aux_error_sizes, 0) + + # Update workspace max sizes + b.max_main_buffer = max(b.max_main_buffer, length(write_indices)) + b.max_aux_buffer = max(b.max_aux_buffer, length(aux_write_indices)) + b.max_error_buffer = max(b.max_error_buffer, error_size) +end + +"""Append a numerical step to the builder.""" +function push_numerical_step!(b::NSSSSolverBuilder; + solve_block::ss_solve_block, + block_index::Int, + write_indices::Vector{Int}, + param_gather_indices::Vector{Int}, + var_gather_indices::Vector{Int}, + lbs::Vector{Float64}, + ubs::Vector{Float64}, + aux_func!::Function = NOOP_NSSS_FUNC!, + aux_write_indices::Vector{Int} = Int[], + aux_error_func!::Function = NOOP_NSSS_FUNC!, + aux_error_size::Int = 0, + description::String = "") + push!(b.step_types, NUMERICAL_STEP) + push!(b.descriptions, description) + push!(b.block_indices, block_index) + + # Functions + push!(b.aux_funcs, aux_func!) + push!(b.error_funcs, aux_error_func!) # numerical steps use error_funcs slot for aux_error + push!(b.eval_funcs, NOOP_NSSS_EVAL!) + push!(b.solve_blocks, solve_block) + + # Write indices + off = length(b.write_indices) + append!(b.write_indices, write_indices) + push!(b.write_ranges, (off+1):(off+length(write_indices))) + + # Aux write indices + off = length(b.aux_write_indices) + append!(b.aux_write_indices, aux_write_indices) + push!(b.aux_write_ranges, (off+1):(off+length(aux_write_indices))) + + # Param/var gather indices + off = length(b.param_gather_indices) + append!(b.param_gather_indices, param_gather_indices) + push!(b.param_gather_ranges, (off+1):(off+length(param_gather_indices))) + + off = length(b.var_gather_indices) + append!(b.var_gather_indices, var_gather_indices) + push!(b.var_gather_ranges, (off+1):(off+length(var_gather_indices))) + + # No analytical bounds + push!(b.bounds_ranges, 1:0) + + # Numerical bounds + off = length(b.numerical_lbs) + append!(b.numerical_lbs, lbs) + append!(b.numerical_ubs, ubs) + push!(b.numerical_bounds_ranges, (off+1):(off+length(lbs))) + + # Error sizes + push!(b.error_sizes, 0) + push!(b.aux_error_sizes, aux_error_size) + + # Update workspace max sizes + gather_size = length(param_gather_indices) + length(var_gather_indices) + b.max_main_buffer = max(b.max_main_buffer, gather_size) + b.max_aux_buffer = max(b.max_aux_buffer, length(aux_write_indices)) + b.max_error_buffer = max(b.max_error_buffer, aux_error_size) + b.max_guess_buffer = max(b.max_guess_buffer, length(write_indices)) +end + +"""Assign the solver functions, constants, and workspace from builder data into `๐“‚`.""" +function build_nsss_solver!(๐“‚::โ„ณ, b::NSSSSolverBuilder, param_prep!::Union{Nothing,Function}) + n = length(b.step_types) + n_ext_params = length(๐“‚.constants.post_complete_parameters.parameters) + length(๐“‚.equations.calibration_no_var) + ๐“‚.functions.nsss_solver = NSSSSolverFunctions( + b.aux_funcs, b.error_funcs, b.eval_funcs, b.solve_blocks, + ) + ๐“‚.functions.nsss_param_prep! = param_prep! + ๐“‚.constants.nsss_solver = NSSSSolverConstants( + n, + n_ext_params, + b.step_types, b.descriptions, b.block_indices, + b.write_indices, b.write_ranges, + b.aux_write_indices, b.aux_write_ranges, + b.param_gather_indices, b.param_gather_ranges, + b.var_gather_indices, b.var_gather_ranges, + b.lower_bounds, b.upper_bounds, b.has_bounds, b.bounds_ranges, + b.numerical_lbs, b.numerical_ubs, b.numerical_bounds_ranges, + b.error_sizes, b.aux_error_sizes, + ) + ๐“‚.workspaces.nsss_solver = NSSSSolverWorkspace( + zeros(Float64, max(b.max_main_buffer, 1)), + zeros(Float64, max(b.max_aux_buffer, 1)), + zeros(Float64, max(b.max_error_buffer, 1)), + zeros(Float64, max(๐“‚.constants.nsss_solver.n_ext_params, 1)), + Float64[], + Float64[], + zeros(Float64, max(b.max_guess_buffer, 1)), + [zeros(Float64, max(b.max_guess_buffer, 1)), Float64[Inf]], + zeros(Float64, max(b.max_main_buffer, 1)), + zeros(Float64, max(b.max_guess_buffer, 1)), + zeros(Float64, max(b.max_guess_buffer, 1)), + Float64[], + CircularBuffer{Vector{Vector{Float64}}}(1), + 1, + zeros(Float64, length(๐“‚.equations.steady_state) + length(๐“‚.equations.calibration)), + ) + return nothing +end + +@unstable begin + function replace_symbols(exprs, remap::AbstractDict{Symbol, <:Any}) + postwalk(node -> + (node isa Symbol && haskey(remap, node)) ? remap[node] : node, + exprs, + ) + end +end + +function write_block_solution!(๐“‚, + vars_to_solve, + eqs_to_solve, + relevant_pars_across, + nsss_solver_cache_init_tmp, + eq_idx_in_block_to_solve, + atoms_in_equations_list, + solved_vars, + solved_vals; + block_index::Int, + cse = true, + skipzeros = true, + density_threshold::Float64 = .1, + nnz_parallel_threshold::Int = 1000000, + min_length::Int = 10000) + + unique_โž•_eqs = Dict{Union{Expr,Symbol},Symbol}() + + vars_to_exclude = [vcat(Symbol.(vars_to_solve), ๐“‚.constants.post_model_macro.โž•_vars),Symbol[]] + + rewritten_eqs, ss_and_aux_equations, ss_and_aux_equations_dep, ss_and_aux_equations_error, ss_and_aux_equations_error_dep = make_equation_robust_to_domain_errors(Meta.parse.(string.(eqs_to_solve)), vars_to_exclude, ๐“‚.constants.post_parameters_macro.bounds, ๐“‚.constants.post_model_macro.โž•_vars, unique_โž•_eqs) + + push!(solved_vars, Symbol.(vars_to_solve)) + push!(solved_vals, rewritten_eqs) + + syms_in_eqs = Set{Symbol}() + for i in vcat(ss_and_aux_equations_dep, ss_and_aux_equations, rewritten_eqs) + push!(syms_in_eqs, get_symbols(i)...) + end + + setdiff!(syms_in_eqs,๐“‚.constants.post_model_macro.โž•_vars) + + syms_in_eqs2 = Set{Symbol}() + for i in ss_and_aux_equations + push!(syms_in_eqs2, get_symbols(i)...) + end + + โž•_vars_alread_in_eqs = intersect(๐“‚.constants.post_model_macro.โž•_vars,reduce(union,get_symbols.(Meta.parse.(string.(eqs_to_solve))))) + + union!(syms_in_eqs, intersect(union(โž•_vars_alread_in_eqs, syms_in_eqs2), ๐“‚.constants.post_model_macro.โž•_vars)) + + push!(atoms_in_equations_list,setdiff(syms_in_eqs, solved_vars[end])) + + calib_pars_input = Symbol[] + + relevant_pars = union(intersect(reduce(union, vcat(๐“‚.constants.post_model_macro.par_list_aux_SS, ๐“‚.constants.post_parameters_macro.par_calib_list)[eq_idx_in_block_to_solve]), syms_in_eqs),intersect(syms_in_eqs, ๐“‚.constants.post_model_macro.โž•_vars)) + union!(relevant_pars_across, relevant_pars) + + sorted_vars = sort(Symbol.(vars_to_solve)) + + iii = 1 + for parss in union(๐“‚.constants.post_complete_parameters.parameters, ๐“‚.constants.post_parameters_macro.parameters_as_function_of_parameters) + if :($parss) โˆˆ relevant_pars + push!(calib_pars_input, :($parss)) + iii += 1 + end + end + + other_vrs_eliminated_by_sympy = Set{Symbol}() + for (i,val) in enumerate(solved_vals[end]) + if eq_idx_in_block_to_solve[i] โˆˆ ๐“‚.constants.post_model_macro.ss_equations_with_aux_variables + val = vcat(๐“‚.equations.steady_state_aux, ๐“‚.equations.calibration)[eq_idx_in_block_to_solve[i]] + push!(other_vrs_eliminated_by_sympy, val.args[2]) + end + end + + solved_vals_local = Union{Expr, Symbol}[] + for (i,val) in enumerate(rewritten_eqs) + push!(solved_vals_local, postwalk(x -> x isa Expr ? x.args[1] == :conjugate ? x.args[2] : x : x, val)) + end + + other_vars_input = Symbol[] + other_vrs = intersect( setdiff( union(๐“‚.constants.post_model_macro.var, ๐“‚.equations.calibration_parameters, ๐“‚.constants.post_model_macro.โž•_vars), + sort(solved_vars[end]) ), + union(syms_in_eqs, other_vrs_eliminated_by_sympy ) ) + + for var in other_vrs + push!(other_vars_input,:($(var))) + iii += 1 + end + + parameters_and_solved_vars = vcat(calib_pars_input, other_vrs) + + ng = length(sorted_vars) + np = length(parameters_and_solved_vars) + nd = length(ss_and_aux_equations_dep) + nx = iii - 1 + + Symbolics.@variables ๐”Š[1:ng] ๐”“[1:np] + + parameter_dict = Dict{Symbol, Symbol}() + back_to_array_dict = Dict{Symbolics.Num, Symbolics.Num}() + aux_vars = Symbol[] + aux_expr = [] + + for (i,v) in enumerate(sorted_vars) + push!(parameter_dict, v => :($(Symbol("๐”Š_$i")))) + push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”Š_$i"))), @__MODULE__) => ๐”Š[i]) + end + + for (i,v) in enumerate(parameters_and_solved_vars) + push!(parameter_dict, v => :($(Symbol("๐”“_$i")))) + push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("๐”“_$i"))), @__MODULE__) => ๐”“[i]) + end + + for (i,v) in enumerate(ss_and_aux_equations_dep) + push!(aux_vars, v.args[1]) + push!(aux_expr, v.args[2]) + end + + aux_replacements = Dict{Symbol, Union{Expr, Symbol, Number}}() + for (i,x) in enumerate(aux_vars) + replacement = Dict{Symbol, Union{Expr, Symbol, Number}}(x => aux_expr[i]) + for ii in i+1:length(aux_vars) + aux_expr[ii] = replace_symbols(aux_expr[ii], replacement) + end + push!(aux_replacements, x => aux_expr[i]) + end + + replaced_solved_vals = solved_vals_local |> + x -> replace_symbols.(x, Ref(aux_replacements)) |> + x -> replace_symbols.(x, Ref(parameter_dict)) |> + x -> Symbolics.parse_expr_to_symbolic.(x, Ref(@__MODULE__)) |> + x -> Symbolics.substitute.(x, Ref(back_to_array_dict)) + + lennz = length(replaced_solved_vals) + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end + + _, calc_block! = Symbolics.build_function(replaced_solved_vals, ๐”Š, ๐”“, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + ฯตหข = zeros(Symbolics.Num, ng) + ฯต = zeros(ng) + + โˆ‚block_โˆ‚parameters_and_solved_vars = Symbolics.sparsejacobian(replaced_solved_vals, ๐”Š) + + lennz = nnz(โˆ‚block_โˆ‚parameters_and_solved_vars) + if (lennz / length(โˆ‚block_โˆ‚parameters_and_solved_vars) > density_threshold) || (length(โˆ‚block_โˆ‚parameters_and_solved_vars) < min_length) + derivatives_mat = convert(Matrix, โˆ‚block_โˆ‚parameters_and_solved_vars) + buffer = zeros(Float64, size(โˆ‚block_โˆ‚parameters_and_solved_vars)) + else + derivatives_mat = โˆ‚block_โˆ‚parameters_and_solved_vars + buffer = similar(โˆ‚block_โˆ‚parameters_and_solved_vars, Float64) + buffer.nzval .= 1 + end + + chol_buff = buffer * buffer' + chol_buff += โ„’.I + + prob = ๐’ฎ.LinearProblem(chol_buff, ฯต) + chol_buffer = ๐’ฎ.init(prob, ๐’ฎ.CholeskyFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + + lu_factorization = issparse(buffer) ? ๐’ฎ.LUFactorization() : ๐’ฎ.FastLUFactorization() + prob = ๐’ฎ.LinearProblem(buffer, ฯต) + lu_buffer = ๐’ฎ.init(prob, lu_factorization, verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end + + _, func_exprs = Symbolics.build_function(derivatives_mat, ๐”Š, ๐”“, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + Symbolics.@variables ๐”Š[1:ng+nx] + + ext_diff = Symbolics.Num[] + for i in 1:nx + push!(ext_diff, ๐”“[i] - ๐”Š[ng + i]) + end + replaced_solved_vals_ext = vcat(replaced_solved_vals, ext_diff) + + _, calc_ext_block! = Symbolics.build_function(replaced_solved_vals_ext, ๐”Š, ๐”“, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + ฯตแต‰ = zeros(ng + nx) + โˆ‚ext_block_โˆ‚parameters_and_solved_vars = Symbolics.sparsejacobian(replaced_solved_vals_ext, ๐”Š) + + lennz = nnz(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) + if (lennz / length(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) > density_threshold) || (length(โˆ‚ext_block_โˆ‚parameters_and_solved_vars) < min_length) + derivatives_mat_ext = convert(Matrix, โˆ‚ext_block_โˆ‚parameters_and_solved_vars) + ext_buffer = zeros(Float64, size(โˆ‚ext_block_โˆ‚parameters_and_solved_vars)) + else + derivatives_mat_ext = โˆ‚ext_block_โˆ‚parameters_and_solved_vars + ext_buffer = similar(โˆ‚ext_block_โˆ‚parameters_and_solved_vars, Float64) + ext_buffer.nzval .= 1 + end + + ext_chol_buff = ext_buffer * ext_buffer' + ext_chol_buff += โ„’.I + + prob = ๐’ฎ.LinearProblem(ext_chol_buff, ฯตแต‰) + ext_chol_buffer = ๐’ฎ.init(prob, ๐’ฎ.CholeskyFactorization(), verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + + ext_lu_factorization = issparse(ext_buffer) ? ๐’ฎ.LUFactorization() : ๐’ฎ.FastLUFactorization() + prob = ๐’ฎ.LinearProblem(ext_buffer, ฯตแต‰) + ext_lu_buffer = ๐’ฎ.init(prob, ext_lu_factorization, verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end + + _, ext_func_exprs = Symbolics.build_function(derivatives_mat_ext, ๐”Š, ๐”“, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + push!(nsss_solver_cache_init_tmp, [haskey(๐“‚.constants.post_parameters_macro.guess, v) ? ๐“‚.constants.post_parameters_macro.guess[v] : Inf for v in sorted_vars]) + push!(nsss_solver_cache_init_tmp, [Inf]) + + lbs = Float64[] + ubs = Float64[] + limit_boundaries = 1e12 + + for i in vcat(sorted_vars, calib_pars_input, other_vars_input) + if haskey(๐“‚.constants.post_parameters_macro.bounds,i) + push!(lbs,๐“‚.constants.post_parameters_macro.bounds[i][1]) + push!(ubs,๐“‚.constants.post_parameters_macro.bounds[i][2]) + else + push!(lbs,-limit_boundaries) + push!(ubs, limit_boundaries) + end + end + + n_block = block_index + + workspace = Nonlinear_solver_workspace(ฯต, buffer, chol_buffer, lu_buffer) + ext_workspace = Nonlinear_solver_workspace(ฯตแต‰, ext_buffer, ext_chol_buffer, ext_lu_buffer) + + solve_block = ss_solve_block( + function_and_jacobian(calc_block!::Function, func_exprs::Function, workspace), + function_and_jacobian(calc_ext_block!::Function, ext_func_exprs::Function, ext_workspace) + ) + + return (sorted_vars = sorted_vars, + calib_pars_input = Symbol.(calib_pars_input), + other_vars_input = Symbol.(other_vars_input), + lbs = lbs, + ubs = ubs, + n_block = n_block, + solve_block = solve_block, + ss_and_aux_equations = ss_and_aux_equations, + ss_and_aux_equations_error = ss_and_aux_equations_error) +end + +struct PartialSolveResult{T,E} + remaining_vars::Vector{T} + solved_vars::Vector{T} + remaining_eqs::Vector{E} + solved_exprs::Vector{E} + remaining_var_indices::Vector{Int} + solved_var_indices::Vector{Int} + remaining_eq_indices::Vector{Int} + solved_eq_indices::Vector{Int} +end + +function partial_solve(eqs_to_solve::Vector{E}, vars_to_solve::Vector{T}, incidence_matrix_subset; avoid_solve::Bool = false)::PartialSolveResult{T,E} where {E, T} + for n in length(eqs_to_solve)-1:-1:2 + for eq_combo in combinations(1:length(eqs_to_solve), n) + var_indices_to_select_from = findall([sum(incidence_matrix_subset[:,eq_combo],dims = 2)...] .> 0) + var_indices_in_remaining_eqs = findall([sum(incidence_matrix_subset[:,setdiff(1:length(eqs_to_solve),eq_combo)],dims = 2)...] .> 0) + + for var_combo in combinations(var_indices_to_select_from, n) + remaining_vars_in_remaining_eqs = setdiff(var_indices_in_remaining_eqs, var_combo) + if length(remaining_vars_in_remaining_eqs) == length(eqs_to_solve) - n + if avoid_solve || count_ops(Meta.parse(string(eqs_to_solve[eq_combo]))) > 15 + soll = nothing + else + soll = solve_symbolically(eqs_to_solve[eq_combo], vars_to_solve[var_combo]) + end + + if !(isnothing(soll) || isempty(soll)) + soll_collected = E.(collect(values(soll))) + solved_var_indices = Int[var_combo...] + remaining_var_indices = [i for i in 1:length(eqs_to_solve) if i โˆ‰ solved_var_indices] + solved_eq_indices = Int[eq_combo...] + remaining_eq_indices = [i for i in 1:length(eqs_to_solve) if i โˆ‰ solved_eq_indices] + + return PartialSolveResult( + vars_to_solve[remaining_var_indices], + vars_to_solve[solved_var_indices], + eqs_to_solve[remaining_eq_indices], + soll_collected, + remaining_var_indices, + solved_var_indices, + remaining_eq_indices, + solved_eq_indices, + ) + end + end + end + end + end + + return PartialSolveResult(T[], T[], E[], E[], Int[], Int[], Int[], Int[]) +end + +function make_equation_robust_to_domain_errors(eqs, + vars_to_exclude::Vector{Vector{Symbol}}, + bounds::Dict{Symbol,Tuple{Float64,Float64}}, + โž•_vars::Vector{Symbol}, + unique_โž•_eqs; + precompile::Bool = false) + ss_and_aux_equations = Expr[] + ss_and_aux_equations_dep = Expr[] + ss_and_aux_equations_error = Expr[] + ss_and_aux_equations_error_dep = Expr[] + rewritten_eqs = Union{Expr,Symbol}[] + for eq in eqs + if eq isa Symbol + push!(rewritten_eqs, eq) + elseif eq isa Expr + rewritten_eq = postwalk(x -> + x isa Expr ? + x.head == :call ? + x.args[1] == :* ? + x.args[2] isa Int ? + x.args[3] isa Int ? + x : + Expr(:call, :*, x.args[3:end]..., x.args[2]) : + x : + x.args[1] โˆˆ [:^] ? + !(x.args[3] isa Int) ? + x.args[2] isa Symbol ? + x.args[2] โˆˆ vars_to_exclude[1] ? + begin + bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 1e12)) : (eps(), 1e12) + x + end : + begin + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if x.args[2] in vars_to_exclude[1] + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + + :($(replacement) ^ $(x.args[3])) + end : + x.args[2] isa Number ? + x : + x.args[2].head == :call ? + begin + if precompile + replacement = trivial_simplify(x.args[2]) + else + replacement = simplify(x.args[2]) + end + + if !(replacement isa Int) + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + end + + :($(replacement) ^ $(x.args[3])) + end : + x : + x : + x.args[2] isa Float64 ? + x : + x.args[1] โˆˆ [:log] ? + x.args[2] isa Symbol ? + x.args[2] โˆˆ vars_to_exclude[1] ? + begin + bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 1e12)) : (eps(), 1e12) + x + end : + begin + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if x.args[2] in vars_to_exclude[1] + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x.args[2].head == :call ? + begin + if precompile + replacement = trivial_simplify(x.args[2]) + else + replacement = simplify(x.args[2]) + end + + if !(replacement isa Int) + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1e12,max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1e12)) : (eps(), 1e12) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x : + x.args[1] โˆˆ [:norminvcdf, :norminv, :qnorm] ? + x.args[2] isa Symbol ? + x.args[2] โˆˆ vars_to_exclude[1] ? + begin + bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 1-eps())) : (eps(), 1 - eps()) + x + end : + begin + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if x.args[2] in vars_to_exclude[1] + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1 - eps())) : (eps(), 1 - eps()) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x.args[2].head == :call ? + begin + if precompile + replacement = trivial_simplify(x.args[2]) + else + replacement = simplify(x.args[2]) + end + + if !(replacement isa Int) + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(1-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 1 - eps())) : (eps(), 1 - eps()) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x : + x.args[1] โˆˆ [:exp] ? + x.args[2] isa Symbol ? + x.args[2] โˆˆ vars_to_exclude[1] ? + begin + bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], -1e12), min(bounds[x.args[2]][2], 600)) : (-1e12, 600) + x + end : + begin + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if x.args[2] in vars_to_exclude[1] + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], -1e12), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 600)) : (-1e12, 600) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x.args[2].head == :call ? + begin + if precompile + replacement = trivial_simplify(x.args[2]) + else + replacement = simplify(x.args[2]) + end + + if !(replacement isa Int) + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(600,max(-1e12,$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], -1e12), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 600)) : (-1e12, 600) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x : + x.args[1] โˆˆ [:erfcinv] ? + x.args[2] isa Symbol ? + x.args[2] โˆˆ vars_to_exclude[1] ? + begin + bounds[x.args[2]] = haskey(bounds, x.args[2]) ? (max(bounds[x.args[2]][1], eps()), min(bounds[x.args[2]][2], 2 - eps())) : (eps(), 2 - eps()) + x + end : + begin + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if x.args[2] in vars_to_exclude[1] + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 2 - eps())) : (eps(), 2 - eps()) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x.args[2].head == :call ? + begin + if precompile + replacement = trivial_simplify(x.args[2]) + else + replacement = simplify(x.args[2]) + end + + if !(replacement isa Int) + if haskey(unique_โž•_eqs, x.args[2]) + replacement = unique_โž•_eqs[x.args[2]] + else + if isempty(intersect(get_symbols(x.args[2]), vars_to_exclude[1])) + push!(ss_and_aux_equations, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + else + push!(ss_and_aux_equations_dep, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1)))) = min(2-eps(),max(eps(),$(x.args[2]))))) + push!(ss_and_aux_equations_error_dep, Expr(:call,:abs, Expr(:call,:-, :($(Symbol("โž•" * sub(string(length(โž•_vars)+1))))), x.args[2]))) + end + + bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))] = haskey(bounds, Symbol("โž•" * sub(string(length(โž•_vars)+1)))) ? (max(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][1], eps()), min(bounds[Symbol("โž•" * sub(string(length(โž•_vars)+1)))][2], 2 - eps())) : (eps(), 2 - eps()) + push!(โž•_vars,Symbol("โž•" * sub(string(length(โž•_vars)+1)))) + replacement = Symbol("โž•" * sub(string(length(โž•_vars)))) + + unique_โž•_eqs[x.args[2]] = replacement + end + end + + :($(Expr(:call, x.args[1], replacement))) + end : + x : + x : + x : + x, + eq) + push!(rewritten_eqs,rewritten_eq) + else + @assert typeof(eq) in [Symbol, Expr] + end + end + + vars_to_exclude_from_block = vcat(vars_to_exclude...) + found_new_dependecy = true + + while found_new_dependecy + found_new_dependecy = false + for ssauxdep in ss_and_aux_equations_dep + push!(vars_to_exclude_from_block, ssauxdep.args[1]) + end + + for (iii, ssaux) in enumerate(ss_and_aux_equations) + if !isempty(intersect(get_symbols(ssaux), vars_to_exclude_from_block)) + found_new_dependecy = true + push!(vars_to_exclude_from_block, ssaux.args[1]) + push!(ss_and_aux_equations_dep, ssaux) + push!(ss_and_aux_equations_error_dep, ss_and_aux_equations_error[iii]) + deleteat!(ss_and_aux_equations, iii) + deleteat!(ss_and_aux_equations_error, iii) + end + end + end + + return rewritten_eqs, ss_and_aux_equations, ss_and_aux_equations_dep, ss_and_aux_equations_error, ss_and_aux_equations_error_dep +end + +function compile_exprs_to_func(exprs::Vector, ๐”–, ๐”“_ext, placeholder_dict, back_to_array_dict; + cse = true, skipzeros = true, nnz_parallel_threshold::Int = 1000000) + sym_exprs = Symbolics.Num[] + for expr in exprs + if expr isa Number + push!(sym_exprs, Symbolics.Num(expr)) + else + clean_expr = postwalk(x -> x isa Expr && length(x.args) >= 2 && x.args[1] == :conjugate ? x.args[2] : x, expr) + replaced = replace_symbols(clean_expr, placeholder_dict) + sym = Symbolics.parse_expr_to_symbolic(replaced, @__MODULE__) + sym = Symbolics.substitute(sym, back_to_array_dict) + push!(sym_exprs, sym) + end + end + + lennz = length(sym_exprs) + parallel = lennz > nnz_parallel_threshold ? + Symbolics.ShardedForm(1500, 4) : Symbolics.SerialForm() + + _, func! = Symbolics.build_function(sym_exprs, ๐”–, ๐”“_ext, + cse = cse, skipzeros = skipzeros, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + return func! +end + +function append_numerical_step!(builder::NSSSSolverBuilder, block_meta, sol_name_to_index, ext_param_to_index, + ๐”–, ๐”“_ext, placeholder_dict, back_to_array_dict, + global_solvetime_aux_sub::Dict{Symbol, Union{Symbol, Expr}} = Dict{Symbol, Union{Symbol, Expr}}()) + write_indices = [sol_name_to_index[v] for v in block_meta.sorted_vars] + param_gather_indices = [ext_param_to_index[p] for p in block_meta.calib_pars_input] + var_gather_indices = [sol_name_to_index[v] for v in block_meta.other_vars_input] + + aux_func! = NOOP_NSSS_FUNC! + aux_write_indices = Int[] + aux_error_func! = NOOP_NSSS_FUNC! + aux_error_size = 0 + + if !isempty(block_meta.ss_and_aux_equations) + model_aux_names = Symbol[] + model_aux_rhs = Any[] + model_aux_sub = Dict{Symbol, Any}() + for eq in block_meta.ss_and_aux_equations + if eq isa Expr && eq.head == :(=) + lhs = eq.args[1] + rhs = eq.args[2] + expanded_rhs = isempty(global_solvetime_aux_sub) ? rhs : replace_symbols(rhs, global_solvetime_aux_sub) + expanded_rhs = isempty(model_aux_sub) ? expanded_rhs : replace_symbols(expanded_rhs, model_aux_sub) + if haskey(sol_name_to_index, lhs) + push!(model_aux_names, lhs) + push!(model_aux_rhs, expanded_rhs) + model_aux_sub[lhs] = expanded_rhs + else + global_solvetime_aux_sub[lhs] = expanded_rhs + end + end + end + if !isempty(model_aux_rhs) + aux_write_indices = [sol_name_to_index[v] for v in model_aux_names] + aux_func! = compile_exprs_to_func(model_aux_rhs, ๐”–, ๐”“_ext, placeholder_dict, back_to_array_dict) + end + end + + if !isempty(block_meta.ss_and_aux_equations_error) + inlined_errors = isempty(global_solvetime_aux_sub) ? block_meta.ss_and_aux_equations_error : [replace_symbols(e, global_solvetime_aux_sub) for e in block_meta.ss_and_aux_equations_error] + aux_error_size = length(inlined_errors) + aux_error_func! = compile_exprs_to_func(inlined_errors, + ๐”–, ๐”“_ext, placeholder_dict, back_to_array_dict) + end + + desc = "Numerical block $(block_meta.n_block): $(join(string.(block_meta.sorted_vars), ", "))" + + push_numerical_step!(builder; + solve_block = block_meta.solve_block, + block_index = block_meta.n_block, + write_indices = write_indices, + param_gather_indices = param_gather_indices, + var_gather_indices = var_gather_indices, + lbs = block_meta.lbs, + ubs = block_meta.ubs, + aux_func! = aux_func!, + aux_write_indices = aux_write_indices, + aux_error_func! = aux_error_func!, + aux_error_size = aux_error_size, + description = desc, + ) +end + +function write_steady_state_solver_function!(๐“‚::โ„ณ, symbolic_enabled::Bool = false, symbolics_data::Union{Nothing, symbolics} = nothing; + verbose::Bool = false, + avoid_solve::Bool = false) + symbolic_enabled = symbolic_enabled && (symbolics_data !== nothing) + + unknowns = if symbolics_data === nothing + union(๐“‚.constants.post_model_macro.vars_in_ss_equations, ๐“‚.equations.calibration_parameters) + else + union(symbolics_data.calibration_equations_parameters, symbolics_data.vars_in_ss_equations) + end + + n_equations_total = if symbolics_data === nothing + length(๐“‚.equations.steady_state_aux) + length(๐“‚.equations.calibration) + else + length(symbolics_data.ss_equations) + length(symbolics_data.calibration_equations) + end + @assert length(unknowns) <= n_equations_total "Unable to solve steady state. More unknowns than equations." + + incidence_matrix = spzeros(Int, length(unknowns), length(unknowns)) + + eq_list = if symbolics_data === nothing + empty_var_redundant_list = [Symbol[] for _ in eachindex(๐“‚.constants.post_model_macro.var_list_aux_SS)] + vcat( + union.( + setdiff.( + union.( + ๐“‚.constants.post_model_macro.var_list_aux_SS, + ๐“‚.constants.post_model_macro.ss_list_aux_SS, + ), + empty_var_redundant_list, + ), + ๐“‚.constants.post_model_macro.par_list_aux_SS, + ), + union.( + ๐“‚.constants.post_parameters_macro.ss_calib_list, + ๐“‚.constants.post_parameters_macro.par_calib_list, + ), + ) + else + vcat( + union.( + setdiff.( + union.( + symbolics_data.var_list_aux_SS, + symbolics_data.ss_list_aux_SS, + ), + symbolics_data.var_redundant_list, + ), + symbolics_data.par_list_aux_SS, + ), + union.( + symbolics_data.ss_calib_list, + symbolics_data.par_calib_list, + ), + ) + end + + for (i,u) in enumerate(unknowns) + for (k,e) in enumerate(eq_list) + incidence_matrix[i,k] = u โˆˆ e + end + end + + Q, P, R, nmatch, n_blocks = BlockTriangularForm.order(incidence_matrix) + Rฬ‚ = Int[] + for i in 1:n_blocks + [push!(Rฬ‚, n_blocks - i + 1) for ii in R[i]:R[i+1] - 1] + end + push!(Rฬ‚,1) + + vars = hcat(P, Rฬ‚)' + eqs = hcat(Q, Rฬ‚)' + + @assert all(eqs[1,:] .> 0) "Could not solve system of steady state and calibration equations. Number of redundant equations: " * repr(sum(eqs[1,:] .< 0)) * ". Try defining some steady state values as parameters (e.g. r[ss] -> rฬ„). Nonstationary variables are not supported as of now." + + n = n_blocks + + ss_equations = if symbolics_data === nothing + vcat(๐“‚.equations.steady_state_aux, ๐“‚.equations.calibration) + else + vcat(symbolics_data.ss_equations, symbolics_data.calibration_equations) + end + + output_var_names = unique(Symbol.(replace.(string.(sort(union( + ๐“‚.constants.post_model_macro.var, + ๐“‚.constants.post_model_macro.exo_past, + ๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => ""))) + calib_param_names = ๐“‚.equations.calibration_parameters + plus_var_names = Symbol.(๐“‚.constants.post_model_macro.โž•_vars) + all_sol_names = vcat(output_var_names, calib_param_names, plus_var_names) + n_sol = length(all_sol_names) + sol_name_to_index = Dict(name => i for (i, name) in enumerate(all_sol_names)) + plus_var_count_at_start = length(plus_var_names) + + for d in union(๐“‚.constants.post_model_macro.var, ๐“‚.constants.post_model_macro.exo_past, ๐“‚.constants.post_model_macro.exo_future) + raw_name = Symbol(d) + stripped_name = Symbol(replace(string(d), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => "")) + if raw_name != stripped_name && haskey(sol_name_to_index, stripped_name) + sol_name_to_index[raw_name] = sol_name_to_index[stripped_name] + end + end + + output_names_full = vcat( + Symbol.(replace.(string.(sort(union( + ๐“‚.constants.post_model_macro.var, + ๐“‚.constants.post_model_macro.exo_past, + ๐“‚.constants.post_model_macro.exo_future))), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => "")), + calib_param_names + ) + output_indices = [sol_name_to_index[name] for name in output_names_full] + + raw_param_names = collect(๐“‚.constants.post_complete_parameters.parameters) + n_raw_params = length(raw_param_names) + calib_no_var_names = Symbol[expr.args[1] for expr in ๐“‚.equations.calibration_no_var] + ext_param_names = vcat(raw_param_names, calib_no_var_names) + n_ext_params = length(ext_param_names) + ext_param_to_index = Dict(name => i for (i, name) in enumerate(ext_param_names)) + + exo_zero_indices = Int[] + for d in union(๐“‚.constants.post_model_macro.exo_past, ๐“‚.constants.post_model_macro.exo_future) + dns = Symbol(replace(string(d), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => "")) + if haskey(sol_name_to_index, dns) + push!(exo_zero_indices, sol_name_to_index[dns]) + end + end + + n_sol_max = n_sol + 2 * length(ss_equations) + MacroModelling.Symbolics.@variables ๐”–[1:n_sol_max] ๐”“_ext[1:n_ext_params] + + global_placeholder = Dict{Symbol, Symbol}() + global_back_to_array = Dict{MacroModelling.Symbolics.Num, MacroModelling.Symbolics.Num}() + + for (name, idx) in sol_name_to_index + sym = Symbol("๐”–_$idx") + global_placeholder[name] = sym + global_back_to_array[MacroModelling.Symbolics.parse_expr_to_symbolic(sym, @__MODULE__)] = ๐”–[idx] + end + for (name, idx) in ext_param_to_index + sym = Symbol("๐”“e_$idx") + global_placeholder[name] = sym + global_back_to_array[MacroModelling.Symbolics.parse_expr_to_symbolic(sym, @__MODULE__)] = ๐”“_ext[idx] + end + + MacroModelling.Symbolics.@variables P_raw[1:n_raw_params] + + pp_back = Dict{MacroModelling.Symbolics.Num, MacroModelling.Symbolics.Num}() + for i in 1:n_raw_params + sym = Symbol("Praw_$i") + pp_back[MacroModelling.Symbolics.parse_expr_to_symbolic(sym, @__MODULE__)] = P_raw[i] + end + + bounded_param_exprs_for_sub = Dict{Symbol, Union{Symbol, Expr}}() + for (i, par) in enumerate(raw_param_names) + if haskey(๐“‚.constants.post_parameters_macro.bounds, par) + lb, ub = ๐“‚.constants.post_parameters_macro.bounds[par] + bounded_param_exprs_for_sub[par] = :(min(max($(Symbol("Praw_$i")), $lb), $ub)) + else + bounded_param_exprs_for_sub[par] = Symbol("Praw_$i") + end + end + + ext_param_sym_exprs = MacroModelling.Symbolics.Num[] + for (i, par) in enumerate(raw_param_names) + if haskey(๐“‚.constants.post_parameters_macro.bounds, par) + lb, ub = ๐“‚.constants.post_parameters_macro.bounds[par] + push!(ext_param_sym_exprs, min(max(P_raw[i], lb), ub)) + else + push!(ext_param_sym_exprs, P_raw[i]) + end + end + + calib_expr_replacements = Dict{Symbol, Union{Symbol, Expr}}() + for expr in ๐“‚.equations.calibration_no_var + lhs = expr.args[1] + rhs = expr.args[2] + rhs_expanded = replace_symbols(rhs, calib_expr_replacements) + rhs_final = replace_symbols(rhs_expanded, bounded_param_exprs_for_sub) + calib_expr_replacements[lhs] = rhs_final + + sym_expr = MacroModelling.Symbolics.parse_expr_to_symbolic(rhs_final, @__MODULE__) + sym_expr = MacroModelling.Symbolics.substitute(sym_expr, pp_back) + push!(ext_param_sym_exprs, sym_expr) + end + + _, param_prep_func! = MacroModelling.Symbolics.build_function(ext_param_sym_exprs, P_raw, + cse = true, skipzeros = true, + parallel = MacroModelling.Symbolics.SerialForm(), + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + atoms_in_equations = Set{Symbol}() + atoms_in_equations_list = [] + relevant_pars_across = Symbol[] + nsss_solver_cache_init_tmp = [] + + solved_vars = [] + solved_vals = [] + + min_max_error_exprs = [] + unique_โž•_eqs = Dict{Union{Expr,Symbol},Symbol}() + global_solvetime_aux_sub = Dict{Symbol, Union{Symbol, Expr}}() + builder = NSSSSolverBuilder() + numerical_block_count = 0 + + while n > 0 + if length(eqs[:,eqs[2,:] .== n]) == 2 + var_to_solve_for = unknowns[vars[:,vars[2,:] .== n][1]] + + eq_to_solve = ss_equations[eqs[:,eqs[2,:] .== n][1]] + minmax_rewritten = false + + parsed_eq_to_solve_for = eq_to_solve |> string |> Meta.parse + + minmax_fixed_eqs = postwalk(x -> + x isa Expr ? + x.head == :call ? + x.args[1] โˆˆ [:Max,:Min] ? + Symbol(var_to_solve_for) โˆˆ get_symbols(x.args[2]) ? + x.args[2] : + Symbol(var_to_solve_for) โˆˆ get_symbols(x.args[3]) ? + x.args[3] : + x : + x : + x : + x, + parsed_eq_to_solve_for) + + if parsed_eq_to_solve_for != minmax_fixed_eqs + [push!(atoms_in_equations, a) for a in setdiff(get_symbols(parsed_eq_to_solve_for), get_symbols(minmax_fixed_eqs))] + push!(min_max_error_exprs, parsed_eq_to_solve_for) + eq_to_solve = minmax_fixed_eqs + minmax_rewritten = true + end + + if symbolics_data === nothing || avoid_solve || minmax_rewritten || count_ops(Meta.parse(string(eq_to_solve))) > 15 + soll = nothing + else + if eq_to_solve isa SPyPyC.Sym{PythonCall.Core.Py} && var_to_solve_for isa SPyPyC.Sym{PythonCall.Core.Py} + soll = solve_symbolically(eq_to_solve, var_to_solve_for) + else + soll = nothing + end + end + + if isnothing(soll) || isempty(soll) + if verbose && symbolic_enabled + println("Failed finding solution symbolically for: ",var_to_solve_for," in: ",eq_to_solve) + end + + eq_idx_in_block_to_solve = eqs[:,eqs[2,:] .== n][1,:] + + numerical_block_count += 1 + block_meta = write_block_solution!(๐“‚, [var_to_solve_for], [eq_to_solve], relevant_pars_across, nsss_solver_cache_init_tmp, eq_idx_in_block_to_solve, atoms_in_equations_list, solved_vars, solved_vals, block_index = numerical_block_count) + + current_plus_count = length(๐“‚.constants.post_model_macro.โž•_vars) + if current_plus_count > plus_var_count_at_start + for pvi in (plus_var_count_at_start + 1):current_plus_count + pv = Symbol(๐“‚.constants.post_model_macro.โž•_vars[pvi]) + if !haskey(sol_name_to_index, pv) + push!(all_sol_names, pv) + idx = length(all_sol_names) + sol_name_to_index[pv] = idx + sym = Symbol("๐”–_$idx") + global_placeholder[pv] = sym + global_back_to_array[MacroModelling.Symbolics.parse_expr_to_symbolic(sym, @__MODULE__)] = ๐”–[idx] + end + end + plus_var_count_at_start = current_plus_count + end + + append_numerical_step!(builder, block_meta, sol_name_to_index, ext_param_to_index, + ๐”–, ๐”“_ext, global_placeholder, global_back_to_array, global_solvetime_aux_sub) + + elseif soll[1].is_number == true + if var_to_solve_for isa SPyPyC.Sym{PythonCall.Core.Py} && soll[1] isa SPyPyC.Sym{PythonCall.Core.Py} + ss_equations = [eq isa SPyPyC.Sym{PythonCall.Core.Py} ? replace_symbolic(eq, var_to_solve_for, soll[1]) : eq for eq in ss_equations] + end + + push!(solved_vars, Symbol(var_to_solve_for)) + push!(solved_vals, Meta.parse(string(soll[1]))) + push!(atoms_in_equations_list, []) + + var_name = solved_vars[end] + val = solved_vals[end] + widx = sol_name_to_index[var_name] + + if var_name โˆˆ ๐“‚.constants.post_model_macro.โž•_vars + step_expr = :(max(eps(), $val)) + eval_func! = compile_exprs_to_func([step_expr], ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + else + constant_value = Float64(soll[1]) + eval_func! = let constant_value = constant_value + (out, _sol_vec, _params_vec) -> begin + out[1] = constant_value + return nothing + end + end + end + + push_analytical_step!(builder; + eval_func! = eval_func!, + write_indices = [widx], + description = "Constant: $var_name = $val", + ) + + else + push!(solved_vars, Symbol(var_to_solve_for)) + push!(solved_vals, Meta.parse(string(soll[1]))) + + [push!(atoms_in_equations, Symbol(a)) for a in soll[1].atoms()] + push!(atoms_in_equations_list, Set(union(setdiff(get_symbols(parsed_eq_to_solve_for), get_symbols(minmax_fixed_eqs)), Symbol.(soll[1].atoms())))) + + var_name = solved_vars[end] + val_expr = solved_vals[end] + widx = sol_name_to_index[var_name] + + if var_name โˆˆ ๐“‚.constants.post_model_macro.โž•_vars + bounds_tuple = get(๐“‚.constants.post_parameters_macro.bounds, var_name, (eps(), 1e12)) + lb, ub = Float64(bounds_tuple[1]), Float64(bounds_tuple[2]) + + eval_func! = compile_exprs_to_func([val_expr], ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + + push_analytical_step!(builder; + eval_func! = eval_func!, + write_indices = [widx], + lower_bounds = [lb], + upper_bounds = [ub], + has_bounds = trues(1), + description = "Analytical โž•: $var_name", + ) + + unique_โž•_eqs[val_expr] = var_name + else + vars_to_exclude = [vcat(Symbol.(var_to_solve_for), ๐“‚.constants.post_model_macro.โž•_vars), Symbol[]] + + rewritten_eqs, ss_and_aux_equations, ss_and_aux_equations_dep, ss_and_aux_equations_error, ss_and_aux_equations_error_dep = make_equation_robust_to_domain_errors([val_expr], vars_to_exclude, ๐“‚.constants.post_parameters_macro.bounds, ๐“‚.constants.post_model_macro.โž•_vars, unique_โž•_eqs) + + current_plus_count = length(๐“‚.constants.post_model_macro.โž•_vars) + if current_plus_count > plus_var_count_at_start + for pvi in (plus_var_count_at_start + 1):current_plus_count + pv = Symbol(๐“‚.constants.post_model_macro.โž•_vars[pvi]) + if !haskey(sol_name_to_index, pv) + push!(all_sol_names, pv) + idx = length(all_sol_names) + sol_name_to_index[pv] = idx + sym = Symbol("๐”–_$idx") + global_placeholder[pv] = sym + global_back_to_array[MacroModelling.Symbolics.parse_expr_to_symbolic(sym, @__MODULE__)] = ๐”–[idx] + end + end + plus_var_count_at_start = current_plus_count + end + + all_aux_eqs = vcat(ss_and_aux_equations, ss_and_aux_equations_dep) + all_aux_errors = vcat(ss_and_aux_equations_error, ss_and_aux_equations_error_dep) + + aux_func! = NOOP_NSSS_FUNC! + aux_write_indices = Int[] + error_func! = NOOP_NSSS_FUNC! + error_size = 0 + + model_aux_names = Symbol[] + model_aux_rhs = Any[] + model_aux_sub = Dict{Symbol, Any}() + + for eq in all_aux_eqs + if eq isa Expr && eq.head == :(=) + lhs = eq.args[1] + rhs = eq.args[2] + expanded_rhs = isempty(global_solvetime_aux_sub) ? rhs : replace_symbols(rhs, global_solvetime_aux_sub) + expanded_rhs = isempty(model_aux_sub) ? expanded_rhs : replace_symbols(expanded_rhs, model_aux_sub) + if haskey(sol_name_to_index, lhs) + push!(model_aux_names, lhs) + push!(model_aux_rhs, expanded_rhs) + model_aux_sub[lhs] = expanded_rhs + else + global_solvetime_aux_sub[lhs] = expanded_rhs + end + end + end + + if !isempty(model_aux_rhs) + aux_write_indices = [sol_name_to_index[v] for v in model_aux_names] + aux_func! = compile_exprs_to_func(model_aux_rhs, ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + end + + main_expr = isempty(global_solvetime_aux_sub) ? rewritten_eqs[1] : replace_symbols(rewritten_eqs[1], global_solvetime_aux_sub) + eval_func! = compile_exprs_to_func([main_expr], ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + + if !isempty(all_aux_errors) + inlined_errors = isempty(global_solvetime_aux_sub) ? all_aux_errors : [replace_symbols(e, global_solvetime_aux_sub) for e in all_aux_errors] + error_size = length(inlined_errors) + error_func! = compile_exprs_to_func(inlined_errors, ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + end + + has_user_bounds = haskey(๐“‚.constants.post_parameters_macro.bounds, var_name) && var_name โˆ‰ ๐“‚.constants.post_model_macro.โž•_vars + if has_user_bounds + lb = Float64(๐“‚.constants.post_parameters_macro.bounds[var_name][1]) + ub = Float64(๐“‚.constants.post_parameters_macro.bounds[var_name][2]) + push_analytical_step!(builder; + aux_func! = aux_func!, + aux_write_indices = aux_write_indices, + error_func! = error_func!, + error_size = error_size, + eval_func! = eval_func!, + write_indices = [widx], + lower_bounds = [lb], + upper_bounds = [ub], + has_bounds = trues(1), + description = "Analytical bounded: $var_name", + ) + else + push_analytical_step!(builder; + aux_func! = aux_func!, + aux_write_indices = aux_write_indices, + error_func! = error_func!, + error_size = error_size, + eval_func! = eval_func!, + write_indices = [widx], + description = "Analytical: $var_name", + ) + end + end + end + else + vars_to_solve = unknowns[vars[:,vars[2,:] .== n][1,:]] + eqs_to_solve = ss_equations[eqs[:,eqs[2,:] .== n][1,:]] + + numerical_sol = false + + if symbolic_enabled + if avoid_solve || count_ops(Meta.parse(string(eqs_to_solve))) > 15 + soll = nothing + else + soll = solve_symbolically(eqs_to_solve::Vector{SPyPyC.Sym{PythonCall.Core.Py}}, vars_to_solve::Vector{SPyPyC.Sym{PythonCall.Core.Py}}) + end + + if isnothing(soll) || isempty(soll) || length(intersect((union(SPyPyC.free_symbols.(collect(values(soll)))...) .|> SPyPyC.:โ†“),(vars_to_solve .|> SPyPyC.:โ†“))) > 0 + if verbose println("Failed finding solution symbolically for: ",vars_to_solve," in: ",eqs_to_solve,". Solving numerically.") end + numerical_sol = true + else + if verbose println("Solved: ",string.(eqs_to_solve)," for: ",Symbol.(vars_to_solve), " symbolically.") end + + atoms = reduce(union,map(x->x.atoms(),collect(values(soll)))) + for a in atoms push!(atoms_in_equations, Symbol(a)) end + + step_exprs = [] + step_write_indices = Int[] + + for v in vars_to_solve + push!(solved_vars, Symbol(v)) + push!(solved_vals, Meta.parse(string(soll[v]))) + push!(atoms_in_equations_list, Set(Symbol.(soll[v].atoms()))) + push!(step_exprs, solved_vals[end]) + push!(step_write_indices, sol_name_to_index[Symbol(v)]) + end + + eval_func! = compile_exprs_to_func(step_exprs, ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + + push_analytical_step!(builder; + eval_func! = eval_func!, + write_indices = step_write_indices, + description = "Analytical multi: $(join(string.(Symbol.(vars_to_solve)), ", "))", + ) + end + end + + eq_idx_in_block_to_solve = eqs[:,eqs[2,:] .== n][1,:] + incidence_matrix_subset = incidence_matrix[vars[:,vars[2,:] .== n][1,:], eq_idx_in_block_to_solve] + + if numerical_sol || !symbolic_enabled + vars_to_solve_reduced = vars_to_solve + eqs_to_solve_reduced = eqs_to_solve + eq_idx_in_block_to_solve_reduced = eq_idx_in_block_to_solve + + numerical_block_count += 1 + block_meta = write_block_solution!(๐“‚, vars_to_solve_reduced, eqs_to_solve_reduced, relevant_pars_across, nsss_solver_cache_init_tmp, eq_idx_in_block_to_solve_reduced, atoms_in_equations_list, solved_vars, solved_vals, block_index = numerical_block_count) + + if !isnothing(block_meta) + current_plus_count = length(๐“‚.constants.post_model_macro.โž•_vars) + if current_plus_count > plus_var_count_at_start + for pvi in (plus_var_count_at_start + 1):current_plus_count + pv = Symbol(๐“‚.constants.post_model_macro.โž•_vars[pvi]) + if !haskey(sol_name_to_index, pv) + push!(all_sol_names, pv) + idx = length(all_sol_names) + sol_name_to_index[pv] = idx + sym = Symbol("๐”–_$idx") + global_placeholder[pv] = sym + global_back_to_array[MacroModelling.Symbolics.parse_expr_to_symbolic(sym, @__MODULE__)] = ๐”–[idx] + end + end + plus_var_count_at_start = current_plus_count + end + + append_numerical_step!(builder, block_meta, sol_name_to_index, ext_param_to_index, + ๐”–, ๐”“_ext, global_placeholder, global_back_to_array, global_solvetime_aux_sub) + end + + if !symbolic_enabled && verbose + println("Solved: ",string.(eqs_to_solve)," for: ",Symbol.(vars_to_solve), " numerically.") + end + end + end + n -= 1 + end + + push!(nsss_solver_cache_init_tmp, fill(Inf, length(๐“‚.constants.post_complete_parameters.parameters))) + push!(๐“‚.caches.solver, nsss_solver_cache_init_tmp) + + parameters_only_in_par_defs = Set() + if length(๐“‚.equations.calibration_no_var) > 0 + atoms = reduce(union, get_symbols.(๐“‚.equations.calibration_no_var)) + [push!(atoms_in_equations, a) for a in atoms] + [push!(parameters_only_in_par_defs, a) for a in atoms] + end + + dependencies = [] + for (i, a) in enumerate(atoms_in_equations_list) + push!(dependencies, solved_vars[i] => intersect(a, union(๐“‚.constants.post_model_macro.var, ๐“‚.constants.post_complete_parameters.parameters))) + end + + push!(dependencies, :SS_relevant_calibration_parameters => intersect(reduce(union, atoms_in_equations_list), ๐“‚.constants.post_complete_parameters.parameters)) + if !isempty(min_max_error_exprs) + minmax_error_func! = compile_exprs_to_func(min_max_error_exprs, ๐”–, ๐”“_ext, global_placeholder, global_back_to_array) + n_errors = length(min_max_error_exprs) + push_analytical_step!(builder; + error_func! = minmax_error_func!, + error_size = n_errors, + eval_func! = compile_exprs_to_func([0.0], ๐”–, ๐”“_ext, global_placeholder, global_back_to_array), + write_indices = Int[], + description = "Min/Max validation", + ) + end + + # Patch bounds on โž• steps in the builder's flat arrays + if !isempty(๐“‚.constants.post_parameters_macro.bounds) + for i in 1:length(builder.step_types) + if builder.step_types[i] == ANALYTICAL_STEP && startswith(builder.descriptions[i], "Analytical โž•:") + wr = builder.write_ranges[i] + br = builder.bounds_ranges[i] + for (j_local, j_wr) in enumerate(wr) + widx = builder.write_indices[j_wr] + name = all_sol_names[widx] + if haskey(๐“‚.constants.post_parameters_macro.bounds, name) + bt = ๐“‚.constants.post_parameters_macro.bounds[name] + j_br = br[j_local] + builder.lower_bounds[j_br] = Float64(bt[1]) + builder.upper_bounds[j_br] = Float64(bt[2]) + builder.has_bounds[j_br] = true + end + end + end + end + end + + build_nsss_solver!(๐“‚, builder, param_prep_func!) + n_sol = length(all_sol_names) + ๐“‚.constants.post_complete_parameters = update_post_complete_parameters( + ๐“‚.constants.post_complete_parameters; + nsss_dependencies = dependencies, + nsss_n_sol = n_sol, + nsss_output_indices = output_indices, + nsss_n_ext_params = n_ext_params, + nsss_sol_names = all_sol_names, + nsss_exo_zero_indices = exo_zero_indices, + nsss_param_names_ext = ext_param_names, + ) + + return nothing +end + +function find_closest_solution(cache, initial_parameters::Vector{Float64}, expected_length::Int) + current_best = Inf + closest_solution = cache[end] + + target_parameters_norm_squared = 0.0 + @inbounds for i in eachindex(initial_parameters) + pi = initial_parameters[i] + target_parameters_norm_squared += pi * pi + end + + @inbounds for idx in length(cache):-1:1 + pars = cache[idx] + if length(pars) < expected_length || !(pars[end] isa Vector{Float64}) || length(pars[end]) != length(initial_parameters) + continue + end + + cached_parameters = pars[end] + squared_distance = 0.0 + cached_parameters_norm_squared = 0.0 + for i in eachindex(initial_parameters) + ci = cached_parameters[i] + d = ci - initial_parameters[i] + squared_distance += d * d + cached_parameters_norm_squared += ci * ci + end + + normalisation_norm_squared = max(target_parameters_norm_squared, cached_parameters_norm_squared) + relative_parameter_distance_squared = squared_distance / normalisation_norm_squared + + if relative_parameter_distance_squared < eps() + return squared_distance, pars + end + + if squared_distance < current_best + current_best = squared_distance + closest_solution = pars + end + end + + if !isfinite(current_best) + if (closest_solution[end] isa Vector{Float64}) && (length(closest_solution[end]) == length(initial_parameters)) + cached_parameters = closest_solution[end] + current_best = 0.0 + @inbounds for i in eachindex(initial_parameters) + d = cached_parameters[i] - initial_parameters[i] + current_best += d * d + end + else + current_best = Inf + end + end + + return current_best, closest_solution +end + +""" + execute_step!(step_idx, sol_vec, params_vec, closest_solution, ๐“‚, ...) + +Execute a single NSSS solve step. +Dispatches on `๐“‚.constants.nsss_solver.step_types[step_idx]` (ANALYTICAL_STEP or NUMERICAL_STEP). + +Uses shared workspace buffers for scratch computations, avoiding per-step allocation. + +Returns: (error, iterations, cache_entries::Vector{Vector{Float64}}) +""" +function execute_step!(step_idx::Int, + sol_vec::Vector{Float64}, params_vec::Vector{Float64}, + closest_solution, ๐“‚, tol, fail_fast_solvers_only, + cold_start, solver_parameters, preferred_solver_parameter_idx::Int, verbose) + + c = ๐“‚.constants.nsss_solver + f = ๐“‚.functions.nsss_solver + w = ๐“‚.workspaces.nsss_solver + step_type = c.step_types[step_idx] + + error = 0.0 + + # Phase 1: Compute auxiliary variables (shared across both step types) + aux_wr = c.aux_write_ranges[step_idx] + n_aux = length(aux_wr) + if n_aux > 0 + aux_buf = @view w.aux_buffer[1:n_aux] + f.aux_funcs[step_idx](aux_buf, sol_vec, params_vec) + @inbounds for j in 1:n_aux + sol_vec[c.aux_write_indices[aux_wr[j]]] = aux_buf[j] + end + end + + if step_type == ANALYTICAL_STEP + # Error check (analytical domain-safety) + err_n = c.error_sizes[step_idx] + if err_n > 0 + err_buf = @view w.error_buffer[1:err_n] + f.error_funcs[step_idx](err_buf, sol_vec, params_vec) + error += sum(abs, err_buf) + end + + # Main evaluation + wr = c.write_ranges[step_idx] + n_write = length(wr) + if n_write > 0 + main_buf = @view w.main_buffer[1:n_write] + f.eval_funcs[step_idx](main_buf, sol_vec, params_vec) + br = c.bounds_ranges[step_idx] + @inbounds for j in 1:n_write + raw = main_buf[j] + widx = c.write_indices[wr[j]] + if !isempty(br) && c.has_bounds[br[j]] + clamped = clamp(raw, c.lower_bounds[br[j]], c.upper_bounds[br[j]]) + error += abs(clamped - raw) + sol_vec[widx] = clamped + else + sol_vec[widx] = raw + end + end + else + # Min/Max validation step: no writes but eval_func exists + f.eval_funcs[step_idx](@view(w.main_buffer[1:1]), sol_vec, params_vec) + end + + return error, 0, EMPTY_NSSS_STEP_CACHE + + else # NUMERICAL_STEP + # Gather params_and_solved_vars into shared main_buffer + pgr = c.param_gather_ranges[step_idx] + vgr = c.var_gather_ranges[step_idx] + n_params = length(pgr) + n_vars = length(vgr) + gather_size = n_params + n_vars + + params_and_solved_vars = w.params_and_solved_vars_buffer + resize!(params_and_solved_vars, gather_size) + @inbounds for j in 1:n_params + params_and_solved_vars[j] = params_vec[c.param_gather_indices[pgr[j]]] + end + @inbounds for j in 1:n_vars + params_and_solved_vars[n_params + j] = sol_vec[c.var_gather_indices[vgr[j]]] + end + + # Build initial guesses + block_idx = c.block_indices[step_idx] + cache_sol_idx = 2*(block_idx-1)+1 + cache_par_idx = 2*block_idx + cache_sol = cache_sol_idx <= length(closest_solution) ? closest_solution[cache_sol_idx] : Float64[] + cache_par = cache_par_idx <= length(closest_solution) ? closest_solution[cache_par_idx] : Float64[Inf] + + wr = c.write_ranges[step_idx] + n_write = length(wr) + nbr = c.numerical_bounds_ranges[step_idx] + guess_len = min(n_write, length(nbr)) + + guess_buf = @view w.guess_buffer[1:guess_len] + copy_len = min(length(cache_sol), guess_len) + @inbounds for i in 1:copy_len + guess_buf[i] = clamp(cache_sol[i], c.numerical_lbs[nbr[i]], c.numerical_ubs[nbr[i]]) + end + @inbounds for i in (copy_len + 1):guess_len + guess_buf[i] = clamp(0.5 * (c.numerical_lbs[nbr[i]] + c.numerical_ubs[nbr[i]]), + c.numerical_lbs[nbr[i]], c.numerical_ubs[nbr[i]]) + end + + # Use workspace inits container + resize!(w.inits[1], guess_len) + if guess_len > 0 + copyto!(w.inits[1], 1, guess_buf, 1, guess_len) + end + w.inits[2] = cache_par + + lbs = w.lbs_buffer + ubs = w.ubs_buffer + n_bounds = length(nbr) + resize!(lbs, n_bounds) + resize!(ubs, n_bounds) + @inbounds for i in 1:n_bounds + lbs[i] = c.numerical_lbs[nbr[i]] + ubs[i] = c.numerical_ubs[nbr[i]] + end + + # Call block solver + solve_block = f.solve_blocks[step_idx] + if solve_block === nothing + if verbose + println("Missing numerical solve block for step $(step_idx)") + end + return Inf, 0, EMPTY_NSSS_STEP_CACHE + end + + solution = block_solver( + params_and_solved_vars, + block_idx, + solve_block, + w.inits, + lbs, + ubs, + solver_parameters, + preferred_solver_parameter_idx, + fail_fast_solvers_only, + cold_start, + verbose + ) + + error += solution[2][1] + iters = solution[2][2] + if error > tol.nsss.acceptance_tol + if verbose + println("Failed after solving block with error $error") + end + return error, iters, EMPTY_NSSS_STEP_CACHE + end + + # Domain safety error check after block solve + err_n = c.aux_error_sizes[step_idx] + if err_n > 0 + err_buf = @view w.error_buffer[1:err_n] + f.error_funcs[step_idx](err_buf, sol_vec, params_vec) + error += sum(abs, err_buf) + if error > tol.nsss.acceptance_tol + if verbose + println("Failed for aux variables with error $error") + end + return error, iters, EMPTY_NSSS_STEP_CACHE + end + end + + # Write results to solution vector + sol = solution[1] + @inbounds for j in 1:n_write + sol_vec[c.write_indices[wr[j]]] = sol[j] + end + + # Build cache entries for this block + cache_entries = [ + typeof(sol) == Vector{Float64} ? copy(sol) : โ„ฑ.value.(sol), + typeof(params_and_solved_vars) == Vector{Float64} ? copy(params_and_solved_vars) : โ„ฑ.value.(params_and_solved_vars) + ] + + return error, iters, cache_entries + end +end + + +# ============================================================================ +# Orchestrator: solve_nsss_steps +# ============================================================================ + +""" + solve_nsss_steps(parameters, ๐“‚, tol, verbose, fail_fast_solvers_only, + closest_solution, cold_start, solver_params) + +Solve the NSSS by executing pipeline steps in a single pass. + +Steps are dispatched via `execute_step!` which uses the pipeline's shared +workspace buffers. Steps are executed in order, filling the solution vector +progressively. + +Returns: (SS_and_pars, (solution_error, iters), nsss_solver_cache_tmp) +""" +function solve_nsss_steps( + parameters::Vector{Float64}, + ๐“‚::โ„ณ, + tol::Tolerances, + verbose::Bool, + fail_fast_solvers_only::Bool, + closest_solution, + cold_start::Bool, + solver_params::Vector{solver_parameters}, + preferred_solver_parameter_idx::Int +) + nsss_n_ext_params = ๐“‚.constants.post_complete_parameters.nsss_n_ext_params + nsss_n_sol = ๐“‚.constants.post_complete_parameters.nsss_n_sol + nsss_output_indices = ๐“‚.constants.post_complete_parameters.nsss_output_indices + nsss_consts = ๐“‚.constants.nsss_solver + nsss_ws = ๐“‚.workspaces.nsss_solver + + # Prepare extended parameter vector (raw params โ†’ bounded + calibration_no_var) + params_vec = nsss_ws.params_vec_buffer + if length(params_vec) != nsss_n_ext_params + resize!(params_vec, nsss_n_ext_params) + end + ๐“‚.functions.nsss_param_prep!(params_vec, parameters) + + # Initialize solution vector from workspace buffer + sol_vec = nsss_ws.sol_vec_buffer + if length(sol_vec) != nsss_n_sol + resize!(sol_vec, nsss_n_sol) + end + fill!(sol_vec, 0.0) + + # Single pass through all steps + nsss_solver_cache_tmp = Vector{Float64}[] + solution_error = 0.0 + iters = 0 + + n_steps = nsss_consts.n_steps + for step_idx in 1:n_steps + step_error, step_iters, step_cache = execute_step!( + step_idx, sol_vec, params_vec, closest_solution, ๐“‚, tol, + fail_fast_solvers_only, cold_start, solver_params, preferred_solver_parameter_idx, verbose + ) + + solution_error += step_error + iters += step_iters + if !isempty(step_cache) + append!(nsss_solver_cache_tmp, step_cache) + end + + if solution_error > tol.nsss.acceptance_tol + if verbose + println("Step '$(nsss_consts.descriptions[step_idx])' failed with accumulated error $solution_error") + end + break + end + end + + # Build SS_and_pars from solution vector into reusable output buffer + SS_and_pars = nsss_ws.output_buffer + n_output = length(nsss_output_indices) + if length(SS_and_pars) != n_output + resize!(SS_and_pars, n_output) + end + + if solution_error >= tol.nsss.acceptance_tol + fill!(SS_and_pars, 0.0) + else + @inbounds for i in 1:n_output + SS_and_pars[i] = sol_vec[nsss_output_indices[i]] + end + end + + # Append parameters to cache + parameters_copy = copy(parameters) + push!(nsss_solver_cache_tmp, parameters_copy) + + return SS_and_pars, (solution_error, iters), nsss_solver_cache_tmp +end + + +# ============================================================================ +# Wrapper: solve_nsss_wrapper (handles cache + continuation method) +# ============================================================================ + +""" + solve_nsss_wrapper( + parameter_values::Vector{<:Real}, + ๐“‚::โ„ณ, + tol::Tolerances, + verbose::Bool, + cold_start::Bool, + solver_params::Vector{solver_parameters} + )::Tuple{Vector, Tuple{Real, Int}} + +Normal Julia function wrapper for NSSS solving. + +This function handles cache management and continuation scaling for solving +the non-stochastic steady state using the step-based NSSS solver. + +The continuation method gradually transitions from a cached solution to the +target parameters using a scaling approach, which improves convergence. + +# Arguments +- `parameter_values`: Parameter values to solve at +- `๐“‚`: Model structure +- `tol`: Tolerance settings +- `verbose`: Whether to print verbose output +- `cold_start`: Whether this is a cold start (limits iterations to 1) +- `solver_params`: Solver configuration + +# Keyword arguments +- `continuation_cache_capacity`: Size of local continuation cache buffer +- `continuation_max_iters`: Maximum continuation iterations for warm starts +- `stall_tolerance`: Threshold to stop when continuation scale no longer moves +- `cache_push_distance_tol`: Distance threshold before pushing solved cache to model cache +- `scale_snap_threshold`: Scale above which continuation snaps directly to `1.0` +- `scale_success_weight`: Weight on current scale after successful continuation step +- `scale_failure_weight`: Weight on current scale after failed continuation step + +# Returns +- Tuple of (solution_vector, (solution_error, iterations)) +""" +function solve_nsss_wrapper( + parameter_values::Vector{<:Real}, + ๐“‚::โ„ณ, + tol::Tolerances, + verbose::Bool, + cold_start::Bool, + solver_params::Vector{solver_parameters} + ; + continuation_cache_capacity::Int = 500, + continuation_max_iters::Int = 500, + stall_tolerance::Float64 = 1e-2, + cache_push_distance_tol::Float64 = 1e-8, + scale_snap_threshold::Float64 = 0.95, + scale_success_weight::Float64 = 0.4, + scale_failure_weight::Float64 = 0.3, + preferred_solver_parameter_idx::Int = 1, +)::Tuple{Vector, Tuple{Real, Int}} + + n_numerical_steps = count(==(NUMERICAL_STEP), ๐“‚.constants.nsss_solver.step_types) + + # Type conversion for AD compatibility + initial_parameters = parameter_values isa Vector{Float64} ? + parameter_values : + โ„ฑ.value.(parameter_values) + + # Find closest cached solution as starting point + expected_cache_length = 2 * n_numerical_steps + 1 + _, closest_solution_init = find_closest_solution(๐“‚.caches.solver, initial_parameters, expected_cache_length) + + # Initialize continuation method variables + range_iters = 0 + solution_error = 1.0 + solved_scale = 0.0 + scale = 1.0 + SS_and_pars = Float64[] + + nsss_ws = ๐“‚.workspaces.nsss_solver + if nsss_ws.continuation_capacity != continuation_cache_capacity + nsss_ws.continuation = CircularBuffer{Vector{Vector{Float64}}}(continuation_cache_capacity) + nsss_ws.continuation_capacity = continuation_cache_capacity + else + empty!(nsss_ws.continuation) + end + + continuation_cache = nsss_ws.continuation + push!(continuation_cache, closest_solution_init) + scaled_parameters = nsss_ws.scaled_parameters_buffer + if length(scaled_parameters) != length(initial_parameters) + resize!(scaled_parameters, length(initial_parameters)) + end + + # Continuation method: iterate with scaling to gradually approach target + max_iters = cold_start ? 1 : continuation_max_iters + n_solver_parameters = length(solver_params) + @assert n_solver_parameters > 0 "At least one steady-state solver parameter set is required." + preferred_idx = clamp(preferred_solver_parameter_idx, 1, n_solver_parameters) + + while range_iters <= max_iters && !(solution_error < tol.nsss.acceptance_tol && solved_scale == 1) + range_iters += 1 + fail_fast_solvers_only = range_iters > 1 + + # Stall detection: stop if scale hasn't moved + if abs(solved_scale - scale) < stall_tolerance + break + end + + # Find closest solution from local intermediate cache + current_best, closest_solution = find_closest_solution(continuation_cache, initial_parameters, expected_cache_length) + + # Interpolate parameters between target and cached solution + if all(isfinite, closest_solution[end]) && initial_parameters != closest_solution_init[end] + @inbounds for i in eachindex(initial_parameters) + scaled_parameters[i] = scale * initial_parameters[i] + (1 - scale) * closest_solution_init[end][i] + end + parameters = scaled_parameters + else + parameters = initial_parameters + end + + # Call step-based solver + SS_and_pars, (solution_error, iters), nsss_solver_cache_tmp = solve_nsss_steps( + parameters, + ๐“‚, + tol, + verbose, + fail_fast_solvers_only, + closest_solution, + cold_start, + solver_params, + preferred_idx + ) + + # Check convergence and update scaling + if solution_error < tol.nsss.acceptance_tol + solved_scale = scale + + if scale == 1 + if current_best > cache_push_distance_tol + push!(๐“‚.caches.solver, nsss_solver_cache_tmp) + end + return SS_and_pars, (solution_error, iters) + end + + # Cache intermediate result for warm starts + push!(continuation_cache, nsss_solver_cache_tmp) + + # Advance scale toward 1.0 + if scale > scale_snap_threshold + scale = 1.0 + else + scale = scale * scale_success_weight + (1 - scale_success_weight) + end + else + # Failed: pull scale back toward last successful scale + scale = scale * scale_failure_weight + solved_scale * (1 - scale_failure_weight) + end + end + + # Failed to converge - return zeros with matching output length + n_output = length(๐“‚.constants.post_complete_parameters.nsss_output_indices) + SS_and_pars = nsss_ws.output_buffer + if length(SS_and_pars) != n_output + resize!(SS_and_pars, n_output) + end + fill!(SS_and_pars, 0.0) + + return SS_and_pars, (1.0, 0) +end diff --git a/src/options_and_caches.jl b/src/options_and_caches.jl index eac752f4b..a3d6a76c9 100644 --- a/src/options_and_caches.jl +++ b/src/options_and_caches.jl @@ -9,14 +9,27 @@ See [`second_order_indices`](@ref) for field documentation. """ function Second_order_indices() empty_sparse_int = SparseMatrixCSC{Int, Int64}(โ„’.I, 0, 0) + empty_sparse_bool = spzeros(Bool, 0, 0) empty_sparse_float = spzeros(Float64, 0, 0) empty_matrix_float = Matrix{Float64}(undef, 0, 0) return second_order_indices( - # Auxiliary matrices (๐›”, ๐‚โ‚‚, ๐”โ‚‚, ๐”โˆ‡โ‚‚) + # Auxiliary matrices (๐›”, ๐›”_sym, ๐›”cโ‚‚, ๐›”๐‚โ‚‚, ๐‚โ‚‚, ๐”โ‚‚, ๐”โˆ‡โ‚‚, ๐ˆโ‚™โ‚Š, ๐ˆโ‚™โ‚‹) empty_sparse_int, empty_sparse_int, empty_sparse_int, empty_sparse_int, + empty_sparse_int, + empty_sparse_int, + empty_sparse_int, + empty_sparse_int, + empty_sparse_int, + Int[], # โˆ‡โ‚‚_nonempty_col_as_kron_rowmask + Int[], # ๐›”๐‚โ‚‚_nonempty_row_as_kron_colmask + # Pre-transposed constants for rrule pullback + empty_sparse_int, # ๐›”แต€ + empty_sparse_int, # ๐‚โ‚‚แต€ + empty_sparse_int, # ๐”โ‚‚แต€ + empty_sparse_int, # ๐”โˆ‡โ‚‚แต€ # Computational index caches (BitVectors) BitVector(), # s_in_sโบ BitVector(), # s_in_s @@ -63,13 +76,15 @@ function Third_order_indices() empty_sparse_int = SparseMatrixCSC{Int, Int64}(โ„’.I, 0, 0) empty_matrix_float = Matrix{Float64}(undef, 0, 0) return third_order_indices( - # Auxiliary matrices (๐‚โ‚ƒ, ๐”โ‚ƒ, ๐ˆโ‚ƒ, ๐‚โˆ‡โ‚ƒ, ๐”โˆ‡โ‚ƒ, ๐, ๐โ‚โ‚—, ๐โ‚แตฃ, ...) + # Auxiliary matrices (๐‚โ‚ƒ, ๐”โ‚ƒ, ๐ˆโ‚ƒ, ๐‚โˆ‡โ‚ƒ, ๐”โˆ‡โ‚ƒ, ๐, ๐๐‚โ‚ƒ, ๐โ‚โ‚—, ๐โ‚แตฃ, ...) empty_sparse_int, # ๐‚โ‚ƒ empty_sparse_int, # ๐”โ‚ƒ Dict{Vector{Int}, Int}(), # ๐ˆโ‚ƒ empty_sparse_int, # ๐‚โˆ‡โ‚ƒ empty_sparse_int, # ๐”โˆ‡โ‚ƒ + Int[], # โˆ‡โ‚ƒ_rowmask empty_sparse_int, # ๐ + empty_sparse_int, # ๐๐‚โ‚ƒ empty_sparse_int, # ๐โ‚โ‚— empty_sparse_int, # ๐โ‚แตฃ empty_sparse_int, # ๐โ‚โ‚—ฬ‚ @@ -79,6 +94,16 @@ function Third_order_indices() empty_sparse_int, # ๐โ‚แตฃฬƒ empty_sparse_int, # ๐โ‚‚แตฃฬƒ empty_sparse_int, # ๐’๐ + # Pre-transposed constants for rrule pullback + empty_sparse_int, # ๐‚โ‚ƒแต€ + empty_sparse_int, # ๐”โ‚ƒแต€ + empty_sparse_int, # ๐๐‚โ‚ƒแต€ + empty_sparse_int, # ๐โ‚โ‚—แต€ + empty_sparse_int, # ๐โ‚แตฃแต€ + empty_sparse_int, # ๐โ‚โ‚—ฬ„แต€ + empty_sparse_int, # ๐โ‚‚โ‚—ฬ„แต€ + empty_sparse_int, # ๐โ‚แตฃฬƒแต€ + empty_sparse_int, # ๐โ‚‚แตฃฬƒแต€ # Conditional forecast index caches Int[], # var_volยณ_idxs Int[], # shock_idxs2 @@ -151,6 +176,7 @@ function Sylvester_workspace(;S::Type = Float64, T::Type = Float64) zeros(S,0,0), # ๐‚ยน (doubling) zeros(S,0,0), # ๐‚B (doubling) Krylov_workspace(S = S), + zeros(S,0,0), # P (stable primal cache) # ForwardDiff partials buffers zeros(T,0,0), # Pฬƒ zeros(T,0,0), # Aฬƒ_fd @@ -187,8 +213,12 @@ function Higher_order_workspace(;T::Type = Float64, S::Type = Float64) (Int[], Int[], T[], Int[], Int[], Int[], T[]), (Int[], Int[], T[], Int[], Int[], Int[], T[]), (Int[], Int[], T[], Int[], Int[], Int[], T[]), + (Int[], Int[], T[], Int[], Int[], Int[], T[]), + zeros(T,0,0), # ๐’โ‚ + zeros(T,0,0), # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ zeros(T,0,0), Sylvester_workspace(S = S), + zeros(T,0), # โˆ‚โˆ‡_vec # Second order pullback gradient buffers (lazily allocated) zeros(T,0,0), # โˆ‚โˆ‡โ‚‚ zeros(T,0,0), # โˆ‚โˆ‡โ‚ @@ -201,19 +231,125 @@ function Higher_order_workspace(;T::Type = Float64, S::Type = Float64) zeros(T,0,0), # โˆ‚โˆ‡โ‚_3rd zeros(T,0,0), # โˆ‚๐’โ‚_3rd zeros(T,0,0), # โˆ‚spinv_3rd + zeros(T,0,0), # โˆ‚โˆ‡โ‚‚_3rd + zeros(T,0,0), # โˆ‚โˆ‡โ‚ƒ_3rd + zeros(T,0,0), # โˆ‚๐’โ‚‚_3rd + zeros(T,0,0), # โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd + zeros(T,0,0), # โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd + zeros(T,0,0), # โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd + # Third order pullback temporary buffers + zeros(T,0,0), # โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd + zeros(T,0,0), # โˆ‚R_c_3rd + zeros(T,0,0), # โˆ‚L_c_3rd + zeros(T,0,0), # โˆ‚L_d_3rd + zeros(T,0,0), # โˆ‚R_d_3rd + zeros(T,0,0), # โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd + zeros(T,0,0), # โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd + zeros(T,0,0), # โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd + zeros(T,0,0), # โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tk0_3rd + zeros(T,0,0), # โˆ‚tmpkron0_ฯƒ_3rd + zeros(T,0,0), # โˆ‚aux_3rd + zeros(T,0,0), # โˆ‚S1S1_from_ck_3rd + zeros(T,0,0), # โˆ‚S1p0_kron_sigma_3rd + zeros(T,0,0), # โˆ‚S1p0_left_3rd + zeros(T,0,0), # โˆ‚S1p0_right_3rd + # Third order pullback intermediate product buffers (for mul!) + zeros(T,0,0), # โˆ‚A_3rd + zeros(T,0,0), # โˆ‚B_sylv_3rd + zeros(T,0,0), # โˆ‚๐—โ‚ƒ_3rd + zeros(T,0,0), # โˆ‚๐—โ‚ƒ_pre_3rd + zeros(T,0,0), # โˆ‚out2_3rd + zeros(T,0,0), # โˆ‚โˆ‡โ‚โ‚Š_3rd + zeros(T,0,0), # โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd + zeros(T,0,0), # โˆ‡โ‚‚t_โˆ‚out2_3rd + zeros(T,0,0), # mul_tmp_3rd # ForwardDiff partials buffers for stochastic steady state (accessed via model struct) zeros(S,0,0), # โˆ‚x_second_order zeros(S,0,0)) # โˆ‚x_third_order end +function ensure_higher_order_solution_buffers!(ws::higher_order_workspace{S,G,H}, n::Int, nโ‚‘โ‚‹::Int) where {S <: Real, G <: AbstractFloat, H <: Real} + size(ws.๐’โ‚) == (n, nโ‚‘โ‚‹) || (ws.๐’โ‚ = zeros(S, n, nโ‚‘โ‚‹)) + size(ws.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) == (nโ‚‘โ‚‹, nโ‚‘โ‚‹) || (ws.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = zeros(S, nโ‚‘โ‚‹, nโ‚‘โ‚‹)) + return ws +end + +""" + First_order_workspace(; T::Type = Float64, S::Type = Float64) + +Create a pre-allocated workspace for first-order perturbation and related AD paths. +""" +function First_order_workspace(; T::Type{TT} = Float64, S::Type{SS} = Float64) where {TT <: AbstractFloat, SS <: Real} + empty_qr_factors = zeros(TT, 0, 0) + empty_qr_ws::FastLapackInterface.QRWs = FastLapackInterface.QRWs(empty_qr_factors) + empty_qr_rhs = zeros(TT, 0, 0) + empty_qr_orm_ws::FastLapackInterface.QROrmWs = FastLapackInterface.QROrmWs(empty_qr_ws, 'L', 'T', empty_qr_factors, empty_qr_rhs) + empty_lu_factors = zeros(TT, 0, 0) + empty_lu_ws = FastLapackInterface.LUWs(empty_lu_factors) + empty_sparse = spzeros(TT, 0, 0) + empty_sparse_rhs = zeros(TT, 0) + empty_sparse_prob = ๐’ฎ.LinearProblem(empty_sparse, empty_sparse_rhs) + empty_sparse_lu::๐’ฎ.LinearCache = ๐’ฎ.init(empty_sparse_prob, + ๐’ฎ.LUFactorization(), + verbose = isdefined(๐’ฎ, :LinearVerbosity) ? ๐’ฎ.LinearVerbosity(๐’ฎ.SciMLLogging.Minimal()) : false) + + first_order_workspace( + Sylvester_workspace(S = TT, T = SS), # sylvester + # ForwardDiff partials buffers + zeros(SS, 0, 0), # Xฬƒ_first_order + zeros(SS, 0, 0), # p_tmp + zeros(SS, 0, 0), # โˆ‚SS_and_pars + zeros(TT, 0), # โˆ‚โˆ‡โ‚_vec + # First-order perturbation workspaces (primal) + zeros(TT, 0, 0), # ๐งโ‚šโ‚‹ + zeros(TT, 0, 0), # ๐Œ + zeros(TT, 0, 0), # ๐€โ‚Š + zeros(TT, 0, 0), # ๐€โ‚€ + zeros(TT, 0, 0), # ๐€โ‚‹ + zeros(TT, 0, 0), # ๐€ฬƒโ‚Š + zeros(TT, 0, 0), # ๐€ฬƒโ‚€ + zeros(TT, 0, 0), # ๐€ฬƒโ‚‹ + zeros(TT, 0, 0), # ๐€ฬ„โ‚€แตค + zeros(TT, 0, 0), # ๐€โ‚Šแตค + zeros(TT, 0, 0), # ๐€ฬƒโ‚€แตค + zeros(TT, 0, 0), # ๐€โ‚‹แตค + zeros(TT, 0, 0), # ๐€ + zeros(TT, 0, 0), # โˆ‡โ‚€ + zeros(TT, 0, 0), # โˆ‡โ‚‘ + # FastLapackInterface QR workspaces + empty_qr_factors, + empty_qr_ws, + empty_qr_orm_ws, + (0, 0, 0), + empty_qr_orm_ws, + (0, 0, 0), + empty_qr_orm_ws, + (0, 0, 0), + # FastLapackInterface LU workspaces + empty_lu_ws, + (0, 0), + empty_lu_ws, + (0, 0), + # Dedicated FastLapackInterface LU workspace for NSSS implicit derivatives + empty_lu_ws, + (0, 0), + empty_sparse_lu, + zeros(TT, 0), + zeros(TT, 0, 0)) +end + """ - Qme_workspace(n::Int; T::Type = Float64) + Qme_doubling_workspace(n::Int; T::Type = Float64, S::Type = Float64) Create a pre-allocated workspace for the quadratic matrix equation doubling algorithm. `n` is the dimension of the square matrices (nVars - nPresent_only). """ -function Qme_workspace(n::Int; T::Type = Float64, S::Type = Float64, nPast::Int = 0) - qme_workspace( zeros(T, n, n), # E +function Qme_doubling_workspace(n::Int; T::Type = Float64, S::Type = Float64) + empty_lu_factors = zeros(T, 0, 0) + empty_lu_ws = FastLapackInterface.LUWs(empty_lu_factors) + + qme_doubling_workspace( + zeros(T, n, n), # E zeros(T, n, n), # F zeros(T, n, n), # X zeros(T, n, n), # Y @@ -226,15 +362,67 @@ function Qme_workspace(n::Int; T::Type = Float64, S::Type = Float64, nPast::Int zeros(T, n, n), # temp3 zeros(T, n, n), # Bฬ„ zeros(T, n, n), # AXX - Sylvester_workspace(S = T), # sylvester_ws + Sylvester_workspace(S = T, T = S), # sylvester # ForwardDiff partials buffers zeros(S, 0, 0), # Xฬƒ - zeros(S, 0, 0), # Xฬƒ_first_order - zeros(S, 0, 0), # p_tmp - zeros(S, 0, 0), # โˆ‚SS_and_pars - # Pre-computed identity matrices (Diagonal{Bool} - supports indexing) - โ„’.I(n), # I_n - โ„’.I(nPast)) # I_nPast + # FastLapackInterface LU workspaces + empty_lu_ws, + (0, 0), + empty_lu_ws, + (0, 0)) +end + +function ensure_first_order_fast_qr_workspace!(ws::first_order_workspace{T}, qr_mat::AbstractMatrix) where {T <: Union{Float32, Float64}} + if size(ws.fast_qr_factors) != size(qr_mat) + ws.fast_qr_factors = zeros(T, size(qr_mat, 1), size(qr_mat, 2)) + ws.fast_qr_ws = FastLapackInterface.QRWs(ws.fast_qr_factors) + end + copyto!(ws.fast_qr_factors, qr_mat) + + return ws.fast_qr_factors, ws.fast_qr_ws +end + +""" + Schur_workspace(n::Int, nMixed::Int, nPfm::Int, nFnpm::Int; T::Type = Float64) + +Create a pre-allocated workspace for the schur-based quadratic matrix equation solver. +Dimensions: +- `n` = nVars - nPresent_only (dynamic variables) +- `nMixed` = number of mixed timing variables +- `nPfm` = nPast_not_future_and_mixed +- `nFnpm` = nFuture_not_past_and_mixed +""" +function Schur_workspace(n::Int, nMixed::Int, nPfm::Int, nFnpm::Int; T::Type = Float64) + companion_size = n + nMixed + nComb = nPfm + nFnpm # comb = union(future_not_past_and_mixed, past_not_future) + qz_seed_size = max(companion_size, 1) + qz_seed = zeros(T, qz_seed_size, qz_seed_size) + qz_ws = FastLapackInterface.GeneralizedSchurWs(qz_seed) + lu_seed_size = max(nPfm, 1) + lu_seed = zeros(T, lu_seed_size, lu_seed_size) + empty_lu_ws = FastLapackInterface.LUWs(lu_seed) + schur_workspace( + zeros(T, companion_size, companion_size), # D + zeros(T, companion_size, companion_size), # E + zeros(T, n, nPfm), # รƒโ‚‹ + zeros(T, n, nFnpm), # รƒโ‚€โ‚Š + zeros(T, n, nPfm), # รƒโ‚€โ‚‹ + zeros(T, nPfm, nPfm), # Zโ‚โ‚ + zeros(T, nFnpm, nPfm), # Zโ‚‚โ‚ + zeros(T, nPfm, nPfm), # Sโ‚โ‚ + zeros(T, nPfm, nPfm), # Tโ‚โ‚ + zeros(T, n, nPfm), # sol + zeros(T, n, n), # temp_X2 + zeros(T, n, n), # AXX + Vector{Bool}(undef, companion_size), # eigenselect + qz_ws, + (0, 0), + empty_lu_ws, + (0, 0), + empty_lu_ws, + (0, 0), + zeros(T, nPfm, nFnpm), # fast_lu_rhs_t_z21 + zeros(T, nPfm, nPfm)) # fast_lu_rhs_t_s11 end """ @@ -255,8 +443,12 @@ function Lyapunov_workspace(n::Int; T::Type = Float64) zeros(T, 0, 0), # tmpฬ„ (Krylov) zeros(T, 0, 0), # ๐— (Krylov) zeros(T, 0), # b (Krylov) - Krylov.BicgstabWorkspace(0, 0, Vector{T}), # bicgstab_workspace - Krylov.GmresWorkspace(0, 0, Vector{T}; memory = 20), # gmres_workspace + Krylov.BicgstabWorkspace(0, 0, Vector{T}), # bicgstab + Krylov.GmresWorkspace(0, 0, Vector{T}; memory = 20), # gmres + zeros(T, 0), # b_vech (vech-space Krylov) + Krylov.BicgstabWorkspace(0, 0, Vector{T}), # bicgstab_vech + Krylov.GmresWorkspace(0, 0, Vector{T}; memory = 20), # gmres_vech + zeros(T, 0, 0), # P (stable primal cache) # ForwardDiff partials buffers zeros(T, 0, 0), # Pฬƒ zeros(T, 0, 0), # Aฬƒ_fd @@ -309,30 +501,61 @@ function ensure_lyapunov_krylov_buffers!(ws::lyapunov_workspace{T}) where T end """ - ensure_lyapunov_bicgstab_solver!(ws::lyapunov_workspace{T}) where T + ensure_lyapunov_krylov_solver!(ws::lyapunov_workspace{T}, algorithm::Symbol) where T -Ensure the bicgstab solver workspace is allocated. +Ensure Krylov method buffers and the requested solver workspace are allocated. +Supported algorithms are `:bicgstab` and `:gmres`. """ -function ensure_lyapunov_bicgstab_solver!(ws::lyapunov_workspace{T}) where T +function ensure_lyapunov_krylov_solver!(ws::lyapunov_workspace{T}, algorithm::Symbol) where T ensure_lyapunov_krylov_buffers!(ws) n = ws.n - if length(ws.bicgstab_workspace.x) != n * n && n > 0 - ws.bicgstab_workspace = Krylov.BicgstabWorkspace(n * n, n * n, Vector{T}) + if n == 0 + return ws end + + if algorithm == :bicgstab + if length(ws.bicgstab.x) != n * n + ws.bicgstab = Krylov.BicgstabWorkspace(n * n, n * n, Vector{T}) + end + elseif algorithm == :gmres + if length(ws.gmres.x) != n * n + ws.gmres = Krylov.GmresWorkspace(n * n, n * n, Vector{T}; memory = 20) + end + else + error("Invalid Krylov algorithm: $algorithm. Must be :bicgstab or :gmres") + end + return ws end """ - ensure_lyapunov_gmres_solver!(ws::lyapunov_workspace{T}) where T + ensure_lyapunov_krylov_vech_solver!(ws::lyapunov_workspace{T}, algorithm::Symbol) where T -Ensure the gmres solver workspace is allocated. +Ensure vech-space Krylov buffers and solver workspace are allocated for symmetric Lyapunov equations. +The vech dimension is n(n+1)/2 instead of nยฒ. """ -function ensure_lyapunov_gmres_solver!(ws::lyapunov_workspace{T}) where T +function ensure_lyapunov_krylov_vech_solver!(ws::lyapunov_workspace{T}, algorithm::Symbol) where T ensure_lyapunov_krylov_buffers!(ws) n = ws.n - if length(ws.gmres_workspace.x) != n * n && n > 0 - ws.gmres_workspace = Krylov.GmresWorkspace(n * n, n * n, Vector{T}; memory = 20) + if n == 0 + return ws + end + n_vech = n * (n + 1) รท 2 + + if length(ws.b_vech) != n_vech + ws.b_vech = zeros(T, n_vech) + end + + if algorithm == :bicgstab + if length(ws.bicgstab_vech.x) != n_vech + ws.bicgstab_vech = Krylov.BicgstabWorkspace(n_vech, n_vech, Vector{T}) + end + elseif algorithm == :gmres + if length(ws.gmres_vech.x) != n_vech + ws.gmres_vech = Krylov.GmresWorkspace(n_vech, n_vech, Vector{T}; memory = 20) + end end + return ws end @@ -464,6 +687,20 @@ function Inversion_workspace(;T::Type = Float64) zeros(T, 0), # state_vol (n_past+1) zeros(T, 0), # aug_stateโ‚ (n_past+1+n_exo) zeros(T, 0), # aug_stateโ‚‚ (n_past+1+n_exo) + # Estimation loop temporaries + 0, # n_cond_var + zeros(T, 0), # shock_independent (n_cond_var) + zeros(T, 0), # init_guess (n_exo) + zeros(T, 0, 0), # Si_buffer (n_cond_var ร— n_exo) + zeros(T, 0, 0), # jacc_buffer (n_cond_var ร— n_exo) + zeros(T, 0, 0), # Si2e_buffer (n_cond_var ร— n_exo^2) + zeros(T, 0), # y_obs (n_cond_var) + zeros(T, 0), # x_shocks (n_exo) + zeros(T, 0), # state_concat (n_past + n_exo) + zeros(T, 0), # aug_stateโ‚ƒ (n_past+1+n_exo) + zeros(T, 0), # aug_stateโ‚ฬ‚ (n_past+1+n_exo) + zeros(T, 0), # stateยฒโป_vol (n_past+1) + zeros(T, 0), # kronstate_volยณ ((n_past+1)^3) # Pullback buffers (for reverse-mode AD) zeros(T, 0, 0), # โˆ‚_tmp1 (n_exo ร— n_past+n_exo) zeros(T, 0, 0), # โˆ‚_tmp2 (n_past ร— n_past+n_exo) @@ -542,6 +779,70 @@ function ensure_inversion_buffers!(ws::inversion_workspace{T}, n_exo::Int, n_pas ws.aug_stateโ‚‚ = zeros(T, n_aug) end + # Estimation loop temporaries (init_guess depends only on n_exo) + if length(ws.init_guess) != n_exo + ws.init_guess = zeros(T, n_exo) + end + if length(ws.x_shocks) != n_exo + ws.x_shocks = zeros(T, n_exo) + end + if length(ws.state_concat) != n_past + n_exo + ws.state_concat = zeros(T, n_past + n_exo) + end + + # Augmented state buffers for pruned third-order + if third_order + if length(ws.aug_stateโ‚ƒ) != n_aug + ws.aug_stateโ‚ƒ = zeros(T, n_aug) + end + if length(ws.aug_stateโ‚ฬ‚) != n_aug + ws.aug_stateโ‚ฬ‚ = zeros(T, n_aug) + end + if length(ws.stateยฒโป_vol) != n_state_vol + ws.stateยฒโป_vol = zeros(T, n_state_vol) + end + if length(ws.kronstate_volยณ) != n_state_vol^3 + ws.kronstate_volยณ = zeros(T, n_state_vol^3) + end + end + + return ws +end + + +""" + ensure_inversion_estimation_buffers!(ws::inversion_workspace{T}, n_exo::Int, n_cond_var::Int) where T + +Ensure observation-dimension-dependent estimation buffers are allocated. +Call after ensure_inversion_buffers! when the number of conditioning variables (observables) is known. +""" +function ensure_inversion_estimation_buffers!(ws::inversion_workspace{T}, n_exo::Int, n_cond_var::Int; third_order::Bool = false) where T + if ws.n_cond_var == n_cond_var && length(ws.shock_independent) == n_cond_var && + size(ws.Si_buffer) == (n_cond_var, n_exo) + return ws + end + + ws.n_cond_var = n_cond_var + + if length(ws.shock_independent) != n_cond_var + ws.shock_independent = zeros(T, n_cond_var) + end + if length(ws.y_obs) != n_cond_var + ws.y_obs = zeros(T, n_cond_var) + end + if size(ws.Si_buffer) != (n_cond_var, n_exo) + ws.Si_buffer = zeros(T, n_cond_var, n_exo) + end + if size(ws.jacc_buffer) != (n_cond_var, n_exo) + ws.jacc_buffer = zeros(T, n_cond_var, n_exo) + end + if third_order + n_exoยฒ = n_exo^2 + if size(ws.Si2e_buffer) != (n_cond_var, n_exoยฒ) + ws.Si2e_buffer = zeros(T, n_cond_var, n_exoยฒ) + end + end + return ws end @@ -550,9 +851,12 @@ end Kalman_workspace(;T::Type = Float64) Create a workspace for Kalman filter computations with lazy buffer allocation. -All buffers are initialized to 0-dimensional objects and resized on-demand via ensure_kalman_buffers!. +All buffers are initialized to 0-dimensional objects and resized on-demand via ensure_kalman_workspaces!. """ function Kalman_workspace(;T::Type = Float64) + empty_lu_factors = zeros(T, 1, 1) + empty_lu_ws = FastLapackInterface.LUWs(empty_lu_factors) + kalman_workspace{T}( 0, 0, # n_obs, n_states dimensions zeros(T, 0), # u (n_states) @@ -560,19 +864,26 @@ function Kalman_workspace(;T::Type = Float64) zeros(T, 0), # ztmp (n_obs) zeros(T, 0), # utmp (n_states) zeros(T, 0, 0), # Ctmp (n_obs ร— n_states) + zeros(T, 0, 0), # ๐ (n_states ร— n_states) zeros(T, 0, 0), # F (n_obs ร— n_obs) zeros(T, 0, 0), # K (n_states ร— n_obs) zeros(T, 0, 0), # tmp (n_states ร— n_states) - zeros(T, 0, 0)) # Ptmp (n_states ร— n_states) + zeros(T, 0, 0), # Ptmp (n_states ร— n_states) + empty_lu_ws, + (0, 0), + zeros(T, 0, 0)) # fast_lu_rhs_t_k (n_obs ร— n_states) end """ - ensure_kalman_buffers!(ws::kalman_workspace{T}, n_obs::Int, n_states::Int) where T + ensure_kalman_workspaces!(workspaces::workspaces, n_obs::Int, n_states::Int) -Ensure the Kalman workspaces are allocated for the given dimensions. +Ensure the Kalman workspace inside `workspaces` is allocated for the given dimensions and return it. """ -function ensure_kalman_buffers!(ws::kalman_workspace{T}, n_obs::Int, n_states::Int) where T +function ensure_kalman_workspaces!(workspaces::workspaces, n_obs::Int, n_states::Int) + ws = workspaces.kalman + T = eltype(ws.u) + # Check if dimensions changed if ws.n_obs == n_obs && ws.n_states == n_states return ws @@ -599,6 +910,9 @@ function ensure_kalman_buffers!(ws::kalman_workspace{T}, n_obs::Int, n_states::I if size(ws.Ctmp, 1) != n_obs || size(ws.Ctmp, 2) != n_states ws.Ctmp = zeros(T, n_obs, n_states) end + if size(ws.๐, 1) != n_states || size(ws.๐, 2) != n_states + ws.๐ = zeros(T, n_states, n_states) + end if size(ws.F, 1) != n_obs || size(ws.F, 2) != n_obs ws.F = zeros(T, n_obs, n_obs) end @@ -611,23 +925,31 @@ function ensure_kalman_buffers!(ws::kalman_workspace{T}, n_obs::Int, n_states::I if size(ws.Ptmp, 1) != n_states || size(ws.Ptmp, 2) != n_states ws.Ptmp = zeros(T, n_states, n_states) end + if size(ws.fast_lu_rhs_t_k, 1) != n_obs || size(ws.fast_lu_rhs_t_k, 2) != n_states + ws.fast_lu_rhs_t_k = zeros(T, n_obs, n_states) + end return ws end -function Workspaces(;T::Type = Float64, S::Type = Float64) +function Workspaces(;T::Type{Float64} = Float64, S::Type{Float64} = Float64) workspaces(Higher_order_workspace(T = T, S = S), Higher_order_workspace(T = T, S = S), Float64[], - Qme_workspace(0, T = T), # Initialize with size 0, will be resized when needed + First_order_workspace(T = T, S = S), # Initialize with size 0, will be resized when needed + Qme_doubling_workspace(0, T = T, S = S), # Initialize with size 0, will be resized when needed + Schur_workspace(0, 0, 0, 0, T = T), # Initialize with size 0, will be resized when needed Lyapunov_workspace(0, T = T), # 1st order - will be resized Lyapunov_workspace(0, T = T), # 2nd order - will be resized Lyapunov_workspace(0, T = T), # 3rd order - will be resized + Lyapunov_workspace(0, T = T), # block-triangular inner - will be resized Sylvester_workspace(S = S), # 1st order sylvester - will be resized + Sylvester_workspace(S = S), # block-triangular sylvester - will be resized Find_shocks_workspace(T = T), # conditional forecast - will be resized Inversion_workspace(T = T), # inversion filter - will be resized - Kalman_workspace(T = T)) # Kalman filter - will be resized + Kalman_workspace(T = T), # Kalman filter - will be resized + NSSSSolverWorkspace()) # NSSS solver scratch buffers end function Constants(model_struct; T::Type = Float64, S::Type = Float64) @@ -635,7 +957,9 @@ function Constants(model_struct; T::Type = Float64, S::Type = Float64) post_parameters_macro( Symbol[], false, - true, + :single_equation, + :ESCH, + 120.0, Dict{Symbol, Float64}(), Set{Symbol}[], Set{Symbol}[], @@ -668,6 +992,8 @@ function Constants(model_struct; T::Type = Float64, S::Type = Float64) spzeros(Float64, 0, 0), Symbol[], Symbol[], + Int[], + Int[], Symbol[], # Symbol[], Int[], @@ -680,13 +1006,31 @@ function Constants(model_struct; T::Type = Float64, S::Type = Float64) Int[], Int[], โ„’.I(0), + โ„’.I(0), 1:0, 1:0, 1, zeros(Bool, 0, 0), - zeros(Bool, 0, 0)), + zeros(Bool, 0, 0), + Int[], + Int[], # indices_past_not_future_in_comb + zeros(Bool, 0, 0), # I_nPast_not_mixed + zeros(Bool, 0, 0), # Ir_past_selector + zeros(Bool, 0, 0), # schur_Zโ‚Š + zeros(Bool, 0, 0), # schur_Iโ‚Š + zeros(Bool, 0, 0), # schur_Zโ‚‹ + zeros(Bool, 0, 0), # schur_Iโ‚‹ + nothing, + 0, + Int[], + 0, + Symbol[], + Int[], + Symbol[], + 1), Second_order_indices(), - Third_order_indices()) + Third_order_indices(), + NSSSSolverConstants()) end function _axis_has_string(axis) @@ -752,6 +1096,8 @@ function update_post_complete_parameters(p::post_complete_parameters; kwargs...) get(kwargs, :custom_ss_expand_matrix, p.custom_ss_expand_matrix), get(kwargs, :vars_in_ss_equations, p.vars_in_ss_equations), get(kwargs, :vars_in_ss_equations_with_aux, p.vars_in_ss_equations_with_aux), + get(kwargs, :ss_var_idx_in_var_and_calib, p.ss_var_idx_in_var_and_calib), + get(kwargs, :calib_idx_in_var_and_calib, p.calib_idx_in_var_and_calib), get(kwargs, :SS_and_pars_names_lead_lag, p.SS_and_pars_names_lead_lag), # get(kwargs, :SS_and_pars_names_no_exo, p.SS_and_pars_names_no_exo), get(kwargs, :SS_and_pars_no_exo_idx, p.SS_and_pars_no_exo_idx), @@ -764,20 +1110,38 @@ function update_post_complete_parameters(p::post_complete_parameters; kwargs...) get(kwargs, :future_not_past_and_mixed_in_comb, p.future_not_past_and_mixed_in_comb), get(kwargs, :past_not_future_and_mixed_in_comb, p.past_not_future_and_mixed_in_comb), get(kwargs, :Ir, p.Ir), + get(kwargs, :I_n, hasfield(typeof(p), :I_n) ? p.I_n : โ„’.I(0)), get(kwargs, :nabla_zero_cols, p.nabla_zero_cols), get(kwargs, :nabla_minus_cols, p.nabla_minus_cols), get(kwargs, :nabla_e_start, p.nabla_e_start), get(kwargs, :expand_future, p.expand_future), get(kwargs, :expand_past, p.expand_past), + get(kwargs, :past_not_future_and_mixed_in_present_but_not_only, + hasfield(typeof(p), :past_not_future_and_mixed_in_present_but_not_only) ? p.past_not_future_and_mixed_in_present_but_not_only : Int[]), + get(kwargs, :indices_past_not_future_in_comb, hasfield(typeof(p), :indices_past_not_future_in_comb) ? p.indices_past_not_future_in_comb : Int[]), + get(kwargs, :I_nPast_not_mixed, hasfield(typeof(p), :I_nPast_not_mixed) ? p.I_nPast_not_mixed : Matrix{Bool}(undef, 0, 0)), + get(kwargs, :Ir_past_selector, hasfield(typeof(p), :Ir_past_selector) ? p.Ir_past_selector : Matrix{Bool}(undef, 0, 0)), + get(kwargs, :schur_Zโ‚Š, hasfield(typeof(p), :schur_Zโ‚Š) ? p.schur_Zโ‚Š : Matrix{Bool}(undef, 0, 0)), + get(kwargs, :schur_Iโ‚Š, hasfield(typeof(p), :schur_Iโ‚Š) ? p.schur_Iโ‚Š : Matrix{Bool}(undef, 0, 0)), + get(kwargs, :schur_Zโ‚‹, hasfield(typeof(p), :schur_Zโ‚‹) ? p.schur_Zโ‚‹ : Matrix{Bool}(undef, 0, 0)), + get(kwargs, :schur_Iโ‚‹, hasfield(typeof(p), :schur_Iโ‚‹) ? p.schur_Iโ‚‹ : Matrix{Bool}(undef, 0, 0)), + get(kwargs, :nsss_dependencies, p.nsss_dependencies), + get(kwargs, :nsss_n_sol, p.nsss_n_sol), + get(kwargs, :nsss_output_indices, p.nsss_output_indices), + get(kwargs, :nsss_n_ext_params, p.nsss_n_ext_params), + get(kwargs, :nsss_sol_names, p.nsss_sol_names), + get(kwargs, :nsss_exo_zero_indices, p.nsss_exo_zero_indices), + get(kwargs, :nsss_param_names_ext, p.nsss_param_names_ext), + get(kwargs, :nsss_fastest_solver_parameter_idx, p.nsss_fastest_solver_parameter_idx), ) end # Initialize all commonly used constants at once (call at entry points) # This reduces repeated ensure_*! calls throughout the codebase function initialise_constants!(๐“‚) - ensure_computational_constants!(๐“‚) + ensure_computational_constants!(๐“‚.constants) ensure_name_display_constants!(๐“‚) - ensure_first_order_constants!(๐“‚) + ensure_first_order_constants!(๐“‚.constants) return ๐“‚.constants end @@ -862,61 +1226,6 @@ function set_up_name_display_cache(T::post_model_macro, calibration_equations_pa end -function ensure_computational_constants!(๐“‚) - constants = ๐“‚.constants - so = constants.second_order - if isempty(so.s_in_sโบ) - # Use timings from constants if available, otherwise from model - T = constants.post_model_macro - nแต‰ = T.nExo - nหข = T.nPast_not_future_and_mixed - - s_in_sโบ = BitVector(vcat(ones(Bool, nหข + 1), zeros(Bool, nแต‰))) - s_in_s = BitVector(vcat(ones(Bool, nหข), zeros(Bool, nแต‰ + 1))) - - kron_sโบ_sโบ = โ„’.kron(s_in_sโบ, s_in_sโบ) - kron_sโบ_s = โ„’.kron(s_in_sโบ, s_in_s) - - kron_sโบ_sโบ_sโบ = โ„’.kron(s_in_sโบ, kron_sโบ_sโบ) - kron_s_sโบ_sโบ = โ„’.kron(kron_sโบ_sโบ, s_in_s) - - e_in_sโบ = BitVector(vcat(zeros(Bool, nหข + 1), ones(Bool, nแต‰))) - v_in_sโบ = BitVector(vcat(zeros(Bool, nหข), 1, zeros(Bool, nแต‰))) - - kron_s_s = โ„’.kron(s_in_sโบ, s_in_sโบ) - kron_e_e = โ„’.kron(e_in_sโบ, e_in_sโบ) - kron_v_v = โ„’.kron(v_in_sโบ, v_in_sโบ) - kron_e_s = โ„’.kron(e_in_sโบ, s_in_sโบ) - - # Compute sparse index patterns for filter operations - shockvar_idxs = sparse(โ„’.kron(e_in_sโบ, s_in_sโบ)).nzind - shock_idxs = sparse(โ„’.kron(e_in_sโบ, zero(e_in_sโบ) .+ 1)).nzind - shock_idxs2 = sparse(โ„’.kron(zero(e_in_sโบ) .+ 1, e_in_sโบ)).nzind - shockยฒ_idxs = sparse(โ„’.kron(e_in_sโบ, e_in_sโบ)).nzind - var_volยฒ_idxs = sparse(โ„’.kron(s_in_sโบ, s_in_sโบ)).nzind - - so.s_in_sโบ = s_in_sโบ - so.s_in_s = s_in_s - so.kron_sโบ_sโบ = kron_sโบ_sโบ - so.kron_sโบ_s = kron_sโบ_s - so.kron_sโบ_sโบ_sโบ = kron_sโบ_sโบ_sโบ - so.kron_s_sโบ_sโบ = kron_s_sโบ_sโบ - so.e_in_sโบ = e_in_sโบ - so.v_in_sโบ = v_in_sโบ - so.kron_s_s = kron_s_s - so.kron_e_e = kron_e_e - so.kron_v_v = kron_v_v - so.kron_e_s = kron_e_s - so.shockvar_idxs = shockvar_idxs - so.shock_idxs = shock_idxs - so.shock_idxs2 = shock_idxs2 - so.shockยฒ_idxs = shockยฒ_idxs - so.var_volยฒ_idxs = var_volยฒ_idxs - end - - return constants.second_order -end - function ensure_computational_constants!(constants::constants) so = constants.second_order if isempty(so.s_in_sโบ) @@ -971,56 +1280,6 @@ function ensure_computational_constants!(constants::constants) return constants.second_order end -function ensure_conditional_forecast_constants!(๐“‚; third_order::Bool = false) - constants = ๐“‚.constants - so = ensure_computational_constants!(๐“‚) - - if isempty(so.varยฒ_idxs) - s_in_sโบ = so.s_in_s - e_in_sโบ = so.e_in_sโบ - - shock_idxs = so.shock_idxs - shockยฒ_idxs = so.shockยฒ_idxs - shockvarยฒ_idxs = setdiff(shock_idxs, shockยฒ_idxs) - var_volยฒ_idxs = so.var_volยฒ_idxs - varยฒ_idxs = sparse(โ„’.kron(s_in_sโบ, s_in_sโบ)).nzind - so.varยฒ_idxs = varยฒ_idxs - so.shockvarยฒ_idxs = shockvarยฒ_idxs - so.var_volยฒ_idxs = var_volยฒ_idxs - end - - if third_order - to = constants.third_order - if isempty(to.var_volยณ_idxs) - sv_in_sโบ = so.s_in_sโบ - e_in_sโบ = so.e_in_sโบ - ones_e = zero(e_in_sโบ) .+ 1 - - var_volยณ_idxs = sparse(โ„’.kron(sv_in_sโบ, โ„’.kron(sv_in_sโบ, sv_in_sโบ))).nzind - shock_idxs2 = sparse(โ„’.kron(โ„’.kron(e_in_sโบ, ones_e), ones_e)).nzind - shock_idxs3 = sparse(โ„’.kron(โ„’.kron(e_in_sโบ, e_in_sโบ), ones_e)).nzind - shockยณ_idxs = sparse(โ„’.kron(e_in_sโบ, โ„’.kron(e_in_sโบ, e_in_sโบ))).nzind - shockvar1_idxs = sparse(โ„’.kron(ones_e, โ„’.kron(e_in_sโบ, e_in_sโบ))).nzind - shockvar2_idxs = sparse(โ„’.kron(e_in_sโบ, โ„’.kron(ones_e, e_in_sโบ))).nzind - shockvar3_idxs = sparse(โ„’.kron(e_in_sโบ, โ„’.kron(e_in_sโบ, ones_e))).nzind - shockvarยณ2_idxs = setdiff(shock_idxs2, shockยณ_idxs, shockvar1_idxs, shockvar2_idxs, shockvar3_idxs) - shockvarยณ_idxs = setdiff(shock_idxs3, shockยณ_idxs) - - to.var_volยณ_idxs = var_volยณ_idxs - to.shock_idxs2 = shock_idxs2 - to.shock_idxs3 = shock_idxs3 - to.shockยณ_idxs = shockยณ_idxs - to.shockvar1_idxs = shockvar1_idxs - to.shockvar2_idxs = shockvar2_idxs - to.shockvar3_idxs = shockvar3_idxs - to.shockvarยณ2_idxs = shockvarยณ2_idxs - to.shockvarยณ_idxs = shockvarยณ_idxs - end - end - - return so -end - function ensure_conditional_forecast_constants!(constants::constants; third_order::Bool = false) so = ensure_computational_constants!(constants) @@ -1109,6 +1368,33 @@ function build_first_order_index_cache(T, I_nVars) expand_future = I_nVars[T.future_not_past_and_mixed_idx,:] expand_past = I_nVars[T.past_not_future_and_mixed_idx,:] + past_not_future_and_mixed_in_present_but_not_only_tmp = indexin(T.past_not_future_and_mixed_idx, T.present_but_not_only_idx) + if any(isnothing.(past_not_future_and_mixed_in_present_but_not_only_tmp)) + past_not_future_and_mixed_in_present_but_not_only = Int[] + else + past_not_future_and_mixed_in_present_but_not_only = Int.(past_not_future_and_mixed_in_present_but_not_only_tmp) + end + + # Schur QME cached indices and constant matrices + indices_past_not_future_in_comb_tmp = indexin(T.past_not_future_idx, comb) + if any(isnothing.(indices_past_not_future_in_comb_tmp)) + indices_past_not_future_in_comb = Int[] + else + indices_past_not_future_in_comb = Int.(indices_past_not_future_in_comb_tmp) + end + + I_nPast = โ„’.I(T.nPast_not_future_and_mixed) + I_nPast_not_mixed = Matrix{Bool}(I_nPast[T.not_mixed_in_past_idx, :]) + Ir_past_selector = Matrix{Bool}(Ir[past_not_future_and_mixed_in_comb, :]) + I_n = โ„’.I(T.nVars - T.nPresent_only) + + schur_Zโ‚Š = zeros(Bool, T.nMixed, T.nFuture_not_past_and_mixed) + I_nFuture = โ„’.I(T.nFuture_not_past_and_mixed) + schur_Iโ‚Š = Matrix{Bool}(I_nFuture[T.mixed_in_future_idx, :]) + + schur_Zโ‚‹ = zeros(Bool, T.nMixed, T.nPast_not_future_and_mixed) + schur_Iโ‚‹ = Matrix{Bool}(I_nPast[T.mixed_in_past_idx, :]) + return ( initialized = true, dyn_index = dyn_index, @@ -1117,18 +1403,26 @@ function build_first_order_index_cache(T, I_nVars) future_not_past_and_mixed_in_comb = future_not_past_and_mixed_in_comb, past_not_future_and_mixed_in_comb = past_not_future_and_mixed_in_comb, Ir = Ir, + I_n = I_n, nabla_zero_cols = nabla_zero_cols, nabla_minus_cols = nabla_minus_cols, nabla_e_start = nabla_e_start, expand_future = expand_future, expand_past = expand_past, + past_not_future_and_mixed_in_present_but_not_only = past_not_future_and_mixed_in_present_but_not_only, + indices_past_not_future_in_comb = indices_past_not_future_in_comb, + I_nPast_not_mixed = I_nPast_not_mixed, + Ir_past_selector = Ir_past_selector, + schur_Zโ‚Š = schur_Zโ‚Š, + schur_Iโ‚Š = schur_Iโ‚Š, + schur_Zโ‚‹ = schur_Zโ‚‹, + schur_Iโ‚‹ = schur_Iโ‚‹, ) end -function ensure_first_order_constants!(๐“‚) - constants = ๐“‚.constants +function ensure_first_order_constants!(constants::constants) if !constants.post_complete_parameters.initialized - # Use timings from constants if available, otherwise from model + # Use timings from constants if available T = constants.post_model_macro diag_nVars = constants.post_complete_parameters.diag_nVars if size(diag_nVars, 1) == 0 @@ -1145,86 +1439,182 @@ function ensure_first_order_constants!(๐“‚) future_not_past_and_mixed_in_comb = cache.future_not_past_and_mixed_in_comb, past_not_future_and_mixed_in_comb = cache.past_not_future_and_mixed_in_comb, Ir = cache.Ir, + I_n = cache.I_n, nabla_zero_cols = cache.nabla_zero_cols, nabla_minus_cols = cache.nabla_minus_cols, nabla_e_start = cache.nabla_e_start, expand_future = cache.expand_future, expand_past = cache.expand_past, + past_not_future_and_mixed_in_present_but_not_only = cache.past_not_future_and_mixed_in_present_but_not_only, + indices_past_not_future_in_comb = cache.indices_past_not_future_in_comb, + I_nPast_not_mixed = cache.I_nPast_not_mixed, + Ir_past_selector = cache.Ir_past_selector, + schur_Zโ‚Š = cache.schur_Zโ‚Š, + schur_Iโ‚Š = cache.schur_Iโ‚Š, + schur_Zโ‚‹ = cache.schur_Zโ‚‹, + schur_Iโ‚‹ = cache.schur_Iโ‚‹, ) end return constants.post_complete_parameters end -function ensure_first_order_constants!(constants::constants) - if !constants.post_complete_parameters.initialized - # Use timings from constants if available - T = constants.post_model_macro - diag_nVars = constants.post_complete_parameters.diag_nVars - if size(diag_nVars, 1) == 0 - diag_nVars = โ„’.I(T.nVars) - end - cache = build_first_order_index_cache(T, diag_nVars) - constants.post_complete_parameters = update_post_complete_parameters( - constants.post_complete_parameters; - diag_nVars = diag_nVars, - initialized = cache.initialized, - dyn_index = cache.dyn_index, - reverse_dynamic_order = cache.reverse_dynamic_order, - comb = cache.comb, - future_not_past_and_mixed_in_comb = cache.future_not_past_and_mixed_in_comb, - past_not_future_and_mixed_in_comb = cache.past_not_future_and_mixed_in_comb, - Ir = cache.Ir, - nabla_zero_cols = cache.nabla_zero_cols, - nabla_minus_cols = cache.nabla_minus_cols, - nabla_e_start = cache.nabla_e_start, - expand_future = cache.expand_future, - expand_past = cache.expand_past, - ) + +""" + ensure_qme_doubling_workspace!(workspaces, n) + +Ensure the QME doubling workspace has dimension `n`. +If the workspace is the wrong size, it is reallocated. +""" +function ensure_qme_doubling_workspace!(workspaces::workspaces, n::Int) + ws = workspaces.qme_doubling + if size(ws.E, 1) != n + workspaces.qme_doubling = Qme_doubling_workspace(n) end - return constants.post_complete_parameters + return workspaces.qme_doubling end +""" + ensure_third_order_pullback_workspaces!(โ„‚, S, T, Mโ‚‚, Mโ‚ƒ) + +Ensure workspace buffers for the third-order pullback are allocated with correct dimensions. +Only dense intermediate-product temporaries are workspace-backed; gradient accumulators for +โˆ‡โ‚‚, โˆ‡โ‚ƒ, ๐’โ‚‚ and "may be sparse" matrices are freshly allocated via `zero()` inside the +pullback to preserve their sparse/dense format. +""" +function ensure_third_order_pullback_workspaces!(โ„‚::higher_order_workspace, ::Type{S}, T, Mโ‚‚, Mโ‚ƒ) where S + n = T.nVars + nโ‚Š = T.nFuture_not_past_and_mixed + nโ‚‹ = T.nPast_not_future_and_mixed + nโ‚‘ = T.nExo + nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + n_stack = nโ‚Š + n + nโ‚‹ + nโ‚‘ + + # Structural dimensions from constants + n_โˆ‡โ‚‚ = size(Mโ‚‚.๐”โˆ‡โ‚‚, 2) + n_๐‚โ‚ƒ_r = size(Mโ‚ƒ.๐‚โ‚ƒ, 1) + n_๐‚โ‚ƒ = size(Mโ‚ƒ.๐‚โ‚ƒ, 2) + ฯƒ_c = size(Mโ‚‚.๐›”, 2) + n_out2_c = ฯƒ_c * nโ‚‘โ‚‹ + + # Dense workspace: always-dense gradient accumulators (matches main branch) + size(โ„‚.โˆ‚spinv_3rd) == (n, n) || (โ„‚.โˆ‚spinv_3rd = zeros(S, n, n)) + size(โ„‚.โˆ‚โˆ‡โ‚_3rd) == (n, n_stack) || (โ„‚.โˆ‚โˆ‡โ‚_3rd = zeros(S, n, n_stack)) + size(โ„‚.โˆ‚๐’โ‚_3rd) == (n, nโ‚‘โ‚‹) || (โ„‚.โˆ‚๐’โ‚_3rd = zeros(S, n, nโ‚‘โ‚‹)) + + # Dense workspace: intermediate-product temporaries (overwritten by mul! each call) + size(โ„‚.โˆ‚A_3rd) == (n, n) || (โ„‚.โˆ‚A_3rd = zeros(S, n, n)) + size(โ„‚.โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd) == (n, n) || (โ„‚.โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd = zeros(S, n, n)) + size(โ„‚.mul_tmp_3rd) == (n, n) || (โ„‚.mul_tmp_3rd = zeros(S, n, n)) + size(โ„‚.โˆ‚B_sylv_3rd) == (n_๐‚โ‚ƒ, n_๐‚โ‚ƒ) || (โ„‚.โˆ‚B_sylv_3rd = zeros(S, n_๐‚โ‚ƒ, n_๐‚โ‚ƒ)) + size(โ„‚.โˆ‚๐—โ‚ƒ_3rd) == (n, n_๐‚โ‚ƒ) || (โ„‚.โˆ‚๐—โ‚ƒ_3rd = zeros(S, n, n_๐‚โ‚ƒ)) + size(โ„‚.โˆ‚๐—โ‚ƒ_pre_3rd) == (n, n_๐‚โ‚ƒ_r) || (โ„‚.โˆ‚๐—โ‚ƒ_pre_3rd = zeros(S, n, n_๐‚โ‚ƒ_r)) + size(โ„‚.โˆ‚out2_3rd) == (n, n_out2_c) || (โ„‚.โˆ‚out2_3rd = zeros(S, n, n_out2_c)) + size(โ„‚.โˆ‡โ‚‚t_โˆ‚out2_3rd) == (n_โˆ‡โ‚‚, n_out2_c) || (โ„‚.โˆ‡โ‚‚t_โˆ‚out2_3rd = zeros(S, n_โˆ‡โ‚‚, n_out2_c)) + + # Pullback gradient accumulator buffers (zeroed at start of each pullback call) + size(โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd) == (n_stack, nโ‚‘โ‚‹^2) || (โ„‚.โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd = zeros(S, n_stack, nโ‚‘โ‚‹^2)) + size(โ„‚.โˆ‚L_c_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚L_c_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚R_c_3rd) == (n_stack, nโ‚‘โ‚‹^2) || (โ„‚.โˆ‚R_c_3rd = zeros(S, n_stack, nโ‚‘โ‚‹^2)) + size(โ„‚.โˆ‚L_d_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚L_d_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚R_d_3rd) == (n_stack, nโ‚‘โ‚‹^2) || (โ„‚.โˆ‚R_d_3rd = zeros(S, n_stack, nโ‚‘โ‚‹^2)) + size(โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd) == (nโ‚‘โ‚‹, nโ‚‘โ‚‹) || (โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd = zeros(S, nโ‚‘โ‚‹, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd) == (nโ‚‘โ‚‹, nโ‚‘โ‚‹^2) || (โ„‚.โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd = zeros(S, nโ‚‘โ‚‹, nโ‚‘โ‚‹^2)) + size(โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd) == (nโ‚‘โ‚‹, nโ‚‘โ‚‹) || (โ„‚.โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd = zeros(S, nโ‚‘โ‚‹, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tk0_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tk0_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚aux_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚aux_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚tmpkron0_ฯƒ_3rd) == (nโ‚‘โ‚‹^2, nโ‚‘โ‚‹^2) || (โ„‚.โˆ‚tmpkron0_ฯƒ_3rd = zeros(S, nโ‚‘โ‚‹^2, nโ‚‘โ‚‹^2)) + size(โ„‚.โˆ‚โˆ‡โ‚โ‚Š_3rd) == (n, n) || (โ„‚.โˆ‚โˆ‡โ‚โ‚Š_3rd = zeros(S, n, n)) + size(โ„‚.โˆ‚S1S1_from_ck_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚S1S1_from_ck_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚S1p0_kron_sigma_3rd) == (n_stack^2, ฯƒ_c) || (โ„‚.โˆ‚S1p0_kron_sigma_3rd = zeros(S, n_stack^2, ฯƒ_c)) + size(โ„‚.โˆ‚S1p0_left_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚S1p0_left_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + size(โ„‚.โˆ‚S1p0_right_3rd) == (n_stack, nโ‚‘โ‚‹) || (โ„‚.โˆ‚S1p0_right_3rd = zeros(S, n_stack, nโ‚‘โ‚‹)) + + return โ„‚ +end """ - ensure_qme_workspace!(๐“‚) - ensure_qme_workspace!(workspaces, n) + ensure_first_order_workspace_buffers!(ws, T, n_dyn, n_comb) -Ensure the QME (quadratic matrix equation) workspace is properly sized for the model. -The workspace dimension is `n = nVars - nPresent_only` (the size of the QME matrices). -If the workspace is the wrong size, it will be reallocated. +Ensure all first-order perturbation buffers in `first_order_workspace` are allocated with +the correct dimensions. """ -function ensure_qme_workspace!(๐“‚) - T = ๐“‚.constants.post_model_macro - n = T.nVars - T.nPresent_only - nPast = T.nPast_not_future_and_mixed - return ensure_qme_workspace!(๐“‚.workspaces, n, nPast) +function ensure_first_order_workspace_buffers!(ws::first_order_workspace{R,S}, T, n_dyn::Int, n_comb::Int) where {R <: Real, S <: Real} + n = T.nVars + nโ‚Š = T.nFuture_not_past_and_mixed + nโ‚‹ = T.nPast_not_future_and_mixed + nโ‚‘ = T.nExo + nแตค = T.nPresent_only + nโ‚€แตค = length(T.present_but_not_only_idx) + + size(ws.๐€โ‚Š) == (n, nโ‚Š) || (ws.๐€โ‚Š = zeros(R, n, nโ‚Š)) + size(ws.๐€โ‚€) == (n, n) || (ws.๐€โ‚€ = zeros(R, n, n)) + size(ws.๐€โ‚‹) == (n, nโ‚‹) || (ws.๐€โ‚‹ = zeros(R, n, nโ‚‹)) + + size(ws.๐€ฬƒโ‚Š) == (n_dyn, n_comb) || (ws.๐€ฬƒโ‚Š = zeros(R, n_dyn, n_comb)) + size(ws.๐€ฬƒโ‚€) == (n_dyn, n_comb) || (ws.๐€ฬƒโ‚€ = zeros(R, n_dyn, n_comb)) + size(ws.๐€ฬƒโ‚‹) == (n_dyn, n_comb) || (ws.๐€ฬƒโ‚‹ = zeros(R, n_dyn, n_comb)) + + size(ws.๐€ฬ„โ‚€แตค) == (nแตค, nแตค) || (ws.๐€ฬ„โ‚€แตค = zeros(R, nแตค, nแตค)) + size(ws.๐€โ‚Šแตค) == (nแตค, nโ‚Š) || (ws.๐€โ‚Šแตค = zeros(R, nแตค, nโ‚Š)) + size(ws.๐€ฬƒโ‚€แตค) == (nแตค, nโ‚€แตค) || (ws.๐€ฬƒโ‚€แตค = zeros(R, nแตค, nโ‚€แตค)) + size(ws.๐€โ‚‹แตค) == (nแตค, nโ‚‹) || (ws.๐€โ‚‹แตค = zeros(R, nแตค, nโ‚‹)) + + size(ws.๐งโ‚šโ‚‹) == (nแตค, nโ‚‹) || (ws.๐งโ‚šโ‚‹ = zeros(R, nแตค, nโ‚‹)) + size(ws.๐Œ) == (nโ‚Š, n) || (ws.๐Œ = zeros(R, nโ‚Š, n)) + size(ws.๐€) == (n, nโ‚‹) || (ws.๐€ = zeros(R, n, nโ‚‹)) + size(ws.โˆ‡โ‚€) == (n, n) || (ws.โˆ‡โ‚€ = zeros(R, n, n)) + size(ws.โˆ‡โ‚‘) == (n, nโ‚‘) || (ws.โˆ‡โ‚‘ = zeros(R, n, nโ‚‘)) + + return ws end -function ensure_qme_workspace!(workspaces::workspaces, n::Int, nPast::Int = 0) - ws = workspaces.qme - # Check if workspace needs to be resized (either n or nPast changed) - if size(ws.E, 1) != n || size(ws.I_nPast, 1) != nPast - workspaces.qme = Qme_workspace(n, nPast = nPast) +function ensure_first_order_cotangent_buffer!(ws::first_order_workspace{T}, n::Int) where T <: Real + if length(ws.โˆ‚โˆ‡โ‚_vec) != n + ws.โˆ‚โˆ‡โ‚_vec = zeros(T, n) end - return workspaces.qme + + return ws.โˆ‚โˆ‡โ‚_vec +end + +function ensure_higher_order_cotangent_buffer!(ws::higher_order_workspace{T}, n::Int) where T <: Real + if length(ws.โˆ‚โˆ‡_vec) != n + ws.โˆ‚โˆ‡_vec = zeros(T, n) + end + + return ws.โˆ‚โˆ‡_vec end """ - ensure_sylvester_1st_order_workspace!(๐“‚) - ensure_sylvester_1st_order_workspace!(workspaces) + ensure_schur_workspace!(workspaces, n, nMixed, nPfm, nFnpm) -Return the first-order sylvester workspace from the model or workspaces. -The workspace is lazily sized by the sylvester solver when needed. +Ensure the schur workspace is properly sized for the model. +Dimensions are: +- `n = nVars - nPresent_only` (dynamic variables) +- `nMixed` (mixed timing variables) +- `nPfm = nPast_not_future_and_mixed` +- `nFnpm = nFuture_not_past_and_mixed` + +If the workspace is the wrong size, it will be reallocated. """ -function ensure_sylvester_1st_order_workspace!(๐“‚) - return ๐“‚.workspaces.sylvester_1st_order +function ensure_schur_workspace!(workspaces::workspaces, n::Int, nMixed::Int, nPfm::Int, nFnpm::Int) + workspaces.schur = ensure_schur_workspace!(workspaces.schur, n, nMixed, nPfm, nFnpm) + return workspaces.schur end -function ensure_sylvester_1st_order_workspace!(workspaces::workspaces) - return workspaces.sylvester_1st_order +function ensure_schur_workspace!(ws::schur_workspace{T}, n::Int, nMixed::Int, nPfm::Int, nFnpm::Int) where T + companion_size = n + nMixed + if size(ws.D, 1) != companion_size || + size(ws.sol) != (n, nPfm) || + size(ws.Zโ‚โ‚) != (nPfm, nPfm) || + size(ws.Zโ‚‚โ‚) != (nFnpm, nPfm) + return Schur_workspace(n, nMixed, nPfm, nFnpm, T = T) + end + return ws end - """ ensure_lyapunov_workspace!(workspaces, n, order::Symbol) @@ -1253,50 +1643,17 @@ function ensure_lyapunov_workspace!(workspaces::workspaces, n::Int, order::Symbo workspaces.lyapunov_3rd_order = Lyapunov_workspace(n) end return workspaces.lyapunov_3rd_order + elseif order == :block + ws = workspaces.lyapunov_block + if ws.n != n + workspaces.lyapunov_block = Lyapunov_workspace(n) + end + return workspaces.lyapunov_block else - error("Invalid order: $order. Must be :first_order, :second_order, or :third_order") + error("Invalid order: $order. Must be :first_order, :second_order, :third_order, or :block") end end -""" - ensure_lyapunov_workspace_1st_order!(๐“‚) - -Ensure the first-order Lyapunov workspace is properly sized for the model. -The dimension is `nVars` (size of the covariance matrix). -""" -function ensure_lyapunov_workspace_1st_order!(๐“‚) - T = ๐“‚.constants.post_model_macro - n = T.nVars - return ensure_lyapunov_workspace!(๐“‚.workspaces, n, :first_order) -end - - -""" - ensure_inversion_workspace!(๐“‚; third_order::Bool = false) - -Ensure the inversion filter workspace is properly sized for the model. -Dimensions are based on nExo (number of shocks) and nPast_not_future_and_mixed. -""" -function ensure_inversion_workspace!(๐“‚; third_order::Bool = false) - T = ๐“‚.constants.post_model_macro - n_exo = T.nExo - n_past = T.nPast_not_future_and_mixed - ensure_inversion_buffers!(๐“‚.workspaces.inversion, n_exo, n_past; third_order = third_order) - return ๐“‚.workspaces.inversion -end - - -""" - ensure_kalman_workspace!(๐“‚) - -Ensure the Kalman filter workspace is available. Returns the workspace for use. -Actual buffer resizing happens lazily in ensure_kalman_buffers! when dimensions are known. -""" -function ensure_kalman_workspace!(๐“‚) - return ๐“‚.workspaces.kalman -end - - function create_selector_matrix(target::Vector{Symbol}, source::Vector{Symbol}) selector = spzeros(Float64, length(target), length(source)) idx = indexin(target, source) @@ -1337,6 +1694,9 @@ function ensure_model_structure_constants!(constants::constants, calibration_par vars_in_ss_equations = T.vars_in_ss_equations_no_aux vars_in_ss_equations_with_aux = T.vars_in_ss_equations + vars_and_calib = vcat(T.var, calibration_parameters) + ss_var_idx_in_var_and_calib = Int.(indexin(vars_in_ss_equations, vars_and_calib)) + calib_idx_in_var_and_calib = Int.(indexin(calibration_parameters, vars_and_calib)) extended_SS_and_pars = vcat(map(x -> Symbol(replace(string(x), r"แดธโฝโป?[โฐยนยฒยณโดโตโถโทโธโน]+โพ" => "")), T.var), calibration_parameters) custom_ss_expand_matrix = create_selector_matrix(extended_SS_and_pars, vcat(vars_in_ss_equations, calibration_parameters)) @@ -1360,6 +1720,8 @@ function ensure_model_structure_constants!(constants::constants, calibration_par custom_ss_expand_matrix = custom_ss_expand_matrix, vars_in_ss_equations = vars_in_ss_equations, vars_in_ss_equations_with_aux = vars_in_ss_equations_with_aux, + ss_var_idx_in_var_and_calib = ss_var_idx_in_var_and_calib, + calib_idx_in_var_and_calib = calib_idx_in_var_and_calib, SS_and_pars_names_lead_lag = SS_and_pars_names_lead_lag, # SS_and_pars_names_no_exo = SS_and_pars_names_no_exo, SS_and_pars_no_exo_idx = SS_and_pars_no_exo_idx, @@ -1375,28 +1737,31 @@ function compute_e4(nแต‰::Int) if nแต‰ == 0 return Float64[] end - E_e4 = zeros(nแต‰ * (nแต‰ + 1)รท2 * (nแต‰ + 2)รท3 * (nแต‰ + 3)รท4) - quadrup = multiplicate(nแต‰, 4) - comb4 = reduce(vcat, generateSumVectors(nแต‰, 4)) - comb4 = comb4 isa Int64 ? reshape([comb4], 1, 1) : comb4 - for j = 1:size(comb4, 1) - E_e4[j] = product_moments(โ„’.I(nแต‰), 1:nแต‰, comb4[j, :]) + # Isserlis' theorem for i.i.d. standard normal shocks: + # E[ฮต_a ฮต_b ฮต_c ฮต_d] = ฮด_ab ฮด_cd + ฮด_ac ฮด_bd + ฮด_ad ฮด_bc + e4 = zeros(nแต‰^4) + for d in 1:nแต‰, c in 1:nแต‰, b in 1:nแต‰, a in 1:nแต‰ + e4[a + nแต‰*(b-1) + nแต‰^2*(c-1) + nแต‰^3*(d-1)] = Float64((a==b)*(c==d) + (a==c)*(b==d) + (a==d)*(b==c)) end - return quadrup * E_e4 + return e4 end function compute_e6(nแต‰::Int) if nแต‰ == 0 return Float64[] end - E_e6 = zeros(nแต‰ * (nแต‰ + 1)รท2 * (nแต‰ + 2)รท3 * (nแต‰ + 3)รท4 * (nแต‰ + 4)รท5 * (nแต‰ + 5)รท6) - sextup = multiplicate(nแต‰, 6) - comb6 = reduce(vcat, generateSumVectors(nแต‰, 6)) - comb6 = comb6 isa Int64 ? reshape([comb6], 1, 1) : comb6 - for j = 1:size(comb6, 1) - E_e6[j] = product_moments(โ„’.I(nแต‰), 1:nแต‰, comb6[j, :]) - end - return sextup * E_e6 + # Isserlis' theorem for i.i.d. standard normal shocks: + # E[ฮต_a ฮต_b ฮต_c ฮต_d ฮต_e ฮต_f] = sum over all 15 perfect matchings + e6 = zeros(nแต‰^6) + for f in 1:nแต‰, e in 1:nแต‰, d in 1:nแต‰, c in 1:nแต‰, b in 1:nแต‰, a in 1:nแต‰ + e6[a + nแต‰*(b-1) + nแต‰^2*(c-1) + nแต‰^3*(d-1) + nแต‰^4*(e-1) + nแต‰^5*(f-1)] = Float64( + (a==b)*((c==d)*(e==f) + (c==e)*(d==f) + (c==f)*(d==e)) + + (a==c)*((b==d)*(e==f) + (b==e)*(d==f) + (b==f)*(d==e)) + + (a==d)*((b==c)*(e==f) + (b==e)*(c==f) + (b==f)*(c==e)) + + (a==e)*((b==c)*(d==f) + (b==d)*(c==f) + (b==f)*(c==d)) + + (a==f)*((b==c)*(d==e) + (b==d)*(c==e) + (b==e)*(c==d))) + end + return e6 end function ensure_moments_constants!(constants::constants) @@ -1459,7 +1824,51 @@ function ensure_moments_substate_indices!(๐“‚, nหข::Int) e_ss = sparse(reshape(โ„’.kron(vec(โ„’.I(nแต‰)), โ„’.I(nหข^2)), nแต‰ * nหข^2, nแต‰ * nหข^2)) ss_s = sparse(reshape(โ„’.kron(vec(โ„’.I(nหข^2)), โ„’.I(nหข)), nหข^3, nหข^3)) s_s = sparse(reshape(โ„’.kron(vec(โ„’.I(nหข)), โ„’.I(nหข)), nหข^2, nหข^2)) - to.substate_indices[nหข] = moments_substate_indices(I_plus_s_s, e_es, e_ss, ss_s, s_s) + + # Second-order duplication/elimination matrices (Dโ‚‚หข: nหขยฒ ร— nหข(nหข+1)/2, Lโ‚‚หข: nหข(nหข+1)/2 ร— nหขยฒ) + # Dโ‚‚หข * vech(M) = vec(M) for symmetric M; Lโ‚‚หข * vec(M) = vech(M) + # vech ordering: (1,1), (1,2), (2,2), (1,3), (2,3), (3,3), ... (upper triangle, col-major) + canonical2 = [nหข * (i-1) + k for i in 1:nหข for k in 1:i] # canonical vec positions + rows2 = Int[]; cols2 = Int[] + col_idx = 0 + for i in 1:nหข + for k in 1:i + col_idx += 1 + push!(rows2, nหข * (i-1) + k) # M_{k,i} position + push!(cols2, col_idx) + if i != k + push!(rows2, nหข * (k-1) + i) # M_{i,k} symmetric duplicate + push!(cols2, col_idx) + end + end + end + Dโ‚‚หข = sparse(rows2, cols2, 1.0, nหข^2, col_idx) + Lโ‚‚หข = sparse(1:length(canonical2), canonical2, 1.0, length(canonical2), nหข^2) + + # Third-order duplication/elimination matrices (Dโ‚ƒหข: nหขยณ ร— nหข(nหข+1)(nหข+2)/6, Lโ‚ƒหข: inverse) + # Dโ‚ƒหข * vechโ‚ƒ(T) = vec(T) for symmetric 3-tensor T; Lโ‚ƒหข * vec(T) = vechโ‚ƒ(T) + canonical3 = [nหข^2 * (i-1) + nหข * (k-1) + l for i in 1:nหข for k in 1:i for l in 1:k] + rows3 = Int[]; cols3 = Int[] + col_idx = 0 + for i in 1:nหข + for k in 1:i + for l in 1:k + col_idx += 1 + perms = Set{Tuple{Int,Int,Int}}() + for p in ((i,k,l), (i,l,k), (k,i,l), (k,l,i), (l,i,k), (l,k,i)) + push!(perms, p) + end + for (a, b, c) in perms + push!(rows3, nหข^2 * (a-1) + nหข * (b-1) + c) + push!(cols3, col_idx) + end + end + end + end + Dโ‚ƒหข = sparse(rows3, cols3, 1.0, nหข^3, col_idx) + Lโ‚ƒหข = sparse(1:length(canonical3), canonical3, 1.0, length(canonical3), nหข^3) + + to.substate_indices[nหข] = moments_substate_indices(I_plus_s_s, e_es, e_ss, ss_s, s_s, Dโ‚‚หข, Lโ‚‚หข, Dโ‚ƒหข, Lโ‚ƒหข) end return to.substate_indices[nหข] end @@ -1469,7 +1878,7 @@ function ensure_moments_dependency_kron_indices!(๐“‚, dependencies::Vector{Symb to = constants.third_order key = Tuple(dependencies) if !haskey(to.dependency_kron_indices, key) - so = ensure_computational_constants!(๐“‚) + so = ensure_computational_constants!(constants) to.dependency_kron_indices[key] = moments_dependency_kron_indices( โ„’.kron(s_in_sโบ, s_in_sโบ), โ„’.kron(s_in_sโบ, so.e_in_sโบ), @@ -1480,24 +1889,188 @@ function ensure_moments_dependency_kron_indices!(๐“‚, dependencies::Vector{Symb end -struct Tolerances - NSSS_acceptance_tol::AbstractFloat - NSSS_xtol::AbstractFloat - NSSS_ftol::AbstractFloat - NSSS_rel_xtol::AbstractFloat +""" + SolverTolerances + +Tolerance settings for a single numerical equation solver (Sylvester, Lyapunov, or QME). + +# Fields +- `atol::Float64`: absolute convergence tolerance (used by Krylov solvers). +- `rtol::Float64`: relative convergence tolerance (used by iterative stopping checks). +- `initial_guess_acceptance_tol::Float64`: if an initial guess achieves a relative + residual below this threshold it is accepted immediately, skipping the full solve. +- `acceptance_tol::Float64`: result is accepted when the relative residual falls below + this threshold; otherwise the dispatcher retries with a fallback algorithm. - qme_tol::AbstractFloat - qme_acceptance_tol::AbstractFloat +Construct via `SolverTolerances(; atol, rtol, initial_guess_acceptance_tol, acceptance_tol)`. +Default values differ by solver type and are set by the enclosing tolerance hierarchy; +see [`Tolerances`](@ref) and [`FirstOrderTolerances`](@ref) / [`HigherOrderTolerances`](@ref). +""" +struct SolverTolerances + atol::Float64 + rtol::Float64 + initial_guess_acceptance_tol::Float64 + acceptance_tol::Float64 +end - sylvester_tol::AbstractFloat - sylvester_acceptance_tol::AbstractFloat +function SolverTolerances(; atol::Float64 = 1e-14, + rtol::Float64 = 1e-14, + initial_guess_acceptance_tol::Float64 = 1e-10, + acceptance_tol::Float64 = 1e-10) + return SolverTolerances(atol, rtol, initial_guess_acceptance_tol, acceptance_tol) +end - lyapunov_tol::AbstractFloat - lyapunov_acceptance_tol::AbstractFloat +""" + NsssTolerances + +Tolerance settings for the non-stochastic steady state (NSSS) solver. + +# Fields +- `acceptance_tol::Float64` [Default: `1e-12`]: solution is accepted when the residual + norm falls below this value. +- `initial_guess_acceptance_tol::Float64` [Default: `1e-12`]: an initial guess is reused + when its residual is below this threshold. +- `xtol::Float64` [Default: `1e-12`]: absolute step-size tolerance. +- `ftol::Float64` [Default: `1e-14`]: absolute function-value tolerance. +- `rel_xtol::Float64` [Default: `eps()`]: relative step-size tolerance. + +Construct via `NsssTolerances(; acceptance_tol, initial_guess_acceptance_tol, xtol, ftol, rel_xtol)`. +""" +struct NsssTolerances + acceptance_tol::Float64 + initial_guess_acceptance_tol::Float64 + xtol::Float64 + ftol::Float64 + rel_xtol::Float64 +end + +function NsssTolerances(; acceptance_tol::Float64 = 1e-12, + initial_guess_acceptance_tol::Float64 = 1e-12, + xtol::Float64 = 1e-12, + ftol::Float64 = 1e-14, + rel_xtol::Float64 = eps()) + return NsssTolerances(acceptance_tol, initial_guess_acceptance_tol, xtol, ftol, rel_xtol) +end + +""" + AdTolerances + +Tolerance settings passed to the automatic differentiation (AD) paths of each equation +solver. Each field is a [`SolverTolerances`](@ref) that controls the corresponding solver +when it is called inside a ForwardDiff dual-number overload or a ChainRulesCore rrule. + +# Fields +- `qme::SolverTolerances`: tolerances for the quadratic matrix equation (QME) derivative solve. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-8`, `acceptance_tol=1e-8`. +- `sylvester::SolverTolerances`: tolerances for the Sylvester equation derivative solve. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-10`, `acceptance_tol=1e-10`. +- `lyapunov::SolverTolerances`: tolerances for the Lyapunov equation derivative solve. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-12`, `acceptance_tol=1e-12`. + +Construct via `AdTolerances(; qme, sylvester, lyapunov)`. +""" +struct AdTolerances + qme::SolverTolerances + sylvester::SolverTolerances + lyapunov::SolverTolerances +end + +function AdTolerances(; qme::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-8, + acceptance_tol = 1e-8), + sylvester::SolverTolerances = SolverTolerances(), + lyapunov::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12)) + return AdTolerances(qme, sylvester, lyapunov) +end + +""" + FirstOrderTolerances + +Tolerance settings for the first-order perturbation solution and its AD pathways. + +# Fields +- `qme::SolverTolerances`: tolerances for the quadratic matrix equation solver. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-8`, `acceptance_tol=1e-8`. +- `lyapunov::SolverTolerances`: tolerances for the Lyapunov equation solver used to + compute first-order covariance matrices. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-12`, `acceptance_tol=1e-12`. +- `droptol::Float64` [Default: `1e-14`]: entries smaller than this threshold in solution + matrices are dropped (set to zero) to reduce sparsity fill-in. +- `dependencies_tol::Float64` [Default: `1e-12`]: threshold for determining variable + dependencies when isolating subsystems for covariance statistics. +- `ad::AdTolerances`: tolerances used in the AD derivative evaluation paths. + +Construct via `FirstOrderTolerances(; qme, lyapunov, droptol, dependencies_tol, ad)`. +""" +struct FirstOrderTolerances + qme::SolverTolerances + lyapunov::SolverTolerances + droptol::Float64 + dependencies_tol::Float64 + ad::AdTolerances +end + +function FirstOrderTolerances(; qme::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-8, + acceptance_tol = 1e-8), + lyapunov::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), + droptol::Float64 = 1e-14, + dependencies_tol::Float64 = 1e-12, + ad::AdTolerances = AdTolerances()) + return FirstOrderTolerances(qme, lyapunov, droptol, dependencies_tol, ad) +end + +""" + HigherOrderTolerances + +Tolerance settings for second- and third-order perturbation solutions and their AD pathways. + +# Fields +- `sylvester::SolverTolerances`: tolerances for the Sylvester equation solver. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-10`, `acceptance_tol=1e-10`. +- `lyapunov::SolverTolerances`: tolerances for the Lyapunov equation solver used to + compute higher-order covariance matrices. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-12`, `acceptance_tol=1e-12`. +- `droptol::Float64` [Default: `1e-14`]: entries smaller than this threshold in solution + matrices are dropped (set to zero) to reduce sparsity fill-in. +- `dependencies_tol::Float64` [Default: `1e-12`]: threshold for determining variable + dependencies when isolating subsystems for covariance statistics. +- `ad::AdTolerances`: tolerances used in the AD derivative evaluation paths. + +Construct via `HigherOrderTolerances(; sylvester, lyapunov, droptol, dependencies_tol, ad)`. +""" +struct HigherOrderTolerances + sylvester::SolverTolerances + lyapunov::SolverTolerances + droptol::Float64 + dependencies_tol::Float64 + ad::AdTolerances +end - droptol::AbstractFloat +function HigherOrderTolerances(; sylvester::SolverTolerances = SolverTolerances(), + lyapunov::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), + droptol::Float64 = 1e-14, + dependencies_tol::Float64 = 1e-12, + ad::AdTolerances = AdTolerances()) + return HigherOrderTolerances(sylvester, lyapunov, droptol, dependencies_tol, ad) +end - dependencies_tol::AbstractFloat +struct Tolerances + nsss::NsssTolerances + first_order::FirstOrderTolerances + second_order::HigherOrderTolerances + third_order::HigherOrderTolerances end struct CalculationOptions @@ -1515,61 +2088,203 @@ end @stable default_mode = "disable" begin """ $(SIGNATURES) -Function to manually define tolerances for the solvers of various problems: non-stochastic steady state solver (NSSS), Sylvester equations, Lyapunov equation, and quadratic matrix equation (qme). + +Define tolerances for the numerical solvers used throughout model solution and estimation. +Tolerances are organised in a two-level hierarchy: + +``` +Tolerances +โ”œโ”€โ”€ nsss :: NsssTolerances โ€” non-stochastic steady state solver +โ”œโ”€โ”€ first_order :: FirstOrderTolerances โ€” first-order perturbation solution +โ”‚ โ”œโ”€โ”€ qme :: SolverTolerances โ€” quadratic matrix equation (QME) +โ”‚ โ”œโ”€โ”€ lyapunov :: SolverTolerances โ€” Lyapunov equation +โ”‚ โ”œโ”€โ”€ droptol โ€” zero-threshold for solution matrices +โ”‚ โ”œโ”€โ”€ dependencies_tol โ€” subsystem isolation threshold +โ”‚ โ””โ”€โ”€ ad :: AdTolerances โ€” AD derivative paths +โ”‚ โ”œโ”€โ”€ qme :: SolverTolerances +โ”‚ โ”œโ”€โ”€ sylvester:: SolverTolerances +โ”‚ โ””โ”€โ”€ lyapunov :: SolverTolerances +โ”œโ”€โ”€ second_order :: HigherOrderTolerances โ€” second-order perturbation solution +โ”‚ โ”œโ”€โ”€ sylvester :: SolverTolerances โ€” Sylvester equation +โ”‚ โ”œโ”€โ”€ lyapunov :: SolverTolerances โ€” Lyapunov equation +โ”‚ โ”œโ”€โ”€ droptol / dependencies_tol +โ”‚ โ””โ”€โ”€ ad :: AdTolerances +โ””โ”€โ”€ third_order :: HigherOrderTolerances โ€” third-order perturbation solution + โ””โ”€โ”€ (same structure as second_order) +``` + +Each [`SolverTolerances`](@ref) carries four values: +- `atol`: absolute convergence tolerance used by Krylov solvers. +- `rtol`: relative convergence tolerance used by iterative stopping checks. +- `initial_guess_acceptance_tol`: accept an initial guess without re-solving if its + residual is already below this threshold. +- `acceptance_tol`: accept the final result when the residual falls below this threshold; + otherwise the dispatcher retries with a fallback algorithm. # Keyword Arguments -- `NSSS_acceptance_tol` [Default: `1e-12`, Type: `AbstractFloat`]: Acceptance tolerance for non-stochastic steady state solver. -- `NSSS_xtol` [Default: `1e-12`, Type: `AbstractFloat`]: Absolute tolerance for solver steps for non-stochastic steady state solver. -- `NSSS_ftol` [Default: `1e-14`, Type: `AbstractFloat`]: Absolute tolerance for solver function values for non-stochastic steady state solver. -- `NSSS_rel_xtol` [Default: `eps()`, Type: `AbstractFloat`]: Relative tolerance for solver steps for non-stochastic steady state solver. +- `nsss` [Default: `NsssTolerances()`]: tolerances for the non-stochastic steady state + solver. See [`NsssTolerances`](@ref). +- `first_order` [Default: `FirstOrderTolerances()`]: tolerances for the first-order + solution and its AD paths. See [`FirstOrderTolerances`](@ref). +- `second_order` [Default: `HigherOrderTolerances()`]: tolerances for the second-order + solution and its AD paths. See [`HigherOrderTolerances`](@ref). +- `third_order` [Default: `HigherOrderTolerances()`]: tolerances for the third-order + solution and its AD paths. See [`HigherOrderTolerances`](@ref). + +# Examples +```julia +# use defaults +tol = Tolerances() + +# tighten the NSSS solver +tol = Tolerances(nsss = NsssTolerances(xtol = 1e-14)) + +# tighten second- and third-order Sylvester/Lyapunov solvers +tight = SolverTolerances(acceptance_tol = 1e-14) +tol = Tolerances( + second_order = HigherOrderTolerances(sylvester = tight, lyapunov = tight), + third_order = HigherOrderTolerances(sylvester = tight, lyapunov = tight), +) +``` +""" +function Tolerances(; nsss::NsssTolerances = NsssTolerances(), + first_order::FirstOrderTolerances = FirstOrderTolerances(), + second_order::HigherOrderTolerances = HigherOrderTolerances(), + third_order::HigherOrderTolerances = HigherOrderTolerances()) + return Tolerances(nsss, first_order, second_order, third_order) +end -- `qme_tol` [Default: `1e-14`, Type: `AbstractFloat`]: Tolerance for quadratic matrix equation solver. -- `qme_acceptance_tol` [Default: `1e-8`, Type: `AbstractFloat`]: Acceptance tolerance for quadratic matrix equation solver. -- `sylvester_tol` [Default: `1e-14`, Type: `AbstractFloat`]: Tolerance for Sylvester equation solver. -- `sylvester_acceptance_tol` [Default: `1e-10`, Type: `AbstractFloat`]: Acceptance tolerance for Sylvester equation solver. +const HIGHER_ORDER_ALGORITHMS = (:second_order, :pruned_second_order, :third_order, :pruned_third_order) +const THIRD_ORDER_ALGORITHMS = (:third_order, :pruned_third_order) + +""" + solver_tol_to_dict(st::SolverTolerances) -> Dict{Symbol,Any} + +Convert a [`SolverTolerances`](@ref) struct to a flat `Dict`. +""" +function solver_tol_to_dict(st::SolverTolerances) + return Dict{Symbol,Any}( + :atol => st.atol, + :rtol => st.rtol, + :initial_guess_acceptance_tol => st.initial_guess_acceptance_tol, + :acceptance_tol => st.acceptance_tol, + ) +end -- `lyapunov_tol` [Default: `1e-14`, Type: `AbstractFloat`]: Tolerance for Lyapunov equation solver. -- `lyapunov_acceptance_tol` [Default: `1e-12`, Type: `AbstractFloat`]: Acceptance tolerance for Lyapunov equation solver. +""" + nsss_tol_to_dict(nt::NsssTolerances) -> Dict{Symbol,Any} -- `droptol` [Default: `1e-14`, Type: `AbstractFloat`]: Tolerance below which matrix entries are considered 0. +Convert a [`NsssTolerances`](@ref) struct to a flat `Dict`. +""" +function nsss_tol_to_dict(nt::NsssTolerances) + return Dict{Symbol,Any}( + :acceptance_tol => nt.acceptance_tol, + :initial_guess_acceptance_tol => nt.initial_guess_acceptance_tol, + :xtol => nt.xtol, + :ftol => nt.ftol, + :rel_xtol => nt.rel_xtol, + ) +end -- `dependencies_tol` [Default: `1e-12`, Type: `AbstractFloat`]: tolerance for the effect of a variable on the variable of interest when isolating part of the system for calculating covariance related statistics """ -function Tolerances(;NSSS_acceptance_tol::AbstractFloat = 1e-12, - NSSS_xtol::AbstractFloat = 1e-12, - NSSS_ftol::AbstractFloat = 1e-14, - NSSS_rel_xtol::AbstractFloat = eps(), - - qme_tol::AbstractFloat = 1e-14, - qme_acceptance_tol::AbstractFloat = 1e-8, + tol_to_dict(tol::Tolerances, algorithm::Symbol; needs_covariance::Bool = false) -> Dict{Symbol,Any} - sylvester_tol::AbstractFloat = 1e-14, - sylvester_acceptance_tol::AbstractFloat = 1e-10, +Build a nested `Dict` of tolerance values that are **relevant** for the given +`algorithm` and covariance requirement. Irrelevant sub-trees (e.g. third-order +tolerances when running a first-order solve) are omitted so that +`compare_args_and_kwargs` never reports spurious differences in unused settings. - lyapunov_tol::AbstractFloat = 1e-14, - lyapunov_acceptance_tol::AbstractFloat = 1e-12, +AD sub-tolerances are always excluded (too internal for plot annotations). +""" +function tol_to_dict(tol::Tolerances, algorithm::Symbol; needs_covariance::Bool = false) + d = Dict{Symbol,Any}() + + # NSSS โ€” always relevant + d[:nsss] = nsss_tol_to_dict(tol.nsss) + + # First-order โ€” always relevant + fo = Dict{Symbol,Any}(:qme => solver_tol_to_dict(tol.first_order.qme), + :droptol => tol.first_order.droptol) + if needs_covariance + fo[:lyapunov] = solver_tol_to_dict(tol.first_order.lyapunov) + fo[:dependencies_tol] = tol.first_order.dependencies_tol + end + d[:first_order] = fo + + # Second-order โ€” only for higher-order algorithms + if algorithm in HIGHER_ORDER_ALGORITHMS + so = Dict{Symbol,Any}(:sylvester => solver_tol_to_dict(tol.second_order.sylvester), + :droptol => tol.second_order.droptol) + if needs_covariance + so[:lyapunov] = solver_tol_to_dict(tol.second_order.lyapunov) + so[:dependencies_tol] = tol.second_order.dependencies_tol + end + d[:second_order] = so + end - droptol::AbstractFloat = 1e-14, + # Third-order โ€” only for third-order algorithms + if algorithm in THIRD_ORDER_ALGORITHMS + to = Dict{Symbol,Any}(:sylvester => solver_tol_to_dict(tol.third_order.sylvester), + :droptol => tol.third_order.droptol) + if needs_covariance + to[:lyapunov] = solver_tol_to_dict(tol.third_order.lyapunov) + to[:dependencies_tol] = tol.third_order.dependencies_tol + end + d[:third_order] = to + end - dependencies_tol::AbstractFloat = 1e-12) - - return Tolerances(NSSS_acceptance_tol, - NSSS_xtol, - NSSS_ftol, - NSSS_rel_xtol, - qme_tol, - qme_acceptance_tol, - sylvester_tol, - sylvester_acceptance_tol, - lyapunov_tol, - lyapunov_acceptance_tol, - droptol, - dependencies_tol) -end - - -function merge_calculation_options(;quadratic_matrix_equation_algorithm::Symbol = :schur, + return d +end + +""" + warn_irrelevant_tol(tol::Tolerances, algorithm::Symbol; needs_covariance::Bool = false) + +Emit `@info` messages when `tol` contains non-default values in sub-trees that +have **no effect** for the given `algorithm` and covariance setting. This gives +users immediate feedback that their custom tolerances are being ignored. +""" +function warn_irrelevant_tol(tol::Tolerances, algorithm::Symbol; needs_covariance::Bool = false) + defaults = Tolerances() + + # --- order-based irrelevance --- + if algorithm โˆ‰ HIGHER_ORDER_ALGORITHMS + if tol.second_order != defaults.second_order + @info "Second-order tolerances have no effect with algorithm = :$algorithm and are ignored." + end + end + + if algorithm โˆ‰ THIRD_ORDER_ALGORITHMS + if tol.third_order != defaults.third_order + @info "Third-order tolerances have no effect with algorithm = :$algorithm and are ignored." + end + end + + # --- covariance-based irrelevance --- + if !needs_covariance + if tol.first_order.lyapunov != defaults.first_order.lyapunov || + tol.first_order.dependencies_tol != defaults.first_order.dependencies_tol + @info "First-order Lyapunov/dependencies tolerances have no effect without covariance computation (current operation does not require it) and are ignored." + end + + if algorithm in HIGHER_ORDER_ALGORITHMS + if tol.second_order.lyapunov != defaults.second_order.lyapunov || + tol.second_order.dependencies_tol != defaults.second_order.dependencies_tol + @info "Second-order Lyapunov/dependencies tolerances have no effect without covariance computation (current operation does not require it) and are ignored." + end + end + + if algorithm in THIRD_ORDER_ALGORITHMS + if tol.third_order.lyapunov != defaults.third_order.lyapunov || + tol.third_order.dependencies_tol != defaults.third_order.dependencies_tol + @info "Third-order Lyapunov/dependencies tolerances have no effect without covariance computation (current operation does not require it) and are ignored." + end + end + end +end + + +function merge_calculation_options(;quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, sylvester_algorithmยฒ::Symbol = :doubling, sylvester_algorithmยณ::Symbol = :bicgstab, lyapunov_algorithm::Symbol = :doubling, diff --git a/src/perturbation.jl b/src/perturbation.jl index 1fbe42d05..609878719 100644 --- a/src/perturbation.jl +++ b/src/perturbation.jl @@ -2,10 +2,23 @@ function calculate_first_order_solution(โˆ‡โ‚::Matrix{R}, constants::constants, - qme_ws::qme_workspace{R,S}, - sylv_ws::sylvester_workspace{R,S}; + workspaces::workspaces, + cache::caches; opts::CalculationOptions = merge_calculation_options(), - initial_guess::AbstractMatrix{R} = zeros(0,0))::Tuple{Matrix{R}, Matrix{R}, Bool} where {R <: AbstractFloat, S <: Real} + use_fastlapack_qr::Bool = true, + use_fastlapack_lu::Bool = true, + initial_guess::AbstractMatrix{R} = zeros(0,0), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true)::Tuple{Matrix{R}, Matrix{R}, Bool} where {R <: AbstractFloat} + # Cache hit: return cached first-order solution if valid for current parameters + if caching && R === Float64 && !isempty(parameter_values) && + cache_valid_for_parameters(cache.valid_for.first_order_solution, parameter_values) + Sโ‚_cached = cache.first_order_solution_matrix + qme_cached = cache.qme_solution + if Sโ‚_cached isa Matrix{R} && !isempty(Sโ‚_cached) && qme_cached isa Matrix{R} && !isempty(qme_cached) + return Sโ‚_cached, qme_cached, true + end + end # @timeit_debug timer "Calculate 1st order solution" begin # @timeit_debug timer "Preprocessing" begin @@ -17,102 +30,200 @@ function calculate_first_order_solution(โˆ‡โ‚::Matrix{R}, comb = idx_constants.comb future_not_past_and_mixed_in_comb = idx_constants.future_not_past_and_mixed_in_comb past_not_future_and_mixed_in_comb = idx_constants.past_not_future_and_mixed_in_comb + past_not_future_and_mixed_in_present_but_not_only = idx_constants.past_not_future_and_mixed_in_present_but_not_only Ir = idx_constants.Ir - - โˆ‡โ‚Š = โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] - โˆ‡โ‚€ = โˆ‡โ‚[:,idx_constants.nabla_zero_cols] - โˆ‡โ‚‹ = โˆ‡โ‚[:,idx_constants.nabla_minus_cols] - โˆ‡โ‚‘ = โˆ‡โ‚[:,idx_constants.nabla_e_start:end] + + qme_ws = workspaces.first_order + + ensure_first_order_workspace_buffers!(qme_ws, T, length(dynIndex), length(comb)) + + โˆ‡โ‚Š = @view โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed] + โˆ‡โ‚€ = qme_ws.โˆ‡โ‚€ + copyto!(โˆ‡โ‚€, @view(โˆ‡โ‚[:,idx_constants.nabla_zero_cols])) + โˆ‡โ‚‹ = @view โˆ‡โ‚[:,idx_constants.nabla_minus_cols] + โˆ‡โ‚‘ = qme_ws.โˆ‡โ‚‘ + copyto!(โˆ‡โ‚‘, @view(โˆ‡โ‚[:,idx_constants.nabla_e_start:end])) # end # timeit_debug # @timeit_debug timer "Invert โˆ‡โ‚€" begin - Q = โ„’.qr!(โˆ‡โ‚€[:,T.present_only_idx]) - - Aโ‚Š = Q.Q' * โˆ‡โ‚Š - Aโ‚€ = Q.Q' * โˆ‡โ‚€ - Aโ‚‹ = Q.Q' * โˆ‡โ‚‹ + Aโ‚Š = qme_ws.๐€โ‚Š + Aโ‚€ = qme_ws.๐€โ‚€ + Aโ‚‹ = qme_ws.๐€โ‚‹ + โˆ‡โ‚€_present = @view โˆ‡โ‚€[:, T.present_only_idx] + # Legacy readable flow (before allocation-focused refactor): + # Q = qr!(โˆ‡โ‚€[:, T.present_only_idx]) + # Aโ‚Š = Q.Q' * โˆ‡โ‚Š; Aโ‚€ = Q.Q' * โˆ‡โ‚€; Aโ‚‹ = Q.Q' * โˆ‡โ‚‹ + # Current code performs the same transforms using reusable QR/ORM workspaces. + qr_factors, qr_ws = ensure_first_order_fast_qr_workspace!(qme_ws, โˆ‡โ‚€_present) + Q = factorize_qr!(โˆ‡โ‚€_present, qr_factors, qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + + qme_ws.fast_qr_orm_ws_plus, qme_ws.fast_qr_orm_dims_plus = apply_qr_transpose_left!(Aโ‚Š, โˆ‡โ‚Š, Q, + qme_ws.fast_qr_orm_ws_plus, + qme_ws.fast_qr_orm_dims_plus, + qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + qme_ws.fast_qr_orm_ws_zero, qme_ws.fast_qr_orm_dims_zero = apply_qr_transpose_left!(Aโ‚€, โˆ‡โ‚€, Q, + qme_ws.fast_qr_orm_ws_zero, + qme_ws.fast_qr_orm_dims_zero, + qr_ws; + use_fastlapack_qr = use_fastlapack_qr) + qme_ws.fast_qr_orm_ws_minus, qme_ws.fast_qr_orm_dims_minus = apply_qr_transpose_left!(Aโ‚‹, โˆ‡โ‚‹, Q, + qme_ws.fast_qr_orm_ws_minus, + qme_ws.fast_qr_orm_dims_minus, + qr_ws; + use_fastlapack_qr = use_fastlapack_qr) # end # timeit_debug # @timeit_debug timer "Sort matrices" begin - Aฬƒโ‚Š = Aโ‚Š[dynIndex,:] * Ir[future_not_past_and_mixed_in_comb,:] - Aฬƒโ‚€ = Aโ‚€[dynIndex, comb] - Aฬƒโ‚‹ = Aโ‚‹[dynIndex,:] * Ir[past_not_future_and_mixed_in_comb,:] + Aฬƒโ‚Š = qme_ws.๐€ฬƒโ‚Š + โ„’.mul!(Aฬƒโ‚Š, @view(Aโ‚Š[dynIndex,:]), @view(Ir[future_not_past_and_mixed_in_comb,:])) + + Aฬƒโ‚€ = qme_ws.๐€ฬƒโ‚€ + copyto!(Aฬƒโ‚€, @view(Aโ‚€[dynIndex, comb])) + + Aฬƒโ‚‹ = qme_ws.๐€ฬƒโ‚‹ + โ„’.mul!(Aฬƒโ‚‹, @view(Aโ‚‹[dynIndex,:]), @view(Ir[past_not_future_and_mixed_in_comb,:])) # end # timeit_debug # @timeit_debug timer "Quadratic matrix equation solve" begin - sol, solved = solve_quadratic_matrix_equation(Aฬƒโ‚Š, Aฬƒโ‚€, Aฬƒโ‚‹, constants, qme_ws; + sol, solved = solve_quadratic_matrix_equation(Aฬƒโ‚Š, Aฬƒโ‚€, Aฬƒโ‚‹, constants, workspaces, cache; initial_guess = initial_guess, quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - tol = opts.tol.qme_tol, - acceptance_tol = opts.tol.qme_acceptance_tol, + use_fastlapack_lu = use_fastlapack_lu, + tol = opts.tol.first_order.qme, verbose = opts.verbose) if !solved if opts.verbose println("Quadratic matrix equation solution failed.") end - return zeros(R, T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false + return fill(R(NaN), T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false end # end # timeit_debug # @timeit_debug timer "Postprocessing" begin # @timeit_debug timer "Setup matrices" begin - sol_compact = sol[reverse_dynamic_order, past_not_future_and_mixed_in_comb] + sol_compact = @view sol[reverse_dynamic_order, past_not_future_and_mixed_in_comb] + + n_dyn = length(reverse_dynamic_order) + ๐ƒ = @view sol[@view(reverse_dynamic_order[n_dyn - T.nFuture_not_past_and_mixed + 1:n_dyn]), past_not_future_and_mixed_in_comb] + + L = @view sol[past_not_future_and_mixed_in_present_but_not_only, past_not_future_and_mixed_in_comb] + + Aฬ„โ‚€แตค = qme_ws.๐€ฬ„โ‚€แตค + copyto!(Aฬ„โ‚€แตค, @view(Aโ‚€[1:T.nPresent_only, T.present_only_idx])) - D = sol_compact[end - T.nFuture_not_past_and_mixed + 1:end, :] + Aโ‚Šแตค = qme_ws.๐€โ‚Šแตค + copyto!(Aโ‚Šแตค, @view(Aโ‚Š[1:T.nPresent_only,:])) - L = sol[indexin(T.past_not_future_and_mixed_idx, T.present_but_not_only_idx), past_not_future_and_mixed_in_comb] + Aฬƒโ‚€แตค = qme_ws.๐€ฬƒโ‚€แตค + copyto!(Aฬƒโ‚€แตค, @view(Aโ‚€[1:T.nPresent_only, T.present_but_not_only_idx])) - Aฬ„โ‚€แตค = Aโ‚€[1:T.nPresent_only, T.present_only_idx] - Aโ‚Šแตค = Aโ‚Š[1:T.nPresent_only,:] - Aฬƒโ‚€แตค = Aโ‚€[1:T.nPresent_only, T.present_but_not_only_idx] - Aโ‚‹แตค = Aโ‚‹[1:T.nPresent_only,:] + Aโ‚‹แตค = qme_ws.๐€โ‚‹แตค + copyto!(Aโ‚‹แตค, @view(Aโ‚‹[1:T.nPresent_only,:])) # end # timeit_debug # @timeit_debug timer "Invert Aฬ„โ‚€แตค" begin - Aฬ„ฬ‚โ‚€แตค = โ„’.lu!(Aฬ„โ‚€แตค, check = false) + qme_ws.fast_lu_ws_a0u, qme_ws.fast_lu_dims_a0u, solved_Aฬ„โ‚€แตค, Aฬ„ฬ‚โ‚€แตค = factorize_lu!(Aฬ„โ‚€แตค, + qme_ws.fast_lu_ws_a0u, + qme_ws.fast_lu_dims_a0u; + use_fastlapack_lu = use_fastlapack_lu) - if !โ„’.issuccess(Aฬ„ฬ‚โ‚€แตค) + if !solved_Aฬ„โ‚€แตค if opts.verbose println("Factorisation of Aฬ„โ‚€แตค failed") end - return zeros(R, T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false + return fill(R(NaN), T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false end # A = vcat(-(Aฬ„ฬ‚โ‚€แตค \ (Aโ‚Šแตค * D * L + Aฬƒโ‚€แตค * sol[T.dynamic_order,:] + Aโ‚‹แตค)), sol) if T.nPresent_only > 0 - โ„’.mul!(Aโ‚‹แตค, Aฬƒโ‚€แตค, sol[:,past_not_future_and_mixed_in_comb], 1, 1) - nโ‚šโ‚‹ = Aโ‚Šแตค * D + โ„’.mul!(Aโ‚‹แตค, Aฬƒโ‚€แตค, @view(sol[:,past_not_future_and_mixed_in_comb]), 1, 1) + nโ‚šโ‚‹ = qme_ws.๐งโ‚šโ‚‹ + โ„’.mul!(nโ‚šโ‚‹, Aโ‚Šแตค, ๐ƒ) โ„’.mul!(Aโ‚‹แตค, nโ‚šโ‚‹, L, 1, 1) - โ„’.ldiv!(Aฬ„ฬ‚โ‚€แตค, Aโ‚‹แตค) + solve_lu_left!(Aฬ„โ‚€แตค, Aโ‚‹แตค, qme_ws.fast_lu_ws_a0u, Aฬ„ฬ‚โ‚€แตค; + use_fastlapack_lu = use_fastlapack_lu) โ„’.rmul!(Aโ‚‹แตค, -1) end + + A = qme_ws.๐€ + # Legacy readable flow: + # A = vcat(Aโ‚‹แตค, sol_compact)[T.reorder, :] + # Expanded loop below writes into preallocated `A` without temporary concatenation. + n_cols = size(A, 2) - A = vcat(Aโ‚‹แตค, sol_compact)[T.reorder,:] + for i in 1:T.nVars + src = T.reorder[i] + if src <= T.nPresent_only + for j in 1:n_cols + @inbounds A[i, j] = Aโ‚‹แตค[src, j] + end + else + src_idx = src - T.nPresent_only + for j in 1:n_cols + @inbounds A[i, j] = sol_compact[src_idx, j] + end + end + end # end # timeit_debug # end # timeit_debug # @timeit_debug timer "Exogenous part solution" begin - M = A[T.future_not_past_and_mixed_idx,:] * idx_constants.expand_past + M = qme_ws.๐Œ + # Legacy readable flow: + # M = A[T.future_not_past_and_mixed_idx, :] * expand_past + # โˆ‡โ‚€ = โˆ‡โ‚[:, 1:T.nFuture_not_past_and_mixed] * M + โˆ‡โ‚€ + โ„’.mul!(M, @view(A[T.future_not_past_and_mixed_idx,:]), idx_constants.expand_past) - โ„’.mul!(โˆ‡โ‚€, โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed], M, 1, 1) + โ„’.mul!(โˆ‡โ‚€, @view(โˆ‡โ‚[:,1:T.nFuture_not_past_and_mixed]), M, 1, 1) - C = โ„’.lu!(โˆ‡โ‚€, check = false) - - if !โ„’.issuccess(C) + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_โˆ‡โ‚€, C = factorize_lu!(โˆ‡โ‚€, + qme_ws.fast_lu_ws_nabla0, + qme_ws.fast_lu_dims_nabla0; + use_fastlapack_lu = use_fastlapack_lu) + + if !solved_โˆ‡โ‚€ if opts.verbose println("Factorisation of โˆ‡โ‚€ failed") end - return zeros(R, T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false + return fill(R(NaN), T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false end - - โ„’.ldiv!(C, โˆ‡โ‚‘) + + solve_lu_left!(โˆ‡โ‚€, โˆ‡โ‚‘, qme_ws.fast_lu_ws_nabla0, C; + use_fastlapack_lu = use_fastlapack_lu) โ„’.rmul!(โˆ‡โ‚‘, -1) # end # timeit_debug # end # timeit_debug - return hcat(A, โˆ‡โ‚‘), sol, true + n_rows = size(A, 1) + n_cols_A = size(A, 2) + n_cols_ฯต = size(โˆ‡โ‚‘, 2) + total_cols = n_cols_A + n_cols_ฯต + + Sโ‚ = if caching + Sโ‚_existing = cache.first_order_solution_matrix + if Sโ‚_existing isa Matrix{R} && size(Sโ‚_existing) == (n_rows, total_cols) + copyto!(@view(Sโ‚_existing[:, 1:n_cols_A]), A) + copyto!(@view(Sโ‚_existing[:, n_cols_A+1:total_cols]), โˆ‡โ‚‘) + Sโ‚_existing + else + Sโ‚_tmp = hcat(A, โˆ‡โ‚‘) + cache.first_order_solution_matrix = Sโ‚_tmp + Sโ‚_tmp + end + else + hcat(A, โˆ‡โ‚‘) + end + + # Stamp cache validity for current parameters + if caching && !isempty(parameter_values) + cache.valid_for.first_order_solution = eltype(parameter_values) <: โ„ฑ.Dual ? Float64.(โ„ฑ.value.(parameter_values)) : Float64.(parameter_values) + end + + return Sโ‚, sol, true end @@ -120,9 +231,20 @@ function calculate_second_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order โˆ‡โ‚‚::SparseMatrixCSC{S}, #second order derivatives ๐‘บโ‚::AbstractMatrix{S},#first order solution constants::constants, - workspaces::workspaces; + workspaces::workspaces, + cache::caches; initial_guess::AbstractMatrix{R} = zeros(0,0), - opts::CalculationOptions = merge_calculation_options())::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {R <: Real, S <: Real} + opts::CalculationOptions = merge_calculation_options(), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true)::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {R <: Real, S <: Real} + # Cache hit: return cached second-order solution if valid for current parameters + if caching && S === Float64 && !isempty(parameter_values) && + cache_valid_for_parameters(cache.valid_for.second_order_solution, parameter_values) + cached = cache.second_order_solution + if cached isa Matrix{S} && !isempty(cached) + return cached, true + end + end if !(eltype(workspaces.second_order.Sฬ‚) == S) workspaces.second_order = Higher_order_workspace(T = S) end @@ -143,13 +265,29 @@ function calculate_second_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order n = T.nVars nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + ensure_higher_order_solution_buffers!(โ„‚, n, nโ‚‘โ‚‹) + + initial_guess_sylv = if length(initial_guess) == 0 + zeros(S, 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{S} ? initial_guess : Matrix{S}(initial_guess) + else + zeros(S, 0, 0) + end + # @timeit_debug timer "Setup matrices" begin # 1st order solution - ๐’โ‚ = @views [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]]# |> sparse + ๐’โ‚ = โ„‚.๐’โ‚::Matrix{S} + copyto!(@view(๐’โ‚[:,1:nโ‚‹]), @view(๐‘บโ‚[:,1:nโ‚‹])) + fill!(@view(๐’โ‚[:,nโ‚‹+1]), zero(S)) + copyto!(@view(๐’โ‚[:,nโ‚‹+2:end]), @view(๐‘บโ‚[:,nโ‚‹+1:end])) # droptol!(๐’โ‚,tol) - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) โ„’.I(nโ‚‘ + 1)[1,:] zeros(nโ‚‘ + 1, nโ‚‘)]# |> sparse + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} + copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹,:]), @view(๐’โ‚[iโ‚‹,:])) + fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1:end,:]), zero(S)) + @inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1,nโ‚‹+1] = one(S) # droptol!(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘,tol) ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0) @@ -162,7 +300,7 @@ function calculate_second_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)]# |> sparse # droptol!(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹,tol) - โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.I(n)[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] + โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * Mโ‚‚.๐ˆโ‚™โ‚‹ - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] # end # timeit_debug @@ -179,19 +317,35 @@ function calculate_second_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # spinv = choose_matrix_format(spinv) # end # timeit_debug - # @timeit_debug timer "Setup second order matrices" begin # @timeit_debug timer "A" begin - โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * โ„’.I(n)[iโ‚Š,:] + โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š A = โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu \ โˆ‡โ‚โ‚Š # end # timeit_debug # @timeit_debug timer "C" begin - # โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = โˆ‡โ‚‚ * (โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›”) * Mโ‚‚.๐‚โ‚‚ - โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, Mโ‚‚.๐‚โ‚‚) + mat_mult_kron(โˆ‡โ‚‚, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, Mโ‚‚.๐›” * Mโ‚‚.๐‚โ‚‚) + # Build first forcing term directly in compressed Hessian space: + # โˆ‡โ‚‚ * compressed_kronยฒ(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹) + # This skips explicit right-compression by Mโ‚‚.๐‚โ‚‚ for this term. + kron_compressed = compressed_kronยฒ(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + rowmask = Mโ‚‚.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask, + sparse_preallocation = โ„‚.tmp_sparse_prealloc2) + + term1 = โˆ‡โ‚‚ * kron_compressed + + # Build second forcing term in compressed Hessian space with extra pruning. + # We only keep compressed-kron columns that can survive right multiplication by ฯƒcโ‚‚. + kron_sigma_compressed = compressed_kronยฒ(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, + rowmask = Mโ‚‚.โˆ‡โ‚‚_nonempty_col_as_kron_rowmask, + colmask = Mโ‚‚.๐›”๐‚โ‚‚_nonempty_row_as_kron_colmask, + sparse_preallocation = โ„‚.tmp_sparse_prealloc3) + + term2 = (โˆ‡โ‚‚ * kron_sigma_compressed) * Mโ‚‚.๐›”cโ‚‚ + + โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน = term1 + term2 C = โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu \ โˆ‡โ‚‚โŽธkโŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹โž•๐›”k๐’โ‚โ‚Šโ•ฑ๐ŸŽโŽน @@ -199,19 +353,17 @@ function calculate_second_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # @timeit_debug timer "B" begin # ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0) - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0) - B = mat_mult_kron(Mโ‚‚.๐”โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐‚โ‚‚) + Mโ‚‚.๐”โ‚‚ * Mโ‚‚.๐›” * Mโ‚‚.๐‚โ‚‚ + B = compressed_kronยฒ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, sparse_preallocation = โ„‚.tmp_sparse_prealloc1) + Mโ‚‚.๐›”cโ‚‚ # end # timeit_debug # end # timeit_debug # @timeit_debug timer "Solve sylvester equation" begin ๐’โ‚‚, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, - initial_guess = initial_guess, + initial_guess = initial_guess_sylv, sylvester_algorithm = opts.sylvester_algorithmยฒ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, + tol = opts.tol.second_order.sylvester, verbose = opts.verbose) # end # timeit_debug @@ -238,6 +390,23 @@ function calculate_second_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # end # timeit_debug # end # timeit_debug + if solved && caching + if ๐’โ‚‚ isa Matrix{S} && cache.second_order_solution isa Matrix{S} && size(cache.second_order_solution) == size(๐’โ‚‚) + copyto!(cache.second_order_solution, ๐’โ‚‚) + elseif ๐’โ‚‚ isa SparseMatrixCSC{S, Int} && cache.second_order_solution isa SparseMatrixCSC{S, Int} && + size(cache.second_order_solution) == size(๐’โ‚‚) && + cache.second_order_solution.colptr == ๐’โ‚‚.colptr && + cache.second_order_solution.rowval == ๐’โ‚‚.rowval + copyto!(cache.second_order_solution.nzval, ๐’โ‚‚.nzval) + else + cache.second_order_solution = copy(๐’โ‚‚) + end + if !isempty(parameter_values) + cache.valid_for.second_order_solution = eltype(parameter_values) <: โ„ฑ.Dual ? Float64.(โ„ฑ.value.(parameter_values)) : Float64.(parameter_values) + cache.valid_for.pruned_second_order_solution = Float64[] + end + end + return ๐’โ‚‚, solved end @@ -246,11 +415,22 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order โˆ‡โ‚‚::SparseMatrixCSC{S}, #second order derivatives โˆ‡โ‚ƒ::SparseMatrixCSC{S}, #third order derivatives ๐‘บโ‚::AbstractMatrix{S}, #first order solution - ๐’โ‚‚::SparseMatrixCSC{S}, #second order solution + ๐’โ‚‚::AbstractMatrix{S}, #second order solution (compressed) constants::constants, - workspaces::workspaces; + workspaces::workspaces, + cache::caches; initial_guess::AbstractMatrix{R} = zeros(0,0), - opts::CalculationOptions = merge_calculation_options())::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {S <: Real,R <: Real} + opts::CalculationOptions = merge_calculation_options(), + parameter_values::AbstractVector{<:Real} = Float64[], + caching::Bool = true)::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {S <: Real,R <: Real} + # Cache hit: return cached third-order solution if valid for current parameters + if caching && S === Float64 && !isempty(parameter_values) && + cache_valid_for_parameters(cache.valid_for.third_order_solution, parameter_values) + cached = cache.third_order_solution + if cached isa Matrix{S} && !isempty(cached) + return cached, true + end + end if !(eltype(workspaces.third_order.Sฬ‚) == S) workspaces.third_order = Higher_order_workspace(T = S) end @@ -259,6 +439,13 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order Mโ‚ƒ = constants.third_order T = constants.post_model_macro # @timeit_debug timer "Calculate third order solution" begin + + # Expand compressed hessian to full space + โˆ‡โ‚‚ = โˆ‡โ‚‚ * Mโ‚‚.๐”โˆ‡โ‚‚ + + # Expand compressed second-order solution to full space + ๐’โ‚‚ = sparse(๐’โ‚‚ * Mโ‚‚.๐”โ‚‚)::SparseMatrixCSC{S, Int} + # inspired by Levintal # Indices and number of variables @@ -271,14 +458,30 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order n = T.nVars nโ‚‘โ‚‹ = nโ‚‹ + 1 + nโ‚‘ + ensure_higher_order_solution_buffers!(โ„‚, n, nโ‚‘โ‚‹) + + initial_guess_sylv = if length(initial_guess) == 0 + zeros(S, 0, 0) + elseif eltype(initial_guess) <: AbstractFloat + initial_guess isa Matrix{S} ? initial_guess : Matrix{S}(initial_guess) + else + zeros(S, 0, 0) + end + # @timeit_debug timer "Setup matrices" begin # 1st order solution - ๐’โ‚ = @views [๐‘บโ‚[:,1:nโ‚‹] zeros(n) ๐‘บโ‚[:,nโ‚‹+1:end]]# |> sparse + ๐’โ‚ = โ„‚.๐’โ‚::Matrix{S} + copyto!(@view(๐’โ‚[:,1:nโ‚‹]), @view(๐‘บโ‚[:,1:nโ‚‹])) + fill!(@view(๐’โ‚[:,nโ‚‹+1]), zero(S)) + copyto!(@view(๐’โ‚[:,nโ‚‹+2:end]), @view(๐‘บโ‚[:,nโ‚‹+1:end])) - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = @views [๐’โ‚[iโ‚‹,:]; zeros(nโ‚‘ + 1, nโ‚‹) โ„’.I(nโ‚‘ + 1)[1,:] zeros(nโ‚‘ + 1, nโ‚‘)] + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„‚.๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{S} + copyto!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[1:nโ‚‹,:]), @view(๐’โ‚[iโ‚‹,:])) + fill!(@view(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1:end,:]), zero(S)) + @inbounds ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘[nโ‚‹+1,nโ‚‹+1] = one(S) - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹ = @views [(๐’โ‚ * ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)[iโ‚Š,:] ๐’โ‚ @@ -286,7 +489,7 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚[iโ‚Š,:] zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹)]# |> sparse - ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + ๐’โ‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€ = @views -โˆ‡โ‚[:,1:nโ‚Š] * ๐’โ‚[iโ‚Š,1:nโ‚‹] * โ„’.I(n)[iโ‚‹,:] - โˆ‡โ‚[:,range(1,n) .+ nโ‚Š] @@ -305,39 +508,28 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # end # timeit_debug - โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * โ„’.I(n)[iโ‚Š,:] + โˆ‡โ‚โ‚Š = @views โˆ‡โ‚[:,1:nโ‚Š] * Mโ‚‚.๐ˆโ‚™โ‚Š A = โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€lu \ โˆ‡โ‚โ‚Š # @timeit_debug timer "Setup B" begin # @timeit_debug timer "Add tmpkron" begin - tmpkron = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”) kron๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - - B = tmpkron - - # end # timeit_debug - # @timeit_debug timer "Step 1" begin - - B += Mโ‚ƒ.๐โ‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚แตฃฬƒ - - # end # timeit_debug - # @timeit_debug timer "Step 2" begin - - B += Mโ‚ƒ.๐โ‚‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚‚แตฃฬƒ - - # end # timeit_debug - # @timeit_debug timer "Mult" begin - - B *= Mโ‚ƒ.๐‚โ‚ƒ - B = choose_matrix_format(Mโ‚ƒ.๐”โ‚ƒ * B, tol = opts.tol.droptol, multithreaded = false) + # tmpkron = โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”) + # B = tmpkron + Mโ‚ƒ.๐โ‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚แตฃฬƒ + Mโ‚ƒ.๐โ‚‚โ‚—ฬ„ * tmpkron * Mโ‚ƒ.๐โ‚‚แตฃฬƒ + # B *= Mโ‚ƒ.๐‚โ‚ƒ + # B = choose_matrix_format(Mโ‚ƒ.๐”โ‚ƒ * B, tol = opts.tol.third_order.droptol, multithreaded = false) + # println("size(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) = ",size(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)) + B = compressed_permuted_mixed_kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, Mโ‚‚.๐›”, + sparse_preallocation = โ„‚.tmp_sparse_prealloc7)#, timer = timer) + # println("size(B) = ",size(B)) # end # timeit_debug # @timeit_debug timer "3rd Kronecker power" begin # B += mat_mult_kron(Mโ‚ƒ.๐”โ‚ƒ, collect(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘), collect(โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘)), Mโ‚ƒ.๐‚โ‚ƒ) # slower than direct compression - B += compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc1)#, timer = timer) + B += compressed_kronยณ(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, tol = opts.tol.third_order.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc1)#, timer = timer) # end # timeit_debug # end # timeit_debug @@ -348,7 +540,7 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order ๐’โ‚‚ zeros(nโ‚‹ + nโ‚‘, nโ‚‘โ‚‹^2)]; - โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, density_threshold = 0.0, min_length = 10, tol = opts.tol.droptol) + โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ = choose_matrix_format(โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, density_threshold = 0.0, min_length = 10, tol = opts.tol.third_order.droptol) ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = @views [๐’โ‚‚[iโ‚Š,:] zeros(nโ‚‹ + n + nโ‚‘, nโ‚‘โ‚‹^2)]; @@ -359,86 +551,69 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # end # timeit_debug # @timeit_debug timer "โˆ‡โ‚ƒ" begin - if length(โ„‚.tmpkron0) > 0 && eltype(โ„‚.tmpkron0) == S - โ„’.kron!(โ„‚.tmpkron0, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - else - โ„‚.tmpkron0 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) - end + # if length(โ„‚.tmpkron0) > 0 && eltype(โ„‚.tmpkron0) == S + # โ„’.kron!(โ„‚.tmpkron0, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + # else + # โ„‚.tmpkron0 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + # end - if length(โ„‚.tmpkron22) > 0 && eltype(โ„‚.tmpkron22) == S - โ„’.kron!(โ„‚.tmpkron22, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) - else - โ„‚.tmpkron22 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) - end + # if length(โ„‚.tmpkron22) > 0 && eltype(โ„‚.tmpkron22) == S + # โ„’.kron!(โ„‚.tmpkron22, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) + # else + # โ„‚.tmpkron22 = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„‚.tmpkron0 * Mโ‚‚.๐›”) + # end - # tmpkron = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›”) + # # tmpkron = โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) * Mโ‚‚.๐›”) - ๐”โˆ‡โ‚ƒ = โˆ‡โ‚ƒ * Mโ‚ƒ.๐”โˆ‡โ‚ƒ + # ๐”โˆ‡โ‚ƒ = โˆ‡โ‚ƒ * Mโ‚ƒ.๐”โˆ‡โ‚ƒ + + # ๐—โ‚ƒ = ๐”โˆ‡โ‚ƒ * โ„‚.tmpkron22 + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚โ‚—ฬ‚ * โ„‚.tmpkron22 * Mโ‚ƒ.๐โ‚แตฃฬƒ + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚‚โ‚—ฬ‚ * โ„‚.tmpkron22 * Mโ‚ƒ.๐โ‚‚แตฃฬƒ - ๐—โ‚ƒ = ๐”โˆ‡โ‚ƒ * โ„‚.tmpkron22 + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚โ‚—ฬ‚ * โ„‚.tmpkron22 * Mโ‚ƒ.๐โ‚แตฃฬƒ + ๐”โˆ‡โ‚ƒ * Mโ‚ƒ.๐โ‚‚โ‚—ฬ‚ * โ„‚.tmpkron22 * Mโ‚ƒ.๐โ‚‚แตฃฬƒ - # end # timeit_debug # @timeit_debug timer "โˆ‡โ‚‚ & โˆ‡โ‚โ‚Š" begin - ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ = choose_matrix_format(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) - if length(โ„‚.tmpkron1) > 0 && eltype(โ„‚.tmpkron1) == S - โ„’.kron!(โ„‚.tmpkron1, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) - else - โ„‚.tmpkron1 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) - end - - if length(โ„‚.tmpkron2) > 0 && eltype(โ„‚.tmpkron2) == S - โ„’.kron!(โ„‚.tmpkron2, Mโ‚‚.๐›”, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - else - โ„‚.tmpkron2 = โ„’.kron(Mโ‚‚.๐›”, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘) - end - - โˆ‡โ‚โ‚Š = choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold = 1.0, min_length = 10, tol = opts.tol.droptol) + โˆ‡โ‚โ‚Š = choose_matrix_format(โˆ‡โ‚โ‚Š, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ = [๐’โ‚‚[iโ‚‹,:] ; zeros(size(๐’โ‚)[2] - nโ‚‹, nโ‚‘โ‚‹^2)] - # @timeit_debug timer "Step 1" begin - out2 = โˆ‡โ‚‚ * โ„‚.tmpkron1 * โ„‚.tmpkron2 # this help - - # end # timeit_debug - # @timeit_debug timer "Step 2" begin - - # end # timeit_debug - # @timeit_debug timer "Step 3" begin - - out2 += โˆ‡โ‚‚ * โ„‚.tmpkron1 * Mโ‚ƒ.๐โ‚โ‚— * โ„‚.tmpkron2 * Mโ‚ƒ.๐โ‚แตฃ# |> findnz - - # end # timeit_debug - # @timeit_debug timer "Step 4" begin + # Terms (a)+(b): โˆ‡โ‚‚ * kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) * [tmpkron2 + ๐โ‚โ‚— * tmpkron2 * ๐โ‚แตฃ] * ๐๐‚โ‚ƒ + # Compute D_ab to avoid materializing kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ) + tmpkron2_sp = โ„’.kron(Mโ‚‚.๐›”, choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.third_order.droptol)) + D_ab = (tmpkron2_sp + Mโ‚ƒ.๐โ‚โ‚— * tmpkron2_sp * Mโ‚ƒ.๐โ‚แตฃ) * Mโ‚ƒ.๐๐‚โ‚ƒ - out2 += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc2)# |> findnz + ๐—โ‚ƒ = mat_mult_kron(โˆ‡โ‚‚, collect(๐’โ‚โ‚Šโ•ฑ๐ŸŽ), collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ), D_ab, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc2) - # out2 += โˆ‡โ‚‚ * โ„’.kron(โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›”)# |> findnz - out2 += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›”), sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc3)# |> findnz + # Term (c): โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, โŽธ๐’โ‚‚k..โŽน) * ๐๐‚โ‚ƒ + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, โŽธ๐’โ‚‚k๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โž•๐’โ‚๐’โ‚‚โ‚‹โŽนโ•ฑ๐’โ‚‚โ•ฑ๐ŸŽ, Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc3) - # end # timeit_debug - # @timeit_debug timer "Step 5" begin - # out2 += โˆ‡โ‚โ‚Š * mat_mult_kron(๐’โ‚‚, collect(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘), collect(๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ)) - # out2 += mat_mult_kron(โˆ‡โ‚โ‚Š * ๐’โ‚‚, collect(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘), collect(๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ)) - # out2 += โˆ‡โ‚โ‚Š * ๐’โ‚‚ * โ„’.kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) + # Term (d): โˆ‡โ‚‚ * kron(โŽธ๐’โ‚..โŽน, ๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ*๐›”) * ๐๐‚โ‚ƒ + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚‚, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, collect(๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ * Mโ‚‚.๐›”), Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) - ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.droptol) - out2 += โˆ‡โ‚โ‚Š * mat_mult_kron(๐’โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, sparse = true, sparse_preallocation = โ„‚.tmp_sparse_prealloc4) + # Term (e): โˆ‡โ‚โ‚Š * ๐’โ‚‚ * kron(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ) * ๐๐‚โ‚ƒ + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘ = choose_matrix_format(๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, density_threshold = 0.0, tol = opts.tol.third_order.droptol) - # end # timeit_debug - # @timeit_debug timer "Mult" begin - # โ„’.mul!(๐—โ‚ƒ, out2, Mโ‚ƒ.๐, 1, 1) # less memory but way slower; .+= also more memory and slower - ๐—โ‚ƒ += out2 * Mโ‚ƒ.๐ + ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚โ‚Š * ๐’โ‚‚, ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘, ๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ, Mโ‚ƒ.๐๐‚โ‚ƒ, sparse = true) - ๐—โ‚ƒ *= Mโ‚ƒ.๐‚โ‚ƒ + if length(โ„‚.tmpkron0) > 0 && eltype(โ„‚.tmpkron0) == S + โ„’.kron!(โ„‚.tmpkron0, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + else + โ„‚.tmpkron0 = โ„’.kron(๐’โ‚โ‚Šโ•ฑ๐ŸŽ, ๐’โ‚โ‚Šโ•ฑ๐ŸŽ) + end + + โ„‚.tmpkron0 *= Mโ‚‚.๐›” + # โ„’.rmul!(โ„‚.tmpkron0, Mโ‚‚.๐›”) + ๐—โ‚ƒ += mul_compressed_permuted_mixed_kron(โˆ‡โ‚ƒ, โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹, + โ„‚.tmpkron0, + sparse_preallocation = โ„‚.tmp_sparse_prealloc6) # end # timeit_debug # end # timeit_debug # @timeit_debug timer "3rd Kronecker power" begin # ๐—โ‚ƒ += mat_mult_kron(โˆ‡โ‚ƒ, collect(aux), collect(โ„’.kron(aux, aux)), Mโ‚ƒ.๐‚โ‚ƒ) # slower than direct compression - ๐—โ‚ƒ += โˆ‡โ‚ƒ * compressed_kronยณ(aux, rowmask = unique(findnz(โˆ‡โ‚ƒ)[2]), tol = opts.tol.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc5) #, timer = timer) + ๐—โ‚ƒ += mul_compressed_kronยณ(โˆ‡โ‚ƒ, aux, tol = opts.tol.third_order.droptol, sparse_preallocation = โ„‚.tmp_sparse_prealloc5) #, timer = timer) # end # timeit_debug # @timeit_debug timer "Mult 2" begin @@ -450,10 +625,9 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # @timeit_debug timer "Solve sylvester equation" begin ๐’โ‚ƒ, solved = solve_sylvester_equation(A, B, C, โ„‚.sylvester_workspace, - initial_guess = initial_guess, + initial_guess = initial_guess_sylv, sylvester_algorithm = opts.sylvester_algorithmยณ, - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, + tol = opts.tol.third_order.sylvester, verbose = opts.verbose) # end # timeit_debug @@ -477,11 +651,28 @@ function calculate_third_order_solution(โˆ‡โ‚::AbstractMatrix{S}, #first order # ๐’โ‚ƒ *= Mโ‚ƒ.๐”โ‚ƒ - ๐’โ‚ƒ = choose_matrix_format(๐’โ‚ƒ, multithreaded = false, tol = opts.tol.droptol) + ๐’โ‚ƒ = choose_matrix_format(๐’โ‚ƒ, multithreaded = false, tol = opts.tol.third_order.droptol) # end # timeit_debug # end # timeit_debug + if solved && caching + if ๐’โ‚ƒ isa Matrix{S} && cache.third_order_solution isa Matrix{S} && size(cache.third_order_solution) == size(๐’โ‚ƒ) + copyto!(cache.third_order_solution, ๐’โ‚ƒ) + elseif ๐’โ‚ƒ isa SparseMatrixCSC{S, Int} && cache.third_order_solution isa SparseMatrixCSC{S, Int} && + size(cache.third_order_solution) == size(๐’โ‚ƒ) && + cache.third_order_solution.colptr == ๐’โ‚ƒ.colptr && + cache.third_order_solution.rowval == ๐’โ‚ƒ.rowval + copyto!(cache.third_order_solution.nzval, ๐’โ‚ƒ.nzval) + else + cache.third_order_solution = copy(๐’โ‚ƒ) + end + if !isempty(parameter_values) + cache.valid_for.third_order_solution = eltype(parameter_values) <: โ„ฑ.Dual ? Float64.(โ„ฑ.value.(parameter_values)) : Float64.(parameter_values) + cache.valid_for.pruned_third_order_solution = Float64[] + end + end + return ๐’โ‚ƒ, solved end diff --git a/src/structures.jl b/src/structures.jl index 62659cc92..3b0f69c08 100644 --- a/src/structures.jl +++ b/src/structures.jl @@ -74,7 +74,8 @@ # # 2. WORKSPACES (๐“‚.workspaces) - Pre-allocated temporary buffers that are # reused across function calls to avoid repeated allocations: -# - qme: Quadratic matrix equation solver workspace +# - first_order: First-order perturbation solver workspace +# - qme_doubling: Quadratic matrix equation doubling solver workspace # - sylvester_*: Sylvester equation solver workspaces # - lyapunov_*: Lyapunov equation solver workspaces # - second_order/third_order: Higher order perturbation workspaces @@ -86,10 +87,10 @@ # - non_stochastic_steady_state: NSSS solution values # - jacobian/hessian/third_order_derivatives: Perturbation derivatives # - first_order_solution_matrix/second_order_solution/etc.: Solved policy matrices -# - outdated: Flags indicating which caches need recomputation +# - valid_for: Parameter vectors for which each cache entry is valid # # 4. FUNCTIONS (๐“‚.functions) - Compiled model functions: -# - NSSS_solve/check: Steady state solvers +# - NSSS_check + step-based NSSS solver functions # - jacobian/hessian/third_order_derivatives: Derivative functions # - state_update functions: Policy function evaluators # @@ -97,7 +98,7 @@ # @model macro โ†’ post_model_macro (constants) # @parameters macro โ†’ post_parameters_macro, post_complete_parameters (constants) # solve!() โ†’ populates caches using workspaces, guided by constants -# get_irf/simulate/etc โ†’ reads from caches, may trigger solve!() if outdated +# get_irf/simulate/etc โ†’ reads from caches, may trigger solve!() if not valid_for current parameters # # ============================================================================= @@ -146,6 +147,7 @@ struct post_model_macro nMixed::Int nFuture_not_past_and_mixed::Int nPast_not_future_and_mixed::Int + I_nPast::โ„’.Diagonal{Bool, Vector{Bool}} # nPresent_but_not_only::Int nVars::Int nExo::Int @@ -239,6 +241,11 @@ struct moments_substate_indices e_ss::SparseMatrixCSC{Float64, Int} ss_s::SparseMatrixCSC{Float64, Int} s_s::SparseMatrixCSC{Float64, Int} + # Duplication/elimination matrices for symmetric Kronecker compression + Dโ‚‚หข::SparseMatrixCSC{Float64, Int} # nหขยฒ ร— nหข(nหข+1)/2 duplication + Lโ‚‚หข::SparseMatrixCSC{Float64, Int} # nหข(nหข+1)/2 ร— nหขยฒ elimination + Dโ‚ƒหข::SparseMatrixCSC{Float64, Int} # nหขยณ ร— nหข(nหข+1)(nหข+2)/6 duplication + Lโ‚ƒหข::SparseMatrixCSC{Float64, Int} # nหข(nหข+1)(nหข+2)/6 ร— nหขยณ elimination end struct moments_dependency_kron_indices @@ -253,7 +260,7 @@ Second-order perturbation auxiliary matrices and index caches. These are computed once when the model structure is known and reused across solutions. Contains three categories of data: -1. **Auxiliary matrices** (๐›”, ๐‚โ‚‚, ๐”โ‚‚, ๐”โˆ‡โ‚‚): Sparse integer matrices for second-order +1. **Auxiliary matrices** (๐›”, ๐‚โ‚‚, ๐”โ‚‚, ๐”โˆ‡โ‚‚): Sparse matrices for second-order perturbation solution. Populated by `create_second_order_auxiliary_matrices` during `write_functions_mapping!`. @@ -271,9 +278,21 @@ mutable struct second_order_indices # Triggered by: write_functions_mapping! โ† solve! # ========================================================================= ๐›”::SparseMatrixCSC{Int} # Commutation matrix + ๐›”_sym::SparseMatrixCSC{Int} # Symmetrised volatility: ๐›” + P_swap * ๐›” * P_swap + ๐›”cโ‚‚::SparseMatrixCSC{Int} # Compressed volatility: ๐”โ‚‚ * ๐›” * ๐‚โ‚‚ + ๐›”๐‚โ‚‚::SparseMatrixCSC{Int} # Product ๐›” * ๐‚โ‚‚ (precomputed) ๐‚โ‚‚::SparseMatrixCSC{Int} # Duplication matrix for 2nd order ๐”โ‚‚::SparseMatrixCSC{Int} # Unique elements selector for 2nd order ๐”โˆ‡โ‚‚::SparseMatrixCSC{Int} # Gradient unique elements selector + ๐ˆโ‚™โ‚Š::SparseMatrixCSC{Int} # Future-state row selector from I(nVars) + ๐ˆโ‚™โ‚‹::SparseMatrixCSC{Int} # Past-state row selector from I(nVars) + โˆ‡โ‚‚_nonempty_col_as_kron_rowmask::Vector{Int} # Non-empty columns of โˆ‡โ‚‚, mapped to rowmask in compressed_kronยฒ + ๐›”๐‚โ‚‚_nonempty_row_as_kron_colmask::Vector{Int} # Non-empty rows of ฯƒcโ‚‚, mapped to colmask in compressed_kronยฒ + # Pre-transposed constants for rrule pullback + ๐›”แต€::SparseMatrixCSC{Int} # ๐›”' + ๐‚โ‚‚แต€::SparseMatrixCSC{Int} # ๐‚โ‚‚' + ๐”โ‚‚แต€::SparseMatrixCSC{Int} # ๐”โ‚‚' + ๐”โˆ‡โ‚‚แต€::SparseMatrixCSC{Int} # ๐”โˆ‡โ‚‚' # ========================================================================= # COMPUTATIONAL CONSTANTS (for efficient sparse operations) @@ -351,7 +370,12 @@ mutable struct third_order_indices ๐ˆโ‚ƒ::Dict{Vector{Int}, Int} # Index mapping for 3rd order terms ๐‚โˆ‡โ‚ƒ::SparseMatrixCSC{Int} # Gradient duplication matrix ๐”โˆ‡โ‚ƒ::SparseMatrixCSC{Int} # Gradient unique selector + โˆ‡โ‚ƒ_rowmask::Vector{Int} # Structural nonzero compressed gradient columns ๐::SparseMatrixCSC{Int} # Permutation matrix + + + + ๐๐‚โ‚ƒ::SparseMatrixCSC{Int} # Cached product ๐ * ๐‚โ‚ƒ ๐โ‚โ‚—::SparseMatrixCSC{Int} # Left permutation 1 ๐โ‚แตฃ::SparseMatrixCSC{Int} # Right permutation 1 ๐โ‚โ‚—ฬ‚::SparseMatrixCSC{Int} # Modified left permutation 1 @@ -362,6 +386,17 @@ mutable struct third_order_indices ๐โ‚‚แตฃฬƒ::SparseMatrixCSC{Int} # Alternative right permutation 2 ๐’๐::SparseMatrixCSC{Int} # Combined selection-permutation + # Pre-transposed constants (computed once, reused by rrule pullback) + ๐‚โ‚ƒแต€::SparseMatrixCSC{Int} # ๐‚โ‚ƒ' + ๐”โ‚ƒแต€::SparseMatrixCSC{Int} # ๐”โ‚ƒ' + ๐๐‚โ‚ƒแต€::SparseMatrixCSC{Int} # ๐๐‚โ‚ƒ' + ๐โ‚โ‚—แต€::SparseMatrixCSC{Int} # ๐โ‚โ‚—' + ๐โ‚แตฃแต€::SparseMatrixCSC{Int} # ๐โ‚แตฃ' + ๐โ‚โ‚—ฬ„แต€::SparseMatrixCSC{Int} # ๐โ‚โ‚—ฬ„' + ๐โ‚‚โ‚—ฬ„แต€::SparseMatrixCSC{Int} # ๐โ‚‚โ‚—ฬ„' + ๐โ‚แตฃฬƒแต€::SparseMatrixCSC{Int} # ๐โ‚แตฃฬƒ' + ๐โ‚‚แตฃฬƒแต€::SparseMatrixCSC{Int} # ๐โ‚‚แตฃฬƒ' + # ========================================================================= # CONDITIONAL FORECAST CONSTANTS # Filled by ensure_conditional_forecast_constants! (options_and_caches.jl) @@ -472,7 +507,10 @@ mutable struct sylvester_workspace{G <: AbstractFloat, H <: Real} ๐‚B::Matrix{G} # nร—m temporary for C*B multiplication # Krylov solver state (lazily allocated) - krylov_workspace::krylov_workspace{G} + krylov::krylov_workspace{G} + + # Stable primal solution cache for AD/rrule pullbacks + P::Matrix{G} # ForwardDiff partials buffers (for forward-mode AD) Pฬƒ::Matrix{H} # For sylvester equation partials @@ -483,25 +521,70 @@ end """ -Pre-allocated workspace matrices for the quadratic matrix equation doubling algorithm. -All matrices are square with dimension n = size(A,1) = size(B,1) = size(C,1). +Pre-allocated workspace matrices for first-order perturbation and related AD paths. + +Contains temporary matrices and factorization workspaces reused by +`calculate_first_order_solution` and first-order derivative routines. +""" +mutable struct first_order_workspace{T <: Real, R <: Real} + # Sylvester workspace for ForwardDiff path + sylvester::sylvester_workspace{T, R} + + # ForwardDiff partials buffers (for forward-mode AD) + Xฬƒ_first_order::Matrix{R} # For first order solution partials + p_tmp::Matrix{R} # For calculate_first_order_solution + โˆ‚SS_and_pars::Matrix{R} # For NSSS partials in get_NSSS_and_parameters + โˆ‚โˆ‡โ‚_vec::Vector{T} # Flattened cotangent buffer for calculate_jacobian pullback + + # First-order perturbation workspaces (primal) + ๐งโ‚šโ‚‹::Matrix{T} # nโ‚šโ‚‹ = Aโ‚Šแตค * D + ๐Œ::Matrix{T} # M = A_future * expand_past + ๐€โ‚Š::Matrix{T} # Aโ‚Š + ๐€โ‚€::Matrix{T} # Aโ‚€ + ๐€โ‚‹::Matrix{T} # Aโ‚‹ + ๐€ฬƒโ‚Š::Matrix{T} # Aฬƒโ‚Š + ๐€ฬƒโ‚€::Matrix{T} # Aฬƒโ‚€ + ๐€ฬƒโ‚‹::Matrix{T} # Aฬƒโ‚‹ + ๐€ฬ„โ‚€แตค::Matrix{T} # Aฬ„โ‚€แตค + ๐€โ‚Šแตค::Matrix{T} # Aโ‚Šแตค + ๐€ฬƒโ‚€แตค::Matrix{T} # Aฬƒโ‚€แตค + ๐€โ‚‹แตค::Matrix{T} # Aโ‚‹แตค + ๐€::Matrix{T} # A + โˆ‡โ‚€::Matrix{T} # copy of โˆ‡โ‚€ block (mutable workspace buffer) + โˆ‡โ‚‘::Matrix{T} # copy of โˆ‡โ‚‘ block (mutable workspace buffer) + + # FastLapackInterface QR workspaces for first-order solution + fast_qr_factors::Matrix{T} + fast_qr_ws::FastLapackInterface.QRWs{T} + fast_qr_orm_ws_plus::FastLapackInterface.QROrmWs{T} + fast_qr_orm_dims_plus::NTuple{3, Int} + fast_qr_orm_ws_zero::FastLapackInterface.QROrmWs{T} + fast_qr_orm_dims_zero::NTuple{3, Int} + fast_qr_orm_ws_minus::FastLapackInterface.QROrmWs{T} + fast_qr_orm_dims_minus::NTuple{3, Int} + + # FastLapackInterface LU workspaces for first-order solve + fast_lu_ws_a0u::FastLapackInterface.LUWs + fast_lu_dims_a0u::NTuple{2, Int} + fast_lu_ws_nabla0::FastLapackInterface.LUWs + fast_lu_dims_nabla0::NTuple{2, Int} + + # Dedicated FastLapackInterface LU workspace for NSSS implicit derivatives + fast_lu_ws_nsss::FastLapackInterface.LUWs + fast_lu_dims_nsss::NTuple{2, Int} + nsss_sparse_lu_buffer::๐’ฎ.LinearCache + nsss_sparse_rhs::Vector{T} + nsss_jvp_rhs::Matrix{T} +end -Used by `solve_quadratic_matrix_equation` with `Val{:doubling}` in quadratic_matrix_equation.jl. -Also used by stochastic steady state calculations in `calculate_second_order_stochastic_steady_state` -and `calculate_third_order_stochastic_steady_state`. -Avoids per-call allocations for temporary matrices in the iterative doubling algorithm. -Fields: -- `E`, `F`: Working matrices for the doubling recurrence -- `X`, `Y`: Current iteration solution matrices -- `X_new`, `Y_new`, `E_new`, `F_new`: Next iteration matrices -- `temp1`, `temp2`, `temp3`: Temporary matrices for intermediate computations -- `Bฬ„`: Copy of B for LU factorization (modified in-place) -- `AXX`: Temporary for residual computation (A * Xยฒ + B * X + C) -- `I_n`: Pre-computed identity matrix for QME doubling (UniformScaling) -- `I_nPast`: Pre-computed identity matrix for stochastic steady state (UniformScaling) """ -mutable struct qme_workspace{T <: Real, R <: Real} +Pre-allocated workspace matrices for quadratic matrix equation doubling and dual QME differentiation. + +All matrices are square with dimension n = size(A,1) = size(B,1) = size(C,1). +Used by `solve_quadratic_matrix_equation` with `Val{:doubling}`. +""" +mutable struct qme_doubling_workspace{T <: Real, R <: Real} # Doubling algorithm working matrices E::Matrix{T} F::Matrix{T} @@ -511,30 +594,78 @@ mutable struct qme_workspace{T <: Real, R <: Real} Y_new::Matrix{T} E_new::Matrix{T} F_new::Matrix{T} - + # Temporary matrices for intermediate operations temp1::Matrix{T} temp2::Matrix{T} temp3::Matrix{T} - - # LU factorization buffer + + # LU factorization and residual buffers Bฬ„::Matrix{T} - - # Residual computation buffer AXX::Matrix{T} - + # Sylvester workspace for ForwardDiff path - sylvester_ws::sylvester_workspace{T, R} - + sylvester::sylvester_workspace{T, R} + # ForwardDiff partials buffers (for forward-mode AD) Xฬƒ::Matrix{R} # For QME solution partials - Xฬƒ_first_order::Matrix{R} # For first order solution partials - p_tmp::Matrix{R} # For calculate_first_order_solution - โˆ‚SS_and_pars::Matrix{R} # For NSSS partials in get_NSSS_and_parameters - - # Pre-computed identity matrices (Diagonal{Bool} - supports indexing for schur algorithm) - I_n::โ„’.Diagonal{Bool, Vector{Bool}} # Identity for QME doubling (dimension n = nVars - nPresent_only) - I_nPast::โ„’.Diagonal{Bool, Vector{Bool}} # Identity for schur & stochastic steady state (dimension nPast_not_future_and_mixed) + + # FastLapackInterface LU workspaces for QME doubling solve + fast_lu_ws_qme_a::FastLapackInterface.LUWs + fast_lu_dims_qme_a::NTuple{2, Int} + fast_lu_ws_qme_b::FastLapackInterface.LUWs + fast_lu_dims_qme_b::NTuple{2, Int} +end + + +""" +Pre-allocated workspace matrices for the schur-based quadratic matrix equation solver. + +The schur method solves A*Xยฒ + B*X + C = 0 by forming a companion linearization +and computing its generalized Schur decomposition. All temporary matrices are +pre-allocated here to avoid per-call allocations. + +Fields: +- `D`, `E`: Companion form matrices (n+nMixed) ร— (nPfm+nFnpm), overwritten by schur! +- `รƒโ‚‹`, `รƒโ‚€โ‚Š`: Negated slices from C and B (need owned copies for rmul!) +- `รƒโ‚€โ‚‹`: Product B[:,indices_past_not_future_in_comb] * I_nPast[not_mixed_in_past_idx,:] +- `Zโ‚‚โ‚`, `Sโ‚โ‚`, `Tโ‚โ‚`: Schur decomposition result blocks (need owned copies for lu!) +- `sol`: Assembled solution before reordering (nPfm+nFnpm) ร— nPfm +- `temp_X2`: Buffer for Xยฒ in residual check +- `AXX`: Buffer for A*Xยฒ + B*X + C residual +- `eigenselect`: Boolean vector for eigenvalue selection +""" +mutable struct schur_workspace{T <: Real} + # Companion form matrices (overwritten by schur!) + D::Matrix{T} + E::Matrix{T} + # Slices that need negation (owned copies) + รƒโ‚‹::Matrix{T} + รƒโ‚€โ‚Š::Matrix{T} + รƒโ‚€โ‚‹::Matrix{T} + # Schur decomposition result blocks (owned copies for lu!) + Zโ‚โ‚::Matrix{T} + Zโ‚‚โ‚::Matrix{T} + Sโ‚โ‚::Matrix{T} + Tโ‚โ‚::Matrix{T} + # Solution assembly buffers + sol::Matrix{T} + # Residual check buffers + temp_X2::Matrix{T} + AXX::Matrix{T} + # Eigenvalue selection + eigenselect::Vector{Bool} + # FastLapack generalized Schur workspace + fast_qz_ws::FastLapackInterface.GeneralizedSchurWs{T} + fast_qz_dims::NTuple{2, Int} + # FastLapack LU workspaces for schur post-processing + fast_lu_ws_z11::FastLapackInterface.LUWs + fast_lu_dims_z11::NTuple{2, Int} + fast_lu_ws_s11::FastLapackInterface.LUWs + fast_lu_dims_s11::NTuple{2, Int} + # Scratch buffers for right-side solves (store transposed RHS) + fast_lu_rhs_t_z21::Matrix{T} + fast_lu_rhs_t_s11::Matrix{T} end @@ -577,10 +708,16 @@ mutable struct lyapunov_workspace{T <: Real, R <: Real} b::Vector{T} # Krylov solver state (lazily allocated, can be reused across calls) - bicgstab_workspace::Krylov.BicgstabWorkspace{T, T, Vector{T}} - gmres_workspace::Krylov.GmresWorkspace{T, T, Vector{T}} + bicgstab::Krylov.BicgstabWorkspace{T, T, Vector{T}} + gmres::Krylov.GmresWorkspace{T, T, Vector{T}} + + # vech-space Krylov buffers (for symmetric C, dimension n(n+1)/2) + b_vech::Vector{T} + bicgstab_vech::Krylov.BicgstabWorkspace{T, T, Vector{T}} + gmres_vech::Krylov.GmresWorkspace{T, T, Vector{T}} # ForwardDiff partials buffers (for forward-mode AD) + P::Matrix{T} # Stable primal solution cache for AD/rrule pullbacks Pฬƒ::Matrix{R} # For lyapunov equation partials Aฬƒ_fd::Matrix{R} # Temporary for ForwardDiff partials of A Cฬƒ_fd::Matrix{R} # Temporary for ForwardDiff partials of C @@ -592,83 +729,213 @@ struct ss_solve_block extended_ss_problem::function_and_jacobian end -mutable struct non_stochastic_steady_state - solve_blocks_in_place::Vector{ss_solve_block} - dependencies::Any + +# ============================================================================ +# NSSS Solver Pipeline โ€” struct-of-arrays design +# +# Steps are stored as parallel vectors of per-step data, with shared +# workspaces for scratch buffers and separated caches for past results. +# +# Step types are encoded as UInt8 flags: +const ANALYTICAL_STEP = 0x01 +const NUMERICAL_STEP = 0x02 +# ============================================================================ + +""" +Per-step compiled functions, stored as parallel vectors indexed by step number. + +Each step has an optional `aux_func!` (pre-step domain-safety computation), +an optional `error_func!` (domain-safety error check), and a main function +which is either `eval_func!` (analytical) or dispatched via `solve_block` (numerical). +""" +struct NSSSSolverFunctions + # Per-step compiled functions (indexed by step number) + aux_funcs::Vector{Function} # f!(out, sol_vec, params_vec) โ€” optional pre-step aux + error_funcs::Vector{Function} # g!(out, sol_vec, params_vec) โ€” optional error check + eval_funcs::Vector{Function} # f!(out, sol_vec, params_vec) โ€” main eval (analytical only) + solve_blocks::Vector{Union{Nothing, ss_solve_block}} # compiled residual/Jacobian (numerical only) end + """ -Tracks which cache elements are outdated and need recalculation. +Per-step immutable configuration: indices, bounds, and metadata. -When parameters change (via `๐“‚.parameter_values = ...`), all fields are set to `true` (outdated). -When a cache is computed (e.g., by `solve!()`), its corresponding field is set to `false` (up to date). +Index arrays are stored in flat contiguous vectors, with per-step `UnitRange{Int}` +providing zero-copy views into the flat storage. This reduces heap allocations +and improves cache locality compared to per-step `Vector{Int}` fields. +""" +struct NSSSSolverConstants + # Step metadata + n_steps::Int + n_ext_params::Int + step_types::Vector{UInt8} # ANALYTICAL_STEP or NUMERICAL_STEP per step + descriptions::Vector{String} # debug description per step + block_indices::Vector{Int} # numerical block index (0 for analytical) + + # Flat index arrays + per-step ranges + write_indices::Vector{Int} # flat: which sol_vec positions to write + write_ranges::Vector{UnitRange{Int}} # per-step range into write_indices + aux_write_indices::Vector{Int} # flat: aux write positions + aux_write_ranges::Vector{UnitRange{Int}} # per-step range into aux_write_indices + param_gather_indices::Vector{Int} # flat: numerical param gather (0-length for analytical) + param_gather_ranges::Vector{UnitRange{Int}} # per-step range + var_gather_indices::Vector{Int} # flat: numerical var gather (0-length for analytical) + var_gather_ranges::Vector{UnitRange{Int}} # per-step range + + # Flat bounds arrays + per-step ranges (analytical bounds for clamping) + lower_bounds::Vector{Float64} + upper_bounds::Vector{Float64} + has_bounds::BitVector + bounds_ranges::Vector{UnitRange{Int}} # per-step range into lower/upper/has_bounds + + # Flat bounds arrays for numerical block solver + numerical_lbs::Vector{Float64} + numerical_ubs::Vector{Float64} + numerical_bounds_ranges::Vector{UnitRange{Int}} # per-step range into numerical_lbs/ubs + + # Flat error buffer sizing per step + error_sizes::Vector{Int} # size of error output for each step + aux_error_sizes::Vector{Int} # size of aux error output (numerical steps) +end + + +""" +Shared scratch buffers reused across all steps during a single solve pass. -This enables lazy evaluation: caches are only recomputed when actually needed AND outdated. +All buffers are pre-allocated to the maximum size needed across all steps, +avoiding per-step allocation. Steps use `@view` slices into these buffers. """ -mutable struct outdated_caches - # Non-stochastic steady state - non_stochastic_steady_state::Bool - # Perturbation derivative buffers - jacobian::Bool - hessian::Bool - third_order_derivatives::Bool - # Perturbation solution buffers - first_order_solution::Bool - second_order_solution::Bool - pruned_second_order_solution::Bool - third_order_solution::Bool - pruned_third_order_solution::Bool +mutable struct NSSSSolverWorkspace + main_buffer::Vector{Float64} # for eval_func! output or params_and_solved_vars gather + aux_buffer::Vector{Float64} # for aux_func! output + error_buffer::Vector{Float64} # for error_func! / aux_error_func! output + params_vec_buffer::Vector{Float64} # extended parameter vector (bounded + calibration_no_var) + sol_vec_buffer::Vector{Float64} # solution vector across NSSS steps + output_buffer::Vector{Float64} # returned NSSS output (subset view materialized into reusable buffer) + guess_buffer::Vector{Float64} # for initial_guess in numerical steps + inits::Vector{Vector{Float64}} # 2-element container: [clamped_guess, cached_params] + params_and_solved_vars_buffer::Vector{Float64} # gathered block inputs (params + solved vars) + lbs_buffer::Vector{Float64} # numerical lower bounds for current block + ubs_buffer::Vector{Float64} # numerical upper bounds for current block + scaled_parameters_buffer::Vector{Float64} # continuation interpolation scratch + continuation::CircularBuffer{Vector{Vector{Float64}}} # continuation warm-start cache + continuation_capacity::Int + check_residual::Vector{Float64} # for NSSS_check in get_NSSS_and_parameters (n_equations + n_calibration) end +"""Construct an empty `NSSSSolverFunctions` with no steps.""" +NSSSSolverFunctions() = NSSSSolverFunctions( + Function[], + Function[], + Function[], + Union{Nothing,ss_solve_block}[], +) + +"""Construct an empty `NSSSSolverConstants` with no steps.""" +NSSSSolverConstants() = NSSSSolverConstants( + 0, + 0, + UInt8[], String[], Int[], + Int[], UnitRange{Int}[], + Int[], UnitRange{Int}[], + Int[], UnitRange{Int}[], + Int[], UnitRange{Int}[], + Float64[], Float64[], BitVector(), UnitRange{Int}[], + Float64[], Float64[], UnitRange{Int}[], + Int[], Int[], +) + +"""Construct an empty `NSSSSolverWorkspace` with no buffers.""" +NSSSSolverWorkspace() = NSSSSolverWorkspace( + Float64[], Float64[], Float64[], Float64[], Float64[], Float64[], Float64[], + [Float64[], Float64[Inf]], + Float64[], Float64[], Float64[], + Float64[], CircularBuffer{Vector{Vector{Float64}}}(1), 1, + Float64[], +) + +mutable struct valid_for_caches + non_stochastic_steady_state::Vector{Float64} + jacobian::Vector{Float64} + hessian::Vector{Float64} + third_order_derivatives::Vector{Float64} + first_order_solution::Vector{Float64} + first_order_obc_solution::Vector{Float64} + second_order_solution::Vector{Float64} + pruned_second_order_solution::Vector{Float64} + second_order_stochastic_steady_state::Vector{Float64} + pruned_second_order_stochastic_steady_state::Vector{Float64} + third_order_solution::Vector{Float64} + pruned_third_order_solution::Vector{Float64} + third_order_stochastic_steady_state::Vector{Float64} + pruned_third_order_stochastic_steady_state::Vector{Float64} +end + + +valid_for_caches() = valid_for_caches( + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], +) + + """ Stored computation results that can be reused across function calls. Caches store the final outputs of expensive computations (steady state, perturbation solutions). -They are invalidated when parameters change (tracked by `outdated` flags) and recomputed -lazily when needed by get_* functions. +Each cache is reused only when marked valid for the active parameter vector in `valid_for`. Purpose: Avoid recomputation when the same result is needed multiple times. Fields: -- `outdated`: Flags indicating which caches need recomputation (see [`outdated_caches`](@ref)) +- `valid_for`: Parameter vectors for which each cache entry is valid - Perturbation derivatives (`jacobian`, `hessian`, `third_order_derivatives`): Model derivative matrices evaluated at steady state - Perturbation solutions (`first_order_solution_matrix`, `second_order_solution`, etc.): Policy function coefficient matrices - `non_stochastic_steady_state`: NSSS solution values -- `solver_cache`: Recent solver guesses for warm-starting +- `solver`: Recent solver guesses for warm-starting Relationship to other structs: - Caches are computed using `constants` (for dimensions/structure) and `workspaces` (for temporary buffers) - Caches are read by get_* functions (get_irf, simulate, etc.) -- Caches are invalidated when `parameter_values` changes +- Caches are reused only when `valid_for` matches current `parameter_values` """ mutable struct caches - # ========================================================================= - # CACHE INVALIDATION FLAGS - # ========================================================================= - outdated::outdated_caches + valid_for::valid_for_caches # ========================================================================= # PERTURBATION DERIVATIVE CACHES # Computed by model derivative functions, used by perturbation solvers # ========================================================================= jacobian::AbstractMatrix{<: Real} # โˆ‡f at SS - jacobian_parameters::AbstractMatrix{<: Real} # โˆ‚โˆ‡f/โˆ‚ฮธ - jacobian_SS_and_pars::AbstractMatrix{<: Real} # โˆ‚โˆ‡f/โˆ‚(SS,ฮธ) + jacobian_parameters::AbstractMatrix{<: Real} # โˆ‚โˆ‡f/โˆ‚ฮธ, stored as (targets ร— vec(โˆ‡f)) + jacobian_SS_and_pars::AbstractMatrix{<: Real} # โˆ‚โˆ‡f/โˆ‚(SS,ฮธ), stored as (targets ร— vec(โˆ‡f)) hessian::AbstractMatrix{<: Real} # โˆ‡ยฒf at SS - hessian_parameters::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยฒf/โˆ‚ฮธ - hessian_SS_and_pars::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยฒf/โˆ‚(SS,ฮธ) + hessian_parameters::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยฒf/โˆ‚ฮธ, stored as (targets ร— vec(โˆ‡ยฒf)) + hessian_SS_and_pars::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยฒf/โˆ‚(SS,ฮธ), stored as (targets ร— vec(โˆ‡ยฒf)) third_order_derivatives::AbstractMatrix{<: Real} # โˆ‡ยณf at SS - third_order_derivatives_parameters::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยณf/โˆ‚ฮธ - third_order_derivatives_SS_and_pars::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยณf/โˆ‚(SS,ฮธ) + third_order_derivatives_parameters::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยณf/โˆ‚ฮธ, stored as (targets ร— vec(โˆ‡ยณf)) + third_order_derivatives_SS_and_pars::AbstractMatrix{<: Real} # โˆ‚โˆ‡ยณf/โˆ‚(SS,ฮธ), stored as (targets ร— vec(โˆ‡ยณf)) # ========================================================================= # PERTURBATION SOLUTION CACHES # Policy function coefficient matrices (๐’โ‚, ๐’โ‚‚, ๐’โ‚ƒ) # ========================================================================= first_order_solution_matrix::Matrix{<: Real} # ๐’โ‚ - first order policy + first_order_obc_solution_matrix::Matrix{<: Real} # ลœโ‚ - first order OBC policy qme_solution::Matrix{<: Real} # Quadratic matrix eqn solution second_order_stochastic_steady_state::Vector{<: Real} # E[x] deviation from NSSS (2nd) second_order_solution::AbstractMatrix{<: Real} # ๐’โ‚‚ - second order policy @@ -681,9 +948,9 @@ mutable struct caches # STEADY STATE CACHES # ========================================================================= non_stochastic_steady_state::Vector{<: Real} # NSSS values - solver_cache::CircularBuffer{Vector{Vector{Float64}}} # Recent solver guesses - โˆ‚equations_โˆ‚parameters::AbstractMatrix{<: Real} # SS sensitivity to params - โˆ‚equations_โˆ‚SS_and_pars::AbstractMatrix{<: Real} # SS Jacobian + solver::CircularBuffer{Vector{Vector{Float64}}} # Recent solver guesses + NSSS_โˆ‚equations_โˆ‚parameters::AbstractMatrix{<: Real} # Dedicated NSSS SS sensitivity + NSSS_โˆ‚equations_โˆ‚SS_and_pars::AbstractMatrix{<: Real} # Dedicated NSSS SS Jacobian end # Structs for perturbation derivative functions (used for AD) @@ -707,26 +974,16 @@ end mutable struct model_functions # NSSS-related functions - NSSS_solve::Function NSSS_check::Function NSSS_custom::Union{Nothing, Function} NSSS_โˆ‚equations_โˆ‚parameters::Function NSSS_โˆ‚equations_โˆ‚SS_and_pars::Function + nsss_solver::NSSSSolverFunctions + nsss_param_prep!::Union{Nothing, Function} # Perturbation derivative functions jacobian::jacobian_functions hessian::hessian_functions third_order_derivatives::third_order_derivatives_functions - # State update functions for perturbation solutions - first_order_state_update::Function - first_order_state_update_obc::Function - second_order_state_update::Function - second_order_state_update_obc::Function - pruned_second_order_state_update::Function - pruned_second_order_state_update_obc::Function - third_order_state_update::Function - third_order_state_update_obc::Function - pruned_third_order_state_update::Function - pruned_third_order_state_update_obc::Function # OBC-related functions obc_violation::Function # Whether all functions have been written/compiled @@ -787,6 +1044,24 @@ mutable struct inversion_workspace{T <: Real} aug_stateโ‚::Vector{T} # n_past+1+n_exo aug_stateโ‚‚::Vector{T} # n_past+1+n_exo + # Estimation loop temporaries (lazily allocated via ensure_inversion_estimation_buffers!) + n_cond_var::Int # number of conditioning variables (observables) + shock_independent::Vector{T} # n_cond_var - shock-independent residual + init_guess::Vector{T} # n_exo - initial guess for find_shocks + Si_buffer::Matrix{T} # (n_cond_var, n_exo) - effective Jacobian ๐’โฑ workspace + jacc_buffer::Matrix{T} # (n_cond_var, n_exo) - Jacobian for logdet + Si2e_buffer::Matrix{T} # (n_cond_var, n_exo^2) - ๐’โฑยฒแต‰ workspace for 3rd order + # First-order inversion filter buffers + y_obs::Vector{T} # n_cond_var - observation prediction + x_shocks::Vector{T} # n_exo - recovered shocks + state_concat::Vector{T} # n_past + n_exo - for vcat-free concatenation in 1st order + # Pruned third-order augmented state buffers + aug_stateโ‚ƒ::Vector{T} # n_past+1+n_exo - third state component + aug_stateโ‚ฬ‚::Vector{T} # n_past+1+n_exo - hat state (vol=0) + stateยฒโป_vol::Vector{T} # n_past+1 - second-order state with volatility slot + # Third-order state kron buffers + kronstate_volยณ::Vector{T} # (n_past+1)^3 - triple kron of state_vol + # Pullback buffers (for reverse-mode AD in rrule) โˆ‚_tmp1::Matrix{T} # (n_exo, n_past + n_exo) โˆ‚_tmp2::Matrix{T} # (n_past, n_past + n_exo) @@ -804,7 +1079,7 @@ end """ Workspace for Kalman filter computations. Contains pre-allocated buffers for state estimates, covariances, and matrix operations. -Buffers are lazily allocated and resized as needed via ensure_kalman_buffers!. +Buffers are lazily allocated and resized as needed via ensure_kalman_workspaces!. """ mutable struct kalman_workspace{T <: Real} # Dimensions (for reallocation checks) @@ -819,10 +1094,16 @@ mutable struct kalman_workspace{T <: Real} # Matrix buffers Ctmp::Matrix{T} # (n_obs, n_states) - C*P buffer + ๐::Matrix{T} # (n_states, n_states) - B*B' buffer F::Matrix{T} # (n_obs, n_obs) - innovation covariance K::Matrix{T} # (n_states, n_obs) - Kalman gain tmp::Matrix{T} # (n_states, n_states) - temp for P Ptmp::Matrix{T} # (n_states, n_states) - temp for P + + # FastLapackInterface LU workspace for F factorization/solves + fast_lu_ws_f::FastLapackInterface.LUWs + fast_lu_dims_f::NTuple{2, Int} + fast_lu_rhs_t_k::Matrix{T} # (n_obs, n_states) scratch for right solves end @@ -839,8 +1120,12 @@ mutable struct higher_order_workspace{F <: Real, G <: AbstractFloat, H <: Real} tmp_sparse_prealloc4::Tuple{Vector{Int}, Vector{Int}, Vector{F}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{F}} tmp_sparse_prealloc5::Tuple{Vector{Int}, Vector{Int}, Vector{F}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{F}} tmp_sparse_prealloc6::Tuple{Vector{Int}, Vector{Int}, Vector{F}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{F}} + tmp_sparse_prealloc7::Tuple{Vector{Int}, Vector{Int}, Vector{F}, Vector{Int}, Vector{Int}, Vector{Int}, Vector{F}} + ๐’โ‚::Matrix{F} + ๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘::Matrix{F} Sฬ‚::Matrix{F} sylvester_workspace::sylvester_workspace{G, H} + โˆ‚โˆ‡_vec::Vector{F} # Flattened cotangent buffer for low-level higher-order derivative pullbacks # Pullback gradient buffers (lazily allocated, used in rrule pullback functions) # Second order pullback buffers โˆ‚โˆ‡โ‚‚::Matrix{F} @@ -854,6 +1139,38 @@ mutable struct higher_order_workspace{F <: Real, G <: AbstractFloat, H <: Real} โˆ‚โˆ‡โ‚_3rd::Matrix{F} # separate from 2nd order since dimensions differ โˆ‚๐’โ‚_3rd::Matrix{F} # separate from 2nd order since dimensions differ โˆ‚spinv_3rd::Matrix{F} # separate from 2nd order since dimensions differ + โˆ‚โˆ‡โ‚‚_3rd::Matrix{F} + โˆ‚โˆ‡โ‚ƒ_3rd::Matrix{F} + โˆ‚๐’โ‚‚_3rd::Matrix{F} + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_3rd::Matrix{F} + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_3rd::Matrix{F} + โˆ‚โŽธ๐’โ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘โŽนโ•ฑ๐’โ‚โ•ฑ๐Ÿโ‚‘โ‚‹_3rd::Matrix{F} + # Third order pullback temporary buffers (reused across calls) + โˆ‚๐’โ‚‚โ‚Šโ•ฑ๐ŸŽ_3rd::Matrix{F} + โˆ‚R_c_3rd::Matrix{F} + โˆ‚L_c_3rd::Matrix{F} + โˆ‚L_d_3rd::Matrix{F} + โˆ‚R_d_3rd::Matrix{F} + โˆ‚๐’โ‚‚โ‚‹โ•ฑ๐ŸŽ_3rd::Matrix{F} + โˆ‚๐’โ‚โ‚‹โ•ฑ๐Ÿโ‚‘_t8_3rd::Matrix{F} + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tmp_3rd::Matrix{F} + โˆ‚๐’โ‚โ‚Šโ•ฑ๐ŸŽ_tk0_3rd::Matrix{F} + โˆ‚tmpkron0_ฯƒ_3rd::Matrix{F} + โˆ‚aux_3rd::Matrix{F} + โˆ‚S1S1_from_ck_3rd::Matrix{F} + โˆ‚S1p0_kron_sigma_3rd::Matrix{F} + โˆ‚S1p0_left_3rd::Matrix{F} + โˆ‚S1p0_right_3rd::Matrix{F} + # Third order pullback intermediate product buffers (for mul!) + โˆ‚A_3rd::Matrix{F} + โˆ‚B_sylv_3rd::Matrix{F} + โˆ‚๐—โ‚ƒ_3rd::Matrix{F} + โˆ‚๐—โ‚ƒ_pre_3rd::Matrix{F} + โˆ‚out2_3rd::Matrix{F} + โˆ‚โˆ‡โ‚โ‚Š_3rd::Matrix{F} + โˆ‚โˆ‡โ‚โ‚Š๐’โ‚โž•โˆ‡โ‚โ‚€_3rd::Matrix{F} + โˆ‡โ‚‚t_โˆ‚out2_3rd::Matrix{F} + mul_tmp_3rd::Matrix{F} # ForwardDiff partials buffers for stochastic steady state (accessed via model struct) โˆ‚x_second_order::Matrix{H} # For second order SSS partials โˆ‚x_third_order::Matrix{H} # For third order SSS partials @@ -871,8 +1188,9 @@ Purpose: Speed up computation by eliminating allocation overhead in hot loops. Fields: - `second_order/third_order`: Higher-order perturbation solution workspaces -- `custom_steady_state_buffer`: Buffer for custom steady state evaluation -- `qme`: Quadratic matrix equation solver workspace +- `custom_steady_state`: Buffer for custom steady state evaluation +- `first_order`: First-order perturbation solver workspace +- `qme_doubling`: Quadratic matrix equation doubling solver workspace - `lyapunov_*`: Lyapunov equation solver workspaces (1st, 2nd, 3rd order) - `sylvester_*`: Sylvester equation solver workspace - `find_shocks`: Conditional forecast shock finding workspace @@ -889,17 +1207,23 @@ mutable struct workspaces second_order::higher_order_workspace # Kronecker products, sparse preallocs third_order::higher_order_workspace # Separate workspace for 3rd order # Steady state buffer - custom_steady_state_buffer::Vector{Float64} # For custom SS function evaluation + custom_steady_state::Vector{Float64} # For custom SS function evaluation # Matrix equation solver workspaces - qme::qme_workspace{Float64, Float64} # Quadratic matrix equation (1st order) + first_order::first_order_workspace{Float64, Float64} # First-order perturbation solver + qme_doubling::qme_doubling_workspace{Float64, Float64} # QME doubling solver + schur::schur_workspace{Float64} # Schur-based QME solver lyapunov_1st_order::lyapunov_workspace{Float64, Float64} # Covariance (1st order moments) lyapunov_2nd_order::lyapunov_workspace{Float64, Float64} # Covariance (2nd order moments) lyapunov_3rd_order::lyapunov_workspace{Float64, Float64} # Covariance (3rd order moments) + lyapunov_block::lyapunov_workspace{Float64, Float64} # Block-triangular inner Lyapunov sylvester_1st_order::sylvester_workspace{Float64, Float64} # Sylvester equation + sylvester_block::sylvester_workspace{Float64, Float64} # Block-triangular Sylvester # Filter workspaces find_shocks::find_shocks_workspace{Float64} # Conditional forecast shock finding inversion::inversion_workspace{Float64} # Inversion filter kalman::kalman_workspace{Float64} # Kalman filter + # NSSS solver shared scratch buffers + nsss_solver::NSSSSolverWorkspace end @@ -907,7 +1231,9 @@ end struct post_parameters_macro parameters_as_function_of_parameters::Vector{Symbol} precompile::Bool - simplify::Bool + ss_symbolic_mode::Symbol + ss_solver_parameters_algorithm::Symbol + ss_solver_parameters_maxtime::Float64 guess::Dict{Symbol, Float64} ss_calib_list::Vector{Set{Symbol}} par_calib_list::Vector{Set{Symbol}} @@ -941,6 +1267,8 @@ struct post_complete_parameters{S <: Union{Symbol, String}} custom_ss_expand_matrix::SparseMatrixCSC{Float64, Int} vars_in_ss_equations::Vector{Symbol} vars_in_ss_equations_with_aux::Vector{Symbol} + ss_var_idx_in_var_and_calib::Vector{Int} + calib_idx_in_var_and_calib::Vector{Int} SS_and_pars_names_lead_lag::Vector{Symbol} # SS_and_pars_names_no_exo::Vector{Symbol} SS_and_pars_no_exo_idx::Vector{Int} @@ -953,11 +1281,29 @@ struct post_complete_parameters{S <: Union{Symbol, String}} future_not_past_and_mixed_in_comb::Vector{Int} past_not_future_and_mixed_in_comb::Vector{Int} Ir::โ„’.Diagonal{Bool, Vector{Bool}} + I_n::โ„’.Diagonal{Bool, Vector{Bool}} nabla_zero_cols::UnitRange{Int} nabla_minus_cols::UnitRange{Int} nabla_e_start::Int expand_future::Matrix{Bool} expand_past::Matrix{Bool} + past_not_future_and_mixed_in_present_but_not_only::Vector{Int} + # Schur QME cached indices and constant matrices + indices_past_not_future_in_comb::Vector{Int} + I_nPast_not_mixed::Matrix{Bool} # I_nPast[not_mixed_in_past_idx,:] + Ir_past_selector::Matrix{Bool} # Ir[past_not_future_and_mixed_in_comb,:] + schur_Zโ‚Š::Matrix{Bool} # zeros(nMixed, nFuture_not_past_and_mixed) + schur_Iโ‚Š::Matrix{Bool} # I(nFuture_not_past_and_mixed)[mixed_in_future_idx,:] + schur_Zโ‚‹::Matrix{Bool} # zeros(nMixed, nPast_not_future_and_mixed) + schur_Iโ‚‹::Matrix{Bool} # I_nPast[mixed_in_past_idx,:] + nsss_dependencies::Any + nsss_n_sol::Int + nsss_output_indices::Vector{Int} + nsss_n_ext_params::Int + nsss_sol_names::Vector{Symbol} + nsss_exo_zero_indices::Vector{Int} + nsss_param_names_ext::Vector{Symbol} + nsss_fastest_solver_parameter_idx::Int end """ @@ -990,6 +1336,8 @@ mutable struct constants#{F <: Real, G <: AbstractFloat} second_order::second_order_indices # Third-order perturbation auxiliary matrices and indices third_order::third_order_indices + # NSSS solver step constants (indices, bounds, metadata) + nsss_solver::NSSSSolverConstants end mutable struct solver_parameters @@ -1099,11 +1447,6 @@ mutable struct โ„ณ model_name::Any # Model identifier parameter_values::Vector{Float64} # Current parameter values (mutable) - # ========================================================================= - # STEADY STATE SOLVER INFRASTRUCTURE - # ========================================================================= - NSSS::non_stochastic_steady_state # Steady state solver blocks - # ========================================================================= # MODEL EQUATIONS (various representations) # ========================================================================= diff --git a/tasks/repro_higher_order_irf_test.jl b/tasks/repro_higher_order_irf_test.jl new file mode 100644 index 000000000..0223cae4b --- /dev/null +++ b/tasks/repro_higher_order_irf_test.jl @@ -0,0 +1,108 @@ +using MacroModelling +using Random +using Test +import LinearAlgebra as LA + +# Isolated reproduction of the higher-order IRF assertions from +# test/test_standalone_function.jl (without running the full test file). + +include("../test/models/RBC_CME.jl") + +Random.seed!(3) + +SS_and_pars, _ = MacroModelling.get_NSSS_and_parameters(m, m.parameter_values) +get_irf(m, algorithm = :third_order) +get_irf(m, algorithm = :pruned_third_order) +get_irf(m, algorithm = :pruned_second_order) + +โˆ‡โ‚ = calculate_jacobian(m.parameter_values, SS_and_pars, m.caches, m.functions.jacobian, m.workspaces) +โˆ‡โ‚‚ = calculate_hessian(m.parameter_values, SS_and_pars, m.caches, m.functions.hessian, m.workspaces) +โˆ‡โ‚ƒ = calculate_third_order_derivatives(m.parameter_values, SS_and_pars, m.caches, m.functions.third_order_derivatives, m.workspaces) + +T = m.constants.post_model_macro + +first_order_solution, _, _ = calculate_first_order_solution(โˆ‡โ‚, m.constants, m.workspaces, m.caches) +second_order_solution, _ = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, first_order_solution, m.constants, m.workspaces, m.caches) +third_order_solution, _ = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, โˆ‡โ‚ƒ, first_order_solution, second_order_solution, m.constants, m.workspaces, m.caches) + +second_order_solution = sparse(second_order_solution * m.constants.second_order.๐”โ‚‚) +third_order_solution = sparse(third_order_solution * m.constants.third_order.๐”โ‚ƒ) + +Tz = [first_order_solution[:, 1:T.nPast_not_future_and_mixed] zeros(T.nVars) first_order_solution[:, T.nPast_not_future_and_mixed+1:end]] + +second_order_state_update = function(state::Vector{Float64}, shock::Vector{Float64}) + aug_state = [state[T.past_not_future_and_mixed_idx] + 1 + shock] + return Tz * aug_state + second_order_solution * kron(aug_state, aug_state) / 2 +end + +third_order_state_update = function(state::Vector{Float64}, shock::Vector{Float64}) + aug_state = [state[T.past_not_future_and_mixed_idx] + 1 + shock] + return Tz * aug_state + + second_order_solution * kron(aug_state, aug_state) / 2 + + third_order_solution * kron(kron(aug_state, aug_state), aug_state) / 6 +end + +pruned_second_order_state_update = function(pruned_states::Vector{Vector{Float64}}, shock::Vector{Float64}) + aug_stateโ‚ = [pruned_states[1][m.constants.post_model_macro.past_not_future_and_mixed_idx]; 1; shock] + aug_stateโ‚‚ = [pruned_states[2][m.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] + return [Tz * aug_stateโ‚, + Tz * aug_stateโ‚‚ + second_order_solution * LA.kron(aug_stateโ‚, aug_stateโ‚) / 2] +end + +pruned_third_order_state_update = function(pruned_states::Vector{Vector{Float64}}, shock::Vector{Float64}) + aug_stateโ‚ = [pruned_states[1][m.constants.post_model_macro.past_not_future_and_mixed_idx]; 1; shock] + aug_stateโ‚ฬ‚ = [pruned_states[1][m.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; shock] + aug_stateโ‚‚ = [pruned_states[2][m.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] + aug_stateโ‚ƒ = [pruned_states[3][m.constants.post_model_macro.past_not_future_and_mixed_idx]; 0; zero(shock)] + + kron_aug_stateโ‚ = LA.kron(aug_stateโ‚, aug_stateโ‚) + return [Tz * aug_stateโ‚, + Tz * aug_stateโ‚‚ + second_order_solution * kron_aug_stateโ‚ / 2, + Tz * aug_stateโ‚ƒ + second_order_solution * LA.kron(aug_stateโ‚ฬ‚, aug_stateโ‚‚) + third_order_solution * LA.kron(kron_aug_stateโ‚, aug_stateโ‚) / 6] +end + +# Reproduce exactly the four IRF checks that are currently failing in CI. +SSS_delta_2 = m.caches.non_stochastic_steady_state[1:length(m.constants.post_model_macro.var)] - m.caches.second_order_stochastic_steady_state +initial_state_2 = zeros(m.constants.post_model_macro.nVars) - SSS_delta_2 +iirrff2 = irf(second_order_state_update, initial_state_2 + SSS_delta_2, zeros(T.nVars), m.constants) + +SSS_delta_3 = m.caches.non_stochastic_steady_state[1:length(m.constants.post_model_macro.var)] - m.caches.third_order_stochastic_steady_state +initial_state_3 = zeros(m.constants.post_model_macro.nVars) - SSS_delta_3 +iirrff3 = irf(third_order_state_update, initial_state_3 + SSS_delta_3, zeros(T.nVars), m.constants) + +iirrffp2 = irf(pruned_second_order_state_update, + [zeros(m.constants.post_model_macro.nVars), zeros(m.constants.post_model_macro.nVars)], + zeros(T.nVars), + m.constants) + +iirrffp3 = irf(pruned_third_order_state_update, + [zeros(m.constants.post_model_macro.nVars), zeros(m.constants.post_model_macro.nVars), zeros(m.constants.post_model_macro.nVars)], + zeros(T.nVars), + m.constants) + +expected_iirrff2 = [-0.0004547347878067665, 0.0020831426377533636] +expected_iirrff3 = [-0.00045473149068020854, 0.002083198241302615] +expected_iirrffp2 = [-0.00045473478780675195, 0.002083142637753389] +expected_iirrffp3 = [-0.0004547315171573783, 0.0020831990353127696] + +actual_iirrff2 = vec(iirrff2[4, 1, :]) +actual_iirrff3 = vec(iirrff3[4, 1, :]) +actual_iirrffp2 = vec(iirrffp2[4, 1, :]) +actual_iirrffp3 = vec(iirrffp3[4, 1, :]) + +println("Higher-order IRF isolated repro") +println("iirrff2 actual=$(actual_iirrff2) expected=$(expected_iirrff2)") +println("iirrff3 actual=$(actual_iirrff3) expected=$(expected_iirrff3)") +println("iirrffp2 actual=$(actual_iirrffp2) expected=$(expected_iirrffp2)") +println("iirrffp3 actual=$(actual_iirrffp3) expected=$(expected_iirrffp3)") + +@test isapprox(actual_iirrff2, expected_iirrff2, rtol = 1e-6) +@test isapprox(actual_iirrff3, expected_iirrff3, rtol = 1e-6) +@test isapprox(actual_iirrffp2, expected_iirrffp2, rtol = 1e-6) +@test isapprox(actual_iirrffp3, expected_iirrffp3, rtol = 1e-6) + +println("HIGHER_ORDER_IRF_REPRO=PASS") diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index 873de5a12..a479dd981 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -185,7 +185,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for quadratic_matrix_equation_algorithm in qme_algorithms for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms - for tol in [MacroModelling.Tolerances(), MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(), MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] clear_solution_caches!(m, algorithm) plot_model_estimates(m, data, @@ -220,7 +220,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for quadratic_matrix_equation_algorithm in qme_algorithms for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms - for tol in [MacroModelling.Tolerances(NSSS_xtol = 1e-14), MacroModelling.Tolerances()] + for tol in [MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14)), MacroModelling.Tolerances()] if i % 4 == 0 plot_model_estimates(m, data_in_levels, algorithm = algorithm, @@ -467,7 +467,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end for variables in vars - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -723,7 +723,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) shock_mat3 = KeyedArray(randn(m.constants.post_model_macro.nExo,10),Shocks = string.(m.constants.post_model_macro.exo), Periods = 1:10) for parameters in params - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -747,7 +747,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) i = 1 for parameters in params - for tol in [MacroModelling.Tolerances(NSSS_xtol = 1e-14), MacroModelling.Tolerances()] + for tol in [MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14)), MacroModelling.Tolerances()] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -919,7 +919,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms clear_solution_caches!(m, algorithm) @@ -1144,7 +1144,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) # end - for tol in [MacroModelling.Tolerances(), MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(), MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -1179,7 +1179,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) i = 1 - for tol in [MacroModelling.Tolerances(NSSS_xtol = 1e-14), MacroModelling.Tolerances()] + for tol in [MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14)), MacroModelling.Tolerances()] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -1588,7 +1588,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end for parameters in params - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] get_shock_decomposition(m, data, parameters = parameters, algorithm = algorithm, @@ -1677,7 +1677,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end for parameters in params - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] get_estimated_variable_standard_deviations(m, data, parameters = parameters, data_in_levels = false, @@ -1697,7 +1697,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for initial_covariance in [:diagonal, :theoretical] for verbose in [false] # [true, false] for parameter_values in [old_params, old_params .* exp.(rand(length(old_params))*1e-4)] - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] llh = get_loglikelihood(m, data_in_levels, parameter_values, algorithm = algorithm, filter = filter, @@ -1706,36 +1706,36 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) tol = tol, verbose = verbose) - clear_solution_caches!(m, algorithm) - - zyg_grad_llh = Zygote.gradient(x -> get_loglikelihood(m, data_in_levels, x, - algorithm = algorithm, - filter = filter, - presample_periods = presample_periods, - initial_covariance = initial_covariance, - tol = tol, - verbose = verbose), parameter_values) - - if algorithm == :first_order && filter == :kalman - for i in 1:100 - local fin_grad_llh = FiniteDifferences.grad(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-3), - x -> begin - clear_solution_caches!(m, algorithm) - - get_loglikelihood(m, data_in_levels, x, - algorithm = algorithm, - filter = filter, - presample_periods = presample_periods, - initial_covariance = initial_covariance, - tol = tol, - verbose = verbose) - end, parameter_values) - if isfinite(โ„’.norm(fin_grad_llh[1])) - @test isapprox(fin_grad_llh[1], zyg_grad_llh[1], rtol = 1e-5) - break + clear_solution_caches!(m, algorithm) + + zyg_grad_llh = Zygote.gradient(x -> get_loglikelihood(m, data_in_levels, x, + algorithm = algorithm, + filter = filter, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + tol = tol, + verbose = verbose), parameter_values) + + if algorithm == :first_order && filter == :kalman + for i in 1:100 + local fin_grad_llh = FiniteDifferences.grad(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + + get_loglikelihood(m, data_in_levels, x, + algorithm = algorithm, + filter = filter, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + tol = tol, + verbose = verbose) + end, parameter_values) + if isfinite(โ„’.norm(fin_grad_llh[1])) + @test isapprox(fin_grad_llh[1], zyg_grad_llh[1], rtol = 1e-5) + break + end end end - end for quadratic_matrix_equation_algorithm in qme_algorithms for lyapunov_algorithm in lyapunov_algorithms @@ -1755,20 +1755,20 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) verbose = verbose) @test isapprox(llh, LLH, rtol = 1e-8) - clear_solution_caches!(m, algorithm) - - ZYG_grad_llh = Zygote.gradient(x -> get_loglikelihood(m, data_in_levels, x, - algorithm = algorithm, - filter = filter, - presample_periods = presample_periods, - initial_covariance = initial_covariance, - tol = tol, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - lyapunov_algorithm = lyapunov_algorithm, - sylvester_algorithm = sylvester_algorithm, - verbose = verbose), parameter_values) - - @test isapprox(ZYG_grad_llh[1], zyg_grad_llh[1], rtol = 1e-6) + clear_solution_caches!(m, algorithm) + + ZYG_grad_llh = Zygote.gradient(x -> get_loglikelihood(m, data_in_levels, x, + algorithm = algorithm, + filter = filter, + presample_periods = presample_periods, + initial_covariance = initial_covariance, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + lyapunov_algorithm = lyapunov_algorithm, + sylvester_algorithm = sylvester_algorithm, + verbose = verbose), parameter_values) + + @test isapprox(ZYG_grad_llh[1], zyg_grad_llh[1], rtol = 1e-6) end end end @@ -1872,7 +1872,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for variables in vars for levels in [true, false] for verbose in [false] # [true, false] - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -2044,7 +2044,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for verbose in [false] # [true, false] - for tol in [MacroModelling.Tolerances(), MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(), MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms for lyapunov_algorithm in lyapunov_algorithms @@ -2127,7 +2127,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for verbose in [false] # [true, false] - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms for sylvester_algorithm in sylvester_algorithms clear_solution_caches!(m, algorithm) @@ -2161,37 +2161,39 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sol = get_solution(m, parameter_values, algorithm = algorithm) - clear_solution_caches!(m, algorithm) + deriv_sol = nothing + deriv_sol_zyg = nothing + clear_solution_caches!(m, algorithm) - deriv_sol = [] - for i in 1:length(sol)-2 - push!(deriv_sol, ForwardDiff.jacobian(x->get_solution(m, x, algorithm = algorithm)[i], parameter_values)) - end + deriv_sol = [] + for i in 1:length(sol)-2 + push!(deriv_sol, ForwardDiff.jacobian(x->get_solution(m, x, algorithm = algorithm)[i], parameter_values)) + end - clear_solution_caches!(m, algorithm) + clear_solution_caches!(m, algorithm) - deriv_sol_fin = [] - for i in 1:length(sol)-2 - push!(deriv_sol_fin, FiniteDifferences.jacobian(FiniteDifferences.forward_fdm(3,1, max_range = 1e-3), - x -> begin - clear_solution_caches!(m, algorithm) - - get_solution(m, x, algorithm = algorithm)[i] - end, parameter_values)[1]) - end + deriv_sol_fin = [] + for i in 1:length(sol)-2 + push!(deriv_sol_fin, FiniteDifferences.jacobian(FiniteDifferences.forward_fdm(3,1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + + get_solution(m, x, algorithm = algorithm)[i] + end, parameter_values)[1]) + end - clear_solution_caches!(m, algorithm) + clear_solution_caches!(m, algorithm) - deriv_sol_zyg = [] - for i in 1:length(sol)-2 - push!(deriv_sol_zyg, Zygote.jacobian(x->get_solution(m, x, algorithm = algorithm)[i], parameter_values)[1]) - end + deriv_sol_zyg = [] + for i in 1:length(sol)-2 + push!(deriv_sol_zyg, Zygote.jacobian(x->get_solution(m, x, algorithm = algorithm)[i], parameter_values)[1]) + end - @test isapprox(deriv_sol_zyg, deriv_sol_fin, rtol = 1e-5) - - @test isapprox(deriv_sol, deriv_sol_fin, rtol = 1e-5) + @test isapprox(deriv_sol_zyg, deriv_sol_fin, rtol = 1e-5) + + @test isapprox(deriv_sol, deriv_sol_fin, rtol = 1e-5) - for tol in [MacroModelling.Tolerances(lyapunov_acceptance_tol = 1e-14, sylvester_acceptance_tol = 1e-14), MacroModelling.Tolerances(lyapunov_acceptance_tol = 1e-14, sylvester_acceptance_tol = 1e-14, NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)))] for quadratic_matrix_equation_algorithm in qme_algorithms for sylvester_algorithm in sylvester_algorithms clear_solution_caches!(m, algorithm) @@ -2202,29 +2204,29 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) @test isapprox([s for s in sol[1:end-1]], [S for S in SOL[1:end-1]], rtol = 1e-8) - clear_solution_caches!(m, algorithm) + clear_solution_caches!(m, algorithm) - DERIV_SOL = [] - for i in 1:length(sol)-2 - push!(DERIV_SOL, ForwardDiff.jacobian(x->get_solution(m, x, algorithm = algorithm, - tol = tol, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - sylvester_algorithm = sylvester_algorithm)[i], parameter_values)) - end + DERIV_SOL = [] + for i in 1:length(sol)-2 + push!(DERIV_SOL, ForwardDiff.jacobian(x->get_solution(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + sylvester_algorithm = sylvester_algorithm)[i], parameter_values)) + end - @test isapprox(deriv_sol, DERIV_SOL, rtol = 1e-8) + @test isapprox(deriv_sol, DERIV_SOL, rtol = 1e-8) - clear_solution_caches!(m, algorithm) + clear_solution_caches!(m, algorithm) - DERIV_SOL_zyg = [] - for i in 1:length(sol)-2 - push!(DERIV_SOL_zyg, Zygote.jacobian(x->get_solution(m, x, algorithm = algorithm, - tol = tol, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - sylvester_algorithm = sylvester_algorithm)[i], parameter_values)[1]) - end + DERIV_SOL_zyg = [] + for i in 1:length(sol)-2 + push!(DERIV_SOL_zyg, Zygote.jacobian(x->get_solution(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + sylvester_algorithm = sylvester_algorithm)[i], parameter_values)[1]) + end - @test isapprox(deriv_sol_zyg, DERIV_SOL_zyg, rtol = 1e-8) + @test isapprox(deriv_sol_zyg, DERIV_SOL_zyg, rtol = 1e-8) end end end @@ -2319,12 +2321,65 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) get_irf(m, x, initial_state = initial_state)[:,1,1] end, parameter_values) if isfinite(โ„’.norm(deriv_fin[1])) - @test isapprox(deriv_for, deriv_fin[1], rtol = 1e-5) + @test isapprox(deriv_for, deriv_fin[1], rtol = 1e-5, atol = 1e-8) + break + end + end + + clear_solution_caches!(m, algorithm) + + deriv_zyg = Zygote.jacobian(x -> get_irf(m, x, initial_state = initial_state)[:,1,1], parameter_values)[1] + + for i in 1:100 + local deriv_fin_zyg = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-4), + x -> begin + clear_solution_caches!(m, algorithm) + + get_irf(m, x, initial_state = initial_state)[:,1,1] + end, parameter_values) + if isfinite(โ„’.norm(deriv_fin_zyg[1])) + @test isapprox(deriv_zyg, deriv_fin_zyg[1], rtol = 1e-5, atol = 1e-8) break end end - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + # Last period derivative tests (ForwardDiff) + clear_solution_caches!(m, algorithm) + + deriv_for_last = ForwardDiff.jacobian(x->get_irf(m, x, initial_state = initial_state)[:,end,1], parameter_values) + + for i in 1:100 + local deriv_fin_last = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-4), + x -> begin + clear_solution_caches!(m, algorithm) + + get_irf(m, x, initial_state = initial_state)[:,end,1] + end, parameter_values) + if isfinite(โ„’.norm(deriv_fin_last[1])) + @test isapprox(deriv_for_last, deriv_fin_last[1], rtol = 1e-5, atol = 1e-8) + break + end + end + + # Last period derivative tests (Zygote) + clear_solution_caches!(m, algorithm) + + deriv_zyg_last = Zygote.jacobian(x -> get_irf(m, x, initial_state = initial_state)[:,end,1], parameter_values)[1] + + for i in 1:100 + local deriv_fin_zyg_last = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-4), + x -> begin + clear_solution_caches!(m, algorithm) + + get_irf(m, x, initial_state = initial_state)[:,end,1] + end, parameter_values) + if isfinite(โ„’.norm(deriv_fin_zyg_last[1])) + @test isapprox(deriv_zyg_last, deriv_fin_zyg_last[1], rtol = 1e-5, atol = 1e-8) + break + end + end + + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms clear_solution_caches!(m, algorithm) @@ -2388,7 +2443,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) clear_solution_caches!(m, algorithm) stats = get_statistics(m, parameter_values, algorithm = algorithm, - # tol = MacroModelling.Tolerances(lyapunov_acceptance_tol = 1e-14, sylvester_acceptance_tol = 1e-14, NSSS_xtol = 1e-14), + # tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), non_stochastic_steady_state = :all, mean = (algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), standard_deviation = (algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), @@ -2396,7 +2451,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) covariance = (algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), autocorrelation = (algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[])) - for tol in [MacroModelling.Tolerances(lyapunov_acceptance_tol = 1e-14, sylvester_acceptance_tol = 1e-14),MacroModelling.Tolerances(lyapunov_acceptance_tol = 1e-14, sylvester_acceptance_tol = 1e-14,NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)))] for quadratic_matrix_equation_algorithm in qme_algorithms for sylvester_algorithm in sylvester_algorithms for lyapunov_algorithm in lyapunov_algorithms @@ -2436,42 +2491,42 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end - clear_solution_caches!(m, algorithm) + clear_solution_caches!(m, algorithm) - deriv1 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, - non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) + deriv1 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, + non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) - deriv1_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, - non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) + deriv1_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, + non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) - for i in 1:100 - local deriv1_fin = FiniteDifferences.jacobian(FiniteDifferences.forward_fdm(3,1, max_range = 1e-3), - x -> begin - clear_solution_caches!(m, algorithm) - - get_statistics(m, x, - algorithm = algorithm, - non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state] - end, old_params) - if isfinite(โ„’.norm(deriv1_fin[1])) - # โ„’.norm(deriv1 - deriv1_fin[1]) / max(โ„’.norm(deriv1), โ„’.norm(deriv1_fin[1])) - # โ„’.norm(deriv1 - deriv1_zyg[1]) / max(โ„’.norm(deriv1), โ„’.norm(deriv1_zyg[1])) - - @test isapprox(deriv1_zyg[1], deriv1_fin[1], rtol = 1e-5) + for i in 1:100 + local deriv1_fin = FiniteDifferences.jacobian(FiniteDifferences.forward_fdm(3,1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) - @test isapprox(deriv1, deriv1_fin[1], rtol = 1e-5) - break + get_statistics(m, x, + algorithm = algorithm, + non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state] + end, old_params) + if isfinite(โ„’.norm(deriv1_fin[1])) + # โ„’.norm(deriv1 - deriv1_fin[1]) / max(โ„’.norm(deriv1), โ„’.norm(deriv1_fin[1])) + # โ„’.norm(deriv1 - deriv1_zyg[1]) / max(โ„’.norm(deriv1), โ„’.norm(deriv1_zyg[1])) + + @test isapprox(deriv1_zyg[1], deriv1_fin[1], rtol = 1e-5) + + @test isapprox(deriv1, deriv1_fin[1], rtol = 1e-5) + break + end end - end - if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] clear_solution_caches!(m, algorithm) deriv2 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, mean = :all_excluding_obc)[:mean], old_params) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] deriv2_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, mean = :all_excluding_obc)[:mean], old_params) end @@ -2487,7 +2542,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end, old_params) if isfinite(โ„’.norm(deriv2_fin[1])) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] @test isapprox(deriv2_zyg[1], deriv2_fin[1], rtol = 1e-5) end @@ -2501,7 +2556,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) deriv3 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, standard_deviation = :all_excluding_obc)[:standard_deviation], old_params) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] deriv3_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, standard_deviation = :all_excluding_obc)[:standard_deviation], old_params) end @@ -2515,7 +2570,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end, old_params) if isfinite(โ„’.norm(deriv3_fin[1])) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] @test isapprox(deriv3_zyg[1], deriv3_fin[1], rtol = 1e-5) end @@ -2529,7 +2584,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) deriv4 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, variance = :all_excluding_obc)[:variance], old_params) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] deriv4_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, variance = :all_excluding_obc)[:variance], old_params) end @@ -2542,7 +2597,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) get_statistics(m, x, algorithm = algorithm, variance = :all_excluding_obc)[:variance] end, old_params) if isfinite(โ„’.norm(deriv4_fin[1])) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] @test isapprox(deriv4_zyg[1], deriv4_fin[1], rtol = 1e-5) end @test isapprox(deriv4, deriv4_fin[1], rtol = 1e-5) @@ -2553,14 +2608,12 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) clear_solution_caches!(m, algorithm) deriv5 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, - tol = MacroModelling.Tolerances(NSSS_xtol = 1e-14, lyapunov_acceptance_tol = 1e-14, - sylvester_acceptance_tol = 1e-14), + tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), covariance = :all_excluding_obc)[:covariance], old_params) - if algorithm == :first_order_ + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] deriv5_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, - tol = MacroModelling.Tolerances(NSSS_xtol = 1e-14, lyapunov_acceptance_tol = 1e-14, - sylvester_acceptance_tol = 1e-14), + tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), covariance = :all_excluding_obc)[:covariance], old_params) end @@ -2570,12 +2623,11 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) clear_solution_caches!(m, algorithm) get_statistics(m, x, algorithm = algorithm, - tol = MacroModelling.Tolerances(NSSS_xtol = 1e-14, lyapunov_acceptance_tol = 1e-14, - sylvester_acceptance_tol = 1e-14), + tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), covariance = :all_excluding_obc)[:covariance] end, old_params) if isfinite(โ„’.norm(deriv5_fin[1])) - if algorithm == :first_order_ + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] @test isapprox(deriv5_zyg[1], deriv5_fin[1], rtol = 1e-4) end @@ -2584,36 +2636,86 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) break end end - end + + clear_solution_caches!(m, algorithm) + + deriv6 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, + autocorrelation = :all_excluding_obc)[:autocorrelation], old_params) + + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + deriv6_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, + autocorrelation = :all_excluding_obc)[:autocorrelation], old_params) + end + + for i in 1:100 + local deriv6_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + + get_statistics(m, x, algorithm = algorithm, autocorrelation = :all_excluding_obc)[:autocorrelation] + end, old_params) + if isfinite(โ„’.norm(deriv6_fin[1])) + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + @test isapprox(deriv6_zyg[1], deriv6_fin[1], rtol = 1e-4) + end + @test isapprox(deriv6, deriv6_fin[1], rtol = 1e-4) + break + end + end + + if algorithm == :pruned_third_order + var_obj = x -> begin + Zygote.ChainRulesCore.@ignore_derivatives clear_solution_caches!(m, algorithm) + get_statistics(m, x, algorithm = algorithm, variance = :all_excluding_obc)[:variance] |> sum + end + + autocorr_obj = x -> begin + Zygote.ChainRulesCore.@ignore_derivatives clear_solution_caches!(m, algorithm) + get_statistics(m, x, algorithm = algorithm, autocorrelation = :all_excluding_obc)[:autocorrelation] |> sum + end + + var_grad_zyg = Zygote.gradient(var_obj, old_params)[1] + var_grad_fin = FiniteDifferences.grad(FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), var_obj, old_params)[1] + @test all(isfinite, var_grad_zyg) + @test all(isfinite, var_grad_fin) + @test โ„’.norm(var_grad_zyg - var_grad_fin) / max(โ„’.norm(var_grad_fin), eps()) < 1e-4 + + autocorr_grad_zyg = Zygote.gradient(autocorr_obj, old_params)[1] + autocorr_grad_fin = FiniteDifferences.grad(FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), autocorr_obj, old_params)[1] + @test all(isfinite, autocorr_grad_zyg) + @test all(isfinite, autocorr_grad_fin) + @test โ„’.norm(autocorr_grad_zyg - autocorr_grad_fin) / max(โ„’.norm(autocorr_grad_fin), eps()) < 1e-4 + end + end - for tol in [MacroModelling.Tolerances(NSSS_xtol = 1e-14, lyapunov_acceptance_tol = 1e-14, sylvester_acceptance_tol = 1e-14)] - for quadratic_matrix_equation_algorithm in qme_algorithms - for sylvester_algorithm in sylvester_algorithms - for lyapunov_algorithm in lyapunov_algorithms - clear_solution_caches!(m, algorithm) + for tol in [MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)))] + for quadratic_matrix_equation_algorithm in qme_algorithms + for sylvester_algorithm in sylvester_algorithms + for lyapunov_algorithm in lyapunov_algorithms + clear_solution_caches!(m, algorithm) - DERIV1 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, - tol = tol, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - lyapunov_algorithm = lyapunov_algorithm, - sylvester_algorithm = sylvester_algorithm, - non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) - @test isapprox(deriv1, DERIV1, rtol = 1e-8) - - DERIV1_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, - tol = tol, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - lyapunov_algorithm = lyapunov_algorithm, - sylvester_algorithm = sylvester_algorithm, - non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) - @test isapprox(deriv1_zyg[1], DERIV1_zyg[1], rtol = 1e-8) + DERIV1 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + lyapunov_algorithm = lyapunov_algorithm, + sylvester_algorithm = sylvester_algorithm, + non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) + @test isapprox(deriv1, DERIV1, rtol = 1e-8) + + DERIV1_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + lyapunov_algorithm = lyapunov_algorithm, + sylvester_algorithm = sylvester_algorithm, + non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], old_params) + @test isapprox(deriv1_zyg[1], DERIV1_zyg[1], rtol = 1e-8) - if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] - clear_solution_caches!(m, algorithm) + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + clear_solution_caches!(m, algorithm) DERIV2 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, tol = tol, @@ -2623,7 +2725,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) mean = :all_excluding_obc)[:mean], old_params) @test isapprox(deriv2, DERIV2, rtol = 1e-8) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] clear_solution_caches!(m, algorithm) DERIV2_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, @@ -2645,7 +2747,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) standard_deviation = :all_excluding_obc)[:standard_deviation], old_params) @test isapprox(deriv3, DERIV3, rtol = 1e-8) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] clear_solution_caches!(m, algorithm) DERIV3_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, @@ -2654,7 +2756,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, standard_deviation = :all_excluding_obc)[:standard_deviation], old_params) - @test isapprox(deriv3_zyg[1], DERIV3_zyg[1], rtol = 1e-8) + @test isapprox(deriv3_zyg[1], DERIV3_zyg[1], rtol = 1e-6) end clear_solution_caches!(m, algorithm) @@ -2667,7 +2769,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) variance = :all_excluding_obc)[:variance], old_params) @test isapprox(deriv4, DERIV4, rtol = 1e-8) - if algorithm == :first_order + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] clear_solution_caches!(m, algorithm) DERIV4_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, @@ -2690,7 +2792,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) # println(โ„’.norm(deriv5 - DERIV5) / max(โ„’.norm(deriv5), โ„’.norm(DERIV5))) @test isapprox(deriv5, DERIV5, rtol = 1e-4) - if algorithm == :first_order_ + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] clear_solution_caches!(m, algorithm) DERIV5_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, @@ -2701,11 +2803,33 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) covariance = :all_excluding_obc)[:covariance], old_params) @test isapprox(deriv5_zyg[1], DERIV5_zyg[1], rtol = 1e-4) end + + clear_solution_caches!(m, algorithm) + + DERIV6 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + lyapunov_algorithm = lyapunov_algorithm, + sylvester_algorithm = sylvester_algorithm, + autocorrelation = :all_excluding_obc)[:autocorrelation], old_params) + @test isapprox(deriv6, DERIV6, rtol = 1e-4) + + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + clear_solution_caches!(m, algorithm) + + DERIV6_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + lyapunov_algorithm = lyapunov_algorithm, + sylvester_algorithm = sylvester_algorithm, + autocorrelation = :all_excluding_obc)[:autocorrelation], old_params) + @test isapprox(deriv6_zyg[1], DERIV6_zyg[1], rtol = 1e-4) + end + end end end end end - end end @@ -2806,36 +2930,75 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for parameter_derivatives in param_derivs - get_moments(m, - algorithm = algorithm, - non_stochastic_steady_state = true, - mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - parameter_derivatives = parameter_derivatives, - derivatives = true) + get_moments(m, + algorithm = algorithm, + non_stochastic_steady_state = true, + mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + parameter_derivatives = parameter_derivatives, + derivatives = true) end for variables in vars - get_moments(m, - algorithm = algorithm, - variables = variables, - non_stochastic_steady_state = true, - mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - derivatives = true) + get_moments(m, + algorithm = algorithm, + variables = variables, + non_stochastic_steady_state = true, + mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + derivatives = true) end for parameters in params - for derivatives in [true, false] + # derivatives=false: sweep all solver combos to verify numerical consistency + clear_solution_caches!(m, algorithm) + + moms = get_moments(m, + algorithm = algorithm, + parameters = parameters, + non_stochastic_steady_state = true, + mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + derivatives = false) + + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] + for quadratic_matrix_equation_algorithm in qme_algorithms + for sylvester_algorithm in sylvester_algorithms + for lyapunov_algorithm in lyapunov_algorithms + clear_solution_caches!(m, algorithm) + + MOMS = get_moments(m, + algorithm = algorithm, + parameters = parameters, + non_stochastic_steady_state = true, + mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + derivatives = false, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + lyapunov_algorithm = lyapunov_algorithm, + sylvester_algorithm = sylvester_algorithm) + + @test isapprox([v for (k,v) in moms], [v for (k,v) in MOMS], rtol = 1e-8) + end + end + end + end + + # derivatives=true: only test one representative solver combo (derivatives don't depend on solver choice) clear_solution_caches!(m, algorithm) - - moms = get_moments(m, + + moms_d = get_moments(m, algorithm = algorithm, parameters = parameters, non_stochastic_steady_state = true, @@ -2843,33 +3006,140 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - derivatives = derivatives) - - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] - for quadratic_matrix_equation_algorithm in qme_algorithms - for sylvester_algorithm in sylvester_algorithms - for lyapunov_algorithm in lyapunov_algorithms - clear_solution_caches!(m, algorithm) - - MOMS = get_moments(m, - algorithm = algorithm, - parameters = parameters, - non_stochastic_steady_state = true, - mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], - derivatives = derivatives, - tol = tol, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - lyapunov_algorithm = lyapunov_algorithm, - sylvester_algorithm = sylvester_algorithm) + derivatives = true) - @test isapprox([v for (k,v) in moms], [v for (k,v) in MOMS], rtol = 1e-8) - end - end + clear_solution_caches!(m, algorithm) + + MOMS_d = get_moments(m, + algorithm = algorithm, + parameters = parameters, + non_stochastic_steady_state = true, + mean = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + standard_deviation = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + variance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + covariance = algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order], + derivatives = true, + tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14)), + quadratic_matrix_equation_algorithm = :doubling, + lyapunov_algorithm = :doubling, + sylvester_algorithm = :doubling) + + @test isapprox([v for (k,v) in moms_d], [v for (k,v) in MOMS_d], rtol = 1e-8) + end + + # FD parity for get_moments derivative columns (rrule-based VJP Jacobians) + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + # NSSS derivatives + clear_solution_caches!(m, algorithm) + mom_nsss = get_moments(m, algorithm = algorithm, non_stochastic_steady_state = true, standard_deviation = false, derivatives = true) + nsss_jac = collect(mom_nsss[:non_stochastic_steady_state])[:, 2:end] + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + collect(get_moments(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, non_stochastic_steady_state = true, standard_deviation = false, derivatives = false)[:non_stochastic_steady_state]) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(nsss_jac, fd[1], rtol = 1e-5) + break + end + end + m.parameter_values .= old_params + + # Variance derivatives + clear_solution_caches!(m, algorithm) + mom_var = get_moments(m, algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, variance = true, derivatives = true) + var_jac = collect(mom_var[:variance])[:, 2:end] + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + collect(get_moments(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, variance = true, derivatives = false)[:variance]) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(var_jac, fd[1], rtol = 1e-4) + break + end + end + m.parameter_values .= old_params + + # Standard deviation derivatives + clear_solution_caches!(m, algorithm) + mom_std = get_moments(m, algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = true, variance = false, derivatives = true) + std_jac = collect(mom_std[:standard_deviation])[:, 2:end] + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + collect(get_moments(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = true, variance = false, derivatives = false)[:standard_deviation]) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(std_jac, fd[1], rtol = 1e-4) + break + end + end + m.parameter_values .= old_params + + # Covariance derivatives + clear_solution_caches!(m, algorithm) + mom_cov = get_moments(m, algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, covariance = true, + tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), + derivatives = true) + cov_ka = collect(mom_cov[:covariance]) + n_cv = size(cov_ka, 1) + cov_jac = reshape(cov_ka[:, :, 2:end], n_cv * n_cv, :) + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + vec(collect(get_moments(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, covariance = true, + tol = MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14), second_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14)), third_order = MacroModelling.HigherOrderTolerances(sylvester = MacroModelling.SolverTolerances(acceptance_tol = 1e-14), lyapunov = MacroModelling.SolverTolerances(acceptance_tol = 1e-14))), + derivatives = false)[:covariance])) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(cov_jac, fd[1], rtol = 1e-4) + break + end + end + m.parameter_values .= old_params + + # Mean derivatives (for algorithms that support it) + if algorithm โˆˆ [:pruned_second_order, :pruned_third_order] + clear_solution_caches!(m, algorithm) + mom_mean = get_moments(m, algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, mean = true, derivatives = true) + mean_jac = collect(mom_mean[:mean])[:, 2:end] + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + collect(get_moments(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, mean = true, derivatives = false)[:mean]) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(mean_jac, fd[1], rtol = 1e-4) + break end end + m.parameter_values .= old_params end end end @@ -2936,7 +3206,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) ignore_obc = true, initial_state = initial_state) - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for quadratic_matrix_equation_algorithm in qme_algorithms # for lyapunov_algorithm in lyapunov_algorithms for sylvester_algorithm in sylvester_algorithms @@ -2985,7 +3255,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] for parameters in params clear_solution_caches!(m, algorithm) @@ -3089,7 +3359,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for parameter_derivatives in param_derivs for parameters in params - for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(NSSS_xtol = 1e-14)] + for tol in [MacroModelling.Tolerances(),MacroModelling.Tolerances(nsss = MacroModelling.NsssTolerances(xtol = 1e-14))] clear_solution_caches!(m, algorithm) nsss = get_steady_state(m, @@ -3101,6 +3371,51 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end end end + + # FD parity for get_steady_state derivative columns (rrule-based VJP Jacobians) + # NSSS derivatives + clear_solution_caches!(m, algorithm) + nsss_d = get_steady_state(m, algorithm = algorithm, stochastic = false, derivatives = true, return_variables_only = true) + nsss_jac = collect(nsss_d)[:, 2:end] + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + collect(get_steady_state(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, stochastic = false, derivatives = false, return_variables_only = true)) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(nsss_jac, fd[1], rtol = 1e-5) + break + end + end + m.parameter_values .= old_params + + # Stochastic SS derivatives (non-first-order only) + if algorithm != :first_order + clear_solution_caches!(m, algorithm) + sss_d = get_steady_state(m, algorithm = algorithm, stochastic = true, derivatives = true, return_variables_only = true) + sss_jac = collect(sss_d)[:, 2:end] + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + collect(get_steady_state(m, + parameters = m.constants.post_complete_parameters.parameters .=> x, + algorithm = algorithm, stochastic = true, derivatives = false, return_variables_only = true)) + end, old_params) + if isfinite(โ„’.norm(fd[1])) + @test isapprox(sss_jac, fd[1], rtol = 1e-4) + break + end + end + m.parameter_values .= old_params + end end GC.gc() diff --git a/test/runtests.jl b/test/runtests.jl index bcc17b78e..013e57ceb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,9 +14,8 @@ import Zygote, FiniteDifferences, ForwardDiff import StatsPlots, Turing, Optim # has to come before Aqua, otherwise exports are not recognised using Aqua import LinearAlgebra as โ„’ -using CSV, DataFrames +using DelimitedFiles using Dates -using RuntimeGeneratedFunctions function quarterly_dates(start_date::Date, len::Int) dates = Vector{Date}(undef, len) @@ -368,10 +367,12 @@ if test_set == "plots_5" include("../models/Smets_Wouters_2007.jl") # load data - dat = CSV.read("data/usmodel.csv", DataFrame) + dat, header = readdlm("data/usmodel.csv", ',', header = true) + dat = Float64.(dat) + names = vec(Symbol.(strip.(header))) # load data - data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) + data = KeyedArray(dat', Variable = names, Time = axes(dat, 1)) # declare observables as written in csv file observables_old = [:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model @@ -432,12 +433,14 @@ if test_set == "plots_5" include("../models/FS2000.jl") # load data - dat = CSV.read("data/FS2000_data.csv", DataFrame) - dataFS2000 = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) + dat, header = readdlm("data/FS2000_data.csv", ',', header = true) + dat = Float64.(dat) + names = vec(header) + dataFS2000 = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) dataFS2000 = log.(dataFS2000) # declare observables - observables = sort(Symbol.("log_".*names(dat))) + observables = sort(Symbol.("log_".*names)) # subset observables in data dataFS2000 = dataFS2000(observables,:) @@ -796,16 +799,6 @@ if test_set == "basic" plots = false # test_higher_order = false - @testset verbose = true "Test equation filtering" begin - include("test_filter_equations.jl") - end - GC.gc() - - @testset verbose = true "Standalone functions" begin - include("test_standalone_function.jl") - end - GC.gc() - function rbc_steady_state(params) std_z, rho, delta, alpha, beta = params @@ -926,6 +919,16 @@ if test_set == "basic" @test isapprox(irf_nopre, irf_pre) end + @testset verbose = true "Test equation filtering" begin + include("test_filter_equations.jl") + end + GC.gc() + + @testset verbose = true "Standalone functions" begin + include("test_standalone_function.jl") + end + GC.gc() + @testset "Custom steady state assignment" begin @model RBC_switch begin 1 / c[0] = (beta / c[1]) * (alpha * exp(z[1]) * k[0]^(alpha - 1) + (1 - delta)) @@ -966,7 +969,7 @@ if test_set == "basic" @test inplace_calls[] > 0 @test isapprox(inplace_result(:,:Steady_state), rbc_steady_state(RBC_switch.parameter_values), rtol = 1e-10) expected_cache_length = length(RBC_switch.constants.post_model_macro.vars_in_ss_equations_no_aux) + length(RBC_switch.equations.calibration_parameters) - @test length(RBC_switch.workspaces.custom_steady_state_buffer) == expected_cache_length + @test length(RBC_switch.workspaces.custom_steady_state) == expected_cache_length # @test_throws ArgumentError get_steady_state(RBC_switch, steady_state_function = bad_ss) # @test bad_calls[] > 0 @@ -1005,16 +1008,18 @@ if test_set == "basic" beta = 0.95 end - @test !(RBC_macro_switch.functions.NSSS_solve isa RuntimeGeneratedFunction) + @test RBC_macro_switch.functions.NSSS_custom isa Function + @test RBC_macro_switch.constants.nsss_solver.n_steps == 0 _ = get_steady_state(RBC_macro_switch) @test macro_calls[] > 0 - @test !(RBC_macro_switch.functions.NSSS_solve isa RuntimeGeneratedFunction) + @test RBC_macro_switch.functions.NSSS_custom isa Function + @test RBC_macro_switch.constants.nsss_solver.n_steps == 0 MacroModelling.set_custom_steady_state_function!(RBC_macro_switch, nothing) _ = get_steady_state(RBC_macro_switch) @test isnothing(RBC_macro_switch.functions.NSSS_custom) - @test RBC_macro_switch.functions.NSSS_solve isa RuntimeGeneratedFunction + @test RBC_macro_switch.constants.nsss_solver.n_steps != 0 calls_before = macro_calls[] _ = get_steady_state(RBC_macro_switch) @@ -2165,15 +2170,17 @@ if test_set == "basic" end # write the parameters from NAWM_EAUS_2008 to a csv file - using CSV - using DataFrames - - df = DataFrame(Parameter = NAWM_EAUS_2008.constants.post_complete_parameters.parameters, Value = NAWM_EAUS_2008.parameter_values) - CSV.write("NAWM_EAUS_2008_parameters.csv", df) + open("NAWM_EAUS_2008_parameters.csv", "w") do io + println(io, "Parameter,Value") + for (param, val) in zip(NAWM_EAUS_2008.constants.post_complete_parameters.parameters, NAWM_EAUS_2008.parameter_values) + println(io, string(param), ",", val) + end + end # read the parameters from the csv file as a Dict and update NAWM_EAUS_2008_incomplete - param_df = CSV.read("NAWM_EAUS_2008_parameters.csv", DataFrame) - param_dict = Dict(row.Parameter => row.Value for row in eachrow(param_df)) + param_vals, param_header = readdlm("NAWM_EAUS_2008_parameters.csv", ',', header = true) + @assert vec(param_header) == ["Parameter", "Value"] + param_dict = Dict(Symbol(param_vals[i, 1]) => Float64(param_vals[i, 2]) for i in axes(param_vals, 1)) sol1 = get_solution(NAWM_EAUS_2008_incomplete, parameters = param_dict) sol2 = get_solution(NAWM_EAUS_2008) @@ -3341,7 +3348,7 @@ if test_set == "basic" end - @parameters RBC_CME symbolic = true verbose = true begin + @parameters RBC_CME ss_symbolic_mode = :full verbose = true begin # alpha | k[ss] / (4 * y[ss]) = cap_share # cap_share = 1.66 alpha = .157 @@ -3482,7 +3489,7 @@ if test_set == "basic" end - @parameters RBC_CME symbolic = true verbose = true begin + @parameters RBC_CME ss_symbolic_mode = :full verbose = true begin alpha | k[ss] / (4 * y[ss]) = cap_share cap_share = 1.66 # alpha = .157 diff --git a/test/test_1st_order_inversion_filter_estimation.jl b/test/test_1st_order_inversion_filter_estimation.jl index 1360cbcec..ae118c8d9 100644 --- a/test/test_1st_order_inversion_filter_estimation.jl +++ b/test/test_1st_order_inversion_filter_estimation.jl @@ -3,17 +3,19 @@ import Turing import Turing: NUTS, sample, logpdf import ADTypes: AutoZygote import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -61,13 +63,27 @@ modeFS2000i = Turing.maximum_a_posteriori(FS2000_loglikelihood_function(data, FS println("Mode variable values: $(modeFS2000i.values); Mode loglikelihood: $(modeFS2000i.lp)") +@testset "Zygote vs FiniteDifferences gradient (1st order inversion)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(FS2000, data, x, filter = :inversion), FS2000.parameter_values) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), x -> get_loglikelihood(FS2000, data, x, filter = :inversion), FS2000.parameter_values) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end # # estimate highly nonlinear model # # load data -# dat = CSV.read("data/usmodel.csv", DataFrame) +# dat, header = readdlm("data/usmodel.csv", ',', header = true) # data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) # # declare observables diff --git a/test/test_1st_order_inversion_filter_estimation_pigeons.jl b/test/test_1st_order_inversion_filter_estimation_pigeons.jl index b03365c7d..3aae84afb 100644 --- a/test/test_1st_order_inversion_filter_estimation_pigeons.jl +++ b/test/test_1st_order_inversion_filter_estimation_pigeons.jl @@ -3,18 +3,20 @@ using Test import Turing import Pigeons import Turing: logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -52,6 +54,7 @@ end FS2000_lp = Pigeons.TuringLogPotential(FS2000_loglikelihood_function(data, FS2000, :inversion, -floatmax(Float64)+1e10)) #, verbose = true)) init_params = FS2000.parameter_values +const PIGEONS_SEED = 30 const FS2000_LP = typeof(FS2000_lp) @@ -64,12 +67,13 @@ function Pigeons.initialization(target::FS2000_LP, rng::AbstractRNG, _::Int64) return result end -pt = Pigeons.pigeons(target = FS2000_lp, n_rounds = 0, n_chains = 1) +pt = Pigeons.pigeons(target = FS2000_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) pt = @time Pigeons.pigeons(target = FS2000_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], n_chains = 2, n_rounds = 10, + seed = PIGEONS_SEED, multithreaded = false) # tests fail on multithreaded samps = MCMCChains.Chains(pt) diff --git a/test/test_2nd_order_estimation.jl b/test/test_2nd_order_estimation.jl index 412d51385..a2ae6a59c 100644 --- a/test/test_2nd_order_estimation.jl +++ b/test/test_2nd_order_estimation.jl @@ -3,17 +3,19 @@ import Turing import ADTypes: AutoZygote import Turing: NUTS, sample, logpdf import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -44,7 +46,7 @@ end Random.seed!(30) -n_samples = 500 +n_samples = 1000 samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :second_order, -Inf), NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) @@ -53,13 +55,27 @@ println("Mean variable values (Zygote): $(mean(samps).nt.mean)") sample_nuts = mean(samps).nt.mean +@testset "Zygote vs FiniteDifferences gradient (2nd order)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(FS2000, data, x, algorithm = :second_order), FS2000.parameter_values) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), x -> get_loglikelihood(FS2000, data, x, algorithm = :second_order), FS2000.parameter_values) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end # # estimate highly nonlinear model # # load data -# dat = CSV.read("data/usmodel.csv", DataFrame) +# dat, header = readdlm("data/usmodel.csv", ',', header = true) # data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) # # declare observables diff --git a/test/test_2nd_order_estimation_pigeons.jl b/test/test_2nd_order_estimation_pigeons.jl index 5246f80f1..70b16bdce 100644 --- a/test/test_2nd_order_estimation_pigeons.jl +++ b/test/test_2nd_order_estimation_pigeons.jl @@ -3,18 +3,20 @@ using Test import Turing import Pigeons import Turing: logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -45,7 +47,7 @@ Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_fail end -Random.seed!(30) +const PIGEONS_SEED = 30 # generate a Pigeons log potential FS2000_2nd_lp = Pigeons.TuringLogPotential(FS2000_loglikelihood_function(data, FS2000, :second_order, -floatmax(Float64)+1e10)) @@ -66,9 +68,9 @@ if isfinite(LLH) return result end - pt = Pigeons.pigeons(target = FS2000_2nd_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = FS2000_2nd_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) else - pt = Pigeons.pigeons(target = FS2000_2nd_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = FS2000_2nd_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) replica = pt.replicas[end] XMAX = deepcopy(replica.state) LPmax = FS2000_2nd_lp(XMAX) @@ -93,6 +95,7 @@ pt = @time Pigeons.pigeons(target = FS2000_2nd_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], n_chains = 1, n_rounds = 9, + seed = PIGEONS_SEED, multithreaded = false) # tests fail on multithreaded samps = MCMCChains.Chains(pt) diff --git a/test/test_3rd_order_estimation.jl b/test/test_3rd_order_estimation.jl index bd82a5758..d1bfd92eb 100644 --- a/test/test_3rd_order_estimation.jl +++ b/test/test_3rd_order_estimation.jl @@ -3,13 +3,15 @@ import Turing import ADTypes: AutoZygote import Turing: NUTS, sample, logpdf, PG, IS import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys # estimate highly nonlinear model # load data -dat = CSV.read("data/usmodel.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) +dat, header = readdlm("data/usmodel.csv", ',', header = true) +dat = Float64.(dat) +names = vec(Symbol.(strip.(header))) +data = KeyedArray(dat', Variable = names, Time = axes(dat, 1)) # declare observables observables = [:dy]#, :dinve, :labobs, :pinfobs, :dw, :robs] @@ -81,19 +83,32 @@ println("Mode variable values (L-BFGS): $init_params") n_samples = 100 -samps = sample(Caldara_et_al_2012_loglikelihood, NUTS(250, 0.65, adtype = AutoZygote()), n_samples, progress = true, initial_params = init_params) +samps = sample(Caldara_et_al_2012_loglikelihood, NUTS(1000, 0.65, adtype = AutoZygote()), n_samples, progress = true, initial_params = init_params) println("Mean variable values (Zygote): $(mean(samps).nt.mean)") sample_nuts = mean(samps).nt.mean - +@testset "Zygote vs FiniteDifferences gradient (3rd order)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(Caldara_et_al_2012_estim, data, x, algorithm = :third_order), init_params) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1, max_range = 1e-3), x -> get_loglikelihood(Caldara_et_al_2012_estim, data, x, algorithm = :third_order), init_params) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end # include("../models/FS2000.jl") # # load data -# dat = CSV.read("data/FS2000_data.csv", DataFrame) +# dat, header = readdlm("data/FS2000_data.csv", ',', header = true) # data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) # data = log.(data) diff --git a/test/test_3rd_order_estimation_pigeons.jl b/test/test_3rd_order_estimation_pigeons.jl index 164fa6a10..ab7d46788 100644 --- a/test/test_3rd_order_estimation_pigeons.jl +++ b/test/test_3rd_order_estimation_pigeons.jl @@ -3,14 +3,16 @@ using Test import Turing import Pigeons import Turing: logpdf, PG, IS -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL # estimate highly nonlinear model # load data -dat = CSV.read("data/usmodel.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) +dat, header = readdlm("data/usmodel.csv", ',', header = true) +dat = Float64.(dat) +names = vec(Symbol.(strip.(header))) +data = KeyedArray(dat', Variable = names, Time = axes(dat, 1)) # declare observables observables = [:dy]#, :dinve, :labobs, :pinfobs, :dw, :robs] @@ -57,7 +59,7 @@ Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_fai end -Random.seed!(3) +const PIGEONS_SEED = 3 Caldara_et_al_2012_loglikelihood = Caldara_et_al_2012_loglikelihood_function(data, Caldara_et_al_2012_estim, -Inf) @@ -82,9 +84,9 @@ if isfinite(LLH) return result end - pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) else - pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) replica = pt.replicas[end] XMAX = deepcopy(replica.state) @@ -108,8 +110,9 @@ end pt = @time Pigeons.pigeons(target = Caldara_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], - n_chains = 1, + n_chains = 4, n_rounds = 8, + seed = PIGEONS_SEED, multithreaded = false) # tests fail on multithreaded samps = MCMCChains.Chains(pt) diff --git a/test/test_estimation.jl b/test/test_estimation.jl index 42e14144d..64c51cce1 100644 --- a/test/test_estimation.jl +++ b/test/test_estimation.jl @@ -3,18 +3,20 @@ import Turing import ADTypes: AutoZygote import Turing: NUTS, sample, logpdf import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import Zygote include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -90,6 +92,21 @@ println("Mode variable values: $(modeFS2000.values); Mode loglikelihood: $(modeF @test isapprox(sample_nuts, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) end +@testset "Zygote vs FiniteDifferences gradient (1st order Kalman)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(FS2000, data, x), FS2000.parameter_values) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), x -> get_loglikelihood(FS2000, data, x), FS2000.parameter_values) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end + plot_model_estimates(FS2000, data, parameters = sample_nuts) plot_shock_decomposition(FS2000, data) diff --git a/test/test_estimation_pigeons.jl b/test/test_estimation_pigeons.jl index 634eb2e91..374f41a02 100644 --- a/test/test_estimation_pigeons.jl +++ b/test/test_estimation_pigeons.jl @@ -3,18 +3,20 @@ using Test import Turing, Pigeons import ADTypes: AutoZygote import Turing: NUTS, sample, logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -53,6 +55,7 @@ end FS2000_lp = Pigeons.TuringLogPotential(FS2000_loglikelihood_function(data, FS2000, -floatmax(Float64)+1e10)) init_params = FS2000.parameter_values +const PIGEONS_SEED = 30 const FS2000_LP = typeof(FS2000_lp) @@ -64,12 +67,13 @@ function Pigeons.initialization(target::FS2000_LP, rng::AbstractRNG, _::Int64) return result end -pt = Pigeons.pigeons(target = FS2000_lp, n_rounds = 0, n_chains = 1) +pt = Pigeons.pigeons(target = FS2000_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) pt = @time Pigeons.pigeons(target = FS2000_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], n_chains = 1, n_rounds = 10, + seed = PIGEONS_SEED, multithreaded = false) # tests fail on multithreaded samps = MCMCChains.Chains(pt) diff --git a/test/test_pruned_2nd_order_estimation.jl b/test/test_pruned_2nd_order_estimation.jl index cb99e67ea..12cc02ac4 100644 --- a/test/test_pruned_2nd_order_estimation.jl +++ b/test/test_pruned_2nd_order_estimation.jl @@ -3,17 +3,19 @@ import Turing import ADTypes: AutoZygote import Turing: NUTS, sample, logpdf import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -49,7 +51,7 @@ end Random.seed!(30) -n_samples = 500 +n_samples = 1000 samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :pruned_second_order, -Inf), NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) @@ -58,12 +60,27 @@ println("Mean variable values (Zygote): $(mean(samps).nt.mean)") sample_nuts = mean(samps).nt.mean +@testset "Zygote vs FiniteDifferences gradient (pruned 2nd order)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(FS2000, data, x, algorithm = :pruned_second_order), FS2000.parameter_values) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), x -> get_loglikelihood(FS2000, data, x, algorithm = :pruned_second_order), FS2000.parameter_values) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end + # # estimate highly nonlinear model # # load data -# dat = CSV.read("data/usmodel.csv", DataFrame) +# dat, header = readdlm("data/usmodel.csv", ',', header = true) # data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) # # declare observables diff --git a/test/test_pruned_2nd_order_estimation_pigeons.jl b/test/test_pruned_2nd_order_estimation_pigeons.jl index 1e1087cf3..a3807a449 100644 --- a/test/test_pruned_2nd_order_estimation_pigeons.jl +++ b/test/test_pruned_2nd_order_estimation_pigeons.jl @@ -3,18 +3,20 @@ using Test import Turing import Pigeons import Turing: logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL include("../models/FS2000.jl") # load data -dat = CSV.read("data/FS2000_data.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) +dat, header = readdlm("data/FS2000_data.csv", ',', header = true) +dat = Float64.(dat) +names = vec(header) +data = KeyedArray(dat', Variable = Symbol.("log_".*names), Time = axes(dat, 1)) data = log.(data) # declare observables -observables = sort(Symbol.("log_".*names(dat))) +observables = sort(Symbol.("log_".*names)) # subset observables in data data = data(observables,:) @@ -50,7 +52,7 @@ Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_fail end -Random.seed!(30) +const PIGEONS_SEED = 30 # generate a Pigeons log potential FS2000_pruned2nd_lp = Pigeons.TuringLogPotential(FS2000_loglikelihood_function(data, FS2000, :pruned_second_order, -floatmax(Float64)+1e10)) #, verbose = true)) @@ -71,9 +73,9 @@ if isfinite(LLH) return result end - pt = Pigeons.pigeons(target = FS2000_pruned2nd_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = FS2000_pruned2nd_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) else - pt = Pigeons.pigeons(target = FS2000_pruned2nd_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = FS2000_pruned2nd_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) replica = pt.replicas[end] XMAX = deepcopy(replica.state) LPmax = FS2000_pruned2nd_lp(XMAX) @@ -98,6 +100,7 @@ pt = @time Pigeons.pigeons(target = FS2000_pruned2nd_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], n_chains = 1, n_rounds = 8, + seed = PIGEONS_SEED, multithreaded = false) samps = MCMCChains.Chains(pt) diff --git a/test/test_pruned_3rd_order_estimation.jl b/test/test_pruned_3rd_order_estimation.jl index 1956cff99..848435edb 100644 --- a/test/test_pruned_3rd_order_estimation.jl +++ b/test/test_pruned_3rd_order_estimation.jl @@ -3,13 +3,15 @@ import Turing import ADTypes: AutoZygote import Turing: NUTS, sample, logpdf, PG, IS import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys # estimate highly nonlinear model # load data -dat = CSV.read("data/usmodel.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) +dat, header = readdlm("data/usmodel.csv", ',', header = true) +dat = Float64.(dat) +names = vec(Symbol.(strip.(header))) +data = KeyedArray(dat', Variable = names, Time = axes(dat, 1)) # declare observables observables = [:dy]#, :dinve, :labobs, :pinfobs, :dw, :robs] @@ -89,20 +91,35 @@ println("Mode variable values (L-BFGS): $init_params") n_samples = 100 -samps = @time sample(Caldara_et_al_2012_loglikelihood, NUTS(250, 0.65, adtype = AutoZygote()), n_samples, progress = true, initial_params = init_params) +samps = @time sample(Caldara_et_al_2012_loglikelihood, NUTS(1000, 0.65, adtype = AutoZygote()), n_samples, progress = true, initial_params = init_params) println("Mean variable values (Zygote): $(mean(samps).nt.mean)") sample_nuts = mean(samps).nt.mean +@testset "Zygote vs FiniteDifferences gradient (pruned 3rd order)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(Caldara_et_al_2012_estim, data, x, algorithm = :pruned_third_order), init_params) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1, max_range = 1e-3), x -> get_loglikelihood(Caldara_et_al_2012_estim, data, x, algorithm = :pruned_third_order), init_params) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end + # include("../models/FS2000.jl") # # load data -# dat = CSV.read("data/FS2000_data.csv", DataFrame) +# dat, header = readdlm("data/FS2000_data.csv", ',', header = true) # data = KeyedArray(Array(dat)',Variable = Symbol.("log_".*names(dat)),Time = 1:size(dat)[1]) # data = log.(data) diff --git a/test/test_pruned_3rd_order_estimation_pigeons.jl b/test/test_pruned_3rd_order_estimation_pigeons.jl index 7570e2bb6..60adc82b5 100644 --- a/test/test_pruned_3rd_order_estimation_pigeons.jl +++ b/test/test_pruned_3rd_order_estimation_pigeons.jl @@ -3,14 +3,16 @@ using Test import Turing import Pigeons import Turing: logpdf, PG, IS -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL # estimate highly nonlinear model # load data -dat = CSV.read("data/usmodel.csv", DataFrame) -data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) +dat, header = readdlm("data/usmodel.csv", ',', header = true) +dat = Float64.(dat) +names = vec(Symbol.(strip.(header))) +data = KeyedArray(dat', Variable = names, Time = axes(dat, 1)) # declare observables observables = [:dy]#, :dinve, :labobs, :pinfobs, :dw, :robs] @@ -62,7 +64,7 @@ Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_fai end -Random.seed!(3) +const PIGEONS_SEED = 3 Caldara_et_al_2012_loglikelihood = Caldara_et_al_2012_loglikelihood_function(data, Caldara_et_al_2012_estim, -Inf) @@ -90,9 +92,9 @@ if isfinite(LLH) return result end - pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) else - pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1) + pt = Pigeons.pigeons(target = Caldara_lp, n_rounds = 0, n_chains = 1, seed = PIGEONS_SEED) replica = pt.replicas[end] XMAX = deepcopy(replica.state) @@ -116,8 +118,9 @@ end pt = @time Pigeons.pigeons(target = Caldara_lp, record = [Pigeons.traces; Pigeons.round_trip; Pigeons.record_default()], - n_chains = 1, + n_chains = 4, n_rounds = 8, + seed = PIGEONS_SEED, multithreaded = false) # tests fail on multithreaded samps = MCMCChains.Chains(pt) diff --git a/test/test_standalone_function.jl b/test/test_standalone_function.jl index 7094ba2f5..5c5c434d0 100644 --- a/test/test_standalone_function.jl +++ b/test/test_standalone_function.jl @@ -2,7 +2,7 @@ using SparseArrays using MacroModelling using Random using Test -import MacroModelling: post_model_macro, get_NSSS_and_parameters, ensure_qme_workspace!, ensure_sylvester_1st_order_workspace! +import MacroModelling: post_model_macro, get_NSSS_and_parameters using ForwardDiff import LinearAlgebra as โ„’ using FiniteDifferences, Zygote @@ -66,23 +66,19 @@ get_irf(RBC_CME, algorithm = :third_order) get_irf(RBC_CME, algorithm = :pruned_third_order) get_irf(RBC_CME, algorithm = :pruned_second_order) -โˆ‡โ‚ = calculate_jacobian(RBC_CME.parameter_values, SS_and_pars, RBC_CME.caches, RBC_CME.functions.jacobian)# |> Matrix -โˆ‡โ‚‚ = calculate_hessian(RBC_CME.parameter_values, SS_and_pars, RBC_CME.caches, RBC_CME.functions.hessian)# * RBC_CME.constants.second_order.๐”โˆ‡โ‚‚ -โˆ‡โ‚ƒ = calculate_third_order_derivatives(RBC_CME.parameter_values, SS_and_pars, RBC_CME.caches, RBC_CME.functions.third_order_derivatives)# * RBC_CME.constants.third_order.๐”โˆ‡โ‚ƒ +โˆ‡โ‚ = calculate_jacobian(RBC_CME.parameter_values, SS_and_pars, RBC_CME.caches, RBC_CME.functions.jacobian, RBC_CME.workspaces)# |> Matrix +โˆ‡โ‚‚ = calculate_hessian(RBC_CME.parameter_values, SS_and_pars, RBC_CME.caches, RBC_CME.functions.hessian, RBC_CME.workspaces)# * RBC_CME.constants.second_order.๐”โˆ‡โ‚‚ +โˆ‡โ‚ƒ = calculate_third_order_derivatives(RBC_CME.parameter_values, SS_and_pars, RBC_CME.caches, RBC_CME.functions.third_order_derivatives, RBC_CME.workspaces)# * RBC_CME.constants.third_order.๐”โˆ‡โ‚ƒ #SS = get_steady_state(RBC_CME, derivatives = false) T = RBC_CME.constants.post_model_macro -qme_ws = ensure_qme_workspace!(RBC_CME) -sylv_ws = ensure_sylvester_1st_order_workspace!(RBC_CME) -first_order_solution, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, RBC_CME.constants, qme_ws, sylv_ws)# |> Matrix{Float32} +first_order_solution, qme_sol, solved = calculate_first_order_solution(โˆ‡โ‚, RBC_CME.constants, RBC_CME.workspaces, RBC_CME.caches)# |> Matrix{Float32} -second_order_solution, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, first_order_solution, RBC_CME.constants, RBC_CME.workspaces) +second_order_solution, solved2 = calculate_second_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, first_order_solution, RBC_CME.constants, RBC_CME.workspaces, RBC_CME.caches) - -# second_order_solution *= RBC_CME.constants.second_order_auxiliary_matrices.๐”โ‚‚ - -second_order_solution = sparse(second_order_solution * RBC_CME.constants.second_order.๐”โ‚‚) +# second_order_solution is now compressed (bโ‚‚ columns); pass compressed to third-order +# (both functions expand internally) third_order_solution, solved3 = calculate_third_order_solution(โˆ‡โ‚, โˆ‡โ‚‚, @@ -90,7 +86,11 @@ third_order_solution, solved3 = calculate_third_order_solution(โˆ‡โ‚, first_order_solution, second_order_solution, RBC_CME.constants, - RBC_CME.workspaces) + RBC_CME.workspaces, + RBC_CME.caches) + +# Expand second_order_solution to full space for comparison +second_order_solution = sparse(second_order_solution * RBC_CME.constants.second_order.๐”โ‚‚) # third_order_solution *= RBC_CME.constants.third_order_auxiliary_matrices.๐”โ‚ƒ @@ -164,7 +164,7 @@ third_order_solution = sparse(third_order_solution * RBC_CME.constants.third_ord -0.0226 0.0021014511165327685 -0.0021014511165327685],7,225) - @test isapprox(โˆ‡โ‚‚,hessian2,rtol = eps(Float32)) + @test isapprox(โˆ‡โ‚‚ * RBC_CME.constants.second_order.๐”โˆ‡โ‚‚, hessian2,rtol = eps(Float32)) third_order_derivatives2 = sparse(vec([ 2 2 2 2 3 3 3 3 3 3 3 3 2 2 3 3 3 2 3 2 3 3 2 3 3 2 2 2 1 5 4 3 3 3 3 2 3 2 2 2 2 2 2 2 2 1 5 1 5 1 5]), @@ -564,7 +564,7 @@ end [0,0.95,0,0], [1,1,1,2], [.16, .999,.022,1], Optim.Fminbox(Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3))); autodiff = :forward) - get_statistics(RBC_CME, sol.minimizer, parameters = RBC_CME.constants.post_complete_parameters.parameters[1:4], mean = RBC_CME.constants.post_model_macro.var[[4,6]], standard_deviation = RBC_CME.constants.post_model_macro.var[4:5], autocorrelation = RBC_CME.constants.post_model_macro.var[[3,5]], autocorrelation_periods = 1:1, algorithm = :pruned_third_order) + out = get_statistics(RBC_CME, sol.minimizer, parameters = RBC_CME.constants.post_complete_parameters.parameters[1:4], mean = RBC_CME.constants.post_model_macro.var[[4,6]], standard_deviation = RBC_CME.constants.post_model_macro.var[4:5], autocorrelation = RBC_CME.constants.post_model_macro.var[[3,5]], autocorrelation_periods = 1:1, algorithm = :pruned_third_order) @test isapprox([out[:mean], out[:standard_deviation], out[:autocorrelation], sol.minimizer[3]], [[1.2,1.4],[.013,.2],[.955,.997][:,:],.0215], diff --git a/test/test_sw07_estimation.jl b/test/test_sw07_estimation.jl index 46363a0ed..55a34f5f8 100644 --- a/test/test_sw07_estimation.jl +++ b/test/test_sw07_estimation.jl @@ -3,13 +3,15 @@ import ADTypes: AutoZygote import Turing import Turing: NUTS, sample, logpdf import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys # load data -dat = CSV.read("data/usmodel.csv", DataFrame) +dat, header = readdlm("data/usmodel.csv", ',', header = true) +dat = Float64.(dat) +names = vec(Symbol.(strip.(header))) # load data -data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) +data = KeyedArray(dat', Variable = names, Time = axes(dat, 1)) # declare observables as written in csv file observables_old = [:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model @@ -119,6 +121,21 @@ samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoZygote()), n_s println(samps) println("Mean variable values (linear): $(mean(samps).nt.mean)") +@testset "Zygote vs FiniteDifferences gradient (SW07 linear)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(Smets_Wouters_2007_linear, data(observables), x, presample_periods = 4, initial_covariance = :diagonal, filter = :kalman), Smets_Wouters_2007_linear.parameter_values) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), x -> get_loglikelihood(Smets_Wouters_2007_linear, data(observables), x, presample_periods = 4, initial_covariance = :diagonal, filter = :kalman), Smets_Wouters_2007_linear.parameter_values) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end + # estimate nonlinear model include("../models/Smets_Wouters_2007.jl") @@ -154,4 +171,19 @@ samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoZygote()), n_s progress = true) println(samps) -println("Mean variable values (nonlinear): $(mean(samps).nt.mean)") \ No newline at end of file +println("Mean variable values (nonlinear): $(mean(samps).nt.mean)") + +@testset "Zygote vs FiniteDifferences gradient (SW07 nonlinear)" begin + back_grad = Zygote.gradient(x -> get_loglikelihood(Smets_Wouters_2007, data(observables), x, presample_periods = 4, initial_covariance = :diagonal, filter = :kalman), Smets_Wouters_2007.parameter_values) + @test !isnothing(back_grad[1]) + @test all(isfinite, back_grad[1]) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), x -> get_loglikelihood(Smets_Wouters_2007, data(observables), x, presample_periods = 4, initial_covariance = :diagonal, filter = :kalman), Smets_Wouters_2007.parameter_values) + if isfinite(โ„’.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + break + end + end +end \ No newline at end of file