Add KnownJacobianSparsityDetector and KnownHessianSparsityDetector (

#81) * Add `KnownSparsityDetector` * Split Jacobians and Hessians * Add tests * More tests * Update docs * Throw `ArgumentError` * Add more testsets
SciML · Sep 4, 2024 · f206682 · f206682
1 parent c27ebd3
commit f206682
Show file tree

Hide file tree

Showing 4 changed files with 174 additions and 18 deletions.
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -69,6 +69,8 @@ ADTypes.AbstractSparsityDetector
 ADTypes.jacobian_sparsity
 ADTypes.hessian_sparsity
 ADTypes.NoSparsityDetector
+ADTypes.KnownJacobianSparsityDetector
+ADTypes.KnownHessianSparsityDetector
 ```
 
 ### Coloring algorithm

diff --git a/src/sparse.jl b/src/sparse.jl
@@ -42,6 +42,67 @@ jacobian_sparsity(f, x, ::NoSparsityDetector) = trues(length(f(x)), length(x))
 jacobian_sparsity(f!, y, x, ::NoSparsityDetector) = trues(length(y), length(x))
 hessian_sparsity(f, x, ::NoSparsityDetector) = trues(length(x), length(x))
 
+"""
+    KnownJacobianSparsityDetector(jacobian_sparsity::AbstractMatrix) <: AbstractSparsityDetector
+
+Trivial sparsity detector used to return a known Jacobian sparsity pattern.
+
+# See also
+
+  - [`AbstractSparsityDetector`](@ref)
+  - [`KnownHessianSparsityDetector`](@ref)
+"""
+struct KnownJacobianSparsityDetector{J <: AbstractMatrix} <: AbstractSparsityDetector
+    jacobian_sparsity::J
+end
+
+function jacobian_sparsity(f, x, sd::KnownJacobianSparsityDetector)
+    sz = size(sd.jacobian_sparsity)
+    sz_expected = (length(f(x)), length(x))
+    sz != sz_expected &&
+        throw(DimensionMismatch("Jacobian size $sz of KnownJacobianSparsityDetector doesn't match expected size $sz_expected."))
+    return sd.jacobian_sparsity
+end
+function jacobian_sparsity(f!, y, x, sd::KnownJacobianSparsityDetector)
+    sz = size(sd.jacobian_sparsity)
+    sz_expected = (length(y), length(x))
+    sz != sz_expected &&
+        throw(DimensionMismatch("Jacobian size $sz of KnownJacobianSparsityDetector doesn't match expected size $sz_expected."))
+    return sd.jacobian_sparsity
+end
+function hessian_sparsity(f, x, sd::KnownJacobianSparsityDetector)
+    throw(ArgumentError("KnownJacobianSparsityDetector can't be used to compute Hessian sparsity."))
+end
+
+"""
+    KnownHessianSparsityDetector(hessian_sparsity::AbstractMatrix) <: AbstractSparsityDetector
+
+Trivial sparsity detector used to return a known Hessian sparsity pattern.
+
+# See also
+
+  - [`AbstractSparsityDetector`](@ref)
+  - [`KnownJacobianSparsityDetector`](@ref)
+"""
+struct KnownHessianSparsityDetector{H <: AbstractMatrix} <: AbstractSparsityDetector
+    hessian_sparsity::H
+end
+
+function hessian_sparsity(f, x, sd::KnownHessianSparsityDetector)
+    sz = size(sd.hessian_sparsity)
+    sz_expected = (length(x), length(x))
+    sz != sz_expected &&
+        throw(DimensionMismatch("Hessian size $sz of KnownHessianSparsityDetector doesn't match expected size $sz_expected."))
+    return sd.hessian_sparsity
+end
+
+function jacobian_sparsity(f, x, sd::KnownHessianSparsityDetector)
+    throw(ArgumentError("KnownHessianSparsityDetector can't be used to compute Jacobian sparsity."))
+end
+function jacobian_sparsity(f!, y, x, sd::KnownHessianSparsityDetector)
+    throw(ArgumentError("KnownHessianSparsityDetector can't be used to compute Jacobian sparsity."))
+end
+
 ## Coloring algorithm
 
 """

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -7,6 +7,8 @@ using ADTypes: AbstractADType,
                SymbolicMode
 using ADTypes: dense_ad,
                NoSparsityDetector,
+               KnownJacobianSparsityDetector,
+               KnownHessianSparsityDetector,
                sparsity_detector,
                jacobian_sparsity,
                hessian_sparsity,

diff --git a/test/sparse.jl b/test/sparse.jl
@@ -18,33 +18,124 @@
 end
 
 @testset "Sparsity detector" begin
-    sd = NoSparsityDetector()
-
     f_jac1(x) = vcat(x, x)
     f_jac2(x) = hcat(x, x)
     function f_jac! end
     f_hess(x) = sum(x)
 
-    for x in (rand(2), rand(2, 3)), f in (f_jac1, f_jac2)
-        y = f(x)
-        Js = jacobian_sparsity(f, x, sd)
-        @test Js isa AbstractMatrix
-        @test size(Js) == (length(y), length(x))
-        @test all(isone, Js)
+    @testset "NoSparsityDetector" begin
+        sd = NoSparsityDetector()
+
+        for x in (rand(2), rand(2, 3)), f in (f_jac1, f_jac2)
+            y = f(x)
+            Js = jacobian_sparsity(f, x, sd)
+            @test Js isa AbstractMatrix
+            @test size(Js) == (length(y), length(x))
+            @test all(isone, Js)
+        end
+
+        for x in (rand(2), rand(2, 3)), y in (rand(5), rand(5, 6))
+            Js = jacobian_sparsity(f_jac!, y, x, sd)
+            @test Js isa AbstractMatrix
+            @test size(Js) == (length(y), length(x))
+            @test all(isone, Js)
+        end
+
+        for x in (rand(2), rand(2, 3))
+            Hs = hessian_sparsity(f_hess, x, sd)
+            @test Hs isa AbstractMatrix
+            @test size(Hs) == (length(x), length(x))
+            @test all(isone, Hs)
+        end
     end
 
-    for x in (rand(2), rand(2, 3)), y in (rand(5), rand(5, 6))
-        Js = jacobian_sparsity(f_jac!, y, x, sd)
-        @test Js isa AbstractMatrix
-        @test size(Js) == (length(y), length(x))
-        @test all(isone, Js)
+    @testset "KnownJacobianSparsityDetector" begin
+        @testset "Jacobian sparsity detection" begin
+            @testset "Out-of-place functions" begin
+                for sx in ((2,), (2, 3)), f in (f_jac1, f_jac2)
+                    x = rand(sx...)
+                    nx = length(x)
+                    Jref = rand(Bool, 2 * nx, nx)
+                    sd = KnownJacobianSparsityDetector(Jref)
+                    Js = jacobian_sparsity(f, x, sd)
+                    @test Js isa AbstractMatrix
+                    @test size(Js) == (2 * nx, nx)
+                    @test Js === Jref
+                end
+            end
+            @testset "In-place functions" begin
+                for sx in ((2,), (2, 3)), sy in ((5,), (5, 6))
+                    x, y = rand(sx...), rand(sy...)
+                    nx, ny = length(x), length(y)
+                    Jref = rand(Bool, ny, nx)
+                    sd = KnownJacobianSparsityDetector(Jref)
+                    Js = jacobian_sparsity(f_jac!, y, x, sd)
+                    @test Js isa AbstractMatrix
+                    @test size(Js) == (ny, nx)
+                    @test Js === Jref
+                end
+            end
+        end
+        @testset "Exceptions: Hessian sparsity detection not supported" begin
+            for sx in ((2,), (2, 3))
+                x = rand(sx...)
+                nx = length(x)
+                Href = rand(Bool, nx, nx)
+                sd = KnownJacobianSparsityDetector(Href)
+                @test_throws ArgumentError hessian_sparsity(f_hess, x, sd)
+            end
+        end
+        @testset "Exceptions: DimensionMismatch" begin
+            sd = KnownJacobianSparsityDetector(rand(Bool, 6, 7)) # wrong Jacobian size
+            for x in (rand(2), rand(2, 3)), f in (f_jac1, f_jac2)
+                @test_throws DimensionMismatch jacobian_sparsity(f, x, sd)
+            end
+            for x in (rand(2), rand(2, 3)), y in (rand(5), rand(5, 6))
+                @test_throws DimensionMismatch jacobian_sparsity(f_jac!, y, x, sd)
+            end
+        end
     end
 
-    for x in (rand(2), rand(2, 3))
-        Hs = hessian_sparsity(f_hess, x, sd)
-        @test Hs isa AbstractMatrix
-        @test size(Hs) == (length(x), length(x))
-        @test all(isone, Hs)
+    @testset "KnownHessianSparsityDetector" begin
+        @testset "Hessian sparsity detection" begin
+            for sx in ((2,), (2, 3))
+                x = rand(sx...)
+                nx = length(x)
+                Href = rand(Bool, nx, nx)
+                sd = KnownHessianSparsityDetector(Href)
+
+                Hs = hessian_sparsity(f_hess, x, sd)
+                @test Hs isa AbstractMatrix
+                @test size(Hs) == (nx, nx)
+                @test Hs === Href
+            end
+        end
+        @testset "Exceptions: Jacobian sparsity detection not supported" begin
+            @testset "Out-of-place functions" begin
+                for sx in ((2,), (2, 3)), f in (f_jac1, f_jac2)
+                    x = rand(sx...)
+                    nx = length(x)
+                    Jref = rand(Bool, 2 * nx, nx)
+                    sd = KnownHessianSparsityDetector(Jref)
+                    @test_throws ArgumentError jacobian_sparsity(f, x, sd)
+                end
+            end
+            @testset "In-place functions" begin
+                for sx in ((2,), (2, 3)), sy in ((5,), (5, 6))
+                    x, y = rand(sx...), rand(sy...)
+                    nx, ny = length(x), length(y)
+                    Jref = rand(Bool, ny, nx)
+                    sd = KnownHessianSparsityDetector(Jref)
+                    @test_throws ArgumentError jacobian_sparsity(f_jac!, y, x, sd)
+                end
+            end
+        end
+        @testset "Exceptions: DimensionMismatch" begin
+            sd = KnownHessianSparsityDetector(rand(Bool, 2, 3)) #wrong Hessian size
+            for x in (rand(2), rand(2, 3))
+                @test_throws DimensionMismatch hessian_sparsity(f_hess, x, sd)
+            end
+        end
     end
 end