54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
|
"""Order-independent named NumPy ``SeedSequence`` substreams."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
from typing import Iterable, Mapping
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .hashing import canonical_json_bytes
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_STREAMS = ("trajectory", "sensor", "model", "network")
|
||
|
|
|
||
|
|
|
||
|
|
def _label_words(*labels: object) -> list[int]:
|
||
|
|
digest = hashlib.sha256(canonical_json_bytes(labels)).digest()
|
||
|
|
return [
|
||
|
|
int.from_bytes(digest[offset : offset + 4], "little")
|
||
|
|
for offset in range(0, 16, 4)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def named_seed_record(
|
||
|
|
root_seed: int,
|
||
|
|
namespace: object,
|
||
|
|
stream_names: Iterable[str] = DEFAULT_STREAMS,
|
||
|
|
) -> dict[str, list[int]]:
|
||
|
|
"""Create stable named seed states independent of request order."""
|
||
|
|
root_seed = int(root_seed)
|
||
|
|
if root_seed < 0:
|
||
|
|
raise ValueError("root_seed must be non-negative")
|
||
|
|
record: dict[str, list[int]] = {}
|
||
|
|
for name in sorted(set(stream_names)):
|
||
|
|
if not name:
|
||
|
|
raise ValueError("stream names must be non-empty")
|
||
|
|
entropy = [root_seed & 0xFFFFFFFF, (root_seed >> 32) & 0xFFFFFFFF]
|
||
|
|
entropy.extend(_label_words(namespace, name))
|
||
|
|
state = np.random.SeedSequence(entropy).generate_state(4, dtype=np.uint32)
|
||
|
|
record[name] = [int(word) for word in state]
|
||
|
|
return record
|
||
|
|
|
||
|
|
|
||
|
|
def generator_from_record(
|
||
|
|
record: Mapping[str, list[int]],
|
||
|
|
stream_name: str,
|
||
|
|
) -> np.random.Generator:
|
||
|
|
if stream_name not in record:
|
||
|
|
raise KeyError(f"Unknown random stream {stream_name!r}")
|
||
|
|
words = [int(word) for word in record[stream_name]]
|
||
|
|
if len(words) != 4 or any(word < 0 or word > 0xFFFFFFFF for word in words):
|
||
|
|
raise ValueError(f"Invalid seed state for stream {stream_name!r}")
|
||
|
|
return np.random.default_rng(np.random.SeedSequence(words))
|