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(stats): xtreg fixed effects and ivregress 2sls

- parser: (endog = instruments) varlist groups -> ZQIVSpec; digit-led
  sub-commands (2sls) reassembled from number+identifier tokens via
  column adjacency
- ZQFixedEffects: within estimator with Stata conventions (add-back
  means, reported _cons, df = N-K-G), within R-squared, panel-clustered
  VCE with G/(G-1) and t on G-1 df; v0.2 restriction: cluster variable
  must equal the panel variable
- ZQIV: 2SLS via thin-Q projection of the instrument matrix (Z'Z never
  formed), residuals from original regressors, Stata 'small' inference,
  classical/HC1/cluster VCE built on projected regressors
- engine: xtreg (requires xtset + fe), ivregress 2sls with dedicated
  listwise deletion across depvar/exog/endog/instruments; shared
  coefficient-table renderer extracted
- fixtures: z1/z2 instrument columns (drawn after existing draws, earlier
  golden values bit-identical), manual within/2SLS algebra in R
- 75 tests green (swift test and xcodebuild with GPU suites)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 6 days ago (Aug 5, 2026) parent 02932d6

Showing 13 changed files with +1,044 and −73

modified MetrikaKit/Sources/ZQEngine/Session.swift +224 −0
@@ -126,6 +126,8 @@ public actor ZQSession {
126 126 case "count": return try handleCount(command)
127 127 case "list": return try handleList(command)
128 128 case "regress": return try handleRegress(command)
129 + case "xtreg": return try handleXTReg(command)
130 + case "ivregress": return try handleIVRegress(command)
129 131 case "logit": return try handleGLM(command, family: .logit)
130 132 case "probit": return try handleGLM(command, family: .probit)
131 133 case "poisson": return try handleGLM(command, family: .poisson)
@@ -729,6 +731,228 @@ public actor ZQSession {
729 731 return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
730 732 }
731 733
734 + private func handleXTReg(_ command: ZQCommand) throws -> ZQResult {
735 + guard command.hasOption("fe") else {
736 + throw ZQEngineError("xtreg: only the fixed-effects estimator is implemented — add ', fe'")
737 + }
738 + guard let panelVariable else {
739 + throw ZQEngineError("xtreg: declare the panel first with 'xtset panelvar [timevar]'")
740 + }
741 + // v0.2 restriction: the cluster variable must be the panel variable.
742 + var clustered = false
743 + if let clusterName = command.option("cluster")?.firstArgument
744 + ?? clusterFromVCE(command) {
745 + guard clusterName == panelVariable else {
746 + throw ZQEngineError(
747 + "xtreg: clustering on a variable other than the panel (\(panelVariable)) is not supported yet"
748 + )
749 + }
750 + clustered = true
751 + }
752 +
753 + // Panel codes ride the same listwise-deletion path as cluster codes.
754 + let sample = try buildRegressionSample(command, clusterVariable: panelVariable)
755 + guard let panelCodes = sample.clusterLabels else {
756 + throw ZQEngineError("xtreg: panel variable missing from sample")
757 + }
758 +
759 + let result = try ZQFixedEffects.fitWithin(
760 + y: sample.y,
761 + predictors: sample.predictors,
762 + groups: panelCodes,
763 + clustered: clustered,
764 + confidenceLevel: try confidenceLevel(command)
765 + )
766 +
767 + var header = ["Fixed-effects (within) regression"]
768 + if sample.droppedMissing > 0 {
769 + header.append("(\(sample.droppedMissing) observations dropped due to missing values)")
770 + }
771 + var stats: [(String, String)] = [
772 + ("Number of obs", "\(result.observationCount)"),
773 + ("Number of groups", "\(result.groupCount)"),
774 + ("R-squared (within)", TableFormatter.fixed(result.rSquaredWithin, decimals: 4)),
775 + ("Root MSE", TableFormatter.general(result.rootMSE, significant: 6)),
776 + ]
777 + if let f = result.fStatistic, let p = result.fPValue {
778 + stats.append(("F", TableFormatter.fixed(f, decimals: 2)))
779 + stats.append(("Prob > F", TableFormatter.fixed(p, decimals: 4)))
780 + }
781 + for (label, value) in stats {
782 + header.append(
783 + TableFormatter.pad(label, 46, right: false) + "= " +
784 + TableFormatter.pad(value, 12)
785 + )
786 + }
787 +
788 + var scalars: [String: Double] = [
789 + "N": Double(result.observationCount),
790 + "N_g": Double(result.groupCount),
791 + "r2_w": result.rSquaredWithin,
792 + "rmse": result.rootMSE,
793 + "df_r": result.inferenceDF,
794 + ]
795 + let lines = header + [""] + coefficientTable(
796 + response: sample.responseName,
797 + coefficients: result.coefficients,
798 + statisticLabel: "t",
799 + level: try confidenceLevel(command),
800 + scalars: &scalars
801 + )
802 + return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
803 + }
804 +
805 + private func handleIVRegress(_ command: ZQCommand) throws -> ZQResult {
806 + guard command.subverb == "2sls" else {
807 + throw ZQEngineError("ivregress: only the 2sls estimator is implemented")
808 + }
809 + guard let ivSpec = command.ivSpec else {
810 + throw ZQEngineError("ivregress: syntax is 'ivregress 2sls depvar [exog] (endog = instruments)'")
811 + }
812 + guard !frame.isEmpty else { throw ZQEngineError("no data in memory") }
813 + guard let first = command.varlist.first, case .simple(let response) = first else {
814 + throw ZQEngineError("ivregress: dependent variable required")
815 + }
816 + let exogenousNames = command.varlist.dropFirst().flatMap(\.referencedNames)
817 + let clusterVariable = command.option("cluster")?.firstArgument
818 + ?? clusterFromVCE(command)
819 +
820 + // Listwise deletion across every variable involved.
821 + let allNames = [response] + exogenousNames + ivSpec.endogenous + ivSpec.instruments
822 + let sources = try allNames.map { try frame.requireNumeric($0) }
823 + var clusterColumn: ZQColumn?
824 + if let clusterVariable {
825 + clusterColumn = try frame.requireColumn(clusterVariable)
826 + }
827 + let mask = try observationMask(command)
828 + var keptRows: [Int] = []
829 + var dropped = 0
830 + for i in 0..<frame.rowCount where mask[i] {
831 + var complete = sources.allSatisfy { !$0.missing[i] }
832 + if complete, let clusterColumn {
833 + switch clusterColumn.data {
834 + case .float64(_, let missing): complete = !missing[i]
835 + case .string(let values): complete = values[i] != nil
836 + }
837 + }
838 + if complete { keptRows.append(i) } else { dropped += 1 }
839 + }
840 + guard !keptRows.isEmpty else { throw ZQEngineError("no observations") }
841 +
842 + func gather(_ names: [String]) throws -> [(name: String, values: [Double])] {
843 + try names.map { name in
844 + let (values, _) = try frame.requireNumeric(name)
845 + return (name, keptRows.map { values[$0] })
846 + }
847 + }
848 + let y = try frame.requireNumeric(response).values
849 + let clusterLabels: [Int]? = try clusterVariable.map { name in
850 + let column = try frame.requireColumn(name)
851 + var mapping: [String: Int] = [:]
852 + var codes: [Int] = []
853 + for row in keptRows {
854 + let key: String
855 + switch column.data {
856 + case .float64(let values, _): key = "\(values[row])"
857 + case .string(let values): key = values[row] ?? ""
858 + }
859 + if let code = mapping[key] {
860 + codes.append(code)
861 + } else {
862 + mapping[key] = mapping.count
863 + codes.append(mapping.count - 1)
864 + }
865 + }
866 + return codes
867 + }
868 +
869 + let result = try ZQIV.fit2SLS(
870 + y: keptRows.map { y[$0] },
871 + endogenous: try gather(ivSpec.endogenous),
872 + exogenous: try gather(exogenousNames),
873 + instruments: try gather(ivSpec.instruments),
874 + includeConstant: !command.hasOption("noconstant"),
875 + variance: try varianceEstimator(command, clusterLabels: clusterLabels),
876 + confidenceLevel: try confidenceLevel(command)
877 + )
878 +
879 + var header = ["Instrumental variables (2SLS) regression"]
880 + if dropped > 0 {
881 + header.append("(\(dropped) observations dropped due to missing values)")
882 + }
883 + header.append("Instrumented: \(ivSpec.endogenous.joined(separator: " "))")
884 + header.append("Instruments: \(ivSpec.instruments.joined(separator: " "))")
885 + var stats: [(String, String)] = [
886 + ("Number of obs", "\(result.observationCount)"),
887 + ("R-squared", TableFormatter.fixed(result.rSquared, decimals: 4)),
888 + ("Root MSE", TableFormatter.general(result.rootMSE, significant: 6)),
889 + ]
890 + if let f = result.fStatistic, let p = result.fPValue {
891 + stats.insert(("F", TableFormatter.fixed(f, decimals: 2)), at: 1)
892 + stats.insert(("Prob > F", TableFormatter.fixed(p, decimals: 4)), at: 2)
893 + }
894 + if let g = result.clusterCount { stats.append(("Clusters", "\(g)")) }
895 + for (label, value) in stats {
896 + header.append(
897 + TableFormatter.pad(label, 46, right: false) + "= " +
898 + TableFormatter.pad(value, 12)
899 + )
900 + }
901 +
902 + var scalars: [String: Double] = [
903 + "N": Double(result.observationCount),
904 + "r2": result.rSquared,
905 + "rmse": result.rootMSE,
906 + "df_r": result.inferenceDF,
907 + ]
908 + if let f = result.fStatistic { scalars["F"] = f }
909 + let lines = header + [""] + coefficientTable(
910 + response: response,
911 + coefficients: result.coefficients,
912 + statisticLabel: "t",
913 + level: try confidenceLevel(command),
914 + scalars: &scalars
915 + )
916 + return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
917 + }
918 +
919 + /// Shared coefficient-table rendering (regress/xtreg/ivregress/GLMs).
920 + private func coefficientTable(
921 + response: String,
922 + coefficients: [ZQCoefficient],
923 + statisticLabel: String,
924 + level: Double,
925 + scalars: inout [String: Double]
926 + ) -> [String] {
927 + let widths = [12, 12, 11, 8, 8, 22]
928 + var lines = [
929 + TableFormatter.pad(response, widths[0]) + " | " +
930 + TableFormatter.pad("Coefficient", widths[1]) + " " +
931 + TableFormatter.pad("Std. err.", widths[2]) + " " +
932 + TableFormatter.pad(statisticLabel, widths[3]) + " " +
933 + TableFormatter.pad("P>|\(statisticLabel)|", widths[4]) + " " +
934 + TableFormatter.pad("[\(Int(level * 100))% conf. interval]", widths[5]),
935 + TableFormatter.rule(widths),
936 + ]
937 + for coefficient in coefficients {
938 + lines.append(
939 + TableFormatter.pad(coefficient.name, widths[0]) + " | " +
940 + TableFormatter.pad(TableFormatter.general(coefficient.estimate), widths[1]) + " " +
941 + TableFormatter.pad(TableFormatter.general(coefficient.standardError), widths[2]) + " " +
942 + TableFormatter.pad(TableFormatter.fixed(coefficient.tStatistic, decimals: 2), widths[3]) + " " +
943 + TableFormatter.pad(TableFormatter.fixed(coefficient.pValue, decimals: 3), widths[4]) + " " +
944 + TableFormatter.pad(
945 + TableFormatter.general(coefficient.confidenceLower) + " " +
946 + TableFormatter.general(coefficient.confidenceUpper),
947 + widths[5]
948 + )
949 + )
950 + scalars["b_\(coefficient.name)"] = coefficient.estimate
951 + scalars["se_\(coefficient.name)"] = coefficient.standardError
952 + }
953 + return lines
954 + }
955 +
732 956 private func handleGLM(
733 957 _ command: ZQCommand, family: ZQGLMFamily
734 958 ) throws -> ZQResult {
modified MetrikaKit/Sources/ZQParser/AST.swift +17 −0
@@ -84,6 +84,18 @@ public struct ZQOption: Equatable, Sendable {
84 84 public var firstArgument: String? { arguments.first }
85 85 }
86 86
87 +/// Instrumental-variables group in a varlist:
88 +/// `ivregress 2sls y exog (endog1 endog2 = instr1 instr2)`.
89 +public struct ZQIVSpec: Equatable, Sendable {
90 + public var endogenous: [String]
91 + public var instruments: [String]
92 +
93 + public init(endogenous: [String], instruments: [String]) {
94 + self.endogenous = endogenous
95 + self.instruments = instruments
96 + }
97 +}
98 +
87 99 /// Assignment payload for `gen`/`replace`-style commands:
88 100 /// `gen log_rev = ln(revenue)`.
89 101 public struct ZQAssignment: Equatable, Sendable {
@@ -105,6 +117,9 @@ public struct ZQCommand: Equatable, Sendable {
105 117 /// Sub-verb for compound commands (`graph scatter …` → "scatter").
106 118 public var subverb: String?
107 119 public var varlist: [ZQVarSpec]
120 + /// Instrumented group for IV estimators, parsed from
121 + /// `(endog… = instruments…)` inside the varlist.
122 + public var ivSpec: ZQIVSpec?
108 123 /// `gen`/`replace` assignment, mutually exclusive with a plain varlist.
109 124 public var assignment: ZQAssignment?
110 125 /// Raw argument such as a file path (`use sales.parquet`).
@@ -127,6 +142,7 @@ public struct ZQCommand: Equatable, Sendable {
127 142 verb: String,
128 143 subverb: String? = nil,
129 144 varlist: [ZQVarSpec] = [],
145 + ivSpec: ZQIVSpec? = nil,
130 146 assignment: ZQAssignment? = nil,
131 147 argument: String? = nil,
132 148 condition: ZQExpression? = nil,
@@ -138,6 +154,7 @@ public struct ZQCommand: Equatable, Sendable {
138 154 self.verb = verb
139 155 self.subverb = subverb
140 156 self.varlist = varlist
157 + self.ivSpec = ivSpec
141 158 self.assignment = assignment
142 159 self.argument = argument
143 160 self.condition = condition
modified MetrikaKit/Sources/ZQParser/Grammar/CommandParser.swift +73 −9
@@ -185,16 +185,32 @@ public struct ZQCommandParser: Sendable {
185 185 _ = cursor.advance()
186 186 var command = ZQCommand(verb: verb)
187 187
188 − // Compound verbs: `graph scatter y x`
188 + // Compound verbs: `graph scatter y x`, `ivregress 2sls …`. A
189 + // digit-led sub-command like `2sls` lexes as number+identifier;
190 + // reassemble when the tokens are column-adjacent.
189 191 if verbTable.compoundVerbs.contains(verb) {
190 − guard case .identifier(let sub) = cursor.current.kind else {
192 + switch cursor.current.kind {
193 + case .identifier(let sub):
194 + _ = cursor.advance()
195 + command.subverb = sub
196 + case .number(let value) where value == value.rounded() && value >= 0:
197 + let numberToken = cursor.advance()
198 + let prefix = String(Int(value))
199 + guard case .identifier(let rest) = cursor.current.kind,
200 + cursor.current.column == numberToken.column + prefix.count else {
201 + throw ZQParseError(
202 + message: "'\(verb)' requires a sub-command",
203 + column: numberToken.column
204 + )
205 + }
206 + _ = cursor.advance()
207 + command.subverb = prefix + rest
208 + default:
191 209 throw ZQParseError(
192 210 message: "'\(verb)' requires a sub-command",
193 211 column: cursor.current.column
194 212 )
195 213 }
196 − _ = cursor.advance()
197 − command.subverb = sub
198 214 }
199 215
200 216 // `set seed 42` → subverb "seed", argument "42"
@@ -239,7 +255,9 @@ public struct ZQCommandParser: Sendable {
239 255 let expression = try cursor.parseExpression()
240 256 command.assignment = ZQAssignment(target: target, expression: expression)
241 257 } else {
242 − command.varlist = try parseVarlist(&cursor)
258 + let (varlist, ivSpec) = try parseVarlist(&cursor)
259 + command.varlist = varlist
260 + command.ivSpec = ivSpec
243 261 }
244 262
245 263 // Qualifiers may appear in any sensible order; Stata fixes the
@@ -264,13 +282,59 @@ public struct ZQCommandParser: Sendable {
264 282
265 283 private func parseVarlist(
266 284 _ cursor: inout ExpressionParser
267 − ) throws(ZQParseError) -> [ZQVarSpec] {
285 + ) throws(ZQParseError) -> ([ZQVarSpec], ZQIVSpec?) {
268 286 var specs: [ZQVarSpec] = []
287 + var ivSpec: ZQIVSpec?
288 + loop: while true {
289 + switch cursor.current.kind {
290 + case .identifier(let name):
291 + if name == "if" || name == "in" { break loop }
292 + specs.append(try parseVarSpec(&cursor))
293 + case .lparen:
294 + // IV group: (endog… = instruments…)
295 + guard ivSpec == nil else {
296 + throw ZQParseError(
297 + message: "only one (endogenous = instruments) group is allowed",
298 + column: cursor.current.column
299 + )
300 + }
301 + ivSpec = try parseIVGroup(&cursor)
302 + default:
303 + break loop
304 + }
305 + }
306 + return (specs, ivSpec)
307 + }
308 +
309 + private func parseIVGroup(
310 + _ cursor: inout ExpressionParser
311 + ) throws(ZQParseError) -> ZQIVSpec {
312 + try cursor.expect(.lparen)
313 + var endogenous: [String] = []
269 314 while case .identifier(let name) = cursor.current.kind {
270 − if name == "if" || name == "in" { break }
271 − specs.append(try parseVarSpec(&cursor))
315 + endogenous.append(name)
316 + _ = cursor.advance()
317 + }
318 + guard !endogenous.isEmpty else {
319 + throw ZQParseError(
320 + message: "expected endogenous variable names before '='",
321 + column: cursor.current.column
322 + )
323 + }
324 + try cursor.expect(.op("="))
325 + var instruments: [String] = []
326 + while case .identifier(let name) = cursor.current.kind {
327 + instruments.append(name)
328 + _ = cursor.advance()
329 + }
330 + guard !instruments.isEmpty else {
331 + throw ZQParseError(
332 + message: "expected instrument names after '='",
333 + column: cursor.current.column
334 + )
272 335 }
273 − return specs
336 + try cursor.expect(.rparen)
337 + return ZQIVSpec(endogenous: endogenous, instruments: instruments)
274 338 }
275 339
276 340 private func parseVarSpec(
modified MetrikaKit/Sources/ZQParser/KnownVerbs.swift +1 −1
@@ -68,7 +68,7 @@ public struct ZQVerbTable: Sendable {
68 68 fileVerbs: ["use", "save", "import", "export", "log"],
69 69 assignmentVerbs: ["generate", "replace", "egen"],
70 70 prefixVerbs: ["bootstrap", "permute", "jackknife"],
71 − compoundVerbs: ["graph"]
71 + compoundVerbs: ["graph", "ivregress"]
72 72 )
73 73
74 74 public init(
added MetrikaKit/Sources/ZQStats/FixedEffects.swift +217 −0
@@ -0,0 +1,217 @@
1 +//
2 +// FixedEffects.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 +/// Panel fixed-effects (within) estimator for `xtreg, fe`.
13 +///
14 +/// The within transformation subtracts group means and adds back grand
15 +/// means (Stata's convention, so `_cons` is reported), then solves by the
16 +/// usual QR path. Degrees of freedom absorb the G group effects:
17 +/// df = N − K − G. Cluster-robust VCE clusters on the panel variable with
18 +/// the G/(G−1) factor and t statistics on G−1 df.
19 +public struct ZQFEResult: Equatable, Sendable {
20 + public var coefficients: [ZQCoefficient]
21 + public var observationCount: Int
22 + public var groupCount: Int
23 + public var degreesOfFreedomResidual: Int
24 + public var inferenceDF: Double
25 + /// Within R²: fit of the demeaned regression.
26 + public var rSquaredWithin: Double
27 + public var rootMSE: Double
28 + public var fStatistic: Double?
29 + public var fPValue: Double?
30 + public var clustered: Bool
31 +
32 + public init(
33 + coefficients: [ZQCoefficient], observationCount: Int, groupCount: Int,
34 + degreesOfFreedomResidual: Int, inferenceDF: Double,
35 + rSquaredWithin: Double, rootMSE: Double,
36 + fStatistic: Double?, fPValue: Double?, clustered: Bool
37 + ) {
38 + self.coefficients = coefficients
39 + self.observationCount = observationCount
40 + self.groupCount = groupCount
41 + self.degreesOfFreedomResidual = degreesOfFreedomResidual
42 + self.inferenceDF = inferenceDF
43 + self.rSquaredWithin = rSquaredWithin
44 + self.rootMSE = rootMSE
45 + self.fStatistic = fStatistic
46 + self.fPValue = fPValue
47 + self.clustered = clustered
48 + }
49 +}
50 +
51 +public enum ZQFixedEffects {
52 +
53 + /// - Parameters:
54 + /// - groups: dense panel codes (0..<G), one per observation.
55 + /// - clustered: cluster the VCE on the panel variable.
56 + public static func fitWithin(
57 + y: [Double],
58 + predictors: [(name: String, values: [Double])],
59 + groups: [Int],
60 + clustered: Bool = false,
61 + confidenceLevel: Double = 0.95
62 + ) throws -> ZQFEResult {
63 + let n = y.count
64 + let slopes = predictors.count
65 + guard slopes > 0 else { throw ZQStatsError("xtreg: regressors required") }
66 + guard groups.count == n else {
67 + throw ZQStatsError("panel variable has wrong length")
68 + }
69 + let groupCount = Set(groups).count
70 + let k = slopes + 1 // slopes + reported constant
71 + let dfResidual = n - slopes - groupCount
72 + guard dfResidual > 0 else {
73 + throw ZQStatsError("insufficient observations: N=\(n), K=\(slopes), G=\(groupCount)")
74 + }
75 +
76 + // Group means and grand means.
77 + func withinTransform(_ values: [Double]) -> (transformed: [Double], demeaned: [Double]) {
78 + var sums = [Double](repeating: 0, count: groupCount)
79 + var counts = [Double](repeating: 0, count: groupCount)
80 + for i in 0..<n {
81 + sums[groups[i]] += values[i]
82 + counts[groups[i]] += 1
83 + }
84 + let grand = values.reduce(0, +) / Double(n)
85 + var transformed = [Double](repeating: 0, count: n)
86 + var demeaned = [Double](repeating: 0, count: n)
87 + for i in 0..<n {
88 + let groupMean = sums[groups[i]] / counts[groups[i]]
89 + demeaned[i] = values[i] - groupMean
90 + transformed[i] = demeaned[i] + grand
91 + }
92 + return (transformed, demeaned)
93 + }
94 +
95 + let (yTransformed, yDemeaned) = withinTransform(y)
96 + var design = [Double]()
97 + design.reserveCapacity(n * k)
98 + var demeanedColumns: [[Double]] = []
99 + for column in predictors {
100 + guard column.values.count == n else {
101 + throw ZQStatsError("regressor '\(column.name)' has wrong length")
102 + }
103 + let (transformed, demeaned) = withinTransform(column.values)
104 + design.append(contentsOf: transformed)
105 + demeanedColumns.append(demeaned)
106 + }
107 + design.append(contentsOf: [Double](repeating: 1, count: n))
108 +
109 + let qr: LinearAlgebra.QR
110 + let beta: [Double]
111 + let xtxInverse: [Double]
112 + do {
113 + qr = try LinearAlgebra.QR(matrix: design, rows: n, cols: k)
114 + beta = try qr.solve(rhs: yTransformed)
115 + xtxInverse = try qr.crossProductInverse()
116 + } catch {
117 + throw ZQStatsError("xtreg: design matrix is rank deficient after demeaning — \(error)")
118 + }
119 +
120 + // Within residuals and fit.
121 + let fitted = LinearAlgebra.multiply(matrix: design, rows: n, cols: k, vector: beta)
122 + var residuals = [Double](repeating: 0, count: n)
123 + var rss = 0.0
124 + for i in 0..<n {
125 + residuals[i] = yTransformed[i] - fitted[i]
126 + rss += residuals[i] * residuals[i]
127 + }
128 + let tssWithin = yDemeaned.reduce(0) { $0 + $1 * $1 }
129 + let sigma2 = rss / Double(dfResidual)
130 + let rSquaredWithin = tssWithin > 0 ? 1 - rss / tssWithin : .nan
131 +
132 + // Covariance.
133 + let vce: [Double]
134 + let inferenceDF: Double
135 + if clustered {
136 + var scores: [Int: [Double]] = [:]
137 + for i in 0..<n {
138 + var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)
139 + for j in 0..<slopes {
140 + u[j] += demeanedColumns[j][i] * residuals[i]
141 + }
142 + u[slopes] += residuals[i] // constant column
143 + scores[groups[i]] = u
144 + }
145 + guard groupCount > 1 else {
146 + throw ZQStatsError("cluster-robust VCE needs at least 2 groups")
147 + }
148 + var meat = [Double](repeating: 0, count: k * k)
149 + for u in scores.values {
150 + for j in 0..<k {
151 + for i in 0..<k {
152 + meat[j * k + i] += u[i] * u[j]
153 + }
154 + }
155 + }
156 + let scale = Double(groupCount) / Double(groupCount - 1)
157 + for index in meat.indices { meat[index] *= scale }
158 + vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)
159 + inferenceDF = Double(groupCount - 1)
160 + } else {
161 + vce = xtxInverse.map { $0 * sigma2 }
162 + inferenceDF = Double(dfResidual)
163 + }
164 +
165 + let tCritical = ZQDistributions.studentTQuantile(
166 + 0.5 + confidenceLevel / 2, df: inferenceDF
167 + )
168 + var names = predictors.map(\.name)
169 + names.append("_cons")
170 + var coefficients: [ZQCoefficient] = []
171 + for j in 0..<k {
172 + let se = vce[j * k + j].squareRoot()
173 + let t = beta[j] / se
174 + coefficients.append(ZQCoefficient(
175 + name: names[j],
176 + estimate: beta[j],
177 + standardError: se,
178 + tStatistic: t,
179 + pValue: ZQDistributions.tTestPValue(t, df: inferenceDF),
180 + confidenceLower: beta[j] - tCritical * se,
181 + confidenceUpper: beta[j] + tCritical * se
182 + ))
183 + }
184 +
185 + // Wald F on the slopes.
186 + var fStatistic: Double?
187 + var fPValue: Double?
188 + var subVce = [Double](repeating: 0, count: slopes * slopes)
189 + var subBeta = [Double](repeating: 0, count: slopes)
190 + for j in 0..<slopes {
191 + subBeta[j] = beta[j]
192 + for i in 0..<slopes {
193 + subVce[j * slopes + i] = vce[j * k + i]
194 + }
195 + }
196 + if let solved = try? LinearAlgebra.solveSymmetric(subVce, k: slopes, rhs: subBeta) {
197 + let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }
198 + fStatistic = wald / Double(slopes)
199 + fPValue = ZQDistributions.fTestPValue(
200 + wald / Double(slopes), df1: Double(slopes), df2: inferenceDF
201 + )
202 + }
203 +
204 + return ZQFEResult(
205 + coefficients: coefficients,
206 + observationCount: n,
207 + groupCount: groupCount,
208 + degreesOfFreedomResidual: dfResidual,
209 + inferenceDF: inferenceDF,
210 + rSquaredWithin: rSquaredWithin,
211 + rootMSE: sigma2.squareRoot(),
212 + fStatistic: fStatistic,
213 + fPValue: fPValue,
214 + clustered: clustered
215 + )
216 + }
217 +}
added MetrikaKit/Sources/ZQStats/IV.swift +234 −0
@@ -0,0 +1,234 @@
1 +//
2 +// IV.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 Accelerate
11 +import Foundation
12 +
13 +/// Two-stage least squares for `ivregress 2sls`.
14 +///
15 +/// The projection uses the thin Q of the instrument matrix (X̂ = Q·QᵀX),
16 +/// never forming Z'Z. Inference follows Stata's `small` convention:
17 +/// σ² = u'u/(N−K) with t statistics; residuals come from the ORIGINAL
18 +/// regressors (u = y − Xβ), not the projected ones.
19 +public enum ZQIV {
20 +
21 + public static func fit2SLS(
22 + y: [Double],
23 + endogenous: [(name: String, values: [Double])],
24 + exogenous: [(name: String, values: [Double])],
25 + instruments: [(name: String, values: [Double])],
26 + includeConstant: Bool = true,
27 + variance: ZQVarianceEstimator = .classical,
28 + confidenceLevel: Double = 0.95
29 + ) throws -> ZQOLSResult {
30 + let n = y.count
31 + guard !endogenous.isEmpty else {
32 + throw ZQStatsError("ivregress: at least one endogenous regressor required")
33 + }
34 + let k = endogenous.count + exogenous.count + (includeConstant ? 1 : 0)
35 + let m = instruments.count + exogenous.count + (includeConstant ? 1 : 0)
36 + guard m >= k else {
37 + throw ZQStatsError(
38 + "ivregress: order condition fails — \(instruments.count) instruments for \(endogenous.count) endogenous regressors"
39 + )
40 + }
41 + guard n > k else {
42 + throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")
43 + }
44 +
45 + func columnMajor(_ columns: [[Double]]) -> [Double] {
46 + var flat = [Double]()
47 + flat.reserveCapacity(n * columns.count)
48 + for column in columns { flat.append(contentsOf: column) }
49 + return flat
50 + }
51 + let ones = [Double](repeating: 1, count: n)
52 +
53 + // X = [endog, exog, 1], Z = [instruments, exog, 1].
54 + var xColumns = endogenous.map(\.values) + exogenous.map(\.values)
55 + var zColumns = instruments.map(\.values) + exogenous.map(\.values)
56 + if includeConstant {
57 + xColumns.append(ones)
58 + zColumns.append(ones)
59 + }
60 + for column in xColumns + zColumns where column.count != n {
61 + throw ZQStatsError("ivregress: variable has wrong length")
62 + }
63 + let x = columnMajor(xColumns)
64 + let z = columnMajor(zColumns)
65 +
66 + // First stage: X̂ = Q·(QᵀX) with Q the thin Q of Z.
67 + let projected: [Double]
68 + do {
69 + let zQR = try LinearAlgebra.QR(matrix: z, rows: n, cols: m)
70 + let q = try zQR.thinQ()
71 + var qtx = [Double](repeating: 0, count: m * k)
72 + cblas_dgemm(
73 + CblasColMajor, CblasTrans, CblasNoTrans,
74 + Int32(m), Int32(k), Int32(n),
75 + 1.0, q, Int32(n), x, Int32(n), 0.0, &qtx, Int32(m)
76 + )
77 + var xhat = [Double](repeating: 0, count: n * k)
78 + cblas_dgemm(
79 + CblasColMajor, CblasNoTrans, CblasNoTrans,
80 + Int32(n), Int32(k), Int32(m),
81 + 1.0, q, Int32(n), qtx, Int32(m), 0.0, &xhat, Int32(n)
82 + )
83 + projected = xhat
84 + } catch let error as LinearAlgebra.Failure {
85 + throw ZQStatsError("ivregress: instrument matrix is rank deficient — \(error)")
86 + }
87 +
88 + // Second stage.
89 + let secondStage: LinearAlgebra.QR
90 + let beta: [Double]
91 + let bread: [Double]
92 + do {
93 + secondStage = try LinearAlgebra.QR(matrix: projected, rows: n, cols: k)
94 + beta = try secondStage.solve(rhs: y)
95 + bread = try secondStage.crossProductInverse()
96 + } catch {
97 + throw ZQStatsError(
98 + "ivregress: projected design is rank deficient (weak or collinear instruments?)"
99 + )
100 + }
101 +
102 + // Residuals from the ORIGINAL regressors.
103 + let fitted = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)
104 + var residuals = [Double](repeating: 0, count: n)
105 + var rss = 0.0
106 + for i in 0..<n {
107 + residuals[i] = y[i] - fitted[i]
108 + rss += residuals[i] * residuals[i]
109 + }
110 + let dfResidual = n - k
111 + let sigma2 = rss / Double(dfResidual)
112 + let meanY = y.reduce(0, +) / Double(n)
113 + let tss = y.reduce(0) { $0 + ($1 - meanY) * ($1 - meanY) }
114 + let rSquared = tss > 0 ? 1 - rss / tss : .nan // can be negative for IV
115 + let dfModel = k - (includeConstant ? 1 : 0)
116 +
117 + // Covariance: sandwich pieces built on the PROJECTED regressors.
118 + var clusterCount: Int?
119 + let vce: [Double]
120 + switch variance {
121 + case .classical:
122 + vce = bread.map { $0 * sigma2 }
123 + case .hc0, .hc1, .hc2, .hc3:
124 + var meat = [Double](repeating: 0, count: k * k)
125 + for row in 0..<n {
126 + let weight = residuals[row] * residuals[row]
127 + for j in 0..<k {
128 + for i in j..<k {
129 + let value = weight * projected[i * n + row] * projected[j * n + row]
130 + meat[j * k + i] += value
131 + }
132 + }
133 + }
134 + for j in 0..<k {
135 + for i in 0..<j { meat[j * k + i] = meat[i * k + j] }
136 + }
137 + if variance == .hc1 {
138 + let scale = Double(n) / Double(dfResidual)
139 + for index in meat.indices { meat[index] *= scale }
140 + }
141 + vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)
142 + case .cluster(let clusters):
143 + guard clusters.count == n else {
144 + throw ZQStatsError("cluster variable has wrong length")
145 + }
146 + var scores: [Int: [Double]] = [:]
147 + for i in 0..<n {
148 + var u = scores[clusters[i]] ?? [Double](repeating: 0, count: k)
149 + for j in 0..<k {
150 + u[j] += projected[j * n + i] * residuals[i]
151 + }
152 + scores[clusters[i]] = u
153 + }
154 + let g = scores.count
155 + guard g > 1 else {
156 + throw ZQStatsError("cluster variable must define at least 2 groups")
157 + }
158 + clusterCount = g
159 + var meat = [Double](repeating: 0, count: k * k)
160 + for u in scores.values {
161 + for j in 0..<k {
162 + for i in 0..<k {
163 + meat[j * k + i] += u[i] * u[j]
164 + }
165 + }
166 + }
167 + let scale = Double(g) / Double(g - 1) * Double(n - 1) / Double(dfResidual)
168 + for index in meat.indices { meat[index] *= scale }
169 + vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)
170 + }
171 +
172 + let inferenceDF = clusterCount.map { Double($0 - 1) } ?? Double(dfResidual)
173 + let tCritical = ZQDistributions.studentTQuantile(
174 + 0.5 + confidenceLevel / 2, df: inferenceDF
175 + )
176 +
177 + var names = endogenous.map(\.name) + exogenous.map(\.name)
178 + if includeConstant { names.append("_cons") }
179 + var coefficients: [ZQCoefficient] = []
180 + for j in 0..<k {
181 + let se = vce[j * k + j].squareRoot()
182 + let t = beta[j] / se
183 + coefficients.append(ZQCoefficient(
184 + name: names[j],
185 + estimate: beta[j],
186 + standardError: se,
187 + tStatistic: t,
188 + pValue: ZQDistributions.tTestPValue(t, df: inferenceDF),
189 + confidenceLower: beta[j] - tCritical * se,
190 + confidenceUpper: beta[j] + tCritical * se
191 + ))
192 + }
193 +
194 + // Wald F on the slopes.
195 + var fStatistic: Double?
196 + var fPValue: Double?
197 + var fDF: (Double, Double)?
198 + if dfModel > 0 {
199 + var subVce = [Double](repeating: 0, count: dfModel * dfModel)
200 + var subBeta = [Double](repeating: 0, count: dfModel)
201 + for j in 0..<dfModel {
202 + subBeta[j] = beta[j]
203 + for i in 0..<dfModel {
204 + subVce[j * dfModel + i] = vce[j * k + i]
205 + }
206 + }
207 + if let solved = try? LinearAlgebra.solveSymmetric(
208 + subVce, k: dfModel, rhs: subBeta
209 + ) {
210 + let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }
211 + fStatistic = wald / Double(dfModel)
212 + fDF = (Double(dfModel), inferenceDF)
213 + fPValue = ZQDistributions.fTestPValue(
214 + wald / Double(dfModel), df1: Double(dfModel), df2: inferenceDF
215 + )
216 + }
217 + }
218 +
219 + return ZQOLSResult(
220 + coefficients: coefficients,
221 + observationCount: n,
222 + degreesOfFreedomResidual: dfResidual,
223 + inferenceDF: inferenceDF,
224 + rSquared: rSquared,
225 + adjustedRSquared: 1 - (1 - rSquared) * Double(n - 1) / Double(dfResidual),
226 + rootMSE: sigma2.squareRoot(),
227 + fStatistic: fStatistic,
228 + fPValue: fPValue,
229 + fDF: fDF,
230 + clusterCount: clusterCount,
231 + residuals: residuals
232 + )
233 + }
234 +}
modified MetrikaKit/Tests/MetrikaKitTests/DTATests.swift +2 −1
@@ -30,7 +30,8 @@ struct DTATests {
30 30 private func expectMatchesCSV(_ frame: ZQDataFrame) throws {
31 31 #expect(frame.rowCount == 60)
32 32 #expect(frame.columnNames == [
33 − "revenue", "price", "region", "firm_id", "purchase", "orders", "firm_name",
33 + "revenue", "price", "region", "firm_id", "purchase", "orders",
34 + "z1", "z2", "firm_name",
34 35 ])
35 36
36 37 // Numeric columns: bit-identical to the CSV load (both sides parse
modified MetrikaKit/Tests/MetrikaKitTests/Fixtures/expected.tsv +18 −0
@@ -56,6 +56,24 @@ pois_b_cons 1.3512956388566333
56 56 pois_se_price 0.025531826392571365
57 57 pois_ll -81.84353537859181
58 58 pois_se_hc0_price 0.026876084479165823
59 +fe_N 57
60 +fe_G 12
61 +fe_b_price -0.07570509136847102
62 +fe_b_cons 4.0500845138104182
63 +fe_se_price 0.0053912121336584509
64 +fe_r2_within 0.81756861559354577
65 +fe_rmse 0.16821768102868412
66 +fe_se_cluster_price 0.0064468155075312625
67 +fe_p_cluster_price 1.4527243012182743e-07
68 +iv_N 57
69 +iv_b_price -0.076626429435850574
70 +iv_b_cons 4.0624770490720685
71 +iv_se_price 0.0051412992800199363
72 +iv_se_cons 0.072512179823191303
73 +iv_p_price 5.7814855186358344e-21
74 +iv_r2 0.81943731431549249
75 +iv_rmse 0.16468550892743888
76 +iv_se_hc1_price 0.005144683256992212
59 77 corr_rev_price -0.88919780339930277
60 78 corr_rev_logrev 0.98089325161602359
61 79 dist_pchisq_3p8_1 0.94874741714263044
modified MetrikaKit/Tests/MetrikaKitTests/Fixtures/regression.csv +61 −61
@@ -1,61 +1,61 @@
1 −revenue,price,region,firm_id,purchase,orders
2 −13.073013,18.7221,1,1,1,0
3 −14.603161,19.0561,2,1,1,1
4 −37.036281,9.2921,3,1,1,0
5 −12.331106,17.4567,1,1,0,0
6 −20.199281,14.6262,2,1,1,1
7 −18.532705,12.7864,3,2,1,1
8 −13.442576,,1,2,1,0
9 −30.288376,7.02,2,2,0,1
10 −14.146638,14.8549,3,2,0,1
11 −15.789387,15.576,1,2,0,3
12 −24.085548,11.8661,2,3,0,1
13 −17.866099,15.7867,3,3,1,0
14 −13.358007,19.0201,1,3,0,0
15 −26.695153,8.8314,2,3,1,2
16 −20.905369,11.9344,3,3,0,0
17 −12.640513,19.1002,1,4,0,2
18 −11.071996,19.6734,2,4,0,0
19 −48.213111,6.7623,3,4,1,3
20 −19.400193,12.125,1,4,1,2
21 −22.781459,13.405,2,4,1,2
22 −15.854431,18.5605,3,5,1,1
23 −27.549049,7.0807,1,5,1,1
24 −15.637737,,2,5,0,0
25 −15.806717,19.2,3,5,0,3
26 −33.600334,6.2366,1,5,1,3
27 −22.746857,12.7132,2,6,0,0
28 −30.989393,10.8531,3,6,1,3
29 −12.510734,18.5861,1,6,0,0
30 −15.099717,11.7045,2,6,0,4
31 −17.10762,17.5401,3,6,0,2
32 −14.293473,16.0639,1,7,0,0
33 −15.713512,17.1658,2,7,0,2
34 −30.616712,10.8216,3,7,0,2
35 −19.841208,15.2775,1,7,0,0
36 −36.095565,5.0592,2,7,1,5
37 −20.003055,17.4937,3,8,1,0
38 −38.151955,5.11,1,8,1,2
39 −36.839992,8.1149,2,8,1,2
40 −17.29105,18.599,3,8,0,1
41 −19.569965,14.1767,1,8,1,1
42 −21.934484,,2,9,0,2
43 −26.14202,11.5366,3,9,0,0
44 −38.421147,5.5615,1,9,1,1
45 −10.899566,19.6031,2,9,1,3
46 −24.544112,11.4763,3,9,0,1
47 −12.655109,19.3636,1,10,0,0
48 −15.641427,18.3163,2,10,0,1
49 −22.233252,14.5997,3,10,0,0
50 −9.993891,19.5645,1,10,0,0
51 −16.320549,14.2826,2,10,0,0
52 −37.592103,10.0014,3,11,0,4
53 −25.093077,10.2012,1,11,0,3
54 −25.408455,10.9773,2,11,0,1
55 −17.119771,16.7704,3,11,0,1
56 −29.198899,5.584,1,11,0,2
57 −18.051847,16.2319,2,12,1,0
58 −19.195829,15.1592,3,12,0,2
59 −28.993336,7.569,1,12,0,3
60 −34.011476,8.9163,2,12,0,1
61 −27.274985,12.7162,3,12,0,0
1 +revenue,price,region,firm_id,purchase,orders,z1,z2
2 +13.073013,18.7221,1,1,1,0,18.6407,5.7574
3 +14.603161,19.0561,2,1,1,1,15.953,15.639
4 +37.036281,9.2921,3,1,1,0,11.6264,4.9694
5 +12.331106,17.4567,1,1,0,0,16.9094,8.476
6 +20.199281,14.6262,2,1,1,1,13.6905,8.8
7 +18.532705,12.7864,3,2,1,1,10.3099,6.5054
8 +13.442576,,1,2,1,0,16.0333,7.6281
9 +30.288376,7.02,2,2,0,1,5.4194,7.9404
10 +14.146638,14.8549,3,2,0,1,13.7879,6.7764
11 +15.789387,15.576,1,2,0,3,18.1514,3.9372
12 +24.085548,11.8661,2,3,0,1,11.515,7.0901
13 +17.866099,15.7867,3,3,1,0,13.6431,6.8388
14 +13.358007,19.0201,1,3,0,0,19.3465,7.9447
15 +26.695153,8.8314,2,3,1,2,8.1059,1.2113
16 +20.905369,11.9344,3,3,0,0,13.1144,7.2523
17 +12.640513,19.1002,1,4,0,2,21.965,9.028
18 +11.071996,19.6734,2,4,0,0,17.688,11.3837
19 +48.213111,6.7623,3,4,1,3,7.6716,2.6781
20 +19.400193,12.125,1,4,1,2,12.2948,4.087
21 +22.781459,13.405,2,4,1,2,15.1961,10.4532
22 +15.854431,18.5605,3,5,1,1,18.1009,8.465
23 +27.549049,7.0807,1,5,1,1,8.7539,6.3842
24 +15.637737,,2,5,0,0,16.3433,6.312
25 +15.806717,19.2,3,5,0,3,22.5789,8.2017
26 +33.600334,6.2366,1,5,1,3,7.9662,2.3102
27 +22.746857,12.7132,2,6,0,0,12.4116,5.1837
28 +30.989393,10.8531,3,6,1,3,7.9551,9.4727
29 +12.510734,18.5861,1,6,0,0,19.8721,9.2248
30 +15.099717,11.7045,2,6,0,4,12.6709,6.5849
31 +17.10762,17.5401,3,6,0,2,17.5274,5.9429
32 +14.293473,16.0639,1,7,0,0,16.3668,5.8443
33 +15.713512,17.1658,2,7,0,2,15.9976,11.5771
34 +30.616712,10.8216,3,7,0,2,11.5592,9.1862
35 +19.841208,15.2775,1,7,0,0,15.8668,11.3853
36 +36.095565,5.0592,2,7,1,5,4.5007,-1.6123
37 +20.003055,17.4937,3,8,1,0,14.8212,14.8967
38 +38.151955,5.11,1,8,1,2,6.5115,5.6056
39 +36.839992,8.1149,2,8,1,2,9.2233,3.9773
40 +17.29105,18.599,3,8,0,1,16.9264,11.4103
41 +19.569965,14.1767,1,8,1,1,10.9875,4.1742
42 +21.934484,,2,9,0,2,11.1033,2.0582
43 +26.14202,11.5366,3,9,0,0,10.8464,5.9155
44 +38.421147,5.5615,1,9,1,1,6.0667,-0.8147
45 +10.899566,19.6031,2,9,1,3,17.0151,10.3716
46 +24.544112,11.4763,3,9,0,1,9.558,9.6313
47 +12.655109,19.3636,1,10,0,0,21.5351,6.5802
48 +15.641427,18.3163,2,10,0,1,19.1238,6.9428
49 +22.233252,14.5997,3,10,0,0,15.7727,7.4395
50 +9.993891,19.5645,1,10,0,0,23.195,6.7295
51 +16.320549,14.2826,2,10,0,0,14.5402,5.9914
52 +37.592103,10.0014,3,11,0,4,5.9995,7.619
53 +25.093077,10.2012,1,11,0,3,10.8688,8.0092
54 +25.408455,10.9773,2,11,0,1,13.32,6.6402
55 +17.119771,16.7704,3,11,0,1,20.8895,2.8305
56 +29.198899,5.584,1,11,0,2,2.8303,2.63
57 +18.051847,16.2319,2,12,1,0,13.9302,11.3103
58 +19.195829,15.1592,3,12,0,2,13.7476,10.0192
59 +28.993336,7.569,1,12,0,3,5.4609,3.2121
60 +34.011476,8.9163,2,12,0,1,7.6248,-3.6416
61 +27.274985,12.7162,3,12,0,0,12.3454,6.541
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.

added MetrikaKit/Tests/MetrikaKitTests/PanelIVTests.swift +121 −0
@@ -0,0 +1,121 @@
1 +//
2 +// PanelIVTests.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 ZQParser
14 +
15 +/// xtreg FE and ivregress 2SLS against R fixtures, end-to-end through the
16 +/// session (which exercises the parser's IV-group grammar too).
17 +@Suite("Panel FE and IV", .serialized)
18 +struct PanelIVTests {
19 + let fixtures: Fixtures
20 + let session: ZQSession
21 +
22 + init() async throws {
23 + self.fixtures = try Fixtures()
24 + self.session = try ZQSession(discoverUserCommands: false)
25 + _ = try await session.execute("use \(fixtures.datasetURL.path)")
26 + _ = try await session.execute("gen log_rev = ln(revenue)")
27 + _ = try await session.execute("xtset firm_id")
28 + }
29 +
30 + @Test("ivregress parser grammar")
31 + func parserIVGroup() throws {
32 + let parser = ZQCommandParser()
33 + let command = try #require(
34 + try parser.parse("ivregress 2sls y x1 (p q = z1 z2 z3), robust")
35 + )
36 + #expect(command.verb == "ivregress")
37 + #expect(command.subverb == "2sls")
38 + #expect(command.varlist == [.simple("y"), .simple("x1")])
39 + #expect(command.ivSpec == ZQIVSpec(
40 + endogenous: ["p", "q"], instruments: ["z1", "z2", "z3"]
41 + ))
42 + #expect(command.hasOption("robust"))
43 + }
44 +
45 + @Test("xtreg fe matches the R within estimator")
46 + func fixedEffects() async throws {
47 + let result = try await session.execute("xtreg log_rev price, fe")
48 + #expect(result.scalars["N"] == fixtures["fe_N"])
49 + #expect(result.scalars["N_g"] == fixtures["fe_G"])
50 + expectClose(try #require(result.scalars["b_price"]), fixtures["fe_b_price"], "b[price]")
51 + expectClose(try #require(result.scalars["b__cons"]), fixtures["fe_b_cons"], "b[_cons]")
52 + expectClose(
53 + try #require(result.scalars["se_price"]),
54 + fixtures["fe_se_price"], "se[price]"
55 + )
56 + expectClose(try #require(result.scalars["r2_w"]), fixtures["fe_r2_within"], "within R²")
57 + expectClose(try #require(result.scalars["rmse"]), fixtures["fe_rmse"], "root MSE")
58 + #expect(result.text.contains("Fixed-effects (within) regression"))
59 + }
60 +
61 + @Test("xtreg fe with panel-clustered SE matches R")
62 + func fixedEffectsClustered() async throws {
63 + let result = try await session.execute("xtreg log_rev price, fe cluster(firm_id)")
64 + expectClose(
65 + try #require(result.scalars["se_price"]),
66 + fixtures["fe_se_cluster_price"], "cluster se[price]"
67 + )
68 + #expect(result.scalars["df_r"] == fixtures["fe_G"] - 1)
69 + }
70 +
71 + @Test("xtreg without xtset or fe errors clearly")
72 + func xtregValidation() async throws {
73 + let fresh = try ZQSession(discoverUserCommands: false)
74 + _ = try await fresh.execute("use \(fixtures.datasetURL.path)")
75 + await #expect(throws: ZQEngineError.self) {
76 + _ = try await fresh.execute("xtreg revenue price, fe") // no xtset
77 + }
78 + await #expect(throws: ZQEngineError.self) {
79 + _ = try await session.execute("xtreg log_rev price") // no fe
80 + }
81 + }
82 +
83 + @Test("ivregress 2sls matches the R manual 2SLS")
84 + func twoStageLeastSquares() async throws {
85 + let result = try await session.execute(
86 + "ivregress 2sls log_rev (price = z1 z2)"
87 + )
88 + #expect(result.scalars["N"] == fixtures["iv_N"])
89 + expectClose(try #require(result.scalars["b_price"]), fixtures["iv_b_price"], "b[price]")
90 + expectClose(try #require(result.scalars["b__cons"]), fixtures["iv_b_cons"], "b[_cons]")
91 + expectClose(
92 + try #require(result.scalars["se_price"]), fixtures["iv_se_price"], "se[price]"
93 + )
94 + expectClose(
95 + try #require(result.scalars["se__cons"]), fixtures["iv_se_cons"], "se[_cons]"
96 + )
97 + expectClose(try #require(result.scalars["r2"]), fixtures["iv_r2"], "R²")
98 + expectClose(try #require(result.scalars["rmse"]), fixtures["iv_rmse"], "root MSE")
99 + #expect(result.text.contains("Instrumented: price"))
100 + }
101 +
102 + @Test("ivregress robust SE matches R HC1-small")
103 + func ivRobust() async throws {
104 + let result = try await session.execute(
105 + "ivregress 2sls log_rev (price = z1 z2), robust"
106 + )
107 + expectClose(
108 + try #require(result.scalars["se_price"]),
109 + fixtures["iv_se_hc1_price"], "robust se[price]"
110 + )
111 + }
112 +
113 + @Test("ivregress order condition is enforced")
114 + func orderCondition() async throws {
115 + await #expect(throws: (any Error).self) {
116 + _ = try await session.execute(
117 + "ivregress 2sls log_rev (price orders = z1)"
118 + )
119 + }
120 + }
121 +}
modified Tests/Fixtures/generate.R +76 −1
@@ -37,6 +37,11 @@ revenue <- round(exp(log_rev), 6)
37 37 purchase <- rbinom(n, 1, plogis(2 - 0.20 * price))
38 38 orders <- rpois(n, exp(1.6 - 0.09 * price))
39 39
40 +# Instruments for the 2SLS fixture (correlated with price by
41 +# construction), drawn after everything above.
42 +z1 <- round(price + rnorm(n, 0, 2), 4)
43 +z2 <- round(0.5 * price + rnorm(n, 0, 3), 4)
44 +
40 45 # Inject missing values to exercise listwise deletion.
41 46 price_missing <- price
42 47 price_missing[c(7, 23, 41)] <- NA
@@ -47,7 +52,9 @@ data <- data.frame(
47 52 region = region,
48 53 firm_id = firm_id,
49 54 purchase = purchase,
50 − orders = orders
55 + orders = orders,
56 + z1 = z1,
57 + z2 = z2
51 58 )
52 59 csv_path <- file.path(out_dir, "regression.csv")
53 60 write.csv(data, csv_path, row.names = FALSE, quote = FALSE, na = "")
@@ -216,6 +223,74 @@ emit("pois_se_price", sqrt(glm_bread(ps)["price", "price"]))
216 223 emit("pois_ll", as.numeric(logLik(ps)))
217 224 emit("pois_se_hc0_price", glm_bread_sandwich(ps, d$orders - fitted(ps))["price"])
218 225
226 +# ----------------------------------------------------- fixed effects (FE)
227 +# Within estimator, Stata conventions: demean within firm, add back grand
228 +# means (so _cons is reported), df = N − K − G. Cluster VCE on the panel
229 +# with the G/(G−1) factor and t on G−1 df.
230 +fe_within <- function(v, g) {
231 + gm <- ave(v, g) # group means
232 + v - gm + mean(v)
233 +}
234 +g <- d$firm_id
235 +Nw <- nrow(d)
236 +yt <- fe_within(d$log_rev, g)
237 +xt_p <- fe_within(d$price, g)
238 +Xw <- cbind(price = xt_p, `_cons` = 1)
239 +G <- length(unique(g))
240 +Kw <- 1
241 +df_fe <- Nw - Kw - G
242 +
243 +fe_fit <- lm.fit(Xw, yt)
244 +b_fe <- fe_fit$coefficients
245 +res_fe <- fe_fit$residuals
246 +XtXinv_w <- solve(t(Xw) %*% Xw)
247 +sigma2_fe <- sum(res_fe^2) / df_fe
248 +emit("fe_N", Nw)
249 +emit("fe_G", G)
250 +emit("fe_b_price", b_fe[["price"]])
251 +emit("fe_b_cons", b_fe[["_cons"]])
252 +emit("fe_se_price", sqrt(sigma2_fe * XtXinv_w["price", "price"]))
253 +yd <- d$log_rev - ave(d$log_rev, g)
254 +emit("fe_r2_within", 1 - sum(res_fe^2) / sum(yd^2))
255 +emit("fe_rmse", sqrt(sigma2_fe))
256 +
257 +# Cluster on the panel: scores use DEMEANED x (constant column as-is).
258 +xd <- d$price - ave(d$price, g)
259 +u_fe <- rowsum(cbind(xd * res_fe, res_fe), g)
260 +meat_fe <- (G / (G - 1)) * (t(u_fe) %*% u_fe)
261 +V_fe_cl <- XtXinv_w %*% meat_fe %*% XtXinv_w
262 +emit("fe_se_cluster_price", sqrt(V_fe_cl[1, 1]))
263 +t_fe_cl <- b_fe[["price"]] / sqrt(V_fe_cl[1, 1])
264 +emit("fe_p_cluster_price", 2 * pt(-abs(t_fe_cl), G - 1))
265 +
266 +# ------------------------------------------------------------- 2SLS (IV)
267 +# ivregress 2sls log_rev (price = z1 z2), Stata `small` convention:
268 +# residuals from the ORIGINAL regressors, sigma2 = u'u/(N-K), t stats.
269 +Xiv <- cbind(price = d$price, `_cons` = 1)
270 +Ziv <- cbind(z1 = d$z1, z2 = d$z2, `_cons` = 1)
271 +PZ <- Ziv %*% solve(t(Ziv) %*% Ziv) %*% t(Ziv)
272 +Xhat <- PZ %*% Xiv
273 +XhXhinv <- solve(t(Xhat) %*% Xhat)
274 +b_iv <- XhXhinv %*% t(Xhat) %*% d$log_rev
275 +u_iv <- d$log_rev - Xiv %*% b_iv
276 +df_iv <- Nw - 2
277 +sigma2_iv <- sum(u_iv^2) / df_iv
278 +V_iv <- sigma2_iv * XhXhinv
279 +emit("iv_N", Nw)
280 +emit("iv_b_price", b_iv[1, 1])
281 +emit("iv_b_cons", b_iv[2, 1])
282 +emit("iv_se_price", sqrt(V_iv[1, 1]))
283 +emit("iv_se_cons", sqrt(V_iv[2, 2]))
284 +t_iv <- b_iv[1, 1] / sqrt(V_iv[1, 1])
285 +emit("iv_p_price", 2 * pt(-abs(t_iv), df_iv))
286 +emit("iv_r2", 1 - sum(u_iv^2) / sum((d$log_rev - mean(d$log_rev))^2))
287 +emit("iv_rmse", sqrt(sigma2_iv))
288 +
289 +# Robust (HC1 with the small factor N/(N-K)) built on projected X.
290 +meat_iv <- t(Xhat * as.vector(u_iv^2)) %*% Xhat * (Nw / df_iv)
291 +V_iv_r <- XhXhinv %*% meat_iv %*% XhXhinv
292 +emit("iv_se_hc1_price", sqrt(V_iv_r[1, 1]))
293 +
219 294 # -------------------------------------------------------------- correlate
220 295 emit("corr_rev_price", cor(d$revenue, d$price))
221 296 emit("corr_rev_logrev", cor(d$revenue, d$log_rev))
222 297