feat: add pass which removes equations incident on a single variable

Author: AayushSabharwal

src/systems/alias_elimination.jl

@@ -148,6 +148,91 @@ function find_perfect_aliases!(
     return aliases
 end
 
+"""
+    $TYPEDSIGNATURES
+
+Analytically remove any equations in `state` which are only incident on a single variable.
+"""
+function remove_constant_variables!(state::TearingState; allow_parameter::Bool = true, kwargs...)
+    StateSelection.complete!(state.structure)
+    (; additional_observed, original_eqs, fullvars, structure, sys) = state
+    (; graph, var_to_diff) = structure
+    eqs = collect(equations(state))
+    eqs_to_rm = Int[]
+    vars_to_rm = Int[]
+    vars_to_rm_set = BitSet()
+    eqs_to_rm_set = BitSet()
+    fullvars_set = Set{SymbolicT}(fullvars)
+    subs = Dict{SymbolicT, SymbolicT}()
+    eqs_to_substitute = Set{Int}()
+    param_der_subber = SU.Substituter{false}(state.param_derivative_map)
+
+    snbors = Int[]
+    for _ in 1:4
+        removed_eq = false
+        for ieq in 𝑠vertices(graph)
+            empty!(snbors)
+            append!(snbors, 𝑠neighbors(graph, ieq))
+            setdiff!(snbors, vars_to_rm_set)
+            length(snbors) == 1 || continue
+            ivar = first(snbors)
+            eq = eqs[ieq]
+            var = fullvars[ivar]
+            lex = Symbolics.LinearExpander(var; strict = true)
+            a, b, islin = lex(eq.rhs)
+            islin || continue
+            # `allow_symbolic = true` since we know this equation (directly or indirectly) only
+            # depends on `var`. Any variables present in it can only be ones we've already
+            # eliminated in `vars_to_rm`.
+            if !MTKTearing._check_allow_symbolic_parameter(
+                    state, a, true, allow_parameter; fullvars_set
+                )
+                continue
+            end
+            removed_eq = true
+            push!(eqs_to_rm, ieq)
+            push!(eqs_to_rm_set, ieq)
+            push!(vars_to_rm, ivar)
+            push!(vars_to_rm_set, ivar)
+            # `a` typically is faster to negate, since it usually is a constant or small expression
+            rhs = b / -a
+            push!(additional_observed, var ~ rhs)
+            subs[var] = rhs
+            union!(eqs_to_substitute, 𝑑neighbors(graph, ivar))
+
+            v = var_to_diff[ivar]
+            while v isa Int
+                # We're only looking at equations incident on a single variable, so `rhs` will only ever
+                # involve parameters. The derivative is thus going to be zero.
+                rhs = param_der_subber(Symbolics.derivative(rhs, get_iv(sys)::SymbolicT; throw_no_derivative = true))
+                subs[fullvars[v]] = rhs
+                push!(additional_observed, default_toterm(fullvars[v]) ~ rhs)
+                union!(eqs_to_substitute, 𝑑neighbors(graph, v))
+                push!(vars_to_rm, v)
+                push!(vars_to_rm_set, v)
+                v = var_to_diff[v]
+            end
+        end
+
+        removed_eq || break
+    end
+
+    subber = SU.Substituter{false}(subs)
+    for ieq in eqs_to_substitute
+        ieq in eqs_to_rm_set && continue
+        eqs[ieq] = subber(eqs[ieq])
+        original_eqs[ieq] = subber(original_eqs[ieq])
+    end
+
+    @set! sys.eqs = eqs
+    state.sys = sys
+    old_to_new_eq, old_to_new_var = StateSelection.rm_eqs_vars!(
+        state, eqs_to_rm, vars_to_rm
+    )
+
+    return length(eqs_to_rm)
+end
+
 function alias_elimination!(state::TearingState; fully_determined = true,
                             print_underconstrained_variables = false, kwargs...)
     StateSelection.complete!(state.structure)

src/systems/systemstructure.jl

@@ -239,6 +239,7 @@ function _mtkcompile!(
     state = ModelingToolkit.inputs_to_parameters!(state, discrete_inputs, OrderedSet{SymbolicT}())
     state = ModelingToolkit.inputs_to_parameters!(state, inputs, outputs)
     eliminate_perfect_aliases!(state)
+    remove_constant_variables!(state; kwargs...)
     StateSelection.trivial_tearing!(state)
     sys, mm = ModelingToolkit.alias_elimination!(state; fully_determined, kwargs...)
     if check_consistency

test/reduction.jl

@@ -353,3 +353,13 @@ ss = mtkcompile(sys)
     @mtkcompile sys = System([D(x) ~ 2x, y ~ x], t; state_priorities = [y => 10])
     @test isequal(only(unknowns(sys)), y)
 end
+
+@testset "Constant equations are removed" begin
+    @variables x(t) y(t) z(t)
+    @named sys = System([0 ~ 2x + 3t + 4, 0 ~ x * y + 2, 0 ~ D(x) + D(z) + 2z], t)
+    ts = TearingState(sys)
+    ModelingToolkit.remove_constant_variables!(ts)
+    dx = ModelingToolkit.default_toterm(unwrap(D(x)))
+    @test isequal(ts.additional_observed, [x ~ (3t + 4) / -2, dx ~ (-3//2), y ~ 2 / (-x)])
+    @test isequal(equations(ts), [0 ~ D(z) + 2z - 3/2])
+end