Fix format

Author: blegat

src/ArrayDiff.jl

@@ -17,11 +17,18 @@ import SparseMatrixColorings
 
 Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
 """
-struct Mode{C<:SparseMatrixColorings.GreedyColoringAlgorithm} <: MOI.Nonlinear.AbstractAutomaticDifferentiation
+struct Mode{C<:SparseMatrixColorings.GreedyColoringAlgorithm} <:
+       MOI.Nonlinear.AbstractAutomaticDifferentiation
     coloring_algorithm::C
 end
 
-Mode() = Mode(SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution))
+function Mode()
+    return Mode(
+        SparseMatrixColorings.GreedyColoringAlgorithm(;
+            decompression = :substitution,
+        ),
+    )
+end
 
 function MOI.Nonlinear.Evaluator(
     model::MOI.Nonlinear.Model,

src/coloring_compat.jl

@@ -48,7 +48,7 @@ function _hessian_color_preprocess(
     end
     local_indices = sort!(collect(seen_idx))
     empty!(seen_idx)
-    
+
     # Handle empty case (no edges in Hessian)
     if isempty(local_indices)
         # Return empty structure - no variables to color
@@ -57,24 +57,30 @@ function _hessian_color_preprocess(
         # For the result, we'll create a minimal valid structure with a diagonal element
         # Note: This case should rarely occur in practice
         S = SparseArrays.spdiagm(0 => [true])
-        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        problem = SparseMatrixColorings.ColoringProblem(;
+            structure = :symmetric,
+            partition = :column,
+        )
         tree_result = SparseMatrixColorings.coloring(S, problem, algo)
         result = ColoringResult(tree_result, Int[])
         return I, J, result
     end
-    
+
     # Also handle case where we have vertices but no edges (diagonal-only Hessian)
     if isempty(I)
         # Create identity matrix pattern (diagonal only)
         n = length(local_indices)
         S = SparseArrays.spdiagm(0 => trues(n))
-        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        problem = SparseMatrixColorings.ColoringProblem(;
+            structure = :symmetric,
+            partition = :column,
+        )
         tree_result = SparseMatrixColorings.coloring(S, problem, algo)
         result = ColoringResult(tree_result, local_indices)
         # I and J are already empty, which is correct for no off-diagonal elements
         return I, J, result
     end
-    
+
     global_to_local_idx = seen_idx.nzidx # steal for storage
     for k in eachindex(local_indices)
         global_to_local_idx[local_indices[k]] = k
@@ -84,7 +90,7 @@ function _hessian_color_preprocess(
         I[k] = global_to_local_idx[I[k]]
         J[k] = global_to_local_idx[J[k]]
     end
-    
+
     # Create sparsity pattern matrix
     n = length(local_indices)
     S = SparseArrays.spzeros(Bool, n, n)
@@ -93,39 +99,42 @@ function _hessian_color_preprocess(
         S[i, j] = true
         S[j, i] = true  # symmetric
     end
-    
+
     # Perform coloring using SparseMatrixColorings
-    problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+    problem = SparseMatrixColorings.ColoringProblem(;
+        structure = :symmetric,
+        partition = :column,
+    )
     tree_result = SparseMatrixColorings.coloring(S, problem, algo)
-    
+
     # Reconstruct I and J from the tree structure (matching original _indirect_recover_structure)
     # First add all diagonal elements
     N = length(local_indices)
-    
+
     # Count off-diagonal elements from tree structure
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_result
     nnz_offdiag = 0
     for tree_idx in 1:nt
         first = tree_edge_indices[tree_idx]
-        last = tree_edge_indices[tree_idx + 1] - 1
+        last = tree_edge_indices[tree_idx+1] - 1
         nnz_offdiag += (last - first + 1)
     end
-    
+
     I_new = Vector{Int}(undef, N + nnz_offdiag)
     J_new = Vector{Int}(undef, N + nnz_offdiag)
     k = 0
-    
+
     # Add all diagonal elements
     for i in 1:N
         k += 1
         I_new[k] = local_indices[i]
         J_new[k] = local_indices[i]
     end
-    
+
     # Then add off-diagonal elements from the tree structure
     for tree_idx in 1:nt
         first = tree_edge_indices[tree_idx]
-        last = tree_edge_indices[tree_idx + 1] - 1
+        last = tree_edge_indices[tree_idx+1] - 1
         for pos in first:last
             (i_local, j_local) = reverse_bfs_orders[pos]
             # Convert from local to global indices and normalize (lower triangle)
@@ -139,9 +148,9 @@ function _hessian_color_preprocess(
             J_new[k] = j_global
         end
     end
-    
+
     @assert k == length(I_new)
-    
+
     # Wrap result with local_indices
     result = ColoringResult(tree_result, local_indices)
     return I_new, J_new, result
@@ -203,7 +212,7 @@ function _recover_from_matmat!(
     # V contains N diagonal elements + nnz_offdiag off-diagonal elements
     nnz_offdiag = length(V) - N
     @assert length(stored_values) >= N
-    
+
     # Recover diagonal elements
     k = 0
     for i in 1:N
@@ -214,23 +223,23 @@ function _recover_from_matmat!(
             V[k] = zero(T)
         end
     end
-    
+
     # Recover off-diagonal elements using tree structure
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_result
     fill!(stored_values, zero(T))
-    
+
     for tree_idx in 1:nt
         first = tree_edge_indices[tree_idx]
-        last = tree_edge_indices[tree_idx + 1] - 1
-        
+        last = tree_edge_indices[tree_idx+1] - 1
+
         # Reset stored_values for vertices in this tree
         for pos in first:last
             (vertex, _) = reverse_bfs_orders[pos]
             stored_values[vertex] = zero(T)
         end
         (_, root) = reverse_bfs_orders[last]
         stored_values[root] = zero(T)
-        
+
         # Recover edge values
         for pos in first:last
             (i, j) = reverse_bfs_orders[pos]
@@ -244,7 +253,7 @@ function _recover_from_matmat!(
             V[k] = value
         end
     end
-    
+
     @assert k == length(V)
     return
 end

src/forward_over_reverse.jl

@@ -65,13 +65,12 @@ function _eval_hessian(
         )
     end
     # TODO(odow): consider reverting to a view.
-    output_slice = _UnsafeVectorView{Float64}(nzcount, length(ex.hess_I), pointer(H))::_UnsafeVectorView{Float64}
-    _recover_from_matmat!(
-        output_slice,
-        ex.seed_matrix,
-        ex.rinfo,
-        d.output_ϵ,
-    )
+    output_slice = _UnsafeVectorView{Float64}(
+        nzcount,
+        length(ex.hess_I),
+        pointer(H),
+    )::_UnsafeVectorView{Float64}
+    _recover_from_matmat!(output_slice, ex.seed_matrix, ex.rinfo, d.output_ϵ)
     for i in 1:length(output_slice)
         output_slice[i] *= scale
     end

src/mathoptinterface_api.jl

@@ -19,7 +19,10 @@ function MOI.features_available(d::NLPEvaluator)
     return [:Grad, :Jac, :JacVec, :Hess, :HessVec]
 end
 
-function MOI.initialize(d::NLPEvaluator{R}, requested_features::Vector{Symbol}) where {R}
+function MOI.initialize(
+    d::NLPEvaluator{R},
+    requested_features::Vector{Symbol},
+) where {R}
     # Check that we support the features requested by the user.
     available_features = MOI.features_available(d)
     for feature in requested_features

src/types.jl

@@ -78,7 +78,10 @@ struct _FunctionStorage{R<:SparseMatrixColorings.AbstractColoringResult}
         expr::_SubexpressionStorage,
         num_variables,
         coloring_storage::IndexedSet,
-        coloring_algorithm::Union{Nothing,SparseMatrixColorings.GreedyColoringAlgorithm},
+        coloring_algorithm::Union{
+            Nothing,
+            SparseMatrixColorings.GreedyColoringAlgorithm,
+        },
         subexpressions::Vector{_SubexpressionStorage},
         dependent_subexpressions,
         subexpression_edgelist,
@@ -147,7 +150,10 @@ interface.
 !!! warning
     Before using, you must initialize the evaluator using `MOI.initialize`.
 """
