switch to SparseConnectivityTracer and DifferentiationInterface

CHANGELOG.md

@@ -1,5 +1,13 @@
 # CHANGES
 
+## v1.5.0
+
+### Added
+  - new parameter `autodiff_backend` for NonlinearOperator to change differentiation backend
+
+### Changed
+  - sparsity patterns of local jacobians are now resolved by SparseConnectivityTracer
+  - local jacobians are now prepared by DifferentiationInterface 
 
 ## v1.4.0
 

Project.toml

@@ -1,11 +1,13 @@
 name = "ExtendableFEM"
 uuid = "a722555e-65e0-4074-a036-ca7ce79a4aed"
-version = "1.4"
+version = "1.5"
 authors = ["Christian Merdon <[email protected]>", "Patrick Jaap <[email protected]>"]
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 ExtendableFEMBase = "12fb9182-3d4c-4424-8fd1-727a0899810c"
 ExtendableGrids = "cfc395e8-590f-11e8-1f13-43a2532b2fa8"
@@ -17,15 +19,17 @@ LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
-Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 
 [compat]
+ADTypes = "1.16.0"
 Aqua = "0.8"
 CommonSolve = "0.2"
 DiffResults = "1"
+DifferentiationInterface = "0.7.4"
 DocStringExtensions = "0.8,0.9"
 ExampleJuggler = "2.2.1"
 ExplicitImports = "1"
@@ -41,10 +45,11 @@ Metis = "1.5.0"
 OrdinaryDiffEqRosenbrock = "1.4.0"
 OrdinaryDiffEqSDIRK = "1.2.0"
 Printf = "1.9"
-SciMLBase = "2.6"
+SciMLBase = "2.99.0"
 SimplexGridFactory = "2.4.0"
 SparseArrays = "1.9"
-SparseDiffTools = "^1.19,2"
+SparseConnectivityTracer = "1.0.0"
+SparseMatrixColorings = "0.4.21"
 StaticArrays = "1.9.13"
 Symbolics = "4.2,5,6"
 Test = "1"
@@ -64,9 +69,10 @@ OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
 SimplexGridFactory = "57bfcd06-606e-45d6-baf4-4ba06da0efd5"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TetGen = "c5d3f3f7-f850-59f6-8a2e-ffc6dc1317ea"
 Triangulate = "f7e6ffb2-c36d-4f8f-a77e-16e897189344"
 
 [targets]
-test = ["Aqua", "ExampleJuggler", "ExplicitImports", "IncompleteLU", "Metis", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "SimplexGridFactory", "StaticArrays", "Test", "TetGen", "Triangulate"]
+test = ["Aqua", "ExampleJuggler", "ExplicitImports", "IncompleteLU", "Metis", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "SimplexGridFactory", "StaticArrays", "Symbolics", "Test", "TetGen", "Triangulate"]

docs/src/index.md

@@ -23,7 +23,8 @@
 - [GridVisualize.jl](https://github.com/WIAS-PDELib/GridVisualize.jl)
 - [DocStringExtensions.jl](https://github.com/JuliaDocs/DocStringExtensions.jl)
 - [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)
-- [DiffResults.jl](https://github.com/JuliaDiff/DiffResults.jl)
+- [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl)
+- [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl)
 
 ## Getting Started
 

src/ExtendableFEM.jl

@@ -66,11 +66,13 @@ using LinearAlgebra: LinearAlgebra, copyto!, isposdef, mul!, norm
 using LinearSolve: LinearSolve, LinearProblem, UMFPACKFactorization, deleteat!,
     init, solve
 using Printf: Printf, @printf, @sprintf
-using SparseArrays: SparseArrays, AbstractSparseArray, SparseMatrixCSC, findnz, nnz,
+using SparseArrays: SparseArrays, AbstractSparseArray, findnz, nnz,
     nzrange, rowvals, sparse, SparseVector
-using SparseDiffTools: SparseDiffTools, ForwardColorJacCache,
-    forwarddiff_color_jacobian!, matrix_colors
-using Symbolics: Symbolics
+using ADTypes: ADTypes, KnownJacobianSparsityDetector
+using SparseConnectivityTracer: SparseConnectivityTracer, TracerSparsityDetector, jacobian_sparsity
+using DifferentiationInterface: DifferentiationInterface,
+    AutoSparse, AutoForwardDiff, prepare_jacobian
+using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern
 using SciMLBase: SciMLBase
 using TimerOutputs: TimerOutput, print_timer, @timeit
 using UnicodePlots: UnicodePlots

src/common_operators/bilinear_operator.jl

@@ -408,7 +408,8 @@ function build_assembler!(A::AbstractMatrix, O::BilinearOperator{Tv}, FE_test, F
         coupling_matrix::Matrix{Bool} = ones(Bool, nansatz, ntest)
         if use_sparsity_pattern
             kernel_params = (result, input) -> (O.kernel(result, input, O.QP_infos[1]))
-            sparsity_pattern = SparseMatrixCSC{Float64, Int}(Symbolics.jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end])))
+            detector = TracerSparsityDetector()
+            sparsity_pattern = jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end]), detector)
 
             ## find out which test and ansatz functions couple
             for id in 1:nansatz
