WIP: upgrade to SymbolicUtils w/ fast terms

src/ModelingToolkit.jl

@@ -23,7 +23,9 @@ RuntimeGeneratedFunctions.init(@__MODULE__)
 using RecursiveArrayTools
 
 import SymbolicUtils
-import SymbolicUtils: Term, Sym, to_symbolic, FnType, @rule, Rewriters, substitute, similarterm
+import SymbolicUtils: Term, Add, Mul, Pow, Sym, to_symbolic, FnType, @rule, Rewriters, substitute, similarterm
+
+import SymbolicUtils.Rewriters: Chain, Postwalk, Prewalk, Fixpoint
 
 using LinearAlgebra: LU, BlasInt
 
@@ -72,13 +74,23 @@ Base.show(io::IO, n::Num) = show_numwrap[] ? print(io, :(Num($(value(n))))) : Ba
 
 Base.promote_rule(::Type{<:Number}, ::Type{<:Num}) = Num
 Base.promote_rule(::Type{<:Symbolic{<:Number}}, ::Type{<:Num}) = Num
-Base.getproperty(t::Term, f::Symbol) = f === :op ? operation(t) : f === :args ? arguments(t) : getfield(t, f)
+function Base.getproperty(t::Union{Add, Mul, Pow, Term}, f::Symbol)
+    if f === :op
+        Base.depwarn("`x.op` is deprecated, use `operation(x)` instead", :getproperty, force=true)
+        operation(t)
+    elseif f === :args
+        Base.depwarn("`x.args` is deprecated, use `arguments(x)` instead", :getproperty, force=true)
+        arguments(t)
+    else
+        getfield(t, f)
+    end
+end
 <ₑ(s::Num, x) = value(s) <ₑ value(x)
 <ₑ(s, x::Num) = value(s) <ₑ value(x)
 <ₑ(s::Num, x::Num) = value(s) <ₑ value(x)
 
 for T in (Integer, Rational)
-    @eval Base.:(^)(n::Num, i::$T) = Num(Term{symtype(n)}(^, [value(n),i]))
+    @eval Base.:(^)(n::Num, i::$T) = Num(value(n)^i)
 end
 
 macro num_method(f, expr, Ts=nothing)

src/build_function.jl

@@ -479,25 +479,25 @@ get_varnumber(varop, vars::Vector) =  findfirst(x->isequal(x,varop),vars)
 function numbered_expr(O::Union{Term,Sym},args...;varordering = args[1],offset = 0,
                        lhsname=gensym("du"),rhsnames=[gensym("MTK") for i in 1:length(args)])
     O = value(O)
-  if O isa Sym || isa(O.op, Sym)
+  if O isa Sym || isa(operation(O), Sym)
 	for j in 1:length(args)
 		i = get_varnumber(O,args[j])
 		if i !== nothing
 			return :($(rhsnames[j])[$(i+offset)])
 		end
 	end
   end
-  return Expr(:call, O isa Sym ? tosymbol(O, escape=false) : Symbol(O.op),
+  return Expr(:call, O isa Sym ? tosymbol(O, escape=false) : Symbol(operation(O)),
          [numbered_expr(x,args...;offset=offset,lhsname=lhsname,
-                        rhsnames=rhsnames,varordering=varordering) for x in O.args]...)
+                        rhsnames=rhsnames,varordering=varordering) for x in arguments(O)]...)
 end
 
 function numbered_expr(de::ModelingToolkit.Equation,args...;varordering = args[1],
                        lhsname=gensym("du"),rhsnames=[gensym("MTK") for i in 1:length(args)],offset=0)
 
     varordering = value.(args[1])
     var = var_from_nested_derivative(de.lhs)[1]
