feat: additional simplification types and generality

Author: MilesCranmer

src/Simplify.jl

@@ -4,105 +4,122 @@ import ..NodeModule: AbstractExpressionNode, constructorof, Node, copy_node, set
 import ..NodeUtilsModule: tree_mapreduce, is_node_constant
 import ..OperatorEnumModule: AbstractOperatorEnum
 import ..ValueInterfaceModule: is_valid
+import ..EvaluateModule: get_nbin
 
 _una_op_kernel(f::F, l::T) where {F,T} = f(l)
 _bin_op_kernel(f::F, l::T, r::T) where {F,T} = f(l, r)
 
+# Operator traits
 is_commutative(::typeof(*)) = true
 is_commutative(::typeof(+)) = true
 is_commutative(_) = false
 
-is_subtraction(::typeof(-)) = true
-is_subtraction(_) = false
+# Zero-related traits
+has_right_identity_zero(::typeof(+)) = true
+has_right_identity_zero(::typeof(-)) = true
+has_right_identity_zero(_) = false
+
+has_left_identity_zero(::typeof(+)) = true
+has_left_identity_zero(_) = false
+
+absorbed_by_zero(::typeof(*)) = true
+absorbed_by_zero(_) = false
+
+# One-related traits
+has_identity_one(::typeof(*)) = true
+has_identity_one(_) = false
+
+# Self-operation traits
+simplifies_given_equal_operands(::typeof(/)) = true
+simplifies_given_equal_operands(_) = false
 
 combine_operators(tree::AbstractExpressionNode, ::AbstractOperatorEnum) = tree
-# This is only defined for `Node` as it is not possible for, e.g.,
-# `GraphNode`.
+
 function combine_operators(tree::Node{T}, operators::AbstractOperatorEnum) where {T}
-    # NOTE: (const (+*-) const) already accounted for. Call simplify_tree! before.
-    # ((const + var) + const) => (const + var)
-    # ((const * var) * const) => (const * var)
-    # ((const - var) - const) => (const - var)
-    # (want to add anything commutative!)
-    # TODO - need to combine plus/sub if they are both there.
-    if tree.degree == 0
-        return tree
-    elseif tree.degree == 1
-        tree.l = combine_operators(tree.l, operators)
-    elseif tree.degree == 2
-        tree.l = combine_operators(tree.l, operators)
-        tree.r = combine_operators(tree.r, operators)
+    deg = tree.degree
+    deg == 0 && return tree
+    tree.l = combine_operators(tree.l, operators)
+    deg == 1 && return tree
+    tree.r = combine_operators(tree.r, operators)
+    return dispatch_deg2_simplify(tree, operators)
+end
+@generated function dispatch_deg2_simplify(
+    tree::Node{T}, operators::AbstractOperatorEnum
+) where {T}
+    nbin = get_nbin(operators)
+    quote
+        op_idx = tree.op
+        return Base.Cartesian.@nif(
+            $nbin, i -> i == op_idx, i -> _combine_operators_on(operators.binops[i], tree)
+        )
     end
+end
 
-    top_level_constant =
-        tree.degree == 2 && (is_node_constant(tree.l) || is_node_constant(tree.r))
-    if tree.degree == 2 && is_commutative(operators.binops[tree.op]) && top_level_constant
-        # TODO: Does this break SymbolicRegression.jl due to the different names of operators?
-        op = tree.op
-        # Put the constant in r. Need to assume var in left for simplification assumption.
-        if is_node_constant(tree.l)
-            tmp = tree.r
-            tree.r = tree.l
-            tree.l = tmp
+function _combine_operators_on(f::F, tree::Node{T}) where {F,T}
+    # NOTE: This assumes tree.degree == 2 and tree.op corresponds to f
+    # Handle basic simplifications first
+    if is_node_constant(tree.r)
+        rval = tree.r.val
+        if rval == zero(T)
+            # Operations where right zero is identity (x + 0 -> x, x - 0 -> x)
+            if has_right_identity_zero(f)
+                return tree.l
+            end
+            # Operations that are absorbed by zero (x * 0 -> 0)
+            if absorbed_by_zero(f)
+                return tree.r
+            end
+        elseif rval == one(T)
+            # x * 1 -> x
+            if has_identity_one(f)
+                return tree.l
+            end
         end
-        topconstant = tree.r.val
-        # Simplify down first
-        below = tree.l
-        if below.degree == 2 && below.op == op
-            if is_node_constant(below.l)
-                tree = below
-                tree.l.val = _bin_op_kernel(operators.binops[op], tree.l.val, topconstant)
-            elseif is_node_constant(below.r)
-                tree = below
-                tree.r.val = _bin_op_kernel(operators.binops[op], tree.r.val, topconstant)
+    end
+
+    if is_node_constant(tree.l)
+        lval = tree.l.val
+        if lval == zero(T)
+            # Operations where left zero is identity (0 + x -> x)
+            if has_left_identity_zero(f)
+                return tree.r
             end
