More tests for sparse Jacobians and Hessians (#189)

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl

@@ -148,7 +148,7 @@ end
 function DI.prepare_jacobian(f, backend::AnyAutoSymbolics, x)
     x_var = variables(:x, axes(x)...)
     jac_var = if issparse(backend)
-        sparsejacobian(f(x_var), x_var)
+        sparsejacobian(vec(f(x_var)), vec(x_var))
     else
         jacobian(f(x_var), x_var)
     end
@@ -192,7 +192,7 @@ end
 
 function DI.prepare_hessian(f, backend::AnyAutoSymbolics, x)
     x_var = variables(:x, axes(x)...)
-    # Symbolic.gradient only accepts vectors
+    # Symbolic.hessian only accepts vectors
     hess_var = if issparse(backend)
         sparsehessian(f(x_var), vec(x_var))
     else

DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/twoarg.jl

@@ -124,7 +124,7 @@ function DI.prepare_jacobian(f!, y, backend::AnyAutoSymbolics, x)
     y_var = variables(:y, axes(y)...)
     f!(y_var, x_var)
     jac_var = if issparse(backend)
-        sparsejacobian(y_var, x_var)
+        sparsejacobian(vec(y_var), vec(x_var))
     else
         jacobian(y_var, x_var)
     end

DifferentiationInterface/test/sparsity.jl

@@ -20,15 +20,15 @@ end
 
 test_differentiation(
     sparse_backends,
-    sparse_scenarios(rand(5));
+    sparse_scenarios();
     sparsity=true,
     second_order=false,
     logging=get(ENV, "CI", "false") == "false",
 )
 
 test_differentiation(
     sparse_second_order_backends,
-    sparse_scenarios(rand(5));
+    sparse_scenarios();
     sparsity=true,
     first_order=false,
     logging=get(ENV, "CI", "false") == "false",

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -21,7 +21,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ADTypes = "0.2.7"
 Chairmarks = "1.2.1"
 ComponentArrays = "0.15"
-DifferentiationInterface = "0.2.0"
+DifferentiationInterface = "0.2.1"
 DocStringExtensions = "0.9"
 JET = "0.8"
 JLArrays = "0.1"

DifferentiationInterfaceTest/src/scenarios/sparse.jl

@@ -1,52 +1,239 @@
 ## Vector to vector
 
-abs2diff(x) = abs2.(diff(x))
+diffsquare(x::AbstractVector)::AbstractVector = diff(x) .^ 2
+diffcube(x::AbstractVector)::AbstractVector = diff(x) .^ 3
 
-function abs2diff!(y, x)
-    y .= abs2.(diff(x))
+function diffsquare!(y::AbstractVector, x::AbstractVector)
+    x1 = @view x[1:(end - 1)]
+    x2 = @view x[2:end]
+    y .= x2 .- x1
+    y .^= 2
     return nothing
 end
 
-function abs2diff_jacobian(x)
+function diffcube!(y::AbstractVector, x::AbstractVector)
+    x1 = @view x[1:(end - 1)]
+    x2 = @view x[2:end]
+    y .= x2 .- x1
+    y .^= 3
+    return nothing
+end
+
+function diffsquare_jacobian(x)
     n = length(x)
     return spdiagm(n - 1, n, 0 => -2 * diff(x), 1 => 2 * diff(x))
 end
 
-## Vector to scalar
-
-sumabs2diff(x) = sum(abs2diff(x))
+function diffcube_jacobian(x)
+    n = length(x)
+    return spdiagm(n - 1, n, 0 => -3 * diff(x) .^ 2, 1 => 3 * diff(x) .^ 2)
+end
 
-function sumabs2diff_hessian(x)
-    T = eltype(x)
+function sparse_vec_to_vec_scenarios(x::AbstractVector)
     n = length(x)
-    return spdiagm(
-        0 => vcat(2 * one(T), fill(4 * one(T), n - 2), 2 * one(T)),
-        1 => fill(-2 * one(T), n - 1),
-        -1 => fill(-2 * one(T), n - 1),
-    )
+    scens = AbstractScenario[]
+    for op in (:outofplace, :inplace)
+        append!(
+            scens,
+            [
+                JacobianScenario(diffsquare; x=x, ref=diffsquare_jacobian, operator=op),
+                JacobianScenario(
+                    diffsquare!;
+                    x=x,
+                    y=similar(x, n - 1),
+                    ref=diffsquare_jacobian,
+                    operator=op,
+                ),
+            ],
+        )
+    end
+    return scens
 end
 
-## Gather
+## Matrix to vector
 
-"""
-    sparse_scenarios()
+function diffsquarecube_matvec(x::AbstractMatrix)::AbstractVector
+    return vcat(diffsquare(vec(x)), diffcube(vec(x)))
+end
 
-Create a vector of [`AbstractScenario`](@ref)s with sparse array types, focused on sparse Jacobians and Hessians.
-"""
-function sparse_scenarios(x::AbstractVector)
+function diffsquarecube_matvec!(y::AbstractVector, x::AbstractMatrix)
+    m, n = size(x)
+    diffsquare!(view(y, 1:(m * n - 1)), vec(x))
+    diffcube!(view(y, (m * n):(2(m * n) - 2)), vec(x))
+    return nothing
+end
+
+function diffsquarecube_matvec_jacobian(x)
+    return vcat(diffsquare_jacobian(vec(x)), diffcube_jacobian(vec(x)))
+end
+
+function sparse_mat_to_vec_scenarios(x::AbstractMatrix)
+    m, n = size(x)
+    scens = AbstractScenario[]
+    for op in (:outofplace, :inplace)
+        append!(
+            scens,
+            [
+                JacobianScenario(
+                    diffsquarecube_matvec;
+                    x=x,
+                    ref=diffsquarecube_matvec_jacobian,
+                    operator=op,
+                ),
+                JacobianScenario(
+                    diffsquarecube_matvec!;
+                    x=x,
+                    y=similar(x, 2(m * n) - 2),
+                    ref=diffsquarecube_matvec_jacobian,
+                    operator=op,
+                ),
+            ],
+        )
+    end
+    return scens
+end
+
+## Vector to matrix
+
+diffsquarecube_vecmat(x::AbstractVector)::AbstractMatrix = hcat(diffsquare(x), diffcube(x))
+
+function diffsquarecube_vecmat!(y::AbstractMatrix, x::AbstractVector)
+    diffsquare!(view(y, :, 1), x)
+    diffcube!(view(y, :, 2), x)
+    return nothing
+end
+
+function diffsquarecube_vecmat_jacobian(x::AbstractVector)
+    return vcat(diffsquare_jacobian(x), diffcube_jacobian(x))
+end
+
+function sparse_vec_to_mat_scenarios(x::AbstractVector)
     n = length(x)
     scens = AbstractScenario[]
     for op in (:outofplace, :inplace)
         append!(
             scens,
             [
-                JacobianScenario(abs2diff; x=x, ref=abs2diff_jacobian, operator=op),
                 JacobianScenario(
-                    abs2diff!; x=x, y=similar(x, n - 1), ref=abs2diff_jacobian, operator=op
+                    diffsquarecube_vecmat;
+                    x=x,
+                    ref=diffsquarecube_vecmat_jacobian,
+                    operator=op,
+                ),
+                JacobianScenario(
+                    diffsquarecube_vecmat!;
+                    x=x,
+                    y=similar(x, n - 1, 2),
+                    ref=diffsquarecube_vecmat_jacobian,
+                    operator=op,
                 ),
-                HessianScenario(sumabs2diff; x=x, ref=sumabs2diff_hessian, operator=op),
             ],
         )
     end
     return scens
 end
+
+## Matrix to matrix
+
+function diffsquarecube_matmat(x::AbstractMatrix)::AbstractMatrix
+    return hcat(diffsquare(vec(x)), diffcube(vec(x)))
+end
+
+function diffsquarecube_matmat!(y::AbstractMatrix, x::AbstractMatrix)
+    diffsquare!(view(y, :, 1), vec(x))
+    diffcube!(view(y, :, 2), vec(x))
+    return nothing
+end
+
+function diffsquarecube_matmat_jacobian(x::AbstractMatrix)
+    return vcat(diffsquare_jacobian(vec(x)), diffcube_jacobian(vec(x)))
+end
+
+function sparse_mat_to_mat_scenarios(x::AbstractMatrix)
+    m, n = size(x)
+    scens = AbstractScenario[]
+    for op in (:outofplace, :inplace)
+        append!(
+            scens,
+            [
+                JacobianScenario(
+                    diffsquarecube_matmat;
+                    x=x,
+                    ref=diffsquarecube_matmat_jacobian,
+                    operator=op,
+                ),
+                JacobianScenario(
+                    diffsquarecube_matmat!;
+                    x=x,
+                    y=similar(x, m * n - 1, 2),
+                    ref=diffsquarecube_matmat_jacobian,
+                    operator=op,
+                ),
+            ],
+        )
+    end
+    return scens
+end
+
+## Vector to scalar
+
+sumdiffcube(x::AbstractVector)::Number = sum(diffcube(x))
+
+function sumdiffcube_hessian(x)
+    T = eltype(x)
+    d = 6 * diff(x)
+    return spdiagm(0 => vcat(d, zero(T)) + vcat(zero(T), d), 1 => -d, -1 => -d)
+end
+
+function sparse_vec_to_num_scenarios(x::AbstractVector)
+    scens = AbstractScenario[]
+    for op in (:outofplace, :inplace)
+        append!(
+            scens, [HessianScenario(sumdiffcube; x=x, ref=sumdiffcube_hessian, operator=op)]
+        )
+    end
+    return scens
+end
+
+## Matrix to scalar
+
+sumdiffcube_mat(x::AbstractMatrix)::Number = sum(diffcube(vec(x)))
+
+function sumdiffcube_mat_hessian(x::AbstractMatrix)
+    T = eltype(x)
+    d = 6 * diff(vec(x))
+    return spdiagm(0 => vcat(d, zero(T)) + vcat(zero(T), d), 1 => -d, -1 => -d)
+end
+
+function sparse_mat_to_num_scenarios(x::AbstractMatrix)
+    scens = AbstractScenario[]
+    for op in (:outofplace, :inplace)
+        append!(
+            scens,
+            [
+                HessianScenario(
+                    sumdiffcube_mat; x=x, ref=sumdiffcube_mat_hessian, operator=op
+                ),
+            ],
+        )
+    end
+    return scens
+end
+
+## Gather
+
+"""
+    sparse_scenarios()
+
+Create a vector of [`AbstractScenario`](@ref)s with sparse array types, focused on sparse Jacobians and Hessians.
+"""
+function sparse_scenarios()
+    return vcat(
+        sparse_vec_to_vec_scenarios(randn(6)),
+        sparse_vec_to_mat_scenarios(randn(6)),
+        sparse_mat_to_vec_scenarios(randn(2, 3)),
+        sparse_mat_to_mat_scenarios(randn(2, 3)),
+        sparse_vec_to_num_scenarios(randn(6)),
+        sparse_mat_to_num_scenarios(randn(2, 3)),
+    )
+end