python/dpt/kernels/random.py

1"""Counter-addressed Philox4x32-10 shared by transport and observations.23The counter is (identity low, identity high, event, domain); the key is the464-bit seed. All coordinates are explicit: scheduling and rejection in another5history cannot change this history's stream. Domain 0 belongs to transport.6Domains 1..0x7fffffff belong to observations; transport rejection streams use7domains >= 0x80000000. Counter8coordinates must never wrap; the host boundary rejects exhausted identities.910Constants and round schedule: Random123, Salmon et al. (SC11), Philox4x32-10,11https://github.com/DEShawResearch/random123/blob/main/include/Random123/philox.h.12This is the published algorithm expressed in Warp, not a bundled copy of a13third-party implementation. These functions return no persistent RNG state.14"""1516# Warp annotations are executable DSL expressions; host interfaces remain strict.17# The optional GPU import is resolved only when an operator is prepared.18# pyright: reportInvalidTypeForm=false, reportUnknownParameterType=false19# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false20# pyright: reportUnknownVariableType=false, reportUntypedFunctionDecorator=false21# pyright: reportMissingImports=false, reportUntypedClassDecorator=false2223import warp as wp242526# region book:counter-random-stream27@wp.func28def random4(29    seed: wp.uint64,30    identity: wp.uint64,31    event: wp.uint32,32    domain: wp.uint32,33) -> wp.vec4d:34    """Return four open-interval uniforms; each has 32 random mantissa bits."""35    c0 = wp.uint32(identity & wp.uint64(0xFFFFFFFF))36    c1 = wp.uint32(identity >> wp.uint64(32))37    c2 = event38    c3 = domain39    k0 = wp.uint32(seed & wp.uint64(0xFFFFFFFF))40    k1 = wp.uint32(seed >> wp.uint64(32))41    for _ in range(10):42        product0 = wp.uint64(0xD2511F53) * wp.uint64(c0)43        product1 = wp.uint64(0xCD9E8D57) * wp.uint64(c2)44        low0 = wp.uint32(product0 & wp.uint64(0xFFFFFFFF))45        low1 = wp.uint32(product1 & wp.uint64(0xFFFFFFFF))46        high0 = wp.uint32(product0 >> wp.uint64(32))47        high1 = wp.uint32(product1 >> wp.uint64(32))48        c0 = high1 ^ c1 ^ k049        c1 = low150        c2 = high0 ^ c3 ^ k151        c3 = low052        k0 = k0 + wp.uint32(0x9E3779B9)53        k1 = k1 + wp.uint32(0xBB67AE85)54    scale = wp.float64(1.0 / 4294967296.0)55    return wp.vec4d(56        (wp.float64(c0) + wp.float64(0.5)) * scale,57        (wp.float64(c1) + wp.float64(0.5)) * scale,58        (wp.float64(c2) + wp.float64(0.5)) * scale,59        (wp.float64(c3) + wp.float64(0.5)) * scale,60    )616263# endregion book:counter-random-stream64