From 0bce2a80948dc368617c80a0a8264a1c0114a99e Mon Sep 17 00:00:00 2001
From: Paras Puneet Singh <136245940+ParasPuneetSingh@users.noreply.github.com>
Date: Fri, 26 Jul 2024 21:39:20 +0530
Subject: [PATCH] Update loss function in OptimizationBBO.jl

Updated the __loss function to handle MOOFunction results.
---
 lib/OptimizationBBO/src/OptimizationBBO.jl | 35 +++++++++++++++-------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/lib/OptimizationBBO/src/OptimizationBBO.jl b/lib/OptimizationBBO/src/OptimizationBBO.jl
index f4487c0ba..4b6c23431 100644
--- a/lib/OptimizationBBO/src/OptimizationBBO.jl
+++ b/lib/OptimizationBBO/src/OptimizationBBO.jl
@@ -142,17 +142,32 @@ function SciMLBase.__solve(cache::Optimization.OptimizationCache{
     maxtime = Optimization._check_and_convert_maxtime(cache.solver_args.maxtime)
 
     _loss = function (θ)
-        if cache.callback === Optimization.DEFAULT_CALLBACK &&
-           cache.data === Optimization.DEFAULT_DATA
-            return first(cache.f(θ, cache.p))
-        elseif cache.callback === Optimization.DEFAULT_CALLBACK
-            return first(cache.f(θ, cache.p, cur...))
-        elseif cache.data !== Optimization.DEFAULT_DATA
-            x = cache.f(θ, cache.p)
-            return first(x)
+        if isa(f, MultiObjectiveOptimizationFunction)
+            if cache.callback === Optimization.DEFAULT_CALLBACK &&
+               cache.data === Optimization.DEFAULT_DATA
+                return cache.f(θ, cache.p)
+            elseif cache.callback === Optimization.DEFAULT_CALLBACK
+                return cache.f(θ, cache.p, cur...)
+            elseif cache.data !== Optimization.DEFAULT_DATA
+                x = cache.f(θ, cache.p)
+                return x
+            else
+                x = cache.f(θ, cache.p, cur...)
+                return first(x)
+            end
         else
-            x = cache.f(θ, cache.p, cur...)
-            return first(x)
+            if cache.callback === Optimization.DEFAULT_CALLBACK &&
+               cache.data === Optimization.DEFAULT_DATA
+                return first(cache.f(θ, cache.p))
+            elseif cache.callback === Optimization.DEFAULT_CALLBACK
+                return first(cache.f(θ, cache.p, cur...))
+            elseif cache.data !== Optimization.DEFAULT_DATA
+                x = cache.f(θ, cache.p)
+                return first(x)
+            else
+                x = cache.f(θ, cache.p, cur...)
+                return first(x)
+            end
         end
     end