create state from seed

Author: katxiao

Diff for: ctgan/synthesizers/base.py

@@ -7,24 +7,34 @@
 
 
 @contextlib.contextmanager
-def random_seed(seed):
-    """Context manager for managing the random seed.
+def set_random_states(random_state, set_model_random_state):
+    """Context manager for managing the random state.
 
     Args:
-        seed (int):
-            The random seed.
+        random_state (int or tuple):
+            The random seed or a tuple of (numpy.random.RandomState, torch.Generator).
+        set_model_random_state (function):
+            Function to set the random state on the model.
     """
-    state = np.random.get_state()
-    torch_state = torch.get_rng_state()
+    original_np_state = np.random.get_state()
+    original_torch_state = torch.get_rng_state()
 
-    np.random.seed(seed)
-    torch.manual_seed(seed)
+    random_np_state, random_torch_state = random_state
+
+    np.random.set_state(random_np_state.get_state())
+    torch.set_rng_state(random_torch_state.get_state())
 
     try:
         yield
     finally:
-        np.random.set_state(state)
-        torch.set_rng_state(torch_state)
+        current_np_state = np.random.RandomState()
+        current_np_state.set_state(np.random.get_state())
+        current_torch_state = torch.Generator()
+        current_torch_state.set_state(torch.get_rng_state())
+        set_model_random_state((current_np_state, current_torch_state))
+
+        np.random.set_state(original_np_state)
+        torch.set_rng_state(original_torch_state)
 
 
 def random_state(function):
@@ -36,11 +46,11 @@ def random_state(function):
     """
 
     def wrapper(self, *args, **kwargs):
-        if self._random_seed is None:
+        if self.random_states is None:
             return function(self, *args, **kwargs)
 
         else:
-            with random_seed(self._random_seed):
+            with set_random_states(self.random_states, self.set_random_state):
                 return function(self, *args, **kwargs)
 
     return wrapper
@@ -52,7 +62,7 @@ class BaseSynthesizer:
     This should contain the save/load methods.
     """
 
-    _random_seed = None
+    random_states = None
 
     def save(self, path):
         """Save the model in the passed `path`."""
@@ -69,11 +79,26 @@ def load(cls, path):
         model.set_device(device)
         return model
 
-    def set_random_seed(self, random_seed):
-        """Set the random seed.
+    def set_random_state(self, random_state):
+        """Set the random state.
 
         Args:
-            random_seed (int):
-                Seed for the random generator.
+            random_state (tuple or int):
+                Either a tuple containing the (numpy.random.RandomState, torch.Generator)
+                or an int representing the random seed to use for both random states.
         """
-        self._random_seed = random_seed
+        if isinstance(random_state, int):
+            self.random_states = (
+                np.random.RandomState(seed=random_state),
+                torch.Generator().manual_seed(random_state),
+            )
+        elif (
+            isinstance(random_state, tuple) and
+            isinstance(random_state[0], np.random.RandomState) and
+            isinstance(random_state[1], torch.Generator)
+        ):
+            self.random_states = random_state
+        else:
+            raise ValueError(
+                '`random_state` {random_state} expected to be an int or a tuple of '
+                '(`np.random.RandomState`, `torch.Generator`)')

Diff for: tests/integration/synthesizer/test_ctgan.py

@@ -215,13 +215,20 @@ def test_fixed_random_seed():
 
     # Run
     sampled_random = ctgan.sample(10)
-    ctgan.set_random_seed(0)
-    sampled_0 = ctgan.sample(10)
-    sampled_1 = ctgan.sample(10)
+
+    ctgan.set_random_state(0)
+    sampled_0_0 = ctgan.sample(10)
+    sampled_0_1 = ctgan.sample(10)
+
+    ctgan.set_random_state(0)
+    sampled_1_0 = ctgan.sample(10)
+    sampled_1_1 = ctgan.sample(10)
 
     # Assert
-    assert not np.array_equal(sampled_random, sampled_0)
-    np.testing.assert_array_equal(sampled_0, sampled_1)
+    assert not np.array_equal(sampled_random, sampled_0_0)
+    assert not np.array_equal(sampled_random, sampled_0_1)
+    np.testing.assert_array_equal(sampled_0_0, sampled_1_0)
+    np.testing.assert_array_equal(sampled_0_1, sampled_1_1)
 
 
 # Below are CTGAN tests that should be implemented in the future

Diff for: tests/integration/synthesizer/test_tvae.py

@@ -115,10 +115,17 @@ def test_fixed_random_seed():
 
     # Run
     sampled_random = tvae.sample(10)
