spb/metrika Public

1//2//  Philox.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910/// Philox4x32-10 counter-based random number generator (Salmon et al.,11/// SC'11). Every 128-bit counter maps to four independent 32-bit outputs,12/// so draws are addressable by index — the property that makes CPU and13/// GPU streams identical for the same seed (CLAUDE.md §5: reproducibility14/// is a release blocker). This pure-Swift implementation is the reference15/// the MLX kernel path must match bit-for-bit.16public struct Philox4x32: Sendable {17    private static let multiplier0: UInt32 = 0xD251_1F5318    private static let multiplier1: UInt32 = 0xCD9E_8D5719    private static let weyl0: UInt32 = 0x9E37_79B920    private static let weyl1: UInt32 = 0xBB67_AE8521    private static let rounds = 102223    public let key: (UInt32, UInt32)2425    /// `set seed` maps the 64-bit seed directly onto the Philox key.26    public init(seed: UInt64) {27        self.key = (UInt32(truncatingIfNeeded: seed),28                    UInt32(truncatingIfNeeded: seed >> 32))29    }3031    /// One Philox block: 128-bit counter → four 32-bit words.32    public func block(counter: (UInt32, UInt32, UInt32, UInt32))33        -> (UInt32, UInt32, UInt32, UInt32) {34        var c = counter35        var k = key36        for _ in 0..<Self.rounds {37            let product0 = UInt64(Self.multiplier0) * UInt64(c.0)38            let product1 = UInt64(Self.multiplier1) * UInt64(c.2)39            let high0 = UInt32(truncatingIfNeeded: product0 >> 32)40            let low0 = UInt32(truncatingIfNeeded: product0)41            let high1 = UInt32(truncatingIfNeeded: product1 >> 32)42            let low1 = UInt32(truncatingIfNeeded: product1)43            c = (high1 ^ c.1 ^ k.0, low1, high0 ^ c.3 ^ k.1, low0)44            k.0 = k.0 &+ Self.weyl045            k.1 = k.1 &+ Self.weyl146        }47        return c48    }4950    /// The i-th 32-bit word of the stream (block = i/4, lane = i%4).51    public func word(at index: UInt64) -> UInt32 {52        let blockIndex = index / 453        let lane = Int(index % 4)54        let output = block(counter: (55            UInt32(truncatingIfNeeded: blockIndex),56            UInt32(truncatingIfNeeded: blockIndex >> 32),57            0, 058        ))59        switch lane {60        case 0: return output.061        case 1: return output.162        case 2: return output.263        default: return output.364        }65    }6667    /// Uniform Double in [0, 1) from the i-th draw, using 53 bits built68    /// from two consecutive 32-bit words.69    public func uniform(at index: UInt64) -> Double {70        let high = UInt64(word(at: index &* 2))71        let low = UInt64(word(at: index &* 2 &+ 1))72        let bits53 = ((high << 32) | low) >> 1173        return Double(bits53) * (1.0 / 9_007_199_254_740_992.0)  // 2^-5374    }7576    /// Uniform integer in 0..<bound from the i-th draw (used for77    /// bootstrap resampling indices). Bounded by 64-bit multiply-shift;78    /// the tiny modulo bias (< 2⁻³² for realistic n) is acceptable and,79    /// critically, is computed identically on CPU and GPU.80    public func integer(at index: UInt64, bound: Int) -> Int {81        precondition(bound > 0)82        let word = UInt64(self.word(at: index))83        return Int((word &* UInt64(bound)) >> 32)84    }85}8687/// Bootstrap resampling indices: replicate r draws its n indices from88/// dedicated counter positions, so any subset of replicates can be89/// regenerated independently — on either backend.90public enum ZQResampling {91    public static func pairsBootstrapIndices(92        replicate: Int, sampleSize: Int, generator: Philox4x3293    ) -> [Int] {94        let base = UInt64(replicate) &* UInt64(sampleSize)95        var indices = [Int](repeating: 0, count: sampleSize)96        for i in 0..<sampleSize {97            indices[i] = generator.integer(at: base &+ UInt64(i), bound: sampleSize)98        }99        return indices100    }101102    /// A random permutation of 0..<n for one replicate, defined as the103    /// argsort of n 64-bit Philox keys. Unlike Fisher–Yates this is104    /// order-free — each key is counter-addressable — so a future GPU105    /// argsort produces the identical permutation. Key collisions occur106    /// with probability ~n²·2⁻⁶⁴ and would only perturb tie order.107    ///108    /// Counter space: permutations use a dedicated stream offset (bit 63)109    /// so they never collide with bootstrap draws under the same seed.110    public static func permutationIndices(111        replicate: Int, sampleSize: Int, generator: Philox4x32112    ) -> [Int] {113        let streamOffset = UInt64(1) << 63114        let base = UInt64(replicate) &* UInt64(sampleSize)115        var keys = [UInt64](repeating: 0, count: sampleSize)116        for i in 0..<sampleSize {117            // Word positions: offset applied after the ×2 expansion so it118            // survives (offset·2 would wrap to 0 mod 2⁶⁴).119            let wordIndex = streamOffset &+ (base &+ UInt64(i)) &* 2120            let high = UInt64(generator.word(at: wordIndex))121            let low = UInt64(generator.word(at: wordIndex &+ 1))122            keys[i] = (high << 32) | low123        }124        return (0..<sampleSize).sorted { keys[$0] < keys[$1] }125    }126}127