Add tests

Author: odow

src/Nonlinear/ReverseAD/Coloring/Coloring.jl

@@ -6,7 +6,7 @@
 
 module Coloring
 
-include("Forest.jl")
+include("IntDisjointSet.jl")
 include("topological_sort.jl")
 
 """
@@ -176,7 +176,7 @@ function _grow_star(v, w, e_idx, firstNeighbor, color, S)
     return
 end
 
-function _merge_trees(S::_Forest, eg::Int, eg1::Int)
+function _merge_trees(S::_IntDisjointSet, eg::Int, eg1::Int)
     if _find_root!(S, eg) != _find_root!(S, eg1)
         _root_union!(S, eg, eg1)
     end
@@ -202,7 +202,7 @@ function acyclic_coloring(g::UndirectedGraph)
     firstNeighbor = _Edge[]
     firstVisitToTree = fill(_Edge(0, 0, 0), _num_edges(g))
     color = fill(0, _num_vertices(g))
-    S = _Forest(_num_edges(g))
+    S = _IntDisjointSet(_num_edges(g))
     @inbounds for v in 1:_num_vertices(g)
         n_neighbor = _num_neighbors(v, g)
         start_neighbor = _start_neighbors(v, g)

…c/Nonlinear/ReverseAD/Coloring/Forest.jl …ear/ReverseAD/Coloring/IntDisjointSet.jlsrc/Nonlinear/ReverseAD/Coloring/Forest.jl renamed to src/Nonlinear/ReverseAD/Coloring/IntDisjointSet.jl src/Nonlinear/ReverseAD/Coloring/Forest.jl renamed to src/Nonlinear/ReverseAD/Coloring/IntDisjointSet.jl

@@ -6,46 +6,42 @@
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
 # The code in this file was taken from
-# https://github.com/gdalle/SparseMatrixColorings.jl/blob/main/src/forest.jl
+# https://github.com/gdalle/SparseMatrixColorings.jl/blob/main/src/Forest.jl
 #
 # It was copied at the suggestion of Alexis in his JuMP-dev 2025 talk.
 #
 # @odow made minor changes to match MOI coding styles.
 #
 # x-ref https://github.com/gdalle/SparseMatrixColorings.jl/pull/190
 
-mutable struct _Forest
-    # current number of distinct trees in the forest
+mutable struct _IntDisjointSet
+    # current number of distinct trees in the S
     number_of_trees::Int
     # vector storing the index of a parent in the tree for each edge, used in
     # union-find operations
     parents::Vector{Int}
     # vector approximating the depth of each tree to optimize path compression
     ranks::Vector{Int}
 
-    _Forest(n::Integer) = new(n, collect(Base.OneTo(n)), zeros(Int, n))
+    _IntDisjointSet(n::Integer) = new(n, collect(1:n), zeros(Int, n))
 end
 
-function _find_root!(parents::Vector{Int}, index_edge::Integer)
-    p = parents[index_edge]
-    if parents[p] != p
-        parents[index_edge] = p = _find_root!(parents, p)
+function _find_root!(S::_IntDisjointSet, x::Integer)
+    p = S.parents[x]
+    if S.parents[p] != p
+        S.parents[x] = p = _find_root!(S, p)
     end
     return p
 end
 
-function _find_root!(forest::_Forest, index_edge::Integer)
-    return _find_root!(forest.parents, index_edge)
-end
-
-function _root_union!(forest::_Forest, index_edge1::Int, index_edge2::Int)
-    rank1, rank2 = forest.ranks[index_edge1], forest.ranks[index_edge2]
+function _root_union!(S::_IntDisjointSet, x::Int, y::Int)
+    rank1, rank2 = S.ranks[x], S.ranks[y]
     if rank1 < rank2
-        index_edge1, index_edge2 = index_edge2, index_edge1
+        x, y = y, x
     elseif rank1 == rank2
-        forest.ranks[index_edge1] += 1
+        S.ranks[x] += 1
     end
-    forest.parents[index_edge2] = index_edge1
-    forest.number_of_trees -= 1
+    S.parents[y] = x
+    S.number_of_trees -= 1
     return
 end

test/Nonlinear/ReverseAD.jl

@@ -11,9 +11,9 @@ import LinearAlgebra
 import MathOptInterface as MOI
 import SparseArrays
 
-const Nonlinear = MOI.Nonlinear
-const ReverseAD = Nonlinear.ReverseAD
-const Coloring = ReverseAD.Coloring
+import MathOptInterface.Nonlinear
+import MathOptInterface.Nonlinear.ReverseAD
+import MathOptInterface.Nonlinear.ReverseAD.Coloring
 
 function runtests()
     for name in names(@__MODULE__; all = true)
@@ -1421,6 +1421,33 @@ function test_hessian_reinterpret_unsafe()
     return
 end
 
+function test_IntDisjointSet()
+    for case in [
+        [(1, 2) => [1, 1, 3], (1, 3) => [1, 1, 1]],
+        [(1, 2) => [1, 1, 3], (3, 1) => [1, 1, 1]],
+        [(2, 1) => [2, 2, 3], (1, 3) => [2, 2, 2]],
+        [(2, 1) => [2, 2, 3], (3, 1) => [3, 2, 3]],
+        [(1, 3) => [1, 2, 1], (2, 3) => [1, 2, 2]],
+        [(1, 3) => [1, 2, 1], (3, 2) => [1, 1, 1]],
+        [(3, 1) => [3, 2, 3], (2, 3) => [3, 3, 3]],
+        [(3, 1) => [3, 2, 3], (3, 2) => [3, 3, 3]],
+        [(2, 3) => [1, 2, 2], (1, 3) => [1, 2, 1]],
+        [(2, 3) => [1, 2, 2], (3, 1) => [2, 2, 2]],
+        [(3, 2) => [1, 3, 3], (1, 3) => [3, 3, 3]],
+        [(3, 2) => [1, 3, 3], (3, 1) => [3, 3, 3]],
+    ]
+        S = Coloring._IntDisjointSet(3)
+        @test Coloring._find_root!.((S,), [1, 2, 3]) == [1, 2, 3]
+        @test S.number_of_trees == 3
+        for (i, (union, result)) in enumerate(case)
+            Coloring._root_union!(S, union[1], union[2])
+            @test Coloring._find_root!.((S,), [1, 2, 3]) == result
+            @test S.number_of_trees == 3 - i
+        end
+    end
+    return
+end
+
 end  # module
 
 TestReverseAD.runtests()