Proper Optim interface of results struct

Author: MilesCranmer

ext/DynamicExpressionsOptimExt.jl

@@ -3,16 +3,51 @@ module DynamicExpressionsOptimExt
 using DynamicExpressions: AbstractExpressionNode, eval_tree_array
 using Compat: @inline
 
-import Optim: optimize
+import Optim: Optim, OptimizationResults
+
+#! format: off
+"""
+    ExpressionOptimizationResults{R,N<:AbstractExpressionNode}
+
+Optimization results for an expression, which wraps the base optimization results
+on a vector of constants.
+"""
+struct ExpressionOptimizationResults{R<:OptimizationResults,N<:AbstractExpressionNode} <: OptimizationResults
+    _results::R # The raw results from Optim.
+    tree::N # The final expression tree
+end
+#! format: on
+function Base.getproperty(r::ExpressionOptimizationResults, s::Symbol)
+    if s == :tree || s == :minimizer
+        return getfield(r, :tree)
+    else
+        return getproperty(getfield(r, :_results), s)
+    end
+end
+function Base.propertynames(r::ExpressionOptimizationResults)
+    return (:tree, propertynames(getfield(r, :_results))...)
+end
+function base_results(r::ExpressionOptimizationResults)
+    return getfield(r, :_results)
+end
+function Optim.minimizer(r::ExpressionOptimizationResults)
+    return r.tree
+end
+
+function set_constant_nodes!(
+    constant_nodes::AbstractArray{N}, x
+) where {T,N<:AbstractExpressionNode{T}}
+    for (ci, xi) in zip(constant_nodes, x)
+        ci.val::T = xi::T
+    end
+end
 
 """Wrap f with insertion of values of the constant nodes."""
 function get_wrapped_f(
     f::F, tree::N, constant_nodes::AbstractArray{N}
 ) where {F,T,N<:AbstractExpressionNode{T}}
     function wrapped_f(x)
-        for (ci, xi) in zip(constant_nodes, x)
-            ci.val::T = xi::T
-        end
+        set_constant_nodes!(constant_nodes, x)
         return @inline(f(tree))
     end
     return wrapped_f
@@ -23,31 +58,52 @@ function get_wrapped_gh!(
     gh!::GH, tree::N, constant_nodes::AbstractArray{N}
 ) where {GH,T,N<:AbstractExpressionNode{T}}
     function wrapped_gh!(G, x)
-        for (ci, xi) in zip(constant_nodes, x)
-            ci.val::T = xi::T
-        end
+        set_constant_nodes!(constant_nodes, x)
         @inline(gh!(G, tree))
         return nothing
     end
     return wrapped_gh!
 end
 
-function optimize(f::F, g!::G, h!::H, tree::AbstractExpressionNode{T}, args...; kwargs...) where {F,G,H,T}
+"""
+    optimize(f, [g!, [h!,]] tree, args...; kwargs...)
+
+Optimize an expression tree with respect to the constants in the tree.
+Returns an `ExpressionOptimizationResults` object, which wraps the base
+optimization results on a vector of constants. You may use `res.minimizer`
+to view the optimized expression tree.
+"""
+function Optim.optimize(
+    f::F, tree::AbstractExpressionNode, args...; kwargs...
+) where {F<:Function}
+    return Optim.optimize(f, nothing, tree, args...; kwargs...)
+end
+function Optim.optimize(
+    f::F, g!::G, tree::AbstractExpressionNode, args...; kwargs...
+) where {F,G<:Union{Function,Nothing}}
+    return Optim.optimize(f, g!, nothing, tree, args...; kwargs...)
+end
+function Optim.optimize(
+    f::F, g!::G, h!::H, tree::AbstractExpressionNode{T}, args...; make_copy=true, kwargs...
+) where {F,G<:Union{Function,Nothing},H<:Union{Function,Nothing},T}
+    if make_copy
+        tree = copy(tree)
+    end
     constant_nodes = filter(t -> t.degree == 0 && t.constant, tree)
     x0 = T[t.val::T for t in constant_nodes]
-    if g! === nothing
+    base_res = if g! === nothing
         @assert h! === nothing
-        return optimize(get_wrapped_f(f, tree, constant_nodes), x0, args...; kwargs...)
+        Optim.optimize(get_wrapped_f(f, tree, constant_nodes), x0, args...; kwargs...)
     elseif h! === nothing
-        return optimize(
+        Optim.optimize(
             get_wrapped_f(f, tree, constant_nodes),
             get_wrapped_gh!(g!, tree, constant_nodes),
             x0,
             args...;
             kwargs...,
         )
     else
-        return optimize(
+        Optim.optimize(
             get_wrapped_f(f, tree, constant_nodes),
             get_wrapped_gh!(g!, tree, constant_nodes),
             get_wrapped_gh!(h!, tree, constant_nodes),
@@ -56,12 +112,8 @@ function optimize(f::F, g!::G, h!::H, tree::AbstractExpressionNode{T}, args...;
             kwargs...,
         )
     end
-end
-function optimize(f::F, g!::G, tree::AbstractExpressionNode, args...; kwargs...) where {F,G}
-    return optimize(f, g!, nothing, tree, args...; kwargs...)
-end
-function optimize(f::F, tree::AbstractExpressionNode, args...; kwargs...) where {F}
-    return optimize(f, nothing, tree, args...; kwargs...)
+    set_constant_nodes!(constant_nodes, Optim.minimizer(base_res))
+    return ExpressionOptimizationResults(base_res, tree)
 end
 
 end

test/test_optim.jl

@@ -0,0 +1,16 @@
+using DynamicExpressions, Optim, Zygote
+using Random: Xoshiro
+
+operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(sin, cos))
+x1, x2, x3 = (i -> Node(Float64; feature=i)).(1:3);
+
+X = randn(Xoshiro(0), Float64, 3, 100);
+y = @. cos(X[1, :] * 2.1 - 0.9) + X[3, :] * -0.9
+
+original_tree = cos(x1 * 0.8 - 0.0) + 5.2 * x3
+tree = copy(original_tree)
+
+res = optimize(t -> sum(abs2, t(X, operators) .- y), tree)
+
+# Should be unchanged by default
+@test tree == original_tree

test/unittest.jl

@@ -4,6 +4,10 @@ VERSION >= v"1.9" && @safetestset "Test Aqua.jl" begin
     include("test_aqua.jl")
 end
 
+@safetestset "Test Optim.jl" begin
+    include("test_optim.jl")
+end
+
 @safetestset "Initial error handling test" begin
     include("test_initial_errors.jl")
 end