Repository: secondmind-labs/trieste

Issue:
Pareto set: n objectives
As a user, I want functionality available to calculate the cell bounds of a Pareto front for the case of any number of objectives, so that I can easily find the cell bounds when defining multi-objective acquisition functionality

This should be valid for any number of objectives, including 1 and 2


Discussion:

Resolving patch:
diff --git a/trieste/utils/__init__.py b/trieste/utils/__init__.py
--- a/trieste/utils/__init__.py
+++ b/trieste/utils/__init__.py
@@ -17,18 +17,7 @@
 from typing_extensions import Final
 
 from . import objectives, pareto
-from .misc import Err, Ok, Result, T_co, jit, shapes_equal, to_numpy
-
-
-class DEFAULTS:
-    """ Default constants used in Trieste. """
-
-    JITTER: Final[float] = 1e-6
-    """
-    The default jitter, typically used to stabilise computations near singular points, such as in
-    Cholesky decomposition.
-    """
-
+from .misc import DEFAULTS, Err, Ok, Result, T_co, jit, shapes_equal, to_numpy
 
 K = TypeVar("K")
 """ An unbound type variable. """
diff --git a/trieste/utils/misc.py b/trieste/utils/misc.py
--- a/trieste/utils/misc.py
+++ b/trieste/utils/misc.py
@@ -19,7 +19,7 @@
 
 import numpy as np
 import tensorflow as tf
-from typing_extensions import final
+from typing_extensions import Final, final
 
 from ..type import TensorType
 
@@ -171,3 +171,13 @@ def unwrap(self) -> NoReturn:
         :raise Exception: Always. Raises the wrapped exception.
         """
         raise self._exc
+
+
+class DEFAULTS:
+    """ Default constants used in Trieste. """
+
+    JITTER: Final[float] = 1e-6
+    """
+    The default jitter, typically used to stabilise computations near singular points, such as in
+    Cholesky decomposition.
+    """
diff --git a/trieste/utils/pareto.py b/trieste/utils/pareto.py
--- a/trieste/utils/pareto.py
+++ b/trieste/utils/pareto.py
@@ -13,12 +13,15 @@
 # limitations under the License.
 """ This module contains functions and classes for Pareto based multi-objective optimization. """
 
+from __future__ import annotations
+
 from typing import Tuple
 
 import tensorflow as tf
 from typing_extensions import Final
 
 from ..type import TensorType
+from .misc import DEFAULTS
 
 
 def non_dominated(observations: TensorType) -> Tuple[TensorType, TensorType]:
@@ -71,27 +74,38 @@ class Pareto:
     A :class:`Pareto` Construct a Pareto set.
     Stores a Pareto set and calculates the cell bounds covering the non-dominated region.
     The latter is needed for certain multiobjective acquisition functions.
+
+    For hypervolume-based multiobjective optimisation with n>2 objectives, this class
+    implements branch and bound procedure algorithm. a divide and conquer method introduced
+    in :cite:`Couckuyt2012`.
     """
 
-    def __init__(self, observations: TensorType):
+    def __init__(self, observations: TensorType, *, jitter: float = DEFAULTS.JITTER):
         """
-        :param observations: The observations for all objectives, with shape [N, 2].
+        :param observations: The observations for all objectives, with shape [N, D].
         :raise ValueError (or InvalidArgumentError): If ``observations`` has an invalid shape.
         """
-        tf.debugging.assert_shapes([(observations, [None, 2])])
+        tf.debugging.assert_rank(observations, 2)
+        tf.debugging.assert_greater_equal(tf.shape(observations)[-1], 2)
 
