integrate SymbolicUtils

Author: shashi

src/ModelingToolkit.jl

@@ -58,6 +58,7 @@ include("variables.jl")
 include("context_dsl.jl")
 include("operations.jl")
 include("differentials.jl")
+include("symutils.jl")
 
 function Base.convert(::Type{Variable},x::Operation)
     if x.op isa Variable

src/simplify.jl

@@ -1,97 +1,51 @@
-function simplify_constants(O::Operation, shorten_tree)
-    while true
-        O′ = _simplify_constants(O, shorten_tree)
-        if is_operation(O′)
-            O′ = Operation(O′.op, simplify_constants.(O′.args, shorten_tree))
-        end
-        isequal(O, O′) && return O
-        O = O′
+import SymbolicUtils
+import SymbolicUtils: FnType
+
+# ModelingToolkit -> SymbolicUtils
+SymbolicUtils.istree(x::Operation) = true
+function SymbolicUtils.operation(x::Operation)
+    if x.op isa Variable
+        T = FnType{NTuple{length(x.args), Any}, vartype(x.op)}
+        SymbolicUtils.Variable{T}(x.op.name)
+    else
+        x.op
     end
 end
-simplify_constants(x, shorten_tree) = x
 
-"""
-simplify_constants(x::Operation)
+# This is required to infer the right type for
+# Operation(Variable{Parameter{Number}}(:foo), [])
+# While keeping the metadata that the variable is a parameter.
+SymbolicUtils.promote_symtype(f::SymbolicUtils.Variable{FnType{X,Parameter{Y}}},
+                              xs...) where {X, Y} = Y
 
-Simplifies the constants within an expression, for example removing equations
-multiplied by a zero and summing constant values.
-"""
-simplify_constants(x) = simplify_constants(x, true)
+SymbolicUtils.arguments(x::Operation) = x.args
 
-Base.isone(x::Operation) = x.op == one || x.op == Constant && isone(x.args)
-const AC_OPERATORS = (*, +)
+# SymbolicUtils wants raw numbers
+SymbolicUtils.to_symbolic(x::Constant) = x.value
+SymbolicUtils.to_symbolic(x::Variable{T}) where {T} = SymbolicUtils.Variable{T}(x.name)
 
-function _simplify_constants(O::Operation, shorten_tree)
-    # Tree shrinking
-    if shorten_tree && O.op ∈ AC_OPERATORS
-        # Flatten tree
-        idxs = findall(x -> is_operation(x) && x.op === O.op, O.args)
-        if !isempty(idxs)
-            keep_idxs = eachindex(O.args) .∉ (idxs,)
-            args = Vector{Expression}[O.args[i].args for i in idxs]
-            push!(args, O.args[keep_idxs])
-            return Operation(O.op, vcat(args...))
-        end
+# Optional types of vars
+# Once converted to SymbolicUtils Variable, a Parameter needs to hide its metadata
+_vartype(x::Variable{<:Parameter{T}}) where {T} = T
+_vartype(x::Variable{T}) where {T} = T
+SymbolicUtils.symtype(x::Variable) = _vartype(x) # needed for a()
+SymbolicUtils.symtype(x::SymbolicUtils.Variable{<:Parameter{T}}) where {T} = T
 
-        # Collapse constants
-        idxs = findall(is_constant, O.args)
-        if length(idxs) > 1
-            other_idxs = eachindex(O.args) .∉ (idxs,)
-            new_const = Constant(mapreduce(get, O.op, O.args[idxs]))
-            args = push!(O.args[other_idxs], new_const)
+# returning Any causes SymbolicUtils to infer the type using `promote_symtype`
+# But we are OK with Number here for now I guess
+SymbolicUtils.symtype(x::Expression) = Number
 
-            length(args) == 1 && return first(args)
-            return Operation(O.op, args)
-        end
-    end
-
-    if O.op === (*)
-        # If any variable is `Constant(0)`, zero the whole thing
-        any(iszero, O.args) && return Constant(0)
-
-        # If any variable is `Constant(1)`, remove that `Constant(1)` unless
-        # they are all `Constant(1)`, in which case simplify to a single variable
-        if any(isone, O.args)
-            args = filter(!isone, O.args)
-
-            isempty(args)     && return Constant(1)
-            length(args) == 1 && return first(args)
-            return Operation(O.op, args)
-        end
-
-        return O
-    end
-
-    if O.op === (^) && length(O.args) == 2 && iszero(O.args[2])
-        return Constant(1)
-    end
 
-    if O.op === (^) && length(O.args) == 2 && isone(O.args[2])
-        return O.args[1]
-    end
-
-    if O.op === (+) && any(iszero, O.args)
-        # If there are Constant(0)s in a big `+` expression, get rid of them
-        args = filter(!iszero, O.args)
-
-        isempty(args)     && return Constant(0)
-        length(args) == 1 && return first(args)
-        return Operation(O.op, args)
-    end
+# SymbolicUtils -> ModelingToolkit
 
-    if (O.op === (-) || O.op === (+) || O.op === (*)) && all(is_constant, O.args) && !isempty(O.args)
-        v = O.args[1].value
-        for i in 2:length(O.args)
-            v = O.op(v, O.args[i].value)
-        end
-        return Constant(v)
-    end
-
-    (O.op, length(O.args)) === (identity, 1) && return O.args[1]
-
-    (O.op, length(O.args)) === (-, 1) && return Operation(*, Expression[-1, O.args[1]])
+function simplify_constants(expr)
+    SymbolicUtils.simplify(expr) |> to_mtk
+end
 
-    return O
+to_mtk(x) = x
+to_mtk(v::SymbolicUtils.Variable{T}) where {T} = Variable{T}(nameof(v))
+to_mtk(v::SymbolicUtils.Variable{FnType{X,Y}}) where {X,Y} = Variable{Y}(nameof(v))
+function to_mtk(expr::SymbolicUtils.Term)
+    Operation(to_mtk(SymbolicUtils.operation(expr)),
+              map(to_mtk, SymbolicUtils.arguments(expr)))
 end
-_simplify_constants(x, shorten_tree) = x
-_simplify_constants(x) = _simplify_constants(x, true)