PolicyEngine Social Security microdata regression investigation

README.md

Social Security microdata regression investigation

Reproducibility appendix for the PolicyEngine postmortem on social_security and social_security_retirement being dropped or ignored in microsimulation.

Files:

• social_security_microdata_regression_investigation.ipynb: executed notebook with source-control, GitHub Actions, and PyPI evidence.
• reproduce_social_security_microdata_regression.py: plain Python script with the same checks.

The investigation is metadata-only: it does not download large H5 data files. It expects local clones of PolicyEngine/policyengine-us-data and PolicyEngine/policyengine-us; set US_DATA_REPO and US_REPO if needed.

reproduce_social_security_microdata_regression.py

"""Reproduce the Social Security microdata regression timeline.

Requirements:
  - local clones of PolicyEngine/policyengine-us-data and PolicyEngine/policyengine-us
  - git, gh, jq
  - network access for PyPI metadata checks

Set US_DATA_REPO and US_REPO if your clones are in different locations.
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import tarfile
import tempfile
import urllib.request
from datetime import datetime
from pathlib import Path


US_DATA = Path(os.environ.get("US_DATA_REPO", "/Users/maxghenis/PolicyEngine/policyengine-us-data"))
US = Path(os.environ.get("US_REPO", "/Users/maxghenis/PolicyEngine/policyengine-us"))


def run(cmd: str, cwd: Path | None = None) -> str:
    result = subprocess.run(
        cmd,
        cwd=str(cwd) if cwd else None,
        shell=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        check=True,
    )
    return result.stdout.strip()


def git(repo: Path, args: str) -> str:
    return run(f"git {args}", cwd=repo)


def policyengine_us_version_at(data_sha: str) -> str:
    lock = git(US_DATA, f"show {data_sha}:uv.lock")
    lines = lock.splitlines()
    for idx, line in enumerate(lines):
        if line.strip() == 'name = "policyengine-us"':
            for following in lines[idx + 1 : idx + 8]:
                m = re.search(r'version = "([^"]+)"', following)
                if m:
                    return m.group(1)
    return "not found"


def pyproject_version_at(repo: Path, sha: str) -> str:
    pyproject = git(repo, f"show {sha}:pyproject.toml")
    for line in pyproject.splitlines():
        if line.startswith("version = "):
            return line.split('"')[1]
    return "not found"


def commit_row(repo: Path, sha: str) -> dict[str, str]:
    out = git(repo, f"show --format='%cI%x09%h%x09%s' --no-patch {sha}")
    date, short, subject = out.split("\t", 2)
    return {"sha": short, "date": date, "subject": subject}


def has_retirement_formula_in_sdist(version: str) -> bool:
    with urllib.request.urlopen(f"https://pypi.org/pypi/policyengine-us/{version}/json", timeout=30) as response:
        payload = json.load(response)
    sdist_url = next(
        file["url"]
        for file in payload["urls"]
        if file["packagetype"] == "sdist"
    )
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / f"policyengine_us-{version}.tar.gz"
        urllib.request.urlretrieve(sdist_url, path)
        with tarfile.open(path) as tar:
            member = next(
                m
                for m in tar.getmembers()
                if m.name.endswith(
                    "policyengine_us/variables/gov/ssa/ss/social_security_retirement.py"
                )
            )
            text = tar.extractfile(member).read().decode()
    return "def formula(" in text


def duration(start: str, end: str) -> str:
    delta = datetime.fromisoformat(end) - datetime.fromisoformat(start)
    days = delta.days
    hours, rem = divmod(delta.seconds, 3600)
    minutes = rem // 60
    return f"{days}d {hours}h {minutes}m"


def print_table(rows: list[dict[str, str]], columns: list[str]) -> None:
    widths = {
        col: max(len(col), *(len(str(row.get(col, ""))) for row in rows))
        for col in columns
    }
    print(" | ".join(col.ljust(widths[col]) for col in columns))
    print("-+-".join("-" * widths[col] for col in columns))
    for row in rows:
        print(" | ".join(str(row.get(col, "")).ljust(widths[col]) for col in columns))


def main() -> None:
    print("Repositories")
    print(f"policyengine-us-data: {US_DATA}")
    print(f"policyengine-us:      {US}")
    print()

    print("1. Key source-control events")
    key_events = [
        ("policyengine-us-data", US_DATA, "7025f6a2", "merged formula/adds/subtracts export pruning"),
        ("policyengine-us-data", US_DATA, "6216d02c", "first data package update after #554"),
        ("policyengine-us-data", US_DATA, "dd3455a0", "QRF-impute CPS-only variables including SS subcomponents"),
        ("policyengine-us", US, "c5f1b5da1b", "made social_security_retirement formula-backed"),
        ("policyengine-us-data", US_DATA, "491ac09c", "first bad dependency lock, failed pipeline"),
        ("policyengine-us-data", US_DATA, "a2f3bb36", "first completed bad generated data run"),
        ("policyengine-us-data", US_DATA, "f14931eb", "first promoted/live bad data run found"),
        ("policyengine-us", US, "06a52825cf", "restored social_security_retirement as canonical input"),
        ("policyengine-us-data", US_DATA, "61a43e95", "first data package update after the fix"),
    ]
    rows = []
    for repo_name, repo, sha, note in key_events:
        row = commit_row(repo, sha)
        row["repo"] = repo_name
        row["note"] = note
        if repo_name == "policyengine-us-data":
            row["data_version"] = pyproject_version_at(repo, sha)
            row["policyengine_us"] = policyengine_us_version_at(sha)
        else:
            row["data_version"] = ""
            row["policyengine_us"] = pyproject_version_at(repo, sha)
        rows.append(row)
    print_table(rows, ["date", "repo", "sha", "data_version", "policyengine_us", "note"])
    print()

    print("2. PyPI check: first policyengine-us release where social_security_retirement had a formula")
    formula_rows = []
    for version in ["1.642.0", "1.643.0", "1.644.0", "1.691.3"]:
        formula_rows.append(
            {
                "policyengine_us": version,
                "retirement_has_formula": str(has_retirement_formula_in_sdist(version)),
            }
        )
    print_table(formula_rows, ["policyengine_us", "retirement_has_formula"])
    print()

    print("3. Incident windows used in the postmortem")
    windows = [
        {
            "window": "total social_security direct-input risk",
            "start": "2026-03-04T15:23:42+00:00",
            "end": "2026-03-14T13:41:06+00:00",
            "duration": duration(
                "2026-03-04T15:23:42+00:00",
                "2026-03-14T13:41:06+00:00",
            ),
            "basis": "#554 data release -> #589 source fix",
        },
        {
            "window": "retirement, generated artifacts",
            "start": "2026-04-30T20:51:13+00:00",
            "end": "2026-05-12T17:02:15+00:00",
            "duration": duration(
                "2026-04-30T20:51:13+00:00",
                "2026-05-12T17:02:15+00:00",
            ),
            "basis": "first completed bad run -> first fixed Run Pipeline",
        },
        {
            "window": "retirement, live/promoted artifacts",
            "start": "2026-05-04T03:08:33+00:00",
            "end": "2026-05-12T17:02:15+00:00",
            "duration": duration(
                "2026-05-04T03:08:33+00:00",
                "2026-05-12T17:02:15+00:00",
            ),
            "basis": "first promoted bad run found -> first fixed Run Pipeline",
        },
    ]
    print_table(windows, ["window", "start", "end", "duration", "basis"])
    print()

    print("4. GitHub Actions evidence for first fixed data update")
    runs = run(
        "gh run list --repo PolicyEngine/policyengine-us-data "
        "--workflow 'Run Pipeline' --created '2026-05-12..2026-05-13' "
        "--limit 10 --json createdAt,conclusion,headSha,url "
        "| jq -r '.[] | [.createdAt,.conclusion,.headSha,.url] | @tsv'"
    )
    print(runs)
    print()

    print("Conclusions")
    print("- social_security was at risk when policyengine-us-data stopped exporting computed variables in #554.")
    print("- The safer data contract is: export leaf inputs; calculate formula/adds/subtracts variables in policyengine-us.")
    print("- That contract failed again when policyengine-us made social_security_retirement formula-backed without preserving the canonical input path.")
    print("- The May 12 fix restored social_security_retirement as the canonical input and moved reported/data-only paths out of public formulas.")


if __name__ == "__main__":
    main()

social_security_microdata_regression_investigation.ipynb

ss_by_release_2024.png

ss_by_release_three_year.png

ss_plot.py

"""Plot total Social Security from H5 contents by data release date.

