Flesh out sparsity and secondorder things more, change the extension …

…structure
SciML · Jul 15, 2024 · 4fe1f54 · 4fe1f54
1 parent ad7ca08
commit 4fe1f54
Show file tree

Hide file tree

Showing 10 changed files with 112 additions and 113 deletions.
diff --git a/Project.toml b/Project.toml
@@ -7,6 +7,7 @@ version = "1.3.3"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
@@ -21,18 +22,20 @@ SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 
 [weakdeps]
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
-OptimizationDIExt = ["DifferentiationInterface", "ForwardDiff", "ReverseDiff"]
+OptimizationForwardDiffExt = "ForwardDiff"
+OptimizationFiniteDiffExt = "FiniteDiff"
+OptimizationReverseDiffExt = "ReverseDiff"
 OptimizationEnzymeExt = "Enzyme"
 OptimizationMTKExt = "ModelingToolkit"
-OptimizationZygoteExt = ["Zygote", "DifferentiationInterface"]
+OptimizationZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "1.3"
@@ -44,11 +47,9 @@ ModelingToolkit = "9"
 Reexport = "1.2"
 Requires = "1"
 SciMLBase = "2"
-SparseDiffTools = "2.14"
 SymbolicAnalysis = "0.1, 0.2"
 SymbolicIndexingInterface = "0.3"
 Symbolics = "5.12"
-Tracker = "0.2.29"
 Zygote = "0.6.67"
 julia = "1.10"
 

diff --git a/ext/OptimizationFiniteDiffExt.jl b/ext/OptimizationFiniteDiffExt.jl
@@ -0,0 +1,5 @@
+module OptimizationFiniteDiffExt
+
+using DifferentiationInterface, FiniteDiff
+
+end
diff --git a/ext/OptimizationForwardDiffExt.jl b/ext/OptimizationForwardDiffExt.jl
@@ -0,0 +1,5 @@
+module OptimizationForwardDiffExt
+
+using DifferentiationInterface, ForwardDiff
+
+end
diff --git a/ext/OptimizationMTKExt.jl b/ext/OptimizationMTKExt.jl
@@ -7,8 +7,8 @@ import OptimizationBase.ADTypes: AutoModelingToolkit, AutoSymbolics, AutoSparse
 using ModelingToolkit
 
 function OptimizationBase.instantiate_function(
-        f::OptimizationFunction{true}, x, adtype::AutoSparse{<:AutoSymbolics, S, C}, p,
-        num_cons = 0) where {S, C}
+        f::OptimizationFunction{true}, x, adtype::AutoSparse{<:AutoSymbolics}, p,
+        num_cons = 0)
     p = isnothing(p) ? SciMLBase.NullParameters() : p
 
     sys = complete(ModelingToolkit.modelingtoolkitize(OptimizationProblem(f, x, p;
@@ -53,7 +53,7 @@ function OptimizationBase.instantiate_function(
 end
 
 function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache,
-        adtype::AutoSparse{<:AutoSymbolics, S, C}, num_cons = 0) where {S, C}
+        adtype::AutoSparse{<:AutoSymbolics}, num_cons = 0)
     p = isnothing(cache.p) ? SciMLBase.NullParameters() : cache.p
 
     sys = complete(ModelingToolkit.modelingtoolkitize(OptimizationProblem(f, cache.u0,

diff --git a/ext/OptimizationReverseDiffExt.jl b/ext/OptimizationReverseDiffExt.jl
@@ -0,0 +1,5 @@
+module OptimizationReverseDiffExt
+
+using DifferentiationInterface, ReverseDiff
+
+end
diff --git a/src/OptimizationBase.jl b/src/OptimizationBase.jl
@@ -32,6 +32,8 @@ Base.length(::NullData) = 0
 include("adtypes.jl")
 include("cache.jl")
 include("function.jl")
+include("OptimizationDIExt.jl")
+include("OptimizationDISparseExt.jl")
 
 export solve, OptimizationCache, DEFAULT_CALLBACK, DEFAULT_DATA
 

diff --git a/ext/OptimizationDIExt.jl → src/OptimizationDIExt.jl b/ext/OptimizationDIExt.jl → src/OptimizationDIExt.jl
@@ -1,20 +1,18 @@
-module OptimizationDIExt
-
-import OptimizationBase, OptimizationBase.ArrayInterface
+using OptimizationBase
+import OptimizationBase.ArrayInterface
 import OptimizationBase.SciMLBase: OptimizationFunction
 import OptimizationBase.LinearAlgebra: I
 import DifferentiationInterface
 import DifferentiationInterface: prepare_gradient, prepare_hessian, prepare_hvp, prepare_jacobian, 
                                 gradient!, hessian!, hvp!, jacobian!, gradient, hessian, hvp, jacobian
-using ADTypes
-import ForwardDiff, ReverseDiff
+using ADTypes, SciMLBase
 
 function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x, adtype::ADTypes.AbstractADType, p = SciMLBase.NullParameters(), num_cons = 0)
     _f = (θ, args...) -> first(f.f(θ, p, args...))
 
-    if ADTypes.mode(adtype) isa ADTypes.ForwardMode
-        soadtype = DifferentiationInterface.SecondOrder(adtype, AutoReverseDiff())
-    elseif ADTypes.mode(adtype) isa ADTypes.ReverseMode
+    if !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ForwardMode
+        soadtype = DifferentiationInterface.SecondOrder(adtype, AutoReverseDiff()) #make zygote?
+    elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
         soadtype = DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype)
     end
 
@@ -32,7 +30,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
     if f.hess === nothing
         extras_hess = prepare_hessian(_f, soadtype, x)
         function hess(res, θ, args...)
-            hessian!(_f, res, adtype, θ, extras_hess)
+            hessian!(_f, res, soadtype, θ, extras_hess)
         end
     else
         hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
@@ -79,7 +77,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
 
         function cons_h(H, θ)
             for i in 1:num_cons
-                hessian!(fncs[i], H[i], adtype, θ, extras_cons_hess[i])
+                hessian!(fncs[i], H[i], soadtype, θ, extras_cons_hess[i])
             end
         end
     else
@@ -106,7 +104,12 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, ca
     x = cache.u0
     p = cache.p
     _f = (θ, args...) -> first(f.f(θ, p, args...))
-    soadtype = DifferentiationInterface.SecondOrder(adtype, adtype)
+
+    if !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ForwardMode
+        soadtype = DifferentiationInterface.SecondOrder(adtype, AutoReverseDiff()) #make zygote?
+    elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+        soadtype = DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype)
+    end
 
     if f.grad === nothing
         extras_grad = prepare_gradient(_f, adtype, x)
@@ -169,7 +172,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, ca
 
         function cons_h(H, θ)
             for i in 1:num_cons
-                hessian!(fncs[i], H[i], adtype, θ, extras_cons_hess[i])
+                hessian!(fncs[i], H[i], soadtype, θ, extras_cons_hess[i])
             end
         end
     else
@@ -195,7 +198,12 @@ end
 
 function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x, adtype::ADTypes.AbstractADType, p = SciMLBase.NullParameters(), num_cons = 0)
     _f = (θ, args...) -> first(f.f(θ, p, args...))
-    soadtype = DifferentiationInterface.SecondOrder(adtype, adtype)
+
+    if !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ForwardMode
+        soadtype = DifferentiationInterface.SecondOrder(adtype, AutoReverseDiff()) #make zygote?
+    elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+        soadtype = DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype)
+    end
 
     if f.grad === nothing
         extras_grad = prepare_gradient(_f, adtype, x)
