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

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "# Social Security microdata regression investigation\n\nThis notebook captures the metadata-only investigation behind the PolicyEngine Social Security postmortem.\nIt avoids downloading H5 datasets; instead it reconstructs the incident timeline from local git history,\nGitHub Actions metadata, and PyPI source distributions.\n\nTo rerun on another machine, set `US_DATA_REPO` and `US_REPO` to local clones of\n`PolicyEngine/policyengine-us-data` and `PolicyEngine/policyengine-us`."
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "import os, sys\nprint(sys.version)\nprint('US_DATA_REPO=', os.environ.get('US_DATA_REPO', '/Users/maxghenis/PolicyEngine/policyengine-us-data'))\nprint('US_REPO=', os.environ.get('US_REPO', '/Users/maxghenis/PolicyEngine/policyengine-us'))",
      "execution_count": 1,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "3.9.6 (default, Apr 17 2026, 18:15:52) \n[Clang 21.0.0 (clang-2100.1.1.101)]\nUS_DATA_REPO= /Users/maxghenis/PolicyEngine/policyengine-us-data\nUS_REPO= /Users/maxghenis/PolicyEngine/policyengine-us\n"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Reproducible investigation script\n\nThe cell below is the same Python code included as `reproduce_social_security_microdata_regression.py` in this gist."
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "\n\"\"\"Reproduce the Social Security microdata regression timeline.\n\nRequirements:\n  - local clones of PolicyEngine/policyengine-us-data and PolicyEngine/policyengine-us\n  - git, gh, jq\n  - network access for PyPI metadata checks\n\nSet US_DATA_REPO and US_REPO if your clones are in different locations.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport os\nimport re\nimport subprocess\nimport tarfile\nimport tempfile\nimport urllib.request\nfrom datetime import datetime\nfrom pathlib import Path\n\n\nUS_DATA = Path(os.environ.get(\"US_DATA_REPO\", \"/Users/maxghenis/PolicyEngine/policyengine-us-data\"))\nUS = Path(os.environ.get(\"US_REPO\", \"/Users/maxghenis/PolicyEngine/policyengine-us\"))\n\n\ndef run(cmd: str, cwd: Path | None = None) -> str:\n    result = subprocess.run(\n        cmd,\n        cwd=str(cwd) if cwd else None,\n        shell=True,\n        text=True,\n        stdout=subprocess.PIPE,\n        stderr=subprocess.STDOUT,\n        check=True,\n    )\n    return result.stdout.strip()\n\n\ndef git(repo: Path, args: str) -> str:\n    return run(f\"git {args}\", cwd=repo)\n\n\ndef policyengine_us_version_at(data_sha: str) -> str:\n    lock = git(US_DATA, f\"show {data_sha}:uv.lock\")\n    lines = lock.splitlines()\n    for idx, line in enumerate(lines):\n        if line.strip() == 'name = \"policyengine-us\"':\n            for following in lines[idx + 1 : idx + 8]:\n                m = re.search(r'version = \"([^\"]+)\"', following)\n                if m:\n                    return m.group(1)\n    return \"not found\"\n\n\ndef pyproject_version_at(repo: Path, sha: str) -> str:\n    pyproject = git(repo, f\"show {sha}:pyproject.toml\")\n    for line in pyproject.splitlines():\n        if line.startswith(\"version = \"):\n            return line.split('\"')[1]\n    return \"not found\"\n\n\ndef commit_row(repo: Path, sha: str) -> dict[str, str]:\n    out = git(repo, f\"show --format='%cI%x09%h%x09%s' --no-patch {sha}\")\n    date, short, subject = out.split(\"\\t\", 2)\n    return {\"sha\": short, \"date\": date, \"subject\": subject}\n\n\ndef has_retirement_formula_in_sdist(version: str) -> bool:\n    with urllib.request.urlopen(f\"https://pypi.org/pypi/policyengine-us/{version}/json\", timeout=30) as response:\n        payload = json.load(response)\n    sdist_url = next(\n        file[\"url\"]\n        for file in payload[\"urls\"]\n        if file[\"packagetype\"] == \"sdist\"\n    )\n    with tempfile.TemporaryDirectory() as tmp:\n        path = Path(tmp) / f\"policyengine_us-{version}.tar.gz\"\n        urllib.request.urlretrieve(sdist_url, path)\n        with tarfile.open(path) as tar:\n            member = next(\n                m\n                for m in tar.getmembers()\n                if m.name.endswith(\n                    \"policyengine_us/variables/gov/ssa/ss/social_security_retirement.py\"\n                )\n            )\n            text = tar.extractfile(member).read().decode()\n    return \"def formula(\" in text\n\n\ndef duration(start: str, end: str) -> str:\n    delta = datetime.fromisoformat(end) - datetime.fromisoformat(start)\n    days = delta.days\n    hours, rem = divmod(delta.seconds, 3600)\n    minutes = rem // 60\n    return f\"{days}d {hours}h {minutes}m\"\n\n\ndef print_table(rows: list[dict[str, str]], columns: list[str]) -> None:\n    widths = {\n        col: max(len(col), *(len(str(row.get(col, \"\"))) for row in rows))\n        for col in columns\n    }\n    print(\" | \".join(col.ljust(widths[col]) for col in columns))\n    print(\"-+-\".join(\"-\" * widths[col] for col in columns))\n    for row in rows:\n        print(\" | \".join(str(row.get(col, \"\")).ljust(widths[col]) for col in columns))\n\n\ndef main() -> None:\n    print(\"Repositories\")\n    print(f\"policyengine-us-data: {US_DATA}\")\n    print(f\"policyengine-us:      {US}\")\n    print()\n\n    print(\"1. Key source-control events\")\n    key_events = [\n        (\"policyengine-us-data\", US_DATA, \"7025f6a2\", \"merged formula/adds/subtracts export pruning\"),\n        (\"policyengine-us-data\", US_DATA, \"6216d02c\", \"first data package update after #554\"),\n        (\"policyengine-us-data\", US_DATA, \"dd3455a0\", \"QRF-impute CPS-only variables including SS subcomponents\"),\n        (\"policyengine-us\", US, \"c5f1b5da1b\", \"made social_security_retirement formula-backed\"),\n        (\"policyengine-us-data\", US_DATA, \"491ac09c\", \"first bad dependency lock, failed pipeline\"),\n        (\"policyengine-us-data\", US_DATA, \"a2f3bb36\", \"first completed bad generated data run\"),\n        (\"policyengine-us-data\", US_DATA, \"f14931eb\", \"first promoted/live bad data run found\"),\n        (\"policyengine-us\", US, \"06a52825cf\", \"restored social_security_retirement as canonical input\"),\n        (\"policyengine-us-data\", US_DATA, \"61a43e95\", \"first data package update after the fix\"),\n    ]\n    rows = []\n    for repo_name, repo, sha, note in key_events:\n        row = commit_row(repo, sha)\n        row[\"repo\"] = repo_name\n        row[\"note\"] = note\n        if repo_name == \"policyengine-us-data\":\n            row[\"data_version\"] = pyproject_version_at(repo, sha)\n            row[\"policyengine_us\"] = policyengine_us_version_at(sha)\n        else:\n            row[\"data_version\"] = \"\"\n            row[\"policyengine_us\"] = pyproject_version_at(repo, sha)\n        rows.append(row)\n    print_table(rows, [\"date\", \"repo\", \"sha\", \"data_version\", \"policyengine_us\", \"note\"])\n    print()\n\n    print(\"2. PyPI check: first policyengine-us release where social_security_retirement had a formula\")\n    formula_rows = []\n    for version in [\"1.642.0\", \"1.643.0\", \"1.644.0\", \"1.691.3\"]:\n        formula_rows.append(\n            {\n                \"policyengine_us\": version,\n                \"retirement_has_formula\": str(has_retirement_formula_in_sdist(version)),\n            }\n        )\n    print_table(formula_rows, [\"policyengine_us\", \"retirement_has_formula\"])\n    print()\n\n    print(\"3. Incident windows used in the postmortem\")\n    windows = [\n        {\n            \"window\": \"total social_security direct-input risk\",\n            \"start\": \"2026-03-04T15:23:42+00:00\",\n            \"end\": \"2026-03-14T13:41:06+00:00\",\n            \"duration\": duration(\n                \"2026-03-04T15:23:42+00:00\",\n                \"2026-03-14T13:41:06+00:00\",\n            ),\n            \"basis\": \"#554 data release -> #589 source fix\",\n        },\n        {\n            \"window\": \"retirement, generated artifacts\",\n            \"start\": \"2026-04-30T20:51:13+00:00\",\n            \"end\": \"2026-05-12T17:02:15+00:00\",\n            \"duration\": duration(\n                \"2026-04-30T20:51:13+00:00\",\n                \"2026-05-12T17:02:15+00:00\",\n            ),\n            \"basis\": \"first completed bad run -> first fixed Run Pipeline\",\n        },\n        {\n            \"window\": \"retirement, live/promoted artifacts\",\n            \"start\": \"2026-05-04T03:08:33+00:00\",\n            \"end\": \"2026-05-12T17:02:15+00:00\",\n            \"duration\": duration(\n                \"2026-05-04T03:08:33+00:00\",\n                \"2026-05-12T17:02:15+00:00\",\n            ),\n            \"basis\": \"first promoted bad run found -> first fixed Run Pipeline\",\n        },\n    ]\n    print_table(windows, [\"window\", \"start\", \"end\", \"duration\", \"basis\"])\n    print()\n\n    print(\"4. GitHub Actions evidence for first fixed data update\")\n    runs = run(\n        \"gh run list --repo PolicyEngine/policyengine-us-data \"\n        \"--workflow 'Run Pipeline' --created '2026-05-12..2026-05-13' \"\n        \"--limit 10 --json createdAt,conclusion,headSha,url \"\n        \"| jq -r '.[] | [.createdAt,.conclusion,.headSha,.url] | @tsv'\"\n    )\n    print(runs)\n    print()\n\n    print(\"Conclusions\")\n    print(\"- social_security was at risk when policyengine-us-data stopped exporting computed variables in #554.\")\n    print(\"- The safer data contract is: export leaf inputs; calculate formula/adds/subtracts variables in policyengine-us.\")\n    print(\"- That contract failed again when policyengine-us made social_security_retirement formula-backed without preserving the canonical input path.\")\n    print(\"- The May 12 fix restored social_security_retirement as the canonical input and moved reported/data-only paths out of public formulas.\")\n\n\nif __name__ == \"__main__\":\n    main()\n",
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Execute the investigation"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "main()",
      "execution_count": 3,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "Repositories\npolicyengine-us-data: /Users/maxghenis/PolicyEngine/policyengine-us-data\npolicyengine-us:      /Users/maxghenis/PolicyEngine/policyengine-us\n\n1. Key source-control events\ndate                      | repo                 | sha        | data_version | policyengine_us | note                                                    \n--------------------------+----------------------+------------+--------------+-----------------+---------------------------------------------------------\n2026-03-04T10:23:08-05:00 | policyengine-us-data | 7025f6a2   | 1.71.1       | 1.587.0         | merged formula/adds/subtracts export pruning            \n2026-03-04T15:23:42Z      | policyengine-us-data | 6216d02c   | 1.71.2       | 1.587.0         | first data package update after #554                    \n2026-03-14T06:41:06-07:00 | policyengine-us-data | dd3455a0   | 1.73.0       | 1.587.0         | QRF-impute CPS-only variables including SS subcomponents\n2026-04-17T14:36:52-04:00 | policyengine-us      | c5f1b5da1b |              | 1.643.0         | made social_security_retirement formula-backed          \n2026-04-29T21:04:17Z      | policyengine-us-data | 491ac09c   | 1.88.3       | 1.674.1         | first bad dependency lock, failed pipeline              \n2026-04-30T20:47:53Z      | policyengine-us-data | a2f3bb36   | 1.89.1       | 1.678.0         | first completed bad generated data run                  \n2026-05-04T03:05:29Z      | policyengine-us-data | f14931eb   | 1.90.1       | 1.678.0         | first promoted/live bad data run found                  \n2026-05-12T12:24:28-04:00 | policyengine-us      | 06a52825cf |              | 1.691.2         | restored social_security_retirement as canonical input  \n2026-05-12T17:00:28Z      | policyengine-us-data | 61a43e95   | 1.112.2      | 1.691.3         | first data package update after the fix                 \n\n2. PyPI check: first policyengine-us release where social_security_retirement had a formula\npolicyengine_us | retirement_has_formula\n----------------+-----------------------\n1.642.0         | False                 \n1.643.0         | False                 \n1.644.0         | True                  \n1.691.3         | False                 \n\n3. Incident windows used in the postmortem\nwindow                                  | start                     | end                       | duration    | basis                                                   \n----------------------------------------+---------------------------+---------------------------+-------------+---------------------------------------------------------\ntotal social_security direct-input risk | 2026-03-04T15:23:42+00:00 | 2026-03-14T13:41:06+00:00 | 9d 22h 17m  | #554 data release -> #589 source fix                    \nretirement, generated artifacts         | 2026-04-30T20:51:13+00:00 | 2026-05-12T17:02:15+00:00 | 11d 20h 11m | first completed bad run -> first fixed Run Pipeline     \nretirement, live/promoted artifacts     | 2026-05-04T03:08:33+00:00 | 2026-05-12T17:02:15+00:00 | 8d 13h 53m  | first promoted bad run found -> first fixed Run Pipeline\n\n4. GitHub Actions evidence for first fixed data update\n2026-05-13T00:51:51Z\tsuccess\ta3370145c92c84b8f01873323650be2e98f1c1c7\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25771234091\n2026-05-12T18:29:29Z\tsuccess\ta7440e58f6016491620810e7b8b6b95bd3d4629b\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25754210398\n2026-05-12T17:42:58Z\tsuccess\td4871933cdd908d8f124c4fdee66e4d4821c7e9e\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25751771946\n2026-05-12T17:02:15Z\tsuccess\t61a43e952dceaa0194547ab3da34736f78962bdc\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25749650363\n2026-05-12T14:53:27Z\tsuccess\tc3a4624adf47486f26a76f5a1f52a0307c377b98\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25742516772\n2026-05-12T13:04:15Z\tsuccess\ta7e80740a8e8d84f39e0de66fabd4332bbf43053\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25736280595\n2026-05-12T11:49:31Z\tsuccess\t1ef3060e8eac6de49956a03cdb1563f7015c0481\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25732498265\n2026-05-12T07:11:16Z\tsuccess\t889ab4383fe77815be08ef0d8cd970bcdc0d9e74\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25719311417\n2026-05-12T06:39:51Z\tsuccess\t89cd9efe9e858d7de40c8421333009b567022bdd\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25717997086\n2026-05-12T02:59:02Z\tsuccess\t49445139f749a20176164b59e0ca4ff60d238895\thttps://github.com/PolicyEngine/policyengine-us-data/actions/runs/25710573876\n\nConclusions\n- social_security was at risk when policyengine-us-data stopped exporting computed variables in #554.\n- The safer data contract is: export leaf inputs; calculate formula/adds/subtracts variables in policyengine-us.\n- That contract failed again when policyengine-us made social_security_retirement formula-backed without preserving the canonical input path.\n- The May 12 fix restored social_security_retirement as the canonical input and moved reported/data-only paths out of public formulas.\n"
        }
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "pygments_lexer": "ipython3"
    },
    "generated_at": "2026-05-13T12:49:07.160798+00:00"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}

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