Making autodiff Jacobian reusable (#211)

* jacobian code WIP

* Separating Jacobian from tangent_rule

* inserting new jacobian back into tangent_rule

* tests

* import not needed

* docstring

* changelog and version change

Author: JonasKoziorek

CHANGELOG.md

@@ -1,3 +1,8 @@
+# v3.9.0
+
+A new function `jacobian` that generates Jacobian rule for any `ds<:CoreDynamicalSystem` via 
+automatic differentiation with `ForwardDiff.jl`.
+
 # v3.8.0
 
 Even more amazing level of integration with ModelingToolkit.jl!

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicalSystemsBase"
 uuid = "6e36e845-645a-534a-86f2-f5d4aa5a06b4"
 repo = "https://github.com/JuliaDynamics/DynamicalSystemsBase.jl.git"
-version = "3.8.3"
+version = "3.9.0"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

docs/src/index.md

@@ -74,6 +74,7 @@ CoreDynamicalSystem
 TangentDynamicalSystem
 current_deviations
 set_deviations!
+jacobian
 orthonormal
 ```
 

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,46 @@
+export jacobian
+
+import ForwardDiff
+
+"""
+    jacobian(ds::CoreDynamicalSystem)
+
+Construct the Jacobian rule for the dynamical system `ds`.
+This is done via automatic differentiation using module 
+[`ForwardDiff`](https://github.com/JuliaDiff/ForwardDiff.jl).
+
+## Description
+
+For out-of-place systems, `jacobian` returns the Jacobian rule as a 
+function `Jf(u, p, t) -> J0::SMatrix`. Calling `Jf(u, p, t)` will compute the Jacobian 
+at the state `u`, parameters `p` and time `t` and return the result as `J0`.
+For in-place systems, `jacobian` returns the Jacobian rule as a function 
+`Jf!(J0, u, p, t)`. Calling `Jf!(J0, u, p, t)` will compute the Jacobian 
+at the state `u`, parameters `p` and time `t` and save the result in `J0`.
+"""
+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
@@ -137,7 +138,6 @@ correct_matrix_type(::Val{true}, Q::AbstractMatrix) = ismutable(Q) ? Q : Array(Q
 ###########################################################################################
 # Creation of tangent rule
 ###########################################################################################
-import ForwardDiff
 # IIP Tangent space dynamics
 function tangent_rule(f::F, J::JAC, J0, ::Val{true}, ::Val{k}, u0) where {F, JAC, k}
     tangentf = (du, u, p, t) -> begin
@@ -149,37 +149,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