Repository: secondmind-labs/trieste

Issue:
More space-filling sampling method than tf.uniform 
**Describe the feature you'd like**

In the `Box` class the sampling method is `tf.random.uniform`. I found it less space-filling than for example 
`tf.math.sobol_sample`. This will make the metamodel better approximate the function at the initial phase of optimization. it might be good if the user can choose what kind of sampler to use in the `Box` class. 
as a comparison, botorch had a set of samplers we can select ([link](https://botorch.org/api/sampling.html)) and in the GPflowOpt we have `Design` class ([link](https://github.com/GPflow/GPflowOpt/blob/master/gpflowopt/design.py)) . 

**Additional context**
to give some context what's available there: [link](https://como.ceb.cam.ac.uk/media/preprints/c4e-Preprint-182.pdf)


Discussion:
henrymoss: Hi @satrialoka  and @icouckuy, this PR is definitely in the right direction.
henrymoss: @satrialoka  please could you explain why you chose to have a separate design of experiment class. It feels like a lot of unneeded code. Why is this better than just having some extra methods to the Box search space (e.g. .sample_halton) ?
satrialoka: Yeah seems it's indeed add unneeded code. My reason to add that is to make users easily implement their own DoE/Sampler method by having a base abstract class for it. Do you think adding function in BOX is enough? If so, I'll be happy to change the current implementation. 
icouckuy: Indeed, we were also thinking about lhd and other designs. If we restrict to Sobol/Halton we can indeed simplify things
henrymoss: I think its worth just adding the functionality into box for now. :)
henrymoss: Brill. Please can you add you descriptions in the code to say what these functions do, what inputs/outputs they have and what the inputs do, i.e. following the format of other functions in that class.

You could also add some unit tests, like we have for Box.sample(), e.g. check the size of the output is as expected, check that it only works for valid number of samples, check that we get deterministic samples back for fixed choices of  seed/skip
satrialoka: for the test, I am just copying what they did in test_box_sampling.
please check!

Patch:
diff --git a/trieste/space.py b/trieste/space.py
index c3913ff9..fe2a794b 100644
--- a/trieste/space.py
+++ b/trieste/space.py
@@ -15,9 +15,10 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Sequence, TypeVar, overload
+from typing import Optional, Sequence, TypeVar, overload
 
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 from .type import TensorType
 from .utils import shapes_equal
@@ -173,7 +174,11 @@ class Box(SearchSpace):
     def __init__(self, lower: TensorType, upper: TensorType):
         ...
 
-    def __init__(self, lower: Sequence[float] | TensorType, upper: Sequence[float] | TensorType):
+    def __init__(
+        self,
+        lower: Sequence[float] | TensorType,
+        upper: Sequence[float] | TensorType,
+    ):
         r"""
         If ``lower`` and ``upper`` are `Sequence`\ s of floats (such as lists or tuples),
         they will be converted to tensors of dtype `tf.float64`.
@@ -188,6 +193,7 @@ class Box(SearchSpace):
             - ``lower`` and ``upper`` do not have the same floating point type.
             - ``upper`` is not greater than ``lower`` across all dimensions.
         """
+
         tf.debugging.assert_shapes([(lower, ["D"]), (upper, ["D"])])
         tf.assert_rank(lower, 1)
         tf.assert_rank(upper, 1)
@@ -241,11 +247,62 @@ class Box(SearchSpace):
         return tf.reduce_all(value >= self._lower) and tf.reduce_all(value <= self._upper)
 
     def sample(self, num_samples: int) -> TensorType:
+        """
+        Sample randomly from the space.
+
+        :param num_samples: The number of points to sample from this search space.
+        :return: ``num_samples`` i.i.d. random points, sampled uniformly, and without replacement,
+            from this search space with shape '[num_samples, D]' , where D is the search space
+            dimension.
+        """
+        tf.debugging.assert_non_negative(num_samples)
+
         dim = tf.shape(self._lower)[-1]
         return tf.random.uniform(
             (num_samples, dim), minval=self._lower, maxval=self._upper, dtype=self._lower.dtype
         )
 
+    def sample_halton(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
+        """
+        Sample from the space using a Halton sequence. The resulting samples are guaranteed to be
+        diverse and are reproducible by using the same choice of ``seed``.
+
+        :param num_samples: The number of points to sample from this search space.
+        :param seed: Random seed for the halton sequence
+        :return: ``num_samples`` of points, using halton sequence with shape '[num_samples, D]' ,
+            where D is the search space dimension.
+        """
+
+        tf.debugging.assert_non_negative(num_samples)
+        if num_samples == 0:
+            return tf.constant([])
+        if seed is not None:  # ensure reproducibility
+            tf.random.set_seed(seed)
+        dim = tf.shape(self._lower)[-1]
+        return (self._upper - self._lower) * tfp.mcmc.sample_halton_sequence(
+            dim=dim, num_results=num_samples, dtype=self._lower.dtype, seed=seed
+        ) + self._lower
+
+    def sample_sobol(self, num_samples: int, skip: Optional[int] = None) -> TensorType:
+        """
+        Sample a diverse set from the space using a Sobol sequence.
+        If ``skip`` is specified, then the resulting samples are reproducible.
+
+        :param num_samples: The number of points to sample from this search space.
+        :param skip: The number of initial points of the Sobol sequence to skip
+        :return: ``num_samples`` of points, using sobol sequence with shape '[num_samples, D]' ,
+            where D is the search space dimension.
+        """
+        tf.debugging.assert_non_negative(num_samples)
+        if num_samples == 0:
+            return tf.constant([])
+        if skip is None:  # generate random skip
+            skip = tf.random.uniform([1], maxval=2 ** 16, dtype=tf.int32)[0]
+        dim = tf.shape(self._lower)[-1]
+        return (self._upper - self._lower) * tf.math.sobol_sample(
+            dim=dim, num_results=num_samples, dtype=self._lower.dtype, skip=skip
+        ) + self._lower
+
     def discretize(self, num_samples: int) -> DiscreteSearchSpace:
         """
         :param num_samples: The number of points in the :class:`DiscreteSearchSpace`.


Test patch:
diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py
index 4ae18547..f7b55c06 100644
--- a/tests/unit/test_space.py
+++ b/tests/unit/test_space.py
@@ -302,12 +302,77 @@ def test_box_contains_raises_on_point_of_different_shape(
 
 
 @pytest.mark.parametrize("num_samples", [0, 1, 10])
-def test_box_sampling(num_samples: int) -> None:
+def test_box_sampling_returns_correct_shape(num_samples: int) -> None:
     box = Box(tf.zeros((3,)), tf.ones((3,)))
     samples = box.sample(num_samples)
     _assert_correct_number_of_unique_constrained_samples(num_samples, box, samples)
 
 
+@pytest.mark.parametrize("num_samples", [0, 1, 10])
+def test_box_sobol_sampling_returns_correct_shape(num_samples: int) -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    sobol_samples = box.sample_sobol(num_samples)
+    _assert_correct_number_of_unique_constrained_samples(num_samples, box, sobol_samples)
+
+
+@pytest.mark.parametrize("num_samples", [0, 1, 10])
+def test_box_halton_sampling_returns_correct_shape(num_samples: int) -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    halton_samples = box.sample_halton(num_samples)
+    _assert_correct_number_of_unique_constrained_samples(num_samples, box, halton_samples)
+
+
+@pytest.mark.parametrize("num_samples", [-1, -10])
+def test_box_sampling_raises_for_invalid_sample_size(num_samples: int) -> None:
+    with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
+        box = Box(tf.zeros((3,)), tf.ones((3,)))
+        box.sample(num_samples)
+
+
+@pytest.mark.parametrize("num_samples", [-1, -10])
+def test_box_sobol_sampling_raises_for_invalid_sample_size(num_samples: int) -> None:
+    with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
+        box = Box(tf.zeros((3,)), tf.ones((3,)))
+        box.sample_sobol(num_samples)
+
+
+@pytest.mark.parametrize("num_samples", [-1, -10])
+def test_box_halton_sampling_raises_for_invalid_sample_size(num_samples: int) -> None:
+    with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
+        box = Box(tf.zeros((3,)), tf.ones((3,)))
+        box.sample_halton(num_samples)
+
+
+@pytest.mark.parametrize("skip", [1, 10, 100])
+def test_box_sobol_sampling_returns_same_points_for_same_skip(skip: int) -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    sobol_samples_1 = box.sample_sobol(num_samples=100, skip=skip)
+    sobol_samples_2 = box.sample_sobol(num_samples=100, skip=skip)
+    npt.assert_allclose(sobol_samples_1, sobol_samples_2)
+
+
+@pytest.mark.parametrize("seed", [1, 42, 123])
+def test_box_halton_sampling_returns_same_points_for_same_seed(seed: int) -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    halton_samples_1 = box.sample_halton(num_samples=100, seed=seed)
+    halton_samples_2 = box.sample_halton(num_samples=100, seed=seed)
+    npt.assert_allclose(halton_samples_1, halton_samples_2)
+
+
+def test_box_sobol_sampling_returns_different_points_for_different_call() -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    sobol_samples_1 = box.sample_sobol(num_samples=100)
+    sobol_samples_2 = box.sample_sobol(num_samples=100)
+    npt.assert_raises(AssertionError, npt.assert_allclose, sobol_samples_1, sobol_samples_2)
+
+
+def test_box_haltom_sampling_returns_different_points_for_different_call() -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    halton_samples_1 = box.sample_halton(num_samples=100)
+    halton_samples_2 = box.sample_halton(num_samples=100)
+    npt.assert_raises(AssertionError, npt.assert_allclose, halton_samples_1, halton_samples_2)
+
+
 @pytest.mark.parametrize("num_samples", [0, 1, 10])
 def test_box_discretize_returns_search_space_with_only_points_contained_within_box(
     num_samples: int,