refactor: optim tests

Author: MilesCranmer

ext/DynamicExpressionsOptimExt.jl

@@ -42,7 +42,7 @@ function wrap_func(
     f::F, tree::N, refs
 ) where {F<:Function,T,N<:Union{AbstractExpressionNode{T},AbstractExpression{T}}}
     function wrapped_f(args::Vararg{Any,M}) where {M}
-        first_args = args[begin:end-1]
+        first_args = args[begin:(end - 1)]
         x = args[end]
         set_constants!(tree, x, refs)
         return @inline(f(first_args..., tree))

test/test_optim.jl

@@ -1,24 +1,7 @@
-using DynamicExpressions, Optim, Zygote
-using Random: MersenneTwister as RNG
-using Test
+@testitem "Basic optimization" begin
+    using DynamicExpressions, Optim
 
-operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(exp,))
-x1, x2 = (i -> Node(Float64; feature=i)).(1:2)
-
-X = rand(RNG(0), Float64, 2, 100)
-y = @. exp(X[1, :] * 2.1 - 0.9) + X[2, :] * -0.9
-
-original_tree = exp(x1 * 0.8 - 0.0) + 5.2 * x2
-target_tree = exp(x1 * 2.1 - 0.9) + -0.9 * x2
-
-f(tree) = sum(abs2, tree(X, operators) .- y)
-function g!(G, tree)
-    dy = only(gradient(f, tree))
-    G .= dy.gradient
-    return nothing
-end
-
-@testset "Basic optimization" begin
+    include("test_optim_setup.jl")
     tree = copy(original_tree)
     res = optimize(f, tree)
 
@@ -33,7 +16,11 @@ end
     )
 end
 
-@testset "With gradients, using Zygote" begin
+@testitem "With gradients, using Zygote" begin
+    using DynamicExpressions, Optim, Zygote
+
+    include("test_optim_setup.jl")
+
     tree = copy(original_tree)
     res = optimize(f, g!, tree, BFGS())
     @test tree == original_tree
@@ -42,7 +29,11 @@ end
     )
 end
 
-@testset "With gradients, manually" begin
+@testitem "With gradients, manually" begin
+    using DynamicExpressions, Optim, Zygote
+
+    include("test_optim_setup.jl")
+
     tree = copy(original_tree)
     did_i_run = Ref(false)
     # Now, try with gradients too (via Zygote and our hand-rolled forward-mode AD)
@@ -79,7 +70,10 @@ end
 end
 
 # Now, try combined
-@testset "Combined evaluation with gradient" begin
+@testitem "Combined evaluation with gradient" begin
+    using DynamicExpressions, Optim, Zygote
+    include("test_optim_setup.jl")
+
     tree = copy(original_tree)
     did_i_run_2 = Ref(false)
     fg!(F, G, tree) =

test/test_optim_setup.jl

@@ -0,0 +1,18 @@
+using DynamicExpressions
+using Random: MersenneTwister as RNG
+
+operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(exp,))
+x1, x2 = (i -> Node(Float64; feature=i)).(1:2)
+
+X = rand(RNG(0), Float64, 2, 100)
+y = @. exp(X[1, :] * 2.1 - 0.9) + X[2, :] * -0.9
+
+original_tree = exp(x1 * 0.8 - 0.0) + 5.2 * x2
+target_tree = exp(x1 * 2.1 - 0.9) + -0.9 * x2
+
+f(tree) = sum(abs2, tree(X, operators) .- y)
+function g!(G, tree)
+    dy = only(gradient(f, tree))
+    G .= dy.gradient
+    return nothing
+end

test/unittest.jl

@@ -16,9 +16,7 @@ using Zygote, SymbolicUtils, LoopVectorization, Bumper, Optim
     include("test_deprecations.jl")
 end
 
-@testitem "Test Optim.jl" begin
-    include("test_optim.jl")
-end
+include("test_optim.jl")
 
 @testitem "Test tree construction and scoring" begin
     include("test_tree_construction.jl")