Two views:
  1. Total $B at sim period 2024: stored aggregate vs subcomponent-sum (shows pre-#554 discontinuity)
  2. Total $B at 2024, 2025, 2026 (uprating the subcomponent path with SSA COLA)

Annotates key events: #554 merge, #589 merge, #8040 model break, May 12 fix.
"""
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


CSV = Path("/tmp/ss_sweep_results.csv")
OUT_PRIMARY = Path("/tmp/ss_by_release.png")
OUT_THREE_YEAR = Path("/tmp/ss_by_release_3yr.png")

# SSA COLA factors actually applied to benefits:
# - 2024 COLA = 3.2% (raised benefits paid in 2024 vs 2023)
# - 2025 COLA = 2.5% (raised benefits paid in 2025 vs 2024)
# Policyengine-us uprates SS benefits ~ COLA. Use the published 2025 COLA as 2024→2025,
# and a CBO-typical ~2.5% for 2025→2026. Note: approximate.
COLA_2024_TO_2025 = 1.025
COLA_2025_TO_2026 = 1.025


KEY_EVENTS = [
    ("2026-03-04T15:23Z", "#554 merged", "darkred"),
    ("2026-03-14T13:41Z", "#589 merged", "darkgreen"),
    ("2026-04-17T18:36Z", "policyengine-us#8040", "darkorange"),
    ("2026-05-12T17:02Z", "Fix promoted", "darkblue"),
]


def main():
    df = pd.read_csv(CSV)
    df = df[df["status"] == "ok"].copy()
    df["date"] = pd.to_datetime(df["commit_date"])
    df = df.sort_values("date")
    print(f"{len(df)} ok rows from {df['date'].min()} to {df['date'].max()}")

    # 1) primary: stored vs sum-of-subs at 2024
    fig, ax = plt.subplots(figsize=(11, 5.5))
    stored = df.copy()
    stored.loc[stored["stored_ss_present"] == 0, "stored_ss_total_B"] = float("nan")
    ax.plot(
        stored["date"],
        stored["stored_ss_total_B"],
        marker="o",
        markersize=4,
        label="Stored social_security/2024 (only when present in H5)",
        color="C3",
    )
    ax.plot(
        df["date"],
        df["subs_sum_total_B"],
        marker="o",
        markersize=4,
        label="Sum of subcomponents/2024",
        color="C0",
    )
    for ts, label, color in KEY_EVENTS:
        d = pd.Timestamp(ts)
        ax.axvline(d, color=color, linestyle="--", alpha=0.6)
        ax.text(
            d,
            ax.get_ylim()[1] * 0.95 if ax.get_ylim()[1] > 0 else 1,
            f" {label}",
            color=color,
            rotation=90,
            va="top",
            fontsize=8,
        )
    ax.set_title("Social Security in policyengine-us-data H5 by release date (sim period 2024)")
    ax.set_ylabel("Total Social Security, weighted ($B)")
    ax.set_xlabel("Release date (HF tag)")
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=1))
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m"))
    fig.autofmt_xdate()
    ax.grid(alpha=0.3)
    ax.legend(loc="lower left")
    fig.tight_layout()
    fig.savefig(OUT_PRIMARY, dpi=140)
    print(f"wrote {OUT_PRIMARY}")

    # 2) three-year view (uprate subcomponent path)
    df["ss_2024_effective_B"] = df.apply(
        lambda r: r["stored_ss_total_B"] if r["stored_ss_present"] == 1 else r["subs_sum_total_B"],
        axis=1,
    )
    df["ss_2025_implied_B"] = df["subs_sum_total_B"] * COLA_2024_TO_2025
    df["ss_2026_implied_B"] = df["subs_sum_total_B"] * COLA_2024_TO_2025 * COLA_2025_TO_2026

    fig, ax = plt.subplots(figsize=(11, 5.5))
    ax.plot(df["date"], df["ss_2024_effective_B"], marker="o", markersize=4, label="2024 (sim base year): uses stored if present, else formula", color="C3")
    ax.plot(df["date"], df["ss_2025_implied_B"], marker="o", markersize=4, label="2025 (sim formula × 2.5% COLA)", color="C0")
    ax.plot(df["date"], df["ss_2026_implied_B"], marker="o", markersize=4, label="2026 (sim formula × 2.5% × 2.5%)", color="C2")
    for ts, label, color in KEY_EVENTS:
        d = pd.Timestamp(ts)
        ax.axvline(d, color=color, linestyle="--", alpha=0.6)
    ax.set_title("Implied total Social Security by sim year, by data release date")
    ax.set_ylabel("Total Social Security, weighted ($B)")
    ax.set_xlabel("Release date (HF tag)")
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=1))
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m"))
    fig.autofmt_xdate()
    ax.grid(alpha=0.3)
    ax.legend(loc="lower left", fontsize=9)
    fig.tight_layout()
    fig.savefig(OUT_THREE_YEAR, dpi=140)
    print(f"wrote {OUT_THREE_YEAR}")


if __name__ == "__main__":
    main()

ss_sweep.py

"""Sweep every HF tag of policyengine/policyengine-us-data, compute weighted Social Security
aggregates from the enhanced_cps_2024.h5, then delete the file before fetching the next.

Outputs /tmp/ss_sweep_results.csv with columns:
  version, commit_date, n_persons, sum_weights,
  stored_ss_present, stored_ss_recipients_M, stored_ss_total_B,
  subs_sum_recipients_M, subs_sum_total_B,
  retirement_B, disability_B, survivors_B, dependents_B,
  retirement_present, disability_present, survivors_present, dependents_present
"""
import csv
import os
import sys
import time
import urllib.request
from pathlib import Path

import h5py
import numpy as np


H5_PATH = Path("/tmp/h5_samples/_current.h5")
RESULTS = Path("/tmp/ss_sweep_results.csv")
TAGS_CSV = Path("/tmp/us_data_tags.csv")


def parse_version_key(v):
    try:
        return tuple(int(p) for p in v.split("."))
    except Exception:
        return (0,)


def load_versions():
    out = []
    with TAGS_CSV.open() as f:
        next(f)
        for line in f:
            parts = line.rstrip("\n").split(",")
            ver, sha, date = parts[0], parts[1], parts[2]
            if date.startswith("ERR") or date == "NOT_FOUND":
                continue
            out.append((ver, sha, date))
    out.sort(key=lambda r: parse_version_key(r[0]))
    return out


def download_h5(version, dest):
    url = f"https://huggingface.co/policyengine/policyengine-us-data/resolve/{version}/enhanced_cps_2024.h5"
    req = urllib.request.Request(url, headers={"User-Agent": "ss-sweep/1.0"})
    with urllib.request.urlopen(req, timeout=180) as r, dest.open("wb") as out:
        while True:
            chunk = r.read(1 << 20)
            if not chunk:
                break
            out.write(chunk)


SUBS = ["retirement", "disability", "survivors", "dependents"]


def analyze(path):
    with h5py.File(path, "r") as f:
        if "person_household_id/2024" not in f:
            return None
        person_hh = f["person_household_id/2024"][:]
        hh_id = f["household_id/2024"][:]
        hh_w = f["household_weight/2024"][:]
        weight_map = dict(zip(hh_id.tolist(), hh_w.tolist()))
        person_w = np.array([weight_map[int(h)] for h in person_hh], dtype=np.float64)

        row = {
            "n_persons": int(person_w.size),
            "sum_weights": float(person_w.sum()),
        }

        stored_present = "social_security/2024" in f
        row["stored_ss_present"] = int(stored_present)
        if stored_present:
            ss = f["social_security/2024"][:]
            row["stored_ss_recipients_M"] = float(((ss > 0) * person_w).sum() / 1e6)
            row["stored_ss_total_B"] = float((ss * person_w).sum() / 1e9)
        else:
            row["stored_ss_recipients_M"] = float("nan")
            row["stored_ss_total_B"] = float("nan")

        subs_sum = np.zeros_like(person_w)
        for sub in SUBS:
            key = f"social_security_{sub}/2024"
            present = key in f
            row[f"{sub}_present"] = int(present)
            if present:
                vals = f[key][:].astype(np.float64)
                row[f"{sub}_B"] = float((vals * person_w).sum() / 1e9)
                subs_sum += vals
            else:
                row[f"{sub}_B"] = float("nan")

        row["subs_sum_recipients_M"] = float(((subs_sum > 0) * person_w).sum() / 1e6)
        row["subs_sum_total_B"] = float((subs_sum * person_w).sum() / 1e9)
        return row


def main():
    H5_PATH.parent.mkdir(parents=True, exist_ok=True)
    versions = load_versions()
    print(f"{len(versions)} versions to process", flush=True)

    cols = [
        "version",
        "commit_date",
        "n_persons",
        "sum_weights",
        "stored_ss_present",
        "stored_ss_recipients_M",
        "stored_ss_total_B",
        "subs_sum_recipients_M",
        "subs_sum_total_B",
        "retirement_B",
        "disability_B",
        "survivors_B",
        "dependents_B",
        "retirement_present",
        "disability_present",
        "survivors_present",
        "dependents_present",
        "status",
        "error",
    ]
    write_header = not RESULTS.exists()
    fout = RESULTS.open("a", newline="")
    w = csv.DictWriter(fout, fieldnames=cols)
    if write_header:
        w.writeheader()

    done = set()
    if not write_header:
        with RESULTS.open() as f:
            r = csv.DictReader(f)
            for row in r:
                if row["status"] == "ok":
                    done.add(row["version"])
        print(f"resume: {len(done)} already processed", flush=True)

    t0 = time.time()
    for idx, (ver, sha, date) in enumerate(versions, 1):
        if ver in done:
            continue
        row = {"version": ver, "commit_date": date, "status": "?", "error": ""}
        try:
            if H5_PATH.exists():
                H5_PATH.unlink()
            t = time.time()
            download_h5(ver, H5_PATH)
            dt_dl = time.time() - t
            t = time.time()
            res = analyze(H5_PATH)
            dt_an = time.time() - t
            if res is None:
                row["status"] = "no_2024"
            else:
                row.update(res)
                row["status"] = "ok"
            print(
                f"[{idx}/{len(versions)}] {ver}  dl={dt_dl:.1f}s an={dt_an:.1f}s  status={row['status']}  stored={row.get('stored_ss_total_B')}  subs={row.get('subs_sum_total_B')}",
                flush=True,
            )
        except Exception as e:
            row["status"] = "error"
            row["error"] = str(e)[:200]
            print(f"[{idx}/{len(versions)}] {ver}  ERR {e}", flush=True)
        finally:
            try:
                if H5_PATH.exists():
                    H5_PATH.unlink()
            except Exception:
                pass
            w.writerow(row)
            fout.flush()

    fout.close()
    print(f"DONE in {time.time()-t0:.0f}s. Output: {RESULTS}", flush=True)


if __name__ == "__main__":
    main()

ss_sweep_results.csv