Merge pull request #394 from SciML/myb/autosys

Automatic state detection and some performance optimization

Author: YingboMa

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "3.6.2"
+version = "3.6.3"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/ModelingToolkit.jl

@@ -6,6 +6,7 @@ using Latexify, Unitful, ArrayInterface
 using MacroTools
 using UnPack: @unpack
 using DiffEqJump
+using DataStructures: OrderedDict, OrderedSet
 
 using Base.Threads
 import MacroTools: splitdef, combinedef, postwalk, striplines

src/systems/diffeqs/first_order_transform.jl

@@ -1,4 +1,3 @@
-using DataStructures: OrderedDict
 function lower_varname(var::Variable, idv, order)
     order == 0 && return var
     name = Symbol(var.name, :ˍ, string(idv.name)^order)

src/systems/diffeqs/odesystem.jl

@@ -74,22 +74,50 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     ODESystem(deqs, iv′, dvs′, ps′, tgrad, jac, Wfact, Wfact_t, name, systems)
 end
 
-var_from_nested_derivative(x) = var_from_nested_derivative(x,0)
 var_from_nested_derivative(x::Constant) = (missing, missing)
-var_from_nested_derivative(x,i) = x.op isa Differential ? var_from_nested_derivative(x.args[1],i+1) : (x.op,i)
+var_from_nested_derivative(x,i=0) = x.op isa Differential ? var_from_nested_derivative(x.args[1],i+1) : (x.op,i)
+
 iv_from_nested_derivative(x) = x.op isa Differential ? iv_from_nested_derivative(x.args[1]) : x.args[1].op
 iv_from_nested_derivative(x::Constant) = missing
 
 function ODESystem(eqs; kwargs...)
-    ivs = unique(skipmissing(iv_from_nested_derivative(eq.lhs) for eq ∈ eqs))
-    length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
-    iv = first(ivs)
-
-    dvs = unique(skipmissing(var_from_nested_derivative(eq.lhs)[1] for eq ∈ eqs))
-    ps = filter(vars(eq.rhs for eq ∈ eqs)) do x
-        isparameter(x) & !isequal(x, iv)
-    end |> collect
-    ODESystem(eqs, iv, dvs, ps; kwargs...)
+    # NOTE: this assumes that the order of algebric equations doesn't matter
+    diffvars = OrderedSet{Variable}()
+    allstates = OrderedSet{Variable}()
+    ps = OrderedSet{Variable}()
+    # reorder equations such that it is in the form of `diffeq, algeeq`
+    diffeq = Equation[]
+    algeeq = Equation[]
+    # initial loop for finding `iv`
+    iv = nothing
+    for eq in eqs
+        if !(eq.lhs isa Constant) # assume eq.lhs is either Differential or Constant
+            iv = iv_from_nested_derivative(eq.lhs)
+        end
+    end
+    iv === nothing && throw(ArgumentError("No differential variable detected."))
+    for eq in eqs
+        for var in vars(eq.rhs for eq ∈ eqs)
+            var isa Variable || continue
+            if isparameter(var)
+                isequal(var, iv) || push!(ps, var)
+            else
+                push!(allstates, var)
+            end
+        end
+        if eq.lhs isa Constant
+            push!(algeeq, eq)
+        else
+            diffvar = first(var_from_nested_derivative(eq.lhs))
+            iv == iv_from_nested_derivative(eq.lhs) || throw(ArgumentError("An ODESystem can only have one independent variable."))
+            diffvar in diffvars && throw(ArgumentError("The differential variable $diffvar is not unique in the system of equations."))
+            push!(diffvars, diffvar)
+            push!(diffeq, eq)
+        end
+    end
+    algevars = setdiff(allstates, diffvars)
+    # the orders here are very important!
+    return ODESystem(append!(diffeq, algeeq), iv, vcat(collect(diffvars), collect(algevars)), ps; kwargs...)
 end
 
 Base.:(==)(sys1::ODESystem, sys2::ODESystem) =

test/odesystem.jl

@@ -1,4 +1,5 @@
 using ModelingToolkit, StaticArrays, LinearAlgebra
+using OrdinaryDiffEq
 using DiffEqBase
 using Test
 
@@ -192,3 +193,28 @@ prob = ODEProblem(lotka,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])
 
 de = modelingtoolkitize(prob)
 ODEFunction(de)(similar(prob.u0), prob.u0, prob.p, 0.1)
+
+# automatic state detection for DAEs
+@parameters t k₁ k₂ k₃
+@variables y₁(t) y₂(t) y₃(t)
+@derivatives D'~t
+# reorder the system just to be a little spicier
+eqs = [D(y₁) ~ -k₁*y₁+k₃*y₂*y₃,
+       0     ~  y₁ + y₂ + y₃ - 1,
+       D(y₂) ~  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃]
+sys = ODESystem(eqs)
+u0 = [y₁ => 1.0,
+      y₂ => 0.0,
+      y₃ => 0.0]
+p  = [k₁ => 0.04,
+      k₂ => 3e7,
+      k₃ => 1e4]
+tspan = (0.0,100000.0)
+prob1 = ODEProblem(sys,u0,tspan,p)
+prob2 = ODEProblem(sys,u0,tspan,p,jac=true)
+# Wfact version is not very stable because of the lack of pivoting
+prob3 = ODEProblem(sys,u0,tspan,p,Wfact=true,Wfact_t=true)
+for (prob, atol) in [(prob1, 1e-12), (prob2, 1e-12), (prob3, 0.1)]
+    sol = solve(prob, Rodas5())
+    @test all(x->≈(sum(x), 1.0, atol=atol), sol.u)
+end