Add more sparse AD test scenarios (#578)

* Add more sparse AD test scenarios

* No test with fixed matrices on Zygote

* Increase max bands beyond ForwardDiff max chunk

* More bandwidths, fix Enzyme

* Speed up symbolic

Author: gdalle

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -90,5 +90,8 @@ test_differentiation(
 ## Sparse
 
 test_differentiation(
-    MyAutoSparse.(AutoEnzyme()), sparse_scenarios(); sparsity=true, logging=LOGGING
+    MyAutoSparse.(AutoEnzyme(; function_annotation=Enzyme.Const)),
+    sparse_scenarios();
+    sparsity=true,
+    logging=LOGGING,
 );

DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl

@@ -16,7 +16,7 @@ test_differentiation(AutoFastDifferentiation(); logging=LOGGING);
 
 test_differentiation(
     AutoSparse(AutoFastDifferentiation()),
-    sparse_scenarios();
+    sparse_scenarios(; band_sizes=0:-1);
     sparsity=true,
     logging=LOGGING,
 );

DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl

@@ -15,5 +15,8 @@ end
 test_differentiation(AutoSymbolics(); logging=LOGGING);
 
 test_differentiation(
-    AutoSparse(AutoSymbolics()), sparse_scenarios(); sparsity=true, logging=LOGGING
+    AutoSparse(AutoSymbolics()),
+    sparse_scenarios(; band_sizes=0:-1);
+    sparsity=true,
+    logging=LOGGING,
 );

DifferentiationInterface/test/Back/Zygote/test.jl

@@ -41,7 +41,7 @@ test_differentiation(
 
 test_differentiation(
     MyAutoSparse.(vcat(backends, second_order_backends)),
-    sparse_scenarios();
+    sparse_scenarios(; band_sizes=0:-1);
     sparsity=true,
     logging=LOGGING,
 )

DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl

@@ -56,7 +56,8 @@ using JET: @test_opt
 using LinearAlgebra: Adjoint, Diagonal, Transpose, dot, parent
 using ProgressMeter: ProgressUnknown, next!
 using Random: AbstractRNG, default_rng, rand!
-using SparseArrays: SparseArrays, AbstractSparseMatrix, SparseMatrixCSC, nnz, spdiagm
+using SparseArrays:
+    SparseArrays, AbstractSparseMatrix, SparseMatrixCSC, nnz, sparse, spdiagm
 using Test: @testset, @test
 
 """

DifferentiationInterfaceTest/src/scenarios/sparse.jl

@@ -178,7 +178,6 @@ end
 
 function sparse_vec_to_num_scenarios(x::AbstractVector)
     f = sumdiffcube
-    y = f(x)
     grad = sumdiffcube_gradient(x)
     hess = sumdiffcube_hessian(x)
 
@@ -203,7 +202,6 @@ end
 
 function sparse_mat_to_num_scenarios(x::AbstractMatrix)
     f = sumdiffcube_mat
-    y = f(x)
     grad = sumdiffcube_mat_gradient(x)
     hess = sumdiffcube_mat_hessian(x)
 
@@ -214,14 +212,120 @@ function sparse_mat_to_num_scenarios(x::AbstractMatrix)
     return scens
 end
 
+## Various matrices
+
+function banded_matrix(m, n, b)
+    pairs = [k => rand(min(m, n) - k) for k in 0:b]
+    return spdiagm(m, n, pairs...)
+end
+
+### Linear map
+
+struct SquareLinearMap{M<:AbstractMatrix}
+    A::M
+end
+
+function Base.show(io::IO, s::SquareLinearMap{M}) where {M}
+    return print(io, "SquareLinearMap{$M - $(size(s.A)) with $(mynnz(s.A)) nonzeros}")
+end
+
+function (s::SquareLinearMap)(x::AbstractArray)
+    return s.A * abs2.(vec(x))
+end
+
+function (s::SquareLinearMap)(y::AbstractArray, x::AbstractArray)
+    vec(y) .= s.A * abs2.(vec(x))
+    return nothing
+end
+
+function squarelinearmap_jacobian(x::AbstractArray, A::AbstractMatrix)
+    return 2 .* A .* transpose(vec(x))
+end
+
+function squarelinearmap_scenarios(x::AbstractVector, band_sizes)
+    n = length(x)
+    scens = Scenario[]
+    for A in vcat(banded_matrix.(2n, n, band_sizes), banded_matrix.(n ÷ 2, n, band_sizes))
+        f = SquareLinearMap(A)
+        f! = f
+        y = f(x)
+        jac = sparse(squarelinearmap_jacobian(x, A))
+        for pl_op in (:out, :in)
+            append!(
+                scens,
+                [
+                    Scenario{:jacobian,pl_op}(f, x; res1=jac),
+                    Scenario{:jacobian,pl_op}(f!, y, x; res1=jac),
+                ],
+            )
+        end
+    end
+    return scens
+end
+
+### Quadratic form
+
+struct SquareQuadraticForm{M<:AbstractMatrix}
+    A::M
+end
+
+function Base.show(io::IO, s::SquareQuadraticForm{M}) where {M}
+    return print(io, "SquareQuadraticForm{$M - $(size(s.A)) with $(mynnz(s.A)) nonzeros}")
+end
+
+function (s::SquareQuadraticForm)(x::AbstractArray)
+    v = abs2.(vec(x))
+    return dot(v, s.A, v)
+end
+
+function squarequadraticform_gradient(x::AbstractArray, A::AbstractMatrix)
+    g = similar(x)
+    for i in eachindex(g)
+        g[i] =
+            4 * A[i, i] * x[i]^3 +
+            2 * sum((A[i, j] + A[j, i]) * x[i] * x[j]^2 for j in eachindex(g) if j != i)
+    end
+    return g
+end
+
+function squarequadraticform_hessian(x::AbstractArray, A::AbstractMatrix)
+    H = similar(x, length(x), length(x))
+    for i in axes(H, 1), j in axes(H, 2)
+        if i == j
+            H[i, i] =
+                12 * A[i, i] * x[i]^2 +
+                2 * sum((A[i, j2] + A[j2, i]) * x[j2]^2 for j2 in axes(H, 2) if j2 != i)
+        else
+            H[i, j] = 4 * (A[i, j] + A[j, i]) * x[i] * x[j]
+        end
+    end
+    return H
+end
+
+function squarequadraticform_scenarios(x::AbstractVector, band_sizes)
+    n = length(x)
+    scens = Scenario[]
+    for A in banded_matrix.(n, n, band_sizes)
+        f = SquareQuadraticForm(A)
+        grad = squarequadraticform_gradient(x, A)
+        hess = sparse(squarequadraticform_hessian(x, A))
+        for pl_op in (:out, :in)
+            push!(scens, Scenario{:hessian,pl_op}(f, x; res1=grad, res2=hess))
+        end
+    end
+    return scens
+end
+
 ## Gather
 
 """
     sparse_scenarios(rng=Random.default_rng())
 
 Create a vector of [`Scenario`](@ref)s with sparse array types, focused on sparse Jacobians and Hessians.
 """
-function sparse_scenarios(rng::AbstractRNG=default_rng(); include_constantified=false)
+function sparse_scenarios(
+    rng::AbstractRNG=default_rng(); band_sizes=0:4:36, include_constantified=false
+)
     scens = vcat(
         sparse_vec_to_vec_scenarios(rand(rng, 6)),
         sparse_vec_to_mat_scenarios(rand(rng, 6)),
@@ -230,6 +334,10 @@ function sparse_scenarios(rng::AbstractRNG=default_rng(); include_constantified=
         sparse_vec_to_num_scenarios(rand(rng, 6)),
         sparse_mat_to_num_scenarios(rand(rng, 2, 3)),
     )
+    if !isempty(band_sizes)
+        append!(scens, squarelinearmap_scenarios(rand(rng, 100), band_sizes))
+        append!(scens, squarequadraticform_scenarios(rand(rng, 100), band_sizes))
+    end
     include_constantified && append!(scens, constantify(scens))
     return scens
 end

DifferentiationInterfaceTest/test/standard.jl

@@ -23,6 +23,7 @@ sparse_backend = AutoSparse(
     sparsity_detector=TracerSparsityDetector(),
     coloring_algorithm=GreedyColoringAlgorithm(),
 )
+
 test_differentiation(
     sparse_backend,
     sparse_scenarios(; include_constantified=true);