Merge pull request #994 from ChrisRackauckas-Claude/fix-store-trace-optimjl

Fix error when store_trace=true with OptimizationOptimJL

Author: ChrisRackauckas

lib/OptimizationOptimJL/src/OptimizationOptimJL.jl

@@ -139,16 +139,17 @@ function SciMLBase.__solve(cache::OptimizationCache{
         error("Use OptimizationFunction to pass the derivatives or automatically generate them with one of the autodiff backends")
 
     function _cb(trace)
-        metadata = decompose_trace(trace).metadata
+        trace_state = decompose_trace(trace)
+        metadata = trace_state.metadata
         θ = metadata[cache.opt isa Optim.NelderMead ? "centroid" : "x"]
-        opt_state = Optimization.OptimizationState(iter = trace.iteration,
+        opt_state = Optimization.OptimizationState(iter = trace_state.iteration,
             u = θ,
             p = cache.p,
-            objective = trace.value,
+            objective = trace_state.value,
             grad = get(metadata, "g(x)", nothing),
             hess = get(metadata, "h(x)", nothing),
             original = trace)
-        cb_call = cache.callback(opt_state, trace.value)
+        cb_call = cache.callback(opt_state, trace_state.value)
         if !(cb_call isa Bool)
             error("The callback should return a boolean `halt` for whether to stop the optimization process.")
         end
@@ -257,18 +258,19 @@ function SciMLBase.__solve(cache::OptimizationCache{
     local x, cur, state
 
     function _cb(trace)
-        metadata = decompose_trace(trace).metadata
+        trace_state = decompose_trace(trace)
+        metadata = trace_state.metadata
         θ = !(cache.opt isa Optim.SAMIN) && cache.opt.method == Optim.NelderMead() ?
             metadata["centroid"] :
             metadata["x"]
-        opt_state = Optimization.OptimizationState(iter = trace.iteration,
+        opt_state = Optimization.OptimizationState(iter = trace_state.iteration,
             u = θ,
             p = cache.p,
-            objective = trace.value,
+            objective = trace_state.value,
             grad = get(metadata, "g(x)", nothing),
             hess = get(metadata, "h(x)", nothing),
             original = trace)
-        cb_call = cache.callback(opt_state, trace.value)
+        cb_call = cache.callback(opt_state, trace_state.value)
         if !(cb_call isa Bool)
             error("The callback should return a boolean `halt` for whether to stop the optimization process.")
         end

lib/OptimizationOptimJL/test/runtests.jl

@@ -213,4 +213,27 @@ end
         sol = Optimization.solve!(cache)
         @test sol.u≈[2.0] atol=1e-3
     end
+
+    @testset "store_trace=true" begin
+        # Test that store_trace=true works without throwing errors (issue #990)
+        rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+        x0 = zeros(2)
+        _p = [1.0, 100.0]
+
+        # Test with NelderMead
+        prob = OptimizationProblem(rosenbrock, x0, _p)
+        sol = solve(prob, NelderMead(), store_trace = true)
+        @test sol isa Any  # just test it doesn't throw
+
+        # Test with Fminbox(NelderMead)
+        optprob = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff())
+        prob = OptimizationProblem(optprob, x0, _p, lb = [-1.0, -1.0], ub = [0.8, 0.8])
+        sol = solve(prob, Optim.Fminbox(NelderMead()), store_trace = true)
+        @test sol isa Any  # just test it doesn't throw
+
+        # Test with BFGS
+        prob = OptimizationProblem(optprob, x0, _p)
+        sol = solve(prob, BFGS(), store_trace = true)
+        @test sol isa Any  # just test it doesn't throw
+    end
 end