@@ -211,7 +219,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x
     if f.hess === nothing
         extras_hess = prepare_hessian(_f, soadtype, x) #placeholder logic, can be made much better
         function hess(θ, args...)
-            hessian(_f, adtype, θ, extras_hess)
+            hessian(_f, soadtype, θ, extras_hess)
         end
     else
         hess = (θ, args...) -> f.hess(θ, p, args...)
@@ -259,7 +267,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x
 
         function cons_h(θ)
             H = map(1:num_cons) do i
-                hessian(fncs[i], adtype, θ, extras_cons_hess[i])
+                hessian(fncs[i], soadtype, θ, extras_cons_hess[i])
             end
             return H
         end
@@ -287,7 +295,12 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, c
     x = cache.u0
     p = cache.p
     _f = (θ, args...) -> first(f.f(θ, p, args...))
-    soadtype = DifferentiationInterface.SecondOrder(adtype, adtype)
+
+    if !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ForwardMode
+        soadtype = DifferentiationInterface.SecondOrder(adtype, AutoReverseDiff()) #make zygote?
+    elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+        soadtype = DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype)
+    end
 
     if f.grad === nothing
         extras_grad = prepare_gradient(_f, adtype, x)
@@ -351,7 +364,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, c
 
         function cons_h(θ)
             H = map(1:num_cons) do i
-                hessian(fncs[i], adtype, θ, extras_cons_hess[i])
+                hessian(fncs[i], soadtype, θ, extras_cons_hess[i])
             end
             return H
         end
@@ -374,5 +387,3 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, c
         cons_hess_colorvec = conshess_colors,
         lag_h, f.lag_hess_prototype)
 end
