Fix codegen bugs and update tests

Author: YingboMa

src/structural_transformation/codegen.jl

@@ -419,8 +419,13 @@ function build_observed_function(
     if !isempty(subset)
         eqs = equations(sys)
 
-        torn_eqs  = map(i->map(v->eqs[var_eq_matching[v]], var_sccs[i]), subset)
-        torn_vars = map(i->map(v->fullvars[v], var_sccs[i]), subset)
+        nested_torn_vars_idxs = []
+        for iscc in subset
+            torn_vars_idxs = Int[var for var in var_sccs[iscc] if var_eq_matching[var] !== unassigned]
+            isempty(torn_vars_idxs) || push!(nested_torn_vars_idxs, torn_vars_idxs)
+        end
+        torn_eqs = [[eqs[var_eq_matching[i]] for i in idxs] for idxs in nested_torn_vars_idxs]
+        torn_vars = [fullvars[idxs] for idxs in nested_torn_vars_idxs]
         u0map = defaults(sys)
         assignments = copy(assignments)
         solves = map(zip(torn_eqs, torn_vars)) do (eqs, vars)

src/structural_transformation/symbolics_tearing.jl

@@ -153,7 +153,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify=false
         # convert it into the mass matrix form.
         # We cannot solve the differential variable like D(x)
         if isdiffvar(iv)
-            push!(diffeqs, solve_equation(eqs[ieq], fullvars[iv], simplify))
+            push!(diffeqs, solve_equation(neweqs[ieq], fullvars[iv], simplify))
             continue
         end
         push!(solved_equations, ieq); push!(solved_variables, iv)
@@ -174,25 +174,17 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify=false
             length(dterms) == 0 && return 0 ~ rhs
             new_rhs = rhs
             new_lhs = 0
-            nnegative = 0
             for iv in dterms
                 var = fullvars[iv]
+                # 0 ~ a * D(x) + b
+                # D(x) ~ -b/a
                 a, b, islinear = linear_expansion(new_rhs, var)
                 au = unwrap(a)
-                if !islinear || (au isa Symbolic) || isinput(var) || !(au isa Number)
+                if !islinear
                     return 0 ~ rhs
                 end
-                if -au < 0
-                    nnegative += 1
-                end
-                new_lhs -= a*var
-                new_rhs = b
-            end
-            # If most of the terms are negative, just multiply through by -1
-            # to make the equations looks slightly nicer.
-            if nnegative > div(length(dterms), 2)
-                new_lhs = -new_lhs
-                new_rhs = -new_rhs
+                new_lhs += var
+                new_rhs = -b/a
             end
             return new_lhs ~ new_rhs
         else # a number

test/reduction.jl

@@ -256,13 +256,11 @@ eq = [
 sys = structural_simplify(sys0)
 @test length(equations(sys)) == 1
 eq = equations(tearing_substitution(sys))[1]
-@test isequal(eq.lhs, 0)
+@test isequal(eq.lhs, D(v25))
 dv25 = ModelingToolkit.value(ModelingToolkit.derivative(eq.rhs, v25))
-ddv25 = ModelingToolkit.value(ModelingToolkit.derivative(eq.lhs, D(v25)))
 dt = ModelingToolkit.value(ModelingToolkit.derivative(eq.rhs, sin(10t)))
-@test dv25 ≈ -0.3
-@test ddv25 == 0.005
-@test dt == 0.1
+@test dv25 ≈ -60
+@test dt ≈ 20
 
 # Don't reduce inputs
 @parameters t σ ρ β

test/structural_transformation/index_reduction.jl

@@ -35,26 +35,6 @@ state = TearingState(pendulum)
 @unpack graph, var_to_diff = state.structure
 @test StructuralTransformations.maximal_matching(graph, eq->true, v->var_to_diff[v] === nothing) == map(x -> x == 0 ? StructuralTransformations.unassigned : x, [1, 2, 3, 4, 0, 0, 0, 0, 0])
 
-sys, var_eq_matching, eq_to_diff = StructuralTransformations.pantelides!(pendulum)
-state = TearingState(sys)
-@unpack graph, var_to_diff = state.structure
-@test graph.fadjlist == [[1, 7], [2, 8], [3, 5, 9], [4, 6, 9], [5, 6], [1, 2, 5, 6], [1, 3, 7, 10], [2, 4, 8, 11], [1, 2, 5, 6, 10, 11]]
-let N=nothing;
-    @test var_to_diff == [10, 11, N, N, 1, 2, 3, 4, N, N, N];
-#1: D(x) ~ w
-#2: D(y) ~ z
-#3: D(w) ~ T*x
-#4: D(z) ~ T*y - g
-#5: 0 ~ x^2 + y^2 - L^2
-# ----
-#6: D(eq:5) -> 0 ~ 2xx'+ 2yy'
-#7: D(eq:1) -> D(D(x)) ~ D(w) -> D(xˍt) ~ D(w) -> D(xˍt) ~ T*x
-#8: D(eq:2) -> D(D(y)) ~ D(z) -> D(y_t) ~ T*y - g
-#9: D(eq:6) -> 0 ~ 2xx'' + 2x'x' + 2yy'' + 2y'y'
-#                 [1, 2, 3, 4, 5, 6, 7, 8, 9]
-    @test eq_to_diff == [7, 8, N, N, 6, 9, N, N, N]
-end
-
 using ModelingToolkit
 @parameters t L g
 @variables x(t) y(t) w(t) z(t) T(t) xˍt(t) yˍt(t)

test/structural_transformation/tearing.jl

@@ -61,7 +61,7 @@ graph2vars(graph) = map(is->Set(map(i->int2var[i], is)), graph.fadjlist)
 state = TearingState(tearing(sys))
 let sss = state.structure
     @unpack graph = sss
-    @test graph2vars(graph) == [Set([u5])]
+    @test graph2vars(graph) == [Set([u1, u2, u5])]
 end
 
 # Before:

test/structural_transformation/utils.jl

@@ -19,10 +19,9 @@ pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)
 state = TearingState(pendulum)
 StructuralTransformations.find_solvables!(state)
 sss = state.structure
-@unpack graph, solvable_graph, fullvars, var_to_diff = sss
-@test isequal(fullvars, [D(x), D(y), D(w), D(z), x, y, w, z, T])
+@unpack graph, solvable_graph, var_to_diff = sss
 @test graph.fadjlist == [[1, 7], [2, 8], [3, 5, 9], [4, 6, 9], [5, 6]]
-@test graph.badjlist == 9 == length(fullvars)
+@test graph.badjlist == 9
 @test ne(graph) == nnz(incidence_matrix(graph)) == 12
 @test nv(solvable_graph) == 9 + 5
 let N = nothing