Repository: secondmind-labs/trieste

Issue:
Seed argument on space's random sample method
**Describe the feature you'd like**
I would like to have set seed functionality in `Box` and `DiscreteSearchSpace`'s `.sample()` method 

**Is your feature request related to a problem? Please describe.**
This will help reproducibility for a research purposes

**Describe alternatives you've considered**
Simply add seed argument to the `.sample()` method  and pass it to the `tf.random` function that generate the random points

Discussion:
satrialoka: I add `seed: Optional[int]` argument in `sample` method inside `SearchSpace` abstract base class. This assume that we always want seed option in `sample`. Since the other sample (non uniform)  method that possibly doesn't have seed arg, can have different name. for example `sample_sobol`, `sample_lhd` etc
hstojic: ready to merge?
satrialoka: Yes please!

Patch:
diff --git a/trieste/space.py b/trieste/space.py
index 62e761d6..bd442fa9 100644
--- a/trieste/space.py
+++ b/trieste/space.py
@@ -35,9 +35,10 @@ class SearchSpace(ABC):
     """
 
     @abstractmethod
-    def sample(self, num_samples: int) -> TensorType:
+    def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
         """
         :param num_samples: The number of points to sample from this search space.
+        :param seed: Random seed for reproducibility.
         :return: ``num_samples`` i.i.d. random points, sampled uniformly from this search space.
         """
 
@@ -176,17 +177,21 @@ class DiscreteSearchSpace(SearchSpace):
         tf.debugging.assert_shapes([(value, self.points.shape[1:])])
         return tf.reduce_any(tf.reduce_all(value == self._points, axis=1))
 
-    def sample(self, num_samples: int) -> TensorType:
+    def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
         """
         :param num_samples: The number of points to sample from this search space.
+        :param seed: Random seed for reproducibility.
         :return: ``num_samples`` i.i.d. random points, sampled uniformly,
             from this search space.
         """
+        if seed is not None:  # ensure reproducibility
+            tf.random.set_seed(seed)
+
         if num_samples == 0:
             return self.points[:0, :]
         else:
             sampled_indices = tf.random.categorical(
-                tf.ones((1, tf.shape(self.points)[0])), num_samples
+                tf.ones((1, tf.shape(self.points)[0])), num_samples, seed=seed
             )
             return tf.gather(self.points, sampled_indices)[0, :, :]  # [num_samples, D]
 
@@ -324,20 +329,26 @@ class Box(SearchSpace):
 
         return tf.reduce_all(value >= self._lower) and tf.reduce_all(value <= self._upper)
 
-    def sample(self, num_samples: int) -> TensorType:
+    def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
         """
         Sample randomly from the space.
 
         :param num_samples: The number of points to sample from this search space.
+        :param seed: Random seed for reproducibility.
         :return: ``num_samples`` i.i.d. random points, sampled uniformly,
             from this search space with shape '[num_samples, D]' , where D is the search space
             dimension.
         """
         tf.debugging.assert_non_negative(num_samples)
-
+        if seed is not None:  # ensure reproducibility
+            tf.random.set_seed(seed)
         dim = tf.shape(self._lower)[-1]
         return tf.random.uniform(
-            (num_samples, dim), minval=self._lower, maxval=self._upper, dtype=self._lower.dtype
+            (num_samples, dim),
+            minval=self._lower,
+            maxval=self._upper,
+            dtype=self._lower.dtype,
+            seed=seed,
         )
 
     def sample_halton(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
@@ -583,7 +594,7 @@ class TaggedProductSearchSpace(SearchSpace):
         ]
         return tf.reduce_all(in_each_subspace)
 
-    def sample(self, num_samples: int) -> TensorType:
+    def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType:
         """
         Sample randomly from the space by sampling from each subspace
         and concatenating the resulting samples.
