import numpy as np from numpy.random import Generator, Philox class JumpablePRNG: """ Jumpable random number generator Support jumping to any position, then use it like a normal Generator """ def __init__(self, seed=42): self.seed = seed self.values_per_counter = 4 self._position = 0 self._bit_gen = Philox(seed=seed) self.rng = Generator(self._bit_gen) # Record the number of consumed random numbers self._consumed = 0 def jump_to(self, position): """Jump to the position-th random number""" if position == self._position: return self # Re-create the generator self._bit_gen = Philox(seed=self.seed) # Calculate the needed counter value counter_value = position // self.values_per_counter position_in_block = position % self.values_per_counter # advance to the correct counter if counter_value > 0: # Batch advance to avoid overflow max_advance = 2**30 # Safe value while counter_value > 0: step = min(counter_value, max_advance) self._bit_gen.advance(step) counter_value -= step # Create new Generator self.rng = Generator(self._bit_gen) # Consume the values before the position in the block for _ in range(position_in_block): self.rng.random() self._position = position self._consumed = position return self def get_position(self): """Get the current position""" return self._position # def test_jumpable_generator(): # """Test jumpable generator""" # seed = 42 # print("=== Step 1: Generate reference data ===") # # Generate reference list # ref_gen = Generator(Philox(seed=seed)) # list_a = [ref_gen.random() for _ in range(100)] # print(f"First 5 values: {list_a[:5]}") # print("\n=== Step 2: Test jumpable functionality ===") # # Create jumpable generator # jgen = JumpablePRNG(seed=seed) # # Test sequential access # sequential_values = [] # for i in range(10): # jgen.jump_to(i) # sequential_values.append(jgen.rng.random()) # print(f"Sequential access first 10 values: {sequential_values}") # print(f"Matches reference: {all(abs(a - b) < 1e-10 for a, b in zip(sequential_values, list_a[:10]))}") # print("\n=== Step 3: Test random access ===") # # Test random access # test_indices = [20, 1, 90, 0, 50, 99, 33, 77, 15, 88] # for k in test_indices: # jgen.jump_to(k) # value = jgen.rng.random() # ref_value = list_a[k] # match = abs(value - ref_value) < 1e-10 # print(f"Position {k:2d}: {value:.6f} vs {ref_value:.6f} - {'✓' if match else '✗'}") # assert match, f"Position {k} does not match" # print("\n=== Step 4: Use jumpable generator ===") # # Jump to position 10 and continue generating # jgen.jump_to(10) # # Generate the next 5 values # next_values = [jgen.rng.random() for _ in range(5)] # print(f"From position 10: {next_values}") # print(f"Reference values(10-14): {list_a[10:15]}") # print(f"Match: {all(abs(a - b) < 1e-10 for a, b in zip(next_values, list_a[10:15]))}") # print("\n=== Step 5: Test various random methods ===") # # Create new reference generator # ref_gen2 = Generator(Philox(seed=seed)) # # Skip the first 50 values # for _ in range(50): # ref_gen2.random() # # Generate various types of random numbers # ref_integers = [ref_gen2.integers(0, 100) for _ in range(5)] # ref_normal = ref_gen2.normal(0, 1, 5) # ref_choice = ref_gen2.choice(['A', 'B', 'C', 'D'], size=5, p=[0.1, 0.2, 0.3, 0.4]) # # Use jumpable generator # jgen.jump_to(50) # test_integers = [jgen.rng.integers(0, 100) for _ in range(5)] # test_normal = jgen.rng.normal(0, 1, 5) # test_choice = jgen.rng.choice(['A', 'B', 'C', 'D'], size=5, p=[0.1, 0.2, 0.3, 0.4]) # print(f"Integers: {test_integers} vs {ref_integers}") # print(f"Integers match: {all(a == b for a, b in zip(test_integers, ref_integers))}") # print(f"Normal distribution match: {np.allclose(test_normal, ref_normal)}") # print(f"Choice match: {all(a == b for a, b in zip(test_choice, ref_choice))}") # print("\n=== Step 6: Large position test ===") # # Test larger positions # large_positions = [1000, 10000, 100000] # Useful wrapper class class CheckpointablePRNG: """ A random number generator that supports checkpointing. Not all the random methods are supported. """ def __init__(self, seed=42): self.seed = seed self.generator = JumpablePRNG(seed=seed) self.position = 0 @property def rng(self): """Get the underlying NumPy Generator""" return self.generator.rng def checkpoint(self): """Create a checkpoint""" return self.position def restore(self, checkpoint_pos): """Restore to checkpoint""" self.generator.jump_to(checkpoint_pos) self.position = checkpoint_pos def random(self, size=None): """Generate random numbers and update position""" # self.generator.jump_to(self.position) result = self.generator.rng.random(size) if size is None: self.position += 1 else: self.position += int(np.prod(size)) return result def uniform(self, low, high, size=None): # self.generator.jump_to(self.position) result = self.generator.rng.uniform(low, high, size) if size is None: self.position += 1 else: self.position += int(np.prod(size)) return result def integers(self, low, high=None, size=None): raise NotImplementedError("integers method is not supported") def normal(self, loc=0.0, scale=1.0, size=None): raise NotImplementedError("normal method is not supported") def choice(self, a, size=None, replace=True, p=None): raise NotImplementedError("choice method is not supported") def shuffle(self, x: list) -> list: """In-place shuffle using Fisher-Yates algorithm with deterministic consumption""" n = len(x) # Fisher-Yates shuffle - each iteration consumes exactly 1 random number for i in range(n - 1, 0, -1): # Generate random float [0, 1) and map to integer [0, i] r = self.random() # Consumes exactly 1 random number j = int(r * (i + 1)) # Map to [0, i] range # Swap elements x[i], x[j] = x[j], x[i]