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.
Files changed (65) hide show
  1. {diff_diff-2.7.5 → diff_diff-2.7.6}/PKG-INFO +1 -1
  2. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/__init__.py +8 -1
  3. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/_backend.py +20 -19
  4. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/bacon.py +32 -18
  5. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/bootstrap_utils.py +274 -17
  6. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/continuous_did.py +235 -80
  7. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/continuous_did_bspline.py +7 -5
  8. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/continuous_did_results.py +2 -18
  9. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/datasets.py +179 -72
  10. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/diagnostics.py +62 -101
  11. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did.py +50 -18
  12. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_bootstrap.py +29 -4
  13. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_results.py +2 -18
  14. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/estimators.py +39 -24
  15. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/imputation.py +15 -8
  16. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/imputation_bootstrap.py +79 -9
  17. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/imputation_results.py +2 -18
  18. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/linalg.py +236 -28
  19. diff_diff-2.7.6/diff_diff/practitioner.py +869 -0
  20. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/prep.py +75 -92
  21. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/prep_dgp.py +69 -60
  22. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/results.py +53 -36
  23. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/stacked_did.py +9 -0
  24. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/stacked_did_results.py +2 -18
  25. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered.py +173 -53
  26. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered_aggregation.py +56 -6
  27. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered_bootstrap.py +186 -93
  28. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/staggered_results.py +2 -18
  29. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/sun_abraham.py +311 -26
  30. diff_diff-2.7.6/diff_diff/survey.py +1646 -0
  31. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/synthetic_did.py +390 -102
  32. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/triple_diff.py +57 -25
  33. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop.py +182 -63
  34. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop_global.py +311 -48
  35. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop_local.py +294 -47
  36. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/trop_results.py +53 -32
  37. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/twfe.py +16 -1
  38. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/two_stage.py +10 -7
  39. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/two_stage_bootstrap.py +73 -16
  40. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/two_stage_results.py +2 -18
  41. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/utils.py +7 -0
  42. {diff_diff-2.7.5 → diff_diff-2.7.6}/pyproject.toml +1 -1
  43. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/Cargo.lock +7 -7
  44. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/Cargo.toml +1 -1
  45. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/trop.rs +43 -8
  46. diff_diff-2.7.5/diff_diff/survey.py +0 -706
  47. {diff_diff-2.7.5 → diff_diff-2.7.6}/README.md +0 -0
  48. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_covariates.py +0 -0
  49. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/efficient_did_weights.py +0 -0
  50. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/honest_did.py +0 -0
  51. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/power.py +0 -0
  52. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/pretrends.py +0 -0
  53. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/__init__.py +0 -0
  54. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_common.py +0 -0
  55. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_continuous.py +0 -0
  56. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_diagnostic.py +0 -0
  57. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_event_study.py +0 -0
  58. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_power.py +0 -0
  59. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_staggered.py +0 -0
  60. {diff_diff-2.7.5 → diff_diff-2.7.6}/diff_diff/visualization/_synthetic.py +0 -0
  61. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/build.rs +0 -0
  62. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/bootstrap.rs +0 -0
  63. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/lib.rs +0 -0
  64. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/linalg.rs +0 -0
  65. {diff_diff-2.7.5 → diff_diff-2.7.6}/rust/src/weights.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diff-diff
3
- Version: 2.7.5
3
+ Version: 2.7.6
4
4
  Classifier: Development Status :: 5 - Production/Stable
5
5
  Classifier: Intended Audience :: Science/Research
6
6
  Classifier: Operating System :: OS Independent
@@ -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.5"
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('DIFF_DIFF_BACKEND', 'auto').lower()
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 == 'python':
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 == 'rust':
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
- 'HAS_RUST_BACKEND',
115
- 'rust_backend_info',
116
- '_rust_bootstrap_weights',
117
- '_rust_synthetic_weights',
118
- '_rust_project_simplex',
119
- '_rust_solve_ols',
120
- '_rust_compute_robust_vcov',
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
- '_rust_unit_distance_matrix',
123
- '_rust_loocv_grid_search',
124
- '_rust_bootstrap_trop_variance',
123
+ "_rust_unit_distance_matrix",
124
+ "_rust_loocv_grid_search",
125
+ "_rust_bootstrap_trop_variance",
125
126
  # TROP estimator acceleration (global method)
126
- '_rust_loocv_grid_search_global',
127
- '_rust_bootstrap_trop_variance_global',
127
+ "_rust_loocv_grid_search_global",
128
+ "_rust_bootstrap_trop_variance_global",
128
129
  # SDID weights (Frank-Wolfe matching R's synthdid)
129
- '_rust_sdid_unit_weights',
130
- '_rust_compute_time_weights',
131
- '_rust_compute_noise_level',
132
- '_rust_sc_weight_fw',
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
- -np.sqrt(3 / 2), -np.sqrt(2 / 2), -np.sqrt(1 / 2),
59
- np.sqrt(1 / 2), np.sqrt(2 / 2), np.sqrt(3 / 2)
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
- -np.sqrt(3 / 2), -np.sqrt(2 / 2), -np.sqrt(1 / 2),
138
- np.sqrt(1 / 2), np.sqrt(2 / 2), np.sqrt(3 / 2)
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], bootstrap_matrix[:, j], alpha=alpha,
396
- context=f"effect {j}"
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