1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
|
#------------------------------------------------------------------------------
# Copyright (c) 2005, Enthought, Inc.
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in enthought/LICENSE.txt and may be redistributed only
# under the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
# Thanks for using Enthought open source!
#
# Author: Enthought, Inc.
# Description: <Enthought util package component>
#------------------------------------------------------------------------------
""" The two classes here help keep up with the current "state" of the
random number generators in scipy.stats and random. They can be
used to save the current state of the random number generators to
be used in the future to ensure identical results between calls.
The RandomStateManager works as a stack with save_state, set_state
and restore_state methods that allow you to keep a stack of states
available. A common usage is as follows:
>>> # save the current state in a variable named "old_state"
>>> old_state = RandomState()
Perform some stochastic calculation...
>>> # Now if you want to have stochastic_calculation() return identical
>>> # results without disrupting other calculations that use random values
>>> # do the following:
>>> rsm = RandomStateManager()
>>> rsm.save_state()
>>> rsm.set_state(old_state)
Perform some stochastic calculation...
>>> rsm.restore_state()
Note that these routines currently only support the state of random and
scipy.stats. If you use other random number generators, their states
will not be managed correctly.
"""
import random
from numpy.random import get_state as get_seed
from numpy.random import set_state as set_seed
class RandomState:
def __init__(self):
self.update()
def update(self):
self.stats_seed = get_seed()
self.random_state = random.getstate()
class RandomStateManager:
# todo - does it make any sense to use a stack structure?
# we currently store the seeds elsewhere anyway so the stack only
# ever has one element in it.
#
def __init__(self):
self.state_stack = []
def save_state(self):
current_state = RandomState()
self.state_stack.append(current_state)
def set_state(self, random_state):
seed = random_state.stats_seed
set_seed(seed)
state = random_state.random_state
random.setstate(state)
def restore_state(self):
try:
previous_state = self.state_stack.pop(-1)
self.set_state(previous_state)
except:
raise IndexError("trying to call restore_state without matching"
" call to save_state")
if __name__ == '__main__':
import doctest
doctest.testmod()
|
