From 0891eda3342e227bca66a4861533d2364da5d573 Mon Sep 17 00:00:00 2001 From: Orjan Ameye Date: Thu, 14 May 2026 08:59:08 +0200 Subject: [PATCH 1/2] move to smaller sublibrary of StochasticDiffEq and move away from extension --- Project.toml | 11 +- docs/make.jl | 6 +- ext/StochasticSystemsBase.jl | 19 - ext/src/CoupledSDEs.jl | 315 ---------------- ext/src/classification.jl | 145 -------- src/core_systems/continuous_time_sde.jl | 460 ++++++++++++++++++++++++ test/stochastic.jl | 6 +- 7 files changed, 466 insertions(+), 496 deletions(-) delete mode 100644 ext/StochasticSystemsBase.jl delete mode 100644 ext/src/CoupledSDEs.jl delete mode 100644 ext/src/classification.jl diff --git a/Project.toml b/Project.toml index e17642af..f34b8af7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "DynamicalSystemsBase" uuid = "6e36e845-645a-534a-86f2-f5d4aa5a06b4" repo = "https://github.com/JuliaDynamics/DynamicalSystemsBase.jl.git" -version = "3.16.0" +version = "3.17.0" [deps] ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -14,14 +14,9 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StateSpaceSets = "40b095a5-5852-4c12-98c7-d43bf788e795" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +StochasticDiffEqHighOrder = "0520c28c-50fd-4d16-9c96-902fc80b3bab" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" -[weakdeps] -StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" - -[extensions] -StochasticSystemsBase = "StochasticDiffEq" - [compat] ForwardDiff = "0.10, 1" OrdinaryDiffEqTsit5 = "2" @@ -30,6 +25,6 @@ Roots = "2, 3" SciMLBase = "3" StateSpaceSets = "2.5" Statistics = "1" -StochasticDiffEq = "7" +StochasticDiffEqHighOrder = "2" SymbolicIndexingInterface = "0.3.4" julia = "1.10" diff --git a/docs/make.jl b/docs/make.jl index 5bb56f06..21ddebfd 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,9 +1,7 @@ cd(@__DIR__) using DynamicalSystemsBase -using StochasticDiffEq, DiffEqNoiseProcess # to enable extention -# We need this because Documenter doesn't know where to get the docstring from otherwise -StochasticSystemsBase = Base.get_extension(DynamicalSystemsBase, :StochasticSystemsBase) +using StochasticDiffEq, DiffEqNoiseProcess pages = [ "index.md", @@ -20,5 +18,5 @@ include("build_docs_with_style.jl") build_docs_with_style( pages, - DynamicalSystemsBase, SciMLBase, StateSpaceSets, StochasticSystemsBase + DynamicalSystemsBase, SciMLBase, StateSpaceSets ) diff --git a/ext/StochasticSystemsBase.jl b/ext/StochasticSystemsBase.jl deleted file mode 100644 index 8ccd9570..00000000 --- a/ext/StochasticSystemsBase.jl +++ /dev/null @@ -1,19 +0,0 @@ -module StochasticSystemsBase - -using DynamicalSystemsBase: DynamicalSystemsBase, SciMLBase, correct_state, CoupledODEs, - CoupledSDEs, StateSpaceSets, isinplace, _delete, set_parameter!, - set_parameters!, set_state!, dynamic_rule, isdeterministic, current_state, - DynamicalSystemsBase, _set_parameter!, derivative_discontinuity!, - additional_details, referrenced_sciml_prob, DEFAULT_DIFFEQ, - SVector, SMatrix, current_parameters, initial_state, initial_time -using SciMLBase: SDEProblem, AbstractSDEIntegrator, __init, SDEFunction, step! -using StochasticDiffEq: SOSRA, setup_next_step! -using LinearAlgebra -import Random - -include("src/CoupledSDEs.jl") -include("src/classification.jl") - -export CoupledSDEs, CoupledODEs - -end diff --git a/ext/src/CoupledSDEs.jl b/ext/src/CoupledSDEs.jl deleted file mode 100644 index 8e31ecf3..00000000 --- a/ext/src/CoupledSDEs.jl +++ /dev/null @@ -1,315 +0,0 @@ -using LinearAlgebra: LinearAlgebra -import DynamicalSystemsBase: CoupledSDEs - -########################################################################################### -# DiffEq options -########################################################################################### -# SOSRA is only applicable for additive noise and must be adaptive -# default sciml tolerance is 1e-2 -const DEFAULT_SDE_SOLVER = SOSRA() -const DEFAULT_STOCH_DIFFEQ_KWARGS = (abstol = 1.0e-2, reltol = 1.0e-2, dt = 0.1) -const DEFAULT_STOCH_DIFFEQ = (alg = DEFAULT_SDE_SOLVER, DEFAULT_STOCH_DIFFEQ_KWARGS...) - -# Function from user `@xlxs4`, see -# https://github.com/JuliaDynamics/jl/pull/153 -function _decompose_into_sde_solver_and_remaining(diffeq) - if haskey(diffeq, :alg) - return (diffeq[:alg], _delete(diffeq, :alg)) - else - return (DEFAULT_SDE_SOLVER, diffeq) - end -end - -########################################################################################### -# Constructor functions -########################################################################################### - -function DynamicalSystemsBase.CoupledSDEs( - f, - u0, - p = SciMLBase.NullParameters(); - g = nothing, - noise_strength = 1.0, - covariance = nothing, - t0 = 0.0, - diffeq = DEFAULT_STOCH_DIFFEQ, - noise_prototype = nothing, - noise_process = nothing, - seed = rand(UInt64) - ) - IIP = isinplace(f, 4) # from SciMLBase - if !isnothing(g) - @assert IIP == isinplace(g, 4) "f and g must both be in-place or out-of-place" - end - - noise_type, cov = find_noise_type(g, u0, p, t0, noise_process, covariance, noise_prototype, IIP) - g, noise_prototype = construct_diffusion_function( - g, cov, noise_prototype, noise_strength, length(u0), IIP - ) - - s = correct_state(Val{IIP}(), u0) - T = eltype(s) - prob = SDEProblem{IIP}( - f, - g, - s, - (T(t0), T(1.0e11)), - p; - noise_rate_prototype = noise_prototype, - noise = noise_process, - seed = seed - ) - return CoupledSDEs(prob, diffeq, noise_type) -end - -function DynamicalSystemsBase.CoupledSDEs( - prob::SDEProblem, diffeq = DEFAULT_STOCH_DIFFEQ, noise_type = nothing; special_kwargs... - ) - if haskey(special_kwargs, :diffeq) - throw( - ArgumentError( - "`diffeq` is given as positional argument when an SDEProblem is provided." - ), - ) - end - IIP = isinplace(prob) # from SciMLBase - D = length(prob.u0) - P = typeof(prob.p) - if prob.tspan === (nothing, nothing) - # If the problem was made via MTK, it is possible to not have a default timespan. - U = eltype(prob.u0) - prob = SciMLBase.remake(prob; tspan = (U(0), U(Inf))) - end - if isnothing(noise_type) - noise_type, _ = find_noise_type(prob, IIP) - end - - solver, remaining = _decompose_into_sde_solver_and_remaining(diffeq) - # The default `dtmin` from SciML scales with `tspan`. With our open-ended - # `tspan = (0, 1e11)` it becomes ~1e-5, which is too coarse for the SDE - # adaptive controller and causes spurious `DtLessThanMin` aborts. - remaining = haskey(remaining, :dtmin) ? remaining : merge((dtmin = 0.0,), remaining) - integ = __init( - prob, - solver; - remaining..., - # Integrators are used exclusively iteratively. There is no reason to save anything. - save_start = false, - save_end = false, - save_everystep = false, - # DynamicalSystems.jl operates on integrators and `step!` exclusively, - # so there is no reason to limit the maximum time evolution - maxiters = Inf - ) - return CoupledSDEs{IIP, D, typeof(integ), P}( - integ, deepcopy(prob.p), diffeq, noise_type - ) -end -# This preserves the referrenced MTK system and the originally passed diffeq kwargs -CoupledSDEs(ds::CoupledSDEs, diffeq) = CoupledSDEs(SDEProblem(ds), merge(ds.diffeq, diffeq)) - -""" - reinit!(ds::CoupledSDEs, u = initial_state(ds); kwargs...) → ds - -Re-initialize a [`CoupledSDEs`](@ref). In addition to the keywords accepted by -`reinit!` for any [`DynamicalSystem`](@ref), the following is supported: - -- `seed::UInt64`: re-seed the noise process random number generator. Defaults to - `rand(UInt64)`, so by default every `reinit!` produces a fresh, independent noise - realization. Pass an explicit `seed` to obtain a reproducible noise stream. -""" -function SciMLBase.reinit!( - ds::CoupledSDEs, u::AbstractArray = DynamicalSystemsBase.initial_state(ds); - p = DynamicalSystemsBase.current_parameters(ds), - t0 = DynamicalSystemsBase.initial_time(ds), - seed::UInt64 = rand(UInt64), kw... - ) - DynamicalSystemsBase.set_parameters!(ds, p) - # Adaptive solvers carry the previous run's last `dt`; resetting it ensures - # `seed` reproducibility. Non-adaptive solvers must keep their user-specified `dt`. - reset_dt = ds.integ.opts.adaptive - SciMLBase.reinit!(ds.integ, u; reinit_dae = false, reset_dt, t0, kw...) - Random.seed!(ds.integ.W.rng, seed) - setup_next_step!(ds.integ.W, ds.integ.u, ds.integ.p) - return ds -end - -""" - CoupledSDEs(ds::CoupledODEs, p = current_parameters(ds); kwargs...) - -Converts a [`CoupledODEs`](@ref) into a [`CoupledSDEs`](@ref). -While `p` is optional, it is likely to change as the -diffusion (noise) function `g` is added. -""" -function DynamicalSystemsBase.CoupledSDEs( - ds::CoupledODEs, - p = current_parameters(ds); # the parameter is likely changed as the diffusion function g is added. - g = nothing, - noise_strength = 1.0, - covariance = nothing, - diffeq = DEFAULT_STOCH_DIFFEQ, - noise_prototype = nothing, - noise_process = nothing, - seed = rand(UInt64) - ) - 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 - -""" - CoupledODEs(ds::CoupledSDEs; kwargs...) - -Converts a [`CoupledSDEs`](@ref) into a [`CoupledODEs`](@ref) by extracting the -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( - prob.f, SVector{length(sys.integ.u)}(sys.integ.u), sys.p0; diffeq = diffeq, t0 = t0 - ) -end - -# Pretty print -function DynamicalSystemsBase.additional_details(ds::CoupledSDEs) - solver, remaining = _decompose_into_sde_solver_and_remaining(ds.diffeq) - return [ - "SDE solver" => string(nameof(typeof(solver))), - "SDE kwargs" => remaining, - "Noise type" => ds.noise_type, - ] -end - -########################################################################################### -# API - obtaining information from the system -########################################################################################### - -SciMLBase.isinplace(::CoupledSDEs{IIP}) where {IIP} = IIP -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.set_state!(ds::CoupledSDEs, u::AbstractArray) - (set_state!(ds.integ, u); ds) -end -SciMLBase.step!(ds::CoupledSDEs, args...) = (step!(ds.integ, args...); ds) - -# TODO We have to check if for SDEIntegrators a different step ReturnCode is possible. -function DynamicalSystemsBase.successful_step(integ::AbstractSDEIntegrator) - rcode = integ.sol.retcode - return rcode == SciMLBase.ReturnCode.Default || rcode == SciMLBase.ReturnCode.Success -end - -# This is here to ensure that `derivative_discontinuity!` is called -function DynamicalSystemsBase.set_parameter!(ds::CoupledSDEs, args...) - _set_parameter!(ds, args...) - derivative_discontinuity!(ds.integ, true) - return nothing -end - -DynamicalSystemsBase.referrenced_sciml_prob(ds::CoupledSDEs) = ds.integ.sol.prob - -""" - diffusion_matrix(ds::CoupledSDEs) - -Returns the diffusion matrix of the stochastic term of the [`CoupledSDEs`](@ref) `ds`, -provided that the diffusion function `g` can be expressed as a constant invertible matrix. -If this is not the case, returns `nothing`. - -Note: The diffusion matrix ``Σ`` is the square root of the noise covariance matrix ``Q`` (see -[`covariance_matrix`](@ref)), defined via the Cholesky decomposition ``Q = Σ Σ^\\top``. -""" -function diffusion_matrix(ds::CoupledSDEs{IIP, D})::AbstractMatrix where {IIP, D} - if ds.noise_type[:invertible] - diffusion = diffusion_function(ds) - A = diffusion(zeros(D), current_parameters(ds), 0.0) - A = A isa AbstractMatrix ? A : Diagonal(A) - else - @warn """ - The diffusion function of the `CoupledSDEs` cannot be expressed as a constant - invertible matrix. - """ - A = nothing - end - return A -end - -""" - covariance_matrix(ds::CoupledSDEs) - -Returns the covariance matrix of the stochastic term of the [`CoupledSDEs`](@ref) `ds`, -provided that the diffusion function `g` can be expressed as a constant invertible matrix. -If this is not the case, returns `nothing`. - -See also [`diffusion_matrix`](@ref). -""" -function covariance_matrix(ds::CoupledSDEs)::AbstractMatrix - A = diffusion_matrix(ds) - return (A == nothing) ? nothing : A * A' -end - -jacobian(sde::CoupledSDEs) = DynamicalSystemsBase.jacobian(CoupledODEs(sde)) - -########################################################################################### -# Utilities -########################################################################################### - -""" - construct_diffusion_function(g, covariance, noise_prototype, noise_strength, D, IIP) - -Constructs the noise function `g` based on the keyword arguments -`g`, `covariance`, `noise_prototype`, and `noise_strength` -specified by the user when defining a `CoupledSDEs`. - -Here `D` is the system dimension and `IIP` indicated whether the function `g` is in-place -(`IIP = true`) or out-of-place (`false`). - -Returns `g, noise_prototype`. -""" -function construct_diffusion_function( - g, covariance, noise_prototype, noise_strength, D, IIP - ) - if isnothing(g) # diagonal additive noise - cov = isnothing(covariance) ? LinearAlgebra.I(D) : covariance - size(cov) != (D, D) && - throw(ArgumentError("Covariance matrix must be of size $((D, D))")) - A = sqrt(cov) - if IIP - if isdiag(cov) - diag_const = collect(noise_strength .* diag(A)) - g = let diag_const = diag_const - (du, u, p, t) -> (du .= diag_const; nothing) - end - else - A_const = collect(noise_strength .* A) - g = let A_const = A_const - (du, u, p, t) -> (du .= A_const; nothing) - end - noise_prototype = zeros(size(A)) - # ^ we could make this sparse to make it more performant - end - else - if isdiag(cov) - diag_const = SVector{D, eltype(A)}(diag(noise_strength .* A)) - g = let diag_const = diag_const - (u, p, t) -> diag_const - end - else - A_const = SMatrix{size(A)..., eltype(A)}(noise_strength .* A) - g = let A_const = A_const - (u, p, t) -> A_const - end - noise_prototype = zeros(size(A)) - end - end - end - return g, noise_prototype -end diff --git a/ext/src/classification.jl b/ext/src/classification.jl deleted file mode 100644 index 9b4d54b3..00000000 --- a/ext/src/classification.jl +++ /dev/null @@ -1,145 +0,0 @@ -using LinearAlgebra - -""" -Checks whether the function g depends on u based on 10 random points around the given u. -""" -function is_state_independent(g, u, p, t) - rdm_states = [u .+ rand(eltype(u), length(u)) .- 0.5 for _ in 1:10] - val = map(u -> g(u, p, t), rdm_states) - return length(unique(val)) == 1 -end - -""" -Checks whether g depends explicitly on time for select points on the interval [t0, t0+101]. -""" -function is_time_independent(g, u, p, t0) - trange = t0 .+ [0.0, 0.101, 1.01, 10.1, 101.0] - val = map(t -> g(u, p, t), trange) - return length(unique(val)) == 1 -end - -""" -Checks whether a matrix x is invertible by verifying that it has nonzero determinant. -""" -function is_invertible(x; tol = 1.0e-10) - F = lu(x, check = false) - det = abs(prod(diag(F.U))) - return det > tol -end - -""" -Checks whether the function f is linear by checking additivity and scalar multiplication -for two points x and y. -""" -function is_linear(f, x, y, c) - check1 = f(x + y) == f(x) + f(y) - check2 = f(c * x) == c * f(x) - return check1 && check2 -end - -function diffusion_function(g, IIP, noise_prototype) - return function diffusion(u, p, t) - if IIP - du = deepcopy(isnothing(noise_prototype) ? u : noise_prototype) - g(du, u, p, t) - return du - else - return g(u, p, t) - end - end -end - -function diffusion_function(ds::CoupledSDEs{IIP}) where {IIP} - prob = referrenced_sciml_prob(ds) - return diffusion_function(prob.g, IIP, prob.noise_rate_prototype) -end - -""" -Classifies the noise type of the CoupledSDEs given by the user. -Returns a named tuple of noise properties and the noise covariance matrix (if applicable). -""" -function find_noise_type(g, u0, p, t0, noise, covariance, noise_prototype, IIP) - noise_size = isnothing(noise_prototype) ? nothing : size(noise_prototype) - noise_cov = isnothing(noise) ? nothing : noise.covariance - D = length(u0) - - if !isnothing(noise_cov) - throw( - ArgumentError("CoupledSDEs does not support correlation between noise processes through DiffEqNoiseProcess.jl interface. Instead, use the `covariance` kwarg of `CoupledSDEs`.") - ) - end - - isadditive = false - isautonomous = false - islinear = false - isinvertible = false - - diffusion = diffusion_function(g, IIP, noise_prototype) - - if isnothing(g) - isadditive = true - isautonomous = true - islinear = true - if isnothing(covariance) - covariance = LinearAlgebra.I(D) - isinvertible = true - else - isinvertible = is_invertible(covariance) - end - elseif !isnothing(covariance) - throw( - ArgumentError("Both `g` and `covariance` are provided. Instead opt to encode the covariance in the difussion function `g` with the `noise_prototype` kwarg.") - ) - else - time_independent = is_time_independent(diffusion, rand(D), p, t0) - state_independent = is_state_independent(diffusion, u0, p, t0) - - # additive noise is equal to state independent noise - isadditive = state_independent - isautonomous = time_independent - - islinear = true - if !state_independent - for i in 1:10 - check = is_linear( - u -> diffusion(u, p, t0), - u0 + i .* rand(D), u0 + i .* rand(D), 2.0 - ) - check ? nothing : islinear = false - end - end - - # Previous formulation: - #islinear = !state_independent ? - # is_linear(u -> diffusion(u, p, t0), rand(D), rand(D), 2.0) : true - - if time_independent && state_independent - if !isnothing(noise_size) && isequal(noise_size...) - A = diffusion(zeros(D), p, 0.0) - covariance = A * A' - isinvertible = is_invertible(covariance) - elseif !isnothing(noise_size) && !isequal(noise_size...) - isinvertible = false - covariance = nothing - else - isinvertible = true - covariance = LinearAlgebra.I(D) - end - else - covariance = nothing - end - end - - noise_type = ( - additive = isadditive, autonomous = isautonomous, - linear = islinear, invertible = isinvertible, - ) - return noise_type, covariance -end - -function find_noise_type(prob::SDEProblem, IIP) - return find_noise_type( - prob.g, prob.u0, prob.p, prob.tspan[1], prob.noise, - nothing, prob.noise_rate_prototype, IIP - ) -end diff --git a/src/core_systems/continuous_time_sde.jl b/src/core_systems/continuous_time_sde.jl index 152e8882..b64d5ef9 100644 --- a/src/core_systems/continuous_time_sde.jl +++ b/src/core_systems/continuous_time_sde.jl @@ -1,5 +1,11 @@ export CoupledSDEs +using SciMLBase: SDEProblem, AbstractSDEIntegrator, SDEFunction, __init +using StochasticDiffEqHighOrder: SOSRA +using StochasticDiffEqHighOrder.StochasticDiffEqCore: setup_next_step! +using LinearAlgebra +import Random + """ CoupledSDEs <: ContinuousTimeDynamicalSystem CoupledSDEs(f, u0 [, p]; kwargs...) @@ -112,3 +118,457 @@ struct CoupledSDEs{IIP, D, I, P} <: ContinuousTimeDynamicalSystem diffeq # isn't parameterized because it is only used for display noise_type::NamedTuple end + +########################################################################################### +# DiffEq options +########################################################################################### +# SOSRA is only applicable for additive noise and must be adaptive +# default sciml tolerance is 1e-2 +const DEFAULT_SDE_SOLVER = SOSRA() +const DEFAULT_STOCH_DIFFEQ_KWARGS = (abstol = 1.0e-2, reltol = 1.0e-2, dt = 0.1) +const DEFAULT_STOCH_DIFFEQ = (alg = DEFAULT_SDE_SOLVER, DEFAULT_STOCH_DIFFEQ_KWARGS...) + +# Function from user `@xlxs4`, see +# https://github.com/JuliaDynamics/jl/pull/153 +function _decompose_into_sde_solver_and_remaining(diffeq) + if haskey(diffeq, :alg) + return (diffeq[:alg], _delete(diffeq, :alg)) + else + return (DEFAULT_SDE_SOLVER, diffeq) + end +end + +########################################################################################### +# Constructor functions +########################################################################################### + +function CoupledSDEs( + f, + u0, + p = SciMLBase.NullParameters(); + g = nothing, + noise_strength = 1.0, + covariance = nothing, + t0 = 0.0, + diffeq = DEFAULT_STOCH_DIFFEQ, + noise_prototype = nothing, + noise_process = nothing, + seed = rand(UInt64) + ) + IIP = isinplace(f, 4) # from SciMLBase + if !isnothing(g) + @assert IIP == isinplace(g, 4) "f and g must both be in-place or out-of-place" + end + + noise_type, cov = find_noise_type(g, u0, p, t0, noise_process, covariance, noise_prototype, IIP) + g, noise_prototype = construct_diffusion_function( + g, cov, noise_prototype, noise_strength, length(u0), IIP + ) + + s = correct_state(Val{IIP}(), u0) + T = eltype(s) + prob = SDEProblem{IIP}( + f, + g, + s, + (T(t0), T(1.0e11)), + p; + noise_rate_prototype = noise_prototype, + noise = noise_process, + seed = seed + ) + return CoupledSDEs(prob, diffeq, noise_type) +end + +function CoupledSDEs( + prob::SDEProblem, diffeq = DEFAULT_STOCH_DIFFEQ, noise_type = nothing; special_kwargs... + ) + if haskey(special_kwargs, :diffeq) + throw( + ArgumentError( + "`diffeq` is given as positional argument when an SDEProblem is provided." + ), + ) + end + IIP = isinplace(prob) # from SciMLBase + D = length(prob.u0) + P = typeof(prob.p) + if prob.tspan === (nothing, nothing) + # If the problem was made via MTK, it is possible to not have a default timespan. + U = eltype(prob.u0) + prob = SciMLBase.remake(prob; tspan = (U(0), U(Inf))) + end + if isnothing(noise_type) + noise_type, _ = find_noise_type(prob, IIP) + end + + solver, remaining = _decompose_into_sde_solver_and_remaining(diffeq) + # The default `dtmin` from SciML scales with `tspan`. With our open-ended + # `tspan = (0, 1e11)` it becomes ~1e-5, which is too coarse for the SDE + # adaptive controller and causes spurious `DtLessThanMin` aborts. + remaining = haskey(remaining, :dtmin) ? remaining : merge((dtmin = 0.0,), remaining) + integ = __init( + prob, + solver; + remaining..., + # Integrators are used exclusively iteratively. There is no reason to save anything. + save_start = false, + save_end = false, + save_everystep = false, + # DynamicalSystems.jl operates on integrators and `step!` exclusively, + # so there is no reason to limit the maximum time evolution + maxiters = Inf + ) + return CoupledSDEs{IIP, D, typeof(integ), P}( + integ, deepcopy(prob.p), diffeq, noise_type + ) +end +# This preserves the referrenced MTK system and the originally passed diffeq kwargs +CoupledSDEs(ds::CoupledSDEs, diffeq) = CoupledSDEs(SDEProblem(ds), merge(ds.diffeq, diffeq)) + +""" + reinit!(ds::CoupledSDEs, u = initial_state(ds); kwargs...) → ds + +Re-initialize a [`CoupledSDEs`](@ref). In addition to the keywords accepted by +`reinit!` for any [`DynamicalSystem`](@ref), the following is supported: + +- `seed::UInt64`: re-seed the noise process random number generator. Defaults to + `rand(UInt64)`, so by default every `reinit!` produces a fresh, independent noise + realization. Pass an explicit `seed` to obtain a reproducible noise stream. +""" +function SciMLBase.reinit!( + ds::CoupledSDEs, u::AbstractArray = initial_state(ds); + p = current_parameters(ds), + t0 = initial_time(ds), + seed::UInt64 = rand(UInt64), kw... + ) + set_parameters!(ds, p) + # Adaptive solvers carry the previous run's last `dt`; resetting it ensures + # `seed` reproducibility. Non-adaptive solvers must keep their user-specified `dt`. + reset_dt = ds.integ.opts.adaptive + SciMLBase.reinit!(ds.integ, u; reinit_dae = false, reset_dt, t0, kw...) + Random.seed!(ds.integ.W.rng, seed) + setup_next_step!(ds.integ.W, ds.integ.u, ds.integ.p) + return ds +end + +""" + CoupledSDEs(ds::CoupledODEs, p = current_parameters(ds); kwargs...) + +Converts a [`CoupledODEs`](@ref) into a [`CoupledSDEs`](@ref). +While `p` is optional, it is likely to change as the +diffusion (noise) function `g` is added. +""" +function CoupledSDEs( + ds::CoupledODEs, + p = current_parameters(ds); # the parameter is likely changed as the diffusion function g is added. + g = nothing, + noise_strength = 1.0, + covariance = nothing, + diffeq = DEFAULT_STOCH_DIFFEQ, + noise_prototype = nothing, + noise_process = nothing, + seed = rand(UInt64) + ) + 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 + +""" + CoupledODEs(ds::CoupledSDEs; kwargs...) + +Converts a [`CoupledSDEs`](@ref) into a [`CoupledODEs`](@ref) by extracting the +deterministic part of `ds`. +""" +function 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( + prob.f, SVector{length(sys.integ.u)}(sys.integ.u), sys.p0; diffeq = diffeq, t0 = t0 + ) +end + +# Pretty print +function additional_details(ds::CoupledSDEs) + solver, remaining = _decompose_into_sde_solver_and_remaining(ds.diffeq) + return [ + "SDE solver" => string(nameof(typeof(solver))), + "SDE kwargs" => remaining, + "Noise type" => ds.noise_type, + ] +end + +########################################################################################### +# API - obtaining information from the system +########################################################################################### + +SciMLBase.isinplace(::CoupledSDEs{IIP}) where {IIP} = IIP +StateSpaceSets.dimension(::CoupledSDEs{IIP, D}) where {IIP, D} = D +current_state(ds::CoupledSDEs) = current_state(ds.integ) +isdeterministic(ds::CoupledSDEs) = false + +set_state!(ds::CoupledSDEs, u::AbstractArray) = (set_state!(ds.integ, u); ds) +SciMLBase.step!(ds::CoupledSDEs, args...) = (step!(ds.integ, args...); ds) + +# TODO We have to check if for SDEIntegrators a different step ReturnCode is possible. +function successful_step(integ::AbstractSDEIntegrator) + rcode = integ.sol.retcode + return rcode == SciMLBase.ReturnCode.Default || rcode == SciMLBase.ReturnCode.Success +end + +# This is here to ensure that `derivative_discontinuity!` is called +function set_parameter!(ds::CoupledSDEs, args...) + _set_parameter!(ds, args...) + derivative_discontinuity!(ds.integ, true) + return nothing +end + +referrenced_sciml_prob(ds::CoupledSDEs) = ds.integ.sol.prob + +""" + diffusion_matrix(ds::CoupledSDEs) + +Returns the diffusion matrix of the stochastic term of the [`CoupledSDEs`](@ref) `ds`, +provided that the diffusion function `g` can be expressed as a constant invertible matrix. +If this is not the case, returns `nothing`. + +Note: The diffusion matrix ``Σ`` is the square root of the noise covariance matrix ``Q`` (see +[`covariance_matrix`](@ref)), defined via the Cholesky decomposition ``Q = Σ Σ^\\top``. +""" +function diffusion_matrix(ds::CoupledSDEs{IIP, D})::AbstractMatrix where {IIP, D} + if ds.noise_type[:invertible] + diffusion = diffusion_function(ds) + A = diffusion(zeros(D), current_parameters(ds), 0.0) + A = A isa AbstractMatrix ? A : Diagonal(A) + else + @warn """ + The diffusion function of the `CoupledSDEs` cannot be expressed as a constant + invertible matrix. + """ + A = nothing + end + return A +end + +""" + covariance_matrix(ds::CoupledSDEs) + +Returns the covariance matrix of the stochastic term of the [`CoupledSDEs`](@ref) `ds`, +provided that the diffusion function `g` can be expressed as a constant invertible matrix. +If this is not the case, returns `nothing`. + +See also [`diffusion_matrix`](@ref). +""" +function covariance_matrix(ds::CoupledSDEs)::AbstractMatrix + A = diffusion_matrix(ds) + return (A == nothing) ? nothing : A * A' +end + +########################################################################################### +# Utilities +########################################################################################### + +""" + construct_diffusion_function(g, covariance, noise_prototype, noise_strength, D, IIP) + +Constructs the noise function `g` based on the keyword arguments +`g`, `covariance`, `noise_prototype`, and `noise_strength` +specified by the user when defining a `CoupledSDEs`. + +Here `D` is the system dimension and `IIP` indicated whether the function `g` is in-place +(`IIP = true`) or out-of-place (`false`). + +Returns `g, noise_prototype`. +""" +function construct_diffusion_function( + g, covariance, noise_prototype, noise_strength, D, IIP + ) + if isnothing(g) # diagonal additive noise + cov = isnothing(covariance) ? LinearAlgebra.I(D) : covariance + size(cov) != (D, D) && + throw(ArgumentError("Covariance matrix must be of size $((D, D))")) + A = sqrt(cov) + if IIP + if isdiag(cov) + diag_const = collect(noise_strength .* diag(A)) + g = let diag_const = diag_const + (du, u, p, t) -> (du .= diag_const; nothing) + end + else + A_const = collect(noise_strength .* A) + g = let A_const = A_const + (du, u, p, t) -> (du .= A_const; nothing) + end + noise_prototype = zeros(size(A)) + # ^ we could make this sparse to make it more performant + end + else + if isdiag(cov) + diag_const = SVector{D, eltype(A)}(diag(noise_strength .* A)) + g = let diag_const = diag_const + (u, p, t) -> diag_const + end + else + A_const = SMatrix{size(A)..., eltype(A)}(noise_strength .* A) + g = let A_const = A_const + (u, p, t) -> A_const + end + noise_prototype = zeros(size(A)) + end + end + end + return g, noise_prototype +end + + +""" +Checks whether the function g depends on u based on 10 random points around the given u. +""" +function is_state_independent(g, u, p, t) + rdm_states = [u .+ rand(eltype(u), length(u)) .- 0.5 for _ in 1:10] + val = map(u -> g(u, p, t), rdm_states) + return length(unique(val)) == 1 +end + +""" +Checks whether g depends explicitly on time for select points on the interval [t0, t0+101]. +""" +function is_time_independent(g, u, p, t0) + trange = t0 .+ [0.0, 0.101, 1.01, 10.1, 101.0] + val = map(t -> g(u, p, t), trange) + return length(unique(val)) == 1 +end + +""" +Checks whether a matrix x is invertible by verifying that it has nonzero determinant. +""" +function is_invertible(x; tol = 1.0e-10) + F = lu(x, check = false) + det = abs(prod(diag(F.U))) + return det > tol +end + +""" +Checks whether the function f is linear by checking additivity and scalar multiplication +for two points x and y. +""" +function is_linear(f, x, y, c) + check1 = f(x + y) == f(x) + f(y) + check2 = f(c * x) == c * f(x) + return check1 && check2 +end + +function diffusion_function(g, IIP, noise_prototype) + return function diffusion(u, p, t) + if IIP + du = deepcopy(isnothing(noise_prototype) ? u : noise_prototype) + g(du, u, p, t) + return du + else + return g(u, p, t) + end + end +end + +function diffusion_function(ds::CoupledSDEs{IIP}) where {IIP} + prob = referrenced_sciml_prob(ds) + return diffusion_function(prob.g, IIP, prob.noise_rate_prototype) +end + +""" +Classifies the noise type of the CoupledSDEs given by the user. +Returns a named tuple of noise properties and the noise covariance matrix (if applicable). +""" +function find_noise_type(g, u0, p, t0, noise, covariance, noise_prototype, IIP) + noise_size = isnothing(noise_prototype) ? nothing : size(noise_prototype) + noise_cov = isnothing(noise) ? nothing : noise.covariance + D = length(u0) + + if !isnothing(noise_cov) + throw( + ArgumentError("CoupledSDEs does not support correlation between noise processes through DiffEqNoiseProcess.jl interface. Instead, use the `covariance` kwarg of `CoupledSDEs`.") + ) + end + + isadditive = false + isautonomous = false + islinear = false + isinvertible = false + + diffusion = diffusion_function(g, IIP, noise_prototype) + + if isnothing(g) + isadditive = true + isautonomous = true + islinear = true + if isnothing(covariance) + covariance = LinearAlgebra.I(D) + isinvertible = true + else + isinvertible = is_invertible(covariance) + end + elseif !isnothing(covariance) + throw( + ArgumentError("Both `g` and `covariance` are provided. Instead opt to encode the covariance in the difussion function `g` with the `noise_prototype` kwarg.") + ) + else + time_independent = is_time_independent(diffusion, rand(D), p, t0) + state_independent = is_state_independent(diffusion, u0, p, t0) + + # additive noise is equal to state independent noise + isadditive = state_independent + isautonomous = time_independent + + islinear = true + if !state_independent + for i in 1:10 + check = is_linear( + u -> diffusion(u, p, t0), + u0 + i .* rand(D), u0 + i .* rand(D), 2.0 + ) + check ? nothing : islinear = false + end + end + + # Previous formulation: + #islinear = !state_independent ? + # is_linear(u -> diffusion(u, p, t0), rand(D), rand(D), 2.0) : true + + if time_independent && state_independent + if !isnothing(noise_size) && isequal(noise_size...) + A = diffusion(zeros(D), p, 0.0) + covariance = A * A' + isinvertible = is_invertible(covariance) + elseif !isnothing(noise_size) && !isequal(noise_size...) + isinvertible = false + covariance = nothing + else + isinvertible = true + covariance = LinearAlgebra.I(D) + end + else + covariance = nothing + end + end + + noise_type = ( + additive = isadditive, autonomous = isautonomous, + linear = islinear, invertible = isinvertible, + ) + return noise_type, covariance +end + +function find_noise_type(prob::SDEProblem, IIP) + return find_noise_type( + prob.g, prob.u0, prob.p, prob.tspan[1], prob.noise, + nothing, prob.noise_rate_prototype, IIP + ) +end diff --git a/test/stochastic.jl b/test/stochastic.jl index ac015a93..f32c8415 100644 --- a/test/stochastic.jl +++ b/test/stochastic.jl @@ -3,8 +3,7 @@ using OrdinaryDiffEqTsit5: Tsit5 using OrdinaryDiffEqTsit5.OrdinaryDiffEqCore: None using StochasticDiffEq: SDEProblem, SRA, SOSRA, LambaEM, CorrelatedWienerProcess, EM -StochasticSystemsBase = Base.get_extension(DynamicalSystemsBase, :StochasticSystemsBase) -diffusion_matrix = StochasticSystemsBase.diffusion_matrix +using DynamicalSystemsBase: diffusion_matrix # Creation of lorenz @inbounds function lorenz_rule(u, p, t) @@ -133,9 +132,6 @@ end end @testset "utilities" begin - StochasticSystemsBase = Base.get_extension(DynamicalSystemsBase, :StochasticSystemsBase) - diffusion_matrix = StochasticSystemsBase.diffusion_matrix - @testset "diffusion_matrix" begin Γ = [1.0 0.3 0.0; 0.3 1 0.5; 0.0 0.5 1.0] A = sqrt(Γ) From b3ee73f420cce945e1604869bc842e9f292f8078 Mon Sep 17 00:00:00 2001 From: Orjan Ameye Date: Thu, 14 May 2026 14:39:24 +0200 Subject: [PATCH 2/2] add changelog --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e3ae8887..2ad5314a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,8 @@ +# v3.17 + +- `CoupledSDEs` (along with `diffusion_matrix` and related utilities) is now part of the core package instead of living behind the `StochasticSystemsBase` package extension. Users no longer need to load `StochasticDiffEq` to construct or use `CoupledSDEs`, and `Base.get_extension` is no longer required to access `diffusion_matrix`. +- The `StochasticDiffEq` dependency has been replaced with the smaller `StochasticDiffEqHighOrder` sublibrary, reducing load time and dependency surface for users of `CoupledSDEs`. Compile time of `DynamicalSystemsBase` increases slightly (about 2.5s to 3s) in exchange. + # v3.16 - `reinit!(ds::CoupledSDEs)` now accepts a `seed::UInt64` keyword that re-seeds the noise process random number generator. The default is `rand(UInt64)`, so each `reinit!` produces a fresh, independent noise realization unless an explicit seed is given.