Merge upstream main into capital gains branch

Author: MaxGhenis

.github/publication_candidates/usdata-gha26379540291-a1/changelog.d/fsla-overtime-premium.added

@@ -0,0 +1 @@
+Export a data-backed FLSA overtime premium proxy from CPS and enhanced CPS.

.github/publication_candidates/usdata-gha26379540291-a1/publication_scope.json

@@ -0,0 +1,7 @@
+{
+  "base_release_version": "1.115.5",
+  "candidate_scope": "1.115.5-minor",
+  "release_bump": "minor",
+  "run_id": "usdata-gha26379540291-a1",
+  "would_release_as_at_build_time": "1.116.0"
+}

.github/publication_scope.json

@@ -1,7 +1,7 @@
 {
   "base_release_version": "1.115.5",
-  "candidate_scope": "1.115.5-patch",
-  "release_bump": "patch",
-  "run_id": "usdata-gha26360054055-a1",
-  "would_release_as_at_build_time": "1.115.6"
+  "candidate_scope": "1.115.5-minor",
+  "release_bump": "minor",
+  "run_id": "usdata-gha26379540291-a1",
+  "would_release_as_at_build_time": "1.116.0"
 }

.github/scripts/check_policyengine_us_dependency.py

@@ -17,9 +17,9 @@
 REPO_ROOT = Path(__file__).resolve().parents[2]
 PYPI_JSON_TIMEOUT_SECONDS = 20
 POLICYENGINE_US = "policyengine-us"
+POLICYENGINE_US_GITHUB_REPO = "github.com/PolicyEngine/policyengine-us"
 STALE_LOCK_PREFIX = "uv.lock has policyengine-us "
-LOCK_GIT_REF_PREFIX = "uv.lock resolves policyengine-us from a Git ref."
-PROJECT_GIT_REF_PREFIX = "pyproject.toml pins policyengine-us to a Git ref."
+GIT_REF_PREFIX = "uv.lock resolves policyengine-us from a Git ref"
 
 
 def _annotation(level: str, message: str) -> str:
@@ -85,11 +85,39 @@ def _latest_pypi_version() -> str:
     return version
 
 
+def _is_policyengine_us_git_source(source: dict[str, object]) -> bool:
+    git_source = source.get("git")
+    return isinstance(git_source, str) and POLICYENGINE_US_GITHUB_REPO in git_source
+
+
+def _is_policyengine_us_git_dependency(dependency: str) -> bool:
+    return (
+        dependency.startswith(f"{POLICYENGINE_US} @ git+")
+        and POLICYENGINE_US_GITHUB_REPO in dependency
+        and re.search(r"@[0-9a-f]{40}$", dependency) is not None
+    )
+
+
+def _allows_temporary_git_ref(
+    locked_version: str,
+    source: dict[str, object],
+    project_dependency: str,
+    latest_version: str | None,
+) -> bool:
+    return (
+        latest_version is not None
+        and _compare_versions(locked_version, latest_version) > 0
+        and _is_policyengine_us_git_source(source)
+        and _is_policyengine_us_git_dependency(project_dependency)
+    )
+
+
 def check_dependency(root: Path, latest_version: str | None = None) -> list[str]:
     locked_version, source = _locked_policyengine_us(root)
     project_dependency = _project_policyengine_us_dependency(root)
-    lock_uses_git_ref = "git" in source
-    project_uses_git_ref = "@" in project_dependency and "git+" in project_dependency
+    git_ref_allowed = _allows_temporary_git_ref(
+        locked_version, source, project_dependency, latest_version
+    )
 
     violations: list[str] = []
     if (
@@ -103,40 +131,31 @@ def check_dependency(root: Path, latest_version: str | None = None) -> list[str]
         )
 
     expected_dependency = f"{POLICYENGINE_US}=={locked_version}"
-    if not project_uses_git_ref and project_dependency != expected_dependency:
+    if project_dependency != expected_dependency and not git_ref_allowed:
         violations.append(
             f"pyproject.toml must pin {expected_dependency} to match uv.lock; "
             f"found {project_dependency!r}."
         )
 
