SparseReverseADJacobian

Author: amontoison

docs/src/backend.md

@@ -10,7 +10,7 @@ The functions used internally to define the NLPModel API and the possible backen
 | Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend |
 | --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- |
 | `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeReverseADGradient` | -- |
-| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `SparseEnzymeADJacobian` | `SparseADJacobian` |
+| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `SparseEnzymeADJacobian` | `SparseADJacobian`/`SparseReverseADJacobian` |
 | `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | `SparseEnzymeADHessian` | `SparseADHessian`/`SparseReverseADHessian` |
 | `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | `EnzymeReverseADJprod` | -- |
 | `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | `EnzymeReverseADJtprod` | -- |

docs/src/sparse.md

@@ -34,7 +34,7 @@ H = hess(nlp, x)
 
 ## Options for sparsity pattern detection and coloring
 
-The backends available for sparse derivatives (`SparseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, `SparseReverseADHessian`, and `SparseEnzymeADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
+The backends available for sparse derivatives (`SparseADJacobian`, `SparseReverseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, `SparseReverseADHessian`, and `SparseEnzymeADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
 These arguments specify the sparsity pattern detector and the coloring algorithm, respectively.
 
 - A **`detector`** must be of type `ADTypes.AbstractSparsityDetector`.
@@ -53,7 +53,7 @@ set_adbackend!(
 ```
 
 - A **`coloring_algorithm`** must be of type `SparseMatrixColorings.GreedyColoringAlgorithm`.
-  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian` and `SparseEnzymeADHessian`.
+  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseADReverseJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian` and `SparseEnzymeADHessian`.
   These algorithms are provided by the package `SparseMatrixColorings.jl`.
 
 ```@example ex1

src/sparse_jacobian.jl

@@ -11,6 +11,18 @@ struct SparseADJacobian{Tag, R, T, C, S} <: ADBackend
   cz::Vector{ForwardDiff.Dual{Tag, T, 1}}
 end
 
+struct SparseReverseADJacobian{T, C, S, Tape} <: ADBackend
+    nvar::Int
+    ncon::Int
+    rowval::Vector{Int}
+    colptr::Vector{Int}
+    nzval::Vector{T}
+    result_coloring::C
+    compressed_jacobian::S
+    seed::BitVector
+    tape::Tape
+end
+
 function SparseADJacobian(
   nvar,
   f,
@@ -30,6 +42,25 @@ function SparseADJacobian(
   SparseADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, show_time, kwargs...)
 end
 
+function SparseReverseADJacobian(
+    nvar,
+    f,
+    ncon,
+    c!;
+    x0::AbstractVector{T} = rand(nvar),
+    coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
+    detector::AbstractSparsityDetector = TracerSparsityDetector(),
+    show_time::Bool = false,
+    kwargs...,
+) where {T}
+    timer = @elapsed begin
+        output = similar(x0, ncon)
+        J = compute_jacobian_sparsity(c!, output, x0, detector = detector)
+    end
+    show_time && println("  • Sparsity pattern detection of the Jacobian: $timer seconds.")
+    SparseReverseADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, show_time, kwargs...)
+end
+
 function SparseADJacobian(
   nvar,
   f,
@@ -42,7 +73,7 @@ function SparseADJacobian(
   kwargs...,
 ) where {T}
   timer = @elapsed begin
-    # We should support :row and :bidirectional in the future
+    # We should support :bidirectional in the future with DI.jl
     problem = ColoringProblem{:nonsymmetric, :column}()
     result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
 
@@ -75,12 +106,56 @@ function SparseADJacobian(
   )
 end
 
-function get_nln_nnzj(b::SparseADJacobian, nvar, ncon)
+function SparseReverseADJacobian(
+    nvar,
+    f,
+    ncon,
+    c!,
+    J::SparseMatrixCSC{Bool, Int};
+    x0::AbstractVector{T} = rand(nvar),
+    coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
+    show_time::Bool = false,
+    kwargs...,
+) where {T}
+    timer = @elapsed begin
+        # We should support :bidirectional in the future with DI.jl
+        problem = ColoringProblem{:nonsymmetric, :row}()
+        result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
+
+        rowval = J.rowval
+        colptr = J.colptr
+        nzval = T.(J.nzval)
+        compressed_jacobian = zeros(T, nvar)
+        seed = BitVector(undef, ncon)
+    end
+    show_time && println("  • Coloring of the sparse Jacobian: $timer seconds.")
+
+    timer = @elapsed begin
+        trace_input = copy(x0)
+        tape = ReverseDiff.GradientTape(x -> c!(similar(x0, ncon), x), trace_input)
+        ReverseDiff.compile!(tape)
+    end
+    show_time && println("  • Compilation of the ReverseDiff tape: $timer seconds.")
+
+    return SparseReverseADJacobian(
+        nvar,
+        ncon,
+        rowval,
+        colptr,
+        nzval,
+        result_coloring,
+        compressed_jacobian,
+        seed,
+        tape,
+    )
+end
+
+function get_nln_nnzj(b::Union{SparseADJacobian, SparseReverseADJacobian}, nvar, ncon)
   length(b.rowval)
 end
 
 function NLPModels.jac_structure!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseReverseADJacobian},
   nlp::ADModel,
   rows::AbstractVector{<:Integer},
   cols::AbstractVector{<:Integer},
@@ -122,8 +197,40 @@ function sparse_jac_coord!(
   return vals
 end
 
+function sparse_jac_coord!(
+    ℓ!::Function,
+    b::SparseReverseADJacobian{T},
+    x::AbstractVector,
+    vals::AbstractVector,
+) where {T}
+    # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
+    A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval)
+
+    groups = row_groups(b.result_coloring)
+    for (irow, rows) in enumerate(groups)
+    # Update the seed
+        b.seed .= false
+        for row in rows
+            b.seed[row] = true
+        end
+
+        ReverseDiff.gradient!(
+            b.compressed_jacobian,
+            b.tape,
+            x;
+            output_seed = b.seed,
+        )
+
+        # Update the rows of the Jacobian that have the color `irow`
+        decompress_single_color!(A, b.compressed_jacobian, irow, b.result_coloring)
+    end
+
+    vals .= b.nzval
+    return vals
+end
+
 function NLPModels.jac_coord!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseReverseADJacobian},
   nlp::ADModel,
   x::AbstractVector,
   vals::AbstractVector,
@@ -133,7 +240,7 @@ function NLPModels.jac_coord!(
 end
 
 function NLPModels.jac_structure_residual!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseReverseADJacobian},
   nls::AbstractADNLSModel,
   rows::AbstractVector{<:Integer},
   cols::AbstractVector{<:Integer},
@@ -148,7 +255,7 @@ function NLPModels.jac_structure_residual!(
 end
 
 function NLPModels.jac_coord_residual!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseReverseADJacobian},
   nls::AbstractADNLSModel,
   x::AbstractVector,
   vals::AbstractVector,

test/runtests.jl

@@ -29,7 +29,9 @@ include("sparse_hessian.jl")
 include("sparse_hessian_nls.jl")
 
 list_sparse_jac_backend =
-  ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()))
+  ((ADNLPModels.SparseADJacobian, Dict()),
+   (ADNLPModels.SparseReverseADJacobian, Dict()),
+   (ADNLPModels.ForwardDiffADJacobian, Dict()))
 
 @testset "Sparse Jacobian" begin
   for (backend, kw) in list_sparse_jac_backend