+            # Operations that are absorbed by zero (0 * x -> 0)
+            if absorbed_by_zero(f)
+                return tree.l
+            end
+        elseif lval == one(T) && has_identity_one(f)
+            # 1 * x -> x
+            return tree.r
         end
     end
 
-    if tree.degree == 2 && is_subtraction(operators.binops[tree.op]) && top_level_constant
+    # x/x -> 1, or other self-simplifying operations
+    if simplifies_given_equal_operands(f) && tree.l == tree.r
+        return constructorof(typeof(tree))(T; val=one(T))
+    end
 
-        # Currently just simplifies subtraction. (can't assume both plus and sub are operators)
-        # Not commutative, so use different op.
+    # Handle commutative operations with constants
+    if is_commutative(f) && (is_node_constant(tree.l) || is_node_constant(tree.r))
+        # Put constant on right for consistent handling
         if is_node_constant(tree.l)
-            if tree.r.degree == 2 && tree.op == tree.r.op
-                if is_node_constant(tree.r.l)
-                    #(const - (const - var)) => (var - const)
-                    l = tree.l
-                    r = tree.r
-                    simplified_const = (r.l.val - l.val) #neg(sub(l.val, r.l.val))
-                    tree.l = tree.r.r
-                    tree.r = l
-                    tree.r.val = simplified_const
-                elseif is_node_constant(tree.r.r)
-                    #(const - (var - const)) => (const - var)
-                    l = tree.l
-                    r = tree.r
-                    simplified_const = l.val + r.r.val #plus(l.val, r.r.val)
-                    tree.r = tree.r.l
-                    tree.l.val = simplified_const
-                end
-            end
-        else #tree.r is a constant
-            if tree.l.degree == 2 && tree.op == tree.l.op
-                if is_node_constant(tree.l.l)
-                    #((const - var) - const) => (const - var)
-                    l = tree.l
-                    r = tree.r
-                    simplified_const = l.l.val - r.val#sub(l.l.val, r.val)
-                    tree.r = tree.l.r
-                    tree.l = r
-                    tree.l.val = simplified_const
-                elseif is_node_constant(tree.l.r)
-                    #((var - const) - const) => (var - const)
-                    l = tree.l
-                    r = tree.r
-                    simplified_const = r.val + l.r.val #plus(r.val, l.r.val)
-                    tree.l = tree.l.l
-                    tree.r.val = simplified_const
-                end
+            tree.l, tree.r = tree.r, tree.l
+        end
+
+        # Now we know tree.r is constant
+        below = tree.l
+        if below.degree == 2 && below.op == tree.op
+            # Combine nested operations with constants: ((a * x) * b) -> ((a*b) * x)
+            if is_node_constant(below.l)
+                below.l.val = _bin_op_kernel(f, below.l.val, tree.r.val)
+                return below
+            elseif is_node_constant(below.r)
+                below.r.val = _bin_op_kernel(f, below.r.val, tree.r.val)
+                return below
             end
         end
     end
+
     return tree
 end
 
