[WIP] refactor: update to SymbolicUtils@4 and Symbolics@7

Project.toml

@@ -8,7 +8,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -33,7 +32,6 @@ FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 ImplicitDiscreteSolve = "3263718b-31ed-49cf-8a0f-35a466e8af96"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -68,18 +66,22 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 [weakdeps]
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 CasADi = "c49709b8-5c63-11e9-2fb2-69db5844192f"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816"
 
 [extensions]
 MTKBifurcationKitExt = "BifurcationKit"
 MTKCasADiDynamicOptExt = "CasADi"
+MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
+MTKJuliaFormatterExt = "JuliaFormatter"
 MTKLabelledArraysExt = "LabelledArrays"
 MTKPyomoDynamicOptExt = "Pyomo"
 

docs/src/basics/FAQ.md

@@ -192,7 +192,7 @@ p, replace, alias = SciMLStructures.canonicalize(Tunable(), prob.p)
 # changes to the array will be reflected in parameter values
 ```
 
-# ERROR: ArgumentError: SymbolicUtils.BasicSymbolic{Real}[xˍt(t)] are missing from the variable map.
+# ERROR: ArgumentError: `[xˍt(t)]` are missing from the variable map.
 
 This error can come up after running `mtkcompile` on a system that generates dummy derivatives (i.e. variables with `ˍt`).  For example, here even though all the variables are defined with initial values, the `ODEProblem` generation will throw an error that defaults are missing from the variable map.
 

ext/MTKCasADiDynamicOptExt.jl

@@ -122,7 +122,7 @@ end
 
 function MTK.lowered_var(m::CasADiModel, uv, i, t)
     X = getfield(m, uv)
-    t isa Union{Num, Symbolics.Symbolic} ? X.u[i, :] : X(t)[i]
+    t isa Union{Num, SymbolicT} ? X.u[i, :] : X(t)[i]
 end
 
 function MTK.lowered_integral(model::CasADiModel, expr, lo, hi)

src/adjoints.jl → ext/MTKChainRulesCoreExt.jl

@@ -1,3 +1,13 @@
+module MTKChainRulesCoreExt
+
+import ChainRulesCore
+import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
+using ModelingToolkit: MTKParameters, NONNUMERIC_PORTION, AbstractSystem
+import ModelingToolkit as MTK
+import SciMLStructures
+import SymbolicIndexingInterface: remake_buffer
+import SciMLBase: AbstractNonlinearProblem, remake
+
 function ChainRulesCore.rrule(::Type{MTKParameters}, tunables, args...)
     function mtp_pullback(dt)
         dt = unthunk(dt)
@@ -104,3 +114,11 @@ function ChainRulesCore.rrule(
 end
 
 ChainRulesCore.@non_differentiable Base.getproperty(sys::AbstractSystem, x::Symbol)
+
+function ModelingToolkit.update_initializeprob!(initprob::AbstractNonlinearProblem, prob)
+    pgetter = ChainRulesCore.@ignore_derivatives MTK.get_scimlfn(prob).initialization_data.metadata.oop_reconstruct_u0_p.pgetter
+    p = pgetter(prob, initprob)
+    return remake(initprob; p)
+end
+
+end

ext/MTKInfiniteOptExt.jl

@@ -122,7 +122,7 @@ end
 
 function MTK.lowered_var(m::InfiniteOptModel, uv, i, t)
     X = getfield(m, uv)
-    t isa Union{Num, Symbolics.Symbolic} ? X[i] : X[i](t)
+    t isa Union{Num, SymbolicT} ? X[i] : X[i](t)
 end
 
 function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
@@ -256,13 +256,13 @@ for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqr
 end
 
 # JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both.
-function Base.isequal(::SymbolicUtils.Symbolic,
+function Base.isequal(::SymbolicT,
         ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr})
     false
 end
 function Base.isequal(
         ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr},
-        ::SymbolicUtils.Symbolic)
+        ::SymbolicT)
     false
 end
 end

ext/MTKJuliaFormatterExt.jl

@@ -0,0 +1,12 @@
+module MTKJuliaFormatterExt
+
+import ModelingToolkit: readable_code, _readable_code, rec_remove_macro_linenums!
+import JuliaFormatter
+
+function readable_code(expr::Expr)
+    expr = Base.remove_linenums!(_readable_code(expr))
+    rec_remove_macro_linenums!(expr)
+    JuliaFormatter.format_text(string(expr), JuliaFormatter.SciMLStyle())
+end
+
+end

ext/MTKPyomoDynamicOptExt.jl

@@ -53,7 +53,7 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
     end
 end
 
-function pysym_getproperty(s::Union{Num, Symbolics.Symbolic}, name::Symbol)
+function pysym_getproperty(s::Union{Num, SymbolicT}, name::Symbol)
     Symbolics.wrap(SymbolicUtils.term(
         _getproperty, Symbolics.unwrap(s), Val{name}(), type = Symbolics.Struct{PyomoVar}))
 end
@@ -112,7 +112,7 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1)
         Symbolics.unwrap(expr), SPECIAL_FUNCTIONS_DICT, fold = false)
 
     cons_sym = Symbol("cons", hash(cons))
-    if occursin(Symbolics.unwrap(t_sym), expr)
+    if SU.query!(isequal(Symbolics.unwrap(t_sym)), expr)
         f = eval(Symbolics.build_function(expr, model_sym, t_sym))
         setproperty!(model, cons_sym, pyomo.Constraint(model.t, rule = Pyomo.pyfunc(f)))
     else
@@ -124,7 +124,7 @@ end
 function MTK.set_objective!(pmodel::PyomoDynamicOptModel, expr)
     @unpack model, model_sym, t_sym, dummy_sym = pmodel
     expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT, fold = false)
-    if occursin(Symbolics.unwrap(t_sym), expr)
+    if SU.query!(isequal(Symbolics.unwrap(t_sym)), expr)
         f = eval(Symbolics.build_function(expr, model_sym, t_sym))
         model.obj = pyomo.Objective(model.t, rule = Pyomo.pyfunc(f))
     else
@@ -165,7 +165,7 @@ end
 
 function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t)
     X = Symbolics.value(pysym_getproperty(m.model_sym, uv))
-    var = t isa Union{Num, Symbolics.Symbolic} ? X[i, m.t_sym] : X[i, t]
+    var = t isa Union{Num, SymbolicT} ? X[i, m.t_sym] : X[i, t]
     Symbolics.unwrap(var)
 end
 

src/ModelingToolkit.jl

@@ -9,16 +9,18 @@ using PrecompileTools, Reexport
 end
 
 import SymbolicUtils
+import SymbolicUtils as SU
 import SymbolicUtils: iscall, arguments, operation, maketerm, promote_symtype,
-                      Symbolic, isadd, ismul, ispow, issym, FnType,
-                      @rule, Rewriters, substitute, metadata, BasicSymbolic,
-                      Sym, Term
+                      isadd, ismul, ispow, issym, FnType, isconst, BSImpl,
+                      @rule, Rewriters, substitute, metadata, BasicSymbolic
 using SymbolicUtils.Code
 import SymbolicUtils.Code: toexpr
 import SymbolicUtils.Rewriters: Chain, Postwalk, Prewalk, Fixpoint
 using DocStringExtensions
 using SpecialFunctions, NaNMath
-using DiffEqCallbacks
+@recompile_invalidations begin
+    using DiffEqCallbacks
+end
 using Graphs
 import ExprTools: splitdef, combinedef
 import OrderedCollections
@@ -48,7 +50,6 @@ using SciMLBase: StandardODEProblem, StandardNonlinearProblem, handle_varmap, Ti
                  PeriodicClock, Clock, SolverStepClock, ContinuousClock, OverrideInit,
                  NoInit
 using Distributed
-import JuliaFormatter
 using MLStyle
 import Moshi
 using Moshi.Data: @data
@@ -62,19 +63,16 @@ import BlockArrays: BlockArray, BlockedArray, Block, blocksize, blocksizes, bloc
 using OffsetArrays: Origin
 import CommonSolve
 import EnumX
-import ChainRulesCore
-import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
 
 using RuntimeGeneratedFunctions
 using RuntimeGeneratedFunctions: drop_expr
 
-using Symbolics: degree
-using Symbolics: _parse_vars, value, @derivatives, get_variables,
-                 exprs_occur_in, symbolic_linear_solve, build_expr, unwrap, wrap,
+using Symbolics: degree, VartypeT, SymbolicT
+using Symbolics: parse_vars, value, @derivatives, get_variables,
+                 exprs_occur_in, symbolic_linear_solve, unwrap, wrap,
                  VariableSource, getname, variable,
-                 NAMESPACE_SEPARATOR, set_scalar_metadata, setdefaultval,
-                 hasnode, fixpoint_sub, fast_substitute,
-                 CallWithMetadata, CallWithParent
+                 NAMESPACE_SEPARATOR, setdefaultval,
+                 hasnode, fixpoint_sub, CallAndWrap, SArgsT, SSym, STerm
 const NAMESPACE_SEPARATOR_SYMBOL = Symbol(NAMESPACE_SEPARATOR)
 import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
                   jacobian_sparsity, isaffine, islinear, _iszero, _isone,
@@ -83,7 +81,7 @@ import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
                   ParallelForm, SerialForm, MultithreadedForm, build_function,
                   rhss, lhss, prettify_expr, gradient,
                   jacobian, hessian, derivative, sparsejacobian, sparsehessian,
-                  substituter, scalarize, getparent, hasderiv, hasdiff
+                  scalarize, hasderiv
 
 import DiffEqBase: @add_kwonly
 export independent_variables, unknowns, observables, parameters, full_parameters,
@@ -157,6 +155,9 @@ include("parameters.jl")
 include("independent_variables.jl")
 include("constants.jl")
 
+const SymmapT = Dict{SymbolicT, SymbolicT}
+const COMMON_NOTHING = SU.Const{VartypeT}(nothing)
+
 include("utils.jl")
 
 include("systems/index_cache.jl")
@@ -225,7 +226,6 @@ include("structural_transformation/StructuralTransformations.jl")
 @reexport using .StructuralTransformations
 include("inputoutput.jl")
 
-include("adjoints.jl")
 include("deprecations.jl")
 
 const t_nounits = let
@@ -334,6 +334,8 @@ export AbstractCollocation, JuMPCollocation, InfiniteOptCollocation,
        CasADiCollocation, PyomoCollocation
 export DynamicOptSolution
 
+const set_scalar_metadata = setmetadata
+
 @public apply_to_variables, equations_toplevel, unknowns_toplevel, parameters_toplevel
 @public continuous_events_toplevel, discrete_events_toplevel, assertions, is_alg_equation
 @public is_diff_equation, Equality, linearize_symbolic, reorder_unknowns
@@ -356,10 +358,100 @@ for prop in [SYS_PROPS; [:continuous_events, :discrete_events]]
 end
 
 PrecompileTools.@compile_workload begin
-    using ModelingToolkit
+        fold1 = Val{false}()
+        using SymbolicUtils
+        using SymbolicUtils: shape
+        using Symbolics
+        @syms x y f(t) q[1:5]
+        SymbolicUtils.Sym{SymReal}(:a; type = Real, shape = SymbolicUtils.ShapeVecT())
+        x + y
+        x * y
+        x / y
+        x ^ y
+        x ^ 5
+        6 ^ x
+        x - y
+        -y
+        2y
+        z = 2
+        dict = SymbolicUtils.ACDict{VartypeT}()
+        dict[x] = 1
+        dict[y] = 1
+        type::typeof(DataType) = rand() < 0.5 ? Real : Float64
+        nt = (; type, shape, unsafe = true)
+        Base.pairs(nt)
+        BSImpl.AddMul{VartypeT}(1, dict, SymbolicUtils.AddMulVariant.MUL; type, shape = SymbolicUtils.ShapeVecT(), unsafe = true)
+        *(y, z)
+        *(z, y)
+        SymbolicUtils.symtype(y)
+        f(x)
+        (5x / 5)
+        expand((x + y) ^ 2)
+        simplify(x ^ (1//2) + (sin(x) ^ 2 + cos(x) ^ 2) + 2(x + y) - x - y)
+        ex = x + 2y + sin(x)
+        rules1 = Dict(x => y)
+        rules2 = Dict(x => 1)
+        Dx = Differential(x)
+        Differential(y)(ex)
+        uex = unwrap(ex)
+        Symbolics.executediff(Dx, uex)
+        # Running `fold = Val(true)` invalidates the precompiled statements
+        # for `fold = Val(false)` and itself doesn't precompile anyway.
+        # substitute(ex, rules1)
+        substitute(ex, rules1; fold = fold1)
+        substitute(ex, rules2; fold = fold1)
+        @variables foo
+        f(foo)
+        @variables x y f(::Real) q[1:5]
+        x + y
+        x * y
+        x / y
+        x ^ y
+        x ^ 5
+        # 6 ^ x
+        x - y
+        -y
+        2y
+        symtype(y)
+        z = 2
+        *(y, z)
+        *(z, y)
+        f(x)
+        (5x / 5)
+        [x, y]
+        [x, f, f]
+        promote_type(Int, Num)
+        promote_type(Real, Num)
+        promote_type(Float64, Num)
+        # expand((x + y) ^ 2)
+        # simplify(x ^ (1//2) + (sin(x) ^ 2 + cos(x) ^ 2) + 2(x + y) - x - y)
+        ex = x + 2y + sin(x)
+        rules1 = Dict(x => y)
+        # rules2 = Dict(x => 1)
+        # Running `fold = Val(true)` invalidates the precompiled statements
+        # for `fold = Val(false)` and itself doesn't precompile anyway.
+        # substitute(ex, rules1)
+        substitute(ex, rules1; fold = fold1)
+        Symbolics.linear_expansion(ex, y)
+        # substitute(ex, rules2; fold = fold1)
+        # substitute(ex, rules2)
+        # substitute(ex, rules1; fold = fold2)
+        # substitute(ex, rules2; fold = fold2)
+        q[1]
+        q'q
+     using ModelingToolkit
     @variables x(ModelingToolkit.t_nounits)
-    @named sys = System([ModelingToolkit.D_nounits(x) ~ -x], ModelingToolkit.t_nounits)
-    prob = ODEProblem(mtkcompile(sys), [x => 30.0], (0, 100), jac = true)
+    isequal(ModelingToolkit.D_nounits.x, ModelingToolkit.t_nounits)
+    sys = System([ModelingToolkit.D_nounits(x) ~ x], ModelingToolkit.t_nounits, [x], Num[]; name = :sys)
+    complete(sys)
+    @syms p[1:2]
+    ndims(p)
+    size(p)
+    axes(p)
+    length(p)
+    v = [p]
+    isempty(v)
+    # mtkcompile(sys)
     @mtkmodel __testmod__ begin
         @constants begin
             c = 1.0
@@ -390,4 +482,17 @@ PrecompileTools.@compile_workload begin
     end
 end
 
+precompile(Tuple{typeof(Base.merge), NamedTuple{(:f, :args, :metadata, :hash, :hash2, :shape, :type, :id), Tuple{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, SymbolicUtils.SmallVec{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, Array{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, 1}}, Nothing, UInt64, UInt64, SymbolicUtils.SmallVec{Base.UnitRange{Int64}, Array{Base.UnitRange{Int64}, 1}}, DataType, SymbolicUtils.IDType}}, NamedTuple{(:metadata,), Tuple{Base.ImmutableDict{DataType, Any}}}})
+precompile(Tuple{typeof(Base.merge), NamedTuple{(:f, :args, :metadata, :hash, :hash2, :shape, :type, :id), Tuple{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, SymbolicUtils.SmallVec{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, Array{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, 1}}, Base.ImmutableDict{DataType, Any}, UInt64, UInt64, SymbolicUtils.SmallVec{Base.UnitRange{Int64}, Array{Base.UnitRange{Int64}, 1}}, DataType, SymbolicUtils.IDType}}, NamedTuple{(:id, :hash, :hash2), Tuple{Nothing, Int64, Int64}}})
+precompile(Tuple{typeof(Core.kwcall), NamedTuple{(:f, :args, :metadata, :hash, :hash2, :shape, :type, :id), Tuple{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, SymbolicUtils.SmallVec{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, Array{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, 1}}, Base.ImmutableDict{DataType, Any}, Int64, Int64, SymbolicUtils.SmallVec{Base.UnitRange{Int64}, Array{Base.UnitRange{Int64}, 1}}, DataType, Nothing}}, Type{SymbolicUtils.BasicSymbolicImpl.Term{SymbolicUtils.SymReal}}})
+precompile(Tuple{typeof(Symbolics.parse_vars), Symbol, Type, Tuple{Symbol, Symbol}, Function})
+precompile(Tuple{typeof(Base.merge), NamedTuple{(:name, :metadata, :hash, :hash2, :shape, :type, :id), Tuple{Symbol, Base.ImmutableDict{DataType, Any}, UInt64, UInt64, SymbolicUtils.SmallVec{Base.UnitRange{Int64}, Array{Base.UnitRange{Int64}, 1}}, DataType, SymbolicUtils.IDType}}, NamedTuple{(:metadata,), Tuple{Base.ImmutableDict{DataType, Any}}}})
+precompile(Tuple{typeof(Base.vect), Symbolics.Equation, Vararg{Symbolics.Equation}})
+precompile(Tuple{typeof(Core.kwcall), NamedTuple{(:name, :defaults), Tuple{Symbol, Base.Dict{Symbolics.Num, Float64}}}, Type{ModelingToolkit.System}, Array{Symbolics.Equation, 1}, Symbolics.Num, Array{Symbolics.Num, 1}, Array{Symbolics.Num, 1}})
+precompile(Tuple{Type{NamedTuple{(:name, :defaults), T} where T<:Tuple}, Tuple{Symbol, Base.Dict{Symbolics.Num, Float64}}})
+precompile(Tuple{typeof(SymbolicUtils.isequal_somescalar), Float64, Float64})
+precompile(Tuple{Type{NamedTuple{(:name, :defaults, :guesses), T} where T<:Tuple}, Tuple{Symbol, Base.Dict{Symbolics.Num, Float64}, Base.Dict{Symbolics.Num, Float64}}})
+precompile(Tuple{typeof(Core.kwcall), NamedTuple{(:name, :defaults, :guesses), Tuple{Symbol, Base.Dict{Symbolics.Num, Float64}, Base.Dict{Symbolics.Num, Float64}}}, Type{ModelingToolkit.System}, Array{Symbolics.Equation, 1}, Symbolics.Num, Array{Symbolics.Num, 1}, Array{Symbolics.Num, 1}})
+precompile(Tuple{typeof(Core.kwcall), NamedTuple{(:type, :shape), Tuple{DataType, SymbolicUtils.SmallVec{Base.UnitRange{Int64}, Array{Base.UnitRange{Int64}, 1}}}}, typeof(SymbolicUtils.term), Any, SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}})
+
 end # module

src/clock.jl

@@ -45,7 +45,7 @@ has_time_domain(_, x) = has_time_domain(x)
 
 Determine if variable `x` has a time-domain attributed to it.
 """
-function has_time_domain(x::Symbolic)
+function has_time_domain(x::SymbolicT)
     # getmetadata(x, ContinuousClock, nothing) !== nothing ||
     # getmetadata(x, Discrete,   nothing) !== nothing
     getmetadata(x, VariableTimeDomain, nothing) !== nothing
@@ -77,7 +77,7 @@ See also [`is_continuous_domain`](@ref)
 """
 function has_continuous_domain(x)
     issym(x) && return is_continuous_domain(x)
-    hasderiv(x) || hasdiff(x) || hassample(x) || hashold(x)
+    hasderiv(x) || hassample(x) || hashold(x)
 end
 
 """