spb/metrika Public

1//2//  GradientBoosting.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Gradient-boosted regression trees with squared-error loss, following13/// xgboost's exact-greedy algorithm precisely so results cross-validate:14///15///   gain  = ½[G_L²/(H_L+λ) + G_R²/(H_R+λ) − (G_L+G_R)²/(H_L+H_R+λ)] − γ16///   leaf  = −G/(H+λ)17///18/// with g_i = ŷ_i − y_i and h_i = 1 for squared loss. Splits are placed19/// at midpoints between consecutive distinct feature values (xgboost's20/// convention), `x < split` goes left, and missing features follow the21/// default-left rule. No subsampling — training is deterministic.22public struct ZQBoostModel: Equatable, Sendable {23    public struct Node: Equatable, Sendable {24        /// Feature index for internal nodes; nil marks a leaf.25        public var feature: Int?26        public var split: Double27        public var left: Int28        public var right: Int29        public var value: Double   // leaf weight30    }3132    public struct Tree: Equatable, Sendable {33        public var nodes: [Node]3435        public func predict(_ features: [Double]) -> Double {36            var index = 037            while let feature = nodes[index].feature {38                let value = features[feature]39                index = value.isNaN || value < nodes[index].split40                    ? nodes[index].left41                    : nodes[index].right42            }43            return nodes[index].value44        }45    }4647    public var featureNames: [String]48    public var baseScore: Double49    public var learningRate: Double50    public var trees: [Tree]5152    /// Prediction for one observation's feature vector (ordered as53    /// `featureNames`; NaN = missing).54    public func predict(_ features: [Double]) -> Double {55        var result = baseScore56        for tree in trees {57            result += learningRate * tree.predict(features)58        }59        return result60    }61}6263public enum ZQGradientBoosting {6465    /// - Parameters:66    ///   - rounds: number of trees.67    ///   - learningRate: shrinkage η applied to every leaf.68    ///   - maxDepth: maximum tree depth (1 = stumps).69    ///   - lambda: L2 regularization on leaf weights (xgboost default 1).70    ///   - gamma: minimum gain to split (default 0).71    ///   - minChildWeight: minimum hessian sum per child (= observation72    ///     count under squared loss; default 1).73    ///   - baseScore: initial prediction; nil = mean of y.74    public static func fit(75        y: [Double],76        features: [(name: String, values: [Double])],77        rounds: Int,78        learningRate: Double = 0.3,79        maxDepth: Int = 6,80        lambda: Double = 1,81        gamma: Double = 0,82        minChildWeight: Double = 1,83        baseScore: Double? = nil84    ) throws -> ZQBoostModel {85        let n = y.count86        let p = features.count87        guard p > 0 else { throw ZQStatsError("boost: features required") }88        guard n > 1 else { throw ZQStatsError("boost: insufficient observations") }89        guard rounds > 0 else { throw ZQStatsError("boost: rounds() must be positive") }90        guard maxDepth >= 1 else { throw ZQStatsError("boost: maxdepth() must be ≥ 1") }91        for column in features where column.values.count != n {92            throw ZQStatsError("boost: feature '\(column.name)' has wrong length")93        }9495        // Pre-sorted feature indices (missing excluded; they default left).96        var sortedIndices = [[Int]](repeating: [], count: p)97        for j in 0..<p {98            let values = features[j].values99            sortedIndices[j] = (0..<n)100                .filter { !values[$0].isNaN }101                .sorted { values[$0] < values[$1] }102        }103104        let base = baseScore ?? y.reduce(0, +) / Double(n)105        var predictions = [Double](repeating: base, count: n)106        var trees: [ZQBoostModel.Tree] = []107108        for _ in 0..<rounds {109            // Squared loss: g = ŷ − y, h = 1.110            let gradients = (0..<n).map { predictions[$0] - y[$0] }111            var tree = ZQBoostModel.Tree(nodes: [])112            buildNode(113                rows: Array(0..<n),114                depth: 0,115                features: features,116                sortedIndices: sortedIndices,117                gradients: gradients,118                maxDepth: maxDepth,119                lambda: lambda,120                gamma: gamma,121                minChildWeight: minChildWeight,122                into: &tree.nodes123            )124            for i in 0..<n {125                let vector = (0..<p).map { features[$0].values[i] }126                predictions[i] += learningRate * tree.predict(vector)127            }128            trees.append(tree)129        }130131        return ZQBoostModel(132            featureNames: features.map(\.name),133            baseScore: base,134            learningRate: learningRate,135            trees: trees136        )137    }138139    /// Recursively grows one node; returns its index in `nodes`.140    @discardableResult141    private static func buildNode(142        rows: [Int],143        depth: Int,144        features: [(name: String, values: [Double])],145        sortedIndices: [[Int]],146        gradients: [Double],147        maxDepth: Int,148        lambda: Double,149        gamma: Double,150        minChildWeight: Double,151        into nodes: inout [ZQBoostModel.Node]152    ) -> Int {153        let rowSet = Set(rows)154        let gTotal = rows.reduce(0.0) { $0 + gradients[$1] }155        let hTotal = Double(rows.count)156157        let index = nodes.count158        nodes.append(ZQBoostModel.Node(159            feature: nil, split: 0, left: -1, right: -1,160            value: -gTotal / (hTotal + lambda)161        ))162        guard depth < maxDepth, rows.count > 1 else { return index }163164        // Exact greedy split search over every feature.165        let parentScore = gTotal * gTotal / (hTotal + lambda)166        var bestGain = 0.0167        var bestFeature = -1168        var bestSplit = 0.0169        for j in 0..<features.count {170            let values = features[j].values171            let ordered = sortedIndices[j].filter { rowSet.contains($0) }172            guard ordered.count > 1 else { continue }173            let gMissing = gTotal - ordered.reduce(0.0) { $0 + gradients[$1] }174            let hMissing = hTotal - Double(ordered.count)175176            var gLeft = gMissing   // missing rows follow the left child177            var hLeft = hMissing178            for position in 0..<(ordered.count - 1) {179                gLeft += gradients[ordered[position]]180                hLeft += 1181                let current = values[ordered[position]]182                let next = values[ordered[position + 1]]183                guard next > current else { continue }184                let gRight = gTotal - gLeft185                let hRight = hTotal - hLeft186                guard hLeft >= minChildWeight, hRight >= minChildWeight else { continue }187                let gain = 0.5 * (188                    gLeft * gLeft / (hLeft + lambda)189                        + gRight * gRight / (hRight + lambda)190                        - parentScore191                ) - gamma192                if gain > bestGain {193                    bestGain = gain194                    bestFeature = j195                    bestSplit = (current + next) / 2196                }197            }198        }199        guard bestFeature >= 0 else { return index }200201        let values = features[bestFeature].values202        let leftRows = rows.filter { values[$0].isNaN || values[$0] < bestSplit }203        let rightRows = rows.filter { !values[$0].isNaN && values[$0] >= bestSplit }204        guard !leftRows.isEmpty, !rightRows.isEmpty else { return index }205206        nodes[index].feature = bestFeature207        nodes[index].split = bestSplit208        nodes[index].left = buildNode(209            rows: leftRows, depth: depth + 1, features: features,210            sortedIndices: sortedIndices, gradients: gradients,211            maxDepth: maxDepth, lambda: lambda, gamma: gamma,212            minChildWeight: minChildWeight, into: &nodes213        )214        nodes[index].right = buildNode(215            rows: rightRows, depth: depth + 1, features: features,216            sortedIndices: sortedIndices, gradients: gradients,217            maxDepth: maxDepth, lambda: lambda, gamma: gamma,218            minChildWeight: minChildWeight, into: &nodes219        )220        return index221    }222}223