spb/wp20_uqo

1#!/usr/bin/env python32"""Step 02 — Draw the stratified extraction sample from the frame.34Stratification: cma × age_band × price_tier, proportional allocation with a floor per non-empty5stratum, fixed seed. A nested random subsample receives the third (Opus) extractor tier.6Images: the first `max_images_per_listing` URLs in the seller's presentation order (the order is7itself an object of study: image-count sensitivity is simulated by truncation at K = 1, 2, 4, 8, 12).89Outputs:10  data/processed/sample.parquet                 sampled listings with `image_urls` (list) and tier flags11  data/processed/sample_images.parquet          one row per (uid, image_index, url)12  results/raw/sample_summary.json13"""14from __future__ import annotations1516import json17import sys18from pathlib import Path1920import numpy as np21import pandas as pd22import yaml2324ROOT = Path(__file__).resolve().parents[2]25CFG = yaml.safe_load(open(ROOT / "configs" / "data.yaml"))["sample"]262728def main() -> None:29    rng = np.random.default_rng(CFG["seed"])30    df = pd.read_parquet(ROOT / "data" / "processed" / "frame.parquet")31    df["stratum"] = df[CFG["strata"]].astype(str).agg("|".join, axis=1)32    sizes = df["stratum"].value_counts()33    n_target = CFG["n_listings"]34    floor = 2035    # proportional allocation, then enforce floor where the stratum is large enough, then rescale36    alloc = (sizes / sizes.sum() * n_target).round().astype(int)37    alloc = alloc.clip(lower=np.minimum(floor, sizes))38    alloc = np.minimum(alloc, sizes)39    # trim largest strata to hit the target exactly40    while alloc.sum() > n_target:41        k = alloc.idxmax()42        alloc[k] -= 143    while alloc.sum() < n_target:44        room = (sizes - alloc)45        k = room.idxmax()46        alloc[k] += 147    parts = []48    for strat, n in alloc.items():49        sub = df[df["stratum"] == strat]50        idx = rng.choice(sub.index.values, size=int(n), replace=False)51        parts.append(df.loc[idx])52    smp = pd.concat(parts).sample(frac=1.0, random_state=int(rng.integers(1e9))).reset_index(drop=True)53    # design weight = stratum size in frame / stratum size in sample (for frame-representative statistics)54    nsamp = smp["stratum"].value_counts()55    smp["weight"] = smp["stratum"].map(lambda k: sizes[k] / nsamp[k])56    smp["tier_haiku"] = True57    smp["tier_sonnet"] = True58    opus_idx = rng.choice(smp.index.values, size=CFG["subsample_opus"], replace=False)59    smp["tier_opus"] = False60    smp.loc[opus_idx, "tier_opus"] = True61    K = CFG["max_images_per_listing"]62    smp["image_urls_used"] = smp["image_urls"].map(lambda l: list(l)[:K])63    smp["n_images_used"] = smp["image_urls_used"].map(len)64    smp.to_parquet(ROOT / "data" / "processed" / "sample.parquet", index=False)65    rows = [{"uid": r.uid, "image_index": i, "url": u} for r in smp.itertuples() for i, u in enumerate(r.image_urls_used)]66    img = pd.DataFrame(rows)67    img.to_parquet(ROOT / "data" / "processed" / "sample_images.parquet", index=False)68    summary = {69        "seed": CFG["seed"], "n_listings": int(len(smp)), "n_images": int(len(img)),70        "n_strata": int(len(alloc)), "opus_subsample": int(smp["tier_opus"].sum()),71        "by_cma": smp["cma"].value_counts().to_dict(),72        "by_age_band": smp["age_band"].astype(str).value_counts().sort_index().to_dict(),73        "by_price_tier": smp["price_tier"].astype(str).value_counts().sort_index().to_dict(),74        "by_source": smp["source"].value_counts().to_dict(),75        "images_per_listing_used": smp["n_images_used"].describe().to_dict(),76        "frame_vs_sample_price_median": [float(df["price"].median()), float(smp["price"].median())],77    }78    json.dump(summary, open(ROOT / "results" / "raw" / "sample_summary.json", "w"), indent=2, ensure_ascii=False, default=str)79    print(json.dumps(summary, indent=1, ensure_ascii=False, default=str))808182if __name__ == "__main__":83    sys.exit(main())84