Making autodiff Jacobian reusable

src/DynamicalSystemsBase.jl

@@ -24,6 +24,7 @@ include("core_systems/discrete_time_map.jl")
 include("core_systems/continuous_time_ode.jl")
 include("core_systems/arbitrary_steppable.jl")
 include("core_systems/additional_supertypes.jl")
+include("core_systems/jacobian.jl")
 
 include("derived_systems/stroboscopic_map.jl")
 include("derived_systems/parallel_systems.jl")

src/core_systems/jacobian.jl

@@ -0,0 +1,30 @@
+export jacobian
+
+import ForwardDiff
+
+function jacobian(ds::CoreDynamicalSystem{IIP}) where {IIP}
+    _jacobian(ds, Val{IIP}())
+end
+
+function _jacobian(ds, ::Val{true})
+    f = dynamic_rule(ds)
+    u0 = current_state(ds)
+    cfg = ForwardDiff.JacobianConfig(
+        (du, u) -> f(du, u, p, p), deepcopy(u0), deepcopy(u0)
+    )
+    Jf = (J0, u, p, t) -> begin
+        uv = @view u[:, 1]
+        du = copy(u)
+        ForwardDiff.jacobian!(
+            J0, (du, u) -> f(du, u, p, t), view(du, :, 1), uv, cfg, Val{false}()
+        )
+        nothing
+    end
+    return Jf
+end
+
+function _jacobian(ds, ::Val{false})
+    f = dynamic_rule(ds)
+    Jf = (u, p, t) -> ForwardDiff.jacobian((x) -> f(x, p, t), u)
+    return Jf
+end

src/derived_systems/tangent_space.jl

@@ -113,6 +113,7 @@ function TangentDynamicalSystem(ds::CoreDynamicalSystem{IIP};
     u0_correct = correct_state(Val{IIP}(), u0)
     Q0_correct = correct_matrix_type(Val{IIP}(), Q0)
     newstate = hcat(u0_correct, Q0_correct)
+    J = isnothing(J) ? jacobian(ds) : J
     newrule = tangent_rule(f, J, J0, Val{IIP}(), Val{k}(), u0_correct)
 
     # Pass everything to analytic system constructors
@@ -149,37 +150,12 @@ function tangent_rule(f::F, J::JAC, J0, ::Val{true}, ::Val{k}, u0) where {F, JAC
     end
     return tangentf
 end
-# for the case of autodiffed systems, a specialized version is created
-# so that f! is not called twice in ForwardDiff
-function tangent_rule(f::F, ::Nothing, J0, ::Val{true}, ::Val{k}, u0) where {F, k}
-    let
-        cfg = ForwardDiff.JacobianConfig(
-            (du, u) -> f(du, u, p, p), deepcopy(u0), deepcopy(u0)
-        )
-        tangentf = (du, u, p, t) -> begin
-            uv = @view u[:, 1]
-            ForwardDiff.jacobian!(
-                J0, (du, u) -> f(du, u, p, t), view(du, :, 1), uv, cfg, Val{false}()
-            )
-            mul!((@view du[:, 2:k+1]), J0, (@view u[:, 2:k+1]))
-            nothing
-        end
-        return tangentf
-    end
-end
 
 # OOP Tangent space dynamics
 function tangent_rule(f::F, J::JAC, J0, ::Val{false}, ::Val{k}, u0) where {F, JAC, k}
-    # out of place
-    if JAC == Nothing
-        # There is no config needed here
-        Jf = (u, p, t) -> ForwardDiff.jacobian((x) -> f(x, p, t), u)
-    else
-        Jf = J
-    end
     # Initial matrix `J0` is ignored
     ws_index = SVector{k, Int}(2:(k+1)...)
-    tangentf = TangentOOP(f, Jf, ws_index)
+    tangentf = TangentOOP(f, J, ws_index)
     return tangentf
 end
 struct TangentOOP{F, JAC, k} <: Function

test/jacobian.jl

@@ -0,0 +1,29 @@
+using DynamicalSystemsBase, Test
+
+function oop(u, p, t)
+    return p[1] * SVector(u[1], -u[2])
+end
+
+function iip(du, u, p, t)
+    du .= oop(u, p, t)
+    return nothing
+end
+
+#%%
+@testset "IDT=$(IDT), IIP=$(IIP)" for IDT in (true, false), IIP in (false, true)
+    SystemType = IDT ? DeterministicIteratedMap : CoupledODEs
+    rule = IIP ? iip : oop
+    p = 3.0
+    u0 = [1.0, 1.0]
+    result = [p 0.0; 0.0 -p]
+
+    ds = SystemType(rule, u0, p)
+    J0 = zeros(dimension(ds), dimension(ds))
+    J = jacobian(ds)
+    if IIP
+        J(J0, current_state(ds), current_parameters(ds), 0.0)
+        @test J0 == result
+    else
+        @test J(current_state(ds), current_parameters(ds), 0.0) == result
+    end
+end

test/runtests.jl

@@ -16,4 +16,5 @@ testfile(file, testname=defaultname(file)) = @testset "$testname" begin; include
     testfile("successful_step.jl")
     testfile("mtk_integration.jl")
     testfile("trajectory.jl")
+    testfile("jacobian.jl")
 end