@@ -777,7 +778,8 @@ function build_assembler!(A, O::BilinearOperator{Tv}, FE_test, FE_ansatz; time =
         coupling_matrix::Matrix{Bool} = ones(Bool, nansatz, ntest)
         if use_sparsity_pattern
             kernel_params = (result, input) -> (O.kernel(result, input, O.QP_infos[1]))
-            sparsity_pattern = SparseMatrixCSC{Float64, Int}(Symbolics.jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end])))
+            detector = TracerSparsityDetector()
+            sparsity_pattern = jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end]), detector)
 
             ## find out which test and ansatz functions couple
             for id in 1:nansatz

src/common_operators/bilinear_operator_dg.jl

@@ -374,7 +374,8 @@ function build_assembler!(A, O::BilinearOperatorDG{Tv}, FE_test, FE_ansatz, FE_a
         coupling_matrix::Matrix{Bool} = ones(Bool, nansatz, ntest)
         if use_sparsity_pattern
             kernel_params = (result, input) -> (O.kernel(result, input, O.QP_infos[1]))
-            sparsity_pattern = SparseMatrixCSC{Float64, Int}(Symbolics.jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end])))
+            detector = TracerSparsityDetector()
+            sparsity_pattern = jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end]), detector)
 
             ## find out which test and ansatz functions couple
             for id in 1:nansatz
@@ -789,7 +790,8 @@ function build_assembler!(A, O::BilinearOperatorDG{Tv}, FE_test, FE_ansatz; time
         coupling_matrix::Matrix{Bool} = ones(Bool, nansatz, ntest)
         if use_sparsity_pattern
             kernel_params = (result, input) -> (O.kernel(result, input, O.QP_infos[1]))
-            sparsity_pattern = SparseMatrixCSC{Float64, Int}(Symbolics.jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end])))
+            detector = TracerSparsityDetector()
+            sparsity_pattern = jacobian_sparsity(kernel_params, zeros(Tv, op_offsets_test[end]), zeros(Tv, op_offsets_ansatz[end]), detector)
 
             ## find out which test and ansatz functions couple
             for id in 1:nansatz

src/common_operators/nonlinear_operator.jl

@@ -1,11 +1,9 @@
-struct KernelEvaluator{Tv <: Real, QPInfosT <: QPInfos, JT, VT, DRT, CFGT, KPT <: Function}
+struct KernelEvaluator{Tv <: Real, QPInfosT <: QPInfos, DIT, JBT, KPT <: Function}
     input_args::Array{Tv, 1}
     params::QPInfosT
     result_kernel::Array{Tv, 1}
-    jac::JT
-    value::VT
-    Dresult::DRT
-    cfg::CFGT
+    jac_prep::DIT
+    jac_backend::JBT
     kernel::KPT
 end
 
@@ -43,6 +41,7 @@ default_nlop_kwargs() = Dict{Symbol, Tuple{Any, String}}(
     :regions => ([], "subset of regions where operator should be assembly only"),
     :sparse_jacobians => (true, "use sparse jacobians"),
     :sparse_jacobians_pattern => (nothing, "user provided sparsity pattern for the sparse jacobians (in case automatic detection fails)"),
+    :autodiff_backend => (AutoForwardDiff(), "backend for automatic differentiation"),
     :time_dependent => (false, "operator is time-dependent ?"),
     :verbosity => (0, "verbosity level"),
 )
