add coupledSDEs type

.gitignore

@@ -6,4 +6,5 @@ deps/build.log
 Manifest.toml
 *style.jl
 *.scss
-*.css
+*.css
+.vscode

Project.toml

@@ -13,6 +13,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StateSpaceSets = "40b095a5-5852-4c12-98c7-d43bf788e795"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
 [compat]
@@ -23,5 +24,6 @@ Roots = "1, 2"
 SciMLBase = "1.19.5, 2"
 StateSpaceSets = "1"
 Statistics = "1"
+StochasticDiffEq = "6.66.0"
 SymbolicIndexingInterface = "0.3.4"
 julia = "1.9"

src/DynamicalSystemsBase.jl

@@ -22,6 +22,7 @@ include("core/pretty_printing.jl")
 
 include("core_systems/discrete_time_map.jl")
 include("core_systems/continuous_time_ode.jl")
+include("core_systems/continuous_time_sde.jl")
 include("core_systems/arbitrary_steppable.jl")
 include("core_systems/additional_supertypes.jl")
 include("core_systems/jacobian.jl")

src/core_systems/continuous_time_sde.jl

@@ -0,0 +1,245 @@
+using SciMLBase: SDEProblem, AbstractSDEIntegrator, __init, SDEFunction
+using StochasticDiffEq: SOSRA
+export CoupledSDEs
+
+###########################################################################################
+# DiffEq options
+###########################################################################################
+const DEFAULT_SDE_SOLVER = SOSRA() # default recommendation for additive noise
+const DEFAULT_STOCH_DIFFEQ = (alg=DEFAULT_SDE_SOLVER, DEFAULT_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
+
+###########################################################################################
+# Type
+###########################################################################################
+
+# Define CoupledSDEs
+"""
+    CoupledSDEs <: ContinuousTimeDynamicalSystem
+    CoupledSDEs(f, g, u0 [, p, σ]; kwargs...)
+
+A stochastic continuous time dynamical system defined by a set of
+coupled ordinary differential equations as follows:
+```math
+d\\vec{u} = \\vec{f}(\\vec{u}, p, t) dt + \\vec{g}(\\vec{u}, p, t) dW_t
+```
+
+Optionally provide the overall noise strength `σ`, the parameter container `p` and initial time as keyword `t0`. If `σ` is provided, the diffusion function `g` is multiplied by `σ`.
+
+For construction instructions regarding `f, u0` see the [DynamicalSystems.jl tutorial](https://juliadynamics.github.io/DynamicalSystems.jl/latest/tutorial/#CoupledODEs).
+
+The stochastic part of the differential equation is defined by the function `g` and the keyword arguments `noise_rate_prototype` and `noise`. `noise` indicates the noise process applied during generation and defaults to Gaussian white noise. For details on defining various noise processes, refer to the noise process documentation page. `noise_rate_prototype` indicates the prototype type instance for the noise rates, i.e., the output of `g`. It can be any type which overloads `A_mul_B!` with itself being the middle argument. Commonly, this is a matrix or sparse matrix. If this is not given, it defaults to `nothing`, which means the problem should be interpreted as having diagonal noise.
+
+## DifferentialEquations.jl interfacing
+
+The ODEs are evolved via the solvers of DifferentialEquations.jl.
+When initializing a `CoupledODEs`, you can specify the solver that will integrate
+`f` in time, along with any other integration options, using the `diffeq` keyword.
+For example you could use `diffeq = (abstol = 1e-9, reltol = 1e-9)`.
+If you want to specify a solver, do so by using the keyword `alg`, e.g.:
+`diffeq = (alg = Tsit5(), reltol = 1e-6)`. This requires you to have been first
+`using OrdinaryDiffEq` to access the solvers. The default `diffeq` is:
+
+```julia
+$(DEFAULT_STOCH_DIFFEQ)
+```
+
+`diffeq` keywords can also include `callback` for [event handling
+](https://docs.sciml.ai/DiffEqDocs/stable/features/callback_functions/).
+
+Dev note: `CoupledSDEs` is a light wrapper of  `StochasticDiffEq.SDEIntegrator` from DifferentialEquations.jl.
+The integrator is available as the field `integ`, and the `SDEProblem` is `integ.sol.prob`.
+The convenience syntax `SDEProblem(ds::CoupledSDEs, tspan = (t0, Inf))` is available
+to extract the problem.
+"""
+struct CoupledSDEs{IIP,D,I,P,S} <: ContinuousTimeDynamicalSystem
+    # D parametrised by length of u0
+    # S indicated if the noise strength has been added to the diffusion function
+    integ::I
+    # things we can't recover from `integ`
+    p0::P
+    noise_strength
+    diffeq # isn't parameterized because it is only used for display
+end
+
+function CoupledSDEs(
+    f,
+    g,
+    u0,
+    p=SciMLBase.NullParameters(),
+    noise_strength=1.0;
+    t0=0.0,
+    diffeq=DEFAULT_STOCH_DIFFEQ,
+    noise_rate_prototype=nothing,
+    noise=nothing,
+    seed=UInt64(0),
+)
+    IIP = isinplace(f, 4) # from SciMLBase
+    @assert IIP == isinplace(g, 4) "f and g must both be in-place or out-of-place"
+
+    s = correct_state(Val{IIP}(), u0)
+    T = eltype(s)
+    prob = SDEProblem{IIP}(
+        f,
+        g,
+        s,
+        (T(t0), T(Inf)),
+        p;
+        noise_rate_prototype=noise_rate_prototype,
+        noise=noise,
+        seed=seed,
+    )
+    return CoupledSDEs(prob, diffeq, noise_strength)
+end
+
+function CoupledSDEs(
+    prob::SDEProblem, diffeq=DEFAULT_STOCH_DIFFEQ, noise_strength=1.0; special_kwargs...
+    )
+    if haskey(special_kwargs, :diffeq)
+        throw(ArgumentError("`diffeq` is given as positional argument when an ODEProblem 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
+    sde_function = SDEFunction(prob.f, add_noise_strength(noise_strength, prob.g, IIP))
+    prob = SciMLBase.remake(prob; f=sde_function)
+
+    solver, remaining = _decompose_into_sde_solver_and_remaining(diffeq)
+    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,true}(
+        integ, deepcopy(prob.p), noise_strength, diffeq
+    )
+end
+
+"""
+    CoupledSDEs(ds::CoupledODEs, g, p [, σ]; kwargs...)
+
+Converts a [`CoupledODEs`](https://juliadynamics.github.io/DynamicalSystems.jl/stable/tutorial/#CoupledODEs)
+system into a [`CoupledSDEs`](@ref).
+"""
+function CoupledSDEs(
+    ds::CoupledODEs,
+    g,
+    p, # the parameter is likely changed as the diffusion function g is added.
+    noise_strength=1.0;
+    diffeq=DEFAULT_STOCH_DIFFEQ,
+    noise_rate_prototype=nothing,
+    noise=nothing,
+    seed=UInt64(0),
+)
+    return CoupledSDEs(
+        dynamic_rule(ds),
+        g,
+        current_state(ds),
+        p,
+        noise_strength;
+        diffeq=diffeq,
+        noise_rate_prototype=noise_rate_prototype,
+        noise=noise,
+        seed=seed,
+    )
+end
+
+"""
+    CoupledODEs(ds::CoupledSDEs; kwargs...)
+
+Converts a [`CoupledSDEs`](@ref) into [`CoupledODEs`](@ref).
+"""
+function CoupledODEs(
+    sys::CoupledSDEs; diffeq=DEFAULT_STOCH_DIFFEQ, t0=0.0
+)
+    return CoupledODEs(
+        sys.integ.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 [
+        "Noise strength" => ds.noise_strength,
+        "SDE solver" => string(nameof(typeof(solver))),
+        "SDE kwargs" => remaining,
+    ]
+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
+
+function dynamic_rule(sys::CoupledSDEs)
+    f = sys.integ.f
+    while hasfield(typeof(f), :f)
+        f = f.f
+    end
+    return f
+end
+
+# so that `ds` is printed
+set_state!(ds::CoupledSDEs, u::AbstractArray) = (set_state!(ds.integ, u); ds)
+SciMLBase.step!(ds::CoupledSDEs, args...) = (step!(ds.integ, args...); ds)
+
+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 `u_modified!` is called
+function set_parameter!(ds::CoupledSDEs, args...)
+    _set_parameter!(ds, args...)
+    u_modified!(ds.integ, true)
+    return nothing
+end
+
+referrenced_sciml_prob(ds::CoupledSDEs) = ds.integ.sol.prob
+
+###########################################################################################
+# Utilities
+###########################################################################################
+
+function σg(σ, g)
+    return (u, p, t) -> σ .* g(u, p, t)
+end
+
+function σg!(σ, g!)
+    function (du, u, p, t)
+        g!(du, u, p, t)
+        du .*= σ
+        return nothing
+    end
+end
+
+function add_noise_strength(σ, g, IIP)
+    newg = IIP ? σg!(σ, g) : σg(σ, g)
+    return newg
+end

test/Project.toml

@@ -4,6 +4,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

test/runtests.jl

@@ -7,6 +7,7 @@ testfile(file, testname=defaultname(file)) = @testset "$testname" begin; include
 @testset "DynamicalSystemsBase" begin
     testfile("discrete.jl")
     testfile("continuous.jl")
+    testfile("stochastic.jl")
     testfile("arbitrarysteppable.jl")
     testfile("stroboscopic.jl")
     testfile("parallel.jl")

test/stochastic.jl

@@ -0,0 +1,70 @@
+using DynamicalSystemsBase, Test
+using StochasticDiffEq: SDEProblem, SRA, SRI, SOSRA
+
+include("test_system_function.jl")
+
+# Creation of lorenz
+@inbounds function lorenz_rule(u, p, t)
+    σ = p[1]; ρ = p[2]; β = p[3]
+    du1 = σ*(u[2]-u[1])
+    du2 = u[1]*(ρ-u[3]) - u[2]
+    du3 = u[1]*u[2] - β*u[3]
+    return SVector{3}(du1, du2, du3)
+end
+@inbounds function lorenz_rule_iip(du, u, p, t)
+    σ = p[1]; ρ = p[2]; β = p[3]
+    du[1] = σ*(u[2]-u[1])
+    du[2] = u[1]*(ρ-u[3]) - u[2]
+    du[3] = u[1]*u[2] - β*u[3]
+    return nothing
+end
+
+σ = 0.1
+function diagonal_noise!(σ)
+    function (du, u, p, t)
+        du .= σ .* ones(length(u))
+        return nothing
+    end
+end
+diagonal_noise(σ) = (u, p, t) -> SVector{3}(σ, σ, σ)
+
+u0 = [0, 10.0, 0]
+p0 = [10, 28, 8/3]
+
+# diagonal additive noise
+lorenz_oop = CoupledSDEs(lorenz_rule, diagonal_noise(σ), u0, p0)
+lorenz_iip = CoupledSDEs(SDEProblem(lorenz_rule_iip, diagonal_noise!(σ), copy(u0), (0.0, Inf), p0))
+lorenz_SRA = CoupledSDEs(lorenz_rule, diagonal_noise(σ), u0, p0;
+    diffeq = (alg = SRA(), abstol = 1e-6, reltol = 1e-6)
+)
+
+for (ds, iip) in zip((lorenz_oop, lorenz_iip, lorenz_SRA), (false, true, false))
+
+    name = (ds === lorenz_SRA) ? "lorSRA" : "lorenz"
+    @testset "$(name) IIP=$(iip)" begin
+        @test dynamic_rule(ds) == (iip ? lorenz_rule_iip : lorenz_rule)
+        test_dynamical_system(ds, u0, p0; idt = false, iip)
+    end
+end
+
+@testset "correct SDE propagation" begin
+    lorenz_oop = CoupledSDEs(lorenz_rule, diagonal_noise(σ), u0, p0)
+    @test lorenz_oop.integ.alg isa SOSRA
+
+    lorenz_SRA = CoupledSDEs(lorenz_rule, diagonal_noise(σ), u0, p0;
+        diffeq = (alg = SRA(), abstol = 1e-6, reltol = 1e-6, verbose=false)
+    )
+    @test lorenz_SRA.integ.alg isa SRA
+    @test lorenz_SRA.integ.opts.verbose == false
+
+    # # also test SDEproblem creation
+    prob = lorenz_SRA.integ.sol.prob
+
+    ds = CoupledSDEs(prob, (alg=SRI(), abstol=0.0, reltol=1e-6, verbose=false))
+
+    @test ds.integ.alg isa SRI
+    @test ds.integ.opts.verbose == false
+
+    @test_throws ArgumentError CoupledSDEs(prob; diffeq = (alg=SRI(), ))
+
+end

test/test_system_function.jl

@@ -14,7 +14,11 @@ function test_dynamical_system(ds, u0, p0; idt, iip,
         @test current_time(ds) == 0
         @test initial_time(ds) == 0
         @test isinplace(ds) == iip
-        @test isdeterministic(ds) == true
+        if ds isa DynamicalSystemsBase.CoupledSDEs
+            @test isdeterministic(ds) == false
+        else
+            @test isdeterministic(ds) == true
+        end
         @test isdiscretetime(ds) == idt
         @test ds(0) == u0
     end
@@ -138,4 +142,3 @@ function test_dynamical_system(ds, u0, p0; idt, iip,
     end
 
 end
-