Optimize sparse AD

JuliaSmoothOptimizers · Jun 24, 2024 · 12ac3bb · 12ac3bb
1 parent cbeff59
commit 12ac3bb
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 23 deletions.
diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl
@@ -4,6 +4,7 @@ struct SparseADHessian{Tag, GT, S, T} <: ADNLPModels.ADBackend
   colptr::Vector{Int}
   colors::Vector{Int}
   ncolors::Int
+  dcolors::Dict{Int, Vector{Int}}
   res::S
   lz::Vector{ForwardDiff.Dual{Tag, T, 1}}
   glz::Vector{ForwardDiff.Dual{Tag, T, 1}}
@@ -37,6 +38,14 @@ function SparseADHessian(
   rowval = trilH.rowval
   colptr = trilH.colptr
 
+  # The indices of the nonzero elements in `vals` that will be processed by color `c` are stored in `dcolors[c]`.
+  dcolors = Dict{Int, Vector{Int}}(i => Int[] for i=1:ncolors)
+  for (i, color) in enumerate(colors)
+    for k = colptr[i]:(colptr[i + 1] - 1)
+      push!(dcolors[color], k)
+    end
+  end
+
   # prepare directional derivatives
   res = similar(x0)
 
@@ -65,7 +74,7 @@ function SparseADHessian(
   Hvp = fill!(S(undef, ntotal), 0)
   y = fill!(S(undef, ncon), 0)
 
-  return SparseADHessian(d, rowval, colptr, colors, ncolors, res, lz, glz, sol, longv, Hvp, ∇φ!, y)
+  return SparseADHessian(d, rowval, colptr, colors, ncolors, dcolors, res, lz, glz, sol, longv, Hvp, ∇φ!, y)
 end
 
 struct SparseReverseADHessian{T, S, Tagf, F, Tagψ, P} <: ADNLPModels.ADBackend
@@ -74,6 +83,7 @@ struct SparseReverseADHessian{T, S, Tagf, F, Tagψ, P} <: ADNLPModels.ADBackend
   colptr::Vector{Int}
   colors::Vector{Int}
   ncolors::Int
+  dcolors::Dict{Int, Vector{Int}}
   res::S
   z::Vector{ForwardDiff.Dual{Tagf, T, 1}}
   gz::Vector{ForwardDiff.Dual{Tagf, T, 1}}
@@ -109,6 +119,14 @@ function SparseReverseADHessian(
   rowval = trilH.rowval
   colptr = trilH.colptr
 
+  # The indices of the nonzero elements in `vals` that will be processed by color `c` are stored in `dcolors[c]`.
+  dcolors = Dict{Int, Vector{Int}}(i => Int[] for i=1:ncolors)
+  for (i, color) in enumerate(colors)
+    for k = colptr[i]:(colptr[i + 1] - 1)
+      push!(dcolors[color], k)
+    end
+  end
+
   # prepare directional derivatives
   res = similar(x0)
 
@@ -147,6 +165,7 @@ function SparseReverseADHessian(
     colptr,
     colors,
     ncolors,
+    dcolors,
     res,
     z,
     gz,
@@ -213,13 +232,12 @@ function sparse_hess_coord!(
     b.∇φ!(b.glz, b.lz)
     ForwardDiff.extract_derivative!(Tag, b.Hvp, b.glz)
     b.res .= view(b.Hvp, (ncon + 1):(ncon + nvar))
-    for j = 1:nvar
-      if b.colors[j] == icol
-        for k = b.colptr[j]:(b.colptr[j + 1] - 1)
-          i = b.rowval[k]
-          vals[k] = b.res[i]
-        end
-      end
+
+    # Update the vector vals
+    index_vals = b.dcolors[icol]
+    for k in index_vals
+      row = b.rowval[k]
+      vals[k] = b.res[row]
     end
   end
 
@@ -251,13 +269,11 @@ function sparse_hess_coord!(
     ForwardDiff.extract_derivative!(Tagψ, b.Hv_temp, b.gzψ)
     b.res .+= b.Hv_temp
 
-    for j = 1:nvar
-      if b.colors[j] == icol
-        for k = b.colptr[j]:(b.colptr[j + 1] - 1)
-          i = b.rowval[k]
-          vals[k] = b.res[i]
-        end
-      end
+    # Update the vector vals
+    index_vals = b.dcolors[icol]
+    for k in index_vals
+      row = b.rowval[k]
+      vals[k] = b.res[row]
     end
   end
 

diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl
@@ -4,6 +4,7 @@ struct SparseADJacobian{T, Tag, S} <: ADBackend
   colptr::Vector{Int}
   colors::Vector{Int}
   ncolors::Int
+  dcolors::Dict{Int, Vector{Int}}
   z::Vector{ForwardDiff.Dual{Tag, T, 1}}
   cz::Vector{ForwardDiff.Dual{Tag, T, 1}}
   res::S
@@ -31,13 +32,21 @@ function SparseADJacobian(
   rowval = J.rowval
   colptr = J.colptr
 
+  # The indices of the nonzero elements in `vals` that will be processed by color `c` are stored in `dcolors[c]`.
+  dcolors = Dict{Int, Vector{Int}}(i => Int[] for i=1:ncolors)
+  for (i, color) in enumerate(colors)
+    for k = colptr[i]:(colptr[i + 1] - 1)
+      push!(dcolors[color], k)
+    end
+  end
+
   tag = ForwardDiff.Tag{typeof(c!), T}
 
   z = Vector{ForwardDiff.Dual{tag, T, 1}}(undef, nvar)
   cz = similar(z, ncon)
   res = similar(x0, ncon)
 
-  SparseADJacobian(d, rowval, colptr, colors, ncolors, z, cz, res)
+  SparseADJacobian(d, rowval, colptr, colors, ncolors, dcolors, z, cz, res)
 end
 
 function get_nln_nnzj(b::SparseADJacobian, nvar, ncon)
@@ -71,13 +80,12 @@ function sparse_jac_coord!(
     map!(ForwardDiff.Dual{Tag}, b.z, x, b.d) # x + ε * v
     ℓ!(b.cz, b.z) # c!(cz, x + ε * v)
     ForwardDiff.extract_derivative!(Tag, b.res, b.cz) # ∇c!(cx, x)ᵀv
-    for j = 1:nvar
-      if b.colors[j] == icol
-        for k = b.colptr[j]:(b.colptr[j + 1] - 1)
-          i = b.rowval[k]
-          vals[k] = b.res[i]
-        end
-      end
+
+    # Update the vector vals
+    index_vals = b.dcolors[icol]
+    for k in index_vals
+      row = b.rowval[k]
+      vals[k] = b.res[row]
     end
   end
   return vals