Merge pull request #27 from SymbolicML/tree-map

Define functions in Base to treat `Node` as collection

Author: MilesCranmer

Project.toml

@@ -1,25 +1,27 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "0.7.0"
+version = "0.8.0"
 
 [deps]
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+Compat = "3.37, 4"
 LoopVectorization = "0.12"
 MacroTools = "0.4, 0.5"
-Reexport = "1"
 PrecompileTools = "1"
+Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5"
 Zygote = "0.6"
 julia = "1.6"

benchmark/benchmarks.jl

@@ -1,5 +1,5 @@
 using DynamicExpressions, BenchmarkTools, Random
-using DynamicExpressions: copy_node
+using DynamicExpressions.EquationUtilsModule: is_constant
 
 include("benchmark_utils.jl")
 
@@ -73,30 +73,54 @@ end
     PACKAGE_VERSION < v"0.7.0" && return :(copy_node(t; preserve_topology=preserve_sharing))
     return :(copy_node(t; preserve_sharing=preserve_sharing))
 end
+@generated function get_set_constants!(tree)
+    !(@isdefined set_constants!) && return :(set_constants(tree, get_constants(tree)))
+    return :(set_constants!(tree, get_constants(tree)))
+end
 #! format: on
 
+f_tree_op(f::F, tree, operators) where {F} = f(tree, operators)
+f_tree_op(f::F, tree) where {F} = f(tree)
+
 function benchmark_utilities()
     suite = BenchmarkGroup()
+
+    all_funcs = (
+        :copy,
+        :convert,
+        :simplify_tree,
+        :combine_operators,
+        :count_nodes,
+        :count_depth,
+        :count_constants,
+        :has_constants,
+        :has_operators,
+        :is_constant,
+        :get_set_constants!,
+        :index_constants,
+    )
+
     operators = OperatorEnum(; binary_operators=[+, -, /, *], unary_operators=[cos, exp])
-    for func_k in ("copy", "convert", "simplify_tree", "combine_operators")
+
+    for func_k in all_funcs
         suite[func_k] = let s = BenchmarkGroup()
-            for k in ("break_sharing", "preserve_sharing")
-                k == "preserve_sharing" &&
-                    func_k in ("simplify_tree", "combine_operators") &&
-                    continue
+            for k in (:break_sharing, :preserve_sharing)
+                k == :preserve_sharing && !(func_k in (:copy, :convert)) && continue
 
-                f = if func_k == "copy"
-                    tree -> _copy_node(tree; preserve_sharing=(k == "preserve_sharing"))
-                elseif func_k == "convert"
+                f = if func_k == :copy
+                    tree -> _copy_node(tree; preserve_sharing=(k == :preserve_sharing))
+                elseif func_k == :convert
                     tree -> _convert(
                         Node{Float64},
                         tree;
-                        preserve_sharing=(k == "preserve_sharing"),
+                        preserve_sharing=(k == :preserve_sharing),
                     )
-                elseif func_k == "simplify_tree"
-                    tree -> simplify_tree(tree, operators)
-                elseif func_k == "combine_operators"
-                    tree -> combine_operators(tree, operators)
+                elseif func_k in (:simplify_tree, :combine_operators)
+                    g = getfield(@__MODULE__, func_k)
+                    tree -> f_tree_op(g, tree, operators)
+                else
+                    g = getfield(@__MODULE__, func_k)
+                    tree -> f_tree_op(g, tree)
                 end
 
                 #! format: off

src/DynamicExpressions.jl

@@ -12,18 +12,18 @@ include("SimplifyEquation.jl")
 include("OperatorEnumConstruction.jl")
 
 using Reexport
-@reexport import .EquationModule: Node, string_tree, print_tree, copy_node, set_node!
+@reexport import .EquationModule:
+    Node, string_tree, print_tree, copy_node, set_node!, tree_mapreduce, filter_map
 @reexport import .EquationUtilsModule:
     count_nodes,
-    count_nodes_with_stack,
     count_constants,
     count_depth,
     NodeIndex,
     index_constants,
     has_operators,
     has_constants,
     get_constants,
-    set_constants
+    set_constants!
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators
@@ -34,6 +34,8 @@ using Reexport
 @reexport import .SimplifyEquationModule: combine_operators, simplify_tree
 @reexport import .EvaluationHelpersModule
 
+include("deprecated.jl")
+
 import TOML: parsefile
 
 const PACKAGE_VERSION = let

src/Equation.jl

@@ -1,7 +1,7 @@
 module EquationModule
 
 import ..OperatorEnumModule: AbstractOperatorEnum
-import ..UtilsModule: @generate_idmap, @use_idmap
+import ..UtilsModule: @memoize_on, @with_memoize
 
 const DEFAULT_NODE_TYPE = Float32
 
@@ -62,51 +62,7 @@ mutable struct Node{T}
 end
 ################################################################################
 
