SPB Git

spb/metrika Public

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

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

feat(stats): logit/probit/poisson MLE, tabulate, correlate

- ZQGLM: Fisher-scoring IRLS through the LAPACK QR path (X'WX never
  formed), converged to |Δll| < 1e-13; classical, robust (HC0 score
  sandwich, Stata ML convention), and cluster (G/(G−1)) VCE; LR/Wald
  chi2, McFadden pseudo-R²; information matrix re-evaluated at the
  converged beta (R's vcov carries last-iteration weights and is only
  ~1e-6-accurate by its own stopping rule — fixtures compute expected
  information at the optimum explicitly)
- Distributions: regularized incomplete gamma, chi-square CDF/p-value,
  normal quantile
- ZQCorrelate: Pearson matrix on listwise-complete data
- Engine: logit/probit/poisson tables (z, P>|z|, LR chi2, pseudo R²),
  one-way and two-way tabulate with totals, Stata-style lower-triangle
  correlate
- Fixtures: binomial/poisson outcomes drawn after existing draws (earlier
  golden values bit-identical); glm at epsilon 1e-12
- 56 tests green (12 new)

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

Showing 9 changed files with +1,199 and −62

modified MetrikaKit/Sources/ZQEngine/Session.swift +271 −0
@@ -124,6 +124,11 @@ public actor ZQSession {
124 124 case "count": return try handleCount(command)
125 125 case "list": return try handleList(command)
126 126 case "regress": return try handleRegress(command)
127 + case "logit": return try handleGLM(command, family: .logit)
128 + case "probit": return try handleGLM(command, family: .probit)
129 + case "poisson": return try handleGLM(command, family: .poisson)
130 + case "tabulate": return try handleTabulate(command)
131 + case "correlate": return try handleCorrelate(command)
127 132 case "set": return try handleSet(command)
128 133 case "xtset": return try handleXTSet(command)
129 134 case "display": return try handleDisplay(command)
@@ -722,6 +727,272 @@ public actor ZQSession {
722 727 return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
723 728 }
724 729
730 + private func handleGLM(
731 + _ command: ZQCommand, family: ZQGLMFamily
732 + ) throws -> ZQResult {
733 + let clusterVariable = command.option("cluster")?.firstArgument
734 + ?? clusterFromVCE(command)
735 + let sample = try buildRegressionSample(command, clusterVariable: clusterVariable)
736 + let variance = try varianceEstimator(command, clusterLabels: sample.clusterLabels)
737 + let level = try confidenceLevel(command)
738 +
739 + let result = try ZQGLM.fit(
740 + y: sample.y,
741 + predictors: sample.predictors,
742 + family: family,
743 + includeConstant: !command.hasOption("noconstant"),
744 + variance: variance,
745 + confidenceLevel: level
746 + )
747 +
748 + let title: String
749 + switch family {
750 + case .logit: title = "Logistic regression"
751 + case .probit: title = "Probit regression"
752 + case .poisson: title = "Poisson regression"
753 + }
754 +
755 + var header = [title]
756 + if sample.droppedMissing > 0 {
757 + header.append("(\(sample.droppedMissing) observations dropped due to missing values)")
758 + }
759 + var stats: [(String, String)] = [
760 + ("Number of obs", "\(result.observationCount)"),
761 + ]
762 + if let chi2 = result.chiSquared, let p = result.chiSquaredPValue {
763 + let label: String
764 + if case .classical = variance {
765 + label = "LR chi2(\(result.chiSquaredDF))"
766 + } else {
767 + label = "Wald chi2(\(result.chiSquaredDF))"
768 + }
769 + stats.append((label, TableFormatter.fixed(chi2, decimals: 2)))
770 + stats.append(("Prob > chi2", TableFormatter.fixed(p, decimals: 4)))
771 + }
772 + stats.append(("Log likelihood", TableFormatter.fixed(result.logLikelihood, decimals: 4)))
773 + stats.append(("Pseudo R2", TableFormatter.fixed(result.pseudoRSquared, decimals: 4)))
774 + if let g = result.clusterCount {
775 + stats.append(("Clusters", "\(g)"))
776 + }
777 + for (label, value) in stats {
778 + header.append(
779 + TableFormatter.pad(label, 46, right: false) + "= " +
780 + TableFormatter.pad(value, 12)
781 + )
782 + }
783 +
784 + var lines = header
785 + lines.append("")
786 + let widths = [12, 12, 11, 8, 8, 22]
787 + lines.append(
788 + TableFormatter.pad(sample.responseName, widths[0]) + " | " +
789 + TableFormatter.pad("Coefficient", widths[1]) + " " +
790 + TableFormatter.pad("Std. err.", widths[2]) + " " +
791 + TableFormatter.pad("z", widths[3]) + " " +
792 + TableFormatter.pad("P>|z|", widths[4]) + " " +
793 + TableFormatter.pad("[\(Int(level * 100))% conf. interval]", widths[5])
794 + )
795 + lines.append(TableFormatter.rule(widths))
796 +
797 + var scalars: [String: Double] = [
798 + "N": Double(result.observationCount),
799 + "ll": result.logLikelihood,
800 + "ll_0": result.nullLogLikelihood,
801 + "r2_p": result.pseudoRSquared,
802 + ]
803 + if let chi2 = result.chiSquared { scalars["chi2"] = chi2 }
804 + if let g = result.clusterCount { scalars["N_clust"] = Double(g) }
805 +
806 + for coefficient in result.coefficients {
807 + lines.append(
808 + TableFormatter.pad(coefficient.name, widths[0]) + " | " +
809 + TableFormatter.pad(TableFormatter.general(coefficient.estimate), widths[1]) + " " +
810 + TableFormatter.pad(TableFormatter.general(coefficient.standardError), widths[2]) + " " +
811 + TableFormatter.pad(TableFormatter.fixed(coefficient.tStatistic, decimals: 2), widths[3]) + " " +
812 + TableFormatter.pad(TableFormatter.fixed(coefficient.pValue, decimals: 3), widths[4]) + " " +
813 + TableFormatter.pad(
814 + TableFormatter.general(coefficient.confidenceLower) + " " +
815 + TableFormatter.general(coefficient.confidenceUpper),
816 + widths[5]
817 + )
818 + )
819 + scalars["b_\(coefficient.name)"] = coefficient.estimate
820 + scalars["se_\(coefficient.name)"] = coefficient.standardError
821 + }
822 + return ZQResult(text: lines.joined(separator: "\n"), scalars: scalars)
823 + }
824 +
825 + // MARK: - Tabulate & correlate
826 +
827 + private func handleTabulate(_ command: ZQCommand) throws -> ZQResult {
828 + guard !frame.isEmpty else { throw ZQEngineError("no data in memory") }
829 + let names = command.varlist.flatMap(\.referencedNames)
830 + guard names.count == 1 || names.count == 2 else {
831 + throw ZQEngineError("tabulate: syntax is 'tabulate varname [varname2]'")
832 + }
833 + let mask = try observationMask(command)
834 +
835 + /// Rendered level label per observation; nil = missing (excluded
836 + /// unless the `missing` option is given).
837 + func labels(_ name: String) throws -> [String?] {
838 + let column = try frame.requireColumn(name)
839 + switch column.data {
840 + case .float64(let values, let missing):
841 + return (0..<frame.rowCount).map { i in
842 + missing[i] ? nil : TableFormatter.general(values[i])
843 + }
844 + case .string(let values):
845 + return values
846 + }
847 + }
848 +
849 + let includeMissing = command.hasOption("missing")
850 + let first = try labels(names[0])
851 +
852 + if names.count == 1 {
853 + var counts: [String: Int] = [:]
854 + var total = 0
855 + for i in 0..<frame.rowCount where mask[i] {
856 + guard let label = includeMissing ? (first[i] ?? ".") : first[i] else {
857 + continue
858 + }
859 + counts[label, default: 0] += 1
860 + total += 1
861 + }
862 + guard total > 0 else { throw ZQEngineError("no observations") }
863 +
864 + var lines = [
865 + TableFormatter.pad(names[0], 14) + " | " +
866 + TableFormatter.pad("Freq.", 9) + " " +
867 + TableFormatter.pad("Percent", 9) + " " +
868 + TableFormatter.pad("Cum.", 9),
869 + String(repeating: "-", count: 15) + "+" + String(repeating: "-", count: 33),
870 + ]
871 + var cumulative = 0.0
872 + for label in counts.keys.sorted(by: numericAwareLess) {
873 + let count = counts[label]!
874 + let percent = 100 * Double(count) / Double(total)
875 + cumulative += percent
876 + lines.append(
877 + TableFormatter.pad(label, 14) + " | " +
878 + TableFormatter.pad("\(count)", 9) + " " +
879 + TableFormatter.pad(TableFormatter.fixed(percent, decimals: 2), 9) + " " +
880 + TableFormatter.pad(TableFormatter.fixed(cumulative, decimals: 2), 9)
881 + )
882 + }
883 + lines.append(String(repeating: "-", count: 15) + "+" + String(repeating: "-", count: 33))
884 + lines.append(
885 + TableFormatter.pad("Total", 14) + " | " +
886 + TableFormatter.pad("\(total)", 9) + " " +
887 + TableFormatter.pad("100.00", 9)
888 + )
889 + return ZQResult(
890 + text: lines.joined(separator: "\n"),
891 + scalars: ["N": Double(total), "r": Double(counts.count)]
892 + )
893 + }
894 +
895 + // Two-way table.
896 + let second = try labels(names[1])
897 + var cells: [String: [String: Int]] = [:]
898 + var rowTotals: [String: Int] = [:]
899 + var columnTotals: [String: Int] = [:]
900 + var total = 0
901 + for i in 0..<frame.rowCount where mask[i] {
902 + let rowLabel = includeMissing ? (first[i] ?? ".") : first[i]
903 + let columnLabel = includeMissing ? (second[i] ?? ".") : second[i]
904 + guard let rowLabel, let columnLabel else { continue }
905 + cells[rowLabel, default: [:]][columnLabel, default: 0] += 1
906 + rowTotals[rowLabel, default: 0] += 1
907 + columnTotals[columnLabel, default: 0] += 1
908 + total += 1
909 + }
910 + guard total > 0 else { throw ZQEngineError("no observations") }
911 +
912 + let rows = rowTotals.keys.sorted(by: numericAwareLess)
913 + let columns = columnTotals.keys.sorted(by: numericAwareLess)
914 + let width = 9
915 + var lines = [
916 + TableFormatter.pad(names[0], 14) + " | " +
917 + columns.map { TableFormatter.pad($0, width) }.joined(separator: " ") +
918 + " " + TableFormatter.pad("Total", width),
919 + String(repeating: "-", count: 15) + "+" +
920 + String(repeating: "-", count: (width + 1) * (columns.count + 1) + 2),
921 + ]
922 + for row in rows {
923 + let cellsText = columns.map {
924 + TableFormatter.pad("\(cells[row]?[$0] ?? 0)", width)
925 + }.joined(separator: " ")
926 + lines.append(
927 + TableFormatter.pad(row, 14) + " | " + cellsText + " " +
928 + TableFormatter.pad("\(rowTotals[row] ?? 0)", width)
929 + )
930 + }
931 + lines.append(
932 + TableFormatter.pad("Total", 14) + " | " +
933 + columns.map { TableFormatter.pad("\(columnTotals[$0] ?? 0)", width) }
934 + .joined(separator: " ") +
935 + " " + TableFormatter.pad("\(total)", width)
936 + )
937 + return ZQResult(
938 + text: lines.joined(separator: "\n"),
939 + scalars: [
940 + "N": Double(total),
941 + "r": Double(rows.count),
942 + "c": Double(columns.count),
943 + ]
944 + )
945 + }
946 +
947 + /// Sorts numeric labels numerically, everything else lexically.
948 + private func numericAwareLess(_ a: String, _ b: String) -> Bool {
949 + if let x = Double(a), let y = Double(b) { return x < y }
950 + return a < b
951 + }
952 +
953 + private func handleCorrelate(_ command: ZQCommand) throws -> ZQResult {
954 + guard !frame.isEmpty else { throw ZQEngineError("no data in memory") }
955 + var names = command.varlist.flatMap(\.referencedNames)
956 + if names.isEmpty {
957 + names = frame.columns.filter(\.data.isNumeric).map(\.name)
958 + }
959 + guard names.count >= 2 else {
960 + throw ZQEngineError("correlate: at least two numeric variables required")
961 + }
962 + let mask = try observationMask(command)
963 + let sources = try names.map { try frame.requireNumeric($0) }
964 +
965 + // Listwise deletion across the varlist (Stata correlate).
966 + var kept: [Int] = []
967 + for i in 0..<frame.rowCount where mask[i] {
968 + if sources.allSatisfy({ !$0.missing[i] }) { kept.append(i) }
969 + }
970 + guard kept.count > 1 else { throw ZQEngineError("no observations") }
971 +
972 + let columns = sources.map { source in kept.map { source.values[$0] } }
973 + let matrix = ZQCorrelate.matrix(columns: columns)
974 +
975 + var lines = [
976 + "(obs=\(kept.count))",
977 + "",
978 + TableFormatter.pad("", 12) + " | " +
979 + names.map { TableFormatter.pad($0, 9) }.joined(separator: " "),
980 + String(repeating: "-", count: 13) + "+" +
981 + String(repeating: "-", count: (9 + 1) * names.count + 1),
982 + ]
983 + for (i, name) in names.enumerated() {
984 + // Lower triangle only, Stata-style.
985 + let cells = (0...i).map {
986 + TableFormatter.pad(TableFormatter.fixed(matrix[i][$0], decimals: 4), 9)
987 + }.joined(separator: " ")
988 + lines.append(TableFormatter.pad(name, 12) + " | " + cells)
989 + }
990 + return ZQResult(
991 + text: lines.joined(separator: "\n"),
992 + scalars: ["N": Double(kept.count), "rho": matrix[0][1]]
993 + )
994 + }
995 +
725 996 private func clusterFromVCE(_ command: ZQCommand) -> String? {
726 997 guard let vce = command.option("vce"),
727 998 vce.arguments.count == 2,
added MetrikaKit/Sources/ZQStats/Correlate.swift +40 −0
@@ -0,0 +1,40 @@
1 +//
2 +// Correlate.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 +/// Pearson correlation matrix on listwise-complete data (the engine
11 +/// applies listwise deletion before calling, matching Stata `correlate`).
12 +public enum ZQCorrelate {
13 + public static func matrix(columns: [[Double]]) -> [[Double]] {
14 + let k = columns.count
15 + guard k > 0 else { return [] }
16 + let n = columns[0].count
17 +
18 + let means = columns.map { $0.reduce(0, +) / Double(n) }
19 + var centered = columns
20 + for j in 0..<k {
21 + for i in 0..<n { centered[j][i] -= means[j] }
22 + }
23 + let norms = centered.map { column in
24 + column.reduce(0) { $0 + $1 * $1 }.squareRoot()
25 + }
26 +
27 + var result = [[Double]](repeating: [Double](repeating: 1, count: k), count: k)
28 + for a in 0..<k {
29 + for b in (a + 1)..<k {
30 + var dot = 0.0
31 + for i in 0..<n { dot += centered[a][i] * centered[b][i] }
32 + let denominator = norms[a] * norms[b]
33 + let r = denominator > 0 ? dot / denominator : .nan
34 + result[a][b] = r
35 + result[b][a] = r
36 + }
37 + }
38 + return result
39 + }
40 +}
modified MetrikaKit/Sources/ZQStats/Distributions.swift +66 −0
@@ -127,6 +127,72 @@ public enum ZQDistributions {
127 127 return incompleteBeta(a: df2 / 2, b: df1 / 2, x: df2 / (df2 + df1 * f))
128 128 }
129 129
130 + /// Regularized lower incomplete gamma P(a, x), via the series for
131 + /// x < a+1 and the Lentz continued fraction otherwise.
132 + public static func incompleteGamma(a: Double, x: Double) -> Double {
133 + precondition(a > 0, "incompleteGamma requires a > 0")
134 + if x <= 0 { return 0 }
135 +
136 + if x < a + 1 {
137 + // Series representation.
138 + var term = 1.0 / a
139 + var sum = term
140 + var ap = a
141 + for _ in 0..<500 {
142 + ap += 1
143 + term *= x / ap
144 + sum += term
145 + if abs(term) < abs(sum) * 1e-16 { break }
146 + }
147 + return sum * exp(-x + a * log(x) - logGamma(a))
148 + } else {
149 + // Continued fraction for Q(a, x); P = 1 − Q.
150 + let tiny = 1e-300
151 + var b = x + 1 - a
152 + var c = 1 / tiny
153 + var d = 1 / b
154 + var h = d
155 + for i in 1...500 {
156 + let an = -Double(i) * (Double(i) - a)
157 + b += 2
158 + d = an * d + b
159 + if abs(d) < tiny { d = tiny }
160 + c = b + an / c
161 + if abs(c) < tiny { c = tiny }
162 + d = 1 / d
163 + let delta = d * c
164 + h *= delta
165 + if abs(delta - 1) < 1e-16 { break }
166 + }
167 + let q = exp(-x + a * log(x) - logGamma(a)) * h
168 + return 1 - q
169 + }
170 + }
171 +
172 + /// Chi-square CDF with `df` degrees of freedom.
173 + public static func chiSquareCDF(_ x: Double, df: Double) -> Double {
174 + x <= 0 ? 0 : incompleteGamma(a: df / 2, x: x / 2)
175 + }
176 +
177 + /// Upper-tail p-value for a chi-square statistic.
178 + public static func chiSquarePValue(_ x: Double, df: Double) -> Double {
179 + 1 - chiSquareCDF(x, df: df)
180 + }
181 +
182 + /// Standard normal quantile via bisection on the CDF (erfc-based, so
183 + /// accurate to ~1e-15) — used for z confidence intervals.
184 + public static func normalQuantile(_ p: Double) -> Double {
185 + precondition(p > 0 && p < 1, "quantile requires 0 < p < 1")
186 + if abs(p - 0.5) < 1e-15 { return 0 }
187 + var low = -40.0, high = 40.0
188 + for _ in 0..<200 {
189 + let mid = 0.5 * (low + high)
190 + if normalCDF(mid) < p { low = mid } else { high = mid }
191 + if high - low < 1e-15 * max(1, abs(mid)) { break }
192 + }
193 + return 0.5 * (low + high)
194 + }
195 +
130 196 /// Student t quantile (inverse CDF) via bisection refined with Newton
131 197 /// steps — used for confidence intervals.
132 198 public static func studentTQuantile(_ p: Double, df: Double) -> Double {
added MetrikaKit/Sources/ZQStats/GLM.swift +445 −0
@@ -0,0 +1,445 @@
1 +//
2 +// GLM.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 +/// Maximum-likelihood binary/count models via iteratively reweighted
13 +/// least squares (Fisher scoring), CLAUDE.md §5. Each IRLS step solves a
14 +/// weighted least-squares problem through the same LAPACK QR path as OLS,
15 +/// so X'WX is never formed explicitly.
16 +public enum ZQGLMFamily: String, Equatable, Sendable {
17 + case logit, probit, poisson
18 +}
19 +
20 +public struct ZQGLMResult: Equatable, Sendable {
21 + public var family: ZQGLMFamily
22 + public var coefficients: [ZQCoefficient]
23 + public var observationCount: Int
24 + public var logLikelihood: Double
25 + public var nullLogLikelihood: Double
26 + /// McFadden pseudo-R²: 1 − ll/ll₀ (Stata's reported definition).
27 + public var pseudoRSquared: Double
28 + /// LR χ² against the intercept-only model (classical VCE only; under
29 + /// robust/cluster VCE Stata reports a Wald χ², which we mirror).
30 + public var chiSquared: Double?
31 + public var chiSquaredDF: Int
32 + public var chiSquaredPValue: Double?
33 + public var iterations: Int
34 + public var clusterCount: Int?
35 + /// Linear predictor and fitted mean at the optimum.
36 + public var fittedMeans: [Double]
37 +
38 + public init(
39 + family: ZQGLMFamily,
40 + coefficients: [ZQCoefficient],
41 + observationCount: Int,
42 + logLikelihood: Double,
43 + nullLogLikelihood: Double,
44 + pseudoRSquared: Double,
45 + chiSquared: Double?,
46 + chiSquaredDF: Int,
47 + chiSquaredPValue: Double?,
48 + iterations: Int,
49 + clusterCount: Int?,
50 + fittedMeans: [Double]
51 + ) {
52 + self.family = family
53 + self.coefficients = coefficients
54 + self.observationCount = observationCount
55 + self.logLikelihood = logLikelihood
56 + self.nullLogLikelihood = nullLogLikelihood
57 + self.pseudoRSquared = pseudoRSquared
58 + self.chiSquared = chiSquared
59 + self.chiSquaredDF = chiSquaredDF
60 + self.chiSquaredPValue = chiSquaredPValue
61 + self.iterations = iterations
62 + self.clusterCount = clusterCount
63 + self.fittedMeans = fittedMeans
64 + }
65 +}
66 +
67 +public enum ZQGLM {
68 +
69 + /// Fits by Fisher scoring to |Δll| < 1e-13 (tighter than R's default
70 + /// so both land on the same optimum to ≥1e-10).
71 + ///
72 + /// Variance estimators: `.classical` is the inverse expected
73 + /// information; `.hc0`/`.hc1` (and Stata's `robust`) are the score
74 + /// sandwich without df correction; `.cluster` applies Stata's ML
75 + /// factor G/(G−1). `.hc2`/`.hc3` are not defined for ML estimators.
76 + public static func fit(
77 + y: [Double],
78 + predictors: [(name: String, values: [Double])],
79 + family: ZQGLMFamily,
80 + includeConstant: Bool = true,
81 + variance: ZQVarianceEstimator = .classical,
82 + confidenceLevel: Double = 0.95,
83 + maximumIterations: Int = 100
84 + ) throws -> ZQGLMResult {
85 + let n = y.count
86 + var names = predictors.map(\.name)
87 + if includeConstant { names.append("_cons") }
88 + let k = names.count
89 + guard n > k else {
90 + throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")
91 + }
92 + if case .hc2 = variance {
93 + throw ZQStatsError("hc2 is not defined for maximum-likelihood estimators")
94 + }
95 + if case .hc3 = variance {
96 + throw ZQStatsError("hc3 is not defined for maximum-likelihood estimators")
97 + }
98 +
99 + switch family {
100 + case .logit, .probit:
101 + for value in y where value != 0 && value != 1 {
102 + throw ZQStatsError("\(family.rawValue): outcome must be 0/1")
103 + }
104 + case .poisson:
105 + for value in y where value < 0 || value != value.rounded() {
106 + throw ZQStatsError("poisson: outcome must be a nonnegative count")
107 + }
108 + }
109 +
110 + // Design matrix, column-major.
111 + var x = [Double]()
112 + x.reserveCapacity(n * k)
113 + for column in predictors {
114 + guard column.values.count == n else {
115 + throw ZQStatsError("regressor '\(column.name)' has wrong length")
116 + }
117 + x.append(contentsOf: column.values)
118 + }
119 + if includeConstant { x.append(contentsOf: [Double](repeating: 1, count: n)) }
120 +
121 + // IRLS.
122 + var beta = [Double](repeating: 0, count: k)
123 + var eta = [Double](repeating: 0, count: n)
124 + var mu = startingMeans(y: y, family: family)
125 + if family == .poisson || !includeConstant {
126 + eta = mu.map { link(family: family, mean: $0) }
127 + } else {
128 + eta = mu.map { link(family: family, mean: $0) }
129 + }
130 + var logLikelihood = self.logLikelihood(y: y, mu: mu, family: family)
131 + var iterations = 0
132 + var finalQR: LinearAlgebra.QR?
133 + var weights = [Double](repeating: 0, count: n)
134 +
135 + for iteration in 1...maximumIterations {
136 + iterations = iteration
137 +
138 + // Working response and weights.
139 + var z = [Double](repeating: 0, count: n)
140 + for i in 0..<n {
141 + let derivative = meanDerivative(family: family, eta: eta[i], mean: mu[i])
142 + let varMu = varianceFunction(family: family, mean: mu[i])
143 + weights[i] = derivative * derivative / varMu
144 + z[i] = eta[i] + (y[i] - mu[i]) / derivative
145 + }
146 +
147 + // Weighted least squares via QR on √W·X, √W·z.
148 + var xw = [Double](repeating: 0, count: n * k)
149 + var zw = [Double](repeating: 0, count: n)
150 + for i in 0..<n {
151 + let root = weights[i].squareRoot()
152 + zw[i] = z[i] * root
153 + for j in 0..<k {
154 + xw[j * n + i] = x[j * n + i] * root
155 + }
156 + }
157 + let qr: LinearAlgebra.QR
158 + do {
159 + qr = try LinearAlgebra.QR(matrix: xw, rows: n, cols: k)
160 + beta = try qr.solve(rhs: zw)
161 + } catch {
162 + throw ZQStatsError("\(family.rawValue): weighted solve failed — \(error)")
163 + }
164 + finalQR = qr
165 +
166 + // Update linear predictor and means.
167 + eta = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)
168 + mu = eta.map { inverseLink(family: family, eta: $0) }
169 + let newLogLikelihood = self.logLikelihood(y: y, mu: mu, family: family)
170 +
171 + if newLogLikelihood.isNaN || newLogLikelihood.isInfinite {
172 + throw ZQStatsError(
173 + "\(family.rawValue) failed to converge (perfect prediction?)"
174 + )
175 + }
176 + let done = abs(newLogLikelihood - logLikelihood)
177 + < 1e-13 * (abs(newLogLikelihood) + 0.1)
178 + logLikelihood = newLogLikelihood
179 + if done { break }
180 + if iteration == maximumIterations {
181 + throw ZQStatsError(
182 + "\(family.rawValue) did not converge in \(maximumIterations) iterations"
183 + )
184 + }
185 + }
186 +
187 + // Bread: (X'WX)⁻¹ with the weights re-evaluated AT the converged β.
188 + // The last IRLS factorization carries weights from the previous
189 + // iterate (off by ~√tolerance), which would contaminate standard
190 + // errors at ~1e-7 relative — visible against R fixtures.
191 + guard finalQR != nil else {
192 + throw ZQStatsError("\(family.rawValue): no iterations performed")
193 + }
194 + let bread: [Double]
195 + do {
196 + var xw = [Double](repeating: 0, count: n * k)
197 + for i in 0..<n {
198 + let derivative = meanDerivative(family: family, eta: eta[i], mean: mu[i])
199 + let varMu = varianceFunction(family: family, mean: mu[i])
200 + let root = (derivative * derivative / varMu).squareRoot()
201 + for j in 0..<k {
202 + xw[j * n + i] = x[j * n + i] * root
203 + }
204 + }
205 + let informationQR = try LinearAlgebra.QR(matrix: xw, rows: n, cols: k)
206 + bread = try informationQR.crossProductInverse()
207 + } catch {
208 + throw ZQStatsError("\(family.rawValue): information matrix is singular")
209 + }
210 +
211 + // Covariance.
212 + var clusterCount: Int? = nil
213 + let vce: [Double]
214 + switch variance {
215 + case .classical:
216 + vce = bread
217 +
218 + case .hc0, .hc1:
219 + // Score sandwich: u_i = x_i (y−μ) μ′/V(μ). No df correction —
220 + // Stata's `robust` for ML.
221 + let meat = scoreMeat(
222 + x: x, y: y, mu: mu, eta: eta, n: n, k: k, family: family
223 + )
224 + vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)
225 +
226 + case .cluster(let groups):
227 + guard groups.count == n else {
228 + throw ZQStatsError("cluster variable has wrong length")
229 + }
230 + var scores: [Int: [Double]] = [:]
231 + for i in 0..<n {
232 + let scale = scoreScale(family: family, y: y[i], mu: mu[i], eta: eta[i])
233 + var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)
234 + for j in 0..<k {
235 + u[j] += x[j * n + i] * scale
236 + }
237 + scores[groups[i]] = u
238 + }
239 + let g = scores.count
240 + guard g > 1 else {
241 + throw ZQStatsError("cluster variable must define at least 2 groups")
242 + }
243 + clusterCount = g
244 + var meat = [Double](repeating: 0, count: k * k)
245 + for u in scores.values {
246 + for j in 0..<k {
247 + for i in 0..<k {
248 + meat[j * k + i] += u[i] * u[j]
249 + }
250 + }
251 + }
252 + // Stata ML cluster factor: G/(G−1).
253 + let scale = Double(g) / Double(g - 1)
254 + for index in meat.indices { meat[index] *= scale }
255 + vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)
256 +
257 + case .hc2, .hc3:
258 + fatalError("unreachable — rejected above")
259 + }
260 +
261 + // Inference: z statistics (normal), Stata ML convention.
262 + let zCritical = ZQDistributions.normalQuantile(0.5 + confidenceLevel / 2)
263 + var coefficients: [ZQCoefficient] = []
264 + for j in 0..<k {
265 + let se = vce[j * k + j].squareRoot()
266 + let zStat = beta[j] / se
267 + coefficients.append(ZQCoefficient(
268 + name: names[j],
269 + estimate: beta[j],
270 + standardError: se,
271 + tStatistic: zStat,
272 + pValue: 2 * (1 - ZQDistributions.normalCDF(abs(zStat))),
273 + confidenceLower: beta[j] - zCritical * se,
274 + confidenceUpper: beta[j] + zCritical * se
275 + ))
276 + }
277 +
278 + // Null (intercept-only) log likelihood — closed form: constant mean.
279 + let meanY = y.reduce(0, +) / Double(n)
280 + let nullMu = [Double](repeating: meanY, count: n)
281 + let nullLogLikelihood = self.logLikelihood(y: y, mu: nullMu, family: family)
282 +
283 + let dfModel = k - (includeConstant ? 1 : 0)
284 + var chiSquared: Double? = nil
285 + var chiSquaredPValue: Double? = nil
286 + if dfModel > 0, includeConstant {
287 + switch variance {
288 + case .classical:
289 + let lr = 2 * (logLikelihood - nullLogLikelihood)
290 + chiSquared = lr
291 + chiSquaredPValue = ZQDistributions.chiSquarePValue(lr, df: Double(dfModel))
292 + default:
293 + // Wald χ² on the non-constant coefficients under the
294 + // robust/cluster VCE.
295 + var subVce = [Double](repeating: 0, count: dfModel * dfModel)
296 + var subBeta = [Double](repeating: 0, count: dfModel)
297 + for j in 0..<dfModel {
298 + subBeta[j] = beta[j]
299 + for i in 0..<dfModel {
300 + subVce[j * dfModel + i] = vce[j * k + i]
301 + }
302 + }
303 + if let solved = try? LinearAlgebra.solveSymmetric(
304 + subVce, k: dfModel, rhs: subBeta
305 + ) {
306 + let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }
307 + chiSquared = wald
308 + chiSquaredPValue = ZQDistributions.chiSquarePValue(
309 + wald, df: Double(dfModel)
310 + )
311 + }
312 + }
313 + }
314 +
315 + return ZQGLMResult(
316 + family: family,
317 + coefficients: coefficients,
318 + observationCount: n,
319 + logLikelihood: logLikelihood,
320 + nullLogLikelihood: nullLogLikelihood,
321 + pseudoRSquared: 1 - logLikelihood / nullLogLikelihood,
322 + chiSquared: chiSquared,
323 + chiSquaredDF: dfModel,
324 + chiSquaredPValue: chiSquaredPValue,
325 + iterations: iterations,
326 + clusterCount: clusterCount,
327 + fittedMeans: mu
328 + )
329 + }
330 +
331 + // MARK: - Family functions
332 +
333 + private static func startingMeans(y: [Double], family: ZQGLMFamily) -> [Double] {
334 + switch family {
335 + case .logit, .probit:
336 + return y.map { ($0 + 0.5) / 2 }
337 + case .poisson:
338 + return y.map { max($0, 0.1) }
339 + }
340 + }
341 +
342 + private static func link(family: ZQGLMFamily, mean: Double) -> Double {
343 + switch family {
344 + case .logit: return Foundation.log(mean / (1 - mean))
345 + case .probit: return ZQDistributions.normalQuantile(mean)
346 + case .poisson: return Foundation.log(mean)
347 + }
348 + }
349 +
350 + private static func inverseLink(family: ZQGLMFamily, eta: Double) -> Double {
351 + switch family {
352 + case .logit:
353 + let clamped = min(max(eta, -30), 30)
354 + return 1 / (1 + Foundation.exp(-clamped))
355 + case .probit:
356 + let p = ZQDistributions.normalCDF(eta)
357 + return min(max(p, 1e-12), 1 - 1e-12)
358 + case .poisson:
359 + return Foundation.exp(min(eta, 300))
360 + }
361 + }
362 +
363 + /// dμ/dη at the current point.
364 + private static func meanDerivative(
365 + family: ZQGLMFamily, eta: Double, mean: Double
366 + ) -> Double {
367 + switch family {
368 + case .logit:
369 + return max(mean * (1 - mean), 1e-12)
370 + case .probit:
371 + return max(normalDensity(eta), 1e-12)
372 + case .poisson:
373 + return max(mean, 1e-12)
374 + }
375 + }
376 +
377 + /// Var(Y | μ) up to the (unit) dispersion.
378 + private static func varianceFunction(family: ZQGLMFamily, mean: Double) -> Double {
379 + switch family {
380 + case .logit, .probit:
381 + return max(mean * (1 - mean), 1e-12)
382 + case .poisson:
383 + return max(mean, 1e-12)
384 + }
385 + }
386 +
387 + /// Per-observation score scale: (y−μ)·μ′/V(μ), so the score is that
388 + /// times x_i. For canonical links this collapses to y−μ.
389 + private static func scoreScale(
390 + family: ZQGLMFamily, y: Double, mu: Double, eta: Double
391 + ) -> Double {
392 + switch family {
393 + case .logit, .poisson:
394 + return y - mu
395 + case .probit:
396 + let density = normalDensity(eta)
397 + return (y - mu) * density / max(mu * (1 - mu), 1e-12)
398 + }
399 + }
400 +
401 + private static func scoreMeat(
402 + x: [Double], y: [Double], mu: [Double], eta: [Double],
403 + n: Int, k: Int, family: ZQGLMFamily
404 + ) -> [Double] {
405 + var meat = [Double](repeating: 0, count: k * k)
406 + for row in 0..<n {
407 + let scale = scoreScale(family: family, y: y[row], mu: mu[row], eta: eta[row])
408 + for j in 0..<k {
409 + let xj = x[j * n + row] * scale
410 + for i in j..<k {
411 + meat[j * k + i] += xj * x[i * n + row] * scale
412 + }
413 + }
414 + }
415 + for j in 0..<k {
416 + for i in 0..<j {
417 + meat[j * k + i] = meat[i * k + j]
418 + }
419 + }
420 + return meat
421 + }
422 +
423 + private static func logLikelihood(
424 + y: [Double], mu: [Double], family: ZQGLMFamily
425 + ) -> Double {
426 + var total = 0.0
427 + switch family {
428 + case .logit, .probit:
429 + for i in 0..<y.count {
430 + let p = min(max(mu[i], 1e-12), 1 - 1e-12)
431 + total += y[i] == 1 ? Foundation.log(p) : Foundation.log(1 - p)
432 + }
433 + case .poisson:
434 + for i in 0..<y.count {
435 + let m = max(mu[i], 1e-300)
436 + total += y[i] * Foundation.log(m) - m - ZQDistributions.logGamma(y[i] + 1)
437 + }
438 + }
439 + return total
440 + }
441 +
442 + private static func normalDensity(_ z: Double) -> Double {
443 + Foundation.exp(-0.5 * z * z) / (2 * Double.pi).squareRoot()
444 + }
445 +}
modified MetrikaKit/Tests/MetrikaKitTests/EngineTests.swift +51 −0
@@ -132,6 +132,57 @@ struct EngineTests {
132 132 #expect(plot.xLabel == "price")
133 133 }
134 134
135 + @Test("logit through the console matches R")
136 + func logitCommand() async throws {
137 + let result = try await session.execute("logit purchase price")
138 + expectClose(
139 + try #require(result.scalars["b_price"]), fixtures["logit_b_price"], "b[price]"
140 + )
141 + expectClose(try #require(result.scalars["ll"]), fixtures["logit_ll"], "ll")
142 + #expect(result.text.contains("Logistic regression"))
143 + #expect(result.text.contains("Pseudo R2"))
144 + }
145 +
146 + @Test("poisson with robust SE through the console")
147 + func poissonCommand() async throws {
148 + let result = try await session.execute("poisson orders price, robust")
149 + expectClose(
150 + try #require(result.scalars["b_price"]), fixtures["pois_b_price"], "b[price]"
151 + )
152 + expectClose(
153 + try #require(result.scalars["se_price"]),
154 + fixtures["pois_se_hc0_price"], "robust se[price]"
155 + )
156 + }
157 +
158 + @Test("tabulate one-way counts region levels")
159 + func tabulateOneWay() async throws {
160 + let result = try await session.execute("tab region")
161 + #expect(result.scalars["N"] == 60)
162 + #expect(result.scalars["r"] == 3)
163 + #expect(result.text.contains("Total"))
164 + // region cycles 1,2,3 over 60 rows → 20 each.
165 + #expect(result.text.contains("20"))
166 + }
167 +
168 + @Test("tabulate two-way region by purchase")
169 + func tabulateTwoWay() async throws {
170 + let result = try await session.execute("tabulate region purchase")
171 + #expect(result.scalars["r"] == 3)
172 + #expect(result.scalars["c"] == 2)
173 + #expect(result.scalars["N"] == 60)
174 + }
175 +
176 + @Test("correlate matches R and reports listwise obs")
177 + func correlateCommand() async throws {
178 + let result = try await session.execute("correlate revenue price")
179 + expectClose(
180 + try #require(result.scalars["rho"]),
181 + fixtures["corr_rev_price"], "cor(revenue, price)"
182 + )
183 + #expect(result.scalars["N"] == fixtures["ols_n"])
184 + }
185 +
135 186 @Test("display evaluates scalar expressions")
136 187 func display() async throws {
137 188 let result = try await session.execute("display 2 + 2 * 3")
modified MetrikaKit/Tests/MetrikaKitTests/Fixtures/expected.tsv +24 −0
@@ -37,6 +37,30 @@ ols2_b_region2 0.071445724920196046
37 37 ols2_b_region3 0.21232048832355585
38 38 ols2_b_cons 3.9851984093144015
39 39 ols2_r2 0.87452597383166375
40 +logit_b_price -0.11516614916659125
41 +logit_b_cons 1.0575375744587068
42 +logit_se_price 0.063597140127676768
43 +logit_se_cons 0.87352098060763228
44 +logit_ll -36.287874252110491
45 +logit_ll0 -38.013806140556945
46 +logit_chi2 3.4518637768929068
47 +logit_chi2p 0.063180488513405839
48 +logit_se_hc0_price 0.063596861213014763
49 +logit_se_cluster_price 0.069446458133887223
50 +probit_b_price -0.070442868008111323
51 +probit_b_cons 0.64722038725062636
52 +probit_se_price 0.038635906815363288
53 +probit_ll -36.305228252679598
54 +pois_b_price -0.085650626748355135
55 +pois_b_cons 1.3512956388566333
56 +pois_se_price 0.025531826392571365
57 +pois_ll -81.84353537859181
58 +pois_se_hc0_price 0.026876084479165823
59 +corr_rev_price -0.88919780339930277
60 +corr_rev_logrev 0.98089325161602359
61 +dist_pchisq_3p8_1 0.94874741714263044
62 +dist_pchisq_25_4 0.99994969018217694
63 +dist_qnorm_0p975 1.9599639845400534
40 64 dist_pt_2p5_df10 0.98427657788169554
41 65 dist_pt_m1p3_df3 0.14223375436394847
42 66 dist_pt_0p05_df57 0.51985140594489099
modified MetrikaKit/Tests/MetrikaKitTests/Fixtures/regression.csv +61 −61
@@ -1,61 +1,61 @@
1 revenue,price,region,firm_id
2 13.073013,18.7221,1,1
3 14.603161,19.0561,2,1
4 37.036281,9.2921,3,1
5 12.331106,17.4567,1,1
6 20.199281,14.6262,2,1
7 18.532705,12.7864,3,2
8 13.442576,,1,2
9 30.288376,7.02,2,2
10 14.146638,14.8549,3,2
11 15.789387,15.576,1,2
12 24.085548,11.8661,2,3
13 17.866099,15.7867,3,3
14 13.358007,19.0201,1,3
15 26.695153,8.8314,2,3
16 20.905369,11.9344,3,3
17 12.640513,19.1002,1,4
18 11.071996,19.6734,2,4
19 48.213111,6.7623,3,4
20 19.400193,12.125,1,4
21 22.781459,13.405,2,4
22 15.854431,18.5605,3,5
23 27.549049,7.0807,1,5
24 15.637737,,2,5
25 15.806717,19.2,3,5
26 33.600334,6.2366,1,5
27 22.746857,12.7132,2,6
28 30.989393,10.8531,3,6
29 12.510734,18.5861,1,6
30 15.099717,11.7045,2,6
31 17.10762,17.5401,3,6
32 14.293473,16.0639,1,7
33 15.713512,17.1658,2,7
34 30.616712,10.8216,3,7
35 19.841208,15.2775,1,7
36 36.095565,5.0592,2,7
37 20.003055,17.4937,3,8
38 38.151955,5.11,1,8
39 36.839992,8.1149,2,8
40 17.29105,18.599,3,8
41 19.569965,14.1767,1,8
42 21.934484,,2,9
43 26.14202,11.5366,3,9
44 38.421147,5.5615,1,9
45 10.899566,19.6031,2,9
46 24.544112,11.4763,3,9
47 12.655109,19.3636,1,10
48 15.641427,18.3163,2,10
49 22.233252,14.5997,3,10
50 9.993891,19.5645,1,10
51 16.320549,14.2826,2,10
52 37.592103,10.0014,3,11
53 25.093077,10.2012,1,11
54 25.408455,10.9773,2,11
55 17.119771,16.7704,3,11
56 29.198899,5.584,1,11
57 18.051847,16.2319,2,12
58 19.195829,15.1592,3,12
59 28.993336,7.569,1,12
60 34.011476,8.9163,2,12
61 27.274985,12.7162,3,12
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
added MetrikaKit/Tests/MetrikaKitTests/GLMTests.swift +163 −0
@@ -0,0 +1,163 @@
1 +//
2 +// GLMTests.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 ZQData
13 +import ZQStats
14 +
15 +/// GLM estimators vs R glm() fixtures (epsilon 1e-12) at 1e-10 relative
16 +/// tolerance.
17 +@Suite("ZQGLM vs R fixtures")
18 +struct GLMTests {
19 + let fixtures: Fixtures
20 + let purchase: [Double]
21 + let orders: [Double]
22 + let price: [Double]
23 + let firmID: [Int]
24 +
25 + init() async throws {
26 + self.fixtures = try Fixtures()
27 + let store = try ZQDataStore()
28 + let frame = try await store.load(contentsOf: fixtures.datasetURL)
29 + let (purchaseAll, _) = try frame.requireNumeric("purchase")
30 + let (ordersAll, _) = try frame.requireNumeric("orders")
31 + let (priceAll, priceMissing) = try frame.requireNumeric("price")
32 + let (firmAll, _) = try frame.requireNumeric("firm_id")
33 +
34 + var purchase: [Double] = [], orders: [Double] = []
35 + var price: [Double] = [], firm: [Int] = []
36 + for i in 0..<frame.rowCount where !priceMissing[i] {
37 + purchase.append(purchaseAll[i])
38 + orders.append(ordersAll[i])
39 + price.append(priceAll[i])
40 + firm.append(Int(firmAll[i]))
41 + }
42 + self.purchase = purchase
43 + self.orders = orders
44 + self.price = price
45 + self.firmID = firm
46 + }
47 +
48 + @Test("logit matches R glm(binomial)")
49 + func logit() throws {
50 + let result = try ZQGLM.fit(
51 + y: purchase, predictors: [("price", price)], family: .logit
52 + )
53 + expectClose(result.coefficients[0].estimate, fixtures["logit_b_price"], "b[price]")
54 + expectClose(result.coefficients[1].estimate, fixtures["logit_b_cons"], "b[_cons]")
55 + expectClose(
56 + result.coefficients[0].standardError, fixtures["logit_se_price"], "se[price]"
57 + )
58 + expectClose(
59 + result.coefficients[1].standardError, fixtures["logit_se_cons"], "se[_cons]"
60 + )
61 + expectClose(result.logLikelihood, fixtures["logit_ll"], "log likelihood")
62 + expectClose(result.nullLogLikelihood, fixtures["logit_ll0"], "null log likelihood")
63 + expectClose(try #require(result.chiSquared), fixtures["logit_chi2"], "LR chi2")
64 + expectClose(
65 + try #require(result.chiSquaredPValue),
66 + fixtures["logit_chi2p"], rtol: 1e-9, "Prob > chi2"
67 + )
68 + }
69 +
70 + @Test("logit robust and cluster standard errors")
71 + func logitRobust() throws {
72 + let robust = try ZQGLM.fit(
73 + y: purchase, predictors: [("price", price)], family: .logit, variance: .hc0
74 + )
75 + expectClose(
76 + robust.coefficients[0].standardError,
77 + fixtures["logit_se_hc0_price"], "robust se[price]"
78 + )
79 + let clustered = try ZQGLM.fit(
80 + y: purchase, predictors: [("price", price)], family: .logit,
81 + variance: .cluster(firmID)
82 + )
83 + expectClose(
84 + clustered.coefficients[0].standardError,
85 + fixtures["logit_se_cluster_price"], "cluster se[price]"
86 + )
87 + }
88 +
89 + @Test("probit matches R glm(binomial(probit))")
90 + func probit() throws {
91 + let result = try ZQGLM.fit(
92 + y: purchase, predictors: [("price", price)], family: .probit
93 + )
94 + expectClose(result.coefficients[0].estimate, fixtures["probit_b_price"], "b[price]")
95 + expectClose(result.coefficients[1].estimate, fixtures["probit_b_cons"], "b[_cons]")
96 + expectClose(
97 + result.coefficients[0].standardError, fixtures["probit_se_price"], "se[price]"
98 + )
99 + expectClose(result.logLikelihood, fixtures["probit_ll"], "log likelihood")
100 + }
101 +
102 + @Test("poisson matches R glm(poisson)")
103 + func poisson() throws {
104 + let result = try ZQGLM.fit(
105 + y: orders, predictors: [("price", price)], family: .poisson
106 + )
107 + expectClose(result.coefficients[0].estimate, fixtures["pois_b_price"], "b[price]")
108 + expectClose(result.coefficients[1].estimate, fixtures["pois_b_cons"], "b[_cons]")
109 + expectClose(
110 + result.coefficients[0].standardError, fixtures["pois_se_price"], "se[price]"
111 + )
112 + expectClose(result.logLikelihood, fixtures["pois_ll"], "log likelihood")
113 +
114 + let robust = try ZQGLM.fit(
115 + y: orders, predictors: [("price", price)], family: .poisson, variance: .hc0
116 + )
117 + expectClose(
118 + robust.coefficients[0].standardError,
119 + fixtures["pois_se_hc0_price"], "robust se[price]"
120 + )
121 + }
122 +
123 + @Test("logit rejects non-binary outcomes")
124 + func binaryValidation() {
125 + #expect(throws: ZQStatsError.self) {
126 + _ = try ZQGLM.fit(
127 + y: orders, predictors: [("price", price)], family: .logit
128 + )
129 + }
130 + }
131 +
132 + @Test("correlate matches R cor()")
133 + func correlate() async throws {
134 + let store = try ZQDataStore()
135 + let frame = try await store.load(contentsOf: fixtures.datasetURL)
136 + let (revenueAll, _) = try frame.requireNumeric("revenue")
137 + let (priceAll, priceMissing) = try frame.requireNumeric("price")
138 + var revenue: [Double] = [], priceComplete: [Double] = []
139 + for i in 0..<frame.rowCount where !priceMissing[i] {
140 + revenue.append(revenueAll[i])
141 + priceComplete.append(priceAll[i])
142 + }
143 + let matrix = ZQCorrelate.matrix(columns: [revenue, priceComplete])
144 + expectClose(matrix[0][1], fixtures["corr_rev_price"], "cor(revenue, price)")
145 + #expect(matrix[0][0] == 1 && matrix[1][1] == 1)
146 + }
147 +
148 + @Test("chi-square CDF and normal quantile vs R")
149 + func chiSquare() {
150 + expectClose(
151 + ZQDistributions.chiSquareCDF(3.8, df: 1),
152 + fixtures["dist_pchisq_3p8_1"], rtol: 1e-12, "pchisq(3.8, 1)"
153 + )
154 + expectClose(
155 + ZQDistributions.chiSquareCDF(25, df: 4),
156 + fixtures["dist_pchisq_25_4"], rtol: 1e-12, "pchisq(25, 4)"
157 + )
158 + expectClose(
159 + ZQDistributions.normalQuantile(0.975),
160 + fixtures["dist_qnorm_0p975"], rtol: 1e-12, "qnorm(0.975)"
161 + )
162 + }
163 +}
modified Tests/Fixtures/generate.R +78 −1
@@ -32,6 +32,11 @@ noise <- round(rnorm(n, 0, 0.15), 6)
32 32 log_rev <- 4 - 0.08 * price + 0.10 * (region == 2) + 0.20 * (region == 3) + noise
33 33 revenue <- round(exp(log_rev), 6)
34 34
35 +# GLM outcomes — drawn AFTER all draws above so earlier fixture values
36 +# stay bit-identical when regenerating.
37 +purchase <- rbinom(n, 1, plogis(2 - 0.20 * price))
38 +orders <- rpois(n, exp(1.6 - 0.09 * price))
39 +
35 40 # Inject missing values to exercise listwise deletion.
36 41 price_missing <- price
37 42 price_missing[c(7, 23, 41)] <- NA
@@ -40,7 +45,9 @@ data <- data.frame(
40 45 revenue = revenue,
41 46 price = price_missing,
42 47 region = region,
43 firm_id = firm_id
48 + firm_id = firm_id,
49 + purchase = purchase,
50 + orders = orders
44 51 )
45 52 csv_path <- file.path(out_dir, "regression.csv")
46 53 write.csv(data, csv_path, row.names = FALSE, quote = FALSE, na = "")
@@ -146,7 +153,77 @@ emit("ols2_b_region3", b2[["factor(region)3"]])
146 153 emit("ols2_b_cons", b2[["(Intercept)"]])
147 154 emit("ols2_r2", summary(fit2)$r.squared)
148 155
156 +# ------------------------------------------------------------------ GLMs
157 +# Fit tightly (epsilon 1e-12) so R and Swift land on the same optimum to
158 +# well past the 1e-10 test tolerance. Estimation sample = complete cases.
159 +ctrl <- glm.control(epsilon = 1e-12, maxit = 100)
160 +
161 +# Expected information (X'WX)^-1 evaluated AT the converged coefficients.
162 +# R's vcov(fit) instead reuses the weights of the last IRLS step (at the
163 +# second-to-last iterate), so it is only ~1e-6-accurate by its own
164 +# stopping rule — not good enough for 1e-10 fixtures.
165 +glm_bread <- function(fit) {
166 + Xg <- model.matrix(fit)
167 + eta <- as.vector(Xg %*% coef(fit))
168 + fam <- fit$family
169 + mu <- fam$linkinv(eta)
170 + W <- fam$mu.eta(eta)^2 / fam$variance(mu)
171 + solve(t(Xg * W) %*% Xg)
172 +}
173 +
174 +glm_bread_sandwich <- function(fit, score_scale, cl = NULL) {
175 + Xg <- model.matrix(fit)
176 + u <- Xg * score_scale
177 + bread <- glm_bread(fit)
178 + if (is.null(cl)) {
179 + meat <- t(u) %*% u
180 + } else {
181 + ug <- rowsum(u, cl)
182 + G <- nrow(ug)
183 + meat <- (G / (G - 1)) * (t(ug) %*% ug)
184 + }
185 + sqrt(diag(bread %*% meat %*% bread))
186 +}
187 +
188 +lg <- glm(purchase ~ price, data = d, family = binomial(), control = ctrl)
189 +V_lg <- glm_bread(lg)
190 +emit("logit_b_price", coef(lg)[["price"]])
191 +emit("logit_b_cons", coef(lg)[["(Intercept)"]])
192 +emit("logit_se_price", sqrt(V_lg["price", "price"]))
193 +emit("logit_se_cons", sqrt(V_lg["(Intercept)", "(Intercept)"]))
194 +emit("logit_ll", as.numeric(logLik(lg)))
195 +lg0 <- glm(purchase ~ 1, data = d, family = binomial(), control = ctrl)
196 +emit("logit_ll0", as.numeric(logLik(lg0)))
197 +emit("logit_chi2", 2 * (as.numeric(logLik(lg)) - as.numeric(logLik(lg0))))
198 +emit("logit_chi2p", pchisq(
199 + 2 * (as.numeric(logLik(lg)) - as.numeric(logLik(lg0))), 1, lower.tail = FALSE
200 +))
201 +emit("logit_se_hc0_price", glm_bread_sandwich(lg, d$purchase - fitted(lg))["price"])
202 +emit("logit_se_cluster_price",
203 + glm_bread_sandwich(lg, d$purchase - fitted(lg), cl = d$firm_id)["price"])
204 +
205 +pr <- glm(purchase ~ price, data = d,
206 + family = binomial(link = "probit"), control = ctrl)
207 +emit("probit_b_price", coef(pr)[["price"]])
208 +emit("probit_b_cons", coef(pr)[["(Intercept)"]])
209 +emit("probit_se_price", sqrt(glm_bread(pr)["price", "price"]))
210 +emit("probit_ll", as.numeric(logLik(pr)))
211 +
212 +ps <- glm(orders ~ price, data = d, family = poisson(), control = ctrl)
213 +emit("pois_b_price", coef(ps)[["price"]])
214 +emit("pois_b_cons", coef(ps)[["(Intercept)"]])
215 +emit("pois_se_price", sqrt(glm_bread(ps)["price", "price"]))
216 +emit("pois_ll", as.numeric(logLik(ps)))
217 +emit("pois_se_hc0_price", glm_bread_sandwich(ps, d$orders - fitted(ps))["price"])
218 +
219 +# -------------------------------------------------------------- correlate
220 +emit("corr_rev_price", cor(d$revenue, d$price))
221 +emit("corr_rev_logrev", cor(d$revenue, d$log_rev))
222 +
149 223 # --------------------------------------------------------- distributions
224 +emit("dist_pchisq_3p8_1", pchisq(3.8, 1))
225 +emit("dist_pchisq_25_4", pchisq(25, 4))
226 +emit("dist_qnorm_0p975", qnorm(0.975))
150 227 emit("dist_pt_2p5_df10", pt(2.5, 10))
151 228 emit("dist_pt_m1p3_df3", pt(-1.3, 3))
152 229 emit("dist_pt_0p05_df57", pt(0.05, 57))
153 230