Can I create a local numpy random seed?

Question

There is a function, foo, that uses the np.random functionality. I want to control the seed that foo uses, but without actually changing the function itself. How do I do this?

Essentially I want something like this:

bar() # should have normal seed
with np.random.seed(0): # Doesn't work
    foo()
bar() # should have normal seed

---

Solutions like this:

rng = random.Random(42)
number = rng.randint(10, 20)

doesn't work in this case, as I don't have access to the inner workings of foo (or am I missing something??).

Answer 1

You could keep the global random state in a temporary variable and reset it once your function is done:

import contextlib
import numpy as np

@contextlib.contextmanager
def temp_seed(seed):
    state = np.random.get_state()
    np.random.seed(seed)
    try:
        yield
    finally:
        np.random.set_state(state)

Demo:

>>> np.random.seed(0)
>>> np.random.randn(3)
array([1.76405235, 0.40015721, 0.97873798])
>>> np.random.randn(3)
array([ 2.2408932 ,  1.86755799, -0.97727788])

>>> np.random.seed(0)
>>> np.random.randn(3)
array([1.76405235, 0.40015721, 0.97873798])
>>> with temp_seed(5):
...     np.random.randn(3)                                                                                        
array([ 0.44122749, -0.33087015,  2.43077119])
>>> np.random.randn(3)
array([ 2.2408932 ,  1.86755799, -0.97727788])

Answer 2

I assume the idea is that calls to bar() should when given a starting seed always see the same sequence of random numbers; regardless of how many calls to foo()are inserted in-between.

We can do this by creating a random seed from the random state that we use to re-seed when the temporary seeded state is done. This can be wrapped in a context manager:

import numpy as np

class temporary_seed:
    def __init__(self, seed):
        self.seed = seed
        self.backup = None

    def __enter__(self):
        self.backup = np.random.randint(2**32-1, dtype=np.uint32)
        np.random.seed(self.seed)

    def __exit__(self, *_):
        np.random.seed(self.backup)

Let's try this with

def bar():
    print('bar:', np.random.randint(10))

def foo():
    print('foo:', np.random.randint(10))

np.random.seed(999)

bar()  # bar: 0
with temporary_seed(42):
    foo()  # foo: 6
    foo()  # foo: 3
bar()  # bar: 9

So we get bar-sequence [0, 9] and foo-sequence [6, 3].

We try again without re-seeding globally:

bar()  # bar: 1
with temporary_seed(42):
    foo()  # foo: 6
    foo()  # foo: 3
bar()  # bar: 2

New bar-sequence [1, 2] and same foo-sequence again [6, 3].

Once again with same global seed, but a different seed for foo:

np.random.seed(999)

bar()  # bar: 0
with temporary_seed(0):
    foo()  # foo: 5
bar()  # bar: 9

This time we get the first bar-sequence again [0, 9] and a different foo. Nice!

So where is the catch? By entering and leaving the temorary seed part we change the random state. We do so deterministically and the results are repeatable, but if we get a different sequence if we don't call enter temorary_seed:

np.random.seed(999)

bar()  # bar: 0
bar()  # bar: 5

bar-sequence [0, 5] instead of [0, 9]. If you can live with that limitation this approach should work.