spb/metrika Public

1//2//  GPUBootstrapTests.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import Testing12import ZQGPU13import ZQParser14import ZQPlanner15import ZQStats1617/// GPU bootstrap cross-checks (CLAUDE.md §9): same seed ⇒ identical18/// resample indices across backends; estimates within the documented19/// float32 tolerance.20@Suite("GPU bootstrap", .enabled(if: ZQGPUBootstrap.isAvailable))21struct GPUBootstrapTests {2223    @Test("MLX Philox indices are bit-identical to the CPU reference")24    func philoxParity() {25        let seed: UInt64 = 4226        let n = 13727        let generator = Philox4x32(seed: seed)2829        for replicate in [0, 1, 7, 1000] {30            let cpu = ZQResampling.pairsBootstrapIndices(31                replicate: replicate, sampleSize: n, generator: generator32            )33            let gpu = ZQGPUBootstrap.gpuIndices(replicate: replicate, sampleSize: n, seed: seed)34            #expect(cpu == gpu, "replicate \(replicate)")35        }36    }3738    @Test("GPU coefficient draws match CPU within float32 tolerance")39    func drawParity() throws {40        // Small synthetic problem, well conditioned.41        let n = 20042        let generator = Philox4x32(seed: 7)43        let x = (0..<n).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }44        let y = (0..<n).map { i in45            2 + 0.5 * x[i] + (generator.uniform(at: UInt64(n + i)) - 0.5)46        }47        let reps = 504849        let gpuDraws = ZQGPUBootstrap.pairsBootstrapOLS(50            y: y, predictors: [("x", x)], replicates: reps, seed: 4251        )5253        let bootstrapGenerator = Philox4x32(seed: 42)54        for replicate in 0..<reps {55            let indices = ZQResampling.pairsBootstrapIndices(56                replicate: replicate, sampleSize: n, generator: bootstrapGenerator57            )58            let fit = try ZQOLS.fit(59                y: indices.map { y[$0] },60                predictors: [("x", indices.map { x[$0] })]61            )62            for (j, coefficient) in fit.coefficients.enumerated() {63                expectClose(64                    gpuDraws[replicate][j], coefficient.estimate,65                    rtol: 1e-4, "replicate \(replicate) b[\(coefficient.name)]"66                )67            }68        }69    }7071    @Test("GPU bootstrap is deterministic for a fixed seed")72    func determinism() {73        let n = 10074        let generator = Philox4x32(seed: 3)75        let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }76        let y = (0..<n).map { i in x[i] + generator.uniform(at: UInt64(n + i)) }7778        let first = ZQGPUBootstrap.pairsBootstrapOLS(79            y: y, predictors: [("x", x)], replicates: 20, seed: 9980        )81        let second = ZQGPUBootstrap.pairsBootstrapOLS(82            y: y, predictors: [("x", x)], replicates: 20, seed: 9983        )84        #expect(first == second)85    }8687    @Test("chunking does not change results")88    func chunkingInvariance() {89        let n = 8090        let generator = Philox4x32(seed: 5)91        let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }92        let y = (0..<n).map { i in 2 * x[i] + generator.uniform(at: UInt64(n + i)) }9394        let oneChunk = ZQGPUBootstrap.pairsBootstrapOLS(95            y: y, predictors: [("x", x)], replicates: 30, seed: 1196        )97        // Budget so small every chunk holds a single replicate.98        let manyChunks = ZQGPUBootstrap.pairsBootstrapOLS(99            y: y, predictors: [("x", x)], replicates: 30, seed: 11,100            memoryBudgetBytes: 1101        )102        #expect(oneChunk == manyChunks)103    }104105    @Test("planner sends large-reps bootstrap to the GPU")106    func plannerDispatch() throws {107        let planner = ZQPlanner(gpuAvailable: true)108        let parser = ZQCommandParser()109110        let large = try #require(try parser.parse("bootstrap, reps(10000): reg y x"))111        #expect(planner.plan(large, rowCount: 1000).backend == .gpu)112113        let small = try #require(try parser.parse("bootstrap, reps(100): reg y x"))114        #expect(planner.plan(small, rowCount: 1000).backend == .cpu)115116        let noGPU = ZQPlanner(gpuAvailable: false)117        #expect(noGPU.plan(large, rowCount: 1000).backend == .cpu)118    }119}120