diff-diff 3.3.2__tar.gz → 3.3.3__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.3.2 → diff_diff-3.3.3}/PKG-INFO +4 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/README.md +1 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/__init__.py +1 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/bacon.py +4 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/bootstrap_utils.py +132 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/chaisemartin_dhaultfoeuille.py +734 -145
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/chaisemartin_dhaultfoeuille_results.py +193 -50
- diff_diff-3.3.3/diff_diff/conley.py +1006 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/continuous_did_results.py +41 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/diagnostic_report.py +2 -3
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/efficient_did_results.py +23 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/estimators.py +219 -7
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/guides/llms-full.txt +347 -3
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/guides/llms-practitioner.txt +39 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/guides/llms.txt +2 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/had.py +63 -61
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/had_pretests.py +55 -55
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/honest_did.py +11 -4
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/imputation.py +1 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/imputation_results.py +25 -4
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/linalg.py +376 -17
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/practitioner.py +477 -25
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/prep.py +2 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/results.py +84 -24
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/stacked_did_results.py +23 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/staggered.py +5 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/staggered_results.py +23 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/staggered_triple_diff_results.py +23 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/sun_abraham.py +23 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/synthetic_did.py +151 -88
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/triple_diff.py +10 -4
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/trop.py +1 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/trop_results.py +2 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/twfe.py +111 -7
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/two_stage.py +1 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/two_stage_results.py +23 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/wooldridge_results.py +25 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/pyproject.toml +6 -2
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/Cargo.lock +6 -6
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/Cargo.toml +1 -1
- {diff_diff-3.3.2 → diff_diff-3.3.3}/LICENSE +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/_backend.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/_guides_api.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/_nprobust_port.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/_reporting_helpers.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/business_report.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/continuous_did.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/continuous_did_bspline.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/datasets.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/diagnostics.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/efficient_did.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/efficient_did_bootstrap.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/efficient_did_covariates.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/efficient_did_weights.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/guides/__init__.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/guides/llms-autonomous.txt +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/imputation_bootstrap.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/local_linear.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/power.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/prep_dgp.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/pretrends.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/profile.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/stacked_did.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/staggered_aggregation.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/staggered_bootstrap.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/staggered_triple_diff.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/survey.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/trop_global.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/trop_local.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/two_stage_bootstrap.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/utils.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/__init__.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_common.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_continuous.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_diagnostic.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_event_study.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_power.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_staggered.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/visualization/_synthetic.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/diff_diff/wooldridge.py +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/build.rs +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/src/bootstrap.rs +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/src/lib.rs +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/src/linalg.rs +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/src/trop.rs +0 -0
- {diff_diff-3.3.2 → diff_diff-3.3.3}/rust/src/weights.rs +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: diff-diff
|
|
3
|
-
Version: 3.3.
|
|
3
|
+
Version: 3.3.3
|
|
4
4
|
Classifier: Development Status :: 5 - Production/Stable
|
|
5
5
|
Classifier: Intended Audience :: Science/Research
|
|
6
6
|
Classifier: Operating System :: OS Independent
|
|
@@ -28,11 +28,12 @@ Requires-Dist: matplotlib>=3.5 ; extra == 'dev'
|
|
|
28
28
|
Requires-Dist: nbmake>=1.5 ; extra == 'dev'
|
|
29
29
|
Requires-Dist: plotly>=5.0 ; extra == 'dev'
|
|
30
30
|
Requires-Dist: sphinx>=6.0 ; extra == 'docs'
|
|
31
|
-
Requires-Dist: pydata-sphinx-theme>=0.
|
|
31
|
+
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
35
|
Requires-Dist: matplotlib>=3.5 ; extra == 'docs'
|
|
36
|
+
Requires-Dist: ipython>=8.0 ; extra == 'docs'
|
|
36
37
|
Requires-Dist: plotly>=5.0 ; extra == 'plotly'
|
|
37
38
|
Provides-Extra: dev
|
|
38
39
|
Provides-Extra: docs
|
|
@@ -176,6 +177,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
|
|
|
176
177
|
- [Honest DiD](https://diff-diff.readthedocs.io/en/stable/api/honest_did.html) - Rambachan & Roth (2023) sensitivity analysis: robust CI under PT violations, breakdown values
|
|
177
178
|
- [Pre-Trends Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/pretrends.html) - Roth (2022) minimum detectable violation and power curves
|
|
178
179
|
- [Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/power.html) - analytical and simulation-based MDE, sample size, power curves for study design
|
|
180
|
+
- Conley spatial HAC SE (`vcov_type="conley"`) on cross-sectional `LinearRegression` / `compute_robust_vcov` plus panel `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` (with `conley_lag_cutoff` for within-unit Bartlett temporal HAC) - Conley (1999) spatial-correlation-aware SEs with parity vs R `conleyreg` on cross-sectional + panel fixtures, optional combined spatial + cluster product kernel via explicit `cluster=`, auto-activating sparse k-d-tree fast path for `n > 5_000`
|
|
179
181
|
|
|
180
182
|
## Survey Support
|
|
181
183
|
|
|
@@ -124,6 +124,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
|
|
|
124
124
|
- [Honest DiD](https://diff-diff.readthedocs.io/en/stable/api/honest_did.html) - Rambachan & Roth (2023) sensitivity analysis: robust CI under PT violations, breakdown values
|
|
125
125
|
- [Pre-Trends Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/pretrends.html) - Roth (2022) minimum detectable violation and power curves
|
|
126
126
|
- [Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/power.html) - analytical and simulation-based MDE, sample size, power curves for study design
|
|
127
|
+
- Conley spatial HAC SE (`vcov_type="conley"`) on cross-sectional `LinearRegression` / `compute_robust_vcov` plus panel `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` (with `conley_lag_cutoff` for within-unit Bartlett temporal HAC) - Conley (1999) spatial-correlation-aware SEs with parity vs R `conleyreg` on cross-sectional + panel fixtures, optional combined spatial + cluster product kernel via explicit `cluster=`, auto-activating sparse k-d-tree fast path for `n > 5_000`
|
|
127
128
|
|
|
128
129
|
## Survey Support
|
|
129
130
|
|
|
@@ -324,6 +324,7 @@ class BaconDecomposition:
|
|
|
324
324
|
----------
|
|
325
325
|
weights : str, default="approximate"
|
|
326
326
|
Weight calculation method:
|
|
327
|
+
|
|
327
328
|
- "approximate": Fast simplified formula using group shares and
|
|
328
329
|
treatment variance. Good for diagnostic purposes where relative
|
|
329
330
|
weights are sufficient to identify problematic comparisons.
|
|
@@ -401,6 +402,7 @@ class BaconDecomposition:
|
|
|
401
402
|
----------
|
|
402
403
|
weights : str, default="approximate"
|
|
403
404
|
Weight calculation method:
|
|
405
|
+
|
|
404
406
|
- "approximate": Fast simplified formula (default)
|
|
405
407
|
- "exact": Variance-based weights from Goodman-Bacon (2021)
|
|
406
408
|
"""
|
|
@@ -1084,6 +1086,7 @@ def bacon_decompose(
|
|
|
1084
1086
|
Use 0 (or np.inf) for never-treated units.
|
|
1085
1087
|
weights : str, default="approximate"
|
|
1086
1088
|
Weight calculation method:
|
|
1089
|
+
|
|
1087
1090
|
- "approximate": Fast simplified formula (default). Good for
|
|
1088
1091
|
diagnostic purposes where relative weights are sufficient.
|
|
1089
1092
|
- "exact": Variance-based weights from Goodman-Bacon (2021)
|
|
@@ -1093,6 +1096,7 @@ def bacon_decompose(
|
|
|
1093
1096
|
-------
|
|
1094
1097
|
BaconDecompositionResults
|
|
1095
1098
|
Object containing decomposition results with:
|
|
1099
|
+
|
|
1096
1100
|
- twfe_estimate: The overall TWFE coefficient
|
|
1097
1101
|
- comparisons: List of all 2x2 comparisons with estimates and weights
|
|
1098
1102
|
- Weight totals by comparison type
|
|
@@ -17,6 +17,7 @@ __all__ = [
|
|
|
17
17
|
"generate_bootstrap_weights_batch",
|
|
18
18
|
"generate_bootstrap_weights_batch_numpy",
|
|
19
19
|
"generate_survey_multiplier_weights_batch",
|
|
20
|
+
"apply_stratum_centering",
|
|
20
21
|
"generate_rao_wu_weights",
|
|
21
22
|
"generate_rao_wu_weights_batch",
|
|
22
23
|
"compute_percentile_ci",
|
|
@@ -654,6 +655,137 @@ def generate_survey_multiplier_weights_batch(
|
|
|
654
655
|
return weights, psu_ids
|
|
655
656
|
|
|
656
657
|
|
|
658
|
+
def apply_stratum_centering(
|
|
659
|
+
tensor: np.ndarray,
|
|
660
|
+
resolved_survey: "ResolvedSurveyDesign",
|
|
661
|
+
psu_ids: np.ndarray,
|
|
662
|
+
psu_axis: int = 0,
|
|
663
|
+
) -> np.ndarray:
|
|
664
|
+
"""Within-stratum demean + sqrt(n_h/(n_h-1)) Bessel rescale along the
|
|
665
|
+
PSU axis of a tensor. Mutates ``tensor`` in place AND returns it.
|
|
666
|
+
|
|
667
|
+
Shared by the HAD sup-t event-study bootstrap
|
|
668
|
+
(``had._sup_t_multiplier_bootstrap``, PSU axis = 0 on the
|
|
669
|
+
PSU-aggregated influence tensor ``Psi_psu`` of shape ``(n_psu,
|
|
670
|
+
n_horizons)``) AND the HAD Stute survey-bootstrap family
|
|
671
|
+
(``stute_test``, ``stute_joint_pretest``, PSU axis = 1 on the
|
|
672
|
+
multiplier matrix ``psu_mults`` of shape ``(n_bootstrap, n_psu)``).
|
|
673
|
+
Same algebra applied at different points in the two pipelines:
|
|
674
|
+
|
|
675
|
+
- HAD sup-t is a multiplier bootstrap on a precomputed influence
|
|
676
|
+
tensor. The correction is applied to the tensor before
|
|
677
|
+
``perturbations = psu_weights @ Psi_psu`` — see ``had.py:2151-2204``.
|
|
678
|
+
- Stute is a wild residual bootstrap with refit-in-loop and a
|
|
679
|
+
nonlinear functional. The correction is applied to the multipliers
|
|
680
|
+
before the per-obs broadcast ``eta_obs = psu_mults[b,
|
|
681
|
+
psu_col_idx]`` — see ``had_pretests.py:1988-2007``.
|
|
682
|
+
|
|
683
|
+
Locks the algebraic identity architecturally (so future drift in
|
|
684
|
+
one site cannot silently diverge from the other) and makes both
|
|
685
|
+
sites consume the same battle-tested code path.
|
|
686
|
+
|
|
687
|
+
The combined correction makes ``Var_xi[xi @ Psi_psu]`` match the
|
|
688
|
+
analytical Binder-TSL stratified target
|
|
689
|
+
``V = sum_h (1 - f_h) (n_h / (n_h - 1)) sum_j (psi_hj - psi_h_bar)²``
|
|
690
|
+
exactly (the ``(1 - f_h)`` factor is already baked into
|
|
691
|
+
``psu_mults`` by :func:`generate_survey_multiplier_weights_batch`;
|
|
692
|
+
this helper bakes the remaining ``(n_h / (n_h - 1))`` factor and
|
|
693
|
+
enforces the within-stratum-zero centering required for cluster
|
|
694
|
+
wild bootstrap consistency under stratification).
|
|
695
|
+
|
|
696
|
+
See REGISTRY § HeterogeneousAdoptionDiD —
|
|
697
|
+
"Note (Stute stratified survey-bootstrap calibration)" for the
|
|
698
|
+
derivation.
|
|
699
|
+
|
|
700
|
+
Parameters
|
|
701
|
+
----------
|
|
702
|
+
tensor : np.ndarray
|
|
703
|
+
Tensor with the PSU dimension on ``psu_axis``. Modified
|
|
704
|
+
in-place.
|
|
705
|
+
resolved_survey : ResolvedSurveyDesign
|
|
706
|
+
Resolved survey design. Provides ``.psu`` (per-unit PSU IDs;
|
|
707
|
+
may be None for the implicit-PSU case where each unit is its
|
|
708
|
+
own PSU) and ``.strata`` (per-unit stratum labels; may be None
|
|
709
|
+
for the single-implicit-stratum case).
|
|
710
|
+
psu_ids : np.ndarray, shape ``(n_psu,)``
|
|
711
|
+
Unique PSU identifiers aligned to the ``psu_axis`` of
|
|
712
|
+
``tensor``. Output of
|
|
713
|
+
:func:`generate_survey_multiplier_weights_batch`.
|
|
714
|
+
psu_axis : int, default 0
|
|
715
|
+
Axis of ``tensor`` along which PSUs are indexed. 0 for HAD
|
|
716
|
+
sup-t ``Psi_psu`` (shape ``(n_psu, n_horizons)``); 1 for Stute
|
|
717
|
+
``psu_mults`` (shape ``(n_bootstrap, n_psu)``).
|
|
718
|
+
|
|
719
|
+
Returns
|
|
720
|
+
-------
|
|
721
|
+
np.ndarray
|
|
722
|
+
Same object as ``tensor`` (in-place mutation; returned for
|
|
723
|
+
chaining). Singleton strata under ``lonely_psu='remove'`` /
|
|
724
|
+
``'certainty'`` have all-zero entries along ``psu_axis``
|
|
725
|
+
(set by :func:`generate_survey_multiplier_weights_batch`'s
|
|
726
|
+
lonely-PSU handling); the centering here skips them to avoid
|
|
727
|
+
a divide-by-zero on ``sqrt(n_h / 0)``.
|
|
728
|
+
|
|
729
|
+
Notes
|
|
730
|
+
-----
|
|
731
|
+
Under ``strata=None``, the correction is applied uniformly with a
|
|
732
|
+
single implicit stratum (``n_h = n_psu``) — demean across all PSUs
|
|
733
|
+
along ``psu_axis`` and rescale by ``sqrt(n_psu / (n_psu - 1))``.
|
|
734
|
+
This is the standard small-sample correction for an iid cluster
|
|
735
|
+
wild bootstrap (Wu 1986; Liu 1988) and matches the HAD sup-t
|
|
736
|
+
convention at ``had.py:2199-2204``.
|
|
737
|
+
|
|
738
|
+
The Stute call site has historically NOT applied this correction
|
|
739
|
+
(pre-PR Phase 4.5 C). Lifting the gate on stratified designs +
|
|
740
|
+
introducing this shared helper makes the non-strata Stute path
|
|
741
|
+
apply the correction uniformly, which is a deliberate calibration
|
|
742
|
+
improvement (~1-2% p-value shift for typical ``n_psu``) — see
|
|
743
|
+
REGISTRY § "Note (Stute stratified survey-bootstrap calibration)".
|
|
744
|
+
"""
|
|
745
|
+
n_psu = tensor.shape[psu_axis]
|
|
746
|
+
|
|
747
|
+
if resolved_survey.strata is not None:
|
|
748
|
+
strata = np.asarray(resolved_survey.strata)
|
|
749
|
+
# Build PSU -> stratum map. Constant-within-PSU by the
|
|
750
|
+
# SurveyDesign.resolve contract — assert defensively in tests.
|
|
751
|
+
psu_id_to_col = {int(p): c for c, p in enumerate(psu_ids)}
|
|
752
|
+
psu_stratum = np.empty(n_psu, dtype=strata.dtype)
|
|
753
|
+
if resolved_survey.psu is not None:
|
|
754
|
+
unit_psu = np.asarray(resolved_survey.psu)
|
|
755
|
+
seen = np.zeros(n_psu, dtype=bool)
|
|
756
|
+
for i in range(len(unit_psu)):
|
|
757
|
+
col = psu_id_to_col[int(unit_psu[i])]
|
|
758
|
+
if not seen[col]:
|
|
759
|
+
psu_stratum[col] = strata[i]
|
|
760
|
+
seen[col] = True
|
|
761
|
+
else:
|
|
762
|
+
# Each unit is its own PSU; psu_ids = arange(n_psu).
|
|
763
|
+
psu_stratum = strata.copy()
|
|
764
|
+
|
|
765
|
+
for h in np.unique(psu_stratum):
|
|
766
|
+
mask_h = psu_stratum == h
|
|
767
|
+
n_h = int(mask_h.sum())
|
|
768
|
+
if n_h < 2:
|
|
769
|
+
# Singleton / empty stratum: caller's lonely_psu logic
|
|
770
|
+
# has already zeroed the corresponding multipliers (or
|
|
771
|
+
# the contribution is zero by construction). Skip to
|
|
772
|
+
# avoid a divide-by-zero on sqrt(n_h / 0).
|
|
773
|
+
continue
|
|
774
|
+
slc = [slice(None)] * tensor.ndim
|
|
775
|
+
slc[psu_axis] = mask_h
|
|
776
|
+
slc = tuple(slc)
|
|
777
|
+
tensor[slc] -= tensor[slc].mean(axis=psu_axis, keepdims=True)
|
|
778
|
+
tensor[slc] *= np.sqrt(n_h / (n_h - 1))
|
|
779
|
+
else:
|
|
780
|
+
# Single implicit stratum — demean across all PSUs along
|
|
781
|
+
# psu_axis, rescale by sqrt(n_psu / (n_psu - 1)).
|
|
782
|
+
if n_psu >= 2:
|
|
783
|
+
tensor -= tensor.mean(axis=psu_axis, keepdims=True)
|
|
784
|
+
tensor *= np.sqrt(n_psu / (n_psu - 1))
|
|
785
|
+
|
|
786
|
+
return tensor
|
|
787
|
+
|
|
788
|
+
|
|
657
789
|
def generate_rao_wu_weights(
|
|
658
790
|
resolved_survey: "ResolvedSurveyDesign",
|
|
659
791
|
rng: np.random.Generator,
|