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.
Files changed (99) hide show
  1. {diff_diff-3.6.1 → diff_diff-3.7.0}/PKG-INFO +3 -1
  2. {diff_diff-3.6.1 → diff_diff-3.7.0}/README.md +1 -0
  3. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/__init__.py +1 -1
  4. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_backend.py +38 -0
  5. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_nprobust_port.py +13 -0
  6. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_reporting_helpers.py +20 -0
  7. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/bacon.py +15 -0
  8. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/bootstrap_chunking.py +148 -1
  9. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/chaisemartin_dhaultfoeuille.py +10 -2
  10. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/continuous_did.py +871 -139
  11. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/continuous_did_bspline.py +165 -0
  12. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/continuous_did_results.py +33 -5
  13. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/diagnostic_report.py +181 -7
  14. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did.py +171 -65
  15. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_bootstrap.py +76 -43
  16. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_covariates.py +655 -12
  17. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_results.py +4 -0
  18. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/efficient_did_weights.py +175 -4
  19. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/estimators.py +173 -14
  20. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms-autonomous.txt +28 -24
  21. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms-full.txt +62 -25
  22. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms-practitioner.txt +5 -0
  23. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/llms.txt +4 -2
  24. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/had.py +531 -519
  25. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/had_pretests.py +124 -559
  26. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/honest_did.py +255 -68
  27. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/imputation.py +208 -146
  28. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/imputation_results.py +6 -0
  29. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/linalg.py +333 -9
  30. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/practitioner.py +7 -3
  31. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/profile.py +16 -13
  32. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/results.py +5 -3
  33. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/spillover.py +40 -78
  34. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/stacked_did.py +4 -0
  35. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered.py +674 -198
  36. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_aggregation.py +317 -89
  37. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_bootstrap.py +138 -105
  38. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_results.py +9 -0
  39. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_triple_diff.py +3 -2
  40. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/sun_abraham.py +16 -14
  41. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/survey.py +68 -65
  42. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/synthetic_control.py +198 -21
  43. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/synthetic_control_results.py +717 -72
  44. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/twfe.py +48 -53
  45. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/two_stage.py +72 -128
  46. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/utils.py +618 -81
  47. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/wooldridge.py +21 -1
  48. {diff_diff-3.6.1 → diff_diff-3.7.0}/pyproject.toml +2 -1
  49. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/Cargo.lock +37 -103
  50. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/Cargo.toml +5 -1
  51. diff_diff-3.7.0/rust/src/alloc_profile.rs +70 -0
  52. diff_diff-3.7.0/rust/src/batched_solve.rs +353 -0
  53. diff_diff-3.7.0/rust/src/demean.rs +333 -0
  54. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/lib.rs +25 -1
  55. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/linalg.rs +157 -60
  56. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/trop.rs +212 -87
  57. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/src/weights.rs +227 -61
  58. {diff_diff-3.6.1 → diff_diff-3.7.0}/LICENSE +0 -0
  59. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/_guides_api.py +0 -0
  60. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/agent_workflow.py +0 -0
  61. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/balancing.py +0 -0
  62. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/bootstrap_utils.py +0 -0
  63. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/business_report.py +0 -0
  64. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +0 -0
  65. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/chaisemartin_dhaultfoeuille_results.py +0 -0
  66. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/conformal.py +0 -0
  67. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/conley.py +0 -0
  68. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/datasets.py +0 -0
  69. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/diagnostics.py +0 -0
  70. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/guides/__init__.py +0 -0
  71. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/imputation_bootstrap.py +0 -0
  72. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/local_linear.py +0 -0
  73. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/lpdid.py +0 -0
  74. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/lpdid_results.py +0 -0
  75. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/power.py +0 -0
  76. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/prep.py +0 -0
  77. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/prep_dgp.py +0 -0
  78. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/pretrends.py +0 -0
  79. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/stacked_did_results.py +0 -0
  80. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/staggered_triple_diff_results.py +0 -0
  81. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/synthetic_did.py +0 -0
  82. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/triple_diff.py +0 -0
  83. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop.py +0 -0
  84. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop_global.py +0 -0
  85. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop_local.py +0 -0
  86. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/trop_results.py +0 -0
  87. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/two_stage_bootstrap.py +0 -0
  88. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/two_stage_results.py +0 -0
  89. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/__init__.py +0 -0
  90. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_common.py +0 -0
  91. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_continuous.py +0 -0
  92. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_diagnostic.py +0 -0
  93. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_event_study.py +0 -0
  94. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_power.py +0 -0
  95. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_staggered.py +0 -0
  96. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/visualization/_synthetic.py +0 -0
  97. {diff_diff-3.6.1 → diff_diff-3.7.0}/diff_diff/wooldridge_results.py +0 -0
  98. {diff_diff-3.6.1 → diff_diff-3.7.0}/rust/build.rs +0 -0
  99. {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.6.1
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
 
@@ -301,7 +301,7 @@ ETWFE = WooldridgeDiD
301
301
  DCDH = ChaisemartinDHaultfoeuille
302
302
  HAD = HeterogeneousAdoptionDiD
303
303
 
304
- __version__ = "3.6.1"
304
+ __version__ = "3.7.0"
305
305
  __all__ = [
306
306
  # Estimators
307
307
  "DifferenceInDifferences",
@@ -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 typing import Iterator, Optional, Tuple
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
- placebo_plus_t = joiner_avg - stable0_avg
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
- placebo_minus_t = stable1_avg - leaver_avg
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)