update first_order_transform

Co-authored-by: "Shashi Gowda" <gowda@mit.edu>
Co-authored-by: "Yingbo Ma" <mayingbo5@gmail.com>

Author: YingboMa

src/ModelingToolkit.jl

@@ -83,6 +83,7 @@ include("systems/diffeqs/odesystem.jl")
 include("systems/diffeqs/sdesystem.jl")
 include("systems/diffeqs/abstractodesystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
+include("systems/diffeqs/index_reduction.jl")
 include("systems/diffeqs/modelingtoolkitize.jl")
 include("systems/diffeqs/validation.jl")
 

src/systems/diffeqs/first_order_transform.jl

@@ -19,8 +19,8 @@ Takes a Nth order ODESystem and returns a new ODESystem written in first order
 form by defining new variables which represent the N-1 derivatives.
 """
 function ode_order_lowering(sys::ODESystem)
-    eqs_lowered, _ = ode_order_lowering(sys.eqs, sys.iv)
-    ODESystem(eqs_lowered, sys.iv, [var_from_nested_derivative(eq.lhs)[1] for eq in eqs_lowered], sys.ps)
+    eqs_lowered, new_vars = ode_order_lowering(sys.eqs, sys.iv)
+    ODESystem(eqs_lowered, sys.iv, vcat(new_vars, states(sys)), sys.ps)
 end
 
 function ode_order_lowering(eqs, iv)
@@ -30,14 +30,18 @@ function ode_order_lowering(eqs, iv)
     new_vars = Variable[]
 
     for (i, eq) ∈ enumerate(eqs)
-        var, maxorder = var_from_nested_derivative(eq.lhs)
-        if maxorder > get(var_order, var, 0)
-            var_order[var] = maxorder
-            any(isequal(var), vars) || push!(vars, var)
+        if isequal(eq.lhs, Constant(0))
+            push!(new_eqs, eq)
+        else
+            var, maxorder = var_from_nested_derivative(eq.lhs)
+            if maxorder > get(var_order, var, 0)
+                var_order[var] = maxorder
+                any(isequal(var), vars) || push!(vars, var)
+            end
+            var′ = lower_varname(var, iv, maxorder - 1)
+            rhs′ = rename_lower_order(eq.rhs)
+            push!(new_eqs,Differential(iv())(var′(iv())) ~ rhs′)
         end
-        var′ = lower_varname(var, iv, maxorder - 1)
-        rhs′ = rename_lower_order(eq.rhs)
-        push!(new_eqs,Differential(iv())(var′(iv())) ~ rhs′)
     end
 
     for var ∈ vars

src/systems/diffeqs/odesystem.jl

@@ -75,15 +75,17 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
 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)
 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(iv_from_nested_derivative(eq.lhs) for eq ∈ eqs)
+    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(var_from_nested_derivative(eq.lhs)[1] for eq ∈ eqs)
+    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

test/index_reduction.jl

@@ -5,14 +5,22 @@ using Test
 
 # Define some variables
 @parameters t L g
-@variables x(t) y(t) w(t) z(t) T(t)
+@variables x(t) y(t) w(t) z(t) T(t) x_t(t) y_t(t)
 @derivatives D'~t
 
 eqs2 = [D(D(x)) ~ T*x,
         D(D(y)) ~ T*y - g,
         0 ~ x^2 + y^2 - L^2]
-pendulum2 = ODESystem(eqs, t, [x, y, T], [L, g], name=:pendulum)
-ModelingToolkit.ode_order_lowering(pendulum2)
+pendulum2 = ODESystem(eqs2, t, [x, y, T], [L, g], name=:pendulum)
+lowered_sys = ModelingToolkit.ode_order_lowering(pendulum2)
+
+lowered_eqs = [D(x_t) ~ T*x,
+               D(y_t) ~ T*y - g,
+               0 ~ x^2 + y^2 - L^2,
+               D(x) ~ x_t,
+               D(y) ~ x_t]
+@test_skip ODESystem(lowered_eqs) == lowered_sys # not gonna work
+@test_broken isequal(lowered_sys.eqs, lowered_eqs)
 
 # Simple pendulum in cartesian coordinates
 eqs = [D(x) ~ w,