Require raw ACS allocation flags

Author: MaxGhenis

policyengine_us_data/datasets/acs/acs.py

@@ -5,6 +5,7 @@
     construct_tax_units_acs,
 )
 from policyengine_us_data.storage import STORAGE_FOLDER
+from policyengine_us_data.utils.source_quality import require_columns_present
 from pandas import DataFrame
 import numpy as np
 import pandas as pd
@@ -76,21 +77,24 @@ def add_person_variables(
             .loc[person["household_id"]][["RNTP", "TAXAMT"]]
             .values
         )
-        for source_flag, output_flag in [
+        allocation_flag_columns = [
             ("FRNTP", "rent_is_allocated"),
             ("FTAXP", "real_estate_taxes_is_allocated"),
-        ]:
-            if source_flag in household:
-                person[output_flag] = (
-                    household.set_index("household_id")
-                    .loc[person["household_id"]][source_flag]
-                    .fillna(0)
-                    .astype(int)
-                    .ne(0)
-                    .values
-                )
-            else:
-                person[output_flag] = False
+        ]
+        require_columns_present(
+            household.columns,
+            [source_flag for source_flag, _ in allocation_flag_columns],
+            source_name="raw Census ACS household table",
+        )
+        for source_flag, output_flag in allocation_flag_columns:
+            person[output_flag] = (
+                household.set_index("household_id")
+                .loc[person["household_id"]][source_flag]
+                .fillna(0)
+                .astype(int)
+                .ne(0)
+                .values
+            )
         acs["is_household_head"] = person.SPORDER == 1
         factor = person.SPORDER == 1
         person.rent *= factor * 12

tests/unit/datasets/test_acs_tax_unit_construction.py

@@ -235,3 +235,45 @@ def test_acs_add_id_variables_writes_related_to_head_or_spouse():
 
     assert person_tax_unit_id.tolist() == [1, 1]
     assert is_related_to_head_or_spouse.tolist() == [True, False]
+
+
+def test_acs_add_person_variables_requires_raw_allocation_flags():
+    person = _acs_person_fixture(household_id=[0])
+    household = pd.DataFrame(
+        {
+            "household_id": [0],
+            "RNTP": [1_200],
+            "TAXAMT": [500],
+            "TEN": [3],
+        }
+    )
+
+    with h5py.File("memory", mode="w", driver="core", backing_store=False) as acs:
+        with pytest.raises(KeyError, match="FRNTP"):
+            ACS.add_person_variables(acs, person, household)
+
+
+def test_acs_add_person_variables_writes_allocation_flags_for_heads_only():
+    person = _acs_person_fixture(
+        household_id=[0, 0],
+        SPORDER=[1, 2],
+        AGEP=[40, 10],
+    )
+    household = pd.DataFrame(
+        {
+            "household_id": [0],
+            "RNTP": [1_200],
+            "TAXAMT": [500],
+            "FRNTP": [1],
+            "FTAXP": [0],
+            "TEN": [3],
+        }
+    )
+
+    with h5py.File("memory", mode="w", driver="core", backing_store=False) as acs:
+        ACS.add_person_variables(acs, person, household)
+        rent_is_allocated = acs["rent_is_allocated"][:]
+        real_estate_taxes_is_allocated = acs["real_estate_taxes_is_allocated"][:]
+
+    assert rent_is_allocated.tolist() == [True, False]
+    assert real_estate_taxes_is_allocated.tolist() == [False, False]