From 9d781b3be6eecd3eddf0dadfba138722d8e3f0a6 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 7 Jun 2026 18:00:39 +0200 Subject: [PATCH 01/37] Remove try/catch for findnz and improve pullback handling --- Project.toml | 1 + src/MacroModelling.jl | 4 + src/algorithms/lyapunov.jl | 21 +- src/algorithms/sylvester.jl | 20 +- src/rrules.jl | 544 ++++++++++++++++++++---- test/check_third_order_solution_grad.jl | 382 +++++++++++++++++ 6 files changed, 869 insertions(+), 103 deletions(-) create mode 100644 test/check_third_order_solution_grad.jl diff --git a/Project.toml b/Project.toml index 3d5afe3ba..8664e26d9 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ authors = ["Thore Kockerols "] [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" BlockTriangularForm = "adeb47b7-70bf-415a-bb24-c358563e873a" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index f10ccf6b0..a8aff7acd 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -862,6 +862,10 @@ end 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) + + # a = sparse(A) + # droptol!(a, tol) + if multithreaded return ThreadedSparseArrays.ThreadedSparseMatrixCSC(a) else diff --git a/src/algorithms/lyapunov.jl b/src/algorithms/lyapunov.jl index aad87f7b2..35415b77b 100644 --- a/src/algorithms/lyapunov.jl +++ b/src/algorithms/lyapunov.jl @@ -1,6 +1,6 @@ @stable default_mode = "disable" begin -# Available algorithms: +# Available algorithms: # :doubling - fast and precise # :bartels_stewart - fast for small matrices and precise, dense matrices only # :bicgstab - less precise @@ -113,7 +113,7 @@ end denom = max(ℒ.norm(guess), ℒ.norm(C)) reached_tol = denom == 0 ? 0.0 : ℒ.norm(res_buf) / 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 + if verbose println("Lyapunov equation - initial guess achieves relative tol of $reached_tol (initial guess tol: $initial_guess_acceptance_tol) [A: $(typeof(A)), C: $(typeof(C))]") end return choose_matrix_format(guess), true end end @@ -133,7 +133,7 @@ end verbose = verbose) if deflation_solved if verbose - println("Lyapunov equation - solved via Schur deflation (unit roots pre-detected)") + println("Lyapunov equation - solved via Schur deflation (unit roots pre-detected) [A: $(typeof(A_dense)), C: $(typeof(C_dense))]") end return X_deflated, true end @@ -143,7 +143,7 @@ end X, i, reached_tol = solve_lyapunov_equation(A, C, Val(lyapunov_algorithm), workspace; tol = tol) # timer = timer) if verbose - println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: $lyapunov_algorithm") + println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: $lyapunov_algorithm [A: $(typeof(A)), C: $(typeof(C))]") end if reached_tol > acceptance_tol && lyapunov_algorithm ≠ :doubling @@ -152,7 +152,7 @@ end X, i, reached_tol = solve_lyapunov_equation(A, C, Val(:doubling), workspace; tol = tol) # timer = timer) if verbose - println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: doubling") + println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: doubling [A: $(typeof(A)), C: $(typeof(C))]") end end @@ -162,7 +162,7 @@ end X, i, reached_tol = solve_lyapunov_equation(A, C, Val(:bicgstab), workspace; tol = tol) # timer = timer) if verbose - println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bicgstab") + println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bicgstab [A: $(typeof(A)), C: $(typeof(C))]") end end @@ -174,7 +174,7 @@ end X, i, reached_tol = solve_lyapunov_equation(A, C, Val(:bartels_stewart), workspace; tol = tol) # timer = timer) if verbose - println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bartels_stewart") + println("Lyapunov equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bartels_stewart [A: $(typeof(A)), C: $(typeof(C))]") end end @@ -191,7 +191,7 @@ end X = X_deflated reached_tol = zero(T) # signal success if verbose - println("Lyapunov equation - solved via Schur deflation (unit-root subspace set to NaN)") + println("Lyapunov equation - solved via Schur deflation (unit-root subspace set to NaN) [A: $(typeof(A_dense)), C: $(typeof(C_dense))]") end end end @@ -917,7 +917,7 @@ function solve_lyapunov_schur_deflation(A::DenseMatrix{T}, if sub_tol > tol.acceptance_tol if verbose - println("Schur deflation: stable sub-block Lyapunov failed (tol=$sub_tol)") + println("Schur deflation: stable sub-block Lyapunov failed (tol=$sub_tol) [T_ss: $(typeof(T_ss)), C_ss: $(typeof(C_ss))]") end return Matrix{T}(undef, 0, 0), false end @@ -946,7 +946,8 @@ function solve_lyapunov_schur_deflation(A::DenseMatrix{T}, if verbose println("Schur deflation: $n_unstable unstable eigenvalue(s), ", - "$n_stable stable, $(count(unit_root_vars)) variable(s) set to NaN") + "$n_stable stable, $(count(unit_root_vars)) variable(s) set to NaN ", + "[T_ss: $(typeof(T_ss)), C_ss: $(typeof(C_ss))]") end return Σ, true diff --git a/src/algorithms/sylvester.jl b/src/algorithms/sylvester.jl index 99fb88763..827e5ff3c 100644 --- a/src/algorithms/sylvester.jl +++ b/src/algorithms/sylvester.jl @@ -1,6 +1,6 @@ @stable default_mode = "disable" begin -# Available algorithms: +# Available algorithms: # :doubling - fast, expensive part: B^2 # :bartels_stewart - fast, dense matrices only # :bicgstab - fastest for large problems, might not reach desired precision, warm start not always helpful @@ -89,7 +89,7 @@ reached_tol = denom == 0 ? 0.0 : ℒ.norm(res_buf) / denom 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 + if verbose println("Sylvester equation - $guess_name achieves relative tol of $reached_tol (initial guess tol: $initial_guess_acceptance_tol) [A: $(size(a)) $(typeof(a)), B: $(size(b)) $(typeof(b)), C: $(size(c)) $(typeof(c))]") end return choose_matrix_format(guess), true end @@ -106,7 +106,7 @@ verbose = verbose) if verbose && i != 0 - println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: $sylvester_algorithm") + println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: $sylvester_algorithm [A: $(size(a)) $(typeof(a)), B: $(size(b)) $(typeof(b)), C: $(size(c)) $(typeof(c))]") end if (!isfinite(reached_tol) || !(reached_tol < acceptance_tol)) && sylvester_algorithm ≠ :doubling @@ -125,7 +125,7 @@ verbose = verbose) if verbose# && i != 0 - println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: doubling") + println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: doubling [A: $(size(aa)) $(typeof(aa)), B: $(size(b)) $(typeof(b)), C: $(size(cc)) $(typeof(cc))]") end end @@ -148,7 +148,7 @@ verbose = verbose) if verbose && i != 0 - println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bartels_stewart") + println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: bartels_stewart [A: $(size(aa)) $(typeof(aa)), B: $(size(bb)) $(typeof(bb)), C: $(size(cc)) $(typeof(cc))]") end end @@ -173,7 +173,7 @@ end if verbose# && i != 0 - println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $Reached_tol; algorithm: dqgmres (refinement of previous solution)") + println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $Reached_tol; algorithm: dqgmres (refinement of previous solution) [A: $(size(aa)) $(typeof(aa)), B: $(size(b)) $(typeof(b)), C: $(size(cc)) $(typeof(cc))]") end end @@ -193,7 +193,7 @@ verbose = verbose) if verbose# && i != 0 - println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: gmres") + println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $reached_tol; algorithm: gmres [A: $(size(aa)) $(typeof(aa)), B: $(size(b)) $(typeof(b)), C: $(size(cc)) $(typeof(cc))]") end end @@ -205,9 +205,9 @@ cc = 𝕊ℂ.𝐂¹ copyto!(cc, C) - X, i, Reached_tol = solve_sylvester_equation(aa, b, cc, + X, i, Reached_tol = solve_sylvester_equation(aa, b, cc, Val(:dqgmres), 𝕊ℂ, - initial_guess = x, + initial_guess = x, preconditioner = preconditioner, tol = tol, # timer = timer, @@ -218,7 +218,7 @@ end if verbose# && i != 0 - println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $Reached_tol; algorithm: dqgmres (refinement of previous solution)") + println("Sylvester equation - converged to tol $acceptance_tol: $(reached_tol < acceptance_tol); iterations: $i; reached tol: $Reached_tol; algorithm: dqgmres (refinement of previous solution) [A: $(size(aa)) $(typeof(aa)), B: $(size(b)) $(typeof(b)), C: $(size(cc)) $(typeof(cc))]") end end diff --git a/src/rrules.jl b/src/rrules.jl index 07b186a28..0ab6f003c 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -6696,64 +6696,203 @@ 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} +function mul_fill_kron_adjoint!( + ∂A::AbstractMatrix{R}, + ∂B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::SparseMatrixCSC{TA, Ti}, + B::AbstractMatrix{TB}; + tol::Real = 0.0, + block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0), +) where {R <: Real, TA <: Real, TB <: Real, Ti <: Integer} + 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) + size(M1, 1) == n1 * n2 || throw(DimensionMismatch("size(M1, 1) must equal size(B, 1) * size(A, 1)")) + size(M2, 2) == m1 * m2 || throw(DimensionMismatch("size(M2, 2) must equal size(B, 2) * size(A, 2)")) + size(M1, 2) == size(M2, 1) || throw(DimensionMismatch("inner dimensions of M1 and M2 do not match")) nrows = n1 * n2 - - if size(block, 1) == nrows && size(block, 2) >= m2 - blk = view(block, :, 1:m2) + tol_R = R(tol) + + # Transpose M2 so that M2t[(j - 1) * m2 + l, s] is contiguous in l. + M2t = Matrix{R}(transpose(M2)) + + # Row compressed representation of M1. + # Use conditional handling instead of try/catch to obtain non‑zero indices and values. + # If `M1` is already a sparse matrix we can call `findnz` directly; otherwise we + # convert it to a sparse representation first. This satisfies the request to + # avoid `try … catch` while preserving the original behaviour. + if typeof(M1) <: AbstractSparseMatrix + M1_row_ids, M1_col_ids, M1_values = findnz(M1) else - blk = Matrix{R}(undef, nrows, m2) + M1_row_ids, M1_col_ids, M1_values = findnz(sparse(M1)) end - A_colptr = A.colptr - A_rowval = A.rowval - A_nzval = A.nzval + M1_row_counts = zeros(Int, nrows) + + @inbounds for row in M1_row_ids + M1_row_counts[row] += 1 + end + + n_supported_rows = 0 + + @inbounds for row in 1:nrows + if M1_row_counts[row] != 0 + n_supported_rows += 1 + end + end + + row_to_slot = zeros(Int, nrows) + supported_rows = Vector{Int}(undef, n_supported_rows) + + slot = 0 + + @inbounds for row in 1:nrows + if M1_row_counts[row] != 0 + slot += 1 + row_to_slot[row] = slot + supported_rows[slot] = row + end + end + + M1_row_ptr = Vector{Int}(undef, n_supported_rows + 1) + M1_row_ptr[1] = 1 + + @inbounds for slot in 1:n_supported_rows + M1_row_ptr[slot + 1] = M1_row_ptr[slot] + M1_row_counts[supported_rows[slot]] + end + + M1_nz_col = Vector{Int}(undef, length(M1_values)) + M1_nz_val = Vector{R}(undef, length(M1_values)) + next_ptr = copy(M1_row_ptr) + + @inbounds for index in eachindex(M1_values) + row = M1_row_ids[index] + slot = row_to_slot[row] + position = next_ptr[slot] + next_ptr[slot] += 1 + + M1_nz_col[position] = M1_col_ids[index] + M1_nz_val[position] = R(M1_values[index]) + end + + M1_row_k = Vector{Int}(undef, n_supported_rows) + M1_row_i = Vector{Int}(undef, n_supported_rows) + + @inbounds for slot in 1:n_supported_rows + row = supported_rows[slot] + i = (row - 1) ÷ n2 + 1 + k = row - (i - 1) * n2 + + M1_row_k[slot] = k + M1_row_i[slot] = i + end + + # Row compressed representation of A. + A_row_ids, A_col_ids, A_values = findnz(A) + + A_row_counts = zeros(Int, n2) + + @inbounds for row in A_row_ids + A_row_counts[row] += 1 + end + + A_row_ptr = Vector{Int}(undef, n2 + 1) + A_row_ptr[1] = 1 + + @inbounds for row in 1:n2 + A_row_ptr[row + 1] = A_row_ptr[row] + A_row_counts[row] + end + + A_nz_col = Vector{Int}(undef, length(A_values)) + A_nz_val = Vector{R}(undef, length(A_values)) + next_A_ptr = copy(A_row_ptr) + + @inbounds for index in eachindex(A_values) + row = A_row_ids[index] + position = next_A_ptr[row] + next_A_ptr[row] += 1 + + A_nz_col[position] = A_col_ids[index] + A_nz_val[position] = R(A_values[index]) + end + + work = Vector{R}(undef, m2) @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)) + column_offset = (j - 1) * m2 - re_blk = reshape(blk, n2, n1, m2) + for row_slot in 1:n_supported_rows + k = M1_row_k[row_slot] + i = M1_row_i[row_slot] - # ∂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] + bij = R(B[i, j]) + has_B_contribution = abs(bij) > tol_R + + A_row_start = A_row_ptr[k] + A_row_stop = A_row_ptr[k + 1] - 1 + has_A_contribution = A_row_start <= A_row_stop + + if !(has_B_contribution || has_A_contribution) + continue + end + + M1_row_start = M1_row_ptr[row_slot] + M1_row_stop = M1_row_ptr[row_slot + 1] - 1 + + if has_B_contribution + fill!(work, zero(R)) + + for M1_index in M1_row_start:M1_row_stop + M1_value = M1_nz_val[M1_index] + M1_col = M1_nz_col[M1_index] + + @simd for l in 1:m2 + work[l] += M1_value * M2t[column_offset + l, M1_col] 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] + for l in 1:m2 + ∂A[k, l] += bij * work[l] + end + + if has_A_contribution + acc = zero(R) + + for A_index in A_row_start:A_row_stop + l = A_nz_col[A_index] + acc += A_nz_val[A_index] * work[l] + end + + ∂B[i, j] += acc + end + + else + acc = zero(R) + + for M1_index in M1_row_start:M1_row_stop + M1_value = M1_nz_val[M1_index] + M1_col = M1_nz_col[M1_index] + + local_acc = zero(R) + + for A_index in A_row_start:A_row_stop + l = A_nz_col[A_index] + local_acc += A_nz_val[A_index] * M2t[column_offset + l, M1_col] + end + + acc += M1_value * local_acc end + + ∂B[i, j] += acc end - ∂B[i, j] += acc end end + + return ∂A, ∂B end @@ -7683,7 +7822,126 @@ 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 +function compressed_kron³_pullback!( + ∂X::AbstractMatrix{T}, + ∂Y::AbstractSparseMatrix{T, Ti}, + X::AbstractMatrix{T}; + tol::Real = 0.0, +) where {T <: Real, Ti <: Integer} + + Xd = X isa DenseMatrix ? X : collect(X) + + n_x_rows, n_x_cols = size(Xd) + + compressed_rows = n_x_rows * (n_x_rows + 1) * (n_x_rows + 2) ÷ 6 + compressed_cols = n_x_cols * (n_x_cols + 1) * (n_x_cols + 2) ÷ 6 + + size(∂Y, 1) == compressed_rows || throw(DimensionMismatch("size(∂Y, 1) is inconsistent with size(X, 1)")) + size(∂Y, 2) == compressed_cols || throw(DimensionMismatch("size(∂Y, 2) is inconsistent with size(X, 2)")) + + tol_T = T(tol) + + row_values = SparseArrays.rowvals(∂Y) + non_zero_values = nonzeros(∂Y) + + row_first = Vector{Int}(undef, compressed_rows) + row_second = Vector{Int}(undef, compressed_rows) + row_third = Vector{Int}(undef, compressed_rows) + row_scale = Vector{T}(undef, compressed_rows) + + @inbounds for i1 in 1:n_x_rows + for j1 in 1:i1 + for k1 in 1:j1 + row = (i1 - 1) * i1 * (i1 + 1) ÷ 6 + (j1 - 1) * j1 ÷ 2 + k1 + + row_first[row] = i1 + row_second[row] = j1 + row_third[row] = k1 + + if i1 == j1 + divisor = j1 == k1 ? 6 : 2 + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end + + row_scale[row] = inv(T(divisor)) + end + end + end + + col_first = Vector{Int}(undef, compressed_cols) + col_second = Vector{Int}(undef, compressed_cols) + col_third = Vector{Int}(undef, compressed_cols) + + @inbounds for i2 in 1:n_x_cols + for j2 in 1:i2 + for k2 in 1:j2 + col = (i2 - 1) * i2 * (i2 + 1) ÷ 6 + (j2 - 1) * j2 ÷ 2 + k2 + + col_first[col] = i2 + col_second[col] = j2 + col_third[col] = k2 + end + end + end + + @inbounds for col in 1:compressed_cols + i2 = col_first[col] + j2 = col_second[col] + k2 = col_third[col] + + for stored_index in SparseArrays.nzrange(∂Y, col) + row = Int(row_values[stored_index]) + g = non_zero_values[stored_index] + + abs(g) <= tol_T && continue + + i1 = row_first[row] + j1 = row_second[row] + k1 = row_third[row] + + g_d = g * row_scale[row] + + aii = Xd[i1, i2] + aij = Xd[i1, j2] + aik = Xd[i1, k2] + + aji = Xd[j1, i2] + ajj = Xd[j1, j2] + ajk = Xd[j1, k2] + + aki = Xd[k1, i2] + akj = Xd[k1, j2] + akk = Xd[k1, k2] + + q_i2 = aij * ajk + aik * ajj + q_j2 = aik * aji + aii * ajk + q_k2 = aii * ajj + aij * aji + + ∂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 + + return ∂X +end + +function compressed_kron³_pullback!(∂X::AbstractMatrix{T}, + ∂Y::DenseMatrix{T}, + X::AbstractMatrix{T}; + tol::Real = 0.0) where T <: Real + println("size(∂X) = ", size(∂X), " approxzeros/length = ", count(x -> abs(x) <= tol, ∂X) / length(∂X)) + println("size(∂Y) = ", size(∂Y), " approxzeros/length = ", count(x -> abs(x) <= tol, ∂Y) / length(∂Y)) + println("size(X) = ", size(X), " approxzeros/length = ", count(x -> abs(x) <= tol, X) / length(X)) 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 @@ -7746,60 +8004,179 @@ end # 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 +function mul_compressed_kron³_pullback!( + ∂X::AbstractMatrix, + M1::AbstractMatrix, + M2::AbstractMatrix, + X::AbstractMatrix; + tol::Real = 0.0, +) + Xd = X isa DenseMatrix ? X : collect(X) - n_rows, n_cols = size(Xd) - g_col = Vector{T}(undef, size(M1, 1)) + n_x_rows, n_x_cols = size(Xd) - 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 + compressed_rows = n_x_rows * (n_x_rows + 1) * (n_x_rows + 2) ÷ 6 + compressed_cols = n_x_cols * (n_x_cols + 1) * (n_x_cols + 2) ÷ 6 - # Compute g_col = M1 * M2[:, col] lazily for this triple - ℒ.mul!(g_col, M1, view(M2, :, col)) + size(M1, 1) == compressed_rows || throw(DimensionMismatch("size(M1, 1) is inconsistent with size(X, 1)")) + size(M2, 2) == compressed_cols || throw(DimensionMismatch("size(M2, 2) is inconsistent with size(X, 2)")) + size(M1, 2) == size(M2, 1) || throw(DimensionMismatch("inner dimensions of M1 and M2 do not match")) + + T = promote_type(eltype(∂X), eltype(M1), eltype(M2), eltype(Xd)) + tol_T = T(tol) + + # Avoid try/catch: directly handle dense vs sparse `M1`. + if typeof(M1) <: AbstractSparseMatrix + row_ids, col_ids, values = findnz(M1) + else + row_ids, col_ids, values = findnz(sparse(M1)) + end + + row_counts = zeros(Int, compressed_rows) + + @inbounds for row in row_ids + row_counts[row] += 1 + end + + n_supported_rows = 0 + + @inbounds for row in 1:compressed_rows + if row_counts[row] != 0 + n_supported_rows += 1 + end + end + + row_to_slot = zeros(Int, compressed_rows) + supported_rows = Vector{Int}(undef, n_supported_rows) + + slot = 0 + + @inbounds for row in 1:compressed_rows + if row_counts[row] != 0 + slot += 1 + row_to_slot[row] = slot + supported_rows[slot] = row + end + end + + row_ptr = Vector{Int}(undef, n_supported_rows + 1) + row_ptr[1] = 1 + + @inbounds for slot in 1:n_supported_rows + row_ptr[slot + 1] = row_ptr[slot] + row_counts[supported_rows[slot]] + end + + nz_col = Vector{Int}(undef, length(values)) + nz_val = Vector{T}(undef, length(values)) + next_ptr = copy(row_ptr) + + @inbounds for index in eachindex(values) + row = row_ids[index] + slot = row_to_slot[row] + position = next_ptr[slot] + next_ptr[slot] += 1 + + nz_col[position] = col_ids[index] + nz_val[position] = values[index] + end + + row_first = Vector{Int}(undef, n_supported_rows) + row_second = Vector{Int}(undef, n_supported_rows) + row_third = Vector{Int}(undef, n_supported_rows) + row_scale = Vector{T}(undef, n_supported_rows) + + @inbounds for i1 in 1:n_x_rows + for j1 in 1:i1 + for k1 in 1:j1 + row = (i1 - 1) * i1 * (i1 + 1) ÷ 6 + (j1 - 1) * j1 ÷ 2 + k1 + slot = row_to_slot[row] + + if slot != 0 + row_first[slot] = i1 + row_second[slot] = j1 + row_third[slot] = k1 + + if i1 == j1 + divisor = j1 == k1 ? 6 : 2 + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end + + row_scale[slot] = inv(T(divisor)) + end + end + end + end + + work = zeros(T, compressed_cols) + + @inbounds for row_slot in 1:n_supported_rows + fill!(work, zero(T)) + + row_start = row_ptr[row_slot] + row_stop = row_ptr[row_slot + 1] - 1 + + for nz_index in row_start:row_stop + m1_value = nz_val[nz_index] + m1_col = nz_col[nz_index] + + @simd for col in 1:compressed_cols + work[col] += m1_value * M2[m1_col, col] + end + end + + i1 = row_first[row_slot] + j1 = row_second[row_slot] + k1 = row_third[row_slot] + scale = row_scale[row_slot] + + col = 0 + + for i2 in 1:n_x_cols + for j2 in 1:i2 + for k2 in 1:j2 + col += 1 + + g = work[col] + + abs(g) <= tol_T && continue + + g_d = g * scale + + aii = Xd[i1, i2] + aij = Xd[i1, j2] + aik = Xd[i1, k2] + + aji = Xd[j1, i2] + ajj = Xd[j1, j2] + ajk = Xd[j1, k2] + + aki = Xd[k1, i2] + akj = Xd[k1, j2] + akk = Xd[k1, k2] - 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 + + ∂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 + + return ∂X end # ===================================================================================== @@ -8103,7 +8480,8 @@ function rrule(::typeof(calculate_third_order_solution), # ∂B_from_sylv = 𝐒₃_stable' * A' * ∂C_adj — reuse ∂𝐗₃ as temp ℒ.mul!(∂𝐗₃, At, ∂C_adj) ℒ.mul!(∂B_from_sylv, 𝐒₃_stable', ∂𝐗₃) - # ∂B_from_sylv = sparse(𝐒₃_stable' * ∂𝐗₃) + + ∂B_from_sylv = choose_matrix_format(∂B_from_sylv) # ∂𝐗₃ = spinv' * ∂C_adj (overwrite temp with real value) # ℒ.mul!(∂𝐗₃, sxpinv', ∂C_adj) ∂𝐗₃ = choose_matrix_format(spinv' * ∂C_adj) diff --git a/test/check_third_order_solution_grad.jl b/test/check_third_order_solution_grad.jl new file mode 100644 index 000000000..03437c0da --- /dev/null +++ b/test/check_third_order_solution_grad.jl @@ -0,0 +1,382 @@ +using MacroModelling +import DifferentiationInterface +import ADTypes +import Mooncake +import LinearAlgebra as ℒ +using BenchmarkTools +using Random +include("../models/Smets_Wouters_2007.jl") + +# using MacroModelling + +@model Smets_Wouters_2007_SV_EPZ begin + # Epstein Zin + # U[0] = exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) * ((c[0] - c[-1] * chabb / cgamma) / (1 - csigma)) ^ (1 - csigma) + + U[0] = (1 - cbetabar) * exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) / exp((csigma - 1) / (1 + csigl) * (lab[ss] * (curvW + wdot[ss]) / (1 + curvW)) ^ (1 + csigl)) * ((c[0] - c[-1] * chabb / cgamma) / ((c[ss] - c[ss] * chabb / cgamma))) ^ (1 - csigma) + + + # (V[0]) = 1 - cbetabar + cbetabar * ((V[1]) ^ (1 - αᴱ)) ^ (1 / (1 - αᴱ)) + + # (V[0]) = U[0] + cbetabar * V[1] + + (V[0]) = U[0] + cbetabar * ((V[1]) ^ (1 - αᴱ)) ^ (1 / (1 - αᴱ)) + + # V[0] = U[0] - cbetabar * ((-V[1]) ^ (1 - αᴱ)) ^ (1 / (1 - αᴱ)) + + + # E_t_V_tp1_1_minus_alpha[0] = V[1] ^ (1 - αᴱ) + + E_t_V_tp1_1_minus_sigma[0] = (V[1]) ^ (1 - csigma) + + # E_t_V_tp1_1_minus_alpha[0] = (-V[1]) ^ (1 - αᴱ) + + + # M[0] = cbetabar * (exp(V[0]) / E_t_V_tp1_1_minus_alpha[-1] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) + + M[0] = cbetabar * ((V[1]) / E_t_V_tp1_1_minus_sigma[0] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) + + # M[0] = cbetabar * ((V[1]) / E_t_V_tp1_1_minus_alpha[0] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) + + # M[0] = cbetabar * (-V[1] / E_t_V_tp1_1_minus_alpha[0] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) + + + # Stochastic volatility + log(z_ea[0]) = (1 - rho_z_ea) * log(z̄_ea) + rho_z_ea * log(z_ea[-1]) + z_z_ea * e_z_ea[x] + log(z_eb[0]) = (1 - rho_z_eb) * log(z̄_eb) + rho_z_eb * log(z_eb[-1]) + z_z_eb * e_z_eb[x] + log(z_eg[0]) = (1 - rho_z_eg) * log(z̄_eg) + rho_z_eg * log(z_eg[-1]) + z_z_eg * e_z_eg[x] + log(z_em[0]) = (1 - rho_z_em) * log(z̄_em) + rho_z_em * log(z_em[-1]) + z_z_em * e_z_em[x] + log(z_ew[0]) = (1 - rho_z_ew) * log(z̄_ew) + rho_z_ew * log(z_ew[-1]) + z_z_ew * e_z_ew[x] + log(z_eqs[0]) = (1 - rho_z_eqs) * log(z̄_eqs) + rho_z_eqs * log(z_eqs[-1]) + z_z_eqs * e_z_eqs[x] + log(z_epinf[0]) = (1 - rho_z_epinf) * log(z̄_epinf) + rho_z_epinf * log(z_epinf[-1]) + z_z_epinf * e_z_epinf[x] + + + + y[0] = c[0] + inve[0] + y[ss] * gy[0] + afunc[0] * kp[-1] / cgamma + + y[0] * (pdot[0] + curvP) / (1 + curvP) = a[0] * k[0] ^ calfa * lab[0] ^ (1 - calfa) - (cfc - 1) * y[ss] + + k[0] = kp[-1] * zcap[0] / cgamma + + kp[0] = inve[0] * qs[0] * (1 - Sfunc[0]) + kp[-1] * (1 - ctou) / cgamma + + pdot[0] = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + pdot[-1] * cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + + wdot[0] = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) + wdot[-1] * cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) + + 1 = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + + 1 = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) + cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) + + 1 = dp[0] * (1 + pdotl[0] * curvP) / (1 + curvP) + + w[0] = dw[0] * (1 + curvW * wdotl[0]) / (1 + curvW) + + pdotl[0] = (1 - cprobp) * Pratio[0] / dp[0] + cprobp * dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0] * pdotl[-1] + + wdotl[0] = (1 - cprobw) * wnew[0] / dw[0] + cprobw * dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0] * wdotl[-1] + + xi[0] = exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) * (c[0] - c[-1] * chabb / cgamma) ^ (-csigma) + + 1 = qs[0] * pk[0] * (1 - Sfunc[0] - cgamma * inve[0] * SfuncD[0] / inve[-1]) + SfuncD[1] * xi[1] / xi[0] * qsaux[0] * pk[1] * (cgamma * inve[1] / inve[0]) ^ 2 * M[0] + + xi[0] = xi[1] * b[0] * r[0] * M[0] / pinf[1] + + rk[0] = afuncD[0] + + pk[0] = (rk[1] * zcap[1] - afunc[1] + (1 - ctou) * pk[1]) * xi[1] * M[0] / xi[0] + + k[0] = lab[0] * w[0] * calfa / (1 - calfa) / rk[0] + + mc[0] = w[0] ^ (1 - calfa) * rk[0] ^ calfa / (a[0] * calfa ^ calfa * (1 - calfa) ^ (1 - calfa)) + + wnew[0] * gamw1[0] * (1 + curvw * (1 - clandaw)) / (1 + curvW) = clandaw * gamw2[0] + gamw3[0] * curvW * (clandaw - 1) / (1 + curvW) * wnew[0] ^ (1 + clandaw * (1 + curvW) / (clandaw - 1)) + + gamw1[0] = lab[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) + gamw1[1] * (cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * M[0] + + gamw2[0] = (c[0] - c[-1] * chabb / cgamma) * lab[0] * sw[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ csigl + gamw2[1] * (cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * M[0] + + gamw3[0] = lab[0] + gamw3[1] * cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1] * xi[1] / xi[0] * cgamma * cprobw * M[0] + + Pratio[0] * gam1[0] * (1 + curvp * (1 - cfc)) / (1 + curvP) = cfc * gam2[0] + gam3[0] * (cfc - 1) * curvP / (1 + curvP) * Pratio[0] ^ (1 + cfc * (1 + curvP) / (cfc - 1)) + + gam1[0] = y[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam1[1] * xi[1] / xi[0] * cgamma * cprobp * M[0] * (cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + + gam2[0] = y[0] * mc[0] * spinf[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam2[1] * xi[1] / xi[0] * cgamma * cprobp * M[0] * (cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + + gam3[0] = y[0] + gam3[1] * cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1] * xi[1] / xi[0] * cgamma * cprobp * M[0] + + qsaux[0] = qs[1] + + # r[0] = max(1.00025,r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / pinfss) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * (y[0] / yflex[0] / (y[-1] / yflex[-1])) ^ crdy * ms[0]) + + 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] + + afunc[0] = rk[ss] * 1 / cZcap * (exp(cZcap * (zcap[0] - 1)) - 1) + + afuncD[0] = rk[ss] * exp(cZcap * (zcap[0] - 1)) + + Sfunc[0] = csadjcost / 2 * (cgamma * inve[0] / inve[-1] - cgamma) ^ 2 + + SfuncD[0] = csadjcost * (cgamma * inve[0] / inve[-1] - cgamma) + + a[0] = 1 - crhoa + crhoa * a[-1] + z_ea[0] / 100 * ea[x] + + b[0] = 1 - crhob + crhob * b[-1] + z_eb[0] / 100 * SCALE1_eb * eb[x] + + gy[0] - cg = crhog * (gy[-1] - cg) + z_eg[0] / 100 * eg[x] + z_ea[0] / 100 * ea[x] * cgy + + qs[0] = 1 - crhoqs + crhoqs * qs[-1] + z_eqs[0] / 100 * SCALE1_eqs * eqs[x] + + ms[0] = 1 - crhoms + crhoms * ms[-1] + z_em[0] / 100 * em[x] + + spinf[0] = 1 - crhopinf + crhopinf * spinf[-1] + epinfma[0] - cmap * epinfma[-1] + + epinfma[0] = z_epinf[0] / 100 * SCALE1_epinf * epinf[x] + + sw[0] = 1 - crhow + crhow * sw[-1] + ewma[0] - cmaw * ewma[-1] + + ewma[0] = z_ew[0] / 100 * SCALE1_ew * ew[x] + + yflex[0] = cflex[0] + inveflex[0] + gy[0] * yflex[ss] + afuncflex[0] * kpflex[-1] / cgamma + + yflex[0] = a[0] * kflex[0] ^ calfa * labflex[0] ^ (1 - calfa) - (cfc - 1) * yflex[ss] + + kflex[0] = kpflex[-1] * zcapflex[0] / cgamma + + kpflex[0] = inveflex[0] * qs[0] * (1 - Sfuncflex[0]) + kpflex[-1] * (1 - ctou) / cgamma + + xiflex[0] = exp((csigma - 1) / (1 + csigl) * labflex[0] ^ (1 + csigl)) * (cflex[0] - cflex[-1] * chabb / cgamma) ^ (-csigma) + + 1 = qs[0] * pkflex[0] * (1 - Sfuncflex[0] - cgamma * inveflex[0] * SfuncDflex[0] / inveflex[-1]) + SfuncDflex[1] * qsaux[0] * xiflex[1] / xiflex[0] * pkflex[1] * (cgamma * inveflex[1] / inveflex[0]) ^ 2 * M[0] + + xiflex[0] = xiflex[1] * b[0] * rrflex[0] * M[0] + + rkflex[0] = afuncDflex[0] + + pkflex[0] = (rkflex[1] * zcapflex[1] - afuncflex[1] + (1 - ctou) * pkflex[1]) * xiflex[1] * M[0] / xiflex[0] + + kflex[0] = labflex[0] * calfa / (1 - calfa) * wflex[0] / rkflex[0] + + mcflex = wflex[0] ^ (1 - calfa) * rkflex[0] ^ calfa / (a[0] * calfa ^ calfa * (1 - calfa) ^ (1 - calfa)) + + wflex[0] * (1 + curvw * (1 - clandaw)) / (1 + curvW) = sw[ss] * (labflex[0] ^ csigl * clandaw * (cflex[0] - cflex[-1] * chabb / cgamma) + wflex[0] * curvW * (clandaw - 1) / (1 + curvW)) + + # (1 + curvp * (1 - cfc)) / (1 + curvP) = spinf[ss] * cfc * mcflex + spinf[ss] * (cfc - 1) * curvP / (1 + curvP) + + afuncflex[0] = rkflex[ss] * 1 / cZcap * (exp(cZcap * (zcapflex[0] - 1)) - 1) + + afuncDflex[0] = rkflex[ss] * exp(cZcap * (zcapflex[0] - 1)) + + Sfuncflex[0] = csadjcost / 2 * (cgamma * inveflex[0] / inveflex[-1] - cgamma) ^ 2 + + SfuncDflex[0] = csadjcost * (cgamma * inveflex[0] / inveflex[-1] - cgamma) + + ygap[0] = 100 * log(y[0] / yflex[0]) + + dy[0] = ctrend + 100 * (y[0] / y[-1] - 1) + + dc[0] = ctrend + 100 * (c[0] / c[-1] - 1) + + dinve[0] = ctrend + 100 * (inve[0] / inve[-1] - 1) + + pinfobs[0] = 100 * (pinf[0] - 1) + + robs[0] = 100 * (r[0] - 1) + + dwobs[0] = ctrend + 100 * (w[0] / w[-1] - 1) + + labobs[0] = constelab + 100 * (lab[0] / lab[ss] - 1) + + 1 + spreadobs[0] / 40000 = b[0] + constespread + z_spreadobs * e_spreadobs[x] +end + +@parameters Smets_Wouters_2007_SV_EPZ begin + # Epstein-Zin + αᴱ = 0 + + #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) + + # Stochastic volatility + ## steady state volatility + z̄_ea = 0.4618 + z̄_eb = 1.8513 + z̄_eg = 0.6090 + z̄_em = 0.2397 + z̄_ew = 0.2089 + z̄_eqs = 0.6017 + z̄_epinf = 0.1455 + + #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) + + ## stochastic volatility shock size + z_z_ea = 1 + z_z_eb = 1 + z_z_eg = 1 + z_z_em = 1 + z_z_ew = 1 + z_z_eqs = 1 + z_z_epinf = 1 + + #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) + + ## stochastic volatility AR(1) coefficient + rho_z_ea = 0.75 + rho_z_eb = 0.75 + rho_z_eg = 0.75 + rho_z_em = 0.75 + rho_z_ew = 0.75 + rho_z_eqs = 0.75 + rho_z_epinf = 0.75 + + #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) + + SCALE1_eb = -((1 - chabb / cgamma) / (csigma * (1 + chabb / cgamma))) ^ (-1) + + SCALE1_eqs = (cgamma ^ 2 * csadjcost) * (1 + cbeta * cgamma ^ (1 - csigma)) + + SCALE1_epinf = 1 / ((1 / (1 + cbetabar * cgamma * cindp)) * ((1 - cprobp) * (1 - cbetabar * cgamma * cprobp) / cprobp) / ((cfc - 1) * curvp + 1)) + + SCALE1_ew = 1 / ((1 - cprobw) * (1 - cbetabar * cgamma * cprobw) / ((1 + cbetabar * cgamma) * cprobw) * (1 / ((clandaw - 1) * curvw + 1))) + + cgamma = 1 + ctrend / 100 # gross growth rate + + cbeta = 1 / (1 + constebeta / 100) # discount factor + + cZcap = czcap / (1 - czcap) + + curvP = curvp * (1 - clandap) / clandap + + curvW = curvw * (1 - clandaw) / clandaw + + clandap = cfc # fixed cost share/gross price markup + + cbetabar= cbeta * cgamma ^ (-csigma) # growth-adjusted discount factor in Euler equation + + mcflex = mc[ss] | mcflex + + pinf[ss] = 1 + constepinf / 100 | cpie + + ctou = .025 + + clandaw = 1.5 + + cg = 0.18 + + curvp = 10 + + curvw = 10 + + calfa = .24 + + csigma = 1.5 + + cfc = 1.5 + + cgy = 0.51 + + csadjcost = 6.0144 + + chabb = 0.6361 + + cprobw = 0.8087 + + csigl = 1.9423 + + cprobp = 0.6 + + cindw = 0.3243 + + cindp = 0.47 + + czcap = 0.2696 + + crpi = 1.488 + + crr = 0.8762 + + cry = 0.0593 + + crdy = 0.2347 + + crhoa = 0.9977 + + crhob = 0.5799 + + crhog = 0.9957 + + crhoqs = 0.7165 + + crhoms = 0 + + crhopinf = 0 + + crhow = 0 + + cmap = 0 + + cmaw = 0 + + constelab = 0 + + constepinf = 0.7 + + constebeta = 0.7420 + + ctrend = 0.3982 + + constespread = 0 + + z_spreadobs = 0.02 + # z_ea = 0.4618 # technology shock + + # z_eb = 1.8513 # risk-premium shock + + # z_eg = 0.6090 # government shock + + # z_em = 0.2397 # interest rate shock + + # z_ew = 0.2089 # wage mark-up shock + + # z_eqs = 0.6017 # investment-specific shock + + # z_epinf = 0.1455 # price mark-up shock + + 1e-6 > ygap > -1e-6 +end + +# ── Objective: norm of the 3rd order solution matrix ─────────────────────── +# get_solution(model, parameters; algorithm = :third_order) returns +# (SS, [S1, S2, S3], success). We take ‖S3‖_F and differentiate wrt parameters. + +function third_order_norm(params) + SS, sol_matrices, success = get_solution( + Smets_Wouters_2007_SV_EPZ, params; + algorithm = :third_order, + verbose = true, + ) + if !success + return Inf + end + S3 = sol_matrices[3] + return ℒ.norm(S3) +end + +params0 = copy(Smets_Wouters_2007_SV_EPZ.parameter_values) + +println("=== Forward pass ===") +val = third_order_norm(params0) +println("‖S3‖_F = $val") + +println("\n=== Reverse-mode gradient (Mooncake via DI) ===") +backend = ADTypes.AutoMooncake(config = nothing) +grad = DifferentiationInterface.gradient(third_order_norm, backend, params0) +@profview grad = DifferentiationInterface.gradient(third_order_norm, backend, params0) + +println("\n=== Results ===") +println("gradient = $grad") +println("norm of gradient = $(ℒ.norm(grad))") + + +Random.seed!(1234) +@benchmark grad = DifferentiationInterface.gradient(third_order_norm, backend, params0 .+ randn(length(params0)) * 1e-6) \ No newline at end of file From 11ec5652ab9ba4e60ef850209ba4041a4d752576 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 7 Jun 2026 19:35:22 +0200 Subject: [PATCH 02/37] =?UTF-8?q?Remove=20debug=20print=20statements=20fro?= =?UTF-8?q?m=20compressed=5Fkron=C2=B3=5Fpullback!=20and=20delete=20unused?= =?UTF-8?q?=20test=20file=20for=20third=20order=20solution=20gradient?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/rrules.jl | 3 - test/check_third_order_solution_grad.jl | 382 ------------------------ 2 files changed, 385 deletions(-) delete mode 100644 test/check_third_order_solution_grad.jl diff --git a/src/rrules.jl b/src/rrules.jl index 0ab6f003c..9e90075e3 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -7939,9 +7939,6 @@ function compressed_kron³_pullback!(∂X::AbstractMatrix{T}, ∂Y::DenseMatrix{T}, X::AbstractMatrix{T}; tol::Real = 0.0) where T <: Real - println("size(∂X) = ", size(∂X), " approxzeros/length = ", count(x -> abs(x) <= tol, ∂X) / length(∂X)) - println("size(∂Y) = ", size(∂Y), " approxzeros/length = ", count(x -> abs(x) <= tol, ∂Y) / length(∂Y)) - println("size(X) = ", size(X), " approxzeros/length = ", count(x -> abs(x) <= tol, X) / length(X)) 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 diff --git a/test/check_third_order_solution_grad.jl b/test/check_third_order_solution_grad.jl deleted file mode 100644 index 03437c0da..000000000 --- a/test/check_third_order_solution_grad.jl +++ /dev/null @@ -1,382 +0,0 @@ -using MacroModelling -import DifferentiationInterface -import ADTypes -import Mooncake -import LinearAlgebra as ℒ -using BenchmarkTools -using Random -include("../models/Smets_Wouters_2007.jl") - -# using MacroModelling - -@model Smets_Wouters_2007_SV_EPZ begin - # Epstein Zin - # U[0] = exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) * ((c[0] - c[-1] * chabb / cgamma) / (1 - csigma)) ^ (1 - csigma) - - U[0] = (1 - cbetabar) * exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) / exp((csigma - 1) / (1 + csigl) * (lab[ss] * (curvW + wdot[ss]) / (1 + curvW)) ^ (1 + csigl)) * ((c[0] - c[-1] * chabb / cgamma) / ((c[ss] - c[ss] * chabb / cgamma))) ^ (1 - csigma) - - - # (V[0]) = 1 - cbetabar + cbetabar * ((V[1]) ^ (1 - αᴱ)) ^ (1 / (1 - αᴱ)) - - # (V[0]) = U[0] + cbetabar * V[1] - - (V[0]) = U[0] + cbetabar * ((V[1]) ^ (1 - αᴱ)) ^ (1 / (1 - αᴱ)) - - # V[0] = U[0] - cbetabar * ((-V[1]) ^ (1 - αᴱ)) ^ (1 / (1 - αᴱ)) - - - # E_t_V_tp1_1_minus_alpha[0] = V[1] ^ (1 - αᴱ) - - E_t_V_tp1_1_minus_sigma[0] = (V[1]) ^ (1 - csigma) - - # E_t_V_tp1_1_minus_alpha[0] = (-V[1]) ^ (1 - αᴱ) - - - # M[0] = cbetabar * (exp(V[0]) / E_t_V_tp1_1_minus_alpha[-1] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) - - M[0] = cbetabar * ((V[1]) / E_t_V_tp1_1_minus_sigma[0] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) - - # M[0] = cbetabar * ((V[1]) / E_t_V_tp1_1_minus_alpha[0] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) - - # M[0] = cbetabar * (-V[1] / E_t_V_tp1_1_minus_alpha[0] ^ (1 / (1 - αᴱ)) ) ^ (-αᴱ) - - - # Stochastic volatility - log(z_ea[0]) = (1 - rho_z_ea) * log(z̄_ea) + rho_z_ea * log(z_ea[-1]) + z_z_ea * e_z_ea[x] - log(z_eb[0]) = (1 - rho_z_eb) * log(z̄_eb) + rho_z_eb * log(z_eb[-1]) + z_z_eb * e_z_eb[x] - log(z_eg[0]) = (1 - rho_z_eg) * log(z̄_eg) + rho_z_eg * log(z_eg[-1]) + z_z_eg * e_z_eg[x] - log(z_em[0]) = (1 - rho_z_em) * log(z̄_em) + rho_z_em * log(z_em[-1]) + z_z_em * e_z_em[x] - log(z_ew[0]) = (1 - rho_z_ew) * log(z̄_ew) + rho_z_ew * log(z_ew[-1]) + z_z_ew * e_z_ew[x] - log(z_eqs[0]) = (1 - rho_z_eqs) * log(z̄_eqs) + rho_z_eqs * log(z_eqs[-1]) + z_z_eqs * e_z_eqs[x] - log(z_epinf[0]) = (1 - rho_z_epinf) * log(z̄_epinf) + rho_z_epinf * log(z_epinf[-1]) + z_z_epinf * e_z_epinf[x] - - - - y[0] = c[0] + inve[0] + y[ss] * gy[0] + afunc[0] * kp[-1] / cgamma - - y[0] * (pdot[0] + curvP) / (1 + curvP) = a[0] * k[0] ^ calfa * lab[0] ^ (1 - calfa) - (cfc - 1) * y[ss] - - k[0] = kp[-1] * zcap[0] / cgamma - - kp[0] = inve[0] * qs[0] * (1 - Sfunc[0]) + kp[-1] * (1 - ctou) / cgamma - - pdot[0] = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) + pdot[-1] * cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) - - wdot[0] = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) + wdot[-1] * cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) - - 1 = (1 - cprobp) * (Pratio[0] / dp[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) + cprobp * (dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) - - 1 = (1 - cprobw) * (wnew[0] / dw[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) + cprobw * (dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) - - 1 = dp[0] * (1 + pdotl[0] * curvP) / (1 + curvP) - - w[0] = dw[0] * (1 + curvW * wdotl[0]) / (1 + curvW) - - pdotl[0] = (1 - cprobp) * Pratio[0] / dp[0] + cprobp * dp[-1] / dp[0] * pinf[-1] ^ cindp * cpie ^ (1 - cindp) / pinf[0] * pdotl[-1] - - wdotl[0] = (1 - cprobw) * wnew[0] / dw[0] + cprobw * dw[-1] / dw[0] * pinf[-1] ^ cindw * cpie ^ (1 - cindw) / pinf[0] * wdotl[-1] - - xi[0] = exp((csigma - 1) / (1 + csigl) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ (1 + csigl)) * (c[0] - c[-1] * chabb / cgamma) ^ (-csigma) - - 1 = qs[0] * pk[0] * (1 - Sfunc[0] - cgamma * inve[0] * SfuncD[0] / inve[-1]) + SfuncD[1] * xi[1] / xi[0] * qsaux[0] * pk[1] * (cgamma * inve[1] / inve[0]) ^ 2 * M[0] - - xi[0] = xi[1] * b[0] * r[0] * M[0] / pinf[1] - - rk[0] = afuncD[0] - - pk[0] = (rk[1] * zcap[1] - afunc[1] + (1 - ctou) * pk[1]) * xi[1] * M[0] / xi[0] - - k[0] = lab[0] * w[0] * calfa / (1 - calfa) / rk[0] - - mc[0] = w[0] ^ (1 - calfa) * rk[0] ^ calfa / (a[0] * calfa ^ calfa * (1 - calfa) ^ (1 - calfa)) - - wnew[0] * gamw1[0] * (1 + curvw * (1 - clandaw)) / (1 + curvW) = clandaw * gamw2[0] + gamw3[0] * curvW * (clandaw - 1) / (1 + curvW) * wnew[0] ^ (1 + clandaw * (1 + curvW) / (clandaw - 1)) - - gamw1[0] = lab[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) + gamw1[1] * (cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - (1 + curvw * (1 - clandaw))) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * M[0] - - gamw2[0] = (c[0] - c[-1] * chabb / cgamma) * lab[0] * sw[0] * dw[0] ^ (clandaw * (1 + curvW) / (clandaw - 1)) * (lab[0] * (curvW + wdot[0]) / (1 + curvW)) ^ csigl + gamw2[1] * (cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1]) ^ (( - clandaw) * (1 + curvW) / (clandaw - 1)) * xi[1] / xi[0] * cgamma * cprobw * M[0] - - gamw3[0] = lab[0] + gamw3[1] * cpie ^ (1 - cindw) * pinf[0] ^ cindw / pinf[1] * xi[1] / xi[0] * cgamma * cprobw * M[0] - - Pratio[0] * gam1[0] * (1 + curvp * (1 - cfc)) / (1 + curvP) = cfc * gam2[0] + gam3[0] * (cfc - 1) * curvP / (1 + curvP) * Pratio[0] ^ (1 + cfc * (1 + curvP) / (cfc - 1)) - - gam1[0] = y[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam1[1] * xi[1] / xi[0] * cgamma * cprobp * M[0] * (cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - (1 + curvp * (1 - cfc))) / (cfc - 1)) - - gam2[0] = y[0] * mc[0] * spinf[0] * dp[0] ^ (cfc * (1 + curvP) / (cfc - 1)) + gam2[1] * xi[1] / xi[0] * cgamma * cprobp * M[0] * (cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1]) ^ (( - cfc) * (1 + curvP) / (cfc - 1)) - - gam3[0] = y[0] + gam3[1] * cpie ^ (1 - cindp) * pinf[0] ^ cindp / pinf[1] * xi[1] / xi[0] * cgamma * cprobp * M[0] - - qsaux[0] = qs[1] - - # r[0] = max(1.00025,r[ss] ^ (1 - crr) * r[-1] ^ crr * (pinf[0] / pinfss) ^ ((1 - crr) * crpi) * (y[0] / yflex[0]) ^ ((1 - crr) * cry) * (y[0] / yflex[0] / (y[-1] / yflex[-1])) ^ crdy * ms[0]) - - 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] - - afunc[0] = rk[ss] * 1 / cZcap * (exp(cZcap * (zcap[0] - 1)) - 1) - - afuncD[0] = rk[ss] * exp(cZcap * (zcap[0] - 1)) - - Sfunc[0] = csadjcost / 2 * (cgamma * inve[0] / inve[-1] - cgamma) ^ 2 - - SfuncD[0] = csadjcost * (cgamma * inve[0] / inve[-1] - cgamma) - - a[0] = 1 - crhoa + crhoa * a[-1] + z_ea[0] / 100 * ea[x] - - b[0] = 1 - crhob + crhob * b[-1] + z_eb[0] / 100 * SCALE1_eb * eb[x] - - gy[0] - cg = crhog * (gy[-1] - cg) + z_eg[0] / 100 * eg[x] + z_ea[0] / 100 * ea[x] * cgy - - qs[0] = 1 - crhoqs + crhoqs * qs[-1] + z_eqs[0] / 100 * SCALE1_eqs * eqs[x] - - ms[0] = 1 - crhoms + crhoms * ms[-1] + z_em[0] / 100 * em[x] - - spinf[0] = 1 - crhopinf + crhopinf * spinf[-1] + epinfma[0] - cmap * epinfma[-1] - - epinfma[0] = z_epinf[0] / 100 * SCALE1_epinf * epinf[x] - - sw[0] = 1 - crhow + crhow * sw[-1] + ewma[0] - cmaw * ewma[-1] - - ewma[0] = z_ew[0] / 100 * SCALE1_ew * ew[x] - - yflex[0] = cflex[0] + inveflex[0] + gy[0] * yflex[ss] + afuncflex[0] * kpflex[-1] / cgamma - - yflex[0] = a[0] * kflex[0] ^ calfa * labflex[0] ^ (1 - calfa) - (cfc - 1) * yflex[ss] - - kflex[0] = kpflex[-1] * zcapflex[0] / cgamma - - kpflex[0] = inveflex[0] * qs[0] * (1 - Sfuncflex[0]) + kpflex[-1] * (1 - ctou) / cgamma - - xiflex[0] = exp((csigma - 1) / (1 + csigl) * labflex[0] ^ (1 + csigl)) * (cflex[0] - cflex[-1] * chabb / cgamma) ^ (-csigma) - - 1 = qs[0] * pkflex[0] * (1 - Sfuncflex[0] - cgamma * inveflex[0] * SfuncDflex[0] / inveflex[-1]) + SfuncDflex[1] * qsaux[0] * xiflex[1] / xiflex[0] * pkflex[1] * (cgamma * inveflex[1] / inveflex[0]) ^ 2 * M[0] - - xiflex[0] = xiflex[1] * b[0] * rrflex[0] * M[0] - - rkflex[0] = afuncDflex[0] - - pkflex[0] = (rkflex[1] * zcapflex[1] - afuncflex[1] + (1 - ctou) * pkflex[1]) * xiflex[1] * M[0] / xiflex[0] - - kflex[0] = labflex[0] * calfa / (1 - calfa) * wflex[0] / rkflex[0] - - mcflex = wflex[0] ^ (1 - calfa) * rkflex[0] ^ calfa / (a[0] * calfa ^ calfa * (1 - calfa) ^ (1 - calfa)) - - wflex[0] * (1 + curvw * (1 - clandaw)) / (1 + curvW) = sw[ss] * (labflex[0] ^ csigl * clandaw * (cflex[0] - cflex[-1] * chabb / cgamma) + wflex[0] * curvW * (clandaw - 1) / (1 + curvW)) - - # (1 + curvp * (1 - cfc)) / (1 + curvP) = spinf[ss] * cfc * mcflex + spinf[ss] * (cfc - 1) * curvP / (1 + curvP) - - afuncflex[0] = rkflex[ss] * 1 / cZcap * (exp(cZcap * (zcapflex[0] - 1)) - 1) - - afuncDflex[0] = rkflex[ss] * exp(cZcap * (zcapflex[0] - 1)) - - Sfuncflex[0] = csadjcost / 2 * (cgamma * inveflex[0] / inveflex[-1] - cgamma) ^ 2 - - SfuncDflex[0] = csadjcost * (cgamma * inveflex[0] / inveflex[-1] - cgamma) - - ygap[0] = 100 * log(y[0] / yflex[0]) - - dy[0] = ctrend + 100 * (y[0] / y[-1] - 1) - - dc[0] = ctrend + 100 * (c[0] / c[-1] - 1) - - dinve[0] = ctrend + 100 * (inve[0] / inve[-1] - 1) - - pinfobs[0] = 100 * (pinf[0] - 1) - - robs[0] = 100 * (r[0] - 1) - - dwobs[0] = ctrend + 100 * (w[0] / w[-1] - 1) - - labobs[0] = constelab + 100 * (lab[0] / lab[ss] - 1) - - 1 + spreadobs[0] / 40000 = b[0] + constespread + z_spreadobs * e_spreadobs[x] -end - -@parameters Smets_Wouters_2007_SV_EPZ begin - # Epstein-Zin - αᴱ = 0 - - #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) - - # Stochastic volatility - ## steady state volatility - z̄_ea = 0.4618 - z̄_eb = 1.8513 - z̄_eg = 0.6090 - z̄_em = 0.2397 - z̄_ew = 0.2089 - z̄_eqs = 0.6017 - z̄_epinf = 0.1455 - - #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) - - ## stochastic volatility shock size - z_z_ea = 1 - z_z_eb = 1 - z_z_eg = 1 - z_z_em = 1 - z_z_ew = 1 - z_z_eqs = 1 - z_z_epinf = 1 - - #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) - - ## stochastic volatility AR(1) coefficient - rho_z_ea = 0.75 - rho_z_eb = 0.75 - rho_z_eg = 0.75 - rho_z_em = 0.75 - rho_z_ew = 0.75 - rho_z_eqs = 0.75 - rho_z_epinf = 0.75 - - #### do NOT change parameters here in the parameter declaration - everything needs to be recompiled if you change it here (slow) - change it in the function calls instead (fast - see below) - - SCALE1_eb = -((1 - chabb / cgamma) / (csigma * (1 + chabb / cgamma))) ^ (-1) - - SCALE1_eqs = (cgamma ^ 2 * csadjcost) * (1 + cbeta * cgamma ^ (1 - csigma)) - - SCALE1_epinf = 1 / ((1 / (1 + cbetabar * cgamma * cindp)) * ((1 - cprobp) * (1 - cbetabar * cgamma * cprobp) / cprobp) / ((cfc - 1) * curvp + 1)) - - SCALE1_ew = 1 / ((1 - cprobw) * (1 - cbetabar * cgamma * cprobw) / ((1 + cbetabar * cgamma) * cprobw) * (1 / ((clandaw - 1) * curvw + 1))) - - cgamma = 1 + ctrend / 100 # gross growth rate - - cbeta = 1 / (1 + constebeta / 100) # discount factor - - cZcap = czcap / (1 - czcap) - - curvP = curvp * (1 - clandap) / clandap - - curvW = curvw * (1 - clandaw) / clandaw - - clandap = cfc # fixed cost share/gross price markup - - cbetabar= cbeta * cgamma ^ (-csigma) # growth-adjusted discount factor in Euler equation - - mcflex = mc[ss] | mcflex - - pinf[ss] = 1 + constepinf / 100 | cpie - - ctou = .025 - - clandaw = 1.5 - - cg = 0.18 - - curvp = 10 - - curvw = 10 - - calfa = .24 - - csigma = 1.5 - - cfc = 1.5 - - cgy = 0.51 - - csadjcost = 6.0144 - - chabb = 0.6361 - - cprobw = 0.8087 - - csigl = 1.9423 - - cprobp = 0.6 - - cindw = 0.3243 - - cindp = 0.47 - - czcap = 0.2696 - - crpi = 1.488 - - crr = 0.8762 - - cry = 0.0593 - - crdy = 0.2347 - - crhoa = 0.9977 - - crhob = 0.5799 - - crhog = 0.9957 - - crhoqs = 0.7165 - - crhoms = 0 - - crhopinf = 0 - - crhow = 0 - - cmap = 0 - - cmaw = 0 - - constelab = 0 - - constepinf = 0.7 - - constebeta = 0.7420 - - ctrend = 0.3982 - - constespread = 0 - - z_spreadobs = 0.02 - # z_ea = 0.4618 # technology shock - - # z_eb = 1.8513 # risk-premium shock - - # z_eg = 0.6090 # government shock - - # z_em = 0.2397 # interest rate shock - - # z_ew = 0.2089 # wage mark-up shock - - # z_eqs = 0.6017 # investment-specific shock - - # z_epinf = 0.1455 # price mark-up shock - - 1e-6 > ygap > -1e-6 -end - -# ── Objective: norm of the 3rd order solution matrix ─────────────────────── -# get_solution(model, parameters; algorithm = :third_order) returns -# (SS, [S1, S2, S3], success). We take ‖S3‖_F and differentiate wrt parameters. - -function third_order_norm(params) - SS, sol_matrices, success = get_solution( - Smets_Wouters_2007_SV_EPZ, params; - algorithm = :third_order, - verbose = true, - ) - if !success - return Inf - end - S3 = sol_matrices[3] - return ℒ.norm(S3) -end - -params0 = copy(Smets_Wouters_2007_SV_EPZ.parameter_values) - -println("=== Forward pass ===") -val = third_order_norm(params0) -println("‖S3‖_F = $val") - -println("\n=== Reverse-mode gradient (Mooncake via DI) ===") -backend = ADTypes.AutoMooncake(config = nothing) -grad = DifferentiationInterface.gradient(third_order_norm, backend, params0) -@profview grad = DifferentiationInterface.gradient(third_order_norm, backend, params0) - -println("\n=== Results ===") -println("gradient = $grad") -println("norm of gradient = $(ℒ.norm(grad))") - - -Random.seed!(1234) -@benchmark grad = DifferentiationInterface.gradient(third_order_norm, backend, params0 .+ randn(length(params0)) * 1e-6) \ No newline at end of file From ed74538a82e125f2550230c02d6bd57b38c151ef Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Sun, 7 Jun 2026 20:12:34 +0200 Subject: [PATCH 03/37] undo BnechmarkTools in toml --- Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/Project.toml b/Project.toml index 8664e26d9..3d5afe3ba 100644 --- a/Project.toml +++ b/Project.toml @@ -6,7 +6,6 @@ authors = ["Thore Kockerols "] [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" -BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" BlockTriangularForm = "adeb47b7-70bf-415a-bb24-c358563e873a" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" From bf823f5f405cc5b8c5e5700da4a3a6b7c52e26ba Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 7 Jun 2026 22:25:11 +0000 Subject: [PATCH 04/37] Add disambiguation method for sparse-sparse mul_fill_kron_adjoint! Co-authored-by: thorek1 <13523097+thorek1@users.noreply.github.com> --- src/rrules.jl | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/src/rrules.jl b/src/rrules.jl index 9e90075e3..de9c48f31 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -6896,6 +6896,39 @@ function mul_fill_kron_adjoint!( end +# Disambiguation for mixed-sparsity methods when both A and B are sparse +function mul_fill_kron_adjoint!( + ∂A::AbstractMatrix{R}, + ∂B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, + A::SparseMatrixCSC{TA, Ti}, + 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, Ti <: Integer} + return invoke( + mul_fill_kron_adjoint!, + Tuple{ + AbstractMatrix{R}, + AbstractMatrix{R}, + AbstractMatrix, + AbstractMatrix, + SparseMatrixCSC{TA, Ti}, + AbstractMatrix{TB}, + }, + ∂A, + ∂B, + M1, + M2, + A, + B; + tol = tol, + block = block, + ) +end + + # Mixed-sparsity variant: A is dense, B is sparse function mul_fill_kron_adjoint!(∂A::AbstractMatrix{R}, ∂B::AbstractMatrix{R}, From defcf0006f83c8cd5aa64c0b7060c0e8eb23361c Mon Sep 17 00:00:00 2001 From: thorek1 Date: Mon, 8 Jun 2026 18:09:32 +0200 Subject: [PATCH 05/37] Refactor mul_fill_kron_adjoint! to use DenseMatrix for A and B parameters --- src/rrules.jl | 42 ++++-------------------------------------- 1 file changed, 4 insertions(+), 38 deletions(-) diff --git a/src/rrules.jl b/src/rrules.jl index de9c48f31..e0e101fe7 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -6563,8 +6563,8 @@ function mul_fill_kron_adjoint!(∂A::AbstractMatrix{R}, ∂B::AbstractMatrix{R}, M1::AbstractMatrix, M2::AbstractMatrix, - A::AbstractMatrix{TA}, - B::AbstractMatrix{TB}; + A::DenseMatrix{TA}, + B::DenseMatrix{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) @@ -6702,7 +6702,7 @@ function mul_fill_kron_adjoint!( M1::AbstractMatrix, M2::AbstractMatrix, A::SparseMatrixCSC{TA, Ti}, - B::AbstractMatrix{TB}; + B::DenseMatrix{TB}; tol::Real = 0.0, block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0), ) where {R <: Real, TA <: Real, TB <: Real, Ti <: Integer} @@ -6895,46 +6895,12 @@ function mul_fill_kron_adjoint!( return ∂A, ∂B end - -# Disambiguation for mixed-sparsity methods when both A and B are sparse -function mul_fill_kron_adjoint!( - ∂A::AbstractMatrix{R}, - ∂B::AbstractMatrix{R}, - M1::AbstractMatrix, - M2::AbstractMatrix, - A::SparseMatrixCSC{TA, Ti}, - 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, Ti <: Integer} - return invoke( - mul_fill_kron_adjoint!, - Tuple{ - AbstractMatrix{R}, - AbstractMatrix{R}, - AbstractMatrix, - AbstractMatrix, - SparseMatrixCSC{TA, Ti}, - AbstractMatrix{TB}, - }, - ∂A, - ∂B, - M1, - M2, - A, - B; - tol = tol, - block = block, - ) -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}, + A::DenseMatrix{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} From 86fff9f201c3d40f698a40cbeb455c93b51c3795 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Mon, 8 Jun 2026 19:58:13 +0200 Subject: [PATCH 06/37] Add functions for expanding filter-free cotangents and scattering state derivatives --- src/rrules.jl | 325 +++++++++++++++++++++++++++++++------------------- 1 file changed, 199 insertions(+), 126 deletions(-) diff --git a/src/rrules.jl b/src/rrules.jl index e0e101fe7..dd7963abc 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -18789,6 +18789,43 @@ function contract_filter_free_me_std_cotangent(d_me_std::AbstractMatrix, user_me return out end +function expand_filter_free_matrix_cotangent(d_red::AbstractMatrix, + needed::AbstractVector{Int}, + n_vars_full::Int, + n_cols::Int) + d_full = zeros(eltype(d_red), n_vars_full, n_cols) + @inbounds d_full[needed, :] .= d_red + return d_full +end + +function expand_filter_free_state_cotangent(d_red::AbstractVector, + needed::AbstractVector{Int}, + n_vars_full::Int) + d_full = zeros(eltype(d_red), n_vars_full) + @inbounds d_full[needed] .= d_red + return d_full +end + +function expand_filter_free_state_cotangent(d_red::AbstractVector{<:AbstractVector}, + needed::AbstractVector{Int}, + n_vars_full::Int) + d_full = [zeros(eltype(d_red[i]), n_vars_full) for i in eachindex(d_red)] + @inbounds for i in eachindex(d_red) + d_full[i][needed] .= d_red[i] + end + return d_full +end + +function scatter_filter_free_ss_and_pars_cotangent(d_SS_obs::AbstractVector, + obs_indices::AbstractVector{Int}, + n_ss_and_pars::Int) + d_SS_and_pars = zeros(eltype(d_SS_obs), n_ss_and_pars) + @inbounds for k in eachindex(obs_indices) + d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] + end + return d_SS_and_pars +end + # Visible-period filter-free pullbacks run on the reduced `needed` row slice of # the policy matrices. They propagate the state cotangent backward across only # the retained visible sample, accumulate reduced-matrix cotangents, and return @@ -19536,35 +19573,38 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_𝐒₁_full_cot = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁) - @inbounds d_𝐒₁_full_cot[needed, :] .= d_𝐒₁_red - d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) - @inbounds for k in eachindex(obs_indices) - d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] + d_params = @thunk begin + d_𝐒₁_full_cot = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if has_override && !(d_state_red isa Union{NoTangent, AbstractZero}) + d_state_full = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full_for_init) + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state_full + end + end + ss_pb((NoTangent(), d_SS_and_pars, d_𝐒₁_full_cot, NoTangent()))[3] end - # backprop through initial_state_deviations = initial_state - SS_and_pars[1:nVars] - # AND compute d_initial_state (positional tangent for the user-facing - # `initial_state` argument). - d_initial_state = NoTangent() - if has_override && !(d_state_red isa Union{NoTangent, AbstractZero}) - d_state_full_init = zeros(eltype(d_state_red), nVars_full_for_init) - @inbounds d_state_full_init[needed] .= d_state_red - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state_full_init - d_initial_state = d_state_full_init - else - d_is_vec = Vector{Vector{eltype(d_state_red)}}(undef, n_overridden_components) - d_is_vec[1] = d_state_full_init - for k in 2:n_overridden_components - d_is_vec[k] = zeros(eltype(d_state_red), length(initial_state[k])) + d_initial_state = if has_override + @thunk begin + d_𝐒₁_full_cot = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + d_state_full_init = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full_for_init) + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state_full_init + initial_state_pullback_tangent(initial_state, d_state_full_init) + else + d_is_vec = Vector{Vector{eltype(d_state_red)}}(undef, n_overridden_components) + d_is_vec[1] = d_state_full_init + for k in 2:n_overridden_components + d_is_vec[k] = zeros(eltype(d_state_red), length(initial_state[k])) + end + initial_state_pullback_tangent(initial_state, d_is_vec) end - d_initial_state = d_is_vec end + else + NoTangent() end - # first_order ss rrule expects bare 𝐒₁ cotangent and ignores Δstate - ss_grads = ss_pb((NoTangent(), d_SS_and_pars, d_𝐒₁_full_cot, NoTangent())) - d_params = ss_grads[3] - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19621,35 +19661,45 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red - d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red - d_state = zeros(eltype(d_state_red), nVars_full); @inbounds d_state[needed] .= d_state_red - d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) - @inbounds for k in eachindex(obs_indices) - d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] - end - d_initial_state = NoTangent() - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] - d_is = zeros(eltype(d_state), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] - d_initial_state = d_is + d_params = @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] + end + d_state_for_ss = zeros(eltype(d_state), nVars_full) else - d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) - d_is_vec[1] = collect(d_state) - for k in 2:n_overridden_components - d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) + d_state_for_ss = d_state + end + ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent()))[3] + end + d_initial_state = if has_override + @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] + d_is = zeros(eltype(d_state), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] + initial_state_pullback_tangent(initial_state, d_is) + else + d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) + d_is_vec[1] = collect(d_state) + for k in 2:n_overridden_components + d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) + end + initial_state_pullback_tangent(initial_state, d_is_vec) end - d_initial_state = d_is_vec end - d_state_for_ss = zeros(eltype(d_state), nVars_full) else - d_state_for_ss = d_state + NoTangent() end - ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent())) - d_params = ss_grads[3] - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19712,35 +19762,40 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red - d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red - d_state = [zeros(eltype(d_state_red[1]), nVars_full), - zeros(eltype(d_state_red[2]), nVars_full)] - @inbounds d_state[1][needed] .= d_state_red[1] - @inbounds d_state[2][needed] .= d_state_red[2] - d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) - @inbounds for k in eachindex(obs_indices) - d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] - end - d_initial_state = NoTangent() - if has_override - if initial_state_is_levels - # Only first-order component depends on SS (higher-order initialized to zero) - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] - d_is = zeros(eltype(d_state[1]), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] - d_initial_state = d_is + d_params = @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] + end + d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] else - d_initial_state = [collect(d_state[i]) for i in 1:n_overridden_components] + d_state_for_ss = d_state + end + ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent()))[3] + end + d_initial_state = if has_override + @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] + d_is = zeros(eltype(d_state[1]), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] + initial_state_pullback_tangent(initial_state, d_is) + else + initial_state_pullback_tangent(initial_state, [collect(d_state[i]) for i in 1:n_overridden_components]) + end end - # Zero only the components the user overrode; preserve others. - d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] else - d_state_for_ss = d_state + NoTangent() end - ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent())) - d_params = ss_grads[3] - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19803,36 +19858,47 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red - d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red - d_𝐒₃ = zeros(eltype(d_𝐒₃_red), nVars_full, ncols₃); @inbounds d_𝐒₃[needed, :] .= d_𝐒₃_red - d_state = zeros(eltype(d_state_red), nVars_full); @inbounds d_state[needed] .= d_state_red - d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) - @inbounds for k in eachindex(obs_indices) - d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] - end - d_initial_state = NoTangent() - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] - d_is = zeros(eltype(d_state), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] - d_initial_state = d_is + d_params = @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] + end + d_state_for_ss = zeros(eltype(d_state), nVars_full) else - d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) - d_is_vec[1] = collect(d_state) - for k in 2:n_overridden_components - d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) + d_state_for_ss = d_state + end + ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent()))[3] + end + d_initial_state = if has_override + @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] + d_is = zeros(eltype(d_state), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] + initial_state_pullback_tangent(initial_state, d_is) + else + d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) + d_is_vec[1] = collect(d_state) + for k in 2:n_overridden_components + d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) + end + initial_state_pullback_tangent(initial_state, d_is_vec) end - d_initial_state = d_is_vec end - d_state_for_ss = zeros(eltype(d_state), nVars_full) else - d_state_for_ss = d_state + NoTangent() end - ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent())) - d_params = ss_grads[3] - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19911,37 +19977,44 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red - d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red - d_𝐒₃ = zeros(eltype(d_𝐒₃_red), nVars_full, ncols₃); @inbounds d_𝐒₃[needed, :] .= d_𝐒₃_red - d_state = [zeros(eltype(d_state_red[1]), nVars_full), - zeros(eltype(d_state_red[2]), nVars_full), - zeros(eltype(d_state_red[3]), nVars_full)] - @inbounds d_state[1][needed] .= d_state_red[1] - @inbounds d_state[2][needed] .= d_state_red[2] - @inbounds d_state[3][needed] .= d_state_red[3] - d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) - @inbounds for k in eachindex(obs_indices) - d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] - end - d_initial_state = NoTangent() - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] - d_is = zeros(eltype(d_state[1]), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] - d_initial_state = d_is - d_state_for_ss = [zeros(eltype(d_state[1]), nVars_full), d_state[2], level_override_zeroes_third_state ? zeros(eltype(d_state[3]), nVars_full) : d_state[3]] + d_params = @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] + d_state_for_ss = [zeros(eltype(d_state[1]), nVars_full), d_state[2], level_override_zeroes_third_state ? zeros(eltype(d_state[3]), nVars_full) : d_state[3]] + else + d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] + end else - d_initial_state = [collect(d_state[i]) for i in 1:n_overridden_components] - d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] + d_state_for_ss = d_state + end + ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent()))[3] + end + d_initial_state = if has_override + @thunk begin + d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) + d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) + d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) + d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) + d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] + d_is = zeros(eltype(d_state[1]), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] + initial_state_pullback_tangent(initial_state, d_is) + else + initial_state_pullback_tangent(initial_state, [collect(d_state[i]) for i in 1:n_overridden_components]) + end end else - d_state_for_ss = d_state + NoTangent() end - ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent())) - d_params = ss_grads[3] - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state end return isfinite(llh) ? (llh, pullback) : on_failure end From e1add5d90225701c6e354cb1e324c1f906f5d7de Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Tue, 9 Jun 2026 15:27:51 +0000 Subject: [PATCH 07/37] Update mul_fill_kron_adjoint! and Mooncake extension --- ext/MooncakeExt.jl | 161 +++++++++++++++++++++++++++++++++---- src/rrules.jl | 195 ++++++++++++++++++++++++--------------------- 2 files changed, 248 insertions(+), 108 deletions(-) diff --git a/ext/MooncakeExt.jl b/ext/MooncakeExt.jl index ebb08dbfa..ebfe975ff 100644 --- a/ext/MooncakeExt.jl +++ b/ext/MooncakeExt.jl @@ -71,17 +71,147 @@ Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglike Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_irf), MacroModelling.ℳ, Vector{T}} where {T<:Base.IEEEFloat} true -# get_loglikelihood: -# (𝓂, data, parameter_values::Vector{T}, shocks::Matrix{T}, me_std::T_or_Vector{T}) -# Two narrow @from_rrule generations cover scalar and vector me_std cases. -# 5-arg kwarg-only path (no AD through initial_state) -Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, T} where {T<:Base.IEEEFloat} true -Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Vector{T}} where {T<:Base.IEEEFloat} true -Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Matrix{T}} where {T<:Base.IEEEFloat} true -# 6-arg positional path (AD through initial_state) -Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, T, Vector{Float64}} where {T<:Base.IEEEFloat} true -Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Vector{T}, Vector{Float64}} where {T<:Base.IEEEFloat} true -Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Matrix{T}, Vector{Float64}} where {T<:Base.IEEEFloat} true +# ── get_loglikelihood (filter-free, with shocks/me_std): manual rrule!! ── +# Replaces @from_rrule so that CRC @thunks for d_params / d_initial_state are +# NOT forced when the caller differentiates only w.r.t. shocks / me_std. +# When an argument's fdata is inactive the corresponding CRC tangent (which may +# contain an expensive @thunk) is skipped entirely. + +mooncake_tangent_is_inactive(::Mooncake.NoTangent) = true +mooncake_tangent_is_inactive(tangent::Mooncake.PossiblyUninitTangent) = !Mooncake.is_init(tangent) +mooncake_tangent_is_inactive(_) = false + +function increment_rrule_arg_rdata!(arg_cd::CoDual, lazy_rdata, cr_tangent) + cr_tangent isa ChainRulesCore.AbstractZero && return NoRData() + arg_fdata = Mooncake.tangent(arg_cd) + if Mooncake.primal(arg_cd) isa AbstractArray && arg_fdata isa Mooncake.NoFData + return NoRData() + end + mooncake_tangent_is_inactive(arg_fdata) && return NoRData() + return Mooncake.increment_and_get_rdata!( + arg_fdata, + Mooncake.instantiate(lazy_rdata), + cr_tangent, + ) +end + +Mooncake._get_fdata_field(primal::Function, ::Mooncake.NoFData, f) = Mooncake.NoFData() +Mooncake.uninit_fcodual(x::AbstractArray{<:Base.IEEEFloat}) = CoDual(x, Mooncake.NoFData()) + +# ── 5-arg positional (no initial_state) ── +# (𝓂, data, parameter_values, shocks, me_std) +function Mooncake.rrule!!( + f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, + model_cd::CoDual{MacroModelling.ℳ}, + data_cd::CoDual{<:KeyedArray{Float64}}, + params_cd::CoDual{Vector{T}}, + shocks_cd::CoDual{Matrix{T}}, + me_std_cd::CoDual{M}, +) where {T<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} + fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_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_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) + return ntuple(length(fargs)) do i + increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) + end + end + return CoDual(y_primal, y_fdata), pb!! +end + +# ── 5-arg kwcall ── +function Mooncake.rrule!!( + kwcall_cd::CoDual{typeof(Core.kwcall)}, + kwargs_cd::CoDual{<:NamedTuple}, + f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, + model_cd::CoDual{MacroModelling.ℳ}, + data_cd::CoDual{<:KeyedArray{Float64}}, + params_cd::CoDual{Vector{T}}, + shocks_cd::CoDual{Matrix{T}}, + me_std_cd::CoDual{M}, +) where {T<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} + kwargs = Mooncake.primal(kwargs_cd) + fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_cd) + primals = map(Mooncake.primal, fargs) + lazy_rdata = map(Mooncake.lazy_zero_rdata, primals) + y_primal, cr_pb = ChainRulesCore.rrule(primals...; kwargs...) + y_fdata = Mooncake.fdata(Mooncake.zero_tangent(y_primal)) + kwargs_lazy_rdata = Mooncake.lazy_zero_rdata(kwargs) + function pb!!(y_rdata) + cr_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) + kwargs_rdata = Mooncake.increment_and_get_rdata!( + Mooncake.tangent(kwargs_cd), + Mooncake.instantiate(kwargs_lazy_rdata), + ChainRulesCore.NoTangent(), + ) + regular_rdata = ntuple(length(fargs)) do i + increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) + end + return (NoRData(), kwargs_rdata, regular_rdata...) + end + return CoDual(y_primal, y_fdata), pb!! +end + +# ── 6-arg positional (with Vector{Float64} initial_state) ── +function Mooncake.rrule!!( + f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, + model_cd::CoDual{MacroModelling.ℳ}, + data_cd::CoDual{<:KeyedArray{Float64}}, + params_cd::CoDual{Vector{T}}, + shocks_cd::CoDual{Matrix{T}}, + me_std_cd::CoDual{M}, + initial_state_cd::CoDual{Vector{I}}, +) where {T<:Base.IEEEFloat,I<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} + fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_cd, initial_state_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_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) + return ntuple(length(fargs)) do i + increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) + end + end + return CoDual(y_primal, y_fdata), pb!! +end + +# ── 6-arg kwcall ── +function Mooncake.rrule!!( + kwcall_cd::CoDual{typeof(Core.kwcall)}, + kwargs_cd::CoDual{<:NamedTuple}, + f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, + model_cd::CoDual{MacroModelling.ℳ}, + data_cd::CoDual{<:KeyedArray{Float64}}, + params_cd::CoDual{Vector{T}}, + shocks_cd::CoDual{Matrix{T}}, + me_std_cd::CoDual{M}, + initial_state_cd::CoDual{Vector{I}}, +) where {T<:Base.IEEEFloat,I<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} + kwargs = Mooncake.primal(kwargs_cd) + fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_cd, initial_state_cd) + primals = map(Mooncake.primal, fargs) + lazy_rdata = map(Mooncake.lazy_zero_rdata, primals) + y_primal, cr_pb = ChainRulesCore.rrule(primals...; kwargs...) + y_fdata = Mooncake.fdata(Mooncake.zero_tangent(y_primal)) + kwargs_lazy_rdata = Mooncake.lazy_zero_rdata(kwargs) + function pb!!(y_rdata) + cr_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) + kwargs_rdata = Mooncake.increment_and_get_rdata!( + Mooncake.tangent(kwargs_cd), + Mooncake.instantiate(kwargs_lazy_rdata), + ChainRulesCore.NoTangent(), + ) + regular_rdata = ntuple(length(fargs)) do i + increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) + end + return (NoRData(), kwargs_rdata, regular_rdata...) + end + return CoDual(y_primal, y_fdata), pb!! +end # Nested Vector{Vector} initial_state forms are implemented manually below. # ── DynamicPPL compatibility: wider @is_primitive declarations ── @@ -188,6 +318,7 @@ function increment_nested_initial_state_rdata!( ) cr_tangent isa ChainRulesCore.AbstractZero && return NoRData() initial_state_fdata = Mooncake.tangent(initial_state_cd) + mooncake_tangent_is_inactive(initial_state_fdata) && return NoRData() @inbounds for i in eachindex(cr_tangent) component_tangent = cr_tangent[i] component_tangent isa ChainRulesCore.AbstractZero && continue @@ -197,11 +328,9 @@ function increment_nested_initial_state_rdata!( end function increment_rrule_args_with_nested_initial_state(fargs::Tuple, lazy_rdata::Tuple, cr_dfargs::Tuple) - regular_rdata = ntuple(i -> Mooncake.increment_and_get_rdata!( - Mooncake.tangent(fargs[i]), - Mooncake.instantiate(lazy_rdata[i]), - cr_dfargs[i], - ), length(fargs) - 1) + regular_rdata = ntuple(length(fargs) - 1) do i + increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) + end initial_state_rdata = increment_nested_initial_state_rdata!(fargs[end], cr_dfargs[end]) return (regular_rdata..., initial_state_rdata) end diff --git a/src/rrules.jl b/src/rrules.jl index dd7963abc..27647e4d4 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -17,6 +17,12 @@ # - Matrix equations: solve_sylvester_equation, solve_lyapunov_equation # - Filters: calculate_loglikelihood, run_kalman_iterations, find_shocks +# Instrumentation counter: incremented inside the @thunk body for d_params in the +# filter-free get_loglikelihood pullback. Used by tests to verify that the +# expensive ss_pb computation is NOT triggered when differentiating only w.r.t. +# shocks / me_std. +const _params_pullback_counter = Ref{Int}(0) + # 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) @@ -32,7 +38,7 @@ function rrule(::typeof(mat_mult_kron), Y = mat_mult_kron(A, B, C, D) function mat_mult_kron_pullback(Ȳ) - Ȳ = unthunk(Ȳ) + Ȳ = (Ȳ) if Ȳ isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -124,7 +130,7 @@ function rrule(::typeof(sparse_preallocated!), Ŝ::Matrix{T}; ℂ::higher_order_ project_Ŝ = ProjectTo(Ŝ) function sparse_preallocated_pullback(Ω̄) - ΔΩ = unthunk(Ω̄) + ΔΩ = (Ω̄) ΔŜ = project_Ŝ(ΔΩ) return NoTangent(), ΔŜ, NoTangent() end @@ -210,7 +216,7 @@ function rrule(::typeof(solve_stochastic_steady_state_newton), function second_order_stochastic_steady_state_pullback(∂x) # @timeit_debug timer "Calculate SSS - pullback" begin - ∂x₁ = unthunk(∂x[1]) + ∂x₁ = (∂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' @@ -304,7 +310,7 @@ function rrule(::typeof(solve_stochastic_steady_state_newton), ∂𝐒₃ = zero(𝐒₃) function third_order_stochastic_steady_state_pullback(∂x) - ∂x₁ = unthunk(∂x[1]) + ∂x₁ = (∂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' @@ -334,7 +340,7 @@ function rrule(::typeof(calculate_jacobian), return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() end - ∂∇₁u = unthunk(∂∇₁) + ∂∇₁u = (∂∇₁) copyto!(∂∇₁_vec, ∂∇₁u) jacobian_funcs.f_parameters(caches_obj.jacobian_parameters, parameters, SS_and_pars) @@ -363,7 +369,7 @@ function rrule(::typeof(calculate_hessian), return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() end - ∂∇₂u = unthunk(∂∇₂) + ∂∇₂u = (∂∇₂) copyto!(∂∇₂_vec, ∂∇₂u) hessian_funcs.f_parameters(caches_obj.hessian_parameters, parameters, SS_and_pars) @@ -393,7 +399,7 @@ function rrule(::typeof(calculate_third_order_derivatives), return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() end - ∂∇₃u = unthunk(∂∇₃) + ∂∇₃u = (∂∇₃) copyto!(∂∇₃_vec, ∂∇₃u) third_order_derivatives_funcs.f_parameters(caches_obj.third_order_derivatives_parameters, parameters, SS_and_pars) @@ -414,7 +420,7 @@ function incremental_cotangent!(Δ, prev_ref::Base.RefValue) return Δ end - Δu = unthunk(Δ) + Δu = (Δ) prev = prev_ref[] prev_ref[] = copy(Δu) @@ -597,7 +603,7 @@ function rrule(::typeof(get_NSSS_and_parameters), # try block-gmres here function get_non_stochastic_steady_state_pullback(∂SS_and_pars) - ∂SS = unthunk(∂SS_and_pars[1]) + ∂SS = (∂SS_and_pars[1]) if ∂SS isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, size(jvp_no_exo, 2)), NoTangent() end @@ -631,7 +637,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, zeros(S, 0, 0), [state], false) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -671,7 +677,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, zeros(S, 0, 0), [state], false) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -690,7 +696,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, 𝐒₁, [state], true) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -958,7 +964,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -991,7 +997,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1011,7 +1017,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂) pullback = function (Δresult) - Δ = unthunk(Δresult) + Δ = (Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1091,7 +1097,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1120,7 +1126,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, true, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂) pullback = function (Δresult) - Δ = unthunk(Δresult) + Δ = (Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1186,7 +1192,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1222,7 +1228,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1257,7 +1263,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1281,7 +1287,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃̂) pullback = function (Δresult) - Δ = unthunk(Δresult) + Δ = (Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1394,7 +1400,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1430,7 +1436,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) + Δ = (Δ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}) @@ -1458,7 +1464,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, true, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃̂) pullback = function (Δresult) - Δ = unthunk(Δresult) + Δ = (Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1572,7 +1578,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1646,7 +1652,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1720,7 +1726,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1798,7 +1804,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged) pullback = function (ȳ) - Δy = unthunk(ȳ) + Δy = (ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1974,7 +1980,7 @@ function rrule(::typeof(get_loglikelihood), # ── pullback ── pullback = function (∂llh_bar) - ∂llh = unthunk(∂llh_bar) + ∂llh = (∂llh_bar) # backprop through calculate_loglikelihood # returns: (_, _, _, _, ∂𝐒, ∂data_in_deviations, _, ∂state, _) @@ -2544,7 +2550,7 @@ function rrule(::typeof(get_irf), # Pullback (common structure, algorithm-specific parts dispatched) function get_irf_pullback(∂result_bar) - ∂result = unthunk(∂result_bar) + ∂result = (∂result_bar) if ∂result isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, length(parameters)) @@ -2657,10 +2663,10 @@ function rrule(::typeof(calculate_covariance), Δ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) + Δcovar = (Δcovar) + Δsol_ret = (Δsol_ret) + Δ∇₁_ret = (Δ∇₁_ret) + ΔSS_ret = (ΔSS_ret) # Accumulators ∂sol_total = zeros(S, size(sol)) @@ -2767,7 +2773,7 @@ function rrule(::typeof(calculate_mean), mean_of_variables = SS_and_pars[1:nVars] function first_order_mean_pullback(∂out) - ∂mean = unthunk(∂out[1]) + ∂mean = (∂out[1]) if ∂mean isa AbstractZero return NoTangent(), zeros(S, np), NoTangent() end @@ -2880,7 +2886,7 @@ function rrule(::typeof(calculate_mean), # ── Pullback ── function calculate_mean_pullback(∂out) - ∂mean_in = unthunk(∂out[1]) + ∂mean_in = (∂out[1]) if ∂mean_in isa AbstractZero return NoTangent(), zeros(S, np), NoTangent() @@ -3083,15 +3089,15 @@ function rrule(::typeof(calculate_second_order_moments), ∂𝐒₁_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) + ∂μʸ₂_in = (∂μʸ₂_in) + ∂Δμˢ₂_in = (∂Δμˢ₂_in) + ∂Σʸ₁_pass = (∂Σʸ₁_pass) + ∂Σᶻ₁_pass = (∂Σᶻ₁_pass) + ∂SS_pass = (∂SS_pass) + ∂𝐒₁_pass = (∂𝐒₁_pass) + ∂∇₁_pass = (∂∇₁_pass) + ∂𝐒₂_pass = (∂𝐒₂_pass) + ∂∇₂_pass = (∂∇₂_pass) # Accumulators ∂𝐒₁_acc = zeros(S, size(𝐒₁)) @@ -3362,20 +3368,20 @@ function rrule(::typeof(calculate_second_order_moments_with_covariance), ∂𝐒₁_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) + ∂Σʸ₂_in = (∂Σʸ₂_in) + ∂Σᶻ₂_pass = (∂Σᶻ₂_pass) + ∂μʸ₂_in = (∂μʸ₂_in) + ∂Δμˢ₂_in = (∂Δμˢ₂_in) + ∂at_in = (∂at_in) + ∂ŝŝ₂_pass = (∂ŝŝ₂_pass) + ∂ŝy₂_pass = (∂ŝy₂_pass) + ∂Σʸ₁_pass = (∂Σʸ₁_pass) + ∂Σᶻ₁_pass = (∂Σᶻ₁_pass) + ∂SS_pass = (∂SS_pass) + ∂𝐒₁_pass = (∂𝐒₁_pass) + ∂∇₁_pass = (∂∇₁_pass) + ∂𝐒₂_pass = (∂𝐒₂_pass) + ∂∇₂_pass = (∂∇₂_pass) # Accumulators ∂𝐒₁_acc = zeros(S, size(𝐒₁)) @@ -3911,9 +3917,9 @@ function rrule(::typeof(calculate_third_order_moments), function calculate_third_order_moments_pullback(∂out) ∂Σʸ₃_in, ∂μʸ₂_in, ∂SS_in, _ = ∂out - ∂Σʸ₃_in = unthunk(∂Σʸ₃_in) - ∂μʸ₂_in = unthunk(∂μʸ₂_in) - ∂SS_in = unthunk(∂SS_in) + ∂Σʸ₃_in = (∂Σʸ₃_in) + ∂μʸ₂_in = (∂μʸ₂_in) + ∂SS_in = (∂SS_in) n₋ = T_pm.nPast_not_future_and_mixed @@ -4842,10 +4848,10 @@ function rrule(::typeof(calculate_third_order_moments_with_autocorrelation), 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) + ∂Σʸ₃_in = (∂Σʸ₃_in) + ∂μʸ₂_in = (∂μʸ₂_in) + ∂autocorr_in = (∂autocorr_in) + ∂SS_in = (∂SS_in) n₋ = T_pm.nPast_not_future_and_mixed @@ -5742,7 +5748,7 @@ function rrule(::typeof(calculate_first_order_solution), # 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]) + ∂𝐒_mat = (∂𝐒[1]) if ∂𝐒_mat isa Union{NoTangent, AbstractZero} return NoTangent(), zero(∇₁), NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -6047,7 +6053,7 @@ function rrule(::typeof(calculate_second_order_solution), # end # timeit_debug - ∂𝐒₂ = unthunk(∂𝐒₂_solved[1]) + ∂𝐒₂ = (∂𝐒₂_solved[1]) if size(∂𝐒₂, 2) == size(𝐒₂_stable, 2) nothing @@ -8410,7 +8416,7 @@ function rrule(::typeof(calculate_third_order_solution), # PULLBACK # ========================================================================= function third_order_solution_pullback(∂𝐒₃_solved) - ∂𝐒₃ = choose_matrix_format(unthunk(∂𝐒₃_solved[1])) + ∂𝐒₃ = choose_matrix_format((∂𝐒₃_solved[1])) if ℒ.norm(∂𝐒₃) < opts.tol.third_order.ad.sylvester.acceptance_tol return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent()) @@ -8751,7 +8757,7 @@ function rrule(::typeof(solve_sylvester_equation), # pullback function solve_sylvester_equation_pullback(∂P) - ∂P₁ = unthunk(∂P[1]) + ∂P₁ = (∂P[1]) if ℒ.norm(∂P₁) < tol.rtol return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -8835,7 +8841,7 @@ function rrule(::typeof(solve_lyapunov_equation), # pullback # https://arxiv.org/abs/2011.11430 function solve_lyapunov_equation_pullback(∂P) - ∂P₁ = unthunk(∂P[1]) + ∂P₁ = (∂P[1]) if ℒ.norm(∂P₁) < tol.rtol return NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -8924,7 +8930,7 @@ function rrule(::typeof(find_shocks), # ∂𝐒ⁱ²ᵉ = similar(𝐒ⁱ²ᵉ) function find_shocks_pullback(∂x) - ∂x = vcat(unthunk(∂x[1]), zero(λ)) + ∂x = vcat((∂x[1]), zero(λ)) S = -fXλp' \ ∂x @@ -8989,7 +8995,7 @@ function rrule(::typeof(find_shocks), xxλ = ℒ.kron(x,xλ) function find_shocks_pullback(∂x) - ∂x = vcat(unthunk(∂x[1]), zero(λ)) + ∂x = vcat((∂x[1]), zero(λ)) S = -fXλp' \ ∂x @@ -17573,7 +17579,7 @@ end function get_statistics_cotangent_helper(Δret, key::Symbol) - Δ = unthunk(Δret) + Δ = (Δret) if Δ isa Union{NoTangent, AbstractZero} return NoTangent() end @@ -17751,7 +17757,7 @@ function rrule(::typeof(get_statistics), end ∂SS = zeros(T, length(SS)) - ∂SS[SS_var_idx] .+= unthunk(Δnsss) + ∂SS[SS_var_idx] .+= (Δnsss) ∂SS_and_pars = zeros(T, length(SS_and_pars)) ∂SS_and_pars[1:length(SS)] .+= ∂SS @@ -18336,27 +18342,27 @@ function rrule(::typeof(get_solution), result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂_raw], true) pullback_2nd = function (∂result_bar) - Δ = unthunk(∂result_bar) + Δ = (∂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}} + ∂mats = (Δ[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 = (∂mats[1]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end ∂𝐒₂_ext = if ∂mats isa Union{NoTangent, AbstractZero} NoTangent() else - m = unthunk(∂mats[2]) + m = (∂mats[2]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end @@ -18468,33 +18474,33 @@ function rrule(::typeof(get_solution), result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂_raw, 𝐒₃_raw], true) pullback_3rd = function (∂result_bar) - Δ = unthunk(∂result_bar) + Δ = (∂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}} + ∂mats = (Δ[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 = (∂mats[1]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end ∂𝐒₂_ext = if ∂mats isa Union{NoTangent, AbstractZero} NoTangent() else - m = unthunk(∂mats[2]) + m = (∂mats[2]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end ∂𝐒₃_ext = if ∂mats isa Union{NoTangent, AbstractZero} NoTangent() else - m = unthunk(∂mats[3]) + m = (∂mats[3]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end @@ -18577,21 +18583,21 @@ function rrule(::typeof(get_solution), result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁], true) pullback_1st = function (∂result_bar) - Δ = unthunk(∂result_bar) + Δ = (∂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}} + ∂mats = (Δ[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 = (∂mats[1]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end @@ -19551,7 +19557,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = unthunk(Δ) + Δllh = (Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19574,6 +19580,7 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin + _params_pullback_counter[] += 1 d_𝐒₁_full_cot = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) if has_override && !(d_state_red isa Union{NoTangent, AbstractZero}) @@ -19638,7 +19645,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = unthunk(Δ) + Δllh = (Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19662,6 +19669,7 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin + _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) @@ -19739,7 +19747,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = unthunk(Δ) + Δllh = (Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19763,6 +19771,7 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin + _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) @@ -19834,7 +19843,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = unthunk(Δ) + Δllh = (Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19859,6 +19868,7 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin + _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) @@ -19953,7 +19963,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = unthunk(Δ) + Δllh = (Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19978,6 +19988,7 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin + _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) From 92240dffa9999f4ef1507ae702783d63469395a0 Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Wed, 10 Jun 2026 09:39:25 +0000 Subject: [PATCH 08/37] no thunking except at high level in filter free llh calc --- ext/MooncakeExt.jl | 159 ++++----------------------------------------- src/rrules.jl | 11 ---- 2 files changed, 14 insertions(+), 156 deletions(-) diff --git a/ext/MooncakeExt.jl b/ext/MooncakeExt.jl index ebfe975ff..6cfabc44a 100644 --- a/ext/MooncakeExt.jl +++ b/ext/MooncakeExt.jl @@ -71,147 +71,15 @@ Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglike Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_irf), MacroModelling.ℳ, Vector{T}} where {T<:Base.IEEEFloat} true -# ── get_loglikelihood (filter-free, with shocks/me_std): manual rrule!! ── -# Replaces @from_rrule so that CRC @thunks for d_params / d_initial_state are -# NOT forced when the caller differentiates only w.r.t. shocks / me_std. -# When an argument's fdata is inactive the corresponding CRC tangent (which may -# contain an expensive @thunk) is skipped entirely. - -mooncake_tangent_is_inactive(::Mooncake.NoTangent) = true -mooncake_tangent_is_inactive(tangent::Mooncake.PossiblyUninitTangent) = !Mooncake.is_init(tangent) -mooncake_tangent_is_inactive(_) = false - -function increment_rrule_arg_rdata!(arg_cd::CoDual, lazy_rdata, cr_tangent) - cr_tangent isa ChainRulesCore.AbstractZero && return NoRData() - arg_fdata = Mooncake.tangent(arg_cd) - if Mooncake.primal(arg_cd) isa AbstractArray && arg_fdata isa Mooncake.NoFData - return NoRData() - end - mooncake_tangent_is_inactive(arg_fdata) && return NoRData() - return Mooncake.increment_and_get_rdata!( - arg_fdata, - Mooncake.instantiate(lazy_rdata), - cr_tangent, - ) -end - -Mooncake._get_fdata_field(primal::Function, ::Mooncake.NoFData, f) = Mooncake.NoFData() -Mooncake.uninit_fcodual(x::AbstractArray{<:Base.IEEEFloat}) = CoDual(x, Mooncake.NoFData()) - -# ── 5-arg positional (no initial_state) ── -# (𝓂, data, parameter_values, shocks, me_std) -function Mooncake.rrule!!( - f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, - model_cd::CoDual{MacroModelling.ℳ}, - data_cd::CoDual{<:KeyedArray{Float64}}, - params_cd::CoDual{Vector{T}}, - shocks_cd::CoDual{Matrix{T}}, - me_std_cd::CoDual{M}, -) where {T<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} - fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_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_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) - return ntuple(length(fargs)) do i - increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) - end - end - return CoDual(y_primal, y_fdata), pb!! -end - -# ── 5-arg kwcall ── -function Mooncake.rrule!!( - kwcall_cd::CoDual{typeof(Core.kwcall)}, - kwargs_cd::CoDual{<:NamedTuple}, - f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, - model_cd::CoDual{MacroModelling.ℳ}, - data_cd::CoDual{<:KeyedArray{Float64}}, - params_cd::CoDual{Vector{T}}, - shocks_cd::CoDual{Matrix{T}}, - me_std_cd::CoDual{M}, -) where {T<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} - kwargs = Mooncake.primal(kwargs_cd) - fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_cd) - primals = map(Mooncake.primal, fargs) - lazy_rdata = map(Mooncake.lazy_zero_rdata, primals) - y_primal, cr_pb = ChainRulesCore.rrule(primals...; kwargs...) - y_fdata = Mooncake.fdata(Mooncake.zero_tangent(y_primal)) - kwargs_lazy_rdata = Mooncake.lazy_zero_rdata(kwargs) - function pb!!(y_rdata) - cr_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) - kwargs_rdata = Mooncake.increment_and_get_rdata!( - Mooncake.tangent(kwargs_cd), - Mooncake.instantiate(kwargs_lazy_rdata), - ChainRulesCore.NoTangent(), - ) - regular_rdata = ntuple(length(fargs)) do i - increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) - end - return (NoRData(), kwargs_rdata, regular_rdata...) - end - return CoDual(y_primal, y_fdata), pb!! -end - -# ── 6-arg positional (with Vector{Float64} initial_state) ── -function Mooncake.rrule!!( - f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, - model_cd::CoDual{MacroModelling.ℳ}, - data_cd::CoDual{<:KeyedArray{Float64}}, - params_cd::CoDual{Vector{T}}, - shocks_cd::CoDual{Matrix{T}}, - me_std_cd::CoDual{M}, - initial_state_cd::CoDual{Vector{I}}, -) where {T<:Base.IEEEFloat,I<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} - fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_cd, initial_state_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_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) - return ntuple(length(fargs)) do i - increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) - end - end - return CoDual(y_primal, y_fdata), pb!! -end - -# ── 6-arg kwcall ── -function Mooncake.rrule!!( - kwcall_cd::CoDual{typeof(Core.kwcall)}, - kwargs_cd::CoDual{<:NamedTuple}, - f_cd::CoDual{typeof(MacroModelling.get_loglikelihood)}, - model_cd::CoDual{MacroModelling.ℳ}, - data_cd::CoDual{<:KeyedArray{Float64}}, - params_cd::CoDual{Vector{T}}, - shocks_cd::CoDual{Matrix{T}}, - me_std_cd::CoDual{M}, - initial_state_cd::CoDual{Vector{I}}, -) where {T<:Base.IEEEFloat,I<:Base.IEEEFloat,M<:Union{T,Vector{T},Matrix{T}}} - kwargs = Mooncake.primal(kwargs_cd) - fargs = (f_cd, model_cd, data_cd, params_cd, shocks_cd, me_std_cd, initial_state_cd) - primals = map(Mooncake.primal, fargs) - lazy_rdata = map(Mooncake.lazy_zero_rdata, primals) - y_primal, cr_pb = ChainRulesCore.rrule(primals...; kwargs...) - y_fdata = Mooncake.fdata(Mooncake.zero_tangent(y_primal)) - kwargs_lazy_rdata = Mooncake.lazy_zero_rdata(kwargs) - function pb!!(y_rdata) - cr_dfargs = cr_pb(scalar_output_to_cr(y_fdata, y_rdata)) - kwargs_rdata = Mooncake.increment_and_get_rdata!( - Mooncake.tangent(kwargs_cd), - Mooncake.instantiate(kwargs_lazy_rdata), - ChainRulesCore.NoTangent(), - ) - regular_rdata = ntuple(length(fargs)) do i - increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) - end - return (NoRData(), kwargs_rdata, regular_rdata...) - end - return CoDual(y_primal, y_fdata), pb!! -end +# get_loglikelihood: +# (𝓂, data, parameter_values::Vector{T}, shocks::Matrix{T}, me_std::T_or_Vector{T}) +# Two narrow @from_rrule generations cover scalar and vector me_std cases. +# 5-arg kwarg-only path (no AD through initial_state) +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, T} where {T<:Base.IEEEFloat} true +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Vector{T}} where {T<:Base.IEEEFloat} true +# 6-arg positional path (AD through initial_state) +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, T, Vector{Float64}} where {T<:Base.IEEEFloat} true +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Vector{T}, Vector{Float64}} where {T<:Base.IEEEFloat} true # Nested Vector{Vector} initial_state forms are implemented manually below. # ── DynamicPPL compatibility: wider @is_primitive declarations ── @@ -318,7 +186,6 @@ function increment_nested_initial_state_rdata!( ) cr_tangent isa ChainRulesCore.AbstractZero && return NoRData() initial_state_fdata = Mooncake.tangent(initial_state_cd) - mooncake_tangent_is_inactive(initial_state_fdata) && return NoRData() @inbounds for i in eachindex(cr_tangent) component_tangent = cr_tangent[i] component_tangent isa ChainRulesCore.AbstractZero && continue @@ -328,9 +195,11 @@ function increment_nested_initial_state_rdata!( end function increment_rrule_args_with_nested_initial_state(fargs::Tuple, lazy_rdata::Tuple, cr_dfargs::Tuple) - regular_rdata = ntuple(length(fargs) - 1) do i - increment_rrule_arg_rdata!(fargs[i], lazy_rdata[i], cr_dfargs[i]) - end + regular_rdata = ntuple(i -> Mooncake.increment_and_get_rdata!( + Mooncake.tangent(fargs[i]), + Mooncake.instantiate(lazy_rdata[i]), + cr_dfargs[i], + ), length(fargs) - 1) initial_state_rdata = increment_nested_initial_state_rdata!(fargs[end], cr_dfargs[end]) return (regular_rdata..., initial_state_rdata) end diff --git a/src/rrules.jl b/src/rrules.jl index 27647e4d4..54c4a3b76 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -17,12 +17,6 @@ # - Matrix equations: solve_sylvester_equation, solve_lyapunov_equation # - Filters: calculate_loglikelihood, run_kalman_iterations, find_shocks -# Instrumentation counter: incremented inside the @thunk body for d_params in the -# filter-free get_loglikelihood pullback. Used by tests to verify that the -# expensive ss_pb computation is NOT triggered when differentiating only w.r.t. -# shocks / me_std. -const _params_pullback_counter = Ref{Int}(0) - # 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) @@ -19580,7 +19574,6 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin - _params_pullback_counter[] += 1 d_𝐒₁_full_cot = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) if has_override && !(d_state_red isa Union{NoTangent, AbstractZero}) @@ -19669,7 +19662,6 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin - _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) @@ -19771,7 +19763,6 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin - _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) @@ -19868,7 +19859,6 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin - _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) @@ -19988,7 +19978,6 @@ function rrule(::typeof(get_loglikelihood), d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) d_params = @thunk begin - _params_pullback_counter[] += 1 d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) From 8f12e9a3d4ae430632c414eaa2abafcd5e28696f Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Wed, 10 Jun 2026 12:31:38 +0000 Subject: [PATCH 09/37] Add additional rrule definitions for get_loglikelihood with Matrix{T} support --- ext/MooncakeExt.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ext/MooncakeExt.jl b/ext/MooncakeExt.jl index 6cfabc44a..fea87f7f7 100644 --- a/ext/MooncakeExt.jl +++ b/ext/MooncakeExt.jl @@ -77,9 +77,11 @@ Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_irf), M # 5-arg kwarg-only path (no AD through initial_state) Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, T} where {T<:Base.IEEEFloat} true Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Vector{T}} where {T<:Base.IEEEFloat} true +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Matrix{T}} where {T<:Base.IEEEFloat} true # 6-arg positional path (AD through initial_state) Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, T, Vector{Float64}} where {T<:Base.IEEEFloat} true Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Vector{T}, Vector{Float64}} where {T<:Base.IEEEFloat} true +Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_loglikelihood), MacroModelling.ℳ, KeyedArray{Float64}, Vector{T}, Matrix{T}, Matrix{T}, Vector{Float64}} where {T<:Base.IEEEFloat} true # Nested Vector{Vector} initial_state forms are implemented manually below. # ── DynamicPPL compatibility: wider @is_primitive declarations ── @@ -448,4 +450,4 @@ function Mooncake.rrule!!( return CoDual(nothing, Mooncake.NoFData()), pb!! end -end # module MooncakeExt +end # module MooncakeExt \ No newline at end of file From 8003dbc77989030a3db4651fcc547647e2573293 Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Wed, 10 Jun 2026 13:47:00 +0000 Subject: [PATCH 10/37] make gradient checks pass --- src/rrules.jl | 505 +++++++++++++++++++++++--------------------------- 1 file changed, 236 insertions(+), 269 deletions(-) diff --git a/src/rrules.jl b/src/rrules.jl index 54c4a3b76..6e24abec2 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -32,7 +32,7 @@ function rrule(::typeof(mat_mult_kron), Y = mat_mult_kron(A, B, C, D) function mat_mult_kron_pullback(Ȳ) - Ȳ = (Ȳ) + Ȳ = unthunk(Ȳ) if Ȳ isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -124,7 +124,7 @@ function rrule(::typeof(sparse_preallocated!), Ŝ::Matrix{T}; ℂ::higher_order_ project_Ŝ = ProjectTo(Ŝ) function sparse_preallocated_pullback(Ω̄) - ΔΩ = (Ω̄) + ΔΩ = unthunk(Ω̄) ΔŜ = project_Ŝ(ΔΩ) return NoTangent(), ΔŜ, NoTangent() end @@ -210,7 +210,7 @@ function rrule(::typeof(solve_stochastic_steady_state_newton), function second_order_stochastic_steady_state_pullback(∂x) # @timeit_debug timer "Calculate SSS - pullback" begin - ∂x₁ = (∂x[1]) + ∂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' @@ -304,7 +304,7 @@ function rrule(::typeof(solve_stochastic_steady_state_newton), ∂𝐒₃ = zero(𝐒₃) function third_order_stochastic_steady_state_pullback(∂x) - ∂x₁ = (∂x[1]) + ∂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' @@ -334,7 +334,7 @@ function rrule(::typeof(calculate_jacobian), return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() end - ∂∇₁u = (∂∇₁) + ∂∇₁u = unthunk(∂∇₁) copyto!(∂∇₁_vec, ∂∇₁u) jacobian_funcs.f_parameters(caches_obj.jacobian_parameters, parameters, SS_and_pars) @@ -363,7 +363,7 @@ function rrule(::typeof(calculate_hessian), return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() end - ∂∇₂u = (∂∇₂) + ∂∇₂u = unthunk(∂∇₂) copyto!(∂∇₂_vec, ∂∇₂u) hessian_funcs.f_parameters(caches_obj.hessian_parameters, parameters, SS_and_pars) @@ -393,7 +393,7 @@ function rrule(::typeof(calculate_third_order_derivatives), return NoTangent(), zero(parameters), zero(SS_and_pars), NoTangent(), NoTangent(), NoTangent() end - ∂∇₃u = (∂∇₃) + ∂∇₃u = unthunk(∂∇₃) copyto!(∂∇₃_vec, ∂∇₃u) third_order_derivatives_funcs.f_parameters(caches_obj.third_order_derivatives_parameters, parameters, SS_and_pars) @@ -414,7 +414,7 @@ function incremental_cotangent!(Δ, prev_ref::Base.RefValue) return Δ end - Δu = (Δ) + Δu = unthunk(Δ) prev = prev_ref[] prev_ref[] = copy(Δu) @@ -597,7 +597,7 @@ function rrule(::typeof(get_NSSS_and_parameters), # try block-gmres here function get_non_stochastic_steady_state_pullback(∂SS_and_pars) - ∂SS = (∂SS_and_pars[1]) + ∂SS = unthunk(∂SS_and_pars[1]) if ∂SS isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, size(jvp_no_exo, 2)), NoTangent() end @@ -631,7 +631,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, zeros(S, 0, 0), [state], false) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -671,7 +671,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, zeros(S, 0, 0), [state], false) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -690,7 +690,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, 𝐒₁, [state], true) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -958,7 +958,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -991,7 +991,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1011,7 +1011,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂) pullback = function (Δresult) - Δ = (Δresult) + Δ = unthunk(Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1091,7 +1091,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1120,7 +1120,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, true, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂) pullback = function (Δresult) - Δ = (Δresult) + Δ = unthunk(Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1186,7 +1186,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1222,7 +1222,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1257,7 +1257,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1281,7 +1281,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃̂) pullback = function (Δresult) - Δ = (Δresult) + Δ = unthunk(Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1394,7 +1394,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1430,7 +1430,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), 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) - Δ = (Δ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}) @@ -1458,7 +1458,7 @@ function rrule(::typeof(calculate_stochastic_steady_state), result = (sss, true, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃̂) pullback = function (Δresult) - Δ = (Δresult) + Δ = unthunk(Δresult) Δsss = zeros(Float64, length(sss)) ΔSS_and_pars = zeros(Float64, length(SS_and_pars)) Δ∇₁ = zeros(Float64, size(∇₁)) @@ -1572,7 +1572,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1646,7 +1646,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂], state, converged) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1720,7 +1720,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1798,7 +1798,7 @@ function rrule(::typeof(get_relevant_steady_state_and_state_update), y = (𝓂.constants, SS_and_pars, [𝐒₁, 𝐒₂, 𝐒₃], state, converged) pullback = function (ȳ) - Δy = (ȳ) + Δy = unthunk(ȳ) if Δy isa NoTangent || Δy isa AbstractZero return NoTangent(), NoTangent(), zeros(S, length(parameter_values)), NoTangent() end @@ -1974,7 +1974,7 @@ function rrule(::typeof(get_loglikelihood), # ── pullback ── pullback = function (∂llh_bar) - ∂llh = (∂llh_bar) + ∂llh = unthunk(∂llh_bar) # backprop through calculate_loglikelihood # returns: (_, _, _, _, ∂𝐒, ∂data_in_deviations, _, ∂state, _) @@ -2544,7 +2544,7 @@ function rrule(::typeof(get_irf), # Pullback (common structure, algorithm-specific parts dispatched) function get_irf_pullback(∂result_bar) - ∂result = (∂result_bar) + ∂result = unthunk(∂result_bar) if ∂result isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, length(parameters)) @@ -2657,10 +2657,10 @@ function rrule(::typeof(calculate_covariance), Δcovar, Δsol_ret, Δ∇₁_ret, ΔSS_ret, _ = Δout # Materialise any InplaceableThunk / Thunk wrappers - Δcovar = (Δcovar) - Δsol_ret = (Δsol_ret) - Δ∇₁_ret = (Δ∇₁_ret) - ΔSS_ret = (ΔSS_ret) + Δcovar = unthunk(Δcovar) + Δsol_ret = unthunk(Δsol_ret) + Δ∇₁_ret = unthunk(Δ∇₁_ret) + ΔSS_ret = unthunk(ΔSS_ret) # Accumulators ∂sol_total = zeros(S, size(sol)) @@ -2767,7 +2767,7 @@ function rrule(::typeof(calculate_mean), mean_of_variables = SS_and_pars[1:nVars] function first_order_mean_pullback(∂out) - ∂mean = (∂out[1]) + ∂mean = unthunk(∂out[1]) if ∂mean isa AbstractZero return NoTangent(), zeros(S, np), NoTangent() end @@ -2880,7 +2880,7 @@ function rrule(::typeof(calculate_mean), # ── Pullback ── function calculate_mean_pullback(∂out) - ∂mean_in = (∂out[1]) + ∂mean_in = unthunk(∂out[1]) if ∂mean_in isa AbstractZero return NoTangent(), zeros(S, np), NoTangent() @@ -3083,15 +3083,15 @@ function rrule(::typeof(calculate_second_order_moments), ∂𝐒₁_pass, ∂∇₁_pass, ∂𝐒₂_pass, ∂∇₂_pass, _ = ∂out # Materialise any InplaceableThunk / Thunk wrappers - ∂μʸ₂_in = (∂μʸ₂_in) - ∂Δμˢ₂_in = (∂Δμˢ₂_in) - ∂Σʸ₁_pass = (∂Σʸ₁_pass) - ∂Σᶻ₁_pass = (∂Σᶻ₁_pass) - ∂SS_pass = (∂SS_pass) - ∂𝐒₁_pass = (∂𝐒₁_pass) - ∂∇₁_pass = (∂∇₁_pass) - ∂𝐒₂_pass = (∂𝐒₂_pass) - ∂∇₂_pass = (∂∇₂_pass) + ∂μʸ₂_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(𝐒₁)) @@ -3362,20 +3362,20 @@ function rrule(::typeof(calculate_second_order_moments_with_covariance), ∂𝐒₁_pass, ∂∇₁_pass, ∂𝐒₂_pass, ∂∇₂_pass, _ = ∂out # Materialise any InplaceableThunk / Thunk wrappers - ∂Σʸ₂_in = (∂Σʸ₂_in) - ∂Σᶻ₂_pass = (∂Σᶻ₂_pass) - ∂μʸ₂_in = (∂μʸ₂_in) - ∂Δμˢ₂_in = (∂Δμˢ₂_in) - ∂at_in = (∂at_in) - ∂ŝŝ₂_pass = (∂ŝŝ₂_pass) - ∂ŝy₂_pass = (∂ŝy₂_pass) - ∂Σʸ₁_pass = (∂Σʸ₁_pass) - ∂Σᶻ₁_pass = (∂Σᶻ₁_pass) - ∂SS_pass = (∂SS_pass) - ∂𝐒₁_pass = (∂𝐒₁_pass) - ∂∇₁_pass = (∂∇₁_pass) - ∂𝐒₂_pass = (∂𝐒₂_pass) - ∂∇₂_pass = (∂∇₂_pass) + ∂Σʸ₂_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(𝐒₁)) @@ -3911,9 +3911,9 @@ function rrule(::typeof(calculate_third_order_moments), function calculate_third_order_moments_pullback(∂out) ∂Σʸ₃_in, ∂μʸ₂_in, ∂SS_in, _ = ∂out - ∂Σʸ₃_in = (∂Σʸ₃_in) - ∂μʸ₂_in = (∂μʸ₂_in) - ∂SS_in = (∂SS_in) + ∂Σʸ₃_in = unthunk(∂Σʸ₃_in) + ∂μʸ₂_in = unthunk(∂μʸ₂_in) + ∂SS_in = unthunk(∂SS_in) n₋ = T_pm.nPast_not_future_and_mixed @@ -4842,10 +4842,10 @@ function rrule(::typeof(calculate_third_order_moments_with_autocorrelation), function calculate_third_order_moments_with_autocorrelation_pullback(∂out) ∂Σʸ₃_in, ∂μʸ₂_in, ∂autocorr_in, ∂SS_in, _ = ∂out - ∂Σʸ₃_in = (∂Σʸ₃_in) - ∂μʸ₂_in = (∂μʸ₂_in) - ∂autocorr_in = (∂autocorr_in) - ∂SS_in = (∂SS_in) + ∂Σʸ₃_in = unthunk(∂Σʸ₃_in) + ∂μʸ₂_in = unthunk(∂μʸ₂_in) + ∂autocorr_in = unthunk(∂autocorr_in) + ∂SS_in = unthunk(∂SS_in) n₋ = T_pm.nPast_not_future_and_mixed @@ -5742,7 +5742,7 @@ function rrule(::typeof(calculate_first_order_solution), # 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 = (∂𝐒[1]) + ∂𝐒_mat = unthunk(∂𝐒[1]) if ∂𝐒_mat isa Union{NoTangent, AbstractZero} return NoTangent(), zero(∇₁), NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -6047,7 +6047,7 @@ function rrule(::typeof(calculate_second_order_solution), # end # timeit_debug - ∂𝐒₂ = (∂𝐒₂_solved[1]) + ∂𝐒₂ = unthunk(∂𝐒₂_solved[1]) if size(∂𝐒₂, 2) == size(𝐒₂_stable, 2) nothing @@ -6698,12 +6698,12 @@ 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, + ∂B::AbstractMatrix{R}, + M1::AbstractMatrix, + M2::AbstractMatrix, A::SparseMatrixCSC{TA, Ti}, B::DenseMatrix{TB}; - tol::Real = 0.0, + tol::Real = 0.0, block::AbstractMatrix{R} = Matrix{R}(undef, size(M1, 1), 0), ) where {R <: Real, TA <: Real, TB <: Real, Ti <: Integer} @@ -6852,15 +6852,15 @@ function mul_fill_kron_adjoint!( @simd for l in 1:m2 work[l] += M1_value * M2t[column_offset + l, M1_col] - end end + end for l in 1:m2 ∂A[k, l] += bij * work[l] - end + end if has_A_contribution - acc = zero(R) + acc = zero(R) for A_index in A_row_start:A_row_stop l = A_nz_col[A_index] @@ -6882,15 +6882,15 @@ function mul_fill_kron_adjoint!( for A_index in A_row_start:A_row_stop l = A_nz_col[A_index] local_acc += A_nz_val[A_index] * M2t[column_offset + l, M1_col] - end - - acc += M1_value * local_acc end - ∂B[i, j] += acc + acc += M1_value * local_acc end + + ∂B[i, j] += acc end end +end return ∂A, ∂B end @@ -8002,8 +8002,8 @@ end # but avoids the n_compressed³ × n_compressed³ allocation. function mul_compressed_kron³_pullback!( ∂X::AbstractMatrix, - M1::AbstractMatrix, - M2::AbstractMatrix, + M1::AbstractMatrix, + M2::AbstractMatrix, X::AbstractMatrix; tol::Real = 0.0, ) @@ -8083,9 +8083,9 @@ function mul_compressed_kron³_pullback!( row_scale = Vector{T}(undef, n_supported_rows) @inbounds for i1 in 1:n_x_rows - for j1 in 1:i1 + for j1 in 1:i1 for k1 in 1:j1 - row = (i1 - 1) * i1 * (i1 + 1) ÷ 6 + (j1 - 1) * j1 ÷ 2 + k1 + row = (i1 - 1) * i1 * (i1 + 1) ÷ 6 + (j1 - 1) * j1 ÷ 2 + k1 slot = row_to_slot[row] if slot != 0 @@ -8093,11 +8093,11 @@ function mul_compressed_kron³_pullback!( row_second[slot] = j1 row_third[slot] = k1 - if i1 == j1 + if i1 == j1 divisor = j1 == k1 ? 6 : 2 - else - divisor = (j1 == k1 || i1 == k1) ? 2 : 1 - end + else + divisor = (j1 == k1 || i1 == k1) ? 2 : 1 + end row_scale[slot] = inv(T(divisor)) end @@ -8148,29 +8148,29 @@ function mul_compressed_kron³_pullback!( ajj = Xd[j1, j2] ajk = Xd[j1, k2] - aki = Xd[k1, i2] - akj = Xd[k1, j2] - akk = Xd[k1, k2] + aki = Xd[k1, i2] + akj = Xd[k1, j2] + akk = Xd[k1, k2] q_i2 = aij * ajk + aik * ajj q_j2 = aik * aji + aii * ajk q_k2 = aii * ajj + aij * aji - ∂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[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[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 + ∂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 return ∂X end @@ -8410,7 +8410,7 @@ function rrule(::typeof(calculate_third_order_solution), # PULLBACK # ========================================================================= function third_order_solution_pullback(∂𝐒₃_solved) - ∂𝐒₃ = choose_matrix_format((∂𝐒₃_solved[1])) + ∂𝐒₃ = 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()) @@ -8751,7 +8751,7 @@ function rrule(::typeof(solve_sylvester_equation), # pullback function solve_sylvester_equation_pullback(∂P) - ∂P₁ = (∂P[1]) + ∂P₁ = unthunk(∂P[1]) if ℒ.norm(∂P₁) < tol.rtol return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -8835,7 +8835,7 @@ function rrule(::typeof(solve_lyapunov_equation), # pullback # https://arxiv.org/abs/2011.11430 function solve_lyapunov_equation_pullback(∂P) - ∂P₁ = (∂P[1]) + ∂P₁ = unthunk(∂P[1]) if ℒ.norm(∂P₁) < tol.rtol return NoTangent(), NoTangent(), NoTangent(), NoTangent() end @@ -8924,7 +8924,7 @@ function rrule(::typeof(find_shocks), # ∂𝐒ⁱ²ᵉ = similar(𝐒ⁱ²ᵉ) function find_shocks_pullback(∂x) - ∂x = vcat((∂x[1]), zero(λ)) + ∂x = vcat(unthunk(∂x[1]), zero(λ)) S = -fXλp' \ ∂x @@ -8989,7 +8989,7 @@ function rrule(::typeof(find_shocks), xxλ = ℒ.kron(x,xλ) function find_shocks_pullback(∂x) - ∂x = vcat((∂x[1]), zero(λ)) + ∂x = vcat(unthunk(∂x[1]), zero(λ)) S = -fXλp' \ ∂x @@ -17573,7 +17573,7 @@ end function get_statistics_cotangent_helper(Δret, key::Symbol) - Δ = (Δret) + Δ = unthunk(Δret) if Δ isa Union{NoTangent, AbstractZero} return NoTangent() end @@ -17751,7 +17751,7 @@ function rrule(::typeof(get_statistics), end ∂SS = zeros(T, length(SS)) - ∂SS[SS_var_idx] .+= (Δnsss) + ∂SS[SS_var_idx] .+= unthunk(Δnsss) ∂SS_and_pars = zeros(T, length(SS_and_pars)) ∂SS_and_pars[1:length(SS)] .+= ∂SS @@ -18336,27 +18336,27 @@ function rrule(::typeof(get_solution), result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂_raw], true) pullback_2nd = function (∂result_bar) - Δ = (∂result_bar) + Δ = unthunk(∂result_bar) if Δ isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, length(parameters)) end ∂NSSS = Δ[1] - ∂mats = (Δ[2]) # cotangent for Vector{AbstractMatrix{S}} + ∂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 = (∂mats[1]) + m = unthunk(∂mats[1]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end ∂𝐒₂_ext = if ∂mats isa Union{NoTangent, AbstractZero} NoTangent() else - m = (∂mats[2]) + m = unthunk(∂mats[2]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end @@ -18468,33 +18468,33 @@ function rrule(::typeof(get_solution), result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁, 𝐒₂_raw, 𝐒₃_raw], true) pullback_3rd = function (∂result_bar) - Δ = (∂result_bar) + Δ = unthunk(∂result_bar) if Δ isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, length(parameters)) end ∂NSSS = Δ[1] - ∂mats = (Δ[2]) # cotangent for Vector{AbstractMatrix{S}} + ∂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 = (∂mats[1]) + m = unthunk(∂mats[1]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end ∂𝐒₂_ext = if ∂mats isa Union{NoTangent, AbstractZero} NoTangent() else - m = (∂mats[2]) + m = unthunk(∂mats[2]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end ∂𝐒₃_ext = if ∂mats isa Union{NoTangent, AbstractZero} NoTangent() else - m = (∂mats[3]) + m = unthunk(∂mats[3]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end @@ -18577,21 +18577,21 @@ function rrule(::typeof(get_solution), result = (SS_and_pars[1:nVar], AbstractMatrix{S}[𝐒₁], true) pullback_1st = function (∂result_bar) - Δ = (∂result_bar) + Δ = unthunk(∂result_bar) if Δ isa Union{NoTangent, AbstractZero} return NoTangent(), NoTangent(), zeros(S, length(parameters)) end ∂NSSS = Δ[1] - ∂mats = (Δ[2]) # cotangent for Vector{AbstractMatrix{S}} + ∂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 = (∂mats[1]) + m = unthunk(∂mats[1]) m isa Union{NoTangent, AbstractZero} ? NoTangent() : m end @@ -19551,7 +19551,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = (Δ) + Δllh = unthunk(Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19576,7 +19576,7 @@ function rrule(::typeof(get_loglikelihood), d_params = @thunk begin d_𝐒₁_full_cot = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if has_override && !(d_state_red isa Union{NoTangent, AbstractZero}) + if has_override && !(d_state_red isa Union{NoTangent, AbstractZero}) d_state_full = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full_for_init) if initial_state_is_levels @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state_full @@ -19589,18 +19589,18 @@ function rrule(::typeof(get_loglikelihood), d_𝐒₁_full_cot = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) d_state_full_init = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full_for_init) - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state_full_init + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state_full_init initial_state_pullback_tangent(initial_state, d_state_full_init) - else - d_is_vec = Vector{Vector{eltype(d_state_red)}}(undef, n_overridden_components) - d_is_vec[1] = d_state_full_init - for k in 2:n_overridden_components - d_is_vec[k] = zeros(eltype(d_state_red), length(initial_state[k])) - end - initial_state_pullback_tangent(initial_state, d_is_vec) + else + d_is_vec = Vector{Vector{eltype(d_state_red)}}(undef, n_overridden_components) + d_is_vec[1] = d_state_full_init + for k in 2:n_overridden_components + d_is_vec[k] = zeros(eltype(d_state_red), length(initial_state[k])) end + initial_state_pullback_tangent(initial_state, d_is_vec) end + end else NoTangent() end @@ -19638,7 +19638,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = (Δ) + Δllh = unthunk(Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19661,45 +19661,35 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_params = @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] - end - d_state_for_ss = zeros(eltype(d_state), nVars_full) - else - d_state_for_ss = d_state - end - ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent()))[3] + d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red + d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red + d_state = zeros(eltype(d_state_red), nVars_full); @inbounds d_state[needed] .= d_state_red + d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) + @inbounds for k in eachindex(obs_indices) + d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] end - d_initial_state = if has_override - @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] - d_is = zeros(eltype(d_state), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] - initial_state_pullback_tangent(initial_state, d_is) - else - d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) - d_is_vec[1] = collect(d_state) - for k in 2:n_overridden_components - d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) - end - initial_state_pullback_tangent(initial_state, d_is_vec) + d_initial_state = NoTangent() + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] + d_is = zeros(eltype(d_state), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] + d_initial_state = d_is + else + d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) + d_is_vec[1] = collect(d_state) + for k in 2:n_overridden_components + d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) end + d_initial_state = d_is_vec end + d_state_for_ss = zeros(eltype(d_state), nVars_full) else - NoTangent() + d_state_for_ss = d_state end - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state + ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent())) + d_params = ss_grads[3] + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19739,7 +19729,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = (Δ) + Δllh = unthunk(Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19762,40 +19752,35 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_params = @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] - end - d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] - else - d_state_for_ss = d_state - end - ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent()))[3] + d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red + d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red + d_state = [zeros(eltype(d_state_red[1]), nVars_full), + zeros(eltype(d_state_red[2]), nVars_full)] + @inbounds d_state[1][needed] .= d_state_red[1] + @inbounds d_state[2][needed] .= d_state_red[2] + d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) + @inbounds for k in eachindex(obs_indices) + d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] end - d_initial_state = if has_override - @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] - d_is = zeros(eltype(d_state[1]), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] - initial_state_pullback_tangent(initial_state, d_is) - else - initial_state_pullback_tangent(initial_state, [collect(d_state[i]) for i in 1:n_overridden_components]) - end + d_initial_state = NoTangent() + if has_override + if initial_state_is_levels + # Only first-order component depends on SS (higher-order initialized to zero) + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] + d_is = zeros(eltype(d_state[1]), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] + d_initial_state = d_is + else + d_initial_state = [collect(d_state[i]) for i in 1:n_overridden_components] end + # Zero only the components the user overrode; preserve others. + d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] else - NoTangent() + d_state_for_ss = d_state end - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state + ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂], d_state_for_ss, NoTangent())) + d_params = ss_grads[3] + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19834,7 +19819,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = (Δ) + Δllh = unthunk(Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19858,47 +19843,36 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_params = @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] - end - d_state_for_ss = zeros(eltype(d_state), nVars_full) - else - d_state_for_ss = d_state - end - ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent()))[3] + d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red + d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red + d_𝐒₃ = zeros(eltype(d_𝐒₃_red), nVars_full, ncols₃); @inbounds d_𝐒₃[needed, :] .= d_𝐒₃_red + d_state = zeros(eltype(d_state_red), nVars_full); @inbounds d_state[needed] .= d_state_red + d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) + @inbounds for k in eachindex(obs_indices) + d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] end - d_initial_state = if has_override - @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] - d_is = zeros(eltype(d_state), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] - initial_state_pullback_tangent(initial_state, d_is) - else - d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) - d_is_vec[1] = collect(d_state) - for k in 2:n_overridden_components - d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) - end - initial_state_pullback_tangent(initial_state, d_is_vec) + d_initial_state = NoTangent() + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1:nVars_full_for_init] + d_is = zeros(eltype(d_state), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1:nVars_full_for_init] + d_initial_state = d_is + else + d_is_vec = Vector{Vector{eltype(d_state)}}(undef, n_overridden_components) + d_is_vec[1] = collect(d_state) + for k in 2:n_overridden_components + d_is_vec[k] = zeros(eltype(d_state), length(initial_state[k])) end + d_initial_state = d_is_vec end + d_state_for_ss = zeros(eltype(d_state), nVars_full) else - NoTangent() + d_state_for_ss = d_state end - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state + ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent())) + d_params = ss_grads[3] + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) end return isfinite(llh) ? (llh, pullback) : on_failure @@ -19953,7 +19927,7 @@ function rrule(::typeof(get_loglikelihood), if !use_workspaces; 𝓂.workspaces = orig_ws; end pullback = function (Δ) - Δllh = (Δ) + Δllh = unthunk(Δ) if Δllh isa AbstractZero return NoTangent(), NoTangent(), NoTangent(), zeros(S, nP), zero(shocks), me_std_zero_tan, NoTangent() @@ -19977,44 +19951,37 @@ function rrule(::typeof(get_loglikelihood), end d_shocks_full = expand_filter_free_shock_cotangent(d_shocks, shocks, visible_cols, n_warm) d_me_std_full = contract_filter_free_me_std_cotangent(expand_filter_free_me_std_cotangent(d_me_std, measurement_error_std, period_range), user_me_std) - d_params = @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if has_override - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] - d_state_for_ss = [zeros(eltype(d_state[1]), nVars_full), d_state[2], level_override_zeroes_third_state ? zeros(eltype(d_state[3]), nVars_full) : d_state[3]] - else - d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] - end - else - d_state_for_ss = d_state - end - ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent()))[3] + d_𝐒₁ = zeros(eltype(d_𝐒₁_red), nVars_full, ncols₁); @inbounds d_𝐒₁[needed, :] .= d_𝐒₁_red + d_𝐒₂ = zeros(eltype(d_𝐒₂_red), nVars_full, ncols₂); @inbounds d_𝐒₂[needed, :] .= d_𝐒₂_red + d_𝐒₃ = zeros(eltype(d_𝐒₃_red), nVars_full, ncols₃); @inbounds d_𝐒₃[needed, :] .= d_𝐒₃_red + d_state = [zeros(eltype(d_state_red[1]), nVars_full), + zeros(eltype(d_state_red[2]), nVars_full), + zeros(eltype(d_state_red[3]), nVars_full)] + @inbounds d_state[1][needed] .= d_state_red[1] + @inbounds d_state[2][needed] .= d_state_red[2] + @inbounds d_state[3][needed] .= d_state_red[3] + d_SS_and_pars = zeros(eltype(d_SS_obs), length(SS_and_pars)) + @inbounds for k in eachindex(obs_indices) + d_SS_and_pars[obs_indices[k]] += d_SS_obs[k] end - d_initial_state = if has_override - @thunk begin - d_𝐒₁ = expand_filter_free_matrix_cotangent(d_𝐒₁_red, needed, nVars_full, ncols₁) - d_𝐒₂ = expand_filter_free_matrix_cotangent(d_𝐒₂_red, needed, nVars_full, ncols₂) - d_𝐒₃ = expand_filter_free_matrix_cotangent(d_𝐒₃_red, needed, nVars_full, ncols₃) - d_state = expand_filter_free_state_cotangent(d_state_red, needed, nVars_full) - d_SS_and_pars = scatter_filter_free_ss_and_pars_cotangent(d_SS_obs, obs_indices, length(SS_and_pars)) - if initial_state_is_levels - @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] - d_is = zeros(eltype(d_state[1]), length(initial_state)) - @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] - initial_state_pullback_tangent(initial_state, d_is) - else - initial_state_pullback_tangent(initial_state, [collect(d_state[i]) for i in 1:n_overridden_components]) - end + d_initial_state = NoTangent() + if has_override + if initial_state_is_levels + @views d_SS_and_pars[1:nVars_full_for_init] .-= d_state[1][1:nVars_full_for_init] + d_is = zeros(eltype(d_state[1]), length(initial_state)) + @views d_is[1:nVars_full_for_init] .+= d_state[1][1:nVars_full_for_init] + d_initial_state = d_is + d_state_for_ss = [zeros(eltype(d_state[1]), nVars_full), d_state[2], level_override_zeroes_third_state ? zeros(eltype(d_state[3]), nVars_full) : d_state[3]] + else + d_initial_state = [collect(d_state[i]) for i in 1:n_overridden_components] + d_state_for_ss = [i <= n_overridden_components ? zeros(eltype(d_state[i]), nVars_full) : d_state[i] for i in 1:length(d_state)] end else - NoTangent() + d_state_for_ss = d_state end - return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, d_initial_state + ss_grads = ss_pb((NoTangent(), d_SS_and_pars, [d_𝐒₁, d_𝐒₂, d_𝐒₃], d_state_for_ss, NoTangent())) + d_params = ss_grads[3] + return NoTangent(), NoTangent(), NoTangent(), d_params, d_shocks_full, d_me_std_full, initial_state_pullback_tangent(initial_state, d_initial_state) end return isfinite(llh) ? (llh, pullback) : on_failure end @@ -20037,4 +20004,4 @@ function rrule(::typeof(get_loglikelihood), return (t[1], t[2], t[3], t[4], t[5], t[6]) end return y, pb_short -end +end \ No newline at end of file From dcff09addf3d71bf8fdda80e2e72731697f16a9a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 11 Jun 2026 13:33:39 +0000 Subject: [PATCH 11/37] Fix MethodError: unthunk ChainRulesCore.Thunk before iterating in increment_nested_initial_state_rdata! Co-authored-by: thorek1 <13523097+thorek1@users.noreply.github.com> --- ext/MooncakeExt.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/ext/MooncakeExt.jl b/ext/MooncakeExt.jl index fea87f7f7..e1a4912b7 100644 --- a/ext/MooncakeExt.jl +++ b/ext/MooncakeExt.jl @@ -187,6 +187,7 @@ function increment_nested_initial_state_rdata!( cr_tangent, ) cr_tangent isa ChainRulesCore.AbstractZero && return NoRData() + cr_tangent = ChainRulesCore.unthunk(cr_tangent) initial_state_fdata = Mooncake.tangent(initial_state_cd) @inbounds for i in eachindex(cr_tangent) component_tangent = cr_tangent[i] From 133efe17f52257c91d94d963ae5558ebd2917798 Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Fri, 12 Jun 2026 10:32:35 +0000 Subject: [PATCH 12/37] update todos --- docs/src/unfinished_docs/todo.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/src/unfinished_docs/todo.md b/docs/src/unfinished_docs/todo.md index 173396b25..1ae9c3a98 100644 --- a/docs/src/unfinished_docs/todo.md +++ b/docs/src/unfinished_docs/todo.md @@ -1,5 +1,7 @@ # Todo list ## High priority +- [ ] get_solution for third order is off as in the SSS is not on the line +- [ ] higher order solution matrix is set to zero at some point - [ ] 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. - [ ] 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 very long input types From 5749fbaab1203fb66ae90a91119e5aa9a3a7346c Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Fri, 12 Jun 2026 10:32:51 +0000 Subject: [PATCH 13/37] add BGG and GK --- ...rnanke_Gertler_Gilchrist_1999_nonlinear.jl | 124 +++++++ models/Gertler_Karadi_2011_nonlinear.jl | 345 ++++++++++++++++++ 2 files changed, 469 insertions(+) create mode 100644 models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl create mode 100644 models/Gertler_Karadi_2011_nonlinear.jl diff --git a/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl b/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl new file mode 100644 index 000000000..b2e1e7853 --- /dev/null +++ b/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl @@ -0,0 +1,124 @@ +using MacroModelling + +@model Bernanke_Gertler_Gilchrist_1999_nonlinear begin + # Household Euler equation: intertemporal optimality for deposits. + C[0]^(-1) = β * C[1]^(-1) * R[0] / Pi[1] + + # Household intratemporal condition: consumption-leisure tradeoff. + ζ * C[0] / (1 - H[0]) = W[0] + + # Labor aggregation with entrepreneurial labor share Ω. + L[0] = H[0]^(1 - Ω) + + # Final goods technology: Cobb-Douglas production. + Y[0] = A[0] * K[-1]^α * L[0]^(1 - α) + + # Household wage from marginal product of household labor. + W[0] = (1 - α) * (1 - Ω) * Y[0] / (X[0] * H[0]) + + # Entrepreneurial labor income from marginal product of entrepreneurial labor. + W_e[0] = (1 - α) * Ω * Y[0] / X[0] + + # Gross return on capital before external finance costs. + Rk[0] = (α * Y[0] / (X[0] * K[-1]) + (1 - δ) * Q[0]) / Q[-1] + + # Capital accumulation with convex investment adjustment costs. + K[0] = (1 - δ) * K[-1] + (I[0] / K[-1] - ψ_i / 2 * (I[0] / K[-1] - δ)^2) * K[-1] + + # Tobin's Q implied by the marginal adjustment cost of investment. + Q[0] = 1 / (1 - ψ_i * (I[0] / K[-1] - δ)) + + # Aggregate resource constraint including entrepreneurial consumption and monitoring costs. + Y[0] = C[0] + Ce[0] + I[0] + Gov[0] + monitoring_cost[0] + + # Monetary policy rule used in the paper's quantitative model. + R[0] = (Pi_ss / β)^(1 - ρ_R) * R[-1]^ρ_R * (Pi[0] / Pi_ss)^((1 - ρ_R) * ϕ_pi) * (Y[0] / Y[ss])^((1 - ρ_R) * ϕ_y) * exp(std_R * eps_R[x]) + + # Calvo price index. + 1.0 = θ * Pi[0]^(ϵ_p - 1) + (1 - θ) * Pi_star[0]^(1 - ϵ_p) + + # Optimal reset price from the Calvo price-setting FOC. + Pi_star[0] = ϵ_p / (ϵ_p - 1) * price_aux_1[0] / price_aux_2[0] + + # Calvo numerator recursion for the optimal reset price. + price_aux_1[0] = Y[0] / X[0] / C[0] + β * θ * Pi[1]^ϵ_p * price_aux_1[1] + + # Calvo denominator recursion for the optimal reset price. + price_aux_2[0] = Y[0] / C[0] + β * θ * Pi[1]^(ϵ_p - 1) * price_aux_2[1] + + # Aggregate technology process. + log(A[0]) = ρ_A * log(A[-1]) + std_A * eps_A[x] + + # Government spending share with persistent spending disturbance. + Gov[0] = Gov_y * Y[0] * exp(gov_gap[0]) + + # Government spending disturbance. + gov_gap[0] = ρ_G * gov_gap[-1] + std_G * eps_G[x] + + # Money demand equation included in the baseline quantitative model. + M_real[0] = b_m * C[0] / (R[0] - 1) + + # Entrepreneurial net worth before consumption, net of lender monitoring share Γ. + V[0] = (1 - Gamma[0]) * Rk[0] * Q[-1] * K[-1] + + # Entrepreneurial consumption rule with survival probability γ_e. + Ce[0] = (1 - γ_e) * V[0] + + # Entrepreneurial net worth: retained entrepreneurial wealth plus labor income. + N[0] = γ_e * V[0] + W_e[0] + + # Balance sheet identity defining leverage. + leverage[0] * N[0] = Q[0] * K[0] + + # Appendix B lognormal distribution: default probability F(ωbar). + F[0] = normcdf((log(omega_bar[0]) + σ_ω^2 / 2) / σ_ω) + + # Appendix B lognormal distribution: partial expectation G(ωbar). + G[0] = normcdf((log(omega_bar[0]) - σ_ω^2 / 2) / σ_ω) + + # Appendix B contract object Γ(ωbar) = G(ωbar) + ωbar * (1 - F(ωbar)). + Gamma[0] = G[0] + omega_bar[0] * (1 - F[0]) + + # Lender participation condition solved for the external finance premium. + finance_premium[0] = (1 - 1 / leverage[0]) / (Gamma[0] - μ * G[0]) + + # Optimal contract FOC linking leverage and the default threshold. + leverage[0] - 1 = (1 - F[0]) * (Gamma[0] - μ * G[0]) / ((1 - F[0] - μ * omega_bar[0] * normpdf((log(omega_bar[0]) + σ_ω^2 / 2) / σ_ω) / σ_ω) * (1 - Gamma[0])) + + # External finance premium wedge: expected capital return over the safe real rate. + Rk[1] = finance_premium[0] * R[0] / Pi[1] + + # Aggregate monitoring costs paid in default states. + monitoring_cost[0] = μ * G[0] * Rk[0] * Q[-1] * K[-1] + + # Observable/log-output convenience variable. + log_y[0] = log(Y[0]) + + # Annualized external finance premium convenience variable. + premium_ann[0] = 400 * log(finance_premium[0]) +end + +@parameters Bernanke_Gertler_Gilchrist_1999_nonlinear begin + β = 0.99 + α = 0.35 + δ = 0.025 + Ω = 0.015384615384615385 + ζ = 2.0 + ψ_i = 4.0 + μ = 0.12 + σ_ω = 0.28 + γ_e = 0.9728 + Pi_ss = 1.0 + θ = 0.75 + ϵ_p = 11.0 + ρ_R = 0.9 + ϕ_pi = 1.1 + ϕ_y = 0.0 + ρ_A = 0.95 + std_A = 0.01 + Gov_y = 0.20 + ρ_G = 0.95 + std_G = 0.01 + std_R = 0.0025 + b_m = 0.01 +end diff --git a/models/Gertler_Karadi_2011_nonlinear.jl b/models/Gertler_Karadi_2011_nonlinear.jl new file mode 100644 index 000000000..717e14cdb --- /dev/null +++ b/models/Gertler_Karadi_2011_nonlinear.jl @@ -0,0 +1,345 @@ +using MacroModelling + +# function GK2011_intermediary_leverage(spread, β, θ, λ) +# aa = λ * β * θ * spread +# bb = -(1 - θ) * (λ - β * spread) +# cc = 1 - θ +# discriminant = bb^2 - 4 * aa * cc + +# if discriminant < 0 +# return NaN +# end + +# return (-bb - sqrt(discriminant)) / (2 * aa) +# end + +# function GK2011_spread_gap(spread, β, θ, λ, ω) +# R = 1 / β +# φ = GK2011_intermediary_leverage(spread, β, θ, λ) +# z = spread * φ + R + +# return φ * ω / (1 - θ * z) - 1 +# end + +# function GK2011_solve_spread(β, θ, λ, ω) +# lower = 1e-8 +# lower_gap = GK2011_spread_gap(lower, β, θ, λ, ω) + +# upper = lower +# upper_gap = lower_gap +# for trial in exp.(range(log(lower * 1.01), log(0.05), length = 400)) +# trial_gap = GK2011_spread_gap(trial, β, θ, λ, ω) +# if isfinite(trial_gap) +# upper = trial +# upper_gap = trial_gap +# if lower_gap * upper_gap <= 0 +# break +# end +# end +# end + +# if lower_gap * upper_gap > 0 || !isfinite(upper_gap) +# error("Could not bracket the GK2011 steady-state intermediary spread.") +# end + +# midpoint = (lower + upper) / 2 +# for iteration in 1:200 +# midpoint = (lower + upper) / 2 +# midpoint_gap = GK2011_spread_gap(midpoint, β, θ, λ, ω) + +# if abs(midpoint_gap) < 1e-13 +# return midpoint +# elseif lower_gap * midpoint_gap <= 0 +# upper = midpoint +# upper_gap = midpoint_gap +# else +# lower = midpoint +# lower_gap = midpoint_gap +# end +# end + +# return midpoint +# end + +# function GK2011_nonlinear_steady_state!(ss, parameters) +# β, σ, h, χ, φ_l, ζ_u, θ, α, δ, g_y, η_i, ϵ_p, γ_p, γ_p_index, ρ_i, +# κ_pi, ρ_ξ, std_ξ, ρ_A, std_A, ρ_g, std_g, std_Ne, std_i, λ, ω, b_u, +# κ_y = parameters + +# Pm = (ϵ_p - 1) / ϵ_p +# X = 1 / Pm +# R = 1 / β +# spread = GK2011_solve_spread(β, θ, λ, ω) +# Rk = R + spread +# φ = GK2011_intermediary_leverage(spread, β, θ, λ) +# z = spread * φ + R +# x = z +# ν = ((1 - θ) * β * spread) / (1 - β * θ * x) +# η = (1 - θ) / (1 - β * θ * z) + +# K_to_L = (Pm * α / (Rk - 1 + δ))^(1 / (1 - α)) +# consumption_to_labor = (1 - g_y) * K_to_L^α - δ * K_to_L +# labor_constant = +# (1 - β * h) * ((1 - h) * consumption_to_labor)^(-σ) * +# Pm * (1 - α) * K_to_L^α +# L = (labor_constant / χ)^(1 / (φ_l + σ)) +# K = K_to_L * L +# Y = K^α * L^(1 - α) +# Ym = Y +# I = δ * K +# G = g_y * Y +# C = Y - I - G +# varrho = (1 - β * h) * ((1 - h) * C)^(-σ) +# Λ = 1.0 +# Q = 1.0 +# Ξ = 1.0 +# U = 1.0 +# δ_rate = δ +# In = 0.0 +# N = K / φ +# Ne = θ * z * N +# Nn = ω * K +# Keff = K +# w = Pm * (1 - α) * Y / L +# VMPK = Pm * α * Y / K +# D = 1.0 +# F_price = Y * Pm / (1 - β * γ_p) +# Z_price = Y / (1 - β * γ_p) +# Pi = 1.0 +# Pi_star = 1.0 +# i_nom = R +# prem = Rk / R +# A = 1.0 +# g_gap = 0.0 + +# if length(ss) != 38 +# resize!(ss, 38) +# end + +# ss[1] = A +# ss[2] = C +# ss[3] = D +# ss[4] = F_price +# ss[5] = G +# ss[6] = I +# ss[7] = In +# ss[8] = K +# ss[9] = Keff +# ss[10] = L +# ss[11] = Λ +# ss[12] = N +# ss[13] = Ne +# ss[14] = Nn +# ss[15] = Pi +# ss[16] = Pi_star +# ss[17] = Pm +# ss[18] = Q +# ss[19] = R +# ss[20] = Rk +# ss[21] = U +# ss[22] = VMPK +# ss[23] = X +# ss[24] = Ξ +# ss[25] = Y +# ss[26] = Ym +# ss[27] = Z_price +# ss[28] = δ_rate +# ss[29] = η +# ss[30] = g_gap +# ss[31] = i_nom +# ss[32] = ν +# ss[33] = φ +# ss[34] = prem +# ss[35] = varrho +# ss[36] = w +# ss[37] = x +# ss[38] = z + +# return ss +# end + +# function GK2011_nonlinear_steady_state(parameters) +# ss = Vector{Float64}(undef, 38) +# return GK2011_nonlinear_steady_state!(ss, parameters) +# end + + +@model Gertler_Karadi_2011_nonlinear begin + # GK replication eq. 1: marginal utility of consumption with external habit. + varrho[0] = (C[0] - h * C[-1])^(-σ) - β * h * (C[1] - h * C[0])^(-σ) + + # GK replication eq. 2: household Euler equation for the riskless asset. + β * R[0] * Lambda[1] = 1.0 + + # GK replication eq. 3: stochastic discount factor. + Lambda[0] = varrho[0] / varrho[-1] + + # GK replication eq. 4: labor market equilibrium. + χ * L[0]^φ_l = varrho[0] * Pm[0] * (1 - α) * Y[0] / L[0] + + # GK replication eq. 5: marginal value of bankers' capital. + nu[0] = (1 - θ) * β * Lambda[1] * (Rk[1] - R[0]) + β * Lambda[1] * θ * x[1] * nu[1] + + # GK replication eq. 6: marginal value of bankers' net worth. + eta[0] = (1 - θ) + β * Lambda[1] * θ * z[1] * eta[1] + + # GK replication eq. 7: incentive-constraint leverage relation. + phi[0] = eta[0] / (λ - nu[0]) + + # GK replication eq. 8: gross growth rate of bankers' capital. + z[0] = (Rk[0] - R[-1]) * phi[-1] + R[-1] + + # GK replication eq. 9: gross growth rate of bankers' net worth. + x[0] = phi[0] / phi[-1] * z[0] + + # GK replication eq. 10: aggregate intermediary balance sheet. + Q[0] * K[0] = phi[0] * N[0] + + # GK replication eq. 11: aggregate bankers' net worth. + N[0] = Ne[0] + Nn[0] + + # GK replication eq. 12: existing bankers' net worth accumulation. + Ne[0] = θ * z[0] * N[-1] * exp(-std_Ne * eps_Ne[x]) + + # GK replication eq. 13: entering bankers' net worth. + Nn[0] = ω * Q[0] * Xi[0] * K[-1] + + # GK replication eq. 14: gross return on capital. + Rk[0] = (Pm[0] * α * Ym[0] / K[-1] + Xi[0] * (Q[0] - delta_rate[0])) / Q[-1] + + # GK replication eq. 15: intermediate goods production. + Ym[0] = A[0] * (Xi[0] * U[0] * K[-1])^α * L[0]^(1 - α) + + # GK replication eq. 16: capital producer's optimal investment condition. + Q[0] = 1 + η_i / 2 * ((In[0] + I[ss]) / (In[-1] + I[ss]) - 1)^2 + + η_i * ((In[0] + I[ss]) / (In[-1] + I[ss]) - 1) * (In[0] + I[ss]) / (In[-1] + I[ss]) - + β * Lambda[1] * η_i * ((In[1] + I[ss]) / (In[0] + I[ss]) - 1) * ((In[1] + I[ss]) / (In[0] + I[ss]))^2 + + # GK replication eq. 17: utilization-dependent depreciation. + delta_rate[0] = δ + b_u / (1 + ζ_u) * (U[0]^(1 + ζ_u) - 1) + + # GK replication eq. 18: optimal capacity utilization. + Pm[0] * α * Ym[0] / U[0] = b_u * U[0]^ζ_u * Xi[0] * K[-1] + + # GK replication eq. 19: net investment. + In[0] = I[0] - delta_rate[0] * Xi[0] * K[-1] + + # GK replication eq. 20: capital accumulation. + K[0] = Xi[0] * K[-1] + In[0] + + # GK replication eq. 21: government consumption. + G[0] = g_y * Y[ss] * exp(g_gap[0]) + + # GK replication eq. 22: aggregate resource constraint. + Y[0] = C[0] + G[0] + I[0] + η_i / 2 * ((In[0] + I[ss]) / (In[-1] + I[ss]) - 1)^2 * (In[0] + I[ss]) + + # GK replication eq. 23: wholesale and retail output relation. + Ym[0] = Y[0] * D[0] + + # GK replication eq. 24: Calvo price dispersion. + D[0] = γ_p * D[-1] * Pi[-1]^(-γ_p_index * ϵ_p) * Pi[0]^ϵ_p + + (1 - γ_p) * ((1 - γ_p * Pi[-1]^(γ_p_index * (1 - γ_p)) * Pi[0]^(γ_p - 1)) / (1 - γ_p))^(-ϵ_p / (1 - γ_p)) + + # GK replication eq. 25: markup definition. + X[0] = 1 / Pm[0] + + # GK replication eq. 26: optimal reset-price numerator recursion. + F_price[0] = Y[0] * Pm[0] + β * γ_p * Lambda[1] * Pi[1]^ϵ_p * Pi[0]^(-ϵ_p * γ_p_index) * F_price[1] + + # GK replication eq. 27: optimal reset-price denominator recursion. + Z_price[0] = Y[0] + β * γ_p * Lambda[1] * Pi[1]^(ϵ_p - 1) * Pi[0]^(γ_p_index * (1 - ϵ_p)) * Z_price[1] + + # GK replication eq. 28: optimal reset-price inflation. + Pi_star[0] = ϵ_p / (ϵ_p - 1) * F_price[0] / Z_price[0] * Pi[0] + + # GK replication eq. 29: aggregate price index. + Pi[0]^(1 - ϵ_p) = γ_p * Pi[-1]^(γ_p_index * (1 - ϵ_p)) + (1 - γ_p) * Pi_star[0]^(1 - ϵ_p) + + # GK replication eq. 30: Fisher equation. + i_nom[0] = R[0] * Pi[1] + + # GK replication eq. 31: interest-rate rule. + i_nom[0] = i_nom[-1]^ρ_i * ((1 / β) * Pi[0]^κ_pi * (X[0] / (ϵ_p / (ϵ_p - 1)))^κ_y)^(1 - ρ_i) * exp(std_i * eps_i[x]) + + # GK replication eq. 32: TFP process. + A[0] = A[-1]^ρ_A * exp(-std_A * eps_A[x]) + + # GK replication eq. 33: capital quality process. + Xi[0] = Xi[-1]^ρ_ξ * exp(-std_ξ * eps_ξ[x]) + + # GK replication eq. 34: government spending process. + g_gap[0] = ρ_g * g_gap[-1] - std_g * eps_g[x] + + # GK replication eq. 35: effective capital convenience variable. + Keff[0] = Xi[0] * K[-1] + + # GK replication eq. 36: wage convenience variable. + w[0] = Pm[0] * (1 - α) * Y[0] / L[0] + + # GK replication eq. 37: marginal value product of capital convenience variable. + VMPK[0] = Pm[0] * α * Y[0] / (Xi[0] * K[-1]) + + # GK replication eq. 39: premium convenience variable. + prem[0] = Rk[1] / R[0] +end + + +# @parameters Gertler_Karadi_2011_nonlinear steady_state_function = GK2011_nonlinear_steady_state! begin +@parameters Gertler_Karadi_2011_nonlinear begin + β = 0.99 + σ = 1.0 + h = 0.815 + χ = 3.410808502195193 + φ_l = 0.276 + ζ_u = 7.2 + θ = 0.97155955 + α = 0.33 + δ = 0.025 + g_y = 0.2 + η_i = 1.728 + ϵ_p = 4.167 + γ_p = 0.779 + γ_p_index = 0.241 + ρ_i = 0.0 + κ_pi = 1.5 + κ_y = -0.5 / 4 + ρ_ξ = 0.66 + std_ξ = 0.05 + ρ_A = 0.95 + std_A = 0.01 + ρ_g = 0.95 + std_g = 0.01 + std_Ne = 0.01 + std_i = 0.01 + λ = 0.38149498593241726 + ω = 0.0022277804179292593 + b_u = 0.037601010101010155 + + β > 0 + σ > 0 + 0 <= h < 1 + χ > 0 + φ_l > 0 + ζ_u > 0 + 0 < θ < 1 + 0 < α < 1 + 0 < δ < 1 + 0 <= g_y < 1 + η_i >= 0 + ϵ_p > 1 + 0 <= γ_p < 1 + 0 <= γ_p_index < 1 + 0 <= ρ_i < 1 + κ_pi > 1 + 0 <= ρ_ξ < 1 + std_ξ >= 0 + 0 <= ρ_A < 1 + std_A >= 0 + 0 <= ρ_g < 1 + std_g >= 0 + std_Ne >= 0 + std_i >= 0 + λ > 0 + ω > 0 + b_u > 0 +end From b63812e2891efe26c440073c041e6c54439301e8 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 12 Jun 2026 17:19:38 +0200 Subject: [PATCH 14/37] fix irf for second and third order --- src/get_functions.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/get_functions.jl b/src/get_functions.jl index 863a69d2a..569cd0641 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -120,7 +120,7 @@ function adjust_initial_state(initial_state, elseif algorithm == :pruned_third_order return [initial_state - reference_steady_state[1:nVars], zeros(R, nVars) - SSS_delta, zeros(R, nVars)] else - return initial_state - reference_steady_state[1:nVars] + return initial_state - (reference_steady_state[1:nVars] + SSS_delta[1:nVars]) end end From ff2d389a4bdc3ffea475484b3aedd95f46879a4f Mon Sep 17 00:00:00 2001 From: thorek1 Date: Fri, 12 Jun 2026 17:40:11 +0200 Subject: [PATCH 15/37] test tha tno shock irf started from relevant ss remains there --- test/functionality_tests.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index 666f99aae..b8d9c0bb4 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -3834,6 +3834,17 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end end + @testset "no-shock IRF from steady state stays at steady state" begin + periods = 10 + steady_state = SS(m, algorithm = algorithm, stochastic = true, derivatives = false, return_variables_only = true) + irf = get_irf(m, algorithm = algorithm, levels = true, shocks = :none, initial_state = collect(steady_state), periods = periods) + for var in m.constants.post_model_macro.var + for t in 1:periods + @test isapprox(irf(var, t, :)[1], steady_state(var), atol = 1e-10) + end + end + end + @testset "get_non_stochastic_steady_state_residuals" begin stst = SS(m, derivatives = false) From 5dcf3b475aa31d75afc3c690a0e4637ff82c0289 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sat, 13 Jun 2026 00:42:03 +0200 Subject: [PATCH 16/37] follow up from previous get_irf fix --- src/get_functions.jl | 4 ++-- src/rrules.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/get_functions.jl b/src/get_functions.jl index 569cd0641..a2cf09c4c 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -4492,7 +4492,7 @@ function get_loglikelihood(𝓂::ℳ, nVars = 𝓂.constants.post_model_macro.nVars if initial_state isa AbstractVector{<:Real} if length(initial_state) == nVars - state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : -state + state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : zero(state) state = adjust_initial_state(initial_state, algorithm, nVars, state_shift, SS_and_pars[1:nVars]) if algorithm == :first_order state = [state] @@ -4756,7 +4756,7 @@ function get_loglikelihood(𝓂::ℳ, nVars = 𝓂.constants.post_model_macro.nVars if initial_state isa AbstractVector{<:Real} if length(initial_state) == nVars - state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : -state + state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : zero(state) state = adjust_initial_state(initial_state, algorithm, nVars, state_shift, SS_and_pars[1:nVars]) if algorithm == :first_order state = [state] diff --git a/src/rrules.jl b/src/rrules.jl index 6e24abec2..b6673f27a 100644 --- a/src/rrules.jl +++ b/src/rrules.jl @@ -1906,7 +1906,7 @@ function rrule(::typeof(get_loglikelihood), state_overridden = state if initial_state isa AbstractVector{<:Real} if length(initial_state) == nVars - state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : -state + state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : zero(state) state_overridden = adjust_initial_state(initial_state, algorithm, nVars, state_shift, SS_and_pars[1:nVars]) if algorithm == :first_order state_overridden = [state_overridden] @@ -19492,7 +19492,7 @@ function rrule(::typeof(get_loglikelihood), nVars_full_for_init = 𝓂.constants.post_model_macro.nVars if initial_state isa AbstractVector{<:Real} if length(initial_state) == nVars_full_for_init - state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : -state + state_shift = state isa AbstractVector{<:AbstractVector{<:Real}} ? (length(state) == 1 ? zero(state[1]) : -state[2]) : zero(state) state = adjust_initial_state(initial_state, algorithm, nVars_full_for_init, state_shift, SS_and_pars[1:nVars_full_for_init]) if algorithm == :first_order state = [state] From 65803eb7640e9023b5a302286af2e95fd74d1412 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sat, 13 Jun 2026 09:39:39 +0200 Subject: [PATCH 17/37] fix: iterate over steady_state keys in no-shock IRF test to skip auxiliary variables --- test/functionality_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index b8d9c0bb4..6aa3e2ed8 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -3838,7 +3838,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) periods = 10 steady_state = SS(m, algorithm = algorithm, stochastic = true, derivatives = false, return_variables_only = true) irf = get_irf(m, algorithm = algorithm, levels = true, shocks = :none, initial_state = collect(steady_state), periods = periods) - for var in m.constants.post_model_macro.var + for var in axiskeys(steady_state, 1) for t in 1:periods @test isapprox(irf(var, t, :)[1], steady_state(var), atol = 1e-10) end From 0e0d2ddd079aef7480fb27d4ea7b59a82baeccc1 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sat, 13 Jun 2026 17:14:30 +0200 Subject: [PATCH 18/37] =?UTF-8?q?fix:=20use=20full=20S=E2=82=81=20matrix?= =?UTF-8?q?=20in=20state=5Fupdate=20for=20higher-order=20IRF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The state_update in parse_algorithm_to_state_update was zeroing out the constant column of S₁: [S₁[:,1:nPast] zeros(nVars) S₁[:,nPast+1:end]]. Since the augmented state includes 1 as the constant, this dropped the constant contribution from the first-order term, breaking the SSS fixed-point property (no-shock IRF would not stay at SSS). --- src/MacroModelling.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index a8aff7acd..e9716bc77 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -2471,32 +2471,29 @@ end 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 + return S₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 end else # :third_order 𝐒₃ = 𝓂.caches.third_order_solution * 𝓂.constants.third_order.𝐔₃ state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S} aug_state = [state[past_idx]; 1; shock] - return 𝐒₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 𝐒₃ * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 + return S₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 𝐒₃ * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 end end elseif algorithm == :pruned_second_order 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, 𝐒₁, 𝐒₂) + state_update = (state, shock) -> pruned_second_order_state_update(state, shock, past_idx, nVars, S₁, 𝐒₂) 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, 𝐒₁, 𝐒₂, 𝐒₃) + state_update = (state, shock) -> pruned_third_order_state_update(state, shock, past_idx, nVars, S₁, 𝐒₂, 𝐒₃) end end From 2fc12309814d5b7ef749405102349f2390d5a717 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sat, 13 Jun 2026 17:31:49 +0200 Subject: [PATCH 19/37] =?UTF-8?q?Revert=20"fix:=20use=20full=20S=E2=82=81?= =?UTF-8?q?=20matrix=20in=20state=5Fupdate=20for=20higher-order=20IRF"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 0e0d2ddd079aef7480fb27d4ea7b59a82baeccc1. --- src/MacroModelling.jl | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index e9716bc77..a8aff7acd 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -2471,29 +2471,32 @@ end 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 S₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 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 S₁ * aug_state + 𝐒₂ * ℒ.kron(aug_state, aug_state) / 2 + 𝐒₃ * ℒ.kron(ℒ.kron(aug_state,aug_state),aug_state) / 6 + 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, S₁, 𝐒₂) + 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, S₁, 𝐒₂, 𝐒₃) + state_update = (state, shock) -> pruned_third_order_state_update(state, shock, past_idx, nVars, 𝐒₁, 𝐒₂, 𝐒₃) end end From 62a31baf5b1fda9497a4a4c5c1f3ba60813783a7 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sat, 13 Jun 2026 17:52:40 +0200 Subject: [PATCH 20/37] Fix first-order no-shock IRF steady-state test --- test/functionality_tests.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index 6aa3e2ed8..b5b78aabf 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -3836,7 +3836,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) @testset "no-shock IRF from steady state stays at steady state" begin periods = 10 - steady_state = SS(m, algorithm = algorithm, stochastic = true, derivatives = false, return_variables_only = true) + steady_state = SS(m, algorithm = algorithm, stochastic = algorithm != :first_order, derivatives = false, return_variables_only = true) irf = get_irf(m, algorithm = algorithm, levels = true, shocks = :none, initial_state = collect(steady_state), periods = periods) for var in axiskeys(steady_state, 1) for t in 1:periods @@ -4042,4 +4042,4 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) @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 +end From 989708004980419d3fa5d311d8934d0f1f5af20c Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 14 Jun 2026 00:17:16 +0200 Subject: [PATCH 21/37] Fix no-shock IRF test: only test variables in both steady state and IRF output --- test/functionality_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index b5b78aabf..e0b65b544 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -3838,7 +3838,7 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) periods = 10 steady_state = SS(m, algorithm = algorithm, stochastic = algorithm != :first_order, derivatives = false, return_variables_only = true) irf = get_irf(m, algorithm = algorithm, levels = true, shocks = :none, initial_state = collect(steady_state), periods = periods) - for var in axiskeys(steady_state, 1) + for var in intersect(axiskeys(steady_state, 1), axiskeys(irf, 1)) for t in 1:periods @test isapprox(irf(var, t, :)[1], steady_state(var), atol = 1e-10) end From c22e4bafb01c05e0110d83667b7a290519f7c032 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 14 Jun 2026 17:13:10 +0200 Subject: [PATCH 22/37] update todos --- docs/src/unfinished_docs/todo.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/src/unfinished_docs/todo.md b/docs/src/unfinished_docs/todo.md index 1ae9c3a98..173396b25 100644 --- a/docs/src/unfinished_docs/todo.md +++ b/docs/src/unfinished_docs/todo.md @@ -1,7 +1,5 @@ # Todo list ## High priority -- [ ] get_solution for third order is off as in the SSS is not on the line -- [ ] higher order solution matrix is set to zero at some point - [ ] 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. - [ ] 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 very long input types From a5e1bed3fa90b57181c9440f0b98c74884d32a6f Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Mon, 15 Jun 2026 14:15:33 +0000 Subject: [PATCH 23/37] clean up bgg and gk model notation --- ...rnanke_Gertler_Gilchrist_1999_nonlinear.jl | 141 +++---- models/Gertler_Karadi_2011_nonlinear.jl | 356 ++++-------------- src/steady_state/solver_parameters.jl | 4 +- 3 files changed, 155 insertions(+), 346 deletions(-) diff --git a/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl b/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl index b2e1e7853..f6a4238b4 100644 --- a/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl +++ b/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl @@ -1,101 +1,101 @@ using MacroModelling @model Bernanke_Gertler_Gilchrist_1999_nonlinear begin - # Household Euler equation: intertemporal optimality for deposits. - C[0]^(-1) = β * C[1]^(-1) * R[0] / Pi[1] + # [B.3] Household Euler equation: intertemporal optimality for deposits. + C[0]^(-1) = β * C[1]^(-1) * R[0] / π[1] - # Household intratemporal condition: consumption-leisure tradeoff. + # [B.4] Household intratemporal condition: consumption-leisure tradeoff. ζ * C[0] / (1 - H[0]) = W[0] - # Labor aggregation with entrepreneurial labor share Ω. - L[0] = H[0]^(1 - Ω) + # [4.6] Labor aggregation with entrepreneurial labor share Ω (Hᵉ normalized to 1). + L[0] = H[0]^(1 - Ω) * Hᵉ^Ω - # Final goods technology: Cobb-Douglas production. + # [4.1] Final goods technology: Cobb-Douglas production. Y[0] = A[0] * K[-1]^α * L[0]^(1 - α) - # Household wage from marginal product of household labor. + # [4.11] Household labor demand: wage equals marginal product. W[0] = (1 - α) * (1 - Ω) * Y[0] / (X[0] * H[0]) - # Entrepreneurial labor income from marginal product of entrepreneurial labor. - W_e[0] = (1 - α) * Ω * Y[0] / X[0] + # [4.12] Entrepreneurial labor income from marginal product. + Wᵉ[0] = (1 - α) * Ω * Y[0] / X[0] - # Gross return on capital before external finance costs. - Rk[0] = (α * Y[0] / (X[0] * K[-1]) + (1 - δ) * Q[0]) / Q[-1] + # [4.4] Gross return on capital (Rᵏ = R^k in the paper). + Rᵏ[0] = (α * Y[0] / (X[0] * K[-1]) + (1 - δ) * Q[0]) / Q[-1] - # Capital accumulation with convex investment adjustment costs. - K[0] = (1 - δ) * K[-1] + (I[0] / K[-1] - ψ_i / 2 * (I[0] / K[-1] - δ)^2) * K[-1] + # [4.2] Capital accumulation with convex investment adjustment costs. + K[0] = (1 - δ) * K[-1] + (I[0] / K[-1] - ψⁱ / 2 * (I[0] / K[-1] - δ)^2) * K[-1] - # Tobin's Q implied by the marginal adjustment cost of investment. - Q[0] = 1 / (1 - ψ_i * (I[0] / K[-1] - δ)) + # [4.3] Tobin's Q implied by the marginal adjustment cost of investment. + Q[0] = 1 / (1 - ψⁱ * (I[0] / K[-1] - δ)) - # Aggregate resource constraint including entrepreneurial consumption and monitoring costs. - Y[0] = C[0] + Ce[0] + I[0] + Gov[0] + monitoring_cost[0] + # [B.8] Aggregate resource constraint including entrepreneurial consumption and monitoring costs. + Y[0] = C[0] + Cᵉ[0] + I[0] + Gov[0] + monitoring_cost[0] - # Monetary policy rule used in the paper's quantitative model. - R[0] = (Pi_ss / β)^(1 - ρ_R) * R[-1]^ρ_R * (Pi[0] / Pi_ss)^((1 - ρ_R) * ϕ_pi) * (Y[0] / Y[ss])^((1 - ρ_R) * ϕ_y) * exp(std_R * eps_R[x]) + # [4.25] Monetary policy rule used in the paper's quantitative model. + R[0] = (πˢˢ / β)^(1 - ρʳ) * R[-1]^ρʳ * (π[0] / πˢˢ)^((1 - ρʳ) * ϕ_pi) * (Y[0] / Y[ss])^((1 - ρʳ) * ϕʸ) * exp(σʳ * εʳ[x]) - # Calvo price index. - 1.0 = θ * Pi[0]^(ϵ_p - 1) + (1 - θ) * Pi_star[0]^(1 - ϵ_p) + # [B.12] Calvo price index. + 1.0 = θ * π[0]^(ϵᵖ - 1) + (1 - θ) * πstar[0]^(1 - ϵᵖ) - # Optimal reset price from the Calvo price-setting FOC. - Pi_star[0] = ϵ_p / (ϵ_p - 1) * price_aux_1[0] / price_aux_2[0] + # [B.11] Optimal reset price from the Calvo price-setting FOC. + πstar[0] = ϵᵖ / (ϵᵖ - 1) * price_aux_1[0] / price_aux_2[0] - # Calvo numerator recursion for the optimal reset price. - price_aux_1[0] = Y[0] / X[0] / C[0] + β * θ * Pi[1]^ϵ_p * price_aux_1[1] + # [B.11a] Calvo numerator recursion. + price_aux_1[0] = Y[0] / X[0] / C[0] + β * θ * π[1]^ϵᵖ * price_aux_1[1] - # Calvo denominator recursion for the optimal reset price. - price_aux_2[0] = Y[0] / C[0] + β * θ * Pi[1]^(ϵ_p - 1) * price_aux_2[1] + # [B.11b] Calvo denominator recursion. + price_aux_2[0] = Y[0] / C[0] + β * θ * π[1]^(ϵᵖ - 1) * price_aux_2[1] - # Aggregate technology process. - log(A[0]) = ρ_A * log(A[-1]) + std_A * eps_A[x] + # [4.27] Aggregate technology process. + log(A[0]) = ρᵃ * log(A[-1]) + σᵃ * εᵃ[x] - # Government spending share with persistent spending disturbance. - Gov[0] = Gov_y * Y[0] * exp(gov_gap[0]) + # [4.26] Government spending level (level form from the log-linearized process). + Gov[0] = Govʸ * Y[0] * exp(gov_gap[0]) - # Government spending disturbance. - gov_gap[0] = ρ_G * gov_gap[-1] + std_G * eps_G[x] + # [4.26] Government spending AR(1) disturbance. + gov_gap[0] = ρᵍ * gov_gap[-1] + σᵍ * εᵍ[x] - # Money demand equation included in the baseline quantitative model. - M_real[0] = b_m * C[0] / (R[0] - 1) + # [B.5] Money demand equation. + M_real[0] = χ * C[0] * R[0] / (R[0] - 1) - # Entrepreneurial net worth before consumption, net of lender monitoring share Γ. - V[0] = (1 - Gamma[0]) * Rk[0] * Q[-1] * K[-1] + # [4.8] Entrepreneurial equity (net of lender repayment and monitoring). + V[0] = (1 - Γ[0]) * Rᵏ[0] * Q[-1] * K[-1] - # Entrepreneurial consumption rule with survival probability γ_e. - Ce[0] = (1 - γ_e) * V[0] + # [4.8] Entrepreneurial consumption: dying entrepreneurs consume their equity. + Cᵉ[0] = (1 - γᵉ) * V[0] - # Entrepreneurial net worth: retained entrepreneurial wealth plus labor income. - N[0] = γ_e * V[0] + W_e[0] + # [4.7] Entrepreneurial net worth: retained equity plus labor income. + N[0] = γᵉ * V[0] + Wᵉ[0] - # Balance sheet identity defining leverage. + # [3.2] Balance sheet identity defining leverage. leverage[0] * N[0] = Q[0] * K[0] - # Appendix B lognormal distribution: default probability F(ωbar). - F[0] = normcdf((log(omega_bar[0]) + σ_ω^2 / 2) / σ_ω) + # [A.2] Lognormal distribution: default probability F(ω̄). + F[0] = normcdf((log(ω̄[0]) + σ_ω^2 / 2) / σ_ω) - # Appendix B lognormal distribution: partial expectation G(ωbar). - G[0] = normcdf((log(omega_bar[0]) - σ_ω^2 / 2) / σ_ω) + # [A.2] Lognormal distribution: partial expectation G(ω̄). + G[0] = normcdf((log(ω̄[0]) - σ_ω^2 / 2) / σ_ω) - # Appendix B contract object Γ(ωbar) = G(ωbar) + ωbar * (1 - F(ωbar)). - Gamma[0] = G[0] + omega_bar[0] * (1 - F[0]) + # [A.2] Contract object Γ(ω̄) = G(ω̄) + ω̄(1 - F(ω̄)). + Γ[0] = G[0] + ω̄[0] * (1 - F[0]) - # Lender participation condition solved for the external finance premium. - finance_premium[0] = (1 - 1 / leverage[0]) / (Gamma[0] - μ * G[0]) + # [3.5] Lender participation condition (inverted for external finance premium s). + s[0] = (1 - 1 / leverage[0]) / (Γ[0] - μ * G[0]) - # Optimal contract FOC linking leverage and the default threshold. - leverage[0] - 1 = (1 - F[0]) * (Gamma[0] - μ * G[0]) / ((1 - F[0] - μ * omega_bar[0] * normpdf((log(omega_bar[0]) + σ_ω^2 / 2) / σ_ω) / σ_ω) * (1 - Gamma[0])) + # [A.1]–[A.2] Optimal contract FOC linking leverage and the default threshold. + leverage[0] - 1 = (1 - F[0]) * (Γ[0] - μ * G[0]) / ((1 - F[0] - μ * ω̄[0] * normpdf((log(ω̄[0]) + σ_ω^2 / 2) / σ_ω) / σ_ω) * (1 - Γ[0])) - # External finance premium wedge: expected capital return over the safe real rate. - Rk[1] = finance_premium[0] * R[0] / Pi[1] + # [4.5] External finance premium wedge (ex-post). + Rᵏ[1] = s[0] * R[0] / π[1] - # Aggregate monitoring costs paid in default states. - monitoring_cost[0] = μ * G[0] * Rk[0] * Q[-1] * K[-1] + # [B.8] Aggregate monitoring costs paid in default states. + monitoring_cost[0] = μ * G[0] * Rᵏ[0] * Q[-1] * K[-1] # Observable/log-output convenience variable. log_y[0] = log(Y[0]) # Annualized external finance premium convenience variable. - premium_ann[0] = 400 * log(finance_premium[0]) + premium_ann[0] = 400 * log(s[0]) end @parameters Bernanke_Gertler_Gilchrist_1999_nonlinear begin @@ -103,22 +103,23 @@ end α = 0.35 δ = 0.025 Ω = 0.015384615384615385 + Hᵉ = 1.0 ζ = 2.0 - ψ_i = 4.0 + ψⁱ = 4.0 μ = 0.12 σ_ω = 0.28 - γ_e = 0.9728 - Pi_ss = 1.0 + γᵉ = 0.9728 + πˢˢ = 1.0 θ = 0.75 - ϵ_p = 11.0 - ρ_R = 0.9 + ϵᵖ = 11.0 + ρʳ = 0.9 ϕ_pi = 1.1 - ϕ_y = 0.0 - ρ_A = 0.95 - std_A = 0.01 - Gov_y = 0.20 - ρ_G = 0.95 - std_G = 0.01 - std_R = 0.0025 - b_m = 0.01 + ϕʸ = 0.0 + ρᵃ = 0.95 + σᵃ = 0.01 + Govʸ = 0.20 + ρᵍ = 0.95 + σᵍ = 0.01 + σʳ = 0.0025 + χ = 0.0099 end diff --git a/models/Gertler_Karadi_2011_nonlinear.jl b/models/Gertler_Karadi_2011_nonlinear.jl index 717e14cdb..0e12689ca 100644 --- a/models/Gertler_Karadi_2011_nonlinear.jl +++ b/models/Gertler_Karadi_2011_nonlinear.jl @@ -1,290 +1,124 @@ using MacroModelling -# function GK2011_intermediary_leverage(spread, β, θ, λ) -# aa = λ * β * θ * spread -# bb = -(1 - θ) * (λ - β * spread) -# cc = 1 - θ -# discriminant = bb^2 - 4 * aa * cc - -# if discriminant < 0 -# return NaN -# end - -# return (-bb - sqrt(discriminant)) / (2 * aa) -# end - -# function GK2011_spread_gap(spread, β, θ, λ, ω) -# R = 1 / β -# φ = GK2011_intermediary_leverage(spread, β, θ, λ) -# z = spread * φ + R - -# return φ * ω / (1 - θ * z) - 1 -# end - -# function GK2011_solve_spread(β, θ, λ, ω) -# lower = 1e-8 -# lower_gap = GK2011_spread_gap(lower, β, θ, λ, ω) - -# upper = lower -# upper_gap = lower_gap -# for trial in exp.(range(log(lower * 1.01), log(0.05), length = 400)) -# trial_gap = GK2011_spread_gap(trial, β, θ, λ, ω) -# if isfinite(trial_gap) -# upper = trial -# upper_gap = trial_gap -# if lower_gap * upper_gap <= 0 -# break -# end -# end -# end - -# if lower_gap * upper_gap > 0 || !isfinite(upper_gap) -# error("Could not bracket the GK2011 steady-state intermediary spread.") -# end - -# midpoint = (lower + upper) / 2 -# for iteration in 1:200 -# midpoint = (lower + upper) / 2 -# midpoint_gap = GK2011_spread_gap(midpoint, β, θ, λ, ω) - -# if abs(midpoint_gap) < 1e-13 -# return midpoint -# elseif lower_gap * midpoint_gap <= 0 -# upper = midpoint -# upper_gap = midpoint_gap -# else -# lower = midpoint -# lower_gap = midpoint_gap -# end -# end - -# return midpoint -# end - -# function GK2011_nonlinear_steady_state!(ss, parameters) -# β, σ, h, χ, φ_l, ζ_u, θ, α, δ, g_y, η_i, ϵ_p, γ_p, γ_p_index, ρ_i, -# κ_pi, ρ_ξ, std_ξ, ρ_A, std_A, ρ_g, std_g, std_Ne, std_i, λ, ω, b_u, -# κ_y = parameters - -# Pm = (ϵ_p - 1) / ϵ_p -# X = 1 / Pm -# R = 1 / β -# spread = GK2011_solve_spread(β, θ, λ, ω) -# Rk = R + spread -# φ = GK2011_intermediary_leverage(spread, β, θ, λ) -# z = spread * φ + R -# x = z -# ν = ((1 - θ) * β * spread) / (1 - β * θ * x) -# η = (1 - θ) / (1 - β * θ * z) - -# K_to_L = (Pm * α / (Rk - 1 + δ))^(1 / (1 - α)) -# consumption_to_labor = (1 - g_y) * K_to_L^α - δ * K_to_L -# labor_constant = -# (1 - β * h) * ((1 - h) * consumption_to_labor)^(-σ) * -# Pm * (1 - α) * K_to_L^α -# L = (labor_constant / χ)^(1 / (φ_l + σ)) -# K = K_to_L * L -# Y = K^α * L^(1 - α) -# Ym = Y -# I = δ * K -# G = g_y * Y -# C = Y - I - G -# varrho = (1 - β * h) * ((1 - h) * C)^(-σ) -# Λ = 1.0 -# Q = 1.0 -# Ξ = 1.0 -# U = 1.0 -# δ_rate = δ -# In = 0.0 -# N = K / φ -# Ne = θ * z * N -# Nn = ω * K -# Keff = K -# w = Pm * (1 - α) * Y / L -# VMPK = Pm * α * Y / K -# D = 1.0 -# F_price = Y * Pm / (1 - β * γ_p) -# Z_price = Y / (1 - β * γ_p) -# Pi = 1.0 -# Pi_star = 1.0 -# i_nom = R -# prem = Rk / R -# A = 1.0 -# g_gap = 0.0 - -# if length(ss) != 38 -# resize!(ss, 38) -# end - -# ss[1] = A -# ss[2] = C -# ss[3] = D -# ss[4] = F_price -# ss[5] = G -# ss[6] = I -# ss[7] = In -# ss[8] = K -# ss[9] = Keff -# ss[10] = L -# ss[11] = Λ -# ss[12] = N -# ss[13] = Ne -# ss[14] = Nn -# ss[15] = Pi -# ss[16] = Pi_star -# ss[17] = Pm -# ss[18] = Q -# ss[19] = R -# ss[20] = Rk -# ss[21] = U -# ss[22] = VMPK -# ss[23] = X -# ss[24] = Ξ -# ss[25] = Y -# ss[26] = Ym -# ss[27] = Z_price -# ss[28] = δ_rate -# ss[29] = η -# ss[30] = g_gap -# ss[31] = i_nom -# ss[32] = ν -# ss[33] = φ -# ss[34] = prem -# ss[35] = varrho -# ss[36] = w -# ss[37] = x -# ss[38] = z - -# return ss -# end - -# function GK2011_nonlinear_steady_state(parameters) -# ss = Vector{Float64}(undef, 38) -# return GK2011_nonlinear_steady_state!(ss, parameters) -# end - - @model Gertler_Karadi_2011_nonlinear begin - # GK replication eq. 1: marginal utility of consumption with external habit. - varrho[0] = (C[0] - h * C[-1])^(-σ) - β * h * (C[1] - h * C[0])^(-σ) + # [4] Marginal utility of consumption with external habit. + ϱ[0] = (C[0] - h * C[-1])^(-σ) - β * h * (C[1] - h * C[0])^(-σ) - # GK replication eq. 2: household Euler equation for the riskless asset. - β * R[0] * Lambda[1] = 1.0 + # [4'] Household Euler equation for the riskless asset. + β * R[0] * Λ[1] = 1.0 - # GK replication eq. 3: stochastic discount factor. - Lambda[0] = varrho[0] / varrho[-1] + # [--] Stochastic discount factor. + Λ[0] = ϱ[0] / ϱ[-1] - # GK replication eq. 4: labor market equilibrium. - χ * L[0]^φ_l = varrho[0] * Pm[0] * (1 - α) * Y[0] / L[0] + # [3] Labor market equilibrium. + χ * L[0]^φ_l = ϱ[0] * Pₘ[0] * (1 - α) * Y[0] / L[0] - # GK replication eq. 5: marginal value of bankers' capital. - nu[0] = (1 - θ) * β * Lambda[1] * (Rk[1] - R[0]) + β * Lambda[1] * θ * x[1] * nu[1] + # [11] Marginal value of bankers' capital (ν_t). + ν[0] = (1 - θ) * β * Λ[1] * (Rᵏ[1] - R[0]) + β * Λ[1] * θ * x[1] * ν[1] - # GK replication eq. 6: marginal value of bankers' net worth. - eta[0] = (1 - θ) + β * Lambda[1] * θ * z[1] * eta[1] + # [11] Marginal value of bankers' net worth (η_t). + η[0] = (1 - θ) + β * Λ[1] * θ * z[1] * η[1] - # GK replication eq. 7: incentive-constraint leverage relation. - phi[0] = eta[0] / (λ - nu[0]) + # [13] Incentive-constraint leverage relation (φ_t = η_t/(λ - ν_t)). + φ[0] = η[0] / (λ - ν[0]) - # GK replication eq. 8: gross growth rate of bankers' capital. - z[0] = (Rk[0] - R[-1]) * phi[-1] + R[-1] + # [14] Gross growth rate of bankers' capital (z_{t-1,t} = N_t/N_{t-1}). + z[0] = (Rᵏ[0] - R[-1]) * φ[-1] + R[-1] - # GK replication eq. 9: gross growth rate of bankers' net worth. - x[0] = phi[0] / phi[-1] * z[0] + # [--] Gross growth rate of assets (x_{t-1,t} = Q_t S_t/(Q_{t-1} S_{t-1})). + x[0] = φ[0] / φ[-1] * z[0] - # GK replication eq. 10: aggregate intermediary balance sheet. - Q[0] * K[0] = phi[0] * N[0] + # [15] Aggregate intermediary balance sheet. + Q[0] * K[0] = φ[0] * N[0] - # GK replication eq. 11: aggregate bankers' net worth. - N[0] = Ne[0] + Nn[0] + # [16] Aggregate bankers' net worth (survivors + entrants). + N[0] = Nᵉ[0] + Nⁿ[0] - # GK replication eq. 12: existing bankers' net worth accumulation. - Ne[0] = θ * z[0] * N[-1] * exp(-std_Ne * eps_Ne[x]) + # [17] Existing bankers' net worth accumulation. + Nᵉ[0] = θ * z[0] * N[-1] * exp(-σ_Ne * ε_Ne[x]) - # GK replication eq. 13: entering bankers' net worth. - Nn[0] = ω * Q[0] * Xi[0] * K[-1] + # [18] Entering bankers' net worth. + Nⁿ[0] = ω * Q[0] * ξ[0] * K[-1] - # GK replication eq. 14: gross return on capital. - Rk[0] = (Pm[0] * α * Ym[0] / K[-1] + Xi[0] * (Q[0] - delta_rate[0])) / Q[-1] + # [25] Gross return on capital. + Rᵏ[0] = (Pₘ[0] * α * Ym[0] / K[-1] + ξ[0] * (Q[0] - δ_rate[0])) / Q[-1] - # GK replication eq. 15: intermediate goods production. - Ym[0] = A[0] * (Xi[0] * U[0] * K[-1])^α * L[0]^(1 - α) + # [22] Intermediate goods production. + Ym[0] = A[0] * (ξ[0] * U[0] * K[-1])^α * L[0]^(1 - α) - # GK replication eq. 16: capital producer's optimal investment condition. + # [27] Capital producer's optimal investment condition (Tobin's Q). Q[0] = 1 + η_i / 2 * ((In[0] + I[ss]) / (In[-1] + I[ss]) - 1)^2 + η_i * ((In[0] + I[ss]) / (In[-1] + I[ss]) - 1) * (In[0] + I[ss]) / (In[-1] + I[ss]) - - β * Lambda[1] * η_i * ((In[1] + I[ss]) / (In[0] + I[ss]) - 1) * ((In[1] + I[ss]) / (In[0] + I[ss]))^2 + β * Λ[1] * η_i * ((In[1] + I[ss]) / (In[0] + I[ss]) - 1) * ((In[1] + I[ss]) / (In[0] + I[ss]))^2 - # GK replication eq. 17: utilization-dependent depreciation. - delta_rate[0] = δ + b_u / (1 + ζ_u) * (U[0]^(1 + ζ_u) - 1) + # [--] Utilization-dependent depreciation. + δ_rate[0] = δ + b_u / (1 + ζ_u) * (U[0]^(1 + ζ_u) - 1) - # GK replication eq. 18: optimal capacity utilization. - Pm[0] * α * Ym[0] / U[0] = b_u * U[0]^ζ_u * Xi[0] * K[-1] + # [23] Optimal capacity utilization. + Pₘ[0] * α * Ym[0] / U[0] = b_u * U[0]^ζ_u * ξ[0] * K[-1] - # GK replication eq. 19: net investment. - In[0] = I[0] - delta_rate[0] * Xi[0] * K[-1] + # [26] Net investment. + In[0] = I[0] - δ_rate[0] * ξ[0] * K[-1] - # GK replication eq. 20: capital accumulation. - K[0] = Xi[0] * K[-1] + In[0] + # [35] Capital accumulation. + K[0] = ξ[0] * K[-1] + In[0] - # GK replication eq. 21: government consumption. - G[0] = g_y * Y[ss] * exp(g_gap[0]) + # [--] Government consumption. + G[0] = g_y * Y[ss] * exp(ĝ[0]) - # GK replication eq. 22: aggregate resource constraint. + # [34] Aggregate resource constraint. Y[0] = C[0] + G[0] + I[0] + η_i / 2 * ((In[0] + I[ss]) / (In[-1] + I[ss]) - 1)^2 * (In[0] + I[ss]) - # GK replication eq. 23: wholesale and retail output relation. + # [--] Wholesale and retail output relation. Ym[0] = Y[0] * D[0] - # GK replication eq. 24: Calvo price dispersion. - D[0] = γ_p * D[-1] * Pi[-1]^(-γ_p_index * ϵ_p) * Pi[0]^ϵ_p + - (1 - γ_p) * ((1 - γ_p * Pi[-1]^(γ_p_index * (1 - γ_p)) * Pi[0]^(γ_p - 1)) / (1 - γ_p))^(-ϵ_p / (1 - γ_p)) + # [--] Calvo price dispersion. + D[0] = γ_p * D[-1] * π[-1]^(-γ_p_index * ϵ_p) * π[0]^ϵ_p + + (1 - γ_p) * ((1 - γ_p * π[-1]^(γ_p_index * (1 - ϵ_p)) * π[0]^(ϵ_p - 1)) / (1 - γ_p))^(-ϵ_p / (1 - ϵ_p)) - # GK replication eq. 25: markup definition. - X[0] = 1 / Pm[0] + # [--] Markup definition. + X[0] = 1 / Pₘ[0] - # GK replication eq. 26: optimal reset-price numerator recursion. - F_price[0] = Y[0] * Pm[0] + β * γ_p * Lambda[1] * Pi[1]^ϵ_p * Pi[0]^(-ϵ_p * γ_p_index) * F_price[1] + # [31] Optimal reset-price numerator recursion (F_t). + F[0] = Y[0] * Pₘ[0] + β * γ_p * Λ[1] * π[1]^ϵ_p * π[0]^(-ϵ_p * γ_p_index) * F[1] - # GK replication eq. 27: optimal reset-price denominator recursion. - Z_price[0] = Y[0] + β * γ_p * Lambda[1] * Pi[1]^(ϵ_p - 1) * Pi[0]^(γ_p_index * (1 - ϵ_p)) * Z_price[1] + # [31] Optimal reset-price denominator recursion (Z_t). + Z[0] = Y[0] + β * γ_p * Λ[1] * π[1]^(ϵ_p - 1) * π[0]^(γ_p_index * (1 - ϵ_p)) * Z[1] - # GK replication eq. 28: optimal reset-price inflation. - Pi_star[0] = ϵ_p / (ϵ_p - 1) * F_price[0] / Z_price[0] * Pi[0] + # [31]–[32] Optimal reset-price inflation. + π_star[0] = ϵ_p / (ϵ_p - 1) * F[0] / Z[0] * π[0] - # GK replication eq. 29: aggregate price index. - Pi[0]^(1 - ϵ_p) = γ_p * Pi[-1]^(γ_p_index * (1 - ϵ_p)) + (1 - γ_p) * Pi_star[0]^(1 - ϵ_p) + # [33] Aggregate price index. + π[0]^(1 - ϵ_p) = γ_p * π[-1]^(γ_p_index * (1 - ϵ_p)) + (1 - γ_p) * π_star[0]^(1 - ϵ_p) - # GK replication eq. 30: Fisher equation. - i_nom[0] = R[0] * Pi[1] + # [38] Fisher equation. + i[0] = R[0] * π[1] - # GK replication eq. 31: interest-rate rule. - i_nom[0] = i_nom[-1]^ρ_i * ((1 / β) * Pi[0]^κ_pi * (X[0] / (ϵ_p / (ϵ_p - 1)))^κ_y)^(1 - ρ_i) * exp(std_i * eps_i[x]) + # [37] Interest-rate rule (Taylor rule with smoothing). + i[0] = i[-1]^ρ_i * ((1 / β) * π[0]^κ_pi * (X[0] / (ϵ_p / (ϵ_p - 1)))^κ_y)^(1 - ρ_i) * exp(σ_i * ε_i[x]) - # GK replication eq. 32: TFP process. - A[0] = A[-1]^ρ_A * exp(-std_A * eps_A[x]) + # [--] TFP process. + A[0] = A[-1]^ρ_A * exp(-σ_A * ε_A[x]) - # GK replication eq. 33: capital quality process. - Xi[0] = Xi[-1]^ρ_ξ * exp(-std_ξ * eps_ξ[x]) + # [--] Capital quality process. + ξ[0] = ξ[-1]^ρ_ξ * exp(-σ_ξ * ε_ξ[x]) - # GK replication eq. 34: government spending process. - g_gap[0] = ρ_g * g_gap[-1] - std_g * eps_g[x] + # [--] Government spending process. + ĝ[0] = ρ_g * ĝ[-1] - σ_g * ε_g[x] - # GK replication eq. 35: effective capital convenience variable. - Keff[0] = Xi[0] * K[-1] + # [--] Effective capital convenience variable. + Keff[0] = ξ[0] * K[-1] - # GK replication eq. 36: wage convenience variable. - w[0] = Pm[0] * (1 - α) * Y[0] / L[0] + # [--] Wage convenience variable. + w[0] = Pₘ[0] * (1 - α) * Y[0] / L[0] - # GK replication eq. 37: marginal value product of capital convenience variable. - VMPK[0] = Pm[0] * α * Y[0] / (Xi[0] * K[-1]) + # [--] Marginal value product of capital convenience variable. + VMPK[0] = Pₘ[0] * α * Y[0] / (ξ[0] * K[-1]) - # GK replication eq. 39: premium convenience variable. - prem[0] = Rk[1] / R[0] + # [--] Premium convenience variable (Rᵏ_{t+1}/R_{t+1}). + prem[0] = Rᵏ[1] / R[0] end - -# @parameters Gertler_Karadi_2011_nonlinear steady_state_function = GK2011_nonlinear_steady_state! begin @parameters Gertler_Karadi_2011_nonlinear begin β = 0.99 σ = 1.0 @@ -300,46 +134,18 @@ end ϵ_p = 4.167 γ_p = 0.779 γ_p_index = 0.241 - ρ_i = 0.0 + ρ_i = 0.8 κ_pi = 1.5 κ_y = -0.5 / 4 ρ_ξ = 0.66 - std_ξ = 0.05 + σ_ξ = 0.05 ρ_A = 0.95 - std_A = 0.01 + σ_A = 0.01 ρ_g = 0.95 - std_g = 0.01 - std_Ne = 0.01 - std_i = 0.01 + σ_g = 0.01 + σ_Ne = 0.01 + σ_i = 0.01 λ = 0.38149498593241726 ω = 0.0022277804179292593 b_u = 0.037601010101010155 - - β > 0 - σ > 0 - 0 <= h < 1 - χ > 0 - φ_l > 0 - ζ_u > 0 - 0 < θ < 1 - 0 < α < 1 - 0 < δ < 1 - 0 <= g_y < 1 - η_i >= 0 - ϵ_p > 1 - 0 <= γ_p < 1 - 0 <= γ_p_index < 1 - 0 <= ρ_i < 1 - κ_pi > 1 - 0 <= ρ_ξ < 1 - std_ξ >= 0 - 0 <= ρ_A < 1 - std_A >= 0 - 0 <= ρ_g < 1 - std_g >= 0 - std_Ne >= 0 - std_i >= 0 - λ > 0 - ω > 0 - b_u > 0 end diff --git a/src/steady_state/solver_parameters.jl b/src/steady_state/solver_parameters.jl index cde260f68..e4ccdf355 100644 --- a/src/steady_state/solver_parameters.jl +++ b/src/steady_state/solver_parameters.jl @@ -28,5 +28,7 @@ const DEFAULT_SOLVER_PARAMETERS = [ solver_parameters(4.1784912636092235, 1.8166012668623566, 0.5168801279930487, 78.18194336881028, 2.139580134601701, 0.4617967010780055, 33.95219683424897, 17.315839925955242, 2.220446049250313e-16, 12.287343174930065, 2.220446049250313e-16, 6.185479065850274, 88.3014875814592, 36.31304631280673, 5.262437586106421, 2.220446049250313e-16, 2.220446049250313e-16, 6.347784900438273, 0.7130503478600859, 0.6594888633818169, 1, 0.0, 2), - solver_parameters(75.41767502114854, 48.44868207231484, 66.70557675476336, 17.001461658038423, 2.715293366195093, 1.2745896783633328, 37.70721763849395, 50.739875243093444, 66.72525429469775, 54.137861579508154, 12.078847271504216, 19.723433260864525, 2.2648933923720476, 54.735557478829, 98.94783074858547, 17.863177309960086, 85.97559329517274, 64.79678330684743, 26.59975637589043, 24.72319192940016, 1, 0.0, 2) + solver_parameters(75.41767502114854, 48.44868207231484, 66.70557675476336, 17.001461658038423, 2.715293366195093, 1.2745896783633328, 37.70721763849395, 50.739875243093444, 66.72525429469775, 54.137861579508154, 12.078847271504216, 19.723433260864525, 2.2648933923720476, 54.735557478829, 98.94783074858547, 17.863177309960086, 85.97559329517274, 64.79678330684743, 26.59975637589043, 24.72319192940016, 1, 0.0, 2), + + solver_parameters(10.327106226758554, 10.327106226758554, 1.0821485512186284, 52.71196149219814, 2.53193716831508, 1.7586052095638707, 50.434347492548646, 1.7535573088570338, 1.51581318888478, 0.943191530494055, 1.8556283860336538, 6.8221667685490015, 72.18579387599354, 7.7659847231914725, 1.6450426805205853, 1.7020036851998528, 2.352764735403362, 0.24785068495263674, 5.649433169080131, 0.9245175406565802, 1, 0.0, 2) ] From c84a5033bb81cb225f5ee50876eb080a34d67d12 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 18 Jun 2026 22:37:35 +0200 Subject: [PATCH 24/37] edit equation processing to allow ( exp(e[x]) ) --- src/parser/equation_processing.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/parser/equation_processing.jl b/src/parser/equation_processing.jl index 13c793e1b..3873cb70b 100644 --- a/src/parser/equation_processing.jl +++ b/src/parser/equation_processing.jl @@ -257,6 +257,7 @@ function process_model_equations(model_block_in::Expr, max_obc_horizon::Int, pre 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 : + !(x.args[2] isa Expr) ? x : x.args[2].head == :ref ? x.args[2].args[1] isa Symbol ? # nonnegative variables begin @@ -305,6 +306,7 @@ function process_model_equations(model_block_in::Expr, max_obc_horizon::Int, pre 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 : + !(x.args[2] isa Expr) ? x : x.args[2].head == :ref ? x.args[2].args[1] isa Symbol ? # nonnegative variables begin @@ -393,6 +395,7 @@ function process_model_equations(model_block_in::Expr, max_obc_horizon::Int, pre 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 : + !(x.args[2] isa Expr) ? x : # exp of a constant (e.g. after shock → 0 substitution) x.args[2].head == :ref ? x.args[2].args[1] isa Symbol ? # have exp terms bound so they dont go to Inf begin @@ -978,7 +981,7 @@ function process_parameter_definitions(parameter_block_in::Expr, pmm::post_model x : begin diffed = intersect(setdiff([x], ss_tmp), get_symbols(cal_eq)) - if !isempty(diffed) + if !isempty(diffed) && diffed[1] ∉ SYMPYWORKSPACE_RESERVED_NAMES push!(par_tmp,diffed[1]) end end : @@ -1034,7 +1037,7 @@ function process_parameter_definitions(parameter_block_in::Expr, pmm::post_model x : begin diffed = setdiff([x],ss_tmp) - if !isempty(diffed) + if !isempty(diffed) && diffed[1] ∉ SYMPYWORKSPACE_RESERVED_NAMES push!(par_tmp,diffed[1]) end end : From ea69ae1b8d904bda3440a94269a169ca1c091e59 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Thu, 18 Jun 2026 22:52:17 +0200 Subject: [PATCH 25/37] FV baseline dsge model --- models/Fernandez_Villaverde_2010.jl | 134 ++++++++++++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 models/Fernandez_Villaverde_2010.jl diff --git a/models/Fernandez_Villaverde_2010.jl b/models/Fernandez_Villaverde_2010.jl new file mode 100644 index 000000000..a7906a1b8 --- /dev/null +++ b/models/Fernandez_Villaverde_2010.jl @@ -0,0 +1,134 @@ +# Fernández-Villaverde, J. (2010). The econometrics of DSGE models. +# SERIEs, 1(1–2), 3–49. https://doi.org/10.1007/s13209-009-0001-0 + +@model Fernandez_Villaverde_2010 begin + d[0] / (c[0] - h * c[-1] / μᶻ[0]) - h * β * d[1] / (c[1] * μᶻ[1] - h * c[0]) = λ[0] + + λ[0] = β * λ[1] / μᶻ[1] / π[1] * R[0] + + r[0] = γ₁ + γ₂ * (u[0] - 1) + + q[0] = β * λ[1] / λ[0] / μᶻ[1] / μᴵ[1] * ((1 - δ) * q[1] + r[1] * u[1] - (γ₁ * (u[1] - 1) + γ₂ / 2 * (u[1] - 1)^2)) + + 1 = q[0] * (1 - (κ / 2 * (x[0] / x[-1] * μᶻ[0] - Λˣ)^2) - (κ * (x[0] / x[-1] * μᶻ[0] - Λˣ) * x[0] / x[-1] * μᶻ[0])) + β * q[1] * λ[1] / λ[0] / μᶻ[1] * κ * (x[1] / x[0] * μᶻ[1] - Λˣ) * (x[1] / x[0] * μᶻ[1])^2 + + f[0] = (η - 1) / η * w_star[0]^(1 - η) * λ[0] * w[0]^η * lᵈ[0] + β * θʷ * (π[0]^χʷ / π[1])^(1 - η) * (w_star[1] / w_star[0] * μᶻ[1])^(η - 1) * f[1] + + f[0] = ψ * d[0] * φ[0] * π_w_star[0]^(-η * (1 + γ)) * lᵈ[0]^(1 + γ) + β * θʷ * (π[0]^χʷ / π[1])^(-η * (1 + γ)) * (w_star[1] / w_star[0] * μᶻ[1])^(η * (1 + γ)) * f[1] + + g₁[0] = λ[0] * mc[0] * yᵈ[0] + β * θᵖ * (π[0]^χ / π[1])^(-ε) * g₁[1] + + g₂[0] = λ[0] * π_star[0] * yᵈ[0] + β * θᵖ * (π[0]^χ / π[1])^(1 - ε) * π_star[0] / π_star[1] * g₂[1] + + ε * g₁[0] = (ε - 1) * g₂[0] + + u[0] * k[-1] / lᵈ[0] = α / (1 - α) * w[0] / r[0] * μᶻ[0] * μᴵ[0] + + mc[0] = (1 / (1 - α))^(1 - α) * (1 / α)^α * w[0]^(1 - α) * r[0]^α + + 1 = θʷ * (π[-1]^χʷ / π[0])^(1 - η) * (w[-1] / w[0] / μᶻ[0])^(1 - η) + (1 - θʷ) * π_w_star[0]^(1 - η) + + 1 = θᵖ * (π[-1]^χ / π[0])^(1 - ε) + (1 - θᵖ) * π_star[0]^(1 - ε) + + R[0] / R̄ = (R[-1] / R̄)^γᴿ * ((π[0] / π̄)^γ_π * ((yᵈ[0] / yᵈ[-1] * μᶻ[0]) / exp(Λʸᵈ))^γʸ)^(1 - γᴿ) * exp(ϵᵐ[x]) + + yᵈ[0] = c[0] + x[0] + (γ₁ * (u[0] - 1) + γ₂ / 2 * (u[0] - 1)^2) * k[-1] / μᶻ[0] / μᴵ[0] + + yᵈ[0] = (μᴬ[0] / μᶻ[0] * (u[0] * k[-1])^α * lᵈ[0]^(1 - α) - Φ) / vᵖ[0] + + l[0] = vʷ[0] * lᵈ[0] + + vᵖ[0] = θᵖ * (π[-1]^χ / π[0])^(-ε) * vᵖ[-1] + (1 - θᵖ) * π_star[0]^(-ε) + + vʷ[0] = θʷ * (w[-1] / w[0] / μᶻ[0] * π[-1]^χʷ / π[0])^(-η) * vʷ[-1] + (1 - θʷ) * π_w_star[0]^(-η) + + k[0] = (1 - δ) * k[-1] / (μᶻ[0] * μᴵ[0]) + (1 - κ / 2 * (x[0] / x[-1] * μᶻ[0] - Λˣ)^2) * x[0] + + F[0] = yᵈ[0] - 1 / (1 - α) * w[0] * lᵈ[0] + + π_w_star[0] = w_star[0] / w[0] + + log(d[0]) = ρᵈ * log(d[-1]) + σᵈ * ϵᵈ[x] + + log(φ[0]) = ρᵠ * log(φ[-1]) + σᵠ * ϵᵠ[x] + + log(μᴵ[0]) = Λ_μ + σ_μ * ϵ_μ_I[x] + + log(μᴬ[0]) = Λᴬ + σᴬ * ϵᴬ[x] + + μᶻ[0] = μᴬ[0]^(1 / (1 - α)) * μᴵ[0]^(α / (1 - α)) +end + +@parameters Fernandez_Villaverde_2010 begin + # Steady-state growth rates (computed from drift parameters) + Λʸᵈ = (Λᴬ + α * Λ_μ) / (1 - α) + + Λˣ = exp(Λʸᵈ) + + γ₁ = exp(Λʸᵈ) * exp(Λ_μ) / β - (1 - δ) + + R̄ = 1 + (π̄ * exp(Λʸᵈ) / β - 1) + + # Structural parameters + h = 0.97 + + β = 0.998 + + δ = 0.025 + + κ = 9.51 + + η = 10 + + ε = 10 + + ψ = 8.92 + + γ = 1.17 + + χʷ = 0.62 + + χ = 0.63 + + θᵖ = 0.82 + + θʷ = 0.68 + + α = 0.21 + + Φ = 0 + + γ₂ = 0.001 + + # Monetary policy rule + γᴿ = 0.77 + + γʸ = 0.19 + + γ_π = 1.29 + + π̄ = 1.01 + + # Shock persistence + ρᵈ = 0.12 + + ρᵠ = 0.93 + + # Shock drift parameters + Λ_μ = 0.0034 + + Λᴬ = 0.0028 + + # Shock standard deviations (median posterior estimates, FV 2010 Table 1) + σᵈ = exp(-1.51) + + σᵠ = exp(-2.36) + + σ_μ = exp(-5.43) + + σᴬ = exp(-3.97) + + σᵐ = exp(-5.85) + + u > 0 +end From 65b56c2038c1e19c5ad002b04db92475ea249c3e Mon Sep 17 00:00:00 2001 From: Thore Kockerols Date: Fri, 19 Jun 2026 15:30:11 +0000 Subject: [PATCH 26/37] add sigma m to shock --- models/Fernandez_Villaverde_2010.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Fernandez_Villaverde_2010.jl b/models/Fernandez_Villaverde_2010.jl index a7906a1b8..92fd09ebc 100644 --- a/models/Fernandez_Villaverde_2010.jl +++ b/models/Fernandez_Villaverde_2010.jl @@ -30,7 +30,7 @@ 1 = θᵖ * (π[-1]^χ / π[0])^(1 - ε) + (1 - θᵖ) * π_star[0]^(1 - ε) - R[0] / R̄ = (R[-1] / R̄)^γᴿ * ((π[0] / π̄)^γ_π * ((yᵈ[0] / yᵈ[-1] * μᶻ[0]) / exp(Λʸᵈ))^γʸ)^(1 - γᴿ) * exp(ϵᵐ[x]) + R[0] / R̄ = (R[-1] / R̄)^γᴿ * ((π[0] / π̄)^γ_π * ((yᵈ[0] / yᵈ[-1] * μᶻ[0]) / exp(Λʸᵈ))^γʸ)^(1 - γᴿ) * exp(σᵐ * ϵᵐ[x]) yᵈ[0] = c[0] + x[0] + (γ₁ * (u[0] - 1) + γ₂ / 2 * (u[0] - 1)^2) * k[-1] / μᶻ[0] / μᴵ[0] From 439c3d4b44199204df77e50fcfb3cbcaffffa7b8 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Tue, 23 Jun 2026 09:07:22 +0200 Subject: [PATCH 27/37] add citation --- models/Fernandez_Villaverde_2010.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/models/Fernandez_Villaverde_2010.jl b/models/Fernandez_Villaverde_2010.jl index a7906a1b8..3afa8e454 100644 --- a/models/Fernandez_Villaverde_2010.jl +++ b/models/Fernandez_Villaverde_2010.jl @@ -1,3 +1,6 @@ +# Fernández-Villaverde, J. and Rubio-Ramírez, J.F. (2006). A Baseline DSGE Model. +# Manuscript, University of Pennsylvania. +# # Fernández-Villaverde, J. (2010). The econometrics of DSGE models. # SERIEs, 1(1–2), 3–49. https://doi.org/10.1007/s13209-009-0001-0 From 474f5b45520b3547656f583b15eb5d8a8327f286 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 14:24:06 +0200 Subject: [PATCH 28/37] parameters allow for loops and if/else statements in parameters macro --- src/parser/equation_processing.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/parser/equation_processing.jl b/src/parser/equation_processing.jl index 3873cb70b..04a240d06 100644 --- a/src/parser/equation_processing.jl +++ b/src/parser/equation_processing.jl @@ -852,7 +852,10 @@ function process_parameter_definitions(parameter_block_in::Expr, pmm::post_model par_defined_more_than_once = Set() bounded_vars = [] - parameter_definitions = replace_indices(parameter_block_in) + parameter_definitions = parse_for_loops(parameter_block_in) + parameter_definitions = resolve_if_expr(parameter_definitions) + parameter_definitions = remove_nothing(parameter_definitions) + parameter_definitions = replace_indices(parameter_definitions) # parse parameter inputs # label all variables parameters and exogenous variables and timings across all equations From 44b3a1eb13878ceecfbc5d187490d2f4b351b7d6 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 14:25:19 +0200 Subject: [PATCH 29/37] tick off todo for for loop/ if else in parameters macro --- docs/src/unfinished_docs/todo.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/unfinished_docs/todo.md b/docs/src/unfinished_docs/todo.md index 173396b25..962fba35b 100644 --- a/docs/src/unfinished_docs/todo.md +++ b/docs/src/unfinished_docs/todo.md @@ -105,7 +105,7 @@ - [ ] add balanced growth path handling - [x] autocorr and corr with derivatives. return 3d array - [ ] add pydsge and econpizza to overview -- [ ] add for loop parser in @parameters +- [x] add for loop parser in @parameters - [ ] implement more multi country models - [ ] speed benchmarking (focus on ImplicitDiff part) - [ ] for cond forecasting allow less shocks than conditions with a warning. should be svd then From f8caed52ef42bc4dcf64be2e9bd7b914ed7db4ab Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 14:35:00 +0200 Subject: [PATCH 30/37] make BKK92 more general and compatible with more countries --- models/Backus_Kehoe_Kydland_1992.jl | 59 +++++++++++++++++------------ 1 file changed, 35 insertions(+), 24 deletions(-) diff --git a/models/Backus_Kehoe_Kydland_1992.jl b/models/Backus_Kehoe_Kydland_1992.jl index c6eeb4fae..5a222a909 100644 --- a/models/Backus_Kehoe_Kydland_1992.jl +++ b/models/Backus_Kehoe_Kydland_1992.jl @@ -1,44 +1,49 @@ @model Backus_Kehoe_Kydland_1992 begin + # To scale to N countries, add/remove symbols in every [H, F] list below. + # Update the @parameters block rho definitions and steady-state calibration accordingly. 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}) + 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] + 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 + 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] + 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] + 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} + 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} * 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}) + 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} + 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}) + -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})) + 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] + 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] + # Shock process: each country's lambda depends on its own lag and spillovers from all others. + # The inner accumulator sums over co2 != co, handling any number of countries. + for co in [H, F] + (LAMBDA{co}[0] - 1) = rho{co}{co} * (LAMBDA{co}[-1] - 1) + for co2 in [H, F] if co2 != co rho{co}{co2} * (LAMBDA{co2}[-1] - 1) end end + Z_E{co} * E{co}[x] + end - 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 + # World resource constraint: sum of expenditures = sum of outputs + 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 end @parameters Backus_Kehoe_Kydland_1992 begin K_ss = 11 K[ss] = K_ss | beta - + mu = 0.34 gamma = -1.0 alpha = 1 @@ -51,9 +56,15 @@ end 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} + + # To scale to N countries, update the country list in the loops below + for co1 in [H, F] + for co2 in [H, F] + if co1 == co2 + rho{co1}{co2} = 0.906 + else + rho{co1}{co2} = 0.088 / 1 # divide by N - 1 for the model to be stable with more countries + end + end + end end \ No newline at end of file From cd380b867a1cb14f8ce98d1883ba9086d848d9c5 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 17:19:58 +0200 Subject: [PATCH 31/37] rm nonlinear from new model naming --- ...t_1999_nonlinear.jl => Bernanke_Gertler_Gilchrist_1999.jl} | 4 ++-- ...ertler_Karadi_2011_nonlinear.jl => Gertler_Karadi_2011.jl} | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) rename models/{Bernanke_Gertler_Gilchrist_1999_nonlinear.jl => Bernanke_Gertler_Gilchrist_1999.jl} (97%) rename models/{Gertler_Karadi_2011_nonlinear.jl => Gertler_Karadi_2011.jl} (98%) diff --git a/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl b/models/Bernanke_Gertler_Gilchrist_1999.jl similarity index 97% rename from models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl rename to models/Bernanke_Gertler_Gilchrist_1999.jl index f6a4238b4..ddd6e5bc0 100644 --- a/models/Bernanke_Gertler_Gilchrist_1999_nonlinear.jl +++ b/models/Bernanke_Gertler_Gilchrist_1999.jl @@ -1,6 +1,6 @@ using MacroModelling -@model Bernanke_Gertler_Gilchrist_1999_nonlinear begin +@model Bernanke_Gertler_Gilchrist_1999 begin # [B.3] Household Euler equation: intertemporal optimality for deposits. C[0]^(-1) = β * C[1]^(-1) * R[0] / π[1] @@ -98,7 +98,7 @@ using MacroModelling premium_ann[0] = 400 * log(s[0]) end -@parameters Bernanke_Gertler_Gilchrist_1999_nonlinear begin +@parameters Bernanke_Gertler_Gilchrist_1999 begin β = 0.99 α = 0.35 δ = 0.025 diff --git a/models/Gertler_Karadi_2011_nonlinear.jl b/models/Gertler_Karadi_2011.jl similarity index 98% rename from models/Gertler_Karadi_2011_nonlinear.jl rename to models/Gertler_Karadi_2011.jl index 0e12689ca..e1589a782 100644 --- a/models/Gertler_Karadi_2011_nonlinear.jl +++ b/models/Gertler_Karadi_2011.jl @@ -1,6 +1,6 @@ using MacroModelling -@model Gertler_Karadi_2011_nonlinear begin +@model Gertler_Karadi_2011 begin # [4] Marginal utility of consumption with external habit. ϱ[0] = (C[0] - h * C[-1])^(-σ) - β * h * (C[1] - h * C[0])^(-σ) @@ -119,7 +119,7 @@ using MacroModelling prem[0] = Rᵏ[1] / R[0] end -@parameters Gertler_Karadi_2011_nonlinear begin +@parameters Gertler_Karadi_2011 begin β = 0.99 σ = 1.0 h = 0.815 From 8cd93df1b7931170991d2902086ebea90dcffd07 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 17:20:05 +0200 Subject: [PATCH 32/37] update todos --- docs/src/unfinished_docs/todo.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/src/unfinished_docs/todo.md b/docs/src/unfinished_docs/todo.md index 962fba35b..1b0d2aa11 100644 --- a/docs/src/unfinished_docs/todo.md +++ b/docs/src/unfinished_docs/todo.md @@ -52,7 +52,6 @@ - [ ] implement forwarddiff for find_shocks - [ ] redo inversion filter 1st order rrule based on the higher order ones. the accumulated matmul might not be necessary at all - [ ] inversion filter: use subset of observables and states when propagating states (see kalman filter) -- [x] start filter from initial values provided by user - [ ] 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 @@ -73,9 +72,6 @@ - [ ] do proper testing of ss solver with random set of params, equal across configs - [ ] load create parts of derivatives later and not directly after parameters block - [ ] fix model estimate plot. data not above estimate (should be red but is blue) -- [x] implement higher order (pruned) variance decomposition -- [x] add `marginal_contribution` (Shapley) option to `get_variance_decomposition` for pruned higher-order solutions -- [x] add `marginal_contribution` (Shapley) option to `get_shock_decomposition`/`plot_shock_decomposition` for pruned higher-order solutions (allocates the `Nonlinearities` term across shocks) - [ ] try slicesampler instead of pigeons - [ ] 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. @@ -103,9 +99,7 @@ - [ ] 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 -- [x] autocorr and corr with derivatives. return 3d array - [ ] add pydsge and econpizza to overview -- [x] add for loop parser in @parameters - [ ] implement more multi country models - [ ] speed benchmarking (focus on ImplicitDiff part) - [ ] for cond forecasting allow less shocks than conditions with a warning. should be svd then @@ -129,6 +123,12 @@ - [ ] figure out combinations for inputs (parameters and variables in different formats for get_irf for example) - [ ] weed out SS solver and saved objects +- [x] start filter from initial values provided by user +- [x] implement higher order (pruned) variance decomposition +- [x] add `marginal_contribution` (Shapley) option to `get_variance_decomposition` for pruned higher-order solutions +- [x] add `marginal_contribution` (Shapley) option to `get_shock_decomposition`/`plot_shock_decomposition` for pruned higher-order solutions (allocates the `Nonlinearities` term across shocks) +- [x] autocorr and corr with derivatives. return 3d array +- [x] add for loop parser in @parameters - [x] add filter free estimation to one of the estimation tests (think about filter free filtering instead of smoothing) - [x] check dispatch doctor coverage in new functions - [x] dont use variable, parameter, function names starting with _ From 8f5dbaa580ffb3b6e1a3add927ca478ed4086e9b Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 17:37:39 +0200 Subject: [PATCH 33/37] model todos --- models/Bernanke_Gertler_Gilchrist_1999.jl | 3 ++- models/Gertler_Karadi_2011.jl | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/models/Bernanke_Gertler_Gilchrist_1999.jl b/models/Bernanke_Gertler_Gilchrist_1999.jl index ddd6e5bc0..49308915f 100644 --- a/models/Bernanke_Gertler_Gilchrist_1999.jl +++ b/models/Bernanke_Gertler_Gilchrist_1999.jl @@ -1,4 +1,5 @@ -using MacroModelling +# TODO: add reference here +# TODO: add tests and verify against paper / other implementations @model Bernanke_Gertler_Gilchrist_1999 begin # [B.3] Household Euler equation: intertemporal optimality for deposits. diff --git a/models/Gertler_Karadi_2011.jl b/models/Gertler_Karadi_2011.jl index e1589a782..f3af96fa9 100644 --- a/models/Gertler_Karadi_2011.jl +++ b/models/Gertler_Karadi_2011.jl @@ -1,4 +1,5 @@ -using MacroModelling +# TODO: add reference here +# TODO: add tests and verify against paper / other implementations @model Gertler_Karadi_2011 begin # [4] Marginal utility of consumption with external habit. From ab1af02c369257bb565a2f4a6396b17b0bde0b28 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 20:00:51 +0100 Subject: [PATCH 34/37] guard x.args[2].head access against non-Expr values in process_parameter_definitions --- src/parser/equation_processing.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/parser/equation_processing.jl b/src/parser/equation_processing.jl index 04a240d06..b74cdf871 100644 --- a/src/parser/equation_processing.jl +++ b/src/parser/equation_processing.jl @@ -916,7 +916,7 @@ function process_parameter_definitions(parameter_block_in::Expr, pmm::post_model end : x : x : - x.args[2].head == :block ? + x.args[2] isa Expr && x.args[2].head == :block ? x.args[1].args[1] == :| ? x : x.args[2].args[2].args[1] == :| ? @@ -929,7 +929,7 @@ function process_parameter_definitions(parameter_block_in::Expr, pmm::post_model @warn "Invalid parameter input ignored: " * repr(x) x end : - x.args[2].head == :call ? + x.args[2] isa Expr && x.args[2].head == :call ? x.args[1].args[1] == :| ? x : begin # this is calibration by targeting SS values (conditional parameter at the end) From 42e085ba103f6059a8885a8db75a9e572731cb42 Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 20:48:19 +0100 Subject: [PATCH 35/37] Add references and model list entries for GK11 and BGG1999 --- README.md | 2 ++ docs/src/index.md | 2 ++ models/Bernanke_Gertler_Gilchrist_1999.jl | 6 ++++-- models/Gertler_Karadi_2011.jl | 4 ++-- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 73d0efe6c..a679384a9 100644 --- a/README.md +++ b/README.md @@ -236,11 +236,13 @@ The package contains the following models in the `models` folder: - [Ascari and Sbordone (2014)](https://www.aeaweb.org/articles?id=10.1257/jel.52.3.679) `Ascari_Sbordone_2014.jl` - [Backus, Kehoe, and Kydland (1992)](https://www.jstor.org/stable/2138686) `Backus_Kehoe_Kydland_1992.jl` - [Baxter and King (1993)](https://www.jstor.org/stable/2117521) `Baxter_King_1993.jl` +- [Bernanke, Gertler, and Gilchrist (1999)](https://doi.org/10.1016/S1574-0048(99)10034-X) `Bernanke_Gertler_Gilchrist_1999.jl` - [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` +- [Gertler and Karadi (2011)](https://doi.org/10.1016/j.jmoneco.2010.10.004) `Gertler_Karadi_2011.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` - [Jermann and Quadrini (2012)](https://www.aeaweb.org/articles?id=10.1257/aer.102.1.238) - RBC `JQ_2012_RBC.jl` diff --git a/docs/src/index.md b/docs/src/index.md index b7efe3592..2c0fa65d9 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -38,11 +38,13 @@ The package contains the following models in the `models` folder: - [Ascari and Sbordone (2014)](https://www.aeaweb.org/articles?id=10.1257/jel.52.3.679) `Ascari_Sbordone_2014.jl` - [Backus, Kehoe, and Kydland (1992)](https://www.jstor.org/stable/2138686) `Backus_Kehoe_Kydland_1992.jl` - [Baxter and King (1993)](https://www.jstor.org/stable/2117521) `Baxter_King_1993.jl` +- [Bernanke, Gertler, and Gilchrist (1999)](https://doi.org/10.1016/S1574-0048(99)10034-X) `Bernanke_Gertler_Gilchrist_1999.jl` - [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` +- [Gertler and Karadi (2011)](https://doi.org/10.1016/j.jmoneco.2010.10.004) `Gertler_Karadi_2011.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` - [Jermann and Quadrini (2012)](https://www.aeaweb.org/articles?id=10.1257/aer.102.1.238) - RBC `JQ_2012_RBC.jl` diff --git a/models/Bernanke_Gertler_Gilchrist_1999.jl b/models/Bernanke_Gertler_Gilchrist_1999.jl index 49308915f..3c7d51513 100644 --- a/models/Bernanke_Gertler_Gilchrist_1999.jl +++ b/models/Bernanke_Gertler_Gilchrist_1999.jl @@ -1,5 +1,7 @@ -# TODO: add reference here -# TODO: add tests and verify against paper / other implementations +# Bernanke, B., Gertler, M., and Gilchrist, S. (1999). The financial accelerator in a +# quantitative business cycle framework. In J. B. Taylor and M. Woodford (Eds.), +# Handbook of Macroeconomics (Vol. 1, pp. 1341–1393). Elsevier. +# https://doi.org/10.1016/S1574-0048(99)10034-X @model Bernanke_Gertler_Gilchrist_1999 begin # [B.3] Household Euler equation: intertemporal optimality for deposits. diff --git a/models/Gertler_Karadi_2011.jl b/models/Gertler_Karadi_2011.jl index f3af96fa9..28279b561 100644 --- a/models/Gertler_Karadi_2011.jl +++ b/models/Gertler_Karadi_2011.jl @@ -1,5 +1,5 @@ -# TODO: add reference here -# TODO: add tests and verify against paper / other implementations +# Gertler, M. and Karadi, P. (2011). A model of unconventional monetary policy. +# Journal of Monetary Economics, 58(1), 17–34. https://doi.org/10.1016/j.jmoneco.2010.10.004 @model Gertler_Karadi_2011 begin # [4] Marginal utility of consumption with external habit. From ab19cf38037d633505db897abae9810de0166a7f Mon Sep 17 00:00:00 2001 From: thorek1 Date: Sun, 28 Jun 2026 20:49:51 +0100 Subject: [PATCH 36/37] Add paper references at the top of all model files --- models/Aguiar_Gopinath_2007.jl | 3 +++ models/Ascari_Sbordone_2014.jl | 3 +++ models/Backus_Kehoe_Kydland_1992.jl | 3 +++ models/Baxter_King_1993.jl | 3 +++ models/Caldara_et_al_2012.jl | 4 ++++ models/FS2000.jl | 3 +++ models/GNSS_2010.jl | 4 ++++ models/Gali_2015_chapter_3_nonlinear.jl | 3 +++ models/Gali_2015_chapter_3_obc.jl | 3 +++ models/Gali_Monacelli_2005_CITR.jl | 4 ++++ models/Ghironi_Melitz_2005.jl | 4 ++++ models/Guerrieri_Iacoviello_2017.jl | 4 ++++ models/Iacoviello_2005_linear.jl | 4 ++++ models/Ireland_2004.jl | 3 +++ models/JQ_2012_RBC.jl | 3 +++ models/NAWM_EAUS_2008.jl | 4 ++++ models/QUEST3_2009.jl | 4 ++++ models/RBC_baseline.jl | 4 ++++ models/SGU_2003_debt_premium.jl | 4 ++++ models/Smets_Wouters_2003.jl | 4 ++++ models/Smets_Wouters_2003_obc.jl | 6 ++++++ models/Smets_Wouters_2007.jl | 4 ++++ models/Smets_Wouters_2007_linear.jl | 6 ++++++ 23 files changed, 87 insertions(+) diff --git a/models/Aguiar_Gopinath_2007.jl b/models/Aguiar_Gopinath_2007.jl index 04bab5e88..e6477f236 100644 --- a/models/Aguiar_Gopinath_2007.jl +++ b/models/Aguiar_Gopinath_2007.jl @@ -1,3 +1,6 @@ +# Aguiar, M. and Gopinath, G. (2007). Emerging market business cycles: The cycle is the trend. +# Journal of Political Economy, 115(1), 69–102. https://doi.org/10.1086/511283 + @model Aguiar_Gopinath_2007 begin y[0] = (exp(g[0]) * l[0]) ^ alpha * exp(z[0]) * k[-1] ^ (1 - alpha) diff --git a/models/Ascari_Sbordone_2014.jl b/models/Ascari_Sbordone_2014.jl index c246ab450..a73543548 100644 --- a/models/Ascari_Sbordone_2014.jl +++ b/models/Ascari_Sbordone_2014.jl @@ -1,3 +1,6 @@ +# Ascari, G. and Sbordone, A.M. (2014). The macroeconomics of trend inflation. +# Journal of Economic Literature, 52(3), 679–739. https://doi.org/10.1257/jel.52.3.679 + @model Ascari_Sbordone_2014 begin 1 / y[0] ^ sigma = beta * (1 + i[0]) / (pi[1] * y[1] ^ sigma) diff --git a/models/Backus_Kehoe_Kydland_1992.jl b/models/Backus_Kehoe_Kydland_1992.jl index 5a222a909..4870e90ba 100644 --- a/models/Backus_Kehoe_Kydland_1992.jl +++ b/models/Backus_Kehoe_Kydland_1992.jl @@ -1,3 +1,6 @@ +# Backus, D.K., Kehoe, P.J., and Kydland, F.E. (1992). International real business cycles. +# Journal of Political Economy, 100(4), 745–775. https://www.jstor.org/stable/2138686 + @model Backus_Kehoe_Kydland_1992 begin # To scale to N countries, add/remove symbols in every [H, F] list below. # Update the @parameters block rho definitions and steady-state calibration accordingly. diff --git a/models/Baxter_King_1993.jl b/models/Baxter_King_1993.jl index 31fbd8d26..51f310ab6 100644 --- a/models/Baxter_King_1993.jl +++ b/models/Baxter_King_1993.jl @@ -1,3 +1,6 @@ +# Baxter, M. and King, R.G. (1993). Fiscal policy in general equilibrium. +# American Economic Review, 83(3), 315–334. https://www.jstor.org/stable/2117521 + @model Baxter_King_1993 begin uc[0] = c[0]^-1 diff --git a/models/Caldara_et_al_2012.jl b/models/Caldara_et_al_2012.jl index a93eae445..e38ad9ac7 100644 --- a/models/Caldara_et_al_2012.jl +++ b/models/Caldara_et_al_2012.jl @@ -1,3 +1,7 @@ +# Caldara, D., Fernández-Villaverde, J., Rubio-Ramírez, J.F., and Yao, W. (2012). Computing +# DSGE models with recursive preferences and stochastic volatility. +# Review of Economic Dynamics, 15(2), 188–206. https://doi.org/10.1016/j.red.2011.10.001 + @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/FS2000.jl b/models/FS2000.jl index 3e4d86e78..da5d89f60 100644 --- a/models/FS2000.jl +++ b/models/FS2000.jl @@ -1,3 +1,6 @@ +# Schorfheide, F. (2000). Loss function-based evaluation of DSGE models. +# Journal of Applied Econometrics, 15(6), 645–670. https://doi.org/10.1002/jae.582 + @model FS2000 begin dA[0] = exp(gam + z_e_a * e_a[x]) diff --git a/models/GNSS_2010.jl b/models/GNSS_2010.jl index a71337a38..da30f1c00 100644 --- a/models/GNSS_2010.jl +++ b/models/GNSS_2010.jl @@ -1,3 +1,7 @@ +# Gerali, A., Neri, S., Sessa, L., and Signoretti, F.M. (2010). Credit and banking in a DSGE +# model of the euro area. Journal of Money, Credit and Banking, 42(s1), 107–141. +# https://doi.org/10.1111/j.1538-4616.2010.00331.x + @model GNSS_2010 begin (1 - a_i) * exp(ee_z[0]) * (c_p[0] - a_i * c_p[-1]) ^ (-1) = lam_p[0] diff --git a/models/Gali_2015_chapter_3_nonlinear.jl b/models/Gali_2015_chapter_3_nonlinear.jl index db71d8a1f..89bb982c0 100644 --- a/models/Gali_2015_chapter_3_nonlinear.jl +++ b/models/Gali_2015_chapter_3_nonlinear.jl @@ -1,3 +1,6 @@ +# Gali, J. (2015). Monetary Policy, Inflation, and the Business Cycle: An Introduction to the +# New Keynesian Framework and Its Applications (2nd ed.). Princeton University Press. + @model Gali_2015_chapter_3_nonlinear begin W_real[0] = C[0] ^ σ * N[0] ^ φ diff --git a/models/Gali_2015_chapter_3_obc.jl b/models/Gali_2015_chapter_3_obc.jl index 09212af7a..68555fc88 100644 --- a/models/Gali_2015_chapter_3_obc.jl +++ b/models/Gali_2015_chapter_3_obc.jl @@ -1,3 +1,6 @@ +# Gali, J. (2015). Monetary Policy, Inflation, and the Business Cycle: An Introduction to the +# New Keynesian Framework and Its Applications (2nd ed.). Princeton University Press. + @model Gali_2015_chapter_3_obc begin W_real[0] = C[0] ^ σ * N[0] ^ φ diff --git a/models/Gali_Monacelli_2005_CITR.jl b/models/Gali_Monacelli_2005_CITR.jl index 28f810a09..8c75a5b81 100644 --- a/models/Gali_Monacelli_2005_CITR.jl +++ b/models/Gali_Monacelli_2005_CITR.jl @@ -1,3 +1,7 @@ +# Gali, J. and Monacelli, T. (2005). Monetary policy and exchange rate volatility in a small +# open economy. Review of Economic Studies, 72(3), 707–734. +# https://doi.org/10.1111/j.1467-937X.2005.00349.x + @model Gali_Monacelli_2005_CITR begin x[0] = x[1] - σᵃ ^ (-1) * (r[0] - pih[1] - rnat[0]) diff --git a/models/Ghironi_Melitz_2005.jl b/models/Ghironi_Melitz_2005.jl index a66372eb1..34ebc38c1 100644 --- a/models/Ghironi_Melitz_2005.jl +++ b/models/Ghironi_Melitz_2005.jl @@ -1,3 +1,7 @@ +# Ghironi, F. and Melitz, M.J. (2005). International trade and macroeconomic dynamics with +# heterogeneous firms. Quarterly Journal of Economics, 120(3), 865–915. +# https://doi.org/10.1093/qje/120.3.865 + @model Ghironi_Melitz_2005 begin 1 = Nd[0] * ρ̃d[0] ^ (1 - θ) + Nx̄[0] * ρ̃x̄[0] ^ (1 - θ) diff --git a/models/Guerrieri_Iacoviello_2017.jl b/models/Guerrieri_Iacoviello_2017.jl index f8b5412eb..be3f92fa8 100644 --- a/models/Guerrieri_Iacoviello_2017.jl +++ b/models/Guerrieri_Iacoviello_2017.jl @@ -1,3 +1,7 @@ +# Guerrieri, L. and Iacoviello, M. (2017). Collateral constraints and macroeconomic +# asymmetries. Journal of Monetary Economics, 90, 28–49. +# https://doi.org/10.1016/j.jmoneco.2017.06.004 + @model Guerrieri_Iacoviello_2017 begin c[0] + c1[0] + ik[0] = y[0] diff --git a/models/Iacoviello_2005_linear.jl b/models/Iacoviello_2005_linear.jl index 55d405e70..2f15f6ffc 100644 --- a/models/Iacoviello_2005_linear.jl +++ b/models/Iacoviello_2005_linear.jl @@ -1,3 +1,7 @@ +# Iacoviello, M. (2005). House prices, borrowing constraints, and monetary policy in the +# business cycle. American Economic Review, 95(3), 739–764. +# https://doi.org/10.1257/0002828054201477 + @model Iacoviello_2005_linear begin Ŷ[0] = ctoY * ĉ[0] + c1toY * c1̂[0] + c2toY * c2̂[0] + ItoY * Î[0] diff --git a/models/Ireland_2004.jl b/models/Ireland_2004.jl index 5740a53de..77f9c25e3 100644 --- a/models/Ireland_2004.jl +++ b/models/Ireland_2004.jl @@ -1,3 +1,6 @@ +# Ireland, P.N. (2004). Technology shocks in the New Keynesian model. +# Review of Economics and Statistics, 86(4), 923–936. https://doi.org/10.1162/0034653043125185 + @model Ireland_2004 begin a[0] = ρᵃ * a[-1] + σᵃ * ϵᵃ[x] diff --git a/models/JQ_2012_RBC.jl b/models/JQ_2012_RBC.jl index 0c9b29f88..ff5e4b8e5 100644 --- a/models/JQ_2012_RBC.jl +++ b/models/JQ_2012_RBC.jl @@ -1,3 +1,6 @@ +# Jermann, U. and Quadrini, V. (2012). Macroeconomic effects of financial shocks. +# American Economic Review, 102(1), 238–271. https://doi.org/10.1257/aer.102.1.238 + @model JQ_2012_RBC begin w[0] / c[0] ^ σ = α / (1 - n[0]) diff --git a/models/NAWM_EAUS_2008.jl b/models/NAWM_EAUS_2008.jl index 6881ca613..4db37f41e 100644 --- a/models/NAWM_EAUS_2008.jl +++ b/models/NAWM_EAUS_2008.jl @@ -1,3 +1,7 @@ +# European Central Bank. (2008). The New Area-Wide Model of the euro area: A micro-founded +# open-economy model for forecasting and policy analysis. ECB Working Paper No. 944. +# https://www.ecb.europa.eu/pub/pdf/scpwps/ecbwp944.pdf + @model NAWM_EAUS_2008 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] diff --git a/models/QUEST3_2009.jl b/models/QUEST3_2009.jl index eb941fed9..c5f01b730 100644 --- a/models/QUEST3_2009.jl +++ b/models/QUEST3_2009.jl @@ -1,3 +1,7 @@ +# Ratto, M., Roeger, W., and in 't Veld, J. (2009). QUEST III: An estimated open-economy +# DSGE model of the euro area with fiscal and monetary policy. +# Economic Modelling, 26(1), 222–233. https://doi.org/10.1016/j.econmod.2008.06.014 + @model QUEST3_2009 begin interest[0] = ((1 + E_INOM[0]) ^ 4 - interestq_exog ^ 4) / interestq_exog ^ 4 diff --git a/models/RBC_baseline.jl b/models/RBC_baseline.jl index d5fcaeec8..18c949f74 100644 --- a/models/RBC_baseline.jl +++ b/models/RBC_baseline.jl @@ -1,3 +1,7 @@ +# King, R.G., Plosser, C.I., and Rebelo, S.T. (1988). Production, growth and business cycles: +# I. The basic neoclassical model. Journal of Monetary Economics, 21(2–3), 195–232. +# https://doi.org/10.1016/0304-3932(88)90030-X + using MacroModelling @model RBC_baseline begin diff --git a/models/SGU_2003_debt_premium.jl b/models/SGU_2003_debt_premium.jl index 2cbc6b8df..3844eaf39 100644 --- a/models/SGU_2003_debt_premium.jl +++ b/models/SGU_2003_debt_premium.jl @@ -1,3 +1,7 @@ +# Schmitt-Grohé, S. and Uribe, M. (2003). Closing small open economy models. +# Journal of International Economics, 61(1), 163–185. +# https://doi.org/10.1016/S0022-1996(02)00056-9 + @model SGU_2003_debt_premium begin d[0] = (1 + r[-1]) * d[-1] - y[0] + c[0] + i[0] + ϕ / 2 * (k[0] - k[-1]) ^ 2 diff --git a/models/Smets_Wouters_2003.jl b/models/Smets_Wouters_2003.jl index 55751a9ce..f25c7d3b4 100644 --- a/models/Smets_Wouters_2003.jl +++ b/models/Smets_Wouters_2003.jl @@ -1,3 +1,7 @@ +# Smets, F. and Wouters, R. (2003). An estimated dynamic stochastic general equilibrium model +# of the euro area. Journal of the European Economic Association, 1(5), 1123–1175. +# https://doi.org/10.1162/154247603770383415 + @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)) diff --git a/models/Smets_Wouters_2003_obc.jl b/models/Smets_Wouters_2003_obc.jl index 6ecf5ea07..636152dfb 100644 --- a/models/Smets_Wouters_2003_obc.jl +++ b/models/Smets_Wouters_2003_obc.jl @@ -1,3 +1,9 @@ +# Smets, F. and Wouters, R. (2003). An estimated dynamic stochastic general equilibrium model +# of the euro area. Journal of the European Economic Association, 1(5), 1123–1175. +# https://doi.org/10.1162/154247603770383415 +# +# This variant includes an occasionally binding constraint (OBC) on the interest rate. + @model Smets_Wouters_2003_obc 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)) diff --git a/models/Smets_Wouters_2007.jl b/models/Smets_Wouters_2007.jl index e2979d985..6e83b1fdb 100644 --- a/models/Smets_Wouters_2007.jl +++ b/models/Smets_Wouters_2007.jl @@ -1,3 +1,7 @@ +# Smets, F. and Wouters, R. (2007). Shocks and frictions in US business cycles: A Bayesian +# DSGE approach. American Economic Review, 97(3), 586–606. +# https://doi.org/10.1257/aer.97.3.586 + @model Smets_Wouters_2007 begin y[0] = c[0] + inve[0] + y[ss] * gy[0] + afunc[0] * kp[-1] / cgamma diff --git a/models/Smets_Wouters_2007_linear.jl b/models/Smets_Wouters_2007_linear.jl index 40eb7e5f0..0b1bbd2c2 100644 --- a/models/Smets_Wouters_2007_linear.jl +++ b/models/Smets_Wouters_2007_linear.jl @@ -1,3 +1,9 @@ +# Smets, F. and Wouters, R. (2007). Shocks and frictions in US business cycles: A Bayesian +# DSGE approach. American Economic Review, 97(3), 586–606. +# https://doi.org/10.1257/aer.97.3.586 +# +# Linearised variant of the Smets and Wouters (2007) model. + @model Smets_Wouters_2007_linear begin a[0] = calfa * rkf[0] + (1 - calfa) * wf[0] From 3eed3aee8149ac6a2ea2b10cfed833e0f5d1f86a Mon Sep 17 00:00:00 2001 From: thorek1 Date: Mon, 29 Jun 2026 09:24:14 +0100 Subject: [PATCH 37/37] Fix shock ordering: sort shock indices at source via parse_shocks_input_to_index - Move normalize_superscript from StatsPlotsExt.jl to core module - Sort shock_idx by normalized name inside parse_shocks_input_to_index - Import normalize_superscript explicitly in StatsPlotsExt.jl - All downstream callers (get_irf, plot_irf, plot_irf!, girf) now get consistently sorted shock order --- ext/StatsPlotsExt.jl | 42 +++---------------------------------- src/MacroModelling.jl | 48 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 50 insertions(+), 40 deletions(-) diff --git a/ext/StatsPlotsExt.jl b/ext/StatsPlotsExt.jl index 702cc0567..981bcf626 100644 --- a/ext/StatsPlotsExt.jl +++ b/ext/StatsPlotsExt.jl @@ -2,7 +2,7 @@ module StatsPlotsExt using MacroModelling -import MacroModelling: ParameterType, ℳ, Symbol_input, String_input, Tolerances, 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, normalize_presample_periods, trim_informative_sample, adjust_initial_state, SteadyStateFunctionType +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, normalize_presample_periods, trim_informative_sample, adjust_initial_state, SteadyStateFunctionType, normalize_superscript 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 @@ -116,44 +116,6 @@ function apply_custom_name(symbol::R, custom_names::AbstractDict{S, T})::R where 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, @@ -2233,6 +2195,7 @@ function plot_irf(𝓂::ℳ; shock_sort_perm = sortperm(shock_names_display, by = normalize_superscript) shock_idx = shock_idx[shock_sort_perm] shock_names_display = shock_names_display[shock_sort_perm] + Y = Y[:, :, shock_sort_perm] end else shock_names_display = ["shock_matrix"] @@ -2869,6 +2832,7 @@ function plot_irf!(𝓂::ℳ; shock_sort_perm = sortperm(shock_names_display, by = normalize_superscript) shock_idx = shock_idx[shock_sort_perm] shock_names_display = shock_names_display[shock_sort_perm] + Y = Y[:, :, shock_sort_perm] end else shock_names_display = ["shock_matrix"] diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index a8aff7acd..40eb9bb1c 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -386,6 +386,43 @@ Base.show(io::IO, 𝓂::ℳ) = println(io, # ForwardDiffExt extends this for ForwardDiff.Dual numbers. primal(x::Real) = x +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 normalize_filtering_options(filter::Symbol, smooth::Bool, @@ -2319,7 +2356,16 @@ end @warn "Invalid `shocks` argument. Provide a Symbol, Tuple, Vector, Matrix, or one of the documented selectors such as `:all`." shock_idx = Int64[] end - return shock_idx + + if shock_idx isa Integer + return shock_idx + else + shock_idx = unique(collect(shock_idx)) + if length(shock_idx) > 1 + sort!(shock_idx, by = i -> normalize_superscript(string(T.exo[i]))) + end + return shock_idx + end end