Use 2024 SIPP source for imputations

Author: MaxGhenis

.gitignore

@@ -2,9 +2,12 @@
 **/__pycache__
 **/.DS_STORE
 **/*.h5
+**/*.h5.lock
 **/*.npy
 **/*.csv
 **/*.csv.gz
+**/pu*_csv.zip
+**/*.clone_diagnostics.json
 **/_build
 **/*.pkl
 **/*.db

policyengine_us_data/calibration/source_impute.py

@@ -51,6 +51,7 @@
     SSI_DISABILITY_EXPORT_VARIABLES,
     VEHICLE_MODEL_PREDICTORS,
     build_vehicle_training_frame,
+    ensure_sipp_file,
     get_ssi_disability_model,
     predict_ssi_disability_criteria,
     preserve_under_65_ssi_disability_criteria,
@@ -663,16 +664,26 @@ def _impute_sipp(
     Returns:
         Updated data dict.
     """
-    from huggingface_hub import hf_hub_download
-    from policyengine_us_data.storage import STORAGE_FOLDER
-
-    hf_hub_download(
-        repo_id="PolicyEngine/policyengine-us-data",
-        filename="pu2023_slim.csv",
-        repo_type="model",
-        local_dir=STORAGE_FOLDER,
+    tip_cols = (
+        [
+            "SSUID",
+            "MONTHCODE",
+            "WPFINWGT",
+            "TAGE",
+            "TPTOTINC",
+        ]
+        + SIPP_JOB_OCCUPATION_COLUMNS
+        + SIPP_TIP_AMOUNT_COLUMNS
+        + [
+            SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN[column]
+            for column in SIPP_TIP_AMOUNT_COLUMNS
+        ]
+    )
+    sipp_df = pd.read_csv(
+        ensure_sipp_file(),
+        delimiter="|",
+        usecols=tip_cols,
     )
-    sipp_df = pd.read_csv(STORAGE_FOLDER / "pu2023_slim.csv")
 
     tip_amount_columns = [
         column for column in SIPP_TIP_AMOUNT_COLUMNS if column in sipp_df
@@ -788,12 +799,6 @@ def _impute_sipp(
 
     # Asset imputation
     try:
-        hf_hub_download(
-            repo_id="PolicyEngine/policyengine-us-data",
-            filename="pu2023.csv",
-            repo_type="model",
-            local_dir=STORAGE_FOLDER,
-        )
         asset_cols = (
             [
                 "SSUID",
@@ -817,7 +822,7 @@ def _impute_sipp(
             + SIPP_ASSET_ALLOCATION_COLUMNS
         )
         asset_df = pd.read_csv(
-            STORAGE_FOLDER / "pu2023.csv",
+            ensure_sipp_file(),
             delimiter="|",
             usecols=asset_cols,
         )

policyengine_us_data/datasets/sipp/README.md

@@ -18,8 +18,8 @@ SIPP panel wave. These are the canonical reference for every variable
 name, value code, and weighting construct used by the code in this
 folder:
 
-- [SIPP 2023 public-use data dictionary (PDF)](https://www2.census.gov/programs-surveys/sipp/tech-documentation/data-dictionaries/2023/2023_SIPP_Data_Dictionary.pdf)
-- [SIPP 2023 users' guide (PDF, Aug 2026 revision)](https://www2.census.gov/programs-surveys/sipp/tech-documentation/methodology/2023_SIPP_Users_Guide_AUG26.pdf)
+- [SIPP 2024 public-use data dictionary (PDF)](https://www2.census.gov/programs-surveys/sipp/tech-documentation/data-dictionaries/2024/2024_SIPP_Data_Dictionary.pdf)
+- [SIPP 2024 users' guide (PDF)](https://www2.census.gov/programs-surveys/sipp/tech-documentation/methodology/2024_SIPP_Users_Guide.pdf)
 
 See also:
 
@@ -30,15 +30,16 @@ See also:
 ## Data products in this folder
 
 - `sipp.py` — trains and caches QRF imputation models (`get_tip_model`,
-  `get_asset_model`, `get_vehicle_model`) from SIPP 2023 person-month
+  `get_asset_model`, `get_vehicle_model`) from SIPP 2024 person-month
   data. The training frame is filtered to `MONTHCODE == 12` (December)
   so every row represents one person-year rather than twelve annualized
   months.
 
-The raw SIPP CSVs (`pu2023.csv` and the slim variant `pu2023_slim.csv`)
-are mirrored on the `PolicyEngine/policyengine-us-data` HuggingFace model
-repo and downloaded on demand when a training run is needed. They are
-not vendored in this Git repository.
+The raw SIPP CSV (`pu2024.csv`) is downloaded on demand when a training
+run is needed. The downloader first checks the
+`PolicyEngine/policyengine-us-data` HuggingFace model repo for a cached
+copy, then falls back to Census's public `pu2024_csv.zip` archive. The raw
+file is not vendored in this Git repository.
 
 ## Licensing
 

policyengine_us_data/datasets/sipp/sipp.py

@@ -1,14 +1,19 @@
-import pandas as pd
-import numpy as np
-from microimpute.models.qrf import QRF
-from policyengine_us_data.storage import STORAGE_FOLDER
-from policyengine_us_data.utils.randomness import seeded_rng
 import pickle
+from urllib.error import HTTPError, URLError
+from urllib.request import urlretrieve
+from zipfile import ZipFile
+
 from huggingface_hub import hf_hub_download
+import numpy as np
+import pandas as pd
+from microimpute.models.qrf import QRF
+
 from policyengine_us_data.datasets.cps.tipped_occupation import (
     derive_any_treasury_tipped_occupation_code,
     derive_is_tipped_occupation,
 )
+from policyengine_us_data.storage import STORAGE_FOLDER
+from policyengine_us_data.utils.randomness import seeded_rng
 from policyengine_us_data.utils.source_quality import (
     cap_training_sample,
     filter_positive_finite_weight_rows,
@@ -19,6 +24,15 @@
 )
 
 
+SIPP_YEAR = 2024
+SIPP_REFERENCE_YEAR = 2023
+SIPP_FULL_FILE = f"pu{SIPP_YEAR}.csv"
+SIPP_FULL_ZIP_FILE = f"pu{SIPP_YEAR}_csv.zip"
+SIPP_FULL_ZIP_URL = (
+    "https://www2.census.gov/programs-surveys/sipp/data/datasets/"
+    f"{SIPP_YEAR}/{SIPP_FULL_ZIP_FILE}"
+)
+
 SIPP_JOB_OCCUPATION_COLUMNS = [f"TJB{i}_OCC" for i in range(1, 8)]
 SIPP_TIP_AMOUNT_COLUMNS = [f"TJB{i}_TXAMT" for i in range(1, 8)]
 SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN = {
@@ -91,63 +105,88 @@
 ]
 
 
-def train_tip_model():
-    DOWNLOAD_FULL_SIPP = False
+def ensure_sipp_file(filename: str = SIPP_FULL_FILE):
+    """Return a local SIPP public-use file, downloading it if needed."""
+
+    local_path = STORAGE_FOLDER / filename
+    if local_path.exists():
+        return local_path
 
-    if DOWNLOAD_FULL_SIPP:
-        hf_hub_download(
+    try:
+        downloaded_path = hf_hub_download(
             repo_id="PolicyEngine/policyengine-us-data",
-            filename="pu2023.csv",
+            filename=filename,
             repo_type="model",
             local_dir=STORAGE_FOLDER,
         )
-        cols = [
-            "SSUID",
-            "PNUM",
-            "MONTHCODE",
-            "ERESIDENCEID",
-            "ERELRPE",
-            "SPANEL",
-            "SWAVE",
-            "WPFINWGT",
-            "ESEX",
-            "TAGE",
-            "TAGE_EHC",
-            "ERACE",
-            "EORIGIN",
-            "EEDUC",
-            "EDEPCLM",
-            "EMS",
-            "EFSTATUS",
-            "TJB1_TXAMT",
-            "TJB1_MSUM",
-            "TJB1_OCC",
-            "TJB1_IND",
-            "AJB1_TXAMT",
-            "TPTOTINC",
-        ]
+        if downloaded_path:
+            return downloaded_path
+    except Exception:
+        if filename != SIPP_FULL_FILE:
+            raise
+        _download_sipp_full_file_from_census()
+
+    if not local_path.exists():
+        raise FileNotFoundError(f"Could not download {filename}")
+    return local_path
+
+
+def _download_sipp_full_file_from_census():
+    zip_path = STORAGE_FOLDER / SIPP_FULL_ZIP_FILE
+    if not zip_path.exists():
+        try:
+            urlretrieve(SIPP_FULL_ZIP_URL, zip_path)
+        except (HTTPError, URLError) as error:
+            raise FileNotFoundError(
+                f"Could not download {SIPP_FULL_FILE} from HuggingFace or "
+                f"Census at {SIPP_FULL_ZIP_URL}"
+            ) from error
+
+    with ZipFile(zip_path) as archive:
+        if SIPP_FULL_FILE not in archive.namelist():
+            raise FileNotFoundError(
+                f"{SIPP_FULL_ZIP_FILE} does not contain {SIPP_FULL_FILE}"
+            )
+        archive.extract(SIPP_FULL_FILE, STORAGE_FOLDER)
 
-        for col in cols:
-            if "JB1" in col:
-                for i in range(2, 8):
-                    cols.append(col.replace("JB1", f"JB{i}"))
 
-        df = pd.read_csv(
-            STORAGE_FOLDER / "pu2023.csv",
-            delimiter="|",
-            usecols=cols,
-        )
+def train_tip_model():
+    cols = [
+        "SSUID",
+        "PNUM",
+        "MONTHCODE",
+        "ERESIDENCEID",
+        "ERELRPE",
+        "SPANEL",
+        "SWAVE",
+        "WPFINWGT",
+        "ESEX",
+        "TAGE",
+        "TAGE_EHC",
+        "ERACE",
+        "EORIGIN",
+        "EEDUC",
+        "EDEPCLM",
+        "EMS",
+        "EFSTATUS",
+        "TJB1_TXAMT",
+        "TJB1_MSUM",
+        "TJB1_OCC",
+        "TJB1_IND",
+        "AJB1_TXAMT",
+        "TPTOTINC",
+    ]
 
-    else:
-        hf_hub_download(
-            repo_id="PolicyEngine/policyengine-us-data",
-            filename="pu2023_slim.csv",
-            repo_type="model",
-            local_dir=STORAGE_FOLDER,
-        )
-        df = pd.read_csv(
-            STORAGE_FOLDER / "pu2023_slim.csv",
-        )
+    for col in cols.copy():
+        if "JB1" in col:
+            for i in range(2, 8):
+                cols.append(col.replace("JB1", f"JB{i}"))
+
+    df = pd.read_csv(
+        ensure_sipp_file(),
+        delimiter="|",
+        usecols=cols,
+    )
     # Sum tip dollar-amount columns (TJB*_TXAMT) across all jobs.
     # Previously used `str.contains("TXAMT")`, which also picked up
     # AJB*_TXAMT Census allocation flags (small ints 0/1/2 indicating
@@ -255,7 +294,7 @@ def get_tip_model() -> QRF:
     return model
 
 
-# Asset imputation from SIPP 2023
+# Asset imputation from the latest available SIPP public-use file
 # Imputes asset categories separately for policy flexibility
 
 ASSET_JOB_EARNINGS_COLUMNS = [f"TJB{i}_MSUM" for i in range(1, 8)]
@@ -757,7 +796,7 @@ def predict_ssi_disability_criteria(model, receiver_df: pd.DataFrame) -> np.ndar
 
 
 def train_asset_model():
-    """Train QRF model for liquid asset categories using SIPP 2023 data.
+    """Train QRF model for liquid asset categories using SIPP data.
 
     Imputes three asset categories separately:
     - bank_account_assets: checking, savings, money market (TVAL_BANK)
@@ -766,15 +805,8 @@ def train_asset_model():
 
     Policy models can then define countable resources based on rules.
     """
-    hf_hub_download(
-        repo_id="PolicyEngine/policyengine-us-data",
-        filename="pu2023.csv",
-        repo_type="model",
-        local_dir=STORAGE_FOLDER,
-    )
-
     df = pd.read_csv(
-        STORAGE_FOLDER / "pu2023.csv",
+        ensure_sipp_file(),
         delimiter="|",
         usecols=ASSET_COLUMNS,
     )
@@ -843,7 +875,7 @@ def train_asset_model():
 
 def get_asset_model() -> QRF:
     """Get or train the liquid asset imputation model."""
-    model_path = STORAGE_FOLDER / "liquid_assets_v3.pkl"
+    model_path = STORAGE_FOLDER / f"liquid_assets_sipp_{SIPP_YEAR}.pkl"
 
     if not model_path.exists():
         model = train_asset_model()
@@ -859,15 +891,8 @@ def get_asset_model() -> QRF:
 
 def train_ssi_disability_model(time_period: int = 2024):
     """Train a boolean model for likely SSI disability criteria passage."""
-    hf_hub_download(
-        repo_id="PolicyEngine/policyengine-us-data",
-        filename="pu2023.csv",
-        repo_type="model",
-        local_dir=STORAGE_FOLDER,
-    )
-
     df = pd.read_csv(
-        STORAGE_FOLDER / "pu2023.csv",
+        ensure_sipp_file(),
         delimiter="|",
         usecols=SSI_DISABILITY_COLUMNS,
     )
@@ -920,20 +945,16 @@ def get_ssi_disability_model(time_period: int = 2024) -> QRF:
 
 
 def _ssi_disability_model_path(time_period: int):
-    return STORAGE_FOLDER / f"ssi_disability_criteria_{time_period}.pkl"
+    return (
+        STORAGE_FOLDER
+        / f"ssi_disability_criteria_{time_period}_sipp_{SIPP_YEAR}.pkl"
+    )
 
 
 def build_vehicle_training_frame() -> pd.DataFrame:
     """Build a household-level SIPP frame for vehicle asset imputation."""
-    hf_hub_download(
-        repo_id="PolicyEngine/policyengine-us-data",
-        filename="pu2023.csv",
-        repo_type="model",
-        local_dir=STORAGE_FOLDER,
-    )
-
     df = pd.read_csv(
-        STORAGE_FOLDER / "pu2023.csv",
+        ensure_sipp_file(),
         delimiter="|",
         usecols=VEHICLE_COLUMNS,
     )
@@ -1005,7 +1026,7 @@ def build_vehicle_training_frame() -> pd.DataFrame:
 
 
 def train_vehicle_model():
-    """Train a household-level vehicle asset model from SIPP 2023."""
+    """Train a household-level vehicle asset model from SIPP."""
     sipp = build_vehicle_training_frame()
     sipp = sipp[~sipp.isna().any(axis=1)]
     vehicle_vars = [
@@ -1042,7 +1063,7 @@ def train_vehicle_model():
 
 def get_vehicle_model() -> QRF:
     """Get or train the household vehicle imputation model."""
-    model_path = STORAGE_FOLDER / "household_vehicle_assets_v2.pkl"
+    model_path = STORAGE_FOLDER / f"household_vehicle_assets_sipp_{SIPP_YEAR}.pkl"
 
     if not model_path.exists():
         model = train_vehicle_model()