Merge pull request #70 from SymbolicML/type-stability

feat!: simplify expression optimization routine

Author: MilesCranmer

ext/DynamicExpressionsOptimExt.jl

@@ -1,7 +1,6 @@
 module DynamicExpressionsOptimExt
 
-using DynamicExpressions:
-    AbstractExpressionNode, eval_tree_array, get_constant_refs, set_constant_refs!
+using DynamicExpressions: AbstractExpressionNode, filter_map, eval_tree_array
 using Compat: @inline
 
 import Optim: Optim, OptimizationResults, NLSolversBase
@@ -39,7 +38,9 @@ function wrap_func(
     function wrapped_f(args::Vararg{Any,M}) where {M}
         first_args = args[1:(end - 1)]
         x = last(args)
-        set_constant_refs!(constant_refs, x)
+        @inbounds for i in eachindex(constant_refs, x)
+            constant_refs[i][].val = x[i]
+        end
         return @inline(f(first_args..., tree))
     end
     return wrapped_f
@@ -87,8 +88,10 @@ function Optim.optimize(
     if make_copy
         tree = copy(tree)
     end
-    constant_refs = get_constant_refs(tree)
-    x0 = map(t -> t.x, constant_refs)
+    constant_refs = filter_map(
+        t -> t.degree == 0 && t.constant, t -> Ref(t), tree, Ref{typeof(tree)}
+    )
+    x0 = T[copy(t[].val) for t in constant_refs]
     if !isnothing(h!)
         throw(
             ArgumentError(
@@ -108,7 +111,10 @@ function Optim.optimize(
             kwargs...,
         )
     end
-    set_constant_refs!(constant_refs, Optim.minimizer(base_res))
+    minimizer = Optim.minimizer(base_res)
+    @inbounds for i in eachindex(constant_refs, minimizer)
+        constant_refs[i][].val = minimizer[i]
+    end
     return ExpressionOptimizationResults(base_res, tree)
 end
 

src/DynamicExpressions.jl

@@ -35,9 +35,7 @@ import .EquationModule: constructorof, preserve_sharing
     has_operators,
     has_constants,
     get_constants,
-    set_constants!,
-    get_constant_refs,
-    set_constant_refs!
+    set_constants!
 @reexport import .StringsModule: string_tree, print_tree
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:

src/EquationUtils.jl

@@ -98,60 +98,6 @@ function set_constants!(
     return nothing
 end
 
-"""
-    NodeConstantRef{T,N<:AbstractExpressionNode{T}}
-
-A reference to a constant in an expression tree. Use `.x` to access
-the value of the constant for setting or getting.
-"""
-struct NodeConstantRef{T,N<:AbstractExpressionNode{T}}
-    _node::Ref{N}
-
-    function NodeConstantRef(node::_N) where {_T,_N<:AbstractExpressionNode{_T}}
-        return new{_T,_N}(Ref(node))
-    end
-end
-function Base.getproperty(cr::NodeConstantRef{T}, s::Symbol) where {T}
-    s != :x && error("Only :x is a valid property for NodeConstantRef")
-
-    return getfield(cr, :_node).x.val
-end
-function Base.setproperty!(cr::NodeConstantRef{T}, s::Symbol, v) where {T}
-    s != :x && error("Only :x is a valid property for NodeConstantRef")
-
-    return getfield(cr, :_node).x.val = v
-end
-Base.propertynames(::NodeConstantRef) = (:x,)
-
-"""
-    get_constant_refs(tree::AbstractExpressionNode)
-
-Get references to all constants in a tree, in depth-first order. Using the output of this lets
-you quickly modify the constants in the tree in-place.
-"""
-function get_constant_refs(tree::AbstractExpressionNode)
-    return filter_map(
-        is_node_constant,
-        t -> NodeConstantRef(t),
-        tree,
-        NodeConstantRef{eltype(tree),typeof(tree)},
-    )
-end
-
-"""
-    set_constant_refs!(crs::AbstractArray{C}, xs::AbstractArray{T}) where {T,C<:NodeConstantRef{T}}
-
-Set the constants in a tree to the values in a vector.
-"""
-@inline function set_constant_refs!(
-    constant_refs::AbstractArray{C}, xs::AbstractArray{T}
-) where {T,C<:NodeConstantRef{T}}
-    for (cr, x) in zip(constant_refs, xs)
-        cr.x = x
-    end
-    return nothing
-end
-
 ## Assign index to nodes of a tree
 # This will mirror a Node struct, rather
 # than adding a new attribute to Node.

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