Skip to content

Dynamical integrators #545

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

Open
wants to merge 16 commits into
base: master
Choose a base branch
from
+478 −49
Open

Dynamical integrators #545

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
a6347c4
Allow Dynamical problem to have non diagonal noise
Aug 31, 2023
d078c4c
Allow Dynamical problem to have non diagonal noise
Aug 31, 2023
18c09a2
Add ABOBA, start implement OBABO, check to be done on weak order
Aug 31, 2023
e10ebfa
Implement OBABO, remains some error on weak order
Sep 1, 2023
6423413
change variable name
Sep 11, 2023
52beb4c
correct weak order
Sep 11, 2023
f43249b
Actually OBABO need 2 noise processes
Sep 18, 2023
9a80b47
update new SDEProblem interface
Oct 7, 2023
508edf7
tests are now passing
Oct 10, 2023
f836738
even more trajectories to pass the test
Oct 13, 2023
9783fbd
add a testset to test consistency between iip and oop implementations
frankschae Nov 19, 2023
f03a017
Change variable names and split test into interface and convergence t…
Nov 20, 2023
84f2e9e
Merge branch 'SciML:master' into dynamical_integrators
HadrienNU Nov 20, 2023
cb25c7c
move tests and add convergence test to buildkite queue
Nov 20, 2023
3ac8f46
Missing "
Nov 20, 2023
0b3baf1
remove info message about number of montecarlo iterations
Nov 20, 2023
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: 1 addition & 1 deletion src/StochasticDiffEq.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ end

export TauLeaping, CaoTauLeaping

export BAOAB
export BAOAB, ABOBA, OBABO

export StochasticDiffEqRODEAlgorithm, StochasticDiffEqRODEAdaptiveAlgorithm,
StochasticDiffEqRODECompositeAlgorithm
Expand Down
5 changes: 5 additions & 0 deletions src/alg_utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,8 @@ alg_order(alg::TauLeaping) = 1 // 1
alg_order(alg::CaoTauLeaping) = 1 // 1

alg_order(alg::BAOAB) = 1 // 1
alg_order(alg::ABOBA) = 2 // 1
alg_order(alg::OBABO) = 2 // 1

alg_order(alg::SKenCarp) = 2 // 1
alg_order(alg::Union{StochasticDiffEqCompositeAlgorithm,StochasticDiffEqRODECompositeAlgorithm}) = maximum(alg_order.(alg.algs))
Expand Down Expand Up @@ -225,6 +227,8 @@ alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::IIF1M) = true
alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::IIF2M) = true
alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::Union{StochasticDiffEqCompositeAlgorithm,StochasticDiffEqRODECompositeAlgorithm}) = max((alg_compatible(prob, a) for a in alg.algs)...)
alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::BAOAB) = is_diagonal_noise(prob)
alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::ABOBA) = is_diagonal_noise(prob)
alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::OBABO) = is_diagonal_noise(prob)

function alg_compatible(prob::JumpProblem, alg::Union{StochasticDiffEqJumpAdaptiveAlgorithm,StochasticDiffEqJumpAlgorithm})
prob.prob isa DiscreteProblem
Expand Down Expand Up @@ -259,6 +263,7 @@ alg_needs_extra_process(alg::RS2) = true
alg_needs_extra_process(alg::PL1WM) = true
alg_needs_extra_process(alg::NON) = true
alg_needs_extra_process(alg::NON2) = true
alg_needs_extra_process(alg::OBABO) = true

