diff-diff 2.7.5__tar.gz → 2.8.0__tar.gz

{diff_diff-2.7.5 → diff_diff-2.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 2.7.5
+Version: 2.8.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
@@ -45,6 +45,7 @@ Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
 Project-URL: Documentation, https://diff-diff.readthedocs.io
 Project-URL: Homepage, https://github.com/igerber/diff-diff
 Project-URL: Issues, https://github.com/igerber/diff-diff/issues
+Project-URL: Practitioner Guide, https://github.com/igerber/diff-diff/blob/main/docs/llms-practitioner.txt
 Project-URL: Repository, https://github.com/igerber/diff-diff
 
 # diff-diff
@@ -116,6 +117,14 @@ Signif. codes: '***' 0.001, '**' 0.01, '*' 0.05, '.' 0.1
 ======================================================================
 ```
 
+## For AI Agents
+
+If you are an AI agent or LLM using this library, read [`docs/llms.txt`](docs/llms.txt) for a concise API reference with an 8-step practitioner workflow (based on Baker et al. 2025). The workflow ensures rigorous DiD analysis — not just calling `fit()`, but testing assumptions, running sensitivity analysis, and checking robustness.
+
+After estimation, call `practitioner_next_steps(results)` for context-aware guidance on remaining diagnostic steps.
+
+Detailed guide: [`docs/llms-practitioner.txt`](docs/llms-practitioner.txt)
+
 ## Features
 
 - **sklearn-like API**: Familiar `fit()` interface with `get_params()` and `set_params()`

{diff_diff-2.7.5 → diff_diff-2.8.0}/README.md

@@ -67,6 +67,14 @@ Signif. codes: '***' 0.001, '**' 0.01, '*' 0.05, '.' 0.1
 ======================================================================
 ```
 
+## For AI Agents
+
+If you are an AI agent or LLM using this library, read [`docs/llms.txt`](docs/llms.txt) for a concise API reference with an 8-step practitioner workflow (based on Baker et al. 2025). The workflow ensures rigorous DiD analysis — not just calling `fit()`, but testing assumptions, running sensitivity analysis, and checking robustness.
+
+After estimation, call `practitioner_next_steps(results)` for context-aware guidance on remaining diagnostic steps.
+
+Detailed guide: [`docs/llms-practitioner.txt`](docs/llms-practitioner.txt)
+
 ## Features
 
 - **sklearn-like API**: Familiar `fit()` interface with `get_params()` and `set_params()`

{diff_diff-2.7.5 → diff_diff-2.8.0}/diff_diff/__init__.py

@@ -3,6 +3,13 @@ diff-diff: A library for Difference-in-Differences analysis.
 
 This library provides sklearn-like estimators for causal inference
 using the difference-in-differences methodology.
+
+For rigorous analysis, follow the 8-step practitioner workflow in
+docs/llms-practitioner.txt (based on Baker et al. 2025). After
+estimation, call ``practitioner_next_steps(results)`` for context-aware
+guidance on remaining diagnostic steps.
+
+AI agent reference: docs/llms.txt
 """
 
 # Import backend detection from dedicated module (avoids circular imports)
@@ -81,6 +88,7 @@ from diff_diff.prep import (
     generate_factor_data,
     generate_panel_data,
     generate_staggered_data,
+    generate_staggered_ddd_data,
     make_post_indicator,
     make_treatment_indicator,
     rank_control_units,
@@ -95,8 +103,10 @@ from diff_diff.results import (
     SyntheticDiDResults,
 )
 from diff_diff.survey import (
+    DEFFDiagnostics,
     SurveyDesign,
     SurveyMetadata,
+    compute_deff_diagnostics,
 )
 from diff_diff.staggered import (
     CallawaySantAnna,
@@ -131,6 +141,12 @@ from diff_diff.triple_diff import (
     TripleDifferenceResults,
     triple_difference,
 )
+from diff_diff.staggered_triple_diff import (
+    StaggeredTripleDifference,
+)
+from diff_diff.staggered_triple_diff_results import (
+    StaggeredTripleDiffResults,
+)
 from diff_diff.continuous_did import (
     ContinuousDiD,
     ContinuousDiDResults,
@@ -166,6 +182,7 @@ from diff_diff.visualization import (
     plot_staircase,
     plot_synth_weights,
 )
+from diff_diff.practitioner import practitioner_next_steps
 from diff_diff.datasets import (
     clear_cache,
     list_datasets,
@@ -187,11 +204,12 @@ SA = SunAbraham
 BJS = ImputationDiD
 Gardner = TwoStageDiD
 DDD = TripleDifference
+SDDD = StaggeredTripleDifference
 Stacked = StackedDiD
 Bacon = BaconDecomposition
 EDiD = EfficientDiD
 
-__version__ = "2.7.5"
+__version__ = "2.8.0"
 __all__ = [
     # Estimators
     "DifferenceInDifferences",
@@ -217,6 +235,7 @@ __all__ = [
     "BJS",
     "Gardner",
     "DDD",
+    "SDDD",
     "Stacked",
     "Bacon",
     # Bacon Decomposition
@@ -244,6 +263,8 @@ __all__ = [
     "two_stage_did",
     "TripleDifferenceResults",
     "triple_difference",
+    "StaggeredTripleDifference",
+    "StaggeredTripleDiffResults",
     "TROPResults",
     "trop",
     "StackedDiDResults",
@@ -293,6 +314,7 @@ __all__ = [
     "generate_ddd_data",
     "generate_panel_data",
     "generate_event_study_data",
+    "generate_staggered_ddd_data",
     "generate_continuous_did_data",
     "create_event_time",
     "aggregate_to_cohorts",
@@ -327,6 +349,8 @@ __all__ = [
     # Survey support
     "SurveyDesign",
     "SurveyMetadata",
+    "DEFFDiagnostics",
+    "compute_deff_diagnostics",
     # Rust backend
     "HAS_RUST_BACKEND",
     # Linear algebra helpers
@@ -340,4 +364,6 @@ __all__ = [
     "load_dataset",
     "list_datasets",
     "clear_cache",
+    # Practitioner guidance
+    "practitioner_next_steps",
 ]

{diff_diff-2.7.5 → diff_diff-2.8.0}/diff_diff/_backend.py

@@ -13,7 +13,7 @@ import os
 
 # Check for backend override via environment variable
 # DIFF_DIFF_BACKEND can be: 'auto' (default), 'python', or 'rust'
-_backend_env = os.environ.get('DIFF_DIFF_BACKEND', 'auto').lower()
+_backend_env = os.environ.get("DIFF_DIFF_BACKEND", "auto").lower()
 
 # Try to import Rust backend for accelerated operations
 try:
@@ -38,6 +38,7 @@ try:
         # Diagnostics
         rust_backend_info as _rust_backend_info,
     )
+
     _rust_available = True
 except ImportError:
     _rust_available = False
@@ -61,7 +62,7 @@ except ImportError:
     _rust_backend_info = None
 
 # Determine final backend based on environment variable and availability
-if _backend_env == 'python':
+if _backend_env == "python":
     # Force pure Python mode - disable Rust even if available
     HAS_RUST_BACKEND = False
     _rust_bootstrap_weights = None
@@ -82,7 +83,7 @@ if _backend_env == 'python':
     _rust_compute_noise_level = None
     _rust_sc_weight_fw = None
     _rust_backend_info = None
-elif _backend_env == 'rust':
+elif _backend_env == "rust":
     # Force Rust mode - fail if not available
     if not _rust_available:
         raise ImportError(
@@ -111,23 +112,23 @@ def rust_backend_info():
 
 
 __all__ = [
-    'HAS_RUST_BACKEND',
-    'rust_backend_info',
-    '_rust_bootstrap_weights',
-    '_rust_synthetic_weights',
-    '_rust_project_simplex',
-    '_rust_solve_ols',
-    '_rust_compute_robust_vcov',
+    "HAS_RUST_BACKEND",
+    "rust_backend_info",
+    "_rust_bootstrap_weights",
+    "_rust_synthetic_weights",
+    "_rust_project_simplex",
+    "_rust_solve_ols",
+    "_rust_compute_robust_vcov",
     # TROP estimator acceleration (local method)
-    '_rust_unit_distance_matrix',
-    '_rust_loocv_grid_search',
-    '_rust_bootstrap_trop_variance',
+    "_rust_unit_distance_matrix",
+    "_rust_loocv_grid_search",
+    "_rust_bootstrap_trop_variance",
     # TROP estimator acceleration (global method)
-    '_rust_loocv_grid_search_global',
-    '_rust_bootstrap_trop_variance_global',
+    "_rust_loocv_grid_search_global",
+    "_rust_bootstrap_trop_variance_global",
     # SDID weights (Frank-Wolfe matching R's synthdid)
-    '_rust_sdid_unit_weights',
-    '_rust_compute_time_weights',
-    '_rust_compute_noise_level',
-    '_rust_sc_weight_fw',
+    "_rust_sdid_unit_weights",
+    "_rust_compute_time_weights",
+    "_rust_compute_noise_level",
+    "_rust_sc_weight_fw",
 ]

{diff_diff-2.7.5 → diff_diff-2.8.0}/diff_diff/bacon.py

@@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 import pandas as pd
 
+from diff_diff.results import _format_survey_block
 from diff_diff.utils import within_transform as _within_transform_util
 
 
@@ -144,23 +145,7 @@ class BaconDecompositionResults:
         # Add survey design info
         if self.survey_metadata is not None:
             sm = self.survey_metadata
-            lines.extend(
-                [
-                    "-" * 85,
-                    "Survey Design".center(85),
-                    "-" * 85,
-                    f"{'Weight type:':<35} {sm.weight_type:>10}",
-                ]
-            )
-            if sm.n_strata is not None:
-                lines.append(f"{'Strata:':<35} {sm.n_strata:>10}")
-            if sm.n_psu is not None:
-                lines.append(f"{'PSU/Cluster:':<35} {sm.n_psu:>10}")
-            lines.append(f"{'Effective sample size:':<35} {sm.effective_n:>10.1f}")
-            lines.append(f"{'Design effect (DEFF):':<35} {sm.design_effect:>10.2f}")
-            if sm.df_survey is not None:
-                lines.append(f"{'Survey d.f.:':<35} {sm.df_survey:>10}")
-            lines.extend(["-" * 85, ""])
+            lines.extend(_format_survey_block(sm, 85))
 
         lines.extend(
             [
@@ -477,6 +462,13 @@ class BaconDecomposition:
         resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
             _resolve_survey_for_fit(survey_design, data, "analytical")
         )
+        # Reject replicate-weight designs — Bacon decomposition is a
+        # diagnostic that does not compute replicate-based variance
+        if resolved_survey is not None and resolved_survey.uses_replicate_variance:
+            raise NotImplementedError(
+                "BaconDecomposition does not support replicate-weight survey "
+                "designs. Use a TSL-based survey design (strata/psu/fpc)."
+            )
 
         # Validate within-unit constancy for exact survey weights only.
         # The exact-weight path collapses to per-unit weights via groupby().first(),
@@ -593,6 +585,13 @@ class BaconDecomposition:
                 weights=survey_weights,
             )
 
+        if not comparisons:
+            raise ValueError(
+                "No valid 2x2 comparisons remain after filtering. "
+                "All cells have zero effective weight or insufficient data. "
+                "Check subpopulation/domain definition."
+            )
+
         # Normalize weights to sum to 1
         total_weight = sum(c.weight for c in comparisons)
         if total_weight > 0:
@@ -849,6 +848,7 @@ class BaconDecomposition:
         never_post_mask = never_mask & df[time].isin(post_periods)
 
         # Guard against empty cells (unbalanced/filtered panels)
+        # Also check positive weight mass for survey/subpopulation designs
         if not (
             np.any(treated_pre_mask)
             and np.any(treated_post_mask)
@@ -856,6 +856,13 @@ class BaconDecomposition:
             and np.any(never_post_mask)
         ):
             return None
+        if (
+            np.sum(w[treated_pre_mask]) <= 0
+            or np.sum(w[treated_post_mask]) <= 0
+            or np.sum(w[never_pre_mask]) <= 0
+            or np.sum(w[never_post_mask]) <= 0
+        ):
+            return None
 
         treated_pre = np.average(y[treated_pre_mask], weights=w[treated_pre_mask])
         treated_post = np.average(y[treated_post_mask], weights=w[treated_post_mask])
@@ -966,7 +973,7 @@ class BaconDecomposition:
         control_pre_mask = control_mask & df[time].isin(pre_periods)
         control_post_mask = control_mask & df[time].isin(post_periods)
 
-        # Skip if any cell is empty
+        # Skip if any cell is empty or has zero effective weight
         if (
             treated_pre_mask.sum() == 0
             or treated_post_mask.sum() == 0
@@ -974,6 +981,13 @@ class BaconDecomposition:
             or control_post_mask.sum() == 0
         ):
             return None
+        if (
+            np.sum(w[treated_pre_mask]) <= 0
+            or np.sum(w[treated_post_mask]) <= 0
+            or np.sum(w[control_pre_mask]) <= 0
+            or np.sum(w[control_post_mask]) <= 0
+        ):
+            return None
 
         treated_pre = np.average(y[treated_pre_mask], weights=w[treated_pre_mask])
         treated_post = np.average(y[treated_post_mask], weights=w[treated_post_mask])

{diff_diff-2.7.5 → diff_diff-2.8.0}/diff_diff/bootstrap_utils.py

@@ -16,6 +16,9 @@ __all__ = [
     "generate_bootstrap_weights",
     "generate_bootstrap_weights_batch",
     "generate_bootstrap_weights_batch_numpy",
+    "generate_survey_multiplier_weights_batch",
+    "generate_rao_wu_weights",
+    "generate_rao_wu_weights_batch",
     "compute_percentile_ci",
     "compute_bootstrap_pvalue",
     "compute_effect_bootstrap_stats",
@@ -54,15 +57,20 @@ def generate_bootstrap_weights(
         p1 = (sqrt5 + 1) / (2 * sqrt5)
         return rng.choice([val1, val2], size=n_units, p=[p1, 1 - p1])
     elif weight_type == "webb":
-        values = np.array([
-            -np.sqrt(3 / 2), -np.sqrt(2 / 2), -np.sqrt(1 / 2),
-            np.sqrt(1 / 2), np.sqrt(2 / 2), np.sqrt(3 / 2)
-        ])
+        values = np.array(
+            [
+                -np.sqrt(3 / 2),
+                -np.sqrt(2 / 2),
+                -np.sqrt(1 / 2),
+                np.sqrt(1 / 2),
+                np.sqrt(2 / 2),
+                np.sqrt(3 / 2),
+            ]
+        )
         return rng.choice(values, size=n_units)
     else:
         raise ValueError(
-            f"weight_type must be 'rademacher', 'mammen', or 'webb', "
-            f"got '{weight_type}'"
+            f"weight_type must be 'rademacher', 'mammen', or 'webb', " f"got '{weight_type}'"
         )
 
 
@@ -133,15 +141,20 @@ def generate_bootstrap_weights_batch_numpy(
         p1 = (sqrt5 + 1) / (2 * sqrt5)
         return rng.choice([val1, val2], size=(n_bootstrap, n_units), p=[p1, 1 - p1])
     elif weight_type == "webb":
-        values = np.array([
-            -np.sqrt(3 / 2), -np.sqrt(2 / 2), -np.sqrt(1 / 2),
-            np.sqrt(1 / 2), np.sqrt(2 / 2), np.sqrt(3 / 2)
-        ])
+        values = np.array(
+            [
+                -np.sqrt(3 / 2),
+                -np.sqrt(2 / 2),
+                -np.sqrt(1 / 2),
+                np.sqrt(1 / 2),
+                np.sqrt(2 / 2),
+                np.sqrt(3 / 2),
+            ]
+        )
         return rng.choice(values, size=(n_bootstrap, n_units))
     else:
         raise ValueError(
-            f"weight_type must be 'rademacher', 'mammen', or 'webb', "
-            f"got '{weight_type}'"
+            f"weight_type must be 'rademacher', 'mammen', or 'webb', " f"got '{weight_type}'"
         )
 
 
@@ -274,9 +287,7 @@ def compute_effect_bootstrap_stats(
         return np.nan, (np.nan, np.nan), np.nan
 
     ci = compute_percentile_ci(valid_dist, alpha)
-    p_value = compute_bootstrap_pvalue(
-        original_effect, valid_dist, n_valid=len(valid_dist)
-    )
+    p_value = compute_bootstrap_pvalue(original_effect, valid_dist, n_valid=len(valid_dist))
     return se, ci, p_value
 
 
@@ -392,8 +403,10 @@ def compute_effect_bootstrap_stats_batch(
     if np.any(partial_valid):
         for j in np.where(partial_valid)[0]:
             se, ci, pv = compute_effect_bootstrap_stats(
-                original_effects[j], bootstrap_matrix[:, j], alpha=alpha,
-                context=f"effect {j}"
+                original_effects[j],
+                bootstrap_matrix[:, j],
+                alpha=alpha,
+                context=f"effect {j}",
             )
             ses[j] = se
             ci_lowers[j] = ci[0]
@@ -401,3 +414,247 @@ def compute_effect_bootstrap_stats_batch(
             p_values[j] = pv
 
     return ses, ci_lowers, ci_uppers, p_values
+
+
+# ---------------------------------------------------------------------------
+# Survey-aware bootstrap weight generators
+# ---------------------------------------------------------------------------
+
+
+def generate_survey_multiplier_weights_batch(
+    n_bootstrap: int,
+    resolved_survey: "ResolvedSurveyDesign",
+    weight_type: str,
+    rng: np.random.Generator,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Generate PSU-level multiplier weights for survey-aware bootstrap.
+
+    Within each stratum, weights are generated independently.  When FPC
+    is present, weights are scaled by ``sqrt(1 - f_h)`` per stratum so
+    the bootstrap variance matches the TSL variance.
+
+    Parameters
+    ----------
+    n_bootstrap : int
+        Number of bootstrap iterations.
+    resolved_survey : ResolvedSurveyDesign
+        Resolved survey design.
+    weight_type : str
+        Multiplier distribution: ``"rademacher"``, ``"mammen"``, or ``"webb"``.
+    rng : np.random.Generator
+        Random number generator.
+
+    Returns
+    -------
+    weights : np.ndarray
+        Multiplier weights, shape ``(n_bootstrap, n_psu)``.
+    psu_ids : np.ndarray
+        Unique PSU identifiers aligned to columns of *weights*.
+    """
+    psu = resolved_survey.psu
+    strata = resolved_survey.strata
+
+    if resolved_survey.lonely_psu == "adjust":
+        raise NotImplementedError(
+            "lonely_psu='adjust' is not yet supported for survey-aware bootstrap. "
+            "Use lonely_psu='remove' or 'certainty', or use analytical inference."
+        )
+
+    if psu is None:
+        # Each observation is its own PSU
+        n_psu = len(resolved_survey.weights)
+        psu_ids = np.arange(n_psu)
+    else:
+        psu_ids = np.unique(psu)
+        n_psu = len(psu_ids)
+
+    if strata is None:
+        # No stratification — generate a single block of weights
+        if n_psu < 2:
+            # Single PSU — variance unidentified (matches compute_survey_vcov)
+            weights = np.zeros((n_bootstrap, n_psu), dtype=np.float64)
+            return weights, psu_ids
+        weights = generate_bootstrap_weights_batch(n_bootstrap, n_psu, weight_type, rng)
+        # FPC scaling (unstratified)
+        if resolved_survey.fpc is not None:
+            if psu is not None:
+                n_units_for_fpc = n_psu
+            else:
+                n_units_for_fpc = len(resolved_survey.weights)
+            if resolved_survey.fpc[0] < n_units_for_fpc:
+                raise ValueError(
+                    f"FPC ({resolved_survey.fpc[0]}) is less than the number of PSUs "
+                    f"({n_units_for_fpc}). FPC must be >= number of PSUs."
+                )
+            f = n_units_for_fpc / resolved_survey.fpc[0]
+            if f < 1.0:
+                weights = weights * np.sqrt(1.0 - f)
+            else:
+                weights = np.zeros_like(weights)
+    else:
+        # Stratified — generate independently within strata
+        weights = np.empty((n_bootstrap, n_psu), dtype=np.float64)
+
+        # Build PSU → column-index map
+        psu_to_col = {int(p): i for i, p in enumerate(psu_ids)}
+
+        unique_strata = np.unique(strata)
+        for h in unique_strata:
+            mask_h = strata == h
+
+            if psu is not None:
+                psus_in_h = np.unique(psu[mask_h])
+            else:
+                psus_in_h = np.where(mask_h)[0]
+
+            n_h = len(psus_in_h)
+            cols = np.array([psu_to_col[int(p)] for p in psus_in_h])
+
+            if n_h < 2:
+                # Lonely PSU — zero weight (matches remove/certainty behavior)
+                weights[:, cols] = 0.0
+                continue
+
+            # Generate weights for this stratum
+            stratum_weights = generate_bootstrap_weights_batch_numpy(
+                n_bootstrap, n_h, weight_type, rng
+            )
+
+            # FPC scaling
+            if resolved_survey.fpc is not None:
+                N_h = resolved_survey.fpc[mask_h][0]
+                if N_h < n_h:
+                    raise ValueError(
+                        f"FPC ({N_h}) is less than the number of PSUs "
+                        f"({n_h}) in stratum {h}. FPC must be >= n_PSU."
+                    )
+                f_h = n_h / N_h
+                if f_h < 1.0:
+                    stratum_weights = stratum_weights * np.sqrt(1.0 - f_h)
+                else:
+                    stratum_weights = np.zeros_like(stratum_weights)
+
+            weights[:, cols] = stratum_weights
+
+    return weights, psu_ids
+
+
+def generate_rao_wu_weights(
+    resolved_survey: "ResolvedSurveyDesign",
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Generate one set of Rao-Wu (1988) rescaled observation weights.
+
+    Within each stratum *h* with *n_h* PSUs, draw ``m_h`` PSUs with
+    replacement and rescale observation weights by ``(n_h / m_h) * r_hi``
+    where ``r_hi`` is the count of PSU *i* being selected.
+
+    Without FPC: ``m_h = n_h - 1``.
+    With FPC: ``m_h = max(1, round((1 - f_h) * (n_h - 1)))``
+    (Rao, Wu & Yue 1992, Section 3).
+
+    Parameters
+    ----------
+    resolved_survey : ResolvedSurveyDesign
+        Resolved survey design.
+    rng : np.random.Generator
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Rescaled observation weights, shape ``(n_obs,)``.
+    """
+    n_obs = len(resolved_survey.weights)
+    base_weights = resolved_survey.weights
+    psu = resolved_survey.psu
+    strata = resolved_survey.strata
+
+    if resolved_survey.lonely_psu == "adjust":
+        raise NotImplementedError(
+            "lonely_psu='adjust' is not yet supported for survey-aware bootstrap. "
+            "Use lonely_psu='remove' or 'certainty', or use analytical inference."
+        )
+
+    rescaled = np.zeros(n_obs, dtype=np.float64)
+
+    if psu is None:
+        obs_psu = np.arange(n_obs)
+    else:
+        obs_psu = psu
+
+    if strata is None:
+        strata_masks = [np.ones(n_obs, dtype=bool)]
+    else:
+        unique_strata = np.unique(strata)
+        strata_masks = [strata == h for h in unique_strata]
+
+    for mask_h in strata_masks:
+        psu_h = obs_psu[mask_h]
+        unique_psu_h = np.unique(psu_h)
+        n_h = len(unique_psu_h)
+
+        if n_h < 2:
+            # Census / lonely PSU — keep original weights (zero variance)
+            rescaled[mask_h] = base_weights[mask_h]
+            continue
+
+        # Compute resample size
+        if resolved_survey.fpc is not None:
+            N_h = resolved_survey.fpc[mask_h][0]
+            if N_h < n_h:
+                raise ValueError(
+                    f"FPC ({N_h}) is less than the number of PSUs "
+                    f"({n_h}). FPC must be >= number of PSUs."
+                )
+            f_h = n_h / N_h
+            if f_h >= 1.0:
+                # Census stratum — keep original weights (zero variance)
+                rescaled[mask_h] = base_weights[mask_h]
+                continue
+            m_h = max(1, round((1.0 - f_h) * (n_h - 1)))
+        else:
+            m_h = n_h - 1
+
+        # Draw m_h PSUs with replacement
+        drawn_indices = rng.choice(n_h, size=m_h, replace=True)
+        counts = np.bincount(drawn_indices, minlength=n_h)
+
+        # Rescale factor per PSU: (n_h / m_h) * r_hi
+        scale_per_psu = (n_h / m_h) * counts.astype(np.float64)
+
+        # Map PSU → local index for vectorized application
+        psu_to_local = {int(p): i for i, p in enumerate(unique_psu_h)}
+        obs_in_h = np.where(mask_h)[0]
+        local_indices = np.array([psu_to_local[int(obs_psu[idx])] for idx in obs_in_h])
+        rescaled[obs_in_h] = base_weights[obs_in_h] * scale_per_psu[local_indices]
+
+    return rescaled
+
+
+def generate_rao_wu_weights_batch(
+    n_bootstrap: int,
+    resolved_survey: "ResolvedSurveyDesign",
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Generate multiple sets of Rao-Wu rescaled weights.
+
+    Parameters
+    ----------
+    n_bootstrap : int
+        Number of bootstrap iterations.
+    resolved_survey : ResolvedSurveyDesign
+        Resolved survey design.
+    rng : np.random.Generator
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Rescaled weights, shape ``(n_bootstrap, n_obs)``.
+    """
+    n_obs = len(resolved_survey.weights)
+    result = np.empty((n_bootstrap, n_obs), dtype=np.float64)
+    for b in range(n_bootstrap):
+        result[b] = generate_rao_wu_weights(resolved_survey, rng)
+    return result