diff-diff 3.11.1__tar.gz → 3.12.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (129) hide show
  1. {diff_diff-3.11.1 → diff_diff-3.12.0}/PKG-INFO +6 -4
  2. {diff_diff-3.11.1 → diff_diff-3.12.0}/README.md +4 -2
  3. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/__init__.py +7 -1
  4. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_backend.py +5 -5
  5. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_crossfit.py +327 -130
  6. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_dr_scores.py +139 -1
  7. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_reporting_helpers.py +19 -3
  8. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/business_report.py +85 -14
  9. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/chaisemartin_dhaultfoeuille_results.py +11 -5
  10. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/changes_in_changes_results.py +2 -2
  11. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/continuous_did_results.py +19 -4
  12. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/diagnostic_report.py +40 -7
  13. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/diagnostics.py +2 -2
  14. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/dml_did.py +998 -47
  15. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/dml_did_results.py +111 -1
  16. diff_diff-3.12.0/diff_diff/duration_did.py +1313 -0
  17. diff_diff-3.12.0/diff_diff/duration_did_results.py +593 -0
  18. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/efficient_did_results.py +7 -2
  19. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/estimators.py +8 -0
  20. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/guides/llms-autonomous.txt +38 -7
  21. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/guides/llms-full.txt +64 -11
  22. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/guides/llms-practitioner.txt +29 -2
  23. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/guides/llms.txt +7 -3
  24. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/had.py +3 -3
  25. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/honest_did.py +2 -2
  26. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/imputation_results.py +7 -2
  27. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/linalg.py +77 -68
  28. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lpdid_results.py +2 -2
  29. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid.py +4 -23
  30. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid_results.py +2 -2
  31. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid_wild_bootstrap.py +2 -1
  32. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/practitioner.py +207 -3
  33. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/rdd.py +3 -3
  34. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/results.py +110 -19
  35. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/results_base.py +49 -6
  36. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/stacked_did_results.py +7 -2
  37. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/staggered_results.py +7 -2
  38. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/staggered_triple_diff_results.py +2 -2
  39. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/sun_abraham.py +2 -2
  40. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/synthetic_control.py +6 -3
  41. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/synthetic_control_results.py +25 -4
  42. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/synthetic_did.py +205 -21
  43. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/triple_diff.py +12 -5
  44. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/trop_results.py +19 -5
  45. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/two_stage_results.py +7 -2
  46. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/utils.py +2 -1
  47. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_continuous.py +119 -29
  48. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_event_study.py +2 -2
  49. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/wooldridge.py +38 -0
  50. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/wooldridge_results.py +14 -4
  51. {diff_diff-3.11.1 → diff_diff-3.12.0}/pyproject.toml +2 -2
  52. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/Cargo.lock +34 -34
  53. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/Cargo.toml +1 -1
  54. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/lib.rs +7 -0
  55. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/linalg.rs +11 -11
  56. {diff_diff-3.11.1 → diff_diff-3.12.0}/LICENSE +0 -0
  57. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_base.py +0 -0
  58. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_deprecation.py +0 -0
  59. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_guides_api.py +0 -0
  60. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_learners.py +0 -0
  61. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_nprobust_port.py +0 -0
  62. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_rdrobust_port.py +0 -0
  63. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/_staggered_triple_diff_engine.py +0 -0
  64. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/agent_workflow.py +0 -0
  65. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/aggregation.py +0 -0
  66. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/bacon.py +0 -0
  67. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/balancing.py +0 -0
  68. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/bootstrap_chunking.py +0 -0
  69. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/bootstrap_utils.py +0 -0
  70. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/chaisemartin_dhaultfoeuille.py +0 -0
  71. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +0 -0
  72. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/changes_in_changes.py +0 -0
  73. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/conformal.py +0 -0
  74. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/conley.py +0 -0
  75. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/continuous_did.py +0 -0
  76. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/continuous_did_aggregation.py +0 -0
  77. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/continuous_did_bspline.py +0 -0
  78. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/datasets.py +0 -0
  79. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/efficient_did.py +0 -0
  80. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/efficient_did_aggregation.py +0 -0
  81. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/efficient_did_bootstrap.py +0 -0
  82. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/efficient_did_covariates.py +0 -0
  83. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/efficient_did_weights.py +0 -0
  84. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/guides/__init__.py +0 -0
  85. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/had_pretests.py +0 -0
  86. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/imputation.py +0 -0
  87. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/imputation_aggregation.py +0 -0
  88. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/imputation_bootstrap.py +0 -0
  89. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/local_linear.py +0 -0
  90. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lpdid.py +0 -0
  91. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid_randomization.py +0 -0
  92. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid_sensitivity.py +0 -0
  93. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid_staggered.py +0 -0
  94. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/lwdid_visualization.py +0 -0
  95. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/mmm.py +0 -0
  96. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/power.py +0 -0
  97. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/prep.py +0 -0
  98. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/prep_dgp.py +0 -0
  99. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/pretrends.py +0 -0
  100. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/profile.py +0 -0
  101. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/rddensity.py +0 -0
  102. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/rdplot.py +0 -0
  103. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/spillover.py +0 -0
  104. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/stacked_did.py +0 -0
  105. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/staggered.py +0 -0
  106. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/staggered_aggregation.py +0 -0
  107. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/staggered_bootstrap.py +0 -0
  108. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/staggered_triple_diff.py +0 -0
  109. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/survey.py +0 -0
  110. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/trop.py +0 -0
  111. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/trop_global.py +0 -0
  112. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/trop_local.py +0 -0
  113. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/twfe.py +0 -0
  114. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/two_stage.py +0 -0
  115. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/two_stage_aggregation.py +0 -0
  116. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/two_stage_bootstrap.py +0 -0
  117. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/__init__.py +0 -0
  118. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_common.py +0 -0
  119. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_diagnostic.py +0 -0
  120. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_power.py +0 -0
  121. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_staggered.py +0 -0
  122. {diff_diff-3.11.1 → diff_diff-3.12.0}/diff_diff/visualization/_synthetic.py +0 -0
  123. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/build.rs +0 -0
  124. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/alloc_profile.rs +0 -0
  125. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/batched_solve.rs +0 -0
  126. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/bootstrap.rs +0 -0
  127. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/demean.rs +0 -0
  128. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/trop.rs +0 -0
  129. {diff_diff-3.11.1 → diff_diff-3.12.0}/rust/src/weights.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diff-diff
