spb/metrika Public

1//2//  OLS.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Variance estimator selection for linear regression.13public enum ZQVarianceEstimator: Equatable, Sendable {14    /// Classical homoskedastic σ²(X'X)⁻¹.15    case classical16    /// Heteroskedasticity-consistent sandwich estimators. `robust` in17    /// Stata parlance maps to `.hc1`.18    case hc0, hc1, hc2, hc319    /// Cluster-robust: `cluster` holds one group label per observation.20    case cluster([Int])21}2223/// One row of a regression coefficient table.24public struct ZQCoefficient: Equatable, Sendable {25    public var name: String26    public var estimate: Double27    public var standardError: Double28    public var tStatistic: Double29    public var pValue: Double30    public var confidenceLower: Double31    public var confidenceUpper: Double3233    public init(34        name: String, estimate: Double, standardError: Double,35        tStatistic: Double, pValue: Double,36        confidenceLower: Double, confidenceUpper: Double37    ) {38        self.name = name39        self.estimate = estimate40        self.standardError = standardError41        self.tStatistic = tStatistic42        self.pValue = pValue43        self.confidenceLower = confidenceLower44        self.confidenceUpper = confidenceUpper45    }46}4748/// Full OLS fit result.49public struct ZQOLSResult: Equatable, Sendable {50    public var coefficients: [ZQCoefficient]51    public var observationCount: Int52    public var degreesOfFreedomResidual: Int53    /// Degrees of freedom used for t statistics and CIs: n−k classically,54    /// G−1 under clustering (Stata convention).55    public var inferenceDF: Double56    public var rSquared: Double57    public var adjustedRSquared: Double58    public var rootMSE: Double59    public var fStatistic: Double?60    public var fPValue: Double?61    public var fDF: (Double, Double)? {62        get { fDFStorage.map { ($0.first, $0.second) } }63        set { fDFStorage = newValue.map { FDF($0.0, $0.1) } }64    }65    public var clusterCount: Int?66    public var residuals: [Double]67    /// Full covariance matrix of the coefficients, column-major k×k,68    /// aligned with `coefficients` — needed by delta-method69    /// post-estimation (margins).70    public var vce: [Double]7172    private struct FDF: Equatable, Sendable {73        let first: Double74        let second: Double75        init(_ first: Double, _ second: Double) {76            self.first = first77            self.second = second78        }79        var asTuple: (Double, Double) { (first, second) }80    }8182    private var fDFStorage: FDF?8384    public init(85        coefficients: [ZQCoefficient],86        observationCount: Int,87        degreesOfFreedomResidual: Int,88        inferenceDF: Double,89        rSquared: Double,90        adjustedRSquared: Double,91        rootMSE: Double,92        fStatistic: Double?,93        fPValue: Double?,94        fDF: (Double, Double)?,95        clusterCount: Int?,96        residuals: [Double],97        vce: [Double] = []98    ) {99        self.coefficients = coefficients100        self.observationCount = observationCount101        self.degreesOfFreedomResidual = degreesOfFreedomResidual102        self.inferenceDF = inferenceDF103        self.rSquared = rSquared104        self.adjustedRSquared = adjustedRSquared105        self.rootMSE = rootMSE106        self.fStatistic = fStatistic107        self.fPValue = fPValue108        self.fDFStorage = fDF.map { FDF($0.0, $0.1) }109        self.clusterCount = clusterCount110        self.residuals = residuals111        self.vce = vce112    }113}114115public struct ZQStatsError: Error, Equatable, Sendable, CustomStringConvertible {116    public var message: String117    public init(_ message: String) { self.message = message }118    public var description: String { message }119}120121/// Ordinary least squares on fully observed (already listwise-deleted)122/// data. Solved by LAPACK QR (dgeqrf/dormqr) — X'X is never formed for123/// the coefficient path (CLAUDE.md §5).124public enum ZQOLS {125126    /// - Parameters:127    ///   - y: response, length n128    ///   - predictors: named regressor columns (length n each), intercept129    ///     excluded — it is appended automatically as `_cons` unless130    ///     `includeConstant` is false.131    public static func fit(132        y: [Double],133        predictors: [(name: String, values: [Double])],134        includeConstant: Bool = true,135        variance: ZQVarianceEstimator = .classical,136        confidenceLevel: Double = 0.95137    ) throws -> ZQOLSResult {138        let n = y.count139        var names = predictors.map(\.name)140        if includeConstant { names.append("_cons") }141        let k = names.count142        guard n > k else {143            throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")144        }145        for column in predictors where column.values.count != n {146            throw ZQStatsError("regressor '\(column.name)' has wrong length")147        }148149        // Design matrix, column-major.150        var x = [Double]()151        x.reserveCapacity(n * k)152        for column in predictors { x.append(contentsOf: column.values) }153        if includeConstant { x.append(contentsOf: [Double](repeating: 1, count: n)) }154155        let qr: LinearAlgebra.QR156        let beta: [Double]157        let xtxInverse: [Double]158        do {159            qr = try LinearAlgebra.QR(matrix: x, rows: n, cols: k)160            beta = try qr.solve(rhs: y)161            xtxInverse = try qr.crossProductInverse()162        } catch {163            throw ZQStatsError(164                "design matrix is rank deficient or QR failed: \(error)"165            )166        }167168        // Residuals and fit statistics.169        let fitted = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)170        var residuals = [Double](repeating: 0, count: n)171        var rss = 0.0172        for i in 0..<n {173            residuals[i] = y[i] - fitted[i]174            rss += residuals[i] * residuals[i]175        }176        let meanY = y.reduce(0, +) / Double(n)177        let tss = y.reduce(0) { $0 + ($1 - meanY) * ($1 - meanY) }178        let dfResidual = n - k179        let sigma2 = rss / Double(dfResidual)180        let rSquared = tss > 0 ? 1 - rss / tss : .nan181        let dfModel = k - (includeConstant ? 1 : 0)182        let adjustedRSquared =183            1 - (1 - rSquared) * Double(n - (includeConstant ? 1 : 0)) / Double(dfResidual)184185        // Covariance matrix.186        var clusterCount: Int? = nil187        let vce: [Double]188        switch variance {189        case .classical:190            vce = xtxInverse.map { $0 * sigma2 }191192        case .hc0, .hc1, .hc2, .hc3:193            let leverage: [Double]?194            if variance == .hc2 || variance == .hc3 {195                let q = try qr.thinQ()196                var h = [Double](repeating: 0, count: n)197                for j in 0..<k {198                    for i in 0..<n {199                        let value = q[j * n + i]200                        h[i] += value * value201                    }202                }203                leverage = h204            } else {205                leverage = nil206            }207            var weights = [Double](repeating: 0, count: n)208            for i in 0..<n {209                let e2 = residuals[i] * residuals[i]210                switch variance {211                case .hc0: weights[i] = e2212                case .hc1: weights[i] = e2 * Double(n) / Double(dfResidual)213                case .hc2: weights[i] = e2 / (1 - leverage![i])214                case .hc3: weights[i] = e2 / ((1 - leverage![i]) * (1 - leverage![i]))215                default: break216                }217            }218            let meat = weightedCrossProduct(x: x, n: n, k: k, weights: weights)219            vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)220221        case .cluster(let groups):222            guard groups.count == n else {223                throw ZQStatsError("cluster variable has wrong length")224            }225            var scores: [Int: [Double]] = [:]226            for i in 0..<n {227                var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)228                for j in 0..<k {229                    u[j] += x[j * n + i] * residuals[i]230                }231                scores[groups[i]] = u232            }233            let g = scores.count234            guard g > 1 else {235                throw ZQStatsError("cluster variable must define at least 2 groups")236            }237            clusterCount = g238            var meat = [Double](repeating: 0, count: k * k)239            for u in scores.values {240                for j in 0..<k {241                    for i in 0..<k {242                        meat[j * k + i] += u[i] * u[j]243                    }244                }245            }246            // Stata regress small-sample factor: G/(G−1) · (n−1)/(n−k).247            let scale = Double(g) / Double(g - 1) * Double(n - 1) / Double(dfResidual)248            for index in meat.indices { meat[index] *= scale }249            vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)250        }251252        // Inference: t statistics against n−k df, or G−1 under clustering.253        let inferenceDF: Double254        if let g = clusterCount {255            inferenceDF = Double(g - 1)256        } else {257            inferenceDF = Double(dfResidual)258        }259        let tCritical = ZQDistributions.studentTQuantile(260            0.5 + confidenceLevel / 2, df: inferenceDF261        )262263        var coefficients: [ZQCoefficient] = []264        for j in 0..<k {265            let se = vce[j * k + j].squareRoot()266            let t = beta[j] / se267            coefficients.append(ZQCoefficient(268                name: names[j],269                estimate: beta[j],270                standardError: se,271                tStatistic: t,272                pValue: ZQDistributions.tTestPValue(t, df: inferenceDF),273                confidenceLower: beta[j] - tCritical * se,274                confidenceUpper: beta[j] + tCritical * se275            ))276        }277278        // Overall Wald F test on all non-constant coefficients.279        var fStatistic: Double? = nil280        var fPValue: Double? = nil281        var fDF: (Double, Double)? = nil282        if dfModel > 0 {283            let restricted = Array(0..<dfModel)284            var subVce = [Double](repeating: 0, count: dfModel * dfModel)285            var subBeta = [Double](repeating: 0, count: dfModel)286            for (jj, j) in restricted.enumerated() {287                subBeta[jj] = beta[j]288                for (ii, i) in restricted.enumerated() {289                    subVce[jj * dfModel + ii] = vce[j * k + i]290                }291            }292            if let solved = try? LinearAlgebra.solveSymmetric(subVce, k: dfModel, rhs: subBeta) {293                let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }294                let f = wald / Double(dfModel)295                fStatistic = f296                fDF = (Double(dfModel), inferenceDF)297                fPValue = ZQDistributions.fTestPValue(298                    f, df1: Double(dfModel), df2: inferenceDF299                )300            }301        }302303        return ZQOLSResult(304            coefficients: coefficients,305            observationCount: n,306            degreesOfFreedomResidual: dfResidual,307            inferenceDF: inferenceDF,308            rSquared: rSquared,309            adjustedRSquared: adjustedRSquared,310            rootMSE: sigma2.squareRoot(),311            fStatistic: fStatistic,312            fPValue: fPValue,313            fDF: fDF,314            clusterCount: clusterCount,315            residuals: residuals,316            vce: vce317        )318    }319320    /// meat = Σ_i w_i · x_i x_iᵀ for column-major X.321    private static func weightedCrossProduct(322        x: [Double], n: Int, k: Int, weights: [Double]323    ) -> [Double] {324        var meat = [Double](repeating: 0, count: k * k)325        for j in 0..<k {326            for i in j..<k {327                var sum = 0.0328                for row in 0..<n {329                    sum += weights[row] * x[i * n + row] * x[j * n + row]330                }331                meat[j * k + i] = sum332                meat[i * k + j] = sum333            }334        }335        return meat336    }337}338