Cache model-constant computations to minimize allocations across all modules (#234)

Author: Copilot

.github/workflows/ci.yml

@@ -152,7 +152,7 @@ jobs:
             arch: x64
             test_set: "jet"
           - version: 'lts'
-            os: ubuntu-latest
+            os: macOS-latest
             arch: x64
             test_set: "jet"
     steps:
@@ -186,9 +186,9 @@ jobs:
       - name: Set Custom Test Environment Variable (non-Windows)
         if: matrix.os != 'windows-latest'
         run: echo "TEST_SET=${{ matrix.test_set }}" >> $GITHUB_ENV
-      - name: Set JULIA_NUM_THREADS for estimation tests
-        if: (matrix.version == '1' && (matrix.test_set == 'estimation' || matrix.test_set == 'estimate_sw07' || matrix.test_set == '1st_order_inversion_estimation' || matrix.test_set == 'estimation_pigeons' || matrix.test_set == '1st_order_inversion_estimation_pigeons'))
-        run: echo "JULIA_NUM_THREADS=auto" >> $GITHUB_ENV
+      # - name: Set JULIA_NUM_THREADS for estimation tests
+      #   if: (matrix.version == '1' && (matrix.test_set == 'estimation' || matrix.test_set == 'estimate_sw07' || matrix.test_set == '1st_order_inversion_estimation' || matrix.test_set == 'estimation_pigeons' || matrix.test_set == '1st_order_inversion_estimation_pigeons'))
+      #   run: echo "JULIA_NUM_THREADS=auto" >> $GITHUB_ENV
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1

AGENTS.md

@@ -30,6 +30,14 @@ like `[0]` (present), `[1]` (future), `[-1]` (past), and `[x]` (shock).
 
 ## Tests
 
+Before running tests, activate the test environment and instantiate dependencies:
+
+```julia
+using Pkg
+Pkg.activate("test")
+Pkg.instantiate()
+```
+
 Tests are organized by test sets specified via `TEST_SET`:
 
 ```bash

AGENT_PROGRESS.md

@@ -13,6 +13,14 @@ MacroModelling.jl is a Julia package for developing and solving Dynamic Stochast
 
 ### Running Tests
 
+Before running tests, activate the test environment and instantiate dependencies:
+
+```julia
+using Pkg
+Pkg.activate("test")
+Pkg.instantiate()
+```
+
 Tests are organized by test sets specified via the `TEST_SET` environment variable:
 
 ```bash

CLAUDE.md

@@ -16,6 +16,56 @@ import LinearAlgebra as ℒ
 using MacroModelling
 import MacroModelling: clear_solution_caches!, get_NSSS_and_parameters, calculate_jacobian, merge_calculation_options, solve_lyapunov_equation, ℳ
 
+# Check if new workspace API is available (not present in old package versions)
+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!
+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)
+    end
+end
+
+function timings_for_bench(𝓂::ℳ)
+    if hasproperty(𝓂, :timings)
+        out = 𝓂.timings
+    else
+        out = 𝓂.constants.post_model_macro
+    end
+    return out
+end
+
+function first_order_solution_for_bench(∇₁::AbstractMatrix, 𝓂::ℳ; opts = merge_calculation_options())
+    if HAS_WORKSPACE_API
+        qme_ws = ensure_qme_workspace!(𝓂)
+        sylv_ws = ensure_sylvester_1st_order_workspace!(𝓂)
+        out = calculate_first_order_solution(∇₁, 𝓂.constants, qme_ws, sylv_ws; opts = opts)
+    else
+        out = calculate_first_order_solution(∇₁; T = timings_for_bench(𝓂), opts = opts)
+    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)
+    end
+    return out
+end
+
 
 function run_benchmarks!(𝓂::ℳ, SUITE::BenchmarkGroup)
     SUITE[𝓂.model_name] = BenchmarkGroup()
@@ -34,37 +84,46 @@ function run_benchmarks!(𝓂::ℳ, SUITE::BenchmarkGroup)
     SUITE[𝓂.model_name]["NSSS"] = @benchmarkable get_NSSS_and_parameters($𝓂, $𝓂.parameter_values) setup = clear_solution_caches!($𝓂, :first_order)
 
 
