SPB Git

spb/metrika Public

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

macosswiftstatisticsgpu
Swift 92.4% HTML 3.3% R 3% Shell 1.3%

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 6 days ago (Aug 5, 2026) parent c7a8845

Showing 3 changed files with +324 and −6

modified MetrikaKit/Sources/ZQEngine/Session.swift +217 −4
@@ -40,6 +40,7 @@ public actor ZQSession {
40 40 private let pluginRegistry: ZQPluginRegistry
41 41 private var logFileURL: URL?
42 42 private var scriptDepth = 0
43 + private var lastEstimation: EstimationState?
43 44
44 45 /// - Parameters:
45 46 /// - discoverUserCommands: scan the commands directory for `.zyq`
@@ -193,6 +194,7 @@ public actor ZQSession {
193 194 case "regress": return try handleRegress(command)
194 195 case "xtreg": return try handleXTReg(command)
195 196 case "ivregress": return try handleIVRegress(command)
197 + case "predict": return try handlePredict(command)
196 198 case "logit": return try handleGLM(command, family: .logit)
197 199 case "probit": return try handleGLM(command, family: .probit)
198 200 case "poisson": return try handleGLM(command, family: .poisson)
@@ -531,12 +533,37 @@ public actor ZQSession {
531 533
532 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 559 /// Assembled estimation sample after if/in restriction and listwise
535 560 /// deletion (§6: explicit report line for dropped observations).
536 561 struct RegressionSample {
537 562 var responseName: String
538 563 var y: [Double]
539 564 var predictors: [(name: String, values: [Double])]
565 + /// Aligned with `predictors`.
566 + var definitions: [PredictorDefinition]
540 567 var clusterLabels: [Int]?
541 568 var droppedMissing: Int
542 569 }
@@ -606,13 +633,16 @@ public actor ZQSession {
606 633
607 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 638 var predictors: [(name: String, values: [Double])] = []
639 + var definitions: [PredictorDefinition] = []
611 640 for source in sources {
612 641 switch source.spec {
613 642 case .simple(let name):
614 643 let (values, _) = try frame.requireNumeric(name)
615 644 predictors.append((name, keptRows.map { values[$0] }))
645 + definitions.append(.column(name))
616 646 case .factor(let op, let name):
617 647 let (values, missing) = try frame.requireNumeric(name)
618 648 switch op {
@@ -625,31 +655,39 @@ public actor ZQSession {
625 655 guard !expanded.isEmpty else {
626 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 662 case "c":
630 663 predictors.append((name, keptRows.map { values[$0] }))
664 + definitions.append(.column(name))
631 665 default:
632 666 throw ZQEngineError("unsupported factor operator '\(op).'")
633 667 }
634 668 case .interaction(let parts, _):
635 − // v0.1: continuous-by-continuous interactions only.
669 + // Continuous-by-continuous interactions only.
636 670 var label: [String] = []
671 + var factors: [String] = []
637 672 var product = [Double](repeating: 1, count: keptRows.count)
638 673 for part in parts {
639 674 guard case .factor(let op, let name) = part, op == "c" else {
640 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 678 let (values, _) = try frame.requireNumeric(name)
644 679 for (j, row) in keptRows.enumerated() { product[j] *= values[row] }
645 680 label.append(name)
681 + factors.append(name)
646 682 continue
647 683 }
648 684 let (values, _) = try frame.requireNumeric(name)
649 685 for (j, row) in keptRows.enumerated() { product[j] *= values[row] }
650 686 label.append("c.\(name)")
687 + factors.append(name)
651 688 }
652 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 727 responseName: responseName,
690 728 y: y,
691 729 predictors: predictors,
730 + definitions: definitions,
692 731 clusterLabels: clusterLabels,
693 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 760 private func varianceEstimator(
698 761 _ command: ZQCommand, clusterLabels: [Int]?
699 762 ) throws -> ZQVarianceEstimator {
@@ -727,6 +790,13 @@ public actor ZQSession {
727 790 variance: variance,
728 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 801 var header: [String] = []
732 802 switch variance {
@@ -834,6 +904,8 @@ public actor ZQSession {
834 904 clustered: clustered,
835 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 910 var header = ["Fixed-effects (within) regression"]
839 911 if sample.droppedMissing > 0 {
@@ -946,6 +1018,13 @@ public actor ZQSession {
946 1018 variance: try varianceEstimator(command, clusterLabels: clusterLabels),
947 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 1029 var header = ["Instrumental variables (2SLS) regression"]
951 1030 if dropped > 0 {
@@ -1041,6 +1120,13 @@ public actor ZQSession {
1041 1120 variance: variance,
1042 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 1131 let title: String
1046 1132 switch family {
@@ -1305,6 +1391,133 @@ public actor ZQSession {
1305 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 1521 // MARK: - Bootstrap prefix
1309 1522
1310 1523 private func handleBootstrap(
modified MetrikaKit/Sources/ZQPlanner/Planner.swift +4 −2
@@ -90,9 +90,11 @@ public enum ZQFactorExpansion {
90 90
91 91 /// Expands one factor spec into named indicator columns against the
92 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 95 public static func expand(
94 96 name: String, values: [Double], missing: [Bool]
95 − ) -> [(name: String, values: [Double])] {
97 + ) -> [(name: String, level: Double, values: [Double])] {
96 98 let allLevels = levels(values: values, missing: missing)
97 99 guard allLevels.count > 1 else { return [] }
98 100 return allLevels.dropFirst().map { level in
@@ -102,7 +104,7 @@ public enum ZQFactorExpansion {
102 104 let indicator = values.enumerated().map { index, value in
103 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