@@ -121,7 +138,6 @@ function combine_children!(operators, p::N, c::N...) where {T,N<:AbstractExpress
     return p
 end
 
-# Simplify tree
 function simplify_tree!(tree::AbstractExpressionNode, operators::AbstractOperatorEnum)
     return tree_mapreduce(
         identity, (p, c...) -> combine_children!(operators, p, c...), tree, typeof(tree);

test/test_parametric_expression.jl

@@ -161,7 +161,8 @@ end
 
     (init_out, true_out) = map(
         p -> [
-            (X[1, i] * p[2, classes[i]] + X[2, i] + p[1, classes[i]]) for i in 1:size(X, 2)
+            (X[1, i] * p[2, classes[i]] + X[2, i] + p[1, classes[i]]) for
+            i in eachindex(axes(X, 2))
         ],
         (init_parameters, true_parameters),
     )

test/test_simplification.jl

@@ -51,13 +51,14 @@ end
         ((x2 + x2) * ((-0.5982493 / pow_abs2(x1, x2)) / -0.54734415)) + (
             sin(
                 custom_cos(
-                    sin(1.2926733 - 1.6606787) /
-                    sin(((0.14577048 * x1) + ((0.111149654 + x1) - -0.8298334)) - -1.2071426),
+                    sin(1.2926733 - 1.6606787) / sin(
+                        ((0.14577048 * x1) + ((0.111149654 + x1) - -0.8298334)) - -1.2071426
+                    ),
                 ) * (custom_cos(x3 - 2.3201916) + ((x1 - (x1 * x2)) / x2)),
             ) / (0.14854191 - ((custom_cos(x2) * -1.6047639) - 0.023943262))
         )
     )
-    
+
     eqn = convert(Symbolic, tree, operators; index_functions=true)
     tree_copy = convert(Node, eqn, operators)
     tree_copy2 = convert(Node, simplify(eqn), operators)
@@ -89,28 +90,72 @@ end
     @test repr(simplify_tree!(tree, operators)) ≈ "cos(NaN)"
 
     # Nested constant folding
-    tree = Node(1, Node(1, Node(; val=0.1), Node(; val=0.2)) + Node(; val=0.2)) + Node(; val=2.0)
+    tree =
+        Node(1, Node(1, Node(; val=0.1), Node(; val=0.2)) + Node(; val=0.2)) +
+        Node(; val=2.0)
     @test repr(tree) ≈ "(cos((0.1 + 0.2) + 0.2) + 2.0)"
     @test repr(combine_operators(tree, operators)) ≈ "(cos(0.4 + 0.1) + 2.0)"
 end
 
-# (const - (const - var)) => (var - const)
-tree = Node(2, Node(; val=0.5), Node(; val=0.2) - x1)
-@test repr(tree) ≈ "(0.5 - (0.2 - x1))"
-@test repr(combine_operators(tree, operators)) ≈ "(x1 - -0.3)"
-
-# ((const - var) - const) => (const - var)
-tree = Node(2, Node(; val=0.5) - x1, Node(; val=0.2))
-@test repr(tree) ≈ "((0.5 - x1) - 0.2)"
-@test repr(combine_operators(tree, operators)) ≈ "(0.3 - x1)"
-
-# (const - (var - const)) => (const - var)
-tree = Node(2, Node(; val=0.5), x1 - Node(; val=0.2))
-@test repr(tree) ≈ "(0.5 - (x1 - 0.2))"
-@test repr(combine_operators(tree, operators)) ≈ "(0.7 - x1)"
-
-# ((var - const) - const) => (var - const)
-tree = ((x1 - 0.2) - 0.6)
-@test repr(tree) ≈ "((x1 - 0.2) - 0.6)"
-@test repr(combine_operators(tree, operators)) ≈ "(x1 - 0.8)"
-###############################################################################
+@testitem "Basic operator simplifications" begin
+    using DynamicExpressions, Test
+    import DynamicExpressions.SimplifyModule: combine_operators
+
+    operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(cos, sin))
+    x = Node(; feature=1)
+    zero_node = Node(; val=0.0)
+    one_node = Node(; val=1.0)
+    two_node = Node(; val=2.0)
+    three_node = Node(; val=3.0)
+
+    # multiplication by 0
+    tree = zero_node * x
+    @test combine_operators(tree, operators) == zero_node
+    tree = x * zero_node
+    @test combine_operators(tree, operators) == zero_node
+
+    # multiplication by 1
+    tree = one_node * x
+    @test combine_operators(tree, operators) == x
+    tree = x * one_node
+    @test combine_operators(tree, operators) == x
+
+    # addition by 0
+    tree = zero_node + x
+    @test combine_operators(tree, operators) == x
+    tree = x + zero_node
+    @test combine_operators(tree, operators) == x
+
+    # division by self -> 1
+    tree = x / x
+    @test combine_operators(tree, operators).val == 1.0
+
+    # nested multiplication by constants
+    tree1 = (two_node * x) * three_node
+    tree2 = Node(; val=6.0) * x
+    @test combine_operators(tree1, operators) == combine_operators(tree2, operators)
+end
+
+@testitem "Constant combination" begin
+    using DynamicExpressions, Test
+    import DynamicExpressions.SimplifyModule: combine_operators
+
+    operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(cos, sin))
+    x1 = Node(; feature=1)
+
+    # Test commutative constant combination
+    tree = Node(; val=0.5) + (Node(; val=0.2) + x1)
+    @test combine_operators(tree, operators) == x1 + Node(; val=0.7)
+
+    # Test nested multiplication by constants
+    tree = (Node(; val=2.0) * x1) * Node(; val=3.0)
+    @test combine_operators(tree, operators) == x1 * Node(; val=6.0)
+
+    # Test nested addition by constants
+    tree = (Node(; val=2.0) + x1) + Node(; val=3.0)
+    @test combine_operators(tree, operators) == x1 + Node(; val=5.0)
+
+    # Test mixed operations don't combine incorrectly
+    tree = (Node(; val=2.0) * x1) + Node(; val=3.0)
+    @test combine_operators(tree, operators) == tree
+end