-    ∇₁ = calculate_jacobian(𝓂.parameter_values, reference_steady_state, 𝓂)
+    ∇₁ = calculate_jacobian_for_bench(𝓂.parameter_values, reference_steady_state, 𝓂)
 
     clear_solution_caches!(𝓂, :first_order)
 
-    SUITE[𝓂.model_name]["jacobian"] = @benchmarkable calculate_jacobian($𝓂.parameter_values, $reference_steady_state, $𝓂) setup = 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()
-    
-    sol, qme_sol, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings, opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :schur))
-    
+
+    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 calculate_first_order_solution($∇₁; T = $𝓂.timings, opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :schur)) setup = clear_solution_caches!($𝓂, :first_order)
-    
-    SUITE[𝓂.model_name]["qme"]["doubling"] = @benchmarkable calculate_first_order_solution($∇₁; T = $𝓂.timings, opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :doubling)) setup = clear_solution_caches!($𝓂, :first_order)
-    
-    
-    A = @views sol[:, 1:𝓂.timings.nPast_not_future_and_mixed] * ℒ.diagm(ones(𝓂.timings.nVars))[𝓂.timings.past_not_future_and_mixed_idx,:]
-    
-    C = @views sol[:, 𝓂.timings.nPast_not_future_and_mixed+1:end]
-    
+
+    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(𝓂)
+
+    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'
 
-    solve_lyapunov_equation(A, CC)
+    # 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_equation($A, $CC, lyapunov_algorithm = :doubling) # setup = clear_solution_caches!($𝓂, :first_order)
-    SUITE[𝓂.model_name]["lyapunov"]["bartels_stewart"] = @benchmarkable solve_lyapunov_equation($A, $CC, lyapunov_algorithm = :bartels_stewart) # setup = clear_solution_caches!($𝓂, :first_order)
-    SUITE[𝓂.model_name]["lyapunov"]["bicgstab"] = @benchmarkable solve_lyapunov_equation($A, $CC, lyapunov_algorithm = :bicgstab) # setup = clear_solution_caches!($𝓂, :first_order)
-    SUITE[𝓂.model_name]["lyapunov"]["gmres"] = @benchmarkable solve_lyapunov_equation($A, $CC, lyapunov_algorithm = :gmres) # setup = clear_solution_caches!($𝓂, :first_order)
+    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)
@@ -103,7 +162,7 @@ run_benchmarks!(Smets_Wouters_2007, SUITE)
 #     end
 # end
 
-# If a cache of tuned parameters already exists, use it, otherwise, tune and cache
+# 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")
@@ -113,4 +172,4 @@ run_benchmarks!(Smets_Wouters_2007, SUITE)
 # else
 #     tune!(SUITE)
 #     BenchmarkTools.save(paramspath, params(SUITE))
-# end
+# end

benchmark/benchmarks.jl

@@ -6,7 +6,7 @@
 
 using Documenter
 using MacroModelling
