-
Notifications
You must be signed in to change notification settings - Fork 148
Open
Description
Description
This is low priority, but when RVs are not inplacing on the RNG, we have to deepcopy it before using it.
This can be done faster, and specialized for the most common generator type: PCG64
from copy import deepcopy
from numpy.random import default_rng, PCG64, Generator
def custom_deepcopy(rng):
orig_bitgen = rng.bit_generator
if isinstance(orig_bitgen, PCG64):
# `seed=0` is faster than `None`, and we don't care because we override the state next
new_bitgen = PCG64(0)
new_bitgen.state = new_state = orig_bitgen.state.copy()
# Copy the one inner dict manually, all other values are immutable
new_state["state"] = new_state["state"].copy()
else:
new_bitgen = type(orig_bitgen)(0)
new_bitgen.state = deepcopy(orig_bitgen.state)
return Generator(new_bitgen)
rng = default_rng()
assert deepcopy(rng).bit_generator.state == custom_deepcopy(rng).bit_generator.state
%timeit deepcopy(rng) # 66.6 μs ± 1.96 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit custom_deepcopy(rng) # 22.1 μs ± 335 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)