-"""
-    convert(::Type{Node{T1}}, n::Node{T2}) where {T1,T2}
-
-Convert a `Node{T2}` to a `Node{T1}`.
-This will recursively convert all children nodes to `Node{T1}`,
-using `convert(T1, tree.val)` at constant nodes.
-
-# Arguments
-- `::Type{Node{T1}}`: Type to convert to.
-- `tree::Node{T2}`: Node to convert.
-"""
-function Base.convert(
-    ::Type{Node{T1}}, tree::Node{T2}; preserve_sharing::Bool=false
-) where {T1,T2}
-    if T1 == T2
-        return tree
-    end
-    if preserve_sharing
-        @use_idmap(_convert(Node{T1}, tree), IdDict{Node{T2},Node{T1}}())
-    else
-        _convert(Node{T1}, tree)
-    end
-end
-
-@generate_idmap tree function _convert(::Type{Node{T1}}, tree::Node{T2}) where {T1,T2}
-    if tree.degree == 0
-        if tree.constant
-            val = tree.val::T2
-            if !(T2 <: T1)
-                # e.g., we don't want to convert Float32 to Union{Float32,Vector{Float32}}!
-                val = convert(T1, val)
-            end
-            Node(T1, 0, tree.constant, val)
-        else
-            Node(T1, 0, tree.constant, nothing, tree.feature)
-        end
-    elseif tree.degree == 1
-        l = _convert(Node{T1}, tree.l)
-        Node(1, tree.constant, nothing, tree.feature, tree.op, l)
-    else
-        l = _convert(Node{T1}, tree.l)
-        r = _convert(Node{T1}, tree.r)
-        Node(2, tree.constant, nothing, tree.feature, tree.op, l, r)
-    end
-end
+include("base.jl")
 
 """
     Node([::Type{T}]; val=nothing, feature::Int=nothing) where {T}
@@ -224,45 +180,6 @@ function set_node!(tree::Node{T}, new_tree::Node{T}) where {T}
     return nothing
 end
 
-"""
-    copy_node(tree::Node; preserve_sharing::Bool=false)
-
-Copy a node, recursively copying all children nodes.
-This is more efficient than the built-in copy.
-With `preserve_sharing=true`, this will also
-preserve linkage between a node and
-multiple parents, whereas without, this would create
-duplicate child node copies.
-
-id_map is a map from `objectid(tree)` to `copy(tree)`.
-We check against the map before making a new copy; otherwise
-we can simply reference the existing copy.
-[Thanks to Ted Hopp.](https://stackoverflow.com/questions/49285475/how-to-copy-a-full-non-binary-tree-including-loops)
-
-Note that this will *not* preserve loops in graphs.
-"""
-function copy_node(tree::Node{T}; preserve_sharing::Bool=false)::Node{T} where {T}
-    if preserve_sharing
-        @use_idmap(_copy_node(tree), IdDict{Node{T},Node{T}}())
-    else
-        _copy_node(tree)
-    end
-end
-
-@generate_idmap tree function _copy_node(tree::Node{T})::Node{T} where {T}
-    if tree.degree == 0
-        if tree.constant
-            Node(; val=copy(tree.val::T))
-        else
-            Node(T; feature=copy(tree.feature))
-        end
-    elseif tree.degree == 1
-        Node(copy(tree.op), _copy_node(tree.l))
-    else
-        Node(copy(tree.op), _copy_node(tree.l), _copy_node(tree.r))
-    end
-end
-
 const OP_NAMES = Dict(
     "safe_log" => "log",
     "safe_log2" => "log2",
@@ -363,51 +280,4 @@ function print_tree(
     return println(string_tree(tree, operators; varMap=varMap))
 end
 
-function Base.hash(tree::Node{T})::UInt where {T}
-    if tree.degree == 0
-        if tree.constant
-            # tree.val used.
-            return hash((0, tree.val::T))
-        else
-            # tree.feature used.
-            return hash((1, tree.feature))
-        end
-    elseif tree.degree == 1
-        return hash((1, tree.op, hash(tree.l)))
-    else
-        return hash((2, tree.op, hash(tree.l), hash(tree.r)))
-    end
-end
-
-function is_equal(a::Node{T}, b::Node{T})::Bool where {T}
-    if a.degree == 0
-        b.degree != 0 && return false
-        if a.constant
-            !(b.constant) && return false
-            return a.val::T == b.val::T
-        else
-            b.constant && return false
-            return a.feature == b.feature
-        end
-    elseif a.degree == 1
-        b.degree != 1 && return false
-        a.op != b.op && return false
-        return is_equal(a.l, b.l)
-    else
-        b.degree != 2 && return false
-        a.op != b.op && return false
-        return is_equal(a.l, b.l) && is_equal(a.r, b.r)
-    end
-end
-
-function Base.:(==)(a::Node{T}, b::Node{T})::Bool where {T}
-    return is_equal(a, b)
-end
-
-function Base.:(==)(a::Node{T1}, b::Node{T2})::Bool where {T1,T2}
-    T = promote_type(T1, T2)
-    # TODO: Should also have preserve_sharing check...
-    return is_equal(convert(Node{T}, a), convert(Node{T}, b))
-end
-
 end