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.
- {diff_diff-1.4.0 → diff_diff-2.0.0}/PKG-INFO +19 -18
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/__init__.py +13 -1
- diff_diff-2.0.0/diff_diff/_backend.py +64 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/linalg.py +151 -4
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/staggered.py +151 -26
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/utils.py +39 -9
- {diff_diff-1.4.0 → diff_diff-2.0.0}/pyproject.toml +14 -5
- diff_diff-2.0.0/rust/Cargo.lock +2396 -0
- diff_diff-2.0.0/rust/Cargo.toml +39 -0
- diff_diff-2.0.0/rust/src/bootstrap.rs +223 -0
- diff_diff-2.0.0/rust/src/lib.rs +33 -0
- diff_diff-2.0.0/rust/src/linalg.rs +229 -0
- diff_diff-2.0.0/rust/src/weights.rs +220 -0
- diff_diff-1.4.0/diff_diff.egg-info/PKG-INFO +0 -2255
- diff_diff-1.4.0/diff_diff.egg-info/SOURCES.txt +0 -38
- diff_diff-1.4.0/diff_diff.egg-info/dependency_links.txt +0 -1
- diff_diff-1.4.0/diff_diff.egg-info/requires.txt +0 -14
- diff_diff-1.4.0/diff_diff.egg-info/top_level.txt +0 -1
- diff_diff-1.4.0/setup.cfg +0 -4
- diff_diff-1.4.0/tests/test_bacon.py +0 -679
- diff_diff-1.4.0/tests/test_diagnostics.py +0 -674
- diff_diff-1.4.0/tests/test_estimators.py +0 -2719
- diff_diff-1.4.0/tests/test_honest_did.py +0 -699
- diff_diff-1.4.0/tests/test_linalg.py +0 -422
- diff_diff-1.4.0/tests/test_power.py +0 -691
- diff_diff-1.4.0/tests/test_prep.py +0 -794
- diff_diff-1.4.0/tests/test_pretrends.py +0 -813
- diff_diff-1.4.0/tests/test_staggered.py +0 -1358
- diff_diff-1.4.0/tests/test_sun_abraham.py +0 -732
- diff_diff-1.4.0/tests/test_triple_diff.py +0 -869
- diff_diff-1.4.0/tests/test_utils.py +0 -1270
- diff_diff-1.4.0/tests/test_visualization.py +0 -284
- diff_diff-1.4.0/tests/test_wild_bootstrap.py +0 -804
- {diff_diff-1.4.0 → diff_diff-2.0.0}/README.md +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/bacon.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/diagnostics.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/estimators.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/honest_did.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/power.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/prep.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/pretrends.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/results.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/sun_abraham.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/synthetic_did.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/triple_diff.py +0 -0
- {diff_diff-1.4.0 → diff_diff-2.0.0}/diff_diff/twfe.py +0 -0
- {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:
|
|
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
|
-
|
|
34
|
-
|
|
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__ = "
|
|
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
|
|
5
|
-
|
|
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
|
-
|
|
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 =
|
|
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(
|
|
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,
|
|
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
|
-
|
|
1436
|
-
|
|
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
|
-
|
|
1441
|
-
|
|
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(
|
|
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
|
-
#
|
|
1454
|
-
|
|
1455
|
-
|
|
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[
|
|
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[
|
|
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
|
-
|
|
1511
|
-
|
|
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
|
-
|
|
1519
|
-
|
|
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
|
-
(
|
|
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
|
-
|
|
1560
|
-
|
|
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
|
-
|
|
1567
|
-
|
|
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)
|