From 30d09dea32c47a8b00a202f5a26f348d80962280 Mon Sep 17 00:00:00 2001 From: Orjan Ameye Date: Mon, 2 Dec 2024 11:30:52 +0100 Subject: [PATCH] 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 --- ext/src/CoupledSDEs.jl | 23 +++++++++----------- src/core_systems/jacobian.jl | 19 ++++++++++++++--- test/jacobian.jl | 41 ++++++++++++++++++++++++++++++++++-- test/stochastic.jl | 4 ++-- 4 files changed, 67 insertions(+), 20 deletions(-) diff --git a/ext/src/CoupledSDEs.jl b/ext/src/CoupledSDEs.jl index 93b0b47c..55b595f2 100644 --- a/ext/src/CoupledSDEs.jl +++ b/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 ########################################################################################### diff --git a/src/core_systems/jacobian.jl b/src/core_systems/jacobian.jl index c443f25c..41958750 100644 --- a/src/core_systems/jacobian.jl +++ b/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 \ No newline at end of file +end + +jacobian(ds::CoupledSDEs) = jacobian(CoupledODEs(ds)) diff --git a/test/jacobian.jl b/test/jacobian.jl index a36869d8..bf1c42ae 100644 --- a/test/jacobian.jl +++ b/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 \ No newline at end of file +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 diff --git a/test/stochastic.jl b/test/stochastic.jl index 0754a64c..8d2bebd0 100644 --- a/test/stochastic.jl +++ b/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