Create utility functions for benchmarks

Author: MilesCranmer

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,6 +1,7 @@
 using DynamicExpressions, BenchmarkTools, Random
 using DynamicExpressions: copy_node
-using SymbolicRegression: gen_random_tree_fixed_size, Options
+
+include("benchmark_utils.jl")
 
 const v_PACKAGE_VERSION = try
     VersionNumber(DynamicExpressions.PACKAGE_VERSION)
@@ -82,7 +83,7 @@ end
 
 function benchmark_utilities()
     suite = BenchmarkGroup()
-    options = Options(; binary_operators=[+, -, /, *], unary_operators=[cos, exp])
+    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")
@@ -99,20 +100,19 @@ function benchmark_utilities()
                         preserve_sharing=(k == "preserve_sharing"),
                     )
                 elseif func_k == "simplify_tree"
-                    tree -> simplify_tree(tree, options.operators)
+                    tree -> simplify_tree(tree, operators)
                 elseif func_k == "combine_operators"
-                    tree -> combine_operators(tree, options.operators)
+                    tree -> combine_operators(tree, operators)
                 end
 
                 #! format: off
                 s[k] = @benchmarkable(
-                    $(f)(tree),
-                    evals=300,
-                    samples=300,
+                    [$(f)(tree) for tree in trees],
                     seconds=10.0,
                     setup=(
+                        ntrees=100;
                         n=20;
-                        tree=gen_random_tree_fixed_size(n, $options, 5, Float32)
+                        trees=[gen_random_tree_fixed_size(n, $operators, 5, Float32) for _ in 1:ntrees]
                     )
                 )
                 #! format: on