spb/metrika Public

1//2//  FixedEffects.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Panel fixed-effects (within) estimator for `xtreg, fe`.13///14/// The within transformation subtracts group means and adds back grand15/// means (Stata's convention, so `_cons` is reported), then solves by the16/// usual QR path. Degrees of freedom absorb the G group effects:17/// df = N − K − G. Cluster-robust VCE clusters on the panel variable with18/// the G/(G−1) factor and t statistics on G−1 df.19public struct ZQFEResult: Equatable, Sendable {20    public var coefficients: [ZQCoefficient]21    public var observationCount: Int22    public var groupCount: Int23    public var degreesOfFreedomResidual: Int24    public var inferenceDF: Double25    /// Within R²: fit of the demeaned regression.26    public var rSquaredWithin: Double27    public var rootMSE: Double28    public var fStatistic: Double?29    public var fPValue: Double?30    public var clustered: Bool3132    public init(33        coefficients: [ZQCoefficient], observationCount: Int, groupCount: Int,34        degreesOfFreedomResidual: Int, inferenceDF: Double,35        rSquaredWithin: Double, rootMSE: Double,36        fStatistic: Double?, fPValue: Double?, clustered: Bool37    ) {38        self.coefficients = coefficients39        self.observationCount = observationCount40        self.groupCount = groupCount41        self.degreesOfFreedomResidual = degreesOfFreedomResidual42        self.inferenceDF = inferenceDF43        self.rSquaredWithin = rSquaredWithin44        self.rootMSE = rootMSE45        self.fStatistic = fStatistic46        self.fPValue = fPValue47        self.clustered = clustered48    }49}5051public enum ZQFixedEffects {5253    /// - Parameters:54    ///   - groups: dense panel codes (0..<G), one per observation.55    ///   - clustered: cluster the VCE on the panel variable.56    public static func fitWithin(57        y: [Double],58        predictors: [(name: String, values: [Double])],59        groups: [Int],60        clustered: Bool = false,61        confidenceLevel: Double = 0.9562    ) throws -> ZQFEResult {63        let n = y.count64        let slopes = predictors.count65        guard slopes > 0 else { throw ZQStatsError("xtreg: regressors required") }66        guard groups.count == n else {67            throw ZQStatsError("panel variable has wrong length")68        }69        let groupCount = Set(groups).count70        let k = slopes + 1   // slopes + reported constant71        let dfResidual = n - slopes - groupCount72        guard dfResidual > 0 else {73            throw ZQStatsError("insufficient observations: N=\(n), K=\(slopes), G=\(groupCount)")74        }7576        // Group means and grand means.77        func withinTransform(_ values: [Double]) -> (transformed: [Double], demeaned: [Double]) {78            var sums = [Double](repeating: 0, count: groupCount)79            var counts = [Double](repeating: 0, count: groupCount)80            for i in 0..<n {81                sums[groups[i]] += values[i]82                counts[groups[i]] += 183            }84            let grand = values.reduce(0, +) / Double(n)85            var transformed = [Double](repeating: 0, count: n)86            var demeaned = [Double](repeating: 0, count: n)87            for i in 0..<n {88                let groupMean = sums[groups[i]] / counts[groups[i]]89                demeaned[i] = values[i] - groupMean90                transformed[i] = demeaned[i] + grand91            }92            return (transformed, demeaned)93        }9495        let (yTransformed, yDemeaned) = withinTransform(y)96        var design = [Double]()97        design.reserveCapacity(n * k)98        var demeanedColumns: [[Double]] = []99        for column in predictors {100            guard column.values.count == n else {101                throw ZQStatsError("regressor '\(column.name)' has wrong length")102            }103            let (transformed, demeaned) = withinTransform(column.values)104            design.append(contentsOf: transformed)105            demeanedColumns.append(demeaned)106        }107        design.append(contentsOf: [Double](repeating: 1, count: n))108109        let qr: LinearAlgebra.QR110        let beta: [Double]111        let xtxInverse: [Double]112        do {113            qr = try LinearAlgebra.QR(matrix: design, rows: n, cols: k)114            beta = try qr.solve(rhs: yTransformed)115            xtxInverse = try qr.crossProductInverse()116        } catch {117            throw ZQStatsError("xtreg: design matrix is rank deficient after demeaning — \(error)")118        }119120        // Within residuals and fit.121        let fitted = LinearAlgebra.multiply(matrix: design, rows: n, cols: k, vector: beta)122        var residuals = [Double](repeating: 0, count: n)123        var rss = 0.0124        for i in 0..<n {125            residuals[i] = yTransformed[i] - fitted[i]126            rss += residuals[i] * residuals[i]127        }128        let tssWithin = yDemeaned.reduce(0) { $0 + $1 * $1 }129        let sigma2 = rss / Double(dfResidual)130        let rSquaredWithin = tssWithin > 0 ? 1 - rss / tssWithin : .nan131132        // Covariance.133        let vce: [Double]134        let inferenceDF: Double135        if clustered {136            var scores: [Int: [Double]] = [:]137            for i in 0..<n {138                var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)139                for j in 0..<slopes {140                    u[j] += demeanedColumns[j][i] * residuals[i]141                }142                u[slopes] += residuals[i]   // constant column143                scores[groups[i]] = u144            }145            guard groupCount > 1 else {146                throw ZQStatsError("cluster-robust VCE needs at least 2 groups")147            }148            var meat = [Double](repeating: 0, count: k * k)149            for u in scores.values {150                for j in 0..<k {151                    for i in 0..<k {152                        meat[j * k + i] += u[i] * u[j]153                    }154                }155            }156            let scale = Double(groupCount) / Double(groupCount - 1)157            for index in meat.indices { meat[index] *= scale }158            vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)159            inferenceDF = Double(groupCount - 1)160        } else {161            vce = xtxInverse.map { $0 * sigma2 }162            inferenceDF = Double(dfResidual)163        }164165        let tCritical = ZQDistributions.studentTQuantile(166            0.5 + confidenceLevel / 2, df: inferenceDF167        )168        var names = predictors.map(\.name)169        names.append("_cons")170        var coefficients: [ZQCoefficient] = []171        for j in 0..<k {172            let se = vce[j * k + j].squareRoot()173            let t = beta[j] / se174            coefficients.append(ZQCoefficient(175                name: names[j],176                estimate: beta[j],177                standardError: se,178                tStatistic: t,179                pValue: ZQDistributions.tTestPValue(t, df: inferenceDF),180                confidenceLower: beta[j] - tCritical * se,181                confidenceUpper: beta[j] + tCritical * se182            ))183        }184185        // Wald F on the slopes.186        var fStatistic: Double?187        var fPValue: Double?188        var subVce = [Double](repeating: 0, count: slopes * slopes)189        var subBeta = [Double](repeating: 0, count: slopes)190        for j in 0..<slopes {191            subBeta[j] = beta[j]192            for i in 0..<slopes {193                subVce[j * slopes + i] = vce[j * k + i]194            }195        }196        if let solved = try? LinearAlgebra.solveSymmetric(subVce, k: slopes, rhs: subBeta) {197            let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }198            fStatistic = wald / Double(slopes)199            fPValue = ZQDistributions.fTestPValue(200                wald / Double(slopes), df1: Double(slopes), df2: inferenceDF201            )202        }203204        return ZQFEResult(205            coefficients: coefficients,206            observationCount: n,207            groupCount: groupCount,208            degreesOfFreedomResidual: dfResidual,209            inferenceDF: inferenceDF,210            rSquaredWithin: rSquaredWithin,211            rootMSE: sigma2.squareRoot(),212            fStatistic: fStatistic,213            fPValue: fPValue,214            clustered: clustered215        )216    }217}218