From 56eebe6fa4a33bf53d036124ccbb5ee2e29dc246 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 20 Nov 2025 12:16:11 +1300 Subject: [PATCH 01/14] Replace Coloring by SparseMatrixColoring --- .github/workflows/ci.yml | 7 + Project.toml | 4 +- src/ArrayDiff.jl | 37 +- src/Coloring/Coloring.jl | 574 ------------------------------- src/Coloring/topological_sort.jl | 106 ------ src/coloring.jl | 140 ++++++++ src/forward_over_reverse.jl | 14 +- src/graph_tools.jl | 10 +- src/mathoptinterface_api.jl | 20 +- src/types.jl | 60 +++- src/utils.jl | 15 + test/Project.toml | 1 + test/test_ArrayDiff.jl | 28 +- test/test_ReverseAD.jl | 148 ++++---- 14 files changed, 372 insertions(+), 792 deletions(-) delete mode 100644 src/Coloring/Coloring.jl delete mode 100644 src/Coloring/topological_sort.jl create mode 100644 src/coloring.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 79a4c27..247a219 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,6 +25,13 @@ jobs: version: ${{ matrix.version }} arch: ${{ matrix.arch }} - uses: julia-actions/cache@v3 + - name: MOI + shell: julia --project=@. {0} + run: | + using Pkg + Pkg.add([ + PackageSpec(url="https://github.com/blegat/SparseMatrixColorings.jl/", rev="abstractsparse"), + ]) - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 with: diff --git a/Project.toml b/Project.toml index 8cc0a91..7ceed19 100644 --- a/Project.toml +++ b/Project.toml @@ -6,7 +6,6 @@ authors = ["Sophie Lequeu ", "Benoît Legat length(v) - @assert v.nnz == 0 # only resize empty vector - resize!(v.nzidx, n) - resize!(v.empty, n) - fill!(v.empty, true) - end - return -end - -Base.collect(v::IndexedSet) = v.nzidx[1:(v.nnz)] - -function Base.union!(v::IndexedSet, s) - for x in s - push!(v, x) - end - return -end - -""" - struct UndirectedGraph - adjlist::Vector{Int} - edgeindex::Vector{Int} - offsets::Vector{Int} - edges::Vector{Tuple{Int,Int}} - end - -Compact storage for an undirected graph. The number of nodes is given by -`length(offsets) - 1`. The edges of node `u` are given by `edges[e]` for -`e in edgeindex[offsets[u]:(offsets[u+1] - 1)]`. The neighbors are also -stored at `adjlist[offsets[u]:(offsets[u+1] - 1)]`. -""" -struct UndirectedGraph - adjlist::Vector{Int} - edgeindex::Vector{Int} # corresponding edge number, indexed as adjlist - offsets::Vector{Int} - edges::Vector{Tuple{Int,Int}} -end - -_num_vertices(g::UndirectedGraph) = length(g.offsets) - 1 - -_num_edges(g::UndirectedGraph) = length(g.edges) - -_num_neighbors(i::Int, g::UndirectedGraph) = g.offsets[i+1] - g.offsets[i] - -_start_neighbors(i::Int, g::UndirectedGraph) = g.offsets[i] - -function UndirectedGraph(I, J, nel) - adjcount = zeros(Int, nel) - n_edges = 0 - for k in eachindex(I) - i = I[k] - j = J[k] - if i == j - continue - end - n_edges += 1 - adjcount[i] += 1 - adjcount[j] += 1 - end - offsets = Vector{Int}(undef, nel + 1) - offsets[1] = 1 - for k in 1:nel - offsets[k+1] = offsets[k] + adjcount[k] - end - fill!(adjcount, 0) - edges = Vector{Tuple{Int,Int}}(undef, n_edges) - adjlist = Vector{Int}(undef, offsets[nel+1] - 1) - edgeindex = Vector{Int}(undef, length(adjlist)) - edge_count = 0 - for k in eachindex(I) - i = I[k] - j = J[k] - if i == j - continue - end - edge_count += 1 - adjlist[offsets[i]+adjcount[i]] = j - edgeindex[offsets[i]+adjcount[i]] = edge_count - adjcount[i] += 1 - adjlist[offsets[j]+adjcount[j]] = i - edgeindex[offsets[j]+adjcount[j]] = edge_count - adjcount[j] += 1 - edges[edge_count] = (i, j) - end - @assert edge_count == n_edges - return UndirectedGraph(adjlist, edgeindex, offsets, edges) -end - -struct _Edge - index::Int - source::Int - target::Int -end - -function _prevent_cycle( - v, - w, - x, - e_idx1, - e_idx2, - S, - firstVisitToTree, - forbiddenColors, - color, -) - er = DataStructures.find_root!(S, e_idx2) - @inbounds first = firstVisitToTree[er] - p = first.source # but this depends on the order? - q = first.target - @inbounds if p != v - firstVisitToTree[er] = _Edge(e_idx1, v, w) - elseif q != w - forbiddenColors[color[x]] = v - end - return -end - -function _grow_star(v, w, e_idx, firstNeighbor, color, S) - @inbounds e = firstNeighbor[color[w]] - p = e.source - @inbounds if p != v - firstNeighbor[color[w]] = _Edge(e_idx, v, w) - else - union!(S, e_idx, e.index) - end - return -end - -function _merge_trees(eg, eg1, S) - e1 = DataStructures.find_root!(S, eg) - e2 = DataStructures.find_root!(S, eg1) - if e1 != e2 - union!(S, eg, eg1) - end - return -end - -""" - acyclic_coloring(g::UndirectedGraph) - -Implement the acyclic coloring algorithm of Gebremdehin, Tarafdar, Manne, and -Pothen, "New Acyclic and Star Coloring Algorithms with Application to Computing -Hessians." SIAM J. Sci. Comput. 2007. - -Returns `Tuple{Vector{Int},Int}` giving the color index of each node in `g`, as -well as the total number of colors used. -""" -function acyclic_coloring(g::UndirectedGraph) - if _num_edges(g) == 0 - return fill(1, _num_vertices(g)), 1 - end - num_colors = 0 - forbiddenColors = Int[] - firstNeighbor = _Edge[] - firstVisitToTree = fill(_Edge(0, 0, 0), _num_edges(g)) - color = fill(0, _num_vertices(g)) - # disjoint set forest of edges in the graph - S = DataStructures.IntDisjointSet{Int}(_num_edges(g)) - @inbounds for v in 1:_num_vertices(g) - n_neighbor = _num_neighbors(v, g) - start_neighbor = _start_neighbors(v, g) - for k in 0:(n_neighbor-1) - w = g.adjlist[start_neighbor+k] - if color[w] == 0 - continue - end - forbiddenColors[color[w]] = v - end - for k in 0:(n_neighbor-1) - w = g.adjlist[start_neighbor+k] - e_idx = g.edgeindex[start_neighbor+k] - if color[w] == 0 - continue - end - n_neighbor_w = _num_neighbors(w, g) - start_neighbor_w = _start_neighbors(w, g) - for k2 in 0:(n_neighbor_w-1) - x = g.adjlist[start_neighbor_w+k2] - e2_idx = g.edgeindex[start_neighbor_w+k2] - if color[x] == 0 - continue - end - if forbiddenColors[color[x]] != v - _prevent_cycle( - v, - w, - x, - e_idx, - e2_idx, - S, - firstVisitToTree, - forbiddenColors, - color, - ) - end - end - end - # find feasible color - found = false - for k in 1:num_colors - if forbiddenColors[k] != v - color[v] = k - found = true - break - end - end - if !found - num_colors += 1 - push!(forbiddenColors, 0) - push!(firstNeighbor, _Edge(0, 0, 0)) - color[v] = num_colors - end - for k in 0:(n_neighbor-1) - w = g.adjlist[start_neighbor+k] - e_idx = g.edgeindex[start_neighbor+k] - if color[w] == 0 - continue - end - _grow_star(v, w, e_idx, firstNeighbor, color, S) - end - for k in 0:(n_neighbor-1) - w = g.adjlist[start_neighbor+k] - e_idx = g.edgeindex[start_neighbor+k] - if color[w] == 0 - continue - end - n_neighbor_w = _num_neighbors(w, g) - start_neighbor_w = _start_neighbors(w, g) - for k2 in 0:(n_neighbor_w-1) - x = g.adjlist[start_neighbor_w+k2] - e2_idx = g.edgeindex[start_neighbor_w+k2] - if color[x] == 0 || x == v - continue - end - if color[x] == color[v] - _merge_trees(e_idx, e2_idx, S) - end - end - end - end - return color, num_colors -end - -struct RecoveryInfo - vertexmap::Vector{Vector{Int}} - postorder::Vector{Vector{Int}} - parents::Vector{Vector{Int}} - color::Vector{Int} - num_colors::Int - nnz::Int # number of off-diagonal nonzeros - local_indices::Vector{Int} # map back to global indices -end - -function RecoveryInfo() - return RecoveryInfo( - Vector{Vector{Int}}(undef, 0), - Vector{Vector{Int}}(undef, 0), - Vector{Vector{Int}}(undef, 0), - Vector{Int}(undef, 0), - 0, - 0, - Vector{Int}(undef, 0), - ) -end - -function recovery_preprocess( - g::UndirectedGraph, - color, - num_colors, - local_indices, -) - # represent two-color subgraph as: - # list of vertices (with map to global indices) - # adjacency list in a single vector (with list of offsets) - # linear index of pair of colors - twocolorindex = zeros(Int32, num_colors, num_colors) - seen_twocolors = 0 - # count of edges in each subgraph - edge_count = Int[] - for k in eachindex(g.edges) - u, v = g.edges[k] - i = min(color[u], color[v]) - j = max(color[u], color[v]) - if twocolorindex[i, j] == 0 - seen_twocolors += 1 - twocolorindex[i, j] = seen_twocolors - push!(edge_count, 0) - end - idx = twocolorindex[i, j] - edge_count[idx] += 1 - end - # edges sorted by twocolor subgraph - sorted_edges = Vector{Vector{Tuple{Int,Int}}}(undef, seen_twocolors) - for idx in 1:seen_twocolors - sorted_edges[idx] = Tuple{Int,Int}[] - sizehint!(sorted_edges[idx], edge_count[idx]) - end - for k in eachindex(g.edges) - u, v = g.edges[k] - i = min(color[u], color[v]) - j = max(color[u], color[v]) - idx = twocolorindex[i, j] - push!(sorted_edges[idx], (u, v)) - end - # list of unique vertices in each twocolor subgraph - vertexmap = Vector{Vector{Int}}(undef, seen_twocolors) - postorder = Vector{Vector{Int}}(undef, seen_twocolors) - parents = Vector{Vector{Int}}(undef, seen_twocolors) - # temporary lookup map from global index to subgraph index - revmap = zeros(Int, _num_vertices(g)) - adjcount = zeros(Int, _num_vertices(g)) - cmap = zeros(Int, 0) # shared storage for DFS - for idx in 1:seen_twocolors - my_edges = sorted_edges[idx] - vlist = Int[] - # build up the vertex list and adjacency count - for k in eachindex(my_edges) - u, v = my_edges[k] - # seen these vertices yet? - if revmap[u] == 0 - push!(vlist, u) - revmap[u] = length(vlist) - adjcount[u] = 0 - end - if revmap[v] == 0 - push!(vlist, v) - revmap[v] = length(vlist) - adjcount[v] = 0 - end - adjcount[u] += 1 - adjcount[v] += 1 - end - # set up offsets for adjlist - offset = Vector{Int}(undef, length(vlist) + 1) - offset[1] = 1 - for k in eachindex(vlist) - offset[k+1] = offset[k] + adjcount[vlist[k]] - adjcount[vlist[k]] = 0 - end - # adjlist for node u in twocolor idx starts at - # vec[offset[u]] - # u has global index vlist[u] - vec = Vector{Int}(undef, offset[length(vlist)+1] - 1) - # now fill in - for k in eachindex(my_edges) - u, v = my_edges[k] - u_rev = revmap[u] # indices in the subgraph - v_rev = revmap[v] - vec[offset[u_rev]+adjcount[u]] = v_rev - vec[offset[v_rev]+adjcount[v]] = u_rev - adjcount[u] += 1 - adjcount[v] += 1 - end - resize!(cmap, length(vlist)) - order, parent = - reverse_topological_sort_by_dfs(vec, offset, length(vlist), cmap) - for k in eachindex(vlist) - revmap[vlist[k]] = 0 # clear for reuse - end - postorder[idx] = order - parents[idx] = parent - vertexmap[idx] = vlist - end - return RecoveryInfo( - vertexmap, - postorder, - parents, - color, - num_colors, - _num_edges(g), - local_indices, - ) -end - -_normalize(i, j) = (j > i) ? (j, i) : (i, j) - -function _indirect_recover_structure(rinfo::RecoveryInfo) - N = length(rinfo.color) - I = zeros(Int, rinfo.nnz + N) - J = zeros(Int, rinfo.nnz + N) - k = 0 - for i in 1:N - k += 1 - I[k] = J[k] = i - end - for t in eachindex(rinfo.postorder) - vmap = rinfo.vertexmap[t] - order = rinfo.postorder[t] - parent = rinfo.parents[t] - for z in eachindex(order) - v = order[z] - p = parent[v] - if p == 0 - continue - end - k += 1 - I[k], J[k] = _normalize(vmap[v], vmap[p]) - end - end - @assert k == rinfo.nnz + N - return I, J -end - -""" - hessian_color_preprocess( - edgelist, - num_total_var, - seen_idx = IndexedSet(0), - ) - -`edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the -diagonal. -""" -function hessian_color_preprocess( - edgelist, - num_total_var, - seen_idx = IndexedSet(0), -) - resize!(seen_idx, num_total_var) - I, J = Int[], Int[] - for (i, j) in edgelist - push!(seen_idx, i) - push!(seen_idx, j) - push!(I, i) - push!(J, j) - end - local_indices = sort!(collect(seen_idx)) - empty!(seen_idx) - global_to_local_idx = seen_idx.nzidx # steal for storage - for k in eachindex(local_indices) - global_to_local_idx[local_indices[k]] = k - end - # only do the coloring on the local indices - for k in eachindex(I) - I[k] = global_to_local_idx[I[k]] - J[k] = global_to_local_idx[J[k]] - end - g = UndirectedGraph(I, J, length(local_indices)) - color, num_colors = acyclic_coloring(g) - @assert length(color) == _num_vertices(g) - rinfo = recovery_preprocess(g, color, num_colors, local_indices) - I, J = _indirect_recover_structure(rinfo) - # convert back to global indices - for k in eachindex(I) - I[k] = local_indices[I[k]] - J[k] = local_indices[J[k]] - end - return I, J, rinfo -end - -""" - seed_matrix(rinfo::RecoveryInfo) - -Allocate a seed matrix for the Coloring. -""" -function seed_matrix(rinfo::RecoveryInfo) - return Matrix{Float64}(undef, length(rinfo.local_indices), rinfo.num_colors) -end - -""" - prepare_seed_matrix!(R, rinfo::RecoveryInfo) -""" -function prepare_seed_matrix!(R, rinfo::RecoveryInfo) - N = length(rinfo.color) - @assert N == size(R, 1) == length(rinfo.local_indices) - @assert size(R, 2) == rinfo.num_colors - fill!(R, 0.0) - for i in 1:N - R[i, rinfo.color[i]] = 1 - end - return -end - -""" - recover_from_matmat!( - V::AbstractVector{T}, - R::AbstractMatrix{T}, - rinfo::RecoveryInfo, - stored_values::AbstractVector{T}, - ) where {T} - -recover the hessian values from the hessian-matrix solution -stored_values is a temporary vector of length >= length(rinfo.local_indices) -""" -function recover_from_matmat!( - V::AbstractVector{T}, - R::AbstractMatrix{T}, - rinfo::RecoveryInfo, - stored_values::AbstractVector{T}, -) where {T} - N = length(rinfo.color) # number of local variables - @assert length(V) == rinfo.nnz + N - @assert length(stored_values) >= length(rinfo.local_indices) - k = 0 - for i in 1:N - k += 1 - V[k] = R[i, rinfo.color[i]] - end - for t in eachindex(rinfo.vertexmap) - vmap = rinfo.vertexmap[t] - order = rinfo.postorder[t] - parent = rinfo.parents[t] - for z in eachindex(order) - @inbounds stored_values[z] = 0.0 - end - @inbounds for z in eachindex(order) - v = order[z] - p = parent[v] - if p == 0 - continue - end - i, j = vmap[v], vmap[p] - value = R[i, rinfo.color[j]] - stored_values[v] - stored_values[p] += value - k += 1 - V[k] = value - end - end - @assert k == rinfo.nnz + N - return -end - -end # module diff --git a/src/Coloring/topological_sort.jl b/src/Coloring/topological_sort.jl deleted file mode 100644 index 0b8d0ca..0000000 --- a/src/Coloring/topological_sort.jl +++ /dev/null @@ -1,106 +0,0 @@ -# Modified from Graphs.jl. -# Workaround for edge_index not being O(1) on SimpleGraph. -# edge_index was only needed to test for cycles, so -# this implementation skips that check. - -# Graphs.jl is licensed under the MIT License: -# -# Copyright (c) 2012: John Myles White and other contributors. -# -# Permission is hereby granted, free of charge, to any person obtaining -# a copy of this software and associated documentation files (the -# "Software"), to deal in the Software without restriction, including -# without limitation the rights to use, copy, modify, merge, publish, -# distribute, sublicense, and/or sell copies of the Software, and to -# permit persons to whom the Software is furnished to do so, subject to -# the following conditions: -# -# The above copyright notice and this permission notice shall be -# included in all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND -# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE -# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION -# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION -# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - -mutable struct _TopologicalSortVisitor - vertices::Vector{Int} - parents::Vector{Int} - - function _TopologicalSortVisitor(n::Int) - vs = Int[] - sizehint!(vs, n) - return new(vs, zeros(Int, n)) - end -end - -function _traverse_graph( - adjlist, - offsets, - s, - visitor::_TopologicalSortVisitor, - vertexcolormap, - vertex_stack, - index_stack, -) - vertexcolormap[s] = 1 - resize!(vertex_stack, 1) - vertex_stack[1] = s - resize!(index_stack, 1) - index_stack[1] = 1 - while !isempty(vertex_stack) - u = pop!(vertex_stack) - out_idx = pop!(index_stack) - len_uegs = offsets[u+1] - offsets[u] - found_new_vertex = false - while out_idx <= len_uegs && !found_new_vertex - v = adjlist[offsets[u]+out_idx-1] - out_idx += 1 - v_color = vertexcolormap[v] - if v_color == 0 - found_new_vertex = true - vertexcolormap[v] = 1 - push!(vertex_stack, u) - push!(index_stack, out_idx) - visitor.parents[v] = u - push!(vertex_stack, v) - push!(index_stack, 1) - end - end - if !found_new_vertex - push!(visitor.vertices, u) - vertexcolormap[u] = 2 - end - end - return -end - -function reverse_topological_sort_by_dfs( - adjlist, - offsets, - num_vertices, - cmap::Vector{Int}, - vertex_stack::Vector{Int} = Int[], - index_stack::Vector{Int} = Int[], -) - @assert length(cmap) == num_vertices - fill!(cmap, 0) - visitor = _TopologicalSortVisitor(num_vertices) - for s in 1:num_vertices - if cmap[s] == 0 - _traverse_graph( - adjlist, - offsets, - s, - visitor, - cmap, - vertex_stack, - index_stack, - ) - end - end - return visitor.vertices, visitor.parents -end diff --git a/src/coloring.jl b/src/coloring.jl new file mode 100644 index 0000000..32ca7e7 --- /dev/null +++ b/src/coloring.jl @@ -0,0 +1,140 @@ +# Copyright (c) 2017: Miles Lubin and contributors +# Copyright (c) 2017: Google Inc. +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. + +""" + struct ColoringResult + result::SMC.TreeSetColoringResult + local_indices::Vector{Int} # map from local to global indices + end + +Wrapper around TreeSetColoringResult that also stores local_indices mapping. +""" +struct ColoringResult{R<:SMC.AbstractColoringResult} + result::R + local_indices::Vector{Int} # map from local to global indices +end + +""" + _hessian_color_preprocess( + edgelist, + num_total_var, + algo::SMC.GreedyColoringAlgorithm, + seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0), + ) + +`edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the +diagonal. + +Returns `(I, J, result)` where `I` and `J` are the row and column indices +of the Hessian structure, and `result` is a `TreeSetColoringResult` from +SparseMatrixColorings. +""" +function _hessian_color_preprocess( + edgelist, + num_total_var, + algo::SMC.GreedyColoringAlgorithm, + seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0), +) + resize!(seen_idx, num_total_var) + I, J = Int[], Int[] + for (i, j) in edgelist + push!(seen_idx, i) + push!(seen_idx, j) + push!(I, i) + push!(J, j) + if i != j + push!(I, j) + push!(J, i) + end + end + local_indices = sort!(collect(seen_idx)) + empty!(seen_idx) + + global_to_local_idx = seen_idx.nzidx # steal for storage + for k in eachindex(local_indices) + global_to_local_idx[local_indices[k]] = k + end + # only do the coloring on the local indices + for k in eachindex(I) + 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 = SMC.SparsityPatternCSC( + SparseArrays.sparse(I, J, trues(length(I)), n, n, &), + ) + + # Perform coloring using SMC + problem = SMC.ColoringProblem(; structure = :symmetric, partition = :column) + display(S) + tree_result = SMC.coloring(S, problem, algo) + + # Wrap result with local_indices + result = ColoringResult(tree_result, local_indices) + + # 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 + +""" + _seed_matrix(result::ColoringResult) + +Allocate a seed matrix for the coloring result. +""" +function _seed_matrix(result::ColoringResult) + n = length(result.local_indices) + ncolors = SMC.ncolors(result.result) + return Matrix{Float64}(undef, n, ncolors) +end + +""" + _prepare_seed_matrix!(R, result::ColoringResult) + +Prepare the seed matrix R for Hessian computation. +""" +function _prepare_seed_matrix!(R, result::ColoringResult) + color = SMC.column_colors(result.result) + N = length(result.local_indices) + @assert N == size(R, 1) + @assert size(R, 2) == SMC.ncolors(result.result) + fill!(R, 0.0) + for i in 1:N + if color[i] > 0 + R[i, color[i]] = 1 + end + end + return +end + +""" + _recover_from_matmat!( + colptr::AbstractVector, + V::AbstractVector{T}, + R::AbstractMatrix{T}, + result::ColoringResult, + stored_values::AbstractVector{T}, + ) where {T} + +Recover the Hessian values from the Hessian-matrix product H*R_seed. +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} + SMC.decompress_csc!(V, colptr, R, result.result, :U) + return +end diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl index 6b170c2..fd7e8a7 100644 --- a/src/forward_over_reverse.jl +++ b/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 @@ -43,7 +43,7 @@ function _eval_hessian( return 0 end chunk = min(size(ex.seed_matrix, 2), d.max_chunk) - Coloring.prepare_seed_matrix!(ex.seed_matrix, ex.rinfo) + _prepare_seed_matrix!(ex.seed_matrix, ex.rinfo) # Compute hessian-vector products num_products = size(ex.seed_matrix, 2) # number of hessian-vector products num_chunks = div(num_products, chunk) @@ -65,8 +65,14 @@ function _eval_hessian( ) end # TODO(odow): consider reverting to a view. - output_slice = _UnsafeVectorView(nzcount, length(ex.hess_I), pointer(H)) - Coloring.recover_from_matmat!( + N = size(ex.seed_matrix, 1) + output_slice = _UnsafeVectorView{Float64}( + nzcount, + length(ex.hess_I), + pointer(H), + )::_UnsafeVectorView{Float64} + _recover_from_matmat!( + ex.hess_colptr, output_slice, ex.seed_matrix, ex.rinfo, diff --git a/src/graph_tools.jl b/src/graph_tools.jl index 4d20d47..22706c8 100644 --- a/src/graph_tools.jl +++ b/src/graph_tools.jl @@ -244,7 +244,10 @@ function _classify_linearity( end """ - _compute_gradient_sparsity!(indices::Coloring.IndexedSet, f) + _compute_gradient_sparsity!( + indices::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet, + f, + ) Compute the sparsity pattern of the gradient of an expression (that is, a list of which variable indices are present). @@ -255,7 +258,10 @@ For `NODE_VARIABLE_BLOCK` nodes, the block's length is read from `f.sizes` (via `_length`), and the block contributes that many consecutive variable indices starting at `nodes[k].index`. """ -function _compute_gradient_sparsity!(indices::Coloring.IndexedSet, f) +function _compute_gradient_sparsity!( + indices::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet, + f, +) for (k, node) in enumerate(f.nodes) if node.type == NODE_VARIABLE push!(indices, node.index) diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl index b0f54b9..2356010 100644 --- a/src/mathoptinterface_api.jl +++ b/src/mathoptinterface_api.jl @@ -20,9 +20,9 @@ function MOI.features_available(d::NLPEvaluator) end function MOI.initialize( - d::NLPEvaluator{T,S}, + d::NLPEvaluator{T,S,R}, requested_features::Vector{Symbol}, -) where {T<:Real,S<:AbstractVector{T}} +) where {T<:Real,S<:AbstractVector{T},R} # Check that we support the features requested by the user. available_features = MOI.features_available(d) for feature in requested_features @@ -42,7 +42,7 @@ function MOI.initialize( d.residual = nothing d.user_output_buffer = zeros(T, largest_user_input_dimension) d.jac_storage = zeros(T, max(N, largest_user_input_dimension)) - d.constraints = _FunctionStorage{T,S}[] + d.constraints = _FunctionStorage{T,S,R}[] # Allocate `last_x` on the same device as the AD tape (S = Vector for the # CPU path, S = CuVector for GPU). Filling with NaN forces the # short-circuit check `last_x == x` to take the recompute branch on the @@ -50,7 +50,7 @@ function MOI.initialize( d.last_x = fill!(S(undef, N), T(NaN)) d.want_hess = :Hess in requested_features want_hess_storage = (:HessVec in requested_features) || d.want_hess - coloring_storage = Coloring.IndexedSet(N) + coloring_storage = MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(N) max_expr_length = 0 max_expr_with_sub_length = 0 # @@ -126,11 +126,11 @@ function MOI.initialize( d, S, ) - objective = _FunctionStorage( + objective = _FunctionStorage{T,S,R}( subexpr, N, coloring_storage, - d.want_hess, + d.want_hess ? d.coloring_algorithm : nothing, d.subexpressions, individual_order[main_offset], subexpression_edgelist, @@ -151,11 +151,11 @@ function MOI.initialize( d, S, ) - residual = _FunctionStorage( + residual = _FunctionStorage{T,S,R}( subexpr, N, coloring_storage, - d.want_hess, + d.want_hess ? d.coloring_algorithm : nothing, d.subexpressions, individual_order[main_offset], subexpression_edgelist, @@ -178,11 +178,11 @@ function MOI.initialize( ) push!( d.constraints, - _FunctionStorage( + _FunctionStorage{T,S,R}( subexpr, N, coloring_storage, - d.want_hess, + d.want_hess ? d.coloring_algorithm : nothing, d.subexpressions, individual_order[idx], subexpression_edgelist, diff --git a/src/types.jl b/src/types.jl index 4e2ce6f..bf89336 100644 --- a/src/types.jl +++ b/src/types.jl @@ -124,28 +124,33 @@ function _subexpression_and_linearity( linearity end -struct _FunctionStorage{T<:Real,S<:AbstractVector{T}} +struct _FunctionStorage{ + T<:Real, + S<:AbstractVector{T}, + R<:SMC.AbstractColoringResult, +} expr::_SubexpressionStorage{T,S} grad_sparsity::Vector{Int} # Nonzero pattern of Hessian matrix + hess_colptr::Vector{Int} hess_I::Vector{Int} hess_J::Vector{Int} - rinfo::Coloring.RecoveryInfo # coloring info for hessians + rinfo::Union{Nothing,ColoringResult{R}} seed_matrix::Matrix{T} # subexpressions which this function depends on, ordered for forward pass. dependent_subexpressions::Vector{Int} - function _FunctionStorage( + function _FunctionStorage{T,S,R}( expr::_SubexpressionStorage{T,S}, num_variables, - coloring_storage::Coloring.IndexedSet, - want_hess::Bool, + coloring_storage::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet, + coloring_algorithm::Union{Nothing,SMC.GreedyColoringAlgorithm}, subexpressions::Vector{_SubexpressionStorage{T,S}}, dependent_subexpressions, subexpression_edgelist, subexpression_variables, linearity::Vector{Linearity}, - ) where {T<:Real,S<:AbstractVector{T}} + ) where {T<:Real,S<:AbstractVector{T},R<:SMC.AbstractColoringResult} empty!(coloring_storage) _compute_gradient_sparsity!(coloring_storage, expr) for k in dependent_subexpressions @@ -153,7 +158,7 @@ struct _FunctionStorage{T<:Real,S<:AbstractVector{T}} end grad_sparsity = sort!(collect(coloring_storage)) empty!(coloring_storage) - if want_hess + if !isnothing(coloring_algorithm) edgelist = _compute_hessian_sparsity( expr.nodes, expr.adj, @@ -161,15 +166,17 @@ struct _FunctionStorage{T<:Real,S<:AbstractVector{T}} subexpression_edgelist, subexpression_variables, ) - hess_I, hess_J, rinfo = Coloring.hessian_color_preprocess( + hess_colptr, hess_I, hess_J, rinfo = _hessian_color_preprocess( edgelist, num_variables, + coloring_algorithm, coloring_storage, ) - seed_matrix = Coloring.seed_matrix(rinfo) - return new{T,S}( + seed_matrix = _seed_matrix(rinfo) + return new{T,S,R}( expr, grad_sparsity, + hess_colptr, hess_I, hess_J, rinfo, @@ -177,12 +184,12 @@ struct _FunctionStorage{T<:Real,S<:AbstractVector{T}} dependent_subexpressions, ) else - return new{T,S}( + return new{T,S,R}( expr, grad_sparsity, Int[], Int[], - Coloring.RecoveryInfo(), + nothing, Array{T}(undef, 0, 0), dependent_subexpressions, ) @@ -306,14 +313,19 @@ interface. !!! warning Before using, you must initialize the evaluator using `MOI.initialize`. """ -mutable struct NLPEvaluator{T<:Real,S<:AbstractVector{T}} <: - MOI.AbstractNLPEvaluator +mutable struct NLPEvaluator{ + T<:Real, + S<:AbstractVector{T}, + R, + C<:SMC.GreedyColoringAlgorithm, +} <: MOI.AbstractNLPEvaluator data::Model ordered_variables::Vector{MOI.VariableIndex} + coloring_algorithm::C - objective::Union{Nothing,_FunctionStorage{T,S}} - residual::Union{Nothing,_FunctionStorage{T,S}} - constraints::Vector{_FunctionStorage{T,S}} + objective::Union{Nothing,_FunctionStorage{T,S,R}} + residual::Union{Nothing,_FunctionStorage{T,S,R}} + constraints::Vector{_FunctionStorage{T,S,R}} subexpressions::Vector{_SubexpressionStorage{T,S}} subexpression_order::Vector{Int} # Storage for the subexpressions in reverse-mode automatic differentiation. @@ -350,7 +362,19 @@ mutable struct NLPEvaluator{T<:Real,S<:AbstractVector{T}} <: function NLPEvaluator{T,S}( data::Model, ordered_variables::Vector{MOI.VariableIndex}, + coloring_algorithm::SMC.GreedyColoringAlgorithm = SMC.GreedyColoringAlgorithm(; + decompression = :substitution, + ), ) where {T<:Real,S<:AbstractVector{T}} - return new{T,S}(data, ordered_variables) + problem = + SMC.ColoringProblem(; structure = :symmetric, partition = :column) + C = typeof(coloring_algorithm) + R = Base.promote_op( + SMC.coloring, + SMC.SparsityPatternCSC{Int}, + typeof(problem), + C, + ) + return new{T,S,R,C}(data, ordered_variables, coloring_algorithm) end end diff --git a/src/utils.jl b/src/utils.jl index d1d69e4..dbd3c9b 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -203,3 +203,18 @@ 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 diff --git a/test/Project.toml b/test/Project.toml index 9d5d3ee..69b7d6b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -21,6 +21,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Revise = "295af30f-e4ad-537b-8983-00126c2a3abe" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/test_ArrayDiff.jl b/test/test_ArrayDiff.jl index b4bda2e..5b61a5f 100644 --- a/test/test_ArrayDiff.jl +++ b/test/test_ArrayDiff.jl @@ -8,7 +8,6 @@ import MathOptInterface as MOI const Nonlinear = MOI.Nonlinear import ArrayDiff -const Coloring = ArrayDiff.Coloring function runtests() for name in names(@__MODULE__; all = true) @@ -856,6 +855,33 @@ function test_model_typed_float32_evaluator_runs() return end +function test_mode_storage_and_coloring() + algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm( + ArrayDiff.SMC.LargestFirst(); + decompression = :substitution, + ) + for mode in ( + ArrayDiff.Mode{Vector{Float32}}(), + ArrayDiff.Mode{Vector{Float32}}(algorithm), + ArrayDiff.Mode(algorithm), + ) + model = ArrayDiff.model(mode) + T = eltype(model.parameters) + x = MOI.VariableIndex(1) + ArrayDiff.set_objective(model, :($x^2)) + evaluator = ArrayDiff.Evaluator(model, mode, [x]) + MOI.initialize(evaluator, [:Grad]) + @test evaluator.backend.coloring_algorithm === mode.coloring_algorithm + @test evaluator.backend.last_x isa Vector{T} + @test evaluator.backend.objective.expr.forward_storage isa Vector{T} + @test MOI.eval_objective(evaluator, T[3]) == T(9) + g = zeros(T, 1) + MOI.eval_objective_gradient(evaluator, g, T[3]) + @test g == T[6] + end + return +end + function test_residual_with_subexpression() # Residual references a subexpression `e = x1 * x2`, so the evaluator's # subexpression-iteration loops in `_forward_pass_residual!` and diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index ae919e9..05d5c2f 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -14,7 +14,6 @@ import MathOptInterface as MOI const Nonlinear = MOI.Nonlinear import ArrayDiff -const Coloring = ArrayDiff.Coloring function runtests() for name in names(@__MODULE__; all = true) @@ -414,79 +413,90 @@ struct _ColoringGraph edges::Vector{Tuple{Int,Int}} end -function to_adjlist(graph::_ColoringGraph) - I = [i for (i, _) in graph.edges] - J = [j for (_, j) in graph.edges] - return Coloring.UndirectedGraph(I, J, graph.num_vertices) -end - -function test_coloring_edge_free_graph() - graph = _ColoringGraph(10, []) - _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) - @test numcolors == 1 - return -end - -function test_coloring_one_edge_graph() - graph = _ColoringGraph(10, [(2, 4)]) - color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) - @test numcolors == 2 - @test color[2] != color[4] - return -end - -function test_coloring_two_edge_graph() - graph = _ColoringGraph(10, [(2, 4), (2, 3)]) - color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) - @test numcolors == 2 - @test color[3] == color[4] - return -end - -function test_coloring_three_edge_graph() - graph = _ColoringGraph(10, [(2, 4), (2, 3), (3, 4)]) - color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) - @test numcolors == 3 - # TODO: What is this testing? - Coloring.recovery_preprocess(to_adjlist(graph), color, numcolors, Int[]) - return -end - -function test_coloring_two_edge_three_vertex_graph() - graph = _ColoringGraph(3, [(1, 3), (2, 3)]) - _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) - @test numcolors == 2 - return -end - -function test_coloring_four_edge_four_vertex_graph() - graph = _ColoringGraph(4, [(1, 2), (2, 3), (3, 4), (4, 1)]) - _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) - @test numcolors == 3 - return -end - -function test_coloring_topological_sort() - # graph = _ColoringGraph(6, [(1, 2), (1, 3), (1, 6), (2, 4), (2, 5)]) - vec = [3, 6, 2, 1, 4, 5, 1, 2, 2, 1] - offset = [1, 4, 7, 8, 9, 10, 11] - v = Coloring.reverse_topological_sort_by_dfs(vec, offset, 6, zeros(Int, 6)) - @test v[1] == [3, 6, 4, 5, 2, 1] - @test v[2] == [0, 1, 1, 2, 2, 1] - return -end +# Tests for internal Coloring submodule removed - now using SparseMatrixColorings +# The coloring functionality is tested indirectly through the Hessian evaluation tests + +# function to_adjlist(graph::_ColoringGraph) +# I = [i for (i, _) in graph.edges] +# J = [j for (_, j) in graph.edges] +# return Coloring.UndirectedGraph(I, J, graph.num_vertices) +# end + +# function test_coloring_edge_free_graph() +# graph = _ColoringGraph(10, []) +# _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) +# @test numcolors == 1 +# return +# end + +# function test_coloring_one_edge_graph() +# graph = _ColoringGraph(10, [(2, 4)]) +# color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) +# @test numcolors == 2 +# @test color[2] != color[4] +# return +# end + +# function test_coloring_two_edge_graph() +# graph = _ColoringGraph(10, [(2, 4), (2, 3)]) +# color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) +# @test numcolors == 2 +# @test color[3] == color[4] +# return +# end + +# function test_coloring_three_edge_graph() +# graph = _ColoringGraph(10, [(2, 4), (2, 3), (3, 4)]) +# color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) +# @test numcolors == 3 +# # TODO: What is this testing? +# Coloring.recovery_preprocess(to_adjlist(graph), color, numcolors, Int[]) +# return +# end + +# function test_coloring_two_edge_three_vertex_graph() +# graph = _ColoringGraph(3, [(1, 3), (2, 3)]) +# _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) +# @test numcolors == 2 +# return +# end + +# function test_coloring_four_edge_four_vertex_graph() +# graph = _ColoringGraph(4, [(1, 2), (2, 3), (3, 4), (4, 1)]) +# _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) +# @test numcolors == 3 +# return +# end + +# function test_coloring_topological_sort() +# # graph = _ColoringGraph(6, [(1, 2), (1, 3), (1, 6), (2, 4), (2, 5)]) +# vec = [3, 6, 2, 1, 4, 5, 1, 2, 2, 1] +# offset = [1, 4, 7, 8, 9, 10, 11] +# v = Coloring.reverse_topological_sort_by_dfs(vec, offset, 6, zeros(Int, 6)) +# @test v[1] == [3, 6, 4, 5, 2, 1] +# @test v[2] == [0, 1, 1, 2, 2, 1] +# return +# end function test_coloring_end_to_end_hessian_coloring_and_recovery() - I, J, rinfo = Coloring.hessian_color_preprocess(Set([(1, 2)]), 2) - R = Coloring.seed_matrix(rinfo) - Coloring.prepare_seed_matrix!(R, rinfo) + # Test the new coloring API through the compatibility layer + coloring_algorithm = + ArrayDiff.SMC.GreedyColoringAlgorithm(; decompression = :substitution) + colptr, I, J, rinfo = ArrayDiff._hessian_color_preprocess( + Set([(1, 2)]), + 2, + coloring_algorithm, + MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(0), + ) + R = ArrayDiff._seed_matrix(rinfo) + ArrayDiff._prepare_seed_matrix!(R, rinfo) @test I == [1, 2, 2] @test J == [1, 2, 1] @test R == [1.0 0.0; 0.0 1.0] hess = [3.4 2.1; 2.1 1.3] matmat = hess * R V = zeros(3) - Coloring.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 @@ -547,7 +557,7 @@ function test_linearity() nodes = ArrayDiff._replace_moi_variables(expr.nodes, variables) ret = ArrayDiff._classify_linearity(nodes, adj, ArrayDiff.Linearity[]) @test ret[1] == test_value - indexed_set = Coloring.IndexedSet(100) + indexed_set = MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(100) edge_list = ArrayDiff._compute_hessian_sparsity( nodes, adj, @@ -1466,11 +1476,11 @@ function test_hessian_pow_constant_exponent_nonpositive_base() evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), [x, y]) MOI.initialize(evaluator, [:Grad, :Hess]) @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (2, 1)] + [(1, 1), (2, 1), (2, 2)] H = zeros(3) for x_val in (-1.5, 0.0) MOI.eval_hessian_lagrangian(evaluator, H, [x_val, 2.0], 1.0, Float64[]) - @test isapprox(H, [4.0, 2 * x_val, 2 * x_val + 4.0]) + @test isapprox(H, [4.0, 2 * x_val + 4.0, 2 * x_val]) end return end From c02c15ef6da8fffdb020affc61f1d0541c631628 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 25 Feb 2026 16:00:07 +0100 Subject: [PATCH 02/14] Remove coloring tests --- test/test_ReverseAD.jl | 65 ------------------------------------------ 1 file changed, 65 deletions(-) diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index 05d5c2f..03a2102 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -413,71 +413,6 @@ struct _ColoringGraph edges::Vector{Tuple{Int,Int}} end -# Tests for internal Coloring submodule removed - now using SparseMatrixColorings -# The coloring functionality is tested indirectly through the Hessian evaluation tests - -# function to_adjlist(graph::_ColoringGraph) -# I = [i for (i, _) in graph.edges] -# J = [j for (_, j) in graph.edges] -# return Coloring.UndirectedGraph(I, J, graph.num_vertices) -# end - -# function test_coloring_edge_free_graph() -# graph = _ColoringGraph(10, []) -# _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) -# @test numcolors == 1 -# return -# end - -# function test_coloring_one_edge_graph() -# graph = _ColoringGraph(10, [(2, 4)]) -# color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) -# @test numcolors == 2 -# @test color[2] != color[4] -# return -# end - -# function test_coloring_two_edge_graph() -# graph = _ColoringGraph(10, [(2, 4), (2, 3)]) -# color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) -# @test numcolors == 2 -# @test color[3] == color[4] -# return -# end - -# function test_coloring_three_edge_graph() -# graph = _ColoringGraph(10, [(2, 4), (2, 3), (3, 4)]) -# color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) -# @test numcolors == 3 -# # TODO: What is this testing? -# Coloring.recovery_preprocess(to_adjlist(graph), color, numcolors, Int[]) -# return -# end - -# function test_coloring_two_edge_three_vertex_graph() -# graph = _ColoringGraph(3, [(1, 3), (2, 3)]) -# _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) -# @test numcolors == 2 -# return -# end - -# function test_coloring_four_edge_four_vertex_graph() -# graph = _ColoringGraph(4, [(1, 2), (2, 3), (3, 4), (4, 1)]) -# _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph)) -# @test numcolors == 3 -# return -# end - -# function test_coloring_topological_sort() -# # graph = _ColoringGraph(6, [(1, 2), (1, 3), (1, 6), (2, 4), (2, 5)]) -# vec = [3, 6, 2, 1, 4, 5, 1, 2, 2, 1] -# offset = [1, 4, 7, 8, 9, 10, 11] -# v = Coloring.reverse_topological_sort_by_dfs(vec, offset, 6, zeros(Int, 6)) -# @test v[1] == [3, 6, 4, 5, 2, 1] -# @test v[2] == [0, 1, 1, 2, 2, 1] -# return -# end - function test_coloring_end_to_end_hessian_coloring_and_recovery() # Test the new coloring API through the compatibility layer coloring_algorithm = From acf999310c517d9a3b711d6f54119aa0efdb89cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 20 Mar 2026 11:50:50 +0100 Subject: [PATCH 03/14] SMC v0.4.26 --- .github/workflows/ci.yml | 7 ------- Project.toml | 2 +- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 247a219..79a4c27 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,13 +25,6 @@ jobs: version: ${{ matrix.version }} arch: ${{ matrix.arch }} - uses: julia-actions/cache@v3 - - name: MOI - shell: julia --project=@. {0} - run: | - using Pkg - Pkg.add([ - PackageSpec(url="https://github.com/blegat/SparseMatrixColorings.jl/", rev="abstractsparse"), - ]) - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 with: diff --git a/Project.toml b/Project.toml index 7ceed19..b8627d3 100644 --- a/Project.toml +++ b/Project.toml @@ -36,6 +36,6 @@ NaNMath = "1" ONNX = "0.2, 0.3" OrderedCollections = "1.8.1, 2" SparseArrays = "1.10" -SparseMatrixColorings = "0.4" +SparseMatrixColorings = "0.4.26" SpecialFunctions = "2.6.1" julia = "1.10" From 622ea126dfdbb4f4e894013cd18bb7f7c1b43a83 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 07:01:40 +0000 Subject: [PATCH 04/14] Fix CI errors: remove debug display and missing hess_J field - Remove leftover `display(S)` debug statement in coloring.jl - Add missing `Int[]` for `hess_J` field in `_FunctionStorage` else branch https://claude.ai/code/session_01WBu9hZukriWDSSybN9gfBq --- src/coloring.jl | 1 - src/types.jl | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coloring.jl b/src/coloring.jl index 32ca7e7..124cb0d 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -71,7 +71,6 @@ function _hessian_color_preprocess( # Perform coloring using SMC problem = SMC.ColoringProblem(; structure = :symmetric, partition = :column) - display(S) tree_result = SMC.coloring(S, problem, algo) # Wrap result with local_indices diff --git a/src/types.jl b/src/types.jl index bf89336..4fa2353 100644 --- a/src/types.jl +++ b/src/types.jl @@ -189,6 +189,7 @@ struct _FunctionStorage{ grad_sparsity, Int[], Int[], + Int[], nothing, Array{T}(undef, 0, 0), dependent_subexpressions, From 613aebe86787fde95638b31bfbba1c9a85d10938 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 07:36:52 +0000 Subject: [PATCH 05/14] Fix Hessian coloring: use lower triangular CSC format - Build lower triangular sparsity pattern (row >= col) instead of full symmetric matrix, so decompress_csc! with :L works correctly - Always add diagonal entries explicitly (old Coloring module did this implicitly; SMC requires them to be in the sparsity pattern) - Get CSC structure directly from the sparse matrix instead of going through compress/decompress - Switch decompress_csc! from :U to :L to match lower triangular storage - Update test expectations to CSC lower triangular order: J=[1,1,2] (was [1,2,1]), V=[3.4,2.1,1.3] (was [3.4,1.3,2.1]) https://claude.ai/code/session_01WBu9hZukriWDSSybN9gfBq --- src/coloring.jl | 38 +++++++++++++++++++------------------- test/test_ReverseAD.jl | 4 ++-- 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index 124cb0d..3bbc190 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -40,15 +40,12 @@ function _hessian_color_preprocess( ) resize!(seen_idx, num_total_var) I, J = Int[], Int[] - for (i, j) in edgelist - push!(seen_idx, i) - push!(seen_idx, j) - push!(I, i) - push!(J, j) - if i != j - push!(I, j) - push!(J, i) - end + for (ei, ej) in edgelist + push!(seen_idx, ei) + push!(seen_idx, ej) + # Store in lower triangular format: row >= col + push!(I, max(ei, ej)) + push!(J, min(ei, ej)) end local_indices = sort!(collect(seen_idx)) empty!(seen_idx) @@ -63,11 +60,16 @@ function _hessian_color_preprocess( J[k] = global_to_local_idx[J[k]] end - # Create sparsity pattern matrix n = length(local_indices) - S = SMC.SparsityPatternCSC( - SparseArrays.sparse(I, J, trues(length(I)), n, n, &), - ) + # Always include diagonal entries (needed for Hessian recovery) + for k in 1:n + push!(I, k) + push!(J, k) + end + + # Create lower triangular sparsity pattern (including diagonal) + mat = SparseArrays.sparse(I, J, trues(length(I)), n, n, |) + S = SMC.SparsityPatternCSC(mat) # Perform coloring using SMC problem = SMC.ColoringProblem(; structure = :symmetric, partition = :column) @@ -76,12 +78,10 @@ function _hessian_color_preprocess( # Wrap result with local_indices result = ColoringResult(tree_result, local_indices) - # 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) + # Get CSC-ordered indices directly from the sparse matrix + I_sorted, J_sorted, _ = SparseArrays.findnz(mat) - return C.colptr, I_sorted, J_sorted, result + return copy(mat.colptr), I_sorted, J_sorted, result end """ @@ -134,6 +134,6 @@ function _recover_from_matmat!( result::ColoringResult, stored_values::AbstractVector{T}, ) where {T} - SMC.decompress_csc!(V, colptr, R, result.result, :U) + SMC.decompress_csc!(V, colptr, R, result.result, :L) return end diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index 03a2102..976fd93 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -426,13 +426,13 @@ function test_coloring_end_to_end_hessian_coloring_and_recovery() R = ArrayDiff._seed_matrix(rinfo) ArrayDiff._prepare_seed_matrix!(R, rinfo) @test I == [1, 2, 2] - @test J == [1, 2, 1] + @test J == [1, 1, 2] @test R == [1.0 0.0; 0.0 1.0] hess = [3.4 2.1; 2.1 1.3] matmat = hess * R V = zeros(3) ArrayDiff._recover_from_matmat!(colptr, V, matmat, rinfo, zeros(3)) - @test V == [3.4, 1.3, 2.1] + @test V == [3.4, 2.1, 1.3] return end From b8411a6bb6424b5f2eb6d668b1a8d5acf9899555 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 08:39:29 +0000 Subject: [PATCH 06/14] Fix test expectations to match CSC lower triangular ordering The SMC-based Hessian recovery stores nonzeros in lower triangular CSC format (column-major: all entries of column j before column j+1, with diagonal first in each column). Update all multivariate Hessian structure and value expectations to match this ordering instead of the old diagonal-first ordering. https://claude.ai/code/session_01WBu9hZukriWDSSybN9gfBq --- test/test_ReverseAD.jl | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index 976fd93..c9e46ed 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -103,15 +103,15 @@ function test_objective_quadratic_multivariate() g = [NaN, NaN] MOI.eval_objective_gradient(evaluator, g, [1.2, 2.3]) @test g == [2 * 1.2 + 2.3, 1.2 + 2 * 2.3] - @test MOI.hessian_objective_structure(evaluator) == [(1, 1), (2, 2), (2, 1)] + @test MOI.hessian_objective_structure(evaluator) == [(1, 1), (2, 1), (2, 2)] H = [NaN, NaN, NaN] MOI.eval_hessian_objective(evaluator, H, [1.2, 2.3]) - @test H == [2.0, 2.0, 1.0] + @test H == [2.0, 1.0, 2.0] @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (2, 1)] + [(1, 1), (2, 1), (2, 2)] H = [NaN, NaN, NaN] MOI.eval_hessian_lagrangian(evaluator, H, [1.2, 2.3], 1.5, Float64[]) - @test H == 1.5 .* [2.0, 2.0, 1.0] + @test H == 1.5 .* [2.0, 1.0, 2.0] v = [0.3, 0.4] hv = [NaN, NaN] MOI.eval_hessian_lagrangian_product( @@ -143,19 +143,19 @@ function test_objective_quadratic_multivariate_subexpressions() MOI.eval_objective_gradient(evaluator, g, val) @test g == [2 * 1.2 + 2.3, 1.2 + 2 * 2.3] @test 0 == @allocated MOI.eval_objective_gradient(evaluator, g, val) - @test MOI.hessian_objective_structure(evaluator) == [(1, 1), (2, 2), (2, 1)] + @test MOI.hessian_objective_structure(evaluator) == [(1, 1), (2, 1), (2, 2)] H = [NaN, NaN, NaN] MOI.eval_hessian_objective(evaluator, H, val) - @test H == [2.0, 2.0, 1.0] + @test H == [2.0, 1.0, 2.0] @test evaluator.backend.max_chunk == 2 @test 0 == @allocated MOI.eval_hessian_objective(evaluator, H, val) @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (2, 1)] + [(1, 1), (2, 1), (2, 2)] H = [NaN, NaN, NaN] μ = Float64[] MOI.eval_hessian_lagrangian(evaluator, H, val, 1.5, μ) @test 0 == @allocated MOI.eval_hessian_lagrangian(evaluator, H, val, 1.5, μ) - @test H == 1.5 .* [2.0, 2.0, 1.0] + @test H == 1.5 .* [2.0, 1.0, 2.0] v = [0.3, 0.4] hv = [NaN, NaN] MOI.eval_hessian_lagrangian_product(evaluator, hv, val, v, 1.5, μ) @@ -276,10 +276,10 @@ function test_constraint_quadratic_multivariate() MOI.eval_constraint_jacobian(evaluator, J, x_val) @test J == [2 * x_val[1] + x_val[2], x_val[1] + 2 * x_val[2]] @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (2, 1)] + [(1, 1), (2, 1), (2, 2)] H = [NaN, NaN, NaN] MOI.eval_hessian_lagrangian(evaluator, H, x_val, 0.0, [1.5]) - @test H == 1.5 .* [2.0, 2.0, 1.0] + @test H == 1.5 .* [2.0, 1.0, 2.0] return end @@ -313,10 +313,10 @@ function test_constraint_quadratic_multivariate_subexpressions() @test y ≈ wJ[:] # Hessian-Lagrangian @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (2, 1)] + [(1, 1), (2, 1), (2, 2)] H = [NaN, NaN, NaN] MOI.eval_hessian_lagrangian(evaluator, H, x_val, 0.0, [1.5]) - @test H ≈ 1.5 .* [2.0, 2.0, 1.0] + @test H ≈ 1.5 .* [2.0, 1.0, 2.0] return end @@ -342,10 +342,10 @@ function test_hessian_sparsity_registered_function() @test :Hess in MOI.features_available(evaluator) MOI.initialize(evaluator, [:Grad, :Jac, :Hess]) @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (3, 3), (3, 1)] + [(1, 1), (3, 1), (2, 2), (3, 3)] H = fill(NaN, 4) MOI.eval_hessian_lagrangian(evaluator, H, rand(3), 1.5, Float64[]) - @test H == 1.5 .* [2.0, 2.0, 2.0, 0.0] + @test H == 1.5 .* [2.0, 0.0, 2.0, 2.0] return end @@ -372,10 +372,10 @@ function test_hessian_sparsity_registered_rosenbrock() @test :Hess in MOI.features_available(evaluator) MOI.initialize(evaluator, [:Grad, :Jac, :Hess]) @test MOI.hessian_lagrangian_structure(evaluator) == - [(1, 1), (2, 2), (2, 1)] + [(1, 1), (2, 1), (2, 2)] H = fill(NaN, 3) MOI.eval_hessian_lagrangian(evaluator, H, [1.0, 1.0], 1.5, Float64[]) - @test H == 1.5 .* [802.0, 200.0, -400.0] + @test H == 1.5 .* [802.0, -400.0, 200.0] return end From 7ee8d7609e88a8c4d0c35c2bf1d4805b03e2dfff Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 09:07:27 +0000 Subject: [PATCH 07/14] Fix NLPEvaluator type inference by using concrete coloring call instead of Base.promote_op Base.promote_op may return Any for SMC.coloring since it lacks an explicit return type annotation, causing TypeError when constructing _FunctionStorage{Any} because Any <: AbstractColoringResult is false. Use a small test coloring call instead. https://claude.ai/code/session_01WBu9hZukriWDSSybN9gfBq --- src/types.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/types.jl b/src/types.jl index 4fa2353..7cba31d 100644 --- a/src/types.jl +++ b/src/types.jl @@ -370,12 +370,10 @@ mutable struct NLPEvaluator{ problem = SMC.ColoringProblem(; structure = :symmetric, partition = :column) C = typeof(coloring_algorithm) - R = Base.promote_op( - SMC.coloring, - SMC.SparsityPatternCSC{Int}, - typeof(problem), - C, + _S = SMC.SparsityPatternCSC( + SparseArrays.sparse([1], [1], Bool[true], 1, 1), ) + R = typeof(SMC.coloring(_S, problem, coloring_algorithm)) return new{T,S,R,C}(data, ordered_variables, coloring_algorithm) end end From f654a4a757cfc1931a7b8ecf192006ffc97e6f37 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 09:56:15 +0000 Subject: [PATCH 08/14] Fix: map local Hessian indices back to global indices in _hessian_color_preprocess The function was returning I/J in local variable indices but callers use them as global variable indices (for hessian_lagrangian_structure). This caused wrong sparsity patterns when a function depends on non-contiguous variables (e.g., variables 4 and 5 out of 5 total). https://claude.ai/code/session_01WBu9hZukriWDSSybN9gfBq --- src/coloring.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index 3bbc190..bc4058c 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -78,10 +78,12 @@ function _hessian_color_preprocess( # Wrap result with local_indices result = ColoringResult(tree_result, local_indices) - # Get CSC-ordered indices directly from the sparse matrix - I_sorted, J_sorted, _ = SparseArrays.findnz(mat) + # Get CSC-ordered indices from the sparse matrix, then map back to global + I_local, J_local, _ = SparseArrays.findnz(mat) + I_global = [local_indices[i] for i in I_local] + J_global = [local_indices[j] for j in J_local] - return copy(mat.colptr), I_sorted, J_sorted, result + return copy(mat.colptr), I_global, J_global, result end """ From 4c27a73e7001c44e672b9f88d70da0af3ea0b06b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 11:11:41 +0000 Subject: [PATCH 09/14] Fix BoundsError in decompress_csc! by using full symmetric matrix for coloring Pass a full symmetric sparsity pattern to SMC.coloring instead of lower triangular, so the tree edges are correctly oriented. Store full_colptr, lower_pos, and full_buffer in ColoringResult to allow recovery via decompress_csc! with :F, then extract lower-triangular values. This fixes: BoundsError: attempt to access 0-element Vector{Int64} at index [1] in test_coloring_end_to_end_hessian_coloring_and_recovery. https://claude.ai/code/session_01WBu9hZukriWDSSybN9gfBq --- src/coloring.jl | 99 +++++++++++++++++++++++++++++++++---------------- 1 file changed, 68 insertions(+), 31 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index bc4058c..e8b0ce4 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -4,17 +4,25 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. + """ struct ColoringResult - result::SMC.TreeSetColoringResult + result::SMC.AbstractColoringResult local_indices::Vector{Int} # map from local to global indices + full_colptr::Vector{Int} # colptr of the full symmetric matrix used for coloring + lower_pos::Vector{Int} # positions of lower-triangular entries in the full nzval + full_buffer::Vector{Float64} # pre-allocated buffer of size nnz(full symmetric matrix) end -Wrapper around TreeSetColoringResult that also stores local_indices mapping. +Wrapper around AbstractColoringResult that also stores auxiliary data needed +for Hessian recovery from a full symmetric matrix decompression. """ struct ColoringResult{R<:SMC.AbstractColoringResult} result::R - local_indices::Vector{Int} # map from local to global indices + local_indices::Vector{Int} # map from local to global indices + full_colptr::Vector{Int} # colptr of full symmetric matrix used for coloring + lower_pos::Vector{Int} # positions of lower-triangular entries in full nzval + full_buffer::Vector{Float64} # scratch buffer of length nnz(full symmetric matrix) end """ @@ -28,9 +36,9 @@ end `edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the diagonal. -Returns `(I, J, result)` where `I` and `J` are the row and column indices -of the Hessian structure, and `result` is a `TreeSetColoringResult` from -SparseMatrixColorings. +Returns `(colptr, I, J, result)` where `colptr`, `I` and `J` define the lower +triangular CSC sparsity structure of the Hessian (in global variable indices), +and `result` is a `ColoringResult` wrapping an SMC coloring result. """ function _hessian_color_preprocess( edgelist, @@ -39,13 +47,16 @@ function _hessian_color_preprocess( seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0), ) resize!(seen_idx, num_total_var) - I, J = Int[], Int[] + # Collect off-diagonal lower-triangular entries (local coords, filled later) + I_off, J_off = Int[], Int[] for (ei, ej) in edgelist push!(seen_idx, ei) push!(seen_idx, ej) - # Store in lower triangular format: row >= col - push!(I, max(ei, ej)) - push!(J, min(ei, ej)) + if ei != ej + # Store in lower triangular format: row > col + push!(I_off, max(ei, ej)) + push!(J_off, min(ei, ej)) + end end local_indices = sort!(collect(seen_idx)) empty!(seen_idx) @@ -54,36 +65,58 @@ function _hessian_color_preprocess( for k in eachindex(local_indices) global_to_local_idx[local_indices[k]] = k end - # only do the coloring on the local indices - for k in eachindex(I) - I[k] = global_to_local_idx[I[k]] - J[k] = global_to_local_idx[J[k]] + # Map off-diagonal entries to local indices + for k in eachindex(I_off) + I_off[k] = global_to_local_idx[I_off[k]] + J_off[k] = global_to_local_idx[J_off[k]] end n = length(local_indices) - # Always include diagonal entries (needed for Hessian recovery) + + # Build full symmetric matrix: both (i,j) and (j,i) for off-diagonal, plus diagonal + I_full, J_full = Int[], Int[] + for k in eachindex(I_off) + push!(I_full, I_off[k]); push!(J_full, J_off[k]) # lower + push!(I_full, J_off[k]); push!(J_full, I_off[k]) # upper (transpose) + end for k in 1:n - push!(I, k) - push!(J, k) + push!(I_full, k); push!(J_full, k) # diagonal end + mat_sym = SparseArrays.sparse(I_full, J_full, trues(length(I_full)), n, n, |) - # Create lower triangular sparsity pattern (including diagonal) - mat = SparseArrays.sparse(I, J, trues(length(I)), n, n, |) - S = SMC.SparsityPatternCSC(mat) - - # Perform coloring using SMC + # Perform coloring on full symmetric matrix + S = SMC.SparsityPatternCSC(mat_sym) problem = SMC.ColoringProblem(; structure = :symmetric, partition = :column) tree_result = SMC.coloring(S, problem, algo) - # Wrap result with local_indices - result = ColoringResult(tree_result, local_indices) + # Find positions of lower-triangular entries within the full CSC nzval array. + # findnz on a CSC matrix returns elements in CSC (column-major) order, + # matching the nzval layout exactly. + I_nz, J_nz, _ = SparseArrays.findnz(mat_sym) + lower_pos = findall(k -> I_nz[k] >= J_nz[k], eachindex(I_nz)) + + # Lower-triangular CSC-ordered local indices + I_low_csc = I_nz[lower_pos] + J_low_csc = J_nz[lower_pos] - # Get CSC-ordered indices from the sparse matrix, then map back to global - I_local, J_local, _ = SparseArrays.findnz(mat) - I_global = [local_indices[i] for i in I_local] - J_global = [local_indices[j] for j in J_local] + # Map back to global indices for the returned hess_I / hess_J + I_global = [local_indices[i] for i in I_low_csc] + J_global = [local_indices[j] for j in J_low_csc] - return copy(mat.colptr), I_global, J_global, result + # Build lower-triangular sparse matrix to obtain its colptr + mat_low = SparseArrays.sparse(I_low_csc, J_low_csc, trues(length(I_low_csc)), n, n, |) + + full_buffer = Vector{Float64}(undef, SparseArrays.nnz(mat_sym)) + + result = ColoringResult( + tree_result, + local_indices, + copy(mat_sym.colptr), + lower_pos, + full_buffer, + ) + + return copy(mat_low.colptr), I_global, J_global, result end """ @@ -127,7 +160,7 @@ end Recover the Hessian values from the Hessian-matrix product H*R_seed. R is the result of H*R_seed where R_seed is the seed matrix. -`stored_values` is a temporary vector. +`stored_values` is a temporary vector (unused, kept for API compatibility). """ function _recover_from_matmat!( colptr::AbstractVector, @@ -136,6 +169,10 @@ function _recover_from_matmat!( result::ColoringResult, stored_values::AbstractVector{T}, ) where {T} - SMC.decompress_csc!(V, colptr, R, result.result, :L) + # Decompress into the full symmetric buffer, then extract lower-triangular values. + SMC.decompress_csc!(result.full_buffer, result.full_colptr, R, result.result, :F) + for k in eachindex(V) + V[k] = result.full_buffer[result.lower_pos[k]] + end return end From 813e65847fdfde3a89a33a044367a2a470d5456c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 2 Apr 2026 15:51:05 +0200 Subject: [PATCH 10/14] Fix format --- src/coloring.jl | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index e8b0ce4..0684004 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -4,7 +4,6 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. - """ struct ColoringResult result::SMC.AbstractColoringResult @@ -76,13 +75,17 @@ function _hessian_color_preprocess( # Build full symmetric matrix: both (i,j) and (j,i) for off-diagonal, plus diagonal I_full, J_full = Int[], Int[] for k in eachindex(I_off) - push!(I_full, I_off[k]); push!(J_full, J_off[k]) # lower - push!(I_full, J_off[k]); push!(J_full, I_off[k]) # upper (transpose) + push!(I_full, I_off[k]); + push!(J_full, J_off[k]) # lower + push!(I_full, J_off[k]); + push!(J_full, I_off[k]) # upper (transpose) end for k in 1:n - push!(I_full, k); push!(J_full, k) # diagonal + push!(I_full, k); + push!(J_full, k) # diagonal end - mat_sym = SparseArrays.sparse(I_full, J_full, trues(length(I_full)), n, n, |) + mat_sym = + SparseArrays.sparse(I_full, J_full, trues(length(I_full)), n, n, |) # Perform coloring on full symmetric matrix S = SMC.SparsityPatternCSC(mat_sym) @@ -104,7 +107,14 @@ function _hessian_color_preprocess( J_global = [local_indices[j] for j in J_low_csc] # Build lower-triangular sparse matrix to obtain its colptr - mat_low = SparseArrays.sparse(I_low_csc, J_low_csc, trues(length(I_low_csc)), n, n, |) + mat_low = SparseArrays.sparse( + I_low_csc, + J_low_csc, + trues(length(I_low_csc)), + n, + n, + |, + ) full_buffer = Vector{Float64}(undef, SparseArrays.nnz(mat_sym)) @@ -170,7 +180,13 @@ function _recover_from_matmat!( stored_values::AbstractVector{T}, ) where {T} # Decompress into the full symmetric buffer, then extract lower-triangular values. - SMC.decompress_csc!(result.full_buffer, result.full_colptr, R, result.result, :F) + SMC.decompress_csc!( + result.full_buffer, + result.full_colptr, + R, + result.result, + :F, + ) for k in eachindex(V) V[k] = result.full_buffer[result.lower_pos[k]] end From 388ff013bd6146c801edadb54399f7a0d32906e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 16 Oct 2026 08:31:38 +0200 Subject: [PATCH 11/14] Fix docstring --- src/coloring.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index 0684004..96d516f 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -5,8 +5,8 @@ # in the LICENSE.md file or at https://opensource.org/licenses/MIT. """ - struct ColoringResult - result::SMC.AbstractColoringResult + struct ColoringResult{R<:SMC.AbstractColoringResult} + result::R local_indices::Vector{Int} # map from local to global indices full_colptr::Vector{Int} # colptr of the full symmetric matrix used for coloring lower_pos::Vector{Int} # positions of lower-triangular entries in the full nzval From 0e76337193edd034adf38ad8e4af9df6a340416a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 16 Oct 2026 09:04:26 +0200 Subject: [PATCH 12/14] Remove diagonal entries --- src/coloring.jl | 18 ++++++----- test/test_ReverseAD.jl | 70 +++++++++++++++++++++++++++++++----------- 2 files changed, 63 insertions(+), 25 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index 96d516f..f32f1f8 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -33,7 +33,7 @@ end ) `edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the -diagonal. +diagonal. Diagonal entries absent from `edgelist` are not added to the pattern. Returns `(colptr, I, J, result)` where `colptr`, `I` and `J` define the lower triangular CSC sparsity structure of the Hessian (in global variable indices), @@ -48,10 +48,13 @@ function _hessian_color_preprocess( resize!(seen_idx, num_total_var) # Collect off-diagonal lower-triangular entries (local coords, filled later) I_off, J_off = Int[], Int[] + diagonal_indices = Int[] for (ei, ej) in edgelist push!(seen_idx, ei) push!(seen_idx, ej) - if ei != ej + if ei == ej + push!(diagonal_indices, ei) + else # Store in lower triangular format: row > col push!(I_off, max(ei, ej)) push!(J_off, min(ei, ej)) @@ -72,16 +75,17 @@ function _hessian_color_preprocess( n = length(local_indices) - # Build full symmetric matrix: both (i,j) and (j,i) for off-diagonal, plus diagonal + # Build the full symmetric pattern, keeping only the supplied diagonal. I_full, J_full = Int[], Int[] for k in eachindex(I_off) - push!(I_full, I_off[k]); + push!(I_full, I_off[k]) push!(J_full, J_off[k]) # lower - push!(I_full, J_off[k]); + push!(I_full, J_off[k]) push!(J_full, I_off[k]) # upper (transpose) end - for k in 1:n - push!(I_full, k); + for i in diagonal_indices + k = global_to_local_idx[i] + push!(I_full, k) push!(J_full, k) # diagonal end mat_sym = diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index c9e46ed..096f658 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -417,22 +417,57 @@ function test_coloring_end_to_end_hessian_coloring_and_recovery() # Test the new coloring API through the compatibility layer coloring_algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(; decompression = :substitution) - colptr, I, J, rinfo = ArrayDiff._hessian_color_preprocess( - Set([(1, 2)]), - 2, - coloring_algorithm, - MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(0), + for (edges, hess, expected_I, expected_J, expected_V) in ( + ( + Set([(1, 1), (1, 2), (2, 2)]), + [3.4 2.1; 2.1 1.3], + [1, 2, 2], + [1, 1, 2], + [3.4, 2.1, 1.3], + ), + (Set([(1, 2)]), [0.0 2.1; 2.1 0.0], [2], [1], [2.1]), + (Set([(1, 2), (2, 2)]), [0.0 2.1; 2.1 1.3], [2, 2], [1, 2], [2.1, 1.3]), ) - R = ArrayDiff._seed_matrix(rinfo) - ArrayDiff._prepare_seed_matrix!(R, rinfo) - @test I == [1, 2, 2] - @test J == [1, 1, 2] - @test R == [1.0 0.0; 0.0 1.0] - hess = [3.4 2.1; 2.1 1.3] - matmat = hess * R - V = zeros(3) - ArrayDiff._recover_from_matmat!(colptr, V, matmat, rinfo, zeros(3)) - @test V == [3.4, 2.1, 1.3] + colptr, I, J, rinfo = ArrayDiff._hessian_color_preprocess( + edges, + 2, + coloring_algorithm, + MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(0), + ) + R = ArrayDiff._seed_matrix(rinfo) + ArrayDiff._prepare_seed_matrix!(R, rinfo) + @test I == expected_I + @test J == expected_J + @test R == [1.0 0.0; 0.0 1.0] + matmat = hess * R + V = zeros(length(I)) + ArrayDiff._recover_from_matmat!(colptr, V, matmat, rinfo, zeros(2)) + @test V == expected_V + end + return +end + +function test_hessian_structurally_zero_diagonal() + variables = MOI.VariableIndex.(1:4) + x, y = variables[2], variables[4] + for (expr, structure) in + ((:($x * $y), [(4, 2)]), (:($x * $y + $y^3), [(4, 2), (4, 4)])) + model = ArrayDiff.Model() + ArrayDiff.set_objective(model, expr) + evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), variables) + MOI.initialize(evaluator, [:Hess]) + @test MOI.hessian_objective_structure(evaluator) == structure + @test MOI.hessian_lagrangian_structure(evaluator) == structure + for y_value in (0.0, 2.0) + values = [0.0, 3.0, 0.0, y_value] + expected = length(structure) == 1 ? [1.0] : [1.0, 6 * y_value] + H = fill(NaN, length(structure)) + MOI.eval_hessian_objective(evaluator, H, values) + @test H == expected + MOI.eval_hessian_lagrangian(evaluator, H, values, 1.5, Float64[]) + @test H == 1.5 .* expected + end + end return end @@ -1309,9 +1344,8 @@ function test_hessian_reinterpret_unsafe() H = zeros(length(H_s)) x_v = ones(5) MOI.eval_hessian_lagrangian(evaluator, H, x_v, 0.0, [1.0, 1.0]) - @test count(isapprox.(H, 1.0; atol = 1e-8)) == 3 - @test count(isapprox.(H, 0.0; atol = 1e-8)) == 5 - @test sort(H_s[round.(Bool, H)]) == [(3, 1), (3, 2), (5, 4)] + @test H ≈ ones(3) + @test sort(H_s) == [(3, 1), (3, 2), (5, 4)] return end From 95a5564453dbde577a4051d3d974c54061e6021e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 16 Oct 2026 09:27:49 +0200 Subject: [PATCH 13/14] Simplifies --- src/coloring.jl | 125 ++++++++---------------------------- src/forward_over_reverse.jl | 2 - test/test_ReverseAD.jl | 45 ++++++++++++- 3 files changed, 71 insertions(+), 101 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index f32f1f8..10a1dc3 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -8,20 +8,14 @@ struct ColoringResult{R<:SMC.AbstractColoringResult} result::R local_indices::Vector{Int} # map from local to global indices - full_colptr::Vector{Int} # colptr of the full symmetric matrix used for coloring - lower_pos::Vector{Int} # positions of lower-triangular entries in the full nzval - full_buffer::Vector{Float64} # pre-allocated buffer of size nnz(full symmetric matrix) end -Wrapper around AbstractColoringResult that also stores auxiliary data needed -for Hessian recovery from a full symmetric matrix decompression. +Wrapper around AbstractColoringResult that maps local Hessian indices back +to global variable indices. """ struct ColoringResult{R<:SMC.AbstractColoringResult} result::R local_indices::Vector{Int} # map from local to global indices - full_colptr::Vector{Int} # colptr of full symmetric matrix used for coloring - lower_pos::Vector{Int} # positions of lower-triangular entries in full nzval - full_buffer::Vector{Float64} # scratch buffer of length nnz(full symmetric matrix) end """ @@ -29,36 +23,29 @@ end edgelist, num_total_var, algo::SMC.GreedyColoringAlgorithm, - seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0), + seen_idx = MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(0), ) `edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the diagonal. Diagonal entries absent from `edgelist` are not added to the pattern. -Returns `(colptr, I, J, result)` where `colptr`, `I` and `J` define the lower -triangular CSC sparsity structure of the Hessian (in global variable indices), -and `result` is a `ColoringResult` wrapping an SMC coloring result. +Returns `(colptr, I, J, result)` where `colptr` contains the local lower- +triangular CSC column pointers and `I`, `J` give the corresponding entries +in global variable indices. `result` wraps an SMC coloring result. """ function _hessian_color_preprocess( edgelist, num_total_var, algo::SMC.GreedyColoringAlgorithm, - seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0), + seen_idx = MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(0), ) resize!(seen_idx, num_total_var) - # Collect off-diagonal lower-triangular entries (local coords, filled later) - I_off, J_off = Int[], Int[] - diagonal_indices = Int[] - for (ei, ej) in edgelist - push!(seen_idx, ei) - push!(seen_idx, ej) - if ei == ej - push!(diagonal_indices, ei) - else - # Store in lower triangular format: row > col - push!(I_off, max(ei, ej)) - push!(J_off, min(ei, ej)) - end + I, J = Int[], Int[] + for (i, j) in edgelist + push!(seen_idx, i) + push!(seen_idx, j) + push!(I, max(i, j)) + push!(J, min(i, j)) end local_indices = sort!(collect(seen_idx)) empty!(seen_idx) @@ -67,70 +54,25 @@ function _hessian_color_preprocess( for k in eachindex(local_indices) global_to_local_idx[local_indices[k]] = k end - # Map off-diagonal entries to local indices - for k in eachindex(I_off) - I_off[k] = global_to_local_idx[I_off[k]] - J_off[k] = global_to_local_idx[J_off[k]] + for k in eachindex(I) + I[k] = global_to_local_idx[I[k]] + J[k] = global_to_local_idx[J[k]] end n = length(local_indices) - - # Build the full symmetric pattern, keeping only the supplied diagonal. - I_full, J_full = Int[], Int[] - for k in eachindex(I_off) - push!(I_full, I_off[k]) - push!(J_full, J_off[k]) # lower - push!(I_full, J_off[k]) - push!(J_full, I_off[k]) # upper (transpose) - end - for i in diagonal_indices - k = global_to_local_idx[i] - push!(I_full, k) - push!(J_full, k) # diagonal - end - mat_sym = - SparseArrays.sparse(I_full, J_full, trues(length(I_full)), n, n, |) - - # Perform coloring on full symmetric matrix - S = SMC.SparsityPatternCSC(mat_sym) + mat = SparseArrays.sparse(I, J, trues(length(I)), n, n, |) + # Color the full symmetric pattern, but recover only its lower triangle. + S = SMC.SparsityPatternCSC( + SparseArrays.sparse(LinearAlgebra.Symmetric(mat, :L)), + ) problem = SMC.ColoringProblem(; structure = :symmetric, partition = :column) tree_result = SMC.coloring(S, problem, algo) - # Find positions of lower-triangular entries within the full CSC nzval array. - # findnz on a CSC matrix returns elements in CSC (column-major) order, - # matching the nzval layout exactly. - I_nz, J_nz, _ = SparseArrays.findnz(mat_sym) - lower_pos = findall(k -> I_nz[k] >= J_nz[k], eachindex(I_nz)) - - # Lower-triangular CSC-ordered local indices - I_low_csc = I_nz[lower_pos] - J_low_csc = J_nz[lower_pos] - - # Map back to global indices for the returned hess_I / hess_J - I_global = [local_indices[i] for i in I_low_csc] - J_global = [local_indices[j] for j in J_low_csc] - - # Build lower-triangular sparse matrix to obtain its colptr - mat_low = SparseArrays.sparse( - I_low_csc, - J_low_csc, - trues(length(I_low_csc)), - n, - n, - |, - ) - - full_buffer = Vector{Float64}(undef, SparseArrays.nnz(mat_sym)) - - result = ColoringResult( - tree_result, - local_indices, - copy(mat_sym.colptr), - lower_pos, - full_buffer, - ) - - return copy(mat_low.colptr), I_global, J_global, result + I, J, _ = SparseArrays.findnz(mat) + I_global = local_indices[I] + J_global = local_indices[J] + result = ColoringResult(tree_result, local_indices) + return mat.colptr, I_global, J_global, result end """ @@ -169,30 +111,17 @@ end V::AbstractVector{T}, R::AbstractMatrix{T}, result::ColoringResult, - stored_values::AbstractVector{T}, ) where {T} Recover the Hessian values from the Hessian-matrix product H*R_seed. R is the result of H*R_seed where R_seed is the seed matrix. -`stored_values` is a temporary vector (unused, kept for API compatibility). """ function _recover_from_matmat!( colptr::AbstractVector, V::AbstractVector{T}, R::AbstractMatrix{T}, result::ColoringResult, - stored_values::AbstractVector{T}, ) where {T} - # Decompress into the full symmetric buffer, then extract lower-triangular values. - SMC.decompress_csc!( - result.full_buffer, - result.full_colptr, - R, - result.result, - :F, - ) - for k in eachindex(V) - V[k] = result.full_buffer[result.lower_pos[k]] - end + SMC.decompress_csc!(V, colptr, R, result.result, :L) return end diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl index fd7e8a7..c103ebd 100644 --- a/src/forward_over_reverse.jl +++ b/src/forward_over_reverse.jl @@ -65,7 +65,6 @@ 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), @@ -76,7 +75,6 @@ function _eval_hessian( output_slice, ex.seed_matrix, ex.rinfo, - d.output_ϵ, ) for i in 1:length(output_slice) output_slice[i] *= scale diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index 096f658..99bbc17 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -441,7 +441,7 @@ function test_coloring_end_to_end_hessian_coloring_and_recovery() @test R == [1.0 0.0; 0.0 1.0] matmat = hess * R V = zeros(length(I)) - ArrayDiff._recover_from_matmat!(colptr, V, matmat, rinfo, zeros(2)) + ArrayDiff._recover_from_matmat!(colptr, V, matmat, rinfo) @test V == expected_V end return @@ -471,6 +471,49 @@ function test_hessian_structurally_zero_diagonal() return end +function test_coloring_lower_triangular_recovery() + for order in (ArrayDiff.SMC.NaturalOrder(), ArrayDiff.SMC.LargestFirst()) + algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm( + order; + decompression = :substitution, + ) + for edges in ( + Set([(2, 4), (4, 7), (7, 9)]), + Set([(4, 2), (2, 7), (9, 2), (7, 7)]), + Set([(4, 2), (4, 7), (7, 9), (9, 2), (4, 4), (9, 9)]), + Set([(2, 4), (4, 2), (2, 2)]), + Set([(2, 2), (9, 9)]), + Set{Tuple{Int,Int}}(), + ) + H = zeros(9, 9) + for (i, j) in edges + H[i, j] = H[j, i] = 10 * min(i, j) + max(i, j) + end + colptr, I, J, result = + ArrayDiff._hessian_color_preprocess(edges, 9, algorithm) + expected_structure = sort!( + unique([(max(i, j), min(i, j)) for (i, j) in edges]); + by = reverse, + ) + @test collect(zip(I, J)) == expected_structure + R = ArrayDiff._seed_matrix(result) + ArrayDiff._prepare_seed_matrix!(R, result) + local_indices = result.local_indices + compressed = H[local_indices, local_indices] * R + V = fill(NaN, length(I)) + ArrayDiff._recover_from_matmat!(colptr, V, compressed, result) + @test V == [H[i, j] for (i, j) in expected_structure] + @test 0 == @allocated ArrayDiff._recover_from_matmat!( + colptr, + V, + compressed, + result, + ) + end + end + return +end + function test_derivatives() a = MOI.VariableIndex(1) b = MOI.VariableIndex(2) From a43abaf08781a565e64e5adedf314c32d8d29e59 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 16 Oct 2026 10:12:47 +0200 Subject: [PATCH 14/14] Remove unused --- src/utils.jl | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/src/utils.jl b/src/utils.jl index dbd3c9b..d1d69e4 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -203,18 +203,3 @@ 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