OrdinaryDiffEq._alg_autodiff(alg::StochasticDiffEqNewtonAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = Val{AD}()
OrdinaryDiffEq._alg_autodiff(alg::StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = Val{AD}()
Expand Down
13 changes: 13 additions & 0 deletions src/algorithms.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -869,3 +869,16 @@ struct BAOAB{T} <: StochasticDiffEqAlgorithm
scale_noise::Bool
end
BAOAB(;gamma=1.0, scale_noise=true) = BAOAB(gamma, scale_noise)

struct ABOBA{T} <: StochasticDiffEqAlgorithm
Copy link
Member

Choose a reason for hiding this comment

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

maybe add a docstring that contains the SDE form:
Screen Shot 2023-09-19 at 11 44 02 AM

(or to the SDE solver docs https://docs.sciml.ai/DiffEqDocs/stable/solvers/sde_solve/ given that it doesn't match our common diagonal/non-diagonal noise classification exactly?)

Copy link
Member

Choose a reason for hiding this comment

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

Or maybe it needs a different problem type?

Copy link
Author

Choose a reason for hiding this comment

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

There is already the DynamicalSDEProblem (was introduced by the previous implementation of BAOAB) that is supposed to be solved by those integrators

Copy link
Author

Choose a reason for hiding this comment

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

Should I add some alg_compatible() functions to allow only DynamicalSDEProblem to be solved by Dynamical integrators?

Copy link
Member

Choose a reason for hiding this comment

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

yes that would be helpful in order to gate it to the right form.

Copy link
Author

Choose a reason for hiding this comment

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

I just realize that the DynamicalSDEProblem constructor actually return a SDEProblem so alg_comptable functions won't be useful here

Copy link
Member

Choose a reason for hiding this comment

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

But it can dispatch on prob.problem_type

Copy link
Member

Choose a reason for hiding this comment

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

yeah I think we should use DynamicalSDEProblem and it's a good idea to use alg_compatible to make sure that only BAOAB, ABOBA, etc. can be used to solve it.

Copy link
Author

Choose a reason for hiding this comment

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

But it can dispatch on prob.problem_type

Not sure to understand, SDEProblem does not have a field problem_type.

gamma::T
scale_noise::Bool
end
ABOBA(;gamma=1.0, scale_noise=true) = ABOBA(gamma, scale_noise)


struct OBABO{T} <: StochasticDiffEqAlgorithm
gamma::T
scale_noise::Bool
end
OBABO(;gamma=1.0, scale_noise=true) = OBABO(gamma, scale_noise)
176 changes: 161 additions & 15 deletions src/caches/dynamical_caches.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,42 +1,188 @@
abstract type StochasticDynamicalEqConstantCache <: StochasticDiffEqConstantCache end # Pourquoi faire ça, Si c'est pour avoir une seul function de check dans initialize!
abstract type StochasticDynamicalEqMutableCache <: StochasticDiffEqMutableCache end

mutable struct BAOABConstantCache{uType,uEltypeNoUnits} <: StochasticDiffEqConstantCache

mutable struct BAOABConstantCache{uType,uEltypeNoUnits,uCoeffType, uCoeffMType} <: StochasticDynamicalEqConstantCache
k::uType
half::uEltypeNoUnits
c1::uEltypeNoUnits
c2::uEltypeNoUnits
c1::uCoeffType
c2::uCoeffMType
end
@cache struct BAOABCache{uType,uEltypeNoUnits,rateNoiseType,uTypeCombined} <: StochasticDiffEqMutableCache
@cache struct BAOABCache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType, uCoeffMType,uTypeCombined} <: StochasticDynamicalEqMutableCache
utmp::uType
dumid::uType
dutmp::uType
dunoise::uType
k::uType
gtmp::uType
noise::rateNoiseType
gtmp::rateNoiseType
noise::uType
half::uEltypeNoUnits
c1::uEltypeNoUnits
c2::uEltypeNoUnits
c1::uCoeffType
c2::uCoeffMType
tmp::uTypeCombined
end

function alg_cache(alg::BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{false}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
k = zero(rate_prototype.x[1])
c1 = exp(-alg.gamma*dt)
c2 = sqrt(1 - alg.scale_noise*c1^2) # if scale_noise == false, c2 = 1
BAOABConstantCache(k, uEltypeNoUnits(1//2), uEltypeNoUnits(c1), uEltypeNoUnits(c2))
if typeof(alg.gamma) <: AbstractMatrix
c1 = exp(-alg.gamma*dt)
Copy link
Member

Choose a reason for hiding this comment

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

c1 and c2 should be functions. They shouldn't be constants because they depend on dt. So if someone uses a non-uniform time-stepping, this gets wrong.

Copy link
Author

Choose a reason for hiding this comment

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

There are chosen as constant for efficiency in the constant dt case, as it is the most common case. Should I pass them as function anyway? Knowing that I have as plan for the future to implement the case of gamma being a function of the current state, in which case c1 and c2 would be function of both dt and gamma.
The alternative is to add to the documention (to be written) "Fixed timestepping only." as in the SplitODE solvers for example (https://docs.sciml.ai/DiffEqDocs/stable/solvers/split_ode_solve/)

Copy link
Member

Choose a reason for hiding this comment

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

Writing "fixed time-stepping only" might not be sufficient because I think we typically only distinguish between adaptively chosen time steps and fixed time-stepping schemes. In particular, there could be situations where one has a concrete idea when dt should be small and when it should be large apriori. For instance, Frank and Moritz proposed a rescaling of the time stepping of the stochastic process in https://projecteuclid.org/journals/electronic-journal-of-statistics/volume-11/issue-1/Bayesian-estimation-of-discretely-observed-multi-dimensional-diffusion-processes-using/10.1214/17-EJS1290.full conditioned on the observation times. And we allow this behavior for fixed time stepping schemes by passing tstops only to the solve call (https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/).

It also seems relatively cheap to compute c1 and c2 -- do you have an example in mind where this matters in practice?

I think

Knowing that I have as plan for the future to implement the case of gamma being a function of the current state, in which case c1 and c2 would be function of both dt and gamma.

is another very good reason to do it in this PR.

Copy link
Member

Choose a reason for hiding this comment

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

Generally symplectic integrators lose the symplectic property if you adapt the time steps. Even the ODE solver ones are only fixed time step. But indeed we might as well make it safe for future work.

Copy link
Author

Choose a reason for hiding this comment

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

On the efficiency question: If you have a matricial gamma, if you have to do a matrix exponential and a chlolesky at each timestep you increase the computation time by a few percent.

At least on my case, the typical use case of this type of dynamics is to infers it from larger system (molecular dynamics) that are too expensive to run for long trajectories and run the inferred dynamics for very long times, so even a few percent increase of runtime could lead to extra hours of computation.

But ok, I will implement those coefficients as function, I could always add some caching behavior latter one if this becomes a performance issue. To be honest, I have never really thought to the possibility of time-stepping (losing the symplectic property is certainly an issue in my case), thanks for the reference, I will look at it.

Copy link
Member

Choose a reason for hiding this comment

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

oh, then it might also be enough to check this and throw an error (if nobody would ever want to/should change it)? I think

  • fixing it and throwing an error (or documenting it very explicitly maybe) or
  • allowing to adapt the time steps
    both sound good to me.

Copy link
Member

Choose a reason for hiding this comment

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

On the efficiency question: If you have a matricial gamma, if you have to do a matrix exponential and a chlolesky at each timestep you increase the computation time by a few percent.

I can see this point. But the cholesky/exponential also needs to be recomputed for fixed-time steps when gamma depends on the state, right? (so then it only hurts a few percent in the case of large matrix-valued gamma and fixed dt -- I have no intuition if that's a common use case.)

Copy link
Author

Choose a reason for hiding this comment

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

That make sense, I will implement those as function, will do it this week.

c2 = cholesky(I - alg.scale_noise*c1*transpose(c1)).U# if scale_noise == false, c2 = 1
else
c1 = exp.(-alg.gamma*dt)
c2 = sqrt.(1 .- alg.scale_noise*c1.^2)# if scale_noise == false, c2 = 1
end
BAOABConstantCache(k, uEltypeNoUnits(1//2),c1, c2)
end

function alg_cache(alg::BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{true}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
dumid = zero(u.x[1])
dutmp = zero(u.x[1])
dunoise = zero(u.x[1])
utmp = zero(u.x[2])
k = zero(rate_prototype.x[1])

gtmp = zero(noise_rate_prototype)
noise = zero(rate_prototype.x[1])

half = uEltypeNoUnits(1//2)

if typeof(alg.gamma) <: AbstractMatrix
c1 = exp(-alg.gamma*dt)
c2 = cholesky(I - alg.scale_noise*c1*transpose(c1)).U# if scale_noise == false, c2 = 1
else
c1 = exp.(-alg.gamma*dt)
c2 = sqrt.(1 .- alg.scale_noise*c1.^2)# if scale_noise == false, c2 = 1
end

tmp = zero(u)

BAOABCache(utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c1, c2, tmp)
end



mutable struct ABOBAConstantCache{uType,uEltypeNoUnits, uCoeffType, uCoeffMType} <: StochasticDynamicalEqConstantCache
k::uType
half::uEltypeNoUnits
c₂::uCoeffType
σ::uCoeffMType
end
@cache struct ABOBACache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType, uCoeffMType,uTypeCombined} <: StochasticDynamicalEqMutableCache
utmp::uType
dumid::uType
dutmp::uType
dunoise::uType
k::uType
gtmp::rateNoiseType
noise::uType
half::uEltypeNoUnits
c₂::uCoeffType
σ::uCoeffMType
tmp::uTypeCombined
end

function alg_cache(alg::ABOBA,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{false}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
k = zero(rate_prototype.x[1])

if typeof(alg.gamma) <: AbstractMatrix
c₂ = exp(-alg.gamma*dt)
Copy link
Member

Choose a reason for hiding this comment

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

Why not label them c1 and c2 as in the other methods?

Copy link
Author

@HadrienNU HadrienNU Nov 20, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

Good points, no particular reason, I will change this. The coefficients will be c2 for the friction term and σ for the diffusion term

σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
else
c₂ = exp.(-alg.gamma*dt)
σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
end
# if scale_noise == false, c2 = 1
ABOBAConstantCache(k, uEltypeNoUnits(1//2), c₂, σ)
end

function alg_cache(alg::ABOBA,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{true}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
dutmp = zero(u.x[1])
dumid = zero(u.x[1])
dunoise = zero(u.x[1])
utmp = zero(u.x[2])
k = zero(rate_prototype.x[1])

gtmp = zero(rate_prototype.x[1])
gtmp = zero(noise_rate_prototype)
noise = zero(rate_prototype.x[1])

half = uEltypeNoUnits(1//2)
c1 = exp(-alg.gamma*dt)
c2 = sqrt(1 - alg.scale_noise*c1^2) # if scale_noise == false, c2 = 1

if typeof(alg.gamma) <: AbstractMatrix
c₂ = exp(-alg.gamma*dt)
σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
else
c₂ = exp.(-alg.gamma*dt)
σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
end

tmp = zero(u)

ABOBACache(utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c₂, σ, tmp)
end




mutable struct OBABOConstantCache{uType,rateNoiseType, uEltypeNoUnits, uCoeffType, uCoeffMType} <: StochasticDynamicalEqConstantCache
k::uType
gt::rateNoiseType
half::uEltypeNoUnits
c₂::uCoeffType
σ::uCoeffMType
end

@cache struct OBABOCache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType, uCoeffMType,uTypeCombined} <: StochasticDynamicalEqMutableCache
utmp::uType
dumid::uType
dutmp::uType
dunoise::uType
k::uType
gtmp::rateNoiseType
noise::uType
half::uEltypeNoUnits
c₂::uCoeffType
σ::uCoeffMType
tmp::uTypeCombined
end

function alg_cache(alg::OBABO,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{false}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
k = zero(rate_prototype.x[1])
gt = zero(noise_rate_prototype)
half=uEltypeNoUnits(1//2)

if typeof(alg.gamma) <: AbstractMatrix
c₂ = exp(-alg.gamma*half*dt)
σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
else
c₂ = exp.(-alg.gamma*half*dt)
σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
end
# if scale_noise == false, c2 = 1
OBABOConstantCache(k, gt, half, c₂, σ)
end

function alg_cache(alg::OBABO,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{true}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
dutmp = zero(u.x[1])
dumid = zero(u.x[1])
dunoise = zero(u.x[1])
utmp = zero(u.x[2])
k = zero(rate_prototype.x[1])

gtmp = zero(noise_rate_prototype)
noise = zero(rate_prototype.x[1])


half = uEltypeNoUnits(1//2)

if typeof(alg.gamma) <: AbstractMatrix
c₂ = exp(-alg.gamma*half*dt)
σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
else
c₂ = exp.(-alg.gamma*half*dt)
σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
end

tmp = zero(u)

BAOABCache(utmp, dutmp, k, gtmp, noise, half, uEltypeNoUnits(c1), uEltypeNoUnits(c2), tmp)
OBABOCache(utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c₂, σ, tmp)
end
Loading