diff-diff 3.6.1__tar.gz → 3.7.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {diff_diff-3.6.1 → diff_diff-3.7.0}/PKG-INFO +3 -1
- {diff_diff-3.6.1 → diff_diff-3.7.0}/README.md +1 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/__init__.py +1 -1
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_backend.py +38 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_nprobust_port.py +13 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_reporting_helpers.py +20 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/bacon.py +15 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/bootstrap_chunking.py +148 -1
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/chaisemartin_dhaultfoeuille.py +10 -2
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/continuous_did.py +871 -139
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/continuous_did_bspline.py +165 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/continuous_did_results.py +33 -5
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/diagnostic_report.py +181 -7
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did.py +171 -65
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_bootstrap.py +76 -43
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_covariates.py +655 -12
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_results.py +4 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_weights.py +175 -4
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/estimators.py +173 -14
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms-autonomous.txt +28 -24
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms-full.txt +62 -25
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms-practitioner.txt +5 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms.txt +4 -2
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/had.py +531 -519
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/had_pretests.py +124 -559
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/honest_did.py +255 -68
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/imputation.py +208 -146
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/imputation_results.py +6 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/linalg.py +333 -9
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/practitioner.py +7 -3
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/profile.py +16 -13
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/results.py +5 -3
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/spillover.py +40 -78
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/stacked_did.py +4 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered.py +674 -198
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_aggregation.py +317 -89
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_bootstrap.py +138 -105
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_results.py +9 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_triple_diff.py +3 -2
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/sun_abraham.py +16 -14
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/survey.py +68 -65
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/synthetic_control.py +198 -21
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/synthetic_control_results.py +717 -72
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/twfe.py +48 -53
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/two_stage.py +72 -128
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/utils.py +618 -81
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/wooldridge.py +21 -1
- {diff_diff-3.6.1 → diff_diff-3.7.0}/pyproject.toml +2 -1
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/Cargo.lock +37 -103
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/Cargo.toml +5 -1
- diff_diff-3.7.0/rust/src/alloc_profile.rs +70 -0
- diff_diff-3.7.0/rust/src/batched_solve.rs +353 -0
- diff_diff-3.7.0/rust/src/demean.rs +333 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/lib.rs +25 -1
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/linalg.rs +157 -60
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/trop.rs +212 -87
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/weights.rs +227 -61
- {diff_diff-3.6.1 → diff_diff-3.7.0}/LICENSE +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_guides_api.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/agent_workflow.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/balancing.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/bootstrap_utils.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/business_report.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/chaisemartin_dhaultfoeuille_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/conformal.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/conley.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/datasets.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/diagnostics.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/__init__.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/imputation_bootstrap.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/local_linear.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/lpdid.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/lpdid_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/power.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/prep.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/prep_dgp.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/pretrends.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/stacked_did_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_triple_diff_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/synthetic_did.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/triple_diff.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop_global.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop_local.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/two_stage_bootstrap.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/two_stage_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/__init__.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_common.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_continuous.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_diagnostic.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_event_study.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_power.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_staggered.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_synthetic.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/wooldridge_results.py +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/build.rs +0 -0
- {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/bootstrap.rs +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: diff-diff
|
|
3
|
-
Version: 3.
|
|
3
|
+
Version: 3.7.0
|
|
4
4
|
Classifier: Development Status :: 5 - Production/Stable
|
|
5
5
|
Classifier: Intended Audience :: Science/Research
|
|
6
6
|
Classifier: Operating System :: OS Independent
|
|
@@ -32,6 +32,7 @@ Requires-Dist: pydata-sphinx-theme>=0.16.1 ; extra == 'docs'
|
|
|
32
32
|
Requires-Dist: sphinxext-opengraph>=0.9 ; extra == 'docs'
|
|
33
33
|
Requires-Dist: sphinx-sitemap>=2.5 ; extra == 'docs'
|
|
34
34
|
Requires-Dist: nbsphinx>=0.9 ; extra == 'docs'
|
|
35
|
+
Requires-Dist: myst-parser>=2.0 ; extra == 'docs'
|
|
35
36
|
Requires-Dist: matplotlib>=3.5 ; extra == 'docs'
|
|
36
37
|
Requires-Dist: ipython>=8.0 ; extra == 'docs'
|
|
37
38
|
Requires-Dist: plotly>=5.0 ; extra == 'plotly'
|
|
@@ -190,6 +191,7 @@ Most estimators accept an optional `survey_design` parameter (or `survey=` / `we
|
|
|
190
191
|
- **Variance methods**: Taylor Series Linearization (TSL via Binder 1983), replicate weights (BRR / Fay / JK1 / JKn / SDR), survey-aware bootstrap
|
|
191
192
|
- **Diagnostics**: DEFF per coefficient, effective n, subpopulation analysis, weight trimming, CV on estimates
|
|
192
193
|
- **Repeated cross-sections**: `CallawaySantAnna(panel=False)` for BRFSS, ACS, CPS
|
|
194
|
+
- **Weight calibration / raking**: upstream by design - pair with Meta's [balance](https://import-balance.org/) package, whose `balance.interop.diff_diff` adapter hands raked samples straight to diff-diff; see the [composition-drift tutorial](https://diff-diff.readthedocs.io/en/stable/tutorials/26_composition_drift_calibration.html)
|
|
193
195
|
|
|
194
196
|
No other Python or R DiD package offers design-based variance estimation for modern heterogeneity-robust estimators.
|
|
195
197
|
|
|
@@ -137,6 +137,7 @@ Most estimators accept an optional `survey_design` parameter (or `survey=` / `we
|
|
|
137
137
|
- **Variance methods**: Taylor Series Linearization (TSL via Binder 1983), replicate weights (BRR / Fay / JK1 / JKn / SDR), survey-aware bootstrap
|
|
138
138
|
- **Diagnostics**: DEFF per coefficient, effective n, subpopulation analysis, weight trimming, CV on estimates
|
|
139
139
|
- **Repeated cross-sections**: `CallawaySantAnna(panel=False)` for BRFSS, ACS, CPS
|
|
140
|
+
- **Weight calibration / raking**: upstream by design - pair with Meta's [balance](https://import-balance.org/) package, whose `balance.interop.diff_diff` adapter hands raked samples straight to diff-diff; see the [composition-drift tutorial](https://diff-diff.readthedocs.io/en/stable/tutorials/26_composition_drift_calibration.html)
|
|
140
141
|
|
|
141
142
|
No other Python or R DiD package offers design-based variance estimation for modern heterogeneity-robust estimators.
|
|
142
143
|
|
|
@@ -65,6 +65,34 @@ except ImportError:
|
|
|
65
65
|
_rust_sc_weight_fw_weighted_with_convergence = None
|
|
66
66
|
_rust_backend_info = None
|
|
67
67
|
|
|
68
|
+
# FE-absorption MAP demeaning kernel: imported independently so a stale or
|
|
69
|
+
# mixed-version extension missing only this newer symbol degrades to the
|
|
70
|
+
# numpy demeaning engine WITHOUT disabling the older Rust accelerations.
|
|
71
|
+
try:
|
|
72
|
+
from diff_diff._rust_backend import demean_map as _rust_demean_map
|
|
73
|
+
except ImportError:
|
|
74
|
+
_rust_demean_map = None
|
|
75
|
+
|
|
76
|
+
# Batched ridge-regularized SPD solve (EfficientDiD per-unit weights):
|
|
77
|
+
# imported independently for the same mixed-version reason as demean_map.
|
|
78
|
+
try:
|
|
79
|
+
from diff_diff._rust_backend import (
|
|
80
|
+
batched_ridge_chol_solve_ones as _rust_batched_ridge_chol_solve,
|
|
81
|
+
)
|
|
82
|
+
except ImportError:
|
|
83
|
+
_rust_batched_ridge_chol_solve = None
|
|
84
|
+
|
|
85
|
+
# HC2 (leverage-corrected) robust vcov: imported independently for the same
|
|
86
|
+
# mixed-version reason as demean_map (a stale extension missing only this
|
|
87
|
+
# newer symbol degrades HC2 to the NumPy path without disabling the older
|
|
88
|
+
# Rust accelerations).
|
|
89
|
+
try:
|
|
90
|
+
from diff_diff._rust_backend import (
|
|
91
|
+
compute_robust_vcov_hc2 as _rust_compute_robust_vcov_hc2,
|
|
92
|
+
)
|
|
93
|
+
except ImportError:
|
|
94
|
+
_rust_compute_robust_vcov_hc2 = None
|
|
95
|
+
|
|
68
96
|
# Determine final backend based on environment variable and availability
|
|
69
97
|
if _backend_env == "python":
|
|
70
98
|
# Force pure Python mode - disable Rust even if available
|
|
@@ -73,6 +101,12 @@ if _backend_env == "python":
|
|
|
73
101
|
_rust_project_simplex = None
|
|
74
102
|
_rust_solve_ols = None
|
|
75
103
|
_rust_compute_robust_vcov = None
|
|
104
|
+
# FE-absorption MAP demeaning kernel
|
|
105
|
+
_rust_demean_map = None
|
|
106
|
+
# Batched ridge-regularized SPD solve
|
|
107
|
+
_rust_batched_ridge_chol_solve = None
|
|
108
|
+
# HC2 robust vcov
|
|
109
|
+
_rust_compute_robust_vcov_hc2 = None
|
|
76
110
|
# TROP estimator acceleration (local method)
|
|
77
111
|
_rust_unit_distance_matrix = None
|
|
78
112
|
_rust_loocv_grid_search = None
|
|
@@ -124,6 +158,10 @@ __all__ = [
|
|
|
124
158
|
"_rust_project_simplex",
|
|
125
159
|
"_rust_solve_ols",
|
|
126
160
|
"_rust_compute_robust_vcov",
|
|
161
|
+
# FE-absorption MAP demeaning kernel
|
|
162
|
+
"_rust_demean_map",
|
|
163
|
+
# Batched ridge-regularized SPD solve (EfficientDiD per-unit weights)
|
|
164
|
+
"_rust_batched_ridge_chol_solve",
|
|
127
165
|
# TROP estimator acceleration (local method)
|
|
128
166
|
"_rust_unit_distance_matrix",
|
|
129
167
|
"_rust_loocv_grid_search",
|
|
@@ -1361,6 +1361,19 @@ def lprobust(
|
|
|
1361
1361
|
se_cl = float(np.sqrt((deriv_fact**2) * V_Y_cl[deriv, deriv]))
|
|
1362
1362
|
se_rb = float(np.sqrt((deriv_fact**2) * V_Y_bc[deriv, deriv]))
|
|
1363
1363
|
|
|
1364
|
+
# Cluster-robust variance is unidentified when fewer than two clusters
|
|
1365
|
+
# contribute to the ACTIVE kernel window (``eC = cluster[ind]``): the
|
|
1366
|
+
# between-cluster meat is degenerate, so a finite ``se`` here would report
|
|
1367
|
+
# unidentified clustered inference as if identified. NaN both SEs so any
|
|
1368
|
+
# downstream inference (the ``safe_inference`` gate in
|
|
1369
|
+
# ``bias_corrected_local_linear``; HAD's beta-scale rescale) is NaN-coupled.
|
|
1370
|
+
# Unclustered fits (``eC is None``) are unaffected, and a clustered window
|
|
1371
|
+
# with >= 2 distinct clusters is bit-identical, so the DGP-4 golden parity
|
|
1372
|
+
# is preserved.
|
|
1373
|
+
if eC is not None and len(np.unique(eC)) < 2:
|
|
1374
|
+
se_cl = float("nan")
|
|
1375
|
+
se_rb = float("nan")
|
|
1376
|
+
|
|
1364
1377
|
# --- Per-observation influence function for the BIAS-CORRECTED point
|
|
1365
1378
|
# estimate at ``deriv`` (Phase 4.5 survey composition).
|
|
1366
1379
|
# Aligned with ``V_Y_bc`` (NOT ``V_Y_cl``) so survey-composed variance
|
|
@@ -635,6 +635,26 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]:
|
|
|
635
635
|
"reference": "REGISTRY.md Sec. SyntheticControl",
|
|
636
636
|
}
|
|
637
637
|
|
|
638
|
+
if name == "SpilloverDiDResults":
|
|
639
|
+
return {
|
|
640
|
+
"name": "total effect on the treated (Butts spillover-aware ATT)",
|
|
641
|
+
"definition": (
|
|
642
|
+
"The total effect on the treated ``tau_total`` from Butts (2021) "
|
|
643
|
+
"ring-indicator spillover DiD, identified off FAR-AWAY control "
|
|
644
|
+
"observations (``d_it > d_bar``, Assumption 5) rather than any "
|
|
645
|
+
"not-yet-/never-treated pool. The estimator decomposes into the "
|
|
646
|
+
"DIRECT effect on treated units plus per-ring spillover-on-control "
|
|
647
|
+
"effects that relax SUTVA within the treated units' spatial "
|
|
648
|
+
"neighborhood; ``att`` is the headline total effect, while the "
|
|
649
|
+
"per-ring ``spillover_effects`` and (when ``event_study=True``) the "
|
|
650
|
+
"per-event-time direct dynamics are available on the result object "
|
|
651
|
+
"for disaggregated inference."
|
|
652
|
+
),
|
|
653
|
+
"aggregation": "spillover",
|
|
654
|
+
"headline_attribute": "att",
|
|
655
|
+
"reference": "Butts (2021); REGISTRY.md Sec. SpilloverDiD",
|
|
656
|
+
}
|
|
657
|
+
|
|
638
658
|
# Default: unrecognized result class. Fall through with a neutral
|
|
639
659
|
# block — agents / downstream consumers can still dispatch on
|
|
640
660
|
# ``aggregation="unknown"`` and fall back to generic ATT narration.
|
|
@@ -18,6 +18,7 @@ import numpy as np
|
|
|
18
18
|
import pandas as pd
|
|
19
19
|
|
|
20
20
|
from diff_diff.results import _format_survey_block
|
|
21
|
+
from diff_diff.utils import pre_demean_norms, snap_absorbed_regressors
|
|
21
22
|
from diff_diff.utils import within_transform as _within_transform_util
|
|
22
23
|
|
|
23
24
|
|
|
@@ -795,6 +796,7 @@ class BaconDecomposition:
|
|
|
795
796
|
) -> float:
|
|
796
797
|
"""Compute TWFE estimate using within-transformation."""
|
|
797
798
|
# Apply two-way within transformation (weighted if survey weights provided)
|
|
799
|
+
_pre_norms = pre_demean_norms(df, [treat_col], weights=weights)
|
|
798
800
|
df_dm = _within_transform_util(
|
|
799
801
|
df,
|
|
800
802
|
[outcome, treat_col],
|
|
@@ -803,6 +805,19 @@ class BaconDecomposition:
|
|
|
803
805
|
suffix="_within",
|
|
804
806
|
weights=weights,
|
|
805
807
|
)
|
|
808
|
+
# Snap an FE-spanned treatment to exact zero: the d_var == 0 guard
|
|
809
|
+
# below then returns its deterministic 0.0 (with the cause warning)
|
|
810
|
+
# instead of an arbitrary junk/junk division.
|
|
811
|
+
snap_absorbed_regressors(
|
|
812
|
+
df_dm,
|
|
813
|
+
[treat_col],
|
|
814
|
+
_pre_norms,
|
|
815
|
+
absorbed_desc=f"unit '{unit}' and time '{time}' fixed effects",
|
|
816
|
+
group_vars=[unit, time],
|
|
817
|
+
suffix="_within",
|
|
818
|
+
display_names={treat_col: "treatment"},
|
|
819
|
+
weights=weights,
|
|
820
|
+
)
|
|
806
821
|
|
|
807
822
|
# Extract within-transformed values
|
|
808
823
|
y_within = df_dm[f"{outcome}_within"].values
|
|
@@ -23,11 +23,27 @@ Bit-identity of the weight *generation* is preserved on **both** backends:
|
|
|
23
23
|
- The **NumPy** fallback draws the matrix row-major from the ``Generator``
|
|
24
24
|
stream, so consuming it in contiguous, in-order blocks from the same generator
|
|
25
25
|
reproduces the identical sequence.
|
|
26
|
+
|
|
27
|
+
The *influence* side can be tiled too. When a fit needs many perturbation
|
|
28
|
+
columns (per-cell IFs, the overall combined IF, per-event-time combined IFs),
|
|
29
|
+
computing them one column at a time degenerates into per-column slicing/GEMV
|
|
30
|
+
over the weight blocks -- memory-bandwidth-bound, not FLOP-bound. Instead,
|
|
31
|
+
:func:`tiled_if_matmul` scatters the columns into an ``(n_units, tile_cols)``
|
|
32
|
+
buffer and runs one BLAS GEMM per (weight block, column tile). The column axis
|
|
33
|
+
is tiled under a byte cap, and each column tile makes its own full pass over
|
|
34
|
+
the weight stream via :class:`ReplayableWeightStream`, which snapshots the
|
|
35
|
+
generator state at construction and restores it per pass -- so every pass sees
|
|
36
|
+
the bit-identical weight stream on both backends (Rust re-draws the same
|
|
37
|
+
``base_seed`` from the restored state; the NumPy fallback replays the stream
|
|
38
|
+
itself; the survey stratified branch regenerates its full PSU matrix from the
|
|
39
|
+
restored state). GEMM results match the per-column GEMV path to within BLAS
|
|
40
|
+
reassociation, same as the draw-axis chunking above.
|
|
26
41
|
"""
|
|
27
42
|
|
|
28
43
|
from __future__ import annotations
|
|
29
44
|
|
|
30
|
-
from
|
|
45
|
+
from collections import abc
|
|
46
|
+
from typing import Callable, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
|
|
31
47
|
|
|
32
48
|
import numpy as np
|
|
33
49
|
|
|
@@ -183,3 +199,134 @@ def iter_survey_multiplier_weight_blocks(
|
|
|
183
199
|
yield chunk_start, block
|
|
184
200
|
|
|
185
201
|
return psu_ids, _generated()
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
# Byte ceiling for the ``(n_units, tile_cols)`` float64 scatter buffer used by
|
|
205
|
+
# ``tiled_if_matmul``. Sibling of ``_TARGET_BLOCK_BYTES``; together they bound
|
|
206
|
+
# the kernel's live intermediates at ~384 MB regardless of n_units or the
|
|
207
|
+
# number of perturbation columns. Read at CALL time inside ``tiled_if_matmul``
|
|
208
|
+
# (never bound as a def-time default) so tests can monkeypatch it.
|
|
209
|
+
_TARGET_TILE_BYTES = 128 * 1024 * 1024
|
|
210
|
+
|
|
211
|
+
# One perturbation column of ``tiled_if_matmul``'s influence matrix: a list of
|
|
212
|
+
# (index, values) contributions. ``index=None`` assigns the full-width dense
|
|
213
|
+
# ``values``; otherwise ``values`` is scattered at ``index``. A column may also
|
|
214
|
+
# be a zero-argument callable returning such a list, materialized only when its
|
|
215
|
+
# tile is filled (keeps peak memory at one O(n_units) temporary for callers
|
|
216
|
+
# whose columns are computed on the fly).
|
|
217
|
+
IFColumn = List[Tuple[Optional[np.ndarray], np.ndarray]]
|
|
218
|
+
IFColumnSpec = Union[IFColumn, Callable[[], IFColumn]]
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
class ReplayableWeightStream:
|
|
222
|
+
"""Multiplier-weight block stream that can be re-iterated bit-identically.
|
|
223
|
+
|
|
224
|
+
Wraps a weight-block iterator factory together with the ``Generator`` that
|
|
225
|
+
feeds it, snapshotting the generator state at construction. Every
|
|
226
|
+
``__iter__`` restores the snapshot before invoking the factory, so each
|
|
227
|
+
pass yields the exact same ``(chunk_start, block)`` sequence:
|
|
228
|
+
|
|
229
|
+
- **Rust backend**: ``iter_weight_blocks`` draws its per-run ``base_seed``
|
|
230
|
+
from the rng lazily at the first ``next()``; the restored state
|
|
231
|
+
reproduces the same base seed, and rows are seeded absolutely from it.
|
|
232
|
+
- **NumPy fallback**: generation consumes the rng stream directly; the
|
|
233
|
+
restored state replays the identical stream.
|
|
234
|
+
- **Survey stratified branch**: ``iter_survey_multiplier_weight_blocks``
|
|
235
|
+
generates its full PSU matrix eagerly at factory-call time -- from the
|
|
236
|
+
restored state, so per-pass recreation is also bit-identical.
|
|
237
|
+
|
|
238
|
+
The snapshot MUST precede any rng-consuming call: construct this object
|
|
239
|
+
before probing generators that draw at call time (the survey stratified
|
|
240
|
+
branch draws eagerly; resolve ``psu_ids`` from the design's ``psu`` array
|
|
241
|
+
directly instead of via a probe call). After the final pass the rng ends
|
|
242
|
+
exactly where a single pass would leave it, so downstream consumers of the
|
|
243
|
+
same generator are unaffected by the number of passes.
|
|
244
|
+
"""
|
|
245
|
+
|
|
246
|
+
def __init__(
|
|
247
|
+
self,
|
|
248
|
+
make_iter: Callable[[np.random.Generator], Iterator[Tuple[int, np.ndarray]]],
|
|
249
|
+
rng: np.random.Generator,
|
|
250
|
+
) -> None:
|
|
251
|
+
self._make_iter = make_iter
|
|
252
|
+
self._rng = rng
|
|
253
|
+
self._state = rng.bit_generator.state
|
|
254
|
+
|
|
255
|
+
def __iter__(self) -> Iterator[Tuple[int, np.ndarray]]:
|
|
256
|
+
self._rng.bit_generator.state = self._state
|
|
257
|
+
return self._make_iter(self._rng)
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
def tiled_if_matmul(
|
|
261
|
+
weight_stream: Iterable[Tuple[int, np.ndarray]],
|
|
262
|
+
n_bootstrap: int,
|
|
263
|
+
n_units: int,
|
|
264
|
+
columns: Sequence[IFColumnSpec],
|
|
265
|
+
*,
|
|
266
|
+
tile_bytes: Optional[int] = None,
|
|
267
|
+
) -> np.ndarray:
|
|
268
|
+
"""Compute ``W @ M`` for the multiplier bootstrap without materializing
|
|
269
|
+
``W`` (streamed in row blocks) or ``M`` (scattered in column tiles).
|
|
270
|
+
|
|
271
|
+
``W`` is the ``(n_bootstrap, n_units)`` weight matrix yielded by
|
|
272
|
+
``weight_stream`` as ``(chunk_start, block)`` row blocks; ``M`` is the
|
|
273
|
+
``(n_units, len(columns))`` influence matrix described by ``columns``
|
|
274
|
+
(see :data:`IFColumnSpec`). Returns the dense ``(n_bootstrap, n_cols)``
|
|
275
|
+
product. Columns are tiled under ``tile_bytes`` (default: module attribute
|
|
276
|
+
``_TARGET_TILE_BYTES``, resolved at call time); each tile is scattered
|
|
277
|
+
once into an F-order buffer and then consumed by one BLAS GEMM per weight
|
|
278
|
+
block. With more than one tile, ``weight_stream`` is iterated once per
|
|
279
|
+
tile and therefore must be re-iterable with an identical stream --
|
|
280
|
+
:class:`ReplayableWeightStream` provides exactly that.
|
|
281
|
+
|
|
282
|
+
Scatter contract: contributions are written by plain assignment, so index
|
|
283
|
+
arrays must be unique within each contribution and the contributions of a
|
|
284
|
+
single column must be mutually disjoint (duplicate or overlapping indices
|
|
285
|
+
keep the last write instead of summing like ``W[:, idx] @ values`` would).
|
|
286
|
+
All in-repo callers satisfy this by construction (treated/control unit
|
|
287
|
+
sets are disjoint; dense columns are single-contribution).
|
|
288
|
+
|
|
289
|
+
A non-finite influence value poisons only its own output column (every
|
|
290
|
+
output element is an independent dot product), matching the per-column
|
|
291
|
+
GEMV behavior this kernel replaces. Results match the per-column GEMV
|
|
292
|
+
path to within BLAS reassociation (~1 ULP), not bit-for-bit.
|
|
293
|
+
"""
|
|
294
|
+
n_cols = len(columns)
|
|
295
|
+
out = np.empty((n_bootstrap, n_cols))
|
|
296
|
+
if n_cols == 0 or n_bootstrap == 0:
|
|
297
|
+
return out
|
|
298
|
+
if tile_bytes is None:
|
|
299
|
+
tile_bytes = _TARGET_TILE_BYTES
|
|
300
|
+
tile_cols = int(max(1, min(n_cols, tile_bytes // (max(1, n_units) * 8))))
|
|
301
|
+
|
|
302
|
+
# A single-pass iterator silently exhausts after the first tile, leaving
|
|
303
|
+
# uninitialized rows in `out` -- fail loudly instead. isinstance (not an
|
|
304
|
+
# iter() probe) so re-iterable streams see no extra pass.
|
|
305
|
+
if tile_cols < n_cols and isinstance(weight_stream, abc.Iterator):
|
|
306
|
+
raise ValueError(
|
|
307
|
+
"tiled_if_matmul needs more than one column tile but weight_stream "
|
|
308
|
+
"is a single-pass iterator; pass a re-iterable stream "
|
|
309
|
+
"(ReplayableWeightStream) so each tile can replay the weights."
|
|
310
|
+
)
|
|
311
|
+
|
|
312
|
+
buf = np.zeros((n_units, tile_cols), order="F")
|
|
313
|
+
for lo in range(0, n_cols, tile_cols):
|
|
314
|
+
hi = min(lo + tile_cols, n_cols)
|
|
315
|
+
width = hi - lo
|
|
316
|
+
if lo > 0:
|
|
317
|
+
buf[:, :width] = 0.0
|
|
318
|
+
for c in range(lo, hi):
|
|
319
|
+
spec = columns[c]
|
|
320
|
+
if callable(spec):
|
|
321
|
+
spec = spec()
|
|
322
|
+
col = buf[:, c - lo]
|
|
323
|
+
for idx, values in spec:
|
|
324
|
+
if idx is None:
|
|
325
|
+
np.copyto(col, values)
|
|
326
|
+
else:
|
|
327
|
+
col[idx] = values
|
|
328
|
+
tile = buf[:, :width]
|
|
329
|
+
with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
|
|
330
|
+
for chunk_start, w_block in weight_stream:
|
|
331
|
+
out[chunk_start : chunk_start + w_block.shape[0], lo:hi] = w_block @ tile
|
|
332
|
+
return out
|
|
@@ -4665,7 +4665,12 @@ def _compute_placebo(
|
|
|
4665
4665
|
else:
|
|
4666
4666
|
joiner_avg = float((y_prev[joiner_mask] - y_pre_prev[joiner_mask]).mean())
|
|
4667
4667
|
stable0_avg = float((y_prev[stable0_mask] - y_pre_prev[stable0_mask]).mean())
|
|
4668
|
-
|
|
4668
|
+
# Backward-difference convention (pre-period minus reference,
|
|
4669
|
+
# Y_{t-2} - Y_{t-1}), matching the multi-horizon placebo path
|
|
4670
|
+
# `_compute_multi_horizon_placebos` (`switcher_change - ctrl_avg`
|
|
4671
|
+
# with `Y_bwd - Y_ref`, times switch_direction=+1 for joiners) and
|
|
4672
|
+
# R `did_multiplegt_dyn`. Equivalently `stable0_avg - joiner_avg`.
|
|
4673
|
+
placebo_plus_t = stable0_avg - joiner_avg
|
|
4669
4674
|
|
|
4670
4675
|
# Leavers side: symmetric A11 distinction
|
|
4671
4676
|
if n_01 == 0:
|
|
@@ -4678,7 +4683,10 @@ def _compute_placebo(
|
|
|
4678
4683
|
else:
|
|
4679
4684
|
stable1_avg = float((y_prev[stable1_mask] - y_pre_prev[stable1_mask]).mean())
|
|
4680
4685
|
leaver_avg = float((y_prev[leaver_mask] - y_pre_prev[leaver_mask]).mean())
|
|
4681
|
-
|
|
4686
|
+
# Backward-difference convention times switch_direction=-1 for
|
|
4687
|
+
# leavers (see the joiners-side note above): the multi-horizon path
|
|
4688
|
+
# contributes `leaver_avg - stable1_avg` here.
|
|
4689
|
+
placebo_minus_t = leaver_avg - stable1_avg
|
|
4682
4690
|
|
|
4683
4691
|
placebo_plus_per_t.append(placebo_plus_t)
|
|
4684
4692
|
placebo_minus_per_t.append(placebo_minus_t)
|