SPB Git
spboucher.ai

Repository Public

spbmetrika

Commit

feat(stats): Bayesian linear regression via Gibbs sampling (bayes prefix)

- PhiloxStream in ZQGPU: sequential variate stream over a dedicated
  counter block (bit 62 + per-stream 2^48 words, disjoint from bootstrap
  and permutation streams); Box-Muller normals, Marsaglia-Tsang gammas
  (moments verified against theory at 200k draws)
- ZQBayesianRegression: semi-conjugate Gibbs with Stata bayes default
  priors (coefficients N(0, 10000), variance InvGamma(.01, .01)); exact
  full conditionals via dense Cholesky; posterior mean/sd and
  equal-tailed 95% credible intervals; ZQStats now depends on ZQGPU for
  the shared RNG
- engine: 'bayes [, mcmcsize() burnin() seed() normalprior()]: reg …'
  with reproducible chains; manual entry included
- validation: with diffuse priors the posterior reproduces OLS (mean
  within 5% of a posterior SD, sd ratio in [0.9, 1.15], CrI brackets the
  estimate, sigma recovers the DGP); tight priors shrink toward zero;
  chains bit-reproducible per seed
- 111 tests green (swift test and xcodebuild with GPU suites)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed · Aug 5, 2026 parent a3f7d61
8 changed files +561 −3 Browse files
8 changed files
Unified Split
modified MetrikaKit/Package.swift +1 −1
@@ -48,7 +48,7 @@ let package = Package(
48 ), 48 ),
49 .target( 49 .target(
50 name: "ZQStats", 50 name: "ZQStats",
51− dependencies: ["ZQData"], 51+ dependencies: ["ZQData", "ZQGPU"],
52 swiftSettings: strictConcurrency 52 swiftSettings: strictConcurrency
53 ), 53 ),
54 .target( 54 .target(
modified MetrikaKit/Sources/ZQEngine/CommandHelp.swift +13 −0
@@ -230,6 +230,19 @@ public enum ZQCommandReference {
230 examples: ["bootstrap, reps(10000) seed(42): reg log_rev price"], 230 examples: ["bootstrap, reps(10000) seed(42): reg log_rev price"],
231 notes: "Replicates are counter-addressable: any subset recomputes identically regardless of chunking or backend. The planner dispatches ≥500 reps to the GPU when available." 231 notes: "Replicates are counter-addressable: any subset recomputes identically regardless of chunking or backend. The planner dispatches ≥500 reps to the GPU when available."
232 ), 232 ),
233+ ZQCommandDoc(
234+ verb: "bayes", category: "Resampling & simulation",
235+ summary: "Bayesian linear regression by Gibbs sampling: posterior means, standard deviations, and 95% credible intervals.",
236+ syntax: "bayes [, mcmcsize(#) burnin(#) seed(#) normalprior(#)]: regress depvar indepvars",
237+ options: [
238+ ("mcmcsize(#)", "posterior draws after burn-in (default 10000)"),
239+ ("burnin(#)", "discarded warm-up iterations (default 2500)"),
240+ ("seed(#)", "Philox seed — chains are exactly reproducible"),
241+ ("normalprior(#)", "prior variance of the N(0, #) coefficient priors (default 10000)"),
242+ ],
243+ examples: ["bayes, mcmcsize(20000) seed(42): reg log_rev price"],
244+ notes: "Priors: coefficients N(0, normalprior), variance InvGamma(0.01, 0.01) — Stata's bayes defaults. With diffuse priors the posterior reproduces OLS."
245+ ),
233 ZQCommandDoc( 246 ZQCommandDoc(
234 verb: "permute", category: "Resampling & simulation", 247 verb: "permute", category: "Resampling & simulation",
235 summary: "Permutation test: the response is permuted, the model refit, and empirical two-sided p-values reported per coefficient.", 248 summary: "Permutation test: the response is permuted, the model refit, and empirical two-sided p-values reported per coefficient.",
modified MetrikaKit/Sources/ZQEngine/Session.swift +96 −0
@@ -208,6 +208,7 @@ public actor ZQSession {
208 case "display": return try handleDisplay(command) 208 case "display": return try handleDisplay(command)
209 case "bootstrap": return try await handleBootstrap(command, backend: backend) 209 case "bootstrap": return try await handleBootstrap(command, backend: backend)
210 case "permute": return try await handlePermute(command) 210 case "permute": return try await handlePermute(command)
211+ case "bayes": return try handleBayes(command)
211 case "graph": return try handleGraph(command) 212 case "graph": return try handleGraph(command)
212 case "histogram", "scatter", "kdensity": 213 case "histogram", "scatter", "kdensity":
213 var promoted = command 214 var promoted = command
@@ -2046,6 +2047,101 @@ public actor ZQSession {
2046 return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars) 2047 return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
2047 } 2048 }
2048 2049
2050+ // MARK: - Bayesian estimation
2051+
2052+ /// `bayes [, mcmcsize(#) burnin(#) seed(#) normalprior(#)]: reg y x…`
2053+ /// — Gibbs-sampled Bayesian linear regression with Stata-style default
2054+ /// priors: N(0, 10000) on coefficients, InvGamma(0.01, 0.01) on σ².
2055+ private func handleBayes(_ command: ZQCommand) throws -> ZQResult {
2056+ guard let body = command.body else {
2057+ throw ZQEngineError("bayes: syntax is 'bayes [, options]: regress …'")
2058+ }
2059+ guard body.verb == "regress" else {
2060+ throw ZQEngineError("bayes currently supports 'regress' bodies only")
2061+ }
2062+ if let seedText = command.option("seed")?.firstArgument {
2063+ guard let value = UInt64(seedText) else {
2064+ throw ZQEngineError("bayes: invalid seed")
2065+ }
2066+ seed = value
2067+ }
2068+ func intOption(_ name: String, default defaultValue: Int) throws -> Int {
2069+ guard let text = command.option(name)?.firstArgument else { return defaultValue }
2070+ guard let value = Int(text), value > 0 else {
2071+ throw ZQEngineError("bayes: invalid \(name)()")
2072+ }
2073+ return value
2074+ }
2075+ let mcmcSize = try intOption("mcmcsize", default: 10_000)
2076+ let burnIn = try intOption("burnin", default: 2_500)
2077+ var priorVariance = 10_000.0
2078+ if let text = command.option("normalprior")?.firstArgument {
2079+ guard let value = Double(text), value > 0 else {
2080+ throw ZQEngineError("bayes: invalid normalprior()")
2081+ }
2082+ priorVariance = value
2083+ }
2084+
2085+ let sample = try buildRegressionSample(body, clusterVariable: nil)
2086+ let result = try ZQBayesianRegression.fitGibbs(
2087+ y: sample.y,
2088+ predictors: sample.predictors,
2089+ includeConstant: !body.hasOption("noconstant"),
2090+ mcmcSize: mcmcSize,
2091+ burnIn: burnIn,
2092+ seed: seed,
2093+ coefficientPriorVariance: priorVariance
2094+ )
2095+ lastEstimation = nil // predict after bayes needs posterior draws
2096+
2097+ var lines = ["Bayesian linear regression (Gibbs)"]
2098+ if sample.droppedMissing > 0 {
2099+ lines.append("(\(sample.droppedMissing) observations dropped due to missing values)")
2100+ }
2101+ for (label, value) in [
2102+ ("Number of obs", "\(result.observationCount)"),
2103+ ("MCMC iterations", "\(result.mcmcSize)"),
2104+ ("Burn-in", "\(result.burnIn)"),
2105+ ("Priors", "b ~ N(0, \(TableFormatter.general(priorVariance))), sigma2 ~ IG(.01, .01)"),
2106+ ] {
2107+ lines.append(
2108+ TableFormatter.pad(label, 46, right: false) + "= " +
2109+ TableFormatter.pad(value, 24)
2110+ )
2111+ }
2112+ lines.append("")
2113+
2114+ let widths = [12, 12, 12, 24]
2115+ lines.append(
2116+ TableFormatter.pad(sample.responseName, widths[0]) + " | " +
2117+ TableFormatter.pad("Mean", widths[1]) + " " +
2118+ TableFormatter.pad("Std. dev.", widths[2]) + " " +
2119+ TableFormatter.pad("[95% cred. interval]", widths[3])
2120+ )
2121+ lines.append(TableFormatter.rule(widths))
2122+
2123+ var scalars: [String: Double] = [
2124+ "N": Double(result.observationCount),
2125+ "mcmcsize": Double(result.mcmcSize),
2126+ "burnin": Double(result.burnIn),
2127+ ]
2128+ for coefficient in result.coefficients + [result.sigma] {
2129+ lines.append(
2130+ TableFormatter.pad(coefficient.name, widths[0]) + " | " +
2131+ TableFormatter.pad(TableFormatter.general(coefficient.posteriorMean), widths[1]) + " " +
2132+ TableFormatter.pad(TableFormatter.general(coefficient.posteriorSD), widths[2]) + " " +
2133+ TableFormatter.pad(
2134+ TableFormatter.general(coefficient.credibleLower) + " " +
2135+ TableFormatter.general(coefficient.credibleUpper),
2136+ widths[3]
2137+ )
2138+ )
2139+ scalars["b_\(coefficient.name)"] = coefficient.posteriorMean
2140+ scalars["sd_\(coefficient.name)"] = coefficient.posteriorSD
2141+ }
2142+ return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
2143+ }
2144+
2049 // MARK: - Permutation test 2145 // MARK: - Permutation test
2050 2146
2051 /// `permute, reps(#) [seed(#)]: reg y x…` — permutes the response 2147 /// `permute, reps(#) [seed(#)]: reg y x…` — permutes the response
added MetrikaKit/Sources/ZQGPU/PhiloxStream.swift +77 −0
@@ -0,0 +1,77 @@
1+//
2+// PhiloxStream.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+import Foundation
11+
12+/// Sequential random-variate stream over the Philox counter space, for
13+/// samplers whose draws are inherently ordered (MCMC chains). Unlike the
14+/// bootstrap's counter-addressable draws, a stream consumes positions one
15+/// by one — rejection samplers use a variable number — but the sequence
16+/// is fully determined by (seed, stream id), so chains are reproducible
17+/// and independent chains never overlap.
18+///
19+/// Stream ids occupy bit 62 of the word-counter space plus a 2⁴⁸-word
20+/// block per id, disjoint from bootstrap draws (low positions) and
21+/// permutation keys (bit 63).
22+public struct PhiloxStream: Sendable {
23+ private let generator: Philox4x32
24+ private let base: UInt64
25+ private var position: UInt64 = 0
26+
27+ public init(seed: UInt64, stream: UInt64 = 0) {
28+ self.generator = Philox4x32(seed: seed)
29+ self.base = (UInt64(1) << 62) &+ stream &* (UInt64(1) << 48)
30+ }
31+
32+ /// Uniform in (0, 1) — endpoints excluded so inverse-CDF transforms
33+ /// and logs stay finite.
34+ public mutating func nextUniform() -> Double {
35+ position &+= 1
36+ let u = generator.uniform(at: base &+ position)
37+ return min(max(u, 5e-324), 1 - 2.2e-16)
38+ }
39+
40+ /// Standard normal via Box–Muller (two uniforms per pair, the spare
41+ /// is cached).
42+ private var cachedNormal: Double?
43+ public mutating func nextNormal() -> Double {
44+ if let cached = cachedNormal {
45+ cachedNormal = nil
46+ return cached
47+ }
48+ let u1 = nextUniform()
49+ let u2 = nextUniform()
50+ let radius = (-2 * Foundation.log(u1)).squareRoot()
51+ let angle = 2 * Double.pi * u2
52+ cachedNormal = radius * Foundation.sin(angle)
53+ return radius * Foundation.cos(angle)
54+ }
55+
56+ /// Gamma(shape, rate) via Marsaglia–Tsang squeeze (with the standard
57+ /// boost for shape < 1).
58+ public mutating func nextGamma(shape: Double, rate: Double) -> Double {
59+ precondition(shape > 0 && rate > 0)
60+ if shape < 1 {
61+ // Gamma(a) = Gamma(a+1) · U^(1/a)
62+ let boosted = nextGamma(shape: shape + 1, rate: rate)
63+ return boosted * Foundation.pow(nextUniform(), 1 / shape)
64+ }
65+ let d = shape - 1.0 / 3.0
66+ let c = 1 / (9 * d).squareRoot()
67+ while true {
68+ let z = nextNormal()
69+ let v = (1 + c * z) * (1 + c * z) * (1 + c * z)
70+ guard v > 0 else { continue }
71+ let u = nextUniform()
72+ if Foundation.log(u) < 0.5 * z * z + d - d * v + d * Foundation.log(v) {
73+ return d * v / rate
74+ }
75+ }
76+ }
77+}
modified MetrikaKit/Sources/ZQParser/KnownVerbs.swift +2 −1
@@ -47,6 +47,7 @@ public struct ZQVerbTable: Sendable {
47 "areg": 4, 47 "areg": 4,
48 "elasticnet": 7, 48 "elasticnet": 7,
49 "lasso": 5, 49 "lasso": 5,
50+ "bayes": 5,
50 "bootstrap": 9, 51 "bootstrap": 9,
51 "permute": 7, 52 "permute": 7,
52 "jackknife": 9, 53 "jackknife": 9,
@@ -71,7 +72,7 @@ public struct ZQVerbTable: Sendable {
71 ], 72 ],
72 fileVerbs: ["use", "save", "import", "export", "log"], 73 fileVerbs: ["use", "save", "import", "export", "log"],
73 assignmentVerbs: ["generate", "replace", "egen"], 74 assignmentVerbs: ["generate", "replace", "egen"],
74− prefixVerbs: ["bootstrap", "permute", "jackknife"], 75+ prefixVerbs: ["bayes", "bootstrap", "permute", "jackknife"],
75 compoundVerbs: ["graph", "ivregress"] 76 compoundVerbs: ["graph", "ivregress"]
76 ) 77 )
77 78
added MetrikaKit/Sources/ZQStats/BayesianRegression.swift +223 −0
@@ -0,0 +1,223 @@
1+//
2+// BayesianRegression.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+import Foundation
11+import ZQGPU
12+
13+/// Bayesian linear regression by Gibbs sampling (CLAUDE.md §1 MCMC
14+/// module). Semi-conjugate model with Stata `bayes` default priors:
15+///
16+/// y | β, σ² ~ N(Xβ, σ²I)
17+/// β_j ~ N(0, τ²), τ² = 10 000 by default
18+/// σ² ~ InvGamma(a₀, b₀), a₀ = b₀ = 0.01 by default
19+///
20+/// Full conditionals are exact (Gibbs, acceptance rate 1):
21+/// β | σ², y ~ N(Vₙ X'y/σ², Vₙ), Vₙ = (I/τ² + X'X/σ²)⁻¹
22+/// σ² | β, y ~ InvGamma(a₀ + n/2, b₀ + ‖y − Xβ‖²/2)
23+///
24+/// Draws come from the Philox stream, so a (seed, chain) pair fully
25+/// determines the chain.
26+public struct ZQBayesCoefficient: Equatable, Sendable {
27+ public var name: String
28+ public var posteriorMean: Double
29+ public var posteriorSD: Double
30+ public var credibleLower: Double // equal-tailed 2.5%
31+ public var credibleUpper: Double // 97.5%
32+}
33+
34+public struct ZQBayesResult: Equatable, Sendable {
35+ public var coefficients: [ZQBayesCoefficient]
36+ /// Posterior summary of σ (the residual standard deviation).
37+ public var sigma: ZQBayesCoefficient
38+ public var observationCount: Int
39+ public var mcmcSize: Int
40+ public var burnIn: Int
41+}
42+
43+public enum ZQBayesianRegression {
44+
45+ public static func fitGibbs(
46+ y: [Double],
47+ predictors: [(name: String, values: [Double])],
48+ includeConstant: Bool = true,
49+ mcmcSize: Int = 10_000,
50+ burnIn: Int = 2_500,
51+ seed: UInt64,
52+ coefficientPriorVariance: Double = 10_000,
53+ sigmaPriorShape: Double = 0.01,
54+ sigmaPriorRate: Double = 0.01
55+ ) throws -> ZQBayesResult {
56+ let n = y.count
57+ var names = predictors.map(\.name)
58+ if includeConstant { names.append("_cons") }
59+ let k = names.count
60+ guard n > k else {
61+ throw ZQStatsError("bayes: insufficient observations: n=\(n), k=\(k)")
62+ }
63+ guard mcmcSize > 10 else { throw ZQStatsError("bayes: mcmcsize too small") }
64+
65+ // Design (column-major) and sufficient statistics.
66+ var x = [Double]()
67+ x.reserveCapacity(n * k)
68+ for column in predictors {
69+ guard column.values.count == n else {
70+ throw ZQStatsError("regressor '\(column.name)' has wrong length")
71+ }
72+ x.append(contentsOf: column.values)
73+ }
74+ if includeConstant { x.append(contentsOf: [Double](repeating: 1, count: n)) }
75+
76+ var xtx = [Double](repeating: 0, count: k * k)
77+ var xty = [Double](repeating: 0, count: k)
78+ for j in 0..<k {
79+ for i in j..<k {
80+ var sum = 0.0
81+ for row in 0..<n { sum += x[i * n + row] * x[j * n + row] }
82+ xtx[j * k + i] = sum
83+ xtx[i * k + j] = sum
84+ }
85+ var sum = 0.0
86+ for row in 0..<n { sum += x[j * n + row] * y[row] }
87+ xty[j] = sum
88+ }
89+
90+ // Start at OLS-ish values via a ridge solve to be safe.
91+ var stream = PhiloxStream(seed: seed)
92+ var beta = [Double](repeating: 0, count: k)
93+ var sigma2 = max(y.reduce(0) { $0 + $1 * $1 } / Double(n), 1e-8)
94+
95+ let total = burnIn + mcmcSize
96+ var betaDraws = [[Double]](repeating: [], count: mcmcSize)
97+ var sigmaDraws = [Double](repeating: 0, count: mcmcSize)
98+ let priorPrecision = 1 / coefficientPriorVariance
99+
100+ for iteration in 0..<total {
101+ // β | σ²: precision Λ = I/τ² + X'X/σ², posterior N(Λ⁻¹X'y/σ², Λ⁻¹).
102+ var lambda = [Double](repeating: 0, count: k * k)
103+ for j in 0..<k {
104+ for i in 0..<k {
105+ lambda[j * k + i] = xtx[j * k + i] / sigma2
106+ }
107+ lambda[j * k + j] += priorPrecision
108+ }
109+ let rhs = xty.map { $0 / sigma2 }
110+
111+ // Cholesky of Λ (lower L, column-major): solve for the mean and
112+ // draw β = mean + L⁻ᵀ z (since Λ = L Lᵀ ⇒ Var = L⁻ᵀ L⁻¹).
113+ let chol = try choleskyLower(lambda, k: k)
114+ let mean = try choleskySolve(chol, k: k, rhs: rhs)
115+ var z = [Double](repeating: 0, count: k)
116+ for j in 0..<k { z[j] = stream.nextNormal() }
117+ let noise = try backSolveTransposed(chol, k: k, rhs: z)
118+ for j in 0..<k { beta[j] = mean[j] + noise[j] }
119+
120+ // σ² | β: InvGamma(a₀ + n/2, b₀ + RSS/2).
121+ var rss = 0.0
122+ for row in 0..<n {
123+ var fitted = 0.0
124+ for j in 0..<k { fitted += beta[j] * x[j * n + row] }
125+ rss += (y[row] - fitted) * (y[row] - fitted)
126+ }
127+ let precision = stream.nextGamma(
128+ shape: sigmaPriorShape + Double(n) / 2,
129+ rate: sigmaPriorRate + rss / 2
130+ )
131+ sigma2 = 1 / precision
132+
133+ if iteration >= burnIn {
134+ betaDraws[iteration - burnIn] = beta
135+ sigmaDraws[iteration - burnIn] = sigma2.squareRoot()
136+ }
137+ }
138+
139+ // Posterior summaries.
140+ func summarize(_ name: String, _ draws: [Double]) -> ZQBayesCoefficient {
141+ let m = draws.reduce(0, +) / Double(draws.count)
142+ let variance = draws.reduce(0) { $0 + ($1 - m) * ($1 - m) }
143+ / Double(draws.count - 1)
144+ let sorted = draws.sorted()
145+ func quantile(_ p: Double) -> Double {
146+ let position = p * Double(sorted.count - 1)
147+ let lower = Int(position)
148+ let fraction = position - Double(lower)
149+ if lower + 1 < sorted.count {
150+ return sorted[lower] * (1 - fraction) + sorted[lower + 1] * fraction
151+ }
152+ return sorted[lower]
153+ }
154+ return ZQBayesCoefficient(
155+ name: name,
156+ posteriorMean: m,
157+ posteriorSD: variance.squareRoot(),
158+ credibleLower: quantile(0.025),
159+ credibleUpper: quantile(0.975)
160+ )
161+ }
162+
163+ var coefficients: [ZQBayesCoefficient] = []
164+ for j in 0..<k {
165+ coefficients.append(summarize(names[j], betaDraws.map { $0[j] }))
166+ }
167+ return ZQBayesResult(
168+ coefficients: coefficients,
169+ sigma: summarize("sigma", sigmaDraws),
170+ observationCount: n,
171+ mcmcSize: mcmcSize,
172+ burnIn: burnIn
173+ )
174+ }
175+
176+ // MARK: - Small dense Cholesky helpers (k×k, column-major)
177+
178+ private static func choleskyLower(_ a: [Double], k: Int) throws -> [Double] {
179+ var l = [Double](repeating: 0, count: k * k)
180+ for j in 0..<k {
181+ var diagonal = a[j * k + j]
182+ for p in 0..<j { diagonal -= l[p * k + j] * l[p * k + j] }
183+ guard diagonal > 0 else {
184+ throw ZQStatsError("bayes: posterior precision is not positive definite")
185+ }
186+ let root = diagonal.squareRoot()
187+ l[j * k + j] = root
188+ for i in (j + 1)..<k {
189+ var value = a[j * k + i]
190+ for p in 0..<j { value -= l[p * k + i] * l[p * k + j] }
191+ l[j * k + i] = value / root
192+ }
193+ }
194+ return l
195+ }
196+
197+ /// Solves Λx = b given the lower Cholesky factor (forward then back).
198+ private static func choleskySolve(
199+ _ l: [Double], k: Int, rhs: [Double]
200+ ) throws -> [Double] {
201+ var z = [Double](repeating: 0, count: k)
202+ for i in 0..<k {
203+ var value = rhs[i]
204+ for p in 0..<i { value -= l[p * k + i] * z[p] }
205+ z[i] = value / l[i * k + i]
206+ }
207+ return try backSolveTransposed(l, k: k, rhs: z)
208+ }
209+
210+ /// Solves Lᵀx = b (also the covariance-square-root transform for
211+ /// drawing from N(0, Λ⁻¹)).
212+ private static func backSolveTransposed(
213+ _ l: [Double], k: Int, rhs: [Double]
214+ ) throws -> [Double] {
215+ var solution = [Double](repeating: 0, count: k)
216+ for i in stride(from: k - 1, through: 0, by: -1) {
217+ var value = rhs[i]
218+ for p in (i + 1)..<k { value -= l[i * k + p] * solution[p] }
219+ solution[i] = value / l[i * k + i]
220+ }
221+ return solution
222+ }
223+}
added MetrikaKit/Tests/MetrikaKitTests/BayesTests.swift +148 −0
@@ -0,0 +1,148 @@
1+//
2+// BayesTests.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+import Foundation
11+import Testing
12+import ZQEngine
13+import ZQGPU
14+import ZQStats
15+
16+@Suite("Bayesian regression", .serialized)
17+struct BayesTests {
18+
19+ @Test("Philox stream normals and gammas match their theoretical moments")
20+ func variateGenerators() {
21+ var stream = PhiloxStream(seed: 42)
22+ let n = 200_000
23+
24+ var normalSum = 0.0, normalSquares = 0.0
25+ for _ in 0..<n {
26+ let z = stream.nextNormal()
27+ normalSum += z
28+ normalSquares += z * z
29+ }
30+ #expect(abs(normalSum / Double(n)) < 0.01)
31+ #expect(abs(normalSquares / Double(n) - 1) < 0.02)
32+
33+ // Gamma(3, rate 2): mean 1.5, variance 0.75.
34+ var gammaSum = 0.0, gammaSquares = 0.0
35+ for _ in 0..<n {
36+ let g = stream.nextGamma(shape: 3, rate: 2)
37+ gammaSum += g
38+ gammaSquares += g * g
39+ }
40+ let gammaMean = gammaSum / Double(n)
41+ let gammaVariance = gammaSquares / Double(n) - gammaMean * gammaMean
42+ #expect(abs(gammaMean - 1.5) < 0.01, "gamma mean \(gammaMean)")
43+ #expect(abs(gammaVariance - 0.75) < 0.02, "gamma variance \(gammaVariance)")
44+
45+ // Shape < 1 boost path: Gamma(0.5, rate 1): mean 0.5.
46+ var smallSum = 0.0
47+ for _ in 0..<n { smallSum += stream.nextGamma(shape: 0.5, rate: 1) }
48+ #expect(abs(smallSum / Double(n) - 0.5) < 0.01)
49+ }
50+
51+ @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 = 300
56+ let x = (0..<n).map { _ in 5 + 10 * stream.nextUniform() }
57+ let y = x.map { 2 + 0.5 * $0 + 0.8 * stream.nextNormal() }
58+
59+ 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: 1e8
64+ )
65+
66+ for (posterior, frequentist) in zip(bayes.coefficients, ols.coefficients) {
67+ // Monte-Carlo error with 40k (autocorrelated) draws: assert
68+ // within 5% of a posterior SD.
69+ let toleranceMean = 0.05 * frequentist.standardError
70+ #expect(
71+ abs(posterior.posteriorMean - frequentist.estimate) < toleranceMean,
72+ "mean[\(posterior.name)]: \(posterior.posteriorMean) vs \(frequentist.estimate)"
73+ )
74+ let sdRatio = posterior.posteriorSD / frequentist.standardError
75+ #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+ }
86+
87+ @Test("a tight prior shrinks coefficients toward zero")
88+ func priorShrinkage() throws {
89+ var stream = PhiloxStream(seed: 9)
90+ let n = 60
91+ let x = (0..<n).map { _ in stream.nextUniform() * 10 }
92+ let y = x.map { 3 * $0 + stream.nextNormal() }
93+
94+ let flat = try ZQBayesianRegression.fitGibbs(
95+ y: y, predictors: [("x", x)],
96+ mcmcSize: 5_000, burnIn: 1_000, seed: 42,
97+ coefficientPriorVariance: 1e6
98+ )
99+ let tight = try ZQBayesianRegression.fitGibbs(
100+ y: y, predictors: [("x", x)],
101+ mcmcSize: 5_000, burnIn: 1_000, seed: 42,
102+ coefficientPriorVariance: 0.01
103+ )
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+ }
110+
111+ @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() }
116+
117+ let a = try ZQBayesianRegression.fitGibbs(
118+ y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 42
119+ )
120+ let b = try ZQBayesianRegression.fitGibbs(
121+ y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 42
122+ )
123+ let c = try ZQBayesianRegression.fitGibbs(
124+ y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 43
125+ )
126+ #expect(a == b)
127+ #expect(a != c)
128+ }
129+
130+ @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)")
136+
137+ 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+}
modified MetrikaKit/Tests/MetrikaKitTests/EngineTests.swift +1 −1
@@ -206,7 +206,7 @@ struct EngineTests {
206 "replace", "drop", "keep", "summarize", "tabulate", "correlate", 206 "replace", "drop", "keep", "summarize", "tabulate", "correlate",
207 "regress", "logit", "probit", "poisson", "ivregress", "xtreg", 207 "regress", "logit", "probit", "poisson", "ivregress", "xtreg",
208 "xtset", "lasso", "elasticnet", "bootstrap", "permute", "predict", 208 "xtset", "lasso", "elasticnet", "bootstrap", "permute", "predict",
209− "margins", "scatter", "histogram", "kdensity", "graph", "display", 209+ "margins", "bayes", "scatter", "histogram", "kdensity", "graph", "display",
210 "set", "log", "help", "zscore", 210 "set", "log", "help", "zscore",
211 ] { 211 ] {
212 #expect(documented.contains(verb), Comment(rawValue: "missing manual entry: \(verb)")) 212 #expect(documented.contains(verb), Comment(rawValue: "missing manual entry: \(verb)"))
213 213
Menu

Theme
↑↓navigate ↵open escclose /⌘Kanywhere