Merge pull request #283 from ErikQQY/qqy/di

Use DifferentiationInterface for Jacobian stuff

Author: ChrisRackauckas

Project.toml

@@ -1,6 +1,6 @@
 name = "BoundaryValueDiffEq"
 uuid = "764a87c0-6b3e-53db-9096-fe964310641d"
-version = "5.15.0"
+version = "5.16.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -25,7 +25,7 @@ ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c"
 BoundaryValueDiffEqODEInterfaceExt = "ODEInterface"
 
 [compat]
-ADTypes = "1.11"
+ADTypes = "1.13"
 Aqua = "0.8.9"
 ArrayInterface = "7.18"
 BoundaryValueDiffEqAscher = "1"

docs/src/basics/autodiff.md

@@ -2,7 +2,7 @@
 
 !!! note
 
-    We support all backends supported by DifferentiationInterface.jl. Please refer to
+    We support ForwardDiff.jl, FiniteDiff.jl and PolyesterForwardDiff.jl(PolyesterForwardDiff only for collocation methods) via DifferentiationInterface.jl. Please refer to
     the [backends page](https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/backends/)
     for more information.
 
@@ -17,8 +17,6 @@ In BoundaryValueDiffEq.jl, we require AD to obtain the Jacobian of the loss func
 
   - [`AutoFiniteDiff`](@extref ADTypes.AutoFiniteDiff): Finite differencing using
     `FiniteDiff.jl`, not optimal but always applicable.
-  - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences): Finite differencing
-    using `FiniteDifferences.jl`, not optimal but always applicable.
 
 ## Summary of Forward Mode AD Backends
 
@@ -27,10 +25,3 @@ In BoundaryValueDiffEq.jl, we require AD to obtain the Jacobian of the loss func
   - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff): Might be faster
     than [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff) for large problems. Requires
     `PolyesterForwardDiff.jl` to be installed and loaded.
-
-## Summary of Reverse Mode AD Backends
-
-  - [`AutoZygote`](@extref ADTypes.AutoZygote): The fastest choice for non-mutating
-    array-based (BLAS) functions.
-  - [`AutoEnzyme`](@extref ADTypes.AutoEnzyme): Uses `Enzyme.jl` Reverse Mode and works for
-    both in-place and out-of-place functions.

lib/BoundaryValueDiffEqAscher/Project.toml

@@ -1,22 +1,21 @@
 name = "BoundaryValueDiffEqAscher"
 uuid = "7227322d-7511-4e07-9247-ad6ff830280e"
 authors = ["Qingyu Qu <[email protected]>"]
-version = "1.4.0"
+version = "1.5.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AlmostBlockDiagonals = "a95523ee-d6da-40b5-98cc-27bc505739d5"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BoundaryValueDiffEqCore = "56b672f2-a5fe-4263-ab2d-da677488eb3a"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
-Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
@@ -25,26 +24,24 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 
 [compat]
 ADTypes = "1.11"
-Adapt = "4.1.1"
 AlmostBlockDiagonals = "0.1.10"
 ArrayInterface = "7.18"
 BandedMatrices = "1.7.5"
 BoundaryValueDiffEqCore = "1"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.158.3"
 DiffEqDevTools = "2.44"
+DifferentiationInterface = "0.6.42"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.38"
 Hwloc = "3"
 InteractiveUtils = "<0.0.1, 1"
 JET = "0.9.12"
 LinearAlgebra = "1.10"
 LinearSolve = "2.36.2, 3"
-Logging = "1.10"
 PreallocationTools = "0.4.24"
 PrecompileTools = "1.2"
 Preferences = "1.4"
@@ -55,7 +52,6 @@ Reexport = "1.2"
 SciMLBase = "2.71"
 Setfield = "1.1.1"
 SparseArrays = "1.10"
-SparseDiffTools = "2.23"
 StaticArrays = "1.9.8"
 Test = "1.10"
 julia = "1.10"

lib/BoundaryValueDiffEqAscher/src/BoundaryValueDiffEqAscher.jl

@@ -7,9 +7,10 @@ using BoundaryValueDiffEqCore: BVPJacobianAlgorithm, __extract_problem_details,
                                __concrete_nonlinearsolve_algorithm,
                                __internal_nlsolve_problem, BoundaryValueDiffEqAlgorithm,
                                __vec, __vec_f, __vec_f!, __vec_bc, __vec_bc!,
-                               __extract_mesh, get_dense_ad, __sparse_jacobian_cache
+                               __extract_mesh, get_dense_ad, __get_bcresid_prototype
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase
+using DifferentiationInterface: DifferentiationInterface, Constant, prepare_jacobian
 using FastClosures: @closure
 using ForwardDiff: ForwardDiff, Dual
 using LinearAlgebra
@@ -19,8 +20,8 @@ using Reexport: @reexport
 using SciMLBase: SciMLBase, AbstractDiffEqInterpolation, StandardBVProblem, __solve,
                  _unwrap_val
 using Setfield: @set!
