Avoid discarded objective evaluations in OptimizationOptimisers

lib/OptimizationOptimisers/Project.toml

@@ -24,6 +24,7 @@ Reexport = "1.2.2"
 SciMLBase = "2.122.1, 3"
 julia = "1.10"
 ComponentArrays = "0.15"
+CommonSolve = "0.2"
 ForwardDiff = "0.10, 1"
 Lux = "1"
 MLDataDevices = "1"
@@ -38,6 +39,7 @@ SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
@@ -48,4 +50,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["SciMLTesting", "ComponentArrays", "ForwardDiff", "Lux", "MLDataDevices", "MLUtils", "Random", "Test", "Zygote", "Printf", "SafeTestsets", "Pkg"]
+test = ["SciMLTesting", "ComponentArrays", "CommonSolve", "ForwardDiff", "Lux", "MLDataDevices", "MLUtils", "Random", "Test", "Zygote", "Printf", "SafeTestsets", "Pkg"]

lib/OptimizationOptimisers/src/OptimizationOptimisers.jl

@@ -61,7 +61,7 @@ function SciMLBase.__solve(cache::OptimizationCache{O}) where {O <: AbstractRule
     G = copy(θ)
 
     local x, min_err, min_θ
-    x = cache.f(cache.u0, first(data))
+    x = convert(eltype(real(cache.u0)), Inf)
     min_err = typemax(eltype(real(cache.u0))) #dummy variables
     min_opt = 1
     min_θ = cache.u0
@@ -151,6 +151,10 @@ function SciMLBase.__solve(cache::OptimizationCache{O}) where {O <: AbstractRule
             )
         end
     end
+    if iszero(iterations)
+        x = cache.f(cache.u0, first(data))
+        fevals += 1
+    end
     cache.progress && @logmsg(
         LogLevel(-1), "Optimization",
         _id = progress_id, message = "Done", progress = 1.0

lib/OptimizationOptimisers/test/core_tests.jl

@@ -1,4 +1,7 @@
 using OptimizationOptimisers, ForwardDiff, OptimizationBase
+using CommonSolve: solve
+using Optimisers
+using SciMLBase: OptimizationFunction, OptimizationProblem
 using Test
 using Zygote
 using Lux, MLUtils, Random, ComponentArrays, Printf, MLDataDevices
@@ -117,6 +120,59 @@ using Lux, MLUtils, Random, ComponentArrays, Printf, MLDataDevices
         )
     end
 
+    @testset "Objective evaluations" begin
+        objective_calls = Ref(0)
+        fg_calls = Ref(0)
+        function counted_objective(x, p)
+            objective_calls[] += 1
+            return sum(abs2, x)
+        end
+        function counted_fg!(G, x, p)
+            fg_calls[] += 1
+            G .= 2 .* x
+            return sum(abs2, x)
+        end
+
+        optf = OptimizationFunction(counted_objective; fg = counted_fg!)
+        prob = OptimizationProblem(optf, ones(2), nothing)
+        sol = solve(prob, Optimisers.Adam(); maxiters = 3, save_best = false)
+
+        @test sol.stats.iterations == 3
+        @test fg_calls[] == sol.stats.iterations
+        @test objective_calls[] == 0
+
+        objective_calls[] = 0
+        fg_calls[] = 0
+        sol = solve(prob, Optimisers.Adam(); maxiters = 0.4, save_best = false)
+
+        @test sol.stats.iterations == 0
+        @test sol.stats.fevals == objective_calls[] == 1
+        @test fg_calls[] == 0
+
+        inf_objective_calls = Ref(0)
+        inf_fg_calls = Ref(0)
+        function inf_objective(x, p)
+            inf_objective_calls[] += 1
+            return Inf32
+        end
+        function inf_fg!(G, x, p)
+            inf_fg_calls[] += 1
+            G .= 0
+            return Inf32
+        end
+
+        inf_optf = OptimizationFunction(inf_objective; fg = inf_fg!)
+        inf_prob = OptimizationProblem(inf_optf, ones(Float32, 2), nothing)
+        inf_sol = solve(inf_prob, Optimisers.Adam(); maxiters = 1, save_best = false)
+
+        @test isinf(inf_sol.objective)
+        @test inf_sol.objective isa Float32
+        @test inf_sol.stats.iterations == 1
+        @test inf_sol.stats.fevals == 1
+        @test inf_objective_calls[] == 0
+        @test inf_fg_calls[] == 1
+    end
+
     @test_throws ArgumentError sol = solve(prob, Optimisers.Adam())
 end