-mutable struct NLPEvaluator{R,C<:SparseMatrixColorings.GreedyColoringAlgorithm} <: MOI.AbstractNLPEvaluator
+mutable struct NLPEvaluator{
+    R,
+    C<:SparseMatrixColorings.GreedyColoringAlgorithm,
+} <: MOI.AbstractNLPEvaluator
     data::Nonlinear.Model
     ordered_variables::Vector{MOI.VariableIndex}
     coloring_algorithm::C
@@ -187,9 +193,14 @@ mutable struct NLPEvaluator{R,C<:SparseMatrixColorings.GreedyColoringAlgorithm}
     function NLPEvaluator(
         data::Nonlinear.Model,
         ordered_variables::Vector{MOI.VariableIndex},
-        coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution),
+        coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(;
+            decompression = :substitution,
+        ),
     )
-        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        problem = SparseMatrixColorings.ColoringProblem(;
+            structure = :symmetric,
+            partition = :column,
+        )
         C = typeof(coloring_algorithm)
         R = Base.promote_op(
             SparseMatrixColorings.coloring,

test/ReverseAD.jl

@@ -480,8 +480,16 @@ end
 
 function test_coloring_end_to_end_hessian_coloring_and_recovery()
     # Test the new coloring API through the compatibility layer
-    coloring_algorithm = ArrayDiff.SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
-    I, J, rinfo = ArrayDiff._hessian_color_preprocess(Set([(1, 2)]), 2, coloring_algorithm, ArrayDiff.IndexedSet(0))
+    coloring_algorithm =
+        ArrayDiff.SparseMatrixColorings.GreedyColoringAlgorithm(;
+            decompression = :substitution,
+        )
+    I, J, rinfo = ArrayDiff._hessian_color_preprocess(
+        Set([(1, 2)]),
+        2,
+        coloring_algorithm,
+        ArrayDiff.IndexedSet(0),
+    )
     R = ArrayDiff._seed_matrix(rinfo)
     ArrayDiff._prepare_seed_matrix!(R, rinfo)
     @test I == [1, 2, 2]