From eb558f496d5dc8a41b781f786ed0493644c9418e Mon Sep 17 00:00:00 2001
From: Alexis Montoison <alexis.montoison@polymtl.ca>
Date: Mon, 24 Jun 2024 18:33:16 -0400
Subject: [PATCH] Use UnitRange{Int} in dcolors to use less storage

---
 src/sparse_hessian.jl  | 40 ++++++++++++++++++++--------------------
 src/sparse_jacobian.jl | 20 ++++++++++----------
 2 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl
index 4b9486bb..f9412a24 100644
--- 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
index 43d6320a..f5091479 100644
--- 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