Merge pull request #861 from SciML/myb/ex

Handle mass matrix in modelingtoolkitize

Author: ChrisRackauckas

Project.toml

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

src/structural_transformation/pantelides.jl

@@ -63,16 +63,11 @@ function pantelides_reassemble(sys, eqassoc, assign)
             lhsarg1 = arguments(eq.lhs)[1]
             @assert !(lhsarg1 isa Differential) "The equation $eq is not first order"
             i = get(d_dict, lhsarg1, nothing)
-            if i !== nothing
-                lhs = D(out_vars[varassoc[i]])
-                if lhs in lhss
-                    # check only trivial equations are removed
-                    @assert isequal(diff2term(D(eq.rhs)), diff2term(lhs)) "The duplicate equation is not trivial: $eq"
-                    lhs = Num(nothing)
-                end
-                lhs
-            else
+            if i === nothing
                 D(eq.lhs)
+            else
+                # remove clashing equations
+                lhs = Num(nothing)
             end
         else
             D(eq.lhs)
@@ -86,8 +81,10 @@ function pantelides_reassemble(sys, eqassoc, assign)
     final_vars = unique(filter(x->!(operation(x) isa Differential), fullvars))
     final_eqs = map(identity, filter(x->value(x.lhs) !== nothing, out_eqs[sort(filter(x->x != UNASSIGNED, assign))]))
 
-    # remove clashing equations (from order lowering vs index reduction)
-    return ODESystem(final_eqs, independent_variable(sys), final_vars, parameters(sys))
+    @set! sys.eqs = final_eqs
+    @set! sys.states = final_vars
+    @set! sys.structure = nothing
+    return sys
 end
 
 """

src/systems/diffeqs/modelingtoolkitize.jl

@@ -18,27 +18,46 @@ function modelingtoolkitize(prob::DiffEqBase.ODEProblem)
     vars = ArrayInterface.restructure(prob.u0,[var(:x, i)(ModelingToolkit.value(t)) for i in eachindex(prob.u0)])
     params = p isa DiffEqBase.NullParameters ? [] :
              reshape([Num(Sym{Real}(nameof(Variable(:α, i)))) for i in eachindex(p)],size(p))
+    var_set = Set(vars)
 
     D = Differential(t)
+    mm = prob.f.mass_matrix
 
-    rhs = [D(var) for var in vars]
+    if mm === I
+        lhs = map(v->D(v), vars)
+    else
+        lhs = map(mm * vars) do v
+            if iszero(v)
+                0
+            elseif v in var_set
+                D(v)
+            else
+                error("Non-permuation mass matrix is not supported.")
+            end
+        end
+    end
 
     if DiffEqBase.isinplace(prob)
-        lhs = similar(vars, Num)
-        prob.f(lhs, vars, params, t)
+        rhs = similar(vars, Num)
+        prob.f(rhs, vars, params, t)
     else
-        lhs = prob.f(vars, params, t)
+        rhs = prob.f(vars, params, t)
     end
 
-    eqs = vcat([rhs[i] ~ lhs[i] for i in eachindex(prob.u0)]...)
+    eqs = vcat([lhs[i] ~ rhs[i] for i in eachindex(prob.u0)]...)
 
+    sts = vec(collect(vars))
     params = if ndims(params) == 0
         [params[1]]
     else
-        Vector(vec(params))
+        vec(collect(params))
     end
 
-    de = ODESystem(eqs,t,Vector(vec(vars)),params)
+    de = ODESystem(
+        eqs, t, sts, params,
+        default_u0=Dict(sts .=> vec(collect(prob.u0))),
+        default_p=Dict(params .=> vec(collect(prob.p))),
+    )
 
     de
 end

test/modelingtoolkitize.jl

@@ -160,3 +160,28 @@ prob = ODEProblem(f, rv0, (0.0, Δt), μ)
 sol = solve(prob, Vern8())
 
 modelingtoolkitize(prob)
+
+# Index reduction and mass matrix handling
+using LinearAlgebra
+function pendulum!(du, u, p, t)
+    x, dx, y, dy, T = u
+    g, L = p
+    du[1] = dx
+    du[2] = T*x
+    du[3] = dy
+    du[4] = T*y - g
+    du[5] = x^2 + y^2 - L^2
+    return nothing
+end
+pendulum_fun! = ODEFunction(pendulum!, mass_matrix=Diagonal([1,1,1,1,0]))
+u0 = [1.0, 0, 0, 0, 0]
+p = [9.8, 1]
+tspan = (0, 10.0)
+pendulum_prob = ODEProblem(pendulum_fun!, u0, tspan, p)
+pendulum_sys_org = modelingtoolkitize(pendulum_prob)
+sts = states(pendulum_sys_org)
+pendulum_sys = dae_index_lowering(pendulum_sys_org)
+prob = ODEProblem(pendulum_sys, Pair[], tspan)
+sol = solve(prob, Rodas4())
+l2 = sol[sts[1]].^2 + sol[sts[3]].^2
+@test all(l->abs(sqrt(l) - 1) < 0.05, l2)