@@ -230,37 +229,25 @@ function build_assembler!(A::AbstractMatrix, b::AbstractVector, O::NonlinearOper
         Kj = Array{KernelEvaluator, 1}([])
 
         sparse_jacobians = O.parameters[:sparse_jacobians]
-        sparsity_pattern = O.parameters[:sparse_jacobians_pattern]
         use_autodiff = O.jacobian === nothing
         for EG in EGs
             ## prepare parameters
             QPj = QPInfos(xgrid; time = time, x = ones(Tv, size(xgrid[Coordinates], 1)), params = O.parameters[:params])
             kernel_params = (result, input) -> (O.kernel(result, input, QPj))
-            if sparse_jacobians
-                input_args = zeros(Tv, op_offsets_args[end] + O.parameters[:extra_inputsize])
-                result_kernel = zeros(Tv, op_offsets_test[end])
-                if sparsity_pattern === nothing
-                    sparsity_pattern = Symbolics.jacobian_sparsity(kernel_params, result_kernel, input_args)
-                end
-                jac = Float64.(sparse(sparsity_pattern))
-                value = zeros(Tv, op_offsets_test[end])
-                colors = matrix_colors(jac)
-                Dresult = nothing
-                cfg = ForwardColorJacCache(
-                    kernel_params, input_args, nothing;
-                    dx = nothing,
-                    colorvec = colors,
-                    sparsity = sparsity_pattern
-                )
+            input_args = zeros(Tv, op_offsets_args[end] + O.parameters[:extra_inputsize])
+            result_kernel = zeros(Tv, op_offsets_test[end])
+            if O.parameters[:sparse_jacobians_pattern] === nothing
+                sparsity_detector = TracerSparsityDetector()
             else
-                input_args = zeros(Tv, op_offsets_args[end] + O.parameters[:extra_inputsize])
-                result_kernel = zeros(Tv, op_offsets_test[end])
-                Dresult = DiffResults.JacobianResult(result_kernel, input_args)
-                jac = DiffResults.jacobian(Dresult)
-                value = DiffResults.value(Dresult)
-                cfg = ForwardDiff.JacobianConfig(kernel_params, result_kernel, input_args, ForwardDiff.Chunk{op_offsets_args[end]}())
+                sparsity_detector = KnownJacobianSparsityDetector(O.parameters[:sparse_jacobians_pattern])
             end
-            push!(Kj, KernelEvaluator(input_args, QPj, result_kernel, jac, value, Dresult, cfg, kernel_params))
+            sparse_forward_backend = AutoSparse(
+                O.parameters[:autodiff_backend];
+                sparsity_detector = sparsity_detector,
+                coloring_algorithm = GreedyColoringAlgorithm()
+            )
+            jac_prep = prepare_jacobian(kernel_params, result_kernel, sparse_forward_backend, input_args)
+            push!(Kj, KernelEvaluator(input_args, QPj, result_kernel, jac_prep, sparse_forward_backend, kernel_params))
         end
 
         ## prepare parallel assembly
@@ -301,12 +288,20 @@ function build_assembler!(A::AbstractMatrix, b::AbstractVector, O::NonlinearOper
             ## extract kernel properties
             params = K.params
             input_args = K.input_args
-            result_kernel = K.result_kernel
-            cfg = K.cfg
-            Dresult = K.Dresult
-            jac = K.jac
-            value = K.value
+            value = K.result_kernel
+            jac_prep = K.jac_prep
+            jac_backend = K.jac_backend
             kernel_params = K.kernel
+            if O.parameters[:sparse_jacobians]
+                if O.parameters[:sparse_jacobians_pattern] === nothing
+                    jac_sparsity_pattern = sparsity_pattern(jac_prep)
+                else
+                    jac_sparsity_pattern = O.parameters[:sparse_jacobians_pattern]
+                end
+                jac = Tv.(sparse(jac_sparsity_pattern))
+            else
+                jac = zeros(Tv, length(value), length(input_args))
+            end
             params.time = time
 
             ndofs_test::Array{Int, 1} = [get_ndofs(ON_CELLS, FE, EG) for FE in FETypes_test]
@@ -359,12 +354,14 @@ function build_assembler!(A::AbstractMatrix, b::AbstractVector, O::NonlinearOper
                     ## get global x for quadrature point
                     eval_trafo!(params.x, L2G, xref[qp])
                     if use_autodiff
-                        if sparse_jacobians
-                            forwarddiff_color_jacobian!(jac, kernel_params, input_args, cfg)
-                            kernel_params(value, input_args)
-                        else
-                            ForwardDiff.chunk_mode_jacobian!(Dresult, kernel_params, result_kernel, input_args, cfg)
-                        end
+                        DifferentiationInterface.value_and_jacobian!(
+                            kernel_params,
+                            value,
+                            jac,
+                            jac_prep,
+                            jac_backend,
+                            input_args
+                        )
                     else
                         O.jacobian(jac, input_args, params)
                         O.kernel(value, input_args, params)

test/runtests.jl

@@ -8,6 +8,7 @@ using Metis
 using Aqua
 using Triangulate
 using SimplexGridFactory
+using Symbolics
 
 
 include("test_dgblf.jl")