refactor: use fixpoint_sub and fast_substitute from Symbolics

Author: AayushSabharwal

Project.toml

@@ -104,7 +104,7 @@ SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicIndexingInterface = "0.3.11"
 SymbolicUtils = "1.0"
-Symbolics = "5.24"
+Symbolics = "5.26"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"

src/ModelingToolkit.jl

@@ -56,7 +56,7 @@ using PrecompileTools, Reexport
                      VariableSource, getname, variable, Connection, connect,
                      NAMESPACE_SEPARATOR, set_scalar_metadata, setdefaultval,
                      initial_state, transition, activeState, entry,
-                     ticksInState, timeInState
+                     ticksInState, timeInState, fixpoint_sub, fast_substitute
     import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
                       jacobian_sparsity, isaffine, islinear, _iszero, _isone,
                       tosymbol, lower_varname, diff2term, var_from_nested_derivative,

src/structural_transformation/StructuralTransformations.jl

@@ -3,7 +3,7 @@ module StructuralTransformations
 using Setfield: @set!, @set
 using UnPack: @unpack
 
-using Symbolics: unwrap, linear_expansion
+using Symbolics: unwrap, linear_expansion, fast_substitute
 using SymbolicUtils
 using SymbolicUtils.Code
 using SymbolicUtils.Rewriters
@@ -23,7 +23,7 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
                        invalidate_cache!, Substitutions, get_or_construct_tearing_state,
                        filter_kwargs, lower_varname, setio, SparseMatrixCLIL,
-                       fast_substitute, get_fullvars, has_equations, observed,
+                       get_fullvars, has_equations, observed,
                        Schedule
 
 using ModelingToolkit.BipartiteGraphs

src/structural_transformation/symbolics_tearing.jl

@@ -81,7 +81,7 @@ function eq_derivative!(ts::TearingState{ODESystem}, ieq::Int)
 end
 
 function tearing_sub(expr, dict, s)
-    expr = ModelingToolkit.fixpoint_sub(expr, dict)
+    expr = Symbolics.fixpoint_sub(expr, dict)
     s ? simplify(expr) : expr
 end
 
@@ -439,7 +439,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
                 order, lv = var_order(iv)
                 dx = D(simplify_shifts(lower_varname_withshift(
                     fullvars[lv], idep, order - 1)))
-                eq = dx ~ simplify_shifts(ModelingToolkit.fixpoint_sub(
+                eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.solve_for(neweqs[ieq],
                         fullvars[iv]),
                     total_sub; operator = ModelingToolkit.Shift))
@@ -467,7 +467,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
                 @warn "Tearing: solving $eq for $var is singular!"
             else
                 rhs = -b / a
-                neweq = var ~ ModelingToolkit.fixpoint_sub(
+                neweq = var ~ Symbolics.fixpoint_sub(
                     simplify ?
                     Symbolics.simplify(rhs) : rhs,
                     total_sub; operator = ModelingToolkit.Shift)
@@ -481,7 +481,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
             if !(eq.lhs isa Number && eq.lhs == 0)
                 rhs = eq.rhs - eq.lhs
             end
-            push!(alge_eqs, 0 ~ ModelingToolkit.fixpoint_sub(rhs, total_sub))
+            push!(alge_eqs, 0 ~ Symbolics.fixpoint_sub(rhs, total_sub))
             push!(algeeq_idxs, ieq)
         end
     end

src/systems/alias_elimination.jl

@@ -462,13 +462,3 @@ function observed2graph(eqs, unknowns)
 
     return graph, assigns
 end
-
-function fixpoint_sub(x, dict; operator = Nothing)
-    y = fast_substitute(x, dict; operator)
-    while !isequal(x, y)
-        y = x
-        x = fast_substitute(y, dict; operator)
-    end
-
-    return x
-end

src/utils.jl

@@ -799,86 +799,6 @@ function fold_constants(ex)
     end
 end
 
-# Symbolics needs to call unwrap on the substitution rules, but most of the time
-# we don't want to do that in MTK.
-const Eq = Union{Equation, Inequality}
-function fast_substitute(eq::Eq, subs; operator = Nothing)
-    if eq isa Inequality
-        Inequality(fast_substitute(eq.lhs, subs; operator),
-            fast_substitute(eq.rhs, subs; operator),
-            eq.relational_op)
-    else
-        Equation(fast_substitute(eq.lhs, subs; operator),
-            fast_substitute(eq.rhs, subs; operator))
-    end
-end
-function fast_substitute(eq::T, subs::Pair; operator = Nothing) where {T <: Eq}
-    T(fast_substitute(eq.lhs, subs; operator), fast_substitute(eq.rhs, subs; operator))
-end
-function fast_substitute(eqs::AbstractArray, subs; operator = Nothing)
-    fast_substitute.(eqs, (subs,); operator)
-end
-function fast_substitute(eqs::AbstractArray, subs::Pair; operator = Nothing)
-    fast_substitute.(eqs, (subs,); operator)
-end
-for (exprType, subsType) in Iterators.product((Num, Symbolics.Arr), (Any, Pair))
-    @eval function fast_substitute(expr::$exprType, subs::$subsType; operator = Nothing)
-        fast_substitute(value(expr), subs; operator)
-    end
-end
-function fast_substitute(expr, subs; operator = Nothing)
-    if (_val = get(subs, expr, nothing)) !== nothing
-        return _val
-    end
-    istree(expr) || return expr
-    op = fast_substitute(operation(expr), subs; operator)
-    args = SymbolicUtils.unsorted_arguments(expr)
-    if !(op isa operator)
-        canfold = Ref(!(op isa Symbolic))
-        args = let canfold = canfold
-            map(args) do x
-                x′ = fast_substitute(x, subs; operator)
-                canfold[] = canfold[] && !(x′ isa Symbolic)
-                x′
-            end
-        end
-        canfold[] && return op(args...)
-    end
-    similarterm(expr,
-        op,
-        args,
-        symtype(expr);
-        metadata = metadata(expr))
-end
-function fast_substitute(expr, pair::Pair; operator = Nothing)
-    a, b = pair
-    isequal(expr, a) && return b
-    if a isa AbstractArray
-        for (ai, bi) in zip(a, b)
-            expr = fast_substitute(expr, ai => bi; operator)
-        end
-    end
-    istree(expr) || return expr
-    op = fast_substitute(operation(expr), pair; operator)
-    args = SymbolicUtils.unsorted_arguments(expr)
-    if !(op isa operator)
-        canfold = Ref(!(op isa Symbolic))
-        args = let canfold = canfold
-            map(args) do x
-                x′ = fast_substitute(x, pair; operator)
-                canfold[] = canfold[] && !(x′ isa Symbolic)
-                x′
-            end
-        end
-        canfold[] && return op(args...)
-    end
-    similarterm(expr,
-        op,
-        args,
-        symtype(expr);
-        metadata = metadata(expr))
-end
-
 normalize_to_differential(s) = s
 
 function restrict_array_to_union(arr)