Skip to content

Commit ff7134d

Browse files
igerberclaude
andcommitted
Cover stratified-no-PSU branch in scaffolding equivalence tests
Addresses the P2 test-coverage gap from the CI review on PR #338: _precompute_psu_scaffolding() has a dedicated strata-with-no-PSU branch (survey.py:1458-1466) where each observation is its own PSU within its stratum, but TestAggregateSurveyScaffolding did not exercise it despite the docstring claiming every supported design mode was covered. Adds two parametrized cases to test_fast_path_equals_legacy: - stratified_no_psu — SurveyDesign(weights=, strata=) - stratified_no_psu_fpc — same plus stratum-level FPC lookup (the fpc-on-this-branch path goes through the same per-stratum first-obs FPC read as stratified+PSU, so both variants matter). Both pass assert_allclose(atol=1e-14, rtol=1e-14) equivalence with the legacy path across all cells. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 6b7ba9f commit ff7134d

1 file changed

Lines changed: 20 additions & 0 deletions

File tree

tests/test_prep.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3483,6 +3483,24 @@ def _build_microdata(self, mode, seed=42):
34833483
sd = SurveyDesign(weights="wt", strata="stratum", psu="psu")
34843484
return df, sd
34853485

3486+
if mode == "stratified_no_psu":
3487+
# strata present, psu absent — each observation is its own
3488+
# PSU within its stratum. This is a distinct scaffolding
3489+
# branch (survey.py:_precompute_psu_scaffolding, else clause
3490+
# of the `if psu is not None` block).
3491+
df = df_base.copy()
3492+
df["stratum"] = rng.integers(0, 4, n)
3493+
sd = SurveyDesign(weights="wt", strata="stratum")
3494+
return df, sd
3495+
3496+
if mode == "stratified_no_psu_fpc":
3497+
# Same branch as above plus stratum-level FPC lookup.
3498+
df = df_base.copy()
3499+
df["stratum"] = rng.integers(0, 4, n)
3500+
df["fpc"] = 1000.0 # well above per-stratum obs count
3501+
sd = SurveyDesign(weights="wt", strata="stratum", fpc="fpc")
3502+
return df, sd
3503+
34863504
if mode == "psu_only":
34873505
df = df_base.copy()
34883506
df["psu"] = rng.integers(0, 12, n)
@@ -3532,6 +3550,8 @@ def _assert_panels_equivalent(p_fast, p_legacy, outcome="y"):
35323550
[
35333551
"stratified_fpc",
35343552
"stratified_no_fpc",
3553+
"stratified_no_psu",
3554+
"stratified_no_psu_fpc",
35353555
"psu_only",
35363556
"weights_only",
35373557
"lonely_remove",

0 commit comments

Comments
 (0)