From ddbb5ef497859e7e382df0fbc6594960008347e3 Mon Sep 17 00:00:00 2001
From: Marc Becker <33069354+be-marc@users.noreply.github.com>
Date: Thu, 7 Nov 2024 11:49:45 +0100
Subject: [PATCH] refactor: pass extra information of the result in the extra
 parameter (#168)

* refactor: pass extra information of the result in the extra parameter

* remove remotes
---
 DESCRIPTION                 |  2 +-
 R/ResultAssignerArchive.R   | 15 ++++++++++-----
 R/ResultAssignerSurrogate.R |  3 ++-
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index 70563fe6..0424c1fb 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -41,7 +41,7 @@ BugReports: https://github.com/mlr-org/mlr3mbo/issues
 Depends:
     R (>= 3.1.0)
 Imports:
-    bbotk (>= 1.1.1),
+    bbotk (>= 1.2.0),
     checkmate (>= 2.0.0),
     data.table,
     lgr (>= 0.3.4),
diff --git a/R/ResultAssignerArchive.R b/R/ResultAssignerArchive.R
index c9266fdb..12370377 100644
--- a/R/ResultAssignerArchive.R
+++ b/R/ResultAssignerArchive.R
@@ -30,14 +30,19 @@ ResultAssignerArchive = R6Class("ResultAssignerArchive",
     #'   The [bbotk::OptimInstance] the final result should be assigned to.
     assign_result = function(instance) {
       xydt = instance$archive$best()
-      xdt = xydt[, instance$search_space$ids(), with = FALSE]
+      cols_x = instance$archive$cols_x
+      cols_y = instance$archive$cols_y
+
+      xdt = xydt[, cols_x, with = FALSE]
+      extra = xydt[, !c(cols_x, cols_y), with = FALSE]
+
       if (inherits(instance, "OptimInstanceBatchMultiCrit")) {
-        ydt = xydt[, instance$archive$cols_y, with = FALSE]
-        instance$assign_result(xdt, ydt, xydt = xydt)
+        ydt = xydt[, cols_y, with = FALSE]
+        instance$assign_result(xdt, ydt, extra = extra)
       }
       else {
-        y = unlist(xydt[, instance$archive$cols_y, with = FALSE])
-        instance$assign_result(xdt, y, xydt = xydt)
+        y = unlist(xydt[, cols_y, with = FALSE])
+        instance$assign_result(xdt, y, extra = extra)
       }
     }
   ),
diff --git a/R/ResultAssignerSurrogate.R b/R/ResultAssignerSurrogate.R
index 4870947a..59c3aa95 100644
--- a/R/ResultAssignerSurrogate.R
+++ b/R/ResultAssignerSurrogate.R
@@ -60,6 +60,7 @@ ResultAssignerSurrogate = R6Class("ResultAssignerSurrogate",
       archive_tmp = archive$clone(deep = TRUE)
       archive_tmp$data[, self$surrogate$cols_y := means]
       xydt = archive_tmp$best()
+      extra = xydt[, !c(archive_tmp$cols_x, archive_tmp$cols_y), with = FALSE]
       best = xydt[, archive_tmp$cols_x, with = FALSE]
 
       # ys are still the ones originally evaluated
@@ -68,7 +69,7 @@ ResultAssignerSurrogate = R6Class("ResultAssignerSurrogate",
       } else if (inherits(instance, "OptimInstanceBatchMultiCrit")) {
         archive$data[best, on = archive$cols_x][, archive$cols_y, with = FALSE]
       }
-      instance$assign_result(xdt = best, best_y, xydt = xydt)
+      instance$assign_result(xdt = best, best_y, extra = extra)
     }
   ),