diff --git a/Project.toml b/Project.toml index 014259e..5b92302 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ArrayDiff" uuid = "c45fa1ca-6901-44ac-ae5b-5513a4852d50" -authors = ["Benoît Legat "] version = "0.1.0" +authors = ["Benoît Legat "] [deps] DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" @@ -9,6 +9,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" [compat] DataStructures = "0.18, 0.19" @@ -16,4 +17,5 @@ ForwardDiff = "1" MathOptInterface = "1.40" NaNMath = "1" SparseArrays = "1.10" +SparseMatrixColorings = "0.4" julia = "1.10" diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl index 1f08c1e..2b78193 100644 --- a/src/ArrayDiff.jl +++ b/src/ArrayDiff.jl @@ -6,26 +6,38 @@ module ArrayDiff +import SparseArrays +import SparseMatrixColorings as SMC import ForwardDiff import MathOptInterface as MOI const Nonlinear = MOI.Nonlinear -import SparseArrays """ - Mode() <: AbstractAutomaticDifferentiation + Mode(coloring_algorithm::SMC.GreedyColoringAlgorithm) <: AbstractAutomaticDifferentiation Fork of `MOI.Nonlinear.SparseReverseMode` to add array support. """ -struct Mode <: MOI.Nonlinear.AbstractAutomaticDifferentiation end +struct Mode{C<:SMC.GreedyColoringAlgorithm} <: + MOI.Nonlinear.AbstractAutomaticDifferentiation + coloring_algorithm::C +end + +function Mode() + return Mode( + SMC.GreedyColoringAlgorithm(; + decompression = :substitution, + ), + ) +end function MOI.Nonlinear.Evaluator( model::MOI.Nonlinear.Model, - ::Mode, + mode::Mode, ordered_variables::Vector{MOI.VariableIndex}, ) return MOI.Nonlinear.Evaluator( model, - NLPEvaluator(model, ordered_variables), + NLPEvaluator(model, ordered_variables, mode.coloring_algorithm), ) end @@ -48,7 +60,7 @@ import NaNMath: pow, sqrt -include("Coloring/Coloring.jl") +include("coloring.jl") include("graph_tools.jl") include("sizes.jl") include("types.jl") diff --git a/src/Coloring/Coloring.jl b/src/Coloring/Coloring.jl deleted file mode 100644 index 1ef8dc5..0000000 --- a/src/Coloring/Coloring.jl +++ /dev/null @@ -1,574 +0,0 @@ -# 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. - -module Coloring - -import DataStructures - -include("topological_sort.jl") - -""" - mutable struct IndexedSet - nzidx::Vector{Int} - empty::BitVector - nnz::Int - end - -Represent the set `nzidx[1:nnz]` by additionally setting `empty[i]` to `false` -for each element `i` of the set for fast membership check. -""" -mutable struct IndexedSet - nzidx::Vector{Int} - empty::BitVector - nnz::Int -end - -IndexedSet(n::Integer) = IndexedSet(zeros(Int, n), trues(n), 0) - -function Base.push!(v::IndexedSet, i::Integer) - if v.empty[i] # new index - v.nzidx[v.nnz+=1] = i - v.empty[i] = false - end - return -end - -function Base.empty!(v::IndexedSet) - for i in 1:v.nnz - v.empty[v.nzidx[i]] = true - end - v.nnz = 0 - return v -end - -# Returns the maximum index that the set can contain, -# not the cardinality of the set like `length(::Base.Set)` -Base.length(v::IndexedSet) = length(v.nzidx) - -function Base.resize!(v::IndexedSet, n::Integer) - if n > 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..e23c3fa --- /dev/null +++ b/src/coloring.jl @@ -0,0 +1,243 @@ +# 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) + + # Handle empty case (no edges in Hessian) + if isempty(local_indices) + # Return empty structure - no variables to color + # We still need to return a valid ColoringResult, but with empty local_indices + # The I and J vectors are already empty, which is correct + # For the result, we'll create a minimal valid structure with a diagonal element + # Note: This case should rarely occur in practice + S = SMC.SparsityPatternCSC(SparseArrays.spdiagm(0 => [true])) + problem = SMC.ColoringProblem(; + structure = :symmetric, + partition = :column, + ) + tree_result = SMC.coloring(S, problem, algo) + result = ColoringResult(tree_result, Int[]) + 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 + 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, + ) + 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 +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!( + 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!( + 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) + return +end diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl index ea4c7e2..5f3be96 100644 --- a/src/forward_over_reverse.jl +++ b/src/forward_over_reverse.jl @@ -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,13 +65,12 @@ 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!( - 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 diff --git a/src/graph_tools.jl b/src/graph_tools.jl index 7e9c366..f813016 100644 --- a/src/graph_tools.jl +++ b/src/graph_tools.jl @@ -151,7 +151,7 @@ end """ _compute_gradient_sparsity!( - indices::Coloring.IndexedSet, + indices::IndexedSet, nodes::Vector{Nonlinear.Node}, ) @@ -159,7 +159,7 @@ Compute the sparsity pattern of the gradient of an expression (that is, a list o which variable indices are present). """ function _compute_gradient_sparsity!( - indices::Coloring.IndexedSet, + indices::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet, nodes::Vector{Nonlinear.Node}, ) for node in nodes @@ -180,7 +180,7 @@ end nodes::Vector{Nonlinear.Node}, adj, input_linearity::Vector{Linearity}, - indexedset::Coloring.IndexedSet, + indexedset::IndexedSet, subexpression_edgelist::Vector{Set{Tuple{Int,Int}}}, subexpression_variables::Vector{Vector{Int}}, ) @@ -207,7 +207,7 @@ function _compute_hessian_sparsity( nodes::Vector{Nonlinear.Node}, adj, input_linearity::Vector{Linearity}, - indexedset::Coloring.IndexedSet, + indexedset::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet, subexpression_edgelist::Vector{Set{Tuple{Int,Int}}}, subexpression_variables::Vector{Vector{Int}}, ) diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl index 21c6804..943ea61 100644 --- a/src/mathoptinterface_api.jl +++ b/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, requested_features::Vector{Symbol}) +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 @@ -38,11 +41,11 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) d.objective = nothing d.user_output_buffer = zeros(largest_user_input_dimension) d.jac_storage = zeros(max(N, largest_user_input_dimension)) - d.constraints = _FunctionStorage[] + d.constraints = _FunctionStorage{R}[] d.last_x = fill(NaN, N) 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 # @@ -111,11 +114,11 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) shared_partials_storage_ϵ, d, ) - objective = _FunctionStorage( + objective = _FunctionStorage{R}( subexpr, N, coloring_storage, - d.want_hess, + d.want_hess ? d.coloring_algorithm : nothing, d.subexpressions, individual_order[1], subexpression_edgelist, @@ -137,11 +140,11 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) ) push!( d.constraints, - _FunctionStorage( + _FunctionStorage{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 4f1f19a..d7373f9 100644 --- a/src/types.jl +++ b/src/types.jl @@ -63,28 +63,31 @@ function _subexpression_and_linearity( linearity end -struct _FunctionStorage +struct _FunctionStorage{R<:SMC.AbstractColoringResult} expr::_SubexpressionStorage grad_sparsity::Vector{Int} # Nonzero pattern of Hessian matrix hess_I::Vector{Int} hess_J::Vector{Int} - rinfo::Coloring.RecoveryInfo # coloring info for hessians + rinfo::Union{Nothing,ColoringResult{R}} seed_matrix::Matrix{Float64} # subexpressions which this function depends on, ordered for forward pass. dependent_subexpressions::Vector{Int} - function _FunctionStorage( + function _FunctionStorage{R}( expr::_SubexpressionStorage, 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}, dependent_subexpressions, subexpression_edgelist, subexpression_variables, linearity::Vector{Linearity}, - ) + ) where {R} empty!(coloring_storage) _compute_gradient_sparsity!(coloring_storage, expr.nodes) for k in dependent_subexpressions @@ -95,7 +98,7 @@ struct _FunctionStorage 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, @@ -104,13 +107,14 @@ struct _FunctionStorage subexpression_edgelist, subexpression_variables, ) - hess_I, hess_J, rinfo = Coloring.hessian_color_preprocess( + hess_I, hess_J, rinfo = _hessian_color_preprocess( edgelist, num_variables, + coloring_algorithm, coloring_storage, ) - seed_matrix = Coloring.seed_matrix(rinfo) - return new( + seed_matrix = _seed_matrix(rinfo) + return new{R}( expr, grad_sparsity, hess_I, @@ -120,12 +124,12 @@ struct _FunctionStorage dependent_subexpressions, ) else - return new( + return new{R}( expr, grad_sparsity, Int[], Int[], - Coloring.RecoveryInfo(), + nothing, Array{Float64}(undef, 0, 0), dependent_subexpressions, ) @@ -137,6 +141,7 @@ end NLPEvaluator( model::Nonlinear.Model, ordered_variables::Vector{MOI.VariableIndex}, + coloring_algorithm::SMC.AbstractColoringAlgorithm = SMC.GreedyColoringAlgorithm(; decompression=:substitution), ) Return an `NLPEvaluator` object that implements the `MOI.AbstractNLPEvaluator` @@ -145,12 +150,16 @@ interface. !!! warning Before using, you must initialize the evaluator using `MOI.initialize`. """ -mutable struct NLPEvaluator <: MOI.AbstractNLPEvaluator +mutable struct NLPEvaluator{ + R, + C<:SMC.GreedyColoringAlgorithm, +} <: MOI.AbstractNLPEvaluator data::Nonlinear.Model ordered_variables::Vector{MOI.VariableIndex} + coloring_algorithm::C - objective::Union{Nothing,_FunctionStorage} - constraints::Vector{_FunctionStorage} + objective::Union{Nothing,_FunctionStorage{R}} + constraints::Vector{_FunctionStorage{R}} subexpressions::Vector{_SubexpressionStorage} subexpression_order::Vector{Int} # Storage for the subexpressions in reverse-mode automatic differentiation. @@ -184,7 +193,21 @@ mutable struct NLPEvaluator <: MOI.AbstractNLPEvaluator function NLPEvaluator( data::Nonlinear.Model, ordered_variables::Vector{MOI.VariableIndex}, + coloring_algorithm::SMC.GreedyColoringAlgorithm = SMC.GreedyColoringAlgorithm(; + decompression = :substitution, + ), ) - return new(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{R,C}(data, ordered_variables, coloring_algorithm) end end diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl index 4b90f3b..5e894ec 100644 --- a/test/ArrayDiff.jl +++ b/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) diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl index 552ff3f..e2605e0 100644 --- a/test/ReverseAD.jl +++ b/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,92 @@ 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, + ) + 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!(V, matmat, rinfo, zeros(3)) @test V == [3.4, 1.3, 2.1] return end @@ -547,7 +559,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,