Skip to content

Faster deepcopy of RNG in RandomVariables #1708

@ricardoV94

Description

@ricardoV94

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)

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions