style: cleaner test code

Author: MilesCranmer

test/test_optim.jl

@@ -80,29 +80,26 @@ end
 
     tree = copy(original_tree)
     did_i_run_2 = Ref(false)
-    fg!(F, G, tree) =
-        let
-            if G !== nothing
-                ŷ, dŷ_dconstants, _ = eval_grad_tree_array(
-                    tree, X, operators; variable=false
+    function my_fg!(F, G, tree)
+        if G !== nothing
+            ŷ, dŷ_dconstants, _ = eval_grad_tree_array(tree, X, operators; variable=false)
+            dresult_dŷ = @. 2 * (ŷ - y)
+            for i in eachindex(G)
+                G[i] = sum(
+                    j -> dresult_dŷ[j] * dŷ_dconstants[i, j],
+                    eachindex(axes(dŷ_dconstants, 2), axes(dresult_dŷ, 1)),
                 )
-                dresult_dŷ = @. 2 * (ŷ - y)
-                for i in eachindex(G)
-                    G[i] = sum(
-                        j -> dresult_dŷ[j] * dŷ_dconstants[i, j],
-                        eachindex(axes(dŷ_dconstants, 2), axes(dresult_dŷ, 1)),
-                    )
-                end
-                if F !== nothing
-                    did_i_run_2[] = true
-                    return sum(abs2, ŷ .- y)
-                end
-            elseif F !== nothing
-                # Only f
-                return sum(abs2, tree(X, operators) .- y)
             end
+            if F !== nothing
+                did_i_run_2[] = true
+                return sum(abs2, ŷ .- y)
+            end
+        elseif F !== nothing
+            # Only f
+            return sum(abs2, tree(X, operators) .- y)
         end
-    res = optimize(Optim.only_fg!(fg!), tree, BFGS())
+    end
+    res = optimize(Optim.only_fg!(my_fg!), tree, BFGS())
 
     @test did_i_run_2[]
     @test isapprox(

test/test_optim_setup.jl

@@ -1,16 +1,19 @@
 using DynamicExpressions
 using Random: MersenneTwister as RNG
 
-operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(exp,))
+const 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
+const X = rand(RNG(0), Float64, 2, 100)
+const 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 f(tree)
+    out = tree(X, operators)
+    return sum(i -> abs2(out[i] - y[i]), eachindex(out, y))
+end
 function g!(G, tree)
     dy = only(gradient(f, tree))
     G .= dy.gradient