@@ -594,8 +605,9 @@ class TaggedProductSearchSpace(SearchSpace):
             dimension.
         """
         tf.debugging.assert_non_negative(num_samples)
-
-        subspace_samples = [self._spaces[tag].sample(num_samples) for tag in self._tags]
+        if seed is not None:  # ensure reproducibility
+            tf.random.set_seed(seed)
+        subspace_samples = [self._spaces[tag].sample(num_samples, seed=seed) for tag in self._tags]
         return tf.concat(subspace_samples, -1)
 
     def product(self, other: TaggedProductSearchSpace) -> TaggedProductSearchSpace:


Test patch:
diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py
index 351cadd2..537ad017 100644
--- a/tests/unit/test_space.py
+++ b/tests/unit/test_space.py
@@ -19,6 +19,7 @@ import operator
 from collections import Sequence
 from collections.abc import Container
 from functools import reduce
+from typing import Optional
 
 import numpy.testing as npt
 import pytest
@@ -43,8 +44,8 @@ class Integers(SearchSpace):
     def upper(self) -> None:
         pass
 
-    def sample(self, num_samples: int) -> tf.Tensor:
-        return tf.random.shuffle(tf.range(self.limit))[:num_samples]
+    def sample(self, num_samples: int, seed: Optional[int] = None) -> tf.Tensor:
+        return tf.random.shuffle(tf.range(self.limit), seed=seed)[:num_samples]
 
     def __contains__(self, point: tf.Tensor) -> bool | TensorType:
         tf.debugging.assert_integer(point)
@@ -167,6 +168,21 @@ def test_discrete_search_space_sampling(num_samples: int) -> None:
     assert len(samples) == num_samples
 
 
+@pytest.mark.parametrize("seed", [1, 42, 123])
+def test_discrete_search_space_sampling_returns_same_points_for_same_seed(seed: int) -> None:
+    search_space = DiscreteSearchSpace(_points_in_2D_search_space())
+    random_samples_1 = search_space.sample(num_samples=100, seed=seed)
+    random_samples_2 = search_space.sample(num_samples=100, seed=seed)
+    npt.assert_allclose(random_samples_1, random_samples_2)
+
+
+def test_discrete_search_space_sampling_returns_different_points_for_different_call() -> None:
+    search_space = DiscreteSearchSpace(_points_in_2D_search_space())
+    random_samples_1 = search_space.sample(num_samples=100)
+    random_samples_2 = search_space.sample(num_samples=100)
+    npt.assert_raises(AssertionError, npt.assert_allclose, random_samples_1, random_samples_2)
+
+
 def test_discrete_search_space___mul___points_is_the_concatenation_of_original_points() -> None:
     dss1 = DiscreteSearchSpace(tf.constant([[-1.0, -1.4], [-1.5, -3.6], [-0.5, -0.6]]))
     dss2 = DiscreteSearchSpace(tf.constant([[1.0, 1.4], [1.5, 3.6]]))
@@ -416,6 +432,14 @@ def test_box_halton_sampling_raises_for_invalid_sample_size(num_samples: int) ->
         box.sample_halton(num_samples)
 
 
+@pytest.mark.parametrize("seed", [1, 42, 123])
+def test_box_sampling_returns_same_points_for_same_seed(seed: int) -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    random_samples_1 = box.sample(num_samples=100, seed=seed)
+    random_samples_2 = box.sample(num_samples=100, seed=seed)
+    npt.assert_allclose(random_samples_1, random_samples_2)
+
+
 @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,)))
@@ -432,6 +456,13 @@ def test_box_halton_sampling_returns_same_points_for_same_seed(seed: int) -> Non
     npt.assert_allclose(halton_samples_1, halton_samples_2)
 
 
+def test_box_sampling_returns_different_points_for_different_call() -> None:
+    box = Box(tf.zeros((3,)), tf.ones((3,)))
+    random_samples_1 = box.sample(num_samples=100)
+    random_samples_2 = box.sample(num_samples=100)
+    npt.assert_raises(AssertionError, npt.assert_allclose, random_samples_1, random_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)
@@ -439,7 +470,7 @@ def test_box_sobol_sampling_returns_different_points_for_different_call() -> Non
     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:
+def test_box_halton_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)
@@ -796,6 +827,25 @@ def test_product_space_discretize_returns_search_space_with_correct_number_of_po
     assert len(samples) == num_samples
 
 
+@pytest.mark.parametrize("seed", [1, 42, 123])
+def test_product_space_sampling_returns_same_points_for_same_seed(seed: int) -> None:
+    space_A = Box([-1], [2])
+    space_B = DiscreteSearchSpace(tf.random.uniform([100, 2], dtype=tf.float64, seed=42))
+    product_space = TaggedProductSearchSpace(spaces=[space_A, space_B])
+    random_samples_1 = product_space.sample(num_samples=100, seed=seed)
+    random_samples_2 = product_space.sample(num_samples=100, seed=seed)
+    npt.assert_allclose(random_samples_1, random_samples_2)
+
+
+def test_product_space_sampling_returns_different_points_for_different_call() -> None:
+    space_A = Box([-1], [2])
+    space_B = DiscreteSearchSpace(tf.random.uniform([100, 2], dtype=tf.float64, seed=42))
+    product_space = TaggedProductSearchSpace(spaces=[space_A, space_B])
+    random_samples_1 = product_space.sample(num_samples=100)
+    random_samples_2 = product_space.sample(num_samples=100)
+    npt.assert_raises(AssertionError, npt.assert_allclose, random_samples_1, random_samples_2)
+
+
 def test_product_space___mul___() -> None:
     space_A = Box([-1], [2])
     space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64))