-    if lock_uses_git_ref:
+    if "git" in source and not git_ref_allowed:
         violations.append(
-            f"{LOCK_GIT_REF_PREFIX} Prefer an exact "
+            f"{GIT_REF_PREFIX}. Prefer an exact "
             f"PyPI release pin once policyengine-us {locked_version} is published."
         )
 
-    if project_uses_git_ref:
+    if (
+        "@" in project_dependency
+        and "git+" in project_dependency
+        and not git_ref_allowed
+    ):
         violations.append(
-            f"{PROJECT_GIT_REF_PREFIX} Prefer an exact "
+            "pyproject.toml pins policyengine-us to a Git ref. Prefer an exact "
             "PyPI release pin for production data builds."
         )
 
     return violations
 
 
-def _is_unreleased_git_ref_violation(
-    violation: str,
-    locked_version: str,
-    latest_version: str | None,
-) -> bool:
-    if latest_version is None:
-        return False
-    git_ref_violation = violation.startswith(
-        LOCK_GIT_REF_PREFIX
-    ) or violation.startswith(PROJECT_GIT_REF_PREFIX)
-    return git_ref_violation and _compare_versions(locked_version, latest_version) > 0
-
-
 def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -176,34 +195,35 @@ def main() -> int:
         return 0
 
     if not violations:
-        locked_version, _source = _locked_policyengine_us(REPO_ROOT)
+        locked_version, source = _locked_policyengine_us(REPO_ROOT)
         print(f"policyengine-us dependency is current at {locked_version}.")
+        if _allows_temporary_git_ref(
+            locked_version,
+            source,
+            _project_policyengine_us_dependency(REPO_ROOT),
+            latest_version,
+        ):
+            print(
+                _annotation(
+                    "warning",
+                    f"policyengine-us {locked_version} is temporarily pinned to "
+                    "GitHub because it is newer than the latest PyPI release. "
+                    "Replace it with an exact PyPI release pin once published.",
+                )
+            )
         return 0
 
-    locked_version, _source = _locked_policyengine_us(REPO_ROOT)
     has_blocking_violation = False
     allowed_stale_version = False
-    allowed_unreleased_git_ref = False
     for violation in violations:
         stale_version_violation = violation.startswith(STALE_LOCK_PREFIX)
         allowed_by_override = allow_stale and stale_version_violation
-        allowed_git_ref = _is_unreleased_git_ref_violation(
-            violation,
-            locked_version,
-            latest_version,
-        )
-        level = (
-            "warning"
-            if args.mode == "warn" or allowed_by_override or allowed_git_ref
-            else "error"
-        )
+        level = "warning" if args.mode == "warn" or allowed_by_override else "error"
         print(_annotation(level, violation))
-        if args.mode == "fail" and not allowed_by_override and not allowed_git_ref:
+        if args.mode == "fail" and not allowed_by_override:
             has_blocking_violation = True
         if allowed_by_override:
             allowed_stale_version = True
-        if allowed_git_ref:
-            allowed_unreleased_git_ref = True
 
     if allowed_stale_version:
         print(
@@ -213,18 +233,10 @@ def main() -> int:
                 "policyengine-us lagging the latest PyPI release.",
             )
         )
-    if allowed_unreleased_git_ref:
-        print(
-            _annotation(
-                "warning",
-                "policyengine-us is pinned to an unreleased Git ref; switch to "
-                f"policyengine-us=={locked_version} once that PyPI release exists.",
-            )
-        )
 
     if has_blocking_violation:
         return 1
-    if allowed_stale_version or allowed_unreleased_git_ref:
+    if allowed_stale_version:
         return 0
 
     return 1 if args.mode == "fail" else 0

docs/generated/pipeline_api.json

@@ -61,7 +61,7 @@
     "docstring": "\"Add auto loan balance, interest and net_worth variable.",
     "id": "add_auto_loan",
     "kind": "function",
-    "line": 2951,
+    "line": 3074,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_auto_loan_interest_and_net_worth"
@@ -88,7 +88,7 @@
     "docstring": "Populate household-level geography variables used by PolicyEngine US.\n\nArgs:\n    cps: Output CPS H5 group receiving derived household variables.\n    household: Raw CPS household table.",
     "id": "add_household_variables",
     "kind": "function",
