Skip to content

Commit

Permalink
feat: add ability to convert SDESystem to equivalent ODESystem
Browse files Browse the repository at this point in the history
  • Loading branch information
AayushSabharwal committed Dec 11, 2024
1 parent 6d6383a commit a2acbe5
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 0 deletions.
41 changes: 41 additions & 0 deletions src/systems/diffeqs/sdesystem.jl
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,47 @@ function Base.:(==)(sys1::SDESystem, sys2::SDESystem)
all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
end

"""
function ODESystem(sys::SDESystem)
Convert an `SDESystem` to the equivalent `ODESystem` using `@brownian` variables instead
of noise equations. The returned system will not be `iscomplete` and will not have an
index cache, regardless of `iscomplete(sys)`.
"""
function ODESystem(sys::SDESystem)
neqs = get_noiseeqs(sys)
eqs = equations(sys)
is_scalar_noise = get_is_scalar_noise(sys)
nbrownian = if is_scalar_noise
length(neqs)
else
size(neqs, 2)
end
brownvars = map(1:nbrownian) do i
name = gensym(Symbol(:brown_, i))
only(@brownian $name)
end
if is_scalar_noise
brownterms = reduce(+, neqs .* brownvars; init = 0)
neweqs = map(eqs) do eq
eq.lhs ~ eq.rhs + brownterms
end
else
if neqs isa AbstractVector
neqs = reshape(neqs, (length(neqs), 1))
end
brownterms = neqs * brownvars
neweqs = map(eqs, brownterms) do eq, brown
eq.lhs ~ eq.rhs + brown
end
end
newsys = ODESystem(neweqs, get_iv(sys), unknowns(sys), parameters(sys);
parameter_dependencies = parameter_dependencies(sys), defaults = defaults(sys),
continuous_events = continuous_events(sys), discrete_events = discrete_events(sys),
name = nameof(sys), description = description(sys), metadata = get_metadata(sys))
@set newsys.parent = sys
end

function __num_isdiag_noise(mat)
for i in axes(mat, 1)
nnz = 0
Expand Down
49 changes: 49 additions & 0 deletions test/sdesystem.jl
Original file line number Diff line number Diff line change
Expand Up @@ -809,3 +809,52 @@ end
prob = SDEProblem(sys, [x => 1.0, y => 1.0], (0.0, 1.0))
@test prob[z] 2.0
end

@testset "SDESystem to ODESystem" begin
@variables x(t) y(t) z(t)
@testset "Scalar noise" begin
@named sys = SDESystem([D(x) ~ x, D(y) ~ y, z ~ x + y], [x, y, 3],
t, [x, y, z], [], is_scalar_noise = true)
odesys = ODESystem(sys)
@test odesys isa ODESystem
vs = ModelingToolkit.vars(equations(odesys))
nbrownian = count(
v -> ModelingToolkit.getvariabletype(v) == ModelingToolkit.BROWNIAN, vs)
@test nbrownian == 3
for eq in equations(odesys)
ModelingToolkit.isdiffeq(eq) || continue
@test length(arguments(eq.rhs)) == 4
end
end

@testset "Non-scalar vector noise" begin
@named sys = SDESystem([D(x) ~ x, D(y) ~ y, z ~ x + y], [x, y, 0],
t, [x, y, z], [], is_scalar_noise = false)
odesys = ODESystem(sys)
@test odesys isa ODESystem
vs = ModelingToolkit.vars(equations(odesys))
nbrownian = count(
v -> ModelingToolkit.getvariabletype(v) == ModelingToolkit.BROWNIAN, vs)
@test nbrownian == 1
for eq in equations(odesys)
ModelingToolkit.isdiffeq(eq) || continue
@test length(arguments(eq.rhs)) == 2
end
end

@testset "Matrix noise" begin
noiseeqs = [x+y y+z z+x
2y 2z 2x
z+1 x+1 y+1]
@named sys = SDESystem([D(x) ~ x, D(y) ~ y, D(z) ~ z], noiseeqs, t, [x, y, z], [])
odesys = ODESystem(sys)
@test odesys isa ODESystem
vs = ModelingToolkit.vars(equations(odesys))
nbrownian = count(
v -> ModelingToolkit.getvariabletype(v) == ModelingToolkit.BROWNIAN, vs)
@test nbrownian == 3
for eq in equations(odesys)
@test length(arguments(eq.rhs)) == 4
end
end
end

0 comments on commit a2acbe5

Please sign in to comment.