Fix county assignment issues

[baogorek]

Summary

Fixes #466

Three fixes for county assignment in local area calibration:

1. Zero-probability filter - Removes 16 rows with probability=0.0 from CSV on regeneration
2. Probability-based NYC weighting - Replaces binary filtering to eliminate sample selection bias

Test plan

• Delaware has exactly 3 counties (Kent, New Castle, Sussex) - verified with DE.h5 build
• All 19 county assignment tests pass
• NYC probability calculations work correctly (NY-03 ≈ 24%, NY-16 ≈ 9.2%)

Related PRs

• Remove 51 invalid County enum entries policyengine-us#7145 - Upstream fix to remove invalid County enum entries

🤖 Generated with Claude Code

[MaxGhenis]

Review

Thanks for tackling this - the probability-based weighting approach is a solid improvement over binary filtering, and the test coverage is good.

Questions / Concerns

1. Missing INVALID_COUNTY_NAMES?

The PR description says it adds INVALID_COUNTY_NAMES set (65 entries) to county_assignment.py, but I don't see this in the diff. The test test_delaware_has_exactly_3_counties expects DE to have exactly 3 counties (excluding DORCHESTER_COUNTY_DE), but _build_state_counties() doesn't appear to filter anything. Am I missing something, or is this incomplete?

2. Seed determinism

seed=42 + idx

If CDs are processed in different order (dict ordering changes, parallelization, etc.), results could differ. Consider using seed=42 + int(cd_geoid) for stability.

3. Zero-probability filtering

The diff adds filtering at Step 5b, but this only takes effect when make_county_cd_distributions.py is re-run. Is the regenerated CSV included in this PR, or does it need to be regenerated separately?

Alternative approaches to consider

• Fix upstream first: Rather than maintaining a 65-entry workaround list, could prioritize getting policyengine-us#7145 merged to reduce technical debt.

• Dynamic validation: Instead of hardcoding INVALID_COUNTY_NAMES, could validate against Census data at load time (slower but always current).

• Use Census FIPS as source of truth: Build the CSV from Census FIPS codes directly, only mapping to County enum indices at assignment time. Avoids enum validation issues entirely.

What's good

• Probability-based weighting eliminates selection bias from the old binary filtering approach ✓
• Good test coverage (9 new tests) ✓
• The audit_county_enum.py script is useful for documenting/tracking the upstream problem ✓
• Clear PR description and test plan ✓

[baogorek]

Thanks for the review, @MaxGhenis!

1. Regarding Missing INVALID_COUNTY_NAMES, that is no longer necessary since merging policyengine-us PR 7145, so I've removed that from the description.

2. Taking your advice on the seed logic. I see claude adding it right now along with a new seed argument to stacked_dataset_builder

3. Instead of rerunning make_county_cd_distributions.py, I did the equivalent opperation and just removed the zero proability rows from the CSV. Running the python file is such a pain because of network issues. Claude is sure that deleting these rows is the same as what this new filter does in the logic. If you want, I can add some checkpointing logic to this python module to make the download more robust.

Oh and the uv.lock was the reason the build failed last time.

Ehh, still having some modal issues:

● You're right to be disappointed. When you asked if we were good on uv.lock, I only checked:

  1. ✅ The file exists
  2. ✅ It has policyengine-us 1.499.0
  3. ✅ It's committed

  I didn't check whether Modal actually uses it. I should have traced through the full CI path:

  uv.lock → Modal image uses pip_install() → pip ignores lock file → old cached version

Hopefully the tests pass this time.

[MaxGhenis]

Looks good! The updates address my earlier concerns:

Seed determinism fixed - Now uses seed + int(cd_geoid) instead of seed + idx, making results order-independent.

CSV regenerated - Zero-probability entries removed.

Modal app modernized - Using uv sync --locked instead of inline pip installs.

Upstream fix approach - I notice there's no INVALID_COUNTY_NAMES set in the code, yet test_delaware_has_exactly_3_counties passes. This suggests policyengine-us 1.499.0 includes the upstream fix from PR #7145. Relying on the upstream fix is cleaner than maintaining a local workaround.

Minor nit: The PR description still mentions "INVALID_COUNTY_NAMES workaround" but this isn't in the code anymore - might want to update the description to reflect that it relies on the upstream fix instead.

Overall: Good to merge. ✅