feat: MTK Jacobian for CoupledDEs

ext/src/CoupledSDEs.jl

@@ -224,6 +224,8 @@ function covariance_matrix(ds::CoupledSDEs)::AbstractMatrix
     (A == nothing) ? nothing : A * A'
 end
 
+jacobian(sde::CoupledSDEs) = DynamicalSystemsBase.jacobian(CoupledODEs(sde))
+
 ###########################################################################################
 # Utilities
 ###########################################################################################

src/core_systems/additional_supertypes.jl

@@ -9,4 +9,4 @@ which are the core systems whose dynamic rule `f` is known analytically.
 This type is used for deciding whether a creation of a [`TangentDynamicalSystem`](@ref)
 is possible or not.
 """
-CoreDynamicalSystem{IIP} = Union{CoupledODEs{IIP}, DeterministicIteratedMap{IIP}}
+CoreDynamicalSystem{IIP} = Union{CoupledSDEs{IIP}, CoupledODEs{IIP}, DeterministicIteratedMap{IIP}}

src/core_systems/jacobian.jl

@@ -6,20 +6,26 @@ import ForwardDiff
     jacobian(ds::CoreDynamicalSystem)
 
 Construct the Jacobian rule for the dynamical system `ds`.
-This is done via automatic differentiation using module 
+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 
+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 
+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}())
+    if hasproperty(ds, :integ) &&
+            ds.integ.f isa SciMLBase.AbstractDiffEqFunction && !isnothing(ds.integ.f.jac)
+        jac = ds.integ.f.jac
+    else
+        jac = _jacobian(ds, Val{IIP}())
+    end
+    return jac
 end
 
 function _jacobian(ds, ::Val{true})
@@ -43,4 +49,4 @@ function _jacobian(ds, ::Val{false})
     f = dynamic_rule(ds)
     Jf = (u, p, t) -> ForwardDiff.jacobian((x) -> f(x, p, t), u)
     return Jf
-end
+end

test/jacobian.jl

@@ -26,4 +26,52 @@ 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)
+    p = 3.0
+
+    eqs = [D(u[1]) ~ p * u[1],
+        D(u[2]) ~ -p * u[2]]
+
+    @named sys = ODESystem(eqs, t)
+    sys = structural_simplify(sys)
+
+    jac = calculate_jacobian(sys)
+    @test jac isa Matrix{Num}
+
+    prob = ODEProblem(sys, [1.0, 1.0], (0.0, 1.0); jac=true)
+    ode = CoupledODEs(prob)
+    @test ode.integ.f.jac.jac_oop isa RuntimeGeneratedFunction
+    @test ode.integ.f.jac([1.0, 1.0], [3.0], 0.0) isa Matrix{Float64}
+
+    jac = jacobian(ode)
+    @test jac.jac_oop isa RuntimeGeneratedFunction
+    @test jac([1.0, 1.0], [3.0], 0.0) isa Matrix{Float64}
+
+    @testset "CoupledSDEs" begin
+        using StochasticDiffEq
+        @brownian β
+        eqs = [D(u[1]) ~ p * u[1]+ β,
+                D(u[2]) ~ -p * u[2] + β]
+        @mtkbuild sys = System(eqs, t)
+
+        jac = calculate_jacobian(sys)
+        @test jac isa Matrix{Num}
+
+        prob = SDEProblem(sys, [1.0, 1.0], (0.0, 1.0), jac=true)
+        sde = CoupledSDEs(prob)
+        @test sde.integ.f.jac.jac_oop isa RuntimeGeneratedFunction
+        @test sde.integ.f.jac([1.0, 1.0], [3.0], 0.0) isa Matrix{Float64}
+
+        jac = jacobian(ode)
+        @test jac.jac_oop isa RuntimeGeneratedFunction
+        @test jac([1.0, 1.0], [3.0], 0.0) isa Matrix{Float64}
+    end
+end