Merge pull request #3547 from AayushSabharwal/as/lex-sort-eqs

feat: lexicographically sort equations in `structural_simplify`

Author: ChrisRackauckas

src/structural_transformation/tearing.jl

@@ -63,7 +63,8 @@ function algebraic_variables_scc(state::TearingState)
         all(v -> !isdervar(state.structure, v),
             𝑠neighbors(graph, eq))
     end))
-    var_eq_matching = complete(maximal_matching(graph, e -> e in algeqs, v -> v in algvars))
+    var_eq_matching = complete(
+        maximal_matching(graph, e -> e in algeqs, v -> v in algvars), ndsts(graph))
     var_sccs = find_var_sccs(complete(graph), var_eq_matching)
 
     return var_eq_matching, var_sccs

src/systems/systems.jl

@@ -24,6 +24,7 @@ topological sort of the observed equations in `sys`.
 + When `simplify=true`, the `simplify` function will be applied during the tearing process.
 + `allow_symbolic=false`, `allow_parameter=true`, and `conservative=false` limit the coefficient types during tearing. In particular, `conservative=true` limits tearing to only solve for trivial linear systems where the coefficient has the absolute value of ``1``.
 + `fully_determined=true` controls whether or not an error will be thrown if the number of equations don't match the number of inputs, outputs, and equations.
