Construct CPS tax units from household records (#824)

* Construct CPS tax units from household records

* Fix CPS tax unit construction CI

Author: MaxGhenis

changelog.d/cps-tax-unit-construction.added

@@ -0,0 +1 @@
+Construct CPS tax units from ASEC household relationships instead of using Census tax-unit assignments.

policyengine_us_data/datasets/cps/census_cps.py

@@ -5,6 +5,9 @@
 from zipfile import ZipFile
 import pandas as pd
 from policyengine_us_data.storage import STORAGE_FOLDER
+from policyengine_us_data.datasets.cps.tax_unit_construction import (
+    construct_tax_units,
+)
 
 
 OPTIONAL_PERSON_COLUMNS = {
@@ -26,6 +29,7 @@
     "NOW_CHAMPVA",
     "NOW_VACARE",
     "NOW_IHSFLG",
+    "PTOTVAL",
 }
 
 
@@ -59,6 +63,9 @@ class CensusCPS(Dataset):
     time_period: int
     """Year of the dataset."""
 
+    tax_unit_construction_mode: str = "policyengine"
+    """Mode used when constructing tax units from CPS person records."""
+
     def generate(self):
         if self._cps_download_url is None:
             raise ValueError(f"No raw CPS data URL known for year {self.time_period}.")
@@ -117,6 +124,7 @@ def generate(self):
                         usecols=person_usecols,
                     ).fillna(0)
                 person = _fill_missing_optional_person_columns(person)
+                tax_unit = self._create_tax_unit_table(person)
                 storage["person"] = person
                 with zipfile.open(f"{file_prefix}ffpub{file_year_code}.csv") as f:
                     person_family_id = person.PH_SEQ * 10 + person.PF_SEQ
@@ -130,7 +138,7 @@ def generate(self):
                     household_id = household.H_SEQ
                     household = household[household_id.isin(person_household_id)]
                     storage["household"] = household
-                storage["tax_unit"] = self._create_tax_unit_table(person)
+                storage["tax_unit"] = tax_unit
                 storage["spm_unit"] = self._create_spm_unit_table(
                     person, self.time_period
                 )
@@ -139,10 +147,20 @@ def generate(self):
     def _cps_download_url(self) -> str:
         return CPS_URL_BY_YEAR.get(self.time_period)
 
-    def _create_tax_unit_table(self, person: pd.DataFrame) -> pd.DataFrame:
-        tax_unit_df = person[TAX_UNIT_COLUMNS].groupby(person.TAX_ID).sum()
-        tax_unit_df["TAX_ID"] = tax_unit_df.index
-        return tax_unit_df
+    def _create_tax_unit_table(
+        self,
+        person: pd.DataFrame,
+        mode: str | None = None,
+    ) -> pd.DataFrame:
+        person["CENSUS_TAX_ID"] = person["TAX_ID"]
+        mode = mode or self.tax_unit_construction_mode
+        constructed_person, tax_unit_df = construct_tax_units(
+            person=person,
+            year=self.time_period,
+            mode=mode,
+        )
+        person["TAX_ID"] = constructed_person["TAX_ID"].values
+        return tax_unit_df[["TAX_ID"]]
 
     def _create_spm_unit_table(
         self, person: pd.DataFrame, time_period: int
@@ -282,12 +300,18 @@ class CensusCPS_2018(CensusCPS):
     "PF_SEQ",
     "P_SEQ",
     "TAX_ID",
+    "PECOHAB",
     "SPM_ID",
     "A_FNLWGT",
     "A_LINENO",
     "A_SPOUSE",
+    "A_EXPRRP",
+    "A_FAMREL",
+    "A_FAMTYP",
     "A_AGE",
     "A_SEX",
+    "A_ENRLW",
+    "A_FTPT",
     "PEDISEYE",
     "NOW_COV",
     "NOW_DIR",
@@ -318,6 +342,7 @@ class CensusCPS_2018(CensusCPS):
     "LKWEEKS",  # Weeks looking for work during the year (Census variable)
     "ANN_VAL",
     "PNSN_VAL",
+    "PTOTVAL",
     "OI_OFF",
     "OI_VAL",
     "CSP_VAL",

policyengine_us_data/datasets/cps/cps.py

@@ -3,7 +3,16 @@
 from policyengine_core.data import Dataset
 from policyengine_us_data.storage import STORAGE_FOLDER, DOCS_FOLDER
 import h5py
-from policyengine_us_data.datasets.cps.census_cps import *
+from policyengine_us_data.datasets.cps.census_cps import (
+    CensusCPS,
+    CensusCPS_2018,
+    CensusCPS_2019,
+    CensusCPS_2020,
+    CensusCPS_2021,
+    CensusCPS_2022,
+    CensusCPS_2023,
+    CensusCPS_2024,
+)
 from pandas import DataFrame, Series
 import numpy as np
 import pandas as pd
@@ -138,6 +147,7 @@ def generate(self):
             person, tax_unit, family, spm_unit, household = [
                 raw_data[entity] for entity in ENTITIES
             ]
+        _validate_raw_cps_schema(person, tax_unit, self.raw_cps.name)
 
         logging.info("Adding ID variables")
         add_id_variables(cps, person, tax_unit, family, spm_unit, household)
@@ -562,6 +572,33 @@ def uprate_cps_data(data, from_period, to_period):
     return data
 
 
+def _validate_raw_cps_schema(
+    person: DataFrame,
+    tax_unit: DataFrame,
+    raw_cps_name: str,
+) -> None:
+    required_person_columns = {
+        "CENSUS_TAX_ID",
+    }
+    required_tax_unit_columns = set()
+
+    missing_person = sorted(required_person_columns - set(person.columns))
+    missing_tax_unit = sorted(required_tax_unit_columns - set(tax_unit.columns))
+    if not missing_person and not missing_tax_unit:
+        return
+
+    missing_parts = []
+    if missing_person:
+        missing_parts.append("person: " + ", ".join(missing_person))
+    if missing_tax_unit:
+        missing_parts.append("tax_unit: " + ", ".join(missing_tax_unit))
+
+    raise ValueError(
+        f"Raw CPS dataset {raw_cps_name} is stale and must be regenerated; "
+        f"missing constructed tax-unit columns ({'; '.join(missing_parts)})."
+    )
+
+
 def add_id_variables(
     cps: h5py.File,
     person: DataFrame,
@@ -717,8 +754,12 @@ def children_per_parent(col: str) -> pd.DataFrame:
     cps["is_surviving_spouse"] = person.A_MARITL == 4
     cps["is_separated"] = person.A_MARITL == 6
     # High school or college/university enrollment status.
-    cps["is_full_time_college_student"] = person.A_HSCOL == 2
-
+    if "A_FTPT" in person.columns:
+        cps["is_full_time_college_student"] = (person.A_HSCOL == 2) & (
+            person.A_FTPT == 1
+        )
+    else:
+        cps["is_full_time_college_student"] = person.A_HSCOL == 2
     cps["detailed_occupation_recode"] = person.POCCU2
     cps["treasury_tipped_occupation_code"] = derive_treasury_tipped_occupation_code(
         person.PEIOOCC