refactor: update tree_gen_utils.jl

Author: MilesCranmer

src/base.jl

@@ -22,8 +22,9 @@ import Base:
     mapreduce,
     reduce,
     sum
-import Compat: @inline, Returns
-import ..UtilsModule: @memoize_on, @with_memoize, Undefined
+
+using Compat: @inline, Returns
+using ..UtilsModule: @memoize_on, @with_memoize, Undefined
 
 """
     tree_mapreduce(

test/tree_gen_utils.jl

@@ -1,61 +1,66 @@
-import DynamicExpressions:
-    Node, copy_node, set_node!, count_nodes, has_constants, has_operators
+using DynamicExpressions:
+    AbstractExpressionNode,
+    AbstractNode,
+    Node,
+    NodeSampler,
+    constructorof,
+    set_node!,
+    count_nodes
+using Random: AbstractRNG, default_rng
 
-# This code is copied from SymbolicRegression.jl and modified
+"""
+    random_node(tree::AbstractNode; filter::F=Returns(true))
 
-# 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 = count_nodes(tree.l)
-    c = if tree.degree == 2
-        count_nodes(tree.r)
-    else
-        0
-    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)
+Return a random node from the tree. You may optionally
+filter the nodes matching some condition before sampling.
+"""
+function random_node(
+    tree::AbstractNode, rng::AbstractRNG=default_rng(); filter::F=Returns(true)
+) where {F<:Function}
+    Base.depwarn(
+        "Instead of `random_node(tree, filter)`, use `rand(NodeSampler(; tree, filter))`",
+        :random_node,
+    )
+    return rand(rng, NodeSampler(; tree, filter))
 end
 
-function make_random_leaf(nfeatures::Integer, ::Type{T})::Node{T} where {T}
-    if rand() > 0.5
-        return Node(; val=randn(T))
+function make_random_leaf(
+    nfeatures::Int, ::Type{T}, ::Type{N}, rng::AbstractRNG=default_rng()
+) where {T,N<:AbstractExpressionNode}
+    if rand(rng, Bool)
+        return constructorof(N)(; val=randn(rng, T))
     else
-        return Node(T; feature=rand(1:nfeatures))
+        return constructorof(N)(T; feature=rand(rng, 1:nfeatures))
     end
 end
 
-# Add a random unary/binary operation to the end of a tree
+"""Add a random unary/binary operation to the end of a tree"""
 function append_random_op(
-    tree::Node{T}, operators, nfeatures::Integer; makeNewBinOp::Union{Bool,Nothing}=nothing
-)::Node{T} where {T}
+    tree::AbstractExpressionNode{T},
+    operators,
+    nfeatures::Int,
+    rng::AbstractRNG=default_rng();
+    makeNewBinOp::Union{Bool,Nothing}=nothing,
+) where {T}
+    node = rand(rng, NodeSampler(; tree, filter=t -> t.degree == 0))
     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()
+        choice = rand(rng)
         makeNewBinOp = choice < nbin / (nuna + nbin)
     end
 
     if makeNewBinOp
-        newnode = Node(
-            rand(1:nbin), make_random_leaf(nfeatures, T), make_random_leaf(nfeatures, T)
+        newnode = constructorof(typeof(tree))(
+            rand(rng, 1:nbin),
+            make_random_leaf(nfeatures, T, typeof(tree), rng),
+            make_random_leaf(nfeatures, T, typeof(tree), rng),
         )
     else
-        newnode = Node(rand(1:nuna), make_random_leaf(nfeatures, T))
+        newnode = constructorof(typeof(tree))(
+            rand(rng, 1:nuna), make_random_leaf(nfeatures, T, typeof(tree), rng)
+        )
     end
 
     set_node!(node, newnode)
@@ -64,16 +69,21 @@ function append_random_op(
 end
 
 function gen_random_tree_fixed_size(
-    node_count::Integer, operators, nfeatures::Integer, ::Type{T}
-)::Node{T} where {T}
-    tree = make_random_leaf(nfeatures, T)
+    node_count::Int,
+    operators,
+    nfeatures::Int,
+    ::Type{T},
+    node_type=Node,
+    rng::AbstractRNG=default_rng(),
+) where {T}
+    tree = make_random_leaf(nfeatures, T, node_type, rng)
     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)
+            tree = append_random_op(tree, operators, nfeatures, rng; makeNewBinOp=false)
         else
-            tree = append_random_op(tree, operators, nfeatures)
+            tree = append_random_op(tree, operators, nfeatures, rng)
         end
         cur_size = count_nodes(tree)
     end