spb/metrika Public

1//2//  ElasticNet.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Elastic-net linear regression by cyclic coordinate descent, following13/// glmnet's conventions exactly so results are cross-checkable:14///15///   min (1/2n)·Σᵢ(yᵢ − b₀ − xᵢ'β)² + λ·[α‖β‖₁ + ((1−α)/2)‖β‖₂²]16///17/// Predictors are standardized internally (mean 0, variance 1 with the18/// 1/n denominator — glmnet's `standardize = TRUE`), the penalty applies19/// on the standardized scale, and coefficients are reported back on the20/// original scale with an unpenalized intercept.21public struct ZQElasticNetResult: Equatable, Sendable {22    public var coefficients: [(name: String, value: Double)]23    public var intercept: Double24    public var alpha: Double25    public var lambda: Double26    public var nonzeroCount: Int27    public var rSquared: Double28    public var iterations: Int2930    public static func == (lhs: ZQElasticNetResult, rhs: ZQElasticNetResult) -> Bool {31        lhs.coefficients.elementsEqual(rhs.coefficients, by: { $0 == $1 })32            && lhs.intercept == rhs.intercept33            && lhs.alpha == rhs.alpha34            && lhs.lambda == rhs.lambda35            && lhs.nonzeroCount == rhs.nonzeroCount36            && lhs.rSquared == rhs.rSquared37            && lhs.iterations == rhs.iterations38    }39}4041public enum ZQElasticNet {4243    /// - Parameters:44    ///   - alpha: 1 = lasso, 0 = ridge, in between = elastic net.45    ///   - lambda: penalty strength on glmnet's scale.46    public static func fit(47        y: [Double],48        predictors: [(name: String, values: [Double])],49        alpha: Double,50        lambda: Double,51        maximumIterations: Int = 100_000,52        tolerance: Double = 1e-1453    ) throws -> ZQElasticNetResult {54        let n = y.count55        let p = predictors.count56        guard p > 0 else { throw ZQStatsError("elasticnet: regressors required") }57        guard n > 1 else { throw ZQStatsError("elasticnet: insufficient observations") }58        guard (0...1).contains(alpha) else {59            throw ZQStatsError("elasticnet: alpha() must be in [0, 1]")60        }61        guard lambda >= 0 else {62            throw ZQStatsError("elasticnet: lambda() must be nonnegative")63        }6465        // Standardize (1/n variance, glmnet convention); center y.66        let meanY = y.reduce(0, +) / Double(n)67        var shifts = [Double](repeating: 0, count: p)68        var scales = [Double](repeating: 1, count: p)69        var standardized = [[Double]](repeating: [], count: p)70        for (j, column) in predictors.enumerated() {71            guard column.values.count == n else {72                throw ZQStatsError("regressor '\(column.name)' has wrong length")73            }74            let mean = column.values.reduce(0, +) / Double(n)75            let variance = column.values.reduce(0) {76                $0 + ($1 - mean) * ($1 - mean)77            } / Double(n)78            guard variance > 0 else {79                throw ZQStatsError("elasticnet: '\(column.name)' is constant")80            }81            shifts[j] = mean82            scales[j] = variance.squareRoot()83            standardized[j] = column.values.map { ($0 - mean) / scales[j] }84        }85        let centeredY = y.map { $0 - meanY }8687        // Cyclic coordinate descent with residual updates. On the88        // standardized scale each column has (1/n)Σx² = 1, so89        //   β_j ← S(ρ_j, λα) / (1 + λ₂),  ρ_j = (1/n)x_j'r + β_j.90        //91        // glmnet quirk, matched deliberately: for gaussian models glmnet92        // standardizes y internally, which leaves the L1 penalty invariant93        // but scales the effective ridge penalty by 1/sd(y) — so94        // λ₂ = λ(1−α)/sd(y) (1/n variance). Ridge solutions therefore95        // depend on the scale of y, exactly as glmnet's do.96        let varY = centeredY.reduce(0) { $0 + $1 * $1 } / Double(n)97        guard varY > 0 else {98            throw ZQStatsError("elasticnet: the response is constant")99        }100        var beta = [Double](repeating: 0, count: p)101        var residuals = centeredY102        let l1 = lambda * alpha103        let l2 = lambda * (1 - alpha) / varY.squareRoot()104        var iterations = 0105106        for iteration in 1...maximumIterations {107            iterations = iteration108            var maxChange = 0.0109            for j in 0..<p {110                var rho = beta[j]111                let column = standardized[j]112                var dot = 0.0113                for i in 0..<n { dot += column[i] * residuals[i] }114                rho += dot / Double(n)115116                let updated = softThreshold(rho, l1) / (1 + l2)117                let change = updated - beta[j]118                if change != 0 {119                    for i in 0..<n { residuals[i] -= change * column[i] }120                    beta[j] = updated121                    maxChange = max(maxChange, abs(change))122                }123            }124            if maxChange < tolerance * max(1, beta.map(abs).max() ?? 1) {125                break126            }127            if iteration == maximumIterations {128                throw ZQStatsError("elasticnet: did not converge — increase lambda or iterations")129            }130        }131132        // Back to the original scale.133        var coefficients: [(name: String, value: Double)] = []134        var intercept = meanY135        var nonzero = 0136        for j in 0..<p {137            let value = beta[j] / scales[j]138            coefficients.append((predictors[j].name, value))139            intercept -= shifts[j] * value140            if value != 0 { nonzero += 1 }141        }142143        // Fit on the original data.144        var rss = 0.0145        var tss = 0.0146        for i in 0..<n {147            var fitted = intercept148            for j in 0..<p { fitted += coefficients[j].value * predictors[j].values[i] }149            rss += (y[i] - fitted) * (y[i] - fitted)150            tss += (y[i] - meanY) * (y[i] - meanY)151        }152153        return ZQElasticNetResult(154            coefficients: coefficients,155            intercept: intercept,156            alpha: alpha,157            lambda: lambda,158            nonzeroCount: nonzero,159            rSquared: tss > 0 ? 1 - rss / tss : .nan,160            iterations: iterations161        )162    }163164    /// The smallest λ at which every coefficient is zero (for α > 0):165    /// λ_max = max_j |(1/n)·x_j'y| / α on the standardized scale.166    public static func lambdaMax(167        y: [Double], predictors: [(name: String, values: [Double])], alpha: Double168    ) -> Double {169        let n = y.count170        let meanY = y.reduce(0, +) / Double(n)171        var best = 0.0172        for column in predictors {173            let mean = column.values.reduce(0, +) / Double(n)174            let variance = column.values.reduce(0) {175                $0 + ($1 - mean) * ($1 - mean)176            } / Double(n)177            guard variance > 0 else { continue }178            let scale = variance.squareRoot()179            var dot = 0.0180            for i in 0..<n {181                dot += (column.values[i] - mean) / scale * (y[i] - meanY)182            }183            best = max(best, abs(dot) / Double(n))184        }185        return alpha > 0 ? best / alpha : best * 1000186    }187188    private static func softThreshold(_ value: Double, _ threshold: Double) -> Double {189        if value > threshold { return value - threshold }190        if value < -threshold { return value + threshold }191        return 0192    }193}194