Add quick_cancel in alias_elimination

Author: YingboMa

src/systems/alias_elimination.jl

@@ -3,7 +3,7 @@ using SymbolicUtils: Rewriters
 const KEEP = typemin(Int)
 
 function alias_elimination(sys)
-    sys = initialize_system_structure(sys)
+    sys = initialize_system_structure(sys; quick_cancel=true)
     s = structure(sys)
     is_linear_equations, eadj, cadj = find_linear_equations(sys)
 

src/systems/systemstructure.jl

@@ -3,6 +3,7 @@ module SystemStructures
 using DataStructures
 using Symbolics: linear_expansion, unwrap
 using SymbolicUtils: istree, operation, arguments, Symbolic
+using SymbolicUtils: quick_cancel, similarterm
 using ..ModelingToolkit
 import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten,
     value, InvalidSystemException, isdifferential, _iszero, isparameter,
@@ -13,6 +14,14 @@ using UnPack
 using Setfield
 using SparseArrays
 
+quick_cancel_expr(expr) = Rewriters.Postwalk(
+    quick_cancel,
+    similarterm=(x, f, args) -> similarterm(
+        x, f, args, SymbolicUtils.symtype(x);
+        metadata=SymbolicUtils.metadata(x)
+    )
+)(expr)
+
 #=
 When we don't do subsitution, variable information is split into two different
 places, i.e. `states` and the right-hand-side of `observed`.
@@ -86,7 +95,7 @@ algvars_range(s::SystemStructure) = Iterators.filter(Base.Fix1(isalgvar, s), eac
 isalgeq(s::SystemStructure, eq::Integer) = s.algeqs[eq]
 isdiffeq(s::SystemStructure, eq::Integer) = !isalgeq(s, eq)
 
-function initialize_system_structure(sys)
+function initialize_system_structure(sys; quick_cancel=false)
     sys = flatten(sys)
     ivs = independent_variables(sys)
     eqs = copy(equations(sys))
@@ -109,9 +118,12 @@ function initialize_system_structure(sys)
     vars = OrderedSet()
     for (i, eq′) in enumerate(eqs)
         if _iszero(eq′.lhs)
+            rhs = quick_cancel ? quick_cancel_expr(eq′.rhs) : eq′.rhs
             eq = eq′
         else
-            eq = 0 ~ eq′.rhs - eq′.lhs
+            lhs = quick_cancel ? quick_cancel_expr(eq′.lhs) : eq′.lhs
+            rhs = quick_cancel ? quick_cancel_expr(eq′.rhs) : eq′.rhs
+            eq = 0 ~ rhs - lhs
         end
         vars!(vars, eq.rhs)
         isalgeq = true
@@ -141,6 +153,8 @@ function initialize_system_structure(sys)
         algeqs[i] = isalgeq
         if isalgeq
             eqs[i] = eq
+        else
+            eqs[i] = eqs[i].lhs ~ rhs
         end
     end