Numpy restoration of state using default_rng instead of RandomState

Viewed 21

This is how I am current restoring the state of a random generator incase my script gets stopped in the middle by accident. So in order to continue, the script I save the state of the random state and load it when I need it to generate samples from where it left off.

import os
import numpy as np
from numpy.random import RandomState

rndgen = RandomState(42)

i1 = rndgen.random_integers(low=0, high=1, size=10)
print(i1)

A = np.array(rndgen.get_state(), dtype=object)

np.save('saved_state.npy', A)

i2 = rndgen.random_integers(low=0, high=1, size=10)
print(i2)

ar = np.load('saved_state.npy', allow_pickle=True)
B = rndgen.set_state(tuple(ar))

print(rndgen.random_integers(low=0, high=1, size=10))

How do I implement this same functionality but with default_rng()? instead of RandomState? Is it even possible?

Output:

[0 1 0 0 0 1 0 0 0 1] # 1st generation of numbers

[0 0 0 0 1 0 1 1 1 0] # 2nd generation of numbers

[0 0 0 0 1 0 1 1 1 0] # Same 2nd sequence of numbers generated after loading the state

1 Answers

Just for reference, this is what I did based on André's comment:

import numpy as np
import pickle
from numpy.random import default_rng

rndgen = default_rng(42)
print(rndgen.integers(low=0, high=1, size=10, endpoint=True))

S = rndgen.bit_generator.state
with open('saved_state.pickle', 'wb') as handle:
    pickle.dump(S, handle, protocol=pickle.HIGHEST_PROTOCOL)

print(rndgen.integers(low=0, high=1, size=10, endpoint=True))

with open('saved_state.pickle', 'rb') as handle:
    rndgen.bit_generator.state = pickle.load(handle)

print(rndgen.integers(low=0, high=1, size=10, endpoint=True))

The output as required is:

[0 1 1 0 0 1 0 1 0 0]

[1 1 1 1 1 1 1 0 1 0]

[1 1 1 1 1 1 1 0 1 0]

Related