-    i = findfirst(x->isequal(tosymbol(x isa Sym ? x : x.op, escape=false), tosymbol(var, escape=false)),varordering)
+    i = findfirst(x->isequal(tosymbol(x isa Sym ? x : operation(x), escape=false), tosymbol(var, escape=false)),varordering)
     :($lhsname[$(i+offset)] = $(numbered_expr(de.rhs,args...;offset=offset,
 											  varordering = varordering,
 											  lhsname = lhsname,

src/context_dsl.jl

@@ -1,10 +1,14 @@
-import SymbolicUtils: symtype
+import SymbolicUtils: symtype, term
 struct Parameter{T} end
 
 isparameter(x) = false
 isparameter(::Sym{<:Parameter}) = true
 isparameter(::Sym{<:FnType{<:Any, <:Parameter}}) = true
 
+SymbolicUtils.@number_methods(Sym{Parameter{Real}},
+                              term(f, a),
+                              term(f, a, b), skipbasics)
+
 SymbolicUtils.symtype(s::Symbolic{Parameter{T}}) where T = T
 SymbolicUtils.similarterm(t::Term{T}, f, args) where {T<:Parameter} = Term{T}(f, args)
 

src/differentials.jl

@@ -35,13 +35,14 @@ Base.show(io::IO, D::Differential) = print(io, "(D'~", D.x, ")")
 Base.:(==)(D1::Differential, D2::Differential) = isequal(D1.x, D2.x)
 
 _isfalse(occ::Bool) = occ === false
-_isfalse(occ::Term) = _isfalse(occ.op)
+_isfalse(occ::Term) = _isfalse(operation(occ))
 
-function occursin_info(x, expr::Term)
+function occursin_info(x, expr)
+    !istree(expr) && return false
     if isequal(x, expr)
         true
     else
-        args = map(a->occursin_info(x, a), expr.args)
+        args = map(a->occursin_info(x, a), arguments(expr))
         if all(_isfalse, args)
             return false
         end
@@ -52,66 +53,65 @@ function occursin_info(x, expr::Sym)
     isequal(x, expr)
 end
 
-hasderiv(O::Term) = O.op isa Differential || any(hasderiv, O.args)
-hasderiv(O) = false
-
-occursin_info(x, y) = false
+function hasderiv(O)
+    istree(O) ? operation(O) isa Differential || any(hasderiv, arguments(O)) : false
+end
 """
 $(SIGNATURES)
 
 TODO
 """
-function expand_derivatives(O::Term, simplify=true; occurances=nothing)
-    if isa(O.op, Differential)
-        @assert length(O.args) == 1
-        arg = expand_derivatives(O.args[1], false)
+function expand_derivatives(O::Symbolic, simplify=false; occurances=nothing)
+    if istree(O) && isa(operation(O), Differential)
+        @assert length(arguments(O)) == 1
+        arg = expand_derivatives(arguments(O)[1], false)
 
         if occurances == nothing
-            occurances = occursin_info(O.op.x, arg)
+            occurances = occursin_info(operation(O).x, arg)
         end
 
         _isfalse(occurances) && return 0
         occurances isa Bool && return 1 # means it's a `true`
 
-        (D, o) = (O.op, arg)
+        D = operation(O)
 
-        if !isa(o, Term)
-            return O # Cannot expand
-        elseif isa(o.op, Sym)
-            return O # Cannot expand
-        elseif isa(o.op, Differential)
+        if !istree(arg)
+            return D(arg) # Cannot expand
+        elseif isa(operation(arg), Sym)
+            return D(arg) # Cannot expand
+        elseif isa(operation(arg), Differential)
             # The recursive expand_derivatives was not able to remove
             # a nested Differential. We can attempt to differentiate the
             # inner expression wrt to the outer iv. And leave the
             # unexpandable Differential outside.
-            if isequal(o.op.x, D.x)
-                return O
+            if isequal(operation(arg).x, D.x)
+                return D(arg)
             else
-                inner = expand_derivatives(D(o.args[1]), false)
+                inner = expand_derivatives(D(arguments(arg)[1]), false)
                 # if the inner expression is not expandable either, return
-                if inner isa Term && operation(inner) isa Differential
-                    return O
+                if istree(inner) && operation(inner) isa Differential
+                    return D(arg)
                 else
-                    return expand_derivatives(o.op(inner), simplify)
+                    return expand_derivatives(operation(arg)(inner), simplify)
                 end
             end
         end
 
-        l = length(o.args)
+        l = length(arguments(arg))
         exprs = []
         c = 0
 
         for i in 1:l
-            t2 = expand_derivatives(D(o.args[i]),false, occurances=occurances.args[i])
+            t2 = expand_derivatives(D(arguments(arg)[i]),false, occurances=arguments(occurances)[i])
 
             x = if _iszero(t2)
                 t2
             elseif _isone(t2)
-                d = derivative_idx(o, i)
-                d isa NoDeriv ? D(o) : d
+                d = derivative_idx(arg, i)
+                d isa NoDeriv ? D(arg) : d
             else
-                t1 = derivative_idx(o, i)
-                t1 = t1 isa NoDeriv ? D(o) : t1
+                t1 = derivative_idx(arg, i)
+                t1 = t1 isa NoDeriv ? D(arg) : t1
                 make_operation(*, [t1, t2])
             end
 
@@ -136,8 +136,8 @@ function expand_derivatives(O::Term, simplify=true; occurances=nothing)
     elseif !hasderiv(O)
         return O
     else
-        args = map(a->expand_derivatives(a, false), O.args)
-        O1 = make_operation(O.op, args)
+        args = map(a->expand_derivatives(a, false), arguments(O))
+        O1 = make_operation(operation(O), args)
         return simplify ? SymbolicUtils.simplify(O1) : O1
     end
 end
@@ -176,7 +176,7 @@ chain rule is not applied:
 julia> myop = sin(x) * y^2
 sin(x()) * y() ^ 2
 
-julia> typeof(myop.op)  # Op is multiplication function
+julia> typeof(operation(myop))  # Op is multiplication function
 typeof(*)
 
 julia> ModelingToolkit.derivative_idx(myop, 1)  # wrt. sin(x)
@@ -187,7 +187,9 @@ sin(x())
 ```
 """
 derivative_idx(O::Any, ::Any) = 0
-derivative_idx(O::Term, idx) = derivative(O.op, (O.args...,), Val(idx))
+function derivative_idx(O::Symbolic, idx)
+    istree(O) ? derivative(operation(O), (arguments(O)...,), Val(idx)) : 0
+end
 
 # Indicate that no derivative is defined.
 struct NoDeriv

src/direct.jl

@@ -128,7 +128,8 @@ let
 
     _scalar = one(TermCombination)
 
-    linearity_propagator = [
+    simterm(t, f, args) = Term{Any}(f, args)
+    linearity_rules = [
           @rule +(~~xs) => reduce(+, filter(isidx, ~~xs), init=_scalar)
           @rule *(~~xs) => reduce(*, filter(isidx, ~~xs), init=_scalar)
           @rule (~f)(~x::(!isidx)) => _scalar
@@ -146,7 +147,8 @@ let
               else
                   error("Function of unknown linearity used: ", ~f)
               end
-          end] |> Rewriters.Chain |> Rewriters.Postwalk |> Rewriters.Fixpoint
+          end]
+    linearity_propagator = Fixpoint(Postwalk(Chain(linearity_rules); similarterm=simterm))
 
     global hessian_sparsity
 
@@ -164,9 +166,9 @@ let
         u = map(value, u)
         idx(i) = TermCombination(Set([Dict(i=>1)]))
         dict = Dict(u .=> idx.(1:length(u)))
-        found = []
-        f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x)(f)
-        _sparse(linearity_propagator(f), length(u))
+        f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x; similarterm=simterm)(f)
+        lp = linearity_propagator(f)
+        _sparse(lp, length(u))
     end
 end
 
@@ -213,22 +215,22 @@ function sparsehessian(O, vars::AbstractVector; simplify = true)
     return H
 end
 
-function toexpr(O::Term)
-  if isa(O.op, Differential)
-     return :(derivative($(toexpr(O.args[1])),$(toexpr(O.op.x))))
-  elseif isa(O.op, Sym)
-    isempty(O.args) && return O.op.name
-    return Expr(:call, toexpr(O.op), toexpr.(O.args)...)
-  end
-  if O.op === (^)
-      if length(O.args) > 1  && O.args[2] isa Number && O.args[2] < 0
-          return Expr(:call, ^, Expr(:call, inv, toexpr(O.args[1])), -(O.args[2]))
-      end
-  end
-  return Expr(:call, O.op, toexpr.(O.args)...)
+function toexpr(O)
+    !istree(O) && return O
+    if isa(operation(O), Differential)
+        return :(derivative($(toexpr(arguments(O)[1])),$(toexpr(operation(O).x))))
+    elseif isa(operation(O), Sym)
+        isempty(arguments(O)) && return operation(O).name
+        return Expr(:call, toexpr(operation(O)), toexpr.(arguments(O))...)
+    end
+    if operation(O) === (^)
+        if length(arguments(O)) > 1  && arguments(O)[2] isa Number && arguments(O)[2] < 0
+            return Expr(:call, ^, Expr(:call, inv, toexpr(arguments(O)[1])), -(arguments(O)[2]))
+        end
+    end
+    return Expr(:call, operation(O), toexpr.(arguments(O))...)
 end
 toexpr(s::Sym) = nameof(s)
-toexpr(s) = s
 
 function toexpr(eq::Equation)
     Expr(:(=), toexpr(eq.lhs), toexpr(eq.rhs))

src/latexify_recipes.jl

@@ -12,15 +12,15 @@ prettify_expr(expr::Expr) = Expr(expr.head, prettify_expr.(expr.args)...)
 
     rhs = getfield.(eqs, :rhs)
     rhs = prettify_expr.(toexpr.(rhs))
-    rhs = [postwalk(x -> x isa Expr && length(x.args) == 1 ? x.args[1] : x, eq) for eq in rhs]
-    rhs = [postwalk(x -> x isa Expr && x.args[1] == :derivative && length(x.args[2].args) == 2 ? :($(Symbol(:d, x.args[2]))/($(Symbol(:d, x.args[2].args[2])))) : x, eq) for eq in rhs]
-    rhs = [postwalk(x -> x isa Expr && x.args[1] == :derivative ? "\\frac{d\\left($(Latexify.latexraw(x.args[2]))\\right)}{d$(Latexify.latexraw(x.args[3]))}" : x, eq) for eq in rhs]
+    rhs = [postwalk(x -> x isa Expr && length(arguments(x)) == 1 ? arguments(x)[1] : x, eq) for eq in rhs]
+    rhs = [postwalk(x -> x isa Expr && arguments(x)[1] == :derivative && length(arguments(x)[2].args) == 2 ? :($(Symbol(:d, arguments(x)[2]))/($(Symbol(:d, arguments(x)[2].args[2])))) : x, eq) for eq in rhs]
+    rhs = [postwalk(x -> x isa Expr && arguments(x)[1] == :derivative ? "\\frac{d\\left($(Latexify.latexraw(arguments(x)[2]))\\right)}{d$(Latexify.latexraw(arguments(x)[3]))}" : x, eq) for eq in rhs]
 
     lhs = getfield.(eqs, :lhs)
     lhs = prettify_expr.(toexpr.(lhs))
-    lhs = [postwalk(x -> x isa Expr && length(x.args) == 1 ? x.args[1] : x, eq) for eq in lhs]
-    lhs = [postwalk(x -> x isa Expr && x.args[1] == :derivative && length(x.args[2].args) == 2 ? :($(Symbol(:d, x.args[2]))/($(Symbol(:d, x.args[2].args[2])))) : x, eq) for eq in lhs]
-    lhs = [postwalk(x -> x isa Expr && x.args[1] == :derivative ? "\\frac{d\\left($(Latexify.latexraw(x.args[2]))\\right)}{d$(Latexify.latexraw(x.args[3]))}" : x, eq) for eq in lhs]
+    lhs = [postwalk(x -> x isa Expr && length(arguments(x)) == 1 ? arguments(x)[1] : x, eq) for eq in lhs]
+    lhs = [postwalk(x -> x isa Expr && arguments(x)[1] == :derivative && length(arguments(x)[2].args) == 2 ? :($(Symbol(:d, arguments(x)[2]))/($(Symbol(:d, arguments(x)[2].args[2])))) : x, eq) for eq in lhs]
+    lhs = [postwalk(x -> x isa Expr && arguments(x)[1] == :derivative ? "\\frac{d\\left($(Latexify.latexraw(arguments(x)[2]))\\right)}{d$(Latexify.latexraw(arguments(x)[3]))}" : x, eq) for eq in lhs]
 
     return lhs, rhs
 end

src/linearity.jl

@@ -78,6 +78,14 @@ function Base.:(==)(comb1::TermCombination, comb2::TermCombination)
 end
 =#
 
+# to make Mul and Add work
+Base.:*(::Number, comb::TermCombination) = comb
+function Base.:^(comb::TermCombination, ::Number)
+    isone(comb) && return comb
+    iszero(comb) && return _scalar
+    return comb *  comb
+end
+
 function Base.:+(comb1::TermCombination, comb2::TermCombination)
     if isone(comb1) && !iszero(comb2)
         return comb2

src/systems/abstractsystem.jl

@@ -118,7 +118,7 @@ Generate a function to evaluate the system's equations.
 function generate_function end
 
 getname(x::Sym) = nameof(x)
-getname(t::Term) = t.op isa Sym ? getname(t.op) : error("Cannot get name of $t")
+getname(t::Term) = operation(t) isa Sym ? getname(operation(t)) : error("Cannot get name of $t")
 
 function Base.getproperty(sys::AbstractSystem, name::Symbol)
 
@@ -161,7 +161,7 @@ function renamespace(namespace, x::Sym)
 end
 
 function renamespace(namespace, x::Term)
-    renamespace(namespace, x.op)(x.args...)
+    renamespace(namespace, operation(x))(arguments(x)...)
 end
 
 function namespace_variables(sys::AbstractSystem)
@@ -189,10 +189,10 @@ function namespace_expr(O::Sym,name,ivname)
 end
 
 function namespace_expr(O::Term{T},name,ivname) where {T}
-    if O.op isa Sym
-        Term{T}(rename(O.op,renamespace(name,O.op.name)),namespace_expr.(O.args,name,ivname))
+    if operation(O) isa Sym
+        Term{T}(rename(operation(O),renamespace(name,operation(O).name)),namespace_expr.(arguments(O),name,ivname))
     else
-        Term{T}(O.op,namespace_expr.(O.args,name,ivname))
+        Term{T}(operation(O),namespace_expr.(arguments(O),name,ivname))
     end
 end
 namespace_expr(O,name,ivname) = O
@@ -278,11 +278,11 @@ struct AbstractSysToExpr
 end
 AbstractSysToExpr(sys) = AbstractSysToExpr(sys,states(sys))
 function (f::AbstractSysToExpr)(O::Term)
-    any(isequal(O), f.states) && return O.op.name  # variables
-    if isa(O.op, Sym)
-        return build_expr(:call, Any[O.op.name; f.(O.args)])
+    any(isequal(O), f.states) && return operation(O).name  # variables
+    if isa(operation(O), Sym)
+        return build_expr(:call, Any[operation(O).name; f.(arguments(O))])
     end
-    return build_expr(:call, Any[O.op; f.(O.args)])
+    return build_expr(:call, Any[operation(O); f.(arguments(O))])
 end
 (f::AbstractSysToExpr)(x) = toexpr(x)