up

Author: vyudu

src/systems/diffeqs/odesystem.jl

@@ -296,9 +296,38 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
 end
 
 function ODESystem(eqs, iv; kwargs...)
-    param_deps = get(kwargs, :parameter_dependencies, Equation[])
-    eqs, dvs, ps = process_equations_DESystem(eqs, param_deps, iv)
-    return ODESystem(eqs, iv, dvs, ps; kwargs...)
+    diffvars, allunknowns, ps, eqs = process_equations(eqs, iv)
+
+    for eq in get(kwargs, :parameter_dependencies, Equation[])
+        collect_vars!(allunknowns, ps, eq, iv)
+    end
+
+    for ssys in get(kwargs, :systems, ODESystem[])
+        collect_scoped_vars!(allunknowns, ps, ssys, iv)
+    end
+
+    for v in allunknowns
+        isdelay(v, iv) || continue
+        collect_vars!(allunknowns, ps, arguments(v)[1], iv)
+    end
+
+    new_ps = OrderedSet()
+    for p in ps
+        if iscall(p) && operation(p) === getindex
+            par = arguments(p)[begin]
+            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
+               all(par[i] in ps for i in eachindex(par))
+                push!(new_ps, par)
+            else
+                push!(new_ps, p)
+            end
+        else
+            push!(new_ps, p)
+        end
+    end
+    algevars = setdiff(allunknowns, diffvars)
+
+    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars))), collect(new_ps); kwargs...)
 end
 
 # NOTE: equality does not check cached Jacobian

src/systems/diffeqs/sdesystem.jl

@@ -274,18 +274,46 @@ function SDESystem(sys::ODESystem, neqs; kwargs...)
 end
 
 function SDESystem(eqs::Vector{Equation}, noiseeqs::AbstractArray, iv; kwargs...) 
-    param_deps = get(kwargs, :parameter_dependencies, Equation[])
-    eqs, dvs, ps = process_equations_DESystem(eqs, param_deps, iv)
+    diffvars, allunknowns, ps, eqs = process_equations(eqs, iv)
+
+    for eq in get(kwargs, :parameter_dependencies, Equation[])
+        collect_vars!(allunknowns, ps, eq, iv)
+    end
+
+    for ssys in get(kwargs, :systems, ODESystem[])
+        collect_scoped_vars!(allunknowns, ps, ssys, iv)
+    end
+
+    for v in allunknowns
+        isdelay(v, iv) || continue
+        collect_vars!(allunknowns, ps, arguments(v)[1], iv)
+    end
+
+    new_ps = OrderedSet()
+    for p in ps
+        if iscall(p) && operation(p) === getindex
+            par = arguments(p)[begin]
+            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
+               all(par[i] in ps for i in eachindex(par))
+                push!(new_ps, par)
+            else
+                push!(new_ps, p)
+            end
+        else
+            push!(new_ps, p)
+        end
+    end
 
     # validate noise equations
     noisedvs = OrderedSet()
     noiseps = OrderedSet()
     collect_vars!(noisedvs, noiseps, noiseeqs, iv)
     for dv in noisedvs
-        var ∈ Set(dvs) || throw(ArgumentError("Variable $dv in noise equations is not an unknown of the system."))
+        var ∈ allunknowns || throw(ArgumentError("Variable $dv in noise equations is not an unknown of the system."))
     end
+    algevars = setdiff(allunknowns, diffvars)
 
-    return SDESystem(eqs, noiseeqs, iv, dvs, [ps; collect(noiseps)]; kwargs...)
+    return SDESystem(eqs, noiseeqs, iv, Iterators.flatten((diffvars, algevars)), [ps; collect(noiseps)]; kwargs...)
 end
 
 SDESystem(eq::Equation, noiseeqs::AbstractArray, args...; kwargs...) = SDESystem([eq], noiseeqs, args...; kwargs...)

src/utils.jl

@@ -1189,9 +1189,9 @@ end
 """
     $(TYPEDSIGNATURES)
 
-Find all the unknowns and parameters from the equations and parameter dependencies of an ODESystem or SDESystem. Return re-ordered equations, unknowns, and parameters.
+Find all the unknowns and parameters from the equations of a SDESystem or ODESystem. Return re-ordered equations, differential variables, all variables, and parameters.
 """
-function process_equations_DESystem(eqs, param_deps, iv)
+function process_equations(eqs, iv)
     eqs = collect(eqs)
 
     diffvars = OrderedSet()
@@ -1233,25 +1233,6 @@ function process_equations_DESystem(eqs, param_deps, iv)
             push!(algeeq, eq)
         end
     end
-    for eq in param_deps
-        collect_vars!(allunknowns, ps, eq, iv)
-    end
-
-    new_ps = OrderedSet()
-    for p in ps
-        if iscall(p) && operation(p) === getindex
-            par = arguments(p)[begin]
-            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
-               all(par[i] in ps for i in eachindex(par))
-                push!(new_ps, par)
-            else
-                push!(new_ps, p)
-            end
-        else
-            push!(new_ps, p)
-        end
-    end
-    algevars = setdiff(allunknowns, diffvars)
 
-    Equation[diffeq; algeeq; compressed_eqs], collect(Iterators.flatten((diffvars, algevars))), collect(new_ps)
+    diffvars, allunknowns, ps, Equation[diffeq; algeeq; compressed_eqs]
 end