spb/metrika Public

1//2//  BayesianRegression.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import ZQGPU1213/// Bayesian linear regression by Gibbs sampling (CLAUDE.md §1 MCMC14/// module). Semi-conjugate model with Stata `bayes` default priors:15///16///   y | β, σ²  ~  N(Xβ, σ²I)17///   β_j        ~  N(0, τ²),        τ² = 10 000 by default18///   σ²         ~  InvGamma(a₀, b₀), a₀ = b₀ = 0.01 by default19///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 fully25/// determines the chain.26public struct ZQBayesCoefficient: Equatable, Sendable {27    public var name: String28    public var posteriorMean: Double29    public var posteriorSD: Double30    public var credibleLower: Double   // equal-tailed 2.5%31    public var credibleUpper: Double   // 97.5%32}3334public struct ZQBayesResult: Equatable, Sendable {35    public var coefficients: [ZQBayesCoefficient]36    /// Posterior summary of σ (the residual standard deviation).37    public var sigma: ZQBayesCoefficient38    public var observationCount: Int39    public var mcmcSize: Int40    public var burnIn: Int41}4243public enum ZQBayesianRegression {4445    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.0155    ) throws -> ZQBayesResult {56        let n = y.count57        var names = predictors.map(\.name)58        if includeConstant { names.append("_cons") }59        let k = names.count60        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") }6465        // 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)) }7576        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.081                for row in 0..<n { sum += x[i * n + row] * x[j * n + row] }82                xtx[j * k + i] = sum83                xtx[i * k + j] = sum84            }85            var sum = 0.086            for row in 0..<n { sum += x[j * n + row] * y[row] }87            xty[j] = sum88        }8990        // 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)9495        let total = burnIn + mcmcSize96        var betaDraws = [[Double]](repeating: [], count: mcmcSize)97        var sigmaDraws = [Double](repeating: 0, count: mcmcSize)98        let priorPrecision = 1 / coefficientPriorVariance99100        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] / sigma2106                }107                lambda[j * k + j] += priorPrecision108            }109            let rhs = xty.map { $0 / sigma2 }110111            // Cholesky of Λ (lower L, column-major): solve for the mean and112            // 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] }119120            // σ² | β: InvGamma(a₀ + n/2, b₀ + RSS/2).121            var rss = 0.0122            for row in 0..<n {123                var fitted = 0.0124                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 / 2130            )131            sigma2 = 1 / precision132133            if iteration >= burnIn {134                betaDraws[iteration - burnIn] = beta135                sigmaDraws[iteration - burnIn] = sigma2.squareRoot()136            }137        }138139        // 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] * fraction151                }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        }162163        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: burnIn173        )174    }175176    // MARK: - Small dense Cholesky helpers (k×k, column-major)177178    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] = root188            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 / root192            }193        }194        return l195    }196197    /// 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    }209210    /// Solves Lᵀx = b (also the covariance-square-root transform for211    /// 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 solution222    }223}224