SPB Git
spboucher.ai

spb/metrika

Public

Stata-class statistics, GPU-accelerated by Apple Silicon. Native Swift — no Electron, no Python runtime, no compromises.

macosswiftstatisticsgpu

feat(engine): predict with estimation state (e() analog)

- PredictorDefinition (column / indicator / product / constant) records
  how each fitted regressor is recomputed from the dataset; threaded
  through buildRegressionSample so factor expansions and interactions
  carry their recipes
- EstimationState stored after regress, logit/probit/poisson, and
  ivregress; cleared after xtreg (predict there needs the estimated u_i)
- predict newvar [, xb | residuals | pr | n]: evaluates over ALL current
  observations with missing propagation; kind-aware defaults (xb for
  ols/iv, pr for logit/probit, n for poisson) and statistic validation
- ZQFactorExpansion.expand now returns the numeric level alongside each
  indicator column
- tests: xb + residuals reconstruct the response exactly (plain and
  factor models), mean fitted probability/count equals ybar (score
  equations), option validation; 93 green

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed (Aug 5, 2026) parent c7a8845 Browse files

3 changed files +324 −6

3 changed files
Unified Split
modified MetrikaKit/Sources/ZQEngine/Session.swift +217 −4
@@ -40,6 +40,7 @@ public actor ZQSession {
40 private let pluginRegistry: ZQPluginRegistry 40 private let pluginRegistry: ZQPluginRegistry
41 private var logFileURL: URL? 41 private var logFileURL: URL?
42 private var scriptDepth = 0 42 private var scriptDepth = 0
43+ private var lastEstimation: EstimationState?
43 44
44 /// - Parameters: 45 /// - Parameters:
45 /// - discoverUserCommands: scan the commands directory for `.zyq` 46 /// - discoverUserCommands: scan the commands directory for `.zyq`
@@ -193,6 +194,7 @@ public actor ZQSession {
193 case "regress": return try handleRegress(command) 194 case "regress": return try handleRegress(command)
194 case "xtreg": return try handleXTReg(command) 195 case "xtreg": return try handleXTReg(command)
195 case "ivregress": return try handleIVRegress(command) 196 case "ivregress": return try handleIVRegress(command)
197+ case "predict": return try handlePredict(command)
196 case "logit": return try handleGLM(command, family: .logit) 198 case "logit": return try handleGLM(command, family: .logit)
197 case "probit": return try handleGLM(command, family: .probit) 199 case "probit": return try handleGLM(command, family: .probit)
198 case "poisson": return try handleGLM(command, family: .poisson) 200 case "poisson": return try handleGLM(command, family: .poisson)
@@ -531,12 +533,37 @@ public actor ZQSession {
531 533
532 // MARK: - Estimation 534 // MARK: - Estimation
533 535
536+ /// How a fitted regressor is recomputed from the working dataset —
537+ /// the estimation-state piece that lets `predict` (and later
538+ /// `margins`) evaluate xb on arbitrary observations.
539+ enum PredictorDefinition: Equatable, Sendable {
540+ case column(String)
541+ case indicator(variable: String, level: Double)
542+ case product([String])
543+ case constant
544+ }
545+
546+ /// Stata's e() analog: what the last estimation command fitted.
547+ struct EstimationState: Sendable {
548+ enum Kind: Equatable, Sendable {
549+ case ols, iv
550+ case glm(ZQGLMFamily)
551+ }
552+
553+ var kind: Kind
554+ var responseName: String
555+ /// Coefficients aligned with their recomputation recipes.
556+ var coefficients: [(name: String, value: Double, definition: PredictorDefinition)]
557+ }
558+
534 /// Assembled estimation sample after if/in restriction and listwise 559 /// Assembled estimation sample after if/in restriction and listwise
535 /// deletion (§6: explicit report line for dropped observations). 560 /// deletion (§6: explicit report line for dropped observations).
536 struct RegressionSample { 561 struct RegressionSample {
537 var responseName: String 562 var responseName: String
538 var y: [Double] 563 var y: [Double]
539 var predictors: [(name: String, values: [Double])] 564 var predictors: [(name: String, values: [Double])]
565+ /// Aligned with `predictors`.
566+ var definitions: [PredictorDefinition]
540 var clusterLabels: [Int]? 567 var clusterLabels: [Int]?
541 var droppedMissing: Int 568 var droppedMissing: Int
542 } 569 }
@@ -606,13 +633,16 @@ public actor ZQSession {
606 633
607 let y = keptRows.map { yAll[$0] } 634 let y = keptRows.map { yAll[$0] }
608 635
609− // Expand specs on the estimation sample (plan-time expansion). 636+ // Expand specs on the estimation sample (plan-time expansion),
637+ // recording each column's recomputation recipe for predict.
610 var predictors: [(name: String, values: [Double])] = [] 638 var predictors: [(name: String, values: [Double])] = []
639+ var definitions: [PredictorDefinition] = []
611 for source in sources { 640 for source in sources {
612 switch source.spec { 641 switch source.spec {
613 case .simple(let name): 642 case .simple(let name):
614 let (values, _) = try frame.requireNumeric(name) 643 let (values, _) = try frame.requireNumeric(name)
615 predictors.append((name, keptRows.map { values[$0] })) 644 predictors.append((name, keptRows.map { values[$0] }))
645+ definitions.append(.column(name))
616 case .factor(let op, let name): 646 case .factor(let op, let name):
617 let (values, missing) = try frame.requireNumeric(name) 647 let (values, missing) = try frame.requireNumeric(name)
618 switch op { 648 switch op {
@@ -625,31 +655,39 @@ public actor ZQSession {
625 guard !expanded.isEmpty else { 655 guard !expanded.isEmpty else {
626 throw ZQEngineError("factor variable '\(name)' has a single level") 656 throw ZQEngineError("factor variable '\(name)' has a single level")
627 } 657 }
628− predictors.append(contentsOf: expanded) 658+ for term in expanded {
659+ predictors.append((term.name, term.values))
660+ definitions.append(.indicator(variable: name, level: term.level))
661+ }
629 case "c": 662 case "c":
630 predictors.append((name, keptRows.map { values[$0] })) 663 predictors.append((name, keptRows.map { values[$0] }))
664+ definitions.append(.column(name))
631 default: 665 default:
632 throw ZQEngineError("unsupported factor operator '\(op).'") 666 throw ZQEngineError("unsupported factor operator '\(op).'")
633 } 667 }
634 case .interaction(let parts, _): 668 case .interaction(let parts, _):
635− // v0.1: continuous-by-continuous interactions only. 669+ // Continuous-by-continuous interactions only.
636 var label: [String] = [] 670 var label: [String] = []
671+ var factors: [String] = []
637 var product = [Double](repeating: 1, count: keptRows.count) 672 var product = [Double](repeating: 1, count: keptRows.count)
638 for part in parts { 673 for part in parts {
639 guard case .factor(let op, let name) = part, op == "c" else { 674 guard case .factor(let op, let name) = part, op == "c" else {
640 guard case .simple(let name) = part else { 675 guard case .simple(let name) = part else {
641− throw ZQEngineError("factor interactions beyond c.#c. arrive with xtreg (v0.2)") 676+ throw ZQEngineError("factor interactions beyond c.#c. are not supported yet")
642 } 677 }
643 let (values, _) = try frame.requireNumeric(name) 678 let (values, _) = try frame.requireNumeric(name)
644 for (j, row) in keptRows.enumerated() { product[j] *= values[row] } 679 for (j, row) in keptRows.enumerated() { product[j] *= values[row] }
645 label.append(name) 680 label.append(name)
681+ factors.append(name)
646 continue 682 continue
647 } 683 }
648 let (values, _) = try frame.requireNumeric(name) 684 let (values, _) = try frame.requireNumeric(name)
649 for (j, row) in keptRows.enumerated() { product[j] *= values[row] } 685 for (j, row) in keptRows.enumerated() { product[j] *= values[row] }
650 label.append("c.\(name)") 686 label.append("c.\(name)")
687+ factors.append(name)
651 } 688 }
652 predictors.append((label.joined(separator: "#"), product)) 689 predictors.append((label.joined(separator: "#"), product))
690+ definitions.append(.product(factors))
653 } 691 }
654 } 692 }
655 693
@@ -689,11 +727,36 @@ public actor ZQSession {
689 responseName: responseName, 727 responseName: responseName,
690 y: y, 728 y: y,
691 predictors: predictors, 729 predictors: predictors,
730+ definitions: definitions,
692 clusterLabels: clusterLabels, 731 clusterLabels: clusterLabels,
693 droppedMissing: droppedMissing 732 droppedMissing: droppedMissing
694 ) 733 )
695 } 734 }
696 735
736+ /// Records the estimation state (Stata's e()) for post-estimation
737+ /// commands. The constant, when fitted, is the last coefficient.
738+ private func recordEstimation(
739+ kind: EstimationState.Kind,
740+ responseName: String,
741+ coefficients: [ZQCoefficient],
742+ definitions: [PredictorDefinition],
743+ includeConstant: Bool
744+ ) {
745+ var recipes = definitions
746+ if includeConstant { recipes.append(.constant) }
747+ guard recipes.count == coefficients.count else {
748+ lastEstimation = nil
749+ return
750+ }
751+ lastEstimation = EstimationState(
752+ kind: kind,
753+ responseName: responseName,
754+ coefficients: zip(coefficients, recipes).map {
755+ (name: $0.name, value: $0.estimate, definition: $1)
756+ }
757+ )
758+ }
759+
697 private func varianceEstimator( 760 private func varianceEstimator(
698 _ command: ZQCommand, clusterLabels: [Int]? 761 _ command: ZQCommand, clusterLabels: [Int]?
699 ) throws -> ZQVarianceEstimator { 762 ) throws -> ZQVarianceEstimator {
@@ -727,6 +790,13 @@ public actor ZQSession {
727 variance: variance, 790 variance: variance,
728 confidenceLevel: level 791 confidenceLevel: level
729 ) 792 )
793+ recordEstimation(
794+ kind: .ols,
795+ responseName: sample.responseName,
796+ coefficients: result.coefficients,
797+ definitions: sample.definitions,
798+ includeConstant: !command.hasOption("noconstant")
799+ )
730 800
731 var header: [String] = [] 801 var header: [String] = []
732 switch variance { 802 switch variance {
@@ -834,6 +904,8 @@ public actor ZQSession {
834 clustered: clustered, 904 clustered: clustered,
835 confidenceLevel: try confidenceLevel(command) 905 confidenceLevel: try confidenceLevel(command)
836 ) 906 )
907+ // predict after xtreg needs the estimated u_i — not stored yet.
908+ lastEstimation = nil
837 909
838 var header = ["Fixed-effects (within) regression"] 910 var header = ["Fixed-effects (within) regression"]
839 if sample.droppedMissing > 0 { 911 if sample.droppedMissing > 0 {
@@ -946,6 +1018,13 @@ public actor ZQSession {
946 variance: try varianceEstimator(command, clusterLabels: clusterLabels), 1018 variance: try varianceEstimator(command, clusterLabels: clusterLabels),
947 confidenceLevel: try confidenceLevel(command) 1019 confidenceLevel: try confidenceLevel(command)
948 ) 1020 )
1021+ recordEstimation(
1022+ kind: .iv,
1023+ responseName: response,
1024+ coefficients: result.coefficients,
1025+ definitions: (ivSpec.endogenous + exogenousNames).map { .column($0) },
1026+ includeConstant: !command.hasOption("noconstant")
1027+ )
949 1028
950 var header = ["Instrumental variables (2SLS) regression"] 1029 var header = ["Instrumental variables (2SLS) regression"]
951 if dropped > 0 { 1030 if dropped > 0 {
@@ -1041,6 +1120,13 @@ public actor ZQSession {
1041 variance: variance, 1120 variance: variance,
1042 confidenceLevel: level 1121 confidenceLevel: level
1043 ) 1122 )
1123+ recordEstimation(
1124+ kind: .glm(family),
1125+ responseName: sample.responseName,
1126+ coefficients: result.coefficients,
1127+ definitions: sample.definitions,
1128+ includeConstant: !command.hasOption("noconstant")
1129+ )
1044 1130
1045 let title: String 1131 let title: String
1046 switch family { 1132 switch family {
@@ -1305,6 +1391,133 @@ public actor ZQSession {
1305 return value / 100 1391 return value / 100
1306 } 1392 }
1307 1393
1394+ // MARK: - Post-estimation
1395+
1396+ /// `predict newvar [, xb | residuals | pr | n]` — evaluates the last
1397+ /// estimation over ALL current observations (missing inputs yield
1398+ /// missing predictions). Defaults: xb after regress/ivregress, pr
1399+ /// after logit/probit, n (mean count) after poisson.
1400+ private func handlePredict(_ command: ZQCommand) throws -> ZQResult {
1401+ guard let estimation = lastEstimation else {
1402+ throw ZQEngineError("predict: no estimation results — run a regression first")
1403+ }
1404+ guard !frame.isEmpty else { throw ZQEngineError("no data in memory") }
1405+ let names = command.varlist.flatMap(\.referencedNames)
1406+ guard names.count == 1, let target = names.first else {
1407+ throw ZQEngineError("predict: syntax is 'predict newvar [, statistic]'")
1408+ }
1409+
1410+ // Statistic selection and validation against the model kind.
1411+ enum Statistic { case xb, residuals, probability, mean }
1412+ var statistic: Statistic
1413+ switch estimation.kind {
1414+ case .ols, .iv: statistic = .xb
1415+ case .glm(.logit), .glm(.probit): statistic = .probability
1416+ case .glm(.poisson): statistic = .mean
1417+ }
1418+ for option in command.options {
1419+ switch option.name {
1420+ case "xb": statistic = .xb
1421+ case "residuals", "resid": statistic = .residuals
1422+ case "pr":
1423+ switch estimation.kind {
1424+ case .glm(.logit), .glm(.probit): statistic = .probability
1425+ default:
1426+ throw ZQEngineError("predict: 'pr' requires logit or probit results")
1427+ }
1428+ case "n":
1429+ guard case .glm(.poisson) = estimation.kind else {
1430+ throw ZQEngineError("predict: 'n' requires poisson results")
1431+ }
1432+ statistic = .mean
1433+ default:
1434+ throw ZQEngineError("predict: unknown statistic '\(option.name)'")
1435+ }
1436+ }
1437+
1438+ // Linear predictor over the full dataset.
1439+ let n = frame.rowCount
1440+ var xb = [Double](repeating: 0, count: n)
1441+ var missing = [Bool](repeating: false, count: n)
1442+ for coefficient in estimation.coefficients {
1443+ switch coefficient.definition {
1444+ case .constant:
1445+ for i in 0..<n { xb[i] += coefficient.value }
1446+ case .column(let name):
1447+ let (values, columnMissing) = try frame.requireNumeric(name)
1448+ for i in 0..<n {
1449+ if columnMissing[i] { missing[i] = true } else {
1450+ xb[i] += coefficient.value * values[i]
1451+ }
1452+ }
1453+ case .indicator(let variable, let level):
1454+ let (values, columnMissing) = try frame.requireNumeric(variable)
1455+ for i in 0..<n {
1456+ if columnMissing[i] { missing[i] = true }
1457+ else if values[i] == level { xb[i] += coefficient.value }
1458+ }
1459+ case .product(let variables):
1460+ let columns = try variables.map { try frame.requireNumeric($0) }
1461+ for i in 0..<n {
1462+ var product = 1.0
1463+ for column in columns {
1464+ if column.missing[i] { missing[i] = true; break }
1465+ product *= column.values[i]
1466+ }
1467+ if !missing[i] { xb[i] += coefficient.value * product }
1468+ }
1469+ }
1470+ }
1471+
1472+ // Transform.
1473+ var result = [Double](repeating: .nan, count: n)
1474+ var label: String
1475+ switch statistic {
1476+ case .xb:
1477+ label = "linear prediction"
1478+ for i in 0..<n where !missing[i] { result[i] = xb[i] }
1479+ case .residuals:
1480+ label = "residuals"
1481+ let (yValues, yMissing) = try frame.requireNumeric(estimation.responseName)
1482+ for i in 0..<n where !missing[i] && !yMissing[i] {
1483+ switch estimation.kind {
1484+ case .ols, .iv:
1485+ result[i] = yValues[i] - xb[i]
1486+ case .glm(.logit):
1487+ result[i] = yValues[i] - 1 / (1 + Foundation.exp(-xb[i]))
1488+ case .glm(.probit):
1489+ result[i] = yValues[i] - ZQDistributions.normalCDF(xb[i])
1490+ case .glm(.poisson):
1491+ result[i] = yValues[i] - Foundation.exp(xb[i])
1492+ }
1493+ }
1494+ for i in 0..<n where yMissing[i] { missing[i] = true }
1495+ case .probability:
1496+ label = "predicted probability"
1497+ for i in 0..<n where !missing[i] {
1498+ if case .glm(.probit) = estimation.kind {
1499+ result[i] = ZQDistributions.normalCDF(xb[i])
1500+ } else {
1501+ result[i] = 1 / (1 + Foundation.exp(-xb[i]))
1502+ }
1503+ }
1504+ case .mean:
1505+ label = "predicted mean count"
1506+ for i in 0..<n where !missing[i] { result[i] = Foundation.exp(xb[i]) }
1507+ }
1508+ for i in 0..<n where missing[i] { result[i] = .nan }
1509+
1510+ try frame.addColumn(ZQColumn(
1511+ name: target,
1512+ data: .float64(values: result, missing: missing)
1513+ ))
1514+ let missingCount = missing.count { $0 }
1515+ var note = "(\(label) '\(target)' generated"
1516+ if missingCount > 0 { note += ", \(missingCount) missing values" }
1517+ note += ")"
1518+ return ZQResult(text: note, scalars: ["N": Double(n - missingCount)])
1519+ }
1520+
1308 // MARK: - Bootstrap prefix 1521 // MARK: - Bootstrap prefix
1309 1522
1310 private func handleBootstrap( 1523 private func handleBootstrap(
modified MetrikaKit/Sources/ZQPlanner/Planner.swift +4 −2
@@ -90,9 +90,11 @@ public enum ZQFactorExpansion {
90 90
91 /// Expands one factor spec into named indicator columns against the 91 /// Expands one factor spec into named indicator columns against the
92 /// given data. The lowest level is the omitted base (Stata default). 92 /// given data. The lowest level is the omitted base (Stata default).
93+ /// The numeric level rides along so post-estimation commands can
94+ /// recompute the indicator on other observations.
93 public static func expand( 95 public static func expand(
94 name: String, values: [Double], missing: [Bool] 96 name: String, values: [Double], missing: [Bool]
95− ) -> [(name: String, values: [Double])] { 97+ ) -> [(name: String, level: Double, values: [Double])] {
96 let allLevels = levels(values: values, missing: missing) 98 let allLevels = levels(values: values, missing: missing)
97 guard allLevels.count > 1 else { return [] } 99 guard allLevels.count > 1 else { return [] }
98 return allLevels.dropFirst().map { level in 100 return allLevels.dropFirst().map { level in
@@ -102,7 +104,7 @@ public enum ZQFactorExpansion {
102 let indicator = values.enumerated().map { index, value in 104 let indicator = values.enumerated().map { index, value in
103 (!missing[index] && value == level) ? 1.0 : 0.0 105 (!missing[index] && value == level) ? 1.0 : 0.0
104 } 106 }
105− return (name: "\(rendered).\(name)", values: indicator) 107+ return (name: "\(rendered).\(name)", level: level, values: indicator)
106 } 108 }
107 } 109 }
108 } 110 }
added MetrikaKit/Tests/MetrikaKitTests/PredictTests.swift +103 −0
@@ -0,0 +1,103 @@
1+//
2+// PredictTests.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+
14+@Suite("predict", .serialized)
15+struct PredictTests {
16+ let fixtures: Fixtures
17+ let session: ZQSession
18+
19+ init() async throws {
20+ self.fixtures = try Fixtures()
21+ self.session = try ZQSession(discoverUserCommands: false)
22+ _ = try await session.execute("use \(fixtures.datasetURL.path)")
23+ _ = try await session.execute("gen log_rev = ln(revenue)")
24+ }
25+
26+ @Test("xb + residuals reconstruct the response after regress")
27+ func fittedPlusResiduals() async throws {
28+ _ = try await session.execute("reg log_rev price")
29+ _ = try await session.execute("predict yhat")
30+ _ = try await session.execute("predict e, residuals")
31+ _ = try await session.execute("gen check = yhat + e - log_rev")
32+ let summary = try await session.execute("summarize check")
33+ #expect(abs(try #require(summary.scalars["min"])) < 1e-12)
34+ #expect(abs(try #require(summary.scalars["max"])) < 1e-12)
35+ // Missing prices → missing predictions.
36+ #expect(summary.scalars["N"] == 57)
37+ _ = try await session.execute("drop yhat e check")
38+ }
39+
40+ @Test("predict evaluates factor-variable models on all observations")
41+ func factorModel() async throws {
42+ _ = try await session.execute("reg log_rev price i.region")
43+ _ = try await session.execute("predict fhat")
44+ _ = try await session.execute("predict fe2, residuals")
45+ _ = try await session.execute("gen fcheck = fhat + fe2 - log_rev")
46+ let summary = try await session.execute("summarize fcheck")
47+ #expect(abs(try #require(summary.scalars["max"])) < 1e-12)
48+ _ = try await session.execute("drop fhat fe2 fcheck")
49+ }
50+
51+ @Test("logit default is predicted probability; fitted mean matches ybar")
52+ func logitProbability() async throws {
53+ _ = try await session.execute("logit purchase price")
54+ _ = try await session.execute("predict p")
55+ let pSummary = try await session.execute("summarize p")
56+ let ySummary = try await session.execute("summarize purchase if !missing(price)")
57+ // With an intercept, mean fitted probability equals the sample
58+ // mean of the outcome (logit score equation).
59+ expectClose(
60+ try #require(pSummary.scalars["mean"]),
61+ try #require(ySummary.scalars["mean"]),
62+ rtol: 1e-9, "mean fitted probability"
63+ )
64+ #expect(try #require(pSummary.scalars["min"]) > 0)
65+ #expect(try #require(pSummary.scalars["max"]) < 1)
66+ _ = try await session.execute("drop p")
67+ }
68+
69+ @Test("poisson default is the predicted mean count")
70+ func poissonMean() async throws {
71+ _ = try await session.execute("poisson orders price")
72+ _ = try await session.execute("predict mu")
73+ let muSummary = try await session.execute("summarize mu")
74+ let ySummary = try await session.execute("summarize orders if !missing(price)")
75+ // Poisson with intercept: mean fitted count equals ybar.
76+ expectClose(
77+ try #require(muSummary.scalars["mean"]),
78+ try #require(ySummary.scalars["mean"]),
79+ rtol: 1e-9, "mean fitted count"
80+ )
81+ _ = try await session.execute("drop mu")
82+ }
83+
84+ @Test("predict before any estimation errors clearly")
85+ func requiresEstimation() async throws {
86+ let fresh = try ZQSession(discoverUserCommands: false)
87+ _ = try await fresh.execute("use \(fixtures.datasetURL.path)")
88+ await #expect(throws: ZQEngineError.self) {
89+ _ = try await fresh.execute("predict yhat")
90+ }
91+ }
92+
93+ @Test("statistic options are validated against the model kind")
94+ func statisticValidation() async throws {
95+ _ = try await session.execute("reg log_rev price")
96+ await #expect(throws: ZQEngineError.self) {
97+ _ = try await session.execute("predict p, pr") // pr needs logit/probit
98+ }
99+ await #expect(throws: ZQEngineError.self) {
100+ _ = try await session.execute("predict m, n") // n needs poisson
101+ }
102+ }
103+}
104
Menu

Theme
↑↓ navigate ↵ open esc close / or ⌘K anywhere