Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: MTK Jacobian for CoupledDEs #229

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: MTK Jacobian for CoupledDEs #229

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
65bd3c5
feat: MTK jacobian for CoupledDEs
oameye Nov 23, 2024
8fd4b0e
add CoupledSDE to CoreDynamicalSystem
oameye Nov 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions ext/src/CoupledSDEs.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,8 @@ function covariance_matrix(ds::CoupledSDEs)::AbstractMatrix
(A == nothing) ? nothing : A * A'
end

jacobian(sde::CoupledSDEs) = DynamicalSystemsBase.jacobian(CoupledODEs(sde))

###########################################################################################
# Utilities
###########################################################################################
Expand Down
2 changes: 1 addition & 1 deletion src/core_systems/additional_supertypes.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this change made? SDEs are not a core system. Tangent space makes no sense for SDEs, and the tangent space is the only reason for the "Core" hierarchy.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just want Jacobian to work for CoupledSDEs, I could also define jacobian(Union{CoupledSDEs, CoreDynamicalSystem})

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is better to define jacobian(sde::CoupledSDEs) = jacobian(CoupledODEs(sde)). Which is also more transparent. It makes it clear that you want the jacobian of the drift, and not of the noise function g which also has a jacobian.

20 changes: 13 additions & 7 deletions src/core_systems/jacobian.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand All @@ -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
50 changes: 49 additions & 1 deletion test/jacobian.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +51 to +52
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tests must ONLY use the public API. Not access internal fields that may or may not exist. You have to re-write the tests to use the jacobian function exclusively after you create the ode type.

@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)
Comment on lines +65 to +66
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is this test testing? calculate_jacobian comes from what package?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it comes from MTK

Copy link
Member

@Datseris Datseris Nov 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Then we shouldn't be testing it. We should only test functionality from our package. This is easier to maintain and also less confusing. So, the jacobian function and the matrix it gives, which is always a matrix of real numbers.

If you think this function from MTK needs additional testing it is much better to contribute these tests directly to MTK (although I doubt it).

@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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally speaking testing the type of things is not of much value. We test functionality and interfaces. Instead of testing whether something "is what it should be", test instead that it "does what it should do".

@test jac.jac_oop isa RuntimeGeneratedFunction
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ALl of these tests should instead test explicitly the analytic jacobian. That the matrix we get is exactly, numerically, the matrix we expect.

@test jac([1.0, 1.0], [3.0], 0.0) isa Matrix{Float64}
end
end
Loading