Use UnitRange{Int} in dcolors to use less storage

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

Hide file tree

Showing 2 changed files with 30 additions and 30 deletions.
diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl
@@ -4,7 +4,7 @@ struct SparseADHessian{Tag, GT, S, T} <: ADNLPModels.ADBackend
   colptr::Vector{Int}
   colors::Vector{Int}
   ncolors::Int
-  dcolors::Dict{Int, Vector{Int}}
+  dcolors::Dict{Int, Vector{UnitRange{Int}}}
   res::S
   lz::Vector{ForwardDiff.Dual{Tag, T, 1}}
   glz::Vector{ForwardDiff.Dual{Tag, T, 1}}
@@ -39,11 +39,10 @@ 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 = Dict{Int, Vector{Int}}(i => Int[] for i=1:ncolors)
+  dcolors = Dict(i => UnitRange{Int}[] for i=1:ncolors)
   for (i, color) in enumerate(colors)
-    for k = colptr[i]:(colptr[i + 1] - 1)
-      push!(dcolors[color], k)
-    end
+    range_vals = colptr[i]:(colptr[i + 1] - 1)
+    push!(dcolors[color], range_vals)
   end
 
   # prepare directional derivatives
@@ -83,7 +82,7 @@ struct SparseReverseADHessian{T, S, Tagf, F, Tagψ, P} <: ADNLPModels.ADBackend
   colptr::Vector{Int}
   colors::Vector{Int}
   ncolors::Int
-  dcolors::Dict{Int, Vector{Int}}
+  dcolors::Dict{Int, Vector{UnitRange{Int}}}
   res::S
   z::Vector{ForwardDiff.Dual{Tagf, T, 1}}
   gz::Vector{ForwardDiff.Dual{Tagf, T, 1}}
@@ -120,11 +119,10 @@ 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 = Dict{Int, Vector{Int}}(i => Int[] for i=1:ncolors)
+  dcolors = Dict(i => UnitRange{Int}[] for i=1:ncolors)
   for (i, color) in enumerate(colors)
-    for k = colptr[i]:(colptr[i + 1] - 1)
-      push!(dcolors[color], k)
-    end
+    range_vals = colptr[i]:(colptr[i + 1] - 1)
+    push!(dcolors[color], range_vals)
   end
 
   # prepare directional derivatives
@@ -233,11 +231,12 @@ function sparse_hess_coord!(
     ForwardDiff.extract_derivative!(Tag, b.Hvp, b.glz)
     b.res .= view(b.Hvp, (ncon + 1):(ncon + nvar))
 
-    # Update the vector vals
-    index_vals = b.dcolors[icol]
-    for k in index_vals
-      row = b.rowval[k]
-      vals[k] = b.res[row]
+    # Store in `vals` the nonzeros of each column of the Hessian computed with color `icol`
+    for range_vals in b.dcolors[icol]
+      for k in range_vals
+        row = b.rowval[k]
+        vals[k] = b.res[row]
+      end
     end
   end
 
@@ -269,11 +268,12 @@ function sparse_hess_coord!(
     ForwardDiff.extract_derivative!(Tagψ, b.Hv_temp, b.gzψ)
     b.res .+= b.Hv_temp
 
-    # Update the vector vals
-    index_vals = b.dcolors[icol]
-    for k in index_vals
-      row = b.rowval[k]
-      vals[k] = b.res[row]
+    # Store in `vals` the nonzeros of each column of the Hessian computed with color `icol`
+    for range_vals in b.dcolors[icol]
+      for k in range_vals
+        row = b.rowval[k]
+        vals[k] = b.res[row]
+      end
     end
   end
 

diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl
@@ -4,7 +4,7 @@ struct SparseADJacobian{T, Tag, S} <: ADBackend
   colptr::Vector{Int}
   colors::Vector{Int}
   ncolors::Int
-  dcolors::Dict{Int, Vector{Int}}
+  dcolors::Dict{Int, Vector{UnitRange{Int}}}
   z::Vector{ForwardDiff.Dual{Tag, T, 1}}
   cz::Vector{ForwardDiff.Dual{Tag, T, 1}}
   res::S
@@ -33,11 +33,10 @@ function SparseADJacobian(
   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)
+  dcolors = Dict(i => UnitRange{Int}[] for i=1:ncolors)
   for (i, color) in enumerate(colors)
-    for k = colptr[i]:(colptr[i + 1] - 1)
-      push!(dcolors[color], k)
-    end
+    range_vals = colptr[i]:(colptr[i + 1] - 1)
+    push!(dcolors[color], range_vals)
   end
 
   tag = ForwardDiff.Tag{typeof(c!), T}
@@ -81,11 +80,12 @@ function sparse_jac_coord!(
     ℓ!(b.cz, b.z) # c!(cz, x + ε * v)
     ForwardDiff.extract_derivative!(Tag, b.res, b.cz) # ∇c!(cx, x)ᵀv
 
-    # Update the vector vals
-    index_vals = b.dcolors[icol]
-    for k in index_vals
-      row = b.rowval[k]
-      vals[k] = b.res[row]
+    # Store in `vals` the nonzeros of each column of the Jacobian computed with color `icol`
+    for range_vals in b.dcolors[icol]
+      for k in range_vals
+        row = b.rowval[k]
+        vals[k] = b.res[row]
+      end
     end
   end
   return vals