Add per-column zero-rate breakdown + embedding-PRDC validation script

MaxGhenis · claude · MaxGhenis · commit 225eb3610438 · 2026-04-17T07:28:34.000-04:00
ScaleUpResult now includes zero_rate_per_column: for every column, the
real zero-rate, synthetic zero-rate, and absolute difference. Lets the
stage-1 doc identify which specific columns drive each method's
overall zero-rate MAE — the pilot/stage-1 result showed every method
drives disabled_ssdi to 0, but aggregate MAE of 0.18+ implies many
other columns also diverge.

scripts/embedding_prdc_compare.py: one-off validation script that
fits a 16-dim autoencoder on the holdout, encodes real and synthetic
to latent space, and reports PRDC both in the raw 50-dim feature
space and in the learned 16-dim embedding. Settles whether the
stage-1 ordering (ZI-QRF &gt; ZI-QDNN &gt; ZI-MAF) is a metric artifact
from PRDC-in-high-dimensions or a genuine method difference.

Usage:
    uv run python scripts/embedding_prdc_compare.py --n-rows 40000

Tests still pass (7/7).

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/scripts/embedding_prdc_compare.py b/scripts/embedding_prdc_compare.py
@@ -0,0 +1,269 @@
+"""Compare raw-feature PRDC vs learned-embedding PRDC on the stage-1 methods.
+
+The scale-up-protocol doc flagged that PRDC in ~50 dimensions may be
+degenerate (curse of dimensionality: k-NN distances concentrate and the
+metric becomes noise-dominated). This script settles the question.
+
+Procedure:
+
+1. Fit each of (ZI-QRF, ZI-MAF, ZI-QDNN) on 40k x 50 real ECPS.
+2. Generate synthetic records from each.
+3. Train a 16-dim autoencoder on the holdout's raw features only.
+4. Compute PRDC in the raw 50-dim feature space (unchanged from stage 1).
+5. Compute PRDC in the 16-dim learned latent space.
+6. Report both side-by-side. If the ordering changes, the stage-1
+   finding was metric-driven not method-driven; if it's preserved, the
+   finding is robust.
+
+Usage:
+    uv run python scripts/embedding_prdc_compare.py \
+        --output artifacts/embedding_prdc_compare.json
+
+Runs in ~5 minutes on 40 k rows x 50 cols (driven by ZI-MAF fit time).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from prdc import compute_prdc
+from sklearn.preprocessing import StandardScaler
+
+from microplex.eval.benchmark import ZIMAFMethod, ZIQDNNMethod, ZIQRFMethod
+from microplex_us.bakeoff import (
+    DEFAULT_CONDITION_COLS,
+    DEFAULT_TARGET_COLS,
+    ScaleUpRunner,
+    ScaleUpStageConfig,
+    stage1_config,
+)
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Autoencoder(nn.Module):
+    """Tiny autoencoder for dimensionality reduction on tabular features."""
+
+    def __init__(self, n_features: int, latent_dim: int = 16, hidden: int = 64) -> None:
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Linear(n_features, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, latent_dim),
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear(latent_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, n_features),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.encoder(x))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x)
+
+
+def fit_autoencoder(
+    x: np.ndarray, latent_dim: int = 16, epochs: int = 200, lr: float = 1e-3
+) -> Autoencoder:
+    """Fit an autoencoder on standardized features."""
+    n_features = x.shape[1]
+    model = Autoencoder(n_features=n_features, latent_dim=latent_dim)
+    x_t = torch.tensor(x, dtype=torch.float32)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    batch_size = 256
+    ds = torch.utils.data.TensorDataset(x_t)
+    g = torch.Generator()
+    g.manual_seed(42)
+    loader = torch.utils.data.DataLoader(ds, batch_size=batch_size, shuffle=True, generator=g)
+
+    model.train()
+    for epoch in range(epochs):
+        total = 0.0
+        for (batch,) in loader:
+            optimizer.zero_grad()
+            recon = model(batch)
+            loss = ((recon - batch) ** 2).mean()
+            loss.backward()
+            optimizer.step()
+            total += loss.item() * len(batch)
+        if (epoch + 1) % 50 == 0:
+            LOGGER.info("  AE epoch %d loss=%.4f", epoch + 1, total / len(x))
+    model.eval()
+    return model
+
+
+def encode(model: Autoencoder, x: np.ndarray) -> np.ndarray:
+    with torch.no_grad():
+        return model.encode(torch.tensor(x, dtype=torch.float32)).numpy()
+
+
+def compute_prdc_both_spaces(
+    real: pd.DataFrame,
+    synthetic: pd.DataFrame,
+    encoder: Autoencoder,
+    scaler: StandardScaler,
+    k: int = 5,
+    max_samples: int = 15_000,
+    seed: int = 42,
+) -> dict:
+    """Return {raw: ..., embed: ...} PRDC tuples."""
+    rng = np.random.default_rng(seed)
+    cols = [c for c in real.columns if c in synthetic.columns]
+    r = real[cols].to_numpy(dtype=np.float64)
+    s = synthetic[cols].to_numpy(dtype=np.float64)
+    if len(r) > max_samples:
+        r = r[rng.choice(len(r), size=max_samples, replace=False)]
+    if len(s) > max_samples:
+        s = s[rng.choice(len(s), size=max_samples, replace=False)]
+
+    raw_r = scaler.transform(r)
+    raw_s = scaler.transform(s)
+    raw_metrics = compute_prdc(raw_r, raw_s, nearest_k=k)
+
+    emb_r = encode(encoder, raw_r.astype(np.float32))
+    emb_s = encode(encoder, raw_s.astype(np.float32))
+    emb_metrics = compute_prdc(emb_r, emb_s, nearest_k=k)
+
+    return {
+        "raw": {k: float(v) for k, v in raw_metrics.items()},
+        "embed": {k: float(v) for k, v in emb_metrics.items()},
+    }
+
+
+def build_method(name: str):
+    registry = {
+        "ZI-QRF": ZIQRFMethod,
+        "ZI-MAF": ZIMAFMethod,
+        "ZI-QDNN": ZIQDNNMethod,
+    }
+    return registry[name]()
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--n-rows", type=int, default=40_000)
+    parser.add_argument(
+        "--methods", nargs="+", default=["ZI-QRF", "ZI-MAF", "ZI-QDNN"]
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=Path("artifacts/embedding_prdc_compare.json"),
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--latent-dim", type=int, default=16)
+    parser.add_argument("--ae-epochs", type=int, default=200)
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(
+        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
+    )
+
+    base = stage1_config()
+    cfg = ScaleUpStageConfig(
+        stage="embedding_prdc",
+        n_rows=args.n_rows,
+        methods=tuple(args.methods),
+        condition_cols=DEFAULT_CONDITION_COLS,
+        target_cols=DEFAULT_TARGET_COLS,
+        holdout_frac=0.2,
+        seed=args.seed,
+        k=5,
+        data_path=base.data_path,
+        year=base.year,
+        rare_cell_checks=(),
+        prdc_max_samples=15_000,
+    )
+
+    runner = ScaleUpRunner(cfg)
+    df = runner.load_frame()
+    train, holdout = runner.split(df)
+    LOGGER.info(
+        "loaded: train=%d holdout=%d cols=%d", len(train), len(holdout), len(df.columns)
+    )
+
+    scaler = StandardScaler().fit(holdout.to_numpy(dtype=np.float64))
+
+    LOGGER.info("fitting autoencoder on holdout...")
+    t0 = time.time()
+    encoder = fit_autoencoder(
+        scaler.transform(holdout.to_numpy(dtype=np.float64)).astype(np.float32),
+        latent_dim=args.latent_dim,
+        epochs=args.ae_epochs,
+    )
+    LOGGER.info("  autoencoder fit=%.1fs", time.time() - t0)
+
+    results = []
+    for method_name in args.methods:
+        LOGGER.info("== %s ==", method_name)
+        method = build_method(method_name)
+        t0 = time.time()
+        method.fit(sources={"ecps": train.copy()}, shared_cols=list(DEFAULT_CONDITION_COLS))
+        fit_s = time.time() - t0
+
+        t0 = time.time()
+        synth = method.generate(len(train), seed=args.seed)
+        gen_s = time.time() - t0
+
+        metrics = compute_prdc_both_spaces(
+            holdout, synth, encoder, scaler, k=5, seed=args.seed
+        )
+        LOGGER.info(
+            "  raw:   prec=%.3f dens=%.3f cov=%.3f",
+            metrics["raw"]["precision"],
+            metrics["raw"]["density"],
+            metrics["raw"]["coverage"],
+        )
+        LOGGER.info(
+            "  embed: prec=%.3f dens=%.3f cov=%.3f  (fit=%.1fs gen=%.1fs)",
+            metrics["embed"]["precision"],
+            metrics["embed"]["density"],
+            metrics["embed"]["coverage"],
+            fit_s,
+            gen_s,
+        )
+        results.append(
+            {
+                "method": method_name,
+                "fit_wall_seconds": fit_s,
+                "generate_wall_seconds": gen_s,
+                **metrics,
+            }
+        )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(results, indent=2, default=str))
+
+    print()
+    print("== Raw-feature PRDC (50-dim) ==")
+    for r in sorted(results, key=lambda x: -x["raw"]["coverage"]):
+        print(
+            f"  {r['method']:8s}: cov={r['raw']['coverage']:.3f} "
+            f"prec={r['raw']['precision']:.3f} dens={r['raw']['density']:.3f}"
+        )
+    print()
+    print(f"== Learned-embedding PRDC ({args.latent_dim}-dim) ==")
+    for r in sorted(results, key=lambda x: -x["embed"]["coverage"]):
+        print(
+            f"  {r['method']:8s}: cov={r['embed']['coverage']:.3f} "
+            f"prec={r['embed']['precision']:.3f} dens={r['embed']['density']:.3f}"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/src/microplex_us/bakeoff/scale_up.py b/src/microplex_us/bakeoff/scale_up.py
@@ -202,6 +202,7 @@ class ScaleUpResult:
     coverage: float
     rare_cell_ratios: dict[str, float]
     zero_rate_mae: float
+    zero_rate_per_column: dict[str, dict[str, float]] = field(default_factory=dict)
     notes: str = ""
 
     def to_dict(self) -> dict[str, Any]:
@@ -407,6 +408,23 @@ def _compute_zero_rate_mae(real: pd.DataFrame, synthetic: pd.DataFrame) -> float
     return float(np.mean(errs)) if errs else 0.0
 
 
+def _compute_zero_rate_per_column(
+    real: pd.DataFrame, synthetic: pd.DataFrame
+) -> dict[str, dict[str, float]]:
+    """Per-column {real_zero_rate, synth_zero_rate, abs_diff} breakdown."""
+    cols = [c for c in real.columns if c in synthetic.columns]
+    out: dict[str, dict[str, float]] = {}
+    for c in cols:
+        r_zero = float((real[c] == 0).mean())
+        s_zero = float((synthetic[c] == 0).mean())
+        out[c] = {
+            "real": r_zero,
+            "synth": s_zero,
+            "abs_diff": abs(r_zero - s_zero),
+        }
+    return out
+
+
 def _compute_prdc(
     real: pd.DataFrame,
     synthetic: pd.DataFrame,
@@ -614,6 +632,7 @@ def run(
                 holdout, synthetic, self.config.rare_cell_checks
             )
             zero_mae = _compute_zero_rate_mae(holdout, synthetic)
+            zero_per_col = _compute_zero_rate_per_column(holdout, synthetic)
 
             result = ScaleUpResult(
                 stage=self.config.stage,
@@ -630,6 +649,7 @@ def run(
                 coverage=coverage,
                 rare_cell_ratios=rare,
                 zero_rate_mae=zero_mae,
+                zero_rate_per_column=zero_per_col,
                 notes="",
             )
             results.append(result)
