From b98bda80c1d150a8a37eb1be054616b714296374 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 17 Jul 2024 20:35:51 -0400 Subject: [PATCH] Update the decompression of sparse Hessian --- src/sparse_hessian.jl | 8 ++--- src/sparsity_pattern.jl | 70 ++++++++++------------------------------- 2 files changed, 19 insertions(+), 59 deletions(-) diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl index 267aedd0..f4cd883b 100644 --- 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 index 5edde823..4c345fc3 100644 --- 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