spb/metrika Public

1//2//  BayesTests.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import Testing12import ZQEngine13import ZQGPU14import ZQStats1516@Suite("Bayesian regression", .serialized)17struct BayesTests {1819    @Test("Philox stream normals and gammas match their theoretical moments")20    func variateGenerators() {21        var stream = PhiloxStream(seed: 42)22        let n = 200_0002324        var normalSum = 0.0, normalSquares = 0.025        for _ in 0..<n {26            let z = stream.nextNormal()27            normalSum += z28            normalSquares += z * z29        }30        #expect(abs(normalSum / Double(n)) < 0.01)31        #expect(abs(normalSquares / Double(n) - 1) < 0.02)3233        // Gamma(3, rate 2): mean 1.5, variance 0.75.34        var gammaSum = 0.0, gammaSquares = 0.035        for _ in 0..<n {36            let g = stream.nextGamma(shape: 3, rate: 2)37            gammaSum += g38            gammaSquares += g * g39        }40        let gammaMean = gammaSum / Double(n)41        let gammaVariance = gammaSquares / Double(n) - gammaMean * gammaMean42        #expect(abs(gammaMean - 1.5) < 0.01, "gamma mean \(gammaMean)")43        #expect(abs(gammaVariance - 0.75) < 0.02, "gamma variance \(gammaVariance)")4445        // Shape < 1 boost path: Gamma(0.5, rate 1): mean 0.5.46        var smallSum = 0.047        for _ in 0..<n { smallSum += stream.nextGamma(shape: 0.5, rate: 1) }48        #expect(abs(smallSum / Double(n) - 0.5) < 0.01)49    }5051    @Test("flat priors recover OLS: posterior mean ≈ b̂, sd ≈ SE")52    func flatPriorAgreement() throws {53        // Synthetic data with a known DGP.54        var stream = PhiloxStream(seed: 7)55        let n = 30056        let x = (0..<n).map { _ in 5 + 10 * stream.nextUniform() }57        let y = x.map { 2 + 0.5 * $0 + 0.8 * stream.nextNormal() }5859        let ols = try ZQOLS.fit(y: y, predictors: [("x", x)])60        let bayes = try ZQBayesianRegression.fitGibbs(61            y: y, predictors: [("x", x)],62            mcmcSize: 40_000, burnIn: 4_000, seed: 42,63            coefficientPriorVariance: 1e864        )6566        for (posterior, frequentist) in zip(bayes.coefficients, ols.coefficients) {67            // Monte-Carlo error with 40k (autocorrelated) draws: assert68            // within 5% of a posterior SD.69            let toleranceMean = 0.05 * frequentist.standardError70            #expect(71                abs(posterior.posteriorMean - frequentist.estimate) < toleranceMean,72                "mean[\(posterior.name)]: \(posterior.posteriorMean) vs \(frequentist.estimate)"73            )74            let sdRatio = posterior.posteriorSD / frequentist.standardError75            #expect(76                sdRatio > 0.9 && sdRatio < 1.15,77                "sd[\(posterior.name)] ratio \(sdRatio)"78            )79            // The 95% credible interval brackets the OLS estimate.80            #expect(posterior.credibleLower < frequentist.estimate)81            #expect(posterior.credibleUpper > frequentist.estimate)82        }83        // σ posterior around the DGP value 0.8.84        #expect(abs(bayes.sigma.posteriorMean - 0.8) < 0.1)85    }8687    @Test("a tight prior shrinks coefficients toward zero")88    func priorShrinkage() throws {89        var stream = PhiloxStream(seed: 9)90        let n = 6091        let x = (0..<n).map { _ in stream.nextUniform() * 10 }92        let y = x.map { 3 * $0 + stream.nextNormal() }9394        let flat = try ZQBayesianRegression.fitGibbs(95            y: y, predictors: [("x", x)],96            mcmcSize: 5_000, burnIn: 1_000, seed: 42,97            coefficientPriorVariance: 1e698        )99        let tight = try ZQBayesianRegression.fitGibbs(100            y: y, predictors: [("x", x)],101            mcmcSize: 5_000, burnIn: 1_000, seed: 42,102            coefficientPriorVariance: 0.01103        )104        #expect(105            abs(tight.coefficients[0].posteriorMean)106                < abs(flat.coefficients[0].posteriorMean)107        )108        #expect(abs(flat.coefficients[0].posteriorMean - 3) < 0.2)109    }110111    @Test("chains are reproducible for a fixed seed")112    func determinism() throws {113        var stream = PhiloxStream(seed: 3)114        let x = (0..<50).map { _ in stream.nextUniform() }115        let y = x.map { $0 + 0.1 * stream.nextNormal() }116117        let a = try ZQBayesianRegression.fitGibbs(118            y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 42119        )120        let b = try ZQBayesianRegression.fitGibbs(121            y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 42122        )123        let c = try ZQBayesianRegression.fitGibbs(124            y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 43125        )126        #expect(a == b)127        #expect(a != c)128    }129130    @Test("bayes prefix through the console")131    func consoleCommand() async throws {132        let fixtures = try Fixtures()133        let session = try ZQSession(discoverUserCommands: false)134        _ = try await session.execute("use \(fixtures.datasetURL.path)")135        _ = try await session.execute("gen log_rev = ln(revenue)")136137        let ols = try await session.execute("reg log_rev price")138        let bayes = try await session.execute(139            "bayes, mcmcsize(20000) burnin(2000) seed(42) normalprior(100000000): reg log_rev price"140        )141        let posterior = try #require(bayes.scalars["b_price"])142        let frequentist = try #require(ols.scalars["b_price"])143        let se = try #require(ols.scalars["se_price"])144        #expect(abs(posterior - frequentist) < 0.05 * se)145        #expect(bayes.text.contains("Bayesian linear regression"))146        #expect(bayes.scalars["b_sigma"] != nil)147    }148}149