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},
@@ -9060,7 +9060,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 +9080,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 +9112,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 vector of covariance matrices. 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 => Matrix{Float64}[[...2x2 covariance matrix for c and k...], [...1x1 matrix for q...]]
 ```
 """
 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,31 @@ 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)
+            if solved
+                # Return a vector of covariance matrices, one for each group
+                covar_result = Vector{Matrix{T}}(undef, length(covar_groups))
+
+                for (i, group) in enumerate(covar_groups)
+                    covar_result[i] = covar_dcmp_sp[group, group]
+                end
+
+                ret[:covariance] = covar_result
+            else
+                # Return vector of Inf-filled matrices
+                covar_result = Vector{Matrix{T}}(undef, length(covar_groups))
+                for (i, group) in enumerate(covar_groups)
+                    group_size = length(group)
+                    covar_result[i] = fill(Inf * sum(abs2,parameter_values), group_size, group_size)
+                end
+                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,:] )

test/functionality_tests.jl

@@ -2622,6 +2622,38 @@ 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 Vector
+            @test length(stats_grouped[:covariance]) == 2
+            @test size(stats_grouped[:covariance][1]) == (2, 2)
+            @test size(stats_grouped[:covariance][2]) == (2, 2)
+
+            # Compare with non-grouped version for validation
+            stats_non_grouped = get_statistics(m, old_params,
+                                              algorithm = algorithm,
+                                              covariance = vars[2:3])
+
+            @test isapprox(stats_grouped[:covariance][1], stats_non_grouped[:covariance], rtol = 1e-10)
+
+            # Test with different group sizes
+            stats_varied = get_statistics(m, old_params,
+                                         algorithm = algorithm,
+                                         covariance = [[vars[2]], vars[3:5]])
+
+            @test length(stats_varied[:covariance]) == 2
+            @test size(stats_varied[:covariance][1]) == (1, 1)
+            @test size(stats_varied[:covariance][2]) == (3, 3)
+        end
+    end
+
 
     @testset "get_moments" begin
         for non_stochastic_steady_state in [true, false]