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.
Files changed (52) hide show
  1. {diff_diff-2.7.0 → diff_diff-2.7.1}/PKG-INFO +53 -2
  2. {diff_diff-2.7.0 → diff_diff-2.7.1}/README.md +52 -1
  3. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/__init__.py +1 -1
  4. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/linalg.py +217 -27
  5. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered.py +276 -253
  6. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered_results.py +87 -66
  7. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/sun_abraham.py +9 -15
  8. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/triple_diff.py +241 -258
  9. {diff_diff-2.7.0 → diff_diff-2.7.1}/pyproject.toml +3 -3
  10. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/Cargo.lock +1 -1
  11. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/Cargo.toml +1 -1
  12. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/_backend.py +0 -0
  13. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/bacon.py +0 -0
  14. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/bootstrap_utils.py +0 -0
  15. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/continuous_did.py +0 -0
  16. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/continuous_did_bspline.py +0 -0
  17. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/continuous_did_results.py +0 -0
  18. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/datasets.py +0 -0
  19. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/diagnostics.py +0 -0
  20. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did.py +0 -0
  21. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did_bootstrap.py +0 -0
  22. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did_results.py +0 -0
  23. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/efficient_did_weights.py +0 -0
  24. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/estimators.py +0 -0
  25. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/honest_did.py +0 -0
  26. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/imputation.py +0 -0
  27. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/imputation_bootstrap.py +0 -0
  28. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/imputation_results.py +0 -0
  29. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/power.py +0 -0
  30. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/prep.py +0 -0
  31. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/prep_dgp.py +0 -0
  32. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/pretrends.py +0 -0
  33. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/results.py +0 -0
  34. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/stacked_did.py +0 -0
  35. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/stacked_did_results.py +0 -0
  36. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered_aggregation.py +0 -0
  37. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/staggered_bootstrap.py +0 -0
  38. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/synthetic_did.py +0 -0
  39. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/trop.py +0 -0
  40. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/trop_results.py +0 -0
  41. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/twfe.py +0 -0
  42. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/two_stage.py +0 -0
  43. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/two_stage_bootstrap.py +0 -0
  44. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/two_stage_results.py +0 -0
  45. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/utils.py +0 -0
  46. {diff_diff-2.7.0 → diff_diff-2.7.1}/diff_diff/visualization.py +0 -0
  47. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/build.rs +0 -0
  48. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/bootstrap.rs +0 -0
  49. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/lib.rs +0 -0
  50. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/linalg.rs +0 -0
  51. {diff_diff-2.7.0 → diff_diff-2.7.1}/rust/src/trop.rs +0 -0
  52. {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.0
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), and Stacked DiD (Wing, Freedman & Hollingsworth 2024) estimators for heterogeneous treatment timing
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), and Stacked DiD (Wing, Freedman & Hollingsworth 2024) estimators for heterogeneous treatment timing
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).
@@ -179,7 +179,7 @@ Stacked = StackedDiD
179
179
  Bacon = BaconDecomposition
180
180
  EDiD = EfficientDiD
181
181
 
182
- __version__ = "2.7.0"
182
+ __version__ = "2.7.1"
183
183
  __all__ = [
184
184
  # Estimators
185
185
  "DifferenceInDifferences",
@@ -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='economic', pivoting=True)
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[Union[
255
- Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
256
- Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
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, y,
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, y,
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, y,
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, y,
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, y,
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, y,
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, y,
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, y,
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
  """