feat(stats): gradient-boosted regression trees (boost command)
- ZQGradientBoosting: exact-greedy trees cloning xgboost's algorithm — gain 1/2[GL2/(HL+l) + GR2/(HR+l) - G2/(H+l)] - gamma, leaf -G/(H+l), midpoint splits between consecutive distinct values, missing rows default left, pre-sorted feature indices; squared loss, deterministic (no subsampling) - engine: 'boost y x…, rounds(#) [eta() maxdepth() lambda()]' reporting training R2/RMSE with an in-sample caveat; model stored in the estimation state (Kind.boost) so predict routes through the trees (missing features follow the default direction); margins and GLM statistics refused after boost - validation: per-observation prediction parity with R xgboost 3.2 (exact method, base_score = mean) at 1e-4 (xgboost is float32 internally); a stump finds the exact midpoint split with lambda 0; training loss decreases monotonically in rounds - manual entry + coverage test - 115 tests green (swift test and xcodebuild with GPU suites) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
10 changed files +524 −1
modified
MetrikaKit/Sources/ZQEngine/CommandHelp.swift
+14 −0
@@ -218,6 +218,20 @@ public enum ZQCommandReference { | ||
| 218 | 218 | notes: "Matches R glmnet, including its gaussian y-standardization convention (ridge penalties scale with sd of the response)." |
| 219 | 219 | ), |
| 220 | 220 | |
| 221 | + ZQCommandDoc( | |
| 222 | + verb: "boost", category: "Machine learning", | |
| 223 | + summary: "Gradient-boosted regression trees (xgboost-style exact greedy, squared loss). Deterministic — no subsampling.", | |
| 224 | + syntax: "boost depvar features, rounds(#) [eta(#) maxdepth(#) lambda(#)]", | |
| 225 | + options: [ | |
| 226 | + ("rounds(#)", "number of trees"), | |
| 227 | + ("eta(#)", "learning rate, default 0.3"), | |
| 228 | + ("maxdepth(#)", "tree depth, default 6"), | |
| 229 | + ("lambda(#)", "L2 regularization on leaf weights, default 1"), | |
| 230 | + ], | |
| 231 | + examples: ["boost log_rev z1 z2 orders, rounds(100) eta(0.1) maxdepth(3)", "predict yhat"], | |
| 232 | + notes: "Matches R xgboost predictions on identical settings. Training R² is in-sample — expect it to be optimistic." | |
| 233 | + ), | |
| 234 | + | |
| 221 | 235 | // ------------------------------------------ Resampling & simulation |
| 222 | 236 | ZQCommandDoc( |
| 223 | 237 | verb: "bootstrap", category: "Resampling & simulation", |
modified
MetrikaKit/Sources/ZQEngine/Session.swift
+140 −0
@@ -197,6 +197,7 @@ public actor ZQSession { | ||
| 197 | 197 | case "predict": return try handlePredict(command) |
| 198 | 198 | case "margins": return try handleMargins(command) |
| 199 | 199 | case "elasticnet": return try handleElasticNet(command, defaultAlpha: nil) |
| 200 | + case "boost": return try handleBoost(command) | |
| 200 | 201 | case "lasso": return try handleElasticNet(command, defaultAlpha: 1) |
| 201 | 202 | case "logit": return try handleGLM(command, family: .logit) |
| 202 | 203 | case "probit": return try handleGLM(command, family: .probit) |
@@ -603,6 +604,7 @@ public actor ZQSession { | ||
| 603 | 604 | enum Kind: Equatable, Sendable { |
| 604 | 605 | case ols, iv |
| 605 | 606 | case glm(ZQGLMFamily) |
| 607 | + case boost(ZQBoostModel) | |
| 606 | 608 | } |
| 607 | 609 | |
| 608 | 610 | var kind: Kind |
@@ -1502,6 +1504,14 @@ public actor ZQSession { | ||
| 1502 | 1504 | throw ZQEngineError("predict: syntax is 'predict newvar [, statistic]'") |
| 1503 | 1505 | } |
| 1504 | 1506 | |
| 1507 | + // Boosted models predict through their trees, not a linear xb. | |
| 1508 | + if case .boost(let model) = estimation.kind { | |
| 1509 | + return try predictBoost( | |
| 1510 | + model: model, target: target, command: command, | |
| 1511 | + responseName: estimation.responseName | |
| 1512 | + ) | |
| 1513 | + } | |
| 1514 | + | |
| 1505 | 1515 | // Statistic selection and validation against the model kind. |
| 1506 | 1516 | enum Statistic { case xb, residuals, probability, mean } |
| 1507 | 1517 | var statistic: Statistic |
@@ -1509,6 +1519,7 @@ public actor ZQSession { | ||
| 1509 | 1519 | case .ols, .iv: statistic = .xb |
| 1510 | 1520 | case .glm(.logit), .glm(.probit): statistic = .probability |
| 1511 | 1521 | case .glm(.poisson): statistic = .mean |
| 1522 | + case .boost: fatalError("handled above") | |
| 1512 | 1523 | } |
| 1513 | 1524 | for option in command.options { |
| 1514 | 1525 | switch option.name { |
@@ -1584,6 +1595,8 @@ public actor ZQSession { | ||
| 1584 | 1595 | result[i] = yValues[i] - ZQDistributions.normalCDF(xb[i]) |
| 1585 | 1596 | case .glm(.poisson): |
| 1586 | 1597 | result[i] = yValues[i] - Foundation.exp(xb[i]) |
| 1598 | + case .boost: | |
| 1599 | + break // handled by predictBoost | |
| 1587 | 1600 | } |
| 1588 | 1601 | } |
| 1589 | 1602 | for i in 0..<n where yMissing[i] { missing[i] = true } |
@@ -1720,6 +1733,89 @@ public actor ZQSession { | ||
| 1720 | 1733 | return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars) |
| 1721 | 1734 | } |
| 1722 | 1735 | |
| 1736 | + /// `boost y x…, rounds(#) [eta(#) maxdepth(#) lambda(#)]` — | |
| 1737 | + /// gradient-boosted regression trees (xgboost exact-greedy, squared | |
| 1738 | + /// loss). Deterministic: no subsampling. | |
| 1739 | + private func handleBoost(_ command: ZQCommand) throws -> ZQResult { | |
| 1740 | + let sample = try buildRegressionSample(command, clusterVariable: nil) | |
| 1741 | + func numberOption(_ name: String, default defaultValue: Double) throws -> Double { | |
| 1742 | + guard let text = command.option(name)?.firstArgument else { return defaultValue } | |
| 1743 | + guard let value = Double(text), value > 0 else { | |
| 1744 | + throw ZQEngineError("boost: invalid \(name)()") | |
| 1745 | + } | |
| 1746 | + return value | |
| 1747 | + } | |
| 1748 | + guard let roundsText = command.option("rounds")?.firstArgument, | |
| 1749 | + let rounds = Int(roundsText), rounds > 0 else { | |
| 1750 | + throw ZQEngineError("boost: rounds(#) required") | |
| 1751 | + } | |
| 1752 | + let model = try ZQGradientBoosting.fit( | |
| 1753 | + y: sample.y, | |
| 1754 | + features: sample.predictors, | |
| 1755 | + rounds: rounds, | |
| 1756 | + learningRate: try numberOption("eta", default: 0.3), | |
| 1757 | + maxDepth: Int(try numberOption("maxdepth", default: 6)), | |
| 1758 | + lambda: try numberOption("lambda", default: 1) | |
| 1759 | + ) | |
| 1760 | + | |
| 1761 | + // Training fit. | |
| 1762 | + let n = sample.y.count | |
| 1763 | + var rss = 0.0 | |
| 1764 | + var tss = 0.0 | |
| 1765 | + let meanY = sample.y.reduce(0, +) / Double(n) | |
| 1766 | + for i in 0..<n { | |
| 1767 | + let vector = sample.predictors.map { $0.values[i] } | |
| 1768 | + let predicted = model.predict(vector) | |
| 1769 | + rss += (sample.y[i] - predicted) * (sample.y[i] - predicted) | |
| 1770 | + tss += (sample.y[i] - meanY) * (sample.y[i] - meanY) | |
| 1771 | + } | |
| 1772 | + let rSquared = tss > 0 ? 1 - rss / tss : .nan | |
| 1773 | + let rmse = (rss / Double(n)).squareRoot() | |
| 1774 | + | |
| 1775 | + recordEstimation( | |
| 1776 | + kind: .boost(model), | |
| 1777 | + responseName: sample.responseName, | |
| 1778 | + coefficients: sample.predictors.map { | |
| 1779 | + ZQCoefficient( | |
| 1780 | + name: $0.name, estimate: .nan, standardError: .nan, | |
| 1781 | + tStatistic: .nan, pValue: .nan, | |
| 1782 | + confidenceLower: .nan, confidenceUpper: .nan | |
| 1783 | + ) | |
| 1784 | + }, | |
| 1785 | + definitions: sample.definitions, | |
| 1786 | + includeConstant: false, | |
| 1787 | + sampleSize: n | |
| 1788 | + ) | |
| 1789 | + | |
| 1790 | + var lines = ["Gradient-boosted trees (squared loss)"] | |
| 1791 | + if sample.droppedMissing > 0 { | |
| 1792 | + lines.append("(\(sample.droppedMissing) observations dropped due to missing values)") | |
| 1793 | + } | |
| 1794 | + for (label, value) in [ | |
| 1795 | + ("Number of obs", "\(n)"), | |
| 1796 | + ("Trees", "\(rounds)"), | |
| 1797 | + ("Learning rate (eta)", TableFormatter.general(model.learningRate)), | |
| 1798 | + ("Max depth", command.option("maxdepth")?.firstArgument ?? "6"), | |
| 1799 | + ("Training R-squared", TableFormatter.fixed(rSquared, decimals: 4)), | |
| 1800 | + ("Training RMSE", TableFormatter.general(rmse, significant: 6)), | |
| 1801 | + ] { | |
| 1802 | + lines.append( | |
| 1803 | + TableFormatter.pad(label, 46, right: false) + "= " + | |
| 1804 | + TableFormatter.pad(value, 12) | |
| 1805 | + ) | |
| 1806 | + } | |
| 1807 | + lines.append("") | |
| 1808 | + lines.append("Use 'predict newvar' for fitted values. Training fit is in-sample —") | |
| 1809 | + lines.append("expect it to be optimistic relative to held-out data.") | |
| 1810 | + return ZQResult( | |
| 1811 | + text: lines.joined(separator: "\n"), | |
| 1812 | + scalars: [ | |
| 1813 | + "N": Double(n), "rounds": Double(rounds), | |
| 1814 | + "r2": rSquared, "rmse": rmse, | |
| 1815 | + ] | |
| 1816 | + ) | |
| 1817 | + } | |
| 1818 | + | |
| 1723 | 1819 | /// `margins, dydx(varlist)` — average marginal effects with |
| 1724 | 1820 | /// delta-method standard errors. OLS/IV effects are the coefficients |
| 1725 | 1821 | /// themselves; GLM effects average dμ/dx over the estimation sample: |
@@ -1763,6 +1859,8 @@ public actor ZQSession { | ||
| 1763 | 1859 | let beta = estimation.coefficients[j].value |
| 1764 | 1860 | |
| 1765 | 1861 | switch estimation.kind { |
| 1862 | + case .boost: | |
| 1863 | + throw ZQEngineError("margins: not available after boost") | |
| 1766 | 1864 | case .ols, .iv: |
| 1767 | 1865 | // Linear model: the marginal effect is the coefficient. |
| 1768 | 1866 | effects.append(Effect( |
@@ -1882,6 +1980,48 @@ public actor ZQSession { | ||
| 1882 | 1980 | return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars) |
| 1883 | 1981 | } |
| 1884 | 1982 | |
| 1983 | + /// Boosted-tree predictions: trees route missing features down the | |
| 1984 | + /// default (left) branch, so predictions exist for every observation. | |
| 1985 | + private func predictBoost( | |
| 1986 | + model: ZQBoostModel, target: String, command: ZQCommand, | |
| 1987 | + responseName: String | |
| 1988 | + ) throws -> ZQResult { | |
| 1989 | + for option in command.options | |
| 1990 | + where !["xb", "residuals", "resid"].contains(option.name) { | |
| 1991 | + throw ZQEngineError("predict: '\(option.name)' is not available after boost") | |
| 1992 | + } | |
| 1993 | + let wantsResiduals = command.hasOption("residuals") || command.hasOption("resid") | |
| 1994 | + | |
| 1995 | + let n = frame.rowCount | |
| 1996 | + let columns = try model.featureNames.map { try frame.requireNumeric($0) } | |
| 1997 | + var values = [Double](repeating: .nan, count: n) | |
| 1998 | + var missing = [Bool](repeating: false, count: n) | |
| 1999 | + var yValues = [Double](repeating: 0, count: n) | |
| 2000 | + var yMissing = [Bool](repeating: false, count: n) | |
| 2001 | + if wantsResiduals { | |
| 2002 | + (yValues, yMissing) = try frame.requireNumeric(responseName) | |
| 2003 | + } | |
| 2004 | + | |
| 2005 | + for i in 0..<n { | |
| 2006 | + let vector = columns.map { $0.missing[i] ? Double.nan : $0.values[i] } | |
| 2007 | + let predicted = model.predict(vector) | |
| 2008 | + if wantsResiduals { | |
| 2009 | + if yMissing[i] { | |
| 2010 | + missing[i] = true | |
| 2011 | + } else { | |
| 2012 | + values[i] = yValues[i] - predicted | |
| 2013 | + } | |
| 2014 | + } else { | |
| 2015 | + values[i] = predicted | |
| 2016 | + } | |
| 2017 | + } | |
| 2018 | + try frame.addColumn(ZQColumn( | |
| 2019 | + name: target, data: .float64(values: values, missing: missing) | |
| 2020 | + )) | |
| 2021 | + let label = wantsResiduals ? "residuals" : "boosted prediction" | |
| 2022 | + return ZQResult(text: "(\(label) '\(target)' generated)") | |
| 2023 | + } | |
| 2024 | + | |
| 1885 | 2025 | // MARK: - Bootstrap prefix |
| 1886 | 2026 | |
| 1887 | 2027 | private func handleBootstrap( |
modified
MetrikaKit/Sources/ZQParser/KnownVerbs.swift
+1 −0
@@ -45,6 +45,7 @@ public struct ZQVerbTable: Sendable { | ||
| 45 | 45 | "xtreg": 5, |
| 46 | 46 | "xtset": 5, |
| 47 | 47 | "areg": 4, |
| 48 | + "boost": 5, | |
| 48 | 49 | "elasticnet": 7, |
| 49 | 50 | "lasso": 5, |
| 50 | 51 | "bayes": 5, |
added
MetrikaKit/Sources/ZQStats/GradientBoosting.swift
+222 −0
@@ -0,0 +1,222 @@ | ||
| 1 | +// | |
| 2 | +// GradientBoosting.swift | |
| 3 | +// Metrika | |
| 4 | +// | |
| 5 | +// Author: Simon-Pierre Boucher | |
| 6 | +// Contact: contact@spboucher.ai | |
| 7 | +// Copyright © 2026 Simon-Pierre Boucher. All rights reserved. | |
| 8 | +// | |
| 9 | + | |
| 10 | +import Foundation | |
| 11 | + | |
| 12 | +/// Gradient-boosted regression trees with squared-error loss, following | |
| 13 | +/// 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 placed | |
| 19 | +/// at midpoints between consecutive distinct feature values (xgboost's | |
| 20 | +/// convention), `x < split` goes left, and missing features follow the | |
| 21 | +/// default-left rule. No subsampling — training is deterministic. | |
| 22 | +public 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: Double | |
| 27 | + public var left: Int | |
| 28 | + public var right: Int | |
| 29 | + public var value: Double // leaf weight | |
| 30 | + } | |
| 31 | + | |
| 32 | + public struct Tree: Equatable, Sendable { | |
| 33 | + public var nodes: [Node] | |
| 34 | + | |
| 35 | + public func predict(_ features: [Double]) -> Double { | |
| 36 | + var index = 0 | |
| 37 | + while let feature = nodes[index].feature { | |
| 38 | + let value = features[feature] | |
| 39 | + index = value.isNaN || value < nodes[index].split | |
| 40 | + ? nodes[index].left | |
| 41 | + : nodes[index].right | |
| 42 | + } | |
| 43 | + return nodes[index].value | |
| 44 | + } | |
| 45 | + } | |
| 46 | + | |
| 47 | + public var featureNames: [String] | |
| 48 | + public var baseScore: Double | |
| 49 | + public var learningRate: Double | |
| 50 | + public var trees: [Tree] | |
| 51 | + | |
| 52 | + /// Prediction for one observation's feature vector (ordered as | |
| 53 | + /// `featureNames`; NaN = missing). | |
| 54 | + public func predict(_ features: [Double]) -> Double { | |
| 55 | + var result = baseScore | |
| 56 | + for tree in trees { | |
| 57 | + result += learningRate * tree.predict(features) | |
| 58 | + } | |
| 59 | + return result | |
| 60 | + } | |
| 61 | +} | |
| 62 | + | |
| 63 | +public enum ZQGradientBoosting { | |
| 64 | + | |
| 65 | + /// - 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 (= observation | |
| 72 | + /// 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? = nil | |
| 84 | + ) throws -> ZQBoostModel { | |
| 85 | + let n = y.count | |
| 86 | + let p = features.count | |
| 87 | + 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 | + } | |
| 94 | + | |
| 95 | + // 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].values | |
| 99 | + sortedIndices[j] = (0..<n) | |
| 100 | + .filter { !values[$0].isNaN } | |
| 101 | + .sorted { values[$0] < values[$1] } | |
| 102 | + } | |
| 103 | + | |
| 104 | + let base = baseScore ?? y.reduce(0, +) / Double(n) | |
| 105 | + var predictions = [Double](repeating: base, count: n) | |
| 106 | + var trees: [ZQBoostModel.Tree] = [] | |
| 107 | + | |
| 108 | + 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.nodes | |
| 123 | + ) | |
| 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 | + } | |
| 130 | + | |
| 131 | + return ZQBoostModel( | |
| 132 | + featureNames: features.map(\.name), | |
| 133 | + baseScore: base, | |
| 134 | + learningRate: learningRate, | |
| 135 | + trees: trees | |
| 136 | + ) | |
| 137 | + } | |
| 138 | + | |
| 139 | + /// Recursively grows one node; returns its index in `nodes`. | |
| 140 | + @discardableResult | |
| 141 | + 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) | |
| 156 | + | |
| 157 | + let index = nodes.count | |
| 158 | + 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 } | |
| 163 | + | |
| 164 | + // Exact greedy split search over every feature. | |
| 165 | + let parentScore = gTotal * gTotal / (hTotal + lambda) | |
| 166 | + var bestGain = 0.0 | |
| 167 | + var bestFeature = -1 | |
| 168 | + var bestSplit = 0.0 | |
| 169 | + for j in 0..<features.count { | |
| 170 | + let values = features[j].values | |
| 171 | + 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) | |
| 175 | + | |
| 176 | + var gLeft = gMissing // missing rows follow the left child | |
| 177 | + var hLeft = hMissing | |
| 178 | + for position in 0..<(ordered.count - 1) { | |
| 179 | + gLeft += gradients[ordered[position]] | |
| 180 | + hLeft += 1 | |
| 181 | + let current = values[ordered[position]] | |
| 182 | + let next = values[ordered[position + 1]] | |
| 183 | + guard next > current else { continue } | |
| 184 | + let gRight = gTotal - gLeft | |
| 185 | + let hRight = hTotal - hLeft | |
| 186 | + guard hLeft >= minChildWeight, hRight >= minChildWeight else { continue } | |
| 187 | + let gain = 0.5 * ( | |
| 188 | + gLeft * gLeft / (hLeft + lambda) | |
| 189 | + + gRight * gRight / (hRight + lambda) | |
| 190 | + - parentScore | |
| 191 | + ) - gamma | |
| 192 | + if gain > bestGain { | |
| 193 | + bestGain = gain | |
| 194 | + bestFeature = j | |
| 195 | + bestSplit = (current + next) / 2 | |
| 196 | + } | |
| 197 | + } | |
| 198 | + } | |
| 199 | + guard bestFeature >= 0 else { return index } | |
| 200 | + | |
| 201 | + let values = features[bestFeature].values | |
| 202 | + 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 } | |
| 205 | + | |
| 206 | + nodes[index].feature = bestFeature | |
| 207 | + nodes[index].split = bestSplit | |
| 208 | + 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: &nodes | |
| 213 | + ) | |
| 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: &nodes | |
| 219 | + ) | |
| 220 | + return index | |
| 221 | + } | |
| 222 | +} | |
added
MetrikaKit/Tests/MetrikaKitTests/BoostTests.swift
+113 −0
@@ -0,0 +1,113 @@ | ||
| 1 | +// | |
| 2 | +// BoostTests.swift | |
| 3 | +// Metrika | |
| 4 | +// | |
| 5 | +// Author: Simon-Pierre Boucher | |
| 6 | +// Contact: contact@spboucher.ai | |
| 7 | +// Copyright © 2026 Simon-Pierre Boucher. All rights reserved. | |
| 8 | +// | |
| 9 | + | |
| 10 | +import Foundation | |
| 11 | +import Testing | |
| 12 | +import ZQEngine | |
| 13 | +import ZQGPU | |
| 14 | +import ZQStats | |
| 15 | + | |
| 16 | +/// Gradient boosting vs R xgboost (exact greedy, no subsampling — | |
| 17 | +/// deterministic). xgboost computes in float32, so per-observation | |
| 18 | +/// prediction agreement is asserted at 1e-4 relative. | |
| 19 | +@Suite("Gradient boosting", .serialized) | |
| 20 | +struct BoostTests { | |
| 21 | + let fixtures: Fixtures | |
| 22 | + let session: ZQSession | |
| 23 | + | |
| 24 | + init() async throws { | |
| 25 | + self.fixtures = try Fixtures() | |
| 26 | + self.session = try ZQSession(discoverUserCommands: false) | |
| 27 | + _ = try await session.execute("use \(fixtures.datasetURL.path)") | |
| 28 | + _ = try await session.execute("gen log_rev = ln(revenue)") | |
| 29 | + } | |
| 30 | + | |
| 31 | + @Test("matches xgboost predictions observation by observation") | |
| 32 | + func xgboostParity() async throws { | |
| 33 | + let result = try await session.execute( | |
| 34 | + "boost log_rev z1 z2 orders, rounds(20) eta(0.3) maxdepth(3) lambda(1)" | |
| 35 | + ) | |
| 36 | + #expect(result.scalars["N"] == 60) | |
| 37 | + expectClose( | |
| 38 | + try #require(result.scalars["rmse"]), | |
| 39 | + fixtures["gbm_rmse"], rtol: 1e-4, "training RMSE" | |
| 40 | + ) | |
| 41 | + | |
| 42 | + _ = try await session.execute("predict bhat") | |
| 43 | + let frame = await session.frame | |
| 44 | + let (predictions, _) = try frame.requireNumeric("bhat") | |
| 45 | + for (row, key) in [(0, "gbm_pred_1"), (16, "gbm_pred_17"), | |
| 46 | + (41, "gbm_pred_42"), (56, "gbm_pred_57")] { | |
| 47 | + expectClose( | |
| 48 | + predictions[row], fixtures[key], rtol: 1e-4, | |
| 49 | + "prediction[\(row + 1)]" | |
| 50 | + ) | |
| 51 | + } | |
| 52 | + _ = try await session.execute("drop bhat") | |
| 53 | + } | |
| 54 | + | |
| 55 | + @Test("a single stump finds the obvious split") | |
| 56 | + func stumpSplit() throws { | |
| 57 | + // Step function: y = 0 for x < 5, y = 10 for x ≥ 5. | |
| 58 | + let x = (0..<20).map(Double.init) | |
| 59 | + let y = x.map { $0 < 5 ? 0.0 : 10.0 } | |
| 60 | + let model = try ZQGradientBoosting.fit( | |
| 61 | + y: y, features: [("x", x)], | |
| 62 | + rounds: 1, learningRate: 1, maxDepth: 1, lambda: 0 | |
| 63 | + ) | |
| 64 | + let tree = model.trees[0] | |
| 65 | + #expect(tree.nodes[0].feature == 0) | |
| 66 | + #expect(tree.nodes[0].split == 4.5) | |
| 67 | + // Leaf values are exact means of the residuals under lambda 0. | |
| 68 | + #expect(abs(model.predict([2]) - 0) < 1e-12) | |
| 69 | + #expect(abs(model.predict([9]) - 10) < 1e-12) | |
| 70 | + } | |
| 71 | + | |
| 72 | + @Test("training loss decreases every round") | |
| 73 | + func monotoneTrainingLoss() throws { | |
| 74 | + var stream = PhiloxStream(seed: 5) | |
| 75 | + let n = 150 | |
| 76 | + let x1 = (0..<n).map { _ in stream.nextUniform() * 10 } | |
| 77 | + let x2 = (0..<n).map { _ in stream.nextUniform() * 10 } | |
| 78 | + let y = (0..<n).map { i in | |
| 79 | + Foundation.sin(x1[i]) + 0.3 * x2[i] + 0.2 * stream.nextNormal() | |
| 80 | + } | |
| 81 | + | |
| 82 | + var previous = Double.infinity | |
| 83 | + for rounds in [1, 5, 20, 80] { | |
| 84 | + let model = try ZQGradientBoosting.fit( | |
| 85 | + y: y, features: [("x1", x1), ("x2", x2)], | |
| 86 | + rounds: rounds, learningRate: 0.3, maxDepth: 3 | |
| 87 | + ) | |
| 88 | + var rss = 0.0 | |
| 89 | + for i in 0..<n { | |
| 90 | + let predicted = model.predict([x1[i], x2[i]]) | |
| 91 | + rss += (y[i] - predicted) * (y[i] - predicted) | |
| 92 | + } | |
| 93 | + #expect(rss < previous, "rounds \(rounds): rss \(rss)") | |
| 94 | + previous = rss | |
| 95 | + } | |
| 96 | + // Deep boosting on smooth data should fit very well in-sample. | |
| 97 | + #expect(previous / Double(n) < 0.05) | |
| 98 | + } | |
| 99 | + | |
| 100 | + @Test("boost requires rounds() and margins refuses afterwards") | |
| 101 | + func validation() async throws { | |
| 102 | + await #expect(throws: ZQEngineError.self) { | |
| 103 | + _ = try await session.execute("boost log_rev z1") | |
| 104 | + } | |
| 105 | + _ = try await session.execute("boost log_rev z1 z2, rounds(5)") | |
| 106 | + await #expect(throws: ZQEngineError.self) { | |
| 107 | + _ = try await session.execute("margins, dydx(z1)") | |
| 108 | + } | |
| 109 | + await #expect(throws: ZQEngineError.self) { | |
| 110 | + _ = try await session.execute("predict p, pr") | |
| 111 | + } | |
| 112 | + } | |
| 113 | +} | |
modified
MetrikaKit/Tests/MetrikaKitTests/EngineTests.swift
+1 −1
@@ -205,7 +205,7 @@ struct EngineTests { | ||
| 205 | 205 | "use", "save", "clear", "describe", "list", "count", "generate", |
| 206 | 206 | "replace", "drop", "keep", "summarize", "tabulate", "correlate", |
| 207 | 207 | "regress", "logit", "probit", "poisson", "ivregress", "xtreg", |
| 208 | − "xtset", "lasso", "elasticnet", "bootstrap", "permute", "predict", | |
| 208 | + "xtset", "lasso", "elasticnet", "boost", "bootstrap", "permute", "predict", | |
| 209 | 209 | "margins", "bayes", "scatter", "histogram", "kdensity", "graph", "display", |
| 210 | 210 | "set", "log", "help", "zscore", |
| 211 | 211 | ] { |
modified
MetrikaKit/Tests/MetrikaKitTests/Fixtures/expected.tsv
+5 −0
@@ -88,6 +88,11 @@ enet_b_price -0.058669179277376367 | ||
| 88 | 88 | enet_b_z1 -0.014365113727136758 |
| 89 | 89 | enet_b_z2 9.5044046897303132e-05 |
| 90 | 90 | enet_b_orders 0 |
| 91 | +gbm_rmse 0.074954829650842755 | |
| 92 | +gbm_pred_1 2.6102960109710693 | |
| 93 | +gbm_pred_17 2.4943184852600098 | |
| 94 | +gbm_pred_42 3.2452964782714844 | |
| 95 | +gbm_pred_57 2.9414722919464111 | |
| 91 | 96 | corr_rev_price -0.88919780339930277 |
| 92 | 97 | corr_rev_logrev 0.98089325161602359 |
| 93 | 98 | dist_pchisq_3p8_1 0.94874741714263044 |
modified
MetrikaKit/Tests/MetrikaKitTests/Fixtures/regression_v117.dta
+0 −0
Binary file not shown.
modified
MetrikaKit/Tests/MetrikaKitTests/Fixtures/regression_v118.dta
+0 −0
Binary file not shown.
modified
Tests/Fixtures/generate.R
+28 −0
@@ -352,6 +352,34 @@ if (requireNamespace("glmnet", quietly = TRUE)) { | ||
| 352 | 352 | cat("glmnet not installed — skipping elastic-net fixtures\n") |
| 353 | 353 | } |
| 354 | 354 | |
| 355 | +# ------------------------------------------------------- gradient boosting | |
| 356 | +# xgboost exact-greedy at fixed settings, no subsampling — deterministic. | |
| 357 | +# Complete-case features only (missing-value default-direction search | |
| 358 | +# differs between implementations). | |
| 359 | +if (requireNamespace("xgboost", quietly = TRUE)) { | |
| 360 | + # Full sample: these variables have no missing values, so the engine's | |
| 361 | + # listwise deletion keeps all 60 rows. | |
| 362 | + Xgb <- as.matrix(data[, c("z1", "z2", "orders")]) | |
| 363 | + ygb <- log(data$revenue) | |
| 364 | + base <- mean(ygb) | |
| 365 | + booster <- xgboost::xgboost( | |
| 366 | + x = Xgb, y = ygb, | |
| 367 | + nrounds = 20, max_depth = 3, learning_rate = 0.3, | |
| 368 | + reg_lambda = 1, reg_alpha = 0, min_child_weight = 1, | |
| 369 | + subsample = 1, colsample_bytree = 1, | |
| 370 | + tree_method = "exact", base_score = base, | |
| 371 | + nthread = 1, verbosity = 0 | |
| 372 | + ) | |
| 373 | + pred <- predict(booster, Xgb) | |
| 374 | + emit("gbm_rmse", sqrt(mean((ygb - pred)^2))) | |
| 375 | + emit("gbm_pred_1", pred[1]) | |
| 376 | + emit("gbm_pred_17", pred[17]) | |
| 377 | + emit("gbm_pred_42", pred[42]) | |
| 378 | + emit("gbm_pred_57", pred[57]) | |
| 379 | +} else { | |
| 380 | + cat("xgboost not installed — skipping boosting fixtures\n") | |
| 381 | +} | |
| 382 | + | |
| 355 | 383 | # -------------------------------------------------------------- correlate |
| 356 | 384 | emit("corr_rev_price", cor(d$revenue, d$price)) |
| 357 | 385 | emit("corr_rev_logrev", cor(d$revenue, d$log_rev)) |
| 358 | 386 | |