Merge pull request #156 from thorek1/copilot/fix-20a9a78f-efdd-429a-8977-682bd45f4534

Author: thorek1

src/get_functions.jl

@@ -2725,11 +2725,11 @@ function get_moments(𝓂::ℳ;
         @info "Most of the time is spent calculating derivatives wrt parameters. If they are not needed, add `derivatives = false` as an argument to the function call." maxlog = DEFAULT_MAXLOG
     end 
 
-    if (!variance && !standard_deviation && !non_stochastic_steady_state && !mean)
+    if (!variance && !standard_deviation && !non_stochastic_steady_state && !mean && !covariance)
         derivatives = false
     end
 
-    if parameter_derivatives != :all && (variance || standard_deviation || non_stochastic_steady_state || mean)
+    if parameter_derivatives != :all && (variance || standard_deviation || non_stochastic_steady_state || mean || covariance)
         derivatives = true
     end
 
@@ -2878,14 +2878,30 @@ function get_moments(𝓂::ℳ;
 
 
         if covariance
+            axis3 = vcat(:Covariance, 𝓂.parameters[param_idx])
+        
+            if any(x -> contains(string(x), "◖"), axis3)
+                axis3_decomposed = decompose_name.(axis3)
+                axis3 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis3_decomposed]
+            end
+
             if algorithm == :pruned_second_order
                 covar_dcmp, Σᶻ₂, state_μ, Δμˢ₂, autocorr_tmp, ŝ_to_ŝ₂, ŝ_to_y₂, Σʸ₁, Σᶻ₁, SS_and_pars, 𝐒₁, ∇₁, 𝐒₂, ∇₂, solved = calculate_second_order_moments_with_covariance(𝓂.parameter_values, 𝓂, opts = opts)
+
+                # Compute covariance derivatives
+                dcovariance = 𝒟.jacobian(x -> vec(calculate_second_order_moments_with_covariance(x, 𝓂, opts = opts)[1]), backend, 𝓂.parameter_values)[:,param_idx]
             elseif algorithm == :pruned_third_order
                 covar_dcmp, state_μ, _, solved = calculate_third_order_moments(𝓂.parameter_values, :full_covar, 𝓂, opts = opts)
+
+                # Compute covariance derivatives
+                dcovariance = 𝒟.jacobian(x -> vec(calculate_third_order_moments(x, :full_covar, 𝓂, opts = opts)[1]), backend, 𝓂.parameter_values)[:,param_idx]
             else
                 covar_dcmp, ___, __, _, solved = calculate_covariance(𝓂.parameter_values, 𝓂, opts = opts)
 
                 @assert solved "Could not find covariance matrix."
+
+                # Compute covariance derivatives
+                dcovariance = 𝒟.jacobian(x -> vec(calculate_covariance(x, 𝓂, opts = opts)[1]), backend, 𝓂.parameter_values)[:,param_idx]
             end
         end
 
@@ -3045,8 +3061,38 @@ function get_moments(𝓂::ℳ;
             axis1 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis1_decomposed]
         end
 
-        # push!(ret,KeyedArray(covar_dcmp[var_idx, var_idx]; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1))
-        ret[:covariance] = KeyedArray(covar_dcmp[var_idx, var_idx]; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1)
+        if derivatives
+            # Reshape the flattened derivatives back to n x n x p structure
+            n_vars = length(var_idx)
+            n_params = length(param_idx)
+            
+            # Create array to hold covariance and derivatives: n x n x (1 + p)
+            covar_with_derivs = zeros(n_vars, n_vars, 1 + n_params)
+            
+            # First slice is the covariance matrix
+            covar_with_derivs[:, :, 1] = covar_dcmp[var_idx, var_idx]
+            
+            # Subsequent slices are derivatives wrt each parameter
+            for i in 1:n_params
+                # dcovariance[:,i] is vectorized, need to reshape to n x n
+                covar_with_derivs[:, :, i+1] = reshape(dcovariance[:, i], n_vars, n_vars)
+            end
+            
+            # Create axis names
+            if !@isdefined axis3
+                axis3 = vcat(:Covariance, 𝓂.parameters[param_idx])
+            
+                if any(x -> contains(string(x), "◖"), axis3)
+                    axis3_decomposed = decompose_name.(axis3)
+                    axis3 = [length(a) > 1 ? string(a[1]) * "{" * join(a[2],"}{") * "}" * (a[end] isa Symbol ? string(a[end]) : "") : string(a[1]) for a in axis3_decomposed]
+                end
+            end
+            
+            ret[:covariance] = KeyedArray(covar_with_derivs; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1, Covariance_and_∂covariance∂parameter = axis3)
+        else
+            # push!(ret,KeyedArray(covar_dcmp[var_idx, var_idx]; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1))
+            ret[:covariance] = KeyedArray(covar_dcmp[var_idx, var_idx]; Variables = axis1, 𝑉𝑎𝑟𝑖𝑎𝑏𝑙𝑒𝑠 = axis1)
+        end
     end
 
     return ret