-    "line": 1531,
+    "line": 1667,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_household_variables"
@@ -115,7 +115,7 @@
     "docstring": "Add basic ID and weight variables.\n\nArgs:\n    cps (h5py.File): The CPS dataset file.\n    person (DataFrame): The person table of the ASEC.\n    tax_unit (DataFrame): The tax unit table created from the person table\n        of the ASEC.\n    family (DataFrame): The family table of the ASEC.\n    spm_unit (DataFrame): The SPM unit table created from the person table\n        of the ASEC.\n    household (DataFrame): The household table of the ASEC.",
     "id": "add_id_variables",
     "kind": "function",
-    "line": 997,
+    "line": 1043,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_id_variables"
@@ -139,15 +139,15 @@
     "source_file": "policyengine_us_data/datasets/cps/cps.py"
   },
   "add_org_inputs": {
-    "docstring": "Impute ORG-derived wage and union inputs onto CPS persons.",
+    "docstring": "Impute ORG-derived labor-market inputs and derive overtime premium.",
     "id": "add_org_inputs",
     "kind": "function",
-    "line": 2835,
+    "line": 2974,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_org_labor_market_inputs"
       ],
-      "description": "Impute hourly wage, hourly-pay status, and union coverage from CPS ORG donors.",
+      "description": "Impute hourly wage, hourly-pay status, and union coverage from CPS ORG donors, then derive FLSA overtime premium.",
       "id": "add_org_inputs",
       "label": "ORG Labor-Market Inputs",
       "node_type": "library",
@@ -162,14 +162,14 @@
       ]
     },
     "object_path": "policyengine_us_data.datasets.cps.cps.add_org_labor_market_inputs",
