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.