++ `sort_eqs=true` controls whether equations are sorted lexicographically before simplification or not.
 """
 function structural_simplify(
         sys::AbstractSystem, io = nothing; additional_passes = [], simplify = false, split = true,
@@ -69,10 +70,11 @@ function __structural_simplify(sys::SDESystem, args...; kwargs...)
     return __structural_simplify(ODESystem(sys), args...; kwargs...)
 end
 
-function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = false,
+function __structural_simplify(
+        sys::AbstractSystem, io = nothing; simplify = false, sort_eqs = true,
         kwargs...)
     sys = expand_connections(sys)
-    state = TearingState(sys)
+    state = TearingState(sys; sort_eqs)
 
     @unpack structure, fullvars = state
     @unpack graph, var_to_diff, var_types = structure

src/systems/systemstructure.jl

@@ -260,7 +260,7 @@ function is_time_dependent_parameter(p, iv)
             (args = arguments(p); length(args)) == 1 && isequal(only(args), iv))
 end
 
-function TearingState(sys; quick_cancel = false, check = true)
+function TearingState(sys; quick_cancel = false, check = true, sort_eqs = true)
     sys = flatten(sys)
     ivs = independent_variables(sys)
     iv = length(ivs) == 1 ? ivs[1] : nothing
@@ -381,6 +381,14 @@ function TearingState(sys; quick_cancel = false, check = true)
     neqs = length(eqs)
     symbolic_incidence = symbolic_incidence[eqs_to_retain]
 
+    if sort_eqs
+        # sort equations lexicographically to reduce simplification issues
+        # depending on order due to NP-completeness of tearing.
+        sortidxs = Base.sortperm(eqs, by = string)
+        eqs = eqs[sortidxs]
+        symbolic_incidence = symbolic_incidence[sortidxs]
+    end
+
     ### Handle discrete variables
     lowest_shift = Dict()
     for var in fullvars

test/clock.jl

@@ -72,8 +72,9 @@ sss, = ModelingToolkit._structural_simplify!(deepcopy(tss[1]), (inputs[1], ()))
 @test isempty(equations(sss))
 d = Clock(dt)
 k = ShiftIndex(d)
-@test observed(sss) == [yd(k + 1) ~ Sample(dt)(y); r(k + 1) ~ 1.0;
-       ud(k + 1) ~ kp * (r(k + 1) - yd(k + 1))]
+@test issetequal(observed(sss),
+    [yd(k + 1) ~ Sample(dt)(y); r(k + 1) ~ 1.0;
+     ud(k + 1) ~ kp * (r(k + 1) - yd(k + 1))])
 
 d = Clock(dt)
 # Note that TearingState reorders the equations

test/discrete_system.jl

@@ -3,7 +3,7 @@
 - https://github.com/epirecipes/sir-julia/blob/master/markdown/function_map/function_map.md
 - https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#Deterministic_versus_stochastic_epidemic_models
 =#
-using ModelingToolkit, Test
+using ModelingToolkit, SymbolicIndexingInterface, Test
 using ModelingToolkit: t_nounits as t
 using ModelingToolkit: get_metadata, MTKParameters
 
@@ -37,13 +37,15 @@ syss = structural_simplify(sys)
 df = DiscreteFunction(syss)
 # iip
 du = zeros(3)
-u = collect(1:3)
+u = ModelingToolkit.better_varmap_to_vars(Dict([S => 1, I => 2, R => 3]), unknowns(syss))
 p = MTKParameters(syss, [c, nsteps, δt, β, γ] .=> collect(1:5))
 df.f(du, u, p, 0)
-@test du ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
+reorderer = getu(syss, [S, I, R])
+@test reorderer(du) ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
 
 # oop
-@test df.f(u, p, 0) ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
+@test reorderer(df.f(u, p, 0)) ≈
+      [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
 
 # Problem
 u0 = [S(k - 1) => 990.0, I(k - 1) => 10.0, R(k - 1) => 0.0]
@@ -98,12 +100,12 @@ function sir_map!(u_diff, u, p, t)
     end
     nothing
 end;
-u0 = prob_map2.u0;
+u0 = prob_map2[[S, I, R]];
 p = [0.05, 10.0, 0.25, 0.1];
 prob_map = DiscreteProblem(sir_map!, u0, tspan, p);
 sol_map2 = solve(prob_map, FunctionMap());
 
-@test Array(sol_map) ≈ Array(sol_map2)
+@test reduce(hcat, sol_map[[S, I, R]]) ≈ Array(sol_map2)
 
 # Delayed difference equation
 # @variables x(..) y(..) z(t)
@@ -317,9 +319,9 @@ end
 
     import ModelingToolkit: shift2term
     # unknowns(de) = xₜ₋₁, x, zₜ₋₁, xₜ₋₂, z
-    vars = ModelingToolkit.value.(unknowns(de))
-    @test isequal(shift2term(Shift(t, 1)(vars[1])), vars[2])
-    @test isequal(shift2term(Shift(t, 1)(vars[4])), vars[1])
-    @test isequal(shift2term(Shift(t, -1)(vars[5])), vars[3])
-    @test isequal(shift2term(Shift(t, -2)(vars[2])), vars[4])
+    vars = sort(ModelingToolkit.value.(unknowns(de)); by = string)
+    @test isequal(shift2term(Shift(t, 1)(vars[2])), vars[1])
+    @test isequal(shift2term(Shift(t, 1)(vars[3])), vars[2])
+    @test isequal(shift2term(Shift(t, -1)(vars[4])), vars[5])
+    @test isequal(shift2term(Shift(t, -2)(vars[1])), vars[3])
 end

test/downstream/linearize.jl

@@ -125,8 +125,9 @@ lsys = ModelingToolkit.reorder_unknowns(lsys0, unknowns(ssys), desired_order)
 @test lsys.C == [400 -4000]
 @test lsys.D == [4400 -4400]
 
-lsyss, _ = ModelingToolkit.linearize_symbolic(pid, [reference.u, measurement.u],
+lsyss0, ssys2 = ModelingToolkit.linearize_symbolic(pid, [reference.u, measurement.u],
     [ctr_output.u])
+lsyss = ModelingToolkit.reorder_unknowns(lsyss0, unknowns(ssys2), desired_order)
 
 @test ModelingToolkit.fixpoint_sub(
     lsyss.A, ModelingToolkit.defaults_and_guesses(pid)) == lsys.A
@@ -138,7 +139,7 @@ lsyss, _ = ModelingToolkit.linearize_symbolic(pid, [reference.u, measurement.u],
     lsyss.D, ModelingToolkit.defaults_and_guesses(pid)) == lsys.D
 
 # Test with the reverse desired unknown order as well to verify that similarity transform and reoreder_unknowns really works
-lsys = ModelingToolkit.reorder_unknowns(lsys, unknowns(ssys), reverse(desired_order))
+lsys = ModelingToolkit.reorder_unknowns(lsys, desired_order, reverse(desired_order))
 
 @test lsys.A == [-10 0; 0 0]
 @test lsys.B == [10 -10; 2 -2]

test/initializationsystem.jl

@@ -801,12 +801,17 @@ end
         end
 
         @parameters p=2.0 q=missing [guess = 1.0] c=1.0
-        @variables x=1.0 y=2.0 z=3.0
-
-        eqs = [0 ~ p * (y - x),
-            0 ~ x * (q - z) - y,
-            0 ~ x * y - c * z]
-        @mtkbuild sys = NonlinearSystem(eqs; initialization_eqs = [p^2 + q^2 + 2p * q ~ 0])
+        @variables x=1.0 z=3.0
+
+        # eqs = [0 ~ p * (y - x),
+        #     0 ~ x * (q - z) - y,
+        #     0 ~ x * y - c * z]
+        # specifically written this way due to
+        # https://github.com/SciML/NonlinearSolve.jl/issues/586
+        eqs = [0 ~ -c * z + (q - z) * (x^2)
+               0 ~ p * (-x + (q - z) * x)]
+        @named sys = NonlinearSystem(eqs; initialization_eqs = [p^2 + q^2 + 2p * q ~ 0])
+        sys = complete(sys)
         # @mtkbuild sys = NonlinearSystem(
         #     [p * x^2 + q * y^3 ~ 0, x - q ~ 0]; defaults = [q => missing],
         #     guesses = [q => 1.0], initialization_eqs = [p^2 + q^2 + 2p * q ~ 0])

test/input_output_handling.jl

@@ -414,6 +414,8 @@ matrices, ssys = linearize(augmented_sys,
         augmented_sys.d
     ], outs;
     op = [augmented_sys.u => 0.0, augmented_sys.input.u[2] => 0.0, augmented_sys.d => 0.0])
+matrices = ModelingToolkit.reorder_unknowns(
+    matrices, unknowns(ssys), [ssys.x[2], ssys.integrator.x[1], ssys.x[1]])
 @test matrices.A ≈ [A [1; 0]; zeros(1, 2) -0.001]
 @test matrices.B == I
 @test matrices.C == [C zeros(2)]

test/lowering_solving.jl

@@ -59,8 +59,7 @@ u0 = [lorenz1.x => 1.0,
     lorenz1.z => 0.0,
     lorenz2.x => 0.0,
     lorenz2.y => 1.0,
-    lorenz2.z => 0.0,
-    α => 2.0]
+    lorenz2.z => 0.0]
 
 p = [lorenz1.σ => 10.0,
     lorenz1.ρ => 28.0,
@@ -73,5 +72,5 @@ p = [lorenz1.σ => 10.0,
 tspan = (0.0, 100.0)
 prob = ODEProblem(connected, u0, tspan, p)
 sol = solve(prob, Rodas5())
-@test maximum(sol[2, :] + sol[6, :] + 2sol[1, :]) < 1e-12
+@test maximum(sol[lorenz1.x] + sol[lorenz2.y] + 2sol[α]) < 1e-12
 #using Plots; plot(sol,idxs=(:α,Symbol(lorenz1.x),Symbol(lorenz2.y)))

test/odesystem.jl

@@ -927,7 +927,7 @@ let
     sys_simp = structural_simplify(sys_con)
     true_eqs = [D(sys.x) ~ sys.v
                 D(sys.v) ~ ctrl.kv * sys.v + ctrl.kx * sys.x]
-    @test isequal(full_equations(sys_simp), true_eqs)
+    @test issetequal(full_equations(sys_simp), true_eqs)
 end
 
 let