Refactor take2 (#238)

Author: thorek1

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

@@ -0,0 +1,71 @@
+## Summary
+- Reworked first-order index cache to avoid `Union{Nothing,...}` and added cached identity-slice matrices for `expand` usage.
+- Switched `calculate_first_order_solution` to positional `𝓂` wrappers (Real and Dual) and updated call sites to stop passing caches explicitly.
+- Added a runnable script for RBC baseline pruned third-order standard deviations and updated the standalone cache test to use the new initializer and positional wrapper.
+- Fixed pruned mean calculation to use state-only kron indices and added a runnable script for RBC baseline pruned second-order mean.
+- Fixed `get_steady_state` axis labels for models with calibration equations and added a runnable script to exercise calibration equations.
+- Added a moments cache for kron selectors, `I_plus_s_s`, and `e⁴`/`e⁶`, plus substate/dependency caches, and switched `moments.jl` to use cached values.
+- Extended model structure caching with steady-state selector matrices, full NSSS display names, and cached variable indices; updated steady-state and variable parsing code to use cached values.
+- Added a targeted cache validation script for steady-state expansion, custom steady state mapping, and cached variable indices.
+- Moved cache creation helpers into `src/options_and_caches.jl` (name display, computational constants, model structure, first-order indices, selector matrix, moments) and removed duplicate definitions from `src/MacroModelling.jl` and `src/moments.jl`.
+- Refactored cache usage to call ensure helpers once per high-level call and use `𝓂.caches` in `src/moments.jl`, `src/perturbation.jl`, `src/get_functions.jl`, and filter paths; updated steady-state/model-structure callers to use cached selectors.
+- Fixed package load/precompile failures caused by include-order cycles (`timings` and `ℳ`) by introducing `src/timings.jl` and removing the `::ℳ` annotation from `initialize_caches!`.
+- Removed obsolete `T = 𝓂.timings` keywords from second/third-order solution call sites and fixed a typo (`ℂC` → `cache`) in `calculate_third_order_solution`.
+- Made `get_loglikelihood(..., algorithm = :pruned_third_order)` work by adding `get_computational_constants(::timings)` and fixing the inversion filter’s pruned third-order Jacobian construction (state² term needed an explicit volatility slot).
+- Added `cache::caches` overloads for `determine_efficient_order` and switched `src/moments.jl` call sites from `𝓂.timings` to `𝓂.caches`.
+- Replaced cache→timings wrapper overloads by making cache-carrying methods the concrete implementations for `determine_efficient_order`, `create_*_order_auxiliary_matrices`, `solve_quadratic_matrix_equation`, and `parse_*_input_to_index`/`parse_covariance_groups`.
+- Added a targeted `get_std` script test for pruned second/third order on RBC_CME models with/without calibration equations.
+- Fixed `calculate_first_order_solution` (Dual path) to use `cache.first_order_index_cache` fields (`nabla_zero_cols`, `nabla_e_start`) after cache refactor.
+- Added a targeted script test covering the three `run_rbc_*.jl` entrypoints.
+- Restored `apply_custom_name` and `normalize_superscript` helpers needed by `StatsPlotsExt` so `plot_model_estimates` works again.
+- Added a targeted script test for `plot_model_estimates` / `plot_model_estimates!`.
+- Fixed `name_display_cache` axis fields to be concrete (`Vector{Symbol}`) so `get_var_axis(::ℳ)` is type-stable and doesn’t trip DispatchDoctor.
+- Added a targeted script test for `get_var_axis` / `get_exo_axis` return types.
+- Added ChainRules `rrule` for `mat_mult_kron` to make Zygote gradients work (avoids mutation errors from `mul!` into views).
+- Fixed the `rrule` signature for `calculate_second_order_solution` to match current argument ordering so Zygote doesn’t trace into LAPACK solves.
+- Added a targeted script test for Zygote gradients of `get_loglikelihood` at `:third_order` and `:pruned_third_order`.
+- Moved second/third-order auxiliary matrices into `caches`, introduced a `workspaces` struct for higher-order scratch buffers, and removed `nothing`-based timings initialization (added `Empty_timings` + `is_empty_timings`).
+- Fixed `ensure_name_display_cache!` to initialize `calib_axis` correctly when curly-brace formatting is used.
+- Moved auxiliary indices into `caches`, added a dedicated `workspaces` field on `ℳ`, and updated higher-order solution calls to pass workspaces explicitly.
+- Updated `run_cache_steady_state_checks.jl` to only call mean/std for supported algorithms.
+
+## Tests
+- `TEST_SET=basic julia --project -e 'using Pkg; Pkg.test()'` failed: unsatisfiable requirements during Pkg resolve (DynamicPPL/Pigeons/Turing/JET constraints).
+- `julia --project -e 'using Test; include("test/test_standalone_function.jl")'` failed: `FiniteDifferences` not found in current environment.
+- `julia --project run_rbc_pruned_third_order_std.jl` succeeded (prints `get_std` output for `RBC_baseline`, algorithm `:pruned_third_order`).
+- Re-ran `julia --project run_rbc_pruned_second_order_mean.jl`: succeeded (prints `get_mean` output for `RBC_baseline`, algorithm `:pruned_second_order`).
+- Re-ran `julia --project run_rbc_pruned_third_order_std.jl`: succeeded (prints `get_std` output for `RBC_baseline`, algorithm `:pruned_third_order`).
+- Re-ran `julia --project run_rbc_pruned_third_order_std.jl`: succeeded (prints `get_std` output for `RBC_baseline`, algorithm `:pruned_third_order`).
+- `julia --project run_rbc_pruned_second_order_mean.jl` succeeded (prints `get_mean` output for `RBC_baseline`, algorithm `:pruned_second_order`).
+- `julia --project run_rbc_pruned_third_order_std.jl` succeeded (prints `get_std` output for `RBC_baseline`, algorithm `:pruned_third_order`).
+- `julia --project run_cache_steady_state_checks.jl` failed: `get_relevant_steady_states` not exported (used unqualified call).
+- `julia --project run_cache_steady_state_checks.jl` succeeded after qualifying `MacroModelling.get_relevant_steady_states`.
+- `julia --project run_rbc_calibration_equations_ss.jl` succeeded (prints calibration equations, calibrated parameter values, and steady state for `RBC_CME_calibration_equations`).
+- `julia --project -e 'using MacroModelling; include("models/RBC_baseline.jl"); MacroModelling.calculate_second_order_moments(RBC_baseline.parameter_values, RBC_baseline);'` failed: `calculate_hessian` method mismatch (hessian functions not initialized for second-order).
+- `julia --project -e 'using MacroModelling; include("models/RBC_baseline.jl"); solve!(RBC_baseline); MacroModelling.calculate_second_order_moments(RBC_baseline.parameter_values, RBC_baseline);'` failed: same `calculate_hessian` method mismatch.
+- `julia --project -e 'using MacroModelling; include("models/RBC_baseline.jl"); solve!(RBC_baseline; algorithm = :second_order); MacroModelling.calculate_second_order_moments(RBC_baseline.parameter_values, RBC_baseline);'` succeeded.
+- `julia --project run_rbc_pruned_second_order_mean.jl` succeeded (prints `get_mean` output for `RBC_baseline`, algorithm `:pruned_second_order`).
+- `julia --project run_rbc_pruned_second_order_mean.jl` succeeded (prints `get_mean` output for `RBC_baseline`, algorithm `:pruned_second_order`).
+- `julia --project run_rbc_pruned_third_order_std.jl` succeeded (prints `get_std` output for `RBC_baseline`, algorithm `:pruned_third_order`).
+- `julia --project run_cache_steady_state_checks.jl` succeeded (prints steady-state cache checks for `RBC_baseline` and `RBC_switch`).
+- `julia --project -e 'using MacroModelling; println(\"loaded\")'` succeeded.
+- `julia --project scripts/test_get_loglikelihood_pruned_third.jl` succeeded (prints finite pruned third-order loglikelihood).
+- `julia --project scripts/test_get_std_pruned_second_third.jl` succeeded (pruned second/third order `get_std` for RBC_CME models with/without calibration equations).
+- Re-ran `julia --project scripts/test_get_loglikelihood_pruned_third.jl`: succeeded.
+- Re-ran `julia --project scripts/test_get_std_pruned_second_third.jl`: succeeded after cache-propagation refactor.
+- `julia --project run_rbc_pruned_second_order_mean.jl` failed: `caches` has no field `nabla_zero_cols`.
+- `julia --project run_rbc_pruned_third_order_std.jl` failed: `caches` has no field `nabla_zero_cols`.
+- `julia --project run_rbc_pruned_second_order_mean.jl` succeeded after fixing Dual first-order cache indexing.
+- `julia --project run_rbc_pruned_third_order_std.jl` succeeded after fixing Dual first-order cache indexing.
+- `julia --project scripts/test_run_rbc_scripts.jl` succeeded (pruned second/third moments + calibration equations steady state).
+- `julia --project scripts/test_get_std_pruned_second_third.jl` succeeded.
+- `julia --project scripts/test_get_loglikelihood_pruned_third.jl` succeeded.
+- `julia --project -e 'using MacroModelling, StatsPlots; include(\"test/models/RBC_CME.jl\"); sim = simulate(m); plot_model_estimates(m, sim([:k], :,:simulate); show_plots=false, save_plots=false)'` succeeded after restoring plotting helpers.
+- `julia --project scripts/test_plot_model_estimates.jl` succeeded.
+- `julia --project scripts/test_get_var_axis.jl` succeeded.
+- `julia --project scripts/test_zygote_loglikelihood_third_pruned_third.jl` succeeded (Zygote gradients for `:third_order` and `:pruned_third_order`).
+- `julia --project run_cache_steady_state_checks.jl` succeeded.
+- `julia --project run_cache_steady_state_checks.jl` succeeded after updating the script’s mean/std algorithm list.
+
+## Remaining
+- Resolve test environment dependencies and rerun tests.

CLAUDE.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

benchmark/benchmarks.jl

@@ -43,18 +43,18 @@ function run_benchmarks!(𝓂::ℳ, SUITE::BenchmarkGroup)
 
     SUITE[𝓂.model_name]["qme"] = BenchmarkGroup()
 
-    sol, qme_sol, solved = calculate_first_order_solution(∇₁; T = 𝓂.timings, opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :schur))
+    sol, qme_sol, solved = calculate_first_order_solution(∇₁, 𝓂.caches; opts = merge_calculation_options(quadratic_matrix_equation_algorithm = :schur))
 
     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"]["schur"] = @benchmarkable calculate_first_order_solution($∇₁, $𝓂.caches; 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)
+    SUITE[𝓂.model_name]["qme"]["doubling"] = @benchmarkable calculate_first_order_solution($∇₁, $𝓂.caches; 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,:]
+    A = @views sol[:, 1:𝓂.caches.timings.nPast_not_future_and_mixed] * ℒ.diagm(ones(𝓂.caches.timings.nVars))[𝓂.caches.timings.past_not_future_and_mixed_idx,:]
 
-    C = @views sol[:, 𝓂.timings.nPast_not_future_and_mixed+1:end]
+    C = @views sol[:, 𝓂.caches.timings.nPast_not_future_and_mixed+1:end]
 
     CC = C * C'
 
@@ -103,7 +103,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")