Merge pull request #2051 from SciML/myb/sde

Add `System` that handles ODESystem and SDESystem

Author: YingboMa

src/ModelingToolkit.jl

@@ -147,6 +147,7 @@ include("discretedomain.jl")
 include("systems/systemstructure.jl")
 using .SystemStructures
 include("systems/clock_inference.jl")
+include("systems/systems.jl")
 
 include("debugging.jl")
 include("systems/alias_elimination.jl")
@@ -162,8 +163,9 @@ end
 
 export AbstractTimeDependentSystem, AbstractTimeIndependentSystem,
        AbstractMultivariateSystem
+
 export ODESystem, ODEFunction, ODEFunctionExpr, ODEProblemExpr, convert_system,
-       add_accumulations
+       add_accumulations, System
 export DAEFunctionExpr, DAEProblemExpr
 export SDESystem, SDEFunction, SDEFunctionExpr, SDEProblemExpr
 export SystemStructure
@@ -212,7 +214,7 @@ export simplify, substitute
 export build_function
 export modelingtoolkitize
 
-export @variables, @parameters, @constants
+export @variables, @parameters, @constants, @brownian
 export @named, @nonamespace, @namespace, extend, compose, complete
 export debug_system
 

src/parameters.jl

@@ -1,24 +1,26 @@
 import SymbolicUtils: symtype, term, hasmetadata, issym
-struct MTKParameterCtx end
+@enum VariableType VARIABLE PARAMETER BROWNIAN
+struct MTKVariableTypeCtx end
+
+getvariabletype(x, def = VARIABLE) = getmetadata(unwrap(x), MTKVariableTypeCtx, def)
 
 function isparameter(x)
     x = unwrap(x)
 
-    if x isa Symbolic && (isp = getmetadata(x, MTKParameterCtx, nothing)) !== nothing
-        return isp
+    if x isa Symbolic && (varT = getvariabletype(x, nothing)) !== nothing
+        return varT === PARAMETER
         #TODO: Delete this branch
     elseif x isa Symbolic && Symbolics.getparent(x, false) !== false
         p = Symbolics.getparent(x)
         isparameter(p) ||
             (hasmetadata(p, Symbolics.VariableSource) &&
              getmetadata(p, Symbolics.VariableSource)[1] == :parameters)
     elseif istree(x) && operation(x) isa Symbolic
-        getmetadata(x, MTKParameterCtx, false) ||
-            isparameter(operation(x))
+        varT === PARAMETER || isparameter(operation(x))
     elseif istree(x) && operation(x) == (getindex)
         isparameter(arguments(x)[1])
     elseif x isa Symbolic
-        getmetadata(x, MTKParameterCtx, false)
+        varT === PARAMETER
     else
         false
     end
@@ -35,7 +37,7 @@ function toparam(s)
     elseif s isa AbstractArray
         map(toparam, s)
     else
-        setmetadata(s, MTKParameterCtx, true)
+        setmetadata(s, MTKVariableTypeCtx, PARAMETER)
     end
 end
 toparam(s::Num) = wrap(toparam(value(s)))
@@ -45,13 +47,13 @@ toparam(s::Num) = wrap(toparam(value(s)))
 
 Maps the variable to a state.
 """
-tovar(s::Symbolic) = setmetadata(s, MTKParameterCtx, false)
+tovar(s::Symbolic) = setmetadata(s, MTKVariableTypeCtx, VARIABLE)
 tovar(s::Num) = Num(tovar(value(s)))
 
 """
 $(SIGNATURES)
 
