spb/wp3_uqo

1"""Rebuilt analytical sample == stored pickle (WP03-DATA-01 / ADVE-07 / ADVE-10 / THEO-15)."""2import json34import numpy as np5import pandas as pd6import pytest78from wp3 import build, config910SAMPLE_V2 = config.DATA_PROCESSED / "analytical_sample_v2.parquet"11BUILD = config.RESULTS_DIR / "build"121314@pytest.fixture(scope="module")15def sample():16    if not SAMPLE_V2.exists():17        pytest.skip("analytical_sample_v2.parquet not built (run scripts/00_build_sample.py)")18    return pd.read_parquet(SAMPLE_V2)192021def test_schema(sample):22    assert len(sample) == 788_84223    assert list(sample.columns[:68]) == build.SAMPLE_COLUMNS24    assert sample["zpid"].is_unique25    assert sample["address_state"].nunique() == 5126    assert sample["zip3"].nunique() == 886 and sample["address_zipcode"].isna().sum() == 1627    assert sample["region"].isin(list(build.REGION)).all()28    assert int(sample["implausible"].sum()) == 4_970293031def test_lot_reconciliation(sample):32    vc = sample["lot_class"].value_counts()33    assert vc["null"] == 123_942 and vc["zero"] == 6_115 and vc["one"] == 211 and vc["negative"] == 334    assert int((sample["ln_lot"] == 0).sum()) == 130_271 == 123_942 + 6_115 + 211 + 3353637def test_duplicates(sample):38    d = pd.read_csv(BUILD / "duplicates.csv").iloc[0]39    assert d["n_rows_in_duplicate_groups"] == 6_787 and d["n_duplicate_groups"] == 3_15240    assert int((sample["dup_group_id"] >= 0).sum()) == 6_787414243def test_attrition():44    a = pd.read_csv(BUILD / "attrition.csv")45    assert a["n_remaining"].tolist() == [839_313, 817_473, 797_380, 790_799, 789_199, 788_842]464748@pytest.mark.skipif(not config.ANALYTICAL_SAMPLE.exists(), reason="stored pickle not available")49def test_equals_stored_pickle(sample):50    pk = pd.read_pickle(config.ANALYTICAL_SAMPLE)51    assert (pk["zpid"].values == sample["zpid"].values).all()52    for c in build.SAMPLE_COLUMNS:53        a, b = pk[c], sample[c]54        if a.dtype.kind in "biufc" or str(a.dtype).startswith("Int"):55            aa, bb = a.astype(float).values, b.astype(float).values56            assert (np.isnan(aa) == np.isnan(bb)).all(), c57            m = ~np.isnan(aa)58            assert np.abs(aa[m] - bb[m]).max() < 1e-9, c59        else:60            assert ((a.isna() & b.isna()) | (a == b)).all(), c616263def test_build_meta():64    meta = json.loads((BUILD / "build_meta.json").read_text())65    assert meta["n_sample"] == 788_842 and meta["n_trimmed_sample"] == 783_87266    assert meta["data_window"]["max_event_date"].startswith("2026-01-08")67