Simplify

Author: blegat

src/coloring.jl

@@ -67,7 +67,7 @@ function _hessian_color_preprocess(
         )
         tree_result = SMC.coloring(S, problem, algo)
         result = ColoringResult(tree_result, Int[])
-        return I, J, result
+        return ones(length(local_indices)), I, J, result
     end
 
     global_to_local_idx = seen_idx.nzidx # steal for storage
@@ -93,53 +93,15 @@ function _hessian_color_preprocess(
     )
     tree_result = SMC.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
-        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
-        for pos in first:last
-            (i_local, j_local) = reverse_bfs_orders[pos]
-            # Convert from local to global indices and normalize (lower triangle)
-            i_global = local_indices[i_local]
-            j_global = local_indices[j_local]
-            if j_global > i_global
-                i_global, j_global = j_global, i_global
-            end
-            k += 1
-            I_new[k] = i_global
-            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
+
+    # SparseMatrixColorings assumes that `I` and `J` are CSC-ordered
+    B = SMC.compress(S, tree_result)
+    C = SMC.decompress(B, tree_result)
+    I_sorted, J_sorted = SparseArrays.findnz(C)
+
+    return C.colptr, I_sorted, J_sorted, result
 end
 
 """
@@ -174,6 +136,7 @@ end
 
 """
     _recover_from_matmat!(
+        colptr::AbstractVector,
         V::AbstractVector{T},
         R::AbstractMatrix{T},
         result::ColoringResult,
@@ -185,59 +148,20 @@ R is the result of H*R_seed where R_seed is the seed matrix.
 `stored_values` is a temporary vector.
 """
 function _recover_from_matmat!(
+    colptr::AbstractVector,
     V::AbstractVector{T},
     R::AbstractMatrix{T},
     result::ColoringResult,
     stored_values::AbstractVector{T},
 ) where {T}
     tree_result = result.result
-    color = SMC.column_colors(tree_result)
     N = length(result.local_indices)
-    # Compute number of off-diagonal nonzeros from the length of V
-    # V contains N diagonal elements + nnz_offdiag off-diagonal elements
-    @assert length(stored_values) >= N
-
-    # Recover diagonal elements
-    k = 0
-    for i in 1:N
-        k += 1
-        if color[i] > 0
-            V[k] = R[i, color[i]]
-        else
-            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
-
-        # 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]
-            if color[j] > 0
-                value = R[i, color[j]] - stored_values[i]
-            else
-                value = zero(T)
-            end
-            stored_values[j] += value
-            k += 1
-            V[k] = value
-        end
-    end
-
-    @assert k == length(V)
+    S = _SparseMatrixValuesCSC(
+        N,
+        N,
+        colptr,
+        V,
+    )
+    SMC.decompress!(S, R, tree_result)
     return
 end

src/forward_over_reverse.jl

@@ -21,7 +21,7 @@ const MAX_CHUNK = 10
         f::_FunctionStorage,
         H::AbstractVector{Float64},
         λ::Float64,
-        offset::Int,
+        nzcount::Int,
     )::Int
 
 Evaluate the hessian matrix of the function `f` and store the result, scaled by
@@ -65,12 +65,13 @@ function _eval_hessian(
         )
     end
     # TODO(odow): consider reverting to a view.
+    N = size(ex.seed_matrix, 1)
     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_ϵ)
+    _recover_from_matmat!(ex.hess_colptr, output_slice, ex.seed_matrix, ex.rinfo, d.output_ϵ)
     for i in 1:length(output_slice)
         output_slice[i] *= scale
     end

src/types.jl

@@ -67,6 +67,7 @@ struct _FunctionStorage{R<:SMC.AbstractColoringResult}
     expr::_SubexpressionStorage
     grad_sparsity::Vector{Int}
     # Nonzero pattern of Hessian matrix
+    hess_colptr::Vector{Int}
     hess_I::Vector{Int}
     hess_J::Vector{Int}
     rinfo::Union{Nothing,ColoringResult{R}}
@@ -107,7 +108,7 @@ struct _FunctionStorage{R<:SMC.AbstractColoringResult}
                 subexpression_edgelist,
                 subexpression_variables,
             )
-            hess_I, hess_J, rinfo = _hessian_color_preprocess(
+            hess_colptr, hess_I, hess_J, rinfo = _hessian_color_preprocess(
                 edgelist,
                 num_variables,
                 coloring_algorithm,
@@ -117,6 +118,7 @@ struct _FunctionStorage{R<:SMC.AbstractColoringResult}
             return new{R}(
                 expr,
                 grad_sparsity,
+                hess_colptr,
                 hess_I,
                 hess_J,
                 rinfo,

src/utils.jl

@@ -203,3 +203,13 @@ function _reinterpret_unsafe(::Type{T}, x::Vector{R}) where {T,R}
     p = reinterpret(Ptr{T}, pointer(x))
     return _UnsafeVectorView(0, div(len, sizeof(T)), p)
 end
+
+struct _SparseMatrixValuesCSC{Tv,Ti<:Integer,CT<:AbstractVector{Ti},VT<:AbstractVector{Tv}} <: SparseArrays.AbstractSparseMatrixCSC{Tv,Ti}
+    m::Int          # Number of rows
+    n::Int          # Number of columns
+    colptr::CT      # Column i is in colptr[i]:(colptr[i+1]-1)
+    nzval::VT       # Stored values, typically nonzeros
+end
+
+Base.size(A::_SparseMatrixValuesCSC) = (A.m, A.n)
+SparseArrays.nonzeros(A::_SparseMatrixValuesCSC) = A.nzval

test/Project.toml

@@ -3,4 +3,5 @@ ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/ReverseAD.jl

@@ -484,7 +484,7 @@ function test_coloring_end_to_end_hessian_coloring_and_recovery()
         ArrayDiff.SMC.GreedyColoringAlgorithm(;
             decompression = :substitution,
         )
-    I, J, rinfo = ArrayDiff._hessian_color_preprocess(
+    colptr, I, J, rinfo = ArrayDiff._hessian_color_preprocess(
         Set([(1, 2)]),
         2,
         coloring_algorithm,
@@ -498,7 +498,7 @@ function test_coloring_end_to_end_hessian_coloring_and_recovery()
     hess = [3.4 2.1; 2.1 1.3]
     matmat = hess * R
     V = zeros(3)
-    ArrayDiff._recover_from_matmat!(V, matmat, rinfo, zeros(3))
+    ArrayDiff._recover_from_matmat!(colptr, V, matmat, rinfo, zeros(3))
     @test V == [3.4, 1.3, 2.1]
     return
 end