-import Turing, StatsPlots
+import Optim, StatsPlots, Turing
 using DocumenterCitations
 
 bib = CitationBibliography(
@@ -23,7 +23,12 @@ makedocs(
     # doctest = false,
     # draft = true,
     format = Documenter.HTML(size_threshold = 204800*10),
-    modules = [MacroModelling],
+    modules = [
+        MacroModelling,
+        Base.get_extension(MacroModelling, :OptimExt),
+        Base.get_extension(MacroModelling, :StatsPlotsExt),
+        Base.get_extension(MacroModelling, :TuringExt),
+    ],
     pages = [
         "Introduction" => "index.md",
         "Tutorials" => [

docs/make.jl

@@ -1,4 +1,9 @@
 ```@autodocs
-Modules = [MacroModelling]
+Modules = [
+    MacroModelling,
+    Base.get_extension(MacroModelling, :OptimExt),
+    Base.get_extension(MacroModelling, :StatsPlotsExt),
+    Base.get_extension(MacroModelling, :TuringExt),
+]
 Order   = [:function, :macro, :module]
-```
+```

docs/src/api.md

@@ -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
+import MacroModelling: find_shocks_conditional_forecast, find_SS_solver_parameters!, Tolerances, ℳ, calculate_SS_solver_runtime_and_loglikelihood, solver_parameters, find_shocks_workspace
 import Optim
 
 # Helper function for LBFGS optimization objective
@@ -24,8 +24,8 @@ end
                                     conditions::Vector{Float64},
                                     cond_var_idx::Vector{Int},
                                     free_shock_idx::Vector{Int},
-                                    pruning::Bool,
-                                    S₁, S₂, S₃, timings; verbose::Bool = false)
+                                    state_update::Function,
+                                    S₁, S₂, S₃, constants, ws; verbose::Bool = false)
 
 Find shocks that satisfy conditional forecast constraints using LBFGS optimizer.
 
@@ -37,10 +37,11 @@ Note: This is the Optim-based implementation. It requires the Optim.jl extension
 - `conditions`: Target values for conditioned variables
 - `cond_var_idx`: Indices of conditioned variables
 - `free_shock_idx`: Indices of free shocks to be determined
-- `pruning`: Whether pruning is used
+- `state_update`: State update function for the selected algorithm
 - `S₁`: First-order solution matrix
 - `S₂`, `S₃`: Higher-order perturbation matrices (not used in LBFGS, for compatibility only)
-- `timings`: Model timings structure
+- `constants`: Model constants (unused for LBFGS)
+- `ws`: Find shocks workspace (unused for LBFGS)
 
 # Returns
 - `x`: Vector of optimal shock values
@@ -55,7 +56,9 @@ function find_shocks_conditional_forecast(::Val{:LBFGS},
                                          free_shock_idx::Vector{Int},
                                          state_update::Function,
                                         #  pruning::Bool,
-                                         S₁, S₂, S₃, timings; verbose::Bool = false)
+                                         S₁, S₂, S₃, constants, 
+                                         ws::find_shocks_workspace{Float64}; 
+                                         verbose::Bool = false)
 
     pruning = typeof(initial_state) <: Vector{Vector{Float64}}
     precision_factor = 1.0
@@ -94,7 +97,7 @@ function find_shocks_conditional_forecast(::Val{:LBFGS},
 end
 
 """
-    find_SS_solver_parameters!(::Val{:SAMIN}, 𝓂::ℳ; maxtime::Int = 120, maxiter::Int = 2500000, tol::Tolerances = Tolerances(), verbosity = 0)
+    find_SS_solver_parameters!(::Val{:SAMIN}, 𝒂::ℳ; maxtime::Real = 120, maxiter::Int = 2500000, tol::Tolerances = Tolerances(), verbosity = 0)
 
 Find optimal steady state solver parameters using Optim's SAMIN algorithm.
 
@@ -109,7 +112,7 @@ It uses Simulated Annealing with Metropolis acceptance (SAMIN) from Optim.jl.
 - `verbosity`: Verbosity level for output
 """
 function find_SS_solver_parameters!(::Val{:SAMIN}, 𝓂::ℳ; 
-                                                    maxtime::Int = 120, 
+                                                    maxtime::Real = 120, 
                                                     maxiter::Int = 2500000, 
                                                     tol::Tolerances = Tolerances(), 
                                                     verbosity = 0)
@@ -131,10 +134,10 @@ function find_SS_solver_parameters!(::Val{:SAMIN}, 𝓂::ℳ;
 
     par_inputs = solver_parameters(pars..., 1, 0.0, 2)
 
-    SS_and_pars, (solution_error, iters) = 𝓂.SS_solve_func(𝓂.parameter_values, 𝓂, tol, false, true, [par_inputs])
+    SS_and_pars, (solution_error, iters) = 𝓂.functions.NSSS_solve(𝓂.parameter_values, 𝓂, tol, false, true, [par_inputs])
 
     if solution_error < tol.NSSS_acceptance_tol
-        push!(𝓂.solver_parameters, par_inputs)
+        push!(MacroModelling.DEFAULT_SOLVER_PARAMETERS, par_inputs)
         return true
     else 
         return false