Fix Hessian recovery with star coloring

src/coloring.jl

@@ -31,7 +31,8 @@ diagonal. Diagonal entries absent from `edgelist` are not added to the pattern.
 
 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.
+in global variable indices. `result` wraps an SMC coloring result, or is
+`nothing` if the pattern is empty.
 """
 function _hessian_color_preprocess(
     edgelist,
@@ -49,6 +50,11 @@ function _hessian_color_preprocess(
     end
     local_indices = sort!(collect(seen_idx))
     empty!(seen_idx)
+    # SMC's star coloring does not support an empty pattern.
+    # TODO: remove when https://github.com/JuliaDiff/SparseMatrixColorings.jl/pull/352 lands
+    if isempty(local_indices)
+        return [1], Int[], Int[], nothing
+    end
 
     global_to_local_idx = seen_idx.nzidx # steal for storage
     for k in eachindex(local_indices)
@@ -86,6 +92,10 @@ function _seed_matrix(result::ColoringResult)
     return Matrix{Float64}(undef, n, ncolors)
 end
 
+function _seed_matrix(::Nothing)
+    return Matrix{Float64}(undef, 0, 0)
+end
+
 """
     _prepare_seed_matrix!(R, result::ColoringResult)
 
@@ -125,3 +135,26 @@ function _recover_from_matmat!(
     SMC.decompress_csc!(V, colptr, R, result.result, :L)
     return
 end
+
+function _recover_from_matmat!(
+    colptr::AbstractVector,
+    V::AbstractVector{T},
+    R::AbstractMatrix{T},
+    result::ColoringResult{<:SMC.StarSetColoringResult},
+) where {T}
+    # SMC.decompress_csc! only supports substitution decompression. Mirror
+    # SMC.decompress! for a lower-triangular CSC matrix without constructing it.
+    S = SMC.sparsity_pattern(result.result)
+    rows = SparseArrays.rowvals(S)
+    compressed_indices = result.result.compressed_indices
+    l = 0
+    for j in axes(S, 2)
+        for k in SparseArrays.nzrange(S, j)
+            if rows[k] >= j
+                l += 1
+                V[l] = R[compressed_indices[k]]
+            end
+        end
+    end
+    return
+end

src/forward_over_reverse.jl

@@ -38,8 +38,7 @@ function _eval_hessian(
     scale::Float64,
     nzcount::Int,
 )::Int
-    if ex.expr.linearity == LINEAR
-        @assert length(ex.hess_I) == 0
+    if isempty(ex.hess_I)
         return 0
     end
     chunk = min(size(ex.seed_matrix, 2), d.max_chunk)

test/test_ReverseAD.jl

@@ -450,17 +450,32 @@ 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)]))
+    for decompression in (:direct, :substitution),
+        (expr, structure) in (
+            (:($x * $y), [(4, 2)]),
+            (:($x * $y + $y^3), [(4, 2), (4, 4)]),
+            (:($x + $y), Tuple{Int,Int}[]),
+            (:(3.0), Tuple{Int,Int}[]),
+            (:(ifelse($x < 0, $x, $y)), Tuple{Int,Int}[]),
+        )
+
         model = ArrayDiff.Model()
         ArrayDiff.set_objective(model, expr)
-        evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), variables)
+        algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(; decompression)
+        evaluator =
+            ArrayDiff.Evaluator(model, ArrayDiff.Mode(algorithm), 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]
+            expected = if isempty(structure)
+                Float64[]
+            elseif length(structure) == 1
+                [1.0]
+            else
+                [1.0, 6 * y_value]
+            end
             H = fill(NaN, length(structure))
             MOI.eval_hessian_objective(evaluator, H, values)
             @test H == expected
@@ -472,11 +487,11 @@ function test_hessian_structurally_zero_diagonal()
 end
 
 function test_coloring_lower_triangular_recovery()
-    for order in (ArrayDiff.SMC.NaturalOrder(), ArrayDiff.SMC.LargestFirst())
-        algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(
-            order;
-            decompression = :substitution,
-        )
+    for (T, decompression) in
+        ((Float32, :direct), (Float64, :direct), (Float64, :substitution)),
+        order in (ArrayDiff.SMC.NaturalOrder(), ArrayDiff.SMC.LargestFirst())
+
+        algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(order; decompression)
         for edges in (
             Set([(2, 4), (4, 7), (7, 9)]),
             Set([(4, 2), (2, 7), (9, 2), (7, 7)]),
@@ -485,7 +500,7 @@ function test_coloring_lower_triangular_recovery()
             Set([(2, 2), (9, 9)]),
             Set{Tuple{Int,Int}}(),
         )
-            H = zeros(9, 9)
+            H = zeros(T, 9, 9)
             for (i, j) in edges
                 H[i, j] = H[j, i] = 10 * min(i, j) + max(i, j)
             end
@@ -496,11 +511,17 @@ function test_coloring_lower_triangular_recovery()
                 by = reverse,
             )
             @test collect(zip(I, J)) == expected_structure
-            R = ArrayDiff._seed_matrix(result)
+            R = Matrix{T}(ArrayDiff._seed_matrix(result))
+            if isempty(expected_structure)
+                @test result === nothing
+                @test colptr == [1]
+                @test size(R) == (0, 0)
+                continue
+            end
             ArrayDiff._prepare_seed_matrix!(R, result)
             local_indices = result.local_indices
             compressed = H[local_indices, local_indices] * R
-            V = fill(NaN, length(I))
+            V = fill(T(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!(
@@ -514,6 +535,34 @@ function test_coloring_lower_triangular_recovery()
     return
 end
 
+function test_hessian_star_coloring()
+    x, y = MOI.VariableIndex.(1:2)
+    model = ArrayDiff.Model()
+    ArrayDiff.set_objective(model, :((1 - $x)^2 + 100 * ($y - $x^2)^2))
+    ArrayDiff.add_constraint(model, :($x * $y), MOI.LessThan(2.0))
+    ArrayDiff.add_constraint(model, :($y^2), MOI.LessThan(2.0))
+    algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(; decompression = :direct)
+    evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(algorithm), [x, y])
+    MOI.initialize(evaluator, [:Hess])
+    @test MOI.hessian_lagrangian_structure(evaluator) ==
+          [(1, 1), (2, 1), (2, 2), (2, 1), (2, 2)]
+    H = zeros(5)
+    μ = [2.0, 3.0]
+    for xv in ([1.0, 1.0], [-1.2, 1.0])
+        MOI.eval_hessian_lagrangian(evaluator, H, xv, 1.5, μ)
+        @test H ≈ [
+            1.5 * (2 + 1200 * xv[1]^2 - 400 * xv[2]),
+            1.5 * (-400 * xv[1]),
+            300.0,
+            2.0,
+            6.0,
+        ]
+        @test 0 ==
+              @allocated MOI.eval_hessian_lagrangian(evaluator, H, xv, 1.5, μ)
+    end
+    return
+end
+
 function test_derivatives()
     a = MOI.VariableIndex(1)
     b = MOI.VariableIndex(2)