I managed to get the sparsity structure detection (via SparseDiffTools.j) of the rhs! function (the RHS of the ODE problem fed into OrdinaryDiffEq.jl) for our code Trixi.jl to run.
Can I supply the
cache = sparse_jacobian_cache(...)
to the implicit methods to speed things up, i.e., prevent the code from trying to construct the derivatives naively?
Can you share an MWE?
Perhaps this tutorial will be helpful: Solving Large Stiff Equations · DifferentialEquations.jl
That tutorial looks helpful, good to know that the sparsity pattern can be fed into the ODE problem itself.
The current code (with some dirty hacks) looks like this:
using LinearAlgebra, SparseArrays
using ForwardDiff
using Trixi
using SparseDiffTools, Symbolics
using BenchmarkTools
using Plots
###############################################################################
### Hacks ###
# Required for setting up the Lobatto Legendre basis for abstract `Real` type
function Trixi.eps(::Union{Type{Real}, Int}, RealT = Float64)
return eps(RealT)
end
import Base: * # For overloading
# Multiplying two Matrix{Real}s gives a Matrix{Any}.
# This causes problems when instantiating the Legendre basis.
# Called in `calc_{forward,reverse}_{upper, lower}`
function *(A::Matrix{Real}, B::Matrix{Real})::Matrix{Real}
m, n = size(A, 1), size(B, 2)
kA = size(A, 2)
kB = size(B, 1)
@assert kA == kB "Matrix dimensions must match for multiplication"
C = Matrix{Real}(undef, m, n)
for i in 1:m, j in 1:n
#acc::Real = zero(promote_type(typeof(A[i,1]), typeof(B[1,j])))
acc = zero(Real)
for k in 1:kA
acc += A[i,k] * B[k,j]
end
C[i,j] = acc
end
return C
end
###############################################################################
### semidiscretization of the linear advection equation ###
advection_velocity = 1.0
equations = LinearScalarAdvectionEquation1D(advection_velocity)
# Create DG solver with polynomial degree = 3 and (local) Lax-Friedrichs/Rusanov flux as surface flux
d = 3
# SD for SpareDiff, FD for ForwardDiff
# `RealT = Real` requires fewer overloads than the more explicit `RealT = Num`
solver_SD = DGSEM(polydeg = d, surface_flux = flux_lax_friedrichs, RealT = Real)
solver_FD = DGSEM(polydeg = d, surface_flux = flux_lax_friedrichs)
coordinates_min = -1.0 # minimum coordinate
coordinates_max = 1.0 # maximum coordinate
# Create a uniformly refined mesh with periodic boundaries
RefinementLevel = 3
mesh = TreeMesh(coordinates_min, coordinates_max,
initial_refinement_level = RefinementLevel,
n_cells_max = 30_000) # set maximum capacity of tree data structu0testre
# A semidiscretization collects data structures and functions for the spatial discretization
initial_condition = initial_condition_convergence_test
semi_SD = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver_SD)
semi_FD = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver_FD)
# 2^RefinementLevel = 16 elements, (d+1) local polynomial coefficients per element
N = Int(2^RefinementLevel * (d+1)) #64
u0_ode = zeros(N)
du_ode = 80 * ones(N) # initialize with something
t0 = 0.0
###############################################################################
### Compute the Jacobian with SparseDiffTools ###
rhs = (du_ode,u0_ode)->Trixi.rhs!(du_ode, u0_ode, semi_SD, t0) # Match an in_place function in SparseDiffTools
# From the example to detect the pattern and choose how to do the AutoDiff automatically
sd = SymbolicsSparsityDetection()
adtype = AutoSparseFiniteDiff()
println("Sparsity detection took:")
#From the example provided in SparseDiffTools. cache will reduce calculation time when Jacobian will be calculated multiple times
cache = @btime sparse_jacobian_cache(adtype, sd, rhs, du_ode, u0_ode)
println("Applying the sparse Jacobian took:")
J_SD = @btime sparse_jacobian(adtype, cache, rhs, du_ode, u0_ode)
###############################################################################
### Compute the Jacobian with ForwardDiff in order to compare with SparseDiff Jacobian ###
config = ForwardDiff.JacobianConfig(nothing, du_ode, u0_ode)
# See https://juliadiff.org/ForwardDiff.jl/v0.7/dev/how_it_works.html#Dual-Number-Implementation-1
uEltype = eltype(config)
# Change type of u0_ode and du_ode to automatic-differentiation type
u0_ode_FD = Vector{uEltype}(u0_ode)
du_ode_FD = Vector{uEltype}(du_ode)
println("ForwardDiff took:")
J_FD = @btime jacobian_ad_forward(semi_FD)
###############################################################################
### Compare Jacobians ###
println(J_SD==J_FD)
It is the general consensus that SparseDiffTools.jl should be deprecated in favor of the new sparse autodiff ecosystem built around DifferentiationInterface.jl + SparseConnectivityTracer.jl + SparseMatrixColorings.jl.
If you want to use sparse Jacobians in the framework of ordinary differential equations, the best way to do that is probably to follow the tutorial I linked above, and the rest of the OrdinaryDiffEq.jl docs. If you need more manual control for some reason, check out this demo for standalone computation of sparse Jacobian matrices.
Indeed this. If you give it the sparse matrix then it will do the coloring and sparse differentiation. You could instead pass a jac function but, if you’re using one of the auto diff libraries to do sparse diff, then it’s just more code to do what it would do automatically.
The only case I know of right now where that’s not the case is if you need GPU support, and that is because DI currently doesn’t support GPU sparse matrices while SparseDiffTools does. Am I guessing right that’s the actually issue here? Given the Trixi PDE mention it would make sense but it wasn’t made clear by the post.
As discussed on GitHub, can you please provide an MWE for GPU support so that I know what to fix?
In fact I am at this point not worried about GPUs yet. Thanks for clarifying the roadmap / way to go for the future!
I did some experimental tests on sparse GPU matrices in CUDA a while back and they seem to work ok, but if anyone actually relies on that functionality we will add it to the DI test suite. And if there’s a bug, please report it so it can get fixed.
Test notebook: Google Colab
It’s still the same case from the original PR that brought in DI, since DI broke the test suite 
The problem is DI uses the ForwardDiff.jacobian function which is not GPU compatible. I’ve been saying it should avoid it and use the same trick as SparseDiffTools, or finally patch ForwardDiff.jacobian (though there’s some trade-offs there).
Issue ForwardDiff scalar indexing with GPU arrays · Issue #820 · JuliaDiff/DifferentiationInterface.jl · GitHub, but I haven’t had a proper GPU on me for a few weeks now so I cannot do the step to make it just the DI call. But it should be pretty clear why it’s doing it since the ForwardDiff.jacobian call scalar indexes the partials which is exactly the operation that isn’t allowed and is pointed to in the error message.
You don’t check for scalar indexing. It’s slow because of scalar indexing.