diff --git a/CHANGELOG.md b/CHANGELOG.md index dcb4684a..863276ab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 still defers cluster with a warning). ### Changed +- **Internal: consolidated the panel-to-unit survey-design collapse** shared by `ContinuousDiD` + and `EfficientDiD` into two `diff_diff/survey.py` helpers — + `ResolvedSurveyDesign.subset_to_units_by_row_idx` (folds the index-and-recount preamble around + the existing `subset_to_units`) and `build_unit_first_row_index` (first-panel-row index per + unit). Removes four hand-rolled copies of the collapse plus a slow `df.iterrows()` index build. + No behavior change: survey-weighted SEs, `df_survey`, and design-effect metadata are + bit-identical (locked by an oracle test against the old inline logic). - **`DiagnosticReport` now routes `SpilloverDiDResults`.** Previously a fitted `SpilloverDiD` (Butts 2021) result matched no `_APPLICABILITY` / `_PT_METHOD` entry, so `DiagnosticReport(spillover_result)` reported every check as `not_applicable`. It now routes diff --git a/TODO.md b/TODO.md index fb2d4130..dff1090a 100644 --- a/TODO.md +++ b/TODO.md @@ -28,7 +28,6 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| -| Survey-design resolution / collapse patterns are inconsistent across panel estimators — `ContinuousDiD` rebuilds unit-level design in SE code, `EfficientDiD` builds once in `fit()`, `StackedDiD` re-resolves on stacked data. Extract shared helpers for panel-to-unit collapse, post-filter re-resolution, and metadata recomputation. | `continuous_did.py`, `efficient_did.py`, `stacked_did.py` | #226 | Mid | Low | | `SyntheticControl` remaining ADH-2015 §4 items: the regression-weight `W^reg = X_0'(X_0 X_0')^{-1} X_1` extrapolation diagnostic (flag implied OLS weights outside `[0,1]`) and sparse-SC subset search (`l < J`, holding `V` fixed). LOO, in-time placebo, CV `V`-selection, and inverse-variance `V` have landed; these two are the deferred tail. | `synthetic_control.py`, `synthetic_control_results.py` | ADH-2015 | Mid | Low | | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low | | `ContinuousDiD` CGBS-2024 extensions (matches R `contdid` v0.1.0 deferral set): (a) `covariates=` kwarg; (b) discrete-treatment saturated regression (integer dose currently warned, not routed to per-level coefficients); (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0`. REGISTRY `## ContinuousDiD` → Implementation Checklist marks these `[ ]`. | `continuous_did.py` | CGBS-2024 | Heavy | Low | @@ -123,6 +122,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. | Issue | Location | PR | Priority | |-------|----------|----|----------| +| `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low | | dCDH parity-test SE/CI assertions only cover pure-direction scenarios; mixed-direction SE comparison is structurally apples-to-oranges (cell-count vs obs-count weighting). | `test_chaisemartin_dhaultfoeuille_parity.py` | #294 | Low | | `HeterogeneousAdoptionDiD` survey-design API consolidation (**scheduled: next minor bump**): drop the deprecated `survey=` / `weights=` kwargs on all 8 HAD surfaces; only `survey_design=` remains. Also fold the legacy back-end `weights=` routing into the unified `_resolve_survey_for_fit` path. DeprecationWarning has shipped; the removal is ~50 LoC gated on the semver bump. | `had.py`, `had_pretests.py` | next minor | Medium | | `HeterogeneousAdoptionDiD` joint cross-horizon covariance / sup-t bands: per-horizon SEs use independent sandwiches (paper-faithful pointwise CIs per Pierce-Schott Fig 2). Follow-ups (low demand): IF-based stacking for joint cross-horizon inference; analytical H×H covariance on the weighted ES path; a sup-t band on the unweighted ES path. | `had.py::_fit_event_study` | Phase 2b / 4.5 B | Low | diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index afed61f7..cb4f8e4e 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -31,9 +31,9 @@ ) from diff_diff.linalg import _rank_guarded_inv, solve_ols from diff_diff.survey import ( - ResolvedSurveyDesign, _resolve_survey_for_fit, _validate_unit_constant_survey, + build_unit_first_row_index, compute_survey_vcov, ) from diff_diff.utils import safe_inference @@ -413,8 +413,7 @@ def fit( # Filter out NaN cells (e.g., from zero effective survey mass) gt_results = { - gt: r for gt, r in gt_results.items() - if np.isfinite(r.get("att_glob", np.nan)) + gt: r for gt, r in gt_results.items() if np.isfinite(r.get("att_glob", np.nan)) } if len(gt_results) == 0: @@ -573,9 +572,12 @@ def fit( # Survey df for t-distribution inference (unit-level, not panel-level) _survey_df = analytic.get("df_survey") # Guard: replicate design with undefined df → NaN inference - if (_survey_df is None and resolved_survey is not None - and hasattr(resolved_survey, 'uses_replicate_variance') - and resolved_survey.uses_replicate_variance): + if ( + _survey_df is None + and resolved_survey is not None + and hasattr(resolved_survey, "uses_replicate_variance") + and resolved_survey.uses_replicate_variance + ): _survey_df = 0 # Recompute survey_metadata from unit-level design so reported @@ -589,8 +591,7 @@ def fit( # Propagate replicate df override to survey_metadata for display # (but not the df=0 sentinel — keep metadata as None for undefined df) - if (_survey_df is not None and _survey_df != 0 - and survey_metadata is not None): + if _survey_df is not None and _survey_df != 0 and survey_metadata is not None: if survey_metadata.df_survey != _survey_df: survey_metadata.df_survey = _survey_df @@ -624,30 +625,8 @@ def fit( unit_resolved_es = None if resolved_survey is not None: row_idx = precomp["unit_first_panel_row"] - uw = ( - precomp.get("unit_survey_weights") - if precomp.get("unit_survey_weights") is not None - else np.ones(n_units) - ) - us = ( - resolved_survey.strata[row_idx] - if resolved_survey.strata is not None - else None - ) - up = ( - resolved_survey.psu[row_idx] - if resolved_survey.psu is not None - else None - ) - uf = ( - resolved_survey.fpc[row_idx] - if resolved_survey.fpc is not None - else None - ) - n_strata_u = len(np.unique(us)) if us is not None else 0 - n_psu_u = len(np.unique(up)) if up is not None else 0 - unit_resolved_es = resolved_survey.subset_to_units( - row_idx, uw, us, up, uf, n_strata_u, n_psu_u, + unit_resolved_es = resolved_survey.subset_to_units_by_row_idx( + row_idx, unit_weights=precomp.get("unit_survey_weights") ) for e_val, info_e in event_study_effects.items(): @@ -711,13 +690,21 @@ def fit( # Score-scale: psi = w * if_es (matches TSL bread) psi_es = unit_resolved_es.weights * if_es - variance, _nv = compute_replicate_if_variance(psi_es, unit_resolved_es) - es_se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan + variance, _nv = compute_replicate_if_variance( + psi_es, unit_resolved_es + ) + es_se = ( + float(np.sqrt(max(variance, 0.0))) + if np.isfinite(variance) + else np.nan + ) else: X_ones_es = np.ones((n_units, 1)) tsl_scale_es = float(unit_resolved_es.weights.sum()) if_es_tsl = if_es * tsl_scale_es - vcov_es = compute_survey_vcov(X_ones_es, if_es_tsl, unit_resolved_es) + vcov_es = compute_survey_vcov( + X_ones_es, if_es_tsl, unit_resolved_es + ) es_se = float(np.sqrt(np.abs(vcov_es[0, 0]))) else: es_se = float(np.sqrt(np.sum(if_es**2))) @@ -831,15 +818,11 @@ def _precompute_structures( unit_cohorts[i] = unit_first.loc[u, first_treat] dose_vector[i] = unit_first.loc[u, dose] - # Build unit-to-first-panel-row mapping (for subsetting panel-level arrays) - # This maps each unit index to the positional index of its first row in df. - unit_first_panel_row = np.zeros(n_units, dtype=int) - seen_units: set = set() - for pos_idx, (_, row) in enumerate(df.iterrows()): - u = row[unit] - if u not in seen_units: - seen_units.add(u) - unit_first_panel_row[unit_to_idx[u]] = pos_idx + # Build unit-to-first-panel-row mapping (for subsetting panel-level + # arrays): the positional index of each unit's first row in df, aligned + # to ``all_units`` (== ``unit_to_idx`` order since + # ``unit_to_idx = {u: i for i, u in enumerate(all_units)}``). + unit_first_panel_row = build_unit_first_row_index(df[unit].values, all_units) # Per-unit survey weights (take first obs per unit from panel data) unit_survey_weights = None @@ -949,8 +932,10 @@ def _compute_dose_response_gt( # Guard against zero effective mass (e.g., after subpopulation) if np.sum(w_treated) <= 0 or np.sum(w_control) <= 0: return { - "att_glob": np.nan, "acrt_glob": np.nan, - "n_treated": 0, "n_control": 0, + "att_glob": np.nan, + "acrt_glob": np.nan, + "n_treated": 0, + "n_control": 0, "att_d": np.full(len(dvals), np.nan), "acrt_d": np.full(len(dvals), np.nan), } @@ -1293,23 +1278,8 @@ def _compute_analytical_se( # but influence functions are unit-level (n_units). Build a unit-level # ResolvedSurveyDesign by subsetting to one obs per unit. row_idx = precomp["unit_first_panel_row"] - unit_weights = precomp.get("unit_survey_weights") - if unit_weights is None: - unit_weights = np.ones(n_units) - - unit_strata = ( - resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None - ) - unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None - unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None - - # Count unique strata/PSU in the unit-level subset - n_strata_unit = len(np.unique(unit_strata)) if unit_strata is not None else 0 - n_psu_unit = len(np.unique(unit_psu)) if unit_psu is not None else 0 - - unit_resolved = resolved_survey.subset_to_units( - row_idx, unit_weights, unit_strata, unit_psu, unit_fpc, - n_strata_unit, n_psu_unit, + unit_resolved = resolved_survey.subset_to_units_by_row_idx( + row_idx, unit_weights=precomp.get("unit_survey_weights") ) X_ones = np.ones((n_units, 1)) @@ -1370,7 +1340,11 @@ def _rep_se(if_vals): # Return unit-level survey df and resolved design for metadata recomputation # Only override with n_valid-based df when replicates were actually dropped - if resolved_survey is not None and hasattr(resolved_survey, 'uses_replicate_variance') and resolved_survey.uses_replicate_variance: + if ( + resolved_survey is not None + and hasattr(resolved_survey, "uses_replicate_variance") + and resolved_survey.uses_replicate_variance + ): if _rep_n_valid < unit_resolved.n_replicates: unit_df_survey = _rep_n_valid - 1 if _rep_n_valid > 1 else None else: @@ -1415,7 +1389,11 @@ def _run_bootstrap( # Reject replicate-weight designs for bootstrap — replicate variance # is an analytical alternative to bootstrap, not compatible with it - if resolved_survey is not None and hasattr(resolved_survey, "uses_replicate_variance") and resolved_survey.uses_replicate_variance: + if ( + resolved_survey is not None + and hasattr(resolved_survey, "uses_replicate_variance") + and resolved_survey.uses_replicate_variance + ): raise NotImplementedError( "ContinuousDiD bootstrap (n_bootstrap > 0) is not supported " "with replicate-weight survey designs. Replicate weights provide " @@ -1429,22 +1407,9 @@ def _run_bootstrap( # Build unit-level ResolvedSurveyDesign for survey-aware bootstrap unit_resolved = None if resolved_survey is not None: - from diff_diff.survey import ResolvedSurveyDesign - row_idx = precomp["unit_first_panel_row"] - unit_weights = precomp.get("unit_survey_weights") - if unit_weights is None: - unit_weights = np.ones(n_units) - unit_strata = ( - resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None - ) - unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None - unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None - n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0 - n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0 - unit_resolved = resolved_survey.subset_to_units( - row_idx, unit_weights, unit_strata, unit_psu, unit_fpc, - n_strata_u, n_psu_u, + unit_resolved = resolved_survey.subset_to_units_by_row_idx( + row_idx, unit_weights=precomp.get("unit_survey_weights") ) # Generate bootstrap weights — PSU-level when survey design is present @@ -1682,7 +1647,7 @@ def _bootstrap_gt_cell(gt, info): boot_es[e], alpha=self.alpha, context=f"event study e={e}", - ) + ) es_se[e] = se_e es_ci[e] = ci_e es_p[e] = p_e diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index 59d1e749..0c0fadd1 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -521,7 +521,7 @@ def fit( ) # Resolve survey design if provided - from diff_diff.survey import _resolve_survey_for_fit + from diff_diff.survey import _resolve_survey_for_fit, build_unit_first_row_index resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, "analytical") @@ -655,33 +655,14 @@ def fit( # order. The previous approach (groupby cumcount == 0) yielded # first-appearance order which can differ from sorted order when the # input DataFrame is not pre-sorted by unit. - first_pos: Dict[Any, int] = {} - for i, u in enumerate(df[unit].values): - if u not in first_pos: - first_pos[u] = i - self._unit_first_panel_row = np.array([first_pos[u] for u in all_units]) + self._unit_first_panel_row = build_unit_first_row_index(df[unit].values, all_units) # Build unit-level ResolvedSurveyDesign once (avoids repeated # construction in _compute_survey_eif_se and ensures consistent # unit-level df for safe_inference t-distribution). if resolved_survey is not None: - row_idx = self._unit_first_panel_row - unit_weights_s = resolved_survey.weights[row_idx] - unit_strata = ( - resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None - ) - unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None - unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None - n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0 - n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0 - self._unit_resolved_survey = resolved_survey.subset_to_units( - row_idx, - unit_weights_s, - unit_strata, - unit_psu, - unit_fpc, - n_strata_u, - n_psu_u, + self._unit_resolved_survey = resolved_survey.subset_to_units_by_row_idx( + self._unit_first_panel_row ) # Use unit-level df (not panel-level) for t-distribution self._survey_df = self._unit_resolved_survey.df_survey diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index 20dd2f47..003e26ff 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -773,6 +773,10 @@ def _refit_stacked(w_r): # Re-resolve survey design on the stacked data so that strata/PSU # arrays have the correct length for TSL variance estimation. + # (Unlike ContinuousDiD/EfficientDiD, StackedDiD does NOT collapse + # to unit level via ``subset_to_units_by_row_idx``: control units are + # duplicated across sub-experiments, so the design must be resolved + # at the stacked granularity, not one row per original unit.) resolved_stacked = survey_design.resolve(stacked_df) # Create a copy with composed weights (normalized to sum=n_stacked) diff --git a/diff_diff/survey.py b/diff_diff/survey.py index 89a75a60..900cf38d 100644 --- a/diff_diff/survey.py +++ b/diff_diff/survey.py @@ -16,7 +16,7 @@ import warnings from dataclasses import dataclass, field, replace -from typing import Callable, List, Optional, Tuple +from typing import Any, Callable, List, Optional, Sequence, Tuple import numpy as np import pandas as pd @@ -645,7 +645,10 @@ def subset_to_units( """Create a unit-level copy preserving replicate metadata. Used by panel estimators (ContinuousDiD, EfficientDiD) that collapse - panel-level survey info to one row per unit. + panel-level survey info to one row per unit. Callers that only have a + ``row_idx`` (first panel row per unit) should prefer the + :meth:`subset_to_units_by_row_idx` convenience wrapper, which folds the + index-and-recount preamble. Parameters ---------- @@ -679,6 +682,46 @@ def subset_to_units( nest=self.nest, ) + def subset_to_units_by_row_idx( + self, + row_idx: np.ndarray, + unit_weights: Optional[np.ndarray] = None, + ) -> "ResolvedSurveyDesign": + """Collapse this panel-level design to unit level given a ``row_idx``. + + Convenience wrapper over :meth:`subset_to_units` that folds the + index-and-recount preamble panel estimators (ContinuousDiD, + EfficientDiD) would otherwise hand-roll: index ``strata`` / ``psu`` / + ``fpc`` at ``row_idx`` (one panel row per unit) and recount + ``n_strata`` / ``n_psu`` from the collapsed arrays. + + Parameters + ---------- + row_idx : np.ndarray + First panel-row position for each unit (see + :func:`build_unit_first_row_index`). + unit_weights : np.ndarray, optional + Pre-collapsed unit-level weights. When ``None`` (default), uses + ``self.weights[row_idx]`` — identical to the caller-side + ``survey_weights[row_idx]`` collapse. + """ + if unit_weights is None: + unit_weights = self.weights[row_idx] + unit_strata = self.strata[row_idx] if self.strata is not None else None + unit_psu = self.psu[row_idx] if self.psu is not None else None + unit_fpc = self.fpc[row_idx] if self.fpc is not None else None + n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0 + n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0 + return self.subset_to_units( + row_idx, + unit_weights, + unit_strata, + unit_psu, + unit_fpc, + n_strata_u, + n_psu_u, + ) + @property def needs_survey_vcov(self) -> bool: """Whether survey vcov (not generic sandwich) should be used.""" @@ -1092,6 +1135,26 @@ def _resolve_pweight_only(resolved_survey, estimator_name): ) +def build_unit_first_row_index(unit_values: np.ndarray, unit_order: Sequence[Any]) -> np.ndarray: + """Positional index of each unit's first row, aligned to ``unit_order``. + + Panel estimators that collapse a panel-level survey design to unit level + (ContinuousDiD, EfficientDiD) need, for each unit, the position of its first + row in the fit DataFrame so panel-length survey arrays can be indexed down + to one row per unit. ``unit_values`` is the fit frame's unit column + (``df[unit].values``, in row order); ``unit_order`` is the estimator's + canonical unit ordering (typically ``all_units = sorted(df[unit].unique())``). + + Returns an ``int`` array ``idx`` where ``idx[i]`` is the first row position + of ``unit_order[i]``. + """ + first_pos: dict = {} + for i, u in enumerate(unit_values): + if u not in first_pos: + first_pos[u] = i + return np.array([first_pos[u] for u in unit_order], dtype=int) + + def collapse_survey_to_unit_level(resolved_survey, df, unit_col, all_units): """Collapse observation-level ResolvedSurveyDesign to unit level. diff --git a/tests/test_survey.py b/tests/test_survey.py index 7298686e..95fddbb3 100644 --- a/tests/test_survey.py +++ b/tests/test_survey.py @@ -3690,3 +3690,98 @@ def test_well_conditioned_X_no_warning(self): ) assert vcov.shape == (3, 3) assert np.all(np.isfinite(vcov)) + + +class TestUnitCollapseHelpers: + """Shared panel-to-unit survey-collapse helpers (#226 consolidation). + + ``build_unit_first_row_index`` and + ``ResolvedSurveyDesign.subset_to_units_by_row_idx`` replace the row_idx->unit + collapse hand-rolled in ContinuousDiD / EfficientDiD. These tests lock the + wrappers to bit-identity against the OLD inline preamble (the oracle). + """ + + def test_build_unit_first_row_index_unsorted_panel(self): + from diff_diff.survey import build_unit_first_row_index + + # Unsorted 3-unit x 3-period panel (unit column neither grouped nor sorted). + unit_values = np.array([2, 0, 1, 2, 0, 1, 2, 0, 1]) + unit_order = [0, 1, 2] # sorted(unique) + idx = build_unit_first_row_index(unit_values, unit_order) + # First occurrence: unit 0 -> pos 1, unit 1 -> pos 2, unit 2 -> pos 0. + np.testing.assert_array_equal(idx, np.array([1, 2, 0])) + assert idx.dtype == np.dtype(int) + + @staticmethod + def _oracle_collapse(resolved, row_idx, unit_weights): + # Frozen copy of the OLD inline preamble + subset_to_units call that the + # new subset_to_units_by_row_idx wrapper replaces. This is the oracle. + unit_strata = resolved.strata[row_idx] if resolved.strata is not None else None + unit_psu = resolved.psu[row_idx] if resolved.psu is not None else None + unit_fpc = resolved.fpc[row_idx] if resolved.fpc is not None else None + n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0 + n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0 + return resolved.subset_to_units( + row_idx, unit_weights, unit_strata, unit_psu, unit_fpc, n_strata_u, n_psu_u + ) + + @staticmethod + def _panel_design(with_replicates=False): + n_obs = 6 # 3 units x 2 periods, panel-level + kwargs = dict( + weights=np.array([1.0, 1.0, 2.0, 2.0, 3.0, 3.0]), + weight_type="pweight", + strata=np.array([0, 0, 0, 0, 1, 1]), + psu=np.array([0, 0, 1, 1, 2, 2]), + fpc=np.array([10.0, 10.0, 10.0, 10.0, 20.0, 20.0]), + n_strata=2, + n_psu=3, + lonely_psu="remove", + ) + if with_replicates: + kwargs["replicate_weights"] = np.arange(n_obs * 4, dtype=float).reshape(n_obs, 4) + kwargs["replicate_method"] = "bootstrap" + kwargs["n_replicates"] = 4 + return ResolvedSurveyDesign(**kwargs) + + @staticmethod + def _assert_designs_equal(a, b): + np.testing.assert_array_equal(a.weights, b.weights) + np.testing.assert_array_equal(a.strata, b.strata) + np.testing.assert_array_equal(a.psu, b.psu) + np.testing.assert_array_equal(a.fpc, b.fpc) + assert a.n_strata == b.n_strata + assert a.n_psu == b.n_psu + if a.replicate_weights is None: + assert b.replicate_weights is None + else: + np.testing.assert_array_equal(a.replicate_weights, b.replicate_weights) + + def test_subset_by_row_idx_matches_oracle_explicit_weights(self): + resolved = self._panel_design() + row_idx = np.array([0, 2, 4]) + unit_weights = resolved.weights[row_idx] + got = resolved.subset_to_units_by_row_idx(row_idx, unit_weights=unit_weights) + want = self._oracle_collapse(resolved, row_idx, unit_weights) + self._assert_designs_equal(got, want) + # n_strata/n_psu recounted from the collapsed arrays (3 units -> 2 strata, 3 PSU). + assert got.n_strata == 2 and got.n_psu == 3 + + def test_subset_by_row_idx_default_weights_path(self): + # unit_weights=None must default to self.weights[row_idx] (review LOW #3). + resolved = self._panel_design() + row_idx = np.array([0, 2, 4]) + got = resolved.subset_to_units_by_row_idx(row_idx) + want = self._oracle_collapse(resolved, row_idx, resolved.weights[row_idx]) + self._assert_designs_equal(got, want) + np.testing.assert_array_equal(got.weights, resolved.weights[row_idx]) + + def test_subset_by_row_idx_preserves_replicate_rows(self): + # Replicate-weight design: the (row_idx, :) R-column subset must survive + # the wrapper unchanged (review LOW #3). + resolved = self._panel_design(with_replicates=True) + row_idx = np.array([0, 2, 4]) + got = resolved.subset_to_units_by_row_idx(row_idx) + want = self._oracle_collapse(resolved, row_idx, resolved.weights[row_idx]) + self._assert_designs_equal(got, want) + np.testing.assert_array_equal(got.replicate_weights, resolved.replicate_weights[row_idx, :])