spb/metrika Public

1//2//  GLM.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Maximum-likelihood binary/count models via iteratively reweighted13/// least squares (Fisher scoring), CLAUDE.md §5. Each IRLS step solves a14/// weighted least-squares problem through the same LAPACK QR path as OLS,15/// so X'WX is never formed explicitly.16public enum ZQGLMFamily: String, Equatable, Sendable {17    case logit, probit, poisson18}1920public struct ZQGLMResult: Equatable, Sendable {21    public var family: ZQGLMFamily22    public var coefficients: [ZQCoefficient]23    public var observationCount: Int24    public var logLikelihood: Double25    public var nullLogLikelihood: Double26    /// McFadden pseudo-R²: 1 − ll/ll₀ (Stata's reported definition).27    public var pseudoRSquared: Double28    /// LR χ² against the intercept-only model (classical VCE only; under29    /// robust/cluster VCE Stata reports a Wald χ², which we mirror).30    public var chiSquared: Double?31    public var chiSquaredDF: Int32    public var chiSquaredPValue: Double?33    public var iterations: Int34    public var clusterCount: Int?35    /// Linear predictor and fitted mean at the optimum.36    public var fittedMeans: [Double]37    /// Full covariance matrix, column-major k×k, aligned with38    /// `coefficients`.39    public var vce: [Double]4041    public init(42        family: ZQGLMFamily,43        coefficients: [ZQCoefficient],44        observationCount: Int,45        logLikelihood: Double,46        nullLogLikelihood: Double,47        pseudoRSquared: Double,48        chiSquared: Double?,49        chiSquaredDF: Int,50        chiSquaredPValue: Double?,51        iterations: Int,52        clusterCount: Int?,53        fittedMeans: [Double],54        vce: [Double] = []55    ) {56        self.family = family57        self.coefficients = coefficients58        self.observationCount = observationCount59        self.logLikelihood = logLikelihood60        self.nullLogLikelihood = nullLogLikelihood61        self.pseudoRSquared = pseudoRSquared62        self.chiSquared = chiSquared63        self.chiSquaredDF = chiSquaredDF64        self.chiSquaredPValue = chiSquaredPValue65        self.iterations = iterations66        self.clusterCount = clusterCount67        self.fittedMeans = fittedMeans68        self.vce = vce69    }70}7172public enum ZQGLM {7374    /// Fits by Fisher scoring to |Δll| < 1e-13 (tighter than R's default75    /// so both land on the same optimum to ≥1e-10).76    ///77    /// Variance estimators: `.classical` is the inverse expected78    /// information; `.hc0`/`.hc1` (and Stata's `robust`) are the score79    /// sandwich without df correction; `.cluster` applies Stata's ML80    /// factor G/(G−1). `.hc2`/`.hc3` are not defined for ML estimators.81    public static func fit(82        y: [Double],83        predictors: [(name: String, values: [Double])],84        family: ZQGLMFamily,85        includeConstant: Bool = true,86        variance: ZQVarianceEstimator = .classical,87        confidenceLevel: Double = 0.95,88        maximumIterations: Int = 10089    ) throws -> ZQGLMResult {90        let n = y.count91        var names = predictors.map(\.name)92        if includeConstant { names.append("_cons") }93        let k = names.count94        guard n > k else {95            throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")96        }97        if case .hc2 = variance {98            throw ZQStatsError("hc2 is not defined for maximum-likelihood estimators")99        }100        if case .hc3 = variance {101            throw ZQStatsError("hc3 is not defined for maximum-likelihood estimators")102        }103104        switch family {105        case .logit, .probit:106            for value in y where value != 0 && value != 1 {107                throw ZQStatsError("\(family.rawValue): outcome must be 0/1")108            }109        case .poisson:110            for value in y where value < 0 || value != value.rounded() {111                throw ZQStatsError("poisson: outcome must be a nonnegative count")112            }113        }114115        // Design matrix, column-major.116        var x = [Double]()117        x.reserveCapacity(n * k)118        for column in predictors {119            guard column.values.count == n else {120                throw ZQStatsError("regressor '\(column.name)' has wrong length")121            }122            x.append(contentsOf: column.values)123        }124        if includeConstant { x.append(contentsOf: [Double](repeating: 1, count: n)) }125126        // IRLS.127        var beta = [Double](repeating: 0, count: k)128        var eta = [Double](repeating: 0, count: n)129        var mu = startingMeans(y: y, family: family)130        if family == .poisson || !includeConstant {131            eta = mu.map { link(family: family, mean: $0) }132        } else {133            eta = mu.map { link(family: family, mean: $0) }134        }135        var logLikelihood = self.logLikelihood(y: y, mu: mu, family: family)136        var iterations = 0137        var finalQR: LinearAlgebra.QR?138        var weights = [Double](repeating: 0, count: n)139140        for iteration in 1...maximumIterations {141            iterations = iteration142143            // Working response and weights.144            var z = [Double](repeating: 0, count: n)145            for i in 0..<n {146                let derivative = meanDerivative(family: family, eta: eta[i], mean: mu[i])147                let varMu = varianceFunction(family: family, mean: mu[i])148                weights[i] = derivative * derivative / varMu149                z[i] = eta[i] + (y[i] - mu[i]) / derivative150            }151152            // Weighted least squares via QR on √W·X, √W·z.153            var xw = [Double](repeating: 0, count: n * k)154            var zw = [Double](repeating: 0, count: n)155            for i in 0..<n {156                let root = weights[i].squareRoot()157                zw[i] = z[i] * root158                for j in 0..<k {159                    xw[j * n + i] = x[j * n + i] * root160                }161            }162            let qr: LinearAlgebra.QR163            do {164                qr = try LinearAlgebra.QR(matrix: xw, rows: n, cols: k)165                beta = try qr.solve(rhs: zw)166            } catch {167                throw ZQStatsError("\(family.rawValue): weighted solve failed — \(error)")168            }169            finalQR = qr170171            // Update linear predictor and means.172            eta = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)173            mu = eta.map { inverseLink(family: family, eta: $0) }174            let newLogLikelihood = self.logLikelihood(y: y, mu: mu, family: family)175176            if newLogLikelihood.isNaN || newLogLikelihood.isInfinite {177                throw ZQStatsError(178                    "\(family.rawValue) failed to converge (perfect prediction?)"179                )180            }181            let done = abs(newLogLikelihood - logLikelihood)182                < 1e-13 * (abs(newLogLikelihood) + 0.1)183            logLikelihood = newLogLikelihood184            if done { break }185            if iteration == maximumIterations {186                throw ZQStatsError(187                    "\(family.rawValue) did not converge in \(maximumIterations) iterations"188                )189            }190        }191192        // Bread: (X'WX)⁻¹ with the weights re-evaluated AT the converged β.193        // The last IRLS factorization carries weights from the previous194        // iterate (off by ~√tolerance), which would contaminate standard195        // errors at ~1e-7 relative — visible against R fixtures.196        guard finalQR != nil else {197            throw ZQStatsError("\(family.rawValue): no iterations performed")198        }199        let bread: [Double]200        do {201            var xw = [Double](repeating: 0, count: n * k)202            for i in 0..<n {203                let derivative = meanDerivative(family: family, eta: eta[i], mean: mu[i])204                let varMu = varianceFunction(family: family, mean: mu[i])205                let root = (derivative * derivative / varMu).squareRoot()206                for j in 0..<k {207                    xw[j * n + i] = x[j * n + i] * root208                }209            }210            let informationQR = try LinearAlgebra.QR(matrix: xw, rows: n, cols: k)211            bread = try informationQR.crossProductInverse()212        } catch {213            throw ZQStatsError("\(family.rawValue): information matrix is singular")214        }215216        // Covariance.217        var clusterCount: Int? = nil218        let vce: [Double]219        switch variance {220        case .classical:221            vce = bread222223        case .hc0, .hc1:224            // Score sandwich: u_i = x_i (y−μ) μ′/V(μ). No df correction —225            // Stata's `robust` for ML.226            let meat = scoreMeat(227                x: x, y: y, mu: mu, eta: eta, n: n, k: k, family: family228            )229            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)230231        case .cluster(let groups):232            guard groups.count == n else {233                throw ZQStatsError("cluster variable has wrong length")234            }235            var scores: [Int: [Double]] = [:]236            for i in 0..<n {237                let scale = scoreScale(family: family, y: y[i], mu: mu[i], eta: eta[i])238                var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)239                for j in 0..<k {240                    u[j] += x[j * n + i] * scale241                }242                scores[groups[i]] = u243            }244            let g = scores.count245            guard g > 1 else {246                throw ZQStatsError("cluster variable must define at least 2 groups")247            }248            clusterCount = g249            var meat = [Double](repeating: 0, count: k * k)250            for u in scores.values {251                for j in 0..<k {252                    for i in 0..<k {253                        meat[j * k + i] += u[i] * u[j]254                    }255                }256            }257            // Stata ML cluster factor: G/(G−1).258            let scale = Double(g) / Double(g - 1)259            for index in meat.indices { meat[index] *= scale }260            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)261262        case .hc2, .hc3:263            fatalError("unreachable — rejected above")264        }265266        // Inference: z statistics (normal), Stata ML convention.267        let zCritical = ZQDistributions.normalQuantile(0.5 + confidenceLevel / 2)268        var coefficients: [ZQCoefficient] = []269        for j in 0..<k {270            let se = vce[j * k + j].squareRoot()271            let zStat = beta[j] / se272            coefficients.append(ZQCoefficient(273                name: names[j],274                estimate: beta[j],275                standardError: se,276                tStatistic: zStat,277                pValue: 2 * (1 - ZQDistributions.normalCDF(abs(zStat))),278                confidenceLower: beta[j] - zCritical * se,279                confidenceUpper: beta[j] + zCritical * se280            ))281        }282283        // Null (intercept-only) log likelihood — closed form: constant mean.284        let meanY = y.reduce(0, +) / Double(n)285        let nullMu = [Double](repeating: meanY, count: n)286        let nullLogLikelihood = self.logLikelihood(y: y, mu: nullMu, family: family)287288        let dfModel = k - (includeConstant ? 1 : 0)289        var chiSquared: Double? = nil290        var chiSquaredPValue: Double? = nil291        if dfModel > 0, includeConstant {292            switch variance {293            case .classical:294                let lr = 2 * (logLikelihood - nullLogLikelihood)295                chiSquared = lr296                chiSquaredPValue = ZQDistributions.chiSquarePValue(lr, df: Double(dfModel))297            default:298                // Wald χ² on the non-constant coefficients under the299                // robust/cluster VCE.300                var subVce = [Double](repeating: 0, count: dfModel * dfModel)301                var subBeta = [Double](repeating: 0, count: dfModel)302                for j in 0..<dfModel {303                    subBeta[j] = beta[j]304                    for i in 0..<dfModel {305                        subVce[j * dfModel + i] = vce[j * k + i]306                    }307                }308                if let solved = try? LinearAlgebra.solveSymmetric(309                    subVce, k: dfModel, rhs: subBeta310                ) {311                    let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }312                    chiSquared = wald313                    chiSquaredPValue = ZQDistributions.chiSquarePValue(314                        wald, df: Double(dfModel)315                    )316                }317            }318        }319320        return ZQGLMResult(321            family: family,322            coefficients: coefficients,323            observationCount: n,324            logLikelihood: logLikelihood,325            nullLogLikelihood: nullLogLikelihood,326            pseudoRSquared: 1 - logLikelihood / nullLogLikelihood,327            chiSquared: chiSquared,328            chiSquaredDF: dfModel,329            chiSquaredPValue: chiSquaredPValue,330            iterations: iterations,331            clusterCount: clusterCount,332            fittedMeans: mu,333            vce: vce334        )335    }336337    // MARK: - Family functions338339    private static func startingMeans(y: [Double], family: ZQGLMFamily) -> [Double] {340        switch family {341        case .logit, .probit:342            return y.map { ($0 + 0.5) / 2 }343        case .poisson:344            return y.map { max($0, 0.1) }345        }346    }347348    private static func link(family: ZQGLMFamily, mean: Double) -> Double {349        switch family {350        case .logit: return Foundation.log(mean / (1 - mean))351        case .probit: return ZQDistributions.normalQuantile(mean)352        case .poisson: return Foundation.log(mean)353        }354    }355356    private static func inverseLink(family: ZQGLMFamily, eta: Double) -> Double {357        switch family {358        case .logit:359            let clamped = min(max(eta, -30), 30)360            return 1 / (1 + Foundation.exp(-clamped))361        case .probit:362            let p = ZQDistributions.normalCDF(eta)363            return min(max(p, 1e-12), 1 - 1e-12)364        case .poisson:365            return Foundation.exp(min(eta, 300))366        }367    }368369    /// dμ/dη at the current point.370    private static func meanDerivative(371        family: ZQGLMFamily, eta: Double, mean: Double372    ) -> Double {373        switch family {374        case .logit:375            return max(mean * (1 - mean), 1e-12)376        case .probit:377            return max(normalDensity(eta), 1e-12)378        case .poisson:379            return max(mean, 1e-12)380        }381    }382383    /// Var(Y | μ) up to the (unit) dispersion.384    private static func varianceFunction(family: ZQGLMFamily, mean: Double) -> Double {385        switch family {386        case .logit, .probit:387            return max(mean * (1 - mean), 1e-12)388        case .poisson:389            return max(mean, 1e-12)390        }391    }392393    /// Per-observation score scale: (y−μ)·μ′/V(μ), so the score is that394    /// times x_i. For canonical links this collapses to y−μ.395    private static func scoreScale(396        family: ZQGLMFamily, y: Double, mu: Double, eta: Double397    ) -> Double {398        switch family {399        case .logit, .poisson:400            return y - mu401        case .probit:402            let density = normalDensity(eta)403            return (y - mu) * density / max(mu * (1 - mu), 1e-12)404        }405    }406407    private static func scoreMeat(408        x: [Double], y: [Double], mu: [Double], eta: [Double],409        n: Int, k: Int, family: ZQGLMFamily410    ) -> [Double] {411        var meat = [Double](repeating: 0, count: k * k)412        for row in 0..<n {413            let scale = scoreScale(family: family, y: y[row], mu: mu[row], eta: eta[row])414            for j in 0..<k {415                let xj = x[j * n + row] * scale416                for i in j..<k {417                    meat[j * k + i] += xj * x[i * n + row] * scale418                }419            }420        }421        for j in 0..<k {422            for i in 0..<j {423                meat[j * k + i] = meat[i * k + j]424            }425        }426        return meat427    }428429    private static func logLikelihood(430        y: [Double], mu: [Double], family: ZQGLMFamily431    ) -> Double {432        var total = 0.0433        switch family {434        case .logit, .probit:435            for i in 0..<y.count {436                let p = min(max(mu[i], 1e-12), 1 - 1e-12)437                total += y[i] == 1 ? Foundation.log(p) : Foundation.log(1 - p)438            }439        case .poisson:440            for i in 0..<y.count {441                let m = max(mu[i], 1e-300)442                total += y[i] * Foundation.log(m) - m - ZQDistributions.logGamma(y[i] + 1)443            }444        }445        return total446    }447448    private static func normalDensity(_ z: Double) -> Double {449        Foundation.exp(-0.5 * z * z) / (2 * Double.pi).squareRoot()450    }451}452