Remove unnecessary specializations, simplify_constants, and add fold_constants

Author: YingboMa

src/structural_transformation/utils.jl

@@ -179,6 +179,8 @@ function find_eq_solvables!(state::TearingState, ieq, to_rm = Int[], coeffs = no
         a, b, islinear = linear_expansion(term, var)
         a, b = unwrap(a), unwrap(b)
         islinear || (all_int_vars = false; continue)
+        a = ModelingToolkit.fold_constants(a)
+        b = ModelingToolkit.fold_constants(b)
         if a isa Symbolic
             all_int_vars = false
             if !allow_symbolic

src/utils.jl

@@ -19,7 +19,7 @@ end
 function detime_dvs(op)
     if !istree(op)
         op
-    elseif operation(op) isa Sym
+    elseif issym(operation(op))
         Sym{Real}(nameof(operation(op)))
     else
         similarterm(op, operation(op), detime_dvs.(arguments(op)))
@@ -60,7 +60,7 @@ function states_to_sym(states::Set)
         elseif istree(O)
             op = operation(O)
             args = arguments(O)
-            if op isa Sym
+            if issym(op)
                 O in states && return tosymbol(O)
                 # dependent variables
                 return build_expr(:call, Any[nameof(op); _states_to_sym.(args)])
@@ -511,7 +511,7 @@ function collect_constants(x)
     return constants
 end
 
-function collect_constants!(constants, arr::AbstractArray{T}) where {T}
+function collect_constants!(constants, arr::AbstractArray)
     for el in arr
         collect_constants!(constants, el)
     end
@@ -526,8 +526,8 @@ collect_constants!(constants, x::Num) = collect_constants!(constants, unwrap(x))
 collect_constants!(constants, x::Real) = nothing
 collect_constants(n::Nothing) = Symbolics.Sym[]
 
-function collect_constants!(constants, expr::Symbolics.Symbolic{T}) where {T}
-    if expr isa Sym && isconstant(expr)
+function collect_constants!(constants, expr::Symbolics.Symbolic)
+    if issym(expr) && isconstant(expr)
         push!(constants, expr)
     else
         evars = vars(expr)
@@ -542,8 +542,7 @@ function collect_constants!(constants, expr::Symbolics.Symbolic{T}) where {T}
 end
 
 """ Replace symbolic constants with their literal values """
-function eliminate_constants(eqs::AbstractArray{<:Union{Equation, Symbolic}},
-                             cs::Vector{Sym})
+function eliminate_constants(eqs, cs)
     cmap = Dict(x => getdefault(x) for x in cs)
     return substitute(eqs, cmap)
 end
@@ -807,6 +806,17 @@ function jacobian_wrt_vars(pf::F, p, input_idxs, chunk::C) where {F, C}
     ForwardDiff.jacobian(p_closure, p_small, cfg, Val(false))
 end
 
+function fold_constants(ex)
+    if istree(ex)
+        similarterm(ex, operations(ex), map(fold_constants, arguments(ex)),
+                    symtype(expr); metadata = metadata(expr))
+    elseif issym(ex) && isconstant(ex)
+        getdefault(ex)
+    else
+        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.
 function fast_substitute(eq::Equation, subs)

test/constants.jl

@@ -19,15 +19,9 @@ newsys = MT.eliminate_constants(sys)
 eqs = [D(x) ~ 1,
     w ~ a]
 @named sys = ODESystem(eqs)
-simp = structural_simplify(sys, simplify_constants = false);
-@test isequal(simp.substitutions.subs[1], eqs[2])
-@test isequal(equations(simp)[1], eqs[1])
-prob = ODEProblem(simp, [0], [0.0, 1.0], [])
-sol = solve(prob, Tsit5())
-@test sol[w][1] == 1
 # Now eliminate the constants first
-simp = structural_simplify(sys, simplify_constants = true);
-@test isequal(simp.substitutions.subs[1], w ~ 1)
+simp = structural_simplify(sys)
+@test equations(simp) == [D(x) ~ 1.0]
 
 #Constant with units
 @constants β=1 [unit = u"m/s"]
@@ -43,4 +37,4 @@ sys = ODESystem(eqs, name = :sys)
 simp = structural_simplify(sys)
 @test_throws MT.ValidationError MT.check_units(simp.eqs...)
 
-@test MT.collect_constants(nothing) == Symbolics.Sym[]
+@test isempty(MT.collect_constants(nothing))