diff --git a/src/MacroModelling.jl b/src/MacroModelling.jl index 37e1e18f..c62bb4b2 100644 --- a/src/MacroModelling.jl +++ b/src/MacroModelling.jl @@ -71,8 +71,8 @@ Reexport.@reexport import SparseArrays: sparse, spzeros, droptol!, sparsevec, sp # Type definitions -const Symbol_input = Union{Symbol,Vector{Symbol},Matrix{Symbol},Tuple{Symbol,Vararg{Symbol}}} -const String_input = Union{String,Vector{String},Matrix{String},Tuple{String,Vararg{String}}} +const Symbol_input = Union{Symbol,Vector{Symbol},Matrix{Symbol},Tuple{Symbol,Vararg{Symbol}},Vector{Vector{Symbol}},Vector{Tuple{Symbol,Vararg{Symbol}}},Tuple{Tuple{Symbol,Vararg{Symbol}},Vararg{Tuple{Symbol,Vararg{Symbol}}}}} +const String_input = Union{String,Vector{String},Matrix{String},Tuple{String,Vararg{String}},Vector{Vector{String}},Vector{Tuple{String,Vararg{String}}},Tuple{Tuple{String,Vararg{String}},Vararg{Tuple{String,Vararg{String}}}}} const ParameterType = Union{Nothing, Pair{Symbol, Float64}, Pair{String, Float64}, @@ -2092,12 +2092,21 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, observables = T.var[var_idx] end + # Precompute state indices to avoid repeated indexin calls + state_idx_in_var = indexin(T.past_not_future_and_mixed, T.var) .|> Int + ๐’โ‚_states = ๐’โ‚[state_idx_in_var, 1:nหข] + for obs in observables obs_in_var_idx = indexin([obs],T.var) .|> Int dependencies_in_states = vec(sum(abs, ๐’โ‚[obs_in_var_idx,1:nหข], dims=1) .> tol) .> 0 - while dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[indexin(T.past_not_future_and_mixed, T.var),1:nหข]) .> tol) != dependencies_in_states - dependencies_in_states = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[indexin(T.past_not_future_and_mixed, T.var),1:nหข]) .> tol) + # Iterative propagation without redundant allocations + while true + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states) .> tol) + if new_deps == dependencies_in_states + break + end + dependencies_in_states = new_deps end dependencies = T.past_not_future_and_mixed[dependencies_in_states] @@ -2117,8 +2126,13 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, obs_in_var_idx = indexin([covar_var], T.var) .|> Int dependencies_in_states = vec(sum(abs, ๐’โ‚[obs_in_var_idx,1:nหข], dims=1) .> tol) .> 0 - while dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[indexin(T.past_not_future_and_mixed, T.var),1:nหข]) .> tol) != dependencies_in_states - dependencies_in_states = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[indexin(T.past_not_future_and_mixed, T.var),1:nหข]) .> tol) + # Iterative propagation without redundant allocations + while true + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states) .> tol) + if new_deps == dependencies_in_states + break + end + dependencies_in_states = new_deps end dependencies = T.past_not_future_and_mixed[dependencies_in_states] @@ -2126,12 +2140,20 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, end end + # Build lookup dictionary for faster searches + var_to_idx = Dict{Symbol, Int}() + for (idx, order) in enumerate(orders) + for var in order.first + var_to_idx[var] = idx + end + end + # Add entries for all pairs of covariance variables for i in 1:length(covariance_vars) for j in (i+1):length(covariance_vars) - # Find dependencies for both variables (they should exist now) - idx_i = findfirst(x -> covariance_vars[i] in x.first, orders) - idx_j = findfirst(x -> covariance_vars[j] in x.first, orders) + # Find dependencies for both variables using lookup dictionary + idx_i = var_to_idx[covariance_vars[i]] + idx_j = var_to_idx[covariance_vars[j]] deps_i = orders[idx_i].second deps_j = orders[idx_j].second @@ -2173,8 +2195,10 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, # Kronecker product indices for state-state interactions kron_s_s = โ„’.kron(s_in_sโบ, s_in_sโบ) - # State variables indices in full state vector + # Precompute state indices and matrix slices to avoid repeated operations state_idx_in_var = indexin(T.past_not_future_and_mixed, T.var) .|> Int + ๐’โ‚_states = ๐’โ‚[state_idx_in_var, 1:nหข] + ๐’โ‚‚_states = nnz(๐’โ‚‚) > 0 ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing for obs in observables obs_in_var_idx = indexin([obs],T.var) .|> Int @@ -2187,47 +2211,35 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_var_idx, kron_s_s] # Check which state variable pairs have influence - # The selected columns correspond to state-state products in order (1,1), (1,2), ..., (1,nหข), (2,1), ..., (nหข,nหข) - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if sum(abs, s_s_to_yโ‚‚[:, col_idx]) > tol - dependencies_in_states[i] = true - dependencies_in_states[j] = true - end - col_idx += 1 - end - end + # Vectorized approach: reshape to nหขร—nหข and check column/row sums + s_s_matrix = reshape(vec(sum(abs, s_s_to_yโ‚‚, dims=1) .> tol), nหข, nหข) + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_matrix, dims=2) .> 0) .| vec(sum(s_s_matrix, dims=1) .> 0) end # Propagate dependencies through the system (iterative closure) # considering both first and second order propagation while true - prev_dependencies = copy(dependencies_in_states) + prev_dependencies = dependencies_in_states # First order propagation - dependencies_in_states = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[state_idx_in_var, 1:nหข]) .> tol) + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states) .> tol) # Second order propagation: if state i and state j are dependencies, # their product can affect states - if nnz(๐’โ‚‚) > 0 - ๐’โ‚‚_states = ๐’โ‚‚[state_idx_in_var, kron_s_s] - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if prev_dependencies[i] && prev_dependencies[j] - # Check which states are affected by this product - affected = vec(sum(abs, ๐’โ‚‚_states[:, col_idx:col_idx], dims=2) .> tol) - dependencies_in_states = dependencies_in_states .| affected - end - col_idx += 1 - end + if !isnothing(๐’โ‚‚_states) + # Generate selector vector for columns where both states are dependencies + selector = vec(โ„’.kron(prev_dependencies, prev_dependencies)) + if any(selector) + # Check which states are affected by the selected products + affected = vec(sum(abs, ๐’โ‚‚_states[:, selector], dims=2) .> tol) + new_deps = new_deps .| affected end end - if dependencies_in_states == prev_dependencies + if new_deps == dependencies_in_states break end + dependencies_in_states = new_deps end dependencies = T.past_not_future_and_mixed[dependencies_in_states] @@ -2252,44 +2264,36 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, # Second order dependencies from quadratic terms (s โŠ— s) if nnz(๐’โ‚‚) > 0 s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_var_idx, kron_s_s] - - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if sum(abs, s_s_to_yโ‚‚[:, col_idx]) > tol - dependencies_in_states[i] = true - dependencies_in_states[j] = true - end - col_idx += 1 - end - end + # Vectorized approach: reshape to nหขร—nหข and check column/row sums + s_s_matrix = reshape(vec(sum(abs, s_s_to_yโ‚‚, dims=1) .> tol), nหข, nหข) + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_matrix, dims=2) .> 0) .| vec(sum(s_s_matrix, dims=1) .> 0) end # Propagate dependencies through the system + # Precompute matrix slices + ๐’โ‚_states_local = ๐’โ‚[state_idx_in_var, 1:nหข] + ๐’โ‚‚_states_local = nnz(๐’โ‚‚) > 0 ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing + while true - prev_dependencies = copy(dependencies_in_states) + prev_dependencies = dependencies_in_states # First order propagation - dependencies_in_states = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[state_idx_in_var, 1:nหข]) .> tol) + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states_local) .> tol) # Second order propagation - if nnz(๐’โ‚‚) > 0 - ๐’โ‚‚_states = ๐’โ‚‚[state_idx_in_var, kron_s_s] - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if prev_dependencies[i] && prev_dependencies[j] - affected = vec(sum(abs, ๐’โ‚‚_states[:, col_idx:col_idx], dims=2) .> tol) - dependencies_in_states = dependencies_in_states .| affected - end - col_idx += 1 - end + if !isnothing(๐’โ‚‚_states_local) + # Generate selector vector for columns where both states are dependencies + selector = vec(โ„’.kron(prev_dependencies, prev_dependencies)) + if any(selector) + affected = vec(sum(abs, ๐’โ‚‚_states_local[:, selector], dims=2) .> tol) + new_deps = new_deps .| affected end end - if dependencies_in_states == prev_dependencies + if new_deps == dependencies_in_states break end + dependencies_in_states = new_deps end dependencies = T.past_not_future_and_mixed[dependencies_in_states] @@ -2346,8 +2350,11 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, kron_s_s = โ„’.kron(s_in_sโบ, s_in_sโบ) kron_s_s_s = โ„’.kron(kron_s_s, s_in_sโบ) - # State variables indices in full state vector + # Precompute state indices and matrix slices state_idx_in_var = indexin(T.past_not_future_and_mixed, T.var) .|> Int + ๐’โ‚_states = ๐’โ‚[state_idx_in_var, 1:nหข] + ๐’โ‚‚_states = nnz(๐’โ‚‚) > 0 ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing + ๐’โ‚ƒ_states = nnz(๐’โ‚ƒ) > 0 ? ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] : nothing for obs in observables obs_in_var_idx = indexin([obs],T.var) .|> Int @@ -2358,83 +2365,53 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, # Second order dependencies from quadratic terms (s โŠ— s) if nnz(๐’โ‚‚) > 0 s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_var_idx, kron_s_s] - - # The selected columns correspond to state-state products in order (1,1), (1,2), ..., (nหข,nหข) - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if sum(abs, s_s_to_yโ‚‚[:, col_idx]) > tol - dependencies_in_states[i] = true - dependencies_in_states[j] = true - end - col_idx += 1 - end - end + # Vectorized approach: reshape and check row/column sums + s_s_matrix = reshape(vec(sum(abs, s_s_to_yโ‚‚, dims=1) .> tol), nหข, nหข) + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_matrix, dims=2) .> 0) .| vec(sum(s_s_matrix, dims=1) .> 0) end # Third order dependencies from cubic terms (s โŠ— s โŠ— s) if nnz(๐’โ‚ƒ) > 0 s_s_s_to_yโ‚ƒ = ๐’โ‚ƒ[obs_in_var_idx, kron_s_s_s] - - # The selected columns correspond to state-state-state products in order (1,1,1), (1,1,2), ..., (nหข,nหข,nหข) - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - for k in 1:nหข - if sum(abs, s_s_s_to_yโ‚ƒ[:, col_idx]) > tol - dependencies_in_states[i] = true - dependencies_in_states[j] = true - dependencies_in_states[k] = true - end - col_idx += 1 - end - end - end + # Vectorized approach: reshape to 3D and check along dimensions + s_s_s_tensor = reshape(vec(sum(abs, s_s_s_to_yโ‚ƒ, dims=1) .> tol), nหข, nหข, nหข) + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_s_tensor, dims=(2,3)) .> 0) .| + vec(sum(s_s_s_tensor, dims=(1,3)) .> 0) .| + vec(sum(s_s_s_tensor, dims=(1,2)) .> 0) end # Propagate dependencies through the system (iterative closure) # considering first, second, and third order propagation while true - prev_dependencies = copy(dependencies_in_states) + prev_dependencies = dependencies_in_states # First order propagation - dependencies_in_states = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[state_idx_in_var, 1:nหข]) .> tol) + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states) .> tol) # Second order propagation - if nnz(๐’โ‚‚) > 0 - ๐’โ‚‚_states = ๐’โ‚‚[state_idx_in_var, kron_s_s] - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if prev_dependencies[i] && prev_dependencies[j] - affected = vec(sum(abs, ๐’โ‚‚_states[:, col_idx:col_idx], dims=2) .> tol) - dependencies_in_states = dependencies_in_states .| affected - end - col_idx += 1 - end + if !isnothing(๐’โ‚‚_states) + # Generate selector vector for columns where both states are dependencies + selector = vec(โ„’.kron(prev_dependencies, prev_dependencies)) + if any(selector) + affected = vec(sum(abs, ๐’โ‚‚_states[:, selector], dims=2) .> tol) + new_deps = new_deps .| affected end end # Third order propagation - if nnz(๐’โ‚ƒ) > 0 - ๐’โ‚ƒ_states = ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - for k in 1:nหข - if prev_dependencies[i] && prev_dependencies[j] && prev_dependencies[k] - affected = vec(sum(abs, ๐’โ‚ƒ_states[:, col_idx:col_idx], dims=2) .> tol) - dependencies_in_states = dependencies_in_states .| affected - end - col_idx += 1 - end - end + if !isnothing(๐’โ‚ƒ_states) + # Generate selector vector for columns where all three states are dependencies + selector = vec(โ„’.kron(โ„’.kron(prev_dependencies, prev_dependencies), prev_dependencies)) + if any(selector) + affected = vec(sum(abs, ๐’โ‚ƒ_states[:, selector], dims=2) .> tol) + new_deps = new_deps .| affected end end - if dependencies_in_states == prev_dependencies + if new_deps == dependencies_in_states break end + dependencies_in_states = new_deps end dependencies = T.past_not_future_and_mixed[dependencies_in_states] @@ -2458,81 +2435,58 @@ function determine_efficient_order(๐’โ‚::Matrix{<: Real}, # Second order dependencies from quadratic terms (s โŠ— s) if nnz(๐’โ‚‚) > 0 - s_s_to_yโ‚‚ = ๐’โ‚‚[:, kron_s_s] - - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if sum(abs, s_s_to_yโ‚‚[obs_in_var_idx,col_idx]) > tol - dependencies_in_states[i] = true - dependencies_in_states[j] = true - end - col_idx += 1 - end - end + s_s_to_yโ‚‚ = ๐’โ‚‚[obs_in_var_idx, kron_s_s] + # Vectorized approach: reshape to nหขร—nหข and check column/row sums + s_s_matrix = reshape(vec(sum(abs, s_s_to_yโ‚‚, dims=1) .> tol), nหข, nหข) + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_matrix, dims=2) .> 0) .| vec(sum(s_s_matrix, dims=1) .> 0) end # Third order dependencies from cubic terms (s โŠ— s โŠ— s) if nnz(๐’โ‚ƒ) > 0 s_s_s_to_yโ‚ƒ = ๐’โ‚ƒ[obs_in_var_idx, kron_s_s_s] - - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - for k in 1:nหข - if sum(abs, s_s_s_to_yโ‚ƒ[:, col_idx]) > tol - dependencies_in_states[i] = true - dependencies_in_states[j] = true - dependencies_in_states[k] = true - end - col_idx += 1 - end - end - end + # Vectorized approach: reshape to 3D and check along dimensions + s_s_s_tensor = reshape(vec(sum(abs, s_s_s_to_yโ‚ƒ, dims=1) .> tol), nหข, nหข, nหข) + dependencies_in_states = dependencies_in_states .| vec(sum(s_s_s_tensor, dims=(2,3)) .> 0) .| + vec(sum(s_s_s_tensor, dims=(1,3)) .> 0) .| + vec(sum(s_s_s_tensor, dims=(1,2)) .> 0) end # Propagate dependencies through the system + # Precompute matrix slices + ๐’โ‚_states_local = ๐’โ‚[state_idx_in_var, 1:nหข] + ๐’โ‚‚_states_local = nnz(๐’โ‚‚) > 0 ? ๐’โ‚‚[state_idx_in_var, kron_s_s] : nothing + ๐’โ‚ƒ_states_local = nnz(๐’โ‚ƒ) > 0 ? ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] : nothing + while true - prev_dependencies = copy(dependencies_in_states) + prev_dependencies = dependencies_in_states # First order propagation - dependencies_in_states = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚[state_idx_in_var, 1:nหข]) .> tol) + new_deps = dependencies_in_states .| vec(abs.(dependencies_in_states' * ๐’โ‚_states_local) .> tol) # Second order propagation - if nnz(๐’โ‚‚) > 0 - ๐’โ‚‚_states = ๐’โ‚‚[state_idx_in_var, kron_s_s] - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - if prev_dependencies[i] && prev_dependencies[j] - affected = vec(sum(abs, ๐’โ‚‚_states[:, col_idx:col_idx], dims=2) .> tol) - dependencies_in_states = dependencies_in_states .| affected - end - col_idx += 1 - end + if !isnothing(๐’โ‚‚_states_local) + # Generate selector vector for columns where both states are dependencies + selector = vec(โ„’.kron(prev_dependencies, prev_dependencies)) + if any(selector) + affected = vec(sum(abs, ๐’โ‚‚_states_local[:, selector], dims=2) .> tol) + new_deps = new_deps .| affected end end # Third order propagation - if nnz(๐’โ‚ƒ) > 0 - ๐’โ‚ƒ_states = ๐’โ‚ƒ[state_idx_in_var, kron_s_s_s] - col_idx = 1 - for i in 1:nหข - for j in 1:nหข - for k in 1:nหข - if prev_dependencies[i] && prev_dependencies[j] && prev_dependencies[k] - affected = vec(sum(abs, ๐’โ‚ƒ_states[:, col_idx:col_idx], dims=2) .> tol) - dependencies_in_states = dependencies_in_states .| affected - end - col_idx += 1 - end - end + if !isnothing(๐’โ‚ƒ_states_local) + # Generate selector vector for columns where all three states are dependencies + selector = vec(โ„’.kron(โ„’.kron(prev_dependencies, prev_dependencies), prev_dependencies)) + if any(selector) + affected = vec(sum(abs, ๐’โ‚ƒ_states_local[:, selector], dims=2) .> tol) + new_deps = new_deps .| affected end end - if dependencies_in_states == prev_dependencies + if new_deps == dependencies_in_states break end + dependencies_in_states = new_deps end dependencies = T.past_not_future_and_mixed[dependencies_in_states] @@ -9060,7 +9014,12 @@ end function parse_variables_input_to_index(variables::Union{Symbol_input,String_input}, T::timings)::Union{UnitRange{Int}, Vector{Int}} - variables = variables isa String_input ? variables .|> Meta.parse .|> replace_indices : variables + # Handle nested vector conversion separately + if variables isa Vector{Vector{String}} + variables = [group .|> Meta.parse .|> replace_indices for group in variables] + elseif variables isa String_input + variables = variables .|> Meta.parse .|> replace_indices + end if variables == :all_excluding_auxiliary_and_obc return Int.(indexin(setdiff(T.var[.!contains.(string.(T.var),"แต’แต‡แถœ")],union(T.aux, T.exo_present)),sort(union(T.var,T.aux,T.exo_present)))) @@ -9075,6 +9034,30 @@ function parse_variables_input_to_index(variables::Union{Symbol_input,String_inp return Int[] end return getindex(1:length(T.var),convert(Vector{Bool},vec(sum(variables .== T.var,dims= 2)))) + elseif variables isa Vector{Vector{Symbol}} + # For grouped inputs, return union of all variables + all_vars = reduce(vcat, variables) + if length(setdiff(all_vars,T.var)) > 0 + @warn "The following variables are not part of the model: " * join(string.(setdiff(all_vars,T.var)),", ") * ". Use `get_variables(๐“‚)` to list valid names." + return Int[] + end + return Int.(indexin(unique(all_vars), T.var)) + elseif variables isa Vector{Tuple{Symbol,Vararg{Symbol}}} + # For grouped inputs with tuples, return union of all variables + all_vars = reduce(vcat, [collect(group) for group in variables]) + if length(setdiff(all_vars,T.var)) > 0 + @warn "The following variables are not part of the model: " * join(string.(setdiff(all_vars,T.var)),", ") * ". Use `get_variables(๐“‚)` to list valid names." + return Int[] + end + return Int.(indexin(unique(all_vars), T.var)) + elseif variables isa Tuple{Tuple{Symbol,Vararg{Symbol}},Vararg{Tuple{Symbol,Vararg{Symbol}}}} + # For grouped inputs with tuple of tuples, return union of all variables + all_vars = reduce(vcat, [collect(group) for group in variables]) + if length(setdiff(all_vars,T.var)) > 0 + @warn "The following variables are not part of the model: " * join(string.(setdiff(all_vars,T.var)),", ") * ". Use `get_variables(๐“‚)` to list valid names." + return Int[] + end + return Int.(indexin(unique(all_vars), T.var)) elseif variables isa Vector{Symbol} if length(setdiff(variables,T.var)) > 0 @warn "The following variables are not part of the model: " * join(string.(setdiff(variables,T.var)),", ") * ". Use `get_variables(๐“‚)` to list valid names." @@ -9099,6 +9082,47 @@ function parse_variables_input_to_index(variables::Union{Symbol_input,String_inp end end +# Helper function to check if input is grouped covariance format +function is_grouped_covariance_input(variables::Union{Symbol_input,String_input})::Bool + # Check if it's a nested structure (vector of vectors, vector of tuples, or tuple of tuples) + return variables isa Vector{Vector{Symbol}} || variables isa Vector{Vector{String}} || + variables isa Vector{Tuple{Symbol,Vararg{Symbol}}} || variables isa Vector{Tuple{String,Vararg{String}}} || + variables isa Tuple{Tuple{Symbol,Vararg{Symbol}},Vararg{Tuple{Symbol,Vararg{Symbol}}}} || + variables isa Tuple{Tuple{String,Vararg{String}},Vararg{Tuple{String,Vararg{String}}}} +end + +# Function to parse grouped covariance input into groups of indices +function parse_covariance_groups(variables::Union{Symbol_input,String_input}, T::timings)::Vector{Vector{Int}} + # Convert String_input to Symbol_input for nested structures + if variables isa Vector{Vector{String}} + variables = [group .|> Meta.parse .|> replace_indices for group in variables] + elseif variables isa Vector{Tuple{String,Vararg{String}}} + variables = [Tuple(group .|> Meta.parse .|> replace_indices) for group in variables] + elseif variables isa Tuple{Tuple{String,Vararg{String}},Vararg{Tuple{String,Vararg{String}}}} + variables = Tuple(Tuple(group .|> Meta.parse .|> replace_indices) for group in variables) + end + + if !is_grouped_covariance_input(variables) + # Not grouped, return single group + idx = parse_variables_input_to_index(variables, T) + return [collect(idx)] + end + + # Parse each group (convert tuples to vectors for uniform handling) + groups = Vector{Vector{Int}}() + for group in variables + group_vec = group isa Tuple ? collect(group) : group + if length(setdiff(group_vec, T.var)) > 0 + @warn "The following variables are not part of the model: " * join(string.(setdiff(group_vec,T.var)),", ") * ". Use `get_variables(๐“‚)` to list valid names." + push!(groups, Int[]) + else + push!(groups, Int.(indexin(group_vec, T.var))) + end + end + + return groups +end + function parse_shocks_input_to_index(shocks::Union{Symbol_input, String_input}, T::timings)#::Union{UnitRange{Int64}, Int64, Vector{Int64}} shocks = shocks isa String_input ? shocks .|> Meta.parse .|> replace_indices : shocks diff --git a/src/get_functions.jl b/src/get_functions.jl index 69764e4a..d606dd4f 100644 --- a/src/get_functions.jl +++ b/src/get_functions.jl @@ -3207,7 +3207,7 @@ If occasionally binding constraints are present in the model, they are not taken - `mean` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the mean of selected variables (the mean for the linearised solution is the NSSS). Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or \"y\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. - `standard_deviation` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the standard deviation of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or \"y\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. - `variance` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the variance of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or \"y\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. -- `covariance` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the covariance of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or \"y\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. +- `covariance` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the covariance of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or \"y\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. For grouped covariance computation, pass a `Vector` of `Vector`s (e.g. `[[:y, :c], [:k, :i]]`) to compute covariances only within each group, returning a single covariance matrix where cross-group covariances are set to zero. This allows more granular control over which covariances to compute. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all variables less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all variables less those related to occasionally binding constraints. `:all` will contain all variables. - `autocorrelation` [Default: `Symbol[]`, Type: `Union{Symbol_input,String_input}`]: variables for which to show the autocorrelation of selected variables. Inputs can be a variable name passed on as either a `Symbol` or `String` (e.g. `:y` or \"y\"), or `Tuple`, `Matrix` or `Vector` of `String` or `Symbol`. Any variables not part of the model will trigger a warning. `:all_excluding_auxiliary_and_obc` contains all shocks less those related to auxiliary variables and related to occasionally binding constraints (obc). `:all_excluding_obc` contains all shocks less those related to auxiliary variables. `:all` will contain all variables. - `autocorrelation_periods` [Default: `1:5`, Type = `UnitRange{Int}`]: periods for which to return the autocorrelation of selected variables - $ALGORITHMยฎ @@ -3243,6 +3243,12 @@ get_statistics(RBC, RBC.parameter_values, parameters = RBC.parameters, standard_ # output Dict{Symbol, AbstractArray{Float64}} with 1 entry: :standard_deviation => [0.0266642, 0.264677, 0.0739325, 0.0102062] + +# For grouped covariance (computing covariances only within specified groups): +get_statistics(RBC, RBC.parameter_values, covariance = [[:c, :k], [:y, :i]]) +# output +Dict{Symbol, AbstractArray{Float64}} with 1 entry: + :covariance => [...4x4 matrix with c-k covariances filled, y-i covariances filled, and cross-group elements set to zero...] ``` """ function get_statistics(๐“‚, @@ -3283,6 +3289,9 @@ function get_statistics(๐“‚, var_var_idx = @ignore_derivatives parse_variables_input_to_index(variance, ๐“‚.timings) covar_var_idx = @ignore_derivatives parse_variables_input_to_index(covariance, ๐“‚.timings) + + # Parse covariance groups if input is grouped format + covar_groups = @ignore_derivatives is_grouped_covariance_input(covariance) ? parse_covariance_groups(covariance, ๐“‚.timings) : nothing autocorr_var_idx = @ignore_derivatives parse_variables_input_to_index(autocorrelation, ๐“‚.timings) @@ -3405,8 +3414,38 @@ function get_statistics(๐“‚, # droptol!(covar_dcmp_sp,eps(Float64)) - # push!(ret,covar_dcmp_sp[covar_var_idx,covar_var_idx]) - ret[:covariance] = solved ? covar_dcmp_sp[covar_var_idx,covar_var_idx] : fill(Inf * sum(abs2,parameter_values),isnothing(covar_var_idx) ? 0 : length(covar_var_idx), isnothing(covar_var_idx) ? 0 : length(covar_var_idx)) + if !isnothing(covar_groups) + # Extract only the specified covariance groups (block diagonal structure) + # Return a single matrix with zeros for non-computed covariances + if solved + # Initialize matrix with zeros + covar_result = zeros(T, length(covar_var_idx), length(covar_var_idx)) + + # Fill in only the specified groups + for group in covar_groups + for (i_idx, i) in enumerate(group) + for (j_idx, j) in enumerate(group) + # Find position in covar_var_idx + i_pos = findfirst(==(i), covar_var_idx) + j_pos = findfirst(==(j), covar_var_idx) + if !isnothing(i_pos) && !isnothing(j_pos) + covar_result[i_pos, j_pos] = covar_dcmp_sp[i, j] + end + end + end + end + + ret[:covariance] = covar_result + else + # Return matrix with Inf-filled diagonal and zeros elsewhere + covar_result = fill(Inf * sum(abs2,parameter_values), length(covar_var_idx), length(covar_var_idx)) + ret[:covariance] = covar_result + end + else + # Original behavior for non-grouped input + # push!(ret,covar_dcmp_sp[covar_var_idx,covar_var_idx]) + ret[:covariance] = solved ? covar_dcmp_sp[covar_var_idx,covar_var_idx] : fill(Inf * sum(abs2,parameter_values),isnothing(covar_var_idx) ? 0 : length(covar_var_idx), isnothing(covar_var_idx) ? 0 : length(covar_var_idx)) + end end if !(autocorrelation == Symbol[]) # push!(ret,autocorr[autocorr_var_idx,:] ) diff --git a/test/functionality_tests.jl b/test/functionality_tests.jl index f60c252a..9c650a2f 100644 --- a/test/functionality_tests.jl +++ b/test/functionality_tests.jl @@ -2622,6 +2622,49 @@ function functionality_test(m, m2; algorithm = :first_order, plots = true) end + @testset "get_statistics - grouped covariance" begin + # Test grouped covariance functionality + if algorithm โˆˆ [:first_order, :pruned_second_order, :pruned_third_order] + # Test with 2 groups + stats_grouped = get_statistics(m, old_params, + algorithm = algorithm, + covariance = [vars[2:3], vars[4:5]]) + + @test haskey(stats_grouped, :covariance) + @test stats_grouped[:covariance] isa Matrix + @test size(stats_grouped[:covariance]) == (4, 4) + + # Compare with non-grouped version for validation + stats_non_grouped_1 = get_statistics(m, old_params, + algorithm = algorithm, + covariance = vars[2:3]) + + stats_non_grouped_2 = get_statistics(m, old_params, + algorithm = algorithm, + covariance = vars[4:5]) + + # Check that within-group covariances match + @test isapprox(stats_grouped[:covariance][1:2, 1:2], stats_non_grouped_1[:covariance], rtol = 1e-10) + @test isapprox(stats_grouped[:covariance][3:4, 3:4], stats_non_grouped_2[:covariance], rtol = 1e-10) + + # Check that cross-group covariances are zero + @test all(stats_grouped[:covariance][1:2, 3:4] .== 0) + @test all(stats_grouped[:covariance][3:4, 1:2] .== 0) + + # Test with different group sizes + stats_varied = get_statistics(m, old_params, + algorithm = algorithm, + covariance = [[vars[2]], vars[3:5]]) + + @test stats_varied[:covariance] isa Matrix + @test size(stats_varied[:covariance]) == (4, 4) + # First group is 1x1, second group is 3x3 + @test stats_varied[:covariance][1, 1] != 0 # within first group + @test all(stats_varied[:covariance][1, 2:4] .== 0) # cross-group + @test all(stats_varied[:covariance][2:4, 1] .== 0) # cross-group + end + end + @testset "get_moments" begin for non_stochastic_steady_state in [true, false]