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:
I assume you're particularly looking at how the space is used in the acquisition optimizer. If so, it may also be good to let users customize the acquisition optimizer (which would indirectly allow users to customize the sampling in the optimizer, depending on how that's done). This is something we intend to do anyway
yes. additionally, it might be good to have that sampling customization for generating initial data before the optimization iteration. 
I have the custom sampler for the initial DoE at the moment. I added a [Design class ](https://github.com/satrialoka/trieste/commit/88a6f02050530df077516d775fadf8d7f57219b9#diff-b491084d9dee913da25dc734df837ed88169ba70299461477fc4c20640aae52c) for customizing the sampler/DoE method. Haven't check how it will interact with the optimizer though. The demo can be checked at [link](https://github.com/satrialoka/trieste/blob/custom-sampling/docs/notebooks/custom_sampler.pct.py). What do you think? @henrymoss 

for the Latin Hypercube Design, we have NumPy implementation in UGhent internal repo, If the current way to incorporate the custom sampler is ok for you, I can port it to TensorFlow. Or do you have another idea?

That looks great! I'm happy with the Sobol/Halton. Do you feel that the LHD is needed?
personally, I wouldn't include the LHD and stick with Sobol/Halton. Some BO library also doesn't have the LHD. If the user wants LHD, it can be implemented by using the Design class.   
Perfect. Are you happy to do a full PR with this functionality? 
Okay :ok_hand:

Resolving patch:
diff --git a/trieste/space.py b/trieste/space.py
--- 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 @@ def __init__(self, lower: Sequence[float], upper: Sequence[float]):
     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 @@ def __init__(self, lower: Sequence[float] | TensorType, upper: Sequence[float] |
             - ``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 @@ def __contains__(self, value: TensorType) -> bool | TensorType:
         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
--- 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,