diff-diff 2.7.0__tar.gz → 2.7.1__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-2.7.0 → diff_diff-2.7.1}/PKG-INFO +53 -2
- {diff_diff-2.7.0 → diff_diff-2.7.1}/README.md +52 -1
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/__init__.py +1 -1
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/linalg.py +217 -27
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered.py +276 -253
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered_results.py +87 -66
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/sun_abraham.py +9 -15
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/triple_diff.py +241 -258
- {diff_diff-2.7.0 → diff_diff-2.7.1}/pyproject.toml +3 -3
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/Cargo.lock +1 -1
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/Cargo.toml +1 -1
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/_backend.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/bacon.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/bootstrap_utils.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/continuous_did.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/continuous_did_bspline.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/continuous_did_results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/datasets.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/diagnostics.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did_bootstrap.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did_results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did_weights.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/estimators.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/honest_did.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/imputation.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/imputation_bootstrap.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/imputation_results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/power.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/prep.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/prep_dgp.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/pretrends.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/stacked_did.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/stacked_did_results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered_aggregation.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered_bootstrap.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/synthetic_did.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/trop.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/trop_results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/twfe.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/two_stage.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/two_stage_bootstrap.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/two_stage_results.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/utils.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/visualization.py +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/build.rs +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/bootstrap.rs +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/lib.rs +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/linalg.rs +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/trop.rs +0 -0
- {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/weights.rs +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: diff-diff
|
|
3
|
-
Version: 2.7.
|
|
3
|
+
Version: 2.7.1
|
|
4
4
|
Classifier: Development Status :: 5 - Production/Stable
|
|
5
5
|
Classifier: Intended Audience :: Science/Research
|
|
6
6
|
Classifier: Operating System :: OS Independent
|
|
@@ -108,7 +108,7 @@ Signif. codes: '***' 0.001, '**' 0.01, '*' 0.05, '.' 0.1
|
|
|
108
108
|
- **Wild cluster bootstrap**: Valid inference with few clusters (<50) using Rademacher, Webb, or Mammen weights
|
|
109
109
|
- **Panel data support**: Two-way fixed effects estimator for panel designs
|
|
110
110
|
- **Multi-period analysis**: Event-study style DiD with period-specific treatment effects
|
|
111
|
-
- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022),
|
|
111
|
+
- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022), Stacked DiD (Wing, Freedman & Hollingsworth 2024), and Efficient DiD (Chen, Sant'Anna & Xie 2025) estimators for heterogeneous treatment timing
|
|
112
112
|
- **Triple Difference (DDD)**: Ortiz-Villavicencio & Sant'Anna (2025) estimators with proper covariate handling
|
|
113
113
|
- **Synthetic DiD**: Combined DiD with synthetic control for improved robustness
|
|
114
114
|
- **Triply Robust Panel (TROP)**: Factor-adjusted DiD with synthetic weights (Athey et al. 2025)
|
|
@@ -163,6 +163,7 @@ We provide Jupyter notebook tutorials in `docs/tutorials/`:
|
|
|
163
163
|
| `11_imputation_did.ipynb` | Imputation DiD (Borusyak et al. 2024), pre-trend test, efficiency comparison |
|
|
164
164
|
| `12_two_stage_did.ipynb` | Two-Stage DiD (Gardner 2022), GMM sandwich variance, per-observation effects |
|
|
165
165
|
| `13_stacked_did.ipynb` | Stacked DiD (Wing et al. 2024), Q-weights, sub-experiment inspection, trimming, clean control definitions |
|
|
166
|
+
| `15_efficient_did.ipynb` | Efficient DiD (Chen et al. 2025), optimal weighting, PT-All vs PT-Post, efficiency gains, bootstrap inference |
|
|
166
167
|
|
|
167
168
|
## Data Preparation
|
|
168
169
|
|
|
@@ -1109,6 +1110,56 @@ results = stacked_did(
|
|
|
1109
1110
|
)
|
|
1110
1111
|
```
|
|
1111
1112
|
|
|
1113
|
+
### Efficient DiD (Chen, Sant'Anna & Xie 2025)
|
|
1114
|
+
|
|
1115
|
+
Efficient DiD achieves the semiparametric efficiency bound for ATT estimation in staggered adoption designs. It optimally weights across all valid comparison groups and baselines via the inverse covariance matrix Omega*, producing tighter confidence intervals than standard estimators like Callaway-Sant'Anna when the stronger PT-All assumption holds.
|
|
1116
|
+
|
|
1117
|
+
```python
|
|
1118
|
+
from diff_diff import EfficientDiD, generate_staggered_data
|
|
1119
|
+
|
|
1120
|
+
# Generate sample data
|
|
1121
|
+
data = generate_staggered_data(n_units=300, n_periods=10,
|
|
1122
|
+
cohort_periods=[4, 6, 8], seed=42)
|
|
1123
|
+
|
|
1124
|
+
# Fit with PT-All (overidentified, tighter SEs)
|
|
1125
|
+
edid = EfficientDiD(pt_assumption="all")
|
|
1126
|
+
results = edid.fit(data, outcome='outcome', unit='unit',
|
|
1127
|
+
time='period', first_treat='first_treat',
|
|
1128
|
+
aggregate='all')
|
|
1129
|
+
results.print_summary()
|
|
1130
|
+
|
|
1131
|
+
# PT-Post mode (matches CS for post-treatment effects)
|
|
1132
|
+
edid_post = EfficientDiD(pt_assumption="post")
|
|
1133
|
+
results_post = edid_post.fit(data, outcome='outcome', unit='unit',
|
|
1134
|
+
time='period', first_treat='first_treat')
|
|
1135
|
+
```
|
|
1136
|
+
|
|
1137
|
+
**Parameters:**
|
|
1138
|
+
|
|
1139
|
+
```python
|
|
1140
|
+
EfficientDiD(
|
|
1141
|
+
pt_assumption='all', # 'all' (overidentified) or 'post' (matches CS post-treatment ATT)
|
|
1142
|
+
alpha=0.05, # Significance level
|
|
1143
|
+
n_bootstrap=0, # Bootstrap iterations (0 = analytical only)
|
|
1144
|
+
bootstrap_weights='rademacher', # 'rademacher', 'mammen', or 'webb'
|
|
1145
|
+
seed=None, # Random seed
|
|
1146
|
+
anticipation=0, # Anticipation periods
|
|
1147
|
+
)
|
|
1148
|
+
```
|
|
1149
|
+
|
|
1150
|
+
> **Note:** Phase 1 supports the no-covariates path only. Use CallawaySantAnna with
|
|
1151
|
+
> `estimation_method='dr'` if you need covariate adjustment.
|
|
1152
|
+
|
|
1153
|
+
**When to use Efficient DiD vs Callaway-Sant'Anna:**
|
|
1154
|
+
|
|
1155
|
+
| Aspect | Efficient DiD | Callaway-Sant'Anna |
|
|
1156
|
+
|--------|--------------|-------------------|
|
|
1157
|
+
| Approach | Optimal EIF-based weighting | Separate 2x2 DiD aggregation |
|
|
1158
|
+
| PT assumption | PT-All (stronger) or PT-Post | Conditional PT |
|
|
1159
|
+
| Efficiency | Achieves semiparametric bound | Not efficient |
|
|
1160
|
+
| Covariates | Not yet (Phase 2) | Supported (OR, IPW, DR) |
|
|
1161
|
+
| When to choose | Maximum efficiency, PT-All credible | Covariates needed, weaker PT |
|
|
1162
|
+
|
|
1112
1163
|
### Triple Difference (DDD)
|
|
1113
1164
|
|
|
1114
1165
|
Triple Difference (DDD) is used when treatment requires satisfying two criteria: belonging to a treated **group** AND being in an eligible **partition**. The `TripleDifference` class implements the methodology from Ortiz-Villavicencio & Sant'Anna (2025), which correctly handles covariate adjustment (unlike naive implementations).
|
|
@@ -70,7 +70,7 @@ Signif. codes: '***' 0.001, '**' 0.01, '*' 0.05, '.' 0.1
|
|
|
70
70
|
- **Wild cluster bootstrap**: Valid inference with few clusters (<50) using Rademacher, Webb, or Mammen weights
|
|
71
71
|
- **Panel data support**: Two-way fixed effects estimator for panel designs
|
|
72
72
|
- **Multi-period analysis**: Event-study style DiD with period-specific treatment effects
|
|
73
|
-
- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022),
|
|
73
|
+
- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022), Stacked DiD (Wing, Freedman & Hollingsworth 2024), and Efficient DiD (Chen, Sant'Anna & Xie 2025) estimators for heterogeneous treatment timing
|
|
74
74
|
- **Triple Difference (DDD)**: Ortiz-Villavicencio & Sant'Anna (2025) estimators with proper covariate handling
|
|
75
75
|
- **Synthetic DiD**: Combined DiD with synthetic control for improved robustness
|
|
76
76
|
- **Triply Robust Panel (TROP)**: Factor-adjusted DiD with synthetic weights (Athey et al. 2025)
|
|
@@ -125,6 +125,7 @@ We provide Jupyter notebook tutorials in `docs/tutorials/`:
|
|
|
125
125
|
| `11_imputation_did.ipynb` | Imputation DiD (Borusyak et al. 2024), pre-trend test, efficiency comparison |
|
|
126
126
|
| `12_two_stage_did.ipynb` | Two-Stage DiD (Gardner 2022), GMM sandwich variance, per-observation effects |
|
|
127
127
|
| `13_stacked_did.ipynb` | Stacked DiD (Wing et al. 2024), Q-weights, sub-experiment inspection, trimming, clean control definitions |
|
|
128
|
+
| `15_efficient_did.ipynb` | Efficient DiD (Chen et al. 2025), optimal weighting, PT-All vs PT-Post, efficiency gains, bootstrap inference |
|
|
128
129
|
|
|
129
130
|
## Data Preparation
|
|
130
131
|
|
|
@@ -1071,6 +1072,56 @@ results = stacked_did(
|
|
|
1071
1072
|
)
|
|
1072
1073
|
```
|
|
1073
1074
|
|
|
1075
|
+
### Efficient DiD (Chen, Sant'Anna & Xie 2025)
|
|
1076
|
+
|
|
1077
|
+
Efficient DiD achieves the semiparametric efficiency bound for ATT estimation in staggered adoption designs. It optimally weights across all valid comparison groups and baselines via the inverse covariance matrix Omega*, producing tighter confidence intervals than standard estimators like Callaway-Sant'Anna when the stronger PT-All assumption holds.
|
|
1078
|
+
|
|
1079
|
+
```python
|
|
1080
|
+
from diff_diff import EfficientDiD, generate_staggered_data
|
|
1081
|
+
|
|
1082
|
+
# Generate sample data
|
|
1083
|
+
data = generate_staggered_data(n_units=300, n_periods=10,
|
|
1084
|
+
cohort_periods=[4, 6, 8], seed=42)
|
|
1085
|
+
|
|
1086
|
+
# Fit with PT-All (overidentified, tighter SEs)
|
|
1087
|
+
edid = EfficientDiD(pt_assumption="all")
|
|
1088
|
+
results = edid.fit(data, outcome='outcome', unit='unit',
|
|
1089
|
+
time='period', first_treat='first_treat',
|
|
1090
|
+
aggregate='all')
|
|
1091
|
+
results.print_summary()
|
|
1092
|
+
|
|
1093
|
+
# PT-Post mode (matches CS for post-treatment effects)
|
|
1094
|
+
edid_post = EfficientDiD(pt_assumption="post")
|
|
1095
|
+
results_post = edid_post.fit(data, outcome='outcome', unit='unit',
|
|
1096
|
+
time='period', first_treat='first_treat')
|
|
1097
|
+
```
|
|
1098
|
+
|
|
1099
|
+
**Parameters:**
|
|
1100
|
+
|
|
1101
|
+
```python
|
|
1102
|
+
EfficientDiD(
|
|
1103
|
+
pt_assumption='all', # 'all' (overidentified) or 'post' (matches CS post-treatment ATT)
|
|
1104
|
+
alpha=0.05, # Significance level
|
|
1105
|
+
n_bootstrap=0, # Bootstrap iterations (0 = analytical only)
|
|
1106
|
+
bootstrap_weights='rademacher', # 'rademacher', 'mammen', or 'webb'
|
|
1107
|
+
seed=None, # Random seed
|
|
1108
|
+
anticipation=0, # Anticipation periods
|
|
1109
|
+
)
|
|
1110
|
+
```
|
|
1111
|
+
|
|
1112
|
+
> **Note:** Phase 1 supports the no-covariates path only. Use CallawaySantAnna with
|
|
1113
|
+
> `estimation_method='dr'` if you need covariate adjustment.
|
|
1114
|
+
|
|
1115
|
+
**When to use Efficient DiD vs Callaway-Sant'Anna:**
|
|
1116
|
+
|
|
1117
|
+
| Aspect | Efficient DiD | Callaway-Sant'Anna |
|
|
1118
|
+
|--------|--------------|-------------------|
|
|
1119
|
+
| Approach | Optimal EIF-based weighting | Separate 2x2 DiD aggregation |
|
|
1120
|
+
| PT assumption | PT-All (stronger) or PT-Post | Conditional PT |
|
|
1121
|
+
| Efficiency | Achieves semiparametric bound | Not efficient |
|
|
1122
|
+
| Covariates | Not yet (Phase 2) | Supported (OR, IPW, DR) |
|
|
1123
|
+
| When to choose | Maximum efficiency, PT-All credible | Covariates needed, weaker PT |
|
|
1124
|
+
|
|
1074
1125
|
### Triple Difference (DDD)
|
|
1075
1126
|
|
|
1076
1127
|
Triple Difference (DDD) is used when treatment requires satisfying two criteria: belonging to a treated **group** AND being in an eligible **partition**. The `TripleDifference` class implements the methodology from Ortiz-Villavicencio & Sant'Anna (2025), which correctly handles covariate adjustment (unlike naive implementations).
|
|
@@ -116,7 +116,7 @@ def _detect_rank_deficiency(
|
|
|
116
116
|
|
|
117
117
|
# Compute pivoted QR decomposition: X @ P = Q @ R
|
|
118
118
|
# P is a permutation matrix, represented as pivot indices
|
|
119
|
-
Q, R, pivot = qr(X, mode=
|
|
119
|
+
Q, R, pivot = qr(X, mode="economic", pivoting=True)
|
|
120
120
|
|
|
121
121
|
# Determine rank tolerance
|
|
122
122
|
# R's qr() uses tol = 1e-07 by default, which is sqrt(eps) ≈ 1.49e-08
|
|
@@ -169,8 +169,7 @@ def _format_dropped_columns(
|
|
|
169
169
|
return ""
|
|
170
170
|
|
|
171
171
|
if column_names is not None:
|
|
172
|
-
names = [column_names[i] if i < len(column_names) else f"column {i}"
|
|
173
|
-
for i in dropped_cols]
|
|
172
|
+
names = [column_names[i] if i < len(column_names) else f"column {i}" for i in dropped_cols]
|
|
174
173
|
if len(names) == 1:
|
|
175
174
|
return f"'{names[0]}'"
|
|
176
175
|
elif len(names) <= 5:
|
|
@@ -251,10 +250,12 @@ def _solve_ols_rust(
|
|
|
251
250
|
cluster_ids: Optional[np.ndarray] = None,
|
|
252
251
|
return_vcov: bool = True,
|
|
253
252
|
return_fitted: bool = False,
|
|
254
|
-
) -> Optional[
|
|
255
|
-
|
|
256
|
-
|
|
257
|
-
]]
|
|
253
|
+
) -> Optional[
|
|
254
|
+
Union[
|
|
255
|
+
Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
|
|
256
|
+
Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
|
|
257
|
+
]
|
|
258
|
+
]:
|
|
258
259
|
"""
|
|
259
260
|
Rust backend implementation of solve_ols for full-rank matrices.
|
|
260
261
|
|
|
@@ -447,8 +448,7 @@ def solve_ols(
|
|
|
447
448
|
raise ValueError(f"y must be 1-dimensional, got shape {y.shape}")
|
|
448
449
|
if X.shape[0] != y.shape[0]:
|
|
449
450
|
raise ValueError(
|
|
450
|
-
f"X and y must have same number of observations: "
|
|
451
|
-
f"{X.shape[0]} vs {y.shape[0]}"
|
|
451
|
+
f"X and y must have same number of observations: " f"{X.shape[0]} vs {y.shape[0]}"
|
|
452
452
|
)
|
|
453
453
|
|
|
454
454
|
n, k = X.shape
|
|
@@ -484,7 +484,8 @@ def solve_ols(
|
|
|
484
484
|
if skip_rank_check:
|
|
485
485
|
if HAS_RUST_BACKEND and _rust_solve_ols is not None:
|
|
486
486
|
result = _solve_ols_rust(
|
|
487
|
-
X,
|
|
487
|
+
X,
|
|
488
|
+
y,
|
|
488
489
|
cluster_ids=cluster_ids,
|
|
489
490
|
return_vcov=return_vcov,
|
|
490
491
|
return_fitted=return_fitted,
|
|
@@ -494,7 +495,8 @@ def solve_ols(
|
|
|
494
495
|
# Fall through to NumPy on numerical instability
|
|
495
496
|
# Fall through to Python without rank check (user guarantees full rank)
|
|
496
497
|
return _solve_ols_numpy(
|
|
497
|
-
X,
|
|
498
|
+
X,
|
|
499
|
+
y,
|
|
498
500
|
cluster_ids=cluster_ids,
|
|
499
501
|
return_vcov=return_vcov,
|
|
500
502
|
return_fitted=return_fitted,
|
|
@@ -521,7 +523,8 @@ def solve_ols(
|
|
|
521
523
|
# - No Rust → Python backend (works for all cases)
|
|
522
524
|
if HAS_RUST_BACKEND and _rust_solve_ols is not None and not is_rank_deficient:
|
|
523
525
|
result = _solve_ols_rust(
|
|
524
|
-
X,
|
|
526
|
+
X,
|
|
527
|
+
y,
|
|
525
528
|
cluster_ids=cluster_ids,
|
|
526
529
|
return_vcov=return_vcov,
|
|
527
530
|
return_fitted=return_fitted,
|
|
@@ -531,7 +534,8 @@ def solve_ols(
|
|
|
531
534
|
# signaled us to fall back to Python backend
|
|
532
535
|
if result is None:
|
|
533
536
|
return _solve_ols_numpy(
|
|
534
|
-
X,
|
|
537
|
+
X,
|
|
538
|
+
y,
|
|
535
539
|
cluster_ids=cluster_ids,
|
|
536
540
|
return_vcov=return_vcov,
|
|
537
541
|
return_fitted=return_fitted,
|
|
@@ -555,7 +559,8 @@ def solve_ols(
|
|
|
555
559
|
# and SVD disagreed about rank. Python's QR will re-detect and
|
|
556
560
|
# apply R-style NaN handling for dropped columns.
|
|
557
561
|
return _solve_ols_numpy(
|
|
558
|
-
X,
|
|
562
|
+
X,
|
|
563
|
+
y,
|
|
559
564
|
cluster_ids=cluster_ids,
|
|
560
565
|
return_vcov=return_vcov,
|
|
561
566
|
return_fitted=return_fitted,
|
|
@@ -569,7 +574,8 @@ def solve_ols(
|
|
|
569
574
|
# Use NumPy implementation for rank-deficient cases (R-style NA handling)
|
|
570
575
|
# or when Rust backend is not available
|
|
571
576
|
return _solve_ols_numpy(
|
|
572
|
-
X,
|
|
577
|
+
X,
|
|
578
|
+
y,
|
|
573
579
|
cluster_ids=cluster_ids,
|
|
574
580
|
return_vcov=return_vcov,
|
|
575
581
|
return_fitted=return_fitted,
|
|
@@ -834,9 +840,7 @@ def _compute_robust_vcov_numpy(
|
|
|
834
840
|
n_clusters = len(unique_clusters)
|
|
835
841
|
|
|
836
842
|
if n_clusters < 2:
|
|
837
|
-
raise ValueError(
|
|
838
|
-
f"Need at least 2 clusters for cluster-robust SEs, got {n_clusters}"
|
|
839
|
-
)
|
|
843
|
+
raise ValueError(f"Need at least 2 clusters for cluster-robust SEs, got {n_clusters}")
|
|
840
844
|
|
|
841
845
|
# Small-sample adjustment
|
|
842
846
|
adjustment = (n_clusters / (n_clusters - 1)) * ((n - 1) / (n - k))
|
|
@@ -871,6 +875,193 @@ def _compute_robust_vcov_numpy(
|
|
|
871
875
|
return vcov
|
|
872
876
|
|
|
873
877
|
|
|
878
|
+
# Empirical threshold: coefficients above this magnitude suggest near-separation
|
|
879
|
+
# in the logistic model (predicted probabilities collapse to 0/1).
|
|
880
|
+
_LOGIT_SEPARATION_COEF_THRESHOLD = 10
|
|
881
|
+
_LOGIT_SEPARATION_PROB_THRESHOLD = 1e-5
|
|
882
|
+
|
|
883
|
+
|
|
884
|
+
def solve_logit(
|
|
885
|
+
X: np.ndarray,
|
|
886
|
+
y: np.ndarray,
|
|
887
|
+
max_iter: int = 25,
|
|
888
|
+
tol: float = 1e-8,
|
|
889
|
+
check_separation: bool = True,
|
|
890
|
+
rank_deficient_action: str = "warn",
|
|
891
|
+
) -> Tuple[np.ndarray, np.ndarray]:
|
|
892
|
+
"""
|
|
893
|
+
Fit logistic regression via IRLS (Fisher scoring).
|
|
894
|
+
|
|
895
|
+
Matches R's ``glm(family=binomial)`` algorithm: iteratively reweighted
|
|
896
|
+
least squares with working weights ``mu*(1-mu)`` and working response
|
|
897
|
+
``eta + (y-mu)/(mu*(1-mu))``.
|
|
898
|
+
|
|
899
|
+
Parameters
|
|
900
|
+
----------
|
|
901
|
+
X : np.ndarray
|
|
902
|
+
Feature matrix (n_samples, n_features). Intercept added automatically.
|
|
903
|
+
y : np.ndarray
|
|
904
|
+
Binary outcome (0/1).
|
|
905
|
+
max_iter : int, default 25
|
|
906
|
+
Maximum IRLS iterations (R's ``glm`` default).
|
|
907
|
+
tol : float, default 1e-8
|
|
908
|
+
Convergence tolerance on coefficient change (R's ``glm`` default).
|
|
909
|
+
check_separation : bool, default True
|
|
910
|
+
Whether to check for near-separation and emit warnings.
|
|
911
|
+
rank_deficient_action : str, default "warn"
|
|
912
|
+
How to handle rank-deficient design matrices:
|
|
913
|
+
- "warn": Emit warning and drop columns (default)
|
|
914
|
+
- "error": Raise ValueError
|
|
915
|
+
- "silent": Drop columns silently
|
|
916
|
+
|
|
917
|
+
Returns
|
|
918
|
+
-------
|
|
919
|
+
beta : np.ndarray
|
|
920
|
+
Fitted coefficients (including intercept as element 0).
|
|
921
|
+
probs : np.ndarray
|
|
922
|
+
Predicted probabilities.
|
|
923
|
+
"""
|
|
924
|
+
n, p = X.shape
|
|
925
|
+
X_with_intercept = np.column_stack([np.ones(n), X])
|
|
926
|
+
k = p + 1 # number of parameters including intercept
|
|
927
|
+
|
|
928
|
+
# Validate rank_deficient_action
|
|
929
|
+
valid_actions = {"warn", "error", "silent"}
|
|
930
|
+
if rank_deficient_action not in valid_actions:
|
|
931
|
+
raise ValueError(
|
|
932
|
+
f"rank_deficient_action must be one of {valid_actions}, "
|
|
933
|
+
f"got '{rank_deficient_action}'"
|
|
934
|
+
)
|
|
935
|
+
|
|
936
|
+
# Check rank deficiency once before iterating
|
|
937
|
+
rank_info = _detect_rank_deficiency(X_with_intercept)
|
|
938
|
+
rank, dropped_cols, _ = rank_info
|
|
939
|
+
if len(dropped_cols) > 0:
|
|
940
|
+
col_desc = _format_dropped_columns(dropped_cols)
|
|
941
|
+
if rank_deficient_action == "error":
|
|
942
|
+
raise ValueError(
|
|
943
|
+
f"Rank-deficient design matrix in logistic regression: "
|
|
944
|
+
f"dropping {col_desc}. Propensity score estimates may be unreliable."
|
|
945
|
+
)
|
|
946
|
+
elif rank_deficient_action == "warn":
|
|
947
|
+
warnings.warn(
|
|
948
|
+
f"Rank-deficient design matrix in logistic regression: "
|
|
949
|
+
f"dropping {col_desc}. Propensity score estimates may be unreliable.",
|
|
950
|
+
UserWarning,
|
|
951
|
+
stacklevel=2,
|
|
952
|
+
)
|
|
953
|
+
kept_cols = np.array([i for i in range(k) if i not in dropped_cols])
|
|
954
|
+
X_solve = X_with_intercept[:, kept_cols]
|
|
955
|
+
else:
|
|
956
|
+
kept_cols = np.arange(k)
|
|
957
|
+
X_solve = X_with_intercept
|
|
958
|
+
|
|
959
|
+
# IRLS (Fisher scoring)
|
|
960
|
+
beta_solve = np.zeros(X_solve.shape[1])
|
|
961
|
+
converged = False
|
|
962
|
+
|
|
963
|
+
for iteration in range(max_iter):
|
|
964
|
+
eta = X_solve @ beta_solve
|
|
965
|
+
# Clip to prevent overflow in exp
|
|
966
|
+
eta = np.clip(eta, -500, 500)
|
|
967
|
+
mu = 1.0 / (1.0 + np.exp(-eta))
|
|
968
|
+
# Clip mu to prevent zero working weights
|
|
969
|
+
mu = np.clip(mu, 1e-10, 1 - 1e-10)
|
|
970
|
+
|
|
971
|
+
# Working weights and working response
|
|
972
|
+
w = mu * (1.0 - mu)
|
|
973
|
+
z = eta + (y - mu) / w
|
|
974
|
+
|
|
975
|
+
# Weighted least squares: solve (X'WX) beta = X'Wz
|
|
976
|
+
sqrt_w = np.sqrt(w)
|
|
977
|
+
Xw = X_solve * sqrt_w[:, None]
|
|
978
|
+
zw = z * sqrt_w
|
|
979
|
+
beta_new, _, _, _ = np.linalg.lstsq(Xw, zw, rcond=None)
|
|
980
|
+
|
|
981
|
+
# Check convergence
|
|
982
|
+
if np.max(np.abs(beta_new - beta_solve)) < tol:
|
|
983
|
+
beta_solve = beta_new
|
|
984
|
+
converged = True
|
|
985
|
+
break
|
|
986
|
+
beta_solve = beta_new
|
|
987
|
+
|
|
988
|
+
# Final predicted probabilities
|
|
989
|
+
eta_final = X_solve @ beta_solve
|
|
990
|
+
eta_final = np.clip(eta_final, -500, 500)
|
|
991
|
+
probs = 1.0 / (1.0 + np.exp(-eta_final))
|
|
992
|
+
|
|
993
|
+
# Warnings
|
|
994
|
+
if not converged:
|
|
995
|
+
warnings.warn(
|
|
996
|
+
f"Logistic regression did not converge in {max_iter} iterations. "
|
|
997
|
+
f"Propensity score estimates may be unreliable.",
|
|
998
|
+
UserWarning,
|
|
999
|
+
stacklevel=2,
|
|
1000
|
+
)
|
|
1001
|
+
|
|
1002
|
+
if check_separation:
|
|
1003
|
+
if np.max(np.abs(beta_solve)) > _LOGIT_SEPARATION_COEF_THRESHOLD:
|
|
1004
|
+
warnings.warn(
|
|
1005
|
+
"Large coefficients detected in propensity score model "
|
|
1006
|
+
f"(max|beta| > {_LOGIT_SEPARATION_COEF_THRESHOLD}), "
|
|
1007
|
+
"suggesting potential separation.",
|
|
1008
|
+
UserWarning,
|
|
1009
|
+
stacklevel=2,
|
|
1010
|
+
)
|
|
1011
|
+
n_extreme = int(
|
|
1012
|
+
np.sum(
|
|
1013
|
+
(probs < _LOGIT_SEPARATION_PROB_THRESHOLD)
|
|
1014
|
+
| (probs > 1 - _LOGIT_SEPARATION_PROB_THRESHOLD)
|
|
1015
|
+
)
|
|
1016
|
+
)
|
|
1017
|
+
if n_extreme > 0:
|
|
1018
|
+
warnings.warn(
|
|
1019
|
+
f"Near-separation detected in propensity score model: "
|
|
1020
|
+
f"{n_extreme} of {n} observations have predicted probabilities "
|
|
1021
|
+
f"within {_LOGIT_SEPARATION_PROB_THRESHOLD} of 0 or 1. ATT estimates may be sensitive to "
|
|
1022
|
+
f"model specification.",
|
|
1023
|
+
UserWarning,
|
|
1024
|
+
stacklevel=2,
|
|
1025
|
+
)
|
|
1026
|
+
|
|
1027
|
+
# Expand beta back to full size if columns were dropped
|
|
1028
|
+
if len(dropped_cols) > 0:
|
|
1029
|
+
beta_full = np.zeros(k)
|
|
1030
|
+
beta_full[kept_cols] = beta_solve
|
|
1031
|
+
else:
|
|
1032
|
+
beta_full = beta_solve
|
|
1033
|
+
|
|
1034
|
+
return beta_full, probs
|
|
1035
|
+
|
|
1036
|
+
|
|
1037
|
+
def _check_propensity_diagnostics(
|
|
1038
|
+
pscore: np.ndarray,
|
|
1039
|
+
trim_bound: float = 0.01,
|
|
1040
|
+
) -> None:
|
|
1041
|
+
"""
|
|
1042
|
+
Warn if propensity scores are extreme.
|
|
1043
|
+
|
|
1044
|
+
Parameters
|
|
1045
|
+
----------
|
|
1046
|
+
pscore : np.ndarray
|
|
1047
|
+
Predicted probabilities.
|
|
1048
|
+
trim_bound : float, default 0.01
|
|
1049
|
+
Trimming threshold.
|
|
1050
|
+
"""
|
|
1051
|
+
n_extreme = int(np.sum((pscore < trim_bound) | (pscore > 1 - trim_bound)))
|
|
1052
|
+
if n_extreme > 0:
|
|
1053
|
+
n_total = len(pscore)
|
|
1054
|
+
pct = 100.0 * n_extreme / n_total
|
|
1055
|
+
warnings.warn(
|
|
1056
|
+
f"Propensity scores for {n_extreme} of {n_total} observations "
|
|
1057
|
+
f"({pct:.1f}%) were outside [{trim_bound}, {1 - trim_bound}] "
|
|
1058
|
+
f"and will be trimmed. This may indicate near-separation in "
|
|
1059
|
+
f"the propensity score model.",
|
|
1060
|
+
UserWarning,
|
|
1061
|
+
stacklevel=2,
|
|
1062
|
+
)
|
|
1063
|
+
|
|
1064
|
+
|
|
874
1065
|
def compute_r_squared(
|
|
875
1066
|
y: np.ndarray,
|
|
876
1067
|
residuals: np.ndarray,
|
|
@@ -1149,7 +1340,8 @@ class LinearRegression:
|
|
|
1149
1340
|
if self.robust or effective_cluster_ids is not None:
|
|
1150
1341
|
# Use solve_ols with robust/cluster SEs
|
|
1151
1342
|
coefficients, residuals, fitted, vcov = solve_ols(
|
|
1152
|
-
X,
|
|
1343
|
+
X,
|
|
1344
|
+
y,
|
|
1153
1345
|
cluster_ids=effective_cluster_ids,
|
|
1154
1346
|
return_fitted=True,
|
|
1155
1347
|
return_vcov=compute_vcov,
|
|
@@ -1158,7 +1350,8 @@ class LinearRegression:
|
|
|
1158
1350
|
else:
|
|
1159
1351
|
# Classical OLS - compute vcov separately
|
|
1160
1352
|
coefficients, residuals, fitted, _ = solve_ols(
|
|
1161
|
-
X,
|
|
1353
|
+
X,
|
|
1354
|
+
y,
|
|
1162
1355
|
return_fitted=True,
|
|
1163
1356
|
return_vcov=False,
|
|
1164
1357
|
rank_deficient_action=self.rank_deficient_action,
|
|
@@ -1294,6 +1487,7 @@ class LinearRegression:
|
|
|
1294
1487
|
# Handle zero or negative SE (indicates perfect fit or numerical issues)
|
|
1295
1488
|
if se <= 0:
|
|
1296
1489
|
import warnings
|
|
1490
|
+
|
|
1297
1491
|
warnings.warn(
|
|
1298
1492
|
f"Standard error is zero or negative (se={se}) for coefficient at index {index}. "
|
|
1299
1493
|
"This may indicate perfect multicollinearity or numerical issues.",
|
|
@@ -1319,6 +1513,7 @@ class LinearRegression:
|
|
|
1319
1513
|
# Warn if df is non-positive and fall back to normal distribution
|
|
1320
1514
|
if effective_df is not None and effective_df <= 0:
|
|
1321
1515
|
import warnings
|
|
1516
|
+
|
|
1322
1517
|
warnings.warn(
|
|
1323
1518
|
f"Degrees of freedom is non-positive (df={effective_df}). "
|
|
1324
1519
|
"Using normal distribution instead of t-distribution for inference.",
|
|
@@ -1396,10 +1591,7 @@ class LinearRegression:
|
|
|
1396
1591
|
Inference results for each coefficient in order.
|
|
1397
1592
|
"""
|
|
1398
1593
|
self._check_fitted()
|
|
1399
|
-
return [
|
|
1400
|
-
self.get_inference(i, alpha=alpha, df=df)
|
|
1401
|
-
for i in range(len(self.coefficients_))
|
|
1402
|
-
]
|
|
1594
|
+
return [self.get_inference(i, alpha=alpha, df=df) for i in range(len(self.coefficients_))]
|
|
1403
1595
|
|
|
1404
1596
|
def r_squared(self, adjusted: bool = False) -> float:
|
|
1405
1597
|
"""
|
|
@@ -1424,9 +1616,7 @@ class LinearRegression:
|
|
|
1424
1616
|
self._check_fitted()
|
|
1425
1617
|
# Use effective params for adjusted R² to match df correction
|
|
1426
1618
|
n_params = self.n_params_effective_ if adjusted else self.n_params_
|
|
1427
|
-
return compute_r_squared(
|
|
1428
|
-
self._y, self.residuals_, adjusted=adjusted, n_params=n_params
|
|
1429
|
-
)
|
|
1619
|
+
return compute_r_squared(self._y, self.residuals_, adjusted=adjusted, n_params=n_params)
|
|
1430
1620
|
|
|
1431
1621
|
def predict(self, X: np.ndarray) -> np.ndarray:
|
|
1432
1622
|
"""
|