diff-diff 3.5.0__tar.gz → 3.5.2__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 (94) hide show
  1. {diff_diff-3.5.0 → diff_diff-3.5.2}/PKG-INFO +3 -3
  2. {diff_diff-3.5.0 → diff_diff-3.5.2}/README.md +2 -2
  3. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/__init__.py +1 -1
  4. diff_diff-3.5.2/diff_diff/balancing.py +209 -0
  5. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/business_report.py +3 -0
  6. diff_diff-3.5.2/diff_diff/conformal.py +399 -0
  7. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/diagnostic_report.py +226 -4
  8. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/efficient_did.py +96 -32
  9. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/efficient_did_covariates.py +238 -62
  10. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/efficient_did_results.py +5 -2
  11. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/estimators.py +70 -3
  12. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/guides/llms-full.txt +18 -10
  13. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/guides/llms.txt +2 -2
  14. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/had.py +77 -0
  15. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/imputation.py +130 -101
  16. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/practitioner.py +36 -0
  17. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/results.py +44 -9
  18. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/stacked_did.py +275 -1
  19. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/stacked_did_results.py +10 -0
  20. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/synthetic_control.py +1063 -26
  21. diff_diff-3.5.2/diff_diff/synthetic_control_results.py +2449 -0
  22. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/synthetic_did.py +11 -11
  23. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/twfe.py +39 -1
  24. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/utils.py +151 -1
  25. {diff_diff-3.5.0 → diff_diff-3.5.2}/pyproject.toml +1 -1
  26. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/Cargo.lock +5 -5
  27. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/Cargo.toml +1 -1
  28. diff_diff-3.5.0/diff_diff/synthetic_control_results.py +0 -729
  29. {diff_diff-3.5.0 → diff_diff-3.5.2}/LICENSE +0 -0
  30. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/_backend.py +0 -0
  31. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/_guides_api.py +0 -0
  32. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/_nprobust_port.py +0 -0
  33. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/_reporting_helpers.py +0 -0
  34. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/agent_workflow.py +0 -0
  35. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/bacon.py +0 -0
  36. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/bootstrap_utils.py +0 -0
  37. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/chaisemartin_dhaultfoeuille.py +0 -0
  38. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +0 -0
  39. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/chaisemartin_dhaultfoeuille_results.py +0 -0
  40. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/conley.py +0 -0
  41. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/continuous_did.py +0 -0
  42. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/continuous_did_bspline.py +0 -0
  43. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/continuous_did_results.py +0 -0
  44. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/datasets.py +0 -0
  45. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/diagnostics.py +0 -0
  46. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/efficient_did_bootstrap.py +0 -0
  47. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/efficient_did_weights.py +0 -0
  48. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/guides/__init__.py +0 -0
  49. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/guides/llms-autonomous.txt +0 -0
  50. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/guides/llms-practitioner.txt +0 -0
  51. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/had_pretests.py +0 -0
  52. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/honest_did.py +0 -0
  53. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/imputation_bootstrap.py +0 -0
  54. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/imputation_results.py +0 -0
  55. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/linalg.py +0 -0
  56. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/local_linear.py +0 -0
  57. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/power.py +0 -0
  58. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/prep.py +0 -0
  59. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/prep_dgp.py +0 -0
  60. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/pretrends.py +0 -0
  61. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/profile.py +0 -0
  62. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/spillover.py +0 -0
  63. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/staggered.py +0 -0
  64. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/staggered_aggregation.py +0 -0
  65. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/staggered_bootstrap.py +0 -0
  66. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/staggered_results.py +0 -0
  67. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/staggered_triple_diff.py +0 -0
  68. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/staggered_triple_diff_results.py +0 -0
  69. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/sun_abraham.py +0 -0
  70. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/survey.py +0 -0
  71. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/triple_diff.py +0 -0
  72. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/trop.py +0 -0
  73. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/trop_global.py +0 -0
  74. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/trop_local.py +0 -0
  75. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/trop_results.py +0 -0
  76. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/two_stage.py +0 -0
  77. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/two_stage_bootstrap.py +0 -0
  78. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/two_stage_results.py +0 -0
  79. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/__init__.py +0 -0
  80. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_common.py +0 -0
  81. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_continuous.py +0 -0
  82. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_diagnostic.py +0 -0
  83. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_event_study.py +0 -0
  84. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_power.py +0 -0
  85. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_staggered.py +0 -0
  86. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/visualization/_synthetic.py +0 -0
  87. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/wooldridge.py +0 -0
  88. {diff_diff-3.5.0 → diff_diff-3.5.2}/diff_diff/wooldridge_results.py +0 -0
  89. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/build.rs +0 -0
  90. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/src/bootstrap.rs +0 -0
  91. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/src/lib.rs +0 -0
  92. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/src/linalg.rs +0 -0
  93. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/src/trop.rs +0 -0
  94. {diff_diff-3.5.0 → diff_diff-3.5.2}/rust/src/weights.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diff-diff
3
- Version: 3.5.0
3
+ Version: 3.5.2
4
4
  Classifier: Development Status :: 5 - Production/Stable
5
5
  Classifier: Intended Audience :: Science/Research
6
6
  Classifier: Operating System :: OS Independent