-
-end
diff --git a/ext/OptimizationDISparseExt.jl → src/OptimizationDISparseExt.jl b/ext/OptimizationDISparseExt.jl → src/OptimizationDISparseExt.jl
@@ -1,6 +1,5 @@
-module OptimizationDIExt
-
-import OptimizationBase, OptimizationBase.ArrayInterface
+using OptimizationBase
+import OptimizationBase.ArrayInterface
 import OptimizationBase.SciMLBase: OptimizationFunction
 import OptimizationBase.LinearAlgebra: I
 import DifferentiationInterface
@@ -9,21 +8,48 @@ import DifferentiationInterface: prepare_gradient, prepare_hessian, prepare_hvp,
 using ADTypes
 using SparseConnectivityTracer, SparseMatrixColorings
 
-function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x, adtype::ADTypes.AutoSparse, p = SciMLBase.NullParameters(), num_cons = 0)
-    _f = (θ, args...) -> first(f.f(θ, p, args...))
-
+function generate_sparse_adtype(adtype)
     if adtype.sparsity_detector isa ADTypes.NoSparsityDetector && adtype.coloring_algorithm isa ADTypes.NoColoringAlgorithm
         adtype = AutoSparse(adtype.dense_ad; sparsity_detector = TracerLocalSparsityDetector(), coloring_algorithm = GreedyColoringAlgorithm())
-    elseif adtype.sparsity_detector isa ADTypes.NoSparsityDetector && !(adtype.coloring_algorithm isa AbstractADTypes.NoColoringAlgorithm)
+        if !(adtype.dense_ad isa SciMLBase.NoAD) && ADTypes.mode(adtype.dense_ad) isa ADTypes.ForwardMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(adtype.dense_ad, AutoReverseDiff()), sparsity_detector = TracerLocalSparsityDetector(), coloring_algorithm = GreedyColoringAlgorithm()) #make zygote?
+        elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype), sparsity_detector = TracerLocalSparsityDetector(), coloring_algorithm = GreedyColoringAlgorithm())
+        end
+    elseif adtype.sparsity_detector isa ADTypes.NoSparsityDetector && !(adtype.coloring_algorithm isa ADTypes.NoColoringAlgorithm)
         adtype = AutoSparse(adtype.dense_ad; sparsity_detector = TracerLocalSparsityDetector(), coloring_algorithm = adtype.coloring_algorithm)
+        if !(adtype.dense_ad isa SciMLBase.NoAD) && ADTypes.mode(adtype.dense_ad) isa ADTypes.ForwardMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(adtype.dense_ad, AutoReverseDiff()), sparsity_detector = TracerLocalSparsityDetector(), coloring_algorithm = adtype.coloring_algorithm)
+        elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype), sparsity_detector = TracerLocalSparsityDetector(), coloring_algorithm = adtype.coloring_algorithm)
+        end
     elseif !(adtype.sparsity_detector isa ADTypes.NoSparsityDetector) && adtype.coloring_algorithm isa ADTypes.NoColoringAlgorithm
         adtype = AutoSparse(adtype.dense_ad; sparsity_detector = adtype.sparsity_detector, coloring_algorithm = GreedyColoringAlgorithm())
+        if !(adtype.dense_ad isa SciMLBase.NoAD) && ADTypes.mode(adtype.dense_ad) isa ADTypes.ForwardMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(adtype.dense_ad, AutoReverseDiff()), sparsity_detector = adtype.sparsity_detector, coloring_algorithm = GreedyColoringAlgorithm())
+        elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype), sparsity_detector = adtype.sparsity_detector, coloring_algorithm = GreedyColoringAlgorithm())
+        end
+    else
+        if !(adtype.dense_ad isa SciMLBase.NoAD) && ADTypes.mode(adtype.dense_ad) isa ADTypes.ForwardMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(adtype.dense_ad, AutoReverseDiff()), sparsity_detector = adtype.sparsity_detector, coloring_algorithm = adtype.coloring_algorithm)
+        elseif !(adtype isa SciMLBase.NoAD) && ADTypes.mode(adtype) isa ADTypes.ReverseMode
+            soadtype = AutoSparse(DifferentiationInterface.SecondOrder(AutoForwardDiff(), adtype), sparsity_detector = adtype.sparsity_detector, coloring_algorithm = adtype.coloring_algorithm)
+        end
     end
+    return adtype,soadtype
+end
+
+
+function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x, adtype::ADTypes.AutoSparse{<:AbstractADType}, p = SciMLBase.NullParameters(), num_cons = 0)
+    _f = (θ, args...) -> first(f.f(θ, p, args...))
+
+    adtype, soadtype = generate_sparse_adtype(adtype)
 
     if f.grad === nothing
