Merge pull request #17 from SymbolicML/cleanup-iddict

Clean up IdDict use with macro

Project.toml

@@ -1,11 +1,12 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "0.6.1"
+version = "0.7.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -16,6 +17,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 LoopVectorization = "0.12"
+MacroTools = "0.4, 0.5"
 Reexport = "1"
 PrecompileTools = "1"
 SymbolicUtils = "0.19, ^1.0.5"

benchmark/Project.toml

@@ -1,4 +1,3 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

benchmark/benchmark_utils.jl

@@ -0,0 +1,83 @@
+import DynamicExpressions:
+    Node, copy_node, set_node!, count_nodes, has_constants, has_operators
+
+# This code is copied from SymbolicRegression.jl and modified
+
+# Return a random node from the tree
+function random_node(tree::Node{T})::Node{T} where {T}
+    if tree.degree == 0
+        return tree
+    end
+    b = 0
+    c = 0
+    if tree.degree >= 1
+        b = count_nodes(tree.l)
+    end
+    if tree.degree == 2
+        c = count_nodes(tree.r)
+    end
+
+    i = rand(1:(1 + b + c))
+    if i <= b
+        return random_node(tree.l)
+    elseif i == b + 1
+        return tree
+    end
+
+    return random_node(tree.r)
+end
+
+function make_random_leaf(nfeatures::Int, ::Type{T})::Node{T} where {T}
+    if rand() > 0.5
+        return Node(; val=randn(T))
+    else
+        return Node(T; feature=rand(1:nfeatures))
+    end
+end
+
+# Add a random unary/binary operation to the end of a tree
+function append_random_op(
+    tree::Node{T}, operators, nfeatures::Int; makeNewBinOp::Union{Bool,Nothing}=nothing
+)::Node{T} where {T}
+    nuna = length(operators.unaops)
+    nbin = length(operators.binops)
+
+    node = random_node(tree)
+    while node.degree != 0
+        node = random_node(tree)
+    end
+
+    if makeNewBinOp === nothing
+        choice = rand()
+        makeNewBinOp = choice < nbin / (nuna + nbin)
+    end
+
+    if makeNewBinOp
+        newnode = Node(
+            rand(1:nbin), make_random_leaf(nfeatures, T), make_random_leaf(nfeatures, T)
+        )
+    else
+        newnode = Node(rand(1:nuna), make_random_leaf(nfeatures, T))
+    end
+
+    set_node!(node, newnode)
+
+    return tree
+end
+
+function gen_random_tree_fixed_size(
+    node_count::Int, operators, nfeatures::Int, ::Type{T}
+)::Node{T} where {T}
+    tree = make_random_leaf(nfeatures, T)
+    cur_size = count_nodes(tree)
+    while cur_size < node_count
+        if cur_size == node_count - 1  # only unary operator allowed.
+            length(operators.unaops) == 0 && break # We will go over the requested amount, so we must break.
+            tree = append_random_op(tree, operators, nfeatures; makeNewBinOp=false)
+        else
+            tree = append_random_op(tree, operators, nfeatures)
+        end
+        cur_size = count_nodes(tree)
+    end
+    return tree
+end

benchmark/benchmarks.jl

@@ -1,71 +1,122 @@
 using DynamicExpressions, BenchmarkTools, Random
+using DynamicExpressions: copy_node
 
-const v_PACKAGE_VERSION = try
-    VersionNumber(PACKAGE_VERSION)
-catch
-    VersionNumber("v0.0.0")
-end
+include("benchmark_utils.jl")
 
 const SUITE = BenchmarkGroup()
 
-SUITE["OperatorEnum"] = BenchmarkGroup()
-
-operators = OperatorEnum(;
-    binary_operators=[+, -, /, *], unary_operators=[cos, exp], enable_autodiff=true
-)
-simple_tree = Node(
-    2,
-    Node(
-        1,
-        Node(3, Node(1, Node(; val=1.0f0), Node(; feature=2)), Node(2, Node(; val=-1.0f0))),
-        Node(1, Node(; feature=3), Node(; feature=4)),
-    ),
-    Node(
-        4,
-        Node(3, Node(1, Node(; val=1.0f0), Node(; feature=2)), Node(2, Node(; val=-1.0f0))),
-        Node(1, Node(; feature=3), Node(; feature=4)),
-    ),
-)
-for T in (ComplexF32, ComplexF64, Float32, Float64)
-    if !(T <: Real) && v_PACKAGE_VERSION < v"0.5.0" && v_PACKAGE_VERSION != v"0.0.0"
-        continue
-    end
-    evals = 10
-    samples = 1_000
-    n = 1_000
-
-    #! format: off
-    if !haskey(SUITE["OperatorEnum"], T)
-        SUITE["OperatorEnum"][T] = BenchmarkGroup()
-    end
-
-    for turbo in (false, true)
-        if turbo && !(T in (Float32, Float64))
+function benchmark_evaluation()
+    suite = BenchmarkGroup()
+    operators = OperatorEnum(;
+        binary_operators=[+, -, /, *], unary_operators=[cos, exp], enable_autodiff=true
+    )
+    for T in (ComplexF32, ComplexF64, Float32, Float64)
+        if !(T <: Real) && PACKAGE_VERSION < v"0.5.0" && PACKAGE_VERSION != v"0.0.0"
             continue
         end
