Add stacked dataset builder and P(county|CD) distributions (#457)

* Add sparse matrix builder for local area calibration

Core components:
- sparse_matrix_builder.py: Database-driven approach for building calibration matrices
- calibration_utils.py: Shared utilities (cache clearing, constraints, geo helpers)
- matrix_tracer.py: Debugging utility for tracing through sparse matrices
- create_stratified_cps.py: Create stratified sample preserving high-income households
- test_sparse_matrix_builder.py: 6 verification tests for matrix correctness

Data pipeline changes:
- Add GEO_STACKING env var to cps.py and puf.py for geo-stacking data generation
- Add GEO_STACKING_MODE env var to extended_cps.py
- Add CPS_2024_Full, PUF_2023, ExtendedCPS_2023 classes
- Add policy_data.db download to prerequisites
- Add 'make data-geo' target for geo-stacking data pipeline

CI/CD:
- Add geo-stacking dataset build step to workflow
- Add sparse matrix builder test step after geo data generation

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add changelog entry and format code

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Refactor tests and fix enum encoding, minimize PR scope

- Move sparse matrix tests to tests/test_local_area_calibration/
- Split large test file into focused modules (column indexing, same-state,
  cross-state, geo masking)
- Fix small_enhanced_cps.py enum encoding (decode_to_str before astype)
- Fix create_stratified_cps.py to use local storage instead of HuggingFace
- Remove CPS_2024_Full to keep PR minimal
- Revert ExtendedCPS_2024 to use CPS_2024

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Rename GEO_STACKING to LOCAL_AREA_CALIBRATION and restore tracer functionality

- Rename GEO_STACKING to LOCAL_AREA_CALIBRATION in cps.py, puf.py, extended_cps.py
- Rename data-geo to data-local-area in Makefile and workflow
- Add create_target_groups function to calibration_utils.py
- Enhance MatrixTracer with get_group_rows method and variable_desc in row catalog
- Add TARGET GROUPS section to print_matrix_structure output
- Add local_area_calibration_setup.ipynb documentation notebook

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Clear notebook outputs for Myst compatibility

* Pin mystmd>=1.7.0 to fix notebook rendering in docs

* Add population-weighted P(county|CD) distributions and stacked dataset builder

- Add make_county_cd_distributions.py to compute P(county|CD) from Census block data
- Add county_cd_distributions.csv with distributions for all 436 CDs
- Add county_assignment.py module for assigning counties to households
- Add stacked_dataset_builder.py for creating CD-stacked H5 datasets
- Add tests for county assignment functionality
- Update calibration_utils.py with state/CD mapping utilities

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add local area H5 publishing workflow

- New GitHub Actions workflow (local_area_publish.yaml) that:
  - Triggers on local_area_calibration/ changes, repository_dispatch, or manual
  - Downloads calibration inputs from HF calibration/ folder
  - Builds 51 state + 436 district H5 files with checkpointing
  - Uploads to GCP and HF states/ and districts/ subdirectories

- New publish_local_area.py script with:
  - Per-state and per-district checkpointing for spot instance resilience
  - Immediate upload after each file is built
  - Support for --states-only, --districts-only, --skip-download flags

- Added upload_local_area_file() to data_upload.py for subdirectory uploads
- Added download_calibration_inputs() to huggingface.py
- Added publish-local-area Makefile target

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Update paths to use calibration/ subdirectory for policy_data.db

- download_private_prerequisites.py: Download from calibration/policy_data.db
- calibration_utils.py: Look for db in storage/calibration/
- conftest.py: Update test fixture path
- huggingface.py: Fix download_calibration_inputs to return correct paths

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* documentation updates

* Add deterministic test fixture and tests for stacked_dataset_builder

Create a minimal 50-household H5 fixture with known values for stable testing
of the stacked dataset builder without relying on sampled stratified CPS data.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Fix dtype warning in stacked_dataset_builder person ID assignment

Cast np.arange output to int32 to match column dtype.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Format test files with black

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add spm_unit_tenure_type and fix input_variables detection in stratification

- Add spm_unit_tenure_type mapping from SPM_TENMORTSTATUS in add_spm_variables
- Fix create_stratified_cps.py to use source sim's input_variables instead of empty sim
- Fix stacked_dataset_builder.py to use base_sim's input_variables instead of sparse_sim

The input_variables fix ensures variables like spm_unit_tenure_type are preserved
when creating stratified/stacked datasets, since input_variables is only populated
from variables that have actual data in the loaded dataset.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add real rent-based SPM geographic adjustments to stacked dataset builder

- Add spm-calculator integration for SPM threshold calculation
- Replace random placeholder geoadj with real values from Census ACS rent data
- Add load_cd_geoadj_values() to compute geoadj from median 2BR rents
- Add calculate_spm_thresholds_for_cd() to calculate SPM thresholds per CD
- Add CD rent data CSV and fetch script (requires CENSUS_API_KEY)
- Update .gitignore to track rent CSV

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* NYC workflow

* Add batched HuggingFace uploads and fix at-large district geoadj lookup

- Add upload_local_area_batch_to_hf() to batch multiple files per commit
- Add skip_hf parameter to upload_local_area_file() for GCP-only uploads
- Modify publish_local_area.py to batch HF uploads (10 files per commit)
- Fix at-large district geoadj lookup (XX01 -> XX00 mapping for AK, DE, etc.)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Filter pseudo-input variables from H5 output and add checkpoint files to gitignore

Pseudo-inputs are variables with adds/subtracts that aggregate formula-based
components. Saving their stale pre-computed values corrupts calculations when
the dataset is reloaded.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add spm-calculator as a dependency

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Fix database path and add download dependency

- Update notebook to use correct db path: storage/calibration/policy_data.db
- Add download as dependency of data target in Makefile

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Fix test fixture path to use absolute path

Use os.path.dirname(__file__) instead of relative path so tests
work regardless of working directory.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Trigger CI after runner restart

* Add uv.lock to pin dependency versions

- Add uv.lock file with all pinned dependencies
- Update all workflows to use `uv sync --dev` instead of pip install
- Add lock freshness check to PR workflow
- Narrow Python version to >=3.12 (required by microimpute)

This prevents stale cached packages on the self-hosted runner from
causing test failures (e.g., missing spm_unit_tenure_type variable).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Use uv run for all Python commands in workflows

uv sync creates a virtual environment, but commands were running
with system Python which still had stale cached packages.

All make/python/pytest commands now use `uv run` to execute within
the virtual environment where the locked dependencies are installed.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Fix family_id reindexing and vectorize entity ID assignment

- Add family_id/person_family_id to entity reindexing loop
  (was missing, causing ID collisions when same household appears in multiple CDs)
- Vectorize entity reindexing using groupby().ngroup() instead of O(n²) nested loops
- Add comment explaining why fresh Microsimulation per CD is necessary
- Add tests for entity ID uniqueness across stacked CDs

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Update uv.lock with latest package versions

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Format with Black

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Max Ghenis <mghenis@gmail.com>

Author: 3 people

.github/workflows/code_changes.yaml

@@ -37,9 +37,9 @@ jobs:
       - name: Install uv
         uses: astral-sh/setup-uv@v5
       - name: Install package
-        run: uv pip install -e .[dev] --system
+        run: uv sync --dev
       - name: Build package
-        run: python -m build
+        run: uv run python -m build
       - name: Publish a Python distribution to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:

.github/workflows/local_area_publish.yaml

@@ -0,0 +1,70 @@
+name: Publish Local Area H5 Files
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - 'policyengine_us_data/datasets/cps/local_area_calibration/**'
+      - '.github/workflows/local_area_publish.yaml'
+  repository_dispatch:
+    types: [calibration-updated]
+  workflow_dispatch:
+
+# Trigger strategy:
+# 1. Automatic: Code changes to local_area_calibration/ pushed to main
+# 2. repository_dispatch: Calibration workflow triggers after uploading new weights
+# 3. workflow_dispatch: Manual trigger when you update weights/data on HF yourself
+
+jobs:
+  publish-local-area:
+    runs-on: self-hosted
+    permissions:
+      contents: read
+      id-token: write
+    env:
+      HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
+
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.13'
+
+      - name: Authenticate to Google Cloud
+        uses: google-github-actions/auth@v2
+        with:
+          workload_identity_provider: "projects/322898545428/locations/global/workloadIdentityPools/policyengine-research-id-pool/providers/prod-github-provider"
+          service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
+
+      - name: Install package
+        run: uv sync --dev
+
+      - name: Download checkpoint (if exists)
+        continue-on-error: true
+        run: |
+          gsutil cp gs://policyengine-us-data/checkpoints/completed_states.txt . || true
+          gsutil cp gs://policyengine-us-data/checkpoints/completed_districts.txt . || true
+          gsutil cp gs://policyengine-us-data/checkpoints/completed_cities.txt . || true
+
+      - name: Build and publish local area H5 files
+        run: uv run make publish-local-area
+
+      - name: Upload checkpoint
+        if: always()
+        run: |
+          gsutil cp completed_states.txt gs://policyengine-us-data/checkpoints/ || true
+          gsutil cp completed_districts.txt gs://policyengine-us-data/checkpoints/ || true
+          gsutil cp completed_cities.txt gs://policyengine-us-data/checkpoints/ || true
+
+      - name: Clean up checkpoints on success
+        if: success()
+        run: |
+          gsutil rm gs://policyengine-us-data/checkpoints/completed_states.txt || true
+          gsutil rm gs://policyengine-us-data/checkpoints/completed_districts.txt || true
+          gsutil rm gs://policyengine-us-data/checkpoints/completed_cities.txt || true

.github/workflows/pr_code_changes.yaml

@@ -30,8 +30,28 @@ jobs:
           fi
           echo "✅ PR is from the correct repository"
 
-  Lint:
+  check-lock-freshness:
+    name: Check uv.lock freshness
+    runs-on: ubuntu-latest
     needs: check-fork
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.13'
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Check lock file is up-to-date
+        run: |
+          uv lock --upgrade
+          git diff --exit-code uv.lock || {
+            echo "::error::uv.lock is outdated. Run 'uv lock --upgrade' and commit the changes."
+            exit 1
+          }
+
+  Lint:
+    needs: [check-fork, check-lock-freshness]
     uses: ./.github/workflows/reusable_lint.yaml
 
   SmokeTestForMultipleVersions:

.github/workflows/reusable_test.yaml

@@ -57,11 +57,11 @@ jobs:
           service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
 
       - name: Install package
-        run: uv pip install -e .[dev] --system
+        run: uv sync --dev
 
       - name: Download data inputs
         if: inputs.full_suite
-        run: make download
+        run: uv run make download
 
       # Temporarily disabled - database target causing issues
       # - name: Create and load calibration targets database 
@@ -70,22 +70,22 @@ jobs:
 
       - name: Build datasets
         if: inputs.full_suite
-        run: make data
+        run: uv run make data
         env:
           TEST_LITE: ${{ !inputs.upload_data }}
           PYTHON_LOG_LEVEL: INFO
 
       - name: Build datasets for local area calibration
         if: inputs.full_suite
         run: |
-          LOCAL_AREA_CALIBRATION=true python policyengine_us_data/datasets/cps/cps.py
-          LOCAL_AREA_CALIBRATION=true python policyengine_us_data/datasets/puf/puf.py
-          LOCAL_AREA_CALIBRATION=true python policyengine_us_data/datasets/cps/extended_cps.py
-          python policyengine_us_data/datasets/cps/local_area_calibration/create_stratified_cps.py 10500
+          LOCAL_AREA_CALIBRATION=true uv run python policyengine_us_data/datasets/cps/cps.py
+          LOCAL_AREA_CALIBRATION=true uv run python policyengine_us_data/datasets/puf/puf.py
+          LOCAL_AREA_CALIBRATION=true uv run python policyengine_us_data/datasets/cps/extended_cps.py
+          uv run python policyengine_us_data/datasets/cps/local_area_calibration/create_stratified_cps.py 10500
 
       - name: Run local area calibration tests
         if: inputs.full_suite
-        run: pytest policyengine_us_data/tests/test_local_area_calibration/ -v
+        run: uv run pytest policyengine_us_data/tests/test_local_area_calibration/ -v
 
       - name: Save calibration log
         if: inputs.full_suite
@@ -95,14 +95,14 @@ jobs:
           path: calibration_log.csv
 
       - name: Run tests
-        run: pytest
+        run: uv run pytest
 
       - name: Upload data
         if: inputs.upload_data
-        run: make upload
+        run: uv run make upload
 
       - name: Test documentation builds
-        run: make documentation
+        run: uv run make documentation
         env:
           BASE_URL: ${{ inputs.deploy_docs && '/policyengine-us-data' || '' }}
 

.gitignore

@@ -2,6 +2,7 @@
 **/__pycache__
 **/.DS_STORE
 **/*.h5
+**/*.npy
 **/*.csv
 **/_build
 **/*.pkl
@@ -23,4 +24,11 @@ node_modules
 !soi_targets.csv
 !policyengine_us_data/storage/social_security_aux.csv
 !policyengine_us_data/storage/SSPopJul_TR2024.csv
+!policyengine_us_data/storage/national_and_district_rents_2023.csv
 docs/.ipynb_checkpoints/
+
+## Batch processing checkpoints
+completed_*.txt
+
+## Test fixtures
+!policyengine_us_data/tests/test_local_area_calibration/test_fixture_50hh.h5

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all format test install download upload docker documentation data clean build paper clean-paper presentations
+.PHONY: all format test install download upload docker documentation data data-local-area publish-local-area clean build paper clean-paper presentations
 
 all: data test
 
@@ -62,7 +62,7 @@ database:
 	python policyengine_us_data/db/etl_irs_soi.py
 	python policyengine_us_data/db/validate_database.py
 
-data:
+data: download
 	python policyengine_us_data/utils/uprating.py
 	python policyengine_us_data/datasets/acs/acs.py
 	python policyengine_us_data/datasets/cps/cps.py
@@ -80,6 +80,9 @@ data-local-area: data
 	LOCAL_AREA_CALIBRATION=true python policyengine_us_data/datasets/cps/extended_cps.py
 	python policyengine_us_data/datasets/cps/local_area_calibration/create_stratified_cps.py 10500
 
+publish-local-area:
+	python policyengine_us_data/datasets/cps/local_area_calibration/publish_local_area.py
+
 clean:
 	rm -f policyengine_us_data/storage/*.h5
 	rm -f policyengine_us_data/storage/*.db

changelog_entry.yaml

@@ -0,0 +1,10 @@
+- bump: minor
+  changes:
+    added:
+    - Sparse matrix builder for local area calibration with database-driven constraints
+    - Local area calibration data pipeline (make data-local-area)
+    - ExtendedCPS_2023 and PUF_2023 dataset classes
+    - Stratified CPS sampling to preserve high-income households
+    - Matrix verification tests for local area calibration
+    - Population-weighted P(county|CD) distributions from Census block data
+    - County assignment module for stacked dataset builder