-        extras_grad = prepare_gradient(_f, adtype, x)
+        extras_grad = prepare_gradient(_f, adtype.dense_ad, x)
         function grad(res, θ)
-            gradient!(_f, res, adtype, θ, extras_grad)
+            gradient!(_f, res, adtype.dense_ad, θ, extras_grad)
         end
     else
         grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
@@ -34,7 +60,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
     if f.hess === nothing
         extras_hess = prepare_hessian(_f, soadtype, x) #placeholder logic, can be made much better
         function hess(res, θ, args...)
-            hessian!(_f, res, adtype, θ, extras_hess)
+            hessian!(_f, res, soadtype, θ, extras_hess)
         end
     else
         hess = (H, θ, args...) -> f.hess(H, θ, p, args...)
@@ -81,7 +107,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
 
         function cons_h(H, θ)
             for i in 1:num_cons
-                hessian!(fncs[i], H[i], adtype, θ, extras_cons_hess[i])
+                hessian!(fncs[i], H[i], soadtype, θ, extras_cons_hess[i])
             end
         end
     else
@@ -104,11 +130,12 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
         lag_h, f.lag_hess_prototype)
 end
 
-function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache, adtype::ADTypes.AbstractADType, num_cons = 0)
+function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache, adtype::ADTypes.AutoSparse{<:AbstractADType}, num_cons = 0)
     x = cache.u0
     p = cache.p
     _f = (θ, args...) -> first(f.f(θ, p, args...))
-    soadtype = DifferentiationInterface.SecondOrder(adtype, adtype)
+
+    adtype, soadtype = generate_sparse_adtype(adtype)
 
     if f.grad === nothing
         extras_grad = prepare_gradient(_f, adtype, x)
@@ -171,7 +198,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, ca
 
         function cons_h(H, θ)
             for i in 1:num_cons
-                hessian!(fncs[i], H[i], adtype, θ, extras_cons_hess[i])
+                hessian!(fncs[i], H[i], soadtype, θ, extras_cons_hess[i])
             end
         end
     else
@@ -195,9 +222,10 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, ca
 end
 
 
-function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x, adtype::ADTypes.AbstractADType, p = SciMLBase.NullParameters(), num_cons = 0)
+function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x, adtype::ADTypes.AutoSparse{<:AbstractADType}, p = SciMLBase.NullParameters(), num_cons = 0)
     _f = (θ, args...) -> first(f.f(θ, p, args...))
-    soadtype = DifferentiationInterface.SecondOrder(adtype, adtype)
+
+    adtype, soadtype = generate_sparse_adtype(adtype)
 
     if f.grad === nothing
         extras_grad = prepare_gradient(_f, adtype, x)
@@ -213,7 +241,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x
     if f.hess === nothing
         extras_hess = prepare_hessian(_f, soadtype, x) #placeholder logic, can be made much better
         function hess(θ, args...)
-            hessian(_f, adtype, θ, extras_hess)
+            hessian(_f, soadtype, θ, extras_hess)
         end
     else
         hess = (θ, args...) -> f.hess(θ, p, args...)
@@ -261,7 +289,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x
 
         function cons_h(θ)
             H = map(1:num_cons) do i
-                hessian(fncs[i], adtype, θ, extras_cons_hess[i])
+                hessian(fncs[i], soadtype, θ, extras_cons_hess[i])
             end
             return H
         end
@@ -285,11 +313,12 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x
         lag_h, f.lag_hess_prototype)
 end
 
-function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, cache::OptimizationBase.ReInitCache, adtype::ADTypes.AbstractADType, num_cons = 0)
+function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, cache::OptimizationBase.ReInitCache, adtype::ADTypes.AutoSparse{<:AbstractADType}, num_cons = 0)
     x = cache.u0
     p = cache.p
     _f = (θ, args...) -> first(f.f(θ, p, args...))
-    soadtype = DifferentiationInterface.SecondOrder(adtype, adtype)
+
+    adtype, soadtype = generate_sparse_adtype(adtype)
 
     if f.grad === nothing
         extras_grad = prepare_gradient(_f, adtype, x)
@@ -353,7 +382,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, c
 
         function cons_h(θ)
             H = map(1:num_cons) do i
-                hessian(fncs[i], adtype, θ, extras_cons_hess[i])
+                hessian(fncs[i], soadtype, θ, extras_cons_hess[i])
             end
             return H
         end
@@ -376,5 +405,3 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, c
         cons_hess_colorvec = conshess_colors,
         lag_h, f.lag_hess_prototype)
 end
-
-end