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/copilot-instructions.md b/.github/copilot-instructions.md index d41677b1c..a87f68adb 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -2,4 +2,4 @@ See [AGENTS.md](../AGENTS.md) for the complete agent guide. -This file exists for GitHub Copilot compatibility. All instructions are maintained in AGENTS.md. +This file exists for GitHub Copilot compatibility. All instructions are maintained in AGENTS.md. \ No newline at end of file diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index a9d066651..56d264458 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -15,6 +15,8 @@ jobs: steps: - uses: actions/checkout@v6 + with: + ref: ${{ github.event_name == 'pull_request_target' && github.event.pull_request.head.sha || github.sha }} - uses: julia-actions/setup-julia@v2 with: version: "1" @@ -39,7 +41,7 @@ jobs: echo $PATH ls -l ~/.julia/bin mkdir results - benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --url=${{ github.event.repository.clone_url }} --bench-on="${{github.event.repository.default_branch}}" --output-dir=results/ + benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --url=${{ github.event.repository.clone_url }} --bench-on="${{ github.event.pull_request.head.sha }}" --script="$PWD/benchmark/benchmarks.jl" --add="MatrixEquations" --output-dir=results/ - name: Create plots from benchmarks run: | mkdir -p plots diff --git a/.github/workflows/benchmark_push.yml b/.github/workflows/benchmark_push.yml index 60ccfccac..5be9bc737 100644 --- a/.github/workflows/benchmark_push.yml +++ b/.github/workflows/benchmark_push.yml @@ -36,11 +36,11 @@ jobs: echo $PATH ls -l ~/.julia/bin mkdir results - benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="v0.1.39,v0.1.40,v0.1.41,$GITHUB_SHA" --url=${{ github.event.repository.clone_url }} --bench-on="$GITHUB_SHA" --output-dir=results/ + benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="v0.1.39,v0.1.40,v0.1.41,v0.1.46,$GITHUB_SHA" --url=${{ github.event.repository.clone_url }} --bench-on="$GITHUB_SHA" --script="$PWD/benchmark/benchmarks.jl" --add="MatrixEquations" --output-dir=results/ - name: Create plots from benchmarks run: | mkdir -p plots - benchpkgplot ${{ steps.extract-package-name.outputs.package_name }} --rev="v0.1.39,v0.1.40,v0.1.41,$GITHUB_SHA" --npart=10 --format=png --input-dir=results/ --output-dir=plots/ + benchpkgplot ${{ steps.extract-package-name.outputs.package_name }} --rev="v0.1.39,v0.1.40,v0.1.41,v0.1.46,$GITHUB_SHA" --npart=10 --format=png --input-dir=results/ --output-dir=plots/ - name: Upload plot as artifact uses: actions/upload-artifact@v7 with: @@ -48,7 +48,7 @@ jobs: path: plots - name: Create markdown table from benchmarks run: | - benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="v0.1.39,v0.1.40,v0.1.41,$GITHUB_SHA" --input-dir=results/ --ratio --mode "time,memory" > table.md + benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="v0.1.39,v0.1.40,v0.1.41,v0.1.46,$GITHUB_SHA" --input-dir=results/ --ratio --mode "time,memory" > table.md echo '### Benchmark Results' > body.md echo '' >> body.md echo '' >> body.md diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5aa841395..559a0269f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -129,10 +129,10 @@ jobs: os: ubuntu-latest arch: x64 test_set: "pruned_3rd_order_estimation_pigeons" - - version: 'min' - os: ubuntu-latest - arch: x64 - test_set: "basic" + # - version: 'min' + # os: ubuntu-latest + # arch: x64 + # test_set: "basic" - version: 'lts' os: ubuntu-latest arch: x64 @@ -151,10 +151,28 @@ jobs: os: ubuntu-latest arch: x64 test_set: "jet" + allow_failure: true - version: 'lts' os: macOS-latest arch: x64 test_set: "jet" + allow_failure: true + - version: '1' + os: ubuntu-latest + arch: x64 + test_set: "jet_hot_paths" + - version: '1' + os: ubuntu-latest + arch: x64 + test_set: "estimate_sw07_nested_sampling" + - version: '1' + os: ubuntu-latest + arch: x64 + test_set: "gradient_checks" + - version: '1' + os: ubuntu-latest + arch: x64 + test_set: "system_prior_estimation" steps: - uses: actions/checkout@v6 - uses: julia-actions/setup-julia@v2 @@ -168,18 +186,68 @@ jobs: sed -i.bak \ -e '/^\[targets\]/,$ s/,[[:space:]]*"Pigeons"//g' \ -e '/^\[targets\]/,$ s/"Pigeons",[[:space:]]*//g' \ + -e '/^\[compat\]/,/^\[/ s/^Pigeons[[:space:]]*=.*$//g' \ + -e '/^\[extras\]/,/^\[/ s/^Pigeons[[:space:]]*=.*$//g' \ + Project.toml + rm -f Project.toml.bak + + - name: Remove MCMCChains from non-pigeons runs + if: contains(matrix.test_set, 'pigeons') == false + shell: bash + run: | + sed -i.bak \ + -e '/^\[targets\]/,$ s/,[[:space:]]*"MCMCChains"//g' \ + -e '/^\[targets\]/,$ s/"MCMCChains",[[:space:]]*//g' \ + -e '/^\[compat\]/,/^\[/ s/^MCMCChains[[:space:]]*=.*$//g' \ + -e '/^\[extras\]/,/^\[/ s/^MCMCChains[[:space:]]*=.*$//g' \ 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: Restrict DynamicPPL to 0.35 for pigeons runs + if: contains(matrix.test_set, 'pigeons') shell: bash run: | + sed -i.bak \ + -e '/^\[compat\]/,/^\[/ s/^DynamicPPL[[:space:]]*=.*$/DynamicPPL = "0.35"/g' \ + Project.toml + rm -f Project.toml.bak + + - name: Remove Mooncake from pigeons runs + if: contains(matrix.test_set, 'pigeons') + shell: bash + run: | + sed -i.bak \ + -e '/^\[targets\]/,$ s/,[[:space:]]*"Mooncake"//g' \ + -e '/^\[targets\]/,$ s/"Mooncake",[[:space:]]*//g' \ + -e '/^\[compat\]/,/^\[/ s/^Mooncake[[:space:]]*=.*$//g' \ + -e '/^\[extras\]/,/^\[/ s/^Mooncake[[:space:]]*=.*$//g' \ + Project.toml + rm -f Project.toml.bak + + - name: Remove FlexiChains from pigeons runs + if: contains(matrix.test_set, 'pigeons') + shell: bash + run: | + sed -i.bak \ + -e '/^\[targets\]/,$ s/,[[:space:]]*"FlexiChains"//g' \ + -e '/^\[targets\]/,$ s/"FlexiChains",[[:space:]]*//g' \ + -e '/^\[compat\]/,/^\[/ s/^FlexiChains[[:space:]]*=.*$//g' \ + -e '/^\[extras\]/,/^\[/ s/^FlexiChains[[:space:]]*=.*$//g' \ + Project.toml + rm -f Project.toml.bak + + - name: Remove JET from non-jet runs + if: "!startsWith(matrix.test_set, 'jet')" + shell: bash + run: | + # `JET` is only needed for the dedicated `jet` rows. Keeping it in the + # shared test target adds resolver constraints to unrelated jobs. sed -i.bak \ -e '/^\[targets\]/,$ s/,[[:space:]]*"JET"//g' \ -e '/^\[targets\]/,$ s/"JET",[[:space:]]*//g' \ Project.toml - rm -f Project.toml.bak + rm -f Project.toml.bak + - name: Set Custom Test Environment Variable (Windows) if: matrix.os == 'windows-latest' run: echo "TEST_SET=${{ matrix.test_set }}" | Out-File -Append -FilePath $env:GITHUB_ENV -Encoding utf8 @@ -199,3 +267,32 @@ jobs: with: token: ${{ secrets.CODECOV_TOKEN }} file: lcov.info + + dynare-comparison: + name: Dynare Comparison + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: julia-actions/setup-julia@v2 + with: + version: '1' + - uses: julia-actions/cache@v3 + - uses: julia-actions/julia-buildpkg@v1 + + # Phase 1: Generate Julia results and .mod files + - name: Generate Julia results + env: + JULIA_NUM_THREADS: '1' + run: julia --threads=1 --project=. test/dynare_comparison/generate_julia_results.jl + + # Phase 2: Run Dynare in Docker (debian:testing) + - name: Build Dynare Docker image + run: docker build -t dynare-runner test/dynare_comparison/ + - name: Run Dynare models + run: docker run --rm --user "$(id -u):$(id -g)" -e THREAD_COUNT=1 -e OMP_NUM_THREADS=1 -e OMP_THREAD_LIMIT=1 -e OPENBLAS_NUM_THREADS=1 -e MKL_NUM_THREADS=1 -v "${{ github.workspace }}/test/dynare_comparison/output:/work/output" dynare-runner + + # Phase 3: Compare results + - name: Compare Julia vs Dynare results + env: + JULIA_NUM_THREADS: '1' + run: julia --threads=1 --project=. test/dynare_comparison/compare_results.jl diff --git a/.gitignore b/.gitignore index 5141f731c..6a16558ac 100644 --- a/.gitignore +++ b/.gitignore @@ -76,4 +76,10 @@ 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 +tasks/ +test/dynare_comparison/_dynare_work/ +test/dynare_comparison/output* +test/dynare_comparison/output_thread_sweep* diff --git a/AGENTS.md b/AGENTS.md index db0fcca89..4952a0ead 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,340 +1,59 @@ # 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). +- Performance-critical code should live inside functions, not global scope. +- Avoid untyped global variables and abstractly typed containers in hot code paths. +- 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 for control flow. Use explicit checks and validation; reserve try-catch for unavoidable numerical failures. +- **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 +All companion guides live in `docs/agent-guides/`: -### 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` +- Code style conventions, naming, formatting, performance patterns: `docs/agent-guides/STYLE_GUIDE.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/CLAUDE.md b/CLAUDE.md index f89fb040b..41123a425 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -3,3 +3,12 @@ See [AGENTS.md](AGENTS.md) for the complete agent guide. This file exists for Claude Code (claude.ai/code) compatibility. All instructions are maintained in AGENTS.md. + +## Companion Guides + +Detailed references in `docs/agent-guides/`: + +- [development-workflow.md](docs/agent-guides/development-workflow.md) — Julia setup, Revise workflow, testing, docs, benchmarking +- [project-context.md](docs/agent-guides/project-context.md) — Project overview, repo structure, model syntax, design context +- [task-runbook.md](docs/agent-guides/task-runbook.md) — Task flows, orchestration heuristics, common change points +- [STYLE_GUIDE.md](docs/agent-guides/STYLE_GUIDE.md) — Code style conventions, naming, formatting, performance patterns diff --git a/Project.toml b/Project.toml index 4ebce38e3..fbaaab41b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "MacroModelling" uuid = "687ffad2-3618-405e-ac50-e0f7b9c75e44" -authors = ["Thore Kockerols "] version = "0.1.46" +authors = ["Thore Kockerols "] [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" @@ -11,11 +11,10 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" -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" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +FastLapackInterface = "29a986be-02c6-4525-aec4-84b980013641" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" @@ -24,13 +23,11 @@ LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" -MatrixEquations = "99c1a7ee-ab34-5fd5-8076-27c950a045f4" NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd" 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" @@ -44,11 +41,17 @@ ThreadedSparseArrays = "59d54670-b8ac-4d81-ab7a-bb56233e17ab" Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" [weakdeps] +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +MatrixEquations = "99c1a7ee-ab34-5fd5-8076-27c950a045f4" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Optim = "429524aa-4258-5aef-a3af-852621145aeb" StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [extensions] +ForwardDiffExt = "ForwardDiff" +MatrixEquationsExt = "MatrixEquations" +MooncakeExt = "Mooncake" OptimExt = "Optim" StatsPlotsExt = "StatsPlots" TuringExt = "Turing" @@ -58,19 +61,22 @@ ADTypes = "1" Accessors = "0.1" Aqua = "0.8" AxisKeys = "0.2" +BenchmarkTools = "1" BlockTriangularForm = "0.1" -CSV = "0.10" ChainRulesCore = "1" Combinatorics = "1" -DataFrames = "1" +CondaPkg = "0.2" DataStructures = "0.18, 0.19" Dates = "1" -DifferentiationInterface = "0.6,0.7" +DelimitedFiles = "1" +DifferentiationInterface = "0.6, 0.7" DispatchDoctor = "0.4" DocStringExtensions = "0.8, 0.9" -DynamicPPL = "0.35 - 0.38" +DynamicPPL = "0.35 - 0.41" DynarePreprocessor_jll = "6" +FastLapackInterface = "2" FiniteDifferences = "0.12" +FlexiChains = "0.6" ForwardDiff = "0.10, 1" JET = "0.07 - 0.11" JSON = "0.21, 1" @@ -84,6 +90,7 @@ LoopVectorization = "0.12" MCMCChains = "6, 7" MacroTools = "0.5" MatrixEquations = "2" +Mooncake = "0.5.25" NLopt = "0.6, 1" Optim = "1" Pigeons = "0.3, 0.4" @@ -92,7 +99,6 @@ Preferences = "1" PythonCall = "0.9" REPL = "1" Random = "1" -RecursiveFactorization = "0.2" Reexport = "1" RuntimeGeneratedFunctions = "0.5" Showoff = "1" @@ -105,7 +111,7 @@ SymPyPythonCall = "0.2 - 0.5" Symbolics = "5 - 7" Test = "1" ThreadedSparseArrays = "0.2.3" -Turing = "0.30 - 0.39" +Turing = "0.30 - 0.45" Unicode = "1" Zygote = "0.6, 0.7" julia = "1.10" @@ -113,21 +119,27 @@ 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" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Pigeons = "0eb8d820-af6a-4919-95ae-11206f830c31" Preferences = "21216c6a-2e73-6563-6e65-726566657250" +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" 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", "BenchmarkTools", "CondaPkg", "PythonCall", "JET", "Dates", "DelimitedFiles", "DifferentiationInterface", "DynamicPPL", "ForwardDiff", "Mooncake", "FlexiChains", "MCMCChains", "LineSearches", "Optim", "MatrixEquations", "Test", "Turing", "Pigeons", "FiniteDifferences", "StatsPlots", "Preferences", "Zygote"] diff --git a/README.md b/README.md index cc1a3af9e..b345d908a 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,8 @@ As of now the package can: - calibrate parameters using (non-stochastic) steady state relationships - **match model moments** (also for pruned **higher order** solutions) - estimate the model on data (Kalman filter using first order perturbation) with **gradient based samplers** (e.g. NUTS, HMC) or **estimate nonlinear models** using the inversion filter -- **differentiate** (forward AD) the model solution, Kalman filter loglikelihood (forward and reverse-mode AD), model moments, steady state, **with respect to the parameters** +- **differentiate** the model solution, loglikelihood (Kalman and inversion filters), model moments, and steady state **with respect to the parameters** using forward-mode AD ([ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)) and reverse-mode AD ([Mooncake.jl](https://github.com/compintell/Mooncake.jl) recommended; other ChainRules-compatible backends such as Zygote.jl also work via custom rrules) +- **modify a model after it has been defined** — model and calibration equations can be updated, added, or removed in place (`update_equations!`, `add_equation!`, `remove_equation!`, and the `*_calibration_equation!` variants) without having to re-run the `@model` / `@parameters` macros. A chronological revision history is kept (`get_revision_history`); this mirrors the equation-revision workflow familiar from `TROLL` The package is not: @@ -49,6 +50,22 @@ Once set up `MacroModelling.jl` can be installed (and `StatsPlots` in order to p using Pkg; Pkg.add(["MacroModelling", "StatsPlots"]) ``` +### Optional extensions + +`MacroModelling.jl` uses Julia's package extension mechanism to provide additional functionality when certain packages are loaded. Install the ones relevant for the intended workflow: + +```julia +using Pkg; Pkg.add(["Turing", "Mooncake"]) # Bayesian estimation with gradient-based samplers +using Pkg; Pkg.add("ForwardDiff") # Forward-mode AD for derivatives of loglikelihood, solutions, IRFs, and moments +using Pkg; Pkg.add("MatrixEquations") # Bartels-Stewart algorithm for Sylvester/Lyapunov equations +using Pkg; Pkg.add("Optim") # LBFGS for conditional forecasts; SAMIN for steady state solver tuning +``` + +**Automatic differentiation backends:** + +- **[Mooncake.jl](https://github.com/compintell/Mooncake.jl)** (reverse-mode) is the recommended backend for gradient-based estimation with Turing.jl (NUTS, HMC). Custom ChainRules `rrule` definitions ensure efficient reverse-mode differentiation through all solvers and filters. Other ChainRules-compatible backends (e.g. Zygote.jl) also work through these same rrules. +- **[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)** (forward-mode) is supported via a package extension and provides `ForwardDiff.jacobian` / `ForwardDiff.gradient` compatibility for `get_solution`, `get_irf`, and `get_statistics` (steady state, mean, variance, standard deviation, covariance, correlation, autocorrelation) across all perturbation orders (first, second, third, and pruned variants), as well as `get_loglikelihood` (Kalman filter only; the inversion filter is not supported with ForwardDiff). + ### Example See below an implementation of a simple RBC model. You can find more detailed tutorials in the [documentation](https://thorek1.github.io/MacroModelling.jl/stable). @@ -151,6 +168,51 @@ Parameter values can also be supplied later (delayed parameter definition) as il See the documentation for more details on the [steady state](https://thorek1.github.io/MacroModelling.jl/stable/steady_state/). +#### Modifying a model after definition + +Model and calibration equations can be edited in place after the `@model` / +`@parameters` block has been evaluated, without having to re-declare the +model. The `update_equations!`, `add_equation!`, and `remove_equation!` +functions (plus their `*_calibration_equation!` counterparts) mutate the +model object, invalidate cached solver results, and re-solve the +non-stochastic steady state. Every change is appended to a revision log +exposed via `get_revision_history`. + +```julia +using MacroModelling + +@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; + +# Replace the technology-shock process with a more persistent one +update_equations!(RBC, :(z[0] = ρ * z[-1] + std_z * eps_z[x]), + :(z[0] = 0.9 * z[-1] + std_z * eps_z[x])) + +# Append an auxiliary equation defining log output +add_equation!(RBC, :(log_q[0] = log(q[0]))) + +# Drop it again +remove_equation!(RBC, :(log_q[0] = log(q[0]))) + +# Inspect what has been changed +get_revision_history(RBC) +``` + +See the [how-to guide on modifying models](https://thorek1.github.io/MacroModelling.jl/stable/how-to/modify_equations/) +for details, and calibration-equation examples. + ## Models The package contains the following models in the `models` folder: @@ -162,6 +224,7 @@ The package contains the following models in the `models` folder: - [Caldara et al. (2012)](https://www.sciencedirect.com/science/article/abs/pii/S1094202511000433) `Caldara_et_al_2012.jl` - [Gali (2015)](https://press.princeton.edu/books/hardcover/9780691164786/monetary-policy-inflation-and-the-business-cycle) - Chapter 3 `Gali_2015_chapter_3_nonlinear.jl` - [Gali and Monacelli (2005)](https://crei.cat/wp-content/uploads/users/pages/roes8739.pdf) - CPI inflation-based Taylor rule `Gali_Monacelli_2005_CITR.jl` +- [Federal Reserve Board U.S. model (FRB/US) - LINVER (2024))](https://www.federalreserve.gov/econres/us-models-about.htm) `FRBUS.jl` - [Gerali, Neri, Sessa, and Signoretti (2010)](https://onlinelibrary.wiley.com/doi/abs/10.1111/j.1538-4616.2010.00331.x) `GNSS_2010.jl` - [Ghironi and Melitz (2005)](https://faculty.washington.edu/ghiro/GhiroMeliQJE0805.pdf) `Ghironi_Melitz_2005.jl` - [Ireland (2004)](http://irelandp.com/pubs/tshocksnk.pdf) `Ireland_2004.jl` diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 0f522c9f7..bdad81703 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -1,175 +1,162 @@ - using BenchmarkTools BenchmarkTools.DEFAULT_PARAMETERS.evals = 10 BenchmarkTools.DEFAULT_PARAMETERS.samples = 1000 BenchmarkTools.DEFAULT_PARAMETERS.seconds = 10 -# Define a parent BenchmarkGroup to contain our SUITE const SUITE = BenchmarkGroup() -# Add some child groups to our benchmark SUITE. - - -# SUITE["FS2000"]["load_time"] = @elapsed using MacroModelling import LinearAlgebra as ℒ using MacroModelling -import MacroModelling: clear_solution_caches!, get_NSSS_and_parameters, calculate_jacobian, merge_calculation_options, solve_lyapunov_equation, ℳ +import MacroModelling: clear_solution_caches!, get_NSSS_and_parameters, solve_lyapunov_equation, ℳ, merge_calculation_options -# Check if new workspace API is available (not present in old package versions) +# Workspace-enabled Lyapunov API exists in v0.1.46+. 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 + import MatrixEquations end -# Version-aware wrapper for solve_lyapunov_equation benchmarking -# For new API: uses pre-allocated workspace for true benchmark of workspace reuse -# For old API: calls without workspace argument -function solve_lyapunov_for_bench(A, C, lyap_ws; lyapunov_algorithm::Symbol = :doubling) - if HAS_WORKSPACE_API - # New API - reuse pre-allocated workspace (shows benefit of workspace caching) - return solve_lyapunov_equation(A, C, lyap_ws; lyapunov_algorithm = lyapunov_algorithm) - else - # Old API - no workspace argument - return solve_lyapunov_equation(A, C; lyapunov_algorithm = lyapunov_algorithm) +# Timings live in different places across versions. +function get_timings(𝓂::ℳ) + if hasproperty(𝓂, :constants) && hasproperty(𝓂.constants, :post_model_macro) + return 𝓂.constants.post_model_macro end + return 𝓂.timings end -function timings_for_bench(𝓂::ℳ) - if hasproperty(𝓂, :timings) - out = 𝓂.timings - else - out = 𝓂.constants.post_model_macro +has_model_field(𝓂::ℳ, field::Symbol) = hasfield(typeof(𝓂), field) +has_nested_field(obj, field::Symbol) = hasfield(typeof(obj), field) + +# Dispatch to the matching jacobian API by model layout first. +function calculate_jacobian_for_bench(parameters, SS_and_pars, 𝓂::ℳ) + if has_model_field(𝓂, :caches) && has_model_field(𝓂, :functions) && + has_nested_field(getfield(𝓂, :functions), :jacobian) + caches_obj = getfield(𝓂, :caches) + jacobian_funcs = getfield(getfield(𝓂, :functions), :jacobian) + + if has_model_field(𝓂, :workspaces) + workspaces_obj = getfield(𝓂, :workspaces) + if hasmethod(calculate_jacobian, + Tuple{typeof(parameters), typeof(SS_and_pars), typeof(caches_obj), typeof(jacobian_funcs), typeof(workspaces_obj)}) + clear_solution_caches!(𝓂, :first_order) + return calculate_jacobian(parameters, + SS_and_pars, + caches_obj, + jacobian_funcs, + workspaces_obj; + caching = false) + end + end + + if hasmethod(calculate_jacobian, + Tuple{typeof(parameters), typeof(SS_and_pars), typeof(caches_obj), typeof(jacobian_funcs)}) + clear_solution_caches!(𝓂, :first_order) + return calculate_jacobian(parameters, SS_and_pars, caches_obj, jacobian_funcs) + end end - return out + + if hasmethod(calculate_jacobian, Tuple{typeof(parameters), typeof(SS_and_pars), typeof(𝓂)}) + clear_solution_caches!(𝓂, :first_order) + return calculate_jacobian(parameters, SS_and_pars, 𝓂) + end + + error("No supported calculate_jacobian benchmark API found for $(typeof(𝓂)).") end +# Dispatch to the matching first-order API by model layout first, then helper availability. 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!(𝓂) - out = calculate_first_order_solution(∇₁, 𝓂.constants, qme_ws, sylv_ws; opts = opts) - else - out = calculate_first_order_solution(∇₁; T = timings_for_bench(𝓂), opts = opts) + if has_model_field(𝓂, :constants) && has_model_field(𝓂, :workspaces) && has_model_field(𝓂, :caches) + constants_obj = getfield(𝓂, :constants) + workspaces_obj = getfield(𝓂, :workspaces) + caches_obj = getfield(𝓂, :caches) + + if hasmethod(calculate_first_order_solution, + Tuple{typeof(∇₁), typeof(constants_obj), typeof(workspaces_obj), typeof(caches_obj)}) + return calculate_first_order_solution(∇₁, + constants_obj, + workspaces_obj, + caches_obj; + opts = opts, + caching = false) + end end - return out -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) + if has_model_field(𝓂, :constants) && isdefined(MacroModelling, :ensure_qme_workspace!) && isdefined(MacroModelling, :ensure_sylvester_1st_order_workspace!) + constants_obj = getfield(𝓂, :constants) + qme_ws_fn = getfield(MacroModelling, :ensure_qme_workspace!) + sylv_ws_fn = getfield(MacroModelling, :ensure_sylvester_1st_order_workspace!) + qme_ws = qme_ws_fn(𝓂) + sylv_ws = sylv_ws_fn(𝓂) + if hasmethod(calculate_first_order_solution, + Tuple{typeof(∇₁), typeof(constants_obj), typeof(qme_ws), typeof(sylv_ws)}) + return calculate_first_order_solution(∇₁, constants_obj, qme_ws, sylv_ws; opts = opts) + end end - return out + + T = get_timings(𝓂) + return calculate_first_order_solution(∇₁; T = T, opts = opts) end +if HAS_WORKSPACE_API + function solve_lyapunov_for_bench(A, C, lyap_ws; lyapunov_algorithm::Symbol = :doubling) + return solve_lyapunov_equation(A, C, lyap_ws; lyapunov_algorithm = lyapunov_algorithm) + end +else + function solve_lyapunov_for_bench(A, C, ::Nothing; lyapunov_algorithm::Symbol = :doubling) + return solve_lyapunov_equation(A, C; lyapunov_algorithm = lyapunov_algorithm) + end +end function run_benchmarks!(𝓂::ℳ, SUITE::BenchmarkGroup) SUITE[𝓂.model_name] = BenchmarkGroup() get_irf(𝓂) - # SUITE[𝓂.model_name]["ttfx_irf"] = BenchmarkTools.Trial(BenchmarkTools.Parameters(seconds=0,samples=1,evals=1,overhead=0,gctrial=false,gcsample=false),[@elapsed get_irf(𝓂)],zeros(1),0,0) - clear_solution_caches!(𝓂, :first_order) - SUITE[𝓂.model_name]["irf"] = @benchmarkable get_irf($𝓂) setup = clear_solution_caches!($𝓂, :first_order) - + reference_steady_state, (solution_error, iters) = get_NSSS_and_parameters(𝓂, 𝓂.parameter_values) - clear_solution_caches!(𝓂, :first_order) - SUITE[𝓂.model_name]["NSSS"] = @benchmarkable get_NSSS_and_parameters($𝓂, $𝓂.parameter_values) setup = clear_solution_caches!($𝓂, :first_order) - - + ∇₁ = calculate_jacobian_for_bench(𝓂.parameter_values, reference_steady_state, 𝓂) - clear_solution_caches!(𝓂, :first_order) - SUITE[𝓂.model_name]["jacobian"] = @benchmarkable calculate_jacobian_for_bench($𝓂.parameter_values, $reference_steady_state, $𝓂) setup = clear_solution_caches!($𝓂, :first_order) - - + SUITE[𝓂.model_name]["qme"] = BenchmarkGroup() qme_schur_opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :schur) qme_doubling_opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :doubling) sol, qme_sol, solved = first_order_solution_for_bench(∇₁, 𝓂; opts = qme_schur_opts) - clear_solution_caches!(𝓂, :first_order) SUITE[𝓂.model_name]["qme"]["schur"] = @benchmarkable first_order_solution_for_bench($∇₁, $𝓂; opts = $qme_schur_opts) setup = clear_solution_caches!($𝓂, :first_order) SUITE[𝓂.model_name]["qme"]["doubling"] = @benchmarkable first_order_solution_for_bench($∇₁, $𝓂; opts = $qme_doubling_opts) setup = clear_solution_caches!($𝓂, :first_order) - T = timings_for_bench(𝓂) - + T = get_timings(𝓂) A = @views sol[:, 1:T.nPast_not_future_and_mixed] * ℒ.diagm(ones(T.nVars))[T.past_not_future_and_mixed_idx,:] - C = @views sol[:, T.nPast_not_future_and_mixed+1:end] - CC = C * C' - - # Create workspace once before benchmarks (new API) or use nothing (old API) - # For new API: pre-allocated workspace shows benefit of workspace caching - # For old API: workspace is not used + lyap_ws = HAS_WORKSPACE_API ? Lyapunov_workspace(size(A, 1)) : nothing - - # Warm up call solve_lyapunov_for_bench(A, CC, lyap_ws) - + SUITE[𝓂.model_name]["lyapunov"] = BenchmarkGroup() SUITE[𝓂.model_name]["lyapunov"]["doubling"] = @benchmarkable solve_lyapunov_for_bench($A, $CC, $lyap_ws, lyapunov_algorithm = :doubling) SUITE[𝓂.model_name]["lyapunov"]["bartels_stewart"] = @benchmarkable solve_lyapunov_for_bench($A, $CC, $lyap_ws, lyapunov_algorithm = :bartels_stewart) SUITE[𝓂.model_name]["lyapunov"]["bicgstab"] = @benchmarkable solve_lyapunov_for_bench($A, $CC, $lyap_ws, lyapunov_algorithm = :bicgstab) SUITE[𝓂.model_name]["lyapunov"]["gmres"] = @benchmarkable solve_lyapunov_for_bench($A, $CC, $lyap_ws, lyapunov_algorithm = :gmres) - - + clear_solution_caches!(𝓂, :first_order) - SUITE[𝓂.model_name]["covariance"] = @benchmarkable get_covariance($𝓂) setup = clear_solution_caches!($𝓂, :first_order) end - include("../models/FS2000.jl") run_benchmarks!(FS2000, SUITE) - include("../models/NAWM_EAUS_2008.jl") run_benchmarks!(NAWM_EAUS_2008, SUITE) - include("../models/Smets_Wouters_2007.jl") run_benchmarks!(Smets_Wouters_2007, SUITE) - - -# SUITE["trig"] = BenchmarkGroup(["math", "triangles"]) -# SUITE["dot"] = BenchmarkGroup(["broadcast", "elementwise"]) - -# This string will be the same every time because we're seeding the RNG -# teststr = join(rand(MersenneTwister(1), 'a':'d', 10^4)) - -# Add some benchmarks to the "string" group -# SUITE["string"]["replace"] = @benchmarkable replace($teststr, "a", "b") seconds = Float64(π) -# SUITE["string"]["join"] = @benchmarkable join($teststr, $teststr) samples = 42 - -# Add some benchmarks to the "trig"/"dot" group -# for f in (sin, cos, tan) -# for x in (0.0, pi) -# SUITE["trig"][string(f), x] = @benchmarkable $(f)($x) -# SUITE["dot"][string(f), x] = @benchmarkable $(f).([$x, $x, $x]) -# end -# end - -# If a caches of tuned parameters already exists, use it, otherwise, tune and caches -# the benchmark parameters. Reusing cached parameters is faster and more reliable -# than re-tuning `SUITE` every time the file is included. -# paramspath = joinpath(dirname(@__FILE__), "params.json") - -# if isfile(paramspath) -# loadparams!(SUITE, BenchmarkTools.load(paramspath)[1], :evals) -# else -# tune!(SUITE) -# BenchmarkTools.save(paramspath, params(SUITE)) -# end diff --git a/docs/Project.toml b/docs/Project.toml index f1110d575..a2d812bbf 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -8,12 +8,14 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" -MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" -MCMCChainsStorage = "51a256e2-afd8-4c38-88d8-a98ba8ad53ca" MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0" MacroModelling = "687ffad2-3618-405e-ac50-e0f7b9c75e44" +MatrixEquations = "99c1a7ee-ab34-5fd5-8076-27c950a045f4" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" @@ -25,5 +27,3 @@ MacroModelling = {path = ".."} [compat] Documenter = "1" -LineSearches = "7.4" -Turing = "0.39" \ No newline at end of file diff --git a/docs/agent-guides/STYLE_GUIDE.md b/docs/agent-guides/STYLE_GUIDE.md new file mode 100644 index 000000000..b97dc7fb5 --- /dev/null +++ b/docs/agent-guides/STYLE_GUIDE.md @@ -0,0 +1,880 @@ +# 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 + +Prefer lines <= 100 characters for general code and docs. For dense mathematical expressions and long signatures, longer lines are acceptable when wrapping harms readability. + +### 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 +``` + +Avoid `try/catch` in performance-critical loops and for normal control flow. + +--- + +## 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 + +### Julia Performance Checklist + +- Put performance-critical logic inside functions (function barriers), not global scope. +- Avoid untyped globals; use `const` for global constants. +- Keep containers concrete (`Vector{T}`, `Dict{K,V}`), avoid abstract element types (`Vector{Any}`, `Vector{Real}`) in hot paths. +- Keep return types predictable from input types; verify with `@code_warntype` on critical kernels. +- Minimize allocations in iterative solvers by pre-allocation and in-place updates. +- Access arrays in memory order when possible; prefer column-wise traversal for column-major arrays. +- Measure before and after optimization (`BenchmarkTools.@btime`, `@allocated`). + +### `@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 (e.g., `Vector{Any}`, `Dict{Symbol,Any}`) +- Prefer concrete element types for arrays and dictionaries in solver kernels + +Example: + +```julia +# Avoid +cache = Dict() + +# Prefer +cache::Dict{Symbol, Float64} = Dict() +``` + +### 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 +``` + +Keep `@ignore_derivatives` scopes narrow; do not wrap solver math that should contribute to gradients. + +### `@inline` and `@noinline` + +Use `@inline` sparingly for tiny, frequently called helper methods when profiling shows a benefit. +Use `@noinline` to keep large diagnostics or cold paths out of hot kernels. +Do not add either annotation by default without evidence. + +--- + +## 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/docs/agent-guides/development-workflow.md b/docs/agent-guides/development-workflow.md new file mode 100644 index 000000000..3dfd647de --- /dev/null +++ b/docs/agent-guides/development-workflow.md @@ -0,0 +1,98 @@ +# 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(...)`). + +## 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) +``` + +For targeted performance validation during development: + +```julia +using BenchmarkTools + +@code_warntype kernel(args...) +@btime kernel($args...) +@allocated kernel(args...) +``` + +- Use `@code_warntype` to detect type instabilities in critical kernels. +- Use `@btime` for robust timing (avoid relying on first-run `@time`). +- Use `@allocated` to confirm allocation reductions from pre-allocation/in-place changes. diff --git a/docs/agent-guides/project-context.md b/docs/agent-guides/project-context.md new file mode 100644 index 000000000..84625764a --- /dev/null +++ b/docs/agent-guides/project-context.md @@ -0,0 +1,60 @@ +# 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` +- `common_docstrings.jl` (shared docstring constants with `®` suffix) +- `nsss_solver.jl` (non-stochastic steady state solver) +- `obc.jl` (occasionally binding constraints) +- `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..dc92fc736 100644 --- a/docs/generate_plots.jl +++ b/docs/generate_plots.jl @@ -5,8 +5,10 @@ ## Setup # using Revise using MacroModelling +import MatrixEquations import StatsPlots using AxisKeys +using Statistics: mean import Random; Random.seed!(10) # For reproducibility of :simulate @@ -624,7 +626,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 @@ -2911,14 +2917,10 @@ using StatsPlots using CSV, DataFrames, AxisKeys import DynamicPPL import Turing -import Turing: NUTS, sample, logpdf, replacenames +import Turing: NUTS, sample, logpdf import ADTypes: AutoZygote # import Zygote -import MCMCChains: Chains - -using HDF5 - -using MCMCChainsStorage +using Serialization @model FS2000 begin dA[0] = exp(gam + z_e_a * e_a[x]) @@ -2991,7 +2993,7 @@ prior_distributions = [ ] Turing.@model function FS2000_loglikelihood_function(prior_distributions, data, m; verbose = false) - parameters ~ Turing.arraydist(prior_distributions) + parameters ~ Turing.product_distribution(prior_distributions) # if DynamicPPL.leafcontext(DynamicPPL.__context__) !== DynamicPPL.PriorContext() Turing.@addlogprob! get_loglikelihood(m, @@ -3005,38 +3007,39 @@ FS2000_loglikelihood = FS2000_loglikelihood_function(prior_distributions, data, # n_samples = 100 -# chain_NUTS = sample(FS2000_loglikelihood, NUTS(), n_samples, progress = false, initial_params = FS2000.parameter_values) +# chain_NUTS = sample(FS2000_loglikelihood, NUTS(), n_samples, progress = false, initial_params = Turing.InitFromParams((; parameters = FS2000.parameter_values))) -# h5open("docs/src/assets/chain_NUTS.h5", "w") do f -# write(f, chain_NUTS) +# open("docs/src/assets/chain_NUTS.jls", "w") do io +# serialize(io, chain_NUTS) # end -chain_NUTS = h5open("docs/src/assets/chain_NUTS.h5", "r") do f read(f, Chains) end - - -chain_NUTS_rn = replacenames(chain_NUTS, Dict(["parameters[$i]" for i in 1:length(FS2000.parameters)] .=> FS2000.parameters)) - -chain_NUTS = replacenames(chain_NUTS, Dict(FS2000.parameters .=> ["parameters[$i]" for i in 1:length(FS2000.parameters)])) +chain_path = "docs/src/assets/chain_NUTS.jls" +chain_NUTS = if isfile(chain_path) + open(deserialize, chain_path) +else + n_samples = 1000 + chain = sample(FS2000_loglikelihood, NUTS(), n_samples, progress = false, initial_params = Turing.InitFromParams((; parameters = FS2000.parameter_values))) + open(chain_path, "w") do io + serialize(io, chain) + end + chain +end # ensure output directory exists and save the chain plot as PNG -p = plot(chain_NUTS_rn) +p = plot(chain_NUTS) savefig(p, joinpath("./docs/src/assets", "FS2000_chain_NUTS.png")) # ![NUTS chain](../assets/FS2000_chain_NUTS.png) -using ComponentArrays, MCMCChains -import DynamicPPL: logjoint - -parameter_mean = mean(chain_NUTS) - -pars = ComponentArray([parameter_mean.nt[2]], Axis(:parameters)); +parameter_mean = collect(values(mean(chain_NUTS); parameters_only = true)) -logjoint(FS2000_loglikelihood, pars) +DynamicPPL.logjoint(FS2000_loglikelihood, (; parameters = parameter_mean)) -function calculate_log_probability(par1, par2, pars_syms, orig_pars, model) - orig_pars[1][pars_syms] = [par1, par2] - logjoint(model, orig_pars) +function calculate_log_probability(par1, par2, pars_syms, orig_pars, likelihood_model) + p = copy(orig_pars) + p[pars_syms] = [par1, par2] + DynamicPPL.logjoint(likelihood_model, (; parameters = p)) end granularity = 32; @@ -3044,25 +3047,31 @@ granularity = 32; par1 = :del; par2 = :gam; -paridx1 = indexin([par1], FS2000.parameters)[1]; -paridx2 = indexin([par2], FS2000.parameters)[1]; +paridx1 = indexin([par1], Symbol.(get_parameters(FS2000)))[1]; +paridx2 = indexin([par2], Symbol.(get_parameters(FS2000)))[1]; -par_range1 = collect(range(minimum(chain_NUTS[Symbol("parameters[$paridx1]")]), stop = maximum(chain_NUTS[Symbol("parameters[$paridx1]")]), length = granularity)); -par_range2 = collect(range(minimum(chain_NUTS[Symbol("parameters[$paridx2]")]), stop = maximum(chain_NUTS[Symbol("parameters[$paridx2]")]), length = granularity)); +param_data = chain_NUTS[:parameters] +n_iter_chain = size(param_data, 1) +n_ch_chain = size(param_data, 2) +all_samples = [param_data[i, c] for i in 1:n_iter_chain for c in 1:n_ch_chain] +parameter_samples_matrix = reduce(hcat, all_samples)' + +par_range1 = collect(range(minimum(parameter_samples_matrix[:, paridx1]), stop = maximum(parameter_samples_matrix[:, paridx1]), length = granularity)); +par_range2 = collect(range(minimum(parameter_samples_matrix[:, paridx2]), stop = maximum(parameter_samples_matrix[:, paridx2]), length = granularity)); p = surface(par_range1, par_range2, - (x,y) -> calculate_log_probability(x, y, [paridx1, paridx2], pars, FS2000_loglikelihood), + (x,y) -> calculate_log_probability(x, y, [paridx1, paridx2], parameter_mean, FS2000_loglikelihood), camera=(30, 65), colorbar=false, color=:inferno); -joint_loglikelihood = [logjoint(FS2000_loglikelihood, ComponentArray([reduce(hcat, get(chain_NUTS, :parameters)[1])[s,:]], Axis(:parameters))) for s in 1:length(chain_NUTS)]; +joint_loglikelihood = vec([DynamicPPL.logjoint(FS2000_loglikelihood, (; parameters = all_samples[j])) for j in eachindex(all_samples)]); -scatter3d!(vec(collect(chain_NUTS[Symbol("parameters[$paridx1]")])), - vec(collect(chain_NUTS[Symbol("parameters[$paridx2]")])), +scatter3d!(parameter_samples_matrix[:, paridx1], + parameter_samples_matrix[:, paridx2], joint_loglikelihood, mc = :viridis, - marker_z = collect(1:length(chain_NUTS)), + marker_z = collect(1:length(joint_loglikelihood)), msw = 0, legend = false, colorbar = false, @@ -3080,9 +3089,9 @@ savefig(p, joinpath("./docs/src/assets", "FS2000_posterior_surface.png")) modeFS2000 = Turing.maximum_a_posteriori(FS2000_loglikelihood, # adtype = AutoZygote(), - initial_params = FS2000.parameter_values) + initial_params = Turing.InitFromParams((; parameters = FS2000.parameter_values))) -get_estimated_shocks(FS2000, data, parameters = collect(modeFS2000.values)) +get_estimated_shocks(FS2000, data, parameters = first(modeFS2000.params.data)) plot_model_estimates(FS2000, data, save_plots = true, @@ -3403,3 +3412,47 @@ plot_conditional_forecast(Smets_Wouters_2003,conditions, shocks = shocks, plots_ # ![Smets_Wouters_2003 conditional forecast 1](../assets/conditional_fcst__SW03__conditional_forecast__1.png) # ![Smets_Wouters_2003 conditional forecast 2](../assets/conditional_fcst__SW03__conditional_forecast__2.png) + +## Clean up: remove files from docs/src/assets/ that are not referenced in the documentation + +let + assets_dir = joinpath(@__DIR__, "src", "assets") + docs_src = joinpath(@__DIR__, "src") + repo_root = dirname(@__DIR__) + + # Collect all filenames referenced in markdown files under docs/src/ + referenced = Set{String}() + for (root, _, files) in walkdir(docs_src) + for f in files + endswith(f, ".md") || continue + for line in eachline(joinpath(root, f)) + for m in eachmatch(r"assets/([^\)\"]+\.\w+)", line) + push!(referenced, m.captures[1]) + end + end + end + end + + # Also scan markdown files in the repo root (README.md, etc.) + for f in readdir(repo_root) + endswith(f, ".md") || continue + for line in eachline(joinpath(repo_root, f)) + for m in eachmatch(r"docs/src/assets/([^\)\"]+\.\w+)", line) + push!(referenced, m.captures[1]) + end + end + end + + # Also keep files used as inputs by this script + push!(referenced, "chain_NUTS.jls") + + # Delete unreferenced files + n_deleted = 0 + for f in readdir(assets_dir) + if f ∉ referenced + rm(joinpath(assets_dir, f)) + n_deleted += 1 + end + end + println("Cleaned up $n_deleted unreferenced files from docs/src/assets/ (kept $(length(referenced)) referenced files)") +end diff --git a/docs/make.jl b/docs/make.jl index dd10d6fe9..f64f0616e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -6,7 +6,7 @@ using Documenter using MacroModelling -import Optim, StatsPlots, Turing +import MatrixEquations, ForwardDiff, Mooncake, Optim, StatsPlots, Turing using DocumenterCitations bib = CitationBibliography( @@ -25,6 +25,9 @@ makedocs( format = Documenter.HTML(size_threshold = 204800*10), modules = [ MacroModelling, + Base.get_extension(MacroModelling, :ForwardDiffExt), + Base.get_extension(MacroModelling, :MatrixEquationsExt), + Base.get_extension(MacroModelling, :MooncakeExt), Base.get_extension(MacroModelling, :OptimExt), Base.get_extension(MacroModelling, :StatsPlotsExt), Base.get_extension(MacroModelling, :TuringExt), @@ -50,6 +53,8 @@ makedocs( "How-to guides" => [ "Programmatic model writing using for-loops" => "how-to/loops.md", "Occasionally binding constraints" => "how-to/obc.md", + "Calibration equations" => "how-to/calibration_equations.md", + "Modifying a model after definition" => "how-to/modify_equations.md", # "how_to.md" ], # "Model syntax" => "dsl.md", diff --git a/docs/src/api.md b/docs/src/api.md index 5c3a72e15..14cb32156 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -1,6 +1,9 @@ ```@autodocs Modules = [ MacroModelling, + Base.get_extension(MacroModelling, :ForwardDiffExt), + Base.get_extension(MacroModelling, :MatrixEquationsExt), + Base.get_extension(MacroModelling, :MooncakeExt), Base.get_extension(MacroModelling, :OptimExt), Base.get_extension(MacroModelling, :StatsPlotsExt), Base.get_extension(MacroModelling, :TuringExt), diff --git a/docs/src/assets/FS2000_chain_NUTS.png b/docs/src/assets/FS2000_chain_NUTS.png index c11969cb9..4f2861daf 100644 Binary files a/docs/src/assets/FS2000_chain_NUTS.png and b/docs/src/assets/FS2000_chain_NUTS.png differ diff --git a/docs/src/assets/FS2000_posterior_surface.png b/docs/src/assets/FS2000_posterior_surface.png index 1f4f235e8..73cf252a4 100644 Binary files a/docs/src/assets/FS2000_posterior_surface.png and b/docs/src/assets/FS2000_posterior_surface.png differ diff --git a/docs/src/assets/Gali_2015_chapter_3_obc__eps_z.png b/docs/src/assets/Gali_2015_chapter_3_obc__eps_z.png deleted file mode 100644 index 0fbfe3ea2..000000000 Binary files a/docs/src/assets/Gali_2015_chapter_3_obc__eps_z.png and /dev/null differ diff --git a/docs/src/assets/Gali_2015_chapter_3_obc__shock_matrix__1.png b/docs/src/assets/Gali_2015_chapter_3_obc__shock_matrix__1.png deleted file mode 100644 index f9b2e6b0a..000000000 Binary files a/docs/src/assets/Gali_2015_chapter_3_obc__shock_matrix__1.png and /dev/null differ diff --git a/docs/src/assets/Gali_2015_chapter_3_obc__simulation__1.png b/docs/src/assets/Gali_2015_chapter_3_obc__simulation__1.png deleted file mode 100644 index 54cf82da6..000000000 Binary files a/docs/src/assets/Gali_2015_chapter_3_obc__simulation__1.png and /dev/null differ diff --git a/docs/src/assets/Gali_2015_chapter_3_obc__simulation__2.png b/docs/src/assets/Gali_2015_chapter_3_obc__simulation__2.png deleted file mode 100644 index 696336195..000000000 Binary files a/docs/src/assets/Gali_2015_chapter_3_obc__simulation__2.png and /dev/null differ diff --git a/docs/src/assets/Gali_2015_chapter_3_obc__simulation__no.png b/docs/src/assets/Gali_2015_chapter_3_obc__simulation__no.png deleted file mode 100644 index d14e4c7ee..000000000 Binary files a/docs/src/assets/Gali_2015_chapter_3_obc__simulation__no.png and /dev/null differ diff --git a/docs/src/assets/borrowing_constraint__no_obc.png b/docs/src/assets/borrowing_constraint__no_obc.png deleted file mode 100644 index b6f8e67ff..000000000 Binary files a/docs/src/assets/borrowing_constraint__no_obc.png and /dev/null differ diff --git a/docs/src/assets/borrowing_constraint__obc.png b/docs/src/assets/borrowing_constraint__obc.png deleted file mode 100644 index 08d59ca65..000000000 Binary files a/docs/src/assets/borrowing_constraint__obc.png and /dev/null differ diff --git "a/docs/src/assets/borrowing_constraint__\316\265_neg.png" "b/docs/src/assets/borrowing_constraint__\316\265_neg.png" deleted file mode 100644 index e8deaee18..000000000 Binary files "a/docs/src/assets/borrowing_constraint__\316\265_neg.png" and /dev/null differ diff --git "a/docs/src/assets/borrowing_constraint__\316\265_pos.png" "b/docs/src/assets/borrowing_constraint__\316\265_pos.png" deleted file mode 100644 index 774880725..000000000 Binary files "a/docs/src/assets/borrowing_constraint__\316\265_pos.png" and /dev/null differ diff --git a/docs/src/assets/chain_NUTS.h5 b/docs/src/assets/chain_NUTS.h5 deleted file mode 100644 index bac889f89..000000000 Binary files a/docs/src/assets/chain_NUTS.h5 and /dev/null differ diff --git a/docs/src/assets/chain_NUTS.jls b/docs/src/assets/chain_NUTS.jls new file mode 100644 index 000000000..97d2d5be4 Binary files /dev/null and b/docs/src/assets/chain_NUTS.jls differ diff --git a/docs/src/assets/cnd_fcst_color__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_color__Gali_2015_chapter_3_nonlinear__2.png index d9018056c..7a7e50f1a 100644 Binary files a/docs/src/assets/cnd_fcst_color__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_color__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/cnd_fcst_higher_order_combine__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/cnd_fcst_higher_order_combine__Gali_2015_chapter_3_nonlinear__1.png index e79c1c7db..34b234388 100644 Binary files a/docs/src/assets/cnd_fcst_higher_order_combine__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/cnd_fcst_higher_order_combine__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/cnd_fcst_init_state_3rd_order__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/cnd_fcst_init_state_3rd_order__Gali_2015_chapter_3_nonlinear__1.png index 3d55388d3..6b6905b7d 100644 Binary files a/docs/src/assets/cnd_fcst_init_state_3rd_order__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/cnd_fcst_init_state_3rd_order__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/cnd_fcst_init_state_compare_orders__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_init_state_compare_orders__Gali_2015_chapter_3_nonlinear__2.png index d59ba8ee5..7a196f45e 100644 Binary files a/docs/src/assets/cnd_fcst_init_state_compare_orders__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_init_state_compare_orders__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/cnd_fcst_no_levels_second_order__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_no_levels_second_order__Gali_2015_chapter_3_nonlinear__2.png index 63779ccb2..0d1f42d73 100644 Binary files a/docs/src/assets/cnd_fcst_no_levels_second_order__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_no_levels_second_order__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/cnd_fcst_plot_stack__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_plot_stack__Gali_2015_chapter_3_nonlinear__2.png deleted file mode 100644 index a431bf20c..000000000 Binary files a/docs/src/assets/cnd_fcst_plot_stack__Gali_2015_chapter_3_nonlinear__2.png and /dev/null differ diff --git a/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__1.png deleted file mode 100644 index d10dbb584..000000000 Binary files a/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__1.png and /dev/null differ diff --git a/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__2.png index bf3778a5b..2429f48af 100644 Binary files a/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_second_order__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/cnd_fcst_second_order_combine__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_second_order_combine__Gali_2015_chapter_3_nonlinear__2.png index 93a732b94..521da76bb 100644 Binary files a/docs/src/assets/cnd_fcst_second_order_combine__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_second_order_combine__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/cnd_fcst_stack__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_stack__Gali_2015_chapter_3_nonlinear__2.png index 0b53f2bfe..61a0c21af 100644 Binary files a/docs/src/assets/cnd_fcst_stack__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_stack__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/cnd_fcst_two_diff__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/cnd_fcst_two_diff__Gali_2015_chapter_3_nonlinear__2.png index d1f1beca0..7b0c79bbf 100644 Binary files a/docs/src/assets/cnd_fcst_two_diff__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/cnd_fcst_two_diff__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/compare_obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png b/docs/src/assets/compare_obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png index 1ce43dcc7..227ab8925 100644 Binary files a/docs/src/assets/compare_obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png and b/docs/src/assets/compare_obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png differ diff --git a/docs/src/assets/conditional_fcst__RBC__conditional_forecast__1.png b/docs/src/assets/conditional_fcst__RBC__conditional_forecast__1.png deleted file mode 100644 index 16926ee6d..000000000 Binary files a/docs/src/assets/conditional_fcst__RBC__conditional_forecast__1.png and /dev/null differ diff --git a/docs/src/assets/conditional_fcst__SW03__conditional_forecast__1.png b/docs/src/assets/conditional_fcst__SW03__conditional_forecast__1.png deleted file mode 100644 index 8b049b15b..000000000 Binary files a/docs/src/assets/conditional_fcst__SW03__conditional_forecast__1.png and /dev/null differ diff --git a/docs/src/assets/conditional_fcst__SW03__conditional_forecast__2.png b/docs/src/assets/conditional_fcst__SW03__conditional_forecast__2.png deleted file mode 100644 index 51bac45b0..000000000 Binary files a/docs/src/assets/conditional_fcst__SW03__conditional_forecast__2.png and /dev/null differ diff --git a/docs/src/assets/estimates_2_per_page__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/estimates_2_per_page__Gali_2015_chapter_3_nonlinear__1.png index 22023505a..583c2c5db 100644 Binary files a/docs/src/assets/estimates_2_per_page__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/estimates_2_per_page__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/estimates__FS2000__3.png b/docs/src/assets/estimates__FS2000__3.png index 88e142156..d3960c0d7 100644 Binary files a/docs/src/assets/estimates__FS2000__3.png and b/docs/src/assets/estimates__FS2000__3.png differ diff --git a/docs/src/assets/estimates_all__Gali_2015_chapter_3_obc__4.png b/docs/src/assets/estimates_all__Gali_2015_chapter_3_obc__4.png index aa65bcb40..0023f95b0 100644 Binary files a/docs/src/assets/estimates_all__Gali_2015_chapter_3_obc__4.png and b/docs/src/assets/estimates_all__Gali_2015_chapter_3_obc__4.png differ diff --git a/docs/src/assets/estimates_all_excluding_obc__FS2000__1.png b/docs/src/assets/estimates_all_excluding_obc__FS2000__1.png index cb5aa06ca..5e3daa4e3 100644 Binary files a/docs/src/assets/estimates_all_excluding_obc__FS2000__1.png and b/docs/src/assets/estimates_all_excluding_obc__FS2000__1.png differ diff --git a/docs/src/assets/estimates_beta_95__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_beta_95__Gali_2015_chapter_3_nonlinear__2.png index 6fe39db82..2bbfa46cc 100644 Binary files a/docs/src/assets/estimates_beta_95__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_beta_95__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_beta_95_vs_99__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/estimates_beta_95_vs_99__Gali_2015_chapter_3_nonlinear__1.png index 309c8a616..0a9341fd6 100644 Binary files a/docs/src/assets/estimates_beta_95_vs_99__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/estimates_beta_95_vs_99__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/estimates_color__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/estimates_color__Gali_2015_chapter_3_nonlinear__1.png index 74c4453af..c8a930e79 100644 Binary files a/docs/src/assets/estimates_color__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/estimates_color__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/estimates_compare_beta_and_orders__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_compare_beta_and_orders__Gali_2015_chapter_3_nonlinear__2.png index 47b857564..4c91c8873 100644 Binary files a/docs/src/assets/estimates_compare_beta_and_orders__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_compare_beta_and_orders__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_first_and_second_order__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_first_and_second_order__Gali_2015_chapter_3_nonlinear__2.png index 8087880e9..8411af23d 100644 Binary files a/docs/src/assets/estimates_first_and_second_order__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_first_and_second_order__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_font__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/estimates_font__Gali_2015_chapter_3_nonlinear__1.png index a68f8a956..e88de7ffb 100644 Binary files a/docs/src/assets/estimates_font__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/estimates_font__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/estimates_forecast_compare__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_forecast_compare__Gali_2015_chapter_3_nonlinear__2.png index 3c73831c4..c19d7e9bb 100644 Binary files a/docs/src/assets/estimates_forecast_compare__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_forecast_compare__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_labels__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_labels__Gali_2015_chapter_3_nonlinear__2.png index c7f99e1e3..9bbd825f2 100644 Binary files a/docs/src/assets/estimates_labels__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_labels__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_labels_symbol__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_labels_symbol__Gali_2015_chapter_3_nonlinear__2.png deleted file mode 100644 index f2e9f4e6f..000000000 Binary files a/docs/src/assets/estimates_labels_symbol__Gali_2015_chapter_3_nonlinear__2.png and /dev/null differ diff --git a/docs/src/assets/estimates_labels_value__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_labels_value__Gali_2015_chapter_3_nonlinear__2.png deleted file mode 100644 index 699b9b98a..000000000 Binary files a/docs/src/assets/estimates_labels_value__Gali_2015_chapter_3_nonlinear__2.png and /dev/null differ diff --git a/docs/src/assets/estimates_multi_params_2__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_multi_params_2__Gali_2015_chapter_3_nonlinear__2.png index 00d16ba2f..b57ea1fb1 100644 Binary files a/docs/src/assets/estimates_multi_params_2__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_multi_params_2__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_multi_params_3__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_multi_params_3__Gali_2015_chapter_3_nonlinear__2.png deleted file mode 100644 index aef8cbb04..000000000 Binary files a/docs/src/assets/estimates_multi_params_3__Gali_2015_chapter_3_nonlinear__2.png and /dev/null differ diff --git a/docs/src/assets/estimates_multi_params__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/estimates_multi_params__Gali_2015_chapter_3_nonlinear__1.png index f72497afa..70cc8fa86 100644 Binary files a/docs/src/assets/estimates_multi_params__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/estimates_multi_params__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/estimates_multiple_orders__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_multiple_orders__Gali_2015_chapter_3_nonlinear__2.png index adcb7dda9..8daa8fb64 100644 Binary files a/docs/src/assets/estimates_multiple_orders__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_multiple_orders__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimates_rekey__FS2000__3.png b/docs/src/assets/estimates_rekey__FS2000__3.png index d77163724..da018eeab 100644 Binary files a/docs/src/assets/estimates_rekey__FS2000__3.png and b/docs/src/assets/estimates_rekey__FS2000__3.png differ diff --git a/docs/src/assets/estimates_rename_dict__Gali_2015_chapter_3_nonlinear__1.png b/docs/src/assets/estimates_rename_dict__Gali_2015_chapter_3_nonlinear__1.png index dc637f155..cd4ea7b8d 100644 Binary files a/docs/src/assets/estimates_rename_dict__Gali_2015_chapter_3_nonlinear__1.png and b/docs/src/assets/estimates_rename_dict__Gali_2015_chapter_3_nonlinear__1.png differ diff --git a/docs/src/assets/estimates_rename_dict_multiple_models__multiple_models__1.png b/docs/src/assets/estimates_rename_dict_multiple_models__multiple_models__1.png index b119373ef..1421b982c 100644 Binary files a/docs/src/assets/estimates_rename_dict_multiple_models__multiple_models__1.png and b/docs/src/assets/estimates_rename_dict_multiple_models__multiple_models__1.png differ diff --git a/docs/src/assets/estimates_shock_decomp_true__Gali_2015_chapter_3_nonlinear__2.png b/docs/src/assets/estimates_shock_decomp_true__Gali_2015_chapter_3_nonlinear__2.png index 43e28ee5e..4939085f1 100644 Binary files a/docs/src/assets/estimates_shock_decomp_true__Gali_2015_chapter_3_nonlinear__2.png and b/docs/src/assets/estimates_shock_decomp_true__Gali_2015_chapter_3_nonlinear__2.png differ diff --git a/docs/src/assets/estimation__m__2.png b/docs/src/assets/estimation__m__2.png deleted file mode 100644 index 1abaff080..000000000 Binary files a/docs/src/assets/estimation__m__2.png and /dev/null differ diff --git a/docs/src/assets/estimation_shock_decomp__m__2.png b/docs/src/assets/estimation_shock_decomp__m__2.png deleted file mode 100644 index d9e3d587c..000000000 Binary files a/docs/src/assets/estimation_shock_decomp__m__2.png and /dev/null differ diff --git a/docs/src/assets/estimation_tutorial__FS2000__2.png b/docs/src/assets/estimation_tutorial__FS2000__2.png deleted file mode 100644 index 8fa919966..000000000 Binary files a/docs/src/assets/estimation_tutorial__FS2000__2.png and /dev/null differ diff --git a/docs/src/assets/fevd__SW03__1.png b/docs/src/assets/fevd__SW03__1.png deleted file mode 100644 index bafe787c1..000000000 Binary files a/docs/src/assets/fevd__SW03__1.png and /dev/null differ diff --git a/docs/src/assets/girf_1000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png b/docs/src/assets/girf_1000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png index c59e91d8f..f752f86c2 100644 Binary files a/docs/src/assets/girf_1000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png and b/docs/src/assets/girf_1000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png differ diff --git a/docs/src/assets/girf_2nd_irf__Gali_2015_chapter_3_nonlinear__eps_a__1.png b/docs/src/assets/girf_2nd_irf__Gali_2015_chapter_3_nonlinear__eps_a__1.png index d4254d98c..7a852b0ac 100644 Binary files a/docs/src/assets/girf_2nd_irf__Gali_2015_chapter_3_nonlinear__eps_a__1.png and b/docs/src/assets/girf_2nd_irf__Gali_2015_chapter_3_nonlinear__eps_a__1.png differ diff --git a/docs/src/assets/girf_5000_500_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png b/docs/src/assets/girf_5000_500_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png index 33c7f67f3..3874be0fe 100644 Binary files a/docs/src/assets/girf_5000_500_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png and b/docs/src/assets/girf_5000_500_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png differ diff --git a/docs/src/assets/girf_5000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png b/docs/src/assets/girf_5000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png index ca16685c1..5fb8ea37b 100644 Binary files a/docs/src/assets/girf_5000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png and b/docs/src/assets/girf_5000_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png differ diff --git a/docs/src/assets/girf_compare_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png b/docs/src/assets/girf_compare_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png index f1a08a8a4..2b9a087ca 100644 Binary files a/docs/src/assets/girf_compare_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png and b/docs/src/assets/girf_compare_irf__Gali_2015_chapter_3_nonlinear__eps_a__2.png differ diff --git a/docs/src/assets/irf__RBC__shock_matrix__1.png b/docs/src/assets/irf__RBC__shock_matrix__1.png deleted file mode 100644 index c670a6750..000000000 Binary files a/docs/src/assets/irf__RBC__shock_matrix__1.png and /dev/null differ diff --git a/docs/src/assets/irf__RBC_new__eps_z__1.png b/docs/src/assets/irf__RBC_new__eps_z__1.png deleted file mode 100644 index 8c9e7f615..000000000 Binary files a/docs/src/assets/irf__RBC_new__eps_z__1.png and /dev/null differ diff --git a/docs/src/assets/irf__RBC_sim__eps_z__1.png b/docs/src/assets/irf__RBC_sim__eps_z__1.png deleted file mode 100644 index c7a735505..000000000 Binary files a/docs/src/assets/irf__RBC_sim__eps_z__1.png and /dev/null differ diff --git a/docs/src/assets/irf__SW03__eta_R__1.png b/docs/src/assets/irf__SW03__eta_R__1.png deleted file mode 100644 index 3b065fdf3..000000000 Binary files a/docs/src/assets/irf__SW03__eta_R__1.png and /dev/null differ diff --git a/docs/src/assets/irf__SW03__shock_matrix__1.png b/docs/src/assets/irf__SW03__shock_matrix__1.png deleted file mode 100644 index c309b4dcc..000000000 Binary files a/docs/src/assets/irf__SW03__shock_matrix__1.png and /dev/null differ diff --git a/docs/src/assets/irf__SW03__simulation__1.png b/docs/src/assets/irf__SW03__simulation__1.png deleted file mode 100644 index ca96ef35d..000000000 Binary files a/docs/src/assets/irf__SW03__simulation__1.png and /dev/null differ diff --git a/docs/src/assets/irf__SW03_new__eta_R__1.png b/docs/src/assets/irf__SW03_new__eta_R__1.png deleted file mode 100644 index c8970048f..000000000 Binary files a/docs/src/assets/irf__SW03_new__eta_R__1.png and /dev/null differ diff --git a/docs/src/assets/obc_all_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png b/docs/src/assets/obc_all_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png index 9de93172a..a81114a99 100644 Binary files a/docs/src/assets/obc_all_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png and b/docs/src/assets/obc_all_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png differ diff --git a/docs/src/assets/obc_girf_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png b/docs/src/assets/obc_girf_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png index b2d4370be..55f611153 100644 Binary files a/docs/src/assets/obc_girf_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png and b/docs/src/assets/obc_girf_compare_irf__Gali_2015_chapter_3_obc__eps_z__1.png differ diff --git a/docs/src/assets/obc_girf_irf__Gali_2015_chapter_3_obc__eps_z__1.png b/docs/src/assets/obc_girf_irf__Gali_2015_chapter_3_obc__eps_z__1.png index 276b77eb3..960ae8ffd 100644 Binary files a/docs/src/assets/obc_girf_irf__Gali_2015_chapter_3_obc__eps_z__1.png and b/docs/src/assets/obc_girf_irf__Gali_2015_chapter_3_obc__eps_z__1.png differ diff --git a/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png b/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png index 5ea8848d9..e6907fdaa 100644 Binary files a/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png and b/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__eps_z__1.png differ diff --git a/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__shock_matrix__1.png b/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__shock_matrix__1.png index 97bc3a508..b669d20e4 100644 Binary files a/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__shock_matrix__1.png and b/docs/src/assets/obc_irf__Gali_2015_chapter_3_obc__shock_matrix__1.png differ diff --git a/docs/src/assets/obc_variables__Gali_2015_chapter_3_obc__A__4.png b/docs/src/assets/obc_variables__Gali_2015_chapter_3_obc__A__4.png index 6b7a7aa6d..6bae011e7 100644 Binary files a/docs/src/assets/obc_variables__Gali_2015_chapter_3_obc__A__4.png and b/docs/src/assets/obc_variables__Gali_2015_chapter_3_obc__A__4.png differ diff --git a/docs/src/assets/sim_obc__Gali_2015_chapter_3_obc__simulation__1.png b/docs/src/assets/sim_obc__Gali_2015_chapter_3_obc__simulation__1.png index 376afb93b..4c21fb6ea 100644 Binary files a/docs/src/assets/sim_obc__Gali_2015_chapter_3_obc__simulation__1.png and b/docs/src/assets/sim_obc__Gali_2015_chapter_3_obc__simulation__1.png differ diff --git a/docs/src/assets/sim_obc_elb__Gali_2015_chapter_3_obc__simulation__1.png b/docs/src/assets/sim_obc_elb__Gali_2015_chapter_3_obc__simulation__1.png index 44c9f58cd..1eb125030 100644 Binary files a/docs/src/assets/sim_obc_elb__Gali_2015_chapter_3_obc__simulation__1.png and b/docs/src/assets/sim_obc_elb__Gali_2015_chapter_3_obc__simulation__1.png differ diff --git a/docs/src/assets/solution__RBC__1.png b/docs/src/assets/solution__RBC__1.png deleted file mode 100644 index 685ab1e85..000000000 Binary files a/docs/src/assets/solution__RBC__1.png and /dev/null differ diff --git a/docs/src/assets/solution__SW03__1.png b/docs/src/assets/solution__SW03__1.png deleted file mode 100644 index 24ef07caa..000000000 Binary files a/docs/src/assets/solution__SW03__1.png and /dev/null differ diff --git a/docs/src/assets/stacked_init_irf__Gali_2015_chapter_3_nonlinear__multiple_shocks__1.png b/docs/src/assets/stacked_init_irf__Gali_2015_chapter_3_nonlinear__multiple_shocks__1.png index fabb65c6d..34f047898 100644 Binary files a/docs/src/assets/stacked_init_irf__Gali_2015_chapter_3_nonlinear__multiple_shocks__1.png and b/docs/src/assets/stacked_init_irf__Gali_2015_chapter_3_nonlinear__multiple_shocks__1.png differ diff --git a/docs/src/assets/tutorial_cond_fcst__Smets_Wouters_2003__2.png b/docs/src/assets/tutorial_cond_fcst__Smets_Wouters_2003__2.png index 769a2fa4d..1fdd5076f 100644 Binary files a/docs/src/assets/tutorial_cond_fcst__Smets_Wouters_2003__2.png and b/docs/src/assets/tutorial_cond_fcst__Smets_Wouters_2003__2.png differ diff --git a/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__2.png b/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__2.png deleted file mode 100644 index 232ad15a4..000000000 Binary files a/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__2.png and /dev/null differ diff --git a/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__3.png b/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__3.png index efe0ecc54..0bd1ff7f2 100644 Binary files a/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__3.png and b/docs/src/assets/with_obc_vars_irf__Gali_2015_chapter_3_obc__eps_z__3.png differ diff --git a/docs/src/how-to/calibration_equations.md b/docs/src/how-to/calibration_equations.md new file mode 100644 index 000000000..cce8e6f60 --- /dev/null +++ b/docs/src/how-to/calibration_equations.md @@ -0,0 +1,135 @@ +# Calibration equations + +Calibration equations let a parameter be determined implicitly by a steady-state target rather than set to a fixed value. Instead of choosing a number for a parameter, a condition on the model's steady state is specified and the solver finds the parameter value that satisfies it. + +## When to use calibration equations + +Calibration equations are useful when an empirical target is easier to observe than the structural parameter itself. Common cases include: + +- Pinning steady-state labour supply to a data average (e.g. one-third of available time) +- Matching a capital-output or investment-output ratio +- Fixing the steady-state gross interest rate to imply a specific discount factor +- Ensuring a government-spending-to-output ratio matches national accounts data + +## Syntax + +A calibration equation links one free parameter to one steady-state condition. Two equivalent forms are supported in the `@parameters` block: + +```julia +# Form 1: parameter on the left +param | steady_state_equation = target + +# Form 2: parameter on the right +steady_state_equation = target | param +``` + +Both forms tell the solver: "find the value of `param` such that `steady_state_equation = target` holds in the non-stochastic steady state." + +Calibration equations are mixed freely with ordinary parameter assignments: + +```julia +@parameters model_name begin + σ = 1 # fixed value + ψ | l[ss] = 1/3 # calibrated: ψ adjusts so l_ss = 1/3 + δ = 0.025 # fixed value + β | R[ss] = 1.0035 # calibrated: β adjusts so R_ss = 1.0035 +end +``` + +## Examples from included models + +### Targeting steady-state labour supply + +In `RBC_baseline.jl` the disutility-of-labour parameter `ψ` is pinned so that steady-state hours equal one-third of the time endowment: + +```julia +ψ | l[ss] = 1/3 +``` + +The same pattern appears in `Ascari_Sbordone_2014.jl` (`d_n | N[ss] = 1/3`) and `Caldara_et_al_2012.jl` (`l[ss] = 1/3 | ν`). + +### Targeting a steady-state ratio + +In `RBC_baseline.jl` the steady-state government spending level is set to match a spending-to-output ratio: + +```julia +g_y = 0.2038 +ḡ | ḡ = g_y * y[ss] +``` + +A more involved ratio target appears in `JQ_2012_RBC.jl`, where the debt-to-output ratio pins `ξ̄`: + +```julia +b[ss] / (y[ss] * (1 + r[ss])) = BY_ratio | ξ̄ +``` + +### Targeting a steady-state price or rate + +In `Smets_Wouters_2003.jl` two auxiliary parameters ensure that steady-state inflation equals its target: + +```julia +calibr_pi_obj | 1 = pi_obj[ss] +calibr_pi | pi[ss] = pi_obj[ss] +``` + +In `Backus_Kehoe_Kydland_1992.jl` the discount factor `beta` is determined by a target level for the steady-state capital stock, using the alternative syntax: + +```julia +K_ss = 11 +K[ss] = K_ss | beta +``` + +### Multiple calibration equations in one block + +The test model `RBC_CME_calibration_equations_and_parameter_definitions.jl` shows several calibration equations alongside ordinary definitions: + +```julia +@parameters m begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + + beta | R[ss] = R_ss + R_ss = 1.0035 + + delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = R_ss - Pi_real + Pi_real = 1/1000 + + phi_pi = 1.5 + rhoz = 9 / 10 + std_eps = .0068 + rho_z_delta = rhoz + std_z_delta = .005 +end +``` + +Here `alpha`, `beta`, and `Pibar` are all calibrated while `delta`, `phi_pi`, and the remaining parameters are set directly. Targets such as `cap_share` and `R_ss` are themselves defined as parameters, keeping the block self-documenting. + +## Common pitfalls + +**No solution exists.** If the target is inconsistent with the model structure the steady-state solver will fail. For example, requesting a capital-output ratio that implies a negative depreciation rate has no valid solution. Review the target value and the model equations when this happens. + +**Solver convergence.** The nonlinear solver needs a reasonable starting region. Providing bounds and initial guesses for calibrated parameters can help: + +```julia +@parameters model_name guess = Dict(:α => 0.3) begin + α | k[ss] / (4 * y[ss]) = 1.5 + 0 < α < 1 +end +``` + +Bounds must be written as standalone comparison statements (e.g. `0 < α < 1` or `α > 0` on separate lines), not appended to the calibration equation. The `guess` keyword provides an initial value for the solver. + +**One equation per parameter.** Each calibration equation pins exactly one parameter. Adding a second calibration equation for the same parameter, or using one equation for two parameters, will cause an error. + +## Verifying calibrated values + +After defining the model, call [`get_steady_state`](@ref) to inspect the solved steady state and confirm that the targets are met: + +```julia +get_steady_state(model_name) +``` + +The output table shows all steady-state variable values and calibrated parameter values. Check that the targeted variables match the specified values. diff --git a/docs/src/how-to/modify_equations.md b/docs/src/how-to/modify_equations.md new file mode 100644 index 000000000..1bee3cfaa --- /dev/null +++ b/docs/src/how-to/modify_equations.md @@ -0,0 +1,153 @@ +# Modifying a model after definition + +`MacroModelling.jl` lets you edit a model in place after the `@model` and +`@parameters` blocks have been evaluated. Equations (and calibration +equations) can be replaced, appended, or removed without re-declaring the +model: the package re-runs the equation-processing pipeline, and recomputes the non-stochastic steady state. Each +change is recorded in a chronological revision log. + +This is useful for iterating on model variants in the REPL, programmatically +generating model alternatives, swapping in observation equations for +estimation, or testing the impact of a single equation change without +rebuilding the entire model definition. + +## API overview + +| Operation | Model equations | Calibration equations | +|:--------------------------|:-----------------------------|:---------------------------------------| +| Replace one (or many) | `update_equations!` | `update_calibration_equations!` | +| Append | `add_equation!` | `add_calibration_equation!` | +| Remove | `remove_equation!` | `remove_calibration_equation!` | + +Helpers: + +- `get_revision_history(𝓂)` — inspect the chronological log of changes. +- `write_julia_model_file(𝓂, path)` — serialise the current model state to a + Julia file that re-creates the (possibly heavily revised) model when + `include`d. + +Each modifying function accepts: + +- The target equation as a 1-based **index**, an `Expr`, or a `String` + (matching is canonical, so whitespace/parenthesisation is ignored). +- Either a single update or a `Vector` / `Tuple` of `(target, new)` pairs + (or, for `add_*`, a vector of new equations) for batch operations. +- A `parameters` keyword that is forwarded to the re-solve step. + +All of them mutate the model object (note the trailing `!`) and return +`nothing`. + +## Working example + +The examples below use a small RBC model. + +```julia +using MacroModelling + +@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; +``` + +### Replace an equation by index + +`get_equations` shows the equations in their stored order; that order is +what the index-based API refers to. + +```julia +get_equations(RBC) + +# Replace the AR(1) shock process (the 4th equation) with a more +# persistent one. +update_equations!(RBC, 4, :(z[0] = 0.9 * z[-1] + std_z * eps_z[x])) +``` + +### Replace an equation by matching `Expr` / `String` + +You don't need to know the index — pass the old equation literally. The +matcher is canonicalising, so the spacing of the input does not matter. + +```julia +update_equations!(RBC, + :(q[0] = exp(z[0]) * k[-1]^α), + :(q[0] = exp(z[0]) * k[-1]^α * l[0]^(1 - α))) # add labour to production + +# Strings work too: +update_equations!(RBC, + "q[0] = exp(z[0]) * k[-1]^α * l[0]^(1 - α)", + "q[0] = exp(z[0]) * k[-1]^α") # revert +``` + +### Add and remove equations + +```julia +# Append a definitional equation +add_equation!(RBC, :(log_q[0] = log(q[0]))) + +# Remove it again (by Expr, by String, or by index) +remove_equation!(RBC, :(log_q[0] = log(q[0]))) +``` + +### Batch updates + +`update_equations!`, `add_equation!`, and `remove_equation!` all accept a +vector to apply several changes in a single re-solve: + +```julia +update_equations!(RBC, [ + (4, :(z[0] = 0.95 * z[-1] + std_z * eps_z[x])), + (:(c[0] + k[0] = (1 - δ) * k[-1] + q[0]), + :(c[0] + k[0] + g[0] = (1 - δ) * k[-1] + q[0])), +]) +``` + +### Modify calibration equations + +Calibration equations use the `lhs = rhs | param` syntax and are edited via +their own functions. The parameter on the right of `|` must already be +declared in the model. + +```julia +# Replace the calibration target for δ +update_calibration_equations!(RBC, 1, :(k[ss] / q[ss] = 10.0 | δ)) + +# Add a new calibration equation +add_calibration_equation!(RBC, :(c[ss] / q[ss] = 0.7 | β)) + +# Remove it again, fixing the freed parameter to a chosen value +remove_calibration_equation!(RBC, :(c[ss] / q[ss] = 0.7 | β), + parameters = :β => 0.95) +``` + +### Inspect the revision history + +Every modification appends an entry to the revision log. Each entry is a +`NamedTuple` with `timestamp`, `action`, `equation_index`, `old_equation`, +and `new_equation` fields. + +```julia +for entry in get_revision_history(RBC) + println(entry.action, " @ ", entry.equation_index, + " : ", entry.old_equation, " => ", entry.new_equation) +end +``` + +### Persist the revised model + +Once happy with the changes, the current state can be written back to a +Julia source file that re-creates the (revised) model when `include`d: + +```julia +write_julia_model_file(RBC, "RBC_revised.jl"; overwrite = true) +``` diff --git a/docs/src/index.md b/docs/src/index.md index 95dde3628..cc9b7209d 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -18,7 +18,8 @@ As of now the package can: - calibrate parameters using (non-stochastic) steady state relationships - **match model moments** (also for pruned **higher order** solutions) - estimate the model on data (Kalman filter using first order perturbation; see [durbin2012time](@citet)) with **gradient based samplers** (e.g. NUTS, HMC) or **estimate nonlinear models** using the inversion filter -- **differentiate** (forward AD) the model solution, Kalman filter loglikelihood (forward and reverse-mode AD), model moments, steady state, **with respect to the parameters** +- **differentiate** the model solution, loglikelihood (Kalman and inversion filters), model moments, and steady state **with respect to the parameters** using forward-mode AD ([ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)) and reverse-mode AD ([Mooncake.jl](https://github.com/compintell/Mooncake.jl) recommended; other ChainRules-compatible backends such as Zygote.jl also work via custom rrules) +- **modify a model after it has been defined** — model and calibration equations can be updated, added, or removed in place ([`update_equations!`](@ref), [`add_equation!`](@ref), [`remove_equation!`](@ref), and the `*_calibration_equation!` variants) without re-running the [`@model`](@ref) / [`@parameters`](@ref) macros. A chronological revision history is kept ([`get_revision_history`](@ref)) and the revised model can be serialised to a Julia source file ([`write_julia_model_file`](@ref)). See the [how-to guide on modifying models](@ref "Modifying a model after definition"). The package is not: @@ -38,6 +39,7 @@ The package contains the following models in the `models` folder: - [Caldara et al. (2012)](https://www.sciencedirect.com/science/article/abs/pii/S1094202511000433) `Caldara_et_al_2012.jl` - [Gali (2015)](https://press.princeton.edu/books/hardcover/9780691164786/monetary-policy-inflation-and-the-business-cycle) - Chapter 3 `Gali_2015_chapter_3_nonlinear.jl` - [Gali and Monacelli (2005)](https://crei.cat/wp-content/uploads/users/pages/roes8739.pdf) - CPI inflation-based Taylor rule `Gali_Monacelli_2005_CITR.jl` +- [Federal Reserve Board U.S. model (FRB/US) - LINVER (2024))](https://www.federalreserve.gov/econres/us-models-about.htm) `FRBUS.jl` - [Gerali, Neri, Sessa, and Signoretti (2010)](https://onlinelibrary.wiley.com/doi/abs/10.1111/j.1538-4616.2010.00331.x) `GNSS_2010.jl` - [Ghironi and Melitz (2005)](https://faculty.washington.edu/ghiro/GhiroMeliQJE0805.pdf) `Ghironi_Melitz_2005.jl` - [Ireland (2004)](http://irelandp.com/pubs/tshocksnk.pdf) `Ireland_2004.jl` 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/docs/src/speed.md b/docs/src/speed.md new file mode 100644 index 000000000..53c325b04 --- /dev/null +++ b/docs/src/speed.md @@ -0,0 +1,63 @@ +# Speed Benchmarks + +The highlighted Windows entries are the conservative cases in this comparison, so they provide a useful lower bound for the overall result. + +- Across the highlighted Windows solve rows, MacroModelling.jl still usually remains faster, from `1.1x` on the largest stress case `FRBUS` (`316/428`) up to `4.8x` on the `Smets_Wouters_2007` (`26/66`) second-order solve. In absolute terms, that spans from `111.79 ms` versus `125.05 ms` for the `FRBUS` first-order solve down to `18.3 μs` versus `395.2 μs` for the `Caldara_et_al_2012` (`3/12`) second-order solve. Third-order bundled timings remain decisively in MacroModelling.jl's favour on Windows at `80.2x` to `115.6x`, corresponding to `736.3 μs` versus `59.02 ms` for `Gali_2015_chapter_3_nonlinear` and `176.6 μs` versus `20.41 ms` for `Caldara_et_al_2012`. +- On Ubuntu 24 and macOS 26, the largest gains appear in derivative construction, and the effect is visible from very small to fairly large models. Small systems such as `Caldara_et_al_2012` (`3/12`), `FS2000` (`4/16`), and `Gali_2015_chapter_3_nonlinear` (`4/23`) build Jacobians in `0.5-1.2 μs` and Hessians in `1.2-2.8 μs` under MacroModelling.jl, while Dynare needs roughly `280 μs-2.5 ms` for Jacobians and `891 μs-2.51 ms` for Hessians. The same pattern extends to larger systems such as `GNSS_2010` (`38/66`), `QUEST3_2009` (`58/107`), and `NAWM_EAUS_2008` (`106/224`), where MacroModelling.jl still stays in the `3.0-64.3 μs` Jacobian range while Dynare takes `1.18-9.0 ms`. +- As model size grows, absolute solve times move from tens of microseconds for the smallest models to milliseconds and then low hundreds of milliseconds for the largest ones, but the ordering remains broadly stable across operating systems. The main exception is the Windows `FRBUS` Jacobian, where MATLAB mex files reduce Dynare's derivative cost to `49.1 μs` versus `722.8 μs` for MacroModelling.jl. Outside that case, Linux and macOS show the clearest speedups, while Windows narrows the gap without changing the overall picture that MacroModelling.jl scales better across the benchmark set. + +All timings reported below are single-thread measurements. Separate multithreaded thread-count sweeps increased runtime across the board for both MacroModelling.jl and Dynare, so only the one-thread results are shown here: in those multithreaded runs, orchestration and parallelisation overhead outweighed any computational gains from additional threads. + +The table is organised in perturbation-order blocks. The opening rows summarise the full machine and software stack, and each order block then repeats a compact header so the per-platform triplets remain readable in plain markdown. Within every block, the speedup columns and the MacroModelling.jl and Dynare timing columns all run left-to-right as Linux, macOS, and Windows. The Windows Dynare results also benefit from MATLAB mex files, which is most visible in the derivative timings where the gap narrows relative to the Octave runs. + +## Benchmark Timings + +Speedup columns report how many times faster MacroModelling.jl is than Dynare for that component. + +| Perturbation Order | Model | Component | Speedup | Speedup | Speedup | States / Variables | MacroModelling.jl | MacroModelling.jl | MacroModelling.jl | Dynare | Dynare | Dynare | +| --- | --- | --- | ---: | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | +|   |   | OS | Ubuntu 24 | macOS 26 | Windows 11 |   | Ubuntu 24 | macOS 26 | Windows 11 | Ubuntu 24 | macOS 26 | Windows 11 | +|   |   | CPU | AMD EPYC 7R13 | Apple M2 | Intel Alder Lake |   | AMD EPYC 7R13 | Apple M2 | Intel Alder Lake | AMD EPYC 7R13 | Apple M2 | Intel Alder Lake | +|   |   | CPU Architecture | x86_64 | aarch64 | x86_64 |   | x86_64 | aarch64 | x86_64 | x86_64 | aarch64 | x86_64 | +|   |   | Package version |   |   |   |   | 0.1.47 | 0.1.47 | 0.1.47 | 7.0.1 | 7.0.1 | 7.0 | +|   |   | Language |   |   |   |   | Julia 1.12.6 | Julia 1.12.6 | Julia 1.12.6 | Octave 11.1.0 | Octave 11.1.0 | MATLAB R2024b Update 6 | +|   |   | BLAS/LAPACK |   |   |   |   | OpenBLAS 0.3.29 | OpenBLAS 0.3.29 | OpenBLAS 0.3.29 | OpenBLAS 0.3.33 | OpenBLAS 0.3.33 | MKL 2024.1 / LAPACK 3.11.0 | +|   |   |   |   |   |   |   |   |   |   |   |   |   | +| **Perturbation Order** | **Model** | **Component** | **Speedup** | **Speedup** | **Speedup** | **States / Variables** | **MacroModelling.jl** | **MacroModelling.jl** | **MacroModelling.jl** | **Dynare** | **Dynare** | **Dynare** | +|   |   | OS | Ubuntu 24 | macOS 26 | Windows 11 |   | Ubuntu 24 | macOS 26 | Windows 11 | Ubuntu 24 | macOS 26 | Windows 11 | +| First | Caldara_et_al_2012 | Jacobian | 466.2x | 422.7x | 20.1x | 3/12 | 1.1 μs | 0.7 μs | 1.5 μs | 512.8 μs | 295.9 μs | 30.2 μs | +|   |   | First-order solve | 17.5x | 11.7x | 7.0x | 3/12 | 21.5 μs | 15.4 μs | 13.3 μs | 376.9 μs | 180.0 μs | 92.7 μs | +|   | FRBUS | Jacobian | 14.8x | 52.5x | 0.1x | 316/428 | 195.6 μs | 34.3 μs | 722.8 μs | 2.9 ms | 1.8 ms | 49.1 μs | +|   |   | First-order solve | 3.5x | 4.0x | **1.1x** | 316/428 | 102.65 ms | 84.2 ms | 111.79 ms | 356.89 ms | 337.31 ms | 125.05 ms | +|   | FS2000 | Jacobian | 617.5x | 558.4x | 17.7x | 4/16 | 0.8 μs | 0.5 μs | 0.6 μs | 494.0 μs | 279.2 μs | 10.6 μs | +|   |   | First-order solve | 9.3x | 7.2x | 2.8x | 4/16 | 43.2 μs | 26.7 μs | 34.6 μs | 402.0 μs | 191.9 μs | 97.7 μs | +|   | GNSS_2010 | Jacobian | 274.7x | 453.3x | 5.1x | 38/66 | 8.3 μs | 3.0 μs | 9.4 μs | 2.28 ms | 1.36 ms | 48.4 μs | +|   |   | First-order solve | 1.9x | 1.7x | 2.2x | 38/66 | 1.24 ms | 967.9 μs | 1.13 ms | 2.33 ms | 1.6 ms | 2.48 ms | +|   | Gali_2015_chapter_3_nonlinear | Jacobian | 420.0x | 351.1x | 13.4x | 4/23 | 1.2 μs | 0.8 μs | 0.8 μs | 504.0 μs | 280.9 μs | 10.7 μs | +|   |   | First-order solve | 7.8x | 5.7x | 3.9x | 4/23 | 57.1 μs | 37.0 μs | 37.5 μs | 448.0 μs | 212.0 μs | 145.1 μs | +|   | NAWM_EAUS_2008 | Jacobian | 140.0x | 154.5x | 1.2x | 106/224 | 64.3 μs | 24.6 μs | 75.2 μs | 9.0 ms | 3.8 ms | 92.2 μs | +|   |   | First-order solve | 2.7x | 2.9x | **1.5x** | 106/224 | 16.59 ms | 13.29 ms | 13.67 ms | 44.07 ms | 39.1 ms | 21.06 ms | +|   | QUEST3_2009 | Jacobian | 115.7x | 295.0x | 1.6x | 58/107 | 15.3 μs | 4.0 μs | 18.1 μs | 1.77 ms | 1.18 ms | 28.6 μs | +|   |   | First-order solve | 2.2x | 2.6x | **2.3x** | 58/107 | 2.45 ms | 1.68 ms | 2.07 ms | 5.37 ms | 4.41 ms | 4.85 ms | +|   | Smets_Wouters_2003 | Jacobian | 319.7x | 308.9x | 3.8x | 19/54 | 7.1 μs | 4.5 μs | 7.1 μs | 2.27 ms | 1.39 ms | 26.7 μs | +|   |   | First-order solve | 1.9x | 1.8x | **2.2x** | 19/54 | 698.5 μs | 474.9 μs | 589.4 μs | 1.34 ms | 831.8 μs | 1.3 ms | +|   | Smets_Wouters_2007 | Jacobian | 207.5x | 154.5x | 1.9x | 26/66 | 9.3 μs | 7.7 μs | 10.5 μs | 1.93 ms | 1.19 ms | 19.9 μs | +|   |   | First-order solve | 1.8x | 2.0x | **2.1x** | 26/66 | 1.08 ms | 657.0 μs | 779.1 μs | 1.91 ms | 1.29 ms | 1.61 ms | +|   |   |   |   |   |   |   |   |   |   |   |   |   | +| **Perturbation Order** | **Model** | **Component** | **Speedup** | **Speedup** | **Speedup** | **States / Variables** | **MacroModelling.jl** | **MacroModelling.jl** | **MacroModelling.jl** | **Dynare** | **Dynare** | **Dynare** | +|   |   | OS | Ubuntu 24 | macOS 26 | Windows 11 |   | Ubuntu 24 | macOS 26 | Windows 11 | Ubuntu 24 | macOS 26 | Windows 11 | +| Second | Caldara_et_al_2012 | Hessian | 525.0x | 524.1x | 65.9x | 3/12 | 2.8 μs | 1.7 μs | 1.6 μs | 1.47 ms | 891.0 μs | 105.4 μs | +|   |   | Second-order solve | 16.0x | 5.1x | 21.6x | 3/12 | 27.9 μs | 40.5 μs | 18.3 μs | 446.1 μs | 206.9 μs | 395.2 μs | +|   | FS2000 | Hessian | 1321.1x | 1291.7x | 22.1x | 4/16 | 1.9 μs | 1.2 μs | 2.4 μs | 2.51 ms | 1.55 ms | 53.0 μs | +|   |   | Second-order solve | 10.0x | 4.4x | 8.0x | 4/16 | 53.0 μs | 57.6 μs | 70.6 μs | 531.0 μs | 255.1 μs | 564.6 μs | +|   | Gali_2015_chapter_3_nonlinear | Hessian | 595.8x | 605.4x | 33.0x | 4/23 | 2.4 μs | 1.5 μs | 1.6 μs | 1.43 ms | 908.1 μs | 52.8 μs | +|   |   | Second-order solve | 8.1x | 3.6x | 6.4x | 4/23 | 78.4 μs | 90.5 μs | 67.2 μs | 635.5 μs | 329.0 μs | 429.5 μs | +|   | Smets_Wouters_2007 | Hessian | 1337.9x | 1251.1x | 36.5x | 26/66 | 6.6 μs | 4.5 μs | 4.2 μs | 8.83 ms | 5.63 ms | 153.1 μs | +|   |   | Second-order solve | 6.5x | 7.9x | **4.8x** | 26/66 | 3.11 ms | 2.45 ms | 3.79 ms | 20.29 ms | 19.44 ms | 18.19 ms | +|   |   |   |   |   |   |   |   |   |   |   |   |   | +| **Perturbation Order** | **Model** | **Component** | **Speedup** | **Speedup** | **Speedup** | **States / Variables** | **MacroModelling.jl** | **MacroModelling.jl** | **MacroModelling.jl** | **Dynare** | **Dynare** | **Dynare** | +|   |   | OS | Ubuntu 24 | macOS 26 | Windows 11 |   | Ubuntu 24 | macOS 26 | Windows 11 | Ubuntu 24 | macOS 26 | Windows 11 | +| Third | Caldara_et_al_2012 | Third-order bundled | 75.6x | 52.8x | **115.6x** | 3/12 | 275.4 μs | 235.9 μs | 176.6 μs | 20.81 ms | 12.46 ms | 20.41 ms | +|   | Gali_2015_chapter_3_nonlinear | Third-order bundled | 42.6x | 33.8x | 80.2x | 4/23 | 828.2 μs | 597.3 μs | 736.3 μs | 35.3 ms | 20.18 ms | 59.02 ms | + +Third-order bundled is not third-order-only on the MacroModelling side. In this harness it is first-order solve + Hessian + second-order solve + third-order derivatives + third-order solve, while Dynare reports the direct bundled `k_order_pert` timing. diff --git a/docs/src/tutorials/calibration.md b/docs/src/tutorials/calibration.md index 7430b0865..bc994cc3c 100644 --- a/docs/src/tutorials/calibration.md +++ b/docs/src/tutorials/calibration.md @@ -193,7 +193,7 @@ The package provides functions specialised for the use with gradient based code get_statistics(Gali_2015, Gali_2015.parameter_values, parameters = Gali_2015.parameters, mean = [:W_real], standard_deviation = [:Pi]) ``` -First the model object is passed on, followed by the parameter values and the parameter names the values correspond to. Then the desired outputs are defined: for the mean real wages are wanted and for the standard deviation inflation is wanted. Outputs for variance, covariance, or autocorrelation can also be obtained the same way as for the mean and standard deviation. +First the model object is passed on, followed by the parameter values and the parameter names the values correspond to. Then the desired outputs are defined: for the mean real wages are wanted and for the standard deviation inflation is wanted. Outputs for variance, covariance, correlation, or autocorrelation can also be obtained the same way as for the mean and standard deviation. Next, a function measuring how close the model is to the target for given values of `:α` and `:std_a` can be defined: diff --git a/docs/src/tutorials/estimation.md b/docs/src/tutorials/estimation.md index e246e5e31..1de9bf436 100644 --- a/docs/src/tutorials/estimation.md +++ b/docs/src/tutorials/estimation.md @@ -96,13 +96,13 @@ data = data(observables,:) ## Define bayesian model -Next the parameter priors are defined using the Turing package. The `@model` macro of the Turing package allows defining the prior distributions over the parameters and combining it with the (Kalman filter) loglikelihood of the model and parameters given the data with the help of the `get_loglikelihood` function. The prior distributions are defined in an array and passed on to the `arraydist` function inside the `@model` macro from the Turing package. It is also possible to define the prior distributions inside the macro but especially for reverse mode auto differentiation the `arraydist` function is substantially faster. When defining the prior distributions the distribution implemented in the Distributions package can be relied upon. Note that the `μσ` parameter allows handing over the moments (`μ` and `σ`) of the distribution as parameters in case of the non-normal distributions (Gamma, Beta, InverseGamma), and upper and lower bounds truncating the distribution can also be defined as third and fourth arguments to the distribution functions. Last but not least, the loglikelihood is defined and added to the posterior loglikelihood with the help of the `@addlogprob!` macro. +Next the parameter priors are defined using the Turing package. The `@model` macro of the Turing package allows defining the prior distributions over the parameters and combining it with the (Kalman filter) loglikelihood of the model and parameters given the data with the help of the `get_loglikelihood` function. The prior distributions are defined in an array and passed on to the `product_distribution` function inside the `@model` macro from the Turing package. It is also possible to define the prior distributions inside the macro but especially for reverse mode auto differentiation the `product_distribution` function is substantially faster. When defining the prior distributions the distribution implemented in the Distributions package can be relied upon. Note that the `μσ` parameter allows handing over the moments (`μ` and `σ`) of the distribution as parameters in case of the non-normal distributions (Gamma, Beta, InverseGamma), and upper and lower bounds truncating the distribution can also be defined as third and fourth arguments to the distribution functions. Last but not least, the loglikelihood is defined and added to the posterior loglikelihood with the help of the `@addlogprob!` macro. ```@repl tutorial_2 import Turing -import Turing: NUTS, sample, logpdf, replacenames -import ADTypes: AutoZygote -import Zygote +import Turing: NUTS, sample, logpdf +import ADTypes: AutoMooncake +import Mooncake prior_distributions = [ Beta(0.356, 0.02, μσ = true), # alp @@ -117,7 +117,7 @@ prior_distributions = [ ] Turing.@model function FS2000_loglikelihood_function(prior_distributions, data, m; verbose = false) - parameters ~ Turing.arraydist(prior_distributions) + parameters ~ Turing.product_distribution(prior_distributions) Turing.@addlogprob! get_loglikelihood(m, data, parameters) end @@ -127,6 +127,8 @@ end The No-U-Turn Sampler (NUTS) is used to obtain the posterior distribution of the parameters. It exploits gradients of the posterior log‑likelihood with respect to model parameters to navigate the parameter space efficiently. NUTS is regarded as robust and fast, and it simplifies tuning by automatically adapting its hyperparameters. +[Mooncake.jl](https://github.com/compintell/Mooncake.jl) is the recommended reverse-mode automatic differentiation backend for gradient-based sampling. The package provides custom `rrule` definitions (via ChainRulesCore) for all solvers and filters, so other ChainRules-compatible backends (e.g. Zygote.jl) also work. For forward-mode AD (e.g. computing Jacobians of solutions or moments), [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) is supported via a package extension. + First the loglikelihood model is defined with the specific data, and model. Next, 1000 samples are drawn from the model: ```@repl tutorial_2 @@ -134,19 +136,17 @@ FS2000_loglikelihood = FS2000_loglikelihood_function(prior_distributions, data, n_samples = 1000 -chain_NUTS = sample(FS2000_loglikelihood, NUTS(), n_samples, progress = false, initial_params = FS2000.parameter_values) +chain_NUTS = sample(FS2000_loglikelihood, NUTS(), n_samples, progress = false, initial_params = Turing.InitFromParams((; parameters = FS2000.parameter_values))) ``` ### Inspect posterior In order to understand the posterior distribution and the sequence of samples they are plotted: -```@repl tutorial_2; setup = :(using HDF5; using MCMCChainsStorage; chain_NUTS = h5open("../assets/chain_NUTS.h5", "r") do f read(f, Chains) end) +```@repl tutorial_2 using StatsPlots -chain_NUTS_rn = replacenames(chain_NUTS, Dict(["parameters[$i]" for i in 1:length(FS2000.parameters)] .=> FS2000.parameters)) - -plot(chain_NUTS_rn); +plot(chain_NUTS); ``` ![NUTS chain](../assets/FS2000_chain_NUTS.png) @@ -154,12 +154,12 @@ plot(chain_NUTS_rn); Next, the posterior loglikelihood is plotted along two parameters dimensions, with the other parameters kept at the posterior mean, and the samples are added to the visualisation. This visualisation allows understanding the curvature of the posterior and puts the samples in context. ```@repl tutorial_2 -using ComponentArrays, MCMCChains +using ComponentArrays import DynamicPPL: logjoint -parameter_mean = mean(chain_NUTS) +parameter_mean = collect(values(mean(chain_NUTS); parameters_only = true)) -pars = ComponentArray([parameter_mean.nt[2]], Axis(:parameters)); +pars = ComponentArray([parameter_mean], Axis(:parameters)); logjoint(FS2000_loglikelihood, pars) @@ -176,8 +176,10 @@ par2 = :gam; paridx1 = indexin([par1], FS2000.parameters)[1]; paridx2 = indexin([par2], FS2000.parameters)[1]; -par_range1 = collect(range(minimum(chain_NUTS[Symbol("parameters[$paridx1]")]), stop = maximum(chain_NUTS[Symbol("parameters[$paridx1]")]), length = granularity)); -par_range2 = collect(range(minimum(chain_NUTS[Symbol("parameters[$paridx2]")]), stop = maximum(chain_NUTS[Symbol("parameters[$paridx2]")]), length = granularity)); +parameter_samples = chain_NUTS[:parameters, stack = true] + +par_range1 = collect(range(minimum(parameter_samples[:, :, paridx1]), stop = maximum(parameter_samples[:, :, paridx1]), length = granularity)); +par_range2 = collect(range(minimum(parameter_samples[:, :, paridx2]), stop = maximum(parameter_samples[:, :, paridx2]), length = granularity)); p = surface(par_range1, par_range2, (x,y) -> calculate_log_probability(x, y, [paridx1, paridx2], pars, FS2000_loglikelihood), @@ -185,13 +187,13 @@ p = surface(par_range1, par_range2, colorbar=false, color=:inferno); -joint_loglikelihood = [logjoint(FS2000_loglikelihood, ComponentArray([reduce(hcat, get(chain_NUTS, :parameters)[1])[s,:]], Axis(:parameters))) for s in 1:length(chain_NUTS)]; +joint_loglikelihood = vec(collect(logjoint(FS2000_loglikelihood, chain_NUTS))); -scatter3d!(vec(collect(chain_NUTS[Symbol("parameters[$paridx1]")])), - vec(collect(chain_NUTS[Symbol("parameters[$paridx2]")])), +scatter3d!(vec(collect(parameter_samples[:, :, paridx1])), + vec(collect(parameter_samples[:, :, paridx2])), joint_loglikelihood, mc = :viridis, - marker_z = collect(1:length(chain_NUTS)), + marker_z = collect(1:length(joint_loglikelihood)), msw = 0, legend = false, colorbar = false, @@ -211,8 +213,8 @@ Other than the mean and median of the posterior distribution the mode can also b ```@repl tutorial_2 modeFS2000 = Turing.maximum_a_posteriori(FS2000_loglikelihood, - adtype = AutoZygote(), - initial_params = FS2000.parameter_values) + adtype = AutoMooncake(; config=nothing), + initial_params = Turing.InitFromParams((; parameters = FS2000.parameter_values))) ``` ## Model estimates given the data and the model solution @@ -249,4 +251,4 @@ plot_model_estimates(FS2000, data) ![Model estimates](../assets/estimates__FS2000__3.png) -shows the variables of the model (blue), data (red), the shock decomposition for each endogenous variable and in the last panel the estimated shocks used to estimate the model. \ No newline at end of file +shows the variables of the model (blue), data (red), the shock decomposition for each endogenous variable and in the last panel the estimated shocks used to estimate the model. diff --git a/docs/src/tutorials/rbc.md b/docs/src/tutorials/rbc.md index 051aef99d..d977a70d7 100644 --- a/docs/src/tutorials/rbc.md +++ b/docs/src/tutorials/rbc.md @@ -171,7 +171,7 @@ import StatsPlots plot_irf(RBC) ``` -![RBC IRF](../assets/tutorial_irf__RBC____1.png) +![RBC IRF](../assets/tutorial_irf__RBC__ϵᶻ__1.png) When the model is solved the first time (in this case by calling [`plot_irf`](@ref)), the package breaks down the steady state problem into independent blocks and first attempts to solve them symbolically and if that fails numerically. diff --git a/docs/src/unfinished_docs/todo.md b/docs/src/unfinished_docs/todo.md index ceecd7285..294ffda50 100644 --- a/docs/src/unfinished_docs/todo.md +++ b/docs/src/unfinished_docs/todo.md @@ -1,11 +1,16 @@ # Todo list - ## High priority -- [ ] allow not to define all parameters in @parameters and enter them later in subsequent calls. so you can do things like loading them from a file and putting them in. internally he would need to delay the solution until all parameters are defined -- [ ] add FRB US model +- [ ] check dispatch doctor coverage in new functions +- [ ] dont use variable, parameter, function names starting with _ +- [ ] add higher order variance decomposition +- [ ] let filters work with relevant subspace of solution only and also allow for missing obs +- [ ] make package work with semi structural expectations models +- [ ] implement stochastic mean fixed point algorithm from Kliem and Meyer Gohde, and then use first order perturbation + kalman filter to estimate around that point. check whether an approximaion around that point actually capures higher order effects such as stochastic vol. +- [ ] implement NestedSampling with NormalisingFlows in Julia +- [ ] Move mean, second-order moments, and third-order moments to compressed space end-to-end; they still use uncompressed-space matrices, and the third-order solution still uses uncompressed second-order matrices. - [ ] check allocations of jacobian in sparse case (NAWM) -- [ ] write another seciton in the docs explaining the parameters macro and what you can do (calibration equations, parameters as functions of other parameters, delayed definition of parameters) +- [ ] write another section in the docs explaining the parameters macro and what you can do (calibration equations, parameters as functions of other parameters, delayed definition of parameters) - [ ] more informative errors when something goes wrong when writing a model - [ ] error when parsing expression of the form: XYZ[0] = 0 @@ -13,14 +18,11 @@ - [ ] have parser accept rss | (r[ss] - 1) * 400 = rss - [ ] allow to define y[ss] = 1 in parameters block -- [ ] add caches to lyapunov krylov solvers -- [ ] eliminiate last elements of factorisation calls not using linearsolvers.jl, check whether they can be done with linearsolvers in case of a matrix as RHS (otherwise consider mumps for sparse matrix RHS) - [ ] separate docs from main package as all the plots get too big - [ ] write tests/docs/technical details for nonlinear obc, forecasting, (non-linear) solution algorithms, SS solver, obc solver, and other algorithms - [ ] collect helper function only used in statsplots extension in that script - [ ] collect the argument wrangling functions in functions instead of them being in function bodies - [ ] apply sort by normalised superscript across functions -- [ ] make package mooncake compatible. write custom pullback functions where necessary (all in one for llh) - [ ] print out the OBC shocks as auxilliary shocks - [ ] generalised higher order IRF is around mean not SSS. plot mean line? - [ ] set irrelevant arguments back to default and inform user @@ -28,17 +30,13 @@ - [ ] consider making sympy an extension or try to partially replace with Symbolics - [ ] make sympy optional (maybe even an extension) and use Symbolics where possible - [ ] switch from sympy to Symbolics -- [ ] replace RF with LinearSolve codes (RF has too many dependencies) - [ ] check again return value when NSSS not found, maybe NaN is better here - [ ] use isfresh flag on dense linear solves - [ ] implement check for plots, that they always return a plot (shocks = :none didn't return a plot) - [ ] cache sparse kron, sylvester solution, and compressed kron in order to avoid allocs; check sparse kron! call - [ ] prettify plotlyjs plots - [ ] fix findiff and zygote consistency for llh derivatives of inversion filter -- [ ] add correlation and other moments to get statistics - [ ] get irf with parameters for higher order and make it zygote compatible -- [ ] implement rrule for higher order moments -- [ ] add derivatives wrt covariance in get_moments - [ ] recheck function examples and docs (include output description) - [ ] Docs: document outputs and associated functions to work with function - [ ] write documentation/docstrings using copilot @@ -53,17 +51,13 @@ - [ ] higher order estimation should start from mean not the stochastic steady state as the mean is the most likely starting point - [ ] large models will need functions to be compiled individually as done for higher order; when tackling that, also separate steady state related equations from the steady state, so that speed issue is addresses due to replacing parameters with the steady state equations from the parameter block; also creat non allocating (residuals) steady state function - [ ] check tols throughout. adopt max(abs,rel*norm) tols -- [ ] redo diffs (DiffInt or ForwardDiff or FastDiff) - [ ] optimize second order estim with SW07 or NAWM - [ ] optimize third order with smaller model - [ ] fix higher order shock finder (3rd order) and check results for pruned second order. are the right state values taken for 1st and second order subprocesses? -- [ ] take analytical derivatives of NSSS funcs to reduce allocation and speed up the NSSS solver - [ ] in the docs make it clear that for estimation you need to have variables which have the name of the observables in the dataframe and the parameters must be handed over to the get_loglikelihood function in the same order as declared. check with get_parameters - [ ] check out dense sparse matmul on transposed matrices -- [ ] check out DiffInterface for NSSS solver - [ ] write plotting callback for NSSS solver - [ ] time NSSS solver and estimation codes -- [ ] move korn_s_s_s to higher order aux variables - [ ] write own interior point solver - [ ] write more tests for the plots - [ ] add background part in docs on NSSS solver (use material from presentation) @@ -76,7 +70,6 @@ - [ ] fix model estimate plot. data not above estimate (should be red but is blue) - [ ] implement higher order (pruned) variance decomposition - [ ] try slicesampler instead of pigeons -- [ ] use faster derivatives for SS solver (currently forward diff) - [ ] speed up sensitivity by caching matrix inversion from implicit diff with LRUcache - [ ] fix this inference errors for large functions. they are slow. fix derivatives in general. - [ ] check downgrade tests @@ -84,14 +77,9 @@ - [ ] take apart solve_matrix_equation for various cases - [ ] try static arrays in KF - [ ] check derivatives of erfcinv with Symbolics. seems off -- [ ] have a workspace in the model object. to be accessed for example by the riccati solver at each run (instead of initialising values at each function call) -- [ ] check why PG samples are off - [ ] optimise vanilla loglikelihood calculation and gradient thereof (incl comp time) - [ ] checkout dynamic perturbation for obc solution: https://www.southampton.ac.uk/~alexmen/dynamic_perturbation.pdf - [ ] checkout schedule free ADAM for global methods: https://github.com/facebookresearch/schedule_free -- [ ] figure out why PG and IS return basically the prior | related to context but also that they need to be somewhat close to the posterior, if they aren't the sampler has a hard time finding it -- [ ] allow external functions to calculate the steady state (and hand it over via SS or get_loglikelihood function) - need to use the check function for implicit derivatives and cannot use it to get him a guess from which he can use internal solver going forward -- [ ] go through custom SS solver once more and try to find parameters and logic that achieves best results - [ ] SS solver with less equations than variables - [ ] improve docs: timing in first sentence seems off; have something more general in first sentence; why is the syntax user friendly? give an example; make the former and the latter a footnote - [ ] change docs to reflect that the output of irfs include aux vars and also the model info Base.show includes aux vars @@ -100,17 +88,15 @@ - [ ] check whether its possible to run parameters macro/block without rerunning model block - [ ] eliminate possible log, ^ terms in parameters block equations - because of nonnegativity errors - [ ] throw error when equations appear more than once -- [ ] make SS calc faster (func and optim, maybe inplace ops) - [ ] check obc once more - [ ] rm obc vars from get_SS - [ ] check why warmup_iterations = 0 makes estimated shocks larger -- [ ] use analytical derivatives also for shocks matching optim (and HMC - implicit diff) - [ ] info on when what filter is used and chosen options are overridden - [ ] check warnings, errors throughout. check suppress not interfering with pigeons - [ ] functions to reverse state_update (input: previous shock and current state, output previous state), find shocks corresponding to bringing one state to the next - [ ] cover nested case: min(50,a+b+max(c,10)) - [ ] add balanced growth path handling -- [ ] autocorr and covariance with derivatives. return 3d array +- [ ] autocorr and corr with derivatives. return 3d array - [ ] add pydsge and econpizza to overview - [ ] add for loop parser in @parameters - [ ] implement more multi country models @@ -120,8 +106,6 @@ - [ ] sampler returned negative std. investigate and come up with solution ensuring sampler can continue - [ ] have get_std take variables as an input - [ ] initial state accept keyed array, SS and SSS as arguments -- [ ] kick out unused parameters from m.parameters -- [ ] use cache for gradient calc in estimation (see DifferentiableStateSpaceModels) - [ ] write functions to debug (fix_SS.jl...) - [ ] model compression (speed up 2nd moment calc (derivatives) for large models; gradient loglikelihood is very slow due to large matmuls) -> model setup as maximisation problem (gEcon) -> HANK models - [ ] implement global solution methods - Julien Pascal, QuantEcon @@ -138,6 +122,31 @@ - [ ] figure out combinations for inputs (parameters and variables in different formats for get_irf for example) - [ ] weed out SS solver and saved objects +- [x] add correlation and other moments to get statistics +- [x] add FRB US model +- [x] include dynare computation in tests +- [x] try nested samplers for estimation (use the python packages) at first, then implement in Julia if they are promising +- [x] kick out unused parameters from m.parameters +- [x] use cache for gradient calc in estimation (see DifferentiableStateSpaceModels) +- [x] use analytical derivatives also for shocks matching optim (and HMC - implicit diff) +- [x] make SS calc faster (func and optim, maybe inplace ops) +- [x] figure out why PG and IS return basically the prior | related to context but also that they need to be somewhat close to the posterior, if they aren't the sampler has a hard time finding it +- [x] allow external functions to calculate the steady state (and hand it over via SS or get_loglikelihood function) - need to use the check function for implicit derivatives and cannot use it to get him a guess from which he can use internal solver going forward +- [x] go through custom SS solver once more and try to find parameters and logic that achieves best results +- [x] have a workspace in the model object. to be accessed for example by the riccati solver at each run (instead of initialising values at each function call) +- [x] check why PG samples are off +- [x] use faster derivatives for SS solver (currently forward diff) +- [x] move korn_s_s_s to higher order aux variables +- [x] check out DiffInterface for NSSS solver +- [x] take analytical derivatives of NSSS funcs to reduce allocation and speed up the NSSS solver +- [x] redo diffs (DiffInt or ForwardDiff or FastDiff) +- [x] implement rrule for higher order moments +- [x] add derivatives wrt covariance in get_moments +- [x] replace RF with LinearSolve codes (RF has too many dependencies) +- [x] make package mooncake compatible. write custom pullback functions where necessary (all in one for llh) +- [x] add caches to lyapunov krylov solvers +- [x] eliminiate last elements of factorisation calls not using linearsolvers.jl, check whether they can be done with linearsolvers in case of a matrix as RHS (otherwise consider mumps for sparse matrix RHS) +- [x] allow not to define all parameters in @parameters and enter them later in subsequent calls. so you can do things like loading them from a file and putting them in. internally he would need to delay the solution until all parameters are defined - [x] separate estimation test using Pigeons from normal tests so that newest version of Turing can be tested and maintained - [x] fix borrowing_constraint how-to - [x] append forecast (no shocks) after estimated variables diff --git a/ext/ForwardDiffExt.jl b/ext/ForwardDiffExt.jl new file mode 100644 index 000000000..8b195872e --- /dev/null +++ b/ext/ForwardDiffExt.jl @@ -0,0 +1,1281 @@ +module ForwardDiffExt + +# ForwardDiff Dual number specializations for forward-mode automatic differentiation +# +# Strategy for each function: +# 1. Extract Float64 values from Dual numbers using ℱ.value.(...) +# 2. Compute the function result on Float64 values +# 3. Compute partials using implicit differentiation or chain rule +# 4. Reconstruct Dual numbers by combining values and partials + +import MacroModelling +import MacroModelling: + # Types + ℳ, constants, workspaces, caches, CalculationOptions, + higher_order_workspace, sylvester_workspace, lyapunov_workspace, + SolverTolerances, AdTolerances, + # Constructors / factories + Higher_order_workspace, merge_calculation_options, + # Functions being specialized + sparse_preallocated!, solve_stochastic_steady_state_newton, + get_NSSS_and_parameters, calculate_first_order_solution, + solve_quadratic_matrix_equation, solve_sylvester_equation, + solve_lyapunov_equation, calculate_loglikelihood, primal, + # Internal helpers + initialise_constants!, ensure_computational_constants!, + ensure_model_structure_constants!, ensure_first_order_constants!, + ensure_first_order_workspace_buffers!, ensure_sylvester_krylov_buffers!, + ensure_sylvester_doubling_buffers!, ensure_qme_doubling_workspace!, + ensure_lyapunov_workspace!, evaluate_custom_steady_state_function, + solve_nsss_wrapper, update_ss_counter!, factorize_lu!, solve_lu_left!, + get_initial_covariance, find_shocks, + # Constants + DEFAULT_SOLVER_PARAMETERS, DEFAULT_QME_ALGORITHM + +import ForwardDiff +const ℱ = ForwardDiff + +import LinearAlgebra as ℒ +import LinearSolve as 𝒮 +import SparseArrays: SparseMatrixCSC, SparseVector, AbstractSparseMatrix, sparse, sparsevec + +# ── Extend primal() for ForwardDiff.Dual ── +MacroModelling.primal(x::ℱ.Dual) = ℱ.value(x) + + +# ── sparse_preallocated! ── + +function MacroModelling.sparse_preallocated!(Ŝ::Matrix{ℱ.Dual{Z,S,N}}; ℂ::higher_order_workspace = Higher_order_workspace()) where {Z,S,N} + sparse(Ŝ) +end + + +# ── solve_stochastic_steady_state_newton (2nd order) ── + +function MacroModelling.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.(𝐒₂) + x̂ = ℱ.value.(x) + + # Get cached computational constants + constants = initialise_constants!(𝓂) + so = constants.second_order + ℂ = 𝓂.workspaces.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 = 𝐒₂̂[T.past_not_future_and_mixed_idx,kron_s⁺_s] + B̂ = 𝐒₂̂[T.past_not_future_and_mixed_idx,kron_s⁺_s⁺] + + # Allocate or reuse workspace for partials and SSS kron buffers. + # NOTE: when this overload is called from a higher-level ForwardDiff path, + # `ℂ` may have been mutated to a `Dual`-typed workspace by the upstream + # perturbation solver. Since the SSS Newton iter here is intentionally + # carried out on the primal (`S`) values only, we allocate fresh `S`-typed + # local buffers whenever the cached ones are not `S`-typed. + nPast = length(x̂) + MacroModelling.ensure_sss_kron_buffers!(ℂ, nPast; third_order=false) + if size(ℂ.∂x_second_order) != (nPast, N) || eltype(ℂ.∂x_second_order) !== S + ℂ.∂x_second_order = zeros(S, nPast, N) + else + fill!(ℂ.∂x_second_order, zero(S)) + end + ∂x̄ = ℂ.∂x_second_order + n_aug = nPast + 1 + if eltype(ℂ.x_aug_buf) === S + x_aug = ℂ.x_aug_buf + kron_x_aug = ℂ.kron_x_aug_xx + kron_x_aug_I = ℂ.kron_x_aug_I + else + x_aug = zeros(S, n_aug) + kron_x_aug = zeros(S, n_aug^2) + kron_x_aug_I = zeros(S, n_aug * nPast, nPast) + end + x_aug[end] = one(S) + + max_iters = 100 + for i in 1:max_iters + copyto!(x_aug, 1, x̂, 1, nPast) + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ∂x = (A + B * kron_x_aug_I - I_nPast) + + ℒ.kron!(kron_x_aug, x_aug, x_aug) + Δx = A * x̂ + B̂ * kron_x_aug / 2 - x̂ + ∂x_lu = ℒ.lu(∂x, check = false) + ℒ.issuccess(∂x_lu) || break + Δx = ∂x_lu \ Δx + + if i > 5 && isapprox(A * x̂ + B̂ * kron_x_aug / 2, x̂, rtol = tol) + break + end + + ℒ.axpy!(-1, Δx, x̂) + end + + copyto!(x_aug, 1, x̂, 1, nPast) + ℒ.kron!(kron_x_aug, x_aug, x_aug) + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + solved = isapprox(A * x̂ + B̂ * kron_x_aug / 2, x̂, rtol = tol) + + if solved + TMP = A + B * kron_x_aug_I - I_nPast + for i in 1:N + ∂𝐒₁ = ℱ.partials.(𝐒₁, i) + ∂𝐒₂ = ℱ.partials.(𝐒₂, i) + + ∂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⁺] + + tmp = ∂A * x̂ + ∂B̂ * kron_x_aug / 2 + + ∂x̄[:,i] = -TMP \ tmp + end + end + + return reshape(map(x̂, eachrow(∂x̄)) do v, p + ℱ.Dual{Z}(v, p...) # Z is the tag + end, size(x̂)), solved +end + + +# ── solve_stochastic_steady_state_newton (3rd order) ── + +function MacroModelling.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!(𝓂.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 = 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⁺] + + # Allocate or reuse workspace for partials and SSS kron buffers. + # See note in the `:second_order` overload above — fall back to fresh + # `S`-typed local buffers when the cached workspace got mutated to a + # `Dual`-typed one upstream. + nPast = length(x̂) + MacroModelling.ensure_sss_kron_buffers!(ℂ, nPast; third_order=true) + if size(ℂ.∂x_third_order) != (nPast, N) || eltype(ℂ.∂x_third_order) !== S + ℂ.∂x_third_order = zeros(S, nPast, N) + else + fill!(ℂ.∂x_third_order, zero(S)) + end + ∂x̄ = ℂ.∂x_third_order + n_aug = nPast + 1 + if eltype(ℂ.x_aug_buf) === S + x_aug = ℂ.x_aug_buf + kron_x_aug = ℂ.kron_x_aug_xx + kron_x_kron = ℂ.kron_x_aug_x_kron + kron_x_aug_I = ℂ.kron_x_aug_I + kron_x_kron_I = ℂ.kron_x_kron_I + else + x_aug = zeros(S, n_aug) + kron_x_aug = zeros(S, n_aug^2) + kron_x_kron = zeros(S, n_aug^3) + kron_x_aug_I = zeros(S, n_aug * nPast, nPast) + kron_x_kron_I = zeros(S, n_aug^2 * nPast, nPast) + end + x_aug[end] = one(S) + + max_iters = 100 + for i in 1:max_iters + copyto!(x_aug, 1, x̂, 1, nPast) + ℒ.kron!(kron_x_aug, x_aug, x_aug) + ℒ.kron!(kron_x_kron, x_aug, kron_x_aug) + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ℒ.kron!(kron_x_kron_I, kron_x_aug, I_nPast) + ∂x = (A + B * kron_x_aug_I + C * kron_x_kron_I / 2 - I_nPast) + + Δx = A * x̂ + B̂ * kron_x_aug / 2 + Ĉ * kron_x_kron / 6 - x̂ + ∂x_lu = ℒ.lu(∂x, check = false) + ℒ.issuccess(∂x_lu) || break + Δx = ∂x_lu \ Δx + + if i > 5 && isapprox(A * x̂ + B̂ * kron_x_aug / 2 + Ĉ * kron_x_kron / 6, x̂, rtol = tol) + break + end + + ℒ.axpy!(-1, Δx, x̂) + end + + copyto!(x_aug, 1, x̂, 1, nPast) + ℒ.kron!(kron_x_aug, x_aug, x_aug) + ℒ.kron!(kron_x_kron, x_aug, kron_x_aug) + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ℒ.kron!(kron_x_kron_I, kron_x_aug, I_nPast) + solved = isapprox(A * x̂ + B̂ * kron_x_aug / 2 + Ĉ * kron_x_kron / 6, x̂, rtol = tol) + + if solved + TMP = A + B * kron_x_aug_I + C * kron_x_kron_I / 2 - I_nPast + for i in 1:N + ∂𝐒₁ = ℱ.partials.(𝐒₁, i) + ∂𝐒₂ = ℱ.partials.(𝐒₂, i) + ∂𝐒₃ = ℱ.partials.(𝐒₃, i) + + ∂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⁺] + ∂Ĉ = ∂𝐒₃[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx,kron_s⁺_s⁺_s⁺] + + tmp = ∂A * x̂ + ∂B̂ * kron_x_aug / 2 + ∂Ĉ * kron_x_kron / 6 + + ∂x̄[:,i] = -TMP \ tmp + end + end + + return reshape(map(x̂, eachrow(∂x̄)) do v, p + ℱ.Dual{Z}(v, p...) # Z is the tag + end, size(x̂)), solved +end + + +# ── separate_values_and_partials_from_sparsevec_dual (internal helper) ── + +function separate_values_and_partials_from_sparsevec_dual(V::SparseVector{ℱ.Dual{Z,S,N}}; tol::AbstractFloat = eps()) where {Z,S,N} + nrows = length(V) + ncols = length(V.nzval[1].partials) + + rows = Int[] + cols = Int[] + + prtls = Float64[] + + for (i,v) in enumerate(V.nzind) + for (k,w) in enumerate(V.nzval[i].partials) + if abs(w) > tol + push!(rows,v) + push!(cols,k) + push!(prtls,w) + end + end + end + + vvals = sparsevec(V.nzind,[i.value for i in V.nzval],nrows) + ps = sparse(rows,cols,prtls,nrows,ncols) + + return vvals, ps +end + + +# ── get_NSSS_and_parameters ── + +function MacroModelling.get_NSSS_and_parameters(𝓂::ℳ, + parameter_values_dual::Vector{ℱ.Dual{Z,S,N}}; + opts::CalculationOptions = merge_calculation_options(), + cold_start::Bool = false, + estimation::Bool = false, + caching::Bool = true)::Tuple{Vector{ℱ.Dual{Z,S,N}}, Tuple{S, Int}} where {Z, S <: AbstractFloat, N} + parameter_values = ℱ.value.(parameter_values_dual) + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + 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 + 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 = 𝓂.workspaces.nsss_solver.check_residual + fill!(residual, 0.0) + + 𝓂.functions.NSSS_check(residual, parameter_values, SS_and_pars_tmp) + + solution_error = ℒ.norm(residual) + + iters = 0 + + 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 + + # Allocate or reuse workspace for partials + if size(qme_ws.∂SS_and_pars) != (length(SS_and_pars), N) + qme_ws.∂SS_and_pars = zeros(S, length(SS_and_pars), N) + else + fill!(qme_ws.∂SS_and_pars, zero(S)) + end + ∂SS_and_pars = qme_ws.∂SS_and_pars + + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) + if opts.verbose println("Failed to find NSSS") end + + # Update failed counter + update_ss_counter!(𝓂.counters, false, estimation = estimation) + + solution_error = S(10.0) + else + # Update success counter + update_ss_counter!(𝓂.counters, true, estimation = estimation) + + custom_ss_expand_matrix = ms.custom_ss_expand_matrix + + + ∂ = parameter_values + C = SS_and_pars[ms.SS_and_pars_no_exo_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 + if jac_cache isa SparseMatrixCSC + jac_cache.nzval .= 0 + else + fill!(jac_cache, zero(eltype(jac_cache))) + end + end + + 𝓂.functions.NSSS_∂equations_∂parameters(jac_cache, ∂, C) + + ∂SS_equations_∂parameters = jac_cache + + + 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_cache = zeros(eltype(SS_and_pars), size(𝓂.caches.NSSS_∂equations_∂SS_and_pars)) + end + else + 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_cache, ∂, C) + + ∂SS_equations_∂SS_and_pars = jac_cache + + 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 + + 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) + + 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 + + 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!(Val(:FastLapack), ∂SS_equations_∂SS_and_pars, + qme_ws.fast_lu_ws_nsss, + qme_ws.fast_lu_dims_nsss) + + 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) +end + + +# ── calculate_first_order_solution ── + +function MacroModelling.calculate_first_order_solution(∇₁::Matrix{ℱ.Dual{Z,S,N}}, + constants::constants, + workspaces::workspaces, + cache::caches; + opts::CalculationOptions = merge_calculation_options(), + 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 + ensure_first_order_constants!(constants) + idx_constants = constants.post_complete_parameters + 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 = 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 + + 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, workspaces, cache; opts = opts, initial_guess = initial_guess_value, caching = caching) + + if !solved + return ∇₁, qme_sol, false + end + + ℒ.mul!(A, @view(∇̂₁[:,1:T.nFuture_not_past_and_mixed]), expand_future) + copyto!(B, @view(∇̂₁[:,idx_constants.nabla_zero_cols])) + + ℒ.mul!(X, @view(𝐒₁[:,1:end-T.nExo]), expand_past) + + copyto!(AXB, B) + ℒ.mul!(AXB, A, X, 1, 1) + + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_AXB, AXBfact = factorize_lu!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), AXB, + qme_ws.fast_lu_ws_nabla0, + qme_ws.fast_lu_dims_nabla0) + + 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 + 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 + fill!(qme_ws.X̃_first_order, zero(eltype(qme_ws.X̃_first_order))) + end + X̃ = qme_ws.X̃_first_order + + p = ∇̂₁ + + copyto!(B_sylv, X) + ℒ.rmul!(B_sylv, -1) + + initial_guess = zeros(eltype(X), size(X, 1), size(X, 2)) + + prev_capture = sylv_ws.pow_capture + sylv_ws.pow_iters = 0 + sylv_ws.pow_capture = true + sylv_ws.pow_transposed = false + sylv_cache_captured = false + + # https://arxiv.org/abs/2011.11430 + for i in 1:N + p .= ℱ.partials.(∇₁, i) + + ℒ.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 + + solve_lu_left!(AXB, CC, qme_ws.fast_lu_ws_nabla0, AXBfact; + use_fastlapack_lu = use_fastlapack_lu) + + if ℒ.norm(CC) < eps() continue end + + ℒ.rmul!(CC, -1) + + dX, slvd = solve_sylvester_equation(AA, B_sylv, CC, sylv_ws, + initial_guess = initial_guess, + sylvester_algorithm = opts.sylvester_algorithm², + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.first_order.ad.sylvester, + verbose = opts.verbose) + + if !sylv_cache_captured + sylv_ws.pow_capture = false # captured A^(2^k) on first solve; reuse for subsequent + sylv_cache_captured = true + end + + if !slvd + fill!(view(X̃, :, i), NaN) + solved = false + continue + end + + # initial_guess = dX + + @views copyto!(X̃[:,i],dX[:,T.past_not_future_and_mixed_idx]) + end + sylv_ws.pow_capture = prev_capture + sylv_ws.pow_iters = 0 + + x = reshape(map(𝐒₁[:,1:end-T.nExo], eachrow(X̃)) do v, p + ℱ.Dual{Z}(v, p...) # Z is the tag + end, size(𝐒₁[:,1:end-T.nExo])) + + Jm = expand_past + + ∇₊ = ∇₁[:,1:T.nFuture_not_past_and_mixed] * expand_future + ∇₀ = ∇₁[:,idx_constants.nabla_zero_cols] + ∇ₑ = ∇₁[:,idx_constants.nabla_e_start:end] + + B = -((∇₊ * x * Jm + ∇₀) \ ∇ₑ) + + 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 = Float64.(MacroModelling.primal.(parameter_values)) + end + + return S₁, qme_sol, solved +end + + +# ── solve_quadratic_matrix_equation ── + +function MacroModelling.solve_quadratic_matrix_equation(A::AbstractMatrix{ℱ.Dual{Z,S,N}}, + B::AbstractMatrix{ℱ.Dual{Z,S,N}}, + C::AbstractMatrix{ℱ.Dual{Z,S,N}}, + constants::constants, + workspaces::workspaces, + cache::caches; + initial_guess::AbstractMatrix{<:Real} = zeros(0,0), + 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 + Â = ℱ.value.(A) + B̂ = ℱ.value.(B) + Ĉ = ℱ.value.(C) + + 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_ws = ensure_qme_doubling_workspace!(workspaces, + T.nVars - T.nPresent_only) + + X, solved = solve_quadratic_matrix_equation(Â, B̂, Ĉ, + constants, + workspaces, + cache; + tol = tol.qme, + initial_guess = initial_guess_value, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + verbose = verbose, + caching = caching) + + AXB = Â * X + B̂ + + AXBfact = ℒ.lu(AXB, check = false) + + if !ℒ.issuccess(AXBfact) + AXBfact = ℒ.svd(AXB) + end + + invAXB = inv(AXBfact) + + AA = invAXB * Â + + X² = X * X + + # Allocate or reuse workspace for partials + if size(qme_ws.X̃) != (length(X), N) + qme_ws.X̃ = zeros(length(X), N) + else + fill!(qme_ws.X̃, zero(eltype(qme_ws.X̃))) + end + X̃ = qme_ws.X̃ + + sws = qme_ws.sylvester + prev_capture = sws.pow_capture + sws.pow_iters = 0 + sws.pow_capture = true + sws.pow_transposed = false + qme_sylv_cache_captured = false + + # https://arxiv.org/abs/2011.11430 + for i in 1:N + dA = ℱ.partials.(A, i) + dB = ℱ.partials.(B, i) + dC = ℱ.partials.(C, i) + + CC = invAXB * (dA * X² + dB * X + dC) + + if ℒ.norm(CC) < eps() continue end + + dX, slvd = solve_sylvester_equation(AA, -X, -CC, qme_ws.sylvester, + sylvester_algorithm = :doubling, + tol = tol.sylvester) + + if !qme_sylv_cache_captured + sws.pow_capture = false # captured A^(2^k) on first solve; reuse for subsequent + qme_sylv_cache_captured = true + end + + solved = Bool(solved) && Bool(slvd) + + if !slvd + fill!(view(X̃, :, i), NaN) + continue + end + + X̃[:,i] = vec(dX) + end + sws.pow_capture = prev_capture + sws.pow_iters = 0 + + return reshape(map(X, eachrow(X̃)) do v, p + ℱ.Dual{Z}(v, p...) # Z is the tag + end, size(X)), solved +end + + +# ── solve_sylvester_equation ── + +function MacroModelling.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{<:Real} = zeros(0,0), + sylvester_algorithm::Symbol = :doubling, + preconditioner::Symbol = :ilu, + tol::SolverTolerances = SolverTolerances(), + verbose::Bool = false)::Tuple{Matrix{ℱ.Dual{Z,S,N}}, Bool} where {Z,S,N} + # Extract Float64 values from Dual numbers + Â = ℱ.value.(A) + B̂ = ℱ.value.(B) + Ĉ = ℱ.value.(C) + + 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 + + # Capture A^(2^k), B^(2^k) sequence from primal so the partial-loop solves can replay them. + prev_capture = 𝕊ℂ.pow_capture + 𝕊ℂ.pow_iters = 0 + 𝕊ℂ.pow_capture = true + 𝕊ℂ.pow_transposed = false + + P̂, solved = solve_sylvester_equation(Â, B̂, Ĉ, 𝕊ℂ, + sylvester_algorithm = sylvester_algorithm, + preconditioner = preconditioner, + tol = tol, + verbose = verbose, + initial_guess = initial_guess_value) + + 𝕊ℂ.pow_capture = false + + 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(Â) + 𝕊ℂ.Ã_fd = copy(Â) + else + copyto!(𝕊ℂ.Ã_fd, Â) + end + Ã = 𝕊ℂ.Ã_fd + + if size(𝕊ℂ.B̃_fd) != size(B̂) + 𝕊ℂ.B̃_fd = copy(B̂) + else + copyto!(𝕊ℂ.B̃_fd, B̂) + end + B̃ = 𝕊ℂ.B̃_fd + + if size(𝕊ℂ.C̃_fd) != size(Ĉ) + 𝕊ℂ.C̃_fd = copy(Ĉ) + else + copyto!(𝕊ℂ.C̃_fd, Ĉ) + end + C̃ = 𝕊ℂ.C̃_fd + + # Allocate or reuse workspace for partials + if size(𝕊ℂ.P̃) != (length(P̂), N) + 𝕊ℂ.P̃ = zeros(S, length(P̂), N) + else + fill!(𝕊ℂ.P̃, zero(S)) + end + P̃ = 𝕊ℂ.P̃ + + for i in 1:N + Ã .= ℱ.partials.(A, i) + B̃ .= ℱ.partials.(B, i) + C̃ .= ℱ.partials.(C, i) + + X = Ã * P̂_stable * B̂ + Â * P̂_stable * B̃ + C̃ + + if ℒ.norm(X) < eps() continue end + + P, slvd = solve_sylvester_equation(Â, B̂, X, 𝕊ℂ, + sylvester_algorithm = sylvester_algorithm, + preconditioner = preconditioner, + tol = tol, + verbose = verbose) + + solved = solved && slvd + + if !slvd + fill!(view(P̃, :, i), NaN) + continue + end + + P̃[:,i] = vec(P) + end + 𝕊ℂ.pow_capture = prev_capture + 𝕊ℂ.pow_iters = 0 + + return reshape(map(P̂_stable, eachrow(P̃)) do v, p + ℱ.Dual{Z}(v, p...) # Z is the tag + end, size(P̂_stable)), solved +end + + +# ── solve_lyapunov_equation ── + +function MacroModelling.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::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), + verbose::Bool = false, + has_unit_roots::Bool = false)::Tuple{Matrix{ℱ.Dual{Z,S,N}}, Bool} where {Z,S,N} + # Extract Float64 values from Dual numbers + Â = ℱ.value.(A) + Ĉ = ℱ.value.(C) + + 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 + + # Capture A^(2^k) sequence from primal so the partial-loop solves can replay them. + prev_capture = workspace.pow_capture + workspace.pow_iters = 0 + workspace.pow_capture = true + workspace.pow_transposed = false + + P̂, solved = solve_lyapunov_equation(Â, Ĉ, workspace; + lyapunov_algorithm = lyapunov_algorithm, + initial_guess = initial_guess_value, + tol = tol, + verbose = verbose, + has_unit_roots = has_unit_roots) + + workspace.pow_capture = false + + 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 + if size(workspace.Ã_fd) != size(Â) + workspace.Ã_fd = copy(Â) + else + copyto!(workspace.Ã_fd, Â) + end + Ã = workspace.Ã_fd + + if size(workspace.C̃_fd) != size(Ĉ) + workspace.C̃_fd = copy(Ĉ) + else + copyto!(workspace.C̃_fd, Ĉ) + end + C̃ = workspace.C̃_fd + + # Allocate or reuse workspace for partials + if size(workspace.P̃) != (length(P̂), N) + workspace.P̃ = zeros(length(P̂), N) + else + fill!(workspace.P̃, zero(eltype(workspace.P̃))) + end + P̃ = workspace.P̃ + + # https://arxiv.org/abs/2011.11430 + for i in 1:N + Ã .= ℱ.partials.(A, i) + C̃ .= ℱ.partials.(C, i) + + X = Ã * P̂_stable * Â' + Â * P̂_stable * Ã' + C̃ + + if ℒ.norm(X) < eps() continue end + + P, slvd = solve_lyapunov_equation(Â, X, workspace; + lyapunov_algorithm = lyapunov_algorithm, + tol = tol, + verbose = verbose, + has_unit_roots = has_unit_roots) + + solved = solved && slvd + + if !slvd + fill!(view(P̃, :, i), NaN) + continue + end + + P̃[:,i] = vec(P) + end + workspace.pow_capture = prev_capture + workspace.pow_iters = 0 + + return reshape(map(P̂_stable, eachrow(P̃)) do v, p + ℱ.Dual{Z}(v, p...) # Z is the tag + end, size(P̂_stable)), solved +end + + +# ── calculate_loglikelihood (Kalman filter for Dual numbers) ── + +function MacroModelling.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, + 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 + + 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 = @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] + + C = @views I_nVars[observables_sorted, observables_and_states] + 𝐁 = B * B' + + P = get_initial_covariance(Val(initial_covariance), A, 𝐁, lyap_ws, opts = opts) + + 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 + + u = zeros(eltype(A), size(C, 2)) + z = C * u + loglik = zero(eltype(A)) + + # Pre-allocate Dual-typed loop buffers + DT = eltype(A) + ns = size(A, 1) + no = size(C, 1) + 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) + + 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 + + @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 opts.verbose println("KF factorisation failed step $t") end + return on_failure_loglikelihood + end + + Fdet = ℒ.det(luF) + if Fdet < eps(Float64) + if opts.verbose println("KF factorisation failed step $t") end + return on_failure_loglikelihood + end + + if t > presample_periods + ℒ.ldiv!(w, luF, v) + loglik += log(Fdet) + ℒ.dot(v, w) + end + + 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) + ℒ.mul!(AP, A, PmKCP) + ℒ.mul!(P, AP, A') + ℒ.axpy!(1, 𝐁, P) + + # u = A * (u + K * v) + ℒ.mul!(Kv, K, v) + copyto!(uKv, u) + ℒ.axpy!(1, Kv, uKv) + ℒ.mul!(u, A, uKv) + ℒ.mul!(z, C, u) + end + + return -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 +end + + +# ── find_shocks (LagrangeNewton, 2nd order) for Dual numbers ── +# Iterative solvers diverge with Dual numbers due to generic LU vs LAPACK +# numerical differences. Solve with Float64 primals, then compute partials +# via the implicit function theorem. + +function MacroModelling.find_shocks(::Val{:LagrangeNewton}, + initial_guess::Vector{ℱ.Dual{Z,V,N}}, + kron_buffer::Vector{ℱ.Dual{Z,V,N}}, + kron_buffer2::AbstractMatrix{ℱ.Dual{Z,V,N}}, + J::ℒ.Diagonal{Bool, Vector{Bool}}, + 𝐒ⁱ::AbstractMatrix{ℱ.Dual{Z,V,N}}, + 𝐒ⁱ²ᵉ::AbstractMatrix{ℱ.Dual{Z,V,N}}, + shock_independent::Vector{ℱ.Dual{Z,V,N}}; + kwargs...) where {Z,V,N} + + # Extract Float64 primals + ig_f = ℱ.value.(initial_guess) + kb_f = ℱ.value.(kron_buffer) + kb2_f = ℱ.value.(kron_buffer2) + Si_f = ℱ.value.(𝐒ⁱ) + Si2e_f = ℱ.value.(𝐒ⁱ²ᵉ) + si_f = ℱ.value.(shock_independent) + + # Solve the primal LagrangeNewton on Float64. + x_f, matched = find_shocks(Val(:LagrangeNewton), + ig_f, kb_f, kb2_f, J, Si_f, Si2e_f, si_f; kwargs...) + + if !matched + return ℱ.Dual{Z,V,N}.(x_f), false + end + + # Propagate partials through the linearised KKT system at the optimum. + # Implicit differentiation through the linearised KKT block. + # Build fXλp = [A tmp'; -tmp 0] once, factor it, and solve for each + # parameter direction. RHS is differentiation of the KKT residual: + # g_x = tmp'·λ - 2x → d g_x = (d_Si + 2·d_Si2e·kron(I,x))' · λ + # g_λ = si - Si·x - Si2e·kron(x,x) + n_x = length(x_f) + n_obs = size(Si_f, 1) + kIx = ℒ.kron(J, x_f) + tmp = Si_f + 2 * Si2e_f * kIx + λ = tmp' \ (2 .* x_f) + A_mat = reshape(2 * Si2e_f' * λ, n_x, n_x) - 2 * J + kxx = ℒ.kron(x_f, x_f) + + fXλp = [A_mat tmp'; + -tmp zeros(V, n_obs, n_obs)] + fXλp_lu = ℒ.lu(fXλp, check = false) + if !ℒ.issuccess(fXλp_lu) + return ℱ.Dual{Z,V,N}.(x_f), false + end + + partials_matrix = zeros(V, n_x, N) + + for k in 1:N + d_si = V[ℱ.partials(shock_independent[i])[k] for i in eachindex(shock_independent)] + d_Si = V[ℱ.partials(𝐒ⁱ[i])[k] for i in eachindex(𝐒ⁱ)] + d_Si2e = V[ℱ.partials(𝐒ⁱ²ᵉ[i])[k] for i in eachindex(𝐒ⁱ²ᵉ)] + + d_Si_mat = reshape(d_Si, size(Si_f)) + d_Si2e_mat = reshape(d_Si2e, size(Si2e_f)) + + dtmp = d_Si_mat + 2 * d_Si2e_mat * kIx + d_g_x = dtmp' * λ + d_g_λ = d_si - d_Si_mat * x_f - d_Si2e_mat * kxx + + sol = fXλp_lu \ vcat(-d_g_x, -d_g_λ) + partials_matrix[:, k] = sol[1:n_x] + end + + x_dual = Vector{ℱ.Dual{Z,V,N}}(undef, n_x) + for i in 1:n_x + x_dual[i] = ℱ.Dual{Z,V,N}(x_f[i], + ℱ.Partials{N,V}(NTuple{N,V}(partials_matrix[i, k] for k in 1:N))) + end + + return x_dual, matched +end + + +# ── find_shocks (LagrangeNewton, 3rd order) for Dual numbers ── +# Same implicit-differentiation strategy as the 2nd-order variant. +# Residual: g(x) = si - Si*x - Si2e*kron(x,x) - Si3e*kron(x,kron(x,x)) = 0 +# Jacobian: Si + 2*Si2e*kron(I,x) + 3*Si3e*kron(I,kron(x,x)) + +function MacroModelling.find_shocks(::Val{:LagrangeNewton}, + initial_guess::Vector{ℱ.Dual{Z,V,N}}, + kron_buffer::Vector{ℱ.Dual{Z,V,N}}, + kron_buffer²::Vector{ℱ.Dual{Z,V,N}}, + kron_buffer2::AbstractMatrix{ℱ.Dual{Z,V,N}}, + kron_buffer3::AbstractMatrix{ℱ.Dual{Z,V,N}}, + kron_buffer4::AbstractMatrix{ℱ.Dual{Z,V,N}}, + J::ℒ.Diagonal{Bool, Vector{Bool}}, + 𝐒ⁱ::AbstractMatrix{ℱ.Dual{Z,V,N}}, + 𝐒ⁱ²ᵉ::AbstractMatrix{ℱ.Dual{Z,V,N}}, + 𝐒ⁱ³ᵉ::AbstractMatrix{ℱ.Dual{Z,V,N}}, + shock_independent::Vector{ℱ.Dual{Z,V,N}}; + kwargs...) where {Z,V,N} + + # Extract Float64 primals + ig_f = ℱ.value.(initial_guess) + kb_f = ℱ.value.(kron_buffer) + kb²_f = ℱ.value.(kron_buffer²) + kb2_f = ℱ.value.(kron_buffer2) + kb3_f = ℱ.value.(kron_buffer3) + kb4_f = ℱ.value.(kron_buffer4) + Si_f = ℱ.value.(𝐒ⁱ) + Si2e_f = ℱ.value.(𝐒ⁱ²ᵉ) + Si3e_f = ℱ.value.(𝐒ⁱ³ᵉ) + si_f = ℱ.value.(shock_independent) + + x_f, matched = find_shocks(Val(:LagrangeNewton), + ig_f, kb_f, kb²_f, kb2_f, kb3_f, kb4_f, J, Si_f, Si2e_f, Si3e_f, si_f; kwargs...) + + if !matched + return ℱ.Dual{Z,V,N}.(x_f), false + end + + # Implicit differentiation through the linearised KKT block. + # fXλp = [A tmp'; -tmp 0] with + # A = reshape((2·Si2e + 6·Si3e·kron(I,kIx))'·λ, n_x, n_x) - 2I + # tmp = Si + 2·Si2e·kron(I,x) + 3·Si3e·kron(I,kron(x,x)) + n_x = length(x_f) + n_obs = size(Si_f, 1) + kxx = ℒ.kron(x_f, x_f) + kxxx = ℒ.kron(x_f, kxx) + kIx = ℒ.kron(J, x_f) + kIxx = ℒ.kron(J, kxx) + + tmp = Si_f + 2 * Si2e_f * kIx + 3 * Si3e_f * kIxx + λ = tmp' \ (2 .* x_f) + A_mat = reshape((2 * Si2e_f + 6 * Si3e_f * ℒ.kron(J, kIx))' * λ, n_x, n_x) - 2 * J + + fXλp = [A_mat tmp'; + -tmp zeros(V, n_obs, n_obs)] + fXλp_lu = ℒ.lu(fXλp, check = false) + if !ℒ.issuccess(fXλp_lu) + return ℱ.Dual{Z,V,N}.(x_f), false + end + + partials_matrix = zeros(V, n_x, N) + + for k in 1:N + d_si = V[ℱ.partials(shock_independent[i])[k] for i in eachindex(shock_independent)] + d_Si = V[ℱ.partials(𝐒ⁱ[i])[k] for i in eachindex(𝐒ⁱ)] + d_Si2e = V[ℱ.partials(𝐒ⁱ²ᵉ[i])[k] for i in eachindex(𝐒ⁱ²ᵉ)] + d_Si3e = V[ℱ.partials(𝐒ⁱ³ᵉ[i])[k] for i in eachindex(𝐒ⁱ³ᵉ)] + + d_Si_mat = reshape(d_Si, size(Si_f)) + d_Si2e_mat = reshape(d_Si2e, size(Si2e_f)) + d_Si3e_mat = reshape(d_Si3e, size(Si3e_f)) + + dtmp = d_Si_mat + 2 * d_Si2e_mat * kIx + 3 * d_Si3e_mat * kIxx + d_g_x = dtmp' * λ + d_g_λ = d_si - d_Si_mat * x_f - d_Si2e_mat * kxx - d_Si3e_mat * kxxx + + sol = fXλp_lu \ vcat(-d_g_x, -d_g_λ) + partials_matrix[:, k] = sol[1:n_x] + end + + x_dual = Vector{ℱ.Dual{Z,V,N}}(undef, n_x) + for i in 1:n_x + x_dual[i] = ℱ.Dual{Z,V,N}(x_f[i], + ℱ.Partials{N,V}(NTuple{N,V}(partials_matrix[i, k] for k in 1:N))) + end + + return x_dual, matched +end + + +end # module ForwardDiffExt diff --git a/ext/MatrixEquationsExt.jl b/ext/MatrixEquationsExt.jl new file mode 100644 index 000000000..4f3a8e8d3 --- /dev/null +++ b/ext/MatrixEquationsExt.jl @@ -0,0 +1,92 @@ +module MatrixEquationsExt +# good overview: https://cscproxy.mpi-magdeburg.mpg.de/mpcsc/benner/talks/Benner-Melbourne2019.pdf +import MacroModelling +import MacroModelling: + sylvester_workspace, lyapunov_workspace, SolverTolerances, + solve_sylvester_equation, solve_lyapunov_equation, + ensure_sylvester_krylov_buffers!, ensure_lyapunov_doubling_buffers!, + BARTELS_STEWART_AVAILABLE + +import MatrixEquations +import LinearAlgebra as ℒ + +function __init__() + BARTELS_STEWART_AVAILABLE[] = true +end + +# ── Sylvester: Bartels-Stewart via MatrixEquations.sylvd ── + +function MacroModelling.solve_sylvester_equation(A::DenseMatrix{T}, + B::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMatrix{T}}, + C::DenseMatrix{T}, + ::Val{:bartels_stewart}, + 𝕊ℂ::sylvester_workspace; + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), + preconditioner::Symbol = :none, + verbose::Bool = false, + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + + if length(initial_guess) == 0 + initial_guess = zero(C) + end + + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_krylov_buffers!(𝕊ℂ, n, m) + + 𝐂¹ = 𝕊ℂ.𝐂 + tmp̄ = 𝕊ℂ.tmp + + # 𝐂¹ = A * initial_guess * B + C - initial_guess + ℒ.mul!(tmp̄, initial_guess, B) + ℒ.mul!(𝐂¹, A, tmp̄) + ℒ.axpy!(1, C, 𝐂¹) + ℒ.axpy!(-1, initial_guess, 𝐂¹) + + 𝐂 = try + MatrixEquations.sylvd(-A, B, 𝐂¹)::Matrix{T} + catch + return C, 0, 1.0 + end + + 𝐂 += initial_guess + + ℒ.mul!(tmp̄, 𝐂, B) + ℒ.mul!(𝐂¹, A, tmp̄) + ℒ.axpy!(1, C, 𝐂¹) + ℒ.axpy!(-1, 𝐂, 𝐂¹) + + reached_tol = ℒ.norm(𝐂¹) / max(ℒ.norm(𝐂), ℒ.norm(C)) + + return 𝐂, -1, reached_tol +end + +# ── Lyapunov: Bartels-Stewart via MatrixEquations.lyapd ── + +function MacroModelling.solve_lyapunov_equation(A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMatrix{T}}, + C::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMatrix{T}}, + ::Val{:bartels_stewart}, + workspace::lyapunov_workspace; + tol::SolverTolerances = SolverTolerances(), + has_unit_roots::Bool = false)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + + 𝐂 = try + MatrixEquations.lyapd(A, C)::Matrix{T} + catch + return C, 0, 1.0 + end + + 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(𝐂¹_tmp) / ℒ.norm(𝐂) + + return 𝐂, 0, reached_tol +end + +end # module diff --git a/ext/MooncakeExt.jl b/ext/MooncakeExt.jl new file mode 100644 index 000000000..88be29650 --- /dev/null +++ b/ext/MooncakeExt.jl @@ -0,0 +1,219 @@ +module MooncakeExt + +import MacroModelling +import Mooncake +import Mooncake: CoDual, NoRData, @is_primitive +import ChainRulesCore +import AxisKeys: KeyedArray + +# ── Inference short-circuit for Mooncake primitives ── +# Mooncake's default abstract_call_gf_by_type runs full type inference on every +# call BEFORE checking whether it is a primitive. For large primitives like +# get_statistics (thousands of transitive callees) this cascade takes 10+ min +# through Mooncake's fresh inference cache. +# +# Fix: check the primitive predicate FIRST. When a call is recognised as a +# primitive, return a conservative CallMeta (return type Any, unknown effects) +# immediately, skipping the expensive inference cascade. Correctness is +# preserved because Mooncake's AD tape uses CoDual values with concrete runtime +# types, and make_ad_stmts! already handles imprecise inferred types. +# +# Defined in __init__ to avoid "method overwriting during precompilation" error. +# function __init__() +# @static if VERSION >= v"1.12-" +# CC = Core.Compiler +# @eval begin +# function $CC.abstract_call_gf_by_type( +# interp::Mooncake.MooncakeInterpreter{C,M}, +# @nospecialize(f), +# arginfo::$CC.ArgInfo, +# si::$CC.StmtInfo, +# @nospecialize(atype), +# sv::$CC.AbsIntState, +# max_methods::Int, +# ) where {C,M} +# argtypes = arginfo.argtypes +# matches = $CC.find_method_matches(interp, argtypes, atype; max_methods) +# if !isa(matches, $CC.FailedMethodMatch) +# (; applicable) = matches +# if Mooncake.any_matches_primitive(applicable, C, M, interp.world) +# info = Mooncake.NoInlineCallInfo($CC.NoCallInfo(), atype) +# cm = $CC.CallMeta(Any, Any, $CC.Effects(), info) +# return $CC.Future(cm) +# end +# end +# return @invoke $CC.abstract_call_gf_by_type( +# interp::$CC.AbstractInterpreter, +# f::Any, +# arginfo::$CC.ArgInfo, +# si::$CC.StmtInfo, +# atype::Any, +# sv::$CC.AbsIntState, +# max_methods::Int, +# ) +# end +# end +# end +# end + +Mooncake.tangent_type(::Type{MacroModelling.ℳ}) = Mooncake.NoTangent + +# ── Scalar/Array-returning functions: @from_rrule works directly ── + +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_solution), MacroModelling.ℳ, Vector{T}} where {T<:Base.IEEEFloat} true + +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}} where {T<:Base.IEEEFloat} true + +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_irf), MacroModelling.ℳ, Vector{T}} where {T<:Base.IEEEFloat} true + +# ── DynamicPPL compatibility: wider @is_primitive declarations ── +# Inside a Turing @model evaluated through DynamicPPL.logdensity_at, +# tilde_assume!! returns Any, so Julia's type inference widens the params +# argument to Any at the call site. The narrow signatures generated by +# @from_rrule do not match during Mooncake's abstract interpretation, +# causing it to trace into the full function body (~700s). +# These wider declarations ensure the primitive is recognized. +# At runtime, CoDual carries concrete types, so the narrow rrule!! methods +# auto-generated by @from_rrule still dispatch correctly. +@is_primitive Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, Any, Any} +@is_primitive Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), <:NamedTuple, typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, Any, Any} + +@is_primitive Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_solution), MacroModelling.ℳ, Any} +@is_primitive Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), <:NamedTuple, typeof(MacroModelling.get_solution), MacroModelling.ℳ, Any} + +@is_primitive Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_irf), MacroModelling.ℳ, Any} +@is_primitive Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), <:NamedTuple, typeof(MacroModelling.get_irf), MacroModelling.ℳ, Any} + +# ── get_statistics: manual rrule!! ── +# Returns Dict{Symbol,...} whose MutableTangent cannot be converted by to_cr_tangent. +# We mirror rrule_wrapper but reconstruct the Dict cotangent from MutableTangent fields. + +# Convert MutableTangent (Dict internals: slots/keys/vals/...) → actual Dict cotangent +function mooncake_dict_to_cr_tangent(primal_dict::Dict, mt::Mooncake.MutableTangent) + result = Dict{Symbol,Any}() + raw_vals = mt.fields.vals + vals_tangent = if raw_vals isa Mooncake.PossiblyUninitTangent + Mooncake.is_init(raw_vals) ? raw_vals.tangent : return result + else + raw_vals + end + for (k, _) in primal_dict + idx = Base.ht_keyindex(primal_dict, k) + idx > 0 || continue + isassigned(vals_tangent, idx) || continue + vt = vals_tangent[idx] + cr_vt = val_to_cr(vt) + cr_vt isa ChainRulesCore.AbstractZero && continue + result[k] = cr_vt + end + return result +end +mooncake_dict_to_cr_tangent(::Dict, ::Mooncake.NoTangent) = ChainRulesCore.NoTangent() + +val_to_cr(x::AbstractArray{<:AbstractFloat}) = x +val_to_cr(::Mooncake.NoTangent) = ChainRulesCore.ZeroTangent() +val_to_cr(x::Mooncake.PossiblyUninitTangent) = + Mooncake.is_init(x) ? val_to_cr(x.tangent) : ChainRulesCore.ZeroTangent() +val_to_cr(x) = Mooncake.to_cr_tangent(x) + +# Positional: get_statistics(model, params) +@is_primitive Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_statistics), MacroModelling.ℳ, Vector{T}} where {T<:Base.IEEEFloat} +@is_primitive Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_statistics), MacroModelling.ℳ, Any} + +function Mooncake.rrule!!( + f_cd::CoDual{typeof(MacroModelling.get_statistics)}, + model_cd::CoDual{MacroModelling.ℳ}, + params_cd::CoDual{Vector{T}} +) where {T<:Base.IEEEFloat} + fargs = (f_cd, model_cd, params_cd) + primals = map(Mooncake.primal, fargs) + lazy_rdata = map(Mooncake.lazy_zero_rdata, primals) + y_primal, cr_pb = ChainRulesCore.rrule(primals...) + y_fdata = Mooncake.fdata(Mooncake.zero_tangent(y_primal)) + function pb!!(y_rdata) + cr_tangent = mooncake_dict_to_cr_tangent(y_primal, Mooncake.tangent(y_fdata, y_rdata)) + cr_dfargs = cr_pb(cr_tangent) + return map(fargs, lazy_rdata, cr_dfargs) do x, lr, cr_dx + Mooncake.increment_and_get_rdata!(Mooncake.tangent(x), Mooncake.instantiate(lr), cr_dx) + end + end + return CoDual(y_primal, y_fdata), pb!! +end + +# Keyword: Core.kwcall(kwargs, get_statistics, model, params) +@is_primitive Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), <:NamedTuple, typeof(MacroModelling.get_statistics), MacroModelling.ℳ, Vector{T}} where {T<:Base.IEEEFloat} +@is_primitive Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), <:NamedTuple, typeof(MacroModelling.get_statistics), MacroModelling.ℳ, Any} + +function Mooncake.rrule!!( + kwcall_cd::CoDual{typeof(Core.kwcall)}, + kwargs_cd::CoDual{<:NamedTuple}, + f_cd::CoDual{typeof(MacroModelling.get_statistics)}, + model_cd::CoDual{MacroModelling.ℳ}, + params_cd::CoDual{Vector{T}} +) where {T<:Base.IEEEFloat} + kw = Mooncake.primal(kwargs_cd) + model = Mooncake.primal(model_cd) + params = Mooncake.primal(params_cd) + # Call ChainRules rrule directly with kwargs (Core.kwcall has no rrule) + y_primal, cr_pb = ChainRulesCore.rrule(MacroModelling.get_statistics, model, params; kw...) + y_fdata = Mooncake.fdata(Mooncake.zero_tangent(y_primal)) + kwargs_lazy_rdata = Mooncake.lazy_zero_rdata(kw) + inner_fargs = (f_cd, model_cd, params_cd) + lazy_rdata = map(cd -> Mooncake.lazy_zero_rdata(Mooncake.primal(cd)), inner_fargs) + function pb!!(y_rdata) + cr_tangent = mooncake_dict_to_cr_tangent(y_primal, Mooncake.tangent(y_fdata, y_rdata)) + cr_dfargs = cr_pb(cr_tangent) + kwargs_rdata = Mooncake.increment_and_get_rdata!( + Mooncake.tangent(kwargs_cd), + Mooncake.instantiate(kwargs_lazy_rdata), + ChainRulesCore.NoTangent(), + ) + inner_rdata = map(inner_fargs, lazy_rdata, cr_dfargs) do x, lr, cr_dx + Mooncake.increment_and_get_rdata!(Mooncake.tangent(x), Mooncake.instantiate(lr), cr_dx) + end + return (NoRData(), kwargs_rdata, inner_rdata...) + end + return CoDual(y_primal, y_fdata), pb!! +end + +# ── Dict getindex primitive ── +# Without this, Mooncake tries to compile a tape through Dict's hash table internals +# (hashing, slot probing, Memory access), which takes extremely long. +# For mutable containers, fdata is accumulated in-place so the pullback is a no-op. +@is_primitive Mooncake.DefaultCtx Tuple{typeof(Base.getindex), <:Dict{Symbol}, Symbol} + +function Mooncake.rrule!!( + ::CoDual{typeof(Base.getindex)}, + dict_cd::CoDual{<:Dict{Symbol}}, + key_cd::CoDual{Symbol} +) + dict = Mooncake.primal(dict_cd) + key = Mooncake.primal(key_cd) + val = dict[key] + dict_fdata = Mooncake.tangent(dict_cd) + idx = Base.ht_keyindex(dict, key) + val_fdata = dict_fdata.fields.vals[idx] + function pb!!(::NoRData) + return NoRData(), NoRData(), NoRData() + end + return CoDual(val, val_fdata), pb!! +end + + +# ── clear_solution_caches! primitive ── +# This function mutates model caches and has no differentiable effect on outputs. +# Registering it as a primitive with zero gradient prevents Mooncake from tracing +# through its internals when it is called inside differentiable closures. +@is_primitive Mooncake.DefaultCtx Tuple{typeof(MacroModelling.clear_solution_caches!), MacroModelling.ℳ, Symbol} + +function Mooncake.rrule!!( + ::CoDual{typeof(MacroModelling.clear_solution_caches!)}, + model_cd::CoDual{MacroModelling.ℳ}, + alg_cd::CoDual{Symbol} +) + MacroModelling.clear_solution_caches!(Mooncake.primal(model_cd), Mooncake.primal(alg_cd)) + pb!!(::NoRData) = (NoRData(), NoRData(), NoRData()) + return CoDual(nothing, Mooncake.NoFData()), pb!! +end + +end # module MooncakeExt 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..2d3f56d86 100644 --- a/ext/StatsPlotsExt.jl +++ b/ext/StatsPlotsExt.jl @@ -2,8 +2,8 @@ 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: 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 MacroModelling: ParameterType, ℳ, Symbol_input, String_input, Tolerances, NsssTolerances, SolverTolerances, 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, SteadyStateFunctionType +import MacroModelling: DEFAULT_CACHING, DEFAULT_USE_WORKSPACES, 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 @@ -17,12 +17,615 @@ import Showoff import DataStructures: OrderedSet import SparseArrays: SparseMatrixCSC import NLopt +import Dates using DispatchDoctor -import MacroModelling: plot_irfs, plot_irf, plot_IRF, plot_simulations, plot_simulation, plot_solution, plot_girf, plot_conditional_forecast, plot_conditional_variance_decomposition, plot_forecast_error_variance_decomposition, plot_fevd, plot_model_estimates, plot_shock_decomposition, plotlyjs_backend, gr_backend, compare_args_and_kwargs, get_irf +import MacroModelling: plot_irfs, plot_irf, plot_IRF, plot_simulations, plot_simulation, plot_solution, plot_girf, plot_conditional_forecast, plot_conditional_variance_decomposition, plot_forecast_error_variance_decomposition, plot_fevd, plot_model_estimates, plot_shock_decomposition, plotlyjs_backend, gr_backend, get_irf import MacroModelling: plot_irfs!, plot_irf!, plot_IRF!, plot_girf!, plot_simulations!, plot_simulation!, plot_conditional_forecast!, plot_model_estimates!, plot_solution! + +# ────────────────────────────────────────────────────────────────────────────── +# Helper functions moved from core (only used by the plotting extension) +# ────────────────────────────────────────────────────────────────────────────── + +const HIGHER_ORDER_ALGORITHMS = (:second_order, :pruned_second_order, :third_order, :pruned_third_order) +const THIRD_ORDER_ALGORITHMS = (:third_order, :pruned_third_order) + +const TOL_DISPLAY_NAMES = Dict{Symbol,String}( + :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 accept. tol", + :acceptance_tol => "accept. tol", + :xtol => "xtol", + :ftol => "ftol", + :rel_xtol => "rel. xtol", + :droptol => "droptol", + :dependencies_tol => "dep. tol", +) + +function infer_step(x_axis::AbstractVector{T}) where {T<:Number} + x_axis[end] - x_axis[end-1] +end + +function infer_step(x_axis::AbstractVector{T}) where {T<:Dates.TimeType} + d1 = x_axis[end-1] + d2 = x_axis[end] + + # try to infer a monthly step if aligned by day-of-month + if Dates.day(d1) == Dates.day(d2) + m1 = 12 * Dates.year(d1) + Dates.month(d1) + m2 = 12 * Dates.year(d2) + Dates.month(d2) + mstep = m2 - m1 + if mstep != 0 + return Dates.Month(mstep) + end + end + + # fall back to the raw difference (in days, milliseconds, …) + return d2 - d1 +end + +function apply_custom_name(symbol::R, custom_names::AbstractDict{S, T})::R where {R <: Union{Symbol, String}, S, T} + # First, check for an exact match with the original symbol + if haskey(custom_names, symbol) + return R(custom_names[symbol]) + end + + # Handle cross-type check for exact match (String vs Symbol) + if symbol isa Symbol && haskey(custom_names, String(replace_indices_in_symbol(symbol))) + return R(custom_names[String(replace_indices_in_symbol(symbol))]) + elseif symbol isa String && haskey(custom_names, Symbol(symbol)) + return R(custom_names[Symbol(symbol)]) + end + + # If no exact match, strip lag operators and compare base names. + s_str = string(symbol) + lag_regex = r"^(.*)(ᴸ⁽.*⁾)$" + m = match(lag_regex, s_str) + + base_symbol_str, lag_part = if m !== nothing + (m.captures[1], m.captures[2]) + else + (s_str, "") + end + + for (key, value) in custom_names + key_str = string(key) + key_m = match(lag_regex, key_str) + + base_key_str = if key_m !== nothing + key_m.captures[1] + else + key_str + end + + if base_key_str == base_symbol_str + return R(string(value) * lag_part) + end + end + + return symbol +end + +function normalize_superscript(x::Symbol) + return normalize_superscript(string(x)) +end + +function normalize_superscript(x::AbstractString) + sub_map = Dict( + '₀' => '0', '₁' => '1', '₂' => '2', '₃' => '3', '₄' => '4', + '₅' => '5', '₆' => '6', '₇' => '7', '₈' => '8', '₉' => '9', + '₊' => '+', '₋' => '-', '₌' => '=', '₍' => '(', '₎' => ')', + 'ₐ' => 'a', 'ₑ' => 'e', 'ₕ' => 'h', 'ᵢ' => 'i', 'ⱼ' => 'j', + 'ₖ' => 'k', 'ₗ' => 'l', 'ₘ' => 'm', 'ₙ' => 'n', 'ₒ' => 'o', + 'ₚ' => 'p', 'ᵣ' => 'r', 'ₛ' => 's', 'ₜ' => 't', 'ᵤ' => 'u', + 'ᵥ' => 'v', 'ₓ' => 'x' + ) + super_map = Dict( + '⁰' => '0', '¹' => '1', '²' => '2', '³' => '3', '⁴' => '4', + '⁵' => '5', '⁶' => '6', '⁷' => '7', '⁸' => '8', '⁹' => '9', + '⁺' => '+', '⁻' => '-', '⁼' => '=', '⁽' => '(', '⁾' => ')', + 'ᵃ' => 'a', 'ᵇ' => 'b', 'ᶜ' => 'c', 'ᵈ' => 'd', 'ᵉ' => 'e', + 'ᶠ' => 'f', 'ᵍ' => 'g', 'ʰ' => 'h', 'ᶦ' => 'i', 'ʲ' => 'j', + 'ᵏ' => 'k', 'ˡ' => 'l', 'ᵐ' => 'm', 'ⁿ' => 'n', 'ᵒ' => 'o', + 'ᵖ' => 'p', 'ʳ' => 'r', 'ˢ' => 's', 'ᵗ' => 't', 'ᵘ' => 'u', + 'ᵛ' => 'v', 'ʷ' => 'w', 'ˣ' => 'x', 'ʸ' => 'y', 'ᶻ' => 'z' + ) + + buf = IOBuffer() + for c in x + if haskey(sub_map, c) + write(buf, sub_map[c]) + elseif haskey(super_map, c) + write(buf, super_map[c]) + else + write(buf, c) + end + end + return String(take!(buf)) +end + +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 + +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 + +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 + + # 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 + + return d +end + +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 _flatten_tol_dict(d::Dict; + names::Dict{Symbol,String} = TOL_DISPLAY_NAMES, + prefix::String = "") + result = Dict{String,Any}() + for (k, v) in d + seg = get(names, k, String(k)) + label = isempty(prefix) ? seg : prefix * " " * seg + if v isa Dict + merge!(result, _flatten_tol_dict(v; names = names, prefix = label)) + else + result[label] = v + end + end + return result +end + +function compare_args_and_kwargs(dicts::Vector{S}) where S <: Dict + N = length(dicts) + + if N ≤ 1 + 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}() + + all_keys = reduce(union, keys.(dicts)) + + for k in all_keys + if k in [:plot_data, :plot_type] + continue + end + + 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) + nested = compare_args_and_kwargs(vals) + if !isempty(nested) + diffs[k] = nested + end + + elseif all(v -> v isa KeyedArray, vals) + base = vals[1] + identical = all(v -> length(v) == length(base) && all(collect(v) .== collect(base)), vals[2:end]) + if !identical + diffs[k] = vals + end + + elseif all(v -> v isa AbstractArray, vals) + base = vals[1] + identical = all(v -> length(v) == length(base) && all(v .== base), vals[2:end]) + if !identical + diffs[k] = vals + end + + else + identical = all(v -> v == vals[1], vals[2:end]) + if !identical + diffs[k] = vals + end + end + end + + return diffs +end + +function flatten_tol_diff(diff::Dict; + names::Dict{Symbol,String} = TOL_DISPLAY_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 setup_plot_attributes(plot_attributes::Dict) + gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() + attrbts = !gr_back ? merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) : merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) + attributes = merge(attrbts, plot_attributes) + attributes_redux = copy(attributes) + delete!(attributes_redux, :framestyle) + return gr_back, attributes, attributes_redux +end + +function build_extended_palette(attributes_redux::Dict; total_pal_len::Int = 100, alpha_reduction_factor::Float64 = 0.7) + orig_pal = StatsPlots.palette(attributes_redux[:palette]) + mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette +end + +function process_rename_dictionary(rename_dictionary::AbstractDict, 𝓂::ℳ) + relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort + processed = Any[] + for k in relevant_keys + push!(processed, k => rename_dictionary[k]) + end + return processed +end + +function compute_diffdict(container::Vector{Dict}, ref_keys; include_label_in_reduced::Bool = true) + label_keys = include_label_in_reduced ? [:run_id, :label] : [:run_id] + reduced_vector = [ + Dict(k => d[k] for k in vcat(label_keys..., keys(DEFAULT_ARGS_AND_KWARGS_NAMES)...) if haskey(d, k)) + for d in container + ] + + diffdict = compare_args_and_kwargs(reduced_vector) + + grouped_by_model = Dict{Any, Vector{Dict}}() + + for d in container + model = d[:model_name] + d_sub = Dict(k => d[k] for k in setdiff(ref_keys, keys(DEFAULT_ARGS_AND_KWARGS_NAMES), [:tol]) if haskey(d, k)) + push!(get!(grouped_by_model, model, Vector{Dict}()), d_sub) + end + + model_names = unique([d[:model_name] for d in container]) + + for model in model_names + if length(grouped_by_model[model]) > 1 + diffdict_grouped = compare_args_and_kwargs(grouped_by_model[model]) + diffdict = merge_by_runid(diffdict, diffdict_grouped) + end + end + + return diffdict +end + +function annotate_param_diff!(annotate_diff_input, diffdict) + if haskey(diffdict, :parameters) + param_nms = diffdict[:parameters] |> keys |> collect |> sort + for param in param_nms + result = [x === nothing ? "" : x for x in diffdict[:parameters][param]] + push!(annotate_diff_input, String(param) => result) + end + end +end + +function annotate_rename_dict_diff!(annotate_diff_input, diffdict) + if haskey(diffdict, :rename_dictionary) + non_nothing_dicts = [d for d in diffdict[:rename_dictionary] if !isnothing(d) && length(d) > 0] + unique_dicts = unique(non_nothing_dicts) + rename_idx = Int[] + + for init in diffdict[:rename_dictionary] + if isnothing(init) || length(init) == 0 + push!(rename_idx, 0) + continue + end + + for (i,u) in enumerate(unique_dicts) + if u == init + push!(rename_idx, i) + continue + end + end + end + + push!(annotate_diff_input, "Rename dictionary" => [i > 0 ? "#$i" : "nothing" for i in rename_idx]) + end +end + +function annotate_tol_diff!(annotate_diff_input, container) + if length(container) > 1 + flat_tols = [_flatten_tol_dict(d[:tol]) for d in container] + shared_tol_keys = reduce(intersect, keys.(flat_tols)) + for fk in sort(collect(shared_tol_keys)) + fvals = [ft[fk] for ft in flat_tols] + if !all(v -> v == fvals[1], fvals[2:end]) + push!(annotate_diff_input, fk => fvals) + end + end + end +end + +function should_use_label_switch(annotate_diff_input, container) + ((length(annotate_diff_input) > 2) || (Dict(annotate_diff_input)["Plot label"] != collect(1:length(container)))) && length(container) > 1 +end + +function adjust_initial_state(initial_state, algorithm, 𝓂, SSS_delta, reference_steady_state) + nVars = 𝓂.constants.post_model_macro.nVars + unspecified_initial_state = initial_state == [0.0] + + if unspecified_initial_state + if algorithm == :pruned_second_order + return [zeros(nVars), zeros(nVars) - SSS_delta] + elseif algorithm == :pruned_third_order + return [zeros(nVars), zeros(nVars) - SSS_delta, zeros(nVars)] + else + return zeros(nVars) - SSS_delta + end + else + if initial_state isa Vector{Float64} + if algorithm == :pruned_second_order + return [initial_state - reference_steady_state[1:nVars], zeros(nVars) - SSS_delta] + elseif algorithm == :pruned_third_order + return [initial_state - reference_steady_state[1:nVars], zeros(nVars) - SSS_delta, zeros(nVars)] + else + return initial_state - reference_steady_state[1:nVars] + end + else + if algorithm ∉ [:pruned_second_order, :pruned_third_order] + @assert initial_state isa Vector{Float64} "The solution algorithm has one state vector: initial_state must be a Vector{Float64}." + end + return initial_state + end + end +end + +function push_if_no_duplicate!(container, args_and_kwargs, specific_keys; collect_compare_keys = Symbol[]) + no_duplicate = all( + !(all(( + all(get(dict, k, nothing) == args_and_kwargs[k] for k in specific_keys), + all( + k in collect_compare_keys ? + collect(get(dict, k, nothing)) == collect(get(args_and_kwargs, k, nothing)) : + get(dict, k, nothing) == get(args_and_kwargs, k, nothing) + for k in setdiff(keys(DEFAULT_ARGS_AND_KWARGS_NAMES), [:label]) + ) + ))) + for dict in container + ) + + if no_duplicate + push!(container, args_and_kwargs) + else + @info "Plot with same parameters already exists. Using previous plot data to create plot." + end +end + +function check_and_remove_duplicate!(container, specific_keys) + if length(container) > 1 + ref = container[end] + no_duplicate = all( + !(all(( + all(get(dict, k, nothing) == ref[k] for k in specific_keys), + all(get(dict, k, nothing) == get(ref, k, nothing) for k in setdiff(keys(DEFAULT_ARGS_AND_KWARGS_NAMES), [:label])) + ))) + for dict in container[1:end-1] + ) + + if !no_duplicate + @info "Plot with same parameters already exists. Using previous plot data to create plot." + pop!(container) + end + end +end + +function annotate_default_kwarg_diffs!(annotate_diff_input, args_and_kwargs, diffdict, exclude_keys) + has_shock_direction_diff = false + for k in setdiff(keys(args_and_kwargs), exclude_keys) + if haskey(diffdict, k) + push!(annotate_diff_input, DEFAULT_ARGS_AND_KWARGS_NAMES[k] => reduce(vcat, diffdict[k])) + if k == :negative_shock + has_shock_direction_diff = true + end + end + end + return has_shock_direction_diff +end + +function assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + default_model_name; + title_extra::String = "", + filename_extra::String = "", + legend_height = length(annotate_diff_input), + show_diff_table::Bool = false, + annotate_ss = nothing, + annotate_ss_page = nothing, + plt_lab_switch::Bool = false, + is_tail::Bool = false +) + ppp = StatsPlots.plot(pp...; attributes...) + + if haskey(diffdict, :model_name) + model_string = "multiple models" + model_string_filename = "multiple_models" + else + model_string = string(default_model_name) + model_string_filename = string(default_model_name) + end + + plot_title = "Model: " * model_string * title_extra * " (" * string(pane) * "/" * string(Int(ceil(n_subplots / plots_per_page))) * ")" + + plot_elements = [ppp, legend_plot] + layout_heights = [15, legend_height] + + if annotate_ss !== nothing && annotate_ss_page !== nothing + if plt_lab_switch + annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") + ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) + push!(plot_elements, ppp_input_diff) + push!(layout_heights, 5) + pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) + else + pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) + end + + push!(annotate_ss, annotate_ss_page) + + if length(annotate_ss[pane]) > 1 + annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady States") + ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) + push!(plot_elements, ppp_ss) + push!(layout_heights, 5) + end + else + if show_diff_table + annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") + ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) + push!(plot_elements, ppp_input_diff) + push!(layout_heights, 5) + end + end + + p = StatsPlots.plot(plot_elements..., + layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), + plot_title = plot_title; + attributes_redux...) + + push!(return_plots, p) + + if show_plots + display(p) + end + + if save_plots + if !isdir(save_plots_path) mkpath(save_plots_path) end + fn = save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename + if !isempty(filename_extra) + fn *= "__" * string(filename_extra) + end + fn *= "__" * string(pane) * "." * string(save_plots_format) + StatsPlots.savefig(p, fn) + end + + if !is_tail + pane += 1 + empty!(pp) + if annotate_ss_page !== nothing + empty!(annotate_ss_page) + end + end + + return pane +end + @stable default_mode = "disable" begin """ @@ -156,28 +759,23 @@ function plot_model_estimates(𝓂::ℳ, tol::Tolerances = Tolerances(), 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) + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], 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 - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) # write_parameters_input!(𝓂, parameters, verbose = verbose) @@ -227,13 +825,7 @@ function plot_model_estimates(𝓂::ℳ, shock_names_display = shock_names_display[shock_sort_perm] end - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) legend_columns = 1 @@ -297,7 +889,9 @@ function plot_model_estimates(𝓂::ℳ, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, tol = tol, - verbose = verbose) + verbose = verbose, + caching = caching, + use_workspaces = use_workspaces) forecast_data = collect(forecast_irf) @@ -308,13 +902,7 @@ function plot_model_estimates(𝓂::ℳ, extended_x_axis = vcat(x_axis, [last_x + i * period for i in 1:forecast_periods]) end - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 - - alpha_reduction_factor = 0.7 - - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) estimate_color = :navy @@ -340,18 +928,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, @@ -673,6 +1250,8 @@ function plot_model_estimates(𝓂::ℳ, end end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end @@ -809,28 +1388,23 @@ function plot_model_estimates!(𝓂::ℳ, tol::Tolerances = Tolerances(), 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) + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], 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) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + warn_irrelevant_tol(tol, algorithm; needs_covariance = filter == :kalman) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) # write_parameters_input!(𝓂, parameters, verbose = verbose) @@ -880,13 +1454,7 @@ function plot_model_estimates!(𝓂::ℳ, shock_names_display = shock_names_display[shock_sort_perm] end - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) legend_columns = 1 @@ -949,7 +1517,9 @@ function plot_model_estimates!(𝓂::ℳ, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, tol = tol, - verbose = verbose) + verbose = verbose, + caching = caching, + use_workspaces = use_workspaces) forecast_data = collect(forecast_irf) @@ -960,13 +1530,7 @@ function plot_model_estimates!(𝓂::ℳ, extended_x_axis = vcat(x_axis, [last_x + i * period for i in 1:forecast_periods]) end - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 - - alpha_reduction_factor = 0.7 - - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) estimate_color = :navy @@ -988,18 +1552,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, @@ -1019,54 +1572,11 @@ function plot_model_estimates!(𝓂::ℳ, :rename_dictionary => processed_rename_dictionary ) - no_duplicate = all( - !(all(( - get(dict, :parameters, nothing) == args_and_kwargs[:parameters], - get(dict, :rename_dictionary, nothing) == args_and_kwargs[:rename_dictionary], - # get(dict, :warmup_iterations, nothing) == args_and_kwargs[:warmup_iterations], - # get(dict, :smooth, nothing) == args_and_kwargs[:smooth], - all(k == :data ? collect(get(dict, k, nothing)) == collect(get(args_and_kwargs, k, nothing)) : get(dict, k, nothing) == get(args_and_kwargs, k, nothing) for k in setdiff(keys(DEFAULT_ARGS_AND_KWARGS_NAMES),[:label])) - ))) - for dict in model_estimates_active_plot_container - ) # "New plot must be different from previous plot. Use the version without ! to plot." - - if no_duplicate - push!(model_estimates_active_plot_container, args_and_kwargs) - else - @info "Plot with same parameters already exists. Using previous plot data to create plot." - end - - # 1. Keep only certain keys from each dictionary - reduced_vector = [ - Dict(k => d[k] for k in vcat(:run_id, keys(DEFAULT_ARGS_AND_KWARGS_NAMES)...) if haskey(d, k)) - for d in model_estimates_active_plot_container - ] - - diffdict = compare_args_and_kwargs(reduced_vector) - - # 2. Group the original vector by :model_name. Check difference for keys where they matter between models. Two different models might have different shocks so that difference is less important, but the same model with different shocks is a difference to highlight. - grouped_by_model = Dict{Any, Vector{Dict}}() - - for d in model_estimates_active_plot_container - model = d[:model_name] - d_sub = Dict(k => d[k] for k in setdiff(keys(args_and_kwargs), keys(DEFAULT_ARGS_AND_KWARGS_NAMES)) if haskey(d, k)) - push!(get!(grouped_by_model, model, Vector{Dict}()), d_sub) - end - - model_names = [] - - for d in model_estimates_active_plot_container - push!(model_names, d[:model_name]) - end - - model_names = unique(model_names) + push_if_no_duplicate!(model_estimates_active_plot_container, args_and_kwargs, + [:parameters, :rename_dictionary, :tol]; + collect_compare_keys = [:data]) - for model in model_names - if length(grouped_by_model[model]) > 1 - diffdict_grouped = compare_args_and_kwargs(grouped_by_model[model]) - diffdict = merge_by_runid(diffdict, diffdict_grouped) - end - end + diffdict = compute_diffdict(model_estimates_active_plot_container, keys(args_and_kwargs), include_label_in_reduced = false) annotate_ss = Vector{Pair{String, Any}}[] @@ -1078,13 +1588,7 @@ function plot_model_estimates!(𝓂::ℳ, len_diff = length(model_estimates_active_plot_container) - if haskey(diffdict, :parameters) - param_nms = diffdict[:parameters] |> keys |> collect |> sort - for param in param_nms - result = [x === nothing ? "" : x for x in diffdict[:parameters][param]] - push!(annotate_diff_input, String(param) => result) - end - end + annotate_param_diff!(annotate_diff_input, diffdict) common_axis = [] @@ -1105,28 +1609,7 @@ function plot_model_estimates!(𝓂::ℳ, push!(annotate_diff_input, "Data" => ["#$i" for i in data_idx]) end - rename_idx = Int[] - - if haskey(diffdict, :rename_dictionary) - non_nothing_dicts = [d for d in diffdict[:rename_dictionary] if !isnothing(d) && length(d) > 0] - unique_dicts = unique(non_nothing_dicts) - - for init in diffdict[:rename_dictionary] - if isnothing(init) || length(init) == 0 - push!(rename_idx, 0) - continue - end - - for (i,u) in enumerate(unique_dicts) - if u == init - push!(rename_idx,i) - continue - end - end - end - - push!(annotate_diff_input, "Rename dictionary" => [i > 0 ? "#$i" : "nothing" for i in rename_idx]) - end + annotate_rename_dict_diff!(annotate_diff_input, diffdict) # Determine common and combined x axis common_axis = mapreduce(k -> k[:x_axis], intersect, model_estimates_active_plot_container) @@ -1152,7 +1635,7 @@ function plot_model_estimates!(𝓂::ℳ, last_combined = combined_x_axis[end] if required_last_x > last_combined - xs = deepcopy(combined_x_axis) + xs = collect(combined_x_axis) next_x = last_combined while next_x < required_last_x @@ -1185,22 +1668,15 @@ function plot_model_estimates!(𝓂::ℳ, end end - for k in setdiff(keys(args_and_kwargs), - [ - :run_id, :parameters, :data, :data_in_levels, - :decomposition, :variables_to_plot, :data_in_deviations,:shocks_to_plot, :reference_steady_state, :x_axis, - :tol, :label, #:presample_periods, - :shocks, :shock_names, - :variables, :variable_names, - :rename_dictionary, :forecast_periods, :forecast_data, :extended_x_axis - # :periods, :quadratic_matrix_equation_algorithm, :sylvester_algorithm, :lyapunov_algorithm, - ] - ) + annotate_default_kwarg_diffs!(annotate_diff_input, args_and_kwargs, diffdict, + [:run_id, :parameters, :data, :data_in_levels, + :decomposition, :variables_to_plot, :data_in_deviations, :shocks_to_plot, :reference_steady_state, :x_axis, + :tol, :label, + :shocks, :shock_names, + :variables, :variable_names, + :rename_dictionary, :forecast_periods, :forecast_data, :extended_x_axis]) - if haskey(diffdict, k) - push!(annotate_diff_input, DEFAULT_ARGS_AND_KWARGS_NAMES[k] => reduce(vcat, diffdict[k])) - end - end + annotate_tol_diff!(annotate_diff_input, model_estimates_active_plot_container) if haskey(diffdict, :shock_names) if all(length.(diffdict[:shock_names]) .== 1) @@ -1215,7 +1691,7 @@ function plot_model_estimates!(𝓂::ℳ, joint_shocks = OrderedSet{String}() joint_variables = OrderedSet{String}() - plt_lab_switch = ((length(annotate_diff_input) > 2) || (Dict(annotate_diff_input)["Plot label"] != collect(1:length(model_estimates_active_plot_container)))) && length(model_estimates_active_plot_container) > 1 + plt_lab_switch = should_use_label_switch(annotate_diff_input, model_estimates_active_plot_container) for (i,k) in enumerate(model_estimates_active_plot_container) StatsPlots.plot!(legend_plot, [NaN], @@ -1537,138 +2013,37 @@ function plot_model_estimates!(𝓂::ℳ, else plot_count = 1 - ppp = StatsPlots.plot(pp...; attributes...) - - pl = StatsPlots.plot(framestyle = :none) - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = 𝓂.model_name - model_string_filename = 𝓂.model_name - end - - plot_title = "Model: "*model_string*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" - - plot_elements = [ppp, legend_plot] - - layout_heights = [15, length(annotate_diff_input)] - - if plt_lab_switch - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_input_diff) - - push!(layout_heights, 5) - - pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) - else - pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) - end - - push!(annotate_ss, annotate_ss_page) - - if length(annotate_ss[pane]) > 1 - annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady State") - - ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_ss) - - push!(layout_heights, 5) - end - - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux...) - - push!(return_plots,p) - - if show_plots - display(p) - end - - if save_plots - if !isdir(save_plots_path) mkpath(save_plots_path) end - - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * string(pane) * "." * string(save_plots_format)) - end - - pane += 1 - - annotate_ss_page = Pair{String,Any}[] - - pp = [] + pane = assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + 𝓂.model_name; + annotate_ss = annotate_ss, + annotate_ss_page = annotate_ss_page, + plt_lab_switch = plt_lab_switch, + ) end end if length(pp) > 0 - ppp = StatsPlots.plot(pp...; attributes...) - - pl = StatsPlots.plot(framestyle = :none) - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = 𝓂.model_name - model_string_filename = 𝓂.model_name - end - - plot_title = "Model: "*model_string*" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" - - plot_elements = [ppp, legend_plot] - - layout_heights = [15, length(annotate_diff_input)] - - if plt_lab_switch - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_input_diff) - - push!(layout_heights, 5) - - pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) - else - pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) - end - - push!(annotate_ss, annotate_ss_page) - - if length(annotate_ss[pane]) > 1 - annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady States") - - ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_ss) - - push!(layout_heights, 5) - end - - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux...) - - push!(return_plots,p) - - if show_plots - display(p) - end - - if save_plots - if !isdir(save_plots_path) mkpath(save_plots_path) end - - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * string(pane) * "." * string(save_plots_format)) - end + assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + 𝓂.model_name; + annotate_ss = annotate_ss, + annotate_ss_page = annotate_ss_page, + plt_lab_switch = plt_lab_switch, + is_tail = true, + ) end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end @@ -1765,27 +2140,22 @@ function plot_irf(𝓂::ℳ; 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(𝓂)) + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], 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]) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + warn_irrelevant_tol(tol, algorithm; needs_covariance = false) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) shocks, negative_shock, shock_size, periods_extended, shock_idx, shock_history = process_shocks_input(shocks, negative_shock, shock_size, periods, 𝓂) @@ -1809,31 +2179,7 @@ function plot_irf(𝓂::ℳ; initial_state_input = copy(initial_state) - unspecified_initial_state = initial_state == [0.0] - - if unspecified_initial_state - if algorithm == :pruned_second_order - initial_state = [zeros(𝓂.constants.post_model_macro.nVars), zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta] - elseif algorithm == :pruned_third_order - initial_state = [zeros(𝓂.constants.post_model_macro.nVars), zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta, zeros(𝓂.constants.post_model_macro.nVars)] - else - initial_state = zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta - end - else - if initial_state isa Vector{Float64} - if algorithm == :pruned_second_order - initial_state = [initial_state - reference_steady_state[1:𝓂.constants.post_model_macro.nVars], zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta] - elseif algorithm == :pruned_third_order - initial_state = [initial_state - reference_steady_state[1:𝓂.constants.post_model_macro.nVars], zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta, zeros(𝓂.constants.post_model_macro.nVars)] - else - initial_state = initial_state - reference_steady_state[1:𝓂.constants.post_model_macro.nVars] - end - else - if algorithm ∉ [:pruned_second_order, :pruned_third_order] - @assert initial_state isa Vector{Float64} "The solution algorithm has one state vector: initial_state must be a Vector{Float64}." - end - end - end + initial_state = adjust_initial_state(initial_state, algorithm, 𝓂, SSS_delta, reference_steady_state) if occasionally_binding_constraints @@ -1903,14 +2249,9 @@ function plot_irf(𝓂::ℳ; var_idx = var_idx[var_sort_perm] variable_names_display = variable_names_display[var_sort_perm] - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort Y = Y[var_sort_perm, :, :] - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) while length(irf_active_plot_container) > 0 pop!(irf_active_plot_container) @@ -1933,16 +2274,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, @@ -1956,13 +2288,7 @@ function plot_irf(𝓂::ℳ; push!(irf_active_plot_container, args_and_kwargs) - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 - - alpha_reduction_factor = 0.7 - - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) return_plots = [] @@ -2057,6 +2383,8 @@ function plot_irf(𝓂::ℳ; end end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end @@ -2457,37 +2785,26 @@ function plot_irf!(𝓂::ℳ; 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(𝓂)) + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + @assert plot_type ∈ [:compare, :stack] "plot_type must be either :compare or :stack" 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]) - - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) - - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 + 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]) - alpha_reduction_factor = 0.7 + warn_irrelevant_tol(tol, algorithm; needs_covariance = false) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) shocks, negative_shock, shock_size, periods_extended, shock_idx, shock_history = process_shocks_input(shocks, negative_shock, shock_size, periods, 𝓂) @@ -2511,31 +2828,7 @@ function plot_irf!(𝓂::ℳ; initial_state_input = copy(initial_state) - unspecified_initial_state = initial_state == [0.0] - - if unspecified_initial_state - if algorithm == :pruned_second_order - initial_state = [zeros(𝓂.constants.post_model_macro.nVars), zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta] - elseif algorithm == :pruned_third_order - initial_state = [zeros(𝓂.constants.post_model_macro.nVars), zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta, zeros(𝓂.constants.post_model_macro.nVars)] - else - initial_state = zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta - end - else - if initial_state isa Vector{Float64} - if algorithm == :pruned_second_order - initial_state = [initial_state - reference_steady_state[1:𝓂.constants.post_model_macro.nVars], zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta] - elseif algorithm == :pruned_third_order - initial_state = [initial_state - reference_steady_state[1:𝓂.constants.post_model_macro.nVars], zeros(𝓂.constants.post_model_macro.nVars) - SSS_delta, zeros(𝓂.constants.post_model_macro.nVars)] - else - initial_state = initial_state - reference_steady_state[1:𝓂.constants.post_model_macro.nVars] - end - else - if algorithm ∉ [:pruned_second_order, :pruned_third_order] - @assert initial_state isa Vector{Float64} "The solution algorithm has one state vector: initial_state must be a Vector{Float64}." - end - end - end + initial_state = adjust_initial_state(initial_state, algorithm, 𝓂, SSS_delta, reference_steady_state) if occasionally_binding_constraints @@ -2593,13 +2886,7 @@ function plot_irf!(𝓂::ℳ; variable_names_display = variable_names_display[var_sort_perm] Y = Y[var_sort_perm, :, :] - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) args_and_kwargs = Dict(:run_id => length(irf_active_plot_container) + 1, :model_name => 𝓂.model_name, @@ -2618,16 +2905,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, @@ -2638,55 +2916,10 @@ function plot_irf!(𝓂::ℳ; :rename_dictionary => processed_rename_dictionary ) - no_duplicate = all( - !(all(( - get(dict, :parameters, nothing) == args_and_kwargs[:parameters], - get(dict, :shock_names, nothing) == args_and_kwargs[:shock_names], - get(dict, :rename_dictionary, nothing) == args_and_kwargs[:rename_dictionary], - get(dict, :shocks, nothing) == args_and_kwargs[:shocks], - get(dict, :initial_state, nothing) == args_and_kwargs[:initial_state], - all(get(dict, k, nothing) == get(args_and_kwargs, k, nothing) for k in setdiff(keys(DEFAULT_ARGS_AND_KWARGS_NAMES),[:label])) - ))) - for dict in irf_active_plot_container - )# "New plot must be different from previous plot. Use the version without ! to plot." - - if no_duplicate - push!(irf_active_plot_container, args_and_kwargs) - else - @info "Plot with same parameters already exists. Using previous plot data to create plot." - end - - # 1. Keep only certain keys from each dictionary - reduced_vector = [ - Dict(k => d[k] for k in vcat(:run_id, :label, keys(DEFAULT_ARGS_AND_KWARGS_NAMES)...) if haskey(d, k)) - for d in irf_active_plot_container - ] - - diffdict = compare_args_and_kwargs(reduced_vector) - - # 2. Group the original vector by :model_name - grouped_by_model = Dict{Any, Vector{Dict}}() - - for d in irf_active_plot_container - model = d[:model_name] - d_sub = Dict(k => d[k] for k in setdiff(keys(args_and_kwargs), keys(DEFAULT_ARGS_AND_KWARGS_NAMES)) if haskey(d, k)) - push!(get!(grouped_by_model, model, Vector{Dict}()), d_sub) - end - - model_names = [] - - for d in irf_active_plot_container - push!(model_names, d[:model_name]) - end - - model_names = unique(model_names) + push_if_no_duplicate!(irf_active_plot_container, args_and_kwargs, + [:parameters, :shock_names, :rename_dictionary, :shocks, :initial_state, :tol]) - for model in model_names - if length(grouped_by_model[model]) > 1 - diffdict_grouped = compare_args_and_kwargs(grouped_by_model[model]) - diffdict = merge_by_runid(diffdict, diffdict_grouped) - end - end + diffdict = compute_diffdict(irf_active_plot_container, keys(args_and_kwargs)) # @assert haskey(diffdict, :parameters) || haskey(diffdict, :shock_names) || haskey(diffdict, :initial_state) || any(haskey.(Ref(diffdict), keys(DEFAULT_ARGS_AND_KWARGS_NAMES))) "New plot must be different from previous plot. Use the version without ! to plot." @@ -2700,13 +2933,7 @@ function plot_irf!(𝓂::ℳ; len_diff = length(irf_active_plot_container) - if haskey(diffdict, :parameters) - param_nms = diffdict[:parameters] |> keys |> collect |> sort - for param in param_nms - result = [x === nothing ? "" : x for x in diffdict[:parameters][param]] - push!(annotate_diff_input, String(param) => result) - end - end + annotate_param_diff!(annotate_diff_input, diffdict) if haskey(diffdict, :shocks) # Build labels where matrices receive stable indices by content @@ -2775,50 +3002,20 @@ function plot_irf!(𝓂::ℳ; push!(annotate_diff_input, "Initial state" => labels) end - rename_idx = Int[] - - if haskey(diffdict, :rename_dictionary) - non_nothing_dicts = [d for d in diffdict[:rename_dictionary] if !isnothing(d) && length(d) > 0] - unique_dicts = unique(non_nothing_dicts) - - for init in diffdict[:rename_dictionary] - if isnothing(init) || length(init) == 0 - push!(rename_idx, 0) - continue - end - - for (i,u) in enumerate(unique_dicts) - if u == init - push!(rename_idx,i) - continue - end - end - end - - push!(annotate_diff_input, "Rename dictionary" => [i > 0 ? "#$i" : "nothing" for i in rename_idx]) - end + annotate_rename_dict_diff!(annotate_diff_input, diffdict) same_shock_direction = true - for k in setdiff(keys(args_and_kwargs), - [ - :run_id, :parameters, :plot_data, :tol, :reference_steady_state, :initial_state, :label, - :shocks, :shock_names, - :variables, :variable_names, - :rename_dictionary, - # :periods, :quadratic_matrix_equation_algorithm, :sylvester_algorithm, :lyapunov_algorithm, - ] - ) - - if haskey(diffdict, k) - push!(annotate_diff_input, DEFAULT_ARGS_AND_KWARGS_NAMES[k] => reduce(vcat,diffdict[k])) - - if k == :negative_shock - same_shock_direction = false - end - end + if annotate_default_kwarg_diffs!(annotate_diff_input, args_and_kwargs, diffdict, + [:run_id, :parameters, :plot_data, :tol, :reference_steady_state, :initial_state, :label, + :shocks, :shock_names, + :variables, :variable_names, + :rename_dictionary]) + same_shock_direction = false end + annotate_tol_diff!(annotate_diff_input, irf_active_plot_container) + legend_plot = StatsPlots.plot(framestyle = :none, @@ -2830,7 +3027,7 @@ function plot_irf!(𝓂::ℳ; single_shock_per_irf = true max_periods = 0 - plt_lab_switch = ((length(annotate_diff_input) > 2) || (Dict(annotate_diff_input)["Plot label"] != collect(1:length(irf_active_plot_container)))) && length(irf_active_plot_container) > 1 + plt_lab_switch = should_use_label_switch(annotate_diff_input, irf_active_plot_container) for (i,k) in enumerate(irf_active_plot_container) if plot_type == :stack StatsPlots.bar!(legend_plot, @@ -2964,70 +3161,20 @@ function plot_irf!(𝓂::ℳ; shock_name = shock end - ppp = StatsPlots.plot(pp...; attributes...) - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = 𝓂.model_name - model_string_filename = 𝓂.model_name - end - - plot_title = "Model: "*model_string*" " * shock_dir * shock_string *" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" - - plot_elements = [ppp, legend_plot] - - layout_heights = [15,1] - - if plt_lab_switch - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_input_diff) - - push!(layout_heights, 5) - - pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) - else - pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) - end - - push!(annotate_ss, annotate_ss_page) - - if length(annotate_ss[pane]) > 1 - annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady States") - - ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_ss) - - push!(layout_heights, 5) - end - - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux...) - - push!(return_plots,p) - - if show_plots - display(p) - end - - if save_plots - if !isdir(save_plots_path) mkpath(save_plots_path) end - - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * shock_name * "__" * string(pane) * "." * string(save_plots_format)) - end - - pane += 1 - - annotate_ss_page = Pair{String,Any}[] - - pp = [] + pane = assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + 𝓂.model_name; + title_extra = " " * shock_dir * shock_string, + filename_extra = shock_name, + legend_height = 1, + annotate_ss = annotate_ss, + annotate_ss_page = annotate_ss_page, + plt_lab_switch = plt_lab_switch, + ) end end @@ -3055,64 +3202,21 @@ function plot_irf!(𝓂::ℳ; shock_name = shock end - ppp = StatsPlots.plot(pp...; attributes...) - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = 𝓂.model_name - model_string_filename = 𝓂.model_name - end - - plot_title = "Model: "*model_string*" " * shock_dir * shock_string *" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" - - plot_elements = [ppp, legend_plot] - - layout_heights = [15,1] - - if plt_lab_switch - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_input_diff) - - push!(layout_heights, 5) - - pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) - else - pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) - end - - push!(annotate_ss, annotate_ss_page) - - if length(annotate_ss[pane]) > 1 - annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady States") - - ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_ss) - - push!(layout_heights, 5) - end - - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux...) - - push!(return_plots,p) - - if show_plots - display(p) - end - - if save_plots - if !isdir(save_plots_path) mkpath(save_plots_path) end - - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * shock_name * "__" * string(pane) * "." * string(save_plots_format)) - end + assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + 𝓂.model_name; + title_extra = " " * shock_dir * shock_string, + filename_extra = shock_name, + legend_height = 1, + annotate_ss = annotate_ss, + annotate_ss_page = annotate_ss_page, + plt_lab_switch = plt_lab_switch, + is_tail = true, + ) end annotate_ss = Vector{Pair{String, Any}}[] @@ -3120,6 +3224,8 @@ function plot_irf!(𝓂::ℳ; annotate_ss_page = Pair{String,Any}[] end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end @@ -3501,25 +3607,19 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; extra_legend_space::Float64 = DEFAULT_EXTRA_LEGEND_SPACE, verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM) + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) fevds = get_conditional_variance_decomposition(𝓂, periods = 1:periods, @@ -3527,7 +3627,9 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; steady_state_function = steady_state_function, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - tol = tol) + tol = tol, + caching = caching, + use_workspaces = use_workspaces) variables = variables isa String_input ? variables .|> Meta.parse .|> replace_indices : variables @@ -3567,13 +3669,7 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; end end - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 - - alpha_reduction_factor = 0.7 - - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) n_subplots = length(var_idx) pp = [] @@ -3665,6 +3761,8 @@ function plot_conditional_variance_decomposition(𝓂::ℳ; end end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end @@ -3772,28 +3870,23 @@ function plot_solution(𝓂::ℳ, tol::Tolerances = Tolerances(), 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) + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], 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) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) state = state isa Symbol ? state : state |> Meta.parse |> replace_indices @@ -3814,7 +3907,9 @@ function plot_solution(𝓂::ℳ, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, tol = tol, - verbose = verbose) + verbose = verbose, + caching = caching, + use_workspaces = use_workspaces) SS_and_std[:non_stochastic_steady_state] = SS_and_std[:non_stochastic_steady_state] isa KeyedArray ? axiskeys(SS_and_std[:non_stochastic_steady_state],1) isa Vector{String} ? rekey(SS_and_std[:non_stochastic_steady_state], 1 => axiskeys(SS_and_std[:non_stochastic_steady_state],1).|> x->Symbol.(replace.(x, "{" => "◖", "}" => "◗"))) : SS_and_std[:non_stochastic_steady_state] : SS_and_std[:non_stochastic_steady_state] @@ -3837,13 +3932,7 @@ function plot_solution(𝓂::ℳ, vars_sort_perm = sortperm(variable_names_display, by = normalize_superscript) vars_to_plot = vars_to_plot[vars_sort_perm] - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) state_range = collect(range(-SS_and_std[:standard_deviation](state), SS_and_std[:standard_deviation](state), 100)) * σ @@ -3864,7 +3953,9 @@ function plot_solution(𝓂::ℳ, tol = opts.tol, verbose = opts.verbose, quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) + sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³], + caching = caching, + use_workspaces = use_workspaces) full_SS_current = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : relevant_SS(s) for s in full_NSSS] @@ -3873,7 +3964,9 @@ function plot_solution(𝓂::ℳ, tol = opts.tol, verbose = opts.verbose, quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) + sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³], + caching = caching, + use_workspaces = use_workspaces) NSSS = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : NSSS_SS(s) for s in full_NSSS] @@ -3891,7 +3984,7 @@ function plot_solution(𝓂::ℳ, initial_state = collect(full_SS_current) .+ state_selector * x end - push!(var_state_range, get_irf(𝓂, algorithm = algorithm, periods = 1, ignore_obc = ignore_obc, initial_state = initial_state, shocks = :none, levels = true, variables = :all)[:,1,1] |> collect) + push!(var_state_range, get_irf(𝓂, algorithm = algorithm, periods = 1, ignore_obc = ignore_obc, initial_state = initial_state, shocks = :none, levels = true, variables = :all, caching = caching, use_workspaces = use_workspaces)[:,1,1] |> collect) end var_state_range = hcat(var_state_range...) @@ -3924,6 +4017,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, @@ -3935,6 +4029,8 @@ function plot_solution(𝓂::ℳ, push!(solution_active_plot_container, args_and_kwargs) # Generate plots from container + if !use_workspaces 𝓂.workspaces = orig_ws end + return _plot_solution_from_container(; show_plots = show_plots, save_plots = save_plots, @@ -3971,22 +4067,9 @@ function _plot_solution_from_container(; push!(joint_states, string(apply_custom_name.(container[:state], Ref(Dict(container[:rename_dictionary]))))) end - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - attributes_redux = copy(attributes) - delete!(attributes_redux, :framestyle) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - total_pal_len = 100 - alpha_reduction_factor = 0.7 - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) # Create comparison of containers to detect differences # Keep relevant keys for comparison: model_name, state, parameters, algorithm, ignore_obc, label @@ -3994,57 +4077,13 @@ function _plot_solution_from_container(; diffdict = Dict{Symbol,Any}() if length(solution_active_plot_container) > 1 - no_duplicate = all( - !(all(( - get(dict, :parameters, nothing) == solution_active_plot_container[end][:parameters], - get(dict, :model_name, nothing) == solution_active_plot_container[end][:model_name], - get(dict, :algorithm, nothing) == solution_active_plot_container[end][:algorithm], - get(dict, :ignore_obc, nothing) == solution_active_plot_container[end][:ignore_obc], - all(get(dict, k, nothing) == get(solution_active_plot_container[end], k, nothing) for k in setdiff(keys(DEFAULT_ARGS_AND_KWARGS_NAMES),[:label])) - ))) - for dict in solution_active_plot_container[1:end-1] - ) # "New plot must be different from previous plot. Use the version without ! to plot." - - if !no_duplicate - @info "Plot with same parameters already exists. Using previous plot data to create plot." - - pop!(solution_active_plot_container) - end + check_and_remove_duplicate!(solution_active_plot_container, + [:parameters, :model_name, :algorithm, :ignore_obc, :tol]) if length(solution_active_plot_container) == 0 diffdict[:label] = [solution_active_plot_container[1][:label]] else - # 1. Keep only certain keys from each dictionary - reduced_vector = [ - Dict(k => d[k] for k in vcat(:run_id, :label, keys(DEFAULT_ARGS_AND_KWARGS_NAMES)...) if haskey(d, k)) - for d in solution_active_plot_container - ] - - diffdict = compare_args_and_kwargs(reduced_vector) - - # 2. Group the original vector by :model_name - grouped_by_model = Dict{Any, Vector{Dict}}() - - for d in solution_active_plot_container#[1:end-1] - model = d[:model_name] - d_sub = Dict(k => d[k] for k in setdiff(keys(solution_active_plot_container[end]), keys(DEFAULT_ARGS_AND_KWARGS_NAMES)) if haskey(d, k)) - push!(get!(grouped_by_model, model, Vector{Dict}()), d_sub) - end - - model_names = [] - - for d in solution_active_plot_container - push!(model_names, d[:model_name]) - end - - model_names = unique(model_names) - - for model in model_names - if length(grouped_by_model[model]) > 1 - diffdict_grouped = compare_args_and_kwargs(grouped_by_model[model]) - diffdict = merge_by_runid(diffdict, diffdict_grouped) - end - end + diffdict = compute_diffdict(solution_active_plot_container, keys(solution_active_plot_container[end])) end else # For single container, create a diffdict with just the label @@ -4066,49 +4105,24 @@ function _plot_solution_from_container(; push!(annotate_diff_input, "State" => reduce(vcat, diffdict[:state])) end - # Add algorithm if different - if haskey(diffdict, :algorithm) - algo_labels = [String(a) for a in diffdict[:algorithm]] - push!(annotate_diff_input, "Algorithm" => algo_labels) - end - - # Add parameters if different - if haskey(diffdict, :parameters) - param_nms = diffdict[:parameters] |> keys |> collect |> sort - for param in param_nms - result = [x === nothing ? "" : x for x in diffdict[:parameters][param]] - push!(annotate_diff_input, String(param) => result) - end - end - - rename_idx = Int[] - - if haskey(diffdict, :rename_dictionary) - non_nothing_dicts = [d for d in diffdict[:rename_dictionary] if !isnothing(d) && length(d) > 0] - unique_dicts = unique(non_nothing_dicts) - - for init in diffdict[:rename_dictionary] - if isnothing(init) || length(init) == 0 - push!(rename_idx, 0) - continue - end - - for (i,u) in enumerate(unique_dicts) - if u == init - push!(rename_idx,i) - continue - end - end - end - - push!(annotate_diff_input, "Rename dictionary" => [i > 0 ? "#$i" : "nothing" for i in rename_idx]) - end + # Add algorithm if different + if haskey(diffdict, :algorithm) + algo_labels = [String(a) for a in diffdict[:algorithm]] + push!(annotate_diff_input, "Algorithm" => algo_labels) + end + + # Add parameters if different + annotate_param_diff!(annotate_diff_input, diffdict) + + annotate_rename_dict_diff!(annotate_diff_input, diffdict) # Add ignore_obc if different if haskey(diffdict, :ignore_obc) push!(annotate_diff_input, "Ignore OBC" => reduce(vcat, diffdict[:ignore_obc])) end + annotate_tol_diff!(annotate_diff_input, solution_active_plot_container) + # 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) @@ -4117,7 +4131,7 @@ function _plot_solution_from_container(; # Create legend with 2 columns so dynamics and steady state entries are side by side legend_plot = StatsPlots.plot(framestyle = :none, legend = :inside, legend_columns = 2) - plt_lab_switch = ((length(annotate_diff_input) > 2) || (Dict(annotate_diff_input)["Plot label"] != collect(1:length(solution_active_plot_container)))) && length(solution_active_plot_container) > 1 + plt_lab_switch = should_use_label_switch(annotate_diff_input, solution_active_plot_container) if plt_lab_switch # Multiple differences - use custom labels or plot labels for (i, container) in enumerate(solution_active_plot_container) @@ -4300,104 +4314,42 @@ function _plot_solution_from_container(; if !(plot_count % plots_per_page == 0) plot_count += 1 else - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = solution_active_plot_container[1][:model_name] - model_string_filename = solution_active_plot_container[1][:model_name] - end - plot_count = 1 - - ppp = StatsPlots.plot(pp...; attributes...) - - # Build plot elements array - plot_elements = [ppp, legend_plot] - layout_heights = [15, length(annotate_diff_input)] - - # Add relevant input differences table if multiple inputs differ - if plt_lab_switch || (any_custom_label && len_diff > 1) - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - push!(plot_elements, ppp_input_diff) - push!(layout_heights, 5) - end - - # Create plot title including state info + state_string = length(joint_states) > 1 ? " State: " * replace_indices_in_symbol(Symbol(state)) : "" - plot_title = "Model: " * model_string * state_string * " (" * string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page))) * ")" - - # Create final plot with appropriate layout - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux... + state_name = replace_indices_in_symbol(Symbol(state)) + + pane = assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + solution_active_plot_container[1][:model_name]; + title_extra = state_string, + filename_extra = string(state_name), + show_diff_table = plt_lab_switch || (any_custom_label && len_diff > 1), ) - - push!(return_plots, p) - - if show_plots - display(p) - end - - if save_plots - if !isdir(save_plots_path) mkpath(save_plots_path) end - state_name = replace_indices_in_symbol(Symbol(state)) - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * state_name * "__" * string(pane) * "." * string(save_plots_format)) - end - - pane += 1 - pp = [] end end # Handle remaining plots for this state if length(pp) > 0 - ppp = StatsPlots.plot(pp...; attributes...) - - # Build plot elements array - plot_elements = [ppp, legend_plot] - layout_heights = [15, length(annotate_diff_input)] - - # Add relevant input differences table if multiple inputs differ - if plt_lab_switch || (any_custom_label && len_diff > 1) - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - push!(plot_elements, ppp_input_diff) - push!(layout_heights, 5) - end - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = solution_active_plot_container[1][:model_name] - model_string_filename = solution_active_plot_container[1][:model_name] - end - - # Create plot title including state info state_string = length(joint_states) > 1 ? " State: " * replace_indices_in_symbol(Symbol(state)) : "" - plot_title = "Model: " * model_string * state_string * " (" * string(pane) * "/" * string(Int(ceil(n_subplots/plots_per_page))) * ")" - - # Create final plot with appropriate layout - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux... + state_name = replace_indices_in_symbol(Symbol(state)) + + assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + solution_active_plot_container[1][:model_name]; + title_extra = state_string, + filename_extra = string(state_name), + show_diff_table = plt_lab_switch || (any_custom_label && len_diff > 1), + is_tail = true, ) - - push!(return_plots, p) - - if show_plots - display(p) - end - - if save_plots - if !isdir(save_plots_path) mkpath(save_plots_path) end - state_name = replace_indices_in_symbol(Symbol(state)) - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * state_name * "__" * string(pane) * "." * string(save_plots_format)) - end end end # End of state loop @@ -4498,9 +4450,15 @@ function plot_solution!(𝓂::ℳ, tol::Tolerances = Tolerances(), 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) + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + # Do NOT clear container - add to existing opts = merge_calculation_options(tol = tol, verbose = verbose, @@ -4509,19 +4467,8 @@ 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) - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) state = state isa Symbol ? state : state |> Meta.parse |> replace_indices @@ -4542,7 +4489,9 @@ function plot_solution!(𝓂::ℳ, sylvester_algorithm = sylvester_algorithm, lyapunov_algorithm = lyapunov_algorithm, tol = tol, - verbose = verbose) + verbose = verbose, + caching = caching, + use_workspaces = use_workspaces) SS_and_std[:non_stochastic_steady_state] = SS_and_std[:non_stochastic_steady_state] isa KeyedArray ? axiskeys(SS_and_std[:non_stochastic_steady_state],1) isa Vector{String} ? rekey(SS_and_std[:non_stochastic_steady_state], 1 => axiskeys(SS_and_std[:non_stochastic_steady_state],1).|> x->Symbol.(replace.(x, "{" => "◖", "}" => "◗"))) : SS_and_std[:non_stochastic_steady_state] : SS_and_std[:non_stochastic_steady_state] @@ -4565,13 +4514,7 @@ function plot_solution!(𝓂::ℳ, vars_sort_perm = sortperm(variable_names_display, by = normalize_superscript) vars_to_plot = vars_to_plot[vars_sort_perm] - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) state_range = collect(range(-SS_and_std[:standard_deviation](state), SS_and_std[:standard_deviation](state), 100)) * σ @@ -4587,7 +4530,9 @@ function plot_solution!(𝓂::ℳ, tol = opts.tol, verbose = opts.verbose, quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) + sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³], + caching = caching, + use_workspaces = use_workspaces) full_SS_current = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : relevant_SS(s) for s in full_NSSS] @@ -4596,7 +4541,9 @@ function plot_solution!(𝓂::ℳ, tol = opts.tol, verbose = opts.verbose, quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) + sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³], + caching = caching, + use_workspaces = use_workspaces) NSSS = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : NSSS_SS(s) for s in full_NSSS] @@ -4614,7 +4561,7 @@ function plot_solution!(𝓂::ℳ, initial_state = collect(full_SS_current) .+ state_selector * x end - push!(var_state_range, get_irf(𝓂, algorithm = algorithm, periods = 1, ignore_obc = ignore_obc, initial_state = initial_state, shocks = :none, levels = true, variables = :all)[:,1,1] |> collect) + push!(var_state_range, get_irf(𝓂, algorithm = algorithm, periods = 1, ignore_obc = ignore_obc, initial_state = initial_state, shocks = :none, levels = true, variables = :all, caching = caching, use_workspaces = use_workspaces)[:,1,1] |> collect) end var_state_range = hcat(var_state_range...) @@ -4647,6 +4594,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, @@ -4657,6 +4605,8 @@ function plot_solution!(𝓂::ℳ, push!(solution_active_plot_container, args_and_kwargs) + if !use_workspaces 𝓂.workspaces = orig_ws end + # Generate plots from container return _plot_solution_from_container(; show_plots = show_plots, @@ -4784,22 +4734,16 @@ function plot_conditional_forecast(𝓂::ℳ, 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(𝓂)) + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) initial_state_input = copy(initial_state) @@ -4823,7 +4767,10 @@ function plot_conditional_forecast(𝓂::ℳ, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm = sylvester_algorithm, tol = tol, - verbose = verbose) + verbose = verbose, + caching = caching, + use_workspaces = use_workspaces) + warn_irrelevant_tol(tol, algorithm; needs_covariance = true) periods += max(size(conditions,2), isnothing(shocks) ? 1 : size(shocks,2)) @@ -4851,7 +4798,9 @@ function plot_conditional_forecast(𝓂::ℳ, tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - sylvester_algorithm = sylvester_algorithm) + sylvester_algorithm = sylvester_algorithm, + caching = caching, + use_workspaces = use_workspaces) relevant_SS = relevant_SS isa KeyedArray ? axiskeys(relevant_SS,1) isa Vector{String} ? rekey(relevant_SS, 1 => axiskeys(relevant_SS,1) .|> Meta.parse .|> replace_indices) : relevant_SS : relevant_SS @@ -4942,13 +4891,7 @@ function plot_conditional_forecast(𝓂::ℳ, full_shock_sort_perm = sortperm(full_shock_names_display, by = normalize_superscript) # Process rename dictionary to only include relevant keys in sorted order - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) # Combine sorted indices combined_sort_perm = vcat(var_sort_perm, (length(variable_names_display) .+ (1:length(shock_names_display)))[shock_sort_perm]) @@ -4981,16 +4924,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, @@ -5004,13 +4938,7 @@ function plot_conditional_forecast(𝓂::ℳ, push!(conditional_forecast_active_plot_container, args_and_kwargs) - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 - - alpha_reduction_factor = 0.7 - - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) n_subplots = length(var_idx) pp = [] @@ -5123,6 +5051,8 @@ function plot_conditional_forecast(𝓂::ℳ, end end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end @@ -5245,24 +5175,18 @@ function plot_conditional_forecast!(𝓂::ℳ, 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(𝓂)) + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching invalidate_cache_validity!(𝓂) end + orig_ws = 𝓂.workspaces + if !use_workspaces 𝓂.workspaces = fresh_workspaces(orig_ws) end + @assert plot_type ∈ [:compare, :stack] "plot_type must be either :compare or :stack" - gr_back = StatsPlots.backend() == StatsPlots.Plots.GRBackend() - - if !gr_back - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict(:framestyle => :box)) - else - attrbts = merge(DEFAULT_PLOT_ATTRIBUTES, Dict()) - end - - attributes = merge(attrbts, plot_attributes) - - attributes_redux = copy(attributes) - - delete!(attributes_redux, :framestyle) + gr_back, attributes, attributes_redux = setup_plot_attributes(plot_attributes) initial_state_input = copy(initial_state) @@ -5286,8 +5210,11 @@ function plot_conditional_forecast!(𝓂::ℳ, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm = sylvester_algorithm, tol = tol, - verbose = verbose) + verbose = verbose, + caching = caching, + use_workspaces = use_workspaces) + 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)) @@ -5314,7 +5241,9 @@ function plot_conditional_forecast!(𝓂::ℳ, tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, - sylvester_algorithm = sylvester_algorithm) + sylvester_algorithm = sylvester_algorithm, + caching = caching, + use_workspaces = use_workspaces) relevant_SS = relevant_SS isa KeyedArray ? axiskeys(relevant_SS,1) isa Vector{String} ? rekey(relevant_SS, 1 => axiskeys(relevant_SS,1) .|> Meta.parse .|> replace_indices) : relevant_SS : relevant_SS @@ -5401,13 +5330,7 @@ function plot_conditional_forecast!(𝓂::ℳ, full_shock_sort_perm = sortperm(full_shock_names_display, by = normalize_superscript) # Process rename dictionary to only include relevant keys in sorted order - relevant_keys = [k for k in keys(rename_dictionary) if (k isa String ? replace_indices(k) : k) in vcat(𝓂.constants.post_model_macro.var, 𝓂.constants.post_model_macro.exo)] |> sort - - processed_rename_dictionary = Any[] - - for k in relevant_keys - push!(processed_rename_dictionary, k => rename_dictionary[k]) - end + processed_rename_dictionary = process_rename_dictionary(rename_dictionary, 𝓂) # Combine sorted indices combined_sort_perm = vcat(var_sort_perm, (length(variable_names_display) .+ (1:length(shock_names_display)))[shock_sort_perm]) @@ -5426,13 +5349,7 @@ function plot_conditional_forecast!(𝓂::ℳ, # sorted_variable_names_display = sort(variable_names_display) sorted_shock_names_display = sort(shock_names_display) - orig_pal = StatsPlots.palette(attributes_redux[:palette]) - - total_pal_len = 100 - - alpha_reduction_factor = 0.7 - - pal = mapreduce(x -> StatsPlots.coloralpha.(orig_pal, alpha_reduction_factor ^ x), vcat, 0:(total_pal_len ÷ length(orig_pal)) - 1) |> StatsPlots.palette + pal = build_extended_palette(attributes_redux) args_and_kwargs = Dict(:run_id => length(conditional_forecast_active_plot_container) + 1, :model_name => 𝓂.model_name, @@ -5448,16 +5365,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, @@ -5469,55 +5377,10 @@ function plot_conditional_forecast!(𝓂::ℳ, :rename_dictionary => processed_rename_dictionary ) - no_duplicate = all( - !(all(( - get(dict, :parameters, nothing) == args_and_kwargs[:parameters], - get(dict, :rename_dictionary, nothing) == args_and_kwargs[:rename_dictionary], - get(dict, :conditions, nothing) == args_and_kwargs[:conditions], - get(dict, :shocks, nothing) == args_and_kwargs[:shocks], - get(dict, :initial_state, nothing) == args_and_kwargs[:initial_state], - all(get(dict, k, nothing) == get(args_and_kwargs, k, nothing) for k in setdiff(keys(DEFAULT_ARGS_AND_KWARGS_NAMES),[:label])) - ))) - for dict in conditional_forecast_active_plot_container - ) # "New plot must be different from previous plot. Use the version without ! to plot." - - if no_duplicate - push!(conditional_forecast_active_plot_container, args_and_kwargs) - else - @info "Plot with same parameters already exists. Using previous plot data to create plot." - end - - # 1. Keep only certain keys from each dictionary - reduced_vector = [ - Dict(k => d[k] for k in vcat(:run_id, :label, keys(DEFAULT_ARGS_AND_KWARGS_NAMES)...) if haskey(d, k)) - for d in conditional_forecast_active_plot_container - ] - - diffdict = compare_args_and_kwargs(reduced_vector) - - # 2. Group the original vector by :model_name - grouped_by_model = Dict{Any, Vector{Dict}}() - - for d in conditional_forecast_active_plot_container - model = d[:model_name] - d_sub = Dict(k => d[k] for k in setdiff(keys(args_and_kwargs), keys(DEFAULT_ARGS_AND_KWARGS_NAMES)) if haskey(d, k)) - push!(get!(grouped_by_model, model, Vector{Dict}()), d_sub) - end - - model_names = [] - - for d in conditional_forecast_active_plot_container - push!(model_names, d[:model_name]) - end + push_if_no_duplicate!(conditional_forecast_active_plot_container, args_and_kwargs, + [:parameters, :rename_dictionary, :conditions, :shocks, :initial_state, :tol]) - model_names = unique(model_names) - - for model in model_names - if length(grouped_by_model[model]) > 1 - diffdict_grouped = compare_args_and_kwargs(grouped_by_model[model]) - diffdict = merge_by_runid(diffdict, diffdict_grouped) - end - end + diffdict = compute_diffdict(conditional_forecast_active_plot_container, keys(args_and_kwargs)) annotate_ss = Vector{Pair{String, Any}}[] @@ -5529,13 +5392,7 @@ function plot_conditional_forecast!(𝓂::ℳ, len_diff = length(conditional_forecast_active_plot_container) - if haskey(diffdict, :parameters) - param_nms = diffdict[:parameters] |> keys |> collect |> sort - for param in param_nms - result = [x === nothing ? "" : x for x in diffdict[:parameters][param]] - push!(annotate_diff_input, String(param) => result) - end - end + annotate_param_diff!(annotate_diff_input, diffdict) if haskey(diffdict, :shocks) shocks = diffdict[:shocks] @@ -5662,50 +5519,20 @@ function plot_conditional_forecast!(𝓂::ℳ, push!(annotate_diff_input, "Initial state" => labels) end - rename_idx = Int[] - - if haskey(diffdict, :rename_dictionary) - non_nothing_dicts = [d for d in diffdict[:rename_dictionary] if !isnothing(d) && length(d) > 0] - unique_dicts = unique(non_nothing_dicts) - - for init in diffdict[:rename_dictionary] - if isnothing(init) || length(init) == 0 - push!(rename_idx, 0) - continue - end - - for (i,u) in enumerate(unique_dicts) - if u == init - push!(rename_idx,i) - continue - end - end - end - - push!(annotate_diff_input, "Rename dictionary" => [i > 0 ? "#$i" : "nothing" for i in rename_idx]) - end + annotate_rename_dict_diff!(annotate_diff_input, diffdict) same_shock_direction = true - for k in setdiff(keys(args_and_kwargs), - [ - :run_id, :parameters, :plot_data, :tol, :reference_steady_state, :initial_state, :conditions, :conditions_in_levels, :label, - :shocks, :shock_names, - :variables, :variable_names, :var_idx, - :rename_dictionary, - # :periods, :quadratic_matrix_equation_algorithm, :sylvester_algorithm, :lyapunov_algorithm, - ] - ) - - if haskey(diffdict, k) - push!(annotate_diff_input, DEFAULT_ARGS_AND_KWARGS_NAMES[k] => reduce(vcat,diffdict[k])) - - if k == :negative_shock - same_shock_direction = false - end - end + if annotate_default_kwarg_diffs!(annotate_diff_input, args_and_kwargs, diffdict, + [:run_id, :parameters, :plot_data, :tol, :reference_steady_state, :initial_state, :conditions, :conditions_in_levels, :label, + :shocks, :shock_names, + :variables, :variable_names, :var_idx, + :rename_dictionary]) + same_shock_direction = false end + annotate_tol_diff!(annotate_diff_input, conditional_forecast_active_plot_container) + 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])) @@ -5722,7 +5549,7 @@ function plot_conditional_forecast!(𝓂::ℳ, single_shock_per_irf = true max_periods = 0 - plt_lab_switch = ((length(annotate_diff_input) > 2) || (Dict(annotate_diff_input)["Plot label"] != collect(1:length(conditional_forecast_active_plot_container)))) && length(conditional_forecast_active_plot_container) > 1 + plt_lab_switch = should_use_label_switch(annotate_diff_input, conditional_forecast_active_plot_container) for (i,k) in enumerate(conditional_forecast_active_plot_container) if plot_type == :stack StatsPlots.bar!(legend_plot, @@ -5884,138 +5711,39 @@ function plot_conditional_forecast!(𝓂::ℳ, else plot_count = 1 - shock_string = "Conditional forecast" - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = 𝓂.model_name - model_string_filename = 𝓂.model_name - end - - plot_title = "Model: "*model_string*" " * shock_string *" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" - - ppp = StatsPlots.plot(pp...; attributes...) - - plot_elements = [ppp, legend_plot] - - layout_heights = [15, length(annotate_diff_input)] - - if plt_lab_switch - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_input_diff) - - push!(layout_heights, 5) - - pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) - else - pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) - end - - push!(annotate_ss, annotate_ss_page) - - if length(annotate_ss[pane]) > 1 - annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady States") - - ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_ss) - - push!(layout_heights, 5) - end - - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux...) - - push!(return_plots,p) - - if show_plots# & (length(pp) > 0) - display(p) - end - - if save_plots# & (length(pp) > 0) - if !isdir(save_plots_path) mkpath(save_plots_path) end - - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * string(pane) * "." * string(save_plots_format)) - end - - pane += 1 - - annotate_ss_page = Pair{String,Any}[] - - pp = [] + pane = assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + 𝓂.model_name; + title_extra = " Conditional forecast", + annotate_ss = annotate_ss, + annotate_ss_page = annotate_ss_page, + plt_lab_switch = plt_lab_switch, + ) end end if length(pp) > 0 - shock_string = "Conditional forecast" - - if haskey(diffdict, :model_name) - model_string = "multiple models" - model_string_filename = "multiple_models" - else - model_string = 𝓂.model_name - model_string_filename = 𝓂.model_name - end - - plot_title = "Model: "*model_string*" " * shock_string *" ("*string(pane)*"/"*string(Int(ceil(n_subplots/plots_per_page)))*")" - - ppp = StatsPlots.plot(pp...; attributes...) - - plot_elements = [ppp, legend_plot] - - layout_heights = [15, length(annotate_diff_input)] - - if plt_lab_switch - annotate_diff_input_plot = plot_df(annotate_diff_input; fontsize = attributes[:annotationfontsize], title = "Relevant Input Differences") - - ppp_input_diff = StatsPlots.plot(annotate_diff_input_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_input_diff) - - push!(layout_heights, 5) - - pushfirst!(annotate_ss_page, "Plot label" => reduce(vcat, diffdict[:label])) - else - pushfirst!(annotate_ss_page, annotate_diff_input[2][1] => annotate_diff_input[2][2]) - end - - push!(annotate_ss, annotate_ss_page) - - if length(annotate_ss[pane]) > 1 - annotate_ss_plot = plot_df(annotate_ss[pane]; fontsize = attributes[:annotationfontsize], title = "Relevant Steady States") - - ppp_ss = StatsPlots.plot(annotate_ss_plot; attributes..., framestyle = :box) - - push!(plot_elements, ppp_ss) - - push!(layout_heights, 5) - end - - p = StatsPlots.plot(plot_elements..., - layout = StatsPlots.grid(length(layout_heights), 1, heights = layout_heights ./ sum(layout_heights)), - plot_title = plot_title; - attributes_redux...) - - push!(return_plots,p) - - if show_plots# & (length(pp) > 0) - display(p) - end - - if save_plots# & (length(pp) > 0) - if !isdir(save_plots_path) mkpath(save_plots_path) end - - StatsPlots.savefig(p, save_plots_path * "/" * string(save_plots_name) * "__" * model_string_filename * "__" * string(pane) * "." * string(save_plots_format)) - end + assemble_and_emit_page!( + return_plots, pp, legend_plot, + annotate_diff_input, diffdict, + attributes, attributes_redux, + pane, n_subplots, plots_per_page, + show_plots, save_plots, save_plots_path, save_plots_name, save_plots_format, + 𝓂.model_name; + title_extra = " Conditional forecast", + annotate_ss = annotate_ss, + annotate_ss_page = annotate_ss_page, + plt_lab_switch = plt_lab_switch, + is_tail = true, + ) end + if !use_workspaces 𝓂.workspaces = orig_ws end + return return_plots end 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/models/FRBUS.jl b/models/FRBUS.jl new file mode 100644 index 000000000..2de1c718d --- /dev/null +++ b/models/FRBUS.jl @@ -0,0 +1,3061 @@ +# FRB/US LINVER (2024) model +# Source: https://www.federalreserve.gov/econres/us-models-linver.htm +# Settings: +# expvers = "mcap" +# mprule = "intay" +# elb_imposed = "no" + +@model FRBUS begin + delrff[0] = rff[0] - rff[-1] + + dpadj[0] = dpadj[-1] + dpgap[-1] + + dpgap[0] = y_dpgap_1 * pipxnc[0] + y_dpgap_2 * phr_l[0] - pxp_l[0] + y_dpgap_3 * phr_l[-1] + pxp_l[-1] + y_dpgap_4 * pbfir_l[0] + y_dpgap_5 * pbfir_l[-1] + y_dpgap_6 * pegfr_l[0] + y_dpgap_7 * pegfr_l[-1] + y_dpgap_8 * pegsr_l[0] + y_dpgap_9 * pegsr_l[-1] + y_dpgap_10 * pxr_l[0] + y_dpgap_11 * pxr_l[-1] + + ebfi_l[0] = y_ebfi_l_8 * hgpbfir[-1] + y_ebfi_l_6 * xb_l[-1] + y_ebfi_l_5 * zebfi[0] + y_ebfi_l_1 * ebfi_l[-1] + ebfi_l_aerr[x] + y_ebfi_l_2 * qebfi_l[-1] + y_ebfi_l_3 * ebfi_l[-2] + y_ebfi_l_4 * ebfi_l[-3] + y_ebfi_l_7 * xb_l[-2] + + ebfin_l[0] = pxp_l[0] + pbfir_l[0] + ebfi_l[0] + + # ec_l[0] = ec_l[-1] + y_ec_l_1 * eco_l[0] + y_ec_l_2 * eco_l[-1] + y_ec_l_3 * ech_l[0] + y_ec_l_4 * ech_l[-1] + y_ec_l_5 * yhpcd_l[0] + y_ec_l_6 * jkcd_l[0] + y_ec_l_7 * yhpcd_l[-1] + y_ec_l_8 * jkcd_l[-1] + + ecd_l[0] = y_ecd_l_4 * zgapc2[0] + zecd[0] + y_ecd_l_1 * ecd_l[-1] + ecd_l_aerr[x] + y_ecd_l_2 * qecd_l[-1] + y_ecd_l_3 * ecd_l[-2] + + ech_l[0] = y_ech_l_3 * ech_l_aerr[x] + kh_l[-1] + ech_l[-1] * y_ech_l_1 + y_ech_l_2 * kh_l[-2] + y_ech_l_4 * ech_l[-2] + y_ech_l_5 * kh_l[-3] + + ecnia_l[0] = ecnia_l[-1] + eco_l[0] * y_ecnia_l_1 + eco_l[-1] * y_ecnia_l_2 + ecd_l[0] * y_ecnia_l_3 + ecd_l[-1] * y_ecnia_l_4 + ech_l[0] * y_ecnia_l_5 + ech_l[-1] * y_ecnia_l_6 + + ecnian_l[0] = ecnia_l[0] + pcnia_l[0] + + eco_l[0] = y_eco_l_8 * yht_l[-1] + y_eco_l_7 * yhl_l[-1] + y_eco_l_6 * yht_l[0] + y_eco_l_5 * yhl_l[0] + y_eco_l_4 * zeco[0] + eco_l[-1] * y_eco_l_1 + eco_l_aerr[x] + y_eco_l_2 * qeco_l[-1] + y_eco_l_3 * eco_l[-2] + + egfe_l[0] = fiscal_egfe * fiscal[0] + y_egfe_l_7 * xgap2[-1] + y_egfe_l_6 * xgap2[0] + y_egfe_l_5 * egfet_l[0] + y_egfe_l_1 * egfe_l[-1] + egfe_l_aerr[x] + y_egfe_l_2 * egfet_l[-1] + y_egfe_l_3 * egfe_l[-2] + y_egfe_l_4 * egfe_l[-3] + + egfen_l[0] = egfe_l[0] + pxp_l[0] + pegfr_l[0] + + egfet_l[0] = egfet_l[-1] * y_egfet_l_1 + pegfr_l[-1] * y_egfet_l_2 + pxp_l[-1] * y_egfet_l_3 + y_egfet_l_4 * xgdptn_l[-1] + y_egfet_l_5 * hggdpt[0] + y_egfet_l_6 * hggdpt[-1] + y_egfet_l_7 * hggdpt[-2] + y_egfet_l_8 * hggdpt[-3] + + egfl_l[0] = fiscal_egfl * fiscal[0] + xgap2[-1] * y_egfl_l_7 + xgap2[0] * y_egfl_l_6 + y_egfl_l_5 * egflt_l[0] + y_egfl_l_1 * egfl_l[-1] + egfl_l_aerr[x] + y_egfl_l_2 * egflt_l[-1] + y_egfl_l_3 * egfl_l[-2] + y_egfl_l_4 * egfl_l[-3] + + egfln_l[0] = egfl_l[0] + pgfl_l[0] + + egflt_l[0] = egflt_l[-1] * y_egflt_l_1 + y_egflt_l_2 * pgfl_l[-1] + xgdptn_l[-1] * y_egflt_l_3 + hggdpt[0] * y_egflt_l_4 + hggdpt[-1] * y_egflt_l_5 + y_egflt_l_6 * hggdpt[-2] + y_egflt_l_7 * hggdpt[-3] + + egse_l[0] = xgap2[-1] * y_egse_l_7 + xgap2[0] * y_egse_l_6 + y_egse_l_5 * egset_l[0] + y_egse_l_1 * egse_l[-1] + egse_l_aerr[x] + y_egse_l_2 * egset_l[-1] + y_egse_l_3 * egse_l[-2] + y_egse_l_4 * egse_l[-3] + + egsen_l[0] = egse_l[0] + pxp_l[0] + pegsr_l[0] + + egset_l[0] = egset_l[-1] * y_egset_l_1 + pegsr_l[-1] * y_egset_l_2 + pxp_l[-1] * y_egset_l_3 + xgdptn_l[-1] * y_egset_l_4 + hggdpt[0] * y_egset_l_5 + hggdpt[-1] * y_egset_l_6 + y_egset_l_7 * hggdpt[-2] + y_egset_l_8 * hggdpt[-3] + + egsl_l[0] = xgap2[-1] * y_egsl_l_7 + xgap2[0] * y_egsl_l_6 + y_egsl_l_5 * egslt_l[0] + y_egsl_l_1 * egsl_l[-1] + egsl_l_aerr[x] + y_egsl_l_2 * egslt_l[-1] + y_egsl_l_3 * egsl_l[-2] + y_egsl_l_4 * egsl_l[-3] + + egsln_l[0] = egsl_l[0] + pgsl_l[0] + + egslt_l[0] = egslt_l[-1] * y_egslt_l_1 + y_egslt_l_2 * pgsl_l[-1] + xgdptn_l[-1] * y_egslt_l_3 + hggdpt[0] * y_egslt_l_4 + hggdpt[-1] * y_egslt_l_5 + y_egslt_l_6 * hggdpt[-2] + y_egslt_l_7 * hggdpt[-3] + + eh_l[0] = y_eh_l_7 * d83[x] + y_eh_l_5 * rme[-1] + zeh[0] + y_eh_l_1 * eh_l[-1] + eh_l_aerr[x] + y_eh_l_2 * qeh_l[-1] + y_eh_l_3 * eh_l[-2] + y_eh_l_4 * eh_l[-3] + y_eh_l_6 * rme[-2] + + ehn_l[0] = eh_l[0] + phr_l[0] + pxp_l[0] + + # em_l[0] = em_l[-1] + y_em_l_1 * emon_l[0] + y_em_l_2 * emn_l[0] + y_em_l_3 * emon_l[-1] + y_em_l_4 * emn_l[-1] + y_em_l_5 * emo_l[0] + y_em_l_6 * emo_l[-1] + y_em_l_7 * empn_l[0] + y_em_l_8 * empn_l[-1] + y_em_l_9 * emp_l[0] + y_em_l_10 * emp_l[-1] + + emn_l[0] = emon_l[0] * y_emn_l_2 + empn_l[0] * y_emn_l_3 + + emo_l[0] = y_emo_l_9 * ddockm[x-1] + y_emo_l_8 * ddockm[x] + y_emo_l_7 * xgap2[-2] + xgap2[-1] * y_emo_l_6 + xgap2[0] * y_emo_l_5 + y_emo_l_4 * xgdpn_l[-1] + emo_l[-1] * y_emo_l_1 + emo_l̃[0] + y_emo_l_2 * pmo_l[-1] + y_emo_l_3 * uemot[x-1] + + emo_l̃[0] = (1 - rho_emo_l) * emo_l̄ + rho_emo_l * emo_l̃[-1] + emo_l_aerr[x] + + emon_l[0] = emo_l[0] + pmo_l[0] + + emp_l[0] = xgdp_l[0] + emp_l_aerr[x] + y_emp_l_1 * emptrt[x] + y_emp_l_2 * pmp_l[0] + y_emp_l_3 * pxb_l[0] + y_emp_l_4 * pmp_l[-1] + y_emp_l_5 * pxb_l[-1] + xgap2[-1] * y_emp_l_6 + + empn_l[0] = emp_l[0] + pmp_l[0] + + ex_l[0] = y_ex_l_10 * ddockx[x] + y_ex_l_1 * ex_l[-1] + ex_l_aerr[x] + pxr_l[-1] * y_ex_l_2 + pxp_l[-1] * y_ex_l_3 + y_ex_l_4 * fpx_l[-1] + y_ex_l_5 * fgdp_l[-1] + y_ex_l_6 * fpc_l[-1] + y_ex_l_7 * fxgap[0] + y_ex_l_8 * fxgap[-1] + y_ex_l_9 * fxgap[-2] + + exn_l[0] = ex_l[0] + pxp_l[0] + pxr_l[0] + + fcbn_l[0] = exn_l[0] * y_fcbn_l_2 + emn_l[0] * y_fcbn_l_3 + y_fcbn_l_4 * fynicn_l[0] + y_fcbn_l_5 * fyniln_l[0] + y_fcbn_l_6 * ufcbr[x] + pxb_l[0] * y_fcbn_l_7 + y_fcbn_l_8 * xbt_l[0] + + fgdp_l[0] = fgdpt_l[0] + fxgap[0] * y_fgdp_l_2 + + fgdpt_l[0] = y_fgdpt_l_1 * fgdpt_l[-1] + y_fgdpt_l_2 * xgdpt_l[-1] + hggdpt[0] * y_fgdpt_l_3 + hggdpt[-1] * y_fgdpt_l_4 + y_fgdpt_l_5 * hggdpt[-2] + y_fgdpt_l_6 * hggdpt[-3] + + fnicn_l[0] = y_fnicn_l_1 * fnicn_l[-1] + y_fnicn_l_2 * xgdptn_l[0] + y_fnicn_l_4 * fpc_l[0] + fpc_l[-1] * y_fnicn_l_5 + y_fnicn_l_6 * fpx_l[0] + fpx_l[-1] * y_fnicn_l_7 + y_fnicn_l_8 * rfnict[x] + + fniln_l[0] = y_fniln_l_1 * fniln_l[-1] + rfnict[x] * y_fniln_l_3 + xgdptn_l[0] * y_fniln_l_4 + fcbn_l[0] * y_fniln_l_5 + y_fniln_l_6 * pgdp_l[0] + y_fniln_l_7 * pgdp_l[-1] + fpx_l[0] * y_fniln_l_8 + fpx_l[-1] * y_fniln_l_9 + y_fniln_l_10 * fnirn_l[0] + + fnirn_l[0] = y_fnirn_l_2 * ufnir[x] + xgdpn_l[0] + + fpc_l[0] = fpc_l[-1] + y_fpc_l_2 * fpic[0] + + fpi10[0] = fxgap[-1] * y_fpi10_6 + y_fpi10_5 * fpitrg[x] + y_fpi10_1 * fpi10[-1] + y_fpi10_2 * fpi10[-2] + y_fpi10_3 * fpi10[-3] + y_fpi10_4 * fpi10[-4] + + fpi10t[0] = y_fpi10t_1 * fpi10t[-1] + fpi10[0] * y_fpi10t_2 + + fpic[0] = fpi10[0] * y_fpic_1 + y_fpic_2 * fpic[-1] + + fpx_l[0] = fpc_l[0] + fpxr_l[0] - pcpi_l[0] + + fpxr_l[0] = fpxrr_l[0] + y_fpxr_l_1 * rg10[0] + y_fpxr_l_2 * zpi10f[0] + y_fpxr_l_3 * frl10[0] + fpi10t[0] * y_fpxr_l_4 + fnicn_l[0] * y_fpxr_l_5 + fniln_l[0] * y_fpxr_l_6 + xgdpn_l[0] * y_fpxr_l_7 + + fpxrr_l[0] = y_fpxrr_l_4 * fpxrrt[x] + y_fpxrr_l_3 * fpxrr_l[-2] + y_fpxrr_l_1 * fpxrr_l[-1] + fpxrr_l̃[0] + y_fpxrr_l_2 * fpxrrt[x-1] + + fpxrr_l̃[0] = (1 - rho_fpxrr_l) * fpxrr_l̄ + rho_fpxrr_l * fpxrr_l̃[-1] + fpxrr_l_aerr[x] + + frl10[0] = fxgap[-1] * y_frl10_6 + fxgap[0] * y_frl10_5 + y_frl10_4 * frs10[0] + y_frl10_1 * frl10[-1] + y_frl10_2 * frs10[-1] + y_frl10_3 * frl10[-2] + + frs10[0] = rfrs10[x] + fxgap[0] * y_frs10_8 + fpitrg[x] * y_frs10_7 + y_frs10_1 * dfmprr[x] + y_frs10_2 * frstar[-1] + fpi10[0] * y_frs10_3 + fpi10[-1] * y_frs10_4 + y_frs10_5 * fpi10[-2] + y_frs10_6 * fpi10[-3] + + frstar[0] = frstar[-1] * y_frstar_1 + frs10[0] * y_frstar_2 + fpi10[0] * y_frstar_3 + fpi10[-1] * y_frstar_4 + y_frstar_5 * fpi10[-2] + y_frstar_6 * fpi10[-3] + + ftcin_l[0] = y_ftcin_l_2 * uftcin[x] + ynicpn_l[0] + + fxgap[0] = xgap2[-1] * y_fxgap_13 + frstar[0] * y_fxgap_12 + fpi10[-1] * y_fxgap_4 + frs10[-1] * y_fxgap_3 + fxgap_aerr[x] + fxgap[-1] * y_fxgap_1 + y_fxgap_2 * fxgap[-2] + y_fxgap_5 * fpi10[-2] + y_fxgap_6 * fpi10[-3] + y_fxgap_7 * fpi10[-4] + y_fxgap_8 * frs10[-2] + y_fxgap_9 * fpi10[-5] + y_fxgap_10 * frs10[-3] + y_fxgap_11 * fpi10[-6] + + fynicn_l[0] = fnicn_l[-1] + y_fynicn_l_2 * rfynic[0] + + fyniln_l[0] = fniln_l[-1] + y_fyniln_l_2 * rfynil[0] + + gfdbtnp_l[0] = ugfdbtp_l[0] + y_gfdbtnp_l_2 * gfdbtnp_l[-1] + y_gfdbtnp_l_3 * gfexpn_l[0] + y_gfdbtnp_l_4 * gfrecn_l[0] + + gfdbtn_l[0] = gfdbtnp_l[0] + ugfdbt_l[x] + + ugfdbtp_l[0] = (1 - rho_ugfdbtp_l) * ugfdbtp_l̄ + rho_ugfdbtp_l * ugfdbtp_l[-1] + ugfdbtp_lerr[x] + + ugfsrp[0] = y_ugfsrp_1 * ugfsrp[-1] + + uleg_l[0] = uleg_l[-1] + y_uleg_l_1 * leg_l[-1] + y_uleg_l_2 * lep_l[-1] + y_uleg_l_3 * adjlegrt[x] + + gfexpn_l[0] = egfln_l[0] * y_gfexpn_l_2 + egfen_l[0] * y_gfexpn_l_3 + y_gfexpn_l_4 * gtn_l[0] + y_gfexpn_l_5 * gfintn_l[0] + + gfintn_l[0] = y_gfintn_l_2 * rgfint[0] + gfdbtn_l[-1] + + gfrecn_l[0] = y_gfrecn_l_2 * tpn_l[0] + y_gfrecn_l_3 * tcin_l[0] + ugfsrp[0] * y_gfrecn_l_4 + xgdpn_l[0] * y_gfrecn_l_5 + + gtn_l[0] = pgdp_l[0] + gtr_l[0] + + gtr_l[0] = y_gtr_l_2 * gtrd[0] + y_gtr_l_3 * gtrt[x] + xgdpt_l[0] + + gtrd[0] = .0014 * (fiscalav[0] - y_gtrd_6 * fiscalav[-1]) + y_gtrd_6 * gtrd[-1] + gtrd_aerr[x] + xgap2[0] * y_gtrd_1 + xgap2[-1] * y_gtrd_2 + y_gtrd_3 * xgap2[-2] + y_gtrd_4 * xgap2[-3] + y_gtrd_5 * xgap2[-4] + y_gtrd_7 * xgap2[-5] + + hgemp[0] = y_hgemp_1 * hgemp[-1] + emp_l[0] * y_hgemp_2 + emp_l[-1] * y_hgemp_3 + + hggdp[0] = xgdp_l[0] * y_hggdp_1 + y_hggdp_2 * xgdp_l[-1] + + hggdpt[0] = hxbt[0] + huxb[0] + + hgpbfir[0] = hgpbfir[-1] * y_hgpbfir_1 + pbfir_l[0] * y_hgpbfir_2 + pxp_l[0] * y_hgpbfir_3 + pxb_l[0] * y_hgpbfir_4 + pbfir_l[-1] * y_hgpbfir_5 + pxp_l[-1] * y_hgpbfir_6 + pxb_l[-1] * y_hgpbfir_7 + + hgpkir[0] = y_hgpkir_1 * hgpkir[-1] + y_hgpkir_2 * pkir[x] + y_hgpkir_3 * pkir[x-1] + + hgynid[0] = ynicpn_l[0] * y_hgynid_1 + tcin_l[0] * y_hgynid_2 + pxb_l[0] * y_hgynid_3 + y_hgynid_4 * ynicpn_l[-1] + y_hgynid_5 * tcin_l[-1] + pxb_l[-1] * y_hgynid_6 + + hks[0] = y_hks_1 * kbfi_l[0] + y_hks_2 * kbfi_l[-1] + y_hks_3 * ki_l[0] + y_hks_4 * ki_l[-1] + hksr[x] + + hlept[0] = y_hlept_1 * hqlfpr[0] + y_hlept_2 * n16_l[x] + y_hlept_3 * n16_l[x-1] + + hlprdt[0] = hxbt[0] - hlept[0] - hqlww[0] + + hmfpt[0] = hmfpt_aerr[x] + y_hmfpt_1 * hmfpt[-1] + + hqlfpr[0] = hqlfpr_aerr[x] + y_hqlfpr_1 * hqlfpr[-1] + + hqlww[0] = hqlww_aerr[x] + y_hqlww_1 * hqlww[-1] + + huqpct[0] = y_huqpct_1 * huqpct[-1] + + huxb[0] = y_huxb_1 * dglprd[x] + y_huxb_2 * huxb[-1] + + hxbt[0] = hmfpt[0] + hks[0] * y_hxbt_5 + hlept[0] * y_hxbt_1 + hqlww[0] * y_hxbt_2 + y_hxbt_3 * lqualt_l[x] + y_hxbt_4 * lqualt_l[x-1] + + # hxbtr[0] = 0 + + jccan_l[0] = xgdpn_l[0] + y_jccan_l_2 * jccan_l[-1] + xgdpn_l[-1] * y_jccan_l_3 + y_jccan_l_4 * pkbfir[-1] + kbfi_l[-1] * y_jccan_l_5 + y_jccan_l_6 * jrbfi[x] + pxp_l[-1] * y_jccan_l_7 + + jkcd_l[0] = y_jkcd_l_2 * jrcd[x] + kcd_l[-1] + + kbfi_l[0] = pbfir_l[0] * y_kbfi_l_2 + y_kbfi_l_3 * pkbfir[0] + ebfi_l[0] * y_kbfi_l_4 + jrbfi[x] * y_kbfi_l_5 + kbfi_l[-1] * y_kbfi_l_6 + + kcd_l[0] = ecd_l[0] * y_kcd_l_2 + jrcd[x] * y_kcd_l_3 + kcd_l[-1] * y_kcd_l_4 + + kh_l[0] = eh_l[0] * y_kh_l_2 + y_kh_l_3 * jrh[x] + kh_l[-1] * y_kh_l_4 + + ki_l[0] = ki_l[-1] * y_ki_l_1 + ki_l_aerr[x] + y_ki_l_2 * qkir_l[0] + y_ki_l_3 * xfs_l[-1] + y_ki_l_4 * ki_l[-2] + y_ki_l_5 * xfs_l[-2] + y_ki_l_6 * xfs_l[-3] + + ks_l[0] = ks_l[-1] + hks[0] * y_ks_l_1 + + leg_l[0] = uleg_l[0] + egfl_l[0] * y_leg_l_1 + egsl_l[0] * y_leg_l_2 - lprdt_l[0] + + leh_l[0] = y_leh_l_2 * lep_l[0] + leg_l[0] * y_leh_l_3 + y_leh_l_4 * leo_l[0] + + leo_l[0] = xgap2[-1] * y_leo_l_5 + y_leo_l_4 * qlf_l[-1] + leo_l_aerr[x] + y_leo_l_1 * qleor[x] + qlf_l[0] + y_leo_l_2 * leo_l[-1] + y_leo_l_3 * qleor[x-1] + + lep_l[0] = lhp_l[0] - lww_l[0] + + leppot_l[0] = qlf_l[0] + y_leppot_l_2 * lurnat[0] + qleor[x] * y_leppot_l_3 + adjlegrt[x] * y_leppot_l_4 + + lf_l[0] = n16_l[x] + y_lf_l_2 * lfpr[0] + + lfpr[0] = hqlfpr[0] + y_lfpr_1 * lfpr[-1] + lfpr_aerr[x] + y_lfpr_2 * qlfpr[-1] + y_lfpr_3 * lur[-1] + y_lfpr_4 * lurnat[-1] + + lhp_l[0] = y_lhp_l_7 * hlprdt[-1] + y_lhp_l_6 * xbo_l[-1] + y_lhp_l_5 * xbo_l[0] + y_lhp_l_4 * zlhp[0] + y_lhp_l_1 * lhp_l[-1] + lhp_l_aerr[x] + y_lhp_l_2 * qlhp_l[-1] + y_lhp_l_3 * lhp_l[-2] + y_lhp_l_8 * xbo_l[-2] + y_lhp_l_9 * hlprdt[-2] + + lprdt_l[0] = xbt_l[0] - leppot_l[0] - qlww_l[0] + + lur[0] = leh_l[0] * y_lur_1 + lf_l[0] * y_lur_2 + + lurnat[0] = lurnat_aerr[x] + lurnat[-1] * y_lurnat_1 + + lww_l[0] = y_lww_l_1 * lww_l[-1] + hqlww[0] * y_lww_l_2 + lww_l_aerr[x] + y_lww_l_3 * qlww_l[-1] + lhp_l[0] * y_lww_l_4 + lhp_l[-1] * y_lww_l_5 + hlept[0] * y_lww_l_6 + + # mei_l[0] = mei_l[-1] + + # mep_l[0] = mep_l[-1] + + mfpt_l[0] = mfpt_l_aerr[x] + mfpt_l[-1] + hmfpt[0] * y_mfpt_l_1 + + pbfir_l[0] = pxp_l[-1] + dpadj[0] + pbfir_l[-1] + pbfir_l_aerr[x] + pipxnc[0] * y_pbfir_l_1 - pxp_l[0] + + pcdr_l[0] = y_pcdr_l_1 * pcdr_l[-1] + y_pcdr_l_2 * pcdr_l[-2] + + pcer_l[0] = pcer_l[-1] + pcer_l_aerr[x] + pmp_l[0] * y_pcer_l_1 + y_pcer_l_2 * pcxfe_l[0] + pmp_l[-1] * y_pcer_l_3 + y_pcer_l_4 * pcxfe_l[-1] + + pcfr_l[0] = y_pcfr_l_6 * pcfrt[x] + y_pcfr_l_5 * pcfr_l[-4] + y_pcfr_l_4 * pcfr_l[-3] + y_pcfr_l_3 * pcfr_l[-2] + y_pcfr_l_1 * pcfr_l[-1] + pcfr_l_aerr[x] + y_pcfr_l_2 * pcfrt[x-1] + + pchr_l[0] = y_pchr_l_1 * pchr_l[-1] + y_pchr_l_2 * pchr_l[-2] + + pcnia_l[0] = pcnia_l[-1] + y_pcnia_l_1 * picnia[0] + + pcor_l[0] = pcor_l[-1] + pcdr_l[0] * y_pcor_l_1 + pcdr_l[-1] * y_pcor_l_2 + pchr_l[0] * y_pcor_l_3 + pchr_l[-1] * y_pcor_l_4 + + pcpi_l[0] = pcnia_l[0] + y_pcpi_l_2 * upcpi[x] + + pcpix_l[0] = pcxfe_l[0] + y_pcpix_l_2 * upcpix[x] + + pcxfe_l[0] = pcxfe_l[-1] + y_pcxfe_l_1 * picxfe[0] + + pegfr_l[0] = pxp_l[-1] + dpadj[0] + pegfr_l[-1] + pegfr_l_aerr[x] + pipxnc[0] * y_pegfr_l_1 - pxp_l[0] + + pegsr_l[0] = pxp_l[-1] + dpadj[0] + pegsr_l[-1] + pegsr_l_aerr[x] + pipxnc[0] * y_pegsr_l_1 - pxp_l[0] + + pgdp_l[0] = xgdpn_l[0] - xgdp_l[0] + + pgfl_l[0] = y_pgfl_l_1 * upgfl[x] + pl_l[0] - lprdt_l[0] + + pgsl_l[0] = pl_l[0] + y_pgsl_l_1 * upgsl[x] - lprdt_l[0] + + phouse_l[0] = pcnia_l[-1] * y_phouse_l_4 + pchr_l[-1] * y_phouse_l_3 + y_phouse_l_1 * phouse_l[-1] + phouse_l_aerr[x] + y_phouse_l_2 * phouse_l[-2] + + phr_l[0] = pxp_l[-1] + dpadj[0] + phr_l[-1] + phr_l_aerr[x] + pipxnc[0] * y_phr_l_1 - pxp_l[0] + + pic4[0] = pcnia_l[0] * y_pic4_1 + y_pic4_2 * pcnia_l[-4] + + picnia[0] = picxfe[0] + pcer_l[0] * y_picnia_1 + pcer_l[-1] * y_picnia_2 + pcfr_l[0] * y_picnia_3 + pcfr_l[-1] * y_picnia_4 + + picx4[0] = pcxfe_l[0] * y_picx4_1 + y_picx4_2 * pcxfe_l[-4] + + picxfe[0] = picxfe_aerr[x] + y_picxfe_1 * picxfe[-1] + y_picxfe_2 * zpicxfe[0] + y_picxfe_3 * ptr[-1] + y_picxfe_4 * qpcnia_l[-1] + pcnia_l[-1] * y_picxfe_5 + + pieci[0] = y_pieci_12 * pl_l[-1] + y_pieci_11 * qpl_l[-1] + lurnat[-1] * y_pieci_10 + lur[-1] * y_pieci_9 + huqpct[-1] * y_pieci_8 + hlprdt[-1] * y_pieci_7 + ptr[-1] * y_pieci_6 + y_pieci_5 * zpieci[0] + pieci_aerr[x] + y_pieci_1 * pieci[-1] + y_pieci_2 * pieci[-2] + y_pieci_3 * pieci[-3] + y_pieci_4 * pieci[-4] + + pigdp[0] = pgdp_l[0] * y_pigdp_1 + pgdp_l[-1] * y_pigdp_2 + + pipl[0] = pieci[0] + + pipxnc[0] = y_pipxnc_11 * pxnc_l[-1] + y_pipxnc_10 * qpxnc_l[-1] + y_pipxnc_9 * fpxr_l[-1] + fpxr_l[0] * y_pipxnc_8 + picnia[0] + huqpct[0] * y_pipxnc_1 + y_pipxnc_2 * pipxnc[-1] + y_pipxnc_3 * picnia[-1] + huqpct[-1] * y_pipxnc_4 + y_pipxnc_5 * pipxnc[-2] + y_pipxnc_6 * picnia[-2] + y_pipxnc_7 * huqpct[-2] + + pkbfir[0] = y_pkbfir_1 * upkbfir[x] + pbfir_l[0] * y_pkbfir_2 + + pl_l[0] = pl_l[-1] + pipl[0] * y_pl_l_1 + + pmo_l[0] = pmo_l[-1] * y_pmo_l_1 + pmo_l̃[0] + y_pmo_l_2 * qpmo_l + fpc_l[-1] * y_pmo_l_3 + fpx_l[-1] * y_pmo_l_4 + pxb_l[-1] * y_pmo_l_5 + fpc_l[0] * y_pmo_l_6 + fpx_l[0] * y_pmo_l_7 + pxb_l[0] * y_pmo_l_8 + + pmo_l̃[0] = (1 - rho_pmo_l) * pmo_l̄ + rho_pmo_l * pmo_l̃[-1] + pmo_l_aerr[x] + + pmp_l[0] = y_pmp_l_2 * upmp[x] + poil_l[0] + + poil_l[0] = pxb_l[0] + poilr_l[0] + + poilr_l[0] = y_poilr_l_4 * poilrt[x] + y_poilr_l_3 * poilr_l[-2] + y_poilr_l_1 * poilr_l[-1] + poilr_l_aerr[x] + y_poilr_l_2 * poilrt[x-1] + + ptr[0] = ptr[-1] * y_ptr_1 + picxfe[-1] * y_ptr_2 + y_ptr_3 * pitarg[x-1] + + pxb_l[0] = pgdp_l[0] + y_pxb_l_2 * upxb[x] + + pxnc_l[0] = pxnc_l[-1] + pipxnc[0] * y_pxnc_l_1 + + pxp_l[0] = pxp_l[-1] + pcnia_l[0] * y_pxp_l_1 + pcnia_l[-1] * y_pxp_l_2 + pxnc_l[0] * y_pxp_l_3 + pxnc_l[-1] * y_pxp_l_4 + + pxr_l[0] = pxp_l[-1] + dpadj[0] + pxr_l[-1] + pxr_l_aerr[x] + pipxnc[0] * y_pxr_l_1 - pxp_l[0] + + qebfi_l[0] = xb_l[0] + y_qebfi_l_2 * vbfi[0] + hxbt[0] * y_qebfi_l_3 + hgpbfir[0] * y_qebfi_l_4 + jrbfi[x] * y_qebfi_l_5 + + qec_l[0] = y_qec_l_1 * zyh_l[0] + y_qec_l_2 * zyht_l[0] + y_qec_l_3 * zyhp_l[0] + y_qec_l_4 * wpo_l[0] + y_qec_l_5 * wps_l[0] + + qecd_l[0] = y_qecd_l_13 * rccd[0] + pcdr_l[0] * y_qecd_l_12 + y_qecd_l_11 * hgpcdr[x] + qec_l[0] + jrcd[x] * y_qecd_l_2 + hggdpt[0] * y_qecd_l_3 + hggdpt[-1] * y_qecd_l_4 + y_qecd_l_5 * hggdpt[-2] + y_qecd_l_6 * hggdpt[-3] + y_qecd_l_7 * hggdpt[-4] + y_qecd_l_8 * hggdpt[-5] + y_qecd_l_9 * hggdpt[-6] + y_qecd_l_10 * hggdpt[-7] + + qeco_l[0] = qec_l[0] - pcor_l[0] + + qeh_l[0] = y_qeh_l_19 * rcch[0] + pcnia_l[0] + qec_l[0] + jrh[x] * y_qeh_l_2 + hggdpt[0] * y_qeh_l_3 + hggdpt[-1] * y_qeh_l_4 + y_qeh_l_5 * hggdpt[-2] + y_qeh_l_6 * hggdpt[-3] + y_qeh_l_7 * hggdpt[-4] + y_qeh_l_8 * hggdpt[-5] + y_qeh_l_9 * hggdpt[-6] + y_qeh_l_10 * hggdpt[-7] + y_qeh_l_11 * hggdpt[-8] + y_qeh_l_12 * hggdpt[-9] + y_qeh_l_13 * hggdpt[-10] + y_qeh_l_14 * hggdpt[-11] + y_qeh_l_15 * hggdpt[-12] + y_qeh_l_16 * hggdpt[-13] + y_qeh_l_17 * hggdpt[-14] + y_qeh_l_18 * hggdpt[-15] - phr_l[0] - pxp_l[0] + + qkir_l[0] = dglprd[x] * y_qkir_l_1 + rho_qkir_l * qkir_l[-1] + + qlf_l[0] = n16_l[x] + y_qlf_l_2 * qlfpr[0] + + qlfpr[0] = hqlfpr[0] + qlfpr[-1] + + qlhp_l[0] = xbo_l[0] - lprdt_l[0] + + qlww_l[0] = qlww_l[-1] + hqlww[-1] * y_qlww_l_1 + + qpcnia_l[0] = qpxp_l[0] + uqpct_l[0] + + qpl_l[0] = pxb_l[0] + pl_l[0] - qpxb_l[0] + + # qpmo_l[0] = qpmo_l[-1] + + qpxb_l[0] = pl_l[0] + pwstar_l[x] - lprdt_l[0] + + qpxnc_l[0] = pxnc_l[0] + qpxp_l[0] * y_qpxnc_l_1 + pxp_l[0] * y_qpxnc_l_2 + qpcnia_l[0] * y_qpxnc_l_3 + pcnia_l[0] * y_qpxnc_l_4 + + qpxp_l[0] = pxp_l[0] + qpxb_l[0] * y_qpxp_l_1 + pxb_l[0] * y_qpxp_l_2 + + qynidn_l[0] = y_qynidn_l_1 * d79a[x] + ynicpn_l[0] * y_qynidn_l_2 + tcin_l[0] * y_qynidn_l_3 + + rbbb[0] = rg10[0] + rbbbp[0] + + rbbbp[0] = rbbbp_aerr[x] + y_rbbbp_1 * zgap10[0] + y_rbbbp_2 * rbbbp[-1] + y_rbbbp_3 * zgap10[-1] + + rbfi[0] = y_rbfi_1 * trfcim[x] + y_rbfi_2 * rg5[0] + rbbb[0] * y_rbfi_3 + rg10[0] * y_rbfi_4 + y_rbfi_5 * zpib5[0] + y_rbfi_6 * req[0] + + rcar[0] = rcar_aerr[x] + d79a[x] * y_rcar_1 + y_rcar_2 * t47[x] + y_rcar_3 * rcar[-1] + rg5[0] * y_rcar_4 + y_rcar_5 * rg5[-1] + + rccd[0] = rcar[0] + jrcd[x] * y_rccd_1 - zpi5[0] + + rcch[0] = jrh[x] * y_rcch_1 + y_rcch_2 * trfpm[x] + y_rcch_3 * rme[0] + y_rcch_4 * trspp[x] - zpi10[0] + + rcgain[0] = picx4[0] + rcgain_aerr[x] + xgap2[0] * y_rcgain_1 + y_rcgain_2 * rcgain[-1] + y_rcgain_3 * picx4[-1] + + req[0] = rg30[0] - zpic30[0] + reqp[0] + + reqp[0] = reqp_aerr[x] + rbbbp[0] * y_reqp_1 + y_reqp_2 * reqp[-1] + rbbbp[-1] * y_reqp_3 + + rfynic[0] = rfynil[0] * y_rfynic_4 + y_rfynic_1 * rfynic[-1] + rfynic_aerr[x] + y_rfynic_2 * rfynil[-1] + y_rfynic_3 * rfynic[-2] + + rfynil[0] = reqp[0] * y_rfynil_8 + y_rfynil_7 * rtb[0] + rg10[0] * y_rfynil_6 + rfynil[-1] * y_rfynil_1 + rfynil_aerr[x] + y_rfynil_2 * rg10[-1] + y_rfynil_3 * rtb[-1] + reqp[-1] * y_rfynil_4 + y_rfynil_5 * rfynil[-2] + + rg10[0] = zrff10[0] + rg10p[0] + + rg10p[0] = rg10p_aerr[x] + zgap10[0] * y_rg10p_1 + y_rg10p_2 * d8095[x] + y_rg10p_3 * rg10p[-1] + zgap10[-1] * y_rg10p_4 + y_rg10p_5 * d8095[x-1] + + rg30[0] = zrff30[0] + rg30p[0] + + rg30p[0] = rg30p_aerr[x] + y_rg30p_1 * zgap30[0] + d8095[x] * y_rg30p_2 + y_rg30p_3 * rg30p[-1] + y_rg30p_4 * zgap30[-1] + y_rg30p_5 * d8095[x-1] + + rg5[0] = zrff5[0] + rg5p[0] + + rg5p[0] = rg5p_aerr[x] + y_rg5p_1 * zgap05[0] + y_rg5p_2 * rg5p[-1] + y_rg5p_3 * zgap05[-1] + + rgfint[0] = gfdbtn_l[-1] * y_rgfint_4 + rgfint_aerr[x] + y_rgfint_1 * rgfint[-1] + y_rgfint_2 * rgw[-1] + y_rgfint_3 * gfdbtn_l[-2] + + rgw[0] = rtb[0] * y_rgw_1 + rg5[0] * y_rgw_2 + rg10[0] * y_rgw_3 + rg30[0] * y_rgw_4 + + rme[0] = rme[-1] * y_rme_1 + rme_aerr[x] + rg10[0] * y_rme_2 + rg10[-1] * y_rme_3 + y_rme_4 * d87[x] + + rrff[0] = rff[0] + picxfe[0] * y_rrff_1 + picxfe[-1] * y_rrff_2 + y_rrff_3 * picxfe[-2] + y_rrff_4 * picxfe[-3] + + rrtr[0] = y_rrtr_1 * rrtr[-1] + rrff[0] * y_rrtr_2 + + rspnia[0] = y_rspnia_1 * yhsn_l[0] + y_rspnia_2 * ydn_l[0] + + # rstar[0] = rstar[-1] + + rtb[0] = rff[-1] * y_rtb_4 + rff[0] * y_rtb_3 + rtb[-1] * y_rtb_1 + y_rtb_2 * rtb[-2] + + rtbfi_l[0] = pxp_l[0] + rbfi[0] * y_rtbfi_l_2 + jrbfi[x] * y_rtbfi_l_3 + hgpbfir[0] * y_rtbfi_l_4 + y_rtbfi_l_5 * tritc[x] + trfcim[x] * y_rtbfi_l_6 + y_rtbfi_l_7 * tapddp[x] + y_rtbfi_l_8 * tdpv[x] + pkbfir[0] * y_rtbfi_l_9 - pxb_l[0] + + rtinv[0] = pxb_l[0] * y_rtinv_7 + rbfi[0] * y_rtinv_1 + hgpkir[0] * y_rtinv_2 + pxp_l[0] * y_rtinv_3 + pkir[x] * y_rtinv_4 + pxp_l[-1] * y_rtinv_5 + y_rtinv_6 * pkir[x-1] + + rtr[0] = ptr[0] + rrtr[0] + + tcin_l[0] = ynicpn_l[0] + y_tcin_l_2 * trci[0] + + tpn_l[0] = y_tpn_l_2 * trp[0] + y_tpn_l_3 * ypn_l[0] + gtn_l[0] * y_tpn_l_4 + + trci[0] = xgap2[-1] * y_trci_4 + trci_aerr[x] + trcit[x] + xgap2[0] * y_trci_1 + y_trci_2 * trci[-1] + y_trci_3 * trcit[x-1] + + # trp[0] = xgap2[0] * y_trp_5 + trp_aerr[x] + trpt[0] + y_trp_1 * trp[-1] + y_trp_2 * trpt[-1] + y_trp_3 * trp[-2] + y_trp_4 * trpt[-2] + trp[0] = xgap2[0] * y_trp_5 + trp_a[0] + trpt[0] + y_trp_1 * trp[-1] + y_trp_2 * trpt[-1] + y_trp_3 * trp[-2] + y_trp_4 * trpt[-2] + + trp_a[0] = (1 - rho_trp_a) * trp_ā + rho_trp_a * trp_a[-1] + trp_aerr[x] + + trpt[0] = trpts[0] + + trptd[0] = y_trptd_6 * gfdrt[x-2] + y_trptd_5 * xgdpn_l[-2] + y_trptd_4 * gfdbtn_l[-2] + trpt[-1] + gfdbtnp_l[-1] * y_trptd_1 + xgdpn_l[-1] * y_trptd_2 + y_trptd_3 * gfdrt[x-1] + + trpts[0] = xgap2[-1] * y_trpts_5 + trpt[-1] + y_trpts_1 * gfrecn_l[-1] + y_trpts_2 * gfexpn_l[-1] + xgdpn_l[-1] * y_trpts_3 + y_trpts_4 * gfsrt[-1] + + gfsrt[0] = rho_gfsrt * gfsrt[-1] + gfsrt_err[x] + + tryh[0] = tpn_l[0] * y_tryh_1 + y_tryh_2 * yhln_l[0] + y_tryh_3 * yhptn_l[0] + + uqpct_l[0] = huqpct[0] + uqpct_l[-1] + + uxbt_l[0] = uxbt_l[-1] + huxb[0] * y_uxbt_l_1 + + uynicpnr[0] = y_uynicpnr_1 * uynicpnr[-1] + + vbfi[0] = y_vbfi_1 * uvbfi[x] + pkbfir[0] * y_vbfi_2 + pbfir_l[0] * y_vbfi_3 + rtbfi_l[0] * y_vbfi_4 + + wpo_l[0] = wpon_l[0] - pcnia_l[0] + + wpon_l[0] = y_wpon_l_2 * wpon_l[-1] + rcgain[0] * y_wpon_l_3 + phouse_l[0] * y_wpon_l_4 + phouse_l[-1] * y_wpon_l_5 + ydn_l[0] * y_wpon_l_6 + ecnian_l[0] * y_wpon_l_7 + y_wpon_l_8 * yhibn_l[0] + pcdr_l[0] * y_wpon_l_9 + pcnia_l[0] * y_wpon_l_10 + ecd_l[0] * y_wpon_l_11 + jkcd_l[0] * y_wpon_l_12 + + wps_l[0] = wpsn_l[0] - pcnia_l[0] + + wpsn_l[0] = ynicpn_l[0] * y_wpsn_l_1 + tcin_l[0] * y_wpsn_l_2 + req[0] * y_wpsn_l_3 + y_wpsn_l_4 * zdivgr[0] + + xb_l[0] = y_xb_l_2 * xbn_l[0] + pxb_l[0] * y_xb_l_3 + + xbn_l[0] = pxb_l[0] * y_xbn_l_2 + xbo_l[0] * y_xbn_l_3 + xgdpn_l[0] * y_xbn_l_4 + y_xbn_l_5 * xgdo_l[0] + pgdp_l[0] * y_xbn_l_6 + + xbo_l[0] = xbt_l[0] + xgap2[0] * y_xbo_l_1 + + xbt_l[0] = mfpt_l[0] + leppot_l[0] * y_xbt_l_1 + qlww_l[0] * y_xbt_l_2 + lqualt_l[x] * y_xbt_l_3 + ks_l[0] * y_xbt_l_4 + xbtr_l[0] + + xbtr_l[0] = y_xbtr_l_1 * xbtr_l[-1] + + xfs_l[0] = xfs_l[-1] + ecnia_l[0] * y_xfs_l_1 + ecnia_l[-1] * y_xfs_l_2 + eh_l[0] * y_xfs_l_3 + eh_l[-1] * y_xfs_l_4 + ebfi_l[0] * y_xfs_l_5 + ebfi_l[-1] * y_xfs_l_6 + egfe_l[0] * y_xfs_l_7 + egfe_l[-1] * y_xfs_l_8 + egfl_l[0] * y_xfs_l_9 + egfl_l[-1] * y_xfs_l_10 + egse_l[0] * y_xfs_l_11 + egse_l[-1] * y_xfs_l_12 + egsl_l[0] * y_xfs_l_13 + egsl_l[-1] * y_xfs_l_14 + ex_l[0] * y_xfs_l_15 + ex_l[-1] * y_xfs_l_16 + emo_l[0] * y_xfs_l_17 + emo_l[-1] * y_xfs_l_18 + emp_l[0] * y_xfs_l_19 + emp_l[-1] * y_xfs_l_20 + + xfsn_l[0] = xgdpn_l[0] * y_xfsn_l_2 + pkir[x] * y_xfsn_l_3 + pxp_l[0] * y_xfsn_l_4 + ki_l[0] * y_xfsn_l_5 + ki_l[-1] * y_xfsn_l_6 + + xgap[0] = xbo_l[0] * y_xgap_1 + xbt_l[0] * y_xgap_2 + + xgap2[0] = xgdo_l[0] * y_xgap2_1 + xgdpt_l[0] * y_xgap2_2 + + xgdi_l[0] = mei_l + xgdo_l[0] + + xgdin_l[0] = pgdp_l[0] + xgdi_l[0] + + xgdo_l[0] = xgdp_l[0] - mep_l + + xgdp_l[0] = xgdp_l[-1] + xfs_l[0] * y_xgdp_l_1 + xfs_l[-1] * y_xgdp_l_2 + ki_l[0] * y_xgdp_l_3 + ki_l[-1] * y_xgdp_l_4 + y_xgdp_l_5 * ki_l[-2] + + xgdpn_l[0] = y_xgdpn_l_2 * xpn_l[0] + egfln_l[0] * y_xgdpn_l_3 + egsln_l[0] * y_xgdpn_l_4 + emn_l[0] * y_xgdpn_l_5 + pkir[x] * y_xgdpn_l_6 + pxp_l[0] * y_xgdpn_l_7 + ki_l[0] * y_xgdpn_l_8 + ki_l[-1] * y_xgdpn_l_9 + + xgdpt_l[0] = xbt_l[0] + uxbt_l[0] + + xgdptn_l[0] = pgdp_l[0] + xgdpt_l[0] + + xp_l[0] = xp_l[-1] + ecnia_l[0] * y_xp_l_1 + ecnia_l[-1] * y_xp_l_2 + eh_l[0] * y_xp_l_3 + eh_l[-1] * y_xp_l_4 + ebfi_l[0] * y_xp_l_5 + ebfi_l[-1] * y_xp_l_6 + egfe_l[0] * y_xp_l_7 + egfe_l[-1] * y_xp_l_8 + egse_l[0] * y_xp_l_9 + egse_l[-1] * y_xp_l_10 + ex_l[0] * y_xp_l_11 + ex_l[-1] * y_xp_l_12 + + xpn_l[0] = pxp_l[0] + xp_l[0] + + ydn_l[0] = y_ydn_l_2 * uyd[x] + ypn_l[0] * y_ydn_l_3 + tpn_l[0] * y_ydn_l_4 + + yh_l[0] = yhl_l[0] * y_yh_l_2 + yht_l[0] * y_yh_l_3 + y_yh_l_4 * yhp_l[0] + + yhgap[0] = y_yhgap_1 * yhshr_l[0] + y_yhgap_2 * zyhst_l[0] + + yhibn_l[0] = xgdpn_l[0] + y_yhibn_l_2 * uyhibn[x] + + yhl_l[0] = yhln_l[0] + tryh[0] * y_yhl_l_2 - pcnia_l[0] + + yhln_l[0] = y_yhln_l_2 * uyhln[x] + yniln_l[0] + + yhp_l[0] = tryh[0] * y_yhp_l_2 + yhptn_l[0] * y_yhp_l_3 + y_yhp_l_4 * yhpntn_l[0] - pcnia_l[0] + + yhpcd_l[0] = kcd_l[-1] + + yhpgap[0] = y_yhpgap_1 * yhpshr_l[0] + y_yhpgap_2 * zyhpst_l[0] + + yhpntn_l[0] = pcnia_l[0] * y_yhpntn_l_2 + pcdr_l[0] * y_yhpntn_l_3 + yhpcd_l[0] * y_yhpntn_l_4 + yhibn_l[0] * y_yhpntn_l_5 + ynicpn_l[0] * y_yhpntn_l_6 + tcin_l[0] * y_yhpntn_l_7 + y_yhpntn_l_8 * ynidn_l[0] + zpi10[0] * y_yhpntn_l_9 + gfdbtn_l[0] * y_yhpntn_l_10 + + yhpshr_l[0] = yhp_l[0] - yh_l[0] + + yhptn_l[0] = y_yhptn_l_2 * uyhptn[x] + y_yhptn_l_3 * ynirn_l[0] + gfintn_l[0] * y_yhptn_l_4 + ynidn_l[0] * y_yhptn_l_5 + yhibn_l[0] * y_yhptn_l_6 + + yhshr_l[0] = yh_l[0] * y_yhshr_l_2 + xgdp_l[0] * y_yhshr_l_3 + + yhsn_l[0] = yhln_l[0] * y_yhsn_l_2 + y_yhsn_l_3 * yhtn_l[0] + yhptn_l[0] * y_yhsn_l_4 + tpn_l[0] * y_yhsn_l_5 + ecnian_l[0] * y_yhsn_l_6 + yhibn_l[0] * y_yhsn_l_7 + y_yhsn_l_8 * uyhsn[x] + xgdptn_l[0] * y_yhsn_l_9 + + yht_l[0] = yhtn_l[0] - pcnia_l[0] + + yhtgap[0] = y_yhtgap_1 * yhtshr_l[0] + y_yhtgap_2 * zyhtst_l[0] + + yhtn_l[0] = gtn_l[0] + y_yhtn_l_2 * uyhtn[x] + + yhtshr_l[0] = yht_l[0] - yh_l[0] + + ykbfin_l[0] = pxb_l[0] + rtbfi_l[0] + kbfi_l[0] * y_ykbfin_l_2 + kbfi_l[-1] * y_ykbfin_l_3 + + ykin_l[0] = pxb_l[0] + rtinv[0] * y_ykin_l_2 + ki_l[0] * y_ykin_l_3 + ki_l[-1] * y_ykin_l_4 + + ynicpn_l[0] = y_ynicpn_l_2 * ynin_l[0] + yniln_l[0] * y_ynicpn_l_3 + ynirn_l[0] * y_ynicpn_l_4 + uynicpnr[0] * y_ynicpn_l_5 + xgdpn_l[0] * y_ynicpn_l_6 + + ynidn_l[0] = zynid[0] + y_ynidn_l_8 * pxb_l[-2] + y_ynidn_l_7 * ymsdn[x-2] + y_ynidn_l_6 * ynidn_l[-2] + y_ynidn_l_5 * qynidn_l[-1] + ynidn_l_aerr[x] + pxb_l[-1] * y_ynidn_l_4 + pxb_l[0] + y_ynidn_l_1 * ymsdn[x] + y_ynidn_l_2 * ynidn_l[-1] + y_ynidn_l_3 * ymsdn[x-1] + + yniln_l[0] = y_yniln_l_2 * uyl[x] + pl_l[0] * y_yniln_l_3 + lhp_l[0] * y_yniln_l_4 + pgfl_l[0] * y_yniln_l_5 + egfl_l[0] * y_yniln_l_6 + pgsl_l[0] * y_yniln_l_7 + egsl_l[0] * y_yniln_l_8 + + ynin_l[0] = y_ynin_l_2 * uyni[x] + xgdin_l[0] * y_ynin_l_3 + fynicn_l[0] * y_ynin_l_4 + fyniln_l[0] * y_ynin_l_5 + jccan_l[0] * y_ynin_l_6 + + ynirn_l[0] = xgdpn_l[0] + y_ynirn_l_1 * ynirn_l_aerr[x] + y_ynirn_l_2 * ynirn_l[-1] + xgdpn_l[-1] * y_ynirn_l_3 + rbbb[0] * y_ynirn_l_4 + y_ynirn_l_5 * rbbb[-1] + + ypn_l[0] = y_ypn_l_2 * uyp[x] + yhln_l[0] * y_ypn_l_3 + yhtn_l[0] * y_ypn_l_4 + yhptn_l[0] * y_ypn_l_5 + + zdivgr[0] = y_zdivgr_1 * hgynid[1] + y_zdivgr_2 * zdivgr[1] + + zebfi[0] = hgpbfir[-1] * y_zebfi_21 + y_zebfi_20 * hxbt[-1] + qebfi_l[-1] * y_zebfi_15 + y_zebfi_11 * xgap[-1] + ptr[-1] * y_zebfi_10 + y_zebfi_9 * rtr[-1] + rff[-1] * y_zebfi_5 + picnia[-1] * y_zebfi_1 + y_zebfi_2 * picnia[-2] + y_zebfi_3 * picnia[-3] + y_zebfi_4 * picnia[-4] + y_zebfi_6 * rff[-2] + y_zebfi_7 * rff[-3] + y_zebfi_8 * rff[-4] + y_zebfi_12 * xgap[-2] + y_zebfi_13 * xgap[-3] + y_zebfi_14 * xgap[-4] + y_zebfi_16 * qebfi_l[-2] + y_zebfi_17 * qebfi_l[-3] + y_zebfi_18 * qebfi_l[-4] + y_zebfi_19 * qebfi_l[-5] + + zecd[0] = y_zecd_33 * qecd_l[-5] + y_zecd_32 * qecd_l[-4] + y_zecd_31 * qecd_l[-3] + y_zecd_30 * qecd_l[-2] + qecd_l[-1] * y_zecd_29 + hggdpt[-1] * y_zecd_27 + y_zecd_23 * yhpgap[-1] + y_zecd_19 * yhtgap[-1] + y_zecd_15 * yhgap[-1] + rtr[-1] * y_zecd_14 + ptr[-1] * y_zecd_13 + xgap2[-1] * y_zecd_9 + rff[-1] * y_zecd_5 + picnia[-1] * y_zecd_1 + y_zecd_2 * picnia[-2] + y_zecd_3 * picnia[-3] + y_zecd_4 * picnia[-4] + y_zecd_6 * rff[-2] + y_zecd_7 * rff[-3] + y_zecd_8 * rff[-4] + y_zecd_10 * xgap2[-2] + y_zecd_11 * xgap2[-3] + y_zecd_12 * xgap2[-4] + y_zecd_16 * yhgap[-2] + y_zecd_17 * yhgap[-3] + y_zecd_18 * yhgap[-4] + y_zecd_20 * yhtgap[-2] + y_zecd_21 * yhtgap[-3] + y_zecd_22 * yhtgap[-4] + y_zecd_24 * yhpgap[-2] + y_zecd_25 * yhpgap[-3] + y_zecd_26 * yhpgap[-4] + y_zecd_28 * hgpcdr[x-1] + + zeco[0] = qeco_l[-1] * y_zeco_28 + hggdpt[-1] * y_zeco_27 + yhpgap[-1] * y_zeco_23 + yhtgap[-1] * y_zeco_19 + yhgap[-1] * y_zeco_15 + rtr[-1] * y_zeco_14 + ptr[-1] * y_zeco_13 + xgap2[-1] * y_zeco_9 + rff[-1] * y_zeco_5 + picnia[-1] * y_zeco_1 + y_zeco_2 * picnia[-2] + y_zeco_3 * picnia[-3] + y_zeco_4 * picnia[-4] + y_zeco_6 * rff[-2] + y_zeco_7 * rff[-3] + y_zeco_8 * rff[-4] + y_zeco_10 * xgap2[-2] + y_zeco_11 * xgap2[-3] + y_zeco_12 * xgap2[-4] + y_zeco_16 * yhgap[-2] + y_zeco_17 * yhgap[-3] + y_zeco_18 * yhgap[-4] + y_zeco_20 * yhtgap[-2] + y_zeco_21 * yhtgap[-3] + y_zeco_22 * yhtgap[-4] + y_zeco_24 * yhpgap[-2] + y_zeco_25 * yhpgap[-3] + y_zeco_26 * yhpgap[-4] + y_zeco_29 * qeco_l[-2] + y_zeco_30 * qeco_l[-3] + y_zeco_31 * qeco_l[-4] + y_zeco_32 * qeco_l[-5] + + zeh[0] = qeh_l[-1] * y_zeh_28 + hggdpt[-1] * y_zeh_27 + yhpgap[-1] * y_zeh_23 + yhtgap[-1] * y_zeh_19 + yhgap[-1] * y_zeh_15 + rtr[-1] * y_zeh_14 + ptr[-1] * y_zeh_13 + xgap2[-1] * y_zeh_9 + rff[-1] * y_zeh_5 + picnia[-1] * y_zeh_1 + y_zeh_2 * picnia[-2] + y_zeh_3 * picnia[-3] + y_zeh_4 * picnia[-4] + y_zeh_6 * rff[-2] + y_zeh_7 * rff[-3] + y_zeh_8 * rff[-4] + y_zeh_10 * xgap2[-2] + y_zeh_11 * xgap2[-3] + y_zeh_12 * xgap2[-4] + y_zeh_16 * yhgap[-2] + y_zeh_17 * yhgap[-3] + y_zeh_18 * yhgap[-4] + y_zeh_20 * yhtgap[-2] + y_zeh_21 * yhtgap[-3] + y_zeh_22 * yhtgap[-4] + y_zeh_24 * yhpgap[-2] + y_zeh_25 * yhpgap[-3] + y_zeh_26 * yhpgap[-4] + y_zeh_29 * qeh_l[-2] + y_zeh_30 * qeh_l[-3] + y_zeh_31 * qeh_l[-4] + y_zeh_32 * qeh_l[-5] + + zgap05[0] = xgap[0] * y_zgap05_1 + y_zgap05_2 * zgap05[1] + + zgap10[0] = xgap[0] * y_zgap10_1 + y_zgap10_2 * zgap10[1] + + zgap30[0] = xgap[0] * y_zgap30_1 + y_zgap30_2 * zgap30[1] + + zgapc2[0] = rtr[-1] * y_zgapc2_14 + ptr[-1] * y_zgapc2_13 + xgap2[-1] * y_zgapc2_9 + rff[-1] * y_zgapc2_5 + picnia[-1] * y_zgapc2_1 + y_zgapc2_2 * picnia[-2] + y_zgapc2_3 * picnia[-3] + y_zgapc2_4 * picnia[-4] + y_zgapc2_6 * rff[-2] + y_zgapc2_7 * rff[-3] + y_zgapc2_8 * rff[-4] + y_zgapc2_10 * xgap2[-2] + y_zgapc2_11 * xgap2[-3] + y_zgapc2_12 * xgap2[-4] + + zlhp[0] = hqlww[-1] * y_zlhp_20 + y_zlhp_19 * hlept[-1] + y_zlhp_17 * lprdt_l[-1] + xbo_l[-1] * y_zlhp_15 + xgap[-1] * y_zlhp_11 + ptr[-1] * y_zlhp_10 + rtr[-1] * y_zlhp_9 + rff[-1] * y_zlhp_5 + picnia[-1] * y_zlhp_1 + y_zlhp_2 * picnia[-2] + y_zlhp_3 * picnia[-3] + y_zlhp_4 * picnia[-4] + y_zlhp_6 * rff[-2] + y_zlhp_7 * rff[-3] + y_zlhp_8 * rff[-4] + y_zlhp_12 * xgap[-2] + y_zlhp_13 * xgap[-3] + y_zlhp_14 * xgap[-4] + y_zlhp_16 * xbo_l[-2] + y_zlhp_18 * lprdt_l[-2] + + zpi10[0] = picnia[0] * y_zpi10_1 + y_zpi10_2 * zpi10[1] + + zpi10f[0] = picnia[0] * y_zpi10f_1 + y_zpi10f_2 * zpi10f[1] + + zpi5[0] = xgap[-1] * y_zpi5_11 + ptr[-1] * y_zpi5_10 + rtr[-1] * y_zpi5_9 + rff[-1] * y_zpi5_5 + picnia[-1] * y_zpi5_1 + y_zpi5_2 * picnia[-2] + y_zpi5_3 * picnia[-3] + y_zpi5_4 * picnia[-4] + y_zpi5_6 * rff[-2] + y_zpi5_7 * rff[-3] + y_zpi5_8 * rff[-4] + y_zpi5_12 * xgap[-2] + y_zpi5_13 * xgap[-3] + y_zpi5_14 * xgap[-4] + + zpib5[0] = pxb_l[0] * y_zpib5_1 + pxb_l[-1] * y_zpib5_2 + y_zpib5_3 * zpib5[1] + + zpic30[0] = picnia[0] * y_zpic30_1 + y_zpic30_2 * zpic30[1] + + zpic58[0] = pic4[8] + + zpicxfe[0] = lurnat[-1] * y_zpicxfe_26 + lur[-1] * y_zpicxfe_25 + huqpct[-1] * y_zpicxfe_24 + hlprdt[-1] * y_zpicxfe_23 + pl_l[-1] * y_zpicxfe_22 + qpl_l[-1] * y_zpicxfe_21 + pcnia_l[-1] * y_zpicxfe_20 + qpcnia_l[-1] * y_zpicxfe_19 + ptr[-1] * y_zpicxfe_18 + rtr[-1] * y_zpicxfe_17 + xgap2[-1] * y_zpicxfe_13 + rff[-1] * y_zpicxfe_9 + pieci[-1] * y_zpicxfe_5 + picxfe[-1] * y_zpicxfe_1 + y_zpicxfe_2 * picxfe[-2] + y_zpicxfe_3 * picxfe[-3] + y_zpicxfe_4 * picxfe[-4] + y_zpicxfe_6 * pieci[-2] + y_zpicxfe_7 * pieci[-3] + y_zpicxfe_8 * pieci[-4] + y_zpicxfe_10 * rff[-2] + y_zpicxfe_11 * rff[-3] + y_zpicxfe_12 * rff[-4] + y_zpicxfe_14 * xgap2[-2] + y_zpicxfe_15 * xgap2[-3] + y_zpicxfe_16 * xgap2[-4] + y_zpicxfe_27 * lur[-2] + y_zpicxfe_28 * lurnat[-2] + + zpieci[0] = lurnat[-1] * y_zpieci_26 + lur[-1] * y_zpieci_25 + huqpct[-1] * y_zpieci_24 + hlprdt[-1] * y_zpieci_23 + pl_l[-1] * y_zpieci_22 + qpl_l[-1] * y_zpieci_21 + pcnia_l[-1] * y_zpieci_20 + qpcnia_l[-1] * y_zpieci_19 + ptr[-1] * y_zpieci_18 + rtr[-1] * y_zpieci_17 + xgap2[-1] * y_zpieci_13 + rff[-1] * y_zpieci_9 + pieci[-1] * y_zpieci_5 + picxfe[-1] * y_zpieci_1 + y_zpieci_2 * picxfe[-2] + y_zpieci_3 * picxfe[-3] + y_zpieci_4 * picxfe[-4] + y_zpieci_6 * pieci[-2] + y_zpieci_7 * pieci[-3] + y_zpieci_8 * pieci[-4] + y_zpieci_10 * rff[-2] + y_zpieci_11 * rff[-3] + y_zpieci_12 * rff[-4] + y_zpieci_14 * xgap2[-2] + y_zpieci_15 * xgap2[-3] + y_zpieci_16 * xgap2[-4] + y_zpieci_27 * lur[-2] + y_zpieci_28 * lurnat[-2] + + zrff10[0] = rff[0] * y_zrff10_1 + y_zrff10_2 * zrff10[1] + + zrff30[0] = rff[0] * y_zrff30_1 + y_zrff30_2 * zrff30[1] + + zrff5[0] = rff[0] * y_zrff5_1 + y_zrff5_2 * zrff5[1] + + zyh_l[0] = xgdpt_l[0] + zyhst_l[0] + yhgap[-1] * y_zyh_l_16 + yhgap[0] * y_zyh_l_15 + rtr[0] * y_zyh_l_14 + ptr[0] * y_zyh_l_13 + xgap2[-1] * y_zyh_l_10 + xgap2[0] * y_zyh_l_9 + rff[-1] * y_zyh_l_6 + rff[0] * y_zyh_l_5 + picnia[0] * y_zyh_l_1 + picnia[-1] * y_zyh_l_2 + y_zyh_l_3 * picnia[-2] + y_zyh_l_4 * picnia[-3] + y_zyh_l_7 * rff[-2] + y_zyh_l_8 * rff[-3] + y_zyh_l_11 * xgap2[-2] + y_zyh_l_12 * xgap2[-3] + y_zyh_l_17 * yhgap[-2] + y_zyh_l_18 * yhgap[-3] + + zyhp_l[0] = xgdpt_l[0] + zyhst_l[0] + zyhpst_l[0] + yhpgap[-1] * y_zyhp_l_20 + yhpgap[0] * y_zyhp_l_19 + yhgap[-1] * y_zyhp_l_16 + yhgap[0] * y_zyhp_l_15 + rtr[0] * y_zyhp_l_14 + ptr[0] * y_zyhp_l_13 + xgap2[-1] * y_zyhp_l_10 + xgap2[0] * y_zyhp_l_9 + rff[-1] * y_zyhp_l_6 + rff[0] * y_zyhp_l_5 + picnia[0] * y_zyhp_l_1 + picnia[-1] * y_zyhp_l_2 + y_zyhp_l_3 * picnia[-2] + y_zyhp_l_4 * picnia[-3] + y_zyhp_l_7 * rff[-2] + y_zyhp_l_8 * rff[-3] + y_zyhp_l_11 * xgap2[-2] + y_zyhp_l_12 * xgap2[-3] + y_zyhp_l_17 * yhgap[-2] + y_zyhp_l_18 * yhgap[-3] + y_zyhp_l_21 * yhpgap[-2] + y_zyhp_l_22 * yhpgap[-3] + + zyhpst_l[0] = zyhpst_l[-1] + yhpgap[-1] * y_zyhpst_l_1 + + zyhst_l[0] = zyhst_l[-1] + yhgap[-1] * y_zyhst_l_1 + + zyht_l[0] = xgdpt_l[0] + zyhst_l[0] + zyhtst_l[0] + yhtgap[-1] * y_zyht_l_20 + yhtgap[0] * y_zyht_l_19 + yhgap[-1] * y_zyht_l_16 + yhgap[0] * y_zyht_l_15 + rtr[0] * y_zyht_l_14 + ptr[0] * y_zyht_l_13 + xgap2[-1] * y_zyht_l_10 + xgap2[0] * y_zyht_l_9 + rff[-1] * y_zyht_l_6 + rff[0] * y_zyht_l_5 + picnia[0] * y_zyht_l_1 + picnia[-1] * y_zyht_l_2 + y_zyht_l_3 * picnia[-2] + y_zyht_l_4 * picnia[-3] + y_zyht_l_7 * rff[-2] + y_zyht_l_8 * rff[-3] + y_zyht_l_11 * xgap2[-2] + y_zyht_l_12 * xgap2[-3] + y_zyht_l_17 * yhgap[-2] + y_zyht_l_18 * yhgap[-3] + y_zyht_l_21 * yhtgap[-2] + y_zyht_l_22 * yhtgap[-3] + + zyhtst_l[0] = zyhtst_l[-1] + yhtgap[-1] * y_zyhtst_l_1 + + zynid[0] = hggdpt[-1] * y_zynid_25 + pxb_l[-1] * y_zynid_16 + qynidn_l[-1] * y_zynid_15 + xgap[-1] * y_zynid_11 + ptr[-1] * y_zynid_10 + rtr[-1] * y_zynid_9 + rff[-1] * y_zynid_5 + picnia[-1] * y_zynid_1 + y_zynid_2 * picnia[-2] + y_zynid_3 * picnia[-3] + y_zynid_4 * picnia[-4] + y_zynid_6 * rff[-2] + y_zynid_7 * rff[-3] + y_zynid_8 * rff[-4] + y_zynid_12 * xgap[-2] + y_zynid_13 * xgap[-3] + y_zynid_14 * xgap[-4] + y_zynid_17 * qynidn_l[-2] + y_zynid_18 * pxb_l[-2] + y_zynid_19 * qynidn_l[-3] + y_zynid_20 * pxb_l[-3] + y_zynid_21 * qynidn_l[-4] + y_zynid_22 * pxb_l[-4] + y_zynid_23 * qynidn_l[-5] + y_zynid_24 * pxb_l[-5] + + ugap[0] = lur[0] - lurnat[0] + + rff[0] = rule[0] + eradd[x] + + rule[0] = rff[-1] * .85 + rstar * .15 + picx4[0] * .225 - 0.075 * pitarg[x] + xgap2[0] * .15 + + fiscal[0] = (1 - rho_fiscal) * f̄iscal + rho_fiscal * fiscal[-1] + fiscal_aerr[x] + + fiscalav[0] = av * fiscal[0] + fiscalav[-1] * rho_fiscalav + + + gov_exp_share[0] = egfe_l[0] * y_xfs_l_7 * 100 + + income_tax_share_of_gdp[0] = 100 * (y_yh_l_2 * (-y_yhl_l_2 - 1) + y_yh_l_4 * (-y_yhp_l_2 - 1)) * tryh[0] + + debt_to_gdp[0] = - gfdbtnp_l[0] * y_gfdbtnp_l_4 * y_gfrecn_l_4 * y_gfrecn_l_5 * 100 +end + + +@parameters FRBUS begin + mep_l = 0 + + mei_l = 0 + + qpmo_l = 0 + + rstar = 0 + + rho_qkir_l = 0.8 + + y_dpgap_1 = 0.0025 + + y_dpgap_2 = (-0.103649883938) + + y_dpgap_3 = 0.103649883938 + + y_dpgap_4 = (-0.341041547027) + + y_dpgap_5 = 0.341041547027 + + y_dpgap_6 = (-0.121366054939) + + y_dpgap_7 = 0.121366054939 + + y_dpgap_8 = (-0.104958882473) + + y_dpgap_9 = 0.104958882473 + + y_dpgap_10 = (-0.328983631622) + + y_dpgap_11 = 0.328983631622 + + y_ebfi_l_1 = 1.27660626172 + + y_ebfi_l_2 = 0.0453619253429 + + y_ebfi_l_3 = (-0.135655771316) + + y_ebfi_l_4 = (-0.18631241575) + + y_ebfi_l_5 = 0.616485384319 + + y_ebfi_l_6 = 0.383514615681 + + y_ebfi_l_7 = (-0.383514615681) + + y_ebfi_l_8 = (-0.000958786539202) + + y_ebfin_l_1 = 0.000349694902126 + + y_ec_l_1 = 0.7310605131 + + y_ec_l_2 = (-0.7310605131) + + y_ec_l_3 = 0.157421136 + + y_ec_l_4 = (-0.157421136) + + y_ec_l_5 = 0.0223688796433 + + y_ec_l_6 = 0.0891494712567 + + y_ec_l_7 = (-0.0223688796433) + + y_ec_l_8 = (-0.0891494712567) + + y_ecd_l_1 = 0.78385727975 + + y_ecd_l_2 = 0.156149940356 + + y_ecd_l_3 = 0.0599927798938 + + y_ecd_l_4 = 0.0296796460069 + + y_ech_l_1 = 1.71348425234 + + y_ech_l_2 = (-1.71348425234) + + y_ech_l_3 = 9.76051187168 + + y_ech_l_4 = (-0.718706571642) + + y_ech_l_5 = 0.718706571642 + + y_ecnia_l_1 = 0.735 + + y_ecnia_l_2 = (-0.735) + + y_ecnia_l_3 = 0.1055 + + y_ecnia_l_4 = (-0.1055) + + y_ecnia_l_5 = 0.1595 + + y_ecnia_l_6 = (-0.1595) + + y_ecnian_l_1 = 7.05661360558e-05 + + y_eco_l_1 = 1.17546755467 + + y_eco_l_2 = 0.109703169694 + + y_eco_l_3 = (-0.285170724366) + + y_eco_l_4 = 0.692476259501 + + y_eco_l_5 = 0.229572174835 + + y_eco_l_6 = 0.0779515656641 + + y_eco_l_7 = (-0.229612885136) + + y_eco_l_8 = (-0.0779108553636) + + y_egfe_l_1 = 0.726276173623 + + y_egfe_l_2 = (-1.38339974044) + + y_egfe_l_3 = 0.0497143719338 + + y_egfe_l_4 = 0.103593759929 + + y_egfe_l_5 = 1.50381543495 + + y_egfe_l_6 = (-0.000983552448045) + + y_egfe_l_7 = 0.000725681212301 + + y_egfen_l_1 = 0.0010878350668 + + y_egfet_l_1 = 0.9 + + y_egfet_l_2 = (-0.1) + + y_egfet_l_3 = (-0.1) + + y_egfet_l_4 = 0.1 + + y_egfet_l_5 = 0.000625 + + y_egfet_l_6 = 0.000625 + + y_egfet_l_7 = 0.000625 + + y_egfet_l_8 = 0.000625 + + y_egfl_l_1 = 1.16197632264 + + y_egfl_l_2 = (-1.12731388567) + + y_egfl_l_3 = (-0.302868541805) + + y_egfl_l_4 = 0.0613337937414 + + y_egfl_l_5 = 1.2068723111 + + y_egfl_l_6 = (-0.00250725401078) + + y_egfl_l_7 = 0.00235067489642 + + y_egfln_l_1 = 0.00218479904218 + + y_egflt_l_1 = 0.9 + + y_egflt_l_2 = (-0.1) + + y_egflt_l_3 = 0.1 + + y_egflt_l_4 = 0.000625 + + y_egflt_l_5 = 0.000625 + + y_egflt_l_6 = 0.000625 + + y_egflt_l_7 = 0.000625 + + y_egse_l_1 = 1.00049378528 + + y_egse_l_2 = (-0.797614647892) + + y_egse_l_3 = (-0.128950321813) + + y_egse_l_4 = (-0.00262964990773) + + y_egse_l_5 = 0.928700834331 + + y_egse_l_6 = 0.00158066587876 + + y_egse_l_7 = (-0.000853766092194) + + y_egsen_l_1 = 0.00117247778411 + + y_egset_l_1 = 0.9 + + y_egset_l_2 = (-0.1) + + y_egset_l_3 = (-0.1) + + y_egset_l_4 = 0.1 + + y_egset_l_5 = 0.000625 + + y_egset_l_6 = 0.000625 + + y_egset_l_7 = 0.000625 + + y_egset_l_8 = 0.000625 + + y_egsl_l_1 = 1.04483163655 + + y_egsl_l_2 = (-0.633546297018) + + y_egsl_l_3 = (-0.134688612832) + + y_egsl_l_4 = (-0.0215581541096) + + y_egsl_l_5 = 0.744961427412 + + y_egsl_l_6 = (-0.00143256549309) + + y_egsl_l_7 = 0.00176517379444 + + y_egsln_l_1 = 0.000707659055882 + + y_egslt_l_1 = 0.9 + + y_egslt_l_2 = (-0.1) + + y_egslt_l_3 = 0.1 + + y_egslt_l_4 = 0.000625 + + y_egslt_l_5 = 0.000625 + + y_egslt_l_6 = 0.000625 + + y_egslt_l_7 = 0.000625 + + y_eh_l_1 = 1.3576278254 + + y_eh_l_2 = 0.0130993143616 + + y_eh_l_3 = (-0.164666195693) + + y_eh_l_4 = (-0.206060944067) + + y_eh_l_5 = (-0.0282729007489) + + y_eh_l_6 = 0.0282729007489 + + y_eh_l_7 = (-0.000786966438108) + + y_ehn_l_1 = 0.00124036373046 + + y_em_l_1 = 0.0012598389126 + + y_em_l_2 = (-0.000856800245907) + + y_em_l_3 = 0.00125870045787 + + y_em_l_4 = (-0.000781093277566) + + y_em_l_5 = 0.928320853989 + + y_em_l_6 = (-0.928320853989) + + y_em_l_7 = (-0.000403038666697) + + y_em_l_8 = (-0.000477607180303) + + y_em_l_9 = 0.0716791460112 + + y_em_l_10 = (-0.0716791460112) + + y_emn_l_1 = 0.000320220965275 + + y_emn_l_2 = 0.928554219554 + + y_emn_l_3 = 0.0714457804463 + + y_emo_l_1 = 0.819289500318 + + y_emo_l_2 = (-0.180710499682) + + y_emo_l_3 = 1.31018224516 + + y_emo_l_4 = 0.180710499682 + + y_emo_l_5 = 0.0135818692772 + + y_emo_l_6 = 0.00278890259237 + + y_emo_l_7 = (-0.0163707718696) + + y_emo_l_8 = 0.723524924437 + + y_emo_l_9 = (-0.404694213855) + + y_emon_l_1 = 0.000344859738432 + + y_emp_l_1 = 40.1856146542 + + y_emp_l_2 = 0.048026 + + y_emp_l_3 = (-0.048026) + + y_emp_l_4 = (-0.048026) + + y_emp_l_5 = 0.048026 + + y_emp_l_6 = 0.022115 + + y_empn_l_1 = 0.00448201367911 + + y_ex_l_1 = 0.892272127137 + + y_ex_l_2 = (-0.107727872863) + + y_ex_l_3 = (-0.107727872863) + + y_ex_l_4 = (-0.107727872863) + + y_ex_l_5 = 0.107727872863 + + y_ex_l_6 = 0.107727872863 + + y_ex_l_7 = 0.0148164224533 + + y_ex_l_8 = (-0.0045419370785) + + y_ex_l_9 = (-0.0102744853748) + + y_ex_l_10 = 1.01585705046 + + y_exn_l_1 = 0.000395785791626 + + y_fcbn_l_1 = (-0.00219688240418) + + y_fcbn_l_2 = (-5.55068537239) + + y_fcbn_l_3 = 6.86052021077 + + y_fcbn_l_4 = (-2.52909715822) + + y_fcbn_l_5 = 1.9133340876 + + y_fcbn_l_6 = (-35.2463013113) + + y_fcbn_l_7 = 0.305928232246 + + y_fcbn_l_8 = 0.305928232246 + + y_fgdp_l_1 = 0.00843835585766 + + y_fgdp_l_2 = 0.01 + + y_fgdpt_l_1 = 0.9 + + y_fgdpt_l_2 = 0.1 + + y_fgdpt_l_3 = 0.000625 + + y_fgdpt_l_4 = 0.000625 + + y_fgdpt_l_5 = 0.000625 + + y_fgdpt_l_6 = 0.000625 + + y_fnicn_l_1 = 0.993277528339 + + y_fnicn_l_2 = 0.00672247166135 + + y_fnicn_l_3 = 0.892965336399 + + y_fnicn_l_4 = 0.537028034851 + + y_fnicn_l_5 = (-0.537028034851) + + y_fnicn_l_6 = (-0.66631256176) + + y_fnicn_l_7 = 0.66631256176 + + y_fnicn_l_8 = 0.892965336399 + + y_fniln_l_1 = 0.982046754178 + + y_fniln_l_2 = 0.985640979746 + + y_fniln_l_3 = 0.692942512139 + + y_fniln_l_4 = 0.0100008124223 + + y_fniln_l_5 = 0.00373870870246 + + y_fniln_l_6 = 0.315405113519 + + y_fniln_l_7 = (-0.315405113519) + + y_fniln_l_8 = (-0.0591384587847) + + y_fniln_l_9 = 0.0591384587847 + + y_fniln_l_10 = 0.00421372469752 + + y_fnirn_l_1 = (-0.00807177556398) + + y_fnirn_l_2 = (-169.102652771) + + y_fpc_l_1 = 0.00879025119382 + + y_fpc_l_2 = 0.0025 + + y_fpi10_1 = 0.156993726433 + + y_fpi10_2 = 0.156993726433 + + y_fpi10_3 = 0.156993726433 + + y_fpi10_4 = 0.156993726433 + + y_fpi10_5 = 0.372025094268 + + y_fpi10_6 = 0.32214582784 + + y_fpi10t_1 = 0.95 + + y_fpi10t_2 = 0.05 + + y_fpic_1 = 0.678829880162 + + y_fpic_2 = 0.321170119838 + + y_fpx_l_1 = 0.00804862709227 + + y_fpxr_l_1 = 0.048 + + y_fpxr_l_2 = (-0.048) + + y_fpxr_l_3 = (-0.048) + + y_fpxr_l_4 = 0.048 + + y_fpxr_l_5 = 0.563832456119 + + y_fpxr_l_6 = (-0.726654492224) + + y_fpxr_l_7 = 0.162822036105 + + y_fpxrr_l_1 = 1.18364909386 + + y_fpxrr_l_2 = (-0.00291888934318) + + y_fpxrr_l_3 = (-0.211089676177) + + y_fpxrr_l_4 = 0.00302407543125 + + y_frl10_1 = 0.988458285734 + + y_frl10_2 = (-0.29200997295) + + y_frl10_3 = (-0.0655047670227) + + y_frl10_4 = 0.369056454239 + + y_frl10_5 = 0.12455118125 + + y_frl10_6 = (-0.12455118125) + + y_frs10_1 = 4.78434763861 + + y_frs10_2 = 0 + + y_frs10_3 = 0.25 + + y_frs10_4 = 0.25 + + y_frs10_5 = 0.25 + + y_frs10_6 = 0.25 + + y_frs10_7 = 0 + + y_frs10_8 = 0 + + y_frstar_1 = 0.95 + + y_frstar_2 = 0.05 + + y_frstar_3 = (-0.0125) + + y_frstar_4 = (-0.0125) + + y_frstar_5 = (-0.0125) + + y_frstar_6 = (-0.0125) + + y_ftcin_l_1 = 0.0814929508598 + + y_ftcin_l_2 = 190.397828213 + + y_fxgap_1 = 1.29072367633 + + y_fxgap_2 = (-0.468009114875) + + y_fxgap_3 = (-0.0166666666667) + + y_fxgap_4 = 0.00416666666667 + + y_fxgap_5 = 0.00833333333333 + + y_fxgap_6 = 0.0125 + + y_fxgap_7 = 0.0125 + + y_fxgap_8 = (-0.0166666666667) + + y_fxgap_9 = 0.00833333333333 + + y_fxgap_10 = (-0.0166666666667) + + y_fxgap_11 = 0.00416666666667 + + y_fxgap_12 = 0.05 + + y_fxgap_13 = 0.0373455901902 + + y_fynicn_l_1 = 0.000868642945186 + + y_fynicn_l_2 = 0.203972136271 + + y_fyniln_l_1 = 0.00114819592586 + + y_fyniln_l_2 = 0.344642504397 + + y_gfdbtnp_l_1 = 6.19935005084e-05 + + y_gfdbtnp_l_2 = 0.984645217482 + + y_gfdbtnp_l_3 = 0.0737924242446 + + y_gfdbtnp_l_4 = (-0.0584376417269) + + y_gfdbtn_l_1 = 5.5810037311e-05 + + y_ugfsrp_1 = 0.947688 + + y_uleg_l_1 = (-0.0162972181781) + + y_uleg_l_2 = 0.0162972181781 + + y_uleg_l_3 = 0.1 + + y_gfexpn_l_1 = 0.000210646994344 + + y_gfexpn_l_2 = 0.0964148144871 + + y_gfexpn_l_3 = 0.19363872408 + + y_gfexpn_l_4 = 0.600944699108 + + y_gfexpn_l_5 = 0.109001762325 + + y_gfintn_l_1 = 0.00193250998745 + + y_gfintn_l_2 = 34.038852147 + + y_gfrecn_l_1 = 0.000265992685534 + + y_gfrecn_l_2 = 0.5764571204 + + y_gfrecn_l_3 = 0.0743675317358 + + y_gfrecn_l_4 = 5.57251231588 + + y_gfrecn_l_5 = 0.349175347864 + + y_gtn_l_1 = 0.000350526420578 + + y_gtr_l_1 = 0.000390220355331 + + y_gtr_l_2 = 7.39501037898 + + y_gtr_l_3 = 7.39501037898 + + y_gtrd_1 = (-0.000176387604876) + + y_gtrd_2 = (-0.000206546235356) + + y_gtrd_3 = (-4.93246174231e-05) + + y_gtrd_4 = (-4.93246174231e-05) + + y_gtrd_5 = (-4.93246174231e-05) + + y_gtrd_6 = 0.862481931486 + + y_gtrd_7 = 0.000309352740077 + + y_hgemp_1 = 0.9 + + y_hgemp_2 = 40 + + y_hgemp_3 = (-40) + + y_hggdp_1 = 400 + + y_hggdp_2 = (-400) + + y_hgpbfir_1 = 0.975 + + y_hgpbfir_2 = 10 + + y_hgpbfir_3 = 10 + + y_hgpbfir_4 = (-10) + + y_hgpbfir_5 = (-10) + + y_hgpbfir_6 = (-10) + + y_hgpbfir_7 = 10 + + y_hgpkir_1 = 0.9 + + y_hgpkir_2 = 43.1298484247 + + y_hgpkir_3 = (-43.0591386594) + + y_hgynid_1 = 454.348916939 + + y_hgynid_2 = (-54.3489169394) + + y_hgynid_3 = (-400) + + y_hgynid_4 = (-455.23665293) + + y_hgynid_5 = 55.2366529304 + + y_hgynid_6 = 400 + + y_hks_1 = 384.31948476 + + y_hks_2 = (-384.31948476) + + y_hks_3 = 15.68051524 + + y_hks_4 = (-15.68051524) + + y_hlept_1 = 400 + + y_hlept_2 = 400 + + y_hlept_3 = (-400) + + y_hmfpt_1 = 0.95 + + y_hqlfpr_1 = 0.95 + + y_hqlww_1 = 0.95 + + y_huqpct_1 = 0.95 + + y_huxb_1 = 0.324768405324 + + y_huxb_2 = 0.95 + + y_hxbt_1 = 0.725 + + y_hxbt_2 = 0.725 + + y_hxbt_3 = 290 + + y_hxbt_4 = (-290) + + y_hxbt_5 = 0.275 + + y_jccan_l_1 = 6.24582838478 + + y_jccan_l_2 = 0.82051735145 + + y_jccan_l_3 = (-0.948637916333) + + y_jccan_l_4 = 0.121328058188 + + y_jccan_l_5 = 0.128120564883 + + y_jccan_l_6 = 1.35223326447 + + y_jccan_l_7 = 0.128120564883 + + y_jkcd_l_1 = 0.000730359646 + + y_jkcd_l_2 = 4.66817353822 + + y_kbfi_l_1 = 4.49018358914e-05 + + y_kbfi_l_2 = 0.0281084105505 + + y_kbfi_l_3 = (-0.0265200751536) + + y_kbfi_l_4 = 0.0281084105505 + + y_kbfi_l_5 = (-0.248867790412) + + y_kbfi_l_6 = 0.971891589449 + + y_kcd_l_1 = 0.000154373410789 + + y_kcd_l_2 = 0.066147038262 + + y_kcd_l_3 = (-0.246673633735) + + y_kcd_l_4 = 0.933852961738 + + y_kh_l_1 = 5.72922867013e-05 + + y_kh_l_2 = 0.00873032740269 + + y_kh_l_3 = (-0.249249311699) + + y_kh_l_4 = 0.991269672597 + + y_ki_l_1 = 1.44204786648 + + y_ki_l_2 = 0.014692062549 + + y_ki_l_3 = 0.250723990347 + + y_ki_l_4 = (-0.456739929026) + + y_ki_l_5 = 0.0711962153783 + + y_ki_l_6 = (-0.307228143176) + + y_ks_l_1 = 0.0025 + + y_leg_l_1 = 0.248485878175 + + y_leg_l_2 = 0.751514121825 + + y_leh_l_1 = 0.00641807663415 + + y_leh_l_2 = 0.813979789462 + + y_leh_l_3 = 0.132451786431 + + y_leh_l_4 = 0.0535684241064 + + y_leo_l_1 = 20.7652726744 + + y_leo_l_2 = 0.756667597034 + + y_leo_l_3 = (-15.6501866511) + + y_leo_l_4 = (-0.756667597034) + + y_leo_l_5 = (-0.0164258334824) + + y_lep_l_1 = 0.00788481079904 + + y_leppot_l_1 = 0.0079687353701 + + y_leppot_l_2 = (-0.0110028694424) + + y_leppot_l_3 = (-1.10028694424) + + y_leppot_l_4 = (-0.857254870696) + + y_lf_l_1 = 0.00617558165686 + + y_lf_l_2 = 1.58659431972 + + y_lfpr_1 = 0.432392517171 + + y_lfpr_2 = 0.567607482829 + + y_lfpr_3 = (-0.000875189202097) + + y_lfpr_4 = 0.000875189202097 + + y_lhp_l_1 = 1.00059088506 + + y_lhp_l_2 = 0.202289789801 + + y_lhp_l_3 = (-0.202880674857) + + y_lhp_l_4 = 0.372064184885 + + y_lhp_l_5 = 0.627935815115 + + y_lhp_l_6 = (-0.755331857052) + + y_lhp_l_7 = (-0.00156983953779) + + y_lhp_l_8 = 0.127396041937 + + y_lhp_l_9 = 0.000318490104843 + + y_lur_1 = (-96.2208093896) + + y_lur_2 = 96.2208093896 + + y_lurnat_1 = 0.95 + + y_lww_l_1 = 0.804289649347 + + y_lww_l_2 = 0.00170379588201 + + y_lww_l_3 = 0.195710350653 + + y_lww_l_4 = 0.318481647196 + + y_lww_l_5 = (-0.318481647196) + + y_lww_l_6 = (-0.00079620411799) + + y_mfpt_l_1 = 0.0025 + + y_pbfir_l_1 = 0.0025 + + y_pcdr_l_1 = 1.50984819434 + + y_pcdr_l_2 = (-0.509848194342) + + y_pcer_l_1 = 0.248860953365 + + y_pcer_l_2 = (-0.248860953365) + + y_pcer_l_3 = (-0.248860953365) + + y_pcer_l_4 = 0.248860953365 + + y_pcfr_l_1 = 1.21019336782 + + y_pcfr_l_2 = (-0.14928038046) + + y_pcfr_l_3 = (-0.365198296745) + + y_pcfr_l_4 = 0.318574001625 + + y_pcfr_l_5 = (-0.338884189342) + + y_pcfr_l_6 = 0.333798755712 + + y_pchr_l_1 = 1.59806398567 + + y_pchr_l_2 = (-0.598063985667) + + y_pcnia_l_1 = 0.0025 + + y_pcor_l_1 = (-0.1436) + + y_pcor_l_2 = 0.1436 + + y_pcor_l_3 = (-0.217) + + y_pcor_l_4 = 0.217 + + y_pcpi_l_1 = 0.00394679077503 + + y_pcpi_l_2 = 0.43067430272 + + y_pcpix_l_1 = 0.00384060295377 + + y_pcpix_l_2 = 0.426412064374 + + y_pcxfe_l_1 = 0.0025 + + y_pegfr_l_1 = 0.0025 + + y_pegsr_l_1 = 0.0025 + + y_pgdp_l_1 = 0.00898451406694 + + y_pgfl_l_1 = 0.525153490957 + + y_pgsl_l_1 = 0.514419453205 + + y_phouse_l_1 = 1.89031776892 + + y_phouse_l_2 = (-0.901886995515) + + y_phouse_l_3 = 0.0115692265899 + + y_phouse_l_4 = 0.0115692265899 + + y_phr_l_1 = 0.0025 + + y_pic4_1 = 100 + + y_pic4_2 = (-100) + + y_picnia_1 = 15.96 + + y_picnia_2 = (-15.96) + + y_picnia_3 = 29.04 + + y_picnia_4 = (-29.04) + + y_picx4_1 = 100 + + y_picx4_2 = (-100) + + y_picxfe_1 = 0.404860664116 + + y_picxfe_2 = 0.591171818183 + + y_picxfe_3 = 0.00396751770099 + + y_picxfe_4 = 0.462045372577 + + y_picxfe_5 = (-0.462045372577) + + y_pieci_1 = 0.00293156716662 + + y_pieci_2 = 0.00293156716662 + + y_pieci_3 = 0.00293156716662 + + y_pieci_4 = 0.146578358331 + + y_pieci_5 = 0.839226144659 + + y_pieci_6 = 0.00540079551024 + + y_pieci_7 = 0.00540079551024 + + y_pieci_8 = (-2.16031820409) + + y_pieci_9 = (-0.0143209721548) + + y_pieci_10 = 0.0143209721548 + + y_pieci_11 = 0.327959270689 + + y_pieci_12 = (-0.327959270689) + + y_pigdp_1 = 400 + + y_pigdp_2 = (-400) + + y_pipxnc_1 = (-796) + + y_pipxnc_2 = 0.462801 + + y_pipxnc_3 = (-0.462801) + + y_pipxnc_4 = 368.389596 + + y_pipxnc_5 = 0.229745 + + y_pipxnc_6 = (-0.229745) + + y_pipxnc_7 = 182.87702 + + y_pipxnc_8 = (-14.9334031956) + + y_pipxnc_9 = 14.9334031956 + + y_pipxnc_10 = 10 + + y_pipxnc_11 = (-10) + + y_pkbfir_1 = 0.960531663984 + + y_pkbfir_2 = 1.05983283594 + + y_pl_l_1 = 0.0025 + + y_pmo_l_1 = 0.622318401629 + + y_pmo_l_2 = 0.377681598371 + + y_pmo_l_3 = 0.00731956262431 + + y_pmo_l_4 = (-0.00731956262431) + + y_pmo_l_5 = (-0.629637964254) + + y_pmo_l_6 = 0.234396660333 + + y_pmo_l_7 = (-0.234396660333) + + y_pmo_l_8 = 0.765603339667 + + y_pmp_l_1 = 0.0171179383155 + + y_pmp_l_2 = 1.05645668526 + + y_poil_l_1 = 0.0162031615251 + + y_poilr_l_1 = 1.17135063067 + + y_poilr_l_2 = (-0.346197996438) + + y_poilr_l_3 = (-0.390345197801) + + y_poilr_l_4 = 0.79951907837 + + y_ptr_1 = 0.9 + + y_ptr_2 = 0.05 + + y_ptr_3 = 0.05 + + y_pxb_l_1 = 0.00914375584343 + + y_pxb_l_2 = 1.01772402773 + + y_pxnc_l_1 = 0.0025 + + y_pxp_l_1 = 0.6469 + + y_pxp_l_2 = (-0.6469) + + y_pxp_l_3 = 0.3531 + + y_pxp_l_4 = (-0.3531) + + y_pxr_l_1 = 0.0025 + + y_qebfi_l_1 = 0.000358162570912 + + y_qebfi_l_2 = 0.664481948351 + + y_qebfi_l_3 = 0.0787039173848 + + y_qebfi_l_4 = (-0.0787039173848) + + y_qebfi_l_5 = 7.87039173848 + + y_qec_l_1 = 0.935665935123 + + y_qec_l_2 = 0.0166517759473 + + y_qec_l_3 = (-0.139711201786) + + y_qec_l_4 = 0.135400942735 + + y_qec_l_5 = 0.0519925479811 + + y_qecd_l_1 = 0.000593792074211 + + y_qecd_l_2 = 3.98656310426 + + y_qecd_l_3 = 0.00498320388032 + + y_qecd_l_4 = 0.00498320388032 + + y_qecd_l_5 = 0.00498320388032 + + y_qecd_l_6 = 0.00498320388032 + + y_qecd_l_7 = 0.00498320388032 + + y_qecd_l_8 = 0.00498320388032 + + y_qecd_l_9 = 0.00498320388032 + + y_qecd_l_10 = 0.00498320388032 + + y_qecd_l_11 = (-0.0232956396718) + + y_qecd_l_12 = (-0.584353967629) + + y_qecd_l_13 = (-0.0242284661483) + + y_qeh_l_1 = 0.0015504186377 + + y_qeh_l_2 = 24.6010652056 + + y_qeh_l_3 = 0.0153756657535 + + y_qeh_l_4 = 0.0153756657535 + + y_qeh_l_5 = 0.0153756657535 + + y_qeh_l_6 = 0.0153756657535 + + y_qeh_l_7 = 0.0153756657535 + + y_qeh_l_8 = 0.0153756657535 + + y_qeh_l_9 = 0.0153756657535 + + y_qeh_l_10 = 0.0153756657535 + + y_qeh_l_11 = 0.0153756657535 + + y_qeh_l_12 = 0.0153756657535 + + y_qeh_l_13 = 0.0153756657535 + + y_qeh_l_14 = 0.0153756657535 + + y_qeh_l_15 = 0.0153756657535 + + y_qeh_l_16 = 0.0153756657535 + + y_qeh_l_17 = 0.0153756657535 + + y_qeh_l_18 = 0.0153756657535 + + y_qeh_l_19 = (-0.0270350700995) + + y_qkir_l_1 = 0.00188536673771 + + y_qlf_l_1 = 0.00620858571308 + + y_qlf_l_2 = 1.58692282562 + + y_qlhp_l_1 = 0.00465728156706 + + y_qlww_l_1 = 0.0025 + + y_qpxnc_l_1 = 2.98507462687 + + y_qpxnc_l_2 = (-2.98507462687) + + y_qpxnc_l_3 = (-1.98507462687) + + y_qpxnc_l_4 = 1.98507462687 + + y_qpxp_l_1 = 0.7195976338 + + y_qpxp_l_2 = (-0.7195976338) + + y_qynidn_l_1 = 0.354822592523 + + y_qynidn_l_2 = 1.13587229235 + + y_qynidn_l_3 = (-0.135872292349) + + y_rbbbp_1 = (-0.189051) + + y_rbbbp_2 = 0.848879 + + y_rbbbp_3 = 0.160481423829 + + y_rbfi_1 = (-2.21124682364) + + y_rbfi_2 = 0.395 + + y_rbfi_3 = 0.395 + + y_rbfi_4 = (-0.395) + + y_rbfi_5 = (-0.5) + + y_rbfi_6 = 0.5 + + y_rcar_1 = 1.22665328945 + + y_rcar_2 = 0 + + y_rcar_3 = 0.696748171914 + + y_rcar_4 = 0.101669335039 + + y_rcar_5 = 0.201582493047 + + y_rccd_1 = 100 + + y_rcch_1 = 100 + + y_rcch_2 = (-0.0545840410668) + + y_rcch_3 = 0.7953 + + y_rcch_4 = 79.53 + + y_rcgain_1 = 0.32854362351 + + y_rcgain_2 = 0.225785775119 + + y_rcgain_3 = (-0.225785775119) + + y_reqp_1 = 0.808086 + + y_reqp_2 = 0.795819 + + y_reqp_3 = (-0.643090192434) + + y_rfynic_1 = 1.00400815341 + + y_rfynic_2 = (-0.49108746803) + + y_rfynic_3 = (-0.144424360986) + + y_rfynic_4 = 0.631503675605 + + y_rfynil_1 = 0.884413145643 + + y_rfynil_2 = (-0.00726474303036) + + y_rfynil_3 = (-0.171195169347) + + y_rfynil_4 = 0.0265702779079 + + y_rfynil_5 = (-0.132818819092) + + y_rfynil_6 = 0.0876033907073 + + y_rfynil_7 = 0.261434600384 + + y_rfynil_8 = 0.0179349568622 + + y_rg10p_1 = (-0.460658806872) + + y_rg10p_2 = 0.228721864424 + + y_rg10p_3 = 0.920104088065 + + y_rg10p_4 = 0.423854051406 + + y_rg10p_5 = (-0.210447922486) + + y_rg30p_1 = (-0.624829467707) + + y_rg30p_2 = 0.134994250522 + + y_rg30p_3 = 0.938108605708 + + y_rg30p_4 = 0.586157900756 + + y_rg30p_5 = (-0.126639268136) + + y_rg5p_1 = (-0.349564481) + + y_rg5p_2 = 0.90221329312 + + y_rg5p_3 = 0.315381721561 + + y_rgfint_1 = 0.845677566688 + + y_rgfint_2 = 0.154322433312 + + y_rgfint_3 = 0.00556931000493 + + y_rgfint_4 = (-0.00556931000493) + + y_rgw_1 = 0.00495 + + y_rgw_2 = 0.00271 + + y_rgw_3 = 0.00129 + + y_rgw_4 = 0.00105 + + y_rme_1 = 0.660306961037 + + y_rme_2 = 0.884200704474 + + y_rme_3 = (-0.544507665511) + + y_rme_4 = (-0.102549417082) + + y_rrff_1 = (-0.25) + + y_rrff_2 = (-0.25) + + y_rrff_3 = (-0.25) + + y_rrff_4 = (-0.25) + + y_rrtr_1 = 0.97 + + y_rrtr_2 = 0.03 + + y_rspnia_1 = 7.62633280279 + + y_rspnia_2 = (-7.62633280279) + + y_rtb_1 = 0.799718792152 + + y_rtb_2 = 0.11137355158 + + y_rtb_3 = 0.770122562667 + + y_rtb_4 = (-0.681214906399) + + y_rtbfi_l_1 = 5.40790262847 + + y_rtbfi_l_2 = 0.0576949599629 + + y_rtbfi_l_3 = 5.76949599629 + + y_rtbfi_l_4 = (-0.0576949599629) + + y_rtbfi_l_5 = (-0.0123862111793) + + y_rtbfi_l_6 = 0.129531747065 + + y_rtbfi_l_7 = 0 + + y_rtbfi_l_8 = (-0.260110434765) + + y_rtbfi_l_9 = 0.943576128374 + + y_rtinv_1 = 0.00912489842966 + + y_rtinv_2 = (-0.00912489842966) + + y_rtinv_3 = 0.0330561789534 + + y_rtinv_4 = 0.0356516398072 + + y_rtinv_5 = 0.0329826447805 + + y_rtinv_6 = 0.0355066663735 + + y_rtinv_7 = (-0.066038823734) + + y_tcin_l_1 = 0.00357673139507 + + y_tcin_l_2 = 8.35657418879 + + y_tpn_l_1 = 0.000461426663182 + + y_tpn_l_2 = 7.01937651683 + + y_tpn_l_3 = 1.18756132659 + + y_tpn_l_4 = (-0.187561326587) + + y_trci_1 = 0.00706626139452 + + y_trci_2 = 0.810247648208 + + y_trci_3 = (-0.810247648208) + + y_trci_4 = (-0.00572542167653) + + y_trp_1 = 0.603942358608 + + y_trp_2 = (-0.603942358608) + + y_trp_3 = 0.236576213581 + + y_trp_4 = (-0.236576213581) + + y_trp_5 = 0.000630587773923 + + y_trptd_1 = 0.420215062775 + + y_trptd_2 = (-0.420215062775) + + y_trptd_3 = (-0.55) + + y_trptd_4 = (-0.422749789232) + + y_trptd_5 = 0.422749789232 + + y_trptd_6 = (-0.5) + + y_trpts_1 = (-0.0180202713644) + + y_trpts_2 = 0.0225987818683 + + y_trpts_3 = (-0.00457851050393) + + y_trpts_4 = 0.1 + + y_trpts_5 = 0.00075 + + y_tryh_1 = 0.144437010525 + + y_tryh_2 = (-0.0944218552605) + + y_tryh_3 = (-0.0500151552646) + + y_uxbt_l_1 = 0.0025 + + y_uynicpnr_1 = 0.779183 + + y_vbfi_1 = 5.96826486935 + + y_vbfi_2 = 1.41987523928 + + y_vbfi_3 = (-1.50480877253) + + y_vbfi_4 = (-1.50480877253) + + y_wpo_l_1 = 1.36629240684e-05 + + y_wpon_l_1 = 1.2518399289e-05 + + y_wpon_l_2 = 0.99460869287 + + y_wpon_l_3 = 0.00146536714744 + + y_wpon_l_4 = 0.408461833894 + + y_wpon_l_5 = (-0.408461833894) + + y_wpon_l_6 = 0.0498372673814 + + y_wpon_l_7 = (-0.0443498822123) + + y_wpon_l_8 = (-0.00103486242602) + + y_wpon_l_9 = 0.000938784388547 + + y_wpon_l_10 = 0.000938784388547 + + y_wpon_l_11 = 0.00466659001196 + + y_wpon_l_12 = (-0.00372780562341) + + y_wps_l_1 = 3.55473075437e-05 + + y_wpsn_l_1 = 1.13587229235 + + y_wpsn_l_2 = (-0.135872292349) + + y_wpsn_l_3 = (-0.25) + + y_wpsn_l_4 = 0.25 + + y_xb_l_1 = 6.84878469814e-05 + + y_xb_l_2 = 1.0 + + y_xb_l_3 = (-1.0) + + y_xbn_l_1 = 6.26139265285e-05 + + y_xbn_l_2 = 1.0198018271 + + y_xbn_l_3 = 1.0198018271 + + y_xbn_l_4 = 1.31175365227 + + y_xbn_l_5 = (-1.33155547937) + + y_xbn_l_6 = (-1.33155547937) + + y_xbo_l_1 = 0.0132470548943 + + y_xbt_l_1 = 0.725 + + y_xbt_l_2 = 0.725 + + y_xbt_l_3 = 0.725 + + y_xbt_l_4 = 0.275 + + y_xbtr_l_1 = 0.95 + + y_xfs_l_1 = 0.6849 + + y_xfs_l_2 = (-0.6849) + + y_xfs_l_3 = 0.0386 + + y_xfs_l_4 = (-0.0386) + + y_xfs_l_5 = 0.1324 + + y_xfs_l_6 = (-0.1324) + + y_xfs_l_7 = 0.0429 + + y_xfs_l_8 = (-0.0429) + + y_xfs_l_9 = 0.0223 + + y_xfs_l_10 = (-0.0223) + + y_xfs_l_11 = 0.0395 + + y_xfs_l_12 = (-0.0395) + + y_xfs_l_13 = 0.0691 + + y_xfs_l_14 = (-0.0691) + + y_xfs_l_15 = 0.1203 + + y_xfs_l_16 = (-0.1203) + + y_xfs_l_17 = (-0.1399) + + y_xfs_l_18 = 0.1399 + + y_xfs_l_19 = (-0.0101) + + y_xfs_l_20 = 0.0101 + + y_xfsn_l_1 = 4.78939925834e-05 + + y_xfsn_l_2 = 1.00337294235 + + y_xfsn_l_3 = (-0.00363305240167) + + y_xfsn_l_4 = (-0.00337294235067) + + y_xfsn_l_5 = (-0.544852165261) + + y_xfsn_l_6 = 0.541479222911 + + y_xgap_1 = 100 + + y_xgap_2 = (-100) + + y_xgap2_1 = 100 + + y_xgap2_2 = (-100) + + y_xgdi_l_1 = 5.28881809334e-05 + + y_xgdin_l_1 = 4.75117496854e-05 + + y_xgdo_l_1 = 5.23439257193e-05 + + y_xgdp_l_1 = 0.9985 + + y_xgdp_l_2 = (-0.9985) + + y_xgdp_l_3 = 0.6264 + + y_xgdp_l_4 = (-1.2513) + + y_xgdp_l_5 = 0.6249 + + y_xgdpn_l_1 = 4.77329919533e-05 + + y_xgdpn_l_2 = 1.0564013312 + + y_xgdpn_l_3 = 0.021847772281 + + y_xgdpn_l_4 = 0.0674519622926 + + y_xgdpn_l_5 = (-0.149062669624) + + y_xgdpn_l_6 = 0.00362083951871 + + y_xgdpn_l_7 = 0.00336160385466 + + y_xgdpn_l_8 = 0.543020588122 + + y_xgdpn_l_9 = (-0.539658984268) + + y_xgdptn_l_1 = 4.74032737815e-05 + + y_xp_l_1 = 0.6526679404 + + y_xp_l_2 = (-0.6526679404) + + y_xp_l_3 = 0.0361108836 + + y_xp_l_4 = (-0.0361108836) + + y_xp_l_5 = 0.11825695358 + + y_xp_l_6 = (-0.11825695358) + + y_xp_l_7 = 0.04216893278 + + y_xp_l_8 = (-0.04216893278) + + y_xp_l_9 = 0.0365822346 + + y_xp_l_10 = (-0.0365822346) + + y_xp_l_11 = 0.114213055 + + y_xp_l_12 = (-0.114213055) + + y_xpn_l_1 = 4.51845246156e-05 + + y_ydn_l_1 = 6.27963768217e-05 + + y_ydn_l_2 = 0.998336445483 + + y_ydn_l_3 = 1.13631857158 + + y_ydn_l_4 = (-0.136318571582) + + y_yh_l_1 = 6.84845844852e-05 + + y_yh_l_2 = 0.526533658207 + + y_yh_l_3 = 0.178799300517 + + y_yh_l_4 = 0.294667041275 + + y_yhgap_1 = 100 + + y_yhgap_2 = (-100) + + y_yhibn_l_1 = 0.00302416992206 + + y_yhibn_l_2 = 63.3559682371 + + y_yhl_l_1 = 0.000130066869264 + + y_yhl_l_2 = (-1.16884651007) + + y_yhln_l_1 = 0.00010196408965 + + y_yhln_l_2 = 1.14236648073 + + y_yhp_l_1 = 0.000232413452786 + + y_yhp_l_2 = (-1.10635973458) + + y_yhp_l_3 = 0.946560228789 + + y_yhp_l_4 = 0.0534397712107 + + y_yhpgap_1 = 100 + + y_yhpgap_2 = (-100) + + y_yhpntn_l_1 = 0.00398579631637 + + y_yhpntn_l_2 = 1.19125575899 + + y_yhpntn_l_3 = 1.19125575899 + + y_yhpntn_l_4 = 1.19125575899 + + y_yhpntn_l_5 = (-1.31798027859) + + y_yhpntn_l_6 = 9.31230191482 + + y_yhpntn_l_7 = (-1.11436836489) + + y_yhpntn_l_8 = (-5.43026982685) + + y_yhpntn_l_9 = (-0.71417194978) + + y_yhpntn_l_10 = (-1.64093920349) + + y_yhpshr_l_1 = 3.39337481391 + + y_yhptn_l_1 = 0.000192500076699 + + y_yhptn_l_2 = 0.975134043217 + + y_yhptn_l_3 = 0.563603989214 + + y_yhptn_l_4 = 0.102159206272 + + y_yhptn_l_5 = 0.268955804126 + + y_yhptn_l_6 = 0.0652810003888 + + y_yhshr_l_1 = 1.28884347188 + + y_yhshr_l_2 = 0.99999334858 + + y_yhshr_l_3 = (-0.99999334858) + + y_yhsn_l_1 = 0.000823471829286 + + y_yhsn_l_2 = 8.07609651707 + + y_yhsn_l_3 = 2.34637489687 + + y_yhsn_l_4 = 4.27777403212 + + y_yhsn_l_5 = (-1.78462125185) + + y_yhsn_l_6 = (-11.6695043163) + + y_yhsn_l_7 = (-0.272296812186) + + y_yhsn_l_8 = 17.3716235947 + + y_yhsn_l_9 = 0.0261769342453 + + y_yht_l_1 = 0.000383024901591 + + y_yhtgap_1 = 100 + + y_yhtgap_2 = (-100) + + y_yhtn_l_1 = 0.00035095492642 + + y_yhtn_l_2 = 1.00122246375 + + y_yhtshr_l_1 = 5.59327644091 + + y_ykbfin_l_1 = 0.000222515379076 + + y_ykbfin_l_2 = 0.501135005995 + + y_ykbfin_l_3 = 0.498864994005 + + y_ykin_l_1 = 0.0048752071053 + + y_ykin_l_2 = 15.1326363291 + + y_ykin_l_3 = 0.501553881106 + + y_ykin_l_4 = 0.498446118894 + + y_ynicpn_l_1 = 0.000428014077811 + + y_ynicpn_l_2 = 7.69278337234 + + y_ynicpn_l_3 = (-4.79530526339) + + y_ynicpn_l_4 = (-1.22202005096) + + y_ynicpn_l_5 = 8.96683950232 + + y_ynicpn_l_6 = (-0.675458057991) + + y_ynidn_l_1 = 0.000734588128108 + + y_ynidn_l_2 = 0.683167062078 + + y_ynidn_l_3 = (-0.000507771585891) + + y_ynidn_l_4 = (-0.790436568589) + + y_ynidn_l_5 = 0.107269506511 + + y_ynidn_l_6 = 0.209563431411 + + y_ynidn_l_7 = (-0.000157528146717) + + y_ynidn_l_8 = (-0.209563431411) + + y_yniln_l_1 = 8.92568990506e-05 + + y_yniln_l_2 = 0.977159070276 + + y_yniln_l_3 = 0.829114291162 + + y_yniln_l_4 = 0.829114291162 + + y_yniln_l_5 = 0.0418084260005 + + y_yniln_l_6 = 0.0418084260005 + + y_yniln_l_7 = 0.129077282837 + + y_yniln_l_8 = 0.129077282837 + + y_ynin_l_1 = 5.56383895262e-05 + + y_ynin_l_2 = 0.999999935846 + + y_ynin_l_3 = 1.17104491968 + + y_ynin_l_4 = 0.0640520865577 + + y_ynin_l_5 = (-0.0484572291581) + + y_ynin_l_6 = (-0.186639777078) + + y_ynirn_l_1 = 7.33377354801 + + y_ynirn_l_2 = 0.951263114856 + + y_ynirn_l_3 = (-0.951263114856) + + y_ynirn_l_4 = 0.00548690935542 + + y_ynirn_l_5 = (-0.00548690935542) + + y_ypn_l_1 = 5.53550901718e-05 + + y_ypn_l_2 = 0.988173952374 + + y_ypn_l_3 = 0.549385156514 + + y_ypn_l_4 = 0.159614656674 + + y_ypn_l_5 = 0.291000186812 + + y_zdivgr_1 = 0.00975726425743 + + y_zdivgr_2 = 0.990242735743 + + y_zebfi_1 = (-0.000431144211955) + + y_zebfi_2 = (-0.00050714173603) + + y_zebfi_3 = (-3.88181916088e-05) + + y_zebfi_4 = 0.00016798757544 + + y_zebfi_5 = (-0.000975251482943) + + y_zebfi_6 = 0.000417269685018 + + y_zebfi_7 = 9.80402248148e-06 + + y_zebfi_8 = 0.00040254489385 + + y_zebfi_9 = 0.000145632881593 + + y_zebfi_10 = 0.000809116564154 + + y_zebfi_11 = 0.000691481740712 + + y_zebfi_12 = (-0.00152462990113) + + y_zebfi_13 = 0.000182102122415 + + y_zebfi_14 = 0.000170960242897 + + y_zebfi_15 = 0.0142945657655 + + y_zebfi_16 = (-0.00425222899975) + + y_zebfi_17 = (-0.00503049733108) + + y_zebfi_18 = (-0.00112440248315) + + y_zebfi_19 = (-0.0038874369515) + + y_zebfi_20 = 0.00035570453849 + + y_zebfi_21 = (-0.00035570453849) + + y_zecd_1 = (-0.000424433044911) + + y_zecd_2 = (-0.000566112732916) + + y_zecd_3 = (-0.000427835415485) + + y_zecd_4 = 4.27545061866e-06 + + y_zecd_5 = (-0.00133363746841) + + y_zecd_6 = 0.00178510275432 + + y_zecd_7 = (-0.000271474405975) + + y_zecd_8 = 0.000459611864377 + + y_zecd_9 = 0.000428608849069 + + y_zecd_10 = (-0.00111248088805) + + y_zecd_11 = 3.61133130939e-05 + + y_zecd_12 = 7.97590705793e-05 + + y_zecd_13 = 0.00141410574269 + + y_zecd_14 = (-0.000639602744318) + + y_zecd_15 = (-0.00010841426451) + + y_zecd_16 = 0.000210363124201 + + y_zecd_17 = 0.000178061664134 + + y_zecd_18 = 0.000146912749167 + + y_zecd_19 = (-0.000139880426754) + + y_zecd_20 = (-3.38007573296e-05) + + y_zecd_21 = 0.000166975793706 + + y_zecd_22 = 0.000113506936821 + + y_zecd_23 = 0.000124123127885 + + y_zecd_24 = (-0.000203591971486) + + y_zecd_25 = 5.7989188193e-05 + + y_zecd_26 = 0.000114280775871 + + y_zecd_27 = 0.00255088738447 + + y_zecd_28 = (-0.001880611807) + + y_zecd_29 = 0.0308598105755 + + y_zecd_30 = (-0.00201324622316) + + y_zecd_31 = (-0.0365513581269) + + y_zecd_32 = (-0.00465896135484) + + y_zecd_33 = 0.0123637551294 + + y_zeco_1 = (-7.52202049496e-05) + + y_zeco_2 = (-7.94406933181e-05) + + y_zeco_3 = (-2.05931699614e-05) + + y_zeco_4 = 0.000100439779498 + + y_zeco_5 = 2.12832185698e-05 + + y_zeco_6 = 1.70353153588e-05 + + y_zeco_7 = 5.5012376381e-05 + + y_zeco_8 = 3.68085672111e-05 + + y_zeco_9 = (-0.000630171036922) + + y_zeco_10 = 0.000273875586514 + + y_zeco_11 = 0.000133019756131 + + y_zeco_12 = (-3.46619140531e-05) + + y_zeco_13 = 7.48142887307e-05 + + y_zeco_14 = (-0.000130139477521) + + y_zeco_15 = (-0.000574849476126) + + y_zeco_16 = 0.000315791553755 + + y_zeco_17 = 0.000397005436297 + + y_zeco_18 = 2.60636593368e-05 + + y_zeco_19 = (-6.06591388527e-05) + + y_zeco_20 = (-5.86151697491e-06) + + y_zeco_21 = 4.60869299242e-05 + + y_zeco_22 = (-3.67533909379e-05) + + y_zeco_23 = 0.000205501772024 + + y_zeco_24 = (-0.000240937714399) + + y_zeco_25 = (-0.000131812287659) + + y_zeco_26 = (-8.99812036284e-05) + + y_zeco_27 = 0.0011735331967 + + y_zeco_28 = 0.0732239724725 + + y_zeco_29 = (-0.0439002248803) + + y_zeco_30 = (-0.0221752554553) + + y_zeco_31 = 0.00192935493602 + + y_zeco_32 = (-0.00907784707292) + + y_zeh_1 = (-7.82590644729e-05) + + y_zeh_2 = (-5.22177650099e-05) + + y_zeh_3 = (-5.11617748264e-05) + + y_zeh_4 = 1.06118801985e-05 + + y_zeh_5 = 0.00015247021427 + + y_zeh_6 = 0.000115461916411 + + y_zeh_7 = 6.94775905299e-05 + + y_zeh_8 = 6.915284272e-06 + + y_zeh_9 = 0.000501695975955 + + y_zeh_10 = (-0.000531489463237) + + y_zeh_11 = (-0.000152716722883) + + y_zeh_12 = 4.42229725272e-05 + + y_zeh_13 = 0.000171026724111 + + y_zeh_14 = (-0.000344325005483) + + y_zeh_15 = 5.54746256755e-05 + + y_zeh_16 = 4.07823169105e-05 + + y_zeh_17 = (-3.07861940478e-05) + + y_zeh_18 = (-4.58464253978e-06) + + y_zeh_19 = 7.13641520669e-06 + + y_zeh_20 = 1.47073943648e-05 + + y_zeh_21 = 3.83084742308e-05 + + y_zeh_22 = 3.27196755109e-05 + + y_zeh_23 = (-6.99285491718e-05) + + y_zeh_24 = 4.55009661188e-07 + + y_zeh_25 = 5.50461376242e-05 + + y_zeh_26 = 2.86419262708e-05 + + y_zeh_27 = 0.00106943231558 + + y_zeh_28 = 0.00426630302385 + + y_zeh_29 = (-0.00573847474041) + + y_zeh_30 = (-0.00187609012218) + + y_zeh_31 = (-0.000659760712587) + + y_zeh_32 = 0.00400802255132 + + y_zgap05_1 = 0.0547936526434 + + y_zgap05_2 = 0.945206347357 + + y_zgap10_1 = 0.0300745581094 + + y_zgap10_2 = 0.969925441891 + + y_zgap30_1 = 0.014106588982 + + y_zgap30_2 = 0.985893411018 + + y_zgapc2_1 = (-0.0141848331986) + + y_zgapc2_2 = (-0.00438957847118) + + y_zgapc2_3 = (-0.00608986499063) + + y_zgapc2_4 = 0.00127453586676 + + y_zgapc2_5 = (-0.0426889990258) + + y_zgapc2_6 = 0.00775994605046 + + y_zgapc2_7 = 0.0191285668792 + + y_zgapc2_8 = (-0.00220795277592) + + y_zgapc2_9 = 0.194384536968 + + y_zgapc2_10 = (-0.0764007234264) + + y_zgapc2_11 = (-0.0113246023485) + + y_zgapc2_12 = (-0.0155518339662) + + y_zgapc2_13 = 0.0233897407937 + + y_zgapc2_14 = 0.018008438872 + + y_zlhp_1 = (-0.000202321377434) + + y_zlhp_2 = (-6.54709155556e-05) + + y_zlhp_3 = (-0.000172024683014) + + y_zlhp_4 = 3.13937564958e-05 + + y_zlhp_5 = (-0.00104747602255) + + y_zlhp_6 = 0.000259883906459 + + y_zlhp_7 = 0.00050790958016 + + y_zlhp_8 = (-4.19862687488e-05) + + y_zlhp_9 = 0.000321668804678 + + y_zlhp_10 = 0.000408423219508 + + y_zlhp_11 = (-0.00575496472936) + + y_zlhp_12 = 0.00581682769254 + + y_zlhp_13 = (-0.000203485929498) + + y_zlhp_14 = (-0.00027154348171) + + y_zlhp_15 = 0.685794690175 + + y_zlhp_16 = (-0.685794690175) + + y_zlhp_17 = (-0.685794690175) + + y_zlhp_18 = 0.685794690175 + + y_zlhp_19 = 0.000278310168582 + + y_zlhp_20 = 0.000278310168582 + + y_zpi10_1 = 0.0300745581094 + + y_zpi10_2 = 0.969925441891 + + y_zpi10f_1 = 0.0300745581094 + + y_zpi10f_2 = 0.969925441891 + + y_zpi5_1 = 0.0817876274963 + + y_zpi5_2 = 0.0221868418868 + + y_zpi5_3 = 0.0250194521826 + + y_zpi5_4 = (-9.00706244808e-05) + + y_zpi5_5 = (-0.145676547176) + + y_zpi5_6 = (-0.0311377360679) + + y_zpi5_7 = (-0.0294931929574) + + y_zpi5_8 = (-0.0275798582146) + + y_zpi5_9 = 0.233887334416 + + y_zpi5_10 = 0.871096149059 + + y_zpi5_11 = 0.174192252057 + + y_zpi5_12 = (-0.0718402312689) + + y_zpi5_13 = 0.0406637195158 + + y_zpi5_14 = 0.0449446239851 + + y_zpib5_1 = 21.9174610574 + + y_zpib5_2 = (-21.9174610574) + + y_zpib5_3 = 0.945206347357 + + y_zpic30_1 = 0.014106588982 + + y_zpic30_2 = 0.985893411018 + + y_zpicxfe_1 = 0.380818884672 + + y_zpicxfe_2 = 0.00113182715476 + + y_zpicxfe_3 = 0.00146351917605 + + y_zpicxfe_4 = 0.00225729733693 + + y_zpicxfe_5 = 0.0460967342223 + + y_zpicxfe_6 = 0.0338772671906 + + y_zpicxfe_7 = 0.0228924215171 + + y_zpicxfe_8 = 0.0112105032823 + + y_zpicxfe_9 = (-0.0140156100481) + + y_zpicxfe_10 = 0.0011222896601 + + y_zpicxfe_11 = 0.00760121840982 + + y_zpicxfe_12 = (-0.00299260406007) + + y_zpicxfe_13 = 0.0470383710002 + + y_zpicxfe_14 = (-0.0278318348119) + + y_zpicxfe_15 = (-0.00506170904133) + + y_zpicxfe_16 = (-0.00225028901719) + + y_zpicxfe_17 = 0.00828470603822 + + y_zpicxfe_18 = 0.500251545448 + + y_zpicxfe_19 = 11.937795061 + + y_zpicxfe_20 = (-11.937795061) + + y_zpicxfe_21 = 6.84395376806e-05 + + y_zpicxfe_22 = (-6.84395376806e-05) + + y_zpicxfe_23 = (-0.114076926212) + + y_zpicxfe_24 = 45.6307704848 + + y_zpicxfe_25 = (-0.00383816812034) + + y_zpicxfe_26 = 0.00383816812034 + + y_zpicxfe_27 = (-0.000695300677346) + + y_zpicxfe_28 = 0.000695300677346 + + y_zpieci_1 = (-0.026022539351) + + y_zpieci_2 = 0.00320414216918 + + y_zpieci_3 = 0.00402676215955 + + y_zpieci_4 = 0.00650489050087 + + y_zpieci_5 = 0.202430424141 + + y_zpieci_6 = 0.196252633802 + + y_zpieci_7 = 0.195837958296 + + y_zpieci_8 = 0.0246831983934 + + y_zpieci_9 = (-0.0328787076454) + + y_zpieci_10 = 0.00135903909754 + + y_zpieci_11 = 0.0229838005541 + + y_zpieci_12 = (-0.00862383586105) + + y_zpieci_13 = 0.148708914616 + + y_zpieci_14 = (-0.0777266665551) + + y_zpieci_15 = (-0.0137748704693) + + y_zpieci_16 = (-0.00648469451174) + + y_zpieci_17 = 0.0171597038548 + + y_zpieci_18 = 0.393082529889 + + y_zpieci_19 = (-4.49541220961) + + y_zpieci_20 = 4.49541220961 + + y_zpieci_21 = 0.000154587412169 + + y_zpieci_22 = (-0.000154587412169) + + y_zpieci_23 = 0.380795785368 + + y_zpieci_24 = (-152.318314147) + + y_zpieci_25 = (-0.0172443115476) + + y_zpieci_26 = 0.0172443115476 + + y_zpieci_27 = (-0.00416159167724) + + y_zpieci_28 = 0.00416159167724 + + y_zrff10_1 = 0.0300745581094 + + y_zrff10_2 = 0.969925441891 + + y_zrff30_1 = 0.014106588982 + + y_zrff30_2 = 0.985893411018 + + y_zrff5_1 = 0.0547936526434 + + y_zrff5_2 = 0.945206347357 + + y_zyh_l_1 = 8.4030342164e-05 + + y_zyh_l_2 = 0.000702312990074 + + y_zyh_l_3 = 0.00059883124856 + + y_zyh_l_4 = 0.000469068735294 + + y_zyh_l_5 = (-0.00211240097558) + + y_zyh_l_6 = 0.000165660225273 + + y_zyh_l_7 = (-0.000778039304358) + + y_zyh_l_8 = 0.000290681377414 + + y_zyh_l_9 = 0.000905762637463 + + y_zyh_l_10 = 0.00268742554252 + + y_zyh_l_11 = 0.000682359720547 + + y_zyh_l_12 = (-0.00125638870298) + + y_zyh_l_13 = (-0.00185424331619) + + y_zyh_l_14 = 0.00243409867739 + + y_zyh_l_15 = 0.00416166538503 + + y_zyh_l_16 = 0.000750155785081 + + y_zyh_l_17 = (-0.000192123546578) + + y_zyh_l_18 = (-0.000360282907619) + + y_zyhp_l_1 = 0.000863735995891 + + y_zyhp_l_2 = 0.00104399392499 + + y_zyhp_l_3 = 0.000941268897539 + + y_zyhp_l_4 = 0.000458916642717 + + y_zyhp_l_5 = (-0.00154443220559) + + y_zyhp_l_6 = (-0.000640772573525) + + y_zyhp_l_7 = (-0.000877216982815) + + y_zyhp_l_8 = (-4.0750285327e-05) + + y_zyhp_l_9 = (-0.00203693491032) + + y_zyhp_l_10 = 0.00307902287603 + + y_zyhp_l_11 = 0.00153367057932 + + y_zyhp_l_12 = (-0.000979443369808) + + y_zyhp_l_13 = (-0.00330791546136) + + y_zyhp_l_14 = 0.00310317204747 + + y_zyhp_l_15 = 0.00422069758369 + + y_zyhp_l_16 = 0.000130762149267 + + y_zyhp_l_17 = 5.34207087507e-05 + + y_zyhp_l_18 = (-0.00083494424337) + + y_zyhp_l_19 = 0.00225147864855 + + y_zyhp_l_20 = 0.000201405002604 + + y_zyhp_l_21 = (-0.000427498545256) + + y_zyhp_l_22 = (-0.000252078634627) + + y_zyhpst_l_1 = 0.0005 + + y_zyhst_l_1 = 0.0005 + + y_zyht_l_1 = (-0.000334830912493) + + y_zyht_l_2 = 0.000473468268016 + + y_zyht_l_3 = 0.000279807258553 + + y_zyht_l_4 = 0.000327960879431 + + y_zyht_l_5 = (-0.00250027976398) + + y_zyht_l_6 = 0.00088556815263 + + y_zyht_l_7 = (-0.00121488126811) + + y_zyht_l_8 = 8.52432970988e-05 + + y_zyht_l_9 = 0.00217284254807 + + y_zyht_l_10 = 0.00313028920932 + + y_zyht_l_11 = 0.00211420687194 + + y_zyht_l_12 = 0.000248144569686 + + y_zyht_l_13 = (-0.000746405493548) + + y_zyht_l_14 = 0.00274434958239 + + y_zyht_l_15 = 0.00270824568765 + + y_zyht_l_16 = 0.000679378874986 + + y_zyht_l_17 = 0.000308986169352 + + y_zyht_l_18 = (-6.93511660157e-05) + + y_zyht_l_19 = 0.00183760811329 + + y_zyht_l_20 = 0.000509649917682 + + y_zyht_l_21 = 9.47084299939e-05 + + y_zyht_l_22 = 0.0002426388839 + + y_zyhtst_l_1 = 0.0005 + + y_zynid_1 = (-0.000102077846072) + + y_zynid_2 = 0.000348695205252 + + y_zynid_3 = 0.000252250306328 + + y_zynid_4 = 0.00020597993691 + + y_zynid_5 = 0.000352649102887 + + y_zynid_6 = (-0.00091171528933) + + y_zynid_7 = (-0.000833252281803) + + y_zynid_8 = 0.000222028205174 + + y_zynid_9 = 0.00117029026307 + + y_zynid_10 = (-0.000704847602418) + + y_zynid_11 = (-0.00509852446865) + + y_zynid_12 = 0.00166112741624 + + y_zynid_13 = 0.000634556266817 + + y_zynid_14 = 0.0013968199402 + + y_zynid_15 = 0.00251911216766 + + y_zynid_16 = (-0.00251911216766) + + y_zynid_17 = 0.0129976207113 + + y_zynid_18 = (-0.0129976207113) + + y_zynid_19 = (-0.00311724318294) + + y_zynid_20 = 0.00311724318294 + + y_zynid_21 = (-0.0193734997949) + + y_zynid_22 = 0.0193734997949 + + y_zynid_23 = 0.00697401009889 + + y_zynid_24 = (-0.00697401009889) + + y_zynid_25 = 0.00296525365916 + + rho_trp_a = 0.0 + + trp_ā = 0.0 + + rho_fiscal = 0.97 + + rho_fiscalav = 0.9 + + fiscal_egfe = .01 + + fiscal_egfl = .01 + + av = 1 + + f̄iscal = 0 + + fpxrr_l̄ = 0.0 + + rho_fpxrr_l = 0.0 + + pmo_l̄ = 0.0 + + rho_pmo_l = 0.0 + + emo_l̄ = 0.0 + + rho_emo_l = 0.0 + + ugfdbtp_l̄ = 0.0 + + rho_ugfdbtp_l = 0.0 + + rho_gfsrt = 0.0 +end diff --git a/release_notes.md b/release_notes.md new file mode 100644 index 000000000..5f4dc1067 --- /dev/null +++ b/release_notes.md @@ -0,0 +1,50 @@ +compat with mooncake and switch to mooncake in docs +all functions meant to be used with derivatives now can be used with forward and reverse mode autodiff across all options (higher order, filters, etc.) +moved Bartels-Stewart sylvester and lyapunov equations solvers from MatrixEquations.jl to an extension +forwarddiff is an extension now, as internal derivatives are now done using to rrules based on analytical derivatives +Sylvester large systems solved with krylov methods now benefit from an ilu preconditioner +fix correctness issue in inversion filter (second order) +QME solver switches to doubling for large systems by default +use preallocated BLAS/LAPACK calls throughout for better performance and reduced allocations (QZ,LU,QR) +much more detailed tolerance settings for all solvers, with more robust defaults +Lyapunov solver accepts initial guess, has early termination, better handles unstable systems, the Krylov solver now works with the upper triangular system (implicitly forcing symmetry of the solution), added dqgmres support +overall much reduced allocations and better performance +many indices precomputed. moved constants to a separate struct. workspaces reduce allocations and caches are used for repeated solves +preallocated workspaces for all solvers, with better reuse and reduced allocations. use of LinearSolve and FastLapackInterface for matrix solves +write_mod_file (dynare) now allows to modify order, pruning and irf length +allow equation modification like in Troll +custom steady state function support (including in-place functions) +`get_equations` and `get_calibration_equations` accept a `filter` keyword argument +equations are now returned as expressions instead of strings +add correlation to `get_statistics` +counters for steady state and perturbation solves +replace MCMCChains with FlexiChains in estimation +shock decomposition includes calibration parameters +Lyapunov solver supports `has_unit_roots` parameter for unit root covariance handling +added warmup iterations for first-order inversion loglikelihood rrule, fixed gradient accuracy +NSSS solver refactored: struct dissolved into constants, functions, caches, and workspaces; `solve_nsss_wrapper` introduced as API layer +DispatchDoctor type stability coverage expanded across numerical source files +compat with Turing 0.45 +added FRBUS model +get_irf with parameters now also works with higher order +removed RecursiveFactorization and DifferentiationInterface direct dependency + +JET test on less functions (hot paths) so we get some coverage at least +see that all test scripts are actually run +there are various approaches of dealing with operations on sparse matrices and constructing them manually throughout th epackage. take stock of what approaches are out there, evaluate them in terms of performance gains and specific challenge they tackle and then use the best in class throughout appropriate applications + +analytical OBC jacobian instead of ForwardDiff.jl + +follow-ups: +update plots (the actual images) in docs to reflect changes in API and outputs as well as the examples in docstrings and docs +revisit func test tols +inegrate speed section in docs with benchmarks +do triage of helpers. either make sure they are used across the package where applciable but then also check that there is no more consice or already existing implmentation in the ecosystem. +check that we need this BARTELS_STEWART_AVAILABLE thing. its a weird construct to me. check alternatives +with these large models being used, make returns that are scaling with the number of variables and shocks output only the selected variables and shocks in order not to bloat memory and speed up computations. this includes IRFs, variance decompositions, etc. +rethink these crazy long input types +more DD coverage and fixes as well as getting JET to work again on the whole package +checkout StaticArrays for filter, if they speed things up +time filter step as well (SW07) for speed docs +eliminate this nested spaghetti code in nsss_solver +get iterative SSS/mean analytically \ No newline at end of file diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index d1942d6bb..64254ff17 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -23,9 +23,7 @@ import Dates # import MadNLP # import Ipopt # import AbstractDifferentiation as 𝒜 -import DifferentiationInterface as 𝒟 -import ForwardDiff as ℱ -backend = 𝒟.AutoForwardDiff() +# import ForwardDiff as ℱ # moved to ForwardDiffExt # import Diffractor: DiffractorForwardBackend # 𝒷 = 𝒜.ForwardDiffBackend # 𝒷 = Diffractor.DiffractorForwardBackend @@ -34,9 +32,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 @@ -48,19 +48,18 @@ import Krylov import Krylov: GmresWorkspace, DqgmresWorkspace, BicgstabWorkspace import LinearOperators import DataStructures: CircularBuffer, OrderedDict -import MacroTools: unblock, postwalk, prewalk, @capture, flatten +import MacroTools: unblock, postwalk, prewalk, @capture, flatten, rmlines # import SpeedMapping: speedmapping import Suppressor: @suppress import REPL import Unicode -import MatrixEquations # good overview: https://cscproxy.mpi-magdeburg.mpg.de/mpcsc/benner/talks/Benner-Melbourne2019.pdf # import NLboxsolve: nlboxsolve # 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__) @@ -155,21 +154,35 @@ const ParameterType = Union{Nothing, const SteadyStateFunctionType = Union{Nothing, Function, Missing} using DispatchDoctor -# @stable default_mode = "disable" begin # Imports include("default_options.jl") include("common_docstrings.jl") include("structures.jl") -include("solver_parameters.jl") +include("./steady_state/solver_parameters.jl") include("options_and_caches.jl") -include("macros.jl") +include("./steady_state/nsss_solver.jl") +include("occasionally_binding_constraints.jl") +include("./parser/macros.jl") +include("./parser/equation_processing.jl") +include("./parser/model_setup.jl") +include("./parser/equation_modification.jl") include("get_functions.jl") include("dynare.jl") include("inspect.jl") include("moments.jl") -include("perturbation.jl") - +include("./algorithms/fast_lapack_wrappers.jl") +include("./perturbation/derivatives.jl") +include("./perturbation/solution.jl") +include("./steady_state/stochastic_steady_state.jl") +include("impulse_response_function.jl") + +# Sentinel for MatrixEquations extension (bartels_stewart algorithm). +# Set to `true` by MatrixEquationsExt.__init__() when the package is loaded. +const BARTELS_STEWART_AVAILABLE = Ref(false) +has_bartels_stewart() = BARTELS_STEWART_AVAILABLE[] + +include("./algorithms/preconditioner.jl") include("./algorithms/sylvester.jl") include("./algorithms/lyapunov.jl") include("./algorithms/nonlinear_solver.jl") @@ -180,9 +193,6 @@ include("./filter/inversion.jl") include("./filter/kalman.jl") -# end # dispatch_doctor - - export @model, @parameters, solve! export plot_irfs, plot_irf, plot_IRF, plot_simulations, plot_solution, plot_simulation, plot_girf #, plot @@ -203,12 +213,14 @@ 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 export get_equations, get_steady_state_equations, get_dynamic_equations, get_calibration_equations, get_parameters, get_calibrated_parameters, get_parameters_in_equations, get_parameters_defined_by_parameters, get_parameters_defining_parameters, get_calibration_equation_parameters, get_variables, get_nonnegativity_auxiliary_variables, get_dynamic_auxiliary_variables, get_shocks, get_state_variables, get_jump_variables, get_missing_parameters, has_missing_parameters, get_solution_counts, print_solution_counts +export write_julia_model_file, replace_equations!, replace_calibration_equations! +export update_equations!, update_calibration_equations!, add_equation!, add_calibration_equation!, remove_equation!, remove_calibration_equation!, get_revision_history # Internal export irf, girf @@ -283,7 +295,6 @@ Symbolics.@register_symbolic normcdf(z) Symbolics.@register_symbolic pnorm(p) Symbolics.@register_symbolic dnorm(p) -end # dispatch_doctor # ── norminvcdf, norminv & qnorm ── # d/dp (norminvcdf(p)) = 1 / normpdf(norminvcdf(p)) @@ -341,7 +352,6 @@ else Symbolics.derivative(normcdf, args, Val{1}()) end -@stable default_mode = "disable" begin Base.show(io::IO, 𝓂::ℳ) = println(io, @@ -367,93 +377,18 @@ 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 -check_for_dynamic_variables(ex::Float64) = false -check_for_dynamic_variables(ex::Symbol) = occursin(r"₍₁₎|₍₀₎|₍₋₁₎",string(ex)) - # end # dispatch_doctor -function compare_args_and_kwargs(dicts::Vector{S}) where S <: Dict - N = length(dicts) - @assert N ≥ 2 "Need at least two dictionaries to compare" - - diffs = Dict{Symbol,Any}() - - # assume all dictionaries share the same set of keys - for k in keys(dicts[1]) - if k in [:plot_data, :plot_type] - # skip keys that are not relevant for comparison - continue - end - - vals = [d[k] for d in dicts] - - if all(v -> v isa Dict, vals) - # recurse into nested dictionaries - nested = compare_args_and_kwargs(vals) - if !isempty(nested) - diffs[k] = nested - end - - elseif all(v -> v isa KeyedArray, vals) - # compare by length and elementwise equality - base = vals[1] - identical = all(v -> length(v) == length(base) && all(collect(v) .== collect(base)), vals[2:end]) - if !identical - diffs[k] = vals - end - - elseif all(v -> v isa AbstractArray, vals) - # compare by length and elementwise equality - base = vals[1] - identical = all(v -> length(v) == length(base) && all(v .== base), vals[2:end]) - if !identical - diffs[k] = vals - end - - else - # scalar or other types - identical = all(v -> v == vals[1], vals[2:end]) - if !identical - diffs[k] = vals - end - end - end - - return diffs -end - - -function mul_reverse_AD!( C::Matrix{S}, - A::AbstractMatrix{M}, - B::AbstractMatrix{N}) where {S <: Real, M <: Real, N <: Real} - ℒ.mul!(C,A,B) -end -function check_for_dynamic_variables(ex::Expr) - dynamic_indicator = Bool[] +# Generic primal-value extraction — identity for plain reals. +# ForwardDiffExt extends this for ForwardDiff.Dual numbers. +primal(x::Real) = x - postwalk(x -> - x isa Expr ? - x.head == :ref ? - occursin(r"^(ss|stst|steady|steadystate|steady_state){1}$"i,string(x.args[2])) ? - x : - begin - push!(dynamic_indicator,true) - x - end : - x : - x, - ex) - - any(dynamic_indicator) -end function normalize_filtering_options(filter::Symbol, smooth::Bool, @@ -522,77 +457,8 @@ function adjust_generalised_irf_flag(generalised_irf::Bool, return generalised_irf end -end # dispatch_doctor - -function transform_expression(expr::Expr) - # Dictionary to store the transformations for reversing - reverse_transformations = Dict{Symbol, Expr}() - # Counter for generating unique placeholders - unique_counter = Ref(0) - - # Step 1: Replace min/max calls and record their original form - function replace_min_max(expr) - if expr isa Expr && expr.head == :call && (expr.args[1] == :min || expr.args[1] == :max) - # Replace min/max functions with a placeholder - # placeholder = Symbol("minimal__P", unique_counter[]) - placeholder = :minmax__P - unique_counter[] += 1 - - # Store the original min/max call for reversal - reverse_transformations[placeholder] = expr - - return placeholder - else - return expr - end - end - - # Step 2: Transform :ref fields in the rest of the expression - function transform_ref_fields(expr) - if expr isa Expr && expr.head == :ref && isa(expr.args[1], Symbol) - # Handle :ref expressions - if isa(expr.args[2], Number) || isa(expr.args[2], Symbol) - if expr.args[2] < 0 - new_symbol = Symbol(expr.args[1], "__", abs(expr.args[2])) - else - new_symbol = Symbol(expr.args[1], "_", expr.args[2]) - end - else - # Generate a unique placeholder for complex :ref - unique_counter[] += 1 - placeholder = Symbol("__placeholder", unique_counter[]) - new_symbol = placeholder - end - - # Record the reverse transformation - reverse_transformations[new_symbol] = expr - - return new_symbol - else - return expr - end - end - - - # Replace equality sign with minus - function replace_equality_with_minus(expr) - if expr isa Expr && expr.head == :(=) - return Expr(:call, :-, expr.args...) - else - return expr - end - end - - # Apply transformations - expr = postwalk(replace_min_max, expr) - expr = postwalk(transform_ref_fields, expr) - transformed_expr = postwalk(replace_equality_with_minus, expr) - - return transformed_expr, reverse_transformations -end - -function process_shocks_input(shocks::Union{Symbol_input, String_input, Matrix{Float64}, KeyedArray{Float64}}, +@unstable function process_shocks_input(shocks::Union{Symbol_input, String_input, Matrix{Float64}, KeyedArray{Float64}}, negative_shock::Bool, shock_size::Real, periods::Int, @@ -632,7 +498,7 @@ function process_shocks_input(shocks::Union{Symbol_input, String_input, Matrix{F shock_idx = 1 elseif shocks isa Expr error("Expressions are not a valid input for shocks. Please provide a Symbol, Vector of Symbols, Matrix of Float64, KeyedArray of Float64, or :none.") - elseif (typeof(shocks) <: Symbol_input) || (typeof(shocks) <: String_input) + else shock_history = zeros(𝓂.constants.post_model_macro.nExo, periods) periods_extended = periods @@ -657,373 +523,120 @@ function process_shocks_input(shocks::Union{Symbol_input, String_input, Matrix{F return shocks, negative_shock, shock_size, periods_extended, shock_idx, shock_history end -@stable default_mode = "disable" begin - -function process_ignore_obc_flag(shocks, - ignore_obc::Bool, - 𝓂::ℳ; - maxlog::Int = DEFAULT_MAXLOG) - stochastic_model = length(𝓂.constants.post_model_macro.exo) > 0 - obc_model = length(𝓂.equations.obc_violation) > 0 - - obc_shocks_included = false - - if stochastic_model && obc_model - if shocks isa Matrix{Float64} - obc_indices = contains.(string.(𝓂.constants.post_model_macro.exo), "ᵒᵇᶜ") - if any(obc_indices) - obc_shocks_included = sum(abs2, shocks[obc_indices, :]) > 1e-10 - end - elseif shocks isa KeyedArray{Float64} - shock_axis = collect(axiskeys(shocks, 1)) - shock_axis = shock_axis isa Vector{String} ? shock_axis .|> Meta.parse .|> replace_indices : shock_axis - - obc_shocks = 𝓂.constants.post_model_macro.exo[contains.(string.(𝓂.constants.post_model_macro.exo), "ᵒᵇᶜ")] - relevant_shocks = intersect(obc_shocks, shock_axis) - - if !isempty(relevant_shocks) - obc_shocks_included = sum(abs2, shocks(relevant_shocks, :)) > 1e-10 - end - else - shock_idx = parse_shocks_input_to_index(shocks, 𝓂.constants) - - selected_shocks = if (shock_idx isa Vector) || (shock_idx isa UnitRange) - length(shock_idx) > 0 ? 𝓂.constants.post_model_macro.exo[shock_idx] : Symbol[] - else - [𝓂.constants.post_model_macro.exo[shock_idx]] - end - - obc_shocks = 𝓂.constants.post_model_macro.exo[contains.(string.(𝓂.constants.post_model_macro.exo), "ᵒᵇᶜ")] - obc_shocks_included = !isempty(intersect(selected_shocks, obc_shocks)) - end - end - - ignore_obc_flag = ignore_obc - - if ignore_obc_flag && !obc_model - @info "`ignore_obc = true` has no effect because $(𝓂.model_name) has no occasionally binding constraints. Setting `ignore_obc = false`." maxlog = maxlog - ignore_obc_flag = false - end - - if ignore_obc_flag && obc_shocks_included - @warn "`ignore_obc = true` cannot be applied because shocks affecting occasionally binding constraints are included. Enforcing the constraints instead and setting `ignore_obc = false`." maxlog = maxlog - ignore_obc_flag = false - end - - occasionally_binding_constraints = obc_model && !ignore_obc_flag - - return ignore_obc_flag, occasionally_binding_constraints, obc_shocks_included -end - - - -function reverse_transformation(transformed_expr::Expr, reverse_dict::Dict{Symbol, Expr}) - # Function to replace the transformed symbols with their original form - function revert_symbol(expr) - if expr isa Symbol && haskey(reverse_dict, expr) - return reverse_dict[expr] - else - return expr - end - end - - # Revert the expression using postwalk - reverted_expr = postwalk(revert_symbol, transformed_expr) - return reverted_expr -end -function replace_with_one(equation::SPyPyC.Sym{PythonCall.Core.Py}, variable::SPyPyC.Sym{PythonCall.Core.Py})::SPyPyC.Sym{PythonCall.Core.Py} - # equation.subs(variable, 1).replace(SPyPyC.Sym(ℯ), exp(1)) - tmp = SPyPyC.subs(equation, variable, 1) - return replace_e(tmp) +function invalidate_cache_validity!(𝓂::ℳ) + 𝓂.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[] + 𝓂.caches.valid_for.covariance_first_order = Float64[] + 𝓂.caches.valid_for.covariance_second_order = Float64[] + 𝓂.caches.valid_for.covariance_third_order = Float64[] + 𝓂.caches.valid_for.covariance_third_order_obs_key = Int[] + 𝓂.caches.valid_for.covariance_third_order_autocorr = Float64[] + 𝓂.caches.valid_for.covariance_third_order_autocorr_obs_key = Int[] + 𝓂.caches.valid_for.covariance_third_order_autocorr_periods = Int[] + return nothing end -function replace_e(equation::SPyPyC.Sym{PythonCall.Core.Py})::SPyPyC.Sym{PythonCall.Core.Py} - outraw = SPyPyC.subs(equation, SPyPyC.Sym(ℯ), exp(1)) - - if outraw isa SPyPyC.Sym{PythonCall.Core.Py} - out = outraw - else - out = collect(outraw)[1] - end - - return out -end -function replace_symbolic(equation::SPyPyC.Sym{PythonCall.Core.Py}, variable::SPyPyC.Sym{PythonCall.Core.Py}, replacement::SPyPyC.Sym{PythonCall.Core.Py})::SPyPyC.Sym{PythonCall.Core.Py} - # equation.subs(variable, replacement) - return SPyPyC.subs(equation, variable, replacement) -end +function reset_nsss_solver_cache!(𝓂::ℳ) + empty!(𝓂.caches.solver) -function solve_symbolically(equation::SPyPyC.Sym{PythonCall.Core.Py}, variable::SPyPyC.Sym{PythonCall.Core.Py})::Union{Nothing,Vector{SPyPyC.Sym{PythonCall.Core.Py}}} - soll = try SPyPyC.solve(equation, variable) - catch - end + c = 𝓂.constants.nsss_solver + ms = 𝓂.constants.post_complete_parameters + seed = Vector{Vector{Float64}}() - return soll -end + for step_idx in 1:c.n_steps + if c.step_types[step_idx] == NUMERICAL_STEP + wr = c.write_ranges[step_idx] + nbr = c.numerical_bounds_ranges[step_idx] + guess_len = min(length(wr), length(nbr)) + guesses = Vector{Float64}(undef, guess_len) -function solve_symbolically(equations::Vector{SPyPyC.Sym{PythonCall.Core.Py}}, variables::Vector{SPyPyC.Sym{PythonCall.Core.Py}})::Union{Nothing,Dict{SPyPyC.Sym{PythonCall.Core.Py}, SPyPyC.Sym{PythonCall.Core.Py}}} - soll = try SPyPyC.solve(equations, variables) - catch + for i in 1:guess_len + sol_idx = c.write_indices[wr[i]] + sol_name = sol_idx <= length(ms.nsss_sol_names) ? ms.nsss_sol_names[sol_idx] : Symbol("") + guesses[i] = get(𝓂.constants.post_parameters_macro.guess, sol_name, Inf) end - if soll == Any[] - soll = Dict{SPyPyC.Sym{PythonCall.Core.Py}, SPyPyC.Sym{PythonCall.Core.Py}}() - elseif soll isa Vector - soll = Dict{SPyPyC.Sym{PythonCall.Core.Py}, SPyPyC.Sym{PythonCall.Core.Py}}(variables .=> soll[1]) - end - - return soll -end - -function transform_obc(ex::Expr; avoid_solve::Bool = false) - transformed_expr, reverse_dict = transform_expression(ex) - - for symbs in get_symbols(transformed_expr) - sym_value = SPyPyC.symbols(string(symbs), real = true, finite = true) - Core.eval(SymPyWorkspace, :($symbs = $sym_value)) - end - - eq = Core.eval(SymPyWorkspace, transformed_expr) - - if avoid_solve || count_ops(Meta.parse(string(eq))) > 15 - soll = nothing - else - soll = solve_symbolically(eq, Core.eval(SymPyWorkspace, :minmax__P)) - end - - if !isempty(soll) - sorted_minmax = Expr(:call, reverse_dict[:minmax__P].args[1], :($(reverse_dict[:minmax__P].args[2]) - $(Meta.parse(string(soll[1])))), :($(reverse_dict[:minmax__P].args[3]) - $(Meta.parse(string(soll[1]))))) - return reverse_transformation(sorted_minmax, reverse_dict) - else - @error "Occasionally binding constraint not well-defined. See documentation for examples." - end -end - - -function obc_constraint_optim_fun(res::Vector{S}, X::Vector{S}, jac::Matrix{S}, p) where S - 𝓂 = p[4] - - if length(jac) > 0 - # jac .= 𝒜.jacobian(𝒷(), xx -> 𝓂.functions.obc_violation(xx, p), X)[1]' - jac .= 𝒟.jacobian(xx -> 𝓂.functions.obc_violation(xx, p), backend, X)' - end - - res .= 𝓂.functions.obc_violation(X, p) - - return nothing -end - -function obc_objective_optim_fun(X::Vector{S}, grad::Vector{S})::S where S - if length(grad) > 0 - grad .= 2 .* X - end - - sum(abs2, X) -end - -function set_up_obc_violation_function!(𝓂) - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) - present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎$"))) - - sort!(present_varss ,by = x->replace(string(x),r"₍₀₎$"=>"")) - - # write indices in auxiliary objects - dyn_var_present_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₀₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎"))) - - dyn_var_present = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_var_present_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) - - SS_and_pars_names = ms.SS_and_pars_names - - dyn_var_present_idx = indexin(dyn_var_present , SS_and_pars_names) - - alll = [] - for (i,var) in enumerate(present_varss) - if !(match(r"^χᵒᵇᶜ", string(var)) === nothing) - push!(alll,:($var = Y[$(dyn_var_present_idx[i]),1:max(periods, 1)])) + push!(seed, guesses) + push!(seed, Float64[Inf]) end end - calc_obc_violation = :(function calculate_obc_violation(x, p) - state, state_update, reference_steady_state, 𝓂, algorithm, periods, shock_values = p - - T = 𝓂.constants.post_model_macro - - Y = zeros(typeof(x[1]), T.nVars, periods+1) - - shock_values = convert(typeof(x), shock_values) - - shock_values[contains.(string.(T.exo),"ᵒᵇᶜ")] .= x - - zero_shock = zero(shock_values) + push!(seed, fill(Inf, length(ms.parameters))) + push!(𝓂.caches.solver, seed) - if algorithm ∈ [:pruned_second_order, :pruned_third_order] - states = state_update(state, shock_values) - Y[:,1] = sum(states) - else - Y[:,1] = state_update(state, shock_values) - end - - for t in 1:periods - if algorithm ∈ [:pruned_second_order, :pruned_third_order] - states = state_update(states, zero_shock) - Y[:,t+1] = sum(states) - else - Y[:,t+1] = state_update(Y[:,t], zero_shock) - end - end + return nothing +end - Y .+= reference_steady_state[1:T.nVars] - $(alll...) +function clear_solution_caches!(𝓂::ℳ, algorithm::Symbol) + reset_nsss_solver_cache!(𝓂) - constraint_values = Vector[] + 𝓂.caches.first_order_solution_matrix = zeros(0,0) + 𝓂.caches.first_order_obc_solution_matrix = zeros(0,0) + 𝓂.caches.qme_solution = zeros(0,0) + 𝓂.caches.has_unit_roots = false + 𝓂.caches.second_order_solution = spzeros(0,0) + 𝓂.caches.third_order_solution = spzeros(0,0) - $(𝓂.equations.obc_violation...) + 𝓂.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[] - return vcat(constraint_values...) - end) + resize!(𝓂.caches.non_stochastic_steady_state, 0) - 𝓂.functions.obc_violation = @RuntimeGeneratedFunction(calc_obc_violation) + invalidate_cache_validity!(𝓂) return nothing end -function check_for_minmax(expr) - contains_minmax = Bool[] +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, + :covariance_first_order, + :covariance_second_order, + :covariance_third_order, +) - postwalk(x -> - x isa Expr ? - x.head == :call ? - x.args[1] ∈ [:max,:min] ? - begin - push!(contains_minmax,true) - x - end : - x : - x : - x, - expr) - any(contains_minmax) -end - - -function write_obc_violation_equations(𝓂) - eqs = Expr[] - for (i,eq) in enumerate(𝓂.equations.dynamic) - if check_for_minmax(eq) - minmax_fixed_eqs = postwalk(x -> - x isa Expr ? - x.head == :call ? - length(x.args) == 3 ? - x.args[3] isa Expr ? - x.args[3].args[1] ∈ [:Min, :min, :Max, :max] ? - begin - plchldr = Symbol(replace(string(x.args[2]), "₍₀₎" => "")) - - ineq_plchldr_1 = x.args[3].args[2] isa Symbol ? Symbol(replace(string(x.args[3].args[2]), "₍₀₎" => "")) : x.args[3].args[2] - - arg1 = x.args[3].args[2] - arg2 = x.args[3].args[3] - - dyn_1 = check_for_dynamic_variables(x.args[3].args[2]) - dyn_2 = check_for_dynamic_variables(x.args[3].args[3]) - - cond1 = Expr[] - cond2 = Expr[] - - maximisation = contains(string(plchldr), "⁺") - - # if dyn_1 - # if maximisation - # push!(cond1, :(push!(constraint_values, $(x.args[3].args[2])))) - # # push!(cond2, :(push!(constraint_values, $(x.args[3].args[2])))) - # else - # push!(cond1, :(push!(constraint_values, -$(x.args[3].args[2])))) - # # push!(cond2, :(push!(constraint_values, -$(x.args[3].args[2])))) # RBC - # end - # end - - # if dyn_2 - # if maximisation - # push!(cond1, :(push!(constraint_values, $(x.args[3].args[3])))) - # # push!(cond2, :(push!(constraint_values, $(x.args[3].args[3])))) # testmax - # else - # push!(cond1, :(push!(constraint_values, -$(x.args[3].args[3])))) - # # push!(cond2, :(push!(constraint_values, -$(x.args[3].args[3])))) # RBC - # end - # end - - - if maximisation - push!(cond1, :(push!(constraint_values, [sum($(x.args[3].args[2]) .* $(x.args[3].args[3]))]))) - push!(cond1, :(push!(constraint_values, $(x.args[3].args[2])))) - push!(cond1, :(push!(constraint_values, $(x.args[3].args[3])))) - # push!(cond1, :(push!(constraint_values, max.($(x.args[3].args[2]), $(x.args[3].args[3]))))) - else - push!(cond1, :(push!(constraint_values, [sum($(x.args[3].args[2]) .* $(x.args[3].args[3]))]))) - push!(cond1, :(push!(constraint_values, -$(x.args[3].args[2])))) - push!(cond1, :(push!(constraint_values, -$(x.args[3].args[3])))) - # push!(cond1, :(push!(constraint_values, min.($(x.args[3].args[2]), $(x.args[3].args[3]))))) - end - - # if maximisation - # push!(cond1, :(push!(shock_sign_indicators, true))) - # # push!(cond2, :(push!(shock_sign_indicators, true))) - # else - # push!(cond1, :(push!(shock_sign_indicators, false))) - # # push!(cond2, :(push!(shock_sign_indicators, false))) - # end - - # :(if isapprox($plchldr, $ineq_plchldr_1, atol = 1e-12) - # $(Expr(:block, cond1...)) - # else - # $(Expr(:block, cond2...)) - # end) - :($(Expr(:block, cond1...))) - end : - x : - x : - x : - x : - x, - eq) - - push!(eqs, minmax_fixed_eqs) +@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 eqs -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) - end - - 𝓂.caches.qme_solution = zeros(0,0) - 𝓂.caches.second_order_solution = spzeros(0,0) - 𝓂.caches.third_order_solution = spzeros(0,0) - - return nothing + return true end @@ -1107,35 +720,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 @@ -1143,375 +731,104 @@ end -""" - infer_step(x_axis) - -Infer the step for an axis. - -For dates, if the last two points share the same day-of-month, the step is -inferred in whole months (e.g. Month(1), Month(3), …). Otherwise the raw -difference is used. For non time types, uses the plain difference. -""" -function infer_step(x_axis::AbstractVector{T}) where {T<:Number} - x_axis[end] - x_axis[end-1] -end - -function infer_step(x_axis::AbstractVector{T}) where {T<:Dates.TimeType} - d1 = x_axis[end-1] - d2 = x_axis[end] - - # try to infer a monthly step if aligned by day-of-month - if Dates.day(d1) == Dates.day(d2) - m1 = 12 * Dates.year(d1) + Dates.month(d1) - m2 = 12 * Dates.year(d2) + Dates.month(d2) - mstep = m2 - m1 - if mstep != 0 - return Dates.Month(mstep) +# Helper to convert dense matrix to sparse using I,J,V format (avoids Julia 1.12 SparseArrays bug) +function dense_to_sparse(A::DenseMatrix{S}, tol::R) where {S <: Real, R <: AbstractFloat} + m, n = size(A) + I = Int[] + J = Int[] + V = S[] + @inbounds for j in 1:n + for i in 1:m + v = A[i,j] + if abs(v) > tol + push!(I, i) + push!(J, j) + push!(V, v) + end end end - - # fall back to the raw difference (in days, milliseconds, …) - return d2 - d1 + return sparse(I, J, V, m, n) 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 +function choose_matrix_format(A::ℒ.Diagonal{S, Vector{S}}; + density_threshold::Float64 = .1, + 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 length(A) < 100 + a = convert(Matrix, A) 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 + if multithreaded + a = A |> sparse |> ThreadedSparseArrays.ThreadedSparseMatrixCSC + else + a = A |> sparse 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 + return a 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 +@unstable function choose_matrix_format(A::ℒ.Adjoint{S, M}; + density_threshold::Float64 = .1, + 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}} + 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}}; +# density_threshold::Float64 = .1, +# 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} +# choose_matrix_format(convert(typeof(transpose(A)),A), +# density_threshold = density_threshold, +# min_length = min_length, +# multithreaded = multithreaded, +# tol = tol) +# 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 +@unstable function choose_matrix_format(A::DenseMatrix{S}; + density_threshold::Float64 = .1, + 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 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 + return ThreadedSparseArrays.ThreadedSparseMatrixCSC(a) + else + return a end + else + return convert(Matrix, A) end end +@unstable function choose_matrix_format(A::AbstractSparseMatrix{S}; + density_threshold::Float64 = .1, + 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} + droptol!(A, tol) - -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}}; - density_threshold::Float64 = .1, - 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 length(A) < 100 - a = convert(Matrix, A) - else - if multithreaded - a = A |> sparse |> ThreadedSparseArrays.ThreadedSparseMatrixCSC - else - a = A |> sparse - end - end - - return a -end - - -function choose_matrix_format(A::ℒ.Adjoint{S, M}; - density_threshold::Float64 = .1, - 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) -end - -# function choose_matrix_format(A::ℒ.Adjoint{S, <: AbstractSparseMatrix{S}}; -# density_threshold::Float64 = .1, -# 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} -# choose_matrix_format(convert(typeof(transpose(A)),A), -# density_threshold = density_threshold, -# min_length = min_length, -# multithreaded = multithreaded, -# tol = tol) -# end - -# Helper to convert dense matrix to sparse using I,J,V format (avoids Julia 1.12 SparseArrays bug) -function dense_to_sparse(A::DenseMatrix{S}, tol::R) where {S <: Real, R <: AbstractFloat} - m, n = size(A) - I = Int[] - J = Int[] - V = S[] - @inbounds for j in 1:n - for i in 1:m - v = A[i,j] - if abs(v) > tol - push!(I, i) - push!(J, j) - push!(V, v) - end - end - end - return sparse(I, J, V, m, n) -end - -function choose_matrix_format(A::DenseMatrix{S}; - density_threshold::Float64 = .1, - 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 - # Use dense_to_sparse to avoid Julia 1.12 SparseArrays bug in SparseMatrixCSC(::Matrix) - a = dense_to_sparse(A, tol) - if multithreaded - return ThreadedSparseArrays.ThreadedSparseMatrixCSC(a) - else - return a - end - else - return convert(Matrix, A) - end -end - -function choose_matrix_format(A::AbstractSparseMatrix{S}; - density_threshold::Float64 = .1, - 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} - droptol!(A, tol) - - lennz = nnz(A) + lennz = nnz(A) if lennz / length(A) > density_threshold || length(A) < min_length a = convert(Matrix, A) @@ -1534,277 +851,10 @@ function choose_matrix_format(A::AbstractSparseMatrix{S}; return a end -function mat_mult_kron(A::AbstractSparseMatrix{R}, - B::AbstractMatrix{T}, - C::AbstractMatrix{T}, - D::AbstractMatrix{S}) 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)) - - # vals = T[] - # rows = Int[] - # cols = Int[] - - Ā = zeros(T, n_rowC, n_rowB) - ĀB = zeros(T, n_rowC, n_colB) - CĀB = zeros(T, n_colC, n_colB) - vCĀB = zeros(T, n_colB * n_colC) - # vCĀBD = zeros(size(D,2)) - - rv = unique(A isa SparseMatrixCSC ? A.rowval : A.A.rowval) - - # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique - @inbounds for row in rv - @views copyto!(Ā, A[row, :]) - ℒ.mul!(ĀB, Ā, B) - ℒ.mul!(CĀB, C', ĀB) - copyto!(vCĀB, CĀB) - @views ℒ.mul!(X[row,:], D', vCĀB) - end - - return choose_matrix_format(X) - # ℒ.mul!(vCĀBD, D', vCĀB) - - # for (i,v) in enumerate(vCĀBD) - # if abs(v) > eps() - # push!(rows, row) - # push!(cols, i) - # push!(vals, v) - # end - # end - # end - - # 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 -end - - - - -function mat_mult_kron(A::DenseMatrix{R}, - B::AbstractMatrix{T}, - C::AbstractMatrix{T}, - D::AbstractMatrix{S}) 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)) - - # vals = T[] - # rows = Int[] - # cols = Int[] - - Ā = zeros(T, n_rowC, n_rowB) - ĀB = zeros(T, n_rowC, n_colB) - CĀB = zeros(T, n_colC, n_colB) - vCĀB = zeros(T, n_colB * n_colC) - # vCĀBD = zeros(size(D,2)) - - # rv = A isa SparseMatrixCSC ? A.rowval : A.A.rowval - - # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique - r = 1 - @inbounds for row in eachrow(A) - @views copyto!(Ā, row) - ℒ.mul!(ĀB, Ā, B) - ℒ.mul!(CĀB, C', ĀB) - copyto!(vCĀB, CĀB) - @views ℒ.mul!(X[row,:], D', vCĀB) - r += 1 - end - - return choose_matrix_format(X) - # ℒ.mul!(vCĀBD, D', vCĀB) - - # for (i,v) in enumerate(vCĀBD) - # if abs(v) > eps() - # push!(rows, row) - # push!(cols, i) - # push!(vals, v) - # end - # end - # end - - # 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 -end - -function mat_mult_kron(A::AbstractSparseMatrix{R}, - B::AbstractMatrix{T}, - C::AbstractMatrix{T}; - 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} - n_rowB = size(B,1) - n_colB = size(B,2) - - n_rowC = size(C,1) - n_colC = size(C,2) - - estimated_nnz = 0 - I = Vector{Int}() - J = Vector{Int}() - V = Vector{T}() - X = zeros(T, 0, 0) - - if sparse - nnzA = nnz(A) - nnzB = sum(abs.(B) .> eps()) - nnzC = sum(abs.(C) .> eps()) - - p = nnzA * nnzB * nnzC / (length(A) * length(B) * length(C)) - - if length(sparse_preallocation[1]) == 0 - estimated_nnz = Int(ceil((1-(1-p)^size(A,1))*size(A,1) * n_colB * n_colC)) - - 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) - - I = sparse_preallocation[1] - J = sparse_preallocation[2] - V = sparse_preallocation[3] - end - else - X = zeros(T, size(A,1), n_colB * n_colC) - end - - Ā = zeros(T, n_rowC, n_rowB) - ĀB = zeros(T, n_rowC, n_colB) - CĀB = zeros(T, n_colC, n_colB) - - rv = A isa SparseMatrixCSC ? A.rowval : A.A.rowval - - α = .7 # speed of Vector increase - k = 0 - - # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique - @inbounds for row in rv |> unique - @views copyto!(Ā, A[row, :]) - ℒ.mul!(ĀB, Ā, B) - ℒ.mul!(CĀB, C', ĀB) - - if sparse - for (i,v) in enumerate(CĀB) - if abs(v) > eps() - k += 1 - - if k > estimated_nnz - estimated_nnz += min(size(A,1) * n_colB * n_colC, max(10000, Int(ceil((α - 1) * estimated_nnz + (1 - α) * size(A,1) * n_colB * n_colC)))) - - 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,:], CĀB) - end - end - - if sparse - resize!(I, k) - resize!(J, k) - resize!(V, k) - - klasttouch = sparse_preallocation[4] # Vector{Ti}(undef, n) - csrrowptr = sparse_preallocation[5] # Vector{Ti}(undef, m + 1) - csrcolval = sparse_preallocation[6] # Vector{Ti}(undef, length(I)) - csrnzval = sparse_preallocation[7] # Vector{Tv}(undef, length(I)) - - resize!(klasttouch, n_colB * n_colC) - resize!(csrrowptr, size(A, 1) + 1) - 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) - # out = sparse!(I, J, V, size(A, 1), n_colB * n_colC) - else - out = choose_matrix_format(X) - end - - return out -end - - - - -function mat_mult_kron(A::DenseMatrix{R}, - B::AbstractMatrix{T}, - C::AbstractMatrix{T}) where {R <: Real, T <: 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), n_colB * n_colC) - - # vals = T[] - # rows = Int[] - # cols = Int[] - - Ā = zeros(T, n_rowC, n_rowB) - ĀB = zeros(T, n_rowC, n_colB) - CĀB = zeros(T, n_colC, n_colB) - - # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique - r = 1 - @inbounds for row in eachrow(A) - @views copyto!(Ā, row) - ℒ.mul!(ĀB, Ā, B) - ℒ.mul!(CĀB, C', ĀB) - - @views copyto!(X[r,:], CĀB) - r += 1 - end - - return choose_matrix_format(X) - # for (i,v) in enumerate(CĀB) - # if abs(v) > eps() - # push!(rows, row) - # push!(cols, i) - # push!(vals, v) - # end - # end - # end - - # if VERSION >= v"1.10" - # return sparse!(rows,cols,vals,size(A,1),n_colB*n_colC) - # else - # return sparse(rows,cols,vals,size(A,1),n_colB*n_colC) - # end -end function sparse_preallocated!(Ŝ::Matrix{T}; ℂ::higher_order_workspace{T,F,H} = Higher_order_workspace()) where {T <: Real, F <: AbstractFloat, H <: Real} if !(eltype(ℂ.tmp_sparse_prealloc6[3]) == T) - ℂ.tmp_sparse_prealloc6 = Higher_order_workspace(T = T, S = F) + ℂ.tmp_sparse_prealloc6 = Higher_order_workspace(T, F) end I = ℂ.tmp_sparse_prealloc6[1] @@ -1856,247 +906,167 @@ function sparse_preallocated!(Ŝ::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(), - 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 - â = copy(a') - a = sparse(a') - - rmask = colmask - colmask = rowmask - rowmask = rmask - elseif typeof(a) <: DenseMatrix{T} - â = copy(a) - a = sparse(a) + + +# Dead code: compressed_kron (2-arg) — never called anywhere; rrule also dead +#= +# 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 - â = convert(Matrix, a) # Convert to dense matrix for faster access + Vector{T}(undef, length(spV_untyped)) 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 + 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 - # 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 + II = spI + JJ = spJ + VV = spV - m3_exp = (length(colmask) > 0 || length(rowmask) > 0) ? 3 : 4 + cnt = 0 - if length(sparse_preallocation[1]) == 0 - estimated_nnz = floor(Int, max(m3_r * m3_c * (lennz / length(a)) ^ m3_exp, 10000)) + # 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 - resize!(sparse_preallocation[1], estimated_nnz) - resize!(sparse_preallocation[2], estimated_nnz) - resize!(sparse_preallocation[3], estimated_nnz) + for β in 1:α + for γ in 1:β + σ_col = (β - 1) * nᶜ + γ + rng_σ = SparseArrays.nzrange(σs, σ_col) + isempty(rng_σ) && continue - I = sparse_preallocation[1] - J = sparse_preallocation[2] - V = sparse_preallocation[3] - else - estimated_nnz = length(sparse_preallocation[3]) + col = (α - 1) * α * (α + 1) ÷ 6 + (β - 1) * β ÷ 2 + γ - resize!(sparse_preallocation[1], estimated_nnz) + @inbounds for pA in rng_A + i = rv_A[pA] + a_val = nzv_A[pA] - I = sparse_preallocation[1] - J = sparse_preallocation[2] - V = sparse_preallocation[3] - end + for pσ in rng_σ + s = rv_σ[pσ] + σ_val = nzv_σ[pσ] - # k = Threads.Atomic{Int}(0) # Counter for non-zero entries - # k̄ = Threads.Atomic{Int}(0) # effectively slower than the non-threaded version + val = a_val * σ_val + abs(val) > tol || continue - k = 0 + # Decompose σ row: s = (j-1)*nᵣ + k + j = (s - 1) ÷ nᵣ + 1 + k = (s - 1) % nᵣ + 1 - # end # timeit_debug + # 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 - # @timeit_debug timer "findnz" begin - - # Find unique non-zero row and column indices - rowinds, colinds, _ = findnz(a) - ui = unique(rowinds) - uj = unique(colinds) - - # end # timeit_debug + row = (i₁ - 1) * i₁ * (i₁ + 1) ÷ 6 + (j₁ - 1) * j₁ ÷ 2 + k₁ - # @timeit_debug timer "Loop" begin - # Triple nested loops for (i1 ≤ j1 ≤ k1) and (i2 ≤ j2 ≤ k2) - # Polyester.@batch threadlocal=(Vector{Int}(), Vector{Int}(), Vector{T}()) for i1 in ui - # Polyester.@batch minbatch = 10 for i1 in ui - # Threads.@threads for i1 in ui - norowmask = length(rowmask) == 0 - nocolmask = length(colmask) == 0 - - for i1 in ui - for j1 in ui - if j1 ≤ i1 - for k1 in ui - if k1 ≤ j1 - - row = (i1-1) * i1 * (i1+1) ÷ 6 + (j1-1) * j1 ÷ 2 + k1 - - if norowmask || row in rowmask - for i2 in uj - for j2 in uj - if j2 ≤ i2 - for k2 in uj - if k2 ≤ j2 - - col = (i2-1) * i2 * (i2+1) ÷ 6 + (j2-1) * j2 ÷ 2 + k2 - - if nocolmask || col in colmask - # @timeit_debug timer "Multiplication" begin - @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] - - # Compute the six unique products - # val = 0.0 - # val += aii * ajj * akk - # val += aij * aji * akk - # val += aik * ajj * aki - # val += aij * ajk * aki - # val += aik * aji * akj - # val += aii * ajk * akj - - val = aii * (ajj * akk + ajk * akj) + aij * (aji * akk + ajk * aki) + aik * (aji * akj + ajj * aki) - # end # timeit_debug - - # @timeit_debug timer "Save in vector" begin - - # Only add non-zero values to the sparse matrix - if abs(val) > tol - # Threads.atomic_add!(k, 1) - # Threads.atomic_max!(k̄, k[]) - - if i1 == j1 - if i1 == k1 - divisor = 6 - else - divisor = 2 - end - else - if i1 ≠ k1 && j1 ≠ k1 - divisor = 1 - else - divisor = 2 - end - end - # push!(threadlocal[1],row) - # push!(threadlocal[2],col) - # push!(threadlocal[3],val / divisor) - # I[k[]] = row - # J[k[]] = col - # V[k[]] = val / divisor - - k += 1 - - if k > estimated_nnz - estimated_nnz += Int(ceil(max(1000, estimated_nnz * .1))) - estimated_nnz = min(m3_cols * m3_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 # timeit_debug - end - end - end - end - end - end + 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 - # end # timeit_debug - - # @timeit_debug timer "Resize" begin - - # out = map(fetch, threadlocal) - - # I = mapreduce(v -> v[1], vcat, out) - # J = mapreduce(v -> v[2], vcat, out) - # V = mapreduce(v -> v[3], vcat, out) + resize!(II, cnt) + resize!(JJ, cnt) + resize!(VV, cnt) - # # Resize the index and value arrays to the actual number of entries - # resize!(I, k̄[]) - # resize!(J, k̄[]) - # resize!(V, k̄[]) - resize!(I, k) - resize!(J, k) - resize!(V, k) - - # end # timeit_debug - # end # timeit_debug - - # Create the sparse matrix from the collected indices and values - if a_is_adjoint - klasttouch = sparse_preallocation[4] # Vector{Ti}(undef, n) - csrrowptr = sparse_preallocation[5] # Vector{Ti}(undef, m + 1) - csrcolval = sparse_preallocation[6] # Vector{Ti}(undef, length(I)) - csrnzval = sparse_preallocation[7] # Vector{Tv}(undef, length(I)) + # 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, m3_rows) - resize!(csrrowptr, m3_cols + 1) - resize!(csrcolval, length(J)) - resize!(csrnzval, length(J)) + resize!(klasttouch, m₃ᶜ) + resize!(csrrowptr, m₃ᵣ + 1) + resize!(csrcolval, length(II)) + resize!(csrnzval, length(II)) - out = sparse!(J, I, V, m3_cols, m3_rows, +, klasttouch, csrrowptr, csrcolval, csrnzval, J, I, V) - # out = sparse!(J, I, V, m3_cols, m3_rows) + out = if cnt >= m₃ᶜ + 1 + sparse!(II, JJ, VV, m₃ᵣ, m₃ᶜ, +, klasttouch, csrrowptr, csrcolval, csrnzval, II, JJ, VV) else - klasttouch = sparse_preallocation[4] # Vector{Ti}(undef, n) - csrrowptr = sparse_preallocation[5] # Vector{Ti}(undef, m + 1) - csrcolval = sparse_preallocation[6] # Vector{Ti}(undef, length(I)) - csrnzval = sparse_preallocation[7] # Vector{Tv}(undef, length(I)) - - resize!(klasttouch, m3_cols) - resize!(csrrowptr, m3_rows + 1) - resize!(csrcolval, length(I)) - resize!(csrnzval, length(I)) - - out = sparse!(I, J, V, m3_rows, m3_cols, +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) - # out = sparse!(I, J, V, m3_rows, m3_cols) + SparseArrays.sparse(II, JJ, VV, m₃ᵣ, m₃ᶜ) + end + + if tol > 0 + droptol!(out, tol) end return out end +=# + + # function kron³(A::AbstractSparseMatrix{T}, M₃::third_order) where T <: Real @@ -2160,7148 +1130,896 @@ end # return sparse!(result_rows, result_cols, result_vals, size(M₃.𝐂₃, 2), size(M₃.𝐔₃, 1)) # end -function A_mult_kron_power_3_B(A::AbstractSparseMatrix{R}, - B::Union{ℒ.Adjoint{T,Matrix{T}},DenseMatrix{T}}; - tol::AbstractFloat = eps()) where {R <: Real, T <: Real} - n_row = size(B,1) - n_col = size(B,2) - - vals = T[] - rows = Int[] - cols = Int[] - - Ar, Ac, Av = findnz(A) - - for row in unique(Ar) - idx_mat, vals_mat = A[row,:] |> findnz - - for col in 1:size(B,2)^3 - col_1, col_3 = divrem((col - 1) % (n_col^2), n_col) .+ 1 - col_2 = ((col - 1) ÷ (n_col^2)) + 1 - - mult_val = 0.0 - - for (i,idx) in enumerate(idx_mat) - i_1, i_3 = divrem((idx - 1) % (n_row^2), n_row) .+ 1 - i_2 = ((idx - 1) ÷ (n_row^2)) + 1 - @inbounds mult_val += vals_mat[i] * B[i_1,col_1] * B[i_2,col_2] * B[i_3,col_3] - end +# Dead code: A_mult_kron_power_3_B — never called anywhere +# function A_mult_kron_power_3_B(A::AbstractSparseMatrix{R}, +# B::Union{ℒ.Adjoint{T,Matrix{T}},DenseMatrix{T}}; +# tol::AbstractFloat = eps()) where {R <: Real, T <: Real} +# n_row = size(B,1) +# n_col = size(B,2) +# +# vals = T[] +# rows = Int[] +# cols = Int[] +# +# Ar, Ac, Av = findnz(A) +# +# for row in unique(Ar) +# idx_mat, vals_mat = A[row,:] |> findnz +# +# for col in 1:size(B,2)^3 +# col_1, col_3 = divrem((col - 1) % (n_col^2), n_col) .+ 1 +# col_2 = ((col - 1) ÷ (n_col^2)) + 1 +# +# mult_val = 0.0 +# +# for (i,idx) in enumerate(idx_mat) +# i_1, i_3 = divrem((idx - 1) % (n_row^2), n_row) .+ 1 +# i_2 = ((idx - 1) ÷ (n_row^2)) + 1 +# @inbounds mult_val += vals_mat[i] * B[i_1,col_1] * B[i_2,col_2] * B[i_3,col_3] +# end +# +# if abs(mult_val) > tol +# push!(vals,mult_val) +# push!(rows,row) +# push!(cols,col) +# end +# end +# end +# +# sparse(rows,cols,vals,size(A,1),size(B,2)^3) +# end - if abs(mult_val) > tol - push!(vals,mult_val) - push!(rows,row) - push!(cols,col) - end - end - end - sparse(rows,cols,vals,size(A,1),size(B,2)^3) -end +function get_and_check_observables(T::post_model_macro, data::KeyedArray{Float64})::Vector{Symbol} + @assert size(data,1) <= T.nExo "Cannot estimate model with more observables than exogenous shocks. Have at least as many shocks as observable variables." + observables = collect(axiskeys(data,1)) -function translate_symbol_to_ascii(x::Symbol) - ss = Unicode.normalize(replace(string(x), "◖" => "__", "◗" => "__"), :NFD) + @assert observables isa Vector{String} || observables isa Vector{Symbol} "Make sure that the data has variables names as rows. They can be either Strings or Symbols." - outstr = "" + observables_symbols = observables isa String_input ? observables .|> Meta.parse .|> replace_indices : observables - for i in ss - out = REPL.symbol_latex(string(i))[2:end] - if out == "" - outstr *= string(i) - else - outstr *= replace(out, - r"\!" => s"_", - r"\(" => s"_", - r"\)" => s"_", - r"\^" => s"_", - r"\_\^" => s"_", - r"\+" => s"plus", - r"\-" => s"minus", - r"\*" => s"times") - if i != ss[end] - outstr *= "_" - end - end - end + @assert length(setdiff(observables_symbols, T.var)) == 0 "The following symbols in the first axis of the conditions matrix are not part of the model: " * repr(setdiff(observables_symbols, T.var)) - return outstr + sort!(observables_symbols) + + return observables_symbols end +function x_kron_II!(buffer::Matrix{T}, x::Vector{T}) where T + n = length(x) + m = size(buffer,2) -function translate_expression_to_ascii(exp::Expr) - postwalk(x -> - x isa Symbol ? - begin - x_tmp = translate_symbol_to_ascii(x) - - if x_tmp == string(x) - x - else - Symbol(x_tmp) - end - end : - x, - exp) -end - + # @assert size(buffer, 1) == n^3 "Buffer must have n^2 rows." + # @assert size(buffer, 2) == n^2 "Buffer must have n columns." -function combine_pairs(v::Vector{Pair{Vector{Symbol}, Vector{Symbol}}}) - i = 1 - while i <= length(v) - subset_found = false - for j in i+1:length(v) - # Check if v[i].second and v[j].second are equal or if one is subset of the other - if v[i].second == v[j].second - # Exact match: combine first elements and remove duplicate - v[i] = v[i].first ∪ v[j].first => v[i].second - deleteat!(v, j) - subset_found = true - break - elseif all(elem -> elem in v[j].second, v[i].second) || all(elem -> elem in v[i].second, v[j].second) - # One is subset of the other: combine the first elements and assign to the one with the larger second element - if length(v[i].second) > length(v[j].second) - v[i] = v[i].first ∪ v[j].first => v[i].second - deleteat!(v, j) - else - v[j] = v[i].first ∪ v[j].first => v[j].second - deleteat!(v, i) - end - subset_found = true - break - end - end - # If no subset was found for v[i], move to the next element - if !subset_found - i += 1 + @inbounds for j in 1:m + for i in 1:n + buffer[(j - 1) * n + i, j] = x[i] end end - return v end -function determine_efficient_order(𝐒₁::Matrix{<: Real}, - constants::constants, - variables::Union{Symbol_input,String_input}; - covariance::Union{Symbol_input,String_input} = Symbol[], - tol::AbstractFloat = eps()) - - T = constants.post_model_macro - +# Dead code: bivariate_moment, product_moments, multiplicate, generateSumVectors — never called anywhere +# function bivariate_moment(moment::Vector{Int}, rho::Int)::Int +# if (moment[1] + moment[2]) % 2 == 1 +# return 0 +# end +# +# result = 1 +# coefficient = 1 +# odd_value = 2 * (moment[1] % 2) +# +# for j = 1:min(moment[1] ÷ 2, moment[2] ÷ 2) +# coefficient *= 2 * (moment[1] ÷ 2 + 1 - j) * (moment[2] ÷ 2 + 1 - j) * rho^2 / (j * (2 * j - 1 + odd_value)) +# result += coefficient +# end +# +# if odd_value == 2 +# result *= rho +# end +# +# result *= prod(1:2:moment[1]) * prod(1:2:moment[2]) +# +# return result +# end +# +# +# function product_moments(V, ii, nu)::Int +# s = sum(nu) +# +# if s == 0 +# return 1 +# elseif isodd(s) +# return 0 +# end +# +# mask = .!(nu .== 0) +# nu = nu[mask] +# ii = ii[mask] +# V = V[ii, ii] +# +# m, s2 = length(ii), s / 2 +# +# if m == 1 +# return (V^s2 * prod(1:2:s-1))[1] +# elseif m == 2 +# if V[1,1]==0 || V[2,2]==0 +# return 0 +# end +# rho = V[1, 2] / sqrt(V[1, 1] * V[2, 2]) +# return (V[1, 1]^(nu[1] / 2) * V[2, 2]^(nu[2] / 2) * bivariate_moment(nu, Int(rho)))[1] +# end +# +# inu = sortperm(nu, rev=true) +# +# sort!(nu, rev=true) +# +# V = V[inu, inu] +# +# x = zeros(Int, 1, m) +# V = V / 2 +# nu2 = nu' / 2 +# p = 2 +# q = nu2 * V * nu2' +# y = 0 +# +# for _ in 1:round(Int, prod(nu .+ 1) / 2) +# y += p * q^s2 +# for j in 1:m +# if x[j] < nu[j] +# x[j] += 1 +# p = -round(p * (nu[j] + 1 - x[j]) / x[j]) +# q -= (2 * (nu2 - x) * V[:, j] .+ V[j, j])[1] +# break +# else +# x[j] = 0 +# p = isodd(nu[j]) ? -p : p +# q += (2 * nu[j] * (nu2 - x) * V[:, j] .- nu[j]^2 * V[j, j])[1] +# end +# end +# end +# +# return y / prod(1:s2) +# end +# +# +# function multiplicate(p::Int, order::Int) +# # precompute p powers +# pⁿ = [p^i for i in 0:order-1] +# +# DP = spzeros(Bool, p^order, prod(p - 1 .+ (1:order)) ÷ factorial(order)) +# +# binom_p_ord = binomial(p + order - 1, order) +# +# # Initialize index and binomial arrays +# indexes = ones(Int, order) # Vector to hold current indexes +# binomials = zeros(Int, order) # Vector to hold binomial values +# +# # Helper function to handle the nested loops +# function loop(level::Int) +# for i=1:p +# indexes[level] = i +# binomials[level] = binomial(p + level - 1 - i, level) +# +# if level < order # If not at innermost loop yet, continue nesting +# loop(level + 1) +# else # At innermost loop, perform calculation +# n = sum((indexes[k] - 1) * pⁿ[k] for k in 1:order) +# m = binom_p_ord - sum(binomials[k] for k in 1:order) +# DP[n+1, m] = 1 # Arrays are 1-indexed in Julia +# end +# end +# end +# +# loop(1) # Start the recursive loop +# +# return DP +# end +# +# +# function generateSumVectors(vectorLength::Int, totalSum::Int)::Union{Vector{Int}, Vector{ℒ.Adjoint{Int, Vector{Int}}}} +# # Base case: if vectorLength is 1, return totalSum +# if vectorLength == 1 +# return [totalSum] +# end +# +# # Recursive case: generate all possible vectors for smaller values of vectorLength and totalSum +# return [[currentInt; smallerVector...]' for currentInt in totalSum:-1:0 for smallerVector in generateSumVectors(vectorLength-1, totalSum-currentInt)] +# end - orders = Pair{Vector{Symbol}, Vector{Symbol}}[] - nˢ = T.nPast_not_future_and_mixed - - if variables == :full_covar - return [T.var => T.past_not_future_and_mixed] - else - var_idx = MacroModelling.parse_variables_input_to_index(variables, constants) |> sort - observables = T.var[var_idx] - end - # Precompute state indices to avoid repeated indexin calls - state_idx_in_var = indexin(T.past_not_future_and_mixed, T.var) .|> Int - 𝐒₁_states = 𝐒₁[state_idx_in_var, 1:nˢ] - for obs in observables - obs_in_var_idx = indexin([obs],T.var) .|> Int - dependencies_in_states = vec(sum(abs, 𝐒₁[obs_in_var_idx,1:nˢ], dims=1) .> tol) .> 0 - # Iterative propagation without redundant allocations - while true - new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * 𝐒₁_states) .> tol) - if new_deps == dependencies_in_states - break - end - dependencies_in_states = new_deps - end - dependencies = T.past_not_future_and_mixed[dependencies_in_states] - push!(orders,[obs] => sort(dependencies)) +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 covariance variables are specified, compute dependencies and add entries for those pairs - if !(covariance == Symbol[]) - covar_var_idx = MacroModelling.parse_variables_input_to_index(covariance, constants) |> sort - covariance_vars = T.var[covar_var_idx] - - # Compute dependencies for covariance variables (if not already computed) - for covar_var in covariance_vars - # 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 - dependencies_in_states = vec(sum(abs, 𝐒₁[obs_in_var_idx,1:nˢ], dims=1) .> tol) .> 0 - - # Iterative propagation without redundant allocations - while true - new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * 𝐒₁_states) .> tol) - if new_deps == dependencies_in_states - break - end - dependencies_in_states = new_deps - end - - dependencies = T.past_not_future_and_mixed[dependencies_in_states] - push!(orders,[covar_var] => sort(dependencies)) - end - end - - # Build lookup dictionary for faster searches - var_to_idx = Dict{Symbol, Int}() - for (idx, order) in enumerate(orders) - for var in order.first - var_to_idx[var] = idx - end - end - - # Add entries for all pairs of covariance variables - for i in 1:length(covariance_vars) - for j in (i+1):length(covariance_vars) - # Find dependencies for both variables using lookup dictionary - idx_i = var_to_idx[covariance_vars[i]] - idx_j = var_to_idx[covariance_vars[j]] - - deps_i = orders[idx_i].second - deps_j = orders[idx_j].second - # Union of dependencies for covariance computation - combined_deps = sort(union(deps_i, deps_j)) - push!(orders, [covariance_vars[i], covariance_vars[j]] => combined_deps) - end + if !(𝓂.functions.NSSS_custom isa Function) + if !silent + print("Find non-stochastic steady state:\t\t\t\t\t") end end - - sort!(orders, by = x -> length(x[2]), rev = true) - - return combine_pairs(orders) -end - - -function determine_efficient_order(𝐒₁::Matrix{<: Real}, - 𝐒₂::AbstractMatrix{<: Real}, - constants::constants, - variables::Union{Symbol_input,String_input}; - covariance::Union{Symbol_input,String_input} = Symbol[], - tol::AbstractFloat = eps()) - - T = constants.post_model_macro - - - orders = Pair{Vector{Symbol}, Vector{Symbol}}[] - - nˢ = T.nPast_not_future_and_mixed - nᵉ = T.nExo - if variables == :full_covar - return [T.var => T.past_not_future_and_mixed] - else - var_idx = MacroModelling.parse_variables_input_to_index(variables, constants) |> sort - observables = T.var[var_idx] - end - - # Build selector for state variables in the augmented state vector [states; 1; shocks] - s_in_s⁺ = BitVector(vcat(ones(Bool, nˢ), zeros(Bool, nᵉ + 1))) + 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 - # Kronecker product indices for state-state interactions - kron_s_s = ℒ.kron(s_in_s⁺, s_in_s⁺) + found_solution = true - # Precompute state indices and matrix slices to avoid repeated operations - 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 - - for obs in observables - obs_in_var_idx = indexin([obs],T.var) .|> Int - - # First order dependencies - dependencies_in_states = vec(sum(abs, 𝐒₁[obs_in_var_idx,1:nˢ], dims=1) .> tol) .> 0 + if !(𝓂.functions.NSSS_custom isa Function) + select_fastest_SS_solver_parameters!(𝓂, tol = opts.tol) - # Second order dependencies from quadratic terms (s ⊗ s) - if nnz(𝐒₂) > 0 - s_s_to_y₂ = 𝐒₂[obs_in_var_idx, kron_s_s] - - # Check which state variable pairs have influence - # 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 - - # Propagate dependencies through the system (iterative closure) - # considering both first and second 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 state i and state j are dependencies, - # their product can affect states - if !isnothing(𝐒₂_states) - # Generate selector vector for columns where both states are dependencies - selector = vec(ℒ.kron(prev_dependencies, prev_dependencies)) - if any(selector) - # Check which states are affected by the selected products - affected = vec(sum(abs, 𝐒₂_states[:, selector], dims=2) .> tol) - new_deps = new_deps .| affected - end - end + 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 new_deps == dependencies_in_states - break + if found_solution + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(𝓂, 𝓂.parameter_values, opts = opts, cold_start = true) end - dependencies_in_states = new_deps end - - dependencies = T.past_not_future_and_mixed[dependencies_in_states] - - push!(orders,[obs] => sort(dependencies)) end - # If covariance variables are specified, compute dependencies and add entries for those pairs - if !(covariance == Symbol[]) - covar_var_idx = MacroModelling.parse_variables_input_to_index(covariance, constants) |> sort - covariance_vars = T.var[covar_var_idx] - - # Compute dependencies for covariance variables (if not already computed) - for covar_var in covariance_vars - # 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 - - # 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 - - 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 - - if new_deps == dependencies_in_states - break - end - dependencies_in_states = new_deps - end - - dependencies = T.past_not_future_and_mixed[dependencies_in_states] - push!(orders,[covar_var] => sort(dependencies)) - end - end - - # Add entries for all pairs of covariance variables - for i in 1:length(covariance_vars) - for j in (i+1):length(covariance_vars) - # Find dependencies for both variables (they should exist now) - idx_i = findfirst(x -> covariance_vars[i] in x.first, orders) - idx_j = findfirst(x -> covariance_vars[j] in x.first, orders) - - deps_i = orders[idx_i].second - deps_j = orders[idx_j].second - # Union of dependencies for covariance computation - combined_deps = sort(union(deps_i, deps_j)) - push!(orders, [covariance_vars[i], covariance_vars[j]] => combined_deps) - 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 + + return SS_and_pars, solution_error, found_solution +end - sort!(orders, by = x -> length(x[2]), rev = true) +# Centralised helper to write symbolic derivatives and map functions - return combine_pairs(orders) -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 + σ = 5 + log_lik -= -log(σ * sqrt(2 * π)) - (pars[20]^2 / (2 * σ^2)) # logpdf of a normal dist with mean = 0 and variance = 5^2 -function determine_efficient_order(𝐒₁::Matrix{<: Real}, - 𝐒₂::AbstractMatrix{<: Real}, - 𝐒₃::AbstractMatrix{<: Real}, - constants::constants, - variables::Union{Symbol_input,String_input}; - covariance::Union{Symbol_input,String_input} = Symbol[], - tol::AbstractFloat = eps()) + pars[1:2] = sort(pars[1:2], rev = true) - T = constants.post_model_macro - + par_inputs = solver_parameters(pars..., 1, 0.0, 2) - orders = Pair{Vector{Symbol}, Vector{Symbol}}[] + reset_nsss_solver_cache!(𝓂) - nˢ = T.nPast_not_future_and_mixed - nᵉ = T.nExo - - if variables == :full_covar - return [T.var => T.past_not_future_and_mixed] - else - var_idx = MacroModelling.parse_variables_input_to_index(variables, constants) |> sort - observables = T.var[var_idx] - end + runtime = @elapsed outmodel = try solve_nsss_wrapper(𝓂.parameter_values, 𝓂, tol, false, true, [par_inputs]) catch end - # Build selectors for state variables in the augmented state vector [states; 1; shocks] - s_in_s⁺ = BitVector(vcat(ones(Bool, nˢ), zeros(Bool, nᵉ + 1))) - - # Kronecker product indices for interactions - kron_s_s = ℒ.kron(s_in_s⁺, s_in_s⁺) - kron_s_s_s = ℒ.kron(kron_s_s, s_in_s⁺) - - # 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 - - for obs in observables - obs_in_var_idx = indexin([obs],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 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 - 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 + runtime = outmodel isa Tuple{Vector{Float64}, Tuple{Float64, Int64}} ? + (outmodel[2][1] > tol.nsss.acceptance_tol) || !isfinite(outmodel[2][1]) ? + 10 : + runtime : + 10 - # 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 + return log_lik / 1e4 + runtime * 1e3 +end - dependencies = T.past_not_future_and_mixed[dependencies_in_states] +""" + find_SS_solver_parameters!(::Val{:ESCH}, 𝓂::ℳ; maxtime::Real = 120, maxiter::Int = 2500000, tol::Tolerances = Tolerances(), verbosity = 0) - push!(orders,[obs] => sort(dependencies)) - end - - # If covariance variables are specified, compute dependencies and add entries for those pairs - if !(covariance == Symbol[]) - covar_var_idx = MacroModelling.parse_variables_input_to_index(covariance, constants) |> sort - covariance_vars = T.var[covar_var_idx] - - # Compute dependencies for covariance variables (if not already computed) - for covar_var in covariance_vars - # 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 +Find optimal steady state solver parameters using NLopt's ESCH algorithm. - # 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 +This function optimizes solver parameters to minimize runtime while maintaining solver accuracy. +It uses the ESCH global optimization algorithm from the NLopt package. - dependencies = T.past_not_future_and_mixed[dependencies_in_states] - push!(orders,[covar_var] => sort(dependencies)) - end - end - - # Add entries for all pairs of covariance variables - for i in 1:length(covariance_vars) - for j in (i+1):length(covariance_vars) - # Find dependencies for both variables (they should exist now) - idx_i = findfirst(x -> covariance_vars[i] in x.first, orders) - idx_j = findfirst(x -> covariance_vars[j] in x.first, orders) - - deps_i = orders[idx_i].second - deps_j = orders[idx_j].second - # Union of dependencies for covariance computation - combined_deps = sort(union(deps_i, deps_j)) - push!(orders, [covariance_vars[i], covariance_vars[j]] => combined_deps) - end - end - end +# Arguments +- `𝓂`: Model structure +- `maxtime`: Maximum time in seconds for optimization +- `maxiter`: Maximum number of iterations +- `tol`: Tolerance structure +- `verbosity`: Verbosity level for output +""" +function find_SS_solver_parameters!(::Val{:ESCH}, 𝓂::ℳ; maxtime::Real = 120, maxiter::Int = 2500000, tol::Tolerances = Tolerances(), verbosity = 0) + pars = rand(20) .+ 1 + pars[20] -= 1 - sort!(orders, by = x -> length(x[2]), rev = true) + lbs = fill(eps(), length(pars)) + lbs[20] = -20 - return combine_pairs(orders) -end + ubs = fill(100.0, length(pars)) + + opt = NLopt.Opt(NLopt.:GN_ESCH, length(pars)) + opt.min_objective = (x,p) -> calculate_SS_solver_runtime_and_loglikelihood(x, 𝓂, tol = tol) -function get_and_check_observables(T::post_model_macro, data::KeyedArray{Float64})::Vector{Symbol} - @assert size(data,1) <= T.nExo "Cannot estimate model with more observables than exogenous shocks. Have at least as many shocks as observable variables." + NLopt.lower_bounds!(opt, lbs) + NLopt.upper_bounds!(opt, ubs) - observables = collect(axiskeys(data,1)) + opt.xtol_abs = eps(Float32) + opt.ftol_abs = eps(Float32) - @assert observables isa Vector{String} || observables isa Vector{Symbol} "Make sure that the data has variables names as rows. They can be either Strings or Symbols." + # opt.maxeval = maxiter + opt.maxtime = maxtime - observables_symbols = observables isa String_input ? observables .|> Meta.parse .|> replace_indices : observables + (minf,x,ret) = NLopt.optimize(opt, pars) - @assert length(setdiff(observables_symbols, T.var)) == 0 "The following symbols in the first axis of the conditions matrix are not part of the model: " * repr(setdiff(observables_symbols, T.var)) + verbosity > 0 && @info "NLopt return code: $ret" - sort!(observables_symbols) - - return observables_symbols -end + pars = x -function x_kron_II!(buffer::Matrix{T}, x::Vector{T}) where T - n = length(x) - m = size(buffer,2) + par_inputs = solver_parameters(pars..., 1, 0.0, 2) - # @assert size(buffer, 1) == n^3 "Buffer must have n^2 rows." - # @assert size(buffer, 2) == n^2 "Buffer must have n columns." + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(𝓂.parameter_values, 𝓂, tol, false, true, [par_inputs]) - @turbo for j in 1:m - for i in 1:n - buffer[(j - 1) * n + i, j] = x[i] - end + 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 end end -function bivariate_moment(moment::Vector{Int}, rho::Int)::Int - if (moment[1] + moment[2]) % 2 == 1 - return 0 - end - - result = 1 - coefficient = 1 - odd_value = 2 * (moment[1] % 2) - for j = 1:min(moment[1] ÷ 2, moment[2] ÷ 2) - coefficient *= 2 * (moment[1] ÷ 2 + 1 - j) * (moment[2] ÷ 2 + 1 - j) * rho^2 / (j * (2 * j - 1 + odd_value)) - result += coefficient - end +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." - if odd_value == 2 - result *= rho - end + best_idx = 1 + best_score = Inf - result *= prod(1:2:moment[1]) * prod(1:2:moment[2]) + solved = false - return result -end + solved_NSSS = 𝓂.caches.solver[end] + for (i_param, p) in enumerate(DEFAULT_SOLVER_PARAMETERS) + times = Vector{Float64}(undef, n_samples) + valid = true + + for i in 1:n_samples + start_time = time() -function product_moments(V, ii, nu)::Int - s = sum(nu) + reset_nsss_solver_cache!(𝓂) - if s == 0 - return 1 - elseif isodd(s) - return 0 - end + SS_and_pars, (solution_error, iters) = solve_nsss_wrapper(𝓂.parameter_values, 𝓂, tol, false, true, [p]) - mask = .!(nu .== 0) - nu = nu[mask] - ii = ii[mask] - V = V[ii, ii] + elapsed_time = time() - start_time - m, s2 = length(ii), s / 2 + times[i] = elapsed_time - if m == 1 - return (V^s2 * prod(1:2:s-1))[1] - elseif m == 2 - if V[1,1]==0 || V[2,2]==0 - return 0 + if solution_error > tol.nsss.acceptance_tol + valid = false + break + end end - rho = V[1, 2] / sqrt(V[1, 1] * V[2, 2]) - return (V[1, 1]^(nu[1] / 2) * V[2, 2]^(nu[2] / 2) * bivariate_moment(nu, Int(rho)))[1] - end - inu = sortperm(nu, rev=true) + if valid + sort!(times) + score = times[n_samples ÷ 2] - sort!(nu, rev=true) + if !isfinite(best_score) || score < best_score + best_score = score + best_idx = i_param + end - V = V[inu, inu] + solved = true + end + end - x = zeros(Int, 1, m) - V = V / 2 - nu2 = nu' / 2 - p = 2 - q = nu2 * V * nu2' - y = 0 + empty!(𝓂.caches.solver) + push!(𝓂.caches.solver, solved_NSSS) - for _ in 1:round(Int, prod(nu .+ 1) / 2) - y += p * q^s2 - for j in 1:m - if x[j] < nu[j] - x[j] += 1 - p = -round(p * (nu[j] + 1 - x[j]) / x[j]) - q -= (2 * (nu2 - x) * V[:, j] .+ V[j, j])[1] - break - else - x[j] = 0 - p = isodd(nu[j]) ? -p : p - q += (2 * nu[j] * (nu2 - x) * V[:, j] .- nu[j]^2 * V[j, j])[1] - end - end + if solved + 𝓂.constants.post_complete_parameters = update_post_complete_parameters( + 𝓂.constants.post_complete_parameters; + nsss_fastest_solver_parameter_idx = best_idx, + ) end - return y / prod(1:s2) + return nothing end -function multiplicate(p::Int, order::Int) - # precompute p powers - pⁿ = [p^i for i in 0:order-1] - - DP = spzeros(Bool, p^order, prod(p - 1 .+ (1:order)) ÷ factorial(order)) - binom_p_ord = binomial(p + order - 1, order) - # Initialize index and binomial arrays - indexes = ones(Int, order) # Vector to hold current indexes - binomials = zeros(Int, order) # Vector to hold binomial values - # Helper function to handle the nested loops - function loop(level::Int) - for i=1:p - indexes[level] = i - binomials[level] = binomial(p + level - 1 - i, level) - if level < order # If not at innermost loop yet, continue nesting - loop(level + 1) - else # At innermost loop, perform calculation - n = sum((indexes[k] - 1) * pⁿ[k] for k in 1:order) - m = binom_p_ord - sum(binomials[k] for k in 1:order) - DP[n+1, m] = 1 # Arrays are 1-indexed in Julia - end - end - end - - loop(1) # Start the recursive loop +function solve!(𝓂::ℳ; + parameters::ParameterType = nothing, + steady_state_function::SteadyStateFunctionType = missing, + dynamics::Bool = false, + algorithm::Symbol = :first_order, + opts::CalculationOptions = merge_calculation_options(), + obc::Bool = false, + silent::Bool = false) #, + # quadratic_matrix_equation_algorithm::Symbol = :schur, + # verbose::Bool = false, + # timer::TimerOutput = TimerOutput(), + # tol::AbstractFloat = 1e-12) - return DP -end + @assert algorithm ∈ all_available_algorithms + + # Initialize constants at entry point + constants = initialise_constants!(𝓂) + + # Handle steady_state_function argument + set_custom_steady_state_function!(𝓂, steady_state_function) + + # @timeit_debug timer "Write parameter inputs" begin + write_parameters_input!(𝓂, parameters, verbose = opts.verbose) + + if 𝓂.functions.functions_written && + isnothing(𝓂.functions.NSSS_custom) && + 𝓂.constants.nsss_solver.n_steps == 0 -function generateSumVectors(vectorLength::Int, totalSum::Int)::Union{Vector{Int}, Vector{ℒ.Adjoint{Int, Vector{Int}}}} - # Base case: if vectorLength is 1, return totalSum - if vectorLength == 1 - return [totalSum] + 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 + verbose = opts.verbose + + perturbation_order = 1 - # Recursive case: generate all possible vectors for smaller values of vectorLength and totalSum - return [[currentInt; smallerVector...]' for currentInt in totalSum:-1:0 for smallerVector in generateSumVectors(vectorLength-1, totalSum-currentInt)] -end - - -function match_pattern(strings::Union{Set,Vector}, pattern::Regex) - return filter(r -> match(pattern, string(r)) !== nothing, strings) -end - - -function count_ops(expr)::Int - op_count = 0 - postwalk(x -> begin - if x isa Expr && x.head == :call - op_count += 1 - end - x - end, expr) - return op_count -end - -# try: run optim only if there is a violation / capture case with small shocks and set them to zero -function parse_occasionally_binding_constraints(equations_block; max_obc_horizon::Int = 40, avoid_solve::Bool = false) - # precision_factor = 1e #factor to force the optimiser to have non-relevatn shocks at zero - - eqs = [] - obc_shocks = Expr[] - - for arg in equations_block.args - if isa(arg,Expr) - if check_for_minmax(arg) - arg_trans = transform_obc(arg) - else - arg_trans = arg - end - - eq = postwalk(x -> - x isa Expr ? - x.head == :call ? - x.args[1] == :max ? - begin - - obc_vars_left = Expr(:ref, Meta.parse("χᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝˡ" ), 0) - obc_vars_right = Expr(:ref, Meta.parse("χᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝʳ" ), 0) - - if !(x.args[2] isa Symbol) && check_for_dynamic_variables(x.args[2]) - push!(eqs, :($obc_vars_left = $(x.args[2]))) - else - obc_vars_left = x.args[2] - end + 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, + 𝓂.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) - if !(x.args[3] isa Symbol) && check_for_dynamic_variables(x.args[3]) - push!(eqs, :($obc_vars_right = $(x.args[3]))) - else - obc_vars_right = x.args[3] - end + 𝓂.functions.functions_written = true + end - obc_inequality = Expr(:ref, Meta.parse("Χᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ" ), 0) + # Check for missing parameters after processing input + if !isempty(𝓂.constants.post_complete_parameters.missing_parameters) + error("Cannot solve model: missing parameter values for $(𝓂.constants.post_complete_parameters.missing_parameters). Provide them via the `parameters` keyword argument (e.g., `parameters = [:α => 0.3, :β => 0.99]`).") + end - push!(eqs, :($obc_inequality = $(Expr(x.head, x.args[1], obc_vars_left, obc_vars_right)))) + # end # timeit_debug - obc_shock = Expr(:ref, Meta.parse("ϵᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ"), 0) + if 𝓂.constants.second_order.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && + algorithm ∈ [:second_order, :pruned_second_order] + start_time = time() + if !silent print("Take symbolic derivatives up to second order:\t\t\t\t") end + write_functions_mapping!(𝓂, 2) + if !silent println(round(time() - start_time, digits = 3), " seconds") end + elseif 𝓂.constants.third_order.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && algorithm ∈ [:third_order, :pruned_third_order] + start_time = time() + if !silent print("Take symbolic derivatives up to third order:\t\t\t\t") end + write_functions_mapping!(𝓂, 3) + if !silent println(round(time() - start_time, digits = 3), " seconds") end + end - push!(obc_shocks, obc_shock) + if dynamics + if algorithm == :first_order + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(𝓂, 𝓂.parameter_values, opts = opts) - :($obc_inequality - $obc_shock) - end : - x.args[1] == :min ? - begin - obc_vars_left = Expr(:ref, Meta.parse("χᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝˡ" ), 0) - obc_vars_right = Expr(:ref, Meta.parse("χᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝʳ" ), 0) + @assert solution_error < opts.tol.nsss.acceptance_tol "Could not find non-stochastic steady state." - if !(x.args[2] isa Symbol) && check_for_dynamic_variables(x.args[2]) - push!(eqs, :($obc_vars_left = $(x.args[2]))) - else - obc_vars_left = x.args[2] - end + ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂.caches, 𝓂.functions.jacobian, 𝓂.workspaces) - if !(x.args[3] isa Symbol) && check_for_dynamic_variables(x.args[3]) - push!(eqs, :($obc_vars_right = $(x.args[3]))) - else - obc_vars_right = x.args[3] - end + S₁, qme_sol, solved = calculate_first_order_solution(∇₁, + constants, + 𝓂.workspaces, + 𝓂.caches; + opts = opts, + initial_guess = 𝓂.caches.qme_solution, + parameter_values = 𝓂.parameter_values) - obc_inequality = Expr(:ref, Meta.parse("Χᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ" ), 0) + update_perturbation_counter!(𝓂.counters, solved, order = 1) - push!(eqs, :($obc_inequality = $(Expr(x.head, x.args[1], obc_vars_left, obc_vars_right)))) + @assert solved "Could not find stable first order solution." - obc_shock = Expr(:ref, Meta.parse("ϵᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ"), 0) + 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 - push!(obc_shocks, obc_shock) + 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 - :($obc_inequality - $obc_shock) - end : - x : - x : - x, - arg_trans) + 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 - push!(eqs, eq) + 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 - end - for obc in obc_shocks - # push!(eqs, :($(obc) = $(Expr(:ref, obc.args[1], -1)) * 0.3 + $(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(max_obc_horizon)) * "⁾"), 0)))) - push!(eqs, :($(obc) = $(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(max_obc_horizon)) * "⁾"), 0)))) - - push!(eqs, :($(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻⁰⁾"), 0)) = activeᵒᵇᶜshocks * $(Expr(:ref, Meta.parse(string(obc.args[1]) * "⁽" * super(string(max_obc_horizon)) * "⁾"), :x)))) - - for i in 1:max_obc_horizon - push!(eqs, :($(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(i)) * "⁾"), 0)) = $(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(i-1)) * "⁾"), -1)) + activeᵒᵇᶜshocks * $(Expr(:ref, Meta.parse(string(obc.args[1]) * "⁽" * super(string(max_obc_horizon-i)) * "⁾"), :x)))) + if obc + calculate_first_order_obc_solution!(𝓂, constants, opts) end - end - return Expr(:block, eqs...) + end + + return nothing 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) - full_NSSS = ms.full_NSSS_display - - relevant_SS = get_steady_state(𝓂, algorithm = algorithm, - stochastic = algorithm != :first_order, - return_variables_only = true, - derivatives = false, - verbose = opts.verbose, - tol = opts.tol, - quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) - - reference_steady_state = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : relevant_SS(s) for s in full_NSSS] - - relevant_NSSS = get_steady_state(𝓂, algorithm = :first_order, - stochastic = false, - return_variables_only = true, - derivatives = false, - verbose = opts.verbose, - tol = opts.tol, - quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, - sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) - - NSSS = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : relevant_NSSS(s) for s in full_NSSS] - SSS_delta = NSSS - reference_steady_state - - return reference_steady_state, NSSS, SSS_delta -end - -# compatibility with SymPy -Max = max -Min = min +write_parameters_input!(𝓂::ℳ, parameters::Nothing; verbose::Bool = true) = return parameters +write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Pair{S,Float64}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}(parameters[1] |> Meta.parse |> replace_indices => parameters[2]), verbose = verbose) -function simplify(ex::Expr)::Union{Expr,Symbol,Int} - ex_ss = convert_to_ss_equation(ex) - for x in get_symbols(ex_ss) - sym_value = SPyPyC.symbols(string(x), real = true, finite = true) - Core.eval(SymPyWorkspace, :($x = $sym_value)) - end - parsed = ex_ss |> x -> Core.eval(SymPyWorkspace, x) |> string |> Meta.parse +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Float64},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict(parameters), verbose = verbose) +# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Union{Symbol,AbstractString},Union{Float64,Int}},Vararg{Pair{Union{Symbol,AbstractString},Union{Float64,Int}}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{AbstractString,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{S,Float64},Vararg{Pair{S,Float64}}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]) +, verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol, Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{S, Float64}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol, Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Dict{S,Float64}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((keys(parameters) .|> Meta.parse .|> replace_indices) .=> values(parameters)), verbose = verbose) - postwalk(x -> x isa Expr ? - x.args[1] == :conjugate ? - x.args[2] : - x : - x, parsed) -end -function convert_to_ss_equation(eq::Expr)::Expr - 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 : - x.args[1] : - x.head == :call ? - x.args[1] == :* ? - x.args[2] isa Int ? - x.args[3] isa Int ? - x : - :($(x.args[3]) * $(x.args[2])) : # avoid 2X syntax. doesn't work with sympy - x : - x : - unblock(x) : - x, - eq) -end +write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(parameters[1])) => parameters[2]]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Pair{S,Int}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{Symbol,Int}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{S,Int},Vararg{Pair{S,Int}}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{S, Int}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Dict{S,Int}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((keys(parameters) .|> Meta.parse .|> replace_indices) .=> values(parameters)), verbose = verbose) -function resolve_if_expr(ex::Expr) - prewalk(ex) do node - if node isa Expr && (node.head === :if || node.head === :elseif) - cond = node.args[1] - then_blk = node.args[2] - if length(node.args) == 3 - else_blk = node.args[3] - end - val = evaluate_conditions(unblock(cond)) - - if val === true - # recurse into the selected branch - return resolve_if_expr(unblock(then_blk)) - elseif val === false && length(node.args) == 3 - return resolve_if_expr(unblock(else_blk)) - elseif val === false && length(node.args) == 2 - return nothing - elseif val === false && node.head === :elseif - return resolve_if_expr(unblock(else_blk)) - end - end - return node - end -end -# function remove_nothing(ex::Expr) -# postwalk(ex) do node -# # Only consider call-nodes with exactly two arguments -# if node isa Expr && node.head === :call && length(node.args) == 3 -# fn, lhs, rhs = node.args -# lhs2 = unblock(lhs) -# rhs2 = unblock(rhs) - -# if rhs2 === :(nothing) -# # strip the call and recurse to clean deeper -# return remove_nothing(lhs2) -# elseif lhs2 === :(nothing) -# return remove_nothing(rhs2) -# # else -# # return remove_nothing(node.args) -# end -# end -# return node -# end -# end +write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(parameters[1])) => parameters[2]]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Pair{S,Real}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Dict{S,Real}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((keys(parameters) .|> Meta.parse .|> replace_indices) .=> values(parameters)), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Real},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{S,Real},Vararg{Pair{S,Float64}}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{S, Real}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -end # dispatch_doctor -function evaluate_conditions(cond) - if cond isa Bool - return cond - elseif cond isa Expr && cond.head == :call - a, b = cond.args[2], cond.args[3] - if typeof(a) ∉ [Symbol, Number] - a = eval(a) +function write_parameters_input!(𝓂::ℳ, parameters::D; verbose::Bool = true) where D <: AbstractDict{Symbol,Float64} + # Handle missing parameters - add them if they are in the missing_parameters list + p = 𝓂.constants.post_complete_parameters + missing_params_provided = intersect(collect(keys(parameters)), p.missing_parameters) + + if !isempty(missing_params_provided) + + # Remove the provided missing params from the missing list + remaining_missing = setdiff(p.missing_parameters, missing_params_provided) + + # If all missing parameters are now provided, print a message + if !isempty(remaining_missing) + @info "Remaining missing parameters: ", remaining_missing end - if typeof(b) ∉ [Symbol, Number] - b = eval(b) - end + # Amend parameter order by provided missing params + # declared_params = parameters that were never missing (have non-NaN values) + # We identify them as parameters that are not in the union of missing_params_provided and still-missing params + all_missing = union(missing_params_provided, remaining_missing) + declared_params = setdiff(p.parameters, all_missing) + + # Get the current parameter values for declared params + declared_param_indices = indexin(declared_params, p.parameters) + declared_values = 𝓂.parameter_values[declared_param_indices] + + # Get values for the newly provided missing params (currently NaN in parameter_values) + # We'll set them later after the bounds check + missing_values = fill(NaN, length(missing_params_provided)) - if cond.args[1] == :(==) - return a == b - elseif cond.args[1] == :(!=) - return a != b - elseif cond.args[1] == :(<) - return a < b - elseif cond.args[1] == :(<=) - return a <= b - elseif cond.args[1] == :(>) - return a > b - elseif cond.args[1] == :(>=) - return a >= b + # Get values for the remaining missing params (still NaN) + remaining_missing_values = fill(NaN, length(remaining_missing)) + + # Reorder both parameters and parameter_values arrays + new_parameters = vcat(declared_params, collect(missing_params_provided), remaining_missing) + 𝓂.constants.post_complete_parameters = update_post_complete_parameters( + p; + parameters = new_parameters, + missing_parameters = remaining_missing, + ) + 𝓂.parameter_values = vcat(declared_values, missing_values, remaining_missing_values) + + # 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 end - return nothing -end - -function contains_equation(expr) - found = false - postwalk(expr) do x - if x isa Expr && x.head == :(=) - found = true + + # Handle remaining parameters (not missing ones) + p = 𝓂.constants.post_complete_parameters + if length(setdiff(collect(keys(parameters)), p.parameters))>0 + @warn("Parameters not part of the model are ignored: $(setdiff(collect(keys(parameters)),p.parameters))") + for kk in setdiff(collect(keys(parameters)), p.parameters) + delete!(parameters,kk) end - return x end - return found -end -function remove_nothing(ex::Expr) - postwalk(ex) do node - # Only consider call-expressions - if node isa Expr && node.head === :call && any(node.args .=== nothing) - fn = node.args[1] - # Unblock and collect all the operands - # raw_args = map(arg -> unblock(arg), node.args[2:end]) - # Drop any nothing - kept = filter(arg -> !(unblock(arg) === nothing), node.args[2:end]) - if isempty(kept) - return nothing - elseif length(kept) == 1 - return kept[1] - else - # elseif length(kept) < length(raw_args) - return Expr(:call, fn, kept...) - # else - # return node + bounds_broken = false + + for (par,val) in parameters + 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") + bounds_broken = true + continue + end + if val < 𝓂.constants.post_parameters_macro.bounds[par][1] + @warn("Calibration is out of bounds for $par > $(𝓂.constants.post_parameters_macro.bounds[par][1])\t parameter value: $val") + bounds_broken = true + continue end end - return node end -end - -@stable default_mode = "disable" begin -function replace_indices_inside_for_loop(exxpr,index_variable,indices,concatenate, operator) - @assert operator ∈ [:+,:*] "Only :+ and :* allowed as operators in for loops." - calls = [] - indices = indices.args[1] == :(:) ? eval(indices) : [indices.args...] - for idx in indices - push!(calls, postwalk(x -> begin - x isa Expr ? - x.head == :ref ? - @capture(x, name_{index_}[time_]) ? - index == index_variable ? - :($(Expr(:ref, Symbol(string(name) * "{" * string(idx) * "}"),time))) : - time isa Expr || time isa Symbol ? - index_variable ∈ get_symbols(time) ? - :($(Expr(:ref, Expr(:curly,name,index), Meta.parse(replace(string(time), string(index_variable) => idx))))) : - x : - x : - @capture(x, name_[time_]) ? - time isa Expr || time isa Symbol ? - index_variable ∈ get_symbols(time) ? - :($(Expr(:ref, name, Meta.parse(replace(string(time), string(index_variable) => idx))))) : - # occursin("{" * string(index_variable) * "}", string(name)) ? - # Expr(:ref, Symbol(replace(string(name), "{" * string(index_variable) * "}" => "◖" * string(idx) * "◗")), time) : - x : - # occursin("{" * string(index_variable) * "}", string(name)) ? - # Expr(:ref, Symbol(replace(string(name), "{" * string(index_variable) * "}" => "◖" * string(idx) * "◗")), time) : - x : - x : - x.head == :if ? - length(x.args) > 2 ? - Expr(:if, postwalk(x -> x == index_variable ? idx : x, x.args[1]), - replace_indices_inside_for_loop(x.args[2],index_variable,:([$idx]),false,:+) |> unblock, - replace_indices_inside_for_loop(x.args[3],index_variable,:([$idx]),false,:+) |> unblock) : - Expr(:if, postwalk(x -> x == index_variable ? idx : x, x.args[1]), - replace_indices_inside_for_loop(x.args[2],index_variable,:([$idx]),false,:+) |> unblock) : - @capture(x, name_{index_}) ? - index == index_variable ? - :($(Symbol(string(name) * "{" * string(idx) * "}"))) : - x : - x : - @capture(x, name_) ? - name == index_variable && idx isa Int ? - :($idx) : - x isa Symbol ? - occursin("{" * string(index_variable) * "}", string(x)) ? - Symbol(replace(string(x), "{" * string(index_variable) * "}" => "{" * string(idx) * "}")) : - x : - x : - x - end, - exxpr)) - end - - if concatenate - return :($(Expr(:call, operator, calls...))) + if bounds_broken + @warn("Parameters unchanged.") else - return :($(Expr(:block, calls...))) - # return :($calls...) - # return calls - end -end + ntrsct_idx = map(x-> getindex(1:length(𝓂.parameter_values), p.parameters .== x)[1], collect(keys(parameters))) + # ntrsct_idx = indexin(collect(keys(parameters)), p.parameters) + + if !all(𝓂.parameter_values[ntrsct_idx] .== collect(values(parameters))) && !(p.parameters[ntrsct_idx] == [:activeᵒᵇᶜshocks]) + if verbose println("Parameter changes: ") end + end + + for i in 1:length(parameters) + if 𝓂.parameter_values[ntrsct_idx[i]] != collect(values(parameters))[i] + 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] + end + end + end -replace_indices(x::Symbol) = x + return nothing +end -replace_indices_special(x::Symbol) = x -replace_indices(x::String) = Symbol(replace(x, "{" => "◖", "}" => "◗")) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Int,Vararg{Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Matrix{Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) -replace_indices_in_symbol(x::Symbol) = replace(string(x), "◖" => "{", "◗" => "}") +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Float64,Vararg{Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, vec(collect(parameters)), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Matrix{Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, vec(collect(parameters)), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Tuple{Real,Vararg{Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) +write_parameters_input!(𝓂::ℳ, parameters::Matrix{Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) -""" - apply_custom_name(symbol::Symbol, custom_names::Dict{Symbol, String}) -Apply custom name from dictionary if available, otherwise use default name. -""" -function apply_custom_name(symbol::R, custom_names::AbstractDict{S, T})::R where {R <: Union{Symbol, String}, S, T} - # First, check for an exact match with the original symbol - if haskey(custom_names, symbol) - return R(custom_names[symbol]) - end - - # Handle cross-type check for exact match (String vs Symbol) - if symbol isa Symbol && haskey(custom_names, String(replace_indices_in_symbol(symbol))) - return R(custom_names[String(replace_indices_in_symbol(symbol))]) - elseif symbol isa String && haskey(custom_names, Symbol(symbol)) - return R(custom_names[Symbol(symbol)]) - end - # If no exact match, strip lag operators and compare base names. - s_str = string(symbol) - lag_regex = r"^(.*)(ᴸ⁽.*⁾)$" - m = match(lag_regex, s_str) +function write_parameters_input!(𝓂::ℳ, parameters::Vector{Float64}; verbose::Bool = true) + if length(parameters) > length(𝓂.parameter_values) + @warn "Model has $(length(𝓂.parameter_values)) parameters. $(length(parameters)) were provided. The following will be ignored: $(join(parameters[length(𝓂.parameter_values)+1:end], " "))" - base_symbol_str, lag_part = if m !== nothing - (m.captures[1], m.captures[2]) - else - (s_str, "") + parameters = parameters[1:length(𝓂.parameter_values)] end - for (key, value) in custom_names - key_str = string(key) - key_m = match(lag_regex, key_str) - - base_key_str = if key_m !== nothing - key_m.captures[1] - else - key_str - end + bounds_broken = false + parameters_dict = Dict(𝓂.constants.post_complete_parameters.parameters .=> parameters) - if base_key_str == base_symbol_str - return R(string(value) * lag_part) + 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") + bounds_broken = true + continue + end + if val < 𝓂.constants.post_parameters_macro.bounds[par][1] + @warn("Calibration is out of bounds for $par > $(𝓂.constants.post_parameters_macro.bounds[par][1])\t parameter value: $val") + bounds_broken = true + continue + end end end - return symbol -end - -function normalize_superscript(x::Symbol) - return normalize_superscript(string(x)) -end - -function normalize_superscript(x::AbstractString) - sub_map = Dict( - '₀' => '0', '₁' => '1', '₂' => '2', '₃' => '3', '₄' => '4', - '₅' => '5', '₆' => '6', '₇' => '7', '₈' => '8', '₉' => '9', - '₊' => '+', '₋' => '-', '₌' => '=', '₍' => '(', '₎' => ')', - 'ₐ' => 'a', 'ₑ' => 'e', 'ₕ' => 'h', 'ᵢ' => 'i', 'ⱼ' => 'j', - 'ₖ' => 'k', 'ₗ' => 'l', 'ₘ' => 'm', 'ₙ' => 'n', 'ₒ' => 'o', - 'ₚ' => 'p', 'ᵣ' => 'r', 'ₛ' => 's', 'ₜ' => 't', 'ᵤ' => 'u', - 'ᵥ' => 'v', 'ₓ' => 'x' - ) - super_map = Dict( - '⁰' => '0', '¹' => '1', '²' => '2', '³' => '3', '⁴' => '4', - '⁵' => '5', '⁶' => '6', '⁷' => '7', '⁸' => '8', '⁹' => '9', - '⁺' => '+', '⁻' => '-', '⁼' => '=', '⁽' => '(', '⁾' => ')', - 'ᵃ' => 'a', 'ᵇ' => 'b', 'ᶜ' => 'c', 'ᵈ' => 'd', 'ᵉ' => 'e', - 'ᶠ' => 'f', 'ᵍ' => 'g', 'ʰ' => 'h', 'ᶦ' => 'i', 'ʲ' => 'j', - 'ᵏ' => 'k', 'ˡ' => 'l', 'ᵐ' => 'm', 'ⁿ' => 'n', 'ᵒ' => 'o', - 'ᵖ' => 'p', 'ʳ' => 'r', 'ˢ' => 's', 'ᵗ' => 't', 'ᵘ' => 'u', - 'ᵛ' => 'v', 'ʷ' => 'w', 'ˣ' => 'x', 'ʸ' => 'y', 'ᶻ' => 'z' - ) - - buf = IOBuffer() - for c in x - if haskey(sub_map, c) - write(buf, sub_map[c]) - elseif haskey(super_map, c) - write(buf, super_map[c]) - else - write(buf, c) - end - end - return String(take!(buf)) -end - -function replace_indices(exxpr::Expr)::Union{Expr,Symbol} - postwalk(x -> begin - x isa Symbol ? - replace_indices(string(x)) : - x isa Expr ? - x.head == :curly ? - Symbol(string(x.args[1]) * "◖" * string(x.args[2]) * "◗") : - x : - x - end, exxpr) -end - -function replace_indices_special(exxpr::Expr)::Union{Expr,Symbol} - postwalk(x -> begin - x isa Symbol ? - replace_indices(string(x)) : - x isa Expr ? - x.head == :curly ? - Symbol(string(x.args[1]) * "◖" * string(x.args[2]) * "◗") : - x.head == :call ? - x.args[1] == :(*) ? - Symbol(string(x.args[2]), string(x.args[3])) : - x : - x : - x - end, exxpr) -end - -function write_out_for_loops(arg::Expr)::Expr - postwalk(x -> begin - x = flatten(unblock(x)) - x isa Expr ? - x.head == :for ? - x.args[2] isa Array ? - length(x.args[2]) >= 1 ? - x.args[1].head == :block ? - # begin println("here"); - [replace_indices_inside_for_loop(X, Symbol(x.args[1].args[2].args[1]), (x.args[1].args[2].args[2]), false, x.args[1].args[1].args[2].value) for X in x.args[2]] : # end : - # begin println("here2"); - [replace_indices_inside_for_loop(X, Symbol(x.args[1].args[1]), (x.args[1].args[2]), false, :+) for X in x.args[2]] : # end : - x : - x.args[2].head ∉ [:(=), :block] ? - x.args[1].head == :block ? - # begin println("here3"); - replace_indices_inside_for_loop(unblock(x.args[2]), - Symbol(x.args[1].args[2].args[1]), - (x.args[1].args[2].args[2]), - true, - x.args[1].args[1].args[2].value) : # end : # for loop part of equation - x.args[2].head == :if ? - contains_equation(x.args[2]) ? - # begin println("here5"); println(x) - replace_indices_inside_for_loop(unblock(x.args[2]), - Symbol(x.args[1].args[1]), - (x.args[1].args[2]), - false, - :+) : # end : # for loop part of equation - # begin println("here6"); println(x) - replace_indices_inside_for_loop(unblock(x.args[2]), - Symbol(x.args[1].args[1]), - (x.args[1].args[2]), - true, - :+) : # end : # for loop part of equation - # begin println("here4"); println(x) - replace_indices_inside_for_loop(unblock(x.args[2]), - Symbol(x.args[1].args[1]), - (x.args[1].args[2]), - true, - :+) : # end : # for loop part of equation - x.args[1].head == :block ? - # begin println("here5"); - replace_indices_inside_for_loop(unblock(x.args[2]), - Symbol(x.args[1].args[2].args[1]), - (x.args[1].args[2].args[2]), - false, - x.args[1].args[1].args[2].value) : # end : - # end - # : # for loop part of equation - # begin println(x); - # begin println("here7"); println(x) - replace_indices_inside_for_loop(unblock(x.args[2]), - Symbol(x.args[1].args[1]), - (x.args[1].args[2]), - false, - :+) : # end : - # println(out); - # return out end - # : - x : - x - end, - arg) #|> unblock |> flatten -end - -# function parse_for_loops(equations_block) -# eqs = Expr[] # Initialize an empty array to collect expressions - -# # Define a helper recursive function -# function recurse(arg) -# if arg isa Expr -# if arg.head == :block -# for b in arg.args -# if b isa Expr -# # If the result is an Expr, process and add to eqs -# push!(eqs, unblock(replace_indices(b))) -# elseif b isa Array -# recurse(b) -# end -# end -# end -# elseif arg isa Array -# # If the result is an Array, iterate and recurse -# for B in arg -# println((B)) -# recurse(B) -# end -# end -# end - -# for arg in equations_block.args -# if isa(arg,Expr) -# parsed_eqs = write_out_for_loops(arg) -# recurse(parsed_eqs) -# end -# end - -# # Return the collected expressions as a block -# return Expr(:block, eqs...) -# end - - -function parse_for_loops(equations_block)::Expr - eqs = Expr[] - for arg in equations_block.args - if isa(arg,Expr) - parsed_eqs = write_out_for_loops(arg) - # println(parsed_eqs) - if parsed_eqs isa Expr - push!(eqs,unblock(replace_indices(parsed_eqs))) - elseif parsed_eqs isa Array - for B in parsed_eqs - if B isa Array - for b in B - push!(eqs,unblock(replace_indices(b))) - end - elseif B isa Expr - if B.head == :block - for b in B.args - if b isa Expr - push!(eqs,replace_indices(b)) - end - end - else - push!(eqs,unblock(replace_indices(B))) - end - else - push!(eqs,unblock(replace_indices(B))) - end - end - end - - end - end - return Expr(:block,eqs...) |> flatten -end - - - -function decompose_name(name::Symbol) - name = string(name) - matches = eachmatch(r"◖([\p{L}\p{N}]+)◗|([\p{L}\p{N}]+[^◖◗]*)", name) - - result = [] - nested = [] - - for m in matches - if m.captures[1] !== nothing - push!(nested, m.captures[1]) - else - if !isempty(nested) - push!(result, Symbol.(nested)) - nested = [] - end - push!(result, Symbol(m.captures[2])) - end - end - - if !isempty(nested) - push!(result, (nested)) - end - - 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_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 = [] - - for i in indices - if length(i) > 1 - push!(indexset, Symbol.(i[2])...) - end - end - - return indexset -end - - - -function expand_calibration_equations(calibration_equation_parameters::Vector{Symbol}, calibration_equations::Vector{Expr}, ss_calib_list::Vector, par_calib_list::Vector, all_names::Vector{Symbol}) - expanded_parameters = Symbol[] - expanded_equations = Expr[] - expanded_ss_var_list = [] - expanded_par_var_list = [] - - for (u,par) in enumerate(calibration_equation_parameters) - indices_in_calibration_equation = Set() - indexed_names = [] - for i in get_symbols(calibration_equations[u]) - indices = get_possible_indices_for_name(i, all_names) - if indices != Any[] - push!(indices_in_calibration_equation, indices) - push!(indexed_names,i) - end - end - - par_indices = get_possible_indices_for_name(par, all_names) - - if length(par_indices) > 0 - push!(indices_in_calibration_equation, par_indices) - end - - @assert length(indices_in_calibration_equation) <= 1 "Calibration equations cannot have more than one index in the equations or for the parameter." - - if length(indices_in_calibration_equation) == 0 - push!(expanded_parameters,par) - push!(expanded_equations,calibration_equations[u]) - push!(expanded_ss_var_list,ss_calib_list[u]) - push!(expanded_par_var_list,par_calib_list[u]) - else - for i in collect(indices_in_calibration_equation)[1] - expanded_ss_var = Set() - expanded_par_var = Set() - push!(expanded_parameters, Symbol(string(par) * "◖" * string(i) * "◗")) - push!(expanded_equations, postwalk(x -> x ∈ indexed_names ? Symbol(string(x) * "◖" * string(i) * "◗") : x, calibration_equations[u])) - for ss in ss_calib_list[u] - if ss ∈ indexed_names - push!(expanded_ss_var,Symbol(string(ss) * "◖" * string(i) * "◗")) - else - push!(expanded_ss_var,ss) - end - end - # Handle parameters from par_calib_list - expand indexed ones, keep non-indexed - for p in par_calib_list[u] - if p ∈ indexed_names - push!(expanded_par_var, Symbol(string(p) * "◖" * string(i) * "◗")) - else - push!(expanded_par_var, p) - end - end - push!(expanded_ss_var_list, expanded_ss_var) - push!(expanded_par_var_list, expanded_par_var) - end - end - end - - return expanded_parameters, expanded_equations, expanded_ss_var_list, expanded_par_var_list -end - - - -function expand_indices(compressed_inputs::Vector{Symbol}, compressed_values::Vector{T}, expanded_list::Vector{Symbol}) where T - expanded_inputs = Symbol[] - expanded_values = T[] - - for (i,par) in enumerate(compressed_inputs) - par_idx = findall(x -> string(par) == x, first.(split.(string.(expanded_list ), "◖"))) - - if length(par_idx) > 1 - for idx in par_idx - push!(expanded_inputs, expanded_list[idx]) - push!(expanded_values, compressed_values[i]) - end - else#if par ∈ expanded_list ## breaks parameters defined in parameter block - push!(expanded_inputs, par) - push!(expanded_values, compressed_values[i]) - end - end - return expanded_inputs, expanded_values -end - - -function expand_steady_state(SS_and_pars::Vector{M}, ms::post_complete_parameters) where M - X = ms.steady_state_expand_matrix - return X * SS_and_pars -end - - - -function create_symbols_eqs!(𝓂::ℳ)::symbolics - # create symbols in SymPyWorkspace to avoid polluting MacroModelling namespace - symbols_in_dynamic_equations = reduce(union, get_symbols.(𝓂.equations.dynamic)) - - symbols_in_dynamic_equations_wo_subscripts = Symbol.(replace.(string.(symbols_in_dynamic_equations), r"₍₋?(₀|₁|ₛₛ|ₓ)₎$"=>"")) - - symbols_in_ss_equations = reduce(union,get_symbols.(𝓂.equations.steady_state_aux)) - - symbols_in_equation = union(𝓂.constants.post_model_macro.parameters_in_equations, - 𝓂.constants.post_complete_parameters.parameters, - 𝓂.constants.post_parameters_macro.parameters_as_function_of_parameters, - symbols_in_dynamic_equations, - symbols_in_dynamic_equations_wo_subscripts, - symbols_in_ss_equations) #, 𝓂.dynamic_variables_future) - - symbols_pos = [] - symbols_neg = [] - symbols_none = [] - - for symb in symbols_in_equation - if haskey(𝓂.constants.post_parameters_macro.bounds, symb) - if 𝓂.constants.post_parameters_macro.bounds[symb][1] >= 0 - push!(symbols_pos, symb) - elseif 𝓂.constants.post_parameters_macro.bounds[symb][2] <= 0 - push!(symbols_neg, symb) - else - push!(symbols_none, symb) - end - else - push!(symbols_none, symb) - end - end - - # Create symbols in SymPyWorkspace instead of MacroModelling namespace - for pos in symbols_pos - sym_value = SPyPyC.symbols(string(pos), real = true, finite = true, positive = true) - Core.eval(SymPyWorkspace, :($pos = $sym_value)) - end - - for neg in symbols_neg - sym_value = SPyPyC.symbols(string(neg), real = true, finite = true, negative = true) - Core.eval(SymPyWorkspace, :($neg = $sym_value)) - end - - for none in symbols_none - sym_value = SPyPyC.symbols(string(none), real = true, finite = true) - Core.eval(SymPyWorkspace, :($none = $sym_value)) - end - - symbolics( - map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.equations.steady_state_aux), - # map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.dyn_equations_future), - - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift_var_present_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift_var_past_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift_var_future_list), - - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift2_var_past_list), - - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_var_present_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_var_past_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_var_future_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_ss_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_exo_list), - - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_exo_future_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_exo_present_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_exo_past_list), - - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_future_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_present_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_past_list), - - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_present_list_aux_SS), - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_past_list_aux_SS), - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_future_list_aux_SS), - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.ss_list_aux_SS), - - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_list_aux_SS), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dynamic_variables_list), - # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dynamic_variables_future_list), - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.par_list_aux_SS), - - map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.equations.calibration), - map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.equations.calibration_parameters), - # map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.constants.post_complete_parameters.parameters), - - # Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.var_present...)]))), - # Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.var_past...)]))), - # Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.var_future...)]))), - Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.vars_in_ss_equations...)]))), - - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_parameters_macro.ss_calib_list), - map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_parameters_macro.par_calib_list), - - [Set() for _ in 1:length(𝓂.equations.steady_state_aux)], - # [Set() for _ in 1:length(𝓂.calibration_equations)], - # [Set() for _ in 1:length(𝓂.equations.steady_state_aux)], - # [Set() for _ in 1:length(𝓂.calibration_equations)] - ) -end - - - -function remove_redundant_SS_vars!(𝓂::ℳ, Symbolics::symbolics; avoid_solve::Bool = false) - ss_equations = Symbolics.ss_equations - - # check variables which appear in two time periods. they might be redundant in steady state - redundant_vars = intersect.( - union.( - intersect.(Symbolics.var_future_list_aux_SS, Symbolics.var_present_list_aux_SS), - intersect.(Symbolics.var_future_list_aux_SS, Symbolics.var_past_list_aux_SS), - intersect.(Symbolics.var_present_list_aux_SS, Symbolics.var_past_list_aux_SS), - intersect.(Symbolics.ss_list_aux_SS, Symbolics.var_present_list_aux_SS), - intersect.(Symbolics.ss_list_aux_SS, Symbolics.var_past_list_aux_SS), - intersect.(Symbolics.ss_list_aux_SS, Symbolics.var_future_list_aux_SS) - ), - Symbolics.var_list_aux_SS) - - redundant_idx = getindex(1:length(redundant_vars), (length.(redundant_vars) .> 0) .& (length.(Symbolics.var_list_aux_SS) .> 1)) - for i in redundant_idx - for var_to_solve_for in redundant_vars[i] - if avoid_solve || count_ops(Meta.parse(string(ss_equations[i]))) > 15 - soll = nothing - else - soll = solve_symbolically(ss_equations[i],var_to_solve_for) - end - - if isnothing(soll) - continue - end - - if isempty(soll) || soll == SPyPyC.Sym{PythonCall.Core.Py}[0] # take out variable if it is redundant from that euation only - push!(Symbolics.var_redundant_list[i],var_to_solve_for) - ss_equations[i] = replace_with_one(ss_equations[i], var_to_solve_for) # replace euler constant as it is not translated to julia properly - end - - end - end - -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)) ) ) - - for var in other_vrs - # push!(other_vars,:($(var) = parameters_and_solved_vars[$iii])) - 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,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) - end - push!(aux_replacements, x => aux_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)) |> - 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()) - - chol_buffer = 𝒮.init(prob, 𝒮.CholeskyFactorization(), verbose = isdefined(𝒮, :LinearVerbosity) ? 𝒮.LinearVerbosity(𝒮.SciMLLogging.Minimal()) : false) - - prob = 𝒮.LinearProblem(buffer, ϵ, 𝒮.LUFactorization()) - - 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 - - _, 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' - - 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} - - - 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 - - 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 - - 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) - - 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) - buffer = zeros(Float64, size(∂SS_equations_∂SS_and_pars)) - else - 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()) - - chol_buffer = 𝒮.init(prob, 𝒮.CholeskyFactorization(), verbose = isdefined(𝒮, :LinearVerbosity) ? 𝒮.LinearVerbosity(𝒮.SciMLLogging.Minimal()) : false) - - prob = 𝒮.LinearProblem(buffer, ϵ, 𝒮.LUFactorization()) - - 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 - - _, 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' - - 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} - - - - push!(NSSS_solver_cache_init_tmp,fill(1.205996189998029, length(sorted_vars))) - push!(NSSS_solver_cache_init_tmp,[Inf]) - - # WARNING: infinite bounds are transformed to 1e12 - lbs = [] - ubs = [] - - 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] == -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 - - 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...)])) - - 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])) - - # 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)))) - - 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) - ) - ) - - n_block += 1 - - n -= 1 - 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)) - - 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,:(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)) - - 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)) - - # 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}, - 𝓂::ℳ, - 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) - - 𝓂.functions.NSSS_solve = @RuntimeGeneratedFunction(solve_exp) - # 𝓂.functions.NSSS_solve = eval(solve_exp) - - 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 - σ = 5 - log_lik -= -log(σ * sqrt(2 * π)) - (pars[20]^2 / (2 * σ^2)) # logpdf of a normal dist with mean = 0 and variance = 5^2 - - pars[1:2] = sort(pars[1:2], rev = true) - - par_inputs = solver_parameters(pars..., 1, 0.0, 2) - - while length(𝓂.caches.solver_cache) > 1 - pop!(𝓂.caches.solver_cache) - end - - runtime = @elapsed outmodel = try 𝓂.functions.NSSS_solve(𝓂.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]) ? - 10 : - runtime : - 10 - - return log_lik / 1e4 + runtime * 1e3 -end - -""" - find_SS_solver_parameters!(::Val{:ESCH}, 𝓂::ℳ; maxtime::Real = 120, maxiter::Int = 2500000, tol::Tolerances = Tolerances(), verbosity = 0) - -Find optimal steady state solver parameters using NLopt's ESCH algorithm. - -This function optimizes solver parameters to minimize runtime while maintaining solver accuracy. -It uses the ESCH global optimization algorithm from the NLopt package. - -# Arguments -- `𝓂`: Model structure -- `maxtime`: Maximum time in seconds for optimization -- `maxiter`: Maximum number of iterations -- `tol`: Tolerance structure -- `verbosity`: Verbosity level for output -""" -function find_SS_solver_parameters!(::Val{:ESCH}, 𝓂::ℳ; maxtime::Real = 120, maxiter::Int = 2500000, tol::Tolerances = Tolerances(), verbosity = 0) - pars = rand(20) .+ 1 - pars[20] -= 1 - - lbs = fill(eps(), length(pars)) - lbs[20] = -20 - - ubs = fill(100.0, length(pars)) - - opt = NLopt.Opt(NLopt.:GN_ESCH, length(pars)) - - opt.min_objective = (x,p) -> calculate_SS_solver_runtime_and_loglikelihood(x, 𝓂, tol = tol) - - NLopt.lower_bounds!(opt, lbs) - NLopt.upper_bounds!(opt, ubs) - - opt.xtol_abs = eps(Float32) - opt.ftol_abs = eps(Float32) - - # opt.maxeval = maxiter - opt.maxtime = maxtime - - (minf,x,ret) = NLopt.optimize(opt, pars) - - verbosity > 0 && @info "NLopt return code: $ret" - - pars = x - - 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]) - - if solution_error < tol.NSSS_acceptance_tol - push!(DEFAULT_SOLVER_PARAMETERS, par_inputs) - return true - else - return false - end -end - - -function select_fastest_SS_solver_parameters!(𝓂::ℳ; tol::Tolerances = Tolerances()) - best_param = DEFAULT_SOLVER_PARAMETERS[1] - - best_time = Inf - - solved = false - - solved_NSSS = 𝓂.caches.solver_cache[end] - - for p in DEFAULT_SOLVER_PARAMETERS - total_time = 0.0 - - for _ in 1:100 - start_time = time() - - while length(𝓂.caches.solver_cache) > 1 - pop!(𝓂.caches.solver_cache) - end - - SS_and_pars, (solution_error, iters) = 𝓂.functions.NSSS_solve(𝓂.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 - break - end - end - - if total_time < best_time - best_time = total_time - best_param = p - end - - solved = true - end - - while length(𝓂.caches.solver_cache) > 1 - pop!(𝓂.caches.solver_cache) - end - - push!(𝓂.caches.solver_cache, solved_NSSS) - - if solved - pushfirst!(DEFAULT_SOLVER_PARAMETERS, best_param) - end -end - - -function solve_ss(SS_optimizer::Function, - # ss_solve_blocks::Function, - SS_solve_block::ss_solve_block, - parameters_and_solved_vars::Vector{T}, - closest_parameters_and_solved_vars::Vector{T}, - lbs::Vector{T}, - ubs::Vector{T}, - tol::Tolerances, - total_iters::Vector{Int}, - n_block::Int, - verbose::Bool, - guess::Vector{T}, - 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 - 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])) - end - - sol_new_tmp, info = SS_optimizer( extended_problem ? SS_solve_block.extended_ss_problem : SS_solve_block.ss_problem, - # if extended_problem - # function ext_function_to_optimize(guesses) - # gss = guesses[1:length(guess)] - - # parameters_and_solved_vars_guess = guesses[length(guess)+1:end] - - # res = ss_solve_blocks(parameters_and_solved_vars, gss) - - # return vcat(res, parameters_and_solved_vars .- parameters_and_solved_vars_guess) - # end - # else - # function function_to_optimize(guesses) ss_solve_blocks(parameters_and_solved_vars, guesses) end - # 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, - parameters_and_solved_vars, - extended_problem ? lbs : lbs[1:length(guess)], - extended_problem ? ubs : ubs[1:length(guess)], - 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)) - - total_iters[1] += info[1] - total_iters[2] += info[2] - - 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 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 - - # max_resid = maximum(abs,ss_solve_blocks(parameters_and_solved_vars, sol_values)) - - 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") - end - - return sol_values, total_iters, rel_sol_minimum, sol_minimum -end - - -function block_solver(parameters_and_solved_vars::Vector{T}, - n_block::Int, - # ss_solve_blocks::Function, - SS_solve_block::ss_solve_block, - # SS_optimizer, - # f::OptimizationFunction, - guess_and_pars_solved_vars::Vector{Vector{T}}, - lbs::Vector{T}, - ubs::Vector{T}, - parameters::Vector{solver_parameters}, - fail_fast_solvers_only::Bool, - cold_start::Bool, - verbose::Bool ; - tol::Tolerances = Tolerances(), - # rtol::AbstractFloat = sqrt(eps()), - # timeout = 120, - # starting_points::Vector{Float64} = [1.205996189998029, 0.7688, 0.897, 1.2],#, 0.9, 0.75, 1.5, -0.5, 2.0, .25] - # verbose::Bool = false - )::Tuple{Vector{T},Tuple{T, Int}} where T <: AbstractFloat - - # tol = parameters[1].ftol - # rtol = parameters[1].rel_xtol - - solved_yet = false - - guess = guess_and_pars_solved_vars[1] - - sol_values = guess - - closest_parameters_and_solved_vars = sum(abs, guess_and_pars_solved_vars[2]) == Inf ? parameters_and_solved_vars : guess_and_pars_solved_vars[2] - - # res = ss_solve_blocks(parameters_and_solved_vars, guess) - - SS_solve_block.ss_problem.func(SS_solve_block.ss_problem.workspace.func_buffer, guess, parameters_and_solved_vars) - - res = SS_solve_block.ss_problem.workspace.func_buffer - - sol_minimum = ℒ.norm(res) - - if !cold_start - 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) - - SS_solve_block.ss_problem.jac(SS_solve_block.ss_problem.workspace.jac_buffer, guess, parameters_and_solved_vars) - - ∇ = 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) - else - rel_sol_minimum = 1.0 - end - else - rel_sol_minimum = 0.0 - end - else - rel_sol_minimum = 1.0 - end - - if isfinite(sol_minimum) && sol_minimum < tol.NSSS_acceptance_tol - solved_yet = true - - if verbose - println("Block: $n_block, - Solved using previous solution; residual norm: $sol_minimum") - end - end - - total_iters = [0,0] - - SS_optimizer = 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]) - - 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 s in start_vals - 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, - # sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(SS_optimizer, ss_solve_blocks, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, total_iters, n_block, verbose, - g, - p, - ext, - s) - - if isfinite(sol_minimum) && sol_minimum < tol.NSSS_acceptance_tol - solved_yet = true - end - end - end - end - end - 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 - 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, - # sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(algo, ss_solve_blocks, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, - total_iters, - n_block, - false, # verbose - guess, - p, - # parameters[1], - false, # ext - # false) - s) - 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 - - - # if sol_minimum > tol# || rel_sol_minimum > rtol - # for p in unique(parameters)#[1:3] # take unique because some parameters might appear more than once - # # for s in [p.starting_value, 1.206, 1.5, 0.7688, 2.0, 0.897]#, .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 - # if sol_minimum > tol# || rel_sol_minimum > rtol - # sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(SS_optimizer, ss_solve_blocks, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, total_iters, n_block, - # false, # verbose - # guess, - # p, - # false, - # false) - # # s) - # if !solved_yet && sol_minimum < tol# || rel_sol_minimum > rtol) - # solved_yet = true - # if verbose - # loop1 = unique(parameters)#[1:3] - # loop2 = [p.starting_value, 1.206, 1.5, 0.7688, 2.0, 0.897] - # p_in_loop1 = findfirst(x -> x == p, loop1) - # s_in_loop2 = findfirst(x -> x == s, loop2) - # if p_in_loop1 isa Nothing - # p_in_loop1 = 1 - # end - # if s_in_loop2 isa Nothing - # s_in_loop2 = 1 - # end - # n1 = (p_in_loop1 - 1) * length(loop2) + s_in_loop2 - # println("Block: $n_block, - Solved with modified Levenberg-Marquardt - $n1/$(length(loop2) *length(loop1)) - $sol_minimum - $rel_sol_minimum - $total_iters") - # end - # end - # end - # # end - # # end - # end - # end - end - - if verbose - if !solved_yet - println("Block: $n_block, - Solution not found after $(total_iters[1]) gradient evaluations and $(total_iters[2]) function evaluations; reltol: $rel_sol_minimum - tol: $sol_minimum") - end - end - - return sol_values, (sol_minimum, total_iters[1]) -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 - 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) - - # 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) - end - - ms = @ignore_derivatives 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 - - # @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) - - # end # timeit_debug - - 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) - 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 - - @ignore_derivatives update_perturbation_counter!(𝓂.counters, solved2, estimation = estimation, order = 2) - - 𝐒₂ = sparse(𝐒₂ * 𝓂.constants.second_order.𝐔₂)::SparseMatrixCSC{M, Int} - - # end # timeit_debug - - 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) - end - - # @timeit_debug timer "Calculate SSS" begin - - 𝐒₁ = [𝐒₁[:,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 - - tmp = (I_nPast - 𝐒₁[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx,1:𝓂.constants.post_model_macro.nPast_not_future_and_mixed]) - - tmp̄ = @ignore_derivatives ℒ.lu(tmp, check = false) - - if !ℒ.issuccess(tmp̄) - 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) - end - - SSSstates = collect(tmp \ (𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2)[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx]) - - 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) - 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(𝐒₂)], 𝓂 = 𝓂) - end - - # end # timeit_debug - - # all_variables = sort(union(𝓂.constants.post_model_macro.var,𝓂.constants.post_model_macro.aux,𝓂.constants.post_model_macro.exo_present)) - - # 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 - - # 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, ∇₁, ∇₂, 𝐒₁, 𝐒₂ -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 - # @timeit_debug timer "Setup matrices" 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 = 𝐒₂[T.past_not_future_and_mixed_idx,kron_s⁺_s] - B̂ = 𝐒₂[T.past_not_future_and_mixed_idx,kron_s⁺_s⁺] - - max_iters = 100 - # SSS .= 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 𝐒₃ * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 - - # 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) - - ∂x̂ = ℒ.lu!(∂x, check = false) - - if !ℒ.issuccess(∂x̂) - return x, false - end - - x̂ = A * x + B̂ * ℒ.kron(vcat(x,1), vcat(x,1)) / 2 - - Δx = ∂x̂ \ (x̂ - x) - - if i > 3 && isapprox(x̂, x, rtol = tol) - break - end - - # x += Δx - ℒ.axpy!(-1, Δx, x) - end - - # 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) - - ∇₁ = 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) - - @ignore_derivatives 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), 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} - - ∇₃ = 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) - - @ignore_derivatives update_perturbation_counter!(𝓂.counters, solved3, estimation = estimation, order = 3) - - if !solved3 - if opts.verbose println("3rd 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 - - if eltype(𝐒₃) == Float64 && solved3 𝓂.caches.third_order_solution = 𝐒₃ end - - if length(𝓂.workspaces.third_order.Ŝ) == 0 || !(eltype(𝐒₃) == eltype(𝓂.workspaces.third_order.Ŝ)) - 𝓂.workspaces.third_order.Ŝ = 𝐒₃ * 𝓂.constants.third_order.𝐔₃ - else - mul_reverse_AD!(𝓂.workspaces.third_order.Ŝ, 𝐒₃, 𝓂.constants.third_order.𝐔₃) - end - - Ŝ = 𝓂.workspaces.third_order.Ŝ - - 𝐒₃̂ = sparse_preallocated!(Ŝ, ℂ = 𝓂.workspaces.third_order)::SparseMatrixCSC{M, Int} - - # 𝐒₃ *= 𝓂.constants.third_order.𝐔₃ - # 𝐒₃ = sparse_preallocated!(𝐒₃, ℂ = 𝓂.workspaces.third_order) - - # 𝐒₃ = sparse(Ŝ) # * 𝓂.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 - - tmp = (I_nPast - 𝐒₁[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx, 1:𝓂.constants.post_model_macro.nPast_not_future_and_mixed]) - - tmp̄ = @ignore_derivatives ℒ.lu(tmp, check = false) - - if !ℒ.issuccess(tmp̄) - 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) - end - - SSSstates = collect(tmp \ (𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2)[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx]) - - 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⁺] - Ĉ = 𝐒₃̂[:,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) - end - - state = A * SSSstates + B̂ * ℒ.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 + Ĉ * ℒ.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(𝐒₃)], 𝓂 = 𝓂) - end - - # all_variables = sort(union(𝓂.constants.post_model_macro.var,𝓂.constants.post_model_macro.aux,𝓂.constants.post_model_macro.exo_present)) - - # 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 - - # 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, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃̂ -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) - # 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 - - return x, 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) -end - - - -function set_up_steady_state_solver!(𝓂::ℳ; verbose::Bool, silent::Bool, avoid_solve::Bool = false, symbolic::Bool = false) - if !𝓂.constants.post_parameters_macro.precompile - start_time = time() - - if !silent print("Remove redundant variables in non-stochastic steady state problem:\t") end - - symbolics = create_symbols_eqs!(𝓂) - - remove_redundant_SS_vars!(𝓂, symbolics, avoid_solve = avoid_solve) - - if !silent println(round(time() - start_time, digits = 3), " seconds") end - - start_time = time() - - if !silent print("Set up non-stochastic steady state problem:\t\t\t\t") end - - write_ss_check_function!(𝓂) - - write_steady_state_solver_function!(𝓂, symbolic, symbolics, verbose = verbose, avoid_solve = avoid_solve) - - 𝓂.equations.obc_violation = write_obc_violation_equations(𝓂) - - set_up_obc_violation_function!(𝓂) - - if !silent println(round(time() - start_time, digits = 3), " seconds") end - else - start_time = time() - - if !silent print("Set up non-stochastic steady state problem:\t\t\t\t") end - - write_ss_check_function!(𝓂) - - write_steady_state_solver_function!(𝓂, verbose = verbose) - - if !silent println(round(time() - start_time, digits = 3), " seconds") end - end - - return nothing -end - -function solve!(𝓂::ℳ; - parameters::ParameterType = nothing, - steady_state_function::SteadyStateFunctionType = missing, - dynamics::Bool = false, - algorithm::Symbol = :first_order, - opts::CalculationOptions = merge_calculation_options(), - obc::Bool = false, - silent::Bool = false) #, - # quadratic_matrix_equation_algorithm::Symbol = :schur, - # verbose::Bool = false, - # timer::TimerOutput = TimerOutput(), - # tol::AbstractFloat = 1e-12) - - @assert algorithm ∈ all_available_algorithms - - # Initialize constants at entry point - constants = initialise_constants!(𝓂) - - # Handle steady_state_function argument - set_custom_steady_state_function!(𝓂, steady_state_function) - - # @timeit_debug timer "Write parameter inputs" begin - - write_parameters_input!(𝓂, parameters, verbose = opts.verbose) - - if 𝓂.functions.functions_written && - isnothing(𝓂.functions.NSSS_custom) && - !(𝓂.functions.NSSS_solve isa RuntimeGeneratedFunctions.RuntimeGeneratedFunction) - - set_up_steady_state_solver!(𝓂, verbose = opts.verbose, silent = silent) - end - - if !𝓂.functions.functions_written - verbose = opts.verbose - - perturbation_order = 1 - - set_up_steady_state_solver!(𝓂, verbose = verbose, silent = silent, avoid_solve = false) - - SS_and_pars, solution_error, found_solution = solve_steady_state!(𝓂, opts, :ESCH, 120.0, silent = silent) - - write_symbolic_derivatives!(𝓂; perturbation_order = perturbation_order, silent = silent) - - 𝓂.functions.functions_written = true - end - - # Check for missing parameters after processing input - if !isempty(𝓂.constants.post_complete_parameters.missing_parameters) - error("Cannot solve model: missing parameter values for $(𝓂.constants.post_complete_parameters.missing_parameters). Provide them via the `parameters` keyword argument (e.g., `parameters = [:α => 0.3, :β => 0.99]`).") - end - - # end # timeit_debug - - if 𝓂.constants.second_order.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && - algorithm ∈ [:second_order, :pruned_second_order] - start_time = time() - if !silent print("Take symbolic derivatives up to second order:\t\t\t\t") end - write_functions_mapping!(𝓂, 2) - if !silent println(round(time() - start_time, digits = 3), " seconds") end - elseif 𝓂.constants.third_order.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && algorithm ∈ [:third_order, :pruned_third_order] - start_time = time() - if !silent print("Take symbolic derivatives up to third order:\t\t\t\t") end - write_functions_mapping!(𝓂, 3) - if !silent println(round(time() - start_time, digits = 3), " seconds") end - 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 - - ∇₁ = calculate_jacobian(𝓂.parameter_values, SS_and_pars, 𝓂.caches, 𝓂.functions.jacobian)# |> Matrix - - # end # timeit_debug - - # @timeit_debug timer "Calculate first order solution" begin - - 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; - opts = opts, - initial_guess = 𝓂.caches.qme_solution) - - if solved 𝓂.caches.qme_solution = qme_sol end - - update_perturbation_counter!(𝓂.counters, solved, order = 1) - - # end # timeit_debug - - @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 - - Ŝ₁, 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 Ŝ₁ * 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 - Ŝ₁̂ = [Ŝ₁[:,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]] - - 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 - 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 - - if obc - Ŝ₁̂ = [Ŝ₁[:,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]] - - 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 - else - state_update₂̂ = (x,y)->nothing - 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₂̂ - - 𝓂.caches.outdated.pruned_second_order_solution = false - 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 - Ŝ₁̂ = [Ŝ₁[:,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]] - - 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 - 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 - 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 - Ŝ₁̂ = [Ŝ₁[:,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]] - - 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] # 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 -end - - - - -function create_second_order_auxiliary_matrices(constants::constants) - T = constants.post_model_macro - - - # Indices and number of variables - n₋ = T.nPast_not_future_and_mixed - nₑ = T.nExo - - # setup compression matrices for hessian matrix - nₑ₋ = T.nPast_not_future_and_mixed + T.nVars + T.nFuture_not_past_and_mixed + T.nExo - 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) - - # set up vector to capture volatility effect - nₑ₋ = n₋ + 1 + nₑ - 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) - - so = constants.second_order - so.𝛔 = 𝛔 - so.𝐂₂ = 𝐂₂ - so.𝐔₂ = 𝐔₂ - so.𝐔∇₂ = 𝐔∇₂ - return so -end - - - -function add_sparse_entries!(P, perm) - n = size(P, 1) - for i in 1:n - P[perm[i], i] += 1.0 - end -end - - -function create_third_order_auxiliary_matrices(constants::constants, ∇₃_col_indices::Vector{Int}) - T = constants.post_model_macro - - - # Indices and number of variables - n₋ = T.nPast_not_future_and_mixed - n₊ = T.nFuture_not_past_and_mixed - n = T.nVars - nₑ = T.nExo - - n̄ = n₋ + n + n₊ + nₑ - - # compression matrices for third order derivatives matrix - nₑ₋ = T.nPast_not_future_and_mixed + T.nVars + T.nFuture_not_past_and_mixed + T.nExo - colls3 = [nₑ₋^2 * (i-1) + nₑ₋ * (k-1) + l for i in 1:nₑ₋ for k in 1:i for l in 1:k] - 𝐂∇₃ = sparse(colls3, 1:length(colls3) , 1.0) - - idxs = Int[] - for k in 1:nₑ₋ - for j in 1:nₑ₋ - for i in 1:nₑ₋ - sorted_ids = sort([k,j,i]) - push!(idxs, (sorted_ids[3] - 1) * nₑ₋ ^ 2 + (sorted_ids[2] - 1) * nₑ₋ + sorted_ids[1]) - end - end - end - - 𝐔∇₃ = 𝐂∇₃' * sparse(idxs,1:nₑ₋ ^ 3, 1) - - # compression matrices for third order transition matrix - nₑ₋ = n₋ + 1 + nₑ - colls3 = [nₑ₋^2 * (i-1) + nₑ₋ * (k-1) + l for i in 1:nₑ₋ for k in 1:i for l in 1:k] - 𝐂₃ = sparse(colls3, 1:length(colls3) , 1.0) - - idxs = Int[] - for k in 1:nₑ₋ - for j in 1:nₑ₋ - for i in 1:nₑ₋ - sorted_ids = sort([k,j,i]) - push!(idxs, (sorted_ids[3] - 1) * nₑ₋ ^ 2 + (sorted_ids[2] - 1) * nₑ₋ + sorted_ids[1]) - end - end - end - - 𝐔₃ = 𝐂₃' * sparse(idxs,1:nₑ₋ ^ 3, 1) - - # Precompute 𝐈₃ - 𝐈₃ = Dict{Vector{Int}, Int}() - idx = 1 - for i in 1:nₑ₋ - for k in 1:i - for l in 1:k - 𝐈₃[[i,k,l]] = idx - idx += 1 - end - end - end - - # permutation matrices - M = reshape(1:nₑ₋^3,1,nₑ₋,nₑ₋,nₑ₋) - - 𝐏 = spzeros(nₑ₋^3, nₑ₋^3) # Preallocate the sparse matrix - - # Create the permutations directly - add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 4, 2, 3))) - add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 2, 4, 3))) - add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 2, 3, 4))) - - # 𝐏 = @views sparse(reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 4, 2, 3])],nₑ₋^3,nₑ₋^3) - # + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 4, 3])],nₑ₋^3,nₑ₋^3) - # + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 3, 4])],nₑ₋^3,nₑ₋^3)) - - 𝐏₁ₗ = sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(2,1,3))),:]) - 𝐏₁ᵣ = sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(2,1,3)))]) - - 𝐏₁ₗ̂ = @views sparse(spdiagm(ones(n̄^3))[vec(permutedims(reshape(1:n̄^3,n̄,n̄,n̄),(1,3,2))),:]) - 𝐏₂ₗ̂ = @views sparse(spdiagm(ones(n̄^3))[vec(permutedims(reshape(1:n̄^3,n̄,n̄,n̄),(3,1,2))),:]) - - 𝐏₁ₗ̄ = @views sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(1,3,2))),:]) - 𝐏₂ₗ̄ = @views sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(3,1,2))),:]) - - - 𝐏₁ᵣ̃ = @views sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(1,3,2)))]) - 𝐏₂ᵣ̃ = @views sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(3,1,2)))]) - - ∇₃_col_indices_extended = findnz(sparse(ones(Int,length(∇₃_col_indices)),∇₃_col_indices,ones(Int,length(∇₃_col_indices)),1,size(𝐔∇₃,1)) * 𝐔∇₃)[2] - - nonnull_columns = Set{Int}() - for i in 1:n̄ - for j in i:n̄ - for k in j:n̄ - if n̄^2 * (i - 1) + n̄ * (j - 1) + k in ∇₃_col_indices_extended - push!(nonnull_columns,i) - push!(nonnull_columns,j) - push!(nonnull_columns,k) - end - end - end - end - - 𝐒𝐏 = sparse(collect(nonnull_columns), collect(nonnull_columns), 1, n̄, n̄) - - to = constants.third_order - to.𝐂₃ = 𝐂₃ - to.𝐔₃ = 𝐔₃ - to.𝐈₃ = 𝐈₃ - to.𝐂∇₃ = 𝐂∇₃ - to.𝐔∇₃ = 𝐔∇₃ - to.𝐏 = 𝐏 - to.𝐏₁ₗ = 𝐏₁ₗ - to.𝐏₁ᵣ = 𝐏₁ᵣ - to.𝐏₁ₗ̂ = 𝐏₁ₗ̂ - to.𝐏₂ₗ̂ = 𝐏₂ₗ̂ - to.𝐏₁ₗ̄ = 𝐏₁ₗ̄ - to.𝐏₂ₗ̄ = 𝐏₂ₗ̄ - to.𝐏₁ᵣ̃ = 𝐏₁ᵣ̃ - to.𝐏₂ᵣ̃ = 𝐏₂ᵣ̃ - to.𝐒𝐏 = 𝐒𝐏 - return to -end - -function take_nth_order_derivatives( - dyn_equations::Vector{T}, - 𝔙::Symbolics.Arr, - 𝔓::Symbolics.Arr, - SS_mapping::Dict{T, T}, - nps::Int, - nxs::Int; - max_perturbation_order::Int = 1, - output_compressed::Bool = true # Controls compression for X derivatives (order >= 2) -)::Vector{Tuple{SparseMatrixCSC{T, Int}, SparseMatrixCSC{T, Int}}} where T <: Symbolics.Num#, Tuple{Symbolics.Arr{Symbolics.Num, 1}, Symbolics.Arr{Symbolics.Num, 1}}} - - nx = BigInt(length(𝔙)::Int) - # np = length(𝔓)::BigInt - nϵ = length(dyn_equations)::Int - - if max_perturbation_order < 1 - throw(ArgumentError("max_perturbation_order must be at least 1")) - end - - results = [] # To store pairs of sparse matrices (X_matrix, P_matrix) for each order - - # --- Order 1 --- - # Compute the 1st order derivative with respect to X (Jacobian) - spX_order_1 = Symbolics.sparsejacobian(dyn_equations, 𝔙) # nϵ x nx - - - spX_order_1_sub = copy(spX_order_1) - - # spX_order_1_sub.nzval .= Symbolics.fast_substitute(spX_order_1_sub.nzval, Dict(Symbolics.scalarize(𝔛𝔛) .=> 𝔙)) - spX_order_1_sub.nzval .= Symbolics.substitute(spX_order_1_sub.nzval, SS_mapping) - - # Compute the derivative of the non-zeros of the 1st X-derivative w.r.t. P - # This is an intermediate step. The final P matrix will be built from this. - spP_of_flatX_nzval_order_1 = Symbolics.sparsejacobian(spX_order_1_sub.nzval, vcat(𝔓[1:nps], 𝔙[1:nxs])) # nnz(spX_order_1) x np - - # Determine dimensions for the Order 1 P matrix - X_nrows_1 = nϵ - X_ncols_1 = nx - P_nrows_1 = X_nrows_1 * X_ncols_1 - P_ncols_1 = nps + nxs - - # Build the Order 1 P matrix (dimensions nϵ*nx x np) - sparse_rows_1_P = Int[] # Row index in the flattened space of spX_order_1 - sparse_cols_1_P = Int[] # Column index for parameters (1 to np) - sparse_vals_1_P = Symbolics.Num[] - - # Map linear index in spX_order_1.nzval to its (row, col) in spX_order_1 - nz_lin_to_rc_1 = Dict{Int, Tuple{Int, Int}}() - k_lin = 1 - for j = 1:size(spX_order_1, 2) # col - for ptr = spX_order_1.colptr[j]:(spX_order_1.colptr[j+1]-1) - r = spX_order_1.rowval[ptr] # row - nz_lin_to_rc_1[k_lin] = (r, j) - k_lin += 1 - end - end - - - # Iterate through the non-zero entries of spP_of_flatX_nzval_order_1 - k_temp_P = 1 # linear index counter for nzval - for p_col = 1:size(spP_of_flatX_nzval_order_1, 2) # Parameter index - for i_ptr_temp_P = spP_of_flatX_nzval_order_1.colptr[p_col]:(spP_of_flatX_nzval_order_1.colptr[p_col+1]-1) - temp_row = spP_of_flatX_nzval_order_1.rowval[i_ptr_temp_P] # Row index in spP_of_flatX_nzval (corresponds to temp_row-th nzval of spX_order_1) - p_val = spP_of_flatX_nzval_order_1.nzval[i_ptr_temp_P] # Derivative value w.r.t. parameter - - # Get the (row, col) in spX_order_1 corresponding to this derivative - r_X1, c_X1 = nz_lin_to_rc_1[temp_row] - - # Calculate the row index in spP_order_1 (flattened index of spX_order_1) - # P_row_idx = (r_X1 - 1) * X_ncols_1 + c_X1 - P_row_idx = (c_X1 - 1) * X_nrows_1 + r_X1 - P_col_idx = p_col # Parameter column index - - push!(sparse_rows_1_P, P_row_idx) - push!(sparse_cols_1_P, P_col_idx) - push!(sparse_vals_1_P, p_val) - - k_temp_P += 1 - end - end - - spP_order_1 = sparse!(sparse_rows_1_P, sparse_cols_1_P, sparse_vals_1_P, P_nrows_1, P_ncols_1) - - - # Store the pair for order 1 - push!(results, (spX_order_1_sub, spP_order_1)) - - if max_perturbation_order > 1 - # --- Prepare for higher orders (Order 2 to max_perturbation_order) --- - # Initialize map for Order 1: linear index in spX_order_1.nzval -> (row, (v1,)) - # This map is needed to trace indices for Order 2 - # We already built nz_lin_to_rc_1 above, reuse it and wrap the variable index in a Tuple - nz_to_indices_prev = Dict{Int, Tuple{Int, Tuple{Int}}}() - k_lin = 1 - for j = 1:size(spX_order_1, 2) - for ptr = spX_order_1.colptr[j]:(spX_order_1.colptr[j+1]-1) - r = spX_order_1.rowval[ptr] - nz_to_indices_prev[k_lin] = (r, (j,)) # Store (equation row, (v1,)) - k_lin += 1 - end - end - - nzvals_prev = spX_order_1.nzval # nzvals from Order 1 X-matrix - - # --- Iterate for orders n = 2, 3, ..., max_perturbation_order --- - for n = 2:max_perturbation_order - - # Compute the Jacobian of the previous level's nzval w.r.t. 𝔛 - # This gives a flat matrix where rows correspond to non-zeros from order n-1 X-matrix - # and columns correspond to the n-th variable we differentiate by (x_vn). - sp_flat_curr_X_rn = Symbolics.sparsejacobian(nzvals_prev, 𝔙) # nnz(spX_order_(n-1)) x nx - - sp_flat_curr_X = copy(sp_flat_curr_X_rn) - - sp_flat_curr_X.nzval .= Symbolics.substitute(sp_flat_curr_X.nzval, SS_mapping) - - # Build the nz_to_indices map for the *current* level (order n) - # Map: linear index in sp_flat_curr_X.nzval -> (original_row_f, (v_1, ..., v_n)) - nz_to_indices_curr = Dict{Int, Tuple{Int, Tuple{Vararg{Int}}}}() - k_lin_curr = 1 # linear index counter for nzval of sp_flat_curr_X - # Iterate through the non-zeros of the current flat Jacobian - for col_curr = 1:size(sp_flat_curr_X, 2) # Column index in sp_flat_curr_X (corresponds to v_n) - for ptr_curr = sp_flat_curr_X.colptr[col_curr]:(sp_flat_curr_X.colptr[col_curr+1]-1) - row_curr = sp_flat_curr_X.rowval[ptr_curr] # Row index in sp_flat_curr_X (corresponds to the row_curr-th nzval of previous level) - - # Get previous indices info from the map of order n-1 - prev_info = nz_to_indices_prev[row_curr] - orig_row_f = prev_info[1] # Original equation row - vars_prev = prev_info[2] # Tuple of variables from previous order (v_1, ..., v_{n-1}) - - # Append the current variable index (v_n) - vars_curr = (vars_prev..., col_curr) # Full tuple (v_1, ..., v_n) - - # Store info for the current level's non-zero - nz_to_indices_curr[k_lin_curr] = (orig_row_f, vars_curr) - k_lin_curr += 1 - end - end - - # --- Construct the X-derivative sparse matrix for order n (compressed or uncompressed) --- - local spX_order_n # Declare variable to hold the resulting X matrix - local X_ncols_n # Number of columns in the resulting spX_order_n matrix - - if output_compressed - # COMPRESSED output: nϵ x binomial(nx + n - 1, n) - sparse_rows_n = Int[] - sparse_cols_n = Int[] # This will store the compressed column index - sparse_vals_n = Symbolics.Num[] - - # Calculate the total number of compressed columns for order n - X_ncols_n = Int(binomial(nx + n - 1, n)) - - # Iterate through the non-zero entries of the current flat Jacobian (sp_flat_curr_X) - k_flat_curr = 1 # linear index counter for nzval of sp_flat_curr_X - for col_flat_curr = 1:size(sp_flat_curr_X, 2) # This corresponds to the n-th variable (v_n) - for i_ptr_flat_curr = sp_flat_curr_X.colptr[col_flat_curr]:(sp_flat_curr_X.colptr[col_flat_curr+1]-1) - # row_flat_curr = sp_flat_curr_X.rowval[i_ptr_flat_curr] # Row index in sp_flat_curr_X - val = sp_flat_curr_X.nzval[i_ptr_flat_curr] # The derivative value - - # Get the full info for this non-zero from the map - # The linear index in sp_flat_curr_X.nzval is k_flat_curr - orig_row_f, var_indices_full = nz_to_indices_curr[k_flat_curr] # (v_1, ..., v_n) - - # Check the compression rule: v_n <= v_{n-1} <= ... <= v_1 - is_compressed = true - for k_rule = 1:(n-1) - # Check v_{n-k_rule+1} <= v_{n-k_rule} - if var_indices_full[n-k_rule+1] > var_indices_full[n-k_rule] - is_compressed = false - break - end - end - - if is_compressed - # Calculate the compressed column index c_n for the tuple (v_1, ..., v_n) - # using the derived formula: c_n = sum_{k=1}^{n-1} binomial(v_k + n - k - 1, n - k + 1) + v_n - compressed_col_idx = 0 - for k_formula = 1:(n-1) - term = binomial(var_indices_full[k_formula] + n - k_formula - 1, n - k_formula + 1) - compressed_col_idx += term - end - # Add the last term: v_n (var_indices_full[n]) - compressed_col_idx += var_indices_full[n] - - push!(sparse_rows_n, orig_row_f) - push!(sparse_cols_n, compressed_col_idx) - push!(sparse_vals_n, val) - end - - k_flat_curr += 1 # Increment linear index counter for sp_flat_curr_X.nzval - end - end - # Construct the compressed sparse matrix for order n - spX_order_n = sparse!(sparse_rows_n, sparse_cols_n, sparse_vals_n, X_nrows_1, X_ncols_n) - - else # output_compressed == false - # UNCOMPRESSED output: nϵ x nx^n - sparse_rows_n_uncomp = Int[] - sparse_cols_n_uncomp = Int[] # Uncompressed column index (1 to nx^n) - sparse_vals_n_uncomp = Symbolics.Num[] - - # Total number of uncompressed columns - X_ncols_n = nx^n # Use BigInt for the power calculation, cast to Int - - # Iterate through the non-zero entries of the current flat Jacobian (sp_flat_curr_X) - k_flat_curr = 1 # linear index counter for nzval of sp_flat_curr_X - for col_flat_curr = 1:size(sp_flat_curr_X, 2) # This corresponds to the n-th variable (v_n) - for i_ptr_flat_curr = sp_flat_curr_X.colptr[col_flat_curr]:(sp_flat_curr_X.colptr[col_flat_curr+1]-1) - # row_flat_curr = sp_flat_curr_X.rowval[i_ptr_flat_curr] # Row index in sp_flat_curr_X - val = sp_flat_curr_X.nzval[i_ptr_flat_curr] # The derivative value - - # Get the full info for this non-zero from the map - # The linear index in sp_flat_curr_X.nzval is k_flat_curr - orig_row_f, var_indices_full = nz_to_indices_curr[k_flat_curr] # (v_1, ..., v_n) - - # Calculate the UNCOMPRESSED column index for the tuple (v_1, ..., v_n) - # This maps the tuple (v1, ..., vn) to a unique index from 1 to nx^n - # Formula: 1 + (v1-1)*nx^(n-1) + (v2-1)*nx^(n-2) + ... + (vn-1)*nx^0 - uncompressed_col_idx = 1 # 1-based - power_of_nx = nx^(n-1) # Start with nx^(n-1) for v1 term - for i = 1:n - uncompressed_col_col_idx_term = (var_indices_full[i] - 1) * power_of_nx - # Check for overflow before adding - # if (uncompressed_col_idx > 0 && uncompressed_col_col_idx_term > 0 && uncompressed_col_idx + uncompressed_col_col_idx_term <= uncompressed_col_idx) || - # (uncompressed_col_idx < 0 && uncompressed_col_col_idx_term < 0 && uncompressed_col_idx + uncompressed_col_col_idx_term >= uncompressed_col_idx) - # error("Integer overflow calculating uncompressed column index") - # end - uncompressed_col_idx += uncompressed_col_col_idx_term - - if i < n # Avoid nx^-1 - power_of_nx = div(power_of_nx, nx) # Integer division - end - end - - push!(sparse_rows_n_uncomp, orig_row_f) - push!(sparse_cols_n_uncomp, Int(uncompressed_col_idx)) # Cast to Int - push!(sparse_vals_n_uncomp, val) - - k_flat_curr += 1 # Increment linear index counter for sp_flat_curr_X.nzval - end - end - # Construct the uncompressed sparse matrix for order n - spX_order_n = sparse!(sparse_rows_n_uncomp, sparse_cols_n_uncomp, sparse_vals_n_uncomp, X_nrows_1, X_ncols_n) - - end # End of if output_compressed / else - - - # --- Compute the P-derivative sparse matrix for order n --- - # This is the Jacobian of the nzval of the intermediate flat X-Jacobian (sp_flat_curr_X) w.r.t. 𝔓. - # sp_flat_curr_X.nzval contains expressions for d^n f_i / (dx_v1 ... dx_vn) for all - # non-zero such values that were propagated from the previous step. - spP_of_flatX_nzval_curr = Symbolics.sparsejacobian(sp_flat_curr_X.nzval, vcat(𝔓[1:nps], 𝔙[1:nxs])) # nnz(sp_flat_curr_X) x np - - # Determine the desired dimensions of spP_order_n - # Dimensions are (rows of spX_order_n * cols of spX_order_n) x np - P_nrows_n = nϵ * X_ncols_n - P_ncols_n = nps + nxs - - sparse_rows_n_P = Int[] # Row index in the flattened space of spX_order_n (1 to P_nrows_n) - sparse_cols_n_P = Int[] # Column index for parameters (1 to np) - sparse_vals_n_P = Symbolics.Num[] - - # Iterate through the non-zero entries of spP_of_flatX_nzval_curr - # Its rows correspond to the non-zeros in sp_flat_curr_X - k_temp_P = 1 # linear index counter for nzval of spP_of_flatX_nzval_curr - for p_col = 1:size(spP_of_flatX_nzval_curr, 2) # Column index in spP_of_flatX_nzval_curr (corresponds to parameter index) - for i_ptr_temp_P = spP_of_flatX_nzval_curr.colptr[p_col]:(spP_of_flatX_nzval_curr.colptr[p_col+1]-1) - temp_row = spP_of_flatX_nzval_curr.rowval[i_ptr_temp_P] # Row index in spP_of_flatX_nzval_curr (corresponds to the temp_row-th nzval of sp_flat_curr_X) - p_val = spP_of_flatX_nzval_curr.nzval[i_ptr_temp_P] # The derivative w.r.t. parameter value - - # Get the full info for the X-derivative term that this P-derivative is from - # temp_row is the linear index in sp_flat_curr_X.nzval - # This corresponds to the derivative d^n f_orig_row_f / (dx_v1 ... dx_vn) - orig_row_f, var_indices_full = nz_to_indices_curr[temp_row] # (v_1, ..., v_n) - - # We need to find the column index (X_col_idx) this term corresponds to - # in the final spX_order_n matrix (which might be compressed or uncompressed) - local X_col_idx # Column index in the final spX_order_n matrix (1 to X_ncols_n) - - if output_compressed - # Calculate the compressed column index - compressed_col_idx = 0 - for k_formula = 1:(n-1) - term = binomial(var_indices_full[k_formula] + n - k_formula - 1, n - k_formula + 1) - compressed_col_idx += term - end - compressed_col_idx += var_indices_full[n] - X_col_idx = compressed_col_idx # The column in spX_order_n is the compressed one - - else # output_compressed == false - # Calculate the uncompressed column index - uncompressed_col_idx = 1 - power_of_nx = nx^(n-1) - for i = 1:n - uncompressed_col_idx += (var_indices_full[i] - 1) * power_of_nx - if i < n - power_of_nx = div(power_of_nx, nx) - end - end - X_col_idx = Int(uncompressed_col_idx) # The column in spX_order_n is the uncompressed one - end - - # Calculate the row index in spP_order_n - # This maps the (orig_row_f, X_col_idx) pair in spX_order_n's grid to a linear index - # Formula: (row_in_X - 1) * num_cols_in_X + col_in_X - # P_row_idx = (orig_row_f - 1) * X_ncols_n + X_col_idx - P_row_idx = (X_col_idx - 1) * nϵ + orig_row_f - - # The column index in spP_order_n is the parameter index - P_col_idx = p_col - - push!(sparse_rows_n_P, P_row_idx) - push!(sparse_cols_n_P, P_col_idx) - push!(sparse_vals_n_P, p_val) - - k_temp_P += 1 # Increment linear index counter for spP_of_flatX_nzval_curr.nzval - end - end - - # Construct the P-derivative sparse matrix for order n - # Dimensions are (rows of spX_order_n * cols of spX_order_n) x np - spP_order_n = sparse!(sparse_rows_n_P, sparse_cols_n_P, sparse_vals_n_P, P_nrows_n, P_ncols_n) - - # Store the pair (X-matrix, P-matrix) for order n - push!(results, (spX_order_n, spP_order_n)) - - - # Prepare for the next iteration (order n+1) - # The nzvals for the next X-Jacobian step are the nzvals of the current flat X-Jacobian - nzvals_prev = sp_flat_curr_X_rn.nzval - # The map for the next step should provide info for order n derivatives - nz_to_indices_prev = nz_to_indices_curr - - end # End of loop for orders n = 2 to max_perturbation_order - end - - return results #, (𝔛, 𝔓) # Return results as a tuple of (X_matrix, P_matrix) pairs -end - - -function write_functions_mapping!(𝓂::ℳ, max_perturbation_order::Int; - density_threshold::Float64 = .1, - min_length::Int = 1000, - nnz_parallel_threshold::Int = 1000000, - # parallel = Symbolics.SerialForm(), - # parallel = Symbolics.ShardedForm(1500,4), - cse = true, - skipzeros = true) - - future_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₁₎$"))) - present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎$"))) - past_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₋₁₎$"))) - shock_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₓ₎$"))) - ss_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₛₛ₎$"))) - - sort!(future_varss ,by = x->replace(string(x),r"₍₁₎$"=>"")) #sort by name without time index because otherwise eps_zᴸ⁽⁻¹⁾₍₋₁₎ comes before eps_z₍₋₁₎ - sort!(present_varss ,by = x->replace(string(x),r"₍₀₎$"=>"")) - sort!(past_varss ,by = x->replace(string(x),r"₍₋₁₎$"=>"")) - sort!(shock_varss ,by = x->replace(string(x),r"₍ₓ₎$"=>"")) - sort!(ss_varss ,by = x->replace(string(x),r"₍ₛₛ₎$"=>"")) - - dyn_future_list = collect(reduce(union, 𝓂.constants.post_model_macro.dyn_future_list)) - dyn_present_list = collect(reduce(union, 𝓂.constants.post_model_macro.dyn_present_list)) - dyn_past_list = collect(reduce(union, 𝓂.constants.post_model_macro.dyn_past_list)) - dyn_exo_list = collect(reduce(union,𝓂.constants.post_model_macro.dyn_exo_list)) - dyn_ss_list = Symbol.(string.(collect(reduce(union,𝓂.constants.post_model_macro.dyn_ss_list))) .* "₍ₛₛ₎") - - future = map(x -> Symbol(replace(string(x), r"₍₁₎" => "")),string.(dyn_future_list)) - present = map(x -> Symbol(replace(string(x), r"₍₀₎" => "")),string.(dyn_present_list)) - past = map(x -> Symbol(replace(string(x), r"₍₋₁₎" => "")),string.(dyn_past_list)) - exo = map(x -> Symbol(replace(string(x), r"₍ₓ₎" => "")),string.(dyn_exo_list)) - stst = map(x -> Symbol(replace(string(x), r"₍ₛₛ₎" => "")),string.(dyn_ss_list)) - - vars_raw = vcat(dyn_future_list[indexin(sort(future),future)], - dyn_present_list[indexin(sort(present),present)], - dyn_past_list[indexin(sort(past),past)], - dyn_exo_list[indexin(sort(exo),exo)]) - - dyn_var_future_idx = 𝓂.constants.post_complete_parameters.dyn_var_future_idx - dyn_var_present_idx = 𝓂.constants.post_complete_parameters.dyn_var_present_idx - dyn_var_past_idx = 𝓂.constants.post_complete_parameters.dyn_var_past_idx - dyn_ss_idx = 𝓂.constants.post_complete_parameters.dyn_ss_idx - - dyn_var_idxs = vcat(dyn_var_future_idx, dyn_var_present_idx, dyn_var_past_idx) - - pars_ext = vcat(𝓂.constants.post_complete_parameters.parameters, 𝓂.equations.calibration_parameters) - parameters_and_SS = vcat(pars_ext, dyn_ss_list[indexin(sort(stst),stst)]) - - np = length(parameters_and_SS) - nv = length(vars_raw) - nc = length(𝓂.equations.calibration) - nps = length(𝓂.constants.post_complete_parameters.parameters) - nxs = maximum(dyn_var_idxs) + nc - - Symbolics.@variables 𝔓[1:np] 𝔙[1:nv] - - parameter_dict = Dict{Symbol, Symbol}() - back_to_array_dict = Dict{Symbolics.Num, Symbolics.Num}() - calib_vars = Symbol[] - calib_expr = [] - SS_mapping = Dict{Symbolics.Num, Symbolics.Num}() - - - for (i,v) in enumerate(parameters_and_SS) - push!(parameter_dict, v => :($(Symbol("𝔓_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("𝔓_$i"))), @__MODULE__) => 𝔓[i]) - if i > nps - if i > length(pars_ext) - push!(SS_mapping, 𝔓[i] => 𝔙[dyn_ss_idx[i-length(pars_ext)]]) - else - push!(SS_mapping, 𝔓[i] => 𝔙[nxs + i - nps - nc]) - end - end - end - - for (i,v) in enumerate(vars_raw) - push!(parameter_dict, v => :($(Symbol("𝔙_$i")))) - push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("𝔙_$i"))), @__MODULE__) => 𝔙[i]) - if i <= length(dyn_var_idxs) - push!(SS_mapping, 𝔙[i] => 𝔙[dyn_var_idxs[i]]) - else - push!(SS_mapping, 𝔙[i] => 0) - end - end - - - for v in 𝓂.equations.calibration_no_var - push!(calib_vars, v.args[1]) - push!(calib_expr, v.args[2]) - 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 - - - dyn_equations = 𝓂.equations.dynamic |> - 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)) - - derivatives = take_nth_order_derivatives(dyn_equations, 𝔙, 𝔓, SS_mapping, nps, nxs) - - - ∇₁_dyn = derivatives[1][1] - - lennz = nnz(∇₁_dyn) - - if (lennz / length(∇₁_dyn) > density_threshold) || (length(∇₁_dyn) < min_length) - derivatives_mat = convert(Matrix, ∇₁_dyn) - buffer = zeros(Float64, size(∇₁_dyn)) - else - derivatives_mat = ∇₁_dyn - buffer = similar(∇₁_dyn, 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, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.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 - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_∇₁_parameters = Symbolics.build_function(∇₁_parameters_mat, 𝔓, 𝔙, - cse = cse, - skipzeros = skipzeros, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.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 - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_∇₁_SS_and_pars = Symbolics.build_function(∇₁_SS_and_pars_mat, 𝔓, 𝔙, - cse = cse, - skipzeros = skipzeros, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.caches.jacobian_SS_and_pars = buffer_SS_and_pars - - # Create jacobian_functions struct with all three functions - 𝓂.functions.jacobian = jacobian_functions(func_exprs, func_∇₁_parameters, func_∇₁_SS_and_pars) - - - - - # if max_perturbation_order >= 1 - # 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(eqs) - - # 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) - # println(ϵˢ) - # ∂SS_equations_∂parameters = Symbolics.sparsejacobian(ϵˢ, 𝔛¹) # 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, - # parallel = parallel, - # # nanmath = false, - # 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(ϵˢ, 𝔓¹) # nϵ x nx - - # 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) - # buffer = zeros(Float64, size(∂SS_equations_∂SS_and_pars)) - # else - # 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, - # parallel = parallel, - # # nanmath = false, - # 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) - - if 𝓂.constants.second_order.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) - 𝓂.constants.second_order = create_second_order_auxiliary_matrices(𝓂.constants) - - ∇₂_dyn = derivatives[2][1] - - lennz = nnz(∇₂_dyn) - - if (lennz / length(∇₂_dyn) > density_threshold) || (length(∇₂_dyn) < min_length) - derivatives_mat = convert(Matrix, ∇₂_dyn) - buffer = zeros(Float64, size(∇₂_dyn)) - else - derivatives_mat = ∇₂_dyn - buffer = similar(∇₂_dyn, 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, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.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 - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_∇₂_parameters = Symbolics.build_function(∇₂_parameters_mat, 𝔓, 𝔙, - cse = cse, - skipzeros = skipzeros, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.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 - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_∇₂_SS_and_pars = Symbolics.build_function(∇₂_SS_and_pars_mat, 𝔓, 𝔙, - cse = cse, - skipzeros = skipzeros, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.caches.hessian_SS_and_pars = buffer_SS_and_pars - - # Create hessian_functions struct with all three functions - 𝓂.functions.hessian = hessian_functions(func_exprs, func_∇₂_parameters, func_∇₂_SS_and_pars) - end - end - - if max_perturbation_order == 3 - derivatives = take_nth_order_derivatives(dyn_equations, 𝔙, 𝔓, SS_mapping, nps, nxs; max_perturbation_order = max_perturbation_order, output_compressed = true) - # third order - if 𝓂.constants.third_order.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) - I,J,V = findnz(derivatives[3][1]) - 𝓂.constants.third_order = create_third_order_auxiliary_matrices(𝓂.constants, unique(J)) - - ∇₃_dyn = derivatives[3][1] - - lennz = nnz(∇₃_dyn) - - if (lennz / length(∇₃_dyn) > density_threshold) || (length(∇₃_dyn) < min_length) - derivatives_mat = convert(Matrix, ∇₃_dyn) - buffer = zeros(Float64, size(∇₃_dyn)) - else - derivatives_mat = ∇₃_dyn - buffer = similar(∇₃_dyn, 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, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.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 - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_∇₃_parameters = Symbolics.build_function(∇₃_parameters_mat, 𝔓, 𝔙, - cse = cse, - skipzeros = skipzeros, - parallel = parallel, - # nanmath = false, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.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 - - if lennz > nnz_parallel_threshold - parallel = Symbolics.ShardedForm(1500,4) - else - parallel = Symbolics.SerialForm() - end - - _, func_∇₃_SS_and_pars = Symbolics.build_function(∇₃_SS_and_pars_mat, 𝔓, 𝔙, - cse = cse, - skipzeros = skipzeros, - # nanmath = false, - parallel = parallel, - expression_module = @__MODULE__, - expression = Val(false))::Tuple{<:Function, <:Function} - - 𝓂.caches.third_order_derivatives_SS_and_pars = buffer_SS_and_pars - - # Create third_order_derivatives_functions struct with all three functions - 𝓂.functions.third_order_derivatives = third_order_derivatives_functions(func_exprs, func_∇₃_parameters, func_∇₃_SS_and_pars) - end - end - - return nothing -end - - -function write_auxiliary_indices!(𝓂::ℳ) - # write indices in auxiliary objects - dyn_var_future_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₁₎"))) - dyn_var_present_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₀₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎"))) - dyn_var_past_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₋₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₋₁₎"))) - dyn_exo_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₓ₎"))) - dyn_ss_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₛₛ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₛₛ₎"))) - - dyn_var_future = Symbol.(string.(sort(collect(reduce(union,dyn_var_future_list))))) - dyn_var_present = Symbol.(string.(sort(collect(reduce(union,dyn_var_present_list))))) - dyn_var_past = Symbol.(string.(sort(collect(reduce(union,dyn_var_past_list))))) - dyn_exo = Symbol.(string.(sort(collect(reduce(union,dyn_exo_list))))) - dyn_ss = Symbol.(string.(sort(collect(reduce(union,dyn_ss_list))))) - - SS_and_pars_names = vcat(Symbol.(string.(sort(union(𝓂.constants.post_model_macro.var,𝓂.constants.post_model_macro.exo_past,𝓂.constants.post_model_macro.exo_future)))), 𝓂.equations.calibration_parameters) - - dyn_var_future_idx = indexin(dyn_var_future , SS_and_pars_names) - dyn_var_present_idx = indexin(dyn_var_present , SS_and_pars_names) - dyn_var_past_idx = indexin(dyn_var_past , SS_and_pars_names) - dyn_ss_idx = indexin(dyn_ss , SS_and_pars_names) - - shocks_ss = zeros(length(dyn_exo)) - - 𝓂.constants.post_complete_parameters = update_post_complete_parameters( - 𝓂.constants.post_complete_parameters; - dyn_var_future_idx = dyn_var_future_idx, - dyn_var_present_idx = dyn_var_present_idx, - dyn_var_past_idx = dyn_var_past_idx, - dyn_ss_idx = dyn_ss_idx, - shocks_ss = shocks_ss, - ) - - return nothing -end - -write_parameters_input!(𝓂::ℳ, parameters::Nothing; verbose::Bool = true) = return parameters -write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Pair{S,Float64}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}(parameters[1] |> Meta.parse |> replace_indices => parameters[2]), verbose = verbose) - - - -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Float64},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict(parameters), verbose = verbose) -# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Union{Symbol,AbstractString},Union{Float64,Int}},Vararg{Pair{Union{Symbol,AbstractString},Union{Float64,Int}}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -# write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{AbstractString,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Dict(parameters), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{S,Float64},Vararg{Pair{S,Float64}}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]) -, verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol, Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{S, Float64}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol, Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Dict{S,Float64}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((keys(parameters) .|> Meta.parse .|> replace_indices) .=> values(parameters)), verbose = verbose) - - -write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(parameters[1])) => parameters[2]]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Pair{S,Int}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Int},Vararg{Pair{Symbol,Int}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{S,Int},Vararg{Pair{S,Int}}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, Dict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{S, Int}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Dict{S,Int}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((keys(parameters) .|> Meta.parse .|> replace_indices) .=> values(parameters)), verbose = verbose) - - -write_parameters_input!(𝓂::ℳ, parameters::Pair{Symbol,Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(parameters[1])) => parameters[2]]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Pair{S,Real}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((parameters[1] |> Meta.parse |> replace_indices) => parameters[2]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Dict{S,Real}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}((keys(parameters) .|> Meta.parse .|> replace_indices) .=> values(parameters)), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{Symbol,Real},Vararg{Pair{Symbol,Float64}}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Pair{S,Real},Vararg{Pair{S,Float64}}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{Symbol, Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([replace_indices(string(i[1])) => i[2] for i in parameters]), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Vector{Pair{S, Real}}; verbose::Bool = true) where S <: AbstractString = write_parameters_input!(𝓂::ℳ, OrderedDict{Symbol,Float64}([i[1] |> Meta.parse |> replace_indices => i[2] for i in parameters]), verbose = verbose) - - - -function write_parameters_input!(𝓂::ℳ, parameters::D; verbose::Bool = true) where D <: AbstractDict{Symbol,Float64} - # Handle missing parameters - add them if they are in the missing_parameters list - p = 𝓂.constants.post_complete_parameters - missing_params_provided = intersect(collect(keys(parameters)), p.missing_parameters) - - if !isempty(missing_params_provided) - - # 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 - end - - # Amend parameter order by provided missing params - # declared_params = parameters that were never missing (have non-NaN values) - # We identify them as parameters that are not in the union of missing_params_provided and still-missing params - all_missing = union(missing_params_provided, remaining_missing) - declared_params = setdiff(p.parameters, all_missing) - - # Get the current parameter values for declared params - declared_param_indices = indexin(declared_params, p.parameters) - declared_values = 𝓂.parameter_values[declared_param_indices] - - # Get values for the newly provided missing params (currently NaN in parameter_values) - # We'll set them later after the bounds check - missing_values = fill(NaN, length(missing_params_provided)) - - # Get values for the remaining missing params (still NaN) - remaining_missing_values = fill(NaN, length(remaining_missing)) - - # Reorder both parameters and parameter_values arrays - new_parameters = vcat(declared_params, collect(missing_params_provided), remaining_missing) - 𝓂.constants.post_complete_parameters = update_post_complete_parameters( - p; - parameters = new_parameters, - missing_parameters = remaining_missing, - ) - 𝓂.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) - end - end - - # Handle remaining parameters (not missing ones) - p = 𝓂.constants.post_complete_parameters - if length(setdiff(collect(keys(parameters)), p.parameters))>0 - @warn("Parameters not part of the model are ignored: $(setdiff(collect(keys(parameters)),p.parameters))") - for kk in setdiff(collect(keys(parameters)), p.parameters) - delete!(parameters,kk) - end - end - - bounds_broken = false - - for (par,val) in parameters - 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") - bounds_broken = true - continue - end - if val < 𝓂.constants.post_parameters_macro.bounds[par][1] - @warn("Calibration is out of bounds for $par > $(𝓂.constants.post_parameters_macro.bounds[par][1])\t parameter value: $val") - bounds_broken = true - continue - end - end - end - - if bounds_broken - @warn("Parameters unchanged.") - else - ntrsct_idx = map(x-> getindex(1:length(𝓂.parameter_values), p.parameters .== x)[1], collect(keys(parameters))) - # ntrsct_idx = indexin(collect(keys(parameters)), p.parameters) - - 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] - end - end - end - - if 𝓂.caches.outdated.non_stochastic_steady_state == true && verbose println("New parameters changed the steady state.") end - - return nothing -end - - -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Int,Vararg{Int}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Matrix{Int}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) - -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Float64,Vararg{Float64}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, vec(collect(parameters)), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Matrix{Float64}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, vec(collect(parameters)), verbose = verbose) - -write_parameters_input!(𝓂::ℳ, parameters::Tuple{Real,Vararg{Real}}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) -write_parameters_input!(𝓂::ℳ, parameters::Matrix{Real}; verbose::Bool = true) = write_parameters_input!(𝓂::ℳ, Float64.(vec(collect(parameters))), verbose = verbose) - - - -function write_parameters_input!(𝓂::ℳ, parameters::Vector{Float64}; verbose::Bool = true) - if length(parameters) > length(𝓂.parameter_values) - @warn "Model has $(length(𝓂.parameter_values)) parameters. $(length(parameters)) were provided. The following will be ignored: $(join(parameters[length(𝓂.parameter_values)+1:end], " "))" - - parameters = parameters[1:length(𝓂.parameter_values)] - end - - bounds_broken = false - - for (par,val) in Dict(𝓂.constants.post_complete_parameters.parameters .=> parameters) - 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") - bounds_broken = true - continue - end - if val < 𝓂.constants.post_parameters_macro.bounds[par][1] - @warn("Calibration is out of bounds for $par > $(𝓂.constants.post_parameters_macro.bounds[par][1])\t parameter value: $val") - bounds_broken = true - continue - end - end - end - - if bounds_broken - @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] - push!(match_idx,i) - end - end - - changed_vals = parameters[match_idx] - changed_pars = 𝓂.constants.post_complete_parameters.parameters[match_idx] - - if verbose - println("Parameter changes: ") - for (i,m) in enumerate(match_idx) - println("\t",changed_pars[i],"\tfrom ",𝓂.parameter_values[m],"\tto ",changed_vals[i]) - end - end - - 𝓂.parameter_values[match_idx] = parameters[match_idx] - end - end - - if 𝓂.caches.outdated.non_stochastic_steady_state == true && verbose println("New parameters changed the steady state.") end - - return nothing -end - - -# function create_timings_for_estimation!(𝓂::ℳ, observables::Vector{Symbol}) -# dyn_equations = 𝓂.dyn_equations - -# vars_to_exclude = setdiff(𝓂.constants.post_model_macro.present_only, observables) - -# # Mapping variables to their equation index -# variable_to_equation = Dict{Symbol, Vector{Int}}() -# for var in vars_to_exclude -# for (eq_idx, vars_set) in enumerate(𝓂.dyn_var_present_list) -# # for var in vars_set -# if var in vars_set -# if haskey(variable_to_equation, var) -# push!(variable_to_equation[var],eq_idx) -# else -# variable_to_equation[var] = [eq_idx] -# end -# end -# end -# end - -# # cols_to_exclude = indexin(𝓂.constants.post_model_macro.var, setdiff(𝓂.constants.post_model_macro.present_only, observables)) -# cols_to_exclude = indexin(setdiff(𝓂.constants.post_model_macro.present_only, observables), 𝓂.constants.post_model_macro.var) - -# present_idx = 𝓂.constants.post_model_macro.nFuture_not_past_and_mixed .+ (setdiff(range(1, 𝓂.constants.post_model_macro.nVars), cols_to_exclude)) - -# dyn_var_future_list = deepcopy(𝓂.dyn_var_future_list) -# dyn_var_present_list = deepcopy(𝓂.dyn_var_present_list) -# dyn_var_past_list = deepcopy(𝓂.dyn_var_past_list) -# dyn_exo_list = deepcopy(𝓂.dyn_exo_list) -# dyn_ss_list = deepcopy(𝓂.dyn_ss_list) - -# rows_to_exclude = Int[] - -# for vidx in values(variable_to_equation) -# for v in vidx -# if v ∉ rows_to_exclude -# push!(rows_to_exclude, v) - -# for vv in vidx -# dyn_var_future_list[vv] = union(dyn_var_future_list[vv], dyn_var_future_list[v]) -# dyn_var_present_list[vv] = union(dyn_var_present_list[vv], dyn_var_present_list[v]) -# dyn_var_past_list[vv] = union(dyn_var_past_list[vv], dyn_var_past_list[v]) -# dyn_exo_list[vv] = union(dyn_exo_list[vv], dyn_exo_list[v]) -# dyn_ss_list[vv] = union(dyn_ss_list[vv], dyn_ss_list[v]) -# end - -# break -# end -# end -# end - -# rows_to_include = setdiff(1:𝓂.constants.post_model_macro.nVars, rows_to_exclude) - -# all_symbols = setdiff(reduce(union,collect.(get_symbols.(dyn_equations)))[rows_to_include], vars_to_exclude) -# parameters_in_equations = sort(setdiff(all_symbols, match_pattern(all_symbols,r"₎$"))) - -# dyn_var_future = sort(setdiff(collect(reduce(union,dyn_var_future_list[rows_to_include])), vars_to_exclude)) -# dyn_var_present = sort(setdiff(collect(reduce(union,dyn_var_present_list[rows_to_include])), vars_to_exclude)) -# dyn_var_past = sort(setdiff(collect(reduce(union,dyn_var_past_list[rows_to_include])), vars_to_exclude)) -# dyn_var_ss = sort(setdiff(collect(reduce(union,dyn_ss_list[rows_to_include])), vars_to_exclude)) - -# all_dyn_vars = union(dyn_var_future, dyn_var_present, dyn_var_past) - -# @assert length(setdiff(dyn_var_ss, all_dyn_vars)) == 0 "The following variables are (and cannot be) defined only in steady state (`[ss]`): $(setdiff(dyn_var_ss, all_dyn_vars))" - -# all_vars = union(all_dyn_vars, dyn_var_ss) - -# present_only = sort(setdiff(dyn_var_present,union(dyn_var_past,dyn_var_future))) -# future_not_past = sort(setdiff(dyn_var_future, dyn_var_past)) -# past_not_future = sort(setdiff(dyn_var_past, dyn_var_future)) -# mixed = sort(setdiff(dyn_var_present, union(present_only, future_not_past, past_not_future))) -# future_not_past_and_mixed = sort(union(future_not_past,mixed)) -# past_not_future_and_mixed = sort(union(past_not_future,mixed)) -# present_but_not_only = sort(setdiff(dyn_var_present,present_only)) -# mixed_in_past = sort(intersect(dyn_var_past, mixed)) -# not_mixed_in_past = sort(setdiff(dyn_var_past,mixed_in_past)) -# mixed_in_future = sort(intersect(dyn_var_future, mixed)) -# exo = sort(collect(reduce(union,dyn_exo_list))) -# var = sort(dyn_var_present) -# aux_tmp = sort(filter(x->occursin(r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾",string(x)), dyn_var_present)) -# aux = aux_tmp[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∉ exo, aux_tmp)] -# exo_future = dyn_var_future[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_future)] -# exo_present = dyn_var_present[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_present)] -# exo_past = dyn_var_past[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_past)] - -# nPresent_only = length(present_only) -# nMixed = length(mixed) -# nFuture_not_past_and_mixed = length(future_not_past_and_mixed) -# nPast_not_future_and_mixed = length(past_not_future_and_mixed) -# nPresent_but_not_only = length(present_but_not_only) -# nVars = length(all_vars) -# nExo = length(collect(exo)) - -# present_only_idx = indexin(present_only,var) -# present_but_not_only_idx = indexin(present_but_not_only,var) -# future_not_past_and_mixed_idx = indexin(future_not_past_and_mixed,var) -# past_not_future_and_mixed_idx = indexin(past_not_future_and_mixed,var) -# mixed_in_future_idx = indexin(mixed_in_future,dyn_var_future) -# mixed_in_past_idx = indexin(mixed_in_past,dyn_var_past) -# not_mixed_in_past_idx = indexin(not_mixed_in_past,dyn_var_past) -# past_not_future_idx = indexin(past_not_future,var) - -# reorder = indexin(var, [present_only; past_not_future; future_not_past_and_mixed]) -# dynamic_order = indexin(present_but_not_only, [past_not_future; future_not_past_and_mixed]) - -# @assert length(intersect(union(var,exo),parameters_in_equations)) == 0 "Parameters and variables cannot have the same name. This is the case for: " * repr(sort([intersect(union(var,exo),parameters_in_equations)...])) - -# T = timings(present_only, -# future_not_past, -# past_not_future, -# mixed, -# future_not_past_and_mixed, -# past_not_future_and_mixed, -# present_but_not_only, -# mixed_in_past, -# not_mixed_in_past, -# mixed_in_future, -# exo, -# var, -# aux, -# exo_present, - -# nPresent_only, -# nMixed, -# nFuture_not_past_and_mixed, -# nPast_not_future_and_mixed, -# nPresent_but_not_only, -# nVars, -# nExo, - -# present_only_idx, -# present_but_not_only_idx, -# future_not_past_and_mixed_idx, -# not_mixed_in_past_idx, -# past_not_future_and_mixed_idx, -# mixed_in_past_idx, -# mixed_in_future_idx, -# past_not_future_idx, - -# reorder, -# dynamic_order) - -# push!(𝓂.estimation_helper, observables => T) - -# return nothing -# end - - - -function calculate_jacobian(parameters::Vector{M}, - SS_and_pars::Vector{N}, - caches_obj::caches, - jacobian_funcs::jacobian_functions)::Matrix{M} where {M,N} - if eltype(caches_obj.jacobian) != M - if caches_obj.jacobian isa SparseMatrixCSC - jac_buffer = similar(caches_obj.jacobian,M) - jac_buffer.nzval .= 0 - else - jac_buffer = zeros(M, size(caches_obj.jacobian)) - end - else - jac_buffer = caches_obj.jacobian - end - - jacobian_funcs.f(jac_buffer, parameters, SS_and_pars) - - return jac_buffer -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} - if eltype(caches_obj.hessian) != M - if caches_obj.hessian isa SparseMatrixCSC - hes_buffer = similar(caches_obj.hessian,M) - hes_buffer.nzval .= 0 - else - hes_buffer = zeros(M, size(caches_obj.hessian)) - end - else - hes_buffer = caches_obj.hessian - end - - hessian_funcs.f(hes_buffer, parameters, SS_and_pars) - - return hes_buffer -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} - 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) - third_buffer.nzval .= 0 - else - third_buffer = zeros(M, size(caches_obj.third_order_derivatives)) - end - else - third_buffer = caches_obj.third_order_derivatives - end - - third_order_derivatives_funcs.f(third_buffer, parameters, SS_and_pars) - - return third_buffer -end - - - -function compute_irf_responses(𝓂::ℳ, - state_update::Function, - initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, - level::Vector{Float64}; - periods::Int, - shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}}, - variables::Union{Symbol_input,String_input}, - shock_size::Real, - negative_shock::Bool, - generalised_irf::Bool, - generalised_irf_warmup_iterations::Int, - generalised_irf_draws::Int, - enforce_obc::Bool, - algorithm::Symbol) - - if enforce_obc - function obc_state_update(present_states, present_shocks::Vector{R}, state_update::Function) where R <: Float64 - unconditional_forecast_horizon = 𝓂.constants.post_model_macro.max_obc_horizon - - reference_ss = 𝓂.caches.non_stochastic_steady_state - - obc_shock_idx = contains.(string.(𝓂.constants.post_model_macro.exo),"ᵒᵇᶜ") - - periods_per_shock = 𝓂.constants.post_model_macro.max_obc_horizon + 1 - - num_shocks = sum(obc_shock_idx) ÷ periods_per_shock - - p = (present_states, state_update, reference_ss, 𝓂, algorithm, unconditional_forecast_horizon, present_shocks) - - constraints_violated = any(𝓂.functions.obc_violation(zeros(num_shocks*periods_per_shock), p) .> eps(Float32)) - - if constraints_violated - opt = NLopt.Opt(NLopt.:LD_SLSQP, num_shocks*periods_per_shock) - - opt.min_objective = obc_objective_optim_fun - - opt.xtol_abs = eps(Float32) - opt.ftol_abs = eps(Float32) - opt.maxeval = 500 - - upper_bounds = fill(eps(), 1 + 2*(max(num_shocks*periods_per_shock-1, 1))) - - NLopt.inequality_constraint!(opt, (res, x, jac) -> obc_constraint_optim_fun(res, x, jac, p), upper_bounds) - - (minf,x,ret) = NLopt.optimize(opt, zeros(num_shocks*periods_per_shock)) - - present_shocks[contains.(string.(𝓂.constants.post_model_macro.exo),"ᵒᵇᶜ")] .= x - - constraints_violated = any(𝓂.functions.obc_violation(x, p) .> eps(Float32)) - - solved = !constraints_violated - else - solved = true - end - - present_states = state_update(present_states, present_shocks) - - return present_states, present_shocks, solved - end - - if generalised_irf - return girf(state_update, - obc_state_update, - initial_state, - level, - 𝓂.constants; - periods = periods, - shocks = shocks, - shock_size = shock_size, - variables = variables, - negative_shock = negative_shock, - warmup_periods = generalised_irf_warmup_iterations, - draws = generalised_irf_draws) - else - return irf(state_update, - obc_state_update, - initial_state, - level, - 𝓂.constants; - periods = periods, - shocks = shocks, - shock_size = shock_size, - variables = variables, - negative_shock = negative_shock) - end - else - if generalised_irf - return girf(state_update, - initial_state, - level, - 𝓂.constants; - periods = periods, - shocks = shocks, - shock_size = shock_size, - variables = variables, - negative_shock = negative_shock, - warmup_periods = generalised_irf_warmup_iterations, - draws = generalised_irf_draws) - else - return irf(state_update, - initial_state, - level, - 𝓂.constants; - periods = periods, - shocks = shocks, - shock_size = shock_size, - variables = variables, - negative_shock = negative_shock) - end - end -end - - -function irf(state_update::Function, - obc_state_update::Function, - initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, - level::Vector{Float64}, - constants::constants; - periods::Int = 40, - shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, - variables::Union{Symbol_input,String_input} = :all, - shock_size::Real = 1, - negative_shock::Bool = false)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} - T = constants.post_model_macro - - pruning = initial_state isa Vector{Vector{Float64}} - - shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks - - shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks - - if shocks isa Matrix{Float64} - @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model." - - # periods += size(shocks)[2] - - shock_history = zeros(T.nExo, periods) - - shock_history[:,1:size(shocks)[2]] = shocks - - shock_idx = 1 - elseif shocks isa KeyedArray{Float64} - shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) - - # periods += size(shocks)[2] - - @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." - - shock_history = zeros(T.nExo, periods) - - shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks - - shock_idx = 1 - else - shock_idx = parse_shocks_input_to_index(shocks,constants) - end - - var_idx = parse_variables_input_to_index(variables, constants) |> sort - - axis1 = T.var[var_idx] - - if any(x -> contains(string(x), "◖"), axis1) - axis1_decomposed = decompose_name.(axis1) - 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 - - always_solved = true - - if shocks == :simulate - shock_history = randn(T.nExo,periods) * shock_size - - shock_history[contains.(string.(T.exo),"ᵒᵇᶜ"),:] .= 0 - - Y = zeros(T.nVars,periods,1) - - past_states = initial_state - - for t in 1:periods - past_states, past_shocks, solved = obc_state_update(past_states, shock_history[:,t], state_update) - - if !solved @warn "No solution in period: $t" end#. Possible reasons: 1. infeasability 2. too long spell of binding constraint. To address the latter try setting max_obc_horizon to a larger value (default: 40): @model max_obc_horizon=40 begin ... end" end - - always_solved = always_solved && solved - - if !always_solved break end - - Y[:,t,1] = pruning ? sum(past_states) : past_states - - shock_history[:,t] = past_shocks - end - - return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:simulate]) - elseif shocks == :none - Y = zeros(T.nVars,periods,1) - - shck = T.nExo == 0 ? Vector{Float64}(undef, 0) : zeros(T.nExo) - - past_states = initial_state - - for t in 1:periods - past_states, _, solved = obc_state_update(past_states, shck, state_update) - - if !solved @warn "No solution in period: $t" end#. Possible reasons: 1. infeasability 2. too long spell of binding constraint. To address the latter try setting max_obc_horizon to a larger value (default: 40): @model max_obc_horizon=40 begin ... end" end - - always_solved = always_solved && solved - - if !always_solved break end - - Y[:,t,1] = pruning ? sum(past_states) : past_states - end - - return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:none]) - else - Y = zeros(T.nVars,periods,length(shock_idx)) - - for (i,ii) in enumerate(shock_idx) - if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input,String_input} - shock_history = zeros(T.nExo,periods) - shock_history[ii,1] = negative_shock ? -shock_size : shock_size - end - - past_states = initial_state - - for t in 1:periods - past_states, past_shocks, solved = obc_state_update(past_states, shock_history[:,t], state_update) - - if !solved @warn "No solution in period: $t" end#. Possible reasons: 1. infeasability 2. too long spell of binding constraint. To address the latter try setting max_obc_horizon to a larger value (default: 40): @model max_obc_horizon=40 begin ... end" end - - always_solved = always_solved && solved - - if !always_solved break end - - Y[:,t,i] = pruning ? sum(past_states) : past_states - - shock_history[:,t] = past_shocks - end - end - - axis2 = shocks isa Union{Symbol_input,String_input} ? - shock_idx isa Int ? - [T.exo[shock_idx]] : - T.exo[shock_idx] : - [:Shock_matrix] - - if any(x -> contains(string(x), "◖"), axis2) - axis2_decomposed = decompose_name.(axis2) - axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] - end - - return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) - end -end - - - - -function irf(state_update::Function, - initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, - level::Vector{Float64}, - constants::constants; - periods::Int = 40, - shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, - variables::Union{Symbol_input,String_input} = :all, - shock_size::Real = 1, - negative_shock::Bool = false)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} - T = constants.post_model_macro - - pruning = initial_state isa Vector{Vector{Float64}} - - shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks - - shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks - - if shocks isa Matrix{Float64} - @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model." - - # periods += size(shocks)[2] - - shock_history = zeros(T.nExo, periods) - - shock_history[:,1:size(shocks)[2]] = shocks - - shock_idx = 1 - elseif shocks isa KeyedArray{Float64} - shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) - - # periods += size(shocks)[2] - - @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." - - shock_history = zeros(T.nExo, periods) - - shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks - - shock_idx = 1 - else - shock_idx = parse_shocks_input_to_index(shocks,constants) - end - - var_idx = parse_variables_input_to_index(variables, constants) |> sort - - axis1 = T.var[var_idx] - - if any(x -> contains(string(x), "◖"), axis1) - axis1_decomposed = decompose_name.(axis1) - 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 - - if shocks == :simulate - shock_history = randn(T.nExo,periods) * shock_size - - shock_history[contains.(string.(T.exo),"ᵒᵇᶜ"),:] .= 0 - - Y = zeros(T.nVars,periods,1) - - initial_state = state_update(initial_state,shock_history[:,1]) - - Y[:,1,1] = pruning ? sum(initial_state) : initial_state - - for t in 1:periods-1 - initial_state = state_update(initial_state,shock_history[:,t+1]) - - Y[:,t+1,1] = pruning ? sum(initial_state) : initial_state - end - - return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:simulate]) - elseif shocks == :none - Y = zeros(T.nVars,periods,1) - - shck = T.nExo == 0 ? Vector{Float64}(undef, 0) : zeros(T.nExo) - - initial_state = state_update(initial_state, shck) - - Y[:,1,1] = pruning ? sum(initial_state) : initial_state - - for t in 1:periods-1 - initial_state = state_update(initial_state, shck) - - Y[:,t+1,1] = pruning ? sum(initial_state) : initial_state - end - - return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:none]) - else - Y = zeros(T.nVars,periods,length(shock_idx)) - - for (i,ii) in enumerate(shock_idx) - initial_state_copy = deepcopy(initial_state) - - if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input,String_input} - shock_history = zeros(T.nExo,periods) - shock_history[ii,1] = negative_shock ? -shock_size : shock_size - end - - initial_state_copy = state_update(initial_state_copy, shock_history[:,1]) - - Y[:,1,i] = pruning ? sum(initial_state_copy) : initial_state_copy - - for t in 1:periods-1 - initial_state_copy = state_update(initial_state_copy, shock_history[:,t+1]) - - Y[:,t+1,i] = pruning ? sum(initial_state_copy) : initial_state_copy - end - end - - axis2 = shocks isa Union{Symbol_input,String_input} ? - shock_idx isa Int ? - [T.exo[shock_idx]] : - T.exo[shock_idx] : - [:Shock_matrix] - - if any(x -> contains(string(x), "◖"), axis2) - axis2_decomposed = decompose_name.(axis2) - axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] - end - - return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) - end -end - - - -function girf(state_update::Function, - initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, - level::Vector{Float64}, - constants::constants; - periods::Int = 40, - shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, - variables::Union{Symbol_input,String_input} = :all, - shock_size::Real = 1, - negative_shock::Bool = false, - warmup_periods::Int = 100, - draws::Int = 50)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} - T = constants.post_model_macro - - pruning = initial_state isa Vector{Vector{Float64}} - - shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks - - shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks - - if shocks isa Matrix{Float64} - @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model (model has $(T.nExo) shocks)." - - # periods += size(shocks)[2] - - shock_history = zeros(T.nExo, periods) - - shock_history[:,1:size(shocks)[2]] = shocks - - shock_idx = 1 - elseif shocks isa KeyedArray{Float64} - shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) - - # periods += size(shocks)[2] - - @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." - - shock_history = zeros(T.nExo, periods + 1) - - shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks - - shock_idx = 1 - elseif shocks == :simulate - shock_history = randn(T.nExo,periods) * shock_size - - shock_idx = 1 + if bounds_broken + @warn("Parameters unchanged.") else - shock_idx = parse_shocks_input_to_index(shocks,constants) - end - - var_idx = parse_variables_input_to_index(variables, constants) |> sort - - Y = zeros(T.nVars, periods + 1, length(shock_idx)) - - for (i,ii) in enumerate(shock_idx) - initial_state_copy = deepcopy(initial_state) - - accepted_draws = 0 - - for draw in 1:draws - ok = true - - initial_state_copy² = deepcopy(initial_state_copy) - - for i in 1:warmup_periods - initial_state_copy² = state_update(initial_state_copy², randn(T.nExo)) - if any(!isfinite, [x for v in initial_state_copy² for x in v]) - # @warn "No solution in warmup period: $i" - ok = false - break + if !all(parameters .== 𝓂.parameter_values[1:length(parameters)]) + match_idx = [] + for (i, v) in enumerate(parameters) + if v != 𝓂.parameter_values[i] + push!(match_idx,i) end end - if !ok continue end - - Y₁ = zeros(T.nVars, periods + 1) - Y₂ = zeros(T.nVars, periods + 1) - - baseline_noise = randn(T.nExo) - - if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input,String_input} - shock_history = zeros(T.nExo,periods) - shock_history[ii,1] = negative_shock ? -shock_size : shock_size - end - - if pruning - initial_state_copy² = state_update(initial_state_copy², baseline_noise) - - if any(!isfinite, [x for v in initial_state_copy² for x in v]) continue end - - initial_state₁ = deepcopy(initial_state_copy²) - initial_state₂ = deepcopy(initial_state_copy²) - - Y₁[:,1] = initial_state_copy² |> sum - Y₂[:,1] = initial_state_copy² |> sum - else - Y₁[:,1] = state_update(initial_state_copy², baseline_noise) - - if any(!isfinite, Y₁[:,1]) continue end - - Y₂[:,1] = state_update(initial_state_copy², baseline_noise) - - if any(!isfinite, Y₂[:,1]) continue end - end - - for t in 1:periods - baseline_noise = randn(T.nExo) - - if pruning - initial_state₁ = state_update(initial_state₁, baseline_noise) - - if any(!isfinite, [x for v in initial_state₁ for x in v]) - ok = false - break - end - - initial_state₂ = state_update(initial_state₂, baseline_noise + shock_history[:,t]) - - if any(!isfinite, [x for v in initial_state₂ for x in v]) - ok = false - break - end - - Y₁[:,t+1] = initial_state₁ |> sum - Y₂[:,t+1] = initial_state₂ |> sum - else - Y₁[:,t+1] = state_update(Y₁[:,t],baseline_noise) - - if any(!isfinite, Y₁[:,t+1]) - ok = false - break - end - - Y₂[:,t+1] = state_update(Y₂[:,t],baseline_noise + shock_history[:,t]) + changed_vals = parameters[match_idx] + changed_pars = 𝓂.constants.post_complete_parameters.parameters[match_idx] - if any(!isfinite, Y₂[:,t+1]) - ok = false - break - end + if verbose + println("Parameter changes: ") + for (i,m) in enumerate(match_idx) + println("\t",changed_pars[i],"\tfrom ",𝓂.parameter_values[m],"\tto ",changed_vals[i]) end end - if !ok continue end - - Y[:,:,i] += Y₂ - Y₁ - - accepted_draws += 1 - end - - if accepted_draws == 0 - @warn "No draws accepted. Results are empty." - elseif accepted_draws < draws - # average over accepted draws, if desired - @info "$accepted_draws of $draws draws accepted for shock: $(shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input, String_input} ? T.exo[ii] : :Shock_matrix)" - Y[:, :, i] ./= accepted_draws - else - Y[:, :, i] ./= accepted_draws + 𝓂.parameter_values[match_idx] = parameters[match_idx] end end - - axis1 = T.var[var_idx] - - if any(x -> contains(string(x), "◖"), axis1) - axis1_decomposed = decompose_name.(axis1) - axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] - end - axis2 = shocks isa Union{Symbol_input,String_input} ? - shock_idx isa Int ? - [T.exo[shock_idx]] : - T.exo[shock_idx] : - [:Shock_matrix] - - if any(x -> contains(string(x), "◖"), axis2) - axis2_decomposed = decompose_name.(axis2) - axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] + if 𝓂.caches.valid_for.non_stochastic_steady_state != 𝓂.parameter_values && verbose + println("New parameters changed the steady state.") end - return KeyedArray(Y[var_idx,2:end,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) + return nothing end -function girf(state_update::Function, - obc_state_update::Function, - initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, - level::Vector{Float64}, - constants::constants; - periods::Int = 40, - shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, - variables::Union{Symbol_input,String_input} = :all, - shock_size::Real = 1, - negative_shock::Bool = false, - warmup_periods::Int = 100, - draws::Int = 50)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} - T = constants.post_model_macro - - pruning = initial_state isa Vector{Vector{Float64}} - - shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks - - shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks - - if shocks isa Matrix{Float64} - @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model." - - # periods += size(shocks)[2] - - shock_history = zeros(T.nExo, periods) - - shock_history[:,1:size(shocks)[2]] = shocks - - shock_idx = 1 - elseif shocks isa KeyedArray{Float64} - shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) - - # periods += size(shocks)[2] - - @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." +# function create_timings_for_estimation!(𝓂::ℳ, observables::Vector{Symbol}) +# dyn_equations = 𝓂.dyn_equations - shock_history = zeros(T.nExo, periods + 1) +# vars_to_exclude = setdiff(𝓂.constants.post_model_macro.present_only, observables) - shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks +# # Mapping variables to their equation index +# variable_to_equation = Dict{Symbol, Vector{Int}}() +# for var in vars_to_exclude +# for (eq_idx, vars_set) in enumerate(𝓂.dyn_var_present_list) +# # for var in vars_set +# if var in vars_set +# if haskey(variable_to_equation, var) +# push!(variable_to_equation[var],eq_idx) +# else +# variable_to_equation[var] = [eq_idx] +# end +# end +# end +# end - shock_idx = 1 - elseif shocks == :simulate - shock_history = randn(T.nExo,periods) * shock_size - - shock_history[contains.(string.(T.exo),"ᵒᵇᶜ"),:] .= 0 +# # cols_to_exclude = indexin(𝓂.constants.post_model_macro.var, setdiff(𝓂.constants.post_model_macro.present_only, observables)) +# cols_to_exclude = indexin(setdiff(𝓂.constants.post_model_macro.present_only, observables), 𝓂.constants.post_model_macro.var) - shock_idx = 1 - else - shock_idx = parse_shocks_input_to_index(shocks,constants) - end +# present_idx = 𝓂.constants.post_model_macro.nFuture_not_past_and_mixed .+ (setdiff(range(1, 𝓂.constants.post_model_macro.nVars), cols_to_exclude)) - var_idx = parse_variables_input_to_index(variables, constants) |> sort +# dyn_var_future_list = deepcopy(𝓂.dyn_var_future_list) +# dyn_var_present_list = deepcopy(𝓂.dyn_var_present_list) +# dyn_var_past_list = deepcopy(𝓂.dyn_var_past_list) +# dyn_exo_list = deepcopy(𝓂.dyn_exo_list) +# dyn_ss_list = deepcopy(𝓂.dyn_ss_list) - Y = zeros(T.nVars, periods + 1, length(shock_idx)) +# rows_to_exclude = Int[] - for (i,ii) in enumerate(shock_idx) - initial_state_copy = deepcopy(initial_state) +# for vidx in values(variable_to_equation) +# for v in vidx +# if v ∉ rows_to_exclude +# push!(rows_to_exclude, v) - accepted_draws = 0 +# for vv in vidx +# dyn_var_future_list[vv] = union(dyn_var_future_list[vv], dyn_var_future_list[v]) +# dyn_var_present_list[vv] = union(dyn_var_present_list[vv], dyn_var_present_list[v]) +# dyn_var_past_list[vv] = union(dyn_var_past_list[vv], dyn_var_past_list[v]) +# dyn_exo_list[vv] = union(dyn_exo_list[vv], dyn_exo_list[v]) +# dyn_ss_list[vv] = union(dyn_ss_list[vv], dyn_ss_list[v]) +# end - for draw in 1:draws - ok = true +# break +# end +# end +# end - initial_state_copy² = deepcopy(initial_state_copy) +# rows_to_include = setdiff(1:𝓂.constants.post_model_macro.nVars, rows_to_exclude) - warmup_shocks = randn(T.nExo) - warmup_shocks[contains.(string.(T.exo), "ᵒᵇᶜ")] .= 0 +# all_symbols = setdiff(reduce(union,collect.(get_symbols.(dyn_equations)))[rows_to_include], vars_to_exclude) +# parameters_in_equations = sort(setdiff(all_symbols, match_pattern(all_symbols,r"₎$"))) + +# dyn_var_future = sort(setdiff(collect(reduce(union,dyn_var_future_list[rows_to_include])), vars_to_exclude)) +# dyn_var_present = sort(setdiff(collect(reduce(union,dyn_var_present_list[rows_to_include])), vars_to_exclude)) +# dyn_var_past = sort(setdiff(collect(reduce(union,dyn_var_past_list[rows_to_include])), vars_to_exclude)) +# dyn_var_ss = sort(setdiff(collect(reduce(union,dyn_ss_list[rows_to_include])), vars_to_exclude)) - # --- warmup --- - for i_w in 1:warmup_periods - initial_state_copy², _, solved = obc_state_update(initial_state_copy², warmup_shocks, state_update) - if !solved - # @warn "No solution in warmup period: $i_w" - ok = false - break - end - end - - if !ok continue end +# all_dyn_vars = union(dyn_var_future, dyn_var_present, dyn_var_past) - Y₁ = zeros(T.nVars, periods + 1) - Y₂ = zeros(T.nVars, periods + 1) +# @assert length(setdiff(dyn_var_ss, all_dyn_vars)) == 0 "The following variables are (and cannot be) defined only in steady state (`[ss]`): $(setdiff(dyn_var_ss, all_dyn_vars))" - baseline_noise = randn(T.nExo) - baseline_noise[contains.(string.(T.exo), "ᵒᵇᶜ")] .= 0 +# all_vars = union(all_dyn_vars, dyn_var_ss) - if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input, String_input} - shock_history = zeros(T.nExo, periods) - shock_history[ii, 1] = negative_shock ? -shock_size : shock_size - end +# present_only = sort(setdiff(dyn_var_present,union(dyn_var_past,dyn_var_future))) +# future_not_past = sort(setdiff(dyn_var_future, dyn_var_past)) +# past_not_future = sort(setdiff(dyn_var_past, dyn_var_future)) +# mixed = sort(setdiff(dyn_var_present, union(present_only, future_not_past, past_not_future))) +# future_not_past_and_mixed = sort(union(future_not_past,mixed)) +# past_not_future_and_mixed = sort(union(past_not_future,mixed)) +# present_but_not_only = sort(setdiff(dyn_var_present,present_only)) +# mixed_in_past = sort(intersect(dyn_var_past, mixed)) +# not_mixed_in_past = sort(setdiff(dyn_var_past,mixed_in_past)) +# mixed_in_future = sort(intersect(dyn_var_future, mixed)) +# exo = sort(collect(reduce(union,dyn_exo_list))) +# var = sort(dyn_var_present) +# aux_tmp = sort(filter(x->occursin(r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾",string(x)), dyn_var_present)) +# aux = aux_tmp[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∉ exo, aux_tmp)] +# exo_future = dyn_var_future[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_future)] +# exo_present = dyn_var_present[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_present)] +# exo_past = dyn_var_past[map(x->Symbol(replace(string(x),r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) ∈ exo, dyn_var_past)] - # --- period 1 --- - if pruning - initial_state_copy², _, solved = obc_state_update(initial_state_copy², baseline_noise, state_update) - if !solved continue end +# nPresent_only = length(present_only) +# nMixed = length(mixed) +# nFuture_not_past_and_mixed = length(future_not_past_and_mixed) +# nPast_not_future_and_mixed = length(past_not_future_and_mixed) +# nPresent_but_not_only = length(present_but_not_only) +# nVars = length(all_vars) +# nExo = length(collect(exo)) - initial_state₁ = deepcopy(initial_state_copy²) - initial_state₂ = deepcopy(initial_state_copy²) +# present_only_idx = indexin(present_only,var) +# present_but_not_only_idx = indexin(present_but_not_only,var) +# future_not_past_and_mixed_idx = indexin(future_not_past_and_mixed,var) +# past_not_future_and_mixed_idx = indexin(past_not_future_and_mixed,var) +# mixed_in_future_idx = indexin(mixed_in_future,dyn_var_future) +# mixed_in_past_idx = indexin(mixed_in_past,dyn_var_past) +# not_mixed_in_past_idx = indexin(not_mixed_in_past,dyn_var_past) +# past_not_future_idx = indexin(past_not_future,var) - Y₁[:, 1] = initial_state_copy² |> sum - Y₂[:, 1] = initial_state_copy² |> sum - else - Y₁[:, 1], _, solved = obc_state_update(initial_state_copy², baseline_noise, state_update) - if !solved continue end +# reorder = indexin(var, [present_only; past_not_future; future_not_past_and_mixed]) +# dynamic_order = indexin(present_but_not_only, [past_not_future; future_not_past_and_mixed]) - Y₂[:, 1], _, solved = obc_state_update(initial_state_copy², baseline_noise, state_update) - if !solved continue end - end +# @assert length(intersect(union(var,exo),parameters_in_equations)) == 0 "Parameters and variables cannot have the same name. This is the case for: " * repr(sort([intersect(union(var,exo),parameters_in_equations)...])) - # --- remaining periods --- - for t in 1:periods - baseline_noise = randn(T.nExo) - baseline_noise[contains.(string.(T.exo), "ᵒᵇᶜ")] .= 0 - - if pruning - initial_state₁, _, solved = obc_state_update(initial_state₁, baseline_noise, state_update) - if !solved - # @warn "No solution in period: $t" - ok = false - break - end +# T = timings(present_only, +# future_not_past, +# past_not_future, +# mixed, +# future_not_past_and_mixed, +# past_not_future_and_mixed, +# present_but_not_only, +# mixed_in_past, +# not_mixed_in_past, +# mixed_in_future, +# exo, +# var, +# aux, +# exo_present, - initial_state₂, _, solved = obc_state_update(initial_state₂, baseline_noise + shock_history[:, t], state_update) - if !solved - # @warn "No solution in period: $t" - ok = false - break - end +# nPresent_only, +# nMixed, +# nFuture_not_past_and_mixed, +# nPast_not_future_and_mixed, +# nPresent_but_not_only, +# nVars, +# nExo, - Y₁[:, t + 1] = initial_state₁ |> sum - Y₂[:, t + 1] = initial_state₂ |> sum - else - Y₁[:, t + 1], _, solved = obc_state_update(Y₁[:, t], baseline_noise, state_update) - if !solved - # @warn "No solution in period: $t" - ok = false - break - end +# present_only_idx, +# present_but_not_only_idx, +# future_not_past_and_mixed_idx, +# not_mixed_in_past_idx, +# past_not_future_and_mixed_idx, +# mixed_in_past_idx, +# mixed_in_future_idx, +# past_not_future_idx, - Y₂[:, t + 1], _, solved = obc_state_update(Y₂[:, t], baseline_noise + shock_history[:, t], state_update) - if !solved - # @warn "No solution in period: $t" - ok = false - break - end - end - end +# reorder, +# dynamic_order) - if !ok continue end +# push!(𝓂.estimation_helper, observables => T) - # Note: replace `i` if your outer scope uses another index - Y[:, :, i] .+= (Y₂ .- Y₁) - accepted_draws += 1 - end +# return nothing +# end - if accepted_draws == 0 - @warn "No draws accepted. Results are empty." - elseif accepted_draws < draws - # average over accepted draws, if desired - @info "$accepted_draws of $draws draws accepted for shock: $(shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input, String_input} ? T.exo[ii] : :Shock_matrix)" - Y[:, :, i] ./= accepted_draws - else - Y[:, :, i] ./= accepted_draws - end - end - - axis1 = T.var[var_idx] - - if any(x -> contains(string(x), "◖"), axis1) - axis1_decomposed = decompose_name.(axis1) - axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] - end - axis2 = shocks isa Union{Symbol_input,String_input} ? - shock_idx isa Int ? - [T.exo[shock_idx]] : - T.exo[shock_idx] : - [:Shock_matrix] - if any(x -> contains(string(x), "◖"), axis2) - axis2_decomposed = decompose_name.(axis2) - axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] - end - return KeyedArray(Y[var_idx,2:end,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) -end -function parse_variables_input_to_index(variables::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}}}}}, 𝓂::ℳ)::Union{UnitRange{Int}, Vector{Int}} - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) +@unstable function parse_variables_input_to_index(variables::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}}}}}, 𝓂::ℳ)::Union{UnitRange{Int}, Vector{Int}} + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters if variables == :all_excluding_auxiliary_and_obc return ms.vars_idx_excluding_aux_obc elseif variables == :all_excluding_obc @@ -9311,7 +2029,7 @@ function parse_variables_input_to_index(variables::Union{Symbol_input, String_in return parse_variables_input_to_index(variables, 𝓂.constants) end -function parse_variables_input_to_index(variables::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}}}}}, constants::constants)::Union{UnitRange{Int}, Vector{Int}} +@unstable function parse_variables_input_to_index(variables::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}}}}}, constants::constants)::Union{UnitRange{Int}, Vector{Int}} T = constants.post_model_macro @@ -9430,6 +2148,46 @@ function parse_covariance_groups(variables::Union{Symbol_input,String_input, Vec end +function symmetrise_covariance_upper(covariance::AbstractMatrix{T}) where T <: Real + covariance_upper = ℒ.triu(covariance) + return covariance_upper + covariance_upper' - ℒ.Diagonal(ℒ.diag(covariance_upper)) +end + + +function covariance_to_correlation(covariance::AbstractMatrix{T}) where T <: Real + covariance_symmetric = symmetrise_covariance_upper(covariance) + diag_covariance = convert(Vector{T}, ℒ.diag(covariance_symmetric)) + max_diag = maximum(d -> d > 0 ? d : zero(T), diag_covariance; init = zero(T)) + degenerate_tol = max(eps(T), eps(T) * max_diag) + std_corr = Vector{T}(undef, length(diag_covariance)) + + @inbounds for i in eachindex(diag_covariance) + diag_entry = diag_covariance[i] + std_corr[i] = diag_entry > degenerate_tol ? sqrt(diag_entry) : convert(T, NaN) + end + + correlation = covariance_symmetric ./ (std_corr * std_corr') + + # Clamp machine-precision noise to zero for clean display output. + # Skipped for AD element types (e.g. ForwardDiff.Dual) where + # replacing with zero(T) would destroy derivative partials. + if T <: AbstractFloat + noise_tol = eps(T)^(T(2)/T(3)) + n = size(correlation, 1) + @inbounds for j in 1:n, i in 1:n + if i != j + c = correlation[i, j] + if !isnan(c) && abs(c) < noise_tol + correlation[i, j] = zero(T) + end + end + end + end + + return correlation, covariance_symmetric, diag_covariance, std_corr +end + + function parse_shocks_input_to_index(shocks::Expr, constants::constants) @@ -9459,7 +2217,7 @@ function parse_shocks_input_to_index(shocks::BitMatrix, constants::constants) return getindex(1:T.nExo, vec(sum(shocks, dims = 2) .> 0)) end -function parse_shocks_input_to_index(shocks::Union{Symbol_input, String_input}, constants::constants) +@unstable function parse_shocks_input_to_index(shocks::Union{Symbol_input, String_input}, constants::constants) T = constants.post_model_macro @@ -9572,51 +2330,132 @@ end # return [𝐒₁ * aug_state₁̃, 𝐒₁ * aug_state₂̃ + 𝐒₂ * kron_aug_state₁ / 2, 𝐒₁ * aug_state₃̃ + 𝐒₂ * ℒ.kron(aug_state₁̂, aug_state₂) + 𝐒₃ * ℒ.kron(kron_aug_state₁,aug_state₁) / 6] # 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} +@unstable 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 + Ŝ₁ = 𝓂.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 Ŝ₁ * aug_state + end + elseif algorithm ∈ [:second_order, :third_order] + 𝐒₂ = 𝓂.caches.second_order_solution * 𝓂.constants.second_order.𝐔₂ + Ŝ₁̂ = [Ŝ₁[:,1:nPast] zeros(nVars) Ŝ₁[:,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 Ŝ₁̂ * 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 + 𝐒₂ = 𝓂.caches.second_order_solution * 𝓂.constants.second_order.𝐔₂ + Ŝ₁̂ = [Ŝ₁[:,1:nPast] zeros(nVars) Ŝ₁[:,nPast+1:end]] + state_update = (state, shock) -> pruned_second_order_state_update(state, shock, past_idx, nVars, Ŝ₁̂, 𝐒₂) + elseif algorithm == :pruned_third_order + 𝐒₂ = 𝓂.caches.second_order_solution * 𝓂.constants.second_order.𝐔₂ + 𝐒₃ = 𝓂.caches.third_order_solution * 𝓂.constants.third_order.𝐔₃ + Ŝ₁̂ = [Ŝ₁[:,1:nPast] zeros(nVars) Ŝ₁[:,nPast+1:end]] + state_update = (state, shock) -> pruned_third_order_state_update(state, shock, past_idx, nVars, Ŝ₁̂, 𝐒₂, 𝐒₃) 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,12 +2472,13 @@ 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 𝓂.functions.NSSS_custom(output, parameter_values) catch + fill!(output, S(NaN)) end return output elseif applicable(𝓂.functions.NSSS_custom, parameter_values) @@ -9668,36 +2508,38 @@ end # @stable default_mode = "disable" begin -function find_variables_to_exclude(𝓂::ℳ, observables::Vector{Symbol}) - # reduce system - vars_to_exclude = setdiff(𝓂.constants.post_model_macro.present_only, observables) - - # Mapping variables to their equation index - variable_to_equation = Dict{Symbol, Vector{Int}}() - for var in vars_to_exclude - for (eq_idx, vars_set) in enumerate(𝓂.constants.post_model_macro.dyn_var_present_list) - # for var in vars_set - if var in vars_set - if haskey(variable_to_equation, var) - push!(variable_to_equation[var],eq_idx) - else - variable_to_equation[var] = [eq_idx] - end - end - end - end - - return variable_to_equation -end +# Dead code: find_variables_to_exclude — never called anywhere +# function find_variables_to_exclude(𝓂::ℳ, observables::Vector{Symbol}) +# # reduce system +# vars_to_exclude = setdiff(𝓂.constants.post_model_macro.present_only, observables) +# +# # Mapping variables to their equation index +# variable_to_equation = Dict{Symbol, Vector{Int}}() +# for var in vars_to_exclude +# for (eq_idx, vars_set) in enumerate(𝓂.constants.post_model_macro.dyn_var_present_list) +# # for var in vars_set +# if var in vars_set +# if haskey(variable_to_equation, var) +# push!(variable_to_equation[var],eq_idx) +# else +# variable_to_equation[var] = [eq_idx] +# end +# end +# end +# end +# +# return variable_to_equation +# end -function create_broadcaster(indices::Vector{Int}, n::Int) - broadcaster = spzeros(n, length(indices)) - for (i, vid) in enumerate(indices) - broadcaster[vid,i] = 1.0 - end - return broadcaster -end +# Dead code: create_broadcaster — never called anywhere +# function create_broadcaster(indices::Vector{Int}, n::Int) +# broadcaster = spzeros(n, length(indices)) +# for (i, vid) in enumerate(indices) +# broadcaster[vid,i] = 1.0 +# end +# return broadcaster +# end """ update_perturbation_counter!(counters::SolveCounters, solved::Bool; estimation::Bool = false, order::Int = 1) @@ -9743,6 +2585,7 @@ function update_perturbation_counter!(counters::SolveCounters, solved::Bool; est end end end + return nothing end """ @@ -9763,17 +2606,26 @@ function update_ss_counter!(counters::SolveCounters, solved::Bool; estimation::B counters.ss_solves_failed += 1 end end + return nothing end 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) + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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 +2638,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 +2647,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,164 +2670,43 @@ function get_NSSS_and_parameters(𝓂::ℳ, end # end # timeit_debug - 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 - - 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 min(max(parameter_values[indexin([k], 𝓂.constants.post_complete_parameters.parameters)][1], v[1]), v[2]) != parameter_values[indexin([k], 𝓂.constants.post_complete_parameters.parameters)][1] - return true - end - end + # 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 = Float64.(primal.(parameter_values)) + else + 𝓂.caches.valid_for.non_stochastic_steady_state = Float64[] end end - return false -end - -function get_relevant_steady_state_and_state_update(::Val{:second_order}, - parameter_values::Vector{S}, - 𝓂::ℳ; - 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, - - 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 - - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) - all_SS = expand_steady_state(SS_and_pars, ms) - - state = collect(sss) - all_SS - - return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged -end - - - -function get_relevant_steady_state_and_state_update(::Val{:pruned_second_order}, - parameter_values::Vector{S}, - 𝓂::ℳ; - 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, - - 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 - - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) - all_SS = expand_steady_state(SS_and_pars, ms) - - state = [zeros(𝓂.constants.post_model_macro.nVars), collect(sss) - all_SS] - - return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged -end - - - -function get_relevant_steady_state_and_state_update(::Val{:third_order}, - parameter_values::Vector{S}, - 𝓂::ℳ; - 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, - - 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 - - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) - all_SS = expand_steady_state(SS_and_pars, ms) - - state = collect(sss) - all_SS - - return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged -end - - - -function get_relevant_steady_state_and_state_update(::Val{:pruned_third_order}, - parameter_values::Vector{S}, - 𝓂::ℳ; - 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, - - 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 - - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) - all_SS = expand_steady_state(SS_and_pars, ms) - - state = [zeros(𝓂.constants.post_model_macro.nVars), collect(sss) - all_SS, zeros(𝓂.constants.post_model_macro.nVars)] - - return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged + return SS_and_pars, (solution_error, iters) end -function get_relevant_steady_state_and_state_update(::Val{:first_order}, - parameter_values::Vector{S}, - 𝓂::ℳ; - opts::CalculationOptions = merge_calculation_options(), - estimation::Bool = false)::Tuple{constants, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{Vector{Float64}}, Bool} where S <: Real - # timer::TimerOutput = TimerOutput(), - # 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, - - 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 - # println("NSSS not found") - 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 - - 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, - initial_guess = 𝓂.caches.qme_solution, - opts = opts) - if solved 𝓂.caches.qme_solution = qme_sol end +end # dispatch_doctor - @ignore_derivatives update_perturbation_counter!(𝓂.counters, solved, estimation = estimation, order = 1) +@setup_workload begin + @compile_workload begin + @model RBC_for_precompile precompile = true begin + 1 / c[0] = (0.95 / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ)) + c[0] + k[0] = (1 - δ) * k[-1] + exp(z[0]) * k[-1]^α + z[0] = 0.2 * z[-1] + 0.01 * eps_z[x] + end - if !solved - # println("NSSS not found") - return 𝓂.constants, SS_and_pars, zeros(S, 0, 0), [state], solved + @parameters RBC_for_precompile silent = true precompile = true begin + δ = 0.02 + α = 0.5 + end end - - return 𝓂.constants, SS_and_pars, 𝐒₁, [state], solved end - -end # dispatch_doctor - -# @setup_workload begin # # Putting some things in `setup` can reduce the size of the # # precompile file and potentially make loading faster. # @model FS2000 precompile = true begin @@ -10047,12 +2781,10 @@ end # dispatch_doctor # end # end -# Include ForwardDiff Dual specializations for forward-mode AD -# Must be at the end of the module because they depend on function definitions -include("./custom_autodiff_rules/forwarddiff.jl") +# ForwardDiff Dual specializations moved to ext/ForwardDiffExt.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("./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..652a032b0 --- /dev/null +++ b/src/algorithms/fast_lapack_wrappers.jl @@ -0,0 +1,245 @@ +@stable default_mode = "disable" begin + + +# ─── QR factorization ───────────────────────────────────────────────────────── + +function factorize_qr!(::Val{:FastLapack}, + qr_mat::AbstractMatrix, + qr_factors::AbstractMatrix{R}, + qr_ws::FastLapackInterface.QRWs{R}) where {R <: Union{Float32, Float64}} + copyto!(qr_factors, qr_mat) + ℒ.LAPACK.geqrf!(qr_ws, qr_factors; resize = true) + return qr_factors +end + +function factorize_qr!(::Val{:Julia}, + qr_mat::AbstractMatrix, + qr_factors::AbstractMatrix{R}, + qr_ws) where {R <: AbstractFloat} + copyto!(qr_factors, qr_mat) + return ℒ.qr!(qr_factors) +end + +# Old way (≤v0.1.42): dest = Q' * src — allocates intermediate Q.Q' and result +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 + +# Fallback: dest = Q' * src (uses standard mul! when FastLapackInterface is not active) +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) # dest = Q' * src + return qr_orm_ws, qr_orm_dims +end + +# ─── LU factorization ────────────────────────────────────────────────────────── + +function factorize_lu!(::Val{:FastLapack}, + A::AbstractMatrix{R}, + lu_ws, + lu_dims::NTuple{2, Int}) where {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 +end + +function factorize_lu!(::Val{:Julia}, + A::AbstractMatrix{R}, + lu_ws, + lu_dims::NTuple{2, Int}) where {R <: AbstractFloat} + lu = ℒ.lu!(A, check = false) + return lu_ws, lu_dims, ℒ.issuccess(lu), lu +end + +# Old way (≤v0.1.42): X = A \ B — solves A * X = B, allocates result +function solve_lu_left!(A::AbstractMatrix{R}, + B::AbstractVecOrMat{R}, + lu_ws, + lu; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + # B ← A \ B (overwrites B in-place) + if use_fastlapack_lu && R <: Union{Float32, Float64} + ℒ.LAPACK.getrs!(lu_ws, 'N', A, B) + else + ℒ.ldiv!(lu, B) # B = A \ B + end + return B +end + +# B ← A \ B (Nothing-dispatch variant, always uses LAPACK) +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) # B = A \ B + return B +end + +# B ← A' \ B (overwrites B in-place) +function solve_lu_left_transpose!(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, 'T', A, B) + else + ℒ.ldiv!(lu', B) + end + return B +end + +# B ← A' \ B (Nothing-dispatch variant, always uses LAPACK) +function solve_lu_left_transpose!(A::AbstractMatrix{R}, + B::AbstractVecOrMat{R}, + lu_ws, + lu::Nothing; + use_fastlapack_lu::Bool = true) where {R <: AbstractFloat} + ℒ.LAPACK.getrs!(lu_ws, 'T', A, B) + return B +end + +# Old way (≤v0.1.42): X = B / A — solves X * A = B, allocates result +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} + # B ← B / A (overwrites B in-place) + 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) # B = B / A + end + return B +end + +# B ← B / A (Nothing-dispatch variant, always uses LAPACK) +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 + +# ─── Generalized Schur factorization ────────────────────────────────────────── + +# Returns NamedTuple{(:S,:T,:Z)} for the decomposition (or dummy values on failure). +function factorize_generalized_schur!(::Val{:FastLapack}, + D::AbstractMatrix{R}, + E::AbstractMatrix{R}, + qz_ws, + qz_dims::NTuple{2, Int}, + eigenselect::AbstractVector{Bool}; + unit_root_tol::Float64 = 1e-6) where {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 + # FastLapackInterface.ed selects |λ|² ≥ criterium, putting those eigenvalues + # into the leading (unstable) block of the generalized Schur factorization. + # Pushing the cutoff well inside the unit circle (Dynare uses 1e-6) keeps + # eigenvalues clustered near unity entirely on one side of the boundary, + # which prevents LAPACK reordering failures and the resulting Z₁₁ ill- + # conditioning that otherwise blows up the QME residual (e.g. on FRB/US, + # where eigenvalues sit ~3e-8 from unity). A smaller offset such as + # sqrt(eps) is too tight and forces a costly fallback to the doubling + # algorithm for these models. + S, T, α, β, _, Z = ℒ.LAPACK.gges!(qz_ws, 'N', 'V', D, E; + select = FastLapackInterface.ed, + criterium = (1.0 - unit_root_tol)^2, + resize = true) + has_ur = detect_unit_roots(α, β, unit_root_tol) + return qz_ws, qz_dims, (S = S, T = T, Z = Z), true, has_ur + catch + return qz_ws, qz_dims, (S = D, T = E, Z = D), false, false + end +end + +# Returns NamedTuple{(:S,:T,:Z)} for the decomposition (or dummy values on failure). +function factorize_generalized_schur!(::Val{:Julia}, + D::AbstractMatrix{R}, + E::AbstractMatrix{R}, + qz_ws, + qz_dims::NTuple{2, Int}, + eigenselect::AbstractVector{Bool}; + unit_root_tol::Float64 = 1e-6) where {R <: AbstractFloat} + schdcmp = try + ℒ.schur!(D, E) + catch + return qz_ws, qz_dims, (S = D, T = E, Z = D), false, false + end + + # Match the fast-path criterium: classify any eigenvalue with |λ| > 1 - unit_root_tol + # (including near-unit eigenvalues) into the leading "unstable" block, which keeps + # the stationary Z₁₁ block well-conditioned and avoids ordschur failures on + # exactly-unit blocks. + @. eigenselect = abs(schdcmp.β / schdcmp.α) < 1 / (1 - unit_root_tol) + + try + ℒ.ordschur!(schdcmp, eigenselect) + catch + return qz_ws, qz_dims, (S = D, T = E, Z = D), false, false + end + + has_ur = detect_unit_roots(schdcmp.α, schdcmp.β, unit_root_tol) + return qz_ws, qz_dims, (S = schdcmp.S, T = schdcmp.T, Z = schdcmp.Z), true, has_ur +end + +# Detect unit root eigenvalues from generalized Schur eigenvalue vectors. +# Returns true if any |α[i]/β[i]| is within tol of 1.0. +function detect_unit_roots(α::AbstractVector, β::AbstractVector, tol::Float64)::Bool + for i in eachindex(α, β) + βi = abs(β[i]) + βi == 0 && continue + eig_mag = abs(α[i]) / βi + if abs(eig_mag - 1) ≤ tol + return true + end + end + return false +end + + +end # @stable diff --git a/src/algorithms/lyapunov.jl b/src/algorithms/lyapunov.jl index 36ee6f71b..1caaec85a 100644 --- a/src/algorithms/lyapunov.jl +++ b/src/algorithms/lyapunov.jl @@ -1,23 +1,82 @@ +@stable default_mode = "disable" begin + # Available algorithms: # :doubling - fast and precise # :bartels_stewart - fast for small matrices and precise, dense matrices only # :bicgstab - less precise # :gmres - less precise +# :dqgmres - less precise +# Tested column-ILU and triangular-sweep Krylov preconditioners did not improve + # :iterative - slow and precise # :speedmapping - slow and very precise # solves: A * X * A' + C = X -@stable default_mode = "disable" begin -function solve_lyapunov_equation(A::AbstractMatrix{T}, +# 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 + +@unstable 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, - verbose::Bool = false)::Union{Tuple{Matrix{T}, Bool}, Tuple{ThreadedSparseArrays.ThreadedSparseMatrixCSC{T, Int, SparseMatrixCSC{T, Int}}, Bool}} where T <: Float64 + tol::SolverTolerances = SolverTolerances(atol = 1e-14, + rtol = 1e-14, + initial_guess_acceptance_tol = 1e-12, + acceptance_tol = 1e-12), + verbose::Bool = false, + has_unit_roots::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 @@ -27,6 +86,10 @@ function solve_lyapunov_equation(A::AbstractMatrix{T}, # @timeit_debug timer "Solve lyapunov equation" begin # @timeit_debug timer "Choose matrix formats" begin + if lyapunov_algorithm == :bartels_stewart && !has_bartels_stewart() + error("The :bartels_stewart algorithm requires the MatrixEquations package. Run `using MatrixEquations` to enable it.") + end + if lyapunov_algorithm ≠ :bartels_stewart A = choose_matrix_format(A) else @@ -36,10 +99,51 @@ 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 + # Fast path: when unit roots are known from QME solve, skip directly to Schur deflation + # instead of wasting O(n³) on solvers guaranteed to fail. + if has_unit_roots + A_dense = collect(A) + C_dense = collect(C) + + X_deflated, deflation_solved = solve_lyapunov_schur_deflation(A_dense, C_dense, workspace; + tol = tol, + verbose = verbose) + if deflation_solved + if verbose + println("Lyapunov equation - solved via Schur deflation (unit roots pre-detected)") + end + return X_deflated, true + end + # If deflation failed despite the flag, fall through to standard solvers + end + X, i, reached_tol = solve_lyapunov_equation(A, C, Val(lyapunov_algorithm), workspace; tol = tol) # timer = timer) if verbose @@ -66,7 +170,7 @@ function solve_lyapunov_equation(A::AbstractMatrix{T}, end end - if !(reached_tol < acceptance_tol) && lyapunov_algorithm ≠ :bartels_stewart && length(C) < 5e7 # try sylvester if previous one didn't solve it + if !(reached_tol < acceptance_tol) && lyapunov_algorithm ≠ :bartels_stewart && length(C) < 5e7 && has_bartels_stewart() # try bartels_stewart if previous one didn't solve it A = collect(A) C = collect(C) @@ -77,42 +181,38 @@ function solve_lyapunov_equation(A::AbstractMatrix{T}, println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bartels_stewart") end end - # end # timeit_debug - # end # timeit_debug - - # if (reached_tol > tol) println("Lyapunov failed: $reached_tol") end - - return X, reached_tol < acceptance_tol -end - - -function solve_lyapunov_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMatrix{T}}, - C::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMatrix{T}}, - ::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 - 𝐂 = try - MatrixEquations.lyapd(A, C)::Matrix{T} - catch - return C, 0, 1.0 + # Schur deflation fallback: when all standard solvers fail, check for unit-root + # eigenvalues and solve only the stationary subspace. + if !(reached_tol < acceptance_tol) + A_dense = collect(A) + C_dense = collect(C) + + X_deflated, deflation_solved = solve_lyapunov_schur_deflation(A_dense, C_dense, workspace; + tol = tol, + verbose = verbose) + if deflation_solved + X = X_deflated + reached_tol = zero(T) # signal success + if verbose + println("Lyapunov equation - solved via Schur deflation (unit-root subspace set to NaN)") + end + end end - - # 𝐂¹ = A * 𝐂 * A' + C - - # denom = max(ℒ.norm(𝐂), ℒ.norm(𝐂¹)) - # reached_tol = denom == 0 ? 0.0 : ℒ.norm(𝐂¹ - 𝐂) / denom - - reached_tol = ℒ.norm(A * 𝐂 * A' + C - 𝐂) / ℒ.norm(𝐂) + return X, reached_tol < acceptance_tol +end - # if reached_tol > tol - # println("Lyapunov: lyapunov $reached_tol") - # end - return 𝐂, 0, reached_tol # return info on convergence +# Keep the low-level bartels-stewart signature available in core so fallback +# paths remain well-typed when MatrixEquations is not loaded. +function solve_lyapunov_equation(A::AbstractMatrix{T}, + C::AbstractMatrix{T}, + ::Val{:bartels_stewart}, + workspace::lyapunov_workspace; + tol::SolverTolerances = SolverTolerances(), + has_unit_roots::Bool = false)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + return Matrix(C), 0, T(Inf) end @@ -122,25 +222,45 @@ 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 - # Note: workspace is unused for sparse matrices but accepted for API consistency + max_iter::Int = 50, + tol::SolverTolerances = SolverTolerances())::Tuple{<:AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned sparse storage created locally in this method. + # Note: workspace was unused for sparse matrices but is now used for AD power capture 𝐂 = copy(C) 𝐀 = copy(A) - max_iter = 500 + if workspace.pow_capture + cache_set!(workspace.𝐀_pow, 1, A, workspace.pow_transposed) + workspace.pow_iters = 1 + end iters = max_iter for i in 1:max_iter 𝐂¹ = 𝐀 * 𝐂 * 𝐀' + 𝐂 - 𝐀 = 𝐀^2 + if workspace.pow_iters >= i + 1 + cached = workspace.𝐀_pow[i + 1] + if issparse(cached) && size(𝐀) == size(cached) && eltype(𝐀) == eltype(cached) + 𝐀 = cached + else + 𝐀 = convert(typeof(𝐀), cached) + end + else + 𝐀 = 𝐀^2 - droptol!(𝐀, eps()) + droptol!(𝐀, eps()) + + if workspace.pow_capture + target_k = i + 1 + cache_set!(workspace.𝐀_pow, target_k, 𝐀, workspace.pow_transposed) + workspace.pow_iters = target_k + end + end 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,28 +291,48 @@ 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 - # Note: workspace is unused for sparse matrices but accepted for API consistency + max_iter::Int = 50, + tol::SolverTolerances = SolverTolerances())::Tuple{<:AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned sparse storage created locally in this method. + # Note: workspace was unused for sparse matrices but is now used for AD power capture 𝐂 = copy(C) 𝐀 = copy(A) 𝐀² = similar(𝐀) - max_iter = 500 + if workspace.pow_capture + cache_set!(workspace.𝐀_pow, 1, A, workspace.pow_transposed) + workspace.pow_iters = 1 + end iters = max_iter for i in 1:max_iter 𝐂¹ = 𝐀 * 𝐂 * 𝐀' + 𝐂 - ℒ.mul!(𝐀², 𝐀, 𝐀) - copyto!(𝐀, 𝐀²) + if workspace.pow_iters >= i + 1 + cached = workspace.𝐀_pow[i + 1] + if typeof(cached) === typeof(𝐀) && size(cached) == size(𝐀) + copyto!(𝐀, cached) + else + copyto!(𝐀, convert(typeof(𝐀), cached)) + end + else + ℒ.mul!(𝐀², 𝐀, 𝐀) + copyto!(𝐀, 𝐀²) - # droptol!(𝐀, eps()) + # droptol!(𝐀, eps()) + + if workspace.pow_capture + target_k = i + 1 + cache_set!(workspace.𝐀_pow, target_k, 𝐀, workspace.pow_transposed) + workspace.pow_iters = target_k + end + end 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,33 +363,56 @@ 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 - # Note: workspace is unused for sparse matrices but accepted for API consistency + max_iter::Int = 50, + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns owned dense storage created locally in this method. + # Note: workspace was unused for sparse matrices but is now used for AD power capture 𝐂 = copy(C) 𝐀 = copy(A) 𝐂A = collect(𝐀) 𝐂¹ = copy(C) - max_iter = 500 + if workspace.pow_capture + cache_set!(workspace.𝐀_pow, 1, A, workspace.pow_transposed) + workspace.pow_iters = 1 + end 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) - # 𝐀 *= 𝐀 - 𝐀 = 𝐀^2 # faster than A *= A - # copyto!(𝐂A,𝐀) - # 𝐀 = sparse(𝐀 * 𝐂A) - # 𝐀 = sparse(𝐂A * 𝐀) # faster than sparse-dense matmul but slower than sparse sparse matmul - - droptol!(𝐀, eps()) + if workspace.pow_iters >= i + 1 + cached = workspace.𝐀_pow[i + 1] + if issparse(cached) && size(𝐀) == size(cached) && eltype(𝐀) == eltype(cached) + 𝐀 = cached + else + 𝐀 = convert(typeof(𝐀), cached) + end + else + # 𝐀 *= 𝐀 + 𝐀 = 𝐀^2 # faster than A *= A + # copyto!(𝐂A,𝐀) + # 𝐀 = sparse(𝐀 * 𝐂A) + # 𝐀 = sparse(𝐂A * 𝐀) # faster than sparse-dense matmul but slower than sparse sparse matmul + + droptol!(𝐀, eps()) + + if workspace.pow_capture + target_k = i + 1 + cache_set!(workspace.𝐀_pow, target_k, 𝐀, workspace.pow_transposed) + workspace.pow_iters = target_k + end + end 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 +423,12 @@ function solve_lyapunov_equation( A::AbstractSparseMatrix{T}, # 𝐂 = 𝐂¹ 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(𝐂) - - # 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 +441,9 @@ 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 + max_iter::Int = 50, + 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) @@ -302,20 +458,46 @@ function solve_lyapunov_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMat copyto!(𝐂¹, C) copyto!(𝐀, A) - max_iter = 500 + if workspace.pow_capture + if workspace.pow_iters < 1 + cache_set!(workspace.𝐀_pow, 1, A, workspace.pow_transposed) + workspace.pow_iters = 1 + end + end 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 workspace.pow_iters >= i + 1 + cached = workspace.𝐀_pow[i + 1] + if size(𝐀) == size(cached) && eltype(𝐀) == eltype(cached) && !issparse(cached) + copyto!(𝐀, cached) + else + 𝐀 = Matrix{eltype(𝐀)}(cached) + workspace.𝐀 = 𝐀 + 𝐀² = workspace.𝐀² + end + else + ℒ.mul!(𝐀², 𝐀, 𝐀) + copyto!(𝐀, 𝐀²) + # Capture power for AD reuse: 𝐀_pow[i+1] = A^(2^i) + if workspace.pow_capture + target_k = i + 1 + cache_set!(workspace.𝐀_pow, target_k, 𝐀, workspace.pow_transposed) + workspace.pow_iters = target_k + end + end + 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,24 +507,23 @@ function solve_lyapunov_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMat copyto!(𝐂, 𝐂¹) end - # ℒ.mul!(𝐂A, 𝐂, A') - # ℒ.mul!(𝐂¹, A, 𝐂A) - # ℒ.axpy!(1, C, 𝐂¹) + ℒ.mul!(𝐂A, 𝐂, A') + ℒ.mul!(𝐂¹, A, 𝐂A) + ℒ.axpy!(1, C, 𝐂¹) + ℒ.axpy!(-1, 𝐂, 𝐂¹) + + reached_tol = ℒ.norm(𝐂¹) / ℒ.norm(𝐂) + + return 𝐂, iters, reached_tol # return info on convergence +end + + + - # denom = max(ℒ.norm(𝐂), ℒ.norm(𝐂¹)) - # ℒ.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 - return copy(𝐂), iters, reached_tol # return info on convergence -end @@ -352,48 +533,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 +607,146 @@ 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(𝐗)) + 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_op = LinearOperators.LinearOperator(Float64, n_vech, n_vech, true, true, lyapunov_vech_gmres!) + + vech!(b_vech, C) + + Krylov.gmres!(workspace.gmres_vech, lyapunov_op, b_vech, rtol = tol.rtol, atol = tol.atol) + + fill_symmetric_from_vech!(𝐗, workspace.gmres_vech.x) + + # Allocation-free residual + ensure_lyapunov_doubling_buffers!(workspace) ℒ.mul!(tmp̄, 𝐗, A') - ℒ.mul!(𝐗, A, tmp̄, -1, 1) - copyto!(sol, 𝐗) - # sol = @view reshape(𝐗, size(sol)) + ℒ.mul!(workspace.𝐂¹, A, tmp̄) + ℒ.axpy!(1, C, workspace.𝐂¹) + ℒ.axpy!(-1, 𝐗, workspace.𝐂¹) + reached_tol = ℒ.norm(workspace.𝐂¹) / ℒ.norm(𝐗) + + 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 + + lyapunov_op = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov_gmres!) + + copyto!(b, vec(C)) + Krylov.gmres!(workspace.gmres, lyapunov_op, b, rtol = tol.rtol, atol = tol.atol) + copyto!(𝐗, workspace.gmres.x) + + # 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 𝐗, workspace.gmres.stats.niter, reached_tol end +end - lyapunov = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov!) - # Use vectorized C in workspace - copyto!(b, vec(C)) +function solve_lyapunov_equation(A::AbstractMatrix{T}, + C::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMatrix{T}}, + ::Val{:dqgmres}, + workspace::lyapunov_workspace; + # timer::TimerOutput = TimerOutput(), + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + # Ownership: returns workspace-backed dense Krylov buffer workspace.𝐗. - # Use pre-allocated solver - Krylov.gmres!(workspace.gmres_workspace, lyapunov, b, rtol = tol, atol = tol) + if is_approx_symmetric(C) + # vech-space Krylov: solve for n(n+1)/2 unique elements only + ensure_lyapunov_krylov_vech_solver!(workspace, :dqgmres) + tmp̄ = workspace.tmp̄ + 𝐗 = workspace.𝐗 + n = size(A, 1) + n_vech = n * (n + 1) ÷ 2 + b_vech = workspace.b_vech + + function lyapunov_vech_dqgmres!(sol, 𝐱) + fill_symmetric_from_vech!(𝐗, 𝐱) + ℒ.mul!(tmp̄, 𝐗, A') + ℒ.mul!(𝐗, A, tmp̄, -1, 1) + vech!(sol, 𝐗) + end - copyto!(𝐗, workspace.gmres_workspace.x) + lyapunov_op = LinearOperators.LinearOperator(Float64, n_vech, n_vech, true, true, lyapunov_vech_dqgmres!) - # ℒ.mul!(tmp̄, A, 𝐗 * A') - # ℒ.axpy!(1, C, tmp̄) + vech!(b_vech, C) - # denom = max(ℒ.norm(𝐗), ℒ.norm(tmp̄)) + Krylov.dqgmres!(workspace.dqgmres_vech, lyapunov_op, b_vech, rtol = tol.rtol, atol = tol.atol) - # ℒ.axpy!(-1, 𝐗, tmp̄) + fill_symmetric_from_vech!(𝐗, workspace.dqgmres_vech.x) - # reached_tol = denom == 0 ? 0.0 : ℒ.norm(tmp̄) / denom + # 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(𝐗) - reached_tol = ℒ.norm(A * 𝐗 * A' + C - 𝐗) / ℒ.norm(𝐗) + return 𝐗, workspace.dqgmres_vech.stats.niter, reached_tol + else + # Standard full-space Krylov + ensure_lyapunov_krylov_solver!(workspace, :dqgmres) + tmp̄ = workspace.tmp̄ + 𝐗 = workspace.𝐗 + b = workspace.b + + function lyapunov_dqgmres!(sol,𝐱) + copyto!(𝐗, 𝐱) + ℒ.mul!(tmp̄, 𝐗, A') + ℒ.mul!(𝐗, A, tmp̄, -1, 1) + copyto!(sol, 𝐗) + end - # if reached_tol > tol - # println("Lyapunov: gmres $reached_tol") - # end + lyapunov_op = LinearOperators.LinearOperator(Float64, length(C), length(C), true, true, lyapunov_dqgmres!) + + copyto!(b, vec(C)) + Krylov.dqgmres!(workspace.dqgmres, lyapunov_op, b, rtol = tol.rtol, atol = tol.atol) + copyto!(𝐗, workspace.dqgmres.x) - return copy(𝐗), workspace.gmres_workspace.stats.niter, reached_tol + # 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 𝐗, workspace.dqgmres.stats.niter, reached_tol + end end @@ -524,4 +825,121 @@ end # return 𝐂, soll.maps, reached_tol # end -end # dispatch_doctor +# Schur deflation for Lyapunov equations with unit-root eigenvalues. +# +# When A has eigenvalues on or outside the unit circle, the standard Lyapunov +# equation A*X*A' + C = X has no finite solution. This function decomposes A via +# real Schur factorization, reorders so that unstable eigenvalues (|λ| ≥ 1 - unit_root_tol) +# come first, then solves the Lyapunov equation only for the stationary (lower-right) +# block. Original-basis entries whose variance is contaminated by unit-root directions +# are set to NaN. +# +# Returns (X, solved::Bool) where X is n×n with NaN for unit-root-affected entries. + +# Type-stable wrapper for ordered Schur decomposition via LAPACK gees!. +# gees! returns a union type (eigenvalue vector is Float64 or ComplexF64), +# so this barrier function isolates the type instability and returns only +# the concrete types needed by callers: (T_matrix, Z_vectors, n_selected). +function ordered_schur!(A_work::Matrix{T}, unit_root_tol::Float64, + schur_ws::FastLapackInterface.SchurWs{T}) where T <: AbstractFloat + ℒ.LAPACK.gees!(schur_ws, 'V', A_work; + select = FastLapackInterface.ed, + criterium = (1 - unit_root_tol)^2, + resize = true) + vs = schur_ws.vs::Matrix{T} + n_sel = schur_ws.sdim[]::Int + return (A_work, vs, n_sel) +end + +function solve_lyapunov_schur_deflation(A::DenseMatrix{T}, + C::DenseMatrix{T}, + workspace::lyapunov_workspace; + tol::SolverTolerances = SolverTolerances(), + verbose::Bool = false, + unit_root_tol::Float64 = 1e-8)::Tuple{Matrix{T}, Bool} where T <: AbstractFloat + n = size(A, 1) + + # Real Schur decomposition with eigenvalue reordering in one step via LAPACK gees!. + # FastLapackInterface.ed selects eigenvalues on the exterior of the disk (|λ|² ≥ criterium), + # placing unstable eigenvalues in the top-left block. + # After: Tmat = [T_uu T_us; 0 T_ss] where T_ss is the stable block. + A_work = copy(A) + Tmat, U, n_unstable = ordered_schur!(A_work, unit_root_tol, workspace.schur_ws) + + if n_unstable == 0 + # No unit roots found — deflation not applicable, signal failure so caller + # does not silently accept a potentially incorrect result + return Matrix{T}(undef, 0, 0), false + end + + if n_unstable == n + # All eigenvalues are unit roots — no stationary subspace + return fill(T(NaN), n, n), true + end + + n_stable = n - n_unstable + stable_range = (n_unstable + 1):n + + T_ss = Tmat[stable_range, stable_range] + + # Transform noise covariance to Schur basis + C_schur = U' * C * U + C_ss = C_schur[stable_range, stable_range] + + # Symmetrize (numerical noise from rotation can break symmetry) + C_ss = (C_ss + C_ss') / 2 + + # Solve the reduced Lyapunov equation: X_ss = T_ss * X_ss * T_ss' + C_ss + # This converges because all eigenvalues of T_ss are strictly inside the unit circle. + # Try multiple algorithms directly (not via dispatch, to avoid recursion into Schur deflation). + ws_stable = Lyapunov_workspace(n_stable) + X_ss_result, sub_iters, sub_tol = solve_lyapunov_equation(T_ss, C_ss, Val(:doubling), ws_stable; tol = tol) + + if sub_tol > tol.acceptance_tol + X_ss_result, sub_iters, sub_tol = solve_lyapunov_equation(T_ss, C_ss, Val(:bicgstab), ws_stable; tol = tol) + end + + if sub_tol > tol.acceptance_tol && has_bartels_stewart() && length(C_ss) < 5e7 + X_ss_result, sub_iters, sub_tol = solve_lyapunov_equation(T_ss, C_ss, Val(:bartels_stewart), ws_stable; tol = tol) + end + + if sub_tol > tol.acceptance_tol + if verbose + println("Schur deflation: stable sub-block Lyapunov failed (tol=$sub_tol)") + end + return Matrix{T}(undef, 0, 0), false + end + + X_ss = collect(X_ss_result) + + # Map back to original coordinates. + # Only the stationary component contributes finite variance: + # Σ_stationary = U_s * X_ss * U_s' + U_s = U[:, stable_range] + Σ = U_s * X_ss * U_s' + + # Identify which original variables have any loading on unstable Schur vectors. + # These variables have infinite unconditional variance → set to NaN. + U_u = @view U[:, 1:n_unstable] + unstable_loading = vec(sum(abs2, U_u; dims = 2)) # ‖U_u[i,:]‖² + unit_root_vars = unstable_loading .> unit_root_tol + + # Set rows and columns of unit-root-affected variables to NaN + for i in 1:n + if unit_root_vars[i] + Σ[i, :] .= T(NaN) + Σ[:, i] .= T(NaN) + end + end + + if verbose + println("Schur deflation: $n_unstable unstable eigenvalue(s), ", + "$n_stable stable, $(count(unit_root_vars)) variable(s) set to NaN") + end + + return Σ, true +end + + + +end # @stable diff --git a/src/algorithms/nonlinear_solver.jl b/src/algorithms/nonlinear_solver.jl index 0878f87c2..c773b87e8 100644 --- a/src/algorithms/nonlinear_solver.jl +++ b/src/algorithms/nonlinear_solver.jl @@ -1,5 +1,6 @@ @stable default_mode = "disable" begin + function levenberg_marquardt( fnj::function_and_jacobian, # f::Function, @@ -12,9 +13,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 +182,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 @@ -351,6 +352,15 @@ function levenberg_marquardt( return best_current_guess, (grad_iter, func_iter, largest_relative_step, largest_residual)#, f(best_guess)) end +function scale_columns!(A::Matrix{T}, v::Vector{T}) where T <: Real + @turbo for j in 1:size(A, 2) + for i in 1:size(A, 1) + A[i, j] *= v[j] + end + end + return A +end + function scale_columns!(A::AbstractMatrix{T}, v::AbstractVector{T}) where T @inbounds for j in 1:size(A, 2) for i in 1:size(A, 1) @@ -374,6 +384,8 @@ end function update_∇̂!(∇̂::AbstractMatrix{T}, μ¹s::T, μ²::T, p²::T) where T <: Real n = size(∇̂, 1) # hoist size lookup + # Note: @turbo not used here — non-contiguous diagonal stride [i,i] and x^p² (pow) + # are unlikely to benefit from SIMD vectorization @inbounds for i in 1:n x = ∇̂[i,i] # read once x += μ¹s @@ -421,9 +433,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 +496,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 +546,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) @@ -601,7 +615,7 @@ end function minmax!(x::Vector{Float64},lb::Vector{Float64},ub::Vector{Float64}) - @inbounds for i in eachindex(x) + @turbo for i in eachindex(x) x[i] = max(lb[i], min(x[i], ub[i])) end end @@ -683,4 +697,5 @@ end # return x # end -end # dispatch_doctor \ No newline at end of file + +end # @stable diff --git a/src/algorithms/preconditioner.jl b/src/algorithms/preconditioner.jl new file mode 100644 index 000000000..8d1ca6ced --- /dev/null +++ b/src/algorithms/preconditioner.jl @@ -0,0 +1,256 @@ +@stable default_mode = "disable" begin + + +# ─── Incomplete LU (ILU) preconditioner for Krylov Sylvester solvers ───────── +# +# Crout ILU(τ) factorisation and Sylvester block-diagonal preconditioner builder. +# Adapted from KrylovPreconditioners.jl (MPL-2.0, © 2023 Alexis Montoison) +# which incorporated IncompleteLU.jl. Only the subset needed here is retained. + +# ─── Sparse-vector accumulator (Gustavson's O(1)-reset technique) ──────────── +# +# Dense-length vectors with a generation counter avoid O(n) zeroing each step. +# An entry is "live" when occupied[i] == gen. + +mutable struct SparseAccum{Tv,Ti} + occupied::Vector{Ti} + nzind::Vector{Ti} + nzval::Vector{Tv} + nnz::Ti + gen::Ti + + SparseAccum{Tv,Ti}(n::Integer) where {Tv,Ti} = + new(zeros(Ti, n), Vector{Ti}(undef, n), Vector{Tv}(undef, n), zero(Ti), one(Ti)) +end + +@inline function _scatter!(v::SparseAccum, a, idx) + @inbounds if v.occupied[idx] == v.gen + v.nzval[idx] += a + else + v.nnz += 1 + v.occupied[idx] = v.gen + v.nzval[idx] = a + v.nzind[v.nnz] = idx + end + nothing +end + +# Drop entries below threshold, sort survivors, append as column j of A, reset. +function flush_column!(A::SparseMatrixCSC, v::SparseAccum, j::Integer, + drop, scale = one(eltype(A))) + total = 0 + @inbounds for i = 1:v.nnz + row = v.nzind[i] + if abs(v.nzval[row]) >= drop || row == j + total += 1 + v.nzind[total] = row + end + end + sort!(v.nzind, 1, total, Base.Sort.QuickSort, Base.Order.Forward) + @inbounds for i = 1:total + row = v.nzind[i] + push!(A.rowval, row) + push!(A.nzval, scale * v.nzval[row]) + end + @inbounds A.colptr[j + 1] = A.colptr[j] + total + v.gen += 1; v.nnz = 0 + nothing +end + +# ─── ILUFactorization ─────────────────────────────────────────────────────── + +struct ILUFactorization{Tv,Ti} <: ℒ.Factorization{Tv} + L::SparseMatrixCSC{Tv,Ti} + U::SparseMatrixCSC{Tv,Ti} +end + +function ℒ.ldiv!(F::ILUFactorization, y::AbstractVecOrMat) + # Forward substitution (unit lower-triangular L, no stored diagonal) + L = F.L + @inbounds for col = 1 : L.n - 1 + for idx = L.colptr[col] : L.colptr[col + 1] - 1 + y[L.rowval[idx]] -= L.nzval[idx] * y[col] + end + end + # Backward substitution (upper-triangular U with stored diagonal) + U = F.U + @inbounds for col = U.n : -1 : 1 + for idx = U.colptr[col + 1] - 1 : -1 : U.colptr[col] + 1 + y[col] -= U.nzval[idx] * y[U.rowval[idx]] + end + y[col] /= U.nzval[U.colptr[col]] + end + y +end + +function ℒ.ldiv!(y::AbstractVector, F::ILUFactorization, x::AbstractVector) + y .= x + ℒ.ldiv!(F, y) +end + +# ─── Crout ILU(τ) factorisation ───────────────────────────────────────────── +# +# Row access to all matrices (A, L, U) uses linked-list indices over the +# CSC structure. Direct .colptr/.rowval/.nzval field access is used +# throughout for minimal overhead in tight loops. + +function ilu(A::SparseMatrixCSC{ATv,Ti}; τ = 1e-3) where {ATv,Ti} + n = size(A, 1) + Tv = typeof(oneunit(ATv) / (oneunit(ATv) + zero(ATv))) + + L = spzeros(Tv, Ti, n, n) + U = spzeros(Tv, Ti, n, n) + Ur = SparseAccum{Tv,Ti}(n) + Lc = SparseAccum{Tv,Ti}(n) + + # Linked-list row index for A (pre-populated) + A_nxt = A.colptr[1:n] + A_head = zeros(Ti, n); A_rnxt = zeros(Ti, n) + @inbounds for i = Ti(1):Ti(n) + row = A.rowval[A.colptr[i]] + A_head[row], A_rnxt[i] = i, A_head[row] + end + + # Linked-list row index for L and U (start empty, built incrementally) + L_nxt = zeros(Ti, n); L_head = zeros(Ti, n); L_rnxt = zeros(Ti, n) + U_nxt = zeros(Ti, n); U_head = zeros(Ti, n); U_rnxt = zeros(Ti, n) + + @inbounds for k = Ti(1):Ti(n) + + # --- Scatter row k of A into Ur, column k of A into Lc --- + c = A_head[k] + while c != 0 + _scatter!(Ur, A.nzval[A_nxt[c]], c) + nc = A_rnxt[c] + A_nxt[c] += 1 + if A_nxt[c] < A.colptr[c + 1] && A.rowval[A_nxt[c]] <= c + row = A.rowval[A_nxt[c]] + A_head[row], A_rnxt[c] = c, A_head[row] + end + c = nc + end + for idx = A_nxt[k] : A.colptr[k + 1] - 1 + _scatter!(Lc, A.nzval[idx], A.rowval[idx]) + end + + # --- Ur[k:n] -= L[k,i] * U[i, k:n] for i < k --- + c = L_head[k] + while c != 0 + a = -L.nzval[L_nxt[c]] + for idx = U_nxt[c] : U.colptr[c + 1] - 1 + _scatter!(Ur, a * U.nzval[idx], U.rowval[idx]) + end + nc = L_rnxt[c] + L_nxt[c] += 1 + if L_nxt[c] < L.colptr[c + 1] + row = L.rowval[L_nxt[c]] + L_head[row], L_rnxt[c] = c, L_head[row] + end + c = nc + end + + # --- Lc[k+1:n] -= U[i,k] * L[i, k+1:n] for i < k --- + if k < n + c = U_head[k] + while c != 0 + a = -U.nzval[U_nxt[c]] + for idx = L_nxt[c] : L.colptr[c + 1] - 1 + _scatter!(Lc, a * L.nzval[idx], L.rowval[idx]) + end + nc = U_rnxt[c] + U_nxt[c] += 1 + if U_nxt[c] < U.colptr[c + 1] + row = U.rowval[U_nxt[c]] + U_head[row], U_rnxt[c] = c, U_head[row] + end + c = nc + end + end + + # --- Drop small entries, store columns of U and L --- + flush_column!(U, Ur, k, τ) + flush_column!(L, Lc, k, τ, inv(Ur.nzval[k])) + + # Register new entries in row-traversal index + U_nxt[k] = U.colptr[k] + 1 + if U.colptr[k] < U.colptr[k + 1] - 1 + row = U.rowval[U_nxt[k]] + U_head[row], U_rnxt[k] = k, U_head[row] + end + + L_nxt[k] = L.colptr[k] + if L.colptr[k] < L.colptr[k + 1] + row = L.rowval[L_nxt[k]] + L_head[row], L_rnxt[k] = k, L_head[row] + end + end + + ILUFactorization(L, U) +end + +# ─── Sylvester preconditioner builder ──────────────────────────────────────── +# +# Approximates the block-diagonal of the vectorised Sylvester operator +# L(X) = X − AXB +# When B has diagonal entries dⱼ the j-th n×n block is (I − dⱼ A). +# An ILU(τ) factorisation of each unique block serves as a right preconditioner +# for the Krylov solver (bicgstab / dqgmres / gmres). + +const DEFAULT_ILU_TAU = 1e-3 + +to_sparse(B::SparseMatrixCSC) = B +to_sparse(B::ThreadedSparseArrays.ThreadedSparseMatrixCSC) = B.A +to_sparse(B::AbstractMatrix) = sparse(B) + +""" + build_ilu_preconditioner(A, B; τ) → LinearOperator + +Build an ILU(τ) right preconditioner for the vectorised Sylvester operator. +Deduplicates identical diagonal entries of B so only one factorisation per +unique value is computed. Application loops over the m columns of the +solution matrix. +""" +function build_ilu_preconditioner(A::DenseMatrix{T}, + B::AbstractMatrix{T}; + τ::Float64 = DEFAULT_ILU_TAU) where {T <: AbstractFloat} + n = size(A, 1) + B_sp = to_sparse(B) + m = size(B_sp, 2) + diag_B = Vector{T}(undef, m) + @inbounds for j in 1:m + diag_B[j] = B_sp[j, j] + end + + A_sp = sparse(A) + I_n = sparse(one(T) * ℒ.I, n, n) + + # Factorise one block per unique diagonal value + cache = Dict{T, ILUFactorization}() + for d in diag_B + haskey(cache, d) && continue + block = I_n - d .* A_sp + droptol!(block, eps()) + cache[d] = ilu(block; τ = τ) + end + + # Map each column to its factorisation + factors = Vector{ILUFactorization}(undef, m) + @inbounds for j in 1:m + factors[j] = cache[diag_B[j]] + end + + nm = n * m + function precond_ldiv!(y, x) + X = reshape(x, n, m) + Y = reshape(y, n, m) + @inbounds for j in 1:m + ℒ.ldiv!(view(Y, :, j), factors[j], view(X, :, j)) + end + y + end + + LinearOperators.LinearOperator(T, nm, nm, false, false, precond_ldiv!) +end + + +end # @stable diff --git a/src/algorithms/quadratic_matrix_equation.jl b/src/algorithms/quadratic_matrix_equation.jl index da4a88c3a..116e3cebf 100644 --- a/src/algorithms/quadratic_matrix_equation.jl +++ b/src/algorithms/quadratic_matrix_equation.jl @@ -1,3 +1,5 @@ +@stable default_mode = "disable" begin + # Solves A * X ^ 2 + B * X + C = 0 # Algorithms: @@ -6,49 +8,87 @@ # Linear time iteration algorithm (:linear_time_iteration) [ -(A * X + B) \ C = X̂ ] - slow # Quadratic iteration algorithm (:quadratic_iteration) [ B \ A * X ^ 2 + B \ C = X̂ ] - very slow -@stable default_mode = "disable" begin 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-10, + acceptance_tol = 1e-6), + 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 - - AXX = A * X^2 - - AXXnorm = max(ℒ.norm(AXX), ℒ.norm(C)) + X² = qme_ws.temp3 + + # 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²) + norm_AXX = ℒ.norm(qme_ws.AXX) - ℒ.mul!(AXX, B, X, 1, 1) - - ℒ.axpy!(1, C, AXX) + # B*X into X² buffer (no longer needed for X²) + ℒ.mul!(X², B, X) + norm_BX = ℒ.norm(X²) - reached_tol = ℒ.norm(AXX) / AXXnorm + # Accumulate residual: AXX += B*X + C + ℒ.axpy!(1, X², qme_ws.AXX) + ℒ.axpy!(1, C, qme_ws.AXX) + + # Standard relative residual: norm(r) / (norm(A*X²) + norm(B*X) + norm(C)) + reached_tol = ℒ.norm(qme_ws.AXX) / (norm_AXX + norm_BX + ℒ.norm(C)) - 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 +97,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 +136,196 @@ 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[:,future_not_past_and_mixed_in_comb] + idx_constants = constants.post_complete_parameters - Ã₋ = 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 - Ã₀₊ = B[:,future_not_past_and_mixed_in_comb] - - Ã₀₋ = 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(Ã₀₋, Ã₊), hcat(I₋, Z₊)) + # Use views for read-only slices + Ã₊_view = @view A[:, future_not_past_and_mixed_in_comb] - ℒ.rmul!(Ã₋,-1) - ℒ.rmul!(Ã₀₊,-1) - E = vcat(hcat(Ã₋,Ã₀₊), 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, + has_ur = factorize_generalized_schur!((use_fastlapack_schur ? Val(:FastLapack) : Val(:Julia)), D, + E, + schur_ws_local.fast_qz_ws, + schur_ws_local.fast_qz_dims, + schur_ws_local.eigenselect) + + 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 - if verbose println("Quadratic matrix equation solver: schur - converged: false") end - return A, 0, 1.0 + if caching + cache.has_unit_roots = has_ur 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 - - Ẑ₁₁ = ℒ.lu(Z₁₁, check = false) + # 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!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), schur_ws_local.Z₁₁, + schur_ws_local.fast_lu_ws_z11, + schur_ws_local.fast_lu_dims_z11) - if !ℒ.issuccess(Ẑ₁₁) + if !solved_Z₁₁ if verbose println("Quadratic matrix equation solver: schur - converged: false") end return A, 0, 1.0 end - Ŝ₁₁ = ℒ.lu!(S₁₁, 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!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), schur_ws_local.S₁₁, + schur_ws_local.fast_lu_ws_s11, + schur_ws_local.fast_lu_dims_s11) - if !ℒ.issuccess(Ŝ₁₁) + if !solved_S₁₁ if verbose println("Quadratic matrix equation solver: schur - converged: false") end return A, 0, 1.0 end - # end # timeit_debug - # @timeit_debug timer "Matrix divisions" begin - - # D = Z₂₁ / Ẑ₁₁ - ℒ.rdiv!(Z₂₁, Ẑ₁₁) - D = Z₂₁ + # 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) - # L = Z₁₁ * (Ŝ₁₁ \ T₁₁) / Ẑ₁₁ - ℒ.ldiv!(Ŝ₁₁, T₁₁) - ℒ.mul!(S₁₁, Z₁₁, T₁₁) - ℒ.rdiv!(S₁₁, Ẑ₁₁) - 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 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) - AXXnorm = max(ℒ.norm(AXX), ℒ.norm(C)) + # 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₂₁) - ℒ.mul!(AXX, B, X, 1, 1) + # 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 - ℒ.axpy!(1, C, AXX) + ℒ.mul!(X, @view(sol[T.dynamic_order, :]), idx_constants.Ir_past_selector) - reached_tol = ℒ.norm(AXX) / AXXnorm + # 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) - # if reached_tol > tol - # println("QME: schur $reached_tol") - # end - - return X, iter, reached_tol # schur can fail + norm_AXX = ℒ.norm(schur_ws_local.AXX) + + # B*X into temp_X2 buffer (no longer needed for X²) + ℒ.mul!(schur_ws_local.temp_X2, B, X) + norm_BX = ℒ.norm(schur_ws_local.temp_X2) + + # Accumulate residual: AXX += B*X + C + ℒ.axpy!(1, schur_ws_local.temp_X2, schur_ws_local.AXX) + ℒ.axpy!(1, C, schur_ws_local.AXX) + + # Standard relative residual: norm(r) / (norm(A*X²) + norm(B*X) + norm(C)) + reached_tol = ℒ.norm(schur_ws_local.AXX) / (norm_AXX + norm_BX + ℒ.norm(C)) + + return X, 0, reached_tol end @@ -221,13 +334,20 @@ 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 = 50, + caching::Bool = true)::Tuple{Matrix{R}, Int64, R} where {R <: AbstractFloat} T = constants.post_model_macro + ensure_first_order_constants!(constants) + idx_constants = constants.post_complete_parameters + 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 @@ -236,10 +356,10 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, guess_provided = true n = size(A, 1) - if length(initial_guess) == 0 - guess_provided = false - initial_guess = zero(A) - end + # if length(initial_guess) == 0 + # guess_provided = false + # initial_guess = zero(A) + # end # Extract workspaces E = workspace.E @@ -261,29 +381,36 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, copy!(F, A) copy!(B̄, B) - ℒ.mul!(B̄, A, initial_guess, 1, 1) + # ℒ.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!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), B̄, + workspace.fast_lu_ws_qme_a, + workspace.fast_lu_dims_qme_a) + + 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) ℒ.rmul!(X, -1) - ℒ.axpy!(-1, initial_guess, X) + # ℒ.axpy!(-1, initial_guess, X) # Y = -F (in-place) copy!(Y, F) ℒ.rmul!(Y, -1) # 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 +432,14 @@ 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!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), temp1, + workspace.fast_lu_ws_qme_a, + workspace.fast_lu_dims_qme_a) - if !ℒ.issuccess(fEI) + if !solved_EI return A, iter, 1.0 end @@ -315,7 +447,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 +466,14 @@ 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!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), temp2, + workspace.fast_lu_ws_qme_b, + workspace.fast_lu_dims_qme_b) - if !ℒ.issuccess(fFI) + if !solved_FI return A, iter, 1.0 end @@ -342,7 +481,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 +492,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 +508,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 +522,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 @@ -397,27 +540,43 @@ function solve_quadratic_matrix_equation(A::AbstractMatrix{R}, end # end # timeit_debug - ℒ.axpy!(1, initial_guess, X_new) + # ℒ.axpy!(1, initial_guess, X_new) # Compute residual to verify solution quality # AXX = A * X_new^2 (use temp1 for X^2) ℒ.mul!(temp1, X_new, X_new) ℒ.mul!(AXX, A, temp1) - AXXnorm = max(ℒ.norm(AXX), ℒ.norm(C)) - - ℒ.mul!(AXX, B, X_new, 1, 1) + norm_AXX = ℒ.norm(AXX) + + # B*X into temp1 buffer (no longer needed for X²) + ℒ.mul!(temp1, B, X_new) + norm_BX = ℒ.norm(temp1) + # Accumulate residual: AXX += B*X + C + ℒ.axpy!(1, temp1, AXX) ℒ.axpy!(1, C, AXX) - reached_tol = ℒ.norm(AXX) / AXXnorm - + # Standard relative residual: norm(r) / (norm(A*X²) + norm(B*X) + norm(C)) + reached_tol = ℒ.norm(AXX) / (norm_AXX + norm_BX + ℒ.norm(C)) + # 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 @@ -532,4 +691,5 @@ end -end # dispatch_doctor + +end # @stable diff --git a/src/algorithms/sylvester.jl b/src/algorithms/sylvester.jl index c0d739a95..8e3efac23 100644 --- a/src/algorithms/sylvester.jl +++ b/src/algorithms/sylvester.jl @@ -1,3 +1,5 @@ +@stable default_mode = "disable" begin + # Available algorithms: # :doubling - fast, expensive part: B^2 # :bartels_stewart - fast, dense matrices only @@ -9,49 +11,91 @@ # :speedmapping - slow # solves: A * X * B + C = X for X -@stable default_mode = "disable" begin -function solve_sylvester_equation(A::M, +@unstable function 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, + preconditioner::Symbol = :ilu, + 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 + # 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 == :bartels_stewart - b = collect(B) - else - b = choose_matrix_format(B)# |> collect + if sylvester_algorithm == :bartels_stewart && !has_bartels_stewart() + error("The :bartels_stewart algorithm requires the MatrixEquations package. Run `using MatrixEquations` to enable it.") end 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) + + b = choose_matrix_format(B) + # b = B - c = choose_matrix_format(C)# |> sparse + 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) - # X = choose_matrix_format(initial_guess) + 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 - return initial_guess, true + 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 choose_matrix_format(guess), true end end @@ -60,6 +104,7 @@ function solve_sylvester_equation(A::M, x, i, reached_tol = solve_sylvester_equation(a, b, c, Val(sylvester_algorithm), 𝕊ℂ, initial_guess = initial_guess, + preconditioner = preconditioner, tol = tol, # timer = timer, verbose = verbose) @@ -68,16 +113,20 @@ function solve_sylvester_equation(A::M, println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: $sylvester_algorithm") 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) + if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && (sylvester_algorithm ≠ :bartels_stewart) && (length(B) < 5e7) && has_bartels_stewart() # try bartels_stewart if previous one didn't solve it + 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), 𝕊ℂ, initial_guess = zeros(0,0), + preconditioner = preconditioner, tol = tol, # timer = timer, verbose = verbose) @@ -88,13 +137,17 @@ 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), 𝕊ℂ, initial_guess = x, + preconditioner = preconditioner, tol = tol, # timer = timer, verbose = verbose) @@ -109,13 +162,16 @@ 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), 𝕊ℂ, initial_guess = zeros(0,0), + preconditioner = preconditioner, tol = tol, # timer = timer, verbose = verbose) @@ -126,13 +182,17 @@ 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), 𝕊ℂ, initial_guess = x, + preconditioner = preconditioner, tol = tol, # timer = timer, verbose = verbose) @@ -147,6 +207,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) @@ -154,6 +216,7 @@ function solve_sylvester_equation(A::M, x, i, reached_tol = solve_sylvester_equation(aa, b, cc, Val(:doubling), 𝕊ℂ, initial_guess = zeros(0,0), + preconditioner = preconditioner, tol = tol, # timer = timer, verbose = verbose) @@ -228,6 +291,21 @@ function solve_sylvester_equation(A::M, end +# Keep the low-level bartels-stewart signature available in core so fallback +# paths remain well-typed when MatrixEquations is not loaded. +function solve_sylvester_equation(A::DenseMatrix{T}, + B::AbstractMatrix{T}, + C::DenseMatrix{T}, + ::Val{:bartels_stewart}, + 𝕊ℂ::sylvester_workspace; + initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), + preconditioner::Symbol = :none, + verbose::Bool = false, + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + return copy(C), 0, T(Inf) +end + + function solve_sylvester_equation( A::AbstractSparseMatrix{T}, B::AbstractSparseMatrix{T}, @@ -236,39 +314,67 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{AbstractSparseMatrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + 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 - # guess_provided = false - initial_guess = zero(C) - end + # if length(initial_guess) == 0 + # # guess_provided = false + # initial_guess = zero(C) + # end 𝐀 = copy(A) 𝐁 = copy(B) # 𝐂 = length(init) == 0 ? copy(C) : copy(init) - 𝐂 = A * initial_guess * B + C - initial_guess #copy(C) + 𝐂 = copy(C) # - initial_guess #copy(C) - # ℒ.rmul!(𝐂, -1) - - max_iter = 500 + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end + # ℒ.rmul!(𝐂, -1) iters = max_iter for i in 1:max_iter 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - 𝐀 = 𝐀^2 - 𝐁 = 𝐁^2 - - droptol!(𝐀, eps()) - droptol!(𝐁, eps()) + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if issparse(cachedA) && size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) + 𝐀 = cachedA + else + 𝐀 = convert(typeof(𝐀), cachedA) + end + if issparse(cachedB) && size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) + 𝐁 = cachedB + else + 𝐁 = convert(typeof(𝐁), cachedB) + end + else + 𝐀 = 𝐀^2 + 𝐁 = 𝐁^2 + + droptol!(𝐀, eps()) + droptol!(𝐁, eps()) + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 @@ -284,13 +390,12 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, # reached_tol = denom == 0 ? 0.0 : ℒ.norm(𝐂¹ - 𝐂) / denom - 𝐂 += initial_guess - - reached_tol = ℒ.norm(A * 𝐂 * B + C - 𝐂) / max(ℒ.norm(𝐂), ℒ.norm(C)) + # 𝐂 += initial_guess - # 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 @@ -304,29 +409,44 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + 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 - + # 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 - max_iter = 500 + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end + # 𝐂 = A * initial_guess * B + C - initial_guess + fill!(𝐂B, 0) + # ℒ.mul!(𝐂B, initial_guess, B) + ℒ.mul!(𝐂, A, 𝐂B) + ℒ.axpy!(1, C, 𝐂) + # ℒ.axpy!(-1, initial_guess, 𝐂) iters = max_iter for i in 1:max_iter @@ -335,19 +455,43 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, ℒ.axpy!(1, 𝐂, 𝐂¹) # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - ℒ.mul!(𝐀¹,𝐀,𝐀) - copy!(𝐀,𝐀¹) - ℒ.mul!(𝐁¹,𝐁,𝐁) - copy!(𝐁,𝐁¹) - # 𝐀 = 𝐀^2 - # 𝐁 = 𝐁^2 - - droptol!(𝐀, eps()) - droptol!(𝐁, eps()) + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if issparse(cachedA) && size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) + 𝐀 = cachedA + else + 𝐀 = convert(typeof(𝐀), cachedA) + end + if issparse(cachedB) && size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) + 𝐁 = cachedB + else + 𝐁 = convert(typeof(𝐁), cachedB) + end + else + ℒ.mul!(𝐀¹,𝐀,𝐀) + copy!(𝐀,𝐀¹) + ℒ.mul!(𝐁¹,𝐁,𝐁) + copy!(𝐁,𝐁¹) + # 𝐀 = 𝐀^2 + # 𝐁 = 𝐁^2 + + droptol!(𝐀, eps()) + droptol!(𝐁, eps()) + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 +501,14 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, copy!(𝐂,𝐂¹) end - # ℒ.mul!(𝐂B, 𝐂, 𝐁) - # ℒ.mul!(𝐂¹, 𝐀, 𝐂B) - # ℒ.axpy!(1, 𝐂, 𝐂¹) - # # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - - # denom = max(ℒ.norm(𝐂), ℒ.norm(𝐂¹)) + # ℒ.axpy!(1, initial_guess, 𝐂) - # ℒ.axpy!(-1, 𝐂, 𝐂¹) + ℒ.mul!(𝐂B, 𝐂, B) + ℒ.mul!(𝐂¹, A, 𝐂B) + ℒ.axpy!(1, C, 𝐂¹) + ℒ.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 + reached_tol = ℒ.norm(𝐂¹) / max(ℒ.norm(𝐂), ℒ.norm(C)) return 𝐂, iters, reached_tol # return info on convergence end @@ -389,32 +523,48 @@ function solve_sylvester_equation( A::Matrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + 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 # guess_provided = true - if length(initial_guess) == 0 - # guess_provided = false - initial_guess = zero(C) - end - - 𝐀 = copy(A) - 𝐀¹ = copy(A) + # if length(initial_guess) == 0 + # # guess_provided = false + # initial_guess = zero(C) + # end + # 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) - - max_iter = 500 + + 𝐂 = 𝕊ℂ.𝐂_dbl + 𝐂¹ = 𝕊ℂ.𝐂¹ + 𝐂B = 𝕊ℂ.𝐂B + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end + + # 𝐂 = A * initial_guess * B + C - initial_guess + fill!(𝐂B, 0) + # ℒ.mul!(𝐂B, initial_guess, B) + ℒ.mul!(𝐂, A, 𝐂B) + ℒ.axpy!(1, C, 𝐂) + # ℒ.axpy!(-1, initial_guess, 𝐂) iters = max_iter # end # timeit_debug @@ -427,27 +577,51 @@ function solve_sylvester_equation( A::Matrix{T}, # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 # end # timeit_debug - # @timeit_debug timer "Square A" begin - ℒ.mul!(𝐀¹,𝐀,𝐀) - copy!(𝐀,𝐀¹) - # end # timeit_debug - - # 𝐀 = 𝐀^2 - # @timeit_debug timer "Square B" begin - 𝐁 = 𝐁^2 - # ℒ.mul!(𝐁¹,𝐁,𝐁) - # copy!(𝐁,𝐁¹) - # end # timeit_debug - - - # droptol!(𝐀, eps()) - # @timeit_debug timer "droptol B" begin - droptol!(𝐁, eps()) - # end # timeit_debug + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) && !issparse(cachedA) + copyto!(𝐀, cachedA) + else + copyto!(𝐀, convert(typeof(𝐀), cachedA)) + end + if issparse(cachedB) && size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) + 𝐁 = cachedB + else + 𝐁 = convert(typeof(𝐁), cachedB) + end + else + # @timeit_debug timer "Square A" begin + ℒ.mul!(𝐀¹,𝐀,𝐀) + copy!(𝐀,𝐀¹) + # end # timeit_debug + + # 𝐀 = 𝐀^2 + # @timeit_debug timer "Square B" begin + 𝐁 = 𝐁^2 + # ℒ.mul!(𝐁¹,𝐁,𝐁) + # copy!(𝐁,𝐁¹) + # end # timeit_debug + + + # droptol!(𝐀, eps()) + # @timeit_debug timer "droptol B" begin + droptol!(𝐁, eps()) + # end # timeit_debug + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 +633,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, 𝐂, 𝐂¹) + # ℒ.axpy!(1, initial_guess, 𝐂) - # 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 - - # 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 @@ -493,29 +653,45 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{T} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + 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 - if length(initial_guess) == 0 - # guess_provided = false - initial_guess = zero(C) - end + # 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) - - max_iter = 500 + # 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 + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end + + # 𝐂 = A * initial_guess * B + C - initial_guess + fill!(𝐂B, 0) + # ℒ.mul!(𝐂B, initial_guess, B) + ℒ.mul!(𝐂, A, 𝐂B) + ℒ.axpy!(1, C, 𝐂) + # ℒ.axpy!(-1, initial_guess, 𝐂) iters = max_iter for i in 1:max_iter @@ -524,17 +700,41 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, ℒ.axpy!(1, 𝐂, 𝐂¹) # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - 𝐀 = 𝐀^2 - ℒ.mul!(𝐁¹,𝐁,𝐁) - copy!(𝐁,𝐁¹) - # 𝐁 = 𝐁^2 - - droptol!(𝐀, eps()) - # droptol!(𝐁, eps()) + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if issparse(cachedA) && size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) + 𝐀 = cachedA + else + 𝐀 = convert(typeof(𝐀), cachedA) + end + if size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) && !issparse(cachedB) + copyto!(𝐁, cachedB) + else + copyto!(𝐁, convert(typeof(𝐁), cachedB)) + end + else + 𝐀 = 𝐀^2 + ℒ.mul!(𝐁¹,𝐁,𝐁) + copy!(𝐁,𝐁¹) + # 𝐁 = 𝐁^2 + + droptol!(𝐀, eps()) + # droptol!(𝐁, eps()) + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 +744,14 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, copy!(𝐂,𝐂¹) end - # ℒ.mul!(𝐂B, 𝐂, 𝐁) - # ℒ.mul!(𝐂¹, 𝐀, 𝐂B) - # ℒ.axpy!(1, 𝐂, 𝐂¹) - # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - - # denom = max(ℒ.norm(𝐂), ℒ.norm(𝐂¹)) + # ℒ.axpy!(1, initial_guess, 𝐂) - # ℒ.axpy!(-1, 𝐂, 𝐂¹) - - # 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 + ℒ.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 @@ -575,29 +765,38 @@ function solve_sylvester_equation( A::Matrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 # 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) + # if length(initial_guess) == 0 + # # guess_provided = false + # initial_guess = zero(C) + # end - # ℒ.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) + + 𝐂 = copy(C) # - initial_guess #copy(C) 𝐂¹ = similar(𝐂) - # 𝐂B = copy(C) - - max_iter = 500 + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end iters = max_iter for i in 1:max_iter @@ -606,19 +805,44 @@ function solve_sylvester_equation( A::Matrix{T}, # ℒ.axpy!(1, 𝐂, 𝐂¹) 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - ℒ.mul!(𝐀¹,𝐀,𝐀) - copy!(𝐀,𝐀¹) - # 𝐀 = 𝐀^2 - ℒ.mul!(𝐁¹,𝐁,𝐁) - copy!(𝐁,𝐁¹) - # 𝐁 = 𝐁^2 - - # droptol!(𝐀, eps()) - # droptol!(𝐁, eps()) + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) && !issparse(cachedA) + copyto!(𝐀, cachedA) + else + copyto!(𝐀, convert(typeof(𝐀), cachedA)) + end + if size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) && !issparse(cachedB) + copyto!(𝐁, cachedB) + else + copyto!(𝐁, convert(typeof(𝐁), cachedB)) + end + else + ℒ.mul!(𝐀¹,𝐀,𝐀) + copy!(𝐀,𝐀¹) + # 𝐀 = 𝐀^2 + ℒ.mul!(𝐁¹,𝐁,𝐁) + copy!(𝐁,𝐁¹) + # 𝐁 = 𝐁^2 + + # droptol!(𝐀, eps()) + # droptol!(𝐁, eps()) + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 +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 - 𝐂 += 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 @@ -658,29 +876,37 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + tol::SolverTolerances = SolverTolerances())::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat # see doi:10.1016/j.aml.2009.01.012 # guess_provided = true - if length(initial_guess) == 0 - # guess_provided = false - initial_guess = zero(C) - end + # 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) + # Use workspace for dense B buffers + n = size(A, 1) + m = size(B, 2) + ensure_sylvester_doubling_buffers!(𝕊ℂ, n, m) + + 𝐁 = 𝕊ℂ.𝐁 + 𝐁¹ = 𝕊ℂ.𝐁¹ + copyto!(𝐁, B) + + 𝐂 = copy(C) # - initial_guess 𝐂¹ = similar(𝐂) - # 𝐂B = copy(C) - - max_iter = 500 + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end iters = max_iter for i in 1:max_iter @@ -689,19 +915,44 @@ function solve_sylvester_equation( A::AbstractSparseMatrix{T}, # ℒ.axpy!(1, 𝐂, 𝐂¹) 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - # ℒ.mul!(𝐀¹,𝐀,𝐀) - # copy!(𝐀,𝐀¹) - 𝐀 = 𝐀^2 - ℒ.mul!(𝐁¹,𝐁,𝐁) - copy!(𝐁,𝐁¹) - # 𝐁 = 𝐁^2 - - droptol!(𝐀, eps()) - # droptol!(𝐁, eps()) + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if issparse(cachedA) && size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) + 𝐀 = cachedA + else + 𝐀 = convert(typeof(𝐀), cachedA) + end + if size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) && !issparse(cachedB) + copyto!(𝐁, cachedB) + else + copyto!(𝐁, convert(typeof(𝐁), cachedB)) + end + else + # ℒ.mul!(𝐀¹,𝐀,𝐀) + # copy!(𝐀,𝐀¹) + 𝐀 = 𝐀^2 + ℒ.mul!(𝐁¹,𝐁,𝐁) + copy!(𝐁,𝐁¹) + # 𝐁 = 𝐁^2 + + droptol!(𝐀, eps()) + # droptol!(𝐁, eps()) + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 +963,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 + # 𝐂 += 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 @@ -740,29 +985,37 @@ function solve_sylvester_equation( A::Matrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + 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) + # if length(initial_guess) == 0 + # # guess_provided = false + # initial_guess = zero(C) + # end + # 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) + + 𝐂 = copy(C)# - initial_guess 𝐂¹ = similar(𝐂) - # 𝐂B = copy(C) - - max_iter = 500 + if 𝕊ℂ.pow_capture + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end iters = max_iter for i in 1:max_iter @@ -771,19 +1024,44 @@ function solve_sylvester_equation( A::Matrix{T}, # ℒ.axpy!(1, 𝐂, 𝐂¹) 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - ℒ.mul!(𝐀¹,𝐀,𝐀) - copy!(𝐀,𝐀¹) - # 𝐀 = 𝐀^2 - # ℒ.mul!(𝐁¹,𝐁,𝐁) - # copy!(𝐁,𝐁¹) - 𝐁 = 𝐁^2 - - # droptol!(𝐀, eps()) - droptol!(𝐁, eps()) + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) && !issparse(cachedA) + copyto!(𝐀, cachedA) + else + copyto!(𝐀, convert(typeof(𝐀), cachedA)) + end + if issparse(cachedB) && size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) + 𝐁 = cachedB + else + 𝐁 = convert(typeof(𝐁), cachedB) + end + else + ℒ.mul!(𝐀¹,𝐀,𝐀) + copy!(𝐀,𝐀¹) + # 𝐀 = 𝐀^2 + # ℒ.mul!(𝐁¹,𝐁,𝐁) + # copy!(𝐁,𝐁¹) + 𝐁 = 𝐁^2 + + # droptol!(𝐀, eps()) + droptol!(𝐁, eps()) + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end 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 +1072,16 @@ function solve_sylvester_equation( A::Matrix{T}, 𝐂 = 𝐂¹ end - # ℒ.mul!(𝐂B, 𝐂, 𝐁) - # ℒ.mul!(𝐂¹, 𝐀, 𝐂B) - # ℒ.axpy!(1, 𝐂, 𝐂¹) - # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 - - # denom = max(ℒ.norm(𝐂), ℒ.norm(𝐂¹)) + # 𝐂 += initial_guess - # 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 @@ -822,17 +1094,20 @@ function solve_sylvester_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMat 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, verbose::Bool = false, - tol::Float64 = 1e-14)::Tuple{Matrix{T}, Int, T} where T <: AbstractFloat + max_iter::Int = 50, + 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 - if length(initial_guess) == 0 - # guess_provided = false - initial_guess = zero(C) - end + # if length(initial_guess) == 0 + # # guess_provided = false + # initial_guess = zero(C) + # end # Ensure workspaces are allocated n = size(A, 1) @@ -850,15 +1125,21 @@ function solve_sylvester_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMat copyto!(𝐀, A) copyto!(𝐁, B) - + + if 𝕊ℂ.pow_capture + if 𝕊ℂ.pow_iters < 1 + cache_set!(𝕊ℂ.𝐀_pow, 1, A, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, 1, B, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = 1 + end + end + # 𝐂 = A * initial_guess * B + C - initial_guess - ℒ.mul!(𝐂B, initial_guess, B) + fill!(𝐂B, 0) + # ℒ.mul!(𝐂B, initial_guess, B) ℒ.mul!(𝐂, A, 𝐂B) ℒ.axpy!(1, C, 𝐂) - ℒ.axpy!(-1, initial_guess, 𝐂) - - max_iter = 500 - + # ℒ.axpy!(-1, initial_guess, 𝐂) iters = max_iter # end # timeit_debug @@ -871,23 +1152,51 @@ function solve_sylvester_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMat # 𝐂¹ = 𝐀 * 𝐂 * 𝐁 + 𝐂 # end # timeit_debug - # @timeit_debug timer "Square A" begin - ℒ.mul!(𝐀¹,𝐀,𝐀) - copy!(𝐀,𝐀¹) - # end # timeit_debug - # @timeit_debug timer "Square B" begin - ℒ.mul!(𝐁¹,𝐁,𝐁) - copy!(𝐁,𝐁¹) - # end # timeit_debug - # 𝐀 = 𝐀^2 - # 𝐁 = 𝐁^2 + if 𝕊ℂ.pow_iters >= i + 1 + cachedA = 𝕊ℂ.𝐀_pow[i + 1] + cachedB = 𝕊ℂ.𝐁_pow[i + 1] + if size(𝐀) == size(cachedA) && eltype(𝐀) == eltype(cachedA) && !issparse(cachedA) + copyto!(𝐀, cachedA) + else + 𝐀 = Matrix{eltype(𝐀)}(cachedA) + 𝕊ℂ.𝐀 = 𝐀 + 𝐀¹ = 𝕊ℂ.𝐀¹ + end + if size(𝐁) == size(cachedB) && eltype(𝐁) == eltype(cachedB) && !issparse(cachedB) + copyto!(𝐁, cachedB) + else + 𝐁 = Matrix{eltype(𝐁)}(cachedB) + 𝕊ℂ.𝐁 = 𝐁 + 𝐁¹ = 𝕊ℂ.𝐁¹ + end + else + # @timeit_debug timer "Square A" begin + ℒ.mul!(𝐀¹,𝐀,𝐀) + copy!(𝐀,𝐀¹) + # end # timeit_debug + # @timeit_debug timer "Square B" begin + ℒ.mul!(𝐁¹,𝐁,𝐁) + copy!(𝐁,𝐁¹) + # end # timeit_debug + # 𝐀 = 𝐀^2 + # 𝐁 = 𝐁^2 + + if 𝕊ℂ.pow_capture + target_k = i + 1 + cache_set!(𝕊ℂ.𝐀_pow, target_k, 𝐀, 𝕊ℂ.pow_transposed) + cache_set!(𝕊ℂ.𝐁_pow, target_k, 𝐁, 𝕊ℂ.pow_transposed) + 𝕊ℂ.pow_iters = target_k + end + end # droptol!(𝐀, eps()) # 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 @@ -912,79 +1221,24 @@ function solve_sylvester_equation( A::Union{ℒ.Adjoint{T, Matrix{T}}, DenseMat # reached_tol = denom == 0 ? 0.0 : ℒ.norm(𝐂¹) / denom - ℒ.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 - - return copy(𝐂), iters, reached_tol # return info on convergence -end - - -function solve_sylvester_equation(A::DenseMatrix{T}, - B::Union{ℒ.Adjoint{Float64, Matrix{T}}, DenseMatrix{T}}, - C::DenseMatrix{T}, - ::Val{:bartels_stewart}, - 𝕊ℂ::sylvester_workspace; - 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 - # guess_provided = true + # ℒ.axpy!(1, initial_guess, 𝐂) - if length(initial_guess) == 0 - # guess_provided = false - initial_guess = zero(C) - end - - # Ensure workspaces are allocated (reuse Krylov buffers for tmp and 𝐂¹) - n = size(A, 1) - m = size(B, 2) - ensure_sylvester_krylov_buffers!(𝕊ℂ, n, m) - - # Use workspaces - 𝐂¹ = 𝕊ℂ.𝐂 - tmp̄ = 𝕊ℂ.tmp - - # 𝐂¹ = A * initial_guess * B + C - initial_guess - ℒ.mul!(tmp̄, initial_guess, B) - ℒ.mul!(𝐂¹, A, tmp̄) + ℒ.mul!(𝐂B, 𝐂, B) + ℒ.mul!(𝐂¹, A, 𝐂B) ℒ.axpy!(1, C, 𝐂¹) - ℒ.axpy!(-1, initial_guess, 𝐂¹) + ℒ.axpy!(-1, 𝐂, 𝐂¹) + + reached_tol = ℒ.norm(𝐂¹) / max(ℒ.norm(𝐂), ℒ.norm(C)) - 𝐂 = try - MatrixEquations.sylvd(-A, B, 𝐂¹)::Matrix{T} - catch - return C, 0, 1.0 - end + return 𝐂, iters, reached_tol # return info on convergence +end - # 𝐂¹ = A * 𝐂 * B + C - # denom = max(ℒ.norm(𝐂), ℒ.norm(𝐂¹)) +# Adjoint Sylvester via cached doubling powers. +# Solves: X = A' · X · B' + ∂P using cached A^(2^k), B^(2^k) (transposed on the fly). - # reached_tol = denom == 0 ? 0.0 : ℒ.norm(𝐂¹ - 𝐂) / denom - 𝐂 += initial_guess - ℒ.mul!(tmp̄, 𝐂, B) - ℒ.mul!(𝐂¹, A, tmp̄) - ℒ.axpy!(1, C, 𝐂¹) - ℒ.axpy!(-1, 𝐂, 𝐂¹) - - reached_tol = ℒ.norm(𝐂¹) / max(ℒ.norm(𝐂), ℒ.norm(C)) - # reached_tol = ℒ.norm(A * 𝐂 * B + C - 𝐂) / max(ℒ.norm(𝐂), ℒ.norm(C)) - - # if reached_tol > tol - # println("Sylvester: sylvester $reached_tol") - # end - - return 𝐂, -1, reached_tol # return info on convergence -end function solve_sylvester_equation(A::DenseMatrix{T}, @@ -994,8 +1248,10 @@ function solve_sylvester_equation(A::DenseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :ilu, 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 +1335,21 @@ 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, + preconditioner = :ilu + # Build ILU right preconditioner for large problems (opt-in) + N_precond = preconditioner == :ilu ? build_ilu_preconditioner(A, B) : ℒ.I + + Krylov.bicgstab!( 𝕊ℂ.krylov.bicgstab, sylvester, [vec(𝐂¹);], - # [vec(initial_guess);], + N = N_precond, + ldiv = false, 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) # else # 𝐂, info = Krylov.bicgstab(sylvester, [vec(C);], [vec(init);], rtol = tol / 10) # end @@ -1101,7 +1358,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 +1389,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 @@ -1146,8 +1403,10 @@ function solve_sylvester_equation(A::DenseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :ilu, 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 +1490,21 @@ 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, + + # Build ILU right preconditioner for large problems (opt-in) + N_precond = preconditioner == :ilu ? build_ilu_preconditioner(A, B) : ℒ.I + + Krylov.dqgmres!(𝕊ℂ.krylov.dqgmres, sylvester, [vec(𝐂¹);], - # [vec(initial_guess);], + N = N_precond, + ldiv = false, 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) # else # 𝐂, info = Krylov.dqgmres(sylvester, [vec(C);], [vec(init);], rtol = tol / 10) # end @@ -1253,7 +1513,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 +1544,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 @@ -1298,8 +1558,10 @@ function solve_sylvester_equation(A::DenseMatrix{T}, 𝕊ℂ::sylvester_workspace; initial_guess::AbstractMatrix{<:AbstractFloat} = zeros(0,0), # timer::TimerOutput = TimerOutput(), + preconditioner::Symbol = :none, 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 +1645,21 @@ 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, + + # Build ILU right preconditioner for large problems (opt-in) + N_precond = preconditioner == :ilu ? build_ilu_preconditioner(A, B) : ℒ.I + + Krylov.gmres!(𝕊ℂ.krylov.gmres, sylvester, [vec(𝐂¹);], - # [vec(initial_guess);], + N = N_precond, + ldiv = false, 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) # else # 𝐂, info = Krylov.gmres(sylvester, [vec(C);], [vec(init);], rtol = tol / 10) # end @@ -1405,7 +1668,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 +1699,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 @@ -1554,4 +1817,5 @@ end # return 𝐂, soll.maps, reached_tol # end -end # dispatch_doctor + +end # @stable diff --git a/src/common_docstrings.jl b/src/common_docstrings.jl index b0b2ef5ab..f9daa9fbf 100644 --- a/src/common_docstrings.jl +++ b/src/common_docstrings.jl @@ -21,9 +21,9 @@ const PARAMETER_DERIVATIVES® = "`parameter_derivatives` [Default: :all]: parame const DATA® = "`data` [Type: `KeyedArray`]: data matrix with variables (`String` or `Symbol`) in rows and periods in columns. Periods can have any format and will be used for the output. `KeyedArray` is provided by the `AxisKeys` package." const SMOOTH® = "`smooth` [Default: selector that enables smoothing when `filter = $(DEFAULT_FILTER_SELECTOR(DEFAULT_ALGORITHM))` and disables it otherwise, Type: `Bool`]: whether to return smoothed (`true`) or filtered (`false`) shocks/variables. Smoothing is only available for the Kalman filter. The inversion filter only returns filtered shocks/variables, so the default turns smoothing off in that case." const DATA_IN_LEVELS® = "`data_in_levels` [Default: `$(DEFAULT_DATA_IN_LEVELS)`, Type: `Bool`]: indicator whether the data is provided in levels. If `true` the input to the data argument will have the non-stochastic steady state subtracted." -const LYAPUNOV® = "`lyapunov_algorithm` [Default: `$(DEFAULT_LYAPUNOV_ALGORITHM)`, Type: `Symbol`]: algorithm to solve Lyapunov equation (`A * X * A' + C = X`). Available algorithms: `:doubling`, `:bartels_stewart`, `:bicgstab`, `:gmres`" -const SYLVESTER® = "`sylvester_algorithm` [Default: selector that uses `$(DEFAULT_SYLVESTER_ALGORITHM)` for smaller problems and switches to `$(DEFAULT_LARGE_SYLVESTER_ALGORITHM)` for larger problems, Type: `Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}}`]: algorithm to solve the Sylvester equation (`A * X * B + C = X`). Available algorithms: `:doubling`, `:bartels_stewart`, `:bicgstab`, `:dqgmres`, `:gmres`. Input argument can contain up to two elements in a `Vector` or `Tuple`. The first (second) element corresponds to the second (third) order perturbation solutions' Sylvester equation. If only one element is provided it corresponds to the second order perturbation solutions' Sylvester equation." -const QME® = "`quadratic_matrix_equation_algorithm` [Default: `$(DEFAULT_QME_ALGORITHM)`, Type: `Symbol`]: algorithm to solve quadratic matrix equation (`A * X ^ 2 + B * X + C = 0`). Available algorithms: `:schur`, `:doubling`" +const LYAPUNOV® = "`lyapunov_algorithm` [Default: `$(DEFAULT_LYAPUNOV_ALGORITHM)`, Type: `Symbol`]: algorithm to solve Lyapunov equation (`A * X * A' + C = X`). Available algorithms: `:doubling`, `:bartels_stewart` (requires `using MatrixEquations`), `:bicgstab`, `:gmres`, `:dqgmres`" +const SYLVESTER® = "`sylvester_algorithm` [Default: selector that uses `$(DEFAULT_SYLVESTER_ALGORITHM)` for smaller problems and switches to `$(DEFAULT_LARGE_SYLVESTER_ALGORITHM)` for larger problems, Type: `Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}}`]: algorithm to solve the Sylvester equation (`A * X * B + C = X`). Available algorithms: `:doubling`, `:bartels_stewart` (requires `using MatrixEquations`), `:bicgstab`, `:dqgmres`, `:gmres`. Input argument can contain up to two elements in a `Vector` or `Tuple`. The first (second) element corresponds to the second (third) order perturbation solutions' Sylvester equation. If only one element is provided it corresponds to the second order perturbation solutions' Sylvester equation." +const QME® = "`quadratic_matrix_equation_algorithm` [Default: selector that uses `$(DEFAULT_QME_ALGORITHM)` for smaller problems and switches to `$(DEFAULT_LARGE_QME_ALGORITHM)` when the QME problem size exceeds `$(DEFAULT_QME_THRESHOLD)`, Type: `Symbol`]: algorithm to solve quadratic matrix equation (`A * X ^ 2 + B * X + C = 0`). Available algorithms: `:schur`, `:doubling`" const VERBOSE® = "`verbose` [Default: `$(DEFAULT_VERBOSE)`, Type: `Bool`]: print information about results of the different solvers used to solve the model (non-stochastic steady state solver, Sylvester equations, Lyapunov equation, and quadratic matrix equation)." const TOLERANCES® = "`tol` [Default: `Tolerances()`, Type: `Tolerances`]: define various tolerances for the algorithm used to solve the model. See documentation of [`Tolerances`](@ref) for more details: `?Tolerances`." const PLOT_ATTRIBUTES® = "`plot_attributes` [Default: `$(DEFAULT_PLOT_ATTRIBUTES)`, Type: `Dict`]: pass on plot attributes for the top-level plot (see https://docs.juliaplots.org/latest/generated/attributes_plot/). E.g. Dict(:plot_titlefontcolor => :red)." @@ -44,3 +44,5 @@ const INITIAL_STATE®1 = "`initial_state` [Default: `$(DEFAULT_INITIAL_STATE)`, const LABEL® = "`label` [Type: `Union{Real, String, Symbol}`]: label to attribute to this function call in the plots. The default is the number of previous function calls since the last call to the function version with ! + 1." const RENAME_DICTIONARY® = "`rename_dictionary` [Default: `Dict()`, Type: `Dict{Symbol, String}`]: dictionary mapping variable or shock symbols to custom display names in plots. For example: `Dict(:dinve => \"Investment growth\", :c => \"Consumption\")`. Variables/shocks not in the dictionary will use their default names." const CONDITIONS_IN_LEVELS® = "`conditions_in_levels` [Default: `true`, Type: `Bool`]: indicator whether the conditions are provided in levels. If `true` the input to the conditions argument will have the relevant steady state subtracted (non-stochastic or stochastic steady state depending on the solution algorithm)." +const CACHING® = "`caching` [Default: `$(DEFAULT_CACHING)`, Type: `Bool`]: use cached intermediate results from previous calls to speed up computation. Set to `false` to force full recomputation." +const USE_WORKSPACES® = "`use_workspaces` [Default: `$(DEFAULT_USE_WORKSPACES)`, Type: `Bool`]: use pre-allocated workspace buffers for intermediate computations. Set to `false` to allocate fresh buffers, leaving the model's workspace state unchanged." diff --git a/src/custom_autodiff_rules/forwarddiff.jl b/src/custom_autodiff_rules/forwarddiff.jl deleted file mode 100644 index d30db7a19..000000000 --- a/src/custom_autodiff_rules/forwarddiff.jl +++ /dev/null @@ -1,720 +0,0 @@ -# ForwardDiff Dual number specializations for forward-mode automatic differentiation -# -# This file centralizes method specializations for ForwardDiff.Dual types, enabling -# forward-mode AD through the model solution pipeline. -# -# Strategy for each function: -# 1. Extract Float64 values from Dual numbers using ℱ.value.(...) -# 2. Compute the function result on Float64 values -# 3. Compute partials using implicit differentiation or chain rule -# 4. Reconstruct Dual numbers by combining values and partials -# -# Functions covered: -# - sparse_preallocated! -# - calculate_second/third_order_stochastic_steady_state -# - separate_values_and_partials_from_sparsevec_dual -# - get_NSSS_and_parameters -# - calculate_first_order_solution -# - solve_quadratic_matrix_equation -# - solve_sylvester_equation -# - solve_lyapunov_equation - - -function sparse_preallocated!(Ŝ::Matrix{ℱ.Dual{Z,S,N}}; ℂ::higher_order_workspace = Higher_order_workspace()) where {Z,S,N} - sparse(Ŝ) -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} - - 𝐒₁̂ = ℱ.value.(𝐒₁) - 𝐒₂̂ = ℱ.value.(𝐒₂) - x̂ = ℱ.value.(x) - - # Get cached computational constants - constants = initialise_constants!(𝓂) - so = constants.second_order - ℂ = 𝓂.workspaces.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 = 𝐒₂̂[T.past_not_future_and_mixed_idx,kron_s⁺_s] - B̂ = 𝐒₂̂[T.past_not_future_and_mixed_idx,kron_s⁺_s⁺] - - # Allocate or reuse workspace for partials - if size(ℂ.∂x_second_order) != (length(x̂), N) - ℂ.∂x_second_order = zeros(S, length(x̂), N) - else - fill!(ℂ.∂x_second_order, zero(S)) - end - ∂x̄ = ℂ.∂x_second_order - - 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̂) - break - 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) - - if solved - for i in 1:N - ∂𝐒₁ = ℱ.partials.(𝐒₁, i) - ∂𝐒₂ = ℱ.partials.(𝐒₂, i) - - ∂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⁺] - - tmp = ∂A * x̂ + ∂B̂ * ℒ.kron(vcat(x̂,1), vcat(x̂,1)) / 2 - - TMP = A + B * ℒ.kron(vcat(x̂,1), I_nPast) - I_nPast - - ∂x̄[:,i] = -TMP \ tmp - end - end - - return reshape(map(x̂, eachrow(∂x̄)) do v, p - ℱ.Dual{Z}(v, p...) # Z is the tag - 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} - 𝐒₁̂ = ℱ.value.(𝐒₁) - 𝐒₂̂ = ℱ.value.(𝐒₂) - 𝐒₃̂ = ℱ.value.(𝐒₃) - x̂ = ℱ.value.(x) - - # Get cached computational constants - so = ensure_computational_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 - - 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⁺] - - # Allocate or reuse workspace for partials - if size(ℂ.∂x_third_order) != (length(x̂), N) - ℂ.∂x_third_order = zeros(S, length(x̂), N) - else - fill!(ℂ.∂x_third_order, zero(S)) - end - ∂x̄ = ℂ.∂x_third_order - - 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̂) - break - 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) - - if solved - for i in 1:N - ∂𝐒₁ = ℱ.partials.(𝐒₁, i) - ∂𝐒₂ = ℱ.partials.(𝐒₂, i) - ∂𝐒₃ = ℱ.partials.(𝐒₃, i) - - ∂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⁺] - ∂Ĉ = ∂𝐒₃[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx,kron_s⁺_s⁺_s⁺] - - tmp = ∂A * x̂ + ∂B̂ * ℒ.kron(vcat(x̂,1), vcat(x̂,1)) / 2 + ∂Ĉ * ℒ.kron(vcat(x̂,1), ℒ.kron(vcat(x̂,1), vcat(x̂,1))) / 6 - - TMP = A + B * ℒ.kron(vcat(x̂,1), I_nPast) + C * ℒ.kron(ℒ.kron(vcat(x̂,1), vcat(x̂,1)), I_nPast) / 2 - I_nPast - - ∂x̄[:,i] = -TMP \ tmp - end - end - - return reshape(map(x̂, eachrow(∂x̄)) do v, p - ℱ.Dual{Z}(v, p...) # Z is the tag - end, size(x̂)), solved -end - -function separate_values_and_partials_from_sparsevec_dual(V::SparseVector{ℱ.Dual{Z,S,N}}; tol::AbstractFloat = eps()) where {Z,S,N} - nrows = length(V) - ncols = length(V.nzval[1].partials) - - rows = Int[] - cols = Int[] - - prtls = Float64[] - - for (i,v) in enumerate(V.nzind) - for (k,w) in enumerate(V.nzval[i].partials) - if abs(w) > tol - push!(rows,v) - push!(cols,k) - push!(prtls,w) - end - end - end - - vvals = sparsevec(V.nzind,[i.value for i in V.nzval],nrows) - ps = sparse(rows,cols,prtls,nrows,ncols) - - return vvals, ps -end - -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} - # timer::TimerOutput = TimerOutput(), - parameter_values = ℱ.value.(parameter_values_dual) - ms = ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) - qme_ws = 𝓂.workspaces.qme - - 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 - - # Allocate or reuse workspace for partials - if size(qme_ws.∂SS_and_pars) != (length(SS_and_pars), N) - qme_ws.∂SS_and_pars = zeros(S, length(SS_and_pars), N) - else - fill!(qme_ws.∂SS_and_pars, zero(S)) - end - ∂SS_and_pars = qme_ws.∂SS_and_pars - - if solution_error > opts.tol.NSSS_acceptance_tol || isnan(solution_error) - if opts.verbose println("Failed to find NSSS") end - - # Update failed counter - update_ss_counter!(𝓂.counters, false, estimation = estimation) - - solution_error = S(10.0) - else - # 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)) - - - ∂ = 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(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 - 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) - 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),:] - - 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 - - for i in 1:N - parameter_values_partials = ℱ.partials.(parameter_values_dual, i) - - ∂SS_and_pars[:,i] = jvp * parameter_values_partials - end - 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) -end - -function calculate_first_order_solution(∇₁::Matrix{ℱ.Dual{Z,S,N}}, - constants::constants, - qme_ws::qme_workspace, - sylv_ws::sylvester_workspace; - 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.(∇₁) - T = constants.post_model_macro - idx_constants = ensure_first_order_constants!(constants) - - 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] - - 𝐒₁, qme_sol, solved = calculate_first_order_solution(∇̂₁, constants, qme_ws, sylv_ws; opts = opts, initial_guess = initial_guess) - - if !solved - return ∇₁, qme_sol, false - end - - X = 𝐒₁[:,1:end-T.nExo] * expand_past - - AXB = A * X + B - - AXBfact = RF.lu(AXB, check = false) - - if !ℒ.issuccess(AXBfact) - AXBfact = ℒ.svd(AXB) - end - - invAXB = inv(AXBfact) - - AA = invAXB * A - - X² = X * X - - # Allocate or reuse workspace for partials (from qme_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 - fill!(qme_ws.X̃_first_order, zero(eltype(qme_ws.X̃_first_order))) - 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 - - initial_guess = zero(invAXB) - - # 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) - - if ℒ.norm(CC) < eps() continue end - - dX, solved = solve_sylvester_equation(AA, -X, -CC, sylv_ws, - initial_guess = initial_guess, - sylvester_algorithm = opts.sylvester_algorithm², - tol = opts.tol.sylvester_tol, - acceptance_tol = opts.tol.sylvester_acceptance_tol, - verbose = opts.verbose) - - # if !solved - # dX, solved = solve_sylvester_equation(AA, -X, -CC, - # sylvester_algorithm = :bicgstab, # more robust than sylvester - # initial_guess = initial_guess, - # verbose = verbose) - - # if !solved - # return ∇₁, qme_sol, false - # end - # end - - initial_guess = dX - - X̃[:,i] = vec(dX[:,T.past_not_future_and_mixed_idx]) - end - - x = reshape(map(𝐒₁[:,1:end-T.nExo], eachrow(X̃)) do v, p - ℱ.Dual{Z}(v, p...) # Z is the tag - end, size(𝐒₁[:,1:end-T.nExo])) - - Jm = expand_past - - ∇₊ = ∇₁[:,1:T.nFuture_not_past_and_mixed] * expand_future - ∇₀ = ∇₁[:,idx_constants.nabla_zero_cols] - ∇ₑ = ∇₁[:,idx_constants.nabla_e_start:end] - - B = -((∇₊ * x * Jm + ∇₀) \ ∇ₑ) - - return hcat(x, B), 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; - initial_guess::AbstractMatrix{<:Real} = zeros(0,0), - tol::AbstractFloat = 1e-8, - quadratic_matrix_equation_algorithm::Symbol = :schur, - verbose::Bool = false) where {Z,S,N} - T = constants.post_model_macro - # unpack: AoS -> SoA - Â = ℱ.value.(A) - B̂ = ℱ.value.(B) - Ĉ = ℱ.value.(C) - - X, solved = solve_quadratic_matrix_equation(Â, B̂, Ĉ, - Val(quadratic_matrix_equation_algorithm), - constants, - workspace; - tol = tol, - initial_guess = initial_guess, - # timer = timer, - verbose = verbose) - - AXB = Â * X + B̂ - - AXBfact = ℒ.lu(AXB, check = false) - - if !ℒ.issuccess(AXBfact) - AXBfact = ℒ.svd(AXB) - end - - invAXB = inv(AXBfact) - - AA = invAXB * Â - - 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) - else - fill!(workspace.X̃, zero(eltype(workspace.X̃))) - end - X̃ = workspace.X̃ - - # https://arxiv.org/abs/2011.11430 - for i in 1:N - dA = ℱ.partials.(A, i) - dB = ℱ.partials.(B, i) - dC = ℱ.partials.(C, i) - - CC = invAXB * (dA * X² + dB * X + dC) - - if ℒ.norm(CC) < eps() continue end - - dX, slvd = solve_sylvester_equation(AA, -X, -CC, workspace.sylvester_ws, sylvester_algorithm = :doubling) - - solved = Bool(solved) && Bool(slvd) - - X̃[:,i] = vec(dX) - end - - return reshape(map(X, eachrow(X̃)) do v, p - ℱ.Dual{Z}(v, p...) # Z is the tag - end, size(X)), solved -end - -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), - sylvester_algorithm::Symbol = :doubling, - acceptance_tol::AbstractFloat = 1e-10, - tol::AbstractFloat = 1e-14, - verbose::Bool = false)::Tuple{Matrix{ℱ.Dual{Z,S,N}}, Bool} where {Z,S,N} - # Extract Float64 values from Dual numbers - Â = ℱ.value.(A) - B̂ = ℱ.value.(B) - Ĉ = ℱ.value.(C) - - P̂, solved = solve_sylvester_equation(Â, B̂, Ĉ, 𝕊ℂ, - sylvester_algorithm = sylvester_algorithm, - tol = tol, - verbose = verbose, - initial_guess = initial_guess) - - # Allocate or reuse workspaces for temporary copies - if size(𝕊ℂ.Ã_fd) != size(Â) - 𝕊ℂ.Ã_fd = copy(Â) - else - copyto!(𝕊ℂ.Ã_fd, Â) - end - Ã = 𝕊ℂ.Ã_fd - - if size(𝕊ℂ.B̃_fd) != size(B̂) - 𝕊ℂ.B̃_fd = copy(B̂) - else - copyto!(𝕊ℂ.B̃_fd, B̂) - end - B̃ = 𝕊ℂ.B̃_fd - - if size(𝕊ℂ.C̃_fd) != size(Ĉ) - 𝕊ℂ.C̃_fd = copy(Ĉ) - else - copyto!(𝕊ℂ.C̃_fd, Ĉ) - end - C̃ = 𝕊ℂ.C̃_fd - - # Allocate or reuse workspace for partials - if size(𝕊ℂ.P̃) != (length(P̂), N) - 𝕊ℂ.P̃ = zeros(S, length(P̂), N) - else - fill!(𝕊ℂ.P̃, zero(S)) - end - P̃ = 𝕊ℂ.P̃ - - for i in 1:N - Ã .= ℱ.partials.(A, i) - B̃ .= ℱ.partials.(B, i) - C̃ .= ℱ.partials.(C, i) - - X = Ã * P̂ * B̂ + Â * P̂ * B̃ + C̃ - - if ℒ.norm(X) < eps() continue end - - P, slvd = solve_sylvester_equation(Â, B̂, X, 𝕊ℂ, - sylvester_algorithm = sylvester_algorithm, - tol = tol, - verbose = verbose) - - solved = solved && slvd - - P̃[:,i] = vec(P) - end - - return reshape(map(P̂, eachrow(P̃)) do v, p - ℱ.Dual{Z}(v, p...) # Z is the tag - end, size(P̂)), solved -end - -function solve_lyapunov_equation( A::AbstractMatrix{ℱ.Dual{Z,S,N}}, - C::AbstractMatrix{ℱ.Dual{Z,S,N}}, - workspace::lyapunov_workspace; - lyapunov_algorithm::Symbol = :doubling, - tol::AbstractFloat = 1e-14, - acceptance_tol::AbstractFloat = 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) - - # Allocate or reuse workspaces for temporary copies (from lyapunov_workspace) - if size(workspace.Ã_fd) != size(Â) - workspace.Ã_fd = copy(Â) - else - copyto!(workspace.Ã_fd, Â) - end - Ã = workspace.Ã_fd - - if size(workspace.C̃_fd) != size(Ĉ) - workspace.C̃_fd = copy(Ĉ) - else - copyto!(workspace.C̃_fd, Ĉ) - end - C̃ = workspace.C̃_fd - - # Allocate or reuse workspace for partials (from lyapunov_workspace) - if size(workspace.P̃) != (length(P̂), N) - workspace.P̃ = zeros(length(P̂), N) - else - fill!(workspace.P̃, zero(eltype(workspace.P̃))) - end - P̃ = workspace.P̃ - - # https://arxiv.org/abs/2011.11430 - for i in 1:N - Ã .= ℱ.partials.(A, i) - C̃ .= ℱ.partials.(C, i) - - X = Ã * P̂ * Â' + Â * P̂ * Ã' + C̃ - - if ℒ.norm(X) < eps() continue end - - 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 - ℱ.Dual{Z}(v, p...) # Z is the tag - end, size(P̂)), 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; - presample_periods::Int = 0, - 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)) - - z = C * u - - loglik = S(0.0) - - F = similar(C * C') - - K = similar(C') - - 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 - end - - v = data_in_deviations[:, t] - z - - F = C * P * C' - - luF = ℒ.lu(F, check = false) ### - - if !ℒ.issuccess(luF) - if 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 - return on_failure_loglikelihood - end - - invF = inv(luF) ### - - if t > presample_periods - loglik += log(Fdet) + ℒ.dot(v, invF, v)### - end - - K = P * C' * invF - - P = A * (P - K * C * P) * A' + 𝐁 - - u = A * (u + K * v) - - z = C * u - end - - # end # timeit_debug - - 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/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₊[dynIndex,:] * Ir[future_not_past_and_mixed_in_comb,:] - Ã₀ = A₀[dynIndex, comb] - Ã₋ = 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(Ã₊, Ã₀, Ã₋, 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₀[1:T.nPresent_only, T.present_only_idx] - A₊ᵤ = A₊[1:T.nPresent_only,:] - Ã₀ᵤ = A₀[1:T.nPresent_only, T.present_but_not_only_idx] - A₋ᵤ = A₋[1:T.nPresent_only,:] - - # end # timeit_debug - # @timeit_debug timer "Invert Ā₀ᵤ" begin - - Ā̂₀ᵤ = ℒ.lu!(Ā₀ᵤ, check = false) - - if !ℒ.issuccess(Ā̂₀ᵤ) - return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) - end - - # A = vcat(-(Ā̂₀ᵤ \ (A₊ᵤ * D * L + Ã₀ᵤ * sol[T.dynamic_order,:] + A₋ᵤ)), sol) - if T.nPresent_only > 0 - ℒ.mul!(A₋ᵤ, Ã₀ᵤ, sol[:,past_not_future_and_mixed_in_comb], 1, 1) - nₚ₋ = A₊ᵤ * D - ℒ.mul!(A₋ᵤ, nₚ₋, L, 1, 1) - ℒ.ldiv!(Ā̂₀ᵤ, 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) - 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) - 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)) - - ū = zeros(size(C,2)) - - P̄ = deepcopy(P) - - temp_N_N = similar(P) - - PCtmp = similar(C') - - F = similar(C * C') - - u = [similar(ū) for _ in 1:T] # used in backward pass - - P = [copy(P̄) for _ in 1:T] # used in backward pass - - CP = [zero(C) for _ in 1:T] # used in backward pass - - K = [similar(C') for _ in 1:T] # used in backward pass - - invF = [similar(F) for _ in 1:T] # used in backward pass - - v = [zeros(size(data_in_deviations, 1)) for _ in 1:T] # used in backward pass - - loglik = 0.0 - - # @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] + ū[t-1] - ℒ.mul!(u[t], K[t], v[t]) - u[t] .+= ū - - # ū[t] .= A * u[t] - ℒ.mul!(ū, A, u[t]) - - # z[t] .= C * ū[t] - ℒ.mul!(z, C, ū) - end - - llh = -(loglik + ((size(data_in_deviations, 2) - presample_periods) * size(data_in_deviations, 1)) * log(2 * 3.141592653589793)) / 2 - - # initialise derivative variables - ∂A = zero(A) - ∂F = zero(F) - ∂Faccum = zero(F) - ∂P = zero(P̄) - ∂ū = zero(ū) - ∂v = zero(v[1]) - ∂𝐁 = zero(𝐁) - ∂data_in_deviations = zero(data_in_deviations) - vtmp = zero(v[1]) - Ptmp = zero(P[1]) - - # 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!(∂ū, 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∂ū - # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] - # z[t] .= C * ū[t] - # ∂v = (invF[t]' + invF[t]) * v[t] - copy!(invF[1], invF[t]' .+ invF[t]) - # copy!(invF[1], invF[t]) # using invF[1] as temporary storage - # ℒ.axpy!(1, invF[t]', invF[1]) # using invF[1] as temporary storage - ℒ.mul!(∂v, invF[1], v[t]) - # ℒ.mul!(∂ū∂v, C', v[1]) - else - ℒ.rmul!(∂F, 0) - ℒ.rmul!(∂v, 0) - end - - # ∂F∂P - # F[t] .= C * P̄[t-1] * C' - # ∂P += C' * (∂F + ∂Faccum) * C - ℒ.axpy!(1, ∂Faccum, ∂F) - ℒ.mul!(PCtmp, C', ∂F) - ℒ.mul!(∂P, PCtmp, C, 1, 1) - - # ∂ū∂P - # K[t] .= P̄[t-1] * C' * invF[t] - # u[t] .= K[t] * v[t] + ū[t-1] - # ū[t] .= A * u[t] - # ∂P += A' * ∂ū * v[t]' * invF[t]' * C - ℒ.mul!(CP[1], invF[t]', C) # using CP[1] as temporary storage - ℒ.mul!(PCtmp, ∂ū , v[t]') - ℒ.mul!(P[1], PCtmp , CP[1]) # using P[1] as temporary storage - ℒ.mul!(∂P, A', P[1], 1, 1) - - # ∂ū∂data - # v[t] .= data_in_deviations[:, t-1] .- z - # z[t] .= C * ū[t] - # ∂data_in_deviations[:,t-1] = -C * ∂ū - ℒ.mul!(u[1], A', ∂ū) - ℒ.mul!(v[1], K[t]', u[1]) # using v[1] as temporary storage - ℒ.axpy!(1, ∂v, v[1]) - ∂data_in_deviations[:,t-1] .= v[1] - # ℒ.mul!(∂data_in_deviations[:,t-1], C, ∂ū, -1, 0) # cannot assign to columns in matrix, must be whole matrix - - # ∂ū∂ū - # z[t] .= C * ū[t] - # v[t] .= data_in_deviations[:, t-1] .- z - # K[t] .= P̄[t-1] * C' * invF[t] - # u[t] .= K[t] * v[t] + ū[t-1] - # ū[t] .= A * u[t] - # step to next iteration - # ∂ū = A' * ∂ū - C' * K[t]' * A' * ∂ū - ℒ.mul!(u[1], A', ∂ū) # using u[1] as temporary storage - ℒ.mul!(v[1], K[t]', u[1]) # using v[1] as temporary storage - ℒ.mul!(∂ū, C', v[1]) - ℒ.mul!(u[1], C', v[1], -1, 1) - copy!(∂ū, u[1]) - - # ∂llh∂ū - # loglik += logdet(F[t]) + v[t]' * invF[t] * v[t] - # v[t] .= data_in_deviations[:, t-1] .- z - # z[t] .= C * ū[t] - # ∂ū -= ∂ū∂v - ℒ.mul!(u[1], C', ∂v) # using u[1] as temporary storage - ℒ.axpy!(-1, u[1], ∂ū) - - if t > 2 - # ∂ū∂A - # ū[t] .= A * u[t] - # ∂A += ∂ū * u[t-1]' - ℒ.mul!(∂A, ∂ū, u[t-1]', 1, 1) - - # ∂P̄∂A and ∂P̄∂𝐁 - # P̄[t] .= A * P[t] * A' + 𝐁 - # ∂A += ∂P * A * P[t-1]' + ∂P' * A * P[t-1] - ℒ.mul!(P[1], A, P[t-1]') - ℒ.mul!(Ptmp ,∂P, P[1]) - ℒ.mul!(P[1], A, P[t-1]) - ℒ.mul!(Ptmp ,∂P', P[1], 1, 1) - ℒ.axpy!(1, Ptmp, ∂A) - - # ∂𝐁 += ∂P - ℒ.axpy!(1, ∂P, ∂𝐁) - - # ∂P∂P - # P[t] .= P̄[t-1] - K[t] * C * P̄[t-1] - # P̄[t] .= A * P[t] * A' + 𝐁 - # step to next iteration - # ∂P = A' * ∂P * A - ℒ.mul!(P[1], ∂P, A) # using P[1] as temporary storage - ℒ.mul!(∂P, A', P[1]) - - # ∂P̄∂P - # K[t] .= P̄[t-1] * C' * invF[t] - # P[t] .= P̄[t-1] - K[t] * CP[t] - # ∂P -= C' * K[t-1]' * ∂P + ∂P * K[t-1] * C - ℒ.mul!(PCtmp, ∂P, K[t-1]) - ℒ.mul!(CP[1], K[t-1]', ∂P) # using CP[1] as temporary storage - ℒ.mul!(∂P, PCtmp, C, -1, 1) - ℒ.mul!(∂P, C', CP[1], -1, 1) - - # ∂ū∂F - # K[t] .= P̄[t-1] * C' * invF[t] - # u[t] .= K[t] * v[t] + ū[t-1] - # ū[t] .= A * u[t] - # ∂Faccum = -invF[t-1]' * CP[t-1] * A' * ∂ū * v[t-1]' * invF[t-1]' - ℒ.mul!(u[1], A', ∂ū) # using u[1] as temporary storage - ℒ.mul!(v[1], CP[t-1], u[1]) # using v[1] as temporary storage - ℒ.mul!(vtmp, invF[t-1]', v[1], -1, 0) - ℒ.mul!(invF[1], vtmp, v[t-1]') # using invF[1] as temporary storage - ℒ.mul!(∂Faccum, invF[1], invF[t-1]') - - # ∂P∂F - # K[t] .= P̄[t-1] * C' * invF[t] - # P[t] .= P̄[t-1] - K[t] * CP[t] - # ∂Faccum -= invF[t-1]' * CP[t-1] * ∂P * CP[t-1]' * invF[t-1]' - ℒ.mul!(CP[1], invF[t-1]', CP[t-1]) # using CP[1] as temporary storage - ℒ.mul!(PCtmp, CP[t-1]', invF[t-1]') - ℒ.mul!(K[1], ∂P, PCtmp) # using K[1] as temporary storage - ℒ.mul!(∂Faccum, CP[1], K[1], -1, 1) - - end - end - - ℒ.rmul!(∂P, -∂llh/2) - ℒ.rmul!(∂A, -∂llh/2) - ℒ.rmul!(∂𝐁, -∂llh/2) - ℒ.rmul!(∂data_in_deviations, -∂llh/2) - - # 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..430532571 100644 --- a/src/default_options.jl +++ b/src/default_options.jl @@ -56,6 +56,7 @@ const DEFAULT_MEAN_FLAG = false const DEFAULT_STANDARD_DEVIATION_FLAG = true const DEFAULT_VARIANCE_FLAG = false const DEFAULT_COVARIANCE_FLAG = false +const DEFAULT_CORRELATION_FLAG = false const DEFAULT_AUTOCORRELATION_FLAG = false const DEFAULT_DERIVATIVES_FLAG = true const DEFAULT_STOCHASTIC_FLAG = false @@ -65,6 +66,9 @@ const DEFAULT_SILENT_FLAG = false # Solver and tolerance defaults const DEFAULT_VERBOSE = false const DEFAULT_QME_ALGORITHM = :schur +const DEFAULT_QME_THRESHOLD = 1000000 +const DEFAULT_LARGE_QME_ALGORITHM = :doubling +const DEFAULT_QME_SELECTOR = 𝓂 -> (𝓂.constants.post_model_macro.nVars - 𝓂.constants.post_model_macro.nPresent_only)^2 > DEFAULT_QME_THRESHOLD ? DEFAULT_LARGE_QME_ALGORITHM : DEFAULT_QME_ALGORITHM const DEFAULT_LYAPUNOV_ALGORITHM = :doubling const DEFAULT_SYLVESTER_ALGORITHM = :doubling const DEFAULT_SYLVESTER_THRESHOLD = 1000 @@ -105,21 +109,13 @@ 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", ) # Turing distribution wrapper defaults const DEFAULT_TURING_USE_MEAN_STD = false -const DEFAULT_MAXLOG = 3 \ No newline at end of file +const DEFAULT_MAXLOG = 3 + +# Caching and workspace defaults +const DEFAULT_CACHING = true +const DEFAULT_USE_WORKSPACES = true \ No newline at end of file diff --git a/src/dynare.jl b/src/dynare.jl index f4c9b4d74..09b6cee96 100644 --- a/src/dynare.jl +++ b/src/dynare.jl @@ -26,8 +26,9 @@ function translate_mod_file(path_to_mod_file::AbstractString) args = [tmp * "/" * basename(path_to_mod_file), "language=julia", "json=compute"] + current_directory = pwd() + if length(directory) > 0 - current_directory = pwd() cd(directory) end @@ -143,7 +144,7 @@ The recommended workflow is to use this function to write a .mod-file, and then # Arguments - $MODEL® """ -function write_mod_file(𝓂::ℳ) +function write_mod_file(𝓂::ℳ; order::Int = 1, pruning::Bool = false, irf_periods::Int = 40) NSSS = get_SS(𝓂, derivatives = false) index_in_name = NSSS.keys isa Base.RefValue{Vector{String}} @@ -210,7 +211,14 @@ function write_mod_file(𝓂::ℳ) print(io, "\t" * translate_symbol_to_ascii(v) * "\t=\t" * string(NSSS(index_in_name ? replace(string(v), "◖" => "{", "◗" => "}") : v)) * ";\n") end - println(io, "end;\n\nstoch_simul(order = 1, irf = 40);") + stoch_opts = "order = $order, irf = $irf_periods" + if pruning + stoch_opts *= ", pruning" + end + if order > 2 + stoch_opts *= ", k_order_solver" + end + println(io, "end;\n\nstoch_simul($stoch_opts);") end @info "Created " * 𝓂.model_name * ".mod" @@ -252,3 +260,52 @@ write_to_dynare = write_mod_file See [`write_mod_file`](@ref) """ export_model = write_mod_file + + + + +function translate_symbol_to_ascii(x::Symbol) + ss = Unicode.normalize(replace(string(x), "◖" => "__", "◗" => "__"), :NFD) + + outstr = "" + + for i in ss + out = REPL.symbol_latex(string(i))[2:end] + if out == "" + outstr *= string(i) + else + outstr *= replace(out, + r"\!" => s"_", + r"\(" => s"_", + r"\)" => s"_", + r"\^" => s"_", + r"\_\^" => s"_", + r"\+" => s"plus", + r"\-" => s"minus", + r"\*" => s"times") + if i != ss[end] + outstr *= "_" + end + end + end + + return outstr +end + + +function translate_expression_to_ascii(exp::Expr) + postwalk(x -> + x isa Symbol ? + begin + x_tmp = translate_symbol_to_ascii(x) + + if x_tmp == string(x) + x + else + Symbol(x_tmp) + end + end : + x, + exp) +end + diff --git a/src/filter/find_shocks.jl b/src/filter/find_shocks.jl index a014d02c0..acccbc7cf 100644 --- a/src/filter/find_shocks.jl +++ b/src/filter/find_shocks.jl @@ -1,3 +1,5 @@ +@stable default_mode = "disable" begin + # Algorithms # - LagrangeNewton: fast, but no guarantee of convergence to global minimum # - COBYLA: best known chances of convergence to global minimum; ok speed for third order; lower tol on optimality conditions (1e-7) @@ -12,7 +14,6 @@ # conditional forecast constraints (only some variables match target values) # Uses analytical derivatives from perturbation solution matrices (like find_shocks) -@stable default_mode = "disable" begin function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, initial_state::Union{Vector{Float64}, Vector{Vector{Float64}}}, all_shocks::Vector{Float64}, @@ -50,17 +51,15 @@ function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, var²_idxs = so.var²_idxs shockvar_idxs = sparse(ℒ.kron(so.e_in_s⁺, so.s_in_s)).nzind - if third_order - var_vol³_idxs = to.var_vol³_idxs - shock_idxs2 = to.shock_idxs2 - shock_idxs3 = to.shock_idxs3 - shock³_idxs = to.shock³_idxs - shockvar1_idxs = to.shockvar1_idxs - shockvar2_idxs = to.shockvar2_idxs - shockvar3_idxs = to.shockvar3_idxs - shockvar³2_idxs = to.shockvar³2_idxs - shockvar³_idxs = to.shockvar³_idxs - end + var_vol³_idxs = to.var_vol³_idxs + shock_idxs2 = to.shock_idxs2 + shock_idxs3 = to.shock_idxs3 + shock³_idxs = to.shock³_idxs + shockvar1_idxs = to.shockvar1_idxs + shockvar2_idxs = to.shockvar2_idxs + shockvar3_idxs = to.shockvar3_idxs + shockvar³2_idxs = to.shockvar³2_idxs + shockvar³_idxs = to.shockvar³_idxs fixed_shock_idx = setdiff(1:n_exo, free_shock_idx) @@ -70,6 +69,14 @@ function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, J = ℒ.I(n_exo) + nPast = T.nPast_not_future_and_mixed + third_order_pruning = third_order && pruning + ensure_find_shocks_state_buffers!(ws, n_exo, nPast; + third_order = third_order, + third_order_pruning = third_order_pruning) + kron_state_vol = ws.kron_state_vol + kron_I_state = ws.kron_I_state + if isnothing(𝐒₃) # Second order (pruned or non-pruned) if pruning @@ -93,10 +100,10 @@ function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, 𝐒²⁻ᵉ = @views 𝐒₂[cond_var_idx, shockvar²_idxs] 𝐒²ᵉ = @views 𝐒₂[cond_var_idx, shock²_idxs] - kron_state_vol = ℒ.kron(state_vol, state_vol) + ℒ.kron!(kron_state_vol, state_vol, state_vol) ℒ.mul!(shock_independent, 𝐒²⁻ᵛ, kron_state_vol, -1/2, 1) - kron_I_state = ℒ.kron(J, state_vol) + ℒ.kron!(kron_I_state, J, state_vol) 𝐒ⁱ = 𝐒¹ᵉ + 𝐒²⁻ᵉ * kron_I_state 𝐒ⁱ²ᵉ = 𝐒²ᵉ / 2 end @@ -118,17 +125,17 @@ function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, 𝐒²⁻ᵉ = @views 𝐒₂[cond_var_idx, shockvar²_idxs] 𝐒²ᵉ = @views 𝐒₂[cond_var_idx, shock²_idxs] - kron_state_vol = ℒ.kron(state_vol, state_vol) + ℒ.kron!(kron_state_vol, state_vol, state_vol) ℒ.mul!(shock_independent, 𝐒²⁻ᵛ, kron_state_vol, -1/2, 1) - kron_I_state = ℒ.kron(J, state_vol) + ℒ.kron!(kron_I_state, J, state_vol) 𝐒ⁱ = 𝐒¹ᵉ + 𝐒²⁻ᵉ * kron_I_state 𝐒ⁱ²ᵉ = 𝐒²ᵉ / 2 end end 𝐒ⁱ³ᵉ = nothing - elseif third_order + else # third_order # Third order (pruned or non-pruned) II = sparse(ℒ.I(n_exo^2)) @@ -158,18 +165,22 @@ function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, ℒ.mul!(shock_independent, 𝐒¹⁻, state₂, -1, 1) ℒ.mul!(shock_independent, 𝐒¹⁻, state₃, -1, 1) - kron_state_vol = ℒ.kron(state_vol, state_vol) + ℒ.kron!(kron_state_vol, state_vol, state_vol) ℒ.mul!(shock_independent, 𝐒²⁻ᵛ, kron_state_vol, -1/2, 1) - kron_state₁₂ = ℒ.kron(state₁, state₂) + kron_state₁₂ = ws.kron_state₁₂ + ℒ.kron!(kron_state₁₂, state₁, state₂) ℒ.mul!(shock_independent, 𝐒²⁻, kron_state₁₂, -1, 1) - kron_state_vol3 = ℒ.kron(state_vol, kron_state_vol) + kron_state_vol3 = ws.kron_state_vol3 + ℒ.kron!(kron_state_vol3, state_vol, kron_state_vol) ℒ.mul!(shock_independent, 𝐒³⁻ᵛ, kron_state_vol3, -1/6, 1) - kron_I_state = ℒ.kron(J, state_vol) - kron_I_state₂ = ℒ.kron(J, state₂) - kron_I_state_state = ℒ.kron(J, kron_state_vol) + ℒ.kron!(kron_I_state, J, state_vol) + kron_I_state₂ = ws.kron_I_state₂ + ℒ.kron!(kron_I_state₂, J, state₂) + kron_I_state_state = ws.kron_I_state_state + ℒ.kron!(kron_I_state_state, J, kron_state_vol) 𝐒ⁱ = 𝐒¹ᵉ + 𝐒²⁻ᵉ * kron_I_state + @@ -197,14 +208,16 @@ function find_shocks_conditional_forecast(::Val{:LagrangeNewton}, shock_independent = copy(conditions) ℒ.mul!(shock_independent, 𝐒¹⁻ᵛ, state_vol, -1, 1) - kron_state_vol = ℒ.kron(state_vol, state_vol) + ℒ.kron!(kron_state_vol, state_vol, state_vol) ℒ.mul!(shock_independent, 𝐒²⁻ᵛ, kron_state_vol, -1/2, 1) - kron_state_vol3 = ℒ.kron(state_vol, kron_state_vol) + kron_state_vol3 = ws.kron_state_vol3 + ℒ.kron!(kron_state_vol3, state_vol, kron_state_vol) ℒ.mul!(shock_independent, 𝐒³⁻ᵛ, kron_state_vol3, -1/6, 1) - kron_I_state = ℒ.kron(J, state_vol) - kron_I_state_state = ℒ.kron(J, kron_state_vol) + ℒ.kron!(kron_I_state, J, state_vol) + kron_I_state_state = ws.kron_I_state_state + ℒ.kron!(kron_I_state_state, J, kron_state_vol) 𝐒ⁱ = 𝐒¹ᵉ + 𝐒²⁻ᵉ * kron_I_state + @@ -882,19 +895,19 @@ solver started at the origin (not just LagrangeNewton) will converge to the root basin contains the origin rather than guaranteeing the global optimum. """ function find_shocks(::Val{:LagrangeNewton}, - initial_guess::Vector{Float64}, - kron_buffer::Vector{Float64}, - kron_buffer2::AbstractMatrix{Float64}, + initial_guess::Vector{R}, + kron_buffer::Vector{R}, + kron_buffer2::AbstractMatrix{R}, J::ℒ.Diagonal{Bool, Vector{Bool}}, - 𝐒ⁱ::AbstractMatrix{Float64}, - 𝐒ⁱ²ᵉ::AbstractMatrix{Float64}, - shock_independent::Vector{Float64}; + 𝐒ⁱ::AbstractMatrix{R}, + 𝐒ⁱ²ᵉ::AbstractMatrix{R}, + shock_independent::Vector{R}; max_iter::Int = 1000, - tol::Float64 = 1e-13, - verbose::Bool = false) # will fail for higher or lower precision + tol::AbstractFloat = 1e-13, + verbose::Bool = false) where R <: Real x = copy(initial_guess) - λ = zeros(size(𝐒ⁱ, 1)) + λ = zeros(R, size(𝐒ⁱ, 1)) xλ = [ x λ ] @@ -903,23 +916,23 @@ function find_shocks(::Val{:LagrangeNewton}, norm1 = ℒ.norm(shock_independent) - norm2 = 1.0 + norm2 = one(R) - Δnorm = 1e12 + Δnorm = R(1e12) x̂ = copy(shock_independent) - x̄ = zeros(size(𝐒ⁱ,2)) + x̄ = zeros(R, size(𝐒ⁱ,2)) ∂x = zero(𝐒ⁱ) - fxλ = zeros(length(xλ)) + fxλ = zeros(R, length(xλ)) - fxλp = zeros(length(xλ), length(xλ)) + fxλp = zeros(R, length(xλ), length(xλ)) - tmp = zeros(size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) + tmp = zeros(R, size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) - lI = -2 * vec(ℒ.I(size(𝐒ⁱ, 2))) + lI = R(-2) * vec(ℒ.I(size(𝐒ⁱ, 2))) iter = 0 @inbounds for i in 1:max_iter @@ -957,14 +970,13 @@ function find_shocks(::Val{:LagrangeNewton}, # return x, false # end - try - f̂xλp = ℒ.factorize(fxλp) - ℒ.ldiv!(Δxλ, f̂xλp, fxλ) - catch - # ℒ.svd(fxλp) - # println("factorization fails") + f̂xλp = ℒ.lu(fxλp, check = false) + + if !ℒ.issuccess(f̂xλp) return x, false end + + ℒ.ldiv!(Δxλ, f̂xλp, fxλ) if !all(isfinite,Δxλ) break end @@ -988,28 +1000,10 @@ function find_shocks(::Val{:LagrangeNewton}, ℒ.axpby!(1, shock_independent, -1, x̂) if ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol) - # println("LagrangeNewton: $i, Tol reached, $x") break end - - # if i > 500 && ℒ.norm(Δxλ) > 1e-11 && ℒ.norm(Δxλ) > Δnorm - # # println("LagrangeNewton: $i, Norm increase") - # return x, false - # end - # # if i == max_iter - # println("LagrangeNewton: $i, Max iter reached") - # println(ℒ.norm(Δxλ) / ℒ.norm(xλ)) - # end end - # println(λ) - # println("Norm: $(ℒ.norm(x̂) / max(norm1,norm2))") - # println(ℒ.norm(Δxλ)) - # println(ℒ.norm(Δxλ) / ℒ.norm(xλ)) - # if !(ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol)) - # println("Find shocks failed. Norm 1: $(ℒ.norm(x̂) / max(norm1,norm2)); Norm 2: $(ℒ.norm(Δxλ) / ℒ.norm(xλ))") - # end - residual = ℒ.norm(x̂) / max(norm1,norm2) step_norm = ℒ.norm(Δxλ) / ℒ.norm(xλ) matched = residual < tol && step_norm < sqrt(tol) @@ -1019,23 +1013,23 @@ end function 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}, + initial_guess::Vector{R}, + kron_buffer::Vector{R}, + kron_buffer²::Vector{R}, + kron_buffer2::AbstractMatrix{R}, + kron_buffer3::AbstractMatrix{R}, + kron_buffer4::AbstractMatrix{R}, J::ℒ.Diagonal{Bool, Vector{Bool}}, - 𝐒ⁱ::AbstractMatrix{Float64}, - 𝐒ⁱ²ᵉ::AbstractMatrix{Float64}, - 𝐒ⁱ³ᵉ::AbstractMatrix{Float64}, - shock_independent::Vector{Float64}; + 𝐒ⁱ::AbstractMatrix{R}, + 𝐒ⁱ²ᵉ::AbstractMatrix{R}, + 𝐒ⁱ³ᵉ::AbstractMatrix{R}, + shock_independent::Vector{R}; max_iter::Int = 1000, - tol::Float64 = 1e-13, - verbose::Bool = false) # will fail for higher or lower precision + tol::AbstractFloat = 1e-13, + verbose::Bool = false) where R <: Real x = copy(initial_guess) - λ = zeros(size(𝐒ⁱ, 1)) + λ = zeros(R, size(𝐒ⁱ, 1)) xλ = [ x λ ] @@ -1044,33 +1038,35 @@ function find_shocks(::Val{:LagrangeNewton}, norm1 = ℒ.norm(shock_independent) - norm2 = 1.0 + norm2 = one(R) - Δnorm = 1e12 + Δnorm = R(1e12) x̂ = copy(shock_independent) - x̄ = zeros(size(𝐒ⁱ,2)) + x̄ = zeros(R, size(𝐒ⁱ,2)) ∂x = zero(𝐒ⁱ) ∂x̂ = zero(𝐒ⁱ) - fxλ = zeros(length(xλ)) + fxλ = zeros(R, length(xλ)) - fxλp = zeros(length(xλ), length(xλ)) + fxλp = zeros(R, length(xλ), length(xλ)) - tmp = zeros(size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) + tmp = zeros(R, size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) - tmp2 = zeros(size(𝐒ⁱ, 1),size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) + tmp2 = zeros(R, size(𝐒ⁱ, 1),size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) II = sparse(ℒ.I(length(x)^2)) - lI = -2 * vec(ℒ.I(size(𝐒ⁱ, 2))) + lI = R(-2) * vec(ℒ.I(size(𝐒ⁱ, 2))) 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) @@ -1104,14 +1100,13 @@ function find_shocks(::Val{:LagrangeNewton}, # fXλp = [reshape((2 * 𝐒ⁱ²ᵉ + 6 * 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(ℒ.I(length(x)),x)))' * λ, size(𝐒ⁱ, 2), size(𝐒ⁱ, 2)) - 2*ℒ.I(size(𝐒ⁱ, 2)) (𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) + 3 * 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(x, x)))' # -(𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) + 3 * 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(x, x))) zeros(size(𝐒ⁱ, 1),size(𝐒ⁱ, 1))] - try - f̂xλp = ℒ.factorize(fxλp) - ℒ.ldiv!(Δxλ, f̂xλp, fxλ) - catch - # ℒ.svd(fxλp) - # println("factorization fails") + f̂xλp = ℒ.lu(fxλp, check = false) + + if !ℒ.issuccess(f̂xλp) return x, false end + + ℒ.ldiv!(Δxλ, f̂xλp, fxλ) if !all(isfinite,Δxλ) break end @@ -1139,44 +1134,10 @@ function find_shocks(::Val{:LagrangeNewton}, ℒ.axpby!(1, shock_independent, -1, x̂) if ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol) - # println("LagrangeNewton: $i, Tol: $(ℒ.norm(Δxλ) / ℒ.norm(xλ)) reached, x: $x") break end - - # if i > 500 && ℒ.norm(Δxλ) > 1e-11 && ℒ.norm(Δxλ) > Δnorm - # # println(ℒ.norm(Δxλ)) - # # println(ℒ.norm(x̂) / max(norm1,norm2)) - # # println("LagrangeNewton: $i, Norm increase") - # return x, false - # end - # if i == max_iter - # println("LagrangeNewton: $i, Max iter reached") - # # println(ℒ.norm(Δxλ)) - # end end - # λ = (𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) + 3 * 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), kron_buffer))' \ x * 2 - # println("LagrangeNewton: $(ℒ.norm([(𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) + 3 * 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(x, x)))' * λ - 2 * x - # shock_independent - (𝐒ⁱ * x + 𝐒ⁱ²ᵉ * ℒ.kron(x,x) + 𝐒ⁱ³ᵉ * ℒ.kron(x, ℒ.kron(x, x)))]))") - - # println(ℒ.norm(x)) - # println(x) - # println(λ) - # println([(𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) - 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(x, x)))' * λ - 2 * x - # shock_independent - (𝐒ⁱ * x + 𝐒ⁱ²ᵉ * ℒ.kron(x,x) + 𝐒ⁱ³ᵉ * ℒ.kron(x, ℒ.kron(x, x)))]) - # println(fxλp) - # println(reshape(tmp, size(𝐒ⁱ, 2), size(𝐒ⁱ, 2)) - 2*ℒ.I(size(𝐒ⁱ, 2))) - # println([reshape((2 * 𝐒ⁱ²ᵉ - 2 * 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(ℒ.I(length(x)),x)))' * λ, size(𝐒ⁱ, 2), size(𝐒ⁱ, 2)) - 2*ℒ.I(size(𝐒ⁱ, 2)) (𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) - 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(x, x)))' - # -(𝐒ⁱ + 2 * 𝐒ⁱ²ᵉ * ℒ.kron(ℒ.I(length(x)), x) - 𝐒ⁱ³ᵉ * ℒ.kron(ℒ.I(length(x)), ℒ.kron(x, x))) zeros(size(𝐒ⁱ, 1),size(𝐒ⁱ, 1))]) - # println(fxλp) - # println("Norm: $(ℒ.norm(x̂) / max(norm1,norm2))") - # println(ℒ.norm(Δxλ)) - # println(ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < tol) - - # if !(ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol)) - # println("Find shocks failed. Norm 1: $(ℒ.norm(x̂) / max(norm1,norm2)); Norm 2: $(ℒ.norm(Δxλ) / ℒ.norm(xλ))") - # end - residual = ℒ.norm(x̂) / max(norm1,norm2) step_norm = ℒ.norm(Δxλ) / ℒ.norm(xλ) matched = residual < tol && step_norm < sqrt(tol) @@ -1185,7 +1146,6 @@ function find_shocks(::Val{:LagrangeNewton}, end -end # dispatch_doctor @@ -2005,3 +1965,5 @@ end # dispatch_doctor # # println("Norm: $(ℒ.norm(y - shock_independent) / max(norm1,norm2))") # return x, ℒ.norm(y - shock_independent) / max(norm1,norm2) < tol # end + +end # @stable diff --git a/src/filter/inversion.jl b/src/filter/inversion.jl index d627735ce..747433bcc 100644 --- a/src/filter/inversion.jl +++ b/src/filter/inversion.jl @@ -1,5 +1,6 @@ @stable default_mode = "disable" begin + """ Compute log-likelihood using the inversion filter, which calls the find_shocks function to recover shocks that match the observables. For higher-order solutions the global @@ -8,67 +9,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 = R === Float64 ? workspaces.inversion : Inversion_workspace(R) + 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 @@ -84,56 +59,92 @@ function calculate_inversion_filter_loglikelihood(::Val{:first_order}, end end end - - jacdecomp = ℒ.svd(jac) - x = jacdecomp \ data_in_deviations[:,1] - + # Warmup linear solve: LU instead of SVD so ForwardDiff Duals work. + warmup_rhs = data_in_deviations[:,1] + if size(jac,1) == size(jac,2) + warmup_lu = ℒ.lu(jac, check = false) + if !ℒ.issuccess(warmup_lu) + if opts.verbose println("Inversion filter failed") end + return on_failure_loglikelihood + end + x = warmup_lu \ warmup_rhs + else + JJt_w = jac * jac' + JJt_w_lu = ℒ.lu(JJt_w, check = false) + if !ℒ.issuccess(JJt_w_lu) + if opts.verbose println("Inversion filter failed") end + return on_failure_loglikelihood + end + x = jac' * (JJt_w_lu \ warmup_rhs) + end + 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) - logabsdets += ℒ.logabsdet(jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor)[1] + jac_i = jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor + if size(jac_i,1) == size(jac_i,2) + logabsdets += ℒ.logabsdet(jac_i)[1] else - logabsdets += sum(x -> log(abs(x)), ℒ.svdvals(jac[:,(i - 1) * T.nExo+1:i*T.nExo] ./ precision_factor)) + logabsdets += ℒ.logabsdet(jac_i * jac_i')[1] / 2 end end 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) - jacdecomp = ℒ.lu(jac, check = false) + if T.nExo == length(observables_index) + if R <: AbstractFloat + lu_ws = FastLapackInterface.LUWs(jac) + lu_ws, _, ok, lu_handle = factorize_lu!(Val(:FastLapack), jac, lu_ws, size(jac)) - if !ℒ.issuccess(jacdecomp) - if opts.verbose println("Inversion filter failed") end - return on_failure_loglikelihood - end + if !ok + if opts.verbose println("Inversion filter failed") end + return on_failure_loglikelihood + end - logabsdets = ℒ.logabsdet(jac)[1] - invjac = inv(jacdecomp) - else - jacdecomp = try ℒ.svd(jac) - catch - if opts.verbose println("Inversion filter failed") end - return on_failure_loglikelihood + # logabsdet from U-factor diagonal (jac now holds LU factors in place) + logabsdets = zero(R) + @inbounds for k in 1:size(jac,1) + logabsdets += log(abs(jac[k,k])) + end + invjac = Matrix{R}(ℒ.I, size(jac)) + solve_lu_left!(jac, invjac, lu_ws, lu_handle) + else + jacdecomp = ℒ.lu(jac, check = false) + + if !ℒ.issuccess(jacdecomp) + if opts.verbose println("Inversion filter failed") end + return on_failure_loglikelihood + end + + logabsdets = ℒ.logabsdet(jacdecomp)[1] + invjac = inv(jacdecomp) end - - logabsdets = sum(x -> log(abs(x)), ℒ.svdvals(jac)) - invjac = try ℒ.pinv(jac) - catch + else + # Fat jac (n_obs < n_exo): right pseudo-inverse via normal equations. + # LU is AD-friendly; original SVD/pinv have no ForwardDiff.Dual method. + JJt = jac * jac' + JJt_lu = ℒ.lu(JJt, check = false) + if !ℒ.issuccess(JJt_lu) if opts.verbose println("Inversion filter failed") end return on_failure_loglikelihood end + logabsdets = ℒ.logabsdet(JJt_lu)[1] / 2 + invjac = jac' / JJt_lu end logabsdets *= size(data_in_deviations,2) - presample_periods @@ -155,66 +166,66 @@ 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 = R === Float64 ? workspaces.inversion : Inversion_workspace(R) # @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 + shocks² = zero(R) + logabsdets = zero(R) - 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⁺ - 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 = @ignore_derivatives setdiff(shock_idxs, shock²_idxs) + so = ensure_conditional_forecast_constants!(constants) + shock_idxs = so.shock_idxs + shock²_idxs = so.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = so.var_vol²_idxs + var²_idxs = so.var²_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] @@ -232,42 +243,66 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_second_order}, 𝐒²ᵉ = 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₁ = convert(Vector{R}, state[1][T.past_not_future_and_mixed_idx]) + state₂ = convert(Vector{R}, state[2][T.past_not_future_and_mixed_idx]) - # Use workspaces for model-constant allocations - state¹⁻_vol = ws.state_vol - copyto!(state¹⁻_vol, 1, state₁, 1) - state¹⁻_vol[end] = 1 + n_state_vol = n_past + 1 + n_aug = n_past + 1 + n_exo + n_cond = length(cond_var_idx) - aug_state₁ = ws.aug_state₁ - copyto!(aug_state₁, 1, state₁, 1) - aug_state₁[length(state₁) + 1] = 1 - fill!(view(aug_state₁, length(state₁) + 2:length(aug_state₁)), 1) - - aug_state₂ = ws.aug_state₂ - copyto!(aug_state₂, 1, state₂, 1) - aug_state₂[length(state₂) + 1] = 0 - fill!(view(aug_state₂, length(state₂) + 2:length(aug_state₂)), 0) + if R === Float64 + # Use workspaces for model-constant allocations + state¹⁻_vol = ws.state_vol + copyto!(state¹⁻_vol, 1, state₁, 1) + state¹⁻_vol[end] = 1 - kronaug_state₁ = ws.kronaug_state + aug_state₁ = ws.aug_state₁ + copyto!(aug_state₁, 1, state₁, 1) + aug_state₁[length(state₁) + 1] = 1 + fill!(view(aug_state₁, length(state₁) + 2:length(aug_state₁)), 1) + + aug_state₂ = ws.aug_state₂ + copyto!(aug_state₂, 1, state₂, 1) + aug_state₂[length(state₂) + 1] = 0 + fill!(view(aug_state₂, length(state₂) + 2:length(aug_state₂)), 0) - J = ℒ.I(T.nExo) + kronaug_state₁ = ws.kronaug_state - kron_buffer = ws.kron_buffer - kron_buffer2 = ws.kron_buffer2 - kron_buffer3 = ws.kron_buffer_state - kronstate¹⁻_vol = ws.kronstate_vol + kron_buffer = ws.kron_buffer + kron_buffer2 = ws.kron_buffer2 + kron_buffer3 = ws.kron_buffer_state + kronstate¹⁻_vol = ws.kronstate_vol - shock_independent = zeros(size(data_in_deviations,1)) + shock_independent = ws.shock_independent + fill!(shock_independent, zero(R)) - 𝐒ⁱ = 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, zero(R)) + else + # Allocate R-typed buffers for AD compatibility (e.g. ForwardDiff Dual) + state¹⁻_vol = vcat(state₁, one(R)) + aug_state₁ = vcat(state₁, one(R), ones(R, n_exo)) + aug_state₂ = vcat(state₂, zero(R), zeros(R, n_exo)) + kronaug_state₁ = zeros(R, n_aug^2) + kron_buffer = zeros(R, n_exo^2) + kron_buffer2 = zeros(R, n_exo^2, n_exo) + kron_buffer3 = zeros(R, n_exo * n_state_vol, n_exo) + kronstate¹⁻_vol = zeros(R, n_state_vol^2) + shock_independent = zeros(R, n_cond) + 𝐒ⁱ = Matrix{R}(𝐒¹ᵉ) + jacc = Matrix{R}(𝐒¹ᵉ) + init_guess = zeros(R, n_exo) + end + + J = ℒ.I(T.nExo) + + 𝐒ⁱ²ᵉ = 𝐒²ᵉ / 2 # end # timeit_debug # @timeit_debug timer "Loop" begin @@ -378,10 +413,10 @@ 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)) + logabsdets += ℒ.logabsdet(jacc * jacc')[1] / 2 end shocks² += sum(abs2,x) @@ -409,24 +444,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 = R === Float64 ? workspaces.inversion : Inversion_workspace(R) # @timeit_debug timer "2nd - Inversion filter" begin # @timeit_debug timer "Preallocation" begin @@ -434,31 +472,24 @@ 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 + 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 - shocks² = 0.0 - logabsdets = 0.0 + shocks² = zero(R) + logabsdets = zero(R) - # 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⁺ - - 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 + so = ensure_conditional_forecast_constants!(constants) + shock_idxs = cc.shock_idxs + shock²_idxs = cc.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = cc.var_vol²_idxs # tmp = ℒ.kron(s_in_s⁺, s_in_s⁺) |> sparse # var²_idxs = tmp.nzind @@ -480,7 +511,7 @@ function calculate_inversion_filter_loglikelihood(::Val{:second_order}, 𝐒²ᵉ = nnz(𝐒²ᵉ) / length(𝐒²ᵉ) > .1 ? collect(𝐒²ᵉ) : 𝐒²ᵉ 𝐒⁻² = nnz(𝐒⁻²) / length(𝐒⁻²) > .1 ? collect(𝐒⁻²) : 𝐒⁻² - state = state[T.past_not_future_and_mixed_idx] + state = convert(Vector{R}, state[T.past_not_future_and_mixed_idx]) # Use workspaces for model-constant allocations state¹⁻_vol = ws.state_vol @@ -501,17 +532,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, zero(R)) 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, zero(R)) # end # timeit_debug # @timeit_debug timer "Loop" begin @@ -615,10 +651,10 @@ 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 + logabsdets += ℒ.logabsdet(jacc * jacc')[1] / 2 # ./ precision_factor end shocks² += sum(abs2,x) @@ -647,53 +683,54 @@ 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 = R === Float64 ? workspaces.inversion : Inversion_workspace(R) # @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 + precision_factor = one(R) 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 + shocks² = zero(R) + logabsdets = zero(R) - 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⁺ + so = ensure_conditional_forecast_constants!(constants; third_order = true) shockvar_idxs = cc.shockvar_idxs - shock_idxs = cc.shock_idxs - shock_idxs2 = cc.shock_idxs2 shock²_idxs = cc.shock²_idxs - shockvar²_idxs = setdiff(union(shock_idxs), shock²_idxs) + shockvar²_idxs = so.shockvar²_idxs var_vol²_idxs = cc.var_vol²_idxs - - tmp = ℒ.kron(s_in_s⁺, s_in_s⁺) |> sparse - var²_idxs = tmp.nzind + var²_idxs = so.var²_idxs + to = constants.third_order 𝐒⁻¹ = 𝐒[1][T.past_not_future_and_mixed_idx,:] 𝐒¹⁻ = 𝐒[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] @@ -714,30 +751,10 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, 𝐒²ᵉ = 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) + var_vol³_idxs = to.var_vol³_idxs + shock³_idxs = to.shock³_idxs + shockvar³2_idxs = to.shockvar³2_idxs + shockvar³_idxs = to.shockvar³_idxs 𝐒³⁻ᵛ = 𝐒[3][cond_var_idx,var_vol³_idxs] 𝐒³⁻ᵉ² = 𝐒[3][cond_var_idx,shockvar³2_idxs] |> collect @@ -750,65 +767,66 @@ function calculate_inversion_filter_loglikelihood(::Val{:pruned_third_order}, 𝐒³ᵉ = nnz(𝐒³ᵉ) / length(𝐒³ᵉ) > .1 ? collect(𝐒³ᵉ) : 𝐒³ᵉ 𝐒⁻³ = nnz(𝐒⁻³) / length(𝐒⁻³) > .1 ? collect(𝐒⁻³) : 𝐒⁻³ - state[1] = state[1][T.past_not_future_and_mixed_idx] - 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) + # Shadow the input `state` with R-typed local copies so the kernel can + # be driven by ForwardDiff Duals (the input may be Vector{Vector{Float64}}). + state = Vector{R}[ + convert(Vector{R}, state[1][T.past_not_future_and_mixed_idx]), + convert(Vector{R}, state[2][T.past_not_future_and_mixed_idx]), + convert(Vector{R}, state[3][T.past_not_future_and_mixed_idx]), + ] - 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 +834,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,10 +1043,10 @@ 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)) + logabsdets += ℒ.logabsdet(jacc * jacc')[1] / 2 end shocks² += sum(abs2,x) @@ -1040,10 +1056,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 +1106,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 = R === Float64 ? workspaces.inversion : Inversion_workspace(R) # @timeit_debug timer "3rd - Inversion filter" begin # @timeit_debug timer "Preallocation" begin @@ -1100,37 +1134,24 @@ 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 + 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 - shocks² = 0.0 - logabsdets = 0.0 + shocks² = zero(R) + logabsdets = zero(R) 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(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 + so = ensure_conditional_forecast_constants!(constants; third_order = true) + shock²_idxs = cc.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = cc.var_vol²_idxs + var²_idxs = so.var²_idxs + to = constants.third_order 𝐒⁻¹ = 𝐒[1][T.past_not_future_and_mixed_idx,:] 𝐒¹⁻ = 𝐒[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] @@ -1149,32 +1170,12 @@ function calculate_inversion_filter_loglikelihood(::Val{:third_order}, 𝐒²ᵉ = nnz(𝐒²ᵉ) / length(𝐒²ᵉ) > .1 ? collect(𝐒²ᵉ) : 𝐒²ᵉ 𝐒⁻² = nnz(𝐒⁻²) / length(𝐒⁻²) > .1 ? collect(𝐒⁻²) : 𝐒⁻² - state = state[T.past_not_future_and_mixed_idx] - - 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 + state = convert(Vector{R}, state[T.past_not_future_and_mixed_idx]) - 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) + var_vol³_idxs = to.var_vol³_idxs + shock³_idxs = to.shock³_idxs + shockvar³2_idxs = to.shockvar³2_idxs + shockvar³_idxs = to.shockvar³_idxs 𝐒³⁻ᵛ = 𝐒[3][cond_var_idx,var_vol³_idxs] 𝐒³⁻ᵉ² = 𝐒[3][cond_var_idx,shockvar³2_idxs] @@ -1202,31 +1203,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,14 +1392,21 @@ 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)) + logabsdets += ℒ.logabsdet(jacc * jacc')[1] / 2 end shocks² += sum(abs2,x) @@ -1391,23 +1416,27 @@ 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(𝓂::ℳ, +@unstable function filter_data_with_model(𝓂::ℳ, data_in_deviations::KeyedArray{Float64}, ::Val{:first_order}, # algo ::Val{:inversion}; # filter @@ -1425,7 +1454,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 +1463,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 @@ -1498,14 +1523,27 @@ function filter_data_with_model(𝓂::ℳ, jac = 𝐒₁[cond_var_idx, end-T.nExo+1:end] if T.nExo == length(observables) - jacdecomp = ℒ.lu(jac, check = false) + if eltype(jac) <: AbstractFloat + lu_ws = FastLapackInterface.LUWs(jac) + lu_ws, _, ok, lu_handle = factorize_lu!(Val(:FastLapack), jac, lu_ws, size(jac)) - if !ℒ.issuccess(jacdecomp) - @error "Inversion filter failed" - return variables, shocks, zeros(0,0), decomposition - end + if !ok + @error "Inversion filter failed" + return variables, shocks, zeros(0,0), decomposition + end + + invjac = Matrix{eltype(jac)}(ℒ.I, size(jac)) + solve_lu_left!(jac, invjac, lu_ws, lu_handle) + else + jacdecomp = ℒ.lu(jac, check = false) - invjac = inv(jacdecomp) + if !ℒ.issuccess(jacdecomp) + @error "Inversion filter failed" + return variables, shocks, zeros(0,0), decomposition + end + + invjac = inv(jacdecomp) + end else # jacdecomp = ℒ.svd(jac) @@ -1549,7 +1587,7 @@ function filter_data_with_model(𝓂::ℳ, end -function filter_data_with_model(𝓂::ℳ, +@unstable function filter_data_with_model(𝓂::ℳ, data_in_deviations::KeyedArray{Float64}, ::Val{:second_order}, # algo ::Val{:inversion}; # filter @@ -1565,14 +1603,15 @@ 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 - ms = ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ms = constants.post_complete_parameters all_SS = expand_steady_state(SS_and_pars, ms) full_state = collect(sss) - all_SS @@ -1585,20 +1624,12 @@ 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⁺ - - 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 + computational_constants = ensure_computational_constants!(𝓂.constants) + so = ensure_conditional_forecast_constants!(𝓂.constants) + # s_in_s⁺ = computational_constants.s_in_s + shock²_idxs = computational_constants.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = computational_constants.var_vol²_idxs # tmp = ℒ.kron(s_in_s⁺, s_in_s⁺) |> sparse # var²_idxs = tmp.nzind @@ -1767,7 +1798,7 @@ function filter_data_with_model(𝓂::ℳ, end -function filter_data_with_model(𝓂::ℳ, +@unstable function filter_data_with_model(𝓂::ℳ, data_in_deviations::KeyedArray{Float64}, ::Val{:pruned_second_order}, # algo ::Val{:inversion}; # filter @@ -1778,7 +1809,8 @@ function filter_data_with_model(𝓂::ℳ, # Initialize constants at entry point constants = initialise_constants!(𝓂) T = constants.post_model_macro - ms = ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ms = constants.post_complete_parameters variables = zeros(T.nVars, size(data_in_deviations,2)) shocks = zeros(T.nExo, size(data_in_deviations,2)) @@ -1786,9 +1818,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,23 +1837,14 @@ function filter_data_with_model(𝓂::ℳ, cond_var_idx = indexin(observables,sort(union(T.aux,T.var,T.exo_present))) - 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⁺ - e_in_s⁺ = BitVector(vcat(zeros(Bool, T.nPast_not_future_and_mixed + 1), ones(Bool, T.nExo))) + computational_constants = ensure_computational_constants!(𝓂.constants) + so = ensure_conditional_forecast_constants!(𝓂.constants) + sv_in_s⁺ = computational_constants.s_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 + shock²_idxs = computational_constants.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = computational_constants.var_vol²_idxs + var²_idxs = so.var²_idxs 𝐒⁻¹ = 𝐒[1][T.past_not_future_and_mixed_idx, :] 𝐒¹⁻ = 𝐒[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] @@ -2038,7 +2061,7 @@ function filter_data_with_model(𝓂::ℳ, return variables, shocks, zeros(0,0), decomposition end -function filter_data_with_model(𝓂::ℳ, +@unstable function filter_data_with_model(𝓂::ℳ, data_in_deviations::KeyedArray{Float64}, ::Val{:third_order}, # algo ::Val{:inversion}; # filter @@ -2049,16 +2072,17 @@ function filter_data_with_model(𝓂::ℳ, # Initialize constants at entry point constants = initialise_constants!(𝓂) T = constants.post_model_macro - ms = ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ms = constants.post_complete_parameters variables = zeros(T.nVars, size(data_in_deviations,2)) shocks = zeros(T.nExo, size(data_in_deviations,2)) 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,26 +2100,13 @@ 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⁺ - - 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 + computational_constants = ensure_computational_constants!(𝓂.constants) + so = ensure_conditional_forecast_constants!(𝓂.constants; third_order = true) + shock²_idxs = computational_constants.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = computational_constants.var_vol²_idxs + var²_idxs = so.var²_idxs + to = 𝓂.constants.third_order 𝐒⁻¹ = 𝐒[1][T.past_not_future_and_mixed_idx,:] 𝐒¹⁻ = 𝐒[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] @@ -2116,30 +2127,10 @@ function filter_data_with_model(𝓂::ℳ, state = state[T.past_not_future_and_mixed_idx] - 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) + var_vol³_idxs = to.var_vol³_idxs + shock³_idxs = to.shock³_idxs + shockvar³2_idxs = to.shockvar³2_idxs + shockvar³_idxs = to.shockvar³_idxs 𝐒³⁻ᵛ = 𝐒[3][cond_var_idx,var_vol³_idxs] 𝐒³⁻ᵉ² = 𝐒[3][cond_var_idx,shockvar³2_idxs] @@ -2351,7 +2342,7 @@ function filter_data_with_model(𝓂::ℳ, end -function filter_data_with_model(𝓂::ℳ, +@unstable function filter_data_with_model(𝓂::ℳ, data_in_deviations::KeyedArray{Float64}, ::Val{:pruned_third_order}, # algo ::Val{:inversion}; # filter @@ -2362,7 +2353,8 @@ function filter_data_with_model(𝓂::ℳ, # Initialize constants at entry point constants = initialise_constants!(𝓂) T = constants.post_model_macro - ms = ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ms = constants.post_complete_parameters variables = zeros(T.nVars, size(data_in_deviations,2)) shocks = zeros(T.nExo, size(data_in_deviations,2)) @@ -2370,9 +2362,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,29 +2381,19 @@ 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) + so = ensure_conditional_forecast_constants!(𝓂.constants; third_order = true) + s_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 + shockvar_idxs = so.shockvar_no_vol_idxs - 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) + shock²_idxs = computational_constants.shock²_idxs + shockvar²_idxs = so.shockvar²_idxs + var_vol²_idxs = computational_constants.var_vol²_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 + var²_idxs = so.var²_idxs + to = 𝓂.constants.third_order 𝐒⁻¹ = 𝐒[1][T.past_not_future_and_mixed_idx,:] 𝐒¹⁻ = 𝐒[1][cond_var_idx, 1:T.nPast_not_future_and_mixed] @@ -2432,30 +2414,10 @@ function filter_data_with_model(𝓂::ℳ, 𝐒²ᵉ = 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) + var_vol³_idxs = to.var_vol³_idxs + shock³_idxs = to.shock³_idxs + shockvar³2_idxs = to.shockvar³2_idxs + shockvar³_idxs = to.shockvar³_idxs 𝐒³⁻ᵛ = 𝐒[3][cond_var_idx,var_vol³_idxs] 𝐒³⁻ᵉ² = 𝐒[3][cond_var_idx,shockvar³2_idxs] @@ -2764,4 +2726,5 @@ function filter_data_with_model(𝓂::ℳ, return variables, shocks, zeros(0,0), decomposition end -end # dispatch_doctor + +end # @stable diff --git a/src/filter/kalman.jl b/src/filter/kalman.jl index a50f1e6dd..bf5e42b76 100644 --- a/src/filter/kalman.jl +++ b/src/filter/kalman.jl @@ -1,93 +1,39 @@ @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}, +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(), - 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}, - 𝐒::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 +52,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 +66,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 @@ -156,77 +96,81 @@ function run_kalman_iterations(A::Matrix{S}, # Initialize state estimate to zero fill!(u, zero(S)) - ℒ.mul!(z, C, u) + ℒ.mul!(z, C, u) # z = C * u loglik = S(0.0) # @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) - # v = data_in_deviations[:, t] - z + ℒ.axpby!(1, @view(data_in_deviations[:, t]), -1, z) # z = data[:,t] - z (innovation v) - ℒ.mul!(Ctmp, C, P) # use Octavian.jl - ℒ.mul!(F, Ctmp, C') - # F = C * P * C' + ℒ.mul!(Ctmp, C, P) # Ctmp = C * P + ℒ.mul!(F, Ctmp, C') # 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) - # end # timeit_debug + # Old way (≤v0.1.42): luF = lu(F) — allocates new LU each step + ws.fast_lu_ws_f, ws.fast_lu_dims_f, solved_F, luF = factorize_lu!(Val(:FastLapack), F, + ws.fast_lu_ws_f, + ws.fast_lu_dims_f) - if !ℒ.issuccess(luF) + if !solved_F if verbose println("KF factorisation failed step $t") end return on_failure_loglikelihood end - Fdet = ℒ.det(luF) + # Old way (≤v0.1.42): Fdet = det(luF); loglik += log(Fdet) + v' * inv(F) * v + # Current code computes log|det(F)| from the LU diagonal and pivot signs. + 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 - # invF = inv(luF) ### - - # @timeit_debug timer "LU div" begin + # Old way (≤v0.1.42): loglik += log(det(F)) + v' * inv(F) * v if t > presample_periods - ℒ.ldiv!(ztmp, luF, z) - loglik += log(Fdet) + ℒ.dot(z', ztmp) ### - # loglik += log(Fdet) + z' * invF * z### - # loglik += log(Fdet) + v' * invF * v### + copyto!(ztmp, z) + solve_lu_left!(F, ztmp, ws.fast_lu_ws_f, luF) # ztmp = F \ z + loglik += logabsdetF + ℒ.dot(z', ztmp) # loglik += log|det(F)| + z' * (F \ z) end - # ℒ.mul!(Ktmp, P, C') - # ℒ.mul!(K, Ktmp, invF) - ℒ.mul!(K, P, C') - ℒ.rdiv!(K, luF) - # K = P * Ct / luF - # K = P * C' * invF + # Old way (≤v0.1.42): K = P * C' / F — Kalman gain + ℒ.mul!(K, P, C') # K = P * C' + solve_lu_right!(F, K, ws.fast_lu_ws_f, luF, ws.fast_lu_rhs_t_k) # K = K / F # end # timeit_debug # @timeit_debug timer "Matmul" begin - ℒ.mul!(tmp, K, C) - ℒ.mul!(Ptmp, tmp, P) - ℒ.axpy!(-1, Ptmp, P) + # P = A * (P - K * C * P) * A' + B + ℒ.mul!(tmp, K, C) # tmp = K * C + ℒ.mul!(Ptmp, tmp, P) # Ptmp = K * C * P + ℒ.axpy!(-1, Ptmp, P) # P = P - K * C * P - ℒ.mul!(Ptmp, A, P) - ℒ.mul!(P, Ptmp, A') - ℒ.axpy!(1, 𝐁, P) - # P = A * (P - K * C * P) * A' + 𝐁 + ℒ.mul!(Ptmp, A, P) # Ptmp = A * P + ℒ.mul!(P, Ptmp, A') # P = A * P * A' + ℒ.axpy!(1, 𝐁, P) # P = P + B - ℒ.mul!(u, K, z, 1, 1) - ℒ.mul!(utmp, A, u) - u .= utmp # u = A * (u + K * v) + ℒ.mul!(u, K, z, 1, 1) # u = u + K * v + ℒ.mul!(utmp, A, u) # utmp = A * u + u .= utmp # u = A * (u + K * v) - ℒ.mul!(z, C, u) - # z = C * u + ℒ.mul!(z, C, u) # z = C * u # end # timeit_debug end @@ -238,7 +182,7 @@ function run_kalman_iterations(A::Matrix{S}, end -function filter_data_with_model(𝓂::ℳ, +@unstable function filter_data_with_model(𝓂::ℳ, data_in_deviations::KeyedArray{Float64}, ::Val{:first_order}, # algo ::Val{:kalman}; # filter, @@ -284,20 +228,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 +246,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' @@ -314,34 +254,48 @@ function filter_and_smooth(𝓂::ℳ, n_obs = size(data_in_deviations,2) - v = zeros(size(C,1), n_obs) - μ = zeros(size(A,1), n_obs+1) # filtered_states - P = zeros(size(A,1), size(A,1), n_obs+1) # filtered_covariances - σ = zeros(size(A,1), n_obs) # filtered_standard_deviations - iF= zeros(size(C,1), size(C,1), n_obs) - L = zeros(size(A,1), size(A,1), n_obs) + n_obs_C = size(C,1) + n_states = size(A,1) + kalman_ws = ensure_kalman_workspaces!(𝓂.workspaces, n_obs_C, n_states) + + v = zeros(n_obs_C, n_obs) + μ = zeros(n_states, n_obs+1) # filtered_states + P = zeros(n_states, n_states, n_obs+1) # filtered_covariances + σ = zeros(n_states, n_obs) # filtered_standard_deviations + iF= zeros(n_obs_C, n_obs_C, n_obs) + L = zeros(n_states, n_states, n_obs) ϵ = zeros(size(B,2), n_obs) # filtered_shocks P[:, :, 1] = P̄ + F_buf = kalman_ws.F + # Kalman Filter for t in axes(data_in_deviations,2) v[:, t] .= data_in_deviations[:, t] - C * μ[:, t] - F̄ = ℒ.lu(C * P[:, :, t] * C', check = false) + @views F_buf .= C * P[:, :, t] * C' + @views iF_t = iF[:, :, t] + fill!(iF_t, 0.0) + @inbounds for i in 1:n_obs_C + iF_t[i, i] = 1.0 + end - if !ℒ.issuccess(F̄) + kalman_ws.fast_lu_ws_f, kalman_ws.fast_lu_dims_f, solved_F, _ = + factorize_lu!(Val(:FastLapack), F_buf, kalman_ws.fast_lu_ws_f, kalman_ws.fast_lu_dims_f) + + if !solved_F @warn "Kalman filter stopped in period $t due to numerical stabiltiy issues." break end - iF[:, :, t] .= inv(F̄) - PCiF = P[:, :, t] * C' * iF[:, :, t] + solve_lu_left!(F_buf, iF_t, kalman_ws.fast_lu_ws_f, nothing) # iF_t = F̄ \ I + PCiF = P[:, :, t] * C' * iF_t L[:, :, t] .= A - A * PCiF * C P[:, :, t+1].= A * P[:, :, t] * L[:, :, t]' + 𝐁 σ[:, t] .= sqrt.(abs.(ℒ.diag(P[:, :, t+1]))) # small numerical errors in this computation μ[:, t+1] .= A * (μ[:, t] + PCiF * v[:, t]) - ϵ[:, t] .= B' * C' * iF[:, :, t] * v[:, t] + ϵ[:, t] .= B' * C' * iF_t * v[:, t] end @@ -390,4 +344,5 @@ function filter_and_smooth(𝓂::ℳ, return μ̄, σ̄, ϵ̄, smooth_decomposition, μ[:, 2:end], σ, ϵ, filter_decomposition end -end # dispatch_doctor + +end # @stable diff --git a/src/get_functions.jl b/src/get_functions.jl index d8fc003cc..d6e2384c7 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -1,3 +1,6 @@ +@stable default_mode = "disable" begin + + """ $(SIGNATURES) Return the shock decomposition in absolute deviations from the relevant steady state. The non-stochastic steady state (NSSS) is relevant for first order solutions and the stochastic steady state for higher order solutions. The deviations are based on the Kalman smoother or filter (depending on the `smooth` keyword argument) or inversion filter using the provided data and solution of the model. When the defaults are used, the filter is selected automatically—Kalman for first order solutions and inversion otherwise—and smoothing is only enabled when the Kalman filter is active. Data is by default assumed to be in levels unless `data_in_levels` is set to `false`. @@ -36,7 +39,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -44,6 +47,8 @@ end β = 0.95 end +import Random; Random.seed!(3) + simulation = simulate(RBC) get_shock_decomposition(RBC,simulation([:c],:,:simulate)) @@ -55,28 +60,28 @@ get_shock_decomposition(RBC,simulation([:c],:,:simulate)) And data, 4×2×40 Array{Float64, 3}: [showing 3 of 40 slices] [:, :, 1] ~ (:, :, 1): - (:eps_z₍ₓ₎) (:Initial_values) - (:c) 0.000407252 -0.00104779 - (:k) 0.00374808 -0.0104645 - (:q) 0.00415533 -0.000807161 - (:z) 0.000603617 -1.99957e-6 + (:eps_z₍ₓ₎) (:Initial_values) + (:c) 0.00128797 0.00319151 + (:k) 0.0118536 0.0318 + (:q) 0.0131415 0.00335202 + (:z) 0.00190898 0.000146294 [:, :, 21] ~ (:, :, 21): (:eps_z₍ₓ₎) (:Initial_values) - (:c) 0.026511 -0.000433619 - (:k) 0.25684 -0.00433108 - (:q) 0.115858 -0.000328764 - (:z) 0.0150266 0.0 + (:c) -0.0428897 0.00132724 + (:k) -0.425096 0.0132567 + (:q) -0.0721742 0.00100629 + (:z) -0.00622294 1.73472e-18 [:, :, 40] ~ (:, :, 40): - (:eps_z₍ₓ₎) (:Initial_values) - (:c) 0.0437976 -0.000187505 - (:k) 0.4394 -0.00187284 - (:q) 0.00985518 -0.000142164 - (:z) -0.00366442 8.67362e-19 + (:eps_z₍ₓ₎) (:Initial_values) + (:c) -0.0116806 0.000573923 + (:k) -0.116386 0.00573246 + (:q) -0.012256 0.00043514 + (:z) -0.000533531 1.0842e-19 ``` """ -function get_shock_decomposition(𝓂::ℳ, +@unstable function get_shock_decomposition(𝓂::ℳ, data::KeyedArray{Float64}; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, @@ -87,11 +92,17 @@ function get_shock_decomposition(𝓂::ℳ, smooth::Bool = DEFAULT_SMOOTH_SELECTOR(filter), verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM)::KeyedArray + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -143,6 +154,8 @@ function get_shock_decomposition(𝓂::ℳ, decomposition[:,end - 2,:] .-= SSS_delta * (size(decomposition,2) - 4) end + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return KeyedArray(decomposition[:,1:end-1,:]; Variables = axis1, Shocks = axis2, Periods = 1:size(data,2)) end @@ -185,7 +198,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -193,6 +206,8 @@ end β = 0.95 end +import Random; Random.seed!(3) + simulation = simulate(RBC) get_estimated_shocks(RBC,simulation([:c],:,:simulate)) @@ -201,11 +216,11 @@ get_estimated_shocks(RBC,simulation([:c],:,:simulate)) ↓ Shocks ∈ 1-element Vector{Symbol} → Periods ∈ 40-element UnitRange{Int64} And data, 1×40 Matrix{Float64}: - (1) (2) (3) (4) … (37) (38) (39) (40) - (:eps_z₍ₓ₎) 0.0603617 0.614652 -0.519048 0.711454 -0.873774 1.27918 -0.929701 -0.2255 + (1) (2) … (39) (40) + (:eps_z₍ₓ₎) 0.190898 1.24786 -0.676457 -0.00870749 ``` """ -function get_estimated_shocks(𝓂::ℳ, +@unstable function get_estimated_shocks(𝓂::ℳ, data::KeyedArray{Float64}; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, @@ -216,11 +231,17 @@ function get_estimated_shocks(𝓂::ℳ, smooth::Bool = DEFAULT_SMOOTH_SELECTOR(filter), verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM)::KeyedArray + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -260,6 +281,8 @@ function get_estimated_shocks(𝓂::ℳ, ensure_name_display_constants!(𝓂) axis1 = 𝓂.constants.post_complete_parameters.exo_axis_with_subscript + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return KeyedArray(shocks; Shocks = axis1, Periods = 1:size(data,2)) end @@ -305,7 +328,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -313,6 +336,8 @@ end β = 0.95 end +import Random; Random.seed!(3) + simulation = simulate(RBC) get_estimated_variables(RBC,simulation([:c],:,:simulate)) @@ -321,14 +346,14 @@ get_estimated_variables(RBC,simulation([:c],:,:simulate)) ↓ Variables ∈ 4-element Vector{Symbol} → Periods ∈ 40-element UnitRange{Int64} And data, 4×40 Matrix{Float64}: - (1) (2) (3) (4) … (37) (38) (39) (40) - (:c) 5.92901 5.92797 5.92847 5.92048 5.95845 5.95697 5.95686 5.96173 - (:k) 47.3185 47.3087 47.3125 47.2392 47.6034 47.5969 47.5954 47.6402 - (:q) 6.87159 6.86452 6.87844 6.79352 7.00476 6.9026 6.90727 6.95841 - (:z) -0.00109471 -0.00208056 4.43613e-5 -0.0123318 0.0162992 0.000445065 0.00119089 0.00863586 + (1) (2) … (39) (40) + (:c) 5.94073 5.94913 5.9249 5.92515 + (:k) 47.4339 47.5121 47.2781 47.2796 + (:q) 6.90055 6.97596 6.86123 6.87224 + (:z) 0.00205528 0.0128896 -0.00223228 -0.000533531 ``` """ -function get_estimated_variables(𝓂::ℳ, +@unstable function get_estimated_variables(𝓂::ℳ, data::KeyedArray{Float64}; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, @@ -340,11 +365,17 @@ function get_estimated_variables(𝓂::ℳ, smooth::Bool = DEFAULT_SMOOTH_SELECTOR(filter), verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM)::KeyedArray + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -384,7 +415,11 @@ function get_estimated_variables(𝓂::ℳ, ensure_name_display_constants!(𝓂) axis1 = 𝓂.constants.post_complete_parameters.var_axis - return KeyedArray(levels ? variables .+ NSSS[1:length(𝓂.constants.post_model_macro.var)] : variables; Variables = axis1, Periods = 1:size(data,2)) + result = KeyedArray(levels ? variables .+ NSSS[1:length(𝓂.constants.post_model_macro.var)] : variables; Variables = axis1, Periods = 1:size(data,2)) + + if !use_workspaces; 𝓂.workspaces = orig_ws; end + + return result end @@ -429,7 +464,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -437,6 +472,8 @@ end β = 0.95 end +import Random; Random.seed!(3) + simulation = simulate(RBC) get_model_estimates(RBC,simulation([:c],:,:simulate)) @@ -445,15 +482,15 @@ get_model_estimates(RBC,simulation([:c],:,:simulate)) ↓ Variables_and_shocks ∈ 5-element Vector{Symbol} → Periods ∈ 40-element UnitRange{Int64} And data, 5×40 Matrix{Float64}: - (1) (2) (3) (4) … (37) (38) (39) (40) - (:c) 5.94335 5.94676 5.94474 5.95135 5.93773 5.94333 5.94915 5.95473 - (:k) 47.4603 47.4922 47.476 47.5356 47.4079 47.4567 47.514 47.5696 - (:q) 6.89873 6.92782 6.87844 6.96043 6.85055 6.9403 6.95556 6.96064 - (:z) 0.0014586 0.00561728 -0.00189203 0.0101896 -0.00543334 0.00798437 0.00968602 0.00981981 - (:eps_z₍ₓ₎) 0.12649 0.532556 -0.301549 1.0568 … -0.746981 0.907104 0.808914 0.788261 + (1) (2) … (39) (40) + (:c) 5.94073 5.94913 5.9249 5.92515 + (:k) 47.4339 47.5121 47.2781 47.2796 + (:q) 6.90055 6.97596 6.86123 6.87224 + (:z) 0.00205528 0.0128896 -0.00223228 -0.000533531 + (:eps_z₍ₓ₎) 0.190898 1.24786 … -0.676457 -0.00870749 ``` """ -function get_model_estimates(𝓂::ℳ, +@unstable function get_model_estimates(𝓂::ℳ, data::KeyedArray{Float64}; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, @@ -465,9 +502,11 @@ function get_model_estimates(𝓂::ℳ, smooth::Bool = DEFAULT_SMOOTH_SELECTOR(filter), verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM)::KeyedArray + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray vars = get_estimated_variables(𝓂, data; parameters = parameters, @@ -482,7 +521,9 @@ function get_model_estimates(𝓂::ℳ, tol = tol, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm = sylvester_algorithm, - lyapunov_algorithm = lyapunov_algorithm) + lyapunov_algorithm = lyapunov_algorithm, + caching = caching, + use_workspaces = use_workspaces) shks = get_estimated_shocks(𝓂, data; parameters = parameters, @@ -496,7 +537,9 @@ function get_model_estimates(𝓂::ℳ, tol = tol, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm = sylvester_algorithm, - lyapunov_algorithm = lyapunov_algorithm) + lyapunov_algorithm = lyapunov_algorithm, + caching = caching, + use_workspaces = use_workspaces) # Build unified first axis and concatenate data est_labels = vcat(collect(axiskeys(vars, 1)), collect(axiskeys(shks, 1))) @@ -540,7 +583,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -548,6 +591,8 @@ end β = 0.95 end +import Random; Random.seed!(3) + simulation = simulate(RBC) get_estimated_variable_standard_deviations(RBC,simulation([:c],:,:simulate)) @@ -556,14 +601,14 @@ get_estimated_variable_standard_deviations(RBC,simulation([:c],:,:simulate)) ↓ Standard_deviations ∈ 4-element Vector{Symbol} → Periods ∈ 40-element UnitRange{Int64} And data, 4×40 Matrix{Float64}: - (1) (2) (3) (4) … (38) (39) (40) - (:c) 1.23202e-9 1.84069e-10 8.23181e-11 8.23181e-11 8.23181e-11 8.23181e-11 0.0 - (:k) 0.00509299 0.000382934 2.87922e-5 2.16484e-6 1.6131e-9 9.31323e-10 1.47255e-9 - (:q) 0.0612887 0.0046082 0.000346483 2.60515e-5 1.31709e-9 1.31709e-9 9.31323e-10 - (:z) 0.00961766 0.000723136 5.43714e-5 4.0881e-6 3.08006e-10 3.29272e-10 2.32831e-10 + (1) (2) … (39) (40) + (:c) 1.31709e-9 1.16415e-10 8.23181e-11 0.0 + (:k) 0.00509299 0.000382934 9.31323e-10 1.6131e-9 + (:q) 0.0612887 0.0046082 9.31323e-10 9.31323e-10 + (:z) 0.00961766 0.000723136 0.0 1.64636e-10 ``` """ -function get_estimated_variable_standard_deviations(𝓂::ℳ, +@unstable function get_estimated_variable_standard_deviations(𝓂::ℳ, data::KeyedArray{Float64}; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, @@ -571,10 +616,16 @@ function get_estimated_variable_standard_deviations(𝓂::ℳ, smooth::Bool = DEFAULT_SMOOTH_FLAG, verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM) + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, lyapunov_algorithm = lyapunov_algorithm) @@ -610,6 +661,8 @@ function get_estimated_variable_standard_deviations(𝓂::ℳ, ensure_name_display_constants!(𝓂) axis1 = 𝓂.constants.post_complete_parameters.var_axis + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return KeyedArray(standard_deviations; Standard_deviations = axis1, Periods = 1:size(data,2)) end @@ -660,7 +713,7 @@ using SparseArrays, AxisKeys A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] end -@parameters RBC_CME begin +@parameters RBC_CME silent = true begin alpha = .157 beta = .999 delta = .0226 @@ -687,18 +740,20 @@ get_conditional_forecast(RBC_CME, conditions, shocks = shocks, conditions_in_lev ↓ Variables_and_shocks ∈ 9-element Vector{Symbol} → Periods ∈ 42-element UnitRange{Int64} And data, 9×42 Matrix{Float64}: - (1) (2) … (41) (42) - (:A) 0.0313639 0.0134792 0.000221372 0.000199235 - (:Pi) 0.000780257 0.00020929 -0.000146071 -0.000140137 - (:R) 0.00117156 0.00031425 -0.000219325 -0.000210417 - (:c) 0.01 0.00600605 0.00213278 0.00203751 - (:k) 0.034584 0.0477482 … 0.0397631 0.0380482 - (:y) 0.0446375 0.02 0.00129544 0.001222 - (:z_delta) 0.00025 0.000225 3.69522e-6 3.3257e-6 - (:delta_eps) 0.05 0.0 0.0 0.0 - (:eps_z) 4.61234 -2.16887 0.0 0.0 - -# The same can be achieved with the other input formats: + (1) … (41) (42) + (:A) 0.0313639 0.000221372 0.000199235 + (:Pi) 0.000780257 -0.000146071 -0.000140137 + (:R) 0.00117156 -0.000219325 -0.000210417 + (:c) 0.01 0.00213278 0.00203751 + (:k) 0.034584 … 0.0397631 0.0380482 + (:y) 0.0446375 0.00129544 0.001222 + (:z_delta) 0.00025 3.69522e-6 3.3257e-6 + (:delta_eps₍ₓ₎) 0.05 0.0 0.0 + (:eps_z₍ₓ₎) 4.61234 0.0 0.0 +``` + +The same can be achieved with the other input formats: +```julia # conditions = Matrix{Union{Nothing,Float64}}(undef,7,2) # conditions[4,1] = .01 # conditions[6,2] = .02 @@ -716,7 +771,7 @@ And data, 9×42 Matrix{Float64}: # shocks[1,1] = .05 ``` """ -function get_conditional_forecast(𝓂::ℳ, +@unstable function get_conditional_forecast(𝓂::ℳ, conditions::Union{Matrix{Union{Nothing,Float64}}, SparseMatrixCSC{Float64}, KeyedArray{Union{Nothing,Float64}}, KeyedArray{Float64}}; shocks::Union{Matrix{Union{Nothing,Float64}}, SparseMatrixCSC{Float64}, KeyedArray{Union{Nothing,Float64}}, KeyedArray{Float64}, Nothing} = nothing, initial_state::Union{Vector{Vector{Float64}},Vector{Float64}} = DEFAULT_INITIAL_STATE, @@ -729,12 +784,18 @@ function get_conditional_forecast(𝓂::ℳ, levels::Bool = false, verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, - conditional_forecast_solver::Symbol = :LagrangeNewton) + conditional_forecast_solver::Symbol = :LagrangeNewton, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -856,7 +917,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₁[:,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 +929,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 +939,7 @@ function get_conditional_forecast(𝓂::ℳ, cond_var_idx, free_shock_idx, state_update, - S₁, + Ŝ₁, S₂, S₃, 𝓂.constants, @@ -920,7 +981,7 @@ function get_conditional_forecast(𝓂::ℳ, cond_var_idx, free_shock_idx, state_update, - S₁, + Ŝ₁, S₂, S₃, 𝓂.constants, @@ -937,22 +998,29 @@ function get_conditional_forecast(𝓂::ℳ, Y[:,i] = pruning ? sum(initial_state) : initial_state end elseif algorithm == :first_order - C = @views 𝓂.caches.first_order_solution_matrix[:,𝓂.constants.post_model_macro.nPast_not_future_and_mixed+1:end] + C = 𝓂.caches.first_order_solution_matrix[:,𝓂.constants.post_model_macro.nPast_not_future_and_mixed+1:end]::Matrix{Float64} CC = C[cond_var_idx,free_shock_idx] if length(cond_var_idx) == 1 @assert any(CC .!= 0) "Free shocks have no impact on conditioned variable in period 1." + shocks[free_shock_idx,1] .= 0 + shocks[free_shock_idx,1] = CC \ (conditions[cond_var_idx,1] - state_update(initial_state, Float64[shocks[:,1]...])[cond_var_idx]) elseif length(free_shock_idx) == length(cond_var_idx) - CC = RF.lu(CC, check = false) - - @assert ℒ.issuccess(CC) "Numerical stabiltiy issues for restrictions in period 1." + CC_lu_ws = FastLapackInterface.LUWs(CC) + CC_lu_ws, _, ok, CC_lu_handle = factorize_lu!(Val(:FastLapack), CC, CC_lu_ws, size(CC)) + + @assert ok "Numerical stabiltiy issues for restrictions in period 1." + + CC_rhs = conditions[cond_var_idx,1] - state_update(initial_state, Float64[shocks[:,1]...])[cond_var_idx] + solve_lu_left!(CC, CC_rhs, CC_lu_ws, CC_lu_handle) + shocks[free_shock_idx,1] .= 0 + shocks[free_shock_idx,1] = CC_rhs + else + shocks[free_shock_idx,1] .= 0 + shocks[free_shock_idx,1] = CC \ (conditions[cond_var_idx,1] - state_update(initial_state, Float64[shocks[:,1]...])[cond_var_idx]) end - shocks[free_shock_idx,1] .= 0 - - shocks[free_shock_idx,1] = CC \ (conditions[cond_var_idx,1] - state_update(initial_state, Float64[shocks[:,1]...])[cond_var_idx]) - Y[:,1] = state_update(initial_state, Float64[shocks[:,1]...]) for i in 2:size(conditions,2) @@ -971,15 +1039,20 @@ function get_conditional_forecast(𝓂::ℳ, if length(cond_var_idx) == 1 @assert any(CC .!= 0) "Free shocks have no impact on conditioned variable in period " * repr(i) * "." + shocks[free_shock_idx,i] = CC \ (conditions[cond_var_idx,i] - state_update(Y[:,i-1], Float64[shocks[:,i]...])[cond_var_idx]) elseif length(free_shock_idx) == length(cond_var_idx) - - CC = RF.lu(CC, check = false) - - @assert ℒ.issuccess(CC) "Numerical stabiltiy issues for restrictions in period " * repr(i) * "." + CC_lu_ws = FastLapackInterface.LUWs(CC) + CC_lu_ws, _, ok, CC_lu_handle = factorize_lu!(Val(:FastLapack), CC, CC_lu_ws, size(CC)) + + @assert ok "Numerical stabiltiy issues for restrictions in period " * repr(i) * "." + + CC_rhs = conditions[cond_var_idx,i] - state_update(Y[:,i-1], Float64[shocks[:,i]...])[cond_var_idx] + solve_lu_left!(CC, CC_rhs, CC_lu_ws, CC_lu_handle) + shocks[free_shock_idx,i] = CC_rhs + else + shocks[free_shock_idx,i] = CC \ (conditions[cond_var_idx,i] - state_update(Y[:,i-1], Float64[shocks[:,i]...])[cond_var_idx]) end - shocks[free_shock_idx,i] = CC \ (conditions[cond_var_idx,i] - state_update(Y[:,i-1], Float64[shocks[:,i]...])[cond_var_idx]) - Y[:,i] = state_update(Y[:,i-1], Float64[shocks[:,i]...]) end end @@ -994,6 +1067,8 @@ function get_conditional_forecast(𝓂::ℳ, axis1 = [𝓂.constants.post_model_macro.var[var_idx]; map(x->Symbol(string(x) * "₍ₓ₎"), 𝓂.constants.post_model_macro.exo)] end + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return KeyedArray([Y[var_idx,:] .+ (levels ? reference_steady_state + SSS_delta : SSS_delta)[var_idx]; convert(Matrix{Float64}, shocks)]; Variables_and_shocks = axis1, Periods = 1:periods) end @@ -1016,7 +1091,10 @@ If occasionally binding constraints are present in the model, they are not taken - $NEGATIVE_SHOCK® - $INITIAL_STATE®1 - `levels` [Default: `false`, Type: `Bool`]: $LEVELS® +- $ALGORITHM® - $QME® +- $SYLVESTER® +- $LYAPUNOV® - $TOLERANCES® - $VERBOSE® @@ -1034,7 +1112,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -1052,10 +1130,154 @@ get_irf(RBC, RBC.parameter_values) 0.01 0.002 0.0004 8.0e-5 2.74878e-29 5.49756e-30 ``` """ + +# ── IRF helpers: algorithm-dispatched initial state and forward simulation ── + +# Extract/compute initial state for IRF from SSS output +function irf_initial_state(::Val{:first_order}, state, SS_and_pars, initial_state::Vector{Float64}, nVars::Int, ::Type{S}) where S + initial_state == [0.0] ? zeros(S, nVars) : convert(Vector{S}, initial_state) - SS_and_pars[1:nVars] +end + +function irf_initial_state(::Val{:pruned_second_order}, state, SS_and_pars, initial_state::Vector{Float64}, nVars::Int, ::Type{S}) where S + initial_state == [0.0] ? state : [convert(Vector{S}, initial_state) - SS_and_pars[1:nVars], state[2]] +end + +function irf_initial_state(::Val{:pruned_third_order}, state, SS_and_pars, initial_state::Vector{Float64}, nVars::Int, ::Type{S}) where S + initial_state == [0.0] ? state : [convert(Vector{S}, initial_state) - SS_and_pars[1:nVars], state[2], state[3]] +end + +function irf_initial_state(::Val{:second_order}, state, SS_and_pars, initial_state::Vector{Float64}, nVars::Int, ::Type{S}) where S + initial_state == [0.0] ? (state isa Vector{<:Vector} ? state[1] : state) : convert(Vector{S}, initial_state) - SS_and_pars[1:nVars] +end + +function irf_initial_state(::Val{:third_order}, state, SS_and_pars, initial_state::Vector{Float64}, nVars::Int, ::Type{S}) where S + initial_state == [0.0] ? (state isa Vector{<:Vector} ? state[1] : state) : convert(Vector{S}, initial_state) - SS_and_pars[1:nVars] +end + + +# Forward simulation storing intermediate states and shocks +function irf_forward_simulate!(::Val{:first_order}, + Y_all::Array{S,3}, states_store, shocks_store, + init_st, shock_idx, shocks_input, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) where S + sol_mat = 𝐒 + for (si, ii) in enumerate(shock_idx) + shock_hist = zeros(nExo, periods) + if shocks_input isa Union{Symbol_input,String_input} + shocks_input ≠ :none && (shock_hist[ii, 1] = negative_shock ? -1 : 1) + else + shock_hist = shock_history + end + states_store[si, 1] = init_st + for t in 1:periods + shocks_store[si, t] = shock_hist[:, t] + prev = states_store[si, t] + y_t = sol_mat * [prev[past_idx]; shocks_store[si, t]] + states_store[si, t+1] = y_t + Y_all[:, t, si] = y_t + end + end +end + +function irf_forward_simulate!(::Val{:pruned_second_order}, + Y_all::Array{S,3}, states_store, shocks_store, + init_st, shock_idx, shocks_input, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) where S + 𝐒₁, 𝐒₂ = 𝐒 + for (si, ii) in enumerate(shock_idx) + shock_hist = zeros(nExo, periods) + if shocks_input isa Union{Symbol_input,String_input} + shocks_input ≠ :none && (shock_hist[ii, 1] = negative_shock ? -1 : 1) + else + shock_hist = shock_history + end + states_store[si, 1] = init_st + for t in 1:periods + shocks_store[si, t] = shock_hist[:, t] + new_st = pruned_second_order_state_update(states_store[si, t], shocks_store[si, t], past_idx, nVars, 𝐒₁, 𝐒₂) + states_store[si, t+1] = new_st + Y_all[:, t, si] = sum(new_st) + end + end +end + +function irf_forward_simulate!(::Val{:pruned_third_order}, + Y_all::Array{S,3}, states_store, shocks_store, + init_st, shock_idx, shocks_input, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) where S + 𝐒₁, 𝐒₂, 𝐒₃ = 𝐒 + for (si, ii) in enumerate(shock_idx) + shock_hist = zeros(nExo, periods) + if shocks_input isa Union{Symbol_input,String_input} + shocks_input ≠ :none && (shock_hist[ii, 1] = negative_shock ? -1 : 1) + else + shock_hist = shock_history + end + states_store[si, 1] = init_st + for t in 1:periods + shocks_store[si, t] = shock_hist[:, t] + new_st = pruned_third_order_state_update(states_store[si, t], shocks_store[si, t], past_idx, nVars, 𝐒₁, 𝐒₂, 𝐒₃) + states_store[si, t+1] = new_st + Y_all[:, t, si] = sum(new_st) + end + end +end + +function irf_forward_simulate!(::Val{:second_order}, + Y_all::Array{S,3}, states_store, shocks_store, + init_st, shock_idx, shocks_input, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) where S + 𝐒₁, 𝐒₂ = 𝐒 + for (si, ii) in enumerate(shock_idx) + shock_hist = zeros(nExo, periods) + if shocks_input isa Union{Symbol_input,String_input} + shocks_input ≠ :none && (shock_hist[ii, 1] = negative_shock ? -1 : 1) + else + shock_hist = shock_history + end + states_store[si, 1] = init_st + for t in 1:periods + shocks_store[si, t] = shock_hist[:, t] + prev = states_store[si, t] + aug = [prev[past_idx]; one(S); shocks_store[si, t]] + y_t = 𝐒₁ * aug + 𝐒₂ * ℒ.kron(aug, aug) / 2 + states_store[si, t+1] = y_t + Y_all[:, t, si] = y_t + end + end +end + +function irf_forward_simulate!(::Val{:third_order}, + Y_all::Array{S,3}, states_store, shocks_store, + init_st, shock_idx, shocks_input, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) where S + 𝐒₁, 𝐒₂, 𝐒₃ = 𝐒 + for (si, ii) in enumerate(shock_idx) + shock_hist = zeros(nExo, periods) + if shocks_input isa Union{Symbol_input,String_input} + shocks_input ≠ :none && (shock_hist[ii, 1] = negative_shock ? -1 : 1) + else + shock_hist = shock_history + end + states_store[si, 1] = init_st + for t in 1:periods + shocks_store[si, t] = shock_hist[:, t] + prev = states_store[si, t] + aug = [prev[past_idx]; one(S); shocks_store[si, t]] + kaug = ℒ.kron(aug, aug) + y_t = 𝐒₁ * aug + 𝐒₂ * kaug / 2 + 𝐒₃ * ℒ.kron(kaug, aug) / 6 + states_store[si, t+1] = y_t + Y_all[:, t, si] = y_t + end + end +end + + function get_irf(𝓂::ℳ, parameters::Vector{S}; steady_state_function::SteadyStateFunctionType = missing, periods::Int = DEFAULT_PERIODS, + algorithm::Symbol = :first_order, 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, @@ -1063,19 +1285,31 @@ function get_irf(𝓂::ℳ, levels::Bool = false, verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM) where S <: Real + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::Array{S,3} where S <: Real + + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end opts = merge_calculation_options(tol = tol, verbose = verbose, - quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm) + 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 # 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, + algorithm = algorithm) shocks = 𝓂.constants.post_model_macro.nExo == 0 ? :none : shocks @@ -1085,60 +1319,36 @@ function get_irf(𝓂::ℳ, var_idx = parse_variables_input_to_index(variables, 𝓂) |> sort - 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) - return zeros(S, length(var_idx), periods, shocks == :none ? 1 : length(shock_idx)) - end - - ∇₁ = calculate_jacobian(parameters, reference_steady_state, 𝓂.caches, 𝓂.functions.jacobian)# |> Matrix + nVars = 𝓂.constants.post_model_macro.nVars + past_idx = 𝓂.constants.post_model_macro.past_not_future_and_mixed_idx + nPast = 𝓂.constants.post_model_macro.nPast_not_future_and_mixed - # 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) - - @ignore_derivatives update_perturbation_counter!(𝓂.counters, solved, estimation = estimation, order = 1) + constants_obj, SS_and_pars, 𝐒, state, solved = get_relevant_steady_state_and_state_update(Val(algorithm), parameters, 𝓂, opts = opts, estimation = estimation) - if solved - 𝓂.caches.qme_solution = qme_sol - else + if !solved + if !use_workspaces; 𝓂.workspaces = orig_ws; end return zeros(S, length(var_idx), periods, shocks == :none ? 1 : length(shock_idx)) end - state_update = function(state::Vector, shock::Vector) sol_mat * [state[𝓂.constants.post_model_macro.past_not_future_and_mixed_idx]; shock] end + nExo = 𝓂.constants.post_model_macro.nExo + nShocks = shocks == :none ? 1 : length(shock_idx) - initial_state = initial_state == [0.0] ? zeros(𝓂.constants.post_model_macro.nVars) : initial_state - reference_steady_state[1:length(𝓂.constants.post_model_macro.var)] + # Dispatched initial state and forward simulation + val_alg = Val(algorithm) + init_state = irf_initial_state(val_alg, state, SS_and_pars, initial_state, nVars, S) - # Y = zeros(𝓂.constants.post_model_macro.nVars,periods,𝓂.constants.post_model_macro.nExo) - Ŷ = [] + Y_all = zeros(S, nVars, periods, nShocks) + states_store = Array{Any}(undef, nShocks, periods + 1) + shocks_store = Array{Vector{S}}(undef, nShocks, periods) - for ii in shock_idx - Y = [] + irf_forward_simulate!(val_alg, Y_all, states_store, shocks_store, + init_state, shock_idx, shocks, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) - if shocks isa Union{Symbol_input,String_input} - shock_history = zeros(𝓂.constants.post_model_macro.nExo,periods) - if shocks ≠ :none - shock_history[ii,1] = negative_shock ? -1 : 1 - end - end - - push!(Y, state_update(initial_state,shock_history[:,1])) + reference_steady_state = SS_and_pars[1:nVars] + deviations = Y_all[var_idx, :, :] - for t in 1:periods-1 - push!(Y, state_update(Y[end],shock_history[:,t+1])) - end - - push!(Ŷ, reduce(hcat,Y)) - end - - deviations = reshape(reduce(hcat,Ŷ),𝓂.constants.post_model_macro.nVars, periods, shocks == :none ? 1 : length(shock_idx))[var_idx,:,:] + if !use_workspaces; 𝓂.workspaces = orig_ws; end if levels return deviations .+ reference_steady_state[var_idx] @@ -1193,7 +1403,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -1216,7 +1426,7 @@ And data, 4×40×1 Array{Float64, 3}: (:z) 0.01 0.002 2.74878e-29 5.49756e-30 ``` """ -function get_irf(𝓂::ℳ; +@unstable function get_irf(𝓂::ℳ; periods::Int = DEFAULT_PERIODS, algorithm::Symbol = DEFAULT_ALGORITHM, parameters::ParameterType = nothing, @@ -1234,11 +1444,17 @@ function get_irf(𝓂::ℳ; # timer::TimerOutput = TimerOutput(), verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM)::KeyedArray where R <: Real + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray where R <: Real # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -1329,6 +1545,8 @@ function get_irf(𝓂::ℳ; enforce_obc = occasionally_binding_constraints, algorithm = algorithm) + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return responses end @@ -1338,7 +1556,7 @@ end """ See [`get_irf`](@ref) """ -get_irfs = get_irf +@unstable get_irfs = get_irf """ See [`get_irf`](@ref) @@ -1363,22 +1581,22 @@ get_IRF = get_irf """ Wrapper for [`get_irf`](@ref) with `shocks = :simulate`. Function returns values in levels by default. """ -simulate(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., shocks = :simulate, levels = get(kwargs, :levels, true))#[:,:,1] +@unstable simulate(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., shocks = :simulate, levels = get(kwargs, :levels, true))#[:,:,1] """ Wrapper for [`get_irf`](@ref) with `shocks = :simulate`. Function returns values in levels by default. """ -get_simulation(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., shocks = :simulate, levels = get(kwargs, :levels, true))#[:,:,1] +@unstable get_simulation(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., shocks = :simulate, levels = get(kwargs, :levels, true))#[:,:,1] """ Wrapper for [`get_irf`](@ref) with `shocks = :simulate`. Function returns values in levels by default. """ -get_simulations(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., shocks = :simulate, levels = get(kwargs, :levels, true))#[:,:,1] +@unstable get_simulations(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., shocks = :simulate, levels = get(kwargs, :levels, true))#[:,:,1] """ Wrapper for [`get_irf`](@ref) with `generalised_irf = true`. """ -get_girf(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., generalised_irf = true) +@unstable get_girf(𝓂::ℳ; kwargs...) = get_irf(𝓂; kwargs..., generalised_irf = true) @@ -1421,7 +1639,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -1442,7 +1660,7 @@ And data, 4×6 Matrix{Float64}: (:z) 0.0 0.0 0.0 0.0 0.0 0.0 ``` """ -function get_steady_state(𝓂::ℳ; +@unstable function get_steady_state(𝓂::ℳ; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, derivatives::Bool = DEFAULT_DERIVATIVES_FLAG, @@ -1453,9 +1671,15 @@ function get_steady_state(𝓂::ℳ; verbose::Bool = DEFAULT_VERBOSE, silent::Bool = DEFAULT_SILENT_FLAG, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, - sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂))::KeyedArray + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray # @nospecialize # reduce compile time + + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, @@ -1491,7 +1715,7 @@ function get_steady_state(𝓂::ℳ; param_idx = indexin([parameter_derivatives], 𝓂.constants.post_complete_parameters.parameters) length_par = 1 - elseif length(parameter_derivatives) > 1 + else for p in vec(collect(parameter_derivatives)) @assert p ∈ 𝓂.constants.post_complete_parameters.parameters string(p) * " is not part of the free model parameters." end @@ -1501,8 +1725,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 +1749,11 @@ function get_steady_state(𝓂::ℳ; end end - var_idx = indexin([vars_in_ss_equations...], [𝓂.constants.post_model_macro.var...,𝓂.equations.calibration_parameters...]) + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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,77 +1778,66 @@ 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]) + if !use_workspaces; 𝓂.workspaces = orig_ws; end + 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)) # else # return ComponentMatrix(hcat(collect(NSSS), dNSSS)',Axis(vcat(:SS, 𝓂.constants.post_complete_parameters.parameters)),Axis([sort(union(𝓂.constants.post_model_macro.exo_present,var))...,𝓂.calibration_equations_parameters...])) # return NamedArray(hcat(collect(NSSS), dNSSS), ([sort(union(𝓂.constants.post_model_macro.exo_present,var))..., 𝓂.calibration_equations_parameters...], vcat(:Steady_state, 𝓂.constants.post_complete_parameters.parameters)), ("Var. and par.", "∂x/∂y")) + if !use_workspaces; 𝓂.workspaces = orig_ws; end return KeyedArray(hcat(SS[[var_idx...,calib_idx...]],dSS); Variables_and_calibrated_parameters = axis1, Steady_state_and_∂steady_state∂parameter = axis2) # end end else # return ComponentVector(collect(NSSS),Axis([sort(union(𝓂.constants.post_model_macro.exo_present,var))...,𝓂.calibration_equations_parameters...])) # return NamedArray(collect(NSSS), [sort(union(𝓂.constants.post_model_macro.exo_present,var))..., 𝓂.calibration_equations_parameters...], ("Variables and calibrated parameters")) + if !use_workspaces; 𝓂.workspaces = orig_ws; end return KeyedArray(SS[[var_idx...,calib_idx...]]; Variables_and_calibrated_parameters = axis1) end # ComponentVector(non_stochastic_steady_state = ComponentVector(NSSS.non_stochastic_steady_state, Axis(sort(union(𝓂.constants.post_model_macro.exo_present,var)))), # 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 @@ -1630,58 +1845,58 @@ end """ Wrapper for [`get_steady_state`](@ref) with `stochastic = false`. """ -get_non_stochastic_steady_state(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = false) +@unstable get_non_stochastic_steady_state(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = false) """ Wrapper for [`get_steady_state`](@ref) with `stochastic = true`. """ -get_stochastic_steady_state(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) +@unstable get_stochastic_steady_state(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) """ Wrapper for [`get_steady_state`](@ref) with `stochastic = true`. """ -get_SSS(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) +@unstable get_SSS(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) """ Wrapper for [`get_steady_state`](@ref) with `stochastic = true`. """ -SSS(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) +@unstable SSS(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) """ Wrapper for [`get_steady_state`](@ref) with `stochastic = true`. """ -sss(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) +@unstable sss(args...; kwargs...) = get_steady_state(args...; kwargs..., stochastic = true) """ See [`get_steady_state`](@ref) """ -SS = get_steady_state +@unstable SS(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -steady_state = get_steady_state +@unstable steady_state(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -get_SS = get_steady_state +@unstable get_SS(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -get_ss = get_steady_state +@unstable get_ss(args...; kwargs...) = get_steady_state(args...; kwargs...) """ See [`get_steady_state`](@ref) """ -ss(args...; kwargs...) = get_steady_state(args...; kwargs...) +@unstable ss(args...; kwargs...) = get_steady_state(args...; kwargs...) @@ -1717,7 +1932,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -1733,22 +1948,28 @@ get_solution(RBC) And data, 4×4 adjoint(::Matrix{Float64}) with eltype Float64: (:c) (:k) (:q) (:z) (:Steady_state) 5.93625 47.3903 6.88406 0.0 - (:k₍₋₁₎) 0.0957964 0.956835 0.0726316 -0.0 + (:k₍₋₁₎) 0.0957964 0.956835 0.0726316 0.0 (:z₍₋₁₎) 0.134937 1.24187 1.37681 0.2 (:eps_z₍ₓ₎) 0.00674687 0.0620937 0.0688406 0.01 ``` """ -function get_solution(𝓂::ℳ; +@unstable function get_solution(𝓂::ℳ; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, algorithm::Symbol = DEFAULT_ALGORITHM, silent::Bool = DEFAULT_SILENT_FLAG, 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(𝓂))::KeyedArray + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::KeyedArray # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -1762,9 +1983,7 @@ function get_solution(𝓂::ℳ; silent = silent, algorithm = algorithm) - if algorithm == :first_order - solution_matrix = 𝓂.caches.first_order_solution_matrix - end + solution_matrix = 𝓂.caches.first_order_solution_matrix axis1 = [𝓂.constants.post_model_macro.past_not_future_and_mixed; :Volatility; 𝓂.constants.post_model_macro.exo] @@ -1784,6 +2003,8 @@ function get_solution(𝓂::ℳ; 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 !use_workspaces; 𝓂.workspaces = orig_ws; end + if algorithm == :second_order return KeyedArray(permutedims(reshape(𝓂.caches.second_order_solution * 𝓂.constants.second_order.𝐔₂, 𝓂.constants.post_model_macro.nVars, @@ -1836,7 +2057,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 @@ -1846,22 +2070,22 @@ end """ Wrapper for [`get_solution`](@ref) with `algorithm = :first_order`. """ -get_first_order_solution(args...; kwargs...) = get_solution(args...; kwargs..., algorithm = :first_order) +@unstable get_first_order_solution(args...; kwargs...) = get_solution(args...; kwargs..., algorithm = :first_order) """ Wrapper for [`get_solution`](@ref) with `algorithm = :second_order`. """ -get_second_order_solution(args...; kwargs...) = get_solution(args...; kwargs..., algorithm = :second_order) +@unstable get_second_order_solution(args...; kwargs...) = get_solution(args...; kwargs..., algorithm = :second_order) """ Wrapper for [`get_solution`](@ref) with `algorithm = :third_order`. """ -get_third_order_solution(args...; kwargs...) = get_solution(args...; kwargs..., algorithm = :third_order) +@unstable get_third_order_solution(args...; kwargs...) = get_solution(args...; kwargs..., algorithm = :third_order) """ See [`get_solution`](@ref) """ -get_perturbation_solution(args...; kwargs...) = get_solution(args...; kwargs...) +@unstable get_perturbation_solution(args...; kwargs...) = get_solution(args...; kwargs...) @@ -1883,7 +2107,7 @@ Function to use when differentiating IRFs with respect to parameters. - $VERBOSE® # Returns -- `Tuple` consisting of a `Vector` containing the NSSS, followed by a `Matrix` containing the first order solution matrix. In case of higher order solutions, `SparseMatrixCSC` represent the higher order solution matrices. The last element is a `Bool` indicating the correctness of the solution provided. +- `Tuple{Vector, Vector{AbstractMatrix}, Bool}` consisting of a `Vector` containing the NSSS, a `Vector` of solution matrices (one `Matrix` for first order, two for second order, three for third order), and a `Bool` indicating the correctness of the solution provided. # Examples ```jldoctest @@ -1896,7 +2120,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -1906,21 +2130,47 @@ end get_solution(RBC, RBC.parameter_values) # output -([5.936252888048724, 47.39025414828808, 6.884057971014486, 0.0], - [0.09579643002421227 0.1349373930517757 0.006746869652588215; - 0.9568351489231555 1.241874201151121 0.06209371005755664; - 0.07263157894736819 1.376811594202897 0.06884057971014486; - 0.0 0.19999999999999998 0.01], true) +([5.936252888048734, 47.39025414828825, 6.8840579710144985, 0.0], AbstractMatrix{Float64}[[0.09579643002421026 0.1349373930517762 0.006746869652588118; 0.9568351489231574 1.2418742011511228 0.062093710057556865; 0.07263157894736799 1.3768115942028993 0.06884057971014498; 0.0 0.2 0.01]], true) ``` """ + +# Construct a failure return value for get_solution with uniform tuple type. +# When 𝐒₁ is provided, it is included as the first solution matrix placeholder. +function get_solution_fail(algorithm::Symbol, SS::Vector{S}, nVar::Int, ::Type{S}) where S <: Real + placeholder = zeros(S, nVar, 2) + if algorithm in [:second_order, :pruned_second_order] + return SS, AbstractMatrix{S}[placeholder, zeros(S, nVar, 2)], false + elseif algorithm in [:third_order, :pruned_third_order] + return SS, AbstractMatrix{S}[placeholder, zeros(S, nVar, 2), zeros(S, nVar, 2)], false + else + return SS, AbstractMatrix{S}[placeholder], false + end +end + +function get_solution_fail(algorithm::Symbol, SS::Vector{S}, nVar::Int, ::Type{S}, 𝐒₁::AbstractMatrix{S}) where S <: Real + if algorithm in [:second_order, :pruned_second_order] + return SS, AbstractMatrix{S}[𝐒₁, zeros(S, nVar, 2)], false + elseif algorithm in [:third_order, :pruned_third_order] + return SS, AbstractMatrix{S}[𝐒₁, zeros(S, nVar, 2), zeros(S, nVar, 2)], false + else + return SS, AbstractMatrix{S}[𝐒₁], false + end +end + function 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 + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::Tuple{Vector{S}, Vector{AbstractMatrix{S}}, Bool} where S <: Real + + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, @@ -1932,118 +2182,80 @@ function get_solution(𝓂::ℳ, # Initialize constants at entry point constants = initialise_constants!(𝓂) - @ignore_derivatives solve!(𝓂, - opts = opts, - steady_state_function = steady_state_function, - algorithm = algorithm) + nVar = length(𝓂.constants.post_model_macro.var) + + 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] - return -Inf - end - end - end + if check_bounds(parameters, 𝓂) + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return get_solution_fail(algorithm, fill(S(-Inf), nVar), nVar, S) end 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 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] - 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), spzeros(length(𝓂.constants.post_model_macro.var),2), false - else - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], zeros(length(𝓂.constants.post_model_macro.var),2), false - end + if solution_error > tol.nsss.acceptance_tol || isnan(solution_error) + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return get_solution_fail(algorithm, SS_and_pars[1:nVar], nVar, S) 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] - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], 𝐒₁, spzeros(length(𝓂.constants.post_model_macro.var),2), false - elseif algorithm in [:third_order, :pruned_third_order] - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], 𝐒₁, spzeros(length(𝓂.constants.post_model_macro.var),2), spzeros(length(𝓂.constants.post_model_macro.var),2), false - else - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], 𝐒₁, false - end + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return get_solution_fail(algorithm, SS_and_pars[1:nVar], nVar, S, 𝐒₁) 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) - - @ignore_derivatives update_perturbation_counter!(𝓂.counters, solved2, estimation = estimation, order = 2) + opts = opts, parameter_values = parameters) - if eltype(𝐒₂) == Float64 && solved2 𝓂.caches.second_order_solution = 𝐒₂ end + update_perturbation_counter!(𝓂.counters, solved2, estimation = estimation, order = 2) - 𝐒₂ *= 𝓂.constants.second_order.𝐔₂ - - if !(typeof(𝐒₂) <: AbstractSparseMatrix) - 𝐒₂ = sparse(𝐒₂) # * 𝓂.constants.second_order.𝐔₂) - end - - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], 𝐒₁, 𝐒₂, true + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂], 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 + update_perturbation_counter!(𝓂.counters, solved2, estimation = estimation, order = 2) - 𝐒₂ *= 𝓂.constants.second_order.𝐔₂ - - if !(typeof(𝐒₂) <: AbstractSparseMatrix) - 𝐒₂ = sparse(𝐒₂) # * 𝓂.constants.second_order.𝐔₂) - end - - ∇₃ = 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) - - if eltype(𝐒₃) == Float64 && solved3 𝓂.caches.third_order_solution = 𝐒₃ end - - 𝐒₃ *= 𝓂.constants.third_order.𝐔₃ - - if !(typeof(𝐒₃) <: AbstractSparseMatrix) - 𝐒₃ = sparse(𝐒₃) # * 𝓂.constants.third_order.𝐔₃) - end - - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], 𝐒₁, 𝐒₂, 𝐒₃, true + 𝐒₃, solved3 = calculate_third_order_solution(∇₁, ∇₂, ∇₃, + 𝐒₁, 𝐒₂, + 𝓂.constants, + 𝓂.workspaces, + 𝓂.caches; + initial_guess = 𝓂.caches.third_order_solution, + opts = opts, parameter_values = parameters) + + update_perturbation_counter!(𝓂.counters, solved3, estimation = estimation, order = 3) + + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂, 𝐒₃], true else - return SS_and_pars[1:length(𝓂.constants.post_model_macro.var)], 𝐒₁, true + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁], true end end @@ -2082,7 +2294,7 @@ using MacroModelling A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] end -@parameters RBC_CME begin +@parameters RBC_CME silent = true begin alpha = .157 beta = .999 delta = .0226 @@ -2114,7 +2326,7 @@ And data, 7×2×21 Array{Float64, 3}: [:, :, 11] ~ (:, :, 11.0): (:delta_eps) (:eps_z) - (:A) 5.88653e-32 1.0 + (:A) 0.0 1.0 (:Pi) 0.0245641 0.975436 (:R) 0.0245641 0.975436 (:c) 0.0175249 0.982475 @@ -2124,7 +2336,7 @@ And data, 7×2×21 Array{Float64, 3}: [:, :, 21] ~ (:, :, Inf): (:delta_eps) (:eps_z) - (:A) 9.6461e-31 1.0 + (:A) 0.0 1.0 (:Pi) 0.0156771 0.984323 (:R) 0.0156771 0.984323 (:c) 0.0134672 0.986533 @@ -2133,16 +2345,22 @@ And data, 7×2×21 Array{Float64, 3}: (:z_delta) 1.0 0.0 ``` """ -function get_conditional_variance_decomposition(𝓂::ℳ; +@unstable function get_conditional_variance_decomposition(𝓂::ℳ; periods::Union{Vector{Int},Vector{Float64},UnitRange{Int64}} = DEFAULT_CONDITIONAL_VARIANCE_PERIODS, parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM) + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, lyapunov_algorithm = lyapunov_algorithm) @@ -2159,23 +2377,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 +2409,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 +2422,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 @@ -2220,6 +2432,8 @@ function get_conditional_variance_decomposition(𝓂::ℳ; axis1 = 𝓂.constants.post_complete_parameters.var_axis axis2 = 𝓂.constants.post_complete_parameters.exo_axis_plain + if !use_workspaces; 𝓂.workspaces = orig_ws; end + KeyedArray(cond_var_decomp; Variables = axis1, Shocks = axis2, Periods = periods) end @@ -2227,13 +2441,13 @@ end """ See [`get_conditional_variance_decomposition`](@ref) """ -get_fevd = get_conditional_variance_decomposition +@unstable get_fevd = get_conditional_variance_decomposition """ See [`get_conditional_variance_decomposition`](@ref) """ -get_forecast_error_variance_decomposition = get_conditional_variance_decomposition +@unstable get_forecast_error_variance_decomposition = get_conditional_variance_decomposition """ @@ -2277,7 +2491,7 @@ using MacroModelling A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] end -@parameters RBC_CME begin +@parameters RBC_CME silent = true begin alpha = .157 beta = .999 delta = .0226 @@ -2296,7 +2510,7 @@ get_variance_decomposition(RBC_CME) → Shocks ∈ 2-element Vector{Symbol} And data, 7×2 Matrix{Float64}: (:delta_eps) (:eps_z) - (:A) 9.78485e-31 1.0 + (:A) 0.0 1.0 (:Pi) 0.0156771 0.984323 (:R) 0.0156771 0.984323 (:c) 0.0134672 0.986533 @@ -2305,15 +2519,21 @@ And data, 7×2 Matrix{Float64}: (:z_delta) 1.0 0.0 ``` """ -function get_variance_decomposition(𝓂::ℳ; +@unstable function get_variance_decomposition(𝓂::ℳ; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, verbose::Bool = DEFAULT_VERBOSE, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, - lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM) + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, lyapunov_algorithm = lyapunov_algorithm) @@ -2328,47 +2548,96 @@ 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,:] - for i in 1:𝓂.constants.post_model_macro.nExo - C = @views sol[:, 𝓂.constants.post_model_macro.nPast_not_future_and_mixed + i] - - CC = C * C' + # Use pre-detected unit root flag from QME solve when available. + # Only compute Schur decomposition of A when unit roots are present (needed for deflation). + has_unit_roots = 𝓂.caches.has_unit_roots + unit_root_tol = 1e-8 + lyap_ws = ensure_lyapunov_workspace!(𝓂.workspaces, 𝓂.constants.post_model_macro.nVars, :first_order) + + if !has_unit_roots + # Standard path: no unit roots, solve each shock directly + for i in 1:𝓂.constants.post_model_macro.nExo + C = @views sol[:, 𝓂.constants.post_model_macro.nPast_not_future_and_mixed + i] + CC = C * C' + covar_raw, _ = solve_lyapunov_equation(A, CC, lyap_ws, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.first_order.lyapunov, + verbose = opts.verbose) + variances_by_shock[:,i] = ℒ.diag(covar_raw) + end + else + # Unit root path: compute Schur decomposition of A for deflation + A_dense = collect(A) + A_work = copy(A_dense) + Tmat, U_schur, n_unstable = ordered_schur!(A_work, unit_root_tol, lyap_ws.schur_ws) + n = size(A_dense, 1) + + if n_unstable == n + # All eigenvalues unstable — all variances are NaN + variances_by_shock .= NaN + else + n_stable = n - n_unstable + stable_range = (n_unstable + 1):n + T_ss = Tmat[stable_range, stable_range] + U_s = U_schur[:, stable_range] - # Ensure lyapunov workspace is properly sized and get it - lyap_ws = ensure_lyapunov_workspace_1st_order!(𝓂) + # Identify unit-root variables + U_u = @view U_schur[:, 1:n_unstable] + unstable_loading = vec(sum(abs2, U_u; dims = 2)) + unit_root_vars = unstable_loading .> unit_root_tol - 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, - verbose = opts.verbose) + ws_stable = Lyapunov_workspace(n_stable) + + if opts.verbose + println("Variance decomposition: Schur pre-computed ($n_unstable unstable, $n_stable stable eigenvalues)") + end - variances_by_shock[:,i] = ℒ.diag(covar_raw) + for i in 1:𝓂.constants.post_model_macro.nExo + C = @views sol[:, 𝓂.constants.post_model_macro.nPast_not_future_and_mixed + i] + CC = C * C' + + # Transform to Schur basis and extract stable block + CC_schur = U_schur' * CC * U_schur + CC_ss = (CC_schur[stable_range, stable_range] + CC_schur[stable_range, stable_range]') / 2 + + X_ss, _, sub_tol = solve_lyapunov_equation(T_ss, CC_ss, Val(:doubling), ws_stable; + tol = opts.tol.first_order.lyapunov) + + if sub_tol > opts.tol.first_order.lyapunov.acceptance_tol + X_ss, _, sub_tol = solve_lyapunov_equation(T_ss, CC_ss, Val(:bicgstab), ws_stable; + tol = opts.tol.first_order.lyapunov) + end + + X_ss = collect(X_ss) + + # Map back: only need diagonal of U_s * X_ss * U_s' + tmp = X_ss * U_s' + var_i = vec(sum(U_s .* tmp', dims = 2)) + var_i[unit_root_vars] .= NaN + variances_by_shock[:,i] = var_i + end + end end 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 @@ -2378,6 +2647,8 @@ function get_variance_decomposition(𝓂::ℳ; axis1 = 𝓂.constants.post_complete_parameters.var_axis axis2 = 𝓂.constants.post_complete_parameters.exo_axis_plain + if !use_workspaces; 𝓂.workspaces = orig_ws; end + KeyedArray(var_decomp; Variables = axis1, Shocks = axis2) end @@ -2386,7 +2657,7 @@ end """ See [`get_variance_decomposition`](@ref) """ -get_var_decomp = get_variance_decomposition +@unstable get_var_decomp = get_variance_decomposition @@ -2423,7 +2694,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -2444,59 +2715,12 @@ And data, 4×4 Matrix{Float64}: (:z) 0.314562 0.296104 0.965726 1.0 ``` """ -function get_correlation(𝓂::ℳ; - parameters::ParameterType = nothing, - steady_state_function::SteadyStateFunctionType = missing, - 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()) - # @nospecialize # reduce compile time - - 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 algorithm ∈ [:first_order, :pruned_second_order,:pruned_third_order] "Correlation can only be calculated for first order perturbation or second and third order pruned perturbation solutions." - - solve!(𝓂, - parameters = parameters, - steady_state_function = steady_state_function, - opts = opts, - algorithm = algorithm) - - if algorithm == :pruned_third_order - covar_dcmp, state_μ, SS_and_pars, solved = calculate_third_order_moments(𝓂.parameter_values, :full_covar, 𝓂, opts = opts) - elseif algorithm == :pruned_second_order - covar_dcmp, Σᶻ₂, state_μ, Δμˢ₂, autocorr_tmp, ŝ_to_ŝ₂, ŝ_to_y₂, Σʸ₁, Σᶻ₁, SS_and_pars, 𝐒₁, ∇₁, 𝐒₂, ∇₂, solved = calculate_second_order_moments_with_covariance(𝓂.parameter_values, 𝓂, opts = opts) - else - covar_dcmp, sol, _, SS_and_pars, solved = calculate_covariance(𝓂.parameter_values, 𝓂, opts = opts) - - @assert solved "Could not find covariance matrix." - end - - covar_dcmp[abs.(covar_dcmp) .< opts.tol.lyapunov_acceptance_tol] .= 0 - - std = sqrt.(max.(ℒ.diag(covar_dcmp),eps(Float64))) - - corr = covar_dcmp ./ (std * std') - - axis1 = 𝓂.constants.post_model_macro.var - - ensure_name_display_constants!(𝓂) - axis1 = 𝓂.constants.post_complete_parameters.var_axis - - KeyedArray(collect(corr); Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1) -end +@unstable get_correlation(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = false, covariance = false, correlation = true, derivatives = get(kwargs, :derivatives, false))[:correlation] """ See [`get_correlation`](@ref) """ -get_corr = get_correlation +@unstable get_corr = get_correlation """ @@ -2540,7 +2764,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -2561,17 +2785,23 @@ And data, 4×5 Matrix{Float64}: (:z) 0.2 0.04 0.008 0.0016 0.00032 ``` """ -function get_autocorrelation(𝓂::ℳ; +@unstable function get_autocorrelation(𝓂::ℳ; autocorrelation_periods::UnitRange{Int} = DEFAULT_AUTOCORRELATION_PERIODS, parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, algorithm::Symbol = DEFAULT_ALGORITHM, - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), 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()) + tol::Tolerances = Tolerances(), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, @@ -2592,7 +2822,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, ŝ_to_ŝ₂, ŝ_to_y₂, Σʸ₁, Σᶻ₁, SS_and_pars, 𝐒₁, ∇₁, 𝐒₂, ∇₂, solved = calculate_second_order_moments_with_covariance(𝓂.parameter_values, 𝓂, opts = opts) @@ -2600,52 +2830,55 @@ 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(ŝ_to_y₂ * ŝ_to_ŝ₂ⁱ * autocorr_tmp) ./ ℒ.diag(covar_dcmp) ŝ_to_ŝ₂ⁱ *= ŝ_to_ŝ₂ 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) - @assert solved "Could not find covariance matrix." + if !solved + @warn "Could not find covariance matrix. Results may contain NaN for unit-root variables." + end 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 ./ ℒ.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 - axis1 = 𝓂.constants.post_model_macro.var ensure_name_display_constants!(𝓂) axis1 = 𝓂.constants.post_complete_parameters.var_axis + if !use_workspaces; 𝓂.workspaces = orig_ws; end + KeyedArray(collect(autocorr); Variables = axis1, Autocorrelation_periods = autocorrelation_periods) end """ See [`get_autocorrelation`](@ref) """ -get_autocorr = get_autocorrelation +@unstable get_autocorr(args...; kwargs...) = get_autocorrelation(args...; kwargs...) """ See [`get_autocorrelation`](@ref) """ -autocorr = get_autocorrelation +@unstable autocorr(args...; kwargs...) = get_autocorrelation(args...; kwargs...) """ $(SIGNATURES) -Return the first and second moments of endogenous variables using the first, pruned second, or pruned third order perturbation solution. By default returns: non-stochastic steady state (NSSS), and standard deviations, but can optionally return variances, and covariance matrix. Derivatives of the moments (except for covariance) can also be provided by setting `derivatives` to `true`. +Return the first and second moments of endogenous variables using the first, pruned second, or pruned third order perturbation solution. By default returns: non-stochastic steady state (NSSS), and standard deviations, but can optionally return variances, covariance matrix, and correlation matrix. Derivatives of the moments can also be provided by setting `derivatives` to `true`. If occasionally binding constraints are present in the model, they are not taken into account here. @@ -2659,6 +2892,7 @@ If occasionally binding constraints are present in the model, they are not taken - `standard_deviation` [Default: `true`, Type: `Bool`]: switch to return standard deviation of endogenous variables - `variance` [Default: `false`, Type: `Bool`]: switch to return variance of endogenous variables - `covariance` [Default: `false`, Type: `Bool`]: switch to return covariance matrix of endogenous variables +- `correlation` [Default: `false`, Type: `Bool`]: switch to return correlation matrix of endogenous variables - $(VARIABLES®(DEFAULT_VARIABLES_EXCLUDING_OBC)) - $DERIVATIVES® - $PARAMETER_DERIVATIVES® @@ -2670,7 +2904,7 @@ If occasionally binding constraints are present in the model, they are not taken - $VERBOSE® # Returns -- `Dict{Symbol,KeyedArray}` containing the selected moments. All moments have variables as rows and the moment as the first column followed by partial derivatives wrt parameters. The `KeyedArray` type is provided by the `AxisKeys` package. +- `Dict{Symbol,KeyedArray}` containing the selected moments. All moments have variables as rows and the moment as the first column followed by partial derivatives wrt parameters. Covariance and correlation matrices are returned as 2D `KeyedArray`s (or 3D when `derivatives = true`). The `KeyedArray` type is provided by the `AxisKeys` package. # Examples ```jldoctest part1 @@ -2683,7 +2917,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -2720,8 +2954,23 @@ And data, 4×6 Matrix{Float64}: (:q) 0.0739325 7.39325 -0.974722 0.726551 1.08 (:z) 0.0102062 1.02062 0.0 0.0 0.0 ``` + +Correlation matrix (returned when `correlation = true`): +```jldoctest part1 +get_moments(RBC, non_stochastic_steady_state = false, standard_deviation = false, correlation = true, derivatives = false)[:correlation] +# output +2-dimensional KeyedArray(NamedDimsArray(...)) with keys: +↓ Variables ∈ 4-element Vector{Symbol} +→ 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 ∈ 4-element Vector{Symbol} +And data, 4×4 Matrix{Float64}: + (:c) (:k) (:q) (:z) + (:c) 1.0 0.999812 0.550168 0.314562 + (:k) 0.999812 1.0 0.533879 0.296104 + (:q) 0.550168 0.533879 1.0 0.965726 + (:z) 0.314562 0.296104 0.965726 1.0 +``` """ -function get_moments(𝓂::ℳ; +@unstable function get_moments(𝓂::ℳ; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, non_stochastic_steady_state::Bool = DEFAULT_NON_STOCHASTIC_STEADY_STATE_FLAG, @@ -2729,18 +2978,25 @@ function get_moments(𝓂::ℳ; standard_deviation::Bool = DEFAULT_STANDARD_DEVIATION_FLAG, variance::Bool = DEFAULT_VARIANCE_FLAG, covariance::Bool = DEFAULT_COVARIANCE_FLAG, + correlation::Bool = DEFAULT_CORRELATION_FLAG, variables::Union{Symbol_input,String_input} = DEFAULT_VARIABLES_EXCLUDING_OBC, derivatives::Bool = DEFAULT_DERIVATIVES_FLAG, parameter_derivatives::Union{Symbol_input,String_input} = DEFAULT_VARIABLE_SELECTION, algorithm::Symbol = DEFAULT_ALGORITHM, silent::Bool = DEFAULT_SILENT_FLAG, - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), 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())#limit output by selecting pars and vars like for plots and irfs!? + tol::Tolerances = Tolerances(), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)#limit output by selecting pars and vars like for plots and irfs!? # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -2754,7 +3010,7 @@ function get_moments(𝓂::ℳ; opts = opts, silent = silent) - for (moment_name, condition) in [("Mean", mean), ("Standard deviation", standard_deviation), ("Variance", variance), ("Covariance", covariance)] + for (moment_name, condition) in [("Mean", mean), ("Standard deviation", standard_deviation), ("Variance", variance), ("Covariance", covariance), ("Correlation", correlation)] if condition @assert algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] moment_name * " only available for algorithms: `first_order`, `pruned_second_order`, and `pruned_third_order`." end @@ -2785,19 +3041,36 @@ 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." + if solution_error >= tol.nsss.acceptance_tol + @warn "Could not find non-stochastic steady state. Solution error: $solution_error > $(tol.nsss.acceptance_tol)" maxlog = DEFAULT_MAXLOG + if !use_workspaces; 𝓂.workspaces = orig_ws; end + inf_val = Inf * sum(abs2, 𝓂.parameter_values) + var_idx_fail = parse_variables_input_to_index(variables, 𝓂) |> sort + axis1_fail = 𝓂.constants.post_model_macro.var[var_idx_fail] + ret = Dict{Symbol,KeyedArray}() + if non_stochastic_steady_state + axis1_nsss_fail = [axis1_fail..., 𝓂.equations.calibration_parameters...] + ret[:non_stochastic_steady_state] = KeyedArray(fill(inf_val, length(axis1_nsss_fail)); Variables = axis1_nsss_fail) + end + if mean; ret[:mean] = KeyedArray(fill(inf_val, length(axis1_fail)); Variables = axis1_fail); end + if standard_deviation; ret[:standard_deviation] = KeyedArray(fill(inf_val, length(axis1_fail)); Variables = axis1_fail); end + if variance; ret[:variance] = KeyedArray(fill(inf_val, length(axis1_fail)); Variables = axis1_fail); end + if covariance; ret[:covariance] = KeyedArray(fill(inf_val, length(var_idx_fail), length(var_idx_fail)); Variables = axis1_fail, Variables2 = axis1_fail); end + if correlation; ret[:correlation] = KeyedArray(fill(inf_val, length(var_idx_fail), length(var_idx_fail)); Variables = axis1_fail, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1_fail); end + return ret + end 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 end - if (!variance && !standard_deviation && !non_stochastic_steady_state && !mean && !covariance) + if (!variance && !standard_deviation && !non_stochastic_steady_state && !mean && !covariance && !correlation) derivatives = false end - if parameter_derivatives != :all && (variance || standard_deviation || non_stochastic_steady_state || mean || covariance) + if parameter_derivatives != :all && (variance || standard_deviation || non_stochastic_steady_state || mean || covariance || correlation) derivatives = true end @@ -2808,6 +3081,24 @@ function get_moments(𝓂::ℳ; axis1 = 𝓂.constants.post_complete_parameters.var_axis axis2 = 𝓂.constants.post_complete_parameters.exo_axis_plain + # Initialize variables used across derivative/non-derivative branches + # to satisfy JET's definite-assignment analysis + SS = KeyedArray(collect(NSSS)[var_idx]; Variables = 𝓂.constants.post_model_macro.var[var_idx]) + var_means = KeyedArray(collect(NSSS)[var_idx]; Variables = 𝓂.constants.post_model_macro.var[var_idx]) + st_dev = var_means + varrs = var_means + covar_dcmp = zeros(0, 0) + dcovariance = zeros(0, 0) + state_μ = Float64[] + autocorr = zeros(0, 0) + autocorr_tmp = zeros(0, 0) + ŝ_to_ŝ₂ = zeros(0, 0) + ŝ_to_y₂ = zeros(0, 0) + SS_and_pars = Float64[] + _dvariance_full = zeros(0, 0) + _n_cov_tuple = 0 + _cov_pb = nothing + axis3 = Symbol[] if derivatives if non_stochastic_steady_state @@ -2825,8 +3116,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 +3133,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 +3143,42 @@ 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/correlation + if variance || standard_deviation || covariance || correlation + 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 || correlation) ? :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] + if !_cov_result[5] + @warn "Could not find covariance matrix. Results may contain NaN for unit-root variables." + end + _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 +3187,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, ŝ_to_ŝ₂, ŝ_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 +3202,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,31 +3217,30 @@ 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, ŝ_to_ŝ₂, ŝ_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 + if covariance || correlation + # 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 covariance axis3 = vcat(:Covariance, 𝓂.constants.post_complete_parameters.parameters[param_idx]) @@ -2944,25 +3248,6 @@ function get_moments(𝓂::ℳ; axis3_decomposed = decompose_name.(axis3) 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, ŝ_to_ŝ₂, ŝ_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] - end end if mean && algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] @@ -2973,12 +3258,23 @@ 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." + if !mean_result[2] + @warn "Mean not found." maxlog = DEFAULT_MAXLOG + state_μ = fill(NaN, length(state_μ)) + end - # 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] + 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 @@ -3011,10 +3307,14 @@ 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 + var_means = KeyedArray(collect(NSSS)[var_idx]; Variables = 𝓂.constants.post_model_macro.var[var_idx]) + 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 +3332,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 +3362,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 +3388,38 @@ 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 + + if correlation + if algorithm == :pruned_second_order + covar_dcmp, Σᶻ₂, state_μ, Δμˢ₂, autocorr_tmp, ŝ_to_ŝ₂, ŝ_to_y₂, Σʸ₁, Σᶻ₁, SS_and_pars, 𝐒₁, ∇₁, 𝐒₂, ∇₂, solved = calculate_second_order_moments_with_covariance(𝓂.parameter_values, 𝓂, opts = opts) + if mean + var_means = KeyedArray(state_μ[var_idx]; Variables = axis1) + end + elseif algorithm == :pruned_third_order + covar_dcmp, state_μ, _, solved = calculate_third_order_moments(𝓂.parameter_values, :full_covar, 𝓂, opts = opts) + if mean + var_means = KeyedArray(state_μ[var_idx]; Variables = axis1) + end + else + covar_dcmp, ___, __, _, solved = calculate_covariance(𝓂.parameter_values, 𝓂, opts = opts) + + 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 @@ -3170,6 +3499,58 @@ function get_moments(𝓂::ℳ; ret[:covariance] = KeyedArray(covar_dcmp[var_idx, var_idx]; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1) end end + if correlation + axis1 = 𝓂.constants.post_model_macro.var[var_idx] + + if any(x -> contains(string(x), "◖"), axis1) + axis1_decomposed = decompose_name.(axis1) + 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 + + corr_full_mat, covar_sym, diag_cov, std_corr = covariance_to_correlation(covar_dcmp) + + if derivatives + n_full_vars = size(covar_dcmp, 1) + n_reduced_vars = length(var_idx) + n_params = length(param_idx) + + corr_with_derivs = zeros(n_reduced_vars, n_reduced_vars, 1 + n_params) + corr_with_derivs[:, :, 1] = corr_full_mat[var_idx, var_idx] + + for p in 1:n_params + dΣ_full = reshape(dcovariance[:, p], n_full_vars, n_full_vars) + for (ri, i) in enumerate(var_idx) + for (rj, j) in enumerate(var_idx) + σi = std_corr[i] + σj = std_corr[j] + if !isfinite(σi) || !isfinite(σj) || diag_cov[i] <= 0 || diag_cov[j] <= 0 + corr_with_derivs[ri, rj, p+1] = NaN + else + # dC[i,j]/dθ = dΣ[i,j]/(σi*σj) - C[i,j]*(dΣ[i,i]/(2*Σ[i,i]) + dΣ[j,j]/(2*Σ[j,j])) + corr_with_derivs[ri, rj, p+1] = dΣ_full[i,j] / (σi * σj) - corr_full_mat[i,j] * (dΣ_full[i,i] / (2 * diag_cov[i]) + dΣ_full[j,j] / (2 * diag_cov[j])) + end + end + end + end + + axis_corr = vcat(:Correlation, 𝓂.constants.post_complete_parameters.parameters[param_idx]) + + if any(x -> contains(string(x), "◖"), axis_corr) + axis_corr_decomposed = decompose_name.(axis_corr) + axis_corr = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis_corr_decomposed] + end + + ret[:correlation] = KeyedArray(corr_with_derivs; + Variables = axis1, + 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1, + Correlation_and_∂correlation∂parameter = axis_corr + ) + else + ret[:correlation] = KeyedArray(corr_full_mat[var_idx, var_idx]; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1) + end + end + + if !use_workspaces; 𝓂.workspaces = orig_ws; end return ret end @@ -3177,71 +3558,71 @@ end """ Wrapper for [`get_moments`](@ref) with `variance = true` and `non_stochastic_steady_state = false, standard_deviation = false, covariance = false`. """ -get_variance(args...; kwargs...) = get_moments(args...; kwargs..., variance = true, non_stochastic_steady_state = false, standard_deviation = false, covariance = false)[:variance] +@unstable get_variance(args...; kwargs...) = get_moments(args...; kwargs..., variance = true, non_stochastic_steady_state = false, standard_deviation = false, covariance = false, derivatives = get(kwargs, :derivatives, true))[:variance] """ Wrapper for [`get_moments`](@ref) with `variance = true` and `non_stochastic_steady_state = false, standard_deviation = false, covariance = false`. """ -get_var = get_variance +@unstable get_var = get_variance """ Wrapper for [`get_moments`](@ref) with `variance = true` and `non_stochastic_steady_state = false, standard_deviation = false, covariance = false`. """ -var = get_variance +@unstable var = get_variance """ Wrapper for [`get_moments`](@ref) with `standard_deviation = true` and `non_stochastic_steady_state = false, variance = false, covariance = false`. """ -get_standard_deviation(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = true, covariance = false)[:standard_deviation] +@unstable get_standard_deviation(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = true, covariance = false, derivatives = get(kwargs, :derivatives, true))[:standard_deviation] """ Wrapper for [`get_moments`](@ref) with `standard_deviation = true` and `non_stochastic_steady_state = false, variance = false, covariance = false`. """ -get_std = get_standard_deviation +@unstable get_std = get_standard_deviation """ Wrapper for [`get_moments`](@ref) with `standard_deviation = true` and `non_stochastic_steady_state = false, variance = false, covariance = false`. """ -get_stdev = get_standard_deviation +@unstable get_stdev = get_standard_deviation """ Wrapper for [`get_moments`](@ref) with `standard_deviation = true` and `non_stochastic_steady_state = false, variance = false, covariance = false`. """ -stdev = get_standard_deviation +@unstable stdev = get_standard_deviation """ Wrapper for [`get_moments`](@ref) with `standard_deviation = true` and `non_stochastic_steady_state = false, variance = false, covariance = false`. """ -std = get_standard_deviation +@unstable std = get_standard_deviation """ Wrapper for [`get_moments`](@ref) with `covariance = true` and `non_stochastic_steady_state = false, variance = false, standard_deviation = false, derivatives = false`. """ -get_covariance(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = false, covariance = true, derivatives = false)[:covariance] +@unstable get_covariance(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = false, covariance = true, derivatives = get(kwargs, :derivatives, false))[:covariance] """ Wrapper for [`get_moments`](@ref) with `covariance = true` and `non_stochastic_steady_state = false, variance = false, standard_deviation = false`. """ -get_cov = get_covariance +@unstable get_cov = get_covariance """ Wrapper for [`get_moments`](@ref) with `covariance = true` and `non_stochastic_steady_state = false, variance = false, standard_deviation = false`. """ -cov = get_covariance +@unstable cov = get_covariance """ Wrapper for [`get_moments`](@ref) with `mean = true`, and `non_stochastic_steady_state = false, variance = false, standard_deviation = false, covariance = false` """ -get_mean(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = false, covariance = false, mean = true)[:mean] +@unstable get_mean(args...; kwargs...) = get_moments(args...; kwargs..., variance = false, non_stochastic_steady_state = false, standard_deviation = false, covariance = false, mean = true, derivatives = get(kwargs, :derivatives, true))[:mean] # """ @@ -3253,7 +3634,7 @@ get_mean(args...; kwargs...) = get_moments(args...; kwargs..., variance = false """ $(SIGNATURES) -Return the first and second moments of endogenous variables using either the linearised solution or the pruned second or pruned third order perturbation solution. By default returns a `Dict` with: non-stochastic steady state (NSSS), and standard deviations, but can also return variances, and covariance matrix. Values are returned in the order given for the specific moment. +Return the first and second moments of endogenous variables using either the linearised solution or the pruned second or pruned third order perturbation solution. By default returns a `Dict` with: non-stochastic steady state (NSSS), and standard deviations, but can also return variances, covariance matrix, and correlation matrix. Values are returned in the order given for the specific moment. Function to use when differentiating model moments with respect to parameters. If occasionally binding constraints are present in the model, they are not taken into account here. @@ -3268,6 +3649,7 @@ If occasionally binding constraints are present in the model, they are not taken - `standard_deviation` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the standard deviation of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or `\"y\"`), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. - `variance` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the variance of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or `\"y\"`), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. - `covariance` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the covariance of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or `\"y\"`), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. For grouped covariance computation, pass a `Vector` of `Vector`s (e.g. `[[:y, :c], [:k, :i]]`) to compute covariances only within each group, returning a single covariance matrix where cross-group covariances are set to zero. This allows more granular control over which covariances to compute. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all variables less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all variables less those related to occasionally binding constraints. `:all` will contain all variables. +- `correlation` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the correlation matrix of selected variables. Inputs follow the same format as `covariance`, including grouped input (e.g. `[[:y, :c], [:k, :i]]`) which restricts the returned matrix to within-group correlations and sets cross-group entries to zero. Variables with non-positive variance produce `NaN` entries (left unchanged). `:all_excluding_auxiliary_and_obc` contains all variables less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all variables less those related to occasionally binding constraints. `:all` will contain all variables. - `autocorrelation` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the autocorrelation of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or `\"y\"`), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. - `autocorrelation_periods` [Default: `1:5`, Type = `UnitRange{Int}`]: periods for which to return the autocorrelation of selected variables - $STEADY_STATE_FUNCTION® @@ -3279,7 +3661,7 @@ If occasionally binding constraints are present in the model, they are not taken - $VERBOSE® # Returns -- `Dict` with the name of the statistics and the corresponding vectors (NSSS, mean, standard deviation, variance) or matrices (covariance, autocorrelation). +- `Dict` with the name of the statistics and the corresponding vectors (NSSS, mean, standard deviation, variance) or matrices (covariance, correlation, autocorrelation). # Examples ```jldoctest @@ -3292,7 +3674,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -3300,19 +3682,30 @@ end β = 0.95 end -get_statistics(RBC, RBC.parameter_values, parameters = get_parameters(RBC), standard_deviation = RBC.var) +get_statistics(RBC, RBC.parameter_values, standard_deviation = get_variables(RBC)) # output Dict{Symbol, AbstractArray{Float64}} with 1 entry: :standard_deviation => [0.0266642, 0.264677, 0.0739325, 0.0102062] +``` -# For grouped covariance (computing covariances only within specified groups): -get_statistics(RBC, RBC.parameter_values, covariance = [[:c, :k], [:y, :i]]) -# output -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...] +For grouped covariance (computing covariances only within specified groups; cross-group +entries are set to zero): +```julia +get_statistics(RBC, RBC.parameter_values, covariance = [[:c, :k], [:q, :z]]) +# Dict{Symbol, AbstractArray{Float64}} with 1 entry: +# :covariance => [0.00071098 0.00705609 0.0 0.0; 0.0 0.0700541 0.0 0.0; 0.0 0.0… +``` + +For correlation (returns the correlation matrix among the selected variables; +diagonal is 1; supports the same grouped input as `covariance`, with cross-group +entries set to zero): +```julia +get_statistics(RBC, RBC.parameter_values, correlation = [:c, :k]) +# Dict{Symbol, AbstractArray{Float64}} with 1 entry: +# :correlation => [1.0 0.999812; 0.999812 1.0] ``` """ -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, @@ -3321,14 +3714,21 @@ function get_statistics(𝓂, 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[], + correlation::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, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), 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 + tol::Tolerances = Tolerances(), + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) where T + + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, @@ -3338,71 +3738,92 @@ function get_statistics(𝓂, @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 algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] || !(!(standard_deviation == Symbol[]) || !(mean == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(correlation == 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." + @assert !(non_stochastic_steady_state == Symbol[]) || !(standard_deviation == Symbol[]) || !(mean == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(correlation == 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 + + corr_var_idx = parse_variables_input_to_index(correlation, 𝓂) + + corr_groups = is_grouped_covariance_input(correlation) ? parse_covariance_groups(correlation, 𝓂.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] solved = true - if algorithm == :pruned_third_order && !(!(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[])) + if algorithm == :pruned_third_order && !(!(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(correlation == Symbol[]) || !(autocorrelation == Symbol[])) 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[]) + if !(non_stochastic_steady_state == Symbol[]) && (standard_deviation == Symbol[]) && (variance == Symbol[]) && (covariance == Symbol[]) && (correlation == Symbol[]) && (autocorrelation == Symbol[]) SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(𝓂, all_parameters, opts = opts) # timer = timer, 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)) + 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)) + if !use_workspaces; 𝓂.workspaces = orig_ws; end return ret end + # Initialize variables that are conditionally assigned across algorithm branches + # to satisfy JET's definite-assignment analysis. Each is overwritten in the + # relevant branch below before it is actually used. + nVars = 𝓂.constants.post_model_macro.nVars + SS_and_pars = zeros(T, 0) + covar_dcmp = zeros(T, 0, 0) + state_μ = zeros(T, 0) + sol = zeros(T, 0, 0) + autocorr_tmp = zeros(T, 0, 0) + ŝ_to_ŝ₂ = zeros(T, 0, 0) + ŝ_to_y₂ = zeros(T, 0, 0) + autocorr = zeros(T, 0, 0) + varrs = zeros(T, 0) + st_dev = zeros(T, 0) + solved = false + if algorithm == :pruned_third_order if !(autocorrelation == Symbol[]) - second_mom_third_order = union(autocorr_var_idx, std_var_idx, var_var_idx) + second_mom_third_order = union(autocorr_var_idx, std_var_idx, var_var_idx, corr_var_idx) - covar_dcmp, state_μ, autocorr, SS_and_pars, solved = 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) + covar_dcmp, state_μ, autocorr, SS_and_pars, solved = calculate_third_order_moments_with_autocorrelation(all_parameters, 𝓂.constants.post_model_macro.var[second_mom_third_order], 𝓂, covariance = 𝓂.constants.post_model_macro.var[union(covar_var_idx, corr_var_idx)], opts = opts, autocorrelation_periods = autocorrelation_periods) - elseif !(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) + elseif !(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(correlation == Symbol[]) - covar_dcmp, state_μ, SS_and_pars, solved = 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) + covar_dcmp, state_μ, SS_and_pars, solved = calculate_third_order_moments(all_parameters, 𝓂.constants.post_model_macro.var[union(std_var_idx, var_var_idx, corr_var_idx)], 𝓂, covariance = 𝓂.constants.post_model_macro.var[union(covar_var_idx, corr_var_idx)], opts = opts) end elseif algorithm == :pruned_second_order - if !(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(autocorrelation == Symbol[]) + if !(standard_deviation == Symbol[]) || !(variance == Symbol[]) || !(covariance == Symbol[]) || !(correlation == Symbol[]) || !(autocorrelation == Symbol[]) covar_dcmp, Σᶻ₂, state_μ, Δμˢ₂, autocorr_tmp, ŝ_to_ŝ₂, ŝ_to_y₂, Σʸ₁, Σᶻ₁, SS_and_pars, 𝐒₁, ∇₁, 𝐒₂, ∇₂, solved = calculate_second_order_moments_with_covariance(all_parameters, 𝓂, opts = opts) else state_μ, Δμˢ₂, Σʸ₁, Σᶻ₁, SS_and_pars, 𝐒₁, ∇₁, 𝐒₂, ∇₂, solved = calculate_second_order_moments(all_parameters, 𝓂, opts = opts) @@ -3416,14 +3837,14 @@ function get_statistics(𝓂, SS = SS_and_pars[1:end - length(𝓂.equations.calibration)] - if !(variance == Symbol[]) + if solved && !(variance == Symbol[]) varrs = convert(Vector{T},max.(ℒ.diag(covar_dcmp),eps(Float64))) if !(standard_deviation == Symbol[]) st_dev = sqrt.(varrs) end end - if !(autocorrelation == Symbol[]) + if solved && !(autocorrelation == Symbol[]) if algorithm == :pruned_second_order ŝ_to_ŝ₂ⁱ = zero(ŝ_to_ŝ₂) ŝ_to_ŝ₂ⁱ += ℒ.diagm(ones(size(ŝ_to_ŝ₂,1))) @@ -3435,17 +3856,17 @@ function get_statistics(𝓂, ŝ_to_ŝ₂ⁱ *= ŝ_to_ŝ₂ 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 - if !(standard_deviation == Symbol[]) + if solved && !(standard_deviation == Symbol[]) st_dev = sqrt.(abs.(convert(Vector{T}, max.(ℒ.diag(covar_dcmp),eps(Float64))))) end @@ -3512,11 +3933,39 @@ function get_statistics(𝓂, 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 !(correlation == Symbol[]) + if solved + corr_full_mat, _, _, _ = covariance_to_correlation(covar_dcmp) + + if !isnothing(corr_groups) + # Block-grouped correlation: cross-group entries left as zero + corr_result = zeros(T, length(corr_var_idx), length(corr_var_idx)) + for group in corr_groups + for i in group + i_pos = findfirst(==(i), corr_var_idx) + isnothing(i_pos) && continue + for j in group + j_pos = findfirst(==(j), corr_var_idx) + isnothing(j_pos) && continue + corr_result[i_pos, j_pos] = corr_full_mat[i, j] + end + end + end + ret[:correlation] = corr_result + else + ret[:correlation] = corr_full_mat[corr_var_idx, corr_var_idx] + end + else + ret[:correlation] = fill(Inf * sum(abs2, parameter_values), isnothing(corr_var_idx) ? 0 : length(corr_var_idx), isnothing(corr_var_idx) ? 0 : length(corr_var_idx)) + end + end if !(autocorrelation == Symbol[]) # push!(ret,autocorr[autocorr_var_idx,:] ) 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 + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return ret end @@ -3562,7 +4011,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -3570,11 +4019,13 @@ end β = 0.95 end +import Random; Random.seed!(3) + simulated_data = simulate(RBC) get_loglikelihood(RBC, simulated_data([:k], :, :simulate), RBC.parameter_values) # output -58.24780188977981 +53.76735680353869 ``` """ function get_loglikelihood(𝓂::ℳ, @@ -3589,12 +4040,18 @@ function get_loglikelihood(𝓂::ℳ, initial_covariance::Symbol = :theoretical, filter_algorithm::Symbol = :LagrangeNewton, tol::Tolerances = Tolerances(), - quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, + quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_SELECTOR(𝓂), lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), - verbose::Bool = DEFAULT_VERBOSE)::S where {S <: Real, U <: AbstractFloat} + verbose::Bool = DEFAULT_VERBOSE, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES)::S where {S <: Real, U <: AbstractFloat} # timer::TimerOutput = TimerOutput(), + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm² = isa(sylvester_algorithm, Symbol) ? sylvester_algorithm : sylvester_algorithm[1], @@ -3612,26 +4069,27 @@ 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") + if !use_workspaces; 𝓂.workspaces = orig_ws; end 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 @@ -3642,37 +4100,194 @@ function get_loglikelihood(𝓂::ℳ, if !solved # println("Main call: 1st order solution not found") + if !use_workspaces; 𝓂.workspaces = orig_ws; end return on_failure_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 + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return llh end +function check_bounds(parameter_values::Vector{S}, 𝓂::ℳ)::Bool where S <: Real + if !all(isfinite,parameter_values) return true end + + 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 min(max(parameter_values[indexin([k], 𝓂.constants.post_complete_parameters.parameters)][1], v[1]), v[2]) != parameter_values[indexin([k], 𝓂.constants.post_complete_parameters.parameters)][1] + return true + end + end + end + end + + return false +end + + +function get_relevant_steady_state_and_state_update(::Val{:second_order}, + parameter_values::Vector{S}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false)::Tuple{constants, Vector{S}, Vector{AbstractMatrix{S}}, Vector{S}, Bool} where S <: Real + # timer::TimerOutput = TimerOutput(), + 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 opts.verbose println("Could not find 2nd order stochastic steady state") end + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], collect(sss), converged + end + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + all_SS = expand_steady_state(SS_and_pars, ms) + + state = collect(sss) - all_SS + + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged +end + + + +function get_relevant_steady_state_and_state_update(::Val{:pruned_second_order}, + parameter_values::Vector{S}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false)::Tuple{constants, Vector{S}, Vector{AbstractMatrix{S}}, Vector{Vector{S}}, Bool} where S <: Real + # timer::TimerOutput = TimerOutput(), + 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 opts.verbose println("Could not find 2nd order stochastic steady state") end + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], [zeros(S, 𝓂.constants.post_model_macro.nVars), zeros(S, 𝓂.constants.post_model_macro.nVars)], converged + end + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + all_SS = expand_steady_state(SS_and_pars, ms) + + state = [zeros(S, 𝓂.constants.post_model_macro.nVars), collect(sss)::Vector{S} - all_SS] + + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged +end + + + +function get_relevant_steady_state_and_state_update(::Val{:third_order}, + parameter_values::Vector{S}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false)::Tuple{constants, Vector{S}, Vector{AbstractMatrix{S}}, Vector{S}, Bool} where S <: Real + # timer::TimerOutput = TimerOutput(), + 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 opts.verbose println("Could not find 3rd order stochastic steady state") end + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], collect(sss), converged + end + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + all_SS = expand_steady_state(SS_and_pars, ms) + + state = collect(sss) - all_SS + + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged +end + + + +function get_relevant_steady_state_and_state_update(::Val{:pruned_third_order}, + parameter_values::Vector{S}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false)::Tuple{constants, Vector{S}, Vector{AbstractMatrix{S}}, Vector{Vector{S}}, Bool} where S <: Real + # timer::TimerOutput = TimerOutput(), + 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 opts.verbose println("Could not find 3rd order stochastic steady state") end + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], [zeros(S, 𝓂.constants.post_model_macro.nVars), zeros(S, 𝓂.constants.post_model_macro.nVars), zeros(S, 𝓂.constants.post_model_macro.nVars)], converged + end + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + all_SS = expand_steady_state(SS_and_pars, ms) + + state = [zeros(S, 𝓂.constants.post_model_macro.nVars), collect(sss)::Vector{S} - all_SS, zeros(S, 𝓂.constants.post_model_macro.nVars)] + + return 𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged +end + + +function get_relevant_steady_state_and_state_update(::Val{:first_order}, + parameter_values::Vector{S}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false)::Tuple{constants, Vector{S}, Matrix{S}, Vector{Vector{Float64}}, Bool} where S <: Real + # timer::TimerOutput = TimerOutput(), + # 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, + + 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 + # println("NSSS not found") + 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, 𝓂.workspaces) # , timer = timer)# |> Matrix + + 𝐒₁, qme_sol, solved = calculate_first_order_solution(∇₁, + constants_obj, + 𝓂.workspaces, + 𝓂.caches; + opts = opts, + initial_guess = 𝓂.caches.qme_solution, + parameter_values = parameter_values) + + + update_perturbation_counter!(𝓂.counters, solved, estimation = estimation, order = 1) + + if !solved + # println("NSSS not found") + return 𝓂.constants, SS_and_pars, zeros(S, 0, 0), [state], solved + end + + return 𝓂.constants, SS_and_pars, 𝐒₁, [state], solved +end + """ $(SIGNATURES) @@ -3692,7 +4307,7 @@ Calculate the residuals of the non-stochastic steady state equations of the mode - `KeyedArray` (from the `AxisKeys` package) containing the absolute values of the residuals of the non-stochastic steady state equations. # Examples -```jldoctest +```jldoctest; filter = r"(Equation|CalibrationEquation)([^0-9+-]+)\\S+" => s"\\1\\2 0.0" using MacroModelling @model RBC begin @@ -3702,7 +4317,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true begin std_z = 0.01 ρ = 0.2 δ = 0.02 @@ -3722,27 +4337,35 @@ And data, 5-element Vector{Float64}: (:Equation₃) 0.0 (:Equation₄) 0.0 (:CalibrationEquation₁) 0.0 +``` +Passing approximate values returns the residuals at those values: +```julia get_non_stochastic_steady_state_residuals(RBC, [1.1641597, 3.0635781, 1.2254312, 0.0, 0.18157895]) -# output -1-dimensional KeyedArray(NamedDimsArray(...)) with keys: -↓ Equation ∈ 5-element Vector{Symbol} -And data, 5-element Vector{Float64}: - (:Equation₁) 2.7360991250446887e-10 - (:Equation₂) 6.199999980083248e-8 - (:Equation₃) 2.7897102183871425e-8 - (:Equation₄) 0.0 - (:CalibrationEquation₁) 8.160392850342646e-8 +# 1-dimensional KeyedArray(NamedDimsArray(...)) with keys: +# ↓ Equation ∈ 5-element Vector{Symbol} +# And data, 5-element Vector{Float64}: +# (:Equation₁) 2.7360991250446887e-10 +# (:Equation₂) 6.199999980083248e-8 +# (:Equation₃) 2.7897102183871425e-8 +# (:Equation₄) 0.0 +# (:CalibrationEquation₁) 8.160392850342646e-8 ``` """ -function get_non_stochastic_steady_state_residuals(𝓂::ℳ, +@unstable function get_non_stochastic_steady_state_residuals(𝓂::ℳ, values::Union{Vector{Float64}, Dict{Symbol, Float64}, Dict{String, Float64}, KeyedArray{Float64, 1}}; parameters::ParameterType = nothing, steady_state_function::SteadyStateFunctionType = missing, tol::Tolerances = Tolerances(), - verbose::Bool = DEFAULT_VERBOSE) + verbose::Bool = DEFAULT_VERBOSE, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) # @nospecialize # reduce compile time + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + opts = merge_calculation_options(tol = tol, verbose = verbose) solve!(𝓂, @@ -3773,7 +4396,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 @@ -3789,15 +4412,19 @@ function get_non_stochastic_steady_state_residuals(𝓂::ℳ, 𝓂.functions.NSSS_check(residual, 𝓂.parameter_values, vals) + if !use_workspaces; 𝓂.workspaces = orig_ws; end + KeyedArray(abs.(residual), Equation = axis1) end """ See [`get_non_stochastic_steady_state_residuals`](@ref) """ -get_residuals = get_non_stochastic_steady_state_residuals +@unstable get_residuals = get_non_stochastic_steady_state_residuals """ See [`get_non_stochastic_steady_state_residuals`](@ref) """ check_residuals = get_non_stochastic_steady_state_residuals + +end # @stable diff --git a/src/impulse_response_function.jl b/src/impulse_response_function.jl new file mode 100644 index 000000000..7c9a20b4b --- /dev/null +++ b/src/impulse_response_function.jl @@ -0,0 +1,712 @@ +@stable default_mode = "disable" begin + + +@unstable function compute_irf_responses(𝓂::ℳ, + state_update::Function, + initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, + level::Vector{Float64}; + periods::Int, + shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}}, + variables::Union{Symbol_input,String_input}, + shock_size::Real, + negative_shock::Bool, + generalised_irf::Bool, + generalised_irf_warmup_iterations::Int, + generalised_irf_draws::Int, + enforce_obc::Bool, + algorithm::Symbol) + + if enforce_obc + obc_update = (present_states, present_shocks, state_update) -> obc_state_update(present_states, present_shocks, state_update, 𝓂, algorithm) + + if generalised_irf + return girf(state_update, + obc_update, + initial_state, + level, + 𝓂.constants; + periods = periods, + shocks = shocks, + shock_size = shock_size, + variables = variables, + negative_shock = negative_shock, + warmup_periods = generalised_irf_warmup_iterations, + draws = generalised_irf_draws) + else + return irf(state_update, + obc_update, + initial_state, + level, + 𝓂.constants; + periods = periods, + shocks = shocks, + shock_size = shock_size, + variables = variables, + negative_shock = negative_shock) + end + else + if generalised_irf + return girf(state_update, + initial_state, + level, + 𝓂.constants; + periods = periods, + shocks = shocks, + shock_size = shock_size, + variables = variables, + negative_shock = negative_shock, + warmup_periods = generalised_irf_warmup_iterations, + draws = generalised_irf_draws) + else + return irf(state_update, + initial_state, + level, + 𝓂.constants; + periods = periods, + shocks = shocks, + shock_size = shock_size, + variables = variables, + negative_shock = negative_shock) + end + end +end + + +function irf(state_update::Function, + obc_state_update::Function, + initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, + level::Vector{Float64}, + constants::constants; + periods::Int = 40, + shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, + variables::Union{Symbol_input,String_input} = :all, + shock_size::Real = 1, + negative_shock::Bool = false)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} + T = constants.post_model_macro + + pruning = initial_state isa Vector{Vector{Float64}} + + shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks + + shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks + + if shocks isa Matrix{Float64} + @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model." + + # periods += size(shocks)[2] + + shock_history = zeros(T.nExo, periods) + + shock_history[:,1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + elseif shocks isa KeyedArray{Float64} + shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) + + # periods += size(shocks)[2] + + @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." + + shock_history = zeros(T.nExo, periods) + + shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + else + raw_idx = parse_shocks_input_to_index(shocks,constants) + shock_idx = raw_idx isa Integer ? Int[raw_idx] : collect(Int, raw_idx) + shock_history = zeros(T.nExo, periods) + end + + var_idx = parse_variables_input_to_index(variables, constants) |> sort + + axis1 = T.var[var_idx] + + if any(x -> contains(string(x), "◖"), axis1) + axis1_decomposed = decompose_name.(axis1) + 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 + + always_solved = true + + if shocks == :simulate + shock_history = randn(T.nExo,periods) * shock_size + + shock_history[contains.(string.(T.exo),"ᵒᵇᶜ"),:] .= 0 + + Y = zeros(T.nVars,periods,1) + + past_states = initial_state + + for t in 1:periods + past_states, past_shocks, solved = obc_state_update(past_states, shock_history[:,t], state_update) + + if !solved @warn "No solution in period: $t" end#. Possible reasons: 1. infeasability 2. too long spell of binding constraint. To address the latter try setting max_obc_horizon to a larger value (default: 40): @model max_obc_horizon=40 begin ... end" end + + always_solved = always_solved && solved + + if !always_solved break end + + Y[:,t,1] = pruning ? sum(past_states) : past_states + + shock_history[:,t] = past_shocks + end + + return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:simulate]) + elseif shocks == :none + Y = zeros(T.nVars,periods,1) + + shck = T.nExo == 0 ? Vector{Float64}(undef, 0) : zeros(T.nExo) + + past_states = initial_state + + for t in 1:periods + past_states, _, solved = obc_state_update(past_states, shck, state_update) + + if !solved @warn "No solution in period: $t" end#. Possible reasons: 1. infeasability 2. too long spell of binding constraint. To address the latter try setting max_obc_horizon to a larger value (default: 40): @model max_obc_horizon=40 begin ... end" end + + always_solved = always_solved && solved + + if !always_solved break end + + Y[:,t,1] = pruning ? sum(past_states) : past_states + end + + return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:none]) + else + Y = zeros(T.nVars,periods,length(shock_idx)) + + for (i,ii) in enumerate(shock_idx) + if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input,String_input} + shock_history = zeros(T.nExo,periods) + shock_history[ii,1] = negative_shock ? -shock_size : shock_size + end + + past_states = initial_state + + for t in 1:periods + past_states, past_shocks, solved = obc_state_update(past_states, shock_history[:,t], state_update) + + if !solved @warn "No solution in period: $t" end#. Possible reasons: 1. infeasability 2. too long spell of binding constraint. To address the latter try setting max_obc_horizon to a larger value (default: 40): @model max_obc_horizon=40 begin ... end" end + + always_solved = always_solved && solved + + if !always_solved break end + + Y[:,t,i] = pruning ? sum(past_states) : past_states + + shock_history[:,t] = past_shocks + end + end + + axis2 = shocks isa Union{Symbol_input,String_input} ? T.exo[shock_idx] : [:Shock_matrix] + + if any(x -> contains(string(x), "◖"), axis2) + axis2_decomposed = decompose_name.(axis2) + axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] + end + + return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) + end +end + + + + +@unstable function irf(state_update::Function, + initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, + level::Vector{Float64}, + constants::constants; + periods::Int = 40, + shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, + variables::Union{Symbol_input,String_input} = :all, + shock_size::Real = 1, + negative_shock::Bool = false)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} + T = constants.post_model_macro + + pruning = initial_state isa Vector{Vector{Float64}} + + shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks + + shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks + + if shocks isa Matrix{Float64} + @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model." + + # periods += size(shocks)[2] + + shock_history = zeros(T.nExo, periods) + + shock_history[:,1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + elseif shocks isa KeyedArray{Float64} + shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) + + # periods += size(shocks)[2] + + @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." + + shock_history = zeros(T.nExo, periods) + + shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + else + raw_idx = parse_shocks_input_to_index(shocks,constants) + shock_idx = raw_idx isa Integer ? Int[raw_idx] : collect(Int, raw_idx) + shock_history = zeros(T.nExo, periods) + end + + var_idx = parse_variables_input_to_index(variables, constants) |> sort + + axis1 = T.var[var_idx] + + if any(x -> contains(string(x), "◖"), axis1) + axis1_decomposed = decompose_name.(axis1) + 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 + + if shocks == :simulate + shock_history = randn(T.nExo,periods) * shock_size + + shock_history[contains.(string.(T.exo),"ᵒᵇᶜ"),:] .= 0 + + Y = zeros(T.nVars,periods,1) + + initial_state = state_update(initial_state,shock_history[:,1]) + + Y[:,1,1] = pruning ? sum(initial_state) : initial_state + + for t in 1:periods-1 + initial_state = state_update(initial_state,shock_history[:,t+1]) + + Y[:,t+1,1] = pruning ? sum(initial_state) : initial_state + end + + return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:simulate]) + elseif shocks == :none + Y = zeros(T.nVars,periods,1) + + shck = T.nExo == 0 ? Vector{Float64}(undef, 0) : zeros(T.nExo) + + initial_state = state_update(initial_state, shck) + + Y[:,1,1] = pruning ? sum(initial_state) : initial_state + + for t in 1:periods-1 + initial_state = state_update(initial_state, shck) + + Y[:,t+1,1] = pruning ? sum(initial_state) : initial_state + end + + return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = [:none]) + else + Y = zeros(T.nVars,periods,length(shock_idx)) + + for (i,ii) in enumerate(shock_idx) + initial_state_copy = deepcopy(initial_state) + + if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input,String_input} + shock_history = zeros(T.nExo,periods) + shock_history[ii,1] = negative_shock ? -shock_size : shock_size + end + + initial_state_copy = state_update(initial_state_copy, shock_history[:,1]) + + Y[:,1,i] = pruning ? sum(initial_state_copy) : initial_state_copy + + for t in 1:periods-1 + initial_state_copy = state_update(initial_state_copy, shock_history[:,t+1]) + + Y[:,t+1,i] = pruning ? sum(initial_state_copy) : initial_state_copy + end + end + + axis2 = shocks isa Union{Symbol_input,String_input} ? T.exo[shock_idx] : [:Shock_matrix] + + if any(x -> contains(string(x), "◖"), axis2) + axis2_decomposed = decompose_name.(axis2) + axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] + end + + return KeyedArray(Y[var_idx,:,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) + end +end + + + +function girf(state_update::Function, + initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, + level::Vector{Float64}, + constants::constants; + periods::Int = 40, + shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, + variables::Union{Symbol_input,String_input} = :all, + shock_size::Real = 1, + negative_shock::Bool = false, + warmup_periods::Int = 100, + draws::Int = 50)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} + T = constants.post_model_macro + + pruning = initial_state isa Vector{Vector{Float64}} + + shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks + + shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks + + if shocks isa Matrix{Float64} + @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model (model has $(T.nExo) shocks)." + + # periods += size(shocks)[2] + + shock_history = zeros(T.nExo, periods) + + shock_history[:,1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + elseif shocks isa KeyedArray{Float64} + shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) + + # periods += size(shocks)[2] + + @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." + + shock_history = zeros(T.nExo, periods + 1) + + shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + elseif shocks == :simulate + shock_history = randn(T.nExo,periods) * shock_size + + shock_idx = Int[1] + else + raw_idx = parse_shocks_input_to_index(shocks,constants) + shock_idx = raw_idx isa Integer ? Int[raw_idx] : collect(Int, raw_idx) + shock_history = zeros(T.nExo, periods) + end + + var_idx = parse_variables_input_to_index(variables, constants) |> sort + + Y = zeros(T.nVars, periods + 1, length(shock_idx)) + + for (i,ii) in enumerate(shock_idx) + initial_state_copy = deepcopy(initial_state) + + accepted_draws = 0 + + for draw in 1:draws + ok = true + + initial_state_copy² = deepcopy(initial_state_copy) + + for i in 1:warmup_periods + initial_state_copy² = state_update(initial_state_copy², randn(T.nExo)) + if any(!isfinite, [x for v in initial_state_copy² for x in v]) + # @warn "No solution in warmup period: $i" + ok = false + break + end + end + + if !ok continue end + + Y₁ = zeros(T.nVars, periods + 1) + Y₂ = zeros(T.nVars, periods + 1) + + baseline_noise = randn(T.nExo) + + if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input,String_input} + shock_history = zeros(T.nExo,periods) + shock_history[ii,1] = negative_shock ? -shock_size : shock_size + end + + initial_state₁ = initial_state_copy + initial_state₂ = initial_state_copy + + if pruning + initial_state_copy² = state_update(initial_state_copy², baseline_noise) + + initial_state₁ = deepcopy(initial_state_copy²) + initial_state₂ = deepcopy(initial_state_copy²) + + if any(!isfinite, [x for v in initial_state_copy² for x in v]) continue end + + Y₁[:,1] = initial_state_copy² |> sum + Y₂[:,1] = initial_state_copy² |> sum + else + Y₁[:,1] = state_update(initial_state_copy², baseline_noise) + + if any(!isfinite, Y₁[:,1]) continue end + + Y₂[:,1] = state_update(initial_state_copy², baseline_noise) + + if any(!isfinite, Y₂[:,1]) continue end + end + + for t in 1:periods + baseline_noise = randn(T.nExo) + + if pruning + initial_state₁ = state_update(initial_state₁, baseline_noise) + + if any(!isfinite, [x for v in initial_state₁ for x in v]) + ok = false + break + end + + initial_state₂ = state_update(initial_state₂, baseline_noise + shock_history[:,t]) + + if any(!isfinite, [x for v in initial_state₂ for x in v]) + ok = false + break + end + + Y₁[:,t+1] = initial_state₁ |> sum + Y₂[:,t+1] = initial_state₂ |> sum + else + Y₁[:,t+1] = state_update(Y₁[:,t],baseline_noise) + + if any(!isfinite, Y₁[:,t+1]) + ok = false + break + end + + Y₂[:,t+1] = state_update(Y₂[:,t],baseline_noise + shock_history[:,t]) + + if any(!isfinite, Y₂[:,t+1]) + ok = false + break + end + end + end + + if !ok continue end + + Y[:,:,i] += Y₂ - Y₁ + + accepted_draws += 1 + end + + if accepted_draws == 0 + @warn "No draws accepted. Results are empty." + elseif accepted_draws < draws + # average over accepted draws, if desired + @info "$accepted_draws of $draws draws accepted for shock: $(shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input, String_input} ? T.exo[ii] : :Shock_matrix)" + Y[:, :, i] ./= accepted_draws + else + Y[:, :, i] ./= accepted_draws + end + end + + axis1 = T.var[var_idx] + + if any(x -> contains(string(x), "◖"), axis1) + axis1_decomposed = decompose_name.(axis1) + axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] + end + + axis2 = shocks isa Union{Symbol_input,String_input} ? T.exo[shock_idx] : [:Shock_matrix] + + if any(x -> contains(string(x), "◖"), axis2) + axis2_decomposed = decompose_name.(axis2) + axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] + end + + return KeyedArray(Y[var_idx,2:end,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) +end + + +function girf(state_update::Function, + obc_state_update::Function, + initial_state::Union{Vector{Vector{Float64}},Vector{Float64}}, + level::Vector{Float64}, + constants::constants; + periods::Int = 40, + shocks::Union{Symbol_input,String_input,Matrix{Float64},KeyedArray{Float64}} = :all, + variables::Union{Symbol_input,String_input} = :all, + shock_size::Real = 1, + negative_shock::Bool = false, + warmup_periods::Int = 100, + draws::Int = 50)::Union{KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{String}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{String},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{Symbol}}}, KeyedArray{Float64, 3, NamedDimsArray{(:Variables, :Periods, :Shocks), Float64, 3, Array{Float64, 3}}, Tuple{Vector{Symbol},UnitRange{Int},Vector{String}}}} + T = constants.post_model_macro + + pruning = initial_state isa Vector{Vector{Float64}} + + shocks = shocks isa KeyedArray ? axiskeys(shocks,1) isa Vector{String} ? rekey(shocks, 1 => axiskeys(shocks,1) .|> Meta.parse .|> replace_indices) : shocks : shocks + + shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks + + if shocks isa Matrix{Float64} + @assert size(shocks)[1] == T.nExo "Number of rows of provided shock matrix does not correspond to number of shocks. Please provide matrix with as many rows as there are shocks in the model." + + # periods += size(shocks)[2] + + shock_history = zeros(T.nExo, periods) + + shock_history[:,1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + elseif shocks isa KeyedArray{Float64} + shock_input = map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),axiskeys(shocks)[1]) + + # periods += size(shocks)[2] + + @assert length(setdiff(shock_input, T.exo)) == 0 "Provided shocks are not part of the model. Use `get_shocks(𝓂)` to list valid shock names." + + shock_history = zeros(T.nExo, periods + 1) + + shock_history[indexin(shock_input,T.exo),1:size(shocks)[2]] = shocks + + shock_idx = Int[1] + elseif shocks == :simulate + shock_history = randn(T.nExo,periods) * shock_size + + shock_history[contains.(string.(T.exo),"ᵒᵇᶜ"),:] .= 0 + + shock_idx = Int[1] + else + raw_idx = parse_shocks_input_to_index(shocks,constants) + shock_idx = raw_idx isa Integer ? Int[raw_idx] : collect(Int, raw_idx) + shock_history = zeros(T.nExo, periods) + end + + var_idx = parse_variables_input_to_index(variables, constants) |> sort + + Y = zeros(T.nVars, periods + 1, length(shock_idx)) + + for (i,ii) in enumerate(shock_idx) + initial_state_copy = deepcopy(initial_state) + + accepted_draws = 0 + + for draw in 1:draws + ok = true + + initial_state_copy² = deepcopy(initial_state_copy) + + warmup_shocks = randn(T.nExo) + warmup_shocks[contains.(string.(T.exo), "ᵒᵇᶜ")] .= 0 + + # --- warmup --- + for i_w in 1:warmup_periods + initial_state_copy², _, solved = obc_state_update(initial_state_copy², warmup_shocks, state_update) + if !solved + # @warn "No solution in warmup period: $i_w" + ok = false + break + end + end + + if !ok continue end + + Y₁ = zeros(T.nVars, periods + 1) + Y₂ = zeros(T.nVars, periods + 1) + + baseline_noise = randn(T.nExo) + baseline_noise[contains.(string.(T.exo), "ᵒᵇᶜ")] .= 0 + + if shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input, String_input} + shock_history = zeros(T.nExo, periods) + shock_history[ii, 1] = negative_shock ? -shock_size : shock_size + end + + initial_state₁ = initial_state_copy + initial_state₂ = initial_state_copy + + # --- period 1 --- + if pruning + initial_state_copy², _, solved = obc_state_update(initial_state_copy², baseline_noise, state_update) + + initial_state₁ = deepcopy(initial_state_copy²) + initial_state₂ = deepcopy(initial_state_copy²) + + if !solved continue end + + Y₁[:, 1] = initial_state_copy² |> sum + Y₂[:, 1] = initial_state_copy² |> sum + else + Y₁[:, 1], _, solved = obc_state_update(initial_state_copy², baseline_noise, state_update) + if !solved continue end + + Y₂[:, 1], _, solved = obc_state_update(initial_state_copy², baseline_noise, state_update) + if !solved continue end + end + + # --- remaining periods --- + for t in 1:periods + baseline_noise = randn(T.nExo) + baseline_noise[contains.(string.(T.exo), "ᵒᵇᶜ")] .= 0 + + if pruning + initial_state₁, _, solved = obc_state_update(initial_state₁, baseline_noise, state_update) + if !solved + # @warn "No solution in period: $t" + ok = false + break + end + + initial_state₂, _, solved = obc_state_update(initial_state₂, baseline_noise + shock_history[:, t], state_update) + if !solved + # @warn "No solution in period: $t" + ok = false + break + end + + Y₁[:, t + 1] = initial_state₁ |> sum + Y₂[:, t + 1] = initial_state₂ |> sum + else + Y₁[:, t + 1], _, solved = obc_state_update(Y₁[:, t], baseline_noise, state_update) + if !solved + # @warn "No solution in period: $t" + ok = false + break + end + + Y₂[:, t + 1], _, solved = obc_state_update(Y₂[:, t], baseline_noise + shock_history[:, t], state_update) + if !solved + # @warn "No solution in period: $t" + ok = false + break + end + end + end + + if !ok continue end + + # Note: replace `i` if your outer scope uses another index + Y[:, :, i] .+= (Y₂ .- Y₁) + accepted_draws += 1 + end + + if accepted_draws == 0 + @warn "No draws accepted. Results are empty." + elseif accepted_draws < draws + # average over accepted draws, if desired + @info "$accepted_draws of $draws draws accepted for shock: $(shocks ∉ [:simulate, :none] && shocks isa Union{Symbol_input, String_input} ? T.exo[ii] : :Shock_matrix)" + Y[:, :, i] ./= accepted_draws + else + Y[:, :, i] ./= accepted_draws + end + end + + axis1 = T.var[var_idx] + + if any(x -> contains(string(x), "◖"), axis1) + axis1_decomposed = decompose_name.(axis1) + axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed] + end + + axis2 = shocks isa Union{Symbol_input,String_input} ? T.exo[shock_idx] : [:Shock_matrix] + + if any(x -> contains(string(x), "◖"), axis2) + axis2_decomposed = decompose_name.(axis2) + axis2 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis2_decomposed] + end + + return KeyedArray(Y[var_idx,2:end,:] .+ level[var_idx]; Variables = axis1, Periods = 1:periods, Shocks = axis2) +end + + +end # @stable diff --git a/src/inspect.jl b/src/inspect.jl index 0fd7ae6b3..02f2fd96c 100644 --- a/src/inspect.jl +++ b/src/inspect.jl @@ -1,5 +1,6 @@ @stable default_mode = "disable" begin + get_symbols(ex::Symbol) = [ex] get_symbols(ex::Real) = [ex] @@ -27,7 +28,6 @@ function get_symbols(ex::Expr) return par end -end # dispatch_doctor """ replace_curly_braces_in_symbols(expr) -> Union{Expr, Symbol, Any} @@ -85,7 +85,7 @@ end Replace timing subscripts (₍₋₁₎, ₍₀₎, ₍₁₎, ₍ₓ₎) with bracket notation and convert ◖/◗ to curly braces. Transforms symbols like `Symbol("z◖TFP◗₍₀₎")` into `Expr(:ref, Expr(:curly, :z, :TFP), 0)`. """ -function replace_dynamic_symbols(expr) +@unstable function replace_dynamic_symbols(expr) if expr isa Symbol str = string(expr) # First replace timing subscripts @@ -108,14 +108,13 @@ Parse a filter term into (base_symbol, pattern_expr). - `"k[-1]"` → `(:k, :(k[-1]))` matches exact timing - `"eps[x]"` → `(:eps, :(eps[x]))` matches shock at exact timing """ -function parse_filter_term(term::Union{Symbol, String}) +@unstable function parse_filter_term(term::Union{Symbol, String}) term_str = replace(string(term), "{" => "◖", "}" => "◗") m = match(r"^(.+)\[(.+)\]$", term_str) m === nothing && return (Symbol(term_str), nothing) return (Symbol(m.captures[1]), Meta.parse(term_str)) end -@stable default_mode = "disable" begin """ expr_contains(expr, sym::Symbol, pattern) -> Bool @@ -173,7 +172,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -240,7 +239,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -256,8 +255,8 @@ get_steady_state_equations(RBC) :((-β * ((k ^ (α - 1) * α * exp(z{TFP}) - δ * exp(z{δ})) + 1)) / c + 1 / c) :(((c - k * (-δ * exp(z{δ}) + 1)) + k) - q) :(-(k ^ α) * exp(z{TFP}) + q) - :(-z{TFP} * ρ{TFP} + z{TFP}) - :(-z{δ} * ρ{δ} + z{δ}) + :(-(z{TFP}) * ρ{TFP} + z{TFP}) + :(-(z{δ}) * ρ{δ} + z{δ}) :(➕₁ - c / q) :(➕₂ - c / q) :((Δc_share - log(➕₁)) + log(➕₂)) @@ -316,7 +315,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -440,7 +439,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -508,7 +507,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -566,7 +565,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -584,6 +583,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 @@ -612,7 +612,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC_incomplete begin +@parameters RBC_incomplete silent = true report_missing_parameters = false begin std_z = 0.01 ρ = 0.2 # Note: α, β, δ are not defined @@ -652,7 +652,7 @@ using MacroModelling z[0] = ρ * z[-1] + std_z * eps_z[x] end -@parameters RBC begin +@parameters RBC silent = true report_missing_parameters = false begin std_z = 0.01 ρ = 0.2 end @@ -694,7 +694,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -746,7 +746,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -792,7 +792,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -838,7 +838,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -886,7 +886,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -940,7 +940,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -952,9 +952,10 @@ end get_nonnegativity_auxiliary_variables(RBC) # output -2-element Vector{String}: +3-element Vector{String}: "➕₁" "➕₂" + "➕₃" ``` """ function get_nonnegativity_auxiliary_variables(𝓂::ℳ)::Vector{String} @@ -989,7 +990,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -1040,7 +1041,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -1093,7 +1094,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -1151,7 +1152,7 @@ using MacroModelling Δk_4q[0] = log(k[0]) - log(k[-4]) end -@parameters RBC begin +@parameters RBC silent = true begin σ = 0.01 ρ = 0.2 capital_to_output = 1.5 @@ -1173,4 +1174,5 @@ function get_jump_variables(𝓂::ℳ)::Vector{String} 𝓂.constants.post_model_macro.future_not_past_and_mixed |> collect |> sort .|> x -> replace.(string.(x), "◖" => "{", "◗" => "}") end -end # dispatch_doctor \ No newline at end of file + +end # @stable diff --git a/src/moments.jl b/src/moments.jl index 46cf36d19..3cc2368f7 100644 --- a/src/moments.jl +++ b/src/moments.jl @@ -1,5 +1,139 @@ @stable default_mode = "disable" begin + +""" + sparse_ABAt(A::SparseMatrixCSC{T}, B::SparseMatrixCSC{T}; + tol::Real = eps(T)) -> SparseMatrixCSC{T} + +Compute `C = A * B * A'` returning a sparse symmetric matrix, where `B` is +symmetric. Uses a column-by-column sparse-accumulator (SPA) kernel that +computes only the lower triangle with advancing-pointer row pruning, then +emits both `(i,j)` and `(j,i)` entries into COO vectors for direct CSC +assembly. + +`tol` controls the drop tolerance during assembly (entries with `|v| < tol` +are discarded). +""" +function sparse_ABAt(A::SparseMatrixCSC{T}, B::SparseMatrixCSC{T}; + tol::Real = eps(T)) where T <: Real + m, n = size(A) + @assert size(B) == (n, n) "B must be n×n where A is m×n" + + # Linked-list row index for A: avoids allocating sparse(A'). + # row_head[row] -> first nz index in that row; row_next[nz] -> next nz in same row; + # row_col[nz] -> column of that nonzero entry. + A_rv = SparseArrays.rowvals(A) + A_nz = nonzeros(A) + A_cp = SparseArrays.getcolptr(A) + nnzA = nnz(A) + row_head = zeros(Int, m) + row_next = zeros(Int, nnzA) + row_col = Vector{Int}(undef, nnzA) + @inbounds for col in n:-1:1 + for idx in A_cp[col]:(A_cp[col + 1] - 1) + row = A_rv[idx] + row_next[idx] = row_head[row] + row_head[row] = idx + row_col[idx] = col + end + end + + B_rows = SparseArrays.rowvals(B) + B_vals = nonzeros(B) + + A_rows = A_rv + A_vals = A_nz + A_colptr = A_cp + + # SPA workspace: generation-marker pattern avoids zeroing w each column + w = Vector{T}(undef, n) + mark_w = zeros(Int, n) + w_nz = Vector{Int}(undef, n) + + # Dense accumulator for output column (lower triangle only) + c = zeros(T, m) + + # Advancing pointers: A_start[k] tracks first unconsumed position in A[:,k]. + # Monotonic j means pointers only advance; total work bounded by nnz(A). + A_start = A_colptr[1:n] + + # COO output: emit both (i,j) and (j,i) for off-diagonal entries. + # Upper bound: m*(m+1)/2 lower-triangle entries + m*(m-1)/2 mirrors = m² + max_entries = m * m + coo_I = Vector{Int}(undef, max_entries) + coo_J = Vector{Int}(undef, max_entries) + coo_V = Vector{T}(undef, max_entries) + cnt = 0 + + @inbounds for j in 1:m + row_head[j] == 0 && continue + + # Phase A: w = B * Aᵀ[:,j] via SPA (iterate row j of A via linked-list) + w_cnt = 0 + p = row_head[j] + while p != 0 + l = row_col[p] + a_jl = A_nz[p] + for q in SparseArrays.nzrange(B, l) + k = B_rows[q] + b_kl = B_vals[q] + if mark_w[k] != j + mark_w[k] = j + w_cnt += 1 + w_nz[w_cnt] = k + w[k] = b_kl * a_jl + else + w[k] += b_kl * a_jl + end + end + p = row_next[p] + end + + # Sort for sequential A-column access (cache-friendly advancing pointers) + sort!(view(w_nz, 1:w_cnt)) + + # Phase B: c[i] += (A * w)[i] for i ≥ j (lower triangle only) + for idx in 1:w_cnt + k = w_nz[idx] + wk = w[k] + hi_k = A_colptr[k + 1] - 1 + p = A_start[k] + while p ≤ hi_k && A_rows[p] < j + p += 1 + end + A_start[k] = p + for q in p:hi_k + c[A_rows[q]] += A_vals[q] * wk + end + end + + # Phase C: gather nonzeros from c[j:m], emit both (i,j) and (j,i) + for i in j:m + val = c[i] + c[i] = zero(T) + if abs(val) >= tol + cnt += 1 + coo_I[cnt] = i + coo_J[cnt] = j + coo_V[cnt] = val + if i != j + cnt += 1 + coo_I[cnt] = j + coo_J[cnt] = i + coo_V[cnt] = val + end + end + end + end + + resize!(coo_I, cnt) + resize!(coo_J, cnt) + resize!(coo_V, cnt) + + return sparse(coo_I, coo_J, coo_V, m, m) +end + + function calculate_covariance(parameters::Vector{R}, 𝓂::ℳ; opts::CalculationOptions = merge_calculation_options())::Tuple{Matrix{R}, Matrix{R}, Matrix{R}, Vector{R}, Bool} where R <: Real @@ -8,28 +142,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,:] @@ -42,16 +173,42 @@ function calculate_covariance(parameters::Vector{R}, return CC, sol, ∇₁, SS_and_pars, solved end + # Check Lyapunov cache: if valid for current parameters, skip the solve + cached_covar = 𝓂.caches.covariance_first_order + if R === Float64 && !isempty(parameters) && + cache_valid_for_parameters(𝓂.caches.valid_for.covariance_first_order, parameters) && + !isempty(cached_covar) && size(cached_covar) == (T.nVars, T.nVars) + return cached_covar, sol, ∇₁, SS_and_pars, true + 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, + initial_guess = cached_covar, lyapunov_algorithm = opts.lyapunov_algorithm, - tol = opts.tol.lyapunov_tol, - acceptance_tol = opts.tol.lyapunov_acceptance_tol, - verbose = opts.verbose) + tol = opts.tol.first_order.lyapunov, + verbose = opts.verbose, + has_unit_roots = 𝓂.caches.has_unit_roots) + + # Safety net: if Lyapunov result contains NaN (deflation was used), + # ensure the flag is set for subsequent calls even if QME didn't detect it + if solved && any(isnan, covar_raw) + 𝓂.caches.has_unit_roots = true + end + + # Cache the result for reuse + if R === Float64 && solved + if size(𝓂.caches.covariance_first_order) != size(covar_raw) + 𝓂.caches.covariance_first_order = Matrix{Float64}(undef, size(covar_raw)...) + end + copyto!(𝓂.caches.covariance_first_order, covar_raw) + 𝓂.caches.valid_for.covariance_first_order = Float64.(parameters) + end + + covar_stable = covar_raw - return covar_raw, sol , ∇₁, SS_and_pars, solved + return covar_stable, sol , ∇₁, SS_and_pars, solved end @@ -68,52 +225,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 +346,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 +382,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 - ŝ_to_ŝ₂ = [ 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₁ ] - ê_to_ŝ₂ = [ 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₁] - ŝ_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₂ˢ] - ê_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 +449,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 +472,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 +512,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 - ŝ_to_ŝ₂ = [ 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₁ 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₁ ] - ê_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₁] + ê_to_ŝ₂ = [ 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₁] - ŝ_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₂ˢ] ê_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₂ @@ -396,14 +548,34 @@ function calculate_second_order_moments_with_covariance(parameters::Vector{R}, C = ê_to_ŝ₂ * Γ₂ * ê_to_ŝ₂' - # Ensure second-order lyapunov workspace and solve - lyap_ws_2nd = ensure_lyapunov_workspace!(𝓂.workspaces, size(ŝ_to_ŝ₂, 1), :second_order) - - Σᶻ₂, 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, - verbose = opts.verbose) + # Check 2nd-order Lyapunov cache + cached_covar_2nd = 𝓂.caches.covariance_second_order + n_ŝ₂ = size(ŝ_to_ŝ₂, 1) + if R === Float64 && !isempty(parameters) && + cache_valid_for_parameters(𝓂.caches.valid_for.covariance_second_order, parameters) && + !isempty(cached_covar_2nd) && size(cached_covar_2nd) == (n_ŝ₂, n_ŝ₂) + Σᶻ₂ = cached_covar_2nd + info = true + else + # Ensure second-order lyapunov workspace and solve + lyap_ws_2nd = ensure_lyapunov_workspace!(𝓂.workspaces, n_ŝ₂, :second_order) + + Σᶻ₂, info = solve_lyapunov_equation(ŝ_to_ŝ₂, C, lyap_ws_2nd, + initial_guess = cached_covar_2nd, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.second_order.lyapunov, + verbose = opts.verbose, + has_unit_roots = 𝓂.caches.has_unit_roots) + + # Cache the result for reuse + if R === Float64 && info + if size(𝓂.caches.covariance_second_order) != size(Σᶻ₂) + 𝓂.caches.covariance_second_order = Matrix{Float64}(undef, size(Σᶻ₂)...) + end + copyto!(𝓂.caches.covariance_second_order, Σᶻ₂) + 𝓂.caches.valid_for.covariance_second_order = Float64.(parameters) + end + end if info Σʸ₂ = ŝ_to_y₂ * Σᶻ₂ * ŝ_to_y₂' + ê_to_y₂ * Γ₂ * ê_to_y₂' @@ -412,9 +584,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) ŝ_to_ŝ₂ = zeros(R,0,0) @@ -422,9 +595,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) ŝ_to_ŝ₂ = zeros(R,0,0) @@ -432,9 +606,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) ŝ_to_ŝ₂ = zeros(R,0,0) @@ -444,43 +619,170 @@ 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, ŝ_to_ŝ₂, ŝ_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, ŝ_to_ŝ₂, ŝ_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 + # Check 3rd-order autocorrelation cache + nVars_ac = 𝓂.constants.post_model_macro.nVars + obs_key_ac = if observables == :full_covar + collect(1:nVars_ac) + else + obs_idx_ac = parse_variables_input_to_index(observables, 𝓂.constants) |> sort + if covariance == Symbol[] + collect(obs_idx_ac) + else + covar_idx_ac = parse_variables_input_to_index(covariance, 𝓂.constants) |> sort + sort(union(obs_idx_ac, covar_idx_ac)) + end + end + ac_periods = collect(Int, autocorrelation_periods) + cached_covar_ac = 𝓂.caches.covariance_third_order + cached_autocorr = 𝓂.caches.covariance_third_order_autocorr + if T === Float64 && !isempty(parameters) && + cache_valid_for_parameters(𝓂.caches.valid_for.covariance_third_order_autocorr, parameters) && + !isempty(cached_covar_ac) && size(cached_covar_ac) == (nVars_ac, nVars_ac) && + !isempty(cached_autocorr) && size(cached_autocorr, 1) == nVars_ac && + 𝓂.caches.valid_for.covariance_third_order_autocorr_obs_key == obs_key_ac && + 𝓂.caches.valid_for.covariance_third_order_autocorr_periods == ac_periods + return cached_covar_ac, μʸ₂, cached_autocorr, SS_and_pars, true + 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 +790,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 +818,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 +871,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 +902,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 +921,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 - ŝ_to_ŝ₃ = [ 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] + + 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ˢ^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₁)] + ê_to_ŝ₃ = [ 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₁)] - ŝ_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] + ŝ_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] @@ -646,29 +968,61 @@ 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!(ŝ_to_ŝ₃, eps()) - droptol!(ê_to_ŝ₃, eps()) + droptol!(A_UU, eps()) + droptol!(A_LU, eps()) + droptol!(A_LL, eps()) + droptol!(ê_to_ŝ₃, eps()) droptol!(Eᴸᶻ, eps()) droptol!(Γ₃, eps()) - - A = ê_to_ŝ₃ * Eᴸᶻ * ŝ_to_ŝ₃' - droptol!(A, eps()) - C = ê_to_ŝ₃ * Γ₃ * ê_to_ŝ₃' + 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' + eQ = ê_L * Q + C_LL = sparse_ABAt(ê_L, Γ₃) + eQ + eQ' + 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] + + A = ê_to_ŝ₃ * Eᴸᶻ * ŝ_to_ŝ₃' + droptol!(A, eps()) - # Ensure third-order lyapunov workspace and solve - lyap_ws_3rd = ensure_lyapunov_workspace!(𝓂.workspaces, size(ŝ_to_ŝ₃, 1), :third_order) + C = sparse_ABAt(ê_to_ŝ₃, Γ₃) + A + A' + droptol!(C, 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) + 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, + has_unit_roots = 𝓂.caches.has_unit_roots) + end if !info return zeros(T,0,0), zeros(T,0), zeros(T,0,0), zeros(T,0), false @@ -676,7 +1030,7 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T solved_lyapunov = solved_lyapunov && info - Σʸ₃tmp = ŝ_to_y₃ * Σᶻ₃ * ŝ_to_y₃' + ê_to_y₃ * Γ₃ * ê_to_y₃' + ê_to_y₃ * Eᴸᶻ * ŝ_to_y₃' + ŝ_to_y₃ * Eᴸᶻ' * ê_to_y₃' + Σʸ₃tmp = ŝ_to_y₃ * Σᶻ₃ * ŝ_to_y₃' + sparse_ABAt(ê_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,63 +1050,122 @@ function calculate_third_order_moments_with_autocorrelation(parameters::Vector{T Σᶻ₃ⁱ .= ŝ_to_ŝ₃ * Σᶻ₃ⁱ + ê_to_ŝ₃ * 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(ŝ_to_y₃ * Σᶻ₃ⁱ * ŝ_to_y₃' + ŝ_to_y₃ * ŝ_to_ŝ₃ⁱ * autocorr_tmp + ê_to_y₃ * Eᴸᶻ * ŝ_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 ŝ_to_ŝ₃ⁱ *= ŝ_to_ŝ₃ end + + end + + # Compute obs_key for cache storage + nVars_autocorr = 𝓂.constants.post_model_macro.nVars + obs_key_autocorr = if observables == :full_covar + collect(1:nVars_autocorr) + else + obs_idx = parse_variables_input_to_index(observables, 𝓂.constants) |> sort + if covariance == Symbol[] + collect(obs_idx) + else + covar_idx = parse_variables_input_to_index(covariance, 𝓂.constants) |> sort + sort(union(obs_idx, covar_idx)) + end end - return Σʸ₃, μʸ₂, autocorr, SS_and_pars, solved && solved3 && solved_lyapunov + # Cache the 3rd-order covariance for reuse (also benefits calculate_third_order_moments) + all_solved = solved && solved3 && solved_lyapunov + if T === Float64 && all_solved + if size(𝓂.caches.covariance_third_order) != size(Σʸ₃) + 𝓂.caches.covariance_third_order = Matrix{Float64}(undef, size(Σʸ₃)...) + end + copyto!(𝓂.caches.covariance_third_order, Σʸ₃) + 𝓂.caches.valid_for.covariance_third_order = Float64.(parameters) + 𝓂.caches.valid_for.covariance_third_order_obs_key = obs_key_autocorr + + # Cache autocorrelation + if size(𝓂.caches.covariance_third_order_autocorr) != size(autocorr) + 𝓂.caches.covariance_third_order_autocorr = Matrix{Float64}(undef, size(autocorr)...) + end + copyto!(𝓂.caches.covariance_third_order_autocorr, autocorr) + 𝓂.caches.valid_for.covariance_third_order_autocorr = Float64.(parameters) + 𝓂.caches.valid_for.covariance_third_order_autocorr_obs_key = obs_key_autocorr + 𝓂.caches.valid_for.covariance_third_order_autocorr_periods = collect(Int, autocorrelation_periods) + end + + return Σʸ₃, μʸ₂, autocorr, SS_and_pars, all_solved end 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, ŝ_to_ŝ₂, ŝ_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 + # Check 3rd-order covariance cache: if valid for current parameters AND same observables/covariance, skip + nVars_check = 𝓂.constants.post_model_macro.nVars + obs_key = if observables == :full_covar + collect(1:nVars_check) + else + obs_idx = parse_variables_input_to_index(observables, 𝓂.constants) |> sort + if covariance == Symbol[] + collect(obs_idx) + else + covar_idx = parse_variables_input_to_index(covariance, 𝓂.constants) |> sort + sort(union(obs_idx, covar_idx)) + end + end + cached_covar_3rd = 𝓂.caches.covariance_third_order + if T === Float64 && !isempty(parameters) && + cache_valid_for_parameters(𝓂.caches.valid_for.covariance_third_order, parameters) && + !isempty(cached_covar_3rd) && size(cached_covar_3rd) == (nVars_check, nVars_check) && + 𝓂.caches.valid_for.covariance_third_order_obs_key == obs_key + return cached_covar_3rd, μʸ₂, SS_and_pars, true + 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 +1189,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 +1240,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 +1271,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 +1290,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 - ŝ_to_ŝ₃ = [ 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] - ê_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ˢ^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₂ˢ)] - ŝ_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₃ˢ] + + ê_to_ŝ₃ = [ 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₁)] + + ŝ_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] @@ -900,44 +1337,510 @@ 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!(ŝ_to_ŝ₃, eps()) - droptol!(ê_to_ŝ₃, eps()) + droptol!(A_UU, eps()) + droptol!(A_LU, eps()) + droptol!(A_LL, eps()) + droptol!(ê_to_ŝ₃, eps()) droptol!(Eᴸᶻ, eps()) droptol!(Γ₃, eps()) - - A = ê_to_ŝ₃ * Eᴸᶻ * ŝ_to_ŝ₃' - droptol!(A, eps()) - C = ê_to_ŝ₃ * Γ₃ * ê_to_ŝ₃' + 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' + eQ = ê_L * Q + C_LL = sparse_ABAt(ê_L, Γ₃) + eQ + eQ' + 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 = sparse_ABAt(ê_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, + has_unit_roots = 𝓂.caches.has_unit_roots) + 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 = ŝ_to_y₃ * Σᶻ₃ * ŝ_to_y₃' + ê_to_y₃ * Γ₃ * ê_to_y₃' + ê_to_y₃ * Eᴸᶻ * ŝ_to_y₃' + ŝ_to_y₃ * Eᴸᶻ' * ê_to_y₃' - + Σʸ₃tmp = ŝ_to_y₃ * Σᶻ₃ * ŝ_to_y₃' + sparse_ABAt(ê_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), :] end end - return Σʸ₃, μʸ₂, SS_and_pars, solved && solved3 && solved_lyapunov + # Cache the 3rd-order result for reuse + all_solved = solved && solved3 && solved_lyapunov + if T === Float64 && all_solved + if size(𝓂.caches.covariance_third_order) != size(Σʸ₃) + 𝓂.caches.covariance_third_order = Matrix{Float64}(undef, size(Σʸ₃)...) + end + copyto!(𝓂.caches.covariance_third_order, Σʸ₃) + 𝓂.caches.valid_for.covariance_third_order = Float64.(parameters) + 𝓂.caches.valid_for.covariance_third_order_obs_key = obs_key + end + + return Σʸ₃, μʸ₂, SS_and_pars, all_solved +end + + +function determine_efficient_order(𝐒₁::Matrix{<: Real}, + constants::constants, + variables::Union{Symbol_input,String_input}; + covariance::Union{Symbol_input,String_input} = Symbol[], + tol::AbstractFloat = eps()) + + T = constants.post_model_macro + + + orders = Pair{Vector{Symbol}, Vector{Symbol}}[] + + nˢ = T.nPast_not_future_and_mixed + + if variables == :full_covar + return [T.var => T.past_not_future_and_mixed] + else + var_idx = MacroModelling.parse_variables_input_to_index(variables, constants) |> sort + observables = T.var[var_idx] + end + + # Precompute state indices to avoid repeated indexin calls + state_idx_in_var = indexin(T.past_not_future_and_mixed, T.var) .|> Int + 𝐒₁_states = 𝐒₁[state_idx_in_var, 1:nˢ] + + for obs in observables + obs_in_var_idx = indexin([obs],T.var) .|> Int + dependencies_in_states = vec(sum(abs, 𝐒₁[obs_in_var_idx,1:nˢ], dims=1) .> tol) .> 0 + + # Iterative propagation without redundant allocations + while true + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * 𝐒₁_states) .> tol) + if new_deps == dependencies_in_states + break + end + dependencies_in_states = new_deps + end + + dependencies = T.past_not_future_and_mixed[dependencies_in_states] + + push!(orders,[obs] => sort(dependencies)) + end + + # If covariance variables are specified, compute dependencies and add entries for those pairs + if !(covariance == Symbol[]) + covar_var_idx = MacroModelling.parse_variables_input_to_index(covariance, constants) |> sort + covariance_vars = T.var[covar_var_idx] + + # Compute dependencies for covariance variables (if not already computed) + for covar_var in covariance_vars + # 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 + dependencies_in_states = vec(sum(abs, 𝐒₁[obs_in_var_idx,1:nˢ], dims=1) .> tol) .> 0 + + # Iterative propagation without redundant allocations + while true + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * 𝐒₁_states) .> tol) + if new_deps == dependencies_in_states + break + end + dependencies_in_states = new_deps + end + + dependencies = T.past_not_future_and_mixed[dependencies_in_states] + push!(orders,[covar_var] => sort(dependencies)) + end + end + + # Build lookup dictionary for faster searches + var_to_idx = Dict{Symbol, Int}() + for (idx, order) in enumerate(orders) + for var in order.first + var_to_idx[var] = idx + end + end + + # Add entries for all pairs of covariance variables + for i in 1:length(covariance_vars) + for j in (i+1):length(covariance_vars) + # Find dependencies for both variables using lookup dictionary + idx_i = var_to_idx[covariance_vars[i]] + idx_j = var_to_idx[covariance_vars[j]] + + deps_i = orders[idx_i].second + deps_j = orders[idx_j].second + # Union of dependencies for covariance computation + combined_deps = sort(union(deps_i, deps_j)) + push!(orders, [covariance_vars[i], covariance_vars[j]] => combined_deps) + end + end + end + + sort!(orders, by = x -> length(x[2]), rev = true) + + return combine_pairs(orders) end + +function determine_efficient_order(𝐒₁::Matrix{<: Real}, + 𝐒₂::AbstractMatrix{<: Real}, + constants::constants, + variables::Union{Symbol_input,String_input}; + covariance::Union{Symbol_input,String_input} = Symbol[], + tol::AbstractFloat = eps()) + + T = constants.post_model_macro + + + orders = Pair{Vector{Symbol}, Vector{Symbol}}[] + + nˢ = T.nPast_not_future_and_mixed + nᵉ = T.nExo + + if variables == :full_covar + return [T.var => T.past_not_future_and_mixed] + else + var_idx = MacroModelling.parse_variables_input_to_index(variables, constants) |> sort + observables = T.var[var_idx] + end + + # Build selector for state variables in the augmented state vector [states; 1; shocks] + s_in_s⁺ = BitVector(vcat(ones(Bool, nˢ), zeros(Bool, nᵉ + 1))) + + # Kronecker product indices for state-state interactions + kron_s_s = ℒ.kron(s_in_s⁺, s_in_s⁺) + + # Precompute state indices and matrix slices to avoid repeated operations + 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 + + for obs in observables + obs_in_var_idx = indexin([obs],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] + + # Check which state variable pairs have influence + # 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 + + # Propagate dependencies through the system (iterative closure) + # considering both first and second 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 state i and state j are dependencies, + # their product can affect states + if !isnothing(𝐒₂_states) + # Generate selector vector for columns where both states are dependencies + selector = vec(ℒ.kron(prev_dependencies, prev_dependencies)) + if any(selector) + # Check which states are affected by the selected products + 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] + + push!(orders,[obs] => sort(dependencies)) + end + + # If covariance variables are specified, compute dependencies and add entries for those pairs + if !(covariance == Symbol[]) + covar_var_idx = MacroModelling.parse_variables_input_to_index(covariance, constants) |> sort + covariance_vars = T.var[covar_var_idx] + + # Compute dependencies for covariance variables (if not already computed) + for covar_var in covariance_vars + # 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 + + # 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 + + 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 + + if new_deps == dependencies_in_states + break + end + dependencies_in_states = new_deps + end + + dependencies = T.past_not_future_and_mixed[dependencies_in_states] + push!(orders,[covar_var] => sort(dependencies)) + end + end + + # Add entries for all pairs of covariance variables + for i in 1:length(covariance_vars) + for j in (i+1):length(covariance_vars) + # Find dependencies for both variables (they should exist now) + idx_i = findfirst(x -> covariance_vars[i] in x.first, orders) + idx_j = findfirst(x -> covariance_vars[j] in x.first, orders) + + deps_i = orders[idx_i].second + deps_j = orders[idx_j].second + # Union of dependencies for covariance computation + combined_deps = sort(union(deps_i, deps_j)) + push!(orders, [covariance_vars[i], covariance_vars[j]] => combined_deps) + end + end + end + + sort!(orders, by = x -> length(x[2]), rev = true) + + return combine_pairs(orders) end + + +function determine_efficient_order(𝐒₁::Matrix{<: Real}, + 𝐒₂::AbstractSparseMatrix{<: Real}, + 𝐒₃::AbstractSparseMatrix{<: Real}, + constants::constants, + variables::Union{Symbol_input,String_input}; + covariance::Union{Symbol_input,String_input} = Symbol[], + tol::AbstractFloat = eps()) + + T = constants.post_model_macro + + + orders = Pair{Vector{Symbol}, Vector{Symbol}}[] + + nˢ = T.nPast_not_future_and_mixed + nᵉ = T.nExo + + if variables == :full_covar + return [T.var => T.past_not_future_and_mixed] + else + var_idx = MacroModelling.parse_variables_input_to_index(variables, constants) |> sort + observables = T.var[var_idx] + end + + # Build selectors for state variables in the augmented state vector [states; 1; shocks] + s_in_s⁺ = BitVector(vcat(ones(Bool, nˢ), zeros(Bool, nᵉ + 1))) + + # Kronecker product indices for interactions + kron_s_s = ℒ.kron(s_in_s⁺, s_in_s⁺) + kron_s_s_s = ℒ.kron(kron_s_s, s_in_s⁺) + + # 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ˢ] + 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 + + 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 has_S₂ + s_s_to_y₂ = 𝐒₂[obs_in_var_idx, kron_s_s] + 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 has_S₃ + s_s_s_to_y₃ = 𝐒₃[obs_in_var_idx, kron_s_s_s] + 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 (iterative closure) + 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) + 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) + 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 + + 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 + + # If covariance variables are specified, compute dependencies and add entries for those pairs + if !(covariance == Symbol[]) + covar_var_idx = MacroModelling.parse_variables_input_to_index(covariance, constants) |> sort + covariance_vars = T.var[covar_var_idx] + + # Compute dependencies for covariance variables (if not already computed) + for covar_var in covariance_vars + # 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 + dependencies = compute_dependencies(obs_in_var_idx) + push!(orders,[covar_var] => sort(dependencies)) + end + end + + # Add entries for all pairs of covariance variables + for i in 1:length(covariance_vars) + for j in (i+1):length(covariance_vars) + # Find dependencies for both variables (they should exist now) + idx_i = findfirst(x -> covariance_vars[i] in x.first, orders) + idx_j = findfirst(x -> covariance_vars[j] in x.first, orders) + + deps_i = orders[idx_i].second + deps_j = orders[idx_j].second + # Union of dependencies for covariance computation + combined_deps = sort(union(deps_i, deps_j)) + push!(orders, [covariance_vars[i], covariance_vars[j]] => combined_deps) + end + end + end + + sort!(orders, by = x -> length(x[2]), rev = true) + + return combine_pairs(orders) +end + + + +function combine_pairs(v::Vector{Pair{Vector{Symbol}, Vector{Symbol}}}) + i = 1 + while i <= length(v) + subset_found = false + for j in i+1:length(v) + # Check if v[i].second and v[j].second are equal or if one is subset of the other + if v[i].second == v[j].second + # Exact match: combine first elements and remove duplicate + v[i] = v[i].first ∪ v[j].first => v[i].second + deleteat!(v, j) + subset_found = true + break + elseif all(elem -> elem in v[j].second, v[i].second) || all(elem -> elem in v[i].second, v[j].second) + # One is subset of the other: combine the first elements and assign to the one with the larger second element + if length(v[i].second) > length(v[j].second) + v[i] = v[i].first ∪ v[j].first => v[i].second + deleteat!(v, j) + else + v[j] = v[i].first ∪ v[j].first => v[j].second + deleteat!(v, i) + end + subset_found = true + break + end + end + # If no subset was found for v[i], move to the next element + if !subset_found + i += 1 + end + end + return v +end + + +end # @stable diff --git a/src/occasionally_binding_constraints.jl b/src/occasionally_binding_constraints.jl new file mode 100644 index 000000000..8c8197574 --- /dev/null +++ b/src/occasionally_binding_constraints.jl @@ -0,0 +1,907 @@ +@stable default_mode = "disable" begin + + +# ── Occasionally Binding Constraints (OBC) ─────────────────────────────────── +# +# Self-contained OBC functions extracted from MacroModelling.jl. +# Struct definitions remain in structures.jl; default constants in default_options.jl. + + +# ── Parsing & transformation ───────────────────────────────────────────────── + +check_for_dynamic_variables(ex::Int) = false +check_for_dynamic_variables(ex::Float64) = false +check_for_dynamic_variables(ex::Symbol) = occursin(r"₍₁₎|₍₀₎|₍₋₁₎",string(ex)) + +function check_for_dynamic_variables(ex::Expr) + dynamic_indicator = Bool[] + + postwalk(x -> + x isa Expr ? + x.head == :ref ? + occursin(r"^(ss|stst|steady|steadystate|steady_state){1}$"i,string(x.args[2])) ? + x : + begin + push!(dynamic_indicator,true) + x + end : + x : + x, + ex) + + any(dynamic_indicator) +end + + +function transform_expression(expr::Expr) + # Dictionary to store the transformations for reversing + reverse_transformations = Dict{Symbol, Expr}() + + # Counter for generating unique placeholders + unique_counter = Ref(0) + + # Step 1: Replace min/max calls and record their original form + function replace_min_max(expr) + if expr isa Expr && expr.head == :call && (expr.args[1] == :min || expr.args[1] == :max) + # Replace min/max functions with a placeholder + # placeholder = Symbol("minimal__P", unique_counter[]) + placeholder = :minmax__P + unique_counter[] += 1 + + # Store the original min/max call for reversal + reverse_transformations[placeholder] = expr + + return placeholder + else + return expr + end + end + + # Step 2: Transform :ref fields in the rest of the expression + function transform_ref_fields(expr) + if expr isa Expr && expr.head == :ref && isa(expr.args[1], Symbol) + # Handle :ref expressions + if isa(expr.args[2], Number) || isa(expr.args[2], Symbol) + if expr.args[2] < 0 + new_symbol = Symbol(expr.args[1], "__", abs(expr.args[2])) + else + new_symbol = Symbol(expr.args[1], "_", expr.args[2]) + end + else + # Generate a unique placeholder for complex :ref + unique_counter[] += 1 + placeholder = Symbol("__placeholder", unique_counter[]) + new_symbol = placeholder + end + + # Record the reverse transformation + reverse_transformations[new_symbol] = expr + + return new_symbol + else + return expr + end + end + + + # Replace equality sign with minus + function replace_equality_with_minus(expr) + if expr isa Expr && expr.head == :(=) + return Expr(:call, :-, expr.args...) + else + return expr + end + end + + # Apply transformations + expr = postwalk(replace_min_max, expr) + expr = postwalk(transform_ref_fields, expr) + transformed_expr = postwalk(replace_equality_with_minus, expr) + + return transformed_expr, reverse_transformations +end + + +function reverse_transformation(transformed_expr::Expr, reverse_dict::Dict{Symbol, Expr}) + # Function to replace the transformed symbols with their original form + function revert_symbol(expr) + if expr isa Symbol && haskey(reverse_dict, expr) + return reverse_dict[expr] + else + return expr + end + end + + # Revert the expression using postwalk + reverted_expr = postwalk(revert_symbol, transformed_expr) + + return reverted_expr +end + + +function check_for_minmax(expr) + contains_minmax = Bool[] + + postwalk(x -> + x isa Expr ? + x.head == :call ? + x.args[1] ∈ [:max,:min] ? + begin + push!(contains_minmax,true) + x + end : + x : + x : + x, + expr) + + any(contains_minmax) +end + + +function transform_obc(ex::Expr; avoid_solve::Bool = false) + transformed_expr, reverse_dict = transform_expression(ex) + + for symbs in get_symbols(transformed_expr) + sym_value = SPyPyC.symbols(string(symbs), real = true, finite = true) + Core.eval(SymPyWorkspace, :($symbs = $sym_value)) + end + + eq = Core.eval(SymPyWorkspace, transformed_expr) + + if avoid_solve || count_ops(Meta.parse(string(eq))) > 15 + soll = nothing + else + soll = solve_symbolically(eq, Core.eval(SymPyWorkspace, :minmax__P)) + end + + if !isempty(soll) + sorted_minmax = Expr(:call, reverse_dict[:minmax__P].args[1], :($(reverse_dict[:minmax__P].args[2]) - $(Meta.parse(string(soll[1])))), :($(reverse_dict[:minmax__P].args[3]) - $(Meta.parse(string(soll[1]))))) + return reverse_transformation(sorted_minmax, reverse_dict) + else + @error "Occasionally binding constraint not well-defined. See documentation for examples." + end +end + + +# try: run optim only if there is a violation / capture case with small shocks and set them to zero +function parse_occasionally_binding_constraints(equations_block; max_obc_horizon::Int = 40, avoid_solve::Bool = false) + # precision_factor = 1e #factor to force the optimiser to have non-relevatn shocks at zero + + eqs = [] + obc_shocks = Expr[] + + for arg in equations_block.args + if isa(arg,Expr) + if check_for_minmax(arg) + arg_trans = transform_obc(arg) + else + arg_trans = arg + end + + eq = postwalk(x -> + x isa Expr ? + x.head == :call ? + x.args[1] == :max ? + begin + + obc_vars_left = Expr(:ref, Meta.parse("χᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝˡ" ), 0) + obc_vars_right = Expr(:ref, Meta.parse("χᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝʳ" ), 0) + + if !(x.args[2] isa Symbol) && check_for_dynamic_variables(x.args[2]) + push!(eqs, :($obc_vars_left = $(x.args[2]))) + else + obc_vars_left = x.args[2] + end + + if !(x.args[3] isa Symbol) && check_for_dynamic_variables(x.args[3]) + push!(eqs, :($obc_vars_right = $(x.args[3]))) + else + obc_vars_right = x.args[3] + end + + obc_inequality = Expr(:ref, Meta.parse("Χᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ" ), 0) + + push!(eqs, :($obc_inequality = $(Expr(x.head, x.args[1], obc_vars_left, obc_vars_right)))) + + obc_shock = Expr(:ref, Meta.parse("ϵᵒᵇᶜ⁺ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ"), 0) + + push!(obc_shocks, obc_shock) + + :($obc_inequality - $obc_shock) + end : + x.args[1] == :min ? + begin + obc_vars_left = Expr(:ref, Meta.parse("χᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝˡ" ), 0) + obc_vars_right = Expr(:ref, Meta.parse("χᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝʳ" ), 0) + + if !(x.args[2] isa Symbol) && check_for_dynamic_variables(x.args[2]) + push!(eqs, :($obc_vars_left = $(x.args[2]))) + else + obc_vars_left = x.args[2] + end + + if !(x.args[3] isa Symbol) && check_for_dynamic_variables(x.args[3]) + push!(eqs, :($obc_vars_right = $(x.args[3]))) + else + obc_vars_right = x.args[3] + end + + obc_inequality = Expr(:ref, Meta.parse("Χᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ" ), 0) + + push!(eqs, :($obc_inequality = $(Expr(x.head, x.args[1], obc_vars_left, obc_vars_right)))) + + obc_shock = Expr(:ref, Meta.parse("ϵᵒᵇᶜ⁻ꜝ" * super(string(length(obc_shocks) + 1)) * "ꜝ"), 0) + + push!(obc_shocks, obc_shock) + + :($obc_inequality - $obc_shock) + end : + x : + x : + x, + arg_trans) + + push!(eqs, eq) + end + end + + for obc in obc_shocks + # push!(eqs, :($(obc) = $(Expr(:ref, obc.args[1], -1)) * 0.3 + $(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(max_obc_horizon)) * "⁾"), 0)))) + push!(eqs, :($(obc) = $(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(max_obc_horizon)) * "⁾"), 0)))) + + push!(eqs, :($(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻⁰⁾"), 0)) = activeᵒᵇᶜshocks * $(Expr(:ref, Meta.parse(string(obc.args[1]) * "⁽" * super(string(max_obc_horizon)) * "⁾"), :x)))) + + for i in 1:max_obc_horizon + push!(eqs, :($(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(i)) * "⁾"), 0)) = $(Expr(:ref, Meta.parse(string(obc.args[1]) * "ᴸ⁽⁻" * super(string(i-1)) * "⁾"), -1)) + activeᵒᵇᶜshocks * $(Expr(:ref, Meta.parse(string(obc.args[1]) * "⁽" * super(string(max_obc_horizon-i)) * "⁾"), :x)))) + end + end + + return Expr(:block, eqs...) +end + + +function write_obc_violation_equations(𝓂) + eqs = Expr[] + for (i,eq) in enumerate(𝓂.equations.dynamic) + if check_for_minmax(eq) + minmax_fixed_eqs = postwalk(x -> + x isa Expr ? + x.head == :call ? + length(x.args) == 3 ? + x.args[3] isa Expr ? + x.args[3].args[1] ∈ [:Min, :min, :Max, :max] ? + begin + plchldr = Symbol(replace(string(x.args[2]), "₍₀₎" => "")) + + ineq_plchldr_1 = x.args[3].args[2] isa Symbol ? Symbol(replace(string(x.args[3].args[2]), "₍₀₎" => "")) : x.args[3].args[2] + + arg1 = x.args[3].args[2] + arg2 = x.args[3].args[3] + + dyn_1 = check_for_dynamic_variables(x.args[3].args[2]) + dyn_2 = check_for_dynamic_variables(x.args[3].args[3]) + + cond1 = Expr[] + cond2 = Expr[] + + maximisation = contains(string(plchldr), "⁺") + + # if dyn_1 + # if maximisation + # push!(cond1, :(push!(constraint_values, $(x.args[3].args[2])))) + # # push!(cond2, :(push!(constraint_values, $(x.args[3].args[2])))) + # else + # push!(cond1, :(push!(constraint_values, -$(x.args[3].args[2])))) + # # push!(cond2, :(push!(constraint_values, -$(x.args[3].args[2])))) # RBC + # end + # end + + # if dyn_2 + # if maximisation + # push!(cond1, :(push!(constraint_values, $(x.args[3].args[3])))) + # # push!(cond2, :(push!(constraint_values, $(x.args[3].args[3])))) # testmax + # else + # push!(cond1, :(push!(constraint_values, -$(x.args[3].args[3])))) + # # push!(cond2, :(push!(constraint_values, -$(x.args[3].args[3])))) # RBC + # end + # end + + + if maximisation + push!(cond1, :(push!(constraint_values, [sum($(x.args[3].args[2]) .* $(x.args[3].args[3]))]))) + push!(cond1, :(push!(constraint_values, $(x.args[3].args[2])))) + push!(cond1, :(push!(constraint_values, $(x.args[3].args[3])))) + # push!(cond1, :(push!(constraint_values, max.($(x.args[3].args[2]), $(x.args[3].args[3]))))) + else + push!(cond1, :(push!(constraint_values, [sum($(x.args[3].args[2]) .* $(x.args[3].args[3]))]))) + push!(cond1, :(push!(constraint_values, -$(x.args[3].args[2])))) + push!(cond1, :(push!(constraint_values, -$(x.args[3].args[3])))) + # push!(cond1, :(push!(constraint_values, min.($(x.args[3].args[2]), $(x.args[3].args[3]))))) + end + + # if maximisation + # push!(cond1, :(push!(shock_sign_indicators, true))) + # # push!(cond2, :(push!(shock_sign_indicators, true))) + # else + # push!(cond1, :(push!(shock_sign_indicators, false))) + # # push!(cond2, :(push!(shock_sign_indicators, false))) + # end + + # :(if isapprox($plchldr, $ineq_plchldr_1, atol = 1e-12) + # $(Expr(:block, cond1...)) + # else + # $(Expr(:block, cond2...)) + # end) + :($(Expr(:block, cond1...))) + end : + x : + x : + x : + x : + x, + eq) + + push!(eqs, minmax_fixed_eqs) + end + end + + return eqs +end + + +# ── OBC flag processing ────────────────────────────────────────────────────── + +function process_ignore_obc_flag(shocks, + ignore_obc::Bool, + 𝓂::ℳ; + maxlog::Int = DEFAULT_MAXLOG) + stochastic_model = length(𝓂.constants.post_model_macro.exo) > 0 + obc_model = length(𝓂.equations.obc_violation) > 0 + + obc_shocks_included = false + + if stochastic_model && obc_model + if shocks isa Matrix{Float64} + obc_indices = contains.(string.(𝓂.constants.post_model_macro.exo), "ᵒᵇᶜ") + if any(obc_indices) + obc_shocks_included = sum(abs2, shocks[obc_indices, :]) > 1e-10 + end + elseif shocks isa KeyedArray{Float64} + shock_axis = collect(axiskeys(shocks, 1)) + shock_axis = shock_axis isa Vector{String} ? shock_axis .|> Meta.parse .|> replace_indices : shock_axis + + obc_shocks = 𝓂.constants.post_model_macro.exo[contains.(string.(𝓂.constants.post_model_macro.exo), "ᵒᵇᶜ")] + relevant_shocks = intersect(obc_shocks, shock_axis) + + if !isempty(relevant_shocks) + obc_shocks_included = sum(abs2, shocks(relevant_shocks, :)) > 1e-10 + end + else + shock_idx = parse_shocks_input_to_index(shocks, 𝓂.constants) + + selected_shocks = if (shock_idx isa Vector) || (shock_idx isa UnitRange) + length(shock_idx) > 0 ? 𝓂.constants.post_model_macro.exo[shock_idx] : Symbol[] + else + [𝓂.constants.post_model_macro.exo[shock_idx]] + end + + obc_shocks = 𝓂.constants.post_model_macro.exo[contains.(string.(𝓂.constants.post_model_macro.exo), "ᵒᵇᶜ")] + obc_shocks_included = !isempty(intersect(selected_shocks, obc_shocks)) + end + end + + ignore_obc_flag = ignore_obc + + if ignore_obc_flag && !obc_model + @info "`ignore_obc = true` has no effect because $(𝓂.model_name) has no occasionally binding constraints. Setting `ignore_obc = false`." maxlog = maxlog + ignore_obc_flag = false + end + + if ignore_obc_flag && obc_shocks_included + @warn "`ignore_obc = true` cannot be applied because shocks affecting occasionally binding constraints are included. Enforcing the constraints instead and setting `ignore_obc = false`." maxlog = maxlog + ignore_obc_flag = false + end + + occasionally_binding_constraints = obc_model && !ignore_obc_flag + + return ignore_obc_flag, occasionally_binding_constraints, obc_shocks_included +end + + +# ── OBC violation function setup ───────────────────────────────────────────── + +function set_up_obc_violation_function!(𝓂) + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎$"))) + + sort!(present_varss ,by = x->replace(string(x),r"₍₀₎$"=>"")) + + # write indices in auxiliary objects + dyn_var_present_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₀₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎"))) + + dyn_var_present = Symbol.(replace.(string.(sort(collect(reduce(union,dyn_var_present_list)))), r"ᴸ⁽⁻?[⁰¹²³⁴⁵⁶⁷⁸⁹]+⁾" => "")) + + SS_and_pars_names = ms.SS_and_pars_names + + dyn_var_present_idx = indexin(dyn_var_present , SS_and_pars_names) + + alll = [] + for (i,var) in enumerate(present_varss) + if !(match(r"^χᵒᵇᶜ", string(var)) === nothing) + push!(alll,:($var = Y[$(dyn_var_present_idx[i]),1:max(periods, 1)])) + end + end + + calc_obc_violation = :(function calculate_obc_violation(x, p) + state, state_update, reference_steady_state, 𝓂, algorithm, periods, shock_values = p + + T = 𝓂.constants.post_model_macro + + Y = zeros(typeof(x[1]), T.nVars, periods+1) + + shock_values = convert(typeof(x), shock_values) + + shock_values[contains.(string.(T.exo),"ᵒᵇᶜ")] .= x + + zero_shock = zero(shock_values) + + if algorithm ∈ [:pruned_second_order, :pruned_third_order] + states = state_update(state, shock_values) + Y[:,1] = sum(states) + else + Y[:,1] = state_update(state, shock_values) + end + + for t in 1:periods + if algorithm ∈ [:pruned_second_order, :pruned_third_order] + states = state_update(states, zero_shock) + Y[:,t+1] = sum(states) + else + Y[:,t+1] = state_update(Y[:,t], zero_shock) + end + end + + Y .+= reference_steady_state[1:T.nVars] + + $(alll...) + + constraint_values = Vector[] + + $(𝓂.equations.obc_violation...) + + return vcat(constraint_values...) + end) + + 𝓂.functions.obc_violation = @RuntimeGeneratedFunction(calc_obc_violation) + + # ── Extract OBC constraint metadata for the analytical Jacobian ── + # Build mapping: χᵒᵇᶜ variable name (without ₍₀₎) → Y row index + chi_row_map = Dict{String, Int}() + for (i, var) in enumerate(present_varss) + vstr = string(var) + if startswith(vstr, "χᵒᵇᶜ") + name = replace(vstr, "₍₀₎" => "") + chi_row_map[name] = dyn_var_present_idx[i] + end + end + + # Pair left/right χᵒᵇᶜ variables by constraint key + left_vars = Dict{String, String}() + right_vars = Dict{String, String}() + for name in keys(chi_row_map) + if endswith(name, "ˡ") + key = name[1:prevind(name, lastindex(name))] + left_vars[key] = name + elseif endswith(name, "ʳ") + key = name[1:prevind(name, lastindex(name))] + right_vars[key] = name + end + end + + obc_info = Tuple{Int, Int, Float64}[] + for key in sort(collect(keys(left_vars))) + if haskey(right_vars, key) + left_idx = chi_row_map[left_vars[key]] + right_idx = chi_row_map[right_vars[key]] + sign = contains(key, "⁺") ? 1.0 : -1.0 # max → +1, min → −1 + push!(obc_info, (left_idx, right_idx, sign)) + end + end + 𝓂.functions.obc_constraint_info = obc_info + + return nothing +end + + +# ── NLopt objective & constraint callbacks ─────────────────────────────────── + +function obc_objective_optim_fun(X::Vector{S}, grad::Vector{S})::S where S + if length(grad) > 0 + grad .= 2 .* X + end + + sum(abs2, X) +end + +function obc_constraint_optim_fun(res::Vector{S}, X::Vector{S}, jac::Matrix{S}, p) where S + 𝓂 = p[4] + + res .= 𝓂.functions.obc_violation(X, p) + + if length(jac) > 0 + compute_obc_analytical_jacobian!(jac, X, p) + end + + return nothing +end + + +# ── Analytical OBC Jacobian ────────────────────────────────────────────────── +# +# The OBC constraint vector has, per constraint, three blocks: +# 1. [sum(a .* b)] (1 element — complementary slackness) +# 2. sign * a (P elements — left argument) +# 3. sign * b (P elements — right argument) +# where a = Y[left_row, 1:P], b = Y[right_row, 1:P], and +# sign = +1 for max, −1 for min. +# +# Y is the forward path simulated through state_update, which is a known +# function of the perturbation solution matrices. dY/dx is therefore +# computed analytically (exactly for all algorithm orders). + +function compute_obc_analytical_jacobian!(jac::Matrix{S}, X::Vector{S}, p) where S + state, state_update, reference_steady_state, 𝓂, algorithm, periods, shock_values = p + T = 𝓂.constants.post_model_macro + nv = T.nVars + past_idx = T.past_not_future_and_mixed_idx + n_past = T.nPast_not_future_and_mixed + n_x = length(X) + P = max(periods, 1) + + obc_idx = findall(contains.(string.(T.exo), "ᵒᵇᶜ")) + shock_vals = copy(shock_values) + shock_vals[obc_idx] .= X + n_shocks = length(shock_vals) + zero_shock = zero(shock_vals) + + Ŝ₁ = 𝓂.caches.first_order_obc_solution_matrix + + Y = zeros(S, nv, periods + 1) + dYdx = zeros(S, nv, n_x, periods + 1) + + if algorithm == :first_order + obc_dYdx_first_order!(Y, dYdx, state, shock_vals, zero_shock, + past_idx, n_past, obc_idx, Ŝ₁, periods) + + elseif algorithm ∈ [:second_order, :third_order] + obc_dYdx_nonpruned_higher!(Y, dYdx, state, shock_vals, zero_shock, + past_idx, n_past, n_shocks, obc_idx, + Ŝ₁, 𝓂, algorithm, periods) + + elseif algorithm ∈ [:pruned_second_order, :pruned_third_order] + obc_dYdx_pruned!(Y, dYdx, state, shock_vals, zero_shock, + past_idx, n_past, n_shocks, obc_idx, + Ŝ₁, 𝓂, algorithm, periods) + end + + Y .+= @view reference_steady_state[1:nv] + + fill_obc_constraint_jacobian!(jac, Y, dYdx, + 𝓂.functions.obc_constraint_info, n_x, P) + return nothing +end + + +# ── First-order: purely linear propagation ─────────────────────────────────── +function obc_dYdx_first_order!(Y, dYdx, state, shock_vals, zero_shock, + past_idx, n_past, obc_idx, Ŝ₁, periods) + A = @view Ŝ₁[:, 1:n_past] + Y[:, 1] = Ŝ₁ * [state[past_idx]; shock_vals] + dYdx[:, :, 1] .= @view Ŝ₁[:, n_past .+ obc_idx] + for t in 1:periods + Y[:, t+1] = A * Y[past_idx, t] + dYdx[:, :, t+1] = A * dYdx[past_idx, :, t] + end +end + + +# ── Non-pruned second / third order ───────────────────────────────────────── +function obc_dYdx_nonpruned_higher!(Y, dYdx, state, shock_vals, zero_shock, + past_idx, n_past, n_shocks, obc_idx, + Ŝ₁, 𝓂, algorithm, periods) + S = eltype(Y) + nv = size(Y, 1) + n_x = size(dYdx, 2) + 𝐒₂ = 𝓂.caches.second_order_solution * 𝓂.constants.second_order.𝐔₂ + Ŝ₁̂ = [Ŝ₁[:, 1:n_past] zeros(S, nv) Ŝ₁[:, n_past+1:end]] + n_aug = n_past + 1 + n_shocks + + has_third = algorithm == :third_order + 𝐒₃ = has_third ? 𝓂.caches.third_order_solution * 𝓂.constants.third_order.𝐔₃ : nothing + + # ── t = 0 ── + aug = [state[past_idx]; one(S); shock_vals] + kron_aug = ℒ.kron(aug, aug) + Y[:, 1] = Ŝ₁̂ * aug + 𝐒₂ * kron_aug / 2 + if has_third; Y[:, 1] += 𝐒₃ * ℒ.kron(kron_aug, aug) / 6; end + + d_aug = zeros(S, n_aug) + for j in 1:n_x + fill!(d_aug, zero(S)) + d_aug[n_past + 1 + obc_idx[j]] = one(S) + dYdx[:, j, 1] = Ŝ₁̂ * d_aug + + 𝐒₂ * (ℒ.kron(d_aug, aug) + ℒ.kron(aug, d_aug)) / 2 + if has_third + dYdx[:, j, 1] += 𝐒₃ * (ℒ.kron(ℒ.kron(d_aug, aug), aug) + + ℒ.kron(ℒ.kron(aug, d_aug), aug) + + ℒ.kron(kron_aug, d_aug)) / 6 + end + end + + # ── t > 0 ── + d_aug_t = zeros(S, n_aug) + for t in 1:periods + aug_t = [Y[past_idx, t]; one(S); zeros(S, n_shocks)] + kron_aug_t = ℒ.kron(aug_t, aug_t) + Y[:, t+1] = Ŝ₁̂ * aug_t + 𝐒₂ * kron_aug_t / 2 + if has_third; Y[:, t+1] += 𝐒₃ * ℒ.kron(kron_aug_t, aug_t) / 6; end + + for j in 1:n_x + fill!(d_aug_t, zero(S)) + d_aug_t[1:n_past] .= @view dYdx[past_idx, j, t] + dYdx[:, j, t+1] = Ŝ₁̂ * d_aug_t + + 𝐒₂ * (ℒ.kron(d_aug_t, aug_t) + ℒ.kron(aug_t, d_aug_t)) / 2 + if has_third + dYdx[:, j, t+1] += 𝐒₃ * (ℒ.kron(ℒ.kron(d_aug_t, aug_t), aug_t) + + ℒ.kron(ℒ.kron(aug_t, d_aug_t), aug_t) + + ℒ.kron(kron_aug_t, d_aug_t)) / 6 + end + end + end +end + + +# ── Pruned second / third order ───────────────────────────────────────────── +function obc_dYdx_pruned!(Y, dYdx, state, shock_vals, zero_shock, + past_idx, n_past, n_shocks, obc_idx, + Ŝ₁, 𝓂, algorithm, periods) + S = eltype(Y) + nv = size(Y, 1) + n_x = size(dYdx, 2) + 𝐒₂ = 𝓂.caches.second_order_solution * 𝓂.constants.second_order.𝐔₂ + Ŝ₁̂ = [Ŝ₁[:, 1:n_past] zeros(S, nv) Ŝ₁[:, n_past+1:end]] + n_aug = n_past + 1 + n_shocks + + has_third = algorithm == :pruned_third_order + 𝐒₃ = has_third ? 𝓂.caches.third_order_solution * 𝓂.constants.third_order.𝐔₃ : nothing + + # Component vectors + y₁ = state isa AbstractVector{<:AbstractVector} ? state[1] : state + y₂ = state isa AbstractVector{<:AbstractVector} ? state[2] : zeros(S, nv) + y₃ = (has_third && state isa AbstractVector{<:AbstractVector} && length(state) >= 3) ? + state[3] : zeros(S, nv) + + dy₁dx = zeros(S, nv, n_x) + dy₂dx = zeros(S, nv, n_x) + dy₃dx = zeros(S, nv, n_x) + + d_aug = zeros(S, n_aug) + + # ── t = 0 ── + aug₁ = [y₁[past_idx]; one(S); shock_vals] + y₁_new = Ŝ₁̂ * aug₁ + + aug₂ = [y₂[past_idx]; zero(S); zeros(S, n_shocks)] + kron_aug₁ = ℒ.kron(aug₁, aug₁) + y₂_new = Ŝ₁̂ * aug₂ + 𝐒₂ * kron_aug₁ / 2 + + for j in 1:n_x + fill!(d_aug, zero(S)) + d_aug[n_past + 1 + obc_idx[j]] = one(S) + dy₁dx[:, j] = Ŝ₁̂ * d_aug + # dy₂ only depends on aug₁ perturbation (aug₂ initial is independent of x) + dy₂dx[:, j] = 𝐒₂ * (ℒ.kron(d_aug, aug₁) + ℒ.kron(aug₁, d_aug)) / 2 + end + + if has_third + aug₁̂ = [y₁[past_idx]; zero(S); shock_vals] + aug₃ = [y₃[past_idx]; zero(S); zeros(S, n_shocks)] + y₃_new = Ŝ₁̂ * aug₃ + 𝐒₂ * ℒ.kron(aug₁̂, aug₂) + 𝐒₃ * ℒ.kron(kron_aug₁, aug₁) / 6 + + for j in 1:n_x + fill!(d_aug, zero(S)) + d_aug[n_past + 1 + obc_idx[j]] = one(S) + d_aug₁̂ = copy(d_aug); d_aug₁̂[n_past + 1] = zero(S) # hat: zero for the "1" slot + dy₃dx[:, j] = 𝐒₂ * (ℒ.kron(d_aug₁̂, aug₂) + ℒ.kron(aug₁̂, zeros(S, n_aug))) + + 𝐒₃ * (ℒ.kron(ℒ.kron(d_aug, aug₁), aug₁) + + ℒ.kron(ℒ.kron(aug₁, d_aug), aug₁) + + ℒ.kron(kron_aug₁, d_aug)) / 6 + end + y₃ = y₃_new + end + + y₁ = y₁_new + y₂ = y₂_new + Y[:, 1] = y₁ + y₂ + dYdx[:, :, 1] = dy₁dx + dy₂dx + if has_third; Y[:, 1] += y₃; dYdx[:, :, 1] += dy₃dx; end + + # ── t > 0 ── + d_aug_t = zeros(S, n_aug) + for t in 1:periods + aug₁_t = [y₁[past_idx]; one(S); zeros(S, n_shocks)] + kron_aug₁_t = ℒ.kron(aug₁_t, aug₁_t) + + y₁_new = Ŝ₁̂ * aug₁_t + aug₂_t = [y₂[past_idx]; zero(S); zeros(S, n_shocks)] + y₂_new = Ŝ₁̂ * aug₂_t + 𝐒₂ * kron_aug₁_t / 2 + + dy₁dx_new = zeros(S, nv, n_x) + dy₂dx_new = zeros(S, nv, n_x) + + for j in 1:n_x + fill!(d_aug_t, zero(S)) + d_aug_t[1:n_past] .= @view dy₁dx[past_idx, j] + dy₁dx_new[:, j] = Ŝ₁̂ * d_aug_t + + d_aug₂_t = zeros(S, n_aug) + d_aug₂_t[1:n_past] .= @view dy₂dx[past_idx, j] + dy₂dx_new[:, j] = Ŝ₁̂ * d_aug₂_t + + 𝐒₂ * (ℒ.kron(d_aug_t, aug₁_t) + ℒ.kron(aug₁_t, d_aug_t)) / 2 + end + + if has_third + aug₁̂_t = [y₁[past_idx]; zero(S); zeros(S, n_shocks)] + aug₃_t = [y₃[past_idx]; zero(S); zeros(S, n_shocks)] + y₃_new = Ŝ₁̂ * aug₃_t + 𝐒₂ * ℒ.kron(aug₁̂_t, aug₂_t) + 𝐒₃ * ℒ.kron(kron_aug₁_t, aug₁_t) / 6 + + dy₃dx_new = zeros(S, nv, n_x) + for j in 1:n_x + fill!(d_aug_t, zero(S)) + d_aug_t[1:n_past] .= @view dy₁dx[past_idx, j] + d_aug₁̂_t = copy(d_aug_t); d_aug₁̂_t[n_past + 1] = zero(S) + + d_aug₂_t = zeros(S, n_aug) + d_aug₂_t[1:n_past] .= @view dy₂dx[past_idx, j] + + d_aug₃_t = zeros(S, n_aug) + d_aug₃_t[1:n_past] .= @view dy₃dx[past_idx, j] + + dy₃dx_new[:, j] = Ŝ₁̂ * d_aug₃_t + + 𝐒₂ * (ℒ.kron(d_aug₁̂_t, aug₂_t) + ℒ.kron(aug₁̂_t, d_aug₂_t)) + + 𝐒₃ * (ℒ.kron(ℒ.kron(d_aug_t, aug₁_t), aug₁_t) + + ℒ.kron(ℒ.kron(aug₁_t, d_aug_t), aug₁_t) + + ℒ.kron(kron_aug₁_t, d_aug_t)) / 6 + end + y₃ = y₃_new + dy₃dx .= dy₃dx_new + end + + y₁ = y₁_new + y₂ = y₂_new + dy₁dx .= dy₁dx_new + dy₂dx .= dy₂dx_new + + Y[:, t+1] = y₁ + y₂ + dYdx[:, :, t+1] = dy₁dx + dy₂dx + if has_third; Y[:, t+1] += y₃; dYdx[:, :, t+1] += dy₃dx; end + end +end + + +# ── Fill NLopt Jacobian from dY/dx and constraint structure ────────────────── +function fill_obc_constraint_jacobian!(jac, Y, dYdx, constraint_info, n_x, P) + row_offset = 0 + for (left_idx, right_idx, sign) in constraint_info + # Complementary-slackness scalar: sum(Y[left,1:P] .* Y[right,1:P]) + for j in 1:n_x + val = zero(eltype(jac)) + for t in 1:P + val += dYdx[left_idx, j, t] * Y[right_idx, t] + + Y[left_idx, t] * dYdx[right_idx, j, t] + end + jac[j, row_offset + 1] = val + end + + # Left argument: sign * Y[left, 1:P] + for j in 1:n_x + for t in 1:P + jac[j, row_offset + 1 + t] = sign * dYdx[left_idx, j, t] + end + end + + # Right argument: sign * Y[right, 1:P] + for j in 1:n_x + for t in 1:P + jac[j, row_offset + 1 + P + t] = sign * dYdx[right_idx, j, t] + end + end + + row_offset += 1 + 2 * P + end +end + + +# ── First-order OBC solution ───────────────────────────────────────────────── + +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 + + write_parameters_input!(𝓂, :activeᵒᵇᶜshocks => 1, verbose = false) + + ∇̂₁ = calculate_jacobian(𝓂.parameter_values, 𝓂.caches.non_stochastic_steady_state, 𝓂.caches, 𝓂.functions.jacobian, 𝓂.workspaces, caching = false) + + Ŝ₁, qme_sol, solved = calculate_first_order_solution(∇̂₁, + constants, + 𝓂.workspaces, + 𝓂.caches; + opts = opts, + initial_guess = 𝓂.caches.qme_solution, + caching = false) + + update_perturbation_counter!(𝓂.counters, solved, order = 1) + + write_parameters_input!(𝓂, :activeᵒᵇᶜshocks => 0, verbose = false) + + # Cache write + stamp + 𝓂.caches.first_order_obc_solution_matrix = Ŝ₁ + 𝓂.caches.valid_for.first_order_obc_solution = Float64.(𝓂.parameter_values) + + return nothing +end + + +# ── OBC state update (per-period NLopt solver) ─────────────────────────────── + +function obc_state_update(present_states::S, present_shocks::Vector{R}, state_update::F, 𝓂::ℳ, algorithm::Symbol) where {S, R <: Float64, F} + unconditional_forecast_horizon = 𝓂.constants.post_model_macro.max_obc_horizon + + reference_ss = 𝓂.caches.non_stochastic_steady_state + + obc_shock_idx = contains.(string.(𝓂.constants.post_model_macro.exo),"ᵒᵇᶜ") + + periods_per_shock = 𝓂.constants.post_model_macro.max_obc_horizon + 1 + + num_shocks = sum(obc_shock_idx) ÷ periods_per_shock + + p = (present_states, state_update, reference_ss, 𝓂, algorithm, unconditional_forecast_horizon, present_shocks) + + constraints_violated = any(𝓂.functions.obc_violation(zeros(num_shocks*periods_per_shock), p) .> eps(Float32))::Bool + + if constraints_violated + opt = NLopt.Opt(NLopt.:LD_SLSQP, num_shocks*periods_per_shock) + + opt.min_objective = obc_objective_optim_fun + + opt.xtol_abs = eps(Float32) + opt.ftol_abs = eps(Float32) + opt.maxeval = 500 + + upper_bounds = fill(eps(), num_shocks * (1 + 2 * max(unconditional_forecast_horizon, 1))) + + NLopt.inequality_constraint!(opt, (res, x, jac) -> obc_constraint_optim_fun(res, x, jac, p), upper_bounds) + + (minf,x,ret) = NLopt.optimize(opt, zeros(num_shocks*periods_per_shock)) + + present_shocks[contains.(string.(𝓂.constants.post_model_macro.exo),"ᵒᵇᶜ")] .= x + + constraints_violated = any(𝓂.functions.obc_violation(x, p) .> eps(Float32))::Bool + + solved = !constraints_violated + else + solved = true + end + + present_states = state_update(present_states, present_shocks)::S + + return present_states, present_shocks, solved +end + + +end # @stable diff --git a/src/options_and_caches.jl b/src/options_and_caches.jl index eac752f4b..f982f318b 100644 --- a/src/options_and_caches.jl +++ b/src/options_and_caches.jl @@ -1,7 +1,83 @@ +@stable default_mode = "disable" begin + + +# ============================================================================ +# Doubling power-cache helpers +# ============================================================================ +# These are shared by the Lyapunov and Sylvester `:doubling` solvers and their +# adjoint helpers. Capture is gated by `workspace.pow_capture`; the rrule +# wrapper is responsible for setting/clearing the flag and bumping the stamp. +# Slot type is AbstractMatrix so the same vector can hold dense or sparse +# A^(2^k) / B^(2^k) entries (sparse-aware). + + """ - Second_order_indices() + cache_set!(vec, k, M) + +Place `M` into slot `k` of the doubling power cache vector. If the slot does +not yet exist, the vector is extended by one (sequential extension). +The slot is reused in-place when its concrete type and shape match `M`, +otherwise it is replaced with `copy(M)` to preserve the input's concrete type. + +When `transposed=true`, dense matrices are stored as `copy(M)'` (an `Adjoint` +wrapping a sequential copy of the data). This avoids the expensive element +rearrangement of `permutedims` — the solver dispatch handles `Adjoint` matrices +natively. Subsequent writes to the same slot reuse the existing parent buffer +(zero allocation after the first write). Sparse matrices are still materialised +as `SparseMatrixCSC(M')` since CSC format is required for efficient sparse ops. +""" +@inline function cache_set!(vec::Vector{<:AbstractMatrix}, k::Int, M::AbstractMatrix, transposed::Bool = false) + if transposed + if M isa AbstractSparseMatrix + Mt = SparseMatrixCSC(M') + if length(vec) < k + push!(vec, Mt) + else + slot = vec[k] + if typeof(slot) === typeof(Mt) && size(slot) == size(Mt) + copyto!(slot, Mt) + else + vec[k] = Mt + end + end + else + # Dense: store as Adjoint view — sequential copy + zero-cost wrapper. + # On subsequent writes, reuse the parent buffer to avoid allocation. + if length(vec) < k + push!(vec, copy(M)') + else + slot = vec[k] + if slot isa ℒ.Adjoint && typeof(parent(slot)) === typeof(M) && size(parent(slot)) == size(M) + copyto!(parent(slot), M) + else + vec[k] = copy(M)' + end + end + end + else + if length(vec) < k + push!(vec, copy(M)) + else + slot = vec[k] + if typeof(slot) === typeof(M) && size(slot) == size(M) + copyto!(slot, M) + else + vec[k] = copy(M) + end + end + end + return nothing +end + + + + + + + +""" Create an empty `second_order_indices` struct with all fields initialized to empty/zero values. These will be lazily populated by various ensure_*! functions as needed. @@ -9,14 +85,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 @@ -40,6 +129,7 @@ function Second_order_indices() # Conditional forecast indices Int[], # var²_idxs Int[], # shockvar²_idxs + Int[], # shockvar_no_vol_idxs # Moment computation caches BitVector(), # kron_states empty_sparse_float, # I_plus_s_s @@ -63,13 +153,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 +171,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 @@ -131,136 +233,416 @@ function Nonlinear_solver_workspace(func_buffer::Vector{T}, jac_buffer::Abstract end -function Krylov_workspace(;S::Type = Float64) - krylov_workspace( GmresWorkspace(0,0,Vector{S}), - DqgmresWorkspace(0,0,Vector{S}), - BicgstabWorkspace(0,0,Vector{S})) +function Krylov_workspace(::Type{ST} = Float64) where {ST <: AbstractFloat} + krylov_workspace( GmresWorkspace(0,0,Vector{ST}), + DqgmresWorkspace(0,0,Vector{ST}), + BicgstabWorkspace(0,0,Vector{ST})) end -function Sylvester_workspace(;S::Type = Float64, T::Type = Float64) +function Sylvester_workspace(::Type{ST} = Float64, ::Type{TT} = Float64) where {ST <: AbstractFloat, TT <: Real} sylvester_workspace( 0, 0, # n, m dimensions - zeros(S,0,0), # tmp (Krylov) - zeros(S,0,0), # 𝐗 (Krylov) - zeros(S,0,0), # 𝐂 (Krylov) - zeros(S,0,0), # 𝐀 (doubling) - zeros(S,0,0), # 𝐀¹ (doubling) - zeros(S,0,0), # 𝐁 (doubling) - zeros(S,0,0), # 𝐁¹ (doubling) - zeros(S,0,0), # 𝐂_dbl (doubling) - zeros(S,0,0), # 𝐂¹ (doubling) - zeros(S,0,0), # 𝐂B (doubling) - Krylov_workspace(S = S), - # ForwardDiff partials buffers - zeros(T,0,0), # P̃ - zeros(T,0,0), # Ã_fd - zeros(T,0,0), # B̃_fd - zeros(T,0,0)) # C̃_fd + zeros(ST,0,0), # tmp (Krylov) + zeros(ST,0,0), # 𝐗 (Krylov) + zeros(ST,0,0), # 𝐂 (Krylov) + zeros(ST,0,0), # 𝐀 (doubling) + zeros(ST,0,0), # 𝐀¹ (doubling) + zeros(ST,0,0), # 𝐁 (doubling) + zeros(ST,0,0), # 𝐁¹ (doubling) + zeros(ST,0,0), # 𝐂_dbl (doubling) + zeros(ST,0,0), # 𝐂¹ (doubling) + zeros(ST,0,0), # 𝐂B (doubling) + Krylov_workspace(ST), + zeros(ST,0,0), # P (stable primal cache) + # Doubling power cache + Vector{AbstractMatrix{ST}}(), # 𝐀_pow + Vector{AbstractMatrix{ST}}(), # 𝐁_pow + 0, # pow_iters + false, # pow_capture + false, # pow_transposed + zeros(TT,0,0), # P̃ + zeros(TT,0,0), # Ã_fd + zeros(TT,0,0), # B̃_fd + zeros(TT,0,0)) # C̃_fd end """ - Find_shocks_workspace(;T::Type = Float64) + Find_shocks_workspace(::Type{TT} = Float64) Create a workspace for find_shocks conditional forecast with lazy buffer allocation. All buffers are initialized to 0-dimensional objects and resized on-demand via ensure_find_shocks_buffers!. """ -function Find_shocks_workspace(;T::Type = Float64) - find_shocks_workspace{T}( +function Find_shocks_workspace(::Type{TT} = Float64) where {TT <: Real} + find_shocks_workspace{TT}( 0, # n_exo dimension - zeros(T,0), # kron_buffer (n_exo^2) - zeros(T,0,0), # kron_buffer2 (n_exo × n_exo) - zeros(T,0), # kron_buffer² (n_exo^3) - zeros(T,0,0), # kron_buffer3 (n_exo × n_exo^2) - zeros(T,0,0)) # kron_buffer4 (n_exo^2 × n_exo) -end - -function Higher_order_workspace(;T::Type = Float64, S::Type = Float64) - higher_order_workspace(spzeros(T,0,0), - spzeros(T,0,0), - spzeros(T,0,0), - spzeros(T,0,0), - spzeros(T,0,0), - spzeros(T,0,0), - (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[]), - (Int[], Int[], T[], Int[], Int[], Int[], T[]), - (Int[], Int[], T[], Int[], Int[], Int[], T[]), - zeros(T,0,0), - Sylvester_workspace(S = S), + zeros(TT,0), # kron_buffer (n_exo^2) + zeros(TT,0,0), # kron_buffer2 (n_exo × n_exo) + zeros(TT,0), # kron_buffer² (n_exo^3) + zeros(TT,0,0), # kron_buffer3 (n_exo × n_exo^2) + zeros(TT,0,0), # kron_buffer4 (n_exo^2 × n_exo) + 0, # n_past dimension + zeros(TT,0), # kron_state_vol + zeros(TT,0), # kron_state_vol3 + zeros(TT,0), # kron_state₁₂ + zeros(TT,0,0), # kron_I_state + zeros(TT,0,0), # kron_I_state₂ + zeros(TT,0,0)) # kron_I_state_state +end + +function Higher_order_workspace(::Type{TT} = Float64, ::Type{SS} = Float64) where {TT <: Real, SS <: AbstractFloat} + empty_dx_prob = 𝒮.LinearProblem(zeros(Float64, 0, 0), zeros(Float64, 0)) + empty_dx_lu_buffer = 𝒮.init(empty_dx_prob, + 𝒮.FastLUFactorization(), + verbose = isdefined(𝒮, :LinearVerbosity) ? 𝒮.LinearVerbosity(𝒮.SciMLLogging.Minimal()) : false) + empty_lu_factors = zeros(Float64, 0, 0) + empty_lu_ws = FastLapackInterface.LUWs(empty_lu_factors) + higher_order_workspace(spzeros(TT,0,0), + spzeros(TT,0,0), + spzeros(TT,0,0), + spzeros(TT,0,0), + spzeros(TT,0,0), + spzeros(TT,0,0), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + (Int[], Int[], TT[], Int[], Int[], Int[], TT[]), + zeros(TT,0,0), # 𝐒₁ + zeros(TT,0,0), # 𝐒₁₋╱𝟏ₑ + zeros(TT,0,0), + Sylvester_workspace(SS, SS), + zeros(TT,0), # ∂∇_vec # Second order pullback gradient buffers (lazily allocated) - zeros(T,0,0), # ∂∇₂ - zeros(T,0,0), # ∂∇₁ - zeros(T,0,0), # ∂𝐒₁ - zeros(T,0,0), # ∂spinv - zeros(T,0,0), # ∂𝐒₁₋╱𝟏ₑ - zeros(T,0,0), # ∂𝐒₁₊╱𝟎 - zeros(T,0,0), # ∂⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ + zeros(TT,0,0), # ∂∇₂ + zeros(TT,0,0), # ∂∇₁ + zeros(TT,0,0), # ∂𝐒₁ + zeros(TT,0,0), # ∂spinv + zeros(TT,0,0), # ∂𝐒₁₋╱𝟏ₑ + zeros(TT,0,0), # ∂𝐒₁₊╱𝟎 + zeros(TT,0,0), # ∂⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ # Third order pullback gradient buffers (only dense matrices) - zeros(T,0,0), # ∂∇₁_3rd - zeros(T,0,0), # ∂𝐒₁_3rd - zeros(T,0,0), # ∂spinv_3rd + zeros(TT,0,0), # ∂∇₁_3rd + zeros(TT,0,0), # ∂𝐒₁_3rd + zeros(TT,0,0), # ∂spinv_3rd + zeros(TT,0,0), # ∂∇₂_3rd + zeros(TT,0,0), # ∂∇₃_3rd + zeros(TT,0,0), # ∂𝐒₂_3rd + zeros(TT,0,0), # ∂𝐒₁₋╱𝟏ₑ_3rd + zeros(TT,0,0), # ∂𝐒₁₊╱𝟎_3rd + zeros(TT,0,0), # ∂⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋_3rd + # Third order pullback temporary buffers + zeros(TT,0,0), # ∂𝐒₂₊╱𝟎_3rd + zeros(TT,0,0), # ∂R_c_3rd + zeros(TT,0,0), # ∂L_c_3rd + zeros(TT,0,0), # ∂L_d_3rd + zeros(TT,0,0), # ∂R_d_3rd + zeros(TT,0,0), # ∂𝐒₂₋╱𝟎_3rd + zeros(TT,0,0), # ∂𝐒₁₋╱𝟏ₑ_t8_3rd + zeros(TT,0,0), # ∂𝐒₁₊╱𝟎_tmp_3rd + zeros(TT,0,0), # ∂𝐒₁₊╱𝟎_tk0_3rd + zeros(TT,0,0), # ∂tmpkron0_σ_3rd + zeros(TT,0,0), # ∂aux_3rd + zeros(TT,0,0), # ∂S1S1_from_ck_3rd + zeros(TT,0,0), # ∂S1p0_kron_sigma_3rd + zeros(TT,0,0), # ∂S1p0_left_3rd + zeros(TT,0,0), # ∂S1p0_right_3rd + # Third order pullback intermediate product buffers (for mul!) + zeros(TT,0,0), # ∂A_3rd + zeros(TT,0,0), # ∂B_sylv_3rd + zeros(TT,0,0), # ∂𝐗₃_3rd + zeros(TT,0,0), # ∂𝐗₃_pre_3rd + zeros(TT,0,0), # ∂out2_3rd + zeros(TT,0,0), # ∂∇₁₊_3rd + zeros(TT,0,0), # ∂∇₁₊𝐒₁➕∇₁₀_3rd + zeros(TT,0,0), # ∇₂t_∂out2_3rd + zeros(TT,0,0), # mul_tmp_3rd + # LinearSolve cache (FastLUFactorization) for SSS Newton iter ∂x \ Δx + empty_dx_lu_buffer, # dx_lu_buffer + # LinearSolve cache (FastLUFactorization) for SSS common-block tmp \ vec + let p = 𝒮.LinearProblem(zeros(Float64, 0, 0), zeros(Float64, 0)) + 𝒮.init(p, + 𝒮.FastLUFactorization(), + verbose = isdefined(𝒮, :LinearVerbosity) ? 𝒮.LinearVerbosity(𝒮.SciMLLogging.Minimal()) : false) + end, # sss_tmp_lu_buffer + # Dedicated FastLapackInterface LU workspace for SSS pullback transpose solves + empty_lu_ws, # fast_lu_ws_sss_pullback + (0, 0), # fast_lu_dims_sss_pullback + # SSS Newton iter kron! buffers + zeros(TT, 0), # x_aug_buf + zeros(TT, 0), # kron_x_aug_xx + zeros(TT, 0), # kron_x_aug_x_kron + zeros(TT, 0, 0), # kron_x_aug_I + zeros(TT, 0, 0), # kron_x_kron_I # 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 + zeros(SS,0,0), # ∂x_second_order + zeros(SS,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 + +""" + ensure_sss_kron_buffers!(ws, nPast; third_order=false) + +Lazily (re)allocate kron! buffers used by the stochastic-steady-state Newton iter +on `ws` (a `higher_order_workspace`). `nPast` is `T.nPast_not_future_and_mixed`. +The 3rd-order-only buffers are only sized when `third_order=true`. +""" +function ensure_sss_kron_buffers!(ws::higher_order_workspace{S,G,H}, nPast::Int; third_order::Bool=false) where {S <: Real, G <: AbstractFloat, H <: Real} + n_aug = nPast + 1 + length(ws.x_aug_buf) == n_aug || (ws.x_aug_buf = zeros(S, n_aug)) + length(ws.kron_x_aug_xx) == n_aug^2 || (ws.kron_x_aug_xx = zeros(S, n_aug^2)) + size(ws.kron_x_aug_I) == (n_aug * nPast, nPast) || (ws.kron_x_aug_I = zeros(S, n_aug * nPast, nPast)) + if third_order + length(ws.kron_x_aug_x_kron) == n_aug^3 || (ws.kron_x_aug_x_kron = zeros(S, n_aug^3)) + size(ws.kron_x_kron_I) == (n_aug^2 * nPast, nPast) || (ws.kron_x_kron_I = zeros(S, n_aug^2 * nPast, nPast)) + end + return ws end """ - Qme_workspace(n::Int; T::Type = Float64) + ensure_dx_lu_buffer!(ws, ∂x, Δx) + +Ensure the LinearSolve cache `ws.dx_lu_buffer` is sized for the SSS Newton iter +linear system `∂x * y = Δx`. If dimensions match, reuse the cache and just rebind +`A` and `b`; otherwise re-`init` the cache (FastLUFactorization backend). +""" +function ensure_dx_lu_buffer!(ws::higher_order_workspace, ∂x::AbstractMatrix{Float64}, Δx::AbstractVector{Float64}) + cache = ws.dx_lu_buffer + if size(cache.A) != size(∂x) || length(cache.b) != length(Δx) + prob = 𝒮.LinearProblem(∂x, Δx) + ws.dx_lu_buffer = 𝒮.init(prob, + 𝒮.FastLUFactorization(), + verbose = isdefined(𝒮, :LinearVerbosity) ? 𝒮.LinearVerbosity(𝒮.SciMLLogging.Minimal()) : false) + else + cache.A = ∂x + cache.b = Δx + end + return nothing +end + +""" + ensure_sss_tmp_lu_buffer!(ws, tmp, rhs) + +Ensure the LinearSolve cache `ws.sss_tmp_lu_buffer` is sized for the SSS common-block +solve `tmp * y = rhs`. If dimensions match, reuse the cache and just rebind `A` and `b`; +otherwise re-`init` the cache (FastLUFactorization backend). +""" +function ensure_sss_tmp_lu_buffer!(ws::higher_order_workspace, tmp::AbstractMatrix{Float64}, rhs::AbstractVector{Float64}) + cache = ws.sss_tmp_lu_buffer + if size(cache.A) != size(tmp) || length(cache.b) != length(rhs) + prob = 𝒮.LinearProblem(tmp, rhs) + ws.sss_tmp_lu_buffer = 𝒮.init(prob, + 𝒮.FastLUFactorization(), + verbose = isdefined(𝒮, :LinearVerbosity) ? 𝒮.LinearVerbosity(𝒮.SciMLLogging.Minimal()) : false) + else + cache.A = tmp + cache.b = rhs + end + return nothing +end + +function ensure_sss_pullback_fast_lu_workspace!(ws::higher_order_workspace{T}, tmp::AbstractMatrix{T}) where {T <: Union{Float32, Float64}} + dims = (size(tmp, 1), size(tmp, 2)) + if ws.fast_lu_dims_sss_pullback != dims + ws.fast_lu_ws_sss_pullback = FastLapackInterface.LUWs(tmp) + ws.fast_lu_dims_sss_pullback = dims + end + return ws.fast_lu_ws_sss_pullback, ws.fast_lu_dims_sss_pullback +end + +""" + First_order_workspace(::Type{TT} = Float64, ::Type{SS} = Float64) + +Create a pre-allocated workspace for first-order perturbation and related AD paths. +""" +function First_order_workspace(::Type{TT} = Float64, ::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(TT, 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_doubling_workspace(n::Int, ::Type{TT} = Float64, ::Type{SS} = 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 - zeros(T, n, n), # F - zeros(T, n, n), # X - zeros(T, n, n), # Y - zeros(T, n, n), # X_new - zeros(T, n, n), # Y_new - zeros(T, n, n), # E_new - zeros(T, n, n), # F_new - zeros(T, n, n), # temp1 - zeros(T, n, n), # temp2 - zeros(T, n, n), # temp3 - zeros(T, n, n), # B̄ - zeros(T, n, n), # AXX - Sylvester_workspace(S = T), # sylvester_ws +function Qme_doubling_workspace(n::Int, ::Type{TT} = Float64, ::Type{SS} = Float64) where {TT <: AbstractFloat, SS <: Real} + empty_lu_factors = zeros(TT, 0, 0) + empty_lu_ws = FastLapackInterface.LUWs(empty_lu_factors) + + qme_doubling_workspace( + zeros(TT, n, n), # E + zeros(TT, n, n), # F + zeros(TT, n, n), # X + zeros(TT, n, n), # Y + zeros(TT, n, n), # X_new + zeros(TT, n, n), # Y_new + zeros(TT, n, n), # E_new + zeros(TT, n, n), # F_new + zeros(TT, n, n), # temp1 + zeros(TT, n, n), # temp2 + zeros(TT, n, n), # temp3 + zeros(TT, n, n), # B̄ + zeros(TT, n, n), # AXX + Sylvester_workspace(TT, SS), # 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 + zeros(SS, 0, 0), # X̃ + # 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 """ - Lyapunov_workspace(n::Int; T::Type = Float64) + Schur_workspace(n::Int, nMixed::Int, nPfm::Int, nFnpm::Int, ::Type{TT} = 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, ::Type{TT} = Float64) where {TT <: Real} + 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(TT, qz_seed_size, qz_seed_size) + qz_ws = FastLapackInterface.GeneralizedSchurWs(qz_seed) + lu_seed_size = max(nPfm, 1) + lu_seed = zeros(TT, lu_seed_size, lu_seed_size) + empty_lu_ws = FastLapackInterface.LUWs(lu_seed) + schur_workspace( + zeros(TT, companion_size, companion_size), # D + zeros(TT, companion_size, companion_size), # E + zeros(TT, n, nPfm), # Ã₋ + zeros(TT, n, nFnpm), # Ã₀₊ + zeros(TT, n, nPfm), # Ã₀₋ + zeros(TT, nPfm, nPfm), # Z₁₁ + zeros(TT, nFnpm, nPfm), # Z₂₁ + zeros(TT, nPfm, nPfm), # S₁₁ + zeros(TT, nPfm, nPfm), # T₁₁ + zeros(TT, n, nPfm), # sol + zeros(TT, n, n), # temp_X2 + zeros(TT, n, n), # AXX + Vector{Bool}(undef, companion_size), # eigenselect + qz_ws, + (0, 0), + empty_lu_ws, + (0, 0), + empty_lu_ws, + (0, 0), + zeros(TT, nPfm, nFnpm), # fast_lu_rhs_t_z21 + zeros(TT, nPfm, nPfm)) # fast_lu_rhs_t_s11 +end + +""" + Lyapunov_workspace(n::Int, ::Type{TT} = Float64) Create a workspace for the Lyapunov equation solver with lazy buffer allocation. `n` is the dimension of the square matrices. Buffers are initialized to 0-dimensional objects and resized on-demand when the corresponding algorithm is used. """ -function Lyapunov_workspace(n::Int; T::Type = Float64) - lyapunov_workspace{T, T}( +function Lyapunov_workspace(n::Int, ::Type{TT} = Float64) where {TT <: Real} + lyapunov_workspace{TT, TT}( n, # dimension - zeros(T, 0, 0), # 𝐂 (doubling) - zeros(T, 0, 0), # 𝐂¹ (doubling) - zeros(T, 0, 0), # 𝐀 (doubling) - zeros(T, 0, 0), # 𝐂A (doubling) - zeros(T, 0, 0), # 𝐀² (doubling) - 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 + zeros(TT, 0, 0), # 𝐂 (doubling) + zeros(TT, 0, 0), # 𝐂¹ (doubling) + zeros(TT, 0, 0), # 𝐀 (doubling) + zeros(TT, 0, 0), # 𝐂A (doubling) + zeros(TT, 0, 0), # 𝐀² (doubling) + zeros(TT, 0, 0), # tmp̄ (Krylov) + zeros(TT, 0, 0), # 𝐗 (Krylov) + zeros(TT, 0), # b (Krylov) + Krylov.BicgstabWorkspace(0, 0, Vector{TT}), # bicgstab + Krylov.GmresWorkspace(0, 0, Vector{TT}; memory = 20), # gmres + Krylov.DqgmresWorkspace(0, 0, Vector{TT}), # dqgmres + zeros(TT, 0), # b_vech (vech-space Krylov) + Krylov.BicgstabWorkspace(0, 0, Vector{TT}), # bicgstab_vech + Krylov.GmresWorkspace(0, 0, Vector{TT}; memory = 20), # gmres_vech + Krylov.DqgmresWorkspace(0, 0, Vector{TT}), # dqgmres_vech + zeros(TT, 0, 0), # P (stable primal cache) # ForwardDiff partials buffers - zeros(T, 0, 0), # P̃ - zeros(T, 0, 0), # Ã_fd - zeros(T, 0, 0) # C̃_fd + zeros(TT, 0, 0), # P̃ + zeros(TT, 0, 0), # Ã_fd + zeros(TT, 0, 0), # C̃_fd + # Doubling power cache (sparse-aware) + Vector{AbstractMatrix{TT}}(), # 𝐀_pow + 0, # pow_iters + false, # pow_capture + false, # pow_transposed + FastLapackInterface.SchurWs(zeros(TT, 1, 1)) # schur_ws (lazily resized by gees!) ) end @@ -309,30 +691,69 @@ 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`, `:gmres`, and `:dqgmres`. """ -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 + elseif algorithm == :dqgmres + if length(ws.dqgmres.x) != n * n + ws.dqgmres = Krylov.DqgmresWorkspace(n * n, n * n, Vector{T}) + end + else + error("Invalid Krylov algorithm: $algorithm. Must be :bicgstab, :gmres, or :dqgmres") + 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 + elseif algorithm == :dqgmres + if length(ws.dqgmres_vech.x) != n_vech + ws.dqgmres_vech = Krylov.DqgmresWorkspace(n_vech, n_vech, Vector{T}) + end + end + return ws end @@ -442,39 +863,94 @@ function ensure_find_shocks_buffers!(ws::find_shocks_workspace{T}, n_exo::Int; t return ws end +""" + ensure_find_shocks_state_buffers!(ws, n_exo, n_past; third_order=false, third_order_pruning=false) + +Ensure state-related kron buffers used by `find_shocks_conditional_forecast` are sized +for the given dimensions. `n_exo` is `T.nExo`; `n_past` is `T.nPast_not_future_and_mixed`. +The `state_vol` vector has length `n_past+1`. 3rd-order-only and pruning-only buffers +are sized only when those flags are set. +""" +function ensure_find_shocks_state_buffers!(ws::find_shocks_workspace{T}, n_exo::Int, n_past::Int; + third_order::Bool = false, + third_order_pruning::Bool = false) where T + ws.n_past = n_past + n_aug = n_past + 1 + + if length(ws.kron_state_vol) != n_aug^2 + ws.kron_state_vol = zeros(T, n_aug^2) + end + if size(ws.kron_I_state) != (n_exo * n_aug, n_exo) + ws.kron_I_state = zeros(T, n_exo * n_aug, n_exo) + end + + if third_order + if length(ws.kron_state_vol3) != n_aug^3 + ws.kron_state_vol3 = zeros(T, n_aug^3) + end + if size(ws.kron_I_state_state) != (n_exo * n_aug^2, n_exo) + ws.kron_I_state_state = zeros(T, n_exo * n_aug^2, n_exo) + end + if third_order_pruning + if length(ws.kron_state₁₂) != n_past^2 + ws.kron_state₁₂ = zeros(T, n_past^2) + end + if size(ws.kron_I_state₂) != (n_exo * n_past, n_exo) + ws.kron_I_state₂ = zeros(T, n_exo * n_past, n_exo) + end + end + end + + return ws +end + """ - Inversion_workspace(;T::Type = Float64) + Inversion_workspace(::Type{TT} = Float64) Create a workspace for inversion filter computations with lazy buffer allocation. All buffers are initialized to 0-dimensional objects and resized on-demand via ensure_inversion_buffers!. """ -function Inversion_workspace(;T::Type = Float64) - inversion_workspace{T}( +function Inversion_workspace(::Type{TT} = Float64) where {TT <: Real} + inversion_workspace{TT}( 0, 0, # n_exo, n_past dimensions - zeros(T, 0), # kron_buffer (n_exo^2) - zeros(T, 0, 0), # kron_buffer2 (n_exo^2 × n_exo) - zeros(T, 0), # kron_buffer² (n_exo^3) - zeros(T, 0, 0), # kron_buffer3 (n_exo^3 × n_exo) - zeros(T, 0, 0), # kron_buffer4 (n_exo^3 × n_exo^2) - zeros(T, 0, 0), # kron_buffer_state (n_exo × n_past+1) - zeros(T, 0), # kronstate_vol ((n_past+1)^2) - zeros(T, 0), # kronaug_state ((n_past+1+n_exo)^2) - zeros(T, 0), # kron_kron_aug_state ((n_past+1+n_exo)^3) - 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) + zeros(TT, 0), # kron_buffer (n_exo^2) + zeros(TT, 0, 0), # kron_buffer2 (n_exo^2 × n_exo) + zeros(TT, 0), # kron_buffer² (n_exo^3) + zeros(TT, 0, 0), # kron_buffer3 (n_exo^3 × n_exo) + zeros(TT, 0, 0), # kron_buffer4 (n_exo^3 × n_exo^2) + zeros(TT, 0, 0), # kron_buffer_state (n_exo × n_past+1) + zeros(TT, 0), # kronstate_vol ((n_past+1)^2) + zeros(TT, 0), # kronaug_state ((n_past+1+n_exo)^2) + zeros(TT, 0), # kron_kron_aug_state ((n_past+1+n_exo)^3) + zeros(TT, 0), # state_vol (n_past+1) + zeros(TT, 0), # aug_state₁ (n_past+1+n_exo) + zeros(TT, 0), # aug_state₂ (n_past+1+n_exo) + # Estimation loop temporaries + 0, # n_cond_var + zeros(TT, 0), # shock_independent (n_cond_var) + zeros(TT, 0), # init_guess (n_exo) + zeros(TT, 0, 0), # Si_buffer (n_cond_var × n_exo) + zeros(TT, 0, 0), # jacc_buffer (n_cond_var × n_exo) + zeros(TT, 0, 0), # Si2e_buffer (n_cond_var × n_exo^2) + zeros(TT, 0), # y_obs (n_cond_var) + zeros(TT, 0), # x_shocks (n_exo) + zeros(TT, 0), # state_concat (n_past + n_exo) + zeros(TT, 0), # aug_state₃ (n_past+1+n_exo) + zeros(TT, 0), # aug_state₁̂ (n_past+1+n_exo) + zeros(TT, 0), # state²⁻_vol (n_past+1) + zeros(TT, 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) - zeros(T, 0), # ∂_tmp3 (n_past+n_exo) - zeros(T, 0, 0), # ∂𝐒t⁻ (n_past × n_past+n_exo) - zeros(T, 0, 0), # ∂data (n_past × n_periods) + zeros(TT, 0, 0), # ∂_tmp1 (n_exo × n_past+n_exo) + zeros(TT, 0, 0), # ∂_tmp2 (n_past × n_past+n_exo) + zeros(TT, 0), # ∂_tmp3 (n_past+n_exo) + zeros(TT, 0, 0), # ∂𝐒t⁻ (n_past × n_past+n_exo) + zeros(TT, 0, 0), # ∂data (n_past × n_periods) # Pullback buffers for pruned second order - zeros(T, 0, 0), # ∂𝐒ⁱ²ᵉtmp (n_exo × n_exo*n_obs) - zeros(T, 0, 0), # ∂𝐒ⁱ²ᵉtmp2 (n_obs × n_exo^2) - zeros(T, 0), # kronSλ (n_obs * n_exo) - zeros(T, 0)) # kronxS (n_exo * n_obs) + zeros(TT, 0, 0), # ∂𝐒ⁱ²ᵉtmp (n_exo × n_exo*n_obs) + zeros(TT, 0, 0), # ∂𝐒ⁱ²ᵉtmp2 (n_obs × n_exo^2) + zeros(TT, 0), # kronSλ (n_obs * n_exo) + zeros(TT, 0)) # kronxS (n_exo * n_obs) end @@ -542,37 +1018,111 @@ 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 """ - Kalman_workspace(;T::Type = Float64) + Kalman_workspace(::Type{TT} = 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) - kalman_workspace{T}( +function Kalman_workspace(::Type{TT} = Float64) where {TT <: Real} + empty_lu_factors = zeros(TT, 1, 1) + empty_lu_ws = FastLapackInterface.LUWs(empty_lu_factors) + + kalman_workspace{TT}( 0, 0, # n_obs, n_states dimensions - zeros(T, 0), # u (n_states) - zeros(T, 0), # z (n_obs) - 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), # 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(TT, 0), # u (n_states) + zeros(TT, 0), # z (n_obs) + zeros(TT, 0), # ztmp (n_obs) + zeros(TT, 0), # utmp (n_states) + zeros(TT, 0, 0), # Ctmp (n_obs × n_states) + zeros(TT, 0, 0), # 𝐁 (n_states × n_states) + zeros(TT, 0, 0), # F (n_obs × n_obs) + zeros(TT, 0, 0), # K (n_states × n_obs) + zeros(TT, 0, 0), # tmp (n_states × n_states) + zeros(TT, 0, 0), # Ptmp (n_states × n_states) + empty_lu_ws, + (0, 0), + zeros(TT, 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 +1149,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 +1164,43 @@ 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) - workspaces(Higher_order_workspace(T = T, S = S), - Higher_order_workspace(T = T, S = S), +function Workspaces(::Type{T} = Float64, ::Type{S} = Float64) where {T <: Real, S <: Real} + workspaces(Higher_order_workspace(T, S), + Higher_order_workspace(T, S), Float64[], - Qme_workspace(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 - Sylvester_workspace(S = S), # 1st order 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 + First_order_workspace(T, S), # Initialize with size 0, will be resized when needed + Qme_doubling_workspace(0, T, S), # Initialize with size 0, will be resized when needed + Schur_workspace(0, 0, 0, 0, T), # Initialize with size 0, will be resized when needed + Lyapunov_workspace(0, T), # 1st order - will be resized + Lyapunov_workspace(0, T), # 2nd order - will be resized + Lyapunov_workspace(0, T), # 3rd order - will be resized + Lyapunov_workspace(0, T), # block-triangular inner - will be resized + Sylvester_workspace(S), # 1st order sylvester - will be resized + Sylvester_workspace(S), # block-triangular sylvester - will be resized + Find_shocks_workspace(T), # conditional forecast - will be resized + Inversion_workspace(T), # inversion filter - will be resized + Kalman_workspace(T), # Kalman filter - will be resized + NSSSSolverWorkspace()) # NSSS solver scratch buffers +end + +""" +Create a fresh `Workspaces` instance that preserves the NSSS solver workspace from `orig`. + +The NSSS solver workspace contains buffers sized during model compilation that cannot be +lazily re-created, unlike other workspace fields which are lazily allocated via `ensure_*!`. +""" +function fresh_workspaces(orig::workspaces) + ws = Workspaces() + ws.nsss_solver = orig.nsss_solver + return ws end function Constants(model_struct; T::Type = Float64, S::Type = Float64) @@ -635,7 +1208,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 +1243,8 @@ function Constants(model_struct; T::Type = Float64, S::Type = Float64) spzeros(Float64, 0, 0), Symbol[], Symbol[], + Int[], + Int[], Symbol[], # Symbol[], Int[], @@ -680,16 +1257,34 @@ 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) +function axis_has_string(axis) axis === nothing && return false T = eltype(axis) if T === String @@ -702,31 +1297,63 @@ function _axis_has_string(axis) return false end -function _choose_axis_type(var_axis, calib_axis, exo_axis_plain, exo_axis_with_subscript, full_NSSS_display) - return (_axis_has_string(var_axis) || - _axis_has_string(calib_axis) || - _axis_has_string(exo_axis_plain) || - _axis_has_string(exo_axis_with_subscript) || - _axis_has_string(full_NSSS_display)) ? String : Symbol +function choose_axis_type(var_axis, calib_axis, exo_axis_plain, exo_axis_with_subscript, full_NSSS_display) + return (axis_has_string(var_axis) || + axis_has_string(calib_axis) || + axis_has_string(exo_axis_plain) || + axis_has_string(exo_axis_with_subscript) || + axis_has_string(full_NSSS_display)) ? String : Symbol end -function _convert_axis(axis, ::Type{S}) where {S <: Union{Symbol, String}} +function convert_axis(axis, ::Type{S}) where {S <: Union{Symbol, String}} axis === nothing && return Vector{S}() return S === String ? string.(axis) : Symbol.(axis) end +""" + update_post_parameters_macro(p::post_parameters_macro; kwargs...) + +Return a new `post_parameters_macro` where the fields supplied as keyword +arguments replace the corresponding entries on `p`. All other fields are +preserved. Accepted keyword arguments: + +- `parameters_as_function_of_parameters` +- `precompile` +- `ss_symbolic_mode` +- `ss_solver_parameters_algorithm` +- `ss_solver_parameters_maxtime` +- `guess` +- `ss_calib_list` +- `par_calib_list` +- `bounds` +""" +function update_post_parameters_macro(p::post_parameters_macro; kwargs...) + return post_parameters_macro( + get(kwargs, :parameters_as_function_of_parameters, p.parameters_as_function_of_parameters), + get(kwargs, :precompile, p.precompile), + get(kwargs, :ss_symbolic_mode, p.ss_symbolic_mode), + get(kwargs, :ss_solver_parameters_algorithm, p.ss_solver_parameters_algorithm), + get(kwargs, :ss_solver_parameters_maxtime, p.ss_solver_parameters_maxtime), + get(kwargs, :guess, p.guess), + get(kwargs, :ss_calib_list, p.ss_calib_list), + get(kwargs, :par_calib_list, p.par_calib_list), + get(kwargs, :bounds, p.bounds), + ) +end + + function update_post_complete_parameters(p::post_complete_parameters; kwargs...) var_axis_in = get(kwargs, :var_axis, p.var_axis) calib_axis_in = get(kwargs, :calib_axis, p.calib_axis) exo_axis_plain_in = get(kwargs, :exo_axis_plain, p.exo_axis_plain) exo_axis_with_subscript_in = get(kwargs, :exo_axis_with_subscript, p.exo_axis_with_subscript) full_NSSS_display_in = get(kwargs, :full_NSSS_display, p.full_NSSS_display) - S = _choose_axis_type(var_axis_in, calib_axis_in, exo_axis_plain_in, exo_axis_with_subscript_in, full_NSSS_display_in) - var_axis = _convert_axis(var_axis_in, S) - calib_axis = _convert_axis(calib_axis_in, S) - exo_axis_plain = _convert_axis(exo_axis_plain_in, S) - exo_axis_with_subscript = _convert_axis(exo_axis_with_subscript_in, S) - full_NSSS_display = _convert_axis(full_NSSS_display_in, S) + S = choose_axis_type(var_axis_in, calib_axis_in, exo_axis_plain_in, exo_axis_with_subscript_in, full_NSSS_display_in) + var_axis = convert_axis(var_axis_in, S) + calib_axis = convert_axis(calib_axis_in, S) + exo_axis_plain = convert_axis(exo_axis_plain_in, S) + exo_axis_with_subscript = convert_axis(exo_axis_with_subscript_in, S) + full_NSSS_display = convert_axis(full_NSSS_display_in, S) return post_complete_parameters{S}( get(kwargs, :parameters, p.parameters), get(kwargs, :missing_parameters, p.missing_parameters), @@ -752,6 +1379,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 +1393,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 @@ -822,7 +1469,7 @@ function ensure_name_display_constants!(𝓂) ) end - return constants.post_complete_parameters + return nothing end @@ -862,61 +1509,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 +1563,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) @@ -1033,8 +1575,10 @@ function ensure_conditional_forecast_constants!(constants::constants; third_orde 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 + shockvar_no_vol_idxs = sparse(ℒ.kron(e_in_s⁺, s_in_s⁺)).nzind so.var²_idxs = var²_idxs so.shockvar²_idxs = shockvar²_idxs + so.shockvar_no_vol_idxs = shockvar_no_vol_idxs so.var_vol²_idxs = var_vol²_idxs end @@ -1109,6 +1653,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 +1688,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 +1724,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 + return nothing 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. """ - ensure_qme_workspace!(𝓂) - ensure_qme_workspace!(workspaces, n) +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 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. """ -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) + ensure_first_order_workspace_buffers!(ws, T, n_dyn, n_comb) + +Ensure all first-order perturbation buffers in `first_order_workspace` are allocated with +the correct dimensions. +""" +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) + end + return ws end - """ ensure_lyapunov_workspace!(workspaces, n, order::Symbol) @@ -1253,50 +1928,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 +1979,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 +2005,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, @@ -1368,35 +2015,38 @@ function ensure_model_structure_constants!(constants::constants, calibration_par ) end - return constants.post_complete_parameters + return nothing end 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 +2109,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 +2163,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 +2174,253 @@ 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. + +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 + +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 + +# Generic SolverTolerances keyword-constructor defaults, used by resolve_tol to detect +# which fields were left at their generic value and should be replaced by context defaults. +const GENERIC_SOLVER_TOL = SolverTolerances(1e-14, 1e-14, 1e-10, 1e-10) + +# Merge a SolverTolerances with context-specific defaults. Fields that still sit at the +# generic SolverTolerances() defaults are replaced by the context base; explicitly changed +# fields are kept. +function resolve_tol(override::SolverTolerances, base::SolverTolerances) + SolverTolerances( + override.atol != GENERIC_SOLVER_TOL.atol ? override.atol : base.atol, + override.rtol != GENERIC_SOLVER_TOL.rtol ? override.rtol : base.rtol, + override.initial_guess_acceptance_tol != GENERIC_SOLVER_TOL.initial_guess_acceptance_tol ? override.initial_guess_acceptance_tol : base.initial_guess_acceptance_tol, + override.acceptance_tol != GENERIC_SOLVER_TOL.acceptance_tol ? override.acceptance_tol : base.acceptance_tol) +end + +# Merge a NamedTuple of partial overrides with a SolverTolerances base. +function resolve_tol(nt::NamedTuple, base::SolverTolerances) + SolverTolerances( + Float64(get(nt, :atol, base.atol)), + Float64(get(nt, :rtol, base.rtol)), + Float64(get(nt, :initial_guess_acceptance_tol, base.initial_guess_acceptance_tol)), + Float64(get(nt, :acceptance_tol, base.acceptance_tol))) +end + +""" + 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 + +resolve_tol(s::NsssTolerances, ::NsssTolerances) = s +function resolve_tol(nt::NamedTuple, base::NsssTolerances) + NsssTolerances( + Float64(get(nt, :acceptance_tol, base.acceptance_tol)), + Float64(get(nt, :initial_guess_acceptance_tol, base.initial_guess_acceptance_tol)), + Float64(get(nt, :xtol, base.xtol)), + Float64(get(nt, :ftol, base.ftol)), + Float64(get(nt, :rel_xtol, base.rel_xtol))) +end - qme_tol::AbstractFloat - qme_acceptance_tol::AbstractFloat +""" + AdTolerances - sylvester_tol::AbstractFloat - sylvester_acceptance_tol::AbstractFloat +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. - lyapunov_tol::AbstractFloat - lyapunov_acceptance_tol::AbstractFloat +# Fields +- `qme::SolverTolerances`: tolerances for the quadratic matrix equation (QME) derivative solve. + Default: `atol=1e-14`, `rtol=1e-14`, `initial_guess_acceptance_tol=1e-10`, `acceptance_tol=1e-6`. +- `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`. - droptol::AbstractFloat +Construct via `AdTolerances(; qme, sylvester, lyapunov)`. +""" +struct AdTolerances + qme::SolverTolerances + sylvester::SolverTolerances + lyapunov::SolverTolerances +end - dependencies_tol::AbstractFloat +function AdTolerances(; qme = (;), sylvester = (;), lyapunov = (;)) + _base_qme = SolverTolerances(1e-14, 1e-14, 1e-10, 1e-6) + _base_sylv = SolverTolerances(1e-14, 1e-14, 1e-10, 1e-10) + _base_lyap = SolverTolerances(1e-14, 1e-14, 1e-12, 1e-12) + return AdTolerances( + resolve_tol(qme, _base_qme), + resolve_tol(sylvester, _base_sylv), + resolve_tol(lyapunov, _base_lyap)) +end + +resolve_tol(s::AdTolerances, ::AdTolerances) = s +function resolve_tol(nt::NamedTuple, base::AdTolerances) + AdTolerances( + resolve_tol(get(nt, :qme, (;)), base.qme), + resolve_tol(get(nt, :sylvester, (;)), base.sylvester), + resolve_tol(get(nt, :lyapunov, (;)), base.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-10`, `acceptance_tol=1e-6`. +- `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 = (;), + lyapunov = (;), + droptol::Float64 = 1e-14, + dependencies_tol::Float64 = 1e-12, + ad = (;)) + _base_qme = SolverTolerances(1e-14, 1e-14, 1e-10, 1e-6) + _base_lyap = SolverTolerances(1e-14, 1e-14, 1e-12, 1e-12) + _base_ad = AdTolerances() + return FirstOrderTolerances( + resolve_tol(qme, _base_qme), + resolve_tol(lyapunov, _base_lyap), + droptol, dependencies_tol, + resolve_tol(ad, _base_ad)) +end + +resolve_tol(s::FirstOrderTolerances, ::FirstOrderTolerances) = s +function resolve_tol(nt::NamedTuple, base::FirstOrderTolerances) + FirstOrderTolerances( + resolve_tol(get(nt, :qme, (;)), base.qme), + resolve_tol(get(nt, :lyapunov, (;)), base.lyapunov), + Float64(get(nt, :droptol, base.droptol)), + Float64(get(nt, :dependencies_tol, base.dependencies_tol)), + resolve_tol(get(nt, :ad, (;)), base.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 + +function HigherOrderTolerances(; sylvester = (;), + lyapunov = (;), + droptol::Float64 = 1e-14, + dependencies_tol::Float64 = 1e-12, + ad = (;)) + _base_sylv = SolverTolerances(1e-14, 1e-14, 1e-10, 1e-10) + _base_lyap = SolverTolerances(1e-14, 1e-14, 1e-12, 1e-12) + _base_ad = AdTolerances() + return HigherOrderTolerances( + resolve_tol(sylvester, _base_sylv), + resolve_tol(lyapunov, _base_lyap), + droptol, dependencies_tol, + resolve_tol(ad, _base_ad)) +end + +resolve_tol(s::HigherOrderTolerances, ::HigherOrderTolerances) = s +function resolve_tol(nt::NamedTuple, base::HigherOrderTolerances) + HigherOrderTolerances( + resolve_tol(get(nt, :sylvester, (;)), base.sylvester), + resolve_tol(get(nt, :lyapunov, (;)), base.lyapunov), + Float64(get(nt, :droptol, base.droptol)), + Float64(get(nt, :dependencies_tol, base.dependencies_tol)), + resolve_tol(get(nt, :ad, (;)), base.ad)) +end + +struct Tolerances + nsss::NsssTolerances + first_order::FirstOrderTolerances + second_order::HigherOrderTolerances + third_order::HigherOrderTolerances end struct CalculationOptions @@ -1506,72 +2429,113 @@ struct CalculationOptions sylvester_algorithm²::Symbol sylvester_algorithm³::Symbol + sylvester_preconditioner::Symbol + lyapunov_algorithm::Symbol tol::Tolerances verbose::Bool 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). - -# 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. - -- `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. - -- `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. -- `droptol` [Default: `1e-14`, Type: `AbstractFloat`]: Tolerance below which matrix entries are considered 0. +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, except sylvester and lyapunov acceptance_tol default to 1e-8 on primal and AD paths) +``` + +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. -- `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 +# Keyword Arguments +- `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 first-order Lyapunov acceptance_tol while preserving context defaults +tol = Tolerances(first_order = FirstOrderTolerances(lyapunov = SolverTolerances(acceptance_tol = 1e-14))) + +# concise NamedTuple syntax — equivalent to the above +tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),)) + +# tighten second- and third-order Sylvester/Lyapunov solvers +tol = Tolerances( + second_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14), + lyapunov = SolverTolerances(acceptance_tol = 1e-14)), + third_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14), + lyapunov = SolverTolerances(acceptance_tol = 1e-14)), +) +``` """ -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, - - sylvester_tol::AbstractFloat = 1e-14, - sylvester_acceptance_tol::AbstractFloat = 1e-10, +function Tolerances(; nsss = (;), + first_order = (;), + second_order = (;), + third_order = (;)) + _base = Tolerances(NsssTolerances(), FirstOrderTolerances(), HigherOrderTolerances(), + HigherOrderTolerances( + sylvester = SolverTolerances(acceptance_tol = 1e-8), + lyapunov = SolverTolerances(acceptance_tol = 1e-8), + ad = AdTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-8), + lyapunov = SolverTolerances(acceptance_tol = 1e-8)))) + return Tolerances( + resolve_tol(nsss, _base.nsss), + resolve_tol(first_order, _base.first_order), + resolve_tol(second_order, _base.second_order), + resolve_tol(third_order, _base.third_order)) +end - lyapunov_tol::AbstractFloat = 1e-14, - lyapunov_acceptance_tol::AbstractFloat = 1e-12, +resolve_tol(s::Tolerances, ::Tolerances) = s +function resolve_tol(nt::NamedTuple, base::Tolerances) + Tolerances( + resolve_tol(get(nt, :nsss, (;)), base.nsss), + resolve_tol(get(nt, :first_order, (;)), base.first_order), + resolve_tol(get(nt, :second_order, (;)), base.second_order), + resolve_tol(get(nt, :third_order, (;)), base.third_order)) +end - droptol::AbstractFloat = 1e-14, - 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, +function merge_calculation_options(;quadratic_matrix_equation_algorithm::Symbol = DEFAULT_QME_ALGORITHM, sylvester_algorithm²::Symbol = :doubling, sylvester_algorithm³::Symbol = :bicgstab, + sylvester_preconditioner::Symbol = :ilu, lyapunov_algorithm::Symbol = :doubling, tol::Tolerances = Tolerances(), verbose::Bool = false) @@ -1579,9 +2543,11 @@ function merge_calculation_options(;quadratic_matrix_equation_algorithm::Symbol return CalculationOptions(quadratic_matrix_equation_algorithm, sylvester_algorithm², sylvester_algorithm³, + sylvester_preconditioner, lyapunov_algorithm, tol, verbose) end -end # dispatch_doctor + +end # @stable diff --git a/src/parser/equation_modification.jl b/src/parser/equation_modification.jl new file mode 100644 index 000000000..dab70710d --- /dev/null +++ b/src/parser/equation_modification.jl @@ -0,0 +1,1275 @@ +# Equation-modification reprocessing pipeline. +# +# The `@model` and `@parameters` macros build a model struct and compile +# helpers. After an equation modification this file rebuilds the relevant +# internal state surgically — without re-evaluating the macros — by calling +# the pure functions `process_model_equations` and +# `process_parameter_definitions`. +# +# The file is organised as follows: +# 1. User-facing API — exported functions with full docstrings. +# 2. Internal helpers — reprocessing pipeline and small utilities. + +# ======================================================================== +# 1. User-facing API +# ======================================================================== + +const EquationInput = Union{Expr, String} +const EquationOrIndex = Union{Int, Expr, String} + + +""" +$(SIGNATURES) +Replace an existing model equation with a new one. + +The first equation argument selects which equation to update: pass either the +1-based index, the existing equation as an `Expr`, or as a `String`. The new +equation can be passed as an `Expr` or a `String`. + +A batched form is also supported: pass a `Vector` (or `Tuple`) of +`(old_or_index, new_equation)` `Pair`s or 2-tuples to apply several updates +in one rebuild. + +After the update the revision history is appended (see +[`get_revision_history`](@ref)), all caches are invalidated and — if all +parameters are defined — the non-stochastic steady state is resolved and the +first-order symbolic derivatives are rewritten. + +[`replace_equations!`](@ref) is an alias of this function. + +# Arguments +- $MODEL® +- `old_equation_or_index` [Type: `Union{Int, Expr, String}`]: the equation to + replace, identified by its 1-based index, by the equation itself as an + `Expr`, or by the equation as a `String` (parsed before matching). When + matching by `Expr`/`String`, comparison is done on a whitespace- and + brace-insensitive canonical form. +- `new_equation` [Type: `Union{Expr, String}`]: the replacement equation. + +Alternatively a single positional argument can be supplied: +- `updates` [Type: `Union{Vector, Tuple}`]: a collection of + `(old_or_index => new_equation)` `Pair`s or 2-tuples. + +# Keyword Arguments +- $PARAMETERS® +- $VERBOSE® +- `silent` [Default: `true`, Type: `Bool`]: suppress informational warnings + from the rebuild pipeline (e.g. about missing parameters). + +# Returns +- `nothing`. The model `𝓂` is updated in place. + +# Examples +```julia +using MacroModelling + +@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 + +# replace by index +update_equations!(RBC, 3, :(q[0] = exp(z[0]) * k[-1]^α + 0)) + +# replace by old equation expression +update_equations!(RBC, + :(z[0] = ρ * z[-1] + std_z * eps_z[x]), + :(z[0] = ρ * z[-1] + std_z * eps_z[x] + 0)) + +# batched update +update_equations!(RBC, [ + 1 => :(1 / c[0] = (β / c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ))), + 2 => :(c[0] + k[0] = (1 - δ) * k[-1] + q[0]), +]) +``` +""" +function update_equations!(𝓂::ℳ, + old_equation_or_index::EquationOrIndex, + new_equation::EquationInput; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + new_eq = normalize_equation_input(new_equation)::Expr + + originals = copy(𝓂.equations.original) + idx::Int = if old_equation_or_index isa Int + @assert 1 <= old_equation_or_index <= length(originals) "Equation index $(old_equation_or_index) out of bounds (1:$(length(originals)))." + old_equation_or_index + else + target = normalize_equation_input(old_equation_or_index)::Expr + found = find_equation_index(originals, target) + @assert found !== nothing "Equation not found in model: $(target)" + found + end + + old_eq = originals[idx] + originals[idx] = new_eq + + push!(𝓂.revision_history, revision_entry(:update_equation; + equation_index = idx, old_equation = old_eq, new_equation = new_eq)) + + reprocess_model_equations!(𝓂, originals; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + +function update_equations!(𝓂::ℳ, + updates::Union{Vector, Tuple}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + originals = copy(𝓂.equations.original) + history_entries = RevisionEntry[] + for upd in updates + @assert upd isa Union{Tuple, Pair} && length(upd) == 2 "Each update entry must be a (old_or_index, new_equation) pair/tuple." + first_el = upd isa Pair ? upd.first : upd[1] + second_el = upd isa Pair ? upd.second : upd[2] + new_eq = normalize_equation_input(second_el)::Expr + idx::Int = if first_el isa Int + @assert 1 <= first_el <= length(originals) "Equation index $(first_el) out of bounds." + first_el + else + target = normalize_equation_input(first_el)::Expr + found = find_equation_index(originals, target) + @assert found !== nothing "Equation not found in model: $(target)" + found + end + old_eq = originals[idx] + originals[idx] = new_eq + push!(history_entries, revision_entry(:update_equation; + equation_index = idx, old_equation = old_eq, new_equation = new_eq)) + end + append!(𝓂.revision_history, history_entries) + reprocess_model_equations!(𝓂, originals; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + + +""" +$(SIGNATURES) +Append one or more equations to the model. + +Pass either a single equation (`Expr` or `String`) or a `Vector`/`Tuple` of +equations to add several at once. Each addition is recorded in the revision +history (see [`get_revision_history`](@ref)). + +# Arguments +- $MODEL® +- `new_equation` [Type: `Union{Expr, String}`]: the equation to add. + +Alternatively: +- `new_equations` [Type: `Union{Vector, Tuple}`]: a collection of equations + to add in one rebuild. + +# Keyword Arguments +- $PARAMETERS® +- $VERBOSE® +- `silent` [Default: `true`, Type: `Bool`]: suppress informational warnings + from the rebuild pipeline. + +# Returns +- `nothing`. The model `𝓂` is updated in place. + +# Examples +```julia +using MacroModelling + +@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 + +add_equation!(RBC, :(y[0] = c[0] + k[0] - (1 - δ) * k[-1])) + +add_equation!(RBC, [ + :(inv[0] = k[0] - (1 - δ) * k[-1]), + :(log_c[0] = log(c[0])), +]) +``` +""" +function add_equation!(𝓂::ℳ, + new_equation::EquationInput; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + new_eq = normalize_equation_input(new_equation)::Expr + originals = copy(𝓂.equations.original) + push!(originals, new_eq) + push!(𝓂.revision_history, revision_entry(:add_equation; + equation_index = length(originals), old_equation = nothing, new_equation = new_eq)) + reprocess_model_equations!(𝓂, originals; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + +function add_equation!(𝓂::ℳ, + new_equations::Union{Vector, Tuple}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + originals = copy(𝓂.equations.original) + entries = RevisionEntry[] + for ne in new_equations + new_eq = normalize_equation_input(ne)::Expr + push!(originals, new_eq) + push!(entries, revision_entry(:add_equation; + equation_index = length(originals), old_equation = nothing, new_equation = new_eq)) + end + append!(𝓂.revision_history, entries) + reprocess_model_equations!(𝓂, originals; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + + +""" +$(SIGNATURES) +Remove one or more equations from the model. + +Equations can be selected by 1-based index, by their `Expr`, or by their +`String` representation. Pass a `Vector`/`Tuple` of selectors to remove +several equations in one rebuild. The model must always retain at least one +equation. + +# Arguments +- $MODEL® +- `equation_or_index` [Type: `Union{Int, Expr, String}`]: the equation to + remove. When matching by `Expr`/`String`, comparison is done on a + whitespace- and brace-insensitive canonical form. + +Alternatively: +- `removals` [Type: `Union{Vector, Tuple}`]: a collection of selectors. + +# Keyword Arguments +- $PARAMETERS® +- $VERBOSE® +- `silent` [Default: `true`, Type: `Bool`]: suppress informational warnings + from the rebuild pipeline. + +# Returns +- `nothing`. The model `𝓂` is updated in place. + +# Examples +```julia +using MacroModelling + +@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] + y[0] = c[0] + k[0] - (1 - δ) * k[-1] +end + +@parameters RBC begin + std_z = 0.01 + ρ = 0.2 + δ = 0.02 + α = 0.5 + β = 0.95 +end + +# remove by index +remove_equation!(RBC, 5) + +# remove by expression +remove_equation!(RBC, :(q[0] = exp(z[0]) * k[-1]^α)) +``` +""" +function remove_equation!(𝓂::ℳ, + equation_or_index::EquationOrIndex; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + originals = copy(𝓂.equations.original) + @assert length(originals) > 1 "Cannot remove the last equation from the model." + idx::Int = if equation_or_index isa Int + @assert 1 <= equation_or_index <= length(originals) "Equation index $(equation_or_index) out of bounds." + equation_or_index + else + target = normalize_equation_input(equation_or_index)::Expr + found = find_equation_index(originals, target) + @assert found !== nothing "Equation not found in model: $(target)" + found + end + old_eq = originals[idx] + deleteat!(originals, idx) + push!(𝓂.revision_history, revision_entry(:remove_equation; + equation_index = idx, old_equation = old_eq, new_equation = nothing)) + reprocess_model_equations!(𝓂, originals; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + +function remove_equation!(𝓂::ℳ, + removals::Union{Vector, Tuple}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + originals = copy(𝓂.equations.original) + indices = Int[] + old_eqs = Expr[] + for item in removals + @assert length(originals) - length(indices) > 1 "Cannot remove the last equation from the model." + idx::Int = if item isa Int + @assert 1 <= item <= length(originals) "Equation index $(item) out of bounds." + item + else + target = normalize_equation_input(item)::Expr + found = find_equation_index(originals, target) + @assert found !== nothing "Equation not found in model: $(target)" + found + end + push!(indices, idx) + push!(old_eqs, originals[idx]) + end + order = sortperm(indices, rev = true) + updated = copy(originals) + for i in order + deleteat!(updated, indices[i]) + end + entries = RevisionEntry[] + for (i, idx) in enumerate(indices) + push!(entries, revision_entry(:remove_equation; + equation_index = idx, old_equation = old_eqs[i], new_equation = nothing)) + end + append!(𝓂.revision_history, entries) + reprocess_model_equations!(𝓂, updated; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + + +""" +$(SIGNATURES) +Replace an existing calibration equation. + +Calibration equations use the `lhs = rhs | param` syntax and pin down a +parameter so it is solved jointly with the non-stochastic steady state. The +new equation must contain the `| param` marker and reference a parameter +that is already part of the model. + +A batched form is supported: pass a `Vector`/`Tuple` of +`(old_or_index => new_equation)` `Pair`s or 2-tuples. + +[`replace_calibration_equations!`](@ref) is an alias of this function. + +# Arguments +- $MODEL® +- `old_equation_or_index` [Type: `Union{Int, Expr, String}`]: the calibration + equation to replace, identified by its 1-based index in the calibration + list or by its `Expr`/`String` representation. +- `new_equation` [Type: `Union{Expr, String}`]: the replacement calibration + equation. Must contain the `| param` syntax. + +Alternatively: +- `updates` [Type: `Union{Vector, Tuple}`]: a collection of + `(old_or_index => new_equation)` `Pair`s or 2-tuples. + +# Keyword Arguments +- $PARAMETERS® +- $VERBOSE® +- `silent` [Default: `true`, Type: `Bool`]: suppress informational warnings + from the rebuild pipeline. + +# Returns +- `nothing`. The model `𝓂` is updated in place. + +# Examples +```julia +using MacroModelling + +@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.5 + β = 0.95 + k[ss] / (4 * q[ss]) = 1.5 | δ +end + +# retarget the calibration equation +update_calibration_equations!(RBC, 1, :(k[ss] / (4 * q[ss]) = 2.0 | δ)) +``` +""" +function update_calibration_equations!(𝓂::ℳ, + old_equation_or_index::EquationOrIndex, + new_equation::EquationInput; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + new_eq = normalize_equation_input(new_equation)::Expr + @assert extract_calibrated_parameter(new_eq) !== nothing "Calibration equation must contain `| param` syntax." + + calib_orig = copy(𝓂.equations.calibration_original) + idx::Int = if old_equation_or_index isa Int + @assert 1 <= old_equation_or_index <= length(calib_orig) "Calibration index $(old_equation_or_index) out of bounds." + old_equation_or_index + else + target = normalize_equation_input(old_equation_or_index)::Expr + found = find_equation_index(calib_orig, target) + @assert found !== nothing "Calibration equation not found: $(target)" + found + end + + new_param = extract_calibrated_parameter(new_eq) + known_params = Set{Symbol}(𝓂.constants.post_model_macro.parameters_in_equations) + union!(known_params, 𝓂.constants.post_complete_parameters.parameters) + union!(known_params, 𝓂.equations.calibration_parameters) + new_param !== nothing && !(new_param in known_params) && + error("Parameter `$(new_param)` is not part of the model.") + + old_eq = calib_orig[idx] + calib_orig[idx] = new_eq + push!(𝓂.revision_history, revision_entry(:update_calibration_equation; + equation_index = idx, old_equation = old_eq, new_equation = new_eq)) + reprocess_calibration_equations!(𝓂, calib_orig; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + +function update_calibration_equations!(𝓂::ℳ, + updates::Union{Vector, Tuple}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + calib_orig = copy(𝓂.equations.calibration_original) + entries = RevisionEntry[] + for upd in updates + first_el = upd isa Pair ? upd.first : upd[1] + second_el = upd isa Pair ? upd.second : upd[2] + new_eq = normalize_equation_input(second_el)::Expr + @assert extract_calibrated_parameter(new_eq) !== nothing "Calibration equation must contain `| param` syntax." + idx::Int = if first_el isa Int + @assert 1 <= first_el <= length(calib_orig) "Calibration index $(first_el) out of bounds." + first_el + else + target = normalize_equation_input(first_el)::Expr + found = find_equation_index(calib_orig, target) + @assert found !== nothing "Calibration equation not found: $(target)" + found + end + old_eq = calib_orig[idx] + calib_orig[idx] = new_eq + push!(entries, revision_entry(:update_calibration_equation; + equation_index = idx, old_equation = old_eq, new_equation = new_eq)) + end + append!(𝓂.revision_history, entries) + reprocess_calibration_equations!(𝓂, calib_orig; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + + +""" +$(SIGNATURES) +Add a new calibration equation to the model. + +The equation must use the `lhs = rhs | param` syntax, and `param` must +already be part of the model and not yet calibrated. A batched form is +supported by passing a `Vector`/`Tuple` of equations. + +# Arguments +- $MODEL® +- `new_equation` [Type: `Union{Expr, String}`]: the calibration equation to + add. Must contain the `| param` syntax. + +Alternatively: +- `new_equations` [Type: `Union{Vector, Tuple}`]: a collection of + calibration equations to add in one rebuild. + +# Keyword Arguments +- $PARAMETERS® +- $VERBOSE® +- `silent` [Default: `true`, Type: `Bool`]: suppress informational warnings + from the rebuild pipeline. + +# Returns +- `nothing`. The model `𝓂` is updated in place. + +# Examples +```julia +using MacroModelling + +@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 + +# pin down β by targeting a value for c at the steady state +add_calibration_equation!(RBC, :(c[ss] = 1.0 | β)) +``` +""" +function add_calibration_equation!(𝓂::ℳ, + new_equation::EquationInput; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + new_eq = normalize_equation_input(new_equation)::Expr + new_param = extract_calibrated_parameter(new_eq) + new_param === nothing && error("Calibration equation must contain `| param` syntax.") + new_param in 𝓂.equations.calibration_parameters && + error("Parameter `$(new_param)` is already calibrated.") + known_params = Set{Symbol}(𝓂.constants.post_model_macro.parameters_in_equations) + union!(known_params, 𝓂.constants.post_complete_parameters.parameters) + !(new_param in known_params) && error("Parameter `$(new_param)` is not part of the model.") + + calib_orig = copy(𝓂.equations.calibration_original) + push!(calib_orig, new_eq) + push!(𝓂.revision_history, revision_entry(:add_calibration_equation; + equation_index = length(calib_orig), old_equation = nothing, new_equation = new_eq)) + reprocess_calibration_equations!(𝓂, calib_orig; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + +function add_calibration_equation!(𝓂::ℳ, + new_equations::Union{Vector, Tuple}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + calib_orig = copy(𝓂.equations.calibration_original) + entries = RevisionEntry[] + for ne in new_equations + new_eq = normalize_equation_input(ne)::Expr + new_param = extract_calibrated_parameter(new_eq) + new_param === nothing && error("Calibration equation must contain `| param` syntax.") + push!(calib_orig, new_eq) + push!(entries, revision_entry(:add_calibration_equation; + equation_index = length(calib_orig), old_equation = nothing, new_equation = new_eq)) + end + append!(𝓂.revision_history, entries) + reprocess_calibration_equations!(𝓂, calib_orig; parameters = parameters, + verbose = verbose, silent = silent) + return nothing +end + + +""" +$(SIGNATURES) +Remove a calibration equation. + +When a calibration equation is removed, the parameter previously solved by +it becomes a free parameter and needs a numeric value. By default that value +is taken from the parameter's current non-stochastic steady state value. Use +the `parameters` keyword to override this with explicit values. + +A batched form is supported by passing a `Vector`/`Tuple` of selectors. + +# Arguments +- $MODEL® +- `equation_or_index` [Type: `Union{Int, Expr, String}`]: the calibration + equation to remove. + +Alternatively: +- `removals` [Type: `Union{Vector, Tuple}`]: a collection of selectors. + +# Keyword Arguments +- `parameters` [Default: `nothing`]: optional value(s) used as the new fixed + value for the parameter freed by the removal. Accepts the same forms as + the `parameters` argument elsewhere — a `Pair`, a `Dict`, or a + `Vector`/`Tuple` of `Pair`s — but only the entries matching freed + parameters are used. +- $VERBOSE® +- `silent` [Default: `true`, Type: `Bool`]: suppress informational warnings + from the rebuild pipeline. + +# Returns +- `nothing`. The model `𝓂` is updated in place. + +# Examples +```julia +using MacroModelling + +@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.5 + β = 0.95 + k[ss] / (4 * q[ss]) = 1.5 | δ +end + +# remove the calibration equation, fixing δ to 0.02 +remove_calibration_equation!(RBC, 1, parameters = :δ => 0.02) +``` +""" +function remove_calibration_equation!(𝓂::ℳ, + equation_or_index::EquationOrIndex; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + calib_orig = copy(𝓂.equations.calibration_original) + @assert !isempty(calib_orig) "No calibration equations to remove." + idx::Int = if equation_or_index isa Int + @assert 1 <= equation_or_index <= length(calib_orig) "Calibration index $(equation_or_index) out of bounds." + equation_or_index + else + target = normalize_equation_input(equation_or_index)::Expr + found = find_equation_index(calib_orig, target) + @assert found !== nothing "Calibration equation not found: $(target)" + found + end + old_eq = calib_orig[idx] + deleteat!(calib_orig, idx) + push!(𝓂.revision_history, revision_entry(:remove_calibration_equation; + equation_index = idx, old_equation = old_eq, new_equation = nothing)) + + param_overrides = parameters_to_dict(parameters) + reprocess_calibration_equations!(𝓂, calib_orig; parameters = nothing, + parameter_overrides = param_overrides, verbose = verbose, silent = silent) + return nothing +end + +function remove_calibration_equation!(𝓂::ℳ, + removals::Union{Vector, Tuple}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + calib_orig = copy(𝓂.equations.calibration_original) + @assert !isempty(calib_orig) "No calibration equations to remove." + indices = Int[] + old_eqs = Expr[] + for item in removals + idx::Int = if item isa Int + @assert 1 <= item <= length(calib_orig) "Calibration index $(item) out of bounds." + item + else + target = normalize_equation_input(item)::Expr + found = find_equation_index(calib_orig, target) + @assert found !== nothing "Calibration equation not found: $(target)" + found + end + push!(indices, idx) + push!(old_eqs, calib_orig[idx]) + end + updated = copy(calib_orig) + for i in sort(indices, rev = true) + deleteat!(updated, i) + end + entries = RevisionEntry[] + for (i, idx) in enumerate(indices) + push!(entries, revision_entry(:remove_calibration_equation; + equation_index = idx, old_equation = old_eqs[i], new_equation = nothing)) + end + append!(𝓂.revision_history, entries) + + param_overrides = parameters_to_dict(parameters) + reprocess_calibration_equations!(𝓂, updated; parameters = nothing, + parameter_overrides = param_overrides, verbose = verbose, silent = silent) + return nothing +end + + +""" +$(SIGNATURES) +Return the recorded history of equation modifications for the model. + +Each entry is a `NamedTuple` with the fields: +- `timestamp`: the `DateTime` when the modification was applied. +- `action`: one of `:update_equation`, `:add_equation`, `:remove_equation`, + `:update_calibration_equation`, `:add_calibration_equation`, + `:remove_calibration_equation`. +- `equation_index`: the 1-based index that was affected. +- `old_equation`: the previous equation as an `Expr`, or `nothing` for + additions. +- `new_equation`: the new equation as an `Expr`, or `nothing` for removals. + +The list is append-only and ordered chronologically. + +# Arguments +- $MODEL® + +# Returns +- `Vector{RevisionEntry}` — a copy of the model's revision history. + +# Examples +```julia +using MacroModelling + +@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 + +add_equation!(RBC, :(y[0] = c[0] + k[0] - (1 - δ) * k[-1])) +get_revision_history(RBC) +# output +1-element Vector{@NamedTuple{timestamp::Dates.DateTime, action::Symbol, equation_index::Union{Nothing, Int64}, old_equation::Union{Nothing, Expr}, new_equation::Union{Nothing, Expr}}}: + (timestamp = ..., action = :add_equation, equation_index = 5, old_equation = nothing, new_equation = :(y[0] = (c[0] + k[0]) - (1 - δ) * k[-1])) +``` +""" +function get_revision_history(𝓂::ℳ)::Vector{RevisionEntry} + return copy(𝓂.revision_history) +end + + +""" +$(SIGNATURES) +Write the current model equations and parameter block to a Julia source file. + +The generated file uses the `@model` and `@parameters` macros, so +`include`ing it re-creates a model equivalent to the current state of `𝓂` +(equations, calibration, parameter values and bounds). + +# Arguments +- $MODEL® +- `filepath` [Type: `String`]: destination path for the generated `.jl` file. + +# Keyword Arguments +- `overwrite` [Default: `false`, Type: `Bool`]: replace the file if it + already exists. When `false` and the file exists, an error is raised. + +# Returns +- `String` — the path written to (same as `filepath`). + +# Examples +```julia +using MacroModelling + +@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 + +write_julia_model_file(RBC, joinpath(tempdir(), "RBC.jl"), overwrite = true) +``` +""" +function write_julia_model_file(𝓂::ℳ, filepath::String; overwrite::Bool = false)::String + if isfile(filepath) && !overwrite + error("File already exists: $filepath. Pass overwrite = true to replace it.") + end + + parameter_block = reconstruct_parameter_block(𝓂) + model_name = string(𝓂.model_name) + + open(filepath, "w") do io + println(io, "using MacroModelling\n") + println(io, "@model ", model_name, " begin") + for eq in 𝓂.equations.original + println(io, " ", replace(string(eq), "◖" => "{", "◗" => "}")) + end + println(io, "end\n") + + println(io, "@parameters ", model_name, " begin") + for line in parameter_block.args + line isa LineNumberNode && continue + println(io, " ", replace(string(line), "◖" => "{", "◗" => "}")) + end + println(io, "end") + end + + return filepath +end + + +# Aliases +const replace_equations! = update_equations! +const replace_calibration_equations! = update_calibration_equations! + + +# ======================================================================== +# 2. Internal helpers (not exported) +# ======================================================================== + +""" + reset_solver_state!(𝓂::ℳ) + +Invalidate every cached solver result so the next `solve!` call recomputes +from scratch. Also marks the compiled model functions as outdated so the +rebuild pipeline rewrites them. +""" +function reset_solver_state!(𝓂::ℳ) + 𝓂.caches.valid_for = valid_for_caches() + empty!(𝓂.caches.solver) + + # Reset size-dependent cache matrices so downstream code reallocates for + # the (potentially) new model dimensions. + 𝓂.caches.jacobian = zeros(0, 0) + 𝓂.caches.jacobian_parameters = zeros(0, 0) + 𝓂.caches.jacobian_SS_and_pars = zeros(0, 0) + 𝓂.caches.hessian = zeros(0, 0) + 𝓂.caches.hessian_parameters = zeros(0, 0) + 𝓂.caches.hessian_SS_and_pars = zeros(0, 0) + 𝓂.caches.third_order_derivatives = zeros(0, 0) + 𝓂.caches.third_order_derivatives_parameters = zeros(0, 0) + 𝓂.caches.third_order_derivatives_SS_and_pars = zeros(0, 0) + 𝓂.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_stochastic_steady_state = Float64[] + 𝓂.caches.second_order_solution = SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0) + 𝓂.caches.pruned_second_order_stochastic_steady_state = Float64[] + 𝓂.caches.third_order_stochastic_steady_state = Float64[] + 𝓂.caches.third_order_solution = SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0) + 𝓂.caches.pruned_third_order_stochastic_steady_state = Float64[] + 𝓂.caches.non_stochastic_steady_state = Float64[] + 𝓂.caches.covariance_first_order = zeros(0, 0) + 𝓂.caches.covariance_second_order = zeros(0, 0) + 𝓂.caches.covariance_third_order = zeros(0, 0) + 𝓂.caches.covariance_third_order_autocorr = zeros(0, 0) + 𝓂.caches.has_unit_roots = false + + 𝓂.functions.functions_written = false + + # Invalidate derived caches that depend on the equations / calibration. + # Empty axes / name tables are the sentinel used by the `ensure_*!` + # helpers to decide whether to recompute. + 𝓂.constants.post_complete_parameters = update_post_complete_parameters( + 𝓂.constants.post_complete_parameters; + var_axis = Symbol[], + calib_axis = Symbol[], + exo_axis_plain = Symbol[], + exo_axis_with_subscript = Symbol[], + full_NSSS_display = Symbol[], + SS_and_pars_names = Symbol[], + initialized = false, + ) + + # Reset the workspace buffers, forcing `ensure_*!` helpers to resize them + # on next use. + 𝓂.workspaces = Workspaces() + + return nothing +end + + +""" + reconstruct_parameter_block(𝓂; calibration_original_override = nothing, + parameter_overrides = nothing) + +Return a `:block` `Expr` that reproduces a valid `@parameters` body from the +current model state. Optional overrides let callers preview the effect of a +calibration change before committing it to the model. +""" +function reconstruct_parameter_block(𝓂::ℳ; + calibration_original_override::Union{Nothing, Vector{Expr}} = nothing, + parameter_overrides::Union{Nothing, AbstractDict{Symbol, <:Real}} = nothing) + lines = Any[] + + calibration_original = calibration_original_override === nothing ? + 𝓂.equations.calibration_original : calibration_original_override + + new_calib_params = Set{Symbol}() + for eq in calibration_original + p = extract_calibrated_parameter(eq) + if p !== nothing + push!(new_calib_params, p) + end + end + + old_calib_params = Set{Symbol}(𝓂.equations.calibration_parameters) + params_becoming_calibrated = setdiff(new_calib_params, old_calib_params) + params_no_longer_calibrated = setdiff(old_calib_params, new_calib_params) + + fixed_params = 𝓂.constants.post_complete_parameters.parameters + fixed_values = 𝓂.parameter_values + for (p, v) in zip(fixed_params, fixed_values) + (isnan(v) || p in params_becoming_calibrated) && continue + val = (parameter_overrides !== nothing && haskey(parameter_overrides, p)) ? parameter_overrides[p] : v + push!(lines, Expr(:(=), p, val)) + end + + if !isempty(params_no_longer_calibrated) + n_vars = 𝓂.constants.post_model_macro.nVars + old_calib_list = 𝓂.equations.calibration_parameters + for p in params_no_longer_calibrated + if parameter_overrides !== nothing && haskey(parameter_overrides, p) + push!(lines, Expr(:(=), p, parameter_overrides[p])) + continue + end + idx = findfirst(==(p), old_calib_list) + if idx !== nothing && length(𝓂.caches.non_stochastic_steady_state) >= n_vars + idx + val = 𝓂.caches.non_stochastic_steady_state[n_vars + idx] + push!(lines, Expr(:(=), p, val)) + end + end + end + + for eq in 𝓂.equations.calibration_no_var + push!(lines, eq) + end + + for eq in calibration_original + push!(lines, eq) + end + + for (p, (lo, hi)) in 𝓂.constants.post_parameters_macro.bounds + push!(lines, Expr(:comparison, lo, :(<), p, :(<), hi)) + end + + return Expr(:block, lines...) +end + + +""" + extract_calibrated_parameter(eq::Expr) -> Union{Symbol, Nothing} + +Return the parameter on the right of the `|` in a calibration equation such +as `k[ss] = 1.5 | δ`. Returns `nothing` if no calibration marker is found. +""" +function extract_calibrated_parameter(eq::Expr)::Union{Symbol, Nothing} + result = Ref{Union{Symbol, Nothing}}(nothing) + postwalk(eq) do x + if x isa Expr && x.head == :call && !isempty(x.args) && x.args[1] == :| + if length(x.args) >= 3 && x.args[end] isa Symbol + result[] = x.args[end] + end + end + x + end + return result[] +end + + +""" + finalize_model_update!(𝓂; verbose, silent) + +Internal helper that finalises a model update by rebuilding the steady-state +solver and symbolic derivatives for the current `𝓂` state. Called by both +`reprocess_model_equations!` and `reprocess_calibration_equations!`. +""" +function finalize_model_update!(𝓂::ℳ; verbose::Bool = false, silent::Bool = true) + has_missing_parameters = !isempty(𝓂.constants.post_complete_parameters.missing_parameters) + missing_params = 𝓂.constants.post_complete_parameters.missing_parameters + + if !isnothing(𝓂.functions.NSSS_custom) + write_ss_check_function!(𝓂) + else + if !has_missing_parameters + set_up_steady_state_solver!( + 𝓂; + verbose = verbose, + silent = silent, + ss_symbolic_mode = 𝓂.constants.post_parameters_macro.ss_symbolic_mode, + ) + end + end + + if !has_missing_parameters + opts = merge_calculation_options(verbose = verbose) + 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 = 1, silent = silent) + 𝓂.functions.functions_written = true + else + if !silent + @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." + end + end + + return nothing +end + + +""" + reprocess_model_equations!(𝓂, new_equations; parameters, verbose, silent) + +Rebuild the model from an updated equation list while preserving the +revision history. Equivalent to re-running `@model` and `@parameters` on the +current parameter state, but without re-evaluating the macros. +""" +function reprocess_model_equations!(𝓂::ℳ, + new_equations::Vector{Expr}; + parameters::ParameterType = nothing, + verbose::Bool = false, + silent::Bool = true) + if parameters !== nothing + write_parameters_input!(𝓂, parameters, verbose = verbose) + end + + updated_block = Expr(:block, new_equations...) + parameter_block = reconstruct_parameter_block(𝓂) + + T, equations_struct, ℂ, 𝓦 = process_model_equations( + updated_block, + 𝓂.constants.post_model_macro.max_obc_horizon, + 𝓂.constants.post_parameters_macro.precompile, + ) + + 𝓂.constants = ℂ + 𝓂.workspaces = 𝓦 + 𝓂.equations = equations_struct + + reset_solver_state!(𝓂) + + parsed_parameters = process_parameter_definitions(parameter_block, 𝓂.constants.post_model_macro) + + 𝓂.constants.post_parameters_macro = update_post_parameters_macro( + 𝓂.constants.post_parameters_macro; + parameters_as_function_of_parameters = parsed_parameters.calib_parameters_no_var, + ss_calib_list = parsed_parameters.ss_calib_list, + par_calib_list = parsed_parameters.par_calib_list, + bounds = parsed_parameters.bounds, + ) + + 𝓂.equations.calibration = parsed_parameters.equations.calibration + 𝓂.equations.calibration_no_var = parsed_parameters.equations.calibration_no_var + 𝓂.equations.calibration_parameters = parsed_parameters.equations.calibration_parameters + 𝓂.equations.calibration_original = parsed_parameters.equations.calibration_original + + 𝓂.constants.post_complete_parameters = update_post_complete_parameters( + 𝓂.constants.post_complete_parameters; + parameters = parsed_parameters.parameters, + missing_parameters = parsed_parameters.missing_parameters, + ) + + 𝓂.parameter_values = parsed_parameters.parameter_values + + finalize_model_update!(𝓂; verbose = verbose, silent = silent) + + return nothing +end + + +""" + reprocess_calibration_equations!(𝓂, updated_calibration_original; parameters, + parameter_overrides, verbose, silent) + +Rebuild the model's calibration with a modified `calibration_original` list. +`parameter_overrides` lets callers supply replacement fixed values for +parameters that leave the calibration set. +""" +function reprocess_calibration_equations!(𝓂::ℳ, + updated_calibration_original::Vector{Expr}; + parameters::ParameterType = nothing, + parameter_overrides::Dict{Symbol, Float64} = Dict{Symbol, Float64}(), + verbose::Bool = false, + silent::Bool = true) + if parameters !== nothing + write_parameters_input!(𝓂, parameters, verbose = verbose) + end + + parameter_block = reconstruct_parameter_block( + 𝓂; + calibration_original_override = updated_calibration_original, + parameter_overrides = parameter_overrides, + ) + + parsed_parameters = process_parameter_definitions(parameter_block, 𝓂.constants.post_model_macro) + + 𝓂.constants.post_parameters_macro = update_post_parameters_macro( + 𝓂.constants.post_parameters_macro; + parameters_as_function_of_parameters = parsed_parameters.calib_parameters_no_var, + ss_calib_list = parsed_parameters.ss_calib_list, + par_calib_list = parsed_parameters.par_calib_list, + bounds = parsed_parameters.bounds, + ) + + 𝓂.equations.calibration = parsed_parameters.equations.calibration + 𝓂.equations.calibration_no_var = parsed_parameters.equations.calibration_no_var + 𝓂.equations.calibration_parameters = parsed_parameters.equations.calibration_parameters + 𝓂.equations.calibration_original = parsed_parameters.equations.calibration_original + + 𝓂.constants.post_complete_parameters = update_post_complete_parameters( + 𝓂.constants.post_complete_parameters; + parameters = parsed_parameters.parameters, + missing_parameters = parsed_parameters.missing_parameters, + ) + + 𝓂.parameter_values = parsed_parameters.parameter_values + + reset_solver_state!(𝓂) + + finalize_model_update!(𝓂; verbose = verbose, silent = silent) + + return nothing +end + + +""" + normalize_equation_input(eq) + +Normalize a user-provided equation expression: strip line-number nodes and +collapse single-expression `begin ... end` blocks. Accepts `Expr` or +`String` (which is parsed first). +""" +normalize_equation_input(eq::String) = normalize_equation_input(Meta.parse(eq)) +function normalize_equation_input(eq::Expr) + cleaned = rmlines(eq) + return postwalk(cleaned) do node + if @capture(node, begin arg_ end) + arg + else + node + end + end +end + + +""" + expr_contains_dynamic(expr, base_expr, pattern) -> Bool + +Check whether `expr` contains a dynamic reference whose base symbol matches +`base_expr`. When `pattern === nothing` any timing counts as a match; when +it is an `Expr` or `Symbol`, match exact equality. +""" +function expr_contains_dynamic(expr, base_expr, pattern) + base_matches(candidate, base) = candidate == base || + string(candidate) == string(base) + + found = Ref(false) + postwalk(expr) do x + if pattern === nothing + if x isa Expr && x.head == :ref + base_matches(x.args[1], base_expr) && (found[] = true) + elseif x isa Expr && x.head == :curly + base_matches(x, base_expr) && (found[] = true) + elseif x isa Symbol + base_matches(x, base_expr) && (found[] = true) + end + else + x == pattern && (found[] = true) + end + x + end + return found[] +end + + +""" + find_equation_index(equations::Vector{Expr}, target_eq::Expr) -> Union{Int, Nothing} + +Return the 1-based index of `target_eq` inside `equations`, comparing via +the normalized textual form (so whitespace and `{}`/`◖◗` differences are +ignored). Returns `nothing` if no match is found. +""" +function find_equation_index(equations::Vector{Expr}, target_eq::Expr)::Union{Int, Nothing} + target_str = equation_canonical(target_eq) + for (i, eq) in enumerate(equations) + if equation_canonical(eq) == target_str + return i + end + end + return nothing +end + + +""" + equation_canonical(eq) -> String + +Return a whitespace- and brace-insensitive canonical string form of `eq`, +used to compare equations for equality. +""" +function equation_canonical(eq)::String + s = string(eq) + s = replace(s, "◖" => "{", "◗" => "}") + s = replace(s, r"\s+" => "") + return s +end + + +""" + revision_entry(action; equation_index, old_equation, new_equation) + +Construct a `RevisionEntry` `NamedTuple` with the current timestamp. +""" +revision_entry(action::Symbol; + equation_index::Union{Int, Nothing} = nothing, + old_equation::Union{Expr, Nothing} = nothing, + new_equation::Union{Expr, Nothing} = nothing) = + (timestamp = Dates.now(), + action = action, + equation_index = equation_index, + old_equation = old_equation, + new_equation = new_equation) + + +""" + parameters_to_dict(parameters) -> Dict{Symbol, Float64} + +Convert `ParameterType`-like user input into a `Dict{Symbol, Float64}` used +by `reprocess_calibration_equations!`. Unsupported forms fall back to an +empty dict. +""" +function parameters_to_dict(parameters)::Dict{Symbol, Float64} + d = Dict{Symbol, Float64}() + parameters === nothing && return d + if parameters isa Pair + k = parameters.first; v = parameters.second + k_sym = k isa Symbol ? k : Symbol(k) + d[k_sym] = Float64(v) + elseif parameters isa AbstractDict + for (k, v) in parameters + k_sym = k isa Symbol ? k : Symbol(k) + d[k_sym] = Float64(v) + end + elseif parameters isa Union{Tuple, Vector} + for p in parameters + if p isa Pair + k = p.first; v = p.second + k_sym = k isa Symbol ? k : Symbol(k) + d[k_sym] = Float64(v) + end + end + end + return d +end diff --git a/src/macros.jl b/src/parser/equation_processing.jl similarity index 70% rename from src/macros.jl rename to src/parser/equation_processing.jl index 6bda65dec..4f71900f4 100644 --- a/src/macros.jl +++ b/src/parser/equation_processing.jl @@ -1,102 +1,27 @@ -const all_available_algorithms = [:first_order, :second_order, :pruned_second_order, :third_order, :pruned_third_order] - +# Pure-function equation processing helpers used by both the equation +# modification pipeline and (potentially) the model macros. +# +# `process_model_equations` reproduces the work the `@model` macro performs on +# its equation block, returning a `post_model_macro` struct and an `equations` +# struct so the model state can be updated without re-invoking the macro. +# +# `process_parameter_definitions` reproduces the work the `@parameters` macro +# performs on the parameter block. It takes a `post_model_macro` describing the +# current model (used for variable name lookups, index expansion, etc.) and +# returns the pieces needed to update `post_parameters_macro`, the equations +# struct's calibration fields, and `post_complete_parameters`. """ -$(SIGNATURES) -Parses the model equations and assigns them to an object. - -# Arguments -- `𝓂`: name of the object to be created containing the model information. -- `ex`: equations - -# Optional arguments to be placed between `𝓂` and `ex` -- `max_obc_horizon` [Default: `40`, Type: `Int`]: maximum length of anticipated shocks and corresponding unconditional forecast horizon over which the occasionally binding constraint is to be enforced. Increase this number if no solution is found to enforce the constraint. - -Variables must be defined with their time subscript in square brackets. -Endogenous variables can have the following: -- present: `c[0]` -- non-stochastic steady state: `c[ss]` instead of `ss` any of the following is also a valid flag for the non-stochastic steady state: `ss`, `stst`, `steady`, `steadystate`, `steady_state`, and the parser is case-insensitive (`SS` or `sTst` will work as well). -- past: `c[-1]` or any negative Integer: e.g. `c[-12]` -- future: `c[1]` or any positive Integer: e.g. `c[16]` or `c[+16]` -Signed integers are recognised and parsed as such. - -Exogenous variables (shocks) can have the following: -- present: `eps_z[x]` instead of `x` any of the following is also a valid flag for exogenous variables: `ex`, `exo`, `exogenous`, and the parser is case-insensitive (`Ex` or `exoGenous` will work as well). -- past: `eps_z[x-1]` -- future: `eps_z[x+1]` - -Parameters enter the equations without square brackets. - -If an equation contains a `max` or `min` operator, the default dynamic (first order) solution of the model will enforce the occasionally binding constraint. This enforcement can be disabled by setting `ignore_obc = true` in the relevant function calls. - -# Examples -```julia -using MacroModelling - -@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 -``` - -# Programmatic model writing - -Parameters and variables can be indexed using curly braces: e.g. `c{H}[0]`, `eps_z{F}[x]`, or `α{H}`. - -`for` loops can be used to write models programmatically. They can either be used to generate expressions where the time index or the index in curly braces is iterated over: -- generate equation with different indices in curly braces: `for co in [H,F] C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1] end = for co in [H,F] Y{co}[0] end` -- generate multiple equations with different indices in curly braces: `for co in [H, F] K{co}[0] = (1-delta{co}) * K{co}[-1] + S{co}[0] end` -- generate equation with different time indices: `Y_annual[0] = for lag in -3:0 Y[lag] end` or `R_annual[0] = for operator = :*, lag in -3:0 R[lag] end` - -# Returns -- `Nothing`. The macro creates the model `𝓂` in the calling scope. + process_model_equations(model_block::Expr, max_obc_horizon::Int, precompile::Bool) + +Parse a `@model`-style equation block and return `(T, equations_struct)` where +`T::post_model_macro` is the parsed model structure and `equations_struct::equations` +is a freshly constructed equations container with dynamic, steady-state and +original equations populated. Calibration fields on the returned equations +struct are left empty and must be populated by +`process_parameter_definitions` before the model can be solved. """ -macro model(𝓂,ex...) - # parse options - verbose = false - precompile = false - max_obc_horizon = 40 - - for exp in ex[1:end-1] - postwalk(x -> - x isa Expr ? - x.head == :(=) ? - x.args[1] == :verbose && x.args[2] isa Bool ? - verbose = x.args[2] : - x.args[1] == :precompile && x.args[2] isa Bool ? - precompile = x.args[2] : - x.args[1] == :max_obc_horizon && x.args[2] isa Int ? - max_obc_horizon = x.args[2] : - begin - @warn "Invalid option `$(x.args[1])` ignored. See docs: `?@model` for valid options." - x - end : - x : - x, - exp) - end - - # create data containers - parameters = [] - parameter_values = Vector{Float64}(undef,0) - - ss_calib_list = [] - 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_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 - +function process_model_equations(model_block_in::Expr, max_obc_horizon::Int, precompile::Bool) original_equations = [] calibration_equations = [] calibration_equations_parameters = [] @@ -115,7 +40,7 @@ macro model(𝓂,ex...) ss_equations_with_aux_variables = Int[] dyn_eq_aux_ind = Int[] - model_ex = parse_for_loops(ex[end]) + model_ex = parse_for_loops(model_block_in) model_ex = resolve_if_expr(model_ex::Expr)::Expr @@ -704,6 +629,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 +708,7 @@ macro model(𝓂,ex...) nMixed, nFuture_not_past_and_mixed, nPast_not_future_and_mixed, + I_nPast, # nPresent_but_not_only, nVars, nExo, @@ -852,182 +779,51 @@ macro model(𝓂,ex...) end @assert length(duplicate_equations) == 0 "The following equations appear more than once (and should only appear once): \n" * join(["$(original_equations[eq_idxs[1]])" for eq_idxs in duplicate_equations], "\n") - - # default_optimizer = nlboxsolve - # default_optimizer = Optimisers.Adam - # default_optimizer = NLopt.LN_BOBYQA - - #assemble data container - model_name = string(𝓂) - quote - global $𝓂 = ℳ( - $model_name, - # $default_optimizer, - # 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 - ), - zeros(0,0), # jacobian - zeros(0,0), # jacobian_parameters - zeros(0,0), # jacobian_SS_and_pars - zeros(0,0), # hessian - zeros(0,0), # hessian_parameters - zeros(0,0), # hessian_SS_and_pars - zeros(0,0), # third_order_derivatives - 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), # qme_solution - Float64[], # second_order_stochastic_steady_state - SparseMatrixCSC{Float64, Int64}(ℒ.I,0,0), # second_order_solution - Float64[], # pruned_second_order_stochastic_steady_state - Float64[], # third_order_stochastic_steady_state - 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 - ), - # (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 - # ([], Int[], zeros(1,1)), # model_jacobian - # # x->x, # model_jacobian_parameters - # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)), # model_jacobian_SS_and_pars_vars - # # FWrap{Tuple{Vector{Float64}, Vector{Number}, Vector{Float64}}, SparseMatrixCSC{Float64}}(model_jacobian), - # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)),#x->x, # model_hessian - # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)), # model_hessian_SS_and_pars_vars - # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)),#x->x, # model_third_order_derivatives - # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)),#x->x, # model_third_order_derivatives_SS_and_pars_vars - - # $T, - - $ℂ, - $𝓦, - - 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 - 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 - ), - - SolveCounters() - ); - end -end - + ℂ = Constants(T) + 𝓦 = Workspaces() - - - -""" -$(SIGNATURES) -Adds parameter values and calibration equations to the previously defined model. Allows to provide an initial guess for the non-stochastic steady state (NSSS). - -# Arguments -- `𝓂`: name of the object previously created containing the model information. -- `ex`: parameter, parameters values, and calibration equations - -Parameters can be defined in either of the following ways: -- plain number: `δ = 0.02` -- expression containing numbers: `δ = 1/50` -- expression containing other parameters: `δ = 2 * std_z` in this case it is irrelevant if `std_z` is defined before or after. The definitions including other parameters are treated as a system of equations and solved accordingly. -- 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. -- $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 -- `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 - -# Delayed parameter definition -Not all parameters need to be defined in the `@parameters` macro. Calibration equations using the `|` syntax and parameters defined as functions of other parameters must be declared here, but simple parameter value assignments (e.g., `α = 0.5`) can be deferred and provided later by passing them to any function that accepts the `parameters` argument (e.g., [`get_irf`](@ref), [`get_steady_state`](@ref), [`simulate`](@ref)). - -**Parameter ordering:** When some parameters are not defined in `@parameters`, the final parameter vector follows a specific order: first come the parameters defined in `@parameters` (in their declaration order), followed by any missing parameters (in alphabetical order). This ordering is important when passing parameter values by position rather than by name in subsequent function calls. - -# Examples -```julia -using MacroModelling - -@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 verbose = true begin - std_z = 0.01 - ρ = 0.2 - δ = 0.02 - α = 0.5 - β = 0.95 -end - -@model RBC_calibrated 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_calibrated verbose = true guess = Dict(:k => 3) begin - std_z = 0.01 - ρ = 0.2 - δ = 0.02 - k[ss] / q[ss] = 2.5 | α - β = 0.95 + ss_aux_eqs_vec = Expr[e for e in ss_aux_equations] + dyn_eqs_vec = Expr[e for e in dyn_equations] + ss_eqs_vec = Expr[e for e in ss_equations] + orig_eqs_vec = Expr[e for e in original_equations] + calib_eqs_vec = Expr[e for e in calibration_equations] + + equations_struct = equations( + orig_eqs_vec, + dyn_eqs_vec, + ss_eqs_vec, + ss_aux_eqs_vec, + Expr[], # obc_violation + calib_eqs_vec, # calibration (filled later by @parameters) + Expr[], # calibration_no_var + Symbol[], # calibration_parameters + Expr[], # calibration_original + ) + + return T, equations_struct, ℂ, 𝓦 end -``` -# Programmatic model writing -Variables and parameters indexed with curly braces can be either referenced specifically (e.g. `c{H}[ss]`) or generally (e.g. `alpha`). If they are referenced generally the parse assumes all instances (indices) are meant. For example, in a model where `alpha` has two indices `H` and `F`, the expression `alpha = 0.3` is interpreted as two expressions: `alpha{H} = 0.3` and `alpha{F} = 0.3`. The same goes for calibration equations. -# Returns -- `Nothing`. The macro assigns parameter values and calibration equations to `𝓂` in the calling scope. """ -macro parameters(𝓂,ex...) + process_parameter_definitions(parameter_block::Expr, pmm::post_model_macro) -> NamedTuple + +Parse a `@parameters`-style parameter definition block and return the data +needed to populate the `post_parameters_macro`, `post_complete_parameters` +and calibration-related fields on the model's equations struct. + +Returned NamedTuple fields: +- `parameter_values::Vector{Float64}` — values for `parameters`, ordered to match +- `missing_parameters::Vector{Symbol}` +- `parameters::Vector{Symbol}` — parameters in declaration order (with missing appended) +- `calib_parameters_no_var::Vector{Symbol}` — parameters defined as functions of other parameters +- `ss_calib_list::Vector{Set{Symbol}}` +- `par_calib_list::Vector{Set{Symbol}}` +- `bounds::Dict{Symbol,Tuple{Float64,Float64}}` +- `equations::NamedTuple` with fields `calibration`, `calibration_no_var`, + `calibration_parameters`, `calibration_original`. +""" +function process_parameter_definitions(parameter_block_in::Expr, pmm::post_model_macro) calib_equations = [] calib_equations_no_var = [] calib_values_no_var = [] @@ -1053,58 +849,7 @@ macro parameters(𝓂,ex...) par_defined_more_than_once = Set() bounded_vars = [] - - # parse options - verbose = false - silent = false - symbolic = false - 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 - - for exp in ex[1:end-1] - postwalk(x -> - x isa Expr ? - x.head == :(=) ? - (x.args[1] == :symbolic && x.args[2] isa Bool) ? - symbolic = 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) ? - silent = x.args[2] : - (x.args[1] == :report_missing_parameters && x.args[2] isa Bool) ? - report_missing_parameters = x.args[2] : - (x.args[1] == :precompile && x.args[2] isa Bool) ? - precompile = x.args[2] : - (x.args[1] == :perturbation_order && x.args[2] isa Int) ? - perturbation_order = x.args[2] : - (x.args[1] == :guess && (isa(eval(x.args[2]), Dict{Symbol, <:Real}) || isa(eval(x.args[2]), Dict{String, <:Real}))) ? - guess = x.args[2] : - (x.args[1] == :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) ? - ss_solver_parameters_maxtime = x.args[2] : - begin - @warn "Invalid option `$(x.args[1])` ignored. See docs: `?@parameters` for valid options." - x - end : - x : - x, - exp) - end - - @assert ss_solver_parameters_algorithm ∈ [:ESCH, :SAMIN] "ss_solver_parameters_algorithm must be :ESCH or :SAMIN. Got $ss_solver_parameters_algorithm. Using default :ESCH." - - parameter_definitions = replace_indices(ex[end]) + parameter_definitions = replace_indices(parameter_block_in) # parse parameter inputs # label all variables parameters and exogenous variables and timings across all equations @@ -1436,151 +1181,84 @@ macro parameters(𝓂,ex...) x,bound) end - return quote - mod = @__MODULE__ - - if any(contains.(string.(mod.$𝓂.constants.post_model_macro.var), "ᵒᵇᶜ")) - push!($calib_parameters, :activeᵒᵇᶜshocks) - push!($calib_values, 0) - end - - calib_parameters, calib_values = expand_indices($calib_parameters, $calib_values, [mod.$𝓂.constants.post_model_macro.parameters_in_equations; mod.$𝓂.constants.post_model_macro.var]) - calib_eq_parameters, calib_equations_list, ss_calib_list, par_calib_list = expand_calibration_equations($calib_eq_parameters, $calib_equations_list, $ss_calib_list, $par_calib_list, [mod.$𝓂.constants.post_model_macro.parameters_in_equations; mod.$𝓂.constants.post_model_macro.var]) - calib_parameters_no_var, calib_equations_no_var_list = expand_indices($calib_parameters_no_var, $calib_equations_no_var_list, [mod.$𝓂.constants.post_model_macro.parameters_in_equations; mod.$𝓂.constants.post_model_macro.var]) - - # Calculate missing parameters instead of asserting - # Include parameters from: - # 1. par_calib_list - parameters used in calibration equations (e.g., K_ss in "K[ss] = K_ss | beta") - # 2. parameters_in_equations - parameters used in model equations - # 3. par_no_var_calib_list - parameters used in parameter definitions (e.g., rho{H}{H} in "rho{F}{F} = rho{H}{H}") - # Subtract: - # 1. calib_parameters - parameters with explicit values (e.g., "α = 0.5") - # 2. calib_parameters_no_var - parameters defined as functions of other parameters (e.g., "α = alpha_param") - # 3. calib_eq_parameters - parameters determined by calibration equations (e.g., "beta" in "K[ss] = K_ss | beta") - # Start with directly required parameters - all_required_params = union( - reduce(union, par_calib_list, init = Set{Symbol}()), - reduce(union, $par_no_var_calib_rhs_list, init = Set{Symbol}()), - Set{Symbol}(mod.$𝓂.constants.post_model_macro.parameters_in_equations) - ) - - # Add parameters from parameter definitions, but only if the target parameter is needed - # This handles the case where parameter X = f(Y, Z) but X is not used in the model. - # In that case, Y and Z should not be required either. - # We need to check if target is in all_required_params OR in calib_eq_parameters (parameters used in calibration equations) - par_no_var_calib_filtered = mapreduce(i -> $par_no_var_calib_list[i], union, findall(target_param -> target_param ∈ all_required_params, calib_parameters_no_var), init = Set{Symbol}()) - - all_required_params = union(all_required_params, par_no_var_calib_filtered) - - defined_params = union( - Set{Symbol}(calib_parameters), - Set{Symbol}(calib_parameters_no_var), - Set{Symbol}(calib_eq_parameters) - ) - - ignored_params = collect(setdiff(defined_params, all_required_params)) - - if !isempty(ignored_params) @warn "Parameters not part of the model are ignored: $ignored_params" end + # Runtime portion: formerly the quote block in @parameters. + # Variables at hand: calib_parameters, calib_values, calib_parameters_no_var, + # calib_eq_parameters, calib_equations, calib_equations_list, + # calib_equations_no_var_list, ss_calib_list, par_calib_list, + # par_no_var_calib_rhs_list, par_no_var_calib_list, bounds. - missing_params_unsorted = collect(setdiff(all_required_params, defined_params)) - missing_params = sort(missing_params_unsorted) - - has_missing_parameters = length(missing_params) > 0 - - guess_dict = mod.$𝓂.constants.post_parameters_macro.guess - if isa($guess, Dict{String, <:Real}) - guess_dict = Dict{Symbol, Float64}() - for (key, value) in $guess - if key isa String - key = replace_indices(key) - end - guess_dict[replace_indices(key)] = value - end - elseif isa($guess, Dict{Symbol, <:Real}) - guess_dict = $guess - end + if any(contains.(string.(pmm.var), "ᵒᵇᶜ")) + push!(calib_parameters, :activeᵒᵇᶜshocks) + push!(calib_values, 0) + end - bounds_dict = copy(mod.$𝓂.constants.post_parameters_macro.bounds) - for (k,v) in $bounds - bounds_dict[k] = haskey(bounds_dict, k) ? (max(bounds_dict[k][1], v[1]), min(bounds_dict[k][2], v[2])) : (v[1], v[2]) - end - - invalid_bounds = Symbol[] + _pars_and_vars = [pmm.parameters_in_equations; pmm.var] + calib_parameters, calib_values = expand_indices(calib_parameters, calib_values, _pars_and_vars) + calib_eq_parameters, calib_equations_list, ss_calib_list, par_calib_list = expand_calibration_equations(calib_eq_parameters, calib_equations_list, ss_calib_list, par_calib_list, _pars_and_vars) + calib_parameters_no_var, calib_equations_no_var_list = expand_indices(calib_parameters_no_var, calib_equations_no_var_list, _pars_and_vars) - for (k,v) in bounds_dict - if v[1] >= v[2] - push!(invalid_bounds, k) - end - end + all_required_params = union( + reduce(union, par_calib_list, init = Set{Symbol}()), + reduce(union, par_no_var_calib_rhs_list, init = Set{Symbol}()), + Set{Symbol}(pmm.parameters_in_equations) + ) - @assert isempty(invalid_bounds) "Invalid bounds: " * repr(invalid_bounds) - - mod.$𝓂.constants.post_parameters_macro = post_parameters_macro( - calib_parameters_no_var, - $precompile, - $simplify, - guess_dict, - ss_calib_list, - par_calib_list, - # $ss_no_var_calib_list, - # $par_no_var_calib_list, - bounds_dict, - ) - - # Update equations struct with calibration fields - mod.$𝓂.equations.calibration = calib_equations_list - mod.$𝓂.equations.calibration_no_var = calib_equations_no_var_list - mod.$𝓂.equations.calibration_parameters = calib_eq_parameters - - # Keep calib_parameters in declaration order, append missing_params at end - # This preserves declaration order for estimation and method of moments - all_params = vcat(calib_parameters, missing_params) - all_values = vcat(calib_values, fill(NaN, length(missing_params))) - - defined_params_idx = indexin(setdiff(intersect(all_params, defined_params), ignored_params), collect(all_params)) - - mod.$𝓂.constants.post_complete_parameters = update_post_complete_parameters( - mod.$𝓂.constants.post_complete_parameters; - parameters = all_params[defined_params_idx], - 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) - set_custom_steady_state_function!(mod.$𝓂, $steady_state_function) - # end + par_no_var_calib_filtered = mapreduce(i -> par_no_var_calib_list[i], union, findall(target_param -> target_param ∈ all_required_params, calib_parameters_no_var), init = Set{Symbol}()) + all_required_params = union(all_required_params, par_no_var_calib_filtered) - mod.$𝓂.functions.functions_written = false + defined_params = union( + Set{Symbol}(calib_parameters), + Set{Symbol}(calib_parameters_no_var), + Set{Symbol}(calib_eq_parameters), + ) - # time_symbolics = @elapsed - # time_rm_red_SS_vars = @elapsed - if !isnothing($steady_state_function) - 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) - end - end + ignored_params = collect(setdiff(defined_params, all_required_params)) + if !isempty(ignored_params) @warn "Parameters not part of the model are ignored: $ignored_params" end - if !has_missing_parameters - opts = merge_calculation_options(verbose = $verbose) - - SS_and_pars, solution_error, found_solution = solve_steady_state!(mod.$𝓂, opts, $(QuoteNode(ss_solver_parameters_algorithm)), $ss_solver_parameters_maxtime, silent = $silent) - - write_symbolic_derivatives!(mod.$𝓂; perturbation_order = $perturbation_order, silent = $silent) + missing_params_unsorted = collect(setdiff(all_required_params, defined_params)) + missing_params = sort(missing_params_unsorted) - mod.$𝓂.functions.functions_written = true + invalid_bounds = Symbol[] + for (k,v) in bounds + if v[1] >= v[2] + push!(invalid_bounds, k) 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." + end + @assert isempty(invalid_bounds) "Invalid bounds: " * repr(invalid_bounds) + + # Rebuild calibration_original (original "lhs = rhs | param" form) from the raw user-facing + # calibration equation pairs captured during parsing. Use the parameter-at-end form. + _calib_original = Expr[] + for (_eq, _par) in zip(calib_equations, calib_eq_parameters) + if _eq isa Expr && _eq.head == :(=) && length(_eq.args) == 2 + _lhs, _rhs = _eq.args[1], _eq.args[2] + push!(_calib_original, Expr(:(=), _lhs, Expr(:call, :|, _rhs, _par))) end - - if !$silent && $report_missing_parameters Base.show(mod.$𝓂) end - nothing end + + # Keep calib_parameters in declaration order, append missing_params at end + all_params = vcat(calib_parameters, missing_params) + all_values = vcat(calib_values, fill(NaN, length(missing_params))) + defined_params_idx = indexin(setdiff(intersect(all_params, defined_params), ignored_params), collect(all_params)) + + final_parameters = all_params[defined_params_idx] + final_parameter_values = all_values[defined_params_idx] + + eqs_nt = ( + calibration = Expr[e for e in calib_equations_list], + calibration_no_var = Expr[e for e in calib_equations_no_var_list], + calibration_parameters = Symbol[s for s in calib_eq_parameters], + calibration_original = _calib_original, + ) + + return ( + parameter_values = final_parameter_values, + missing_parameters = missing_params, + parameters = final_parameters, + calib_parameters_no_var = Symbol[s for s in calib_parameters_no_var], + ss_calib_list = Vector{Set{Symbol}}([Set{Symbol}(s) for s in ss_calib_list]), + par_calib_list = Vector{Set{Symbol}}([Set{Symbol}(s) for s in par_calib_list]), + bounds = bounds, + equations = eqs_nt, + ) end diff --git a/src/parser/macros.jl b/src/parser/macros.jl new file mode 100644 index 000000000..c305f634f --- /dev/null +++ b/src/parser/macros.jl @@ -0,0 +1,820 @@ +# ======================================================================== +# 1. User-facing macros +# ======================================================================== + +""" +$(SIGNATURES) +Parses the model equations and assigns them to an object. + +# Arguments +- `𝓂`: name of the object to be created containing the model information. +- `ex`: equations + +# Optional arguments to be placed between `𝓂` and `ex` +- `max_obc_horizon` [Default: `40`, Type: `Int`]: maximum length of anticipated shocks and corresponding unconditional forecast horizon over which the occasionally binding constraint is to be enforced. Increase this number if no solution is found to enforce the constraint. + +Variables must be defined with their time subscript in square brackets. +Endogenous variables can have the following: +- present: `c[0]` +- non-stochastic steady state: `c[ss]` instead of `ss` any of the following is also a valid flag for the non-stochastic steady state: `ss`, `stst`, `steady`, `steadystate`, `steady_state`, and the parser is case-insensitive (`SS` or `sTst` will work as well). +- past: `c[-1]` or any negative Integer: e.g. `c[-12]` +- future: `c[1]` or any positive Integer: e.g. `c[16]` or `c[+16]` +Signed integers are recognised and parsed as such. + +Exogenous variables (shocks) can have the following: +- present: `eps_z[x]` instead of `x` any of the following is also a valid flag for exogenous variables: `ex`, `exo`, `exogenous`, and the parser is case-insensitive (`Ex` or `exoGenous` will work as well). +- past: `eps_z[x-1]` +- future: `eps_z[x+1]` + +Parameters enter the equations without square brackets. + +If an equation contains a `max` or `min` operator, the default dynamic (first order) solution of the model will enforce the occasionally binding constraint. This enforcement can be disabled by setting `ignore_obc = true` in the relevant function calls. + +# Examples +```julia +using MacroModelling + +@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 +``` + +# Programmatic model writing + +Parameters and variables can be indexed using curly braces: e.g. `c{H}[0]`, `eps_z{F}[x]`, or `α{H}`. + +`for` loops can be used to write models programmatically. They can either be used to generate expressions where the time index or the index in curly braces is iterated over: +- generate equation with different indices in curly braces: `for co in [H,F] C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1] end = for co in [H,F] Y{co}[0] end` +- generate multiple equations with different indices in curly braces: `for co in [H, F] K{co}[0] = (1-delta{co}) * K{co}[-1] + S{co}[0] end` +- generate equation with different time indices: `Y_annual[0] = for lag in -3:0 Y[lag] end` or `R_annual[0] = for operator = :*, lag in -3:0 R[lag] end` + +# Returns +- `Nothing`. The macro creates the model `𝓂` in the calling scope. +""" +macro model(𝓂, ex...) + # parse options + verbose = false + precompile = false + max_obc_horizon = 40 + + for exp in ex[1:end-1] + postwalk(x -> + x isa Expr ? + x.head == :(=) ? + x.args[1] == :verbose && x.args[2] isa Bool ? + verbose = x.args[2] : + x.args[1] == :precompile && x.args[2] isa Bool ? + precompile = x.args[2] : + x.args[1] == :max_obc_horizon && x.args[2] isa Int ? + max_obc_horizon = x.args[2] : + begin + @warn "Invalid option `$(x.args[1])` ignored. See docs: `?@model` for valid options." + x + end : + x : + x, + exp) + end + + model_name = string(𝓂) + model_block = ex[end] + + # Heavy lifting is delegated to `process_model_equations` in + # src/parser/equation_processing.jl, which is also used by the + # equation-modification reprocess pipeline. Keeping a single source of + # truth avoids drift between the two callers. + return quote + local _T, _eqs, _ℂ, _𝓦 = MacroModelling.process_model_equations( + $(QuoteNode(model_block)), + $max_obc_horizon, + $precompile, + ) + + global $𝓂 = ℳ( + $model_name, + Float64[], # parameter_values, populated by @parameters + _eqs, + caches( + valid_for_caches(), + zeros(0,0), # jacobian + zeros(0,0), # jacobian_parameters + zeros(0,0), # jacobian_SS_and_pars + zeros(0,0), # hessian + zeros(0,0), # hessian_parameters + zeros(0,0), # hessian_SS_and_pars + zeros(0,0), # third_order_derivatives + 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 + Float64[], # pruned_second_order_stochastic_steady_state + Float64[], # third_order_stochastic_steady_state + SparseMatrixCSC{Float64, Int64}(ℒ.I,0,0), # third_order_solution + Float64[], # pruned_third_order_stochastic_steady_state + Float64[], # non_stochastic_steady_state + CircularBuffer{Vector{Vector{Float64}}}(500), # solver + zeros(0,0), # NSSS_∂equations_∂parameters + zeros(0,0), # NSSS_∂equations_∂SS_and_pars + zeros(0,0), # covariance_first_order + zeros(0,0), # covariance_second_order + zeros(0,0), # covariance_third_order + zeros(0,0), # covariance_third_order_autocorr + false, # has_unit_roots + ), + _ℂ, + _𝓦, + model_functions( + x->x, # NSSS_check_func + nothing, # NSSS_custom_function + x->x, # NSSS_∂equations_∂parameters_func + x->x, # NSSS_∂equations_∂SS_and_pars_func + NSSSSolverFunctions(), + nothing, # nsss_param_prep! + jacobian_functions(x->x, x->x, x->x), + hessian_functions(x->x, x->x, x->x), + third_order_derivatives_functions(x->x, x->x, x->x), + x->x, # obc_violation + Tuple{Int,Int,Float64}[], # obc_constraint_info + false, # functions_written + ), + SolveCounters(), + RevisionEntry[], + ); + end +end + + +""" +$(SIGNATURES) +Adds parameter values and calibration equations to the previously defined model. Allows to provide an initial guess for the non-stochastic steady state (NSSS). + +# Arguments +- `𝓂`: name of the object previously created containing the model information. +- `ex`: parameter, parameters values, and calibration equations + +Parameters can be defined in either of the following ways: +- plain number: `δ = 0.02` +- expression containing numbers: `δ = 1/50` +- expression containing other parameters: `δ = 2 * std_z` in this case it is irrelevant if `std_z` is defined before or after. The definitions including other parameters are treated as a system of equations and solved accordingly. +- 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}` 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 +- `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. +- `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 + +# Delayed parameter definition +Not all parameters need to be defined in the `@parameters` macro. Calibration equations using the `|` syntax and parameters defined as functions of other parameters must be declared here, but simple parameter value assignments (e.g., `α = 0.5`) can be deferred and provided later by passing them to any function that accepts the `parameters` argument (e.g., [`get_irf`](@ref), [`get_steady_state`](@ref), [`simulate`](@ref)). + +**Parameter ordering:** When some parameters are not defined in `@parameters`, the final parameter vector follows a specific order: first come the parameters defined in `@parameters` (in their declaration order), followed by any missing parameters (in alphabetical order). This ordering is important when passing parameter values by position rather than by name in subsequent function calls. + +# Examples +```julia +using MacroModelling + +@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 verbose = true begin + std_z = 0.01 + ρ = 0.2 + δ = 0.02 + α = 0.5 + β = 0.95 +end + +@model RBC_calibrated 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_calibrated verbose = true guess = Dict(:k => 3) begin + std_z = 0.01 + ρ = 0.2 + δ = 0.02 + k[ss] / q[ss] = 2.5 | α + β = 0.95 +end +``` + +# Programmatic model writing +Variables and parameters indexed with curly braces can be either referenced specifically (e.g. `c{H}[ss]`) or generally (e.g. `alpha`). If they are referenced generally the parse assumes all instances (indices) are meant. For example, in a model where `alpha` has two indices `H` and `F`, the expression `alpha = 0.3` is interpreted as two expressions: `alpha{H} = 0.3` and `alpha{F} = 0.3`. The same goes for calibration equations. + +# Returns +- `Nothing`. The macro assigns parameter values and calibration equations to `𝓂` in the calling scope. +""" +macro parameters(𝓂, ex...) + # parse options + verbose = false + silent = false + ss_symbolic_mode = :single_equation + precompile = false + report_missing_parameters = true + perturbation_order = 1 + guess = Dict{Symbol,Float64}() + steady_state_function = nothing + ss_solver_parameters_algorithm = :ESCH + ss_solver_parameters_maxtime = 120.0 + + for exp in ex[1:end-1] + postwalk(x -> + x isa Expr ? + x.head == :(=) ? + (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) ? + silent = x.args[2] : + (x.args[1] == :report_missing_parameters && x.args[2] isa Bool) ? + report_missing_parameters = x.args[2] : + (x.args[1] == :precompile && x.args[2] isa Bool) ? + precompile = x.args[2] : + (x.args[1] == :perturbation_order && x.args[2] isa Int) ? + perturbation_order = x.args[2] : + (x.args[1] == :guess && (isa(eval(x.args[2]), Dict{Symbol, <:Real}) || isa(eval(x.args[2]), Dict{String, <:Real}))) ? + guess = x.args[2] : + (x.args[1] == :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] == :steady_state_function && x.args[2] isa Symbol) ? + steady_state_function = esc(x.args[2]) : + (x.args[1] == :ss_solver_parameters_maxtime && x.args[2] isa Real) ? + ss_solver_parameters_maxtime = x.args[2] : + begin + @warn "Invalid option `$(x.args[1])` ignored. See docs: `?@parameters` for valid options." + x + end : + x : + 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." + + parameter_block = ex[end] + + # Parsing of the calibration block is delegated to + # `process_parameter_definitions` in src/parser/equation_processing.jl, + # which is also used by the equation-modification reprocess pipeline. + return quote + mod = @__MODULE__ + + local _parsed = MacroModelling.process_parameter_definitions( + $(QuoteNode(parameter_block)), + mod.$𝓂.constants.post_model_macro, + ) + + # Merge guess option with any guess already on the model. + local _guess_dict = mod.$𝓂.constants.post_parameters_macro.guess + if isa($guess, Dict{String, <:Real}) + _guess_dict = Dict{Symbol, Float64}() + for (key, value) in $guess + if key isa String + key = replace_indices(key) + end + _guess_dict[replace_indices(key)] = value + end + elseif isa($guess, Dict{Symbol, <:Real}) + _guess_dict = $guess + end + + # Merge bounds returned by the parser with bounds already on the model. + local _bounds_dict = copy(mod.$𝓂.constants.post_parameters_macro.bounds) + for (k, v) in _parsed.bounds + _bounds_dict[k] = haskey(_bounds_dict, k) ? + (max(_bounds_dict[k][1], v[1]), min(_bounds_dict[k][2], v[2])) : + (v[1], v[2]) + end + + local _invalid_bounds = Symbol[] + for (k, v) in _bounds_dict + if v[1] >= v[2] + push!(_invalid_bounds, k) + end + end + @assert isempty(_invalid_bounds) "Invalid bounds: " * repr(_invalid_bounds) + + mod.$𝓂.constants.post_parameters_macro = post_parameters_macro( + _parsed.calib_parameters_no_var, + $precompile, + $(QuoteNode(ss_symbolic_mode)), + $(QuoteNode(ss_solver_parameters_algorithm)), + $ss_solver_parameters_maxtime, + _guess_dict, + _parsed.ss_calib_list, + _parsed.par_calib_list, + _bounds_dict, + ) + + mod.$𝓂.equations.calibration = _parsed.equations.calibration + mod.$𝓂.equations.calibration_no_var = _parsed.equations.calibration_no_var + mod.$𝓂.equations.calibration_parameters = _parsed.equations.calibration_parameters + mod.$𝓂.equations.calibration_original = _parsed.equations.calibration_original + + mod.$𝓂.constants.post_complete_parameters = update_post_complete_parameters( + mod.$𝓂.constants.post_complete_parameters; + parameters = _parsed.parameters, + missing_parameters = _parsed.missing_parameters, + ) + mod.$𝓂.parameter_values = _parsed.parameter_values + + local _missing_params = _parsed.missing_parameters + local _has_missing_parameters = !isempty(_missing_params) + + set_custom_steady_state_function!(mod.$𝓂, $steady_state_function) + + mod.$𝓂.functions.functions_written = false + + if !isnothing($steady_state_function) + write_ss_check_function!(mod.$𝓂) + else + if !_has_missing_parameters + set_up_steady_state_solver!(mod.$𝓂, verbose = $verbose, silent = $silent, ss_symbolic_mode = $(QuoteNode(ss_symbolic_mode))) + end + end + + if !_has_missing_parameters + opts = merge_calculation_options(verbose = $verbose) + + SS_and_pars, solution_error, found_solution = solve_steady_state!(mod.$𝓂, opts, $(QuoteNode(ss_solver_parameters_algorithm)), $ss_solver_parameters_maxtime, silent = $silent) + + write_symbolic_derivatives!(mod.$𝓂; perturbation_order = $perturbation_order, silent = $silent) + + mod.$𝓂.functions.functions_written = true + 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_steady_state`, `simulate`, etc." + end + + if !$silent && $report_missing_parameters + Base.show(mod.$𝓂) + end + + nothing + end +end + + +# ======================================================================== +# 2. Internal helpers +# ======================================================================== + +function evaluate_conditions(cond) + if cond isa Bool + return cond + elseif cond isa Expr && cond.head == :call + a, b = cond.args[2], cond.args[3] + + if typeof(a) ∉ [Symbol, Number] + a = eval(a) + end + + if typeof(b) ∉ [Symbol, Number] + b = eval(b) + end + + if cond.args[1] == :(==) + return a == b + elseif cond.args[1] == :(!=) + return a != b + elseif cond.args[1] == :(<) + return a < b + elseif cond.args[1] == :(<=) + return a <= b + elseif cond.args[1] == :(>) + return a > b + elseif cond.args[1] == :(>=) + return a >= b + end + # end + end + return nothing +end + +function resolve_if_expr(ex::Expr) + prewalk(ex) do node + if node isa Expr && (node.head === :if || node.head === :elseif) + cond = node.args[1] + then_blk = node.args[2] + if length(node.args) == 3 + else_blk = node.args[3] + end + val = evaluate_conditions(unblock(cond)) + + if val === true + # recurse into the selected branch + return resolve_if_expr(unblock(then_blk)) + elseif val === false && length(node.args) == 3 + return resolve_if_expr(unblock(else_blk)) + elseif val === false && length(node.args) == 2 + return nothing + elseif val === false && node.head === :elseif + return resolve_if_expr(unblock(else_blk)) + end + end + return node + end +end + +function match_pattern(strings::Union{Set,Vector}, pattern::Regex) + return filter(r -> match(pattern, string(r)) !== nothing, strings) +end + +function contains_equation(expr) + found = false + postwalk(expr) do x + if x isa Expr && x.head == :(=) + found = true + end + return x + end + return found +end + +# function remove_nothing(ex::Expr) +# postwalk(ex) do node +# # Only consider call-nodes with exactly two arguments +# if node isa Expr && node.head === :call && length(node.args) == 3 +# fn, lhs, rhs = node.args +# lhs2 = unblock(lhs) +# rhs2 = unblock(rhs) + +# if rhs2 === :(nothing) +# # strip the call and recurse to clean deeper +# return remove_nothing(lhs2) +# elseif lhs2 === :(nothing) +# return remove_nothing(rhs2) +# # else +# # return remove_nothing(node.args) +# end +# end +# return node +# end +# end + +function remove_nothing(ex::Expr) + postwalk(ex) do node + # Only consider call-expressions + if node isa Expr && node.head === :call && any(node.args .=== nothing) + fn = node.args[1] + # Unblock and collect all the operands + # raw_args = map(arg -> unblock(arg), node.args[2:end]) + # Drop any nothing + kept = filter(arg -> !(unblock(arg) === nothing), node.args[2:end]) + if isempty(kept) + return nothing + elseif length(kept) == 1 + return kept[1] + else + # elseif length(kept) < length(raw_args) + return Expr(:call, fn, kept...) + # else + # return node + end + end + return node + end +end + +function replace_indices_inside_for_loop(exxpr,index_variable,indices,concatenate, operator) + @assert operator ∈ [:+,:*] "Only :+ and :* allowed as operators in for loops." + calls = [] + indices = indices.args[1] == :(:) ? eval(indices) : [indices.args...] + for idx in indices + push!(calls, postwalk(x -> begin + x isa Expr ? + x.head == :ref ? + @capture(x, name_{index_}[time_]) ? + index == index_variable ? + :($(Expr(:ref, Symbol(string(name) * "{" * string(idx) * "}"),time))) : + time isa Expr || time isa Symbol ? + index_variable ∈ get_symbols(time) ? + :($(Expr(:ref, Expr(:curly,name,index), Meta.parse(replace(string(time), string(index_variable) => idx))))) : + x : + x : + @capture(x, name_[time_]) ? + time isa Expr || time isa Symbol ? + index_variable ∈ get_symbols(time) ? + :($(Expr(:ref, name, Meta.parse(replace(string(time), string(index_variable) => idx))))) : + # occursin("{" * string(index_variable) * "}", string(name)) ? + # Expr(:ref, Symbol(replace(string(name), "{" * string(index_variable) * "}" => "◖" * string(idx) * "◗")), time) : + x : + # occursin("{" * string(index_variable) * "}", string(name)) ? + # Expr(:ref, Symbol(replace(string(name), "{" * string(index_variable) * "}" => "◖" * string(idx) * "◗")), time) : + x : + x : + x.head == :if ? + length(x.args) > 2 ? + Expr(:if, postwalk(x -> x == index_variable ? idx : x, x.args[1]), + replace_indices_inside_for_loop(x.args[2],index_variable,:([$idx]),false,:+) |> unblock, + replace_indices_inside_for_loop(x.args[3],index_variable,:([$idx]),false,:+) |> unblock) : + Expr(:if, postwalk(x -> x == index_variable ? idx : x, x.args[1]), + replace_indices_inside_for_loop(x.args[2],index_variable,:([$idx]),false,:+) |> unblock) : + @capture(x, name_{index_}) ? + index == index_variable ? + :($(Symbol(string(name) * "{" * string(idx) * "}"))) : + x : + x : + @capture(x, name_) ? + name == index_variable && idx isa Int ? + :($idx) : + x isa Symbol ? + occursin("{" * string(index_variable) * "}", string(x)) ? + Symbol(replace(string(x), "{" * string(index_variable) * "}" => "{" * string(idx) * "}")) : + x : + x : + x + end, + exxpr)) + end + + if concatenate + return :($(Expr(:call, operator, calls...))) + else + return :($(Expr(:block, calls...))) + # return :($calls...) + # return calls + end +end + +function write_out_for_loops(arg::Expr)::Expr + postwalk(x -> begin + x = flatten(unblock(x)) + x isa Expr ? + x.head == :for ? + x.args[2] isa Array ? + length(x.args[2]) >= 1 ? + x.args[1].head == :block ? + # begin println("here"); + [replace_indices_inside_for_loop(X, Symbol(x.args[1].args[2].args[1]), (x.args[1].args[2].args[2]), false, x.args[1].args[1].args[2].value) for X in x.args[2]] : # end : + # begin println("here2"); + [replace_indices_inside_for_loop(X, Symbol(x.args[1].args[1]), (x.args[1].args[2]), false, :+) for X in x.args[2]] : # end : + x : + x.args[2].head ∉ [:(=), :block] ? + x.args[1].head == :block ? + # begin println("here3"); + replace_indices_inside_for_loop(unblock(x.args[2]), + Symbol(x.args[1].args[2].args[1]), + (x.args[1].args[2].args[2]), + true, + x.args[1].args[1].args[2].value) : # end : # for loop part of equation + x.args[2].head == :if ? + contains_equation(x.args[2]) ? + # begin println("here5"); println(x) + replace_indices_inside_for_loop(unblock(x.args[2]), + Symbol(x.args[1].args[1]), + (x.args[1].args[2]), + false, + :+) : # end : # for loop part of equation + # begin println("here6"); println(x) + replace_indices_inside_for_loop(unblock(x.args[2]), + Symbol(x.args[1].args[1]), + (x.args[1].args[2]), + true, + :+) : # end : # for loop part of equation + # begin println("here4"); println(x) + replace_indices_inside_for_loop(unblock(x.args[2]), + Symbol(x.args[1].args[1]), + (x.args[1].args[2]), + true, + :+) : # end : # for loop part of equation + x.args[1].head == :block ? + # begin println("here5"); + replace_indices_inside_for_loop(unblock(x.args[2]), + Symbol(x.args[1].args[2].args[1]), + (x.args[1].args[2].args[2]), + false, + x.args[1].args[1].args[2].value) : # end : + # end + # : # for loop part of equation + # begin println(x); + # begin println("here7"); println(x) + replace_indices_inside_for_loop(unblock(x.args[2]), + Symbol(x.args[1].args[1]), + (x.args[1].args[2]), + false, + :+) : # end : + # println(out); + # return out end + # : + x : + x + end, + arg) #|> unblock |> flatten +end + +# function parse_for_loops(equations_block) +# eqs = Expr[] # Initialize an empty array to collect expressions + +# # Define a helper recursive function +# function recurse(arg) +# if arg isa Expr +# if arg.head == :block +# for b in arg.args +# if b isa Expr +# # If the result is an Expr, process and add to eqs +# push!(eqs, unblock(replace_indices(b))) +# elseif b isa Array +# recurse(b) +# end +# end +# end +# elseif arg isa Array +# # If the result is an Array, iterate and recurse +# for B in arg +# println((B)) +# recurse(B) +# end +# end +# end + +# for arg in equations_block.args +# if isa(arg,Expr) +# parsed_eqs = write_out_for_loops(arg) +# recurse(parsed_eqs) +# end +# end + +# # Return the collected expressions as a block +# return Expr(:block, eqs...) +# end + + +function parse_for_loops(equations_block)::Expr + eqs = Expr[] + for arg in equations_block.args + if isa(arg,Expr) + parsed_eqs = write_out_for_loops(arg) + # println(parsed_eqs) + if parsed_eqs isa Expr + push!(eqs,unblock(replace_indices(parsed_eqs))) + elseif parsed_eqs isa Array + for B in parsed_eqs + if B isa Array + for b in B + push!(eqs,unblock(replace_indices(b))) + end + elseif B isa Expr + if B.head == :block + for b in B.args + if b isa Expr + push!(eqs,replace_indices(b)) + end + end + else + push!(eqs,unblock(replace_indices(B))) + end + else + push!(eqs,unblock(replace_indices(B))) + end + end + end + + end + end + return Expr(:block,eqs...) |> flatten +end + +function decompose_name(name::Symbol) + name = string(name) + matches = eachmatch(r"◖([\p{L}\p{N}]+)◗|([\p{L}\p{N}]+[^◖◗]*)", name) + + result = [] + nested = [] + + for m in matches + if m.captures[1] !== nothing + push!(nested, m.captures[1]) + else + if !isempty(nested) + push!(result, Symbol.(nested)) + nested = [] + end + push!(result, Symbol(m.captures[2])) + end + end + + if !isempty(nested) + push!(result, (nested)) + end + + return result +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 = [] + + for i in indices + if length(i) > 1 + push!(indexset, Symbol.(i[2])...) + end + end + + return indexset +end + +function expand_calibration_equations(calibration_equation_parameters::Vector{Symbol}, calibration_equations::Vector{Expr}, ss_calib_list::Vector, par_calib_list::Vector, all_names::Vector{Symbol}) + expanded_parameters = Symbol[] + expanded_equations = Expr[] + expanded_ss_var_list = [] + expanded_par_var_list = [] + + for (u,par) in enumerate(calibration_equation_parameters) + indices_in_calibration_equation = Set() + indexed_names = [] + for i in get_symbols(calibration_equations[u]) + indices = get_possible_indices_for_name(i, all_names) + if indices != Any[] + push!(indices_in_calibration_equation, indices) + push!(indexed_names,i) + end + end + + par_indices = get_possible_indices_for_name(par, all_names) + + if length(par_indices) > 0 + push!(indices_in_calibration_equation, par_indices) + end + + @assert length(indices_in_calibration_equation) <= 1 "Calibration equations cannot have more than one index in the equations or for the parameter." + + if length(indices_in_calibration_equation) == 0 + push!(expanded_parameters,par) + push!(expanded_equations,calibration_equations[u]) + push!(expanded_ss_var_list,ss_calib_list[u]) + push!(expanded_par_var_list,par_calib_list[u]) + else + for i in collect(indices_in_calibration_equation)[1] + expanded_ss_var = Set() + expanded_par_var = Set() + push!(expanded_parameters, Symbol(string(par) * "◖" * string(i) * "◗")) + push!(expanded_equations, postwalk(x -> x ∈ indexed_names ? Symbol(string(x) * "◖" * string(i) * "◗") : x, calibration_equations[u])) + for ss in ss_calib_list[u] + if ss ∈ indexed_names + push!(expanded_ss_var,Symbol(string(ss) * "◖" * string(i) * "◗")) + else + push!(expanded_ss_var,ss) + end + end + # Handle parameters from par_calib_list - expand indexed ones, keep non-indexed + for p in par_calib_list[u] + if p ∈ indexed_names + push!(expanded_par_var, Symbol(string(p) * "◖" * string(i) * "◗")) + else + push!(expanded_par_var, p) + end + end + push!(expanded_ss_var_list, expanded_ss_var) + push!(expanded_par_var_list, expanded_par_var) + end + end + end + + return expanded_parameters, expanded_equations, expanded_ss_var_list, expanded_par_var_list +end + +function expand_indices(compressed_inputs::Vector{Symbol}, compressed_values::Vector{T}, expanded_list::Vector{Symbol}) where T + expanded_inputs = Symbol[] + expanded_values = T[] + + for (i,par) in enumerate(compressed_inputs) + par_idx = findall(x -> string(par) == x, first.(split.(string.(expanded_list ), "◖"))) + + if length(par_idx) > 1 + for idx in par_idx + push!(expanded_inputs, expanded_list[idx]) + push!(expanded_values, compressed_values[i]) + end + else#if par ∈ expanded_list ## breaks parameters defined in parameter block + push!(expanded_inputs, par) + push!(expanded_values, compressed_values[i]) + end + end + return expanded_inputs, expanded_values +end + + +const all_available_algorithms = [:first_order, :second_order, :pruned_second_order, :third_order, :pruned_third_order] + diff --git a/src/parser/model_setup.jl b/src/parser/model_setup.jl new file mode 100644 index 000000000..88d556033 --- /dev/null +++ b/src/parser/model_setup.jl @@ -0,0 +1,1433 @@ +function replace_with_one(equation::SPyPyC.Sym{PythonCall.Core.Py}, variable::SPyPyC.Sym{PythonCall.Core.Py})::SPyPyC.Sym{PythonCall.Core.Py} + # equation.subs(variable, 1).replace(SPyPyC.Sym(ℯ), exp(1)) + tmp = SPyPyC.subs(equation, variable, 1) + + return replace_e(tmp) +end + +function replace_e(equation::SPyPyC.Sym{PythonCall.Core.Py})::SPyPyC.Sym{PythonCall.Core.Py} + outraw = SPyPyC.subs(equation, SPyPyC.Sym(ℯ), exp(1)) + + if outraw isa SPyPyC.Sym{PythonCall.Core.Py} + out = outraw + else + out = collect(outraw)[1] + end + + return out +end + +function solve_symbolically(equation::SPyPyC.Sym{PythonCall.Core.Py}, variable::SPyPyC.Sym{PythonCall.Core.Py})::Union{Nothing,Vector{SPyPyC.Sym{PythonCall.Core.Py}}} + soll = try SPyPyC.solve(equation, variable) + catch + end + + return soll +end + +function solve_symbolically(equations::Vector{SPyPyC.Sym{PythonCall.Core.Py}}, variables::Vector{SPyPyC.Sym{PythonCall.Core.Py}})::Union{Nothing,Dict{SPyPyC.Sym{PythonCall.Core.Py}, SPyPyC.Sym{PythonCall.Core.Py}}} + soll = try SPyPyC.solve(equations, variables) + catch + end + + if soll == Any[] + soll = Dict{SPyPyC.Sym{PythonCall.Core.Py}, SPyPyC.Sym{PythonCall.Core.Py}}() + elseif soll isa Vector + soll = Dict{SPyPyC.Sym{PythonCall.Core.Py}, SPyPyC.Sym{PythonCall.Core.Py}}(variables .=> soll[1]) + end + + return soll +end + + +function count_ops(expr)::Int + op_count = 0 + postwalk(x -> begin + if x isa Expr && x.head == :call + op_count += 1 + end + x + end, expr) + return op_count +end + + +function get_relevant_steady_states(𝓂::ℳ, + algorithm::Symbol; + opts::CalculationOptions = merge_calculation_options())::Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}} + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_parameters + full_NSSS = ms.full_NSSS_display + + relevant_SS = get_steady_state(𝓂, algorithm = algorithm, + stochastic = algorithm != :first_order, + return_variables_only = true, + derivatives = false, + verbose = opts.verbose, + tol = opts.tol, + quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, + sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) + + reference_steady_state = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : relevant_SS(s) for s in full_NSSS] + + relevant_NSSS = get_steady_state(𝓂, algorithm = :first_order, + stochastic = false, + return_variables_only = true, + derivatives = false, + verbose = opts.verbose, + tol = opts.tol, + quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, + sylvester_algorithm = [opts.sylvester_algorithm², opts.sylvester_algorithm³]) + + NSSS = [s ∈ 𝓂.constants.post_model_macro.exo_present ? 0.0 : relevant_NSSS(s) for s in full_NSSS] + + SSS_delta = NSSS - reference_steady_state + + return reference_steady_state, NSSS, SSS_delta +end + +# compatibility with SymPy +Max = max +Min = min + +function simplify(ex::Expr)::Union{Expr,Symbol,Int} + ex_ss = convert_to_ss_equation(ex) + + for x in get_symbols(ex_ss) + sym_value = SPyPyC.symbols(string(x), real = true, finite = true) + Core.eval(SymPyWorkspace, :($x = $sym_value)) + end + + parsed = ex_ss |> x -> Core.eval(SymPyWorkspace, x) |> string |> Meta.parse + + postwalk(x -> x isa Expr ? + x.args[1] == :conjugate ? + x.args[2] : + x : + x, parsed) +end + +function convert_to_ss_equation(eq::Expr)::Expr + 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 : + x.args[1] : + x.head == :call ? + x.args[1] == :* ? + x.args[2] isa Int ? + x.args[3] isa Int ? + x : + :($(x.args[3]) * $(x.args[2])) : # avoid 2X syntax. doesn't work with sympy + x : + x : + unblock(x) : + x, + eq) +end + + +replace_indices(x::Symbol) = x + +replace_indices_special(x::Symbol) = x + +replace_indices(x::String) = Symbol(replace(x, "{" => "◖", "}" => "◗")) + +replace_indices_in_symbol(x::Symbol) = replace(string(x), "◖" => "{", "◗" => "}") + +function replace_indices(exxpr::Expr)::Union{Expr,Symbol} + postwalk(x -> begin + x isa Symbol ? + replace_indices(string(x)) : + x isa Expr ? + x.head == :curly ? + Symbol(string(x.args[1]) * "◖" * string(x.args[2]) * "◗") : + x : + x + end, exxpr) +end + +function replace_indices_special(exxpr::Expr)::Union{Expr,Symbol} + postwalk(x -> begin + x isa Symbol ? + replace_indices(string(x)) : + x isa Expr ? + x.head == :curly ? + Symbol(string(x.args[1]) * "◖" * string(x.args[2]) * "◗") : + x.head == :call ? + x.args[1] == :(*) ? + Symbol(string(x.args[2]), string(x.args[3])) : + x : + x : + x + end, exxpr) +end + + +function expand_steady_state(SS_and_pars::Vector{M}, ms::post_complete_parameters) where M + X = ms.steady_state_expand_matrix + return X * SS_and_pars +end + + + +function create_symbols_eqs!(𝓂::ℳ)::symbolics + # create symbols in SymPyWorkspace to avoid polluting MacroModelling namespace + symbols_in_dynamic_equations = reduce(union, get_symbols.(𝓂.equations.dynamic)) + + symbols_in_dynamic_equations_wo_subscripts = Symbol.(replace.(string.(symbols_in_dynamic_equations), r"₍₋?(₀|₁|ₛₛ|ₓ)₎$"=>"")) + + symbols_in_ss_equations = reduce(union,get_symbols.(𝓂.equations.steady_state_aux)) + + symbols_in_equation = union(𝓂.constants.post_model_macro.parameters_in_equations, + 𝓂.constants.post_complete_parameters.parameters, + 𝓂.constants.post_parameters_macro.parameters_as_function_of_parameters, + symbols_in_dynamic_equations, + symbols_in_dynamic_equations_wo_subscripts, + symbols_in_ss_equations) #, 𝓂.dynamic_variables_future) + + symbols_pos = [] + symbols_neg = [] + symbols_none = [] + + for symb in symbols_in_equation + if haskey(𝓂.constants.post_parameters_macro.bounds, symb) + if 𝓂.constants.post_parameters_macro.bounds[symb][1] >= 0 + push!(symbols_pos, symb) + elseif 𝓂.constants.post_parameters_macro.bounds[symb][2] <= 0 + push!(symbols_neg, symb) + else + push!(symbols_none, symb) + end + else + push!(symbols_none, symb) + end + end + + # Create symbols in SymPyWorkspace instead of MacroModelling namespace + for pos in symbols_pos + sym_value = SPyPyC.symbols(string(pos), real = true, finite = true, positive = true) + Core.eval(SymPyWorkspace, :($pos = $sym_value)) + end + + for neg in symbols_neg + sym_value = SPyPyC.symbols(string(neg), real = true, finite = true, negative = true) + Core.eval(SymPyWorkspace, :($neg = $sym_value)) + end + + for none in symbols_none + sym_value = SPyPyC.symbols(string(none), real = true, finite = true) + Core.eval(SymPyWorkspace, :($none = $sym_value)) + end + + symbolics( + map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.equations.steady_state_aux), + # map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.dyn_equations_future), + + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift_var_present_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift_var_past_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift_var_future_list), + + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_shift2_var_past_list), + + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_var_present_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_var_past_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_var_future_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_ss_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_exo_list), + + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_exo_future_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_exo_present_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dyn_exo_past_list), + + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_future_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_present_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.dyn_past_list), + + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_present_list_aux_SS), + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_past_list_aux_SS), + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_future_list_aux_SS), + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.ss_list_aux_SS), + + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.var_list_aux_SS), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dynamic_variables_list), + # map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.dynamic_variables_future_list), + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_model_macro.par_list_aux_SS), + + map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.equations.calibration), + map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.equations.calibration_parameters), + # map(x->Core.eval(SymPyWorkspace, :($x)),𝓂.constants.post_complete_parameters.parameters), + + # Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.var_present...)]))), + # Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.var_past...)]))), + # Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.var_future...)]))), + Set(Core.eval(SymPyWorkspace, :([$(𝓂.constants.post_model_macro.vars_in_ss_equations...)]))), + + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_parameters_macro.ss_calib_list), + map(x->Set(Core.eval(SymPyWorkspace, :([$(x...)]))),𝓂.constants.post_parameters_macro.par_calib_list), + + [Set() for _ in 1:length(𝓂.equations.steady_state_aux)], + # [Set() for _ in 1:length(𝓂.calibration_equations)], + # [Set() for _ in 1:length(𝓂.equations.steady_state_aux)], + # [Set() for _ in 1:length(𝓂.calibration_equations)] + ) +end + + + +function remove_redundant_SS_vars!(𝓂::ℳ, Symbolics::symbolics; avoid_solve::Bool = false) + ss_equations = Symbolics.ss_equations + + # check variables which appear in two time periods. they might be redundant in steady state + redundant_vars = intersect.( + union.( + intersect.(Symbolics.var_future_list_aux_SS, Symbolics.var_present_list_aux_SS), + intersect.(Symbolics.var_future_list_aux_SS, Symbolics.var_past_list_aux_SS), + intersect.(Symbolics.var_present_list_aux_SS, Symbolics.var_past_list_aux_SS), + intersect.(Symbolics.ss_list_aux_SS, Symbolics.var_present_list_aux_SS), + intersect.(Symbolics.ss_list_aux_SS, Symbolics.var_past_list_aux_SS), + intersect.(Symbolics.ss_list_aux_SS, Symbolics.var_future_list_aux_SS) + ), + Symbolics.var_list_aux_SS) + + redundant_idx = getindex(1:length(redundant_vars), (length.(redundant_vars) .> 0) .& (length.(Symbolics.var_list_aux_SS) .> 1)) + for i in redundant_idx + # Cheap Julia-side symbol set for this equation, used to skip SymPy work + # when a candidate variable does not actually appear in the equation. + eq_symbols = Set{Symbol}(get_symbols(Meta.parse(string(ss_equations[i])))) + for var_to_solve_for in redundant_vars[i] + var_sym_candidate = Symbol(var_to_solve_for) + if !(var_sym_candidate in eq_symbols) + # variable already absent (e.g. previous redundancy rewrites removed it) + if var_to_solve_for ∉ Symbolics.var_redundant_list[i] + push!(Symbolics.var_redundant_list[i], var_to_solve_for) + end + continue + end + + if avoid_solve || count_ops(Meta.parse(string(ss_equations[i]))) > 15 + soll = nothing + else + soll = solve_symbolically(ss_equations[i],var_to_solve_for) + end + + if isnothing(soll) + continue + end + + if isempty(soll) || isequal(soll, SPyPyC.Sym{PythonCall.Core.Py}[0]) # take out variable if it is redundant from that euation only + push!(Symbolics.var_redundant_list[i],var_to_solve_for) + ss_equations[i] = replace_with_one(ss_equations[i], var_to_solve_for) # replace euler constant as it is not translated to julia properly + # refresh symbol set since the equation was rewritten + eq_symbols = Set{Symbol}(get_symbols(Meta.parse(string(ss_equations[i])))) + end + + end + end + +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, 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!(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 + + # 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 + + + # 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.NSSS_∂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) + + 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) + buffer = zeros(Float64, size(∂SS_equations_∂SS_and_pars)) + else + 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.NSSS_∂equations_∂SS_and_pars = buffer + 𝓂.functions.NSSS_∂equations_∂SS_and_pars = func_exprs + + return nothing +end + + +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) + + if !silent + println(round(time() - start_time, digits = 3), " seconds") + end + + return nothing +end + + +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) + use_symbolics = !𝓂.constants.post_parameters_macro.precompile + + if use_symbolics + start_time = time() + + if !silent print("Remove redundant variables in non-stochastic steady state problem:\t") end + + symbolics = create_symbols_eqs!(𝓂) + + remove_redundant_SS_vars!(𝓂, symbolics, avoid_solve = avoid_solve) + + if !silent println(round(time() - start_time, digits = 3), " seconds") end + + start_time = time() + + if !silent print("Set up non-stochastic steady state problem:\t\t\t\t") end + + write_ss_check_function!(𝓂) + + write_steady_state_solver_function!(𝓂, symbolic_enabled, symbolics, verbose = verbose, avoid_solve = avoid_solve) + + 𝓂.equations.obc_violation = write_obc_violation_equations(𝓂) + + set_up_obc_violation_function!(𝓂) + + if !silent println(round(time() - start_time, digits = 3), " seconds") end + else + start_time = time() + + if !silent print("Set up non-stochastic steady state problem:\t\t\t\t") end + + write_ss_check_function!(𝓂) + + write_steady_state_solver_function!(𝓂, false, nothing, verbose = verbose, avoid_solve = avoid_solve) + + if !𝓂.constants.post_parameters_macro.precompile + 𝓂.equations.obc_violation = write_obc_violation_equations(𝓂) + set_up_obc_violation_function!(𝓂) + end + + if !silent println(round(time() - start_time, digits = 3), " seconds") end + end + + return nothing +end + + +function take_nth_order_derivatives( + dyn_equations::Vector{T}, + 𝔙::Symbolics.Arr, + 𝔓::Symbolics.Arr, + SS_mapping::Dict{T, T}, + nps::Int, + nxs::Int; + max_perturbation_order::Int = 1, + output_compressed::Bool = true # Controls compression for X derivatives (order >= 2) +)::Vector{Tuple{SparseMatrixCSC{T, Int}, SparseMatrixCSC{T, Int}}} where T <: Symbolics.Num#, Tuple{Symbolics.Arr{Symbolics.Num, 1}, Symbolics.Arr{Symbolics.Num, 1}}} + + nx = BigInt(length(𝔙)::Int) + # np = length(𝔓)::BigInt + nϵ = length(dyn_equations)::Int + + if max_perturbation_order < 1 + throw(ArgumentError("max_perturbation_order must be at least 1")) + end + + results = [] # To store pairs of sparse matrices (X_matrix, P_matrix) for each order + + # --- Order 1 --- + # Compute the 1st order derivative with respect to X (Jacobian) + spX_order_1 = Symbolics.sparsejacobian(dyn_equations, 𝔙) # nϵ x nx + + + spX_order_1_sub = copy(spX_order_1) + + # spX_order_1_sub.nzval .= Symbolics.fast_substitute(spX_order_1_sub.nzval, Dict(Symbolics.scalarize(𝔛𝔛) .=> 𝔙)) + spX_order_1_sub.nzval .= Symbolics.substitute(spX_order_1_sub.nzval, SS_mapping) + + # Compute the derivative of the non-zeros of the 1st X-derivative w.r.t. P + # This is an intermediate step. The final P matrix will be built from this. + spP_of_flatX_nzval_order_1 = Symbolics.sparsejacobian(spX_order_1_sub.nzval, vcat(𝔓[1:nps], 𝔙[1:nxs])) # nnz(spX_order_1) x np + + # Determine dimensions for the Order 1 P matrix + X_nrows_1 = nϵ + X_ncols_1 = nx + P_nrows_1 = X_nrows_1 * X_ncols_1 + P_ncols_1 = nps + nxs + + # Build the Order 1 P matrix (dimensions nϵ*nx x np) + sparse_rows_1_P = Int[] # Row index in the flattened space of spX_order_1 + sparse_cols_1_P = Int[] # Column index for parameters (1 to np) + sparse_vals_1_P = Symbolics.Num[] + + # Map linear index in spX_order_1.nzval to its (row, col) in spX_order_1 + nz_lin_to_rc_1 = Dict{Int, Tuple{Int, Int}}() + k_lin = 1 + for j = 1:size(spX_order_1, 2) # col + for ptr = spX_order_1.colptr[j]:(spX_order_1.colptr[j+1]-1) + r = spX_order_1.rowval[ptr] # row + nz_lin_to_rc_1[k_lin] = (r, j) + k_lin += 1 + end + end + + + # Iterate through the non-zero entries of spP_of_flatX_nzval_order_1 + k_temp_P = 1 # linear index counter for nzval + for p_col = 1:size(spP_of_flatX_nzval_order_1, 2) # Parameter index + for i_ptr_temp_P = spP_of_flatX_nzval_order_1.colptr[p_col]:(spP_of_flatX_nzval_order_1.colptr[p_col+1]-1) + temp_row = spP_of_flatX_nzval_order_1.rowval[i_ptr_temp_P] # Row index in spP_of_flatX_nzval (corresponds to temp_row-th nzval of spX_order_1) + p_val = spP_of_flatX_nzval_order_1.nzval[i_ptr_temp_P] # Derivative value w.r.t. parameter + + # Get the (row, col) in spX_order_1 corresponding to this derivative + r_X1, c_X1 = nz_lin_to_rc_1[temp_row] + + # Calculate the row index in spP_order_1 (flattened index of spX_order_1) + # P_row_idx = (r_X1 - 1) * X_ncols_1 + c_X1 + P_row_idx = (c_X1 - 1) * X_nrows_1 + r_X1 + P_col_idx = p_col # Parameter column index + + push!(sparse_rows_1_P, P_row_idx) + push!(sparse_cols_1_P, P_col_idx) + push!(sparse_vals_1_P, p_val) + + k_temp_P += 1 + end + end + + spP_order_1 = sparse!(sparse_rows_1_P, sparse_cols_1_P, sparse_vals_1_P, P_nrows_1, P_ncols_1) + + + # Store the pair for order 1 + push!(results, (spX_order_1_sub, spP_order_1)) + + if max_perturbation_order > 1 + # --- Prepare for higher orders (Order 2 to max_perturbation_order) --- + # Initialize map for Order 1: linear index in spX_order_1.nzval -> (row, (v1,)) + # This map is needed to trace indices for Order 2 + # We already built nz_lin_to_rc_1 above, reuse it and wrap the variable index in a Tuple + nz_to_indices_prev = Dict{Int, Tuple{Int, Tuple{Int}}}() + k_lin = 1 + for j = 1:size(spX_order_1, 2) + for ptr = spX_order_1.colptr[j]:(spX_order_1.colptr[j+1]-1) + r = spX_order_1.rowval[ptr] + nz_to_indices_prev[k_lin] = (r, (j,)) # Store (equation row, (v1,)) + k_lin += 1 + end + end + + nzvals_prev = spX_order_1.nzval # nzvals from Order 1 X-matrix + + # --- Iterate for orders n = 2, 3, ..., max_perturbation_order --- + for n = 2:max_perturbation_order + + # Compute the Jacobian of the previous level's nzval w.r.t. 𝔛 + # This gives a flat matrix where rows correspond to non-zeros from order n-1 X-matrix + # and columns correspond to the n-th variable we differentiate by (x_vn). + sp_flat_curr_X_rn = Symbolics.sparsejacobian(nzvals_prev, 𝔙) # nnz(spX_order_(n-1)) x nx + + sp_flat_curr_X = copy(sp_flat_curr_X_rn) + + sp_flat_curr_X.nzval .= Symbolics.substitute(sp_flat_curr_X.nzval, SS_mapping) + + # Build the nz_to_indices map for the *current* level (order n) + # Map: linear index in sp_flat_curr_X.nzval -> (original_row_f, (v_1, ..., v_n)) + nz_to_indices_curr = Dict{Int, Tuple{Int, Tuple{Vararg{Int}}}}() + k_lin_curr = 1 # linear index counter for nzval of sp_flat_curr_X + # Iterate through the non-zeros of the current flat Jacobian + for col_curr = 1:size(sp_flat_curr_X, 2) # Column index in sp_flat_curr_X (corresponds to v_n) + for ptr_curr = sp_flat_curr_X.colptr[col_curr]:(sp_flat_curr_X.colptr[col_curr+1]-1) + row_curr = sp_flat_curr_X.rowval[ptr_curr] # Row index in sp_flat_curr_X (corresponds to the row_curr-th nzval of previous level) + + # Get previous indices info from the map of order n-1 + prev_info = nz_to_indices_prev[row_curr] + orig_row_f = prev_info[1] # Original equation row + vars_prev = prev_info[2] # Tuple of variables from previous order (v_1, ..., v_{n-1}) + + # Append the current variable index (v_n) + vars_curr = (vars_prev..., col_curr) # Full tuple (v_1, ..., v_n) + + # Store info for the current level's non-zero + nz_to_indices_curr[k_lin_curr] = (orig_row_f, vars_curr) + k_lin_curr += 1 + end + end + + # --- Construct the X-derivative sparse matrix for order n (compressed or uncompressed) --- + local spX_order_n # Declare variable to hold the resulting X matrix + local X_ncols_n # Number of columns in the resulting spX_order_n matrix + + if output_compressed + # COMPRESSED output: nϵ x binomial(nx + n - 1, n) + sparse_rows_n = Int[] + sparse_cols_n = Int[] # This will store the compressed column index + sparse_vals_n = Symbolics.Num[] + + # Calculate the total number of compressed columns for order n + X_ncols_n = Int(binomial(nx + n - 1, n)) + + # Iterate through the non-zero entries of the current flat Jacobian (sp_flat_curr_X) + k_flat_curr = 1 # linear index counter for nzval of sp_flat_curr_X + for col_flat_curr = 1:size(sp_flat_curr_X, 2) # This corresponds to the n-th variable (v_n) + for i_ptr_flat_curr = sp_flat_curr_X.colptr[col_flat_curr]:(sp_flat_curr_X.colptr[col_flat_curr+1]-1) + # row_flat_curr = sp_flat_curr_X.rowval[i_ptr_flat_curr] # Row index in sp_flat_curr_X + val = sp_flat_curr_X.nzval[i_ptr_flat_curr] # The derivative value + + # Get the full info for this non-zero from the map + # The linear index in sp_flat_curr_X.nzval is k_flat_curr + orig_row_f, var_indices_full = nz_to_indices_curr[k_flat_curr] # (v_1, ..., v_n) + + # Check the compression rule: v_n <= v_{n-1} <= ... <= v_1 + is_compressed = true + for k_rule = 1:(n-1) + # Check v_{n-k_rule+1} <= v_{n-k_rule} + if var_indices_full[n-k_rule+1] > var_indices_full[n-k_rule] + is_compressed = false + break + end + end + + if is_compressed + # Calculate the compressed column index c_n for the tuple (v_1, ..., v_n) + # using the derived formula: c_n = sum_{k=1}^{n-1} binomial(v_k + n - k - 1, n - k + 1) + v_n + compressed_col_idx = 0 + for k_formula = 1:(n-1) + term = binomial(var_indices_full[k_formula] + n - k_formula - 1, n - k_formula + 1) + compressed_col_idx += term + end + # Add the last term: v_n (var_indices_full[n]) + compressed_col_idx += var_indices_full[n] + + push!(sparse_rows_n, orig_row_f) + push!(sparse_cols_n, compressed_col_idx) + push!(sparse_vals_n, val) + end + + k_flat_curr += 1 # Increment linear index counter for sp_flat_curr_X.nzval + end + end + # Construct the compressed sparse matrix for order n + spX_order_n = sparse!(sparse_rows_n, sparse_cols_n, sparse_vals_n, X_nrows_1, X_ncols_n) + + else # output_compressed == false + # UNCOMPRESSED output: nϵ x nx^n + sparse_rows_n_uncomp = Int[] + sparse_cols_n_uncomp = Int[] # Uncompressed column index (1 to nx^n) + sparse_vals_n_uncomp = Symbolics.Num[] + + # Total number of uncompressed columns + X_ncols_n = nx^n # Use BigInt for the power calculation, cast to Int + + # Iterate through the non-zero entries of the current flat Jacobian (sp_flat_curr_X) + k_flat_curr = 1 # linear index counter for nzval of sp_flat_curr_X + for col_flat_curr = 1:size(sp_flat_curr_X, 2) # This corresponds to the n-th variable (v_n) + for i_ptr_flat_curr = sp_flat_curr_X.colptr[col_flat_curr]:(sp_flat_curr_X.colptr[col_flat_curr+1]-1) + # row_flat_curr = sp_flat_curr_X.rowval[i_ptr_flat_curr] # Row index in sp_flat_curr_X + val = sp_flat_curr_X.nzval[i_ptr_flat_curr] # The derivative value + + # Get the full info for this non-zero from the map + # The linear index in sp_flat_curr_X.nzval is k_flat_curr + orig_row_f, var_indices_full = nz_to_indices_curr[k_flat_curr] # (v_1, ..., v_n) + + # Calculate the UNCOMPRESSED column index for the tuple (v_1, ..., v_n) + # This maps the tuple (v1, ..., vn) to a unique index from 1 to nx^n + # Formula: 1 + (v1-1)*nx^(n-1) + (v2-1)*nx^(n-2) + ... + (vn-1)*nx^0 + uncompressed_col_idx = 1 # 1-based + power_of_nx = nx^(n-1) # Start with nx^(n-1) for v1 term + for i = 1:n + uncompressed_col_col_idx_term = (var_indices_full[i] - 1) * power_of_nx + # Check for overflow before adding + # if (uncompressed_col_idx > 0 && uncompressed_col_col_idx_term > 0 && uncompressed_col_idx + uncompressed_col_col_idx_term <= uncompressed_col_idx) || + # (uncompressed_col_idx < 0 && uncompressed_col_col_idx_term < 0 && uncompressed_col_idx + uncompressed_col_col_idx_term >= uncompressed_col_idx) + # error("Integer overflow calculating uncompressed column index") + # end + uncompressed_col_idx += uncompressed_col_col_idx_term + + if i < n # Avoid nx^-1 + power_of_nx = div(power_of_nx, nx) # Integer division + end + end + + push!(sparse_rows_n_uncomp, orig_row_f) + push!(sparse_cols_n_uncomp, Int(uncompressed_col_idx)) # Cast to Int + push!(sparse_vals_n_uncomp, val) + + k_flat_curr += 1 # Increment linear index counter for sp_flat_curr_X.nzval + end + end + # Construct the uncompressed sparse matrix for order n + spX_order_n = sparse!(sparse_rows_n_uncomp, sparse_cols_n_uncomp, sparse_vals_n_uncomp, X_nrows_1, X_ncols_n) + + end # End of if output_compressed / else + + + # --- Compute the P-derivative sparse matrix for order n --- + # This is the Jacobian of the nzval of the intermediate flat X-Jacobian (sp_flat_curr_X) w.r.t. 𝔓. + # sp_flat_curr_X.nzval contains expressions for d^n f_i / (dx_v1 ... dx_vn) for all + # non-zero such values that were propagated from the previous step. + spP_of_flatX_nzval_curr = Symbolics.sparsejacobian(sp_flat_curr_X.nzval, vcat(𝔓[1:nps], 𝔙[1:nxs])) # nnz(sp_flat_curr_X) x np + + # Determine the desired dimensions of spP_order_n + # Dimensions are (rows of spX_order_n * cols of spX_order_n) x np + P_nrows_n = nϵ * X_ncols_n + P_ncols_n = nps + nxs + + sparse_rows_n_P = Int[] # Row index in the flattened space of spX_order_n (1 to P_nrows_n) + sparse_cols_n_P = Int[] # Column index for parameters (1 to np) + sparse_vals_n_P = Symbolics.Num[] + + # Iterate through the non-zero entries of spP_of_flatX_nzval_curr + # Its rows correspond to the non-zeros in sp_flat_curr_X + k_temp_P = 1 # linear index counter for nzval of spP_of_flatX_nzval_curr + for p_col = 1:size(spP_of_flatX_nzval_curr, 2) # Column index in spP_of_flatX_nzval_curr (corresponds to parameter index) + for i_ptr_temp_P = spP_of_flatX_nzval_curr.colptr[p_col]:(spP_of_flatX_nzval_curr.colptr[p_col+1]-1) + temp_row = spP_of_flatX_nzval_curr.rowval[i_ptr_temp_P] # Row index in spP_of_flatX_nzval_curr (corresponds to the temp_row-th nzval of sp_flat_curr_X) + p_val = spP_of_flatX_nzval_curr.nzval[i_ptr_temp_P] # The derivative w.r.t. parameter value + + # Get the full info for the X-derivative term that this P-derivative is from + # temp_row is the linear index in sp_flat_curr_X.nzval + # This corresponds to the derivative d^n f_orig_row_f / (dx_v1 ... dx_vn) + orig_row_f, var_indices_full = nz_to_indices_curr[temp_row] # (v_1, ..., v_n) + + # We need to find the column index (X_col_idx) this term corresponds to + # in the final spX_order_n matrix (which might be compressed or uncompressed) + 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) + term = binomial(var_indices_full[k_formula] + n - k_formula - 1, n - k_formula + 1) + compressed_col_idx += term + end + compressed_col_idx += var_indices_full[n] + X_col_idx = compressed_col_idx # The column in spX_order_n is the compressed one + + else # output_compressed == false + # Calculate the uncompressed column index + uncompressed_col_idx = 1 + power_of_nx = nx^(n-1) + for i = 1:n + uncompressed_col_idx += (var_indices_full[i] - 1) * power_of_nx + if i < n + power_of_nx = div(power_of_nx, nx) + end + end + X_col_idx = Int(uncompressed_col_idx) # The column in spX_order_n is the uncompressed one + end + + # Calculate the row index in spP_order_n + # This maps the (orig_row_f, X_col_idx) pair in spX_order_n's grid to a linear index + # Formula: (row_in_X - 1) * num_cols_in_X + col_in_X + # P_row_idx = (orig_row_f - 1) * X_ncols_n + X_col_idx + P_row_idx = (X_col_idx - 1) * nϵ + orig_row_f + + # The column index in spP_order_n is the parameter index + P_col_idx = p_col + + push!(sparse_rows_n_P, P_row_idx) + push!(sparse_cols_n_P, P_col_idx) + push!(sparse_vals_n_P, p_val) + + k_temp_P += 1 # Increment linear index counter for spP_of_flatX_nzval_curr.nzval + end + end + + # Construct the P-derivative sparse matrix for order n + # Dimensions are (rows of spX_order_n * cols of spX_order_n) x np + spP_order_n = sparse!(sparse_rows_n_P, sparse_cols_n_P, sparse_vals_n_P, P_nrows_n, P_ncols_n) + + # Store the pair (X-matrix, P-matrix) for order n + push!(results, (spX_order_n, spP_order_n)) + + + # Prepare for the next iteration (order n+1) + # The nzvals for the next X-Jacobian step are the nzvals of the current flat X-Jacobian + nzvals_prev = sp_flat_curr_X_rn.nzval + # The map for the next step should provide info for order n derivatives + nz_to_indices_prev = nz_to_indices_curr + + end # End of loop for orders n = 2 to max_perturbation_order + end + + return results #, (𝔛, 𝔓) # Return results as a tuple of (X_matrix, P_matrix) pairs +end + + +function write_functions_mapping!(𝓂::ℳ, max_perturbation_order::Int; + density_threshold::Float64 = .1, + min_length::Int = 1000, + nnz_parallel_threshold::Int = 1000000, + # parallel = Symbolics.SerialForm(), + # parallel = Symbolics.ShardedForm(1500,4), + cse = true, + skipzeros = true) + + future_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₁₎$"))) + present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎$"))) + past_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₋₁₎$"))) + shock_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₓ₎$"))) + ss_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₛₛ₎$"))) + + sort!(future_varss ,by = x->replace(string(x),r"₍₁₎$"=>"")) #sort by name without time index because otherwise eps_zᴸ⁽⁻¹⁾₍₋₁₎ comes before eps_z₍₋₁₎ + sort!(present_varss ,by = x->replace(string(x),r"₍₀₎$"=>"")) + sort!(past_varss ,by = x->replace(string(x),r"₍₋₁₎$"=>"")) + sort!(shock_varss ,by = x->replace(string(x),r"₍ₓ₎$"=>"")) + sort!(ss_varss ,by = x->replace(string(x),r"₍ₛₛ₎$"=>"")) + + dyn_future_list = collect(reduce(union, 𝓂.constants.post_model_macro.dyn_future_list)) + dyn_present_list = collect(reduce(union, 𝓂.constants.post_model_macro.dyn_present_list)) + dyn_past_list = collect(reduce(union, 𝓂.constants.post_model_macro.dyn_past_list)) + dyn_exo_list = collect(reduce(union,𝓂.constants.post_model_macro.dyn_exo_list)) + dyn_ss_list = Symbol.(string.(collect(reduce(union,𝓂.constants.post_model_macro.dyn_ss_list))) .* "₍ₛₛ₎") + + future = map(x -> Symbol(replace(string(x), r"₍₁₎" => "")),string.(dyn_future_list)) + present = map(x -> Symbol(replace(string(x), r"₍₀₎" => "")),string.(dyn_present_list)) + past = map(x -> Symbol(replace(string(x), r"₍₋₁₎" => "")),string.(dyn_past_list)) + exo = map(x -> Symbol(replace(string(x), r"₍ₓ₎" => "")),string.(dyn_exo_list)) + stst = map(x -> Symbol(replace(string(x), r"₍ₛₛ₎" => "")),string.(dyn_ss_list)) + + vars_raw = vcat(dyn_future_list[indexin(sort(future),future)], + dyn_present_list[indexin(sort(present),present)], + dyn_past_list[indexin(sort(past),past)], + dyn_exo_list[indexin(sort(exo),exo)]) + + dyn_var_future_idx = 𝓂.constants.post_complete_parameters.dyn_var_future_idx + dyn_var_present_idx = 𝓂.constants.post_complete_parameters.dyn_var_present_idx + dyn_var_past_idx = 𝓂.constants.post_complete_parameters.dyn_var_past_idx + dyn_ss_idx = 𝓂.constants.post_complete_parameters.dyn_ss_idx + + dyn_var_idxs = vcat(dyn_var_future_idx, dyn_var_present_idx, dyn_var_past_idx) + + pars_ext = vcat(𝓂.constants.post_complete_parameters.parameters, 𝓂.equations.calibration_parameters) + parameters_and_SS = vcat(pars_ext, dyn_ss_list[indexin(sort(stst),stst)]) + + np = length(parameters_and_SS) + nv = length(vars_raw) + nc = length(𝓂.equations.calibration) + nps = length(𝓂.constants.post_complete_parameters.parameters) + nxs = maximum(dyn_var_idxs) + nc + + Symbolics.@variables 𝔓[1:np] 𝔙[1:nv] + + parameter_dict = Dict{Symbol, Symbol}() + back_to_array_dict = Dict{Symbolics.Num, Symbolics.Num}() + calib_vars = Symbol[] + calib_expr = [] + SS_mapping = Dict{Symbolics.Num, Symbolics.Num}() + + + for (i,v) in enumerate(parameters_and_SS) + push!(parameter_dict, v => :($(Symbol("𝔓_$i")))) + push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("𝔓_$i"))), @__MODULE__) => 𝔓[i]) + if i > nps + if i > length(pars_ext) + push!(SS_mapping, 𝔓[i] => 𝔙[dyn_ss_idx[i-length(pars_ext)]]) + else + push!(SS_mapping, 𝔓[i] => 𝔙[nxs + i - nps - nc]) + end + end + end + + for (i,v) in enumerate(vars_raw) + push!(parameter_dict, v => :($(Symbol("𝔙_$i")))) + push!(back_to_array_dict, Symbolics.parse_expr_to_symbolic(:($(Symbol("𝔙_$i"))), @__MODULE__) => 𝔙[i]) + if i <= length(dyn_var_idxs) + push!(SS_mapping, 𝔙[i] => 𝔙[dyn_var_idxs[i]]) + else + push!(SS_mapping, 𝔙[i] => 0) + end + end + + + for v in 𝓂.equations.calibration_no_var + push!(calib_vars, v.args[1]) + push!(calib_expr, v.args[2]) + 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!(calib_replacements, x => calib_expr[i]) + end + + + dyn_equations = 𝓂.equations.dynamic |> + 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)) + + 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) + local nz_count = nnz(transposed) + + if (nz_count / length(transposed) > density_threshold) || (length(transposed) < min_length) + return convert(Matrix, transposed), zeros(Float64, size(transposed)), nz_count + end + + local buf = similar(transposed, Float64) + buf.nzval .= 0 + return transposed, buf, nz_count + end + + + ∇₁_dyn = derivatives[1][1] + + lennz = nnz(∇₁_dyn) + + 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 + derivatives_mat = ∇₁_dyn + buffer = similar(∇₁_dyn, 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, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.jacobian = buffer + + + ∇₁_parameters_mat, buffer_parameters, lennz = prepare_sensitivity_buffer(derivatives[1][2][:,1:nps]) + + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end + + _, func_∇₁_parameters = Symbolics.build_function(∇₁_parameters_mat, 𝔓, 𝔙, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.jacobian_parameters = buffer_parameters + + + ∇₁_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) + else + parallel = Symbolics.SerialForm() + end + + _, func_∇₁_SS_and_pars = Symbolics.build_function(∇₁_SS_and_pars_mat, 𝔓, 𝔙, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.jacobian_SS_and_pars = buffer_SS_and_pars + + # Create jacobian_functions struct with all three functions + 𝓂.functions.jacobian = jacobian_functions(func_exprs, func_∇₁_parameters, func_∇₁_SS_and_pars) + + + + + # if max_perturbation_order >= 1 + # 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(eqs) + + # 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) + # println(ϵˢ) + # ∂SS_equations_∂parameters = Symbolics.sparsejacobian(ϵˢ, 𝔛¹) # 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, + # parallel = parallel, + # # nanmath = false, + # expression_module = @__MODULE__, + # expression = Val(false))::Tuple{<:Function, <:Function} + + # 𝓂.functions.NSSS_∂equations_∂parameters = func_exprs + + + + # ∂SS_equations_∂SS_and_pars = Symbolics.sparsejacobian(ϵˢ, 𝔓¹) # nϵ x nx + + # 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) + # buffer = zeros(Float64, size(∂SS_equations_∂SS_and_pars)) + # else + # 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, + # parallel = parallel, + # # nanmath = false, + # expression_module = @__MODULE__, + # expression = Val(false))::Tuple{<:Function, <:Function} + + # 𝓂.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 = true) + + if 𝓂.constants.second_order.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) + ∇₂_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) + derivatives_mat = convert(Matrix, ∇₂_dyn) + buffer = zeros(Float64, size(∇₂_dyn)) + else + derivatives_mat = ∇₂_dyn + buffer = similar(∇₂_dyn, 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, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.hessian = buffer + + + ∇₂_parameters_mat, buffer_parameters, lennz = prepare_sensitivity_buffer(derivatives[2][2][:,1:nps]) + + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end + + _, func_∇₂_parameters = Symbolics.build_function(∇₂_parameters_mat, 𝔓, 𝔙, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.hessian_parameters = buffer_parameters + + + ∇₂_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) + else + parallel = Symbolics.SerialForm() + end + + _, func_∇₂_SS_and_pars = Symbolics.build_function(∇₂_SS_and_pars_mat, 𝔓, 𝔙, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.hessian_SS_and_pars = buffer_SS_and_pars + + # Create hessian_functions struct with all three functions + 𝓂.functions.hessian = hessian_functions(func_exprs, func_∇₂_parameters, func_∇₂_SS_and_pars) + end + end + + if max_perturbation_order == 3 + derivatives = take_nth_order_derivatives(dyn_equations, 𝔙, 𝔓, SS_mapping, nps, nxs; max_perturbation_order = max_perturbation_order, output_compressed = true) + # third order + if 𝓂.constants.third_order.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) + I,J,V = findnz(derivatives[3][1]) + 𝓂.constants.third_order = create_third_order_auxiliary_matrices(𝓂.constants, unique(J)) + + ∇₃_dyn = derivatives[3][1] + + lennz = nnz(∇₃_dyn) + + if (lennz / length(∇₃_dyn) > density_threshold) || (length(∇₃_dyn) < min_length) + derivatives_mat = convert(Matrix, ∇₃_dyn) + buffer = zeros(Float64, size(∇₃_dyn)) + else + derivatives_mat = ∇₃_dyn + buffer = similar(∇₃_dyn, 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, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.third_order_derivatives = buffer + + + ∇₃_parameters_mat, buffer_parameters, lennz = prepare_sensitivity_buffer(derivatives[3][2][:,1:nps]) + + if lennz > nnz_parallel_threshold + parallel = Symbolics.ShardedForm(1500,4) + else + parallel = Symbolics.SerialForm() + end + + _, func_∇₃_parameters = Symbolics.build_function(∇₃_parameters_mat, 𝔓, 𝔙, + cse = cse, + skipzeros = skipzeros, + parallel = parallel, + # nanmath = false, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.third_order_derivatives_parameters = buffer_parameters + + + ∇₃_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) + else + parallel = Symbolics.SerialForm() + end + + _, func_∇₃_SS_and_pars = Symbolics.build_function(∇₃_SS_and_pars_mat, 𝔓, 𝔙, + cse = cse, + skipzeros = skipzeros, + # nanmath = false, + parallel = parallel, + expression_module = @__MODULE__, + expression = Val(false))::Tuple{<:Function, <:Function} + + 𝓂.caches.third_order_derivatives_SS_and_pars = buffer_SS_and_pars + + # Create third_order_derivatives_functions struct with all three functions + 𝓂.functions.third_order_derivatives = third_order_derivatives_functions(func_exprs, func_∇₃_parameters, func_∇₃_SS_and_pars) + 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 + + +function write_auxiliary_indices!(𝓂::ℳ) + # write indices in auxiliary objects + dyn_var_future_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₁₎"))) + dyn_var_present_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₀₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₀₎"))) + dyn_var_past_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍₋₁₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍₋₁₎"))) + dyn_exo_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₓ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₓ₎"))) + dyn_ss_list = map(x->Set{Symbol}(map(x->Symbol(replace(string(x),"₍ₛₛ₎" => "")),x)),collect.(match_pattern.(get_symbols.(𝓂.equations.dynamic),r"₍ₛₛ₎"))) + + dyn_var_future = Symbol.(string.(sort(collect(reduce(union,dyn_var_future_list))))) + dyn_var_present = Symbol.(string.(sort(collect(reduce(union,dyn_var_present_list))))) + dyn_var_past = Symbol.(string.(sort(collect(reduce(union,dyn_var_past_list))))) + dyn_exo = Symbol.(string.(sort(collect(reduce(union,dyn_exo_list))))) + dyn_ss = Symbol.(string.(sort(collect(reduce(union,dyn_ss_list))))) + + SS_and_pars_names = vcat(Symbol.(string.(sort(union(𝓂.constants.post_model_macro.var,𝓂.constants.post_model_macro.exo_past,𝓂.constants.post_model_macro.exo_future)))), 𝓂.equations.calibration_parameters) + + dyn_var_future_idx = indexin(dyn_var_future , SS_and_pars_names) + dyn_var_present_idx = indexin(dyn_var_present , SS_and_pars_names) + dyn_var_past_idx = indexin(dyn_var_past , SS_and_pars_names) + dyn_ss_idx = indexin(dyn_ss , SS_and_pars_names) + + shocks_ss = zeros(length(dyn_exo)) + + 𝓂.constants.post_complete_parameters = update_post_complete_parameters( + 𝓂.constants.post_complete_parameters; + dyn_var_future_idx = dyn_var_future_idx, + dyn_var_present_idx = dyn_var_present_idx, + dyn_var_past_idx = dyn_var_past_idx, + dyn_ss_idx = dyn_ss_idx, + shocks_ss = shocks_ss, + ) + + return nothing +end diff --git a/src/perturbation.jl b/src/perturbation.jl deleted file mode 100644 index 1fbe42d05..000000000 --- a/src/perturbation.jl +++ /dev/null @@ -1,489 +0,0 @@ -@stable default_mode = "disable" begin - -function 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))::Tuple{Matrix{R}, Matrix{R}, Bool} where {R <: AbstractFloat, S <: Real} - # @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₊[dynIndex,:] * Ir[future_not_past_and_mixed_in_comb,:] - Ã₀ = A₀[dynIndex, comb] - Ã₋ = 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(Ã₊, Ã₀, Ã₋, 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 - 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 - 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₀[1:T.nPresent_only, T.present_only_idx] - A₊ᵤ = A₊[1:T.nPresent_only,:] - Ã₀ᵤ = A₀[1:T.nPresent_only, T.present_but_not_only_idx] - A₋ᵤ = A₋[1:T.nPresent_only,:] - - # end # timeit_debug - # @timeit_debug timer "Invert Ā₀ᵤ" begin - - Ā̂₀ᵤ = ℒ.lu!(Ā₀ᵤ, check = false) - - if !ℒ.issuccess(Ā̂₀ᵤ) - if opts.verbose println("Factorisation of Ā₀ᵤ failed") end - return zeros(R, T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false - end - - # A = vcat(-(Ā̂₀ᵤ \ (A₊ᵤ * D * L + Ã₀ᵤ * sol[T.dynamic_order,:] + A₋ᵤ)), sol) - if T.nPresent_only > 0 - ℒ.mul!(A₋ᵤ, Ã₀ᵤ, sol[:,past_not_future_and_mixed_in_comb], 1, 1) - nₚ₋ = A₊ᵤ * D - ℒ.mul!(A₋ᵤ, nₚ₋, L, 1, 1) - ℒ.ldiv!(Ā̂₀ᵤ, A₋ᵤ) - ℒ.rmul!(A₋ᵤ, -1) - end - - A = vcat(A₋ᵤ, sol_compact)[T.reorder,:] - - # 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 - - ℒ.mul!(∇₀, ∇₁[:,1:T.nFuture_not_past_and_mixed], M, 1, 1) - - C = ℒ.lu!(∇₀, check = false) - - if !ℒ.issuccess(C) - if opts.verbose println("Factorisation of ∇₀ failed") end - return zeros(R, T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false - end - - ℒ.ldiv!(C, ∇ₑ) - ℒ.rmul!(∇ₑ, -1) - - # end # timeit_debug - # end # timeit_debug - - return hcat(A, ∇ₑ), sol, true -end - - -function 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())::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {R <: Real, S <: Real} - if !(eltype(workspaces.second_order.Ŝ) == 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 "Calculate second order solution" 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ₑ)]# |> sparse - # droptol!(𝐒₁₋╱𝟏ₑ,tol) - 𝐒₁₋╱𝟏ₑ = choose_matrix_format(𝐒₁₋╱𝟏ₑ, density_threshold = 1.0) - - ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ = @views [(𝐒₁ * 𝐒₁₋╱𝟏ₑ)[i₊,:] - 𝐒₁ - ℒ.I(nₑ₋)[[range(1,n₋)...,n₋ + 1 .+ range(1,nₑ)...],:]] #|> sparse - # droptol!(⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋,tol) - - 𝐒₁₊╱𝟎 = @views [𝐒₁[i₊,:] - zeros(n₋ + n + nₑ, nₑ₋)]# |> sparse - # droptol!(⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋,tol) - - ∇₁₊𝐒₁➕∇₁₀ = @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 - end - - # spinv = inv(∇₁₊𝐒₁➕∇₁₀) - # 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 = ∇₁₊𝐒₁➕∇₁₀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₂.𝐂₂) - - C = ∇₁₊𝐒₁➕∇₁₀lu \ ∇₂⎸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 "Refine sylvester equation" begin - - # # if !solved && !(sylvester_algorithm == :doubling) - # # 𝐒₂, solved = solve_sylvester_equation(A, B, C, - # # # init = 𝐒₂, - # # # sylvester_algorithm = :gmres, - # # initial_guess = initial_guess, - # # sylvester_algorithm = :doubling, - # # verbose = verbose, - # # # tol = tol, - # # timer = timer) - # # end - - # # end # timeit_debug - # @timeit_debug timer "Post-process" begin - - # 𝐒₂ *= M₂.𝐔₂ - - 𝐒₂ = choose_matrix_format(𝐒₂, multithreaded = false) - - # end # timeit_debug - # end # timeit_debug - - return 𝐒₂, solved -end - - -function 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{R} = zeros(0,0), - opts::CalculationOptions = merge_calculation_options())::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {S <: Real,R <: Real} - 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 - # @timeit_debug timer "Calculate third order solution" 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, tol = opts.tol.droptol) - - ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ = @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, tol = opts.tol.droptol) - - ∇₁₊𝐒₁➕∇₁₀ = @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(), NoTangent(), NoTangent() - end - - # spinv = inv(∇₁₊𝐒₁➕∇₁₀) - # spinv = choose_matrix_format(spinv) - - # end # timeit_debug - - ∇₁₊ = @views ∇₁[:,1:n₊] * ℒ.I(n)[i₊,:] - - 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) - - # 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) - - # 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₃.𝐒𝐏 * ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ - # aux = choose_matrix_format(aux, density_threshold = 1.0, min_length = 10) - - # end # timeit_debug - # @timeit_debug timer "∇₃" begin - - 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 - - # tmpkron = ℒ.kron(⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋, ℒ.kron(𝐒₁₊╱𝟎, 𝐒₁₊╱𝟎) * M₂.𝛔) - - 𝐔∇₃ = ∇₃ * 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)] - - # @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 - out2 += mat_mult_kron(∇₂, ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋, collect(𝐒₂₊╱𝟎 * M₂.𝛔), sparse = true, sparse_preallocation = ℂ.tmp_sparse_prealloc3)# |> findnz - - # end # timeit_debug - # @timeit_debug timer "Step 5" begin - # out2 += ∇₁₊ * mat_mult_kron(𝐒₂, collect(𝐒₁₋╱𝟏ₑ), collect(𝐒₂₋╱𝟎)) - # out2 += mat_mult_kron(∇₁₊ * 𝐒₂, collect(𝐒₁₋╱𝟏ₑ), collect(𝐒₂₋╱𝟎)) - # out2 += ∇₁₊ * 𝐒₂ * ℒ.kron(𝐒₁₋╱𝟏ₑ, 𝐒₂₋╱𝟎) - - 𝐒₁₋╱𝟏ₑ = choose_matrix_format(𝐒₁₋╱𝟏ₑ, density_threshold = 0.0, tol = opts.tol.droptol) - out2 += ∇₁₊ * mat_mult_kron(𝐒₂, 𝐒₁₋╱𝟏ₑ, 𝐒₂₋╱𝟎, sparse = true, sparse_preallocation = ℂ.tmp_sparse_prealloc4) - - # 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₃.𝐏 - - 𝐗₃ *= M₃.𝐂₃ - - # 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) - - # end # timeit_debug - # @timeit_debug timer "Mult 2" begin - - C = ∇₁₊𝐒₁➕∇₁₀lu \ 𝐗₃# * 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 "Refine sylvester equation" begin - - # if !solved - # 𝐒₃, solved = solve_sylvester_equation(A, B, C, - # sylvester_algorithm = :doubling, - # verbose = verbose, - # initial_guess = initial_guess, - # # tol = tol, - # timer = timer) - # end - - # if !solved - # return 𝐒₃, solved - # end - - # # end # timeit_debug - # @timeit_debug timer "Post-process" begin - - # 𝐒₃ *= M₃.𝐔₃ - - 𝐒₃ = choose_matrix_format(𝐒₃, multithreaded = false, tol = opts.tol.droptol) - - # end # timeit_debug - # end # timeit_debug - - return 𝐒₃, solved -end - -end # dispatch_doctor - diff --git a/src/perturbation/derivatives.jl b/src/perturbation/derivatives.jl new file mode 100644 index 000000000..e9e941455 --- /dev/null +++ b/src/perturbation/derivatives.jl @@ -0,0 +1,120 @@ +@stable default_mode = "disable" begin + + +function calculate_jacobian(parameters::Vector{M}, + SS_and_pars::Vector{N}, + caches_obj::caches, + 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) + jac_buffer.nzval .= 0 + else + jac_buffer = zeros(M, size(caches_obj.jacobian)) + end + else + jac_buffer = caches_obj.jacobian + 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 + +function calculate_hessian(parameters::Vector{M}, + SS_and_pars::Vector{N}, + caches_obj::caches, + hessian_funcs::hessian_functions, + workspaces::workspaces; + caching::Bool = true)::SparseMatrixCSC{M, Int} where {M,N} + # Always make sure the higher-order workspace matches the eltype expected by + # downstream consumers (e.g. rrules that grab buffers from it). A previous + # call with a different eltype (e.g. ForwardDiff.Dual) may have replaced the + # workspace; the cache short-circuit below would otherwise leave it stale. + S = promote_type(M, N) + if eltype(workspaces.second_order.Ŝ) != S + workspaces.second_order = Higher_order_workspace(S) + end + + # 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 + + if eltype(caches_obj.hessian) != M + if caches_obj.hessian isa SparseMatrixCSC + hes_buffer = similar(caches_obj.hessian,M) + hes_buffer.nzval .= 0 + else + hes_buffer = zeros(M, size(caches_obj.hessian)) + end + else + hes_buffer = caches_obj.hessian + 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 + + +function calculate_third_order_derivatives(parameters::Vector{M}, + SS_and_pars::Vector{N}, + caches_obj::caches, + third_order_derivatives_funcs::third_order_derivatives_functions, + workspaces::workspaces; + caching::Bool = true)::SparseMatrixCSC{M, Int} where {M,N} + # Always make sure the third-order workspace matches the eltype expected by + # downstream consumers (e.g. rrules that grab buffers from it). A previous + # call with a different eltype (e.g. ForwardDiff.Dual) may have replaced the + # workspace; the cache short-circuit below would otherwise leave it stale. + S = promote_type(M, N) + if eltype(workspaces.third_order.Ŝ) != S + workspaces.third_order = Higher_order_workspace(S) + end + + # 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 + + 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) + third_buffer.nzval .= 0 + else + third_buffer = zeros(M, size(caches_obj.third_order_derivatives)) + end + else + third_buffer = caches_obj.third_order_derivatives + 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 + + +end # @stable diff --git a/src/perturbation/solution.jl b/src/perturbation/solution.jl new file mode 100644 index 000000000..2c7c23a16 --- /dev/null +++ b/src/perturbation/solution.jl @@ -0,0 +1,2448 @@ +@stable default_mode = "disable" begin + + + +function 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)::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 + + T = constants.post_model_macro + ensure_first_order_constants!(constants) + idx_constants = constants.post_complete_parameters + + 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 + + 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] + # Old way (≤v0.1.42): + # Q = qr(∇₀[:, present_only_idx]) + # A₊ = Q' * ∇₊; A₀ = Q' * ∇₀; A₋ = Q' * ∇₋ + # Current code reuses QR/ORM workspaces to avoid allocations. + qr_factors, qr_ws = ensure_first_order_fast_qr_workspace!(qme_ws, ∇₀_present) + Q = factorize_qr!((use_fastlapack_qr ? Val(:FastLapack) : Val(:Julia)), ∇₀_present, qr_factors, qr_ws) # Q = qr(∇₀_present) + + qme_ws.fast_qr_orm_ws_plus, qme_ws.fast_qr_orm_dims_plus = apply_qr_transpose_left!(A₊, ∇₊, Q, # 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, # 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, # 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 + + Ã₊ = qme_ws.𝐀̃₊ + ℒ.mul!(Ã₊, @view(A₊[dynIndex,:]), @view(Ir[future_not_past_and_mixed_in_comb,:])) # Ã₊ = A₊[dynIndex,:] * Ir + + Ã₀ = qme_ws.𝐀̃₀ + copyto!(Ã₀, @view(A₀[dynIndex, comb])) + + Ã₋ = qme_ws.𝐀̃₋ + ℒ.mul!(Ã₋, @view(A₋[dynIndex,:]), @view(Ir[past_not_future_and_mixed_in_comb,:])) # Ã₋ = A₋[dynIndex,:] * Ir + + # end # timeit_debug + # @timeit_debug timer "Quadratic matrix equation solve" begin + + sol, solved = solve_quadratic_matrix_equation(Ã₊, Ã₀, Ã₋, constants, workspaces, cache; + initial_guess = initial_guess, + quadratic_matrix_equation_algorithm = opts.quadratic_matrix_equation_algorithm, + use_fastlapack_lu = use_fastlapack_lu, + tol = opts.tol.first_order.qme, + verbose = opts.verbose, + caching = caching) + + if !solved + if opts.verbose println("Quadratic matrix equation solution failed.") end + return fill(R(NaN), T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false + end + + # Detect unit roots from QME solution eigenvalues when the Schur QME path + # did not already set the flag (e.g. doubling solver was used). + if caching && !cache.has_unit_roots + detect_unit_roots_from_solution!(cache, sol) + 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] + + 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] + + Ā₀ᵤ = qme_ws.𝐀̄₀ᵤ + copyto!(Ā₀ᵤ, @view(A₀[1:T.nPresent_only, T.present_only_idx])) + + A₊ᵤ = qme_ws.𝐀₊ᵤ + copyto!(A₊ᵤ, @view(A₊[1:T.nPresent_only,:])) + + Ã₀ᵤ = qme_ws.𝐀̃₀ᵤ + copyto!(Ã₀ᵤ, @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 Ā₀ᵤ" begin + + qme_ws.fast_lu_ws_a0u, qme_ws.fast_lu_dims_a0u, solved_Ā₀ᵤ, Ā̂₀ᵤ = factorize_lu!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), Ā₀ᵤ, + qme_ws.fast_lu_ws_a0u, + qme_ws.fast_lu_dims_a0u) + + if !solved_Ā₀ᵤ + if opts.verbose println("Factorisation of Ā₀ᵤ failed") end + return fill(R(NaN), T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false + end + + # Old way (≤v0.1.42): A₋ᵤ = -(Ā₀ᵤ \ (A₊ᵤ * D * L + Ã₀ᵤ * sol + A₋ᵤ)) + if T.nPresent_only > 0 + ℒ.mul!(A₋ᵤ, Ã₀ᵤ, @view(sol[:,past_not_future_and_mixed_in_comb]), 1, 1) # A₋ᵤ = A₋ᵤ + Ã₀ᵤ * sol + nₚ₋ = qme_ws.𝐧ₚ₋ + ℒ.mul!(nₚ₋, A₊ᵤ, 𝐃) # nₚ₋ = A₊ᵤ * D + ℒ.mul!(A₋ᵤ, nₚ₋, L, 1, 1) # A₋ᵤ = A₋ᵤ + nₚ₋ * L + solve_lu_left!(Ā₀ᵤ, A₋ᵤ, qme_ws.fast_lu_ws_a0u, Ā̂₀ᵤ; # A₋ᵤ = Ā₀ᵤ \ A₋ᵤ + use_fastlapack_lu = use_fastlapack_lu) + ℒ.rmul!(A₋ᵤ, -1) # A₋ᵤ = -A₋ᵤ + end + + A = qme_ws.𝐀 + # Old way (≤v0.1.42): A = vcat(A₋ᵤ, sol_compact)[reorder, :] + # Expanded loop below writes into preallocated A without temporary concatenation. + n_cols = size(A, 2) + + 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 = qme_ws.𝐌 + # Old way (≤v0.1.42): + # M = A[future_idx, :] * expand_past + # ∇₀ = ∇₊ * M + ∇₀ + ℒ.mul!(M, @view(A[T.future_not_past_and_mixed_idx,:]), idx_constants.expand_past) # M = A[future_idx,:] * expand_past + + ℒ.mul!(∇₀, @view(∇₁[:,1:T.nFuture_not_past_and_mixed]), M, 1, 1) # ∇₀ = ∇₊ * M + ∇₀ + + # Old way (≤v0.1.42): C = lu(∇₀) + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_∇₀, C = factorize_lu!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), ∇₀, + qme_ws.fast_lu_ws_nabla0, + qme_ws.fast_lu_dims_nabla0) + + if !solved_∇₀ + if opts.verbose println("Factorisation of ∇₀ failed") end + return fill(R(NaN), T.nVars, T.nPast_not_future_and_mixed + T.nExo), sol, false + end + + # Old way (≤v0.1.42): ∇ₑ = -(∇₀ \ ∇ₑ) + solve_lu_left!(∇₀, ∇ₑ, qme_ws.fast_lu_ws_nabla0, C; # ∇ₑ = ∇₀ \ ∇ₑ + use_fastlapack_lu = use_fastlapack_lu) + ℒ.rmul!(∇ₑ, -1) + + # end # timeit_debug + # end # timeit_debug + + 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 = Float64.(primal.(parameter_values)) + end + + return S₁, sol, true +end + + +@unstable function 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)::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {R <: Real, S <: Real} + # Always make sure the higher-order workspace matches the current eltype + # before any cache short-circuit, so downstream consumers (e.g. rrules that + # grab buffers from the workspace) never see a stale eltype after a previous + # call with a different eltype (e.g. ForwardDiff.Dual). + if !(eltype(workspaces.second_order.Ŝ) == S) + workspaces.second_order = Higher_order_workspace(S) + end + # 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 + ℂ = workspaces.second_order + M₂ = constants.second_order + T = constants.post_model_macro + # @timeit_debug timer "Calculate second order solution" 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ₑ₋) + + 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 + 𝐒₁ = ℂ.𝐒₁::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) + # droptol!(𝐒₁₋╱𝟏ₑ,tol) + 𝐒₁₋╱𝟏ₑ = choose_matrix_format(𝐒₁₋╱𝟏ₑ, density_threshold = 1.0) + + ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ = @views [(𝐒₁ * 𝐒₁₋╱𝟏ₑ)[i₊,:] + 𝐒₁ + ℒ.I(nₑ₋)[[range(1,n₋)...,n₋ + 1 .+ range(1,nₑ)...],:]] #|> sparse + # droptol!(⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋,tol) + + 𝐒₁₊╱𝟎 = @views [𝐒₁[i₊,:] + zeros(n₋ + n + nₑ, nₑ₋)]# |> sparse + # droptol!(⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋,tol) + + ∇₁₊𝐒₁➕∇₁₀ = collect(@views -∇₁[:,1:n₊] * 𝐒₁[i₊,1:n₋] * M₂.𝐈ₙ₋ - ∇₁[:,range(1,n) .+ n₊]) + + # end # timeit_debug + + # @timeit_debug timer "Invert matrix" begin + + qme_ws = workspaces.first_order + + lu_handle = nothing + ∇₁₊𝐒₁➕∇₁₀lu = nothing + + if S === Float64 + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_∇lu, lu_handle = + factorize_lu!(Val(:FastLapack), ∇₁₊𝐒₁➕∇₁₀, qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0) + + if !solved_∇lu + if opts.verbose println("Second order solution: inversion failed") end + return ∇₁₊𝐒₁➕∇₁₀, false + end + else + ∇₁₊𝐒₁➕∇₁₀lu = ℒ.lu(∇₁₊𝐒₁➕∇₁₀, check = false) + + if !ℒ.issuccess(∇₁₊𝐒₁➕∇₁₀lu) + if opts.verbose println("Second order solution: inversion failed") end + return ∇₁₊𝐒₁➕∇₁₀, false + end + end + + # spinv = inv(∇₁₊𝐒₁➕∇₁₀) + # 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₂.𝐈ₙ₊ + + if S === Float64 + A = ∇₁₊ + solve_lu_left!(∇₁₊𝐒₁➕∇₁₀, A, qme_ws.fast_lu_ws_nabla0, lu_handle) # A = ∇₁₊𝐒₁➕∇₁₀ \ ∇₁₊ + else + A = ∇₁₊𝐒₁➕∇₁₀lu \ ∇₁₊ + end + + # end # timeit_debug + # @timeit_debug timer "C" begin + + # 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 + + if S === Float64 + C = collect(∇₂⎸k⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋➕𝛔k𝐒₁₊╱𝟎⎹) + solve_lu_left!(∇₁₊𝐒₁➕∇₁₀, C, qme_ws.fast_lu_ws_nabla0, lu_handle) + else + C = ∇₁₊𝐒₁➕∇₁₀lu \ ∇₂⎸k⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋➕𝛔k𝐒₁₊╱𝟎⎹ + end + + # 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_sylv, + sylvester_algorithm = opts.sylvester_algorithm², + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.second_order.sylvester, + verbose = opts.verbose) + + # end # timeit_debug + # # @timeit_debug timer "Refine sylvester equation" begin + + # # if !solved && !(sylvester_algorithm == :doubling) + # # 𝐒₂, solved = solve_sylvester_equation(A, B, C, + # # # init = 𝐒₂, + # # # sylvester_algorithm = :gmres, + # # initial_guess = initial_guess, + # # sylvester_algorithm = :doubling, + # # verbose = verbose, + # # # tol = tol, + # # timer = timer) + # # end + + # # end # timeit_debug + # @timeit_debug timer "Post-process" begin + + # 𝐒₂ *= M₂.𝐔₂ + + 𝐒₂ = choose_matrix_format(𝐒₂, multithreaded = false) + + # end # timeit_debug + # end # timeit_debug + + if solved && caching && S === Float64 + 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 = Float64.(primal.(parameter_values)) + cache.valid_for.pruned_second_order_solution = Float64[] + end + end + + return 𝐒₂, solved +end + + +@unstable function calculate_third_order_solution(∇₁::AbstractMatrix{S}, #first order derivatives + ∇₂::SparseMatrixCSC{S}, #second order derivatives + ∇₃::SparseMatrixCSC{S}, #third order derivatives + 𝑺₁::AbstractMatrix{S}, #first order solution + 𝐒₂::AbstractMatrix{S}, #second order solution (compressed) + 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)::Union{Tuple{Matrix{S}, Bool}, Tuple{SparseMatrixCSC{S, Int}, Bool}} where {S <: Real,R <: Real} + # Always make sure the higher-order workspace matches the current eltype + # before any cache short-circuit, so downstream consumers (e.g. rrules that + # grab buffers from the workspace) never see a stale eltype after a previous + # call with a different eltype (e.g. ForwardDiff.Dual). + if !(eltype(workspaces.third_order.Ŝ) == S) + workspaces.third_order = Higher_order_workspace(S) + end + # 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 + ℂ = workspaces.third_order + M₂ = constants.second_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 + 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 + + # @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])) + + 𝐒₁₋╱𝟏ₑ = ℂ.𝐒₁₋╱𝟏ₑ::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ₑ)...],:]] #|> sparse + + 𝐒₁₊╱𝟎 = @views [𝐒₁[i₊,:] + zeros(n₋ + n + nₑ, nₑ₋)]# |> sparse + 𝐒₁₊╱𝟎 = choose_matrix_format(𝐒₁₊╱𝟎, density_threshold = 1.0, min_length = 10, tol = opts.tol.third_order.droptol) + + ∇₁₊𝐒₁➕∇₁₀ = collect(@views -∇₁[:,1:n₊] * 𝐒₁[i₊,1:n₋] * ℒ.I(n)[i₋,:] - ∇₁[:,range(1,n) .+ n₊]) + + # end # timeit_debug + # @timeit_debug timer "Invert matrix" begin + + qme_ws = workspaces.first_order + + lu_handle = nothing + ∇₁₊𝐒₁➕∇₁₀lu = nothing + + if S === Float64 + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_∇lu, lu_handle = + factorize_lu!(Val(:FastLapack), ∇₁₊𝐒₁➕∇₁₀, qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0) + + if !solved_∇lu + if opts.verbose println("Second order solution: inversion failed") end + return (∇₁₊𝐒₁➕∇₁₀, false) + end + else + ∇₁₊𝐒₁➕∇₁₀lu = ℒ.lu(∇₁₊𝐒₁➕∇₁₀, check = false) + + if !ℒ.issuccess(∇₁₊𝐒₁➕∇₁₀lu) + if opts.verbose println("Second order solution: inversion failed") end + return (∇₁₊𝐒₁➕∇₁₀, false) + end + end + + # spinv = inv(∇₁₊𝐒₁➕∇₁₀) + # spinv = choose_matrix_format(spinv) + + # end # timeit_debug + + ∇₁₊ = @views ∇₁[:,1:n₊] * M₂.𝐈ₙ₊ + + if S === Float64 + A = copy(∇₁₊) # solve in-place into a buffer; ∇₁₊ is reused later + solve_lu_left!(∇₁₊𝐒₁➕∇₁₀, A, qme_ws.fast_lu_ws_nabla0, lu_handle) + else + A = ∇₁₊𝐒₁➕∇₁₀lu \ ∇₁₊ + end + + # @timeit_debug timer "Setup B" begin + # @timeit_debug timer "Add tmpkron" begin + + kron𝐒₁₋╱𝟏ₑ = ℒ.kron(𝐒₁₋╱𝟏ₑ, 𝐒₁₋╱𝟏ₑ) + + # 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.third_order.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.third_order.droptol) + + 𝐒₂₊╱𝟎 = @views [𝐒₂[i₊,:] + zeros(n₋ + n + nₑ, nₑ₋^2)]; + + aux = M₃.𝐒𝐏 * ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋ + # aux = choose_matrix_format(aux, density_threshold = 1.0, min_length = 10) + + # end # timeit_debug + # @timeit_debug timer "∇₃" begin + + # 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 + + # # tmpkron = ℒ.kron(⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋, ℒ.kron(𝐒₁₊╱𝟎, 𝐒₁₊╱𝟎) * M₂.𝛔) + + # 𝐔∇₃ = ∇₃ * 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.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 * 𝐏₁ᵣ] * 𝐏𝐂₃ + # 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₃.𝐏𝐂₃ + + 𝐗₃ = 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(⎸𝐒₁..⎹, 𝐒₂₊╱𝟎*𝛔) * 𝐏𝐂₃ + 𝐗₃ += mat_mult_kron(∇₂, ⎸𝐒₁𝐒₁₋╱𝟏ₑ⎹╱𝐒₁╱𝟏ₑ₋, collect(𝐒₂₊╱𝟎 * M₂.𝛔), M₃.𝐏𝐂₃, sparse = true, sparse_preallocation = ℂ.tmp_sparse_prealloc4) + + # Term (e): ∇₁₊ * 𝐒₂ * kron(𝐒₁₋╱𝟏ₑ, 𝐒₂₋╱𝟎) * 𝐏𝐂₃ + 𝐒₁₋╱𝟏ₑ = choose_matrix_format(𝐒₁₋╱𝟏ₑ, density_threshold = 0.0, tol = opts.tol.third_order.droptol) + + 𝐗₃ += mat_mult_kron(∇₁₊ * 𝐒₂, 𝐒₁₋╱𝟏ₑ, 𝐒₂₋╱𝟎, M₃.𝐏𝐂₃, sparse = true) + + 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 + 𝐗₃ += 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 + + if S === Float64 + C = collect(𝐗₃) + solve_lu_left!(∇₁₊𝐒₁➕∇₁₀, C, qme_ws.fast_lu_ws_nabla0, lu_handle) + else + C = ∇₁₊𝐒₁➕∇₁₀lu \ 𝐗₃# * M₃.𝐂₃ + end + + # 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_sylv, + sylvester_algorithm = opts.sylvester_algorithm³, + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.third_order.sylvester, + 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, + # verbose = verbose, + # initial_guess = initial_guess, + # # tol = tol, + # timer = timer) + # end + + # if !solved + # return 𝐒₃, solved + # end + + # # end # timeit_debug + # @timeit_debug timer "Post-process" begin + + # 𝐒₃ *= M₃.𝐔₃ + + 𝐒₃ = choose_matrix_format(𝐒₃, multithreaded = false, tol = opts.tol.third_order.droptol) + + # end # timeit_debug + # end # timeit_debug + + if solved && caching && S === Float64 + 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 = Float64.(primal.(parameter_values)) + cache.valid_for.pruned_third_order_solution = Float64[] + end + end + + return 𝐒₃, solved +end + + + +# ── Compressed Kronecker & matrix utilities (moved from MacroModelling.jl) ── + +# Extract unique nonzero row indices, column indices, and nnz count from a dense +# matrix without allocating a sparse copy. Returns sorted unique indices. +function dense_nz_structure(â::AbstractMatrix{T}) where T + nrows, ncols = size(â) + row_has_nz = falses(nrows) + col_has_nz = falses(ncols) + lennz = 0 + @inbounds for j in 1:ncols + for i in 1:nrows + if !iszero(â[i, j]) + lennz += 1 + row_has_nz[i] = true + col_has_nz[j] = true + end + end + end + ui = findall(row_has_nz) + uj = findall(col_has_nz) + return ui, uj, lennz +end + +function create_second_order_auxiliary_matrices(constants::constants) + T = constants.post_model_macro + + + # Indices and number of variables + n₋ = T.nPast_not_future_and_mixed + n = T.nVars + nₑ = T.nExo + + # setup compression matrices for hessian matrix + nₑ₋ = T.nPast_not_future_and_mixed + T.nVars + T.nFuture_not_past_and_mixed + T.nExo + 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) + + # set up vector to capture volatility effect + nₑ₋ = n₋ + 1 + nₑ + 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 + + + +function add_sparse_entries!(P, perm) + n = size(P, 1) + for i in 1:n + P[perm[i], i] += 1.0 + end +end + + +function create_third_order_auxiliary_matrices(constants::constants, ∇₃_col_indices::Vector{Int}) + T = constants.post_model_macro + + + # Indices and number of variables + n₋ = T.nPast_not_future_and_mixed + n₊ = T.nFuture_not_past_and_mixed + n = T.nVars + nₑ = T.nExo + + n̄ = n₋ + n + n₊ + nₑ + + # compression matrices for third order derivatives matrix + nₑ₋ = T.nPast_not_future_and_mixed + T.nVars + T.nFuture_not_past_and_mixed + T.nExo + colls3 = [nₑ₋^2 * (i-1) + nₑ₋ * (k-1) + l for i in 1:nₑ₋ for k in 1:i for l in 1:k] + 𝐂∇₃ = sparse(colls3, 1:length(colls3) , 1.0) + + idxs = Int[] + for k in 1:nₑ₋ + for j in 1:nₑ₋ + for i in 1:nₑ₋ + sorted_ids = sort([k,j,i]) + push!(idxs, (sorted_ids[3] - 1) * nₑ₋ ^ 2 + (sorted_ids[2] - 1) * nₑ₋ + sorted_ids[1]) + end + end + end + + 𝐔∇₃ = 𝐂∇₃' * sparse(idxs,1:nₑ₋ ^ 3, 1) + + # compression matrices for third order transition matrix + nₑ₋ = n₋ + 1 + nₑ + colls3 = [nₑ₋^2 * (i-1) + nₑ₋ * (k-1) + l for i in 1:nₑ₋ for k in 1:i for l in 1:k] + 𝐂₃ = sparse(colls3, 1:length(colls3) , 1.0) + + idxs = Int[] + for k in 1:nₑ₋ + for j in 1:nₑ₋ + for i in 1:nₑ₋ + sorted_ids = sort([k,j,i]) + push!(idxs, (sorted_ids[3] - 1) * nₑ₋ ^ 2 + (sorted_ids[2] - 1) * nₑ₋ + sorted_ids[1]) + end + end + end + + 𝐔₃ = 𝐂₃' * sparse(idxs,1:nₑ₋ ^ 3, 1) + + # Precompute 𝐈₃ + 𝐈₃ = Dict{Vector{Int}, Int}() + idx = 1 + for i in 1:nₑ₋ + for k in 1:i + for l in 1:k + 𝐈₃[[i,k,l]] = idx + idx += 1 + end + end + end + + # permutation matrices + M = reshape(1:nₑ₋^3,1,nₑ₋,nₑ₋,nₑ₋) + + 𝐏 = spzeros(nₑ₋^3, nₑ₋^3) # Preallocate the sparse matrix + + # Create the permutations directly + add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 4, 2, 3))) + add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 2, 4, 3))) + add_sparse_entries!(𝐏, PermutedDimsArray(M, (1, 2, 3, 4))) + + # 𝐏 = @views sparse(reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 4, 2, 3])],nₑ₋^3,nₑ₋^3) + # + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 4, 3])],nₑ₋^3,nₑ₋^3) + # + reshape(spdiagm(ones(nₑ₋^3))[:,PermutedDimsArray(M,[1, 2, 3, 4])],nₑ₋^3,nₑ₋^3)) + + 𝐏₁ₗ = sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(2,1,3))),:]) + 𝐏₁ᵣ = sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(2,1,3)))]) + + 𝐏₁ₗ̂ = @views sparse(spdiagm(ones(n̄^3))[vec(permutedims(reshape(1:n̄^3,n̄,n̄,n̄),(1,3,2))),:]) + 𝐏₂ₗ̂ = @views sparse(spdiagm(ones(n̄^3))[vec(permutedims(reshape(1:n̄^3,n̄,n̄,n̄),(3,1,2))),:]) + + 𝐏₁ₗ̄ = @views sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(1,3,2))),:]) + 𝐏₂ₗ̄ = @views sparse(spdiagm(ones(nₑ₋^3))[vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(3,1,2))),:]) + + + 𝐏₁ᵣ̃ = @views sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(1,3,2)))]) + 𝐏₂ᵣ̃ = @views sparse(spdiagm(ones(nₑ₋^3))[:,vec(permutedims(reshape(1:nₑ₋^3,nₑ₋,nₑ₋,nₑ₋),(3,1,2)))]) + + ∇₃_col_indices_extended = Set{Int}(findnz(sparse(ones(Int,length(∇₃_col_indices)),∇₃_col_indices,ones(Int,length(∇₃_col_indices)),1,size(𝐔∇₃,1)) * 𝐔∇₃)[2]) + + nonnull_columns = Set{Int}() + for i in 1:n̄ + for j in i:n̄ + for k in j:n̄ + if (n̄^2 * (i - 1) + n̄ * (j - 1) + k) ∈ ∇₃_col_indices_extended + push!(nonnull_columns,i) + push!(nonnull_columns,j) + push!(nonnull_columns,k) + end + end + end + end + + 𝐒𝐏 = sparse(collect(nonnull_columns), collect(nonnull_columns), 1, n̄, n̄) + + to = constants.third_order + to.𝐂₃ = 𝐂₃ + to.𝐔₃ = 𝐔₃ + to.𝐈₃ = 𝐈₃ + to.𝐂∇₃ = 𝐂∇₃ + to.𝐔∇₃ = 𝐔∇₃ + to.∇₃_rowmask = sort!(unique(∇₃_col_indices)) + to.𝐏 = 𝐏 + to.𝐏𝐂₃ = 𝐏 * 𝐂₃ + to.𝐏₁ₗ = 𝐏₁ₗ + to.𝐏₁ᵣ = 𝐏₁ᵣ + to.𝐏₁ₗ̂ = 𝐏₁ₗ̂ + to.𝐏₂ₗ̂ = 𝐏₂ₗ̂ + to.𝐏₁ₗ̄ = 𝐏₁ₗ̄ + to.𝐏₂ₗ̄ = 𝐏₂ₗ̄ + 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 + +@unstable function mat_mult_kron(A::AbstractSparseMatrix{R}, + B::AbstractMatrix{T}, + C::AbstractMatrix{T}, + 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) + + 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 + + Ā = zeros(T, n_rowC, n_rowB) + ĀB = zeros(T, n_rowC, n_colB) + CĀB = zeros(T, n_colC, n_colB) + vCĀB = zeros(T, n_colB * n_colC) + vCĀBD = zeros(T, size(D,2)) + + # Linked-list row index: O(nnz_in_row) per row instead of O(nnz) for A[row,:] + A_csc = A isa SparseMatrixCSC ? A : A.A + A_rv = SparseArrays.rowvals(A_csc) + A_nzv = nonzeros(A_csc) + A_cp = SparseArrays.getcolptr(A_csc) + nnzA_ll = nnz(A_csc) + n_cols_A = size(A_csc, 2) + row_head = zeros(Int, size(A_csc, 1)) + row_next = zeros(Int, nnzA_ll) + nz_col = Vector{Int}(undef, nnzA_ll) + @inbounds for col in n_cols_A:-1:1 + for idx in A_cp[col]:(A_cp[col + 1] - 1) + r = A_rv[idx] + row_next[idx] = row_head[r] + row_head[r] = idx + nz_col[idx] = col + end + end + + α = .7 + k = 0 + + @inbounds for row in eachindex(row_head) + row_head[row] == 0 && continue + fill!(Ā, zero(T)) + p = row_head[row] + while p != 0 + Ā[nz_col[p]] = T(A_nzv[p]) + p = row_next[p] + end + ℒ.mul!(ĀB, Ā, B) + ℒ.mul!(CĀB, C', ĀB) + copyto!(vCĀB, CĀB) + ℒ.mul!(vCĀBD, D', vCĀB) + + if sparse + for (i,v) in enumerate(vCĀ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,:], vCĀBD) + end + end + + if sparse + resize!(I, k) + resize!(J, k) + resize!(V, k) + + klasttouch = sparse_preallocation[4] + csrrowptr = sparse_preallocation[5] + csrcolval = sparse_preallocation[6] + csrnzval = sparse_preallocation[7] + + 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 + + + + +@unstable function mat_mult_kron(A::DenseMatrix{R}, + B::AbstractMatrix{T}, + C::AbstractMatrix{T}, + 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)) + + # vals = T[] + # rows = Int[] + # cols = Int[] + + Ā = zeros(T, n_rowC, n_rowB) + ĀB = zeros(T, n_rowC, n_colB) + CĀB = zeros(T, n_colC, n_colB) + vCĀB = zeros(T, n_colB * n_colC) + # vCĀBD = zeros(size(D,2)) + + # rv = A isa SparseMatrixCSC ? A.rowval : A.A.rowval + + # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique + r = 1 + @inbounds for row in eachrow(A) + @views copyto!(Ā, row) + ℒ.mul!(ĀB, Ā, B) + ℒ.mul!(CĀB, C', ĀB) + copyto!(vCĀB, CĀB) + @views ℒ.mul!(X[row,:], D', vCĀB) + r += 1 + end + + return choose_matrix_format(X) + # ℒ.mul!(vCĀBD, D', vCĀB) + + # for (i,v) in enumerate(vCĀBD) + # if abs(v) > eps() + # push!(rows, row) + # push!(cols, i) + # push!(vals, v) + # end + # end + # end + + # 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 +end + +@unstable function mat_mult_kron(A::AbstractSparseMatrix{R}, + B::AbstractMatrix{T}, + C::AbstractMatrix{T}; + 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} + n_rowB = size(B,1) + n_colB = size(B,2) + + n_rowC = size(C,1) + n_colC = size(C,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()) + + p = nnzA * nnzB * nnzC / (length(A) * length(B) * length(C)) + + if length(sparse_preallocation[1]) == 0 + estimated_nnz = Int(ceil((1-(1-p)^size(A,1))*size(A,1) * n_colB * n_colC)) + + 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), n_colB * n_colC) + end + + Ā = zeros(T, n_rowC, n_rowB) + ĀB = zeros(T, n_rowC, n_colB) + CĀB = zeros(T, n_colC, n_colB) + + # Linked-list row index: O(nnz_in_row) per row instead of O(nnz) for A[row,:] + A_csc = A isa SparseMatrixCSC ? A : A.A + A_rv = SparseArrays.rowvals(A_csc) + A_nzv = nonzeros(A_csc) + A_cp = SparseArrays.getcolptr(A_csc) + nnzA_ll = nnz(A_csc) + n_cols_A = size(A_csc, 2) + row_head = zeros(Int, size(A_csc, 1)) + row_next = zeros(Int, nnzA_ll) + nz_col = Vector{Int}(undef, nnzA_ll) + @inbounds for col in n_cols_A:-1:1 + for idx in A_cp[col]:(A_cp[col + 1] - 1) + r = A_rv[idx] + row_next[idx] = row_head[r] + row_head[r] = idx + nz_col[idx] = col + end + end + + α = .7 # speed of Vector increase + k = 0 + + @inbounds for row in eachindex(row_head) + row_head[row] == 0 && continue + fill!(Ā, zero(T)) + p = row_head[row] + while p != 0 + Ā[nz_col[p]] = T(A_nzv[p]) + p = row_next[p] + end + ℒ.mul!(ĀB, Ā, B) + ℒ.mul!(CĀB, C', ĀB) + + if sparse + for (i,v) in enumerate(CĀB) + if abs(v) > eps() + k += 1 + + if k > estimated_nnz + estimated_nnz += min(size(A,1) * n_colB * n_colC, max(10000, Int(ceil((α - 1) * estimated_nnz + (1 - α) * size(A,1) * n_colB * n_colC)))) + + 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,:], CĀB) + end + end + + if sparse + resize!(I, k) + resize!(J, k) + resize!(V, k) + + klasttouch = sparse_preallocation[4] # Vector{Ti}(undef, n) + csrrowptr = sparse_preallocation[5] # Vector{Ti}(undef, m + 1) + csrcolval = sparse_preallocation[6] # Vector{Ti}(undef, length(I)) + csrnzval = sparse_preallocation[7] # Vector{Tv}(undef, length(I)) + + resize!(klasttouch, n_colB * n_colC) + resize!(csrrowptr, size(A, 1) + 1) + resize!(csrcolval, length(I)) + resize!(csrnzval, length(I)) + + 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) + end + + return out +end + + + + +@unstable function mat_mult_kron(A::DenseMatrix{R}, + B::AbstractMatrix{T}, + C::AbstractMatrix{T}; + 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} + 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), n_colB * n_colC) + + # vals = T[] + # rows = Int[] + # cols = Int[] + + Ā = zeros(T, n_rowC, n_rowB) + ĀB = zeros(T, n_rowC, n_colB) + CĀB = zeros(T, n_colC, n_colB) + + # Polyester.@batch threadlocal = (Vector{T}(), Vector{Int}(), Vector{Int}()) for row in rv |> unique + r = 1 + @inbounds for row in eachrow(A) + @views copyto!(Ā, row) + ℒ.mul!(ĀB, Ā, B) + ℒ.mul!(CĀB, C', ĀB) + + @views copyto!(X[r,:], CĀB) + r += 1 + end + + return choose_matrix_format(X) + # for (i,v) in enumerate(CĀB) + # if abs(v) > eps() + # push!(rows, row) + # push!(cols, i) + # push!(vals, v) + # end + # end + # end + + # if VERSION >= v"1.10" + # return sparse!(rows,cols,vals,size(A,1),n_colB*n_colC) + # else + # return sparse(rows,cols,vals,size(A,1),n_colB*n_colC) + # end +end + +# 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 + + nr = size(A, 1) + nc = size(A, 2) + size(σ) == (nr^2, nc^2) || throw(DimensionMismatch("σ must be $(nr^2)×$(nc^2), got $(size(σ))")) + + # 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) + + 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 = 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 + +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 + â = copy(a') + + rmask = colmask + colmask = rowmask + rowmask = rmask + elseif typeof(a) <: DenseMatrix{T} + â = copy(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(â) + + # 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 + # Extract unique nonzero row/col indices directly from dense matrix + ui, uj, lennz = dense_nz_structure(â) + + 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(â)) ^ 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] + 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 = Threads.Atomic{Int}(0) # Counter for non-zero entries + # k̄ = Threads.Atomic{Int}(0) # effectively slower than the non-threaded version + + k = 0 + + + # @timeit_debug timer "Loop" begin + # Triple nested loops for (i1 ≤ j1 ≤ k1) and (i2 ≤ j2 ≤ k2) + # Polyester.@batch threadlocal=(Vector{Int}(), Vector{Int}(), Vector{T}()) for i1 in ui + # Polyester.@batch minbatch = 10 for i1 in ui + # 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 + + n_ui = length(ui) + n_uj = length(uj) + + for idx_i1 in 1:n_ui + @inbounds i1 = ui[idx_i1] + for idx_j1 in 1:idx_i1 + @inbounds j1 = ui[idx_j1] + for idx_k1 in 1:idx_j1 + @inbounds k1 = ui[idx_k1] + + row = (i1-1) * i1 * (i1+1) ÷ 6 + (j1-1) * j1 ÷ 2 + k1 + + if norowmask || rowmask_lookup[row] + # Divisor depends only on the row triple + if i1 == j1 + divisor = i1 == k1 ? 6 : 2 + else + divisor = (i1 ≠ k1 && j1 ≠ k1) ? 1 : 2 + end + + for idx_i2 in 1:n_uj + @inbounds i2 = uj[idx_i2] + # Hoist i2-dependent reads + @inbounds aii = â[i1, i2] + @inbounds aji = â[j1, i2] + @inbounds aki = â[k1, i2] + + for idx_j2 in 1:idx_i2 + @inbounds j2 = uj[idx_j2] + # Hoist j2-dependent reads + @inbounds aij = â[i1, j2] + @inbounds ajj = â[j1, j2] + @inbounds akj = â[k1, j2] + + # Precompute sub-expressions for the k2 inner loop + p1 = aii * ajj + aij * aji # coefficient of akk + p2 = aii * akj + aij * aki # coefficient of ajk + p3 = aji * akj + ajj * aki # coefficient of aik + col_partial = (i2-1) * i2 * (i2+1) ÷ 6 + (j2-1) * j2 ÷ 2 + + for idx_k2 in 1:idx_j2 + @inbounds k2 = uj[idx_k2] + @inbounds aik = â[i1, k2] + @inbounds ajk = â[j1, k2] + @inbounds akk = â[k1, k2] + + val = akk * p1 + ajk * p2 + aik * p3 + + if abs(val) > tol + col = col_partial + k2 + + if nocolmask || colmask_lookup[col] + k += 1 + + if k > estimated_nnz + estimated_nnz += Int(ceil(max(1000, estimated_nnz * .1))) + estimated_nnz = min(m3_cols * m3_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 + + # end # timeit_debug + + # @timeit_debug timer "Resize" begin + + # out = map(fetch, threadlocal) + + # I = mapreduce(v -> v[1], vcat, out) + # J = mapreduce(v -> v[2], vcat, out) + # V = mapreduce(v -> v[3], vcat, out) + + # # Resize the index and value arrays to the actual number of entries + # resize!(I, k̄[]) + # resize!(J, k̄[]) + # resize!(V, k̄[]) + resize!(I, k) + resize!(J, k) + resize!(V, k) + + # end # timeit_debug + # end # timeit_debug + + # Create the sparse matrix from the collected indices and values + if a_is_adjoint + klasttouch = sparse_preallocation[4] # Vector{Ti}(undef, n) + csrrowptr = sparse_preallocation[5] # Vector{Ti}(undef, m + 1) + csrcolval = sparse_preallocation[6] # Vector{Ti}(undef, length(I)) + csrnzval = sparse_preallocation[7] # Vector{Tv}(undef, length(I)) + + resize!(klasttouch, m3_rows) + resize!(csrrowptr, m3_cols + 1) + resize!(csrcolval, length(J)) + resize!(csrnzval, length(J)) + + out = sparse!(J, I, V, m3_cols, m3_rows, +, klasttouch, csrrowptr, csrcolval, csrnzval, J, I, V) + # out = sparse!(J, I, V, m3_cols, m3_rows) + else + klasttouch = sparse_preallocation[4] # Vector{Ti}(undef, n) + csrrowptr = sparse_preallocation[5] # Vector{Ti}(undef, m + 1) + csrcolval = sparse_preallocation[6] # Vector{Ti}(undef, length(I)) + csrnzval = sparse_preallocation[7] # Vector{Tv}(undef, length(I)) + + resize!(klasttouch, m3_cols) + resize!(csrrowptr, m3_rows + 1) + resize!(csrcolval, length(I)) + resize!(csrnzval, length(I)) + + out = sparse!(I, J, V, m3_rows, m3_cols, +, klasttouch, csrrowptr, csrcolval, csrnzval, I, J, V) + # 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 + + if typeof(a) <: DenseMatrix{T} + â = a + else + â = convert(Matrix, a) + end + + n_rows, n_cols = size(â) + 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 + + size(M, 2) == m3_rows || throw(DimensionMismatch("M must have $m3_rows columns, got $(size(M, 2))")) + + rv_M = SparseArrays.rowvals(M) + nzv_M = nonzeros(M) + + # Extract unique nonzero row/col indices directly from dense matrix + ui, uj, lennz = dense_nz_structure(â) + n_ui = length(ui) + n_uj = length(uj) + + # Row-slice cache: pack â[row, uj[:]] into compact contiguous vectors + # for cache-friendly access in the inner kernel (stride-1 vs stride-nrows) + row_cache_i = Vector{T}(undef, n_uj) + row_cache_j = Vector{T}(undef, n_uj) + row_cache_k = Vector{T}(undef, n_uj) + + # --- sparse IJV buffer management --- + if length(sparse_preallocation[1]) == 0 + estimated_nnz = floor(Int, max(m * m3_cols * (lennz / length(â)) ^ 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; precompute reciprocal + # to replace division with multiplication in the inner loop + if i1 == j1 + divisor = i1 == k1 ? 6 : 2 + else + divisor = (i1 ≠ k1 && j1 ≠ k1) ? 1 : 2 + end + inv_divisor = one(T) / divisor + + # Fill row caches for this (i1, j1, k1) triple — + # sequential stride-1 reads instead of indirect stride-nrows + @inbounds for c in 1:n_uj + col_idx = uj[c] + row_cache_i[c] = â[i1, col_idx] + row_cache_j[c] = â[j1, col_idx] + row_cache_k[c] = â[k1, col_idx] + end + + # Col-inner loop: column triples (i2 ≥ j2 ≥ k2) with bounded ranges + for idx_i2 in 1:n_uj + @inbounds i2 = uj[idx_i2] + # Hoist i2-dependent reads from row cache + @inbounds aii = row_cache_i[idx_i2] + @inbounds aji = row_cache_j[idx_i2] + @inbounds aki = row_cache_k[idx_i2] + + for idx_j2 in 1:idx_i2 + @inbounds j2 = uj[idx_j2] + # Hoist j2-dependent reads from row cache + @inbounds aij = row_cache_i[idx_j2] + @inbounds ajj = row_cache_j[idx_j2] + @inbounds akj = row_cache_k[idx_j2] + + # Precompute sub-expressions for the k2 inner loop + p1 = aii * ajj + aij * aji + p2 = aii * akj + aij * aki + p3 = aji * akj + ajj * aki + col_partial = (i2 - 1) * i2 * (i2 + 1) ÷ 6 + (j2 - 1) * j2 ÷ 2 + + for idx_k2 in 1:idx_j2 + @inbounds k2 = uj[idx_k2] + @inbounds aik = row_cache_i[idx_k2] + @inbounds ajk = row_cache_j[idx_k2] + @inbounds akk = row_cache_k[idx_k2] + + val = akk * p1 + ajk * p2 + aik * p3 + + if abs(val) > tol + scaled_val = val * inv_divisor + col = col_partial + 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') + + rmask = colmask + colmask = rowmask + rowmask = rmask + elseif typeof(a) <: DenseMatrix{T} + â = copy(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(â) + + # 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 + # Extract unique nonzero row/col indices directly from dense matrix + ui, uj, lennz = dense_nz_structure(â) + + 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(â)) ^ 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 + + + 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 + + n_ui = length(ui) + n_uj = length(uj) + + for idx_i1 in 1:n_ui + @inbounds i1 = ui[idx_i1] + for idx_j1 in 1:idx_i1 + @inbounds j1 = ui[idx_j1] + + row = (i1 - 1) * i1 ÷ 2 + j1 + + if norowmask || rowmask_lookup[row] + divisor = i1 == j1 ? 2 : 1 + + for idx_i2 in 1:n_uj + @inbounds i2 = uj[idx_i2] + @inbounds aii = â[i1, i2] + @inbounds aji = â[j1, i2] + + for idx_j2 in 1:idx_i2 + @inbounds j2 = uj[idx_j2] + @inbounds aij = â[i1, j2] + @inbounds ajj = â[j1, j2] + + val = aii * ajj + aij * aji + + if abs(val) > tol + col = (i2 - 1) * i2 ÷ 2 + j2 + + if nocolmask || colmask_lookup[col] + 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 + + 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 + +# Detect unit roots from QME solution without computing eigenvalues. +# If sol has an eigenvalue near 1, (I - sol) is nearly singular. +# Uses LU factorization: exactly singular (info > 0) or smallest absolute pivot < tol. +# Cost: O(n³/3) LU on the small nPfm × nPfm solution matrix. +function detect_unit_roots_from_solution!(cache::caches, sol::AbstractMatrix{R}; + tol::Float64 = 1e-8) where R <: AbstractFloat + n = size(sol, 1) + n == 0 && return nothing + ImA = similar(sol) + @turbo for j in 1:n, i in 1:n + ImA[i, j] = ifelse(i == j, one(R), zero(R)) - sol[i, j] + end + F = ℒ.lu!(ImA; check = false) + if F.info > 0 + cache.has_unit_roots = true + return nothing + end + # Diagonal of packed LU factors = diagonal of U (L has unit diagonal). + # Smallest absolute pivot indicates near-singularity ↔ eigenvalue of sol near 1. + min_abs_pivot = typemax(R) + @inbounds for i in 1:n + v = abs(F.factors[i, i]) + if v < min_abs_pivot + min_abs_pivot = v + end + end + if min_abs_pivot < tol + cache.has_unit_roots = true + end + return nothing +end + + + +end # @stable diff --git a/src/rrules.jl b/src/rrules.jl new file mode 100644 index 000000000..993ade55d --- /dev/null +++ b/src/rrules.jl @@ -0,0 +1,12209 @@ +# 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: 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 + +# clear_solution_caches! is a pure side-effect (cache invalidation) with no +# differentiable outputs, so the pullback is a no-op. +function rrule(::typeof(clear_solution_caches!), 𝓂::ℳ, algorithm::Symbol) + clear_solution_caches!(𝓂, algorithm) + return nothing, _ -> (NoTangent(), NoTangent(), NoTangent()) +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(Ȳ) + Ȳ = unthunk(Ȳ) + 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) + + # Linked-list row index: avoids Dict{Int,Vector{Int}} allocation + n_rows_A = size(A_csc, 1) + row_head = zeros(Int, n_rows_A) + row_next = zeros(Int, nnzA) + @inbounds for col in size(A_csc, 2):-1:1 + for k in A_csc.colptr[col]:(A_csc.colptr[col + 1] - 1) + nz_col[k] = col + r = A_csc.rowval[k] + row_next[k] = row_head[r] + row_head[r] = k + end + end + + ∂A_nz = zeros(G, nnzA) + Abar_vec = zeros(G, size(A_csc, 2)) + + @inbounds for r in 1:n_rows_A + row_head[r] == 0 && continue + + fill!(Abar_vec, zero(G)) + k = row_head[r] + while k != 0 + Abar_vec[nz_col[k]] = A_csc.nzval[k] + k = row_next[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̄) + k = row_head[r] + while k != 0 + ∂A_nz[k] += vecAbar̄[nz_col[k]] + k = row_next[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 + + ℂ = 𝓂.workspaces.second_order + nPast = length(x) + ensure_sss_kron_buffers!(ℂ, nPast; third_order=false) + kron_x_aug_buf = ℂ.kron_x_aug_xx + kron_x_aug_I = ℂ.kron_x_aug_I + + for i in 1:max_iters + copyto!(x_aug, 1, x, 1, nPast) + ℒ.kron!(kron_x_aug_buf, x_aug, x_aug) + + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ∂x = (A + B * kron_x_aug_I - I_nPast) + + Δx = (A * x + B̂ * kron_x_aug_buf / 2 - x) + ensure_dx_lu_buffer!(ℂ, ∂x, Δx) + sol = 𝒮.solve!(ℂ.dx_lu_buffer) + + if sol.retcode != 𝒮.SciMLBase.ReturnCode.Default && !𝒮.SciMLBase.successful_retcode(sol.retcode) + return x, false + end + copyto!(Δx, sol.u) + + if i > 5 && isapprox(A * x + B̂ * kron_x_aug_buf / 2, x, rtol = tol) + break + end + + # x += Δx + ℒ.axpy!(-1, Δx, x) + end + copyto!(x_aug, 1, x, 1, nPast) + # Local kron for closure capture (workspace buffers may be overwritten before pullback runs) + 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 + ∂x₁ = unthunk(∂x[1]) + S = -∂x₁' / (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 + + ℂ = 𝓂.workspaces.third_order + nPast = length(x) + ensure_sss_kron_buffers!(ℂ, nPast; third_order=true) + kron_x_aug_buf = ℂ.kron_x_aug_xx + kron_x_kron_buf = ℂ.kron_x_aug_x_kron + kron_x_aug_I = ℂ.kron_x_aug_I + kron_x_kron_I = ℂ.kron_x_kron_I + + for i in 1:max_iters + copyto!(x_aug, 1, x, 1, nPast) + ℒ.kron!(kron_x_aug_buf, x_aug, x_aug) + ℒ.kron!(kron_x_kron_buf, x_aug, kron_x_aug_buf) + + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ℒ.kron!(kron_x_kron_I, kron_x_aug_buf, I_nPast) + ∂x = (A + B * kron_x_aug_I + C * kron_x_kron_I / 2 - I_nPast) + + Δx = (A * x + B̂ * kron_x_aug_buf / 2 + Ĉ * kron_x_kron_buf / 6 - x) + ensure_dx_lu_buffer!(ℂ, ∂x, Δx) + sol = 𝒮.solve!(ℂ.dx_lu_buffer) + + if sol.retcode != 𝒮.SciMLBase.ReturnCode.Default && !𝒮.SciMLBase.successful_retcode(sol.retcode) + return x, false + end + copyto!(Δx, sol.u) + + if i > 5 && isapprox(A * x + B̂ * kron_x_aug_buf / 2 + Ĉ * kron_x_kron_buf / 6, x, rtol = tol) + break + end + + # x += Δx + ℒ.axpy!(-1, Δx, x) + end + + copyto!(x_aug, 1, x, 1, nPast) + # Local kron for closure capture (workspace buffers may be overwritten before pullback runs) + 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) + ∂x₁ = unthunk(∂x[1]) + S = -∂x₁' / (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 + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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)::Int + rhs_n_cols = size(∂SS_equations_∂parameters, 2)::Int + + 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, 1) != rhs_n_rows || size(qme_ws.nsss_jvp_rhs, 2) != 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, 1) != rhs_n_rows || size(qme_ws.nsss_sparse_lu_buffer.A, 2) != 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 + # Old way (≤v0.1.42): nsss_lu = lu(∂SS/∂SS_and_pars) + qme_ws.fast_lu_ws_nsss, qme_ws.fast_lu_dims_nsss, solved_nsss, nsss_lu = factorize_lu!(Val(:FastLapack), ∂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) + + # JVP = -(∂SS/∂SS_and_pars \ ∂SS/∂parameters) + solve_lu_left!(∂SS_equations_∂SS_and_pars, # rhs ← ∂SS/∂SS_and_pars \ rhs + qme_ws.nsss_jvp_rhs, + qme_ws.fast_lu_ws_nsss, + nsss_lu) + + ℒ.rmul!(qme_ws.nsss_jvp_rhs, -1) # JVP = -JVP + 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 = unthunk(∂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 + + ensure_model_structure_constants!(constants, 𝓂.equations.calibration_parameters) + ms = constants.post_complete_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 = collect(T.I_nPast - 𝐒₁[past_idx, 1:nPast]) + rhs = collect((𝐒₂ * kron_aug1 / 2)[past_idx]) + tmp_for_pullback = copy(tmp) + + ensure_sss_tmp_lu_buffer!(𝓂.workspaces.second_order, tmp, rhs) + tmp_sol = 𝒮.solve!(𝓂.workspaces.second_order.sss_tmp_lu_buffer) + + if tmp_sol.retcode != 𝒮.SciMLBase.ReturnCode.Default && !𝒮.SciMLBase.successful_retcode(tmp_sol.retcode) + 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_sol.u) + tmp_pb_lu_ws, tmp_pb_lu_dims = ensure_sss_pullback_fast_lu_workspace!(𝓂.workspaces.second_order, tmp_for_pullback) + tmp_pb_lu_ws, tmp_pb_lu_dims, solved_tmp_pb_lu, tmp_pb_lu = factorize_lu!(Val(:FastLapack), tmp_for_pullback, tmp_pb_lu_ws, tmp_pb_lu_dims) + 𝓂.workspaces.second_order.fast_lu_ws_sss_pullback = tmp_pb_lu_ws + 𝓂.workspaces.second_order.fast_lu_dims_sss_pullback = tmp_pb_lu_dims + use_fastlapack_tmp_pb = solved_tmp_pb_lu + if !solved_tmp_pb_lu + tmp_pb_lu_ws, tmp_pb_lu_dims, solved_tmp_pb_lu, tmp_pb_lu = + factorize_lu!(Val(:Julia), tmp_for_pullback, tmp_pb_lu_ws, tmp_pb_lu_dims) + @assert solved_tmp_pb_lu "Could not factorize preserved stochastic steady-state pullback matrix." + use_fastlapack_tmp_pb = false + end + ∂rhs_buffer = zeros(Float64, length(SSSstates)) + + 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) + copyto!(∂rhs_buffer, ∂SSSstates) + solve_lu_left_transpose!(tmp_for_pullback, ∂rhs_buffer, tmp_pb_lu_ws, tmp_pb_lu; + use_fastlapack_lu = use_fastlapack_tmp_pb) + ∂tmp = -∂rhs_buffer * SSSstates' + ∂𝐒₁_aug[past_idx, 1:nPast] .-= ∂tmp + ∂𝐒₂_from_rhs = spzeros(Float64, size(𝐒₂)...) + ∂𝐒₂_from_rhs[past_idx, :] += ∂rhs_buffer * 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⁺] + + newton_result, newton_pullback = + rrule(solve_stochastic_steady_state_newton, Val(:second_order), 𝐒₁, 𝐒₂, collect(SSSstates), 𝓂) + SSSstates_final, converged::Bool = newton_result + + 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⁺] + Ĉ = 𝐒₃̂[:,kron_s⁺_s⁺_s⁺] + + newton_result, newton_pullback = + rrule(solve_stochastic_steady_state_newton, Val(:third_order), 𝐒₁, 𝐒₂, 𝐒₃̂, collect(SSSstates), 𝓂) + SSSstates_final, converged::Bool = newton_result + + 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 + Ĉ * 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 += (Ĉ' * ∂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 + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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 + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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 + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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 + + ensure_model_structure_constants!(𝓂.constants, 𝓂.equations.calibration_parameters) + ms = 𝓂.constants.post_complete_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_SELECTOR(𝓂), + 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 + + + +# ── get_irf rrule BPTT helpers ── + +# Backpropagation-through-time (BPTT) pullback: returns +# (∂𝐒_list, ∂state_init, ∂SS_and_pars_from_init) +# These helpers replay the saved IRF simulation backward in time. Each +# algorithm-specific method mirrors the corresponding forward state update while +# accumulating cotangents for the perturbation solution objects, the effective +# IRF initial state, and the steady-state offset induced by a user-supplied +# initial condition. +function irf_bptt(::Val{:first_order}, + ∂Y_all::Array{S,3}, states_store, shocks_store, + nShocks, periods, past_idx, nPast, nVars, nExo, + 𝐒, initial_state, nVar_len) where S + sol_mat = 𝐒 + ∂sol_mat = zeros(S, size(sol_mat)) + ∂state_init = zeros(S, nVars) + ∂SS_from_init = zeros(S, nVar_len) + + for si in 1:nShocks + ∂y_accum = zeros(S, nVars) + + for t in periods:-1:1 + ∂y_t = ∂Y_all[:, t, si] + ∂y_accum + + prev_st = states_store[si, t] + shock_t = shocks_store[si, t] + input_t = [prev_st[past_idx]; shock_t] + + ∂sol_mat .+= ∂y_t * input_t' + ∂input_t = sol_mat' * ∂y_t + + ∂y_accum = zeros(S, nVars) + ∂y_accum[past_idx] .+= ∂input_t[1:nPast] + end + + ∂state_init .+= ∂y_accum + if initial_state != [0.0] + ∂SS_from_init[1:nVar_len] .-= ∂y_accum[1:nVar_len] + end + end + + return [∂sol_mat], ∂state_init, ∂SS_from_init +end + +function irf_bptt(::Val{:pruned_second_order}, + ∂Y_all::Array{S,3}, states_store, shocks_store, + nShocks, periods, past_idx, nPast, nVars, nExo, + 𝐒, initial_state, nVar_len) where S + 𝐒₁, 𝐒₂ = 𝐒 + ∂𝐒₁ = zeros(S, size(𝐒₁)) + ∂𝐒₂ = zeros(S, size(𝐒₂)) + ∂state_init = [zeros(S, nVars), zeros(S, nVars)] + ∂SS_from_init = zeros(S, nVar_len) + n_aug = nPast + 1 + nExo + # Preallocated kron buffer reused across all (si, t) iterations + kaug₁ = Vector{S}(undef, n_aug^2) + + for si in 1:nShocks + ∂y₁_accum = zeros(S, nVars) + ∂δ_accum = zeros(S, nVars) + + for t in periods:-1:1 + ∂out_t = ∂Y_all[:, t, si] + ∂y₁_t = ∂out_t + ∂y₁_accum + ∂δ_t = ∂out_t + ∂δ_accum + + prev_st = states_store[si, t] + shock_t = shocks_store[si, t] + + aug₁ = [prev_st[1][past_idx]; one(S); shock_t] + aug₂ = [prev_st[2][past_idx]; zero(S); zero(shock_t)] + ℒ.kron!(kaug₁, aug₁, aug₁) + + # y₁_new = 𝐒₁ * aug₁ + ∂𝐒₁ .+= ∂y₁_t * aug₁' + ∂aug₁ = 𝐒₁' * ∂y₁_t + + # δ_new = 𝐒₁ * aug₂ + 𝐒₂ * kron(aug₁,aug₁) / 2 + ∂𝐒₁ .+= ∂δ_t * aug₂' + ∂aug₂ = 𝐒₁' * ∂δ_t + ∂𝐒₂ .+= ∂δ_t * kaug₁' / 2 + ∂kaug₁ = 𝐒₂' * ∂δ_t / 2 + ∂kaug₁_mat = reshape(∂kaug₁, n_aug, n_aug) + ∂aug₁ .+= ∂kaug₁_mat' * aug₁ + ∂kaug₁_mat * aug₁ + + ∂y₁_accum = zeros(S, nVars) + ∂δ_accum = zeros(S, nVars) + ∂y₁_accum[past_idx] .+= ∂aug₁[1:nPast] + ∂δ_accum[past_idx] .+= ∂aug₂[1:nPast] + end + + ∂state_init[1] .+= ∂y₁_accum + ∂state_init[2] .+= ∂δ_accum + if initial_state != [0.0] + ∂SS_from_init[1:nVar_len] .-= ∂y₁_accum[1:nVar_len] + end + end + + return [∂𝐒₁, ∂𝐒₂], ∂state_init, ∂SS_from_init +end + +function irf_bptt(::Val{:pruned_third_order}, + ∂Y_all::Array{S,3}, states_store, shocks_store, + nShocks, periods, past_idx, nPast, nVars, nExo, + 𝐒, initial_state, nVar_len) where S + 𝐒₁, 𝐒₂, 𝐒₃ = 𝐒 + ∂𝐒₁ = zeros(S, size(𝐒₁)) + ∂𝐒₂ = zeros(S, size(𝐒₂)) + ∂𝐒₃ = zeros(S, size(𝐒₃)) + ∂state_init = [zeros(S, nVars), zeros(S, nVars), zeros(S, nVars)] + ∂SS_from_init = zeros(S, nVar_len) + n_aug = nPast + 1 + nExo + # Preallocated kron buffers reused across all (si, t) iterations + kaug₁ = Vector{S}(undef, n_aug^2) + kaug₁₁ = Vector{S}(undef, n_aug^3) + k_aug₁̂_aug₂ = Vector{S}(undef, n_aug^2) + + for si in 1:nShocks + ∂y₁_accum = zeros(S, nVars) + ∂δ_accum = zeros(S, nVars) + ∂ξ_accum = zeros(S, nVars) + + for t in periods:-1:1 + ∂out_t = ∂Y_all[:, t, si] + ∂y₁_t = ∂out_t + ∂y₁_accum + ∂δ_t = ∂out_t + ∂δ_accum + ∂ξ_t = ∂out_t + ∂ξ_accum + + prev_st = states_store[si, t] + shock_t = shocks_store[si, t] + + aug₁ = [prev_st[1][past_idx]; one(S); shock_t] + aug₁̂ = [prev_st[1][past_idx]; zero(S); shock_t] + aug₂ = [prev_st[2][past_idx]; zero(S); zero(shock_t)] + aug₃ = [prev_st[3][past_idx]; zero(S); zero(shock_t)] + ℒ.kron!(kaug₁, aug₁, aug₁) + ℒ.kron!(kaug₁₁, kaug₁, aug₁) + + # y₁_new = 𝐒₁ * aug₁ + ∂𝐒₁ .+= ∂y₁_t * aug₁' + ∂aug₁ = 𝐒₁' * ∂y₁_t + + # δ_new = 𝐒₁ * aug₂ + 𝐒₂ * kron(aug₁,aug₁) / 2 + ∂𝐒₁ .+= ∂δ_t * aug₂' + ∂aug₂ = 𝐒₁' * ∂δ_t + ∂𝐒₂ .+= ∂δ_t * kaug₁' / 2 + ∂kaug₁_from_δ = 𝐒₂' * ∂δ_t / 2 + ∂kaug₁_mat = reshape(∂kaug₁_from_δ, n_aug, n_aug) + ∂aug₁ .+= ∂kaug₁_mat' * aug₁ + ∂kaug₁_mat * aug₁ + + # ξ_new = 𝐒₁ * aug₃ + 𝐒₂ * kron(aug₁̂, aug₂) + 𝐒₃ * kron(kaug₁, aug₁) / 6 + ∂𝐒₁ .+= ∂ξ_t * aug₃' + ∂aug₃ = 𝐒₁' * ∂ξ_t + + ℒ.kron!(k_aug₁̂_aug₂, aug₁̂, aug₂) + ∂𝐒₂ .+= ∂ξ_t * k_aug₁̂_aug₂' + ∂k12 = 𝐒₂' * ∂ξ_t + ∂k12_mat = reshape(∂k12, n_aug, n_aug) + ∂aug₁̂ = ∂k12_mat * aug₂ + ∂aug₂ .+= ∂k12_mat' * aug₁̂ + + ∂𝐒₃ .+= ∂ξ_t * kaug₁₁' / 6 + ∂kaug₁₁ = 𝐒₃' * ∂ξ_t / 6 + n_aug2 = n_aug * n_aug + ∂kaug₁₁_mat = reshape(∂kaug₁₁, n_aug2, n_aug) + ∂kaug₁_from_ξ = ∂kaug₁₁_mat * aug₁ + ∂aug₁ .+= ∂kaug₁₁_mat' * kaug₁ + ∂kaug₁_mat2 = reshape(∂kaug₁_from_ξ, n_aug, n_aug) + ∂aug₁ .+= ∂kaug₁_mat2' * aug₁ + ∂kaug₁_mat2 * aug₁ + + # aug₁̂ shares past_idx and shock with aug₁ + ∂aug₁[1:nPast] .+= ∂aug₁̂[1:nPast] + ∂aug₁[nPast+2:end] .+= ∂aug₁̂[nPast+2:end] + + ∂y₁_accum = zeros(S, nVars) + ∂δ_accum = zeros(S, nVars) + ∂ξ_accum = zeros(S, nVars) + ∂y₁_accum[past_idx] .+= ∂aug₁[1:nPast] + ∂δ_accum[past_idx] .+= ∂aug₂[1:nPast] + ∂ξ_accum[past_idx] .+= ∂aug₃[1:nPast] + end + + ∂state_init[1] .+= ∂y₁_accum + ∂state_init[2] .+= ∂δ_accum + ∂state_init[3] .+= ∂ξ_accum + if initial_state != [0.0] + ∂SS_from_init[1:nVar_len] .-= ∂y₁_accum[1:nVar_len] + end + end + + return [∂𝐒₁, ∂𝐒₂, ∂𝐒₃], ∂state_init, ∂SS_from_init +end + +function irf_bptt(::Val{:second_order}, + ∂Y_all::Array{S,3}, states_store, shocks_store, + nShocks, periods, past_idx, nPast, nVars, nExo, + 𝐒, initial_state, nVar_len) where S + 𝐒₁, 𝐒₂ = 𝐒 + ∂𝐒₁ = zeros(S, size(𝐒₁)) + ∂𝐒₂ = zeros(S, size(𝐒₂)) + ∂state_init = zeros(S, nVars) + ∂SS_from_init = zeros(S, nVar_len) + n_aug = nPast + 1 + nExo + kaug = Vector{S}(undef, n_aug^2) + + for si in 1:nShocks + ∂y_accum = zeros(S, nVars) + + for t in periods:-1:1 + ∂y_t = ∂Y_all[:, t, si] + ∂y_accum + + prev_st = states_store[si, t] + shock_t = shocks_store[si, t] + aug = [prev_st[past_idx]; one(S); shock_t] + ℒ.kron!(kaug, aug, aug) + + ∂𝐒₁ .+= ∂y_t * aug' + ∂aug = 𝐒₁' * ∂y_t + ∂𝐒₂ .+= ∂y_t * kaug' / 2 + ∂kaug = 𝐒₂' * ∂y_t / 2 + ∂kaug_mat = reshape(∂kaug, n_aug, n_aug) + ∂aug .+= ∂kaug_mat' * aug + ∂kaug_mat * aug + + ∂y_accum = zeros(S, nVars) + ∂y_accum[past_idx] .+= ∂aug[1:nPast] + end + + ∂state_init .+= ∂y_accum + if initial_state != [0.0] + ∂SS_from_init[1:nVar_len] .-= ∂y_accum[1:nVar_len] + end + end + + return [∂𝐒₁, ∂𝐒₂], ∂state_init, ∂SS_from_init +end + +function irf_bptt(::Val{:third_order}, + ∂Y_all::Array{S,3}, states_store, shocks_store, + nShocks, periods, past_idx, nPast, nVars, nExo, + 𝐒, initial_state, nVar_len) where S + 𝐒₁, 𝐒₂, 𝐒₃ = 𝐒 + ∂𝐒₁ = zeros(S, size(𝐒₁)) + ∂𝐒₂ = zeros(S, size(𝐒₂)) + ∂𝐒₃ = zeros(S, size(𝐒₃)) + ∂state_init = zeros(S, nVars) + ∂SS_from_init = zeros(S, nVar_len) + n_aug = nPast + 1 + nExo + kaug = Vector{S}(undef, n_aug^2) + kaug3 = Vector{S}(undef, n_aug^3) + + for si in 1:nShocks + ∂y_accum = zeros(S, nVars) + + for t in periods:-1:1 + ∂y_t = ∂Y_all[:, t, si] + ∂y_accum + + prev_st = states_store[si, t] + shock_t = shocks_store[si, t] + aug = [prev_st[past_idx]; one(S); shock_t] + ℒ.kron!(kaug, aug, aug) + ℒ.kron!(kaug3, kaug, aug) + + ∂𝐒₁ .+= ∂y_t * aug' + ∂aug = 𝐒₁' * ∂y_t + ∂𝐒₂ .+= ∂y_t * kaug' / 2 + ∂kaug = 𝐒₂' * ∂y_t / 2 + ∂𝐒₃ .+= ∂y_t * kaug3' / 6 + ∂kaug3 = 𝐒₃' * ∂y_t / 6 + + n_aug2 = n_aug * n_aug + ∂kaug3_mat = reshape(∂kaug3, n_aug2, n_aug) + ∂kaug .+= ∂kaug3_mat * aug + ∂aug .+= ∂kaug3_mat' * kaug + + ∂kaug_mat = reshape(∂kaug, n_aug, n_aug) + ∂aug .+= ∂kaug_mat' * aug + ∂kaug_mat * aug + + ∂y_accum = zeros(S, nVars) + ∂y_accum[past_idx] .+= ∂aug[1:nPast] + end + + ∂state_init .+= ∂y_accum + if initial_state != [0.0] + ∂SS_from_init[1:nVar_len] .-= ∂y_accum[1:nVar_len] + end + end + + return [∂𝐒₁, ∂𝐒₂, ∂𝐒₃], ∂state_init, ∂SS_from_init +end + + +# ── Dispatched rrule chain helpers for get_irf ─────────────────────────────────── + +# Forward chain: set up the rrule sub-calls and return standardized output. +# Returns (𝐒, SS_and_pars, state, solved, chain_ctx) or nothing on failure. +function irf_rrule_forward_chain(::Val{:first_order}, parameters::Vector{S}, + 𝓂, constants_obj, opts, tol) where S + nsss_out, nsss_pb = rrule(get_NSSS_and_parameters, 𝓂, parameters; + opts = opts, estimation = true) + reference_steady_state = nsss_out[1] + solution_error = nsss_out[2][1] + + if (solution_error > tol.nsss.acceptance_tol) || isnan(solution_error) + return nothing + end + + ∇₁, jac_pb = rrule(calculate_jacobian, parameters, reference_steady_state, + 𝓂.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 = parameters) + + sol_mat = first_out[1] + solved = first_out[3] + update_perturbation_counter!(𝓂.counters, solved, estimation = true, order = 1) + + return (sol_mat, reference_steady_state, nothing, solved, (nsss_pb, jac_pb, first_pb)) +end + +function irf_rrule_forward_chain(val_alg::Val, parameters::Vector{S}, + 𝓂, constants_obj, opts, tol) where S + ss_rrule = rrule(get_relevant_steady_state_and_state_update, + val_alg, parameters, 𝓂; + opts = opts, estimation = true) + + if ss_rrule === nothing + return nothing + end + + ss_out, ss_pb = ss_rrule + SS_and_pars = ss_out[2] + 𝐒 = ss_out[3] + state = ss_out[4] + solved = ss_out[5] + + return (𝐒, SS_and_pars, state, solved, (ss_pb,)) +end + +# Backward chain: propagate gradients through sub-rrule pullbacks. +# Returns ∂parameters vector. +function irf_rrule_backward_chain(::Val{:first_order}, ∂SS_and_pars, ∂𝐒_list, ∂state_init, chain_ctx) + nsss_pb, jac_pb, first_pb = chain_ctx + ∂sol_mat = ∂𝐒_list[1] + + 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] + + return ∂parameters_from_jac .+ ∂parameters_from_nsss +end + +function irf_rrule_backward_chain(::Val, ∂SS_and_pars, ∂𝐒_list, ∂state_init, chain_ctx) + ss_pb = chain_ctx[1] + ss_grads = ss_pb((NoTangent(), ∂SS_and_pars, ∂𝐒_list, ∂state_init, NoTangent())) + return ss_grads[3] +end + + +function rrule(::typeof(get_irf), + 𝓂::ℳ, + parameters::Vector{S}; + steady_state_function::SteadyStateFunctionType = missing, + periods::Int = DEFAULT_PERIODS, + algorithm::Symbol = :first_order, + 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_SELECTOR(𝓂), + sylvester_algorithm::Union{Symbol,Vector{Symbol},Tuple{Symbol,Vararg{Symbol}}} = DEFAULT_SYLVESTER_SELECTOR(𝓂), + lyapunov_algorithm::Symbol = DEFAULT_LYAPUNOV_ALGORITHM, + caching::Bool = DEFAULT_CACHING, + use_workspaces::Bool = DEFAULT_USE_WORKSPACES) where S <: Real + + val_alg = Val(algorithm) + + # Construct calculation options (sylvester/lyapunov fields ignored for first-order) + nPnExo = 𝓂.constants.post_model_macro.nPast_not_future_and_mixed + 1 + 𝓂.constants.post_model_macro.nExo + syl_sum = sum(k * (k + 1) ÷ 2 for k in 1:nPnExo) + 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) ? + syl_sum > DEFAULT_SYLVESTER_THRESHOLD ? DEFAULT_LARGE_SYLVESTER_ALGORITHM : DEFAULT_SYLVESTER_ALGORITHM : + sylvester_algorithm[2], + lyapunov_algorithm = lyapunov_algorithm) + + # Pre-solve setup + if !caching; invalidate_cache_validity!(𝓂); end + orig_ws = 𝓂.workspaces + if !use_workspaces; 𝓂.workspaces = fresh_workspaces(orig_ws); end + + constants_obj = initialise_constants!(𝓂) + + solve!(𝓂, + steady_state_function = steady_state_function, + opts = opts, + algorithm = algorithm) + + # Common shock/variable processing + 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 + past_idx = 𝓂.constants.post_model_macro.past_not_future_and_mixed_idx + nPast = length(past_idx) + nExo = 𝓂.constants.post_model_macro.nExo + nShocks = shocks == :none ? 1 : length(shock_idx) + nVar_len = length(𝓂.constants.post_model_macro.var) + + zero_result() = zeros(S, length(var_idx), periods, nShocks) + zero_pb(_) = (NoTangent(), NoTangent(), zeros(S, length(parameters))) + + # Dispatched rrule chain forward + chain_result = irf_rrule_forward_chain(val_alg, parameters, 𝓂, constants_obj, opts, tol) + + if chain_result === nothing + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return zero_result(), zero_pb + end + + 𝐒, SS_and_pars, state, solved, chain_ctx = chain_result + + if !solved + if !use_workspaces; 𝓂.workspaces = orig_ws; end + return zero_result(), zero_pb + end + + reference_ss = SS_and_pars[1:nVars] + + # Forward simulation (already dispatched via irf_initial_state / irf_forward_simulate!) + init_st = irf_initial_state(val_alg, state, SS_and_pars, initial_state, nVars, S) + + Y_all = zeros(S, nVars, periods, nShocks) + states_store = Array{Any}(undef, nShocks, periods + 1) + shocks_store = Array{Vector{S}}(undef, nShocks, periods) + + irf_forward_simulate!(val_alg, Y_all, states_store, shocks_store, + init_st, shock_idx, shocks, negative_shock, shock_history, + nExo, periods, past_idx, nVars, 𝐒) + + # Assemble output + deviations = Y_all[var_idx, :, :] + result = levels ? deviations .+ reference_ss[var_idx] : deviations + + if !use_workspaces; 𝓂.workspaces = orig_ws; end + + # Pullback (common structure, algorithm-specific parts dispatched) + function get_irf_pullback(∂result_bar) + ∂result = unthunk(∂result_bar) + + if ∂result isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, length(parameters)) + end + + ∂Y_all = zeros(S, nVars, periods, nShocks) + ∂Y_all[var_idx, :, :] .= ∂result + + ∂SS_and_pars = zeros(S, length(SS_and_pars)) + if levels + ∂SS_and_pars[var_idx] .+= dropdims(sum(∂result, dims = (2, 3)), dims = (2, 3)) + end + + # Dispatched BPTT + ∂𝐒_list, ∂state_init, ∂SS_from_init = irf_bptt(val_alg, + ∂Y_all, states_store, shocks_store, + nShocks, periods, past_idx, nPast, nVars, nExo, + 𝐒, initial_state, nVar_len) + + ∂SS_and_pars[1:nVar_len] .+= ∂SS_from_init + + # Dispatched backward chain through sub-rrule pullbacks + ∂parameters = irf_rrule_backward_chain(val_alg, ∂SS_and_pars, ∂𝐒_list, ∂state_init, chain_ctx) + + return NoTangent(), NoTangent(), ∂parameters + end + + return result, get_irf_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; + initial_guess = 𝓂.caches.covariance_first_order, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.first_order.ad.lyapunov, + verbose = opts.verbose, + has_unit_roots = 𝓂.caches.has_unit_roots) + covar_raw = lyap_out[1] + solved_lyap = lyap_out[2] + + # Cache the Lyapunov result for reuse + if solved_lyap + if size(𝓂.caches.covariance_first_order) != size(covar_raw) + 𝓂.caches.covariance_first_order = Matrix{Float64}(undef, size(covar_raw)...) + end + copyto!(𝓂.caches.covariance_first_order, covar_raw) + 𝓂.caches.valid_for.covariance_first_order = Float64.(parameters) + end + + 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_helper(∂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_helper(∂s₁ks₁, s_to_s₁, s_to_s₁) + ∂e₁_L, ∂e₁_R = kron_vjp_helper(∂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_helper(∂s₁ks₁, s_to_s₁, s_to_s₁) + ∂e₁_L, ∂e₁_R = kron_vjp_helper(∂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; + initial_guess = 𝓂.caches.covariance_second_order, + lyapunov_algorithm = opts.lyapunov_algorithm, + tol = opts.tol.second_order.ad.lyapunov, + verbose = opts.verbose, + has_unit_roots = 𝓂.caches.has_unit_roots) + Σᶻ₂ = lyap_out[1] + info = lyap_out[2] + + # Cache the 2nd-order Lyapunov result for reuse + if info + if size(𝓂.caches.covariance_second_order) != size(Σᶻ₂) + 𝓂.caches.covariance_second_order = Matrix{Float64}(undef, size(Σᶻ₂)...) + end + copyto!(𝓂.caches.covariance_second_order, Σᶻ₂) + 𝓂.caches.valid_for.covariance_second_order = Float64.(parameters) + end + + 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_helper(∂Γ₂_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_helper(∂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_helper(∂e₁ke₁_total, e_to_s₁, e_to_s₁) + ∂s₁_se_L, ∂e₁_se_R = kron_vjp_helper(∂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}(sparse_ABAt(ê_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₃' + sparse_ABAt(ê_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 + + # Cache the 3rd-order covariance for reuse + all_solved_3rd = solved && solved3 && solved_lyapunov + if all_solved_3rd + if size(𝓂.caches.covariance_third_order) != size(Σʸ₃) + 𝓂.caches.covariance_third_order = Matrix{Float64}(undef, size(Σʸ₃)...) + end + copyto!(𝓂.caches.covariance_third_order, Σʸ₃) + 𝓂.caches.valid_for.covariance_third_order = Float64.(parameters) + nVars_rrule = T_pm.nVars + obs_key_rrule = if observables == :full_covar + collect(1:nVars_rrule) + else + obs_idx = parse_variables_input_to_index(observables, 𝓂.constants) |> sort + if covariance == Symbol[] + collect(obs_idx) + else + covar_idx = parse_variables_input_to_index(covariance, 𝓂.constants) |> sort + sort(union(obs_idx, covar_idx)) + end + end + 𝓂.caches.valid_for.covariance_third_order_obs_key = obs_key_rrule + end + + result = (Σʸ₃, μʸ₂, SS_and_pars, all_solved_3rd) + + # ── 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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[gb[3]+1:gb[4], gb[3]+1:gb[4]], Matrix(d.Σ̂ᶻ₁), Ine)[1] + # Row 4: (4,1) kron(Δ̂μˢ₂,Ine) + ∂Δ̂μˢ₂_l .+= vec(kron_vjp_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂EL6[:, sb[1]+1:sb[2]], Matrix(d.Σ̂ᶻ₁), vec_Ie_col)[1] + # Col 4: kron(μˢ₃δμˢ₁', vec_Ie) + ∂μ_T = kron_vjp_helper(∂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_helper(∂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_helper(∂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_helper(∂M3_raw, Matrix(d.Σ̂ᶻ₁), vec_Ie_col)[1] + # Decompose ∂M4 → ∂Δ̂μˢ₂ + ∂Δ̂μˢ₂_l .+= vec(kron_vjp_helper(∂M4_raw, reshape(d.Δ̂μˢ₂, :, 1), Ine)[1]) + # Decompose ∂M6 → ∂Σ̂ᶻ₁ + ∂Σ̂ᶻ₁ .+= reshape(kron_vjp_helper(∂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}(sparse_ABAt(ê_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₃' + sparse_ABAt(ê_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 + + # Cache the 3rd-order covariance for reuse + all_solved_3rd = solved && solved3 && solved_lyapunov + if all_solved_3rd + if size(𝓂.caches.covariance_third_order) != size(Σʸ₃) + 𝓂.caches.covariance_third_order = Matrix{Float64}(undef, size(Σʸ₃)...) + end + copyto!(𝓂.caches.covariance_third_order, Σʸ₃) + 𝓂.caches.valid_for.covariance_third_order = Float64.(parameters) + nVars_rrule = T_pm.nVars + obs_key_rrule = if observables == :full_covar + collect(1:nVars_rrule) + else + obs_idx = parse_variables_input_to_index(observables, 𝓂.constants) |> sort + if covariance == Symbol[] + collect(obs_idx) + else + covar_idx = parse_variables_input_to_index(covariance, 𝓂.constants) |> sort + sort(union(obs_idx, covar_idx)) + end + end + 𝓂.caches.valid_for.covariance_third_order_obs_key = obs_key_rrule + + # Cache autocorrelation + if size(𝓂.caches.covariance_third_order_autocorr) != size(autocorr) + 𝓂.caches.covariance_third_order_autocorr = Matrix{Float64}(undef, size(autocorr)...) + end + copyto!(𝓂.caches.covariance_third_order_autocorr, autocorr) + 𝓂.caches.valid_for.covariance_third_order_autocorr = Float64.(parameters) + 𝓂.caches.valid_for.covariance_third_order_autocorr_obs_key = obs_key_rrule + 𝓂.caches.valid_for.covariance_third_order_autocorr_periods = collect(Int, autocorrelation_periods) + end + + result = (Σʸ₃, μʸ₂, autocorr, SS_and_pars, all_solved_3rd) + + # ── 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_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(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_helper(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_helper(∂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_helper(∂inner66, s₁, s₁²) + ∂s₁_l .+= tmpA + tmpL, tmpR = kron_vjp_helper(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_helper(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_helper(∂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_helper(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_helper(∂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_helper(∂b55, s₁, se₂) + ∂s₁_l .+= tmpA; ∂se₂_l .+= tmpB + ∂k55b = Matrix(d.s_s') * ∂b55 + tmpA, tmpB = kron_vjp_helper(∂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_helper(∂b56, s₁, eeh) + ∂s₁_l .+= tmpA; ∂ee₂_l .+= tmpB ./ 2 + ∂k56b = Matrix(d.s_s') * ∂b56 + tmpA, tmpB = kron_vjp_helper(∂k56b, se₂, e₁) + ∂se₂_l .+= tmpA; ∂e₁_l .+= tmpB + # (5,7) kron(e₁, ee₂/2) + tmpA, tmpB = kron_vjp_helper(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_helper(∂b65, s₁², e₁) + ∂e₁_l .+= tmpB + tmpL, tmpR = kron_vjp_helper(tmpA, s₁, s₁); ∂s₁_l .+= tmpL .+ tmpR + tmpA, tmpB = kron_vjp_helper(∂b65, s₁, ss_s1e1) + ∂s₁_l .+= tmpA + tmpC = Matrix(d.s_s') * tmpB + tmpL, tmpR = kron_vjp_helper(tmpC, s₁, e₁); ∂s₁_l .+= tmpL; ∂e₁_l .+= tmpR + ∂k65c = ∂b65 * Matrix(d.e_ss') + tmpA, tmpB = kron_vjp_helper(∂k65c, e₁, s₁²) + ∂e₁_l .+= tmpA + tmpL, tmpR = kron_vjp_helper(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_helper(∂b66, s₁e₁, e₁) + ∂e₁_l .+= tmpB + tmpL, tmpR = kron_vjp_helper(tmpA, s₁, e₁); ∂s₁_l .+= tmpL; ∂e₁_l .+= tmpR + ∂pre = ∂b66 * Matrix(d.e_es') + tmpA, tmpB = kron_vjp_helper(∂pre, e₁, s₁e₁) + ∂e₁_l .+= tmpA + tmpL, tmpR = kron_vjp_helper(tmpB, s₁, e₁); ∂s₁_l .+= tmpL; ∂e₁_l .+= tmpR + tmpA, tmpB = kron_vjp_helper(∂pre, e₁, ss_s1e1) + ∂e₁_l .+= tmpA + tmpC = Matrix(d.s_s') * tmpB + tmpL, tmpR = kron_vjp_helper(tmpC, s₁, e₁); ∂s₁_l .+= tmpL; ∂e₁_l .+= tmpR + # (6,7) L₃ˢ * kron(e₁, e₁²) — decompress rows + tmpA, tmpB = kron_vjp_helper(Matrix(d.L₃ˢ)' * Matrix(∂ê₃[sb[6]+1:sb[7], eb[7]+1:eb[8]]), e₁, e₁²) + ∂e₁_l .+= tmpA + tmpL, tmpR = kron_vjp_helper(tmpB, e₁, e₁); ∂e₁_l .+= tmpL .+ tmpR + + # ── 3a: Γ₃ disaggregation → ∂Σ̂ᶻ₁, ∂Σ̂ᶻ₂, ∂Δ̂μˢ₂ ── + ∂Γ = Matrix{T}(∂Γ₃_iter) + vΣ = vec(d.Σ̂ᶻ₁) + + # Row 1: (1,4) kron(Δ̂μˢ₂',Ine) + ∂tmp14 = kron_vjp_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[gb[3]+1:gb[4], gb[3]+1:gb[4]], Matrix(d.Σ̂ᶻ₁), Ine)[1] + # Row 4: (4,1) kron(Δ̂μˢ₂,Ine) + ∂Δ̂μˢ₂_l .+= vec(kron_vjp_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂Γ[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_helper(∂EL6[:, sb[1]+1:sb[2]], Matrix(d.Σ̂ᶻ₁), vec_Ie_col)[1] + # Col 4: kron(μˢ₃δμˢ₁', vec_Ie) + ∂μ_T = kron_vjp_helper(∂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_helper(∂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_helper(∂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_helper(∂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_helper(∂M3_raw, Matrix(d.Σ̂ᶻ₁), vec_Ie_col)[1] + # Decompose ∂M4 → ∂Δ̂μˢ₂ + ∂Δ̂μˢ₂_l .+= vec(kron_vjp_helper(∂M4_raw, reshape(d.Δ̂μˢ₂, :, 1), Ine)[1]) + # Decompose ∂M6 → ∂Σ̂ᶻ₁ + ∂Σ̂ᶻ₁ .+= reshape(kron_vjp_helper(∂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 + ensure_first_order_constants!(constants) + idx_constants = constants.post_complete_parameters + + 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] + # Old way (≤v0.1.42): + # Q = qr(∇₀[:, present_only_idx]) + # A₊ = Q' * ∇₊; A₀ = Q' * ∇₀; A₋ = Q' * ∇₋ + # Current code reuses QR workspaces to avoid allocations. + qr_factors, qr_ws = ensure_first_order_fast_qr_workspace!(qme_ws, ∇₀_present) + Q = factorize_qr!((use_fastlapack_qr ? Val(:FastLapack) : Val(:Julia)), ∇₀_present, qr_factors, qr_ws) # Q = qr(∇₀_present) + + qme_ws.fast_qr_orm_ws_plus, qme_ws.fast_qr_orm_dims_plus = apply_qr_transpose_left!(A₊, ∇₊, Q, # 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, # 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, # 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 + + Ã₊ = qme_ws.𝐀̃₊ + ℒ.mul!(Ã₊, @view(A₊[dynIndex,:]), Ir[future_not_past_and_mixed_in_comb,:]) # Ã₊ = A₊[dynIndex,:] * Ir + + Ã₀ = qme_ws.𝐀̃₀ + copyto!(Ã₀, @view(A₀[dynIndex, comb])) + + Ã₋ = qme_ws.𝐀̃₋ + ℒ.mul!(Ã₋, @view(A₋[dynIndex,:]), Ir[past_not_future_and_mixed_in_comb,:]) # Ã₋ = A₋[dynIndex,:] * Ir + + # end # timeit_debug + # @timeit_debug timer "Quadratic matrix equation solve" begin + + sol, solved = solve_quadratic_matrix_equation(Ã₊, Ã₀, Ã₋, 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] + + Ā₀ᵤ = qme_ws.𝐀̄₀ᵤ + copyto!(Ā₀ᵤ, @view(A₀[1:T.nPresent_only, T.present_only_idx])) + + A₊ᵤ = qme_ws.𝐀₊ᵤ + copyto!(A₊ᵤ, @view(A₊[1:T.nPresent_only,:])) + + Ã₀ᵤ = qme_ws.𝐀̃₀ᵤ + copyto!(Ã₀ᵤ, @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 Ā₀ᵤ" begin + + # Old way (≤v0.1.42): Ā̂₀ᵤ = lu(Ā₀ᵤ) + qme_ws.fast_lu_ws_a0u, qme_ws.fast_lu_dims_a0u, solved_Ā₀ᵤ, Ā̂₀ᵤ = factorize_lu!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), Ā₀ᵤ, + qme_ws.fast_lu_ws_a0u, + qme_ws.fast_lu_dims_a0u) + + if !solved_Ā₀ᵤ + return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # A = vcat(-(Ā̂₀ᵤ \ (A₊ᵤ * D * L + Ã₀ᵤ * sol[T.dynamic_order,:] + A₋ᵤ)), sol) + if T.nPresent_only > 0 + ℒ.mul!(A₋ᵤ, Ã₀ᵤ, @view(sol[:,past_not_future_and_mixed_in_comb]), 1, 1) # A₋ᵤ = A₋ᵤ + Ã₀ᵤ * sol + nₚ₋ = qme_ws.𝐧ₚ₋ + ℒ.mul!(nₚ₋, A₊ᵤ, D) # nₚ₋ = A₊ᵤ * D + ℒ.mul!(A₋ᵤ, nₚ₋, L, 1, 1) # A₋ᵤ = A₋ᵤ + nₚ₋ * L + solve_lu_left!(Ā₀ᵤ, A₋ᵤ, qme_ws.fast_lu_ws_a0u, Ā̂₀ᵤ; + use_fastlapack_lu = use_fastlapack_lu) + ℒ.rmul!(A₋ᵤ, -1) # A₋ᵤ = -A₋ᵤ + 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) # Ŝᵗ = Sᵗ * expand_past # Ŝᵗ = Sᵗ * expand_past + + ∇₊ = qme_ws.sylvester.𝐀 + ℒ.mul!(∇₊, @view(∇₁[:,1:T.nFuture_not_past_and_mixed]), expand_future) # ∇₊ = ∇₁[:, future_cols] * expand_future + + ℒ.mul!(∇₀, ∇₊, 𝐒̂ᵗ, 1, 1) # ∇₀ = ∇₊ * Ŝᵗ + ∇₀ # ∇₀ = ∇₊ * Ŝᵗ + ∇₀ + + # Old way (≤v0.1.42): C = lu(∇₀) + # Old way (≤v0.1.42): C = lu(∇₀) + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_∇₀, C = factorize_lu!((use_fastlapack_lu ? Val(:FastLapack) : Val(:Julia)), ∇₀, + qme_ws.fast_lu_ws_nabla0, + qme_ws.fast_lu_dims_nabla0) + + if !solved_∇₀ + return (zeros(T.nVars,T.nPast_not_future_and_mixed + T.nExo), sol, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + # Old way (≤v0.1.42): ∇ₑ = -(∇₀ \ ∇ₑ) + 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 + # Old way (≤v0.1.42): M = ∇₀ \ I (i.e. inv(∇₀)) + solve_lu_left!(∇₀, M, qme_ws.fast_lu_ws_nabla0, C; + use_fastlapack_lu = use_fastlapack_lu) + + tmp2 = qme_ws.sylvester.𝐁 + ℒ.mul!(tmp2, M', ∇₊') # tmp2 = M' * ∇₊' + ℒ.rmul!(tmp2, -1) # tmp2 = -tmp2 + + ∇ₑ = @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. + ∂𝐒_mat = unthunk(∂𝐒[1]) + if ∂𝐒_mat isa Union{NoTangent, AbstractZero} + return NoTangent(), zero(∇₁), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + ∂∇₁ = zero(∇₁) + + ∂𝐒ᵗ = ∂𝐒_mat[:,1:T.nPast_not_future_and_mixed] + ∂𝐒ᵉ = ∂𝐒_mat[:,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², + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.first_order.ad.sylvester, + verbose = opts.verbose) + + if !solved + return 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) + workspaces.second_order = Higher_order_workspace(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ₑ₋)] + + ∇₁₊𝐒₁➕∇₁₀ = collect(@views -∇₁[:,1:n₊] * 𝐒₁[i₊,1:n₋] * M₂.𝐈ₙ₋ - ∇₁[:,range(1,n) .+ n₊]) + + # end # timeit_debug + # @timeit_debug timer "Invert matrix" begin + + qme_ws = workspaces.first_order + + if S === Float64 + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_∇lu, lu_handle = + factorize_lu!(Val(:FastLapack), ∇₁₊𝐒₁➕∇₁₀, qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0) + + if !solved_∇lu + if opts.verbose println("Second order solution: inversion failed") end + return (∇₁₊𝐒₁➕∇₁₀, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + spinv = Matrix{S}(ℒ.I, size(∇₁₊𝐒₁➕∇₁₀)) + solve_lu_left!(∇₁₊𝐒₁➕∇₁₀, spinv, qme_ws.fast_lu_ws_nabla0, lu_handle) + else + ∇₁₊𝐒₁➕∇₁₀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) + end + 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 + + # Doubling power-cache: enable capture so the pullback's adjoint solve can + # reuse A^(2^k), B^(2^k) from this forward pass. + cache_eligible_2nd = opts.sylvester_algorithm² == :doubling + if cache_eligible_2nd + ℂ.sylvester_workspace.pow_iters = 0 + ℂ.sylvester_workspace.pow_capture = true + ℂ.sylvester_workspace.pow_transposed = true + end + 𝐒₂, solved = solve_sylvester_equation(A, B, C, ℂ.sylvester_workspace, + initial_guess = initial_guess, + sylvester_algorithm = opts.sylvester_algorithm², + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.second_order.ad.sylvester, + verbose = opts.verbose) + ℂ.sylvester_workspace.pow_capture = false + pow_iters_captured_2nd = ℂ.sylvester_workspace.pow_iters + ℂ.sylvester_workspace.pow_iters = 0 + 𝐒₂_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 + + ∂𝐒₂ = unthunk(∂𝐒₂_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 + + ws = ℂ.sylvester_workspace + cache_valid = cache_eligible_2nd && + pow_iters_captured_2nd >= 1 && + ws.pow_transposed + saved_capture = ws.pow_capture + if cache_valid + ws.pow_iters = pow_iters_captured_2nd + ws.pow_capture = false + end + ∂C, solved = solve_sylvester_equation(At, Bt, ∂𝐒₂, ws, + sylvester_algorithm = opts.sylvester_algorithm², + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.second_order.ad.sylvester, + verbose = opts.verbose) + ws.pow_capture = saved_capture + ws.pow_iters = 0 + + 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 + +# Dead code: compressed_kron_pullback_2arg! — never called anywhere (2-arg compressed_kron is also dead) +#= +# 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 + # Hoist i1-dependent reads (column indices fixed by outer loop) + @inbounds aii = Xd[i1, i2]; aij = Xd[i1, j2]; aik = Xd[i1, k2] + for j1 in 1:i1 + # Hoist j1-dependent reads + @inbounds aji = Xd[j1, i2]; ajj = Xd[j1, j2]; ajk = Xd[j1, k2] + # Precompute sub-expressions for ∂X[k1, ...] updates + q_i2 = aij * ajk + aik * ajj + q_j2 = aik * aji + aii * ajk + q_k2 = aii * ajj + aij * aji + @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 + if i1 == j1 + divisor = (j1 == k1) ? 6 : 2 + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end + g_d = g / divisor + 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 * q_i2 + ∂X[k1, j2] += g_d * q_j2 + ∂X[k1, k2] += g_d * q_k2 + end + 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 + # Hoist i1-dependent reads + @inbounds aii = Xd[i1, i2]; aij = Xd[i1, j2]; aik = Xd[i1, k2] + for j1 in 1:i1 + # Hoist j1-dependent reads + @inbounds aji = Xd[j1, i2]; ajj = Xd[j1, j2]; ajk = Xd[j1, k2] + # Precompute sub-expressions for ∂X[k1, ...] updates + q_i2 = aij * ajk + aik * ajj + q_j2 = aik * aji + aii * ajk + q_k2 = aii * ajj + aij * aji + @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 + if i1 == j1 + divisor = (j1 == k1) ? 6 : 2 + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end + g_d = g / divisor + 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 * q_i2 + ∂X[k1, j2] += g_d * q_j2 + ∂X[k1, k2] += g_d * q_k2 + end + 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(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) + + ∇₁₊𝐒₁➕∇₁₀ = collect(@views -∇₁[:,1:n₊] * 𝐒₁[i₊,1:n₋] * M₂.𝐈ₙ₋ - ∇₁[:,range(1,n) .+ n₊]) + + qme_ws = workspaces.first_order + + if S === Float64 + qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0, solved_∇lu, lu_handle = + factorize_lu!(Val(:FastLapack), ∇₁₊𝐒₁➕∇₁₀, qme_ws.fast_lu_ws_nabla0, qme_ws.fast_lu_dims_nabla0) + + if !solved_∇lu + return (∇₁₊𝐒₁➕∇₁₀, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + spinv = Matrix{S}(ℒ.I, size(∇₁₊𝐒₁➕∇₁₀)) + solve_lu_left!(∇₁₊𝐒₁➕∇₁₀, spinv, qme_ws.fast_lu_ws_nabla0, lu_handle) + else + ∇₁₊𝐒₁➕∇₁₀lu = ℒ.lu(∇₁₊𝐒₁➕∇₁₀, check = false) + + if !ℒ.issuccess(∇₁₊𝐒₁➕∇₁₀lu) + return (∇₁₊𝐒₁➕∇₁₀, false), x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + spinv = inv(∇₁₊𝐒₁➕∇₁₀lu) + end + 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 = 𝐒₃ ---------------------------------------- + cache_eligible_3rd = opts.sylvester_algorithm³ == :doubling + if cache_eligible_3rd + ℂ.sylvester_workspace.pow_iters = 0 + ℂ.sylvester_workspace.pow_capture = true + ℂ.sylvester_workspace.pow_transposed = true + end + 𝐒₃, solved = solve_sylvester_equation(A, B, C, ℂ.sylvester_workspace, + initial_guess = initial_guess_sylv, + sylvester_algorithm = opts.sylvester_algorithm³, + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.third_order.ad.sylvester, + verbose = opts.verbose) + ℂ.sylvester_workspace.pow_capture = false + pow_iters_captured_3rd = ℂ.sylvester_workspace.pow_iters + ℂ.sylvester_workspace.pow_iters = 0 + + 𝐒₃ = 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(unthunk(∂𝐒₃_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 -------------------- + ws = ℂ.sylvester_workspace + cache_valid = cache_eligible_3rd && + pow_iters_captured_3rd >= 1 && + ws.pow_transposed + saved_capture = ws.pow_capture + if cache_valid + ws.pow_iters = pow_iters_captured_3rd + ws.pow_capture = false + end + ∂C_adj, slvd = solve_sylvester_equation(At, Bt, ∂𝐒₃, ws, + sylvester_algorithm = opts.sylvester_algorithm³, + preconditioner = opts.sylvester_preconditioner, + tol = opts.tol.third_order.ad.sylvester, + verbose = opts.verbose) + ws.pow_capture = saved_capture + ws.pow_iters = 0 + 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, + preconditioner::Symbol = :ilu, + tol::SolverTolerances = SolverTolerances(), + # timer::TimerOutput = TimerOutput(), + verbose::Bool = false) where {M <: AbstractMatrix{Float64}, N <: AbstractMatrix{Float64}, O <: AbstractMatrix{Float64}} + + # Enable doubling-power capture only for the dense-dense :doubling path + # (the dense-dense overload of solve_sylvester_equation populates 𝐀_pow/𝐁_pow). + # Enable doubling-power capture for the :doubling algorithm path. + # The solver overloads populate 𝕊ℂ.𝐀_pow / 𝐁_pow during forward iteration so + # the pullback can skip squaring. With pow_transposed=true, powers are stored + # in transposed form directly, saving a post-hoc transpose pass. + cache_eligible = sylvester_algorithm == :doubling + if cache_eligible + 𝕊ℂ.pow_iters = 0 + 𝕊ℂ.pow_capture = true + 𝕊ℂ.pow_transposed = true + end + P, solved = solve_sylvester_equation(A, B, C, 𝕊ℂ, + sylvester_algorithm = sylvester_algorithm, + preconditioner = preconditioner, + tol = tol, + verbose = verbose, + initial_guess = initial_guess) + 𝕊ℂ.pow_capture = false + pow_iters_captured = 𝕊ℂ.pow_iters + 𝕊ℂ.pow_iters = 0 + + 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)) + + # Precompute transposes once outside the pullback closure: needed for both + # the matmul forming ∂A/∂B (every call) and the fallback adjoint solve + # (when the doubling power cache is unavailable). + At = A' + Bt = B' + + # pullback + function solve_sylvester_equation_pullback(∂P) + ∂P₁ = unthunk(∂P[1]) + if ℒ.norm(∂P₁) < tol.rtol + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + cache_valid = cache_eligible && + pow_iters_captured >= 1 && + 𝕊ℂ.pow_transposed + saved_capture = 𝕊ℂ.pow_capture + if cache_valid + 𝕊ℂ.pow_iters = pow_iters_captured + 𝕊ℂ.pow_capture = false + end + ∂C, slvd = solve_sylvester_equation(At, Bt, ∂P₁, 𝕊ℂ, + sylvester_algorithm = sylvester_algorithm, + preconditioner = preconditioner, + tol = tol, + verbose = verbose) + 𝕊ℂ.pow_capture = saved_capture + 𝕊ℂ.pow_iters = 0 + + solved = solved && slvd + + if !slvd + return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # ∂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 * Bt) * P_cached' + ∂B = (P_cached' * At) * ∂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, + has_unit_roots::Bool = false) + + # Enable doubling-power capture for the :doubling algorithm path. + # With pow_transposed=true, powers are stored in transposed form directly. + if lyapunov_algorithm == :doubling + workspace.pow_iters = 0 + workspace.pow_capture = true + workspace.pow_transposed = true + end + P, solved = solve_lyapunov_equation(A, C, workspace, + initial_guess = initial_guess, + lyapunov_algorithm = lyapunov_algorithm, + tol = tol, + verbose = verbose, + has_unit_roots = has_unit_roots) + workspace.pow_capture = false + pow_iters_captured = workspace.pow_iters + workspace.pow_iters = 0 + 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) + # Precompute Aᵀ once outside the pullback closure: needed by the matmul + # forming ∂A (every call) and by the fallback adjoint solve. + At = A_dense' + + # pullback + # https://arxiv.org/abs/2011.11430 + function solve_lyapunov_equation_pullback(∂P) + ∂P₁ = unthunk(∂P[1]) + if ℒ.norm(∂P₁) < tol.rtol + return NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + # Adjoint Lyapunov: ∂P is generally not symmetric, so issymmetric will route to full-space. + # Prefer the forward dense doubling solver in replay mode against the + # transposed power cache when the forward pass populated workspace.𝐀_pow; + # otherwise fall back to the legacy solver call. + cache_valid = lyapunov_algorithm == :doubling && + pow_iters_captured >= 1 && + workspace.pow_transposed + saved_capture = workspace.pow_capture + if cache_valid + workspace.pow_iters = pow_iters_captured + workspace.pow_capture = false + end + ∂C, slvd = solve_lyapunov_equation(At, Matrix{Float64}(∂P₁), workspace, + lyapunov_algorithm = lyapunov_algorithm, + tol = tol, + verbose = verbose) + workspace.pow_capture = saved_capture + workspace.pow_iters = 0 + + solved = solved && slvd + + if !slvd + return NoTangent(), NoTangent(), NoTangent(), NoTangent() + end + + 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(unthunk(∂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(unthunk(∂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 + + # Warmup forward pass. When `warmup_iterations > 0` we build a + # block-concatenated jacobian, solve a min-norm linear system to recover + # `warmup_iterations` worth of shocks, propagate the state through the + # warmup window, and add the corresponding contributions to `logabsdets` + # and `shocks²`. Intermediates are captured so the pullback can backprop + # through the linear solve, the state propagation and the jacobian + # construction. + warmup_jac = zeros(0, 0) + warmup_x = zeros(0) + warmup_y = zeros(0) # = inv(JJt) * data[:,1] (fat case only) + warmup_state_history = Vector{Vector{Float64}}() + warmup_Sᵉ_powers = Matrix{Float64}[] # [I, Sᵉ, Sᵉ², …, Sᵉ^(N-2)] + warmup_data_first = zeros(length(obs_idx)) + + if warmup_iterations > 0 + warmup_data_first = collect(data_in_deviations[:,1]) + + warmup_jac = 𝐒[obs_idx, end-T.nExo+1:end] + + if warmup_iterations >= 2 + warmup_jac = hcat(𝐒[obs_idx, 1:T.nPast_not_future_and_mixed] * + 𝐒[t⁻, end-T.nExo+1:end], warmup_jac) + push!(warmup_Sᵉ_powers, Matrix{Float64}(ℒ.I, T.nPast_not_future_and_mixed, + T.nPast_not_future_and_mixed)) # Sᵉ^0 + if warmup_iterations >= 3 + Sᵉ_pow = 𝐒[t⁻, 1:T.nPast_not_future_and_mixed] + push!(warmup_Sᵉ_powers, copy(Sᵉ_pow)) # Sᵉ^1 + for e in 1:warmup_iterations-2 + warmup_jac = hcat(𝐒[obs_idx, 1:T.nPast_not_future_and_mixed] * + Sᵉ_pow * 𝐒[t⁻, end-T.nExo+1:end], warmup_jac) + if e < warmup_iterations - 2 + Sᵉ_pow = Sᵉ_pow * 𝐒[t⁻, 1:T.nPast_not_future_and_mixed] + push!(warmup_Sᵉ_powers, copy(Sᵉ_pow)) + end + end + end + end + + # Solve linear system + if size(warmup_jac, 1) == size(warmup_jac, 2) + warmup_lu = ℒ.lu(warmup_jac, check = false) + if !ℒ.issuccess(warmup_lu) + if opts.verbose println("Inversion filter failed (warmup, rrule)") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + warmup_x = warmup_lu \ warmup_data_first + else + JJt_w = warmup_jac * warmup_jac' + JJt_w_lu = ℒ.lu(JJt_w, check = false) + if !ℒ.issuccess(JJt_w_lu) + if opts.verbose println("Inversion filter failed (warmup, rrule)") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + warmup_y = JJt_w_lu \ warmup_data_first + warmup_x = warmup_jac' * warmup_y + end + + warmup_shocks_mat = reshape(warmup_x, T.nExo, warmup_iterations) + + # State propagation across warmup window + st_local = copy(state) + push!(warmup_state_history, copy(st_local)) + for i in 1:warmup_iterations-1 + st_concat = vcat(st_local[t⁻], warmup_shocks_mat[:,i]) + st_local = 𝐒 * st_concat + push!(warmup_state_history, copy(st_local)) + end + state = st_local + + # NOTE: We deliberately do NOT add per-block logabsdets here. The + # primal in `src/filter/inversion.jl` accumulates them at lines 90-97 + # but then unconditionally overwrites `logabsdets` at lines 119/133/145 + # before the main loop scales it by `(n_obs - presample)`. As a result, + # warmup logabsdets contributions never enter `llh`, so the rrule must + # not produce gradients for them either. + + shocks² += sum(abs2, warmup_x) + end + + state = [copy(state) for _ in 1:size(data_in_deviations,2)+1] + + 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) + lu_ws = FastLapackInterface.LUWs(jac) + lu_ws, _, ok, lu_handle = factorize_lu!(Val(:FastLapack), jac, lu_ws, size(jac)) + + if !ok + if opts.verbose println("Inversion filter failed") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + logabsdets = 0.0 + @inbounds for k in 1:size(jac,1) + logabsdets += log(abs(jac[k,k])) + end + invjac = Matrix{Float64}(ℒ.I, size(jac)) + solve_lu_left!(jac, invjac, lu_ws, lu_handle) + 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 + + # ----- Warmup pullback ------------------------------------------------ + # Backprop through the warmup forward. At this point ∂state holds the + # adjoint of the state at the start of the main loop, which equals the + # state at the end of the warmup window (i.e. ∂state_after_warmup). + # We propagate it back through (i) state propagation across the warmup + # window, (ii) the linear-solve recovery of the warmup shocks, and + # (iii) the block-concatenated jacobian construction. + if warmup_iterations > 0 + N = warmup_iterations + nExo = T.nExo + n_pnf = T.nPast_not_future_and_mixed + + # ∂x_warmup gets contributions from (a) shocks² += sum(abs2, x_warmup) + # and (b) the state-propagation backward sweep. + ∂x_warmup = -copy(warmup_x) # from shocks² (∂llh*-1/2 implicit) + + # Backprop state propagation (warmup_iterations-1 evolution steps). + ∂state_local = copy(∂state) # = ∂state_after_warmup + for i in (N-1):-1:1 + state_concat_i = vcat(warmup_state_history[i][t⁻], + warmup_x[(i-1)*nExo+1 : i*nExo]) + ∂𝐒 .+= ∂state_local * state_concat_i' + ∂state_concat = 𝐒' * ∂state_local + # ∂warmup_shocks[:,i] contribution + ∂x_warmup[(i-1)*nExo+1 : i*nExo] .+= ∂state_concat[n_pnf+1:end] + # Reset ∂state and inject t⁻ slots for previous step + ∂state_local = zero(∂state_local) + ∂state_local[t⁻] .= ∂state_concat[1:n_pnf] + end + # After the loop, ∂state_local is the gradient wrt state_initial, + # supported only on t⁻ slots. Override the ∂state we'll return. + ∂state .= ∂state_local + + # ∂jac_concat collects contributions from the linear-solve adjoint + # only. We do NOT add per-block logabsdets contributions because + # the primal silently overwrites the accumulated warmup logabsdets + # before scaling (see note in the forward pass), so they don't + # enter llh and must not enter the gradient. + ∂jac_concat = zeros(size(warmup_jac)) + + # Backprop the linear solve to recover warmup shocks. + ∂data_first = zeros(length(obs_idx)) + if size(warmup_jac, 1) == size(warmup_jac, 2) + # x = jac \ data; ∂data = jac' \ ∂x; ∂jac = -∂data * x' + ∂data_first = warmup_jac' \ ∂x_warmup + ∂jac_concat .-= ∂data_first * warmup_x' + else + # x = jac' * inv(JJt) * data, JJt = jac*jac', y = inv(JJt)*data + # ∂data = inv(JJt) * jac * ∂x + # ∂jac += y * ∂x' - ∂data * x' - y * (jac' * ∂data)' + JJt_w = warmup_jac * warmup_jac' + ∂data_first = JJt_w \ (warmup_jac * ∂x_warmup) + ∂jac_concat .+= warmup_y * ∂x_warmup' + ∂jac_concat .-= ∂data_first * warmup_x' + ∂jac_concat .-= warmup_y * (warmup_jac' * ∂data_first)' + end + ∂data_in_deviations[:,1] .+= ∂data_first + + # Map ∂jac_concat → ∂𝐒. + # Block N is C = 𝐒[obs_idx, end-nExo+1:end]. + ∂𝐒[obs_idx, end-nExo+1:end] .+= ∂jac_concat[:, (N-1)*nExo+1 : N*nExo] + # Blocks 1..N-1 are A * Sᵉ^(N-1-k) * B. + if N >= 2 + A = 𝐒[obs_idx, 1:n_pnf] + B = 𝐒[t⁻, end-nExo+1:end] + Sᵉ = 𝐒[t⁻, 1:n_pnf] + ∂A = zeros(size(A)) + ∂B = zeros(size(B)) + ∂Sᵉ = zeros(size(Sᵉ)) + for k in 1:(N-1) + p = N - 1 - k # power of Sᵉ + M = warmup_Sᵉ_powers[p+1] # Sᵉ^p (1-indexed) + ∂blk = ∂jac_concat[:, (k-1)*nExo+1 : k*nExo] + ∂A .+= ∂blk * (M * B)' + ∂B .+= (A * M)' * ∂blk + if p >= 1 + ∂M = A' * ∂blk * B' + for j in 0:p-1 + Sj = warmup_Sᵉ_powers[j+1] + Spj = warmup_Sᵉ_powers[p-j] # Sᵉ^(p-1-j) → index p-j + ∂Sᵉ .+= Sj' * ∂M * Spj' + end + end + end + ∂𝐒[obs_idx, 1:n_pnf] .+= ∂A + ∂𝐒[t⁻, end-nExo+1:end] .+= ∂B + ∂𝐒[t⁻, 1:n_pnf] .+= ∂Sᵉ + end + end + # ----- end warmup pullback -------------------------------------------- + + # 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]') + + if size(jacct, 1) == size(jacct, 2) + jacc_fact = ℒ.lu(jacct, check = false) + if !ℒ.issuccess(jacc_fact) + 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 + else + jacc_fact = ℒ.qr(jacct) + R = jacc_fact.R + if any(k -> R[k,k] == 0, axes(R, 1)) + 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 + end + + ℒ.ldiv!(λ[i], jacc_fact, x[i]) + + if !all(isfinite, λ[i]) + 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]) — only for i > presample_periods + if i > presample_periods + if i < size(data_in_deviations,2) + ∂x -= copy(x[i]) + else + ∂x += copy(x[i]) + end + end + + # logabsdets += ℒ.logabsdet(jacc ./ precision_factor)[1] — only for i > presample_periods + if i > presample_periods + ∂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 + else + ∂jacc = zero(jacc[i]) + 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]') + + if size(jacct, 1) == size(jacct, 2) + jacc_fact = ℒ.lu(jacct, check = false) + if !ℒ.issuccess(jacc_fact) + 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 + else + jacc_fact = ℒ.qr(jacct) + R = jacc_fact.R + if any(k -> R[k,k] == 0, axes(R, 1)) + 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 + end + + ℒ.ldiv!(λ[i], jacc_fact, x[i]) + + if !all(isfinite, λ[i]) + 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]) — only contributes for i > presample_periods + if i > presample_periods + if i < size(data_in_deviations,2) + ∂x -= copy(x[i]) + else + ∂x += copy(x[i]) + end + end + + # logabsdets += ℒ.logabsdet(jacc ./ precision_factor)[1] — only for i > presample_periods + if i > presample_periods + ∂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 + else + ∂jacc = zero(jacc[i]) + 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]) — only for i > presample_periods + if i > presample_periods + if i < size(data_in_deviations,2) + ∂x -= copy(x[i]) + else + ∂x += copy(x[i]) + end + end + + # logabsdets += ℒ.logabsdet(jacc ./ precision_factor)[1] — only for i > presample_periods + if i > presample_periods + ∂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 + else + ∂jacc = zero(jacc[i]) + 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]) — only for i > presample_periods + if i > presample_periods + if i < size(data_in_deviations,2) + ∂x -= copy(x[i]) + else + ∂x += copy(x[i]) + end + end + + # logabsdets += ℒ.logabsdet(jacc ./ precision_factor)[1] — only for i > presample_periods + if i > presample_periods + ∂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 + else + ∂jacc = zero(jacc[i]) + 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 + lyap_solved = true + 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_solved = lyap_rrule_result[2] + lyap_rrule_result[1] + else + get_initial_covariance(Val(initial_covariance), A, 𝐁, lyap_ws, opts = opts) + end + + if !lyap_solved + if opts.verbose println("KF initial Lyapunov solve failed") end + return on_failure_loglikelihood, x -> (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) + end + + Tt = size(data_in_deviations, 2) + 1 + + z = zeros(size(data_in_deviations, 1)) + ū = zeros(size(C,2)) + P̄ = deepcopy(P) + + temp_N_N = similar(P) + PCtmp = similar(P, size(P, 1), size(C, 1)) + F = similar(P, size(C, 1), size(C, 1)) + + u = [similar(ū) 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 # v[t] = data - C * u_predict + + ℒ.mul!(CP[t], C, P̄) # CP[t] = C * P + ℒ.mul!(F, CP[t], C') # F = CP[t] * C' = C * P * C' + + # Old way (≤v0.1.42): luF = lu(F) + kalman_ws.fast_lu_ws_f, kalman_ws.fast_lu_dims_f, solved_F, luF = factorize_lu!(Val(:FastLapack), 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 + + # Old way (≤v0.1.42): logabsdetF = log(abs(det(luF))) + 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) # invF[t] = F \ I + + if t - 1 > presample_periods + loglik += logabsdetF + ℒ.dot(v[t], invF[t], v[t]) # Old way: loglik += log(det(F)) + v' * inv(F) * v + end + + ℒ.mul!(PCtmp, P̄, C') # PCtmp = P * C' + copyto!(K[t], PCtmp) # K[t] = P * C' (before solving) + solve_lu_right!(F, K[t], kalman_ws.fast_lu_ws_f, luF, kalman_ws.fast_lu_rhs_t_k) # K[t] = P * C' / F + + ℒ.mul!(P_seq[t], K[t], CP[t], -1, 0) # P_seq[t] = -K[t] * CP[t] + P_seq[t] .+= P̄ # P_seq[t] = P - K[t] * C * P + + ℒ.mul!(temp_N_N, P_seq[t], A') # temp = P_seq[t] * A' + ℒ.mul!(P̄, A, temp_N_N) # P = A * P_seq[t] * A' + P̄ .+= 𝐁 # P = A * P_seq[t] * A' + B + + ℒ.mul!(u[t], K[t], v[t]) # u[t] = K[t] * v[t] + u[t] .+= ū # u[t] = K[t] * v[t] + u_predicted + + ℒ.mul!(ū, A, u[t]) # u_predict = A * u[t] + ℒ.mul!(z, C, ū) # z = C * u_predict + 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̄) + ∂ū = zero(ū) + ∂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!(∂ū, 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, ∂ū, v[t]') + ℒ.mul!(P_seq[1], PCtmp, CP[1]) + ℒ.mul!(∂P, A', P_seq[1], 1, 1) + + ℒ.mul!(u[1], A', ∂ū) + ℒ.mul!(v[1], K[t]', u[1]) + ℒ.axpy!(1, ∂v, v[1]) + ∂data_in_deviations[:,t-1] .= v[1] + + ℒ.mul!(u[1], A', ∂ū) + ℒ.mul!(v[1], K[t]', u[1]) + ℒ.mul!(∂ū, C', v[1]) + ℒ.mul!(u[1], C', v[1], -1, 1) + copy!(∂ū, u[1]) + + ℒ.mul!(u[1], C', ∂v) + ℒ.axpy!(-1, u[1], ∂ū) + + if t > 2 + ℒ.mul!(∂A_kf, ∂ū, 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', ∂ū) + ℒ.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_helper(Δ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[], + correlation::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_SELECTOR(𝓂), + 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[]) || !(correlation == 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[]) || !(correlation == 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 + corr_var_idx = parse_variables_input_to_index(correlation, 𝓂) + corr_groups = is_grouped_covariance_input(correlation) ? parse_covariance_groups(correlation, 𝓂.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[]) || !(correlation == 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[]) && (correlation == 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) + ŝ_to_ŝ₂ = zeros(T, 0, 0) + ŝ_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}[] + + corr_full_mat = zeros(T, 0, 0) + diag_C_corr = zeros(T, 0) + s_corr = zeros(T, 0) + corr_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_helper(Δ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, corr_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[union(covar_var_idx, corr_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[]) || !(correlation == 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, corr_var_idx)], + 𝓂; + covariance = 𝓂.constants.post_model_macro.var[union(covar_var_idx, corr_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[]) || !(correlation == 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] + ŝ_to_ŝ₂ = somc_out[6] + ŝ_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(ŝ_to_ŝ₂, 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 = ŝ_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 * ŝ_to_ŝ₂ + 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 + + if !(correlation == Symbol[]) + if size(covar_dcmp, 1) > 0 + corr_full_mat, _, diag_C_corr, s_corr = covariance_to_correlation(covar_dcmp) + end + + if !isnothing(corr_groups) + for group in corr_groups + for i in group + i_pos = findfirst(==(i), corr_var_idx) + isnothing(i_pos) && continue + for j in group + j_pos = findfirst(==(j), corr_var_idx) + isnothing(j_pos) && continue + push!(corr_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 !(correlation == Symbol[]) + if solved + if !isnothing(corr_groups) + corr_result = zeros(T, length(corr_var_idx), length(corr_var_idx)) + for (i_pos, j_pos, i, j) in corr_group_pairs + corr_result[i_pos, j_pos] = corr_full_mat[i, j] + end + ret[:correlation] = corr_result + else + ret[:correlation] = corr_full_mat[corr_var_idx, corr_var_idx] + end + else + ret[:correlation] = fill(Inf * sum(abs2,parameter_values), isnothing(corr_var_idx) ? 0 : length(corr_var_idx), isnothing(corr_var_idx) ? 0 : length(corr_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_Δcorr = 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_helper(Δret, :non_stochastic_steady_state), prev_Δnsss) + Δmean = incremental_cotangent!(get_statistics_cotangent_helper(Δret, :mean), prev_Δmean) + Δstd = incremental_cotangent!(get_statistics_cotangent_helper(Δret, :standard_deviation), prev_Δstd) + Δvar = incremental_cotangent!(get_statistics_cotangent_helper(Δret, :variance), prev_Δvar) + Δcov = incremental_cotangent!(get_statistics_cotangent_helper(Δret, :covariance), prev_Δcov) + Δcorr = incremental_cotangent!(get_statistics_cotangent_helper(Δret, :correlation), prev_Δcorr) + Δautocorr = incremental_cotangent!(get_statistics_cotangent_helper(Δ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_ŝ₂ = size(ŝ_to_ŝ₂, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(ŝ_to_ŝ₂)) + ∂ŝ_to_y₂ = size(ŝ_to_y₂, 1) == 0 ? zeros(T, 0, 0) : zeros(T, size(ŝ_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 !(Δcorr isa Union{NoTangent, AbstractZero}) && !(correlation == Symbol[]) + Δcorr_full = zeros(T, length(corr_var_idx), length(corr_var_idx)) + if !isnothing(corr_groups) + for (i_pos, j_pos, i, j) in corr_group_pairs + Δcorr_full[i_pos, j_pos] += Δcorr[i_pos, j_pos] + end + else + Δcorr_full .+= Δcorr + end + + @inbounds for a_pos in eachindex(corr_var_idx) + a = corr_var_idx[a_pos] + sa = s_corr[a] + isnan(sa) && continue + for b_pos in eachindex(corr_var_idx) + b = corr_var_idx[b_pos] + g = Δcorr_full[a_pos, b_pos] + g == 0 && continue + sb = s_corr[b] + isnan(sb) && continue + sasb = sa * sb + sasb == 0 && continue + corr_ab = corr_full_mat[a, b] + src_a = min(a, b) + src_b = max(a, b) + ∂covar_dcmp[src_a, src_b] += g / sasb + ∂covar_dcmp[a, a] += -g * corr_ab / (2 * diag_C_corr[a]) + ∂covar_dcmp[b, b] += -g * corr_ab / (2 * diag_C_corr[b]) + end + end + 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 = ŝ_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_ŝ₂ .+= second_order_P_seq[i]' * ∂P[i + 1] + ∂P[i] .+= ∂P[i + 1] * ŝ_to_ŝ₂' + 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_ŝ₂ : 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_SELECTOR(𝓂), + 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 check_bounds(parameters, 𝓂) + return get_solution_fail(algorithm, fill(S(-Inf), nVar), nVar, S), zero_pullback + 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) + result = get_solution_fail(algorithm, SS_and_pars[1:nVar], nVar, S) + 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 + result = get_solution_fail(algorithm, SS_and_pars[1:nVar], nVar, S, 𝐒₁) + 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: (NSSS, [𝐒₁, 𝐒₂], solved) + result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂_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] + ∂mats = unthunk(Δ[2]) # cotangent for Vector{AbstractMatrix{S}} + # Δ[3] is ∂solved — not differentiable + + # Extract per-matrix cotangents defensively + ∂𝐒₁_ext = if ∂mats isa Union{NoTangent, AbstractZero} + NoTangent() + else + m = unthunk(∂mats[1]) + m isa Union{NoTangent, AbstractZero} ? NoTangent() : m + end + ∂𝐒₂_ext = if ∂mats isa Union{NoTangent, AbstractZero} + NoTangent() + else + m = unthunk(∂mats[2]) + m isa Union{NoTangent, AbstractZero} ? NoTangent() : m + end + + # ── 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: (NSSS, [𝐒₁, 𝐒₂, 𝐒₃], solved) + result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂_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] + ∂mats = unthunk(Δ[2]) # cotangent for Vector{AbstractMatrix{S}} + # Δ[3] is ∂solved — not differentiable + + # Extract per-matrix cotangents defensively + ∂𝐒₁_ext = if ∂mats isa Union{NoTangent, AbstractZero} + NoTangent() + else + m = unthunk(∂mats[1]) + m isa Union{NoTangent, AbstractZero} ? NoTangent() : m + end + ∂𝐒₂_ext = if ∂mats isa Union{NoTangent, AbstractZero} + NoTangent() + else + m = unthunk(∂mats[2]) + m isa Union{NoTangent, AbstractZero} ? NoTangent() : m + end + ∂𝐒₃_ext = if ∂mats isa Union{NoTangent, AbstractZero} + NoTangent() + else + m = unthunk(∂mats[3]) + m isa Union{NoTangent, AbstractZero} ? NoTangent() : m + end + + # ── 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], AbstractMatrix{S}[𝐒₁], true) + + pullback_1st = function (∂result_bar) + Δ = unthunk(∂result_bar) + + if Δ isa Union{NoTangent, AbstractZero} + return NoTangent(), NoTangent(), zeros(S, length(parameters)) + end + + ∂NSSS = Δ[1] + ∂mats = unthunk(Δ[2]) # cotangent for Vector{AbstractMatrix{S}} + # Δ[3] is ∂solved — not differentiable + + # Extract ∂𝐒₁ defensively + ∂𝐒₁_ext = if ∂mats isa Union{NoTangent, AbstractZero} + NoTangent() + else + m = unthunk(∂mats[1]) + m isa Union{NoTangent, AbstractZero} ? NoTangent() : m + end + + # ── 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/steady_state/nsss_solver.jl b/src/steady_state/nsss_solver.jl new file mode 100644 index 000000000..c4272c4fb --- /dev/null +++ b/src/steady_state/nsss_solver.jl @@ -0,0 +1,2654 @@ +@stable default_mode = "disable" begin + +# 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 + +@unstable @inline function normalize_symbolic_solution(sol::SPyPyC.Sym{PythonCall.Core.Py}) + if sol.is_number == true + return sol + end + num, _ = sol.as_numer_denom() + return num.is_zero == true ? SPyPyC.Sym(0) : sol +end + +@unstable @inline function symbolic_solution_atoms(sol::SPyPyC.Sym{PythonCall.Core.Py}) + sol.is_number == true && return Symbol[] + atoms = Symbol[] + for a in sol.atoms() + if a.is_number != true + push!(atoms, Symbol(a)) + end + end + return atoms +end + +""" +Return `true` if the Julia `Expr`/`Symbol` contains a call that +`make_equation_robust_to_domain_errors` would rewrite with a `➕` auxiliary +variable. Used to gate the rewrite: when no hazard is present the rewrite is +a no-op and can be skipped entirely, saving allocations and SymPy simplify +calls on log-linear models. +""" +@unstable function expression_has_domain_hazards(expr) + hazard = false + postwalk(x -> begin + if x isa Expr && x.head == :call && length(x.args) >= 2 + op = x.args[1] + if op === :log || op === :sqrt || op === :exp || + op === :norminvcdf || op === :norminv || op === :qnorm || + op === :erfcinv + hazard = true + elseif op === :^ + # non-integer exponent triggers the rewrite path + if length(x.args) >= 3 && !(x.args[3] isa Int) + hazard = true + end + end + end + x + end, expr) + return hazard +end + +""" +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{Float64}}} + # 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{Float64}}[], + 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(lower_bounds)), + 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{Float64}, + 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 + +@unstable 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 + +@unstable 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 + +@unstable 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 Float64 ? + x : + x.args[2].head == :call ? + begin + if precompile + replacement = 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 = 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 = 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 = 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 = 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 + +@unstable 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 + +@unstable 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 + +@unstable 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 + + # Precomputed per-equation symbol sets used as a cheap Julia-side filter for + # SymPy solve/subs calls in the analytical branch. Aligned with ss_equations. + eq_symbol_sets = [Set{Symbol}(Symbol.(collect(e))) for e in eq_list] + + 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̂)' + + # Handle unmatched variables: those whose equations became redundant + # (e.g., unit-root variables that cancel in steady state like qlfpr[0] = hqlfpr[0] + qlfpr[-1]) + unmatched_mask = eqs[1,:] .< 0 + n_unmatched = sum(unmatched_mask) + unmatched_var_names = similar(unknowns, 0) + + if n_unmatched > 0 + unmatched_var_indices = vars[1, unmatched_mask] + unmatched_var_names = [unknowns[i] for i in unmatched_var_indices] + + if verbose + println("Indeterminate steady-state variables ($(n_unmatched)): assigned default value (0.0 or user guess)") + for vn in unmatched_var_names + guess_val = haskey(𝓂.constants.post_parameters_macro.guess, Symbol(vn)) ? + 𝓂.constants.post_parameters_macro.guess[Symbol(vn)] : 0.0 + println(" $(Symbol(vn)) = $guess_val") + end + end + + # Keep only matched columns + matched_mask = .!unmatched_mask + vars = vars[:, matched_mask] + eqs = eqs[:, matched_mask] + + # Renumber blocks: the matched entries retain their block numbers from BTF + # but we need contiguous block numbers for the processing loop + if !isempty(vars) + old_blocks = sort(unique(vars[2,:])) + block_remap = Dict(old => new for (new, old) in enumerate(old_blocks)) + for j in axes(vars, 2) + vars[2,j] = block_remap[vars[2,j]] + eqs[2,j] = block_remap[eqs[2,j]] + end + n_blocks = length(old_blocks) + else + n_blocks = 0 + end + end + + 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 + + # Emit analytical steps for unmatched (indeterminate) variables + if n_unmatched > 0 + for vn in unmatched_var_names + var_sym = Symbol(vn) + default_val = haskey(𝓂.constants.post_parameters_macro.guess, var_sym) ? + Float64(𝓂.constants.post_parameters_macro.guess[var_sym]) : 0.0 + + widx = sol_name_to_index[var_sym] + + eval_func! = let cv = default_val + (out, _sol_vec, _params_vec) -> begin + out[1] = cv + return nothing + end + end + + push!(solved_vars, var_sym) + push!(solved_vals, default_val) + push!(atoms_in_equations_list, []) + + push_analytical_step!(builder; + eval_func! = eval_func!, + write_indices = [widx], + description = "Indeterminate: $var_sym = $default_val", + ) + end + end + + 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(parsed_eq_to_solve_for) > 15 + soll = nothing + else + var_sym = Symbol(var_to_solve_for) + eq_idx_for_filter = eqs[:,eqs[2,:] .== n][1] + # Option A: skip SymPy solve when the variable is not a free symbol of the equation + if 1 <= eq_idx_for_filter <= length(eq_symbol_sets) && !(var_sym in eq_symbol_sets[eq_idx_for_filter]) + soll = nothing + elseif 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) + + else + sol_expr = normalize_symbolic_solution(soll[1]) + + if sol_expr.is_number == true + if var_to_solve_for isa SPyPyC.Sym{PythonCall.Core.Py} && sol_expr isa SPyPyC.Sym{PythonCall.Core.Py} + var_sym_for_sub = Symbol(var_to_solve_for) + # Only do the expensive SymPy substitution on equations that actually + # contain the variable (per the precomputed symbol sets). + ss_equations = [ + (eq isa SPyPyC.Sym{PythonCall.Core.Py} && + k <= length(eq_symbol_sets) && + (var_sym_for_sub in eq_symbol_sets[k])) ? + replace_symbolic(eq, var_to_solve_for, sol_expr) : eq + for (k, eq) in enumerate(ss_equations) + ] + end + + push!(solved_vars, Symbol(var_to_solve_for)) + push!(solved_vals, Meta.parse(string(sol_expr))) + 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(sol_expr) + 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(sol_expr))) + + for a in symbolic_solution_atoms(sol_expr) + push!(atoms_in_equations, a) + end + push!(atoms_in_equations_list, Set(union(setdiff(get_symbols(parsed_eq_to_solve_for), get_symbols(minmax_fixed_eqs)), symbolic_solution_atoms(sol_expr)))) + + 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[]] + + if expression_has_domain_hazards(val_expr) + 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) + else + rewritten_eqs = Union{Expr,Symbol}[val_expr] + ss_and_aux_equations = Expr[] + ss_and_aux_equations_dep = Expr[] + ss_and_aux_equations_error = Expr[] + ss_and_aux_equations_error_dep = Expr[] + end + + 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 + 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 + + normalized_solutions = Dict(v => normalize_symbolic_solution(soll[v]) for v in vars_to_solve) + for sol_expr in values(normalized_solutions), a in symbolic_solution_atoms(sol_expr) + push!(atoms_in_equations, a) + end + + step_exprs = [] + step_write_indices = Int[] + + for v in vars_to_solve + sol_expr = normalized_solutions[v] + push!(solved_vars, Symbol(v)) + push!(solved_vals, Meta.parse(string(sol_expr))) + push!(atoms_in_equations_list, Set(symbolic_solution_atoms(sol_expr))) + 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::CircularBuffer{Vector{Vector{Float64}}}, initial_parameters::Vector{Float64}, expected_length::Int)::Tuple{Float64, Vector{Vector{Float64}}} + current_best = Inf + closest_solution = cache[end] + + target_parameters_norm_squared = 0.0 + @turbo 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 + @turbo 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 + @turbo 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 + + +# ============================================================================ +# Block solver helpers (moved from MacroModelling.jl — used only in NSSS pipeline) + +@unstable function replace_symbolic(equation::SPyPyC.Sym{PythonCall.Core.Py}, variable::SPyPyC.Sym{PythonCall.Core.Py}, replacement::SPyPyC.Sym{PythonCall.Core.Py})::SPyPyC.Sym{PythonCall.Core.Py} + # equation.subs(variable, replacement) + return SPyPyC.subs(equation, variable, replacement) +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 + + +function solve_ss(SS_optimizer::F, + SS_solve_block::ss_solve_block{Float64}, + parameters_and_solved_vars::Vector{T}, + closest_parameters_and_solved_vars::Vector{T}, + lbs::Vector{T}, + ubs::Vector{T}, + tol::Tolerances, + total_iters::Vector{Int}, + n_block::Int, + verbose::Bool, + guess::Vector{T}, + solver_params::solver_parameters, + extended_problem::Bool, + separate_starting_value::T)::Tuple{Vector{T}, Vector{Int}, T, T} where {F, T <: AbstractFloat} + ftol = tol.nsss.ftol + n_guess = length(guess) + init_buf = SS_solve_block.ss_problem.workspace.best_previous_guess + use_ssv = !isnan(separate_starting_value) + ssv_val = use_ssv ? 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 + init_buf + end + + sol_new_tmp, info = SS_optimizer( extended_problem ? SS_solve_block.extended_ss_problem : SS_solve_block.ss_problem, + # if extended_problem + # function ext_function_to_optimize(guesses) + # gss = guesses[1:length(guess)] + + # parameters_and_solved_vars_guess = guesses[length(guess)+1:end] + + # res = ss_solve_blocks(parameters_and_solved_vars, gss) + + # return vcat(res, parameters_and_solved_vars .- parameters_and_solved_vars_guess) + # end + # else + # function function_to_optimize(guesses) ss_solve_blocks(parameters_and_solved_vars, guesses) end + # end + + # sol_new_tmp, info = SS_optimizer( extended_problem ? ext_function_to_optimize : function_to_optimize, + optimizer_init, + parameters_and_solved_vars, + 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_minimum = info[4] # isnan(sum(abs, info[4])) ? Inf : ℒ.norm(info[4]) + + rel_sol_minimum = info[3] + + 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] + + if sol_minimum < ftol && verbose + extended_problem_str = extended_problem ? "(extended problem) " : "" + + if isnan(separate_starting_value) + starting_value_str = "" + else + starting_value_str = "and starting point: $separate_starting_value" + 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 + + if all_small_guess && isnan(separate_starting_value) + any_guess_str = "previous solution, " + elseif has_small_guess && isnan(separate_starting_value) + 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) + + 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 +end + + +function block_solver(parameters_and_solved_vars::Vector{T}, + n_block::Int, + SS_solve_block::ss_solve_block{Float64}, + # SS_optimizer, + # f::OptimizationFunction, + guess_and_pars_solved_vars::Vector{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 ; + tol::Tolerances = Tolerances(), + # rtol::AbstractFloat = sqrt(eps()), + # timeout = 120, + # starting_points::Vector{Float64} = [1.205996189998029, 0.7688, 0.897, 1.2],#, 0.9, 0.75, 1.5, -0.5, 2.0, .25] + # verbose::Bool = false + )::Tuple{Vector{T},Tuple{T, Int}} where T <: AbstractFloat + + # tol = parameters[1].ftol + # rtol = parameters[1].rel_xtol + + solved_yet = false + + guess = guess_and_pars_solved_vars[1] + + sol_values = guess + + closest_parameters_and_solved_vars = sum(abs, guess_and_pars_solved_vars[2]) == Inf ? parameters_and_solved_vars : guess_and_pars_solved_vars[2] + + # res = ss_solve_blocks(parameters_and_solved_vars, guess) + + SS_solve_block.ss_problem.func(SS_solve_block.ss_problem.workspace.func_buffer, guess, parameters_and_solved_vars) + + res = SS_solve_block.ss_problem.workspace.func_buffer + + sol_minimum = ℒ.norm(res) + + # Handle removable singularities: if residual contains NaN (from 0/0 forms + # in symbolically simplified equations) but all finite residuals are below + # tolerance, treat NaN entries as zero. This occurs when remove_redundant_SS_vars! + # substitutes variables into equations creating indeterminate forms at the solution. + # Only apply this when we have an actual cached solution for these parameters + # (i.e. guess_and_pars_solved_vars[2] is not the all-Inf placeholder); otherwise + # the guess is just a clamped default that may coincidentally produce a NaN+zeros + # residual (e.g. 1 - R*beta/Pi at [Pi=0, R=0] gives NaN while companion equations + # evaluate to ~0), and accepting it would yield a spurious "solution" of all zeros. + has_cached_solution = sum(abs, guess_and_pars_solved_vars[2]) != Inf + if has_cached_solution && isnan(sol_minimum) + has_nonnan_violation = false + for i in eachindex(res) + v = res[i] + if !isnan(v) && abs(v) > tol.nsss.acceptance_tol + has_nonnan_violation = true + break + end + end + if !has_nonnan_violation + # All finite residuals are zero — NaN entries are removable singularities + for i in eachindex(res) + if isnan(res[i]) + res[i] = zero(T) + end + end + sol_minimum = ℒ.norm(res) + end + end + + if !cold_start + 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) + + SS_solve_block.ss_problem.jac(SS_solve_block.ss_problem.workspace.jac_buffer, guess, parameters_and_solved_vars) + + ∇ = SS_solve_block.ss_problem.workspace.jac_buffer + + 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 = one(T) + else + new_guess = guess - guess_update + rel_sol_minimum = ℒ.norm(guess_update) / max(ℒ.norm(new_guess), sol_minimum) + end + else + rel_sol_minimum = one(T) + end + else + rel_sol_minimum = zero(T) + end + else + rel_sol_minimum = one(T) + end + + if isfinite(sol_minimum) && sol_minimum < tol.nsss.acceptance_tol + solved_yet = true + + if verbose + println("Block: $n_block, - Solved using previous solution; residual norm: $sol_minimum") + end + 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(x -> x < T(1e12), guess) ? [guess, fill(T(1e12), length(guess))] : [guess] # if guess were provided, loop over them, and then the starting points only + start_vals = fail_fast_solvers_only ? (T(NaN),) : (T(NaN), T(0.0), T(1.206), T(1.5), T(0.7688), T(2.0), T(0.897)) + for g in guesses + 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 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, + # sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(SS_optimizer, ss_solve_blocks, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, total_iters, n_block, verbose, + g, + p, + ext, + s) + + if isfinite(sol_minimum) && sol_minimum < tol.nsss.acceptance_tol + solved_yet = true + end + end + end + end + end + end + else !cold_start + + start_vals = Vector{T}(undef, 8) + start_vals[1] = T(NaN) + start_vals[3] = T(0.0) + start_vals[4] = T(1.206) + start_vals[5] = T(1.5) + start_vals[6] = T(0.7688) + start_vals[7] = T(2.0) + start_vals[8] = 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, + # sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(algo, ss_solve_blocks, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, + total_iters, + n_block, + false, # verbose + guess, + p, + # parameters[1], + false, # ext + # false) + s) + 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 + + + # if sol_minimum > tol# || rel_sol_minimum > rtol + # for p in unique(parameters)#[1:3] # take unique because some parameters might appear more than once + # # for s in [p.starting_value, 1.206, 1.5, 0.7688, 2.0, 0.897]#, .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 + # if sol_minimum > tol# || rel_sol_minimum > rtol + # sol_values, total_iters, rel_sol_minimum, sol_minimum = solve_ss(SS_optimizer, ss_solve_blocks, parameters_and_solved_vars, closest_parameters_and_solved_vars, lbs, ubs, tol, total_iters, n_block, + # false, # verbose + # guess, + # p, + # false, + # false) + # # s) + # if !solved_yet && sol_minimum < tol# || rel_sol_minimum > rtol) + # solved_yet = true + # if verbose + # loop1 = unique(parameters)#[1:3] + # loop2 = [p.starting_value, 1.206, 1.5, 0.7688, 2.0, 0.897] + # p_in_loop1 = findfirst(x -> x == p, loop1) + # s_in_loop2 = findfirst(x -> x == s, loop2) + # if p_in_loop1 isa Nothing + # p_in_loop1 = 1 + # end + # if s_in_loop2 isa Nothing + # s_in_loop2 = 1 + # end + # n1 = (p_in_loop1 - 1) * length(loop2) + s_in_loop2 + # println("Block: $n_block, - Solved with modified Levenberg-Marquardt - $n1/$(length(loop2) *length(loop1)) - $sol_minimum - $rel_sol_minimum - $total_iters") + # end + # end + # end + # # end + # # end + # end + # end + end + + if verbose + if !solved_yet + println("Block: $n_block, - Solution not found after $(total_iters[1]) gradient evaluations and $(total_iters[2]) function evaluations; reltol: $rel_sol_minimum - tol: $sol_minimum") + end + end + + return sol_values, (sol_minimum, total_iters[1]) +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::Vector{Vector{Float64}}, 𝓂::ℳ, tol::Tolerances, fail_fast_solvers_only::Bool, + cold_start::Bool, solver_parameters::Vector{solver_parameters}, preferred_solver_parameter_idx::Int, verbose::Bool)::Tuple{Float64, Int, Vector{Vector{Float64}}} + + 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 !isfinite(raw) + raw = sol_vec[widx] + end + 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 + v = cache_sol[i] + if !isfinite(v) + # No prior cached solution; use sol_vec value (starts at 0.0, filled by earlier steps) + v = sol_vec[c.write_indices[wr[i]]] + end + guess_buf[i] = clamp(v, 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 = [ + eltype(sol) === Float64 ? copy(sol) : Float64.(primal.(sol)), + eltype(params_and_solved_vars) === Float64 ? copy(params_and_solved_vars) : Float64.(primal.(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::Vector{Vector{Float64}}, + cold_start::Bool, + solver_params::Vector{solver_parameters}, + preferred_solver_parameter_idx::Int +)::Tuple{Vector{Float64}, Tuple{Float64, Int}, Vector{Vector{Float64}}} + 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 + + @inbounds for i in 1:n_output + SS_and_pars[i] = sol_vec[nsss_output_indices[i]] + end + + if solution_error < tol.nsss.acceptance_tol + if any(x -> !isfinite(x), SS_and_pars) + solution_error = Inf + elseif isempty(𝓂.constants.post_model_macro.➕_vars) + # Cross-check against raw model equations only when no ➕ domain-safety + # rewrites were applied. When ➕ vars exist the step solver evaluates + # max(eps(),x)-substituted equations while NSSS_check evaluates the raw + # model equations — the residuals can legitimately differ (NaN, Inf, or + # large finite values from log/sqrt/power at the same solution point). + residual = nsss_ws.check_residual + fill!(residual, 0.0) + 𝓂.functions.NSSS_check(residual, parameters, SS_and_pars) + residual_error = ℒ.norm(residual) + if isfinite(residual_error) && residual_error > solution_error + solution_error = residual_error + end + end + end + + if solution_error >= tol.nsss.acceptance_tol + fill!(SS_and_pars, 0.0) + 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{Float64}, Tuple{Float64, 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{Float64}, Tuple{Float64, 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 : + Float64.(primal.(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] + closest_params = closest_solution_init[end] + @turbo for i in eachindex(initial_parameters) + scaled_parameters[i] = scale * initial_parameters[i] + (1 - scale) * closest_params[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 + +end # @stable \ No newline at end of file diff --git a/src/solver_parameters.jl b/src/steady_state/solver_parameters.jl similarity index 100% rename from src/solver_parameters.jl rename to src/steady_state/solver_parameters.jl diff --git a/src/steady_state/stochastic_steady_state.jl b/src/steady_state/stochastic_steady_state.jl new file mode 100644 index 000000000..fd12dc63e --- /dev/null +++ b/src/steady_state/stochastic_steady_state.jl @@ -0,0 +1,570 @@ +@stable default_mode = "disable" begin + + +function prepare_stochastic_steady_state_base_terms(parameters::Vector{M}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true)::Tuple{Bool, Vector{M}, Vector{M}, M, Matrix{M}, SparseMatrixCSC{M, Int}, Matrix{M}, SparseMatrixCSC{M, Int}, Vector{M}, constants} where M + C = initialise_constants!(𝓂) + T = C.post_model_macro + + SS_and_pars, (solution_error, iters) = get_NSSS_and_parameters(𝓂, parameters, opts = opts, estimation = estimation, caching = caching) + + if solution_error > opts.tol.nsss.acceptance_tol || isnan(solution_error) + return (false, + zeros(M, 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), + C) + end + + ensure_model_structure_constants!(C, 𝓂.equations.calibration_parameters) + ms = C.post_complete_parameters + all_SS = expand_steady_state(SS_and_pars, ms) + + ∇₁ = calculate_jacobian(parameters, SS_and_pars, 𝓂.caches, 𝓂.functions.jacobian, 𝓂.workspaces, caching = caching) + + 𝐒₁, qme_sol, solved = calculate_first_order_solution(∇₁, + C, + 𝓂.workspaces, + 𝓂.caches; + opts = opts, + initial_guess = 𝓂.caches.qme_solution, + parameter_values = parameters, + caching = caching) + + update_perturbation_counter!(𝓂.counters, solved, estimation = estimation, order = 1) + + if !solved + if opts.verbose println("1st order solution 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), + C) + end + + ∇₂ = calculate_hessian(parameters, SS_and_pars, 𝓂.caches, 𝓂.functions.hessian, 𝓂.workspaces, caching = caching) + + 𝐒₂_raw_untyped, solved2 = calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, 𝓂.constants, 𝓂.workspaces, 𝓂.caches; + initial_guess = 𝓂.caches.second_order_solution, + opts = opts, + parameter_values = parameters, + caching = caching) + + 𝐒₂_raw = sparse(𝐒₂_raw_untyped)::SparseMatrixCSC{M, Int} + + update_perturbation_counter!(𝓂.counters, solved2, estimation = estimation, order = 2) + + if !solved2 + if opts.verbose println("2nd order solution 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), + C) + end + + 𝐒₂ = (𝐒₂_raw * 𝓂.constants.second_order.𝐔₂)::SparseMatrixCSC{M, Int} + + 𝐒₁ = [𝐒₁[:,1:T.nPast_not_future_and_mixed] zeros(M, T.nVars) 𝐒₁[:,T.nPast_not_future_and_mixed+1:end]] + + aug_state₁ = sparse([zeros(M, T.nPast_not_future_and_mixed); one(M); zeros(M, T.nExo)]) + tmp = collect(T.I_nPast - 𝐒₁[T.past_not_future_and_mixed_idx,1:T.nPast_not_future_and_mixed]) + rhs = collect((𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2)[T.past_not_future_and_mixed_idx]) + + if M === Float64 + ensure_sss_tmp_lu_buffer!(𝓂.workspaces.second_order, tmp, rhs) + tmp_sol = 𝒮.solve!(𝓂.workspaces.second_order.sss_tmp_lu_buffer) + + if tmp_sol.retcode != 𝒮.SciMLBase.ReturnCode.Default && !𝒮.SciMLBase.successful_retcode(tmp_sol.retcode) + 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), + C) + end + + SSSstates = collect(tmp_sol.u) + else + tmp̄ = ℒ.lu(tmp, check = false) + + 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), + C) + end + + SSSstates = collect(tmp̄ \ rhs) + end + + SSSstates = SSSstates::Vector{M} + + return (true, + all_SS, + SS_and_pars, + solution_error, + ∇₁, + ∇₂, + 𝐒₁, + 𝐒₂_raw, + SSSstates, + C) +end + +function calculate_stochastic_steady_state(::Val{:second_order}, + parameters::Vector{M}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true)::Tuple{Vector{M}, Bool, Vector{M}, M, Matrix{M}, SparseMatrixCSC{M, Int}, Matrix{M}, SparseMatrixCSC{M, Int}} 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 && M === Float64 𝓂.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 + + # Expand compressed 𝐒₂_raw to full + 𝐒₂ = (𝐒₂_raw * 𝓂.constants.second_order.𝐔₂)::SparseMatrixCSC{M, Int} + + 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⁺] + + SSSstates, converged = solve_stochastic_steady_state_newton(Val(:second_order), 𝐒₁, 𝐒₂, collect(SSSstates), 𝓂) + + if !converged + if opts.verbose println("SSS not found") end + if caching && M === Float64 𝓂.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 + + state = A * SSSstates + B̂ * ℒ.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 + result = all_SS + Vector{M}(state) + + if caching && M === Float64 + 𝓂.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)::Tuple{Vector{M}, Bool, Vector{M}, M, Matrix{M}, SparseMatrixCSC{M, Int}, Matrix{M}, SparseMatrixCSC{M, Int}} 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 + + 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 && M === Float64 𝓂.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 + + # Expand compressed 𝐒₂_raw to full + 𝐒₂ = (𝐒₂_raw * 𝓂.constants.second_order.𝐔₂)::SparseMatrixCSC{M, Int} + + T = 𝓂.constants.post_model_macro + aug_state₁ = sparse([zeros(M, T.nPast_not_future_and_mixed); one(M); zeros(M, T.nExo)]) + state = 𝐒₁[:,1:T.nPast_not_future_and_mixed] * SSSstates + + 𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2 + + result = all_SS + Vector{M}(state) + + if caching && M === Float64 + 𝓂.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 solve_stochastic_steady_state_newton(::Val{:second_order}, + 𝐒₁::Matrix{R}, + 𝐒₂::AbstractSparseMatrix{R}, + x::Vector{R}, + 𝓂::ℳ; + tol::AbstractFloat = 1e-14)::Tuple{Vector{R}, Bool} where R <: AbstractFloat + # @timeit_debug timer "Setup matrices" 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 = 𝐒₂[T.past_not_future_and_mixed_idx,kron_s⁺_s] + B̂ = 𝐒₂[T.past_not_future_and_mixed_idx,kron_s⁺_s⁺] + + max_iters = 100 + # SSS .= 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 𝐒₃ * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 + + ℂ = 𝓂.workspaces.second_order + nPast = length(x) + ensure_sss_kron_buffers!(ℂ, nPast; third_order=false) + x_aug = ℂ.x_aug_buf + x_aug[end] = one(R) + kron_x_aug_xx = ℂ.kron_x_aug_xx + kron_x_aug_I = ℂ.kron_x_aug_I + + for i in 1:max_iters + copyto!(x_aug, 1, x, 1, nPast) + + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ∂x = (A + B * kron_x_aug_I - I_nPast) + + ℒ.kron!(kron_x_aug_xx, x_aug, x_aug) + x̂ = A * x + B̂ * kron_x_aug_xx / 2 + + Δx = x̂ - x + ensure_dx_lu_buffer!(ℂ, ∂x, Δx) + sol = 𝒮.solve!(ℂ.dx_lu_buffer) + + if sol.retcode != 𝒮.SciMLBase.ReturnCode.Default && !𝒮.SciMLBase.successful_retcode(sol.retcode) + return x, false + end + copyto!(Δx, sol.u) + + if i > 3 && isapprox(x̂, x, rtol = tol) + break + end + + # x += Δx + ℒ.axpy!(-1, Δx, x) + end + + # end # timeit_debug + + copyto!(x_aug, 1, x, 1, nPast) + ℒ.kron!(kron_x_aug_xx, x_aug, x_aug) + return x, isapprox(A * x + B̂ * kron_x_aug_xx / 2, x, rtol = tol) +end + + + + + +function calculate_stochastic_steady_state(::Val{:third_order}, + parameters::Vector{M}, + 𝓂::ℳ; + opts::CalculationOptions = merge_calculation_options(), + estimation::Bool = false, + caching::Bool = true)::Tuple{Vector{M}, Bool, Vector{M}, M, Matrix{M}, SparseMatrixCSC{M, Int}, SparseMatrixCSC{M, Int}, Matrix{M}, SparseMatrixCSC{M, Int}, SparseMatrixCSC{M, Int}} 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 + + 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 && M === Float64 𝓂.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 + + # Expand compressed 𝐒₂_raw to full + 𝐒₂ = (𝐒₂_raw * 𝓂.constants.second_order.𝐔₂)::SparseMatrixCSC{M, Int} + + ∇₃ = 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 && M === Float64 𝓂.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 length(𝓂.workspaces.third_order.Ŝ) == 0 || !(eltype(𝐒₃) == eltype(𝓂.workspaces.third_order.Ŝ)) + 𝓂.workspaces.third_order.Ŝ = 𝐒₃ * 𝓂.constants.third_order.𝐔₃ + else + ℒ.mul!(𝓂.workspaces.third_order.Ŝ, 𝐒₃, 𝓂.constants.third_order.𝐔₃) + end + + Ŝ = 𝓂.workspaces.third_order.Ŝ + 𝐒₃̂ = sparse_preallocated!(Ŝ, ℂ = 𝓂.workspaces.third_order)::SparseMatrixCSC{M, Int} + + so = 𝓂.constants.second_order + 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⁺] + Ĉ = 𝐒₃̂[:,kron_s⁺_s⁺_s⁺] + + SSSstates, converged = solve_stochastic_steady_state_newton(Val(:third_order), 𝐒₁, 𝐒₂, 𝐒₃̂, collect(SSSstates), 𝓂) + + if !converged + if opts.verbose println("SSS not found") end + if caching && M === Float64 𝓂.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 + + state = A * SSSstates + B̂ * ℒ.kron(vcat(SSSstates,1), vcat(SSSstates,1)) / 2 + Ĉ * ℒ.kron(vcat(SSSstates,1), ℒ.kron(vcat(SSSstates,1), vcat(SSSstates,1))) / 6 + + + result = all_SS + Vector{M}(state) + + if caching && M === Float64 + 𝓂.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)::Tuple{Vector{M}, Bool, Vector{M}, M, Matrix{M}, SparseMatrixCSC{M, Int}, SparseMatrixCSC{M, Int}, Matrix{M}, SparseMatrixCSC{M, Int}, SparseMatrixCSC{M, Int}} 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 + + 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 && M === Float64 𝓂.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 + + # Expand compressed 𝐒₂_raw to full + 𝐒₂ = (𝐒₂_raw * 𝓂.constants.second_order.𝐔₂)::SparseMatrixCSC{M, Int} + + ∇₃ = 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 && M === Float64 𝓂.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.Ŝ) == 0 || !(eltype(𝐒₃) == eltype(𝓂.workspaces.third_order.Ŝ)) + 𝓂.workspaces.third_order.Ŝ = 𝐒₃ * 𝓂.constants.third_order.𝐔₃ + else + ℒ.mul!(𝓂.workspaces.third_order.Ŝ, 𝐒₃, 𝓂.constants.third_order.𝐔₃) + end + + Ŝ = 𝓂.workspaces.third_order.Ŝ + 𝐒₃̂ = sparse_preallocated!(Ŝ, ℂ = 𝓂.workspaces.third_order)::SparseMatrixCSC{M, Int} + + T = 𝓂.constants.post_model_macro + aug_state₁ = sparse([zeros(M, T.nPast_not_future_and_mixed); one(M); zeros(M, T.nExo)]) + state = 𝐒₁[:,1:T.nPast_not_future_and_mixed] * SSSstates + 𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2 + + result = all_SS + Vector{M}(state) + + if caching && M === Float64 + 𝓂.caches.pruned_third_order_stochastic_steady_state = result + 𝓂.caches.valid_for.pruned_third_order_stochastic_steady_state = Float64.(parameters) + end + + return result, true, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃̂ +end + + +function solve_stochastic_steady_state_newton(::Val{:third_order}, + 𝐒₁::Matrix{Float64}, + 𝐒₂::AbstractSparseMatrix{Float64}, + 𝐒₃::AbstractSparseMatrix{Float64}, + x::Vector{Float64}, + 𝓂::ℳ; + tol::AbstractFloat = 1e-14)::Tuple{Vector{Float64}, Bool} + # 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 + + ℂ = 𝓂.workspaces.third_order + nPast = length(x) + ensure_sss_kron_buffers!(ℂ, nPast; third_order=true) + x_aug = ℂ.x_aug_buf + x_aug[end] = 1.0 + kron_x_aug = ℂ.kron_x_aug_xx + kron_x_kron = ℂ.kron_x_aug_x_kron + kron_x_aug_I = ℂ.kron_x_aug_I + kron_x_kron_I = ℂ.kron_x_kron_I + + for i in 1:max_iters + copyto!(x_aug, 1, x, 1, nPast) + ℒ.kron!(kron_x_aug, x_aug, x_aug) + ℒ.kron!(kron_x_kron, x_aug, kron_x_aug) + + ℒ.kron!(kron_x_aug_I, x_aug, I_nPast) + ℒ.kron!(kron_x_kron_I, kron_x_aug, I_nPast) + ∂x = (A + B * kron_x_aug_I + C * kron_x_kron_I / 2 - I_nPast) + + Δx = (A * x + B̂ * kron_x_aug / 2 + Ĉ * kron_x_kron / 6 - x) + ensure_dx_lu_buffer!(ℂ, ∂x, Δx) + sol = 𝒮.solve!(ℂ.dx_lu_buffer) + + if sol.retcode != 𝒮.SciMLBase.ReturnCode.Default && !𝒮.SciMLBase.successful_retcode(sol.retcode) + return x, false + end + copyto!(Δx, sol.u) + + 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, nPast) + ℒ.kron!(kron_x_aug, x_aug, x_aug) + ℒ.kron!(kron_x_kron, x_aug, kron_x_aug) + return x, isapprox(A * x + B̂ * kron_x_aug / 2 + Ĉ * kron_x_kron / 6, x, rtol = tol) +end + + +end # @stable diff --git a/src/structures.jl b/src/structures.jl index 62659cc92..f392ec349 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 # # ============================================================================= @@ -111,8 +112,23 @@ mutable struct equations calibration::Vector{Expr} calibration_no_var::Vector{Expr} calibration_parameters::Vector{Symbol} + calibration_original::Vector{Expr} end +""" +`RevisionEntry` records a single equation-modification event on a model. + +Fields: +- `timestamp` : when the change was applied +- `action` : one of `:update_equation`, `:add_equation`, `:remove_equation`, + `:update_calibration_equation`, `:add_calibration_equation`, `:remove_calibration_equation` +- `equation_index`: 1-based index of the affected equation (or `nothing`) +- `old_equation` : equation before the change (or `nothing` for add) +- `new_equation` : equation after the change (or `nothing` for remove) +""" +const RevisionEntry = NamedTuple{(:timestamp, :action, :equation_index, :old_equation, :new_equation), + Tuple{Dates.DateTime, Symbol, Union{Int, Nothing}, Union{Expr, Nothing}, Union{Expr, Nothing}}} + struct post_model_macro max_obc_horizon::Int # present_only::Vector{Symbol} @@ -146,6 +162,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 +256,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 +275,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 +293,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) @@ -306,8 +340,9 @@ mutable struct second_order_indices # Filled by ensure_conditional_forecast_constants! (options_and_caches.jl) # Triggered by: get_conditional_forecast, find_shocks # ========================================================================= - var²_idxs::Vector{Int} # Variable² indices + var²_idxs::Vector{Int} # Variable² indices (no-vol: kron(s_in_s, s_in_s)) shockvar²_idxs::Vector{Int} # Shock × variable² indices + shockvar_no_vol_idxs::Vector{Int} # Shock-variable cross indices (no-vol: kron(e_in_s⁺, s_in_s)) # ========================================================================= # MOMENT COMPUTATION CONSTANTS (model-constant values for moments.jl) @@ -351,7 +386,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 +402,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,8 +523,21 @@ 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} + # Doubling power cache (for AD: reuse A^(2^k), B^(2^k) sequences across forward/pullback) + # 𝐀_pow[k] = A^(2^(k-1)) ; 𝐁_pow[k] = B^(2^(k-1)). Only valid when pow_iters > 0. + # Fields are AbstractMatrix so the cache can hold dense Matrix or sparse SparseMatrixCSC entries + # (the doubling overloads dispatched by Sylvester preserve the sparsity of A and B across squaring). + 𝐀_pow::Vector{AbstractMatrix{G}} + 𝐁_pow::Vector{AbstractMatrix{G}} + pow_iters::Int # number of valid entries in 𝐀_pow / 𝐁_pow + pow_capture::Bool # true while solver should populate the cache + pow_transposed::Bool # true when 𝐀_pow / 𝐁_pow store transposes as Adjoint views (for adjoint use) + # ForwardDiff partials buffers (for forward-mode AD) P̃::Matrix{H} # For sylvester equation partials Ã_fd::Matrix{H} # Temporary for ForwardDiff partials of A @@ -483,25 +547,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} # Ã₊ + 𝐀̃₀::Matrix{T} # Ã₀ + 𝐀̃₋::Matrix{T} # Ã₋ + 𝐀̄₀ᵤ::Matrix{T} # Ā₀ᵤ + 𝐀₊ᵤ::Matrix{T} # A₊ᵤ + 𝐀̃₀ᵤ::Matrix{T} # Ã₀ᵤ + 𝐀₋ᵤ::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 +620,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 @@ -552,7 +709,7 @@ Fields for doubling algorithm: - `𝐂A`: Temporary for C * A' - `𝐀²`: Temporary for A * A -Fields for Krylov methods (bicgstab, gmres): +Fields for Krylov methods (bicgstab, gmres, dqgmres): - `tmp̄`: Temporary matrix for linear operator - `𝐗`: Reshape buffer for solution vector - `b`: RHS vector for Krylov solver @@ -577,98 +734,262 @@ 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}} + dqgmres::Krylov.DqgmresWorkspace{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}} + dqgmres_vech::Krylov.DqgmresWorkspace{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 Ã_fd::Matrix{R} # Temporary for ForwardDiff partials of A C̃_fd::Matrix{R} # Temporary for ForwardDiff partials of C + + # Doubling power cache (for AD: reuse A^(2^k) sequence across forward/pullback) + # 𝐀_pow[k] = A^(2^(k-1)). Valid only when pow_iters > 0. + # Slot type is AbstractMatrix so dense and sparse iterations + # share the same storage (sparse-aware capture). + 𝐀_pow::Vector{AbstractMatrix{T}} + pow_iters::Int # number of valid entries in 𝐀_pow + pow_capture::Bool # true while solver should populate the cache + pow_transposed::Bool # true when 𝐀_pow stores transposes as Adjoint views (for adjoint use) + + # FastLapackInterface Schur workspace for unit-root deflation (lazily resized) + schur_ws::FastLapackInterface.SchurWs{T} end -struct ss_solve_block - ss_problem::function_and_jacobian - extended_ss_problem::function_and_jacobian +struct ss_solve_block{T <: Real} + ss_problem::function_and_jacobian{T} + extended_ss_problem::function_and_jacobian{T} 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{Float64}}} # 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{Float64}}[], +) + +"""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} + covariance_first_order::Vector{Float64} + covariance_second_order::Vector{Float64} + covariance_third_order::Vector{Float64} + covariance_third_order_obs_key::Vector{Int} + covariance_third_order_autocorr::Vector{Float64} + covariance_third_order_autocorr_obs_key::Vector{Int} + covariance_third_order_autocorr_periods::Vector{Int} +end + + +valid_for_caches() = valid_for_caches( + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Float64[], + Int[], + Float64[], + Int[], + Int[], +) + + """ 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 +1002,24 @@ 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 + + # ========================================================================= + # LYAPUNOV SOLUTION CACHES + # Covariance matrices from Lyapunov equation solves + # ========================================================================= + covariance_first_order::Matrix{Float64} # 1st order Lyapunov solution + covariance_second_order::Matrix{Float64} # 2nd order Lyapunov solution + covariance_third_order::Matrix{Float64} # 3rd order assembled Σʸ₃ + covariance_third_order_autocorr::Matrix{Float64} # 3rd order autocorrelation matrix + + # ========================================================================= + # UNIT ROOT FLAGS + # Detected during QME solve, used to skip failed Lyapunov solver attempts + # ========================================================================= + has_unit_roots::Bool # true if state transition has eigenvalues ≈ 1 end # Structs for perturbation derivative functions (used for AD) @@ -707,28 +1043,19 @@ 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 + obc_constraint_info::Vector{Tuple{Int, Int, Float64}} # Whether all functions have been written/compiled functions_written::Bool end @@ -754,6 +1081,15 @@ mutable struct find_shocks_workspace{T <: Real} kron_buffer²::Vector{T} # n_exo^3 - for ℒ.kron(x, kron_buffer) kron_buffer3::Matrix{T} # n_exo^3 × n_exo - for ℒ.kron(J, kron_buffer) kron_buffer4::Matrix{T} # n_exo^3 × n_exo^2 - for ℒ.kron(kron(J,J), x) + + # Conditional-forecast (find_shocks_conditional_forecast) state-related kron buffers + n_past::Int + kron_state_vol::Vector{T} # (n_past+1)^2 - for ℒ.kron(state_vol, state_vol) + kron_state_vol3::Vector{T} # (n_past+1)^3 - for ℒ.kron(state_vol, kron_state_vol) (3rd order) + kron_state₁₂::Vector{T} # n_past^2 - for ℒ.kron(state₁, state₂) (3rd order pruned) + kron_I_state::Matrix{T} # n_exo*(n_past+1) × n_exo - for ℒ.kron(J, state_vol) + kron_I_state₂::Matrix{T} # n_exo*n_past × n_exo - for ℒ.kron(J, state₂) (3rd order pruned) + kron_I_state_state::Matrix{T} # n_exo*(n_past+1)^2 × n_exo - for ℒ.kron(J, kron_state_vol) (3rd order) end @@ -787,6 +1123,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 +1158,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 +1173,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 +1199,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} Ŝ::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 +1218,55 @@ 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} + # LinearSolve cache (FastLUFactorization backend) for SSS Newton iter ∂x \ Δx + # Used by primal, rrule forward loop, and ForwardDiffExt (which strips Duals to Float64 first). + # Safe to share across calls because the cache is only mutated during forward Newton iteration; + # rrule pullback does not access this buffer (it builds a fresh LU on a different matrix). + dx_lu_buffer::𝒮.LinearCache + # LinearSolve cache (FastLUFactorization backend) for the SSS common-block solve + # tmp = (I_nPast - 𝐒₁[past, 1:nPast]); collect(tmp \ vec). Shared by primal and rrule. + sss_tmp_lu_buffer::𝒮.LinearCache + # Dedicated FastLapackInterface LU workspace for the SSS pullback transpose solve + fast_lu_ws_sss_pullback::FastLapackInterface.LUWs + fast_lu_dims_sss_pullback::NTuple{2, Int} + # SSS Newton iter kron! buffers (Float64 path; shared by primal, rrule forward loop, and ForwardDiffExt) + x_aug_buf::Vector{F} # length nPast+1, holds [x; 1] + kron_x_aug_xx::Vector{F} # length (nPast+1)^2, holds kron(x_aug, x_aug) + kron_x_aug_x_kron::Vector{F} # length (nPast+1)^3, holds kron(x_aug, kron_x_aug); 3rd order only + kron_x_aug_I::Matrix{F} # size (nPast+1)*nPast × nPast, holds kron(x_aug, I_nPast) + kron_x_kron_I::Matrix{F} # size (nPast+1)^2*nPast × nPast, holds kron(kron_x_aug, I_nPast); 3rd order only # 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 +1284,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 +1303,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 +1327,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 +1363,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 +1377,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 +1432,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 +1543,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) # ========================================================================= @@ -1118,4 +1557,6 @@ mutable struct ℳ functions::model_functions # Compiled model functions counters::SolveCounters # Solve counters (steady state and perturbation) + + revision_history::Vector{RevisionEntry} # log of equation-modification operations end diff --git a/test/dynare_comparison/.gitignore b/test/dynare_comparison/.gitignore new file mode 100644 index 000000000..3161ee4ff --- /dev/null +++ b/test/dynare_comparison/.gitignore @@ -0,0 +1,2 @@ +# Generated by Phase 1 / Phase 2 of the Dynare comparison pipeline +output/ diff --git a/test/dynare_comparison/Dockerfile b/test/dynare_comparison/Dockerfile new file mode 100644 index 000000000..dc5c3f867 --- /dev/null +++ b/test/dynare_comparison/Dockerfile @@ -0,0 +1,14 @@ +FROM debian:testing + +RUN apt-get update \ + && apt-get install -y --no-install-recommends dynare \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /work + +COPY extract_dynare_results.m /work/ +COPY run_all_dynare.sh /work/ + +RUN chmod +x /work/run_all_dynare.sh + +ENTRYPOINT ["/work/run_all_dynare.sh"] diff --git a/test/dynare_comparison/compare_results.jl b/test/dynare_comparison/compare_results.jl new file mode 100644 index 000000000..1940de816 --- /dev/null +++ b/test/dynare_comparison/compare_results.jl @@ -0,0 +1,1112 @@ +# compare_results.jl — Phase 3 of Dynare comparison +# +# Loads Julia and Dynare CSV outputs for each model, compares them, and +# reports pass/fail. Exits non-zero on any failure. +# +# Requires: DelimitedFiles, Test (both available via --project=.) + +using DelimitedFiles +using LinearAlgebra +using Test + +const RTOL = 1e-6 +const ATOL = 1e-6 +const DEFAULT_OUTPUT_ROOT = joinpath(@__DIR__, "output") +const BENCHMARK_ONLY_MODELS = Set(["FRBUS", "NAWM"]) +const N_BENCH_RUNS = 500 +const _BENCH_CACHE = Dict{String, Dict{String, Float64}}() + +# Execution order from generate_julia_results.jl — first entry at each order +# pays the JIT compilation cost for that order's functions. +const EXECUTION_ORDER = [ + "FS2000", + "Ascari_Sbordone_2014", + "Gali_2015_chapter_3_nonlinear", + "Caldara_et_al_2012", + "Smets_Wouters_2007", + "Smets_Wouters_2003", + "GNSS_2010", + "NAWM_EAUS_2008", + "QUEST3_2009", + "FRBUS", + + "FS2000_pruned_2nd", + "Ascari_Sbordone_2014_pruned_2nd", + "Gali_2015_chapter_3_nonlinear_pruned_2nd", + "Caldara_et_al_2012_pruned_2nd", + "Smets_Wouters_2003_pruned_2nd", + "Smets_Wouters_2007_pruned_2nd", + "GNSS_2010_pruned_2nd", + + "FS2000_pruned_3rd", + "Ascari_Sbordone_2014_pruned_3rd", + "Gali_2015_chapter_3_nonlinear_pruned_3rd", + "Caldara_et_al_2012_pruned_3rd", +] + +"""Sort model names by their position in EXECUTION_ORDER; unknowns go last (alphabetically).""" +function sort_by_execution_order(names) + order_map = Dict(n => i for (i, n) in enumerate(EXECUTION_ORDER)) + sentinel = length(EXECUTION_ORDER) + 1 + sort(names; by = n -> (get(order_map, n, sentinel), n)) +end + +function print_usage() + println("Usage: julia --project=. compare_results.jl [--output-root=PATH | PATH]") +end + +function parse_args(args) + output_root = DEFAULT_OUTPUT_ROOT + positional_args = String[] + + for arg in args + if arg in ("-h", "--help") + print_usage() + return nothing + elseif startswith(arg, "--output-root=") + output_root = split(arg, "=", limit = 2)[2] + elseif startswith(arg, "--") + error("Unknown option: $arg") + else + push!(positional_args, arg) + end + end + + if length(positional_args) > 1 + error("Expected at most one positional output-root argument, got $(length(positional_args))") + elseif length(positional_args) == 1 + output_root = positional_args[1] + end + + return abspath(output_root) +end + +# ───────────────────────────────────────────── +# CSV loading helpers +# ───────────────────────────────────────────── +read_names(path) = strip.(readlines(path)) + +function load_benchmarks(dir) + haskey(_BENCH_CACHE, dir) && return _BENCH_CACHE[dir] + bundled = joinpath(dir, "benchmarks.csv") + benchmarks = Dict{String, Float64}() + if isfile(bundled) + raw = readdlm(bundled, ',') + for row in 1:size(raw, 1) + key = strip(string(raw[row, 1])) + isempty(key) && continue + benchmarks[key] = Float64(raw[row, 2]) + end + end + _BENCH_CACHE[dir] = benchmarks + return benchmarks +end + +function read_bench(dir, name) + key = endswith(name, ".csv") ? name[1:end-4] : name + benchmarks = load_benchmarks(dir) + if haskey(benchmarks, key) + return benchmarks[key] + end + legacy = joinpath(dir, key * ".csv") + return isfile(legacy) ? read_vector(legacy)[1] : NaN +end + +function load_runtime_csv(path) + runtimes = Dict{String, Float64}() + isfile(path) || return runtimes + for line in eachline(path) + stripped = strip(line) + isempty(stripped) && continue + stripped == "model,elapsed_seconds" && continue + parts = split(stripped, ',', limit = 2) + length(parts) == 2 || continue + name = strip(parts[1]) + val = tryparse(Float64, strip(parts[2])) + val === nothing && continue + runtimes[name] = val + end + return runtimes +end + +function read_key_value_metadata(path) + metadata = Dict{String, String}() + if !isfile(path) + return metadata + end + + for line in eachline(path) + stripped = strip(line) + isempty(stripped) && continue + idx = findfirst(==('='), stripped) + idx === nothing && continue + key = strip(stripped[begin:prevind(stripped, idx)]) + value = strip(stripped[nextind(stripped, idx):end]) + metadata[key] = value + end + + return metadata +end + +function current_julia_metadata() + blas_lapack = try + string(BLAS.get_config()) + catch + "unknown" + end + + return Dict( + "julia_version" => string(VERSION), + "julia_threads" => string(Threads.nthreads()), + "julia_threads_default" => string(Threads.nthreads(:default)), + "julia_threads_interactive" => string(Threads.nthreads(:interactive)), + "blas_threads" => string(BLAS.get_num_threads()), + "blas_lapack" => blas_lapack, + "hostname" => get(ENV, "COMPUTERNAME", get(ENV, "HOSTNAME", "unknown")), + "kernel" => string(Sys.KERNEL), + "arch" => string(Sys.ARCH), + "cpu_name" => string(Sys.CPU_NAME), + "cpu_threads" => string(Sys.CPU_THREADS), + "word_size" => string(Sys.WORD_SIZE), + "total_memory_bytes" => try + string(Sys.total_memory()) + catch + "unknown" + end, + ) +end + +function format_memory_string(bytes_string) + try + gib = parse(Float64, bytes_string) / 1024.0^3 + return string(round(gib, digits = 2), " GiB") + catch + return bytes_string + end +end + +function print_environment_summary(output_root) + julia_metadata = read_key_value_metadata(joinpath(output_root, "comparison_environment_julia.txt")) + dynare_metadata = read_key_value_metadata(joinpath(output_root, "comparison_environment_dynare.txt")) + julia_source = "phase-1 metadata" + if isempty(julia_metadata) + julia_metadata = current_julia_metadata() + julia_source = "compare runtime fallback" + end + + println("Run Environment") + println(" Julia ($julia_source):") + println(" version: ", get(julia_metadata, "julia_version", "unknown")) + println(" BLAS/LAPACK: ", get(julia_metadata, "blas_lapack", "unknown")) + println(" threads: Julia=", get(julia_metadata, "julia_threads", "unknown"), + " default=", get(julia_metadata, "julia_threads_default", "unknown"), + " interactive=", get(julia_metadata, "julia_threads_interactive", "unknown"), + " BLAS=", get(julia_metadata, "blas_threads", "unknown")) + println(" machine: host=", get(julia_metadata, "hostname", "unknown"), + " kernel=", get(julia_metadata, "kernel", "unknown"), + " arch=", get(julia_metadata, "arch", "unknown"), + " cpu=", get(julia_metadata, "cpu_name", "unknown"), + " cpu_threads=", get(julia_metadata, "cpu_threads", "unknown"), + " memory=", format_memory_string(get(julia_metadata, "total_memory_bytes", "unknown"))) + + println(" Dynare:") + if isempty(dynare_metadata) + println(" metadata unavailable") + else + driver = get(dynare_metadata, "dynare_driver", "unknown") + version = get(dynare_metadata, "dynare_version", "unknown") + blas = get(dynare_metadata, "blas", "unknown") + lapack = get(dynare_metadata, "lapack", "unknown") + println(" driver/version: ", driver, " / ", version) + if haskey(dynare_metadata, "matlab_version") + println(" MATLAB: ", get(dynare_metadata, "matlab_version", "unknown"), + " release=", get(dynare_metadata, "matlab_release", "unknown")) + elseif haskey(dynare_metadata, "octave_version") + println(" Octave: ", get(dynare_metadata, "octave_version", "unknown")) + end + println(" BLAS/LAPACK: ", blas, " / ", lapack) + println(" machine: host=", get(dynare_metadata, "hostname", "unknown"), + " os=", get(dynare_metadata, "os", get(dynare_metadata, "kernel", "unknown")), + " arch=", get(dynare_metadata, "arch", get(dynare_metadata, "computer", "unknown")), + " cpu_threads=", get(dynare_metadata, "cpu_threads", get(dynare_metadata, "max_num_comp_threads", "unknown"))) + println(" threads: requested=", get(dynare_metadata, "thread_count_requested", "unknown"), + " active=", get(dynare_metadata, "max_num_comp_threads", "unknown")) + end +end + +function read_vector(path) + vec(readdlm(path, ',', Float64)) +end + +function read_matrix(path) + readdlm(path, ',', Float64) +end + +# ───────────────────────────────────────────── +# Load results from a directory (julia/ or dynare/) +# ───────────────────────────────────────────── +function load_results(dir) + r = Dict{Symbol, Any}() + + r[:var_names] = read_names(joinpath(dir, "var_names.csv")) + r[:exo_names] = read_names(joinpath(dir, "exo_names.csv")) + r[:state_var_names] = read_names(joinpath(dir, "state_var_names.csv")) + r[:steady_state] = read_vector(joinpath(dir, "steady_state.csv")) + r[:ghx] = read_matrix(joinpath(dir, "ghx.csv")) + r[:ghu] = read_matrix(joinpath(dir, "ghu.csv")) + + policy_alg_path = joinpath(dir, "policy_algorithm.csv") + if isfile(policy_alg_path) + algs = read_names(policy_alg_path) + if !isempty(algs) + r[:policy_algorithm] = algs[1] + end + end + + # IRFs (optional — may not exist if all zero) + irf_fields_path = joinpath(dir, "irf_fields.csv") + if isfile(irf_fields_path) + fields = read_names(irf_fields_path) + irfs = Dict{String, Vector{Float64}}() + bundled_path = joinpath(dir, "irfs.csv") + if isfile(bundled_path) + # Bundled format: matrix with rows = periods, cols = fields (in irf_fields.csv order). + mat = read_matrix(bundled_path) + ncols = min(size(mat, 2), length(fields)) + for j in 1:ncols + irfs[fields[j]] = vec(mat[:, j]) + end + else + # Legacy per-field files (kept for backward compatibility with older outputs). + for f in fields + p = joinpath(dir, "irf_$f.csv") + if isfile(p) + irfs[f] = read_vector(p) + end + end + end + r[:irfs] = irfs + r[:irf_fields] = fields + + irf_alg_path = joinpath(dir, "irf_algorithm.csv") + if isfile(irf_alg_path) + algs = read_names(irf_alg_path) + if !isempty(algs) + r[:irf_algorithm] = algs[1] + end + end + end + + # Variance-covariance + vcov_path = joinpath(dir, "variance_covariance.csv") + if isfile(vcov_path) + r[:variance_covariance] = read_matrix(vcov_path) + end + + # Variance decomposition + vd_path = joinpath(dir, "variance_decomposition.csv") + if isfile(vd_path) + r[:variance_decomposition] = read_matrix(vd_path) + r[:vd_var_names] = read_names(joinpath(dir, "variance_decomposition_var_names.csv")) + r[:vd_exo_names] = read_names(joinpath(dir, "variance_decomposition_exo_names.csv")) + end + + # Higher-order solution matrices (optional) + for key in [:ghxx, :ghxu, :ghuu, :ghs2, + :ghxxx, :ghxxu, :ghxuu, :ghuuu, :ghxss, :ghuss] + p = joinpath(dir, "$(key).csv") + if isfile(p) + r[key] = read_matrix(p) + end + end + + r +end + +# ───────────────────────────────────────────── +# Build index lookup: name → row/col index +# ───────────────────────────────────────────── +name_index(names) = Dict(n => i for (i, n) in enumerate(names)) + +function common_named_indices(jl_names, dy_names) + common_names = intersect(jl_names, dy_names) + jl_idx = name_index(jl_names) + dy_idx = name_index(dy_names) + return common_names, [jl_idx[name] for name in common_names], [dy_idx[name] for name in common_names] +end + +function kron_linear_index(col_names, col_idxs, col_sizes) + linear_index = 1 + stride = prod(col_sizes) + for k in eachindex(col_names) + stride ÷= col_sizes[k] + linear_index += (col_idxs[k][col_names[k]] - 1) * stride + end + return linear_index +end + +function common_kron_column_indices(jl_col_name_vecs::Vector{<:AbstractVector}, + dy_col_name_vecs::Vector{<:AbstractVector}) + jl_col_idxs = [name_index(v) for v in jl_col_name_vecs] + dy_col_idxs = [name_index(v) for v in dy_col_name_vecs] + common_cols = [intersect(jl_col_name_vecs[k], dy_col_name_vecs[k]) for k in eachindex(jl_col_name_vecs)] + jl_col_sizes = [length(v) for v in jl_col_name_vecs] + dy_col_sizes = [length(v) for v in dy_col_name_vecs] + + common_tuples = collect(Iterators.product(common_cols...)) + jl_col_indices = [kron_linear_index(col_names, jl_col_idxs, jl_col_sizes) for col_names in common_tuples] + dy_col_indices = [kron_linear_index(col_names, dy_col_idxs, dy_col_sizes) for col_names in common_tuples] + + return common_cols, jl_col_indices, dy_col_indices +end + +is_nawm_model(model_name) = model_name == "NAWM_EAUS_2008" +is_higher_order_model(model_name) = occursin("_pruned_2nd", model_name) || occursin("_pruned_3rd", model_name) +is_pruned_third_order_model(model_name) = occursin("_pruned_3rd", model_name) +is_excluded_model_dir(model_name) = model_name == "FS2000_pruned_3rd" +is_benchmark_only_model_dir(model_name) = model_name in BENCHMARK_ONLY_MODELS +is_supported_pruned_third_order_variance_model(model_name) = model_name in ( + "Gali_2015_chapter_3_nonlinear_pruned_3rd", +) + +# ───────────────────────────────────────────── +# Comparison functions — first order +# ───────────────────────────────────────────── + +function compare_steady_state(jl, dy; rtol = RTOL, atol = ATOL) + jl_idx = name_index(jl[:var_names]) + dy_idx = name_index(dy[:var_names]) + + for v in jl[:var_names] + if !haskey(dy_idx, v) + @warn "steady state: Variable $v missing from Dynare" + end + end + + common, jl_common_idx, dy_common_idx = common_named_indices(jl[:var_names], dy[:var_names]) + @test length(common) > 0 + @test length(common) >= min(length(jl[:var_names]), length(dy[:var_names])) * 0.5 + @test isapprox(jl[:steady_state][jl_common_idx], dy[:steady_state][dy_common_idx]; rtol = rtol, atol = atol) +end + +function compare_ghx(jl, dy; rtol = RTOL, atol = ATOL) + jl_vidx = name_index(jl[:var_names]) + dy_vidx = name_index(dy[:var_names]) + jl_sidx = name_index(jl[:state_var_names]) + dy_sidx = name_index(dy[:state_var_names]) + + common_vars, jl_var_idx, dy_var_idx = common_named_indices(jl[:var_names], dy[:var_names]) + common_states, jl_state_idx, dy_state_idx = common_named_indices(jl[:state_var_names], dy[:state_var_names]) + @test length(common_vars) > 0 + @test length(common_states) > 0 + + for v in jl[:var_names] + if !haskey(dy_vidx, v) + @warn "ghx: Variable $v missing from Dynare" + end + end + for s in jl[:state_var_names] + if !haskey(dy_sidx, s) + @warn "ghx: State $s missing from Dynare" + end + end + + jl_subset = jl[:ghx][jl_var_idx, jl_state_idx] + dy_subset = dy[:ghx][dy_var_idx, dy_state_idx] + @test isapprox(jl_subset, dy_subset; rtol = rtol, atol = atol) +end + +function compare_ghu(jl, dy; rtol = RTOL, atol = ATOL) + jl_vidx = name_index(jl[:var_names]) + dy_vidx = name_index(dy[:var_names]) + jl_eidx = name_index(jl[:exo_names]) + dy_eidx = name_index(dy[:exo_names]) + + common_vars, jl_var_idx, dy_var_idx = common_named_indices(jl[:var_names], dy[:var_names]) + common_exo, jl_exo_idx, dy_exo_idx = common_named_indices(jl[:exo_names], dy[:exo_names]) + @test length(common_vars) > 0 + @test length(common_exo) > 0 + + for v in jl[:var_names] + if !haskey(dy_vidx, v) + @warn "ghu: Variable $v missing from Dynare" + end + end + for e in jl[:exo_names] + if !haskey(dy_eidx, e) + @warn "ghu: Shock $e missing from Dynare" + end + end + + jl_subset = jl[:ghu][jl_var_idx, jl_exo_idx] + dy_subset = dy[:ghu][dy_var_idx, dy_exo_idx] + @test isapprox(jl_subset, dy_subset; rtol = rtol, atol = atol) +end + +function compare_irfs(jl, dy; model_name = "", rtol = RTOL, atol = ATOL) + haskey(jl, :irfs) && haskey(dy, :irfs) || return + + if is_higher_order_model(model_name) + @info "Skipping IRF comparison for $model_name (higher-order IRFs are convention-dependent; compare moments instead)" + return + end + + # Backward-compatibility guard: for higher-order model directories, compare IRFs + # only when Julia IRFs were explicitly generated at first order. + if is_higher_order_model(model_name) + irf_alg = get(jl, :irf_algorithm, "") + if irf_alg != "first_order" + @info "Skipping IRF comparison for $model_name (IRFs not tagged as first-order; regenerate phase-1 outputs to enable)" + return + end + end + + for f in get(jl, :irf_fields, String[]) + if !haskey(dy[:irfs], f) + @warn "IRF field $f missing from Dynare" + end + end + + common_fields = intersect(keys(jl[:irfs]), keys(dy[:irfs])) + for f in common_fields + jvec = jl[:irfs][f] + dvec = dy[:irfs][f] + n = min(length(jvec), length(dvec)) + @test isapprox(jvec[1:n], dvec[1:n]; rtol = rtol, atol = atol) + end +end + +function compare_variance(jl, dy; rtol = RTOL, atol = ATOL) + haskey(jl, :variance_covariance) && haskey(dy, :variance_covariance) || return + + jl_idx = name_index(jl[:var_names]) + dy_idx = name_index(dy[:var_names]) + common = intersect(jl[:var_names], dy[:var_names]) + valid_common = filter(v -> begin + ji = jl_idx[v] + di = dy_idx[v] + ji <= size(jl[:variance_covariance], 1) && di <= size(dy[:variance_covariance], 1) + end, common) + isempty(valid_common) && return + + jl_variance = [jl[:variance_covariance][jl_idx[v], jl_idx[v]] for v in valid_common] + dy_variance = [dy[:variance_covariance][dy_idx[v], dy_idx[v]] for v in valid_common] + @test isapprox(jl_variance, dy_variance; rtol = rtol, atol = atol) + + jl_std = sqrt.(jl_variance) + dy_std = sqrt.(dy_variance) + @test isapprox(jl_std, dy_std; rtol = rtol, atol = atol) +end + +function compare_variance_decomposition(jl, dy; rtol = RTOL, atol = ATOL) + haskey(jl, :variance_decomposition) && haskey(dy, :variance_decomposition) || return + + jl_vidx = name_index(jl[:vd_var_names]) + dy_vidx = name_index(dy[:vd_var_names]) + jl_eidx = name_index(jl[:vd_exo_names]) + dy_eidx = name_index(dy[:vd_exo_names]) + + common_vars = intersect(jl[:vd_var_names], dy[:vd_var_names]) + common_exo = intersect(jl[:vd_exo_names], dy[:vd_exo_names]) + + valid_vars = filter(v -> begin + ji = jl_vidx[v] + di = dy_vidx[v] + ji <= size(jl[:variance_decomposition], 1) || return false + di <= size(dy[:variance_decomposition], 1) || return false + jl_row_sum = sum(abs, jl[:variance_decomposition][ji, :]) + dy_row_sum = sum(abs, dy[:variance_decomposition][di, :]) + jl_row_sum >= 1.0 && dy_row_sum >= 1.0 + end, common_vars) + valid_exo = filter(e -> begin + jl_eidx[e] <= size(jl[:variance_decomposition], 2) && dy_eidx[e] <= size(dy[:variance_decomposition], 2) + end, common_exo) + isempty(valid_vars) && return + isempty(valid_exo) && return + + jl_var_idx = [jl_vidx[v] for v in valid_vars] + dy_var_idx = [dy_vidx[v] for v in valid_vars] + jl_exo_idx = [jl_eidx[e] for e in valid_exo] + dy_exo_idx = [dy_eidx[e] for e in valid_exo] + + jl_subset = jl[:variance_decomposition][jl_var_idx, jl_exo_idx] + dy_subset = dy[:variance_decomposition][dy_var_idx, dy_exo_idx] + + comparison_atol = max(atol, 0.01) + ok = isapprox(jl_subset, dy_subset; rtol = rtol, atol = comparison_atol) || + all(abs.(jl_subset) .< comparison_atol .&& abs.(dy_subset) .< comparison_atol) + if !ok + diff = maximum(abs.(jl_subset .- dy_subset)) + scale = max(maximum(abs.(jl_subset)), maximum(abs.(dy_subset))) + @warn "Variance decomp mismatch" achieved_atol=diff achieved_rtol=(scale > 0 ? diff / scale : Inf) required_atol=comparison_atol required_rtol=rtol + end + @test ok +end + +# ───────────────────────────────────────────── +# Comparison functions — higher-order matrices +# ───────────────────────────────────────────── + +""" +Compare a Kronecker-product matrix (ghxx, ghuu, ghxxx, ghuuu, etc.) +indexed by kron of name vectors (e.g., state × state for ghxx). +Uses tuple-based column alignment: iterate over common (name₁, name₂[, …]) +tuples and look up elements in each side's matrix via their local indices. +""" +function compare_kron_matrix(jl, dy, mat_key::Symbol, + jl_row_names, dy_row_names, + jl_col_name_vecs::Vector{<:AbstractVector}, + dy_col_name_vecs::Vector{<:AbstractVector}; + rtol = RTOL, atol = ATOL) + haskey(jl, mat_key) && haskey(dy, mat_key) || return + + common_rows, jl_row_idx, dy_row_idx = common_named_indices(jl_row_names, dy_row_names) + @test length(common_rows) > 0 + + common_cols, jl_col_idx, dy_col_idx = common_kron_column_indices(jl_col_name_vecs, dy_col_name_vecs) + for k in eachindex(common_cols) + @test length(common_cols[k]) > 0 + end + + jl_subset = jl[mat_key][jl_row_idx, jl_col_idx] + dy_subset = dy[mat_key][dy_row_idx, dy_col_idx] + @test isapprox(jl_subset, dy_subset; rtol = rtol, atol = atol) +end + +function compare_vector_matrix(jl, dy, mat_key::Symbol, + jl_row_names, dy_row_names, + jl_col_names, dy_col_names; + rtol = RTOL, atol = ATOL) + haskey(jl, mat_key) && haskey(dy, mat_key) || return + + common_rows, jl_row_idx, dy_row_idx = common_named_indices(jl_row_names, dy_row_names) + common_cols, jl_col_idx, dy_col_idx = common_named_indices(jl_col_names, dy_col_names) + + jl_mat = jl[mat_key] + dy_mat = dy[mat_key] + + @test length(common_rows) > 0 + @test length(common_cols) > 0 + @test isapprox(jl_mat[jl_row_idx, jl_col_idx], dy_mat[dy_row_idx, dy_col_idx]; rtol = rtol, atol = atol) +end + +function compare_second_order(jl, dy; rtol = RTOL, atol = ATOL) + sn_jl = jl[:state_var_names]; sn_dy = dy[:state_var_names] + en_jl = jl[:exo_names]; en_dy = dy[:exo_names] + vn_jl = jl[:var_names]; vn_dy = dy[:var_names] + + @testset "ghxx" begin + compare_kron_matrix(jl, dy, :ghxx, vn_jl, vn_dy, + [sn_jl, sn_jl], [sn_dy, sn_dy]; + rtol = rtol, atol = atol) + end + @testset "ghxu" begin + compare_kron_matrix(jl, dy, :ghxu, vn_jl, vn_dy, + [sn_jl, en_jl], [sn_dy, en_dy]; + rtol = rtol, atol = atol) + end + @testset "ghuu" begin + compare_kron_matrix(jl, dy, :ghuu, vn_jl, vn_dy, + [en_jl, en_jl], [en_dy, en_dy]; + rtol = rtol, atol = atol) + end + @testset "ghs2" begin + if haskey(jl, :ghs2) && haskey(dy, :ghs2) + jl_vidx = name_index(vn_jl) + dy_vidx = name_index(vn_dy) + common_vars = intersect(vn_jl, vn_dy) + for v in common_vars + # ghs2 convention differs between MacroModelling and Dynare: + # MacroModelling extracts the (σ,σ) slice of the second-order tensor + # (pure perturbation-parameter² coefficient), whereas Dynare's ghs2 + # absorbs the full shock covariance matrix. These are different + # mathematical objects and cannot be compared element-wise. + # Skipping ghs2 comparison. + end + end + end +end + +function compare_third_order(jl, dy; rtol = RTOL, atol = ATOL) + sn_jl = jl[:state_var_names]; sn_dy = dy[:state_var_names] + en_jl = jl[:exo_names]; en_dy = dy[:exo_names] + vn_jl = jl[:var_names]; vn_dy = dy[:var_names] + + @testset "ghxxx" begin + compare_kron_matrix(jl, dy, :ghxxx, vn_jl, vn_dy, + [sn_jl, sn_jl, sn_jl], [sn_dy, sn_dy, sn_dy]; + rtol = rtol, atol = atol) + end + @testset "ghxxu" begin + compare_kron_matrix(jl, dy, :ghxxu, vn_jl, vn_dy, + [sn_jl, sn_jl, en_jl], [sn_dy, sn_dy, en_dy]; + rtol = rtol, atol = atol) + end + @testset "ghxuu" begin + compare_kron_matrix(jl, dy, :ghxuu, vn_jl, vn_dy, + [sn_jl, en_jl, en_jl], [sn_dy, en_dy, en_dy]; + rtol = rtol, atol = atol) + end + @testset "ghuuu" begin + compare_kron_matrix(jl, dy, :ghuuu, vn_jl, vn_dy, + [en_jl, en_jl, en_jl], [en_dy, en_dy, en_dy]; + rtol = rtol, atol = atol) + end + @testset "ghxss" begin + compare_vector_matrix(jl, dy, :ghxss, vn_jl, vn_dy, sn_jl, sn_dy; + rtol = rtol, atol = atol) + end + @testset "ghuss" begin + compare_vector_matrix(jl, dy, :ghuss, vn_jl, vn_dy, en_jl, en_dy; + rtol = rtol, atol = atol) + end +end + +# ───────────────────────────────────────────── +# Detect whether a model directory has higher-order results +# ───────────────────────────────────────────── +has_second_order(r) = haskey(r, :ghxx) +has_third_order(r) = haskey(r, :ghxxx) + +# ───────────────────────────────────────────── +# Main +# ───────────────────────────────────────────── +function main(args = ARGS) + output_root = parse_args(args) + output_root === nothing && return + + if !isdir(output_root) + error("Output directory not found: $output_root") + end + + model_dirs = filter(d -> isdir(joinpath(output_root, d, "julia")) && + isdir(joinpath(output_root, d, "dynare")) && + !is_excluded_model_dir(d), + readdir(output_root)) + + comparison_model_dirs = filter(d -> !is_benchmark_only_model_dir(d), model_dirs) + + if isempty(model_dirs) + error("No model directories with both julia/ and dynare/ results found in $output_root") + end + + println("Comparison output root: $output_root") + print_environment_summary(output_root) + + benchmark_only_dirs = filter(is_benchmark_only_model_dir, model_dirs) + for mname in sort_by_execution_order(benchmark_only_dirs) + @info "Skipping correctness comparison for benchmark-only model: $mname" + end + + comparison_exception = nothing + try + if !isempty(comparison_model_dirs) + @testset "Dynare Comparison" begin + for mname in sort_by_execution_order(comparison_model_dirs) + julia_dir = joinpath(output_root, mname, "julia") + dynare_dir = joinpath(output_root, mname, "dynare") + + @info "Comparing results for: $mname" + jl = load_results(julia_dir) + dy = load_results(dynare_dir) + + @testset "$mname" begin + moments_only_higher_order = is_higher_order_model(mname) + skip_pruned_third_order = is_pruned_third_order_model(mname) + + @testset "Steady State" begin + compare_steady_state(jl, dy) + end + @testset "Policy Matrix ghx" begin + if skip_pruned_third_order + @info "Skipping ghx comparison for $mname (pruned third-order state representation mismatch)" + elseif moments_only_higher_order && get(jl, :policy_algorithm, "") != "first_order" + @info "Skipping ghx comparison for $mname (policy matrices not tagged as first-order; regenerate phase-1 outputs to enable)" + else + compare_ghx(jl, dy) + end + end + @testset "Policy Matrix ghu" begin + if moments_only_higher_order && get(jl, :policy_algorithm, "") != "first_order" + @info "Skipping ghu comparison for $mname (policy matrices not tagged as first-order; regenerate phase-1 outputs to enable)" + else + compare_ghu(jl, dy) + end + end + @testset "IRFs" begin + compare_irfs(jl, dy; model_name = mname) + end + @testset "Variance" begin + if skip_pruned_third_order && !is_supported_pruned_third_order_variance_model(mname) + @info "Skipping variance comparison for $mname (pruned third-order moment convention mismatch outside the validated benchmark cases)" + else + compare_variance(jl, dy) + end + end + # Higher-order comparisons (when data is present) + if has_second_order(jl) && has_second_order(dy) + @testset "Second Order Matrices" begin + if moments_only_higher_order + @info "Skipping second-order matrix comparison for $mname (higher-order configured as moments-only)" + else + compare_second_order(jl, dy) + end + end + end + if has_third_order(jl) && has_third_order(dy) + @testset "Third Order Matrices" begin + if moments_only_higher_order + @info "Skipping third-order matrix comparison for $mname (higher-order configured as moments-only)" + else + compare_third_order(jl, dy) + end + end + end + end + end + end + else + @info "No correctness-comparison model directories found under $output_root" + end + catch err + if err isa Test.TestSetException + comparison_exception = err + else + rethrow(err) + end + end + + # ── Benchmark comparison ── + # Dynare benchmarks: component-level (Jacobian, first-order solve, Hessian, second-order solve) + # Julia benchmarks: component-level via BenchmarkTools + # Dynare order=3 also exports k_order_pert as an additional bundled reference. + println("\n", "="^100) + println(" Benchmark Comparison: MacroModelling (median of 500 runs) vs Dynare (median of 500 runs)") + println("="^100) + + has_bench(dir, name) = !isnan(read_bench(dir, name)) + + function sum_bench_components(dir, files) + total = 0.0 + for file in files + value = read_bench(dir, file) + if isnan(value) + return NaN + end + total += value + end + return total + end + + # Dynare order=3 runs additionally export a bundled k_order_pert timing. + is_dynare_k_order_dir(dir) = has_bench(dir, "benchmark_k_order_pert.csv") + + # k_order_pert timing is exported explicitly for Dynare order=3 runs. + function read_dynare_k_order_pert(dir) + has_bench(dir, "benchmark_k_order_pert.csv") ? read_bench(dir, "benchmark_k_order_pert.csv") : NaN + end + + function print_bench_table(title, model_dirs, jl_file, dy_file; note = "") + println("\n--- $title ---") + if !isempty(note) + println(" $note") + end + println(rpad("Model", 50), rpad("MacroModelling", 18), rpad("Dynare", 12), "Speedup") + println("-"^100) + for mname in sort_by_execution_order(model_dirs) + jl_time = read_bench(joinpath(output_root, mname, "julia"), jl_file) + dy_time = read_bench(joinpath(output_root, mname, "dynare"), dy_file) + jl_str = isnan(jl_time) ? "N/A" : format_time(jl_time) + dy_str = isnan(dy_time) ? "N/A" : format_time(dy_time) + speedup_str = (!isnan(jl_time) && !isnan(dy_time) && jl_time > 0) ? + string(round(dy_time / jl_time, digits=1), "x") : "N/A" + println(rpad(mname, 50), rpad(jl_str, 18), rpad(dy_str, 12), speedup_str) + end + end + + # Jacobian (Dynare: dynamic_g1) + print_bench_table("Jacobian", model_dirs, + "benchmark_jacobian.csv", "benchmark_jacobian.csv") + + # First-order solve (Julia: direct QME solve; Dynare: dyn_first_order_solver) + print_bench_table("First-Order Solve", model_dirs, + "benchmark_first_order_solve.csv", "benchmark_first_order_solve.csv") + + # First-order total (sum of direct component medians) + println("\n--- First-Order Total (sum of direct Jacobian + solve medians) ---") + println(rpad("Model", 50), rpad("MacroModelling", 18), rpad("Dynare", 12), "Speedup") + println("-"^100) + for mname in sort_by_execution_order(model_dirs) + jl_dir = joinpath(output_root, mname, "julia") + dy_dir = joinpath(output_root, mname, "dynare") + jl_time = sum_bench_components(jl_dir, ["benchmark_jacobian.csv", "benchmark_first_order_solve.csv"]) + dy_time = sum_bench_components(dy_dir, ["benchmark_jacobian.csv", "benchmark_first_order_solve.csv"]) + jl_str = isnan(jl_time) ? "N/A" : format_time(jl_time) + dy_str = isnan(dy_time) ? "N/A" : format_time(dy_time) + speedup_str = (!isnan(jl_time) && !isnan(dy_time) && jl_time > 0) ? + string(round(dy_time / jl_time, digits=1), "x") : "N/A" + println(rpad(mname, 50), rpad(jl_str, 18), rpad(dy_str, 12), speedup_str) + end + + # Hessian / second-order solve + ho_models = filter(d -> has_bench(joinpath(output_root, d, "julia"), "benchmark_hessian.csv"), model_dirs) + dy_decomposable_ho_models = filter(d -> has_bench(joinpath(output_root, d, "dynare"), "benchmark_hessian.csv"), ho_models) + if !isempty(dy_decomposable_ho_models) + print_bench_table("Hessian", dy_decomposable_ho_models, + "benchmark_hessian.csv", "benchmark_hessian.csv") + + print_bench_table("Second-Order Solve", dy_decomposable_ho_models, + "benchmark_second_order_solve.csv", "benchmark_second_order_solve.csv") + + # Second-Order Total (Hessian + Second-Order Solve) + println("\n--- Second-Order Total (Hessian + Second-Order Solve) ---") + println(rpad("Model", 50), rpad("MacroModelling", 18), rpad("Dynare", 12), "Speedup") + println("-"^100) + for mname in sort_by_execution_order(dy_decomposable_ho_models) + jl_dir = joinpath(output_root, mname, "julia") + dy_dir = joinpath(output_root, mname, "dynare") + jl_time = sum_bench_components(jl_dir, ["benchmark_hessian.csv", "benchmark_second_order_solve.csv"]) + dy_time = sum_bench_components(dy_dir, ["benchmark_hessian.csv", "benchmark_second_order_solve.csv"]) + jl_str = isnan(jl_time) ? "N/A" : format_time(jl_time) + dy_str = isnan(dy_time) ? "N/A" : format_time(dy_time) + speedup_str = (!isnan(jl_time) && !isnan(dy_time) && jl_time > 0) ? + string(round(dy_time / jl_time, digits=1), "x") : "N/A" + println(rpad(mname, 50), rpad(jl_str, 18), rpad(dy_str, 12), speedup_str) + end + end + + # Dynare k_order models: report bundled higher-order timing consistently. + k_order_models = filter(d -> is_dynare_k_order_dir(joinpath(output_root, d, "dynare")), model_dirs) + if !isempty(k_order_models) + println("\n--- Higher-Order Bundled (Dynare k_order_pert) ---") + println(" MacroModelling sums directly measured solve-stack components; Dynare reports direct bundled k_order_pert") + println(rpad("Model", 50), rpad("MacroModelling", 18), rpad("Dynare", 12), "Speedup") + println("-"^100) + for mname in sort_by_execution_order(k_order_models) + jl_dir = joinpath(output_root, mname, "julia") + dy_dir = joinpath(output_root, mname, "dynare") + + jl_fo_solve = read_bench(jl_dir, "benchmark_first_order_solve.csv") + jl_hess = read_bench(jl_dir, "benchmark_hessian.csv") + jl_so = read_bench(jl_dir, "benchmark_second_order_solve.csv") + jl_td = read_bench(jl_dir, "benchmark_third_order_derivatives.csv") + jl_ts = read_bench(jl_dir, "benchmark_third_order_solve.csv") + + jl_bundled = jl_fo_solve + isnan(jl_hess) || (jl_bundled += jl_hess) + isnan(jl_so) || (jl_bundled += jl_so) + isnan(jl_td) || (jl_bundled += jl_td) + isnan(jl_ts) || (jl_bundled += jl_ts) + + dy_bundled = read_dynare_k_order_pert(dy_dir) + + jl_str = isnan(jl_bundled) ? "N/A" : format_time(jl_bundled) + dy_str = isnan(dy_bundled) ? "N/A" : format_time(dy_bundled) + speedup_str = (!isnan(jl_bundled) && !isnan(dy_bundled) && jl_bundled > 0) ? + string(round(dy_bundled / jl_bundled, digits=1), "x") : "N/A" + println(rpad(mname, 50), rpad(jl_str, 18), rpad(dy_str, 12), speedup_str) + end + end + + if !isempty(dy_decomposable_ho_models) + println("\n--- Comparable Direct Components Total (Jacobian + FO + Hessian + SO) ---") + println(rpad("Model", 50), rpad("MacroModelling", 18), rpad("Dynare", 12), "Speedup") + println("-"^100) + for mname in sort_by_execution_order(dy_decomposable_ho_models) + jl_dir = joinpath(output_root, mname, "julia") + dy_dir = joinpath(output_root, mname, "dynare") + + jl_total = sum_bench_components(jl_dir, [ + "benchmark_jacobian.csv", + "benchmark_first_order_solve.csv", + "benchmark_hessian.csv", + "benchmark_second_order_solve.csv", + ]) + dy_total = sum_bench_components(dy_dir, [ + "benchmark_jacobian.csv", + "benchmark_first_order_solve.csv", + "benchmark_hessian.csv", + "benchmark_second_order_solve.csv", + ]) + + jl_str = isnan(jl_total) ? "N/A" : format_time(jl_total) + dy_str = isnan(dy_total) ? "N/A" : format_time(dy_total) + speedup_str = (!isnan(jl_total) && !isnan(dy_total) && jl_total > 0) ? + string(round(dy_total / jl_total, digits=1), "x") : "N/A" + println(rpad(mname, 50), rpad(jl_str, 18), rpad(dy_str, 12), speedup_str) + end + end + + # Third-order components (MacroModelling only — Dynare uses k_order_pert for order=3) + to_models = filter(d -> has_bench(joinpath(output_root, d, "julia"), "benchmark_third_order_derivatives.csv"), + model_dirs) + if !isempty(to_models) + println("\n--- Third-Order Components (MacroModelling only — Dynare k_order_pert is bundled) ---") + println(rpad("Model", 50), rpad("3rd Derivs", 15), "3rd Solve") + println("-"^100) + for mname in sort_by_execution_order(to_models) + td_time = read_bench(joinpath(output_root, mname, "julia"), "benchmark_third_order_derivatives.csv") + ts_time = read_bench(joinpath(output_root, mname, "julia"), "benchmark_third_order_solve.csv") + td = isnan(td_time) ? "N/A" : format_time(td_time) + ts = isnan(ts_time) ? "N/A" : format_time(ts_time) + println(rpad(mname, 50), rpad(td, 15), ts) + end + end + + # ── Full Solution Pipeline Summary ── + # Consolidated table showing total solve time per model across all available + # pipeline components, with component breakdown. + println("\n", "="^100) + println(" Full Solution Pipeline Summary: MacroModelling vs Dynare (all components combined)") + println("="^100) + println(rpad("Model", 40), rpad("Order", 7), rpad("MacroModelling", 18), rpad("Dynare", 18), rpad("Speedup", 10), "Components (MacroModelling)") + println("-"^140) + + all_jl_components = [ + "benchmark_jacobian", + "benchmark_first_order_solve", + "benchmark_hessian", + "benchmark_second_order_solve", + "benchmark_third_order_derivatives", + "benchmark_third_order_solve", + ] + component_short_names = Dict( + "benchmark_jacobian" => "Jac", + "benchmark_first_order_solve" => "FO", + "benchmark_hessian" => "Hess", + "benchmark_second_order_solve" => "SO", + "benchmark_third_order_derivatives" => "3rdD", + "benchmark_third_order_solve" => "3rdS", + "benchmark_k_order_pert" => "k_order", + ) + + for mname in sort_by_execution_order(model_dirs) + jl_dir = joinpath(output_root, mname, "julia") + dy_dir = joinpath(output_root, mname, "dynare") + + # Determine model order from available benchmarks + has_3rd = has_bench(jl_dir, "benchmark_third_order_derivatives") + has_2nd = has_bench(jl_dir, "benchmark_hessian") + order_str = has_3rd ? "3rd" : has_2nd ? "2nd" : "1st" + + # Sum all available Julia components + jl_total = 0.0 + jl_parts = String[] + for comp in all_jl_components + val = read_bench(jl_dir, comp) + if !isnan(val) + jl_total += val + push!(jl_parts, "$(component_short_names[comp])=$(format_time(val))") + end + end + jl_total = jl_total > 0.0 ? jl_total : NaN + + # Sum Dynare components — for order 3, Dynare uses bundled k_order_pert + # which includes FO solve + Hessian + SO solve + 3rd-order in one call. + # We use: Jacobian + k_order_pert when available, otherwise sum components. + dy_has_korder = has_bench(dy_dir, "benchmark_k_order_pert") + dy_total = 0.0 + if dy_has_korder + # k_order_pert bundles everything except Jacobian + dy_jac = read_bench(dy_dir, "benchmark_jacobian") + dy_korder = read_bench(dy_dir, "benchmark_k_order_pert") + dy_total = (isnan(dy_jac) ? 0.0 : dy_jac) + (isnan(dy_korder) ? 0.0 : dy_korder) + else + for comp in ["benchmark_jacobian", "benchmark_first_order_solve", + "benchmark_hessian", "benchmark_second_order_solve"] + val = read_bench(dy_dir, comp) + if !isnan(val) + dy_total += val + end + end + end + dy_total = dy_total > 0.0 ? dy_total : NaN + + jl_str = isnan(jl_total) ? "N/A" : format_time(jl_total) + dy_str = isnan(dy_total) ? "N/A" : format_time(dy_total) + speedup_str = (!isnan(jl_total) && !isnan(dy_total) && jl_total > 0) ? + string(round(dy_total / jl_total, digits=1), "x") : "N/A" + parts_str = join(jl_parts, ", ") + + println(rpad(mname, 40), rpad(order_str, 7), rpad(jl_str, 18), rpad(dy_str, 18), rpad(speedup_str, 10), parts_str) + end + + println("="^140) + println(" Note: MacroModelling timings are summed from individually measured components.") + println(" Dynare order≥3 uses bundled k_order_pert (Jacobian measured separately).") + println(" All timings are median of $(N_BENCH_RUNS) runs.") + + # ── Wall-Clock Timing Comparison (single-run, includes compilation) ── + # Shows total elapsed time per model for a single run including JIT + # compilation (Julia) and interpreter startup (Dynare/MATLAB), but excluding + # the repeated benchmark loops. This highlights compilation overhead. + jl_runtime_path = joinpath(output_root, "runtime_julia.csv") + dy_runtime_path = joinpath(output_root, "runtime_dynare.csv") + + if isfile(jl_runtime_path) || isfile(dy_runtime_path) + jl_runtimes = load_runtime_csv(jl_runtime_path) + dy_runtimes = load_runtime_csv(dy_runtime_path) + + # Collect all model names from both sides, in execution order + all_runtime_models = sort_by_execution_order(collect(union( + filter(k -> k != "TOTAL", collect(keys(jl_runtimes))), + filter(k -> k != "TOTAL", collect(keys(dy_runtimes))) + ))) + + if !isempty(all_runtime_models) + println("\n", "="^100) + println(" Wall-Clock Timing Comparison (single run, includes compilation/startup overhead)") + println("="^100) + println(rpad("Model", 50), rpad("MacroModelling", 18), rpad("Dynare", 18), "Speedup") + println("-"^100) + + jl_sum = 0.0 + dy_sum = 0.0 + for mname in all_runtime_models + jl_t = get(jl_runtimes, mname, NaN) + dy_t = get(dy_runtimes, mname, NaN) + if !isnan(jl_t); jl_sum += jl_t; end + if !isnan(dy_t); dy_sum += dy_t; end + jl_str = isnan(jl_t) ? "N/A" : format_time(jl_t) + dy_str = isnan(dy_t) ? "N/A" : format_time(dy_t) + speedup_str = (!isnan(jl_t) && !isnan(dy_t) && jl_t > 0) ? + string(round(dy_t / jl_t, digits=2), "x") : "N/A" + println(rpad(mname, 50), rpad(jl_str, 18), rpad(dy_str, 18), speedup_str) + end + + # Total row + println("-"^100) + jl_total_rt = get(jl_runtimes, "TOTAL", jl_sum) + dy_total_rt = get(dy_runtimes, "TOTAL", dy_sum) + jl_str = jl_total_rt > 0 ? format_time(jl_total_rt) : "N/A" + dy_str = dy_total_rt > 0 ? format_time(dy_total_rt) : "N/A" + speedup_str = (jl_total_rt > 0 && dy_total_rt > 0) ? + string(round(dy_total_rt / jl_total_rt, digits=2), "x") : "N/A" + println(rpad("TOTAL", 50), rpad(jl_str, 18), rpad(dy_str, 18), speedup_str) + + println("="^100) + println(" Note: Wall-clock times measure a single run including JIT compilation (Julia)") + println(" and full model processing including stoch_simul (Dynare). Benchmark loop") + println(" iterations ($(N_BENCH_RUNS) runs) are excluded. This shows the one-shot cost of") + println(" producing results once, highlighting compilation overhead.") + end + end + + println("="^100) + + comparison_exception === nothing || throw(comparison_exception) +end + +function format_time(t) + if t < 1e-3 + string(round(t * 1e6, digits=1), " μs") + elseif t < 1.0 + string(round(t * 1e3, digits=2), " ms") + else + string(round(t, digits=3), " s") + end +end + +if abspath(PROGRAM_FILE) == @__FILE__ + main() +end diff --git a/test/dynare_comparison/compare_thread_sweep_results.jl b/test/dynare_comparison/compare_thread_sweep_results.jl new file mode 100644 index 000000000..46cf91d13 --- /dev/null +++ b/test/dynare_comparison/compare_thread_sweep_results.jl @@ -0,0 +1,467 @@ +include(joinpath(@__DIR__, "compare_results.jl")) + +const DEFAULT_SWEEP_ROOT = joinpath(@__DIR__, "output_thread_sweep") + +function print_usage() + println("Usage: julia --project=. compare_thread_sweep_results.jl [--output-root=PATH | PATH]") +end + +function parse_args(args) + output_root = DEFAULT_SWEEP_ROOT + positional_args = String[] + + for arg in args + if arg in ("-h", "--help") + print_usage() + return nothing + elseif startswith(arg, "--output-root=") + output_root = split(arg, "=", limit = 2)[2] + elseif startswith(arg, "--") + error("Unknown option: $arg") + else + push!(positional_args, arg) + end + end + + if length(positional_args) > 1 + error("Expected at most one positional output-root argument, got $(length(positional_args))") + elseif length(positional_args) == 1 + output_root = positional_args[1] + end + + return abspath(output_root) +end + +function print_environment_summary(thread_counts, thread_dirs) + first_thread = first(thread_counts) + metadata_root = thread_dirs[first_thread] + julia_metadata = read_key_value_metadata(joinpath(metadata_root, "comparison_environment_julia.txt")) + dynare_metadata = read_key_value_metadata(joinpath(metadata_root, "comparison_environment_dynare.txt")) + julia_source = "threads_$(first_thread) metadata" + + if isempty(julia_metadata) + julia_metadata = current_julia_metadata() + julia_source = "compare runtime fallback" + end + + println("Run Environment (metadata source: threads_$(first_thread))") + println(" Julia ($julia_source):") + println(" version: ", get(julia_metadata, "julia_version", "unknown")) + println(" BLAS/LAPACK: ", get(julia_metadata, "blas_lapack", "unknown")) + println(" threads: Julia=", get(julia_metadata, "julia_threads", "unknown"), + " default=", get(julia_metadata, "julia_threads_default", "unknown"), + " interactive=", get(julia_metadata, "julia_threads_interactive", "unknown"), + " BLAS=", get(julia_metadata, "blas_threads", "unknown")) + println(" machine: host=", get(julia_metadata, "hostname", "unknown"), + " kernel=", get(julia_metadata, "kernel", "unknown"), + " arch=", get(julia_metadata, "arch", "unknown"), + " cpu=", get(julia_metadata, "cpu_name", "unknown"), + " cpu_threads=", get(julia_metadata, "cpu_threads", "unknown"), + " memory=", format_memory_string(get(julia_metadata, "total_memory_bytes", "unknown"))) + + println(" Dynare:") + if isempty(dynare_metadata) + println(" metadata unavailable") + else + println(" driver/version: ", get(dynare_metadata, "dynare_driver", "unknown"), + " / ", get(dynare_metadata, "dynare_version", "unknown")) + if haskey(dynare_metadata, "matlab_version") + println(" MATLAB: ", get(dynare_metadata, "matlab_version", "unknown"), + " release=", get(dynare_metadata, "matlab_release", "unknown")) + elseif haskey(dynare_metadata, "octave_version") + println(" Octave: ", get(dynare_metadata, "octave_version", "unknown")) + end + println(" BLAS/LAPACK: ", get(dynare_metadata, "blas", "unknown"), + " / ", get(dynare_metadata, "lapack", "unknown")) + println(" machine: host=", get(dynare_metadata, "hostname", "unknown"), + " os=", get(dynare_metadata, "os", get(dynare_metadata, "kernel", "unknown")), + " arch=", get(dynare_metadata, "arch", get(dynare_metadata, "computer", "unknown")), + " cpu_threads=", get(dynare_metadata, "cpu_threads", get(dynare_metadata, "max_num_comp_threads", "unknown"))) + println(" threads: requested=", get(dynare_metadata, "thread_count_requested", "unknown"), + " active=", get(dynare_metadata, "max_num_comp_threads", "unknown")) + end +end + +function collect_thread_dirs(output_root) + thread_dirs = Dict{Int, String}() + for entry in readdir(output_root) + full_path = joinpath(output_root, entry) + isdir(full_path) || continue + match_result = match(r"^threads_(\d+)$", entry) + match_result === nothing && continue + thread_dirs[parse(Int, match_result.captures[1])] = full_path + end + + isempty(thread_dirs) && error("No thread-sweep directories found under $output_root") + return sort!(collect(keys(thread_dirs))), thread_dirs +end + +function model_names(thread_dir) + filter(model_name -> isdir(joinpath(thread_dir, model_name, "julia")) && + isdir(joinpath(thread_dir, model_name, "dynare")) && + !is_excluded_model_dir(model_name), + readdir(thread_dir)) +end + +function collect_model_names(thread_counts, thread_dirs; include_benchmark_only = true) + model_set = Set{String}() + for thread_count in thread_counts + for model_name in model_names(thread_dirs[thread_count]) + push!(model_set, model_name) + end + end + + model_dir_names = sort!(collect(model_set)) + return include_benchmark_only ? model_dir_names : filter(mname -> !is_benchmark_only_model_dir(mname), model_dir_names) +end + +has_bench(dir, name) = !isnan(read_bench(dir, name)) + +function sum_bench_components(dir, files) + total = 0.0 + for file in files + value = read_bench(dir, file) + if isnan(value) + return NaN + end + total += value + end + return total +end + +function sum_optional_components(dir, files) + total = 0.0 + found_value = false + for file in files + value = read_bench(dir, file) + if !isnan(value) + total += value + found_value = true + end + end + return found_value ? total : NaN +end + +is_dynare_k_order_dir(dir) = has_bench(dir, "benchmark_k_order_pert.csv") +read_dynare_k_order_pert(dir) = has_bench(dir, "benchmark_k_order_pert.csv") ? read_bench(dir, "benchmark_k_order_pert.csv") : NaN + +jacobian_metric(julia_dir, dynare_dir) = ( + read_bench(julia_dir, "benchmark_jacobian.csv"), + read_bench(dynare_dir, "benchmark_jacobian.csv"), +) + +first_order_solve_metric(julia_dir, dynare_dir) = ( + read_bench(julia_dir, "benchmark_first_order_solve.csv"), + read_bench(dynare_dir, "benchmark_first_order_solve.csv"), +) + +first_order_total(julia_dir, dynare_dir) = ( + sum_bench_components(julia_dir, ["benchmark_jacobian.csv", "benchmark_first_order_solve.csv"]), + sum_bench_components(dynare_dir, ["benchmark_jacobian.csv", "benchmark_first_order_solve.csv"]), +) + +hessian_metric(julia_dir, dynare_dir) = ( + read_bench(julia_dir, "benchmark_hessian.csv"), + read_bench(dynare_dir, "benchmark_hessian.csv"), +) + +second_order_solve_metric(julia_dir, dynare_dir) = ( + read_bench(julia_dir, "benchmark_second_order_solve.csv"), + read_bench(dynare_dir, "benchmark_second_order_solve.csv"), +) + +second_order_total(julia_dir, dynare_dir) = ( + sum_bench_components(julia_dir, ["benchmark_hessian.csv", "benchmark_second_order_solve.csv"]), + sum_bench_components(dynare_dir, ["benchmark_hessian.csv", "benchmark_second_order_solve.csv"]), +) + +comparable_direct_total(julia_dir, dynare_dir) = ( + sum_bench_components(julia_dir, [ + "benchmark_jacobian.csv", + "benchmark_first_order_solve.csv", + "benchmark_hessian.csv", + "benchmark_second_order_solve.csv", + ]), + sum_bench_components(dynare_dir, [ + "benchmark_jacobian.csv", + "benchmark_first_order_solve.csv", + "benchmark_hessian.csv", + "benchmark_second_order_solve.csv", + ]), +) + +higher_order_bundled(julia_dir, dynare_dir) = ( + sum_optional_components(julia_dir, [ + "benchmark_first_order_solve.csv", + "benchmark_hessian.csv", + "benchmark_second_order_solve.csv", + "benchmark_third_order_derivatives.csv", + "benchmark_third_order_solve.csv", + ]), + read_dynare_k_order_pert(dynare_dir), +) + +third_order_derivatives_metric(julia_dir, dynare_dir) = ( + read_bench(julia_dir, "benchmark_third_order_derivatives.csv"), + NaN, +) + +third_order_solve_metric(julia_dir, dynare_dir) = ( + read_bench(julia_dir, "benchmark_third_order_solve.csv"), + NaN, +) + +function collect_rows(thread_counts, thread_dirs, metric_fn; model_dir_names = collect_model_names(thread_counts, thread_dirs)) + rows = Vector{Tuple{String, Vector{Float64}, Vector{Float64}}}() + + for model_name in model_dir_names + macro_values = Float64[] + dynare_values = Float64[] + has_any_value = false + + for thread_count in thread_counts + model_dir = joinpath(thread_dirs[thread_count], model_name) + if isdir(joinpath(model_dir, "julia")) && isdir(joinpath(model_dir, "dynare")) + macro_value, dynare_value = metric_fn(joinpath(model_dir, "julia"), joinpath(model_dir, "dynare")) + else + macro_value, dynare_value = NaN, NaN + end + + push!(macro_values, macro_value) + push!(dynare_values, dynare_value) + has_any_value |= !isnan(macro_value) || !isnan(dynare_value) + end + + has_any_value && push!(rows, (model_name, macro_values, dynare_values)) + end + + return rows +end + +function write_summary_csv(path, thread_counts, rows) + ncols = 1 + 2 * length(thread_counts) + table = Matrix{Any}(undef, length(rows) + 1, ncols) + table[1, 1] = "Model" + + column_index = 2 + for thread_count in thread_counts + table[1, column_index] = "MacroModelling_$(thread_count)" + table[1, column_index + 1] = "Dynare_$(thread_count)" + column_index += 2 + end + + for (row_index, (model_name, macro_values, dynare_values)) in enumerate(rows) + table[row_index + 1, 1] = model_name + column_index = 2 + for value_index in eachindex(thread_counts) + table[row_index + 1, column_index] = isnan(macro_values[value_index]) ? "" : macro_values[value_index] + table[row_index + 1, column_index + 1] = isnan(dynare_values[value_index]) ? "" : dynare_values[value_index] + column_index += 2 + end + end + + writedlm(path, table, ',') +end + +function print_summary_table(title, thread_counts, rows; note = "") + println("\n--- $title ---") + if !isempty(note) + println(" $note") + end + if isempty(rows) + println("No benchmark rows found.") + return + end + + header = rpad("Model", 50) + for thread_count in thread_counts + header *= rpad("MacroModelling_$(thread_count)", 18) + header *= rpad("Dynare_$(thread_count)", 18) + end + println(header) + println("-"^length(header)) + + for (model_name, macro_values, dynare_values) in rows + row_text = rpad(model_name, 50) + for value_index in eachindex(thread_counts) + row_text *= rpad(format_time(macro_values[value_index]), 18) + row_text *= rpad(format_time(dynare_values[value_index]), 18) + end + println(row_text) + end +end + +function compare_result_sets(reference_results, candidate_results; model_name) + moments_only_higher_order = is_higher_order_model(model_name) + skip_pruned_third_order = is_pruned_third_order_model(model_name) + + @testset "Steady State" begin + compare_steady_state(reference_results, candidate_results) + end + @testset "Policy Matrix ghx" begin + if skip_pruned_third_order + @info "Skipping ghx comparison for $model_name (pruned third-order state representation mismatch)" + elseif moments_only_higher_order && get(reference_results, :policy_algorithm, "") != "first_order" + @info "Skipping ghx comparison for $model_name (policy matrices not tagged as first-order; regenerate phase-1 outputs to enable)" + else + compare_ghx(reference_results, candidate_results) + end + end + @testset "Policy Matrix ghu" begin + if moments_only_higher_order && get(reference_results, :policy_algorithm, "") != "first_order" + @info "Skipping ghu comparison for $model_name (policy matrices not tagged as first-order; regenerate phase-1 outputs to enable)" + else + compare_ghu(reference_results, candidate_results) + end + end + @testset "IRFs" begin + compare_irfs(reference_results, candidate_results; model_name = model_name) + end + @testset "Variance" begin + if skip_pruned_third_order && !is_supported_pruned_third_order_variance_model(model_name) + @info "Skipping variance comparison for $model_name (pruned third-order moment convention mismatch outside the validated benchmark cases)" + else + compare_variance(reference_results, candidate_results) + end + end + if has_second_order(reference_results) && has_second_order(candidate_results) + @testset "Second Order Matrices" begin + if moments_only_higher_order + @info "Skipping second-order matrix comparison for $model_name (higher-order configured as moments-only)" + else + compare_second_order(reference_results, candidate_results) + end + end + end + if has_third_order(reference_results) && has_third_order(candidate_results) + @testset "Third Order Matrices" begin + if moments_only_higher_order + @info "Skipping third-order matrix comparison for $model_name (higher-order configured as moments-only)" + else + compare_third_order(reference_results, candidate_results) + end + end + end +end + +function compare_thread_consistency(thread_counts, thread_dirs) + if length(thread_counts) <= 1 + @info "Only one thread count detected; skipping cross-thread correctness comparison" + return + end + + reference_thread = first(thread_counts) + all_model_dirs = collect_model_names(thread_counts, thread_dirs) + comparison_model_dirs = collect_model_names(thread_counts, thread_dirs; include_benchmark_only = false) + + for mname in sort(filter(is_benchmark_only_model_dir, all_model_dirs)) + @info "Skipping cross-thread correctness comparison for benchmark-only model: $mname" + end + + isempty(comparison_model_dirs) && return + + comparison_exception = nothing + try + @testset "Cross-Thread Consistency (reference: threads_$reference_thread)" begin + for mname in sort(comparison_model_dirs) + reference_model_dir = joinpath(thread_dirs[reference_thread], mname) + reference_julia_dir = joinpath(reference_model_dir, "julia") + reference_dynare_dir = joinpath(reference_model_dir, "dynare") + + @info "Comparing cross-thread results for: $mname" + @testset "$mname" begin + if !isdir(reference_julia_dir) || !isdir(reference_dynare_dir) + @test isdir(reference_julia_dir) + @test isdir(reference_dynare_dir) + continue + end + + reference_julia_results = load_results(reference_julia_dir) + reference_dynare_results = load_results(reference_dynare_dir) + + for thread_count in thread_counts[2:end] + candidate_model_dir = joinpath(thread_dirs[thread_count], mname) + candidate_julia_dir = joinpath(candidate_model_dir, "julia") + candidate_dynare_dir = joinpath(candidate_model_dir, "dynare") + + @testset "threads_$thread_count" begin + if !isdir(candidate_julia_dir) || !isdir(candidate_dynare_dir) + @test isdir(candidate_julia_dir) + @test isdir(candidate_dynare_dir) + continue + end + + candidate_julia_results = load_results(candidate_julia_dir) + candidate_dynare_results = load_results(candidate_dynare_dir) + + @testset "MacroModelling" begin + compare_result_sets(reference_julia_results, candidate_julia_results; model_name = mname) + end + @testset "Dynare" begin + compare_result_sets(reference_dynare_results, candidate_dynare_results; model_name = mname) + end + end + end + end + end + end + catch err + if err isa Test.TestSetException + comparison_exception = err + else + rethrow(err) + end + end + + comparison_exception === nothing || throw(comparison_exception) +end + +function main(args = ARGS) + output_root = parse_args(args) + output_root === nothing && return + isdir(output_root) || error("Output directory not found: $output_root") + + thread_counts, thread_dirs = collect_thread_dirs(output_root) + + println("Thread sweep output root: $output_root") + println("Detected thread counts: $(join(string.(thread_counts), ", "))") + print_environment_summary(thread_counts, thread_dirs) + compare_thread_consistency(thread_counts, thread_dirs) + + println("\n", "="^100) + println(" Benchmark Summary by Thread: MacroModelling vs Dynare") + println("="^100) + + summaries = [ + ("Jacobian", "benchmark_jacobian_by_thread.csv", jacobian_metric, ""), + ("First-Order Solve", "benchmark_first_order_solve_by_thread.csv", first_order_solve_metric, ""), + ("First-Order Total (sum of direct Jacobian + solve medians)", "benchmark_first_order_total_by_thread.csv", first_order_total, ""), + ("Hessian", "benchmark_hessian_by_thread.csv", hessian_metric, ""), + ("Second-Order Solve", "benchmark_second_order_solve_by_thread.csv", second_order_solve_metric, ""), + ("Second-Order Total (Hessian + Second-Order Solve)", "benchmark_second_order_total_by_thread.csv", second_order_total, ""), + ("Comparable Direct Components Total (Jacobian + FO + Hessian + SO)", "benchmark_comparable_direct_total_by_thread.csv", comparable_direct_total, ""), + ("Higher-Order Bundled (Dynare k_order_pert)", "benchmark_higher_order_bundled_by_thread.csv", higher_order_bundled, + "MacroModelling sums directly measured solve-stack components; Dynare reports direct bundled k_order_pert"), + ("Third-Order Derivatives (MacroModelling only)", "benchmark_third_order_derivatives_by_thread.csv", third_order_derivatives_metric, + "Dynare does not export a directly comparable component-level metric."), + ("Third-Order Solve (MacroModelling only)", "benchmark_third_order_solve_by_thread.csv", third_order_solve_metric, + "Dynare does not export a directly comparable component-level metric."), + ] + + model_dir_names = collect_model_names(thread_counts, thread_dirs) + for (title, file_name, metric_fn, note) in summaries + rows = collect_rows(thread_counts, thread_dirs, metric_fn; model_dir_names = model_dir_names) + print_summary_table(title, thread_counts, rows; note = note) + if !isempty(rows) + summary_path = joinpath(output_root, file_name) + write_summary_csv(summary_path, thread_counts, rows) + println("Wrote $(summary_path)") + end + end + + println("="^100) +end + +if abspath(PROGRAM_FILE) == @__FILE__ + main() +end \ No newline at end of file diff --git a/test/dynare_comparison/extract_dynare_results.m b/test/dynare_comparison/extract_dynare_results.m new file mode 100644 index 000000000..af27e0d75 --- /dev/null +++ b/test/dynare_comparison/extract_dynare_results.m @@ -0,0 +1,346 @@ +% extract_dynare_results.m +% Extracts Dynare results after stoch_simul and saves them as CSV files. +% +% Expects: +% - model_name: string variable set before calling this script +% - output_dir: string variable for the output directory +% - Dynare's oo_, M_, options_ structures populated after stoch_simul +% +% Outputs (all in output_dir/): +% steady_state.csv - steady state values (declaration order) +% var_names.csv - endogenous variable names (declaration order) +% exo_names.csv - exogenous shock names +% state_var_names.csv - state variable names +% ghx.csv - state transition matrix (declaration order rows) +% ghu.csv - shock impact matrix (declaration order rows) +% irf_VARNAME_SHOCKNAME.csv - IRF for each var/shock combination +% variance_covariance.csv - theoretical variance-covariance matrix (declaration order) +% +% Higher-order outputs (when options_.order >= 2): +% ghxx.csv, ghxu.csv, ghuu.csv, ghs2.csv +% Higher-order outputs (when options_.order >= 3): +% ghxxx.csv, ghxxu.csv, ghxuu.csv, ghuuu.csv, ghxss.csv, ghuss.csv + +if ~exist('output_dir', 'var') + output_dir = [model_name '_results']; +end +mkdir(output_dir); + +if ~exist('benchmark_only_mode', 'var') || isempty(benchmark_only_mode) + benchmark_only_mode = false; +end + +n_endo = M_.endo_nbr; +n_exo = M_.exo_nbr; + +if ~benchmark_only_mode +%% --- Variable names --- +fid = fopen(fullfile(output_dir, 'var_names.csv'), 'w'); +for i = 1:n_endo + if iscell(M_.endo_names) + fprintf(fid, '%s\n', M_.endo_names{i}); + else + fprintf(fid, '%s\n', deblank(M_.endo_names(i,:))); + end +end +fclose(fid); + +fid = fopen(fullfile(output_dir, 'exo_names.csv'), 'w'); +for i = 1:n_exo + if iscell(M_.exo_names) + fprintf(fid, '%s\n', M_.exo_names{i}); + else + fprintf(fid, '%s\n', deblank(M_.exo_names(i,:))); + end +end +fclose(fid); + +%% --- Steady state (declaration order) --- +dlmwrite(fullfile(output_dir, 'steady_state.csv'), oo_.steady_state, 'precision', '%.16g'); + +%% --- Policy matrices (convert from DR order to declaration order) --- +ghx_dr = oo_.dr.ghx; +ghu_dr = oo_.dr.ghu; + +ghx_decl = zeros(n_endo, size(ghx_dr, 2)); +ghu_decl = zeros(n_endo, size(ghu_dr, 2)); +ghx_decl(oo_.dr.order_var, :) = ghx_dr; +ghu_decl(oo_.dr.order_var, :) = ghu_dr; + +dlmwrite(fullfile(output_dir, 'ghx.csv'), ghx_decl, 'precision', '%.16g'); +dlmwrite(fullfile(output_dir, 'ghu.csv'), ghu_decl, 'precision', '%.16g'); + +% State variable names +if isfield(oo_, 'dr') && isfield(oo_.dr, 'state_var') && ~isempty(oo_.dr.state_var) + state_var_idx = oo_.dr.state_var; +elseif isfield(M_, 'state_var') && ~isempty(M_.state_var) + % Dynare may store state metadata in M_ for some solver paths. + if isnumeric(M_.state_var) + state_var_idx = M_.state_var; + elseif isstruct(M_.state_var) + if isfield(M_.state_var, 'decl') + state_var_idx = M_.state_var.decl; + elseif isfield(M_.state_var, 'idx') + state_var_idx = M_.state_var.idx; + else + state_var_idx = find(M_.lead_lag_incidence(1, :)); + end + else + state_var_idx = find(M_.lead_lag_incidence(1, :)); + end +else + % Robust fallback: lagged endogenous variables in declaration order. + state_var_idx = find(M_.lead_lag_incidence(1, :)); +end + +state_var_idx = state_var_idx(:); +fid = fopen(fullfile(output_dir, 'state_var_names.csv'), 'w'); +for i = 1:length(state_var_idx) + si = state_var_idx(i); + if iscell(M_.endo_names) + fprintf(fid, '%s\n', M_.endo_names{si}); + else + fprintf(fid, '%s\n', deblank(M_.endo_names(si,:))); + end +end +fclose(fid); + +%% --- Second-order matrices (when order >= 2) --- +if options_.order >= 2 && isfield(oo_.dr, 'ghxx') + ghxx_dr = oo_.dr.ghxx; + ghxu_dr = oo_.dr.ghxu; + ghuu_dr = oo_.dr.ghuu; + ghs2_dr = oo_.dr.ghs2; + + ghxx_decl = zeros(n_endo, size(ghxx_dr, 2)); + ghxu_decl = zeros(n_endo, size(ghxu_dr, 2)); + ghuu_decl = zeros(n_endo, size(ghuu_dr, 2)); + ghs2_decl = zeros(n_endo, 1); + + ghxx_decl(oo_.dr.order_var, :) = ghxx_dr; + ghxu_decl(oo_.dr.order_var, :) = ghxu_dr; + ghuu_decl(oo_.dr.order_var, :) = ghuu_dr; + ghs2_decl(oo_.dr.order_var, :) = ghs2_dr; + + dlmwrite(fullfile(output_dir, 'ghxx.csv'), ghxx_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghxu.csv'), ghxu_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghuu.csv'), ghuu_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghs2.csv'), ghs2_decl, 'precision', '%.16g'); + + fprintf('Extracted second-order matrices: ghxx(%dx%d) ghxu(%dx%d) ghuu(%dx%d) ghs2(%dx1)\n', ... + size(ghxx_decl,1), size(ghxx_decl,2), ... + size(ghxu_decl,1), size(ghxu_decl,2), ... + size(ghuu_decl,1), size(ghuu_decl,2), ... + size(ghs2_decl,1)); +end + +%% --- Third-order matrices (when order >= 3) --- +if options_.order >= 3 && isfield(oo_.dr, 'ghxxx') + ghxxx_dr = oo_.dr.ghxxx; + ghxxu_dr = oo_.dr.ghxxu; + ghxuu_dr = oo_.dr.ghxuu; + ghuuu_dr = oo_.dr.ghuuu; + ghxss_dr = oo_.dr.ghxss; + ghuss_dr = oo_.dr.ghuss; + + ghxxx_decl = zeros(n_endo, size(ghxxx_dr, 2)); + ghxxu_decl = zeros(n_endo, size(ghxxu_dr, 2)); + ghxuu_decl = zeros(n_endo, size(ghxuu_dr, 2)); + ghuuu_decl = zeros(n_endo, size(ghuuu_dr, 2)); + ghxss_decl = zeros(n_endo, size(ghxss_dr, 2)); + ghuss_decl = zeros(n_endo, size(ghuss_dr, 2)); + + ghxxx_decl(oo_.dr.order_var, :) = ghxxx_dr; + ghxxu_decl(oo_.dr.order_var, :) = ghxxu_dr; + ghxuu_decl(oo_.dr.order_var, :) = ghxuu_dr; + ghuuu_decl(oo_.dr.order_var, :) = ghuuu_dr; + ghxss_decl(oo_.dr.order_var, :) = ghxss_dr; + ghuss_decl(oo_.dr.order_var, :) = ghuss_dr; + + dlmwrite(fullfile(output_dir, 'ghxxx.csv'), ghxxx_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghxxu.csv'), ghxxu_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghxuu.csv'), ghxuu_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghuuu.csv'), ghuuu_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghxss.csv'), ghxss_decl, 'precision', '%.16g'); + dlmwrite(fullfile(output_dir, 'ghuss.csv'), ghuss_decl, 'precision', '%.16g'); + + fprintf('Extracted third-order matrices: ghxxx(%dx%d) ghxxu(%dx%d) ghxuu(%dx%d) ghuuu(%dx%d) ghxss(%dx%d) ghuss(%dx%d)\n', ... + size(ghxxx_decl,1), size(ghxxx_decl,2), ... + size(ghxxu_decl,1), size(ghxxu_decl,2), ... + size(ghxuu_decl,1), size(ghxuu_decl,2), ... + size(ghuuu_decl,1), size(ghuuu_decl,2), ... + size(ghxss_decl,1), size(ghxss_decl,2), ... + size(ghuss_decl,1), size(ghuss_decl,2)); +end + +%% --- IRFs (bundled into a single matrix file) --- +if isfield(oo_, 'irfs') + irf_fields = fieldnames(oo_.irfs); + % Determine number of periods from the first non-empty field. + n_periods = 0; + for i = 1:length(irf_fields) + d = oo_.irfs.(irf_fields{i}); + if ~isempty(d) + n_periods = numel(d); + break; + end + end + if n_periods > 0 + irf_matrix = zeros(n_periods, length(irf_fields)); + for i = 1:length(irf_fields) + data = oo_.irfs.(irf_fields{i}); + if ~isempty(data) + irf_matrix(:, i) = data(:); + end + end + dlmwrite(fullfile(output_dir, 'irfs.csv'), irf_matrix, 'precision', '%.16g'); + end + fid = fopen(fullfile(output_dir, 'irf_fields.csv'), 'w'); + for i = 1:length(irf_fields) + fprintf(fid, '%s\n', irf_fields{i}); + end + fclose(fid); +end + +%% --- Variance-covariance matrix (declaration order) --- +if isfield(oo_, 'var') && ~isempty(oo_.var) + dlmwrite(fullfile(output_dir, 'variance_covariance.csv'), oo_.var, 'precision', '%.16g'); +end + +end % if ~benchmark_only_mode + +%% --- Benchmark: component-level timing --- +% Decomposes the solution pipeline into individually timed components. +% For all models: Jacobian, first-order solve, [Hessian, second-order solve] +% For k_order models (order 3): also export bundled k_order_pert as an additional direct reference. +bench_section_tic = tic; +n_bench = 500; + +exo_ss = oo_.exo_steady_state; +if isfield(oo_, 'exo_det_steady_state') + exo_det_ss = oo_.exo_det_steady_state; +else + exo_det_ss = zeros(M_.exo_det_nbr, 1); +end +exo_ss_full = [exo_ss; exo_det_ss]; + +% Decompose stochastic_solvers into individual components for every order. +dyn_endo_ss = repmat(oo_.dr.ys, 3, 1); + +% ── Jacobian (dynamic_g1) ── +bench_times_jac = zeros(1, n_bench); +if options_.order >= 2 + % order >= 2 needs T_order, T outputs for Hessian computation + for i = 1:n_bench + tic; + [g1_bench, T_order_bench, T_bench] = feval([M_.fname '.dynamic_g1'], ... + dyn_endo_ss, exo_ss_full, M_.params, oo_.dr.ys, ... + M_.dynamic_g1_sparse_rowval, M_.dynamic_g1_sparse_colval, ... + M_.dynamic_g1_sparse_colptr); + bench_times_jac(i) = toc; + end +else + for i = 1:n_bench + tic; + g1_bench = feval([M_.fname '.dynamic_g1'], ... + dyn_endo_ss, exo_ss_full, M_.params, oo_.dr.ys, ... + M_.dynamic_g1_sparse_rowval, M_.dynamic_g1_sparse_colval, ... + M_.dynamic_g1_sparse_colptr); + bench_times_jac(i) = toc; + end +end +median_jac = median(bench_times_jac); +bench_keys = {'benchmark_jacobian'}; +bench_values = [median_jac]; + +% ── First-order solve (dyn_first_order_solver) ── +dr_bench = oo_.dr; +bench_times_fo = zeros(1, n_bench); +for i = 1:n_bench + tic; + [dr_bench, ~] = dyn_first_order_solver(g1_bench, M_, dr_bench, options_, 0); + bench_times_fo(i) = toc; +end +median_fo = median(bench_times_fo); +bench_keys{end+1} = 'benchmark_first_order_solve'; +bench_values(end+1) = median_fo; + +median_first_order_total = median_jac + median_fo; +bench_keys{end+1} = 'benchmark_first_order_total'; +bench_values(end+1) = median_first_order_total; +bench_keys{end+1} = 'benchmark_first_order'; +bench_values(end+1) = median_first_order_total; + +fprintf('Benchmark %s (order=%d): Jac=%.1f us, FO_solve=%.1f us', ... + model_name, options_.order, median_jac*1e6, median_fo*1e6); + +if options_.order >= 2 + % ── Hessian (dynamic_g2 + build_two_dim_hessian) ── + bench_times_hess = zeros(1, n_bench); + for i = 1:n_bench + tic; + g2_v_bench = feval([M_.fname '.dynamic_g2'], dyn_endo_ss, exo_ss_full, ... + M_.params, oo_.dr.ys, T_order_bench, T_bench); + g2_bench = build_two_dim_hessian(M_.dynamic_g2_sparse_indices, g2_v_bench, ... + size(g1_bench, 1), size(g1_bench, 2)); + bench_times_hess(i) = toc; + end + median_hess = median(bench_times_hess); + bench_keys{end+1} = 'benchmark_hessian'; + bench_values(end+1) = median_hess; + + % ── Second-order solve (dyn_second_order_solver) ── + bench_times_so = zeros(1, n_bench); + for i = 1:n_bench + tic; + dr_bench = dyn_second_order_solver(g1_bench, g2_bench, dr_bench, M_, ... + options_.threads.kronecker.sparse_hessian_times_B_kronecker_C); + bench_times_so(i) = toc; + end + median_so = median(bench_times_so); + bench_keys{end+1} = 'benchmark_second_order_solve'; + bench_values(end+1) = median_so; + + fprintf(', Hess=%.1f us, SO_solve=%.1f us', median_hess*1e6, median_so*1e6); +end + +if options_.k_order_solver + % k_order_pert remains useful as a directly measured bundled reference for order-3 runs. + dr_korder = struct(); + if isfield(oo_.dr, 'inv_order_var'); dr_korder.inv_order_var = oo_.dr.inv_order_var; end + if isfield(oo_.dr, 'order_var'); dr_korder.order_var = oo_.dr.order_var; end + if isfield(oo_.dr, 'restrict_var_list'); dr_korder.restrict_var_list = oo_.dr.restrict_var_list; end + if isfield(oo_.dr, 'restrict_columns'); dr_korder.restrict_columns = oo_.dr.restrict_columns; end + if isfield(oo_.dr, 'obs_var'); dr_korder.obs_var = oo_.dr.obs_var; end + dr_korder.ys = oo_.dr.ys; + + bench_times_korder = zeros(1, n_bench); + for i = 1:n_bench + dr_tmp = set_state_space(dr_korder, M_); + tic; + [dr_tmp, ~] = k_order_pert(dr_tmp, M_, options_); + bench_times_korder(i) = toc; + end + median_korder = median(bench_times_korder); + bench_keys{end+1} = 'benchmark_k_order_pert'; + bench_values(end+1) = median_korder; + + fprintf(', k_order_pert=%.1f us', median_korder*1e6); +end + +fprintf(', FO_Total=%.1f us over %d runs\n', median_first_order_total*1e6, n_bench); + +%% --- Write bundled benchmarks --- +bench_fid = fopen(fullfile(output_dir, 'benchmarks.csv'), 'w'); +for bi = 1:numel(bench_keys) + fprintf(bench_fid, '%s,%.16g\n', bench_keys{bi}, bench_values(bi)); +end +fclose(bench_fid); + +% Write elapsed time spent on benchmark loops (for wall-clock subtraction) +bench_elapsed_total = toc(bench_section_tic); +bench_elapsed_fid = fopen(fullfile(output_dir, 'bench_elapsed_seconds.csv'), 'w'); +fprintf(bench_elapsed_fid, '%.6f\n', bench_elapsed_total); +fclose(bench_elapsed_fid); + +disp(['Results extracted to: ' output_dir]); diff --git a/test/dynare_comparison/generate_julia_results.jl b/test/dynare_comparison/generate_julia_results.jl new file mode 100644 index 000000000..fdaf1d8bc --- /dev/null +++ b/test/dynare_comparison/generate_julia_results.jl @@ -0,0 +1,731 @@ +# generate_julia_results.jl — Phase 1 of Dynare comparison +# +# Loads models, exports .mod files, and saves Julia-computed results as CSV. +# Output structure: +# output/{model_name}/ +# {model_name}.mod +# julia/ +# var_names.csv, exo_names.csv, state_var_names.csv +# steady_state.csv, ghx.csv, ghu.csv +# irf_fields.csv, irf_{var}_{shock}.csv +# variance_covariance.csv +# +# For higher-order models (pruned 2nd/3rd order): +# output/{model_name}_pruned_2nd/ and output/{model_name}_pruned_3rd/ +# Includes first-order comparable outputs plus higher-order moments: +# steady_state.csv, ghx.csv, ghu.csv, irf_*.csv +# variance_covariance.csv +# Excludes higher-order solution-matrix CSVs (ghxx/ghxu/..., ghxxx/...) + +# On Windows we deliberately switch the LinearAlgebra BLAS backend to MKL so the +# Julia side mirrors what Dynare/MATLAB use. MKL.jl must be loaded BEFORE any +# BLAS calls (including BLAS.set_num_threads) for it to take effect. On non- +# Windows platforms (e.g. Linux CI runners where MKL.jl may not be installed) +# we silently fall back to the default OpenBLAS backend. + +# @static if Sys.iswindows() +# using MKL +# MKL.set_num_threads(Threads.nthreads()) +# @info "Using MKL.jl for BLAS on Windows with $(MKL.get_num_threads()) threads" +# end + +using MacroModelling +using DelimitedFiles +using LinearAlgebra +using Sockets + +const IRF_PERIODS = 40 +const DEFAULT_OUTPUT_ROOT = joinpath(@__DIR__, "output") +const MODELS_DIR = joinpath(@__DIR__, "..", "..", "models") + +# Models to test (first order) +const MODEL_FILES = [ + "FS2000", + "Ascari_Sbordone_2014", + "Gali_2015_chapter_3_nonlinear", + "Caldara_et_al_2012", + "Smets_Wouters_2007", + "Smets_Wouters_2003", + "GNSS_2010", + "NAWM_EAUS_2008", + "QUEST3_2009", + "FRBUS", +] + +# Models to also test at pruned 2nd order +const SECOND_ORDER_MODELS = [ + "FS2000", + "Ascari_Sbordone_2014", + "Gali_2015_chapter_3_nonlinear", + "Caldara_et_al_2012", + "Smets_Wouters_2003", + "Smets_Wouters_2007", + "GNSS_2010", +] + +# Models to also test at pruned 3rd order +const THIRD_ORDER_MODELS = [ + "FS2000", + "Ascari_Sbordone_2014", + "Gali_2015_chapter_3_nonlinear", + "Caldara_et_al_2012", +] + +# Models that skip variance/covariance +const SKIP_MOMENTS_MODELS = Set(["FRBUS", "NAWM"]) + +# Models for which only the benchmark timings are exported (no names, steady state, +# policy matrices, IRFs, or moments). The .mod file is still written so the Dynare +# phase can run and produce its own benchmark CSVs. +const BENCHMARK_ONLY_MODELS = Set(["FRBUS", "NAWM"]) + +# ───────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────── + +function ascii_name(sym::Symbol) + MacroModelling.translate_symbol_to_ascii(sym) +end + +function original_vars(model) + setdiff(model.constants.post_model_macro.vars_in_ss_equations, + model.constants.post_model_macro.➕_vars) +end + +function write_names(path, names) + open(path, "w") do io + for n in names + println(io, n) + end + end +end + +function print_usage() + println("Usage: julia --project=. --threads=N generate_julia_results.jl [--output-root=PATH | PATH]") +end + +function parse_args(args) + output_root = DEFAULT_OUTPUT_ROOT + positional_args = String[] + only_models = String[] + + for arg in args + if arg in ("-h", "--help") + print_usage() + return nothing + elseif startswith(arg, "--output-root=") + output_root = split(arg, "=", limit = 2)[2] + elseif startswith(arg, "--only-models=") + value = split(arg, "=", limit = 2)[2] + for token in split(value, ',') + trimmed = strip(token) + if !isempty(trimmed) + push!(only_models, String(trimmed)) + end + end + elseif startswith(arg, "--") + error("Unknown option: $arg") + else + push!(positional_args, arg) + end + end + + if length(positional_args) > 1 + error("Expected at most one positional output-root argument, got $(length(positional_args))") + elseif length(positional_args) == 1 + output_root = positional_args[1] + end + + env_only = get(ENV, "DYNARE_COMPARE_ONLY_MODELS", "") + if isempty(only_models) && !isempty(env_only) + for token in split(env_only, ',') + trimmed = strip(token) + if !isempty(trimmed) + push!(only_models, String(trimmed)) + end + end + end + + return (abspath(output_root), only_models) +end + +function configure_julia_threads!() + julia_threads = Threads.nthreads() + BLAS.set_num_threads(julia_threads) + blas_threads = BLAS.get_num_threads() + println("Julia thread configuration: julia_threads=$julia_threads blas_threads=$blas_threads") + println(" Threads.nthreads() = ", Threads.nthreads()) + println(" Threads.nthreads(:default) = ", Threads.nthreads(:default)) + println(" Threads.nthreads(:interactive) = ", Threads.nthreads(:interactive)) + println(" BLAS.get_num_threads() = ", BLAS.get_num_threads()) + blas_vendor = try + string(BLAS.get_config()) + catch + "unknown" + end + println(" BLAS vendor: $blas_vendor") + for var in ("JULIA_NUM_THREADS", "OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", + "MKL_NUM_THREADS", "BLIS_NUM_THREADS", "VECLIB_MAXIMUM_THREADS") + println(" ENV $var=", get(ENV, var, "")) + end + @info "Julia thread configuration" julia_threads blas_threads + return julia_threads, blas_threads +end + +function write_thread_configuration(output_root, julia_threads, blas_threads) + open(joinpath(output_root, "julia_thread_configuration.txt"), "w") do io + println(io, "julia_threads=$julia_threads") + println(io, "blas_threads=$blas_threads") + end +end + +sanitize_metadata_value(value) = replace(string(value), r"[\r\n]+" => " | ") + +function machine_hostname() + try + return Sockets.gethostname() + catch + return get(ENV, "COMPUTERNAME", get(ENV, "HOSTNAME", "unknown")) + end +end + +function write_key_value_metadata(path, entries) + open(path, "w") do io + for (key, value) in entries + println(io, key, "=", sanitize_metadata_value(value)) + end + end +end + +function write_julia_environment_metadata(output_root, julia_threads, blas_threads) + blas_lapack = try + string(BLAS.get_config()) + catch + "unknown" + end + total_memory_bytes = try + Sys.total_memory() + catch + "unknown" + end + + entries = [ + "julia_version" => VERSION, + "julia_threads" => julia_threads, + "julia_threads_default" => Threads.nthreads(:default), + "julia_threads_interactive" => Threads.nthreads(:interactive), + "blas_threads" => blas_threads, + "blas_lapack" => blas_lapack, + "hostname" => machine_hostname(), + "kernel" => Sys.KERNEL, + "arch" => Sys.ARCH, + "cpu_name" => Sys.CPU_NAME, + "cpu_threads" => Sys.CPU_THREADS, + "word_size" => Sys.WORD_SIZE, + "total_memory_bytes" => total_memory_bytes, + ] + + write_key_value_metadata(joinpath(output_root, "comparison_environment_julia.txt"), entries) +end + +function export_names_and_steady_state(model, julia_dir, orig, state_vars, exo_vars) + var_names_ascii = [ascii_name(v) for v in orig] + exo_names_ascii = [ascii_name(e) for e in exo_vars] + state_names_ascii = [ascii_name(s) for s in state_vars] + + write_names(joinpath(julia_dir, "var_names.csv"), var_names_ascii) + write_names(joinpath(julia_dir, "exo_names.csv"), exo_names_ascii) + write_names(joinpath(julia_dir, "state_var_names.csv"), state_names_ascii) + + ss = get_SS(model, derivatives = false) + ss_vals = [Float64(ss(v)) for v in orig] + writedlm(joinpath(julia_dir, "steady_state.csv"), ss_vals, ',') + + return var_names_ascii, exo_names_ascii +end + +function export_first_order_matrices(model, julia_dir, orig, state_vars, exo_vars) + sol = get_solution(model, algorithm = :first_order) + + ghx = zeros(length(orig), length(state_vars)) + for (si, s) in enumerate(state_vars) + s_key = Symbol(string(s) * "₍₋₁₎") + for (vi, v) in enumerate(orig) + ghx[vi, si] = Float64(sol(s_key, v)) + end + end + writedlm(joinpath(julia_dir, "ghx.csv"), ghx, ',') + + ghu = zeros(length(orig), length(exo_vars)) + for (ei, e) in enumerate(exo_vars) + e_key = Symbol(string(e) * "₍ₓ₎") + for (vi, v) in enumerate(orig) + ghu[vi, ei] = Float64(sol(e_key, v)) + end + end + writedlm(joinpath(julia_dir, "ghu.csv"), ghu, ',') + write_names(joinpath(julia_dir, "policy_algorithm.csv"), ["first_order"]) +end + +function export_irfs(model, julia_dir, orig, exo_vars; algorithm = :first_order) + irfs = get_irf(model, periods = IRF_PERIODS, algorithm = algorithm) + irf_fields = String[] + n_cols = length(orig) * length(exo_vars) + irf_matrix = Matrix{Float64}(undef, IRF_PERIODS, n_cols) + col = 0 + for v in orig + v_ascii = ascii_name(v) + for e in exo_vars + e_ascii = ascii_name(e) + push!(irf_fields, "$(v_ascii)_$(e_ascii)") + col += 1 + for t in 1:IRF_PERIODS + irf_matrix[t, col] = Float64(irfs(v, t, e)) + end + end + end + # Single bundled file: rows = periods, cols = irf_fields (in same order). + writedlm(joinpath(julia_dir, "irfs.csv"), irf_matrix, ',') + write_names(joinpath(julia_dir, "irf_fields.csv"), irf_fields) + write_names(joinpath(julia_dir, "irf_algorithm.csv"), [String(algorithm)]) +end + +function write_benchmarks(julia_dir, bench::AbstractDict) + keys_sorted = sort(collect(keys(bench))) + table = Matrix{Any}(undef, length(keys_sorted), 2) + for (i, k) in enumerate(keys_sorted) + table[i, 1] = k + table[i, 2] = bench[k] + end + writedlm(joinpath(julia_dir, "benchmarks.csv"), table, ',') +end + +function export_moments(model, julia_dir, orig, exo_vars; + algorithm = :first_order, + var_names_ascii = nothing, + exo_names_ascii = nothing) + moments = get_moments(model, algorithm = algorithm, + derivatives = false, + non_stochastic_steady_state = false, + mean = false, + variance = true, + standard_deviation = false, + covariance = true) + + vcov = zeros(length(orig), length(orig)) + covar_ka = moments[:covariance] + for (ri, rv) in enumerate(orig) + for (ci, cv) in enumerate(orig) + vcov[ri, ci] = Float64(covar_ka(rv, cv)) + end + end + writedlm(joinpath(julia_dir, "variance_covariance.csv"), vcov, ',') +end + +# ───────────────────────────────────────────── +# Benchmark helpers — manual median of N_BENCH runs +# ───────────────────────────────────────────── +const N_BENCH = 500 + +function median_elapsed(f, n = N_BENCH) + times = Vector{Float64}(undef, n) + for i in 1:n + times[i] = @elapsed f() + end + sort!(times) + m = length(times) ÷ 2 + return isodd(length(times)) ? times[m + 1] : (times[m] + times[m + 1]) / 2 +end + +function benchmark_first_order(model, bench::AbstractDict) + params = copy(model.parameter_values) + opts = MacroModelling.merge_calculation_options(verbose = true) + + # Warm up to ensure functions are compiled and reusable inputs are available + MacroModelling.invalidate_cache_validity!(model) + SS_and_pars, _ = MacroModelling.get_NSSS_and_parameters(model, params, opts = opts, caching = false) + ∇₁ = MacroModelling.calculate_jacobian(params, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces, caching = false) + MacroModelling.calculate_first_order_solution(∇₁, model.constants, model.workspaces, model.caches; + opts = opts, initial_guess = model.caches.qme_solution, + parameter_values = params, caching = false) + + opts = MacroModelling.merge_calculation_options() + + # Benchmark Jacobian (given precomputed steady-state inputs) + median_jac = median_elapsed() do + MacroModelling.calculate_jacobian(params, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces, caching = false) + end + bench["benchmark_jacobian"] = median_jac + + # Benchmark first-order solve (given Jacobian) + median_fo = median_elapsed() do + MacroModelling.calculate_first_order_solution(∇₁, model.constants, model.workspaces, model.caches; + opts = opts, initial_guess = model.caches.qme_solution, + parameter_values = params, caching = false) + end + median_fo_total = median_jac + median_fo + bench["benchmark_first_order_solve"] = median_fo + bench["benchmark_first_order_total"] = median_fo_total + bench["benchmark_first_order"] = median_fo_total + + @info "Benchmark $(model.model_name) [first order]:" + @info " Jacobian: $(round(median_jac*1e6, digits=1)) μs" + @info " QME solve: $(round(median_fo*1e6, digits=1)) μs" + @info " Total: $(round(median_fo_total*1e6, digits=1)) μs" +end + +function benchmark_second_order(model, bench::AbstractDict) + params = copy(model.parameter_values) + opts = MacroModelling.merge_calculation_options() + + # Warm up: run full pipeline once + MacroModelling.invalidate_cache_validity!(model) + SS_and_pars, _ = MacroModelling.get_NSSS_and_parameters(model, params, opts = opts, caching = false) + ∇₁ = MacroModelling.calculate_jacobian(params, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces, caching = false) + 𝐒₁, _, _ = MacroModelling.calculate_first_order_solution(∇₁, model.constants, model.workspaces, model.caches; + opts = opts, initial_guess = model.caches.qme_solution, + parameter_values = params, caching = false) + ∇₂ = MacroModelling.calculate_hessian(params, SS_and_pars, model.caches, model.functions.hessian, model.workspaces, caching = false) + MacroModelling.calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, model.constants, model.workspaces, model.caches; + initial_guess = model.caches.second_order_solution, + opts = opts, parameter_values = params, caching = false) + + # Benchmark Hessian + median_hess = median_elapsed() do + MacroModelling.calculate_hessian(params, SS_and_pars, model.caches, model.functions.hessian, model.workspaces, caching = false) + end + bench["benchmark_hessian"] = median_hess + + # Benchmark second-order solve (given first-order solution + Hessian) + median_so = median_elapsed() do + MacroModelling.calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, model.constants, model.workspaces, model.caches; + initial_guess = model.caches.second_order_solution, + opts = opts, parameter_values = params, caching = false) + end + bench["benchmark_second_order_solve"] = median_so + + @info "Benchmark $(model.model_name) [second order]:" + @info " Hessian: $(round(median_hess*1e6, digits=1)) μs" + @info " 2nd order solve: $(round(median_so*1e6, digits=1)) μs" +end + +function benchmark_third_order(model, bench::AbstractDict) + params = copy(model.parameter_values) + opts = MacroModelling.merge_calculation_options() + + # Warm up: run full pipeline once + MacroModelling.invalidate_cache_validity!(model) + SS_and_pars, _ = MacroModelling.get_NSSS_and_parameters(model, params, opts = opts, caching = false) + ∇₁ = MacroModelling.calculate_jacobian(params, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces, caching = false) + 𝐒₁, _, _ = MacroModelling.calculate_first_order_solution(∇₁, model.constants, model.workspaces, model.caches; + opts = opts, initial_guess = model.caches.qme_solution, + parameter_values = params, caching = false) + ∇₂ = MacroModelling.calculate_hessian(params, SS_and_pars, model.caches, model.functions.hessian, model.workspaces, caching = false) + 𝐒₂, _ = MacroModelling.calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, model.constants, model.workspaces, model.caches; + initial_guess = model.caches.second_order_solution, + opts = opts, parameter_values = params, caching = false) + ∇₃ = MacroModelling.calculate_third_order_derivatives(params, SS_and_pars, model.caches, model.functions.third_order_derivatives, model.workspaces, caching = false) + MacroModelling.calculate_third_order_solution(∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, model.constants, model.workspaces, model.caches; + initial_guess = model.caches.third_order_solution, + opts = opts, parameter_values = params, caching = false) + + # Benchmark third-order derivatives + median_d3 = median_elapsed() do + MacroModelling.calculate_third_order_derivatives(params, SS_and_pars, model.caches, model.functions.third_order_derivatives, model.workspaces, caching = false) + end + bench["benchmark_third_order_derivatives"] = median_d3 + + # Benchmark third-order solve + median_to = median_elapsed() do + MacroModelling.calculate_third_order_solution(∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, model.constants, model.workspaces, model.caches; + initial_guess = model.caches.third_order_solution, + opts = opts, parameter_values = params, caching = false) + end + bench["benchmark_third_order_solve"] = median_to + + @info "Benchmark $(model.model_name) [third order]:" + @info " 3rd order derivs: $(round(median_d3*1e6, digits=1)) μs" + @info " 3rd order solve: $(round(median_to*1e6, digits=1)) μs" +end + +# ───────────────────────────────────────────── +# Higher-order solution matrix extraction +# ───────────────────────────────────────────── + +function orig_var_indices(model, orig) + all_vars = model.constants.post_model_macro.var + [findfirst(==(v), all_vars) for v in orig] +end + +function export_second_order_matrices(model, julia_dir, orig) + nPast = model.constants.post_model_macro.nPast_not_future_and_mixed + nExo = model.constants.post_model_macro.nExo + nVars = model.constants.post_model_macro.nVars + n_aug = nPast + 1 + nExo # [states, σ, shocks] + + # Expand compressed solution to full tensor (nVars × n_aug²) + sol2_full = model.caches.second_order_solution * model.constants.second_order.𝐔₂ + sol2_raw = reshape(Matrix(sol2_full), nVars, n_aug, n_aug) + + oi = orig_var_indices(model, orig) + sol2 = sol2_raw[oi, :, :] + nOrig = length(orig) + state_range = 1:nPast + σ_idx = nPast + 1 + exo_range = (nPast + 2):(nPast + 1 + nExo) + + # ghxx: nOrig × nState² — symmetric, column-major reshape matches Dynare kron(x,x) + ghxx = reshape(sol2[:, state_range, state_range], nOrig, nPast * nPast) + writedlm(joinpath(julia_dir, "ghxx.csv"), ghxx, ',') + + # ghxu: nOrig × (nState × nExo) — permutedims to match Dynare kron(x,u) convention + ghxu_block = sol2[:, state_range, exo_range] # (nOrig, nPast, nExo) + ghxu = reshape(permutedims(ghxu_block, (1, 3, 2)), nOrig, nPast * nExo) + writedlm(joinpath(julia_dir, "ghxu.csv"), ghxu, ',') + + # ghuu: nOrig × nExo² — symmetric + ghuu = reshape(sol2[:, exo_range, exo_range], nOrig, nExo * nExo) + writedlm(joinpath(julia_dir, "ghuu.csv"), ghuu, ',') + + # ghs2: nOrig × 1 — volatility correction + ghs2 = sol2[:, σ_idx, σ_idx] + writedlm(joinpath(julia_dir, "ghs2.csv"), ghs2, ',') +end + +function export_third_order_matrices(model, julia_dir, orig) + nPast = model.constants.post_model_macro.nPast_not_future_and_mixed + nExo = model.constants.post_model_macro.nExo + nVars = model.constants.post_model_macro.nVars + n_aug = nPast + 1 + nExo + + sol3_full = model.caches.third_order_solution * model.constants.third_order.𝐔₃ + sol3_raw = reshape(Matrix(sol3_full), nVars, n_aug, n_aug, n_aug) + + oi = orig_var_indices(model, orig) + sol3 = sol3_raw[oi, :, :, :] + nOrig = length(orig) + sr = 1:nPast + σ = nPast + 1 + er = (nPast + 2):(nPast + 1 + nExo) + + # ghxxx: nOrig × nState³ — symmetric, direct reshape + ghxxx = reshape(sol3[:, sr, sr, sr], nOrig, nPast^3) + writedlm(joinpath(julia_dir, "ghxxx.csv"), ghxxx, ',') + + # ghxxu: nOrig × (nState² × nExo) — permutedims [1,4,3,2] to match kron(x,kron(x,u)) + ghxxu = reshape(permutedims(sol3[:, sr, sr, er], (1, 4, 3, 2)), nOrig, nPast^2 * nExo) + writedlm(joinpath(julia_dir, "ghxxu.csv"), ghxxu, ',') + + # ghxuu: nOrig × (nState × nExo²) — same permutation + ghxuu = reshape(permutedims(sol3[:, sr, er, er], (1, 4, 3, 2)), nOrig, nPast * nExo^2) + writedlm(joinpath(julia_dir, "ghxuu.csv"), ghxuu, ',') + + # ghuuu: nOrig × nExo³ — symmetric, direct reshape + ghuuu = reshape(sol3[:, er, er, er], nOrig, nExo^3) + writedlm(joinpath(julia_dir, "ghuuu.csv"), ghuuu, ',') + + # ghxss: nOrig × nState — coefficient for x_i * σ² + ghxss = sol3[:, sr, σ, σ] + writedlm(joinpath(julia_dir, "ghxss.csv"), ghxss, ',') + + # ghuss: nOrig × nExo — coefficient for u_j * σ² + ghuss = sol3[:, er, σ, σ] + writedlm(joinpath(julia_dir, "ghuss.csv"), ghuss, ',') +end + +# ───────────────────────────────────────────── +# Export one model's first-order results +# ───────────────────────────────────────────── +function export_model(model, outdir; include_moments = true, benchmark_only = false) + julia_dir = joinpath(outdir, "julia") + mkpath(julia_dir) + + if benchmark_only + # ── Export .mod file only (needed for the Dynare phase) ── + cd(outdir) do + write_mod_file(model) + end + + # ── Benchmarks ── + bench = Dict{String, Float64}() + bench_start = time() + benchmark_first_order(model, bench) + bench_elapsed = time() - bench_start + write_benchmarks(julia_dir, bench) + + @info "Exported Julia benchmark-only results for $(model.model_name) → $outdir" + return bench_elapsed + end + + orig = original_vars(model) + state_vars = model.constants.post_model_macro.past_not_future_and_mixed + exo_vars = model.constants.post_model_macro.exo + + var_names_ascii, exo_names_ascii = export_names_and_steady_state(model, julia_dir, orig, state_vars, exo_vars) + export_first_order_matrices(model, julia_dir, orig, state_vars, exo_vars) + export_irfs(model, julia_dir, orig, exo_vars, algorithm = :first_order) + + if include_moments + export_moments(model, julia_dir, orig, exo_vars; + algorithm = :first_order, + var_names_ascii = var_names_ascii, + exo_names_ascii = exo_names_ascii) + end + + # ── Export .mod file ── + cd(outdir) do + write_mod_file(model) + end + + # ── Benchmarks ── + bench = Dict{String, Float64}() + bench_start = time() + benchmark_first_order(model, bench) + bench_elapsed = time() - bench_start + write_benchmarks(julia_dir, bench) + + @info "Exported Julia results for $(model.model_name) → $outdir" + return bench_elapsed +end + +# ───────────────────────────────────────────── +# Export one model's higher-order results +# ───────────────────────────────────────────── +function export_higher_order_model(model, outdir, dir_name, order) + julia_dir = joinpath(outdir, "julia") + mkpath(julia_dir) + + algorithm = order == 2 ? :pruned_second_order : :pruned_third_order + + # Trigger solve at the requested order (populates all caches up to that order) + get_solution(model, algorithm = algorithm) + + orig = original_vars(model) + state_vars = model.constants.post_model_macro.past_not_future_and_mixed + exo_vars = model.constants.post_model_macro.exo + + export_names_and_steady_state(model, julia_dir, orig, state_vars, exo_vars) + # Export first-order comparable objects even for higher-order model directories. + export_first_order_matrices(model, julia_dir, orig, state_vars, exo_vars) + export_irfs(model, julia_dir, orig, exo_vars, algorithm = :first_order) + + # Higher-order-specific outputs are moments-only (no higher-order solution matrices). + export_moments(model, julia_dir, orig, exo_vars; + algorithm = algorithm) + + # ── Export .mod file with correct order and pruning, renamed to match directory ── + cd(outdir) do + write_mod_file(model, order = order, pruning = true) + mv("$(model.model_name).mod", "$(dir_name).mod", force = true) + end + + # ── Benchmarks ── + bench = Dict{String, Float64}() + bench_start = time() + benchmark_first_order(model, bench) + benchmark_second_order(model, bench) + if order >= 3 + benchmark_third_order(model, bench) + end + bench_elapsed = time() - bench_start + write_benchmarks(julia_dir, bench) + + @info "Exported Julia higher-order results (order=$order) for $(model.model_name) → $outdir" + return bench_elapsed +end + +# ───────────────────────────────────────────── +# Main +# ───────────────────────────────────────────── +function main(args = ARGS) + parsed = parse_args(args) + parsed === nothing && return + output_root, only_models = parsed + + julia_threads, blas_threads = configure_julia_threads!() + + if isdir(output_root) + rm(output_root, recursive = true) + end + mkpath(output_root) + write_thread_configuration(output_root, julia_threads, blas_threads) + write_julia_environment_metadata(output_root, julia_threads, blas_threads) + + only_set = Set(only_models) + keep(name) = isempty(only_set) || (name in only_set) + + if !isempty(only_set) + @info "Restricting Phase 1 to selected models" only_models + end + + # Runtime tracking (wall-clock time post-startup for each model) + model_runtimes = Vector{Pair{String, Float64}}() + total_start = time() + + # Phase 1a: First-order exports for all models + for mname in MODEL_FILES + keep(mname) || continue + @info "Processing model (first order): $mname" + model_start = time() + include(joinpath(MODELS_DIR, "$mname.jl")) + model = Base.invokelatest(getfield, Main, Symbol(mname)) + outdir = joinpath(output_root, mname) + mkpath(outdir) + bench_elapsed = Base.invokelatest(export_model, model, outdir; + include_moments = !(mname in SKIP_MOMENTS_MODELS), + benchmark_only = mname in BENCHMARK_ONLY_MODELS) + elapsed = time() - model_start - bench_elapsed + push!(model_runtimes, mname => elapsed) + @info " Wall-clock time for $mname: $(round(elapsed, digits=3)) s (excl. $(round(bench_elapsed, digits=3)) s benchmarks)" + end + + # Phase 1b: Second-order exports for selected models + for mname in SECOND_ORDER_MODELS + keep(mname) || continue + dir_name = "$(mname)_pruned_2nd" + @info "Processing model (pruned order 2): $mname → $dir_name" + + model_start = time() + include(joinpath(MODELS_DIR, "$mname.jl")) + model = Base.invokelatest(getfield, Main, Symbol(mname)) + outdir = joinpath(output_root, dir_name) + mkpath(outdir) + bench_elapsed = Base.invokelatest(export_higher_order_model, model, outdir, dir_name, 2) + elapsed = time() - model_start - bench_elapsed + push!(model_runtimes, dir_name => elapsed) + @info " Wall-clock time for $dir_name: $(round(elapsed, digits=3)) s (excl. $(round(bench_elapsed, digits=3)) s benchmarks)" + end + + # Phase 1c: Third-order exports for selected models + for mname in THIRD_ORDER_MODELS + keep(mname) || continue + dir_name = "$(mname)_pruned_3rd" + @info "Processing model (pruned order 3): $mname → $dir_name" + + model_start = time() + include(joinpath(MODELS_DIR, "$mname.jl")) + model = Base.invokelatest(getfield, Main, Symbol(mname)) + outdir = joinpath(output_root, dir_name) + mkpath(outdir) + bench_elapsed = Base.invokelatest(export_higher_order_model, model, outdir, dir_name, 3) + elapsed = time() - model_start - bench_elapsed + push!(model_runtimes, dir_name => elapsed) + @info " Wall-clock time for $dir_name: $(round(elapsed, digits=3)) s (excl. $(round(bench_elapsed, digits=3)) s benchmarks)" + end + + # Write runtime summary (warmup-only times, excluding benchmark loops) + warmup_total = sum(t for (_, t) in model_runtimes) + open(joinpath(output_root, "runtime_julia.csv"), "w") do io + println(io, "model,elapsed_seconds") + for (name, t) in model_runtimes + println(io, "$name,$t") + end + println(io, "TOTAL,$warmup_total") + end + @info "Total wall-clock time (warmup only, excl. benchmarks): $(round(warmup_total, digits=3)) s" + @info "Phase 1 complete. Results in $output_root" +end + +if abspath(PROGRAM_FILE) == @__FILE__ + main() +end diff --git a/test/dynare_comparison/run_all_dynare.sh b/test/dynare_comparison/run_all_dynare.sh new file mode 100644 index 000000000..49cb30613 --- /dev/null +++ b/test/dynare_comparison/run_all_dynare.sh @@ -0,0 +1,162 @@ +#!/bin/bash +# run_all_dynare.sh — Phase 2 of Dynare comparison +# +# Iterates over model directories in /work/output/, runs Dynare on each +# .mod file, and extracts results to CSV via extract_dynare_results.m. +# +# Expected volume mount: -v host_output_dir:/work/output + +set -euo pipefail + +OUTPUT_DIR="/work/output" +EXTRACT_SCRIPT="/work/extract_dynare_results.m" +THREAD_COUNT="${THREAD_COUNT:-1}" + +# Enforce single-thread execution in CI unless explicitly overridden. +export OMP_NUM_THREADS="$THREAD_COUNT" +export OMP_THREAD_LIMIT="$THREAD_COUNT" +export OMP_DYNAMIC="FALSE" +export MKL_NUM_THREADS="$THREAD_COUNT" +export MKL_DYNAMIC="FALSE" +export OPENBLAS_NUM_THREADS="$THREAD_COUNT" +export BLIS_NUM_THREADS="$THREAD_COUNT" +export VECLIB_MAXIMUM_THREADS="$THREAD_COUNT" +export TBB_NUM_THREADS="$THREAD_COUNT" + +# Detect Dynare's Octave path +DYNARE_MATLAB="" +for p in /usr/lib/dynare/matlab /usr/share/dynare/matlab /usr/local/lib/dynare/matlab; do + if [ -d "$p" ]; then + DYNARE_MATLAB="$p" + break + fi +done + +if [ -z "$DYNARE_MATLAB" ]; then + echo "ERROR: Could not find Dynare matlab directory" + exit 1 +fi + +echo "Using Dynare at: $DYNARE_MATLAB" +OCTAVE_VERSION="$(octave --version | head -1)" +DYNARE_VERSION_OUTPUT="$(octave --no-gui --eval "addpath('$DYNARE_MATLAB'); dynare_version();" 2>&1 | tr '\r\n' ' ' | sed 's/[[:space:]]\+/ /g; s/^ //; s/ $//')" +OCTAVE_BLAS="$(octave --no-gui --quiet --eval "try; disp(__octave_config_info__('BLAS_LIBS')); catch; disp('unknown'); end;" 2>&1 | tr '\r\n' ' ' | sed 's/[[:space:]]\+/ /g; s/^ //; s/ $//')" +OCTAVE_LAPACK="$(octave --no-gui --quiet --eval "try; disp(__octave_config_info__('LAPACK_LIBS')); catch; disp('unknown'); end;" 2>&1 | tr '\r\n' ' ' | sed 's/[[:space:]]\+/ /g; s/^ //; s/ $//')" +OCTAVE_LINKED_BLAS="$(ldd /usr/bin/octave-cli 2>/dev/null | grep -Ei 'libblas|openblas|mkl' | tr '\r\n' '|' | sed 's/|$//; s/^ *//')" +OCTAVE_LINKED_LAPACK="$(ldd /usr/bin/octave-cli 2>/dev/null | grep -Ei 'liblapack|mkl' | tr '\r\n' '|' | sed 's/|$//; s/^ *//')" +METADATA_FILE="$OUTPUT_DIR/comparison_environment_dynare.txt" +{ + echo "dynare_driver=Octave" + echo "dynare_matlab_path=$DYNARE_MATLAB" + echo "dynare_version=$DYNARE_VERSION_OUTPUT" + echo "octave_version=$OCTAVE_VERSION" + echo "blas=$OCTAVE_BLAS" + echo "lapack=$OCTAVE_LAPACK" + echo "linked_blas=$OCTAVE_LINKED_BLAS" + echo "linked_lapack=$OCTAVE_LINKED_LAPACK" + echo "hostname=$(hostname 2>/dev/null || echo unknown)" + echo "kernel=$(uname -srmo 2>/dev/null || uname -a)" + echo "arch=$(uname -m 2>/dev/null || echo unknown)" + echo "cpu_threads=$(getconf _NPROCESSORS_ONLN 2>/dev/null || nproc 2>/dev/null || echo unknown)" + echo "thread_count_requested=$THREAD_COUNT" + echo "env_OMP_NUM_THREADS=$OMP_NUM_THREADS" + echo "env_OMP_THREAD_LIMIT=$OMP_THREAD_LIMIT" + echo "env_OMP_DYNAMIC=$OMP_DYNAMIC" + echo "env_MKL_NUM_THREADS=$MKL_NUM_THREADS" + echo "env_MKL_DYNAMIC=$MKL_DYNAMIC" + echo "env_OPENBLAS_NUM_THREADS=$OPENBLAS_NUM_THREADS" + echo "env_BLIS_NUM_THREADS=$BLIS_NUM_THREADS" + echo "env_VECLIB_MAXIMUM_THREADS=$VECLIB_MAXIMUM_THREADS" + echo "env_TBB_NUM_THREADS=$TBB_NUM_THREADS" +} > "$METADATA_FILE" +echo "Octave version: $OCTAVE_VERSION" +echo "Requested thread count: $THREAD_COUNT" +echo "Thread env: OMP_NUM_THREADS=$OMP_NUM_THREADS OPENBLAS_NUM_THREADS=$OPENBLAS_NUM_THREADS MKL_NUM_THREADS=$MKL_NUM_THREADS" +echo "Linked BLAS: ${OCTAVE_LINKED_BLAS:-unknown}" +echo "Linked LAPACK: ${OCTAVE_LINKED_LAPACK:-unknown}" +echo "Dynare version: $DYNARE_VERSION_OUTPUT" + +# Runtime CSV header +RUNTIME_CSV="$OUTPUT_DIR/runtime_dynare.csv" +echo "model,elapsed_seconds" > "$RUNTIME_CSV" +TOTAL_START=$(date +%s%N) + +# Process each model +for model_dir in "$OUTPUT_DIR"/*/; do + model_name=$(basename "$model_dir") + mod_file="$model_dir/${model_name}.mod" + dynare_stub="m" + + if [ ! -f "$mod_file" ]; then + echo "SKIP: No .mod file found for $model_name" + continue + fi + + echo "========================================" + echo "Running Dynare on: $model_name" + echo "========================================" + + dynare_out_dir="$model_dir/dynare" + mkdir -p "$dynare_out_dir" + + # Work in a temporary directory to avoid Dynare file pollution + workdir=$(mktemp -d) + cp "$mod_file" "$workdir/${dynare_stub}.mod" + cp "$EXTRACT_SCRIPT" "$workdir/" + + # Add nograph to stoch_simul to avoid graphics toolkit errors in headless mode + sed -i 's/stoch_simul(/stoch_simul(nograph, /' "$workdir/${dynare_stub}.mod" + + MODEL_START=$(date +%s%N) + + ( + cd "$workdir" + octave --no-gui --eval " + addpath('$DYNARE_MATLAB'); + model_name = '$model_name'; + output_dir = 'dynare_output'; + total_tic = tic; + dynare $dynare_stub noclearall; + extract_dynare_results; + elapsed = toc(total_tic); + fid = fopen(fullfile('dynare_output', 'runtime_seconds.csv'), 'w'); + fprintf(fid, '%.6f\n', elapsed); + fclose(fid); + fprintf('Wall-clock time for %s (post-startup): %.3f s\n', model_name, elapsed); + " + + # Copy results to the mounted output directory + if [ -d "dynare_output" ]; then + cp dynare_output/* "$dynare_out_dir/" + echo "Results copied to $dynare_out_dir" + else + echo "ERROR: No output produced for $model_name" + exit 1 + fi + ) + + MODEL_END=$(date +%s%N) + MODEL_ELAPSED=$(awk "BEGIN {printf \"%.6f\", ($MODEL_END - $MODEL_START) / 1000000000}") + + # Subtract benchmark loop time to get warmup-only elapsed + BENCH_ELAPSED_FILE="$dynare_out_dir/bench_elapsed_seconds.csv" + if [ -f "$BENCH_ELAPSED_FILE" ]; then + BENCH_ELAPSED=$(cat "$BENCH_ELAPSED_FILE") + WARMUP_ELAPSED=$(awk "BEGIN {printf \"%.6f\", $MODEL_ELAPSED - $BENCH_ELAPSED}") + else + WARMUP_ELAPSED="$MODEL_ELAPSED" + BENCH_ELAPSED="0" + fi + echo "$model_name,$WARMUP_ELAPSED" >> "$RUNTIME_CSV" + echo "Done: $model_name (warmup: ${WARMUP_ELAPSED} s, bench: ${BENCH_ELAPSED} s, total: ${MODEL_ELAPSED} s)" + + rm -rf "$workdir" +done + +TOTAL_END=$(date +%s%N) +TOTAL_ELAPSED=$(awk "BEGIN {printf \"%.6f\", ($TOTAL_END - $TOTAL_START) / 1000000000}") +# Sum warmup times from runtime CSV for accurate total (excludes benchmarks) +WARMUP_TOTAL=$(awk -F',' 'NR>1 && $1!="TOTAL" {sum+=$2} END {printf "%.6f", sum}' "$RUNTIME_CSV") +echo "TOTAL,$WARMUP_TOTAL" >> "$RUNTIME_CSV" +echo "Total wall-clock time (warmup only): ${WARMUP_TOTAL} s" +echo "Phase 2 complete." diff --git a/test/dynare_comparison/run_all_dynare_windows.ps1 b/test/dynare_comparison/run_all_dynare_windows.ps1 new file mode 100644 index 000000000..519c65b0d --- /dev/null +++ b/test/dynare_comparison/run_all_dynare_windows.ps1 @@ -0,0 +1,703 @@ +[CmdletBinding()] +param( + [string]$OutputDir, + [string]$ExtractScript, + [string]$DynareMatlabPath, + [string]$MatlabExe, + [ValidateRange(1, 512)] + [int]$ThreadCount = 1, + [ValidateRange(0, 10)] + [int]$MaxLicenseRetries = 1, + [ValidateRange(0, 600)] + [int]$LicenseRetryDelaySeconds = 10, + [string[]]$SkipModels = @(), + [string[]]$OnlyModels = @(), + [string[]]$BenchmarkOnlyModels = @('FRBUS'), + [switch]$ValidateOnly +) + +Set-StrictMode -Version Latest +$ErrorActionPreference = 'Stop' +$scriptRoot = Split-Path -Parent $PSCommandPath + +if (-not $OutputDir) { + $OutputDir = Join-Path $scriptRoot 'output' +} + +if (-not $ExtractScript) { + $ExtractScript = Join-Path $scriptRoot 'extract_dynare_results.m' +} + +function Resolve-ExistingPath { + param( + [string[]]$Candidates, + [string]$Description + ) + + foreach ($candidate in $Candidates) { + if ([string]::IsNullOrWhiteSpace($candidate)) { + continue + } + + if (Test-Path -LiteralPath $candidate) { + return (Resolve-Path -LiteralPath $candidate).Path + } + } + + throw "Could not find $Description. Checked: $($Candidates -join ', ')" +} + +function Invoke-RemoveItemRetry { + # Robust replacement for Remove-Item. Handles transient file locks + # (antivirus, lingering MATLAB/Dynare handles, OneDrive sync) by retrying + # with backoff. Always returns; never throws on missing paths and only + # throws after all attempts fail. + param( + [Parameter(Mandatory)][string]$Path, + [switch]$Recurse, + [int]$MaxAttempts = 6, + [int]$InitialDelayMs = 200 + ) + + if (-not (Test-Path -LiteralPath $Path)) { return } + + $delay = $InitialDelayMs + for ($attempt = 1; $attempt -le $MaxAttempts; $attempt++) { + try { + if ($Recurse) { + Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop + } else { + Remove-Item -LiteralPath $Path -Force -ErrorAction Stop + } + if (-not (Test-Path -LiteralPath $Path)) { return } + } catch { + if ($attempt -eq $MaxAttempts) { + Write-Warning ("Remove-Item failed for {0} after {1} attempts: {2}" -f $Path, $MaxAttempts, $_) + return + } + Start-Sleep -Milliseconds $delay + $delay = $delay * 2 + if ($delay -gt 5000) { + $delay = 5000 + } + } + } +} + +function Get-MatlabExecutable { + param([string]$PreferredPath) + + $candidates = @() + + if ($PreferredPath) { + $candidates += $PreferredPath + } + if ($env:MATLAB_EXE) { + $candidates += $env:MATLAB_EXE + } + + $matlabCommand = Get-Command matlab.exe -ErrorAction SilentlyContinue + if ($matlabCommand) { + $candidates += $matlabCommand.Source + } + + $matlabRoot = 'C:\Program Files\MATLAB' + if (Test-Path -LiteralPath $matlabRoot) { + Get-ChildItem -LiteralPath $matlabRoot -Directory | + Sort-Object Name -Descending | + ForEach-Object { + $candidates += (Join-Path $_.FullName 'bin\matlab.exe') + } + } + + Resolve-ExistingPath -Candidates $candidates -Description 'MATLAB executable' +} + +function Get-DynareMatlabPath { + param([string]$PreferredPath) + + $candidates = @() + + if ($PreferredPath) { + $candidates += $PreferredPath + } + if ($env:DYNARE_MATLAB) { + $candidates += $env:DYNARE_MATLAB + } + if ($env:DYNARE_HOME) { + $candidates += (Join-Path $env:DYNARE_HOME 'matlab') + } + + $candidates += 'D:\CustomTools\dynare-7.0-win\matlab' + $candidates += 'D:\CustomTools\Dynare\7\matlab' + $candidates += 'D:\CustomTools\dynare\7\matlab' + + $resolvedPath = Resolve-ExistingPath -Candidates $candidates -Description 'Dynare matlab directory' + $dynareEntryPoint = Join-Path $resolvedPath 'dynare.m' + if (-not (Test-Path -LiteralPath $dynareEntryPoint)) { + throw "Dynare matlab directory does not contain dynare.m: $resolvedPath" + } + + $resolvedPath +} + +function ConvertTo-MatlabString { + param([string]$Value) + + $Value.Replace('\', '/').Replace("'", "''") +} + +function Set-ThreadEnvironment { + param([int]$RequestedThreadCount) + + $threadValue = [string]$RequestedThreadCount + $threadEnvironment = [ordered]@{ + 'OMP_NUM_THREADS' = $threadValue + 'OMP_THREAD_LIMIT' = $threadValue + 'OMP_DYNAMIC' = 'FALSE' + 'MKL_NUM_THREADS' = $threadValue + 'MKL_DOMAIN_NUM_THREADS' = ('MKL_ALL={0}' -f $threadValue) + 'MKL_DYNAMIC' = 'FALSE' + 'OPENBLAS_NUM_THREADS' = $threadValue + 'BLIS_NUM_THREADS' = $threadValue + 'VECLIB_MAXIMUM_THREADS' = $threadValue + 'TBB_NUM_THREADS' = $threadValue + } + + foreach ($name in $threadEnvironment.Keys) { + Set-Item -Path ("Env:{0}" -f $name) -Value $threadEnvironment[$name] + } + + $threadEnvironment +} + +function Invoke-MatlabBatch { + param( + [string]$Executable, + [string]$WorkingDirectory, + [string]$BatchCommand, + [int]$RequestedThreadCount + ) + + $logPath = Join-Path $WorkingDirectory 'matlab_console.log' + Invoke-RemoveItemRetry -Path $logPath + + $matlabArgs = @() + if ($RequestedThreadCount -eq 1) { + $matlabArgs += '-singleCompThread' + } + $matlabArgs += '-logfile' + $matlabArgs += $logPath + $matlabArgs += '-batch' + $matlabArgs += $BatchCommand + + Write-Host ("Launching MATLAB: {0} {1}" -f $Executable, ($matlabArgs -join ' ')) + Write-Host ("Streaming MATLAB log: {0}" -f $logPath) + + $proc = Start-Process -FilePath $Executable -ArgumentList $matlabArgs -WorkingDirectory $WorkingDirectory -PassThru -NoNewWindow + + $doneFlagPath = Join-Path $WorkingDirectory 'batch_done.flag' + Invoke-RemoveItemRetry -Path $doneFlagPath + + $procId = $proc.Id + $linesPrinted = 0 + # Primary done-signal: the MATLAB driver writes batch_done.flag at the very + # end of run_all_dynare. Poll for that file and (as a backup) check whether + # the MATLAB process is still alive via Get-Process. Avoid method calls on + # the process object so this works in PowerShell Constrained Language Mode. + while ($true) { + Start-Sleep -Milliseconds 1000 + if (Test-Path -LiteralPath $logPath) { + $allLines = @(Get-Content -LiteralPath $logPath -ErrorAction SilentlyContinue) + if ($allLines.Count -gt $linesPrinted) { + for ($idx = $linesPrinted; $idx -lt $allLines.Count; $idx++) { + Write-Host ("[matlab] {0}" -f $allLines[$idx]) + } + $linesPrinted = $allLines.Count + } + } + if (Test-Path -LiteralPath $doneFlagPath) { break } + $alive = Get-Process -Id $procId -ErrorAction SilentlyContinue + if (-not $alive) { break } + } + + # Drain any final log lines once MATLAB is done. + Start-Sleep -Milliseconds 500 + + if (Test-Path -LiteralPath $logPath) { + $allLines = @(Get-Content -LiteralPath $logPath -ErrorAction SilentlyContinue) + if ($allLines.Count -gt $linesPrinted) { + for ($idx = $linesPrinted; $idx -lt $allLines.Count; $idx++) { + Write-Host ("[matlab] {0}" -f $allLines[$idx]) + } + } + } + + if ($null -eq $proc.ExitCode) { + return 0 + } + + return $proc.ExitCode +} + +function Test-MatlabLicenseCheckoutFailure { + param([string]$LogPath) + + if (-not (Test-Path -LiteralPath $LogPath)) { + return $false + } + + $logText = Get-Content -LiteralPath $LogPath -Raw -ErrorAction SilentlyContinue + if (-not $logText) { + return $false + } + + if ($logText -match 'License checkout failed') { + return $true + } + if ($logText -match 'License Manager Error\s*-97') { + return $true + } + + return $false +} + +function Update-StochSimulDirective { + param([string]$ModFilePath) + + $content = Get-Content -LiteralPath $ModFilePath -Raw + $updatedContent = $content -replace 'stoch_simul\s*\((?!\s*nograph\b)', 'stoch_simul(nograph, ' + $updatedContent = $updatedContent -replace 'stoch_simul\s*;', 'stoch_simul(nograph);' + + if ($updatedContent -ne $content) { + Set-Content -LiteralPath $ModFilePath -Value $updatedContent -Encoding ascii -NoNewline + } +} + +function New-WorkDirectory { + param( + [string]$WorkRoot, + [string]$ModelName + ) + + $suffix = '{0}_{1}' -f (Get-Date -Format 'yyyyMMddHHmmssfff'), (Get-Random -Minimum 10000 -Maximum 99999) + $workDir = Join-Path $WorkRoot ("{0}_{1}" -f $ModelName, $suffix) + New-Item -ItemType Directory -Path $workDir | Out-Null + $workDir +} + +if (-not (Test-Path -LiteralPath $OutputDir)) { + throw "Output directory not found: $OutputDir. Run generate_julia_results.jl first or pass -OutputDir." +} + +if (-not (Test-Path -LiteralPath $ExtractScript)) { + throw "Extract script not found: $ExtractScript" +} + +$resolvedOutputDir = (Resolve-Path -LiteralPath $OutputDir).Path +$resolvedExtractScript = (Resolve-Path -LiteralPath $ExtractScript).Path +$resolvedMatlabExe = Get-MatlabExecutable -PreferredPath $MatlabExe +$resolvedDynareMatlabPath = Get-DynareMatlabPath -PreferredPath $DynareMatlabPath +$dynareMatlabLiteral = ConvertTo-MatlabString -Value $resolvedDynareMatlabPath +$threadEnvironment = Set-ThreadEnvironment -RequestedThreadCount $ThreadCount + +$modelDirectories = Get-ChildItem -LiteralPath $resolvedOutputDir -Directory | Sort-Object Name +if (-not $modelDirectories) { + throw "No model directories found under $resolvedOutputDir" +} + +Write-Host "Using MATLAB at: $resolvedMatlabExe" +Write-Host "Using Dynare at: $resolvedDynareMatlabPath" +Write-Host "Output root: $resolvedOutputDir" +Write-Host "Requested thread count: $ThreadCount" +$skipModelSet = @{} +foreach ($skipName in $SkipModels) { + if (-not [string]::IsNullOrWhiteSpace($skipName)) { + $skipModelSet[$skipName] = $true + } +} +if ($skipModelSet.Count -gt 0) { + Write-Host ("Skipping models: {0}" -f (($skipModelSet.Keys | Sort-Object) -join ', ')) +} +$onlyModelSet = @{} +foreach ($onlyName in $OnlyModels) { + if (-not [string]::IsNullOrWhiteSpace($onlyName)) { + $onlyModelSet[$onlyName] = $true + } +} +if ($onlyModelSet.Count -gt 0) { + Write-Host ("Restricting to models: {0}" -f (($onlyModelSet.Keys | Sort-Object) -join ', ')) +} +$benchmarkOnlySet = @{} +foreach ($benchmarkName in $BenchmarkOnlyModels) { + if (-not [string]::IsNullOrWhiteSpace($benchmarkName)) { + $benchmarkOnlySet[$benchmarkName] = $true + } +} +if ($benchmarkOnlySet.Count -gt 0) { + Write-Host ("Benchmark-only models: {0}" -f (($benchmarkOnlySet.Keys | Sort-Object) -join ', ')) +} +Write-Host 'Configured thread environment for MATLAB and MEX libraries:' +foreach ($name in $threadEnvironment.Keys) { + Write-Host (" {0}={1}" -f $name, $threadEnvironment[$name]) +} + +if ($ValidateOnly) { + Write-Host 'Validation only mode enabled.' + foreach ($modelDirectory in $modelDirectories) { + $modelName = $modelDirectory.Name + $modFile = Join-Path $modelDirectory.FullName "$modelName.mod" + if (Test-Path -LiteralPath $modFile) { + Write-Host "READY: $modelName" + } + else { + Write-Warning "SKIP: No .mod file found for $modelName" + } + } + return +} + +$workRoot = Join-Path $scriptRoot '_dynare_work' +New-Item -ItemType Directory -Path $workRoot -Force | Out-Null + +$batchRoot = Join-Path $workRoot ("batch_{0}_{1}" -f (Get-Date -Format 'yyyyMMddHHmmssfff'), (Get-Random -Minimum 10000 -Maximum 99999)) +New-Item -ItemType Directory -Path $batchRoot | Out-Null + +# Prepare per-model working directories under a single batch root so MATLAB can +# iterate through them in one session (avoids per-model license checkouts). +$modelEntries = @() +$failedModels = @() +$dynareStub = 'm' + +foreach ($modelDirectory in $modelDirectories) { + $modelName = $modelDirectory.Name + $modFile = Join-Path $modelDirectory.FullName "$modelName.mod" + + if ($skipModelSet.ContainsKey($modelName)) { + Write-Host "SKIP (configured): $modelName" + continue + } + + if ($onlyModelSet.Count -gt 0 -and -not $onlyModelSet.ContainsKey($modelName)) { + Write-Host "SKIP (not in OnlyModels): $modelName" + continue + } + + if (-not (Test-Path -LiteralPath $modFile)) { + Write-Warning "SKIP: No .mod file found for $modelName" + continue + } + + $dynareOutputDir = Join-Path $modelDirectory.FullName 'dynare' + if (Test-Path -LiteralPath $dynareOutputDir) { + Get-ChildItem -LiteralPath $dynareOutputDir -Force -ErrorAction SilentlyContinue | + ForEach-Object { Invoke-RemoveItemRetry -Path $_.FullName -Recurse } + } + else { + New-Item -ItemType Directory -Path $dynareOutputDir | Out-Null + } + + $modelWorkDir = Join-Path $batchRoot $modelName + New-Item -ItemType Directory -Path $modelWorkDir | Out-Null + $stubModFile = Join-Path $modelWorkDir "$dynareStub.mod" + Copy-Item -LiteralPath $modFile -Destination $stubModFile + Copy-Item -LiteralPath $resolvedExtractScript -Destination $modelWorkDir + Update-StochSimulDirective -ModFilePath $stubModFile + + $modelEntries += @{ + Name = $modelName + WorkDir = $modelWorkDir + DynareOutDir = $dynareOutputDir + BenchmarkOnly = [bool]$benchmarkOnlySet.ContainsKey($modelName) + } +} + +if (-not $modelEntries) { + Write-Warning 'No model entries to process. Phase 2 complete.' + return +} + +# Build the MATLAB driver that runs all models in one session. +$driverScriptPath = Join-Path $batchRoot 'run_all_dynare.m' +$workRootLiteral = ConvertTo-MatlabString -Value $batchRoot +$dynareEnvironmentPath = Join-Path $resolvedOutputDir 'comparison_environment_dynare.txt' +$dynareEnvironmentLiteral = ConvertTo-MatlabString -Value $dynareEnvironmentPath +$matlabExeLiteral = ConvertTo-MatlabString -Value $resolvedMatlabExe + +$modelEntryLines = @() +foreach ($entry in $modelEntries) { + $nameLiteral = ConvertTo-MatlabString -Value $entry.Name + $workLiteral = ConvertTo-MatlabString -Value $entry.WorkDir + $outLiteral = ConvertTo-MatlabString -Value $entry.DynareOutDir + $benchmarkLiteral = if ($entry.BenchmarkOnly) { 'true' } else { 'false' } + $modelEntryLines += "model_entries(end+1) = struct('name', '$nameLiteral', 'work_dir', '$workLiteral', 'output_dir', '$outLiteral', 'benchmark_only', $benchmarkLiteral);" +} +$modelEntriesBlock = ($modelEntryLines -join "`n ") + +$driverScript = @" +diary('matlab_batch.log'); +diary on; +addpath('$dynareMatlabLiteral'); + +requested_threads = $ThreadCount; +thread_env_names = {'OMP_NUM_THREADS', 'OMP_THREAD_LIMIT', 'OMP_DYNAMIC', 'MKL_NUM_THREADS', 'MKL_DOMAIN_NUM_THREADS', 'MKL_DYNAMIC', 'OPENBLAS_NUM_THREADS', 'BLIS_NUM_THREADS', 'VECLIB_MAXIMUM_THREADS', 'TBB_NUM_THREADS'}; +if exist('maxNumCompThreads', 'builtin') || exist('maxNumCompThreads', 'file') + previous_num_comp_threads = maxNumCompThreads(requested_threads); + active_num_comp_threads = maxNumCompThreads(); + fprintf('MATLAB thread configuration: requested=%d active=%d previous=%d\n', requested_threads, active_num_comp_threads, previous_num_comp_threads); +else + fprintf('MATLAB thread configuration: requested=%d active=maxNumCompThreads unavailable\n', requested_threads); +end +for thread_env_idx = 1:numel(thread_env_names) + thread_env_name = thread_env_names{thread_env_idx}; + thread_env_value = getenv(thread_env_name); + if isempty(thread_env_value) + thread_env_value = ''; + end + fprintf('MATLAB thread environment: %s=%s\n', thread_env_name, thread_env_value); +end + +dynare_environment_file = '$dynareEnvironmentLiteral'; +dynare_version_text = ''; +try + dynare_version_value = dynare_version(); + if isstring(dynare_version_value) + dynare_version_text = char(dynare_version_value); + elseif ischar(dynare_version_value) + dynare_version_text = dynare_version_value; + end +catch +end +if isempty(dynare_version_text) + dynare_version_text = strtrim(evalc('dynare_version;')); +end +dynare_version_text = strtrim(dynare_version_text); +dynare_version_text = regexprep(dynare_version_text, '[\r\n]+', ' | '); +if isempty(dynare_version_text) + dynare_version_text = 'unknown'; +end +host_name = getenv('COMPUTERNAME'); +if isempty(host_name) + host_name = getenv('HOSTNAME'); +end +if isempty(host_name) + host_name = 'unknown'; +end +os_name = 'unknown'; +if exist('system_dependent', 'builtin') || exist('system_dependent', 'file') + try + os_name = system_dependent('getos'); + catch + end +end +blas_name = 'unknown'; +lapack_name = 'unknown'; +try + blas_name = version('-blas'); +catch +end +try + lapack_name = version('-lapack'); +catch +end +metadata_fid = fopen(dynare_environment_file, 'w'); +fprintf(metadata_fid, 'dynare_driver=MATLAB\n'); +fprintf(metadata_fid, 'dynare_version=%s\n', dynare_version_text); +fprintf(metadata_fid, 'dynare_matlab_path=%s\n', '$dynareMatlabLiteral'); +fprintf(metadata_fid, 'matlab_executable=%s\n', '$matlabExeLiteral'); +fprintf(metadata_fid, 'matlab_version=%s\n', version); +try + fprintf(metadata_fid, 'matlab_release=%s\n', version('-release')); +catch +end +fprintf(metadata_fid, 'blas=%s\n', regexprep(blas_name, '[\r\n]+', ' | ')); +fprintf(metadata_fid, 'lapack=%s\n', regexprep(lapack_name, '[\r\n]+', ' | ')); +fprintf(metadata_fid, 'hostname=%s\n', host_name); +fprintf(metadata_fid, 'computer=%s\n', computer); +fprintf(metadata_fid, 'os=%s\n', regexprep(os_name, '[\r\n]+', ' | ')); +fprintf(metadata_fid, 'thread_count_requested=%d\n', requested_threads); +if exist('active_num_comp_threads', 'var') + fprintf(metadata_fid, 'max_num_comp_threads=%d\n', active_num_comp_threads); +else + fprintf(metadata_fid, 'max_num_comp_threads=%s\n', 'unknown'); +end +for thread_env_idx = 1:numel(thread_env_names) + thread_env_name = thread_env_names{thread_env_idx}; + thread_env_value = getenv(thread_env_name); + if isempty(thread_env_value) + thread_env_value = ''; + end + fprintf(metadata_fid, 'env_%s=%s\n', thread_env_name, thread_env_value); +end +fclose(metadata_fid); + +batch_root = '$workRootLiteral'; +status_file = fullfile(batch_root, 'model_status.csv'); +status_fid = fopen(status_file, 'w'); +fprintf(status_fid, 'model,status,message\n'); + +runtime_file = fullfile('$( ConvertTo-MatlabString -Value $resolvedOutputDir )', 'runtime_dynare.csv'); +runtime_fid = fopen(runtime_file, 'w'); +fprintf(runtime_fid, 'model,elapsed_seconds\n'); + +model_entries = struct('name', {}, 'work_dir', {}, 'output_dir', {}, 'benchmark_only', {}); + $modelEntriesBlock + +original_dir = pwd; +batch_start_tic = tic; +for entry_idx = 1:numel(model_entries) + entry = model_entries(entry_idx); + fprintf('========================================\n'); + fprintf('[%s] (%d/%d) Running Dynare on: %s\n', datestr(now, 'HH:MM:SS'), entry_idx, numel(model_entries), entry.name); + fprintf('========================================\n'); + cd(entry.work_dir); + model_tic = tic; + try + clearvars -except status_fid runtime_fid model_entries entry_idx entry batch_root original_dir requested_threads thread_env_names previous_num_comp_threads active_num_comp_threads batch_start_tic model_tic; + model_name = entry.name; + output_dir = entry.output_dir; + benchmark_only_mode = entry.benchmark_only; + dynare $dynareStub noclearall; + extract_dynare_results; + elapsed_model = toc(model_tic) - bench_elapsed_total; + fprintf('[%s] OK: %s in %.1f s (excl. %.1f s benchmarks)\n', datestr(now, 'HH:MM:SS'), entry.name, elapsed_model, bench_elapsed_total); + fprintf(status_fid, '%s,ok,\n', entry.name); + fprintf(runtime_fid, '%s,%.6f\n', entry.name, elapsed_model); + catch ME + elapsed_model = toc(model_tic); + report_text = getReport(ME, 'extended', 'hyperlinks', 'off'); + fid = fopen('matlab_error.log', 'w'); + fprintf(fid, '%s\n', report_text); + fclose(fid); + fprintf('[%s] ERROR: %s after %.1f s -- %s\n', datestr(now, 'HH:MM:SS'), entry.name, elapsed_model, ME.message); + disp(report_text); + message = strrep(ME.message, ',', ';'); + message = strrep(message, sprintf('\n'), ' '); + fprintf(status_fid, '%s,error,%s\n', entry.name, message); + fprintf(runtime_fid, '%s,%.6f\n', entry.name, elapsed_model); + end + cd(original_dir); +end +batch_elapsed = toc(batch_start_tic); +fprintf(runtime_fid, 'TOTAL,%.6f\n', batch_elapsed); +fclose(runtime_fid); +fprintf('[%s] Batch finished in %.1f s\n', datestr(now, 'HH:MM:SS'), batch_elapsed); + +fclose(status_fid); +done_fid = fopen('batch_done.flag', 'w'); +fprintf(done_fid, 'done\n'); +fclose(done_fid); +diary off; +exit(0); +"@ + +Set-Content -LiteralPath $driverScriptPath -Value $driverScript -Encoding ascii + +Write-Host '----------------------------------------' +Write-Host ("Launching single MATLAB session for {0} model(s) at thread count {1}..." -f $modelEntries.Count, $ThreadCount) +Write-Host '----------------------------------------' + +$matlabLogPath = Join-Path $batchRoot 'matlab_console.log' +$attempt = 0 +$matlabExitCode = 1 +while ($true) { + $attempt += 1 + if ($attempt -gt 1) { + Write-Warning ("Restarting MATLAB batch after license checkout error (attempt {0}/{1})." -f $attempt, ($MaxLicenseRetries + 1)) + } + + $matlabExitCode = Invoke-MatlabBatch -Executable $resolvedMatlabExe -WorkingDirectory $batchRoot -BatchCommand 'run_all_dynare' -RequestedThreadCount $ThreadCount + + if ($matlabExitCode -eq 0) { + break + } + + $isLicenseFailure = Test-MatlabLicenseCheckoutFailure -LogPath $matlabLogPath + $hasRetryBudget = $attempt -le $MaxLicenseRetries + if (-not $isLicenseFailure -or -not $hasRetryBudget) { + break + } + + if ($LicenseRetryDelaySeconds -gt 0) { + Write-Host ("Waiting {0} seconds before MATLAB restart..." -f $LicenseRetryDelaySeconds) + Start-Sleep -Seconds $LicenseRetryDelaySeconds + } +} + +$statusFile = Join-Path $batchRoot 'model_status.csv' +$statusByModel = @{} +if (Test-Path -LiteralPath $statusFile) { + $statusRows = Import-Csv -LiteralPath $statusFile + foreach ($row in $statusRows) { + $statusByModel[$row.model] = $row + } +} +else { + Write-Warning ("MATLAB status file not produced at {0}; treating all models as failed (MATLAB exit code {1})." -f $statusFile, $matlabExitCode) +} + +$keepBatch = $false +foreach ($entry in $modelEntries) { + $row = $null + if ($statusByModel.ContainsKey($entry.Name)) { + $row = $statusByModel[$entry.Name] + } + # MATLAB writes results directly into $entry.DynareOutDir (no per-file copy). + $matlabOutputDir = $entry.DynareOutDir + + if (-not $row) { + $failedModels += $entry.Name + Write-Warning ("No status recorded for {0}. Likely MATLAB aborted before processing it." -f $entry.Name) + $keepBatch = $true + continue + } + + if ($row.status -ne 'ok') { + $failedModels += $entry.Name + $errMessage = $row.message + if (-not $errMessage) { $errMessage = '' } + Write-Warning ("Dynare failed for {0}: {1}" -f $entry.Name, $errMessage) + $keepBatch = $true + continue + } + + if (-not (Test-Path -LiteralPath $matlabOutputDir)) { + $failedModels += $entry.Name + Write-Warning ("Status reported ok for {0} but no output directory found at {1}." -f $entry.Name, $matlabOutputDir) + $keepBatch = $true + continue + } + + $outputFiles = Get-ChildItem -LiteralPath $matlabOutputDir -File + if (-not $outputFiles) { + $failedModels += $entry.Name + Write-Warning ("Status ok for {0} but no CSV files produced." -f $entry.Name) + $keepBatch = $true + continue + } + + Write-Host ("Done: {0} (results in {1})" -f $entry.Name, $entry.DynareOutDir) +} + +if ($matlabExitCode -ne 0) { + Write-Warning ("MATLAB exited with code {0}. Successfully processed models were still copied." -f $matlabExitCode) + $keepBatch = $true +} + +if ($keepBatch) { + Write-Warning "Keeping batch work directory for inspection: $batchRoot" +} +else { + Invoke-RemoveItemRetry -Path $batchRoot -Recurse +} + +if ($failedModels.Count -gt 0) { + Write-Warning ("Phase 2 finished with failures in: {0}" -f ($failedModels -join ', ')) + throw ("Phase 2 failed for {0} model(s): {1}. Batch kept at {2}" -f $failedModels.Count, ($failedModels -join ', '), $batchRoot) +} + +if ($matlabExitCode -ne 0) { + throw ("Phase 2 failed: MATLAB exited with code {0}. Batch kept at {1}" -f $matlabExitCode, $batchRoot) +} + +Write-Host 'Phase 2 complete.' \ No newline at end of file diff --git a/test/dynare_comparison/run_thread_sweep_macos.sh b/test/dynare_comparison/run_thread_sweep_macos.sh new file mode 100755 index 000000000..94d3f8857 --- /dev/null +++ b/test/dynare_comparison/run_thread_sweep_macos.sh @@ -0,0 +1,391 @@ +#!/usr/bin/env bash +# Example invocation: +# ./test/dynare_comparison/run_thread_sweep_macos.sh \ +# --julia-exe "$HOME/.juliaup/bin/julia" \ +# --thread-counts 1,2,4,8 + +set -euo pipefail + +script_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +repo_root="$(cd -- "$script_root/../.." && pwd)" + +thread_counts_csv="1,2,4,8" +output_root="" +julia_exe="" +generate_julia_script="" +dynare_script="" +compare_script="" +sweep_compare_script="" +extract_script="" +dynare_docker_image_tag="${DYNARE_DOCKER_IMAGE_TAG:-macromodelling-dynare-testing}" +only_models_csv="" +validate_only=0 + +print_usage() { + cat <<'USAGE' +Usage: run_thread_sweep_macos.sh [options] + +Options: + --thread-counts Comma-separated thread counts (default: 1,2,4,8) + --output-root Final output root (default: test/dynare_comparison/output_thread_sweep) + --julia-exe Path to Julia executable + --generate-julia-script Phase-1 Julia script (default: generate_julia_results.jl) + --dynare-script Phase-2 Dynare shell script (default: run_all_dynare.sh) + --extract-script Dynare extract MATLAB script (default: extract_dynare_results.m) + --compare-script Phase-3 Julia script (default: compare_results.jl) + --sweep-compare-script Sweep summary Julia script (default: compare_thread_sweep_results.jl) + --dynare-docker-image-tag Docker image tag (default: macromodelling-dynare-testing) + --only-models Restrict models (forwarded to phase 1) + --validate-only Print resolved plan and exit + -h, --help Show this help +USAGE +} + +resolve_existing_path() { + local description="$1" + shift + + local candidate + for candidate in "$@"; do + if [[ -z "${candidate// }" ]]; then + continue + fi + if [[ -e "$candidate" ]]; then + cd -- "$(dirname -- "$candidate")" + local resolved + resolved="$(pwd)/$(basename -- "$candidate")" + cd -- "$repo_root" + printf '%s\n' "$resolved" + return 0 + fi + done + + printf 'Could not find %s. Checked: %s\n' "$description" "$*" >&2 + return 1 +} + +get_julia_executable() { + local candidates=() + if [[ -n "$julia_exe" ]]; then + candidates+=("$julia_exe") + fi + if [[ -n "${JULIA_EXE:-}" ]]; then + candidates+=("$JULIA_EXE") + fi + if [[ -x "$HOME/.juliaup/bin/julia" ]]; then + candidates+=("$HOME/.juliaup/bin/julia") + fi + if [[ -x "/Applications/Julia-1.12.app/Contents/Resources/julia/bin/julia" ]]; then + candidates+=("/Applications/Julia-1.12.app/Contents/Resources/julia/bin/julia") + fi + + local julia_on_path + julia_on_path="$(command -v julia || true)" + if [[ -n "$julia_on_path" ]]; then + candidates+=("$julia_on_path") + fi + + resolve_existing_path "Julia executable" "${candidates[@]}" +} + +make_staging_output_root() { + local resolved_output_parent="$1" + local output_root_leaf="$2" + local timestamp random_id + timestamp="$(date +%Y%m%d%H%M%S)" + random_id="$RANDOM" + printf '%s/%s.__staging_%s_%s\n' "$resolved_output_parent" "$output_root_leaf" "$timestamp" "$random_id" +} + +publish_staged_output_root() { + local stage_output_root="$1" + local final_output_root="$2" + local resolved_output_parent="$3" + local output_root_leaf="$4" + + if [[ ! -d "$stage_output_root" ]]; then + printf 'Staged sweep output not found: %s\n' "$stage_output_root" >&2 + return 1 + fi + + if [[ -e "$final_output_root" ]]; then + local previous_output_root + previous_output_root="${resolved_output_parent}/${output_root_leaf}.__previous_$(date +%Y%m%d%H%M%S)_$RANDOM" + printf 'Moving existing output root aside: %s -> %s\n' "$final_output_root" "$previous_output_root" + mv "$final_output_root" "$previous_output_root" + fi + + printf 'Publishing staged sweep output: %s -> %s\n' "$stage_output_root" "$final_output_root" + mv "$stage_output_root" "$final_output_root" +} + +invoke_julia_script() { + local executable="$1" + local project_root="$2" + local script_path="$3" + local output_argument="$4" + local description="$5" + local requested_thread_count="$6" + local use_thread_count="$7" + shift 7 + + local julia_args=("--project=${project_root}") + if [[ "$use_thread_count" == "1" ]]; then + julia_args+=("--threads=${requested_thread_count}") + fi + julia_args+=("$script_path") + + local extra + if [[ $# -gt 0 ]]; then + for extra in "$@"; do + if [[ -n "${extra// }" ]]; then + julia_args+=("$extra") + fi + done + fi + + julia_args+=("$output_argument") + + printf 'Running Julia step: %s\n' "$description" + "$executable" "${julia_args[@]}" +} + +invoke_dynare_phase() { + local thread_output_dir="$1" + local requested_thread_count="$2" + local resolved_dynare_script="$3" + local resolved_extract_script="$4" + + if ! docker image inspect "$dynare_docker_image_tag" >/dev/null 2>&1; then + printf 'Docker image %s not found. Build it first with: docker build -t %s test/dynare_comparison/\n' "$dynare_docker_image_tag" "$dynare_docker_image_tag" >&2 + return 1 + fi + + printf 'Running Dynare step for %s thread(s)\n' "$requested_thread_count" + docker run --rm \ + --user "$(id -u):$(id -g)" \ + -e HOME=/tmp \ + -e THREAD_COUNT="$requested_thread_count" \ + -v "$thread_output_dir:/work/output" \ + -v "$resolved_extract_script:/work/extract_dynare_results.m:ro" \ + -v "$resolved_dynare_script:/work/run_all_dynare.sh:ro" \ + --entrypoint /bin/bash \ + "$dynare_docker_image_tag" \ + /work/run_all_dynare.sh +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --thread-counts) + thread_counts_csv="$2" + shift 2 + ;; + --output-root) + output_root="$2" + shift 2 + ;; + --julia-exe) + julia_exe="$2" + shift 2 + ;; + --generate-julia-script) + generate_julia_script="$2" + shift 2 + ;; + --dynare-script) + dynare_script="$2" + shift 2 + ;; + --extract-script) + extract_script="$2" + shift 2 + ;; + --compare-script) + compare_script="$2" + shift 2 + ;; + --sweep-compare-script) + sweep_compare_script="$2" + shift 2 + ;; + --dynare-docker-image-tag) + dynare_docker_image_tag="$2" + shift 2 + ;; + --only-models) + only_models_csv="$2" + shift 2 + ;; + --validate-only) + validate_only=1 + shift + ;; + -h|--help) + print_usage + exit 0 + ;; + *) + printf 'Unknown argument: %s\n\n' "$1" >&2 + print_usage >&2 + exit 1 + ;; + esac +done + +if [[ -z "$output_root" ]]; then + output_root="$script_root/output_thread_sweep" +fi +if [[ -z "$generate_julia_script" ]]; then + generate_julia_script="$script_root/generate_julia_results.jl" +fi +if [[ -z "$dynare_script" ]]; then + dynare_script="$script_root/run_all_dynare.sh" +fi +if [[ -z "$extract_script" ]]; then + extract_script="$script_root/extract_dynare_results.m" +fi +if [[ -z "$compare_script" ]]; then + compare_script="$script_root/compare_results.jl" +fi +if [[ -z "$sweep_compare_script" ]]; then + sweep_compare_script="$script_root/compare_thread_sweep_results.jl" +fi + +requested_output_root="$output_root" +output_root_leaf="$(basename -- "$requested_output_root")" +output_root_parent="$(dirname -- "$requested_output_root")" + +mkdir -p "$output_root_parent" +resolved_output_parent="$(cd -- "$output_root_parent" && pwd)" +resolved_output_root="$resolved_output_parent/$output_root_leaf" +staging_output_root="$(make_staging_output_root "$resolved_output_parent" "$output_root_leaf")" + +cd -- "$repo_root" + +resolved_julia_exe="$(get_julia_executable)" +resolved_generate_julia_script="$(resolve_existing_path 'Julia phase-1 script' "$generate_julia_script")" +resolved_dynare_script="$(resolve_existing_path 'Dynare phase-2 script' "$dynare_script")" +resolved_extract_script="$(resolve_existing_path 'Dynare extract script' "$extract_script")" +resolved_compare_script="$(resolve_existing_path 'Julia phase-3 script' "$compare_script")" +resolved_sweep_compare_script="$(resolve_existing_path 'thread-sweep summary script' "$sweep_compare_script")" + +thread_counts_raw=() +while IFS= read -r token || [[ -n "$token" ]]; do + thread_counts_raw+=("$token") +done < <(printf '%s' "$thread_counts_csv" | tr ',' '\n' | sed 's/^ *//; s/ *$//') + +if [[ "${#thread_counts_raw[@]}" -eq 0 ]]; then + printf 'At least one thread count must be provided.\n' >&2 + exit 1 +fi + +thread_counts=() +for t in "${thread_counts_raw[@]}"; do + if [[ -z "$t" ]]; then + continue + fi + if [[ ! "$t" =~ ^[0-9]+$ ]] || [[ "$t" -lt 1 ]]; then + printf 'Invalid thread count: %s\n' "$t" >&2 + exit 1 + fi + already_seen=0 + for existing_thread in "${thread_counts[@]-}"; do + if [[ "$existing_thread" == "$t" ]]; then + already_seen=1 + break + fi + done + if [[ "$already_seen" -eq 0 ]]; then + thread_counts+=("$t") + fi +done + +if [[ "${#thread_counts[@]}" -eq 0 ]]; then + printf 'At least one valid thread count must be provided.\n' >&2 + exit 1 +fi + +IFS=$'\n' thread_counts=( $(printf '%s\n' "${thread_counts[@]}" | sort -n) ) +unset IFS + +phase1_extra_args=() +if [[ -n "${only_models_csv// }" ]]; then + phase1_extra_args+=("--only-models=${only_models_csv}") + printf 'Restricting sweep to models: %s\n' "$only_models_csv" +fi + +printf 'Repository root: %s\n' "$repo_root" +printf 'Julia executable: %s\n' "$resolved_julia_exe" +printf 'Dynare Docker image tag: %s\n' "$dynare_docker_image_tag" +printf 'Final sweep output root: %s\n' "$resolved_output_root" +printf 'Sweep staging root: %s\n' "$staging_output_root" +printf 'Thread counts: %s\n' "${thread_counts[*]}" + +if [[ "$validate_only" -eq 1 ]]; then + printf 'Validation only mode enabled.\n' + for thread_count in "${thread_counts[@]}"; do + thread_output_dir="$staging_output_root/threads_${thread_count}" + printf 'Planned output directory: %s\n' "$thread_output_dir" + done + exit 0 +fi + +mkdir -p "$staging_output_root" + +cleanup_on_error() { + if [[ -d "$staging_output_root" ]]; then + printf 'Keeping staged sweep output for inspection: %s\n' "$staging_output_root" >&2 + fi +} +trap cleanup_on_error ERR + +for thread_count in "${thread_counts[@]}"; do + thread_output_dir="$staging_output_root/threads_${thread_count}" + + printf '========================================\n' + printf 'Running sweep for thread count: %s\n' "$thread_count" + printf '========================================\n' + + invoke_julia_script \ + "$resolved_julia_exe" \ + "$repo_root" \ + "$resolved_generate_julia_script" \ + "$thread_output_dir" \ + "Phase 1 export for ${thread_count} thread(s)" \ + "$thread_count" \ + 1 \ + "${phase1_extra_args[@]-}" + + invoke_dynare_phase \ + "$thread_output_dir" \ + "$thread_count" \ + "$resolved_dynare_script" \ + "$resolved_extract_script" + + invoke_julia_script \ + "$resolved_julia_exe" \ + "$repo_root" \ + "$resolved_compare_script" \ + "$thread_output_dir" \ + "Phase 3 compare for ${thread_count} thread(s)" \ + "$thread_count" \ + 1 +done + +invoke_julia_script \ + "$resolved_julia_exe" \ + "$repo_root" \ + "$resolved_sweep_compare_script" \ + "$staging_output_root" \ + 'Cross-thread benchmark summary' \ + 1 \ + 0 + +publish_staged_output_root \ + "$staging_output_root" \ + "$resolved_output_root" \ + "$resolved_output_parent" \ + "$output_root_leaf" + +trap - ERR +printf 'Thread sweep complete.\n' \ No newline at end of file diff --git a/test/dynare_comparison/run_thread_sweep_windows.ps1 b/test/dynare_comparison/run_thread_sweep_windows.ps1 new file mode 100644 index 000000000..b4254aff7 --- /dev/null +++ b/test/dynare_comparison/run_thread_sweep_windows.ps1 @@ -0,0 +1,287 @@ +# Example invocation: +# powershell.exe -NoProfile -ExecutionPolicy Bypass -File "D:\CustomTools\MacroModelling.jl\test\dynare_comparison\run_thread_sweep_windows.ps1" -JuliaExe "D:\CustomTools\julia-1.12.6\bin\julia.exe" -DynareMatlabPath "D:\CustomTools\dynare-7.0-win\matlab" -MatlabExe "C:\Program Files\MATLAB\R2024b\bin\matlab.exe" + +[CmdletBinding()] +param( + [int[]]$ThreadCounts = @(1, 2, 4, 8), + [string]$OutputRoot, + [string]$JuliaExe, + [string]$GenerateJuliaScript, + [string]$DynareScript, + [string]$CompareScript, + [string]$SweepCompareScript, + [string]$DynareMatlabPath, + [string]$MatlabExe, + [ValidateRange(0, 10)] + [int]$MaxLicenseRetries = 3, + [ValidateRange(0, 600)] + [int]$LicenseRetryDelaySeconds = 10, + [string[]]$OnlyModels = @(), + [switch]$ValidateOnly +) + +Set-StrictMode -Version Latest +$ErrorActionPreference = 'Stop' + +$scriptRoot = Split-Path -Parent $PSCommandPath +$repoRoot = Split-Path -Parent (Split-Path -Parent $scriptRoot) + +if (-not $OutputRoot) { + $OutputRoot = Join-Path $scriptRoot 'output_thread_sweep' +} +if (-not $GenerateJuliaScript) { + $GenerateJuliaScript = Join-Path $scriptRoot 'generate_julia_results.jl' +} +if (-not $DynareScript) { + $DynareScript = Join-Path $scriptRoot 'run_all_dynare_windows.ps1' +} +if (-not $CompareScript) { + $CompareScript = Join-Path $scriptRoot 'compare_results.jl' +} +if (-not $SweepCompareScript) { + $SweepCompareScript = Join-Path $scriptRoot 'compare_thread_sweep_results.jl' +} + +function Resolve-ExistingPath { + param( + [string[]]$Candidates, + [string]$Description + ) + + foreach ($candidate in $Candidates) { + if ([string]::IsNullOrWhiteSpace($candidate)) { + continue + } + + if (Test-Path -LiteralPath $candidate) { + return (Resolve-Path -LiteralPath $candidate).Path + } + } + + throw "Could not find $Description. Checked: $($Candidates -join ', ')" +} + +function Get-JuliaExecutable { + param([string]$PreferredPath) + + $candidates = @() + + if ($PreferredPath) { + $candidates += $PreferredPath + } + if ($env:JULIA_EXE) { + $candidates += $env:JULIA_EXE + } + + $candidates += 'D:\CustomTools\julia-1.12.6\bin\julia.exe' + $candidates += 'D:\CustomTools\julia-1.12.4\bin\julia.exe' + + if ($env:USERPROFILE) { + $candidates += (Join-Path $env:USERPROFILE '.juliaup\bin\julia.exe') + } + + $juliaCommand = Get-Command julia.exe -ErrorAction SilentlyContinue + if ($juliaCommand) { + $candidates += $juliaCommand.Source + } + + $localPrograms = Join-Path $env:LOCALAPPDATA 'Programs' + if (Test-Path -LiteralPath $localPrograms) { + Get-ChildItem -LiteralPath $localPrograms -Directory -Filter 'Julia*' -ErrorAction SilentlyContinue | + Sort-Object Name -Descending | + ForEach-Object { + $candidates += (Join-Path $_.FullName 'bin\julia.exe') + } + } + + Resolve-ExistingPath -Candidates $candidates -Description 'Julia executable' +} + +function Invoke-JuliaScript { + param( + [string]$Executable, + [string]$ProjectRoot, + [string]$ScriptPath, + [string]$OutputArgument, + [string]$Description, + [int]$RequestedThreadCount, + [string[]]$ExtraScriptArgs, + [switch]$UseThreadCount + ) + + $juliaArgs = @("--project=$ProjectRoot") + if ($UseThreadCount) { + $juliaArgs += "--threads=$RequestedThreadCount" + } + $juliaArgs += $ScriptPath + if ($ExtraScriptArgs) { + foreach ($extraArg in $ExtraScriptArgs) { + if (-not [string]::IsNullOrWhiteSpace($extraArg)) { + $juliaArgs += $extraArg + } + } + } + $juliaArgs += $OutputArgument + + Write-Host "Running Julia step: $Description" + & $Executable @juliaArgs + + if ($null -ne $LASTEXITCODE -and $LASTEXITCODE -ne 0) { + throw "Julia step failed ($Description) with exit code $LASTEXITCODE" + } +} + +function Invoke-DynarePhase { + param( + [string]$ScriptPath, + [string]$ThreadOutputDir, + [int]$RequestedThreadCount, + [string]$PreferredDynareMatlabPath, + [string]$PreferredMatlabExe, + [int]$RequestedMaxLicenseRetries, + [int]$RequestedLicenseRetryDelaySeconds, + [string[]]$RequestedOnlyModels, + [switch]$ValidationOnly + ) + + $dynareParameters = @{ + OutputDir = $ThreadOutputDir + ThreadCount = $RequestedThreadCount + MaxLicenseRetries = $RequestedMaxLicenseRetries + LicenseRetryDelaySeconds = $RequestedLicenseRetryDelaySeconds + } + + if ($PreferredDynareMatlabPath) { + $dynareParameters.DynareMatlabPath = $PreferredDynareMatlabPath + } + if ($PreferredMatlabExe) { + $dynareParameters.MatlabExe = $PreferredMatlabExe + } + if ($RequestedOnlyModels -and $RequestedOnlyModels.Count -gt 0) { + $dynareParameters.OnlyModels = $RequestedOnlyModels + } + if ($ValidationOnly) { + $dynareParameters.ValidateOnly = $true + } + + Write-Host "Running Dynare step for $RequestedThreadCount thread(s)" + + # This invokes another PowerShell script, so rely on terminating errors + # from that script rather than $LASTEXITCODE (which may be stale from a + # previously executed native command). + & $ScriptPath @dynareParameters + + if (-not $?) { + throw "Dynare step failed for $RequestedThreadCount thread(s)." + } +} + +function New-StagingOutputRoot { + param( + [string]$FinalOutputRoot, + [string]$ResolvedOutputParent, + [string]$OutputRootLeaf + ) + + Join-Path $ResolvedOutputParent ("{0}.__staging_{1}_{2}" -f $OutputRootLeaf, (Get-Date -Format 'yyyyMMddHHmmssfff'), (Get-Random -Minimum 10000 -Maximum 99999)) +} + +function Publish-StagedOutputRoot { + param( + [string]$StageOutputRoot, + [string]$FinalOutputRoot, + [string]$ResolvedOutputParent, + [string]$OutputRootLeaf + ) + + if (-not (Test-Path -LiteralPath $StageOutputRoot)) { + throw "Staged sweep output not found: $StageOutputRoot" + } + + if (Test-Path -LiteralPath $FinalOutputRoot) { + $previousOutputRoot = Join-Path $ResolvedOutputParent ("{0}.__previous_{1}_{2}" -f $OutputRootLeaf, (Get-Date -Format 'yyyyMMddHHmmssfff'), (Get-Random -Minimum 10000 -Maximum 99999)) + Write-Host ("Moving existing output root aside: {0} -> {1}" -f $FinalOutputRoot, $previousOutputRoot) + Move-Item -LiteralPath $FinalOutputRoot -Destination $previousOutputRoot + } + + Write-Host ("Publishing staged sweep output: {0} -> {1}" -f $StageOutputRoot, $FinalOutputRoot) + Move-Item -LiteralPath $StageOutputRoot -Destination $FinalOutputRoot +} + +$requestedOutputRoot = $OutputRoot +$outputRootLeaf = Split-Path -Leaf $requestedOutputRoot +$outputRootParent = Split-Path -Parent $requestedOutputRoot +if ([string]::IsNullOrWhiteSpace($outputRootParent)) { + $outputRootParent = '.' +} +if (-not (Test-Path -LiteralPath $outputRootParent)) { + New-Item -ItemType Directory -Path $outputRootParent -Force | Out-Null +} +$resolvedOutputParent = (Resolve-Path -LiteralPath $outputRootParent).Path +$resolvedOutputRoot = Join-Path $resolvedOutputParent $outputRootLeaf +$stagingOutputRoot = New-StagingOutputRoot -FinalOutputRoot $resolvedOutputRoot -ResolvedOutputParent $resolvedOutputParent -OutputRootLeaf $outputRootLeaf +$resolvedJuliaExe = Get-JuliaExecutable -PreferredPath $JuliaExe +$resolvedGenerateJuliaScript = Resolve-ExistingPath -Candidates @($GenerateJuliaScript) -Description 'Julia phase-1 script' +$resolvedDynareScript = Resolve-ExistingPath -Candidates @($DynareScript) -Description 'Dynare phase-2 script' +$resolvedCompareScript = Resolve-ExistingPath -Candidates @($CompareScript) -Description 'Julia phase-3 script' +$resolvedSweepCompareScript = Resolve-ExistingPath -Candidates @($SweepCompareScript) -Description 'thread-sweep summary script' + +$resolvedThreadCounts = $ThreadCounts | Sort-Object -Unique +if (-not $resolvedThreadCounts) { + throw 'At least one thread count must be provided.' +} + +$resolvedOnlyModels = @() +if ($OnlyModels) { + $resolvedOnlyModels = @($OnlyModels | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }) +} +$phase1ExtraArgs = @() +if ($resolvedOnlyModels.Count -gt 0) { + $phase1ExtraArgs += ("--only-models={0}" -f ($resolvedOnlyModels -join ',')) + Write-Host ("Restricting sweep to models: {0}" -f ($resolvedOnlyModels -join ', ')) +} + +Write-Host "Repository root: $repoRoot" +Write-Host "Julia executable: $resolvedJuliaExe" +Write-Host "Final sweep output root: $resolvedOutputRoot" +Write-Host "Sweep staging root: $stagingOutputRoot" +Write-Host ("Thread counts: {0}" -f ($resolvedThreadCounts -join ', ')) + +if ($ValidateOnly) { + Write-Host 'Validation only mode enabled.' + foreach ($threadCount in $resolvedThreadCounts) { + $threadOutputDir = Join-Path $stagingOutputRoot ("threads_{0}" -f $threadCount) + Write-Host ("Planned output directory: {0}" -f $threadOutputDir) + } + return +} + +New-Item -ItemType Directory -Path $stagingOutputRoot -Force | Out-Null + +try { + foreach ($threadCount in $resolvedThreadCounts) { + $threadOutputDir = Join-Path $stagingOutputRoot ("threads_{0}" -f $threadCount) + + Write-Host '========================================' + Write-Host ("Running sweep for thread count: {0}" -f $threadCount) + Write-Host '========================================' + + Invoke-JuliaScript -Executable $resolvedJuliaExe -ProjectRoot $repoRoot -ScriptPath $resolvedGenerateJuliaScript -OutputArgument $threadOutputDir -Description ("Phase 1 export for {0} thread(s)" -f $threadCount) -RequestedThreadCount $threadCount -ExtraScriptArgs $phase1ExtraArgs -UseThreadCount + + Invoke-DynarePhase -ScriptPath $resolvedDynareScript -ThreadOutputDir $threadOutputDir -RequestedThreadCount $threadCount -PreferredDynareMatlabPath $DynareMatlabPath -PreferredMatlabExe $MatlabExe -RequestedMaxLicenseRetries $MaxLicenseRetries -RequestedLicenseRetryDelaySeconds $LicenseRetryDelaySeconds -RequestedOnlyModels $resolvedOnlyModels + + Invoke-JuliaScript -Executable $resolvedJuliaExe -ProjectRoot $repoRoot -ScriptPath $resolvedCompareScript -OutputArgument $threadOutputDir -Description ("Phase 3 compare for {0} thread(s)" -f $threadCount) -RequestedThreadCount $threadCount -UseThreadCount + } + + Invoke-JuliaScript -Executable $resolvedJuliaExe -ProjectRoot $repoRoot -ScriptPath $resolvedSweepCompareScript -OutputArgument $stagingOutputRoot -Description 'Cross-thread benchmark summary' -RequestedThreadCount 1 + + Publish-StagedOutputRoot -StageOutputRoot $stagingOutputRoot -FinalOutputRoot $resolvedOutputRoot -ResolvedOutputParent $resolvedOutputParent -OutputRootLeaf $outputRootLeaf +} +catch { + if (Test-Path -LiteralPath $stagingOutputRoot) { + Write-Warning ("Keeping staged sweep output for inspection: {0}" -f $stagingOutputRoot) + } + throw +} + +Write-Host 'Thread sweep complete.' \ No newline at end of file diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index 873de5a12..e0dbcee9f 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -1,4 +1,30 @@ +import Zygote, FiniteDifferences, ForwardDiff, Mooncake, DifferentiationInterface, ADTypes +import MatrixEquations +import LinearAlgebra as ℒ +import StatsPlots +using Random +Random.seed!(1234) + +# Diagnostic wrapper: prints achieved atol/rtol when isapprox fails +function check_isapprox(a, b; kwargs...) + result = isapprox(a, b; kwargs...) + if !result + d = a .- b + frobenius_diff = ℒ.norm(d) + maxnorm = max(ℒ.norm(a), ℒ.norm(b)) + eff_rtol = maxnorm > 0 ? frobenius_diff / maxnorm : Inf + max_abs = maximum(abs.(d)) + safe_denom = max.(abs.(a), abs.(b), eps()) + max_rel = maximum(abs.(d) ./ safe_denom) + has_nan = any(isnan, a) || any(isnan, b) + has_inf = any(isinf, a) || any(isinf, b) + printstyled(" ⚠ APPROX FAIL: eff_rtol=$(eff_rtol), max_elem_abs=$(max_abs), max_elem_rel=$(max_rel), has_nan=$(has_nan), has_inf=$(has_inf), size=$(size(a))\n", color=:yellow) + end + return result +end + function functionality_test(m, m2; algorithm = :first_order, plots = true) + rndnmbr = rand(max(length(m.parameter_values),2)) old_params = copy(m.parameter_values) old_params2 = copy(m2.parameter_values) @@ -12,20 +38,20 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithms = [:doubling, :bartels_stewart, :bicgstab, :gmres] params = [old_params, - (m.constants.post_complete_parameters.parameters[1] => old_params[1] * exp(rand()*1e-4)), + (m.constants.post_complete_parameters.parameters[1] => old_params[1] * exp(rndnmbr[1]*1e-4)), Tuple(m.constants.post_complete_parameters.parameters[1:2] .=> old_params[1:2] .* 1.0001), m.constants.post_complete_parameters.parameters .=> old_params, (string(m.constants.post_complete_parameters.parameters[1]) => old_params[1] * 1.0001), - Tuple(string.(m.constants.post_complete_parameters.parameters[1:2]) .=> old_params[1:2] .* exp.(rand(2)*1e-4)), + Tuple(string.(m.constants.post_complete_parameters.parameters[1:2]) .=> old_params[1:2] .* exp.(-rndnmbr[1:2]*1e-4)), old_params] params2 = [old_params2, - (m2.constants.post_complete_parameters.parameters[1] => old_params2[1] * exp(rand()*1e-4)), + (m2.constants.post_complete_parameters.parameters[1] => old_params2[1] * exp(rndnmbr[1]*1e-4)), Tuple(m2.constants.post_complete_parameters.parameters[1:2] .=> old_params2[1:2] .* 1.0001), m2.constants.post_complete_parameters.parameters .=> old_params2, (string(m2.constants.post_complete_parameters.parameters[1]) => old_params2[1] * 1.0001), - Tuple(string.(m2.constants.post_complete_parameters.parameters[1:2]) .=> old_params2[1:2] .* exp.(rand(2)*1e-4)), + Tuple(string.(m2.constants.post_complete_parameters.parameters[1:2]) .=> old_params2[1:2] .* exp.(-rndnmbr[1:2]*1e-4)), old_params2] param_derivs = [:all, @@ -185,7 +211,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 +246,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 +493,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 +749,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 +773,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 +945,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 +1170,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 +1205,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 @@ -1497,7 +1523,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(estim1, estim2, rtol = 1e-8) + @test check_isapprox(estim1, estim2, rtol = 1e-8) clear_solution_caches!(m, algorithm) @@ -1522,7 +1548,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(estim1, estim2, rtol = 1e-8) + @test check_isapprox(estim1, estim2, rtol = 1e-8) for levels in [true, false] clear_solution_caches!(m, algorithm) @@ -1550,7 +1576,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(estim1, estim2, rtol = 1e-8) + @test check_isapprox(estim1, estim2, rtol = 1e-8) clear_solution_caches!(m, algorithm) @@ -1578,7 +1604,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(estim1, estim2, rtol = 1e-8) + @test check_isapprox(estim1, estim2, rtol = 1e-8) end end end @@ -1588,7 +1614,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, @@ -1670,14 +1696,14 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, lyapunov_algorithm = lyapunov_algorithm, verbose = verbose) - @test isapprox(estim1,estim2) + @test check_isapprox(estim1,estim2) end end end 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, @@ -1696,8 +1722,8 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) for presample_periods in [0, 3] 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 parameter_values in [old_params, old_params .* exp.(-rndnmbr[1:length(old_params)]*1e-4)] + 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 +1732,47 @@ 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) + + moon_grad_llh = DifferentiationInterface.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), ADTypes.AutoMooncake(config = nothing), parameter_values) + + 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)[1] + + 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 ? 5 : 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 check_isapprox(fin_grad_llh[1], moon_grad_llh, rtol = 1e-4, atol = 1e-6) + @test check_isapprox(fin_grad_llh[1], zyg_grad_llh, rtol = 1e-4, atol = 1e-6) + @test check_isapprox(fin_grad_llh[1], moon_grad_llh, rtol = 1e-4, atol = 1e-6) + @test check_isapprox(fin_grad_llh[1], zyg_grad_llh, rtol = 1e-4, atol = 1e-6) + break + end end end - end for quadratic_matrix_equation_algorithm in qme_algorithms for lyapunov_algorithm in lyapunov_algorithms @@ -1753,22 +1790,34 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(llh, LLH, rtol = 1e-8) + @test check_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) + + MOON_grad_llh = DifferentiationInterface.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), ADTypes.AutoMooncake(config = nothing), parameter_values) + + 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)[1] + + @test check_isapprox(MOON_grad_llh, moon_grad_llh, rtol = 1e-6) + @test check_isapprox(ZYG_grad_llh, zyg_grad_llh, rtol = 1e-6) end end end @@ -1872,7 +1921,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 @@ -1907,7 +1956,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(cond_fcst, cond_fcst_lvl) + @test check_isapprox(cond_fcst, cond_fcst_lvl) clear_solution_caches!(m, algorithm) @@ -1938,7 +1987,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(cond_fcst, cond_fcst_lvl) + @test check_isapprox(cond_fcst, cond_fcst_lvl) end # end end @@ -2044,7 +2093,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 @@ -2057,7 +2106,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, verbose = verbose) - @test isapprox(var_decomp, VAR_DECOMP, rtol = 1e-8) + @test check_isapprox(var_decomp, VAR_DECOMP, rtol = 1e-8, nans = true) clear_solution_caches!(m, algorithm) @@ -2067,7 +2116,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, verbose = verbose) - @test isapprox(cond_var_decomp, COND_VAR_DECOMP, rtol = 1e-8) + @test check_isapprox(cond_var_decomp, COND_VAR_DECOMP, rtol = 1e-8, nans = true) end @@ -2082,7 +2131,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(corrl, CORRL, rtol = 1e-5) + @test check_isapprox(corrl, CORRL, rtol = 1e-5, nans = true) clear_solution_caches!(m, algorithm) @@ -2094,7 +2143,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(autocorr_, AUTOCORR, rtol = 1e-8) + @test check_isapprox(autocorr_, AUTOCORR, rtol = 1e-8, nans = true) end end end @@ -2127,7 +2176,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) @@ -2138,7 +2187,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm = sylvester_algorithm, verbose = verbose) - @test isapprox(sol, SOL)#, rtol = eps(Float32)) + @test check_isapprox(sol, SOL)#, rtol = eps(Float32)) end end end @@ -2146,7 +2195,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end @testset "get_solution with parameter input" begin - for parameter_values in [old_params, old_params .* exp.(rand(length(old_params))*1e-4)] + for parameter_values in [old_params, old_params .* exp.(-rndnmbr[1:length(old_params)]*1e-4)] get_first_order_solution(m, parameter_values) get_perturbation_solution(m, parameter_values) @@ -2161,37 +2210,50 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sol = get_solution(m, parameter_values, algorithm = algorithm) - clear_solution_caches!(m, algorithm) + # Helper to extract element i in flattened order: 1→SS, 2→sol_mats[1], 3→sol_mats[2], ... + _sol_el(s, i) = i == 1 ? s[1] : s[2][i-1] - 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 = nothing + deriv_sol_zyg = nothing + 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 -> _sol_el(get_solution(m, x, algorithm = algorithm), i), parameter_values)) + 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_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) + + _sol_el(get_solution(m, x, algorithm = algorithm), i) + end, 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 + clear_solution_caches!(m, algorithm) - @test isapprox(deriv_sol_zyg, deriv_sol_fin, rtol = 1e-5) - - @test isapprox(deriv_sol, deriv_sol_fin, rtol = 1e-5) + deriv_sol_moon = [] + for i in 1:length(sol[2]) + push!(deriv_sol_moon, DifferentiationInterface.jacobian(x -> _sol_el(get_solution(m, x, algorithm = algorithm), i), ADTypes.AutoMooncake(config = nothing), parameter_values)) + end + + clear_solution_caches!(m, algorithm) - 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)] + deriv_sol_zyg = [] + for i in 1:length(sol[2]) + push!(deriv_sol_zyg, Zygote.jacobian(x -> _sol_el(get_solution(m, x, algorithm = algorithm), i), parameter_values)[1]) + end + + @test check_isapprox(deriv_sol_moon, deriv_sol_fin, rtol = 1e-5) + @test check_isapprox(deriv_sol_zyg, deriv_sol_fin, rtol = 1e-5) + + @test check_isapprox(deriv_sol, deriv_sol_fin, rtol = 1e-5) + + 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) @@ -2200,31 +2262,42 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, sylvester_algorithm = sylvester_algorithm) - @test isapprox([s for s in sol[1:end-1]], [S for S in SOL[1:end-1]], rtol = 1e-8) + @test check_isapprox(vcat([sol[1]], sol[2]), vcat([SOL[1]], SOL[2]), 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 -> _sol_el(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 check_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_moon = [] + for i in 1:length(sol[2]) + push!(DERIV_SOL_moon, DifferentiationInterface.jacobian(x -> _sol_el(get_solution(m, x, algorithm = algorithm, + tol = tol, + quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, + sylvester_algorithm = sylvester_algorithm), i), ADTypes.AutoMooncake(config = nothing), parameter_values)) + end - @test isapprox(deriv_sol_zyg, DERIV_SOL_zyg, rtol = 1e-8) + clear_solution_caches!(m, algorithm) + + DERIV_SOL_zyg = [] + for i in 1:length(sol[2]) + push!(DERIV_SOL_zyg, Zygote.jacobian(x -> _sol_el(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 check_isapprox(DERIV_SOL_moon, DERIV_SOL, rtol = 1e-8) + @test check_isapprox(DERIV_SOL_zyg, DERIV_SOL, rtol = 1e-8) end end end @@ -2274,7 +2347,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) @testset "get_irf with parameter input" begin if algorithm == :first_order - for parameter_values in [old_params, old_params .* exp.(rand(length(old_params))*1e-4)] + for parameter_values in [old_params, old_params .* exp.(-rndnmbr[1:length(old_params)]*1e-4)] for levels in [true,false] for negative_shock in [true,false] for periods in [1,10] @@ -2312,19 +2385,76 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) deriv_for = ForwardDiff.jacobian(x->get_irf(m, x, initial_state = initial_state)[:,1,1], parameter_values) for i in 1:100 - local deriv_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-4), + local deriv_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 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[1])) - @test isapprox(deriv_for, deriv_fin[1], rtol = 1e-5) + @test check_isapprox(deriv_for, deriv_fin[1], rtol = 1e-5) + break + end + end + + clear_solution_caches!(m, algorithm) + + deriv_moon = DifferentiationInterface.jacobian(x -> get_irf(m, x, initial_state = initial_state)[:,1,1], ADTypes.AutoMooncake(config = nothing), parameter_values) + 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 ? 5 : 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 check_isapprox(deriv_moon, deriv_fin_zyg[1], rtol = 1e-5) + @test check_isapprox(deriv_zyg, deriv_fin_zyg[1], rtol = 1e-5) 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 ? 5 : 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 check_isapprox(deriv_for_last, deriv_fin_last[1], rtol = 1e-4) + break + end + end + + # Last period derivative tests (Mooncake) + clear_solution_caches!(m, algorithm) + + deriv_moon_last = DifferentiationInterface.jacobian(x -> get_irf(m, x, initial_state = initial_state)[:,end,1], ADTypes.AutoMooncake(config = nothing), parameter_values) + 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 ? 5 : 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 check_isapprox(deriv_moon_last, deriv_fin_zyg_last[1], rtol = 1e-5) + @test check_isapprox(deriv_zyg_last, deriv_fin_zyg_last[1], rtol = 1e-5) + 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) @@ -2333,12 +2463,12 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) initial_state = initial_state, tol = tol, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm) - @test isapprox(irf_, IRF_, rtol = 1e-8) + @test check_isapprox(irf_, IRF_, rtol = 1e-8) DERIV_for = ForwardDiff.jacobian(x->get_irf(m, x, initial_state = initial_state, tol = tol, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm)[:,1,1], parameter_values) - @test isapprox(deriv_for, DERIV_for, rtol = 1e-8) + @test check_isapprox(deriv_for, DERIV_for, rtol = 1e-8) end end for variables in vars @@ -2355,7 +2485,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) @testset "get_statistics" begin - for parameter_values in [old_params, old_params .* exp.(rand(length(old_params))*1e-4)] + for parameter_values in [old_params, old_params .* exp.(-rndnmbr[1:length(old_params)]*1e-4)] for non_stochastic_steady_state in (Symbol[], vars...) for mean in (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? (Symbol[], vars[1]) : Symbol[]) for standard_deviation in (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? (Symbol[], vars[1]) : Symbol[]) @@ -2372,6 +2502,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) standard_deviation = standard_deviation, variance = variance, covariance = covariance, + correlation = covariance, autocorrelation = autocorrelation ) end @@ -2384,19 +2515,20 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) - for parameter_values in [old_params, old_params .* exp.(rand(length(old_params))*1e-4)] + for parameter_values in [old_params, old_params .* exp.(-rndnmbr[1:length(old_params)]*1e-4)] 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[]), variance = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), covariance = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), + correlation = (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 @@ -2408,6 +2540,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) standard_deviation = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), variance = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), covariance = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), + correlation = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), autocorrelation = (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? :all : Symbol[]), tol = tol, quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, @@ -2420,14 +2553,15 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) # println("standard_deviation: $(ℒ.norm(stats[:standard_deviation] - STATS[:standard_deviation]) / max(ℒ.norm(stats[:standard_deviation]), ℒ.norm(STATS[:standard_deviation])))") # println("covariance: $(ℒ.norm(stats[:covariance] - STATS[:covariance]) / max(ℒ.norm(stats[:covariance]), ℒ.norm(STATS[:covariance])))") # println("autocorrelation (qme: $quadratic_matrix_equation_algorithm, sylv: $sylvester_algorithm, lyap: $lyapunov_algorithm, tol: $tol): $(ℒ.norm(stats[:autocorrelation] - STATS[:autocorrelation]) / max(ℒ.norm(stats[:autocorrelation]), ℒ.norm(STATS[:autocorrelation])))") - @test isapprox(stats[:non_stochastic_steady_state], STATS[:non_stochastic_steady_state], rtol = 1e-8) - @test isapprox(stats[:mean], STATS[:mean], rtol = 1e-8) - @test isapprox(stats[:standard_deviation], STATS[:standard_deviation], rtol = 1e-8) - @test isapprox(stats[:variance], STATS[:variance], rtol = 1e-8) - @test isapprox(stats[:covariance], STATS[:covariance], rtol = 1e-8) - @test isapprox(stats[:autocorrelation], STATS[:autocorrelation], rtol = 1e-8) + @test check_isapprox(stats[:non_stochastic_steady_state], STATS[:non_stochastic_steady_state], rtol = 1e-8) + @test check_isapprox(stats[:mean], STATS[:mean], rtol = 1e-8) + @test check_isapprox(stats[:standard_deviation], STATS[:standard_deviation], rtol = 1e-8) + @test check_isapprox(stats[:variance], STATS[:variance], rtol = 1e-8) + @test check_isapprox(stats[:covariance], STATS[:covariance], rtol = 1e-8, atol = 1e-8) + @test check_isapprox(stats[:correlation], STATS[:correlation], rtol = 1e-8, atol = 1e-8, nans = true) + @test check_isapprox(stats[:autocorrelation], STATS[:autocorrelation], rtol = 1e-8, atol = 1e-8, nans = true) else - @test isapprox(stats[:non_stochastic_steady_state], STATS[:non_stochastic_steady_state], rtol = 1e-8) + @test check_isapprox(stats[:non_stochastic_steady_state], STATS[:non_stochastic_steady_state], rtol = 1e-8) end end end @@ -2436,44 +2570,49 @@ 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_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + non_stochastic_steady_state = :all_excluding_obc)[:non_stochastic_steady_state], ADTypes.AutoMooncake(config = nothing), 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)[1] - 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) / max(ℒ.norm(deriv1), ℒ.norm(deriv1_zyg)) + + @test check_isapprox(deriv1_moon, deriv1_fin[1], rtol = 1e-5) + @test check_isapprox(deriv1_zyg, deriv1_fin[1], rtol = 1e-5) + + @test check_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_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + mean = :all_excluding_obc)[:mean], ADTypes.AutoMooncake(config = nothing), old_params) deriv2_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, - mean = :all_excluding_obc)[:mean], old_params) + mean = :all_excluding_obc)[:mean], old_params)[1] end for i in 1:100 @@ -2487,11 +2626,12 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end, old_params) if isfinite(ℒ.norm(deriv2_fin[1])) - if algorithm == :first_order - @test isapprox(deriv2_zyg[1], deriv2_fin[1], rtol = 1e-5) + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + @test check_isapprox(deriv2_moon, deriv2_fin[1], rtol = 1e-5) + @test check_isapprox(deriv2_zyg, deriv2_fin[1], rtol = 1e-5) end - @test isapprox(deriv2, deriv2_fin[1], rtol = 1e-5) + @test check_isapprox(deriv2, deriv2_fin[1], rtol = 1e-5) break end end @@ -2501,13 +2641,15 @@ 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_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + standard_deviation = :all_excluding_obc)[:standard_deviation], ADTypes.AutoMooncake(config = nothing), old_params) deriv3_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, - standard_deviation = :all_excluding_obc)[:standard_deviation], old_params) + standard_deviation = :all_excluding_obc)[:standard_deviation], old_params)[1] end for i in 1:100 - local deriv3_fin = FiniteDifferences.jacobian(FiniteDifferences.forward_fdm(3,1, max_range = 1e-3), + local deriv3_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 1, max_range = 1e-3), x -> begin clear_solution_caches!(m, algorithm) @@ -2515,11 +2657,12 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end, old_params) if isfinite(ℒ.norm(deriv3_fin[1])) - if algorithm == :first_order - @test isapprox(deriv3_zyg[1], deriv3_fin[1], rtol = 1e-5) + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + @test check_isapprox(deriv3_moon, deriv3_fin[1], rtol = 1e-5, atol = 1e-8) + @test check_isapprox(deriv3_zyg, deriv3_fin[1], rtol = 1e-5, atol = 1e-8) end - @test isapprox(deriv3, deriv3_fin[1], rtol = 1e-5) + @test check_isapprox(deriv3, deriv3_fin[1], rtol = 1e-5, atol = 1e-8) break end end @@ -2529,23 +2672,26 @@ 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_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + variance = :all_excluding_obc)[:variance], ADTypes.AutoMooncake(config = nothing), old_params) deriv4_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, - variance = :all_excluding_obc)[:variance], old_params) + variance = :all_excluding_obc)[:variance], old_params)[1] end for i in 1:100 - local deriv4_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 3 : 4, 1, max_range = 1e-3), + local deriv4_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 1, max_range = 1e-3), x -> begin clear_solution_caches!(m, algorithm) get_statistics(m, x, algorithm = algorithm, variance = :all_excluding_obc)[:variance] end, old_params) if isfinite(ℒ.norm(deriv4_fin[1])) - if algorithm == :first_order - @test isapprox(deriv4_zyg[1], deriv4_fin[1], rtol = 1e-5) + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + @test check_isapprox(deriv4_moon, deriv4_fin[1], rtol = 1e-5, atol = 1e-8) + @test check_isapprox(deriv4_zyg, deriv4_fin[1], rtol = 1e-5, atol = 1e-8) end - @test isapprox(deriv4, deriv4_fin[1], rtol = 1e-5) + @test check_isapprox(deriv4, deriv4_fin[1], rtol = 1e-5, atol = 1e-8) break end end @@ -2553,67 +2699,190 @@ 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_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + 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], ADTypes.AutoMooncake(config = nothing), old_params) 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), - covariance = :all_excluding_obc)[:covariance], old_params) + 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)[1] end for i in 1:100 - local deriv5_fin = FiniteDifferences.jacobian(FiniteDifferences.forward_fdm(3,1, max_range = 1e-3), + local deriv5_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 1, max_range = 1e-3), x -> begin 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_ - @test isapprox(deriv5_zyg[1], deriv5_fin[1], rtol = 1e-4) + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + @test check_isapprox(deriv5_moon, deriv5_fin[1], rtol = 1e-4, atol = 1e-8) + @test check_isapprox(deriv5_zyg, deriv5_fin[1], rtol = 1e-4, atol = 1e-8) end # println(ℒ.norm(deriv5 - deriv5_fin[1]) / max(ℒ.norm(deriv5), ℒ.norm(deriv5_fin[1]))) - @test isapprox(deriv5, deriv5_fin[1], rtol = 1e-4) + @test check_isapprox(deriv5, deriv5_fin[1], rtol = 1e-4, atol = 1e-8) 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_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + autocorrelation = :all_excluding_obc)[:autocorrelation], ADTypes.AutoMooncake(config = nothing), old_params) + deriv6_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, + autocorrelation = :all_excluding_obc)[:autocorrelation], old_params)[1] + end + + for i in 1:100 + local deriv6_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 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 check_isapprox(deriv6_moon, deriv6_fin[1], rtol = 1e-4) + @test check_isapprox(deriv6_zyg, deriv6_fin[1], rtol = 1e-4) + end + @test check_isapprox(deriv6, deriv6_fin[1], rtol = 1e-4) + break + end + end + + clear_solution_caches!(m, algorithm) + + # Restrict the correlation jacobian comparison to non-degenerate + # variables. Degenerate-variance entries produce NaN/0-over-0 + # correlations whose FD jacobian is dominated by perturbation + # noise (huge magnitude), while AD computes the analytic value + # cleanly. Comparing only over non-degenerate entries keeps the + # AD-vs-FD check meaningful without silently masking real bugs. + corr_target_vars_jac = let + _all_vars_jac = m.constants.post_model_macro.var + _sd_jac = get_statistics(m, old_params, algorithm = algorithm, + standard_deviation = _all_vars_jac)[:standard_deviation] + _all_vars_jac[findall(>(1e-6), _sd_jac)] + end + + deriv7 = ForwardDiff.jacobian(x->get_statistics(m, x, algorithm = algorithm, + correlation = corr_target_vars_jac)[:correlation], old_params) + + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + deriv7_moon = DifferentiationInterface.jacobian(x->get_statistics(m, x, algorithm = algorithm, + correlation = corr_target_vars_jac)[:correlation], ADTypes.AutoMooncake(config = nothing), old_params) + deriv7_zyg = Zygote.jacobian(x->get_statistics(m, x, algorithm = algorithm, + correlation = corr_target_vars_jac)[:correlation], old_params)[1] + end + + for i in 1:100 + local deriv7_fin = FiniteDifferences.jacobian(FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 1, max_range = 1e-3), + x -> begin + clear_solution_caches!(m, algorithm) + + get_statistics(m, x, algorithm = algorithm, correlation = corr_target_vars_jac)[:correlation] + end, old_params) + if isfinite(ℒ.norm(deriv7_fin[1])) + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + @test check_isapprox(deriv7_moon, deriv7_fin[1], rtol = 1e-4, atol = 1e-8, nans = true) + @test check_isapprox(deriv7_zyg, deriv7_fin[1], rtol = 1e-4, atol = 1e-8, nans = true) + end + @test check_isapprox(deriv7, deriv7_fin[1], rtol = 1e-4, atol = 1e-8, nans = true) + break + end + end + + if algorithm == :pruned_third_order + var_obj = x -> begin + clear_solution_caches!(m, algorithm) + get_statistics(m, x, algorithm = algorithm, variance = :all_excluding_obc)[:variance] |> sum + end + + autocorr_obj = x -> begin + clear_solution_caches!(m, algorithm) + get_statistics(m, x, algorithm = algorithm, autocorrelation = :all_excluding_obc)[:autocorrelation] |> sum + end + + var_grad_moon = DifferentiationInterface.gradient(var_obj, ADTypes.AutoMooncake(config = nothing), old_params) + 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_moon) + @test all(isfinite, var_grad_zyg) + @test all(isfinite, var_grad_fin) + @test ℒ.norm(var_grad_moon - var_grad_fin) / max(ℒ.norm(var_grad_fin), eps()) < 1e-4 + @test ℒ.norm(var_grad_zyg - var_grad_fin) / max(ℒ.norm(var_grad_fin), eps()) < 1e-4 + + autocorr_grad_moon = DifferentiationInterface.gradient(autocorr_obj, ADTypes.AutoMooncake(config = nothing), old_params) + 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_moon) + @test all(isfinite, autocorr_grad_zyg) + @test all(isfinite, autocorr_grad_fin) + @test ℒ.norm(autocorr_grad_moon - autocorr_grad_fin) / max(ℒ.norm(autocorr_grad_fin), eps()) < 1e-4 + @test ℒ.norm(autocorr_grad_zyg - autocorr_grad_fin) / max(ℒ.norm(autocorr_grad_fin), eps()) < 1e-4 + + corr_obj = x -> begin + clear_solution_caches!(m, algorithm) + get_statistics(m, x, algorithm = algorithm, correlation = corr_target_vars_jac)[:correlation] |> sum + end + + corr_grad_moon = DifferentiationInterface.gradient(corr_obj, ADTypes.AutoMooncake(config = nothing), old_params) + corr_grad_zyg = Zygote.gradient(corr_obj, old_params)[1] + corr_grad_fin = FiniteDifferences.grad(FiniteDifferences.forward_fdm(3, 1, max_range = 1e-3), corr_obj, old_params)[1] + @test all(isfinite, corr_grad_moon) + @test all(isfinite, corr_grad_zyg) + @test all(isfinite, corr_grad_fin) + @test ℒ.norm(corr_grad_moon - corr_grad_fin) / max(ℒ.norm(corr_grad_fin), eps()) < 1e-4 + @test ℒ.norm(corr_grad_zyg - corr_grad_fin) / max(ℒ.norm(corr_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 check_isapprox(deriv1, DERIV1, rtol = 1e-8) + + DERIV1_moon = DifferentiationInterface.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], ADTypes.AutoMooncake(config = nothing), old_params) + 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)[1] + @test check_isapprox(DERIV1_moon, DERIV1, rtol = 1e-8) + @test check_isapprox(DERIV1_zyg, DERIV1, 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, @@ -2621,18 +2890,25 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, mean = :all_excluding_obc)[:mean], old_params) - @test isapprox(deriv2, DERIV2, rtol = 1e-8) + @test check_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_moon = DifferentiationInterface.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, + mean = :all_excluding_obc)[:mean], ADTypes.AutoMooncake(config = nothing), old_params) DERIV2_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, - mean = :all_excluding_obc)[:mean], old_params) - @test isapprox(deriv2_zyg[1], DERIV2_zyg[1], rtol = 1e-8) + mean = :all_excluding_obc)[:mean], old_params)[1] + @test check_isapprox(DERIV2_moon, DERIV2, rtol = 1e-8) + @test check_isapprox(DERIV2_zyg, DERIV2, rtol = 1e-8) end clear_solution_caches!(m, algorithm) @@ -2643,18 +2919,25 @@ 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, DERIV3, rtol = 1e-8) + @test check_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_moon = DifferentiationInterface.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, + standard_deviation = :all_excluding_obc)[:standard_deviation], ADTypes.AutoMooncake(config = nothing), old_params) DERIV3_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, - standard_deviation = :all_excluding_obc)[:standard_deviation], old_params) - @test isapprox(deriv3_zyg[1], DERIV3_zyg[1], rtol = 1e-8) + standard_deviation = :all_excluding_obc)[:standard_deviation], old_params)[1] + @test check_isapprox(DERIV3_moon, DERIV3, rtol = 1e-6) + @test check_isapprox(DERIV3_zyg, DERIV3, rtol = 1e-6) end clear_solution_caches!(m, algorithm) @@ -2665,18 +2948,25 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm, variance = :all_excluding_obc)[:variance], old_params) - @test isapprox(deriv4, DERIV4, rtol = 1e-8) + @test check_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_moon = DifferentiationInterface.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, + variance = :all_excluding_obc)[:variance], ADTypes.AutoMooncake(config = nothing), old_params) DERIV4_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, - variance = :all_excluding_obc)[:variance], old_params) - @test isapprox(deriv4_zyg[1], DERIV4_zyg[1], rtol = 1e-8) + variance = :all_excluding_obc)[:variance], old_params)[1] + @test check_isapprox(DERIV4_moon, DERIV4, rtol = 1e-8) + @test check_isapprox(DERIV4_zyg, DERIV4, rtol = 1e-8) end clear_solution_caches!(m, algorithm) @@ -2688,24 +2978,60 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) sylvester_algorithm = sylvester_algorithm, covariance = :all_excluding_obc)[:covariance], old_params) # println(ℒ.norm(deriv5 - DERIV5) / max(ℒ.norm(deriv5), ℒ.norm(DERIV5))) - @test isapprox(deriv5, DERIV5, rtol = 1e-4) + @test check_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_moon = DifferentiationInterface.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, + covariance = :all_excluding_obc)[:covariance], ADTypes.AutoMooncake(config = nothing), old_params) DERIV5_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, - covariance = :all_excluding_obc)[:covariance], old_params) - @test isapprox(deriv5_zyg[1], DERIV5_zyg[1], rtol = 1e-4) + covariance = :all_excluding_obc)[:covariance], old_params)[1] + @test check_isapprox(DERIV5_moon, DERIV5, rtol = 1e-4) + @test check_isapprox(DERIV5_zyg, DERIV5, 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 check_isapprox(deriv6, DERIV6, rtol = 1e-4) + + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + clear_solution_caches!(m, algorithm) + + DERIV6_moon = DifferentiationInterface.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], ADTypes.AutoMooncake(config = nothing), old_params) + 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)[1] + @test check_isapprox(DERIV6_moon, DERIV6, rtol = 1e-4) + @test check_isapprox(DERIV6_zyg, DERIV6, rtol = 1e-4) + end end end end end end - end end @@ -2731,8 +3057,8 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) covariance = m.constants.post_model_macro.var[4:5]) # Check that within-group covariances match - @test isapprox(stats_grouped[:covariance][1:2, 1:2], stats_non_grouped_1[:covariance], rtol = 1e-10) - @test isapprox(stats_grouped[:covariance][3:4, 3:4], stats_non_grouped_2[:covariance], rtol = 1e-10) + @test check_isapprox(stats_grouped[:covariance][1:2, 1:2], stats_non_grouped_1[:covariance], rtol = 1e-6, nans = true) + @test check_isapprox(stats_grouped[:covariance][3:4, 3:4], stats_non_grouped_2[:covariance], rtol = 1e-6, nans = true) # Check that cross-group covariances are zero @test all(stats_grouped[:covariance][1:2, 3:4] .== 0) @@ -2753,6 +3079,66 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end + @testset "get_statistics - correlation" begin + if algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] + # Pick the first 4 model variables that are non-degenerate (positive variance, + # well above sqrt(eps)). Some models (e.g. Smets_Wouters_2007) have + # near-constant variables in their leading positions which would produce + # NaN/Inf-like correlation entries and break the cov/(sd*sd') cross-check. + _all_vars = m.constants.post_model_macro.var + _all_sd = let s = get_statistics(m, old_params, algorithm = algorithm, + standard_deviation = _all_vars) + s[:standard_deviation] + end + _nondeg_idx = findall(>(1e-6), _all_sd) + vars_corr = _all_vars[_nondeg_idx] + + # Flat input: full correlation matrix among requested variables + stats_corr = get_statistics(m, old_params, algorithm = algorithm, + correlation = vars_corr) + @test haskey(stats_corr, :correlation) + @test stats_corr[:correlation] isa AbstractMatrix + @test size(stats_corr[:correlation]) == (length(vars_corr), length(vars_corr)) + # Diagonal must be 1 (or NaN for degenerate variables) + for i in 1:length(vars_corr) + @test check_isapprox(stats_corr[:correlation][i, i], 1.0, rtol = 1e-6, nans = true) + end + # Symmetric + @test check_isapprox(stats_corr[:correlation], stats_corr[:correlation]', rtol = 1e-6, nans = true) + # All entries in [-1, 1] (or NaN) + @test all(x -> isnan(x) || (-1 - 1e-6 <= x <= 1 + 1e-6), stats_corr[:correlation]) + + # Cross-check correlation = covariance / (std * std') + stats_combo = get_statistics(m, old_params, algorithm = algorithm, + standard_deviation = vars_corr, + covariance = vars_corr, + correlation = vars_corr) + cov_full = stats_combo[:covariance] + stats_combo[:covariance]' - ℒ.Diagonal(stats_combo[:covariance]) + sd = stats_combo[:standard_deviation] + expected_corr = cov_full ./ (sd * sd') + @test check_isapprox(stats_combo[:correlation], expected_corr, rtol = 1e-6, atol = 1e-8, nans = true) + + # Grouped correlation: cross-group entries are zero, within-group preserved + if length(vars_corr) >= 4 + stats_grouped_corr = get_statistics(m, old_params, algorithm = algorithm, + correlation = [vars_corr[1:2], vars_corr[3:4]]) + @test stats_grouped_corr[:correlation] isa Matrix + @test size(stats_grouped_corr[:correlation]) == (4, 4) + # Within-group blocks match unrestricted correlation + stats_block1 = get_statistics(m, old_params, algorithm = algorithm, + correlation = vars_corr[1:2]) + stats_block2 = get_statistics(m, old_params, algorithm = algorithm, + correlation = vars_corr[3:4]) + @test check_isapprox(stats_grouped_corr[:correlation][1:2, 1:2], stats_block1[:correlation], rtol = 1e-6, nans = true) + @test check_isapprox(stats_grouped_corr[:correlation][3:4, 3:4], stats_block2[:correlation], rtol = 1e-6, nans = true) + # Cross-group entries are zero + @test all(stats_grouped_corr[:correlation][1:2, 3:4] .== 0) + @test all(stats_grouped_corr[:correlation][3:4, 1:2] .== 0) + end + end + end + + @testset "get_moments" begin for non_stochastic_steady_state in [true, false] for mean in (algorithm ∈ [:first_order, :pruned_second_order, :pruned_third_order] ? [true, false] : [false]) @@ -2800,42 +3186,89 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) cov(m, algorithm = algorithm) + get_correlation(m, algorithm = algorithm) + + get_corr(m, algorithm = algorithm) + + corr(m, algorithm = algorithm) + get_mean(m, algorithm = algorithm) end 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], + correlation = 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], + correlation = 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 check_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,39 +3276,175 @@ 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 check_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 check_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.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 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 check_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.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 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 check_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.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 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 check_isapprox(cov_jac, fd[1], rtol = 1e-4, nans = true) + 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 check_isapprox(mean_jac, fd[1], rtol = 1e-4) + break end end + m.parameter_values .= old_params + end + + # Correlation derivatives + clear_solution_caches!(m, algorithm) + mom_corr = get_moments(m, algorithm = algorithm, non_stochastic_steady_state = false, standard_deviation = false, correlation = 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) + corr_ka = collect(mom_corr[:correlation]) + n_cr = size(corr_ka, 1) + corr_jac = reshape(corr_ka[:, :, 2:end], n_cr * n_cr, :) + + for i in 1:100 + local fd = FiniteDifferences.jacobian( + FiniteDifferences.central_fdm(length(m.constants.post_complete_parameters.parameters) > 20 ? 5 : 4, 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, correlation = 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)[:correlation])) + end, old_params) + if isfinite(ℒ.norm(fd[1])) + @test check_isapprox(corr_jac, fd[1], rtol = 1e-4, nans = true) + break + end end + m.parameter_values .= old_params end end @testset "get_irf" begin + m.parameter_values .= old_params + clear_solution_caches!(m, algorithm) Random.seed!(123) for ignore_obc in [true,false] @@ -2936,7 +3505,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 @@ -2951,7 +3520,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) quadratic_matrix_equation_algorithm = quadratic_matrix_equation_algorithm, # lyapunov_algorithm = lyapunov_algorithm, sylvester_algorithm = sylvester_algorithm) - @test isapprox(irf_, IRF_, rtol = 1e-6) + @test check_isapprox(irf_, IRF_, rtol = 1e-6) end # end end @@ -2985,7 +3554,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) @@ -2996,7 +3565,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) RES = get_non_stochastic_steady_state_residuals(m, values, tol = tol, verbose = false, parameters = parameters) - @test isapprox(res, RES, rtol = 1e-8) + @test check_isapprox(res, RES, rtol = 1e-8, atol = 1e-8, nans = true) end end @@ -3008,7 +3577,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) res2 = get_non_stochastic_steady_state_residuals(m, stst[1:3], tol = tol, verbose = false) - @test isapprox(res1, res2, rtol = 1e-8) + @test check_isapprox(res1, res2, rtol = 1e-8, atol = 1e-8, nans = true) get_residuals(m, stst) @@ -3078,7 +3647,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) algorithm = algorithm, stochastic = stochastic, derivatives = derivatives) - @test isapprox(NSSS, nsss, rtol = 1e-8) + @test check_isapprox(NSSS, nsss, rtol = 1e-8) end end end @@ -3089,7 +3658,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 +3670,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 check_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 check_isapprox(sss_jac, fd[1], rtol = 1e-4) + break + end + end + m.parameter_values .= old_params + end end GC.gc() @@ -3129,7 +3743,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) if algorithm == :first_order lvl_irfs = get_irf(m, old_params, verbose = true, levels = true, variables = :all) new_sub_lvl_irfs = get_irf(m, old_params, verbose = true, shocks = :none, initial_state = collect(lvl_irfs[:,5,1]), levels = true, variables = :all) - @test isapprox(collect(new_sub_lvl_irfs[:,1,:]), collect(lvl_irfs[:,6,1]),rtol = eps(Float32)) + @test check_isapprox(collect(new_sub_lvl_irfs[:,1,:]), collect(lvl_irfs[:,6,1]),rtol = eps(Float32)) end end \ No newline at end of file diff --git a/test/models/borrowing_constraint.jl b/test/models/borrowing_constraint.jl new file mode 100644 index 000000000..832b78f1f --- /dev/null +++ b/test/models/borrowing_constraint.jl @@ -0,0 +1,18 @@ +@model borrowing_constraint begin + Y[0] + B[0] = C[0] + R * B[-1] + + log(Y[0]) = ρ * log(Y[-1]) + σ * ε[x] + + C[0]^(-γ) = β * R * C[1]^(-γ) + λ[0] + + 0 = max(B[0] - m * Y[0], -λ[0]) +end + +@parameters borrowing_constraint begin + R = 1.05 + β = 0.945 + ρ = 0.9 + σ = 0.05 + m = 1 + γ = 1 +end diff --git a/test/runtests.jl b/test/runtests.jl index bcc17b78e..08a844e68 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4184 +1,78 @@ + # using Revise test_set = ENV["TEST_SET"] using Preferences: set_preferences! -set_preferences!("MacroModelling", "dispatch_doctor_mode" => test_set in ["estimate_sw07", "estimation", "1st_order_inversion_estimation", "pruned_2nd_order_estimation", "2nd_order_estimation", "pruned_3rd_order_estimation", "3rd_order_estimation", "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" +set_preferences!("MacroModelling", "dispatch_doctor_mode" => test_set in ["estimate_sw07", "estimate_sw07_nested_sampling", "estimation", "1st_order_inversion_estimation", "pruned_2nd_order_estimation", "2nd_order_estimation", "pruned_3rd_order_estimation", "3rd_order_estimation", "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", "system_prior_estimation", "gradient_checks", "jet", "jet_hot_paths" ] ? "disable" : "error") set_preferences!("MacroModelling", "dispatch_doctor_union_limit" => 4) -using Test -using MacroModelling -import MacroModelling: clear_solution_caches! -using Random -import SpecialFunctions: erfcinv -using AxisKeys, SparseArrays -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 Dates -using RuntimeGeneratedFunctions - -function quarterly_dates(start_date::Date, len::Int) - dates = Vector{Date}(undef, len) - current_date = start_date - for i in 1:len - dates[i] = current_date - current_date = current_date + Dates.Month(3) - end - return dates -end println("Running test set: $test_set") println("Threads used: ", Threads.nthreads()) -include("functionality_tests.jl") - -# @testset verbose = true "Code formatting (JuliaFormatter.jl)" begin -# @test format(MacroModelling; verbose=true, overwrite=true) -# end - -# if test_set == "solver0" -# transform = 0 -# include("optim_solver_params.jl") -# elseif test_set == "solver1" -# transform = 1 -# include("optim_solver_params.jl") -# elseif test_set == "solver2" -# transform = 2 -# include("optim_solver_params.jl") -# elseif test_set == "solver3" -# transform = 3 -# include("optim_solver_params.jl") -# end - if test_set == "jet" - if VERSION < v"1.13" - using JET - end - - @testset verbose = true "Static checking (JET.jl)" begin - if VERSION < v"1.13" - JET.test_package(MacroModelling; target_defined_modules = true, toplevel_logger = nothing) - end - end -end - -if test_set == "estimate_sw07" + include("test_jet.jl") +elseif test_set == "estimate_sw07" include("test_sw07_estimation.jl") -end - -if test_set == "estimation" +elseif test_set == "estimate_sw07_nested_sampling" + include("test_sw07_estimation_nested_sampling.jl") +elseif test_set == "estimation" include("test_estimation.jl") -end - -if test_set == "1st_order_inversion_estimation" +elseif test_set == "1st_order_inversion_estimation" include("test_1st_order_inversion_filter_estimation.jl") -end - -if test_set == "2nd_order_estimation" +elseif test_set == "2nd_order_estimation" include("test_2nd_order_estimation.jl") -end - -if test_set == "pruned_2nd_order_estimation" +elseif test_set == "pruned_2nd_order_estimation" include("test_pruned_2nd_order_estimation.jl") -end - -if test_set == "3rd_order_estimation" +elseif test_set == "3rd_order_estimation" include("test_3rd_order_estimation.jl") -end - -if test_set == "pruned_3rd_order_estimation" +elseif test_set == "pruned_3rd_order_estimation" include("test_pruned_3rd_order_estimation.jl") -end - -if test_set == "estimation_pigeons" +elseif test_set == "estimation_pigeons" include("test_estimation_pigeons.jl") -end - -if test_set == "1st_order_inversion_estimation_pigeons" +elseif test_set == "1st_order_inversion_estimation_pigeons" include("test_1st_order_inversion_filter_estimation_pigeons.jl") -end - -if test_set == "2nd_order_estimation_pigeons" +elseif test_set == "2nd_order_estimation_pigeons" include("test_2nd_order_estimation_pigeons.jl") -end - -if test_set == "pruned_2nd_order_estimation_pigeons" +elseif test_set == "pruned_2nd_order_estimation_pigeons" include("test_pruned_2nd_order_estimation_pigeons.jl") -end - -if test_set == "3rd_order_estimation_pigeons" +elseif test_set == "3rd_order_estimation_pigeons" include("test_3rd_order_estimation_pigeons.jl") -end - -if test_set == "pruned_3rd_order_estimation_pigeons" +elseif test_set == "pruned_3rd_order_estimation_pigeons" include("test_pruned_3rd_order_estimation_pigeons.jl") -end - - -if test_set == "plots_1" - plots = true - Random.seed!(1) - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "Backus_Kehoe_Kydland_1992" begin - include("../models/Backus_Kehoe_Kydland_1992.jl") - functionality_test(Backus_Kehoe_Kydland_1992, Caldara_et_al_2012_estim, plots = plots) - end - Backus_Kehoe_Kydland_1992 = nothing - GC.gc() - - @testset verbose = true "FS2000" begin - include("../models/FS2000.jl") - functionality_test(FS2000, Caldara_et_al_2012_estim, plots = plots) - end - FS2000 = nothing - GC.gc() -end - -if test_set == "plots_2" - plots = true - Random.seed!(1) - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "Smets and Wouters (2007) nonlinear" begin - include("../models/Smets_Wouters_2007.jl") - functionality_test(Smets_Wouters_2007, Caldara_et_al_2012_estim, plots = plots) - end - Smets_Wouters_2007 = nothing - GC.gc() - - @testset verbose = true "Smets_Wouters_2003 with calibration equations" begin - include("../models/Smets_Wouters_2003.jl") - functionality_test(Smets_Wouters_2003, Caldara_et_al_2012_estim, plots = plots) - end - Smets_Wouters_2003 = nothing - GC.gc() - - @testset verbose = true "Smets and Wouters (2007) linear" begin - include("../models/Smets_Wouters_2007_linear.jl") - functionality_test(Smets_Wouters_2007_linear, Caldara_et_al_2012_estim, plots = plots) - end - Smets_Wouters_2007_linear = nothing - GC.gc() -end - -if test_set == "plots_3" - plots = true - Random.seed!(1) - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "Gali 2015 ELB" begin - include("../models/Gali_2015_chapter_3_obc.jl") - functionality_test(Gali_2015_chapter_3_obc, Caldara_et_al_2012_estim, plots = plots) - end - Gali_2015_chapter_3_obc = nothing - GC.gc() -end - -if test_set == "plots_4" - plots = true - Random.seed!(1) - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables" begin - include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") - functionality_test(m, Caldara_et_al_2012_estim, plots = plots) - - observables = [:R, :k] - - Random.seed!(1) - simulated_data = simulate(m) - - get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) - - back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - for i in 1:100 - local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - if isfinite(ℒ.norm(fin_grad)) - println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - break - end - end - - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables numerical SS" begin - include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve.jl") - - functionality_test(m, Caldara_et_al_2012_estim, plots = plots) - - observables = [:R, :k] - - Random.seed!(1) - simulated_data = simulate(m) - - get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) - - back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - for i in 1:100 - local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - if isfinite(ℒ.norm(fin_grad)) - println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - break - end - end - - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations, parameter definitions, and special functions" begin - include("models/RBC_CME_calibration_equations_and_parameter_definitions_and_specfuns.jl") - functionality_test(m, Caldara_et_al_2012_estim, plots = plots) - - observables = [:R, :k] - - Random.seed!(1) - simulated_data = simulate(m) - - get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) - - back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - for i in 1:100 - local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x, verbose = true), m.parameter_values) - if isfinite(ℒ.norm(fin_grad)) - println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - break - end - end - - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations and parameter definitions" begin - include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") - functionality_test(m, Caldara_et_al_2012_estim, plots = plots) - - observables = [:R, :k] - - Random.seed!(1) - simulated_data = simulate(m) - - get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) - - back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - for i in 1:100 - local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - if isfinite(ℒ.norm(fin_grad)) - println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - break - end - end - - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations" begin - include("models/RBC_CME_calibration_equations.jl") - functionality_test(m, Caldara_et_al_2012_estim, plots = plots) - - observables = [:R, :k] - - Random.seed!(1) - simulated_data = simulate(m) - - get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) - - back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - for i in 1:100 - local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - if isfinite(ℒ.norm(fin_grad)) - println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - break - end - end - - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME" begin - include("models/RBC_CME.jl") - functionality_test(m, Caldara_et_al_2012_estim, plots = plots) - - observables = [:R, :k] - - Random.seed!(1) - simulated_data = simulate(m) - - get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) - - back_grad = Zygote.gradient(x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - - for i in 1:100 - local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) - if isfinite(ℒ.norm(fin_grad)) - println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - break - end - end - - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) - end - m = nothing - GC.gc() - -end - - -if test_set == "plots_5" - Random.seed!(1) - - @testset verbose = true "SW07 estim" begin - include("../models/Smets_Wouters_2007.jl") - - # load data - dat = CSV.read("data/usmodel.csv", DataFrame) - - # load data - data = KeyedArray(Array(dat)',Variable = Symbol.(strip.(names(dat))), Time = 1:size(dat)[1]) - - # declare observables as written in csv file - observables_old = [:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model - - # Subsample - # subset observables in data - sample_idx = 47:230 # 1960Q1-2004Q4 - - data = data(observables_old, sample_idx) - - # declare observables as written in model - observables = [:dy, :dc, :dinve, :labobs, :pinfobs, :dwobs, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model - - data = rekey(data, :Variable => observables) - - data_rekey = rekey(data, :Time => quarterly_dates(Date(1960, 1, 1), size(data,2))) - - - plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 3, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 3, :calfa => 0.28]) - - - plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], filter = :inversion) - - - plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], filter = :inversion) - - plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], smooth = false) - - - plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], smooth = false) - - plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], smooth = false, presample_periods = 50) - - - plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data[:,20:end], parameters = [:csadjcost => 6, :calfa => 0.24]) - - - plot_model_estimates(Smets_Wouters_2007, data_rekey, parameters = [:csadjcost => 6, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data_rekey, parameters = [:csadjcost => 5, :calfa => 0.24]) - - - plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) - - plot_model_estimates!(Smets_Wouters_2007, data_rekey, parameters = [:csadjcost => 5, :calfa => 0.24]) - - # FS2000 model and data - 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]) - dataFS2000 = log.(dataFS2000) - - # declare observables - observables = sort(Symbol.("log_".*names(dat))) - - # subset observables in data - dataFS2000 = dataFS2000(observables,:) - - dataFS2000_rekey = rekey(dataFS2000, :Time => quarterly_dates(Date(1950, 1, 1), size(dataFS2000,2))) - - plot_model_estimates(FS2000, dataFS2000) - - plot_model_estimates(FS2000, dataFS2000_rekey[:,1:10]) - - plot_shock_decomposition(FS2000, dataFS2000_rekey[:,1:10]) - - plot_shock_decomposition(FS2000, dataFS2000_rekey) - - - dataFS2000_rekey2 = rekey(dataFS2000, :Time => 1:1:size(dataFS2000,2)) - - plot_shock_decomposition(FS2000, dataFS2000) - - plot_shock_decomposition(FS2000, dataFS2000_rekey2) - - - plot_model_estimates(FS2000, dataFS2000_rekey, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - plot_model_estimates!(Smets_Wouters_2007, data_rekey) - - - plot_model_estimates(FS2000, dataFS2000_rekey, parameters = :alp => 0.356, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - plot_model_estimates!(Smets_Wouters_2007, data_rekey) - - plot_model_estimates!(FS2000, dataFS2000_rekey, parameters = :alp => 0.3, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - - plot_model_estimates!(Smets_Wouters_2007, data_rekey, parameters = :csigma => 0.3) - - plot_model_estimates(FS2000, dataFS2000_rekey, parameters = :alp => 0.356, shock_decomposition = true, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - - estims = get_estimated_variables(Smets_Wouters_2007, data) - - plot_irf(Smets_Wouters_2007, shocks = :em, shock_size = 10) - - plot_irf!(Smets_Wouters_2007,initial_state = collect(estims[:,end]), shocks = :none, plot_type = :stack) - - plot_irf!(Smets_Wouters_2007, shocks = [:em, :ea], negative_shock = true, plot_type = :stack) - - shock_mat = randn(Smets_Wouters_2007.constants.post_model_macro.nExo,3) - - plot_irf!(Smets_Wouters_2007, shocks = shock_mat, plot_type = :stack) - - plot_irf!(Smets_Wouters_2007, shocks = shock_mat, plot_type = :stack) - - - plot_irf(Smets_Wouters_2007, shocks = :em, periods = 5, variables = [:y, :k, :c]) - - plot_irf!(FS2000, shocks = :e_m, periods = 5, plot_type = :stack, shock_size = 10, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w), variables = [:y, :k, :c]) - - - plot_irf(Smets_Wouters_2007, shocks = :em, periods = 5) - - plot_irf!(FS2000, shocks = :e_m, periods = 5, plot_type = :stack, shock_size = 10, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - plot_irf!(FS2000, shocks = [:e_m, :e_a], shock_size = 20, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - plot_irf!(Smets_Wouters_2007, shocks = [:em, :ea], shock_size = 0.5) - - - - cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) - cndtns_lvl[1,8] = 1.4 - - plot_conditional_forecast(Smets_Wouters_2007, cndtns_lvl, initial_state = collect(estims[:,end])) - - - cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,4), Variables = [:pinfobs], Periods = 1:4) - cndtns_lvl[1,4] = 2 - - plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, plot_type = :stack) - - - - cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) - cndtns_lvl[1,8] = 1.45 - - plot_conditional_forecast!(FS2000, cndtns_lvl, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - - cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,4), Variables = [:y], Periods = 1:4) - cndtns_lvl[1,4] = 2.01 - - plot_conditional_forecast!(FS2000, cndtns_lvl, plot_type = :stack, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - # conditons on #3 is nothing which makes sense since it is not showing - - shock_mat = sprandn(Smets_Wouters_2007.constants.post_model_macro.nExo, 10, .1) - - cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,4), Variables = [:pinfobs], Periods = 1:4) - cndtns_lvl[1,4] = 2 - - plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, shocks = shock_mat, plot_type = :stack) - - - - cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) - cndtns_lvl[1,8] = 1.4 - - shock_mat = sprandn(Smets_Wouters_2007.constants.post_model_macro.nExo, 10, .1) - - plot_conditional_forecast(Smets_Wouters_2007, cndtns_lvl, shocks = shock_mat, label = "SW07 w shocks", variables = [:y, :k, :c]) - - plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, variables = [:y,:w]) - - plot_conditional_forecast!(FS2000, cndtns_lvl, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - shock_mat = sprandn(FS2000.constants.post_model_macro.nExo, 10, .1) - - plot_conditional_forecast!(FS2000, cndtns_lvl, shocks = shock_mat, label = :rand_shocks, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) - - - plot_solution(FS2000, :k) - - plot_solution!(FS2000, :k, algorithm = :second_order) - - - plot_solution(Smets_Wouters_2007, :pinf) - - plot_solution!(Smets_Wouters_2007, :pinf, algorithm = :second_order) - - - plot_solution(FS2000, :y) - - plot_solution!(Smets_Wouters_2007, :y, variables = [:y, :k, :c]) - - plot_solution!(Smets_Wouters_2007, :y, algorithm = :second_order, variables = [:y, :k, :c]) - - end - - # multiple models - @testset verbose = true "Gali 2015 ELB plots" begin - include("../models/Gali_2015_chapter_3_obc.jl") - - - Random.seed!(14) - plot_simulation(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0, ignore_obc = true) - - Random.seed!(14) - plot_simulation!(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0) - - Random.seed!(14) - plot_simulation!(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0025) - - - Random.seed!(13) - plot_simulation(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, - # periods = 40, - parameters = :R̄ => 1.0, ignore_obc = true) - - Random.seed!(13) - plot_simulation!(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, - periods = 40, - parameters = :R̄ => 1.0) - - - plot_irf(Gali_2015_chapter_3_obc, parameters = :R̄ => 1.0) - - plot_irf!(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, parameters = :R̄ => 1.0) - - - plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.0) - - plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.5) - - plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 0.5) - - - plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.0) - - plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, generalised_irf = true) - - plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, ignore_obc = true) - - - plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, algorithm = :pruned_second_order) - - plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, algorithm = :pruned_second_order, ignore_obc = true) - - plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, algorithm = :pruned_second_order, ignore_obc = true, generalised_irf = true) - end - - @testset verbose = true "Caldara et al 2012 plots" begin - include("../models/Caldara_et_al_2012.jl") - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :second_order) - - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_second_order, generalised_irf = true, generalised_irf_draws = 1000) - - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order) - - - plot_irf(Caldara_et_al_2012, algorithm = :second_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :third_order) - - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, generalised_irf = true) - - - plot_irf(Caldara_et_al_2012, algorithm = :third_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :third_order, generalised_irf = true) - - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, shock_size = 2) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, shock_size = 3) - - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = :ψ => 0.8) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = :ψ => 1.5) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = :ψ => 2.5) - - - plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.3]) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.25]) - - plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.35]) - end -end - - -if test_set == "higher_order_1" - plots = true - # test_higher_order = true - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "FS2000 third order" begin - include("../models/FS2000.jl") - functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) - end - FS2000 = nothing - GC.gc() - - @testset verbose = true "FS2000 pruned third order" begin - include("../models/FS2000.jl") - functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :pruned_third_order, plots = plots) - end - FS2000 = nothing - GC.gc() - - @testset verbose = true "FS2000 second order" begin - include("../models/FS2000.jl") - functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) - end - FS2000 = nothing - GC.gc() - - @testset verbose = true "FS2000 pruned second order" begin - include("../models/FS2000.jl") - functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :pruned_second_order, plots = plots) - end - FS2000 = nothing - GC.gc() -end - - -if test_set == "higher_order_2" - Random.seed!(18400875) - plots = true - # test_higher_order = true - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables pruned second order" begin - include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :pruned_second_order, plots = plots) - end - # m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables pruned third order" begin - # include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :pruned_third_order, plots = plots) - end - m = nothing - GC.gc() -end - - -if test_set == "higher_order_3" - plots = true - # test_higher_order = true - - include("models/Caldara_et_al_2012_estim.jl") - - @testset verbose = true "RBC_CME with calibration equations second order" begin - include("models/RBC_CME_calibration_equations.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) - end - # m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations third order" begin - # include("models/RBC_CME_calibration_equations.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME second order" begin - include("models/RBC_CME.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) - end - # m = nothing - GC.gc() - - @testset verbose = true "RBC_CME third order" begin - # include("models/RBC_CME.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) - end - m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations and parameter definitions second order" begin - include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) - end - # m = nothing - GC.gc() - - @testset verbose = true "RBC_CME with calibration equations and parameter definitions third order" begin - # include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") - functionality_test(m, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) - end - m = nothing - GC.gc() -end - - -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 - - k_ss = ((1 / beta - 1 + delta) / alpha)^(1 / (alpha - 1)) - q_ss = k_ss^alpha - c_ss = q_ss - delta * k_ss - z_ss = 0.0 - - return [c_ss, k_ss, q_ss, z_ss] - end - - function make_counted_ss() - calls = Ref(0) - - function ss(params) - calls[] += 1 - return rbc_steady_state(params) - end - - return ss, calls - end - - @testset verbose = true "FS2000 precompile = true" begin - @model FS2000_pre precompile = true begin - dA[0] = exp(gam + z_e_a * e_a[x]) - log(m[0]) = (1 - rho) * log(mst) + rho * log(m[-1]) + z_e_m * e_m[x] - - P[0] / (c[1] * P[1] * m[0]) + bet * P[1] * (alp * exp( - alp * (gam + log(e[1]))) * k[0] ^ (alp - 1) * n[1] ^ (1 - alp) + (1 - del) * exp( - (gam + log(e[1])))) / (c[2] * P[2] * m[1])=0 - W[0] = l[0] / n[0] - - (psi / (1 - psi)) * (c[0] * P[0] / (1 - n[0])) + l[0] / n[0] = 0 - R[0] = P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ ( - alp) / W[0] - 1 / (c[0] * P[0]) - bet * P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) / (m[0] * l[0] * c[1] * P[1]) = 0 - c[0] + k[0] = exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) + (1 - del) * exp( - (gam + z_e_a * e_a[x])) * k[-1] - P[0] * c[0] = m[0] - m[0] - 1 + d[0] = l[0] - e[0] = exp(z_e_a * e_a[x]) - y[0] = k[-1] ^ alp * n[0] ^ (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) - gy_obs[0] = dA[0] * y[0] / y[-1] - gp_obs[0] = (P[0] / P[-1]) * m[-1] / dA[0] - log_gy_obs[0] = log(gy_obs[0]) - log_gp_obs[0] = log(gp_obs[0]) - end - - @parameters FS2000_pre silent = true precompile = true begin - alp = 0.356 - bet = 0.993 - gam = 0.0085 - mst = 1.0002 - rho = 0.129 - psi = 0.65 - del = 0.01 - z_e_a = 0.035449 - z_e_m = 0.008862 - end - - - SS_pre = get_SS(FS2000_pre, silent = true) - SS_change_pre = get_SS(FS2000_pre, parameters = :alp => 0.36, silent = true) - solution_pre = get_solution(FS2000_pre, silent = true) - solution_change_pre = get_solution(FS2000_pre, parameters = :alp => 0.35) - standard_deviation_pre = get_standard_deviation(FS2000_pre) - correlation_pre = get_correlation(FS2000_pre) - autocorrelation_pre = get_autocorrelation(FS2000_pre) - variance_decomposition_pre = get_variance_decomposition(FS2000_pre) - conditional_variance_decomposition_pre = get_conditional_variance_decomposition(FS2000_pre) - irf_pre = get_irf(FS2000_pre) - - @model FS2000 begin - dA[0] = exp(gam + z_e_a * e_a[x]) - log(m[0]) = (1 - rho) * log(mst) + rho * log(m[-1]) + z_e_m * e_m[x] - - P[0] / (c[1] * P[1] * m[0]) + bet * P[1] * (alp * exp( - alp * (gam + log(e[1]))) * k[0] ^ (alp - 1) * n[1] ^ (1 - alp) + (1 - del) * exp( - (gam + log(e[1])))) / (c[2] * P[2] * m[1])=0 - W[0] = l[0] / n[0] - - (psi / (1 - psi)) * (c[0] * P[0] / (1 - n[0])) + l[0] / n[0] = 0 - R[0] = P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ ( - alp) / W[0] - 1 / (c[0] * P[0]) - bet * P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) / (m[0] * l[0] * c[1] * P[1]) = 0 - c[0] + k[0] = exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) + (1 - del) * exp( - (gam + z_e_a * e_a[x])) * k[-1] - P[0] * c[0] = m[0] - m[0] - 1 + d[0] = l[0] - e[0] = exp(z_e_a * e_a[x]) - y[0] = k[-1] ^ alp * n[0] ^ (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) - gy_obs[0] = dA[0] * y[0] / y[-1] - gp_obs[0] = (P[0] / P[-1]) * m[-1] / dA[0] - log_gy_obs[0] = log(gy_obs[0]) - log_gp_obs[0] = log(gp_obs[0]) - end - - @parameters FS2000 begin - alp = 0.356 - bet = 0.993 - gam = 0.0085 - mst = 1.0002 - rho = 0.129 - psi = 0.65 - del = 0.01 - z_e_a = 0.035449 - z_e_m = 0.008862 - end - - SS_nopre = get_SS(FS2000, silent = true) - SS_change_nopre = get_SS(FS2000, parameters = :alp => 0.36, silent = true) - solution_nopre = get_solution(FS2000, silent = true) - solution_change_nopre = get_solution(FS2000, parameters = :alp => 0.35) - standard_deviation_nopre = get_standard_deviation(FS2000) - correlation_nopre = get_correlation(FS2000) - autocorrelation_nopre = get_autocorrelation(FS2000) - variance_decomposition_nopre = get_variance_decomposition(FS2000) - conditional_variance_decomposition_nopre = get_conditional_variance_decomposition(FS2000) - irf_nopre = get_irf(FS2000) - - @test isapprox(SS_nopre, SS_pre) - @test isapprox(SS_change_nopre, SS_change_pre) - @test isapprox(solution_nopre, solution_pre) - @test isapprox(solution_change_nopre, solution_change_pre) - @test isapprox(standard_deviation_nopre, standard_deviation_pre) - @test isapprox(correlation_nopre, correlation_pre) - @test isapprox(autocorrelation_nopre, autocorrelation_pre) - @test isapprox(variance_decomposition_nopre, variance_decomposition_pre) - @test isapprox(conditional_variance_decomposition_nopre, conditional_variance_decomposition_pre) - @test isapprox(irf_nopre, irf_pre) - end - - @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)) - c[0] + k[0] = (1 - delta) * k[-1] + q[0] - q[0] = exp(z[0]) * k[-1]^alpha - z[0] = rho * z[-1] + std_z * eps_z[x] - end - - @parameters RBC_switch begin - std_z = 0.01 - rho = 0.2 - delta = 0.02 - alpha = 0.5 - beta = 0.95 - end - - custom_ss, custom_calls = make_counted_ss() - - inplace_calls = Ref(0) - function inplace_ss!(out, params) - inplace_calls[] += 1 - out .= rbc_steady_state(params) - return nothing - end - - # bad_calls = Ref(0) - # function bad_ss(params) - # bad_calls[] += 1 - # return zeros(4) - # end - - custom_calls[] = 0 - _ = get_steady_state(RBC_switch, steady_state_function = custom_ss) - @test custom_calls[] > 0 - - inplace_calls[] = 0 - inplace_result = get_steady_state(RBC_switch, steady_state_function = inplace_ss!) - @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_throws ArgumentError get_steady_state(RBC_switch, steady_state_function = bad_ss) - # @test bad_calls[] > 0 - - calls_before = custom_calls[] - _ = get_steady_state(RBC_switch, steady_state_function = nothing) - @test custom_calls[] == calls_before - @test isnothing(RBC_switch.functions.NSSS_custom) - - MacroModelling.set_custom_steady_state_function!(RBC_switch, custom_ss) - calls_before = custom_calls[] - _ = get_steady_state(RBC_switch) - @test custom_calls[] > calls_before - - MacroModelling.set_custom_steady_state_function!(RBC_switch, nothing) - calls_before = custom_calls[] - _ = get_steady_state(RBC_switch) - @test custom_calls[] == calls_before - end - - @testset "Macro steady state assignment" begin - macro_ss, macro_calls = make_counted_ss() - - @model RBC_macro_switch begin - 1 / c[0] = (beta / c[1]) * (alpha * exp(z[1]) * k[0]^(alpha - 1) + (1 - delta)) - c[0] + k[0] = (1 - delta) * k[-1] + q[0] - q[0] = exp(z[0]) * k[-1]^alpha - z[0] = rho * z[-1] + std_z * eps_z[x] - end - - @parameters RBC_macro_switch steady_state_function = macro_ss begin - std_z = 0.01 - rho = 0.2 - delta = 0.02 - alpha = 0.5 - beta = 0.95 - end - - @test !(RBC_macro_switch.functions.NSSS_solve isa RuntimeGeneratedFunction) - - _ = get_steady_state(RBC_macro_switch) - @test macro_calls[] > 0 - @test !(RBC_macro_switch.functions.NSSS_solve isa RuntimeGeneratedFunction) - - 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 - - calls_before = macro_calls[] - _ = get_steady_state(RBC_macro_switch) - @test macro_calls[] == calls_before - end - - @testset verbose = true "Custom steady state function" begin - # Test custom steady state function with simple RBC model - @model RBC_custom_ss 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_custom_ss begin - std_z = 0.01 - ρ = 0.2 - δ = 0.02 - α = 0.5 - β = 0.95 - end - - # Get default steady state - default_ss = get_steady_state(RBC_custom_ss) - - # Define custom steady state function - # Variables in order: [:c, :k, :q, :z] (alphabetically sorted) - # Parameters in order: [:std_z, :ρ, :δ, :α, :β] (declaration order) - function my_steady_state_rbc(params) - std_z, ρ, δ, α, β = params - - # Analytical steady state for RBC model - k_ss = ((1/β - 1 + δ) / α)^(1/(α - 1)) - q_ss = k_ss^α - c_ss = q_ss - δ * k_ss - z_ss = 0.0 - - return [c_ss, k_ss, q_ss, z_ss] - end - - # Test custom function directly - custom_result = my_steady_state_rbc(RBC_custom_ss.parameter_values) - @test isapprox(custom_result, default_ss(:,:Steady_state)) - - # Get steady state with custom function - custom_ss = get_steady_state(RBC_custom_ss, steady_state_function = my_steady_state_rbc) - - # Compare with default (should be essentially the same) - @test isapprox(default_ss, custom_ss, rtol = 1e-10) - - # Test that model can be solved with custom SS function - irf_custom = get_irf(RBC_custom_ss, levels = true) - - # Steady state should still work after clearing - after_clear_ss = get_steady_state(RBC_custom_ss, steady_state_function = nothing) - @test isnothing(RBC_custom_ss.functions.NSSS_custom) - @test isapprox(default_ss, after_clear_ss, rtol = 1e-10) - - irf_after_clear = get_irf(RBC_custom_ss, levels = true) - @test isapprox(irf_after_clear, irf_custom, rtol = 1e-10) - - # Test with verbose option (internal function still available but not exported) - MacroModelling.set_custom_steady_state_function!(RBC_custom_ss, my_steady_state_rbc) - @test !isnothing(RBC_custom_ss.functions.NSSS_custom) - - - @model RBC_macro_ss 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_macro_ss steady_state_function = my_steady_state_rbc begin - std_z = 0.01 - ρ = 0.2 - δ = 0.02 - α = 0.5 - β = 0.95 - end - - # Verify macro-defined SS function is set - @test isapprox(RBC_macro_ss.functions.NSSS_custom(RBC_macro_ss.parameter_values), default_ss(:,:Steady_state), rtol = 1e-10) - - macro_ss = get_steady_state(RBC_macro_ss) - @test isapprox(default_ss, macro_ss, rtol = 1e-10) - - RBC_custom_ss = nothing - RBC_macro_ss = nothing - RBC_func_arg = nothing - end - - @testset verbose = true "Non-stochastic steady state guess" begin - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true guess = Dict(:alpha => .2, :beta => .99) begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve - R_ss = 1.0035 - # beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = 1.0025 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - - - # cap_share > 0 - # R_ss > 0 - # Pi_ss > 0 - # I_K_ratio > 0 - - # 0 < alpha < 1 - # 0 < beta < 1 - # 0 < delta < 1 - # 0 < Pibar - # 0 <= rhoz < 1 - # phi_pi > 0 - - # 0 < A < 1 - # 0 < k < 50 - # 0 < y < 10 - # 0 < c < 10 - end - - @test RBC_CME.constants.post_parameters_macro.guess == Dict(:alpha => .2, :beta => .99) - - @test get_steady_state(RBC_CME, verbose = true)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0025, 1.0035, 1.2081023824176236, 9.437411552284384, 1.4212969205027686, 1.0] - - RBC_CME = nothing - end - GC.gc() - - include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve.jl") - - global model = m - - @testset verbose = true "Custom steady state function with calibration equations and lead/lags" begin - # Test custom steady state function with RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve model - - # Get default steady state - default_ss = get_steady_state(model) - - function custom_steady_state(p::Vector{Float64}) - # 1. Unpack parameters - cap_share = p[1] - R_ss_target = p[2] - I_K_ratio = p[3] - phi_pi = p[4] - # std_eps = p[5] - # std_z_d = p[6] - Pi_real = p[7] - # rhoz = p[8] - - # 2. Solve for Deep Parameters and Rates - # Target R: log(R) = R_ss - 1 - R = exp(R_ss_target - 1.0) - - # Target Pi - Pi = R_ss_target - Pi_real - - # Euler Equation: 1 = beta * (R / Pi) - beta = Pi / R - Pibar = (R * beta) ^ (-1/phi_pi) * Pi - - # Ratios - # k / (4 * y) = cap_share - ky_ratio = 4.0 * cap_share - - # c / y = 1 - I_K_ratio - cy_ratio = 1.0 - I_K_ratio - - # Resource Constraint: 1 = c/y + delta * k/y - delta = (1.0 - cy_ratio) / ky_ratio - - # Euler Equation: 1 = beta * (alpha * y/k + 1 - delta) - # alpha = (k/y) * (1/beta - 1 + delta) - alpha = ky_ratio * ((1.0 / beta) - 1.0 + delta) - - # 3. Solve for Levels - A = 1.0 - z_delta = 1.0 - - # Production: y/k = k^(alpha-1) => k = (k/y)^(1/(1-alpha)) - k = (ky_ratio)^(1.0 / (1.0 - alpha)) - y = k^alpha - c = cy_ratio * y - - # Auxiliary variables - ZZ_avg = A - ZZ_avg_fut = A - log_ZZ_avg = 0.0 # log(1.0) - - c_logpdf = (-(abs2(c) + 1.8378770664093453) / 2) # normlogpdf - c_invcdf = (-erfcinv(2*(c - 1.0)) * 1.4142135623730951) # norminvcdf - - # 4. Return Vector - return [ - A, # 7 - Pi, # 4 - R, # 3 - ZZ_avg, # 8 - ZZ_avg_fut, # 9 - c, # 2 - c_invcdf, # 12 - c_logpdf, # 11 - k, # 5 - log_ZZ_avg, # 10 - y, # 1 - z_delta, # 6 - beta, # 14 (Derived) - Pibar, # 15 (Derived) - alpha, # 13 (Derived) - delta # 16 (Derived) - ] - end - - # Get steady state with custom function - custom_ss = get_steady_state(model, steady_state_function = custom_steady_state) - - # Compare key variables with default (should be essentially the same) - @test isapprox(default_ss, custom_ss, rtol = 1e-10) - - # Test that model can be solved with custom SS function - std_custom = get_std(model) - - # Steady state should still work after clearing - after_clear_ss = get_steady_state(model, steady_state_function = nothing) - @test isnothing(model.functions.NSSS_custom) - @test isapprox(default_ss, after_clear_ss, rtol = 1e-10) - - std_after_clear = get_std(model) - @test isapprox(std_after_clear, std_custom, rtol = 1e-10) - end - - @testset verbose = true "Provide parameters later" begin - include("models/Backus_Kehoe_Kydland_1992.jl") - - @model Backus_Kehoe_Kydland_1992_incomplete begin - for co in [H, F] - Y{co}[0] = ((LAMBDA{co}[0] * K{co}[-4]^theta{co} * N{co}[0]^(1-theta{co}))^(-nu{co}) + sigma{co} * Z{co}[-1]^(-nu{co}))^(-1/nu{co}) - - K{co}[0] = (1-delta{co})*K{co}[-1] + S{co}[0] - - X{co}[0] = for lag in (-4+1):0 phi{co} * S{co}[lag] end - - A{co}[0] = (1-eta{co}) * A{co}[-1] + N{co}[0] - - L{co}[0] = 1 - alpha{co} * N{co}[0] - (1-alpha{co})*eta{co} * A{co}[-1] - - U{co}[0] = (C{co}[0]^mu{co}*L{co}[0]^(1-mu{co}))^gamma{co} - - psi{co} * mu{co} / C{co}[0]*U{co}[0] = LGM[0] - - psi{co} * (1-mu{co}) / L{co}[0] * U{co}[0] * (-alpha{co}) = - LGM[0] * (1-theta{co}) / N{co}[0] * (LAMBDA{co}[0] * K{co}[-4]^theta{co}*N{co}[0]^(1-theta{co}))^(-nu{co})*Y{co}[0]^(1+nu{co}) - - for lag in 0:(4-1) - beta{co}^lag * LGM[lag]*phi{co} - end + - for lag in 1:4 - -beta{co}^lag * LGM[lag] * phi{co} * (1-delta{co}) - end = beta{co}^4 * LGM[+4] * theta{co} / K{co}[0] * (LAMBDA{co}[+4] * K{co}[0]^theta{co} * N{co}[+4]^(1-theta{co})) ^ (-nu{co})* Y{co}[+4]^(1+nu{co}) - - LGM[0] = beta{co} * LGM[+1] * (1+sigma{co} * Z{co}[0]^(-nu{co}-1)*Y{co}[+1]^(1+nu{co})) - - NX{co}[0] = (Y{co}[0] - (C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1]))/Y{co}[0] - end - - (LAMBDA{H}[0]-1) = rho{H}{H}*(LAMBDA{H}[-1]-1) + rho{H}{F}*(LAMBDA{F}[-1]-1) + Z_E{H} * E{H}[x] - - (LAMBDA{F}[0]-1) = rho{F}{F}*(LAMBDA{F}[-1]-1) + rho{F}{H}*(LAMBDA{H}[-1]-1) + Z_E{F} * E{F}[x] - - for co in [H,F] C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1] end = for co in [H,F] Y{co}[0] end - - dLGM[0] = LGM[1] / LGM[0] - - dLGM_ann[0] = for operator = :*, lag in -3:0 dLGM[lag] end - end - - @parameters Backus_Kehoe_Kydland_1992_incomplete begin - # K_ss = 11.0148 - # K[ss] = K_ss | beta - # K[ss] = 10 | beta - # F_H_ratio = 1 - K{F}[ss] / K{H}[ss] = F_H_ratio | beta{F} - K{H}[ss] = K_ss | beta{H} - - # beta = 0.99 - # mu = 0.34 - # gamma = -1.0 - # alpha = 1 - # eta = 0.5 - # theta = 0.36 - # nu = 3 - # sigma = 0.01 - # delta = 0.025 - # phi = 1/4 - # psi = 0.5 - - # Z_E = 0.00852 - - # rho{H}{H} = 0.906 - rho{F}{F} = rho{H}{H} - # rho{H}{F} = 0.088 - rho{F}{H} = rho{H}{F} - end - - pars = [ - "F_H_ratio" => 1.0, - "K_ss" => 11.0148, - "Z_E{F}" => 0.00852, - "Z_E{H}" => 0.00852, - "alpha{F}" => 1.0, - "alpha{H}" => 1.0, - "delta{F}" => 0.025, - "delta{H}" => 0.025, - "eta{F}" => 0.5, - "eta{H}" => 0.5, - "gamma{F}" => -1.0, - "gamma{H}" => -1.0, - "mu{F}" => 0.34, - "mu{H}" => 0.34, - "nu{F}" => 3.0, - "nu{H}" => 3.0, - "phi{F}" => 0.25, - "phi{H}" => 0.25, - "psi{F}" => 0.5, - "psi{H}" => 0.5, - "rho{H}{F}" => 0.088, - "rho{H}{H}" => 0.906, - "sigma{F}" => 0.01, - "sigma{H}" => 0.01, - "theta{F}" => 0.36, - "theta{H}" => 0.36 - ] - - cov1 = get_cov(Backus_Kehoe_Kydland_1992_incomplete, parameters = pars) - - cov2 = get_cov(Backus_Kehoe_Kydland_1992) - - @test cov1 ≈ cov2 - - - include("../models/Gali_2015_chapter_3_obc.jl") - - @model Gali_2015_chapter_3_obc_incomplete begin - W_real[0] = C[0] ^ σ * N[0] ^ φ - - Q[0] = β * (C[1] / C[0]) ^ (-σ) * Z[1] / Z[0] / Pi[1] - - R[0] = 1 / Q[0] - - Y[0] = A[0] * (N[0] / S[0]) ^ (1 - α) - - R[0] = Pi[1] * realinterest[0] - - R[0] = max(R̄ , 1 / β * Pi[0] ^ ϕᵖⁱ * (Y[0] / Y[ss]) ^ ϕʸ * exp(nu[0])) - - C[0] = Y[0] - - log(A[0]) = ρ_a * log(A[-1]) + std_a * eps_a[x] - - log(Z[0]) = ρ_z * log(Z[-1]) - std_z * eps_z[x] - - nu[0] = ρ_ν * nu[-1] + std_nu * eps_nu[x] - - MC[0] = W_real[0] / (S[0] * Y[0] * (1 - α) / N[0]) - - 1 = θ * Pi[0] ^ (ϵ - 1) + (1 - θ) * Pi_star[0] ^ (1 - ϵ) - - S[0] = (1 - θ) * Pi_star[0] ^ (( - ϵ) / (1 - α)) + θ * Pi[0] ^ (ϵ / (1 - α)) * S[-1] - - Pi_star[0] ^ (1 + ϵ * α / (1 - α)) = ϵ * x_aux_1[0] / x_aux_2[0] * (1 - τ) / (ϵ - 1) - - x_aux_1[0] = MC[0] * Y[0] * Z[0] * C[0] ^ (-σ) + β * θ * Pi[1] ^ (ϵ + α * ϵ / (1 - α)) * x_aux_1[1] - - x_aux_2[0] = Y[0] * Z[0] * C[0] ^ (-σ) + β * θ * Pi[1] ^ (ϵ - 1) * x_aux_2[1] - - log_y[0] = log(Y[0]) - - log_W_real[0] = log(W_real[0]) - - log_N[0] = log(N[0]) - - pi_ann[0] = 4 * log(Pi[0]) - - i_ann[0] = 4 * log(R[0]) - - r_real_ann[0] = 4 * log(realinterest[0]) - - M_real[0] = Y[0] / R[0] ^ η - - end - - @parameters Gali_2015_chapter_3_obc_incomplete begin - σ = 1 - - φ = 5 - - ϕᵖⁱ = 1.5 - - ϕʸ = 0.125 - - θ = 0.75 - - ρ_ν = 0.5 - - ρ_z = 0.5 - - ρ_a = 0.9 - - β = 0.99 - - η = 3.77 - - α = 0.25 - - ϵ = 9 - - τ = 0 - - std_a = .01 - - std_z = .05 - - std_nu = .0025 - - R > 1.0001 - end - - cov1 = get_cov(Gali_2015_chapter_3_obc_incomplete, parameters = :R̄ => 1.0) - - cov2 = get_cov(Gali_2015_chapter_3_obc) - - @test cov1 ≈ cov2 - - - @model Backus_Kehoe_Kydland_1992_incomplete2 begin - for co in [H, F] - Y{co}[0] = ((LAMBDA{co}[0] * K{co}[-4]^theta{co} * N{co}[0]^(1-theta{co}))^(-nu{co}) + sigma{co} * Z{co}[-1]^(-nu{co}))^(-1/nu{co}) - - K{co}[0] = (1-delta{co})*K{co}[-1] + S{co}[0] - - X{co}[0] = for lag in (-4+1):0 phi{co} * S{co}[lag] end - - A{co}[0] = (1-eta{co}) * A{co}[-1] + N{co}[0] - - L{co}[0] = 1 - alpha{co} * N{co}[0] - (1-alpha{co})*eta{co} * A{co}[-1] - - U{co}[0] = (C{co}[0]^mu{co}*L{co}[0]^(1-mu{co}))^gamma{co} - - psi{co} * mu{co} / C{co}[0]*U{co}[0] = LGM[0] - - psi{co} * (1-mu{co}) / L{co}[0] * U{co}[0] * (-alpha{co}) = - LGM[0] * (1-theta{co}) / N{co}[0] * (LAMBDA{co}[0] * K{co}[-4]^theta{co}*N{co}[0]^(1-theta{co}))^(-nu{co})*Y{co}[0]^(1+nu{co}) - - for lag in 0:(4-1) - beta{co}^lag * LGM[lag]*phi{co} - end + - for lag in 1:4 - -beta{co}^lag * LGM[lag] * phi{co} * (1-delta{co}) - end = beta{co}^4 * LGM[+4] * theta{co} / K{co}[0] * (LAMBDA{co}[+4] * K{co}[0]^theta{co} * N{co}[+4]^(1-theta{co})) ^ (-nu{co})* Y{co}[+4]^(1+nu{co}) - - LGM[0] = beta{co} * LGM[+1] * (1+sigma{co} * Z{co}[0]^(-nu{co}-1)*Y{co}[+1]^(1+nu{co})) - - NX{co}[0] = (Y{co}[0] - (C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1]))/Y{co}[0] - end - - (LAMBDA{H}[0]-1) = rho{H}{H}*(LAMBDA{H}[-1]-1) + rho{H}{F}*(LAMBDA{F}[-1]-1) + Z_E{H} * E{H}[x] - - (LAMBDA{F}[0]-1) = rho{F}{F}*(LAMBDA{F}[-1]-1) + rho{F}{H}*(LAMBDA{H}[-1]-1) + Z_E{F} * E{F}[x] - - for co in [H,F] C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1] end = for co in [H,F] Y{co}[0] end - - dLGM[0] = LGM[1] / LGM[0] - - dLGM_ann[0] = for operator = :*, lag in -3:0 dLGM[lag] end - end - - @parameters Backus_Kehoe_Kydland_1992_incomplete2 begin - K_ss = 11.0148 - kk = 1 - K{F}[ss] / K{H}[ss] = F_H_ratio | beta{F} - K{H}[ss] = K_ss | beta{H} - rho{F}{F} = rho{H}{H} - rho{F}{H} = rho{H}{F} - end - - params_full = get_parameters(Backus_Kehoe_Kydland_1992, values = true) - - std1 = get_std(Backus_Kehoe_Kydland_1992_incomplete2, parameters = params_full) - - std2 = get_std(Backus_Kehoe_Kydland_1992) - - @test std1 ≈ std2 - - - # Define a model with missing parameters - @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 - - # Only define some parameters - @parameters RBC begin - std_z = 0.01 - ρ = 0.2 - α = 0.5 - β = 0.95 - δ = 0.02 - end - - # Define a model with missing parameters - @model RBC_missing_provide 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 - - # Only define some parameters - @parameters RBC_missing_provide begin - std_z = 0.01 - ρ = 0.2 - # α => 0.5 - # β => 0.95 - # δ => 0.02 - end - - # Verify parameters are missing before providing them - @test has_missing_parameters(RBC_missing_provide) - - # Provide missing parameters and get IRF - irf_result = get_irf(RBC_missing_provide, parameters = [:α => 0.5, :β => 0.95, :δ => 0.02]) - - # Provide missing parameters and get IRF - irf_result_no_missing = get_irf(RBC) - - @test irf_result_no_missing ≈ irf_result - - # After providing parameters, they should no longer be missing - @test !has_missing_parameters(RBC_missing_provide) - - - @model m_incomplete begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-2] + eps_z[x+2] + eps_z_s[x]) - ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[ss]) / 4 - ZZ_avg_fut[0] = (A[0] + A[1] + A[2] + A[ss]) / 4 - log_ZZ_avg[0] = log(ZZ_avg[0]/ZZ_avg[ss]) - c_normlogpdf[0]= normlogpdf(c[0]) - c_norminvcdf[0]= norminvcdf(c[0]-1) - end - - - @parameters m_incomplete verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss - # R_ss = 1.0035 - # beta = .999 - - delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = R_ss - Pi_real - # Pi_real = 1/1000 - # Pibar = 1.0008 - - # phi_pi = 1.5 - # rhoz = 9 / 10 - # std_eps = .0068 - rho_z_delta = rhoz - # std_z_delta = .005 - end - - - param_defs = [ - :cap_share => 1.66, - :alpha => .157, - :std_z_delta => .005, - :phi_pi => 1.5, - :rhoz => 9 / 10, - :std_eps => .0068, - :Pi_real => 1/1000, - :Pibar => 1.0008, - :I_K_ratio => .15, - :delta => .0226, - :R_ss => 1.0035, - :beta => .999 - ] - - irfs_m_incomplete = get_irf(m_incomplete, parameters = param_defs) - - std_m_incomplete = get_std(m_incomplete) - - - - @model m begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-2] + eps_z[x+2] + eps_z_s[x]) - ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[ss]) / 4 - ZZ_avg_fut[0] = (A[0] + A[1] + A[2] + A[ss]) / 4 - log_ZZ_avg[0] = log(ZZ_avg[0]/ZZ_avg[ss]) - c_normlogpdf[0]= normlogpdf(c[0]) - c_norminvcdf[0]= norminvcdf(c[0]-1) - end - - @parameters m verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss - R_ss = 1.0035 - # beta = .999 - - delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = R_ss - Pi_real - Pi_real = 1/1000 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = 9 / 10 - std_eps = .0068 - rho_z_delta = rhoz - std_z_delta = .005 - end - - irfs_m = get_irf(m) - - std_m = get_std(m) - - - @test irfs_m ≈ irfs_m_incomplete - - @test std_m(:,:Standard_deviation) ≈ std_m_incomplete(:,:Standard_deviation) - - - - include("../models/NAWM_EAUS_2008.jl") - - - @model NAWM_EAUS_2008_incomplete begin - EA_R[0] ^ 4 - 1 = EA_PHIRR * (EA_R[-1] ^ 4 - 1) + (1 - EA_PHIRR) * (EA_RRSTAR ^ 4 * EA_PI4TARGET - 1 + EA_PHIRPI * (EA_PIC4[0] - EA_PI4TARGET)) + EA_PHIRGY * (EA_Y[0] / EA_Y[-1] - 1) + σ_EA_R * EA_EPSR[x] - - US_R[0] ^ 4 - 1 = US_PHIRR * (US_R[-1] ^ 4 - 1) + (1 - US_PHIRR) * (US_RRSTAR ^ 4 * US_PI4TARGET - 1 + US_PHIRPI * (US_PIC4[0] - US_PI4TARGET)) + US_PHIRGY * (US_Y[0] / US_Y[-1] - 1) + σ_US_R * US_EPSR[x] - - EA_UTILI[0] = 1 / (1 - EA_SIGMA) * (EA_CI[0] - EA_KAPPA * EA_CI[-1]) ^ (1 - EA_SIGMA) - 1 / (1 + EA_ZETA) * EA_NI[0] ^ (1 + EA_ZETA) + EA_BETA * EA_UTILI[1] - - EA_LAMBDAI[0] * (1 + EA_TAUC[0] + EA_GAMMAVI[0] + EA_VI[0] * EA_GAMMAVIDER[0]) = (EA_CI[0] - EA_KAPPA * EA_CI[-1]) ^ (-EA_SIGMA) - - EA_R[0] = EA_LAMBDAI[0] * EA_BETA ^ (-1) / EA_LAMBDAI[1] * EA_PIC[1] - - EA_GAMMAVIDER[0] * EA_VI[0] ^ 2 = 1 - EA_BETA * EA_LAMBDAI[1] / (EA_LAMBDAI[0] * EA_PIC[1]) - - EA_VI[0] = EA_CI[0] * (1 + EA_TAUC[0]) / EA_MI[0] - - EA_GAMMAVI[0] = EA_VI[0] * EA_GAMMAV1 + EA_GAMMAV2 / EA_VI[0] - 2 * (EA_GAMMAV1 * EA_GAMMAV2) ^ 0.5 - - EA_GAMMAVIDER[0] = EA_GAMMAV1 - EA_GAMMAV2 * EA_VI[0] ^ (-2) - - EA_KI[0] = (1 - EA_DELTA) * EA_KI[-1] + (1 - EA_GAMMAI[-1]) * EA_II[-1] - - EA_GAMMAI[0] = EA_GAMMAI1 / 2 * (EA_II[0] / EA_II[-1] - 1) ^ 2 - - EA_GAMMAIDER[0] = EA_GAMMAI1 * (EA_II[0] / EA_II[-1] - 1) / EA_II[-1] - - EA_GAMMAU[0] = ((EA_DELTA + EA_BETA ^ (-1) - 1) * EA_QBAR - EA_DELTA * EA_TAUKBAR * EA_PIBAR) / (EA_PIBAR * (1 - EA_TAUKBAR)) * (EA_U[0] - 1) + EA_GAMMAU2 / 2 * (EA_U[0] - 1) ^ 2 - - EA_GAMMAUDER[0] = ((EA_DELTA + EA_BETA ^ (-1) - 1) * EA_QBAR - EA_DELTA * EA_TAUKBAR * EA_PIBAR) / (EA_PIBAR * (1 - EA_TAUKBAR)) + (EA_U[0] - 1) * EA_GAMMAU2 - - EA_RK[0] = EA_GAMMAUDER[0] * EA_PI[0] - - EA_PI[0] = EA_Q[0] * (1 - EA_GAMMAI[0] - EA_II[0] * EA_GAMMAIDER[0]) + EA_BETA * EA_LAMBDAI[1] / EA_LAMBDAI[0] * EA_Q[1] * EA_GAMMAIDER[1] * EA_II[1] ^ 2 / EA_II[0] - - EA_Q[0] = EA_BETA * EA_LAMBDAI[1] / EA_LAMBDAI[0] * ((1 - EA_TAUK[1]) * (EA_RK[1] * EA_U[1] - EA_GAMMAU[1] * EA_PI[1]) + EA_PI[1] * EA_DELTA * EA_TAUK[1] + (1 - EA_DELTA) * EA_Q[1]) - - EA_WITILDE[0] ^ (1 + EA_ZETA * EA_ETAI) = EA_ETAI / (EA_ETAI - 1) * EA_FI[0] / EA_GI[0] - - EA_FI[0] = EA_WI[0] ^ ((1 + EA_ZETA) * EA_ETAI) * EA_NDI[0] ^ (1 + EA_ZETA) + EA_BETA * EA_XII * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHII * EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)))) ^ ((1 + EA_ZETA) * EA_ETAI) * EA_FI[1] - - EA_GI[0] = EA_NDI[0] * EA_LAMBDAI[0] * (1 - EA_TAUN[0] - EA_TAUWH[0]) * EA_WI[0] ^ EA_ETAI + EA_BETA * EA_XII * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHII * EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)))) ^ (EA_ETAI - 1) * EA_GI[1] - - EA_WI[0] ^ (1 - EA_ETAI) = (1 - EA_XII) * EA_WITILDE[0] ^ (1 - EA_ETAI) + EA_XII * EA_WI[-1] ^ (1 - EA_ETAI) * (EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)) * EA_PIC[-1] ^ EA_CHII / EA_PIC[0]) ^ (1 - EA_ETAI) - - EA_UTILJ[0] = 1 / (1 - EA_SIGMA) * (EA_CJ[0] - EA_KAPPA * EA_CJ[-1]) ^ (1 - EA_SIGMA) - 1 / (1 + EA_ZETA) * EA_NJ[0] ^ (1 + EA_ZETA) + EA_BETA * EA_UTILJ[1] - - EA_CJ[0] * (1 + EA_TAUC[0] + EA_GAMMAVJ[0]) + EA_MJ[0] = EA_NJ[0] * (1 - EA_TAUN[0] - EA_TAUWH[0]) * EA_WJ[0] + EA_TRJ[0] - EA_TJ[0] + EA_MJ[-1] * EA_PIC[0] ^ (-1) - - EA_LAMBDAJ[0] * (1 + EA_TAUC[0] + EA_GAMMAVJ[0] + EA_VJ[0] * EA_GAMMAVJDER[0]) = (EA_CJ[0] - EA_KAPPA * EA_CJ[-1]) ^ (-EA_SIGMA) - - EA_GAMMAVJDER[0] * EA_VJ[0] ^ 2 = 1 - EA_BETA * EA_LAMBDAJ[1] / (EA_PIC[1] * EA_LAMBDAJ[0]) - - EA_VJ[0] = (1 + EA_TAUC[0]) * EA_CJ[0] / EA_MJ[0] - - EA_GAMMAVJ[0] = EA_GAMMAV1 * EA_VJ[0] + EA_GAMMAV2 / EA_VJ[0] - 2 * (EA_GAMMAV1 * EA_GAMMAV2) ^ 0.5 - - EA_GAMMAVJDER[0] = EA_GAMMAV1 - EA_GAMMAV2 * EA_VJ[0] ^ (-2) - - EA_WJTILDE[0] ^ (1 + EA_ZETA * EA_ETAJ) = EA_ETAJ / (EA_ETAJ - 1) * EA_FJ[0] / EA_GJ[0] - - EA_FJ[0] = EA_WJ[0] ^ ((1 + EA_ZETA) * EA_ETAJ) * EA_NDJ[0] ^ (1 + EA_ZETA) + EA_BETA * EA_XIJ * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHIJ * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)))) ^ ((1 + EA_ZETA) * EA_ETAJ) * EA_FJ[1] - - EA_GJ[0] = EA_NDJ[0] * (1 - EA_TAUN[0] - EA_TAUWH[0]) * EA_LAMBDAJ[0] * EA_WJ[0] ^ EA_ETAJ + EA_BETA * EA_XIJ * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHIJ * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)))) ^ (EA_ETAJ - 1) * EA_GJ[1] - - EA_WJ[0] ^ (1 - EA_ETAJ) = (1 - EA_XIJ) * EA_WJTILDE[0] ^ (1 - EA_ETAJ) + EA_XIJ * EA_WJ[-1] ^ (1 - EA_ETAJ) * (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)) * EA_PIC[-1] ^ EA_CHIJ / EA_PIC[0]) ^ (1 - EA_ETAJ) - - EA_YS[0] = EA_Z[0] * EA_KD[0] ^ EA_ALPHA * EA_ND[0] ^ (1 - EA_ALPHA) - EA_PSIBAR - - EA_RK[0] = EA_ALPHA * (EA_YS[0] + EA_PSIBAR) / EA_KD[0] * EA_MC[0] - - EA_MC[0] = 1 / (EA_Z[0] * EA_ALPHA ^ EA_ALPHA * (1 - EA_ALPHA) ^ (1 - EA_ALPHA)) * EA_RK[0] ^ EA_ALPHA * ((1 + EA_TAUWF[0]) * EA_W[0]) ^ (1 - EA_ALPHA) - - EA_NDI[0] = EA_ND[0] * (1 - EA_OMEGA) * (EA_WI[0] / EA_W[0]) ^ (-EA_ETA) - - EA_NDJ[0] = EA_ND[0] * EA_OMEGA * (EA_WJ[0] / EA_W[0]) ^ (-EA_ETA) - - EA_ND[0] ^ (1 - 1 / EA_ETA) = (1 - EA_OMEGA) ^ (1 / EA_ETA) * EA_NDI[0] ^ (1 - 1 / EA_ETA) + EA_OMEGA ^ (1 / EA_ETA) * EA_NDJ[0] ^ (1 - 1 / EA_ETA) - - EA_D[0] = EA_Y[0] * EA_PY[0] - EA_RK[0] * EA_KD[0] - EA_ND[0] * (1 + EA_TAUWF[0]) * EA_W[0] - - EA_PHTILDE[0] / EA_PH[0] = EA_THETA / (EA_THETA - 1) * EA_FH[0] / EA_GH[0] - - EA_FH[0] = EA_MC[0] * EA_H[0] + EA_LAMBDAI[1] * EA_BETA * EA_XIH / EA_LAMBDAI[0] * (EA_PIH[1] / (EA_PIH[0] ^ EA_CHIH * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)))) ^ EA_THETA * EA_FH[1] - - EA_GH[0] = EA_PH[0] * EA_H[0] + EA_LAMBDAI[1] * EA_BETA * EA_XIH / EA_LAMBDAI[0] * (EA_PIH[1] / (EA_PIH[0] ^ EA_CHIH * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)))) ^ (EA_THETA - 1) * EA_GH[1] - - EA_PH[0] ^ (1 - EA_THETA) = (1 - EA_XIH) * EA_PHTILDE[0] ^ (1 - EA_THETA) + EA_XIH * (EA_PH[-1] / EA_PIC[0]) ^ (1 - EA_THETA) * (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)) * EA_PIH[-1] ^ EA_CHIH) ^ (1 - EA_THETA) - - EA_PIH[0] = EA_PIC[0] * EA_PH[0] / EA_PH[-1] - - US_PIMTILDE[0] / US_PIM[0] = EA_THETA / (EA_THETA - 1) * EA_FX[0] / EA_GX[0] - - EA_FX[0] = EA_MC[0] * US_SIZE / EA_SIZE * US_IM[0] + EA_LAMBDAI[1] * EA_BETA * EA_XIX / EA_LAMBDAI[0] * (US_PIIM[1] / (US_PIIM[0] ^ EA_CHIX * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIX)))) ^ EA_THETA * EA_FX[1] - - EA_GX[0] = US_IM[0] * US_SIZE * US_PIM[0] * EAUS_RER[0] / EA_SIZE + EA_LAMBDAI[1] * EA_BETA * EA_XIX / EA_LAMBDAI[0] * (US_PIIM[1] / (US_PIIM[0] ^ EA_CHIX * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIX)))) ^ (EA_THETA - 1) * EA_GX[1] - - US_PIM[0] ^ (1 - EA_THETA) = (1 - EA_XIX) * US_PIMTILDE[0] ^ (1 - EA_THETA) + EA_XIX * (US_PIM[-1] / US_PIC[0]) ^ (1 - EA_THETA) * (US_PIIM[-1] ^ EA_CHIX * US_PI4TARGET ^ (0.25 * (1 - EA_CHIH))) ^ (1 - EA_THETA) - - US_PIIM[0] = US_PIC[0] * US_PIM[0] / US_PIM[-1] - - EAUS_RER[0] = EA_RER[0] / US_RER - - EA_QC[0] ^ ((EA_MUC - 1) / EA_MUC) = EA_NUC ^ (1 / EA_MUC) * EA_HC[0] ^ (1 - 1 / EA_MUC) + (1 - EA_NUC) ^ (1 / EA_MUC) * ((1 - EA_GAMMAIMC[0]) * EA_IMC[0]) ^ (1 - 1 / EA_MUC) - - 1 = EA_NUC * EA_PH[0] ^ (1 - EA_MUC) + (1 - EA_NUC) * (EA_PIM[0] / EA_GAMMAIMCDAG[0]) ^ (1 - EA_MUC) - - EA_HC[0] = EA_QC[0] * EA_NUC * EA_PH[0] ^ (-EA_MUC) - - EA_GAMMAIMC[0] = EA_GAMMAIMC1 / 2 * (EA_IMC[0] / EA_QC[0] / (EA_IMC[-1] / EA_QC[-1]) - 1) ^ 2 - - EA_GAMMAIMCDAG[0] = 1 - EA_GAMMAIMC[0] - EA_IMC[0] * EA_GAMMAIMC1 * (EA_IMC[0] / EA_QC[0] / (EA_IMC[-1] / EA_QC[-1]) - 1) / EA_QC[0] / (EA_IMC[-1] / EA_QC[-1]) - - EA_QI[0] ^ ((EA_MUI - 1) / EA_MUI) = EA_NUI ^ (1 / EA_MUI) * EA_HI[0] ^ (1 - 1 / EA_MUI) + (1 - EA_NUI) ^ (1 / EA_MUI) * ((1 - EA_GAMMAIMI[0]) * EA_IMI[0]) ^ (1 - 1 / EA_MUI) - - EA_PI[0] ^ (1 - EA_MUI) = EA_NUI * EA_PH[0] ^ (1 - EA_MUI) + (1 - EA_NUI) * (EA_PIM[0] / EA_GAMMAIMIDAG[0]) ^ (1 - EA_MUI) - - EA_HI[0] = EA_QI[0] * EA_NUI * (EA_PH[0] / EA_PI[0]) ^ (-EA_MUI) - - EA_GAMMAIMI[0] = EA_GAMMAIMI1 / 2 * (EA_IMI[0] / EA_QI[0] / (EA_IMI[-1] / EA_QI[-1]) - 1) ^ 2 - - EA_GAMMAIMIDAG[0] = 1 - EA_GAMMAIMI[0] - EA_IMI[0] * EA_GAMMAIMI1 * (EA_IMI[0] / EA_QI[0] / (EA_IMI[-1] / EA_QI[0]) - 1) / EA_QI[0] / (EA_IMI[-1] / EA_QI[-1]) - - EA_PH[-1] * EA_G[-1] + EA_TR[-1] + EA_B[-1] * EA_PIC[-1] ^ (-1) + EA_PIC[-1] ^ (-1) * EA_M[-2] = EA_TAUC[-1] * EA_C[-1] + (EA_TAUN[-1] + EA_TAUWH[-1]) * (EA_WI[-1] * EA_NDI[-1] + EA_WJ[-1] * EA_NDJ[-1]) + EA_TAUWF[-1] * EA_W[-1] * EA_ND[-1] + EA_TAUK[-1] * (EA_RK[-1] * EA_U[-1] - (EA_DELTA + EA_GAMMAU[-1]) * EA_PI[-1]) * EA_K[-1] + EA_TAUD[-1] * EA_D[-1] + EA_T[-1] + EA_R[-1] ^ (-1) * EA_B[0] + EA_M[-1] - - EA_PH[0] * EA_G[0] = EA_GY[0] * EA_PYBAR * EA_YBAR - - EA_TR[0] = EA_YBAR * EA_PYBAR * EA_TRY[0] - - EA_T[0] / (EA_PYBAR * EA_YBAR) = EA_PHITB * (EA_B[0] / (EA_PYBAR * EA_YBAR) - EA_BYTARGET) - - EA_TI[0] = EA_T[0] * EA_UPSILONT - - EA_TRI[0] = EA_TR[0] * EA_UPSILONTR - - EA_PIC4[0] = EA_PIC[0] * EA_PIC[-1] * EA_PIC[-2] * EA_PIC[-3] - - EA_RR[0] - 1 = EA_R[0] / EA_PIC[1] - 1 - - EA_C[0] = EA_CI[0] * (1 - EA_OMEGA) + EA_CJ[0] * EA_OMEGA - - EA_M[0] = EA_MI[0] * (1 - EA_OMEGA) + EA_MJ[0] * EA_OMEGA - - EA_K[0] = EA_KI[0] * (1 - EA_OMEGA) - - EA_I[0] = EA_II[0] * (1 - EA_OMEGA) - - EA_TRJ[0] = EA_TR[0] * 1 / EA_OMEGA - EA_TRI[0] * (1 - EA_OMEGA) / EA_OMEGA - - EA_TJ[0] = EA_T[0] * 1 / EA_OMEGA - EA_TI[0] * (1 - EA_OMEGA) / EA_OMEGA - - EA_GAMMAV[0] = EA_GAMMAVI[0] * EA_CI[0] * (1 - EA_OMEGA) + EA_GAMMAVJ[0] * EA_CJ[0] * EA_OMEGA - - EA_NI[0] = EA_NDI[0] * EA_SI[0] - - EA_SI[0] = (1 - EA_XII) * (EA_WITILDE[0] / EA_WI[0]) ^ (-EA_ETAI) + EA_XII * (EA_WI[-1] / EA_WI[0]) ^ (-EA_ETAI) * (EA_PIC[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)) * EA_PIC[-1] ^ EA_CHII)) ^ EA_ETAI * EA_SI[-1] - - EA_NJ[0] = EA_NDJ[0] * EA_SJ[0] - - EA_SJ[0] = (1 - EA_XIJ) * (EA_WJTILDE[0] / EA_WJ[0]) ^ (-EA_ETAJ) + EA_XIJ * (EA_WJ[-1] / EA_WJ[0]) ^ (-EA_ETAJ) * (EA_PIC[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)) * EA_PIC[-1] ^ EA_CHIJ)) ^ EA_ETAJ * EA_SJ[-1] - - EA_U[0] * EA_K[0] = EA_KD[0] - - EA_YS[0] = EA_H[0] * EA_SH[0] + US_IM[0] * US_SIZE * EA_SX[0] / EA_SIZE - - EA_H[0] = EA_G[0] + EA_HC[0] + EA_HI[0] - - EA_IM[0] = EA_IMC[0] + EA_IMI[0] - - EA_SH[0] = (1 - EA_XIH) * (EA_PHTILDE[0] / EA_PH[0]) ^ (-EA_THETA) + EA_XIH * (EA_PIH[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)) * EA_PIH[-1] ^ EA_CHIH)) ^ EA_THETA * EA_SH[-1] - - EA_SX[0] = (1 - EA_XIX) * (US_PIMTILDE[0] / US_PIM[0]) ^ (-EA_THETA) + EA_XIX * (US_PIIM[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)) * US_PIIM[-1] ^ EA_CHIX)) ^ EA_THETA * EA_SX[-1] - - EA_QC[0] = EA_C[0] + EA_GAMMAV[0] - - EA_QI[0] = EA_I[0] + EA_GAMMAU[0] * EA_K[0] - - EA_Y[0] * EA_PY[0] = US_IM[0] * US_SIZE * US_PIM[0] * EAUS_RER[0] / EA_SIZE + EA_PH[0] * EA_G[0] + EA_QC[0] + EA_PI[0] * EA_QI[0] - EA_PIM[0] * ((1 - EA_GAMMAIMC[0]) * EA_IMC[0] / EA_GAMMAIMCDAG[0] + (1 - EA_GAMMAIMI[0]) * EA_IMI[0] / EA_GAMMAIMIDAG[0]) - - EA_Y[0] = EA_YS[0] - - log(EA_Z[0]) = (1 - EA_RHOZ) * log(EA_ZBAR) + EA_RHOZ * log(EA_Z[-1]) + σ_EA_Z * EA_EPSZ[x] - - EA_GY[0] = (1 - EA_RHOG) * EA_GYBAR + EA_RHOG * EA_GY[-1] + σ_EA_G * EA_EPSG[x] - - EA_TRY[0] = (1 - EA_RHOTR) * EA_TRYBAR + EA_RHOTR * EA_TRY[-1] + σ_EA_TR * EA_EPSTR[x] - - EA_TAUC[0] = (1 - EA_RHOTAUC) * EA_TAUCBAR + EA_TAUC[-1] * EA_RHOTAUC + σ_EA_TAUC * EA_EPSTAUC[x] - - EA_TAUD[0] = (1 - EA_RHOTAUD) * EA_TAUDBAR + EA_TAUD[-1] * EA_RHOTAUD + σ_EA_TAUD * EA_EPSTAUD[x] - - EA_TAUK[0] = EA_TAUKBAR * (1 - EA_RHOTAUK) + EA_TAUK[-1] * EA_RHOTAUK + σ_EA_TAUK * EA_EPSTAUK[x] - - EA_TAUN[0] = (1 - EA_RHOTAUN) * EA_TAUNBAR + EA_TAUN[-1] * EA_RHOTAUN + σ_EA_TAUN * EA_EPSTAUN[x] - - EA_TAUWH[0] = (1 - EA_RHOTAUWH) * EA_TAUWHBAR + EA_TAUWH[-1] * EA_RHOTAUWH + σ_EA_TAUWH * EA_EPSTAUWH[x] - - EA_TAUWF[0] = (1 - EA_RHOTAUWF) * EA_TAUWFBAR + EA_TAUWF[-1] * EA_RHOTAUWF + σ_EA_TAUWF * EA_EPSTAUWF[x] - - EA_CY[0] = EA_C[0] / (EA_Y[0] * EA_PY[0]) - - EA_IY[0] = EA_PI[0] * EA_I[0] / (EA_Y[0] * EA_PY[0]) - - EA_IMY[0] = EA_PIM[0] * EA_IM[0] / (EA_Y[0] * EA_PY[0]) - - EA_IMCY[0] = EA_IMC[0] * EA_PIM[0] / (EA_Y[0] * EA_PY[0]) - - EA_IMIY[0] = EA_PIM[0] * EA_IMI[0] / (EA_Y[0] * EA_PY[0]) - - EA_BY[0] = EA_B[0] / (EA_PYBAR * EA_YBAR) - - EA_TY[0] = EA_T[0] / (EA_PYBAR * EA_YBAR) - - EA_YGAP[0] = EA_Y[0] / EA_YBAR - 1 - - EA_YGROWTH[0] = EA_Y[0] / EA_Y[-1] - - EA_YSHARE[0] = EA_Y[0] * EA_PY[0] * EA_SIZE / EA_RER[0] / (EA_Y[0] * EA_PY[0] * EA_SIZE / EA_RER[0] + US_Y[0] * US_SIZE * US_PY[0] / US_RER) - - EA_EPSILONM[0] = ( - 0.125) / (EA_R[0] * (EA_R[0] + EA_R[0] * EA_GAMMAV2 - 1)) - - US_UTILI[0] = 1 / (1 - US_SIGMA) * (US_CI[0] - US_KAPPA * US_CI[-1]) ^ (1 - US_SIGMA) - 1 / (1 + US_ZETA) * US_NI[0] ^ (1 + US_ZETA) + US_BETA * US_UTILI[1] - - US_LAMBDAI[0] * (1 + US_TAUC[0] + US_GAMMAVI[0] + US_VI[0] * US_GAMMAVIDER[0]) = (US_CI[0] - US_KAPPA * US_CI[-1]) ^ (-US_SIGMA) - - US_R[0] = US_LAMBDAI[0] * US_BETA ^ (-1) / US_LAMBDAI[1] * US_PIC[1] - - US_GAMMAVIDER[0] * US_VI[0] ^ 2 = 1 - US_BETA * US_LAMBDAI[1] / (US_LAMBDAI[0] * US_PIC[1]) - - US_VI[0] = US_CI[0] * (1 + US_TAUC[0]) / US_MI[0] - - US_GAMMAVI[0] = US_VI[0] * US_GAMMAV1 + US_GAMMAV2 / US_VI[0] - 2 * (US_GAMMAV1 * US_GAMMAV2) ^ 0.5 - - US_GAMMAVIDER[0] = US_GAMMAV1 - US_GAMMAV2 * US_VI[0] ^ (-2) - - US_KI[0] = (1 - US_DELTA) * US_KI[-1] + (1 - US_GAMMAI[-1]) * US_II[-1] - - US_GAMMAI[0] = US_GAMMAI1 / 2 * (US_II[0] / US_II[-1] - 1) ^ 2 - - US_GAMMAIDER[0] = US_GAMMAI1 * (US_II[0] / US_II[-1] - 1) / US_II[-1] - - US_GAMMAU[0] = ((US_DELTA + US_BETA ^ (-1) - 1) * US_QBAR - US_DELTA * US_TAUKBAR * US_PIBAR) / (US_PIBAR * (1 - US_TAUKBAR)) * (US_U[0] - 1) + US_GAMMAU2 / 2 * (US_U[0] - 1) ^ 2 - - US_GAMMAUDER[0] = ((US_DELTA + US_BETA ^ (-1) - 1) * US_QBAR - US_DELTA * US_TAUKBAR * US_PIBAR) / (US_PIBAR * (1 - US_TAUKBAR)) + (US_U[0] - 1) * US_GAMMAU2 - - US_RK[0] = US_GAMMAUDER[0] * US_PI[0] - - US_PI[0] = US_Q[0] * (1 - US_GAMMAI[0] - US_II[0] * US_GAMMAIDER[0]) + US_BETA * US_LAMBDAI[1] / US_LAMBDAI[0] * US_Q[1] * US_GAMMAIDER[1] * US_II[1] ^ 2 / US_II[0] - - US_Q[0] = US_BETA * US_LAMBDAI[1] / US_LAMBDAI[0] * ((1 - US_TAUK[1]) * (US_RK[1] * US_U[1] - US_GAMMAU[1] * US_PI[1]) + US_PI[1] * US_DELTA * US_TAUK[1] + (1 - US_DELTA) * US_Q[1]) - - US_WITILDE[0] ^ (1 + US_ZETA * US_ETAI) = US_ETAI / (US_ETAI - 1) * US_FI[0] / US_GI[0] - - US_FI[0] = US_WI[0] ^ ((1 + US_ZETA) * US_ETAI) * US_NDI[0] ^ (1 + US_ZETA) + US_BETA * US_XII * (US_PIC[1] / (US_PIC[0] ^ US_CHII * US_PI4TARGET ^ (0.25 * (1 - US_CHII)))) ^ ((1 + US_ZETA) * US_ETAI) * US_FI[1] - - US_GI[0] = US_NDI[0] * US_LAMBDAI[0] * (1 - US_TAUN[0] - US_TAUWH[0]) * US_WI[0] ^ US_ETAI + US_BETA * US_XII * (US_PIC[1] / (US_PIC[0] ^ US_CHII * US_PI4TARGET ^ (0.25 * (1 - US_CHII)))) ^ (US_ETAI - 1) * US_GI[1] - - US_WI[0] ^ (1 - US_ETAI) = (1 - US_XII) * US_WITILDE[0] ^ (1 - US_ETAI) + US_XII * US_WI[-1] ^ (1 - US_ETAI) * (US_PI4TARGET ^ (0.25 * (1 - US_CHII)) * US_PIC[-1] ^ US_CHII / US_PIC[0]) ^ (1 - US_ETAI) - - US_UTILJ[0] = 1 / (1 - US_SIGMA) * (US_CJ[0] - US_KAPPA * US_CJ[-1]) ^ (1 - US_SIGMA) - 1 / (1 + US_ZETA) * US_NJ[0] ^ (1 + US_ZETA) + US_BETA * US_UTILJ[1] - - US_CJ[0] * (1 + US_TAUC[0] + US_GAMMAVJ[0]) + US_MJ[0] = US_NJ[0] * (1 - US_TAUN[0] - US_TAUWH[0]) * US_WJ[0] + US_TRJ[0] - US_TJ[0] + US_MJ[-1] * US_PIC[0] ^ (-1) - - US_LAMBDAJ[0] * (1 + US_TAUC[0] + US_GAMMAVJ[0] + US_VJ[0] * US_GAMMAVJDER[0]) = (US_CJ[0] - US_KAPPA * US_CJ[-1]) ^ (-US_SIGMA) - - US_GAMMAVJDER[0] * US_VJ[0] ^ 2 = 1 - US_BETA * US_LAMBDAJ[1] / (US_PIC[1] * US_LAMBDAJ[0]) - - US_VJ[0] = (1 + US_TAUC[0]) * US_CJ[0] / US_MJ[0] - - US_GAMMAVJ[0] = US_GAMMAV1 * US_VJ[0] + US_GAMMAV2 / US_VJ[0] - 2 * (US_GAMMAV1 * US_GAMMAV2) ^ 0.5 - - US_GAMMAVJDER[0] = US_GAMMAV1 - US_GAMMAV2 * US_VJ[0] ^ (-2) - - US_WJTILDE[0] ^ (1 + US_ZETA * US_ETAJ) = US_ETAJ / (US_ETAJ - 1) * US_FJ[0] / US_GJ[0] - - US_FJ[0] = US_WJ[0] ^ ((1 + US_ZETA) * US_ETAJ) * US_NDJ[0] ^ (1 + US_ZETA) + US_BETA * US_XIJ * (US_PIC[1] / (US_PIC[0] ^ US_CHIJ * US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)))) ^ ((1 + US_ZETA) * US_ETAJ) * US_FJ[1] - - US_GJ[0] = US_NDJ[0] * (1 - US_TAUN[0] - US_TAUWH[0]) * US_LAMBDAJ[0] * US_WJ[0] ^ US_ETAJ + US_BETA * US_XIJ * (US_PIC[1] / (US_PIC[0] ^ US_CHIJ * US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)))) ^ (US_ETAJ - 1) * US_GJ[1] - - US_WJ[0] ^ (1 - US_ETAJ) = (1 - US_XIJ) * US_WJTILDE[0] ^ (1 - US_ETAJ) + US_XIJ * US_WJ[-1] ^ (1 - US_ETAJ) * (US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)) * US_PIC[-1] ^ US_CHIJ / US_PIC[0]) ^ (1 - US_ETAJ) - - US_YS[0] = US_Z[0] * US_KD[0] ^ US_ALPHA * US_ND[0] ^ (1 - US_ALPHA) - US_PSIBAR - - US_RK[0] = US_ALPHA * (US_YS[0] + US_PSIBAR) / US_KD[0] * US_MC[0] - - US_MC[0] = 1 / (US_Z[0] * US_ALPHA ^ US_ALPHA * (1 - US_ALPHA) ^ (1 - US_ALPHA)) * US_RK[0] ^ US_ALPHA * ((1 + US_TAUWF[0]) * US_W[0]) ^ (1 - US_ALPHA) - - US_NDI[0] = US_ND[0] * (1 - US_OMEGA) * (US_WI[0] / US_W[0]) ^ (-US_ETA) - - US_NDJ[0] = US_ND[0] * US_OMEGA * (US_WJ[0] / US_W[0]) ^ (-US_ETA) - - US_ND[0] ^ (1 - 1 / US_ETA) = (1 - US_OMEGA) ^ (1 / US_ETA) * US_NDI[0] ^ (1 - 1 / US_ETA) + US_OMEGA ^ (1 / US_ETA) * US_NDJ[0] ^ (1 - 1 / US_ETA) - - US_D[0] = US_Y[0] * US_PY[0] - US_RK[0] * US_KD[0] - US_ND[0] * (1 + US_TAUWF[0]) * US_W[0] - - US_PHTILDE[0] / US_PH[0] = US_THETA / (US_THETA - 1) * US_FH[0] / US_GH[0] - - US_FH[0] = US_MC[0] * US_H[0] + US_LAMBDAI[1] * US_BETA * US_XIH / US_LAMBDAI[0] * (US_PIH[1] / (US_PIH[0] ^ US_CHIH * US_PI4TARGET ^ (0.25 * (1 - US_CHIH)))) ^ US_THETA * US_FH[1] - - US_GH[0] = US_PH[0] * US_H[0] + US_LAMBDAI[1] * US_BETA * US_XIH / US_LAMBDAI[0] * (US_PIH[1] / (US_PIH[0] ^ US_CHIH * US_PI4TARGET ^ (0.25 * (1 - US_CHIH)))) ^ (US_THETA - 1) * US_GH[1] - - US_PH[0] ^ (1 - US_THETA) = (1 - US_XIH) * US_PHTILDE[0] ^ (1 - US_THETA) + US_XIH * (US_PH[-1] / US_PIC[0]) ^ (1 - US_THETA) * (US_PI4TARGET ^ (0.25 * (1 - US_CHIH)) * US_PIH[-1] ^ US_CHIH) ^ (1 - US_THETA) - - US_PIH[0] = US_PIC[0] * US_PH[0] / US_PH[-1] - - EA_PIMTILDE[0] / EA_PIM[0] = US_THETA / (US_THETA - 1) * US_FX[0] / US_GX[0] - - US_FX[0] = US_MC[0] * EA_IM[0] * EA_SIZE / US_SIZE + US_LAMBDAI[1] * US_BETA * US_XIX / US_LAMBDAI[0] * (EA_PIIM[1] / (EA_PIIM[0] ^ US_CHIX * US_PI4TARGET ^ (0.25 * (1 - US_CHIX)))) ^ US_THETA * US_FX[1] - - US_GX[0] = EA_IM[0] * EA_SIZE * EA_PIM[0] * USEA_RER[0] / US_SIZE + US_LAMBDAI[1] * US_BETA * US_XIX / US_LAMBDAI[0] * (EA_PIIM[1] / (EA_PIIM[0] ^ US_CHIX * US_PI4TARGET ^ (0.25 * (1 - US_CHIX)))) ^ (US_THETA - 1) * US_GX[1] - - EA_PIM[0] ^ (1 - US_THETA) = (1 - US_XIX) * EA_PIMTILDE[0] ^ (1 - US_THETA) + US_XIX * (EA_PIM[-1] / EA_PIC[0]) ^ (1 - US_THETA) * (EA_PIIM[-1] ^ US_CHIX * EA_PI4TARGET ^ (0.25 * (1 - US_CHIH))) ^ (1 - US_THETA) - - EA_PIIM[0] = EA_PIC[0] * EA_PIM[0] / EA_PIM[-1] - - USEA_RER[0] = US_RER / EA_RER[0] - - US_QC[0] ^ ((US_MUC - 1) / US_MUC) = US_NUC ^ (1 / US_MUC) * US_HC[0] ^ (1 - 1 / US_MUC) + (1 - US_NUC) ^ (1 / US_MUC) * ((1 - US_GAMMAIMC[0]) * US_IMC[0]) ^ (1 - 1 / US_MUC) - - 1 = US_NUC * US_PH[0] ^ (1 - US_MUC) + (1 - US_NUC) * (US_PIM[0] / US_GAMMAIMCDAG[0]) ^ (1 - US_MUC) - - US_HC[0] = US_QC[0] * US_NUC * US_PH[0] ^ (-US_MUC) - - US_GAMMAIMC[0] = US_GAMMAIMC1 / 2 * (US_IMC[0] / US_QC[0] / (US_IMC[-1] / US_QC[-1]) - 1) ^ 2 - - US_GAMMAIMCDAG[0] = 1 - US_GAMMAIMC[0] - US_IMC[0] * US_GAMMAIMC1 * (US_IMC[0] / US_QC[0] / (US_IMC[-1] / US_QC[-1]) - 1) / US_QC[0] / (US_IMC[-1] / US_QC[-1]) - - US_QI[0] ^ ((US_MUI - 1) / US_MUI) = US_NUI ^ (1 / US_MUI) * US_HI[0] ^ (1 - 1 / US_MUI) + (1 - US_NUI) ^ (1 / US_MUI) * ((1 - US_GAMMAIMI[0]) * US_IMI[0]) ^ (1 - 1 / US_MUI) - - US_PI[0] ^ (1 - US_MUI) = US_NUI * US_PH[0] ^ (1 - US_MUI) + (1 - US_NUI) * (US_PIM[0] / US_GAMMAIMIDAG[0]) ^ (1 - US_MUI) - - US_HI[0] = US_QI[0] * US_NUI * (US_PH[0] / US_PI[0]) ^ (-US_MUI) - - US_GAMMAIMI[0] = US_GAMMAIMI1 / 2 * (US_IMI[0] / US_QI[0] / (US_IMI[-1] / US_QI[-1]) - 1) ^ 2 - - US_GAMMAIMIDAG[0] = 1 - US_GAMMAIMI[0] - US_IMI[0] * US_GAMMAIMI1 * (US_IMI[0] / US_QI[0] / (US_IMI[-1] / US_QI[0]) - 1) / US_QI[0] / (US_IMI[-1] / US_QI[-1]) - - US_PH[-1] * US_G[-1] + US_TR[-1] + US_B[-1] * US_PIC[-1] ^ (-1) + US_PIC[-1] ^ (-1) * US_M[-2] = US_TAUC[-1] * US_C[-1] + (US_TAUN[-1] + US_TAUWH[-1]) * (US_WI[-1] * US_NDI[-1] + US_WJ[-1] * US_NDJ[-1]) + US_TAUWF[-1] * US_W[-1] * US_ND[-1] + US_TAUK[-1] * (US_RK[-1] * US_U[-1] - (US_DELTA + US_GAMMAU[-1]) * US_PI[-1]) * US_K[-1] + US_TAUD[-1] * US_D[-1] + US_T[-1] + US_R[-1] ^ (-1) * US_B[0] + US_M[-1] - - US_PH[0] * US_G[0] = US_GY[0] * US_PYBAR * US_YBAR - - US_TR[0] = US_YBAR * US_PYBAR * US_TRY[0] - - US_T[0] / (US_PYBAR * US_YBAR) = US_PHITB * (US_B[0] / (US_PYBAR * US_YBAR) - US_BYTARGET) - - US_TI[0] = US_T[0] * US_UPSILONT - - US_TRI[0] = US_TR[0] * US_UPSILONTR - - US_PIC4[0] = US_PIC[0] * US_PIC[-1] * US_PIC[-2] * US_PIC[-3] - - US_RR[0] - 1 = US_R[0] / US_PIC[1] - 1 - - US_C[0] = US_CI[0] * (1 - US_OMEGA) + US_CJ[0] * US_OMEGA - - US_M[0] = US_MI[0] * (1 - US_OMEGA) + US_MJ[0] * US_OMEGA - - US_K[0] = US_KI[0] * (1 - US_OMEGA) - - US_I[0] = US_II[0] * (1 - US_OMEGA) - - US_TRJ[0] = US_TR[0] * 1 / US_OMEGA - US_TRI[0] * (1 - US_OMEGA) / US_OMEGA - - US_TJ[0] = US_T[0] * 1 / US_OMEGA - US_TI[0] * (1 - US_OMEGA) / US_OMEGA - - US_GAMMAV[0] = US_GAMMAVI[0] * US_CI[0] * (1 - US_OMEGA) + US_GAMMAVJ[0] * US_CJ[0] * US_OMEGA - - US_NI[0] = US_NDI[0] * US_SI[0] - - US_SI[0] = (1 - US_XII) * (US_WITILDE[0] / US_WI[0]) ^ (-US_ETAI) + US_XII * (US_WI[-1] / US_WI[0]) ^ (-US_ETAI) * (US_PIC[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHII)) * US_PIC[-1] ^ US_CHII)) ^ US_ETAI * US_SI[-1] - - US_NJ[0] = US_NDJ[0] * US_SJ[0] - - US_SJ[0] = (1 - US_XIJ) * (US_WJTILDE[0] / US_WJ[0]) ^ (-US_ETAJ) + US_XIJ * (US_WJ[-1] / US_WJ[0]) ^ (-US_ETAJ) * (US_PIC[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)) * US_PIC[-1] ^ US_CHIJ)) ^ US_ETAJ * US_SJ[-1] - - US_U[0] * US_K[0] = US_KD[0] - - US_YS[0] = US_H[0] * US_SH[0] + EA_IM[0] * EA_SIZE * US_SX[0] / US_SIZE - - US_H[0] = US_G[0] + US_HC[0] + US_HI[0] - - US_IM[0] = US_IMC[0] + US_IMI[0] - - US_SH[0] = (1 - US_XIH) * (US_PHTILDE[0] / US_PH[0]) ^ (-US_THETA) + US_XIH * (US_PIH[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHIH)) * US_PIH[-1] ^ US_CHIH)) ^ US_THETA * US_SH[-1] - - US_SX[0] = (1 - US_XIX) * (EA_PIMTILDE[0] / EA_PIM[0]) ^ (-US_THETA) + US_XIX * (EA_PIIM[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHIH)) * EA_PIIM[-1] ^ US_CHIX)) ^ US_THETA * US_SX[-1] - - US_QC[0] = US_C[0] + US_GAMMAV[0] - - US_QI[0] = US_I[0] + US_GAMMAU[0] * US_K[0] - - US_Y[0] * US_PY[0] = EA_IM[0] * EA_SIZE * EA_PIM[0] * USEA_RER[0] / US_SIZE + US_PH[0] * US_G[0] + US_QC[0] + US_PI[0] * US_QI[0] - US_PIM[0] * ((1 - US_GAMMAIMC[0]) * US_IMC[0] / US_GAMMAIMCDAG[0] + (1 - US_GAMMAIMI[0]) * US_IMI[0] / US_GAMMAIMIDAG[0]) - - US_Y[0] = US_YS[0] - - log(US_Z[0]) = (1 - US_RHOZ) * log(US_ZBAR) + US_RHOZ * log(US_Z[-1]) + σ_US_Z * US_EPSZ[x] - - US_GY[0] = (1 - US_RHOG) * US_GYBAR + US_RHOG * US_GY[-1] + σ_US_G * US_EPSG[x] - - US_TRY[0] = (1 - US_RHOTR) * US_TRYBAR + US_RHOTR * US_TRY[-1] + σ_US_TR * US_EPSTR[x] - - US_TAUC[0] = (1 - US_RHOTAUC) * US_TAUCBAR + US_TAUC[-1] * US_RHOTAUC + σ_US_TAUC * US_EPSTAUC[x] - - US_TAUD[0] = (1 - US_RHOTAUD) * US_TAUDBAR + US_TAUD[-1] * US_RHOTAUD + σ_US_TAUD * US_EPSTAUD[x] - - US_TAUK[0] = US_TAUKBAR * (1 - US_RHOTAUK) + US_TAUK[-1] * US_RHOTAUK + σ_US_TAUK * US_EPSTAUK[x] - - US_TAUN[0] = (1 - US_RHOTAUN) * US_TAUNBAR + US_TAUN[-1] * US_RHOTAUN + σ_US_TAUN * US_EPSTAUN[x] - - US_TAUWH[0] = (1 - US_RHOTAUWH) * US_TAUWHBAR + US_TAUWH[-1] * US_RHOTAUWH + σ_US_TAUWH * US_EPSTAUWH[x] - - US_TAUWF[0] = (1 - US_RHOTAUWF) * US_TAUWFBAR + US_TAUWF[-1] * US_RHOTAUWF + σ_US_TAUWF * US_EPSTAUWF[x] - - US_CY[0] = US_C[0] / (US_Y[0] * US_PY[0]) - - US_IY[0] = US_PI[0] * US_I[0] / (US_Y[0] * US_PY[0]) - - US_IMY[0] = US_PIM[0] * US_IM[0] / (US_Y[0] * US_PY[0]) - - US_IMCY[0] = US_PIM[0] * US_IMC[0] / (US_Y[0] * US_PY[0]) - - US_IMIY[0] = US_PIM[0] * US_IMI[0] / (US_Y[0] * US_PY[0]) - - US_BY[0] = US_B[0] / (US_PYBAR * US_YBAR) - - US_TY[0] = US_T[0] / (US_PYBAR * US_YBAR) - - US_YGAP[0] = US_Y[0] / US_YBAR - 1 - - US_YGROWTH[0] = US_Y[0] / US_Y[-1] - - US_YSHARE[0] = US_Y[0] * US_SIZE * US_PY[0] / US_RER / (EA_Y[0] * EA_PY[0] * EA_SIZE / EA_RER[0] + US_Y[0] * US_SIZE * US_PY[0] / US_RER) - - US_EPSILONM[0] = ( - 0.125) / (US_R[0] * (US_R[0] + US_R[0] * US_GAMMAV2 - 1)) - - 1 = EA_LAMBDAI[1] * EA_BETA * US_R[0] * (1 - EA_GAMMAB[0]) / EA_LAMBDAI[0] * EA_RERDEP[1] / US_PIC[1] - - EA_GAMMAB[0] = EA_GAMMAB1 * (exp(EA_RER[0] * EA_BF[0] / US_PIC[0] / (EA_Y[0] * EA_PY[0]) - EA_BFYTARGET) - 1) - EA_RP[0] - - EA_RP[0] = EA_RHORP * EA_RP[-1] + σ_EA_RP * EA_EPSRP[x] - - EA_RERDEP[0] = EA_RER[0] / EA_RER[-1] - - EA_TOT[0] = EA_PIM[0] / (US_PIM[0] * EA_RER[0]) - - EA_TB[0] = US_IM[0] * US_SIZE * US_PIM[0] * EA_RER[0] / EA_SIZE - EA_PIM[0] * EA_IM[0] - - EA_BF[0] / US_R[-1] = EA_BF[-1] + EA_TB[-1] / EA_RER[-1] - - EA_SIZE * EA_BF[0] + US_SIZE * US_BF[0] = 0 - - end - - - @parameters NAWM_EAUS_2008_incomplete begin - EA_RRSTAR = 1 / EA_BETA - - US_RRSTAR = 1 / US_BETA - 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) - - # 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)) - - sol1 = get_solution(NAWM_EAUS_2008_incomplete, parameters = param_dict) - sol2 = get_solution(NAWM_EAUS_2008) - - @test isapprox(sol1, sol2, rtol = 1e-7) - end - - @testset verbose = true "Code quality (Aqua.jl)" begin - # Aqua.test_all(MacroModelling) - @testset "Compare Project.toml and test/Project.toml" Aqua.test_project_extras(MacroModelling) - @testset "Stale dependencies" Aqua.test_stale_deps(MacroModelling; ignore = [:Showoff]) - @testset "Unbound type parameters" Aqua.test_unbound_args(MacroModelling) - @testset "Undefined exports" Aqua.test_undefined_exports(MacroModelling) - @testset "Piracy" Aqua.test_piracies(MacroModelling) - @testset "Method ambiguity" Aqua.test_ambiguities(MacroModelling, recursive = false) - @testset "Compat" Aqua.test_deps_compat(MacroModelling)#; ignore = [:Aqua, :JET]) - # @testset "Persistent tasks" Aqua.test_persistent_tasks(MacroModelling) - end - GC.gc() - - # test_higher_order = true - @testset verbose = true "Test various models: NSSS and 1st order solution" begin - include("test_models.jl") - end - GC.gc() - # test_higher_order = false - - @testset verbose = true "for and if loops" begin - include("models/Backus_Kehoe_Kydland_1992_for_if_test.jl") - include("models/Backus_Kehoe_Kydland_1992.jl") - - std1 = get_std(Backus_Kehoe_Kydland_1992) - std2 = get_std(Backus_Kehoe_Kydland_1992_test) - - common_keys1 = intersect(std1.keys[1], std2.keys[1]) - common_keys2 = intersect(std1.keys[2], std2.keys[2]) - - @test isapprox(std2(common_keys1, common_keys2), std1(common_keys1, common_keys2), rtol = 1e-10) - end - - @testset verbose = true "Model without shocks" begin - @model m begin - K[0] = (1 - δ) * K[-1] + I[0] - Z[0] = (1 - ρ) * μ + ρ * Z[-1] - I[1] = ((ρ + δ - Z[0])/(1 - δ)) + ((1 + ρ)/(1 - δ)) * I[0] - end - - @parameters m verbose = true begin - ρ = 0.05 - δ = 0.10 - μ = .17 - σ = .2 - end - - m_ss = get_steady_state(m) - @test isapprox(m_ss(:,:Steady_state),[1/7.5,1/.75,.17],rtol = eps(Float32)) - - m_sol = get_solution(m) - @test isapprox(m_sol(:,:K),[1/.75,.9,.04975124378109454],rtol = eps(Float32)) - - init = m_ss(:,:Steady_state) |> collect - init[2] *= 1.5 - get_irf(m, initial_state = init, shocks = :none) - - plots = plot_irf(m, initial_state = init, shocks = :none) - - @test plots[1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} - - plots! = plot_irf!(m, initial_state = init .* 1.5, shocks = :none) - - @test plots![1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} - end - m = nothing - - - - @testset verbose = true "Distribution functions, general and SS" begin - - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - c_normcdf[0]= normcdf(c[0]) - c_normpdf[0]= normpdf(c[0]) - c_norminvcdf[0]= norminvcdf(c[0]-1) - c_norminv[0]= norminv(c[0]-1) - c_qnorm[0]= qnorm(c[0]-1) - c_dnorm[0]= dnorm(c[0]) - c_pnorm[0]= pnorm(c[0]) - c_normlogpdf[0]= normlogpdf(c[0]) - # c_norm[0]= cdf(Normal(),c[0]) - c_inv[0] = erfcinv(c[0]) - # c_binomlogpdf[0]= binomlogpdf(c[0]) - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - get_solution(RBC_CME) - - - @model finacc begin - R[0] * beta = C[1] / C[0] - C[0] = w[0] * L[0] - B[0] + R[-1] * B[-1] + (1-v) * (Rk[-1] * Q[-1] * K[-1] - (R[-1] + mu * G[0] * Rk[-1] * Q[-1] * K[-1] / (Q[-1] * K[-1] - N[-1])) * (Q[-1] * K[-1] - N[-1])) - We - w[0] = C[0] / (1-L[0]) - K[0] = (1-delta) * K[-1] + I[0] - Q[0] = 1 + chi * (I[0] / K[-1] - delta) - Y[0] = A[0] * K[-1]^alpha * L[0]^(1-alpha) - Rk[-1] = (alpha * Y[0] / K[-1] + Q[0] * (1-delta))/Q[-1] - w[0] = (1-alpha) * Y[0] / L[0] - N[0] = v * (Rk[-1] * Q[-1] * K[-1] - (R[-1] + mu * G[0] * Rk[-1] * Q[-1] * K[-1] / (Q[-1] * K[-1] - N[-1])) * (Q[-1] * K[-1] - N[-1])) + We - 0 = (omegabar[0] * (1 - F[0]) + (1 - mu) * G[0]) * Rk[0] / R[0] * Q[0] * K[0] / N[0] - (Q[0] * K[0] / N[0] - 1) - 0 = (1 - (omegabar[0] * (1 - F[0]) + G[0])) * Rk[0] / R[0] + (1 - F[0]) / (1 - F[0] - omegabar[0] * mu * (normpdf((log(omegabar[0]) + sigma^2/2) / sigma)/ omegabar[0] / sigma)) * ((omegabar[0] * (1 - F[0]) + (1 - mu) * G[0]) * Rk[0] / R[0] - 1) - G[0] = normcdf(((log(omegabar[0])+sigma^2/2)/sigma) - sigma) - F[0] = normcdf((log(omegabar[0])+sigma^2/2)/sigma) - EFP[0] = (mu * G[0] * Rk[-1] * Q[-1] * K[-1] / (Q[-1] * K[-1] - N[-1])) - Y[0] + walras[0] = C[0] + I[0] + EFP[0] * (Q[-1] * K[-1] - N[-1]) - B[0] = Q[0] * K[0] - N[0] - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - end - - - @parameters finacc verbose = true begin - beta = 0.99 - delta = 0.02 - We = 1e-12 - alpha = 0.36 - chi = 0 - v = 0.978 - mu = 0.94 - sigma = 0.2449489742783178 - rhoz = .9 - std_eps = .0068 - - .5 > omegabar > .44 - K > 15 - 0 < L < .45 - end - - # solve!(finacc) - @test isapprox(get_steady_state(finacc,derivatives = false)[:,1],[1.0, 7.004987166460695, 1.2762549358842095, 0.0008293608419033882, 0.0009318065746306208, 0.0003952537570055814, 0.30743973601435376, 15.371986800781423, 0.4435430773517457, 8.366999635233856, 1.0000000000593001, 1.0101010101010102, 1.0172249577970442, 1.5895043340984303, 0.4529051354389826, 2.2935377097663356, -1.4597012487627126e-10], rtol = 1e-6) - end - RBC_CME = nothing - finacc = nothing - - @testset verbose = true "Lead and lag > 1" begin - - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - sol = get_solution(RBC_CME) - - # sol1 = get_solution(RBC_CME, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol, sol1, atol = 1e-4) - - # sol2 = get_solution(RBC_CME, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol, sol2, atol = 1e-4) - - sol3 = get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol, sol3, atol = 1e-4) - - - # exo multi lead/lag >> 1 - @model RBC_CME_exo_mult begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-8] + eps_z[x-4] + eps_z[x+4] + eps_z_s[x]) - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_exo_mult verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_exo_mult - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_exo_mult = nothing - - - # endo/exo multi lead/lag >> 1 - @model RBC_CME_all_mult begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-8] + eps_z[x-4] + eps_z[x+4] + eps_z_s[x]) - ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - ZZ_avg_fut[0] = (A[0] + A[1] + A[2] + A[3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_all_mult verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_all_mult - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_all_mult = nothing - - - # exo lead >> 1 - @model RBC_CME_exo_lead1 begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x+8] - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_exo_lead1 verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_exo_lead1 - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_exo_lead1 = nothing - - - - # exo multi lag >> 1 - @model RBC_CME_exo_lag_mult begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-8] + eps_z[x-4] + eps_z_s[x]) - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_exo_lag_mult verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_exo_lag_mult - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_exo_lag_mult = nothing - - - # exo lag >> 1 - @model RBC_CME_exo_lag1 begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x-8] - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_exo_lag1 verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_exo_lag1 - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_exo_lag1 = nothing - - - # exo lead > 1 - @model RBC_CME_exo_lead begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x+1] - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_exo_lead verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_exo_lead - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_exo_lead = nothing - - - # exo lag > 1 - @model RBC_CME_exo_lag begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x-1] - # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_exo_lag verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_exo_lag - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_exo_lag = nothing - - # Lags > 1 - @model RBC_CME_lag begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - ZZ_dev[0] = log(c[0]/c[ss]) - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_lag verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_lag - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_lag = nothing - - # Leads > 1 - @model RBC_CME_lead begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - ZZ_avg[0] = (A[0] + A[1] + A[2] + A[3]) / 4 - ZZ_dev[0] = log(c[0]/c[ss]) - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - @parameters RBC_CME_lead verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_lead - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_lead = nothing - - - # Leads and lags > 1 - @model RBC_CME_lead_lag begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - ZZ_avg_f[0] = (A[0] + A[1] + A[2] + A[3]) / 4 - ZZ_avg_b[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_lead_lag verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_lead_lag - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_lead_lag = nothing - - - # Leads and lags > 10 - @model RBC_CME_lead_lag10 begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - ZZ_avg_f[0] = (A[0] + A[1] + A[2] + A[3]) / 4 - ZZ_avg_b[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - ZZ_avg_f10[0] = (A[0] + A[10]) / 2 - # ZZ_avg_b10[0] = (A[0] + A[-10]) / 2 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_lead_lag10 verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_lead_lag10 - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_lead_lag10 = nothing - - - # Leads and lags > 10 - @model RBC_CME_lead_lag20 begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - ZZ_avg_f[0] = (A[0] + A[1] + A[2] + A[3]) / 4 - ZZ_avg_b[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - ZZ_avg_ff[0] = (A[0] + A[10]) / 2 - ZZ_avg_bb[0] = (A[0] + A[-10]) / 2 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME_lead_lag20 verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - m = RBC_CME_lead_lag20 - - sol_exo = get_solution(m) - - # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) - - - # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) - - # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) - - sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) - - @test isapprox(sol_exo, sol_exo3, atol = 1e-4) - - m = nothing - RBC_CME_lead_lag20 = nothing - end - - - @testset verbose = true "Steady state RBC CME model" begin - # Basic test - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # A_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 - # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 - # y_avg[0] = log(y[0] / y[-4]) - # y_growth[0] = log(y[1] / y[2]) - # y_growthl[0] = log(y[0] / y[1]) - # y_growthl1[0] = log(y[-1] / y[0]) - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - # get_steady_state(RBC_CME)[1] - # using NLopt - # RBC_CME.SS_optimizer = NLopt.LD_LBFGS - # solve!(RBC_CME) - @test get_steady_state(RBC_CME)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0024019205374952, 1.003405325870413, 1.2092444352939415, 9.467573947982233, 1.42321160651834, 1.0] - # get_moments(RBC_CME)[1] - # irf(RBC_CME) - - RBC_CME = nothing - - - # Symbolic test - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME symbolic = true verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - alpha = .157 - - # beta | R[ss] = R_ss - # R_ss = 1.0035 - beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y - # I_K_ratio = .15 - delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - - # cap_share > 0 - # R_ss > 0 - # Pi_ss > 0 - # I_K_ratio > 0 - - # 0 < alpha < 1 - 0 < beta < 1 - # 0 < delta < 1 - 0 < Pibar - # 0 <= rhoz < 1 - phi_pi > 0 - - # 0 < A < 1 - # 0 < k < 50 - 0 < Pi - 0 < R - end - # get_steady_state(RBC_CME)[1] - # using NLopt - # RBC_CME.SS_optimizer = NLopt.LD_LBFGS - # solve!(RBC_CME,symbolic_SS = true) - @test get_steady_state(RBC_CME)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0024019205374952, 1.003405325870413, 1.2092444352939415, 9.467573947982233, 1.42321160651834, 1.0] - # get_moments(RBC_CME)[1] - - RBC_CME = nothing - - - - - # Numerical test with calibration targets - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve - R_ss = 1.0035 - # beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = 1.0025 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - - - # cap_share > 0 - # R_ss > 0 - # Pi_ss > 0 - # I_K_ratio > 0 - - # 0 < alpha < 1 - # 0 < beta < 1 - # 0 < delta < 1 - # 0 < Pibar - # 0 <= rhoz < 1 - # phi_pi > 0 - - # 0 < A < 1 - # 0 < k < 50 - # 0 < y < 10 - # 0 < c < 10 - end - # get_steady_state(RBC_CME)[1] - # using NLopt - # RBC_CME.SS_optimizer = NLopt.LD_LBFGS - # solve!(RBC_CME, verbose = true) - # RBC_CME.SS_init_guess[1:7] = [1.0, 1.0025, 1.0035, 1.2081023828249515, 9.437411555244328, 1.4212969209705313, 1.0] - # get_steady_state(RBC_CME) - @test get_steady_state(RBC_CME, verbose = true)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0025, 1.0035, 1.2081023824176236, 9.437411552284384, 1.4212969205027686, 1.0] - # get_moments(RBC_CME)[1] - - # RBC_CME.ss_solve_blocks[1]([0.15662344139650963, 1.2081023828249515, 0.02259036144578319, 9.437411555244328, 1.4212969209705313],RBC_CME) - - RBC_CME = nothing - - - # Symbolic test with calibration targets - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME symbolic = true verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss - R_ss = 1.0035 - # beta = .999 - - delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss] / y[ss] = I_K_ratio # this doesnt solve symbolically - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = 1.0025 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - - # cap_share > 0 - # R_ss > 0 - # Pi_ss > 0 - # I_K_ratio > 0 - - # 0 < alpha < .5 - # 0 < beta < 1 - # 0 < delta < 1 - # 0 < Pibar - # 0 <= rhoz < 1 - # phi_pi > 0 - - # 0 < A < 1 - # 0 < k < 50 - # 0 < y < 10 - # 0 < c < 10 - end - # # get_steady_state(RBC_CME)[1] - # # using NLopt - # # RBC_CME.SS_optimizer = NLopt.LD_LBFGS - # # get_steady_state(RBC_CME) - @test isapprox(get_steady_state(RBC_CME, verbose = true)(RBC_CME.constants.post_model_macro.var,:Steady_state), [1.0, 1.0025, 1.0035, 1.2081023828249515, 9.437411555244328, 1.4212969209705313, 1.0],rtol = eps(Float32)) - # get_moments(RBC_CME)[1] - - RBC_CME = nothing - end - - - - - # using MacroModelling: @model, @parameters, get_steady_state, solve! - - @testset verbose = true "Steady state Smets_Wouters_2003 model" begin - - @model Smets_Wouters_2003 begin - -q[0] + beta * ((1 - tau) * q[1] + epsilon_b[1] * (r_k[1] * z[1] - psi^-1 * r_k[ss] * (-1 + exp(psi * (-1 + z[1])))) * (C[1] - h * C[0])^(-sigma_c)) = 0 - -q_f[0] + beta * ((1 - tau) * q_f[1] + epsilon_b[1] * (r_k_f[1] * z_f[1] - psi^-1 * r_k_f[ss] * (-1 + exp(psi * (-1 + z_f[1])))) * (C_f[1] - h * C_f[0])^(-sigma_c)) = 0 - -r_k[0] + alpha * epsilon_a[0] * mc[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^(-1 + alpha) = 0 - -r_k_f[0] + alpha * epsilon_a[0] * mc_f[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^(-1 + alpha) = 0 - -G[0] + T[0] = 0 - -G[0] + G_bar * epsilon_G[0] = 0 - -G_f[0] + T_f[0] = 0 - -G_f[0] + G_bar * epsilon_G[0] = 0 - -L[0] + nu_w[0]^-1 * L_s[0] = 0 - -L_s_f[0] + L_f[0] * (W_i_f[0] * W_f[0]^-1)^(lambda_w^-1 * (-1 - lambda_w)) = 0 - L_s_f[0] - L_f[0] = 0 - L_s_f[0] + lambda_w^-1 * L_f[0] * W_f[0]^-1 * (-1 - lambda_w) * (-W_disutil_f[0] + W_i_f[0]) * (W_i_f[0] * W_f[0]^-1)^(-1 + lambda_w^-1 * (-1 - lambda_w)) = 0 - Pi_ws_f[0] - L_s_f[0] * (-W_disutil_f[0] + W_i_f[0]) = 0 - Pi_ps_f[0] - Y_f[0] * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 - -Q[0] + epsilon_b[0]^-1 * q[0] * (C[0] - h * C[-1])^(sigma_c) = 0 - -Q_f[0] + epsilon_b[0]^-1 * q_f[0] * (C_f[0] - h * C_f[-1])^(sigma_c) = 0 - -W[0] + epsilon_a[0] * mc[0] * (1 - alpha) * L[0]^(-alpha) * (K[-1] * z[0])^alpha = 0 - -W_f[0] + epsilon_a[0] * mc_f[0] * (1 - alpha) * L_f[0]^(-alpha) * (K_f[-1] * z_f[0])^alpha = 0 - -Y_f[0] + Y_s_f[0] = 0 - Y_s[0] - nu_p[0] * Y[0] = 0 - -Y_s_f[0] + Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 - beta * epsilon_b[1] * (C_f[1] - h * C_f[0])^(-sigma_c) - epsilon_b[0] * R_f[0]^-1 * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 - beta * epsilon_b[1] * pi[1]^-1 * (C[1] - h * C[0])^(-sigma_c) - epsilon_b[0] * R[0]^-1 * (C[0] - h * C[-1])^(-sigma_c) = 0 - Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) - lambda_p^-1 * Y_f[0] * (1 + lambda_p) * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-1 - lambda_p^-1 * (1 + lambda_p)) = 0 - epsilon_b[0] * W_disutil_f[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) - omega * epsilon_b[0] * epsilon_L[0] * L_s_f[0]^sigma_l = 0 - -1 + xi_p * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1) + (1 - xi_p) * pi_star[0]^(-lambda_p^-1) = 0 - -1 + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1) + xi_w * (W[-1] * W[0]^-1)^(-lambda_w^-1) * (pi[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1) = 0 - -Phi - Y_s[0] + epsilon_a[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^alpha = 0 - -Phi - Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_a[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^alpha = 0 - eta_b[exo] - log(epsilon_b[0]) + rho_b * log(epsilon_b[-1]) = 0 - -eta_L[exo] - log(epsilon_L[0]) + rho_L * log(epsilon_L[-1]) = 0 - eta_I[exo] - log(epsilon_I[0]) + rho_I * log(epsilon_I[-1]) = 0 - eta_w[exo] - f_1[0] + f_2[0] = 0 - eta_a[exo] - log(epsilon_a[0]) + rho_a * log(epsilon_a[-1]) = 0 - eta_p[exo] - g_1[0] + g_2[0] * (1 + lambda_p) = 0 - eta_G[exo] - log(epsilon_G[0]) + rho_G * log(epsilon_G[-1]) = 0 - -f_1[0] + beta * xi_w * f_1[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1) + epsilon_b[0] * w_star[0] * L[0] * (1 + lambda_w)^-1 * (C[0] - h * C[-1])^(-sigma_c) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) = 0 - -f_2[0] + beta * xi_w * f_2[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) + omega * epsilon_b[0] * epsilon_L[0] * (L[0] * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)))^(1 + sigma_l) = 0 - -g_1[0] + beta * xi_p * pi_star[0] * g_1[1] * pi_star[1]^-1 * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1) + epsilon_b[0] * pi_star[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 - -g_2[0] + beta * xi_p * g_2[1] * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_b[0] * mc[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 - -nu_w[0] + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) + xi_w * nu_w[-1] * (W[-1] * pi[0]^-1 * W[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w)) = 0 - -nu_p[0] + (1 - xi_p) * pi_star[0]^(-lambda_p^-1 * (1 + lambda_p)) + xi_p * nu_p[-1] * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) = 0 - -K[0] + K[-1] * (1 - tau) + I[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2) = 0 - -K_f[0] + K_f[-1] * (1 - tau) + I_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2) = 0 - U[0] - beta * U[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C[0] - h * C[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s[0]^(1 + sigma_l)) = 0 - U_f[0] - beta * U_f[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C_f[0] - h * C_f[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s_f[0]^(1 + sigma_l)) = 0 - -epsilon_b[0] * (C[0] - h * C[-1])^(-sigma_c) + q[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2 - varphi * I[-1]^-1 * epsilon_I[0] * I[0] * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])) + beta * varphi * I[0]^-2 * epsilon_I[1] * q[1] * I[1]^2 * (-1 + I[0]^-1 * epsilon_I[1] * I[1]) = 0 - -epsilon_b[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) + q_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2 - varphi * I_f[-1]^-1 * epsilon_I[0] * I_f[0] * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])) + beta * varphi * I_f[0]^-2 * epsilon_I[1] * q_f[1] * I_f[1]^2 * (-1 + I_f[0]^-1 * epsilon_I[1] * I_f[1]) = 0 - eta_pi[exo] - log(pi_obj[0]) + rho_pi_bar * log(pi_obj[-1]) + log(calibr_pi_obj) * (1 - rho_pi_bar) = 0 - -C[0] - I[0] - T[0] + Y[0] - psi^-1 * r_k[ss] * K[-1] * (-1 + exp(psi * (-1 + z[0]))) = 0 - -calibr_pi + eta_R[exo] - log(R[ss]^-1 * R[0]) + r_Delta_pi * (-log(pi[ss]^-1 * pi[-1]) + log(pi[ss]^-1 * pi[0])) + r_Delta_y * (-log(Y[ss]^-1 * Y[-1]) + log(Y[ss]^-1 * Y[0]) + log(Y_f[ss]^-1 * Y_f[-1]) - log(Y_f[ss]^-1 * Y_f[0])) + rho * log(R[ss]^-1 * R[-1]) + (1 - rho) * (log(pi_obj[0]) + r_pi * (-log(pi_obj[0]) + log(pi[ss]^-1 * pi[-1])) + r_Y * (log(Y[ss]^-1 * Y[0]) - log(Y_f[ss]^-1 * Y_f[0]))) = 0 - -C_f[0] - I_f[0] + Pi_ws_f[0] - T_f[0] + Y_f[0] + L_s_f[0] * W_disutil_f[0] - L_f[0] * W_f[0] - psi^-1 * r_k_f[ss] * K_f[-1] * (-1 + exp(psi * (-1 + z_f[0]))) = 0 - epsilon_b[0] * (K[-1] * r_k[0] - r_k[ss] * K[-1] * exp(psi * (-1 + z[0]))) * (C[0] - h * C[-1])^(-sigma_c) = 0 - epsilon_b[0] * (K_f[-1] * r_k_f[0] - r_k_f[ss] * K_f[-1] * exp(psi * (-1 + z_f[0]))) * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 - end - - - @parameters Smets_Wouters_2003 verbose = true begin - calibr_pi_obj | 1 = pi_obj[ss] - calibr_pi | pi[ss] = pi_obj[ss] - # Phi | Y_s[ss] * .408 = Phi - # Phi = .408 * Y_j[ss] - # (Y_j[ss] + Phi) / Y_j[ss] = 1.408 -> Phi; | this seems problematic because of the parameter - # lambda_p | .6 = C_f[ss] / Y_f[ss] - # lambda_w | L[ss] = .33 - # G_bar | .18 = G[ss] / Y[ss] - # calibr_pi_obj = 0 - # calibr_pi = 1 - lambda_p = .368 - G_bar = .362 - lambda_w = 0.5 - Phi = .819 - - alpha = 0.3 - beta = 0.99 - gamma_w = 0.763 - gamma_p = 0.469 - h = 0.573 - omega = 1 - psi = 0.169 - r_pi = 1.684 - r_Y = 0.099 - r_Delta_pi = 0.14 - r_Delta_y = 0.159 - rho = 0.961 - rho_b = 0.855 - rho_L = 0.889 - rho_I = 0.927 - rho_a = 0.823 - rho_G = 0.949 - rho_pi_bar = 0.924 - sigma_c = 1.353 - sigma_l = 2.4 - tau = 0.025 - varphi = 6.771 - xi_w = 0.737 - xi_p = 0.908 - - # Putting non-negative constraint on first block is enough - # 0 < K - # 0 < I - # 0 < Y_s - # 0 < q - # 0 < r_k - # 5 < f_1 - # 0 < L - # 0 < W - # 30 < g_1 - # 0 < z - # 0 < mc - # 0 < w_star - # 5 < f_2 - # 0 < Y - # 0 < g_2 - # 0 < C - end - - - # solve!(Smets_Wouters_2003, verbose = true) - - - @test isapprox(get_steady_state(Smets_Wouters_2003, verbose = true)(Smets_Wouters_2003.constants.post_model_macro.var,[:Steady_state]), - [ 1.2043777509278788 - 1.2043777484127967 - 0.362 - 0.362 - 0.44153840098985714 - 0.44153839784516097 - 17.66153603957938 - 17.66153591381742 - 1.2891159430437658 - 1.2891159432893282 - 1.289115942962812 - 1.289115943290125 - 0.9999999999999677 - 0.5401411855429173 - 0.482173806623137 - 0.999999999999352 - 1.0000000000002556 - 1.0101010101010102 - 1.0101010101010102 - 0.362 - 0.362 - -427.9858908413812 - -427.98589116567274 - 1.122103431093411 - 0.7480689524203904 - 1.1221034286309022 - 1.122103428630708 - 2.0079161519182205 - 2.0079161462568305 - 2.0079161519185624 - 2.007916146256947 - 1.0 - 1.0 - 1.0 - 1.0 - 1.0 - 8.770699454739315 - 8.770699454739393 - 48.84717584575426 - 35.70699988724729 - 0.7309941520471651 - 0.730994152046712 - 1.0 - 0.9999999999372022 - 1.0 - 1.0 - 1.0 - 2.4590033452182944 - 2.459003352168301 - 0.035101010082054955 - 0.03510101010809239 - 1.1221034311168996 - 0.9999999999999698 - 1.0000000000000089], - rtol = eps(Float32) - ) - - Smets_Wouters_2003 = nothing - # # x = 1 - - - # # Smets_Wouters_2003 = nothing - - # # @model Smets_Wouters_2003 begin - # # -q[0] + beta * ((1 - tau) * q[1] + epsilon_b[1] * (r_k[1] * z[1] - psi^-1 * r_k[ss] * (-1 + exp(psi * (-1 + z[1])))) * (C[1] - h * C[0])^(-sigma_c)) = 0 - # # -q_f[0] + beta * ((1 - tau) * q_f[1] + epsilon_b[1] * (r_k_f[1] * z_f[1] - psi^-1 * r_k_f[ss] * (-1 + exp(psi * (-1 + z_f[1])))) * (C_f[1] - h * C_f[0])^(-sigma_c)) = 0 - # # -r_k[0] + alpha * epsilon_a[0] * mc[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^(-1 + alpha) = 0 - # # -r_k_f[0] + alpha * epsilon_a[0] * mc_f[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^(-1 + alpha) = 0 - # # -G[0] + T[0] = 0 - # # -G[0] + G_bar * epsilon_G[0] = 0 - # # -G_f[0] + T_f[0] = 0 - # # -G_f[0] + G_bar * epsilon_G[0] = 0 - # # -L[0] + nu_w[0]^-1 * L_s[0] = 0 - # # -L_s_f[0] + L_f[0] * (W_i_f[0] * W_f[0]^-1)^(lambda_w^-1 * (-1 - lambda_w)) = 0 - # # L_s_f[0] - L_f[0] = 0 - # # L_s_f[0] + lambda_w^-1 * L_f[0] * W_f[0]^-1 * (-1 - lambda_w) * (-W_disutil_f[0] + W_i_f[0]) * (W_i_f[0] * W_f[0]^-1)^(-1 + lambda_w^-1 * (-1 - lambda_w)) = 0 - # # Pi_ws_f[0] - L_s_f[0] * (-W_disutil_f[0] + W_i_f[0]) = 0 - # # Pi_ps_f[0] - Y_f[0] * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 - # # -Q[0] + epsilon_b[0]^-1 * q[0] * (C[0] - h * C[-1])^(sigma_c) = 0 - # # -Q_f[0] + epsilon_b[0]^-1 * q_f[0] * (C_f[0] - h * C_f[-1])^(sigma_c) = 0 - # # -W[0] + epsilon_a[0] * mc[0] * (1 - alpha) * L[0]^(-alpha) * (K[-1] * z[0])^alpha = 0 - # # -W_f[0] + epsilon_a[0] * mc_f[0] * (1 - alpha) * L_f[0]^(-alpha) * (K_f[-1] * z_f[0])^alpha = 0 - # # -Y_f[0] + Y_s_f[0] = 0 - # # Y_s[0] - nu_p[0] * Y[0] = 0 - # # -Y_s_f[0] + Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 - # # beta * epsilon_b[1] * (C_f[1] - h * C_f[0])^(-sigma_c) - epsilon_b[0] * R_f[0]^-1 * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 - # # beta * epsilon_b[1] * pi[1]^-1 * (C[1] - h * C[0])^(-sigma_c) - epsilon_b[0] * R[0]^-1 * (C[0] - h * C[-1])^(-sigma_c) = 0 - # # Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) - lambda_p^-1 * Y_f[0] * (1 + lambda_p) * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-1 - lambda_p^-1 * (1 + lambda_p)) = 0 - # # epsilon_b[0] * W_disutil_f[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) - omega * epsilon_b[0] * epsilon_L[0] * L_s_f[0]^sigma_l = 0 - # # -1 + xi_p * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1) + (1 - xi_p) * pi_star[0]^(-lambda_p^-1) = 0 - # # -1 + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1) + xi_w * (W[-1] * W[0]^-1)^(-lambda_w^-1) * (pi[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1) = 0 - # # -Phi - Y_s[0] + epsilon_a[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^alpha = 0 - # # -Phi - Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_a[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^alpha = 0 - # # eta_b[exo] - log(epsilon_b[0]) + rho_b * log(epsilon_b[-1]) = 0 - # # -eta_L[exo] - log(epsilon_L[0]) + rho_L * log(epsilon_L[-1]) = 0 - # # eta_I[exo] - log(epsilon_I[0]) + rho_I * log(epsilon_I[-1]) = 0 - # # eta_w[exo] - f_1[0] + f_2[0] = 0 - # # eta_a[exo] - log(epsilon_a[0]) + rho_a * log(epsilon_a[-1]) = 0 - # # eta_p[exo] - g_1[0] + g_2[0] * (1 + lambda_p) = 0 - # # eta_G[exo] - log(epsilon_G[0]) + rho_G * log(epsilon_G[-1]) = 0 - # # -f_1[0] + beta * xi_w * f_1[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1) + epsilon_b[0] * w_star[0] * L[0] * (1 + lambda_w)^-1 * (C[0] - h * C[-1])^(-sigma_c) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) = 0 - # # -f_2[0] + beta * xi_w * f_2[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) + omega * epsilon_b[0] * epsilon_L[0] * (L[0] * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)))^(1 + sigma_l) = 0 - # # -g_1[0] + beta * xi_p * pi_star[0] * g_1[1] * pi_star[1]^-1 * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1) + epsilon_b[0] * pi_star[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 - # # -g_2[0] + beta * xi_p * g_2[1] * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_b[0] * mc[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 - # # -nu_w[0] + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) + xi_w * nu_w[-1] * (W[-1] * pi[0]^-1 * W[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w)) = 0 - # # -nu_p[0] + (1 - xi_p) * pi_star[0]^(-lambda_p^-1 * (1 + lambda_p)) + xi_p * nu_p[-1] * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) = 0 - # # -K[0] + K[-1] * (1 - tau) + I[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2) = 0 - # # -K_f[0] + K_f[-1] * (1 - tau) + I_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2) = 0 - # # U[0] - beta * U[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C[0] - h * C[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s[0]^(1 + sigma_l)) = 0 - # # U_f[0] - beta * U_f[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C_f[0] - h * C_f[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s_f[0]^(1 + sigma_l)) = 0 - # # -epsilon_b[0] * (C[0] - h * C[-1])^(-sigma_c) + q[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2 - varphi * I[-1]^-1 * epsilon_I[0] * I[0] * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])) + beta * varphi * I[0]^-2 * epsilon_I[1] * q[1] * I[1]^2 * (-1 + I[0]^-1 * epsilon_I[1] * I[1]) = 0 - # # -epsilon_b[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) + q_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2 - varphi * I_f[-1]^-1 * epsilon_I[0] * I_f[0] * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])) + beta * varphi * I_f[0]^-2 * epsilon_I[1] * q_f[1] * I_f[1]^2 * (-1 + I_f[0]^-1 * epsilon_I[1] * I_f[1]) = 0 - # # eta_pi[exo] - log(pi_obj[0]) + rho_pi_bar * log(pi_obj[-1]) + log(calibr_pi_obj) * (1 - rho_pi_bar) = 0 - # # -C[0] - I[0] - T[0] + Y[0] - psi^-1 * r_k[ss] * K[-1] * (-1 + exp(psi * (-1 + z[0]))) = 0 - # # -calibr_pi + eta_R[exo] - log(R[ss]^-1 * R[0]) + r_Delta_pi * (-log(pi[ss]^-1 * pi[-1]) + log(pi[ss]^-1 * pi[0])) + r_Delta_y * (-log(Y[ss]^-1 * Y[-1]) + log(Y[ss]^-1 * Y[0]) + log(Y_f[ss]^-1 * Y_f[-1]) - log(Y_f[ss]^-1 * Y_f[0])) + rho * log(R[ss]^-1 * R[-1]) + (1 - rho) * (log(pi_obj[0]) + r_pi * (-log(pi_obj[0]) + log(pi[ss]^-1 * pi[-1])) + r_Y * (log(Y[ss]^-1 * Y[0]) - log(Y_f[ss]^-1 * Y_f[0]))) = 0 - # # -C_f[0] - I_f[0] + Pi_ws_f[0] - T_f[0] + Y_f[0] + L_s_f[0] * W_disutil_f[0] - L_f[0] * W_f[0] - psi^-1 * r_k_f[ss] * K_f[-1] * (-1 + exp(psi * (-1 + z_f[0]))) = 0 - # # epsilon_b[0] * (K[-1] * r_k[0] - r_k[ss] * K[-1] * exp(psi * (-1 + z[0]))) * (C[0] - h * C[-1])^(-sigma_c) = 0 - # # epsilon_b[0] * (K_f[-1] * r_k_f[0] - r_k_f[ss] * K_f[-1] * exp(psi * (-1 + z_f[0]))) * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 - # # end - - - # # @parameters Smets_Wouters_2003 begin - # # calibr_pi_obj | 1 = pi_obj[ss] - # # calibr_pi | pi[ss] = pi_obj[ss] - # # Phi | (Y_s[ss] + Phi) / Y_s[ss] = 1.408 - # # # lambda_p | .6 = C_f[ss] / Y_f[ss] - # # # lambda_w | L[ss] = .33 - # # G_bar | .18 = G[ss] / Y[ss] - - # # lambda_p = .368 - # # # G_bar = .362 - # # lambda_w = 0.5 - # # # Phi = .819 - - # # alpha = 0.3 - # # beta = 0.99 - # # gamma_w = 0.763 - # # gamma_p = 0.469 - # # h = 0.573 - # # omega = 1 - # # psi = 0.169 - # # r_pi = 1.684 - # # r_Y = 0.099 - # # r_Delta_pi = 0.14 - # # r_Delta_y = 0.159 - # # rho = 0.961 - # # rho_b = 0.855 - # # rho_L = 0.889 - # # rho_I = 0.927 - # # rho_a = 0.823 - # # rho_G = 0.949 - # # rho_pi_bar = 0.924 - # # sigma_c = 1.353 - # # sigma_l = 2.4 - # # tau = 0.025 - # # varphi = 6.771 - # # xi_w = 0.737 - # # xi_p = 0.908 - - # # end - - - # # solve!(Smets_Wouters_2003, symbolic_SS = false) - - # # # get_steady_state(Smets_Wouters_2003) - - # # @test get_steady_state(Smets_Wouters_2003)[1] ≈ [ 1.20465991441435 - # # 1.204659917151701 - # # 0.3613478048030788 - # # 0.3613478048030788 - # # 0.4414800855444218 - # # 0.4414800896382151 - # # 17.659203422264238 - # # 17.65920357698873 - # # 1.2889457095271066 - # # 1.2889457096070582 - # # 1.2889457095307755 - # # 1.2889457098239414 - # # 1.0000000000366498 - # # 0.5400259611608715 - # # 0.48211013259048446 - # # 1.00000000000172 - # # 1.000000000127065 - # # 1.0101010101010102 - # # 1.0101010101010102 - # # 0.3613478047907606 - # # 0.3613478048030788 - # # -427.92495898028676 - # # -427.9249587468684 - # # 1.1221034247496608 - # # 0.7480689524616317 - # # 1.122103428477167 - # # 1.1221034282377538 - # # 2.0074878047372287 - # # 2.00748781245403 - # # 2.007487804732286 - # # 2.0074878121606647 - # # 1.0 - # # 1.0 - # # 1.0 - # # 1.0 - # # 1.0 - # # 8.766762166589194 - # # 8.766762166588967 - # # 48.8212791635492 - # # 35.68806956399776 - # # 0.730994152045567 - # # 0.7309941520886629 - # # 1.0 - # # 1.0000000000028464 - # # 1.0 - # # 1.0 - # # 1.0 - # # 2.4582240979093846 - # # 2.4582240906598867 - # # 0.03510101014899653 - # # 0.035101010136073356 - # # 1.1221034247485961 - # # 1.0000000000000178 - # # 0.9999999999583465] - - - end - - - @testset verbose = true "First order perturbation" begin - # Numerical test with calibration targets - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve - R_ss = 1.0035 - # beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = 1.0025 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - get_solution(RBC_CME) - @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]], [ 0.0 0.0068 - 6.73489e-6 0.000168887 - 1.01124e-5 0.000253583 - -0.000365783 0.00217203 - -0.00070019 0.00749279 - 0.0 0.00966482 - 0.005 0.0], atol = 1e-6) - - get_solution(RBC_CME, parameters = :I_K_ratio => .1) - @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 - 3.42408e-6 0.000111417 - 5.14124e-6 0.000167292 - -0.000196196 0.00190741 - -0.000430554 0.0066164 - 0.0 0.00852381 - 0.005 0.0], atol = 1e-6) - - get_solution(RBC_CME, parameters = :cap_share => 1.5) - @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 - 4.00629e-6 0.000118171 - 6.01543e-6 0.000177434 - -0.000207089 0.00201698 - -0.00041124 0.00639229 - 0.0 0.00840927 - 0.005 0.0], atol = 1e-6) - - RBC_CME = nothing - end - - - - - - # @testset verbose = true "First order: linear time iteration" begin - # # Numerical test with calibration targets - # @model RBC_CME begin - # y[0]=A[0]*k[-1]^alpha - # 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - # 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - # R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - # z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # end - - - # @parameters RBC_CME verbose = true begin - # alpha | k[ss] / (4 * y[ss]) = cap_share - # cap_share = 1.66 - # # alpha = .157 - - # beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve - # R_ss = 1.0035 - # # beta = .999 - - # # delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. - # I_K_ratio = .15 - # # delta = .0226 - - # Pibar | Pi[ss] = Pi_ss - # Pi_ss = 1.0025 - # # Pibar = 1.0008 - - # phi_pi = 1.5 - # rhoz = .9 - # std_eps = .0068 - # rho_z_delta = .9 - # std_z_delta = .005 - # end - - # get_solution(RBC_CME, algorithm = :linear_time_iteration) - - # @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]], [ 0.0 0.0068 - # 6.73489e-6 0.000168887 - # 1.01124e-5 0.000253583 - # -0.000365783 0.00217203 - # -0.00070019 0.00749279 - # 0.0 0.00966482 - # 0.005 0.0], atol = 1e-6) - - - # get_solution(RBC_CME, algorithm = :linear_time_iteration, parameters = :I_K_ratio => .1) - - # @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 - # 3.42408e-6 0.000111417 - # 5.14124e-6 0.000167292 - # -0.000196196 0.00190741 - # -0.000430554 0.0066164 - # 0.0 0.00852381 - # 0.005 0.0], atol = 1e-6) - - - # get_solution(RBC_CME, algorithm = :linear_time_iteration, parameters = :cap_share => 1.5) - - # @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 - # 4.00629e-6 0.000118171 - # 6.01543e-6 0.000177434 - # -0.000207089 0.00201698 - # -0.00041124 0.00639229 - # 0.0 0.00840927 - # 0.005 0.0], atol = 1e-6) - - # RBC_CME = nothing - # end - - - - - - @testset verbose = true "First order: doubling" begin - # Numerical test with calibration targets - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve - R_ss = 1.0035 - # beta = .999 - - # delta | c[ss]/y[ss] = 1 - I_K_ratio - delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = 1.0025 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - end - - get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling) - - @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]], [ 0.0 0.0068 - 6.73489e-6 0.000168887 - 1.01124e-5 0.000253583 - -0.000365783 0.00217203 - -0.00070019 0.00749279 - 0.0 0.00966482 - 0.005 0.0], atol = 1e-6) - - - get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling, parameters = :I_K_ratio => .1) - - @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 - 3.42408e-6 0.000111417 - 5.14124e-6 0.000167292 - -0.000196196 0.00190741 - -0.000430554 0.0066164 - 0.0 0.00852381 - 0.005 0.0], atol = 1e-6) - - - get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling, parameters = :cap_share => 1.5) - - @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 - 4.00629e-6 0.000118171 - 6.01543e-6 0.000177434 - -0.000207089 0.00201698 - -0.00041124 0.00639229 - 0.0 0.00840927 - 0.005 0.0], atol = 1e-6) - - RBC_CME = nothing - end - - - @testset verbose = true "Plotting" begin - # Symbolic test with calibration targets - @model RBC_CME begin - y[0]=A[0]*k[-1]^alpha - 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) - 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) - R[0] * beta =(Pi[0]/Pibar)^phi_pi - # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] - A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] - z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] - # z[0]=rhoz*z[-1]+std_eps*eps_z[x] - # A[0]=exp(z[0]) - A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] - # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] - end - - - @parameters RBC_CME verbose = true begin - alpha | k[ss] / (4 * y[ss]) = cap_share - cap_share = 1.66 - # alpha = .157 - - beta | R[ss] = R_ss - R_ss = 1.0035 - # beta = .999 - - delta | c[ss]/y[ss] = 1 - I_K_ratio - # delta | delta * k[ss] / y[ss] = I_K_ratio # this doesnt solve symbolically - I_K_ratio = .15 - # delta = .0226 - - Pibar | Pi[ss] = Pi_ss - Pi_ss = 1.0025 - # Pibar = 1.0008 - - phi_pi = 1.5 - rhoz = .9 - std_eps = .0068 - rho_z_delta = .9 - std_z_delta = .005 - - # cap_share > 0 - # R_ss > 0 - # Pi_ss > 0 - # I_K_ratio > 0 - - # 0 < alpha < 1 - # 0 < beta < 1 - # 0 < delta < 1 - # 0 < Pibar - # 0 <= rhoz < 1 - # phi_pi > 0 - - # 0 < A < 1 - # 0 < k < 50 - # 0 < y < 10 - # 0 < c < 10 - end - plots = plot_irf(RBC_CME) - - @test plots[1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} - - plots! = plot_irf!(RBC_CME, parameters = :rhoz => .8) - - @test plots![1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} - - RBC_CME = nothing +elseif test_set == "plots_1" + include("test_plots_1.jl") +elseif test_set == "plots_2" + include("test_plots_2.jl") +elseif test_set == "plots_3" + include("test_plots_3.jl") +elseif test_set == "plots_4" + include("test_plots_4.jl") +elseif test_set == "plots_5" + include("test_plots_5.jl") +elseif test_set == "higher_order_1" + include("test_higher_order_1.jl") +elseif test_set == "higher_order_2" + include("test_higher_order_2.jl") +elseif test_set == "higher_order_3" + include("test_higher_order_3.jl") +elseif test_set == "basic" + include("test_basic.jl") +elseif test_set == "system_prior_estimation" + include("test_system_prior_estimation.jl") +elseif test_set == "gradient_checks" + include("test_rrule_robustness.jl") + include("test_inversion_filter_gradients.jl") +elseif test_set == "update_equations" + include("test_update_equations.jl") +elseif test_set == "jet_hot_paths" + include("test_jet_hot_paths.jl") +elseif test_set == "dynare_comparison" + # Dynare comparison runs as a standalone 3-phase pipeline (see CI workflow). + # If output/ exists with results, run the comparison script directly. + outdir = joinpath(@__DIR__, "dynare_comparison", "output") + if isdir(outdir) && !isempty(readdir(outdir)) + include(joinpath(@__DIR__, "dynare_comparison", "compare_results.jl")) + else + @warn "No Dynare comparison output found. Run generate_julia_results.jl + Docker first." end - GC.gc() - end diff --git a/test/test_1st_order_inversion_filter_estimation.jl b/test/test_1st_order_inversion_filter_estimation.jl index 1360cbcec..50513dca9 100644 --- a/test/test_1st_order_inversion_filter_estimation.jl +++ b/test/test_1st_order_inversion_filter_estimation.jl @@ -1,19 +1,30 @@ +using Test using MacroModelling import Turing -import Turing: NUTS, sample, logpdf -import ADTypes: AutoZygote +import Mooncake +import Turing: NUTS, sample +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys + +using FlexiChains +include("test_helpers.jl") 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,:) @@ -31,43 +42,58 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, filter, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) llh = get_loglikelihood(m, data, all_params, filter = filter, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end n_samples = 1000 -samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :inversion, -Inf), NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) +samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :inversion, -Inf), NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) -println("Mean variable values (Zygote): $(mean(samps).nt.mean)") +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") -sample_nuts = mean(samps).nt.mean +sample_nuts = collect(values(FlexiChains.mean(samps); parameters_only = true)) modeFS2000i = Turing.maximum_a_posteriori(FS2000_loglikelihood_function(data, FS2000, :inversion, -Inf), Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), - adtype = AutoZygote(), - initial_params = FS2000.parameter_values) - -println("Mode variable values: $(modeFS2000i.values); Mode loglikelihood: $(modeFS2000i.lp)") - + adtype = AutoMooncake(; config=nothing), + initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) + +println("Mode variable values: $(modeFS2000i.params); Mode loglikelihood: $(modeFS2000i.lp)") + +@testset "Mooncake vs FiniteDifferences gradient (1st order inversion)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(FS2000, data, x, filter = :inversion), ADTypes.AutoMooncake(config = nothing), FS2000.parameter_values) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + 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, 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 @@ -112,10 +138,10 @@ println("Mode variable values: $(modeFS2000i.values); Mode loglikelihood: $(mode # n_rounds = 6, # multithreaded = false) -# samps = MCMCChains.Chains(Pigeons.get_sample(pt)) +# samps = pigeons_flexichain(Pigeons.sample_array(pt), Pigeons.sample_names(pt)) -# println(mean(samps).nt.mean) +# println(collect(values(FlexiChains.mean(samps); parameters_only = true))) # Random.seed!(30) @@ -220,7 +246,7 @@ println("Mode variable values: $(modeFS2000i.values); Mode loglikelihood: $(mode # 1 # @testset "Estimation results" begin # @test isapprox(sol.minimum, -1343.7491257498598, rtol = eps(Float32)) -# @test isapprox(mean(samps).nt.mean, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) +# @test isapprox(collect(values(FlexiChains.mean(samps); parameters_only = true)), [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) # end diff --git a/test/test_1st_order_inversion_filter_estimation_pigeons.jl b/test/test_1st_order_inversion_filter_estimation_pigeons.jl index b03365c7d..29a0b5060 100644 --- a/test/test_1st_order_inversion_filter_estimation_pigeons.jl +++ b/test/test_1st_order_inversion_filter_estimation_pigeons.jl @@ -2,19 +2,22 @@ using MacroModelling using Test import Turing import Pigeons -import Turing: logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL +include("test_helpers.jl") + 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,:) @@ -32,7 +35,7 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, filter, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() llh = get_loglikelihood(m, @@ -40,9 +43,7 @@ Turing.@model function FS2000_loglikelihood_function(data, m, filter, on_failure all_params, filter = filter, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end @@ -52,6 +53,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 +66,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..b3c99445b 100644 --- a/test/test_2nd_order_estimation.jl +++ b/test/test_2nd_order_estimation.jl @@ -1,19 +1,30 @@ +using Test using MacroModelling import Turing -import ADTypes: AutoZygote -import Turing: NUTS, sample, logpdf +import Mooncake +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences +import Turing: NUTS, sample import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys + +using FlexiChains +include("test_helpers.jl") 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,:) @@ -32,7 +43,7 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_failure_loglikelihood) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) Turing.@addlogprob! get_loglikelihood(m, data, @@ -44,22 +55,39 @@ end Random.seed!(30) -n_samples = 500 +n_samples = 1000 + +samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :second_order, -Inf), NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) -samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :second_order, -Inf), NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") -println("Mean variable values (Zygote): $(mean(samps).nt.mean)") +sample_nuts = collect(values(FlexiChains.mean(samps); parameters_only = true)) -sample_nuts = mean(samps).nt.mean +@testset "Mooncake vs FiniteDifferences gradient (2nd order)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(FS2000, data, x, algorithm = :second_order), ADTypes.AutoMooncake(config = nothing), FS2000.parameter_values) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + 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, 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 @@ -104,10 +132,10 @@ sample_nuts = mean(samps).nt.mean # n_rounds = 6, # multithreaded = false) -# samps = MCMCChains.Chains(Pigeons.get_sample(pt)) +# samps = pigeons_flexichain(Pigeons.sample_array(pt), Pigeons.sample_names(pt)) -# println(mean(samps).nt.mean) +# println(collect(values(FlexiChains.mean(samps); parameters_only = true))) # Random.seed!(30) @@ -212,7 +240,7 @@ sample_nuts = mean(samps).nt.mean # 1 # @testset "Estimation results" begin # @test isapprox(sol.minimum, -1343.7491257498598, rtol = eps(Float32)) -# @test isapprox(mean(samps).nt.mean, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) +# @test isapprox(collect(values(FlexiChains.mean(samps); parameters_only = true)), [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) # end diff --git a/test/test_2nd_order_estimation_pigeons.jl b/test/test_2nd_order_estimation_pigeons.jl index 5246f80f1..ed6c96e60 100644 --- a/test/test_2nd_order_estimation_pigeons.jl +++ b/test/test_2nd_order_estimation_pigeons.jl @@ -2,19 +2,22 @@ using MacroModelling using Test import Turing import Pigeons -import Turing: logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL +include("test_helpers.jl") + 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,:) @@ -33,7 +36,7 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_failure_loglikelihood) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() Turing.@addlogprob! get_loglikelihood(m, @@ -45,7 +48,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 +69,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 +96,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..427cece2a 100644 --- a/test/test_3rd_order_estimation.jl +++ b/test/test_3rd_order_estimation.jl @@ -1,15 +1,26 @@ +using Test using MacroModelling import Turing -import ADTypes: AutoZygote -import Turing: NUTS, sample, logpdf, PG, IS +import Mooncake +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences +import Turing: NUTS, sample import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys + +using FlexiChains +include("test_helpers.jl") # 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] @@ -44,7 +55,7 @@ dists = [ ] Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_failure_loglikelihood) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) Turing.@addlogprob! get_loglikelihood(m, data, @@ -66,34 +77,50 @@ mode_estimateNM = Turing.maximum_a_posteriori(Caldara_et_al_2012_loglikelihood, Optim.NelderMead(), iterations = 100, # show_trace = true, - initial_params = Caldara_et_al_2012_estim.parameter_values) + initial_params = Turing.InitFromParams((; all_params = Caldara_et_al_2012_estim.parameter_values))) mode_estimateLBFGS = Turing.maximum_a_posteriori(Caldara_et_al_2012_loglikelihood, Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), - adtype = AutoZygote(), + adtype = AutoMooncake(; config=nothing), iterations = 100, # show_trace = true, - initial_params = mode_estimateNM.values) + initial_params = Turing.InitFromParams(mode_estimateNM)) -init_params = mode_estimateLBFGS.values |> collect +init_params = collect(mode_estimateLBFGS.params.data.all_params) -println("Mode variable values (L-BFGS): $init_params") +println("Mode variable values (L-BFGS): $(mode_estimateLBFGS.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 = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams(mode_estimateLBFGS)) -println("Mean variable values (Zygote): $(mean(samps).nt.mean)") +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") -sample_nuts = mean(samps).nt.mean +sample_nuts = collect(values(FlexiChains.mean(samps); parameters_only = true)) +@testset "Mooncake vs FiniteDifferences gradient (3rd order)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(Caldara_et_al_2012_estim, data, x, algorithm = :third_order), ADTypes.AutoMooncake(config = nothing), init_params) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + 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, 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) @@ -127,10 +154,10 @@ sample_nuts = mean(samps).nt.mean # n_rounds = 6, # multithreaded = false) -# samps = MCMCChains.Chains(Pigeons.get_sample(pt)) +# samps = pigeons_flexichain(Pigeons.sample_array(pt), Pigeons.sample_names(pt)) -# println(mean(samps).nt.mean) +# println(collect(values(FlexiChains.mean(samps); parameters_only = true))) # Random.seed!(30) @@ -235,7 +262,7 @@ sample_nuts = mean(samps).nt.mean # 1 # @testset "Estimation results" begin # @test isapprox(sol.minimum, -1343.7491257498598, rtol = eps(Float32)) -# @test isapprox(mean(samps).nt.mean, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) +# @test isapprox(collect(values(FlexiChains.mean(samps); parameters_only = true)), [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) # end diff --git a/test/test_3rd_order_estimation_pigeons.jl b/test/test_3rd_order_estimation_pigeons.jl index 164fa6a10..0a1b7a52b 100644 --- a/test/test_3rd_order_estimation_pigeons.jl +++ b/test/test_3rd_order_estimation_pigeons.jl @@ -2,15 +2,18 @@ using MacroModelling using Test import Turing import Pigeons -import Turing: logpdf, PG, IS -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL +include("test_helpers.jl") + # 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] @@ -45,7 +48,7 @@ dists = [ ] Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_failure_loglikelihood) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() Turing.@addlogprob! get_loglikelihood(m, @@ -57,7 +60,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 +85,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 +111,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_basic.jl b/test/test_basic.jl new file mode 100644 index 000000000..306c99853 --- /dev/null +++ b/test/test_basic.jl @@ -0,0 +1,3427 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +import SpecialFunctions: erfcinv +using AxisKeys, SparseArrays +import StatsPlots, Optim +import MatrixEquations +using Aqua +import LinearAlgebra as ℒ +using DelimitedFiles +using Dates + +Random.seed!(1234) + +include("functionality_tests.jl") + + plots = false + # test_higher_order = false + + function rbc_steady_state(params) + std_z, rho, delta, alpha, beta = params + + k_ss = ((1 / beta - 1 + delta) / alpha)^(1 / (alpha - 1)) + q_ss = k_ss^alpha + c_ss = q_ss - delta * k_ss + z_ss = 0.0 + + return [c_ss, k_ss, q_ss, z_ss] + end + + function make_counted_ss() + calls = Ref(0) + + function ss(params) + calls[] += 1 + return rbc_steady_state(params) + end + + return ss, calls + end + + @testset verbose = true "FS2000 precompile = true" begin + @model FS2000_pre precompile = true begin + dA[0] = exp(gam + z_e_a * e_a[x]) + log(m[0]) = (1 - rho) * log(mst) + rho * log(m[-1]) + z_e_m * e_m[x] + - P[0] / (c[1] * P[1] * m[0]) + bet * P[1] * (alp * exp( - alp * (gam + log(e[1]))) * k[0] ^ (alp - 1) * n[1] ^ (1 - alp) + (1 - del) * exp( - (gam + log(e[1])))) / (c[2] * P[2] * m[1])=0 + W[0] = l[0] / n[0] + - (psi / (1 - psi)) * (c[0] * P[0] / (1 - n[0])) + l[0] / n[0] = 0 + R[0] = P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ ( - alp) / W[0] + 1 / (c[0] * P[0]) - bet * P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) / (m[0] * l[0] * c[1] * P[1]) = 0 + c[0] + k[0] = exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) + (1 - del) * exp( - (gam + z_e_a * e_a[x])) * k[-1] + P[0] * c[0] = m[0] + m[0] - 1 + d[0] = l[0] + e[0] = exp(z_e_a * e_a[x]) + y[0] = k[-1] ^ alp * n[0] ^ (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) + gy_obs[0] = dA[0] * y[0] / y[-1] + gp_obs[0] = (P[0] / P[-1]) * m[-1] / dA[0] + log_gy_obs[0] = log(gy_obs[0]) + log_gp_obs[0] = log(gp_obs[0]) + end + + @parameters FS2000_pre silent = true precompile = true begin + alp = 0.356 + bet = 0.993 + gam = 0.0085 + mst = 1.0002 + rho = 0.129 + psi = 0.65 + del = 0.01 + z_e_a = 0.035449 + z_e_m = 0.008862 + end + + + SS_pre = get_SS(FS2000_pre, silent = true) + SS_change_pre = get_SS(FS2000_pre, parameters = :alp => 0.36, silent = true) + solution_pre = get_solution(FS2000_pre, silent = true) + solution_change_pre = get_solution(FS2000_pre, parameters = :alp => 0.35) + standard_deviation_pre = get_standard_deviation(FS2000_pre) + correlation_pre = get_correlation(FS2000_pre) + autocorrelation_pre = get_autocorrelation(FS2000_pre) + variance_decomposition_pre = get_variance_decomposition(FS2000_pre) + conditional_variance_decomposition_pre = get_conditional_variance_decomposition(FS2000_pre) + irf_pre = get_irf(FS2000_pre) + + @model FS2000 begin + dA[0] = exp(gam + z_e_a * e_a[x]) + log(m[0]) = (1 - rho) * log(mst) + rho * log(m[-1]) + z_e_m * e_m[x] + - P[0] / (c[1] * P[1] * m[0]) + bet * P[1] * (alp * exp( - alp * (gam + log(e[1]))) * k[0] ^ (alp - 1) * n[1] ^ (1 - alp) + (1 - del) * exp( - (gam + log(e[1])))) / (c[2] * P[2] * m[1])=0 + W[0] = l[0] / n[0] + - (psi / (1 - psi)) * (c[0] * P[0] / (1 - n[0])) + l[0] / n[0] = 0 + R[0] = P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ ( - alp) / W[0] + 1 / (c[0] * P[0]) - bet * P[0] * (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) / (m[0] * l[0] * c[1] * P[1]) = 0 + c[0] + k[0] = exp( - alp * (gam + z_e_a * e_a[x])) * k[-1] ^ alp * n[0] ^ (1 - alp) + (1 - del) * exp( - (gam + z_e_a * e_a[x])) * k[-1] + P[0] * c[0] = m[0] + m[0] - 1 + d[0] = l[0] + e[0] = exp(z_e_a * e_a[x]) + y[0] = k[-1] ^ alp * n[0] ^ (1 - alp) * exp( - alp * (gam + z_e_a * e_a[x])) + gy_obs[0] = dA[0] * y[0] / y[-1] + gp_obs[0] = (P[0] / P[-1]) * m[-1] / dA[0] + log_gy_obs[0] = log(gy_obs[0]) + log_gp_obs[0] = log(gp_obs[0]) + end + + @parameters FS2000 begin + alp = 0.356 + bet = 0.993 + gam = 0.0085 + mst = 1.0002 + rho = 0.129 + psi = 0.65 + del = 0.01 + z_e_a = 0.035449 + z_e_m = 0.008862 + end + + SS_nopre = get_SS(FS2000, silent = true) + SS_change_nopre = get_SS(FS2000, parameters = :alp => 0.36, silent = true) + solution_nopre = get_solution(FS2000, silent = true) + solution_change_nopre = get_solution(FS2000, parameters = :alp => 0.35) + standard_deviation_nopre = get_standard_deviation(FS2000) + correlation_nopre = get_correlation(FS2000) + autocorrelation_nopre = get_autocorrelation(FS2000) + variance_decomposition_nopre = get_variance_decomposition(FS2000) + conditional_variance_decomposition_nopre = get_conditional_variance_decomposition(FS2000) + irf_nopre = get_irf(FS2000) + + @test isapprox(SS_nopre, SS_pre) + @test isapprox(SS_change_nopre, SS_change_pre) + @test isapprox(solution_nopre, solution_pre) + @test isapprox(solution_change_nopre, solution_change_pre) + @test isapprox(standard_deviation_nopre, standard_deviation_pre) + @test isapprox(correlation_nopre, correlation_pre) + @test isapprox(autocorrelation_nopre, autocorrelation_pre) + @test isapprox(variance_decomposition_nopre, variance_decomposition_pre) + @test isapprox(conditional_variance_decomposition_nopre, conditional_variance_decomposition_pre) + @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)) + c[0] + k[0] = (1 - delta) * k[-1] + q[0] + q[0] = exp(z[0]) * k[-1]^alpha + z[0] = rho * z[-1] + std_z * eps_z[x] + end + + @parameters RBC_switch begin + std_z = 0.01 + rho = 0.2 + delta = 0.02 + alpha = 0.5 + beta = 0.95 + end + + custom_ss, custom_calls = make_counted_ss() + + inplace_calls = Ref(0) + function inplace_ss!(out, params) + inplace_calls[] += 1 + out .= rbc_steady_state(params) + return nothing + end + + # bad_calls = Ref(0) + # function bad_ss(params) + # bad_calls[] += 1 + # return zeros(4) + # end + + custom_calls[] = 0 + _ = get_steady_state(RBC_switch, steady_state_function = custom_ss) + @test custom_calls[] > 0 + + inplace_calls[] = 0 + inplace_result = get_steady_state(RBC_switch, steady_state_function = inplace_ss!) + @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) == expected_cache_length + + # @test_throws ArgumentError get_steady_state(RBC_switch, steady_state_function = bad_ss) + # @test bad_calls[] > 0 + + calls_before = custom_calls[] + _ = get_steady_state(RBC_switch, steady_state_function = nothing) + @test custom_calls[] == calls_before + @test isnothing(RBC_switch.functions.NSSS_custom) + + MacroModelling.set_custom_steady_state_function!(RBC_switch, custom_ss) + calls_before = custom_calls[] + _ = get_steady_state(RBC_switch) + @test custom_calls[] > calls_before + + MacroModelling.set_custom_steady_state_function!(RBC_switch, nothing) + calls_before = custom_calls[] + _ = get_steady_state(RBC_switch) + @test custom_calls[] == calls_before + end + + @testset "Macro steady state assignment" begin + macro_ss, macro_calls = make_counted_ss() + + @model RBC_macro_switch begin + 1 / c[0] = (beta / c[1]) * (alpha * exp(z[1]) * k[0]^(alpha - 1) + (1 - delta)) + c[0] + k[0] = (1 - delta) * k[-1] + q[0] + q[0] = exp(z[0]) * k[-1]^alpha + z[0] = rho * z[-1] + std_z * eps_z[x] + end + + @parameters RBC_macro_switch steady_state_function = macro_ss begin + std_z = 0.01 + rho = 0.2 + delta = 0.02 + alpha = 0.5 + beta = 0.95 + end + + @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_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.constants.nsss_solver.n_steps != 0 + + calls_before = macro_calls[] + _ = get_steady_state(RBC_macro_switch) + @test macro_calls[] == calls_before + end + + @testset verbose = true "Custom steady state function" begin + # Test custom steady state function with simple RBC model + @model RBC_custom_ss 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_custom_ss begin + std_z = 0.01 + ρ = 0.2 + δ = 0.02 + α = 0.5 + β = 0.95 + end + + # Get default steady state + default_ss = get_steady_state(RBC_custom_ss) + + # Define custom steady state function + # Variables in order: [:c, :k, :q, :z] (alphabetically sorted) + # Parameters in order: [:std_z, :ρ, :δ, :α, :β] (declaration order) + function my_steady_state_rbc(params) + std_z, ρ, δ, α, β = params + + # Analytical steady state for RBC model + k_ss = ((1/β - 1 + δ) / α)^(1/(α - 1)) + q_ss = k_ss^α + c_ss = q_ss - δ * k_ss + z_ss = 0.0 + + return [c_ss, k_ss, q_ss, z_ss] + end + + # Test custom function directly + custom_result = my_steady_state_rbc(RBC_custom_ss.parameter_values) + @test isapprox(custom_result, default_ss(:,:Steady_state)) + + # Get steady state with custom function + custom_ss = get_steady_state(RBC_custom_ss, steady_state_function = my_steady_state_rbc) + + # Compare with default (should be essentially the same) + @test isapprox(default_ss, custom_ss, rtol = 1e-10) + + # Test that model can be solved with custom SS function + irf_custom = get_irf(RBC_custom_ss, levels = true) + + # Steady state should still work after clearing + after_clear_ss = get_steady_state(RBC_custom_ss, steady_state_function = nothing) + @test isnothing(RBC_custom_ss.functions.NSSS_custom) + @test isapprox(default_ss, after_clear_ss, rtol = 1e-10) + + irf_after_clear = get_irf(RBC_custom_ss, levels = true) + @test isapprox(irf_after_clear, irf_custom, rtol = 1e-10) + + # Test with verbose option (internal function still available but not exported) + MacroModelling.set_custom_steady_state_function!(RBC_custom_ss, my_steady_state_rbc) + @test !isnothing(RBC_custom_ss.functions.NSSS_custom) + + + @model RBC_macro_ss 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_macro_ss steady_state_function = my_steady_state_rbc begin + std_z = 0.01 + ρ = 0.2 + δ = 0.02 + α = 0.5 + β = 0.95 + end + + # Verify macro-defined SS function is set + @test isapprox(RBC_macro_ss.functions.NSSS_custom(RBC_macro_ss.parameter_values), default_ss(:,:Steady_state), rtol = 1e-10) + + macro_ss = get_steady_state(RBC_macro_ss) + @test isapprox(default_ss, macro_ss, rtol = 1e-10) + + RBC_custom_ss = nothing + RBC_macro_ss = nothing + RBC_func_arg = nothing + end + + @testset verbose = true "Non-stochastic steady state guess" begin + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true guess = Dict(:alpha => .2, :beta => .99) begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve + R_ss = 1.0035 + # beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = 1.0025 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + + + # cap_share > 0 + # R_ss > 0 + # Pi_ss > 0 + # I_K_ratio > 0 + + # 0 < alpha < 1 + # 0 < beta < 1 + # 0 < delta < 1 + # 0 < Pibar + # 0 <= rhoz < 1 + # phi_pi > 0 + + # 0 < A < 1 + # 0 < k < 50 + # 0 < y < 10 + # 0 < c < 10 + end + + @test RBC_CME.constants.post_parameters_macro.guess == Dict(:alpha => .2, :beta => .99) + + @test get_steady_state(RBC_CME, verbose = true)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0025, 1.0035, 1.2081023824176236, 9.437411552284384, 1.4212969205027686, 1.0] + + RBC_CME = nothing + end + GC.gc() + + include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve.jl") + + global model = m + + @testset verbose = true "Custom steady state function with calibration equations and lead/lags" begin + # Test custom steady state function with RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve model + + # Get default steady state + default_ss = get_steady_state(model) + + function custom_steady_state(p::Vector{Float64}) + # 1. Unpack parameters + cap_share = p[1] + R_ss_target = p[2] + I_K_ratio = p[3] + phi_pi = p[4] + # std_eps = p[5] + # std_z_d = p[6] + Pi_real = p[7] + # rhoz = p[8] + + # 2. Solve for Deep Parameters and Rates + # Target R: log(R) = R_ss - 1 + R = exp(R_ss_target - 1.0) + + # Target Pi + Pi = R_ss_target - Pi_real + + # Euler Equation: 1 = beta * (R / Pi) + beta = Pi / R + Pibar = (R * beta) ^ (-1/phi_pi) * Pi + + # Ratios + # k / (4 * y) = cap_share + ky_ratio = 4.0 * cap_share + + # c / y = 1 - I_K_ratio + cy_ratio = 1.0 - I_K_ratio + + # Resource Constraint: 1 = c/y + delta * k/y + delta = (1.0 - cy_ratio) / ky_ratio + + # Euler Equation: 1 = beta * (alpha * y/k + 1 - delta) + # alpha = (k/y) * (1/beta - 1 + delta) + alpha = ky_ratio * ((1.0 / beta) - 1.0 + delta) + + # 3. Solve for Levels + A = 1.0 + z_delta = 1.0 + + # Production: y/k = k^(alpha-1) => k = (k/y)^(1/(1-alpha)) + k = (ky_ratio)^(1.0 / (1.0 - alpha)) + y = k^alpha + c = cy_ratio * y + + # Auxiliary variables + ZZ_avg = A + ZZ_avg_fut = A + log_ZZ_avg = 0.0 # log(1.0) + + c_logpdf = (-(abs2(c) + 1.8378770664093453) / 2) # normlogpdf + c_invcdf = (-erfcinv(2*(c - 1.0)) * 1.4142135623730951) # norminvcdf + + # 4. Return Vector + return [ + A, # 7 + Pi, # 4 + R, # 3 + ZZ_avg, # 8 + ZZ_avg_fut, # 9 + c, # 2 + c_invcdf, # 12 + c_logpdf, # 11 + k, # 5 + log_ZZ_avg, # 10 + y, # 1 + z_delta, # 6 + beta, # 14 (Derived) + Pibar, # 15 (Derived) + alpha, # 13 (Derived) + delta # 16 (Derived) + ] + end + + # Get steady state with custom function + custom_ss = get_steady_state(model, steady_state_function = custom_steady_state) + + # Compare key variables with default (should be essentially the same) + @test isapprox(default_ss, custom_ss, rtol = 1e-10) + + # Test that model can be solved with custom SS function + std_custom = get_std(model) + + # Steady state should still work after clearing + after_clear_ss = get_steady_state(model, steady_state_function = nothing) + @test isnothing(model.functions.NSSS_custom) + @test isapprox(default_ss, after_clear_ss, rtol = 1e-10) + + std_after_clear = get_std(model) + @test isapprox(std_after_clear, std_custom, rtol = 1e-10) + end + + @testset verbose = true "Provide parameters later" begin + include("models/Backus_Kehoe_Kydland_1992.jl") + + @model Backus_Kehoe_Kydland_1992_incomplete begin + for co in [H, F] + Y{co}[0] = ((LAMBDA{co}[0] * K{co}[-4]^theta{co} * N{co}[0]^(1-theta{co}))^(-nu{co}) + sigma{co} * Z{co}[-1]^(-nu{co}))^(-1/nu{co}) + + K{co}[0] = (1-delta{co})*K{co}[-1] + S{co}[0] + + X{co}[0] = for lag in (-4+1):0 phi{co} * S{co}[lag] end + + A{co}[0] = (1-eta{co}) * A{co}[-1] + N{co}[0] + + L{co}[0] = 1 - alpha{co} * N{co}[0] - (1-alpha{co})*eta{co} * A{co}[-1] + + U{co}[0] = (C{co}[0]^mu{co}*L{co}[0]^(1-mu{co}))^gamma{co} + + psi{co} * mu{co} / C{co}[0]*U{co}[0] = LGM[0] + + psi{co} * (1-mu{co}) / L{co}[0] * U{co}[0] * (-alpha{co}) = - LGM[0] * (1-theta{co}) / N{co}[0] * (LAMBDA{co}[0] * K{co}[-4]^theta{co}*N{co}[0]^(1-theta{co}))^(-nu{co})*Y{co}[0]^(1+nu{co}) + + for lag in 0:(4-1) + beta{co}^lag * LGM[lag]*phi{co} + end + + for lag in 1:4 + -beta{co}^lag * LGM[lag] * phi{co} * (1-delta{co}) + end = beta{co}^4 * LGM[+4] * theta{co} / K{co}[0] * (LAMBDA{co}[+4] * K{co}[0]^theta{co} * N{co}[+4]^(1-theta{co})) ^ (-nu{co})* Y{co}[+4]^(1+nu{co}) + + LGM[0] = beta{co} * LGM[+1] * (1+sigma{co} * Z{co}[0]^(-nu{co}-1)*Y{co}[+1]^(1+nu{co})) + + NX{co}[0] = (Y{co}[0] - (C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1]))/Y{co}[0] + end + + (LAMBDA{H}[0]-1) = rho{H}{H}*(LAMBDA{H}[-1]-1) + rho{H}{F}*(LAMBDA{F}[-1]-1) + Z_E{H} * E{H}[x] + + (LAMBDA{F}[0]-1) = rho{F}{F}*(LAMBDA{F}[-1]-1) + rho{F}{H}*(LAMBDA{H}[-1]-1) + Z_E{F} * E{F}[x] + + for co in [H,F] C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1] end = for co in [H,F] Y{co}[0] end + + dLGM[0] = LGM[1] / LGM[0] + + dLGM_ann[0] = for operator = :*, lag in -3:0 dLGM[lag] end + end + + @parameters Backus_Kehoe_Kydland_1992_incomplete begin + # K_ss = 11.0148 + # K[ss] = K_ss | beta + # K[ss] = 10 | beta + # F_H_ratio = 1 + K{F}[ss] / K{H}[ss] = F_H_ratio | beta{F} + K{H}[ss] = K_ss | beta{H} + + # beta = 0.99 + # mu = 0.34 + # gamma = -1.0 + # alpha = 1 + # eta = 0.5 + # theta = 0.36 + # nu = 3 + # sigma = 0.01 + # delta = 0.025 + # phi = 1/4 + # psi = 0.5 + + # Z_E = 0.00852 + + # rho{H}{H} = 0.906 + rho{F}{F} = rho{H}{H} + # rho{H}{F} = 0.088 + rho{F}{H} = rho{H}{F} + end + + pars = [ + "F_H_ratio" => 1.0, + "K_ss" => 11.0148, + "Z_E{F}" => 0.00852, + "Z_E{H}" => 0.00852, + "alpha{F}" => 1.0, + "alpha{H}" => 1.0, + "delta{F}" => 0.025, + "delta{H}" => 0.025, + "eta{F}" => 0.5, + "eta{H}" => 0.5, + "gamma{F}" => -1.0, + "gamma{H}" => -1.0, + "mu{F}" => 0.34, + "mu{H}" => 0.34, + "nu{F}" => 3.0, + "nu{H}" => 3.0, + "phi{F}" => 0.25, + "phi{H}" => 0.25, + "psi{F}" => 0.5, + "psi{H}" => 0.5, + "rho{H}{F}" => 0.088, + "rho{H}{H}" => 0.906, + "sigma{F}" => 0.01, + "sigma{H}" => 0.01, + "theta{F}" => 0.36, + "theta{H}" => 0.36 + ] + + cov1 = get_cov(Backus_Kehoe_Kydland_1992_incomplete, parameters = pars) + + cov2 = get_cov(Backus_Kehoe_Kydland_1992) + + @test cov1 ≈ cov2 + + + include("../models/Gali_2015_chapter_3_obc.jl") + + @model Gali_2015_chapter_3_obc_incomplete begin + W_real[0] = C[0] ^ σ * N[0] ^ φ + + Q[0] = β * (C[1] / C[0]) ^ (-σ) * Z[1] / Z[0] / Pi[1] + + R[0] = 1 / Q[0] + + Y[0] = A[0] * (N[0] / S[0]) ^ (1 - α) + + R[0] = Pi[1] * realinterest[0] + + R[0] = max(R̄ , 1 / β * Pi[0] ^ ϕᵖⁱ * (Y[0] / Y[ss]) ^ ϕʸ * exp(nu[0])) + + C[0] = Y[0] + + log(A[0]) = ρ_a * log(A[-1]) + std_a * eps_a[x] + + log(Z[0]) = ρ_z * log(Z[-1]) - std_z * eps_z[x] + + nu[0] = ρ_ν * nu[-1] + std_nu * eps_nu[x] + + MC[0] = W_real[0] / (S[0] * Y[0] * (1 - α) / N[0]) + + 1 = θ * Pi[0] ^ (ϵ - 1) + (1 - θ) * Pi_star[0] ^ (1 - ϵ) + + S[0] = (1 - θ) * Pi_star[0] ^ (( - ϵ) / (1 - α)) + θ * Pi[0] ^ (ϵ / (1 - α)) * S[-1] + + Pi_star[0] ^ (1 + ϵ * α / (1 - α)) = ϵ * x_aux_1[0] / x_aux_2[0] * (1 - τ) / (ϵ - 1) + + x_aux_1[0] = MC[0] * Y[0] * Z[0] * C[0] ^ (-σ) + β * θ * Pi[1] ^ (ϵ + α * ϵ / (1 - α)) * x_aux_1[1] + + x_aux_2[0] = Y[0] * Z[0] * C[0] ^ (-σ) + β * θ * Pi[1] ^ (ϵ - 1) * x_aux_2[1] + + log_y[0] = log(Y[0]) + + log_W_real[0] = log(W_real[0]) + + log_N[0] = log(N[0]) + + pi_ann[0] = 4 * log(Pi[0]) + + i_ann[0] = 4 * log(R[0]) + + r_real_ann[0] = 4 * log(realinterest[0]) + + M_real[0] = Y[0] / R[0] ^ η + + end + + @parameters Gali_2015_chapter_3_obc_incomplete begin + σ = 1 + + φ = 5 + + ϕᵖⁱ = 1.5 + + ϕʸ = 0.125 + + θ = 0.75 + + ρ_ν = 0.5 + + ρ_z = 0.5 + + ρ_a = 0.9 + + β = 0.99 + + η = 3.77 + + α = 0.25 + + ϵ = 9 + + τ = 0 + + std_a = .01 + + std_z = .05 + + std_nu = .0025 + + R > 1.0001 + end + + cov1 = get_cov(Gali_2015_chapter_3_obc_incomplete, parameters = :R̄ => 1.0) + + cov2 = get_cov(Gali_2015_chapter_3_obc) + + @test cov1 ≈ cov2 + + + @model Backus_Kehoe_Kydland_1992_incomplete2 begin + for co in [H, F] + Y{co}[0] = ((LAMBDA{co}[0] * K{co}[-4]^theta{co} * N{co}[0]^(1-theta{co}))^(-nu{co}) + sigma{co} * Z{co}[-1]^(-nu{co}))^(-1/nu{co}) + + K{co}[0] = (1-delta{co})*K{co}[-1] + S{co}[0] + + X{co}[0] = for lag in (-4+1):0 phi{co} * S{co}[lag] end + + A{co}[0] = (1-eta{co}) * A{co}[-1] + N{co}[0] + + L{co}[0] = 1 - alpha{co} * N{co}[0] - (1-alpha{co})*eta{co} * A{co}[-1] + + U{co}[0] = (C{co}[0]^mu{co}*L{co}[0]^(1-mu{co}))^gamma{co} + + psi{co} * mu{co} / C{co}[0]*U{co}[0] = LGM[0] + + psi{co} * (1-mu{co}) / L{co}[0] * U{co}[0] * (-alpha{co}) = - LGM[0] * (1-theta{co}) / N{co}[0] * (LAMBDA{co}[0] * K{co}[-4]^theta{co}*N{co}[0]^(1-theta{co}))^(-nu{co})*Y{co}[0]^(1+nu{co}) + + for lag in 0:(4-1) + beta{co}^lag * LGM[lag]*phi{co} + end + + for lag in 1:4 + -beta{co}^lag * LGM[lag] * phi{co} * (1-delta{co}) + end = beta{co}^4 * LGM[+4] * theta{co} / K{co}[0] * (LAMBDA{co}[+4] * K{co}[0]^theta{co} * N{co}[+4]^(1-theta{co})) ^ (-nu{co})* Y{co}[+4]^(1+nu{co}) + + LGM[0] = beta{co} * LGM[+1] * (1+sigma{co} * Z{co}[0]^(-nu{co}-1)*Y{co}[+1]^(1+nu{co})) + + NX{co}[0] = (Y{co}[0] - (C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1]))/Y{co}[0] + end + + (LAMBDA{H}[0]-1) = rho{H}{H}*(LAMBDA{H}[-1]-1) + rho{H}{F}*(LAMBDA{F}[-1]-1) + Z_E{H} * E{H}[x] + + (LAMBDA{F}[0]-1) = rho{F}{F}*(LAMBDA{F}[-1]-1) + rho{F}{H}*(LAMBDA{H}[-1]-1) + Z_E{F} * E{F}[x] + + for co in [H,F] C{co}[0] + X{co}[0] + Z{co}[0] - Z{co}[-1] end = for co in [H,F] Y{co}[0] end + + dLGM[0] = LGM[1] / LGM[0] + + dLGM_ann[0] = for operator = :*, lag in -3:0 dLGM[lag] end + end + + @parameters Backus_Kehoe_Kydland_1992_incomplete2 begin + K_ss = 11.0148 + kk = 1 + K{F}[ss] / K{H}[ss] = F_H_ratio | beta{F} + K{H}[ss] = K_ss | beta{H} + rho{F}{F} = rho{H}{H} + rho{F}{H} = rho{H}{F} + end + + params_full = get_parameters(Backus_Kehoe_Kydland_1992, values = true) + + std1 = get_std(Backus_Kehoe_Kydland_1992_incomplete2, parameters = params_full) + + std2 = get_std(Backus_Kehoe_Kydland_1992) + + @test std1 ≈ std2 + + + # Define a model with missing parameters + @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 + + # Only define some parameters + @parameters RBC begin + std_z = 0.01 + ρ = 0.2 + α = 0.5 + β = 0.95 + δ = 0.02 + end + + # Define a model with missing parameters + @model RBC_missing_provide 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 + + # Only define some parameters + @parameters RBC_missing_provide begin + std_z = 0.01 + ρ = 0.2 + # α => 0.5 + # β => 0.95 + # δ => 0.02 + end + + # Verify parameters are missing before providing them + @test has_missing_parameters(RBC_missing_provide) + + # Provide missing parameters and get IRF + irf_result = get_irf(RBC_missing_provide, parameters = [:α => 0.5, :β => 0.95, :δ => 0.02]) + + # Provide missing parameters and get IRF + irf_result_no_missing = get_irf(RBC) + + @test irf_result_no_missing ≈ irf_result + + # After providing parameters, they should no longer be missing + @test !has_missing_parameters(RBC_missing_provide) + + + @model m_incomplete begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-2] + eps_z[x+2] + eps_z_s[x]) + ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[ss]) / 4 + ZZ_avg_fut[0] = (A[0] + A[1] + A[2] + A[ss]) / 4 + log_ZZ_avg[0] = log(ZZ_avg[0]/ZZ_avg[ss]) + c_normlogpdf[0]= normlogpdf(c[0]) + c_norminvcdf[0]= norminvcdf(c[0]-1) + end + + + @parameters m_incomplete verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss + # R_ss = 1.0035 + # beta = .999 + + delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = R_ss - Pi_real + # Pi_real = 1/1000 + # Pibar = 1.0008 + + # phi_pi = 1.5 + # rhoz = 9 / 10 + # std_eps = .0068 + rho_z_delta = rhoz + # std_z_delta = .005 + end + + + param_defs = [ + :cap_share => 1.66, + :alpha => .157, + :std_z_delta => .005, + :phi_pi => 1.5, + :rhoz => 9 / 10, + :std_eps => .0068, + :Pi_real => 1/1000, + :Pibar => 1.0008, + :I_K_ratio => .15, + :delta => .0226, + :R_ss => 1.0035, + :beta => .999 + ] + + irfs_m_incomplete = get_irf(m_incomplete, parameters = param_defs) + + std_m_incomplete = get_std(m_incomplete) + + + + @model m begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-2] + eps_z[x+2] + eps_z_s[x]) + ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[ss]) / 4 + ZZ_avg_fut[0] = (A[0] + A[1] + A[2] + A[ss]) / 4 + log_ZZ_avg[0] = log(ZZ_avg[0]/ZZ_avg[ss]) + c_normlogpdf[0]= normlogpdf(c[0]) + c_norminvcdf[0]= norminvcdf(c[0]-1) + end + + @parameters m verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss + R_ss = 1.0035 + # beta = .999 + + delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = R_ss - Pi_real + Pi_real = 1/1000 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = 9 / 10 + std_eps = .0068 + rho_z_delta = rhoz + std_z_delta = .005 + end + + irfs_m = get_irf(m) + + std_m = get_std(m) + + + @test irfs_m ≈ irfs_m_incomplete + + @test std_m(:,:Standard_deviation) ≈ std_m_incomplete(:,:Standard_deviation) + + + + include("../models/NAWM_EAUS_2008.jl") + + + @model NAWM_EAUS_2008_incomplete begin + EA_R[0] ^ 4 - 1 = EA_PHIRR * (EA_R[-1] ^ 4 - 1) + (1 - EA_PHIRR) * (EA_RRSTAR ^ 4 * EA_PI4TARGET - 1 + EA_PHIRPI * (EA_PIC4[0] - EA_PI4TARGET)) + EA_PHIRGY * (EA_Y[0] / EA_Y[-1] - 1) + σ_EA_R * EA_EPSR[x] + + US_R[0] ^ 4 - 1 = US_PHIRR * (US_R[-1] ^ 4 - 1) + (1 - US_PHIRR) * (US_RRSTAR ^ 4 * US_PI4TARGET - 1 + US_PHIRPI * (US_PIC4[0] - US_PI4TARGET)) + US_PHIRGY * (US_Y[0] / US_Y[-1] - 1) + σ_US_R * US_EPSR[x] + + EA_UTILI[0] = 1 / (1 - EA_SIGMA) * (EA_CI[0] - EA_KAPPA * EA_CI[-1]) ^ (1 - EA_SIGMA) - 1 / (1 + EA_ZETA) * EA_NI[0] ^ (1 + EA_ZETA) + EA_BETA * EA_UTILI[1] + + EA_LAMBDAI[0] * (1 + EA_TAUC[0] + EA_GAMMAVI[0] + EA_VI[0] * EA_GAMMAVIDER[0]) = (EA_CI[0] - EA_KAPPA * EA_CI[-1]) ^ (-EA_SIGMA) + + EA_R[0] = EA_LAMBDAI[0] * EA_BETA ^ (-1) / EA_LAMBDAI[1] * EA_PIC[1] + + EA_GAMMAVIDER[0] * EA_VI[0] ^ 2 = 1 - EA_BETA * EA_LAMBDAI[1] / (EA_LAMBDAI[0] * EA_PIC[1]) + + EA_VI[0] = EA_CI[0] * (1 + EA_TAUC[0]) / EA_MI[0] + + EA_GAMMAVI[0] = EA_VI[0] * EA_GAMMAV1 + EA_GAMMAV2 / EA_VI[0] - 2 * (EA_GAMMAV1 * EA_GAMMAV2) ^ 0.5 + + EA_GAMMAVIDER[0] = EA_GAMMAV1 - EA_GAMMAV2 * EA_VI[0] ^ (-2) + + EA_KI[0] = (1 - EA_DELTA) * EA_KI[-1] + (1 - EA_GAMMAI[-1]) * EA_II[-1] + + EA_GAMMAI[0] = EA_GAMMAI1 / 2 * (EA_II[0] / EA_II[-1] - 1) ^ 2 + + EA_GAMMAIDER[0] = EA_GAMMAI1 * (EA_II[0] / EA_II[-1] - 1) / EA_II[-1] + + EA_GAMMAU[0] = ((EA_DELTA + EA_BETA ^ (-1) - 1) * EA_QBAR - EA_DELTA * EA_TAUKBAR * EA_PIBAR) / (EA_PIBAR * (1 - EA_TAUKBAR)) * (EA_U[0] - 1) + EA_GAMMAU2 / 2 * (EA_U[0] - 1) ^ 2 + + EA_GAMMAUDER[0] = ((EA_DELTA + EA_BETA ^ (-1) - 1) * EA_QBAR - EA_DELTA * EA_TAUKBAR * EA_PIBAR) / (EA_PIBAR * (1 - EA_TAUKBAR)) + (EA_U[0] - 1) * EA_GAMMAU2 + + EA_RK[0] = EA_GAMMAUDER[0] * EA_PI[0] + + EA_PI[0] = EA_Q[0] * (1 - EA_GAMMAI[0] - EA_II[0] * EA_GAMMAIDER[0]) + EA_BETA * EA_LAMBDAI[1] / EA_LAMBDAI[0] * EA_Q[1] * EA_GAMMAIDER[1] * EA_II[1] ^ 2 / EA_II[0] + + EA_Q[0] = EA_BETA * EA_LAMBDAI[1] / EA_LAMBDAI[0] * ((1 - EA_TAUK[1]) * (EA_RK[1] * EA_U[1] - EA_GAMMAU[1] * EA_PI[1]) + EA_PI[1] * EA_DELTA * EA_TAUK[1] + (1 - EA_DELTA) * EA_Q[1]) + + EA_WITILDE[0] ^ (1 + EA_ZETA * EA_ETAI) = EA_ETAI / (EA_ETAI - 1) * EA_FI[0] / EA_GI[0] + + EA_FI[0] = EA_WI[0] ^ ((1 + EA_ZETA) * EA_ETAI) * EA_NDI[0] ^ (1 + EA_ZETA) + EA_BETA * EA_XII * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHII * EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)))) ^ ((1 + EA_ZETA) * EA_ETAI) * EA_FI[1] + + EA_GI[0] = EA_NDI[0] * EA_LAMBDAI[0] * (1 - EA_TAUN[0] - EA_TAUWH[0]) * EA_WI[0] ^ EA_ETAI + EA_BETA * EA_XII * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHII * EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)))) ^ (EA_ETAI - 1) * EA_GI[1] + + EA_WI[0] ^ (1 - EA_ETAI) = (1 - EA_XII) * EA_WITILDE[0] ^ (1 - EA_ETAI) + EA_XII * EA_WI[-1] ^ (1 - EA_ETAI) * (EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)) * EA_PIC[-1] ^ EA_CHII / EA_PIC[0]) ^ (1 - EA_ETAI) + + EA_UTILJ[0] = 1 / (1 - EA_SIGMA) * (EA_CJ[0] - EA_KAPPA * EA_CJ[-1]) ^ (1 - EA_SIGMA) - 1 / (1 + EA_ZETA) * EA_NJ[0] ^ (1 + EA_ZETA) + EA_BETA * EA_UTILJ[1] + + EA_CJ[0] * (1 + EA_TAUC[0] + EA_GAMMAVJ[0]) + EA_MJ[0] = EA_NJ[0] * (1 - EA_TAUN[0] - EA_TAUWH[0]) * EA_WJ[0] + EA_TRJ[0] - EA_TJ[0] + EA_MJ[-1] * EA_PIC[0] ^ (-1) + + EA_LAMBDAJ[0] * (1 + EA_TAUC[0] + EA_GAMMAVJ[0] + EA_VJ[0] * EA_GAMMAVJDER[0]) = (EA_CJ[0] - EA_KAPPA * EA_CJ[-1]) ^ (-EA_SIGMA) + + EA_GAMMAVJDER[0] * EA_VJ[0] ^ 2 = 1 - EA_BETA * EA_LAMBDAJ[1] / (EA_PIC[1] * EA_LAMBDAJ[0]) + + EA_VJ[0] = (1 + EA_TAUC[0]) * EA_CJ[0] / EA_MJ[0] + + EA_GAMMAVJ[0] = EA_GAMMAV1 * EA_VJ[0] + EA_GAMMAV2 / EA_VJ[0] - 2 * (EA_GAMMAV1 * EA_GAMMAV2) ^ 0.5 + + EA_GAMMAVJDER[0] = EA_GAMMAV1 - EA_GAMMAV2 * EA_VJ[0] ^ (-2) + + EA_WJTILDE[0] ^ (1 + EA_ZETA * EA_ETAJ) = EA_ETAJ / (EA_ETAJ - 1) * EA_FJ[0] / EA_GJ[0] + + EA_FJ[0] = EA_WJ[0] ^ ((1 + EA_ZETA) * EA_ETAJ) * EA_NDJ[0] ^ (1 + EA_ZETA) + EA_BETA * EA_XIJ * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHIJ * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)))) ^ ((1 + EA_ZETA) * EA_ETAJ) * EA_FJ[1] + + EA_GJ[0] = EA_NDJ[0] * (1 - EA_TAUN[0] - EA_TAUWH[0]) * EA_LAMBDAJ[0] * EA_WJ[0] ^ EA_ETAJ + EA_BETA * EA_XIJ * (EA_PIC[1] / (EA_PIC[0] ^ EA_CHIJ * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)))) ^ (EA_ETAJ - 1) * EA_GJ[1] + + EA_WJ[0] ^ (1 - EA_ETAJ) = (1 - EA_XIJ) * EA_WJTILDE[0] ^ (1 - EA_ETAJ) + EA_XIJ * EA_WJ[-1] ^ (1 - EA_ETAJ) * (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)) * EA_PIC[-1] ^ EA_CHIJ / EA_PIC[0]) ^ (1 - EA_ETAJ) + + EA_YS[0] = EA_Z[0] * EA_KD[0] ^ EA_ALPHA * EA_ND[0] ^ (1 - EA_ALPHA) - EA_PSIBAR + + EA_RK[0] = EA_ALPHA * (EA_YS[0] + EA_PSIBAR) / EA_KD[0] * EA_MC[0] + + EA_MC[0] = 1 / (EA_Z[0] * EA_ALPHA ^ EA_ALPHA * (1 - EA_ALPHA) ^ (1 - EA_ALPHA)) * EA_RK[0] ^ EA_ALPHA * ((1 + EA_TAUWF[0]) * EA_W[0]) ^ (1 - EA_ALPHA) + + EA_NDI[0] = EA_ND[0] * (1 - EA_OMEGA) * (EA_WI[0] / EA_W[0]) ^ (-EA_ETA) + + EA_NDJ[0] = EA_ND[0] * EA_OMEGA * (EA_WJ[0] / EA_W[0]) ^ (-EA_ETA) + + EA_ND[0] ^ (1 - 1 / EA_ETA) = (1 - EA_OMEGA) ^ (1 / EA_ETA) * EA_NDI[0] ^ (1 - 1 / EA_ETA) + EA_OMEGA ^ (1 / EA_ETA) * EA_NDJ[0] ^ (1 - 1 / EA_ETA) + + EA_D[0] = EA_Y[0] * EA_PY[0] - EA_RK[0] * EA_KD[0] - EA_ND[0] * (1 + EA_TAUWF[0]) * EA_W[0] + + EA_PHTILDE[0] / EA_PH[0] = EA_THETA / (EA_THETA - 1) * EA_FH[0] / EA_GH[0] + + EA_FH[0] = EA_MC[0] * EA_H[0] + EA_LAMBDAI[1] * EA_BETA * EA_XIH / EA_LAMBDAI[0] * (EA_PIH[1] / (EA_PIH[0] ^ EA_CHIH * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)))) ^ EA_THETA * EA_FH[1] + + EA_GH[0] = EA_PH[0] * EA_H[0] + EA_LAMBDAI[1] * EA_BETA * EA_XIH / EA_LAMBDAI[0] * (EA_PIH[1] / (EA_PIH[0] ^ EA_CHIH * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)))) ^ (EA_THETA - 1) * EA_GH[1] + + EA_PH[0] ^ (1 - EA_THETA) = (1 - EA_XIH) * EA_PHTILDE[0] ^ (1 - EA_THETA) + EA_XIH * (EA_PH[-1] / EA_PIC[0]) ^ (1 - EA_THETA) * (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)) * EA_PIH[-1] ^ EA_CHIH) ^ (1 - EA_THETA) + + EA_PIH[0] = EA_PIC[0] * EA_PH[0] / EA_PH[-1] + + US_PIMTILDE[0] / US_PIM[0] = EA_THETA / (EA_THETA - 1) * EA_FX[0] / EA_GX[0] + + EA_FX[0] = EA_MC[0] * US_SIZE / EA_SIZE * US_IM[0] + EA_LAMBDAI[1] * EA_BETA * EA_XIX / EA_LAMBDAI[0] * (US_PIIM[1] / (US_PIIM[0] ^ EA_CHIX * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIX)))) ^ EA_THETA * EA_FX[1] + + EA_GX[0] = US_IM[0] * US_SIZE * US_PIM[0] * EAUS_RER[0] / EA_SIZE + EA_LAMBDAI[1] * EA_BETA * EA_XIX / EA_LAMBDAI[0] * (US_PIIM[1] / (US_PIIM[0] ^ EA_CHIX * EA_PI4TARGET ^ (0.25 * (1 - EA_CHIX)))) ^ (EA_THETA - 1) * EA_GX[1] + + US_PIM[0] ^ (1 - EA_THETA) = (1 - EA_XIX) * US_PIMTILDE[0] ^ (1 - EA_THETA) + EA_XIX * (US_PIM[-1] / US_PIC[0]) ^ (1 - EA_THETA) * (US_PIIM[-1] ^ EA_CHIX * US_PI4TARGET ^ (0.25 * (1 - EA_CHIH))) ^ (1 - EA_THETA) + + US_PIIM[0] = US_PIC[0] * US_PIM[0] / US_PIM[-1] + + EAUS_RER[0] = EA_RER[0] / US_RER + + EA_QC[0] ^ ((EA_MUC - 1) / EA_MUC) = EA_NUC ^ (1 / EA_MUC) * EA_HC[0] ^ (1 - 1 / EA_MUC) + (1 - EA_NUC) ^ (1 / EA_MUC) * ((1 - EA_GAMMAIMC[0]) * EA_IMC[0]) ^ (1 - 1 / EA_MUC) + + 1 = EA_NUC * EA_PH[0] ^ (1 - EA_MUC) + (1 - EA_NUC) * (EA_PIM[0] / EA_GAMMAIMCDAG[0]) ^ (1 - EA_MUC) + + EA_HC[0] = EA_QC[0] * EA_NUC * EA_PH[0] ^ (-EA_MUC) + + EA_GAMMAIMC[0] = EA_GAMMAIMC1 / 2 * (EA_IMC[0] / EA_QC[0] / (EA_IMC[-1] / EA_QC[-1]) - 1) ^ 2 + + EA_GAMMAIMCDAG[0] = 1 - EA_GAMMAIMC[0] - EA_IMC[0] * EA_GAMMAIMC1 * (EA_IMC[0] / EA_QC[0] / (EA_IMC[-1] / EA_QC[-1]) - 1) / EA_QC[0] / (EA_IMC[-1] / EA_QC[-1]) + + EA_QI[0] ^ ((EA_MUI - 1) / EA_MUI) = EA_NUI ^ (1 / EA_MUI) * EA_HI[0] ^ (1 - 1 / EA_MUI) + (1 - EA_NUI) ^ (1 / EA_MUI) * ((1 - EA_GAMMAIMI[0]) * EA_IMI[0]) ^ (1 - 1 / EA_MUI) + + EA_PI[0] ^ (1 - EA_MUI) = EA_NUI * EA_PH[0] ^ (1 - EA_MUI) + (1 - EA_NUI) * (EA_PIM[0] / EA_GAMMAIMIDAG[0]) ^ (1 - EA_MUI) + + EA_HI[0] = EA_QI[0] * EA_NUI * (EA_PH[0] / EA_PI[0]) ^ (-EA_MUI) + + EA_GAMMAIMI[0] = EA_GAMMAIMI1 / 2 * (EA_IMI[0] / EA_QI[0] / (EA_IMI[-1] / EA_QI[-1]) - 1) ^ 2 + + EA_GAMMAIMIDAG[0] = 1 - EA_GAMMAIMI[0] - EA_IMI[0] * EA_GAMMAIMI1 * (EA_IMI[0] / EA_QI[0] / (EA_IMI[-1] / EA_QI[0]) - 1) / EA_QI[0] / (EA_IMI[-1] / EA_QI[-1]) + + EA_PH[-1] * EA_G[-1] + EA_TR[-1] + EA_B[-1] * EA_PIC[-1] ^ (-1) + EA_PIC[-1] ^ (-1) * EA_M[-2] = EA_TAUC[-1] * EA_C[-1] + (EA_TAUN[-1] + EA_TAUWH[-1]) * (EA_WI[-1] * EA_NDI[-1] + EA_WJ[-1] * EA_NDJ[-1]) + EA_TAUWF[-1] * EA_W[-1] * EA_ND[-1] + EA_TAUK[-1] * (EA_RK[-1] * EA_U[-1] - (EA_DELTA + EA_GAMMAU[-1]) * EA_PI[-1]) * EA_K[-1] + EA_TAUD[-1] * EA_D[-1] + EA_T[-1] + EA_R[-1] ^ (-1) * EA_B[0] + EA_M[-1] + + EA_PH[0] * EA_G[0] = EA_GY[0] * EA_PYBAR * EA_YBAR + + EA_TR[0] = EA_YBAR * EA_PYBAR * EA_TRY[0] + + EA_T[0] / (EA_PYBAR * EA_YBAR) = EA_PHITB * (EA_B[0] / (EA_PYBAR * EA_YBAR) - EA_BYTARGET) + + EA_TI[0] = EA_T[0] * EA_UPSILONT + + EA_TRI[0] = EA_TR[0] * EA_UPSILONTR + + EA_PIC4[0] = EA_PIC[0] * EA_PIC[-1] * EA_PIC[-2] * EA_PIC[-3] + + EA_RR[0] - 1 = EA_R[0] / EA_PIC[1] - 1 + + EA_C[0] = EA_CI[0] * (1 - EA_OMEGA) + EA_CJ[0] * EA_OMEGA + + EA_M[0] = EA_MI[0] * (1 - EA_OMEGA) + EA_MJ[0] * EA_OMEGA + + EA_K[0] = EA_KI[0] * (1 - EA_OMEGA) + + EA_I[0] = EA_II[0] * (1 - EA_OMEGA) + + EA_TRJ[0] = EA_TR[0] * 1 / EA_OMEGA - EA_TRI[0] * (1 - EA_OMEGA) / EA_OMEGA + + EA_TJ[0] = EA_T[0] * 1 / EA_OMEGA - EA_TI[0] * (1 - EA_OMEGA) / EA_OMEGA + + EA_GAMMAV[0] = EA_GAMMAVI[0] * EA_CI[0] * (1 - EA_OMEGA) + EA_GAMMAVJ[0] * EA_CJ[0] * EA_OMEGA + + EA_NI[0] = EA_NDI[0] * EA_SI[0] + + EA_SI[0] = (1 - EA_XII) * (EA_WITILDE[0] / EA_WI[0]) ^ (-EA_ETAI) + EA_XII * (EA_WI[-1] / EA_WI[0]) ^ (-EA_ETAI) * (EA_PIC[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHII)) * EA_PIC[-1] ^ EA_CHII)) ^ EA_ETAI * EA_SI[-1] + + EA_NJ[0] = EA_NDJ[0] * EA_SJ[0] + + EA_SJ[0] = (1 - EA_XIJ) * (EA_WJTILDE[0] / EA_WJ[0]) ^ (-EA_ETAJ) + EA_XIJ * (EA_WJ[-1] / EA_WJ[0]) ^ (-EA_ETAJ) * (EA_PIC[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIJ)) * EA_PIC[-1] ^ EA_CHIJ)) ^ EA_ETAJ * EA_SJ[-1] + + EA_U[0] * EA_K[0] = EA_KD[0] + + EA_YS[0] = EA_H[0] * EA_SH[0] + US_IM[0] * US_SIZE * EA_SX[0] / EA_SIZE + + EA_H[0] = EA_G[0] + EA_HC[0] + EA_HI[0] + + EA_IM[0] = EA_IMC[0] + EA_IMI[0] + + EA_SH[0] = (1 - EA_XIH) * (EA_PHTILDE[0] / EA_PH[0]) ^ (-EA_THETA) + EA_XIH * (EA_PIH[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)) * EA_PIH[-1] ^ EA_CHIH)) ^ EA_THETA * EA_SH[-1] + + EA_SX[0] = (1 - EA_XIX) * (US_PIMTILDE[0] / US_PIM[0]) ^ (-EA_THETA) + EA_XIX * (US_PIIM[0] / (EA_PI4TARGET ^ (0.25 * (1 - EA_CHIH)) * US_PIIM[-1] ^ EA_CHIX)) ^ EA_THETA * EA_SX[-1] + + EA_QC[0] = EA_C[0] + EA_GAMMAV[0] + + EA_QI[0] = EA_I[0] + EA_GAMMAU[0] * EA_K[0] + + EA_Y[0] * EA_PY[0] = US_IM[0] * US_SIZE * US_PIM[0] * EAUS_RER[0] / EA_SIZE + EA_PH[0] * EA_G[0] + EA_QC[0] + EA_PI[0] * EA_QI[0] - EA_PIM[0] * ((1 - EA_GAMMAIMC[0]) * EA_IMC[0] / EA_GAMMAIMCDAG[0] + (1 - EA_GAMMAIMI[0]) * EA_IMI[0] / EA_GAMMAIMIDAG[0]) + + EA_Y[0] = EA_YS[0] + + log(EA_Z[0]) = (1 - EA_RHOZ) * log(EA_ZBAR) + EA_RHOZ * log(EA_Z[-1]) + σ_EA_Z * EA_EPSZ[x] + + EA_GY[0] = (1 - EA_RHOG) * EA_GYBAR + EA_RHOG * EA_GY[-1] + σ_EA_G * EA_EPSG[x] + + EA_TRY[0] = (1 - EA_RHOTR) * EA_TRYBAR + EA_RHOTR * EA_TRY[-1] + σ_EA_TR * EA_EPSTR[x] + + EA_TAUC[0] = (1 - EA_RHOTAUC) * EA_TAUCBAR + EA_TAUC[-1] * EA_RHOTAUC + σ_EA_TAUC * EA_EPSTAUC[x] + + EA_TAUD[0] = (1 - EA_RHOTAUD) * EA_TAUDBAR + EA_TAUD[-1] * EA_RHOTAUD + σ_EA_TAUD * EA_EPSTAUD[x] + + EA_TAUK[0] = EA_TAUKBAR * (1 - EA_RHOTAUK) + EA_TAUK[-1] * EA_RHOTAUK + σ_EA_TAUK * EA_EPSTAUK[x] + + EA_TAUN[0] = (1 - EA_RHOTAUN) * EA_TAUNBAR + EA_TAUN[-1] * EA_RHOTAUN + σ_EA_TAUN * EA_EPSTAUN[x] + + EA_TAUWH[0] = (1 - EA_RHOTAUWH) * EA_TAUWHBAR + EA_TAUWH[-1] * EA_RHOTAUWH + σ_EA_TAUWH * EA_EPSTAUWH[x] + + EA_TAUWF[0] = (1 - EA_RHOTAUWF) * EA_TAUWFBAR + EA_TAUWF[-1] * EA_RHOTAUWF + σ_EA_TAUWF * EA_EPSTAUWF[x] + + EA_CY[0] = EA_C[0] / (EA_Y[0] * EA_PY[0]) + + EA_IY[0] = EA_PI[0] * EA_I[0] / (EA_Y[0] * EA_PY[0]) + + EA_IMY[0] = EA_PIM[0] * EA_IM[0] / (EA_Y[0] * EA_PY[0]) + + EA_IMCY[0] = EA_IMC[0] * EA_PIM[0] / (EA_Y[0] * EA_PY[0]) + + EA_IMIY[0] = EA_PIM[0] * EA_IMI[0] / (EA_Y[0] * EA_PY[0]) + + EA_BY[0] = EA_B[0] / (EA_PYBAR * EA_YBAR) + + EA_TY[0] = EA_T[0] / (EA_PYBAR * EA_YBAR) + + EA_YGAP[0] = EA_Y[0] / EA_YBAR - 1 + + EA_YGROWTH[0] = EA_Y[0] / EA_Y[-1] + + EA_YSHARE[0] = EA_Y[0] * EA_PY[0] * EA_SIZE / EA_RER[0] / (EA_Y[0] * EA_PY[0] * EA_SIZE / EA_RER[0] + US_Y[0] * US_SIZE * US_PY[0] / US_RER) + + EA_EPSILONM[0] = ( - 0.125) / (EA_R[0] * (EA_R[0] + EA_R[0] * EA_GAMMAV2 - 1)) + + US_UTILI[0] = 1 / (1 - US_SIGMA) * (US_CI[0] - US_KAPPA * US_CI[-1]) ^ (1 - US_SIGMA) - 1 / (1 + US_ZETA) * US_NI[0] ^ (1 + US_ZETA) + US_BETA * US_UTILI[1] + + US_LAMBDAI[0] * (1 + US_TAUC[0] + US_GAMMAVI[0] + US_VI[0] * US_GAMMAVIDER[0]) = (US_CI[0] - US_KAPPA * US_CI[-1]) ^ (-US_SIGMA) + + US_R[0] = US_LAMBDAI[0] * US_BETA ^ (-1) / US_LAMBDAI[1] * US_PIC[1] + + US_GAMMAVIDER[0] * US_VI[0] ^ 2 = 1 - US_BETA * US_LAMBDAI[1] / (US_LAMBDAI[0] * US_PIC[1]) + + US_VI[0] = US_CI[0] * (1 + US_TAUC[0]) / US_MI[0] + + US_GAMMAVI[0] = US_VI[0] * US_GAMMAV1 + US_GAMMAV2 / US_VI[0] - 2 * (US_GAMMAV1 * US_GAMMAV2) ^ 0.5 + + US_GAMMAVIDER[0] = US_GAMMAV1 - US_GAMMAV2 * US_VI[0] ^ (-2) + + US_KI[0] = (1 - US_DELTA) * US_KI[-1] + (1 - US_GAMMAI[-1]) * US_II[-1] + + US_GAMMAI[0] = US_GAMMAI1 / 2 * (US_II[0] / US_II[-1] - 1) ^ 2 + + US_GAMMAIDER[0] = US_GAMMAI1 * (US_II[0] / US_II[-1] - 1) / US_II[-1] + + US_GAMMAU[0] = ((US_DELTA + US_BETA ^ (-1) - 1) * US_QBAR - US_DELTA * US_TAUKBAR * US_PIBAR) / (US_PIBAR * (1 - US_TAUKBAR)) * (US_U[0] - 1) + US_GAMMAU2 / 2 * (US_U[0] - 1) ^ 2 + + US_GAMMAUDER[0] = ((US_DELTA + US_BETA ^ (-1) - 1) * US_QBAR - US_DELTA * US_TAUKBAR * US_PIBAR) / (US_PIBAR * (1 - US_TAUKBAR)) + (US_U[0] - 1) * US_GAMMAU2 + + US_RK[0] = US_GAMMAUDER[0] * US_PI[0] + + US_PI[0] = US_Q[0] * (1 - US_GAMMAI[0] - US_II[0] * US_GAMMAIDER[0]) + US_BETA * US_LAMBDAI[1] / US_LAMBDAI[0] * US_Q[1] * US_GAMMAIDER[1] * US_II[1] ^ 2 / US_II[0] + + US_Q[0] = US_BETA * US_LAMBDAI[1] / US_LAMBDAI[0] * ((1 - US_TAUK[1]) * (US_RK[1] * US_U[1] - US_GAMMAU[1] * US_PI[1]) + US_PI[1] * US_DELTA * US_TAUK[1] + (1 - US_DELTA) * US_Q[1]) + + US_WITILDE[0] ^ (1 + US_ZETA * US_ETAI) = US_ETAI / (US_ETAI - 1) * US_FI[0] / US_GI[0] + + US_FI[0] = US_WI[0] ^ ((1 + US_ZETA) * US_ETAI) * US_NDI[0] ^ (1 + US_ZETA) + US_BETA * US_XII * (US_PIC[1] / (US_PIC[0] ^ US_CHII * US_PI4TARGET ^ (0.25 * (1 - US_CHII)))) ^ ((1 + US_ZETA) * US_ETAI) * US_FI[1] + + US_GI[0] = US_NDI[0] * US_LAMBDAI[0] * (1 - US_TAUN[0] - US_TAUWH[0]) * US_WI[0] ^ US_ETAI + US_BETA * US_XII * (US_PIC[1] / (US_PIC[0] ^ US_CHII * US_PI4TARGET ^ (0.25 * (1 - US_CHII)))) ^ (US_ETAI - 1) * US_GI[1] + + US_WI[0] ^ (1 - US_ETAI) = (1 - US_XII) * US_WITILDE[0] ^ (1 - US_ETAI) + US_XII * US_WI[-1] ^ (1 - US_ETAI) * (US_PI4TARGET ^ (0.25 * (1 - US_CHII)) * US_PIC[-1] ^ US_CHII / US_PIC[0]) ^ (1 - US_ETAI) + + US_UTILJ[0] = 1 / (1 - US_SIGMA) * (US_CJ[0] - US_KAPPA * US_CJ[-1]) ^ (1 - US_SIGMA) - 1 / (1 + US_ZETA) * US_NJ[0] ^ (1 + US_ZETA) + US_BETA * US_UTILJ[1] + + US_CJ[0] * (1 + US_TAUC[0] + US_GAMMAVJ[0]) + US_MJ[0] = US_NJ[0] * (1 - US_TAUN[0] - US_TAUWH[0]) * US_WJ[0] + US_TRJ[0] - US_TJ[0] + US_MJ[-1] * US_PIC[0] ^ (-1) + + US_LAMBDAJ[0] * (1 + US_TAUC[0] + US_GAMMAVJ[0] + US_VJ[0] * US_GAMMAVJDER[0]) = (US_CJ[0] - US_KAPPA * US_CJ[-1]) ^ (-US_SIGMA) + + US_GAMMAVJDER[0] * US_VJ[0] ^ 2 = 1 - US_BETA * US_LAMBDAJ[1] / (US_PIC[1] * US_LAMBDAJ[0]) + + US_VJ[0] = (1 + US_TAUC[0]) * US_CJ[0] / US_MJ[0] + + US_GAMMAVJ[0] = US_GAMMAV1 * US_VJ[0] + US_GAMMAV2 / US_VJ[0] - 2 * (US_GAMMAV1 * US_GAMMAV2) ^ 0.5 + + US_GAMMAVJDER[0] = US_GAMMAV1 - US_GAMMAV2 * US_VJ[0] ^ (-2) + + US_WJTILDE[0] ^ (1 + US_ZETA * US_ETAJ) = US_ETAJ / (US_ETAJ - 1) * US_FJ[0] / US_GJ[0] + + US_FJ[0] = US_WJ[0] ^ ((1 + US_ZETA) * US_ETAJ) * US_NDJ[0] ^ (1 + US_ZETA) + US_BETA * US_XIJ * (US_PIC[1] / (US_PIC[0] ^ US_CHIJ * US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)))) ^ ((1 + US_ZETA) * US_ETAJ) * US_FJ[1] + + US_GJ[0] = US_NDJ[0] * (1 - US_TAUN[0] - US_TAUWH[0]) * US_LAMBDAJ[0] * US_WJ[0] ^ US_ETAJ + US_BETA * US_XIJ * (US_PIC[1] / (US_PIC[0] ^ US_CHIJ * US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)))) ^ (US_ETAJ - 1) * US_GJ[1] + + US_WJ[0] ^ (1 - US_ETAJ) = (1 - US_XIJ) * US_WJTILDE[0] ^ (1 - US_ETAJ) + US_XIJ * US_WJ[-1] ^ (1 - US_ETAJ) * (US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)) * US_PIC[-1] ^ US_CHIJ / US_PIC[0]) ^ (1 - US_ETAJ) + + US_YS[0] = US_Z[0] * US_KD[0] ^ US_ALPHA * US_ND[0] ^ (1 - US_ALPHA) - US_PSIBAR + + US_RK[0] = US_ALPHA * (US_YS[0] + US_PSIBAR) / US_KD[0] * US_MC[0] + + US_MC[0] = 1 / (US_Z[0] * US_ALPHA ^ US_ALPHA * (1 - US_ALPHA) ^ (1 - US_ALPHA)) * US_RK[0] ^ US_ALPHA * ((1 + US_TAUWF[0]) * US_W[0]) ^ (1 - US_ALPHA) + + US_NDI[0] = US_ND[0] * (1 - US_OMEGA) * (US_WI[0] / US_W[0]) ^ (-US_ETA) + + US_NDJ[0] = US_ND[0] * US_OMEGA * (US_WJ[0] / US_W[0]) ^ (-US_ETA) + + US_ND[0] ^ (1 - 1 / US_ETA) = (1 - US_OMEGA) ^ (1 / US_ETA) * US_NDI[0] ^ (1 - 1 / US_ETA) + US_OMEGA ^ (1 / US_ETA) * US_NDJ[0] ^ (1 - 1 / US_ETA) + + US_D[0] = US_Y[0] * US_PY[0] - US_RK[0] * US_KD[0] - US_ND[0] * (1 + US_TAUWF[0]) * US_W[0] + + US_PHTILDE[0] / US_PH[0] = US_THETA / (US_THETA - 1) * US_FH[0] / US_GH[0] + + US_FH[0] = US_MC[0] * US_H[0] + US_LAMBDAI[1] * US_BETA * US_XIH / US_LAMBDAI[0] * (US_PIH[1] / (US_PIH[0] ^ US_CHIH * US_PI4TARGET ^ (0.25 * (1 - US_CHIH)))) ^ US_THETA * US_FH[1] + + US_GH[0] = US_PH[0] * US_H[0] + US_LAMBDAI[1] * US_BETA * US_XIH / US_LAMBDAI[0] * (US_PIH[1] / (US_PIH[0] ^ US_CHIH * US_PI4TARGET ^ (0.25 * (1 - US_CHIH)))) ^ (US_THETA - 1) * US_GH[1] + + US_PH[0] ^ (1 - US_THETA) = (1 - US_XIH) * US_PHTILDE[0] ^ (1 - US_THETA) + US_XIH * (US_PH[-1] / US_PIC[0]) ^ (1 - US_THETA) * (US_PI4TARGET ^ (0.25 * (1 - US_CHIH)) * US_PIH[-1] ^ US_CHIH) ^ (1 - US_THETA) + + US_PIH[0] = US_PIC[0] * US_PH[0] / US_PH[-1] + + EA_PIMTILDE[0] / EA_PIM[0] = US_THETA / (US_THETA - 1) * US_FX[0] / US_GX[0] + + US_FX[0] = US_MC[0] * EA_IM[0] * EA_SIZE / US_SIZE + US_LAMBDAI[1] * US_BETA * US_XIX / US_LAMBDAI[0] * (EA_PIIM[1] / (EA_PIIM[0] ^ US_CHIX * US_PI4TARGET ^ (0.25 * (1 - US_CHIX)))) ^ US_THETA * US_FX[1] + + US_GX[0] = EA_IM[0] * EA_SIZE * EA_PIM[0] * USEA_RER[0] / US_SIZE + US_LAMBDAI[1] * US_BETA * US_XIX / US_LAMBDAI[0] * (EA_PIIM[1] / (EA_PIIM[0] ^ US_CHIX * US_PI4TARGET ^ (0.25 * (1 - US_CHIX)))) ^ (US_THETA - 1) * US_GX[1] + + EA_PIM[0] ^ (1 - US_THETA) = (1 - US_XIX) * EA_PIMTILDE[0] ^ (1 - US_THETA) + US_XIX * (EA_PIM[-1] / EA_PIC[0]) ^ (1 - US_THETA) * (EA_PIIM[-1] ^ US_CHIX * EA_PI4TARGET ^ (0.25 * (1 - US_CHIH))) ^ (1 - US_THETA) + + EA_PIIM[0] = EA_PIC[0] * EA_PIM[0] / EA_PIM[-1] + + USEA_RER[0] = US_RER / EA_RER[0] + + US_QC[0] ^ ((US_MUC - 1) / US_MUC) = US_NUC ^ (1 / US_MUC) * US_HC[0] ^ (1 - 1 / US_MUC) + (1 - US_NUC) ^ (1 / US_MUC) * ((1 - US_GAMMAIMC[0]) * US_IMC[0]) ^ (1 - 1 / US_MUC) + + 1 = US_NUC * US_PH[0] ^ (1 - US_MUC) + (1 - US_NUC) * (US_PIM[0] / US_GAMMAIMCDAG[0]) ^ (1 - US_MUC) + + US_HC[0] = US_QC[0] * US_NUC * US_PH[0] ^ (-US_MUC) + + US_GAMMAIMC[0] = US_GAMMAIMC1 / 2 * (US_IMC[0] / US_QC[0] / (US_IMC[-1] / US_QC[-1]) - 1) ^ 2 + + US_GAMMAIMCDAG[0] = 1 - US_GAMMAIMC[0] - US_IMC[0] * US_GAMMAIMC1 * (US_IMC[0] / US_QC[0] / (US_IMC[-1] / US_QC[-1]) - 1) / US_QC[0] / (US_IMC[-1] / US_QC[-1]) + + US_QI[0] ^ ((US_MUI - 1) / US_MUI) = US_NUI ^ (1 / US_MUI) * US_HI[0] ^ (1 - 1 / US_MUI) + (1 - US_NUI) ^ (1 / US_MUI) * ((1 - US_GAMMAIMI[0]) * US_IMI[0]) ^ (1 - 1 / US_MUI) + + US_PI[0] ^ (1 - US_MUI) = US_NUI * US_PH[0] ^ (1 - US_MUI) + (1 - US_NUI) * (US_PIM[0] / US_GAMMAIMIDAG[0]) ^ (1 - US_MUI) + + US_HI[0] = US_QI[0] * US_NUI * (US_PH[0] / US_PI[0]) ^ (-US_MUI) + + US_GAMMAIMI[0] = US_GAMMAIMI1 / 2 * (US_IMI[0] / US_QI[0] / (US_IMI[-1] / US_QI[-1]) - 1) ^ 2 + + US_GAMMAIMIDAG[0] = 1 - US_GAMMAIMI[0] - US_IMI[0] * US_GAMMAIMI1 * (US_IMI[0] / US_QI[0] / (US_IMI[-1] / US_QI[0]) - 1) / US_QI[0] / (US_IMI[-1] / US_QI[-1]) + + US_PH[-1] * US_G[-1] + US_TR[-1] + US_B[-1] * US_PIC[-1] ^ (-1) + US_PIC[-1] ^ (-1) * US_M[-2] = US_TAUC[-1] * US_C[-1] + (US_TAUN[-1] + US_TAUWH[-1]) * (US_WI[-1] * US_NDI[-1] + US_WJ[-1] * US_NDJ[-1]) + US_TAUWF[-1] * US_W[-1] * US_ND[-1] + US_TAUK[-1] * (US_RK[-1] * US_U[-1] - (US_DELTA + US_GAMMAU[-1]) * US_PI[-1]) * US_K[-1] + US_TAUD[-1] * US_D[-1] + US_T[-1] + US_R[-1] ^ (-1) * US_B[0] + US_M[-1] + + US_PH[0] * US_G[0] = US_GY[0] * US_PYBAR * US_YBAR + + US_TR[0] = US_YBAR * US_PYBAR * US_TRY[0] + + US_T[0] / (US_PYBAR * US_YBAR) = US_PHITB * (US_B[0] / (US_PYBAR * US_YBAR) - US_BYTARGET) + + US_TI[0] = US_T[0] * US_UPSILONT + + US_TRI[0] = US_TR[0] * US_UPSILONTR + + US_PIC4[0] = US_PIC[0] * US_PIC[-1] * US_PIC[-2] * US_PIC[-3] + + US_RR[0] - 1 = US_R[0] / US_PIC[1] - 1 + + US_C[0] = US_CI[0] * (1 - US_OMEGA) + US_CJ[0] * US_OMEGA + + US_M[0] = US_MI[0] * (1 - US_OMEGA) + US_MJ[0] * US_OMEGA + + US_K[0] = US_KI[0] * (1 - US_OMEGA) + + US_I[0] = US_II[0] * (1 - US_OMEGA) + + US_TRJ[0] = US_TR[0] * 1 / US_OMEGA - US_TRI[0] * (1 - US_OMEGA) / US_OMEGA + + US_TJ[0] = US_T[0] * 1 / US_OMEGA - US_TI[0] * (1 - US_OMEGA) / US_OMEGA + + US_GAMMAV[0] = US_GAMMAVI[0] * US_CI[0] * (1 - US_OMEGA) + US_GAMMAVJ[0] * US_CJ[0] * US_OMEGA + + US_NI[0] = US_NDI[0] * US_SI[0] + + US_SI[0] = (1 - US_XII) * (US_WITILDE[0] / US_WI[0]) ^ (-US_ETAI) + US_XII * (US_WI[-1] / US_WI[0]) ^ (-US_ETAI) * (US_PIC[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHII)) * US_PIC[-1] ^ US_CHII)) ^ US_ETAI * US_SI[-1] + + US_NJ[0] = US_NDJ[0] * US_SJ[0] + + US_SJ[0] = (1 - US_XIJ) * (US_WJTILDE[0] / US_WJ[0]) ^ (-US_ETAJ) + US_XIJ * (US_WJ[-1] / US_WJ[0]) ^ (-US_ETAJ) * (US_PIC[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHIJ)) * US_PIC[-1] ^ US_CHIJ)) ^ US_ETAJ * US_SJ[-1] + + US_U[0] * US_K[0] = US_KD[0] + + US_YS[0] = US_H[0] * US_SH[0] + EA_IM[0] * EA_SIZE * US_SX[0] / US_SIZE + + US_H[0] = US_G[0] + US_HC[0] + US_HI[0] + + US_IM[0] = US_IMC[0] + US_IMI[0] + + US_SH[0] = (1 - US_XIH) * (US_PHTILDE[0] / US_PH[0]) ^ (-US_THETA) + US_XIH * (US_PIH[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHIH)) * US_PIH[-1] ^ US_CHIH)) ^ US_THETA * US_SH[-1] + + US_SX[0] = (1 - US_XIX) * (EA_PIMTILDE[0] / EA_PIM[0]) ^ (-US_THETA) + US_XIX * (EA_PIIM[0] / (US_PI4TARGET ^ (0.25 * (1 - US_CHIH)) * EA_PIIM[-1] ^ US_CHIX)) ^ US_THETA * US_SX[-1] + + US_QC[0] = US_C[0] + US_GAMMAV[0] + + US_QI[0] = US_I[0] + US_GAMMAU[0] * US_K[0] + + US_Y[0] * US_PY[0] = EA_IM[0] * EA_SIZE * EA_PIM[0] * USEA_RER[0] / US_SIZE + US_PH[0] * US_G[0] + US_QC[0] + US_PI[0] * US_QI[0] - US_PIM[0] * ((1 - US_GAMMAIMC[0]) * US_IMC[0] / US_GAMMAIMCDAG[0] + (1 - US_GAMMAIMI[0]) * US_IMI[0] / US_GAMMAIMIDAG[0]) + + US_Y[0] = US_YS[0] + + log(US_Z[0]) = (1 - US_RHOZ) * log(US_ZBAR) + US_RHOZ * log(US_Z[-1]) + σ_US_Z * US_EPSZ[x] + + US_GY[0] = (1 - US_RHOG) * US_GYBAR + US_RHOG * US_GY[-1] + σ_US_G * US_EPSG[x] + + US_TRY[0] = (1 - US_RHOTR) * US_TRYBAR + US_RHOTR * US_TRY[-1] + σ_US_TR * US_EPSTR[x] + + US_TAUC[0] = (1 - US_RHOTAUC) * US_TAUCBAR + US_TAUC[-1] * US_RHOTAUC + σ_US_TAUC * US_EPSTAUC[x] + + US_TAUD[0] = (1 - US_RHOTAUD) * US_TAUDBAR + US_TAUD[-1] * US_RHOTAUD + σ_US_TAUD * US_EPSTAUD[x] + + US_TAUK[0] = US_TAUKBAR * (1 - US_RHOTAUK) + US_TAUK[-1] * US_RHOTAUK + σ_US_TAUK * US_EPSTAUK[x] + + US_TAUN[0] = (1 - US_RHOTAUN) * US_TAUNBAR + US_TAUN[-1] * US_RHOTAUN + σ_US_TAUN * US_EPSTAUN[x] + + US_TAUWH[0] = (1 - US_RHOTAUWH) * US_TAUWHBAR + US_TAUWH[-1] * US_RHOTAUWH + σ_US_TAUWH * US_EPSTAUWH[x] + + US_TAUWF[0] = (1 - US_RHOTAUWF) * US_TAUWFBAR + US_TAUWF[-1] * US_RHOTAUWF + σ_US_TAUWF * US_EPSTAUWF[x] + + US_CY[0] = US_C[0] / (US_Y[0] * US_PY[0]) + + US_IY[0] = US_PI[0] * US_I[0] / (US_Y[0] * US_PY[0]) + + US_IMY[0] = US_PIM[0] * US_IM[0] / (US_Y[0] * US_PY[0]) + + US_IMCY[0] = US_PIM[0] * US_IMC[0] / (US_Y[0] * US_PY[0]) + + US_IMIY[0] = US_PIM[0] * US_IMI[0] / (US_Y[0] * US_PY[0]) + + US_BY[0] = US_B[0] / (US_PYBAR * US_YBAR) + + US_TY[0] = US_T[0] / (US_PYBAR * US_YBAR) + + US_YGAP[0] = US_Y[0] / US_YBAR - 1 + + US_YGROWTH[0] = US_Y[0] / US_Y[-1] + + US_YSHARE[0] = US_Y[0] * US_SIZE * US_PY[0] / US_RER / (EA_Y[0] * EA_PY[0] * EA_SIZE / EA_RER[0] + US_Y[0] * US_SIZE * US_PY[0] / US_RER) + + US_EPSILONM[0] = ( - 0.125) / (US_R[0] * (US_R[0] + US_R[0] * US_GAMMAV2 - 1)) + + 1 = EA_LAMBDAI[1] * EA_BETA * US_R[0] * (1 - EA_GAMMAB[0]) / EA_LAMBDAI[0] * EA_RERDEP[1] / US_PIC[1] + + EA_GAMMAB[0] = EA_GAMMAB1 * (exp(EA_RER[0] * EA_BF[0] / US_PIC[0] / (EA_Y[0] * EA_PY[0]) - EA_BFYTARGET) - 1) - EA_RP[0] + + EA_RP[0] = EA_RHORP * EA_RP[-1] + σ_EA_RP * EA_EPSRP[x] + + EA_RERDEP[0] = EA_RER[0] / EA_RER[-1] + + EA_TOT[0] = EA_PIM[0] / (US_PIM[0] * EA_RER[0]) + + EA_TB[0] = US_IM[0] * US_SIZE * US_PIM[0] * EA_RER[0] / EA_SIZE - EA_PIM[0] * EA_IM[0] + + EA_BF[0] / US_R[-1] = EA_BF[-1] + EA_TB[-1] / EA_RER[-1] + + EA_SIZE * EA_BF[0] + US_SIZE * US_BF[0] = 0 + + end + + + @parameters NAWM_EAUS_2008_incomplete begin + EA_RRSTAR = 1 / EA_BETA + + US_RRSTAR = 1 / US_BETA + end + + # write the parameters from NAWM_EAUS_2008 to a csv file + 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_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) + + @test isapprox(sol1, sol2, rtol = 1e-7) + end + + @testset verbose = true "Code quality (Aqua.jl)" begin + # Aqua.test_all(MacroModelling) + @testset "Compare Project.toml and test/Project.toml" Aqua.test_project_extras(MacroModelling) + @testset "Stale dependencies" Aqua.test_stale_deps(MacroModelling; ignore = [:Showoff]) + @testset "Unbound type parameters" Aqua.test_unbound_args(MacroModelling) + @testset "Undefined exports" Aqua.test_undefined_exports(MacroModelling) + @testset "Piracy" Aqua.test_piracies(MacroModelling) + @testset "Method ambiguity" Aqua.test_ambiguities(MacroModelling, recursive = false) + @testset "Compat" Aqua.test_deps_compat(MacroModelling)#; ignore = [:Aqua, :JET]) + # @testset "Persistent tasks" Aqua.test_persistent_tasks(MacroModelling) + end + GC.gc() + + # test_higher_order = true + @testset verbose = true "Test various models: NSSS and 1st order solution" begin + include("test_models.jl") + end + GC.gc() + # test_higher_order = false + + @testset verbose = true "for and if loops" begin + include("models/Backus_Kehoe_Kydland_1992_for_if_test.jl") + include("models/Backus_Kehoe_Kydland_1992.jl") + + std1 = get_std(Backus_Kehoe_Kydland_1992) + std2 = get_std(Backus_Kehoe_Kydland_1992_test) + + common_keys1 = intersect(std1.keys[1], std2.keys[1]) + common_keys2 = intersect(std1.keys[2], std2.keys[2]) + + @test isapprox(std2(common_keys1, common_keys2), std1(common_keys1, common_keys2), rtol = 1e-10) + end + + @testset verbose = true "Model without shocks" begin + @model m begin + K[0] = (1 - δ) * K[-1] + I[0] + Z[0] = (1 - ρ) * μ + ρ * Z[-1] + I[1] = ((ρ + δ - Z[0])/(1 - δ)) + ((1 + ρ)/(1 - δ)) * I[0] + end + + @parameters m verbose = true begin + ρ = 0.05 + δ = 0.10 + μ = .17 + σ = .2 + end + + m_ss = get_steady_state(m) + @test isapprox(m_ss(:,:Steady_state),[1/7.5,1/.75,.17],rtol = eps(Float32)) + + m_sol = get_solution(m) + @test isapprox(m_sol(:,:K),[1/.75,.9,.04975124378109454],rtol = eps(Float32)) + + init = m_ss(:,:Steady_state) |> collect + init[2] *= 1.5 + get_irf(m, initial_state = init, shocks = :none) + + plots = plot_irf(m, initial_state = init, shocks = :none) + + @test plots[1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} + + plots! = plot_irf!(m, initial_state = init .* 1.5, shocks = :none) + + @test plots![1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} + end + m = nothing + + @testset verbose = true "Borrowing constraint (OBC model with max)" begin + include("models/borrowing_constraint.jl") + + bc_ss = get_steady_state(borrowing_constraint) + @test size(bc_ss, 1) > 0 + + bc_sol = get_solution(borrowing_constraint) + @test size(bc_sol, 1) > 0 + + bc_irf = get_irf(borrowing_constraint) + @test size(bc_irf, 1) > 0 + end + borrowing_constraint = nothing + GC.gc() + + @testset verbose = true "Distribution functions, general and SS" begin + + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + c_normcdf[0]= normcdf(c[0]) + c_normpdf[0]= normpdf(c[0]) + c_norminvcdf[0]= norminvcdf(c[0]-1) + c_norminv[0]= norminv(c[0]-1) + c_qnorm[0]= qnorm(c[0]-1) + c_dnorm[0]= dnorm(c[0]) + c_pnorm[0]= pnorm(c[0]) + c_normlogpdf[0]= normlogpdf(c[0]) + # c_norm[0]= cdf(Normal(),c[0]) + c_inv[0] = erfcinv(c[0]) + # c_binomlogpdf[0]= binomlogpdf(c[0]) + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + get_solution(RBC_CME) + + + @model finacc begin + R[0] * beta = C[1] / C[0] + C[0] = w[0] * L[0] - B[0] + R[-1] * B[-1] + (1-v) * (Rk[-1] * Q[-1] * K[-1] - (R[-1] + mu * G[0] * Rk[-1] * Q[-1] * K[-1] / (Q[-1] * K[-1] - N[-1])) * (Q[-1] * K[-1] - N[-1])) - We + w[0] = C[0] / (1-L[0]) + K[0] = (1-delta) * K[-1] + I[0] + Q[0] = 1 + chi * (I[0] / K[-1] - delta) + Y[0] = A[0] * K[-1]^alpha * L[0]^(1-alpha) + Rk[-1] = (alpha * Y[0] / K[-1] + Q[0] * (1-delta))/Q[-1] + w[0] = (1-alpha) * Y[0] / L[0] + N[0] = v * (Rk[-1] * Q[-1] * K[-1] - (R[-1] + mu * G[0] * Rk[-1] * Q[-1] * K[-1] / (Q[-1] * K[-1] - N[-1])) * (Q[-1] * K[-1] - N[-1])) + We + 0 = (omegabar[0] * (1 - F[0]) + (1 - mu) * G[0]) * Rk[0] / R[0] * Q[0] * K[0] / N[0] - (Q[0] * K[0] / N[0] - 1) + 0 = (1 - (omegabar[0] * (1 - F[0]) + G[0])) * Rk[0] / R[0] + (1 - F[0]) / (1 - F[0] - omegabar[0] * mu * (normpdf((log(omegabar[0]) + sigma^2/2) / sigma)/ omegabar[0] / sigma)) * ((omegabar[0] * (1 - F[0]) + (1 - mu) * G[0]) * Rk[0] / R[0] - 1) + G[0] = normcdf(((log(omegabar[0])+sigma^2/2)/sigma) - sigma) + F[0] = normcdf((log(omegabar[0])+sigma^2/2)/sigma) + EFP[0] = (mu * G[0] * Rk[-1] * Q[-1] * K[-1] / (Q[-1] * K[-1] - N[-1])) + Y[0] + walras[0] = C[0] + I[0] + EFP[0] * (Q[-1] * K[-1] - N[-1]) + B[0] = Q[0] * K[0] - N[0] + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + end + + + @parameters finacc verbose = true begin + beta = 0.99 + delta = 0.02 + We = 1e-12 + alpha = 0.36 + chi = 0 + v = 0.978 + mu = 0.94 + sigma = 0.2449489742783178 + rhoz = .9 + std_eps = .0068 + + .5 > omegabar > .44 + K > 15 + 0 < L < .45 + end + + # solve!(finacc) + @test isapprox(get_steady_state(finacc,derivatives = false)[:,1],[1.0, 7.004987166460695, 1.2762549358842095, 0.0008293608419033882, 0.0009318065746306208, 0.0003952537570055814, 0.30743973601435376, 15.371986800781423, 0.4435430773517457, 8.366999635233856, 1.0000000000593001, 1.0101010101010102, 1.0172249577970442, 1.5895043340984303, 0.4529051354389826, 2.2935377097663356, -1.4597012487627126e-10], rtol = 1e-6) + end + RBC_CME = nothing + finacc = nothing + + @testset verbose = true "Lead and lag > 1" begin + + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + sol = get_solution(RBC_CME) + + # sol1 = get_solution(RBC_CME, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol, sol1, atol = 1e-4) + + # sol2 = get_solution(RBC_CME, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol, sol2, atol = 1e-4) + + sol3 = get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol, sol3, atol = 1e-4) + + + # exo multi lead/lag >> 1 + @model RBC_CME_exo_mult begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-8] + eps_z[x-4] + eps_z[x+4] + eps_z_s[x]) + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_exo_mult verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_exo_mult + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_exo_mult = nothing + + + # endo/exo multi lead/lag >> 1 + @model RBC_CME_all_mult begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-8] + eps_z[x-4] + eps_z[x+4] + eps_z_s[x]) + ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + ZZ_avg_fut[0] = (A[0] + A[1] + A[2] + A[3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_all_mult verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_all_mult + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_all_mult = nothing + + + # exo lead >> 1 + @model RBC_CME_exo_lead1 begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x+8] + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_exo_lead1 verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_exo_lead1 + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_exo_lead1 = nothing + + + + # exo multi lag >> 1 + @model RBC_CME_exo_lag_mult begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * (eps_z[x-8] + eps_z[x-4] + eps_z_s[x]) + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_exo_lag_mult verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_exo_lag_mult + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_exo_lag_mult = nothing + + + # exo lag >> 1 + @model RBC_CME_exo_lag1 begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x-8] + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_exo_lag1 verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_exo_lag1 + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_exo_lag1 = nothing + + + # exo lead > 1 + @model RBC_CME_exo_lead begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x+1] + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_exo_lead verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_exo_lead + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_exo_lead = nothing + + + # exo lag > 1 + @model RBC_CME_exo_lag begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x-1] + # ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_exo_lag verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_exo_lag + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_exo_lag = nothing + + # Lags > 1 + @model RBC_CME_lag begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + ZZ_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + ZZ_dev[0] = log(c[0]/c[ss]) + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_lag verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_lag + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_lag = nothing + + # Leads > 1 + @model RBC_CME_lead begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + ZZ_avg[0] = (A[0] + A[1] + A[2] + A[3]) / 4 + ZZ_dev[0] = log(c[0]/c[ss]) + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + @parameters RBC_CME_lead verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_lead + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_lead = nothing + + + # Leads and lags > 1 + @model RBC_CME_lead_lag begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + ZZ_avg_f[0] = (A[0] + A[1] + A[2] + A[3]) / 4 + ZZ_avg_b[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_lead_lag verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_lead_lag + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_lead_lag = nothing + + + # Leads and lags > 10 + @model RBC_CME_lead_lag10 begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + ZZ_avg_f[0] = (A[0] + A[1] + A[2] + A[3]) / 4 + ZZ_avg_b[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + ZZ_avg_f10[0] = (A[0] + A[10]) / 2 + # ZZ_avg_b10[0] = (A[0] + A[-10]) / 2 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_lead_lag10 verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_lead_lag10 + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_lead_lag10 = nothing + + + # Leads and lags > 10 + @model RBC_CME_lead_lag20 begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + ZZ_avg_f[0] = (A[0] + A[1] + A[2] + A[3]) / 4 + ZZ_avg_b[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + ZZ_avg_ff[0] = (A[0] + A[10]) / 2 + ZZ_avg_bb[0] = (A[0] + A[-10]) / 2 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME_lead_lag20 verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + m = RBC_CME_lead_lag20 + + sol_exo = get_solution(m) + + # sol_exo1 = get_solution(m, algorithm = :linear_time_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + # @test isapprox(sol_exo, sol_exo1, atol = 1e-4) + + + # sol_exo2 = get_solution(m, algorithm = :quadratic_iteration, verbose = true) + + # @test isapprox(sol(axiskeys(sol, 1)[1:end-1],:), sol_exo2(axiskeys(sol, 1)[1:end-1], axiskeys(sol, 2)), atol = 1e-4) + + sol_exo3 = get_solution(m, quadratic_matrix_equation_algorithm = :doubling, verbose = true) + + @test isapprox(sol_exo, sol_exo3, atol = 1e-4) + + m = nothing + RBC_CME_lead_lag20 = nothing + end + + + @testset verbose = true "Steady state RBC CME model" begin + # Basic test + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # A_avg[0] = (A[0] + A[-1] + A[-2] + A[-3]) / 4 + # A_annual[0] = (A[0] + A[-4] + A[-8] + A[-12]) / 4 + # y_avg[0] = log(y[0] / y[-4]) + # y_growth[0] = log(y[1] / y[2]) + # y_growthl[0] = log(y[0] / y[1]) + # y_growthl1[0] = log(y[-1] / y[0]) + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + # get_steady_state(RBC_CME)[1] + # using NLopt + # RBC_CME.SS_optimizer = NLopt.LD_LBFGS + # solve!(RBC_CME) + @test get_steady_state(RBC_CME)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0024019205374952, 1.003405325870413, 1.2092444352939415, 9.467573947982233, 1.42321160651834, 1.0] + # get_moments(RBC_CME)[1] + # irf(RBC_CME) + + RBC_CME = nothing + + + # Symbolic test + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME ss_symbolic_mode = :full verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + alpha = .157 + + # beta | R[ss] = R_ss + # R_ss = 1.0035 + beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss]/y[ss] = I_K_ratio #check why this doesnt solve for y + # I_K_ratio = .15 + delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + + # cap_share > 0 + # R_ss > 0 + # Pi_ss > 0 + # I_K_ratio > 0 + + # 0 < alpha < 1 + 0 < beta < 1 + # 0 < delta < 1 + 0 < Pibar + # 0 <= rhoz < 1 + phi_pi > 0 + + # 0 < A < 1 + # 0 < k < 50 + 0 < Pi + 0 < R + end + # get_steady_state(RBC_CME)[1] + # using NLopt + # RBC_CME.SS_optimizer = NLopt.LD_LBFGS + # solve!(RBC_CME,symbolic_SS = true) + @test get_steady_state(RBC_CME)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0024019205374952, 1.003405325870413, 1.2092444352939415, 9.467573947982233, 1.42321160651834, 1.0] + # get_moments(RBC_CME)[1] + + RBC_CME = nothing + + + + + # Numerical test with calibration targets + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve + R_ss = 1.0035 + # beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = 1.0025 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + + + # cap_share > 0 + # R_ss > 0 + # Pi_ss > 0 + # I_K_ratio > 0 + + # 0 < alpha < 1 + # 0 < beta < 1 + # 0 < delta < 1 + # 0 < Pibar + # 0 <= rhoz < 1 + # phi_pi > 0 + + # 0 < A < 1 + # 0 < k < 50 + # 0 < y < 10 + # 0 < c < 10 + end + # get_steady_state(RBC_CME)[1] + # using NLopt + # RBC_CME.SS_optimizer = NLopt.LD_LBFGS + # solve!(RBC_CME, verbose = true) + # RBC_CME.SS_init_guess[1:7] = [1.0, 1.0025, 1.0035, 1.2081023828249515, 9.437411555244328, 1.4212969209705313, 1.0] + # get_steady_state(RBC_CME) + @test get_steady_state(RBC_CME, verbose = true)(RBC_CME.constants.post_model_macro.var,:Steady_state) ≈ [1.0, 1.0025, 1.0035, 1.2081023824176236, 9.437411552284384, 1.4212969205027686, 1.0] + # get_moments(RBC_CME)[1] + + # RBC_CME.ss_solve_blocks[1]([0.15662344139650963, 1.2081023828249515, 0.02259036144578319, 9.437411555244328, 1.4212969209705313],RBC_CME) + + RBC_CME = nothing + + + # Symbolic test with calibration targets + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME ss_symbolic_mode = :full verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss + R_ss = 1.0035 + # beta = .999 + + delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss] / y[ss] = I_K_ratio # this doesnt solve symbolically + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = 1.0025 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + + # cap_share > 0 + # R_ss > 0 + # Pi_ss > 0 + # I_K_ratio > 0 + + # 0 < alpha < .5 + # 0 < beta < 1 + # 0 < delta < 1 + # 0 < Pibar + # 0 <= rhoz < 1 + # phi_pi > 0 + + # 0 < A < 1 + # 0 < k < 50 + # 0 < y < 10 + # 0 < c < 10 + end + # # get_steady_state(RBC_CME)[1] + # # using NLopt + # # RBC_CME.SS_optimizer = NLopt.LD_LBFGS + # # get_steady_state(RBC_CME) + @test isapprox(get_steady_state(RBC_CME, verbose = true)(RBC_CME.constants.post_model_macro.var,:Steady_state), [1.0, 1.0025, 1.0035, 1.2081023828249515, 9.437411555244328, 1.4212969209705313, 1.0],rtol = eps(Float32)) + # get_moments(RBC_CME)[1] + + RBC_CME = nothing + end + + + + + # using MacroModelling: @model, @parameters, get_steady_state, solve! + + @testset verbose = true "Steady state Smets_Wouters_2003 model" begin + + @model Smets_Wouters_2003 begin + -q[0] + beta * ((1 - tau) * q[1] + epsilon_b[1] * (r_k[1] * z[1] - psi^-1 * r_k[ss] * (-1 + exp(psi * (-1 + z[1])))) * (C[1] - h * C[0])^(-sigma_c)) = 0 + -q_f[0] + beta * ((1 - tau) * q_f[1] + epsilon_b[1] * (r_k_f[1] * z_f[1] - psi^-1 * r_k_f[ss] * (-1 + exp(psi * (-1 + z_f[1])))) * (C_f[1] - h * C_f[0])^(-sigma_c)) = 0 + -r_k[0] + alpha * epsilon_a[0] * mc[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^(-1 + alpha) = 0 + -r_k_f[0] + alpha * epsilon_a[0] * mc_f[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^(-1 + alpha) = 0 + -G[0] + T[0] = 0 + -G[0] + G_bar * epsilon_G[0] = 0 + -G_f[0] + T_f[0] = 0 + -G_f[0] + G_bar * epsilon_G[0] = 0 + -L[0] + nu_w[0]^-1 * L_s[0] = 0 + -L_s_f[0] + L_f[0] * (W_i_f[0] * W_f[0]^-1)^(lambda_w^-1 * (-1 - lambda_w)) = 0 + L_s_f[0] - L_f[0] = 0 + L_s_f[0] + lambda_w^-1 * L_f[0] * W_f[0]^-1 * (-1 - lambda_w) * (-W_disutil_f[0] + W_i_f[0]) * (W_i_f[0] * W_f[0]^-1)^(-1 + lambda_w^-1 * (-1 - lambda_w)) = 0 + Pi_ws_f[0] - L_s_f[0] * (-W_disutil_f[0] + W_i_f[0]) = 0 + Pi_ps_f[0] - Y_f[0] * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 + -Q[0] + epsilon_b[0]^-1 * q[0] * (C[0] - h * C[-1])^(sigma_c) = 0 + -Q_f[0] + epsilon_b[0]^-1 * q_f[0] * (C_f[0] - h * C_f[-1])^(sigma_c) = 0 + -W[0] + epsilon_a[0] * mc[0] * (1 - alpha) * L[0]^(-alpha) * (K[-1] * z[0])^alpha = 0 + -W_f[0] + epsilon_a[0] * mc_f[0] * (1 - alpha) * L_f[0]^(-alpha) * (K_f[-1] * z_f[0])^alpha = 0 + -Y_f[0] + Y_s_f[0] = 0 + Y_s[0] - nu_p[0] * Y[0] = 0 + -Y_s_f[0] + Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 + beta * epsilon_b[1] * (C_f[1] - h * C_f[0])^(-sigma_c) - epsilon_b[0] * R_f[0]^-1 * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 + beta * epsilon_b[1] * pi[1]^-1 * (C[1] - h * C[0])^(-sigma_c) - epsilon_b[0] * R[0]^-1 * (C[0] - h * C[-1])^(-sigma_c) = 0 + Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) - lambda_p^-1 * Y_f[0] * (1 + lambda_p) * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-1 - lambda_p^-1 * (1 + lambda_p)) = 0 + epsilon_b[0] * W_disutil_f[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) - omega * epsilon_b[0] * epsilon_L[0] * L_s_f[0]^sigma_l = 0 + -1 + xi_p * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1) + (1 - xi_p) * pi_star[0]^(-lambda_p^-1) = 0 + -1 + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1) + xi_w * (W[-1] * W[0]^-1)^(-lambda_w^-1) * (pi[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1) = 0 + -Phi - Y_s[0] + epsilon_a[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^alpha = 0 + -Phi - Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_a[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^alpha = 0 + eta_b[exo] - log(epsilon_b[0]) + rho_b * log(epsilon_b[-1]) = 0 + -eta_L[exo] - log(epsilon_L[0]) + rho_L * log(epsilon_L[-1]) = 0 + eta_I[exo] - log(epsilon_I[0]) + rho_I * log(epsilon_I[-1]) = 0 + eta_w[exo] - f_1[0] + f_2[0] = 0 + eta_a[exo] - log(epsilon_a[0]) + rho_a * log(epsilon_a[-1]) = 0 + eta_p[exo] - g_1[0] + g_2[0] * (1 + lambda_p) = 0 + eta_G[exo] - log(epsilon_G[0]) + rho_G * log(epsilon_G[-1]) = 0 + -f_1[0] + beta * xi_w * f_1[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1) + epsilon_b[0] * w_star[0] * L[0] * (1 + lambda_w)^-1 * (C[0] - h * C[-1])^(-sigma_c) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) = 0 + -f_2[0] + beta * xi_w * f_2[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) + omega * epsilon_b[0] * epsilon_L[0] * (L[0] * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)))^(1 + sigma_l) = 0 + -g_1[0] + beta * xi_p * pi_star[0] * g_1[1] * pi_star[1]^-1 * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1) + epsilon_b[0] * pi_star[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 + -g_2[0] + beta * xi_p * g_2[1] * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_b[0] * mc[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 + -nu_w[0] + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) + xi_w * nu_w[-1] * (W[-1] * pi[0]^-1 * W[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w)) = 0 + -nu_p[0] + (1 - xi_p) * pi_star[0]^(-lambda_p^-1 * (1 + lambda_p)) + xi_p * nu_p[-1] * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) = 0 + -K[0] + K[-1] * (1 - tau) + I[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2) = 0 + -K_f[0] + K_f[-1] * (1 - tau) + I_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2) = 0 + U[0] - beta * U[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C[0] - h * C[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s[0]^(1 + sigma_l)) = 0 + U_f[0] - beta * U_f[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C_f[0] - h * C_f[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s_f[0]^(1 + sigma_l)) = 0 + -epsilon_b[0] * (C[0] - h * C[-1])^(-sigma_c) + q[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2 - varphi * I[-1]^-1 * epsilon_I[0] * I[0] * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])) + beta * varphi * I[0]^-2 * epsilon_I[1] * q[1] * I[1]^2 * (-1 + I[0]^-1 * epsilon_I[1] * I[1]) = 0 + -epsilon_b[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) + q_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2 - varphi * I_f[-1]^-1 * epsilon_I[0] * I_f[0] * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])) + beta * varphi * I_f[0]^-2 * epsilon_I[1] * q_f[1] * I_f[1]^2 * (-1 + I_f[0]^-1 * epsilon_I[1] * I_f[1]) = 0 + eta_pi[exo] - log(pi_obj[0]) + rho_pi_bar * log(pi_obj[-1]) + log(calibr_pi_obj) * (1 - rho_pi_bar) = 0 + -C[0] - I[0] - T[0] + Y[0] - psi^-1 * r_k[ss] * K[-1] * (-1 + exp(psi * (-1 + z[0]))) = 0 + -calibr_pi + eta_R[exo] - log(R[ss]^-1 * R[0]) + r_Delta_pi * (-log(pi[ss]^-1 * pi[-1]) + log(pi[ss]^-1 * pi[0])) + r_Delta_y * (-log(Y[ss]^-1 * Y[-1]) + log(Y[ss]^-1 * Y[0]) + log(Y_f[ss]^-1 * Y_f[-1]) - log(Y_f[ss]^-1 * Y_f[0])) + rho * log(R[ss]^-1 * R[-1]) + (1 - rho) * (log(pi_obj[0]) + r_pi * (-log(pi_obj[0]) + log(pi[ss]^-1 * pi[-1])) + r_Y * (log(Y[ss]^-1 * Y[0]) - log(Y_f[ss]^-1 * Y_f[0]))) = 0 + -C_f[0] - I_f[0] + Pi_ws_f[0] - T_f[0] + Y_f[0] + L_s_f[0] * W_disutil_f[0] - L_f[0] * W_f[0] - psi^-1 * r_k_f[ss] * K_f[-1] * (-1 + exp(psi * (-1 + z_f[0]))) = 0 + epsilon_b[0] * (K[-1] * r_k[0] - r_k[ss] * K[-1] * exp(psi * (-1 + z[0]))) * (C[0] - h * C[-1])^(-sigma_c) = 0 + epsilon_b[0] * (K_f[-1] * r_k_f[0] - r_k_f[ss] * K_f[-1] * exp(psi * (-1 + z_f[0]))) * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 + end + + + @parameters Smets_Wouters_2003 verbose = true begin + calibr_pi_obj | 1 = pi_obj[ss] + calibr_pi | pi[ss] = pi_obj[ss] + # Phi | Y_s[ss] * .408 = Phi + # Phi = .408 * Y_j[ss] + # (Y_j[ss] + Phi) / Y_j[ss] = 1.408 -> Phi; | this seems problematic because of the parameter + # lambda_p | .6 = C_f[ss] / Y_f[ss] + # lambda_w | L[ss] = .33 + # G_bar | .18 = G[ss] / Y[ss] + # calibr_pi_obj = 0 + # calibr_pi = 1 + lambda_p = .368 + G_bar = .362 + lambda_w = 0.5 + Phi = .819 + + alpha = 0.3 + beta = 0.99 + gamma_w = 0.763 + gamma_p = 0.469 + h = 0.573 + omega = 1 + psi = 0.169 + r_pi = 1.684 + r_Y = 0.099 + r_Delta_pi = 0.14 + r_Delta_y = 0.159 + rho = 0.961 + rho_b = 0.855 + rho_L = 0.889 + rho_I = 0.927 + rho_a = 0.823 + rho_G = 0.949 + rho_pi_bar = 0.924 + sigma_c = 1.353 + sigma_l = 2.4 + tau = 0.025 + varphi = 6.771 + xi_w = 0.737 + xi_p = 0.908 + + # Putting non-negative constraint on first block is enough + # 0 < K + # 0 < I + # 0 < Y_s + # 0 < q + # 0 < r_k + # 5 < f_1 + # 0 < L + # 0 < W + # 30 < g_1 + # 0 < z + # 0 < mc + # 0 < w_star + # 5 < f_2 + # 0 < Y + # 0 < g_2 + # 0 < C + end + + + # solve!(Smets_Wouters_2003, verbose = true) + + + @test isapprox(get_steady_state(Smets_Wouters_2003, verbose = true)(Smets_Wouters_2003.constants.post_model_macro.var,[:Steady_state]), + [ 1.2043777509278788 + 1.2043777484127967 + 0.362 + 0.362 + 0.44153840098985714 + 0.44153839784516097 + 17.66153603957938 + 17.66153591381742 + 1.2891159430437658 + 1.2891159432893282 + 1.289115942962812 + 1.289115943290125 + 0.9999999999999677 + 0.5401411855429173 + 0.482173806623137 + 0.999999999999352 + 1.0000000000002556 + 1.0101010101010102 + 1.0101010101010102 + 0.362 + 0.362 + -427.9858908413812 + -427.98589116567274 + 1.122103431093411 + 0.7480689524203904 + 1.1221034286309022 + 1.122103428630708 + 2.0079161519182205 + 2.0079161462568305 + 2.0079161519185624 + 2.007916146256947 + 1.0 + 1.0 + 1.0 + 1.0 + 1.0 + 8.770699454739315 + 8.770699454739393 + 48.84717584575426 + 35.70699988724729 + 0.7309941520471651 + 0.730994152046712 + 1.0 + 0.9999999999372022 + 1.0 + 1.0 + 1.0 + 2.4590033452182944 + 2.459003352168301 + 0.035101010082054955 + 0.03510101010809239 + 1.1221034311168996 + 0.9999999999999698 + 1.0000000000000089], + rtol = eps(Float32) + ) + + Smets_Wouters_2003 = nothing + # # x = 1 + + + # # Smets_Wouters_2003 = nothing + + # # @model Smets_Wouters_2003 begin + # # -q[0] + beta * ((1 - tau) * q[1] + epsilon_b[1] * (r_k[1] * z[1] - psi^-1 * r_k[ss] * (-1 + exp(psi * (-1 + z[1])))) * (C[1] - h * C[0])^(-sigma_c)) = 0 + # # -q_f[0] + beta * ((1 - tau) * q_f[1] + epsilon_b[1] * (r_k_f[1] * z_f[1] - psi^-1 * r_k_f[ss] * (-1 + exp(psi * (-1 + z_f[1])))) * (C_f[1] - h * C_f[0])^(-sigma_c)) = 0 + # # -r_k[0] + alpha * epsilon_a[0] * mc[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^(-1 + alpha) = 0 + # # -r_k_f[0] + alpha * epsilon_a[0] * mc_f[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^(-1 + alpha) = 0 + # # -G[0] + T[0] = 0 + # # -G[0] + G_bar * epsilon_G[0] = 0 + # # -G_f[0] + T_f[0] = 0 + # # -G_f[0] + G_bar * epsilon_G[0] = 0 + # # -L[0] + nu_w[0]^-1 * L_s[0] = 0 + # # -L_s_f[0] + L_f[0] * (W_i_f[0] * W_f[0]^-1)^(lambda_w^-1 * (-1 - lambda_w)) = 0 + # # L_s_f[0] - L_f[0] = 0 + # # L_s_f[0] + lambda_w^-1 * L_f[0] * W_f[0]^-1 * (-1 - lambda_w) * (-W_disutil_f[0] + W_i_f[0]) * (W_i_f[0] * W_f[0]^-1)^(-1 + lambda_w^-1 * (-1 - lambda_w)) = 0 + # # Pi_ws_f[0] - L_s_f[0] * (-W_disutil_f[0] + W_i_f[0]) = 0 + # # Pi_ps_f[0] - Y_f[0] * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 + # # -Q[0] + epsilon_b[0]^-1 * q[0] * (C[0] - h * C[-1])^(sigma_c) = 0 + # # -Q_f[0] + epsilon_b[0]^-1 * q_f[0] * (C_f[0] - h * C_f[-1])^(sigma_c) = 0 + # # -W[0] + epsilon_a[0] * mc[0] * (1 - alpha) * L[0]^(-alpha) * (K[-1] * z[0])^alpha = 0 + # # -W_f[0] + epsilon_a[0] * mc_f[0] * (1 - alpha) * L_f[0]^(-alpha) * (K_f[-1] * z_f[0])^alpha = 0 + # # -Y_f[0] + Y_s_f[0] = 0 + # # Y_s[0] - nu_p[0] * Y[0] = 0 + # # -Y_s_f[0] + Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) = 0 + # # beta * epsilon_b[1] * (C_f[1] - h * C_f[0])^(-sigma_c) - epsilon_b[0] * R_f[0]^-1 * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 + # # beta * epsilon_b[1] * pi[1]^-1 * (C[1] - h * C[0])^(-sigma_c) - epsilon_b[0] * R[0]^-1 * (C[0] - h * C[-1])^(-sigma_c) = 0 + # # Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) - lambda_p^-1 * Y_f[0] * (1 + lambda_p) * (-mc_f[0] + P_j_f[0]) * P_j_f[0]^(-1 - lambda_p^-1 * (1 + lambda_p)) = 0 + # # epsilon_b[0] * W_disutil_f[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) - omega * epsilon_b[0] * epsilon_L[0] * L_s_f[0]^sigma_l = 0 + # # -1 + xi_p * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1) + (1 - xi_p) * pi_star[0]^(-lambda_p^-1) = 0 + # # -1 + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1) + xi_w * (W[-1] * W[0]^-1)^(-lambda_w^-1) * (pi[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1) = 0 + # # -Phi - Y_s[0] + epsilon_a[0] * L[0]^(1 - alpha) * (K[-1] * z[0])^alpha = 0 + # # -Phi - Y_f[0] * P_j_f[0]^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_a[0] * L_f[0]^(1 - alpha) * (K_f[-1] * z_f[0])^alpha = 0 + # # eta_b[exo] - log(epsilon_b[0]) + rho_b * log(epsilon_b[-1]) = 0 + # # -eta_L[exo] - log(epsilon_L[0]) + rho_L * log(epsilon_L[-1]) = 0 + # # eta_I[exo] - log(epsilon_I[0]) + rho_I * log(epsilon_I[-1]) = 0 + # # eta_w[exo] - f_1[0] + f_2[0] = 0 + # # eta_a[exo] - log(epsilon_a[0]) + rho_a * log(epsilon_a[-1]) = 0 + # # eta_p[exo] - g_1[0] + g_2[0] * (1 + lambda_p) = 0 + # # eta_G[exo] - log(epsilon_G[0]) + rho_G * log(epsilon_G[-1]) = 0 + # # -f_1[0] + beta * xi_w * f_1[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1) + epsilon_b[0] * w_star[0] * L[0] * (1 + lambda_w)^-1 * (C[0] - h * C[-1])^(-sigma_c) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) = 0 + # # -f_2[0] + beta * xi_w * f_2[1] * (w_star[0]^-1 * w_star[1])^(lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) * (pi[1]^-1 * pi[0]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w) * (1 + sigma_l)) + omega * epsilon_b[0] * epsilon_L[0] * (L[0] * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)))^(1 + sigma_l) = 0 + # # -g_1[0] + beta * xi_p * pi_star[0] * g_1[1] * pi_star[1]^-1 * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1) + epsilon_b[0] * pi_star[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 + # # -g_2[0] + beta * xi_p * g_2[1] * (pi[1]^-1 * pi[0]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) + epsilon_b[0] * mc[0] * Y[0] * (C[0] - h * C[-1])^(-sigma_c) = 0 + # # -nu_w[0] + (1 - xi_w) * (w_star[0] * W[0]^-1)^(-lambda_w^-1 * (1 + lambda_w)) + xi_w * nu_w[-1] * (W[-1] * pi[0]^-1 * W[0]^-1 * pi[-1]^gamma_w)^(-lambda_w^-1 * (1 + lambda_w)) = 0 + # # -nu_p[0] + (1 - xi_p) * pi_star[0]^(-lambda_p^-1 * (1 + lambda_p)) + xi_p * nu_p[-1] * (pi[0]^-1 * pi[-1]^gamma_p)^(-lambda_p^-1 * (1 + lambda_p)) = 0 + # # -K[0] + K[-1] * (1 - tau) + I[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2) = 0 + # # -K_f[0] + K_f[-1] * (1 - tau) + I_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2) = 0 + # # U[0] - beta * U[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C[0] - h * C[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s[0]^(1 + sigma_l)) = 0 + # # U_f[0] - beta * U_f[1] - epsilon_b[0] * ((1 - sigma_c)^-1 * (C_f[0] - h * C_f[-1])^(1 - sigma_c) - omega * epsilon_L[0] * (1 + sigma_l)^-1 * L_s_f[0]^(1 + sigma_l)) = 0 + # # -epsilon_b[0] * (C[0] - h * C[-1])^(-sigma_c) + q[0] * (1 - 0.5 * varphi * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])^2 - varphi * I[-1]^-1 * epsilon_I[0] * I[0] * (-1 + I[-1]^-1 * epsilon_I[0] * I[0])) + beta * varphi * I[0]^-2 * epsilon_I[1] * q[1] * I[1]^2 * (-1 + I[0]^-1 * epsilon_I[1] * I[1]) = 0 + # # -epsilon_b[0] * (C_f[0] - h * C_f[-1])^(-sigma_c) + q_f[0] * (1 - 0.5 * varphi * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])^2 - varphi * I_f[-1]^-1 * epsilon_I[0] * I_f[0] * (-1 + I_f[-1]^-1 * epsilon_I[0] * I_f[0])) + beta * varphi * I_f[0]^-2 * epsilon_I[1] * q_f[1] * I_f[1]^2 * (-1 + I_f[0]^-1 * epsilon_I[1] * I_f[1]) = 0 + # # eta_pi[exo] - log(pi_obj[0]) + rho_pi_bar * log(pi_obj[-1]) + log(calibr_pi_obj) * (1 - rho_pi_bar) = 0 + # # -C[0] - I[0] - T[0] + Y[0] - psi^-1 * r_k[ss] * K[-1] * (-1 + exp(psi * (-1 + z[0]))) = 0 + # # -calibr_pi + eta_R[exo] - log(R[ss]^-1 * R[0]) + r_Delta_pi * (-log(pi[ss]^-1 * pi[-1]) + log(pi[ss]^-1 * pi[0])) + r_Delta_y * (-log(Y[ss]^-1 * Y[-1]) + log(Y[ss]^-1 * Y[0]) + log(Y_f[ss]^-1 * Y_f[-1]) - log(Y_f[ss]^-1 * Y_f[0])) + rho * log(R[ss]^-1 * R[-1]) + (1 - rho) * (log(pi_obj[0]) + r_pi * (-log(pi_obj[0]) + log(pi[ss]^-1 * pi[-1])) + r_Y * (log(Y[ss]^-1 * Y[0]) - log(Y_f[ss]^-1 * Y_f[0]))) = 0 + # # -C_f[0] - I_f[0] + Pi_ws_f[0] - T_f[0] + Y_f[0] + L_s_f[0] * W_disutil_f[0] - L_f[0] * W_f[0] - psi^-1 * r_k_f[ss] * K_f[-1] * (-1 + exp(psi * (-1 + z_f[0]))) = 0 + # # epsilon_b[0] * (K[-1] * r_k[0] - r_k[ss] * K[-1] * exp(psi * (-1 + z[0]))) * (C[0] - h * C[-1])^(-sigma_c) = 0 + # # epsilon_b[0] * (K_f[-1] * r_k_f[0] - r_k_f[ss] * K_f[-1] * exp(psi * (-1 + z_f[0]))) * (C_f[0] - h * C_f[-1])^(-sigma_c) = 0 + # # end + + + # # @parameters Smets_Wouters_2003 begin + # # calibr_pi_obj | 1 = pi_obj[ss] + # # calibr_pi | pi[ss] = pi_obj[ss] + # # Phi | (Y_s[ss] + Phi) / Y_s[ss] = 1.408 + # # # lambda_p | .6 = C_f[ss] / Y_f[ss] + # # # lambda_w | L[ss] = .33 + # # G_bar | .18 = G[ss] / Y[ss] + + # # lambda_p = .368 + # # # G_bar = .362 + # # lambda_w = 0.5 + # # # Phi = .819 + + # # alpha = 0.3 + # # beta = 0.99 + # # gamma_w = 0.763 + # # gamma_p = 0.469 + # # h = 0.573 + # # omega = 1 + # # psi = 0.169 + # # r_pi = 1.684 + # # r_Y = 0.099 + # # r_Delta_pi = 0.14 + # # r_Delta_y = 0.159 + # # rho = 0.961 + # # rho_b = 0.855 + # # rho_L = 0.889 + # # rho_I = 0.927 + # # rho_a = 0.823 + # # rho_G = 0.949 + # # rho_pi_bar = 0.924 + # # sigma_c = 1.353 + # # sigma_l = 2.4 + # # tau = 0.025 + # # varphi = 6.771 + # # xi_w = 0.737 + # # xi_p = 0.908 + + # # end + + + # # solve!(Smets_Wouters_2003, symbolic_SS = false) + + # # # get_steady_state(Smets_Wouters_2003) + + # # @test get_steady_state(Smets_Wouters_2003)[1] ≈ [ 1.20465991441435 + # # 1.204659917151701 + # # 0.3613478048030788 + # # 0.3613478048030788 + # # 0.4414800855444218 + # # 0.4414800896382151 + # # 17.659203422264238 + # # 17.65920357698873 + # # 1.2889457095271066 + # # 1.2889457096070582 + # # 1.2889457095307755 + # # 1.2889457098239414 + # # 1.0000000000366498 + # # 0.5400259611608715 + # # 0.48211013259048446 + # # 1.00000000000172 + # # 1.000000000127065 + # # 1.0101010101010102 + # # 1.0101010101010102 + # # 0.3613478047907606 + # # 0.3613478048030788 + # # -427.92495898028676 + # # -427.9249587468684 + # # 1.1221034247496608 + # # 0.7480689524616317 + # # 1.122103428477167 + # # 1.1221034282377538 + # # 2.0074878047372287 + # # 2.00748781245403 + # # 2.007487804732286 + # # 2.0074878121606647 + # # 1.0 + # # 1.0 + # # 1.0 + # # 1.0 + # # 1.0 + # # 8.766762166589194 + # # 8.766762166588967 + # # 48.8212791635492 + # # 35.68806956399776 + # # 0.730994152045567 + # # 0.7309941520886629 + # # 1.0 + # # 1.0000000000028464 + # # 1.0 + # # 1.0 + # # 1.0 + # # 2.4582240979093846 + # # 2.4582240906598867 + # # 0.03510101014899653 + # # 0.035101010136073356 + # # 1.1221034247485961 + # # 1.0000000000000178 + # # 0.9999999999583465] + + + end + + + @testset verbose = true "First order perturbation" begin + # Numerical test with calibration targets + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve + R_ss = 1.0035 + # beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = 1.0025 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + get_solution(RBC_CME) + @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]], [ 0.0 0.0068 + 6.73489e-6 0.000168887 + 1.01124e-5 0.000253583 + -0.000365783 0.00217203 + -0.00070019 0.00749279 + 0.0 0.00966482 + 0.005 0.0], atol = 1e-6) + + get_solution(RBC_CME, parameters = :I_K_ratio => .1) + @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 + 3.42408e-6 0.000111417 + 5.14124e-6 0.000167292 + -0.000196196 0.00190741 + -0.000430554 0.0066164 + 0.0 0.00852381 + 0.005 0.0], atol = 1e-6) + + get_solution(RBC_CME, parameters = :cap_share => 1.5) + @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 + 4.00629e-6 0.000118171 + 6.01543e-6 0.000177434 + -0.000207089 0.00201698 + -0.00041124 0.00639229 + 0.0 0.00840927 + 0.005 0.0], atol = 1e-6) + + RBC_CME = nothing + end + + + + + + # @testset verbose = true "First order: linear time iteration" begin + # # Numerical test with calibration targets + # @model RBC_CME begin + # y[0]=A[0]*k[-1]^alpha + # 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + # 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + # R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + # z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # end + + + # @parameters RBC_CME verbose = true begin + # alpha | k[ss] / (4 * y[ss]) = cap_share + # cap_share = 1.66 + # # alpha = .157 + + # beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve + # R_ss = 1.0035 + # # beta = .999 + + # # delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. + # I_K_ratio = .15 + # # delta = .0226 + + # Pibar | Pi[ss] = Pi_ss + # Pi_ss = 1.0025 + # # Pibar = 1.0008 + + # phi_pi = 1.5 + # rhoz = .9 + # std_eps = .0068 + # rho_z_delta = .9 + # std_z_delta = .005 + # end + + # get_solution(RBC_CME, algorithm = :linear_time_iteration) + + # @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]], [ 0.0 0.0068 + # 6.73489e-6 0.000168887 + # 1.01124e-5 0.000253583 + # -0.000365783 0.00217203 + # -0.00070019 0.00749279 + # 0.0 0.00966482 + # 0.005 0.0], atol = 1e-6) + + + # get_solution(RBC_CME, algorithm = :linear_time_iteration, parameters = :I_K_ratio => .1) + + # @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 + # 3.42408e-6 0.000111417 + # 5.14124e-6 0.000167292 + # -0.000196196 0.00190741 + # -0.000430554 0.0066164 + # 0.0 0.00852381 + # 0.005 0.0], atol = 1e-6) + + + # get_solution(RBC_CME, algorithm = :linear_time_iteration, parameters = :cap_share => 1.5) + + # @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 + # 4.00629e-6 0.000118171 + # 6.01543e-6 0.000177434 + # -0.000207089 0.00201698 + # -0.00041124 0.00639229 + # 0.0 0.00840927 + # 0.005 0.0], atol = 1e-6) + + # RBC_CME = nothing + # end + + + + + + @testset verbose = true "First order: doubling" begin + # Numerical test with calibration targets + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss # beta needs to enter into function: block in order to solve + R_ss = 1.0035 + # beta = .999 + + # delta | c[ss]/y[ss] = 1 - I_K_ratio + delta | delta * k[ss] / y[ss] = I_K_ratio #check why this doesnt solve for y; because delta is not recognised as a free parameter here. + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = 1.0025 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + end + + get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling) + + @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]], [ 0.0 0.0068 + 6.73489e-6 0.000168887 + 1.01124e-5 0.000253583 + -0.000365783 0.00217203 + -0.00070019 0.00749279 + 0.0 0.00966482 + 0.005 0.0], atol = 1e-6) + + + get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling, parameters = :I_K_ratio => .1) + + @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 + 3.42408e-6 0.000111417 + 5.14124e-6 0.000167292 + -0.000196196 0.00190741 + -0.000430554 0.0066164 + 0.0 0.00852381 + 0.005 0.0], atol = 1e-6) + + + get_solution(RBC_CME, quadratic_matrix_equation_algorithm = :doubling, parameters = :cap_share => 1.5) + + @test isapprox(RBC_CME.caches.first_order_solution_matrix[:,[(end-RBC_CME.constants.post_model_macro.nExo+1):end...]],[ 0.0 0.0068 + 4.00629e-6 0.000118171 + 6.01543e-6 0.000177434 + -0.000207089 0.00201698 + -0.00041124 0.00639229 + 0.0 0.00840927 + 0.005 0.0], atol = 1e-6) + + RBC_CME = nothing + end + + + @testset verbose = true "Plotting" begin + # Symbolic test with calibration targets + @model RBC_CME begin + y[0]=A[0]*k[-1]^alpha + 1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta)) + 1/c[0]=beta*1/c[1]*(R[0]/Pi[+1]) + R[0] * beta =(Pi[0]/Pibar)^phi_pi + # A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta)*k[-1] + A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1] + z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x] + # z[0]=rhoz*z[-1]+std_eps*eps_z[x] + # A[0]=exp(z[0]) + A[0] = 1 - rhoz + rhoz * A[-1] + std_eps * eps_z[x] + # log(A[0]) = rhoz * log(A[-1]) + std_eps * eps_z[x] + end + + + @parameters RBC_CME verbose = true begin + alpha | k[ss] / (4 * y[ss]) = cap_share + cap_share = 1.66 + # alpha = .157 + + beta | R[ss] = R_ss + R_ss = 1.0035 + # beta = .999 + + delta | c[ss]/y[ss] = 1 - I_K_ratio + # delta | delta * k[ss] / y[ss] = I_K_ratio # this doesnt solve symbolically + I_K_ratio = .15 + # delta = .0226 + + Pibar | Pi[ss] = Pi_ss + Pi_ss = 1.0025 + # Pibar = 1.0008 + + phi_pi = 1.5 + rhoz = .9 + std_eps = .0068 + rho_z_delta = .9 + std_z_delta = .005 + + # cap_share > 0 + # R_ss > 0 + # Pi_ss > 0 + # I_K_ratio > 0 + + # 0 < alpha < 1 + # 0 < beta < 1 + # 0 < delta < 1 + # 0 < Pibar + # 0 <= rhoz < 1 + # phi_pi > 0 + + # 0 < A < 1 + # 0 < k < 50 + # 0 < y < 10 + # 0 < c < 10 + end + plots = plot_irf(RBC_CME) + + @test plots[1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} + + plots! = plot_irf!(RBC_CME, parameters = :rhoz => .8) + + @test plots![1] isa StatsPlots.Plots.Plot{StatsPlots.Plots.GRBackend} + + RBC_CME = nothing + end + GC.gc() + + @testset verbose = true "update_equations! / add_equation! / remove_equation!" begin + include("test_update_equations.jl") + end + GC.gc() + diff --git a/test/test_estimation.jl b/test/test_estimation.jl index 42e14144d..bcbe0f51f 100644 --- a/test/test_estimation.jl +++ b/test/test_estimation.jl @@ -1,20 +1,31 @@ +using Test using MacroModelling import Turing -import ADTypes: AutoZygote -import Turing: NUTS, sample, logpdf +import Mooncake +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences +import Turing: NUTS, sample import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys -import Zygote +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys +import StatsPlots + +using FlexiChains +include("test_helpers.jl") 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,:) @@ -34,15 +45,13 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) llh = get_loglikelihood(m, data, all_params, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh # with Turing >= 0.40 this becomes: Turing.@addlogprob! (; loglikelihood = llh) @@ -53,43 +62,65 @@ FS2000_loglikelihood = FS2000_loglikelihood_function(data, FS2000, -Inf) n_samples = 1000 -samps = @time sample(FS2000_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) -# with Turing >= 0.41 this: initial_params = FS2000.parameter_values becomes: initial_params = InitFromParams(all_params = FS2000.parameter_values,)); # need to import InitFromParams -println("Mean variable values (Zygote): $(mean(samps).nt.mean)") +samps = @time sample(FS2000_loglikelihood, NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") get_steady_state(FS2000, steady_state_function = FS2000_custom_steady_state_function!) -samps = @time sample(FS2000_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) -# with Turing >= 0.41 this: initial_params = FS2000.parameter_values becomes: initial_params = InitFromParams(all_params = FS2000.parameter_values,)); # need to import InitFromParams -println("Mean variable values (Zygote + custom steady state): $(mean(samps).nt.mean)") +samps = @time sample(FS2000_loglikelihood, NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake + custom steady state): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") get_steady_state(FS2000, steady_state_function = nothing) -samps = @time sample(FS2000_loglikelihood, NUTS(), n_samples, progress = true, initial_params = FS2000.parameter_values) +samps = @time sample(FS2000_loglikelihood, NUTS(), n_samples, progress = true, initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) -println("Mean variable values (ForwardDiff): $(mean(samps).nt.mean)") +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (ForwardDiff): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") -sample_nuts = mean(samps).nt.mean +sample_nuts = collect(values(FlexiChains.mean(samps); parameters_only = true)) modeFS2000 = Turing.maximum_a_posteriori(FS2000_loglikelihood, # Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 2)), Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), # Optim.NelderMead(), - adtype = AutoZygote(), + adtype = AutoMooncake(; config=nothing), # maxiters = 100, # lb = [0,0,-10,-10,0,0,0,0,0], # ub = [1,1,10,10,1,1,1,100,100], - initial_params = FS2000.parameter_values) + initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) -println("Mode variable values: $(modeFS2000.values); Mode loglikelihood: $(modeFS2000.lp)") +println("Mode variable values: $(modeFS2000.params); Mode loglikelihood: $(modeFS2000.lp)") @testset "Estimation results" begin # @test isapprox(modeFS2000.lp, 1281.669108730447, rtol = eps(Float32)) @test isapprox(sample_nuts, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) end +@testset "Mooncake vs FiniteDifferences gradient (1st order Kalman)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(FS2000, data, x), ADTypes.AutoMooncake(config = nothing), FS2000.parameter_values) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + 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, 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 9cb7a4fff..bd11ed62a 100644 --- a/test/test_estimation_pigeons.jl +++ b/test/test_estimation_pigeons.jl @@ -1,20 +1,23 @@ using MacroModelling 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 +import StatsPlots + +include("test_helpers.jl") 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,:) @@ -34,16 +37,14 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() llh = get_loglikelihood(m, data, all_params, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end diff --git a/test/test_helpers.jl b/test/test_helpers.jl new file mode 100644 index 000000000..b1bc8a268 --- /dev/null +++ b/test/test_helpers.jl @@ -0,0 +1,19 @@ +using Dates + + +function maybe_print_loglikelihood(verbose::Bool, llh, dists, all_params) + verbose || return nothing + prior_llh = Turing.logpdf(Turing.product_distribution(dists), all_params) + println("Loglikelihood: $(llh) and prior llh: $(prior_llh) with params $(all_params)") + return nothing +end + +function quarterly_dates(start_date::Date, len::Int) + dates = Vector{Date}(undef, len) + current_date = start_date + for i in 1:len + dates[i] = current_date + current_date = current_date + Dates.Month(3) + end + return dates +end \ No newline at end of file diff --git a/test/test_higher_order_1.jl b/test/test_higher_order_1.jl new file mode 100644 index 000000000..87364cdbc --- /dev/null +++ b/test/test_higher_order_1.jl @@ -0,0 +1,40 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +# test_higher_order = true + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "FS2000 third order" begin + include("../models/FS2000.jl") + functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) +end +FS2000 = nothing +GC.gc() + +@testset verbose = true "FS2000 pruned third order" begin + include("../models/FS2000.jl") + functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :pruned_third_order, plots = plots) +end +FS2000 = nothing +GC.gc() + +@testset verbose = true "FS2000 second order" begin + include("../models/FS2000.jl") + functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) +end +FS2000 = nothing +GC.gc() + +@testset verbose = true "FS2000 pruned second order" begin + include("../models/FS2000.jl") + functionality_test(FS2000, Caldara_et_al_2012_estim, algorithm = :pruned_second_order, plots = plots) +end +FS2000 = nothing +GC.gc() diff --git a/test/test_higher_order_2.jl b/test/test_higher_order_2.jl new file mode 100644 index 000000000..a28d1da0b --- /dev/null +++ b/test/test_higher_order_2.jl @@ -0,0 +1,26 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +# test_higher_order = true + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables pruned second order" begin + include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :pruned_second_order, plots = plots) +end +# m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables pruned third order" begin + # include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :pruned_third_order, plots = plots) +end +m = nothing +GC.gc() diff --git a/test/test_higher_order_3.jl b/test/test_higher_order_3.jl new file mode 100644 index 000000000..0d1176285 --- /dev/null +++ b/test/test_higher_order_3.jl @@ -0,0 +1,54 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +# test_higher_order = true + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "RBC_CME with calibration equations second order" begin + include("models/RBC_CME_calibration_equations.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) +end +# m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations third order" begin + # include("models/RBC_CME_calibration_equations.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME second order" begin + include("models/RBC_CME.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) +end +# m = nothing +GC.gc() + +@testset verbose = true "RBC_CME third order" begin + # include("models/RBC_CME.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations and parameter definitions second order" begin + include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :second_order, plots = plots) +end +# m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations and parameter definitions third order" begin + # include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") + functionality_test(m, Caldara_et_al_2012_estim, algorithm = :third_order, plots = plots) +end +m = nothing +GC.gc() diff --git a/test/test_inversion_filter_gradients.jl b/test/test_inversion_filter_gradients.jl new file mode 100644 index 000000000..06365e9ae --- /dev/null +++ b/test/test_inversion_filter_gradients.jl @@ -0,0 +1,255 @@ +using Test +using MacroModelling +using Random +using AxisKeys +import LinearAlgebra as ℒ +import ForwardDiff +import Zygote +import FiniteDifferences + +# ----------------------------------------------------------------------------- +# Verifies that ∂loglik/∂parameters for the inversion filter agrees between +# ForwardDiff, Zygote, and FiniteDifferences across all five perturbation-order +# dispatches (first_order, pruned_second_order, second_order, pruned_third_order, +# third_order). +# +# Models: +# * Gali_2015_chapter_3_nonlinear — all 5 algorithms (under-identified obs; +# first_order also tested with square obs and FULL parameter vector) +# * Smets_Wouters_2007 — first_order (full param vector) + pruned_second_order +# (subset). +# +# Edge cases (where supported by the dispatch): +# * n_obs < n_shocks (under-identified — LagrangeNewton path) +# * n_obs == n_shocks (square — first_order only; higher orders cannot +# invert square systems on this model with LagrangeNewton) +# * warmup_iterations > 0 (first_order only — codebase warns it's first- +# order-only and ignores it otherwise) +# * presample_periods > 0 +# +# Note: MacroModelling enforces n_obs ≤ n_shocks at the API level, so the +# over-identified case is intentionally not exercised. +# ----------------------------------------------------------------------------- + +const RTOL = 1e-5 +const FDM = FiniteDifferences.central_fdm(5, 1) + +# Build a Zygote/ForwardDiff-friendly closure that varies a subset of params. +# The full parameter vector is built via a comprehension (no in-place writes) +# so that reverse-mode AD has no mutation to worry about. +function make_llh_closure(model, data, base_params, idx, algorithm; kwargs...) + n = length(base_params) + pos = zeros(Int, n) # 0 ⇒ use base_params[j], else θ_subset[pos[j]] + @inbounds for (k, j) in enumerate(idx) + pos[j] = k + end + return function(θ_subset) + # Use `map` (Zygote-friendly, no in-place setindex! tracing) to splice + # θ_subset into the base-parameter vector. An explicit element-type + # conversion makes the result eltype-stable so ForwardDiff sees a + # `Vector{Dual}` rather than `Vector{Real}`. + T = eltype(θ_subset) + full = map(j -> pos[j] == 0 ? T(base_params[j]) : θ_subset[pos[j]], 1:n) + return get_loglikelihood(model, data, full; + filter = :inversion, + algorithm = algorithm, + on_failure_loglikelihood = -Inf, + kwargs...) + end +end + +function compare_gradients(label, model, data, base_params, idx, algorithm; + rtol = RTOL, kwargs...) + @testset "$label" begin + f = make_llh_closure(model, data, base_params, idx, algorithm; kwargs...) + θ = base_params[idx] + llh_val = f(θ) + @test isfinite(llh_val) + if !isfinite(llh_val) + @info "Skipping $label: forward loglik not finite at base params" + return + end + + # FiniteDifferences reference (slow, but accurate). + fd_grad = first(FiniteDifferences.grad(FDM, f, θ)) + if !all(isfinite, fd_grad) + @info "Skipping $label: FD reference contains NaN/Inf — model fails under perturbation" + return + end + + # ForwardDiff (forward-mode AD). Wrapped so that a ForwardDiff failure + # does not prevent the Zygote check from running — they exercise + # different code paths (the generic primal vs. the rrules). + @testset "ForwardDiff vs FiniteDifferences" begin + fdiff_grad = nothing + fdiff_err = nothing + try + fdiff_grad = ForwardDiff.gradient(f, θ) + catch err + fdiff_err = err + end + if fdiff_err !== nothing + @error "ForwardDiff threw" exception = (fdiff_err, catch_backtrace()) + @test false + else + @test all(isfinite, fdiff_grad) + @test isapprox(fdiff_grad, fd_grad; rtol = rtol) + end + end + + # Zygote (reverse-mode AD; exercises the rrules in src/rrules.jl). + @testset "Zygote vs FiniteDifferences" begin + zg_grad = nothing + zg_err = nothing + try + zg_grad, = Zygote.gradient(f, θ) + catch err + zg_err = err + end + if zg_err !== nothing + @error "Zygote threw" exception = (zg_err, catch_backtrace()) + @test false + else + @test zg_grad !== nothing + @test all(isfinite, zg_grad) + @test isapprox(zg_grad, fd_grad; rtol = rtol) + end + end + end +end + +# Build data as steady-state level + small Gaussian perturbations. This keeps +# the inversion filter inside its convergence basin at every perturbation +# order, so any test failure reflects an AD problem rather than a numerical +# breakdown of the filter itself. +function ss_perturbed_data(model, observables; periods = 8, σ = 1e-4, seed = 42) + SS = get_steady_state(model) + ss_obs = collect(SS(observables, :Steady_state)) + Random.seed!(seed) + dat = repeat(ss_obs, 1, periods) .+ σ .* randn(length(observables), periods) + return KeyedArray(dat; Variables = observables, Time = 1:periods) +end + +# ============================================================================= +# Gali (2015) Chapter 3 nonlinear NK — all 5 algorithms +# ============================================================================= +include("../models/Gali_2015_chapter_3_nonlinear.jl") +const GALI = Gali_2015_chapter_3_nonlinear + +# 3 shocks total → under-identified = 1 or 2 obs, square = 3 obs. +# Square obs must be linearly independent in the first-order solution. +# log_W_real = σ·log_y + φ·log_N (static identity from W_real = Y^σ·N^φ), +# so [:log_y, :log_W_real, :log_N] is rank-deficient; use i_ann instead. +const GALI_OBS_UNDER = [:log_y, :log_W_real] +const GALI_OBS_SQUARE = [:log_y, :log_N, :i_ann] + +const GALI_PARAM_SUBSET_NAMES = [:σ, :φ, :ϕᵖⁱ, :ρ_a, :ρ_z, :std_a, :std_z] + +function gali_subset_indices() + pnames = GALI.constants.post_complete_parameters.parameters + return [findfirst(==(p), pnames) for p in GALI_PARAM_SUBSET_NAMES] +end + +@testset "inversion filter gradient cross-checks (Gali + SW07)" begin + +@testset "Gali_2015 nonlinear inversion filter — gradient cross-checks" begin + base_params = copy(GALI.parameter_values) + p_subset = gali_subset_indices() + + # --- (a) per-algorithm baseline: subset of params, under-identified obs -- + algorithms = [:first_order, :pruned_second_order, :second_order, + :pruned_third_order, :third_order] + + for algo in algorithms + data = ss_perturbed_data(GALI, GALI_OBS_UNDER; periods = 8, σ = 1e-4, seed = 11) + compare_gradients("Gali :$algo (under-identified, $(length(p_subset)) params)", + GALI, data, base_params, p_subset, algo) + end + + # --- (b) FULL parameter vector — first_order only, under-identified ------ + let algo = :first_order + data = ss_perturbed_data(GALI, GALI_OBS_UNDER; periods = 8, σ = 1e-4, seed = 12) + compare_gradients("Gali :$algo (under-identified, FULL param vector)", + GALI, data, base_params, collect(eachindex(base_params)), algo) + end + + # --- (c) square observables (n_obs == n_shocks): first_order only ------- + # (higher-order inversion does not converge for square systems on this + # model — exercising it would test the failure path, not the gradient.) + let algo = :first_order + # Use a tighter σ so that finite-difference perturbations keep the + # square-system inversion inside its convergence basin. + data = ss_perturbed_data(GALI, GALI_OBS_SQUARE; periods = 6, σ = 1e-6, seed = 13) + compare_gradients("Gali :$algo (square obs)", + GALI, data, base_params, p_subset, algo) + end + + # --- (d) warmup_iterations > 0 (first_order only, per implementation) --- + @testset "Gali :first_order (warmup_iterations=2)" begin + let algo = :first_order + data = ss_perturbed_data(GALI, GALI_OBS_UNDER; periods = 8, σ = 1e-4, seed = 14) + compare_gradients("Gali :$algo (warmup_iterations=2)", + GALI, data, base_params, p_subset, algo; + warmup_iterations = 2) + end + end + + # --- (e) presample_periods > 0 — exercise across all 5 algorithms ------- + for algo in algorithms + data = ss_perturbed_data(GALI, GALI_OBS_UNDER; periods = 10, σ = 1e-4, seed = 15) + compare_gradients("Gali :$algo (presample_periods=3)", + GALI, data, base_params, p_subset, algo; + presample_periods = 3) + end +end # Gali testset + + +# ============================================================================= +# Smets-Wouters 2007 — first_order (full) + pruned_second_order (subset) +# ============================================================================= +include("../models/Smets_Wouters_2007.jl") +SW07 = Smets_Wouters_2007 + +# 7 shocks → under-identified by using 3 obs. +SW07_OBS = [:dy, :dc, :dinve] + +# Representative parameter subset for the higher-order pass. +SW07_SUBSET_PREF = [:crhoa, :crhob, :crhog, :csadjcost, :chabb, + :csigma, :cprobw] + +function sw07_subset_indices() + pnames = SW07.constants.post_complete_parameters.parameters + idx = Int[] + for p in SW07_SUBSET_PREF + j = findfirst(==(p), pnames) + if j !== nothing + push!(idx, j) + end + end + if length(idx) < 5 + idx = collect(1:min(7, length(pnames))) + end + return idx +end + +@testset "Smets-Wouters 2007 inversion filter — gradient cross-checks" begin + base_params = copy(SW07.parameter_values) + + # First-order: FULL parameter vector + let algo = :first_order + data = ss_perturbed_data(SW07, SW07_OBS; periods = 12, σ = 1e-4, seed = 21) + compare_gradients("SW07 :$algo (under-identified, FULL param vector, $(length(base_params)) params)", + SW07, data, base_params, + collect(eachindex(base_params)), algo) + end + + # Pruned-2nd: subset only + let algo = :pruned_second_order + data = ss_perturbed_data(SW07, SW07_OBS; periods = 12, σ = 1e-4, seed = 22) + p_subset = sw07_subset_indices() + compare_gradients("SW07 :$algo (under-identified, $(length(p_subset))-param subset)", + SW07, data, base_params, p_subset, algo) + end +end # SW07 testset + +end # outer wrapping testset diff --git a/test/test_jet.jl b/test/test_jet.jl new file mode 100644 index 000000000..b09ce3d29 --- /dev/null +++ b/test/test_jet.jl @@ -0,0 +1,14 @@ +using Test +using MacroModelling + +if VERSION < v"1.13" + using JET +end + +@testset verbose = true "Static checking (JET.jl)" begin + if VERSION < v"1.11" + JET.test_package(MacroModelling; target_defined_modules = true, toplevel_logger = nothing) + elseif VERSION < v"1.13" + JET.test_package(MacroModelling; target_modules = (MacroModelling,), toplevel_logger = nothing) + end +end diff --git a/test/test_jet_hot_paths.jl b/test/test_jet_hot_paths.jl new file mode 100644 index 000000000..cee21c81e --- /dev/null +++ b/test/test_jet_hot_paths.jl @@ -0,0 +1,827 @@ +using Test +using MacroModelling + +if VERSION < v"1.13" + using JET +end + +import MacroModelling: get_NSSS_and_parameters, calculate_jacobian, calculate_hessian, + calculate_third_order_derivatives, calculate_first_order_solution, + calculate_second_order_solution, calculate_third_order_solution, + calculate_covariance, calculate_mean, + calculate_second_order_moments, calculate_third_order_moments, + calculate_second_order_moments_with_covariance, + calculate_third_order_moments_with_autocorrelation, + get_relevant_steady_state_and_state_update, irf_initial_state, + calculate_stochastic_steady_state, + solve_lyapunov_equation, solve_sylvester_equation, solve_quadratic_matrix_equation, + filter_and_smooth, + prepare_stochastic_steady_state_base_terms, solve_stochastic_steady_state_newton, + merge_calculation_options, initialise_constants!, CalculationOptions, Tolerances, + lyapunov_workspace, sylvester_workspace, ensure_lyapunov_workspace! + +using AxisKeys +using SparseArrays +import LinearAlgebra as ℒ + +# --------------------------------------------------------------------------- +# Set up a small RBC model so we have concrete arguments for @report_call. +# This avoids the OOM problem of JET.test_package on the whole package. +# --------------------------------------------------------------------------- + +@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 + +# Solve at third order to force compilation of all derivative/solution functions +get_solution(RBC, algorithm = :third_order, silent = true) + +params = copy(RBC.parameter_values) +opts = merge_calculation_options() + +# Populate NSSS so downstream calls have valid state +SS_and_pars, _ = get_NSSS_and_parameters(RBC, params, opts = opts) + +# Jacobian +∇₁ = calculate_jacobian(params, SS_and_pars, RBC.caches, RBC.functions.jacobian, RBC.workspaces) + +# First-order solution +constants_obj = initialise_constants!(RBC) +𝐒₁, qme_sol, _ = calculate_first_order_solution(∇₁, constants_obj, RBC.workspaces, RBC.caches; + opts = opts, initial_guess = RBC.caches.qme_solution, parameter_values = params) + +# Hessian (available after second+ order solve) +∇₂ = calculate_hessian(params, SS_and_pars, RBC.caches, RBC.functions.hessian, RBC.workspaces) + +# Second-order solution +𝐒₂, _ = calculate_second_order_solution(∇₁, ∇₂, 𝐒₁, RBC.constants, RBC.workspaces, RBC.caches; + opts = opts, parameter_values = params) + +# Third-order derivatives (available after third order solve) +∇₃ = calculate_third_order_derivatives(params, SS_and_pars, RBC.caches, + RBC.functions.third_order_derivatives, RBC.workspaces) + +# Sparse second/third-order solutions for stochastic steady-state rrules +𝐒₂_sp = sparse(𝐒₂) +𝐒₃_raw = RBC.caches.third_order_solution +𝐒₃_sp = 𝐒₃_raw isa SparseMatrixCSC ? 𝐒₃_raw : sparse(𝐒₃_raw) +x_sss = zeros(RBC.constants.post_model_macro.nVars) + +# --------------------------------------------------------------------------- +# JET analysis targets +# --------------------------------------------------------------------------- + +# Collect ignored modules: filter out packages whose internals we do not own +const JET_TARGET_MODULES = (MacroModelling,) + +# Helper: run JET.@report_call with target_modules filtering. +# On Julia < 1.11 the kwarg was `target_defined_modules`. +function jet_test_call(@nospecialize(f), @nospecialize(argtypes); + broken::Bool = false) + if VERSION >= v"1.13" + @test_skip "JET not supported on Julia ≥ 1.13 yet" + return + end + result = if VERSION < v"1.11" + JET.report_call(f, argtypes; + target_defined_modules = true, + toplevel_logger = nothing) + else + JET.report_call(f, argtypes; + target_modules = JET_TARGET_MODULES, + toplevel_logger = nothing) + end + reports = JET.get_reports(result) + if broken + @test_broken isempty(reports) + else + @test isempty(reports) + end + if !isempty(reports) && !broken + @warn "JET reports for $(f)" reports + end +end + +# --------------------------------------------------------------------------- +# Helper: run the standard first-order JET hot-path battery for any model. +# Expects the model to be solved at least at first order before calling. +# --------------------------------------------------------------------------- +function run_first_order_jet_battery(model; model_name::String) + params = copy(model.parameter_values) + opts = merge_calculation_options() + + SS_and_pars, _ = get_NSSS_and_parameters(model, params, opts = opts) + ∇₁ = calculate_jacobian(params, SS_and_pars, model.caches, model.functions.jacobian, model.workspaces) + constants_obj = initialise_constants!(model) + 𝐒₁, _, _ = calculate_first_order_solution(∇₁, constants_obj, model.workspaces, model.caches; + opts = opts, initial_guess = model.caches.qme_solution, parameter_values = params) + + @testset "get_NSSS_and_parameters" begin + jet_test_call(get_NSSS_and_parameters, + Tuple{typeof(model), Vector{Float64}}) + end + + @testset "calculate_jacobian" begin + jet_test_call(calculate_jacobian, + Tuple{typeof(params), typeof(SS_and_pars), + typeof(model.caches), typeof(model.functions.jacobian), + typeof(model.workspaces)}) + end + + @testset "calculate_first_order_solution" begin + jet_test_call(calculate_first_order_solution, + Tuple{typeof(∇₁), typeof(constants_obj), + typeof(model.workspaces), typeof(model.caches)}) + end + + @testset "calculate_covariance" begin + jet_test_call(calculate_covariance, + Tuple{typeof(params), typeof(model)}) + end + + @testset "calculate_mean" begin + jet_test_call(calculate_mean, + Tuple{typeof(params), typeof(model)}) + end + + @testset "calculate_second_order_moments" begin + jet_test_call(calculate_second_order_moments, + Tuple{typeof(params), typeof(model)}) + end + + @testset "get_relevant_steady_state_and_state_update (first_order)" begin + jet_test_call(get_relevant_steady_state_and_state_update, + Tuple{Val{:first_order}, typeof(params), typeof(model)}) + end + + @testset "get_solution (parameters, first_order)" begin + jet_test_call(MacroModelling.get_solution, + Tuple{typeof(model), typeof(params)}) + end + + @testset "get_irf (parameters)" begin + jet_test_call(MacroModelling.get_irf, + Tuple{typeof(model), typeof(params)}) + end + + @testset "get_loglikelihood" begin + data = KeyedArray(randn(1, 40); + Variables = [model.constants.post_model_macro.var[1]], Periods = 1:40) + jet_test_call(MacroModelling.get_loglikelihood, + Tuple{typeof(model), typeof(data), typeof(params)}) + end + + @testset "get_statistics" begin + jet_test_call(MacroModelling.get_statistics, + Tuple{typeof(model), typeof(params)}) + end + + @testset "get_conditional_variance_decomposition" begin + jet_test_call(MacroModelling.get_conditional_variance_decomposition, + Tuple{typeof(model)}) + end + + @testset "filter_and_smooth" begin + obs_syms = [model.constants.post_model_macro.var[1]] + data_fs = randn(1, 40) + jet_test_call(filter_and_smooth, + Tuple{typeof(model), typeof(data_fs), typeof(obs_syms)}) + end + + @testset "solve_lyapunov_equation" begin + lyap_ws = ensure_lyapunov_workspace!(model.workspaces, model.constants.post_model_macro.nVars, :first_order) + n = model.constants.post_model_macro.nVars + A_test = randn(n, n) * 0.5 + C_test = let X = randn(n, n); X * X'; end + jet_test_call(solve_lyapunov_equation, + Tuple{typeof(A_test), typeof(C_test), typeof(lyap_ws)}) + end + + @testset "solve_quadratic_matrix_equation" begin + constants_qme = initialise_constants!(model) + n = model.constants.post_model_macro.nVars - model.constants.post_model_macro.nPresent_only + A_qme = randn(n, n) + B_qme = randn(n, n) + C_qme = randn(n, n) + jet_test_call(solve_quadratic_matrix_equation, + Tuple{typeof(A_qme), typeof(B_qme), typeof(C_qme), + typeof(model.constants), typeof(model.workspaces), typeof(model.caches)}) + end + + @testset "calculate_loglikelihood (Kalman)" begin + calculate_loglikelihood_fn = MacroModelling.calculate_loglikelihood + obs_idx = [1] + data_dev = randn(1, 40) + state_vec = [zeros(model.constants.post_model_macro.nVars)] + jet_test_call(calculate_loglikelihood_fn, + Tuple{Val{:kalman}, Val{:first_order}, typeof(obs_idx), + Matrix{Float64}, typeof(data_dev), + typeof(model.constants), typeof(state_vec), typeof(model.workspaces)}) + end + + @testset "calculate_loglikelihood (Inversion, first_order)" begin + calculate_loglikelihood_fn = MacroModelling.calculate_loglikelihood + obs_idx = [1] + data_dev = randn(1, 40) + state_vec = [zeros(model.constants.post_model_macro.nVars)] + jet_test_call(calculate_loglikelihood_fn, + Tuple{Val{:inversion}, Val{:first_order}, typeof(obs_idx), + Matrix{Float64}, typeof(data_dev), + typeof(model.constants), typeof(state_vec), typeof(model.workspaces)}) + end + + # rrules + @testset "rrule: calculate_jacobian" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_jacobian), typeof(params), typeof(SS_and_pars), + typeof(model.caches), typeof(model.functions.jacobian), + typeof(model.workspaces)}) + end + + @testset "rrule: get_NSSS_and_parameters" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_NSSS_and_parameters), typeof(model), typeof(params)}) + end + + @testset "rrule: calculate_first_order_solution" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_first_order_solution), typeof(∇₁), + typeof(constants_obj), typeof(model.workspaces), typeof(model.caches)}) + end + + @testset "rrule: calculate_covariance" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_covariance), typeof(params), typeof(model)}) + end + + @testset "rrule: calculate_mean" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_mean), typeof(params), typeof(model)}) + end + + @testset "rrule: calculate_second_order_moments" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_second_order_moments), typeof(params), typeof(model)}) + end + + @testset "rrule: get_relevant_steady_state_and_state_update (first_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_relevant_steady_state_and_state_update), + Val{:first_order}, typeof(params), typeof(model)}) + end + + @testset "rrule: get_irf" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_irf), typeof(model), typeof(params)}) + end + + @testset "rrule: get_loglikelihood" begin + data_rl = KeyedArray(randn(1, 40); + Variables = [model.constants.post_model_macro.var[1]], Periods = 1:40) + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_loglikelihood), typeof(model), typeof(data_rl), typeof(params)}) + end + + @testset "rrule: get_solution" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_solution), typeof(model), typeof(params)}) + end + + @testset "rrule: get_statistics" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_statistics), typeof(model), typeof(params)}) + end + + @testset "rrule: solve_lyapunov_equation" begin + lyap_ws_r = ensure_lyapunov_workspace!(model.workspaces, model.constants.post_model_macro.nVars, :first_order) + n_r = model.constants.post_model_macro.nVars + A_lyap_r = randn(n_r, n_r) * 0.5 + C_lyap_r = let X = randn(n_r, n_r); X * X'; end + jet_test_call(MacroModelling.rrule, + Tuple{typeof(solve_lyapunov_equation), + typeof(A_lyap_r), typeof(C_lyap_r), typeof(lyap_ws_r)}) + end +end + +@testset verbose = true "JET hot-path analysis" begin + + # ------------------------------------------------------------------ + @testset "get_NSSS_and_parameters" begin + jet_test_call(get_NSSS_and_parameters, + Tuple{typeof(RBC), Vector{Float64}}) + end + + # ------------------------------------------------------------------ + @testset "calculate_jacobian" begin + jet_test_call(calculate_jacobian, + Tuple{typeof(params), typeof(SS_and_pars), + typeof(RBC.caches), typeof(RBC.functions.jacobian), + typeof(RBC.workspaces)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_first_order_solution" begin + jet_test_call(calculate_first_order_solution, + Tuple{typeof(∇₁), typeof(constants_obj), + typeof(RBC.workspaces), typeof(RBC.caches)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_hessian" begin + jet_test_call(calculate_hessian, + Tuple{typeof(params), typeof(SS_and_pars), + typeof(RBC.caches), typeof(RBC.functions.hessian), + typeof(RBC.workspaces)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_second_order_solution" begin + jet_test_call(calculate_second_order_solution, + Tuple{typeof(∇₁), typeof(∇₂), typeof(𝐒₁), + typeof(RBC.constants), typeof(RBC.workspaces), + typeof(RBC.caches)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_third_order_derivatives" begin + jet_test_call(calculate_third_order_derivatives, + Tuple{typeof(params), typeof(SS_and_pars), + typeof(RBC.caches), typeof(RBC.functions.third_order_derivatives), + typeof(RBC.workspaces)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_covariance" begin + jet_test_call(calculate_covariance, + Tuple{typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_mean" begin + jet_test_call(calculate_mean, + Tuple{typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "calculate_second_order_moments" begin + jet_test_call(calculate_second_order_moments, + Tuple{typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "get_relevant_steady_state_and_state_update (first_order)" begin + jet_test_call(get_relevant_steady_state_and_state_update, + Tuple{Val{:first_order}, typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "get_relevant_steady_state_and_state_update (second_order)" begin + jet_test_call(get_relevant_steady_state_and_state_update, + Tuple{Val{:second_order}, typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "get_relevant_steady_state_and_state_update (pruned_second_order)" begin + jet_test_call(get_relevant_steady_state_and_state_update, + Tuple{Val{:pruned_second_order}, typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "get_relevant_steady_state_and_state_update (third_order)" begin + jet_test_call(get_relevant_steady_state_and_state_update, + Tuple{Val{:third_order}, typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + @testset "get_relevant_steady_state_and_state_update (pruned_third_order)" begin + jet_test_call(get_relevant_steady_state_and_state_update, + Tuple{Val{:pruned_third_order}, typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # High-level user-facing functions (estimation-style calls with parameters vector) + # ------------------------------------------------------------------ + @testset "get_solution (parameters, first_order)" begin + jet_test_call(MacroModelling.get_solution, + Tuple{typeof(RBC), typeof(params)}) + end + + @testset "get_irf (parameters)" begin + jet_test_call(MacroModelling.get_irf, + Tuple{typeof(RBC), typeof(params)}) + end + + @testset "get_loglikelihood" begin + # Construct minimal fake data with the right shape + data = KeyedArray(randn(1, 40); Variables = [RBC.constants.post_model_macro.var[1]], Periods = 1:40) + jet_test_call(MacroModelling.get_loglikelihood, + Tuple{typeof(RBC), typeof(data), typeof(params)}) + end + + # ------------------------------------------------------------------ + # Third-order solution + # ------------------------------------------------------------------ + @testset "calculate_third_order_solution" begin + jet_test_call(calculate_third_order_solution, + Tuple{typeof(∇₁), typeof(∇₂), typeof(∇₃), + typeof(𝐒₁), typeof(𝐒₂), + typeof(RBC.constants), typeof(RBC.workspaces), + typeof(RBC.caches)}) + end + + # ------------------------------------------------------------------ + # Second-order moments with covariance + # ------------------------------------------------------------------ + @testset "calculate_second_order_moments_with_covariance" begin + jet_test_call(calculate_second_order_moments_with_covariance, + Tuple{typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # Third-order moments + # ------------------------------------------------------------------ + @testset "calculate_third_order_moments" begin + jet_test_call(calculate_third_order_moments, + Tuple{typeof(params), Symbol, typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # Third-order moments with autocorrelation + # ------------------------------------------------------------------ + @testset "calculate_third_order_moments_with_autocorrelation" begin + jet_test_call(calculate_third_order_moments_with_autocorrelation, + Tuple{typeof(params), Symbol, typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # Stochastic steady state (4 algorithm variants) + # ------------------------------------------------------------------ + @testset "calculate_stochastic_steady_state (second_order)" begin + jet_test_call(calculate_stochastic_steady_state, + Tuple{Val{:second_order}, typeof(params), typeof(RBC)}) + end + + @testset "calculate_stochastic_steady_state (pruned_second_order)" begin + jet_test_call(calculate_stochastic_steady_state, + Tuple{Val{:pruned_second_order}, typeof(params), typeof(RBC)}) + end + + @testset "calculate_stochastic_steady_state (third_order)" begin + jet_test_call(calculate_stochastic_steady_state, + Tuple{Val{:third_order}, typeof(params), typeof(RBC)}) + end + + @testset "calculate_stochastic_steady_state (pruned_third_order)" begin + jet_test_call(calculate_stochastic_steady_state, + Tuple{Val{:pruned_third_order}, typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # Lyapunov equation solver (top-level dispatcher, dense doubling path) + # ------------------------------------------------------------------ + @testset "solve_lyapunov_equation (dense, doubling)" begin + lyap_ws = ensure_lyapunov_workspace!(RBC.workspaces, RBC.constants.post_model_macro.nVars, :first_order) + n = RBC.constants.post_model_macro.nVars + A_test = randn(n, n) * 0.5 + C_test = let X = randn(n, n); X * X'; end + jet_test_call(solve_lyapunov_equation, + Tuple{typeof(A_test), typeof(C_test), typeof(lyap_ws)}) + end + + # ------------------------------------------------------------------ + # Sylvester equation solver (top-level dispatcher) + # ------------------------------------------------------------------ + @testset "solve_sylvester_equation" begin + sylv_ws = RBC.workspaces.sylvester_1st_order + n = 3 + A_sylv = randn(n, n) * 0.5 + B_sylv = randn(n, n) * 0.5 + C_sylv = randn(n, n) + jet_test_call(solve_sylvester_equation, + Tuple{typeof(A_sylv), typeof(B_sylv), typeof(C_sylv), typeof(sylv_ws)}) + end + + # ------------------------------------------------------------------ + # Quadratic matrix equation solver + # ------------------------------------------------------------------ + @testset "solve_quadratic_matrix_equation" begin + constants_qme = initialise_constants!(RBC) + n = RBC.constants.post_model_macro.nVars - RBC.constants.post_model_macro.nPresent_only + A_qme = randn(n, n) + B_qme = randn(n, n) + C_qme = randn(n, n) + jet_test_call(solve_quadratic_matrix_equation, + Tuple{typeof(A_qme), typeof(B_qme), typeof(C_qme), + typeof(RBC.constants), typeof(RBC.workspaces), typeof(RBC.caches)}) + end + + # ------------------------------------------------------------------ + # Kalman filter loglikelihood + # ------------------------------------------------------------------ + @testset "calculate_loglikelihood (Kalman)" begin + calculate_loglikelihood_fn = MacroModelling.calculate_loglikelihood + obs_idx = [1] + data_dev = randn(1, 40) + state_vec = [zeros(RBC.constants.post_model_macro.nVars)] + jet_test_call(calculate_loglikelihood_fn, + Tuple{Val{:kalman}, Val{:first_order}, typeof(obs_idx), + Matrix{Float64}, typeof(data_dev), + typeof(RBC.constants), typeof(state_vec), typeof(RBC.workspaces)}) + end + + # ------------------------------------------------------------------ + # Inversion filter loglikelihood (first order) + # ------------------------------------------------------------------ + @testset "calculate_loglikelihood (Inversion, first_order)" begin + calculate_loglikelihood_fn = MacroModelling.calculate_loglikelihood + obs_idx = [1] + data_dev = randn(1, 40) + state_vec = [zeros(RBC.constants.post_model_macro.nVars)] + jet_test_call(calculate_loglikelihood_fn, + Tuple{Val{:inversion}, Val{:first_order}, typeof(obs_idx), + Matrix{Float64}, typeof(data_dev), + typeof(RBC.constants), typeof(state_vec), typeof(RBC.workspaces)}) + end + + # ------------------------------------------------------------------ + # Filter and smooth (Durbin-Koopman) + # ------------------------------------------------------------------ + @testset "filter_and_smooth" begin + obs_syms = [RBC.constants.post_model_macro.var[1]] + data_fs = randn(1, 40) + jet_test_call(filter_and_smooth, + Tuple{typeof(RBC), typeof(data_fs), typeof(obs_syms)}) + end + + # ------------------------------------------------------------------ + # Conditional variance decomposition (user-facing) + # ------------------------------------------------------------------ + @testset "get_conditional_variance_decomposition" begin + jet_test_call(MacroModelling.get_conditional_variance_decomposition, + Tuple{typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # Shock decomposition (user-facing) + # ------------------------------------------------------------------ + @testset "get_shock_decomposition" begin + data_sd = KeyedArray(randn(1, 40); Variables = [RBC.constants.post_model_macro.var[1]], Periods = 1:40) + jet_test_call(MacroModelling.get_shock_decomposition, + Tuple{typeof(RBC), typeof(data_sd)}) + end + + # ------------------------------------------------------------------ + # Conditional forecast (user-facing) + # ------------------------------------------------------------------ + @testset "get_conditional_forecast" begin + cond_mat = Matrix{Union{Nothing,Float64}}(nothing, RBC.constants.post_model_macro.nVars, 5) + cond_mat[1, 1] = 0.01 + jet_test_call(MacroModelling.get_conditional_forecast, + Tuple{typeof(RBC), typeof(cond_mat)}) + end + + # ------------------------------------------------------------------ + # get_solution with higher-order algorithms + # ------------------------------------------------------------------ + @testset "get_solution (parameters, second_order)" begin + jet_test_call(MacroModelling.get_solution, + Tuple{typeof(RBC), typeof(params)}) + end + + # ------------------------------------------------------------------ + # get_statistics (user-facing, used in estimation) + # ------------------------------------------------------------------ + @testset "get_statistics" begin + jet_test_call(MacroModelling.get_statistics, + Tuple{typeof(RBC), typeof(params)}) + end + + # ------------------------------------------------------------------ + # Stochastic steady-state base terms preparation + # ------------------------------------------------------------------ + @testset "prepare_stochastic_steady_state_base_terms" begin + jet_test_call(prepare_stochastic_steady_state_base_terms, + Tuple{typeof(params), typeof(RBC)}) + end + + # ------------------------------------------------------------------ + # ChainRules rrules for key functions (AD hot paths) + # ------------------------------------------------------------------ + @testset "rrule: calculate_jacobian" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_jacobian), typeof(params), typeof(SS_and_pars), + typeof(RBC.caches), typeof(RBC.functions.jacobian), + typeof(RBC.workspaces)}) + end + + @testset "rrule: calculate_hessian" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_hessian), typeof(params), typeof(SS_and_pars), + typeof(RBC.caches), typeof(RBC.functions.hessian), + typeof(RBC.workspaces)}) + end + + @testset "rrule: calculate_third_order_derivatives" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_third_order_derivatives), typeof(params), typeof(SS_and_pars), + typeof(RBC.caches), typeof(RBC.functions.third_order_derivatives), + typeof(RBC.workspaces)}) + end + + @testset "rrule: get_NSSS_and_parameters" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_NSSS_and_parameters), typeof(RBC), typeof(params)}) + end + + @testset "rrule: calculate_first_order_solution" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_first_order_solution), typeof(∇₁), + typeof(constants_obj), typeof(RBC.workspaces), typeof(RBC.caches)}) + end + + @testset "rrule: calculate_second_order_solution" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_second_order_solution), typeof(∇₁), typeof(∇₂), typeof(𝐒₁), + typeof(RBC.constants), typeof(RBC.workspaces), typeof(RBC.caches)}) + end + + @testset "rrule: calculate_covariance" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_covariance), typeof(params), typeof(RBC)}) + end + + @testset "rrule: calculate_mean" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_mean), typeof(params), typeof(RBC)}) + end + + @testset "rrule: calculate_second_order_moments" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_second_order_moments), typeof(params), typeof(RBC)}) + end + + @testset "rrule: get_relevant_steady_state_and_state_update (first_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_relevant_steady_state_and_state_update), + Val{:first_order}, typeof(params), typeof(RBC)}) + end + + @testset "rrule: get_irf" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_irf), typeof(RBC), typeof(params)}) + end + + @testset "rrule: get_loglikelihood" begin + data_rl = KeyedArray(randn(1, 40); Variables = [RBC.constants.post_model_macro.var[1]], Periods = 1:40) + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_loglikelihood), typeof(RBC), typeof(data_rl), typeof(params)}) + end + + # ------------------------------------------------------------------ + # Additional rrules: higher-order solutions, moments, solvers + # ------------------------------------------------------------------ + @testset "rrule: calculate_third_order_solution" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_third_order_solution), typeof(∇₁), typeof(∇₂), typeof(∇₃), + typeof(𝐒₁), typeof(𝐒₂), + typeof(RBC.constants), typeof(RBC.workspaces), typeof(RBC.caches)}) + end + + @testset "rrule: calculate_second_order_moments_with_covariance" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_second_order_moments_with_covariance), + typeof(params), typeof(RBC)}) + end + + @testset "rrule: calculate_third_order_moments" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_third_order_moments), + typeof(params), Symbol, typeof(RBC)}) + end + + @testset "rrule: calculate_third_order_moments_with_autocorrelation" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(calculate_third_order_moments_with_autocorrelation), + typeof(params), Symbol, typeof(RBC)}) + end + + @testset "rrule: prepare_stochastic_steady_state_base_terms" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(prepare_stochastic_steady_state_base_terms), + typeof(params), typeof(RBC)}) + end + + @testset "rrule: solve_stochastic_steady_state_newton (second_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(solve_stochastic_steady_state_newton), + Val{:second_order}, typeof(𝐒₁), typeof(𝐒₂_sp), + typeof(x_sss), typeof(RBC)}) + end + + @testset "rrule: solve_stochastic_steady_state_newton (third_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(solve_stochastic_steady_state_newton), + Val{:third_order}, typeof(𝐒₁), typeof(𝐒₂_sp), typeof(𝐒₃_sp), + typeof(x_sss), typeof(RBC)}) + end + + @testset "rrule: solve_sylvester_equation" begin + sylv_ws_r = RBC.workspaces.sylvester_1st_order + n_r = 3 + A_sylv_r = randn(n_r, n_r) * 0.5 + B_sylv_r = randn(n_r, n_r) * 0.5 + C_sylv_r = randn(n_r, n_r) + jet_test_call(MacroModelling.rrule, + Tuple{typeof(solve_sylvester_equation), + typeof(A_sylv_r), typeof(B_sylv_r), typeof(C_sylv_r), typeof(sylv_ws_r)}) + end + + @testset "rrule: solve_lyapunov_equation" begin + lyap_ws_r = ensure_lyapunov_workspace!(RBC.workspaces, RBC.constants.post_model_macro.nVars, :first_order) + n_r = RBC.constants.post_model_macro.nVars + A_lyap_r = randn(n_r, n_r) * 0.5 + C_lyap_r = let X = randn(n_r, n_r); X * X'; end + jet_test_call(MacroModelling.rrule, + Tuple{typeof(solve_lyapunov_equation), + typeof(A_lyap_r), typeof(C_lyap_r), typeof(lyap_ws_r)}) + end + + @testset "rrule: get_relevant_steady_state_and_state_update (second_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_relevant_steady_state_and_state_update), + Val{:second_order}, typeof(params), typeof(RBC)}) + end + + @testset "rrule: get_relevant_steady_state_and_state_update (pruned_second_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_relevant_steady_state_and_state_update), + Val{:pruned_second_order}, typeof(params), typeof(RBC)}) + end + + @testset "rrule: get_relevant_steady_state_and_state_update (third_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_relevant_steady_state_and_state_update), + Val{:third_order}, typeof(params), typeof(RBC)}) + end + + @testset "rrule: get_relevant_steady_state_and_state_update (pruned_third_order)" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(get_relevant_steady_state_and_state_update), + Val{:pruned_third_order}, typeof(params), typeof(RBC)}) + end + + @testset "rrule: get_solution" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_solution), typeof(RBC), typeof(params)}) + end + + @testset "rrule: get_statistics" begin + jet_test_call(MacroModelling.rrule, + Tuple{typeof(MacroModelling.get_statistics), typeof(RBC), typeof(params)}) + end + +end + +# =========================================================================== +# Backus_Kehoe_Kydland_1992 — 2-country model with loop macros +# =========================================================================== +include("models/Backus_Kehoe_Kydland_1992.jl") +get_solution(Backus_Kehoe_Kydland_1992, algorithm = :first_order, silent = true) + +@testset verbose = true "JET hot-path analysis (Backus_Kehoe_Kydland_1992)" begin + run_first_order_jet_battery(Backus_Kehoe_Kydland_1992; + model_name = "Backus_Kehoe_Kydland_1992") +end + +# =========================================================================== +# NAWM_EAUS_2008 — large-scale 2-region DSGE +# =========================================================================== +include("../models/NAWM_EAUS_2008.jl") +get_solution(NAWM_EAUS_2008, algorithm = :first_order, silent = true) + +@testset verbose = true "JET hot-path analysis (NAWM_EAUS_2008)" begin + run_first_order_jet_battery(NAWM_EAUS_2008; + model_name = "NAWM_EAUS_2008") +end + +# =========================================================================== +# Gali_2015_chapter_3_obc — OBC model with max() and bound constraints +# =========================================================================== +include("../models/Gali_2015_chapter_3_obc.jl") +get_solution(Gali_2015_chapter_3_obc, algorithm = :first_order, silent = true) + +@testset verbose = true "JET hot-path analysis (Gali_2015_chapter_3_obc)" begin + run_first_order_jet_battery(Gali_2015_chapter_3_obc; + model_name = "Gali_2015_chapter_3_obc") +end diff --git a/test/test_models.jl b/test/test_models.jl index 25b20c193..91b6a35ca 100644 --- a/test/test_models.jl +++ b/test/test_models.jl @@ -1,3 +1,5 @@ +import Zygote, FiniteDifferences, Mooncake, DifferentiationInterface, ADTypes + # if !test_higher_order include("../models/Guerrieri_Iacoviello_2017.jl") SSvals = get_SS(Guerrieri_Iacoviello_2017) @@ -35,7 +37,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1), x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = false), model.parameter_values) @@ -43,12 +46,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(SW07_nonlinear) translate_dynare_file("SW07_nonlinear.mod") @@ -76,7 +80,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -84,12 +89,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Backus_Kehoe_Kydland_1992) translate_dynare_file("Backus_Kehoe_Kydland_1992.mod") @@ -118,7 +124,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end]), verbose = true), model.parameter_values[1:10]) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end]), verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values[1:10]) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end]), verbose = true), model.parameter_values[1:10])[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-3),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end]), verbose = true), model.parameter_values[1:10]) @@ -126,12 +133,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-3),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), vcat(x,model.parameter_values[11:end])), model.parameter_values[1:10]) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-2) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-2) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-2) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-2) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-2) write_to_dynare_file(NAWM_EAUS_2008) translate_dynare_file("NAWM_EAUS_2008.mod") @@ -161,7 +169,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -169,12 +178,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Baxter_King_1993) translate_dynare_file("Baxter_King_1993.mod") @@ -201,7 +211,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -209,12 +220,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Ireland_2004) translate_dynare_file("Ireland_2004.mod") @@ -243,7 +255,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(3,1, max_range = 1e-5), x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = false), model.parameter_values) @@ -251,12 +264,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-5), x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-4) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-4) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-5) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-5) write_to_dynare_file(QUEST3_2009) translate_dynare_file("QUEST3_2009.mod") # fix BGADJ1 = 0.001BGADJ2; @@ -285,7 +299,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -293,12 +308,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.forward_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-4) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-4) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-5) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-5) write_to_dynare_file(GNSS_2010) translate_dynare_file("GNSS_2010.mod") @@ -326,7 +342,8 @@ get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -334,12 +351,13 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Gali_Monacelli_2005_CITR) translate_dynare_file("Gali_Monacelli_2005_CITR.mod") @@ -369,7 +387,8 @@ # SS(model, parameters = [:alpha => 0.1, :trend_inflation => 1.5, :var_rho => 0.01]) # avoid the NaN error for finitediff in tests - back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) + back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) + zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # use forward_cdm so that parameter values stay positive. they would return NaN otherwise # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) @@ -379,13 +398,14 @@ local fin_grad = FiniteDifferences.grad(FiniteDifferences.forward_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-4) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-4) break end end # end - # @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-4) + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-4) write_to_dynare_file(Ascari_Sbordone_2014) translate_dynare_file("Ascari_Sbordone_2014.mod") @@ -422,7 +442,8 @@ simulated_data = simulate(model) get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) -back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) +back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) +zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -430,12 +451,13 @@ for i in 1:100 local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end -# @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) +# @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(SGU_2003_debt_premium) translate_dynare_file("SGU_2003_debt_premium.mod") @@ -473,7 +495,8 @@ simulated_data = simulate(model) get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) -back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) +back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) +zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -481,12 +504,13 @@ for i in 1:100 local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end -# @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) +# @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(JQ_2012_RBC) translate_dynare_file("JQ_2012_RBC.mod") @@ -525,7 +549,8 @@ simulated_data = simulate(model) get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) -back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) +back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) +zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -533,12 +558,13 @@ for i in 1:100 local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end -# @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) +# @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Gali_2015_chapter_3_nonlinear) translate_dynare_file("Gali_2015_chapter_3_nonlinear.mod") @@ -578,7 +604,8 @@ simulated_data = simulate(model) get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values, verbose = true) -back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) +back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) +zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -586,12 +613,13 @@ for i in 1:100 local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end -# @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) +# @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Ghironi_Melitz_2005) translate_dynare_file("Ghironi_Melitz_2005.mod") @@ -624,7 +652,8 @@ simulated_data = simulate(model) get_loglikelihood(model, simulated_data(observables, :, :simulate), model.parameter_values) -back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) +back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) +zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -632,12 +661,13 @@ for i in 1:100 local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-3) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-3) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-3) break end end -# @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-3) +# @test isapprox(back_grad, fin_grad[1], rtol = 1e-3) write_to_dynare_file(Caldara_et_al_2012) translate_dynare_file("Caldara_et_al_2012.mod") @@ -676,7 +706,8 @@ simulated_data = simulate(model) get_loglikelihood(model, simulated_data(observables, :, :simulate), verbose = true, model.parameter_values) -back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) +back_grad = DifferentiationInterface.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), ADTypes.AutoMooncake(config = nothing), model.parameter_values) +zygote_back_grad = Zygote.gradient(x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values)[1] # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x, verbose = true), model.parameter_values) @@ -684,15 +715,59 @@ for i in 1:100 local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(model, simulated_data(observables, :, :simulate), x), model.parameter_values) if isfinite(ℒ.norm(fin_grad)) println("Finite differences worked after $i iterations") - @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) break end end -# @test isapprox(back_grad[1], fin_grad[1], rtol = 1e-6) +# @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) write_to_dynare_file(Aguiar_Gopinath_2007) translate_dynare_file("Aguiar_Gopinath_2007.mod") include("Aguiar_Gopinath_2007.jl") get_solution(Aguiar_Gopinath_2007) Aguiar_Gopinath_2007 = nothing + + +include("../models/FRBUS.jl") +SSvals = get_SS(FRBUS, derivatives = false) + +# FRBUS is linearized around zero: all steady-state values should be zero +@test all(SSvals .== 0) + +# Test first-order solution shape and coefficients +sol = get_solution(FRBUS, algorithm = :first_order) + +@test size(sol) == (433, 428) +@test :Steady_state in axiskeys(sol, 1) +@test :rff₍₋₁₎ in axiskeys(sol, 1) +@test :fiscal_aerr₍ₓ₎ in axiskeys(sol, 1) +@test :rff in axiskeys(sol, 2) +@test :eco_l in axiskeys(sol, 2) + +@test isapprox(sol(:rff₍₋₁₎, :rff), 0.84575710864915, rtol = 1e-5) +@test isapprox(sol(:eco_l₍₋₁₎, :eco_l), 1.1848016760901816, rtol = 1e-5) +@test isapprox(sol(:ebfi_l₍₋₁₎, :ebfi_l), 1.27660626172, rtol = 1e-5) +@test isapprox(sol(:ex_l₍₋₁₎, :ex_l), 0.892272127137, rtol = 1e-5) + +# Test impulse responses to fiscal shock (impact and propagation) +irf_result = get_irf(FRBUS, algorithm = :first_order, shocks = [:fiscal_aerr], periods = 5) + +@test isapprox(irf_result(:rff, 1, :fiscal_aerr), 0.0144267064, rtol = 1e-4) +@test isapprox(irf_result(:xgap2, 1, :fiscal_aerr), 0.0961780423, rtol = 1e-4) +@test isapprox(irf_result(:eco_l, 1, :fiscal_aerr), 0.0010262957, rtol = 1e-4) +@test isapprox(irf_result(:debt_to_gdp, 1, :fiscal_aerr), 0.0065268971, rtol = 1e-4) + +@test isapprox(irf_result(:rff, 5, :fiscal_aerr), 0.1445122073, rtol = 1e-4) +@test isapprox(irf_result(:xgap2, 5, :fiscal_aerr), 0.3546553714, rtol = 1e-4) +@test isapprox(irf_result(:debt_to_gdp, 5, :fiscal_aerr), 0.0685009926, rtol = 1e-4) + +# Variance decomposition and loglikelihood tests are skipped for FRBUS: +# the Lyapunov equation does not converge (covariance matrix not found), +# so get_var_decomp, get_moments with standard deviations, and +# get_loglikelihood are not applicable for this model. + +write_to_dynare_file(FRBUS) +translate_dynare_file("FRBUS.mod") +FRBUS = nothing diff --git a/test/test_plots_1.jl b/test/test_plots_1.jl new file mode 100644 index 000000000..755dfad48 --- /dev/null +++ b/test/test_plots_1.jl @@ -0,0 +1,26 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +Random.seed!(1) + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "Backus_Kehoe_Kydland_1992" begin + include("../models/Backus_Kehoe_Kydland_1992.jl") + functionality_test(Backus_Kehoe_Kydland_1992, Caldara_et_al_2012_estim, plots = plots) +end +Backus_Kehoe_Kydland_1992 = nothing +GC.gc() + +@testset verbose = true "FS2000" begin + include("../models/FS2000.jl") + functionality_test(FS2000, Caldara_et_al_2012_estim, plots = plots) +end +FS2000 = nothing +GC.gc() diff --git a/test/test_plots_2.jl b/test/test_plots_2.jl new file mode 100644 index 000000000..2a8264967 --- /dev/null +++ b/test/test_plots_2.jl @@ -0,0 +1,33 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +Random.seed!(1) + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "Smets_Wouters_2003 with calibration equations" begin + include("../models/Smets_Wouters_2003.jl") + functionality_test(Smets_Wouters_2003, Caldara_et_al_2012_estim, plots = plots) +end +Smets_Wouters_2003 = nothing +GC.gc() + +@testset verbose = true "Smets and Wouters (2007) nonlinear" begin + include("../models/Smets_Wouters_2007.jl") + functionality_test(Smets_Wouters_2007, Caldara_et_al_2012_estim, plots = plots) +end +Smets_Wouters_2007 = nothing +GC.gc() + +@testset verbose = true "Smets and Wouters (2007) linear" begin + include("../models/Smets_Wouters_2007_linear.jl") + functionality_test(Smets_Wouters_2007_linear, Caldara_et_al_2012_estim, plots = plots) +end +Smets_Wouters_2007_linear = nothing +GC.gc() diff --git a/test/test_plots_3.jl b/test/test_plots_3.jl new file mode 100644 index 000000000..610527643 --- /dev/null +++ b/test/test_plots_3.jl @@ -0,0 +1,19 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +Random.seed!(1) + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "Gali 2015 ELB" begin + include("../models/Gali_2015_chapter_3_obc.jl") + functionality_test(Gali_2015_chapter_3_obc, Caldara_et_al_2012_estim, plots = plots) +end +Gali_2015_chapter_3_obc = nothing +GC.gc() diff --git a/test/test_plots_4.jl b/test/test_plots_4.jl new file mode 100644 index 000000000..c1ef778e9 --- /dev/null +++ b/test/test_plots_4.jl @@ -0,0 +1,202 @@ +using Test +using MacroModelling +import MacroModelling: clear_solution_caches! +using Random +import SpecialFunctions: erfcinv +using AxisKeys, SparseArrays +import Mooncake, FiniteDifferences, Zygote +import DifferentiationInterface, ADTypes +import LinearAlgebra as ℒ +Random.seed!(1234) + +include("functionality_tests.jl") + +plots = true +Random.seed!(1) + +include("models/Caldara_et_al_2012_estim.jl") + +@testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables" begin + include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags.jl") + functionality_test(m, Caldara_et_al_2012_estim, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), ADTypes.AutoMooncake(config = nothing), m.parameter_values) + zygote_back_grad = Zygote.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values)[1] + + # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences worked after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) + break + end + end + + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations, parameter definitions, special functions, variables in steady state, and leads/lag > 1 on endogenous and exogenous variables numerical SS" begin + include("models/RBC_CME_calibration_equations_and_parameter_definitions_lead_lags_numsolve.jl") + + functionality_test(m, Caldara_et_al_2012_estim, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), ADTypes.AutoMooncake(config = nothing), m.parameter_values) + zygote_back_grad = Zygote.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values)[1] + + # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1, max_range = 1e-4),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences worked after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) + break + end + end + + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations, parameter definitions, and special functions" begin + include("models/RBC_CME_calibration_equations_and_parameter_definitions_and_specfuns.jl") + functionality_test(m, Caldara_et_al_2012_estim, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), ADTypes.AutoMooncake(config = nothing), m.parameter_values) + zygote_back_grad = Zygote.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values)[1] + + # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x, verbose = true), m.parameter_values) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences worked after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) + break + end + end + + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations and parameter definitions" begin + include("models/RBC_CME_calibration_equations_and_parameter_definitions.jl") + functionality_test(m, Caldara_et_al_2012_estim, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), ADTypes.AutoMooncake(config = nothing), m.parameter_values) + zygote_back_grad = Zygote.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values)[1] + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences worked after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) + break + end + end + + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME with calibration equations" begin + include("models/RBC_CME_calibration_equations.jl") + functionality_test(m, Caldara_et_al_2012_estim, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), ADTypes.AutoMooncake(config = nothing), m.parameter_values) + zygote_back_grad = Zygote.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values)[1] + + # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences worked after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) + break + end + end + + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) +end +m = nothing +GC.gc() + +@testset verbose = true "RBC_CME" begin + include("models/RBC_CME.jl") + functionality_test(m, Caldara_et_al_2012_estim, plots = plots) + + observables = [:R, :k] + + Random.seed!(1) + simulated_data = simulate(m) + + get_loglikelihood(m, simulated_data(observables, :, :simulate), m.parameter_values) + + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), ADTypes.AutoMooncake(config = nothing), m.parameter_values) + zygote_back_grad = Zygote.gradient(x -> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values)[1] + + # fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4,1),x-> get_loglikelihood(m, simulated_data(observables, :, :simulate), x), m.parameter_values) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences worked after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) + @test isapprox(zygote_back_grad, fin_grad[1], rtol = 1e-6) + break + end + end + + # @test isapprox(back_grad, fin_grad[1], rtol = 1e-6) +end +m = nothing +GC.gc() diff --git a/test/test_plots_5.jl b/test/test_plots_5.jl new file mode 100644 index 000000000..4fb3f2814 --- /dev/null +++ b/test/test_plots_5.jl @@ -0,0 +1,491 @@ +using Test +using MacroModelling +using Random +import StatsPlots +using DelimitedFiles +using Dates +using AxisKeys, SparseArrays + +include("test_helpers.jl") + +Random.seed!(1) + +@testset verbose = true "SW07 estim" begin + include("../models/Smets_Wouters_2007.jl") + + # load data + dat, header = readdlm("data/usmodel.csv", ',', header = true) + dat = Float64.(dat) + names = vec(Symbol.(strip.(header))) + + # load data + 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 + + # Subsample + # subset observables in data + sample_idx = 47:230 # 1960Q1-2004Q4 + + data = data(observables_old, sample_idx) + + # declare observables as written in model + observables = [:dy, :dc, :dinve, :labobs, :pinfobs, :dwobs, :robs] # note that :dw was renamed to :dwobs in linear model in order to avoid confusion with nonlinear model + + data = rekey(data, :Variable => observables) + + data_rekey = rekey(data, :Time => quarterly_dates(Date(1960, 1, 1), size(data,2))) + + + plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 3, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 3, :calfa => 0.28]) + + + plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], filter = :inversion) + + + plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], filter = :inversion) + + plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], smooth = false) + + + plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], smooth = false) + + plot_model_estimates!(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24], smooth = false, presample_periods = 50) + + + plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data[:,20:end], parameters = [:csadjcost => 6, :calfa => 0.24]) + + + plot_model_estimates(Smets_Wouters_2007, data_rekey, parameters = [:csadjcost => 6, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data_rekey, parameters = [:csadjcost => 5, :calfa => 0.24]) + + + plot_model_estimates(Smets_Wouters_2007, data, parameters = [:csadjcost => 6, :calfa => 0.24]) + + plot_model_estimates!(Smets_Wouters_2007, data_rekey, parameters = [:csadjcost => 5, :calfa => 0.24]) + + # FS2000 model and data + include("../models/FS2000.jl") + + # load data + 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)) + + # subset observables in data + dataFS2000 = dataFS2000(observables,:) + + dataFS2000_rekey = rekey(dataFS2000, :Time => quarterly_dates(Date(1950, 1, 1), size(dataFS2000,2))) + + plot_model_estimates(FS2000, dataFS2000) + + plot_model_estimates(FS2000, dataFS2000_rekey[:,1:10]) + + plot_shock_decomposition(FS2000, dataFS2000_rekey[:,1:10]) + + plot_shock_decomposition(FS2000, dataFS2000_rekey) + + + dataFS2000_rekey2 = rekey(dataFS2000, :Time => 1:1:size(dataFS2000,2)) + + plot_shock_decomposition(FS2000, dataFS2000) + + plot_shock_decomposition(FS2000, dataFS2000_rekey2) + + + plot_model_estimates(FS2000, dataFS2000_rekey, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + plot_model_estimates!(Smets_Wouters_2007, data_rekey) + + + plot_model_estimates(FS2000, dataFS2000_rekey, parameters = :alp => 0.356, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + plot_model_estimates!(Smets_Wouters_2007, data_rekey) + + plot_model_estimates!(FS2000, dataFS2000_rekey, parameters = :alp => 0.3, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + + plot_model_estimates!(Smets_Wouters_2007, data_rekey, parameters = :csigma => 0.3) + + plot_model_estimates(FS2000, dataFS2000_rekey, parameters = :alp => 0.356, shock_decomposition = true, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + + estims = get_estimated_variables(Smets_Wouters_2007, data) + + plot_irf(Smets_Wouters_2007, shocks = :em, shock_size = 10) + + plot_irf!(Smets_Wouters_2007,initial_state = collect(estims[:,end]), shocks = :none, plot_type = :stack) + + plot_irf!(Smets_Wouters_2007, shocks = [:em, :ea], negative_shock = true, plot_type = :stack) + + shock_mat = randn(Smets_Wouters_2007.constants.post_model_macro.nExo,3) + + plot_irf!(Smets_Wouters_2007, shocks = shock_mat, plot_type = :stack) + + plot_irf!(Smets_Wouters_2007, shocks = shock_mat, plot_type = :stack) + + + plot_irf(Smets_Wouters_2007, shocks = :em, periods = 5, variables = [:y, :k, :c]) + + plot_irf!(FS2000, shocks = :e_m, periods = 5, plot_type = :stack, shock_size = 10, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w), variables = [:y, :k, :c]) + + + plot_irf(Smets_Wouters_2007, shocks = :em, periods = 5) + + plot_irf!(FS2000, shocks = :e_m, periods = 5, plot_type = :stack, shock_size = 10, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + plot_irf!(FS2000, shocks = [:e_m, :e_a], shock_size = 20, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + plot_irf!(Smets_Wouters_2007, shocks = [:em, :ea], shock_size = 0.5) + + + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) + cndtns_lvl[1,8] = 1.4 + + plot_conditional_forecast(Smets_Wouters_2007, cndtns_lvl, initial_state = collect(estims[:,end])) + + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,4), Variables = [:pinfobs], Periods = 1:4) + cndtns_lvl[1,4] = 2 + + plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, plot_type = :stack) + + + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) + cndtns_lvl[1,8] = 1.45 + + plot_conditional_forecast!(FS2000, cndtns_lvl, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,4), Variables = [:y], Periods = 1:4) + cndtns_lvl[1,4] = 2.01 + + plot_conditional_forecast!(FS2000, cndtns_lvl, plot_type = :stack, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + # conditons on #3 is nothing which makes sense since it is not showing + + shock_mat = sprandn(Smets_Wouters_2007.constants.post_model_macro.nExo, 10, .1) + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,4), Variables = [:pinfobs], Periods = 1:4) + cndtns_lvl[1,4] = 2 + + plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, shocks = shock_mat, plot_type = :stack) + + + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) + cndtns_lvl[1,8] = 1.4 + + shock_mat = sprandn(Smets_Wouters_2007.constants.post_model_macro.nExo, 10, .1) + + plot_conditional_forecast(Smets_Wouters_2007, cndtns_lvl, shocks = shock_mat, label = "SW07 w shocks", variables = [:y, :k, :c]) + + plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, variables = [:y,:w]) + + plot_conditional_forecast!(FS2000, cndtns_lvl, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + shock_mat = sprandn(FS2000.constants.post_model_macro.nExo, 10, .1) + + plot_conditional_forecast!(FS2000, cndtns_lvl, shocks = shock_mat, label = :rand_shocks, rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + + plot_solution(FS2000, :k) + + plot_solution!(FS2000, :k, algorithm = :second_order) + + + plot_solution(Smets_Wouters_2007, :pinf) + + plot_solution!(Smets_Wouters_2007, :pinf, algorithm = :second_order) + + + plot_solution(FS2000, :y) + + plot_solution!(Smets_Wouters_2007, :y, variables = [:y, :k, :c]) + + plot_solution!(Smets_Wouters_2007, :y, algorithm = :second_order, variables = [:y, :k, :c]) + + + # tol-only and tol-varying tests (struct and NamedTuple formulations) + plot_model_estimates(Smets_Wouters_2007, data, + tol = Tolerances(nsss = NsssTolerances(acceptance_tol = 1e-10))) + + plot_model_estimates!(Smets_Wouters_2007, data, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),))) + + plot_model_estimates!(Smets_Wouters_2007, data, + tol = Tolerances(first_order = FirstOrderTolerances(qme = SolverTolerances(acceptance_tol = 1e-12)))) + + + plot_shock_decomposition(FS2000, dataFS2000_rekey, + tol = Tolerances(first_order = FirstOrderTolerances(lyapunov = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_shock_decomposition(FS2000, dataFS2000_rekey, + tol = Tolerances(nsss = (acceptance_tol = 1e-10,))) + + + plot_irf(Smets_Wouters_2007, shocks = :em, + tol = Tolerances(first_order = FirstOrderTolerances(qme = SolverTolerances(acceptance_tol = 1e-12)))) + + plot_irf!(Smets_Wouters_2007, shocks = :em, + tol = Tolerances(nsss = (xtol = 1e-14,))) + + plot_irf!(FS2000, shocks = :e_m, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),)), + rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) + cndtns_lvl[1,8] = 1.4 + + plot_conditional_forecast(Smets_Wouters_2007, cndtns_lvl, + tol = Tolerances(nsss = NsssTolerances(ftol = 1e-16))) + + plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),))) + + + plot_solution(FS2000, :y, + tol = Tolerances(nsss = NsssTolerances(acceptance_tol = 1e-10))) + + plot_solution!(FS2000, :y, + tol = Tolerances(first_order = FirstOrderTolerances(lyapunov = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_solution!(FS2000, :y, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),))) + + + plot_solution(Smets_Wouters_2007, :y, algorithm = :second_order, + tol = Tolerances(second_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_solution!(Smets_Wouters_2007, :y, algorithm = :second_order, + tol = Tolerances(second_order = (sylvester = (acceptance_tol = 1e-14,), lyapunov = (acceptance_tol = 1e-14,)))) + + + # combined tol + other argument tests + plot_model_estimates(Smets_Wouters_2007, data, + parameters = [:csadjcost => 5, :calfa => 0.22], + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),))) + + plot_model_estimates!(Smets_Wouters_2007, data, + parameters = [:csadjcost => 3, :calfa => 0.28], filter = :inversion, + tol = Tolerances(nsss = NsssTolerances(acceptance_tol = 1e-10))) + + plot_shock_decomposition(FS2000, dataFS2000_rekey, + tol = Tolerances(first_order = FirstOrderTolerances(qme = SolverTolerances(acceptance_tol = 1e-12))), + rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + plot_irf(Smets_Wouters_2007, shocks = :em, periods = 10, variables = [:y, :k, :c], + tol = Tolerances(first_order = FirstOrderTolerances(qme = SolverTolerances(acceptance_tol = 1e-12)))) + + plot_irf!(FS2000, shocks = :e_m, shock_size = 10, periods = 5, + tol = Tolerances(nsss = (xtol = 1e-14,)), + rename_dictionary = Dict(:e_a => :ea, :e_m => :em, :R => :r, :W => :w)) + + cndtns_lvl = KeyedArray(Matrix{Union{Nothing, Float64}}(undef,1,8), Variables = [:y], Periods = 1:8) + cndtns_lvl[1,8] = 1.35 + + plot_conditional_forecast(Smets_Wouters_2007, cndtns_lvl, + initial_state = collect(estims[:,end]), + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),))) + + plot_conditional_forecast!(Smets_Wouters_2007, cndtns_lvl, variables = [:y, :k], + tol = Tolerances(nsss = NsssTolerances(ftol = 1e-16))) + + plot_solution(FS2000, :k, algorithm = :second_order, + tol = Tolerances(second_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_solution!(Smets_Wouters_2007, :pinf, algorithm = :second_order, variables = [:pinf, :y], + tol = Tolerances(second_order = (sylvester = (acceptance_tol = 1e-14,),))) + +end + +# multiple models +@testset verbose = true "Gali 2015 ELB plots" begin + include("../models/Gali_2015_chapter_3_obc.jl") + + + Random.seed!(14) + plot_simulation(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0, ignore_obc = true) + + Random.seed!(14) + plot_simulation!(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0) + + Random.seed!(14) + plot_simulation!(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0025) + + + Random.seed!(13) + plot_simulation(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, + # periods = 40, + parameters = :R̄ => 1.0, ignore_obc = true) + + Random.seed!(13) + plot_simulation!(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, + periods = 40, + parameters = :R̄ => 1.0) + + + plot_irf(Gali_2015_chapter_3_obc, parameters = :R̄ => 1.0) + + plot_irf!(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, parameters = :R̄ => 1.0) + + + plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.0) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.5) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 0.5) + + + plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.0) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, generalised_irf = true) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, ignore_obc = true) + + + plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, algorithm = :pruned_second_order) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, algorithm = :pruned_second_order, ignore_obc = true) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 1.0, algorithm = :pruned_second_order, ignore_obc = true, generalised_irf = true) + + + # tol-only and tol-varying tests + Random.seed!(14) + plot_simulation(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0, + tol = Tolerances(first_order = FirstOrderTolerances(lyapunov = SolverTolerances(acceptance_tol = 1e-14)))) + + Random.seed!(14) + plot_simulation!(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0, + tol = Tolerances(first_order = (qme = (acceptance_tol = 1e-12,),))) + + + plot_irf(Gali_2015_chapter_3_obc, + tol = Tolerances(nsss = NsssTolerances(acceptance_tol = 1e-10))) + + plot_irf!(Gali_2015_chapter_3_obc, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),))) + + + # combined tol + other argument tests + Random.seed!(14) + plot_simulation(Gali_2015_chapter_3_obc, periods = 40, parameters = :R̄ => 1.0025, ignore_obc = true, + tol = Tolerances(first_order = FirstOrderTolerances(lyapunov = SolverTolerances(acceptance_tol = 1e-14)))) + + Random.seed!(13) + plot_simulation!(Gali_2015_chapter_3_obc, algorithm = :pruned_second_order, periods = 40, parameters = :R̄ => 1.0, + tol = Tolerances(first_order = (qme = (acceptance_tol = 1e-12,),))) + + plot_irf(Gali_2015_chapter_3_obc, parameters = :σ => 1.5, + tol = Tolerances(nsss = NsssTolerances(acceptance_tol = 1e-10))) + + plot_irf!(Gali_2015_chapter_3_obc, parameters = :σ => 0.5, algorithm = :pruned_second_order, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),), + second_order = (sylvester = (acceptance_tol = 1e-14,),))) +end + +@testset verbose = true "Caldara et al 2012 plots" begin + include("../models/Caldara_et_al_2012.jl") + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :second_order) + + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_second_order, generalised_irf = true, generalised_irf_draws = 1000) + + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order) + + + plot_irf(Caldara_et_al_2012, algorithm = :second_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :third_order) + + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, generalised_irf = true) + + + plot_irf(Caldara_et_al_2012, algorithm = :third_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :third_order, generalised_irf = true) + + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, shock_size = 2) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, shock_size = 3) + + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = :ψ => 0.8) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = :ψ => 1.5) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = :ψ => 2.5) + + + plot_irf(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.3]) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.25]) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.35]) + + + # tol-only and tol-varying tests + plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order, + tol = Tolerances(second_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_second_order, + tol = Tolerances(second_order = (sylvester = (acceptance_tol = 1e-14,), lyapunov = (acceptance_tol = 1e-14,)))) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, + tol = Tolerances(third_order = (sylvester = (acceptance_tol = 1e-14,),))) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, + tol = Tolerances(nsss = (xtol = 1e-14,), + third_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14), + lyapunov = SolverTolerances(acceptance_tol = 1e-14)))) + + + # combined tol + other argument tests + plot_irf(Caldara_et_al_2012, algorithm = :pruned_second_order, parameters = :ψ => 0.8, + tol = Tolerances(second_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, shock_size = 2, + tol = Tolerances(third_order = (sylvester = (acceptance_tol = 1e-14,), lyapunov = (acceptance_tol = 1e-14,)))) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, parameters = [:ψ => 0.5, :ζ => 0.3], + tol = Tolerances(nsss = (xtol = 1e-14,), + third_order = HigherOrderTolerances(sylvester = SolverTolerances(acceptance_tol = 1e-14)))) + + plot_irf!(Caldara_et_al_2012, algorithm = :pruned_third_order, generalised_irf = true, + tol = Tolerances(first_order = (lyapunov = (acceptance_tol = 1e-14,),), + third_order = (sylvester = (acceptance_tol = 1e-14,),))) +end diff --git a/test/test_pruned_2nd_order_estimation.jl b/test/test_pruned_2nd_order_estimation.jl index cb99e67ea..44696c381 100644 --- a/test/test_pruned_2nd_order_estimation.jl +++ b/test/test_pruned_2nd_order_estimation.jl @@ -1,19 +1,30 @@ +using Test using MacroModelling import Turing -import ADTypes: AutoZygote -import Turing: NUTS, sample, logpdf +import Mooncake +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences +import Turing: NUTS, sample import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys + +using FlexiChains +include("test_helpers.jl") 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,:) @@ -32,16 +43,14 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) llh = get_loglikelihood(m, data, all_params, algorithm = algorithm, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end @@ -49,21 +58,39 @@ end Random.seed!(30) -n_samples = 500 +n_samples = 1000 + +samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :pruned_second_order, -Inf), NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams((; all_params = FS2000.parameter_values))) -samps = @time sample(FS2000_loglikelihood_function(data, FS2000, :pruned_second_order, -Inf), NUTS(adtype = AutoZygote()), n_samples, progress = true, initial_params = FS2000.parameter_values) +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") -println("Mean variable values (Zygote): $(mean(samps).nt.mean)") +sample_nuts = collect(values(FlexiChains.mean(samps); parameters_only = true)) -sample_nuts = mean(samps).nt.mean +@testset "Mooncake vs FiniteDifferences gradient (pruned 2nd order)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(FS2000, data, x, algorithm = :pruned_second_order), ADTypes.AutoMooncake(config = nothing), FS2000.parameter_values) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + 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, 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 @@ -108,10 +135,10 @@ sample_nuts = mean(samps).nt.mean # n_rounds = 6, # multithreaded = false) -# samps = MCMCChains.Chains(Pigeons.get_sample(pt)) +# samps = pigeons_flexichain(Pigeons.sample_array(pt), Pigeons.sample_names(pt)) -# println(mean(samps).nt.mean) +# println(collect(values(FlexiChains.mean(samps); parameters_only = true))) # Random.seed!(30) @@ -216,7 +243,7 @@ sample_nuts = mean(samps).nt.mean # 1 # @testset "Estimation results" begin # @test isapprox(sol.minimum, -1343.7491257498598, rtol = eps(Float32)) -# @test isapprox(mean(samps).nt.mean, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) +# @test isapprox(collect(values(FlexiChains.mean(samps); parameters_only = true)), [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) # end diff --git a/test/test_pruned_2nd_order_estimation_pigeons.jl b/test/test_pruned_2nd_order_estimation_pigeons.jl index 1e1087cf3..cea105ccd 100644 --- a/test/test_pruned_2nd_order_estimation_pigeons.jl +++ b/test/test_pruned_2nd_order_estimation_pigeons.jl @@ -2,19 +2,22 @@ using MacroModelling using Test import Turing import Pigeons -import Turing: logpdf -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL +include("test_helpers.jl") + 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,:) @@ -33,7 +36,7 @@ dists = [ ] Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() llh = get_loglikelihood(m, @@ -41,16 +44,14 @@ Turing.@model function FS2000_loglikelihood_function(data, m, algorithm, on_fail all_params, algorithm = algorithm, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end 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 +72,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 +99,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..eb6f8a6d3 100644 --- a/test/test_pruned_3rd_order_estimation.jl +++ b/test/test_pruned_3rd_order_estimation.jl @@ -1,15 +1,26 @@ +using Test using MacroModelling import Turing -import ADTypes: AutoZygote -import Turing: NUTS, sample, logpdf, PG, IS +import Mooncake +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences +import Turing: NUTS, sample import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys + +using FlexiChains +include("test_helpers.jl") # 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] @@ -44,16 +55,14 @@ dists = [ ] Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) llh = get_loglikelihood(m, data, all_params, algorithm = :pruned_third_order, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end @@ -72,29 +81,47 @@ mode_estimateNM = Turing.maximum_a_posteriori(Caldara_et_al_2012_loglikelihood, Optim.NelderMead(), iterations = 100, # show_trace = true, - initial_params = Caldara_et_al_2012_estim.parameter_values) + initial_params = Turing.InitFromParams((; all_params = Caldara_et_al_2012_estim.parameter_values))) mode_estimateLBFGS = Turing.maximum_a_posteriori(Caldara_et_al_2012_loglikelihood, Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), - adtype = AutoZygote(), + adtype = AutoMooncake(; config=nothing), iterations = 100, # show_trace = true, - initial_params = mode_estimateNM.values) + initial_params = Turing.InitFromParams(mode_estimateNM)) -init_params = mode_estimateLBFGS.values |> collect +init_params = collect(mode_estimateLBFGS.params.data.all_params) -println("Mode variable values (L-BFGS): $init_params") +println("Mode variable values (L-BFGS): $(mode_estimateLBFGS.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 = AutoMooncake(; config=nothing)), n_samples, progress = true, initial_params = Turing.InitFromParams(mode_estimateLBFGS)) + +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (Mooncake): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") -println("Mean variable values (Zygote): $(mean(samps).nt.mean)") +sample_nuts = collect(values(FlexiChains.mean(samps); parameters_only = true)) -sample_nuts = mean(samps).nt.mean +@testset "Mooncake vs FiniteDifferences gradient (pruned 3rd order)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(Caldara_et_al_2012_estim, data, x, algorithm = :pruned_third_order), ADTypes.AutoMooncake(config = nothing), init_params) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + 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, fin_grad[1], rtol = 1e-4) + break + end + end +end @@ -102,7 +129,7 @@ sample_nuts = mean(samps).nt.mean # 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) @@ -136,10 +163,10 @@ sample_nuts = mean(samps).nt.mean # n_rounds = 6, # multithreaded = false) -# samps = MCMCChains.Chains(Pigeons.get_sample(pt)) +# samps = pigeons_flexichain(Pigeons.sample_array(pt), Pigeons.sample_names(pt)) -# println(mean(samps).nt.mean) +# println(collect(values(FlexiChains.mean(samps); parameters_only = true))) # Random.seed!(30) @@ -244,7 +271,7 @@ sample_nuts = mean(samps).nt.mean # 1 # @testset "Estimation results" begin # @test isapprox(sol.minimum, -1343.7491257498598, rtol = eps(Float32)) -# @test isapprox(mean(samps).nt.mean, [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) +# @test isapprox(collect(values(FlexiChains.mean(samps); parameters_only = true)), [0.40248024934137033, 0.9905235783816697, 0.004618184988033483, 1.014268215459915, 0.8459140293740781, 0.6851143053372912, 0.0025570276255960107, 0.01373547787288702, 0.003343985776134218], rtol = 1e-2) # end diff --git a/test/test_pruned_3rd_order_estimation_pigeons.jl b/test/test_pruned_3rd_order_estimation_pigeons.jl index 7570e2bb6..25c95ad72 100644 --- a/test/test_pruned_3rd_order_estimation_pigeons.jl +++ b/test/test_pruned_3rd_order_estimation_pigeons.jl @@ -2,15 +2,18 @@ using MacroModelling using Test import Turing import Pigeons -import Turing: logpdf, PG, IS -using Random, CSV, DataFrames, MCMCChains, AxisKeys +using Random, DelimitedFiles, MCMCChains, AxisKeys import DynamicPPL +include("test_helpers.jl") + # 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] @@ -45,7 +48,7 @@ dists = [ ] Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_failure_loglikelihood; verbose = false) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) if DynamicPPL.leafcontext(__context__) !== DynamicPPL.PriorContext() llh = get_loglikelihood(m, @@ -53,16 +56,14 @@ Turing.@model function Caldara_et_al_2012_loglikelihood_function(data, m, on_fai all_params, algorithm = :pruned_third_order, on_failure_loglikelihood = on_failure_loglikelihood) - if verbose - @info "Loglikelihood: $llh and prior llh: $(Turing.logpdf(Turing.arraydist(dists), all_params)) with params $all_params" - end + maybe_print_loglikelihood(verbose, llh, dists, all_params) Turing.@addlogprob! llh end 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 +91,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 +117,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_rrule_robustness.jl b/test/test_rrule_robustness.jl new file mode 100644 index 000000000..6626fadb5 --- /dev/null +++ b/test/test_rrule_robustness.jl @@ -0,0 +1,224 @@ +using Test +using MacroModelling +import Zygote +import ForwardDiff +import FiniteDifferences +import LinearAlgebra as ℒ + +using Random, AxisKeys + +import MacroModelling: clear_solution_caches! + +# ────────────────────────────────────────────────────────────────────────────── +# Helper: Zygote-compatible full parameter vector builder +# ────────────────────────────────────────────────────────────────────────────── +function make_param_builder(true_params::Vector{Float64}, est_idx::Vector{Int}) + n = length(true_params) + ep = zeros(Int, n) + for (j, i) in enumerate(est_idx) + ep[i] = j + end + fp = copy(true_params) + return function(x) + T = eltype(x) + map(1:n) do i + ep[i] > 0 ? x[ep[i]] : T(fp[i]) + end + end +end + +# ────────────────────────────────────────────────────────────────────────────── +# Load models +# ────────────────────────────────────────────────────────────────────────────── +include("../models/RBC_baseline.jl") +include("../models/FS2000.jl") +include("../models/Ireland_2004.jl") + +# ══════════════════════════════════════════════════════════════════════════════ +# Test 1: IRF rrule gradient across all 5 algorithms (RBC_baseline) +# ══════════════════════════════════════════════════════════════════════════════ +@testset "IRF rrule gradient - all algorithms (RBC_baseline)" begin + m = RBC_baseline + # Parameters: σᶻ(1), σᵍ(2), σ(3), i_y(4), k_y(5), ρᶻ(6), ρᵍ(7), g_y(8), α(9) + est_idx = [1, 2, 6, 7] # σᶻ, σᵍ, ρᶻ, ρᵍ + build_params = make_param_builder(m.parameter_values, est_idx) + test_point = m.parameter_values[est_idx] + + for alg in [:first_order, :pruned_second_order, :pruned_third_order, + :second_order, :third_order] + @testset "$alg" begin + # Invalidate caches to prevent Dual contamination between algorithms + MacroModelling.invalidate_cache_validity!(m) + + f = x -> begin + all_p = build_params(x) + irf_v = get_irf(m, all_p, algorithm = alg, periods = 3) + return sum(irf_v) + end + + zy_grad = Zygote.gradient(f, test_point)[1] + MacroModelling.invalidate_cache_validity!(m) + fd_grad = ForwardDiff.gradient(f, test_point) + MacroModelling.invalidate_cache_validity!(m) + fi_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(5, 1), f, test_point)[1] + + @test all(isfinite, zy_grad) + @test all(isfinite, fd_grad) + @test all(isfinite, fi_grad) + + zy_fd = maximum(abs.(zy_grad .- fd_grad) ./ max.(abs.(fd_grad), 1e-10)) + zy_fi = maximum(abs.(zy_grad .- fi_grad) ./ max.(abs.(fi_grad), 1e-10)) + fd_fi = maximum(abs.(fd_grad .- fi_grad) ./ max.(abs.(fi_grad), 1e-10)) + println(" IRF $alg: Zy-FD=$zy_fd Zy-FI=$zy_fi FD-FI=$fd_fi") + @test zy_fd < 1e-6 + @test zy_fi < 1e-4 # FiniteDiff has lower precision + @test fd_fi < 1e-4 + end + end +end + +# ══════════════════════════════════════════════════════════════════════════════ +# Regression: clearing solution caches must not retain an evicted solver seed +# ══════════════════════════════════════════════════════════════════════════════ +@testset "IRF derivative after NSSS solver cache eviction" begin + include("../models/Smets_Wouters_2003.jl") + + m = Smets_Wouters_2003 + parameters = copy(m.parameter_values) + initial_state = get_irf( + m, + algorithm = :first_order, + shocks = :none, + levels = true, + variables = :all, + periods = 1, + ) |> vec + + get_irf(m, parameters) + for i in 1:(m.caches.solver.capacity + 5) + push!(m.caches.solver, [fill(float(i), 1), fill(float(i), 1), fill(float(i), length(parameters))]) + end + + clear_solution_caches!(m, :first_order) + @test length(m.caches.solver) == 1 + @test all(isinf, m.caches.solver[end][end]) + + deriv_for = ForwardDiff.jacobian(x -> get_irf(m, x, initial_state = initial_state)[:, end, 1], parameters) + deriv_fin = FiniteDifferences.jacobian( + FiniteDifferences.central_fdm(5, 1, max_range = 1e-4), + x -> begin + clear_solution_caches!(m, :first_order) + get_irf(m, x, initial_state = initial_state)[:, end, 1] + end, + parameters, + )[1] + + @test isapprox(deriv_for, deriv_fin, rtol = 1e-5) +end + +# ══════════════════════════════════════════════════════════════════════════════ +# Test 2: Combined objective gradient - pruned_second_order on multiple models +# ══════════════════════════════════════════════════════════════════════════════ +model_configs = [ + ( + name = "RBC_baseline", + model = RBC_baseline, + observables = [:y, :c], + # σᶻ(1), σᵍ(2), ρᶻ(6), ρᵍ(7) + est_idx = [1, 2, 6, 7], + nsss_vars = [:y, :c], + moment_vars = [:y, :c], + ), + ( + name = "FS2000", + model = FS2000, + observables = [:log_gy_obs, :log_gp_obs], + # rho(5), z_e_a(8), z_e_m(9) + est_idx = [5, 8, 9], + nsss_vars = [:y, :c], + moment_vars = [:y, :c], + ), + ( + name = "Ireland_2004", + model = Ireland_2004, + observables = [:ĝ, :ŷ, :π̂, :r̂], + # ρᵃ(9), ρᵉ(10), σʳ(11), σᵃ(12), σᵉ(13), σᶻ(14) + est_idx = [9, 10, 11, 12, 13, 14], + nsss_vars = [:ŷ, :π̂], + moment_vars = [:ŷ, :π̂], + ), +] + +@testset "Combined objective gradient - pruned_second_order" begin + alg = :pruned_second_order + + for cfg in model_configs + @testset "$(cfg.name)" begin + m = cfg.model + MacroModelling.invalidate_cache_validity!(m) + build_params = make_param_builder(m.parameter_values, cfg.est_idx) + test_point = m.parameter_values[cfg.est_idx] + + # Simulate data at true parameter values + Random.seed!(42) + sim = simulate(m, algorithm = alg) + data = sim(cfg.observables, :, :simulate) + + # Compute targets at true parameters + target_stats = get_statistics(m, m.parameter_values, + non_stochastic_steady_state = cfg.nsss_vars, + mean = cfg.moment_vars, + standard_deviation = cfg.moment_vars, + algorithm = alg) + target_nsss = target_stats[:non_stochastic_steady_state] + target_mean = target_stats[:mean] + target_std = target_stats[:standard_deviation] + + target_irf_full = get_irf(m, m.parameter_values, algorithm = alg, periods = 3) + irf_var_idx = sort(MacroModelling.parse_variables_input_to_index(cfg.observables[1], m)) + target_irf = target_irf_full[irf_var_idx, 1, 1] + + # Combined objective exercising all 4 differentiable functions + f = x -> begin + all_p = build_params(x) + + llh = get_loglikelihood(m, data, all_p, + algorithm = alg, on_failure_loglikelihood = -Inf) + + stats_n = get_statistics(m, all_p, + non_stochastic_steady_state = cfg.nsss_vars, algorithm = alg) + llh -= sum((stats_n[:non_stochastic_steady_state] .- target_nsss).^2) + + stats_m = get_statistics(m, all_p, + mean = cfg.moment_vars, standard_deviation = cfg.moment_vars, + algorithm = alg) + llh -= sum((stats_m[:mean] .- target_mean).^2) + llh -= sum((stats_m[:standard_deviation] .- target_std).^2) + + irf_v = get_irf(m, all_p, algorithm = alg, periods = 3) + llh -= sum((irf_v[irf_var_idx, 1, 1] .- target_irf).^2) + + return llh + end + + zy_grad = Zygote.gradient(f, test_point)[1] + MacroModelling.invalidate_cache_validity!(m) + fd_grad = ForwardDiff.gradient(f, test_point) + MacroModelling.invalidate_cache_validity!(m) + fi_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(5, 1), f, test_point)[1] + + @test !isnothing(zy_grad) + @test all(isfinite, zy_grad) + @test all(isfinite, fd_grad) + @test all(isfinite, fi_grad) + + zy_fd = maximum(abs.(zy_grad .- fd_grad) ./ max.(abs.(fd_grad), 1e-10)) + zy_fi = maximum(abs.(zy_grad .- fi_grad) ./ max.(abs.(fi_grad), 1e-10)) + fd_fi = maximum(abs.(fd_grad .- fi_grad) ./ max.(abs.(fi_grad), 1e-10)) + println(" $(cfg.name): Zy-FD=$zy_fd Zy-FI=$zy_fi FD-FI=$fd_fi") + @test zy_fd < 1e-6 + @test zy_fi < 1e-4 + @test fd_fi < 1e-4 + end + end +end diff --git a/test/test_standalone_function.jl b/test/test_standalone_function.jl index 7094ba2f5..1430ce6ea 100644 --- a/test/test_standalone_function.jl +++ b/test/test_standalone_function.jl @@ -2,11 +2,13 @@ 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 +using FiniteDifferences +using Zygote, Mooncake import Optim, LineSearches +import DifferentiationInterface, ADTypes Random.seed!(3) @@ -66,23 +68,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 +88,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 +166,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 +566,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], @@ -622,6 +624,18 @@ RBC_CME = nothing @test isapprox(forw_grad,fin_grad,rtol = 1e-5) + solution_norm_obj = x -> ℒ.norm(get_solution(RBC_CME, x)[2][1]) + forw_grad = ForwardDiff.gradient(solution_norm_obj, Float64.(RBC_CME.parameter_values)) + reverse_grad = DifferentiationInterface.gradient(solution_norm_obj, ADTypes.AutoMooncake(config = nothing), Float64.(RBC_CME.parameter_values)) + zygote_reverse_grad = Zygote.gradient(solution_norm_obj, Float64.(RBC_CME.parameter_values))[1] + fin_grad = FiniteDifferences.grad(central_fdm(4,1), solution_norm_obj, RBC_CME.parameter_values)[1] + + @test isapprox(forw_grad,reverse_grad,rtol = 1e-6) + @test isapprox(forw_grad,zygote_reverse_grad,rtol = 1e-6) + @test isapprox(forw_grad,fin_grad,rtol = 1e-6) + + + Random.seed!(3) @@ -631,11 +645,13 @@ RBC_CME = nothing @test isapprox(425.7689804539224, get_loglikelihood(RBC_CME, data(observables), RBC_CME.parameter_values),rtol = 1e-5) forw_grad = ForwardDiff.gradient(x -> get_loglikelihood(RBC_CME, data(observables), x), Float64.(RBC_CME.parameter_values)) - reverse_grad = Zygote.gradient(x -> get_loglikelihood(RBC_CME, data(observables), x), Float64.(RBC_CME.parameter_values))[1] + reverse_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(RBC_CME, data(observables), x), ADTypes.AutoMooncake(config = nothing), Float64.(RBC_CME.parameter_values)) + zygote_reverse_grad = Zygote.gradient(x -> get_loglikelihood(RBC_CME, data(observables), x), Float64.(RBC_CME.parameter_values))[1] fin_grad = FiniteDifferences.grad(central_fdm(4,1),x -> get_loglikelihood(RBC_CME, data(observables), x), RBC_CME.parameter_values)[1] @test isapprox(forw_grad,fin_grad, rtol = 1e-6) + @test isapprox(forw_grad,zygote_reverse_grad, rtol = 1e-6) @test isapprox(forw_grad,reverse_grad, rtol = 1e-6) RBC_CME = nothing diff --git a/test/test_sw07_estimation.jl b/test/test_sw07_estimation.jl index 46363a0ed..90dee2b73 100644 --- a/test/test_sw07_estimation.jl +++ b/test/test_sw07_estimation.jl @@ -1,15 +1,25 @@ +using Test using MacroModelling -import ADTypes: AutoZygote +import Mooncake +import ADTypes +import ADTypes: AutoMooncake +import DifferentiationInterface +import FiniteDifferences import Turing -import Turing: NUTS, sample, logpdf -import Optim, LineSearches -using Random, CSV, DataFrames, MCMCChains, AxisKeys +import Turing: NUTS +import LinearAlgebra as ℒ +using Random, DelimitedFiles, AxisKeys + +using FlexiChains +include("test_helpers.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 @@ -67,7 +77,7 @@ Normal(0.3, 0.05, 0.01, 1.0), # calfa ] Turing.@model function SW07_loglikelihood_function(data, m, observables, fixed_parameters, filter) - all_params ~ Turing.arraydist(dists) + all_params ~ Turing.product_distribution(dists) z_ea, z_eb, z_eg, z_eqs, z_em, z_epinf, z_ew, crhoa, crhob, crhog, crhoqs, crhoms, crhopinf, crhow, cmap, cmaw, csadjcost, csigma, chabb, cprobw, csigl, cprobp, cindw, cindp, czcap, cfc, crpi, crr, cry, crdy, constepinf, constebeta, constelab, ctrend, cgy, calfa = all_params @@ -101,23 +111,40 @@ SW07_loglikelihood = SW07_loglikelihood_function(data, Smets_Wouters_2007_linear # modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, # Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), -# initial_params = modeSW2007.values) +# initial_params = modeSW2007.params) # modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, # Optim.NelderMead()) -# println("Mode variable values (linear): $(modeSW2007.values); Mode loglikelihood: $(modeSW2007.lp)") +# println("Mode variable values (linear): $(modeSW2007.params); Mode loglikelihood: $(modeSW2007.lp)") # LLH = Turing.logjoint(SW07_loglikelihood, (all_params = inits,)) n_samples = 1000 -samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, +samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, # initial_params = inits, progress = true) -println(samps) -println("Mean variable values (linear): $(mean(samps).nt.mean)") +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (linear): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") + +@testset "Mooncake vs FiniteDifferences gradient (SW07 linear)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(Smets_Wouters_2007_linear, data(observables), x, presample_periods = 4, initial_covariance = :diagonal, filter = :kalman), ADTypes.AutoMooncake(config = nothing), Smets_Wouters_2007_linear.parameter_values) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + 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, fin_grad[1], rtol = 1e-4) + break + end + end +end # estimate nonlinear model @@ -139,19 +166,36 @@ SW07_loglikelihood = SW07_loglikelihood_function(data, Smets_Wouters_2007, obser # modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, # Optim.LBFGS(linesearch = LineSearches.BackTracking(order = 3)), -# initial_params = modeSW2007.values) +# initial_params = modeSW2007.params) # modeSW2007 = Turing.maximum_a_posteriori(SW07_loglikelihood, # Optim.NelderMead(), -# initial_params = modeSW2007.values) +# initial_params = modeSW2007.params) -# println("Mode variable values (linear): $(modeSW2007.values); Mode loglikelihood: $(modeSW2007.lp)") +# println("Mode variable values (linear): $(modeSW2007.params); Mode loglikelihood: $(modeSW2007.lp)") n_samples = 1000 -samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoZygote()), n_samples, +samps = @time Turing.sample(SW07_loglikelihood, NUTS(adtype = AutoMooncake(; config=nothing)), n_samples, # initial_params = inits, progress = true) -println(samps) -println("Mean variable values (nonlinear): $(mean(samps).nt.mean)") \ No newline at end of file +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean variable values (nonlinear): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") + +@testset "Mooncake vs FiniteDifferences gradient (SW07 nonlinear)" begin + back_grad = DifferentiationInterface.gradient(x -> get_loglikelihood(Smets_Wouters_2007, data(observables), x, presample_periods = 4, initial_covariance = :diagonal, filter = :kalman), ADTypes.AutoMooncake(config = nothing), Smets_Wouters_2007.parameter_values) + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + 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, fin_grad[1], rtol = 1e-4) + break + end + end +end \ No newline at end of file diff --git a/test/test_sw07_estimation_nested_sampling.jl b/test/test_sw07_estimation_nested_sampling.jl new file mode 100644 index 000000000..ca12b5867 --- /dev/null +++ b/test/test_sw07_estimation_nested_sampling.jl @@ -0,0 +1,474 @@ +using Test +using MacroModelling +import Turing +using PythonCall +using DelimitedFiles, AxisKeys +using FlexiChains +using FlexiChains: Parameter, FlexiChain +using DataStructures: OrderedDict + +include("test_helpers.jl") + +# ────────────────────────────────────────────────────────────────────────────── +# Configuration switches +# ────────────────────────────────────────────────────────────────────────────── +USE_NESSAI = true +USE_DYNESTY = false +USE_ULTRANEST = false +USE_FLAT_PRIOR = false + +NESSAI_NLIVE = 1500 +NESSAI_FLOW_POOLSIZE = 128 +NESSAI_FLOW_DRAWSIZE = NESSAI_FLOW_POOLSIZE +NESSAI_UNINFORMED_POOLSIZE = NESSAI_FLOW_POOLSIZE +NESSAI_MAXIMUM_UNINFORMED = 2 * NESSAI_NLIVE +NESSAI_LOG_LEVEL = "INFO" +NESSAI_LOGGING_INTERVAL = 500 +NESSAI_IMPORTANCE_NESTED_SAMPLER = false +NESSAI_RESET_FLOW = false +# NESSAI_RETRAIN_ACCEPTANCE = true +# NESSAI_ACCEPTANCE_THRESHOLD = 0.1 + +# NSF flow configuration sized for SW07's 36-dimensional posterior. +# Wrapped in pydict() at the call site so nessai receives native Python dicts. +NESSAI_FLOW_CONFIG = Dict{String,Any}( + "ftype" => "nsf", + # "n_blocks" => 10, + # "n_neurons" => 64, + # "n_layers" => 4, + # "batch_norm_between_layers" => true, + # "use_random_permutations" => true, + # "use_residual_blocks" => true, + # "dropout_probability" => 0.01, + # "activation" => "relu", +) + +# Longer training schedule so the flow can learn the complex posterior shape +# const NESSAI_TRAINING_CONFIG = Dict{String,Any}( +# "max_epochs" => 1000, +# "patience" => 50, +# ) + +DYNESTY_NLIVE_INIT = NESSAI_NLIVE +DYNESTY_NLIVE_BATCH = max(500, DYNESTY_NLIVE_INIT ÷ 2) +DYNESTY_BOUND = "multi" +DYNESTY_SAMPLE = "rslice" +DYNESTY_DLOGZ_INIT = 0.1 +DYNESTY_BOOTSTRAP = 0 +DYNESTY_WEIGHT_PFRAC = 1.0 + +ULTRANEST_MIN_NUM_LIVE_POINTS = 400 + +# ────────────────────────────────────────────────────────────────────────────── +# Install nested-sampling Python packages into PythonCall's Python environment +# ────────────────────────────────────────────────────────────────────────────── +println("Installing nested-sampling Python packages...") +using CondaPkg +USE_NESSAI && CondaPkg.add_pip("nessai") +USE_DYNESTY && CondaPkg.add_pip("dynesty") +USE_ULTRANEST && CondaPkg.add_pip("ultranest") +CondaPkg.resolve() +println("Nested-sampling Python packages installed successfully") + +# ────────────────────────────────────────────────────────────────────────────── +# Load data (identical to test_sw07_estimation.jl) +# ────────────────────────────────────────────────────────────────────────────── +dat, header = readdlm("data/usmodel.csv", ',', header = true) +dat = Float64.(dat) +col_names = vec(Symbol.(strip.(header))) + +data = KeyedArray(dat', Variable = col_names, Time = axes(dat, 1)) + +observables_old = [:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs] +sample_idx = 47:230 # 1960Q1-2004Q4 +data = data(observables_old, sample_idx) + +observables = [:dy, :dc, :dinve, :labobs, :pinfobs, :dwobs, :robs] +data = rekey(data, :Variable => observables) + +# ────────────────────────────────────────────────────────────────────────────── +# Define priors (identical to test_sw07_estimation.jl) +# ────────────────────────────────────────────────────────────────────────────── +informative_dists = [ + InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_ea + InverseGamma(0.1, 2.0, 0.025,5.0, μσ = true), # z_eb + InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_eg + InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_eqs + InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_em + InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_epinf + InverseGamma(0.1, 2.0, 0.01, 3.0, μσ = true), # z_ew + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhoa + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhob + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhog + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhoqs + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhoms + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # crhopinf + Beta(0.5, 0.2, 0.001,0.9999, μσ = true), # crhow + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # cmap + Beta(0.5, 0.2, 0.01, 0.9999, μσ = true), # cmaw + Normal(4.0, 1.5, 2.0, 15.0), # csadjcost + Normal(1.50,0.375, 0.25, 3.0), # csigma + Beta(0.7, 0.1, 0.001, 0.99, μσ = true), # chabb + Beta(0.5, 0.1, 0.3, 0.95, μσ = true), # cprobw + Normal(2.0, 0.75, 0.25, 10.0), # csigl + Beta(0.5, 0.10, 0.5, 0.95, μσ = true), # cprobp + Beta(0.5, 0.15, 0.01, 0.99, μσ = true), # cindw + Beta(0.5, 0.15, 0.01, 0.99, μσ = true), # cindp + Beta(0.5, 0.15, 0.01, 0.99999, μσ = true), # czcap + Normal(1.25, 0.125, 1.0, 3.0), # cfc + Normal(1.5, 0.25, 1.0, 3.0), # crpi + Beta(0.75, 0.10, 0.5, 0.975, μσ = true), # crr + Normal(0.125, 0.05, 0.001, 0.5), # cry + Normal(0.125, 0.05, 0.001, 0.5), # crdy + Gamma(0.625, 0.1, 0.1, 2.0, μσ = true), # constepinf + Gamma(0.25, 0.1, 0.01, 2.0, μσ = true), # constebeta + Normal(0.0, 2.0, -10.0, 10.0), # constelab + Normal(0.4, 0.10, 0.1, 0.8), # ctrend + Normal(0.5, 0.25, 0.01, 2.0), # cgy + Normal(0.3, 0.05, 0.01, 1.0), # calfa +] + +dists = if USE_FLAT_PRIOR + [Turing.Uniform(minimum(d), maximum(d)) for d in informative_dists] +else + informative_dists +end + +# Parameter names in dists order +const param_names = [:z_ea, :z_eb, :z_eg, :z_eqs, :z_em, :z_epinf, :z_ew, + :crhoa, :crhob, :crhog, :crhoqs, :crhoms, :crhopinf, :crhow, + :cmap, :cmaw, + :csadjcost, :csigma, :chabb, :cprobw, :csigl, :cprobp, :cindw, :cindp, :czcap, + :cfc, :crpi, :crr, :cry, :crdy, + :constepinf, :constebeta, :constelab, :ctrend, :cgy, :calfa] + +# ────────────────────────────────────────────────────────────────────────────── +# Include linear model and set up fixed parameters +# ────────────────────────────────────────────────────────────────────────────── +include("../models/Smets_Wouters_2007_linear.jl") + +fixed_parameters = Smets_Wouters_2007_linear.parameter_values[indexin([:ctou, :clandaw, :cg, :curvp, :curvw], Smets_Wouters_2007_linear.constants.post_complete_parameters.parameters)] + +SS(Smets_Wouters_2007_linear, parameters = [:crhoms => 0.01, :crhopinf => 0.01, :crhow => 0.01, :cmap => 0.01, :cmaw => 0.01], derivatives = false) + +# ────────────────────────────────────────────────────────────────────────────── +# Reorder index: maps dists order → parameters_combined order (after fixed) +# parameters_combined = [ctou, clandaw, cg, curvp, curvw, +# calfa, csigma, cfc, cgy, csadjcost, chabb, cprobw, csigl, cprobp, +# cindw, cindp, czcap, crpi, crr, cry, crdy, +# crhoa, crhob, crhog, crhoqs, crhoms, crhopinf, crhow, cmap, cmaw, +# constelab, constepinf, constebeta, ctrend, +# z_ea, z_eb, z_eg, z_em, z_ew, z_eqs, z_epinf] +# ────────────────────────────────────────────────────────────────────────────── +const reorder_idx = [36, 18, 26, 35, 17, 19, 20, 21, 22, 23, 24, 25, + 27, 28, 29, 30, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 33, 31, 32, 34, 1, 2, 3, 5, 7, 4, 6] + +# ────────────────────────────────────────────────────────────────────────────── +# Shared Julia callback functions for nested samplers +# ────────────────────────────────────────────────────────────────────────────── +function sw07_log_prior_density(params::Vector{Float64}) + lp = 0.0 + for i in eachindex(dists) + lp += Turing.logpdf(dists[i], params[i]) + end + return lp +end + +function sw07_log_likelihood(params::Vector{Float64}) + parameters_combined = vcat(fixed_parameters, params[reorder_idx]) + llh = get_loglikelihood(Smets_Wouters_2007_linear, data(observables), parameters_combined, + presample_periods = 4, initial_covariance = :diagonal, + filter = :kalman, on_failure_loglikelihood = -1e10) + return llh +end + +function sw07_prior_transform(unit_params::Vector{Float64}) + transformed_params = Vector{Float64}(undef, length(unit_params)) + for i in eachindex(dists) + transformed_params[i] = Turing.quantile(dists[i], clamp(unit_params[i], eps(Float64), prevfloat(1.0))) + end + return transformed_params +end + +function posterior_matrix_from_named_samples(named_samples) + return reduce(hcat, [ + pyconvert(Vector{Float64}, named_samples[string(name)]) for name in param_names + ]) +end + +function summarize_posterior_matrix(label::String, posterior_matrix::Matrix{Float64}) + n_posterior = size(posterior_matrix, 1) + println("$label number of posterior samples: $n_posterior") + if n_posterior == 0 + println("$label returned no posterior samples") + return n_posterior, nothing + end + + println("$label posterior means:") + for (i, name) in pairs(param_names) + col = @view posterior_matrix[:, i] + println(" $name: $(sum(col) / length(col))") + end + + n_iters, _ = size(posterior_matrix) + symbol_names = Symbol.(collect(param_names)) + chain_data = OrderedDict{FlexiChains.ParameterOrExtra{Symbol}, Matrix{eltype(posterior_matrix)}}() + for (column, name) in pairs(symbol_names) + chain_data[Parameter(name)] = reshape(collect(@view posterior_matrix[:, column]), n_iters, 1) + end + posterior_chain = FlexiChain{Symbol}(n_iters, 1, chain_data) + posterior_summary = FlexiChains.summarystats(posterior_chain) + println("$label FlexiChains summary:") + show(stdout, MIME"text/plain"(), posterior_summary) + println() + return n_posterior, posterior_summary +end + +# Shared Python helpers +names_py = [string(n) for n in param_names] +bounds_py = Dict(string(n) => (Float64(minimum(d)), Float64(maximum(d))) + for (n, d) in zip(param_names, dists)) +np = (USE_DYNESTY || USE_ULTRANEST) ? pyimport("numpy") : nothing + +# ────────────────────────────────────────────────────────────────────────────── +# nessai FlowSampler +# ────────────────────────────────────────────────────────────────────────────── +if USE_NESSAI + + function nessai_log_prior(params_py) + return sw07_log_prior_density(pyconvert(Vector{Float64}, params_py)) + end + + function nessai_log_likelihood(params_py) + return sw07_log_likelihood(pyconvert(Vector{Float64}, params_py)) + end + + nessai_tmpdir = mktempdir() + write(joinpath(nessai_tmpdir, "sw07_nessai_model.py"), """ + import numpy as np + from nessai.model import Model + + class SW07NessaiModel(Model): + # SW07 DSGE model for nessai nested sampling. + + allow_vectorised = False + allow_vectorised_prior = False + likelihood_chunksize = 1 + + def __init__(self, param_names, param_bounds, jl_log_prior, jl_log_likelihood): + self.names = list(param_names) + self.bounds = dict(param_bounds) + self._jl_log_prior = jl_log_prior + self._jl_log_likelihood = jl_log_likelihood + + def _as_points(self, x): + x_array = np.asarray(x) + if x_array.shape == (): + return [x], True + return x_array, False + + def log_prior(self, x): + points, scalar_input = self._as_points(x) + log_p = np.zeros(len(points)) + for i, point in enumerate(points): + params = [float(point[n]) for n in self.names] + log_p[i] = float(self._jl_log_prior(params)) + if scalar_input: + return log_p[0] + return log_p + + def log_likelihood(self, x): + points, scalar_input = self._as_points(x) + log_l = np.zeros(len(points)) + for i, point in enumerate(points): + params = [float(point[n]) for n in self.names] + log_l[i] = float(self._jl_log_likelihood(params)) + if scalar_input: + return log_l[0] + return log_l + """) + + sys_mod = pyimport("sys") + sys_mod.path.insert(0, nessai_tmpdir) + sw07_nessai = pyimport("sw07_nessai_model") + + FlowSampler = pyimport("nessai.flowsampler").FlowSampler + RejectionProposal = pyimport("nessai.proposal").RejectionProposal + configure_nessai_logger = pyimport("nessai.utils.logging").configure_logger + + nessai_model = sw07_nessai.SW07NessaiModel(names_py, bounds_py, nessai_log_prior, nessai_log_likelihood) + nessai_output_dir = pwd() + + println("Running full nessai estimation on SW07 linear model...") + configure_nessai_logger( + output = nessai_output_dir, + label = "", + log_level = NESSAI_LOG_LEVEL, + stream = "stdout", + ) + nessai_fs = FlowSampler(nessai_model; + output = nessai_output_dir, + # importance_nested_sampler = NESSAI_IMPORTANCE_NESTED_SAMPLER, + nlive = NESSAI_NLIVE, + seed = 1234, + # pytorch_threads = 1, + resume = false, + disable_vectorisation = true, + logging_interval = NESSAI_LOGGING_INTERVAL, + log_on_iteration = true, + maximum_uninformed = NESSAI_MAXIMUM_UNINFORMED, + uninformed_proposal = RejectionProposal, + uninformed_proposal_kwargs = pydict(Dict("poolsize" => NESSAI_UNINFORMED_POOLSIZE)), + flow_config = pydict(NESSAI_FLOW_CONFIG), + # training_config = pydict(NESSAI_TRAINING_CONFIG), + # reset_flow = NESSAI_RESET_FLOW, + # retrain_acceptance = NESSAI_RETRAIN_ACCEPTANCE, + # acceptance_threshold = NESSAI_ACCEPTANCE_THRESHOLD, + poolsize = NESSAI_FLOW_POOLSIZE, + drawsize = NESSAI_FLOW_DRAWSIZE, + plot = false, + proposal_plots = false, + ) + nessai_fs.run(plot = false, save = false) + println("nessai estimation completed") + + nessai_log_evidence = pyconvert(Float64, nessai_fs.logZ) + nessai_posterior_samples = nessai_fs.posterior_samples + nessai_n_posterior, nessai_posterior_summary = summarize_posterior_matrix( + "nessai", + posterior_matrix_from_named_samples(nessai_posterior_samples), + ) + + println("nessai log evidence: $nessai_log_evidence") + + @testset "nessai SW07 linear estimation" begin + @test isfinite(nessai_log_evidence) + @test nessai_n_posterior > 0 + @test !isnothing(nessai_posterior_summary) + @test !isnothing(nessai_fs) + end + +end # USE_NESSAI + +# ────────────────────────────────────────────────────────────────────────────── +# dynesty DynamicNestedSampler +# ────────────────────────────────────────────────────────────────────────────── +if USE_DYNESTY + + dynesty = pyimport("dynesty") + + function dynesty_log_likelihood(params_py) + return sw07_log_likelihood(pyconvert(Vector{Float64}, params_py)) + end + + function dynesty_prior_transform(unit_params_py) + transformed_params = sw07_prior_transform(pyconvert(Vector{Float64}, unit_params_py)) + return np.asarray(pylist(transformed_params), dtype = np.float64) + end + + println("Running dynesty dynamic nested sampling on SW07 linear model...") + dynesty_sampler = dynesty.DynamicNestedSampler( + dynesty_log_likelihood, + dynesty_prior_transform, + length(param_names); + bound = DYNESTY_BOUND, + sample = DYNESTY_SAMPLE, + slices = length(param_names) + 3, + bootstrap = DYNESTY_BOOTSTRAP, + queue_size = 1, + ) + dynesty_sampler.run_nested( + nlive_init = DYNESTY_NLIVE_INIT, + nlive_batch = DYNESTY_NLIVE_BATCH, + dlogz_init = DYNESTY_DLOGZ_INIT, + wt_kwargs = Dict("pfrac" => DYNESTY_WEIGHT_PFRAC), + stop_kwargs = Dict("pfrac" => DYNESTY_WEIGHT_PFRAC), + print_progress = true, + save_bounds = false, + ) + println("dynesty dynamic nested sampling completed") + + dynesty_results = dynesty_sampler.results + println("dynesty summary:") + dynesty_results.summary() + dynesty_log_evidence = pyconvert(Vector{Float64}, dynesty_results.logz)[end] + dynesty_posterior_matrix = pyconvert(Matrix{Float64}, dynesty_results.samples_equal()) + dynesty_n_posterior, dynesty_posterior_summary = summarize_posterior_matrix( + "dynesty dynamic", + dynesty_posterior_matrix, + ) + + println("dynesty log evidence: $dynesty_log_evidence") + + @testset "dynesty dynamic SW07 linear estimation" begin + @test isfinite(dynesty_log_evidence) + @test dynesty_n_posterior > 0 + @test !isnothing(dynesty_posterior_summary) + @test !isnothing(dynesty_sampler) + @test !isnothing(dynesty_results) + end + +end # USE_DYNESTY + +# ────────────────────────────────────────────────────────────────────────────── +# UltraNest ReactiveNestedSampler +# ────────────────────────────────────────────────────────────────────────────── +if USE_ULTRANEST + + ultranest = pyimport("ultranest") + ultranest_stepsampler = pyimport("ultranest.stepsampler") + ReactiveNestedSampler = ultranest.ReactiveNestedSampler + + function ultranest_log_likelihood(params_py) + return sw07_log_likelihood(pyconvert(Vector{Float64}, params_py)) + end + + function ultranest_prior_transform(unit_params_py) + transformed_params = sw07_prior_transform(pyconvert(Vector{Float64}, unit_params_py)) + return np.asarray(pylist(transformed_params), dtype = np.float64) + end + + ultranest_log_dir = mktempdir() + + println("Running UltraNest nested sampling on SW07 linear model...") + ultranest_sampler = ReactiveNestedSampler( + pylist(names_py), + ultranest_log_likelihood, + ultranest_prior_transform; + log_dir = ultranest_log_dir, + resume = "overwrite", + ) + + nsteps = length(param_names) + ultranest_sampler.stepsampler = ultranest_stepsampler.SliceSampler(; + nsteps = nsteps, + generate_direction = ultranest_stepsampler.generate_mixture_random_direction, + ) + + ultranest_result = ultranest_sampler.run(; + min_num_live_points = ULTRANEST_MIN_NUM_LIVE_POINTS, + show_status = true, + ) + ultranest_sampler.print_results() + println("UltraNest nested sampling completed") + + ultranest_log_evidence = pyconvert(Float64, ultranest_result["logz"]) + ultranest_posterior_matrix = pyconvert(Matrix{Float64}, ultranest_result["samples"]) + @assert size(ultranest_posterior_matrix, 2) == length(param_names) "UltraNest samples have $(size(ultranest_posterior_matrix, 2)) columns but expected $(length(param_names))" + + ultranest_n_posterior, ultranest_posterior_summary = summarize_posterior_matrix( + "UltraNest", + ultranest_posterior_matrix, + ) + + println("UltraNest log evidence: $ultranest_log_evidence") + + @testset "UltraNest SW07 linear estimation" begin + @test isfinite(ultranest_log_evidence) + @test ultranest_n_posterior > 0 + @test !isnothing(ultranest_posterior_summary) + @test !isnothing(ultranest_result) + end + +end # USE_ULTRANEST diff --git a/test/test_system_prior_estimation.jl b/test/test_system_prior_estimation.jl new file mode 100644 index 000000000..f8948fb66 --- /dev/null +++ b/test/test_system_prior_estimation.jl @@ -0,0 +1,319 @@ +using Test +using MacroModelling +import Turing +import Zygote +import Mooncake +import ForwardDiff +import ADTypes +import ADTypes: AutoZygote, AutoForwardDiff, AutoMooncake +import DifferentiationInterface +import FiniteDifferences +import Turing: NUTS, sample +import LinearAlgebra as ℒ +using Random, AxisKeys + +using FlexiChains +include("test_helpers.jl") + +include("../models/Gali_2015_chapter_3_nonlinear.jl") + +# Gali model parameter order (from @parameters block): +# σ=1, φ=5, ϕᵖⁱ=1.5, ϕʸ=0.125, θ=0.75, ρ_ν=0.5, ρ_z=0.5, ρ_a=0.9, +# β=0.99, η=3.77, α=0.25, ϵ=9, τ=0, std_a=0.01, std_z=0.05, std_nu=0.0025 + +# Simulate data from the model at true parameter values +Random.seed!(42) +simulated_data = simulate(Gali_2015_chapter_3_nonlinear, algorithm = :pruned_second_order) + +# Observables: log_y, pi_ann, i_ann (3 observables matching 3 shocks) +observables = [:log_y, :pi_ann, :i_ann] +data = simulated_data(observables, :, :simulate) + +# Compute target values at true parameters for NSSS, moments, and IRF targeting +true_params = Gali_2015_chapter_3_nonlinear.parameter_values + +nsss_vars = [:Y, :Pi] +moment_vars = [:Y, :Pi] + +target_stats = get_statistics(Gali_2015_chapter_3_nonlinear, true_params, + non_stochastic_steady_state = nsss_vars, + mean = moment_vars, + standard_deviation = moment_vars, + algorithm = :pruned_second_order) + +target_nsss = target_stats[:non_stochastic_steady_state] +target_mean = target_stats[:mean] +target_std = target_stats[:standard_deviation] + +# IRF targets: first period response of log_y to first shock (eps_a) +target_irf_full = get_irf(Gali_2015_chapter_3_nonlinear, true_params, + algorithm = :pruned_second_order, periods = 5) +irf_var_idx = sort(MacroModelling.parse_variables_input_to_index(:log_y, Gali_2015_chapter_3_nonlinear)) +target_irf = target_irf_full[irf_var_idx, 1, 1] + +# Prior distributions for estimated parameters (6 out of 16) +# Estimated: ρ_ν (idx 6), ρ_z (idx 7), ρ_a (idx 8), std_a (idx 14), std_z (idx 15), std_nu (idx 16) +estimated_param_indices = [6, 7, 8, 14, 15, 16] + +# Build index mapping to reconstruct full parameter vector (Zygote-compatible, no mutation) +# Fixed params: indices 1-5 (σ,φ,ϕᵖⁱ,ϕʸ,θ) and 9-13 (β,η,α,ϵ,τ) +# Estimated params: indices 6-8 (ρ_ν,ρ_z,ρ_a) and 14-16 (std_a,std_z,std_nu) +function build_full_params(estimated_vals::AbstractVector{T}) where T + return vcat( + T.(true_params[1:5]), # σ, φ, ϕᵖⁱ, ϕʸ, θ + estimated_vals[1:3], # ρ_ν, ρ_z, ρ_a + T.(true_params[9:13]), # β, η, α, ϵ, τ + estimated_vals[4:6], # std_a, std_z, std_nu + ) +end + +dists = [ + Beta(0.5, 0.15, μσ = true), # ρ_ν + Beta(0.5, 0.15, μσ = true), # ρ_z + Beta(0.9, 0.05, μσ = true), # ρ_a + InverseGamma(0.01, Inf, μσ = true), # std_a + InverseGamma(0.05, Inf, μσ = true), # std_z + InverseGamma(0.0025, Inf, μσ = true), # std_nu +] + +Turing.@model function Gali_estimation(data, m, algorithm, on_failure_llh, + target_nsss, target_mean, target_std, target_irf, + nsss_vars, moment_vars, irf_var_idx; + verbose = false) + estimated_params ~ Turing.product_distribution(dists) + all_params = build_full_params(estimated_params) + + # 1. Log-likelihood from data + llh = get_loglikelihood(m, data, all_params, + algorithm = algorithm, + on_failure_loglikelihood = on_failure_llh) + maybe_print_loglikelihood(verbose, llh, dists, estimated_params) + Turing.@addlogprob! llh + + # 2. NSSS targeting via get_statistics (positional params) + stats_nsss = get_statistics(m, all_params, + non_stochastic_steady_state = nsss_vars, + algorithm = algorithm) + nsss_vals = stats_nsss[:non_stochastic_steady_state] + Turing.@addlogprob! sum(Turing.logpdf.(Turing.Normal.(target_nsss, 0.1), nsss_vals)) + + # 3. Moment targeting via get_statistics (positional params) + stats_moments = get_statistics(m, all_params, + mean = moment_vars, + standard_deviation = moment_vars, + algorithm = algorithm) + mean_vals = stats_moments[:mean] + std_vals = stats_moments[:standard_deviation] + Turing.@addlogprob! sum(Turing.logpdf.(Turing.Normal.(target_mean, 0.1), mean_vals)) + Turing.@addlogprob! sum(Turing.logpdf.(Turing.Normal.(target_std, 0.05), std_vals)) + + # 4. IRF targeting via get_irf (positional params with algorithm) + irf_vals = get_irf(m, all_params, algorithm = algorithm, periods = 5) + Turing.@addlogprob! sum(Turing.logpdf.(Turing.Normal.(target_irf, 0.1), irf_vals[irf_var_idx, 1, 1])) +end + + +# Instantiate the Turing model +gali_model = Gali_estimation(data, Gali_2015_chapter_3_nonlinear, :pruned_second_order, -Inf, + target_nsss, target_mean, target_std, target_irf, + nsss_vars, moment_vars, irf_var_idx) + +Random.seed!(123) + +n_samples = 1000 + +samps = @time sample(gali_model, + NUTS(adtype = AutoForwardDiff()), + n_samples, + progress = true, + initial_params = Turing.InitFromParams((estimated_params = true_params[estimated_param_indices],))) + +posterior_summary = FlexiChains.summarystats(samps) +show(stdout, MIME"text/plain"(), posterior_summary) +println() +println("Mean estimated values (ForwardDiff): $(collect(values(FlexiChains.mean(samps); parameters_only = true)))") + +sample_means = collect(values(FlexiChains.mean(samps); parameters_only = true)) + +@testset "Gali pruned 2nd order estimation results" begin + @test length(sample_means) == 6 + @test all(isfinite, sample_means) + # Means should be in the right ballpark of true values + @test isapprox(sample_means, true_params[estimated_param_indices], rtol = 0.5) +end + +# --------------------------------------------------------------------------- +# Mooncake NUTS sampling +# --------------------------------------------------------------------------- +Random.seed!(123) + +samps_mc = @time sample(gali_model, + NUTS(adtype = AutoMooncake(; config=nothing)), + n_samples, + progress = true, + initial_params = Turing.InitFromParams((estimated_params = true_params[estimated_param_indices],))) + +posterior_summary_mc = FlexiChains.summarystats(samps_mc) +show(stdout, MIME"text/plain"(), posterior_summary_mc) +println() + +sample_means_mc = collect(values(FlexiChains.mean(samps_mc); parameters_only = true)) +println("Mean estimated values (Mooncake): $(sample_means_mc)") + +@testset "Gali pruned 2nd order estimation results (Mooncake)" begin + @test length(sample_means_mc) == 6 + @test all(isfinite, sample_means_mc) + @test isapprox(sample_means_mc, true_params[estimated_param_indices], rtol = 0.5) +end + +@testset "Zygote vs FiniteDifferences gradient (Gali pruned 2nd order)" begin + # Test gradient of combined objective at true parameter values + function combined_objective(x) + all_p = build_full_params(x) + m = Gali_2015_chapter_3_nonlinear + alg = :pruned_second_order + + llh = get_loglikelihood(m, data, all_p, algorithm = alg, on_failure_loglikelihood = -Inf) + + stats_n = get_statistics(m, all_p, non_stochastic_steady_state = nsss_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_nsss, 0.1), stats_n[:non_stochastic_steady_state])) + + stats_m = get_statistics(m, all_p, mean = moment_vars, standard_deviation = moment_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_mean, 0.1), stats_m[:mean])) + llh += sum(Turing.logpdf.(Turing.Normal.(target_std, 0.05), stats_m[:standard_deviation])) + + irf_v = get_irf(m, all_p, algorithm = alg, periods = 5) + llh += sum(Turing.logpdf.(Turing.Normal.(target_irf, 0.1), irf_v[irf_var_idx, 1, 1])) + + return llh + end + + test_point = true_params[estimated_param_indices] + + back_grad = Zygote.gradient(combined_objective, test_point)[1] + @test !isnothing(back_grad) + @test all(isfinite, back_grad) + + for i in 1:100 + local fin_grad = FiniteDifferences.grad(FiniteDifferences.central_fdm(4, 1), combined_objective, test_point) + if isfinite(ℒ.norm(fin_grad)) + println("Finite differences converged after $i iterations") + @test isapprox(back_grad, fin_grad[1], rtol = 1e-4) + break + end + end +end + +# Test Zygote gradient of the full log posterior (likelihood + priors) +# NUTS with AutoZygote fails due to DynamicPPL's bijector using setindex!, +# so we test Zygote differentiation of the combined objective directly. +@testset "Zygote log posterior gradient (Gali pruned 2nd order)" begin + function turing_logjoint(x) + all_p = build_full_params(x) + m = Gali_2015_chapter_3_nonlinear + alg = :pruned_second_order + + llh = get_loglikelihood(m, data, all_p, algorithm = alg, on_failure_loglikelihood = -Inf) + + stats_n = get_statistics(m, all_p, non_stochastic_steady_state = nsss_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_nsss, 0.1), stats_n[:non_stochastic_steady_state])) + + stats_m = get_statistics(m, all_p, mean = moment_vars, standard_deviation = moment_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_mean, 0.1), stats_m[:mean])) + llh += sum(Turing.logpdf.(Turing.Normal.(target_std, 0.05), stats_m[:standard_deviation])) + + irf_v = get_irf(m, all_p, algorithm = alg, periods = 5) + llh += sum(Turing.logpdf.(Turing.Normal.(target_irf, 0.1), irf_v[irf_var_idx, 1, 1])) + + # Add prior log density + for (i, d) in enumerate(dists) + llh += Turing.logpdf(d, x[i]) + end + + return llh + end + + test_pt = true_params[estimated_param_indices] + + zy_grad = Zygote.gradient(turing_logjoint, test_pt)[1] + fd_grad = ForwardDiff.gradient(turing_logjoint, test_pt) + + @test all(isfinite, zy_grad) + @test all(isfinite, fd_grad) + + rel_err = maximum(abs.(zy_grad .- fd_grad) ./ max.(abs.(fd_grad), 1e-10)) + println("Zygote vs ForwardDiff gradient rel err on log posterior: $rel_err") + @test rel_err < 1e-6 +end + +@testset "Mooncake vs ForwardDiff gradient (Gali pruned 2nd order)" begin + function combined_objective_mc(x) + all_p = build_full_params(x) + m = Gali_2015_chapter_3_nonlinear + alg = :pruned_second_order + + llh = get_loglikelihood(m, data, all_p, algorithm = alg, on_failure_loglikelihood = -Inf) + + stats_n = get_statistics(m, all_p, non_stochastic_steady_state = nsss_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_nsss, 0.1), stats_n[:non_stochastic_steady_state])) + + stats_m = get_statistics(m, all_p, mean = moment_vars, standard_deviation = moment_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_mean, 0.1), stats_m[:mean])) + llh += sum(Turing.logpdf.(Turing.Normal.(target_std, 0.05), stats_m[:standard_deviation])) + + irf_v = get_irf(m, all_p, algorithm = alg, periods = 5) + llh += sum(Turing.logpdf.(Turing.Normal.(target_irf, 0.1), irf_v[irf_var_idx, 1, 1])) + + return llh + end + + test_point = true_params[estimated_param_indices] + + mc_grad = DifferentiationInterface.gradient(combined_objective_mc, AutoMooncake(config = nothing), test_point) + fd_grad = ForwardDiff.gradient(combined_objective_mc, test_point) + + @test all(isfinite, mc_grad) + @test all(isfinite, fd_grad) + + rel_err = maximum(abs.(mc_grad .- fd_grad) ./ max.(abs.(fd_grad), 1e-10)) + println("Mooncake vs ForwardDiff gradient rel err: $rel_err") + @test rel_err < 1e-4 +end + +@testset "Mooncake log posterior gradient (Gali pruned 2nd order)" begin + function turing_logjoint_mc(x) + all_p = build_full_params(x) + m = Gali_2015_chapter_3_nonlinear + alg = :pruned_second_order + + llh = get_loglikelihood(m, data, all_p, algorithm = alg, on_failure_loglikelihood = -Inf) + + stats_n = get_statistics(m, all_p, non_stochastic_steady_state = nsss_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_nsss, 0.1), stats_n[:non_stochastic_steady_state])) + + stats_m = get_statistics(m, all_p, mean = moment_vars, standard_deviation = moment_vars, algorithm = alg) + llh += sum(Turing.logpdf.(Turing.Normal.(target_mean, 0.1), stats_m[:mean])) + llh += sum(Turing.logpdf.(Turing.Normal.(target_std, 0.05), stats_m[:standard_deviation])) + + irf_v = get_irf(m, all_p, algorithm = alg, periods = 5) + llh += sum(Turing.logpdf.(Turing.Normal.(target_irf, 0.1), irf_v[irf_var_idx, 1, 1])) + + for (i, d) in enumerate(dists) + llh += Turing.logpdf(d, x[i]) + end + + return llh + end + + test_pt = true_params[estimated_param_indices] + + mc_grad = DifferentiationInterface.gradient(turing_logjoint_mc, AutoMooncake(config = nothing), test_pt) + fd_grad = ForwardDiff.gradient(turing_logjoint_mc, test_pt) + + @test all(isfinite, mc_grad) + @test all(isfinite, fd_grad) + + rel_err = maximum(abs.(mc_grad .- fd_grad) ./ max.(abs.(fd_grad), 1e-10)) + println("Mooncake vs ForwardDiff gradient rel err on log posterior: $rel_err") + @test rel_err < 1e-4 +end diff --git a/test/test_update_equations.jl b/test/test_update_equations.jl new file mode 100644 index 000000000..7ad7086d8 --- /dev/null +++ b/test/test_update_equations.jl @@ -0,0 +1,743 @@ +using MacroModelling +using Test + +# Helper to get a fresh copy of the SW07 model for testing +# We need to reload to get a clean state without revision history +function load_sw07() + # Re-include the model file to get a fresh model instance + Base.invokelatest(include, "../models/Smets_Wouters_2007.jl") + return Base.invokelatest(() -> Smets_Wouters_2007) +end + +function load_RBC_baseline() + # Re-include the model file to get a fresh model instance + Base.invokelatest(include, "../models/RBC_baseline.jl") + return Base.invokelatest(() -> RBC_baseline) +end + +function load_FS2000() + # Re-include the model file to get a fresh model instance + Base.invokelatest(include, "../models/FS2000.jl") + return Base.invokelatest(() -> FS2000) +end + +@testset verbose = true "SW07 update_equations! functionality" begin + + @testset "Update equation by index - modify Taylor rule (remove output growth term)" begin + model = load_sw07() + + # Get original state + original_eqs = get_equations(model) + n_eqs_original = length(original_eqs) + ss_before = get_steady_state(model, derivatives = false) + + # Find the Taylor rule equation (contains r[0] and crpi) + taylor_idx = findfirst(eq -> occursin("r[0]", eq) && occursin("crpi", eq) && occursin("crr", eq), string.(original_eqs)) + @test taylor_idx !== nothing + + # Modify Taylor rule: remove output growth term (crdy term) but keep all parameters + # Original: r[0] = r[ss]^(1-crr) * r[-1]^crr * (pinf[0]/cpie)^((1-crr)*crpi) * (y[0]/yflex[0])^((1-crr)*cry) * (y[0]/yflex[0]/(y[-1]/yflex[-1]))^crdy * ms[0] + # New: remove the output growth response term (crdy term) + new_taylor = :(r[0] = r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / cpie) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * ms[0]) + + update_equations!(model, taylor_idx, new_taylor) + + # Check revision history was recorded + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :update_equation + @test history[1].equation_index == taylor_idx + @test history[1].old_equation !== nothing + @test history[1].new_equation == new_taylor + + # Model should still solve + ss_after = get_steady_state(model, derivatives = false) + @test !any(isnan, ss_after) + + # Number of equations should remain the same + @test length(get_equations(model)) == n_eqs_original + + # Steady state should be unchanged (Taylor rule modification doesn't affect SS) + @test isapprox(collect(ss_before), collect(ss_after), rtol = 1e-10) + + # IRF should work + irf_after = get_irf(model) + @test size(irf_after, 1) > 0 + + model = nothing + end + + @testset "Update equation by matching - change shock process structure" begin + model = load_sw07() + + # Original technology shock: a[0] = 1 - crhoa + crhoa * a[-1] + z_ea / 100 * ea[x] + old_shock = :(a[0] = 1 - crhoa + crhoa * a[-1] + z_ea / 100 * ea[x]) + # Make technology shock more persistent by changing the mean reversion structure + # Use different coefficients but keep crhoa parameter + new_shock = :(a[0] = 1 - crhoa + crhoa * a[-1] + z_ea / 100 * ea[x-1]) + + update_equations!(model, old_shock, new_shock) + + # Check revision history + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :update_equation + + # Model should still solve + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Update equation using string format - modify marginal utility equation" begin + model = load_sw07() + + # Find the marginal utility equation with consumption habit + eqs = get_equations(model) + xi_idx = findfirst(eq -> occursin("xi[0]", eq) && occursin("chabb", eq) && occursin("csigma", eq), string.(eqs)) + @test xi_idx !== nothing + + # Modify to reduce habit formation strength (multiply chabb by 0.5) + new_xi_eq = "xi[0] = exp(((csigma - 1) / (1 + csigl)) * ((lab[0] * (curvW + wdot[0])) / (1 + curvW)) ^ (1 + csigl)) * (c[0] - (c[ss] * chabb) / cgamma) ^ -csigma" + + update_equations!(model, xi_idx, new_xi_eq) + + history = get_revision_history(model) + @test length(history) == 1 + + # Model should still solve + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Update equation by matching old equation string" begin + model = load_sw07() + + # Match monetary policy shock equation using string + old_eq = "ms[0] = 1 - crhoms + crhoms * ms[-1] + z_em / 100 * em[x]" + # Modify structure while keeping parameters + new_eq = "ms[0] = 1 - crhoms + crhoms * ms[-1] + z_em / 100 * em[x-2]" + + update_equations!(model, old_eq, new_eq) + + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :update_equation + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Multiple equation updates with tuple syntax" begin + model = load_sw07() + + ss_before = get_steady_state(model, derivatives = false) + eqs = get_equations(model) + + # Find shock equation indices + a_idx = findfirst(eq -> occursin("a[0]", eq) && occursin("crhoa", eq) && occursin("ea[x]", eq), string.(eqs)) + b_idx = findfirst(eq -> occursin("b[0]", eq) && occursin("crhob", eq) && occursin("eb[x]", eq), string.(eqs)) + @test a_idx !== nothing + @test b_idx !== nothing + + # Update both shock processes at once - double the shock standard deviations + update_equations!(model, [ + (a_idx, :(a[0] = 1 - crhoa + crhoa * a[-2] + z_ea / 100 * ea[x])), + (b_idx, :(b[0] = 1 - crhob + crhob * b[-2] + z_eb / 100 * SCALE1_eb * eb[x])) + ]) + + # Check revision history has 2 entries + history = get_revision_history(model) + @test length(history) == 2 + @test all(h.action == :update_equation for h in history) + + # Model should solve + ss_after = get_steady_state(model, derivatives = false) + @test !any(isnan, ss_after) + + # Steady state unchanged (shock size doesn't affect SS) + @test isapprox(collect(ss_before), collect(ss_after), rtol = 1e-10) + + model = nothing + end + + @testset "Update and revert - round trip with same parameters" begin + model = load_sw07() + + ss_original = get_steady_state(model, derivatives = false) + irf_original = get_irf(model) + + # Original Taylor rule - keep all parameters in both versions + old_taylor = :(r[0] = r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / cpie) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * (y[0] / yflex[0] / (y[-1] / yflex[-1])) ^ crdy * ms[0]) + # Modified Taylor rule - remove output growth term but keep crdy in a benign way + new_taylor = :(r[0] = r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / cpie) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * ms[0]) + + # Update + update_equations!(model, old_taylor, new_taylor) + @test length(get_revision_history(model)) == 1 + + # Revert + update_equations!(model, new_taylor, old_taylor) + @test length(get_revision_history(model)) == 2 + + # Should be back to original + ss_final = get_steady_state(model, derivatives = false, parameters = :crdy => 0.2347) + irf_final = get_irf(model) + + @test isapprox(collect(ss_original), collect(ss_final), rtol = 1e-10) + @test isapprox(collect(irf_original), collect(irf_final), rtol = 1e-10) + + model = nothing + end + + @testset "Error cases" begin + model = load_sw07() + + n_eqs = length(get_equations(model)) + + # Error: index out of bounds + @test_throws AssertionError update_equations!(model, n_eqs + 1, :(x[0] = 1)) + @test_throws AssertionError update_equations!(model, 0, :(x[0] = 1)) + + # Error: equation not found + @test_throws AssertionError update_equations!( + model, + :(nonexistent_variable[0] = some_other[0] + 1), + :(x[0] = 1) + ) + + model = nothing + end +end + + +@testset verbose = true "SW07 add_equation! functionality" begin + + @testset "Add auxiliary output gap variable" begin + model = load_sw07() + + n_eqs_before = length(get_equations(model)) + n_vars_before = length(get_variables(model)) + + # Add an inflation gap measure (deviation from target) + add_equation!(model, :(inflation_gap[0] = pinf[0] - cpie)) + + @test length(get_equations(model)) == n_eqs_before + 1 + @test length(get_variables(model)) == n_vars_before + 1 + @test "inflation_gap" in get_variables(model) + + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :add_equation + @test history[1].equation_index == n_eqs_before + 1 + @test history[1].old_equation === nothing + + ss = get_steady_state(model, derivatives = false) + @test !isnan(ss(:inflation_gap)) + + model = nothing + end + + @testset "Add equation using string format" begin + model = load_sw07() + + n_eqs_before = length(get_equations(model)) + + # Add real interest rate definition + add_equation!(model, "real_rate[0] = r[0] / pinf[1]") + + @test length(get_equations(model)) == n_eqs_before + 1 + @test "real_rate" in get_variables(model) + + model = nothing + end + + @testset "Add multiple equations at once" begin + model = load_sw07() + + n_eqs_before = length(get_equations(model)) + + # Add several auxiliary variables + add_equation!(model, [ + :(nominal_gdp[0] = y[0] * pinf[0]), + :(investment_ratio[0] = inve[0] / y[0]), + :(consumption_ratio[0] = c[0] / y[0]) + ]) + + @test length(get_equations(model)) == n_eqs_before + 3 + @test "nominal_gdp" in get_variables(model) + @test "investment_ratio" in get_variables(model) + @test "consumption_ratio" in get_variables(model) + + history = get_revision_history(model) + @test length(history) == 3 + @test all(h.action == :add_equation for h in history) + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss([:nominal_gdp, :investment_ratio, :consumption_ratio])) + + model = nothing + end + + @testset "Add equation with lagged variables" begin + model = load_sw07() + + # Add output growth measure + add_equation!(model, :(output_growth[0] = y[0] / y[-1] - 1)) + + @test "output_growth" in get_variables(model) + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Add equation with forward-looking variables" begin + model = load_sw07() + + # Add expected inflation measure + add_equation!(model, :(expected_inflation[0] = pinf[1])) + + @test "expected_inflation" in get_variables(model) + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end +end + + +@testset verbose = true "SW07 remove_equation! functionality" begin + + @testset "Remove observable equation by index" begin + model = load_sw07() + + eqs = get_equations(model) + n_eqs_before = length(eqs) + n_vars_before = length(get_variables(model)) + + # Find labobs equation: labobs[0] = constelab + 100 * (lab[0] / lab[ss] - 1) + labobs_idx = findfirst(eq -> occursin("labobs[0]", eq) && occursin("constelab", eq), string.(eqs)) + @test labobs_idx !== nothing + + remove_equation!(model, labobs_idx) + + @test length(get_equations(model)) == n_eqs_before - 1 + @test length(get_variables(model)) == n_vars_before - 1 + @test !("labobs" in get_variables(model)) + + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :remove_equation + @test history[1].new_equation === nothing + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Remove equation by matching" begin + model = load_sw07() + + n_eqs_before = length(get_equations(model)) + + # Remove wage growth observable + remove_equation!(model, :(dwobs[0] = ctrend + 100 * (w[0] / w[-1] - 1))) + + @test length(get_equations(model)) == n_eqs_before - 1 + @test !("dwobs" in get_variables(model)) + + model = nothing + end + + @testset "Remove multiple equations" begin + model = load_sw07() + + eqs = get_equations(model) + n_eqs_before = length(eqs) + + # Find observable equations to remove + labobs_idx = findfirst(eq -> occursin("labobs[0]", eq), string.(eqs)) + dwobs_idx = findfirst(eq -> occursin("dwobs[0]", eq), string.(eqs)) + @test labobs_idx !== nothing + @test dwobs_idx !== nothing + + # Remove in descending order + remove_equation!(model, sort([labobs_idx, dwobs_idx], rev=true)) + + @test length(get_equations(model)) == n_eqs_before - 2 + @test !("labobs" in get_variables(model)) + @test !("dwobs" in get_variables(model)) + + history = get_revision_history(model) + @test length(history) == 2 + + model = nothing + end + + @testset "Add then remove - round trip" begin + model = load_sw07() + + ss_before = get_steady_state(model, derivatives = false) + n_eqs_before = length(get_equations(model)) + + # Add auxiliary variable + add_equation!(model, :(temp_var[0] = y[0] + c[0])) + @test length(get_equations(model)) == n_eqs_before + 1 + + # Remove it + remove_equation!(model, :(temp_var[0] = y[0] + c[0])) + @test length(get_equations(model)) == n_eqs_before + + # Steady state should match original + ss_after = get_steady_state(model, derivatives = false) + @test isapprox(collect(ss_before), collect(ss_after), rtol = 1e-10) + + model = nothing + end + + @testset "Error cases" begin + model = load_sw07() + + n_eqs = length(get_equations(model)) + + # Error: index out of bounds + @test_throws AssertionError remove_equation!(model, n_eqs + 1) + @test_throws AssertionError remove_equation!(model, 0) + + # Error: equation not found + @test_throws AssertionError remove_equation!( + model, + :(nonexistent[0] = 1) + ) + + model = nothing + end +end + + +@testset verbose = true "SW07 update_calibration_equations! functionality" begin + + @testset "Update calibration equation - change target value" begin + model = load_RBC_baseline() + + # SW07 has calibration: mcflex = mc[ss] | mcflex + calib_eqs = get_calibration_equations(model) + @test length(calib_eqs) >= 1 + + calib_params = get_calibrated_parameters(model) + @test "ψ" in calib_params + + # Modify the calibration target slightly (multiply by 1.01) + update_calibration_equations!(model, 2, :(l[ss] = 0.5 | ψ)) + + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :update_calibration_equation + + # mcflex should still be calibrated + @test "ψ" in get_calibrated_parameters(model) + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Update calibration equation using tuple syntax" begin + model = load_RBC_baseline() + + # Update both calibration equations (mcflex and cpie) + update_calibration_equations!(model, [ + (1, :(l[ss] = 2/3 | ψ)), + (2, :(ḡ = 1/4 * y[ss]| ḡ)) + ]) + # TODO: accept also: par | calib_eq + + history = get_revision_history(model) + @test length(history) == 2 + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Error - calibrate non-existent parameter" begin + model = load_sw07() + + @test_throws ErrorException update_calibration_equations!( + model, + 1, + :(mcflex = mc[ss] | nonexistent_param), + silent = true + ) + + model = nothing + end +end + + +@testset verbose = true "SW07 add_calibration_equation! functionality" begin + + @testset "Add calibration for fixed parameter" begin + model = load_FS2000() + + # ctou = 0.025 is a fixed depreciation rate parameter + calib_before = get_calibrated_parameters(model) + @test !("del" in calib_before) + + n_calib_before = length(get_calibration_equations(model)) + + # Add calibration: capital-output ratio determines depreciation + add_calibration_equation!(model, :(k[ss] / y[ss] = 8.0 | del)) + + @test length(get_calibration_equations(model)) == n_calib_before + 1 + @test "del" in get_calibrated_parameters(model) + + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :add_calibration_equation + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Add calibration using vector syntax" begin + model = load_FS2000() + + n_calib_before = length(get_calibration_equations(model)) + + # Add calibration for calfa (capital share) + add_calibration_equation!(model, [ + :(k[ss] / y[ss] = 4 | del) + ]) + + @test length(get_calibration_equations(model)) == n_calib_before + 1 + @test "del" in get_calibrated_parameters(model) + + model = nothing + end + + @testset "Error - add calibration for already calibrated parameter" begin + model = load_RBC_baseline() + + # del is already calibrated + @test "ψ" in get_calibrated_parameters(model) + + @test_throws ErrorException add_calibration_equation!( + model, + :(y[ss] = 1.5 | ψ) + ) + + model = nothing + end + + @testset "Error - add calibration without | syntax" begin + model = load_RBC_baseline() + + @test_throws ErrorException add_calibration_equation!( + model, + :(k[ss] / y[ss] = 8.0) # Missing | ctou + ) + + model = nothing + end + + @testset "Error - calibrate non-existent parameter" begin + model = load_RBC_baseline() + + @test_throws ErrorException add_calibration_equation!( + model, + :(y[ss] = 1.5 | fake_param) + ) + + model = nothing + end +end + + +@testset verbose = true "SW07 remove_calibration_equation! functionality" begin + @testset "Remove calibration with explicit parameters override" begin + model = load_RBC_baseline() + + n_calib_before = length(get_calibration_equations(model)) + calib_params_before = get_calibrated_parameters(model) + @test "ψ" in calib_params_before + + # Remove ψ calibration with explicit value via parameters + remove_calibration_equation!(model, 2, parameters = :ψ => 1.0) + + @test length(get_calibration_equations(model)) == n_calib_before - 1 + @test !("ψ" in get_calibrated_parameters(model)) + @test "ψ" in get_parameters(model) + + history = get_revision_history(model) + @test length(history) == 1 + @test history[1].action == :remove_calibration_equation + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Remove calibration using current value" begin + model = load_RBC_baseline() + + n_calib_before = length(get_calibration_equations(model)) + + # Remove without specifying parameters - uses current calibrated value + remove_calibration_equation!(model, 1) + + @test length(get_calibration_equations(model)) == n_calib_before - 1 + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Remove calibration using vector syntax" begin + model = load_RBC_baseline() + + n_calib_before = length(get_calibration_equations(model)) + @test n_calib_before >= 2 # SW07 has mcflex and cpie calibrations + + # Remove first calibration equation + remove_calibration_equation!(model, [1,2]) + + @test length(get_calibration_equations(model)) == n_calib_before - 2 + + model = nothing + end + + @testset "Add then remove calibration - round trip" begin + model = load_FS2000() + + n_calib_before = length(get_calibration_equations(model)) + + # Add a calibration + add_calibration_equation!(model, :(k[ss] / y[ss] = 8.0 | del), silent = true) + @test length(get_calibration_equations(model)) == n_calib_before + 1 + + ss_calib = get_steady_state(model, derivatives = false) + + # Remove it (last added) + remove_calibration_equation!(model, n_calib_before + 1) + @test length(get_calibration_equations(model)) == n_calib_before + + # Model should solve + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end + + @testset "Error - remove from model with no calibrations" begin + model = load_FS2000() + + # Now trying to remove should error + @test_throws AssertionError remove_calibration_equation!(model, 1) + + model = nothing + end +end + + +@testset verbose = true "SW07 get_revision_history functionality" begin + @testset "Empty history for fresh model" begin + model = load_RBC_baseline() + + history = get_revision_history(model) + @test length(history) == 0 + + model = nothing + end + + @testset "History tracks all operations" begin + model = load_FS2000() + + eqs = get_equations(model) + + # 2. Add equation + add_equation!(model, :(log_y[0] = log(y[0]))) + + # 4. Add calibration + add_calibration_equation!(model, :(k[ss] / y[ss] = 8.0 | del)) + + # Check history + history = get_revision_history(model) + @test length(history) == 2 + # @test history[1].action == :update_equation + @test history[1].action == :add_equation + # @test history[3].action == :update_calibration_equation + @test history[2].action == :add_calibration_equation + + # All entries should have timestamps + @test all(h -> haskey(h, :timestamp), history) + + model = nothing + end +end + + +@testset verbose = true "SW07 complex modification scenarios" begin + @testset "Remove observables and add custom ones" begin + model = load_sw07() + + eqs = get_equations(model) + + # Remove standard observables + labobs_idx = findfirst(eq -> occursin("labobs[0]", eq), string.(eqs)) + remove_equation!(model, labobs_idx) + + # Add custom observable (employment rate instead of hours) + add_equation!(model, :(employment_obs[0] = 100 * log(lab[0] / lab[ss]))) + + @test !("labobs" in get_variables(model)) + @test "employment_obs" in get_variables(model) + + history = get_revision_history(model) + @test length(history) == 2 + + ss = get_steady_state(model, derivatives = false) + @test !any(isnan, ss) + + model = nothing + end +end + + +@testset verbose = true "SW07 dynamics verification" begin + + @testset "Taylor rule modification changes IRF" begin + model = load_sw07() + + # Get original IRF + irf_original = get_irf(model) + + # Modify Taylor rule - remove output growth response + old_taylor = :(r[0] = r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / cpie) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * (y[0] / yflex[0] / (y[-1] / yflex[-1])) ^ crdy * ms[0]) + new_taylor = :(r[0] = r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / cpie) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * ms[0]) + update_equations!(model, old_taylor, new_taylor) + + irf_modified = get_irf(model) + + # IRFs should be different (removed output growth response) + @test !isapprox(collect(irf_original), collect(irf_modified), rtol = 1e-5) + + model = nothing + end +end +