feat: MTK Jacobian for CoupledDEs (#229)

* feat: MTK jacobian for CoupledDEs

* add CoupledSDE to CoreDynamicalSystem

* implement comments

* add docs

* don't use `dynamical_rule` but prob

* make `dynamic_rule` always take the nested f

* simplify `dynamic_rule` to not directly access nested function f

Author: oameye

ext/src/CoupledSDEs.jl

@@ -118,7 +118,10 @@ function DynamicalSystemsBase.CoupledSDEs(
     noise_process=nothing,
     seed=UInt64(0)
 )
-    return CoupledSDEs(dynamic_rule(ds), current_state(ds), p;
+    prob = referrenced_sciml_prob(ds)
+    # we want the symbolic jacobian to be transfered over
+    # dynamic_rule(ds) takes the deepest nested f wich does not have a jac field
+    return CoupledSDEs(prob.f, current_state(ds), p;
         g, noise_strength, covariance, diffeq, noise_prototype, noise_process, seed)
 end
 
@@ -130,9 +133,11 @@ deterministic part of `ds`.
 """
 function DynamicalSystemsBase.CoupledODEs(
     sys::CoupledSDEs; diffeq=DEFAULT_DIFFEQ, t0=0.0)
+    prob = referrenced_sciml_prob(sys)
+    # we want the symbolic jacobian to be transfered over
+    # dynamic_rule(ds) takes the deepest nested f wich does not have a jac field
     return CoupledODEs(
-        dynamic_rule(sys), SVector{length(sys.integ.u)}(sys.integ.u), sys.p0;
-        diffeq=diffeq, t0=t0
+        prob.f, SVector{length(sys.integ.u)}(sys.integ.u), sys.p0; diffeq=diffeq, t0=t0
     )
 end
 
@@ -155,16 +160,6 @@ StateSpaceSets.dimension(::CoupledSDEs{IIP,D}) where {IIP,D} = D
 DynamicalSystemsBase.current_state(ds::CoupledSDEs) = current_state(ds.integ)
 DynamicalSystemsBase.isdeterministic(ds::CoupledSDEs) = false
 
-function DynamicalSystemsBase.dynamic_rule(ds::CoupledSDEs)
-    # with remake it can happen that we have nested SDEFunctions
-    # sciml integrator deals with this internally well
-    f = ds.integ.f
-    while hasfield(typeof(f), :f)
-        f = f.f
-    end
-    return f
-end
-
 function DynamicalSystemsBase.set_state!(ds::CoupledSDEs, u::AbstractArray)
     (set_state!(ds.integ, u); ds)
 end
@@ -224,6 +219,8 @@ function covariance_matrix(ds::CoupledSDEs)::AbstractMatrix
     (A == nothing) ? nothing : A * A'
 end
 
+jacobian(sde::CoupledSDEs) = DynamicalSystemsBase.jacobian(CoupledODEs(sde))
+
 ###########################################################################################
 # Utilities
 ###########################################################################################

src/core_systems/jacobian.jl

@@ -6,7 +6,8 @@ import ForwardDiff
     jacobian(ds::CoreDynamicalSystem)
 
 Construct the Jacobian rule for the dynamical system `ds`.
-This is done via automatic differentiation using module
+If the system already has a Jacobian rule constructed via ModelingToolkit it returns this,
+otherwise it constructs the Jacobian rule with automatic differentiation using module
 [`ForwardDiff`](https://github.com/JuliaDiff/ForwardDiff.jl).
 
 ## Description
@@ -19,7 +20,17 @@ For in-place systems, `jacobian` returns the Jacobian rule as a function
 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}())
+    if ds isa ContinuousTimeDynamicalSystem
+        prob = referrenced_sciml_prob(ds)
+        if prob.f isa SciMLBase.AbstractDiffEqFunction && !isnothing(prob.f.jac)
+            jac = prob.f.jac
+        else
+            jac = _jacobian(ds, Val{IIP}())
+        end
+    else
+        jac = _jacobian(ds, Val{IIP}())
+    end
+    return jac
 end
 
 function _jacobian(ds, ::Val{true})
@@ -43,4 +54,6 @@ function _jacobian(ds, ::Val{false})
     f = dynamic_rule(ds)
     Jf = (u, p, t = 0) -> ForwardDiff.jacobian((x) -> f(x, p, t), u)
     return Jf
-end
+end
+
+jacobian(ds::CoupledSDEs) = jacobian(CoupledODEs(ds))

test/jacobian.jl

@@ -9,7 +9,7 @@ function iip(du, 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
@@ -26,4 +26,41 @@ end
     else
         @test J(current_state(ds), current_parameters(ds), 0.0) == result
     end
-end
+end
+
+@testset "MTK Jacobian" begin
+    using ModelingToolkit
+    using ModelingToolkit: Num, RuntimeGeneratedFunctions.RuntimeGeneratedFunction
+    using DynamicalSystemsBase: SciMLBase
+    @independent_variables t
+    @variables u(t)[1:2]
+    D = Differential(t)
+
+    eqs = [D(u[1]) ~ 3.0 * u[1],
+           D(u[2]) ~ -3.0 * u[2]]
+    @named sys = ODESystem(eqs, t)
+    sys = structural_simplify(sys)
+
+    prob = ODEProblem(sys, [1.0, 1.0], (0.0, 1.0); jac=true)
+    ode = CoupledODEs(prob)
+
+    jac = jacobian(ode)
+    @test jac.jac_oop isa RuntimeGeneratedFunction
+    @test jac([1.0, 1.0], [], 0.0) == [3 0;0 -3]
+
+    @testset "CoupledSDEs" begin
+        # just to check if MTK @brownian does not give any problems
+        using StochasticDiffEq
+        @brownian β
+        eqs = [D(u[1]) ~ 3.0 * u[1]+ β,
+                D(u[2]) ~ -3.0 * u[2] + β]
+        @mtkbuild sys = System(eqs, t)
+
+        prob = SDEProblem(sys, [1.0, 1.0], (0.0, 1.0), jac=true)
+        sde = CoupledSDEs(prob)
+
+        jac = jacobian(sde)
+        @test jac.jac_oop isa RuntimeGeneratedFunction
+        @test jac([1.0, 1.0], [], 0.0) == [3 0;0 -3]
+    end
+end

test/stochastic.jl

@@ -81,12 +81,12 @@ end
 
     # CoupledODEs creation
     ds = CoupledODEs(lorenz_oop)
-    @test dynamic_rule(ds) == lorenz_rule
+    @test dynamic_rule(ds).f == lorenz_rule
     @test ds.integ.alg isa Tsit5
     test_dynamical_system(ds, u0, p0; idt = false, iip = false)
     # and back
     sde = CoupledSDEs(ds, p0)
-    @test dynamic_rule(sde) == lorenz_rule
+    @test dynamic_rule(sde).f.f == lorenz_rule
     @test sde.integ.alg isa SOSRA
 end