Update the decompression of sparse Hessian

JuliaSmoothOptimizers · Jul 18, 2024 · b98bda8 · b98bda8
1 parent 97d9efa
commit b98bda8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 59 deletions.
diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl
@@ -5,7 +5,6 @@ struct SparseADHessian{Tag, GT, S, T} <: ADNLPModels.ADBackend
   colors::Vector{Int}
   ncolors::Int
   dcolors::Dict{Int, Vector{Tuple{Int,Int}}}
-  star_set::SparseMatrixColorings.StarSet
   res::S
   lz::Vector{ForwardDiff.Dual{Tag, T, 1}}
   glz::Vector{ForwardDiff.Dual{Tag, T, 1}}
@@ -39,7 +38,7 @@ function SparseADHessian(
   colptr = trilH.colptr
 
   # The indices of the nonzero elements in `vals` that will be processed by color `c` are stored in `dcolors[c]`.
-  dcolors = nnz_colors(H, trilH, colors, ncolors)
+  dcolors = nnz_colors(trilH, star_set, colors, ncolors)
 
   # prepare directional derivatives
   res = similar(x0)
@@ -76,7 +75,6 @@ function SparseADHessian(
     colors,
     ncolors,
     dcolors,
-    star_set,
     res,
     lz,
     glz,
@@ -95,7 +93,6 @@ struct SparseReverseADHessian{T, S, Tagf, F, Tagψ, P} <: ADNLPModels.ADBackend
   colors::Vector{Int}
   ncolors::Int
   dcolors::Dict{Int, Vector{Tuple{Int,Int}}}
-  star_set::SparseMatrixColorings.StarSet
   res::S
   z::Vector{ForwardDiff.Dual{Tagf, T, 1}}
   gz::Vector{ForwardDiff.Dual{Tagf, T, 1}}
@@ -131,7 +128,7 @@ function SparseReverseADHessian(
   colptr = trilH.colptr
 
   # The indices of the nonzero elements in `vals` that will be processed by color `c` are stored in `dcolors[c]`.
-  dcolors = nnz_colors(H, trilH, colors, ncolors)
+  dcolors = nnz_colors(trilH, star_set, colors, ncolors)
 
   # prepare directional derivatives
   res = similar(x0)
@@ -172,7 +169,6 @@ function SparseReverseADHessian(
     colors,
     ncolors,
     dcolors,
-    star_set,
     res,
     z,
     gz,

diff --git a/src/sparsity_pattern.jl b/src/sparsity_pattern.jl
@@ -53,82 +53,46 @@ function compute_hessian_sparsity(
 end
 
 """
-    dcolors = nnz_colors(trilH, colors, ncolors)
+    dcolors = nnz_colors(trilH, star_set, colors, ncolors)
 
 Determine the coefficients in `trilH` that will be computed by a given color.
-This function leverages the symmetry of the matrix `H` and also stores the row index for a
-given coefficient in the "compressed column".
 
-# Arguments
-- `H::SparseMatrixCSC`: A sparse matrix in CSC format.
-- `trilH::SparseMatrixCSC`: The lower triangular part of `H` in CSC format.
+Arguments:
+- `trilH::SparseMatrixCSC`: The lower triangular part of a symmetric matrix in CSC format.
+- `star_set::StarSet`: A structure `StarSet` returned by the function `symmetric_coloring_detailed` of SparseMatrixColorings.jl.
 - `colors::Vector{Int}`: A vector where the i-th entry represents the color assigned to the i-th column of the matrix.
 - `ncolors::Int`: The number of distinct colors used in the coloring.
 
-# Output
+Output:
 - `dcolors::Dict{Int, Vector{Tuple{Int, Int}}}`: A dictionary where the keys are the color indices (from 1 to `ncolors`),
 and the values are vectors of tuples. Each tuple contains two integers: the first integer is the row index, and the
 second integer is the index in `trilH.nzval` where the non-zero coefficient can be found.
 """
-function nnz_colors(H, trilH, colors, ncolors)
+function nnz_colors(trilH, star_set, colors, ncolors)
   # We want to determine the coefficients in `trilH` that will be computed by a given color.
   # Because we exploit the symmetry, we also need to store the row index for a given coefficient
   # in the "compressed column".
   dcolors = Dict(i => Tuple{Int,Int}[] for i=1:ncolors)
+  star = star_set.star
+  hub = star_set.hub
 
   n = LinearAlgebra.checksquare(trilH)
   for j in 1:n
     for k in trilH.colptr[j]:trilH.colptr[j+1]-1
       i = trilH.rowval[k]
 
-      # Should we use c[i] or c[j] for this coefficient H[i,j]?
-      ci = colors[i]
-      cj = colors[j]
+      star_id = star[(i, j)]
+      h = hub[star_id]
 
-      if i == j
-        # H[i,j] is a coefficient of the diagonal
-        push!(dcolors[ci], (j,k))
-      else # i > j
-        if is_only_color_in_row(H, i, j, n, colors, ci)
-          # column i is the only column of its color c[i] with a non-zero coefficient in row j
-          push!(dcolors[ci], (j,k))
-        else
-          # column j is the only column of its color c[j] with a non-zero coefficient in row i
-          # it is ensured by the star coloring used in `symmetric_coloring`.
-          push!(dcolors[cj], (i,k))
-        end
-      end
+      # pick arbitrary hub
+      (h == 0) && (h = i)
+      c = colors[h]
+      # i is the spoke
+      (h == j) && push!(dcolors[c], (i,k))
+      # j is the spoke
+      (h == i) && push!(dcolors[c], (j,k))
     end
   end
 
   return dcolors
 end
-
-"""
-    flag = is_only_color_in_row(H, i, j, n, colors, ci)
-
-This function returns `true` if the column `i` is the only column of color `ci` with a non-zero coefficient in row `j`.
-It returns `false` otherwise.
-"""
-function is_only_color_in_row(H, i, j, n, colors, ci)
-  column = 1
-  flag = true
-  while flag && column ≤ n
-    # We want to check that all columns (excpect the i-th one)
-    # with color ci doesn't have a non-zero coefficient in row j
-    if (column != i) && (colors[column] == ci)
-      k = H.colptr[column]
-      while (k ≤ H.colptr[column+1]-1) && flag
-        row = H.rowval[k]
-        if row == j
-          # We found a coefficient at row j in a column of color ci
-          flag = false
-        else # row != j
-          k += 1
-        end
-      end
-    end
-    column += 1
-  end
-  return flag
-end