-    tvae.set_random_seed(0)
-    sampled_0 = tvae.sample(10)
-    sampled_1 = tvae.sample(10)
+
+    tvae.set_random_state(0)
+    sampled_0_0 = tvae.sample(10)
+    sampled_0_1 = tvae.sample(10)
+
+    tvae.set_random_state(0)
+    sampled_1_0 = tvae.sample(10)
+    sampled_1_1 = tvae.sample(10)
 
     # Assert
-    assert not np.array_equal(sampled_random, sampled_0)
-    np.testing.assert_array_equal(sampled_0, sampled_1)
+    assert not np.array_equal(sampled_random, sampled_0_0)
+    assert not np.array_equal(sampled_random, sampled_0_1)
+    np.testing.assert_array_equal(sampled_0_0, sampled_1_0)
+    np.testing.assert_array_equal(sampled_0_1, sampled_1_1)

Diff for: tests/unit/synthesizer/test_base.py

@@ -1,27 +1,37 @@
 
 """BaseSynthesizer unit testing module."""
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, call, patch
+
+import numpy as np
+import torch
 
 from ctgan.synthesizers.base import BaseSynthesizer, random_state
 
 
 @patch('ctgan.synthesizers.base.torch')
 @patch('ctgan.synthesizers.base.np.random')
-def test_valid_random_seed(random_mock, torch_mock):
-    """Test the ``random_seed`` attribute with a valid random seed.
+def test_valid_random_state(random_mock, torch_mock):
+    """Test the ``random_state`` attribute with a valid random state.
 
     Expect that the decorated function uses the random_state attribute.
     """
     # Setup
     my_function = MagicMock()
     instance = MagicMock()
-    instance._random_seed = 42
+
+    random_state_mock = MagicMock()
+    random_state_mock.get_state.return_value = 'desired numpy state'
+    torch_generator_mock = MagicMock()
+    torch_generator_mock.get_state.return_value = 'desired torch state'
+    instance.random_states = (random_state_mock, torch_generator_mock)
 
     args = {'some', 'args'}
     kwargs = {'keyword': 'value'}
 
+    random_mock.RandomState.return_value = random_state_mock
     random_mock.get_state.return_value = 'random state'
+    torch_mock.Generator.return_value = torch_generator_mock
     torch_mock.get_rng_state.return_value = 'torch random state'
 
     # Run
@@ -32,25 +42,27 @@ def test_valid_random_seed(random_mock, torch_mock):
     my_function.assert_called_once_with(instance, *args, **kwargs)
 
     instance.assert_not_called
-    random_mock.get_state.assert_called_once_with()
-    torch_mock.get_rng_state.assert_called_once_with()
-    random_mock.seed.assert_called_once_with(42)
-    random_mock.set_state.assert_called_once_with('random state')
-    torch_mock.set_rng_state.assert_called_once_with('torch random state')
+    assert random_mock.get_state.call_count == 2
+    assert torch_mock.get_rng_state.call_count == 2
+    random_mock.RandomState.assert_has_calls(
+        [call().get_state(), call(), call().set_state('random state')])
+    random_mock.set_state.assert_has_calls([call('desired numpy state'), call('random state')])
+    torch_mock.set_rng_state.assert_has_calls(
+        [call('desired torch state'), call('torch random state')])
 
 
 @patch('ctgan.synthesizers.base.torch')
 @patch('ctgan.synthesizers.base.np.random')
 def test_no_random_seed(random_mock, torch_mock):
-    """Test the ``random_seed`` attribute with no random seed.
+    """Test the ``random_state`` attribute with no random state.
 
     Expect that the decorated function calls the original function
     when there is no random state.
     """
     # Setup
     my_function = MagicMock()
     instance = MagicMock()
-    instance._random_seed = None
+    instance.random_states = None
 
     args = {'some', 'args'}
     kwargs = {'keyword': 'value'}
@@ -64,21 +76,24 @@ def test_no_random_seed(random_mock, torch_mock):
 
     instance.assert_not_called
     random_mock.get_state.assert_not_called()
-    random_mock.seed.assert_not_called()
+    random_mock.RandomState.assert_not_called()
     random_mock.set_state.assert_not_called()
     torch_mock.get_rng_state.assert_not_called()
+    torch_mock.Generator.assert_not_called()
     torch_mock.set_rng_state.assert_not_called()
 
 
 class TestBaseSynthesizer:
 
-    def test_set_random_seed(self):
-        """Test ``set_random_seed`` works as expected."""
+    def test_set_random_state(self):
+        """Test ``set_random_state`` works as expected."""
         # Setup
         instance = BaseSynthesizer()
 
         # Run
-        instance.set_random_seed(3)
+        instance.set_random_state(3)
 
         # Assert
-        assert instance._random_seed == 3
+        assert isinstance(instance.random_states, tuple)
+        assert isinstance(instance.random_states[0], np.random.RandomState)
+        assert isinstance(instance.random_states[1], torch.Generator)