-using SparseDiffTools: init_jacobian, sparse_jacobian, sparse_jacobian_cache,
-                       sparse_jacobian!, SymbolicsSparsityDetection, NoSparsityDetection
+
+const DI = DifferentiationInterface
 
 @reexport using ADTypes, BoundaryValueDiffEqCore, SciMLBase
 

lib/BoundaryValueDiffEqAscher/src/ascher.jl

@@ -136,7 +136,7 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractAscher; dt = 0.0,
     end
 
     if prob.f.bcjac === nothing
-        bcjac = construct_bc_jac(prob, bcresid_prototype, prob.problem_type)
+        bcjac = construct_bc_jac(prob)
     else
         bcjac = prob.f.bcjac
     end
@@ -177,7 +177,7 @@ end
 function __perform_ascher_iteration(cache::AscherCache{iip, T}, abstol, adaptive::Bool;
         nlsolve_kwargs = (;), kwargs...) where {iip, T}
     info::ReturnCode.T = ReturnCode.Success
-    nlprob::NonlinearProblem = __construct_nlproblem(cache)
+    nlprob = __construct_nlproblem(cache)
     nlsolve_alg = __concrete_nonlinearsolve_algorithm(nlprob, cache.alg.nlsolve)
     nlsol = __solve(nlprob, nlsolve_alg; abstol, kwargs..., nlsolve_kwargs...)
     error_norm = 2 * abstol
@@ -315,31 +315,50 @@ function __construct_nlproblem(cache::AscherCache{iip, T}) where {iip, T}
     else
         @closure (z, p) -> @views Φ(cache, z, pt)
     end
+
     lz = reduce(vcat, cache.z)
-    sd = alg.jac_alg.diffmode isa AutoSparse ? SymbolicsSparsityDetection() :
-         NoSparsityDetection()
-    ad = alg.jac_alg.diffmode
-    lossₚ = (iip ? __Fix3 : Base.Fix2)(loss, cache.p)
-    jac_cache = __sparse_jacobian_cache(Val(iip), ad, sd, lossₚ, lz, lz)
-    jac_prototype = init_jacobian(jac_cache)
+    resid_prototype = zero(lz)
+    diffmode = if alg.jac_alg.diffmode isa AutoSparse
+        #AutoSparse(get_dense_ad(alg.jac_alg.diffmode);
+        #    sparsity_detector = SparseConnectivityTracer.TracerSparsityDetector(),
+        #    coloring_algorithm = GreedyColoringAlgorithm(LargestFirst()))
+        # Ascher collocation need more generalized collocation to support AutoSparse
+        get_dense_ad(alg.jac_alg.diffmode)
+    else
+        alg.jac_alg.diffmode
+    end
+
+    jac_cache = if iip
+        DI.prepare_jacobian(loss, resid_prototype, diffmode, lz, Constant(cache.p))
+    else
+        DI.prepare_jacobian(loss, diffmode, lz, Constant(cache.p))
+    end
+
+    jac_prototype = if iip
+        DI.jacobian(loss, resid_prototype, jac_cache, diffmode, lz, Constant(cache.p))
+    else
+        DI.jacobian(loss, jac_cache, diffmode, lz, Constant(cache.p))
+    end
+
     jac = if iip
-        @closure (J, u, p) -> __ascher_mpoint_jacobian!(J, u, ad, jac_cache, lossₚ, lz)
+        @closure (J, u, p) -> __ascher_mpoint_jacobian!(
+            J, u, diffmode, jac_cache, loss, lz, cache.p)
     else
-        @closure (u, p) -> __ascher_mpoint_jacobian(jac_prototype, u, ad, jac_cache, lossₚ)
+        @closure (u, p) -> __ascher_mpoint_jacobian(
+            jac_prototype, u, diffmode, jac_cache, loss, cache.p)
     end
-    resid_prototype = zero(lz)
-    _nlf = NonlinearFunction{iip}(
+
+    nlf = NonlinearFunction{iip}(
         loss; jac = jac, resid_prototype = resid_prototype, jac_prototype = jac_prototype)
-    nlprob::NonlinearProblem = NonlinearProblem(_nlf, lz, cache.p)
-    return nlprob
+    return __internal_nlsolve_problem(cache.prob, similar(lz), lz, nlf, lz, cache.p)
 end
 
-function __ascher_mpoint_jacobian!(J, x, diffmode, diffcache, loss, resid)
-    sparse_jacobian!(J, diffmode, diffcache, loss, resid, x)
+function __ascher_mpoint_jacobian!(J, x, diffmode, diffcache, loss, resid, p)
+    DI.jacobian!(loss, resid, J, diffcache, diffmode, x, Constant(p))
     return nothing
 end
-function __ascher_mpoint_jacobian(J, x, diffmode, diffcache, loss)
-    sparse_jacobian!(J, diffmode, diffcache, loss, x)
+function __ascher_mpoint_jacobian(J, x, diffmode, diffcache, loss, p)
+    DI.jacobian!(loss, J, diffcache, diffmode, x, Constant(p))
     return J
 end