SPB Git

spb/metrika Public

Stata-class statistics, GPU-accelerated by Apple Silicon. Native Swift — no Electron, no Python runtime, no compromises.

macosswiftstatisticsgpu
Swift 92.4% HTML 3.3% R 3% Shell 1.3%

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 5 days ago (Aug 5, 2026) parent a3f7d61

Showing 8 changed files with +561 and −3

modified MetrikaKit/Package.swift +1 −1
@@ -48,7 +48,7 @@ let package = Package(
48 48 ),
49 49 .target(
50 50 name: "ZQStats",
51 − dependencies: ["ZQData"],
51 + dependencies: ["ZQData", "ZQGPU"],
52 52 swiftSettings: strictConcurrency
53 53 ),
54 54 .target(
modified MetrikaKit/Sources/ZQEngine/CommandHelp.swift +13 −0
@@ -230,6 +230,19 @@ public enum ZQCommandReference {
230 230 examples: ["bootstrap, reps(10000) seed(42): reg log_rev price"],
231 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 246 ZQCommandDoc(
234 247 verb: "permute", category: "Resampling & simulation",
235 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 208 case "display": return try handleDisplay(command)
209 209 case "bootstrap": return try await handleBootstrap(command, backend: backend)
210 210 case "permute": return try await handlePermute(command)
211 + case "bayes": return try handleBayes(command)
211 212 case "graph": return try handleGraph(command)
212 213 case "histogram", "scatter", "kdensity":
213 214 var promoted = command
@@ -2046,6 +2047,101 @@ public actor ZQSession {
2046 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 2145 // MARK: - Permutation test
2050 2146
2051 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 47 "areg": 4,
48 48 "elasticnet": 7,
49 49 "lasso": 5,
50 + "bayes": 5,
50 51 "bootstrap": 9,
51 52 "permute": 7,
52 53 "jackknife": 9,
@@ -71,7 +72,7 @@ public struct ZQVerbTable: Sendable {
71 72 ],
72 73 fileVerbs: ["use", "save", "import", "export", "log"],
73 74 assignmentVerbs: ["generate", "replace", "egen"],
74 − prefixVerbs: ["bootstrap", "permute", "jackknife"],
75 + prefixVerbs: ["bayes", "bootstrap", "permute", "jackknife"],
75 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 206 "replace", "drop", "keep", "summarize", "tabulate", "correlate",
207 207 "regress", "logit", "probit", "poisson", "ivregress", "xtreg",
208 208 "xtset", "lasso", "elasticnet", "bootstrap", "permute", "predict",
209 − "margins", "scatter", "histogram", "kdensity", "graph", "display",
209 + "margins", "bayes", "scatter", "histogram", "kdensity", "graph", "display",
210 210 "set", "log", "help", "zscore",
211 211 ] {
212 212 #expect(documented.contains(verb), Comment(rawValue: "missing manual entry: \(verb)"))
213 213