diff-diff 2.7.5__tar.gz → 2.7.6__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-2.7.5 → diff_diff-2.7.6}/PKG-INFO +1 -1
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/__init__.py +8 -1
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/_backend.py +20 -19
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/bacon.py +32 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/bootstrap_utils.py +274 -17
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/continuous_did.py +235 -80
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/continuous_did_bspline.py +7 -5
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/continuous_did_results.py +2 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/datasets.py +179 -72
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/diagnostics.py +62 -101
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did.py +50 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_bootstrap.py +29 -4
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_results.py +2 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/estimators.py +39 -24
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/imputation.py +15 -8
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/imputation_bootstrap.py +79 -9
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/imputation_results.py +2 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/linalg.py +236 -28
- diff_diff-2.7.6/diff_diff/practitioner.py +869 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/prep.py +75 -92
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/prep_dgp.py +69 -60
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/results.py +53 -36
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/stacked_did.py +9 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/stacked_did_results.py +2 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered.py +173 -53
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered_aggregation.py +56 -6
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered_bootstrap.py +186 -93
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered_results.py +2 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/sun_abraham.py +311 -26
- diff_diff-2.7.6/diff_diff/survey.py +1646 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/synthetic_did.py +390 -102
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/triple_diff.py +57 -25
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop.py +182 -63
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop_global.py +311 -48
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop_local.py +294 -47
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop_results.py +53 -32
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/twfe.py +16 -1
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/two_stage.py +10 -7
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/two_stage_bootstrap.py +73 -16
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/two_stage_results.py +2 -18
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/utils.py +7 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/pyproject.toml +1 -1
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/Cargo.lock +7 -7
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/Cargo.toml +1 -1
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/trop.rs +43 -8
- diff_diff-2.7.5/diff_diff/survey.py +0 -706
- {diff_diff-2.7.5 → diff_diff-2.7.6}/README.md +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_covariates.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_weights.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/honest_did.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/power.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/pretrends.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/__init__.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_common.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_continuous.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_diagnostic.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_event_study.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_power.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_staggered.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_synthetic.py +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/build.rs +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/bootstrap.rs +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/lib.rs +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/linalg.rs +0 -0
- {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/weights.rs +0 -0
|
@@ -95,8 +95,10 @@ from diff_diff.results import (
|
|
|
95
95
|
SyntheticDiDResults,
|
|
96
96
|
)
|
|
97
97
|
from diff_diff.survey import (
|
|
98
|
+
DEFFDiagnostics,
|
|
98
99
|
SurveyDesign,
|
|
99
100
|
SurveyMetadata,
|
|
101
|
+
compute_deff_diagnostics,
|
|
100
102
|
)
|
|
101
103
|
from diff_diff.staggered import (
|
|
102
104
|
CallawaySantAnna,
|
|
@@ -166,6 +168,7 @@ from diff_diff.visualization import (
|
|
|
166
168
|
plot_staircase,
|
|
167
169
|
plot_synth_weights,
|
|
168
170
|
)
|
|
171
|
+
from diff_diff.practitioner import practitioner_next_steps
|
|
169
172
|
from diff_diff.datasets import (
|
|
170
173
|
clear_cache,
|
|
171
174
|
list_datasets,
|
|
@@ -191,7 +194,7 @@ Stacked = StackedDiD
|
|
|
191
194
|
Bacon = BaconDecomposition
|
|
192
195
|
EDiD = EfficientDiD
|
|
193
196
|
|
|
194
|
-
__version__ = "2.7.
|
|
197
|
+
__version__ = "2.7.6"
|
|
195
198
|
__all__ = [
|
|
196
199
|
# Estimators
|
|
197
200
|
"DifferenceInDifferences",
|
|
@@ -327,6 +330,8 @@ __all__ = [
|
|
|
327
330
|
# Survey support
|
|
328
331
|
"SurveyDesign",
|
|
329
332
|
"SurveyMetadata",
|
|
333
|
+
"DEFFDiagnostics",
|
|
334
|
+
"compute_deff_diagnostics",
|
|
330
335
|
# Rust backend
|
|
331
336
|
"HAS_RUST_BACKEND",
|
|
332
337
|
# Linear algebra helpers
|
|
@@ -340,4 +345,6 @@ __all__ = [
|
|
|
340
345
|
"load_dataset",
|
|
341
346
|
"list_datasets",
|
|
342
347
|
"clear_cache",
|
|
348
|
+
# Practitioner guidance
|
|
349
|
+
"practitioner_next_steps",
|
|
343
350
|
]
|
|
@@ -13,7 +13,7 @@ import os
|
|
|
13
13
|
|
|
14
14
|
# Check for backend override via environment variable
|
|
15
15
|
# DIFF_DIFF_BACKEND can be: 'auto' (default), 'python', or 'rust'
|
|
16
|
-
_backend_env = os.environ.get(
|
|
16
|
+
_backend_env = os.environ.get("DIFF_DIFF_BACKEND", "auto").lower()
|
|
17
17
|
|
|
18
18
|
# Try to import Rust backend for accelerated operations
|
|
19
19
|
try:
|
|
@@ -38,6 +38,7 @@ try:
|
|
|
38
38
|
# Diagnostics
|
|
39
39
|
rust_backend_info as _rust_backend_info,
|
|
40
40
|
)
|
|
41
|
+
|
|
41
42
|
_rust_available = True
|
|
42
43
|
except ImportError:
|
|
43
44
|
_rust_available = False
|
|
@@ -61,7 +62,7 @@ except ImportError:
|
|
|
61
62
|
_rust_backend_info = None
|
|
62
63
|
|
|
63
64
|
# Determine final backend based on environment variable and availability
|
|
64
|
-
if _backend_env ==
|
|
65
|
+
if _backend_env == "python":
|
|
65
66
|
# Force pure Python mode - disable Rust even if available
|
|
66
67
|
HAS_RUST_BACKEND = False
|
|
67
68
|
_rust_bootstrap_weights = None
|
|
@@ -82,7 +83,7 @@ if _backend_env == 'python':
|
|
|
82
83
|
_rust_compute_noise_level = None
|
|
83
84
|
_rust_sc_weight_fw = None
|
|
84
85
|
_rust_backend_info = None
|
|
85
|
-
elif _backend_env ==
|
|
86
|
+
elif _backend_env == "rust":
|
|
86
87
|
# Force Rust mode - fail if not available
|
|
87
88
|
if not _rust_available:
|
|
88
89
|
raise ImportError(
|
|
@@ -111,23 +112,23 @@ def rust_backend_info():
|
|
|
111
112
|
|
|
112
113
|
|
|
113
114
|
__all__ = [
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
115
|
+
"HAS_RUST_BACKEND",
|
|
116
|
+
"rust_backend_info",
|
|
117
|
+
"_rust_bootstrap_weights",
|
|
118
|
+
"_rust_synthetic_weights",
|
|
119
|
+
"_rust_project_simplex",
|
|
120
|
+
"_rust_solve_ols",
|
|
121
|
+
"_rust_compute_robust_vcov",
|
|
121
122
|
# TROP estimator acceleration (local method)
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
123
|
+
"_rust_unit_distance_matrix",
|
|
124
|
+
"_rust_loocv_grid_search",
|
|
125
|
+
"_rust_bootstrap_trop_variance",
|
|
125
126
|
# TROP estimator acceleration (global method)
|
|
126
|
-
|
|
127
|
-
|
|
127
|
+
"_rust_loocv_grid_search_global",
|
|
128
|
+
"_rust_bootstrap_trop_variance_global",
|
|
128
129
|
# SDID weights (Frank-Wolfe matching R's synthdid)
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
130
|
+
"_rust_sdid_unit_weights",
|
|
131
|
+
"_rust_compute_time_weights",
|
|
132
|
+
"_rust_compute_noise_level",
|
|
133
|
+
"_rust_sc_weight_fw",
|
|
133
134
|
]
|
|
@@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Tuple
|
|
|
17
17
|
import numpy as np
|
|
18
18
|
import pandas as pd
|
|
19
19
|
|
|
20
|
+
from diff_diff.results import _format_survey_block
|
|
20
21
|
from diff_diff.utils import within_transform as _within_transform_util
|
|
21
22
|
|
|
22
23
|
|
|
@@ -144,23 +145,7 @@ class BaconDecompositionResults:
|
|
|
144
145
|
# Add survey design info
|
|
145
146
|
if self.survey_metadata is not None:
|
|
146
147
|
sm = self.survey_metadata
|
|
147
|
-
lines.extend(
|
|
148
|
-
[
|
|
149
|
-
"-" * 85,
|
|
150
|
-
"Survey Design".center(85),
|
|
151
|
-
"-" * 85,
|
|
152
|
-
f"{'Weight type:':<35} {sm.weight_type:>10}",
|
|
153
|
-
]
|
|
154
|
-
)
|
|
155
|
-
if sm.n_strata is not None:
|
|
156
|
-
lines.append(f"{'Strata:':<35} {sm.n_strata:>10}")
|
|
157
|
-
if sm.n_psu is not None:
|
|
158
|
-
lines.append(f"{'PSU/Cluster:':<35} {sm.n_psu:>10}")
|
|
159
|
-
lines.append(f"{'Effective sample size:':<35} {sm.effective_n:>10.1f}")
|
|
160
|
-
lines.append(f"{'Design effect (DEFF):':<35} {sm.design_effect:>10.2f}")
|
|
161
|
-
if sm.df_survey is not None:
|
|
162
|
-
lines.append(f"{'Survey d.f.:':<35} {sm.df_survey:>10}")
|
|
163
|
-
lines.extend(["-" * 85, ""])
|
|
148
|
+
lines.extend(_format_survey_block(sm, 85))
|
|
164
149
|
|
|
165
150
|
lines.extend(
|
|
166
151
|
[
|
|
@@ -477,6 +462,13 @@ class BaconDecomposition:
|
|
|
477
462
|
resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
|
|
478
463
|
_resolve_survey_for_fit(survey_design, data, "analytical")
|
|
479
464
|
)
|
|
465
|
+
# Reject replicate-weight designs — Bacon decomposition is a
|
|
466
|
+
# diagnostic that does not compute replicate-based variance
|
|
467
|
+
if resolved_survey is not None and resolved_survey.uses_replicate_variance:
|
|
468
|
+
raise NotImplementedError(
|
|
469
|
+
"BaconDecomposition does not support replicate-weight survey "
|
|
470
|
+
"designs. Use a TSL-based survey design (strata/psu/fpc)."
|
|
471
|
+
)
|
|
480
472
|
|
|
481
473
|
# Validate within-unit constancy for exact survey weights only.
|
|
482
474
|
# The exact-weight path collapses to per-unit weights via groupby().first(),
|
|
@@ -593,6 +585,13 @@ class BaconDecomposition:
|
|
|
593
585
|
weights=survey_weights,
|
|
594
586
|
)
|
|
595
587
|
|
|
588
|
+
if not comparisons:
|
|
589
|
+
raise ValueError(
|
|
590
|
+
"No valid 2x2 comparisons remain after filtering. "
|
|
591
|
+
"All cells have zero effective weight or insufficient data. "
|
|
592
|
+
"Check subpopulation/domain definition."
|
|
593
|
+
)
|
|
594
|
+
|
|
596
595
|
# Normalize weights to sum to 1
|
|
597
596
|
total_weight = sum(c.weight for c in comparisons)
|
|
598
597
|
if total_weight > 0:
|
|
@@ -849,6 +848,7 @@ class BaconDecomposition:
|
|
|
849
848
|
never_post_mask = never_mask & df[time].isin(post_periods)
|
|
850
849
|
|
|
851
850
|
# Guard against empty cells (unbalanced/filtered panels)
|
|
851
|
+
# Also check positive weight mass for survey/subpopulation designs
|
|
852
852
|
if not (
|
|
853
853
|
np.any(treated_pre_mask)
|
|
854
854
|
and np.any(treated_post_mask)
|
|
@@ -856,6 +856,13 @@ class BaconDecomposition:
|
|
|
856
856
|
and np.any(never_post_mask)
|
|
857
857
|
):
|
|
858
858
|
return None
|
|
859
|
+
if (
|
|
860
|
+
np.sum(w[treated_pre_mask]) <= 0
|
|
861
|
+
or np.sum(w[treated_post_mask]) <= 0
|
|
862
|
+
or np.sum(w[never_pre_mask]) <= 0
|
|
863
|
+
or np.sum(w[never_post_mask]) <= 0
|
|
864
|
+
):
|
|
865
|
+
return None
|
|
859
866
|
|
|
860
867
|
treated_pre = np.average(y[treated_pre_mask], weights=w[treated_pre_mask])
|
|
861
868
|
treated_post = np.average(y[treated_post_mask], weights=w[treated_post_mask])
|
|
@@ -966,7 +973,7 @@ class BaconDecomposition:
|
|
|
966
973
|
control_pre_mask = control_mask & df[time].isin(pre_periods)
|
|
967
974
|
control_post_mask = control_mask & df[time].isin(post_periods)
|
|
968
975
|
|
|
969
|
-
# Skip if any cell is empty
|
|
976
|
+
# Skip if any cell is empty or has zero effective weight
|
|
970
977
|
if (
|
|
971
978
|
treated_pre_mask.sum() == 0
|
|
972
979
|
or treated_post_mask.sum() == 0
|
|
@@ -974,6 +981,13 @@ class BaconDecomposition:
|
|
|
974
981
|
or control_post_mask.sum() == 0
|
|
975
982
|
):
|
|
976
983
|
return None
|
|
984
|
+
if (
|
|
985
|
+
np.sum(w[treated_pre_mask]) <= 0
|
|
986
|
+
or np.sum(w[treated_post_mask]) <= 0
|
|
987
|
+
or np.sum(w[control_pre_mask]) <= 0
|
|
988
|
+
or np.sum(w[control_post_mask]) <= 0
|
|
989
|
+
):
|
|
990
|
+
return None
|
|
977
991
|
|
|
978
992
|
treated_pre = np.average(y[treated_pre_mask], weights=w[treated_pre_mask])
|
|
979
993
|
treated_post = np.average(y[treated_post_mask], weights=w[treated_post_mask])
|
|
@@ -16,6 +16,9 @@ __all__ = [
|
|
|
16
16
|
"generate_bootstrap_weights",
|
|
17
17
|
"generate_bootstrap_weights_batch",
|
|
18
18
|
"generate_bootstrap_weights_batch_numpy",
|
|
19
|
+
"generate_survey_multiplier_weights_batch",
|
|
20
|
+
"generate_rao_wu_weights",
|
|
21
|
+
"generate_rao_wu_weights_batch",
|
|
19
22
|
"compute_percentile_ci",
|
|
20
23
|
"compute_bootstrap_pvalue",
|
|
21
24
|
"compute_effect_bootstrap_stats",
|
|
@@ -54,15 +57,20 @@ def generate_bootstrap_weights(
|
|
|
54
57
|
p1 = (sqrt5 + 1) / (2 * sqrt5)
|
|
55
58
|
return rng.choice([val1, val2], size=n_units, p=[p1, 1 - p1])
|
|
56
59
|
elif weight_type == "webb":
|
|
57
|
-
values = np.array(
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
60
|
+
values = np.array(
|
|
61
|
+
[
|
|
62
|
+
-np.sqrt(3 / 2),
|
|
63
|
+
-np.sqrt(2 / 2),
|
|
64
|
+
-np.sqrt(1 / 2),
|
|
65
|
+
np.sqrt(1 / 2),
|
|
66
|
+
np.sqrt(2 / 2),
|
|
67
|
+
np.sqrt(3 / 2),
|
|
68
|
+
]
|
|
69
|
+
)
|
|
61
70
|
return rng.choice(values, size=n_units)
|
|
62
71
|
else:
|
|
63
72
|
raise ValueError(
|
|
64
|
-
f"weight_type must be 'rademacher', 'mammen', or 'webb', "
|
|
65
|
-
f"got '{weight_type}'"
|
|
73
|
+
f"weight_type must be 'rademacher', 'mammen', or 'webb', " f"got '{weight_type}'"
|
|
66
74
|
)
|
|
67
75
|
|
|
68
76
|
|
|
@@ -133,15 +141,20 @@ def generate_bootstrap_weights_batch_numpy(
|
|
|
133
141
|
p1 = (sqrt5 + 1) / (2 * sqrt5)
|
|
134
142
|
return rng.choice([val1, val2], size=(n_bootstrap, n_units), p=[p1, 1 - p1])
|
|
135
143
|
elif weight_type == "webb":
|
|
136
|
-
values = np.array(
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
144
|
+
values = np.array(
|
|
145
|
+
[
|
|
146
|
+
-np.sqrt(3 / 2),
|
|
147
|
+
-np.sqrt(2 / 2),
|
|
148
|
+
-np.sqrt(1 / 2),
|
|
149
|
+
np.sqrt(1 / 2),
|
|
150
|
+
np.sqrt(2 / 2),
|
|
151
|
+
np.sqrt(3 / 2),
|
|
152
|
+
]
|
|
153
|
+
)
|
|
140
154
|
return rng.choice(values, size=(n_bootstrap, n_units))
|
|
141
155
|
else:
|
|
142
156
|
raise ValueError(
|
|
143
|
-
f"weight_type must be 'rademacher', 'mammen', or 'webb', "
|
|
144
|
-
f"got '{weight_type}'"
|
|
157
|
+
f"weight_type must be 'rademacher', 'mammen', or 'webb', " f"got '{weight_type}'"
|
|
145
158
|
)
|
|
146
159
|
|
|
147
160
|
|
|
@@ -274,9 +287,7 @@ def compute_effect_bootstrap_stats(
|
|
|
274
287
|
return np.nan, (np.nan, np.nan), np.nan
|
|
275
288
|
|
|
276
289
|
ci = compute_percentile_ci(valid_dist, alpha)
|
|
277
|
-
p_value = compute_bootstrap_pvalue(
|
|
278
|
-
original_effect, valid_dist, n_valid=len(valid_dist)
|
|
279
|
-
)
|
|
290
|
+
p_value = compute_bootstrap_pvalue(original_effect, valid_dist, n_valid=len(valid_dist))
|
|
280
291
|
return se, ci, p_value
|
|
281
292
|
|
|
282
293
|
|
|
@@ -392,8 +403,10 @@ def compute_effect_bootstrap_stats_batch(
|
|
|
392
403
|
if np.any(partial_valid):
|
|
393
404
|
for j in np.where(partial_valid)[0]:
|
|
394
405
|
se, ci, pv = compute_effect_bootstrap_stats(
|
|
395
|
-
original_effects[j],
|
|
396
|
-
|
|
406
|
+
original_effects[j],
|
|
407
|
+
bootstrap_matrix[:, j],
|
|
408
|
+
alpha=alpha,
|
|
409
|
+
context=f"effect {j}",
|
|
397
410
|
)
|
|
398
411
|
ses[j] = se
|
|
399
412
|
ci_lowers[j] = ci[0]
|
|
@@ -401,3 +414,247 @@ def compute_effect_bootstrap_stats_batch(
|
|
|
401
414
|
p_values[j] = pv
|
|
402
415
|
|
|
403
416
|
return ses, ci_lowers, ci_uppers, p_values
|
|
417
|
+
|
|
418
|
+
|
|
419
|
+
# ---------------------------------------------------------------------------
|
|
420
|
+
# Survey-aware bootstrap weight generators
|
|
421
|
+
# ---------------------------------------------------------------------------
|
|
422
|
+
|
|
423
|
+
|
|
424
|
+
def generate_survey_multiplier_weights_batch(
|
|
425
|
+
n_bootstrap: int,
|
|
426
|
+
resolved_survey: "ResolvedSurveyDesign",
|
|
427
|
+
weight_type: str,
|
|
428
|
+
rng: np.random.Generator,
|
|
429
|
+
) -> Tuple[np.ndarray, np.ndarray]:
|
|
430
|
+
"""Generate PSU-level multiplier weights for survey-aware bootstrap.
|
|
431
|
+
|
|
432
|
+
Within each stratum, weights are generated independently. When FPC
|
|
433
|
+
is present, weights are scaled by ``sqrt(1 - f_h)`` per stratum so
|
|
434
|
+
the bootstrap variance matches the TSL variance.
|
|
435
|
+
|
|
436
|
+
Parameters
|
|
437
|
+
----------
|
|
438
|
+
n_bootstrap : int
|
|
439
|
+
Number of bootstrap iterations.
|
|
440
|
+
resolved_survey : ResolvedSurveyDesign
|
|
441
|
+
Resolved survey design.
|
|
442
|
+
weight_type : str
|
|
443
|
+
Multiplier distribution: ``"rademacher"``, ``"mammen"``, or ``"webb"``.
|
|
444
|
+
rng : np.random.Generator
|
|
445
|
+
Random number generator.
|
|
446
|
+
|
|
447
|
+
Returns
|
|
448
|
+
-------
|
|
449
|
+
weights : np.ndarray
|
|
450
|
+
Multiplier weights, shape ``(n_bootstrap, n_psu)``.
|
|
451
|
+
psu_ids : np.ndarray
|
|
452
|
+
Unique PSU identifiers aligned to columns of *weights*.
|
|
453
|
+
"""
|
|
454
|
+
psu = resolved_survey.psu
|
|
455
|
+
strata = resolved_survey.strata
|
|
456
|
+
|
|
457
|
+
if resolved_survey.lonely_psu == "adjust":
|
|
458
|
+
raise NotImplementedError(
|
|
459
|
+
"lonely_psu='adjust' is not yet supported for survey-aware bootstrap. "
|
|
460
|
+
"Use lonely_psu='remove' or 'certainty', or use analytical inference."
|
|
461
|
+
)
|
|
462
|
+
|
|
463
|
+
if psu is None:
|
|
464
|
+
# Each observation is its own PSU
|
|
465
|
+
n_psu = len(resolved_survey.weights)
|
|
466
|
+
psu_ids = np.arange(n_psu)
|
|
467
|
+
else:
|
|
468
|
+
psu_ids = np.unique(psu)
|
|
469
|
+
n_psu = len(psu_ids)
|
|
470
|
+
|
|
471
|
+
if strata is None:
|
|
472
|
+
# No stratification — generate a single block of weights
|
|
473
|
+
if n_psu < 2:
|
|
474
|
+
# Single PSU — variance unidentified (matches compute_survey_vcov)
|
|
475
|
+
weights = np.zeros((n_bootstrap, n_psu), dtype=np.float64)
|
|
476
|
+
return weights, psu_ids
|
|
477
|
+
weights = generate_bootstrap_weights_batch(n_bootstrap, n_psu, weight_type, rng)
|
|
478
|
+
# FPC scaling (unstratified)
|
|
479
|
+
if resolved_survey.fpc is not None:
|
|
480
|
+
if psu is not None:
|
|
481
|
+
n_units_for_fpc = n_psu
|
|
482
|
+
else:
|
|
483
|
+
n_units_for_fpc = len(resolved_survey.weights)
|
|
484
|
+
if resolved_survey.fpc[0] < n_units_for_fpc:
|
|
485
|
+
raise ValueError(
|
|
486
|
+
f"FPC ({resolved_survey.fpc[0]}) is less than the number of PSUs "
|
|
487
|
+
f"({n_units_for_fpc}). FPC must be >= number of PSUs."
|
|
488
|
+
)
|
|
489
|
+
f = n_units_for_fpc / resolved_survey.fpc[0]
|
|
490
|
+
if f < 1.0:
|
|
491
|
+
weights = weights * np.sqrt(1.0 - f)
|
|
492
|
+
else:
|
|
493
|
+
weights = np.zeros_like(weights)
|
|
494
|
+
else:
|
|
495
|
+
# Stratified — generate independently within strata
|
|
496
|
+
weights = np.empty((n_bootstrap, n_psu), dtype=np.float64)
|
|
497
|
+
|
|
498
|
+
# Build PSU → column-index map
|
|
499
|
+
psu_to_col = {int(p): i for i, p in enumerate(psu_ids)}
|
|
500
|
+
|
|
501
|
+
unique_strata = np.unique(strata)
|
|
502
|
+
for h in unique_strata:
|
|
503
|
+
mask_h = strata == h
|
|
504
|
+
|
|
505
|
+
if psu is not None:
|
|
506
|
+
psus_in_h = np.unique(psu[mask_h])
|
|
507
|
+
else:
|
|
508
|
+
psus_in_h = np.where(mask_h)[0]
|
|
509
|
+
|
|
510
|
+
n_h = len(psus_in_h)
|
|
511
|
+
cols = np.array([psu_to_col[int(p)] for p in psus_in_h])
|
|
512
|
+
|
|
513
|
+
if n_h < 2:
|
|
514
|
+
# Lonely PSU — zero weight (matches remove/certainty behavior)
|
|
515
|
+
weights[:, cols] = 0.0
|
|
516
|
+
continue
|
|
517
|
+
|
|
518
|
+
# Generate weights for this stratum
|
|
519
|
+
stratum_weights = generate_bootstrap_weights_batch_numpy(
|
|
520
|
+
n_bootstrap, n_h, weight_type, rng
|
|
521
|
+
)
|
|
522
|
+
|
|
523
|
+
# FPC scaling
|
|
524
|
+
if resolved_survey.fpc is not None:
|
|
525
|
+
N_h = resolved_survey.fpc[mask_h][0]
|
|
526
|
+
if N_h < n_h:
|
|
527
|
+
raise ValueError(
|
|
528
|
+
f"FPC ({N_h}) is less than the number of PSUs "
|
|
529
|
+
f"({n_h}) in stratum {h}. FPC must be >= n_PSU."
|
|
530
|
+
)
|
|
531
|
+
f_h = n_h / N_h
|
|
532
|
+
if f_h < 1.0:
|
|
533
|
+
stratum_weights = stratum_weights * np.sqrt(1.0 - f_h)
|
|
534
|
+
else:
|
|
535
|
+
stratum_weights = np.zeros_like(stratum_weights)
|
|
536
|
+
|
|
537
|
+
weights[:, cols] = stratum_weights
|
|
538
|
+
|
|
539
|
+
return weights, psu_ids
|
|
540
|
+
|
|
541
|
+
|
|
542
|
+
def generate_rao_wu_weights(
|
|
543
|
+
resolved_survey: "ResolvedSurveyDesign",
|
|
544
|
+
rng: np.random.Generator,
|
|
545
|
+
) -> np.ndarray:
|
|
546
|
+
"""Generate one set of Rao-Wu (1988) rescaled observation weights.
|
|
547
|
+
|
|
548
|
+
Within each stratum *h* with *n_h* PSUs, draw ``m_h`` PSUs with
|
|
549
|
+
replacement and rescale observation weights by ``(n_h / m_h) * r_hi``
|
|
550
|
+
where ``r_hi`` is the count of PSU *i* being selected.
|
|
551
|
+
|
|
552
|
+
Without FPC: ``m_h = n_h - 1``.
|
|
553
|
+
With FPC: ``m_h = max(1, round((1 - f_h) * (n_h - 1)))``
|
|
554
|
+
(Rao, Wu & Yue 1992, Section 3).
|
|
555
|
+
|
|
556
|
+
Parameters
|
|
557
|
+
----------
|
|
558
|
+
resolved_survey : ResolvedSurveyDesign
|
|
559
|
+
Resolved survey design.
|
|
560
|
+
rng : np.random.Generator
|
|
561
|
+
Random number generator.
|
|
562
|
+
|
|
563
|
+
Returns
|
|
564
|
+
-------
|
|
565
|
+
np.ndarray
|
|
566
|
+
Rescaled observation weights, shape ``(n_obs,)``.
|
|
567
|
+
"""
|
|
568
|
+
n_obs = len(resolved_survey.weights)
|
|
569
|
+
base_weights = resolved_survey.weights
|
|
570
|
+
psu = resolved_survey.psu
|
|
571
|
+
strata = resolved_survey.strata
|
|
572
|
+
|
|
573
|
+
if resolved_survey.lonely_psu == "adjust":
|
|
574
|
+
raise NotImplementedError(
|
|
575
|
+
"lonely_psu='adjust' is not yet supported for survey-aware bootstrap. "
|
|
576
|
+
"Use lonely_psu='remove' or 'certainty', or use analytical inference."
|
|
577
|
+
)
|
|
578
|
+
|
|
579
|
+
rescaled = np.zeros(n_obs, dtype=np.float64)
|
|
580
|
+
|
|
581
|
+
if psu is None:
|
|
582
|
+
obs_psu = np.arange(n_obs)
|
|
583
|
+
else:
|
|
584
|
+
obs_psu = psu
|
|
585
|
+
|
|
586
|
+
if strata is None:
|
|
587
|
+
strata_masks = [np.ones(n_obs, dtype=bool)]
|
|
588
|
+
else:
|
|
589
|
+
unique_strata = np.unique(strata)
|
|
590
|
+
strata_masks = [strata == h for h in unique_strata]
|
|
591
|
+
|
|
592
|
+
for mask_h in strata_masks:
|
|
593
|
+
psu_h = obs_psu[mask_h]
|
|
594
|
+
unique_psu_h = np.unique(psu_h)
|
|
595
|
+
n_h = len(unique_psu_h)
|
|
596
|
+
|
|
597
|
+
if n_h < 2:
|
|
598
|
+
# Census / lonely PSU — keep original weights (zero variance)
|
|
599
|
+
rescaled[mask_h] = base_weights[mask_h]
|
|
600
|
+
continue
|
|
601
|
+
|
|
602
|
+
# Compute resample size
|
|
603
|
+
if resolved_survey.fpc is not None:
|
|
604
|
+
N_h = resolved_survey.fpc[mask_h][0]
|
|
605
|
+
if N_h < n_h:
|
|
606
|
+
raise ValueError(
|
|
607
|
+
f"FPC ({N_h}) is less than the number of PSUs "
|
|
608
|
+
f"({n_h}). FPC must be >= number of PSUs."
|
|
609
|
+
)
|
|
610
|
+
f_h = n_h / N_h
|
|
611
|
+
if f_h >= 1.0:
|
|
612
|
+
# Census stratum — keep original weights (zero variance)
|
|
613
|
+
rescaled[mask_h] = base_weights[mask_h]
|
|
614
|
+
continue
|
|
615
|
+
m_h = max(1, round((1.0 - f_h) * (n_h - 1)))
|
|
616
|
+
else:
|
|
617
|
+
m_h = n_h - 1
|
|
618
|
+
|
|
619
|
+
# Draw m_h PSUs with replacement
|
|
620
|
+
drawn_indices = rng.choice(n_h, size=m_h, replace=True)
|
|
621
|
+
counts = np.bincount(drawn_indices, minlength=n_h)
|
|
622
|
+
|
|
623
|
+
# Rescale factor per PSU: (n_h / m_h) * r_hi
|
|
624
|
+
scale_per_psu = (n_h / m_h) * counts.astype(np.float64)
|
|
625
|
+
|
|
626
|
+
# Map PSU → local index for vectorized application
|
|
627
|
+
psu_to_local = {int(p): i for i, p in enumerate(unique_psu_h)}
|
|
628
|
+
obs_in_h = np.where(mask_h)[0]
|
|
629
|
+
local_indices = np.array([psu_to_local[int(obs_psu[idx])] for idx in obs_in_h])
|
|
630
|
+
rescaled[obs_in_h] = base_weights[obs_in_h] * scale_per_psu[local_indices]
|
|
631
|
+
|
|
632
|
+
return rescaled
|
|
633
|
+
|
|
634
|
+
|
|
635
|
+
def generate_rao_wu_weights_batch(
|
|
636
|
+
n_bootstrap: int,
|
|
637
|
+
resolved_survey: "ResolvedSurveyDesign",
|
|
638
|
+
rng: np.random.Generator,
|
|
639
|
+
) -> np.ndarray:
|
|
640
|
+
"""Generate multiple sets of Rao-Wu rescaled weights.
|
|
641
|
+
|
|
642
|
+
Parameters
|
|
643
|
+
----------
|
|
644
|
+
n_bootstrap : int
|
|
645
|
+
Number of bootstrap iterations.
|
|
646
|
+
resolved_survey : ResolvedSurveyDesign
|
|
647
|
+
Resolved survey design.
|
|
648
|
+
rng : np.random.Generator
|
|
649
|
+
Random number generator.
|
|
650
|
+
|
|
651
|
+
Returns
|
|
652
|
+
-------
|
|
653
|
+
np.ndarray
|
|
654
|
+
Rescaled weights, shape ``(n_bootstrap, n_obs)``.
|
|
655
|
+
"""
|
|
656
|
+
n_obs = len(resolved_survey.weights)
|
|
657
|
+
result = np.empty((n_bootstrap, n_obs), dtype=np.float64)
|
|
658
|
+
for b in range(n_bootstrap):
|
|
659
|
+
result[b] = generate_rao_wu_weights(resolved_survey, rng)
|
|
660
|
+
return result
|