-Define one or more known variables.
+Define one or more known parameters.
 """
 macro parameters(xs...)
     Symbolics._parse_vars(:parameters,

src/systems/abstractsystem.jl

@@ -1136,27 +1136,6 @@ function debug_system(sys::AbstractSystem)
     return sys
 end
 
-"""
-$(SIGNATURES)
-
-Structurally simplify algebraic equations in a system and compute the
-topological sort of the observed equations. When `simplify=true`, the `simplify`
-function will be applied during the tearing process. It also takes kwargs
-`allow_symbolic=false` and `allow_parameter=true` which limits the coefficient
-types during tearing.
-
-The optional argument `io` may take a tuple `(inputs, outputs)`.
-This will convert all `inputs` to parameters and allow them to be unconnected, i.e.,
-simplification will allow models where `n_states = n_equations - n_inputs`.
-"""
-function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false,
-                             kwargs...)
-    sys = expand_connections(sys)
-    sys isa DiscreteSystem && return sys
-    state = TearingState(sys)
-    structural_simplify!(state, io; simplify, kwargs...)
-end
-
 function eliminate_constants(sys::AbstractSystem)
     if has_eqs(sys)
         eqs = get_eqs(sys)

src/systems/diffeqs/sdesystem.jl

@@ -134,7 +134,7 @@ struct SDESystem <: AbstractODESystem
     end
 end
 
-function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
+function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dvs, ps;
                    controls = Num[],
                    observed = Num[],
                    systems = SDESystem[],
@@ -190,6 +190,19 @@ function SDESystem(sys::ODESystem, neqs; kwargs...)
     SDESystem(equations(sys), neqs, get_iv(sys), states(sys), parameters(sys); kwargs...)
 end
 
+function Base.:(==)(sys1::SDESystem, sys2::SDESystem)
+    sys1 === sys2 && return true
+    iv1 = get_iv(sys1)
+    iv2 = get_iv(sys2)
+    isequal(iv1, iv2) &&
+        isequal(nameof(sys1), nameof(sys2)) &&
+        isequal(get_eqs(sys1), get_eqs(sys2)) &&
+        isequal(get_noiseeqs(sys1), get_noiseeqs(sys2)) &&
+        _eq_unordered(get_states(sys1), get_states(sys2)) &&
+        _eq_unordered(get_ps(sys1), get_ps(sys2)) &&
+        all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
+end
+
 function generate_diffusion_function(sys::SDESystem, dvs = states(sys),
                                      ps = parameters(sys); kwargs...)
     return build_function(get_noiseeqs(sys),

src/systems/systems.jl

@@ -0,0 +1,88 @@
+function System(eqs::AbstractVector{<:Equation}, iv = nothing, args...; name = nothing,
+                kw...)
+    ODESystem(eqs, iv, args...; name, checks = false)
+end
+
+"""
+$(SIGNATURES)
+
+Structurally simplify algebraic equations in a system and compute the
+topological sort of the observed equations. When `simplify=true`, the `simplify`
+function will be applied during the tearing process. It also takes kwargs
+`allow_symbolic=false` and `allow_parameter=true` which limits the coefficient
+types during tearing.
+
+The optional argument `io` may take a tuple `(inputs, outputs)`.
+This will convert all `inputs` to parameters and allow them to be unconnected, i.e.,
+simplification will allow models where `n_states = n_equations - n_inputs`.
+"""
+function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false,
+                             kwargs...)
+    sys = expand_connections(sys)
+    sys isa DiscreteSystem && return sys
+    state = TearingState(sys)
+
+    @unpack structure, fullvars = state
+    @unpack graph, var_to_diff, var_types = structure
+    eqs = equations(state)
+    brown_vars = Int[]
+    new_idxs = zeros(Int, length(var_types))
+    idx = 0
+    for (i, vt) in enumerate(var_types)
+        if vt === BROWNIAN
+            push!(brown_vars, i)
+        else
+            new_idxs[i] = (idx += 1)
+        end
+    end
+    if isempty(brown_vars)
+        return structural_simplify!(state, io; simplify, kwargs...)
+    else
+        Is = Int[]
+        Js = Int[]
+        vals = Num[]
+        new_eqs = copy(eqs)
+        dvar2eq = Dict{Any, Int}()
+        for (v, dv) in enumerate(var_to_diff)
+            dv === nothing && continue
+            deqs = 𝑑neighbors(graph, dv)
+            if length(deqs) != 1
+                error("$(eqs[deqs]) is not handled.")
+            end
+            dvar2eq[fullvars[dv]] = only(deqs)
+        end
+        for (j, bj) in enumerate(brown_vars), i in 𝑑neighbors(graph, bj)
+            push!(Is, i)
+            push!(Js, j)
+            eq = new_eqs[i]
+            brown = fullvars[bj]
+            (coeff, residual, islinear) = Symbolics.linear_expansion(eq, brown)
+            islinear || error("$brown isn't linear in $eq")
+            new_eqs[i] = 0 ~ residual
+            push!(vals, coeff)
+        end
+        g = Matrix(sparse(Is, Js, vals))
+        sys = state.sys
+        @set! sys.eqs = new_eqs
+        @set! sys.states = [v
+                            for (i, v) in enumerate(fullvars)
+                            if !iszero(new_idxs[i]) && invview(var_to_diff)[i] === nothing]
+        # TODO: IO is not handled.
+        ode_sys = structural_simplify(sys, io; simplify, kwargs...)
+        eqs = equations(ode_sys)
+        sorted_g_rows = zeros(Num, length(eqs), size(g, 2))
+        for (i, eq) in enumerate(eqs)
+            dvar = eq.lhs
+            # differential equations always precede algebraic equations
+            _iszero(dvar) && break
+            g_row = get(dvar2eq, dvar, 0)
+            iszero(g_row) && error("$dvar isn't handled.")
+            g_row > size(g, 1) && continue
+            @views copyto!(sorted_g_rows[i, :], g[g_row, :])
+        end
+
+        return SDESystem(full_equations(ode_sys), sorted_g_rows,
+                         get_iv(ode_sys), states(ode_sys), parameters(ode_sys);
+                         name = nameof(ode_sys))
+    end
+end

src/systems/systemstructure.jl

@@ -9,7 +9,8 @@ import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten,
                           value, InvalidSystemException, isdifferential, _iszero,
                           isparameter, isconstant,
                           independent_variables, SparseMatrixCLIL, AbstractSystem,
-                          equations, isirreducible, input_timedomain, TimeDomain
+                          equations, isirreducible, input_timedomain, TimeDomain,
+                          VariableType, getvariabletype
 using ..BipartiteGraphs
 import ..BipartiteGraphs: invview, complete
 using Graphs
@@ -31,8 +32,6 @@ export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq, algeqs, is_only_discret
 export dervars_range, diffvars_range, algvars_range
 export DiffGraph, complete!
 
-@enum VariableType::Int8 DIFFERENTIAL_VARIABLE ALGEBRAIC_VARIABLE DERIVATIVE_VARIABLE
-
 struct DiffGraph <: Graphs.AbstractGraph{Int}
     primal_to_diff::Vector{Union{Int, Nothing}}
     diff_to_primal::Union{Nothing, Vector{Union{Int, Nothing}}}
@@ -149,13 +148,15 @@ Base.@kwdef mutable struct SystemStructure
     # or as `torn` to assert that tearing has run.
     graph::BipartiteGraph{Int, Nothing}
     solvable_graph::Union{BipartiteGraph{Int, Nothing}, Nothing}
+    var_types::Union{Vector{VariableType}, Nothing}
     only_discrete::Bool
 end
 
 function Base.copy(structure::SystemStructure)
+    var_types = structure.var_types === nothing ? nothing : copy(structure.var_types)
     SystemStructure(copy(structure.var_to_diff), copy(structure.eq_to_diff),
                     copy(structure.graph), copy(structure.solvable_graph),
-                    structure.only_discrete)
+                    var_types, structure.only_discrete)
 end
 
 is_only_discrete(s::SystemStructure) = s.only_discrete
@@ -230,9 +231,11 @@ function TearingState(sys; quick_cancel = false, check = true)
     symbolic_incidence = []
     fullvars = []
     var_counter = Ref(0)
-    addvar! = let fullvars = fullvars, var_counter = var_counter
+    var_types = VariableType[]
+    addvar! = let fullvars = fullvars, var_counter = var_counter, var_types = var_types
         var -> begin get!(var2idx, var) do
             push!(fullvars, var)
+            push!(var_types, getvariabletype(var))
             var_counter[] += 1
         end end
     end
@@ -331,7 +334,10 @@ function TearingState(sys; quick_cancel = false, check = true)
             push!(sorted_fullvars, v)
         end
     end
-    fullvars = collect(sorted_fullvars)
+    new_fullvars = collect(sorted_fullvars)
+    sortperm = indexin(new_fullvars, fullvars)
+    fullvars = new_fullvars
+    var_types = var_types[sortperm]
     var2idx = Dict(fullvars .=> eachindex(fullvars))
     dervaridxs = 1:length(dervaridxs)
 
@@ -358,7 +364,8 @@ function TearingState(sys; quick_cancel = false, check = true)
 
     return TearingState(sys, fullvars,
                         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
-                                        complete(graph), nothing, false), Any[])
+                                        complete(graph), nothing, var_types, false),
+                        Any[])
 end
 
 function lower_order_var(dervar)

src/variables.jl

@@ -390,3 +390,27 @@ function isintegervar(x)
     p === nothing || (x = p)
     return Symbolics.getmetadata(x, VariableInteger, false)
 end
+
+## Brownian
+"""
+    tobrownian(s::Sym)
+
+Maps the brownianiable to a state.
+"""
+tobrownian(s::Symbolic) = setmetadata(s, MTKVariableTypeCtx, BROWNIAN)
+tobrownian(s::Num) = Num(tobrownian(value(s)))
+isbrownian(s) = getvariabletype(s) === BROWNIAN
+
+"""
+$(SIGNATURES)
+
+Define one or more Brownian variables.
+"""
+macro brownian(xs...)
+    all(x -> x isa Symbol || Meta.isexpr(x, :call) && x.args[1] == :$, xs) ||
+        error("@brownian only takes scalar expressions!")
+    Symbolics._parse_vars(:brownian,
+                          Real,
+                          xs,
+                          tobrownian) |> esc
+end

test/sdesystem.jl

@@ -460,7 +460,7 @@ fdif!(du, u0, p, t)
     ]
     sys1 = SDESystem(eqs_short, noiseeqs, t, [x, y, z], [σ, ρ, β], name = :sys1)
     sys2 = SDESystem(eqs_short, noiseeqs, t, [x, y, z], [σ, ρ, β], name = :sys1)
-    @test_throws ArgumentError SDESystem([sys2.y ~ sys1.z], t, [], [], [],
+    @test_throws ArgumentError SDESystem([sys2.y ~ sys1.z], [], t, [], [],
                                          systems = [sys1, sys2], name = :foo)
 end