3
- Version: 3.11.1
3
+ Version: 3.12.0
4
4
  Classifier: Development Status :: 5 - Production/Stable
5
5
  Classifier: Intended Audience :: Science/Research
6
6
  Classifier: Operating System :: OS Independent
@@ -21,7 +21,7 @@ Requires-Dist: pytest>=7.0 ; extra == 'dev'
21
21
  Requires-Dist: pytest-xdist>=3.0 ; extra == 'dev'
22
22
  Requires-Dist: pytest-cov>=4.0 ; extra == 'dev'
23
23
  Requires-Dist: black==26.5.1 ; extra == 'dev'
24
- Requires-Dist: ruff==0.16.4 ; extra == 'dev'
24
+ Requires-Dist: ruff==0.16.5 ; extra == 'dev'
25
25
  Requires-Dist: mypy==2.3.1 ; extra == 'dev'
26
26
  Requires-Dist: maturin>=1.4,<2.0 ; extra == 'dev'
27
27
  Requires-Dist: matplotlib>=3.5 ; extra == 'dev'
@@ -134,7 +134,7 @@ Measuring campaign lift? Evaluating a product launch? Rolling out a policy in wa
134
134
  - [Brand awareness survey tutorial](https://diff-diff.readthedocs.io/en/stable/tutorials/17_brand_awareness_survey.html) - full example with complex survey design, brand funnel analysis, and staggered rollouts
135
135
  - Have BRFSS/ACS/CPS individual records? Use [`aggregate_survey()`](https://diff-diff.readthedocs.io/en/stable/api/prep.html) to roll respondent-level microdata into a geographic-period panel with inverse-variance precision weights for second-stage DiD
136
136
 
137
- `BusinessReport` and `DiagnosticReport` are experimental preview classes that produce plain-English output and a structured `to_dict()` schema from any fitted result - wording and schema will evolve. See [docs/methodology/REPORTING.md](https://github.com/igerber/diff-diff/blob/main/docs/methodology/REPORTING.md) for usage and stability notes.
137
+ `BusinessReport` and `DiagnosticReport` are experimental preview classes that produce plain-English output and a structured `to_dict()` schema from any fitted estimator result except the two rejected by type (`EventStudyResults` surfaces and `DurationDiDResults`) - wording and schema will evolve. See [docs/methodology/REPORTING.md](https://github.com/igerber/diff-diff/blob/main/docs/methodology/REPORTING.md) for usage and stability notes.
138
138
 
139
139
  ## Practitioner Workflow (Baker et al. 2025)
140
140
 
@@ -176,11 +176,13 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
176
176
  - [LPDiD](https://diff-diff.readthedocs.io/en/stable/api/lpdid.html) - Dube, Girardi, Jorda & Taylor (2025) Local Projections DiD: per-horizon long-difference event study on clean controls (no negative weighting), variance- or equally-weighted ATT, for absorbing or non-absorbing (reversible) treatment
177
177
  - [ChangesInChanges](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html) - Athey & Imbens (2006) nonlinear/distributional DiD for the 2x2 design: full counterfactual distribution and quantile treatment effects via CDF transformation, plus the QDiD comparison estimator via `method="qdid"`; bootstrap inference; R qte parity. Alias `CiC`
178
178
  - [LWDiD](https://diff-diff.readthedocs.io/en/stable/api/lwdid.html) - Lee & Wooldridge (2025, 2026) rolling-transformation DiD: unit-specific demean/detrend converts panel to cross-section, staggered adoption, `estimation_method` in `reg`/`ipw`/`dr`/`psm` (the papers' RA/IPW/IPWRA plus propensity-score matching), exact small-N inference on the classical collapsed regression
179
- - [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Chang (2020) double/debiased machine learning DiD: staggered ATT(g,t) with cross-fitted ML nuisance learners (DML2) and Neyman-orthogonal scores, for flexible/high-dimensional covariate adjustment under conditional parallel trends; panel or declared repeated cross sections (`panel=False`); survey/cluster support on both lanes
179
+ - [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Chang (2020) double/debiased machine learning DiD: staggered ATT(g,t) with cross-fitted ML nuisance learners (DML2) and Neyman-orthogonal scores, for flexible/high-dimensional covariate adjustment under conditional parallel trends; panel or declared repeated cross sections (`panel=False`); survey/cluster support on both lanes (bad-control lane: panel only, `cluster=` only); Caetano, Callaway, Payne & Sant'Anna (2026) bad-control score via `fit(bad_control=, bad_control_covariates=)`
180
+ - [DurationDiD](https://diff-diff.readthedocs.io/en/stable/api/duration_did.html) - Deaner & Ku (2026) causal duration DiD for a binary absorbing outcome (spell ended) in a two-group common-timing design: restricts the groups' untreated hazards (`method="cd"` additive gap or `method="ph"` ratio) instead of outcome levels, imputes the treated counterfactual survival, reports the per-date absorption ATT with whole-individual bootstrap pointwise and simultaneous bands plus a fixed-anchor pre-treatment specification test
180
181
  - [BaconDecomposition](https://diff-diff.readthedocs.io/en/stable/api/bacon.html) - Goodman-Bacon (2021) decomposition for diagnosing TWFE bias in staggered settings
181
182
 
182
183
  ## Diagnostics & Sensitivity
183
184
 
185
+ - [Bad-control pre-test](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Caetano, Callaway, Payne & Sant'Anna (2026) Remark 6 `ATT_X(g,t)`: the effect of treatment on a suspected bad control itself, per (g,t) cell (`DMLDiDResults.bad_control_summary()`; analytical SE; pre-period rows pre-test the identifying assumptions MP-5 / MP-8 and should be zero, post-period rows are the Condition-2 check that treatment affects the covariate)
184
186
  - [RD Plots](https://diff-diff.readthedocs.io/en/stable/api/regression_discontinuity.html) - Calonico, Cattaneo & Titiunik (2015) optimal data-driven RD plots (`RDPlot`): all 8 rdrobust `binselect` bin selectors, implied-scale/WIMSE-weight reporting, optional matplotlib rendering
185
187
  - [Manipulation Testing](https://diff-diff.readthedocs.io/en/stable/api/regression_discontinuity.html) - Cattaneo, Jansson & Ma (2020) density-discontinuity test (`RDDensityTest`): rddensity 3.0 parity, robust bias-corrected inference, unrestricted/restricted models, mass-point adjustment
186
188
  - [Parallel Trends Testing](https://diff-diff.readthedocs.io/en/stable/api/diagnostics.html) - simple and Wasserstein-robust parallel trends tests, equivalence testing (TOST)
@@ -79,7 +79,7 @@ Measuring campaign lift? Evaluating a product launch? Rolling out a policy in wa
79
79
  - [Brand awareness survey tutorial](https://diff-diff.readthedocs.io/en/stable/tutorials/17_brand_awareness_survey.html) - full example with complex survey design, brand funnel analysis, and staggered rollouts
80
80
  - Have BRFSS/ACS/CPS individual records? Use [`aggregate_survey()`](https://diff-diff.readthedocs.io/en/stable/api/prep.html) to roll respondent-level microdata into a geographic-period panel with inverse-variance precision weights for second-stage DiD
81
81
 
82
- `BusinessReport` and `DiagnosticReport` are experimental preview classes that produce plain-English output and a structured `to_dict()` schema from any fitted result - wording and schema will evolve. See [docs/methodology/REPORTING.md](https://github.com/igerber/diff-diff/blob/main/docs/methodology/REPORTING.md) for usage and stability notes.
82
+ `BusinessReport` and `DiagnosticReport` are experimental preview classes that produce plain-English output and a structured `to_dict()` schema from any fitted estimator result except the two rejected by type (`EventStudyResults` surfaces and `DurationDiDResults`) - wording and schema will evolve. See [docs/methodology/REPORTING.md](https://github.com/igerber/diff-diff/blob/main/docs/methodology/REPORTING.md) for usage and stability notes.
83
83
 
84
84
  ## Practitioner Workflow (Baker et al. 2025)
85
85
 
@@ -121,11 +121,13 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
121
121
  - [LPDiD](https://diff-diff.readthedocs.io/en/stable/api/lpdid.html) - Dube, Girardi, Jorda & Taylor (2025) Local Projections DiD: per-horizon long-difference event study on clean controls (no negative weighting), variance- or equally-weighted ATT, for absorbing or non-absorbing (reversible) treatment
122
122
  - [ChangesInChanges](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html) - Athey & Imbens (2006) nonlinear/distributional DiD for the 2x2 design: full counterfactual distribution and quantile treatment effects via CDF transformation, plus the QDiD comparison estimator via `method="qdid"`; bootstrap inference; R qte parity. Alias `CiC`
123
123
  - [LWDiD](https://diff-diff.readthedocs.io/en/stable/api/lwdid.html) - Lee & Wooldridge (2025, 2026) rolling-transformation DiD: unit-specific demean/detrend converts panel to cross-section, staggered adoption, `estimation_method` in `reg`/`ipw`/`dr`/`psm` (the papers' RA/IPW/IPWRA plus propensity-score matching), exact small-N inference on the classical collapsed regression
124
- - [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Chang (2020) double/debiased machine learning DiD: staggered ATT(g,t) with cross-fitted ML nuisance learners (DML2) and Neyman-orthogonal scores, for flexible/high-dimensional covariate adjustment under conditional parallel trends; panel or declared repeated cross sections (`panel=False`); survey/cluster support on both lanes
124
+ - [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Chang (2020) double/debiased machine learning DiD: staggered ATT(g,t) with cross-fitted ML nuisance learners (DML2) and Neyman-orthogonal scores, for flexible/high-dimensional covariate adjustment under conditional parallel trends; panel or declared repeated cross sections (`panel=False`); survey/cluster support on both lanes (bad-control lane: panel only, `cluster=` only); Caetano, Callaway, Payne & Sant'Anna (2026) bad-control score via `fit(bad_control=, bad_control_covariates=)`
125
+ - [DurationDiD](https://diff-diff.readthedocs.io/en/stable/api/duration_did.html) - Deaner & Ku (2026) causal duration DiD for a binary absorbing outcome (spell ended) in a two-group common-timing design: restricts the groups' untreated hazards (`method="cd"` additive gap or `method="ph"` ratio) instead of outcome levels, imputes the treated counterfactual survival, reports the per-date absorption ATT with whole-individual bootstrap pointwise and simultaneous bands plus a fixed-anchor pre-treatment specification test
125
126
  - [BaconDecomposition](https://diff-diff.readthedocs.io/en/stable/api/bacon.html) - Goodman-Bacon (2021) decomposition for diagnosing TWFE bias in staggered settings
126
127
 
127
128
  ## Diagnostics & Sensitivity
128
129
 
130
+ - [Bad-control pre-test](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Caetano, Callaway, Payne & Sant'Anna (2026) Remark 6 `ATT_X(g,t)`: the effect of treatment on a suspected bad control itself, per (g,t) cell (`DMLDiDResults.bad_control_summary()`; analytical SE; pre-period rows pre-test the identifying assumptions MP-5 / MP-8 and should be zero, post-period rows are the Condition-2 check that treatment affects the covariate)
129
131
  - [RD Plots](https://diff-diff.readthedocs.io/en/stable/api/regression_discontinuity.html) - Calonico, Cattaneo & Titiunik (2015) optimal data-driven RD plots (`RDPlot`): all 8 rdrobust `binselect` bin selectors, implied-scale/WIMSE-weight reporting, optional matplotlib rendering
130
132
  - [Manipulation Testing](https://diff-diff.readthedocs.io/en/stable/api/regression_discontinuity.html) - Cattaneo, Jansson & Ma (2020) density-discontinuity test (`RDDensityTest`): rddensity 3.0 parity, robust bias-corrected inference, unrestricted/restricted models, mass-point adjustment
131
133
  - [Parallel Trends Testing](https://diff-diff.readthedocs.io/en/stable/api/diagnostics.html) - simple and Wasserstein-robust parallel trends tests, equivalence testing (TOST)
@@ -98,6 +98,8 @@ from diff_diff.diagnostics import (
98
98
  )
99
99
  from diff_diff.dml_did import DMLDiD
100
100
  from diff_diff.dml_did_results import DMLDiDResults
101
+ from diff_diff.duration_did import DurationDiD
102
+ from diff_diff.duration_did_results import DurationDiDPretestResults, DurationDiDResults
101
103
  from diff_diff.efficient_did import (
102
104
  EDiDBootstrapResults,
103
105
  EfficientDiD,
@@ -371,7 +373,7 @@ def __getattr__(name: str) -> _Any:
371
373
  raise AttributeError(f"module 'diff_diff' has no attribute {name!r}")
372
374
 
373
375
 
374
- __version__ = "3.11.1"
376
+ __version__ = "3.12.0"
375
377
  __all__ = [
376
378
  # Estimators
377
379
  "DifferenceInDifferences",
@@ -471,6 +473,10 @@ __all__ = [
471
473
  "DMLDiD",
472
474
  "DMLDiDResults",
473
475
  "SieveLearner",
476
+ # DurationDiD (Deaner & Ku 2026 causal duration DiD)
477
+ "DurationDiD",
478
+ "DurationDiDResults",
479
+ "DurationDiDPretestResults",
474
480
  # Visualization
475
481
  "plot_bacon",
476
482
  "plot_event_study",
@@ -114,13 +114,13 @@ try:
114
114
  except ImportError:
115
115
  _rust_batched_ridge_chol_solve = None
116
116
 
117
- # HC2 (leverage-corrected) robust vcov: imported independently for the same
118
- # mixed-version reason as demean_map (a stale extension missing only this
119
- # newer symbol degrades HC2 to the NumPy path without disabling the older
120
- # Rust accelerations).
117
+ # HC2 requires the v2 fail-closed leverage contract. An older extension can
118
+ # export the original symbol yet return finite covariance at unit leverage,
119
+ # so symbol presence alone is insufficient. Import v2 independently: legacy
120
+ # extensions use NumPy HC2 while retaining every other Rust acceleration.
121
121
  try:
122
122
  from diff_diff._rust_backend import (
123
- compute_robust_vcov_hc2 as _rust_compute_robust_vcov_hc2,
123
+ compute_robust_vcov_hc2_v2 as _rust_compute_robust_vcov_hc2,
124
124
  )
125
125
  except ImportError:
126
126
  _rust_compute_robust_vcov_hc2 = None
@@ -35,7 +35,7 @@ import copy
35
35
  import pickle
36
36
  import warnings
37
37
  from dataclasses import dataclass, field
38
- from typing import Any, Dict, Iterator, Literal, Optional, Tuple, cast, overload
38
+ from typing import Any, Dict, Iterator, Literal, NamedTuple, Optional, Tuple, cast, overload
39
39
 
40
40
  import numpy as np
41
41
  import pandas as pd
@@ -53,12 +53,14 @@ __all__ = [
53
53
  "CrossFitResult",
54
54
  "assign_folds",
55
55
  "cross_fit_predict",
56
+ "FoldFit",
57
+ "iter_fold_fits",
56
58
  ]
57
59
 
58
60
  _LOG_LOSS_CLIP = 1e-15
59
61
 
60
62
 
61
- def _fresh_learner(learner: Any) -> Any:
63
+ def _fresh_learner(learner: Any, *, stacklevel: int) -> Any:
62
64
  """Per-fold learner isolation: a deep copy of the (never-fit) template.
63
65
 
64
66
  ``copy.deepcopy`` of the user's template gives every fold a fully
@@ -70,7 +72,10 @@ def _fresh_learner(learner: Any) -> Any:
70
72
  a container parameter). The template itself is never fit. A copy FAILURE
71
73
  is never silent: the instance is reused with a loud ``UserWarning`` naming
72
74
  the learner and the fit-reset assumption now being relied on
73
- (no-silent-failures rule).
75
+ (no-silent-failures rule). ``stacklevel`` is supplied by the caller so the
76
+ warning is attributed to the entry point's own caller (``cross_fit_predict``
77
+ passes the level that lands on the user's call site; internal consumers
78
+ of ``iter_fold_fits`` pass the level that lands on themselves).
74
79
  """
75
80
  try:
76
81
  return copy.deepcopy(learner)
@@ -80,14 +85,14 @@ def _fresh_learner(learner: Any) -> Any:
80
85
  # and this warning lands in notebook/CI logs (the same boundary as
81
86
  # DMLDiD's persisted-diagnostics sanitization).
82
87
  warnings.warn(
83
- f"cross_fit_predict: could not deep-copy the "
88
+ f"_crossfit: could not deep-copy the "
84
89
  f"{type(learner).__name__} template for this fold "
85
90
  f"({type(exc).__name__}); "
86
91
  "REUSING the same instance and relying on its fit-reset behavior. "
87
92
  "A warm-start/stateful learner in this situation can leak data "
88
93
  "across folds.",
89
94
  UserWarning,
90
- stacklevel=3,
95
+ stacklevel=stacklevel,
91
96
  )
92
97
  return learner
93
98
 
@@ -404,55 +409,45 @@ def assign_folds(
404
409
  )
405
410
 
406
411
 
407
- @overload
408
- def cross_fit_predict(
409
- learner: RegressorLearner,
410
- X: np.ndarray,
411
- y: np.ndarray,
412
- folds: FoldAssignment,
413
- *,
414
- fit_mask: Optional[np.ndarray] = ...,
415
- predict_method: Literal["predict"] = ...,
416
- sample_weight: Optional[np.ndarray] = ...,
417
- context_label: str = ...,
418
- ) -> CrossFitResult: ...
419
-
420
-
421
- @overload
422
- def cross_fit_predict(
423
- learner: ClassifierLearner,
424
- X: np.ndarray,
425
- y: np.ndarray,
426
- folds: FoldAssignment,
427
- *,
428
- fit_mask: Optional[np.ndarray] = ...,
429
- predict_method: Literal["predict_proba"],
430
- sample_weight: Optional[np.ndarray] = ...,
431
- context_label: str = ...,
432
- ) -> CrossFitResult: ...
412
+ class FoldFit(NamedTuple):
413
+ """One fold's fitted learner plus the index sets it was built from.
433
414
 
415
+ Yielded by :func:`iter_fold_fits`. ``X`` / ``y`` / ``sample_weight`` are the
416
+ validated, float64-coerced arrays (so consumers index them positionally
417
+ exactly as :func:`cross_fit_predict` does); ``learner`` is the per-fold
418
+ deep copy already fit on ``fit_idx``; ``train_idx`` is fold ``k``'s full
419
+ training complement and ``test_idx`` the held-out fold.
420
+ """
434
421
 
435
- def cross_fit_predict(
422
+ k: int
423
+ learner: Any
424
+ fit_idx: np.ndarray
425
+ train_idx: np.ndarray
426
+ test_idx: np.ndarray
427
+ n_fit: int
428
+ w_fit: Optional[np.ndarray]
429
+ kind: str
430
+ label: str
431
+ X: np.ndarray
432
+ y: np.ndarray
433
+ sample_weight: Optional[np.ndarray]
434
+
435
+
436
+ def _prepare_cross_fit_inputs(
436
437
  learner: object,
437
438
  X: np.ndarray,
438
439
  y: np.ndarray,
439
440
  folds: FoldAssignment,
440
441
  *,
441
- fit_mask: Optional[np.ndarray] = None,
442
- predict_method: str = "predict",
443
- sample_weight: Optional[np.ndarray] = None,
444
- context_label: str = "",
445
- ) -> CrossFitResult:
446
- """Out-of-fold predictions for every unit.
442
+ fit_mask: Optional[np.ndarray],
443
+ predict_method: str,
444
+ sample_weight: Optional[np.ndarray],
445
+ ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], str]:
446
+ """Argument validation shared by every cross-fit entry point.
447
447
 
448
- Per fold k: fit the learner on ``train_mask(k) & fit_mask`` (e.g. the
449
- untreated units), predict ALL units in fold k. ``context_label``
450
- (e.g. the nuisance name or (g,t) cell) is prefixed into every
451
- ``DegenerateFoldError`` message to identify WHICH cross-fit failed.
448
+ Own-argument failures raise plain ``ValueError`` (the module contract).
452
449
  """
453
450
  n_units = folds.n_units
454
- label = f"{context_label}: " if context_label else ""
455
-
456
451
  if predict_method not in ("predict", "predict_proba"):
457
452
  raise ValueError(
458
453
  f"predict_method must be 'predict' or 'predict_proba', got {predict_method!r}"
@@ -491,20 +486,19 @@ def cross_fit_predict(
491
486
  )
492
487
  fit_mask_arr = raw_mask
493
488
 
489
+ sw: Optional[np.ndarray] = None
494
490
  if sample_weight is not None:
495
- sample_weight = np.asarray(sample_weight, dtype=np.float64)
496
- if sample_weight.ndim != 1:
491
+ sw = np.asarray(sample_weight, dtype=np.float64)
492
+ if sw.ndim != 1:
497
493
  raise ValueError(
498
- f"sample_weight must be 1-dimensional, got ndim={sample_weight.ndim} "
494
+ f"sample_weight must be 1-dimensional, got ndim={sw.ndim} "
499
495
  "(a column vector like (n, 1) is not accepted)"
500
496
  )
501
- if sample_weight.shape[0] != n_units:
502
- raise ValueError(
503
- f"sample_weight has length {sample_weight.shape[0]}, expected {n_units}"
504
- )
505
- if not np.isfinite(sample_weight).all():
497
+ if sw.shape[0] != n_units:
498
+ raise ValueError(f"sample_weight has length {sw.shape[0]}, expected {n_units}")
499
+ if not np.isfinite(sw).all():
506
500
  raise ValueError("sample_weight contains NaN or Inf values")
507
- if np.any(sample_weight < 0):
501
+ if np.any(sw < 0):
508
502
  raise ValueError("sample_weight must be non-negative")
509
503
 
510
504
  if predict_method == "predict_proba" and not np.all((y == 0.0) | (y == 1.0)):
@@ -512,91 +506,294 @@ def cross_fit_predict(
512
506
  "y must be strictly binary 0/1 for predict_method='predict_proba' "
513
507
  "(the logit solver silently saturates on other encodings)"
514
508
  )
509
+ return X, y, fit_mask_arr, sw, kind
515
510
 
516
- oof = np.empty(n_units, dtype=np.float64)
517
- fold_losses = np.empty(folds.n_folds, dtype=np.float64)
518
- n_fit_per_fold = np.empty(folds.n_folds, dtype=np.int64)
519
511
 
520
- for k, train_idx, test_idx in folds.iter_folds():
521
- fit_idx = train_idx[fit_mask_arr[train_idx]]
522
- n_fit = fit_idx.shape[0]
523
- n_fit_per_fold[k] = n_fit
524
- w_fit = None if sample_weight is None else sample_weight[fit_idx]
512
+ def _learner_error(
513
+ exc: ValueError, *, k: int, label: str, n_fit: int, w_fit: Optional[np.ndarray]
514
+ ) -> DegenerateFoldError:
515
+ return DegenerateFoldError(
516
+ f"{label}learner error in fold {k}: {exc}; the fold's fit subset "
517
+ f"has n={n_fit}"
518
+ + (f", n_pos_weight={int(np.sum(w_fit > 0))}" if w_fit is not None else "")
519
+ + ". Reduce n_folds, widen fit_mask, or check the fold's data."
520
+ )
525
521
 
526
- # (a) Universal cheap pre-checks -> DegenerateFoldError directly.
527
- if n_fit == 0:
528
- raise DegenerateFoldError(
529
- f"{label}fold {k}: the fit subset is empty (train size "
530
- f"{train_idx.shape[0]}, fit_mask keeps 0). Reduce n_folds, widen "
531
- "fit_mask, or check the stratify labels."
532
- )
533
- if w_fit is not None and not np.any(w_fit > 0):
522
+
523
+ def _fit_subset(
524
+ learner: object,
525
+ X: np.ndarray,
526
+ y: np.ndarray,
527
+ fit_idx: np.ndarray,
528
+ *,
529
+ n_train: int,
530
+ kind: str,
531
+ sample_weight: Optional[np.ndarray],
532
+ k: int,
533
+ label: str,
534
+ warn_stacklevel: int,
535
+ ) -> Tuple[Any, int, Optional[np.ndarray]]:
536
+ """Fit a fresh deep copy of ``learner`` on ``fit_idx`` (fold ``k``).
537
+
538
+ Universal degeneracy pre-checks raise ``DegenerateFoldError`` directly; a
539
+ learner ``ValueError`` is re-raised as ``DegenerateFoldError`` chained via
540
+ ``raise ... from exc``; a learner-raised ``DegenerateFoldError`` passes
541
+ through unwrapped. Returns ``(fitted_learner, n_fit, w_fit)``.
542
+ """
543
+ n_fit = int(fit_idx.shape[0])
544
+ w_fit = None if sample_weight is None else sample_weight[fit_idx]
545
+
546
+ # (a) Universal cheap pre-checks -> DegenerateFoldError directly.
547
+ if n_fit == 0:
548
+ raise DegenerateFoldError(
549
+ f"{label}fold {k}: the fit subset is empty (train size "
550
+ f"{n_train}, fit_mask keeps 0). Reduce n_folds, widen "
551
+ "fit_mask, or check the stratify labels."
552
+ )
553
+ if w_fit is not None and not np.any(w_fit > 0):
554
+ raise DegenerateFoldError(
555
+ f"{label}fold {k}: all {n_fit} fit rows have zero sample_weight. "
556
+ "Reduce n_folds or check the weights."
557
+ )
558
+ if kind == "classifier":
559
+ labels = y[fit_idx] if w_fit is None else y[fit_idx][w_fit > 0]
560
+ if np.unique(labels).shape[0] < 2:
534
561
  raise DegenerateFoldError(
535
- f"{label}fold {k}: all {n_fit} fit rows have zero sample_weight. "
536
- "Reduce n_folds or check the weights."
562
+ f"{label}fold {k}: the fit subset has a single "
563
+ f"{'positive-weight ' if w_fit is not None else ''}class "
564
+ f"(n_fit={n_fit}). A classifier needs both classes in every "
565
+ "fold's complement; reduce n_folds or stratify by the label."
537
566
  )
538
- if predict_method == "predict_proba":
539
- labels = y[fit_idx] if w_fit is None else y[fit_idx][w_fit > 0]
540
- if np.unique(labels).shape[0] < 2:
541
- raise DegenerateFoldError(
542
- f"{label}fold {k}: the fit subset has a single "
543
- f"{'positive-weight ' if w_fit is not None else ''}class "
544
- f"(n_fit={n_fit}). A classifier needs both classes in every "
545
- "fold's complement; reduce n_folds or stratify by the label."
546
- )
547
567
 
548
- # (b) Learner errors during the fold -> DegenerateFoldError, chained.
549
- try:
550
- fold_learner = _fresh_learner(learner)
551
- # Unweighted path calls fit(X, y) WITHOUT the keyword: the
552
- # advertised duck-typed contract is fit/predict(_proba), so a
553
- # learner whose fit signature is only (X, y) must work when no
554
- # weights are in play. sample_weight= is passed only on
555
- # genuinely weighted paths, where an unsupported signature
556
- # raises TypeError — a caller protocol violation that PROPAGATES
557
- # (the DegenerateFoldError wrapper below catches ValueError
558
- # only; fold-data degeneracy, not signature bugs).
559
- fit_kwargs = {} if w_fit is None else {"sample_weight": w_fit}
560
- if kind == "regressor":
561
- reg = cast(RegressorLearner, fold_learner)
562
- reg.fit(X[fit_idx], y[fit_idx], **fit_kwargs)
563
- raw_pred = reg.predict(X[test_idx])
564
- else:
565
- clf = cast(ClassifierLearner, fold_learner)
566
- clf.fit(X[fit_idx], y[fit_idx], **fit_kwargs)
567
- raw_pred = clf.predict_proba(X[test_idx])
568
- pred = _validate_predictions(
569
- raw_pred,
570
- test_idx.shape[0],
571
- kind=kind,
572
- context=f"{label}fold {k}",
573
- classes=(getattr(fold_learner, "classes_", None) if kind == "classifier" else None),
574
- )
575
- except DegenerateFoldError:
576
- raise
577
- except ValueError as exc:
578
- raise DegenerateFoldError(
579
- f"{label}learner error in fold {k}: {exc}; the fold's fit subset "
580
- f"has n={n_fit}"
581
- + (f", n_pos_weight={int(np.sum(w_fit > 0))}" if w_fit is not None else "")
582
- + ". Reduce n_folds, widen fit_mask, or check the fold's data."
583
- ) from exc
584
-
585
- oof[test_idx] = pred
586
-
587
- # Out-of-fold loss (diagnostic; NaN sentinel for zero-weight folds).
588
- w_test = None if sample_weight is None else sample_weight[test_idx]
589
- if predict_method == "predict":
590
- errs = (y[test_idx] - pred) ** 2
568
+ # (b) Learner errors during the fold -> DegenerateFoldError, chained.
569
+ try:
570
+ fold_learner = _fresh_learner(learner, stacklevel=warn_stacklevel)
571
+ # Unweighted path calls fit(X, y) WITHOUT the keyword: the
572
+ # advertised duck-typed contract is fit/predict(_proba), so a
573
+ # learner whose fit signature is only (X, y) must work when no
574
+ # weights are in play. sample_weight= is passed only on
575
+ # genuinely weighted paths, where an unsupported signature
576
+ # raises TypeError — a caller protocol violation that PROPAGATES
577
+ # (the DegenerateFoldError wrapper below catches ValueError
578
+ # only; fold-data degeneracy, not signature bugs).
579
+ fit_kwargs = {} if w_fit is None else {"sample_weight": w_fit}
580
+ if kind == "regressor":
581
+ cast(RegressorLearner, fold_learner).fit(X[fit_idx], y[fit_idx], **fit_kwargs)
591
582
  else:
592
- p_clip = np.clip(pred, _LOG_LOSS_CLIP, 1.0 - _LOG_LOSS_CLIP)
593
- errs = -(y[test_idx] * np.log(p_clip) + (1.0 - y[test_idx]) * np.log(1.0 - p_clip))
594
- if w_test is None:
595
- fold_losses[k] = float(np.mean(errs))
596
- elif np.sum(w_test) > 0:
597
- fold_losses[k] = float(np.sum(w_test * errs) / np.sum(w_test))
583
+ cast(ClassifierLearner, fold_learner).fit(X[fit_idx], y[fit_idx], **fit_kwargs)
584
+ except DegenerateFoldError:
585
+ raise
586
+ except ValueError as exc:
587
+ raise _learner_error(exc, k=k, label=label, n_fit=n_fit, w_fit=w_fit) from exc
588
+ return fold_learner, n_fit, w_fit
589
+
590
+
591
+ def _predict_subset(
592
+ fold_learner: Any,
593
+ X_rows: np.ndarray,
594
+ *,
595
+ kind: str,
596
+ k: int,
597
+ label: str,
598
+ n_fit: int,
599
+ w_fit: Optional[np.ndarray],
600
+ ) -> np.ndarray:
601
+ """Predict ``X_rows`` with a fitted fold learner and validate the output.
602
+
603
+ Same exception contract as :func:`_fit_subset` (chained
604
+ ``DegenerateFoldError`` on a learner ``ValueError``; pass-through of a
605
+ learner-raised ``DegenerateFoldError``).
606
+ """
607
+ try:
608
+ if kind == "regressor":
609
+ raw_pred = cast(RegressorLearner, fold_learner).predict(X_rows)
598
610
  else:
599
- fold_losses[k] = np.nan
611
+ raw_pred = cast(ClassifierLearner, fold_learner).predict_proba(X_rows)
612
+ return _validate_predictions(
613
+ raw_pred,
614
+ X_rows.shape[0],
615
+ kind=kind,
616
+ context=f"{label}fold {k}",
617
+ classes=(getattr(fold_learner, "classes_", None) if kind == "classifier" else None),
618
+ )
619
+ except DegenerateFoldError:
620
+ raise
621
+ except ValueError as exc:
622
+ raise _learner_error(exc, k=k, label=label, n_fit=n_fit, w_fit=w_fit) from exc
623
+
624
+
625
+ def _fold_loss(
626
+ kind: str, y_test: np.ndarray, pred: np.ndarray, w_test: Optional[np.ndarray]
627
+ ) -> float:
628
+ """Out-of-fold MSE (regressor) or log-loss (classifier); NaN for a zero-weight fold."""
629
+ if kind == "regressor":
630
+ errs = (y_test - pred) ** 2
631
+ else:
632
+ p_clip = np.clip(pred, _LOG_LOSS_CLIP, 1.0 - _LOG_LOSS_CLIP)
633
+ errs = -(y_test * np.log(p_clip) + (1.0 - y_test) * np.log(1.0 - p_clip))
634
+ if w_test is None:
635
+ return float(np.mean(errs))
636
+ if np.sum(w_test) > 0:
637
+ return float(np.sum(w_test * errs) / np.sum(w_test))
638
+ return float(np.nan)
639
+
640
+
641
+ def _iter_fold_fits_impl(
642
+ learner: object,
643
+ X: np.ndarray,
644
+ y: np.ndarray,
645
+ folds: FoldAssignment,
646
+ fit_mask_arr: np.ndarray,
647
+ sample_weight: Optional[np.ndarray],
648
+ kind: str,
649
+ label: str,
650
+ warn_stacklevel: int,
651
+ ) -> Iterator[FoldFit]:
652
+ for k, train_idx, test_idx in folds.iter_folds():
653
+ fit_idx = train_idx[fit_mask_arr[train_idx]]
654
+ fitted, n_fit, w_fit = _fit_subset(
655
+ learner,
656
+ X,
657
+ y,
658
+ fit_idx,
659
+ n_train=int(train_idx.shape[0]),
660
+ kind=kind,
661
+ sample_weight=sample_weight,
662
+ k=k,
663
+ label=label,
664
+ warn_stacklevel=warn_stacklevel,
665
+ )
666
+ yield FoldFit(
667
+ k=k,
668
+ learner=fitted,
669
+ fit_idx=fit_idx,
670
+ train_idx=train_idx,
671
+ test_idx=test_idx,
672
+ n_fit=n_fit,
673
+ w_fit=w_fit,
674
+ kind=kind,
675
+ label=label,
676
+ X=X,
677
+ y=y,
678
+ sample_weight=sample_weight,
679
+ )
680
+
681
+
682
+ def iter_fold_fits(
683
+ learner: object,
684
+ X: np.ndarray,
685
+ y: np.ndarray,
686
+ folds: FoldAssignment,
687
+ *,
688
+ fit_mask: Optional[np.ndarray] = None,
689
+ predict_method: str = "predict",
690
+ sample_weight: Optional[np.ndarray] = None,
691
+ context_label: str = "",
692
+ warn_stacklevel: int = 4,
693
+ ) -> Iterator[FoldFit]:
694
+ """Per-fold fitted learners (the building block under ``cross_fit_predict``).
695
+
696
+ Validation runs EAGERLY at the call (not at the first ``__next__``), then
697
+ the returned iterator yields one :class:`FoldFit` per fold in ascending
698
+ ``k``, each holding a fresh deep copy of ``learner`` fit on
699
+ ``train_mask(k) & fit_mask``. Predictions are the caller's job via
700
+ ``_predict_subset``; this is what lets a consumer fit a NESTED nuisance on
701
+ fold ``k``'s own training units. ``warn_stacklevel`` is the frame the
702
+ deep-copy-failure warning is attributed to (default: the frame advancing
703
+ the iterator).
704
+ """
705
+ X64, y64, fit_mask_arr, sw, kind = _prepare_cross_fit_inputs(
706
+ learner,
707
+ X,
708
+ y,
709
+ folds,
710
+ fit_mask=fit_mask,
711
+ predict_method=predict_method,
712
+ sample_weight=sample_weight,
713
+ )
714
+ label = f"{context_label}: " if context_label else ""
715
+ return _iter_fold_fits_impl(
716
+ learner, X64, y64, folds, fit_mask_arr, sw, kind, label, warn_stacklevel
717
+ )
718
+
719
+
720
+ @overload
721
+ def cross_fit_predict(
722
+ learner: RegressorLearner,
723
+ X: np.ndarray,
724
+ y: np.ndarray,
725
+ folds: FoldAssignment,
726
+ *,
727
+ fit_mask: Optional[np.ndarray] = ...,
728
+ predict_method: Literal["predict"] = ...,
729
+ sample_weight: Optional[np.ndarray] = ...,
730
+ context_label: str = ...,
731
+ ) -> CrossFitResult: ...
732
+
733
+
734
+ @overload
735
+ def cross_fit_predict(
736
+ learner: ClassifierLearner,
737
+ X: np.ndarray,
738
+ y: np.ndarray,
739
+ folds: FoldAssignment,
740
+ *,
741
+ fit_mask: Optional[np.ndarray] = ...,
742
+ predict_method: Literal["predict_proba"],
743
+ sample_weight: Optional[np.ndarray] = ...,
744
+ context_label: str = ...,
745
+ ) -> CrossFitResult: ...
746
+
747
+
748
+ def cross_fit_predict(
749
+ learner: object,
750
+ X: np.ndarray,
751
+ y: np.ndarray,
752
+ folds: FoldAssignment,
753
+ *,
754
+ fit_mask: Optional[np.ndarray] = None,
755
+ predict_method: str = "predict",
756
+ sample_weight: Optional[np.ndarray] = None,
757
+ context_label: str = "",
758
+ ) -> CrossFitResult:
759
+ """Out-of-fold predictions for every unit.
760
+
761
+ Per fold k: fit the learner on ``train_mask(k) & fit_mask`` (e.g. the
762
+ untreated units), predict ALL units in fold k. ``context_label``
763
+ (e.g. the nuisance name or (g,t) cell) is prefixed into every
764
+ ``DegenerateFoldError`` message to identify WHICH cross-fit failed.
765
+ """
766
+ n_units = folds.n_units
767
+ oof = np.empty(n_units, dtype=np.float64)
768
+ fold_losses = np.empty(folds.n_folds, dtype=np.float64)
769
+ n_fit_per_fold = np.empty(folds.n_folds, dtype=np.int64)
770
+
771
+ # stacklevel 5: _fresh_learner -> _fit_subset -> generator frame ->
772
+ # cross_fit_predict -> the user's call site.
773
+ for ff in iter_fold_fits(
774
+ learner,
775
+ X,
776
+ y,
777
+ folds,
778
+ fit_mask=fit_mask,
779
+ predict_method=predict_method,
780
+ sample_weight=sample_weight,
781
+ context_label=context_label,
782
+ warn_stacklevel=5,
783
+ ):
784
+ n_fit_per_fold[ff.k] = ff.n_fit
785
+ pred = _predict_subset(
786
+ ff.learner,
787
+ ff.X[ff.test_idx],
788
+ kind=ff.kind,
789
+ k=ff.k,
790
+ label=ff.label,
791
+ n_fit=ff.n_fit,
792
+ w_fit=ff.w_fit,
793
+ )
794
+ oof[ff.test_idx] = pred
795
+ w_test = None if ff.sample_weight is None else ff.sample_weight[ff.test_idx]
796
+ fold_losses[ff.k] = _fold_loss(ff.kind, ff.y[ff.test_idx], pred, w_test)
600
797
 
601
798
  return CrossFitResult(
602
799
  oof_predictions=oof,