-        pf, _ = non_dominated(observations)
-        self.front: Final[TensorType] = tf.gather_nd(pf, tf.argsort(pf[:, :1], axis=0))
-        self.bounds: Final[BoundedVolumes] = self._bounds_2d(self.front)
+        pfront, _ = non_dominated(observations)
+        self.front: Final[TensorType] = tf.gather_nd(pfront, tf.argsort(pfront[:, :1], axis=0))
+        self.bounds = self._get_bounds(self.front, jitter)
+
+    def _get_bounds(self, front: TensorType, jitter: float) -> BoundedVolumes:
+        if front.shape[-1] > 2:
+            return self._divide_conquer_nd(front, jitter)
+        else:
+            return self._bounds_2d(front)
 
     @staticmethod
     def _bounds_2d(front: TensorType) -> BoundedVolumes:
-
+        # Compute the cells covering the non-dominated region for 2 dimension case
         # this assumes the Pareto set has been sorted in ascending order on the first
         # objective, which implies the second objective is sorted in descending order
         len_front, number_of_objectives = front.shape
 
-        pf_ext_idx = tf.concat(
+        pseudo_front_idx = tf.concat(
             [
                 tf.zeros([1, number_of_objectives], dtype=tf.int32),
                 tf.argsort(front, axis=0) + 1,
@@ -101,10 +115,170 @@ def _bounds_2d(front: TensorType) -> BoundedVolumes:
         )
 
         range_ = tf.range(len_front + 1)[:, None]
-        lower = tf.concat([range_, tf.zeros_like(range_)], axis=-1)
-        upper = tf.concat([range_ + 1, pf_ext_idx[::-1, 1:][: pf_ext_idx[-1, 0]]], axis=-1)
+        lower_result = tf.concat([range_, tf.zeros_like(range_)], axis=-1)
+        upper_result = tf.concat(
+            [range_ + 1, pseudo_front_idx[::-1, 1:][: pseudo_front_idx[-1, 0]]], axis=-1
+        )
+
+        return BoundedVolumes(lower_result, upper_result)
+
+    @classmethod
+    def _divide_conquer_nd(
+        cls, front: TensorType, jitter: float, threshold: TensorType | float = 0
+    ) -> TensorType:
+        # Branch and bound procedure to compute the cells covering the non-dominated region.
+        # Generic version: works for an arbitrary number of objectives.
+
+        len_front, number_of_objectives = front.shape
+        lower_result = tf.zeros([0, number_of_objectives], dtype=tf.int32)
+        upper_result = tf.zeros([0, number_of_objectives], dtype=tf.int32)
+
+        min_front = tf.reduce_min(front, axis=0, keepdims=True) - 1
+        max_front = tf.reduce_max(front, axis=0, keepdims=True) + 1
+        pseudo_front = tf.concat([min_front, front, max_front], axis=0)
+
+        pseudo_front_idx = tf.concat(
+            [
+                tf.zeros([1, number_of_objectives], dtype=tf.int32),
+                tf.argsort(front, axis=0) + 1,  # +1 as index zero is reserved for the ideal point
+                tf.ones([1, number_of_objectives], dtype=tf.int32) * len_front + 1,
+            ],
+            axis=0,
+        )
+
+        divide_conquer_cells = tf.stack(
+            [
+                tf.zeros(number_of_objectives, dtype=tf.int32),
+                (int(pseudo_front_idx.shape[0]) - 1)
+                * tf.ones(number_of_objectives, dtype=tf.int32),
+            ],
+            axis=0,
+        )[None]
+
+        total_size = tf.reduce_prod(max_front - min_front)
+
+        def while_body(
+            divide_conquer_cells: TensorType,
+            lower_result: TensorType,
+            upper_result: TensorType,
+        ) -> Tuple[TensorType, TensorType, TensorType]:
+            divide_conquer_cells_unstacked = tf.unstack(divide_conquer_cells, axis=0)
+            cell = divide_conquer_cells_unstacked[-1]
+            divide_conquer_cells_new = tf.cond(
+                tf.not_equal(tf.size(divide_conquer_cells_unstacked[:-1]), 0),
+                lambda: tf.stack(divide_conquer_cells_unstacked[:-1]),
+                lambda: tf.zeros([0, 2, number_of_objectives], dtype=tf.int32),
+            )
+
+            arr = tf.range(number_of_objectives)
+            lower_idx = tf.gather_nd(pseudo_front_idx, tf.stack((cell[0], arr), -1))
+            upper_idx = tf.gather_nd(pseudo_front_idx, tf.stack((cell[1], arr), -1))
+            lower = tf.gather_nd(pseudo_front, tf.stack((lower_idx, arr), -1))
+            upper = tf.gather_nd(pseudo_front, tf.stack((upper_idx, arr), -1))
+
+            test_accepted = cls._is_test_required((upper - jitter) < front)
+            lower_result_final, upper_result_final = tf.cond(
+                test_accepted,
+                lambda: cls._accepted_test_body(lower_result, upper_result, lower_idx, upper_idx),
+                lambda: (lower_result, upper_result),
+            )
+
+            test_rejected = cls._is_test_required((lower + jitter) < front)
+            divide_conquer_cells_final = tf.cond(
+                tf.logical_and(test_rejected, tf.logical_not(test_accepted)),
+                lambda: cls._rejected_test_body(
+                    cell, lower, upper, divide_conquer_cells_new, total_size, threshold
+                ),
+                lambda: divide_conquer_cells_new,
+            )
+
+            return divide_conquer_cells_final, lower_result_final, upper_result_final
+
+        _, lower_result_final, upper_result_final = tf.while_loop(
+            lambda divide_conquer_cells, lower_result, upper_result: len(divide_conquer_cells) > 0,
+            while_body,
+            loop_vars=[divide_conquer_cells, lower_result, upper_result],
+            shape_invariants=[
+                tf.TensorShape([None, 2, number_of_objectives]),
+                tf.TensorShape([None, number_of_objectives]),
+                tf.TensorShape([None, number_of_objectives]),
+            ],
+        )
+        return BoundedVolumes(lower_result_final, upper_result_final)
+
+    @staticmethod
+    def _is_test_required(smaller: TensorType) -> TensorType:
+        idx_dom_augm = tf.reduce_any(smaller, axis=1)
+        is_dom_augm = tf.reduce_all(idx_dom_augm)
+
+        return is_dom_augm
+
+    @staticmethod
+    def _accepted_test_body(
+        lower_result: TensorType,
+        upper_result: TensorType,
+        lower_idx: TensorType,
+        upper_idx: TensorType,
+    ) -> Tuple[TensorType, TensorType]:
+        lower_result_accepted = tf.concat([lower_result, lower_idx[None]], axis=0)
+        upper_result_accepted = tf.concat([upper_result, upper_idx[None]], axis=0)
+        return lower_result_accepted, upper_result_accepted
+
+    @classmethod
+    def _rejected_test_body(
+        cls,
+        cell: TensorType,
+        lower: TensorType,
+        upper: TensorType,
+        divide_conquer_cells: TensorType,
+        total_size: TensorType,
+        threshold: TensorType,
+    ) -> TensorType:
+
+        divide_conquer_cells_dist = cell[1] - cell[0]
+        hc_size = tf.math.reduce_prod(upper - lower, axis=0, keepdims=True)
+
+        not_unit_cell = tf.reduce_any(divide_conquer_cells_dist > 1)
+        vol_above_thresh = tf.reduce_all((hc_size[0] / total_size) > threshold)
+        divide_conquer_cells_rejected = tf.cond(
+            tf.logical_and(not_unit_cell, vol_above_thresh),
+            lambda: cls._divide_body(divide_conquer_cells, divide_conquer_cells_dist, cell),
+            lambda: tf.identity(divide_conquer_cells),
+        )
+        return divide_conquer_cells_rejected
+
+    @staticmethod
+    def _divide_body(
+        divide_conquer_cells: TensorType,
+        divide_conquer_cells_dist: TensorType,
+        cell: TensorType,
+    ) -> TensorType:
+
+        edge_size = tf.reduce_max(divide_conquer_cells_dist)
+        idx = tf.argmax(divide_conquer_cells_dist)
+        edge_size1 = int(tf.round(tf.cast(edge_size, dtype=tf.float32) / 2.0))
+        edge_size2 = int(edge_size - edge_size1)
+
+        sparse_edge_size1 = tf.concat(
+            [tf.zeros([idx]), edge_size1 * tf.ones([1]), tf.zeros([len(cell[1]) - idx - 1])], axis=0
+        )
+        upper = tf.identity(cell[1]) - tf.cast(sparse_edge_size1, dtype=tf.int32)
+
+        divide_conquer_cells_new = tf.concat(
+            [divide_conquer_cells, tf.stack([tf.identity(cell[0]), upper], axis=0)[None]], axis=0
+        )
+
+        sparse_edge_size2 = tf.concat(
+            [tf.zeros([idx]), edge_size2 * tf.ones([1]), tf.zeros([len(cell[1]) - idx - 1])], axis=0
+        )
+        lower = tf.identity(cell[0]) + tf.cast(sparse_edge_size2, dtype=tf.int32)
+
+        divide_conquer_cells_final = tf.concat(
+            [divide_conquer_cells_new, tf.stack([lower, tf.identity(cell[1])], axis=0)[None]],
+            axis=0,
+        )
 
-        return BoundedVolumes(lower, upper)
+        return divide_conquer_cells_final
 
     def hypervolume_indicator(self, reference: TensorType) -> TensorType:
         """
@@ -132,8 +306,8 @@ def hypervolume_indicator(self, reference: TensorType) -> TensorType:
             ]
         )
 
-        min_pfront = tf.reduce_min(self.front, 0, keepdims=True)
-        pseudo_pfront = tf.concat((min_pfront, self.front, reference[None]), 0)
+        min_front = tf.reduce_min(self.front, 0, keepdims=True)
+        pseudo_front = tf.concat((min_front, self.front, reference[None]), 0)
         N, D = tf.shape(self.bounds.upper_idx)
 
         idx = tf.tile(tf.expand_dims(tf.range(D), -1), [1, N])
@@ -143,8 +317,8 @@ def hypervolume_indicator(self, reference: TensorType) -> TensorType:
         lower_idx = tf.reshape(
             tf.stack([tf.transpose(self.bounds.lower_idx), idx], axis=2), [N * D, 2]
         )
-        upper = tf.reshape(tf.gather_nd(pseudo_pfront, upper_idx), [D, N])
-        lower = tf.reshape(tf.gather_nd(pseudo_pfront, lower_idx), [D, N])
+        upper = tf.reshape(tf.gather_nd(pseudo_front, upper_idx), [D, N])
+        lower = tf.reshape(tf.gather_nd(pseudo_front, lower_idx), [D, N])
         hypervolume = tf.reduce_sum(tf.reduce_prod(upper - lower, 0))
 
-        return tf.reduce_prod(reference[None] - min_pfront) - hypervolume
+        return tf.reduce_prod(reference[None] - min_front) - hypervolume


Test patch:
diff --git a/tests/unit/test_pareto.py b/tests/unit/test_pareto.py
--- a/tests/unit/test_pareto.py
+++ b/tests/unit/test_pareto.py
@@ -84,12 +84,32 @@
             ),
             tf.constant([2, 6, 0, 9, 2, 0, 0, 3, 0, 0]),
         ),
+        (
+            tf.constant(
+                [
+                    [0.90234935, 0.02297473, 0.05389869],
+                    [0.98328614, 0.44182944, 0.6975261],
+                    [0.39555323, 0.3040712, 0.3433497],
+                    [0.72582424, 0.55389977, 0.00330079],
+                    [0.9590585, 0.03233206, 0.2403127],
+                    [0.04540098, 0.22407162, 0.11227596],
+                ]
+            ),
+            tf.constant(
+                [
+                    [0.90234935, 0.02297473, 0.05389869],
+                    [0.72582424, 0.55389977, 0.00330079],
+                    [0.04540098, 0.22407162, 0.11227596],
+                ]
+            ),
+            tf.constant([0, 4, 1, 0, 1, 0]),
+        ),
     ],
 )
 def test_dominated_sort(scores: tf.Tensor, pareto_set: tf.Tensor, dominance: tf.Tensor) -> None:
     ret_pareto_set, ret_dominance = non_dominated(scores)
     npt.assert_allclose(ret_pareto_set, pareto_set)
-    npt.assert_allclose(ret_dominance, dominance)
+    npt.assert_array_equal(ret_dominance, dominance)
 
 
 def test_pareto_2d_bounds() -> None:
@@ -108,8 +128,12 @@ def test_pareto_2d_bounds() -> None:
 
     pareto_2d = Pareto(objectives)
 
-    npt.assert_allclose(pareto_2d.bounds.lower_idx, tf.constant([[0, 0], [1, 0], [2, 0], [3, 0]]))
-    npt.assert_allclose(pareto_2d.bounds.upper_idx, tf.constant([[1, 4], [2, 1], [3, 2], [4, 3]]))
+    npt.assert_array_equal(
+        pareto_2d.bounds.lower_idx, tf.constant([[0, 0], [1, 0], [2, 0], [3, 0]])
+    )
+    npt.assert_array_equal(
+        pareto_2d.bounds.upper_idx, tf.constant([[1, 4], [2, 1], [3, 2], [4, 3]])
+    )
     npt.assert_allclose(
         pareto_2d.front, tf.constant([[0.1419, 0.9340], [0.1576, 0.7922], [0.4854, 0.0357]])
     )
@@ -145,6 +169,7 @@ def test_pareto_hypervolume_indicator_raises_for_reference_below_anti_ideal_poin
             [0.1, -0.6],
             0.37,
         ),
+        ([[2.0, 2.0, 0.0], [2.0, 0.0, 1.0], [3.0, 1.0, 0.0]], [4.0, 4.0, 4.0], 29.0),
     ],
 )
 def test_pareto_hypervolume_indicator(
@@ -152,3 +177,63 @@ def test_pareto_hypervolume_indicator(
 ) -> None:
     pareto = Pareto(tf.constant(objectives))
     npt.assert_allclose(pareto.hypervolume_indicator(tf.constant(reference)), expected)
+
+
+def test_pareto_divide_conquer_nd_three_dimension_case() -> None:
+    objectives = tf.constant(
+        [
+            [9.0, 0.0, 1.0],
+            [10.0, 4.0, 7.0],
+            [4.0, 3.0, 3.0],
+            [7.0, 6.0, 0.0],
+            [10.0, 0.0, 2.0],
+            [0.0, 2.0, 1.0],
+        ]
+    )
+
+    pareto = Pareto(objectives)
+
+    npt.assert_array_equal(
+        pareto.bounds.lower_idx,
+        tf.constant(
+            [
+                [3, 2, 0],
+                [3, 1, 0],
+                [2, 2, 0],
+                [2, 1, 0],
+                [3, 0, 1],
+                [2, 0, 1],
+                [2, 0, 0],
+                [0, 1, 1],
+                [0, 1, 0],
+                [0, 0, 0],
+            ]
+        ),
+    )
+    npt.assert_array_equal(
+        pareto.bounds.upper_idx,
+        tf.constant(
+            [
+                [4, 4, 2],
+                [4, 2, 1],
+                [3, 4, 2],
+                [3, 2, 1],
+                [4, 3, 4],
+                [3, 1, 4],
+                [4, 1, 1],
+                [1, 4, 4],
+                [2, 4, 1],
+                [2, 1, 4],
+            ]
+        ),
+    )
+    npt.assert_allclose(
+        pareto.front,
+        tf.constant(
+            [
+                [0.0, 2.0, 1.0],
+                [7.0, 6.0, 0.0],
+                [9.0, 0.0, 1.0],
+            ]
+        ),
+    )