Merge pull request #570 from PolicyEngine/add-dataset-sanity-tests

Harden data pipeline against corrupted dataset uploads

Author: baogorek

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all format test install download upload docker documentation data calibrate publish-local-area clean build paper clean-paper presentations database database-refresh promote-database promote-dataset
+.PHONY: all format test install download upload docker documentation data validate-data calibrate publish-local-area clean build paper clean-paper presentations database database-refresh promote-database promote-dataset
 
 HF_CLONE_DIR ?= $(HOME)/huggingface/policyengine-us-data
 
@@ -25,7 +25,7 @@ upload:
 
 docker:
 	docker buildx build --platform linux/amd64 . -t policyengine-us-data:latest
-	
+
 documentation:
 	cd docs && \
 	rm -rf _build .jupyter_cache && \
@@ -101,6 +101,9 @@ calibrate: data
 publish-local-area:
 	python policyengine_us_data/datasets/cps/local_area_calibration/publish_local_area.py
 
+validate-data:
+	python -c "from policyengine_us_data.storage.upload_completed_datasets import validate_all_datasets; validate_all_datasets()"
+
 clean:
 	rm -f policyengine_us_data/storage/*.h5
 	rm -f policyengine_us_data/storage/*.db

changelog.d/add-dataset-sanity-tests.added.md

@@ -0,0 +1 @@
+Hardened data pipeline against corrupted dataset uploads: pre-upload validation gate, post-generation assertions in enhanced CPS and sparse builders, CI workflow safety guards, file size checks, and comprehensive sanity tests for all dataset variants (5 layers of defense).

oregon_ctc_analysis.py

@@ -0,0 +1,202 @@
+"""
+Oregon Child Tax Credit Analysis by State Senate District
+
+Calculates the impact of doubling Oregon's Young Child Tax Credit (or_ctc)
+by State Legislative District Upper (SLDU) - i.e., State Senate districts.
+"""
+
+import numpy as np
+import pandas as pd
+from pathlib import Path
+from policyengine_us import Microsimulation
+from policyengine_core.reforms import Reform
+
+# Local imports
+from policyengine_us_data.datasets.cps.local_area_calibration.block_assignment import (
+    assign_geography_for_cd,
+    load_block_crosswalk,
+)
+from policyengine_us_data.storage import STORAGE_FOLDER
+
+# Oregon congressional districts (119th Congress)
+# Oregon has 6 CDs, geoid format: state_fips * 100 + district
+# Oregon FIPS = 41, so: 4101, 4102, 4103, 4104, 4105, 4106
+OREGON_CD_GEOIDS = [4101, 4102, 4103, 4104, 4105, 4106]
+
+
+def load_district_data(cd_geoid: int) -> dict:
+    """Load household data from a district H5 file."""
+    h5_path = STORAGE_FOLDER / "districts" / f"OR-{cd_geoid % 100:02d}.h5"
+    if not h5_path.exists():
+        raise FileNotFoundError(f"District file not found: {h5_path}")
+
+    import h5py
+
+    data = {}
+    with h5py.File(h5_path, "r") as f:
+        # Get key variables we need
+        for var in [
+            "household_weight",
+            "household_id",
+            "person_id",
+            "age",
+            "is_tax_unit_head",
+            "tax_unit_id",
+        ]:
+            if var in f:
+                # Handle year dimension if present
+                arr = f[var][:]
+                if len(arr.shape) > 1:
+                    arr = arr[:, 0]  # Take first year
+                data[var] = arr
+    return data
+
+
+def run_oregon_ctc_analysis():
+    """Run the Oregon CTC analysis by state senate district."""
+    print("=" * 60)
+    print("Oregon Child Tax Credit Analysis by State Senate District")
+    print("=" * 60)
+
+    # Load block crosswalk for SLDU lookups
+    print("\nLoading block crosswalk...")
+    crosswalk = load_block_crosswalk()
+    oregon_blocks = crosswalk[crosswalk["block_geoid"].str[:2] == "41"]
+    print(f"  Oregon blocks: {len(oregon_blocks):,}")
+    print(f"  Unique SLDUs: {oregon_blocks['sldu'].nunique()}")
+
+    # Results accumulator
+    results_by_sldu = {}
+
+    print("\nProcessing Oregon congressional districts...")
+
+    for cd_geoid in OREGON_CD_GEOIDS:
+        cd_name = f"OR-{cd_geoid % 100:02d}"
+        print(f"\n  Processing {cd_name}...")
+
+        # Load district data
+        h5_path = STORAGE_FOLDER / "districts" / f"{cd_name}.h5"
+        if not h5_path.exists():
+            print(f"    Skipping - file not found")
+            continue
+
+        # Run microsimulation for this district
+        # Baseline
+        baseline = Microsimulation(dataset=str(h5_path))
+        baseline_ctc = baseline.calculate("or_ctc", 2024)
+        baseline_weights = baseline.calculate("household_weight", 2024)
+
+        # Reform: double the OR CTC max amounts
+        # or_young_child_tax_credit_max is the parameter
+        def double_or_ctc(parameters):
+            # Double the max credit amount
+            or_ctc = parameters.gov.states.or_.tax.income.credits.ctc
+            or_ctc.amount.update(
+                start=pd.Timestamp("2024-01-01"),
+                stop=pd.Timestamp("2100-12-31"),
+                value=or_ctc.amount("2024-01-01") * 2,
+            )
+            return parameters
+
+        class DoubleORCTC(Reform):
+            def apply(self):
+                self.modify_parameters(double_or_ctc)
+
+        reform = Microsimulation(dataset=str(h5_path), reform=DoubleORCTC)
+        reform_ctc = reform.calculate("or_ctc", 2024)
+
+        # Get number of households for block assignment
+        n_households = len(baseline_weights)
+        print(f"    Households: {n_households:,}")
+
+        # Assign blocks and get SLDU for each household
+        geo = assign_geography_for_cd(
+            cd_geoid=str(cd_geoid),
+            n_households=n_households,
+            seed=cd_geoid,  # Reproducible
+        )
+
+        sldu_assignments = geo["sldu"]
+
+        # Calculate impact per household
+        impact = reform_ctc - baseline_ctc
+
+        # Aggregate by SLDU
+        unique_sldus = np.unique(sldu_assignments[sldu_assignments != ""])
+
+        for sldu in unique_sldus:
+            mask = sldu_assignments == sldu
+            sldu_impact = np.sum(impact[mask] * baseline_weights[mask])
+            sldu_baseline = np.sum(baseline_ctc[mask] * baseline_weights[mask])
+            sldu_reform = np.sum(reform_ctc[mask] * baseline_weights[mask])
+            sldu_hh = np.sum(mask)
+            sldu_weighted_hh = np.sum(baseline_weights[mask])
+
+            if sldu not in results_by_sldu:
+                results_by_sldu[sldu] = {
+                    "baseline_ctc": 0,
+                    "reform_ctc": 0,
+                    "impact": 0,
+                    "households": 0,
+                    "weighted_households": 0,
+                }
+
+            results_by_sldu[sldu]["baseline_ctc"] += sldu_baseline
+            results_by_sldu[sldu]["reform_ctc"] += sldu_reform
+            results_by_sldu[sldu]["impact"] += sldu_impact
+            results_by_sldu[sldu]["households"] += sldu_hh
+            results_by_sldu[sldu]["weighted_households"] += sldu_weighted_hh
+
+    # Create results DataFrame
+    print("\n" + "=" * 60)
+    print("RESULTS: Impact of Doubling Oregon CTC by State Senate District")
+    print("=" * 60)
+
+    df = pd.DataFrame.from_dict(results_by_sldu, orient="index")
+    df.index.name = "sldu"
+    df = df.reset_index()
+
+    # Convert to millions
+    df["baseline_ctc_millions"] = df["baseline_ctc"] / 1e6
+    df["reform_ctc_millions"] = df["reform_ctc"] / 1e6
+    df["impact_millions"] = df["impact"] / 1e6
+
+    # Sort by impact
+    df = df.sort_values("impact_millions", ascending=False)
+
+    # Display results
+    print(
+        f"\n{'SLDU':<8} {'Baseline':>12} {'Reform':>12} {'Impact':>12} {'Households':>12}"
+    )
+    print(f"{'':8} {'($M)':>12} {'($M)':>12} {'($M)':>12} {'(weighted)':>12}")
+    print("-" * 60)
+
+    for _, row in df.iterrows():
+        print(
+            f"{row['sldu']:<8} "
+            f"{row['baseline_ctc_millions']:>12.2f} "
+            f"{row['reform_ctc_millions']:>12.2f} "
+            f"{row['impact_millions']:>12.2f} "
+            f"{row['weighted_households']:>12,.0f}"
+        )
+
+    print("-" * 60)
+    total_baseline = df["baseline_ctc_millions"].sum()
+    total_reform = df["reform_ctc_millions"].sum()
+    total_impact = df["impact_millions"].sum()
+    total_hh = df["weighted_households"].sum()
+    print(
+        f"{'TOTAL':<8} {total_baseline:>12.2f} {total_reform:>12.2f} "
+        f"{total_impact:>12.2f} {total_hh:>12,.0f}"
+    )
+
+    # Save to CSV
+    output_path = Path("oregon_ctc_by_sldu.csv")
+    df.to_csv(output_path, index=False)
+    print(f"\nResults saved to: {output_path}")
+
+    return df
+
+
+if __name__ == "__main__":
+    run_oregon_ctc_analysis()

policyengine_us_data/datasets/cps/enhanced_cps.py

@@ -201,6 +201,31 @@ def generate(self):
             )
             data["household_weight"][year] = optimised_weights
 
+            # Validate dense weights
+            w = optimised_weights
+            if np.any(np.isnan(w)):
+                raise ValueError(
+                    f"Year {year}: household_weight contains NaN values"
+                )
+            if np.any(w < 0):
+                raise ValueError(
+                    f"Year {year}: household_weight contains negative values"
+                )
+            weighted_hh_count = float(np.sum(w))
+            if not (1e8 <= weighted_hh_count <= 2e8):
+                raise ValueError(
+                    f"Year {year}: weighted household count "
+                    f"{weighted_hh_count:,.0f} outside expected range "
+                    f"[100M, 200M]"
+                )
+            logging.info(
+                f"Year {year}: weights validated — "
+                f"{weighted_hh_count:,.0f} weighted households, "
+                f"{int(np.sum(w > 0))} non-zero"
+            )
+
+        logging.info("Post-generation weight validation passed")
+
         self.save_dataset(data)
 
 

policyengine_us_data/datasets/cps/small_enhanced_cps.py

@@ -1,3 +1,4 @@
+import os
 import pandas as pd
 import numpy as np
 import h5py
@@ -17,6 +18,19 @@ def create_small_ecps():
     )
     simulation.subsample(1_000)
 
+    # Basic validation that subsample has reasonable data
+    weights = simulation.calculate("household_weight").values
+    if np.all(weights == 0):
+        raise ValueError(
+            "create_small_ecps: all household weights are zero "
+            "after subsample"
+        )
+    logging.info(
+        f"create_small_ecps: subsample has "
+        f"{len(weights)} households, "
+        f"{int(np.sum(weights > 0))} with non-zero weight"
+    )
+
     data = {}
     for variable in simulation.tax_benefit_system.variables:
         data[variable] = {}
@@ -75,6 +89,16 @@ def create_sparse_ecps():
     h_ids = h_ids[h_weights > 0]
     h_weights = h_weights[h_weights > 0]
 
+    if len(h_ids) < 1000:
+        raise ValueError(
+            f"create_sparse_ecps: only {len(h_ids)} households with "
+            f"non-zero weight (expected > 1000)"
+        )
+    logging.info(
+        f"create_sparse_ecps: {len(h_ids)} households after "
+        f"zero-weight filtering"
+    )
+
     subset_df = df[df[df_household_id_column].isin(h_ids)].copy()
 
     # Update the dataset and rebuild the simulation
@@ -104,12 +128,38 @@ def create_sparse_ecps():
             if len(data[variable]) == 0:
                 del data[variable]
 
-    with h5py.File(STORAGE_FOLDER / "sparse_enhanced_cps_2024.h5", "w") as f:
+    # Validate critical variables exist before writing
+    critical_vars = [
+        "household_weight",
+        "employment_income",
+        "household_id",
+        "person_id",
+    ]
+    missing = [v for v in critical_vars if v not in data]
+    if missing:
+        raise ValueError(
+            f"create_sparse_ecps: missing critical variables: {missing}"
+        )
+    logging.info(f"create_sparse_ecps: data dict has {len(data)} variables")
+
+    output_path = STORAGE_FOLDER / "sparse_enhanced_cps_2024.h5"
+    with h5py.File(output_path, "w") as f:
         for variable, periods in data.items():
             grp = f.create_group(variable)
             for period, values in periods.items():
                 grp.create_dataset(str(period), data=values)
 
+    file_size = os.path.getsize(output_path)
+    if file_size < 1_000_000:
+        raise ValueError(
+            f"create_sparse_ecps: output file only {file_size:,} bytes "
+            f"(expected > 1MB)"
+        )
+    logging.info(
+        f"create_sparse_ecps: wrote {file_size / 1e6:.1f}MB to "
+        f"{output_path}"
+    )
+
 
 if __name__ == "__main__":
     create_small_ecps()