spb/metrika Public

1//2//  IV.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Accelerate11import Foundation1213/// Two-stage least squares for `ivregress 2sls`.14///15/// The projection uses the thin Q of the instrument matrix (X̂ = Q·QᵀX),16/// never forming Z'Z. Inference follows Stata's `small` convention:17/// σ² = u'u/(N−K) with t statistics; residuals come from the ORIGINAL18/// regressors (u = y − Xβ), not the projected ones.19public enum ZQIV {2021    public static func fit2SLS(22        y: [Double],23        endogenous: [(name: String, values: [Double])],24        exogenous: [(name: String, values: [Double])],25        instruments: [(name: String, values: [Double])],26        includeConstant: Bool = true,27        variance: ZQVarianceEstimator = .classical,28        confidenceLevel: Double = 0.9529    ) throws -> ZQOLSResult {30        let n = y.count31        guard !endogenous.isEmpty else {32            throw ZQStatsError("ivregress: at least one endogenous regressor required")33        }34        let k = endogenous.count + exogenous.count + (includeConstant ? 1 : 0)35        let m = instruments.count + exogenous.count + (includeConstant ? 1 : 0)36        guard m >= k else {37            throw ZQStatsError(38                "ivregress: order condition fails — \(instruments.count) instruments for \(endogenous.count) endogenous regressors"39            )40        }41        guard n > k else {42            throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")43        }4445        func columnMajor(_ columns: [[Double]]) -> [Double] {46            var flat = [Double]()47            flat.reserveCapacity(n * columns.count)48            for column in columns { flat.append(contentsOf: column) }49            return flat50        }51        let ones = [Double](repeating: 1, count: n)5253        // X = [endog, exog, 1], Z = [instruments, exog, 1].54        var xColumns = endogenous.map(\.values) + exogenous.map(\.values)55        var zColumns = instruments.map(\.values) + exogenous.map(\.values)56        if includeConstant {57            xColumns.append(ones)58            zColumns.append(ones)59        }60        for column in xColumns + zColumns where column.count != n {61            throw ZQStatsError("ivregress: variable has wrong length")62        }63        let x = columnMajor(xColumns)64        let z = columnMajor(zColumns)6566        // First stage: X̂ = Q·(QᵀX) with Q the thin Q of Z.67        let projected: [Double]68        do {69            let zQR = try LinearAlgebra.QR(matrix: z, rows: n, cols: m)70            let q = try zQR.thinQ()71            var qtx = [Double](repeating: 0, count: m * k)72            cblas_dgemm(73                CblasColMajor, CblasTrans, CblasNoTrans,74                Int32(m), Int32(k), Int32(n),75                1.0, q, Int32(n), x, Int32(n), 0.0, &qtx, Int32(m)76            )77            var xhat = [Double](repeating: 0, count: n * k)78            cblas_dgemm(79                CblasColMajor, CblasNoTrans, CblasNoTrans,80                Int32(n), Int32(k), Int32(m),81                1.0, q, Int32(n), qtx, Int32(m), 0.0, &xhat, Int32(n)82            )83            projected = xhat84        } catch let error as LinearAlgebra.Failure {85            throw ZQStatsError("ivregress: instrument matrix is rank deficient — \(error)")86        }8788        // Second stage.89        let secondStage: LinearAlgebra.QR90        let beta: [Double]91        let bread: [Double]92        do {93            secondStage = try LinearAlgebra.QR(matrix: projected, rows: n, cols: k)94            beta = try secondStage.solve(rhs: y)95            bread = try secondStage.crossProductInverse()96        } catch {97            throw ZQStatsError(98                "ivregress: projected design is rank deficient (weak or collinear instruments?)"99            )100        }101102        // Residuals from the ORIGINAL regressors.103        let fitted = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)104        var residuals = [Double](repeating: 0, count: n)105        var rss = 0.0106        for i in 0..<n {107            residuals[i] = y[i] - fitted[i]108            rss += residuals[i] * residuals[i]109        }110        let dfResidual = n - k111        let sigma2 = rss / Double(dfResidual)112        let meanY = y.reduce(0, +) / Double(n)113        let tss = y.reduce(0) { $0 + ($1 - meanY) * ($1 - meanY) }114        let rSquared = tss > 0 ? 1 - rss / tss : .nan   // can be negative for IV115        let dfModel = k - (includeConstant ? 1 : 0)116117        // Covariance: sandwich pieces built on the PROJECTED regressors.118        var clusterCount: Int?119        let vce: [Double]120        switch variance {121        case .classical:122            vce = bread.map { $0 * sigma2 }123        case .hc0, .hc1, .hc2, .hc3:124            var meat = [Double](repeating: 0, count: k * k)125            for row in 0..<n {126                let weight = residuals[row] * residuals[row]127                for j in 0..<k {128                    for i in j..<k {129                        let value = weight * projected[i * n + row] * projected[j * n + row]130                        meat[j * k + i] += value131                    }132                }133            }134            for j in 0..<k {135                for i in 0..<j { meat[j * k + i] = meat[i * k + j] }136            }137            if variance == .hc1 {138                let scale = Double(n) / Double(dfResidual)139                for index in meat.indices { meat[index] *= scale }140            }141            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)142        case .cluster(let clusters):143            guard clusters.count == n else {144                throw ZQStatsError("cluster variable has wrong length")145            }146            var scores: [Int: [Double]] = [:]147            for i in 0..<n {148                var u = scores[clusters[i]] ?? [Double](repeating: 0, count: k)149                for j in 0..<k {150                    u[j] += projected[j * n + i] * residuals[i]151                }152                scores[clusters[i]] = u153            }154            let g = scores.count155            guard g > 1 else {156                throw ZQStatsError("cluster variable must define at least 2 groups")157            }158            clusterCount = g159            var meat = [Double](repeating: 0, count: k * k)160            for u in scores.values {161                for j in 0..<k {162                    for i in 0..<k {163                        meat[j * k + i] += u[i] * u[j]164                    }165                }166            }167            let scale = Double(g) / Double(g - 1) * Double(n - 1) / Double(dfResidual)168            for index in meat.indices { meat[index] *= scale }169            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)170        }171172        let inferenceDF = clusterCount.map { Double($0 - 1) } ?? Double(dfResidual)173        let tCritical = ZQDistributions.studentTQuantile(174            0.5 + confidenceLevel / 2, df: inferenceDF175        )176177        var names = endogenous.map(\.name) + exogenous.map(\.name)178        if includeConstant { names.append("_cons") }179        var coefficients: [ZQCoefficient] = []180        for j in 0..<k {181            let se = vce[j * k + j].squareRoot()182            let t = beta[j] / se183            coefficients.append(ZQCoefficient(184                name: names[j],185                estimate: beta[j],186                standardError: se,187                tStatistic: t,188                pValue: ZQDistributions.tTestPValue(t, df: inferenceDF),189                confidenceLower: beta[j] - tCritical * se,190                confidenceUpper: beta[j] + tCritical * se191            ))192        }193194        // Wald F on the slopes.195        var fStatistic: Double?196        var fPValue: Double?197        var fDF: (Double, Double)?198        if dfModel > 0 {199            var subVce = [Double](repeating: 0, count: dfModel * dfModel)200            var subBeta = [Double](repeating: 0, count: dfModel)201            for j in 0..<dfModel {202                subBeta[j] = beta[j]203                for i in 0..<dfModel {204                    subVce[j * dfModel + i] = vce[j * k + i]205                }206            }207            if let solved = try? LinearAlgebra.solveSymmetric(208                subVce, k: dfModel, rhs: subBeta209            ) {210                let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }211                fStatistic = wald / Double(dfModel)212                fDF = (Double(dfModel), inferenceDF)213                fPValue = ZQDistributions.fTestPValue(214                    wald / Double(dfModel), df1: Double(dfModel), df2: inferenceDF215                )216            }217        }218219        return ZQOLSResult(220            coefficients: coefficients,221            observationCount: n,222            degreesOfFreedomResidual: dfResidual,223            inferenceDF: inferenceDF,224            rSquared: rSquared,225            adjustedRSquared: 1 - (1 - rSquared) * Double(n - 1) / Double(dfResidual),226            rootMSE: sigma2.squareRoot(),227            fStatistic: fStatistic,228            fPValue: fPValue,229            fDF: fDF,230            clusterCount: clusterCount,231            residuals: residuals,232            vce: vce233        )234    }235}236