def temp_seed(rng: np.random.RandomState,
seed: Optional[Union[int, Tuple[int, ...]]]):
"""A context manager for temporarily adjusting the random seed."""
if seed is None:
try:
yield
finally:
pass
else:
state = rng.get_state()
rng.seed(seed)
try:
yield
finally:
rng.set_state(state)
def __init__(
self,
indices: tp.List[int],
translation_factor: float = 1,
rotation: bool = False,
random_state: np.random.RandomState = None,
) -> None:
dim = len(indices)
assert dim
if random_state is None:
random_state = np.random.RandomState(0)
random_state.set_state(np.random.get_state())
self.indices = np.asarray(indices)
self.translation: np.ndarray = random_state.normal(
0, 1, dim) * translation_factor
self.rotation_matrix: tp.Optional[np.ndarray] = None
if rotation:
self.rotation_matrix = np.linalg.qr(
random_state.normal(0, 1, size=(dim, dim)))[0]