-    "signature": "def add_org_labor_market_inputs(cps: h5py.File) -> None",
+    "signature": "def add_org_labor_market_inputs(cps: h5py.File, time_period: int) -> None",
     "source_file": "policyengine_us_data/datasets/cps/cps.py"
   },
   "add_personal_income_variables": {
     "docstring": "Add income variables.\n\nArgs:\n    cps (h5py.File): The CPS dataset file.\n    person (DataFrame): The CPS person table.\n    year (int): The CPS year",
     "id": "add_personal_income_variables",
     "kind": "function",
-    "line": 1206,
+    "line": 1342,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_personal_income_variables"
@@ -196,7 +196,7 @@
     "docstring": "Add personal demographic variables.\n\nArgs:\n    cps (h5py.File): The CPS dataset file.\n    person (DataFrame): The CPS person table.",
     "id": "add_personal_variables",
     "kind": "function",
-    "line": 1059,
+    "line": 1105,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_personal_variables"
@@ -223,7 +223,7 @@
     "docstring": "",
     "id": "add_previous_year_income",
     "kind": "function",
-    "line": 1573,
+    "line": 1709,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_previous_year_income"
@@ -250,7 +250,7 @@
     "docstring": "",
     "id": "add_rent",
     "kind": "function",
-    "line": 371,
+    "line": 417,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_rent"
@@ -277,7 +277,7 @@
     "docstring": "",
     "id": "add_spm_variables",
     "kind": "function",
-    "line": 1492,
+    "line": 1628,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_spm_variables"
@@ -304,7 +304,7 @@
     "docstring": "Assign SSN card type using PRCITSHP, employment status, and ASEC-UA conditions.\nCodes:\n- 0: \"NONE\" - Likely undocumented immigrants\n- 1: \"CITIZEN\" - US citizens (born or naturalized)\n- 2: \"NON_CITIZEN_VALID_EAD\" - Non-citizens with work/study authorization\n- 3: \"OTHER_NON_CITIZEN\" - Non-citizens with indicators of legal status",
     "id": "add_ssn_card_type",
     "kind": "function",
-    "line": 1679,
+    "line": 1815,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_ssn_card_type"
@@ -331,7 +331,7 @@
     "docstring": "",
     "id": "add_takeup",
     "kind": "function",
-    "line": 519,
+    "line": 565,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_takeup"
@@ -358,7 +358,7 @@
     "docstring": "",
     "id": "add_tips",
     "kind": "function",
-    "line": 2578,
+    "line": 2714,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.add_tips"
@@ -815,7 +815,7 @@
     "docstring": "Replace clone-half person-level feature variables with donor matches.",
     "id": "clone_features",
     "kind": "function",
-    "line": 585,
+    "line": 604,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.extended_cps._splice_clone_feature_predictions"
@@ -878,7 +878,7 @@
     "docstring": "Assert that final exported variables are leaf inputs.",
     "id": "computed_export_contract",
     "kind": "function",
-    "line": 1775,
+    "line": 1795,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.extended_cps.ExtendedCPS._assert_no_computed_variables_exported"
@@ -972,7 +972,7 @@
     "docstring": "Second-stage QRF: train on CPS, predict for PUF clones.\n\nFor the PUF clone half of the extended CPS we need plausible values\nof CPS-only variables (retirement distributions, transfers, hours,\nSPM components, etc.) that are consistent with the clone's\nPUF-imputed income -- not just naively copied from the CPS donor.\n\nWe train a QRF on CPS person-level data where:\n  * predictors = demographics + key income variables\n  * outputs    = CPS-only variables listed in\n                 ``CPS_ONLY_IMPUTED_VARIABLES``\n\nFor PUF clone prediction we use the PUF-imputed income values\nfrom the second half of ``data`` (the clone half, which already\nhas PUF-imputed income from stage 1).\n\nUses ``fit_predict()`` with ``max_train_samples`` instead of\nmanual sampling + separate fit/predict.\n\nArgs:\n    data: Extended dataset dict after ``puf_clone_dataset()`` --\n        already doubled, with PUF-imputed income in the second half.\n    time_period: Tax year.\n    dataset_path: Path to the CPS h5 file for Microsimulation.\n\nReturns:\n    DataFrame with one column per CPS-only variable, containing\n    predicted values for the PUF clone half (person-level).",
     "id": "cps_only",
     "kind": "function",
-    "line": 624,
+    "line": 643,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.extended_cps._impute_cps_only_variables"
@@ -1064,7 +1064,7 @@
     "docstring": "Subsample the loaded CPS dataset and preserve downsampled arrays.\n\nArgs:\n    frac: Fraction of records to retain.",
     "id": "downsample",
     "kind": "function",
-    "line": 338,
+    "line": 384,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.cps.CPS.downsample"
@@ -1325,7 +1325,7 @@
     "docstring": "Check formula-reconstructed housing assistance before export.\n\nThe final H5 must not export formula outputs such as ``housing_assistance``.\nThis guard verifies that the remaining leaf inputs still make those\nformulas produce nonzero values before the export contract strips or\nrejects computed variables.",
     "id": "housing_assistance_microsim_validation",
     "kind": "function",
-    "line": 1545,
+    "line": 1565,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.extended_cps.ExtendedCPS._validate_housing_assistance_microsimulation"
@@ -3216,7 +3216,7 @@
     "docstring": "Replace PUF clone half of CPS-only variables with QRF predictions.\n\nAfter ``puf_clone_dataset()`` the CPS-only variables in the second\nhalf are naive copies of the CPS donor values. This function\nreplaces them with the second-stage QRF predictions that are\nconsistent with the clone's PUF-imputed income.\n\nArgs:\n    data: Extended dataset dict (already doubled).\n    predictions: DataFrame from ``_impute_cps_only_variables()``.\n    time_period: Tax year.\n    dataset_path: Path to CPS h5 file for entity mapping.\n\nReturns:\n    Modified data dict with CPS-only variables spliced in.",
     "id": "qrf_pass2",
     "kind": "function",
-    "line": 1015,
+    "line": 1034,
     "metadata": {
       "api_refs": [
         "policyengine_us_data.datasets.cps.extended_cps._splice_cps_only_predictions"

docs/generated/pipeline_map.json

@@ -2739,7 +2739,7 @@
           "api_refs": [
             "policyengine_us_data.datasets.cps.cps.add_org_labor_market_inputs"
           ],
-          "description": "Impute hourly wage, hourly-pay status, and union coverage from CPS ORG donors.",
+          "description": "Impute hourly wage, hourly-pay status, and union coverage from CPS ORG donors, then derive FLSA overtime premium.",
           "id": "add_org_inputs",
           "label": "ORG Labor-Market Inputs",
           "node_type": "library",