Merge pull request #75 from JuliaDiffEq/hg/fix/simplify

Refactor simplification

Author: HarrisonGrodin

src/simplify.jl

@@ -1,92 +1,73 @@
-function simplify_constants(O::Operation,shorten_tree = true)
-    O_last = nothing
-    _O = O
-    while _O != O_last
-        O_last = _O
-        _O = _simplify_constants(_O,shorten_tree)
-        if typeof(_O) <: Operation
-            _O = Operation(_O.op,simplify_constants.(_O.args,shorten_tree))
+function simplify_constants(O::Operation, shorten_tree = true)
+    while true
+        O′ = _simplify_constants(O, shorten_tree)
+        if is_operation(O′)
+            O′ = Operation(O′.op, simplify_constants.(O′.args, shorten_tree))
         end
+        O == O′ && return O
+        O = O′
     end
-    _O
 end
 
-const TREE_SHRINK_OPS = [:*,:+]
+const AC_OPERATORS = (*, +)
 
-function _simplify_constants(O,shorten_tree = true)
+function _simplify_constants(O, shorten_tree = true)
     # Tree shrinking
-    if shorten_tree
-      for cur_op in TREE_SHRINK_OPS
-        if Symbol(O.op) == cur_op
-          # Shrink tree
-          if any(x->typeof(x)<:Operation && Symbol(x.op) == cur_op ,O.args)
-              idxs = findall(x->typeof(x)<:Operation && Symbol(x.op) == cur_op,O.args)
-              keep_idxs = 1:length(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...))
-          # Collapse constants
-          elseif length(findall(x->typeof(x)<:Variable && x.subtype == :Constant ,O.args)) > 1
-              idxs = findall(x->typeof(x)<:Variable && x.subtype == :Constant ,O.args)
-              other_idxs = 1:length(O.args) .∉ (idxs,)
-              if cur_op == :*
-                new_var = Constant(prod(x->x.value,O.args[idxs]))
-              elseif cur_op == :+
-                new_var = Constant(sum(x->x.value,O.args[idxs]))
-              end
-              new_args = O.args[other_idxs]
-              push!(new_args,new_var)
-              if length(new_args) > 1
-                return Operation(O.op,new_args)
-              else
-                return new_args[1]
-              end
-          end
+    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
+
+        # 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)
+
+            length(args) == 1 && return first(args)
+            return Operation(O.op, args)
         end
-      end
     end
 
-    if Symbol(O.op) == :*
+    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(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
-            return Constant(0)
-        elseif any(x->typeof(x)<:Variable && isequal(x,Constant(1)),O.args)
-            idxs = findall(x->typeof(x)<:Variable && isequal(x,Constant(1)),O.args)
-            _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
-            if isempty(_O.args)
-                return Constant(1)
-            elseif length(_O.args) == 1
-                return _O.args[1]
-            else
-                return _O
-            end
-        else
-            return O
-        end
-    elseif Symbol(O.op) == :+ && any(x->typeof(x)<:Variable &&
-           (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
-        # If there are Constant(0)s in a big `+` expression, get rid of them
-        idxs = findall(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
-        _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
-        if isempty(_O.args)
-            return Constant(0)
-        elseif length(_O.args) == 1
-            return _O.args[1]
-        else
-            return O
+        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
-    elseif O.op == identity
-        return O.args[1]
-    elseif Symbol(O.op) == :- && length(O.args) == 1
-        return Operation(*,Expression[-1,O.args[1]])
-    else
+
         return O
     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
+
+    (O.op, length(O.args)) === (identity, 1) && return O.args[1]
+
+    (O.op, length(O.args)) === (-, 1) && return Operation(*, Expression[-1, O.args[1]])
+
     return O
 end
-simplify_constants(x::Variable,y=false) = x
-_simplify_constants(x::Variable,y=false) = x
+simplify_constants(x::Variable, y=false) = x
+_simplify_constants(x::Variable, y=false) = x
 
 export simplify_constants

src/utils.jl

@@ -19,6 +19,12 @@ end
 
 toexpr(ex) = MacroTools.postwalk(x->x isa Union{Expression,Operation} ? Expr(x) : x, ex)
 
+is_constant(x::Variable) = x.subtype === :Constant
+is_constant(::Any) = false
+
+is_operation(::Operation) = true
+is_operation(::Any) = false
+
 has_dependent(t::Variable) = Base.Fix2(has_dependent, t)
 has_dependent(x::Variable, t::Variable) =
     t ∈ x.dependents || any(has_dependent(t), x.dependents)

src/variables.jl

@@ -63,6 +63,15 @@ end
 export Variable,Parameter,Constant,DependentVariable,IndependentVariable,JumpVariable,NoiseVariable,
        @Var, @DVar, @IVar, @Param, @Const
 
+
+Base.get(x::Variable) = x.value
+
+Base.iszero(::Expression) = false
+Base.iszero(c::Variable) = get(c) isa Number && iszero(get(c))
+Base.isone(::Expression) = false
+Base.isone(c::Variable) = get(c) isa Number && isone(get(c))
+
+
 # Variables use isequal for equality since == is an Operation
 function Base.:(==)(x::Variable,y::Variable)
     x.name == y.name && x.subtype == y.subtype && x.value == y.value &&