src: add AutoSparseDI option for jacobian based on DI

Project.toml

@@ -14,8 +14,16 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+[weakdeps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
+
+[extensions]
+SparseDIExt = ["DifferentiationInterface", "SparseConnectivityTracer", "SparseMatrixColorings"]
+
 [compat]
-BifurcationKit = "^0.5.5"
+BifurcationKit = "^0.5.6"
 DocStringExtensions = "^0.9"
 ForwardDiff = "^0.10, ^0.11"
 NonlinearEigenproblems = "^1.0.1"

ext/SparseDIExt.jl

@@ -0,0 +1,62 @@
+module SparseDIExt
+
+using SparseConnectivityTracer, SparseMatrixColorings
+import DifferentiationInterface as DI
+import DDEBifurcationKit as DDEBK
+import BifurcationKit as BK
+
+coll_residual_for_di!(result, u, coll, pars) = BK.residual!(coll, result, u, pars)
+
+function DDEBK.BK._generate_jacobian(coll::BK.PeriodicOrbitOCollProblem{Tprob},
+                    ::DDEBK.AutoSparseDI, 
+                    orbitguess, 
+                    pars; 
+                    k...) where {Tprob <: DDEBK.AbstractDDEBifurcationProblem}
+    backend = DI.AutoForwardDiff()
+    sparse_forward_backend = DI.AutoSparse(
+        backend;
+        # we use the following instead of TracerSparsityDetector
+        # because of searchsortedfirst in POSolution
+        # sparsity_detector = TracerLocalSparsityDetector(),
+        sparsity_detector = DI.DenseSparsityDetector(backend, atol = 1e-8),
+        coloring_algorithm = GreedyColoringAlgorithm(),
+    )
+    out = copy(orbitguess)
+
+    jac_prep_sparse_nonallocating = DI.prepare_jacobian(coll_residual_for_di!,
+                                    out,
+                                    sparse_forward_backend,
+                                    orbitguess,
+                                    DI.Constant(coll),
+                                    DI.Constant(pars),
+                                    # strict = Val(false)
+                                    )
+
+    jac_buffer = similar(sparsity_pattern(jac_prep_sparse_nonallocating), eltype(ones(length(out))))
+    L1 = copy(jac_buffer)
+    return (DDEBK.AutoSparseDI(),
+            L1,
+            out,
+            jac_prep_sparse_nonallocating,
+            sparse_forward_backend)
+end
+
+function BK.jacobian(coll::BK.PeriodicOrbitOCollProblem, 
+            J::Tuple{DDEBK.AutoSparseDI, T1, T2, T3, T4}, 
+            x, 
+            p) where {T1, T2, T3, T4}
+    L1 = J[2]
+    out = J[3]
+    jac_prep_sparse_nonallocating = J[4]
+    sparse_forward_backend = J[5]
+    DI.jacobian!(coll_residual_for_di!,
+                out,
+                L1,
+                jac_prep_sparse_nonallocating,
+                sparse_forward_backend,
+                x,
+                DI.Constant(p))
+    return L1
+end
+
+end # module

src/periodicorbit/PeriodicOrbitCollocation.jl

@@ -1,4 +1,3 @@
-# TODO use getter from BK
 _get_gauss_nodes(coll) = coll.mesh_cache.gauss_nodes
 
 @views function BK.residual!(coll::PeriodicOrbitOCollProblem{Tprob},
@@ -67,7 +66,8 @@ end
             end
             # udj = VectorOfArray([interp(mod(τ * period - d, period)) for d in _delays])
             for (ind, d) in enumerate(_delays)
-                udj.u[ind] .= interp(τ * period - d)
+                # udj.u[ind] .= interp(τ * period - d)
+                udj.u[ind] .= BK.__interpolate_posolution(coll, τ - d/period, u, 1)
             end
             __po_coll_bc!(coll, outc[:, rg[l]], ∂gj[:, l], gj[:, l], udj, pars, period * dτj, outc[:, end])
         end
@@ -92,7 +92,14 @@ function BK.jacobian(coll::PeriodicOrbitOCollProblem{Tprob},
     return analytical_jacobian_dde_cst(coll, x, p)
 end
 
-# analytical jacobian for constant DDE
+"""
+using DifferentiationInterface to automatically derive the sparse jacobian.
+"""
+struct AutoSparseDI <: BK.AbstractJacobianSparseMatrix end
+
+function BK._generate_jacobian(coll::PeriodicOrbitOCollProblem{Tprob}, ::AutoSparseDI, orbitguess, pars; k...) where {Tprob <: AbstractDDEBifurcationProblem}
+    error("You need to import `DifferentiationInterface, SparseConnectivityTracer, SparseMatrixColorings` in order to use this jacobian")
+end
 ########################################################################################
 # analytical jacobians for constant DDE
 for (fname, floquet) in ((:analytical_jacobian_dde_cst, false), 
@@ -171,8 +178,8 @@ for (fname, floquet) in ((:analytical_jacobian_dde_cst, false),
                                                     (ρD * ∂L[l2, l] - α * L[l2, l] * ρI) * In
                     for (idelay, d) in enumerate(delays_v)
                         # find interval where t-τ/period belongs
-                        t0 = τ * period - d
-                        τd = mod(t0, period) / period
+                        t0 = τ - d/period
+                        τd = mod(t0, 1) / 1
                         index_t = searchsortedfirst(mesh, τd) - 1
                         @assert 1 <= index_t <= Ntst "We have index_t = $index_t, which is out of bounds for mesh of size $(length(mesh)) and τd = $τd. Please open an issue on the website of BifurcationKit.jl"
 
@@ -184,6 +191,9 @@ for (fname, floquet) in ((:analytical_jacobian_dde_cst, false),
                             Jd[idelay][_rgX, rgNy_delay .+ (l2-1)*n] .+= -α .* JacDDE.Jd[idelay] .* β
                         elseif ($(fname == :analytical_jacobian_dde_cst_floquetcoll) && t0 < 0)
                             rgNy_delay = UnitRange(1, n) .+ ((m * n) * (index_t - 1))
+                            # fullmesh = coll.mesh_cache.full_mesh
+                            # index_tau = searchsortedlast(fullmesh, τd)- 2
+                            # rgNy_delay = UnitRange(1, n) .+ index_tau * n
                             Jd[_rgX, rgNy_delay .+ (l2-1)*n] .+= -α .* JacDDE.Jd[idelay] .* β
                         else # case analytical_jacobian_dde_cst
                             J[_rgX, rgNy_delay .+ (l2-1)*n] .+= -α .* JacDDE.Jd[idelay] .* β
@@ -213,4 +223,4 @@ for (fname, floquet) in ((:analytical_jacobian_dde_cst, false),
     end # function end
 
     end # begin
-end # for-loop end
+end # for-loop end

test/test_DI.jl

@@ -0,0 +1,46 @@
+using DDEBifurcationKit, LinearAlgebra, SparseArrays
+using BifurcationKit
+const BK = BifurcationKit
+const DDEBK = DDEBifurcationKit
+
+function Hutchinson(u, ud, p)
+   (;a,d,Δ) = p
+   d .* (Δ*u) .- a .* ud.u[1] .* (1 .+ u)
+end
+
+delaysF(par) = [1.]
+
+# discretisation
+Nx = 100; Lx = pi/2;
+X = -Lx .+ 2Lx/Nx*(0:Nx-1) |> collect
+h = 2Lx/Nx
+Δ = spdiagm(0 => -2ones(Nx), 1 => ones(Nx-1), -1 => ones(Nx-1) ) / h^2; Δ[1,1]=Δ[end,end]=-1/h^2
+
+pars = (a = 0.5, d = 1, τ = 1.0, Δ = Δ, N = Nx)
+x0 = zeros(Nx)
+
+################################################################################
+# case where we specify the jacobian
+function JacHutchinson(u, p)
+	(;a,d,Δ) = p
+   # we compute the jacobian at the steady state
+   J0 = d * Δ .- a .* Diagonal(u)
+   J1 = -a .* Diagonal(1 .+ u)
+   return J0, [J1]
+end
+
+prob = ConstantDDEBifProblem(Hutchinson, delaysF, x0, pars, (@optic _.a); J = JacHutchinson)
+optn = NewtonPar(verbose = true, eigsolver = DDE_DefaultEig())
+opts = ContinuationPar(p_max = 10., p_min = 0., newton_options = optn, ds = 0.01, detect_bifurcation = 3, nev = 5, dsmax = 0.2, n_inversion = 4)
+br = continuation(prob, PALC(), opts; verbosity = 1, plot = false, normC = norminf)
+################################################################################
+# periodic orbits
+using DifferentiationInterface, SparseConnectivityTracer, SparseMatrixColorings
+
+br_pocoll = @time continuation(
+            br, 3, ContinuationPar(br.contparams; detect_bifurcation = 0, max_steps = 3, newton_options = NewtonPar(eigsolver = DDE_DefaultEig(), verbose = true), plot_every_step = 1),
+            PeriodicOrbitOCollProblem(20, 4; jacobian = DDEBK.AutoSparseDI());
+            verbosity = 2,
+            # plot = true,
+            normC = norminf,
+            )