Gotta go faassssttt!!!

Author: YingboMa

src/systems/diffeqs/odesystem.jl

@@ -74,18 +74,28 @@ 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)
 
-function extract_eqs_states_ps(eqs::AbstractArray{<:Equation}, iv)
+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...)
     # 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
@@ -99,25 +109,15 @@ function extract_eqs_states_ps(eqs::AbstractArray{<:Equation}, iv)
             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 append!(diffeq, algeeq), vcat(collect(diffvars), collect(algevars)), ps
-end
-
-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("An ODESystem can only have one independent variable."))
-    iv = first(ivs)
-    eqs, dvs, ps = extract_eqs_states_ps(eqs, iv)
-    return ODESystem(eqs, iv, dvs, ps; kwargs...)
+    return ODESystem(append!(diffeq, algeeq), iv, vcat(collect(diffvars), collect(algevars)), ps; kwargs...)
 end
 
 Base.:(==)(sys1::ODESystem, sys2::ODESystem) =