@@ -161,11 +161,11 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
161
161
  - [TwoStageDiD](https://diff-diff.readthedocs.io/en/stable/api/two_stage.html) - Gardner (2022) two-stage estimator with GMM sandwich variance
162
162
  - [SpilloverDiD](https://diff-diff.readthedocs.io/en/stable/api/spillover.html) - Butts (2021) ring-indicator spillover-aware DiD identifying direct effect on treated + per-ring spillover on near-control units; handles non-staggered and staggered timing; supports survey-design variance under `survey_design=` for HC1 / CR1 (Wave E.1 Binder TSL) and Conley (Wave E.2 panel-aware stratified-Conley sandwich on per-period PSU totals; extended in Wave E.2 follow-up to `conley_lag_cutoff > 0` via panel-block composition with within-PSU serial Bartlett HAC — `lag>0` requires an effective PSU via explicit `survey_design.psu` or injected `cluster=<col>`); `SurveyDesign.subpopulation()` preserves full-design `n_psu` / `df_survey` via zero-padded scores (Wave E.3, R `svyrecvar(subset())` form)
163
163
  - [SyntheticDiD](https://diff-diff.readthedocs.io/en/stable/api/estimators.html) - Synthetic DiD combining standard DiD and synthetic control for few treated units
164
- - [SyntheticControl](https://diff-diff.readthedocs.io/en/stable/api/synthetic_control.html) - Abadie, Diamond & Hainmueller (2010) classic synthetic control for a single treated unit (donor-weight counterfactual, nested/custom V; in-space placebo permutation inference via `in_space_placebo()`)
164
+ - [SyntheticControl](https://diff-diff.readthedocs.io/en/stable/api/synthetic_control.html) - Abadie, Diamond & Hainmueller (2010) classic synthetic control for a single treated unit (donor-weight counterfactual, nested/cv/inverse-variance/custom V; in-space placebo permutation inference via `in_space_placebo()`, plus ADH-2015 `leave_one_out()` + `in_time_placebo()` robustness, Firpo-Possebom (2018) test-inversion confidence sets, and Chernozhukov-Wüthrich-Zhu (2021) conformal inference)
165
165
  - [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html) - triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility
166
166
  - [ContinuousDiD](https://diff-diff.readthedocs.io/en/stable/api/continuous_did.html) - Callaway, Goodman-Bacon & Sant'Anna (2024) continuous treatment DiD with dose-response curves
167
167
  - [HeterogeneousAdoptionDiD](https://diff-diff.readthedocs.io/en/stable/api/had.html) - de Chaisemartin, Ciccia, D'Haultfœuille & Knau (2026) for designs where **no unit remains untreated**; local-linear estimator at the dose support boundary returning Weighted Average Slope (WAS) on Design 1' (`d̲ = 0` / QUG) or `WAS_{d̲}` on Design 1 (`d̲ > 0`, continuous-near-d̲ or mass-point), with a multi-period event-study extension (last-treatment cohort, pointwise CIs). **Panel-only** in this release - repeated cross-sections rejected by the validator. Alias `HAD`.
168
- - [StackedDiD](https://diff-diff.readthedocs.io/en/stable/api/stacked_did.html) - Wing, Freedman & Hollingsworth (2024) stacked DiD with Q-weights and sub-experiments
168
+ - [StackedDiD](https://diff-diff.readthedocs.io/en/stable/api/stacked_did.html) - Wing, Freedman & Hollingsworth (2024) stacked DiD with Q-weights and sub-experiments; optional covariate balancing (Ustyuzhanin 2026)
169
169
  - [EfficientDiD](https://diff-diff.readthedocs.io/en/stable/api/efficient_did.html) - Chen, Sant'Anna & Xie (2025) efficient DiD with optimal weighting for tighter SEs
170
170
  - [TROP](https://diff-diff.readthedocs.io/en/stable/api/trop.html) - Triply Robust Panel estimator (Athey et al. 2025) with nuclear norm factor adjustment
171
171
  - [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference) - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT
@@ -108,11 +108,11 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
108
108
  - [TwoStageDiD](https://diff-diff.readthedocs.io/en/stable/api/two_stage.html) - Gardner (2022) two-stage estimator with GMM sandwich variance
109
109
  - [SpilloverDiD](https://diff-diff.readthedocs.io/en/stable/api/spillover.html) - Butts (2021) ring-indicator spillover-aware DiD identifying direct effect on treated + per-ring spillover on near-control units; handles non-staggered and staggered timing; supports survey-design variance under `survey_design=` for HC1 / CR1 (Wave E.1 Binder TSL) and Conley (Wave E.2 panel-aware stratified-Conley sandwich on per-period PSU totals; extended in Wave E.2 follow-up to `conley_lag_cutoff > 0` via panel-block composition with within-PSU serial Bartlett HAC — `lag>0` requires an effective PSU via explicit `survey_design.psu` or injected `cluster=<col>`); `SurveyDesign.subpopulation()` preserves full-design `n_psu` / `df_survey` via zero-padded scores (Wave E.3, R `svyrecvar(subset())` form)
110
110
  - [SyntheticDiD](https://diff-diff.readthedocs.io/en/stable/api/estimators.html) - Synthetic DiD combining standard DiD and synthetic control for few treated units
111
- - [SyntheticControl](https://diff-diff.readthedocs.io/en/stable/api/synthetic_control.html) - Abadie, Diamond & Hainmueller (2010) classic synthetic control for a single treated unit (donor-weight counterfactual, nested/custom V; in-space placebo permutation inference via `in_space_placebo()`)
111
+ - [SyntheticControl](https://diff-diff.readthedocs.io/en/stable/api/synthetic_control.html) - Abadie, Diamond & Hainmueller (2010) classic synthetic control for a single treated unit (donor-weight counterfactual, nested/cv/inverse-variance/custom V; in-space placebo permutation inference via `in_space_placebo()`, plus ADH-2015 `leave_one_out()` + `in_time_placebo()` robustness, Firpo-Possebom (2018) test-inversion confidence sets, and Chernozhukov-Wüthrich-Zhu (2021) conformal inference)
112
112
  - [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html) - triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility
113
113
  - [ContinuousDiD](https://diff-diff.readthedocs.io/en/stable/api/continuous_did.html) - Callaway, Goodman-Bacon & Sant'Anna (2024) continuous treatment DiD with dose-response curves
114
114
  - [HeterogeneousAdoptionDiD](https://diff-diff.readthedocs.io/en/stable/api/had.html) - de Chaisemartin, Ciccia, D'Haultfœuille & Knau (2026) for designs where **no unit remains untreated**; local-linear estimator at the dose support boundary returning Weighted Average Slope (WAS) on Design 1' (`d̲ = 0` / QUG) or `WAS_{d̲}` on Design 1 (`d̲ > 0`, continuous-near-d̲ or mass-point), with a multi-period event-study extension (last-treatment cohort, pointwise CIs). **Panel-only** in this release - repeated cross-sections rejected by the validator. Alias `HAD`.
115
- - [StackedDiD](https://diff-diff.readthedocs.io/en/stable/api/stacked_did.html) - Wing, Freedman & Hollingsworth (2024) stacked DiD with Q-weights and sub-experiments
115
+ - [StackedDiD](https://diff-diff.readthedocs.io/en/stable/api/stacked_did.html) - Wing, Freedman & Hollingsworth (2024) stacked DiD with Q-weights and sub-experiments; optional covariate balancing (Ustyuzhanin 2026)
116
116
  - [EfficientDiD](https://diff-diff.readthedocs.io/en/stable/api/efficient_did.html) - Chen, Sant'Anna & Xie (2025) efficient DiD with optimal weighting for tighter SEs
117
117
  - [TROP](https://diff-diff.readthedocs.io/en/stable/api/trop.html) - Triply Robust Panel estimator (Athey et al. 2025) with nuclear norm factor adjustment
118
118
  - [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference) - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT
@@ -298,7 +298,7 @@ ETWFE = WooldridgeDiD
298
298
  DCDH = ChaisemartinDHaultfoeuille
299
299
  HAD = HeterogeneousAdoptionDiD
300
300
 
301
- __version__ = "3.5.0"
301
+ __version__ = "3.5.2"
302
302
  __all__ = [
303
303
  # Estimators
304
304
  "DifferenceInDifferences",
@@ -0,0 +1,209 @@
1
+ """Covariate balancing weights (entropy balancing).
2
+
3
+ Implements **entropy balancing** (Hainmueller, J. (2012). "Entropy Balancing for
4
+ Causal Effects: A Multivariate Reweighting Method to Produce Balanced Samples in
5
+ Observational Studies." *Political Analysis*, 20(1), 25-46.
6
+ https://doi.org/10.1093/pan/mpr025).
7
+
8
+ Entropy balancing finds nonnegative control weights ``w_i`` that exactly match a set
9
+ of target covariate moments (here: the treated-group covariate means) while staying as
10
+ close as possible — in the Kullback-Leibler sense — to a set of base weights (uniform by
11
+ default). The solution is obtained from the convex dual
12
+
13
+ minimize over λ: L(λ) = log( Σ_i q_i exp(Z_iᵀ λ) ), Z_i = X_i − target,
14
+
15
+ whose stationary point ``∇L = Σ_i w_i Z_i = 0`` is exactly first-moment balance, with
16
+ ``w_i = q_i exp(Z_iᵀ λ) / Σ_j q_j exp(Z_jᵀ λ)``. ``L`` is convex (log-sum-exp), so a
17
+ damped Newton iteration (gradient = weighted mean of the centered moments, Hessian =
18
+ weighted covariance) converges to the balancing weights whenever the target lies in the
19
+ interior of the control covariate convex hull; otherwise no finite λ balances the
20
+ moments and the problem is **infeasible**.
21
+
22
+ This module is dependency-light (numpy, with an optional scipy L-BFGS fallback) and is
23
+ used by ``StackedDiD`` to construct the within-sub-experiment design weights ``b_{sa}``
24
+ for Covariate-Balanced Weighted Stacked DID (Ustyuzhanin 2026).
25
+ """
26
+
27
+ from __future__ import annotations
28
+
29
+ from typing import Any, Dict, Optional, Tuple
30
+
31
+ import numpy as np
32
+
33
+ __all__ = ["entropy_balance", "BalanceError"]
34
+
35
+
36
+ class BalanceError(ValueError):
37
+ """Raised when entropy balancing fails to achieve first-moment balance.
38
+
39
+ Carries the achieved ``max_residual`` and the per-covariate residual vector so
40
+ callers (e.g. ``StackedDiD``) can attach cohort context and report the worst-
41
+ balanced covariate.
42
+ """
43
+
44
+ def __init__(self, message: str, *, max_residual: float, residuals: np.ndarray):
45
+ super().__init__(message)
46
+ self.max_residual = max_residual
47
+ self.residuals = residuals
48
+
49
+
50
+ def entropy_balance(
51
+ X: np.ndarray,
52
+ target_means: np.ndarray,
53
+ base_weights: Optional[np.ndarray] = None,
54
+ *,
55
+ max_iter: int = 200,
56
+ tol: float = 1e-8,
57
+ ) -> Tuple[np.ndarray, Dict[str, Any]]:
58
+ """Solve entropy balancing for control weights matching ``target_means``.
59
+
60
+ Parameters
61
+ ----------
62
+ X : np.ndarray, shape (n, k)
63
+ Covariate matrix for the ``n`` control units (``k`` covariates).
64
+ target_means : np.ndarray, shape (k,)
65
+ Target first moments to match — for CBWSDID these are the treated-group
66
+ covariate means.
67
+ base_weights : np.ndarray, shape (n,), optional
68
+ Nonnegative base weights ``q_i`` (the KL reference). Defaults to uniform.
69
+ Internally renormalized to sum to one.
70
+ max_iter : int
71
+ Maximum damped-Newton iterations (then a scipy L-BFGS fallback is attempted).
72
+ tol : float
73
+ Convergence tolerance on the maximum absolute (raw-scale) moment residual
74
+ ``max_r |Σ_i w_i X_{i,r} − target_r|``.
75
+
76
+ Returns
77
+ -------
78
+ weights : np.ndarray, shape (n,)
79
+ Nonnegative weights summing to one with ``Σ_i w_i X_i ≈ target_means``.
80
+ info : dict
81
+ ``converged`` (bool), ``max_residual`` (float), ``n_iter`` (int),
82
+ ``ess`` (effective sample size ``1 / Σ_i w_i²``), ``solver`` (str).
83
+
84
+ Raises
85
+ ------
86
+ BalanceError
87
+ If neither the damped-Newton nor the L-BFGS fallback drives the maximum moment
88
+ residual below ``tol`` (the target is outside the control covariate hull, i.e.
89
+ infeasible).
90
+ ValueError
91
+ On malformed inputs (shape mismatch, non-finite, negative base weights).
92
+ """
93
+ X = np.asarray(X, dtype=np.float64)
94
+ if X.ndim != 2:
95
+ raise ValueError(f"X must be 2-D (n, k); got shape {X.shape}")
96
+ n, k = X.shape
97
+ target = np.asarray(target_means, dtype=np.float64).reshape(-1)
98
+ if target.shape[0] != k:
99
+ raise ValueError(f"target_means length {target.shape[0]} != n_covariates {k}")
100
+ if not np.all(np.isfinite(X)) or not np.all(np.isfinite(target)):
101
+ raise ValueError("X and target_means must be finite")
102
+ if n == 0:
103
+ raise ValueError("X has no control rows")
104
+
105
+ if base_weights is None:
106
+ q = np.full(n, 1.0 / n)
107
+ else:
108
+ q = np.asarray(base_weights, dtype=np.float64).reshape(-1)
109
+ if q.shape[0] != n:
110
+ raise ValueError(f"base_weights length {q.shape[0]} != n_control {n}")
111
+ if np.any(q < 0) or not np.all(np.isfinite(q)):
112
+ raise ValueError("base_weights must be nonnegative and finite")
113
+ s = q.sum()
114
+ if s <= 0:
115
+ raise ValueError("base_weights sum to zero")
116
+ q = q / s
117
+
118
+ # Centered moments; standardize columns for conditioning (balance set is invariant
119
+ # to the linear rescaling — it is absorbed into the dual variable λ).
120
+ Z = X - target
121
+ scale = Z.std(axis=0)
122
+ scale[scale < 1e-12] = 1.0
123
+ Zs = Z / scale
124
+
125
+ def weights_at(lam: np.ndarray) -> np.ndarray:
126
+ logits = Zs @ lam
127
+ logits -= logits.max()
128
+ ew = q * np.exp(logits)
129
+ return ew / ew.sum()
130
+
131
+ def dual_loss(lam: np.ndarray) -> float:
132
+ logits = Zs @ lam
133
+ m = logits.max()
134
+ return float(m + np.log(np.sum(q * np.exp(logits - m))))
135
+
136
+ def raw_residual(w: np.ndarray) -> np.ndarray:
137
+ return w @ X - target
138
+
139
+ lam = np.zeros(k)
140
+ solver = "newton"
141
+ n_iter = 0
142
+ for n_iter in range(1, max_iter + 1):
143
+ w = weights_at(lam)
144
+ if np.max(np.abs(raw_residual(w))) < tol:
145
+ break
146
+ g = w @ Zs # gradient of the dual loss (standardized scale)
147
+ Zc = Zs - g
148
+ H = (w[:, None] * Zc).T @ Zc # weighted covariance (PSD)
149
+ ridge = 1e-10 * (np.trace(H) / k + 1e-12)
150
+ try:
151
+ direction = -np.linalg.solve(H + ridge * np.eye(k), g)
152
+ except np.linalg.LinAlgError:
153
+ direction = -np.linalg.lstsq(H, g, rcond=None)[0]
154
+ # Backtracking (Armijo) line search on the convex dual loss.
155
+ base = dual_loss(lam)
156
+ slope = float(g @ direction) # < 0 (descent)
157
+ step = 1.0
158
+ for _ in range(40):
159
+ if dual_loss(lam + step * direction) <= base + 1e-4 * step * slope:
160
+ break
161
+ step *= 0.5
162
+ lam = lam + step * direction
163
+ else:
164
+ w = weights_at(lam)
165
+
166
+ w = weights_at(lam)
167
+ resid = raw_residual(w)
168
+ max_resid = float(np.max(np.abs(resid)))
169
+
170
+ if max_resid >= tol:
171
+ # Fallback: scipy L-BFGS-B on the convex dual (robust to poor Newton scaling).
172
+ try:
173
+ from scipy.optimize import minimize
174
+
175
+ res = minimize(
176
+ dual_loss,
177
+ lam,
178
+ jac=lambda L: weights_at(L) @ Zs,
179
+ method="L-BFGS-B",
180
+ options={"maxiter": 500, "gtol": 1e-12},
181
+ )
182
+ w_lbfgs = weights_at(res.x)
183
+ resid_lbfgs = raw_residual(w_lbfgs)
184
+ if np.max(np.abs(resid_lbfgs)) < max_resid:
185
+ lam, w, resid = res.x, w_lbfgs, resid_lbfgs
186
+ max_resid = float(np.max(np.abs(resid)))
187
+ solver = "lbfgs"
188
+ except Exception: # pragma: no cover - scipy always present, defensive
189
+ pass
190
+
191
+ converged = max_resid < tol
192
+ info: Dict[str, Any] = {
193
+ "converged": converged,
194
+ "max_residual": max_resid,
195
+ "n_iter": n_iter,
196
+ "ess": float(1.0 / np.sum(w**2)),
197
+ "solver": solver,
198
+ }
199
+ if not converged:
200
+ worst = int(np.argmax(np.abs(resid)))
201
+ raise BalanceError(
202
+ "entropy balancing did not converge to first-moment balance "
203
+ f"(max moment residual {max_resid:.3e} >= tol {tol:.1e}; worst covariate "
204
+ f"index {worst}). The target mean is likely outside the convex hull of the "
205
+ "control covariates (infeasible).",
206
+ max_residual=max_resid,
207
+ residuals=resid,
208
+ )
209
+ return w, info
@@ -1019,6 +1019,9 @@ def _lift_robustness(dr: Optional[Dict[str, Any]]) -> Dict[str, Any]:
1019
1019
  native_block["pre_rmspe"] = native.get("pre_rmspe")
1020
1020
  native_block["weight_concentration"] = native.get("weight_concentration")
1021
1021
  native_block["in_space_placebo"] = native.get("in_space_placebo")
1022
+ # ADH-2015 robustness diagnostics (opt-in; "not_run" stub until run).
1023
+ native_block["leave_one_out"] = native.get("leave_one_out")
1024
+ native_block["in_time_placebo"] = native.get("in_time_placebo")
1022
1025
  return {
1023
1026
  "bacon": {
1024
1027
  "status": bacon.get("status"),
@@ -0,0 +1,399 @@
1
+ """Conformal inference for counterfactual / synthetic controls (Chernozhukov,
2
+ Wüthrich & Zhu 2021, *JASA* 116(536):1849–1864).
3
+
4
+ Pure, numpy-only building blocks for the conformal-inference layer surfaced as
5
+ opt-in methods on :class:`SyntheticControlResults`
6
+ (:meth:`~diff_diff.synthetic_control_results.SyntheticControlResults.conformal_test`,
7
+ ``conformal_confidence_intervals``, ``conformal_average_effect``). See
8
+ ``docs/methodology/papers/chernozhukov-wuthrich-zhu-2021-review.md`` and the
9
+ ``## SyntheticControl`` section of ``docs/methodology/REGISTRY.md``.
10
+
11
+ **Method (CWZ §2.2).** Under a sharp null ``H0: θ = θ0`` over the post period,
12
+ impute the counterfactual treated outcomes (``Y^N_{1t} = Y_{1t} − θ0_t`` for
13
+ ``t > T0``; pre-period unchanged), fit a *time-permutation-invariant* proxy on
14
+ **all** periods under that null, take residuals ``û_t = Y^N_{1t} − P̂^N_t``, form
15
+ the statistic ``S_q(û) = ((1/√T*)·Σ_{t>T0}|û_t|^q)^{1/q}`` (high → reject), and
16
+ compute a permutation p-value (eq 2) ``p̂ = (1/|Π|)·#{π∈Π : S(û_π) ≥ S(û)}`` by
17
+ reshuffling residuals over time. ``Π`` contains the identity, so ``p̂ ≥ 1/|Π|``
18
+ automatically (there is NO extra ``+1`` — distinct from the cross-unit placebo
19
+ ``(1+n)/(n+1)``). Inverting the test over a grid of ``θ0`` gives confidence sets.
20
+
21
+ **Proxy.** The proxy is the canonical CWZ constrained-LS synthetic control
22
+ (eqs 3–4): simplex weights minimising ``Σ_{t}(Y^N_{1t} − Σ_j w_j Y^N_{jt})²``
23
+ over **all** periods under the null, ``w ≥ 0, Σ w = 1``, **no V-matrix, no
24
+ intercept, outcomes-only** (footnote 9: "we estimate w under the null based on
25
+ all the data"). This is DISTINCT from the headline ADH V-matrix weights — CWZ's
26
+ exactness theory (Lemma 1, Appendix D exchangeability) requires a time-symmetric
27
+ proxy, which the ADH pre-period V-fit is not. Reuses the Frank-Wolfe simplex
28
+ solver :func:`diff_diff.utils._sc_weight_fw`.
29
+ """
30
+
31
+ import math
32
+ import warnings
33
+ from typing import Any, Dict, List, Optional, Tuple
34
+
35
+ import numpy as np
36
+
37
+ from diff_diff.utils import _sc_weight_fw
38
+
39
+ __all__: List[str] = []
40
+
41
+ _INF = float("inf")
42
+
43
+
44
+ # =============================================================================
45
+ # Proxy fit (CWZ §2.3, eqs 3-4) — canonical constrained-LS synthetic control
46
+ # =============================================================================
47
+
48
+
49
+ def _cwz_proxy_fit(
50
+ y1: np.ndarray,
51
+ Y0: np.ndarray,
52
+ *,
53
+ max_iter: int,
54
+ min_decrease: float,
55
+ init_weights: Optional[np.ndarray] = None,
56
+ ) -> Tuple[np.ndarray, np.ndarray, bool]:
57
+ """Fit the canonical CWZ constrained-LS SC proxy over ALL given periods (eq 4).
58
+
59
+ Minimises ``Σ_t (y1_t − Σ_j w_j Y0_{tj})²`` s.t. ``w ≥ 0, Σ w = 1`` (no
60
+ intercept, no V-matrix). ``y1`` is the (already null-imputed) treated outcome
61
+ vector ``(T,)``; ``Y0`` the donor outcomes ``(T, J)``. Reuses
62
+ :func:`_sc_weight_fw` packed ``(T, J+1)`` (donor columns first, target last),
63
+ ``zeta=0, intercept=False`` — the solver projects onto the unit simplex, so the
64
+ ``w ≥ 0, Σ w = 1`` constraint is delivered with no extra normalization.
65
+
66
+ ``min_decrease`` is the ALREADY-SCALED absolute convergence tolerance; the
67
+ caller must pass a θ0-invariant scale (e.g. the pre-window outcome norm) so CI
68
+ membership does not drift with the grid value via the tolerance.
69
+
70
+ Returns ``(w (J,), resid = y1 − Y0 @ w (T,), converged)``. ``J == 1`` is the
71
+ degenerate single-donor case ``w = [1]`` (no optimisation).
72
+ """
73
+ y1 = np.asarray(y1, dtype=float)
74
+ Y0 = np.asarray(Y0, dtype=float)
75
+ _, J = Y0.shape
76
+ if J == 1:
77
+ w = np.array([1.0], dtype=float)
78
+ return w, y1 - Y0 @ w, True
79
+ packed = np.column_stack([Y0, y1]) # (T, J+1); last column is the target
80
+ with warnings.catch_warnings():
81
+ warnings.filterwarnings("ignore", message=r".*did not converge.*", category=UserWarning)
82
+ w, converged = _sc_weight_fw(
83
+ packed,
84
+ zeta=0.0,
85
+ intercept=False,
86
+ init_weights=init_weights,
87
+ min_decrease=float(min_decrease),
88
+ max_iter=int(max_iter),
89
+ return_convergence=True,
90
+ )
91
+ w = np.asarray(w, dtype=float)
92
+ return w, y1 - Y0 @ w, bool(converged)
93
+
94
+
95
+ # =============================================================================
96
+ # Test statistic (CWZ §2.2) and permutation sets
97
+ # =============================================================================
98
+
99
+
100
+ def _cwz_statistic(u: np.ndarray, post_mask: np.ndarray, q: Any) -> float:
101
+ """``S_q(û) = ((1/√T*)·Σ_{t>T0}|û_t|^q)^{1/q}`` (CWZ §2.2).
102
+
103
+ ``q=∞`` → ``max_{t>T0}|û_t|`` (the ``√T*`` factor vanishes in the sup limit).
104
+ ``post_mask`` is a boolean over the (calendar-ordered) periods selecting the
105
+ post window. Returns NaN if the post window is empty.
106
+ """
107
+ post = np.abs(u[post_mask])
108
+ n_star = post.size
109
+ if n_star == 0:
110
+ return float("nan")
111
+ if q == _INF:
112
+ return float(np.max(post))
113
+ s = float(np.sum(post**q))
114
+ return float((s / np.sqrt(n_star)) ** (1.0 / q))
115
+
116
+
117
+ def _moving_block_perms(m: int) -> np.ndarray:
118
+ """Moving-block permutations ``Π_→`` — ``m`` cyclic shifts (CWZ §2.2).
119
+
120
+ Row ``j`` (``j = 0,…,m−1``) is the index array of ``π_j(i) = (i + j) mod m``
121
+ (the paper's 1-based ``i+j`` wrapped at ``m``); ``û_π = û[row]``. Row 0 is the
122
+ identity. Shape ``(m, m)``. Valid under stationary, weakly-dependent errors
123
+ (Assumption 2.2).
124
+ """
125
+ base = np.arange(m)
126
+ return (base[None, :] + base[:, None]) % m
127
+
128
+
129
+ def _iid_perms(m: int, n_draws: int, rng: np.random.Generator) -> np.ndarray:
130
+ """i.i.d. permutations ``Π_all`` (CWZ §2.2).
131
+
132
+ Exact ``m!`` enumeration ONLY when genuinely small (``m ≤ 12`` and
133
+ ``m! ≤ n_draws``); otherwise ``n_draws`` random draws with the identity
134
+ prepended (the identity must be in ``Π``). Valid under i.i.d. errors
135
+ (Assumption 2.1). p-values below ``1/n_draws`` are unattainable.
136
+ """
137
+ if m <= 12 and math.factorial(m) <= n_draws:
138
+ from itertools import permutations
139
+
140
+ return np.array(list(permutations(range(m))), dtype=int)
141
+ ident = np.arange(m)[None, :]
142
+ if n_draws <= 1:
143
+ return ident
144
+ draws = np.stack([rng.permutation(m) for _ in range(n_draws - 1)])
145
+ return np.vstack([ident, draws]).astype(int)
146
+
147
+
148
+ def _make_perms(m: int, scheme: str, n_iid: int, rng: np.random.Generator) -> np.ndarray:
149
+ """Dispatch to the requested permutation set (both include the identity)."""
150
+ if scheme == "moving_block":
151
+ return _moving_block_perms(m)
152
+ if scheme == "iid":
153
+ return _iid_perms(m, n_iid, rng)
154
+ raise ValueError(f"scheme must be 'moving_block' or 'iid', got {scheme!r}")
155
+
156
+
157
+ # =============================================================================
158
+ # Permutation p-value (CWZ eq 2)
159
+ # =============================================================================
160
+
161
+
162
+ def _cwz_pvalue(
163
+ u: np.ndarray, post_mask: np.ndarray, perms: np.ndarray, q: Any
164
+ ) -> Tuple[float, float, int]:
165
+ """CWZ eq (2) permutation p-value for residuals ``u``.
166
+
167
+ ``p̂ = (1/|Π|)·#{π ∈ Π : S(û_π) ≥ S(û)}``. The residual vector is permuted and
168
+ the statistic always reads the FIXED post-window slots (``post_mask``); the
169
+ identity ``π`` is in ``perms`` so ``p̂ ≥ 1/|Π|`` automatically (no extra ``+1``).
170
+ Ties (``≥``) are counted conservatively. ``perms`` is the ``(|Π|, len(u))``
171
+ integer index array from :func:`_make_perms`. Returns ``(p, S_observed, |Π|)``.
172
+ """
173
+ s_obs = _cwz_statistic(u, post_mask, q)
174
+ post_perm = np.abs(u[perms][:, post_mask]) # (|Π|, T*)
175
+ n_star = post_perm.shape[1]
176
+ if q == _INF:
177
+ s_perm = post_perm.max(axis=1)
178
+ else:
179
+ s_perm = (post_perm**q).sum(axis=1) / np.sqrt(n_star)
180
+ s_perm = s_perm ** (1.0 / q)
181
+ n = perms.shape[0]
182
+ tol = 1e-12 * max(abs(s_obs), 1.0)
183
+ n_ge = int(np.sum(s_perm >= s_obs - tol))
184
+ return n_ge / n, float(s_obs), n
185
+
186
+
187
+ # =============================================================================
188
+ # Single-null evaluation (joint sharp-null test) and grid inversion (CIs)
189
+ # =============================================================================
190
+
191
+
192
+ def _single_null_pvalue(
193
+ y1_obs: np.ndarray,
194
+ Y0: np.ndarray,
195
+ post_mask: np.ndarray,
196
+ effect_on_post: np.ndarray,
197
+ perms: np.ndarray,
198
+ q: Any,
199
+ *,
200
+ max_iter: int,
201
+ min_decrease: float,
202
+ init_weights: Optional[np.ndarray] = None,
203
+ ) -> Dict[str, Any]:
204
+ """Build ``Z(θ0)``, fit the proxy under the null, and return the p-value.
205
+
206
+ ``effect_on_post`` (length = ``post_mask.sum()``, calendar order) is subtracted
207
+ from the observed treated outcomes in the post window. The proxy is fit ONCE on
208
+ the null-imputed data and residuals are permuted (CWZ footnote 7 — the proxy is
209
+ time-permutation-invariant, so permuting residuals ≡ permuting data).
210
+ """
211
+ y1n = np.asarray(y1_obs, dtype=float).copy()
212
+ y1n[post_mask] = y1n[post_mask] - np.asarray(effect_on_post, dtype=float)
213
+ w, resid, conv = _cwz_proxy_fit(
214
+ y1n, Y0, max_iter=max_iter, min_decrease=min_decrease, init_weights=init_weights
215
+ )
216
+ p, s_obs, n = _cwz_pvalue(resid, post_mask, perms, q)
217
+ return {"p_value": p, "s_observed": s_obs, "n_perms": n, "converged": conv, "weights": w}
218
+
219
+
220
+ def _auto_grid(center: float, spread: float, n_grid: int) -> np.ndarray:
221
+ """Symmetric inversion grid centred on the point estimate.
222
+
223
+ Half-width is a generous multiple of the pre-period residual scale (so a valid
224
+ CI lies inside the grid for well-behaved fits); membership outside the grid is
225
+ NOT certified (grid-limited — flagged via ``status`` when the accepted set
226
+ touches an edge).
227
+ """
228
+ half = max(8.0 * spread, 0.5 * abs(center) + 1e-6)
229
+ return np.linspace(center - half, center + half, n_grid)
230
+
231
+
232
+ def _assemble_single_ci(
233
+ thetas: np.ndarray,
234
+ pvals: np.ndarray,
235
+ converged: np.ndarray,
236
+ alpha: float,
237
+ point_estimate: float,
238
+ ) -> Dict[str, Any]:
239
+ """Assemble a confidence interval from a single-post-slot inversion grid.
240
+
241
+ A grid point is **rejected** ONLY when it converged AND ``p ≤ alpha``; a
242
+ non-converged point is **indeterminate** (its p is unreliable) and is therefore
243
+ NOT rejected — keeping it in the set is the conservative choice for an interval
244
+ (excluding tail non-convergence would understate the width, the opposite of
245
+ fail-closed). The confidence set is the complement of the rejections; ``in_set``
246
+ (returned for the grid table) is ``~rejected``. The hull is ``[min, max]`` of the
247
+ in-set ``θ`` with a ``contiguous`` flag (a *rejected* point strictly inside the
248
+ hull ⇒ non-contiguous). ``n_nonconverged`` is surfaced so an interval widened by
249
+ indeterminacy is detectable. ``status`` is one of ``{"ran", "grid_limited",
250
+ "empty"}`` (the granularity-``unbounded`` case is handled by the caller).
251
+ """
252
+ rejected = converged & (pvals <= alpha)
253
+ in_set = ~rejected
254
+ n_nonconv = int(np.sum(~converged))
255
+ acc = thetas[in_set]
256
+ if acc.size == 0:
257
+ return {
258
+ "lower": float("nan"),
259
+ "upper": float("nan"),
260
+ "status": "empty",
261
+ "contiguous": True,
262
+ "point_estimate": float(point_estimate),
263
+ "n_in_set": 0,
264
+ "n_nonconverged": n_nonconv,
265
+ "_in_set": in_set,
266
+ }
267
+ lower, upper = float(acc.min()), float(acc.max())
268
+ inside = (thetas > lower) & (thetas < upper)
269
+ contiguous = bool(not np.any(rejected[inside]))
270
+ touches_edge = bool(in_set[0] or in_set[-1])
271
+ status = "grid_limited" if touches_edge else "ran"
272
+ return {
273
+ "lower": lower,
274
+ "upper": upper,
275
+ "status": status,
276
+ "contiguous": contiguous,
277
+ "point_estimate": float(point_estimate),
278
+ "n_in_set": int(acc.size),
279
+ "n_nonconverged": n_nonconv,
280
+ "_in_set": in_set,
281
+ }
282
+
283
+
284
+ def _invert_single_post(
285
+ y1_obs: np.ndarray,
286
+ Y0: np.ndarray,
287
+ post_idx: int,
288
+ alpha: float,
289
+ perms: np.ndarray,
290
+ *,
291
+ max_iter: int,
292
+ min_decrease: float,
293
+ grid: Optional[np.ndarray] = None,
294
+ n_grid: int = 100,
295
+ ) -> Dict[str, Any]:
296
+ """Invert a single-post-slot conformal test over a grid (CWZ Algorithm 1).
297
+
298
+ The series ``y1_obs`` ``(m,)`` / ``Y0`` ``(m, J)`` has exactly one post slot at
299
+ ``post_idx`` (a per-period CI sub-series ``(pre ∪ {t})`` or a block-collapsed
300
+ average series). For each candidate effect ``θ`` the post outcome is imputed
301
+ (``y1[post_idx] − θ``), the proxy refit (warm-started across the grid — FW is
302
+ convex so warm-starting only affects speed, not the optimum), and the
303
+ permutation p-value recomputed. With one post slot ``S_q`` reduces to
304
+ ``|û_post|`` for every ``q``, so ``q`` is inert here (fixed to 1). Returns the CI
305
+ summary plus the ``grid`` table of ``(θ, p, in_set, converged)`` rows.
306
+
307
+ When ``alpha < 1/|Π|`` every candidate has ``p ≥ 1/|Π| > alpha`` (the identity is
308
+ in ``Π``), so NO value is ever rejected and the confidence set is the whole line:
309
+ this short-circuits to ``status="unbounded"`` with ``(-inf, +inf)`` endpoints and
310
+ an empty grid (the per-point table is uninformative when nothing can be rejected).
311
+ """
312
+ y1_obs = np.asarray(y1_obs, dtype=float)
313
+ Y0 = np.asarray(Y0, dtype=float)
314
+ m = y1_obs.shape[0]
315
+ post_mask = np.zeros(m, dtype=bool)
316
+ post_mask[post_idx] = True
317
+ n_perms = int(perms.shape[0])
318
+
319
+ # Un-absorbed point estimate + grid center from a PRE-ONLY proxy fit: predict the
320
+ # post slot from a proxy fit on the other (pre) slots only. The θ=0 full-series fit
321
+ # would let the proxy soak the effect into the weights, biasing the naive residual
322
+ # toward 0 and mis-centring the auto-grid (CWZ inverts the test precisely to avoid
323
+ # relying on that naive residual). ``spread`` is the pre-fit residual noise scale.
324
+ pre_sel = ~post_mask
325
+ w_pre, resid_pre, _ = _cwz_proxy_fit(
326
+ y1_obs[pre_sel], Y0[pre_sel], max_iter=max_iter, min_decrease=min_decrease
327
+ )
328
+ point_est = float(y1_obs[post_idx] - Y0[post_idx] @ w_pre)
329
+ spread = float(np.std(resid_pre)) if resid_pre.size else 1.0
330
+
331
+ if alpha < 1.0 / n_perms:
332
+ # Granularity-unbounded: every null is accepted (p >= 1/|Π| > alpha).
333
+ return {
334
+ "lower": -float("inf"),
335
+ "upper": float("inf"),
336
+ "status": "unbounded",
337
+ "contiguous": True,
338
+ "point_estimate": point_est,
339
+ "n_in_set": 0,
340
+ "n_nonconverged": 0,
341
+ "grid": [],
342
+ "n_perms": n_perms,
343
+ }
344
+
345
+ if grid is None:
346
+ grid = _auto_grid(point_est, spread, n_grid)
347
+ else:
348
+ grid = np.asarray(grid, dtype=float)
349
+
350
+ pvals = np.empty(grid.shape[0], dtype=float)
351
+ converged = np.empty(grid.shape[0], dtype=bool)
352
+ w_prev = w_pre
353
+ for i, theta in enumerate(grid):
354
+ y1n = y1_obs.copy()
355
+ y1n[post_idx] = y1n[post_idx] - float(theta)
356
+ w, resid, conv = _cwz_proxy_fit(
357
+ y1n, Y0, max_iter=max_iter, min_decrease=min_decrease, init_weights=w_prev
358
+ )
359
+ w_prev = w
360
+ p, _, _ = _cwz_pvalue(resid, post_mask, perms, 1)
361
+ pvals[i] = p
362
+ converged[i] = conv
363
+
364
+ out = _assemble_single_ci(grid, pvals, converged, alpha, point_est)
365
+ in_set = out.pop("_in_set")
366
+ out["grid"] = list(zip(grid.tolist(), pvals.tolist(), in_set.tolist(), converged.tolist()))
367
+ out["n_perms"] = n_perms
368
+ return out
369
+
370
+
371
+ # =============================================================================
372
+ # Block collapse for the average-effect test (CWZ Appendix A.1)
373
+ # =============================================================================
374
+
375
+
376
+ def _block_collapse(
377
+ y1: np.ndarray, Y0: np.ndarray, n_pre: int, n_post: int
378
+ ) -> Tuple[np.ndarray, np.ndarray, int]:
379
+ """Collapse the calendar-ordered panel into non-overlapping ``T*``-blocks (A.1).
380
+
381
+ The post block is the last ``n_post`` periods; pre blocks tile the pre-period
382
+ backwards from ``T0``, so the earliest ``n_pre % n_post`` pre-periods are
383
+ DROPPED to make the pre-block count integral (the paper assumes ``T/T*``
384
+ integer). Each block is the per-unit average over its ``n_post`` periods.
385
+ Returns ``(y1_blocks, Y0_blocks, n_dropped)`` with the post block last. Requires
386
+ at least one full pre-block after the drop (``n_pre ≥ n_post``).
387
+ """
388
+ y1 = np.asarray(y1, dtype=float)
389
+ Y0 = np.asarray(Y0, dtype=float)
390
+ t_star = n_post
391
+ drop = n_pre % t_star
392
+ n_pre_blocks = (n_pre - drop) // t_star
393
+ pre_y1 = y1[drop:n_pre].reshape(n_pre_blocks, t_star).mean(axis=1)
394
+ pre_Y0 = Y0[drop:n_pre].reshape(n_pre_blocks, t_star, Y0.shape[1]).mean(axis=1)
395
+ post_y1 = y1[n_pre:].mean()
396
+ post_Y0 = Y0[n_pre:].mean(axis=0)
397
+ y1_blocks = np.concatenate([pre_y1, [post_y1]])
398
+ Y0_blocks = np.vstack([pre_Y0, post_Y0[None, :]])
399
+ return y1_blocks, Y0_blocks, drop