-        extra_key = turbo ? "_turbo" : ""
-        SUITE["OperatorEnum"][T]["evaluation$(extra_key)"] = @benchmarkable(
-            eval_tree_array(tree, X, $operators; turbo=$turbo),
-            evals=evals,
-            samples=samples,
-            seconds=5.0,
-            setup=(
-                X=randn(MersenneTwister(0), $T, 5, $n);
-                tree=convert(Node{$T}, copy_node($simple_tree))
+        suite[T] = BenchmarkGroup()
+
+        n = 1_000
+
+        #! format: off
+        for turbo in (false, true)
+            if turbo && !(T in (Float32, Float64))
+                continue
+            end
+            extra_key = turbo ? "_turbo" : ""
+            eval_tree_array(
+                gen_random_tree_fixed_size(20, operators, 5, T),
+                randn(MersenneTwister(0), T, 5, n),
+                operators;
+                turbo=turbo
             )
-        )
-        if T <: Real
-            SUITE["OperatorEnum"][T]["derivative$(extra_key)"] = @benchmarkable(
-                eval_grad_tree_array(tree, X, $operators; variable=true, turbo=$turbo),
-                evals=evals,
-                samples=samples,
-                seconds=5.0,
+            suite[T]["evaluation$(extra_key)"] = @benchmarkable(
+                [eval_tree_array(tree, X, $operators; turbo=$turbo) for tree in trees],
                 setup=(
                     X=randn(MersenneTwister(0), $T, 5, $n);
-                    tree=convert(Node{$T}, copy_node($simple_tree))
+                    treesize=20;
+                    ntrees=100;
+                    trees=[gen_random_tree_fixed_size(treesize, $operators, 5, $T) for _ in 1:ntrees]
                 )
             )
+            if T <: Real
+                eval_grad_tree_array(
+                    gen_random_tree_fixed_size(20, operators, 5, T),
+                    randn(MersenneTwister(0), T, 5, n),
+                    operators;
+                    variable=true,
+                    turbo=turbo
+                )
+                suite[T]["derivative$(extra_key)"] = @benchmarkable(
+                    [eval_grad_tree_array(tree, X, $operators; variable=true, turbo=$turbo) for tree in trees],
+                    setup=(
+                        X=randn(MersenneTwister(0), $T, 5, $n);
+                        treesize=20;
+                        ntrees=100;
+                        trees=[gen_random_tree_fixed_size(treesize, $operators, 5, $T) for _ in 1:ntrees]
+                    )
+                )
+            end
+        end
+        #! format: on
+    end
+    return suite
+end
+
+# These macros make the benchmarks work on older versions:
+#! format: off
+@generated function _convert(::Type{N}, t; preserve_sharing) where {N<:Node}
+    PACKAGE_VERSION < v"0.7.0" && return :(convert(N, t))
+    return :(convert(N, t; preserve_sharing=preserve_sharing))
+end
+@generated function _copy_node(t; preserve_sharing)
+    PACKAGE_VERSION < v"0.7.0" && return :(copy_node(t; preserve_topology=preserve_sharing))
+    return :(copy_node(t; preserve_sharing=preserve_sharing))
+end
+#! format: on
+
+function benchmark_utilities()
+    suite = BenchmarkGroup()
+    operators = OperatorEnum(; binary_operators=[+, -, /, *], unary_operators=[cos, exp])
+    for func_k in ("copy", "convert", "simplify_tree", "combine_operators")
+        suite[func_k] = let s = BenchmarkGroup()
+            for k in ("break_sharing", "preserve_sharing")
+                k == "preserve_sharing" &&
+                    func_k in ("simplify_tree", "combine_operators") &&
+                    continue
+
+                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"),
+                    )
+                elseif func_k == "simplify_tree"
+                    tree -> simplify_tree(tree, operators)
+                elseif func_k == "combine_operators"
+                    tree -> combine_operators(tree, operators)
+                end
+
+                #! format: off
+                s[k] = @benchmarkable(
+                    [$(f)(tree) for tree in trees],
+                    seconds=10.0,
+                    setup=(
+                        ntrees=100;
+                        n=20;
+                        trees=[gen_random_tree_fixed_size(n, $operators, 5, Float32) for _ in 1:ntrees]
+                    )
+                )
+                #! format: on
+            end
+            s
         end
     end
-    #! format: on
+
+    return suite
 end
+
+SUITE["eval"] = benchmark_evaluation()
+SUITE["utils"] = benchmark_utilities()