From ee80d5e8302693fe7487409eb681bac9ae81d37e Mon Sep 17 00:00:00 2001
From: Rui Coelho <ruicoelhopedro@gmail.com>
Date: Fri, 29 Mar 2024 17:34:16 +0000
Subject: [PATCH] Expand composite observations to full size

---
 piglot/optimisers/botorch/bayes.py   |  4 +++-
 piglot/optimisers/botorch/dataset.py | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/piglot/optimisers/botorch/bayes.py b/piglot/optimisers/botorch/bayes.py
index 104832f..2b03b1b 100644
--- a/piglot/optimisers/botorch/bayes.py
+++ b/piglot/optimisers/botorch/bayes.py
@@ -252,7 +252,9 @@ def _build_acquisition_composite(
         # Build composite MC objective
         _, y_avg, y_std = dataset.get_obervation_stats()
         mc_objective = GenericMCObjective(
-            lambda vals, X: -self.objective.composition.composition_torch(vals * y_std + y_avg)
+            lambda vals, X: -self.objective.composition.composition_torch(
+                dataset.expand_observations(vals * y_std + y_avg)
+            )
         )
         # Delegate acquisition building
         best = torch.max(dataset.values).item()
diff --git a/piglot/optimisers/botorch/dataset.py b/piglot/optimisers/botorch/dataset.py
index 2ee70be..30c6ea0 100644
--- a/piglot/optimisers/botorch/dataset.py
+++ b/piglot/optimisers/botorch/dataset.py
@@ -113,6 +113,31 @@ def get_obervation_stats(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor
         y_std = y_std[mask]
         return mask, y_avg, y_std
 
+    def expand_observations(self, reduced: torch.Tensor) -> torch.Tensor:
+        """Expand reduced observations to full size.
+        Parameters
+        ----------
+        reduced : torch.Tensor
+            Reduced observations.
+        Returns
+        -------
+        torch.Tensor
+            Expanded observations.
+        """
+        # Infer the shape of the expanded tensor: only modify the last dimension
+        new_shape = list(reduced.shape)
+        new_shape[-1] = self.n_outputs
+        expanded = torch.empty(new_shape, dtype=self.dtype, device=self.device)
+        # Collapse the tensor to a 2D array for indexing the last dimension
+        expanded_flat = expanded.view(-1, self.n_outputs)
+        mask, _, _ = self.get_obervation_stats()
+        expanded_flat[:, mask] = reduced.view(-1, reduced.shape[-1])
+        # Fill the missing values with the average of the observed ones
+        y_avg = torch.mean(self.values, dim=-2)
+        expanded_flat[:, ~mask] = y_avg[~mask]
+        # Note: we are using a view, so the expanded tensor is already modified
+        return expanded
+
     def standardised(self) -> BayesDataset:
         """Return a dataset with unit-cube parameters and standardised outputs.