Add grouped covariance input to get_statistics and optimize determine_efficient_order

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}}}
+const String_input = Union{String,Vector{String},Matrix{String},Tuple{String,Vararg{String}},Vector{Vector{String}}}
 const ParameterType = Union{Nothing,
                             Pair{Symbol, Float64},
                             Pair{String, Float64},
@@ -2212,16 +2212,13 @@ function determine_efficient_order(𝐒₁::Matrix{<: Real},
             # 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
+                # Generate selector vector for columns where both states are dependencies
+                # Kronecker structure: column index = (i-1)*nˢ + j for states i,j
+                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)
+                    dependencies_in_states = dependencies_in_states .| affected
                 end
             end
 
@@ -2275,15 +2272,11 @@ function determine_efficient_order(𝐒₁::Matrix{<: Real},
                     # 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
+                        # 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)
+                            dependencies_in_states = dependencies_in_states .| affected
                         end
                     end
 
@@ -2403,32 +2396,22 @@ function determine_efficient_order(𝐒₁::Matrix{<: Real},
             # 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
+                # 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)
+                    dependencies_in_states = dependencies_in_states .| 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
+                # 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)
+                    dependencies_in_states = dependencies_in_states .| affected
                 end
             end
 
@@ -2501,32 +2484,22 @@ function determine_efficient_order(𝐒₁::Matrix{<: Real},
                     # 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
+                        # 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)
+                            dependencies_in_states = dependencies_in_states .| 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
+                        # 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)
+                            dependencies_in_states = dependencies_in_states .| affected
                         end
                     end
 
@@ -9060,7 +9033,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 +9053,14 @@ 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{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 +9085,39 @@ 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 vector (either String or Symbol)
+    return variables isa Vector{Vector{Symbol}} || variables isa Vector{Vector{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 vectors
+    if variables isa Vector{Vector{String}}
+        variables = [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 (variables is now Vector{Vector{Symbol}})
+    groups = Vector{Vector{Int}}()
+    for group in variables
+        if length(setdiff(group, T.var)) > 0
+            @warn "The following variables are not part of the model: " * join(string.(setdiff(group,T.var)),", ") * ". Use `get_variables(𝓂)` to list valid names."
+            push!(groups, Int[])
+        else
+            push!(groups, Int.(indexin(group, 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

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], [:q]])
+# output
+Dict{Symbol, AbstractArray{Float64}} with 1 entry:
+  :covariance => [...3x3 matrix with c-k covariances filled, q variance filled, and cross-group elements set to zero...]
 ```
 """
 function get_statistics(𝓂,
@@ -3283,6 +3289,10 @@ 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)
+
+    # Check if covariance is grouped and parse groups
+    is_grouped_covar = @ignore_derivatives is_grouped_covariance_input(covariance)
+    covar_groups = @ignore_derivatives is_grouped_covar ? parse_covariance_groups(covariance, 𝓂.timings) : nothing
 
     autocorr_var_idx = @ignore_derivatives parse_variables_input_to_index(autocorrelation, 𝓂.timings)
 
@@ -3405,8 +3415,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 is_grouped_covar
+            # 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,:] )