diff-diff 1.4.0__tar.gz → 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. {diff_diff-1.4.0 → diff_diff-2.0.0}/PKG-INFO +19 -18
  2. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/__init__.py +13 -1
  3. diff_diff-2.0.0/diff_diff/_backend.py +64 -0
  4. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/linalg.py +151 -4
  5. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/staggered.py +151 -26
  6. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/utils.py +39 -9
  7. {diff_diff-1.4.0 → diff_diff-2.0.0}/pyproject.toml +14 -5
  8. diff_diff-2.0.0/rust/Cargo.lock +2396 -0
  9. diff_diff-2.0.0/rust/Cargo.toml +39 -0
  10. diff_diff-2.0.0/rust/src/bootstrap.rs +223 -0
  11. diff_diff-2.0.0/rust/src/lib.rs +33 -0
  12. diff_diff-2.0.0/rust/src/linalg.rs +229 -0
  13. diff_diff-2.0.0/rust/src/weights.rs +220 -0
  14. diff_diff-1.4.0/diff_diff.egg-info/PKG-INFO +0 -2255
  15. diff_diff-1.4.0/diff_diff.egg-info/SOURCES.txt +0 -38
  16. diff_diff-1.4.0/diff_diff.egg-info/dependency_links.txt +0 -1
  17. diff_diff-1.4.0/diff_diff.egg-info/requires.txt +0 -14
  18. diff_diff-1.4.0/diff_diff.egg-info/top_level.txt +0 -1
  19. diff_diff-1.4.0/setup.cfg +0 -4
  20. diff_diff-1.4.0/tests/test_bacon.py +0 -679
  21. diff_diff-1.4.0/tests/test_diagnostics.py +0 -674
  22. diff_diff-1.4.0/tests/test_estimators.py +0 -2719
  23. diff_diff-1.4.0/tests/test_honest_did.py +0 -699
  24. diff_diff-1.4.0/tests/test_linalg.py +0 -422
  25. diff_diff-1.4.0/tests/test_power.py +0 -691
  26. diff_diff-1.4.0/tests/test_prep.py +0 -794
  27. diff_diff-1.4.0/tests/test_pretrends.py +0 -813
  28. diff_diff-1.4.0/tests/test_staggered.py +0 -1358
  29. diff_diff-1.4.0/tests/test_sun_abraham.py +0 -732
  30. diff_diff-1.4.0/tests/test_triple_diff.py +0 -869
  31. diff_diff-1.4.0/tests/test_utils.py +0 -1270
  32. diff_diff-1.4.0/tests/test_visualization.py +0 -284
  33. diff_diff-1.4.0/tests/test_wild_bootstrap.py +0 -804
  34. {diff_diff-1.4.0 → diff_diff-2.0.0}/README.md +0 -0
  35. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/bacon.py +0 -0
  36. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/diagnostics.py +0 -0
  37. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/estimators.py +0 -0
  38. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/honest_did.py +0 -0
  39. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/power.py +0 -0
  40. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/prep.py +0 -0
  41. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/pretrends.py +0 -0
  42. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/results.py +0 -0
  43. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/sun_abraham.py +0 -0
  44. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/synthetic_did.py +0 -0
  45. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/triple_diff.py +0 -0
  46. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/twfe.py +0 -0
  47. {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/visualization.py +0 -0
@@ -1,14 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diff-diff
3
- Version: 1.4.0
4
- Summary: A library for Difference-in-Differences causal inference analysis
5
- Author: diff-diff contributors
6
- License-Expression: MIT
7
- Project-URL: Homepage, https://github.com/igerber/diff-diff
8
- Project-URL: Documentation, https://diff-diff.readthedocs.io
9
- Project-URL: Repository, https://github.com/igerber/diff-diff
10
- Project-URL: Issues, https://github.com/igerber/diff-diff/issues
11
- Keywords: causal-inference,difference-in-differences,econometrics,statistics,treatment-effects
3
+ Version: 2.0.0
12
4
  Classifier: Development Status :: 5 - Production/Stable
13
5
  Classifier: Intended Audience :: Science/Research
14
6
  Classifier: Operating System :: OS Independent
@@ -18,20 +10,28 @@ Classifier: Programming Language :: Python :: 3.10
18
10
  Classifier: Programming Language :: Python :: 3.11
19
11
  Classifier: Programming Language :: Python :: 3.12
20
12
  Classifier: Topic :: Scientific/Engineering :: Mathematics
21
- Requires-Python: >=3.9
22
- Description-Content-Type: text/markdown
23
13
  Requires-Dist: numpy>=1.20.0
24
14
  Requires-Dist: pandas>=1.3.0
25
15
  Requires-Dist: scipy>=1.7.0
16
+ Requires-Dist: pytest>=7.0 ; extra == 'dev'
17
+ Requires-Dist: pytest-cov>=4.0 ; extra == 'dev'
18
+ Requires-Dist: black>=23.0 ; extra == 'dev'
19
+ Requires-Dist: ruff>=0.1.0 ; extra == 'dev'
20
+ Requires-Dist: mypy>=1.0 ; extra == 'dev'
21
+ Requires-Dist: sphinx>=6.0 ; extra == 'docs'
22
+ Requires-Dist: sphinx-rtd-theme>=1.0 ; extra == 'docs'
26
23
  Provides-Extra: dev
27
- Requires-Dist: pytest>=7.0; extra == "dev"
28
- Requires-Dist: pytest-cov>=4.0; extra == "dev"
29
- Requires-Dist: black>=23.0; extra == "dev"
30
- Requires-Dist: ruff>=0.1.0; extra == "dev"
31
- Requires-Dist: mypy>=1.0; extra == "dev"
32
24
  Provides-Extra: docs
33
- Requires-Dist: sphinx>=6.0; extra == "docs"
34
- Requires-Dist: sphinx-rtd-theme>=1.0; extra == "docs"
25
+ Summary: A library for Difference-in-Differences causal inference analysis
26
+ Keywords: causal-inference,difference-in-differences,econometrics,statistics,treatment-effects
27
+ Author: diff-diff contributors
28
+ License-Expression: MIT
29
+ Requires-Python: >=3.9
30
+ Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
31
+ Project-URL: Documentation, https://diff-diff.readthedocs.io
32
+ Project-URL: Homepage, https://github.com/igerber/diff-diff
33
+ Project-URL: Issues, https://github.com/igerber/diff-diff/issues
34
+ Project-URL: Repository, https://github.com/igerber/diff-diff
35
35
 
36
36
  # diff-diff
37
37
 
@@ -2253,3 +2253,4 @@ The `HonestDiD` module implements sensitivity analysis methods for relaxing the
2253
2253
  ## License
2254
2254
 
2255
2255
  MIT License
2256
+
@@ -5,6 +5,16 @@ This library provides sklearn-like estimators for causal inference
5
5
  using the difference-in-differences methodology.
6
6
  """
7
7
 
8
+ # Import backend detection from dedicated module (avoids circular imports)
9
+ from diff_diff._backend import (
10
+ HAS_RUST_BACKEND,
11
+ _rust_bootstrap_weights,
12
+ _rust_compute_robust_vcov,
13
+ _rust_project_simplex,
14
+ _rust_solve_ols,
15
+ _rust_synthetic_weights,
16
+ )
17
+
8
18
  from diff_diff.bacon import (
9
19
  BaconDecomposition,
10
20
  BaconDecompositionResults,
@@ -103,7 +113,7 @@ from diff_diff.visualization import (
103
113
  plot_sensitivity,
104
114
  )
105
115
 
106
- __version__ = "1.4.0"
116
+ __version__ = "2.0.0"
107
117
  __all__ = [
108
118
  # Estimators
109
119
  "DifferenceInDifferences",
@@ -187,4 +197,6 @@ __all__ = [
187
197
  "compute_pretrends_power",
188
198
  "compute_mdv",
189
199
  "plot_pretrends_power",
200
+ # Rust backend
201
+ "HAS_RUST_BACKEND",
190
202
  ]
@@ -0,0 +1,64 @@
1
+ """
2
+ Backend detection and configuration for diff-diff.
3
+
4
+ This module handles:
5
+ 1. Detection of optional Rust backend
6
+ 2. Environment variable configuration (DIFF_DIFF_BACKEND)
7
+ 3. Exports HAS_RUST_BACKEND and Rust function references
8
+
9
+ Other modules should import from here to avoid circular imports with __init__.py.
10
+ """
11
+
12
+ import os
13
+
14
+ # Check for backend override via environment variable
15
+ # DIFF_DIFF_BACKEND can be: 'auto' (default), 'python', or 'rust'
16
+ _backend_env = os.environ.get('DIFF_DIFF_BACKEND', 'auto').lower()
17
+
18
+ # Try to import Rust backend for accelerated operations
19
+ try:
20
+ from diff_diff._rust_backend import (
21
+ generate_bootstrap_weights_batch as _rust_bootstrap_weights,
22
+ compute_synthetic_weights as _rust_synthetic_weights,
23
+ project_simplex as _rust_project_simplex,
24
+ solve_ols as _rust_solve_ols,
25
+ compute_robust_vcov as _rust_compute_robust_vcov,
26
+ )
27
+ _rust_available = True
28
+ except ImportError:
29
+ _rust_available = False
30
+ _rust_bootstrap_weights = None
31
+ _rust_synthetic_weights = None
32
+ _rust_project_simplex = None
33
+ _rust_solve_ols = None
34
+ _rust_compute_robust_vcov = None
35
+
36
+ # Determine final backend based on environment variable and availability
37
+ if _backend_env == 'python':
38
+ # Force pure Python mode - disable Rust even if available
39
+ HAS_RUST_BACKEND = False
40
+ _rust_bootstrap_weights = None
41
+ _rust_synthetic_weights = None
42
+ _rust_project_simplex = None
43
+ _rust_solve_ols = None
44
+ _rust_compute_robust_vcov = None
45
+ elif _backend_env == 'rust':
46
+ # Force Rust mode - fail if not available
47
+ if not _rust_available:
48
+ raise ImportError(
49
+ "DIFF_DIFF_BACKEND=rust but Rust backend is not available. "
50
+ "Install with: pip install diff-diff[rust]"
51
+ )
52
+ HAS_RUST_BACKEND = True
53
+ else:
54
+ # Auto mode - use Rust if available
55
+ HAS_RUST_BACKEND = _rust_available
56
+
57
+ __all__ = [
58
+ 'HAS_RUST_BACKEND',
59
+ '_rust_bootstrap_weights',
60
+ '_rust_synthetic_weights',
61
+ '_rust_project_simplex',
62
+ '_rust_solve_ols',
63
+ '_rust_compute_robust_vcov',
64
+ ]
@@ -1,15 +1,17 @@
1
1
  """
2
2
  Unified linear algebra backend for diff-diff.
3
3
 
4
- This module provides optimized OLS and variance estimation that can be
5
- swapped to a compiled backend (Rust/C++) for maximum performance.
4
+ This module provides optimized OLS and variance estimation with an optional
5
+ Rust backend for maximum performance.
6
6
 
7
7
  The key optimizations are:
8
8
  1. scipy.linalg.lstsq with 'gelsy' driver (QR-based, faster than SVD)
9
9
  2. Vectorized cluster-robust SE via groupby (eliminates O(n*clusters) loop)
10
10
  3. Single interface for all estimators (reduces code duplication)
11
+ 4. Optional Rust backend for additional speedup (when available)
11
12
 
12
- Future: This module can be extended with a Rust backend for additional speedup.
13
+ The Rust backend is automatically used when available, with transparent
14
+ fallback to NumPy/SciPy implementations.
13
15
  """
14
16
 
15
17
  from typing import Optional, Tuple, Union
@@ -18,6 +20,13 @@ import numpy as np
18
20
  import pandas as pd
19
21
  from scipy.linalg import lstsq as scipy_lstsq
20
22
 
23
+ # Import Rust backend if available (from _backend to avoid circular imports)
24
+ from diff_diff._backend import (
25
+ HAS_RUST_BACKEND,
26
+ _rust_compute_robust_vcov,
27
+ _rust_solve_ols,
28
+ )
29
+
21
30
 
22
31
  def solve_ols(
23
32
  X: np.ndarray,
@@ -119,6 +128,87 @@ def solve_ols(
119
128
  "Clean your data or set check_finite=False to skip this check."
120
129
  )
121
130
 
131
+ # Use Rust backend if available
132
+ # Note: Fall back to NumPy if check_finite=False since Rust's LAPACK
133
+ # doesn't support non-finite values
134
+ if HAS_RUST_BACKEND and check_finite:
135
+ # Ensure contiguous arrays for Rust
136
+ X = np.ascontiguousarray(X, dtype=np.float64)
137
+ y = np.ascontiguousarray(y, dtype=np.float64)
138
+
139
+ # Convert cluster_ids to int64 for Rust (if provided)
140
+ cluster_ids_int = None
141
+ if cluster_ids is not None:
142
+ cluster_ids_int = pd.factorize(cluster_ids)[0].astype(np.int64)
143
+
144
+ try:
145
+ coefficients, residuals, vcov = _rust_solve_ols(
146
+ X, y, cluster_ids_int, return_vcov
147
+ )
148
+ except ValueError as e:
149
+ # Translate Rust LAPACK errors to consistent Python error messages
150
+ error_msg = str(e)
151
+ if "Matrix inversion failed" in error_msg or "Least squares failed" in error_msg:
152
+ raise ValueError(
153
+ "Design matrix is rank-deficient (singular X'X matrix). "
154
+ "This indicates perfect multicollinearity. Check your fixed effects "
155
+ "and covariates for linear dependencies."
156
+ ) from e
157
+ raise
158
+
159
+ if return_fitted:
160
+ fitted = X @ coefficients
161
+ return coefficients, residuals, fitted, vcov
162
+ else:
163
+ return coefficients, residuals, vcov
164
+
165
+ # Fallback to NumPy/SciPy implementation
166
+ return _solve_ols_numpy(
167
+ X, y, cluster_ids=cluster_ids, return_vcov=return_vcov, return_fitted=return_fitted
168
+ )
169
+
170
+
171
+ def _solve_ols_numpy(
172
+ X: np.ndarray,
173
+ y: np.ndarray,
174
+ *,
175
+ cluster_ids: Optional[np.ndarray] = None,
176
+ return_vcov: bool = True,
177
+ return_fitted: bool = False,
178
+ ) -> Union[
179
+ Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
180
+ Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
181
+ ]:
182
+ """
183
+ NumPy/SciPy fallback implementation of solve_ols.
184
+
185
+ Uses scipy.linalg.lstsq with 'gelsy' driver (QR with column pivoting)
186
+ for fast and stable least squares solving.
187
+
188
+ Parameters
189
+ ----------
190
+ X : np.ndarray
191
+ Design matrix of shape (n, k).
192
+ y : np.ndarray
193
+ Response vector of shape (n,).
194
+ cluster_ids : np.ndarray, optional
195
+ Cluster identifiers for cluster-robust SEs.
196
+ return_vcov : bool
197
+ Whether to compute variance-covariance matrix.
198
+ return_fitted : bool
199
+ Whether to return fitted values.
200
+
201
+ Returns
202
+ -------
203
+ coefficients : np.ndarray
204
+ OLS coefficients of shape (k,).
205
+ residuals : np.ndarray
206
+ Residuals of shape (n,).
207
+ fitted : np.ndarray, optional
208
+ Fitted values if return_fitted=True.
209
+ vcov : np.ndarray, optional
210
+ Variance-covariance matrix if return_vcov=True.
211
+ """
122
212
  # Solve OLS using scipy's optimized solver
123
213
  # 'gelsy' uses QR with column pivoting, faster than default 'gelsd' (SVD)
124
214
  # Note: gelsy doesn't reliably report rank, so we don't check for deficiency
@@ -131,7 +221,7 @@ def solve_ols(
131
221
  # Compute variance-covariance matrix if requested
132
222
  vcov = None
133
223
  if return_vcov:
134
- vcov = compute_robust_vcov(X, residuals, cluster_ids)
224
+ vcov = _compute_robust_vcov_numpy(X, residuals, cluster_ids)
135
225
 
136
226
  if return_fitted:
137
227
  return coefficients, residuals, fitted, vcov
@@ -176,6 +266,63 @@ def compute_robust_vcov(
176
266
  The cluster-robust computation is vectorized using pandas groupby,
177
267
  which is much faster than a Python loop over clusters.
178
268
  """
269
+ # Use Rust backend if available
270
+ if HAS_RUST_BACKEND:
271
+ X = np.ascontiguousarray(X, dtype=np.float64)
272
+ residuals = np.ascontiguousarray(residuals, dtype=np.float64)
273
+
274
+ cluster_ids_int = None
275
+ if cluster_ids is not None:
276
+ cluster_ids_int = pd.factorize(cluster_ids)[0].astype(np.int64)
277
+
278
+ try:
279
+ return _rust_compute_robust_vcov(X, residuals, cluster_ids_int)
280
+ except ValueError as e:
281
+ # Translate Rust LAPACK errors to consistent Python error messages
282
+ error_msg = str(e)
283
+ if "Matrix inversion failed" in error_msg:
284
+ raise ValueError(
285
+ "Design matrix is rank-deficient (singular X'X matrix). "
286
+ "This indicates perfect multicollinearity. Check your fixed effects "
287
+ "and covariates for linear dependencies."
288
+ ) from e
289
+ raise
290
+
291
+ # Fallback to NumPy implementation
292
+ return _compute_robust_vcov_numpy(X, residuals, cluster_ids)
293
+
294
+
295
+ def _compute_robust_vcov_numpy(
296
+ X: np.ndarray,
297
+ residuals: np.ndarray,
298
+ cluster_ids: Optional[np.ndarray] = None,
299
+ ) -> np.ndarray:
300
+ """
301
+ NumPy fallback implementation of compute_robust_vcov.
302
+
303
+ Computes HC1 (heteroskedasticity-robust) or cluster-robust variance-covariance
304
+ matrix using the sandwich estimator.
305
+
306
+ Parameters
307
+ ----------
308
+ X : np.ndarray
309
+ Design matrix of shape (n, k).
310
+ residuals : np.ndarray
311
+ OLS residuals of shape (n,).
312
+ cluster_ids : np.ndarray, optional
313
+ Cluster identifiers. If None, uses HC1. If provided, uses
314
+ cluster-robust with G/(G-1) small-sample adjustment.
315
+
316
+ Returns
317
+ -------
318
+ vcov : np.ndarray
319
+ Variance-covariance matrix of shape (k, k).
320
+
321
+ Notes
322
+ -----
323
+ Uses vectorized groupby aggregation for cluster-robust SEs to avoid
324
+ the O(n * G) loop that would be required with explicit iteration.
325
+ """
179
326
  n, k = X.shape
180
327
  XtX = X.T @ X
181
328
 
@@ -20,6 +20,9 @@ from diff_diff.utils import (
20
20
  compute_p_value,
21
21
  )
22
22
 
23
+ # Import Rust backend if available (from _backend to avoid circular imports)
24
+ from diff_diff._backend import HAS_RUST_BACKEND, _rust_bootstrap_weights
25
+
23
26
  # Type alias for pre-computed structures
24
27
  PrecomputedData = Dict[str, Any]
25
28
 
@@ -99,6 +102,45 @@ def _generate_bootstrap_weights_batch(
99
102
  rng : np.random.Generator
100
103
  Random number generator.
101
104
 
105
+ Returns
106
+ -------
107
+ np.ndarray
108
+ Array of bootstrap weights with shape (n_bootstrap, n_units).
109
+ """
110
+ # Use Rust backend if available (parallel + fast RNG)
111
+ if HAS_RUST_BACKEND:
112
+ # Get seed from the NumPy RNG for reproducibility
113
+ seed = rng.integers(0, 2**63 - 1)
114
+ return _rust_bootstrap_weights(n_bootstrap, n_units, weight_type, seed)
115
+
116
+ # Fallback to NumPy implementation
117
+ return _generate_bootstrap_weights_batch_numpy(n_bootstrap, n_units, weight_type, rng)
118
+
119
+
120
+ def _generate_bootstrap_weights_batch_numpy(
121
+ n_bootstrap: int,
122
+ n_units: int,
123
+ weight_type: str,
124
+ rng: np.random.Generator,
125
+ ) -> np.ndarray:
126
+ """
127
+ NumPy fallback implementation of _generate_bootstrap_weights_batch.
128
+
129
+ Generates multiplier bootstrap weights for wild cluster bootstrap.
130
+ All weight distributions satisfy E[w] = 0, E[w^2] = 1.
131
+
132
+ Parameters
133
+ ----------
134
+ n_bootstrap : int
135
+ Number of bootstrap iterations.
136
+ n_units : int
137
+ Number of units (clusters) to generate weights for.
138
+ weight_type : str
139
+ Type of weights: "rademacher" (+-1), "mammen" (2-point),
140
+ or "webb" (6-point).
141
+ rng : np.random.Generator
142
+ Random number generator for reproducibility.
143
+
102
144
  Returns
103
145
  -------
104
146
  np.ndarray
@@ -1074,7 +1116,9 @@ class CallawaySantAnna:
1074
1116
  )
1075
1117
 
1076
1118
  # Compute overall ATT (simple aggregation)
1077
- overall_att, overall_se = self._aggregate_simple(group_time_effects, df, unit)
1119
+ overall_att, overall_se = self._aggregate_simple(
1120
+ group_time_effects, influence_func_info, df, unit
1121
+ )
1078
1122
  overall_t = overall_att / overall_se if overall_se > 0 else 0.0
1079
1123
  overall_p = compute_p_value(overall_t)
1080
1124
  overall_ci = compute_confidence_interval(overall_att, overall_se, self.alpha)
@@ -1085,12 +1129,13 @@ class CallawaySantAnna:
1085
1129
 
1086
1130
  if aggregate in ["event_study", "all"]:
1087
1131
  event_study_effects = self._aggregate_event_study(
1088
- group_time_effects, treatment_groups, time_periods, balance_e
1132
+ group_time_effects, influence_func_info,
1133
+ treatment_groups, time_periods, balance_e
1089
1134
  )
1090
1135
 
1091
1136
  if aggregate in ["group", "all"]:
1092
1137
  group_effects = self._aggregate_by_group(
1093
- group_time_effects, treatment_groups
1138
+ group_time_effects, influence_func_info, treatment_groups
1094
1139
  )
1095
1140
 
1096
1141
  # Run bootstrap inference if requested
@@ -1423,6 +1468,7 @@ class CallawaySantAnna:
1423
1468
  def _aggregate_simple(
1424
1469
  self,
1425
1470
  group_time_effects: Dict,
1471
+ influence_func_info: Dict,
1426
1472
  df: pd.DataFrame,
1427
1473
  unit: str,
1428
1474
  ) -> Tuple[float, float]:
@@ -1430,19 +1476,22 @@ class CallawaySantAnna:
1430
1476
  Compute simple weighted average of ATT(g,t).
1431
1477
 
1432
1478
  Weights by group size (number of treated units).
1479
+
1480
+ Standard errors are computed using influence function aggregation,
1481
+ which properly accounts for covariances across (g,t) pairs due to
1482
+ shared control units. This matches R's `did` package approach.
1433
1483
  """
1434
1484
  effects = []
1435
- weights = []
1436
- variances = []
1485
+ weights_list = []
1486
+ gt_pairs = []
1437
1487
 
1438
1488
  for (g, t), data in group_time_effects.items():
1439
1489
  effects.append(data['effect'])
1440
- weights.append(data['n_treated'])
1441
- variances.append(data['se'] ** 2)
1490
+ weights_list.append(data['n_treated'])
1491
+ gt_pairs.append((g, t))
1442
1492
 
1443
1493
  effects = np.array(effects)
1444
- weights = np.array(weights, dtype=float)
1445
- variances = np.array(variances)
1494
+ weights = np.array(weights_list, dtype=float)
1446
1495
 
1447
1496
  # Normalize weights
1448
1497
  weights = weights / np.sum(weights)
@@ -1450,15 +1499,77 @@ class CallawaySantAnna:
1450
1499
  # Weighted average
1451
1500
  overall_att = np.sum(weights * effects)
1452
1501
 
1453
- # Standard error (assuming independence across g,t)
1454
- overall_var = np.sum((weights ** 2) * variances)
1455
- overall_se = np.sqrt(overall_var)
1502
+ # Compute SE using influence function aggregation
1503
+ overall_se = self._compute_aggregated_se(
1504
+ gt_pairs, weights, influence_func_info
1505
+ )
1456
1506
 
1457
1507
  return overall_att, overall_se
1458
1508
 
1509
+ def _compute_aggregated_se(
1510
+ self,
1511
+ gt_pairs: List[Tuple[Any, Any]],
1512
+ weights: np.ndarray,
1513
+ influence_func_info: Dict,
1514
+ ) -> float:
1515
+ """
1516
+ Compute standard error using influence function aggregation.
1517
+
1518
+ This properly accounts for covariances across (g,t) pairs by
1519
+ aggregating unit-level influence functions:
1520
+
1521
+ ψ_i(overall) = Σ_{(g,t)} w_(g,t) × ψ_i(g,t)
1522
+ Var(overall) = (1/n) Σ_i [ψ_i]²
1523
+
1524
+ This matches R's `did` package analytical SE formula.
1525
+ """
1526
+ if not influence_func_info:
1527
+ # Fallback if no influence functions available
1528
+ return 0.0
1529
+
1530
+ # Build unit index mapping from all (g,t) pairs
1531
+ all_units = set()
1532
+ for (g, t) in gt_pairs:
1533
+ if (g, t) in influence_func_info:
1534
+ info = influence_func_info[(g, t)]
1535
+ all_units.update(info['treated_units'])
1536
+ all_units.update(info['control_units'])
1537
+
1538
+ if not all_units:
1539
+ return 0.0
1540
+
1541
+ all_units = sorted(all_units)
1542
+ n_units = len(all_units)
1543
+ unit_to_idx = {u: i for i, u in enumerate(all_units)}
1544
+
1545
+ # Aggregate influence functions across (g,t) pairs
1546
+ psi_overall = np.zeros(n_units)
1547
+
1548
+ for j, (g, t) in enumerate(gt_pairs):
1549
+ if (g, t) not in influence_func_info:
1550
+ continue
1551
+
1552
+ info = influence_func_info[(g, t)]
1553
+ w = weights[j]
1554
+
1555
+ # Treated unit contributions
1556
+ for i, unit_id in enumerate(info['treated_units']):
1557
+ idx = unit_to_idx[unit_id]
1558
+ psi_overall[idx] += w * info['treated_inf'][i]
1559
+
1560
+ # Control unit contributions
1561
+ for i, unit_id in enumerate(info['control_units']):
1562
+ idx = unit_to_idx[unit_id]
1563
+ psi_overall[idx] += w * info['control_inf'][i]
1564
+
1565
+ # Compute variance: Var(θ̄) = (1/n) Σᵢ ψᵢ²
1566
+ variance = np.sum(psi_overall ** 2)
1567
+ return np.sqrt(variance)
1568
+
1459
1569
  def _aggregate_event_study(
1460
1570
  self,
1461
1571
  group_time_effects: Dict,
1572
+ influence_func_info: Dict,
1462
1573
  groups: List[Any],
1463
1574
  time_periods: List[Any],
1464
1575
  balance_e: Optional[int] = None,
@@ -1467,17 +1578,20 @@ class CallawaySantAnna:
1467
1578
  Aggregate effects by relative time (event study).
1468
1579
 
1469
1580
  Computes average effect at each event time e = t - g.
1581
+
1582
+ Standard errors use influence function aggregation to account for
1583
+ covariances across (g,t) pairs.
1470
1584
  """
1471
- # Organize effects by relative time
1472
- effects_by_e: Dict[int, List[Tuple[float, float, int]]] = {}
1585
+ # Organize effects by relative time, keeping track of (g,t) pairs
1586
+ effects_by_e: Dict[int, List[Tuple[Tuple[Any, Any], float, int]]] = {}
1473
1587
 
1474
1588
  for (g, t), data in group_time_effects.items():
1475
1589
  e = t - g # Relative time
1476
1590
  if e not in effects_by_e:
1477
1591
  effects_by_e[e] = []
1478
1592
  effects_by_e[e].append((
1593
+ (g, t), # Keep track of the (g,t) pair
1479
1594
  data['effect'],
1480
- data['se'],
1481
1595
  data['n_treated']
1482
1596
  ))
1483
1597
 
@@ -1490,15 +1604,15 @@ class CallawaySantAnna:
1490
1604
  groups_at_e.add(g)
1491
1605
 
1492
1606
  # Filter effects to only include balanced groups
1493
- balanced_effects: Dict[int, List[Tuple[float, float, int]]] = {}
1607
+ balanced_effects: Dict[int, List[Tuple[Tuple[Any, Any], float, int]]] = {}
1494
1608
  for (g, t), data in group_time_effects.items():
1495
1609
  if g in groups_at_e:
1496
1610
  e = t - g
1497
1611
  if e not in balanced_effects:
1498
1612
  balanced_effects[e] = []
1499
1613
  balanced_effects[e].append((
1614
+ (g, t),
1500
1615
  data['effect'],
1501
- data['se'],
1502
1616
  data['n_treated']
1503
1617
  ))
1504
1618
  effects_by_e = balanced_effects
@@ -1507,16 +1621,19 @@ class CallawaySantAnna:
1507
1621
  event_study_effects = {}
1508
1622
 
1509
1623
  for e, effect_list in sorted(effects_by_e.items()):
1510
- effs = np.array([x[0] for x in effect_list])
1511
- ses = np.array([x[1] for x in effect_list])
1624
+ gt_pairs = [x[0] for x in effect_list]
1625
+ effs = np.array([x[1] for x in effect_list])
1512
1626
  ns = np.array([x[2] for x in effect_list], dtype=float)
1513
1627
 
1514
1628
  # Weight by group size
1515
1629
  weights = ns / np.sum(ns)
1516
1630
 
1517
1631
  agg_effect = np.sum(weights * effs)
1518
- agg_var = np.sum((weights ** 2) * (ses ** 2))
1519
- agg_se = np.sqrt(agg_var)
1632
+
1633
+ # Compute SE using influence function aggregation
1634
+ agg_se = self._compute_aggregated_se(
1635
+ gt_pairs, weights, influence_func_info
1636
+ )
1520
1637
 
1521
1638
  t_stat = agg_effect / agg_se if agg_se > 0 else 0.0
1522
1639
  p_val = compute_p_value(t_stat)
@@ -1536,19 +1653,24 @@ class CallawaySantAnna:
1536
1653
  def _aggregate_by_group(
1537
1654
  self,
1538
1655
  group_time_effects: Dict,
1656
+ influence_func_info: Dict,
1539
1657
  groups: List[Any],
1540
1658
  ) -> Dict[Any, Dict[str, Any]]:
1541
1659
  """
1542
1660
  Aggregate effects by treatment cohort.
1543
1661
 
1544
1662
  Computes average effect for each cohort across all post-treatment periods.
1663
+
1664
+ Standard errors use influence function aggregation to account for
1665
+ covariances across time periods within a cohort.
1545
1666
  """
1546
1667
  group_effects = {}
1547
1668
 
1548
1669
  for g in groups:
1549
1670
  # Get all effects for this group (post-treatment only: t >= g)
1671
+ # Keep track of (g, t) pairs for influence function aggregation
1550
1672
  g_effects = [
1551
- (data['effect'], data['se'], data['n_treated'])
1673
+ ((g, t), data['effect'])
1552
1674
  for (gg, t), data in group_time_effects.items()
1553
1675
  if gg == g and t >= g
1554
1676
  ]
@@ -1556,15 +1678,18 @@ class CallawaySantAnna:
1556
1678
  if not g_effects:
1557
1679
  continue
1558
1680
 
1559
- effs = np.array([x[0] for x in g_effects])
1560
- ses = np.array([x[1] for x in g_effects])
1681
+ gt_pairs = [x[0] for x in g_effects]
1682
+ effs = np.array([x[1] for x in g_effects])
1561
1683
 
1562
1684
  # Equal weight across time periods for a group
1563
1685
  weights = np.ones(len(effs)) / len(effs)
1564
1686
 
1565
1687
  agg_effect = np.sum(weights * effs)
1566
- agg_var = np.sum((weights ** 2) * (ses ** 2))
1567
- agg_se = np.sqrt(agg_var)
1688
+
1689
+ # Compute SE using influence function aggregation
1690
+ agg_se = self._compute_aggregated_se(
1691
+ gt_pairs, weights, influence_func_info
1692
+ )
1568
1693
 
1569
1694
  t_stat = agg_effect / agg_se if agg_se > 0 else 0.0
1570
1695
  p_val = compute_p_value(t_stat)