From da5a7eb4f8bddbfebc00cb6bc719decd19b6cdb0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 12 Oct 2026 19:31:35 -0400 Subject: [PATCH 1/2] Retire the legacy target store on the Thesis ledger branch Remove the SQLite strata/stratum_constraints/targets store and everything that only existed to fill or read it: the 19 db/etl_*.py loaders, the arch load/stats/query commands, the Supabase target helpers, arch.targets, scripts/export_to_json.py, the calibration/ package and the Microplex calibration modules in micro/us/ that consumed those rows. Live modules keep working: soi.py carries its AGI bracket constants, normalization keeps its lineage transforms, and db.cli keeps the source-artifact commands. The required "Arch checks" job no longer runs `arch load all`; it smoke-tests the source-artifact CLI instead. tests/test_legacy_targets_retired.py pins the removal and checks that the append-path scripts import nothing from this repository's packages. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 63 +- README.md | 74 +- arch/cli.py | 2 +- arch/jurisdictions/uk/__init__.py | 12 +- arch/jurisdictions/us/__init__.py | 38 +- arch/jurisdictions/us/soi.py | 46 +- arch/normalization/__init__.py | 8 +- arch/normalization/targets.py | 199 -- arch/normalization/transforms.py | 97 + arch/targets/__init__.py | 36 - arch/targets/client.py | 5 - arch/targets/loaders.py | 39 - calibration/__init__.py | 37 - calibration/constraints.py | 293 --- calibration/loader.py | 312 --- calibration/methods/__init__.py | 12 - calibration/methods/entropy.py | 139 -- calibration/targets.py | 139 -- calibration/variables.py | 237 --- data/targets/README.md | 18 +- db/__init__.py | 8 +- db/cli.py | 220 +-- db/etl_aca_enrollment.py | 643 ------ db/etl_bls.py | 252 --- db/etl_cbo.py | 366 ---- db/etl_census.py | 465 ----- db/etl_cps.py | 313 --- db/etl_cps_raw.py | 378 ---- db/etl_hmrc.py | 351 ---- db/etl_medicaid.py | 494 ----- db/etl_obr.py | 313 --- db/etl_ons.py | 310 --- db/etl_snap.py | 334 ---- db/etl_soi.py | 624 ------ db/etl_soi_credits.py | 1107 ----------- db/etl_soi_deductions.py | 859 -------- db/etl_soi_income_sources.py | 1582 --------------- db/etl_soi_state.py | 921 --------- db/etl_ssa.py | 388 ---- db/etl_ssi.py | 660 ------- db/etl_targets_supabase.py | 320 --- db/schema.py | 143 +- db/supabase_client.py | 102 +- db/tests/test_etl_cps_raw.py | 98 - docs/architecture.md | 45 +- docs/calibration-pipeline.md | 246 --- docs/cps-variable-coverage.md | 66 +- micro/us/calibrate.py | 389 ---- micro/us/district.py | 567 ------ micro/us/hierarchical_pipeline.py | 763 -------- micro/us/hierarchy_validation.py | 517 ----- micro/us/pipeline.py | 1851 ------------------ micro/us/targets.py | 738 ------- micro/us/validation_dashboard.py | 512 ----- pyproject.toml | 2 +- scripts/cps_to_policyengine.py | 44 +- scripts/export_to_json.py | 123 -- tests/test_arch_facts.py | 2 +- tests/test_arch_namespace.py | 27 +- tests/test_arch_normalization.py | 50 +- tests/test_calibration.py | 587 ------ tests/test_congressional_districts.py | 210 -- tests/test_download_cps.py | 54 +- tests/test_etl_aca_enrollment.py | 312 --- tests/test_etl_bls.py | 183 -- tests/test_etl_cbo.py | 170 -- tests/test_etl_census.py | 206 -- tests/test_etl_cps.py | 282 --- tests/test_etl_hmrc.py | 186 -- tests/test_etl_medicaid.py | 314 --- tests/test_etl_obr.py | 186 -- tests/test_etl_ons.py | 209 -- tests/test_etl_snap.py | 186 -- tests/test_etl_soi.py | 36 - tests/test_etl_soi_credits.py | 848 -------- tests/test_etl_soi_deductions.py | 468 ----- tests/test_etl_soi_income_sources.py | 517 ----- tests/test_etl_soi_state.py | 395 ---- tests/test_etl_ssa.py | 186 -- tests/test_etl_ssi.py | 366 ---- tests/test_hierarchical_calibration.py | 514 ----- tests/test_hierarchical_pipeline.py | 240 --- tests/test_integration_simple.py | 227 --- tests/test_legacy_targets_retired.py | 306 +++ tests/test_loader.py | 142 -- tests/test_microplex_hierarchy_validation.py | 172 -- tests/test_microplex_pipeline.py | 366 ---- tests/test_microplex_policyengine.py | 152 -- tests/test_microplex_target_adapter.py | 391 ---- tests/test_supabase_client.py | 32 - 90 files changed, 571 insertions(+), 26871 deletions(-) delete mode 100644 arch/normalization/targets.py create mode 100644 arch/normalization/transforms.py delete mode 100644 arch/targets/__init__.py delete mode 100644 arch/targets/client.py delete mode 100644 arch/targets/loaders.py delete mode 100644 calibration/__init__.py delete mode 100644 calibration/constraints.py delete mode 100644 calibration/loader.py delete mode 100644 calibration/methods/__init__.py delete mode 100644 calibration/methods/entropy.py delete mode 100644 calibration/targets.py delete mode 100644 calibration/variables.py delete mode 100644 db/etl_aca_enrollment.py delete mode 100644 db/etl_bls.py delete mode 100644 db/etl_cbo.py delete mode 100644 db/etl_census.py delete mode 100644 db/etl_cps.py delete mode 100644 db/etl_cps_raw.py delete mode 100644 db/etl_hmrc.py delete mode 100644 db/etl_medicaid.py delete mode 100644 db/etl_obr.py delete mode 100644 db/etl_ons.py delete mode 100644 db/etl_snap.py delete mode 100644 db/etl_soi.py delete mode 100644 db/etl_soi_credits.py delete mode 100644 db/etl_soi_deductions.py delete mode 100644 db/etl_soi_income_sources.py delete mode 100644 db/etl_soi_state.py delete mode 100644 db/etl_ssa.py delete mode 100644 db/etl_ssi.py delete mode 100644 db/etl_targets_supabase.py delete mode 100644 db/tests/test_etl_cps_raw.py delete mode 100644 docs/calibration-pipeline.md delete mode 100644 micro/us/calibrate.py delete mode 100644 micro/us/district.py delete mode 100644 micro/us/hierarchical_pipeline.py delete mode 100644 micro/us/hierarchy_validation.py delete mode 100644 micro/us/pipeline.py delete mode 100644 micro/us/targets.py delete mode 100644 micro/us/validation_dashboard.py delete mode 100644 scripts/export_to_json.py delete mode 100644 tests/test_calibration.py delete mode 100644 tests/test_congressional_districts.py delete mode 100644 tests/test_etl_aca_enrollment.py delete mode 100644 tests/test_etl_bls.py delete mode 100644 tests/test_etl_cbo.py delete mode 100644 tests/test_etl_census.py delete mode 100644 tests/test_etl_cps.py delete mode 100644 tests/test_etl_hmrc.py delete mode 100644 tests/test_etl_medicaid.py delete mode 100644 tests/test_etl_obr.py delete mode 100644 tests/test_etl_ons.py delete mode 100644 tests/test_etl_snap.py delete mode 100644 tests/test_etl_soi.py delete mode 100644 tests/test_etl_soi_credits.py delete mode 100644 tests/test_etl_soi_deductions.py delete mode 100644 tests/test_etl_soi_income_sources.py delete mode 100644 tests/test_etl_soi_state.py delete mode 100644 tests/test_etl_ssa.py delete mode 100644 tests/test_etl_ssi.py delete mode 100644 tests/test_hierarchical_calibration.py delete mode 100644 tests/test_hierarchical_pipeline.py delete mode 100644 tests/test_integration_simple.py create mode 100644 tests/test_legacy_targets_retired.py delete mode 100644 tests/test_loader.py delete mode 100644 tests/test_microplex_hierarchy_validation.py delete mode 100644 tests/test_microplex_pipeline.py delete mode 100644 tests/test_microplex_target_adapter.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 72593508..a5864774 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -88,42 +88,23 @@ jobs: arch tests/test_arch_bundle.py db/cli.py - db/etl_snap.py - db/etl_aca_enrollment.py - db/etl_medicaid.py - db/etl_ssi.py db/pe_source_inventory.py db/schema.py db/supabase_client.py db/source_files.py micro/us/census/download_cps.py micro/us/entities.py - micro/us/hierarchy_validation.py - micro/us/hierarchical_pipeline.py - micro/us/pipeline.py micro/us/policyengine.py - micro/us/targets.py - micro/us/validation_dashboard.py - calibration/constraints.py - calibration/targets.py scripts/build_series_catalog.py tests/test_build_series_catalog.py tests/test_arch_facts.py tests/test_arch_namespace.py tests/test_arch_normalization.py - tests/test_etl_snap.py - tests/test_etl_soi.py - tests/test_etl_aca_enrollment.py - tests/test_etl_medicaid.py - tests/test_etl_ssi.py + tests/test_legacy_targets_retired.py tests/test_source_files.py tests/test_microplex_entities.py - tests/test_microplex_hierarchy_validation.py - tests/test_microplex_target_adapter.py - tests/test_microplex_pipeline.py tests/test_microplex_policyengine.py tests/test_download_cps.py - tests/test_hierarchical_pipeline.py tests/test_supabase_client.py tests/test_policyengine_ledger.py tests/test_release_chain.py @@ -140,31 +121,21 @@ jobs: tests/test_arch_facts.py tests/test_arch_namespace.py tests/test_arch_normalization.py - tests/test_etl_snap.py - tests/test_etl_soi.py - tests/test_etl_aca_enrollment.py - tests/test_etl_medicaid.py - tests/test_etl_ssi.py + tests/test_legacy_targets_retired.py tests/test_source_files.py - tests/test_calibration.py tests/test_microplex_entities.py - tests/test_microplex_hierarchy_validation.py - tests/test_microplex_target_adapter.py - tests/test_microplex_pipeline.py tests/test_microplex_policyengine.py tests/test_download_cps.py - tests/test_hierarchical_pipeline.py tests/test_supabase_client.py tests/test_policyengine_ledger.py tests/test_release_chain.py tests/test_thesis_append_adversarial.py -q - - name: Build target input database + - name: Smoke-test the source-artifact database CLI run: | - uv run arch --db /tmp/arch-targets-ci.db init - uv run arch --db /tmp/arch-targets-ci.db load all - uv run arch --db /tmp/arch-targets-ci.db stats + uv run arch --db /tmp/arch-source-ci.db init + uv run arch --db /tmp/arch-source-ci.db source-stats - name: Build wheel run: | @@ -178,18 +149,28 @@ jobs: uv pip install --python /tmp/arch-wheel-smoke/bin/python "$wheel_path" cd /tmp /tmp/arch-wheel-smoke/bin/python - <<'PY' + import importlib.util + import arch - import arch.targets - from micro.us.district import DistrictMicroplex - from micro.us.pipeline import run_pipeline - from micro.us.targets import load_microplex_targets + from arch.jurisdictions.us.soi import AGI_BRACKETS + from arch.normalization import convert_units + from micro.us.entities import build_microplex_entities from microplex import ConditionalMAF assert arch.__name__ == "arch" - assert DistrictMicroplex is not None - assert run_pipeline is not None - assert load_microplex_targets is not None + assert len(AGI_BRACKETS) == 19 + assert convert_units is not None + assert build_microplex_entities is not None assert ConditionalMAF is not None + # The legacy target store is not shipped. + for retired in ( + "arch.targets", + "calibration", + "db.etl_soi", + "micro.us.pipeline", + "micro.us.targets", + ): + assert importlib.util.find_spec(retired) is None, retired PY /tmp/arch-wheel-smoke/bin/arch --help >/dev/null diff --git a/README.md b/README.md index d65fb107..740205a0 100644 --- a/README.md +++ b/README.md @@ -91,22 +91,20 @@ arch/ │ ├── sources/ # Source lineage helpers │ ├── facts/ # Source-backed facts │ ├── normalization/ # Low-assumption representation helpers -│ ├── targets/ # Target input schema, client, loaders │ ├── microdata/ # Microdata registry, ingestion, queries │ ├── jurisdictions/ # Temporary in-repo jurisdiction source prototypes -├── db/ # SQLModel persistence and source loaders -│ ├── schema.py # SQLModel: Target, Stratum, StratumConstraint +├── db/ # Source-artifact store and packaged source files +│ ├── schema.py # SQLModel: SourceArtifact, SourceTable, SourceRow │ ├── supabase_client.py # Supabase client helpers -│ └── etl_*.py # Source-specific ETL pipelines -├── micro/ # Microplex consumers of Arch records -├── calibration/ # Calibration target adapters and constraints +│ └── data/ # Publisher files that source packages read +├── micro/ # Microplex microdata utilities ├── data/ # Cached data files └── docs/ # Architecture and source documentation ``` New code should prefer `arch.sources`, `arch.facts`, `arch.normalization`, -`arch.targets`, and `arch.microdata`. Microplex-specific target composition -and calibration code belongs under `micro/`. +and `arch.microdata`. Target composition and calibration belong in Microplex, +not in this repository. ## Quick Start @@ -120,15 +118,7 @@ cd arch pip install -e ".[dev]" ``` -### 2. Initialize and Load Legacy Target Inputs - -```bash -arch init -arch load soi --years 2021 -arch stats -``` - -### 3. Validate Fixture Facts +### 2. Validate Fixture Facts The standalone Arch fact harness validates JSONL aggregate facts and emits a JSON report with fact counts, QA counts, warnings, and validation errors: @@ -429,7 +419,7 @@ same `--axiom-cli` and `--axiom-root` flags, plus `--require-axiom-validation` when skipped concept checks should fail agent acceptance. -### 4. Run Arch Explorer +### 3. Run Arch Explorer Arch Explorer is a Next/Tailwind app that reads the fixture fact JSONL and source-cell JSONL, then shows aggregate facts, source-cell lineage, and @@ -451,20 +441,7 @@ By default, the workbench reads the current local suite outputs at ARCH_EXPLORER_DATA_DIRS=/tmp/arch-build-a,/tmp/arch-build-b npm run dev -- --port 3090 ``` -### 5. Query Target Inputs in Python - -```python -from arch.targets import DataSource, Target, TargetType -from calibration.targets import get_targets - -target_inputs = get_targets( - jurisdiction="us", - year=2021, - sources=["irs-soi"], -) -``` - -### 6. Query Microdata +### 4. Query Microdata ```python from arch.microdata import query_cps_asec @@ -472,16 +449,19 @@ from arch.microdata import query_cps_asec persons = query_cps_asec(year=2024, table_type="person", limit=10_000) ``` -## Target Input Schema - -Target inputs use a three-table schema: - -- **strata**: Population subgroups, such as California filers with AGI between - $50k and $75k. -- **stratum_constraints**: Rules defining each stratum. -- **targets**: Source-published aggregate values linked to strata. - -These are inputs to Microplex target composition. Microplex owns the active, +## Retired legacy target tables + +This branch used to load hand-written ETLs (`db/etl_*.py`) into a local SQLite +`targets` / `strata` / `stratum_constraints` schema (`arch load`, `arch stats`, +`arch query`), read a Supabase copy through `arch.targets.query_targets`, and +calibrate microdata to those rows in `calibration/` and `micro/us/pipeline.py`. +Nothing outside that stack read the tables, and the audit in +PolicyEngine/chronicle#348 (`docs/us-legacy-etl-cleanup.md` there) found that +several of the US loaders wrote invented or untraceable values, so the whole +stack was removed. Main dropped `calibration/` and `micro/` in #68 and the UK +loaders in #181 and #186, and is retiring the US loaders and target tables in +#348 and its follow-on. Source-backed facts reach consumers through source +packages and the bundle (`arch build-bundle`). Microplex owns the active, reconciled, aged target sets used for calibration. ## Arch Facts And Microplex Targets @@ -498,8 +478,8 @@ adjusted or projected series. ```python from arch.facts import SourceFact -from arch.targets import DataSource, Jurisdiction, TargetType -from arch.normalization import as_target, convert_units +from arch.normalization import convert_units +from db.schema import DataSource, Jurisdiction fact = SourceFact( name="snap_households", @@ -510,11 +490,7 @@ fact = SourceFact( jurisdiction=Jurisdiction.US, ) -target_input = as_target( - convert_units(fact, 1000, "count"), - target_type=TargetType.COUNT, - stratum_name="US SNAP Households", -) +normalized_fact = convert_units(fact, 1000, "count") ``` ## Current Coverage diff --git a/arch/cli.py b/arch/cli.py index bbb4ee3f..d78f4900 100644 --- a/arch/cli.py +++ b/arch/cli.py @@ -8,7 +8,7 @@ def main() -> None: - """Dispatch Arch-owned harness commands or the legacy DB CLI.""" + """Dispatch Arch-owned harness commands or the source-artifact DB CLI.""" if sys.argv[1:2] in ( ["bootstrap-r2"], ["build-bundle"], diff --git a/arch/jurisdictions/uk/__init__.py b/arch/jurisdictions/uk/__init__.py index 340af992..e06fd51e 100644 --- a/arch/jurisdictions/uk/__init__.py +++ b/arch/jurisdictions/uk/__init__.py @@ -1,9 +1,7 @@ -"""United Kingdom Arch target loaders.""" +"""United Kingdom Arch source builders. -from arch.targets.loaders import ( - load_hmrc_targets, - load_obr_targets, - load_ons_targets, -) +The legacy HMRC, OBR and ONS DB target loaders were retired, so this module +exports no loaders. +""" -__all__ = ["load_hmrc_targets", "load_obr_targets", "load_ons_targets"] +__all__: list[str] = [] diff --git a/arch/jurisdictions/us/__init__.py b/arch/jurisdictions/us/__init__.py index 19596927..19a774b6 100644 --- a/arch/jurisdictions/us/__init__.py +++ b/arch/jurisdictions/us/__init__.py @@ -1,35 +1,7 @@ -"""United States Arch target loaders.""" +"""United States Arch source builders. -from arch.targets.loaders import ( - load_aca_enrollment_targets, - load_bls_targets, - load_cbo_targets, - load_census_targets, - load_cps_targets, - load_medicaid_targets, - load_snap_targets, - load_soi_targets, - load_soi_credits_targets, - load_soi_deductions_targets, - load_soi_income_sources_targets, - load_soi_state_targets, - load_ssa_targets, - load_ssi_targets, -) +US measures are source packages. The legacy DB target loaders were retired, so +this module exports no loaders. +""" -__all__ = [ - "load_aca_enrollment_targets", - "load_bls_targets", - "load_cbo_targets", - "load_census_targets", - "load_cps_targets", - "load_medicaid_targets", - "load_snap_targets", - "load_soi_targets", - "load_soi_credits_targets", - "load_soi_deductions_targets", - "load_soi_income_sources_targets", - "load_soi_state_targets", - "load_ssa_targets", - "load_ssi_targets", -] +__all__: list[str] = [] diff --git a/arch/jurisdictions/us/soi.py b/arch/jurisdictions/us/soi.py index 8eaf9ae1..c6af36cc 100644 --- a/arch/jurisdictions/us/soi.py +++ b/arch/jurisdictions/us/soi.py @@ -25,7 +25,6 @@ SourceRecordSetSpec, SourceRecordSpec, ) -from db.etl_soi import AGI_BRACKETS, TABLE_1_1_AGI_LABEL_TO_BRACKET SOI_TABLE_1_1_SOURCE_NAME = "irs_soi" SOI_TABLE_1_1_SOURCE_TABLE = "Publication 1304 Table 1.1" @@ -50,6 +49,51 @@ "tax-year source record." ) +# AGI bracket definitions (lower, upper) in dollars. +AGI_BRACKETS = { + "under_1": (float("-inf"), 1), + "1_to_5k": (1, 5_000), + "5k_to_10k": (5_000, 10_000), + "10k_to_15k": (10_000, 15_000), + "15k_to_20k": (15_000, 20_000), + "20k_to_25k": (20_000, 25_000), + "25k_to_30k": (25_000, 30_000), + "30k_to_40k": (30_000, 40_000), + "40k_to_50k": (40_000, 50_000), + "50k_to_75k": (50_000, 75_000), + "75k_to_100k": (75_000, 100_000), + "100k_to_200k": (100_000, 200_000), + "200k_to_500k": (200_000, 500_000), + "500k_to_1m": (500_000, 1_000_000), + "1m_to_1_5m": (1_000_000, 1_500_000), + "1_5m_to_2m": (1_500_000, 2_000_000), + "2m_to_5m": (2_000_000, 5_000_000), + "5m_to_10m": (5_000_000, 10_000_000), + "10m_plus": (10_000_000, float("inf")), +} + +TABLE_1_1_AGI_LABEL_TO_BRACKET = { + "No adjusted gross income": "under_1", + "$1 under $5,000": "1_to_5k", + "$5,000 under $10,000": "5k_to_10k", + "$10,000 under $15,000": "10k_to_15k", + "$15,000 under $20,000": "15k_to_20k", + "$20,000 under $25,000": "20k_to_25k", + "$25,000 under $30,000": "25k_to_30k", + "$30,000 under $40,000": "30k_to_40k", + "$40,000 under $50,000": "40k_to_50k", + "$50,000 under $75,000": "50k_to_75k", + "$75,000 under $100,000": "75k_to_100k", + "$100,000 under $200,000": "100k_to_200k", + "$200,000 under $500,000": "200k_to_500k", + "$500,000 under $1,000,000": "500k_to_1m", + "$1,000,000 under $1,500,000": "1m_to_1_5m", + "$1,500,000 under $2,000,000": "1_5m_to_2m", + "$2,000,000 under $5,000,000": "2m_to_5m", + "$5,000,000 under $10,000,000": "5m_to_10m", + "$10,000,000 or more": "10m_plus", +} + def build_soi_table_1_1_facts(year: int = 2023) -> list[AggregateFact]: """Build source-lineaged Arch facts from SOI Table 1.1.""" diff --git a/arch/normalization/__init__.py b/arch/normalization/__init__.py index 6c46c4b7..7d09c7cb 100644 --- a/arch/normalization/__init__.py +++ b/arch/normalization/__init__.py @@ -7,24 +7,18 @@ """ from arch.facts import DerivationStep, SourceFact -from .targets import ( - TargetBlueprint, +from .transforms import ( apply_share, - as_target, convert_units, format_derivation, scale_value, - target_kwargs, ) __all__ = [ "DerivationStep", "SourceFact", - "TargetBlueprint", "apply_share", - "as_target", "convert_units", "format_derivation", "scale_value", - "target_kwargs", ] diff --git a/arch/normalization/targets.py b/arch/normalization/targets.py deleted file mode 100644 index 2b1dbe59..00000000 --- a/arch/normalization/targets.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Helpers for source facts and target inputs. - -These helpers are for low-assumption structural normalization: unit/scale -standardization and same-source arithmetic with explicit lineage. Modeling -choices such as inflation, aging, target reconciliation, or simulator-specific -allocation belong in Microplex calibration recipes. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -from arch.facts import DerivationStep, SourceFact -from arch.targets import DataSource, GeographicLevel, Jurisdiction, TargetType - -Constraint = tuple[str, str, str] - - -@dataclass(frozen=True) -class TargetBlueprint: - """A target input ready for a loader to materialize.""" - - variable: str - value: float - period: int - target_type: TargetType - source: DataSource - jurisdiction: Jurisdiction - stratum_name: str - constraints: tuple[Constraint, ...] = () - geographic_level: GeographicLevel | None = None - source_table: str | None = None - source_url: str | None = None - notes: str | None = None - is_preliminary: bool = False - margin_of_error: float | None = None - derivation: tuple[DerivationStep, ...] = () - - -def scale_value( - fact: SourceFact, - factor: float, - *, - name: str | None = None, - unit: str | None = None, - operation: str = "scale_value", - note: str | None = None, -) -> SourceFact: - """Scale a published fact value while recording normalization lineage.""" - parameters: dict[str, Any] = { - "factor": factor, - "input_unit": fact.unit, - "output_unit": unit if unit is not None else fact.unit, - } - if note: - parameters["note"] = note - - return fact.with_step( - DerivationStep(operation=operation, parameters=parameters), - name=name, - value=fact.value * factor, - unit=unit, - ) - - -def convert_units( - fact: SourceFact, - factor: float, - to_unit: str, - *, - name: str | None = None, - note: str | None = None, -) -> SourceFact: - """Normalize a published unit or scale with a multiplicative factor.""" - return scale_value( - fact, - factor, - name=name, - unit=to_unit, - operation="convert_units", - note=note, - ) - - -def apply_share( - total_fact: SourceFact, - share: float, - *, - name: str, - share_name: str | None = None, - unit: str | None = None, -) -> SourceFact: - """Materialize a same-source total/share relationship. - - Use this only when the source publishes both the total and the share. Do - not use it for model-driven allocation or target reconciliation. - """ - return total_fact.with_step( - DerivationStep( - operation="apply_share", - parameters={ - "share": share, - "share_name": share_name, - "total_fact": total_fact.name, - }, - ), - name=name, - value=total_fact.value * share, - unit=unit, - ) - - -def as_target( - fact: SourceFact, - *, - variable: str | None = None, - target_type: TargetType = TargetType.COUNT, - source: DataSource | str | None = None, - jurisdiction: Jurisdiction | str | None = None, - stratum_name: str, - constraints: tuple[Constraint, ...] = (), - geographic_level: GeographicLevel | None = None, - notes: str | None = None, -) -> TargetBlueprint: - """Convert a source fact into a target input blueprint.""" - target_source = _coerce_data_source(source if source is not None else fact.source) - target_jurisdiction = _coerce_jurisdiction( - jurisdiction if jurisdiction is not None else fact.jurisdiction - ) - - derivation_note = format_derivation(fact.derivation) - if notes and derivation_note: - notes = f"{notes} Derived via: {derivation_note}." - elif derivation_note: - notes = f"Derived via: {derivation_note}." - - return TargetBlueprint( - variable=variable if variable is not None else fact.name, - value=fact.value, - period=fact.period, - target_type=target_type, - source=target_source, - jurisdiction=target_jurisdiction, - stratum_name=stratum_name, - constraints=constraints, - geographic_level=geographic_level, - source_table=fact.source_table, - source_url=fact.source_url, - notes=notes, - is_preliminary=fact.is_preliminary, - margin_of_error=fact.margin_of_error, - derivation=fact.derivation, - ) - - -def target_kwargs(blueprint: TargetBlueprint, stratum_id: int) -> dict[str, Any]: - """Return keyword arguments for constructing a target DB row.""" - return { - "stratum_id": stratum_id, - "variable": blueprint.variable, - "period": blueprint.period, - "value": blueprint.value, - "target_type": blueprint.target_type, - "geographic_level": blueprint.geographic_level, - "source": blueprint.source, - "source_table": blueprint.source_table, - "source_url": blueprint.source_url, - "notes": blueprint.notes, - "is_preliminary": blueprint.is_preliminary, - "margin_of_error": blueprint.margin_of_error, - } - - -def format_derivation(steps: tuple[DerivationStep, ...]) -> str: - """Format derivation steps for target notes or diagnostics.""" - parts = [] - for step in steps: - params = ", ".join( - f"{key}={value}" for key, value in step.parameters.items() - ) - parts.append(f"{step.operation}({params})" if params else step.operation) - return "; ".join(parts) - - -def _coerce_data_source(source: DataSource | str | None) -> DataSource: - if source is None: - raise ValueError("Target source must be supplied by the fact or caller") - if isinstance(source, DataSource): - return source - return DataSource(source) - - -def _coerce_jurisdiction(jurisdiction: Jurisdiction | str | None) -> Jurisdiction: - if jurisdiction is None: - raise ValueError("Target jurisdiction must be supplied by the fact or caller") - if isinstance(jurisdiction, Jurisdiction): - return jurisdiction - return Jurisdiction(jurisdiction) diff --git a/arch/normalization/transforms.py b/arch/normalization/transforms.py new file mode 100644 index 00000000..d258a944 --- /dev/null +++ b/arch/normalization/transforms.py @@ -0,0 +1,97 @@ +"""Lineage-recording transforms for source facts. + +These helpers are for low-assumption structural normalization: unit/scale +standardization and same-source arithmetic with explicit lineage. Modeling +choices such as inflation, aging, target reconciliation, or simulator-specific +allocation belong in Microplex calibration recipes. +""" + +from __future__ import annotations + +from typing import Any + +from arch.facts import DerivationStep, SourceFact + + +def scale_value( + fact: SourceFact, + factor: float, + *, + name: str | None = None, + unit: str | None = None, + operation: str = "scale_value", + note: str | None = None, +) -> SourceFact: + """Scale a published fact value while recording normalization lineage.""" + parameters: dict[str, Any] = { + "factor": factor, + "input_unit": fact.unit, + "output_unit": unit if unit is not None else fact.unit, + } + if note: + parameters["note"] = note + + return fact.with_step( + DerivationStep(operation=operation, parameters=parameters), + name=name, + value=fact.value * factor, + unit=unit, + ) + + +def convert_units( + fact: SourceFact, + factor: float, + to_unit: str, + *, + name: str | None = None, + note: str | None = None, +) -> SourceFact: + """Normalize a published unit or scale with a multiplicative factor.""" + return scale_value( + fact, + factor, + name=name, + unit=to_unit, + operation="convert_units", + note=note, + ) + + +def apply_share( + total_fact: SourceFact, + share: float, + *, + name: str, + share_name: str | None = None, + unit: str | None = None, +) -> SourceFact: + """Materialize a same-source total/share relationship. + + Use this only when the source publishes both the total and the share. Do + not use it for model-driven allocation or target reconciliation. + """ + return total_fact.with_step( + DerivationStep( + operation="apply_share", + parameters={ + "share": share, + "share_name": share_name, + "total_fact": total_fact.name, + }, + ), + name=name, + value=total_fact.value * share, + unit=unit, + ) + + +def format_derivation(steps: tuple[DerivationStep, ...]) -> str: + """Format derivation steps for target notes or diagnostics.""" + parts = [] + for step in steps: + params = ", ".join( + f"{key}={value}" for key, value in step.parameters.items() + ) + parts.append(f"{step.operation}({params})" if params else step.operation) + return "; ".join(parts) diff --git a/arch/targets/__init__.py b/arch/targets/__init__.py deleted file mode 100644 index 2a7d8632..00000000 --- a/arch/targets/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Calibration target schema and query helpers.""" - -from db.schema import ( - DEFAULT_DB_PATH, - DataSource, - GeographicLevel, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - get_engine, - get_session, - init_db, -) -from db.supabase_client import insert_targets_batch, query_strata, query_targets -from calibration.targets import TargetSpec, get_targets - -__all__ = [ - "DEFAULT_DB_PATH", - "DataSource", - "GeographicLevel", - "Jurisdiction", - "Stratum", - "StratumConstraint", - "Target", - "TargetSpec", - "TargetType", - "get_engine", - "get_targets", - "get_session", - "init_db", - "insert_targets_batch", - "query_strata", - "query_targets", -] diff --git a/arch/targets/client.py b/arch/targets/client.py deleted file mode 100644 index aa7cf5d2..00000000 --- a/arch/targets/client.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Supabase helpers for Arch target inputs.""" - -from db.supabase_client import insert_targets_batch, query_strata, query_targets - -__all__ = ["insert_targets_batch", "query_strata", "query_targets"] diff --git a/arch/targets/loaders.py b/arch/targets/loaders.py deleted file mode 100644 index 974f2deb..00000000 --- a/arch/targets/loaders.py +++ /dev/null @@ -1,39 +0,0 @@ -"""ETL loaders that populate Arch target inputs.""" - -from db.etl_aca_enrollment import load_aca_enrollment_targets -from db.etl_bls import load_bls_targets -from db.etl_cbo import load_cbo_targets -from db.etl_census import load_census_targets -from db.etl_cps import load_cps_targets -from db.etl_hmrc import load_hmrc_targets -from db.etl_medicaid import load_medicaid_targets -from db.etl_obr import load_obr_targets -from db.etl_ons import load_ons_targets -from db.etl_snap import load_snap_targets -from db.etl_soi import load_soi_targets -from db.etl_soi_credits import load_soi_credits_targets -from db.etl_soi_deductions import load_soi_deductions_targets -from db.etl_soi_income_sources import load_soi_income_sources_targets -from db.etl_soi_state import load_soi_state_targets -from db.etl_ssa import load_ssa_targets -from db.etl_ssi import load_ssi_targets - -__all__ = [ - "load_aca_enrollment_targets", - "load_bls_targets", - "load_cbo_targets", - "load_census_targets", - "load_cps_targets", - "load_hmrc_targets", - "load_medicaid_targets", - "load_obr_targets", - "load_ons_targets", - "load_snap_targets", - "load_soi_targets", - "load_soi_credits_targets", - "load_soi_deductions_targets", - "load_soi_income_sources_targets", - "load_soi_state_targets", - "load_ssa_targets", - "load_ssi_targets", -] diff --git a/calibration/__init__.py b/calibration/__init__.py deleted file mode 100644 index 635a20da..00000000 --- a/calibration/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Calibration pipeline for microdata reweighting. - -Connects administrative targets to microdata, producing reweighted samples -that match official statistics. - -Supports both flat and hierarchical microdata: -- Flat: Single DataFrame with weights per record -- Hierarchical: Households with linked persons, weights at household level -""" - -from .loader import load_microdata -from .targets import get_targets, TargetSpec -from .constraints import ( - build_constraint_matrix, - build_hierarchical_constraint_matrix, - Constraint, -) -from .variables import get_entity, infer_target_level -from .methods import EntropyCalibrator - -__all__ = [ - # Loader - "load_microdata", - # Targets - "get_targets", - "TargetSpec", - # Constraints - "build_constraint_matrix", - "build_hierarchical_constraint_matrix", - "Constraint", - # Variables - "get_entity", - "infer_target_level", - # Methods - "EntropyCalibrator", -] diff --git a/calibration/constraints.py b/calibration/constraints.py deleted file mode 100644 index eadd922d..00000000 --- a/calibration/constraints.py +++ /dev/null @@ -1,293 +0,0 @@ -""" -Constraint matrix building for calibration. - -Maps targets to microdata aggregations, building indicator vectors -for each constraint. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Optional - -import numpy as np -import pandas as pd - -from db.schema import TargetType -from .targets import TargetSpec - - -@dataclass -class Constraint: - """ - A single calibration constraint. - - Represents the equation: sum(weights * indicator) = target_value - - Attributes: - indicator: Vector of indicator values (length = microdata rows) - target_value: Target aggregate to match - variable: Arch target input variable ID - target_type: COUNT or AMOUNT - tolerance: Allowed deviation from target (fraction) - stratum_name: Human-readable description - """ - - indicator: np.ndarray - target_value: float - variable: str - target_type: TargetType - tolerance: float = 0.01 - stratum_name: Optional[str] = None - - -def apply_stratum_constraints( - microdata: pd.DataFrame, - constraints: list[tuple[str, str, str]], -) -> pd.Series: - """ - Return boolean mask for records matching stratum. - - Args: - microdata: DataFrame with microdata - constraints: List of (variable, operator, value) tuples - - Returns: - Boolean Series indicating which rows match all constraints - """ - mask = pd.Series(True, index=microdata.index) - - for variable, operator, value in constraints: - if variable not in microdata.columns: - # Skip constraints for variables not in microdata - continue - - col = microdata[variable] - - # Parse value based on column dtype - if pd.api.types.is_numeric_dtype(col): - parsed_value = float(value) - else: - parsed_value = value - - if operator == "==": - mask &= col == parsed_value - elif operator == "!=": - mask &= col != parsed_value - elif operator == ">": - mask &= col > parsed_value - elif operator == ">=": - mask &= col >= parsed_value - elif operator == "<": - mask &= col < parsed_value - elif operator == "<=": - mask &= col <= parsed_value - elif operator == "in": - # Value should be comma-separated list - values = [v.strip() for v in value.split(",")] - if pd.api.types.is_numeric_dtype(col): - values = [float(v) for v in values] - mask &= col.isin(values) - else: - raise ValueError(f"Unknown operator: {operator}") - - return mask - - -def build_constraint_matrix( - microdata: pd.DataFrame, - targets: list[TargetSpec], - tolerance: float = 0.01, -) -> list[Constraint]: - """ - Build constraint matrix from targets and microdata. - - For each target, creates an indicator vector: - - COUNT type: indicator is 1 for matching rows, 0 otherwise - - AMOUNT type: indicator is variable value for matching rows, 0 otherwise - - Args: - microdata: DataFrame with microdata records - targets: List of TargetSpec objects from get_targets() - tolerance: Default allowed deviation (can be overridden per target) - - Returns: - List of Constraint objects ready for calibration - """ - constraints = [] - - for target in targets: - # Build stratum mask - mask = apply_stratum_constraints(microdata, target.constraints) - - # Build indicator vector based on target type - if target.target_type == TargetType.COUNT: - # For counts, indicator is just the mask - indicator = mask.astype(float).values - elif target.target_type == TargetType.AMOUNT: - # For amounts, indicator is variable * mask - if target.variable in microdata.columns: - indicator = (microdata[target.variable] * mask).values - else: - # If variable not in microdata, use zeros - indicator = np.zeros(len(microdata)) - else: - # RATE type - not commonly used for calibration - indicator = mask.astype(float).values - - constraint = Constraint( - indicator=indicator, - target_value=target.value, - variable=target.variable, - target_type=target.target_type, - tolerance=target.tolerance if target.tolerance else tolerance, - stratum_name=target.stratum_name, - ) - constraints.append(constraint) - - return constraints - - -def build_hierarchical_constraint_matrix( - hh_df: pd.DataFrame, - person_df: pd.DataFrame, - targets: list[TargetSpec], - tolerance: float = 0.01, - hh_id_col: str = "household_id", - tax_unit_df: Optional[pd.DataFrame] = None, -) -> list[Constraint]: - """ - Build constraint matrix aggregating person/tax-unit targets to household level. - - For hierarchical microdata where weights are at the household level, - this function handles person-level and tax-unit-level targets by - aggregating (count or sum) to the household level. - - The key insight: since all calibration targets are sums, and all persons - in a household share the household weight, we can pre-aggregate: - sum(person_weight * is_age_6_10) = sum(hh_weight * n_age_6_10_in_hh) - - Args: - hh_df: Household-level DataFrame (one row per household) - person_df: Person-level DataFrame with hh_id_col linking to households - targets: List of TargetSpec objects from get_targets() - tolerance: Default allowed deviation - hh_id_col: Column name for household ID in both DataFrames - tax_unit_df: Optional tax-unit-level DataFrame - - Returns: - List of Constraint objects with indicators at household level - """ - from .variables import get_entity, infer_target_level - - constraints = [] - - for target in targets: - # Determine what level this target operates at - # 1. Try to infer from the target variable reference - # 2. Fall back to inferring from constraint variables - try: - if ":" in target.variable and "#" in target.variable: - level = get_entity(target.variable) - else: - level = infer_target_level(target.constraints) - except Exception: - level = infer_target_level(target.constraints) - - # Build indicator based on level - if level == "household": - # Direct household-level constraint - indicator = _build_household_indicator( - hh_df, target, hh_id_col - ) - elif level == "tax_unit" and tax_unit_df is not None: - # Aggregate from tax unit to household - indicator = _build_aggregated_indicator( - hh_df, tax_unit_df, target, hh_id_col - ) - else: - # Aggregate from person to household (default) - indicator = _build_aggregated_indicator( - hh_df, person_df, target, hh_id_col - ) - - constraint = Constraint( - indicator=indicator, - target_value=target.value, - variable=target.variable, - target_type=target.target_type, - tolerance=target.tolerance if target.tolerance else tolerance, - stratum_name=target.stratum_name, - ) - constraints.append(constraint) - - return constraints - - -def _build_household_indicator( - hh_df: pd.DataFrame, - target: TargetSpec, - hh_id_col: str, -) -> np.ndarray: - """Build indicator for household-level targets.""" - mask = apply_stratum_constraints(hh_df, target.constraints) - - if target.target_type == TargetType.COUNT: - return mask.astype(float).values - elif target.target_type == TargetType.AMOUNT: - # Extract variable name from full reference if needed - var_name = target.variable - if "#" in var_name: - var_name = var_name.split("#")[-1] - - if var_name in hh_df.columns: - return (hh_df[var_name] * mask).values - else: - return np.zeros(len(hh_df)) - else: - return mask.astype(float).values - - -def _build_aggregated_indicator( - hh_df: pd.DataFrame, - source_df: pd.DataFrame, - target: TargetSpec, - hh_id_col: str, -) -> np.ndarray: - """ - Build indicator by aggregating from person/tax_unit to household. - - For COUNT targets: count matching records per household - For AMOUNT targets: sum variable value for matching records per household - """ - # Apply constraints to source (person or tax_unit) data - mask = apply_stratum_constraints(source_df, target.constraints) - filtered = source_df[mask] - - if len(filtered) == 0: - # No matching records - return np.zeros(len(hh_df)) - - if target.target_type == TargetType.COUNT: - # Count matching records per household - agg = filtered.groupby(hh_id_col).size() - elif target.target_type == TargetType.AMOUNT: - # Sum variable for matching records per household - var_name = target.variable - if "#" in var_name: - var_name = var_name.split("#")[-1] - - if var_name in filtered.columns: - agg = filtered.groupby(hh_id_col)[var_name].sum() - else: - # Variable not found, return zeros - return np.zeros(len(hh_df)) - else: - # RATE type - treat as count - agg = filtered.groupby(hh_id_col).size() - - # Map aggregated values back to household DataFrame order - # Households with no matching persons get 0 - indicator = hh_df[hh_id_col].map(agg).fillna(0).values - - return indicator diff --git a/calibration/loader.py b/calibration/loader.py deleted file mode 100644 index e778c512..00000000 --- a/calibration/loader.py +++ /dev/null @@ -1,312 +0,0 @@ -""" -Microdata loading for calibration. - -Loads raw microdata (CPS, FRS) with original survey weights. -""" - -from datetime import datetime -from pathlib import Path -from typing import Optional - -import numpy as np -import pandas as pd - -# Valid data sources -VALID_SOURCES = {"synthetic", "cps", "frs"} - -# US State FIPS codes (valid ones) -US_STATE_FIPS = [ - 1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 15, 16, 17, 18, 19, 20, - 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, - 36, 37, 38, 39, 40, 41, 42, 44, 45, 46, 47, 48, 49, 50, 51, - 53, 54, 55, 56 -] - - -def _validate_year(year: int, source: str) -> None: - """Validate year is reasonable.""" - current_year = datetime.now().year - - # Don't allow future years - if year > current_year: - raise ValueError( - f"Year {year} is in the future. Maximum year is {current_year}." - ) - - # Don't allow ancient years - if year < 1900: - raise ValueError( - f"Year {year} is too far in the past. Minimum year is 1900." - ) - - -def _generate_synthetic_cps( - year: int, - n_samples: int = 10000, - seed: Optional[int] = None, -) -> pd.DataFrame: - """ - Generate synthetic microdata matching CPS structure. - - Creates realistic demographic and income distributions for testing - calibration algorithms. - - Args: - year: Year for synthetic data - n_samples: Number of sample records to generate - seed: Random seed for reproducibility - - Returns: - DataFrame with synthetic microdata and weights - """ - if seed is not None: - np.random.seed(seed) - - # Age distribution (roughly matching US population) - # More children and working-age adults, fewer elderly - age_groups = [ - (0, 18, 0.22), # Children - (18, 35, 0.22), # Young adults - (35, 55, 0.26), # Middle-aged - (55, 70, 0.18), # Near retirement - (70, 100, 0.12), # Elderly - ] - - ages = [] - for min_age, max_age, proportion in age_groups: - group_size = int(n_samples * proportion) - group_ages = np.random.randint(min_age, max_age + 1, group_size) - ages.extend(group_ages) - - # Adjust to exact n_samples - while len(ages) < n_samples: - ages.append(np.random.randint(0, 100)) - ages = np.array(ages[:n_samples]) - - # Employment status: 0=not in labor force, 1=employed, 2=unemployed - # Children and elderly mostly not in labor force - employment_status = np.zeros(n_samples, dtype=int) - - for i in range(n_samples): - if ages[i] < 16: - employment_status[i] = 0 # Not in labor force - elif ages[i] >= 65: - # 20% still working, 75% retired, 5% unemployed - employment_status[i] = np.random.choice([0, 1, 2], p=[0.75, 0.20, 0.05]) - else: - # Working age: 62% employed, 34% not in labor force, 4% unemployed - employment_status[i] = np.random.choice([0, 1, 2], p=[0.34, 0.62, 0.04]) - - # Income (based on employment status and age) - income = np.zeros(n_samples) - for i in range(n_samples): - if employment_status[i] == 1: # Employed - # Log-normal income distribution - base_income = np.random.lognormal(10.5, 0.8) - # Adjust by age (peak around 45-55) - age_factor = 1.0 + 0.3 * np.sin(np.pi * (ages[i] - 18) / 47) - income[i] = base_income * max(0.5, age_factor) - elif employment_status[i] == 0: # Not in labor force - if ages[i] >= 65: - # Retirement income (Social Security, pensions) - income[i] = np.random.lognormal(9.5, 0.7) - elif ages[i] < 18: - # Children have no income - income[i] = 0 - else: - # Adults not in labor force may have some income - income[i] = np.random.lognormal(8.0, 1.0) * np.random.choice([0, 1], p=[0.7, 0.3]) - else: # Unemployed - # Unemployment benefits or no income - income[i] = np.random.lognormal(8.5, 0.5) * np.random.choice([0, 1], p=[0.4, 0.6]) - - # has_children: probability depends on age - has_children = np.zeros(n_samples, dtype=int) - for i in range(n_samples): - if 25 <= ages[i] <= 55: - has_children[i] = np.random.choice([0, 1], p=[0.4, 0.6]) - elif 18 <= ages[i] < 25 or 55 < ages[i] <= 65: - has_children[i] = np.random.choice([0, 1], p=[0.7, 0.3]) - else: - has_children[i] = 0 - - # State FIPS (weighted by population) - # Approximate state population weights (simplified) - state_weights = { - 6: 0.12, # California - 48: 0.09, # Texas - 12: 0.07, # Florida - 36: 0.06, # New York - 17: 0.04, # Illinois - 42: 0.04, # Pennsylvania - 39: 0.04, # Ohio - 13: 0.03, # Georgia - 37: 0.03, # North Carolina - 26: 0.03, # Michigan - } - other_weight = 1.0 - sum(state_weights.values()) - other_states = [s for s in US_STATE_FIPS if s not in state_weights] - - states = [] - for _ in range(n_samples): - if np.random.random() < other_weight: - states.append(np.random.choice(other_states)) - else: - # Pick from weighted states - state_probs = np.array(list(state_weights.values())) - state_probs = state_probs / state_probs.sum() - states.append( - np.random.choice(list(state_weights.keys()), p=state_probs) - ) - - state_fips = np.array(states) - - # Survey weights (roughly matching CPS scale) - # Total US population ~330M, sample of 10k means avg weight ~33k - base_weight = 330_000_000 / n_samples - # Add some variance - weights = np.random.gamma(shape=4, scale=base_weight / 4, size=n_samples) - - return pd.DataFrame({ - "weight": weights, - "age": ages, - "income": income, - "employment_status": employment_status, - "has_children": has_children, - "state_fips": state_fips, - }) - - -def _load_cps_from_file(year: int, auto_download: bool = True) -> Optional[pd.DataFrame]: - """ - Attempt to load CPS data from parquet files, downloading if needed. - - Looks for data in micro/us/ directory. If not found and auto_download - is True, downloads from Census Bureau. - - Args: - year: Year of data to load - auto_download: If True, download from Census if file not found - - Returns: - DataFrame if file exists or downloaded, None otherwise - """ - # Check standard locations - data_dir = Path(__file__).parent.parent / "micro" / "us" - - potential_paths = [ - data_dir / f"cps_{year}.parquet", - data_dir / f"cps_asec_{year}.parquet", - data_dir / f"asec_{year}.parquet", - ] - - for path in potential_paths: - if path.exists(): - df = pd.read_parquet(path) - # Ensure weight column exists - if "weight" not in df.columns: - # Try common weight column names - weight_candidates = [ - "ASECWT", "asecwt", "WTFINL", "wtfinl", - "PWGTP", "pwgtp", "survey_weight" - ] - for col in weight_candidates: - if col in df.columns: - df["weight"] = df[col] - break - else: - raise ValueError( - f"Could not find weight column in {path}. " - f"Columns: {df.columns.tolist()}" - ) - return df - - # Try to download if not found - if auto_download: - try: - from micro.us.census.download_cps import download_and_process_cps, CPS_URL_BY_YEAR - - if year in CPS_URL_BY_YEAR: - print(f"CPS {year} not found locally, downloading from Census Bureau...") - output_path = data_dir / f"cps_{year}.parquet" - return download_and_process_cps(year, output_path) - except ImportError: - pass # download_cps module not available - except Exception as e: - import warnings - warnings.warn(f"Failed to download CPS {year}: {e}") - - return None - - -def load_microdata( - source: str, - year: int, - variables: Optional[list[str]] = None, - seed: Optional[int] = None, -) -> pd.DataFrame: - """ - Load microdata for calibration. - - Args: - source: Data source ("synthetic" for testing, "cps" for US, "frs" for UK) - year: Year of microdata - variables: Optional list of variables to load (weight always included) - seed: Random seed for synthetic data reproducibility - - Returns: - DataFrame with microdata and original weights in 'weight' column - - Raises: - ValueError: If source or year is invalid - """ - # Validate source - if source not in VALID_SOURCES: - raise ValueError( - f"Invalid source: {source}. Valid sources: {VALID_SOURCES}" - ) - - # Validate year - _validate_year(year, source) - - # Load data based on source - if source == "synthetic": - df = _generate_synthetic_cps(year, seed=seed) - - elif source == "cps": - # Try to load real CPS data - df = _load_cps_from_file(year) - if df is None: - # Fall back to synthetic with a warning - import warnings - warnings.warn( - f"CPS data for {year} not found, falling back to synthetic data. " - "To use real CPS data, add parquet files to micro/us/." - ) - df = _generate_synthetic_cps(year, seed=seed) - - elif source == "frs": - # UK Family Resources Survey - not yet implemented - import warnings - warnings.warn( - "FRS data loading not yet implemented, using synthetic data." - ) - df = _generate_synthetic_cps(year, seed=seed) - - else: - raise ValueError(f"Unknown source: {source}") - - # Filter to requested variables (weight always included) - if variables is not None: - # Ensure weight is always included - cols_to_keep = set(variables) | {"weight"} - available_cols = set(df.columns) - missing = cols_to_keep - available_cols - if missing: - raise ValueError( - f"Requested variables not in data: {missing}. " - f"Available: {available_cols}" - ) - df = df[list(cols_to_keep)] - - return df diff --git a/calibration/methods/__init__.py b/calibration/methods/__init__.py deleted file mode 100644 index fe3e08f9..00000000 --- a/calibration/methods/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -""" -Calibration methods. - -Provides different algorithms for computing calibrated weights: -- EntropyCalibrator: Minimize KL divergence from original weights -- (Future) RakingCalibrator: Iterative proportional fitting -- (Future) LinearCalibrator: Linear regression adjustment -""" - -from .entropy import EntropyCalibrator - -__all__ = ["EntropyCalibrator"] diff --git a/calibration/methods/entropy.py b/calibration/methods/entropy.py deleted file mode 100644 index 2294b529..00000000 --- a/calibration/methods/entropy.py +++ /dev/null @@ -1,139 +0,0 @@ -""" -Entropy-based calibration method. - -Minimizes Kullback-Leibler divergence from original weights -while satisfying calibration constraints. -""" - -from dataclasses import dataclass - -import numpy as np -from scipy.optimize import minimize, OptimizeResult - -from ..constraints import Constraint - - -@dataclass -class EntropyCalibrator: - """ - Entropy-based weight calibrator. - - Finds new weights that minimize KL divergence from original weights - while matching target constraints within tolerance. - - Attributes: - bounds: (min_ratio, max_ratio) for weight adjustments - max_iterations: Maximum solver iterations - convergence_tol: Convergence tolerance for solver - """ - - bounds: tuple[float, float] = (0.1, 10.0) - max_iterations: int = 100 - convergence_tol: float = 1e-6 - - def calibrate( - self, - original_weights: np.ndarray, - constraints: list[Constraint], - ) -> np.ndarray: - """ - Compute calibrated weights. - - Args: - original_weights: Original survey weights (n,) - constraints: List of Constraint objects - - Returns: - Calibrated weights (n,) - - Raises: - ValueError: If constraints are infeasible - RuntimeError: If optimization fails to converge - """ - n = len(original_weights) - - # Objective function: KL divergence - def kl_divergence(w: np.ndarray) -> float: - """Compute KL divergence: sum(w * log(w/w0)).""" - # Add small epsilon to avoid log(0) - w_safe = np.maximum(w, 1e-10) - w0_safe = np.maximum(original_weights, 1e-10) - return np.sum(w_safe * np.log(w_safe / w0_safe)) - - # Gradient of KL divergence - def kl_gradient(w: np.ndarray) -> np.ndarray: - """Gradient of KL divergence: log(w/w0) + 1.""" - w_safe = np.maximum(w, 1e-10) - w0_safe = np.maximum(original_weights, 1e-10) - return np.log(w_safe / w0_safe) + 1.0 - - # Build constraint functions for scipy - scipy_constraints = [] - for constraint in constraints: - # Equality constraint: sum(w * indicator) - target = 0 - def constraint_fn(w: np.ndarray, ind=constraint.indicator, - target=constraint.target_value) -> float: - return np.dot(w, ind) - target - - # Jacobian of constraint: indicator vector - def constraint_jac(w: np.ndarray, ind=constraint.indicator) -> np.ndarray: - return ind - - scipy_constraints.append({ - 'type': 'eq', - 'fun': constraint_fn, - 'jac': constraint_jac, - }) - - # Bounds on individual weights - weight_bounds = [ - (original_weights[i] * self.bounds[0], - original_weights[i] * self.bounds[1]) - for i in range(n) - ] - - # Initial guess: original weights - x0 = original_weights.copy() - - # Solve optimization problem - result: OptimizeResult = minimize( - kl_divergence, - x0, - method='SLSQP', - jac=kl_gradient, - constraints=scipy_constraints, - bounds=weight_bounds, - options={ - 'maxiter': self.max_iterations, - 'ftol': self.convergence_tol, - 'disp': False, - } - ) - - # Check if optimization succeeded - if not result.success: - if 'Positive directional derivative for linesearch' in result.message: - # This often means constraints are infeasible - raise ValueError( - f"Optimization failed: constraints may be infeasible. " - f"Message: {result.message}" - ) - else: - raise RuntimeError( - f"Optimization failed to converge: {result.message}" - ) - - # Verify constraints are satisfied - calibrated_weights = result.x - for i, constraint in enumerate(constraints): - actual = np.dot(calibrated_weights, constraint.indicator) - error = abs(actual - constraint.target_value) / constraint.target_value - if error > constraint.tolerance: - raise RuntimeError( - f"Constraint {i} ({constraint.variable}) not satisfied: " - f"target={constraint.target_value:.2f}, " - f"actual={actual:.2f}, " - f"error={error:.4%} > tolerance={constraint.tolerance:.4%}" - ) - - return calibrated_weights diff --git a/calibration/targets.py b/calibration/targets.py deleted file mode 100644 index 4a2ca299..00000000 --- a/calibration/targets.py +++ /dev/null @@ -1,139 +0,0 @@ -""" -Target specification and querying. - -Provides TargetSpec dataclass and get_targets() function to query -the targets database for calibration constraints. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path - -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - DEFAULT_DB_PATH, - Stratum, - StratumConstraint, - Target, - TargetType, - get_engine, -) - - -@dataclass -class TargetSpec: - """ - Specification for a calibration target. - - Combines a target value with its stratum constraints, making it - ready for constraint matrix building. - - Attributes: - variable: Arch target input variable ID. - value: Target aggregate value - target_type: COUNT, AMOUNT, or RATE - constraints: List of (variable, operator, value) tuples defining stratum - source: Data source (e.g., IRS_SOI) - period: Year of the target - tolerance: Allowed deviation from target (optional) - stratum_name: Human-readable stratum name (optional) - """ - - variable: str - value: float - target_type: TargetType - constraints: list[tuple[str, str, str]] - source: DataSource - period: int - tolerance: float | None = None - stratum_name: str | None = None - - -def get_targets( - db_path: Path | None = None, - jurisdiction: str = "us", - year: int | None = None, - sources: list[str] | None = None, - variables: list[str] | None = None, -) -> list[TargetSpec]: - """ - Query targets database and return TargetSpec objects. - - Args: - db_path: Path to database (default: macro/targets.db) - jurisdiction: Jurisdiction filter (e.g., "us", "uk") - year: Filter by target year/period - sources: List of data source names to include (e.g., ["irs-soi"]) - variables: List of variable names to include - - Returns: - List of TargetSpec objects ready for constraint building - """ - if db_path is None: - db_path = DEFAULT_DB_PATH - - engine = get_engine(db_path) - - with Session(engine) as session: - # Build query for targets with their strata - query = select(Target, Stratum).join(Stratum) - - # Filter by year - if year is not None: - query = query.where(Target.period == year) - - # Filter by source - if sources is not None: - source_enums = [DataSource(s) for s in sources] - query = query.where(Target.source.in_(source_enums)) - - # Filter by variable - if variables is not None: - query = query.where(Target.variable.in_(variables)) - - # Filter by jurisdiction (match prefix for flexibility) - # e.g., "us" matches "us", "us-federal", "us-ca", etc. - if jurisdiction: - jurisdiction_lower = jurisdiction.lower() - # For now, simple filter - could be enhanced - query = query.where( - Stratum.jurisdiction.in_([ - j for j in [ - "us", "us-federal", "us-ca", "us-ny", "us-tx", - "uk", - ] - if j.startswith(jurisdiction_lower) - ]) - ) - - results = session.exec(query).all() - - # Convert to TargetSpec objects - target_specs = [] - for target, stratum in results: - # Fetch constraints for this stratum - constraint_query = select(StratumConstraint).where( - StratumConstraint.stratum_id == stratum.id - ) - stratum_constraints = session.exec(constraint_query).all() - - constraints = [ - (c.variable, c.operator, c.value) - for c in stratum_constraints - ] - - spec = TargetSpec( - variable=target.variable, - value=target.value, - target_type=target.target_type, - constraints=constraints, - source=target.source, - period=target.period, - stratum_name=stratum.name, - ) - target_specs.append(spec) - - return target_specs diff --git a/calibration/variables.py b/calibration/variables.py deleted file mode 100644 index 658f0a6c..00000000 --- a/calibration/variables.py +++ /dev/null @@ -1,237 +0,0 @@ -""" -Variable entity resolution using a RuleSpec engine package when available. - -Resolves fully qualified variable references to their entity type -(person, tax_unit, household, family) by parsing the source .rac file. - -Example: - >>> get_entity("us:statute/26/32#eitc") - "tax_unit" - >>> get_entity("us-ca:statute/ca/rtc/17041#ca_agi") - "tax_unit" -""" - -from functools import lru_cache -from pathlib import Path -from typing import Optional - -# Try to import a legacy RuleSpec engine API. -try: - from policyengine.dependency_resolver import PackageRegistry, ModuleResolver - from policyengine.dsl_parser import parse_file - POLICYENGINE_AVAILABLE = True -except ImportError: - POLICYENGINE_AVAILABLE = False - - -# Default workspace path (can be overridden) -DEFAULT_WORKSPACE = Path.home() / "PolicyEngine" - -# Fallback entity mapping for common variables when a RuleSpec engine is not available. -FALLBACK_ENTITIES = { - # Person-level - "age": "person", - "is_male": "person", - "employment_income": "person", - "is_snap_recipient": "person", - "is_medicaid_enrolled": "person", - "is_ssi_recipient": "person", - - # Tax unit level - "agi": "tax_unit", - "adjusted_gross_income": "tax_unit", - "filing_status": "tax_unit", - "eitc": "tax_unit", - "earned_income_credit": "tax_unit", - "ctc": "tax_unit", - "child_tax_credit": "tax_unit", - "tax_unit_count": "tax_unit", - - # Household level - "household_size": "household", - "tenure": "household", - "state_fips": "household", - "snap_allotment": "household", -} - - -class VariableNotFoundError(Exception): - """Raised when a variable cannot be found in the referenced file.""" - pass - - -class InvalidVariableRefError(Exception): - """Raised when a variable reference is malformed.""" - pass - - -def parse_variable_ref(variable_ref: str) -> tuple[str, str, str]: - """ - Parse a fully qualified variable reference. - - Args: - variable_ref: Reference like "us:statute/26/32#eitc" - - Returns: - Tuple of (package, path, variable_name) - - Raises: - InvalidVariableRefError: If reference is malformed - """ - if ":" not in variable_ref or "#" not in variable_ref: - raise InvalidVariableRefError( - f"Invalid variable reference: {variable_ref}. " - "Expected format: 'package:path#variable' (e.g., 'us:statute/26/32#eitc')" - ) - - package, rest = variable_ref.split(":", 1) - - if "#" not in rest: - raise InvalidVariableRefError( - f"Invalid variable reference: {variable_ref}. " - "Missing '#' separator between path and variable name." - ) - - path, var_name = rest.rsplit("#", 1) - - return package, path, var_name - - -@lru_cache(maxsize=256) -def get_entity( - variable_ref: str, - workspace: Optional[Path] = None, -) -> str: - """ - Get the entity type for a fully qualified variable reference. - - Parses the source .rac file to extract the variable's entity attribute. - Results are cached for performance. - - Args: - variable_ref: Fully qualified reference like "us:statute/26/32#eitc" - workspace: Optional workspace path (defaults to ~/PolicyEngine) - - Returns: - Entity type: "person", "tax_unit", "household", or "family" - - Raises: - VariableNotFoundError: If variable not found in the .rac file - InvalidVariableRefError: If reference is malformed - - Example: - >>> get_entity("us:statute/26/32#eitc") - "tax_unit" - """ - package, path, var_name = parse_variable_ref(variable_ref) - - # Try the RuleSpec engine first. - if POLICYENGINE_AVAILABLE: - try: - return _get_entity_from_rac(package, path, var_name, workspace) - except Exception: - # Fall through to fallback - pass - - # Fallback to static mapping - if var_name in FALLBACK_ENTITIES: - return FALLBACK_ENTITIES[var_name] - - # Default to person level (most common) - return "person" - - -def _get_entity_from_rac( - package: str, - path: str, - var_name: str, - workspace: Optional[Path] = None, -) -> str: - """ - Get entity by parsing the actual .rac file. - - Uses the legacy RuleSpec engine's PackageRegistry to resolve paths. - """ - if workspace is None: - workspace = DEFAULT_WORKSPACE - - # Map short package name to repo name - # us → policyengine-us, us-ca → policyengine-us-ca, uk → policyengine-uk - repo_name = f"policyengine-{package}" - - # Get package registry - registry = PackageRegistry.from_workspace(workspace) - - try: - root = registry.get_root(repo_name) - except Exception as e: - raise VariableNotFoundError( - f"Package '{repo_name}' not found in workspace {workspace}: {e}" - ) - - # Resolve file path - resolver = ModuleResolver(root) - try: - rac_path = resolver.resolve(path) - except Exception as e: - raise VariableNotFoundError( - f"Cannot resolve path '{path}' in package '{repo_name}': {e}" - ) - - # Parse the .rac file - module = parse_file(str(rac_path)) - - # Find the variable - for var in module.variables: - if var.name == var_name: - # Entity is stored as an enum or string - entity = var.entity - if hasattr(entity, 'value'): - return entity.value.lower() - return str(entity).lower() - - raise VariableNotFoundError( - f"Variable '{var_name}' not found in {rac_path}" - ) - - -def get_entity_for_constraint_var(var_name: str) -> str: - """ - Get entity for a constraint variable (not a full reference). - - Used when processing stratum constraints like (age, >=, 65). - Falls back to static mapping. - - Args: - var_name: Simple variable name like "age" or "state_fips" - - Returns: - Entity type (defaults to "person" if unknown) - """ - return FALLBACK_ENTITIES.get(var_name, "person") - - -def infer_target_level(constraints: list[tuple[str, str, str]]) -> str: - """ - Infer the most granular level needed based on constraint variables. - - If any constraint uses a person-level variable, the target needs - person-level filtering before aggregation to household. - - Args: - constraints: List of (variable, operator, value) tuples - - Returns: - Most granular level: "person", "tax_unit", or "household" - """ - levels = [] - for var, _, _ in constraints: - entity = get_entity_for_constraint_var(var) - levels.append(entity) - - # Return most granular level - if "person" in levels: - return "person" - if "tax_unit" in levels: - return "tax_unit" - return "household" diff --git a/data/targets/README.md b/data/targets/README.md index fa9b718e..bc8034ea 100644 --- a/data/targets/README.md +++ b/data/targets/README.md @@ -133,24 +133,16 @@ weighted_data = reweighter.fit_transform(microdata, targets) - **Update frequency**: Annual - **Notes**: 5-year ACS estimates. Includes demographics and economic characteristics. -## Integration with Calibration Pipeline +## Loading targets -These targets integrate with the calibration pipeline in `calibration/`: +The in-repo `calibration/` package that once consumed these files was retired +with the legacy target store. The parquet outputs stand alone as inputs for +downstream Microplex calibration: ```python -from calibration.loader import load_microdata -from calibration.constraints import build_constraint_matrix -from calibration.methods.entropy import entropy_calibrate - -# Load microdata -df = load_microdata("cps", year=2023) - -# Load state targets from data.targets.build_state_targets import load_state_targets -targets = load_state_targets("income_distribution", years=[2023]) -# Build constraints and calibrate -# ... (see calibration module documentation) +targets = load_state_targets("income_distribution", years=[2023]) ``` ## Extending Targets diff --git a/db/__init__.py b/db/__init__.py index dacca2a1..23147bfc 100644 --- a/db/__init__.py +++ b/db/__init__.py @@ -1,7 +1,9 @@ """ -Database layer for Arch target inputs and calibration data. +Database layer for Arch source artifacts and Supabase registries. Supports two backends: -- SQLite (local development via SQLModel) -- Supabase (production via supabase-py) +- SQLite (local source-artifact store via SQLModel) +- Supabase (hosted source and microdata registries via supabase-py) + +The legacy target-input tables and their ETL loaders were retired. """ diff --git a/db/cli.py b/db/cli.py index 467525c6..cd83fad0 100644 --- a/db/cli.py +++ b/db/cli.py @@ -1,4 +1,4 @@ -"""CLI for managing Arch target input data.""" +"""CLI for Arch's local source-artifact database.""" import argparse from pathlib import Path @@ -6,14 +6,11 @@ from sqlalchemy import func from sqlmodel import Session, select -from .etl_soi import load_soi_targets from .schema import ( DEFAULT_DB_PATH, SourceArtifact, SourceRow, SourceTable, - Stratum, - Target, init_db, get_engine, ) @@ -26,144 +23,6 @@ def cmd_init(args): print(f"Initialized database at {db_path}") -def cmd_load(args): - """Load targets from a source.""" - db_path = Path(args.db) if args.db else DEFAULT_DB_PATH - engine = init_db(db_path) - - with Session(engine) as session: - if args.source == "soi" or args.source == "all": - years = [int(y) for y in args.years.split(",")] if args.years else None - load_soi_targets(session, years=years) - print(f"Loaded SOI targets for years: {years or 'all available'}") - - if args.source == "soi-state" or args.source == "all": - from .etl_soi_state import load_soi_state_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_soi_state_targets(session, years=years) - print(f"Loaded state-level SOI targets for years: {years or 'all available'}") - - if args.source == "soi-credits" or args.source == "all": - from .etl_soi_credits import load_soi_credits_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_soi_credits_targets(session, years=years) - print(f"Loaded SOI credits targets for years: {years or 'all available'}") - - if args.source == "soi-income-sources" or args.source == "all": - from .etl_soi_income_sources import load_soi_income_sources_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_soi_income_sources_targets(session, years=years) - print(f"Loaded SOI income sources targets for years: {years or 'all available'}") - - if args.source == "soi-deductions" or args.source == "all": - from .etl_soi_deductions import load_soi_deductions_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_soi_deductions_targets(session, years=years) - print(f"Loaded SOI deductions targets for years: {years or 'all available'}") - - if args.source == "snap" or args.source == "all": - from .etl_snap import load_snap_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_snap_targets(session, years=years) - print(f"Loaded SNAP targets for years: {years or 'all available'}") - - if args.source == "medicaid" or args.source == "all": - from .etl_medicaid import load_medicaid_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_medicaid_targets(session, years=years) - print(f"Loaded Medicaid targets for years: {years or 'all available'}") - - if args.source == "aca" or args.source == "all": - from .etl_aca_enrollment import load_aca_enrollment_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_aca_enrollment_targets(session, years=years) - print(f"Loaded ACA Marketplace targets for years: {years or 'all available'}") - - if args.source == "hmrc" or args.source == "all": - from .etl_hmrc import load_hmrc_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_hmrc_targets(session, years=years) - print(f"Loaded HMRC targets for years: {years or 'all available'}") - - if args.source == "census" or args.source == "all": - from .etl_census import load_census_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_census_targets(session, years=years) - print(f"Loaded Census targets for years: {years or 'all available'}") - - if args.source == "ssa" or args.source == "all": - from .etl_ssa import load_ssa_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_ssa_targets(session, years=years) - print(f"Loaded SSA targets for years: {years or 'all available'}") - - if args.source == "ssi" or args.source == "all": - from .etl_ssi import load_ssi_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_ssi_targets(session, years=years) - print(f"Loaded SSI targets for years: {years or 'all available'}") - - if args.source == "bls" or args.source == "all": - from .etl_bls import load_bls_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_bls_targets(session, years=years) - print(f"Loaded BLS targets for years: {years or 'all available'}") - - if args.source == "cps" or args.source == "all": - from .etl_cps import load_cps_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_cps_targets(session, years=years) - print(f"Loaded CPS monthly targets for years: {years or 'all available'}") - - if args.source == "cbo" or args.source == "all": - from .etl_cbo import load_cbo_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_cbo_targets(session, years=years) - print(f"Loaded CBO projections for years: {years or 'all available'}") - - if args.source == "obr" or args.source == "all": - from .etl_obr import load_obr_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_obr_targets(session, years=years) - print(f"Loaded OBR projections for years: {years or 'all available'}") - - if args.source == "ons" or args.source == "all": - from .etl_ons import load_ons_targets - years = [int(y) for y in args.years.split(",")] if args.years else None - load_ons_targets(session, years=years) - print(f"Loaded ONS projections for years: {years or 'all available'}") - - -def cmd_stats(args): - """Show database statistics.""" - db_path = Path(args.db) if args.db else DEFAULT_DB_PATH - - if not db_path.exists(): - print(f"Database not found at {db_path}") - return - - engine = get_engine(db_path) - with Session(engine) as session: - strata_count = len(session.exec(select(Stratum)).all()) - targets_count = len(session.exec(select(Target)).all()) - - # Get unique sources - sources = session.exec( - select(Target.source).distinct() - ).all() - - # Get unique years - years = session.exec( - select(Target.period).distinct() - ).all() - - print(f"Database: {db_path}") - print(f"Strata: {strata_count}") - print(f"Targets: {targets_count}") - print(f"Sources: {', '.join(str(s) for s in sources)}") - print(f"Years: {', '.join(str(y) for y in sorted(years))}") - - def cmd_load_source_files(args): """Load full source files into the source artifact tables.""" db_path = Path(args.db) if args.db else DEFAULT_DB_PATH @@ -237,40 +96,9 @@ def cmd_source_stats(args): print(f" {source_id}: {count:,}") -def cmd_query(args): - """Query targets.""" - db_path = Path(args.db) if args.db else DEFAULT_DB_PATH - - if not db_path.exists(): - print(f"Database not found at {db_path}") - return - - engine = get_engine(db_path) - with Session(engine) as session: - query = select(Target, Stratum).join(Stratum) - - if args.variable: - query = query.where(Target.variable == args.variable) - if args.year: - query = query.where(Target.period == int(args.year)) - if args.source: - query = query.where(Target.source == args.source) - - results = session.exec(query).all() - - if not results: - print("No targets found matching criteria") - return - - print(f"{'Stratum':<30} {'Variable':<25} {'Year':<6} {'Value':>15} {'Source':<10}") - print("-" * 90) - for target, stratum in results[:args.limit]: - print(f"{stratum.name:<30} {target.variable:<25} {target.period:<6} {target.value:>15,.0f} {target.source.value:<10}") - - def main(): parser = argparse.ArgumentParser( - description="Manage Arch target input data" + description="Manage Arch's local source-artifact database" ) parser.add_argument( "--db", @@ -283,42 +111,6 @@ def main(): init_parser = subparsers.add_parser("init", help="Initialize database") init_parser.set_defaults(func=cmd_init) - # load - load_parser = subparsers.add_parser("load", help="Load targets from source") - load_parser.add_argument( - "source", - choices=[ - "soi", - "soi-state", - "soi-credits", - "soi-income-sources", - "soi-deductions", - "snap", - "medicaid", - "aca", - "hmrc", - "census", - "ssa", - "ssi", - "bls", - "cps", - "cbo", - "obr", - "ons", - "all", - ], - help="Data source to load" - ) - load_parser.add_argument( - "--years", - help="Comma-separated years to load (default: all)" - ) - load_parser.set_defaults(func=cmd_load) - - # stats - stats_parser = subparsers.add_parser("stats", help="Show database statistics") - stats_parser.set_defaults(func=cmd_stats) - # load-source-files source_parser = subparsers.add_parser( "load-source-files", @@ -358,14 +150,6 @@ def main(): ) source_stats_parser.set_defaults(func=cmd_source_stats) - # query - query_parser = subparsers.add_parser("query", help="Query targets") - query_parser.add_argument("--variable", help="Filter by variable name") - query_parser.add_argument("--year", help="Filter by year") - query_parser.add_argument("--source", help="Filter by source") - query_parser.add_argument("--limit", type=int, default=20, help="Max results") - query_parser.set_defaults(func=cmd_query) - args = parser.parse_args() args.func(args) diff --git a/db/etl_aca_enrollment.py b/db/etl_aca_enrollment.py deleted file mode 100644 index 47f2906b..00000000 --- a/db/etl_aca_enrollment.py +++ /dev/null @@ -1,643 +0,0 @@ -""" -ETL for CMS ACA Marketplace enrollment targets. - -Loads data from CMS Marketplace enrollment reports into the targets database. -Data sources: -- https://www.cms.gov/data-research/statistics-trends-reports/marketplace-products/ -- https://www.kff.org/affordable-care-act/state-indicator/marketplace-enrollment/ -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) -from arch.normalization import SourceFact, apply_share, as_target, target_kwargs - -# State FIPS codes -STATE_FIPS = { - "AL": "01", "AK": "02", "AZ": "04", "AR": "05", "CA": "06", - "CO": "08", "CT": "09", "DE": "10", "DC": "11", "FL": "12", - "GA": "13", "HI": "15", "ID": "16", "IL": "17", "IN": "18", - "IA": "19", "KS": "20", "KY": "21", "LA": "22", "ME": "23", - "MD": "24", "MA": "25", "MI": "26", "MN": "27", "MS": "28", - "MO": "29", "MT": "30", "NE": "31", "NV": "32", "NH": "33", - "NJ": "34", "NM": "35", "NY": "36", "NC": "37", "ND": "38", - "OH": "39", "OK": "40", "OR": "41", "PA": "42", "RI": "44", - "SC": "45", "SD": "46", "TN": "47", "TX": "48", "UT": "49", - "VT": "50", "VA": "51", "WA": "53", "WV": "54", "WI": "55", - "WY": "56", -} - -# ACA Marketplace enrollment data by year -# Sources: -# - CMS Marketplace Open Enrollment Reports -# - KFF State Health Facts: https://www.kff.org/affordable-care-act/state-indicator/marketplace-enrollment/ -# - CMS effectuated enrollment snapshots -ACA_ENROLLMENT_DATA = { - 2025: { - "national": { - # Total marketplace plan selections - "total_enrollment": 24_319_713, - "new_enrollees": 3_900_000, - "returning_enrollees": 20_400_000, - # Subsidy recipients (APTC = Advance Premium Tax Credit) - "aptc_recipients": 22_400_000, # ~92% of enrollees - "aptc_recipients_pct": 0.92, - # Average premiums and subsidies - "avg_monthly_premium_before_aptc": 600, # dollars - "avg_monthly_aptc": 528, # dollars - "avg_monthly_premium_after_aptc": 101, # dollars - # Metal level distribution (national percentages) - "bronze_pct": 0.31, - "silver_pct": 0.54, - "gold_pct": 0.13, - "platinum_pct": 0.01, - "catastrophic_pct": 0.01, - }, - # State-level enrollment (plan selections) - # Source: KFF State Health Facts, CMS OEP reports - "states": { - "AL": {"enrollment": 268_857}, - "AK": {"enrollment": 20_844}, - "AZ": {"enrollment": 423_025}, - "AR": {"enrollment": 137_442}, - "CA": {"enrollment": 1_979_504}, - "CO": {"enrollment": 229_085}, - "CT": {"enrollment": 104_789}, - "DE": {"enrollment": 35_213}, - "DC": {"enrollment": 20_178}, - "FL": {"enrollment": 4_735_415}, - "GA": {"enrollment": 1_510_852}, - "HI": {"enrollment": 27_897}, - "ID": {"enrollment": 120_087}, - "IL": {"enrollment": 459_873}, - "IN": {"enrollment": 285_124}, - "IA": {"enrollment": 99_456}, - "KS": {"enrollment": 136_789}, - "KY": {"enrollment": 112_345}, - "LA": {"enrollment": 295_432}, - "ME": {"enrollment": 78_234}, - "MD": {"enrollment": 189_567}, - "MA": {"enrollment": 322_456}, - "MI": {"enrollment": 531_083}, - "MN": {"enrollment": 142_567}, - "MS": {"enrollment": 198_765}, - "MO": {"enrollment": 367_234}, - "MT": {"enrollment": 78_123}, - "NE": {"enrollment": 123_456}, - "NV": {"enrollment": 118_234}, - "NH": {"enrollment": 62_345}, - "NJ": {"enrollment": 398_567}, - "NM": {"enrollment": 78_234}, - "NY": {"enrollment": 221_534}, - "NC": {"enrollment": 975_110}, - "ND": {"enrollment": 32_145}, - "OH": {"enrollment": 583_443}, - "OK": {"enrollment": 267_345}, - "OR": {"enrollment": 156_789}, - "PA": {"enrollment": 496_661}, - "RI": {"enrollment": 42_345}, - "SC": {"enrollment": 398_765}, - "SD": {"enrollment": 42_678}, - "TN": {"enrollment": 398_765}, - "TX": {"enrollment": 3_966_226}, - "UT": {"enrollment": 287_654}, - "VT": {"enrollment": 32_456}, - "VA": {"enrollment": 456_789}, - "WA": {"enrollment": 267_890}, - "WV": {"enrollment": 52_345}, - "WI": {"enrollment": 267_456}, - "WY": {"enrollment": 28_765}, - }, - }, - 2024: { - "national": { - "total_enrollment": 21_446_150, - "new_enrollees": 5_000_000, - "returning_enrollees": 16_300_000, - "aptc_recipients": 19_500_000, # 93% of effectuated enrollment - "aptc_recipients_pct": 0.93, - "csr_recipients": 10_400_000, # Cost-sharing reduction recipients - "avg_monthly_premium_before_aptc": 592, - "avg_monthly_aptc": 528, - "avg_monthly_premium_after_aptc": 101, - "bronze_pct": 0.31, - "silver_pct": 0.54, - "gold_pct": 0.13, - "platinum_pct": 0.01, - "catastrophic_pct": 0.01, - }, - "states": { - "AL": {"enrollment": 235_678}, - "AK": {"enrollment": 18_234}, - "AZ": {"enrollment": 348_055}, - "AR": {"enrollment": 118_456}, - "CA": {"enrollment": 1_784_653}, - "CO": {"enrollment": 198_765}, - "CT": {"enrollment": 98_234}, - "DE": {"enrollment": 32_456}, - "DC": {"enrollment": 19_987}, - "FL": {"enrollment": 4_211_902}, - "GA": {"enrollment": 1_305_114}, - "HI": {"enrollment": 26_543}, - "ID": {"enrollment": 105_678}, - "IL": {"enrollment": 398_765}, - "IN": {"enrollment": 245_678}, - "IA": {"enrollment": 87_654}, - "KS": {"enrollment": 118_765}, - "KY": {"enrollment": 87_234}, - "LA": {"enrollment": 214_567}, - "ME": {"enrollment": 75_234}, - "MD": {"enrollment": 167_890}, - "MA": {"enrollment": 298_765}, - "MI": {"enrollment": 418_100}, - "MN": {"enrollment": 128_765}, - "MS": {"enrollment": 167_234}, - "MO": {"enrollment": 318_765}, - "MT": {"enrollment": 67_234}, - "NE": {"enrollment": 108_765}, - "NV": {"enrollment": 108_234}, - "NH": {"enrollment": 54_678}, - "NJ": {"enrollment": 308_765}, - "NM": {"enrollment": 68_765}, - "NY": {"enrollment": 288_681}, - "NC": {"enrollment": 1_027_930}, - "ND": {"enrollment": 28_765}, - "OH": {"enrollment": 477_793}, - "OK": {"enrollment": 234_567}, - "OR": {"enrollment": 148_765}, - "PA": {"enrollment": 434_571}, - "RI": {"enrollment": 38_765}, - "SC": {"enrollment": 348_765}, - "SD": {"enrollment": 41_234}, - "TN": {"enrollment": 345_678}, - "TX": {"enrollment": 3_484_632}, - "UT": {"enrollment": 248_765}, - "VT": {"enrollment": 28_765}, - "VA": {"enrollment": 398_765}, - "WA": {"enrollment": 238_765}, - "WV": {"enrollment": 39_876}, - "WI": {"enrollment": 234_567}, - "WY": {"enrollment": 24_567}, - }, - }, - 2023: { - "national": { - "total_enrollment": 16_357_030, - "new_enrollees": 3_600_000, - "returning_enrollees": 12_700_000, - "aptc_recipients": 14_800_000, - "aptc_recipients_pct": 0.92, - "avg_monthly_premium_before_aptc": 593, - "avg_monthly_aptc": 520, - "avg_monthly_premium_after_aptc": 117, - "bronze_pct": 0.32, - "silver_pct": 0.53, - "gold_pct": 0.13, - "platinum_pct": 0.01, - "catastrophic_pct": 0.01, - }, - "states": { - "AL": {"enrollment": 178_456}, - "AK": {"enrollment": 15_678}, - "AZ": {"enrollment": 235_229}, - "AR": {"enrollment": 89_234}, - "CA": {"enrollment": 1_739_368}, - "CO": {"enrollment": 178_234}, - "CT": {"enrollment": 87_654}, - "DE": {"enrollment": 28_765}, - "DC": {"enrollment": 18_765}, - "FL": {"enrollment": 3_225_435}, - "GA": {"enrollment": 879_084}, - "HI": {"enrollment": 24_567}, - "ID": {"enrollment": 89_876}, - "IL": {"enrollment": 298_765}, - "IN": {"enrollment": 187_654}, - "IA": {"enrollment": 67_234}, - "KS": {"enrollment": 89_765}, - "KY": {"enrollment": 67_234}, - "LA": {"enrollment": 148_765}, - "ME": {"enrollment": 65_432}, - "MD": {"enrollment": 145_678}, - "MA": {"enrollment": 267_890}, - "MI": {"enrollment": 322_273}, - "MN": {"enrollment": 108_765}, - "MS": {"enrollment": 118_765}, - "MO": {"enrollment": 245_678}, - "MT": {"enrollment": 54_678}, - "NE": {"enrollment": 87_234}, - "NV": {"enrollment": 98_765}, - "NH": {"enrollment": 45_678}, - "NJ": {"enrollment": 234_567}, - "NM": {"enrollment": 54_321}, - "NY": {"enrollment": 214_052}, - "NC": {"enrollment": 800_850}, - "ND": {"enrollment": 23_456}, - "OH": {"enrollment": 294_644}, - "OK": {"enrollment": 178_234}, - "OR": {"enrollment": 134_567}, - "PA": {"enrollment": 371_516}, - "RI": {"enrollment": 34_567}, - "SC": {"enrollment": 278_234}, - "SD": {"enrollment": 34_567}, - "TN": {"enrollment": 267_890}, - "TX": {"enrollment": 2_410_810}, - "UT": {"enrollment": 198_765}, - "VT": {"enrollment": 23_456}, - "VA": {"enrollment": 312_456}, - "WA": {"enrollment": 212_345}, - "WV": {"enrollment": 28_765}, - "WI": {"enrollment": 187_234}, - "WY": {"enrollment": 19_876}, - }, - }, -} - -# Metal level enrollment by state (2024 data) -# Source: KFF State Health Facts, CMS OEP State, Metal Level PUF -METAL_LEVEL_BY_STATE = { - 2024: { - # Format: {state: {"bronze_pct": x, "silver_pct": x, "gold_pct": x}} - # National averages used where state-specific not available - "FL": {"bronze_pct": 0.31, "silver_pct": 0.63, "gold_pct": 0.05}, - "CA": {"bronze_pct": 0.26, "silver_pct": 0.60, "gold_pct": 0.09}, - "TX": {"bronze_pct": 0.16, "silver_pct": 0.49, "gold_pct": 0.34}, - "GA": {"bronze_pct": 0.28, "silver_pct": 0.58, "gold_pct": 0.12}, - "NC": {"bronze_pct": 0.25, "silver_pct": 0.57, "gold_pct": 0.15}, - "NJ": {"bronze_pct": 0.10, "silver_pct": 0.81, "gold_pct": 0.08}, - "MD": {"bronze_pct": 0.15, "silver_pct": 0.35, "gold_pct": 0.48}, - }, -} - -SOURCE_URL = "https://www.cms.gov/data-research/statistics-trends-reports/marketplace-products/" -KFF_SOURCE_URL = "https://www.kff.org/affordable-care-act/state-indicator/marketplace-enrollment/" - - -def build_aca_metal_enrollment_target( - stratum: Stratum, - *, - total_enrollment: float, - share: float, - share_name: str, - period: int, - source_table: str, - source_url: str, -) -> Target: - """Build ACA metal-level enrollment from total enrollment and a share.""" - total_fact = SourceFact( - name="aca_marketplace_enrollment", - value=total_enrollment, - period=period, - unit="count", - source=DataSource.CMS_ACA, - jurisdiction=stratum.jurisdiction, - source_table=source_table, - source_url=source_url, - ) - metal_fact = apply_share( - total_fact, - share, - name="aca_marketplace_metal_level_enrollment", - share_name=share_name, - unit="count", - ) - blueprint = as_target( - metal_fact, - variable="aca_marketplace_enrollment", - target_type=TargetType.COUNT, - stratum_name=stratum.name, - ) - kwargs = target_kwargs(blueprint, stratum_id=stratum.id) - kwargs["value"] = int(kwargs["value"]) - return Target(**kwargs) - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_aca_enrollment_targets(session: Session, years: list[int] | None = None): - """ - Load ACA Marketplace enrollment targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(ACA_ENROLLMENT_DATA.keys()) - - for year in years: - if year not in ACA_ENROLLMENT_DATA: - continue - - data = ACA_ENROLLMENT_DATA[year] - national_data = data["national"] - - # Create national ACA marketplace stratum - national_stratum = get_or_create_stratum( - session, - name="US ACA Marketplace Enrollees", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("aca_marketplace", "==", "1")], - description="All ACA Marketplace plan selections", - stratum_group_id="aca_national", - ) - - # Total enrollment - session.add( - Target( - stratum_id=national_stratum.id, - variable="aca_marketplace_enrollment", - period=year, - value=national_data["total_enrollment"], - target_type=TargetType.COUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace OEP Report", - source_url=SOURCE_URL, - ) - ) - - # New enrollees - session.add( - Target( - stratum_id=national_stratum.id, - variable="aca_marketplace_new_enrollees", - period=year, - value=national_data["new_enrollees"], - target_type=TargetType.COUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace OEP Report", - source_url=SOURCE_URL, - ) - ) - - # APTC recipients stratum - aptc_stratum = get_or_create_stratum( - session, - name="US ACA APTC Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("aca_marketplace", "==", "1"), - ("receives_aptc", "==", "1"), - ], - description="ACA Marketplace enrollees receiving Advance Premium Tax Credit", - parent_id=national_stratum.id, - stratum_group_id="aca_subsidies", - ) - - session.add( - Target( - stratum_id=aptc_stratum.id, - variable="aca_aptc_recipients", - period=year, - value=national_data["aptc_recipients"], - target_type=TargetType.COUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace Effectuated Enrollment", - source_url=SOURCE_URL, - ) - ) - - # Average premium and subsidy amounts - session.add( - Target( - stratum_id=national_stratum.id, - variable="aca_avg_monthly_premium_gross", - period=year, - value=national_data["avg_monthly_premium_before_aptc"], - target_type=TargetType.AMOUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace Effectuated Enrollment", - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=aptc_stratum.id, - variable="aca_avg_monthly_aptc", - period=year, - value=national_data["avg_monthly_aptc"], - target_type=TargetType.AMOUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace Effectuated Enrollment", - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=national_stratum.id, - variable="aca_avg_monthly_premium_net", - period=year, - value=national_data["avg_monthly_premium_after_aptc"], - target_type=TargetType.AMOUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace Effectuated Enrollment", - source_url=SOURCE_URL, - ) - ) - - # CSR recipients (if available) - if "csr_recipients" in national_data: - csr_stratum = get_or_create_stratum( - session, - name="US ACA CSR Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("aca_marketplace", "==", "1"), - ("receives_csr", "==", "1"), - ], - description="ACA Marketplace enrollees receiving Cost-Sharing Reductions", - parent_id=national_stratum.id, - stratum_group_id="aca_subsidies", - ) - - session.add( - Target( - stratum_id=csr_stratum.id, - variable="aca_csr_recipients", - period=year, - value=national_data["csr_recipients"], - target_type=TargetType.COUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace Effectuated Enrollment", - source_url=SOURCE_URL, - ) - ) - - # Metal level strata - for metal_level in ["bronze", "silver", "gold", "platinum", "catastrophic"]: - pct_key = f"{metal_level}_pct" - if pct_key in national_data: - metal_stratum = get_or_create_stratum( - session, - name=f"US ACA {metal_level.capitalize()} Plan Enrollees", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("aca_marketplace", "==", "1"), - ("aca_metal_level", "==", metal_level), - ], - description=f"ACA Marketplace enrollees in {metal_level} plans", - parent_id=national_stratum.id, - stratum_group_id="aca_metal_levels", - ) - - session.add( - build_aca_metal_enrollment_target( - metal_stratum, - total_enrollment=national_data["total_enrollment"], - share=national_data[pct_key], - share_name=pct_key, - period=year, - source_table="Marketplace OEP Metal Level Report", - source_url=SOURCE_URL, - ) - ) - - # Also store the percentage as a rate - session.add( - Target( - stratum_id=metal_stratum.id, - variable="aca_metal_level_share", - period=year, - value=national_data[pct_key], - target_type=TargetType.RATE, - source=DataSource.CMS_ACA, - source_table="Marketplace OEP Metal Level Report", - source_url=SOURCE_URL, - ) - ) - - # Add state-level targets - for state_abbrev, state_data in data.get("states", {}).items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} ACA Marketplace Enrollees", - jurisdiction=Jurisdiction.US, - constraints=[ - ("aca_marketplace", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"ACA Marketplace enrollees in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="aca_states", - ) - - session.add( - Target( - stratum_id=state_stratum.id, - variable="aca_marketplace_enrollment", - period=year, - value=state_data["enrollment"], - target_type=TargetType.COUNT, - source=DataSource.CMS_ACA, - source_table="Marketplace OEP State Report", - source_url=KFF_SOURCE_URL, - ) - ) - - # Add state-level metal distribution if available - if year in METAL_LEVEL_BY_STATE: - metal_data = METAL_LEVEL_BY_STATE[year].get(state_abbrev) - if metal_data: - for metal_level, pct in metal_data.items(): - if metal_level.endswith("_pct"): - level_name = metal_level.replace("_pct", "") - state_metal_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} ACA {level_name.capitalize()} Plan Enrollees", - jurisdiction=Jurisdiction.US, - constraints=[ - ("aca_marketplace", "==", "1"), - ("state_fips", "==", fips), - ("aca_metal_level", "==", level_name), - ], - description=f"ACA {level_name} plan enrollees in {state_abbrev}", - parent_id=state_stratum.id, - stratum_group_id="aca_state_metal_levels", - ) - - session.add( - build_aca_metal_enrollment_target( - state_metal_stratum, - total_enrollment=state_data["enrollment"], - share=pct, - share_name=metal_level, - period=year, - source_table="Marketplace OEP State Metal Level Report", - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the ACA enrollment ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_aca_enrollment_targets(session) - print(f"Loaded ACA Marketplace enrollment targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_bls.py b/db/etl_bls.py deleted file mode 100644 index 8809ec6a..00000000 --- a/db/etl_bls.py +++ /dev/null @@ -1,252 +0,0 @@ -""" -ETL for BLS (Bureau of Labor Statistics) employment targets. - -Loads data from BLS employment statistics into the targets database. -Data sources: -- https://www.bls.gov/cps/ -- https://www.bls.gov/oes/ -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# BLS employment data by year -# Source: Current Population Survey (CPS) and Current Employment Statistics (CES) -# https://www.bls.gov/cps/tables.htm -BLS_DATA = { - 2023: { - # Labor force (in thousands, converted to actual) - "civilian_labor_force": 167_300_000, - "employed": 161_100_000, - "unemployed": 6_200_000, - "not_in_labor_force": 99_800_000, - - # Rates (percentages stored as decimals) - "unemployment_rate": 3.7, # percent - "labor_force_participation_rate": 62.6, # percent - "employment_population_ratio": 60.3, # percent - - # Earnings - "median_weekly_earnings": 1_145, # dollars, full-time workers - "median_hourly_wage": 23.40, # dollars, all workers - - # By sector (nonfarm payroll, thousands) - "total_nonfarm_employment": 157_000_000, - "private_employment": 133_600_000, - "government_employment": 23_400_000, - }, - 2022: { - "civilian_labor_force": 164_300_000, - "employed": 158_300_000, - "unemployed": 6_000_000, - "not_in_labor_force": 99_600_000, - - "unemployment_rate": 3.6, - "labor_force_participation_rate": 62.2, - "employment_population_ratio": 59.9, - - "median_weekly_earnings": 1_085, - "median_hourly_wage": 22.00, - - "total_nonfarm_employment": 153_000_000, - "private_employment": 130_200_000, - "government_employment": 22_800_000, - }, - 2021: { - "civilian_labor_force": 161_200_000, - "employed": 152_600_000, - "unemployed": 8_600_000, - "not_in_labor_force": 100_000_000, - - "unemployment_rate": 5.4, - "labor_force_participation_rate": 61.7, - "employment_population_ratio": 58.4, - - "median_weekly_earnings": 1_013, - "median_hourly_wage": 20.50, - - "total_nonfarm_employment": 146_100_000, - "private_employment": 124_200_000, - "government_employment": 21_900_000, - }, -} - -SOURCE_URL = "https://www.bls.gov/cps/" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_bls_targets(session: Session, years: list[int] | None = None): - """ - Load BLS employment targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(BLS_DATA.keys()) - - for year in years: - if year not in BLS_DATA: - continue - - data = BLS_DATA[year] - - # Create national labor force stratum - labor_force_stratum = get_or_create_stratum( - session, - name="US Labor Force", - jurisdiction=Jurisdiction.US, - constraints=[("in_labor_force", "==", "1")], - description="US civilian labor force (16+)", - stratum_group_id="bls_national", - ) - - # Employment count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="employed", - period=year, - value=data["employed"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Unemployment count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="unemployed", - period=year, - value=data["unemployed"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Unemployment rate - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="unemployment_rate", - period=year, - value=data["unemployment_rate"], - target_type=TargetType.RATE, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Labor force participation rate - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="labor_force_participation_rate", - period=year, - value=data["labor_force_participation_rate"], - target_type=TargetType.RATE, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Median weekly earnings - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="median_weekly_earnings", - period=year, - value=data["median_weekly_earnings"], - target_type=TargetType.AMOUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Total labor force count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="labor_force_count", - period=year, - value=data["civilian_labor_force"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the BLS ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_bls_targets(session) - print(f"Loaded BLS targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_cbo.py b/db/etl_cbo.py deleted file mode 100644 index 5c5c06dc..00000000 --- a/db/etl_cbo.py +++ /dev/null @@ -1,366 +0,0 @@ -""" -ETL for CBO (Congressional Budget Office) projections. - -Loads budget and economic projections from CBO's Budget and Economic Outlook. -Data source: https://www.cbo.gov/data/budget-economic-data -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# CBO projections data (10-year baseline, 2024-2034) -# Source: CBO Budget and Economic Outlook, February 2024 -# https://www.cbo.gov/publication/59710 -# GDP and budget figures in billions of dollars -# Rates as percentages -CBO_DATA = { - # Fiscal Year 2024 - 2024: { - # Budget (billions) - "gdp": 28_269_000_000_000, # $28.269 trillion - "federal_revenue": 4_869_000_000_000, - "federal_outlays": 6_751_000_000_000, - "federal_deficit": 1_882_000_000_000, - "debt_held_by_public": 28_166_000_000_000, - # Economic indicators (rates as percentages) - "unemployment_rate": 4.0, - "cpi_inflation": 2.9, - "interest_rate_10yr": 4.5, - "real_gdp_growth": 2.0, - # Labor force (thousands) - "labor_force": 167_900_000, - }, - 2025: { - "gdp": 29_458_000_000_000, - "federal_revenue": 5_198_000_000_000, - "federal_outlays": 7_025_000_000_000, - "federal_deficit": 1_827_000_000_000, - "debt_held_by_public": 30_065_000_000_000, - "unemployment_rate": 4.4, - "cpi_inflation": 2.3, - "interest_rate_10yr": 4.2, - "real_gdp_growth": 1.5, - "labor_force": 168_500_000, - }, - 2026: { - "gdp": 30_637_000_000_000, - "federal_revenue": 5_551_000_000_000, - "federal_outlays": 7_297_000_000_000, - "federal_deficit": 1_746_000_000_000, - "debt_held_by_public": 31_887_000_000_000, - "unemployment_rate": 4.4, - "cpi_inflation": 2.3, - "interest_rate_10yr": 4.0, - "real_gdp_growth": 2.1, - "labor_force": 169_100_000, - }, - 2027: { - "gdp": 31_854_000_000_000, - "federal_revenue": 5_852_000_000_000, - "federal_outlays": 7_556_000_000_000, - "federal_deficit": 1_704_000_000_000, - "debt_held_by_public": 33_662_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.9, - "real_gdp_growth": 2.0, - "labor_force": 169_600_000, - }, - 2028: { - "gdp": 33_102_000_000_000, - "federal_revenue": 6_131_000_000_000, - "federal_outlays": 7_901_000_000_000, - "federal_deficit": 1_770_000_000_000, - "debt_held_by_public": 35_501_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 2.0, - "labor_force": 170_100_000, - }, - 2029: { - "gdp": 34_388_000_000_000, - "federal_revenue": 6_397_000_000_000, - "federal_outlays": 8_241_000_000_000, - "federal_deficit": 1_844_000_000_000, - "debt_held_by_public": 37_413_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 1.9, - "labor_force": 170_500_000, - }, - 2030: { - "gdp": 35_714_000_000_000, - "federal_revenue": 6_667_000_000_000, - "federal_outlays": 8_588_000_000_000, - "federal_deficit": 1_921_000_000_000, - "debt_held_by_public": 39_399_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 1.9, - "labor_force": 170_900_000, - }, - 2031: { - "gdp": 37_080_000_000_000, - "federal_revenue": 6_934_000_000_000, - "federal_outlays": 8_960_000_000_000, - "federal_deficit": 2_026_000_000_000, - "debt_held_by_public": 41_488_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 1.9, - "labor_force": 171_200_000, - }, - 2032: { - "gdp": 38_487_000_000_000, - "federal_revenue": 7_212_000_000_000, - "federal_outlays": 9_377_000_000_000, - "federal_deficit": 2_165_000_000_000, - "debt_held_by_public": 43_714_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 1.8, - "labor_force": 171_500_000, - }, - 2033: { - "gdp": 39_934_000_000_000, - "federal_revenue": 7_499_000_000_000, - "federal_outlays": 9_822_000_000_000, - "federal_deficit": 2_323_000_000_000, - "debt_held_by_public": 46_096_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 1.8, - "labor_force": 171_700_000, - }, - 2034: { - "gdp": 41_422_000_000_000, - "federal_revenue": 7_796_000_000_000, - "federal_outlays": 10_289_000_000_000, - "federal_deficit": 2_493_000_000_000, - "debt_held_by_public": 48_645_000_000_000, - "unemployment_rate": 4.3, - "cpi_inflation": 2.3, - "interest_rate_10yr": 3.8, - "real_gdp_growth": 1.8, - "labor_force": 171_900_000, - }, -} - -SOURCE_URL = "https://www.cbo.gov/data/budget-economic-data" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_cbo_targets(session: Session, years: list[int] | None = None): - """ - Load CBO projections into database. - - Args: - session: Database session - years: Years to load (default: all available 2024-2034) - """ - if years is None: - years = list(CBO_DATA.keys()) - - for year in years: - if year not in CBO_DATA: - continue - - data = CBO_DATA[year] - - # Federal budget stratum - budget_stratum = get_or_create_stratum( - session, - name="US Federal Budget", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("sector", "==", "public")], # Government sector - description="US Federal Government budget", - stratum_group_id="cbo_budget", - ) - - # Budget targets - budget_vars = [ - ("gdp", TargetType.AMOUNT), - ("federal_revenue", TargetType.AMOUNT), - ("federal_outlays", TargetType.AMOUNT), - ("federal_deficit", TargetType.AMOUNT), - ("debt_held_by_public", TargetType.AMOUNT), - ] - - for var_name, var_type in budget_vars: - if var_name not in data: - continue - - # Check for existing target - existing = session.exec( - select(Target).where( - Target.stratum_id == budget_stratum.id, - Target.variable == var_name, - Target.period == year, - Target.source == DataSource.CBO, - ) - ).first() - - if existing: - continue - - session.add( - Target( - stratum_id=budget_stratum.id, - variable=var_name, - period=year, - value=data[var_name], - target_type=var_type, - source=DataSource.CBO, - source_url=SOURCE_URL, - is_preliminary=year > 2024, # Projections are preliminary - ) - ) - - # US Economy stratum for macro indicators - economy_stratum = get_or_create_stratum( - session, - name="US Economy", - jurisdiction=Jurisdiction.US, - constraints=[], # Whole economy, not a subset - description="US macroeconomic indicators", - stratum_group_id="cbo_economy", - ) - - # Economic rate targets - rate_vars = [ - "unemployment_rate", - "cpi_inflation", - "interest_rate_10yr", - "real_gdp_growth", - ] - - for var_name in rate_vars: - if var_name not in data: - continue - - existing = session.exec( - select(Target).where( - Target.stratum_id == economy_stratum.id, - Target.variable == var_name, - Target.period == year, - Target.source == DataSource.CBO, - ) - ).first() - - if existing: - continue - - session.add( - Target( - stratum_id=economy_stratum.id, - variable=var_name, - period=year, - value=data[var_name], - target_type=TargetType.RATE, - source=DataSource.CBO, - source_url=SOURCE_URL, - is_preliminary=year > 2024, - ) - ) - - # Labor force count - if "labor_force" in data: - existing = session.exec( - select(Target).where( - Target.stratum_id == economy_stratum.id, - Target.variable == "labor_force", - Target.period == year, - Target.source == DataSource.CBO, - ) - ).first() - - if not existing: - session.add( - Target( - stratum_id=economy_stratum.id, - variable="labor_force", - period=year, - value=data["labor_force"], - target_type=TargetType.COUNT, - source=DataSource.CBO, - source_url=SOURCE_URL, - is_preliminary=year > 2024, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the CBO ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_cbo_targets(session) - print(f"Loaded CBO projections to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_census.py b/db/etl_census.py deleted file mode 100644 index 249d02b4..00000000 --- a/db/etl_census.py +++ /dev/null @@ -1,465 +0,0 @@ -""" -ETL for Census Bureau population targets. - -Loads data from Census Bureau population estimates into the targets database. -Data source: https://www.census.gov/programs-surveys/popest.html -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - GeographicLevel, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes -STATE_FIPS = { - "AL": "01", "AK": "02", "AZ": "04", "AR": "05", "CA": "06", - "CO": "08", "CT": "09", "DE": "10", "DC": "11", "FL": "12", - "GA": "13", "HI": "15", "ID": "16", "IL": "17", "IN": "18", - "IA": "19", "KS": "20", "KY": "21", "LA": "22", "ME": "23", - "MD": "24", "MA": "25", "MI": "26", "MN": "27", "MS": "28", - "MO": "29", "MT": "30", "NE": "31", "NV": "32", "NH": "33", - "NJ": "34", "NM": "35", "NY": "36", "NC": "37", "ND": "38", - "OH": "39", "OK": "40", "OR": "41", "PA": "42", "RI": "44", - "SC": "45", "SD": "46", "TN": "47", "TX": "48", "UT": "49", - "VT": "50", "VA": "51", "WA": "53", "WV": "54", "WI": "55", - "WY": "56", "PR": "72", -} - -STATE_NAMES = { - "AL": "Alabama", "AK": "Alaska", "AZ": "Arizona", "AR": "Arkansas", - "CA": "California", "CO": "Colorado", "CT": "Connecticut", "DE": "Delaware", - "DC": "District of Columbia", "FL": "Florida", "GA": "Georgia", "HI": "Hawaii", - "ID": "Idaho", "IL": "Illinois", "IN": "Indiana", "IA": "Iowa", - "KS": "Kansas", "KY": "Kentucky", "LA": "Louisiana", "ME": "Maine", - "MD": "Maryland", "MA": "Massachusetts", "MI": "Michigan", "MN": "Minnesota", - "MS": "Mississippi", "MO": "Missouri", "MT": "Montana", "NE": "Nebraska", - "NV": "Nevada", "NH": "New Hampshire", "NJ": "New Jersey", "NM": "New Mexico", - "NY": "New York", "NC": "North Carolina", "ND": "North Dakota", "OH": "Ohio", - "OK": "Oklahoma", "OR": "Oregon", "PA": "Pennsylvania", "RI": "Rhode Island", - "SC": "South Carolina", "SD": "South Dakota", "TN": "Tennessee", "TX": "Texas", - "UT": "Utah", "VT": "Vermont", "VA": "Virginia", "WA": "Washington", - "WV": "West Virginia", "WI": "Wisconsin", "WY": "Wyoming", "PR": "Puerto Rico", -} - -# Census data by year -# Source: Census Bureau Population Estimates -# https://www.census.gov/data/tables/time-series/demo/popest/2020s-national-detail.html -CENSUS_DATA = { - 2023: { - "total_population": 334_914_895, - "households": 131_432_000, - "median_age": 38.9, - - # Age groups (5-year brackets matching PolicyEngine) - # Based on Census Bureau Population Estimates - # Source: https://www.census.gov/programs-surveys/popest.html - "age_groups": { - "0_to_4": 19_234_000, - "5_to_9": 20_187_000, - "10_to_14": 20_612_000, - "15_to_19": 20_876_000, - "20_to_24": 21_543_000, - "25_to_29": 23_234_000, - "30_to_34": 23_567_000, - "35_to_39": 22_876_000, - "40_to_44": 20_987_000, - "45_to_49": 20_234_000, - "50_to_54": 19_876_000, - "55_to_59": 21_234_000, - "60_to_64": 20_987_000, - "65_to_69": 18_765_000, - "70_to_74": 16_234_000, - "75_to_79": 11_876_000, - "80_to_84": 7_654_000, - "85_plus": 4_938_895, - }, - - # Top 10 states by population - "states": { - "CA": {"population": 38_965_193, "households": 14_234_000}, - "TX": {"population": 30_503_301, "households": 10_876_000}, - "FL": {"population": 22_610_726, "households": 8_765_000}, - "NY": {"population": 19_571_216, "households": 7_654_000}, - "PA": {"population": 12_961_683, "households": 5_234_000}, - "IL": {"population": 12_549_689, "households": 4_987_000}, - "OH": {"population": 11_785_935, "households": 4_765_000}, - "GA": {"population": 11_029_227, "households": 4_123_000}, - "NC": {"population": 10_835_491, "households": 4_098_000}, - "MI": {"population": 10_037_261, "households": 3_987_000}, - }, - }, - 2022: { - "total_population": 333_287_557, - "households": 130_456_000, - "median_age": 38.8, - - # Age groups (5-year brackets matching PolicyEngine) - "age_groups": { - "0_to_4": 19_112_000, - "5_to_9": 20_034_000, - "10_to_14": 20_487_000, - "15_to_19": 20_743_000, - "20_to_24": 21_398_000, - "25_to_29": 23_076_000, - "30_to_34": 23_398_000, - "35_to_39": 22_701_000, - "40_to_44": 20_834_000, - "45_to_49": 20_098_000, - "50_to_54": 19_743_000, - "55_to_59": 21_098_000, - "60_to_64": 20_834_000, - "65_to_69": 18_612_000, - "70_to_74": 16_098_000, - "75_to_79": 11_743_000, - "80_to_84": 7_587_000, - "85_plus": 4_891_557, - }, - - "states": { - "CA": {"population": 39_029_342, "households": 14_123_000}, - "TX": {"population": 30_029_572, "households": 10_654_000}, - "FL": {"population": 22_244_823, "households": 8_543_000}, - "NY": {"population": 19_677_151, "households": 7_543_000}, - "PA": {"population": 12_972_008, "households": 5_198_000}, - "IL": {"population": 12_582_032, "households": 4_876_000}, - "OH": {"population": 11_756_058, "households": 4_654_000}, - "GA": {"population": 10_912_876, "households": 4_023_000}, - "NC": {"population": 10_698_973, "households": 3_987_000}, - "MI": {"population": 10_034_113, "households": 3_876_000}, - }, - }, - 2021: { - "total_population": 331_893_745, - "households": 129_876_000, - "median_age": 38.6, - - # Age groups (5-year brackets matching PolicyEngine) - "age_groups": { - "0_to_4": 19_045_000, - "5_to_9": 19_932_000, - "10_to_14": 20_398_000, - "15_to_19": 20_643_000, - "20_to_24": 21_287_000, - "25_to_29": 22_943_000, - "30_to_34": 23_267_000, - "35_to_39": 22_567_000, - "40_to_44": 20_698_000, - "45_to_49": 19_987_000, - "50_to_54": 19_643_000, - "55_to_59": 20_987_000, - "60_to_64": 20_698_000, - "65_to_69": 18_487_000, - "70_to_74": 15_987_000, - "75_to_79": 11_643_000, - "80_to_84": 7_532_000, - "85_plus": 4_857_745, - }, - - "states": { - "CA": {"population": 39_142_991, "households": 14_012_000}, - "TX": {"population": 29_527_941, "households": 10_432_000}, - "FL": {"population": 21_781_128, "households": 8_321_000}, - "NY": {"population": 19_835_913, "households": 7_432_000}, - "PA": {"population": 12_964_056, "households": 5_123_000}, - "IL": {"population": 12_671_469, "households": 4_765_000}, - "OH": {"population": 11_780_017, "households": 4_543_000}, - "GA": {"population": 10_799_566, "households": 3_912_000}, - "NC": {"population": 10_551_162, "households": 3_876_000}, - "MI": {"population": 10_050_811, "households": 3_765_000}, - }, - }, -} - -SOURCE_URL = "https://www.census.gov/programs-surveys/popest.html" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_census_targets(session: Session, years: list[int] | None = None): - """ - Load Census population targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(CENSUS_DATA.keys()) - - for year in years: - if year not in CENSUS_DATA: - continue - - data = CENSUS_DATA[year] - - # Create national population stratum - national_stratum = get_or_create_stratum( - session, - name="US Population", - jurisdiction=Jurisdiction.US, - constraints=[], # No constraints = total US - description="Total US resident population", - stratum_group_id="population_national", - ) - - # Total population - session.add( - Target( - stratum_id=national_stratum.id, - variable="population", - period=year, - value=data["total_population"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.NATIONAL, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - ) - ) - - # Households - session.add( - Target( - stratum_id=national_stratum.id, - variable="household_count", - period=year, - value=data["households"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.NATIONAL, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - ) - ) - - # Age group strata (18 brackets matching PolicyEngine) - age_brackets = { - "0_to_4": [("age", ">=", "0"), ("age", "<", "5")], - "5_to_9": [("age", ">=", "5"), ("age", "<", "10")], - "10_to_14": [("age", ">=", "10"), ("age", "<", "15")], - "15_to_19": [("age", ">=", "15"), ("age", "<", "20")], - "20_to_24": [("age", ">=", "20"), ("age", "<", "25")], - "25_to_29": [("age", ">=", "25"), ("age", "<", "30")], - "30_to_34": [("age", ">=", "30"), ("age", "<", "35")], - "35_to_39": [("age", ">=", "35"), ("age", "<", "40")], - "40_to_44": [("age", ">=", "40"), ("age", "<", "45")], - "45_to_49": [("age", ">=", "45"), ("age", "<", "50")], - "50_to_54": [("age", ">=", "50"), ("age", "<", "55")], - "55_to_59": [("age", ">=", "55"), ("age", "<", "60")], - "60_to_64": [("age", ">=", "60"), ("age", "<", "65")], - "65_to_69": [("age", ">=", "65"), ("age", "<", "70")], - "70_to_74": [("age", ">=", "70"), ("age", "<", "75")], - "75_to_79": [("age", ">=", "75"), ("age", "<", "80")], - "80_to_84": [("age", ">=", "80"), ("age", "<", "85")], - "85_plus": [("age", ">=", "85")], - } - - for age_name, constraints in age_brackets.items(): - if age_name not in data.get("age_groups", {}): - continue - - age_stratum = get_or_create_stratum( - session, - name=f"US Population {age_name.replace('_', ' ').title()}", - jurisdiction=Jurisdiction.US, - constraints=constraints, - description=f"US population ages {age_name}", - parent_id=national_stratum.id, - stratum_group_id="age_groups", - ) - - session.add( - Target( - stratum_id=age_stratum.id, - variable="population", - period=year, - value=data["age_groups"][age_name], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.NATIONAL, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - ) - ) - - # State strata - for state_abbrev, state_data in data.get("states", {}).items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - state_name = STATE_NAMES.get(state_abbrev, state_abbrev) - - state_stratum = get_or_create_stratum( - session, - name=f"{state_name} Population", - jurisdiction=Jurisdiction.US, - constraints=[("state_fips", "==", fips)], - description=f"Population of {state_name}", - parent_id=national_stratum.id, - stratum_group_id="state_population", - ) - - session.add( - Target( - stratum_id=state_stratum.id, - variable="population", - period=year, - value=state_data["population"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - ) - ) - - if "households" in state_data: - session.add( - Target( - stratum_id=state_stratum.id, - variable="household_count", - period=year, - value=state_data["households"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def load_congressional_district_targets( - session: Session, year: int, district_data: dict -): - """ - Load congressional district population targets into database. - - Args: - session: Database session - year: Year for the data - district_data: Dict mapping (state_fips, district) to population data - Example: {("06", "01"): {"population": 750000}} - """ - # Get or create national stratum for parent relationship - national_stratum = get_or_create_stratum( - session, - name="US Population", - jurisdiction=Jurisdiction.US, - constraints=[], - description="Total US resident population", - stratum_group_id="population_national", - ) - - for (state_fips, district), data in district_data.items(): - # Create congressional district stratum - cd_stratum = get_or_create_stratum( - session, - name=f"Congressional District {state_fips}-{district}", - jurisdiction=Jurisdiction.US, - constraints=[ - ("state_fips", "==", state_fips), - ("congressional_district", "==", district), - ], - description=f"Population in Congressional District {district} of state {state_fips}", - parent_id=national_stratum.id, - stratum_group_id="congressional_districts", - ) - - # Add population target - session.add( - Target( - stratum_id=cd_stratum.id, - variable="population", - period=year, - value=data["population"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.CONGRESSIONAL_DISTRICT, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - notes="Congressional district boundaries as of current Congress", - ) - ) - - # Add households if available - if "households" in data: - session.add( - Target( - stratum_id=cd_stratum.id, - variable="household_count", - period=year, - value=data["households"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.CONGRESSIONAL_DISTRICT, - source=DataSource.CENSUS_ACS, - source_url=SOURCE_URL, - notes="Congressional district boundaries as of current Congress", - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the Census ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_census_targets(session) - print(f"Loaded Census targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_cps.py b/db/etl_cps.py deleted file mode 100644 index 5ee7e65a..00000000 --- a/db/etl_cps.py +++ /dev/null @@ -1,313 +0,0 @@ -""" -ETL for BLS CPS (Current Population Survey) monthly employment targets. - -Loads monthly employment status, unemployment rates, and labor force participation -data from the Bureau of Labor Statistics Current Population Survey. - -Data sources: -- https://www.bls.gov/cps/tables.htm -- https://www.bls.gov/news.release/empsit.htm -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# CPS monthly employment data -# Source: Current Population Survey Employment Situation Reports -# https://www.bls.gov/cps/tables.htm -# Format: {year: {month: {metrics}}} -# Values in thousands converted to actual counts, rates as percentages - -CPS_MONTHLY_DATA = { - 2025: { - 11: { # November 2025 - "civilian_labor_force": 168_500_000, - "employed": 160_800_000, - "unemployed": 7_700_000, - "not_in_labor_force": 101_200_000, - "unemployment_rate": 4.6, - "labor_force_participation_rate": 62.5, - "employment_population_ratio": 59.6, - }, - 9: { # September 2025 - "civilian_labor_force": 168_600_000, - "employed": 161_000_000, - "unemployed": 7_600_000, - "not_in_labor_force": 101_000_000, - "unemployment_rate": 4.5, - "labor_force_participation_rate": 62.5, - "employment_population_ratio": 59.7, - }, - }, - 2024: { - 12: { # December 2024 - "civilian_labor_force": 168_300_000, - "employed": 161_300_000, - "unemployed": 7_000_000, - "not_in_labor_force": 100_500_000, - "unemployment_rate": 4.2, - "labor_force_participation_rate": 62.6, - "employment_population_ratio": 60.0, - }, - 11: { # November 2024 - "civilian_labor_force": 168_100_000, - "employed": 161_100_000, - "unemployed": 7_000_000, - "not_in_labor_force": 100_400_000, - "unemployment_rate": 4.2, - "labor_force_participation_rate": 62.6, - "employment_population_ratio": 60.0, - }, - 10: { # October 2024 - "civilian_labor_force": 167_900_000, - "employed": 160_800_000, - "unemployed": 7_100_000, - "not_in_labor_force": 100_300_000, - "unemployment_rate": 4.2, - "labor_force_participation_rate": 62.6, - "employment_population_ratio": 60.0, - }, - 9: { # September 2024 - "civilian_labor_force": 167_700_000, - "employed": 160_500_000, - "unemployed": 7_200_000, - "not_in_labor_force": 100_200_000, - "unemployment_rate": 4.3, - "labor_force_participation_rate": 62.6, - "employment_population_ratio": 59.9, - }, - }, - 2023: { - 12: { # December 2023 - "civilian_labor_force": 167_300_000, - "employed": 161_100_000, - "unemployed": 6_200_000, - "not_in_labor_force": 99_800_000, - "unemployment_rate": 3.7, - "labor_force_participation_rate": 62.6, - "employment_population_ratio": 60.3, - }, - 6: { # June 2023 - "civilian_labor_force": 166_700_000, - "employed": 160_400_000, - "unemployed": 6_300_000, - "not_in_labor_force": 99_500_000, - "unemployment_rate": 3.8, - "labor_force_participation_rate": 62.6, - "employment_population_ratio": 60.2, - }, - 1: { # January 2023 - "civilian_labor_force": 165_900_000, - "employed": 159_500_000, - "unemployed": 6_400_000, - "not_in_labor_force": 99_300_000, - "unemployment_rate": 3.9, - "labor_force_participation_rate": 62.4, - "employment_population_ratio": 60.0, - }, - }, -} - -SOURCE_URL = "https://www.bls.gov/cps/tables.htm" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_cps_targets(session: Session, years: list[int] | None = None, months: list[int] | None = None): - """ - Load CPS monthly employment targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - months: Months to load (default: all available for selected years) - """ - if years is None: - years = list(CPS_MONTHLY_DATA.keys()) - - for year in years: - if year not in CPS_MONTHLY_DATA: - continue - - year_data = CPS_MONTHLY_DATA[year] - months_to_load = months if months is not None else list(year_data.keys()) - - for month in months_to_load: - if month not in year_data: - continue - - data = year_data[month] - - # Period as YYYYMM format (e.g., 202311 for November 2023) - period = year * 100 + month - - # Create national labor force stratum for this month - labor_force_stratum = get_or_create_stratum( - session, - name="US Labor Force", - jurisdiction=Jurisdiction.US, - constraints=[("in_labor_force", "==", "1")], - description="US civilian labor force (16+)", - stratum_group_id="cps_monthly", - ) - - # Employment count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="employed", - period=period, - value=data["employed"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Unemployment count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="unemployed", - period=period, - value=data["unemployed"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Unemployment rate - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="unemployment_rate", - period=period, - value=data["unemployment_rate"], - target_type=TargetType.RATE, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Labor force participation rate - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="labor_force_participation_rate", - period=period, - value=data["labor_force_participation_rate"], - target_type=TargetType.RATE, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Employment-population ratio - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="employment_population_ratio", - period=period, - value=data["employment_population_ratio"], - target_type=TargetType.RATE, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Total labor force count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="labor_force_count", - period=period, - value=data["civilian_labor_force"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - # Not in labor force count - session.add( - Target( - stratum_id=labor_force_stratum.id, - variable="not_in_labor_force", - period=period, - value=data["not_in_labor_force"], - target_type=TargetType.COUNT, - source=DataSource.BLS, - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the CPS ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_cps_targets(session) - print(f"Loaded CPS monthly targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_cps_raw.py b/db/etl_cps_raw.py deleted file mode 100644 index 1a964514..00000000 --- a/db/etl_cps_raw.py +++ /dev/null @@ -1,378 +0,0 @@ -""" -ETL for loading raw CPS ASEC microdata to Supabase. - -Loads person, household, and family tables from locally cached parquet files -to the microplex schema in Supabase. - -Table naming: microplex.us_census_cps_asec_{year}_{table_type} - -Usage: - python -m db.etl_cps_raw --year 2024 - python -m db.etl_cps_raw --year 2024 --dry-run - python -m db.etl_cps_raw --year 2024 --method csv # Export to CSV for dashboard import -""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any, Dict, List, Optional - -import pandas as pd - -from .supabase_client import get_supabase_client, get_table_name, register_dataset - - -# Path to local raw cache -CACHE_DIR = Path(__file__).parent.parent / "micro" / "us" / "census" / "raw_cache" - - -def get_raw_cache_dir(year: int) -> Path: - """Get the directory for cached raw CPS data.""" - return CACHE_DIR / f"census_cps_{year}" - - -def get_cps_table_names(year: int) -> Dict[str, str]: - """Get Supabase table names for CPS ASEC tables.""" - return { - "person": get_table_name("us", "census", "cps_asec", year, "person"), - "household": get_table_name("us", "census", "cps_asec", year, "household"), - "family": get_table_name("us", "census", "cps_asec", year, "family"), - } - - -# Key columns to extract as typed columns (rest goes in raw_data JSONB) -PERSON_KEY_COLUMNS = { - "PH_SEQ": "ph_seq", - "PPPOS": "pppos", - "A_AGE": "a_age", - "A_SEX": "a_sex", - "PRDTRACE": "prdtrace", - "PEHSPNON": "pehspnon", - "PRCITSHP": "prcitshp", - "GESTFIPS": "gestfips", - "GTCO": "gtco", - "PTOTVAL": "ptotval", - "PEARNVAL": "pearnval", - "WSAL_VAL": "wsal_val", - "SEMP_VAL": "semp_val", - "MARSUPWT": "marsupwt", -} - -HOUSEHOLD_KEY_COLUMNS = { - "H_SEQ": "h_seq", - "H_NUMPER": "h_numper", - "HH5TO18": "hh5to18", - "HUNDER18": "hunder18", - "GESTFIPS": "gestfips", - "GTCO": "gtco", - "HTOTVAL": "htotval", - "HPROP_VAL": "hprop_val", - "HSUP_WGT": "hsup_wgt", -} - -FAMILY_KEY_COLUMNS = { - "FH_SEQ": "fh_seq", - "FFPOS": "ffpos", - "FKIND": "fkind", - "F_MV_FS": "f_mv_fs", - "FHIP_VAL": "fhip_val", - "FTOTVAL": "ftotval", -} - - -def _prepare_records( - df: pd.DataFrame, - key_columns: Dict[str, str], - include_raw_data: bool = False, -) -> List[Dict[str, Any]]: - """ - Prepare records for Supabase insert. - - Extracts key columns as typed fields. Optionally stores full row in raw_data JSONB. - """ - records = [] - - for _, row in df.iterrows(): - record = {} - - # Extract key columns - for src_col, dest_col in key_columns.items(): - if src_col in df.columns: - val = row[src_col] - # Convert numpy types to Python types - if pd.isna(val): - record[dest_col] = None - elif hasattr(val, 'item'): - record[dest_col] = val.item() - else: - record[dest_col] = val - - # Optionally store full row as raw_data JSONB (heavy, use sparingly) - if include_raw_data: - raw_data = {} - for col in df.columns: - val = row[col] - if pd.isna(val): - raw_data[col] = None - elif hasattr(val, 'item'): - raw_data[col] = val.item() - else: - raw_data[col] = val - record["raw_data"] = raw_data - - records.append(record) - - return records - - -def prepare_person_records(df: pd.DataFrame) -> List[Dict[str, Any]]: - """Prepare person records for Supabase insert.""" - return _prepare_records(df, PERSON_KEY_COLUMNS) - - -def prepare_household_records(df: pd.DataFrame) -> List[Dict[str, Any]]: - """Prepare household records for Supabase insert.""" - return _prepare_records(df, HOUSEHOLD_KEY_COLUMNS) - - -def prepare_family_records(df: pd.DataFrame) -> List[Dict[str, Any]]: - """Prepare family records for Supabase insert.""" - return _prepare_records(df, FAMILY_KEY_COLUMNS) - - -def load_cps_to_supabase( - year: int, - dry_run: bool = False, - chunk_size: int = 200, - limit: Optional[int] = None, - truncate: bool = False, - skip: int = 0, -) -> Dict[str, Any]: - """ - Load CPS ASEC data from local cache to Supabase. - - Args: - year: Data year (e.g., 2024) - dry_run: If True, just return stats without inserting - chunk_size: Records per batch insert (default 200 for stability) - limit: Max records per table (for testing) - truncate: If True, delete existing data before loading - skip: Skip first N records (for resuming interrupted load) - - Returns: - Dict with counts and status - """ - cache_dir = get_raw_cache_dir(year) - person_path = cache_dir / "person.parquet" - household_path = cache_dir / "household.parquet" - family_path = cache_dir / "family.parquet" - - if not person_path.exists(): - raise FileNotFoundError( - f"CPS {year} not found at {cache_dir}. " - f"Run: python micro/us/census/download_cps.py --year {year}" - ) - - # Load data - print(f"Loading CPS ASEC {year} from {cache_dir}...") - person_df = pd.read_parquet(person_path) - household_df = pd.read_parquet(household_path) if household_path.exists() else pd.DataFrame() - family_df = pd.read_parquet(family_path) if family_path.exists() else pd.DataFrame() - - # Apply skip and limit - if skip > 0: - person_df = person_df.iloc[skip:] - # For household/family, skip proportionally (approximate) - hh_skip = int(skip * len(household_df) / len(person_df)) if len(person_df) > 0 else 0 - fam_skip = int(skip * len(family_df) / len(person_df)) if len(person_df) > 0 else 0 - household_df = household_df.iloc[hh_skip:] - family_df = family_df.iloc[fam_skip:] - - if limit: - person_df = person_df.head(limit) - household_df = household_df.head(limit) - family_df = family_df.head(limit) - - result = { - "year": year, - "person_count": len(person_df), - "household_count": len(household_df), - "family_count": len(family_df), - "person_columns": list(person_df.columns), - "dry_run": dry_run, - } - - if dry_run: - print("DRY RUN - would load:") - print(f" {result['person_count']:,} person records") - print(f" {result['household_count']:,} household records") - print(f" {result['family_count']:,} family records") - return result - - # Get Supabase client and table names - client = get_supabase_client() - table_names = get_cps_table_names(year) - - # Truncate existing data if requested - if truncate: - print("Truncating existing data...") - for table_type, table_name in table_names.items(): - try: - client.schema("microplex").table(table_name).delete().neq("id", 0).execute() - print(f" Truncated {table_name}") - except Exception as e: - print(f" Warning: Could not truncate {table_name}: {e}") - - # Load person records - print(f"Loading {len(person_df):,} person records...") - person_records = prepare_person_records(person_df) - _insert_batch(client, table_names["person"], person_records, chunk_size) - - # Load household records - if len(household_df) > 0: - print(f"Loading {len(household_df):,} household records...") - household_records = prepare_household_records(household_df) - _insert_batch(client, table_names["household"], household_records, chunk_size) - - # Load family records - if len(family_df) > 0: - print(f"Loading {len(family_df):,} family records...") - family_records = prepare_family_records(family_df) - _insert_batch(client, table_names["family"], family_records, chunk_size) - - # Update dataset registry - register_dataset( - jurisdiction="us", - institution="census", - dataset="cps_asec", - year=year, - table_type="person", - row_count=len(person_df), - columns=[{"name": c, "dtype": str(person_df[c].dtype)} for c in person_df.columns[:20]], - source_url=f"https://www.census.gov/data/datasets/{year+1}/demo/cps/cps-asec-{year+1}.html", - ) - - result["status"] = "completed" - print(f"Loaded CPS ASEC {year} to Supabase") - return result - - -def _insert_batch( - client, - table_name: str, - records: List[Dict[str, Any]], - chunk_size: int, -) -> int: - """Insert records in batches.""" - total = 0 - - for i in range(0, len(records), chunk_size): - chunk = records[i:i + chunk_size] - client.schema("microplex").table(table_name).insert(chunk).execute() - total += len(chunk) - if (i + chunk_size) % 5000 == 0: - print(f" Inserted {total:,} / {len(records):,} records") - - return total - - -def export_cps_to_csv( - year: int, - output_dir: Optional[Path] = None, -) -> Dict[str, Path]: - """ - Export CPS data to CSV files for bulk import to Supabase. - - Much faster than record-by-record insert. Import via: - - Supabase Dashboard: Table Editor > Import CSV - - psql: \\copy microplex.us_census_cps_asec_2024_person FROM 'person.csv' CSV HEADER - - Args: - year: Data year - output_dir: Directory for CSV files (default: current dir) - - Returns: - Dict of {table_type: csv_path} - """ - cache_dir = get_raw_cache_dir(year) - output_dir = output_dir or Path(".") - - paths = {} - table_names = get_cps_table_names(year) - - # Person - person_path = cache_dir / "person.parquet" - if person_path.exists(): - print("Exporting person records...") - df = pd.read_parquet(person_path) - records = prepare_person_records(df) - out_df = pd.DataFrame(records) - csv_path = output_dir / f"cps_asec_{year}_person.csv" - out_df.to_csv(csv_path, index=False) - paths["person"] = csv_path - print(f" Exported {len(out_df):,} records to {csv_path}") - - # Household - hh_path = cache_dir / "household.parquet" - if hh_path.exists(): - print("Exporting household records...") - df = pd.read_parquet(hh_path) - records = prepare_household_records(df) - out_df = pd.DataFrame(records) - csv_path = output_dir / f"cps_asec_{year}_household.csv" - out_df.to_csv(csv_path, index=False) - paths["household"] = csv_path - print(f" Exported {len(out_df):,} records to {csv_path}") - - # Family - fam_path = cache_dir / "family.parquet" - if fam_path.exists(): - print("Exporting family records...") - df = pd.read_parquet(fam_path) - records = prepare_family_records(df) - out_df = pd.DataFrame(records) - csv_path = output_dir / f"cps_asec_{year}_family.csv" - out_df.to_csv(csv_path, index=False) - paths["family"] = csv_path - print(f" Exported {len(out_df):,} records to {csv_path}") - - print("\nTo import to Supabase:") - print(" 1. Go to Supabase Dashboard > Table Editor") - print(f" 2. Select table (e.g., {table_names['person']})") - print(" 3. Click 'Insert' > 'Import data from CSV'") - print(" 4. Upload the CSV file") - - return paths - - -def main(): - import argparse - - parser = argparse.ArgumentParser(description="Load raw CPS ASEC to Supabase") - parser.add_argument("--year", type=int, default=2024, help="Data year") - parser.add_argument("--dry-run", action="store_true", help="Show stats without inserting") - parser.add_argument("--limit", type=int, help="Limit records per table (for testing)") - parser.add_argument("--chunk-size", type=int, default=200, help="Batch insert size") - parser.add_argument("--truncate", action="store_true", help="Delete existing data before loading") - parser.add_argument("--skip", type=int, default=0, help="Skip first N records (for resuming)") - parser.add_argument("--method", type=str, choices=["api", "csv"], default="api", - help="Method: 'api' for record-by-record, 'csv' for export to CSV") - parser.add_argument("--output-dir", type=str, help="Output directory for CSV export") - args = parser.parse_args() - - if args.method == "csv": - output_dir = Path(args.output_dir) if args.output_dir else None - export_cps_to_csv(args.year, output_dir) - else: - load_cps_to_supabase( - year=args.year, - dry_run=args.dry_run, - limit=args.limit, - chunk_size=args.chunk_size, - truncate=args.truncate, - skip=args.skip, - ) - - -if __name__ == "__main__": - main() diff --git a/db/etl_hmrc.py b/db/etl_hmrc.py deleted file mode 100644 index cc81f9a7..00000000 --- a/db/etl_hmrc.py +++ /dev/null @@ -1,351 +0,0 @@ -""" -ETL for UK HMRC (HM Revenue & Customs) targets. - -Loads data from HMRC statistics into the targets database. -Data sources: -- https://www.gov.uk/government/statistics/income-tax-liabilities-statistics -- https://www.gov.uk/government/collections/national-insurance-contributions-statistics -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# HMRC data by tax year -# Source: HMRC statistics publications -# Values in GBP (pounds sterling) -HMRC_DATA = { - 2022: { - # Tax revenue (billions GBP converted to GBP) - "income_tax": 248_800_000_000, # £248.8bn - "national_insurance": 176_500_000_000, # £176.5bn - "capital_gains_tax": 18_200_000_000, # £18.2bn - "inheritance_tax": 7_100_000_000, # £7.1bn - - # Taxpayer counts - "taxpayers": 34_100_000, - "higher_rate_taxpayers": 6_100_000, - "additional_rate_taxpayers": 629_000, - - # Total income - "total_income": 1_534_000_000_000, # £1.534tn - - # Benefits (DWP data, included here for convenience) - "benefits": { - "universal_credit": { - "recipients": 6_200_000, - "expenditure": 52_300_000_000, # £52.3bn - }, - "child_benefit": { - "recipients": 7_100_000, # families - "expenditure": 12_100_000_000, - }, - "state_pension": { - "recipients": 12_700_000, - "expenditure": 124_300_000_000, - }, - "housing_benefit": { - "recipients": 2_800_000, - "expenditure": 18_400_000_000, - }, - "pension_credit": { - "recipients": 1_400_000, - "expenditure": 5_600_000_000, - }, - }, - }, - 2021: { - "income_tax": 224_700_000_000, - "national_insurance": 157_800_000_000, - "capital_gains_tax": 14_900_000_000, - "inheritance_tax": 6_100_000_000, - - "taxpayers": 32_400_000, - "higher_rate_taxpayers": 5_500_000, - "additional_rate_taxpayers": 538_000, - - "total_income": 1_423_000_000_000, - - "benefits": { - "universal_credit": { - "recipients": 5_800_000, - "expenditure": 44_700_000_000, - }, - "child_benefit": { - "recipients": 7_200_000, - "expenditure": 11_900_000_000, - }, - "state_pension": { - "recipients": 12_500_000, - "expenditure": 105_500_000_000, - }, - "housing_benefit": { - "recipients": 3_100_000, - "expenditure": 20_100_000_000, - }, - "pension_credit": { - "recipients": 1_500_000, - "expenditure": 5_400_000_000, - }, - }, - }, -} - -SOURCE_URL_TAX = "https://www.gov.uk/government/statistics/income-tax-liabilities-statistics" -SOURCE_URL_DWP = "https://www.gov.uk/government/collections/dwp-benefits-statistics" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_hmrc_targets(session: Session, years: list[int] | None = None): - """ - Load HMRC targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(HMRC_DATA.keys()) - - for year in years: - if year not in HMRC_DATA: - continue - - data = HMRC_DATA[year] - - # Create national taxpayer stratum - national_stratum = get_or_create_stratum( - session, - name="UK All Taxpayers", - jurisdiction=Jurisdiction.UK, - constraints=[("is_taxpayer", "==", "1")], # Taxpayers only, not whole population - description="All UK taxpayers", - stratum_group_id="uk_national", - ) - - # Tax revenue targets - session.add( - Target( - stratum_id=national_stratum.id, - variable="income_tax", - period=year, - value=data["income_tax"], - target_type=TargetType.AMOUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - session.add( - Target( - stratum_id=national_stratum.id, - variable="national_insurance", - period=year, - value=data["national_insurance"], - target_type=TargetType.AMOUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - session.add( - Target( - stratum_id=national_stratum.id, - variable="capital_gains_tax", - period=year, - value=data["capital_gains_tax"], - target_type=TargetType.AMOUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - session.add( - Target( - stratum_id=national_stratum.id, - variable="inheritance_tax", - period=year, - value=data["inheritance_tax"], - target_type=TargetType.AMOUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - # Taxpayer counts - session.add( - Target( - stratum_id=national_stratum.id, - variable="taxpayer_count", - period=year, - value=data["taxpayers"], - target_type=TargetType.COUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - session.add( - Target( - stratum_id=national_stratum.id, - variable="total_income", - period=year, - value=data["total_income"], - target_type=TargetType.AMOUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - # Higher rate taxpayers stratum - higher_rate_stratum = get_or_create_stratum( - session, - name="UK Higher Rate Taxpayers", - jurisdiction=Jurisdiction.UK, - constraints=[("income_tax_band", "==", "higher")], - description="UK taxpayers in higher rate band", - parent_id=national_stratum.id, - stratum_group_id="uk_tax_bands", - ) - - session.add( - Target( - stratum_id=higher_rate_stratum.id, - variable="taxpayer_count", - period=year, - value=data["higher_rate_taxpayers"], - target_type=TargetType.COUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - # Additional rate taxpayers stratum - additional_rate_stratum = get_or_create_stratum( - session, - name="UK Additional Rate Taxpayers", - jurisdiction=Jurisdiction.UK, - constraints=[("income_tax_band", "==", "additional")], - description="UK taxpayers in additional rate band", - parent_id=national_stratum.id, - stratum_group_id="uk_tax_bands", - ) - - session.add( - Target( - stratum_id=additional_rate_stratum.id, - variable="taxpayer_count", - period=year, - value=data["additional_rate_taxpayers"], - target_type=TargetType.COUNT, - source=DataSource.HMRC, - source_url=SOURCE_URL_TAX, - ) - ) - - # Benefit targets - for benefit_name, benefit_data in data.get("benefits", {}).items(): - benefit_display = benefit_name.replace("_", " ").title() - - benefit_stratum = get_or_create_stratum( - session, - name=f"UK {benefit_display} Recipients", - jurisdiction=Jurisdiction.UK, - constraints=[(benefit_name, "==", "1")], - description=f"UK {benefit_display} recipients", - stratum_group_id="uk_benefits", - ) - - session.add( - Target( - stratum_id=benefit_stratum.id, - variable=f"{benefit_name}_recipients", - period=year, - value=benefit_data["recipients"], - target_type=TargetType.COUNT, - source=DataSource.DWP, - source_url=SOURCE_URL_DWP, - ) - ) - - session.add( - Target( - stratum_id=benefit_stratum.id, - variable=f"{benefit_name}_expenditure", - period=year, - value=benefit_data["expenditure"], - target_type=TargetType.AMOUNT, - source=DataSource.DWP, - source_url=SOURCE_URL_DWP, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the HMRC ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_hmrc_targets(session) - print(f"Loaded HMRC targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_medicaid.py b/db/etl_medicaid.py deleted file mode 100644 index 854912f4..00000000 --- a/db/etl_medicaid.py +++ /dev/null @@ -1,494 +0,0 @@ -""" -ETL for CMS Medicaid enrollment targets. - -Loads Medicaid and CHIP enrollment data from CMS into the targets database. -Data sources: -- https://www.medicaid.gov/medicaid/national-medicaid-chip-program-information/medicaid-chip-enrollment-data -- https://www.kff.org/affordable-care-act/state-indicator/total-monthly-medicaid-and-chip-enrollment/ -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes -STATE_FIPS = { - "AL": "01", "AK": "02", "AZ": "04", "AR": "05", "CA": "06", - "CO": "08", "CT": "09", "DE": "10", "DC": "11", "FL": "12", - "GA": "13", "HI": "15", "ID": "16", "IL": "17", "IN": "18", - "IA": "19", "KS": "20", "KY": "21", "LA": "22", "ME": "23", - "MD": "24", "MA": "25", "MI": "26", "MN": "27", "MS": "28", - "MO": "29", "MT": "30", "NE": "31", "NV": "32", "NH": "33", - "NJ": "34", "NM": "35", "NY": "36", "NC": "37", "ND": "38", - "OH": "39", "OK": "40", "OR": "41", "PA": "42", "RI": "44", - "SC": "45", "SD": "46", "TN": "47", "TX": "48", "UT": "49", - "VT": "50", "VA": "51", "WA": "53", "WV": "54", "WI": "55", - "WY": "56", "PR": "72", "VI": "78", "GU": "66", -} - -# Medicaid enrollment data by year -# Sources: -# - CMS Medicaid & CHIP Enrollment Data: https://www.medicaid.gov/medicaid/national-medicaid-chip-program-information/medicaid-chip-enrollment-data -# - KFF State Health Facts: https://www.kff.org/affordable-care-act/state-indicator/total-monthly-medicaid-and-chip-enrollment/ -# Values are as of August of each year unless otherwise noted -MEDICAID_DATA = { - 2025: { - "national": { - # August 2025 data from CMS - "total_enrollment": 77_290_050, # Total Medicaid + CHIP - "medicaid_enrollment": 70_067_213, - "chip_enrollment": 7_222_837, - # By age category (from KFF) - "children": 36_100_000, # 48% of total - "adults": 39_400_000, # 52% of total (includes aged/disabled) - }, - # State enrollment from KFF/CMS August 2025 - "states": { - "CA": {"total": 13_226_087}, - "NY": {"total": 6_522_114}, - "TX": {"total": 4_161_101}, - "FL": {"total": 3_674_391}, - "PA": {"total": 3_018_973}, - "IL": {"total": 2_987_000}, - "OH": {"total": 2_654_000}, - "MI": {"total": 2_456_000}, - "NC": {"total": 2_387_000}, - "GA": {"total": 2_123_000}, - "NJ": {"total": 1_987_000}, - "AZ": {"total": 1_876_000}, - "WA": {"total": 1_765_000}, - "MA": {"total": 1_654_000}, - "VA": {"total": 1_543_000}, - }, - }, - 2024: { - "national": { - # Annual average from CMS - "total_enrollment": 81_500_000, - "medicaid_enrollment": 73_800_000, - "chip_enrollment": 7_700_000, - "children": 38_500_000, - "adults": 43_000_000, - }, - "states": { - "CA": {"total": 13_959_148}, - "NY": {"total": 7_123_000}, - "TX": {"total": 4_567_000}, - "FL": {"total": 4_234_000}, - "PA": {"total": 3_345_000}, - "IL": {"total": 3_234_000}, - "OH": {"total": 2_876_000}, - "MI": {"total": 2_654_000}, - "NC": {"total": 2_543_000}, - "GA": {"total": 2_345_000}, - "NJ": {"total": 2_123_000}, - "AZ": {"total": 2_098_000}, - "WA": {"total": 1_987_000}, - "MA": {"total": 1_876_000}, - "VA": {"total": 1_765_000}, - }, - }, - 2023: { - "national": { - # Peak enrollment year (before unwinding completed) - "total_enrollment": 94_000_000, # Record high in March 2023 - "medicaid_enrollment": 85_500_000, - "chip_enrollment": 8_500_000, - "children": 43_000_000, - "adults": 51_000_000, - # By eligibility category from KFF (FY2023 data) - "aged": 10_800_000, # 10% - "disabled": 12_300_000, # 11% - "non_disabled_adults": 21_800_000, # 20% - "aca_expansion_adults": 27_400_000, # 25% - }, - "states": { - "CA": {"total": 15_234_000}, - "NY": {"total": 7_876_000}, - "TX": {"total": 5_234_000}, - "FL": {"total": 4_987_000}, - "PA": {"total": 3_654_000}, - "IL": {"total": 3_543_000}, - "OH": {"total": 3_123_000}, - "MI": {"total": 2_987_000}, - "NC": {"total": 2_765_000}, - "GA": {"total": 2_654_000}, - "NJ": {"total": 2_345_000}, - "AZ": {"total": 2_298_000}, - "WA": {"total": 2_123_000}, - "MA": {"total": 2_098_000}, - "VA": {"total": 1_987_000}, - }, - }, - 2022: { - "national": { - # During continuous enrollment period - "total_enrollment": 91_400_000, - "medicaid_enrollment": 83_100_000, - "chip_enrollment": 8_300_000, - "children": 42_000_000, - "adults": 49_400_000, - }, - "states": { - "CA": {"total": 14_876_000}, - "NY": {"total": 7_654_000}, - "TX": {"total": 5_098_000}, - "FL": {"total": 4_765_000}, - "PA": {"total": 3_543_000}, - "IL": {"total": 3_432_000}, - "OH": {"total": 3_012_000}, - "MI": {"total": 2_876_000}, - "NC": {"total": 2_654_000}, - "GA": {"total": 2_543_000}, - }, - }, - 2021: { - "national": { - # COVID continuous enrollment period - "total_enrollment": 85_000_000, - "medicaid_enrollment": 77_200_000, - "chip_enrollment": 7_800_000, - "children": 39_500_000, - "adults": 45_500_000, - }, - "states": { - "CA": {"total": 14_234_000}, - "NY": {"total": 7_234_000}, - "TX": {"total": 4_876_000}, - "FL": {"total": 4_432_000}, - "PA": {"total": 3_345_000}, - "IL": {"total": 3_234_000}, - "OH": {"total": 2_876_000}, - "MI": {"total": 2_654_000}, - "NC": {"total": 2_432_000}, - "GA": {"total": 2_345_000}, - }, - }, -} - -# Source URLs -SOURCE_URL = "https://www.medicaid.gov/medicaid/national-medicaid-chip-program-information/medicaid-chip-enrollment-data" -KFF_SOURCE_URL = "https://www.kff.org/affordable-care-act/state-indicator/total-monthly-medicaid-and-chip-enrollment/" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_medicaid_targets(session: Session, years: list[int] | None = None): - """ - Load Medicaid enrollment targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(MEDICAID_DATA.keys()) - - for year in years: - if year not in MEDICAID_DATA: - continue - - data = MEDICAID_DATA[year] - national_data = data["national"] - - # Create national Medicaid stratum - national_stratum = get_or_create_stratum( - session, - name="US Medicaid Enrollees", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("medicaid", "==", "1")], - description="All Medicaid and CHIP enrollees in the US", - stratum_group_id="medicaid_national", - ) - - # Total enrollment target - session.add( - Target( - stratum_id=national_stratum.id, - variable="medicaid_total_enrollment", - period=year, - value=national_data["total_enrollment"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="Medicaid & CHIP Enrollment Data", - source_url=SOURCE_URL, - ) - ) - - # Medicaid-only enrollment - session.add( - Target( - stratum_id=national_stratum.id, - variable="medicaid_enrollment", - period=year, - value=national_data["medicaid_enrollment"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="Medicaid & CHIP Enrollment Data", - source_url=SOURCE_URL, - ) - ) - - # CHIP enrollment - chip_stratum = get_or_create_stratum( - session, - name="US CHIP Enrollees", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("chip", "==", "1")], - description="Children's Health Insurance Program enrollees", - parent_id=national_stratum.id, - stratum_group_id="medicaid_chip", - ) - - session.add( - Target( - stratum_id=chip_stratum.id, - variable="chip_enrollment", - period=year, - value=national_data["chip_enrollment"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="Medicaid & CHIP Enrollment Data", - source_url=SOURCE_URL, - ) - ) - - # Child enrollment stratum - if "children" in national_data: - child_stratum = get_or_create_stratum( - session, - name="US Medicaid Children", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("medicaid", "==", "1"), - ("is_child", "==", "1"), - ], - description="Children enrolled in Medicaid or CHIP", - parent_id=national_stratum.id, - stratum_group_id="medicaid_demographics", - ) - - session.add( - Target( - stratum_id=child_stratum.id, - variable="medicaid_child_enrollment", - period=year, - value=national_data["children"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="Medicaid & CHIP Enrollment Data", - source_url=KFF_SOURCE_URL, - ) - ) - - # Adult enrollment stratum - if "adults" in national_data: - adult_stratum = get_or_create_stratum( - session, - name="US Medicaid Adults", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("medicaid", "==", "1"), - ("is_adult", "==", "1"), - ], - description="Adults enrolled in Medicaid (includes aged and disabled)", - parent_id=national_stratum.id, - stratum_group_id="medicaid_demographics", - ) - - session.add( - Target( - stratum_id=adult_stratum.id, - variable="medicaid_adult_enrollment", - period=year, - value=national_data["adults"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="Medicaid & CHIP Enrollment Data", - source_url=KFF_SOURCE_URL, - ) - ) - - # Eligibility category strata (if available for the year) - if "aged" in national_data: - aged_stratum = get_or_create_stratum( - session, - name="US Medicaid Aged", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("medicaid", "==", "1"), - ("age", ">=", "65"), - ], - description="Aged (65+) Medicaid enrollees", - parent_id=national_stratum.id, - stratum_group_id="medicaid_eligibility", - ) - - session.add( - Target( - stratum_id=aged_stratum.id, - variable="medicaid_aged_enrollment", - period=year, - value=national_data["aged"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="KFF Medicaid Enrollees by Enrollment Group", - source_url=KFF_SOURCE_URL, - ) - ) - - if "disabled" in national_data: - disabled_stratum = get_or_create_stratum( - session, - name="US Medicaid Disabled", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("medicaid", "==", "1"), - ("is_disabled", "==", "1"), - ], - description="Disabled Medicaid enrollees (non-aged)", - parent_id=national_stratum.id, - stratum_group_id="medicaid_eligibility", - ) - - session.add( - Target( - stratum_id=disabled_stratum.id, - variable="medicaid_disabled_enrollment", - period=year, - value=national_data["disabled"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="KFF Medicaid Enrollees by Enrollment Group", - source_url=KFF_SOURCE_URL, - ) - ) - - if "aca_expansion_adults" in national_data: - expansion_stratum = get_or_create_stratum( - session, - name="US Medicaid ACA Expansion Adults", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("medicaid", "==", "1"), - ("is_aca_expansion", "==", "1"), - ], - description="Adults enrolled through ACA Medicaid expansion", - parent_id=national_stratum.id, - stratum_group_id="medicaid_eligibility", - ) - - session.add( - Target( - stratum_id=expansion_stratum.id, - variable="medicaid_aca_expansion_enrollment", - period=year, - value=national_data["aca_expansion_adults"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="KFF Medicaid Enrollees by Enrollment Group", - source_url=KFF_SOURCE_URL, - ) - ) - - # State-level enrollment targets - for state_abbrev, state_data in data.get("states", {}).items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} Medicaid Enrollees", - jurisdiction=Jurisdiction.US, - constraints=[ - ("medicaid", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"Medicaid and CHIP enrollees in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="medicaid_states", - ) - - session.add( - Target( - stratum_id=state_stratum.id, - variable="medicaid_total_enrollment", - period=year, - value=state_data["total"], - target_type=TargetType.COUNT, - source=DataSource.CMS_MEDICAID, - source_table="Medicaid State Enrollment", - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the Medicaid enrollment ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_medicaid_targets(session) - print(f"Loaded Medicaid enrollment targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_obr.py b/db/etl_obr.py deleted file mode 100644 index 2ed11875..00000000 --- a/db/etl_obr.py +++ /dev/null @@ -1,313 +0,0 @@ -""" -ETL for OBR (Office for Budget Responsibility) projections. - -Loads UK fiscal and economic projections from OBR's Economic and Fiscal Outlook. -Data source: https://obr.uk/efo/economic-and-fiscal-outlook-march-2024/ -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# OBR projections data (5-year forecast, 2024-2029) -# Source: OBR Economic and Fiscal Outlook, March 2024 -# https://obr.uk/efo/economic-and-fiscal-outlook-march-2024/ -# GDP and budget figures in billions of GBP -# Rates as percentages -OBR_DATA = { - # Fiscal Year 2024-25 - 2024: { - # GDP and public finances (billions GBP) - "gdp": 2_791_000_000_000, # £2.791 trillion - "total_receipts": 1_139_000_000_000, - "total_managed_expenditure": 1_226_000_000_000, - "public_sector_net_borrowing": 87_000_000_000, - "public_sector_net_debt": 2_654_000_000_000, - # Economic indicators (percentages) - "real_gdp_growth": 0.8, - "unemployment_rate": 4.2, - "cpi_inflation": 2.8, - "rpi_inflation": 3.9, - "bank_rate": 4.8, - # Employment (thousands) - "employment": 33_100_000, - }, - 2025: { - "gdp": 2_929_000_000_000, - "total_receipts": 1_207_000_000_000, - "total_managed_expenditure": 1_296_000_000_000, - "public_sector_net_borrowing": 89_000_000_000, - "public_sector_net_debt": 2_783_000_000_000, - "real_gdp_growth": 1.9, - "unemployment_rate": 4.3, - "cpi_inflation": 2.2, - "rpi_inflation": 3.2, - "bank_rate": 4.1, - "employment": 33_200_000, - }, - 2026: { - "gdp": 3_066_000_000_000, - "total_receipts": 1_273_000_000_000, - "total_managed_expenditure": 1_349_000_000_000, - "public_sector_net_borrowing": 76_000_000_000, - "public_sector_net_debt": 2_893_000_000_000, - "real_gdp_growth": 2.0, - "unemployment_rate": 4.1, - "cpi_inflation": 1.8, - "rpi_inflation": 2.8, - "bank_rate": 3.6, - "employment": 33_400_000, - }, - 2027: { - "gdp": 3_199_000_000_000, - "total_receipts": 1_337_000_000_000, - "total_managed_expenditure": 1_399_000_000_000, - "public_sector_net_borrowing": 62_000_000_000, - "public_sector_net_debt": 2_985_000_000_000, - "real_gdp_growth": 1.8, - "unemployment_rate": 4.0, - "cpi_inflation": 1.8, - "rpi_inflation": 2.8, - "bank_rate": 3.4, - "employment": 33_500_000, - }, - 2028: { - "gdp": 3_331_000_000_000, - "total_receipts": 1_399_000_000_000, - "total_managed_expenditure": 1_456_000_000_000, - "public_sector_net_borrowing": 57_000_000_000, - "public_sector_net_debt": 3_070_000_000_000, - "real_gdp_growth": 1.7, - "unemployment_rate": 4.0, - "cpi_inflation": 2.0, - "rpi_inflation": 3.0, - "bank_rate": 3.3, - "employment": 33_600_000, - }, - 2029: { - "gdp": 3_464_000_000_000, - "total_receipts": 1_459_000_000_000, - "total_managed_expenditure": 1_512_000_000_000, - "public_sector_net_borrowing": 53_000_000_000, - "public_sector_net_debt": 3_148_000_000_000, - "real_gdp_growth": 1.7, - "unemployment_rate": 4.0, - "cpi_inflation": 2.0, - "rpi_inflation": 3.0, - "bank_rate": 3.3, - "employment": 33_700_000, - }, -} - -SOURCE_URL = "https://obr.uk/efo/economic-and-fiscal-outlook-march-2024/" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_obr_targets(session: Session, years: list[int] | None = None): - """ - Load OBR projections into database. - - Args: - session: Database session - years: Years to load (default: all available 2024-2029) - """ - if years is None: - years = list(OBR_DATA.keys()) - - for year in years: - if year not in OBR_DATA: - continue - - data = OBR_DATA[year] - - # UK public finances stratum - budget_stratum = get_or_create_stratum( - session, - name="UK Public Finances", - jurisdiction=Jurisdiction.UK, - constraints=[("sector", "==", "public")], # Government sector - description="UK public sector finances", - stratum_group_id="obr_budget", - ) - - # Budget targets - budget_vars = [ - ("gdp", TargetType.AMOUNT), - ("total_receipts", TargetType.AMOUNT), - ("total_managed_expenditure", TargetType.AMOUNT), - ("public_sector_net_borrowing", TargetType.AMOUNT), - ("public_sector_net_debt", TargetType.AMOUNT), - ] - - for var_name, var_type in budget_vars: - if var_name not in data: - continue - - # Check for existing target - existing = session.exec( - select(Target).where( - Target.stratum_id == budget_stratum.id, - Target.variable == var_name, - Target.period == year, - Target.source == DataSource.OBR, - ) - ).first() - - if existing: - continue - - session.add( - Target( - stratum_id=budget_stratum.id, - variable=var_name, - period=year, - value=data[var_name], - target_type=var_type, - source=DataSource.OBR, - source_url=SOURCE_URL, - is_preliminary=year > 2024, - ) - ) - - # UK Economy stratum for macro indicators - economy_stratum = get_or_create_stratum( - session, - name="UK Economy", - jurisdiction=Jurisdiction.UK, - constraints=[], # Whole economy, not a subset - description="UK macroeconomic indicators", - stratum_group_id="obr_economy", - ) - - # Economic rate targets - rate_vars = [ - "real_gdp_growth", - "unemployment_rate", - "cpi_inflation", - "rpi_inflation", - "bank_rate", - ] - - for var_name in rate_vars: - if var_name not in data: - continue - - existing = session.exec( - select(Target).where( - Target.stratum_id == economy_stratum.id, - Target.variable == var_name, - Target.period == year, - Target.source == DataSource.OBR, - ) - ).first() - - if existing: - continue - - session.add( - Target( - stratum_id=economy_stratum.id, - variable=var_name, - period=year, - value=data[var_name], - target_type=TargetType.RATE, - source=DataSource.OBR, - source_url=SOURCE_URL, - is_preliminary=year > 2024, - ) - ) - - # Employment count - if "employment" in data: - existing = session.exec( - select(Target).where( - Target.stratum_id == economy_stratum.id, - Target.variable == "employment", - Target.period == year, - Target.source == DataSource.OBR, - ) - ).first() - - if not existing: - session.add( - Target( - stratum_id=economy_stratum.id, - variable="employment", - period=year, - value=data["employment"], - target_type=TargetType.COUNT, - source=DataSource.OBR, - source_url=SOURCE_URL, - is_preliminary=year > 2024, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the OBR ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_obr_targets(session) - print(f"Loaded OBR projections to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_ons.py b/db/etl_ons.py deleted file mode 100644 index d4619231..00000000 --- a/db/etl_ons.py +++ /dev/null @@ -1,310 +0,0 @@ -""" -ETL for ONS (Office for National Statistics) projections. - -Loads UK population and household projections from ONS. -Data sources: -- National population projections: https://www.ons.gov.uk/peoplepopulationandcommunity/populationandmigration/populationprojections -- Household projections: https://www.ons.gov.uk/peoplepopulationandcommunity/populationandmigration/populationprojections/datasets/householdprojectionsforengland -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# ONS 2022-based projections data (2024-2034) -# Source: ONS National Population Projections (2022-based) -# https://www.ons.gov.uk/peoplepopulationandcommunity/populationandmigration/populationprojections -# Population counts are for the UK as a whole -# Household projections based on ONS household projections data -ONS_DATA = { - # Mid-2024 projections - 2024: { - # Population by broad age groups (counts) - "population_total": 68_265_000, - "population_age_0_15": 12_410_000, # Children - "population_age_16_64": 42_638_000, # Working age - "population_age_65_plus": 13_217_000, # Pension age - # Household projections - "households_total": 28_500_000, - "average_household_size": 2.40, - }, - 2025: { - "population_total": 68_626_000, - "population_age_0_15": 12_448_000, - "population_age_16_64": 42_703_000, - "population_age_65_plus": 13_475_000, - "households_total": 28_750_000, - "average_household_size": 2.39, - }, - 2026: { - "population_total": 68_978_000, - "population_age_0_15": 12_485_000, - "population_age_16_64": 42_759_000, - "population_age_65_plus": 13_734_000, - "households_total": 29_000_000, - "average_household_size": 2.38, - }, - 2027: { - "population_total": 69_321_000, - "population_age_0_15": 12_521_000, - "population_age_16_64": 42_808_000, - "population_age_65_plus": 13_992_000, - "households_total": 29_250_000, - "average_household_size": 2.37, - }, - 2028: { - "population_total": 69_656_000, - "population_age_0_15": 12_555_000, - "population_age_16_64": 42_852_000, - "population_age_65_plus": 14_249_000, - "households_total": 29_500_000, - "average_household_size": 2.36, - }, - 2029: { - "population_total": 69_984_000, - "population_age_0_15": 12_588_000, - "population_age_16_64": 42_891_000, - "population_age_65_plus": 14_505_000, - "households_total": 29_750_000, - "average_household_size": 2.35, - }, - 2030: { - "population_total": 70_306_000, - "population_age_0_15": 12_619_000, - "population_age_16_64": 42_927_000, - "population_age_65_plus": 14_760_000, - "households_total": 30_000_000, - "average_household_size": 2.34, - }, - 2031: { - "population_total": 70_622_000, - "population_age_0_15": 12_649_000, - "population_age_16_64": 42_960_000, - "population_age_65_plus": 15_013_000, - "households_total": 30_250_000, - "average_household_size": 2.34, - }, - 2032: { - "population_total": 70_933_000, - "population_age_0_15": 12_678_000, - "population_age_16_64": 42_991_000, - "population_age_65_plus": 15_264_000, - "households_total": 30_500_000, - "average_household_size": 2.33, - }, - 2033: { - "population_total": 71_240_000, - "population_age_0_15": 12_706_000, - "population_age_16_64": 43_020_000, - "population_age_65_plus": 15_514_000, - "households_total": 30_750_000, - "average_household_size": 2.32, - }, - 2034: { - "population_total": 71_543_000, - "population_age_0_15": 12_733_000, - "population_age_16_64": 43_047_000, - "population_age_65_plus": 15_763_000, - "households_total": 31_000_000, - "average_household_size": 2.31, - }, -} - -SOURCE_URL = "https://www.ons.gov.uk/peoplepopulationandcommunity/populationandmigration/populationprojections" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_ons_targets(session: Session, years: list[int] | None = None): - """ - Load ONS population and household projections into database. - - Args: - session: Database session - years: Years to load (default: all available 2024-2034) - """ - if years is None: - years = list(ONS_DATA.keys()) - - for year in years: - if year not in ONS_DATA: - continue - - data = ONS_DATA[year] - - # UK population stratum - population_stratum = get_or_create_stratum( - session, - name="UK Population", - jurisdiction=Jurisdiction.UK, - constraints=[("entity_type", "==", "person")], # Person-level data - description="UK total population and age group breakdowns", - stratum_group_id="ons_population", - ) - - # Population count targets - population_vars = [ - ("population_total", TargetType.COUNT), - ("population_age_0_15", TargetType.COUNT), - ("population_age_16_64", TargetType.COUNT), - ("population_age_65_plus", TargetType.COUNT), - ] - - for var_name, var_type in population_vars: - if var_name not in data: - continue - - # Check for existing target - existing = session.exec( - select(Target).where( - Target.stratum_id == population_stratum.id, - Target.variable == var_name, - Target.period == year, - Target.source == DataSource.ONS, - ) - ).first() - - if existing: - continue - - session.add( - Target( - stratum_id=population_stratum.id, - variable=var_name, - period=year, - value=data[var_name], - target_type=var_type, - source=DataSource.ONS, - source_url=SOURCE_URL, - is_preliminary=True, # All projections are preliminary - ) - ) - - # UK households stratum - households_stratum = get_or_create_stratum( - session, - name="UK Households", - jurisdiction=Jurisdiction.UK, - constraints=[("entity_type", "==", "household")], # Household-level data - description="UK household projections", - stratum_group_id="ons_households", - ) - - # Household count targets - if "households_total" in data: - existing = session.exec( - select(Target).where( - Target.stratum_id == households_stratum.id, - Target.variable == "households_total", - Target.period == year, - Target.source == DataSource.ONS, - ) - ).first() - - if not existing: - session.add( - Target( - stratum_id=households_stratum.id, - variable="households_total", - period=year, - value=data["households_total"], - target_type=TargetType.COUNT, - source=DataSource.ONS, - source_url=SOURCE_URL, - is_preliminary=True, - ) - ) - - # Average household size (rate) - if "average_household_size" in data: - existing = session.exec( - select(Target).where( - Target.stratum_id == households_stratum.id, - Target.variable == "average_household_size", - Target.period == year, - Target.source == DataSource.ONS, - ) - ).first() - - if not existing: - session.add( - Target( - stratum_id=households_stratum.id, - variable="average_household_size", - period=year, - value=data["average_household_size"], - target_type=TargetType.RATE, - source=DataSource.ONS, - source_url=SOURCE_URL, - is_preliminary=True, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the ONS ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_ons_targets(session) - print(f"Loaded ONS projections to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_snap.py b/db/etl_snap.py deleted file mode 100644 index 7ad1bc60..00000000 --- a/db/etl_snap.py +++ /dev/null @@ -1,334 +0,0 @@ -""" -ETL for USDA SNAP (Supplemental Nutrition Assistance Program) targets. - -Loads data from USDA Food and Nutrition Service into the targets database. -Data source: https://www.fns.usda.gov/pd/supplemental-nutrition-assistance-program-snap -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) -from arch.normalization import SourceFact, as_target, convert_units, target_kwargs - -# State FIPS codes -STATE_FIPS = { - "AL": "01", "AK": "02", "AZ": "04", "AR": "05", "CA": "06", - "CO": "08", "CT": "09", "DE": "10", "DC": "11", "FL": "12", - "GA": "13", "HI": "15", "ID": "16", "IL": "17", "IN": "18", - "IA": "19", "KS": "20", "KY": "21", "LA": "22", "ME": "23", - "MD": "24", "MA": "25", "MI": "26", "MN": "27", "MS": "28", - "MO": "29", "MT": "30", "NE": "31", "NV": "32", "NH": "33", - "NJ": "34", "NM": "35", "NY": "36", "NC": "37", "ND": "38", - "OH": "39", "OK": "40", "OR": "41", "PA": "42", "RI": "44", - "SC": "45", "SD": "46", "TN": "47", "TX": "48", "UT": "49", - "VT": "50", "VA": "51", "WA": "53", "WV": "54", "WI": "55", - "WY": "56", "PR": "72", "VI": "78", "GU": "66", -} - -# SNAP data by year (from USDA FNS) -# Source: https://www.fns.usda.gov/pd/supplemental-nutrition-assistance-program-snap -# Values: households (thousands), benefits (millions of dollars) -SNAP_DATA = { - 2023: { - "national": { - "households": 22_323, # thousands - "participants": 42_104, # thousands - "benefits": 112_848, # millions - "avg_benefit_per_person": 216.99, # dollars/month - }, - # Top 10 states by participation - "states": { - "CA": {"households": 2_891, "participants": 5_123, "benefits": 14_234}, - "TX": {"households": 2_156, "participants": 4_012, "benefits": 9_876}, - "FL": {"households": 1_987, "participants": 3_654, "benefits": 8_234}, - "NY": {"households": 1_654, "participants": 2_987, "benefits": 7_654}, - "PA": {"households": 1_234, "participants": 2_345, "benefits": 5_432}, - "IL": {"households": 1_123, "participants": 2_134, "benefits": 4_987}, - "OH": {"households": 987, "participants": 1_876, "benefits": 4_321}, - "GA": {"households": 912, "participants": 1_765, "benefits": 3_987}, - "MI": {"households": 876, "participants": 1_654, "benefits": 3_654}, - "NC": {"households": 834, "participants": 1_567, "benefits": 3_432}, - }, - }, - 2022: { - "national": { - "households": 21_567, - "participants": 41_234, - "benefits": 119_432, # Higher due to emergency allotments - "avg_benefit_per_person": 234.56, - }, - "states": { - "CA": {"households": 2_765, "participants": 4_987, "benefits": 15_123}, - "TX": {"households": 2_098, "participants": 3_876, "benefits": 10_234}, - "FL": {"households": 1_876, "participants": 3_456, "benefits": 8_765}, - "NY": {"households": 1_598, "participants": 2_876, "benefits": 8_123}, - "PA": {"households": 1_198, "participants": 2_234, "benefits": 5_765}, - "IL": {"households": 1_087, "participants": 2_054, "benefits": 5_234}, - "OH": {"households": 954, "participants": 1_798, "benefits": 4_567}, - "GA": {"households": 887, "participants": 1_698, "benefits": 4_234}, - "MI": {"households": 845, "participants": 1_598, "benefits": 3_876}, - "NC": {"households": 798, "participants": 1_498, "benefits": 3_654}, - }, - }, - 2021: { - "national": { - "households": 21_876, - "participants": 41_987, - "benefits": 113_456, - "avg_benefit_per_person": 228.12, - }, - "states": { - "CA": {"households": 2_834, "participants": 5_076, "benefits": 14_567}, - "TX": {"households": 2_123, "participants": 3_945, "benefits": 9_987}, - "FL": {"households": 1_934, "participants": 3_587, "benefits": 8_456}, - "NY": {"households": 1_623, "participants": 2_934, "benefits": 7_876}, - "PA": {"households": 1_212, "participants": 2_287, "benefits": 5_543}, - "IL": {"households": 1_098, "participants": 2_087, "benefits": 5_098}, - "OH": {"households": 967, "participants": 1_834, "benefits": 4_432}, - "GA": {"households": 898, "participants": 1_723, "benefits": 4_098}, - "MI": {"households": 856, "participants": 1_623, "benefits": 3_765}, - "NC": {"households": 812, "participants": 1_534, "benefits": 3_543}, - }, - }, -} - -SOURCE_URL = "https://www.fns.usda.gov/pd/supplemental-nutrition-assistance-program-snap" - - -def build_snap_target( - stratum: Stratum, - *, - variable: str, - raw_value: float, - period: int, - raw_unit: str, - output_unit: str, - factor: float, - target_type: TargetType, - source_table: str, -) -> Target: - """Build a SNAP target input from a source fact and unit conversion.""" - fact = SourceFact( - name=variable, - value=raw_value, - period=period, - unit=raw_unit, - source=DataSource.USDA_SNAP, - jurisdiction=stratum.jurisdiction, - source_table=source_table, - source_url=SOURCE_URL, - ) - converted = convert_units(fact, factor, output_unit) - blueprint = as_target( - converted, - variable=variable, - target_type=target_type, - stratum_name=stratum.name, - ) - return Target(**target_kwargs(blueprint, stratum_id=stratum.id)) - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_snap_targets(session: Session, years: list[int] | None = None): - """ - Load SNAP targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SNAP_DATA.keys()) - - for year in years: - if year not in SNAP_DATA: - continue - - data = SNAP_DATA[year] - - # Create national SNAP stratum - national_stratum = get_or_create_stratum( - session, - name="US SNAP Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("snap", "==", "1")], - description="All SNAP recipient households/individuals in the US", - stratum_group_id="snap_national", - ) - - # Add national totals - national_data = data["national"] - - session.add( - build_snap_target( - national_stratum, - variable="snap_household_count", - period=year, - raw_value=national_data["households"], - raw_unit="thousands", - output_unit="count", - factor=1000, - target_type=TargetType.COUNT, - source_table="SNAP National Summary", - ) - ) - - session.add( - build_snap_target( - national_stratum, - variable="snap_participant_count", - period=year, - raw_value=national_data["participants"], - raw_unit="thousands", - output_unit="count", - factor=1000, - target_type=TargetType.COUNT, - source_table="SNAP National Summary", - ) - ) - - session.add( - build_snap_target( - national_stratum, - variable="snap_benefits", - period=year, - raw_value=national_data["benefits"], - raw_unit="millions_of_dollars", - output_unit="dollars", - factor=1_000_000, - target_type=TargetType.AMOUNT, - source_table="SNAP National Summary", - ) - ) - - # Add state-level targets - for state_abbrev, state_data in data.get("states", {}).items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} SNAP Recipients", - jurisdiction=Jurisdiction.US, - constraints=[ - ("snap", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"SNAP recipients in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="snap_states", - ) - - session.add( - build_snap_target( - state_stratum, - variable="snap_household_count", - period=year, - raw_value=state_data["households"], - raw_unit="thousands", - output_unit="count", - factor=1000, - target_type=TargetType.COUNT, - source_table="SNAP State Summary", - ) - ) - - session.add( - build_snap_target( - state_stratum, - variable="snap_participant_count", - period=year, - raw_value=state_data["participants"], - raw_unit="thousands", - output_unit="count", - factor=1000, - target_type=TargetType.COUNT, - source_table="SNAP State Summary", - ) - ) - - session.add( - build_snap_target( - state_stratum, - variable="snap_benefits", - period=year, - raw_value=state_data["benefits"], - raw_unit="millions_of_dollars", - output_unit="dollars", - factor=1_000_000, - target_type=TargetType.AMOUNT, - source_table="SNAP State Summary", - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the SNAP ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_snap_targets(session) - print(f"Loaded SNAP targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_soi.py b/db/etl_soi.py deleted file mode 100644 index 914d3eeb..00000000 --- a/db/etl_soi.py +++ /dev/null @@ -1,624 +0,0 @@ -""" -ETL for IRS Statistics of Income (SOI) targets. - -Loads IRS SOI Publication 1304 Table 1.1 into the targets database. -""" - -from __future__ import annotations - -import hashlib -from functools import lru_cache -from importlib.resources import files -from io import BytesIO -from typing import Any, TypedDict -from urllib.request import Request, urlopen - -import pandas as pd -import yaml -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -SOURCE_URL = "https://www.irs.gov/statistics/soi-tax-stats-individual-income-tax-statistics" -TABLE_1_1_SHEET_NAME = "TBL11" -TABLE_1_4_SHEET_NAME = "TBL14" -TABLE_1_1_MONEY_SCALE = 1_000 -TABLE_1_4_MONEY_SCALE = 1_000 -TABLE_1_1_PACKAGE_DIR = "data/irs_soi/table_1_1" -TABLE_1_4_PACKAGE_DIR = "data/irs_soi/table_1_4" -TABLE_1_1_MANIFEST = "manifest.yaml" -TABLE_1_4_MANIFEST = "manifest.yaml" - -# AGI bracket definitions (lower, upper) in dollars. -AGI_BRACKETS = { - "under_1": (float("-inf"), 1), - "1_to_5k": (1, 5_000), - "5k_to_10k": (5_000, 10_000), - "10k_to_15k": (10_000, 15_000), - "15k_to_20k": (15_000, 20_000), - "20k_to_25k": (20_000, 25_000), - "25k_to_30k": (25_000, 30_000), - "30k_to_40k": (30_000, 40_000), - "40k_to_50k": (40_000, 50_000), - "50k_to_75k": (50_000, 75_000), - "75k_to_100k": (75_000, 100_000), - "100k_to_200k": (100_000, 200_000), - "200k_to_500k": (200_000, 500_000), - "500k_to_1m": (500_000, 1_000_000), - "1m_to_1_5m": (1_000_000, 1_500_000), - "1_5m_to_2m": (1_500_000, 2_000_000), - "2m_to_5m": (2_000_000, 5_000_000), - "5m_to_10m": (5_000_000, 10_000_000), - "10m_plus": (10_000_000, float("inf")), -} - -TABLE_1_1_AGI_LABEL_TO_BRACKET = { - "No adjusted gross income": "under_1", - "$1 under $5,000": "1_to_5k", - "$5,000 under $10,000": "5k_to_10k", - "$10,000 under $15,000": "10k_to_15k", - "$15,000 under $20,000": "15k_to_20k", - "$20,000 under $25,000": "20k_to_25k", - "$25,000 under $30,000": "25k_to_30k", - "$30,000 under $40,000": "30k_to_40k", - "$40,000 under $50,000": "40k_to_50k", - "$50,000 under $75,000": "50k_to_75k", - "$75,000 under $100,000": "75k_to_100k", - "$100,000 under $200,000": "100k_to_200k", - "$200,000 under $500,000": "200k_to_500k", - "$500,000 under $1,000,000": "500k_to_1m", - "$1,000,000 under $1,500,000": "1m_to_1_5m", - "$1,500,000 under $2,000,000": "1_5m_to_2m", - "$2,000,000 under $5,000,000": "2m_to_5m", - "$5,000,000 under $10,000,000": "5m_to_10m", - "$10,000,000 or more": "10m_plus", -} - - -class SOITable11Data(TypedDict): - source_url: str - total_returns: int - total_agi: int - total_income_tax: int - returns_by_agi_bracket: dict[str, int] - agi_by_bracket: dict[str, int] - income_tax_by_bracket: dict[str, int] - - -class SOITable14Data(TypedDict): - source_url: str - total_employment_income_returns: int - total_employment_income: int - employment_income_returns_by_agi_bracket: dict[str, int] - employment_income_by_agi_bracket: dict[str, int] - - -@lru_cache(maxsize=1) -def _table_1_1_manifest() -> dict[str, Any]: - manifest_path = files("db").joinpath(TABLE_1_1_PACKAGE_DIR, TABLE_1_1_MANIFEST) - with manifest_path.open("r", encoding="utf-8") as f: - return yaml.safe_load(f) - - -@lru_cache(maxsize=1) -def _table_1_4_manifest() -> dict[str, Any]: - manifest_path = files("db").joinpath(TABLE_1_4_PACKAGE_DIR, TABLE_1_4_MANIFEST) - with manifest_path.open("r", encoding="utf-8") as f: - return yaml.safe_load(f) - - -def available_soi_years() -> list[int]: - """Return SOI Table 1.1 years available in the source manifest.""" - return sorted(int(year) for year in _table_1_1_manifest()["files"]) - - -def available_soi_table_1_4_years() -> list[int]: - """Return SOI Table 1.4 years available in the source manifest.""" - return sorted(int(year) for year in _table_1_4_manifest()["files"]) - - -def soi_table_1_1_source_url(year: int) -> str: - """Return the IRS source URL for a Table 1.1 year.""" - return _table_1_1_file_spec(year)["source_url"] - - -def soi_table_1_4_source_url(year: int) -> str: - """Return the IRS source URL for a Table 1.4 year.""" - return _table_1_4_file_spec(year)["source_url"] - - -def _table_1_1_file_spec(year: int) -> dict[str, str]: - files_by_year = _table_1_1_manifest()["files"] - try: - return files_by_year[year] - except KeyError: - return files_by_year[str(year)] - - -def _table_1_4_file_spec(year: int) -> dict[str, str]: - files_by_year = _table_1_4_manifest()["files"] - try: - return files_by_year[year] - except KeyError: - return files_by_year[str(year)] - - -# Compatibility manifest for older callers. Numeric SOI values are parsed from -# the source files at load time rather than embedded in Python. -SOI_DATA = { - year: {"source_url": soi_table_1_1_source_url(year)} - for year in available_soi_years() -} - - -@lru_cache(maxsize=None) -def _table_1_1_content(year: int) -> bytes: - return _table_content(year, _table_1_1_file_spec, TABLE_1_1_PACKAGE_DIR) - - -@lru_cache(maxsize=None) -def _table_1_4_content(year: int) -> bytes: - return _table_content(year, _table_1_4_file_spec, TABLE_1_4_PACKAGE_DIR) - - -def _table_content(year: int, file_spec_fn, package_dir: str) -> bytes: - spec = file_spec_fn(year) - package_path = files("db").joinpath(package_dir, spec["filename"]) - if package_path.is_file(): - content = package_path.read_bytes() - else: - request = Request( - spec["source_url"], - headers={"User-Agent": "policyengine-arch-data/0.1", "Accept": "*/*"}, - ) - with urlopen(request, timeout=120) as response: - content = response.read() - - expected_sha = spec.get("sha256") - if expected_sha: - actual_sha = hashlib.sha256(content).hexdigest() - if actual_sha != expected_sha: - raise ValueError( - f"SOI source file {year} checksum mismatch: " - f"expected {expected_sha}, got {actual_sha}" - ) - return content - - -@lru_cache(maxsize=None) -def _read_soi_table_1_1_frame(year: int) -> pd.DataFrame: - return pd.read_excel( - BytesIO(_table_1_1_content(year)), - sheet_name=TABLE_1_1_SHEET_NAME, - header=None, - dtype=object, - engine="xlrd", - ) - - -@lru_cache(maxsize=None) -def _read_soi_table_1_4_frame(year: int) -> pd.DataFrame: - return pd.read_excel( - BytesIO(_table_1_4_content(year)), - sheet_name=TABLE_1_4_SHEET_NAME, - header=None, - dtype=object, - engine="xlrd", - ) - - -def load_soi_table_1_1_data(year: int) -> SOITable11Data: - """Parse one year of IRS SOI Publication 1304 Table 1.1.""" - return _parse_soi_table_1_1_frame( - _read_soi_table_1_1_frame(year), - source_url=soi_table_1_1_source_url(year), - ) - - -def load_soi_table_1_4_data(year: int) -> SOITable14Data: - """Parse wage targets from IRS SOI Publication 1304 Table 1.4.""" - return _parse_soi_table_1_4_frame( - _read_soi_table_1_4_frame(year), - source_url=soi_table_1_4_source_url(year), - ) - - -def _clean_label(value: object) -> str: - if value is None or pd.isna(value): - return "" - return " ".join(str(value).replace("\n", " ").split()).strip() - - -def _numeric_value(value: object) -> float: - if value is None or pd.isna(value): - raise ValueError("Expected numeric SOI cell, got empty value") - if isinstance(value, str): - clean = value.replace("$", "").replace(",", "").strip() - if clean.startswith("["): - raise ValueError(f"Expected numeric SOI cell, got footnote marker {value}") - return float(clean) - return float(value) - - -def _count_value(value: object) -> int: - return int(round(_numeric_value(value))) - - -def _money_value(value: object) -> int: - return int(round(_numeric_value(value) * TABLE_1_1_MONEY_SCALE)) - - -def _table_1_4_money_value(value: object) -> int: - return int(round(_numeric_value(value) * TABLE_1_4_MONEY_SCALE)) - - -def _header_text(df: pd.DataFrame, column: int) -> str: - pieces = [_clean_label(df.iat[row, column]) for row in range(min(8, len(df)))] - return " ".join(piece for piece in pieces if piece).lower() - - -def _find_column(df: pd.DataFrame, *phrases: str) -> int: - for column in range(df.shape[1]): - header = _header_text(df, column) - if all(phrase.lower() in header for phrase in phrases): - return column - joined = ", ".join(phrases) - raise ValueError(f"Could not find SOI Table 1.1 column containing: {joined}") - - -def _table_1_1_size_rows(df: pd.DataFrame) -> pd.DataFrame: - labels = df.iloc[:, 0].map(_clean_label) - start_matches = labels[labels == "All returns"] - if start_matches.empty: - raise ValueError("Could not find SOI Table 1.1 'All returns' row") - start = int(start_matches.index[0]) - - end_label = "Accumulated from smallest size of adjusted gross income" - end_matches = labels[(labels == end_label) & (labels.index > start)] - end = int(end_matches.index[0]) if not end_matches.empty else len(df) - return df.loc[start : end - 1] - - -def _table_1_4_size_rows(df: pd.DataFrame) -> pd.DataFrame: - labels = df.iloc[:, 0].map(_clean_label) - start_matches = labels[labels == "All returns, total"] - if start_matches.empty: - raise ValueError("Could not find SOI Table 1.4 'All returns, total' row") - start = int(start_matches.index[0]) - - end_matches = labels[(labels == "Taxable returns, total") & (labels.index > start)] - end = int(end_matches.index[0]) if not end_matches.empty else len(df) - return df.loc[start : end - 1] - - -def _row_by_label(rows: pd.DataFrame, label: str) -> pd.Series: - labels = rows.iloc[:, 0].map(_clean_label) - matches = rows[labels == label] - if matches.empty: - raise ValueError(f"Could not find SOI Table 1.1 row: {label}") - return matches.iloc[0] - - -def _parse_soi_table_1_1_frame( - df: pd.DataFrame, - *, - source_url: str, -) -> SOITable11Data: - count_col = _find_column(df, "number of returns") - agi_col = _find_column(df, "adjusted gross income less deficit", "amount") - income_tax_col = _find_column(df, "total income tax", "amount") - rows = _table_1_1_size_rows(df) - all_returns = _row_by_label(rows, "All returns") - - returns_by_bracket = {} - agi_by_bracket = {} - income_tax_by_bracket = {} - for source_label, bracket_name in TABLE_1_1_AGI_LABEL_TO_BRACKET.items(): - row = _row_by_label(rows, source_label) - returns_by_bracket[bracket_name] = _count_value(row.iat[count_col]) - agi_by_bracket[bracket_name] = _money_value(row.iat[agi_col]) - income_tax_by_bracket[bracket_name] = _money_value(row.iat[income_tax_col]) - - return { - "source_url": source_url, - "total_returns": _count_value(all_returns.iat[count_col]), - "total_agi": _money_value(all_returns.iat[agi_col]), - "total_income_tax": _money_value(all_returns.iat[income_tax_col]), - "returns_by_agi_bracket": returns_by_bracket, - "agi_by_bracket": agi_by_bracket, - "income_tax_by_bracket": income_tax_by_bracket, - } - - -def _parse_soi_table_1_4_frame( - df: pd.DataFrame, - *, - source_url: str, -) -> SOITable14Data: - wage_count_col = _find_column(df, "wages", "number of returns") - wage_amount_col = wage_count_col + 1 - if "amount" not in _header_text(df, wage_amount_col): - raise ValueError("Could not find SOI Table 1.4 wage amount column") - rows = _table_1_4_size_rows(df) - all_returns = _row_by_label(rows, "All returns, total") - - wage_returns_by_bracket = {} - wage_amount_by_bracket = {} - for source_label, bracket_name in TABLE_1_1_AGI_LABEL_TO_BRACKET.items(): - row = _row_by_label(rows, source_label) - wage_returns_by_bracket[bracket_name] = _count_value(row.iat[wage_count_col]) - wage_amount_by_bracket[bracket_name] = _table_1_4_money_value( - row.iat[wage_amount_col] - ) - - return { - "source_url": source_url, - "total_employment_income_returns": _count_value( - all_returns.iat[wage_count_col] - ), - "total_employment_income": _table_1_4_money_value( - all_returns.iat[wage_amount_col] - ), - "employment_income_returns_by_agi_bracket": wage_returns_by_bracket, - "employment_income_by_agi_bracket": wage_amount_by_bracket, - } - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def _add_target( - session: Session, - *, - stratum_id: int, - variable: str, - period: int, - value: float, - target_type: TargetType, - source_url: str, - source_table: str = "Table 1.1", -) -> None: - session.add( - Target( - stratum_id=stratum_id, - variable=variable, - period=period, - value=value, - target_type=target_type, - source=DataSource.IRS_SOI, - source_table=source_table, - source_url=source_url, - ) - ) - - -def load_soi_targets(session: Session, years: list[int] | None = None) -> None: - """ - Load SOI targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = available_soi_years() - - for year in years: - if year not in available_soi_years(): - continue - - data = load_soi_table_1_1_data(year) - source_url = data["source_url"] - table_1_4_data = None - table_1_4_source_url = None - if year in available_soi_table_1_4_years(): - table_1_4_data = load_soi_table_1_4_data(year) - table_1_4_source_url = table_1_4_data["source_url"] - - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - _add_target( - session, - stratum_id=national_stratum.id, - variable="tax_unit_count", - period=year, - value=data["total_returns"], - target_type=TargetType.COUNT, - source_url=source_url, - ) - _add_target( - session, - stratum_id=national_stratum.id, - variable="adjusted_gross_income", - period=year, - value=data["total_agi"], - target_type=TargetType.AMOUNT, - source_url=source_url, - ) - _add_target( - session, - stratum_id=national_stratum.id, - variable="income_tax_liability", - period=year, - value=data["total_income_tax"], - target_type=TargetType.AMOUNT, - source_url=source_url, - ) - - if table_1_4_data is not None and table_1_4_source_url is not None: - _add_target( - session, - stratum_id=national_stratum.id, - variable="employment_income", - period=year, - value=table_1_4_data["total_employment_income_returns"], - target_type=TargetType.COUNT, - source_url=table_1_4_source_url, - source_table="Table 1.4", - ) - _add_target( - session, - stratum_id=national_stratum.id, - variable="employment_income", - period=year, - value=table_1_4_data["total_employment_income"], - target_type=TargetType.AMOUNT, - source_url=table_1_4_source_url, - source_table="Table 1.4", - ) - - for bracket_name, (lower, upper) in AGI_BRACKETS.items(): - constraints = [] - if lower != float("-inf"): - constraints.append(("adjusted_gross_income", ">=", str(lower))) - if upper != float("inf"): - constraints.append(("adjusted_gross_income", "<", str(upper))) - - bracket_stratum = get_or_create_stratum( - session, - name=f"US Filers AGI {bracket_name}", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=constraints, - description=f"Tax filers with AGI in {bracket_name} bracket", - parent_id=national_stratum.id, - stratum_group_id="agi_brackets", - ) - - if bracket_name in data["returns_by_agi_bracket"]: - _add_target( - session, - stratum_id=bracket_stratum.id, - variable="tax_unit_count", - period=year, - value=data["returns_by_agi_bracket"][bracket_name], - target_type=TargetType.COUNT, - source_url=source_url, - ) - - if bracket_name in data["agi_by_bracket"]: - _add_target( - session, - stratum_id=bracket_stratum.id, - variable="adjusted_gross_income", - period=year, - value=data["agi_by_bracket"][bracket_name], - target_type=TargetType.AMOUNT, - source_url=source_url, - ) - - if bracket_name in data["income_tax_by_bracket"]: - _add_target( - session, - stratum_id=bracket_stratum.id, - variable="income_tax_liability", - period=year, - value=data["income_tax_by_bracket"][bracket_name], - target_type=TargetType.AMOUNT, - source_url=source_url, - ) - - if table_1_4_data is not None and table_1_4_source_url is not None: - if ( - bracket_name - in table_1_4_data["employment_income_returns_by_agi_bracket"] - ): - _add_target( - session, - stratum_id=bracket_stratum.id, - variable="employment_income", - period=year, - value=table_1_4_data[ - "employment_income_returns_by_agi_bracket" - ][bracket_name], - target_type=TargetType.COUNT, - source_url=table_1_4_source_url, - source_table="Table 1.4", - ) - - if bracket_name in table_1_4_data["employment_income_by_agi_bracket"]: - _add_target( - session, - stratum_id=bracket_stratum.id, - variable="employment_income", - period=year, - value=table_1_4_data["employment_income_by_agi_bracket"][ - bracket_name - ], - target_type=TargetType.AMOUNT, - source_url=table_1_4_source_url, - source_table="Table 1.4", - ) - - session.commit() - - -def run_etl(db_path=None) -> None: - """Run the SOI ETL pipeline.""" - from pathlib import Path - - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_soi_targets(session) - print(f"Loaded SOI targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_soi_credits.py b/db/etl_soi_credits.py deleted file mode 100644 index 027c59aa..00000000 --- a/db/etl_soi_credits.py +++ /dev/null @@ -1,1107 +0,0 @@ -""" -ETL for IRS Statistics of Income (SOI) state-level tax credit targets. - -Loads state-by-state EITC, CTC, and ACTC data from IRS SOI tables. -Data sources: -- EITC: https://www.irs.gov/credits-deductions/individuals/earned-income-tax-credit/earned-income-tax-credit-statistics -- CTC/ACTC: https://www.irs.gov/statistics/soi-tax-stats-state-data-fy-2022 -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes for all 50 states + DC -STATE_FIPS = { - "AL": "01", - "AK": "02", - "AZ": "04", - "AR": "05", - "CA": "06", - "CO": "08", - "CT": "09", - "DE": "10", - "DC": "11", - "FL": "12", - "GA": "13", - "HI": "15", - "ID": "16", - "IL": "17", - "IN": "18", - "IA": "19", - "KS": "20", - "KY": "21", - "LA": "22", - "ME": "23", - "MD": "24", - "MA": "25", - "MI": "26", - "MN": "27", - "MS": "28", - "MO": "29", - "MT": "30", - "NE": "31", - "NV": "32", - "NH": "33", - "NJ": "34", - "NM": "35", - "NY": "36", - "NC": "37", - "ND": "38", - "OH": "39", - "OK": "40", - "OR": "41", - "PA": "42", - "RI": "44", - "SC": "45", - "SD": "46", - "TN": "47", - "TX": "48", - "UT": "49", - "VT": "50", - "VA": "51", - "WA": "53", - "WV": "54", - "WI": "55", - "WY": "56", -} - -# State-level credit data for 2021 -# EITC: ~25M claims, ~$60B total nationally -# CTC: ~35M claims nationally -# ACTC: ~18M claims, ~$34B total nationally -# Values scaled by state population and poverty rates - -SOI_CREDITS_DATA = { - 2021: { - "AL": { - "eitc_claims": 498_000, - "eitc_amount": 1_345_000_000, - "ctc_claims": 687_000, - "ctc_amount": 1_112_000_000, - "actc_claims": 356_000, - "actc_amount": 534_000_000, - }, - "AK": { - "eitc_claims": 52_000, - "eitc_amount": 142_000_000, - "ctc_claims": 98_000, - "ctc_amount": 159_000_000, - "actc_claims": 43_000, - "actc_amount": 65_000_000, - }, - "AZ": { - "eitc_claims": 598_000, - "eitc_amount": 1_623_000_000, - "ctc_claims": 987_000, - "ctc_amount": 1_598_000_000, - "actc_claims": 498_000, - "actc_amount": 747_000_000, - }, - "AR": { - "eitc_claims": 298_000, - "eitc_amount": 789_000_000, - "ctc_claims": 398_000, - "ctc_amount": 644_000_000, - "actc_claims": 212_000, - "actc_amount": 318_000_000, - }, - "CA": { - "eitc_claims": 2_987_000, - "eitc_amount": 8_234_000_000, - "ctc_claims": 4_987_000, - "ctc_amount": 8_076_000_000, - "actc_claims": 2_134_000, - "actc_amount": 3_201_000_000, - }, - "CO": { - "eitc_claims": 387_000, - "eitc_amount": 1_067_000_000, - "ctc_claims": 698_000, - "ctc_amount": 1_130_000_000, - "actc_claims": 298_000, - "actc_amount": 447_000_000, - }, - "CT": { - "eitc_claims": 243_000, - "eitc_amount": 634_000_000, - "ctc_claims": 387_000, - "ctc_amount": 627_000_000, - "actc_claims": 156_000, - "actc_amount": 234_000_000, - }, - "DE": { - "eitc_claims": 76_000, - "eitc_amount": 198_000_000, - "ctc_claims": 112_000, - "ctc_amount": 181_000_000, - "actc_claims": 52_000, - "actc_amount": 78_000_000, - }, - "DC": { - "eitc_claims": 54_000, - "eitc_amount": 134_000_000, - "ctc_claims": 65_000, - "ctc_amount": 105_000_000, - "actc_claims": 32_000, - "actc_amount": 48_000_000, - }, - "FL": { - "eitc_claims": 1_876_000, - "eitc_amount": 5_123_000_000, - "ctc_claims": 2_654_000, - "ctc_amount": 4_299_000_000, - "actc_claims": 1_287_000, - "actc_amount": 1_931_000_000, - }, - "GA": { - "eitc_claims": 987_000, - "eitc_amount": 2_698_000_000, - "ctc_claims": 1_432_000, - "ctc_amount": 2_320_000_000, - "actc_claims": 698_000, - "actc_amount": 1_047_000_000, - }, - "HI": { - "eitc_claims": 98_000, - "eitc_amount": 267_000_000, - "ctc_claims": 165_000, - "ctc_amount": 267_000_000, - "actc_claims": 76_000, - "actc_amount": 114_000_000, - }, - "ID": { - "eitc_claims": 132_000, - "eitc_amount": 356_000_000, - "ctc_claims": 287_000, - "ctc_amount": 465_000_000, - "actc_claims": 134_000, - "actc_amount": 201_000_000, - }, - "IL": { - "eitc_claims": 987_000, - "eitc_amount": 2_678_000_000, - "ctc_claims": 1_543_000, - "ctc_amount": 2_500_000_000, - "actc_claims": 687_000, - "actc_amount": 1_031_000_000, - }, - "IN": { - "eitc_claims": 543_000, - "eitc_amount": 1_456_000_000, - "ctc_claims": 876_000, - "ctc_amount": 1_419_000_000, - "actc_claims": 412_000, - "actc_amount": 618_000_000, - }, - "IA": { - "eitc_claims": 198_000, - "eitc_amount": 523_000_000, - "ctc_claims": 387_000, - "ctc_amount": 627_000_000, - "actc_claims": 176_000, - "actc_amount": 264_000_000, - }, - "KS": { - "eitc_claims": 198_000, - "eitc_amount": 534_000_000, - "ctc_claims": 376_000, - "ctc_amount": 609_000_000, - "actc_claims": 167_000, - "actc_amount": 251_000_000, - }, - "KY": { - "eitc_claims": 398_000, - "eitc_amount": 1_067_000_000, - "ctc_claims": 543_000, - "ctc_amount": 879_000_000, - "actc_claims": 287_000, - "actc_amount": 431_000_000, - }, - "LA": { - "eitc_claims": 498_000, - "eitc_amount": 1_345_000_000, - "ctc_claims": 654_000, - "ctc_amount": 1_059_000_000, - "actc_claims": 356_000, - "actc_amount": 534_000_000, - }, - "ME": { - "eitc_claims": 87_000, - "eitc_amount": 223_000_000, - "ctc_claims": 132_000, - "ctc_amount": 214_000_000, - "actc_claims": 65_000, - "actc_amount": 98_000_000, - }, - "MD": { - "eitc_claims": 432_000, - "eitc_amount": 1_156_000_000, - "ctc_claims": 687_000, - "ctc_amount": 1_113_000_000, - "actc_claims": 287_000, - "actc_amount": 431_000_000, - }, - "MA": { - "eitc_claims": 398_000, - "eitc_amount": 1_034_000_000, - "ctc_claims": 654_000, - "ctc_amount": 1_059_000_000, - "actc_claims": 265_000, - "actc_amount": 398_000_000, - }, - "MI": { - "eitc_claims": 765_000, - "eitc_amount": 2_045_000_000, - "ctc_claims": 1_176_000, - "ctc_amount": 1_905_000_000, - "actc_claims": 543_000, - "actc_amount": 815_000_000, - }, - "MN": { - "eitc_claims": 345_000, - "eitc_amount": 923_000_000, - "ctc_claims": 654_000, - "ctc_amount": 1_059_000_000, - "actc_claims": 276_000, - "actc_amount": 414_000_000, - }, - "MS": { - "eitc_claims": 343_000, - "eitc_amount": 934_000_000, - "ctc_claims": 432_000, - "ctc_amount": 700_000_000, - "actc_claims": 243_000, - "actc_amount": 365_000_000, - }, - "MO": { - "eitc_claims": 487_000, - "eitc_amount": 1_298_000_000, - "ctc_claims": 743_000, - "ctc_amount": 1_203_000_000, - "actc_claims": 354_000, - "actc_amount": 531_000_000, - }, - "MT": { - "eitc_claims": 76_000, - "eitc_amount": 198_000_000, - "ctc_claims": 132_000, - "ctc_amount": 214_000_000, - "actc_claims": 65_000, - "actc_amount": 98_000_000, - }, - "NE": { - "eitc_claims": 132_000, - "eitc_amount": 356_000_000, - "ctc_claims": 265_000, - "ctc_amount": 429_000_000, - "actc_claims": 121_000, - "actc_amount": 182_000_000, - }, - "NV": { - "eitc_claims": 265_000, - "eitc_amount": 712_000_000, - "ctc_claims": 412_000, - "ctc_amount": 667_000_000, - "actc_claims": 198_000, - "actc_amount": 297_000_000, - }, - "NH": { - "eitc_claims": 76_000, - "eitc_amount": 189_000_000, - "ctc_claims": 143_000, - "ctc_amount": 232_000_000, - "actc_claims": 54_000, - "actc_amount": 81_000_000, - }, - "NJ": { - "eitc_claims": 587_000, - "eitc_amount": 1_534_000_000, - "ctc_claims": 976_000, - "ctc_amount": 1_581_000_000, - "actc_claims": 398_000, - "actc_amount": 597_000_000, - }, - "NM": { - "eitc_claims": 198_000, - "eitc_amount": 534_000_000, - "ctc_claims": 287_000, - "ctc_amount": 465_000_000, - "actc_claims": 154_000, - "actc_amount": 231_000_000, - }, - "NY": { - "eitc_claims": 1_543_000, - "eitc_amount": 4_123_000_000, - "ctc_claims": 2_234_000, - "ctc_amount": 3_619_000_000, - "actc_claims": 987_000, - "actc_amount": 1_481_000_000, - }, - "NC": { - "eitc_claims": 876_000, - "eitc_amount": 2_367_000_000, - "ctc_claims": 1_287_000, - "ctc_amount": 2_085_000_000, - "actc_claims": 612_000, - "actc_amount": 918_000_000, - }, - "ND": { - "eitc_claims": 43_000, - "eitc_amount": 112_000_000, - "ctc_claims": 98_000, - "ctc_amount": 159_000_000, - "actc_claims": 43_000, - "actc_amount": 65_000_000, - }, - "OH": { - "eitc_claims": 943_000, - "eitc_amount": 2_523_000_000, - "ctc_claims": 1_387_000, - "ctc_amount": 2_247_000_000, - "actc_claims": 654_000, - "actc_amount": 981_000_000, - }, - "OK": { - "eitc_claims": 343_000, - "eitc_amount": 923_000_000, - "ctc_claims": 521_000, - "ctc_amount": 844_000_000, - "actc_claims": 265_000, - "actc_amount": 398_000_000, - }, - "OR": { - "eitc_claims": 298_000, - "eitc_amount": 798_000_000, - "ctc_claims": 476_000, - "ctc_amount": 771_000_000, - "actc_claims": 221_000, - "actc_amount": 332_000_000, - }, - "PA": { - "eitc_claims": 976_000, - "eitc_amount": 2_589_000_000, - "ctc_claims": 1_432_000, - "ctc_amount": 2_320_000_000, - "actc_claims": 654_000, - "actc_amount": 981_000_000, - }, - "RI": { - "eitc_claims": 76_000, - "eitc_amount": 198_000_000, - "ctc_claims": 109_000, - "ctc_amount": 177_000_000, - "actc_claims": 54_000, - "actc_amount": 81_000_000, - }, - "SC": { - "eitc_claims": 465_000, - "eitc_amount": 1_256_000_000, - "ctc_claims": 654_000, - "ctc_amount": 1_059_000_000, - "actc_claims": 332_000, - "actc_amount": 498_000_000, - }, - "SD": { - "eitc_claims": 54_000, - "eitc_amount": 145_000_000, - "ctc_claims": 121_000, - "ctc_amount": 196_000_000, - "actc_claims": 54_000, - "actc_amount": 81_000_000, - }, - "TN": { - "eitc_claims": 598_000, - "eitc_amount": 1_612_000_000, - "ctc_claims": 876_000, - "ctc_amount": 1_419_000_000, - "actc_claims": 432_000, - "actc_amount": 648_000_000, - }, - "TX": { - "eitc_claims": 2_654_000, - "eitc_amount": 7_234_000_000, - "ctc_claims": 4_321_000, - "ctc_amount": 6_998_000_000, - "actc_claims": 2_098_000, - "actc_amount": 3_147_000_000, - }, - "UT": { - "eitc_claims": 198_000, - "eitc_amount": 534_000_000, - "ctc_claims": 521_000, - "ctc_amount": 844_000_000, - "actc_claims": 232_000, - "actc_amount": 348_000_000, - }, - "VT": { - "eitc_claims": 43_000, - "eitc_amount": 109_000_000, - "ctc_claims": 65_000, - "ctc_amount": 105_000_000, - "actc_claims": 32_000, - "actc_amount": 48_000_000, - }, - "VA": { - "eitc_claims": 598_000, - "eitc_amount": 1_598_000_000, - "ctc_claims": 987_000, - "ctc_amount": 1_599_000_000, - "actc_claims": 432_000, - "actc_amount": 648_000_000, - }, - "WA": { - "eitc_claims": 487_000, - "eitc_amount": 1_312_000_000, - "ctc_claims": 876_000, - "ctc_amount": 1_419_000_000, - "actc_claims": 376_000, - "actc_amount": 564_000_000, - }, - "WV": { - "eitc_claims": 154_000, - "eitc_amount": 412_000_000, - "ctc_claims": 198_000, - "ctc_amount": 321_000_000, - "actc_claims": 109_000, - "actc_amount": 164_000_000, - }, - "WI": { - "eitc_claims": 387_000, - "eitc_amount": 1_034_000_000, - "ctc_claims": 687_000, - "ctc_amount": 1_113_000_000, - "actc_claims": 298_000, - "actc_amount": 447_000_000, - }, - "WY": { - "eitc_claims": 32_000, - "eitc_amount": 87_000_000, - "ctc_claims": 76_000, - "ctc_amount": 123_000_000, - "actc_claims": 32_000, - "actc_amount": 48_000_000, - }, - }, -} - -SOURCE_URLS = { - "eitc": "https://www.irs.gov/credits-deductions/individuals/earned-income-tax-credit/earned-income-tax-credit-statistics", - "ctc": "https://www.irs.gov/statistics/soi-tax-stats-state-data-fy-2022", -} - -# National EITC data by number of qualifying children (2021) -# Source: IRS SOI EITC Statistics -# https://www.irs.gov/credits-deductions/individuals/earned-income-tax-credit/earned-income-tax-credit-statistics -# -# Representative 2021 data (national distribution): -# - 0 children: ~7M claims, ~$1.5B (max credit ~$543, avg ~$214) -# - 1 child: ~8M claims, ~$20B (max credit ~$3,618, avg ~$2,500) -# - 2 children: ~6M claims, ~$20B (max credit ~$5,980, avg ~$3,333) -# - 3+ children: ~4M claims, ~$15B (max credit ~$6,728, avg ~$3,750) - -EITC_BY_CHILDREN_DATA = { - 2021: { - "0_children": { - "claims": 7_000_000, - "amount": 1_500_000_000, - }, - "1_child": { - "claims": 8_000_000, - "amount": 20_000_000_000, - }, - "2_children": { - "claims": 6_000_000, - "amount": 20_000_000_000, - }, - "3plus_children": { - "claims": 4_000_000, - "amount": 15_000_000_000, - }, - }, -} - -# Child count category definitions for strata -EITC_CHILD_CATEGORIES = { - "0_children": { - "name": "US EITC 0 Children", - "description": "EITC recipients with no qualifying children", - "constraint_operator": "==", - "constraint_value": "0", - }, - "1_child": { - "name": "US EITC 1 Child", - "description": "EITC recipients with one qualifying child", - "constraint_operator": "==", - "constraint_value": "1", - }, - "2_children": { - "name": "US EITC 2 Children", - "description": "EITC recipients with two qualifying children", - "constraint_operator": "==", - "constraint_value": "2", - }, - "3plus_children": { - "name": "US EITC 3+ Children", - "description": "EITC recipients with three or more qualifying children", - "constraint_operator": ">=", - "constraint_value": "3", - }, -} - -# National CTC data by number of qualifying children (2021) -# Source: IRS SOI Statistics -# https://www.irs.gov/statistics/soi-tax-stats-state-data-fy-2022 -# -# Representative 2021 data (national distribution): -# - 1 child: ~14M claims, ~$30B (avg ~$2,143 per claim) -# - 2 children: ~12M claims, ~$38B (avg ~$3,167 per claim) -# - 3+ children: ~9M claims, ~$32B (avg ~$3,556 per claim) -# Total: ~35M claims, ~$100B - -CTC_BY_CHILDREN_DATA = { - 2021: { - "1_child": { - "claims": 14_000_000, - "amount": 30_000_000_000, - }, - "2_children": { - "claims": 12_000_000, - "amount": 38_000_000_000, - }, - "3plus_children": { - "claims": 9_000_000, - "amount": 32_000_000_000, - }, - }, -} - -# CTC child count category definitions (CTC requires at least 1 child) -CTC_CHILD_CATEGORIES = { - "1_child": { - "name": "US CTC 1 Child", - "description": "CTC recipients with one qualifying child", - "constraint_operator": "==", - "constraint_value": "1", - }, - "2_children": { - "name": "US CTC 2 Children", - "description": "CTC recipients with two qualifying children", - "constraint_operator": "==", - "constraint_value": "2", - }, - "3plus_children": { - "name": "US CTC 3+ Children", - "description": "CTC recipients with three or more qualifying children", - "constraint_operator": ">=", - "constraint_value": "3", - }, -} - -# National ACTC data by number of qualifying children (2021) -# Source: IRS SOI Statistics -# https://www.irs.gov/statistics/soi-tax-stats-state-data-fy-2022 -# -# Representative 2021 data (national distribution): -# ACTC is the refundable portion of CTC, taken by lower-income filers -# - 1 child: ~7.5M claims, ~$10B (avg ~$1,333 per claim) -# - 2 children: ~6M claims, ~$11B (avg ~$1,833 per claim) -# - 3+ children: ~4.5M claims, ~$9B (avg ~$2,000 per claim) -# Total: ~18M claims, ~$30B - -ACTC_BY_CHILDREN_DATA = { - 2021: { - "1_child": { - "claims": 7_500_000, - "amount": 10_000_000_000, - }, - "2_children": { - "claims": 6_000_000, - "amount": 11_000_000_000, - }, - "3plus_children": { - "claims": 4_500_000, - "amount": 9_000_000_000, - }, - }, -} - -# ACTC child count category definitions (ACTC requires at least 1 child) -ACTC_CHILD_CATEGORIES = { - "1_child": { - "name": "US ACTC 1 Child", - "description": "ACTC recipients with one qualifying child", - "constraint_operator": "==", - "constraint_value": "1", - }, - "2_children": { - "name": "US ACTC 2 Children", - "description": "ACTC recipients with two qualifying children", - "constraint_operator": "==", - "constraint_value": "2", - }, - "3plus_children": { - "name": "US ACTC 3+ Children", - "description": "ACTC recipients with three or more qualifying children", - "constraint_operator": ">=", - "constraint_value": "3", - }, -} - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - # Check if exists - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - # Create new - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() # Get ID - - # Add constraints - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_soi_credits_targets(session: Session, years: list[int] | None = None): - """ - Load state-level tax credit targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SOI_CREDITS_DATA.keys()) - - for year in years: - if year not in SOI_CREDITS_DATA: - continue - - data = SOI_CREDITS_DATA[year] - - # Get or create national stratum (for parent relationship) - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create state-level strata and targets - for state_abbrev, state_data in data.items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - # Create state stratum - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} All Filers", - jurisdiction=Jurisdiction.US, - constraints=[ - ("is_tax_filer", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"All individual income tax returns filed in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="soi_states", - ) - - # Add EITC claims target - session.add( - Target( - stratum_id=state_stratum.id, - variable="eitc_claims", - period=year, - value=state_data["eitc_claims"], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="EITC Statistics", - source_url=SOURCE_URLS["eitc"], - ) - ) - - # Add EITC amount target - session.add( - Target( - stratum_id=state_stratum.id, - variable="eitc_amount", - period=year, - value=state_data["eitc_amount"], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="EITC Statistics", - source_url=SOURCE_URLS["eitc"], - ) - ) - - # Add CTC claims target - session.add( - Target( - stratum_id=state_stratum.id, - variable="ctc_claims", - period=year, - value=state_data["ctc_claims"], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="State Data FY", - source_url=SOURCE_URLS["ctc"], - ) - ) - - # Add CTC amount target - session.add( - Target( - stratum_id=state_stratum.id, - variable="ctc_amount", - period=year, - value=state_data["ctc_amount"], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="State Data FY", - source_url=SOURCE_URLS["ctc"], - ) - ) - - # Add ACTC claims target - session.add( - Target( - stratum_id=state_stratum.id, - variable="actc_claims", - period=year, - value=state_data["actc_claims"], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="State Data FY", - source_url=SOURCE_URLS["ctc"], - ) - ) - - # Add ACTC amount target - session.add( - Target( - stratum_id=state_stratum.id, - variable="actc_amount", - period=year, - value=state_data["actc_amount"], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="State Data FY", - source_url=SOURCE_URLS["ctc"], - ) - ) - - session.commit() - - -def load_eitc_by_children_targets(session: Session, years: list[int] | None = None): - """ - Load national EITC targets stratified by number of qualifying children. - - Creates strata for: - - 0 children - - 1 child - - 2 children - - 3+ children - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(EITC_BY_CHILDREN_DATA.keys()) - - for year in years: - if year not in EITC_BY_CHILDREN_DATA: - continue - - data = EITC_BY_CHILDREN_DATA[year] - - # Get or create national stratum (parent for child-count strata) - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create strata and targets for each child category - for category_key, category_data in data.items(): - category_def = EITC_CHILD_CATEGORIES[category_key] - - # Create stratum for this child category - stratum = get_or_create_stratum( - session, - name=category_def["name"], - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("is_tax_filer", "==", "1"), - ("eitc_qualifying_children", category_def["constraint_operator"], category_def["constraint_value"]), - ], - description=category_def["description"], - parent_id=national_stratum.id, - stratum_group_id="eitc_by_children", - ) - - # Add EITC claims target - session.add( - Target( - stratum_id=stratum.id, - variable="eitc_claims", - period=year, - value=category_data["claims"], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="EITC Statistics by Number of Qualifying Children", - source_url=SOURCE_URLS["eitc"], - ) - ) - - # Add EITC amount target - session.add( - Target( - stratum_id=stratum.id, - variable="eitc_amount", - period=year, - value=category_data["amount"], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="EITC Statistics by Number of Qualifying Children", - source_url=SOURCE_URLS["eitc"], - ) - ) - - session.commit() - - -def load_ctc_by_children_targets(session: Session, years: list[int] | None = None): - """ - Load national CTC targets stratified by number of qualifying children. - - Creates strata for: - - 1 child - - 2 children - - 3+ children - - Note: CTC requires at least 1 qualifying child (unlike EITC which has a 0-child category). - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(CTC_BY_CHILDREN_DATA.keys()) - - for year in years: - if year not in CTC_BY_CHILDREN_DATA: - continue - - data = CTC_BY_CHILDREN_DATA[year] - - # Get or create national stratum (parent for child-count strata) - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create strata and targets for each child category - for category_key, category_data in data.items(): - category_def = CTC_CHILD_CATEGORIES[category_key] - - # Create stratum for this child category - stratum = get_or_create_stratum( - session, - name=category_def["name"], - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("is_tax_filer", "==", "1"), - ("ctc_qualifying_children", category_def["constraint_operator"], category_def["constraint_value"]), - ], - description=category_def["description"], - parent_id=national_stratum.id, - stratum_group_id="ctc_by_children", - ) - - # Add CTC claims target - session.add( - Target( - stratum_id=stratum.id, - variable="ctc_claims", - period=year, - value=category_data["claims"], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="CTC Statistics by Number of Qualifying Children", - source_url=SOURCE_URLS["ctc"], - ) - ) - - # Add CTC amount target - session.add( - Target( - stratum_id=stratum.id, - variable="ctc_amount", - period=year, - value=category_data["amount"], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="CTC Statistics by Number of Qualifying Children", - source_url=SOURCE_URLS["ctc"], - ) - ) - - session.commit() - - -def load_actc_by_children_targets(session: Session, years: list[int] | None = None): - """ - Load national ACTC targets stratified by number of qualifying children. - - Creates strata for: - - 1 child - - 2 children - - 3+ children - - Note: ACTC (Additional Child Tax Credit) is the refundable portion of CTC, - claimed by lower-income filers who cannot use the full CTC against their tax liability. - Like CTC, it requires at least 1 qualifying child. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(ACTC_BY_CHILDREN_DATA.keys()) - - for year in years: - if year not in ACTC_BY_CHILDREN_DATA: - continue - - data = ACTC_BY_CHILDREN_DATA[year] - - # Get or create national stratum (parent for child-count strata) - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create strata and targets for each child category - for category_key, category_data in data.items(): - category_def = ACTC_CHILD_CATEGORIES[category_key] - - # Create stratum for this child category - stratum = get_or_create_stratum( - session, - name=category_def["name"], - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("is_tax_filer", "==", "1"), - ("actc_qualifying_children", category_def["constraint_operator"], category_def["constraint_value"]), - ], - description=category_def["description"], - parent_id=national_stratum.id, - stratum_group_id="actc_by_children", - ) - - # Add ACTC claims target - session.add( - Target( - stratum_id=stratum.id, - variable="actc_claims", - period=year, - value=category_data["claims"], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="ACTC Statistics by Number of Qualifying Children", - source_url=SOURCE_URLS["ctc"], - ) - ) - - # Add ACTC amount target - session.add( - Target( - stratum_id=stratum.id, - variable="actc_amount", - period=year, - value=category_data["amount"], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="ACTC Statistics by Number of Qualifying Children", - source_url=SOURCE_URLS["ctc"], - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the state-level tax credits ETL pipeline.""" - from pathlib import Path - - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_soi_credits_targets(session) - print(f"Loaded state-level SOI credits targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_soi_deductions.py b/db/etl_soi_deductions.py deleted file mode 100644 index abe7c3b0..00000000 --- a/db/etl_soi_deductions.py +++ /dev/null @@ -1,859 +0,0 @@ -""" -ETL for IRS Statistics of Income (SOI) state-level itemized deduction targets. - -Loads state-by-state deduction data from IRS SOI tables: -- SALT deduction (state and local taxes) -- Mortgage interest deduction -- Charitable contributions -- Medical expenses -- QBI deduction (qualified business income) - -Data source: -https://www.irs.gov/statistics/soi-tax-stats-historic-table-2 -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes for all 50 states + DC -STATE_FIPS = { - "AL": "01", - "AK": "02", - "AZ": "04", - "AR": "05", - "CA": "06", - "CO": "08", - "CT": "09", - "DE": "10", - "DC": "11", - "FL": "12", - "GA": "13", - "HI": "15", - "ID": "16", - "IL": "17", - "IN": "18", - "IA": "19", - "KS": "20", - "KY": "21", - "LA": "22", - "ME": "23", - "MD": "24", - "MA": "25", - "MI": "26", - "MN": "27", - "MS": "28", - "MO": "29", - "MT": "30", - "NE": "31", - "NV": "32", - "NH": "33", - "NJ": "34", - "NM": "35", - "NY": "36", - "NC": "37", - "ND": "38", - "OH": "39", - "OK": "40", - "OR": "41", - "PA": "42", - "RI": "44", - "SC": "45", - "SD": "46", - "TN": "47", - "TX": "48", - "UT": "49", - "VT": "50", - "VA": "51", - "WA": "53", - "WV": "54", - "WI": "55", - "WY": "56", -} - -# Deduction types tracked -DEDUCTION_TYPES = ["salt", "mortgage_interest", "charitable", "medical", "qbi"] - -SOURCE_URL = "https://www.irs.gov/statistics/soi-tax-stats-historic-table-2" - -# State-level deduction data for 2021 -# Representative data by state, scaled by population and income levels: -# - SALT: ~12M claims nationally, ~$150B total -# - Mortgage interest: ~10M claims, ~$80B -# - Charitable: ~25M claims, ~$50B -# - Medical: ~5M claims, ~$30B -# - QBI: ~10M claims, ~$70B - -SOI_DEDUCTIONS_DATA = { - 2021: { - "AL": { - "salt_claims": 132_000, - "salt_amount": 1_650_000_000, - "mortgage_interest_claims": 98_000, - "mortgage_interest_amount": 780_000_000, - "charitable_claims": 398_000, - "charitable_amount": 798_000_000, - "medical_claims": 65_000, - "medical_amount": 390_000_000, - "qbi_claims": 143_000, - "qbi_amount": 1_001_000_000, - }, - "AK": { - "salt_claims": 21_000, - "salt_amount": 262_000_000, - "mortgage_interest_claims": 18_000, - "mortgage_interest_amount": 144_000_000, - "charitable_claims": 65_000, - "charitable_amount": 130_000_000, - "medical_claims": 11_000, - "medical_amount": 66_000_000, - "qbi_claims": 28_000, - "qbi_amount": 196_000_000, - }, - "AZ": { - "salt_claims": 234_000, - "salt_amount": 2_925_000_000, - "mortgage_interest_claims": 198_000, - "mortgage_interest_amount": 1_584_000_000, - "charitable_claims": 598_000, - "charitable_amount": 1_196_000_000, - "medical_claims": 98_000, - "medical_amount": 588_000_000, - "qbi_claims": 232_000, - "qbi_amount": 1_624_000_000, - }, - "AR": { - "salt_claims": 76_000, - "salt_amount": 950_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 243_000, - "charitable_amount": 486_000_000, - "medical_claims": 43_000, - "medical_amount": 258_000_000, - "qbi_claims": 98_000, - "qbi_amount": 686_000_000, - }, - "CA": { - "salt_claims": 1_876_000, - "salt_amount": 23_450_000_000, - "mortgage_interest_claims": 1_543_000, - "mortgage_interest_amount": 12_344_000_000, - "charitable_claims": 3_654_000, - "charitable_amount": 7_308_000_000, - "medical_claims": 543_000, - "medical_amount": 3_258_000_000, - "qbi_claims": 1_234_000, - "qbi_amount": 8_638_000_000, - }, - "CO": { - "salt_claims": 234_000, - "salt_amount": 2_925_000_000, - "mortgage_interest_claims": 198_000, - "mortgage_interest_amount": 1_584_000_000, - "charitable_claims": 543_000, - "charitable_amount": 1_086_000_000, - "medical_claims": 76_000, - "medical_amount": 456_000_000, - "qbi_claims": 198_000, - "qbi_amount": 1_386_000_000, - }, - "CT": { - "salt_claims": 287_000, - "salt_amount": 3_588_000_000, - "mortgage_interest_claims": 198_000, - "mortgage_interest_amount": 1_584_000_000, - "charitable_claims": 343_000, - "charitable_amount": 686_000_000, - "medical_claims": 54_000, - "medical_amount": 324_000_000, - "qbi_claims": 132_000, - "qbi_amount": 924_000_000, - }, - "DE": { - "salt_claims": 54_000, - "salt_amount": 675_000_000, - "mortgage_interest_claims": 43_000, - "mortgage_interest_amount": 344_000_000, - "charitable_claims": 98_000, - "charitable_amount": 196_000_000, - "medical_claims": 15_000, - "medical_amount": 90_000_000, - "qbi_claims": 32_000, - "qbi_amount": 224_000_000, - }, - "DC": { - "salt_claims": 65_000, - "salt_amount": 813_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 76_000, - "charitable_amount": 152_000_000, - "medical_claims": 11_000, - "medical_amount": 66_000_000, - "qbi_claims": 28_000, - "qbi_amount": 196_000_000, - }, - "FL": { - "salt_claims": 876_000, - "salt_amount": 10_950_000_000, - "mortgage_interest_claims": 765_000, - "mortgage_interest_amount": 6_120_000_000, - "charitable_claims": 1_876_000, - "charitable_amount": 3_752_000_000, - "medical_claims": 298_000, - "medical_amount": 1_788_000_000, - "qbi_claims": 654_000, - "qbi_amount": 4_578_000_000, - }, - "GA": { - "salt_claims": 398_000, - "salt_amount": 4_975_000_000, - "mortgage_interest_claims": 343_000, - "mortgage_interest_amount": 2_744_000_000, - "charitable_claims": 876_000, - "charitable_amount": 1_752_000_000, - "medical_claims": 132_000, - "medical_amount": 792_000_000, - "qbi_claims": 343_000, - "qbi_amount": 2_401_000_000, - }, - "HI": { - "salt_claims": 65_000, - "salt_amount": 813_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 132_000, - "charitable_amount": 264_000_000, - "medical_claims": 21_000, - "medical_amount": 126_000_000, - "qbi_claims": 43_000, - "qbi_amount": 301_000_000, - }, - "ID": { - "salt_claims": 54_000, - "salt_amount": 675_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 176_000, - "charitable_amount": 352_000_000, - "medical_claims": 28_000, - "medical_amount": 168_000_000, - "qbi_claims": 65_000, - "qbi_amount": 455_000_000, - }, - "IL": { - "salt_claims": 654_000, - "salt_amount": 8_175_000_000, - "mortgage_interest_claims": 487_000, - "mortgage_interest_amount": 3_896_000_000, - "charitable_claims": 1_176_000, - "charitable_amount": 2_352_000_000, - "medical_claims": 176_000, - "medical_amount": 1_056_000_000, - "qbi_claims": 398_000, - "qbi_amount": 2_786_000_000, - }, - "IN": { - "salt_claims": 198_000, - "salt_amount": 2_475_000_000, - "mortgage_interest_claims": 154_000, - "mortgage_interest_amount": 1_232_000_000, - "charitable_claims": 543_000, - "charitable_amount": 1_086_000_000, - "medical_claims": 87_000, - "medical_amount": 522_000_000, - "qbi_claims": 198_000, - "qbi_amount": 1_386_000_000, - }, - "IA": { - "salt_claims": 98_000, - "salt_amount": 1_225_000_000, - "mortgage_interest_claims": 76_000, - "mortgage_interest_amount": 608_000_000, - "charitable_claims": 287_000, - "charitable_amount": 574_000_000, - "medical_claims": 43_000, - "medical_amount": 258_000_000, - "qbi_claims": 109_000, - "qbi_amount": 763_000_000, - }, - "KS": { - "salt_claims": 98_000, - "salt_amount": 1_225_000_000, - "mortgage_interest_claims": 76_000, - "mortgage_interest_amount": 608_000_000, - "charitable_claims": 265_000, - "charitable_amount": 530_000_000, - "medical_claims": 39_000, - "medical_amount": 234_000_000, - "qbi_claims": 98_000, - "qbi_amount": 686_000_000, - }, - "KY": { - "salt_claims": 132_000, - "salt_amount": 1_650_000_000, - "mortgage_interest_claims": 98_000, - "mortgage_interest_amount": 784_000_000, - "charitable_claims": 376_000, - "charitable_amount": 752_000_000, - "medical_claims": 65_000, - "medical_amount": 390_000_000, - "qbi_claims": 132_000, - "qbi_amount": 924_000_000, - }, - "LA": { - "salt_claims": 132_000, - "salt_amount": 1_650_000_000, - "mortgage_interest_claims": 98_000, - "mortgage_interest_amount": 784_000_000, - "charitable_claims": 398_000, - "charitable_amount": 796_000_000, - "medical_claims": 65_000, - "medical_amount": 390_000_000, - "qbi_claims": 143_000, - "qbi_amount": 1_001_000_000, - }, - "ME": { - "salt_claims": 43_000, - "salt_amount": 538_000_000, - "mortgage_interest_claims": 32_000, - "mortgage_interest_amount": 256_000_000, - "charitable_claims": 132_000, - "charitable_amount": 264_000_000, - "medical_claims": 21_000, - "medical_amount": 126_000_000, - "qbi_claims": 43_000, - "qbi_amount": 301_000_000, - }, - "MD": { - "salt_claims": 398_000, - "salt_amount": 4_975_000_000, - "mortgage_interest_claims": 287_000, - "mortgage_interest_amount": 2_296_000_000, - "charitable_claims": 543_000, - "charitable_amount": 1_086_000_000, - "medical_claims": 76_000, - "medical_amount": 456_000_000, - "qbi_claims": 176_000, - "qbi_amount": 1_232_000_000, - }, - "MA": { - "salt_claims": 432_000, - "salt_amount": 5_400_000_000, - "mortgage_interest_claims": 298_000, - "mortgage_interest_amount": 2_384_000_000, - "charitable_claims": 654_000, - "charitable_amount": 1_308_000_000, - "medical_claims": 87_000, - "medical_amount": 522_000_000, - "qbi_claims": 198_000, - "qbi_amount": 1_386_000_000, - }, - "MI": { - "salt_claims": 343_000, - "salt_amount": 4_288_000_000, - "mortgage_interest_claims": 265_000, - "mortgage_interest_amount": 2_120_000_000, - "charitable_claims": 876_000, - "charitable_amount": 1_752_000_000, - "medical_claims": 143_000, - "medical_amount": 858_000_000, - "qbi_claims": 298_000, - "qbi_amount": 2_086_000_000, - }, - "MN": { - "salt_claims": 243_000, - "salt_amount": 3_038_000_000, - "mortgage_interest_claims": 187_000, - "mortgage_interest_amount": 1_496_000_000, - "charitable_claims": 543_000, - "charitable_amount": 1_086_000_000, - "medical_claims": 76_000, - "medical_amount": 456_000_000, - "qbi_claims": 176_000, - "qbi_amount": 1_232_000_000, - }, - "MS": { - "salt_claims": 65_000, - "salt_amount": 813_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 243_000, - "charitable_amount": 486_000_000, - "medical_claims": 43_000, - "medical_amount": 258_000_000, - "qbi_claims": 87_000, - "qbi_amount": 609_000_000, - }, - "MO": { - "salt_claims": 198_000, - "salt_amount": 2_475_000_000, - "mortgage_interest_claims": 154_000, - "mortgage_interest_amount": 1_232_000_000, - "charitable_claims": 543_000, - "charitable_amount": 1_086_000_000, - "medical_claims": 87_000, - "medical_amount": 522_000_000, - "qbi_claims": 187_000, - "qbi_amount": 1_309_000_000, - }, - "MT": { - "salt_claims": 32_000, - "salt_amount": 400_000_000, - "mortgage_interest_claims": 28_000, - "mortgage_interest_amount": 224_000_000, - "charitable_claims": 109_000, - "charitable_amount": 218_000_000, - "medical_claims": 18_000, - "medical_amount": 108_000_000, - "qbi_claims": 43_000, - "qbi_amount": 301_000_000, - }, - "NE": { - "salt_claims": 65_000, - "salt_amount": 813_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 187_000, - "charitable_amount": 374_000_000, - "medical_claims": 28_000, - "medical_amount": 168_000_000, - "qbi_claims": 65_000, - "qbi_amount": 455_000_000, - }, - "NV": { - "salt_claims": 109_000, - "salt_amount": 1_363_000_000, - "mortgage_interest_claims": 98_000, - "mortgage_interest_amount": 784_000_000, - "charitable_claims": 265_000, - "charitable_amount": 530_000_000, - "medical_claims": 43_000, - "medical_amount": 258_000_000, - "qbi_claims": 109_000, - "qbi_amount": 763_000_000, - }, - "NH": { - "salt_claims": 65_000, - "salt_amount": 813_000_000, - "mortgage_interest_claims": 54_000, - "mortgage_interest_amount": 432_000_000, - "charitable_claims": 132_000, - "charitable_amount": 264_000_000, - "medical_claims": 21_000, - "medical_amount": 126_000_000, - "qbi_claims": 54_000, - "qbi_amount": 378_000_000, - }, - "NJ": { - "salt_claims": 654_000, - "salt_amount": 8_175_000_000, - "mortgage_interest_claims": 432_000, - "mortgage_interest_amount": 3_456_000_000, - "charitable_claims": 787_000, - "charitable_amount": 1_574_000_000, - "medical_claims": 109_000, - "medical_amount": 654_000_000, - "qbi_claims": 265_000, - "qbi_amount": 1_855_000_000, - }, - "NM": { - "salt_claims": 54_000, - "salt_amount": 675_000_000, - "mortgage_interest_claims": 43_000, - "mortgage_interest_amount": 344_000_000, - "charitable_claims": 176_000, - "charitable_amount": 352_000_000, - "medical_claims": 32_000, - "medical_amount": 192_000_000, - "qbi_claims": 65_000, - "qbi_amount": 455_000_000, - }, - "NY": { - "salt_claims": 1_234_000, - "salt_amount": 15_425_000_000, - "mortgage_interest_claims": 876_000, - "mortgage_interest_amount": 7_008_000_000, - "charitable_claims": 1_765_000, - "charitable_amount": 3_530_000_000, - "medical_claims": 265_000, - "medical_amount": 1_590_000_000, - "qbi_claims": 543_000, - "qbi_amount": 3_801_000_000, - }, - "NC": { - "salt_claims": 387_000, - "salt_amount": 4_838_000_000, - "mortgage_interest_claims": 332_000, - "mortgage_interest_amount": 2_656_000_000, - "charitable_claims": 876_000, - "charitable_amount": 1_752_000_000, - "medical_claims": 132_000, - "medical_amount": 792_000_000, - "qbi_claims": 298_000, - "qbi_amount": 2_086_000_000, - }, - "ND": { - "salt_claims": 21_000, - "salt_amount": 263_000_000, - "mortgage_interest_claims": 18_000, - "mortgage_interest_amount": 144_000_000, - "charitable_claims": 65_000, - "charitable_amount": 130_000_000, - "medical_claims": 11_000, - "medical_amount": 66_000_000, - "qbi_claims": 28_000, - "qbi_amount": 196_000_000, - }, - "OH": { - "salt_claims": 398_000, - "salt_amount": 4_975_000_000, - "mortgage_interest_claims": 298_000, - "mortgage_interest_amount": 2_384_000_000, - "charitable_claims": 987_000, - "charitable_amount": 1_974_000_000, - "medical_claims": 165_000, - "medical_amount": 990_000_000, - "qbi_claims": 343_000, - "qbi_amount": 2_401_000_000, - }, - "OK": { - "salt_claims": 98_000, - "salt_amount": 1_225_000_000, - "mortgage_interest_claims": 76_000, - "mortgage_interest_amount": 608_000_000, - "charitable_claims": 343_000, - "charitable_amount": 686_000_000, - "medical_claims": 54_000, - "medical_amount": 324_000_000, - "qbi_claims": 121_000, - "qbi_amount": 847_000_000, - }, - "OR": { - "salt_claims": 154_000, - "salt_amount": 1_925_000_000, - "mortgage_interest_claims": 132_000, - "mortgage_interest_amount": 1_056_000_000, - "charitable_claims": 398_000, - "charitable_amount": 796_000_000, - "medical_claims": 65_000, - "medical_amount": 390_000_000, - "qbi_claims": 143_000, - "qbi_amount": 1_001_000_000, - }, - "PA": { - "salt_claims": 543_000, - "salt_amount": 6_788_000_000, - "mortgage_interest_claims": 387_000, - "mortgage_interest_amount": 3_096_000_000, - "charitable_claims": 1_087_000, - "charitable_amount": 2_174_000_000, - "medical_claims": 176_000, - "medical_amount": 1_056_000_000, - "qbi_claims": 376_000, - "qbi_amount": 2_632_000_000, - }, - "RI": { - "salt_claims": 54_000, - "salt_amount": 675_000_000, - "mortgage_interest_claims": 43_000, - "mortgage_interest_amount": 344_000_000, - "charitable_claims": 98_000, - "charitable_amount": 196_000_000, - "medical_claims": 15_000, - "medical_amount": 90_000_000, - "qbi_claims": 32_000, - "qbi_amount": 224_000_000, - }, - "SC": { - "salt_claims": 176_000, - "salt_amount": 2_200_000_000, - "mortgage_interest_claims": 154_000, - "mortgage_interest_amount": 1_232_000_000, - "charitable_claims": 465_000, - "charitable_amount": 930_000_000, - "medical_claims": 76_000, - "medical_amount": 456_000_000, - "qbi_claims": 154_000, - "qbi_amount": 1_078_000_000, - }, - "SD": { - "salt_claims": 21_000, - "salt_amount": 263_000_000, - "mortgage_interest_claims": 21_000, - "mortgage_interest_amount": 168_000_000, - "charitable_claims": 87_000, - "charitable_amount": 174_000_000, - "medical_claims": 14_000, - "medical_amount": 84_000_000, - "qbi_claims": 32_000, - "qbi_amount": 224_000_000, - }, - "TN": { - "salt_claims": 198_000, - "salt_amount": 2_475_000_000, - "mortgage_interest_claims": 176_000, - "mortgage_interest_amount": 1_408_000_000, - "charitable_claims": 598_000, - "charitable_amount": 1_196_000_000, - "medical_claims": 98_000, - "medical_amount": 588_000_000, - "qbi_claims": 209_000, - "qbi_amount": 1_463_000_000, - }, - "TX": { - "salt_claims": 876_000, - "salt_amount": 10_950_000_000, - "mortgage_interest_claims": 765_000, - "mortgage_interest_amount": 6_120_000_000, - "charitable_claims": 2_098_000, - "charitable_amount": 4_196_000_000, - "medical_claims": 332_000, - "medical_amount": 1_992_000_000, - "qbi_claims": 876_000, - "qbi_amount": 6_132_000_000, - }, - "UT": { - "salt_claims": 98_000, - "salt_amount": 1_225_000_000, - "mortgage_interest_claims": 98_000, - "mortgage_interest_amount": 784_000_000, - "charitable_claims": 343_000, - "charitable_amount": 686_000_000, - "medical_claims": 43_000, - "medical_amount": 258_000_000, - "qbi_claims": 109_000, - "qbi_amount": 763_000_000, - }, - "VT": { - "salt_claims": 28_000, - "salt_amount": 350_000_000, - "mortgage_interest_claims": 21_000, - "mortgage_interest_amount": 168_000_000, - "charitable_claims": 65_000, - "charitable_amount": 130_000_000, - "medical_claims": 11_000, - "medical_amount": 66_000_000, - "qbi_claims": 21_000, - "qbi_amount": 147_000_000, - }, - "VA": { - "salt_claims": 487_000, - "salt_amount": 6_088_000_000, - "mortgage_interest_claims": 376_000, - "mortgage_interest_amount": 3_008_000_000, - "charitable_claims": 765_000, - "charitable_amount": 1_530_000_000, - "medical_claims": 109_000, - "medical_amount": 654_000_000, - "qbi_claims": 265_000, - "qbi_amount": 1_855_000_000, - }, - "WA": { - "salt_claims": 287_000, - "salt_amount": 3_588_000_000, - "mortgage_interest_claims": 265_000, - "mortgage_interest_amount": 2_120_000_000, - "charitable_claims": 654_000, - "charitable_amount": 1_308_000_000, - "medical_claims": 98_000, - "medical_amount": 588_000_000, - "qbi_claims": 243_000, - "qbi_amount": 1_701_000_000, - }, - "WV": { - "salt_claims": 43_000, - "salt_amount": 538_000_000, - "mortgage_interest_claims": 32_000, - "mortgage_interest_amount": 256_000_000, - "charitable_claims": 143_000, - "charitable_amount": 286_000_000, - "medical_claims": 28_000, - "medical_amount": 168_000_000, - "qbi_claims": 54_000, - "qbi_amount": 378_000_000, - }, - "WI": { - "salt_claims": 221_000, - "salt_amount": 2_763_000_000, - "mortgage_interest_claims": 176_000, - "mortgage_interest_amount": 1_408_000_000, - "charitable_claims": 543_000, - "charitable_amount": 1_086_000_000, - "medical_claims": 87_000, - "medical_amount": 522_000_000, - "qbi_claims": 176_000, - "qbi_amount": 1_232_000_000, - }, - "WY": { - "salt_claims": 15_000, - "salt_amount": 188_000_000, - "mortgage_interest_claims": 14_000, - "mortgage_interest_amount": 112_000_000, - "charitable_claims": 54_000, - "charitable_amount": 108_000_000, - "medical_claims": 9_000, - "medical_amount": 54_000_000, - "qbi_claims": 21_000, - "qbi_amount": 147_000_000, - }, - }, -} - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - # Check if exists - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - # Create new - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() # Get ID - - # Add constraints - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_soi_deductions_targets(session: Session, years: list[int] | None = None): - """ - Load state-level itemized deduction targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SOI_DEDUCTIONS_DATA.keys()) - - for year in years: - if year not in SOI_DEDUCTIONS_DATA: - continue - - data = SOI_DEDUCTIONS_DATA[year] - - # Get or create national stratum (for parent relationship) - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create state-level strata and targets - for state_abbrev, state_data in data.items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - # Create state stratum - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} All Filers", - jurisdiction=Jurisdiction.US, - constraints=[ - ("is_tax_filer", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"All individual income tax returns filed in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="soi_states", - ) - - # Add targets for each deduction type - for deduction_type in DEDUCTION_TYPES: - claims_key = f"{deduction_type}_claims" - amount_key = f"{deduction_type}_amount" - - # Add claims target - session.add( - Target( - stratum_id=state_stratum.id, - variable=claims_key, - period=year, - value=state_data[claims_key], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="SOI Individual Returns - Itemized Deductions", - source_url=SOURCE_URL, - ) - ) - - # Add amount target - session.add( - Target( - stratum_id=state_stratum.id, - variable=amount_key, - period=year, - value=state_data[amount_key], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="SOI Individual Returns - Itemized Deductions", - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the state-level deductions ETL pipeline.""" - from pathlib import Path - - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_soi_deductions_targets(session) - print(f"Loaded state-level SOI deduction targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_soi_income_sources.py b/db/etl_soi_income_sources.py deleted file mode 100644 index 1264e1ef..00000000 --- a/db/etl_soi_income_sources.py +++ /dev/null @@ -1,1582 +0,0 @@ -""" -ETL for IRS Statistics of Income (SOI) income by source targets. - -Loads national and state-level income source data from IRS SOI tables including: -- Interest income (taxable and tax-exempt) -- Dividend income (ordinary and qualified) -- Capital gains (short-term and long-term) -- State tax refunds -- Alimony received -- Business income/loss (Schedule C) -- Rental, royalty, partnership, S-corp income (Schedule E) - -Data source: https://www.irs.gov/statistics/soi-tax-stats-individual-income-tax-statistics -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes for all 50 states + DC -STATE_FIPS = { - "AL": "01", - "AK": "02", - "AZ": "04", - "AR": "05", - "CA": "06", - "CO": "08", - "CT": "09", - "DE": "10", - "DC": "11", - "FL": "12", - "GA": "13", - "HI": "15", - "ID": "16", - "IL": "17", - "IN": "18", - "IA": "19", - "KS": "20", - "KY": "21", - "LA": "22", - "ME": "23", - "MD": "24", - "MA": "25", - "MI": "26", - "MN": "27", - "MS": "28", - "MO": "29", - "MT": "30", - "NE": "31", - "NV": "32", - "NH": "33", - "NJ": "34", - "NM": "35", - "NY": "36", - "NC": "37", - "ND": "38", - "OH": "39", - "OK": "40", - "OR": "41", - "PA": "42", - "RI": "44", - "SC": "45", - "SD": "46", - "TN": "47", - "TX": "48", - "UT": "49", - "VT": "50", - "VA": "51", - "WA": "53", - "WV": "54", - "WI": "55", - "WY": "56", -} - -# Detailed income source types to track (aligned with IRS SOI Table 1.4) -# Each source has _returns (count) and _amount (dollars) variants -INCOME_SOURCES = [ - # Interest income - "taxable_interest", - "tax_exempt_interest", - # Dividend income - "ordinary_dividends", - "qualified_dividends", - # Capital gains/losses - "short_term_capital_gains", - "long_term_capital_gains", - # Other income - "state_local_refunds", - "alimony_received", - # Schedule C business income - "schedule_c_income", - # Schedule E income (rental, royalty, partnership, S-corp) - "rental_royalty_income", - "partnership_scorp_income", -] - -# National income by source data for 2021 -# Based on IRS SOI Publication 1304, Table 1.4 -# Source: https://www.irs.gov/statistics/soi-tax-stats-individual-income-tax-returns -# Values are from 2021 tax year (most recent complete data) -# Returns = number of returns reporting this income type -# Amount = total dollars of this income type reported (can be negative for losses) - -SOI_NATIONAL_INCOME_SOURCES_DATA = { - 2021: { - # Interest income (IRS Form 1040 lines 2a and 2b) - "taxable_interest_returns": 54_876_543, - "taxable_interest_amount": 276_543_000_000, - "tax_exempt_interest_returns": 8_765_432, - "tax_exempt_interest_amount": 87_654_000_000, - # Dividend income (IRS Form 1040 lines 3a and 3b) - "ordinary_dividends_returns": 32_109_876, - "ordinary_dividends_amount": 432_109_000_000, - "qualified_dividends_returns": 28_765_432, - "qualified_dividends_amount": 298_765_000_000, - # Capital gains (Schedule D) - "short_term_capital_gains_returns": 12_345_678, - "short_term_capital_gains_amount": 187_654_000_000, - "long_term_capital_gains_returns": 21_098_765, - "long_term_capital_gains_amount": 876_543_000_000, - # Other income - "state_local_refunds_returns": 18_765_432, - "state_local_refunds_amount": 32_109_000_000, - "alimony_received_returns": 234_567, - "alimony_received_amount": 4_321_000_000, - # Schedule C business income (net profit/loss from business) - "schedule_c_income_returns": 28_765_432, - "schedule_c_income_amount": 398_765_000_000, - # Schedule E income - "rental_royalty_income_returns": 12_345_678, - "rental_royalty_income_amount": 98_765_000_000, - "partnership_scorp_income_returns": 9_876_543, - "partnership_scorp_income_amount": 876_543_000_000, - }, -} - -# State-level income by source data for 2021 -# Based on IRS SOI Individual Income Tax State Data -# Source: https://www.irs.gov/statistics/soi-tax-stats-individual-income-tax-state-data -# Values are representative 2021 data scaled by state population and income levels -# Returns = number of returns reporting this income type -# Amount = total dollars of this income type reported - -SOI_STATE_INCOME_SOURCES_DATA = { - 2021: { - "AL": { - "taxable_interest_returns": 543_210, - "taxable_interest_amount": 2_345_000_000, - "tax_exempt_interest_returns": 87_654, - "tax_exempt_interest_amount": 876_000_000, - "ordinary_dividends_returns": 234_567, - "ordinary_dividends_amount": 4_567_000_000, - "qualified_dividends_returns": 198_765, - "qualified_dividends_amount": 3_210_000_000, - "short_term_capital_gains_returns": 98_765, - "short_term_capital_gains_amount": 1_234_000_000, - "long_term_capital_gains_returns": 154_321, - "long_term_capital_gains_amount": 7_654_000_000, - "state_local_refunds_returns": 187_654, - "state_local_refunds_amount": 321_000_000, - "alimony_received_returns": 2_345, - "alimony_received_amount": 43_000_000, - "schedule_c_income_returns": 198_765, - "schedule_c_income_amount": 12_345_000_000, - "rental_royalty_income_returns": 87_654, - "rental_royalty_income_amount": 876_000_000, - "partnership_scorp_income_returns": 65_432, - "partnership_scorp_income_amount": 8_765_000_000, - }, - "AK": { - "taxable_interest_returns": 98_765, - "taxable_interest_amount": 654_000_000, - "tax_exempt_interest_returns": 21_098, - "tax_exempt_interest_amount": 234_000_000, - "ordinary_dividends_returns": 45_678, - "ordinary_dividends_amount": 1_234_000_000, - "qualified_dividends_returns": 38_765, - "qualified_dividends_amount": 876_000_000, - "short_term_capital_gains_returns": 21_098, - "short_term_capital_gains_amount": 321_000_000, - "long_term_capital_gains_returns": 38_765, - "long_term_capital_gains_amount": 2_109_000_000, - "state_local_refunds_returns": 32_109, - "state_local_refunds_amount": 54_000_000, - "alimony_received_returns": 432, - "alimony_received_amount": 8_000_000, - "schedule_c_income_returns": 32_109, - "schedule_c_income_amount": 2_876_000_000, - "rental_royalty_income_returns": 21_098, - "rental_royalty_income_amount": 432_000_000, - "partnership_scorp_income_returns": 15_432, - "partnership_scorp_income_amount": 2_345_000_000, - }, - "AZ": { - "taxable_interest_returns": 765_432, - "taxable_interest_amount": 4_321_000_000, - "tax_exempt_interest_returns": 132_109, - "tax_exempt_interest_amount": 1_432_000_000, - "ordinary_dividends_returns": 387_654, - "ordinary_dividends_amount": 8_765_000_000, - "qualified_dividends_returns": 321_098, - "qualified_dividends_amount": 6_543_000_000, - "short_term_capital_gains_returns": 154_321, - "short_term_capital_gains_amount": 2_345_000_000, - "long_term_capital_gains_returns": 265_432, - "long_term_capital_gains_amount": 16_543_000_000, - "state_local_refunds_returns": 298_765, - "state_local_refunds_amount": 543_000_000, - "alimony_received_returns": 3_876, - "alimony_received_amount": 76_000_000, - "schedule_c_income_returns": 298_765, - "schedule_c_income_amount": 21_098_000_000, - "rental_royalty_income_returns": 143_210, - "rental_royalty_income_amount": 1_654_000_000, - "partnership_scorp_income_returns": 109_876, - "partnership_scorp_income_amount": 18_765_000_000, - }, - "AR": { - "taxable_interest_returns": 321_098, - "taxable_interest_amount": 1_234_000_000, - "tax_exempt_interest_returns": 54_321, - "tax_exempt_interest_amount": 432_000_000, - "ordinary_dividends_returns": 143_210, - "ordinary_dividends_amount": 2_345_000_000, - "qualified_dividends_returns": 121_098, - "qualified_dividends_amount": 1_654_000_000, - "short_term_capital_gains_returns": 54_321, - "short_term_capital_gains_amount": 654_000_000, - "long_term_capital_gains_returns": 98_765, - "long_term_capital_gains_amount": 4_876_000_000, - "state_local_refunds_returns": 109_876, - "state_local_refunds_amount": 187_000_000, - "alimony_received_returns": 1_432, - "alimony_received_amount": 26_000_000, - "schedule_c_income_returns": 132_109, - "schedule_c_income_amount": 7_654_000_000, - "rental_royalty_income_returns": 54_321, - "rental_royalty_income_amount": 543_000_000, - "partnership_scorp_income_returns": 43_210, - "partnership_scorp_income_amount": 5_432_000_000, - }, - "CA": { - "taxable_interest_returns": 4_321_098, - "taxable_interest_amount": 45_678_000_000, - "tax_exempt_interest_returns": 876_543, - "tax_exempt_interest_amount": 12_345_000_000, - "ordinary_dividends_returns": 2_345_678, - "ordinary_dividends_amount": 98_765_000_000, - "qualified_dividends_returns": 1_987_654, - "qualified_dividends_amount": 76_543_000_000, - "short_term_capital_gains_returns": 987_654, - "short_term_capital_gains_amount": 32_109_000_000, - "long_term_capital_gains_returns": 1_654_321, - "long_term_capital_gains_amount": 212_345_000_000, - "state_local_refunds_returns": 1_876_543, - "state_local_refunds_amount": 4_321_000_000, - "alimony_received_returns": 23_456, - "alimony_received_amount": 543_000_000, - "schedule_c_income_returns": 1_543_210, - "schedule_c_income_amount": 187_654_000_000, - "rental_royalty_income_returns": 654_321, - "rental_royalty_income_amount": 12_345_000_000, - "partnership_scorp_income_returns": 543_210, - "partnership_scorp_income_amount": 234_567_000_000, - }, - "CO": { - "taxable_interest_returns": 654_321, - "taxable_interest_amount": 5_432_000_000, - "tax_exempt_interest_returns": 121_098, - "tax_exempt_interest_amount": 1_543_000_000, - "ordinary_dividends_returns": 398_765, - "ordinary_dividends_amount": 12_345_000_000, - "qualified_dividends_returns": 332_109, - "qualified_dividends_amount": 9_876_000_000, - "short_term_capital_gains_returns": 143_210, - "short_term_capital_gains_amount": 2_876_000_000, - "long_term_capital_gains_returns": 254_321, - "long_term_capital_gains_amount": 21_098_000_000, - "state_local_refunds_returns": 265_432, - "state_local_refunds_amount": 543_000_000, - "alimony_received_returns": 3_210, - "alimony_received_amount": 65_000_000, - "schedule_c_income_returns": 265_432, - "schedule_c_income_amount": 23_456_000_000, - "rental_royalty_income_returns": 121_098, - "rental_royalty_income_amount": 1_876_000_000, - "partnership_scorp_income_returns": 98_765, - "partnership_scorp_income_amount": 23_456_000_000, - }, - "CT": { - "taxable_interest_returns": 476_543, - "taxable_interest_amount": 8_765_000_000, - "tax_exempt_interest_returns": 109_876, - "tax_exempt_interest_amount": 2_345_000_000, - "ordinary_dividends_returns": 298_765, - "ordinary_dividends_amount": 18_765_000_000, - "qualified_dividends_returns": 254_321, - "qualified_dividends_amount": 14_321_000_000, - "short_term_capital_gains_returns": 109_876, - "short_term_capital_gains_amount": 3_456_000_000, - "long_term_capital_gains_returns": 187_654, - "long_term_capital_gains_amount": 25_432_000_000, - "state_local_refunds_returns": 154_321, - "state_local_refunds_amount": 387_000_000, - "alimony_received_returns": 2_543, - "alimony_received_amount": 54_000_000, - "schedule_c_income_returns": 154_321, - "schedule_c_income_amount": 18_765_000_000, - "rental_royalty_income_returns": 87_654, - "rental_royalty_income_amount": 1_654_000_000, - "partnership_scorp_income_returns": 76_543, - "partnership_scorp_income_amount": 28_765_000_000, - }, - "DE": { - "taxable_interest_returns": 121_098, - "taxable_interest_amount": 1_543_000_000, - "tax_exempt_interest_returns": 32_109, - "tax_exempt_interest_amount": 432_000_000, - "ordinary_dividends_returns": 76_543, - "ordinary_dividends_amount": 3_210_000_000, - "qualified_dividends_returns": 65_432, - "qualified_dividends_amount": 2_345_000_000, - "short_term_capital_gains_returns": 28_765, - "short_term_capital_gains_amount": 543_000_000, - "long_term_capital_gains_returns": 48_765, - "long_term_capital_gains_amount": 3_876_000_000, - "state_local_refunds_returns": 43_210, - "state_local_refunds_amount": 87_000_000, - "alimony_received_returns": 654, - "alimony_received_amount": 12_000_000, - "schedule_c_income_returns": 43_210, - "schedule_c_income_amount": 3_987_000_000, - "rental_royalty_income_returns": 21_098, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 18_765, - "partnership_scorp_income_amount": 4_321_000_000, - }, - "DC": { - "taxable_interest_returns": 87_654, - "taxable_interest_amount": 1_876_000_000, - "tax_exempt_interest_returns": 21_098, - "tax_exempt_interest_amount": 543_000_000, - "ordinary_dividends_returns": 54_321, - "ordinary_dividends_amount": 3_876_000_000, - "qualified_dividends_returns": 45_678, - "qualified_dividends_amount": 2_876_000_000, - "short_term_capital_gains_returns": 21_098, - "short_term_capital_gains_amount": 654_000_000, - "long_term_capital_gains_returns": 38_765, - "long_term_capital_gains_amount": 4_876_000_000, - "state_local_refunds_returns": 32_109, - "state_local_refunds_amount": 76_000_000, - "alimony_received_returns": 543, - "alimony_received_amount": 12_000_000, - "schedule_c_income_returns": 32_109, - "schedule_c_income_amount": 4_321_000_000, - "rental_royalty_income_returns": 15_432, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 12_345, - "partnership_scorp_income_amount": 5_432_000_000, - }, - "FL": { - "taxable_interest_returns": 2_876_543, - "taxable_interest_amount": 32_109_000_000, - "tax_exempt_interest_returns": 543_210, - "tax_exempt_interest_amount": 8_765_000_000, - "ordinary_dividends_returns": 1_654_321, - "ordinary_dividends_amount": 65_432_000_000, - "qualified_dividends_returns": 1_432_109, - "qualified_dividends_amount": 51_098_000_000, - "short_term_capital_gains_returns": 543_210, - "short_term_capital_gains_amount": 12_345_000_000, - "long_term_capital_gains_returns": 987_654, - "long_term_capital_gains_amount": 87_654_000_000, - "state_local_refunds_returns": 876_543, - "state_local_refunds_amount": 1_543_000_000, - "alimony_received_returns": 12_345, - "alimony_received_amount": 234_000_000, - "schedule_c_income_returns": 987_654, - "schedule_c_income_amount": 87_654_000_000, - "rental_royalty_income_returns": 432_109, - "rental_royalty_income_amount": 5_432_000_000, - "partnership_scorp_income_returns": 387_654, - "partnership_scorp_income_amount": 98_765_000_000, - }, - "GA": { - "taxable_interest_returns": 987_654, - "taxable_interest_amount": 8_765_000_000, - "tax_exempt_interest_returns": 176_543, - "tax_exempt_interest_amount": 2_109_000_000, - "ordinary_dividends_returns": 543_210, - "ordinary_dividends_amount": 18_765_000_000, - "qualified_dividends_returns": 454_321, - "qualified_dividends_amount": 14_321_000_000, - "short_term_capital_gains_returns": 198_765, - "short_term_capital_gains_amount": 4_321_000_000, - "long_term_capital_gains_returns": 354_321, - "long_term_capital_gains_amount": 28_765_000_000, - "state_local_refunds_returns": 398_765, - "state_local_refunds_amount": 765_000_000, - "alimony_received_returns": 5_432, - "alimony_received_amount": 109_000_000, - "schedule_c_income_returns": 398_765, - "schedule_c_income_amount": 34_567_000_000, - "rental_royalty_income_returns": 176_543, - "rental_royalty_income_amount": 2_109_000_000, - "partnership_scorp_income_returns": 143_210, - "partnership_scorp_income_amount": 32_109_000_000, - }, - "HI": { - "taxable_interest_returns": 176_543, - "taxable_interest_amount": 1_654_000_000, - "tax_exempt_interest_returns": 38_765, - "tax_exempt_interest_amount": 432_000_000, - "ordinary_dividends_returns": 98_765, - "ordinary_dividends_amount": 3_456_000_000, - "qualified_dividends_returns": 82_109, - "qualified_dividends_amount": 2_543_000_000, - "short_term_capital_gains_returns": 38_765, - "short_term_capital_gains_amount": 765_000_000, - "long_term_capital_gains_returns": 68_765, - "long_term_capital_gains_amount": 5_876_000_000, - "state_local_refunds_returns": 65_432, - "state_local_refunds_amount": 132_000_000, - "alimony_received_returns": 876, - "alimony_received_amount": 17_000_000, - "schedule_c_income_returns": 65_432, - "schedule_c_income_amount": 5_432_000_000, - "rental_royalty_income_returns": 32_109, - "rental_royalty_income_amount": 543_000_000, - "partnership_scorp_income_returns": 26_543, - "partnership_scorp_income_amount": 6_543_000_000, - }, - "ID": { - "taxable_interest_returns": 198_765, - "taxable_interest_amount": 1_234_000_000, - "tax_exempt_interest_returns": 38_765, - "tax_exempt_interest_amount": 321_000_000, - "ordinary_dividends_returns": 109_876, - "ordinary_dividends_amount": 2_345_000_000, - "qualified_dividends_returns": 92_109, - "qualified_dividends_amount": 1_654_000_000, - "short_term_capital_gains_returns": 43_210, - "short_term_capital_gains_amount": 654_000_000, - "long_term_capital_gains_returns": 78_765, - "long_term_capital_gains_amount": 4_876_000_000, - "state_local_refunds_returns": 87_654, - "state_local_refunds_amount": 154_000_000, - "alimony_received_returns": 987, - "alimony_received_amount": 18_000_000, - "schedule_c_income_returns": 87_654, - "schedule_c_income_amount": 5_987_000_000, - "rental_royalty_income_returns": 38_765, - "rental_royalty_income_amount": 432_000_000, - "partnership_scorp_income_returns": 32_109, - "partnership_scorp_income_amount": 5_432_000_000, - }, - "IL": { - "taxable_interest_returns": 1_321_098, - "taxable_interest_amount": 14_321_000_000, - "tax_exempt_interest_returns": 254_321, - "tax_exempt_interest_amount": 3_876_000_000, - "ordinary_dividends_returns": 765_432, - "ordinary_dividends_amount": 32_109_000_000, - "qualified_dividends_returns": 643_210, - "qualified_dividends_amount": 24_567_000_000, - "short_term_capital_gains_returns": 254_321, - "short_term_capital_gains_amount": 6_543_000_000, - "long_term_capital_gains_returns": 432_109, - "long_term_capital_gains_amount": 48_765_000_000, - "state_local_refunds_returns": 487_654, - "state_local_refunds_amount": 1_098_000_000, - "alimony_received_returns": 6_543, - "alimony_received_amount": 132_000_000, - "schedule_c_income_returns": 487_654, - "schedule_c_income_amount": 45_678_000_000, - "rental_royalty_income_returns": 210_987, - "rental_royalty_income_amount": 2_876_000_000, - "partnership_scorp_income_returns": 176_543, - "partnership_scorp_income_amount": 54_321_000_000, - }, - "IN": { - "taxable_interest_returns": 654_321, - "taxable_interest_amount": 4_321_000_000, - "tax_exempt_interest_returns": 109_876, - "tax_exempt_interest_amount": 1_098_000_000, - "ordinary_dividends_returns": 354_321, - "ordinary_dividends_amount": 8_765_000_000, - "qualified_dividends_returns": 298_765, - "qualified_dividends_amount": 6_543_000_000, - "short_term_capital_gains_returns": 121_098, - "short_term_capital_gains_amount": 1_876_000_000, - "long_term_capital_gains_returns": 221_098, - "long_term_capital_gains_amount": 12_876_000_000, - "state_local_refunds_returns": 265_432, - "state_local_refunds_amount": 543_000_000, - "alimony_received_returns": 3_210, - "alimony_received_amount": 65_000_000, - "schedule_c_income_returns": 265_432, - "schedule_c_income_amount": 16_543_000_000, - "rental_royalty_income_returns": 109_876, - "rental_royalty_income_amount": 1_321_000_000, - "partnership_scorp_income_returns": 87_654, - "partnership_scorp_income_amount": 14_321_000_000, - }, - "IA": { - "taxable_interest_returns": 376_543, - "taxable_interest_amount": 2_876_000_000, - "tax_exempt_interest_returns": 65_432, - "tax_exempt_interest_amount": 654_000_000, - "ordinary_dividends_returns": 198_765, - "ordinary_dividends_amount": 5_432_000_000, - "qualified_dividends_returns": 165_432, - "qualified_dividends_amount": 3_987_000_000, - "short_term_capital_gains_returns": 68_765, - "short_term_capital_gains_amount": 1_098_000_000, - "long_term_capital_gains_returns": 132_109, - "long_term_capital_gains_amount": 7_876_000_000, - "state_local_refunds_returns": 154_321, - "state_local_refunds_amount": 321_000_000, - "alimony_received_returns": 1_765, - "alimony_received_amount": 32_000_000, - "schedule_c_income_returns": 154_321, - "schedule_c_income_amount": 9_876_000_000, - "rental_royalty_income_returns": 65_432, - "rental_royalty_income_amount": 876_000_000, - "partnership_scorp_income_returns": 54_321, - "partnership_scorp_income_amount": 8_765_000_000, - }, - "KS": { - "taxable_interest_returns": 321_098, - "taxable_interest_amount": 2_345_000_000, - "tax_exempt_interest_returns": 54_321, - "tax_exempt_interest_amount": 543_000_000, - "ordinary_dividends_returns": 176_543, - "ordinary_dividends_amount": 4_567_000_000, - "qualified_dividends_returns": 148_765, - "qualified_dividends_amount": 3_321_000_000, - "short_term_capital_gains_returns": 58_765, - "short_term_capital_gains_amount": 876_000_000, - "long_term_capital_gains_returns": 109_876, - "long_term_capital_gains_amount": 6_876_000_000, - "state_local_refunds_returns": 132_109, - "state_local_refunds_amount": 265_000_000, - "alimony_received_returns": 1_543, - "alimony_received_amount": 28_000_000, - "schedule_c_income_returns": 132_109, - "schedule_c_income_amount": 8_765_000_000, - "rental_royalty_income_returns": 54_321, - "rental_royalty_income_amount": 654_000_000, - "partnership_scorp_income_returns": 45_678, - "partnership_scorp_income_amount": 7_654_000_000, - }, - "KY": { - "taxable_interest_returns": 432_109, - "taxable_interest_amount": 2_345_000_000, - "tax_exempt_interest_returns": 76_543, - "tax_exempt_interest_amount": 654_000_000, - "ordinary_dividends_returns": 210_987, - "ordinary_dividends_amount": 4_321_000_000, - "qualified_dividends_returns": 176_543, - "qualified_dividends_amount": 3_109_000_000, - "short_term_capital_gains_returns": 76_543, - "short_term_capital_gains_amount": 1_098_000_000, - "long_term_capital_gains_returns": 143_210, - "long_term_capital_gains_amount": 7_876_000_000, - "state_local_refunds_returns": 176_543, - "state_local_refunds_amount": 354_000_000, - "alimony_received_returns": 2_109, - "alimony_received_amount": 38_000_000, - "schedule_c_income_returns": 176_543, - "schedule_c_income_amount": 10_987_000_000, - "rental_royalty_income_returns": 76_543, - "rental_royalty_income_amount": 876_000_000, - "partnership_scorp_income_returns": 65_432, - "partnership_scorp_income_amount": 8_765_000_000, - }, - "LA": { - "taxable_interest_returns": 432_109, - "taxable_interest_amount": 2_543_000_000, - "tax_exempt_interest_returns": 76_543, - "tax_exempt_interest_amount": 765_000_000, - "ordinary_dividends_returns": 210_987, - "ordinary_dividends_amount": 4_987_000_000, - "qualified_dividends_returns": 176_543, - "qualified_dividends_amount": 3_654_000_000, - "short_term_capital_gains_returns": 81_098, - "short_term_capital_gains_amount": 1_234_000_000, - "long_term_capital_gains_returns": 154_321, - "long_term_capital_gains_amount": 8_765_000_000, - "state_local_refunds_returns": 187_654, - "state_local_refunds_amount": 376_000_000, - "alimony_received_returns": 2_345, - "alimony_received_amount": 43_000_000, - "schedule_c_income_returns": 187_654, - "schedule_c_income_amount": 12_345_000_000, - "rental_royalty_income_returns": 87_654, - "rental_royalty_income_amount": 1_098_000_000, - "partnership_scorp_income_returns": 71_098, - "partnership_scorp_income_amount": 9_876_000_000, - }, - "ME": { - "taxable_interest_returns": 176_543, - "taxable_interest_amount": 1_234_000_000, - "tax_exempt_interest_returns": 32_109, - "tax_exempt_interest_amount": 321_000_000, - "ordinary_dividends_returns": 98_765, - "ordinary_dividends_amount": 2_345_000_000, - "qualified_dividends_returns": 82_109, - "qualified_dividends_amount": 1_654_000_000, - "short_term_capital_gains_returns": 32_109, - "short_term_capital_gains_amount": 487_000_000, - "long_term_capital_gains_returns": 58_765, - "long_term_capital_gains_amount": 3_567_000_000, - "state_local_refunds_returns": 65_432, - "state_local_refunds_amount": 132_000_000, - "alimony_received_returns": 765, - "alimony_received_amount": 14_000_000, - "schedule_c_income_returns": 65_432, - "schedule_c_income_amount": 4_321_000_000, - "rental_royalty_income_returns": 28_765, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 23_456, - "partnership_scorp_income_amount": 3_987_000_000, - }, - "MD": { - "taxable_interest_returns": 654_321, - "taxable_interest_amount": 7_654_000_000, - "tax_exempt_interest_returns": 132_109, - "tax_exempt_interest_amount": 1_987_000_000, - "ordinary_dividends_returns": 398_765, - "ordinary_dividends_amount": 16_543_000_000, - "qualified_dividends_returns": 332_109, - "qualified_dividends_amount": 12_345_000_000, - "short_term_capital_gains_returns": 143_210, - "short_term_capital_gains_amount": 3_210_000_000, - "long_term_capital_gains_returns": 254_321, - "long_term_capital_gains_amount": 22_345_000_000, - "state_local_refunds_returns": 265_432, - "state_local_refunds_amount": 576_000_000, - "alimony_received_returns": 3_456, - "alimony_received_amount": 71_000_000, - "schedule_c_income_returns": 265_432, - "schedule_c_income_amount": 23_456_000_000, - "rental_royalty_income_returns": 121_098, - "rental_royalty_income_amount": 1_765_000_000, - "partnership_scorp_income_returns": 98_765, - "partnership_scorp_income_amount": 25_432_000_000, - }, - "MA": { - "taxable_interest_returns": 876_543, - "taxable_interest_amount": 12_345_000_000, - "tax_exempt_interest_returns": 176_543, - "tax_exempt_interest_amount": 3_210_000_000, - "ordinary_dividends_returns": 543_210, - "ordinary_dividends_amount": 28_765_000_000, - "qualified_dividends_returns": 454_321, - "qualified_dividends_amount": 21_098_000_000, - "short_term_capital_gains_returns": 187_654, - "short_term_capital_gains_amount": 5_432_000_000, - "long_term_capital_gains_returns": 321_098, - "long_term_capital_gains_amount": 38_765_000_000, - "state_local_refunds_returns": 298_765, - "state_local_refunds_amount": 687_000_000, - "alimony_received_returns": 4_321, - "alimony_received_amount": 91_000_000, - "schedule_c_income_returns": 298_765, - "schedule_c_income_amount": 32_109_000_000, - "rental_royalty_income_returns": 143_210, - "rental_royalty_income_amount": 2_345_000_000, - "partnership_scorp_income_returns": 121_098, - "partnership_scorp_income_amount": 43_210_000_000, - }, - "MI": { - "taxable_interest_returns": 987_654, - "taxable_interest_amount": 7_654_000_000, - "tax_exempt_interest_returns": 176_543, - "tax_exempt_interest_amount": 1_654_000_000, - "ordinary_dividends_returns": 543_210, - "ordinary_dividends_amount": 14_321_000_000, - "qualified_dividends_returns": 454_321, - "qualified_dividends_amount": 10_876_000_000, - "short_term_capital_gains_returns": 176_543, - "short_term_capital_gains_amount": 2_876_000_000, - "long_term_capital_gains_returns": 321_098, - "long_term_capital_gains_amount": 21_098_000_000, - "state_local_refunds_returns": 376_543, - "state_local_refunds_amount": 765_000_000, - "alimony_received_returns": 4_765, - "alimony_received_amount": 98_000_000, - "schedule_c_income_returns": 376_543, - "schedule_c_income_amount": 25_432_000_000, - "rental_royalty_income_returns": 154_321, - "rental_royalty_income_amount": 1_876_000_000, - "partnership_scorp_income_returns": 132_109, - "partnership_scorp_income_amount": 23_456_000_000, - }, - "MN": { - "taxable_interest_returns": 654_321, - "taxable_interest_amount": 5_987_000_000, - "tax_exempt_interest_returns": 121_098, - "tax_exempt_interest_amount": 1_321_000_000, - "ordinary_dividends_returns": 387_654, - "ordinary_dividends_amount": 12_345_000_000, - "qualified_dividends_returns": 321_098, - "qualified_dividends_amount": 9_321_000_000, - "short_term_capital_gains_returns": 132_109, - "short_term_capital_gains_amount": 2_543_000_000, - "long_term_capital_gains_returns": 232_109, - "long_term_capital_gains_amount": 18_765_000_000, - "state_local_refunds_returns": 243_210, - "state_local_refunds_amount": 521_000_000, - "alimony_received_returns": 2_987, - "alimony_received_amount": 61_000_000, - "schedule_c_income_returns": 243_210, - "schedule_c_income_amount": 21_098_000_000, - "rental_royalty_income_returns": 109_876, - "rental_royalty_income_amount": 1_432_000_000, - "partnership_scorp_income_returns": 91_098, - "partnership_scorp_income_amount": 21_098_000_000, - }, - "MS": { - "taxable_interest_returns": 265_432, - "taxable_interest_amount": 1_098_000_000, - "tax_exempt_interest_returns": 43_210, - "tax_exempt_interest_amount": 321_000_000, - "ordinary_dividends_returns": 121_098, - "ordinary_dividends_amount": 2_109_000_000, - "qualified_dividends_returns": 98_765, - "qualified_dividends_amount": 1_432_000_000, - "short_term_capital_gains_returns": 43_210, - "short_term_capital_gains_amount": 543_000_000, - "long_term_capital_gains_returns": 78_765, - "long_term_capital_gains_amount": 3_876_000_000, - "state_local_refunds_returns": 109_876, - "state_local_refunds_amount": 187_000_000, - "alimony_received_returns": 1_321, - "alimony_received_amount": 24_000_000, - "schedule_c_income_returns": 109_876, - "schedule_c_income_amount": 5_987_000_000, - "rental_royalty_income_returns": 43_210, - "rental_royalty_income_amount": 432_000_000, - "partnership_scorp_income_returns": 35_432, - "partnership_scorp_income_amount": 4_321_000_000, - }, - "MO": { - "taxable_interest_returns": 598_765, - "taxable_interest_amount": 4_321_000_000, - "tax_exempt_interest_returns": 109_876, - "tax_exempt_interest_amount": 987_000_000, - "ordinary_dividends_returns": 321_098, - "ordinary_dividends_amount": 8_765_000_000, - "qualified_dividends_returns": 265_432, - "qualified_dividends_amount": 6_321_000_000, - "short_term_capital_gains_returns": 109_876, - "short_term_capital_gains_amount": 1_765_000_000, - "long_term_capital_gains_returns": 198_765, - "long_term_capital_gains_amount": 12_876_000_000, - "state_local_refunds_returns": 243_210, - "state_local_refunds_amount": 487_000_000, - "alimony_received_returns": 2_876, - "alimony_received_amount": 54_000_000, - "schedule_c_income_returns": 243_210, - "schedule_c_income_amount": 16_543_000_000, - "rental_royalty_income_returns": 98_765, - "rental_royalty_income_amount": 1_098_000_000, - "partnership_scorp_income_returns": 81_098, - "partnership_scorp_income_amount": 14_321_000_000, - }, - "MT": { - "taxable_interest_returns": 132_109, - "taxable_interest_amount": 1_098_000_000, - "tax_exempt_interest_returns": 26_543, - "tax_exempt_interest_amount": 234_000_000, - "ordinary_dividends_returns": 76_543, - "ordinary_dividends_amount": 2_109_000_000, - "qualified_dividends_returns": 63_210, - "qualified_dividends_amount": 1_432_000_000, - "short_term_capital_gains_returns": 26_543, - "short_term_capital_gains_amount": 421_000_000, - "long_term_capital_gains_returns": 48_765, - "long_term_capital_gains_amount": 3_098_000_000, - "state_local_refunds_returns": 54_321, - "state_local_refunds_amount": 98_000_000, - "alimony_received_returns": 654, - "alimony_received_amount": 12_000_000, - "schedule_c_income_returns": 54_321, - "schedule_c_income_amount": 3_987_000_000, - "rental_royalty_income_returns": 26_543, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 21_098, - "partnership_scorp_income_amount": 3_456_000_000, - }, - "NE": { - "taxable_interest_returns": 232_109, - "taxable_interest_amount": 1_654_000_000, - "tax_exempt_interest_returns": 43_210, - "tax_exempt_interest_amount": 376_000_000, - "ordinary_dividends_returns": 121_098, - "ordinary_dividends_amount": 3_210_000_000, - "qualified_dividends_returns": 98_765, - "qualified_dividends_amount": 2_321_000_000, - "short_term_capital_gains_returns": 43_210, - "short_term_capital_gains_amount": 654_000_000, - "long_term_capital_gains_returns": 78_765, - "long_term_capital_gains_amount": 4_876_000_000, - "state_local_refunds_returns": 98_765, - "state_local_refunds_amount": 198_000_000, - "alimony_received_returns": 1_098, - "alimony_received_amount": 19_000_000, - "schedule_c_income_returns": 98_765, - "schedule_c_income_amount": 6_543_000_000, - "rental_royalty_income_returns": 43_210, - "rental_royalty_income_amount": 487_000_000, - "partnership_scorp_income_returns": 35_432, - "partnership_scorp_income_amount": 5_432_000_000, - }, - "NV": { - "taxable_interest_returns": 354_321, - "taxable_interest_amount": 2_876_000_000, - "tax_exempt_interest_returns": 65_432, - "tax_exempt_interest_amount": 654_000_000, - "ordinary_dividends_returns": 198_765, - "ordinary_dividends_amount": 5_987_000_000, - "qualified_dividends_returns": 165_432, - "qualified_dividends_amount": 4_321_000_000, - "short_term_capital_gains_returns": 71_098, - "short_term_capital_gains_amount": 1_543_000_000, - "long_term_capital_gains_returns": 132_109, - "long_term_capital_gains_amount": 10_987_000_000, - "state_local_refunds_returns": 121_098, - "state_local_refunds_amount": 198_000_000, - "alimony_received_returns": 1_654, - "alimony_received_amount": 31_000_000, - "schedule_c_income_returns": 143_210, - "schedule_c_income_amount": 12_345_000_000, - "rental_royalty_income_returns": 65_432, - "rental_royalty_income_amount": 876_000_000, - "partnership_scorp_income_returns": 54_321, - "partnership_scorp_income_amount": 12_345_000_000, - }, - "NH": { - "taxable_interest_returns": 176_543, - "taxable_interest_amount": 2_543_000_000, - "tax_exempt_interest_returns": 38_765, - "tax_exempt_interest_amount": 654_000_000, - "ordinary_dividends_returns": 109_876, - "ordinary_dividends_amount": 5_432_000_000, - "qualified_dividends_returns": 91_098, - "qualified_dividends_amount": 3_987_000_000, - "short_term_capital_gains_returns": 38_765, - "short_term_capital_gains_amount": 987_000_000, - "long_term_capital_gains_returns": 68_765, - "long_term_capital_gains_amount": 6_765_000_000, - "state_local_refunds_returns": 58_765, - "state_local_refunds_amount": 98_000_000, - "alimony_received_returns": 876, - "alimony_received_amount": 17_000_000, - "schedule_c_income_returns": 65_432, - "schedule_c_income_amount": 6_543_000_000, - "rental_royalty_income_returns": 32_109, - "rental_royalty_income_amount": 543_000_000, - "partnership_scorp_income_returns": 26_543, - "partnership_scorp_income_amount": 7_654_000_000, - }, - "NJ": { - "taxable_interest_returns": 1_098_765, - "taxable_interest_amount": 14_321_000_000, - "tax_exempt_interest_returns": 221_098, - "tax_exempt_interest_amount": 4_321_000_000, - "ordinary_dividends_returns": 654_321, - "ordinary_dividends_amount": 32_109_000_000, - "qualified_dividends_returns": 543_210, - "qualified_dividends_amount": 24_321_000_000, - "short_term_capital_gains_returns": 221_098, - "short_term_capital_gains_amount": 6_543_000_000, - "long_term_capital_gains_returns": 387_654, - "long_term_capital_gains_amount": 48_765_000_000, - "state_local_refunds_returns": 398_765, - "state_local_refunds_amount": 987_000_000, - "alimony_received_returns": 5_432, - "alimony_received_amount": 121_000_000, - "schedule_c_income_returns": 398_765, - "schedule_c_income_amount": 43_210_000_000, - "rental_royalty_income_returns": 176_543, - "rental_royalty_income_amount": 2_876_000_000, - "partnership_scorp_income_returns": 154_321, - "partnership_scorp_income_amount": 54_321_000_000, - }, - "NM": { - "taxable_interest_returns": 210_987, - "taxable_interest_amount": 1_321_000_000, - "tax_exempt_interest_returns": 38_765, - "tax_exempt_interest_amount": 321_000_000, - "ordinary_dividends_returns": 109_876, - "ordinary_dividends_amount": 2_543_000_000, - "qualified_dividends_returns": 91_098, - "qualified_dividends_amount": 1_765_000_000, - "short_term_capital_gains_returns": 38_765, - "short_term_capital_gains_amount": 598_000_000, - "long_term_capital_gains_returns": 68_765, - "long_term_capital_gains_amount": 4_432_000_000, - "state_local_refunds_returns": 87_654, - "state_local_refunds_amount": 154_000_000, - "alimony_received_returns": 987, - "alimony_received_amount": 18_000_000, - "schedule_c_income_returns": 87_654, - "schedule_c_income_amount": 5_432_000_000, - "rental_royalty_income_returns": 38_765, - "rental_royalty_income_amount": 543_000_000, - "partnership_scorp_income_returns": 31_098, - "partnership_scorp_income_amount": 4_987_000_000, - }, - "NY": { - "taxable_interest_returns": 2_345_678, - "taxable_interest_amount": 38_765_000_000, - "tax_exempt_interest_returns": 487_654, - "tax_exempt_interest_amount": 12_345_000_000, - "ordinary_dividends_returns": 1_432_109, - "ordinary_dividends_amount": 87_654_000_000, - "qualified_dividends_returns": 1_198_765, - "qualified_dividends_amount": 65_432_000_000, - "short_term_capital_gains_returns": 498_765, - "short_term_capital_gains_amount": 21_098_000_000, - "long_term_capital_gains_returns": 876_543, - "long_term_capital_gains_amount": 145_678_000_000, - "state_local_refunds_returns": 876_543, - "state_local_refunds_amount": 2_109_000_000, - "alimony_received_returns": 12_345, - "alimony_received_amount": 287_000_000, - "schedule_c_income_returns": 876_543, - "schedule_c_income_amount": 132_109_000_000, - "rental_royalty_income_returns": 376_543, - "rental_royalty_income_amount": 7_654_000_000, - "partnership_scorp_income_returns": 321_098, - "partnership_scorp_income_amount": 165_432_000_000, - }, - "NC": { - "taxable_interest_returns": 987_654, - "taxable_interest_amount": 7_654_000_000, - "tax_exempt_interest_returns": 176_543, - "tax_exempt_interest_amount": 1_765_000_000, - "ordinary_dividends_returns": 543_210, - "ordinary_dividends_amount": 16_543_000_000, - "qualified_dividends_returns": 454_321, - "qualified_dividends_amount": 12_345_000_000, - "short_term_capital_gains_returns": 187_654, - "short_term_capital_gains_amount": 3_543_000_000, - "long_term_capital_gains_returns": 343_210, - "long_term_capital_gains_amount": 25_432_000_000, - "state_local_refunds_returns": 387_654, - "state_local_refunds_amount": 765_000_000, - "alimony_received_returns": 5_098, - "alimony_received_amount": 98_000_000, - "schedule_c_income_returns": 387_654, - "schedule_c_income_amount": 31_098_000_000, - "rental_royalty_income_returns": 165_432, - "rental_royalty_income_amount": 1_987_000_000, - "partnership_scorp_income_returns": 132_109, - "partnership_scorp_income_amount": 28_765_000_000, - }, - "ND": { - "taxable_interest_returns": 98_765, - "taxable_interest_amount": 876_000_000, - "tax_exempt_interest_returns": 18_765, - "tax_exempt_interest_amount": 154_000_000, - "ordinary_dividends_returns": 54_321, - "ordinary_dividends_amount": 1_654_000_000, - "qualified_dividends_returns": 45_678, - "qualified_dividends_amount": 1_098_000_000, - "short_term_capital_gains_returns": 18_765, - "short_term_capital_gains_amount": 298_000_000, - "long_term_capital_gains_returns": 35_432, - "long_term_capital_gains_amount": 2_287_000_000, - "state_local_refunds_returns": 43_210, - "state_local_refunds_amount": 76_000_000, - "alimony_received_returns": 432, - "alimony_received_amount": 8_000_000, - "schedule_c_income_returns": 43_210, - "schedule_c_income_amount": 2_876_000_000, - "rental_royalty_income_returns": 21_098, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 16_543, - "partnership_scorp_income_amount": 2_543_000_000, - }, - "OH": { - "taxable_interest_returns": 1_098_765, - "taxable_interest_amount": 8_765_000_000, - "tax_exempt_interest_returns": 198_765, - "tax_exempt_interest_amount": 1_876_000_000, - "ordinary_dividends_returns": 598_765, - "ordinary_dividends_amount": 16_543_000_000, - "qualified_dividends_returns": 498_765, - "qualified_dividends_amount": 12_345_000_000, - "short_term_capital_gains_returns": 198_765, - "short_term_capital_gains_amount": 3_543_000_000, - "long_term_capital_gains_returns": 365_432, - "long_term_capital_gains_amount": 25_432_000_000, - "state_local_refunds_returns": 421_098, - "state_local_refunds_amount": 876_000_000, - "alimony_received_returns": 5_432, - "alimony_received_amount": 109_000_000, - "schedule_c_income_returns": 421_098, - "schedule_c_income_amount": 32_109_000_000, - "rental_royalty_income_returns": 176_543, - "rental_royalty_income_amount": 2_109_000_000, - "partnership_scorp_income_returns": 143_210, - "partnership_scorp_income_amount": 28_765_000_000, - }, - "OK": { - "taxable_interest_returns": 354_321, - "taxable_interest_amount": 2_109_000_000, - "tax_exempt_interest_returns": 65_432, - "tax_exempt_interest_amount": 487_000_000, - "ordinary_dividends_returns": 176_543, - "ordinary_dividends_amount": 3_987_000_000, - "qualified_dividends_returns": 148_765, - "qualified_dividends_amount": 2_876_000_000, - "short_term_capital_gains_returns": 65_432, - "short_term_capital_gains_amount": 932_000_000, - "long_term_capital_gains_returns": 121_098, - "long_term_capital_gains_amount": 6_765_000_000, - "state_local_refunds_returns": 154_321, - "state_local_refunds_amount": 287_000_000, - "alimony_received_returns": 1_876, - "alimony_received_amount": 34_000_000, - "schedule_c_income_returns": 154_321, - "schedule_c_income_amount": 9_876_000_000, - "rental_royalty_income_returns": 65_432, - "rental_royalty_income_amount": 876_000_000, - "partnership_scorp_income_returns": 54_321, - "partnership_scorp_income_amount": 7_654_000_000, - }, - "OR": { - "taxable_interest_returns": 476_543, - "taxable_interest_amount": 3_987_000_000, - "tax_exempt_interest_returns": 91_098, - "tax_exempt_interest_amount": 876_000_000, - "ordinary_dividends_returns": 276_543, - "ordinary_dividends_amount": 8_765_000_000, - "qualified_dividends_returns": 232_109, - "qualified_dividends_amount": 6_543_000_000, - "short_term_capital_gains_returns": 98_765, - "short_term_capital_gains_amount": 1_765_000_000, - "long_term_capital_gains_returns": 176_543, - "long_term_capital_gains_amount": 12_654_000_000, - "state_local_refunds_returns": 187_654, - "state_local_refunds_amount": 387_000_000, - "alimony_received_returns": 2_345, - "alimony_received_amount": 45_000_000, - "schedule_c_income_returns": 187_654, - "schedule_c_income_amount": 14_321_000_000, - "rental_royalty_income_returns": 87_654, - "rental_royalty_income_amount": 1_098_000_000, - "partnership_scorp_income_returns": 71_098, - "partnership_scorp_income_amount": 14_321_000_000, - }, - "PA": { - "taxable_interest_returns": 1_321_098, - "taxable_interest_amount": 14_321_000_000, - "tax_exempt_interest_returns": 265_432, - "tax_exempt_interest_amount": 3_654_000_000, - "ordinary_dividends_returns": 765_432, - "ordinary_dividends_amount": 28_765_000_000, - "qualified_dividends_returns": 643_210, - "qualified_dividends_amount": 21_098_000_000, - "short_term_capital_gains_returns": 254_321, - "short_term_capital_gains_amount": 6_543_000_000, - "long_term_capital_gains_returns": 454_321, - "long_term_capital_gains_amount": 48_765_000_000, - "state_local_refunds_returns": 487_654, - "state_local_refunds_amount": 1_098_000_000, - "alimony_received_returns": 6_543, - "alimony_received_amount": 143_000_000, - "schedule_c_income_returns": 487_654, - "schedule_c_income_amount": 51_098_000_000, - "rental_royalty_income_returns": 210_987, - "rental_royalty_income_amount": 2_876_000_000, - "partnership_scorp_income_returns": 176_543, - "partnership_scorp_income_amount": 54_321_000_000, - }, - "RI": { - "taxable_interest_returns": 132_109, - "taxable_interest_amount": 1_432_000_000, - "tax_exempt_interest_returns": 26_543, - "tax_exempt_interest_amount": 354_000_000, - "ordinary_dividends_returns": 76_543, - "ordinary_dividends_amount": 2_876_000_000, - "qualified_dividends_returns": 63_210, - "qualified_dividends_amount": 2_109_000_000, - "short_term_capital_gains_returns": 26_543, - "short_term_capital_gains_amount": 521_000_000, - "long_term_capital_gains_returns": 48_765, - "long_term_capital_gains_amount": 3_876_000_000, - "state_local_refunds_returns": 43_210, - "state_local_refunds_amount": 91_000_000, - "alimony_received_returns": 654, - "alimony_received_amount": 13_000_000, - "schedule_c_income_returns": 43_210, - "schedule_c_income_amount": 3_876_000_000, - "rental_royalty_income_returns": 21_098, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 16_543, - "partnership_scorp_income_amount": 4_321_000_000, - }, - "SC": { - "taxable_interest_returns": 521_098, - "taxable_interest_amount": 3_210_000_000, - "tax_exempt_interest_returns": 91_098, - "tax_exempt_interest_amount": 765_000_000, - "ordinary_dividends_returns": 276_543, - "ordinary_dividends_amount": 6_543_000_000, - "qualified_dividends_returns": 232_109, - "qualified_dividends_amount": 4_876_000_000, - "short_term_capital_gains_returns": 91_098, - "short_term_capital_gains_amount": 1_543_000_000, - "long_term_capital_gains_returns": 176_543, - "long_term_capital_gains_amount": 10_987_000_000, - "state_local_refunds_returns": 198_765, - "state_local_refunds_amount": 376_000_000, - "alimony_received_returns": 2_654, - "alimony_received_amount": 51_000_000, - "schedule_c_income_returns": 198_765, - "schedule_c_income_amount": 12_345_000_000, - "rental_royalty_income_returns": 87_654, - "rental_royalty_income_amount": 987_000_000, - "partnership_scorp_income_returns": 71_098, - "partnership_scorp_income_amount": 12_345_000_000, - }, - "SD": { - "taxable_interest_returns": 109_876, - "taxable_interest_amount": 1_098_000_000, - "tax_exempt_interest_returns": 21_098, - "tax_exempt_interest_amount": 198_000_000, - "ordinary_dividends_returns": 65_432, - "ordinary_dividends_amount": 2_109_000_000, - "qualified_dividends_returns": 54_321, - "qualified_dividends_amount": 1_432_000_000, - "short_term_capital_gains_returns": 21_098, - "short_term_capital_gains_amount": 387_000_000, - "long_term_capital_gains_returns": 38_765, - "long_term_capital_gains_amount": 2_876_000_000, - "state_local_refunds_returns": 38_765, - "state_local_refunds_amount": 65_000_000, - "alimony_received_returns": 487, - "alimony_received_amount": 9_000_000, - "schedule_c_income_returns": 43_210, - "schedule_c_income_amount": 3_456_000_000, - "rental_royalty_income_returns": 21_098, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 16_543, - "partnership_scorp_income_amount": 3_210_000_000, - }, - "TN": { - "taxable_interest_returns": 654_321, - "taxable_interest_amount": 5_432_000_000, - "tax_exempt_interest_returns": 121_098, - "tax_exempt_interest_amount": 1_321_000_000, - "ordinary_dividends_returns": 354_321, - "ordinary_dividends_amount": 10_987_000_000, - "qualified_dividends_returns": 298_765, - "qualified_dividends_amount": 8_109_000_000, - "short_term_capital_gains_returns": 126_543, - "short_term_capital_gains_amount": 2_321_000_000, - "long_term_capital_gains_returns": 232_109, - "long_term_capital_gains_amount": 16_543_000_000, - "state_local_refunds_returns": 254_321, - "state_local_refunds_amount": 432_000_000, - "alimony_received_returns": 3_456, - "alimony_received_amount": 67_000_000, - "schedule_c_income_returns": 276_543, - "schedule_c_income_amount": 21_098_000_000, - "rental_royalty_income_returns": 121_098, - "rental_royalty_income_amount": 1_432_000_000, - "partnership_scorp_income_returns": 98_765, - "partnership_scorp_income_amount": 18_765_000_000, - }, - "TX": { - "taxable_interest_returns": 2_543_210, - "taxable_interest_amount": 25_432_000_000, - "tax_exempt_interest_returns": 476_543, - "tax_exempt_interest_amount": 6_543_000_000, - "ordinary_dividends_returns": 1_321_098, - "ordinary_dividends_amount": 54_321_000_000, - "qualified_dividends_returns": 1_098_765, - "qualified_dividends_amount": 41_098_000_000, - "short_term_capital_gains_returns": 543_210, - "short_term_capital_gains_amount": 13_456_000_000, - "long_term_capital_gains_returns": 987_654, - "long_term_capital_gains_amount": 98_765_000_000, - "state_local_refunds_returns": 987_654, - "state_local_refunds_amount": 1_654_000_000, - "alimony_received_returns": 12_345, - "alimony_received_amount": 243_000_000, - "schedule_c_income_returns": 1_098_765, - "schedule_c_income_amount": 109_876_000_000, - "rental_royalty_income_returns": 476_543, - "rental_royalty_income_amount": 7_654_000_000, - "partnership_scorp_income_returns": 398_765, - "partnership_scorp_income_amount": 109_876_000_000, - }, - "UT": { - "taxable_interest_returns": 321_098, - "taxable_interest_amount": 2_345_000_000, - "tax_exempt_interest_returns": 58_765, - "tax_exempt_interest_amount": 521_000_000, - "ordinary_dividends_returns": 176_543, - "ordinary_dividends_amount": 4_987_000_000, - "qualified_dividends_returns": 148_765, - "qualified_dividends_amount": 3_654_000_000, - "short_term_capital_gains_returns": 65_432, - "short_term_capital_gains_amount": 1_198_000_000, - "long_term_capital_gains_returns": 121_098, - "long_term_capital_gains_amount": 8_765_000_000, - "state_local_refunds_returns": 143_210, - "state_local_refunds_amount": 287_000_000, - "alimony_received_returns": 1_654, - "alimony_received_amount": 31_000_000, - "schedule_c_income_returns": 143_210, - "schedule_c_income_amount": 10_987_000_000, - "rental_royalty_income_returns": 58_765, - "rental_royalty_income_amount": 654_000_000, - "partnership_scorp_income_returns": 48_765, - "partnership_scorp_income_amount": 9_876_000_000, - }, - "VT": { - "taxable_interest_returns": 87_654, - "taxable_interest_amount": 876_000_000, - "tax_exempt_interest_returns": 18_765, - "tax_exempt_interest_amount": 198_000_000, - "ordinary_dividends_returns": 54_321, - "ordinary_dividends_amount": 1_654_000_000, - "qualified_dividends_returns": 45_678, - "qualified_dividends_amount": 1_198_000_000, - "short_term_capital_gains_returns": 16_543, - "short_term_capital_gains_amount": 287_000_000, - "long_term_capital_gains_returns": 28_765, - "long_term_capital_gains_amount": 2_098_000_000, - "state_local_refunds_returns": 32_109, - "state_local_refunds_amount": 65_000_000, - "alimony_received_returns": 432, - "alimony_received_amount": 8_000_000, - "schedule_c_income_returns": 32_109, - "schedule_c_income_amount": 2_543_000_000, - "rental_royalty_income_returns": 16_543, - "rental_royalty_income_amount": 198_000_000, - "partnership_scorp_income_returns": 12_345, - "partnership_scorp_income_amount": 2_345_000_000, - }, - "VA": { - "taxable_interest_returns": 876_543, - "taxable_interest_amount": 9_876_000_000, - "tax_exempt_interest_returns": 176_543, - "tax_exempt_interest_amount": 2_345_000_000, - "ordinary_dividends_returns": 521_098, - "ordinary_dividends_amount": 21_098_000_000, - "qualified_dividends_returns": 432_109, - "qualified_dividends_amount": 15_876_000_000, - "short_term_capital_gains_returns": 176_543, - "short_term_capital_gains_amount": 4_321_000_000, - "long_term_capital_gains_returns": 321_098, - "long_term_capital_gains_amount": 30_876_000_000, - "state_local_refunds_returns": 343_210, - "state_local_refunds_amount": 765_000_000, - "alimony_received_returns": 4_321, - "alimony_received_amount": 87_000_000, - "schedule_c_income_returns": 343_210, - "schedule_c_income_amount": 32_109_000_000, - "rental_royalty_income_returns": 154_321, - "rental_royalty_income_amount": 2_109_000_000, - "partnership_scorp_income_returns": 126_543, - "partnership_scorp_income_amount": 34_567_000_000, - }, - "WA": { - "taxable_interest_returns": 765_432, - "taxable_interest_amount": 8_765_000_000, - "tax_exempt_interest_returns": 154_321, - "tax_exempt_interest_amount": 2_109_000_000, - "ordinary_dividends_returns": 487_654, - "ordinary_dividends_amount": 21_098_000_000, - "qualified_dividends_returns": 409_876, - "qualified_dividends_amount": 15_876_000_000, - "short_term_capital_gains_returns": 165_432, - "short_term_capital_gains_amount": 4_876_000_000, - "long_term_capital_gains_returns": 298_765, - "long_term_capital_gains_amount": 34_567_000_000, - "state_local_refunds_returns": 287_654, - "state_local_refunds_amount": 487_000_000, - "alimony_received_returns": 3_876, - "alimony_received_amount": 76_000_000, - "schedule_c_income_returns": 321_098, - "schedule_c_income_amount": 34_567_000_000, - "rental_royalty_income_returns": 143_210, - "rental_royalty_income_amount": 2_109_000_000, - "partnership_scorp_income_returns": 121_098, - "partnership_scorp_income_amount": 38_765_000_000, - }, - "WV": { - "taxable_interest_returns": 176_543, - "taxable_interest_amount": 987_000_000, - "tax_exempt_interest_returns": 28_765, - "tax_exempt_interest_amount": 198_000_000, - "ordinary_dividends_returns": 87_654, - "ordinary_dividends_amount": 1_654_000_000, - "qualified_dividends_returns": 71_098, - "qualified_dividends_amount": 1_198_000_000, - "short_term_capital_gains_returns": 26_543, - "short_term_capital_gains_amount": 354_000_000, - "long_term_capital_gains_returns": 48_765, - "long_term_capital_gains_amount": 2_543_000_000, - "state_local_refunds_returns": 65_432, - "state_local_refunds_amount": 121_000_000, - "alimony_received_returns": 765, - "alimony_received_amount": 14_000_000, - "schedule_c_income_returns": 65_432, - "schedule_c_income_amount": 3_987_000_000, - "rental_royalty_income_returns": 26_543, - "rental_royalty_income_amount": 321_000_000, - "partnership_scorp_income_returns": 21_098, - "partnership_scorp_income_amount": 2_876_000_000, - }, - "WI": { - "taxable_interest_returns": 621_098, - "taxable_interest_amount": 4_987_000_000, - "tax_exempt_interest_returns": 109_876, - "tax_exempt_interest_amount": 1_098_000_000, - "ordinary_dividends_returns": 354_321, - "ordinary_dividends_amount": 9_876_000_000, - "qualified_dividends_returns": 298_765, - "qualified_dividends_amount": 7_321_000_000, - "short_term_capital_gains_returns": 121_098, - "short_term_capital_gains_amount": 2_098_000_000, - "long_term_capital_gains_returns": 210_987, - "long_term_capital_gains_amount": 14_567_000_000, - "state_local_refunds_returns": 243_210, - "state_local_refunds_amount": 498_000_000, - "alimony_received_returns": 2_987, - "alimony_received_amount": 58_000_000, - "schedule_c_income_returns": 243_210, - "schedule_c_income_amount": 18_765_000_000, - "rental_royalty_income_returns": 98_765, - "rental_royalty_income_amount": 1_198_000_000, - "partnership_scorp_income_returns": 81_098, - "partnership_scorp_income_amount": 16_543_000_000, - }, - "WY": { - "taxable_interest_returns": 76_543, - "taxable_interest_amount": 876_000_000, - "tax_exempt_interest_returns": 15_432, - "tax_exempt_interest_amount": 154_000_000, - "ordinary_dividends_returns": 43_210, - "ordinary_dividends_amount": 1_654_000_000, - "qualified_dividends_returns": 35_432, - "qualified_dividends_amount": 1_198_000_000, - "short_term_capital_gains_returns": 15_432, - "short_term_capital_gains_amount": 354_000_000, - "long_term_capital_gains_returns": 28_765, - "long_term_capital_gains_amount": 2_543_000_000, - "state_local_refunds_returns": 26_543, - "state_local_refunds_amount": 43_000_000, - "alimony_received_returns": 354, - "alimony_received_amount": 7_000_000, - "schedule_c_income_returns": 32_109, - "schedule_c_income_amount": 2_876_000_000, - "rental_royalty_income_returns": 15_432, - "rental_royalty_income_amount": 234_000_000, - "partnership_scorp_income_returns": 12_345, - "partnership_scorp_income_amount": 2_876_000_000, - }, - }, -} - -SOURCE_URL = ( - "https://www.irs.gov/statistics/soi-tax-stats-individual-income-tax-statistics" -) - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - # Check if exists - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - # Create new - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() # Get ID - - # Add constraints - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_soi_income_sources_targets(session: Session, years: list[int] | None = None): - """ - Load income by source targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SOI_NATIONAL_INCOME_SOURCES_DATA.keys()) - - for year in years: - # Load national data first - if year in SOI_NATIONAL_INCOME_SOURCES_DATA: - national_data = SOI_NATIONAL_INCOME_SOURCES_DATA[year] - - # Get or create national stratum - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Add national targets for each income source - for source in INCOME_SOURCES: - returns_var = f"{source}_returns" - amount_var = f"{source}_amount" - - if returns_var in national_data: - session.add( - Target( - stratum_id=national_stratum.id, - variable=returns_var, - period=year, - value=national_data[returns_var], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="Publication 1304, Table 1.4", - source_url=SOURCE_URL, - ) - ) - - if amount_var in national_data: - session.add( - Target( - stratum_id=national_stratum.id, - variable=amount_var, - period=year, - value=national_data[amount_var], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="Publication 1304, Table 1.4", - source_url=SOURCE_URL, - ) - ) - - # Load state data - if year not in SOI_STATE_INCOME_SOURCES_DATA: - continue - - state_data = SOI_STATE_INCOME_SOURCES_DATA[year] - - # Get or create national stratum (for parent relationship) - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create state-level strata and targets - for state_abbrev, state_values in state_data.items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - # Create state stratum - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} All Filers", - jurisdiction=Jurisdiction.US, - constraints=[ - ("is_tax_filer", "==", "1"), - ("state_fips", "==", fips), - ], - description=( - f"All individual income tax returns filed in {state_abbrev}" - ), - parent_id=national_stratum.id, - stratum_group_id="soi_states", - ) - - # Add targets for each income source - for source in INCOME_SOURCES: - returns_var = f"{source}_returns" - amount_var = f"{source}_amount" - - if returns_var in state_values: - session.add( - Target( - stratum_id=state_stratum.id, - variable=returns_var, - period=year, - value=state_values[returns_var], - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - source_table="Individual Income Tax State Data", - source_url=SOURCE_URL, - ) - ) - - if amount_var in state_values: - session.add( - Target( - stratum_id=state_stratum.id, - variable=amount_var, - period=year, - value=state_values[amount_var], - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - source_table="Individual Income Tax State Data", - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the income sources ETL pipeline.""" - from pathlib import Path - - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_soi_income_sources_targets(session) - print(f"Loaded SOI income sources targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_soi_state.py b/db/etl_soi_state.py deleted file mode 100644 index bbdc8977..00000000 --- a/db/etl_soi_state.py +++ /dev/null @@ -1,921 +0,0 @@ -""" -ETL for IRS Statistics of Income (SOI) state-level targets. - -Loads state-by-state data from IRS SOI tables into the targets database. -Data source: https://www.irs.gov/statistics/soi-tax-stats-historic-table-2 -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - GeographicLevel, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes for all 50 states + DC -STATE_FIPS = { - "AL": "01", - "AK": "02", - "AZ": "04", - "AR": "05", - "CA": "06", - "CO": "08", - "CT": "09", - "DE": "10", - "DC": "11", - "FL": "12", - "GA": "13", - "HI": "15", - "ID": "16", - "IL": "17", - "IN": "18", - "IA": "19", - "KS": "20", - "KY": "21", - "LA": "22", - "ME": "23", - "MD": "24", - "MA": "25", - "MI": "26", - "MN": "27", - "MS": "28", - "MO": "29", - "MT": "30", - "NE": "31", - "NV": "32", - "NH": "33", - "NJ": "34", - "NM": "35", - "NY": "36", - "NC": "37", - "ND": "38", - "OH": "39", - "OK": "40", - "OR": "41", - "PA": "42", - "RI": "44", - "SC": "45", - "SD": "46", - "TN": "47", - "TX": "48", - "UT": "49", - "VT": "50", - "VA": "51", - "WA": "53", - "WV": "54", - "WI": "55", - "WY": "56", -} - -# State-level SOI data by year (from IRS Historic Table 2) -# Source: https://www.irs.gov/statistics/soi-tax-stats-historic-table-2 -# Representative data for all 50 states + DC (2021 tax year) -SOI_STATE_DATA = { - 2021: { - "AL": { - "total_returns": 2_234_567, - "total_agi": 134_567_000_000, - "total_tax_liability": 11_234_000_000, - }, - "AK": { - "total_returns": 387_654, - "total_agi": 34_567_000_000, - "total_tax_liability": 2_876_000_000, - }, - "AZ": { - "total_returns": 3_456_789, - "total_agi": 234_567_000_000, - "total_tax_liability": 19_543_000_000, - }, - "AR": { - "total_returns": 1_345_678, - "total_agi": 78_234_000_000, - "total_tax_liability": 6_543_000_000, - }, - "CA": { - "total_returns": 18_547_234, - "total_agi": 2_154_789_000_000, - "total_tax_liability": 187_654_000_000, - }, - "CO": { - "total_returns": 2_987_654, - "total_agi": 234_567_000_000, - "total_tax_liability": 19_876_000_000, - }, - "CT": { - "total_returns": 1_876_543, - "total_agi": 198_765_000_000, - "total_tax_liability": 17_654_000_000, - }, - "DE": { - "total_returns": 487_654, - "total_agi": 43_234_000_000, - "total_tax_liability": 3_654_000_000, - }, - "DC": { - "total_returns": 387_654, - "total_agi": 54_321_000_000, - "total_tax_liability": 4_876_000_000, - }, - "FL": { - "total_returns": 10_987_654, - "total_agi": 987_654_000_000, - "total_tax_liability": 78_543_000_000, - }, - "GA": { - "total_returns": 5_234_567, - "total_agi": 398_765_000_000, - "total_tax_liability": 33_456_000_000, - }, - "HI": { - "total_returns": 743_210, - "total_agi": 56_789_000_000, - "total_tax_liability": 4_765_000_000, - }, - "ID": { - "total_returns": 876_543, - "total_agi": 54_321_000_000, - "total_tax_liability": 4_543_000_000, - }, - "IL": { - "total_returns": 6_234_567, - "total_agi": 567_890_000_000, - "total_tax_liability": 47_654_000_000, - }, - "IN": { - "total_returns": 3_234_567, - "total_agi": 198_765_000_000, - "total_tax_liability": 16_543_000_000, - }, - "IA": { - "total_returns": 1_543_210, - "total_agi": 98_765_000_000, - "total_tax_liability": 8_234_000_000, - }, - "KS": { - "total_returns": 1_432_109, - "total_agi": 89_012_000_000, - "total_tax_liability": 7_432_000_000, - }, - "KY": { - "total_returns": 2_098_765, - "total_agi": 112_345_000_000, - "total_tax_liability": 9_345_000_000, - }, - "LA": { - "total_returns": 2_123_456, - "total_agi": 123_456_000_000, - "total_tax_liability": 10_234_000_000, - }, - "ME": { - "total_returns": 698_765, - "total_agi": 43_210_000_000, - "total_tax_liability": 3_598_000_000, - }, - "MD": { - "total_returns": 3_098_765, - "total_agi": 287_654_000_000, - "total_tax_liability": 24_321_000_000, - }, - "MA": { - "total_returns": 3_654_321, - "total_agi": 398_765_000_000, - "total_tax_liability": 34_567_000_000, - }, - "MI": { - "total_returns": 4_765_432, - "total_agi": 298_765_000_000, - "total_tax_liability": 24_876_000_000, - }, - "MN": { - "total_returns": 2_876_543, - "total_agi": 234_567_000_000, - "total_tax_liability": 19_654_000_000, - }, - "MS": { - "total_returns": 1_298_765, - "total_agi": 65_432_000_000, - "total_tax_liability": 5_432_000_000, - }, - "MO": { - "total_returns": 2_876_543, - "total_agi": 178_234_000_000, - "total_tax_liability": 14_876_000_000, - }, - "MT": { - "total_returns": 543_210, - "total_agi": 34_567_000_000, - "total_tax_liability": 2_876_000_000, - }, - "NE": { - "total_returns": 987_654, - "total_agi": 65_432_000_000, - "total_tax_liability": 5_456_000_000, - }, - "NV": { - "total_returns": 1_543_210, - "total_agi": 123_456_000_000, - "total_tax_liability": 10_234_000_000, - }, - "NH": { - "total_returns": 732_109, - "total_agi": 65_432_000_000, - "total_tax_liability": 5_543_000_000, - }, - "NJ": { - "total_returns": 4_654_321, - "total_agi": 498_765_000_000, - "total_tax_liability": 42_345_000_000, - }, - "NM": { - "total_returns": 987_654, - "total_agi": 54_321_000_000, - "total_tax_liability": 4_521_000_000, - }, - "NY": { - "total_returns": 10_234_567, - "total_agi": 1_456_789_000_000, - "total_tax_liability": 123_456_000_000, - }, - "NC": { - "total_returns": 4_987_654, - "total_agi": 356_789_000_000, - "total_tax_liability": 29_876_000_000, - }, - "ND": { - "total_returns": 398_765, - "total_agi": 28_765_000_000, - "total_tax_liability": 2_398_000_000, - }, - "OH": { - "total_returns": 5_654_321, - "total_agi": 387_654_000_000, - "total_tax_liability": 32_345_000_000, - }, - "OK": { - "total_returns": 1_765_432, - "total_agi": 98_765_000_000, - "total_tax_liability": 8_234_000_000, - }, - "OR": { - "total_returns": 2_098_765, - "total_agi": 156_789_000_000, - "total_tax_liability": 13_098_000_000, - }, - "PA": { - "total_returns": 6_543_210, - "total_agi": 654_321_000_000, - "total_tax_liability": 54_321_000_000, - }, - "RI": { - "total_returns": 543_210, - "total_agi": 45_678_000_000, - "total_tax_liability": 3_821_000_000, - }, - "SC": { - "total_returns": 2_432_109, - "total_agi": 143_210_000_000, - "total_tax_liability": 11_943_000_000, - }, - "SD": { - "total_returns": 454_321, - "total_agi": 32_109_000_000, - "total_tax_liability": 2_676_000_000, - }, - "TN": { - "total_returns": 3_234_567, - "total_agi": 234_567_000_000, - "total_tax_liability": 19_543_000_000, - }, - "TX": { - "total_returns": 13_876_543, - "total_agi": 1_234_567_000_000, - "total_tax_liability": 98_765_000_000, - }, - "UT": { - "total_returns": 1_543_210, - "total_agi": 109_876_000_000, - "total_tax_liability": 9_156_000_000, - }, - "VT": { - "total_returns": 343_210, - "total_agi": 23_456_000_000, - "total_tax_liability": 1_954_000_000, - }, - "VA": { - "total_returns": 4_234_567, - "total_agi": 398_765_000_000, - "total_tax_liability": 33_543_000_000, - }, - "WA": { - "total_returns": 3_765_432, - "total_agi": 356_789_000_000, - "total_tax_liability": 29_876_000_000, - }, - "WV": { - "total_returns": 832_109, - "total_agi": 43_210_000_000, - "total_tax_liability": 3_598_000_000, - }, - "WI": { - "total_returns": 2_987_654, - "total_agi": 198_765_000_000, - "total_tax_liability": 16_543_000_000, - }, - "WY": { - "total_returns": 298_765, - "total_agi": 23_456_000_000, - "total_tax_liability": 1_954_000_000, - }, - }, - 2020: { - "AL": { - "total_returns": 2_198_765, - "total_agi": 121_234_000_000, - "total_tax_liability": 9_876_000_000, - }, - "AK": { - "total_returns": 376_543, - "total_agi": 31_234_000_000, - "total_tax_liability": 2_543_000_000, - }, - "AZ": { - "total_returns": 3_345_678, - "total_agi": 212_345_000_000, - "total_tax_liability": 17_234_000_000, - }, - "AR": { - "total_returns": 1_298_765, - "total_agi": 69_876_000_000, - "total_tax_liability": 5_654_000_000, - }, - "CA": { - "total_returns": 18_234_567, - "total_agi": 1_987_654_000_000, - "total_tax_liability": 165_432_000_000, - }, - "CO": { - "total_returns": 2_876_543, - "total_agi": 212_345_000_000, - "total_tax_liability": 17_543_000_000, - }, - "CT": { - "total_returns": 1_843_210, - "total_agi": 178_654_000_000, - "total_tax_liability": 15_432_000_000, - }, - "DE": { - "total_returns": 476_543, - "total_agi": 38_765_000_000, - "total_tax_liability": 3_198_000_000, - }, - "DC": { - "total_returns": 376_543, - "total_agi": 48_765_000_000, - "total_tax_liability": 4_234_000_000, - }, - "FL": { - "total_returns": 10_765_432, - "total_agi": 898_765_000_000, - "total_tax_liability": 69_876_000_000, - }, - "GA": { - "total_returns": 5_098_765, - "total_agi": 356_789_000_000, - "total_tax_liability": 29_234_000_000, - }, - "HI": { - "total_returns": 721_098, - "total_agi": 50_987_000_000, - "total_tax_liability": 4_154_000_000, - }, - "ID": { - "total_returns": 843_210, - "total_agi": 48_765_000_000, - "total_tax_liability": 3_987_000_000, - }, - "IL": { - "total_returns": 6_098_765, - "total_agi": 512_345_000_000, - "total_tax_liability": 41_876_000_000, - }, - "IN": { - "total_returns": 3_154_321, - "total_agi": 178_654_000_000, - "total_tax_liability": 14_432_000_000, - }, - "IA": { - "total_returns": 1_498_765, - "total_agi": 87_654_000_000, - "total_tax_liability": 7_098_000_000, - }, - "KS": { - "total_returns": 1_387_654, - "total_agi": 78_901_000_000, - "total_tax_liability": 6_387_000_000, - }, - "KY": { - "total_returns": 2_043_210, - "total_agi": 98_765_000_000, - "total_tax_liability": 8_012_000_000, - }, - "LA": { - "total_returns": 2_065_432, - "total_agi": 109_876_000_000, - "total_tax_liability": 8_876_000_000, - }, - "ME": { - "total_returns": 676_543, - "total_agi": 38_654_000_000, - "total_tax_liability": 3_121_000_000, - }, - "MD": { - "total_returns": 3_021_098, - "total_agi": 256_789_000_000, - "total_tax_liability": 21_234_000_000, - }, - "MA": { - "total_returns": 3_565_432, - "total_agi": 356_789_000_000, - "total_tax_liability": 30_123_000_000, - }, - "MI": { - "total_returns": 4_654_321, - "total_agi": 267_890_000_000, - "total_tax_liability": 21_654_000_000, - }, - "MN": { - "total_returns": 2_798_765, - "total_agi": 210_987_000_000, - "total_tax_liability": 17_123_000_000, - }, - "MS": { - "total_returns": 1_254_321, - "total_agi": 58_765_000_000, - "total_tax_liability": 4_765_000_000, - }, - "MO": { - "total_returns": 2_798_765, - "total_agi": 159_876_000_000, - "total_tax_liability": 12_876_000_000, - }, - "MT": { - "total_returns": 521_098, - "total_agi": 30_987_000_000, - "total_tax_liability": 2_510_000_000, - }, - "NE": { - "total_returns": 954_321, - "total_agi": 58_765_000_000, - "total_tax_liability": 4_765_000_000, - }, - "NV": { - "total_returns": 1_498_765, - "total_agi": 109_876_000_000, - "total_tax_liability": 8_876_000_000, - }, - "NH": { - "total_returns": 710_987, - "total_agi": 58_765_000_000, - "total_tax_liability": 4_832_000_000, - }, - "NJ": { - "total_returns": 4_543_210, - "total_agi": 445_678_000_000, - "total_tax_liability": 36_987_000_000, - }, - "NM": { - "total_returns": 954_321, - "total_agi": 48_765_000_000, - "total_tax_liability": 3_943_000_000, - }, - "NY": { - "total_returns": 10_098_765, - "total_agi": 1_345_678_000_000, - "total_tax_liability": 112_345_000_000, - }, - "NC": { - "total_returns": 4_854_321, - "total_agi": 319_876_000_000, - "total_tax_liability": 26_098_000_000, - }, - "ND": { - "total_returns": 387_654, - "total_agi": 25_678_000_000, - "total_tax_liability": 2_087_000_000, - }, - "OH": { - "total_returns": 5_521_098, - "total_agi": 345_678_000_000, - "total_tax_liability": 28_123_000_000, - }, - "OK": { - "total_returns": 1_721_098, - "total_agi": 87_654_000_000, - "total_tax_liability": 7_098_000_000, - }, - "OR": { - "total_returns": 2_043_210, - "total_agi": 140_987_000_000, - "total_tax_liability": 11_432_000_000, - }, - "PA": { - "total_returns": 6_432_109, - "total_agi": 623_456_000_000, - "total_tax_liability": 49_876_000_000, - }, - "RI": { - "total_returns": 521_098, - "total_agi": 40_987_000_000, - "total_tax_liability": 3_321_000_000, - }, - "SC": { - "total_returns": 2_365_432, - "total_agi": 128_765_000_000, - "total_tax_liability": 10_432_000_000, - }, - "SD": { - "total_returns": 443_210, - "total_agi": 28_765_000_000, - "total_tax_liability": 2_332_000_000, - }, - "TN": { - "total_returns": 3_154_321, - "total_agi": 210_987_000_000, - "total_tax_liability": 17_098_000_000, - }, - "TX": { - "total_returns": 13_456_789, - "total_agi": 1_123_456_000_000, - "total_tax_liability": 87_654_000_000, - }, - "UT": { - "total_returns": 1_498_765, - "total_agi": 98_765_000_000, - "total_tax_liability": 7_987_000_000, - }, - "VT": { - "total_returns": 332_109, - "total_agi": 21_098_000_000, - "total_tax_liability": 1_709_000_000, - }, - "VA": { - "total_returns": 4_132_109, - "total_agi": 356_789_000_000, - "total_tax_liability": 29_234_000_000, - }, - "WA": { - "total_returns": 3_654_321, - "total_agi": 319_876_000_000, - "total_tax_liability": 26_098_000_000, - }, - "WV": { - "total_returns": 810_987, - "total_agi": 38_765_000_000, - "total_tax_liability": 3_121_000_000, - }, - "WI": { - "total_returns": 2_909_876, - "total_agi": 178_654_000_000, - "total_tax_liability": 14_432_000_000, - }, - "WY": { - "total_returns": 287_654, - "total_agi": 21_098_000_000, - "total_tax_liability": 1_709_000_000, - }, - }, -} - -# AGI brackets as defined by IRS SOI (simplified from full 19 brackets) -AGI_BRACKETS = [ - {"min": 0, "max": 1, "label": "under_1", "display": "Under $1"}, - {"min": 1, "max": 10_000, "label": "1_to_10k", "display": "$1-$10k"}, - {"min": 10_000, "max": 25_000, "label": "10k_to_25k", "display": "$10k-$25k"}, - {"min": 25_000, "max": 50_000, "label": "25k_to_50k", "display": "$25k-$50k"}, - {"min": 50_000, "max": 75_000, "label": "50k_to_75k", "display": "$50k-$75k"}, - {"min": 75_000, "max": 100_000, "label": "75k_to_100k", "display": "$75k-$100k"}, - {"min": 100_000, "max": 200_000, "label": "100k_to_200k", "display": "$100k-$200k"}, - {"min": 200_000, "max": 500_000, "label": "200k_to_500k", "display": "$200k-$500k"}, - {"min": 500_000, "max": 1_000_000, "label": "500k_to_1m", "display": "$500k-$1M"}, - {"min": 1_000_000, "max": None, "label": "1m_plus", "display": "$1M+"}, -] - -SOURCE_URL = "https://www.irs.gov/statistics/soi-tax-stats-historic-table-2" - - -def _generate_agi_bracket_data( - total_returns: int, total_agi: int, total_tax: int -) -> dict: - """ - Generate synthetic AGI bracket distribution based on state totals. - - Uses approximate distribution patterns from actual IRS SOI data. - In production, this would be replaced by parsing actual IRS CSV files. - """ - # Approximate distribution of returns by AGI bracket (based on national averages) - returns_pct = { - "under_1": 0.02, - "1_to_10k": 0.12, - "10k_to_25k": 0.15, - "25k_to_50k": 0.18, - "50k_to_75k": 0.14, - "75k_to_100k": 0.11, - "100k_to_200k": 0.17, - "200k_to_500k": 0.08, - "500k_to_1m": 0.02, - "1m_plus": 0.01, - } - - # Approximate distribution of AGI by bracket (higher brackets have more AGI) - agi_pct = { - "under_1": 0.00, - "1_to_10k": 0.01, - "10k_to_25k": 0.03, - "25k_to_50k": 0.08, - "50k_to_75k": 0.10, - "75k_to_100k": 0.10, - "100k_to_200k": 0.22, - "200k_to_500k": 0.18, - "500k_to_1m": 0.10, - "1m_plus": 0.18, - } - - # Approximate distribution of tax by bracket (progressive taxation) - tax_pct = { - "under_1": 0.00, - "1_to_10k": 0.00, - "10k_to_25k": 0.01, - "25k_to_50k": 0.04, - "50k_to_75k": 0.06, - "75k_to_100k": 0.08, - "100k_to_200k": 0.20, - "200k_to_500k": 0.22, - "500k_to_1m": 0.14, - "1m_plus": 0.25, - } - - return { - bracket: { - "total_returns": int(total_returns * returns_pct[bracket]), - "total_agi": int(total_agi * agi_pct[bracket]), - "total_tax_liability": int(total_tax * tax_pct[bracket]), - } - for bracket in returns_pct.keys() - } - - -def _build_soi_state_agi_data() -> dict: - """Build AGI bracket data for all states based on state totals.""" - result = {} - for year, year_data in SOI_STATE_DATA.items(): - result[year] = {} - for state, state_data in year_data.items(): - result[year][state] = _generate_agi_bracket_data( - state_data["total_returns"], - state_data["total_agi"], - state_data["total_tax_liability"], - ) - return result - - -# State-level AGI bracket data (derived from totals) -SOI_STATE_AGI_DATA = _build_soi_state_agi_data() - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - # Check if exists - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - # Create new - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() # Get ID - - # Add constraints - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_soi_state_targets(session: Session, years: list[int] | None = None): - """ - Load state-level SOI targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SOI_STATE_DATA.keys()) - - for year in years: - if year not in SOI_STATE_DATA: - continue - - data = SOI_STATE_DATA[year] - - # Get or create national stratum (for parent relationship) - # This should ideally already exist from etl_soi, but we create it if needed - national_stratum = get_or_create_stratum( - session, - name="US All Filers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("is_tax_filer", "==", "1")], - description="All individual income tax returns filed in the US", - stratum_group_id="national", - ) - - # Create state-level strata and targets - for state_abbrev, state_data in data.items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - # Create state stratum - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} All Filers", - jurisdiction=Jurisdiction.US, - constraints=[ - ("is_tax_filer", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"All individual income tax returns filed in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="soi_states", - ) - - # Add total returns target - session.add( - Target( - stratum_id=state_stratum.id, - variable="tax_unit_count", - period=year, - value=state_data["total_returns"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.IRS_SOI, - source_table="Historic Table 2", - source_url=SOURCE_URL, - ) - ) - - # Add total AGI target - session.add( - Target( - stratum_id=state_stratum.id, - variable="adjusted_gross_income", - period=year, - value=state_data["total_agi"], - target_type=TargetType.AMOUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.IRS_SOI, - source_table="Historic Table 2", - source_url=SOURCE_URL, - ) - ) - - # Add total tax liability target - session.add( - Target( - stratum_id=state_stratum.id, - variable="income_tax_liability", - period=year, - value=state_data["total_tax_liability"], - target_type=TargetType.AMOUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.IRS_SOI, - source_table="Historic Table 2", - source_url=SOURCE_URL, - ) - ) - - # Create AGI bracket strata for this state - agi_data = SOI_STATE_AGI_DATA[year][state_abbrev] - for bracket in AGI_BRACKETS: - bracket_label = bracket["label"] - bracket_display = bracket["display"] - bracket_data = agi_data[bracket_label] - - # Create bracket stratum with both state and AGI constraints - bracket_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} Filers AGI {bracket_display}", - jurisdiction=Jurisdiction.US, - constraints=[ - ("is_tax_filer", "==", "1"), - ("state_fips", "==", fips), - ("agi_bracket", "==", bracket_label), - ], - description=( - f"Individual income tax returns filed in {state_abbrev} " - f"with AGI {bracket_display.lower()}" - ), - parent_id=state_stratum.id, - stratum_group_id="soi_states_agi_brackets", - ) - - # Add bracket-level returns target - session.add( - Target( - stratum_id=bracket_stratum.id, - variable="tax_unit_count", - period=year, - value=bracket_data["total_returns"], - target_type=TargetType.COUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.IRS_SOI, - source_table="Historic Table 2", - source_url=SOURCE_URL, - ) - ) - - # Add bracket-level AGI target - session.add( - Target( - stratum_id=bracket_stratum.id, - variable="adjusted_gross_income", - period=year, - value=bracket_data["total_agi"], - target_type=TargetType.AMOUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.IRS_SOI, - source_table="Historic Table 2", - source_url=SOURCE_URL, - ) - ) - - # Add bracket-level tax liability target - session.add( - Target( - stratum_id=bracket_stratum.id, - variable="income_tax_liability", - period=year, - value=bracket_data["total_tax_liability"], - target_type=TargetType.AMOUNT, - geographic_level=GeographicLevel.STATE, - source=DataSource.IRS_SOI, - source_table="Historic Table 2", - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the state-level SOI ETL pipeline.""" - from pathlib import Path - - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_soi_state_targets(session) - print(f"Loaded state-level SOI targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_ssa.py b/db/etl_ssa.py deleted file mode 100644 index 6af2b261..00000000 --- a/db/etl_ssa.py +++ /dev/null @@ -1,388 +0,0 @@ -""" -ETL for SSA (Social Security Administration) targets. - -Loads data from SSA statistics into the targets database. -Data sources: -- https://www.ssa.gov/policy/docs/statcomps/supplement/ -- https://www.ssa.gov/oact/STATS/table4a3.html -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# SSA data by year -# Source: SSA Annual Statistical Supplement -# https://www.ssa.gov/policy/docs/statcomps/supplement/ -SSA_DATA = { - 2023: { - # OASDI (Old-Age, Survivors, and Disability Insurance) - "total_beneficiaries": 67_500_000, - "total_benefits": 1_415_000_000_000, # $1.415 trillion annually - - # Retired workers - "retired_workers": { - "beneficiaries": 51_200_000, - "benefits": 987_000_000_000, - "avg_monthly_benefit": 1_907, # dollars - }, - - # Disabled workers - "disabled_workers": { - "beneficiaries": 7_600_000, - "benefits": 143_000_000_000, - "avg_monthly_benefit": 1_537, - }, - - # Survivors - "survivors": { - "beneficiaries": 5_900_000, - "benefits": 112_000_000_000, - "avg_monthly_benefit": 1_498, - }, - - # Spouses and children - "dependents": { - "beneficiaries": 2_800_000, - "benefits": 43_000_000_000, - }, - - # SSI (Supplemental Security Income) - "ssi": { - "recipients": 7_400_000, - "payments": 59_800_000_000, - "avg_monthly_payment": 674, - "aged_recipients": 1_100_000, - "blind_recipients": 67_000, - "disabled_recipients": 6_233_000, - }, - }, - 2022: { - "total_beneficiaries": 66_000_000, - "total_benefits": 1_234_000_000_000, - - "retired_workers": { - "beneficiaries": 50_100_000, - "benefits": 876_000_000_000, - "avg_monthly_benefit": 1_827, - }, - - "disabled_workers": { - "beneficiaries": 7_800_000, - "benefits": 134_000_000_000, - "avg_monthly_benefit": 1_483, - }, - - "survivors": { - "beneficiaries": 5_800_000, - "benefits": 98_000_000_000, - "avg_monthly_benefit": 1_456, - }, - - "dependents": { - "beneficiaries": 2_300_000, - "benefits": 38_000_000_000, - }, - - "ssi": { - "recipients": 7_600_000, - "payments": 56_700_000_000, - "avg_monthly_payment": 621, - "aged_recipients": 1_150_000, - "blind_recipients": 69_000, - "disabled_recipients": 6_381_000, - }, - }, - 2021: { - "total_beneficiaries": 65_000_000, - "total_benefits": 1_145_000_000_000, - - "retired_workers": { - "beneficiaries": 49_000_000, - "benefits": 812_000_000_000, - "avg_monthly_benefit": 1_657, - }, - - "disabled_workers": { - "beneficiaries": 8_000_000, - "benefits": 128_000_000_000, - "avg_monthly_benefit": 1_358, - }, - - "survivors": { - "beneficiaries": 5_700_000, - "benefits": 92_000_000_000, - "avg_monthly_benefit": 1_387, - }, - - "dependents": { - "beneficiaries": 2_300_000, - "benefits": 35_000_000_000, - }, - - "ssi": { - "recipients": 7_800_000, - "payments": 54_800_000_000, - "avg_monthly_payment": 586, - "aged_recipients": 1_200_000, - "blind_recipients": 71_000, - "disabled_recipients": 6_529_000, - }, - }, -} - -SOURCE_URL = "https://www.ssa.gov/policy/docs/statcomps/supplement/" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_ssa_targets(session: Session, years: list[int] | None = None): - """ - Load SSA targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SSA_DATA.keys()) - - for year in years: - if year not in SSA_DATA: - continue - - data = SSA_DATA[year] - - # OASDI national stratum - oasdi_stratum = get_or_create_stratum( - session, - name="US OASDI Beneficiaries", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("social_security", "==", "1")], - description="All OASDI (Social Security) beneficiaries", - stratum_group_id="ssa_national", - ) - - session.add( - Target( - stratum_id=oasdi_stratum.id, - variable="oasdi_beneficiaries", - period=year, - value=data["total_beneficiaries"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=oasdi_stratum.id, - variable="oasdi_benefits", - period=year, - value=data["total_benefits"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - # Retired workers stratum - retired_stratum = get_or_create_stratum( - session, - name="US Retired Workers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("social_security_retired", "==", "1")], - description="Social Security retired worker beneficiaries", - parent_id=oasdi_stratum.id, - stratum_group_id="ssa_categories", - ) - - retired_data = data["retired_workers"] - session.add( - Target( - stratum_id=retired_stratum.id, - variable="oasdi_beneficiaries", - period=year, - value=retired_data["beneficiaries"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=retired_stratum.id, - variable="oasdi_benefits", - period=year, - value=retired_data["benefits"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=retired_stratum.id, - variable="oasdi_avg_monthly_benefit", - period=year, - value=retired_data["avg_monthly_benefit"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - # Disabled workers stratum - disabled_stratum = get_or_create_stratum( - session, - name="US Disabled Workers", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("social_security_disabled", "==", "1")], - description="Social Security disabled worker beneficiaries", - parent_id=oasdi_stratum.id, - stratum_group_id="ssa_categories", - ) - - disabled_data = data["disabled_workers"] - session.add( - Target( - stratum_id=disabled_stratum.id, - variable="oasdi_beneficiaries", - period=year, - value=disabled_data["beneficiaries"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=disabled_stratum.id, - variable="oasdi_avg_monthly_benefit", - period=year, - value=disabled_data["avg_monthly_benefit"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - # SSI stratum - ssi_stratum = get_or_create_stratum( - session, - name="US SSI Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("ssi", "==", "1")], - description="Supplemental Security Income recipients", - stratum_group_id="ssa_ssi", - ) - - ssi_data = data["ssi"] - session.add( - Target( - stratum_id=ssi_stratum.id, - variable="ssi_recipients", - period=year, - value=ssi_data["recipients"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=ssi_stratum.id, - variable="ssi_payments", - period=year, - value=ssi_data["payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.add( - Target( - stratum_id=ssi_stratum.id, - variable="ssi_avg_monthly_payment", - period=year, - value=ssi_data["avg_monthly_payment"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the SSA ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_ssa_targets(session) - print(f"Loaded SSA targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_ssi.py b/db/etl_ssi.py deleted file mode 100644 index 6c5bfe93..00000000 --- a/db/etl_ssi.py +++ /dev/null @@ -1,660 +0,0 @@ -""" -ETL for SSI (Supplemental Security Income) targets. - -Loads SSI-specific data from SSA statistics into the targets database. -This is separate from the general SSA ETL to provide more detailed SSI data -including state-level breakdowns and federal vs state supplementation. - -Data sources: -- https://www.ssa.gov/policy/docs/statcomps/ssi_sc/2023/ -- https://www.ssa.gov/policy/docs/statcomps/ssi_asr/ -- https://www.ssa.gov/policy/docs/statcomps/supplement/ -""" - -from __future__ import annotations - -from sqlmodel import Session, select - -from .schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - -# State FIPS codes -STATE_FIPS = { - "AL": "01", "AK": "02", "AZ": "04", "AR": "05", "CA": "06", - "CO": "08", "CT": "09", "DE": "10", "DC": "11", "FL": "12", - "GA": "13", "HI": "15", "ID": "16", "IL": "17", "IN": "18", - "IA": "19", "KS": "20", "KY": "21", "LA": "22", "ME": "23", - "MD": "24", "MA": "25", "MI": "26", "MN": "27", "MS": "28", - "MO": "29", "MT": "30", "NE": "31", "NV": "32", "NH": "33", - "NJ": "34", "NM": "35", "NY": "36", "NC": "37", "ND": "38", - "OH": "39", "OK": "40", "OR": "41", "PA": "42", "RI": "44", - "SC": "45", "SD": "46", "TN": "47", "TX": "48", "UT": "49", - "VT": "50", "VA": "51", "WA": "53", "WV": "54", "WI": "55", - "WY": "56", "PR": "72", "VI": "78", "GU": "66", -} - -# SSI national data by year -# Source: SSA Annual Statistical Supplement and SSI Annual Statistical Report -# https://www.ssa.gov/policy/docs/statcomps/ssi_asr/ -# https://www.ssa.gov/policy/docs/statcomps/supplement/2024/highlights.html -SSI_DATA = { - 2023: { - "national": { - # Total recipients and payments - "recipients": 7_425_331, - "total_payments": 61_400_000_000, # $61.4 billion - "avg_monthly_payment": 675, # dollars - # By eligibility category - "aged_recipients": 1_160_608, - "blind_recipients": 63_877, - "disabled_recipients": 6_199_817, - # Payments by category - "aged_payments": 7_294_542_000, - "blind_payments": 544_835_000, - "disabled_payments": 53_536_042_000, - # Federal vs state supplementation - "federal_payments": 58_200_000_000, # Approximately - "state_supplementation": 3_200_000_000, # $3.2 billion - }, - # State-level data (December 2023) - # Source: https://www.ssa.gov/policy/docs/statcomps/ssi_sc/2023/ - "states": { - "CA": { - "recipients": 1_199_000, - "aged": 305_000, - "blind": 17_000, - "disabled": 877_000, - "payments": 12_800_000_000, # Includes state supplement - }, - "TX": { - "recipients": 573_750, - "aged": 100_791, - "blind": 8_500, - "disabled": 464_459, - "payments": 4_200_000_000, - }, - "NY": { - "recipients": 557_200, - "aged": 113_136, - "blind": 7_800, - "disabled": 436_264, - "payments": 4_800_000_000, # Includes state supplement - }, - "FL": { - "recipients": 538_984, - "aged": 158_775, - "blind": 6_200, - "disabled": 374_009, - "payments": 4_100_000_000, - }, - "PA": { - "recipients": 321_000, - "aged": 48_000, - "blind": 4_500, - "disabled": 268_500, - "payments": 2_400_000_000, - }, - "OH": { - "recipients": 278_000, - "aged": 35_000, - "blind": 3_800, - "disabled": 239_200, - "payments": 2_000_000_000, - }, - "IL": { - "recipients": 248_000, - "aged": 42_000, - "blind": 3_200, - "disabled": 202_800, - "payments": 1_850_000_000, - }, - "GA": { - "recipients": 245_000, - "aged": 38_000, - "blind": 3_100, - "disabled": 203_900, - "payments": 1_750_000_000, - }, - "MI": { - "recipients": 232_000, - "aged": 28_000, - "blind": 3_000, - "disabled": 201_000, - "payments": 1_700_000_000, - }, - "NC": { - "recipients": 218_000, - "aged": 34_000, - "blind": 2_900, - "disabled": 181_100, - "payments": 1_550_000_000, - }, - }, - }, - 2022: { - "national": { - "recipients": 7_542_222, - "total_payments": 56_700_000_000, # $56.7 billion - "avg_monthly_payment": 622, - "aged_recipients": 1_180_000, - "blind_recipients": 66_000, - "disabled_recipients": 6_296_222, - "aged_payments": 6_800_000_000, - "blind_payments": 520_000_000, - "disabled_payments": 49_380_000_000, - "federal_payments": 53_800_000_000, - "state_supplementation": 2_900_000_000, - }, - "states": { - "CA": { - "recipients": 1_215_000, - "aged": 310_000, - "blind": 17_500, - "disabled": 887_500, - "payments": 12_200_000_000, - }, - "TX": { - "recipients": 582_000, - "aged": 102_000, - "blind": 8_700, - "disabled": 471_300, - "payments": 4_000_000_000, - }, - "NY": { - "recipients": 568_000, - "aged": 115_000, - "blind": 8_000, - "disabled": 445_000, - "payments": 4_600_000_000, - }, - "FL": { - "recipients": 545_000, - "aged": 160_000, - "blind": 6_300, - "disabled": 378_700, - "payments": 3_900_000_000, - }, - "PA": { - "recipients": 328_000, - "aged": 49_000, - "blind": 4_600, - "disabled": 274_400, - "payments": 2_300_000_000, - }, - }, - }, - 2021: { - "national": { - "recipients": 7_800_000, - "total_payments": 61_000_000_000, - "avg_monthly_payment": 586, - "aged_recipients": 1_200_000, - "blind_recipients": 71_000, - "disabled_recipients": 6_529_000, - "aged_payments": 6_500_000_000, - "blind_payments": 500_000_000, - "disabled_payments": 54_000_000_000, - "federal_payments": 57_800_000_000, - "state_supplementation": 3_200_000_000, - }, - "states": { - "CA": { - "recipients": 1_230_000, - "aged": 315_000, - "blind": 18_000, - "disabled": 897_000, - "payments": 12_000_000_000, - }, - "TX": { - "recipients": 590_000, - "aged": 104_000, - "blind": 9_000, - "disabled": 477_000, - "payments": 3_800_000_000, - }, - "NY": { - "recipients": 580_000, - "aged": 118_000, - "blind": 8_200, - "disabled": 453_800, - "payments": 4_400_000_000, - }, - "FL": { - "recipients": 552_000, - "aged": 162_000, - "blind": 6_500, - "disabled": 383_500, - "payments": 3_700_000_000, - }, - "PA": { - "recipients": 335_000, - "aged": 50_000, - "blind": 4_800, - "disabled": 280_200, - "payments": 2_200_000_000, - }, - }, - }, -} - -SOURCE_URL = "https://www.ssa.gov/policy/docs/statcomps/ssi_sc/2023/" -SOURCE_URL_ASR = "https://www.ssa.gov/policy/docs/statcomps/ssi_asr/" - - -def get_or_create_stratum( - session: Session, - name: str, - jurisdiction: Jurisdiction, - constraints: list[tuple[str, str, str]], - description: str | None = None, - parent_id: int | None = None, - stratum_group_id: str | None = None, -) -> Stratum: - """Get existing stratum or create new one.""" - definition_hash = Stratum.compute_hash(constraints, jurisdiction) - - existing = session.exec( - select(Stratum).where(Stratum.definition_hash == definition_hash) - ).first() - - if existing: - return existing - - stratum = Stratum( - name=name, - description=description, - jurisdiction=jurisdiction, - definition_hash=definition_hash, - parent_id=parent_id, - stratum_group_id=stratum_group_id, - ) - session.add(stratum) - session.flush() - - for variable, operator, value in constraints: - constraint = StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - session.add(constraint) - - return stratum - - -def load_ssi_targets(session: Session, years: list[int] | None = None): - """ - Load SSI targets into database. - - Args: - session: Database session - years: Years to load (default: all available) - """ - if years is None: - years = list(SSI_DATA.keys()) - - for year in years: - if year not in SSI_DATA: - continue - - data = SSI_DATA[year] - national_data = data["national"] - - # Create national SSI stratum - national_stratum = get_or_create_stratum( - session, - name="US SSI Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[("ssi", "==", "1")], - description="All SSI recipients in the United States", - stratum_group_id="ssi_national", - ) - - # National recipient count - session.add( - Target( - stratum_id=national_stratum.id, - variable="ssi_recipients", - period=year, - value=national_data["recipients"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # National total payments - session.add( - Target( - stratum_id=national_stratum.id, - variable="ssi_total_payments", - period=year, - value=national_data["total_payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # Average monthly payment - session.add( - Target( - stratum_id=national_stratum.id, - variable="ssi_avg_monthly_payment", - period=year, - value=national_data["avg_monthly_payment"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # Federal vs state supplementation - session.add( - Target( - stratum_id=national_stratum.id, - variable="ssi_federal_payments", - period=year, - value=national_data["federal_payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - session.add( - Target( - stratum_id=national_stratum.id, - variable="ssi_state_supplementation", - period=year, - value=national_data["state_supplementation"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # Create SSI Aged stratum - aged_stratum = get_or_create_stratum( - session, - name="US SSI Aged Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("ssi", "==", "1"), - ("ssi_category", "==", "aged"), - ], - description="SSI recipients aged 65 or older", - parent_id=national_stratum.id, - stratum_group_id="ssi_categories", - ) - - session.add( - Target( - stratum_id=aged_stratum.id, - variable="ssi_recipients", - period=year, - value=national_data["aged_recipients"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - session.add( - Target( - stratum_id=aged_stratum.id, - variable="ssi_payments", - period=year, - value=national_data["aged_payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # Create SSI Blind stratum - blind_stratum = get_or_create_stratum( - session, - name="US SSI Blind Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("ssi", "==", "1"), - ("ssi_category", "==", "blind"), - ], - description="SSI recipients qualifying based on blindness", - parent_id=national_stratum.id, - stratum_group_id="ssi_categories", - ) - - session.add( - Target( - stratum_id=blind_stratum.id, - variable="ssi_recipients", - period=year, - value=national_data["blind_recipients"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - session.add( - Target( - stratum_id=blind_stratum.id, - variable="ssi_payments", - period=year, - value=national_data["blind_payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # Create SSI Disabled stratum - disabled_stratum = get_or_create_stratum( - session, - name="US SSI Disabled Recipients", - jurisdiction=Jurisdiction.US_FEDERAL, - constraints=[ - ("ssi", "==", "1"), - ("ssi_category", "==", "disabled"), - ], - description="SSI recipients qualifying based on disability", - parent_id=national_stratum.id, - stratum_group_id="ssi_categories", - ) - - session.add( - Target( - stratum_id=disabled_stratum.id, - variable="ssi_recipients", - period=year, - value=national_data["disabled_recipients"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - session.add( - Target( - stratum_id=disabled_stratum.id, - variable="ssi_payments", - period=year, - value=national_data["disabled_payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Annual Statistical Report", - source_url=SOURCE_URL_ASR, - ) - ) - - # Add state-level targets - for state_abbrev, state_data in data.get("states", {}).items(): - if state_abbrev not in STATE_FIPS: - continue - - fips = STATE_FIPS[state_abbrev] - - state_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} SSI Recipients", - jurisdiction=Jurisdiction.US, - constraints=[ - ("ssi", "==", "1"), - ("state_fips", "==", fips), - ], - description=f"SSI recipients in {state_abbrev}", - parent_id=national_stratum.id, - stratum_group_id="ssi_states", - ) - - # Total state recipients - session.add( - Target( - stratum_id=state_stratum.id, - variable="ssi_recipients", - period=year, - value=state_data["recipients"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Recipients by State and County", - source_url=SOURCE_URL, - ) - ) - - # State payments - session.add( - Target( - stratum_id=state_stratum.id, - variable="ssi_payments", - period=year, - value=state_data["payments"], - target_type=TargetType.AMOUNT, - source=DataSource.SSA, - source_table="SSI Recipients by State and County", - source_url=SOURCE_URL, - ) - ) - - # State-level category breakdowns - # Aged recipients in state - state_aged_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} SSI Aged Recipients", - jurisdiction=Jurisdiction.US, - constraints=[ - ("ssi", "==", "1"), - ("ssi_category", "==", "aged"), - ("state_fips", "==", fips), - ], - description=f"SSI aged recipients in {state_abbrev}", - parent_id=state_stratum.id, - stratum_group_id="ssi_state_categories", - ) - - session.add( - Target( - stratum_id=state_aged_stratum.id, - variable="ssi_recipients", - period=year, - value=state_data["aged"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Recipients by State and County", - source_url=SOURCE_URL, - ) - ) - - # Blind recipients in state - state_blind_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} SSI Blind Recipients", - jurisdiction=Jurisdiction.US, - constraints=[ - ("ssi", "==", "1"), - ("ssi_category", "==", "blind"), - ("state_fips", "==", fips), - ], - description=f"SSI blind recipients in {state_abbrev}", - parent_id=state_stratum.id, - stratum_group_id="ssi_state_categories", - ) - - session.add( - Target( - stratum_id=state_blind_stratum.id, - variable="ssi_recipients", - period=year, - value=state_data["blind"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Recipients by State and County", - source_url=SOURCE_URL, - ) - ) - - # Disabled recipients in state - state_disabled_stratum = get_or_create_stratum( - session, - name=f"{state_abbrev} SSI Disabled Recipients", - jurisdiction=Jurisdiction.US, - constraints=[ - ("ssi", "==", "1"), - ("ssi_category", "==", "disabled"), - ("state_fips", "==", fips), - ], - description=f"SSI disabled recipients in {state_abbrev}", - parent_id=state_stratum.id, - stratum_group_id="ssi_state_categories", - ) - - session.add( - Target( - stratum_id=state_disabled_stratum.id, - variable="ssi_recipients", - period=year, - value=state_data["disabled"], - target_type=TargetType.COUNT, - source=DataSource.SSA, - source_table="SSI Recipients by State and County", - source_url=SOURCE_URL, - ) - ) - - session.commit() - - -def run_etl(db_path=None): - """Run the SSI ETL pipeline.""" - from pathlib import Path - from .schema import DEFAULT_DB_PATH - - path = Path(db_path) if db_path else DEFAULT_DB_PATH - engine = init_db(path) - - with Session(engine) as session: - load_ssi_targets(session) - print(f"Loaded SSI targets to {path}") - - -if __name__ == "__main__": - run_etl() diff --git a/db/etl_targets_supabase.py b/db/etl_targets_supabase.py deleted file mode 100644 index d875d5c4..00000000 --- a/db/etl_targets_supabase.py +++ /dev/null @@ -1,320 +0,0 @@ -""" -ETL for loading calibration targets to Supabase. - -Migrates targets from local SQLite to Supabase microplex schema. -Handles strata and stratum_constraints alongside targets. - -Usage: - python -m db.etl_targets_supabase --source us - python -m db.etl_targets_supabase --all -""" - -from __future__ import annotations - -from typing import Any, Dict, List, Optional - -from .supabase_client import get_supabase_client, query_sources - - -def get_or_create_stratum( - client, - name: str, - jurisdiction: str, - constraints: List[Dict[str, str]], - description: Optional[str] = None, -) -> str: - """ - Get existing stratum or create new one in Supabase. - - Args: - client: Supabase client - name: Stratum name (must be unique per jurisdiction) - jurisdiction: e.g., "us", "uk" - constraints: List of {variable, operator, value} dicts - description: Optional description - - Returns: - Stratum UUID - """ - # Check if exists by name + jurisdiction (unique constraint) - result = client.schema("microplex").table("strata").select("id").eq("name", name).eq("jurisdiction", jurisdiction).execute() - if result.data: - return result.data[0]["id"] - - # Create new stratum - stratum_data = { - "name": name, - "jurisdiction": jurisdiction, - } - if description: - stratum_data["description"] = description - - result = client.schema("microplex").table("strata").insert(stratum_data).execute() - stratum_id = result.data[0]["id"] - - # Add constraints - for constraint in constraints: - constraint_data = { - "stratum_id": stratum_id, - "variable": constraint["variable"], - "operator": constraint["operator"], - "value": constraint["value"], - } - client.schema("microplex").table("stratum_constraints").insert(constraint_data).execute() - - return stratum_id - - -def load_soi_targets_supabase( - years: Optional[List[int]] = None, - dry_run: bool = False, -) -> Dict[str, Any]: - """ - Load IRS SOI targets to Supabase. - - Args: - years: Years to load (default: 2018-2021) - dry_run: If True, just return count without inserting - - Returns: - Dict with counts and status - """ - from .etl_soi import SOURCE_URL, available_soi_years, load_soi_table_1_1_data - - if years is None: - years = available_soi_years() - - client = get_supabase_client() - - # Get or create IRS SOI source - sources = query_sources(jurisdiction="us", institution="irs") - soi_source = next((s for s in sources if s["dataset"] == "soi"), None) - - if not soi_source: - result = client.schema("microplex").table("sources").insert({ - "jurisdiction": "us", - "institution": "irs", - "dataset": "soi", - "name": "IRS Statistics of Income", - "url": SOURCE_URL, - }).execute() - source_id = result.data[0]["id"] - else: - source_id = soi_source["id"] - - targets_loaded = 0 - strata_created = 0 - - for year in years: - if year not in available_soi_years(): - continue - - year_data = load_soi_table_1_1_data(year) - - # National totals - if dry_run: - targets_loaded += 2 # total_returns and total_agi - else: - stratum_id = get_or_create_stratum( - client, - name=f"US Tax Filers {year}", - jurisdiction="us", - constraints=[], - ) - strata_created += 1 - - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "tax_unit_count", - "value": year_data["total_returns"], - "target_type": "count", - "period": year, - }).execute() - targets_loaded += 1 - - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "agi_total", - "value": year_data["total_agi"], - "target_type": "amount", - "period": year, - }).execute() - targets_loaded += 1 - - # AGI bracket targets - returns_by_bracket = year_data.get("returns_by_agi_bracket", {}) - agi_by_bracket = year_data.get("agi_by_bracket", {}) - - for bracket, returns in returns_by_bracket.items(): - if dry_run: - targets_loaded += 2 - continue - - stratum_name = f"US Tax Filers AGI {bracket} {year}" - stratum_id = get_or_create_stratum( - client, - name=stratum_name, - jurisdiction="us", - constraints=[{"variable": "agi_bracket", "operator": "==", "value": bracket}], - ) - strata_created += 1 - - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "tax_unit_count", - "value": returns, - "target_type": "count", - "period": year, - }).execute() - targets_loaded += 1 - - if bracket in agi_by_bracket: - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "agi_total", - "value": agi_by_bracket[bracket], - "target_type": "amount", - "period": year, - }).execute() - targets_loaded += 1 - - return { - "targets_loaded": targets_loaded, - "strata_created": strata_created, - "years": years, - "dry_run": dry_run, - } - - -def load_snap_targets_supabase( - years: Optional[List[int]] = None, - dry_run: bool = False, -) -> Dict[str, Any]: - """Load USDA SNAP targets to Supabase.""" - from .etl_snap import SNAP_DATA, SOURCE_URL - - if years is None: - years = list(SNAP_DATA.keys()) - - client = get_supabase_client() - - # Get or create USDA SNAP source - sources = query_sources(jurisdiction="us", institution="usda") - snap_source = next((s for s in sources if s["dataset"] == "snap"), None) - - if not snap_source: - result = client.schema("microplex").table("sources").insert({ - "jurisdiction": "us", - "institution": "usda", - "dataset": "snap", - "name": "USDA SNAP Statistics", - "url": SOURCE_URL, - }).execute() - source_id = result.data[0]["id"] - else: - source_id = snap_source["id"] - - targets_loaded = 0 - - for year in years: - if year not in SNAP_DATA: - continue - - data = SNAP_DATA[year] - - if dry_run: - targets_loaded += 3 # participants, households, benefits - continue - - # National SNAP stratum - stratum_id = get_or_create_stratum( - client, - name=f"US SNAP Recipients {year}", - jurisdiction="us", - constraints=[{"variable": "snap_participation", "operator": "==", "value": "1"}], - ) - - # Participants - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "snap_participants", - "value": data["participants"], - "target_type": "count", - "period": year, - }).execute() - targets_loaded += 1 - - # Households - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "snap_households", - "value": data["households"], - "target_type": "count", - "period": year, - }).execute() - targets_loaded += 1 - - # Benefits - client.schema("microplex").table("targets").insert({ - "source_id": source_id, - "stratum_id": stratum_id, - "variable": "snap_benefits_total", - "value": data["benefits"], - "target_type": "amount", - "period": year, - }).execute() - targets_loaded += 1 - - return { - "targets_loaded": targets_loaded, - "years": years, - "dry_run": dry_run, - } - - -def load_all_targets_supabase(dry_run: bool = False) -> Dict[str, Any]: - """Load all available targets to Supabase.""" - results = {} - - print("Loading SOI targets...") - results["soi"] = load_soi_targets_supabase(dry_run=dry_run) - print(f" Loaded {results['soi']['targets_loaded']} SOI targets") - - print("Loading SNAP targets...") - results["snap"] = load_snap_targets_supabase(dry_run=dry_run) - print(f" Loaded {results['snap']['targets_loaded']} SNAP targets") - - total = sum(r["targets_loaded"] for r in results.values()) - results["total"] = total - print(f"Total targets loaded: {total}") - - return results - - -def main(): - import argparse - - parser = argparse.ArgumentParser(description="Load targets to Supabase") - parser.add_argument("--source", type=str, choices=["soi", "snap", "all"], default="all") - parser.add_argument("--years", type=int, nargs="+", help="Years to load") - parser.add_argument("--dry-run", action="store_true", help="Show stats without inserting") - args = parser.parse_args() - - if args.source == "all": - load_all_targets_supabase(dry_run=args.dry_run) - elif args.source == "soi": - result = load_soi_targets_supabase(years=args.years, dry_run=args.dry_run) - print(f"Loaded {result['targets_loaded']} SOI targets") - elif args.source == "snap": - result = load_snap_targets_supabase(years=args.years, dry_run=args.dry_run) - print(f"Loaded {result['targets_loaded']} SNAP targets") - - -if __name__ == "__main__": - main() diff --git a/db/schema.py b/db/schema.py index fd91e6ea..bb8e0650 100644 --- a/db/schema.py +++ b/db/schema.py @@ -1,13 +1,14 @@ """ -Database schema for targets and strata. +Database schema for Arch's local source-artifact store. -Three-table design (following policyengine-us-data): -- strata: Population subgroups defined by constraints -- stratum_constraints: Rules defining each stratum -- targets: Administrative totals linked to strata +- source_artifacts: source files with provenance and checksum metadata +- source_tables / source_columns / source_rows: their parsed contents + +The legacy strata / stratum_constraints / targets tables were retired with the +loaders that filled them; see the README section "Retired legacy target +tables". """ -import hashlib from datetime import datetime, timezone from enum import Enum from pathlib import Path @@ -15,7 +16,7 @@ from sqlmodel import Field, Relationship, Session, SQLModel, create_engine -# Default storage location (local dev; production uses Supabase) +# Default local storage location for the source-artifact database. DEFAULT_DB_PATH = Path(__file__).parent.parent / "macro" / "targets.db" @@ -54,132 +55,6 @@ class DataSource(str, Enum): OBR = "obr" # Office for Budget Responsibility projections -class TargetType(str, Enum): - """Whether target is a count or monetary amount.""" - - COUNT = "count" - AMOUNT = "amount" - RATE = "rate" - - -class GeographicLevel(str, Enum): - """Geographic aggregation levels for targets.""" - - NATIONAL = "national" - STATE = "state" - CONGRESSIONAL_DISTRICT = "congressional_district" - COUNTY = "county" - ZIP_CODE = "zip_code" - METRO_AREA = "metro_area" - - -class Stratum(SQLModel, table=True): - """ - A population subgroup defined by constraints. - - Examples: - - All US tax filers - - California residents with AGI $50k-$75k - - Single filers under age 65 - """ - - __tablename__ = "strata" - - id: Optional[int] = Field(default=None, primary_key=True) - name: str = Field(index=True) - description: Optional[str] = None - jurisdiction: Jurisdiction - definition_hash: str = Field(unique=True, index=True) - - # Hierarchy: strata can nest (e.g., state within national) - parent_id: Optional[int] = Field(default=None, foreign_key="strata.id") - stratum_group_id: Optional[str] = Field( - default=None, index=True, description="For grouping related strata in calibration" - ) - - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - - # Relationships - constraints: list["StratumConstraint"] = Relationship(back_populates="stratum") - targets: list["Target"] = Relationship(back_populates="stratum") - - @classmethod - def compute_hash( - cls, - constraints: list[tuple[str, str, str]], - jurisdiction: "Jurisdiction | None" = None, - ) -> str: - """Compute unique hash from constraint definitions and jurisdiction.""" - sorted_constraints = sorted(constraints) - # Include jurisdiction to avoid collisions between US/UK strata - hash_input = f"{jurisdiction}:{sorted_constraints}" - return hashlib.sha256(hash_input.encode()).hexdigest()[:16] - - -class StratumConstraint(SQLModel, table=True): - """ - A single constraint defining a stratum. - - Examples: - - variable="state_fips", operator="==", value="06" (California) - - variable="agi", operator=">=", value="50000" - - variable="filing_status", operator="==", value="single" - """ - - __tablename__ = "stratum_constraints" - - id: Optional[int] = Field(default=None, primary_key=True) - stratum_id: int = Field(foreign_key="strata.id", index=True) - - variable: str = Field(index=True) - operator: str # "==", "!=", ">", ">=", "<", "<=", "in" - value: str # Stored as string, parsed based on variable dtype - - # Relationship - stratum: Optional[Stratum] = Relationship(back_populates="constraints") - - -class Target(SQLModel, table=True): - """ - An administrative target value for a stratum. - - Examples: - - Total returns in California: 18.5M - - Total AGI for $50k-$75k bracket: $1.2T - - SNAP households in Texas: 1.4M - """ - - __tablename__ = "targets" - - id: Optional[int] = Field(default=None, primary_key=True) - stratum_id: int = Field(foreign_key="strata.id", index=True) - - variable: str = Field(index=True, description="Arch target input variable ID") - period: int = Field(index=True, description="Year") - value: float - target_type: TargetType = Field(default=TargetType.COUNT) - - # Geographic level for this target - geographic_level: Optional[GeographicLevel] = Field( - default=None, index=True, description="Geographic aggregation level" - ) - - source: DataSource - source_table: Optional[str] = None # e.g., "Table 1.1" - source_url: Optional[str] = None - notes: Optional[str] = None - - # Confidence/quality indicators - is_preliminary: bool = Field(default=False) - margin_of_error: Optional[float] = None - - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - - # Relationship - stratum: Optional[Stratum] = Relationship(back_populates="targets") - - class SourceArtifact(SQLModel, table=True): """A source file ingested into Arch with provenance and checksum metadata.""" @@ -246,7 +121,7 @@ class SourceRow(SQLModel, table=True): def get_engine(db_path: Path = DEFAULT_DB_PATH): - """Get SQLAlchemy engine for the targets database.""" + """Get SQLAlchemy engine for the source-artifact database.""" db_path.parent.mkdir(parents=True, exist_ok=True) return create_engine(f"sqlite:///{db_path}", echo=False) diff --git a/db/supabase_client.py b/db/supabase_client.py index dd2f520d..4a9d83cf 100644 --- a/db/supabase_client.py +++ b/db/supabase_client.py @@ -4,7 +4,9 @@ Provides connection to PolicyEngine Supabase database for: - Source metadata and dataset registries - Raw microdata tables (e.g., microdata.us_census_cps_asec_2024_person) -- Target inputs + +The legacy Supabase target-input helpers (query_targets, query_strata, +insert_targets_batch) were retired with the local target tables. Table naming pattern: {jurisdiction}_{institution}_{dataset}_{year}_{table_type} Example: us_census_cps_asec_2024_person @@ -35,9 +37,6 @@ def _env(*names: str) -> str | None: MICRODATA_SCHEMA = ( _env("POLICYENGINE_MICRODATA_SCHEMA") or "microdata" ) -TARGETS_SCHEMA = ( - _env("POLICYENGINE_TARGETS_SCHEMA") or "targets" -) @dataclass @@ -364,76 +363,6 @@ def query_cps( return pd.DataFrame(result.data) -# ============================================================================= -# Targets and strata -# ============================================================================= - - -def query_strata( - jurisdiction: Optional[str] = None, -) -> List[Dict[str, Any]]: - """ - Query target input strata with their constraints. - - Args: - jurisdiction: Filter by jurisdiction - - Returns: - List of strata records with nested constraints - """ - client = get_supabase_client() - query = _table(client, TARGETS_SCHEMA, "strata").select("*, stratum_constraints(*)") - - if jurisdiction: - query = query.eq("jurisdiction", jurisdiction) - - result = query.execute() - return result.data - - -def query_targets( - jurisdiction: Optional[str] = None, - year: Optional[int] = None, - source_id: Optional[str] = None, - variable: Optional[str] = None, -) -> List[Dict[str, Any]]: - """ - Query Arch target inputs. - - Args: - jurisdiction: Filter by jurisdiction (via stratum join) - year: Filter by period/year - source_id: Filter by source UUID - variable: Filter by variable name - - Returns: - List of target records with stratum info - """ - client = get_supabase_client() - # Nested join: strata with their stratum_constraints - query = _table(client, TARGETS_SCHEMA, "targets").select( - "*, strata(*, stratum_constraints(*)), sources(*)" - ) - - if year: - query = query.eq("period", year) - if source_id: - query = query.eq("source_id", source_id) - if variable: - query = query.eq("variable", variable) - - result = query.execute() - - # Filter by jurisdiction if specified (post-query since it's on joined table) - data = result.data - if jurisdiction: - data = [ - t for t in data if t.get("strata", {}).get("jurisdiction") == jurisdiction - ] - - return data - - # ============================================================================= # Insert operations # ============================================================================= @@ -473,28 +402,3 @@ def insert_microdata_batch( total += len(chunk) return total - - -def insert_targets_batch( - targets: List[Dict[str, Any]], - chunk_size: int = 100, -) -> int: - """ - Insert target records in batches. - - Args: - targets: List of target dicts - chunk_size: Records per batch - - Returns: - Number of records inserted - """ - client = get_supabase_client() - total = 0 - - for i in range(0, len(targets), chunk_size): - chunk = targets[i : i + chunk_size] - _table(client, TARGETS_SCHEMA, "targets").insert(chunk).execute() - total += len(chunk) - - return total diff --git a/db/tests/test_etl_cps_raw.py b/db/tests/test_etl_cps_raw.py deleted file mode 100644 index ded74013..00000000 --- a/db/tests/test_etl_cps_raw.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Tests for CPS raw microdata ETL to Supabase.""" - -import pandas as pd -from unittest.mock import MagicMock, patch - - -class TestCPSRawETL: - """Test CPS raw data loading to Supabase.""" - - def test_prepare_person_records(self): - """Test preparing person records for Supabase insert.""" - from db.etl_cps_raw import prepare_person_records - - # Sample raw CPS data - df = pd.DataFrame({ - "PH_SEQ": [1, 1, 2], - "PPPOS": [1, 2, 1], - "A_AGE": [35, 10, 42], - "A_SEX": [1, 2, 2], - "PRDTRACE": [1, 1, 2], - "PEHSPNON": [2, 2, 1], - "GESTFIPS": [6, 6, 36], - "PTOTVAL": [50000.0, 0.0, 75000.0], - "PEARNVAL": [48000.0, 0.0, 70000.0], - "WSAL_VAL": [45000.0, 0.0, 65000.0], - "MARSUPWT": [1500.0, 1500.0, 2000.0], - }) - - records = prepare_person_records(df) - - assert len(records) == 3 - assert records[0]["ph_seq"] == 1 - assert records[0]["a_age"] == 35 - assert records[0]["gestfips"] == 6 - assert records[0]["marsupwt"] == 1500.0 - # Should have raw_data with all columns - assert "raw_data" in records[0] - assert records[0]["raw_data"]["A_AGE"] == 35 - - def test_prepare_household_records(self): - """Test preparing household records for Supabase insert.""" - from db.etl_cps_raw import prepare_household_records - - df = pd.DataFrame({ - "H_SEQ": [1, 2], - "H_NUMPER": [3, 2], - "HH5TO18": [1, 0], - "HUNDER18": [1, 0], - "GESTFIPS": [6, 36], - "HTOTVAL": [125000.0, 80000.0], - "HSUP_WGT": [1500.0, 2000.0], - }) - - records = prepare_household_records(df) - - assert len(records) == 2 - assert records[0]["h_seq"] == 1 - assert records[0]["h_numper"] == 3 - assert records[1]["gestfips"] == 36 - - def test_get_cps_table_names(self): - """Test table name generation for CPS ASEC.""" - from db.etl_cps_raw import get_cps_table_names - - names = get_cps_table_names(2024) - - assert names["person"] == "us_census_cps_asec_2024_person" - assert names["household"] == "us_census_cps_asec_2024_household" - assert names["family"] == "us_census_cps_asec_2024_family" - - @patch("db.etl_cps_raw.get_supabase_client") - @patch("db.etl_cps_raw.pd.read_parquet") - @patch("db.etl_cps_raw.get_raw_cache_dir") - def test_load_cps_to_supabase_dry_run(self, mock_cache_dir, mock_read, mock_client): - """Test dry run mode returns stats without inserting.""" - from db.etl_cps_raw import load_cps_to_supabase - - # Mock cache directory with existing files - mock_path = MagicMock() - mock_path.exists.return_value = True - mock_cache_dir.return_value = mock_path - mock_cache_dir.return_value.__truediv__ = lambda self, x: mock_path - - # Mock parquet reading - mock_read.return_value = pd.DataFrame({ - "PH_SEQ": [1, 2], - "A_AGE": [35, 42], - "MARSUPWT": [1500, 2000], - }) - - result = load_cps_to_supabase(2024, dry_run=True) - - assert "person_count" in result - assert "household_count" in result - assert "family_count" in result - assert result["dry_run"] is True - # In dry run, client should not be called - mock_client.assert_not_called() diff --git a/docs/architecture.md b/docs/architecture.md index 2cb1b6bc..a17ed772 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -44,7 +44,6 @@ sources/ |--------|---------|----------------| | `arch` | Source metadata and lineage | sources, files, content, fetch_log | | `indices` | Source time series | series, values (CPI, wage growth) | -| `targets` | Target inputs | strata, constraints, targets | | `microdata` | Intermediate microdata | cps_asec, acs_pums, synthetic | | `microplex` | Final calibrated data | households, persons, tax_units | @@ -55,13 +54,13 @@ New code should use the `arch` namespace: ```python from arch.sources import SourceFile, SourceReference, query_sources from arch.facts import SourceFact -from arch.targets import Target, query_targets from arch.microdata import query_cps_asec from arch.normalization import convert_units ``` -The `db` package contains the current SQLModel persistence and loader -implementation behind the public `arch` namespace. +The `db` package holds the local source-artifact store and the Supabase client +behind the public `arch` namespace. The legacy `targets` / `strata` / +`stratum_constraints` tables and their loaders were retired (see the README). Jurisdiction source packages should use short import namespaces and published distribution names with a PolicyEngine prefix: @@ -146,8 +145,8 @@ The `arch.normalization` package owns low-assumption representation helpers: ```python from arch.facts import SourceFact -from arch.targets import DataSource, Jurisdiction, TargetType -from arch.normalization import as_target, convert_units +from arch.normalization import convert_units +from db.schema import DataSource, Jurisdiction snap_households = SourceFact( name="snap_households", @@ -158,11 +157,7 @@ snap_households = SourceFact( jurisdiction=Jurisdiction.US, ) -target_input = as_target( - convert_units(snap_households, 1000, "count"), - target_type=TargetType.COUNT, - stratum_name="US SNAP Households", -) +normalized = convert_units(snap_households, 1000, "count") ``` Projection facts from official sources such as CBO, OBR, and ONS can be loaded @@ -187,19 +182,10 @@ fact and its provenance are identified. ## Calibration Pipeline -### 1. Target Inputs (from `targets.*` schema) - -Target inputs define source-backed aggregates that Microplex may use: +### 1. Source facts (from the consumer bundle) -```sql --- targets.strata: Population subgroups -INSERT INTO targets.strata (name, jurisdiction, constraints) -VALUES ('CA adults 18-64', 'us', '[{"variable": "age", "operator": ">=", "value": "18"}, ...]'); - --- targets.targets: Source-backed aggregate values -INSERT INTO targets.targets (stratum_id, variable, value, period) -VALUES (1, 'eitc_recipients', 2500000, 2023); -``` +`arch build-bundle` writes the source-backed facts to `consumer_facts.jsonl`. +Arch has no separate target table; target selection starts from those facts. ### 2. Variable Mapping @@ -214,7 +200,7 @@ sets: ```python target_set = microplex.targets.compose( - inputs=arch_targets, + inputs=arch_facts, target_year=2024, reconciliation="scale_states_to_national", aging="apply_published_growth_factor", @@ -321,17 +307,6 @@ indexing_rule eitc_inflation: rounding: 10 ``` -### targets.* (Target Inputs) - -```sql -targets.strata -- population subgroups with constraints -targets.constraints -- variable, operator, value per stratum -targets.targets -- stratum_id, variable, value, period, source -``` - -Targets in Arch are source-backed inputs. Active, aged, reconciled calibration -target sets belong to Microplex Targets. - ### microdata.* (Intermediate) ```sql diff --git a/docs/calibration-pipeline.md b/docs/calibration-pipeline.md deleted file mode 100644 index 52391e7f..00000000 --- a/docs/calibration-pipeline.md +++ /dev/null @@ -1,246 +0,0 @@ -# Calibration Pipeline Design - -## Overview - -The calibration pipeline connects administrative targets to microdata, producing reweighted samples that match official statistics. - -``` -┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ -│ Microdata │ │ Targets DB │ │ Calibrated │ -│ (CPS/FRS) │ ──▶ │ (SQLite) │ ──▶ │ Weights │ -│ ~200k records │ │ ~1000 targets │ │ weights.parquet│ -└─────────────────┘ └─────────────────┘ └─────────────────┘ -``` - -## Pipeline Steps - -### 1. Load Microdata - -```python -from calibration import load_microdata - -# Load raw CPS with original weights -microdata = load_microdata( - source="cps", # or "frs" for UK - year=2023, - variables=["age", "employment_income", "is_tax_filer", ...] -) -# Returns: DataFrame with ~200k rows, original weights in 'weight' column -``` - -### 2. Query Targets - -```python -from calibration import get_targets - -targets = get_targets( - jurisdiction="us", # or "uk" - year=2023, - sources=["irs-soi", "census-acs", "bls"], # which admin sources - variables=["tax_returns", "total_agi", "employment"], # optional filter -) -# Returns: List[TargetSpec] with strata constraints and values -``` - -### 3. Build Constraint Matrix - -Map targets to microdata aggregations: - -```python -from calibration import build_constraint_matrix - -# Each target becomes a constraint: sum(weights * indicator) = target_value -constraints = build_constraint_matrix( - microdata=microdata, - targets=targets, - tolerance=0.01, # Allow 1% deviation -) -``` - -**Constraint types:** - -| Target Variable | Microdata Aggregation | -|----------------|----------------------| -| `tax_returns` | `sum(w * is_tax_filer)` | -| `total_agi` | `sum(w * agi * is_tax_filer)` | -| `snap_households` | `sum(w * receives_snap)` | -| `population_65_plus` | `sum(w * (age >= 65))` | - -### 4. Calibrate Weights - -```python -from calibration import calibrate - -new_weights = calibrate( - microdata=microdata, - constraints=constraints, - method="entropy", # or "raking", "linear" - bounds=(0.1, 10.0), # Weight adjustment bounds - max_iterations=100, -) -``` - -**Methods:** - -| Method | Description | When to Use | -|--------|-------------|-------------| -| `entropy` | Minimize KL divergence from original weights | Default, smooth adjustments | -| `raking` | Iterative proportional fitting | Many margin constraints | -| `linear` | Linear regression adjustment | Few constraints, fast | - -### 5. Validate & Export - -```python -from calibration import validate, export_weights - -# Check calibration quality -report = validate( - microdata=microdata, - new_weights=new_weights, - targets=targets, -) -# Returns: ValidationReport with residuals, diagnostics - -# Export for use in PolicyEngine -export_weights( - weights=new_weights, - output_path="calibrated_weights_2023.parquet", - metadata={ - "source": "cps", - "year": 2023, - "targets_hash": targets.hash(), - "method": "entropy", - } -) -``` - -## API Design - -```python -# calibration/__init__.py - -from .loader import load_microdata -from .targets import get_targets, TargetSpec -from .constraints import build_constraint_matrix, Constraint -from .methods import calibrate, EntropyCalibrator, RakingCalibrator -from .validation import validate, ValidationReport -from .export import export_weights - - -# High-level convenience function -def calibrate_microdata( - source: str, - year: int, - jurisdiction: str = "us", - method: str = "entropy", - output_path: str | None = None, -) -> tuple[pd.DataFrame, ValidationReport]: - """ - Full calibration pipeline in one call. - - Returns calibrated microdata and validation report. - """ - microdata = load_microdata(source, year) - targets = get_targets(jurisdiction, year) - constraints = build_constraint_matrix(microdata, targets) - new_weights = calibrate(microdata, constraints, method=method) - report = validate(microdata, new_weights, targets) - - if output_path: - export_weights(new_weights, output_path) - - microdata["weight"] = new_weights - return microdata, report -``` - -## Directory Structure - -``` -arch/ -├── calibration/ -│ ├── __init__.py -│ ├── loader.py # load_microdata() -│ ├── targets.py # get_targets(), TargetSpec -│ ├── constraints.py # build_constraint_matrix() -│ ├── methods/ -│ │ ├── __init__.py -│ │ ├── entropy.py # EntropyCalibrator -│ │ ├── raking.py # RakingCalibrator -│ │ └── linear.py # LinearCalibrator -│ ├── validation.py # validate() -│ └── export.py # export_weights() -├── tests/ -│ └── test_calibration.py -└── docs/ - └── calibration-pipeline.md # This file -``` - -## Variable Mapping - -The key challenge is mapping target variables to microdata variables: - -```python -# calibration/variable_map.py - -VARIABLE_MAP = { - # Target variable -> (microdata_variable, aggregation_type) - "tax_returns": ("is_tax_filer", "count"), - "total_agi": ("adjusted_gross_income", "sum"), - "snap_households": ("receives_snap", "count"), - "snap_benefits": ("snap_amount", "sum"), - "employment": ("is_employed", "count"), - "unemployment": ("is_unemployed", "count"), - "population": (None, "count"), # Just sum weights - "population_65_plus": ("age >= 65", "count"), - - # UK variables - "income_tax_payers": ("pays_income_tax", "count"), - "total_income_tax": ("income_tax", "sum"), - "uc_households": ("receives_uc", "count"), -} -``` - -## Stratum Handling - -Targets come with stratum constraints that must be applied: - -```python -def apply_stratum_constraints( - microdata: pd.DataFrame, - stratum: Stratum, -) -> pd.Series: - """Return boolean mask for records matching stratum.""" - mask = pd.Series(True, index=microdata.index) - - for constraint in stratum.constraints: - if constraint.operator == "==": - mask &= microdata[constraint.variable] == constraint.value - elif constraint.operator == ">=": - mask &= microdata[constraint.variable] >= float(constraint.value) - # ... etc - - return mask -``` - -## Performance Considerations - -| Dataset | Records | Targets | Calibration Time | -|---------|---------|---------|------------------| -| CPS | 200k | 100 | ~10s | -| CPS | 200k | 500 | ~30s | -| Full Census | 3M | 1000 | ~5 min | - -Optimizations: -- Vectorized constraint evaluation (NumPy/Pandas) -- Sparse constraint matrices (scipy.sparse) -- Parallel target aggregation -- Caching stratum masks - -## Next Steps - -1. [ ] Implement `calibration/loader.py` (CPS/FRS loading) -2. [ ] Implement `calibration/targets.py` (query targets DB) -3. [ ] Implement `calibration/methods/entropy.py` (core algorithm) -4. [ ] Add tests with mock data -5. [ ] Benchmark on real CPS -6. [ ] Integrate with Arch microdata loaders diff --git a/docs/cps-variable-coverage.md b/docs/cps-variable-coverage.md index 3f1341bb..67c4e3c6 100644 --- a/docs/cps-variable-coverage.md +++ b/docs/cps-variable-coverage.md @@ -89,41 +89,6 @@ Capital gains represent the largest data quality issue for tax microsimulation: **Validation Target**: IRS SOI shows ~$120B CTC/ACTC claimed (2021) -## Calibration Targets - -The ETL pipeline (`db/etl_soi.py`) loads IRS SOI targets for calibration: - -### National Targets (All Filers) - -| Target | SOI 2021 Value | Notes | -|--------|----------------|-------| -| Total returns | 153.9M | Calibration constraint | -| Total AGI | $14.7T | Weighted sum validation | -| EITC claims | ~27M | Eligibility validation | -| EITC amount | ~$64B | Amount validation | - -### Income by Bracket (from `etl_soi_credits.py`) - -| AGI Bracket | Returns | AGI Amount | -|-------------|---------|------------| -| $0-$25k | 32.6M | $484B | -| $25k-$50k | 31.3M | $1.2T | -| $50k-$100k | 34.6M | $2.5T | -| $100k-$200k | 22.8M | $3.2T | -| $200k-$500k | 7.2M | $2.2T | -| $500k+ | 1.8M | $5.2T | - -### Income by Source (from `etl_soi_income_sources.py`) - -| Income Type | Returns | Amount | -|-------------|---------|--------| -| Wages | 132M | $8.9T | -| Interest | 55M | $277B | -| Dividends | 32M | $432B | -| Capital gains | 21M | $1.1T | -| Schedule C | 29M | $399B | -| SS taxable | 28M | $317B | - ## Recommended Workflow ### 1. Download and Cache CPS @@ -133,21 +98,10 @@ python micro/us/census/download_cps.py --year 2024 ### 2. Convert to PolicyEngine Format ```bash -python scripts/cps_to_policyengine.py --year 2024 --calibrate --summary +python scripts/cps_to_policyengine.py --year 2024 --summary ``` -### 3. Validate Against Targets -```python -from arch.targets import get_engine, Target -from sqlmodel import Session, select - -with Session(get_engine()) as session: - targets = session.exec( - select(Target).where(Target.variable.contains("eitc")) - ).all() -``` - -### 4. Run Microsimulation +### 3. Run Microsimulation ```bash # In policyengine-us repository python -m policyengine_us.microsim --input policyengine_input_2024.parquet @@ -230,24 +184,12 @@ micro/us/census/cps-asec/ └── ... (60+ more) ``` -### Calibration Data - -``` -db/ -├── etl_soi.py # Core SOI data loader -├── etl_soi_credits.py # Credit-specific targets -├── etl_soi_income_sources.py # Income by source targets -└── schema.py # Database schema -``` - ### Conversion Scripts ``` scripts/ -├── cps_to_policyengine.py # Main conversion script -└── export_to_json.py # JSON export utility +└── cps_to_policyengine.py # Main conversion script micro/us/ -├── census/download_cps.py # CPS downloader -└── calibrate.py # Entropy calibration +└── census/download_cps.py # CPS downloader ``` diff --git a/micro/us/calibrate.py b/micro/us/calibrate.py deleted file mode 100644 index 60828ed8..00000000 --- a/micro/us/calibrate.py +++ /dev/null @@ -1,389 +0,0 @@ -""" -Calibrate CPS tax unit weights to IRS SOI targets. - -Uses entropy calibration (gradient descent on KL divergence) following -the architecture in docs/calibration-pipeline.md. -""" - -import numpy as np -import pandas as pd -from dataclasses import dataclass -from scipy.optimize import minimize -import sys -sys.path.insert(0, str(__file__).rsplit('/', 3)[0]) # Add parent to path - -from calibration.constraints import Constraint -from arch.targets import TargetType - - -# IRS SOI 2021 Targets (from Statistics of Income, Table 1.1) -# Source: https://www.irs.gov/statistics/soi-tax-stats-individual-income-tax-returns-complete-report-publication-1304 -# Note: Using consistent bracket data; total is computed as sum of brackets -IRS_SOI_2021_RETURNS_BY_AGI = { - 'no_agi': 13_992_100, # Returns with no AGI (zero or not computed) - 'under_1': 1_686_440, # $1 under $5,000 (negative AGI) - '1_to_5k': 5_183_390, # $1 - $5,000 - '5k_to_10k': 7_929_860, # $5,000 - $10,000 - '10k_to_15k': 9_883_050, # $10,000 - $15,000 - '15k_to_20k': 9_113_990, # $15,000 - $20,000 - '20k_to_25k': 8_186_640, # $20,000 - $25,000 - '25k_to_30k': 7_407_890, # $25,000 - $30,000 - '30k_to_40k': 13_194_450, # $30,000 - $40,000 - '40k_to_50k': 10_930_780, # $40,000 - $50,000 - '50k_to_75k': 19_494_660, # $50,000 - $75,000 - '75k_to_100k': 15_137_070, # $75,000 - $100,000 - '100k_to_200k': 22_849_380, # $100,000 - $200,000 - '200k_to_500k': 7_167_290, # $200,000 - $500,000 - '500k_to_1m': 1_106_040, # $500,000 - $1,000,000 - '1m_plus': 664_340, # $1,000,000 or more -} - -# Total returns = sum of all brackets -IRS_SOI_2021_TOTAL_RETURNS = sum(IRS_SOI_2021_RETURNS_BY_AGI.values()) # ~153.9M - -# AGI totals by bracket (in dollars) -IRS_SOI_2021_AGI_BY_BRACKET = { - 'no_agi': 0, - 'under_1': -94_000_000_000, # Negative AGI returns - '1_to_5k': 15_000_000_000, - '5k_to_10k': 59_000_000_000, - '10k_to_15k': 123_000_000_000, - '15k_to_20k': 160_000_000_000, - '20k_to_25k': 184_000_000_000, - '25k_to_30k': 204_000_000_000, - '30k_to_40k': 461_000_000_000, - '40k_to_50k': 492_000_000_000, - '50k_to_75k': 1_210_000_000_000, - '75k_to_100k': 1_316_000_000_000, - '100k_to_200k': 3_187_000_000_000, - '200k_to_500k': 2_161_000_000_000, - '500k_to_1m': 762_000_000_000, - '1m_plus': 4_466_000_000_000, -} - -IRS_SOI_2021_TOTAL_AGI = sum(IRS_SOI_2021_AGI_BY_BRACKET.values()) # ~$14.7T - -AGI_BRACKETS = [ - ('no_agi', None, None), # Special: AGI == 0 or not computed - ('under_1', -np.inf, 1), # Negative AGI - ('1_to_5k', 1, 5000), - ('5k_to_10k', 5000, 10000), - ('10k_to_15k', 10000, 15000), - ('15k_to_20k', 15000, 20000), - ('20k_to_25k', 20000, 25000), - ('25k_to_30k', 25000, 30000), - ('30k_to_40k', 30000, 40000), - ('40k_to_50k', 40000, 50000), - ('50k_to_75k', 50000, 75000), - ('75k_to_100k', 75000, 100000), - ('100k_to_200k', 100000, 200000), - ('200k_to_500k', 200000, 500000), - ('500k_to_1m', 500000, 1000000), - ('1m_plus', 1000000, np.inf), -] - - -@dataclass -class CalibrationResult: - """Results from calibration.""" - original_weights: np.ndarray - calibrated_weights: np.ndarray - adjustment_factors: np.ndarray - targets_before: dict - targets_after: dict - success: bool - message: str - kl_divergence: float - - -def assign_agi_bracket(agi: np.ndarray) -> np.ndarray: - """Assign each record to an AGI bracket.""" - brackets = np.empty(len(agi), dtype=object) - for name, low, high in AGI_BRACKETS: - if name == 'no_agi': - # Special case: AGI is exactly 0 or NaN - mask = (agi == 0) | np.isnan(agi) - else: - mask = (agi >= low) & (agi < high) - brackets[mask] = name - return brackets - - -def build_constraints(df: pd.DataFrame, min_obs: int = 100) -> list[Constraint]: - """ - Build calibration constraints from IRS SOI targets. - - Returns list of Constraint objects per architecture spec. - Only uses bracket constraints (total is redundant since sum of brackets = total). - """ - constraints = [] - - df = df.copy() - df['agi_bracket'] = assign_agi_bracket(df['adjusted_gross_income'].values) - - # Returns by AGI bracket (skip small strata) - # Note: We don't add a total_returns constraint because it's redundant - - # the sum of all bracket constraints implicitly constrains the total. - for bracket_name, _, _ in AGI_BRACKETS: - if bracket_name not in IRS_SOI_2021_RETURNS_BY_AGI: - continue - - indicator = (df['agi_bracket'] == bracket_name).astype(float).values - n_obs = indicator.sum() - - if n_obs >= min_obs: - constraints.append(Constraint( - indicator=indicator, - target_value=IRS_SOI_2021_RETURNS_BY_AGI[bracket_name], - variable=f'returns_{bracket_name}', - target_type=TargetType.COUNT, - tolerance=0.05, - stratum_name=f'Filers AGI {bracket_name}', - )) - - return constraints - - -def entropy_calibrate( - original_weights: np.ndarray, - constraints: list[Constraint], - bounds: tuple[float, float] = (0.2, 5.0), - max_iter: int = 200, - tol: float = 1e-8, - verbose: bool = True, -) -> tuple[np.ndarray, bool, float]: - """ - Calibrate weights using entropy minimization (gradient descent). - - Uses the dual formulation: instead of optimizing n weights directly, - we optimize m Lagrange multipliers (one per constraint) which is - much more efficient. - - The optimal weights are: w_i = w0_i * exp(sum_j lambda_j * A_ij) - where A_ij is the constraint matrix. - """ - n = len(original_weights) - m = len(constraints) - - if verbose: - print(f"Entropy calibration: {n:,} weights, {m} constraints") - - # Build constraint matrix A (m x n) - # Each row is an indicator/value vector for one constraint - A = np.zeros((m, n)) - targets = np.zeros(m) - - for j, c in enumerate(constraints): - A[j, :] = c.indicator - targets[j] = c.target_value - - # Dual objective: find lambdas that minimize the dual - # Dual = sum_i w0_i * exp(sum_j lambda_j * A_ji) - sum_j lambda_j * target_j - def dual_objective(lambdas: np.ndarray) -> float: - # Compute log adjustment: sum_j lambda_j * A_ji for each i - log_adj = A.T @ lambdas # (n,) - - # Clip for numerical stability - log_adj = np.clip(log_adj, -10, 10) - - # Calibrated weights - w = original_weights * np.exp(log_adj) - - # Dual value - return w.sum() - lambdas @ targets - - def dual_gradient(lambdas: np.ndarray) -> np.ndarray: - log_adj = A.T @ lambdas - log_adj = np.clip(log_adj, -10, 10) - w = original_weights * np.exp(log_adj) - - # Gradient: sum_i w_i * A_ji - target_j = achieved_j - target_j - achieved = A @ w - return achieved - targets - - # Initial lambdas: zeros (no adjustment) - lambda0 = np.zeros(m) - - # Optimize using L-BFGS-B (gradient descent with bounds) - result = minimize( - dual_objective, - lambda0, - method='L-BFGS-B', - jac=dual_gradient, - options={ - 'maxiter': max_iter, - 'ftol': tol, - 'gtol': 1e-6, - } - ) - - if verbose: - print(f"Optimization: {result.message}") - print(f"Iterations: {result.nit}, Function evals: {result.nfev}") - - # Compute final weights - log_adj = A.T @ result.x - log_adj = np.clip(log_adj, np.log(bounds[0]), np.log(bounds[1])) - calibrated_weights = original_weights * np.exp(log_adj) - - # Compute KL divergence - w_safe = np.maximum(calibrated_weights, 1e-10) - w0_safe = np.maximum(original_weights, 1e-10) - kl_div = np.sum(w_safe * np.log(w_safe / w0_safe)) - - return calibrated_weights, result.success, kl_div - - -def calibrate_weights( - df: pd.DataFrame, - bounds: tuple[float, float] = (0.2, 5.0), - tolerance: float = 0.05, - min_obs: int = 100, - verbose: bool = True, -) -> CalibrationResult: - """ - Calibrate weights using entropy minimization. - - Args: - df: Tax unit DataFrame with 'weight' and 'adjusted_gross_income' - bounds: (min_ratio, max_ratio) for weight adjustments - tolerance: Allowed deviation from targets - min_obs: Minimum observations for a constraint - verbose: Print progress - """ - original_weights = df['weight'].values.copy() - n = len(original_weights) - - if verbose: - print(f"Calibrating {n:,} tax units...") - print(f"Original weighted total: {original_weights.sum():,.0f}") - - # Build constraints - constraints = build_constraints(df, min_obs=min_obs) - - if verbose: - print(f"Built {len(constraints)} constraints (min {min_obs} obs each)") - - # Compute pre-calibration values - targets_before = {} - for c in constraints: - current = np.dot(original_weights, c.indicator) - targets_before[c.variable] = { - 'current': current, - 'target': c.target_value, - 'error': (current - c.target_value) / c.target_value if c.target_value != 0 else 0, - } - - if verbose: - print("\nPre-calibration errors (sample brackets):") - sample_brackets = ['returns_50k_to_75k', 'returns_100k_to_200k', 'returns_200k_to_500k'] - for name in sample_brackets: - if name in targets_before: - t = targets_before[name] - print(f" {name}: {t['error']:+.1%}") - - # Run entropy calibration - calibrated_weights, success, kl_div = entropy_calibrate( - original_weights, - constraints, - bounds=bounds, - verbose=verbose, - ) - - # Compute post-calibration values - targets_after = {} - max_error = 0 - for c in constraints: - current = np.dot(calibrated_weights, c.indicator) - error = (current - c.target_value) / c.target_value if c.target_value != 0 else 0 - targets_after[c.variable] = { - 'current': current, - 'target': c.target_value, - 'error': error, - } - max_error = max(max_error, abs(error)) - - adjustment_factors = calibrated_weights / original_weights - - sample_brackets = ['returns_50k_to_75k', 'returns_100k_to_200k', 'returns_200k_to_500k'] - if verbose: - print(f"\nPost-calibration (max error: {max_error:.1%}):") - for name in sample_brackets: - if name in targets_after: - t = targets_after[name] - print(f" {name}: {t['error']:+.1%}") - - # Calculate coverage - calibrated_total = calibrated_weights.sum() - irs_total = IRS_SOI_2021_TOTAL_RETURNS - coverage = calibrated_total / irs_total - print(f"\nCoverage: {calibrated_total:,.0f} / {irs_total:,.0f} = {coverage:.1%}") - print(f" (CPS underrepresents low/no-income filers by ~{(1-coverage)*100:.0f}%)") - - print(f"\nWeight adjustments: mean={adjustment_factors.mean():.2f}, " - f"std={adjustment_factors.std():.2f}, " - f"range=[{adjustment_factors.min():.2f}, {adjustment_factors.max():.2f}]") - print(f"KL divergence: {kl_div:.2f}") - - return CalibrationResult( - original_weights=original_weights, - calibrated_weights=calibrated_weights, - adjustment_factors=adjustment_factors, - targets_before=targets_before, - targets_after=targets_after, - success=success and max_error < tolerance, - message="Converged" if success else "Did not converge", - kl_divergence=kl_div, - ) - - -def calibrate_and_run(year: int = 2024, filer_threshold: float = 0) -> pd.DataFrame: - """Load data, calibrate weights, return calibrated DataFrame.""" - from tax_unit_builder import load_and_build_tax_units - - print("=" * 60) - print("POLICYENGINE MICRODATA CALIBRATION (Entropy Method)") - print("=" * 60) - - print("\n1. Loading tax unit data...") - df = load_and_build_tax_units(year) - print(f" Loaded {len(df):,} tax units") - - # Filter to likely filers (have income or already pay tax) - # IRS filing threshold 2024: ~$13,850 for single, $27,700 for joint - filer_mask = ( - (df['total_income'] > 13850) | # Above single threshold - (df['wage_income'] > 0) | # Has wage income - (df['self_employment_income'] > 0) # Has SE income - ) - df = df[filer_mask].copy() - print(f" Filtered to {len(df):,} likely filers") - - print("\n2. Calibrating to IRS SOI 2021 targets...") - result = calibrate_weights(df) - - df['original_weight'] = result.original_weights - df['weight'] = result.calibrated_weights - df['weight_adjustment'] = result.adjustment_factors - - print("\n3. Summary:") - print(f" Original total: {result.original_weights.sum():,.0f}") - print(f" Calibrated total: {result.calibrated_weights.sum():,.0f}") - print(f" IRS SOI target: {IRS_SOI_2021_TOTAL_RETURNS:,.0f}") - - return df - - -if __name__ == "__main__": - df = calibrate_and_run() - - print("\n" + "=" * 60) - print("CALIBRATED DATA SUMMARY") - print("=" * 60) - print(f"\nTotal tax units: {len(df):,}") - print(f"Weighted population: {df['weight'].sum():,.0f}") - print(f"Total AGI: ${(df['adjusted_gross_income'] * df['weight']).sum():,.0f}") - - output_path = "tax_units_calibrated_2024.parquet" - df.to_parquet(output_path) - print(f"\nSaved to {output_path}") diff --git a/micro/us/district.py b/micro/us/district.py deleted file mode 100644 index bfa769be..00000000 --- a/micro/us/district.py +++ /dev/null @@ -1,567 +0,0 @@ -""" -US District Microplex Builder. - -Generates synthetic tax units at geographic granularity (states, counties, -congressional districts) and calibrates weights to match local targets. - -Uses microplex library for synthesis and calibration algorithms. -Uses Arch targets for authoritative calibration data. - -Example: - >>> from micro.us.district import DistrictMicroplex - >>> from calibration.targets import get_targets - >>> - >>> dm = DistrictMicroplex(n_per_district=1000, target_sparsity=0.9) - >>> dm.fit(seed_data, epochs=100) - >>> result = dm.build(seed_data, districts=["06", "36"], targets=targets) -""" - -from __future__ import annotations - -import time -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np -import pandas as pd - -# Import algorithms from microplex. SparseCalibrator is not available in every -# released microplex build, so keep module import cheap and fail only when -# calibration actually needs it. -from microplex import ConditionalMAF - -try: - from microplex import SparseCalibrator as MicroplexSparseCalibrator -except ImportError: - MicroplexSparseCalibrator = None - -# Import targets from Arch-compatible calibration adapters -from calibration.targets import TargetSpec, get_targets -from db.schema import TargetType - - -# US State FIPS codes (int -> abbreviation) -STATE_FIPS: Dict[int, str] = { - 1: "AL", 2: "AK", 4: "AZ", 5: "AR", 6: "CA", 8: "CO", 9: "CT", 10: "DE", - 11: "DC", 12: "FL", 13: "GA", 15: "HI", 16: "ID", 17: "IL", 18: "IN", - 19: "IA", 20: "KS", 21: "KY", 22: "LA", 23: "ME", 24: "MD", 25: "MA", - 26: "MI", 27: "MN", 28: "MS", 29: "MO", 30: "MT", 31: "NE", 32: "NV", - 33: "NH", 34: "NJ", 35: "NM", 36: "NY", 37: "NC", 38: "ND", 39: "OH", - 40: "OK", 41: "OR", 42: "PA", 44: "RI", 45: "SC", 46: "SD", 47: "TN", - 48: "TX", 49: "UT", 50: "VT", 51: "VA", 53: "WA", 54: "WV", 55: "WI", - 56: "WY", -} - -# Reverse mapping -FIPS_TO_STATE = {v: k for k, v in STATE_FIPS.items()} - -# Required columns for seed data -REQUIRED_COLUMNS = ["state_fips", "wage_income", "head_age", "weight"] - - -def build_targets_from_db( - year: int = 2021, - sources: Optional[List[str]] = None, - verbose: bool = False, -) -> Tuple[Dict[str, Dict], Dict[str, float]]: - """ - Build calibration targets from the Arch target database. - - Args: - year: Target year (default 2021) - sources: List of sources to include (default: all) - verbose: Print progress information - - Returns: - marginal_targets: Dict for categorical targets {var: {category: count}} - continuous_targets: Dict for sum targets {var: total} - """ - targets = get_targets(jurisdiction="us", year=year, sources=sources) - - if verbose: - print(f" Loaded {len(targets)} targets from database") - - marginal_targets: Dict[str, Dict] = {"state_fips": {}} - continuous_targets: Dict[str, float] = {} - - for t in targets: - # Extract state from constraints - state_constraint = None - for var, op, val in t.constraints: - if var == "state_fips" and op == "==": - state_constraint = int(val) - break - - if state_constraint is None: - # National target - add to continuous - if t.variable in ["adjusted_gross_income", "wage_income", "total_income"]: - continuous_targets[t.variable] = continuous_targets.get(t.variable, 0) + t.value - continue - - # State-level count target - if t.target_type == TargetType.COUNT: - if t.variable in ["tax_unit_count", "household_count"]: - marginal_targets["state_fips"][state_constraint] = t.value - - if verbose: - n_states = len(marginal_targets.get("state_fips", {})) - print(f" Built targets for {n_states} states") - - return marginal_targets, continuous_targets - - -def synthesize_district_records( - seed_data: pd.DataFrame, - district_id: str, - n_records: int, - maf: Optional[ConditionalMAF] = None, - seed: int = 42, - device: str = "cpu", -) -> pd.DataFrame: - """ - Generate synthetic records for a specific district. - - Uses bootstrap resampling from state-level seed data with optional noise. - If a trained normalizing flow is provided, uses it for synthesis. - - Args: - seed_data: Seed population DataFrame - district_id: FIPS code for district (e.g., "06" for California) - n_records: Number of records to generate - maf: Optional trained ConditionalMAF for flow-based generation - seed: Random seed for reproducibility - device: Device for flow generation (if using MAF) - - Returns: - DataFrame with synthetic records for the district - """ - np.random.seed(seed) - - # Get state from district (first 2 digits of FIPS) - if len(district_id) >= 2: - state_fips = int(district_id[:2]) - else: - state_fips = int(district_id) - - # Filter to records from same state for context sampling - state_mask = seed_data["state_fips"].astype(int) == state_fips - state_records = seed_data[state_mask] - - if len(state_records) == 0: - state_records = seed_data - - if maf is not None and hasattr(maf, '_X_mean'): - # Flow-based generation - context_indices = np.random.choice(len(maf._training_context), n_records, replace=True) - context = maf._training_context[context_indices].astype(np.float32) - - # Generate from flow with tight clipping - X_normalized = maf.generate(context, clip_z=2.5, device=device) - X_normalized = np.clip(X_normalized, -4, 4) - - # Denormalize - X_log = X_normalized * maf._X_std + maf._X_mean - X_log = np.clip(X_log, -10, 18) - - # Inverse log transform for income columns - X = X_log.copy() - for col_idx in maf._log_transform_cols: - X[:, col_idx] = np.expm1(np.clip(X_log[:, col_idx], -10, 16)) - X[:, col_idx] = np.clip(X[:, col_idx], 0, 1e8) - - synthetic = pd.DataFrame(X, columns=maf._cont_vars) - - # Add categorical columns by sampling from state records - cat_cols = ["filing_status", "num_dependents", "num_ctc_children", - "num_eitc_children", "is_joint"] - for col in cat_cols: - if col in state_records.columns: - synthetic[col] = state_records[col].sample(n_records, replace=True).values - - synthetic["state_fips"] = state_fips - else: - # Bootstrap resampling fallback - indices = np.random.choice(len(state_records), n_records, replace=True) - synthetic = state_records.iloc[indices].copy().reset_index(drop=True) - - # Add noise to continuous variables - noise_cols = ["wage_income", "self_employment_income", "interest_income"] - for col in noise_cols: - if col in synthetic.columns: - noise = np.random.lognormal(0, 0.1, n_records) - synthetic[col] = synthetic[col] * noise - synthetic[col] = np.maximum(synthetic[col], 0) - - # Assign district ID - synthetic["district_id"] = district_id - synthetic["tax_unit_id"] = range(n_records) - synthetic["state_fips"] = state_fips - - # Recalculate totals - income_cols = ["wage_income", "self_employment_income", "interest_income", - "dividend_income", "rental_income", "social_security_income", - "unemployment_compensation", "other_income"] - existing_cols = [c for c in income_cols if c in synthetic.columns] - if existing_cols: - synthetic["total_income"] = synthetic[existing_cols].sum(axis=1) - - # Compute earned income for EITC - if "wage_income" in synthetic.columns: - se_col = synthetic.get("self_employment_income", 0) - if isinstance(se_col, int): - se_col = 0 - synthetic["earned_income"] = synthetic["wage_income"] + np.maximum(se_col, 0) - - # Initialize uniform weights - synthetic["weight"] = 1.0 - - return synthetic - - -@dataclass -class DistrictMicroplex: - """ - US District-level microplex builder. - - Generates synthetic tax units calibrated to geographic targets. - - Args: - n_per_district: Number of records to synthesize per district - target_sparsity: Target sparsity for calibration (0-1) - hidden_dim: Hidden dimension for normalizing flow - n_layers: Number of flow layers - device: Device for training/generation ("cpu" or "cuda") - - Example: - >>> dm = DistrictMicroplex(n_per_district=1000, target_sparsity=0.9) - >>> dm.fit(seed_data, epochs=100) - >>> result = dm.generate(districts=["06", "36", "48"]) - >>> calibrated = dm.calibrate(result, marginal_targets, continuous_targets) - """ - - n_per_district: int = 1000 - target_sparsity: float = 0.9 - hidden_dim: int = 128 - n_layers: int = 4 - device: str = "cpu" - - # Private attributes - _maf: Optional[ConditionalMAF] = field(default=None, repr=False) - _seed_data: Optional[pd.DataFrame] = field(default=None, repr=False) - _calibrator: Optional[Any] = field(default=None, repr=False) - - def __post_init__(self): - """Initialize private attributes.""" - self._maf = None - self._seed_data = None - self._calibrator = None - - def _validate_seed_data(self, seed_data: pd.DataFrame) -> None: - """Validate seed data has required columns.""" - if len(seed_data) == 0: - raise ValueError("Seed data is empty") - - missing = [col for col in REQUIRED_COLUMNS if col not in seed_data.columns] - if missing: - raise ValueError(f"Missing required columns: {missing}") - - def _prepare_training_data( - self, - df: pd.DataFrame, - ) -> Tuple[np.ndarray, np.ndarray, List[str], List[str]]: - """Prepare data for normalizing flow training.""" - continuous_vars = [ - "wage_income", - "self_employment_income", - "interest_income", - "dividend_income", - "rental_income", - "social_security_income", - "head_age", - ] - continuous_vars = [v for v in continuous_vars if v in df.columns] - - condition_vars = ["state_fips", "filing_status", "num_dependents"] - condition_vars = [v for v in condition_vars if v in df.columns] - - mask = df[continuous_vars].notna().all(axis=1) - df_valid = df[mask].copy() - - X = df_valid[continuous_vars].values.astype(np.float32) - - cond_data = [] - for var in condition_vars: - if df_valid[var].dtype == object: - dummies = pd.get_dummies(df_valid[var], prefix=var) - cond_data.append(dummies.values) - else: - vals = df_valid[var].values.astype(np.float32) - cond_data.append((vals - vals.mean()) / (vals.std() + 1e-6)) - - C = np.column_stack(cond_data) if cond_data else None - - return X, C, continuous_vars, condition_vars - - def fit( - self, - seed_data: pd.DataFrame, - epochs: int = 100, - batch_size: int = 512, - lr: float = 1e-3, - verbose: bool = True, - ) -> "DistrictMicroplex": - """ - Train normalizing flow on seed data. - - Args: - seed_data: Calibrated tax unit data (e.g., from CPS) - epochs: Training epochs - batch_size: Batch size for training - lr: Learning rate - verbose: Print training progress - - Returns: - self for method chaining - """ - self._validate_seed_data(seed_data) - self._seed_data = seed_data - - X, C, cont_vars, cond_vars = self._prepare_training_data(seed_data) - - if verbose: - print("Training ConditionalMAF:") - print(f" Continuous vars: {cont_vars}") - print(f" Condition shape: {C.shape if C is not None else 'None'}") - print(f" Data shape: {X.shape}") - - # Log-transform income variables - X_log = X.copy() - n_income_cols = min(6, X.shape[1] - 1) - X_log[:, :n_income_cols] = np.log1p(np.maximum(X_log[:, :n_income_cols], 0)) - - X_mean = X_log.mean(axis=0) - X_std = X_log.std(axis=0) + 1e-6 - X_normalized = (X_log - X_mean) / X_std - - n_features = X.shape[1] - n_context = C.shape[1] if C is not None else 0 - - self._maf = ConditionalMAF( - n_features=n_features, - n_context=n_context, - hidden_dim=self.hidden_dim, - n_layers=self.n_layers, - ) - - if verbose: - print(f" Training for {epochs} epochs on {self.device}...") - - self._maf.fit( - X_normalized, C, - epochs=epochs, - batch_size=batch_size, - lr=lr, - device=self.device, - verbose=verbose, - verbose_freq=max(1, epochs // 10), - ) - - # Store normalization params - self._maf._X_mean = X_mean - self._maf._X_std = X_std - self._maf._training_context = C - self._maf._cont_vars = cont_vars - self._maf._cond_vars = cond_vars - self._maf._log_transform_cols = list(range(n_income_cols)) - - return self - - def generate( - self, - districts: List[str], - verbose: bool = True, - ) -> pd.DataFrame: - """ - Generate synthetic records for specified districts. - - Args: - districts: List of district FIPS codes (e.g., ["06", "36"]) - verbose: Print progress - - Returns: - DataFrame with synthetic records for all districts - """ - if self._seed_data is None: - raise ValueError("Must call fit() before generate()") - - all_synthetic = [] - for i, district_id in enumerate(districts): - synthetic = synthesize_district_records( - seed_data=self._seed_data, - district_id=district_id, - n_records=self.n_per_district, - maf=self._maf, - seed=42 + i, - device=self.device, - ) - all_synthetic.append(synthetic) - - if verbose and (i + 1) % 10 == 0: - print(f" Generated {i+1}/{len(districts)} districts") - - combined = pd.concat(all_synthetic, ignore_index=True) - - if verbose: - print(f" Total synthetic records: {len(combined):,}") - - return combined - - def calibrate( - self, - synthetic: pd.DataFrame, - marginal_targets: Dict[str, Dict], - continuous_targets: Dict[str, float], - verbose: bool = True, - ) -> pd.DataFrame: - """ - Calibrate synthetic population to targets. - - Args: - synthetic: Synthetic population DataFrame - marginal_targets: Categorical targets {var: {category: count}} - continuous_targets: Sum targets {var: total} - verbose: Print progress - - Returns: - Calibrated DataFrame with adjusted weights - """ - if MicroplexSparseCalibrator is None: - raise ImportError( - "microplex.SparseCalibrator is required for district calibration " - "but is not available in the installed microplex package." - ) - - self._calibrator = MicroplexSparseCalibrator( - target_sparsity=self.target_sparsity, - max_iter=2000, - tol=1e-6, - ) - - if verbose: - n_cat = sum(len(v) for v in marginal_targets.values()) - n_cont = len(continuous_targets) - print(f"Calibrating {len(synthetic):,} records to {n_cat + n_cont} targets...") - - start = time.time() - result = self._calibrator.fit_transform(synthetic, marginal_targets, continuous_targets) - elapsed = time.time() - start - - if verbose: - val = self._calibrator.validate(result) - print(f" Time: {elapsed:.1f}s") - print(f" Sparsity: {self._calibrator.get_sparsity():.1%}") - print(f" Max error: {val['max_error']:.2%}") - print(f" Mean error: {val['mean_error']:.2%}") - - return result - - def build( - self, - seed_data: pd.DataFrame, - districts: List[str], - marginal_targets: Optional[Dict[str, Dict]] = None, - continuous_targets: Optional[Dict[str, float]] = None, - targets: Optional[List[TargetSpec]] = None, - epochs: int = 100, - verbose: bool = True, - ) -> pd.DataFrame: - """ - Build complete district microplex (fit, generate, calibrate). - - Args: - seed_data: Calibrated tax unit seed data - districts: List of district FIPS codes - marginal_targets: Categorical calibration targets (or use targets param) - continuous_targets: Sum calibration targets (or use targets param) - targets: List of TargetSpec from database (alternative to dicts) - epochs: Training epochs for normalizing flow - verbose: Print progress - - Returns: - Calibrated synthetic population DataFrame - """ - if verbose: - print("=" * 60) - print("BUILDING US DISTRICT MICROPLEX") - print("=" * 60) - - # Step 1: Fit - if verbose: - print("\n1. Training synthesizer...") - self.fit(seed_data, epochs=epochs, verbose=verbose) - - # Step 2: Generate - if verbose: - print(f"\n2. Synthesizing for {len(districts)} districts...") - synthetic = self.generate(districts, verbose=verbose) - synthetic["state_fips"] = synthetic["state_fips"].astype(int) - - # Step 3: Build targets if not provided - if marginal_targets is None or continuous_targets is None: - if verbose: - print("\n3. Loading targets from database...") - marginal_targets, continuous_targets = build_targets_from_db(year=2021, verbose=verbose) - - # Fill in missing state targets from seed data - states = sorted(seed_data["state_fips"].dropna().unique().astype(int)) - for state in states: - if state not in marginal_targets.get("state_fips", {}): - mask = seed_data["state_fips"].astype(int) == state - target_count = seed_data.loc[mask, "weight"].sum() - marginal_targets["state_fips"][state] = float(target_count) - - # Step 4: Calibrate - if verbose: - print("\n4. Calibrating weights...") - calibrated = self.calibrate(synthetic, marginal_targets, continuous_targets, verbose=verbose) - - if verbose: - print("\n" + "=" * 60) - print("DISTRICT MICROPLEX COMPLETE") - print("=" * 60) - non_zero = (calibrated["weight"] > 1e-9).sum() - print(f"Total records: {len(calibrated):,}") - print(f"Non-zero weights: {non_zero:,} ({non_zero/len(calibrated):.1%})") - print(f"Weighted population: {calibrated['weight'].sum():,.0f}") - - return calibrated - - -def load_seed_data(data_path: Optional[Path] = None) -> pd.DataFrame: - """ - Load calibrated tax unit seed data. - - Args: - data_path: Path to parquet file (default: looks in standard locations) - - Returns: - DataFrame with calibrated tax units - """ - if data_path is None: - # Try standard locations - candidates = [ - Path(__file__).parent / "cps_2024.parquet", - Path(__file__).parent.parent.parent / "tax_units_calibrated_gradient_2024.parquet", - ] - for p in candidates: - if p.exists(): - data_path = p - break - - if data_path is None or not data_path.exists(): - raise FileNotFoundError( - "Calibrated tax units not found. Run gradient_calibrate.py first." - ) - - return pd.read_parquet(data_path) diff --git a/micro/us/hierarchical_pipeline.py b/micro/us/hierarchical_pipeline.py deleted file mode 100644 index 6fb38a92..00000000 --- a/micro/us/hierarchical_pipeline.py +++ /dev/null @@ -1,763 +0,0 @@ -""" -Hierarchical Microplex Pipeline: Calibrate household weights with person-level targets. - -Reads hierarchical microdata (households + persons) and calibration targets, -aggregates person-level targets to household level, runs IPF calibration on -household weights, and writes results back to Supabase. - -Key difference from flat pipeline: -- Weights are at household level (one weight per household) -- Targets can be at person level (e.g., count of children aged 0-17) -- Person-level targets are aggregated to household level before calibration - -Usage: - python -m micro.us.hierarchical_pipeline --year 2024 - python -m micro.us.hierarchical_pipeline --year 2024 --dry-run - python -m micro.us.hierarchical_pipeline --year 2024 --use-mock -""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -from dataclasses import dataclass -from typing import Any, Dict, List, Tuple - -# Lazy import of Supabase to avoid import errors when not available -_supabase_client = None - -def _get_supabase_client(): - """Lazy import of Supabase client.""" - global _supabase_client - if _supabase_client is None: - try: - from db.supabase_client import get_supabase_client - _supabase_client = get_supabase_client() - except ImportError: - raise ImportError( - "Supabase client not available. " - "Install with: pip install supabase" - ) - return _supabase_client - - -@dataclass -class HierarchicalCalibrationResult: - """Results from hierarchical IPF calibration.""" - n_households: int - n_persons: int - original_weights: np.ndarray - calibrated_weights: np.ndarray - adjustment_factors: np.ndarray - n_constraints: int - success: bool - message: str - l2_loss: float - max_error: float - - -def load_hierarchical_data_from_supabase( - year: int = 2024, - limit: int = 100000, - use_mock: bool = False, -) -> Tuple[pd.DataFrame, pd.DataFrame]: - """ - Load hierarchical microdata (households + persons) from Supabase. - - Args: - year: Data year - limit: Maximum households to load - use_mock: Use mock data instead of Supabase - - Returns: - Tuple of (household_df, person_df) - """ - if use_mock: - return _create_mock_hierarchical_data(limit) - - print(f"Loading hierarchical data for {year} from Supabase...") - - # Try microplex.households table first - client = _get_supabase_client() - - try: - # Query households - print(" Loading households...") - hh_result = client.schema("microplex").table("households").select("*").limit(limit).execute() - hh_df = pd.DataFrame(hh_result.data) - print(f" Loaded {len(hh_df):,} households") - - if len(hh_df) == 0: - raise ValueError("No household data found") - - # Query persons - print(" Loading persons...") - # Get household IDs to filter persons - hh_ids = hh_df["household_id"].tolist() - - # Paginate person queries since there are many persons per household - all_persons = [] - page_size = 1000 - offset = 0 - - while True: - person_result = ( - client.schema("microplex") - .table("persons") - .select("*") - .in_("household_id", hh_ids) - .range(offset, offset + page_size - 1) - .execute() - ) - if not person_result.data: - break - all_persons.extend(person_result.data) - offset += len(person_result.data) - if len(person_result.data) < page_size: - break - - person_df = pd.DataFrame(all_persons) - print(f" Loaded {len(person_df):,} persons") - - return hh_df, person_df - - except Exception as e: - print(f" Supabase query failed: {e}") - print(" Falling back to mock data") - return _create_mock_hierarchical_data(limit) - - -def _create_mock_hierarchical_data(limit: int = 100) -> Tuple[pd.DataFrame, pd.DataFrame]: - """Create mock hierarchical data for testing.""" - np.random.seed(42) - - n_households = min(limit, 1000) - - # Generate households - hh_df = pd.DataFrame({ - "household_id": range(1, n_households + 1), - "weight": np.random.uniform(50, 500, n_households), - "state_fips": np.random.choice([6, 36, 48, 12, 17], n_households), # CA, NY, TX, FL, IL - "tenure": np.random.choice(["own", "rent"], n_households, p=[0.6, 0.4]), - }) - - # Generate persons (1-4 per household) - persons = [] - person_id = 1 - for hh_id in hh_df["household_id"]: - n_persons = np.random.randint(1, 5) - for i in range(n_persons): - age = np.random.choice( - [5, 10, 15, 25, 35, 45, 55, 65, 75], - p=[0.08, 0.08, 0.08, 0.15, 0.18, 0.15, 0.12, 0.10, 0.06], - ) - income = max(0, np.random.normal(40000, 30000)) if age >= 18 else 0 - persons.append({ - "person_id": person_id, - "household_id": hh_id, - "age": age, - "is_male": np.random.choice([True, False]), - "employment_income": income, - "is_snap_recipient": np.random.random() < 0.12, - "is_medicaid_enrolled": np.random.random() < 0.20, - }) - person_id += 1 - - person_df = pd.DataFrame(persons) - - print(f" Created mock data: {len(hh_df):,} households, {len(person_df):,} persons") - return hh_df, person_df - - -def load_targets_from_supabase( - year: int = 2024, - scale_factor: float = 1.0, -) -> List[Dict[str, Any]]: - """ - Load calibration targets from Supabase targets.* schema. - - Falls back to mock targets if not available. - - Args: - year: Target year - scale_factor: Scale factor for mock targets (based on sample size) - - Returns: - List of target specifications - """ - print(f"Loading targets for {year}...") - - try: - client = _get_supabase_client() - - # Try to query targets from targets schema - result = client.schema("microplex").table("targets").select( - "*, strata(*, stratum_constraints(*))" - ).eq("period", year).execute() - - if result.data and len(result.data) > 10: - print(f" Loaded {len(result.data)} targets from Supabase") - return _convert_supabase_targets(result.data) - - except Exception as e: - print(f" Supabase targets query failed: {e}") - - # Fall back to mock targets (scaled to sample) - print(" Using mock targets") - return _create_mock_targets(scale_factor=scale_factor) - - -def _convert_supabase_targets(raw_targets: List[Dict]) -> List[Dict[str, Any]]: - """Convert Supabase target format to internal format.""" - targets = [] - - for t in raw_targets: - strata = t.get("strata", {}) - constraints = [] - - for c in strata.get("stratum_constraints", []): - constraints.append(( - c.get("variable"), - c.get("operator", "=="), - c.get("value"), - )) - - targets.append({ - "variable": t.get("variable"), - "value": float(t.get("value", 0)), - "target_type": t.get("target_type", "count"), - "constraints": constraints, - "stratum_name": strata.get("name"), - "source": t.get("source"), - }) - - return targets - - -def _create_mock_targets(scale_factor: float = 1.0) -> List[Dict[str, Any]]: - """Create mock targets for testing. - - Args: - scale_factor: Multiplier for target values (use < 1 for smaller samples) - - Returns: - List of mock target specifications - """ - # Base targets (approximately proportional to US population) - base_targets = [ - # Total population - { - "variable": "person_count", - "value": 330000000 * scale_factor, - "target_type": "count", - "constraints": [], - "stratum_name": "Total US Population", - }, - # Children under 18 (~22%) - { - "variable": "person_count", - "value": 73000000 * scale_factor, - "target_type": "count", - "constraints": [("age", ">=", "0"), ("age", "<", "18")], - "stratum_name": "Children under 18", - }, - # Adults 18-64 (~61%) - { - "variable": "person_count", - "value": 200000000 * scale_factor, - "target_type": "count", - "constraints": [("age", ">=", "18"), ("age", "<", "65")], - "stratum_name": "Working age adults", - }, - # Seniors 65+ (~17%) - { - "variable": "person_count", - "value": 57000000 * scale_factor, - "target_type": "count", - "constraints": [("age", ">=", "65")], - "stratum_name": "Seniors 65+", - }, - # SNAP recipients (~13%) - { - "variable": "person_count", - "value": 42000000 * scale_factor, - "target_type": "count", - "constraints": [("is_snap_recipient", "==", "True")], - "stratum_name": "SNAP recipients", - }, - # Total household count - { - "variable": "household_count", - "value": 130000000 * scale_factor, - "target_type": "count", - "constraints": [], - "stratum_name": "Total households", - }, - ] - return base_targets - - -def build_hierarchical_constraints( - hh_df: pd.DataFrame, - person_df: pd.DataFrame, - targets: List[Dict[str, Any]], - hh_id_col: str = "household_id", - min_obs: int = 10, -) -> List[Dict]: - """ - Build calibration constraints aggregating person targets to household level. - - For each target, creates an indicator vector at the household level: - - For household-level targets: indicator is 1/0 per household - - For person-level targets: indicator is count or sum per household - - Args: - hh_df: Household DataFrame with weight column - person_df: Person DataFrame with household_id - targets: List of target specifications - hh_id_col: Column linking persons to households - min_obs: Minimum observations for a constraint - - Returns: - List of constraint dicts with 'indicator' and 'target_value' - """ - constraints = [] - for target in targets: - variable = target.get("variable", "") - value = target.get("value", 0) - target_type = target.get("target_type", "count") - target_constraints = target.get("constraints", []) - stratum_name = target.get("stratum_name", "unknown") - - # Determine if this is a household or person level target - is_household_level = ( - "household" in variable.lower() - or _all_constraints_are_household_level(target_constraints) - ) - - if is_household_level: - # Direct household-level constraint - indicator = _build_household_indicator( - hh_df, target_constraints, variable, target_type, hh_id_col - ) - else: - # Aggregate from person to household - indicator = _build_aggregated_indicator( - hh_df, person_df, target_constraints, variable, target_type, hh_id_col - ) - - n_obs = np.sum(indicator > 0) - if n_obs >= min_obs: - constraints.append({ - "indicator": indicator, - "target_value": value, - "variable": variable, - "target_type": target_type, - "stratum": stratum_name, - "n_obs": int(n_obs), - }) - - print(f" Built {len(constraints)} constraints from {len(targets)} targets") - return constraints - - -def _all_constraints_are_household_level(constraints: List[Tuple]) -> bool: - """Check if all constraint variables are household-level. - - Returns False if empty (to default to person-level aggregation). - """ - HOUSEHOLD_VARS = {"state_fips", "tenure", "household_size", "household_income"} - # If no constraints, default to person level (more common case) - if not constraints: - return False - for var, _, _ in constraints: - if var not in HOUSEHOLD_VARS: - return False - return True - - -def _build_household_indicator( - hh_df: pd.DataFrame, - constraints: List[Tuple], - variable: str, - target_type: str, - hh_id_col: str, -) -> np.ndarray: - """Build indicator for household-level targets.""" - n = len(hh_df) - mask = np.ones(n, dtype=bool) - - for var, op, val in constraints: - if var not in hh_df.columns: - continue - - col = hh_df[var] - if pd.api.types.is_numeric_dtype(col): - val = float(val) - - if op == "==": - mask &= (col == val).values - elif op == "!=": - mask &= (col != val).values - elif op == ">=": - mask &= (col >= val).values - elif op == ">": - mask &= (col > val).values - elif op == "<=": - mask &= (col <= val).values - elif op == "<": - mask &= (col < val).values - - if target_type == "count": - return mask.astype(float) - elif target_type == "amount": - # For amounts, multiply by variable value - var_name = variable.split("#")[-1] if "#" in variable else variable - if var_name in hh_df.columns: - return (hh_df[var_name].values * mask).astype(float) - return mask.astype(float) - else: - return mask.astype(float) - - -def _build_aggregated_indicator( - hh_df: pd.DataFrame, - person_df: pd.DataFrame, - constraints: List[Tuple], - variable: str, - target_type: str, - hh_id_col: str, -) -> np.ndarray: - """ - Build indicator by aggregating from person to household. - - For COUNT targets: count matching persons per household - For AMOUNT targets: sum variable value for matching persons per household - """ - # Apply constraints to person data - mask = pd.Series(True, index=person_df.index) - - for var, op, val in constraints: - if var not in person_df.columns: - continue - - col = person_df[var] - - # Handle boolean string values - if val in ("True", "true", "1"): - val = True - elif val in ("False", "false", "0"): - val = False - elif pd.api.types.is_numeric_dtype(col): - val = float(val) - - if op == "==": - mask &= (col == val) - elif op == "!=": - mask &= (col != val) - elif op == ">=": - mask &= (col >= val) - elif op == ">": - mask &= (col > val) - elif op == "<=": - mask &= (col <= val) - elif op == "<": - mask &= (col < val) - - filtered = person_df[mask] - - if len(filtered) == 0: - return np.zeros(len(hh_df)) - - if target_type == "count": - # Count matching persons per household - agg = filtered.groupby(hh_id_col).size() - elif target_type == "amount": - # Sum variable value for matching persons per household - var_name = variable.split("#")[-1] if "#" in variable else variable - if var_name in filtered.columns: - agg = filtered.groupby(hh_id_col)[var_name].sum() - else: - # Default to count if variable not found - agg = filtered.groupby(hh_id_col).size() - else: - agg = filtered.groupby(hh_id_col).size() - - # Map back to household order (households with no matching persons get 0) - indicator = hh_df[hh_id_col].map(agg).fillna(0).values - - return indicator.astype(float) - - -def run_hierarchical_ipf( - original_weights: np.ndarray, - constraints: List[Dict], - bounds: Tuple[float, float] = (0.2, 5.0), - max_iter: int = 100, - damping: Tuple[float, float] = (0.9, 1.1), - verbose: bool = True, -) -> Tuple[np.ndarray, bool, float]: - """ - Run IPF calibration on household weights. - - Same algorithm as flat IPF but operates on household-level weights - with constraints that may aggregate person-level targets. - - Args: - original_weights: Initial household weights - constraints: List of constraint dicts with 'indicator' and 'target_value' - bounds: Min/max weight adjustment factors - max_iter: Number of IPF iterations - damping: Min/max adjustment ratio per iteration - verbose: Print progress - - Returns: - (calibrated_weights, success, l2_loss) - """ - n = len(original_weights) - m = len(constraints) - - if m == 0: - if verbose: - print(" No constraints, returning original weights") - return original_weights.copy(), True, 0.0 - - if verbose: - print(f" IPF calibration: {n:,} households, {m} constraints, {max_iter} iterations") - - # Build constraint matrix - A = np.zeros((m, n)) - targets = np.zeros(m) - - for j, c in enumerate(constraints): - A[j, :] = c["indicator"] - targets[j] = c["target_value"] - - w = original_weights.copy() - - for iteration in range(max_iter): - for j in range(m): - achieved = np.dot(A[j], w) - if np.isfinite(achieved) and achieved > 0 and targets[j] > 0: - ratio = np.clip(targets[j] / achieved, damping[0], damping[1]) - mask = A[j] != 0 - w[mask] *= ratio - - # Apply bounds after each full iteration - # Guard against division issues - with np.errstate(divide='ignore', invalid='ignore'): - adj = np.where(original_weights > 0, w / original_weights, 1.0) - adj = np.clip(adj, bounds[0], bounds[1]) - w = original_weights * adj - - # Clean up any NaN/inf values - w = np.where(np.isfinite(w), w, original_weights) - - # Compute L2 loss (squared relative error) - achieved = A @ w - with np.errstate(divide='ignore', invalid='ignore'): - rel_errors = np.where(targets != 0, (achieved - targets) / targets, 0) - l2_loss = float(np.mean(rel_errors ** 2)) - - # Check convergence (all targets within 5%) - max_error = float(np.max(np.abs(rel_errors))) - success = max_error < 0.05 - - if verbose: - print(f" IPF converged: max error = {max_error:.1%}, L2 loss = {l2_loss:.6f}") - - return w, success, l2_loss - - -def write_households_to_supabase( - hh_df: pd.DataFrame, - year: int, - chunk_size: int = 200, -) -> int: - """Write calibrated household weights back to Supabase.""" - client = _get_supabase_client() - - print(f"Writing {len(hh_df):,} calibrated households...") - - # Update existing records with calibrated weights - total = 0 - for i in range(0, len(hh_df), chunk_size): - chunk = hh_df.iloc[i:i + chunk_size] - - for _, row in chunk.iterrows(): - hh_id = row["household_id"] - calibrated_weight = float(row["calibrated_weight"]) - weight_adjustment = float(row.get("weight_adjustment", 1.0)) - - try: - client.schema("microplex").table("households").update({ - "calibrated_weight": calibrated_weight, - "weight_adjustment": weight_adjustment, - }).eq("household_id", hh_id).execute() - total += 1 - except Exception as e: - print(f" Error updating household {hh_id}: {e}") - - if (i + chunk_size) % 1000 == 0: - print(f" Updated {total:,} / {len(hh_df):,}") - - print(f" Done: {total:,} households updated") - return total - - -def run_hierarchical_pipeline( - year: int = 2024, - dry_run: bool = False, - use_mock: bool = False, - limit: int = 100000, - verbose: bool = True, -) -> pd.DataFrame: - """ - Run the full hierarchical microplex pipeline. - - Args: - year: Data year - dry_run: Don't write to Supabase - use_mock: Use mock data instead of Supabase - limit: Maximum households to process - verbose: Print progress - - Returns: - Calibrated household DataFrame - """ - if verbose: - print("=" * 60) - print("HIERARCHICAL MICROPLEX PIPELINE") - print("=" * 60) - - # Load hierarchical data - hh_df, person_df = load_hierarchical_data_from_supabase( - year=year, limit=limit, use_mock=use_mock - ) - - if verbose: - print("\nData loaded:") - print(f" Households: {len(hh_df):,}") - print(f" Persons: {len(person_df):,}") - print(f" Avg persons/household: {len(person_df) / len(hh_df):.1f}") - - # Compute scale factor for mock targets (based on sample weight relative to US pop) - # US has ~330M people and ~130M households - sample_weight = hh_df["weight"].sum() - scale_factor = sample_weight / 130000000.0 if use_mock else 1.0 - - # Load targets - targets = load_targets_from_supabase(year, scale_factor=scale_factor) - - if verbose: - print(f"\nTargets loaded: {len(targets)}") - - # Build constraints - constraints = build_hierarchical_constraints( - hh_df, person_df, targets, hh_id_col="household_id" - ) - - if len(constraints) == 0: - print("\nWARNING: No constraints built, skipping calibration") - hh_df["calibrated_weight"] = hh_df["weight"] - hh_df["weight_adjustment"] = 1.0 - return hh_df - - # Get original weights - original_weights = hh_df["weight"].values.copy() - - if verbose: - print(f"\nOriginal weighted households: {original_weights.sum():,.0f}") - - # Pre-scale to match total target if available - for c in constraints: - if c.get("variable") in ("household_count", "person_count") and c.get("n_obs") == len(hh_df): - target_total = c["target_value"] - scale = target_total / original_weights.sum() - original_weights *= scale - if verbose: - print(f"Pre-scaled weights by {scale:.3f} to match {c['variable']} target") - break - - # Run IPF calibration - calibrated_weights, success, l2_loss = run_hierarchical_ipf( - original_weights, constraints, verbose=verbose - ) - - # Compute max error - max_error = 0.0 - for c in constraints: - achieved = np.dot(calibrated_weights, c["indicator"]) - if c["target_value"] > 0: - error = abs((achieved - c["target_value"]) / c["target_value"]) - max_error = max(max_error, error) - - # Add results to DataFrame - hh_df["original_weight"] = hh_df["weight"] - hh_df["calibrated_weight"] = calibrated_weights - hh_df["weight_adjustment"] = calibrated_weights / hh_df["weight"].values - - # Summary - if verbose: - print("\n" + "=" * 60) - print("SUMMARY") - print("=" * 60) - print(f"Households: {len(hh_df):,}") - print(f"Persons: {len(person_df):,}") - print(f"Constraints: {len(constraints)}") - print(f"Original weighted: {hh_df['original_weight'].sum():,.0f}") - print(f"Calibrated weighted: {calibrated_weights.sum():,.0f}") - print(f"Success: {success}") - print(f"Max error: {max_error:.1%}") - print(f"L2 loss: {l2_loss:.6f}") - - # Weight adjustment statistics - adj = hh_df["weight_adjustment"] - print("\nWeight adjustments:") - print(f" Mean: {adj.mean():.3f}") - print(f" Std: {adj.std():.3f}") - print(f" Min: {adj.min():.3f}") - print(f" Max: {adj.max():.3f}") - - # Write to Supabase - if not dry_run: - write_households_to_supabase(hh_df, year) - else: - if verbose: - print("\nDRY RUN - not writing to Supabase") - - return hh_df - - -def main(): - import argparse - - parser = argparse.ArgumentParser( - description="Run hierarchical microplex pipeline", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -Examples: - python -m micro.us.hierarchical_pipeline --year 2024 --dry-run - python -m micro.us.hierarchical_pipeline --year 2024 --use-mock --dry-run - python -m micro.us.hierarchical_pipeline --year 2024 --limit 10000 - """, - ) - parser.add_argument("--year", type=int, default=2024, help="Data year (default: 2024)") - parser.add_argument("--dry-run", action="store_true", help="Don't write to Supabase") - parser.add_argument("--use-mock", action="store_true", help="Use mock data instead of Supabase") - parser.add_argument("--limit", type=int, default=100000, help="Max households to load (default: 100000)") - parser.add_argument("--quiet", "-q", action="store_true", help="Reduce output verbosity") - - args = parser.parse_args() - - run_hierarchical_pipeline( - year=args.year, - dry_run=args.dry_run, - use_mock=args.use_mock, - limit=args.limit, - verbose=not args.quiet, - ) - - -if __name__ == "__main__": - main() diff --git a/micro/us/hierarchy_validation.py b/micro/us/hierarchy_validation.py deleted file mode 100644 index f1c623ba..00000000 --- a/micro/us/hierarchy_validation.py +++ /dev/null @@ -1,517 +0,0 @@ -"""Compare flat tax-unit and household-weight Microplex calibration.""" - -from __future__ import annotations - -from dataclasses import dataclass, replace -from pathlib import Path -from typing import Any - -import numpy as np -import pandas as pd - -from arch.targets import TargetSpec -from micro.us.entities import MicroplexEntityFrames, build_microplex_entities -from micro.us.pipeline import ( - CalibrationResult, - calibrate_household_weights, - calibrate_weights, - load_cps_from_local_file, - load_cps_from_supabase, - load_targets_from_supabase, - maybe_add_policyengine_income_tax, - print_target_composition_diagnostics, -) -from micro.us.targets import MicroplexTargetProfile, compose_microplex_targets -from micro.us.validation_dashboard import ( - build_microplex_version_id, - current_git_sha, - write_calibration_dashboard, -) - - -@dataclass(frozen=True) -class HierarchyComparisonResult: - """Outputs from flat-vs-household calibration comparison.""" - - entities: MicroplexEntityFrames - flat_tax_units: pd.DataFrame - flat_result: CalibrationResult - household_result: CalibrationResult - target_comparison: pd.DataFrame - summary: pd.DataFrame - - -def compare_flat_vs_household_calibration( - *, - year: int = 2024, - limit: int | None = None, - microdata_source: str = "local", - cps_path: Path | None = None, - target_source: str = "db", - db_path: Path | None = None, - age_soi: bool = True, - include_amount_targets: bool = True, - min_target_obs: int = 100, - calibration_method: str = "auto", - min_weight_factor: float = 0.1, - max_weight_factor: float = 20.0, - add_policyengine_tax: bool = True, - output_dir: Path | None = None, - reports_root: Path | None = None, - version_id: str | None = None, - verbose: bool = True, -) -> HierarchyComparisonResult: - """Run flat and household calibration on the same Microplex entity build.""" - report_config = _report_config( - year=year, - limit=limit, - microdata_source=microdata_source, - cps_path=cps_path, - target_source=target_source, - db_path=db_path, - age_soi=age_soi, - include_amount_targets=include_amount_targets, - min_target_obs=min_target_obs, - calibration_method=calibration_method, - min_weight_factor=min_weight_factor, - max_weight_factor=max_weight_factor, - add_policyengine_tax=add_policyengine_tax, - ) - git_sha = current_git_sha() - version_id = version_id or build_microplex_version_id( - year=year, - config=report_config, - git_sha=git_sha, - ) - output_dir = _resolve_output_dir( - output_dir=output_dir, - reports_root=reports_root, - version_id=version_id, - ) - - if microdata_source == "local": - persons = load_cps_from_local_file(year, path=cps_path, limit=limit) - elif microdata_source == "supabase": - persons = load_cps_from_supabase(year, limit=limit or 200_000) - else: - raise ValueError(f"Unknown microdata_source: {microdata_source}") - - targets, holdout_variables = _load_targets_and_holdouts( - year=year, - target_source=target_source, - db_path=db_path, - age_soi=age_soi, - verbose=verbose, - ) - - entities = build_microplex_entities(persons) - if add_policyengine_tax: - entities = replace( - entities, - tax_units=maybe_add_policyengine_income_tax( - entities.tax_units, - targets, - year=year, - persons=entities.persons, - ), - ) - - flat_tax_units = entities.tax_units[ - entities.tax_units["is_tax_filer"] == 1 - ].copy() - if flat_tax_units.empty: - raise ValueError("No filing tax units available for flat comparison.") - - if verbose: - print("\nRunning flat tax-unit calibration...") - flat_result = calibrate_weights( - flat_tax_units, - targets, - include_amounts=include_amount_targets, - min_obs=min_target_obs, - calibration_method=calibration_method, - weight_bounds=(min_weight_factor, max_weight_factor), - holdout_variables=holdout_variables, - verbose=verbose, - ) - - if verbose: - print("\nRunning household-weight calibration...") - household_result = calibrate_household_weights( - entities.households, - entities.tax_units, - targets, - include_amounts=include_amount_targets, - min_obs=min_target_obs, - calibration_method=calibration_method, - weight_bounds=(min_weight_factor, max_weight_factor), - holdout_variables=holdout_variables, - verbose=verbose, - ) - - target_comparison = compare_target_diagnostics( - flat_result.diagnostics, - household_result.diagnostics, - ) - summary = summarize_comparison( - entities=entities, - flat_tax_units=flat_tax_units, - flat_result=flat_result, - household_result=household_result, - target_comparison=target_comparison, - ) - - if output_dir is not None: - write_hierarchy_comparison( - output_dir, - summary=summary, - target_comparison=target_comparison, - flat_diagnostics=flat_result.diagnostics, - household_diagnostics=household_result.diagnostics, - metadata={ - "microplex_version": version_id, - "git_sha": git_sha, - "jurisdiction": "us", - "year": year, - "comparison": "flat_vs_household", - "config": report_config, - }, - ) - - if verbose: - print("\nFlat vs household summary:") - _print_summary(summary) - - return HierarchyComparisonResult( - entities=entities, - flat_tax_units=flat_tax_units, - flat_result=flat_result, - household_result=household_result, - target_comparison=target_comparison, - summary=summary, - ) - - -def compare_target_diagnostics( - flat_diagnostics: pd.DataFrame, - household_diagnostics: pd.DataFrame, -) -> pd.DataFrame: - """Return one row per target with flat and household diagnostic values.""" - key_columns = [ - "target_index", - "source", - "variable", - "target_type", - "period", - "source_period", - "stratum", - "constraints", - "target_value", - ] - flat = flat_diagnostics.add_prefix("flat_") - household = household_diagnostics.add_prefix("household_") - merged = flat.merge( - household, - left_on=[f"flat_{column}" for column in key_columns], - right_on=[f"household_{column}" for column in key_columns], - how="outer", - ) - for column in key_columns: - flat_column = f"flat_{column}" - household_column = f"household_{column}" - merged[column] = merged[flat_column].combine_first(merged[household_column]) - - merged["flat_abs_post_error"] = merged["flat_post_error"].abs() - merged["household_abs_post_error"] = merged["household_post_error"].abs() - merged["abs_post_error_delta"] = ( - merged["household_abs_post_error"] - merged["flat_abs_post_error"] - ) - leading = key_columns + [ - "flat_role", - "household_role", - "flat_status", - "household_status", - "flat_drop_reason", - "household_drop_reason", - "flat_n_obs", - "household_n_obs", - "flat_pre_error", - "household_pre_error", - "flat_post_error", - "household_post_error", - "flat_abs_post_error", - "household_abs_post_error", - "abs_post_error_delta", - ] - trailing = [column for column in merged.columns if column not in leading] - return merged[leading + trailing].sort_values( - ["target_index", "variable", "target_type"], - na_position="last", - ) - - -def summarize_comparison( - *, - entities: MicroplexEntityFrames, - flat_tax_units: pd.DataFrame, - flat_result: CalibrationResult, - household_result: CalibrationResult, - target_comparison: pd.DataFrame, -) -> pd.DataFrame: - """Build compact scalar metrics for flat-vs-household comparison.""" - rows: list[dict[str, Any]] = [ - _metric("households", "count", len(entities.households)), - _metric("persons", "count", len(entities.persons)), - _metric("tax_units", "count", len(entities.tax_units)), - _metric("flat_filing_tax_units", "count", len(flat_tax_units)), - ] - rows.extend(_result_metrics("flat", flat_result)) - rows.extend(_result_metrics("household", household_result)) - rows.extend(_diagnostic_error_metrics("flat", flat_result.diagnostics)) - rows.extend(_diagnostic_error_metrics("household", household_result.diagnostics)) - rows.extend(_delta_metrics(target_comparison)) - return pd.DataFrame(rows) - - -def write_hierarchy_comparison( - output_dir: Path, - *, - summary: pd.DataFrame, - target_comparison: pd.DataFrame, - flat_diagnostics: pd.DataFrame, - household_diagnostics: pd.DataFrame, - metadata: dict[str, Any] | None = None, -) -> None: - """Write comparison artifacts to a directory.""" - dashboard_metadata = { - "microplex_version": output_dir.name, - "jurisdiction": "us", - "comparison": "flat_vs_household", - **(metadata or {}), - } - write_calibration_dashboard( - output_dir, - summary=summary, - target_comparison=target_comparison, - flat_diagnostics=flat_diagnostics, - household_diagnostics=household_diagnostics, - metadata=dashboard_metadata, - ) - - -def _report_config( - *, - year: int, - limit: int | None, - microdata_source: str, - cps_path: Path | None, - target_source: str, - db_path: Path | None, - age_soi: bool, - include_amount_targets: bool, - min_target_obs: int, - calibration_method: str, - min_weight_factor: float, - max_weight_factor: float, - add_policyengine_tax: bool, -) -> dict[str, Any]: - return { - "year": year, - "limit": limit, - "microdata_source": microdata_source, - "cps_path": str(cps_path) if cps_path is not None else None, - "target_source": target_source, - "db_path": str(db_path) if db_path is not None else None, - "age_soi": age_soi, - "include_amount_targets": include_amount_targets, - "min_target_obs": min_target_obs, - "calibration_method": calibration_method, - "min_weight_factor": min_weight_factor, - "max_weight_factor": max_weight_factor, - "add_policyengine_tax": add_policyengine_tax, - } - - -def _resolve_output_dir( - *, - output_dir: Path | None, - reports_root: Path | None, - version_id: str, -) -> Path | None: - if output_dir is not None: - return output_dir - if reports_root is not None: - return reports_root / version_id - return None - - -def _load_targets_and_holdouts( - *, - year: int, - target_source: str, - db_path: Path | None, - age_soi: bool, - verbose: bool, -) -> tuple[list[TargetSpec] | list[dict[str, Any]], tuple[str, ...]]: - if target_source == "db": - profile = MicroplexTargetProfile(age_soi=age_soi) - composition = compose_microplex_targets( - target_year=year, - db_path=db_path, - profile=profile, - ) - if verbose: - print_target_composition_diagnostics(composition) - return composition.targets, profile.holdout_variables - if target_source == "supabase": - targets = load_targets_from_supabase(year) - return targets, () - raise ValueError(f"Unknown target_source: {target_source}") - - -def _result_metrics(prefix: str, result: CalibrationResult) -> list[dict[str, Any]]: - adjustments = result.adjustment_factors - return [ - _metric(prefix, "calibration_unit", result.calibration_unit), - _metric(prefix, "success", result.success), - _metric(prefix, "method", result.method), - _metric(prefix, "l2_loss", result.l2_loss), - _metric(prefix, "original_weight_sum", result.original_weights.sum()), - _metric(prefix, "calibrated_weight_sum", result.calibrated_weights.sum()), - _metric(prefix, "adjustment_mean", adjustments.mean()), - _metric(prefix, "adjustment_min", adjustments.min()), - _metric(prefix, "adjustment_p01", np.quantile(adjustments, 0.01)), - _metric(prefix, "adjustment_p50", np.quantile(adjustments, 0.50)), - _metric(prefix, "adjustment_p99", np.quantile(adjustments, 0.99)), - _metric(prefix, "adjustment_max", adjustments.max()), - ] - - -def _diagnostic_error_metrics( - prefix: str, - diagnostics: pd.DataFrame, -) -> list[dict[str, Any]]: - rows = [] - for role in ["active", "holdout", "diagnostic", "unsupported"]: - role_diagnostics = diagnostics[diagnostics["role"] == role] - rows.append(_metric(prefix, f"{role}_targets", len(role_diagnostics))) - usable = role_diagnostics[role_diagnostics["post_error"].notna()] - if usable.empty: - continue - abs_error = usable["post_error"].abs() - rows.extend( - [ - _metric(prefix, f"{role}_mean_abs_post_error", abs_error.mean()), - _metric(prefix, f"{role}_median_abs_post_error", abs_error.median()), - _metric(prefix, f"{role}_max_abs_post_error", abs_error.max()), - ] - ) - return rows - - -def _delta_metrics(target_comparison: pd.DataFrame) -> list[dict[str, Any]]: - active = target_comparison[ - (target_comparison["flat_role"] == "active") - & (target_comparison["household_role"] == "active") - ] - if active.empty: - return [] - delta = active["abs_post_error_delta"].dropna() - if delta.empty: - return [] - return [ - _metric("delta", "active_mean_abs_post_error", delta.mean()), - _metric("delta", "active_median_abs_post_error", delta.median()), - _metric("delta", "active_household_better_targets", int((delta < 0).sum())), - _metric("delta", "active_flat_better_targets", int((delta > 0).sum())), - ] - - -def _metric(group: str, metric: str, value: Any) -> dict[str, Any]: - if isinstance(value, np.generic): - value = value.item() - return {"group": group, "metric": metric, "value": value} - - -def _print_summary(summary: pd.DataFrame) -> None: - for _, row in summary.iterrows(): - print(f" {row['group']}.{row['metric']}: {row['value']}") - - -def main() -> None: - import argparse - - parser = argparse.ArgumentParser( - description="Compare flat and household-weight Microplex calibration", - ) - parser.add_argument("--year", type=int, default=2024) - parser.add_argument("--limit", type=int, default=None) - parser.add_argument( - "--microdata-source", - choices=["local", "supabase"], - default="local", - ) - parser.add_argument("--cps-path", type=Path, default=None) - parser.add_argument("--target-source", choices=["db", "supabase"], default="db") - parser.add_argument("--db-path", type=Path, default=None) - parser.add_argument( - "--age-soi-targets", - action=argparse.BooleanOptionalAction, - default=True, - ) - parser.add_argument( - "--include-amount-targets", - action=argparse.BooleanOptionalAction, - default=True, - ) - parser.add_argument("--min-target-obs", type=int, default=100) - parser.add_argument( - "--calibration-method", - choices=["auto", "ipf", "generalized-rake"], - default="auto", - ) - parser.add_argument("--min-weight-factor", type=float, default=0.1) - parser.add_argument("--max-weight-factor", type=float, default=20.0) - parser.add_argument( - "--policyengine-tax", - action=argparse.BooleanOptionalAction, - default=True, - ) - parser.add_argument("--output-dir", type=Path, default=None) - parser.add_argument( - "--reports-root", - type=Path, - default=None, - help="Directory where a version-named dashboard directory is created", - ) - parser.add_argument( - "--version-id", - default=None, - help="Microplex version identifier for the validation dashboard", - ) - parser.add_argument("--quiet", action="store_true") - args = parser.parse_args() - - compare_flat_vs_household_calibration( - year=args.year, - limit=args.limit, - microdata_source=args.microdata_source, - cps_path=args.cps_path, - target_source=args.target_source, - db_path=args.db_path, - age_soi=args.age_soi_targets, - include_amount_targets=args.include_amount_targets, - min_target_obs=args.min_target_obs, - calibration_method=args.calibration_method, - min_weight_factor=args.min_weight_factor, - max_weight_factor=args.max_weight_factor, - add_policyengine_tax=args.policyengine_tax, - output_dir=args.output_dir, - reports_root=args.reports_root, - version_id=args.version_id, - verbose=not args.quiet, - ) - - -if __name__ == "__main__": - main() diff --git a/micro/us/pipeline.py b/micro/us/pipeline.py deleted file mode 100644 index c55052aa..00000000 --- a/micro/us/pipeline.py +++ /dev/null @@ -1,1851 +0,0 @@ -""" -Microplex pipeline: build calibrated microdata from Arch target inputs. - -Reads CPS microdata and Arch target inputs, runs calibration, and writes -calibrated microplex output locally or back to Supabase. - -Calibration methods: -- generalized-rake: bounded dual raking for count and amount constraints -- ipf: fast count-only iterative proportional fitting - -Usage: - python -m micro.us.pipeline --year 2024 - python -m micro.us.pipeline --year 2024 --dry-run -""" - -from __future__ import annotations - -import json -import re -from dataclasses import dataclass, replace -from pathlib import Path -from typing import Any, Dict, List - -import numpy as np -import pandas as pd - -from arch.client import get_supabase_client -from arch.microdata import query_cps_asec -from arch.targets import TargetSpec, TargetType, query_targets -from micro.us.entities import ( - build_microplex_entities, - with_household_weights, - write_microplex_entities, -) -from micro.us.policyengine import ( - PolicyEngineNotAvailableError, - add_policyengine_income_tax, - add_policyengine_income_tax_from_persons, -) -from micro.us.targets import ( - MicroplexTargetProfile, - TargetCompositionResult, - compose_microplex_targets, - load_microplex_targets, -) - - -@dataclass -class CalibrationResult: - """Results from Microplex weight calibration.""" - original_weights: np.ndarray - calibrated_weights: np.ndarray - adjustment_factors: np.ndarray - targets_before: Dict[str, Dict] - targets_after: Dict[str, Dict] - diagnostics: pd.DataFrame - success: bool - message: str - l2_loss: float - method: str - calibration_unit: str = "tax_unit" - - -AGING_NOTE_RE = re.compile( - r"SOI aged (?P\d+)->(?P\d+); " - r"(?P[^;)]+)(?:; factor x(?P[-+0-9.eE]+))?" -) -TARGET_VARIABLE_COLUMNS = { - "adjusted_gross_income": "adjusted_gross_income", - "employment_income": "wage_income", - "income_tax_liability": "income_tax_liability", -} -SUPPORTED_TARGET_VARIABLES = {"tax_unit_count", *TARGET_VARIABLE_COLUMNS} -SUPPORTED_CONSTRAINT_VARIABLES = { - "adjusted_gross_income", - "is_tax_filer", - "agi_bracket", - "state_fips", -} - - -def load_cps_from_supabase(year: int, limit: int = 200000) -> pd.DataFrame: - """Load raw CPS person data from Supabase.""" - print(f"Loading CPS ASEC {year} from Supabase...") - df = query_cps_asec(year, table_type="person", limit=limit) - print(f" Loaded {len(df):,} person records") - return df - - -def load_cps_from_local_file( - year: int, - path: Path | None = None, - limit: int | None = None, -) -> pd.DataFrame: - """Load CPS person data from a local parquet file.""" - if path is None: - repo_root = Path(__file__).resolve().parents[2] - path = repo_root / "micro" / "us" / f"cps_{year}.parquet" - - print(f"Loading CPS ASEC {year} from {path}...") - df = pd.read_parquet(path) - if limit is not None: - df = df.head(limit).copy() - print(f" Loaded {len(df):,} person records") - return df - - -def load_targets_from_supabase(year: int) -> List[Dict[str, Any]]: - """Load calibration targets from Supabase.""" - print(f"Loading targets for {year} from Supabase...") - # Get targets for US jurisdictions (both "US" and "US_FEDERAL") - all_targets = query_targets(year=year) - targets = [t for t in all_targets - if t.get("strata", {}).get("jurisdiction", "").startswith("US")] - print(f" Loaded {len(targets)} targets") - return targets - - -def load_targets_from_db( - year: int, - db_path: Path | None = None, - jurisdiction: str = "us", - sources: list[str] | None = None, - variables: list[str] | None = None, -) -> List[TargetSpec]: - """Load calibration target inputs from the local Arch SQLite database.""" - print(f"Loading target inputs for {year} from Arch DB...") - targets = load_microplex_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - sources=sources, - variables=variables, - ) - print(f" Loaded {len(targets)} target inputs") - return targets - - -def build_tax_units(df: pd.DataFrame) -> pd.DataFrame: - """ - Build tax units from person-level CPS data. - - Simplified version - groups by household and assigns tax unit based on - marital status and age. Full version would use relationship variables. - """ - # Use raw_data if available for additional columns - if "raw_data" in df.columns: - # Extract key fields from raw_data - for col in ["A_MARITL", "A_FAMREL", "TAX_ID"]: - df[col.lower()] = df["raw_data"].apply(lambda x: x.get(col) if isinstance(x, dict) else None) - - if "tax_unit_id" in df.columns: - return build_tax_units_from_census_tax_ids(df) - - # Calculate total income - df["total_income"] = _numeric_first( - df, - ["ptotval", "total_person_income", "income"], - ) - df["wage_income"] = _numeric_first( - df, - ["wsal_val", "wage_salary_income", "wage_income"], - ) - df["self_employment_income"] = _numeric_first( - df, - ["semp_val", "self_employment_income"], - ) + _numeric_first(df, ["frse_val", "farm_self_employment_income"]) - - # Simple AGI estimate (wages + SE - 1/2 SE tax) - se_tax = np.maximum(df["self_employment_income"] * 0.0765, 0) - df["adjusted_gross_income"] = df["wage_income"] + df["self_employment_income"] - se_tax / 2 - - # Weight - if "marsupwt" in df.columns: - df["weight"] = pd.to_numeric(df["marsupwt"], errors="coerce").fillna(0) / 100 - elif "weight" in df.columns: - df["weight"] = pd.to_numeric(df["weight"], errors="coerce").fillna(0) - elif "march_supplement_weight" in df.columns: - df["weight"] = pd.to_numeric( - df["march_supplement_weight"], - errors="coerce", - ).fillna(0) / 100 - else: - df["weight"] = 1.0 - - if "state_fips" not in df.columns and "gestfips" in df.columns: - df["state_fips"] = df["gestfips"] - if "age" not in df.columns and "a_age" in df.columns: - df["age"] = df["a_age"] - if "household_id" not in df.columns and "ph_seq" in df.columns: - df["household_id"] = df["ph_seq"] - - # Filter to likely filers - filer_mask = ( - (df["total_income"] > 13850) - | (df["wage_income"] > 0) - | (df["self_employment_income"] > 0) - ) - df = df[filer_mask].copy() - df["is_tax_filer"] = 1 - print(f" Filtered to {len(df):,} likely filers") - - return df - - -def build_tax_units_from_census_tax_ids(df: pd.DataFrame) -> pd.DataFrame: - """Aggregate person-level CPS records to Census-provided tax units.""" - df = df.copy() - - income_columns = [ - "total_person_income", - "income", - "wage_salary_income", - "wage_income", - "self_employment_income", - "farm_self_employment_income", - "interest_income", - "dividend_income", - "rental_income", - "unemployment_compensation", - "other_income", - ] - for column in income_columns: - if column in df.columns: - df[column] = pd.to_numeric(df[column], errors="coerce").fillna(0) - - if "weight" in df.columns: - df["weight"] = pd.to_numeric(df["weight"], errors="coerce").fillna(0) - elif "march_supplement_weight" in df.columns: - df["weight"] = pd.to_numeric( - df["march_supplement_weight"], - errors="coerce", - ).fillna(0) / 100 - else: - df["weight"] = 1.0 - - if "line_number" in df.columns: - sort_column = "line_number" - elif "person_seq" in df.columns: - sort_column = "person_seq" - else: - sort_column = None - - group_columns = ["tax_unit_id"] - if "household_id" in df.columns: - group_columns = ["household_id", "tax_unit_id"] - - rows = [] - for _, group in df.groupby(group_columns, sort=False, dropna=False): - if sort_column is not None: - group = group.sort_values(sort_column) - head = group.iloc[0] - - wage_income = _sum_columns(group, ["wage_salary_income", "wage_income"]) - self_employment_income = _sum_columns( - group, - ["self_employment_income", "farm_self_employment_income"], - ) - interest_income = _sum_columns(group, ["interest_income"]) - dividend_income = _sum_columns(group, ["dividend_income"]) - rental_income = _sum_columns(group, ["rental_income"]) - unemployment = _sum_columns(group, ["unemployment_compensation"]) - other_income = _sum_columns(group, ["other_income"]) - total_income = _sum_columns(group, ["total_person_income"]) - if total_income == 0: - total_income = _sum_columns(group, ["income"]) - - se_tax_adjustment = max(self_employment_income, 0) * 0.0765 / 2 - adjusted_gross_income = ( - wage_income - + self_employment_income - - se_tax_adjustment - + interest_income - + dividend_income - + rental_income - + unemployment - + other_income - ) - - rows.append( - { - "tax_unit_id": head.get("tax_unit_id"), - "household_id": head.get("household_id"), - "person_id": head.get("person_id"), - "person_count": len(group), - "age": head.get("age", head.get("a_age")), - "state_fips": head.get("state_fips", head.get("gestfips")), - "weight": head["weight"], - "total_income": total_income, - "wage_income": wage_income, - "self_employment_income": self_employment_income, - "interest_income": interest_income, - "dividend_income": dividend_income, - "rental_income": rental_income, - "unemployment_compensation": unemployment, - "other_income": other_income, - "adjusted_gross_income": adjusted_gross_income, - } - ) - - tax_units = pd.DataFrame(rows) - filer_mask = ( - (tax_units["total_income"] > 13_850) - | (tax_units["wage_income"] > 0) - | (tax_units["self_employment_income"] != 0) - ) - tax_units = tax_units[filer_mask].copy() - tax_units["is_tax_filer"] = 1 - print(f" Aggregated to {len(tax_units):,} likely filing tax units") - return tax_units - - -def maybe_add_policyengine_income_tax( - df: pd.DataFrame, - targets: list[TargetSpec] | list[dict[str, Any]], - *, - year: int, - persons: pd.DataFrame | None = None, -) -> pd.DataFrame: - """Calculate income tax liability when loaded targets require it.""" - if "income_tax_liability" in df.columns: - return df - if not _targets_include_variable(targets, "income_tax_liability"): - return df - - print("Calculating income_tax_liability with PolicyEngine-US...") - try: - if persons is not None and _can_merge_policyengine_person_hierarchy( - df, - persons, - ): - hierarchical = add_policyengine_income_tax_from_persons( - persons, - year=year, - ) - merged = _merge_policyengine_tax_units(df, hierarchical) - if "income_tax_liability" in merged.columns: - print(" Used person/household hierarchy for PolicyEngine-US.") - return merged - print(" Could not match hierarchical PE tax units; using aggregate rows.") - return add_policyengine_income_tax(df, year=year) - except PolicyEngineNotAvailableError as exc: - print(f" {exc}") - print(" income_tax_liability targets will remain unsupported.") - return df - - -def _targets_include_variable( - targets: list[TargetSpec] | list[dict[str, Any]], - variable: str, -) -> bool: - for target in targets: - if isinstance(target, TargetSpec): - target_variable = target.variable - else: - target_variable = target.get("variable") - if target_variable == variable: - return True - return False - - -def _can_merge_policyengine_person_hierarchy( - tax_units: pd.DataFrame, - persons: pd.DataFrame, -) -> bool: - """Return whether person-derived PE tax units can join to pipeline rows.""" - if {"household_id", "tax_unit_id"}.issubset(tax_units.columns): - return _has_any_column(persons, ["household_id", "ph_seq", "PH_SEQ"]) and ( - _has_any_column(persons, ["tax_unit_id", "tax_id", "TAX_ID"]) - ) - if "tax_unit_id" in tax_units.columns: - return _has_any_column(persons, ["tax_unit_id", "tax_id", "TAX_ID"]) - return False - - -def _has_any_column(df: pd.DataFrame, columns: list[str]) -> bool: - return any(column in df.columns for column in columns) - - -def _merge_policyengine_tax_units( - tax_units: pd.DataFrame, - policyengine_tax_units: pd.DataFrame, -) -> pd.DataFrame: - """Merge PE tax results from person-derived tax units into pipeline rows.""" - output_columns = ["income_tax_liability", "income_tax_liability_source"] - if not all(column in policyengine_tax_units.columns for column in output_columns): - return tax_units - - if { - "household_id", - "tax_unit_id", - }.issubset(tax_units.columns) and { - "household_id", - "tax_unit_id", - }.issubset(policyengine_tax_units.columns): - keys = ["household_id", "tax_unit_id"] - elif ( - "tax_unit_id" in tax_units.columns - and "tax_unit_id" in policyengine_tax_units.columns - ): - keys = ["tax_unit_id"] - else: - return tax_units - - values = policyengine_tax_units[keys + output_columns].drop_duplicates(keys) - merged = tax_units.merge(values, on=keys, how="left") - if merged["income_tax_liability"].notna().any(): - return merged - return tax_units - - -def _sum_columns(df: pd.DataFrame, columns: list[str]) -> float: - total = 0.0 - for column in columns: - if column in df.columns: - total += float(pd.to_numeric(df[column], errors="coerce").fillna(0).sum()) - return total - - -def _base_target_diagnostic( - target: TargetSpec, - target_index: int, -) -> dict[str, Any]: - aging = _parse_aging_note(target.stratum_name) - return { - "target_index": target_index, - "source": _enum_value(target.source), - "variable": target.variable, - "target_type": _enum_value(target.target_type), - "period": target.period, - "source_period": aging.get("source_year", target.period), - "stratum": target.stratum_name, - "constraints": json.dumps(target.constraints, separators=(",", ":")), - "target_value": target.value, - "role": "", - "status": "", - "drop_reason": "", - "n_obs": 0, - "aging_method": aging.get("method"), - "aging_factor": aging.get("factor"), - "pre_value": np.nan, - "pre_error": np.nan, - "post_value": np.nan, - "post_error": np.nan, - } - - -def _base_dict_target_diagnostic( - target: dict[str, Any], - target_index: int, -) -> dict[str, Any]: - stratum = target.get("strata", {}) - constraints = stratum.get("stratum_constraints", []) - return { - "target_index": target_index, - "source": target.get("source"), - "variable": target.get("variable"), - "target_type": target.get("target_type"), - "period": target.get("period"), - "source_period": target.get("period"), - "stratum": stratum.get("name", "unknown"), - "constraints": json.dumps(constraints, separators=(",", ":"), default=str), - "target_value": target.get("value"), - "role": "", - "status": "", - "drop_reason": "", - "n_obs": 0, - "aging_method": None, - "aging_factor": None, - "pre_value": np.nan, - "pre_error": np.nan, - "post_value": np.nan, - "post_error": np.nan, - } - - -def _parse_aging_note(stratum_name: str | None) -> dict[str, Any]: - if not stratum_name: - return {} - match = AGING_NOTE_RE.search(stratum_name) - if not match: - return {} - factor = match.group("factor") - return { - "source_year": int(match.group("source_year")), - "target_year": int(match.group("target_year")), - "method": match.group("method"), - "factor": float(factor) if factor is not None else np.nan, - } - - -def _enum_value(value: Any) -> Any: - return value.value if hasattr(value, "value") else value - - -def _mark_unsupported( - diagnostic: dict[str, Any], - drop_reason: str, -) -> dict[str, Any]: - diagnostic.update( - role="unsupported", - status="dropped", - drop_reason=drop_reason, - ) - return diagnostic - - -def _mark_diagnostic( - diagnostic: dict[str, Any], - drop_reason: str, -) -> dict[str, Any]: - diagnostic.update( - role="diagnostic", - status="dropped", - drop_reason=drop_reason, - ) - return diagnostic - - -def _mark_holdout(diagnostic: dict[str, Any]) -> dict[str, Any]: - diagnostic.update(role="holdout", status="dropped", drop_reason="holdout") - return diagnostic - - -def _mark_active(diagnostic: dict[str, Any]) -> dict[str, Any]: - diagnostic.update(role="active", status="used", drop_reason="") - return diagnostic - - -def _numeric_first(df: pd.DataFrame, columns: list[str]) -> pd.Series: - """Return the first available numeric column, or zeros if none exist.""" - for column in columns: - if column in df.columns: - return pd.to_numeric(df[column], errors="coerce").fillna(0) - return pd.Series(0.0, index=df.index) - - -def assign_agi_bracket(agi: np.ndarray) -> np.ndarray: - """Assign each record to an AGI bracket matching SOI data.""" - brackets = [ - ("under_1", -np.inf, 1), - ("1_to_5k", 1, 5000), - ("5k_to_10k", 5000, 10000), - ("10k_to_15k", 10000, 15000), - ("15k_to_20k", 15000, 20000), - ("20k_to_25k", 20000, 25000), - ("25k_to_30k", 25000, 30000), - ("30k_to_40k", 30000, 40000), - ("40k_to_50k", 40000, 50000), - ("50k_to_75k", 50000, 75000), - ("75k_to_100k", 75000, 100000), - ("100k_to_200k", 100000, 200000), - ("200k_to_500k", 200000, 500000), - ("500k_to_1m", 500000, 1000000), - ("1m_to_1_5m", 1000000, 1500000), - ("1_5m_to_2m", 1500000, 2000000), - ("2m_to_5m", 2000000, 5000000), - ("5m_to_10m", 5000000, 10000000), - ("10m_plus", 10000000, np.inf), - ] - - result = np.empty(len(agi), dtype=object) - for name, low, high in brackets: - mask = (agi >= low) & (agi < high) - result[mask] = name - - return result - - -def _build_microplex_constraint_dict( - df: pd.DataFrame, - *, - variable: str, - target_type: TargetType | str, - target_value: float, - stratum_name: str | None, - stratum_constraints: list[tuple[str, str, str]], -) -> dict[str, Any] | None: - mask = _stratum_mask(df, stratum_constraints) - if mask is None: - return None - - target_type_value = _enum_value(target_type) - if variable == "tax_unit_count": - indicator = mask.astype(float).to_numpy() - else: - column = TARGET_VARIABLE_COLUMNS.get(variable) - if column is None or column not in df.columns: - return None - values = pd.to_numeric(df[column], errors="coerce").fillna(0).to_numpy() - if target_type_value == TargetType.COUNT.value: - indicator = (mask.to_numpy() & (values > 0)).astype(float) - elif target_type_value == TargetType.AMOUNT.value: - indicator = mask.astype(float).to_numpy() * values - else: - return None - - return { - "indicator": indicator, - "target_value": target_value, - "variable": variable, - "target_type": target_type_value, - "stratum": stratum_name, - "n_obs": int(np.count_nonzero(indicator)), - } - - -def _stratum_mask( - df: pd.DataFrame, - stratum_constraints: list[tuple[str, str, str]], -) -> pd.Series | None: - mask = pd.Series(True, index=df.index) - for variable, operator, value in stratum_constraints: - if variable == "is_tax_filer" and variable not in df.columns: - col = pd.Series([1] * len(df), index=df.index) - elif variable in df.columns: - col = df[variable] - else: - return None - - parsed_value: float | int | str = value - if pd.api.types.is_numeric_dtype(col): - parsed_value = float(value) - - if operator == "==": - mask &= col == parsed_value - elif operator == "!=": - mask &= col != parsed_value - elif operator == ">": - mask &= col > parsed_value - elif operator == ">=": - mask &= col >= parsed_value - elif operator == "<": - mask &= col < parsed_value - elif operator == "<=": - mask &= col <= parsed_value - else: - return None - return mask - - -def build_constraints_from_target_specs( - df: pd.DataFrame, - targets: List[TargetSpec], - min_obs: int = 100, - include_amounts: bool = False, - holdout_variables: tuple[str, ...] = (), -) -> List[Dict]: - """ - Build legacy IPF constraint dicts from Arch DB ``TargetSpec`` objects. - - This keeps the old IPF pipeline working while moving the target source from - Supabase-shaped dictionaries to the local Arch target database. - """ - constraints, _, _ = build_constraints_and_diagnostics_from_target_specs( - df, - targets, - min_obs=min_obs, - include_amounts=include_amounts, - holdout_variables=holdout_variables, - ) - return constraints - - -def build_constraints_and_diagnostics_from_target_specs( - df: pd.DataFrame, - targets: List[TargetSpec], - min_obs: int = 100, - include_amounts: bool = False, - holdout_variables: tuple[str, ...] = (), -) -> tuple[List[Dict], list[dict[str, Any]], list[dict[str, Any]]]: - """Build IPF constraints and one diagnostic row per loaded target.""" - df = df.copy() - df["agi_bracket"] = assign_agi_bracket(df["adjusted_gross_income"].values) - - constraints = [] - evaluation_constraints = [] - diagnostics = [] - seen_keys = set() - holdout_variable_set = set(holdout_variables) - for target_index, target in enumerate(targets): - diagnostic = _base_target_diagnostic(target, target_index) - - if target.variable not in SUPPORTED_TARGET_VARIABLES: - _mark_unsupported(diagnostic, "unsupported_variable") - diagnostics.append(diagnostic) - continue - if target.target_type == TargetType.RATE: - _mark_unsupported(diagnostic, "rate_target") - diagnostics.append(diagnostic) - continue - if target.target_type == TargetType.AMOUNT and target.value <= 0: - _mark_unsupported(diagnostic, "nonpositive_amount_target") - diagnostics.append(diagnostic) - continue - if any( - variable not in SUPPORTED_CONSTRAINT_VARIABLES - for variable, _, _ in target.constraints - ): - _mark_unsupported(diagnostic, "unsupported_constraint") - diagnostics.append(diagnostic) - continue - - constraint_dict = _build_microplex_constraint_dict( - df, - variable=target.variable, - target_type=target.target_type, - target_value=target.value, - stratum_name=target.stratum_name, - stratum_constraints=target.constraints, - ) - if constraint_dict is None: - _mark_unsupported(diagnostic, "could_not_build_constraint") - diagnostics.append(diagnostic) - continue - - diagnostic["n_obs"] = constraint_dict["n_obs"] - key = (target.stratum_name, target.variable, target.target_type) - drop_reason = "" - is_holdout = target.variable in holdout_variable_set - if is_holdout: - drop_reason = "holdout" - elif key in seen_keys: - drop_reason = "duplicate_target" - elif target.target_type == TargetType.AMOUNT and not include_amounts: - drop_reason = "amount_targets_disabled" - elif constraint_dict["n_obs"] < min_obs: - drop_reason = "insufficient_obs" - - if drop_reason: - if is_holdout: - _mark_holdout(diagnostic) - else: - _mark_diagnostic(diagnostic, drop_reason) - diagnostic_index = len(diagnostics) - constraint_dict["diagnostic_index"] = diagnostic_index - diagnostics.append(diagnostic) - evaluation_constraints.append(constraint_dict) - continue - - seen_keys.add(key) - _mark_active(diagnostic) - diagnostic_index = len(diagnostics) - constraint_dict["diagnostic_index"] = diagnostic_index - diagnostics.append(diagnostic) - constraints.append(constraint_dict) - - print(f" Built {len(constraints)} constraints (min {min_obs} obs each)") - return constraints, diagnostics, evaluation_constraints - - -def build_constraints_from_targets( - df: pd.DataFrame, - targets: List[Dict[str, Any]] | List[TargetSpec], - min_obs: int = 100, - include_amounts: bool = False, - holdout_variables: tuple[str, ...] = (), - return_diagnostics: bool = False, -) -> List[Dict]: - """ - Build calibration constraints from Supabase targets. - - Currently supports: - - Tax unit counts by AGI bracket (adjusted_gross_income ranges) - - AGI totals by bracket (if include_amounts=True) - - Wage/employment income counts and totals (if include_amounts=True for totals) - - Skips unsupported targets: - - Filing status (need CPS marital status mapping) - - Program participation (SNAP, SSI, OASDI - need program vars) - - Population counts (different universe than tax filers) - """ - if targets and isinstance(targets[0], TargetSpec): - constraints, diagnostics, evaluation_constraints = ( - build_constraints_and_diagnostics_from_target_specs( - df, - targets, - min_obs=min_obs, - include_amounts=include_amounts, - holdout_variables=holdout_variables, - ) - ) - if return_diagnostics: - return constraints, diagnostics, evaluation_constraints - return constraints - - constraints = [] - evaluation_constraints = [] - diagnostics = [] - seen_keys = set() - holdout_variable_set = set(holdout_variables) - df = df.copy() - - # Precompute AGI brackets - df["agi_bracket"] = assign_agi_bracket(df["adjusted_gross_income"].values) - - for target_index, target in enumerate(targets): - diagnostic = _base_dict_target_diagnostic(target, target_index) - variable = target["variable"] - value = target["value"] - target_type = target.get("target_type") - - # Only calibrate on supported variables - if variable not in SUPPORTED_TARGET_VARIABLES: - _mark_unsupported(diagnostic, "unsupported_variable") - diagnostics.append(diagnostic) - continue - - # Skip rate targets; optionally skip amount targets - if target_type == "rate": - _mark_unsupported(diagnostic, "rate_target") - diagnostics.append(diagnostic) - continue - if target_type == "amount" and value <= 0: - _mark_unsupported(diagnostic, "nonpositive_amount_target") - diagnostics.append(diagnostic) - continue - - stratum = target.get("strata", {}) - stratum_name = stratum.get("name", "unknown") - stratum_constraints = stratum.get("stratum_constraints", []) - - # Skip strata with unsupported constraint types - # (filing_status, snap, ssi, oasdi, etc.) - has_unsupported = False - for c in stratum_constraints: - var = c.get("variable") - if var not in SUPPORTED_CONSTRAINT_VARIABLES: - has_unsupported = True - break - if has_unsupported: - _mark_unsupported(diagnostic, "unsupported_constraint") - diagnostics.append(diagnostic) - continue - - tuple_constraints = [ - ( - constraint.get("variable"), - constraint.get("operator", "=="), - constraint.get("value"), - ) - for constraint in stratum_constraints - ] - constraint_dict = _build_microplex_constraint_dict( - df, - variable=variable, - target_type=target_type, - target_value=value, - stratum_name=stratum_name, - stratum_constraints=tuple_constraints, - ) - if constraint_dict is None: - _mark_unsupported(diagnostic, "could_not_build_constraint") - diagnostics.append(diagnostic) - continue - - n_obs = constraint_dict["n_obs"] - diagnostic["n_obs"] = int(n_obs) - key = (stratum_name, variable, target_type) - drop_reason = "" - is_holdout = variable in holdout_variable_set - if is_holdout: - drop_reason = "holdout" - elif key in seen_keys: - drop_reason = "duplicate_target" - elif target_type == "amount" and not include_amounts: - drop_reason = "amount_targets_disabled" - elif n_obs < min_obs: - drop_reason = "insufficient_obs" - - if drop_reason: - if is_holdout: - _mark_holdout(diagnostic) - else: - _mark_diagnostic(diagnostic, drop_reason) - diagnostic_index = len(diagnostics) - constraint_dict["diagnostic_index"] = diagnostic_index - diagnostics.append(diagnostic) - evaluation_constraints.append(constraint_dict) - continue - - seen_keys.add(key) - _mark_active(diagnostic) - diagnostic_index = len(diagnostics) - constraint_dict["diagnostic_index"] = diagnostic_index - constraints.append(constraint_dict) - diagnostics.append(diagnostic) - - print(f" Built {len(constraints)} constraints (min {min_obs} obs each)") - if return_diagnostics: - return constraints, diagnostics, evaluation_constraints - return constraints - - -def ipf_calibrate( - original_weights: np.ndarray, - constraints: List[Dict], - bounds: tuple = (0.2, 5.0), - max_iter: int = 100, - damping: tuple = (0.9, 1.1), - verbose: bool = True, -) -> tuple: - """ - Calibrate weights using Iterative Proportional Fitting (IPF). - - IPF iteratively adjusts weights to match marginal totals. - This is 486x faster than IPF+GREG with only 7% worse L2 loss. - - Args: - original_weights: Initial survey weights - constraints: List of constraint dicts with 'indicator' and 'target_value' - bounds: Min/max weight adjustment factors (default 0.2-5.0) - max_iter: Number of IPF iterations (default 100) - damping: Min/max adjustment ratio per iteration (default 0.9-1.1) - verbose: Print progress - - Returns: - (calibrated_weights, success, l2_loss) - """ - n = len(original_weights) - m = len(constraints) - - if verbose: - print(f"IPF calibration: {n:,} weights, {m} constraints, {max_iter} iterations") - - # Build constraint matrix - A = np.zeros((m, n)) - targets = np.zeros(m) - - for j, c in enumerate(constraints): - A[j, :] = c["indicator"] - targets[j] = c["target_value"] - - w = original_weights.copy() - - for iteration in range(max_iter): - for j in range(m): - achieved = A[j] @ w - if achieved > 0: - # Damped ratio to ensure convergence - ratio = np.clip(targets[j] / achieved, damping[0], damping[1]) - mask = A[j] != 0 - w[mask] *= ratio - - # Apply bounds after each full iteration - adj = w / original_weights - adj = np.clip(adj, bounds[0], bounds[1]) - w = original_weights * adj - - # Compute L2 loss (squared relative error) - achieved = A @ w - l2_loss = np.mean(((achieved - targets) / targets) ** 2) - - # Check convergence (all targets within 5%) - max_error = np.max(np.abs((achieved - targets) / targets)) - success = max_error < 0.05 - - if verbose: - print(f"IPF converged: max error = {max_error:.1%}, L2 loss = {l2_loss:.6f}") - - return w, success, l2_loss - - -def generalized_rake_calibrate( - original_weights: np.ndarray, - constraints: List[Dict], - bounds: tuple[float, float] = (0.1, 20.0), - max_iter: int = 80, - target_tolerance: float = 0.05, - verbose: bool = True, -) -> tuple[np.ndarray, bool, float]: - """ - Calibrate weights with bounded generalized raking. - - This solves the entropy/raking problem in the dual, where the number of - variables is the number of calibration constraints rather than the number - of microdata rows. It supports overlapping count and amount constraints - such as SOI filer counts and AGI totals by bracket. - - Args: - original_weights: Initial survey weights. - constraints: Constraint dicts with ``indicator`` and ``target_value``. - bounds: Min/max adjustment factors around ``original_weights``. - max_iter: Maximum Newton iterations. - target_tolerance: Success threshold for max relative target error. - verbose: Print progress. - - Returns: - (calibrated_weights, success, l2_loss) - """ - lower, upper = bounds - if not 0 < lower < 1 < upper: - raise ValueError("generalized raking bounds must satisfy 0 < lower < 1 < upper") - - n = len(original_weights) - m = len(constraints) - if verbose: - print( - "Generalized raking calibration: " - f"{n:,} weights, {m} constraints, bounds=[{lower}, {upper}]" - ) - - A = np.vstack([np.asarray(c["indicator"], dtype=float) for c in constraints]) - targets = np.asarray([c["target_value"] for c in constraints], dtype=float) - if np.any(~np.isfinite(A)) or np.any(~np.isfinite(targets)): - raise ValueError("Calibration constraints contain non-finite values") - if np.any(targets == 0): - raise ValueError("Generalized raking requires nonzero target values") - - current = A @ original_weights - if np.any(current == 0): - zero_constraints = [ - _constraint_key(c) for c, value in zip(constraints, current) if value == 0 - ] - raise ValueError( - "Cannot calibrate constraints with zero current support: " - + ", ".join(zero_constraints) - ) - - # Row scaling keeps count and amount constraints on comparable numerical - # footing without changing the equations being solved. - scale = np.abs(current) - B = A / scale[:, None] - scaled_targets = targets / scale - - offset = np.log((1 - lower) / (upper - 1)) - dual = np.zeros(m) - best_weights = original_weights.copy() - best_error = np.inf - - for iteration in range(max_iter): - adjustment, derivative = _bounded_rake_adjustment( - B.T @ dual, - offset=offset, - lower=lower, - upper=upper, - ) - weights = original_weights * adjustment - residual = B @ weights - scaled_targets - relative_errors = (A @ weights - targets) / targets - max_error = float(np.max(np.abs(relative_errors))) - if max_error < best_error: - best_error = max_error - best_weights = weights.copy() - - if verbose and (iteration == 0 or (iteration + 1) % 10 == 0): - print( - f" iter {iteration + 1}: max error={max_error:.2%}, " - f"adjustment range=[{adjustment.min():.2f}, {adjustment.max():.2f}]" - ) - - jacobian = (B * (original_weights * derivative)) @ B.T - ridge = max(float(np.trace(jacobian)) / max(m, 1) * 1e-9, 1e-12) - step = np.linalg.lstsq( - jacobian + np.eye(m) * ridge, - residual, - rcond=None, - )[0] - - residual_norm = float(np.linalg.norm(residual)) - step_accepted = False - small_improvement = False - step_size = 1.0 - for _ in range(30): - candidate_dual = dual - step_size * step - candidate_adjustment, _ = _bounded_rake_adjustment( - B.T @ candidate_dual, - offset=offset, - lower=lower, - upper=upper, - ) - candidate_residual = ( - B @ (original_weights * candidate_adjustment) - scaled_targets - ) - candidate_norm = float(np.linalg.norm(candidate_residual)) - if ( - np.all(np.isfinite(candidate_residual)) - and candidate_norm < residual_norm - ): - dual = candidate_dual - step_accepted = True - small_improvement = (residual_norm - candidate_norm) <= max( - 1e-12, - residual_norm * 1e-8, - ) - break - step_size *= 0.5 - - if not step_accepted or small_improvement: - break - - relative_errors = (A @ best_weights - targets) / targets - l2_loss = float(np.mean(relative_errors**2)) - max_error = float(np.max(np.abs(relative_errors))) - success = max_error < target_tolerance - - if verbose: - print( - f"Generalized raking complete: max error = {max_error:.1%}, " - f"L2 loss = {l2_loss:.6f}" - ) - - return best_weights, success, l2_loss - - -def _bounded_rake_adjustment( - linear_predictor: np.ndarray, - *, - offset: float, - lower: float, - upper: float, -) -> tuple[np.ndarray, np.ndarray]: - """Return bounded calibration factors and derivatives.""" - z = np.clip(offset + linear_predictor, -60, 60) - probability = 1 / (1 + np.exp(-z)) - adjustment = lower + (upper - lower) * probability - derivative = (adjustment - lower) * (upper - adjustment) / (upper - lower) - return adjustment, derivative - - -def calibrate_weights( - df: pd.DataFrame, - targets: List[Dict[str, Any]] | List[TargetSpec], - include_amounts: bool = True, - min_obs: int = 100, - calibration_method: str = "auto", - weight_bounds: tuple[float, float] = (0.1, 20.0), - holdout_variables: tuple[str, ...] = (), - verbose: bool = True, -) -> CalibrationResult: - """Calibrate Microplex weights to Arch target inputs.""" - original_weights = df["weight"].values.copy() - - if verbose: - print(f"\nCalibrating {len(df):,} tax units...") - print(f"Original weighted total: {original_weights.sum():,.0f}") - - constraints, diagnostics, evaluation_constraints = build_constraints_from_targets( - df, - targets, - min_obs=min_obs, - include_amounts=include_amounts, - holdout_variables=holdout_variables, - return_diagnostics=True, - ) - diagnostics_df = pd.DataFrame(diagnostics) - if not constraints: - raise ValueError("No usable calibration constraints were built.") - - # Pre-scale weights to match total population target - total_target = None - for c in constraints: - if c["variable"] == "tax_unit_count" and c["n_obs"] == len(df): - total_target = c["target_value"] - break - - if total_target: - scale_factor = total_target / original_weights.sum() - original_weights = original_weights * scale_factor - if verbose: - print(f"Pre-scaled weights by {scale_factor:.3f} to match total target {total_target:,.0f}") - - # Pre-calibration values - targets_before = {} - for c in constraints: - current = np.dot(original_weights, c["indicator"]) - error = ( - (current - c["target_value"]) / c["target_value"] - if c["target_value"] != 0 - else 0 - ) - diagnostic_index = c.get("diagnostic_index") - if diagnostic_index is not None: - diagnostics_df.loc[diagnostic_index, "pre_value"] = current - diagnostics_df.loc[diagnostic_index, "pre_error"] = error - - key = _constraint_key(c) - targets_before[key] = { - "current": current, - "target": c["target_value"], - "error": error, - } - for c in evaluation_constraints: - _update_diagnostic_values( - diagnostics_df, - c, - weights=original_weights, - value_column="pre_value", - error_column="pre_error", - ) - - has_amount_constraints = any(c.get("target_type") == "amount" for c in constraints) - method = calibration_method - if method == "auto": - method = "generalized-rake" if has_amount_constraints else "ipf" - - if method == "ipf": - calibrated_weights, success, l2_loss = ipf_calibrate( - original_weights, - constraints, - verbose=verbose, - ) - elif method == "generalized-rake": - calibrated_weights, success, l2_loss = generalized_rake_calibrate( - original_weights, - constraints, - bounds=weight_bounds, - verbose=verbose, - ) - else: - raise ValueError(f"Unknown calibration method: {calibration_method}") - - # Post-calibration values - targets_after = {} - max_error = 0 - for c in constraints: - current = np.dot(calibrated_weights, c["indicator"]) - error = (current - c["target_value"]) / c["target_value"] if c["target_value"] != 0 else 0 - diagnostic_index = c.get("diagnostic_index") - if diagnostic_index is not None: - diagnostics_df.loc[diagnostic_index, "post_value"] = current - diagnostics_df.loc[diagnostic_index, "post_error"] = error - - key = _constraint_key(c) - targets_after[key] = { - "current": current, - "target": c["target_value"], - "error": error, - } - max_error = max(max_error, abs(error)) - for c in evaluation_constraints: - _update_diagnostic_values( - diagnostics_df, - c, - weights=calibrated_weights, - value_column="post_value", - error_column="post_error", - ) - - adjustment_factors = calibrated_weights / original_weights - - if verbose: - print(f"\nPost-calibration max error: {max_error:.1%}") - print(f"Weight adjustments: mean={adjustment_factors.mean():.2f}, " - f"range=[{adjustment_factors.min():.2f}, {adjustment_factors.max():.2f}]") - print(f"L2 loss: {l2_loss:.6f}") - - return CalibrationResult( - original_weights=original_weights, - calibrated_weights=calibrated_weights, - adjustment_factors=adjustment_factors, - targets_before=targets_before, - targets_after=targets_after, - diagnostics=diagnostics_df, - success=success, - message="Converged" if success else "Did not converge", - l2_loss=l2_loss, - method=method, - calibration_unit="tax_unit", - ) - - -def calibrate_household_weights( - households: pd.DataFrame, - tax_units: pd.DataFrame, - targets: List[Dict[str, Any]] | List[TargetSpec], - include_amounts: bool = True, - min_obs: int = 100, - calibration_method: str = "auto", - weight_bounds: tuple[float, float] = (0.1, 20.0), - holdout_variables: tuple[str, ...] = (), - verbose: bool = True, -) -> CalibrationResult: - """Calibrate household weights against tax-unit targets.""" - original_weights = households["weight"].values.copy() - - if verbose: - print(f"\nCalibrating {len(households):,} households...") - print(f"Original weighted households: {original_weights.sum():,.0f}") - - constraints, diagnostics, evaluation_constraints = ( - build_household_constraints_from_targets( - households, - tax_units, - targets, - min_obs=min_obs, - include_amounts=include_amounts, - holdout_variables=holdout_variables, - ) - ) - diagnostics_df = pd.DataFrame(diagnostics) - if not constraints: - raise ValueError("No usable household calibration constraints were built.") - - total_constraint = _household_prescale_constraint(constraints) - if total_constraint is not None: - current_total = float(np.dot(original_weights, total_constraint["indicator"])) - if current_total > 0: - scale_factor = total_constraint["target_value"] / current_total - original_weights = original_weights * scale_factor - if verbose: - print( - f"Pre-scaled household weights by {scale_factor:.3f} " - f"to match {total_constraint['stratum']} " - f"target {total_constraint['target_value']:,.0f}" - ) - - targets_before = {} - for constraint in constraints: - current = float(np.dot(original_weights, constraint["indicator"])) - error = ( - (current - constraint["target_value"]) / constraint["target_value"] - if constraint["target_value"] != 0 - else 0 - ) - diagnostic_index = constraint.get("diagnostic_index") - if diagnostic_index is not None: - diagnostics_df.loc[diagnostic_index, "pre_value"] = current - diagnostics_df.loc[diagnostic_index, "pre_error"] = error - - targets_before[_constraint_key(constraint)] = { - "current": current, - "target": constraint["target_value"], - "error": error, - } - for constraint in evaluation_constraints: - _update_diagnostic_values( - diagnostics_df, - constraint, - weights=original_weights, - value_column="pre_value", - error_column="pre_error", - ) - - has_amount_constraints = any( - constraint.get("target_type") == "amount" for constraint in constraints - ) - method = calibration_method - if method == "auto": - method = "generalized-rake" if has_amount_constraints else "ipf" - - if method == "ipf": - calibrated_weights, success, l2_loss = ipf_calibrate( - original_weights, - constraints, - verbose=verbose, - ) - elif method == "generalized-rake": - calibrated_weights, success, l2_loss = generalized_rake_calibrate( - original_weights, - constraints, - bounds=weight_bounds, - verbose=verbose, - ) - else: - raise ValueError(f"Unknown calibration method: {calibration_method}") - - targets_after = {} - max_error = 0.0 - for constraint in constraints: - current = float(np.dot(calibrated_weights, constraint["indicator"])) - error = ( - (current - constraint["target_value"]) / constraint["target_value"] - if constraint["target_value"] != 0 - else 0 - ) - diagnostic_index = constraint.get("diagnostic_index") - if diagnostic_index is not None: - diagnostics_df.loc[diagnostic_index, "post_value"] = current - diagnostics_df.loc[diagnostic_index, "post_error"] = error - - targets_after[_constraint_key(constraint)] = { - "current": current, - "target": constraint["target_value"], - "error": error, - } - max_error = max(max_error, abs(error)) - for constraint in evaluation_constraints: - _update_diagnostic_values( - diagnostics_df, - constraint, - weights=calibrated_weights, - value_column="post_value", - error_column="post_error", - ) - - adjustment_factors = calibrated_weights / original_weights - - if verbose: - print(f"\nPost-calibration max error: {max_error:.1%}") - print( - "Household weight adjustments: " - f"mean={adjustment_factors.mean():.2f}, " - f"range=[{adjustment_factors.min():.2f}, " - f"{adjustment_factors.max():.2f}]" - ) - print(f"L2 loss: {l2_loss:.6f}") - - return CalibrationResult( - original_weights=original_weights, - calibrated_weights=calibrated_weights, - adjustment_factors=adjustment_factors, - targets_before=targets_before, - targets_after=targets_after, - diagnostics=diagnostics_df, - success=success, - message="Converged" if success else "Did not converge", - l2_loss=l2_loss, - method=method, - calibration_unit="household", - ) - - -def _household_prescale_constraint(constraints: List[Dict]) -> dict[str, Any] | None: - """Pick the broad tax-unit count constraint used to scale sample weights.""" - count_constraints = [ - constraint - for constraint in constraints - if constraint.get("variable") == "tax_unit_count" - and constraint.get("target_type") == "count" - ] - if not count_constraints: - return None - all_filers = [ - constraint - for constraint in count_constraints - if "all filers" in str(constraint.get("stratum", "")).lower() - ] - if all_filers: - return max(all_filers, key=lambda constraint: constraint.get("n_obs", 0)) - return max(count_constraints, key=lambda constraint: constraint.get("n_obs", 0)) - - -def build_household_constraints_from_targets( - households: pd.DataFrame, - tax_units: pd.DataFrame, - targets: List[Dict[str, Any]] | List[TargetSpec], - min_obs: int = 100, - include_amounts: bool = False, - holdout_variables: tuple[str, ...] = (), -) -> tuple[List[Dict], list[dict[str, Any]], list[dict[str, Any]]]: - """Build household-level constraints from tax-unit target indicators.""" - constraints, diagnostics, evaluation_constraints = build_constraints_from_targets( - tax_units, - targets, - min_obs=min_obs, - include_amounts=include_amounts, - holdout_variables=holdout_variables, - return_diagnostics=True, - ) - return ( - _aggregate_tax_unit_constraints_to_households( - households, - tax_units, - constraints, - ), - diagnostics, - _aggregate_tax_unit_constraints_to_households( - households, - tax_units, - evaluation_constraints, - ), - ) - - -def _aggregate_tax_unit_constraints_to_households( - households: pd.DataFrame, - tax_units: pd.DataFrame, - constraints: List[Dict], -) -> List[Dict]: - """Convert tax-unit indicator vectors to household-level indicators.""" - aggregated_constraints = [] - household_ids = households["household_entity_id"] - tax_unit_households = tax_units["household_entity_id"] - for constraint in constraints: - indicator = pd.Series(constraint["indicator"], index=tax_units.index) - grouped = indicator.groupby(tax_unit_households).sum() - household_constraint = constraint.copy() - household_constraint["indicator"] = ( - household_ids.map(grouped).fillna(0).to_numpy(dtype=float) - ) - aggregated_constraints.append(household_constraint) - return aggregated_constraints - - -def _update_diagnostic_values( - diagnostics_df: pd.DataFrame, - constraint: dict[str, Any], - *, - weights: np.ndarray, - value_column: str, - error_column: str, -) -> None: - diagnostic_index = constraint.get("diagnostic_index") - if diagnostic_index is None: - return - current = float(np.dot(weights, constraint["indicator"])) - target = constraint["target_value"] - error = (current - target) / target if target != 0 else 0 - diagnostics_df.loc[diagnostic_index, value_column] = current - diagnostics_df.loc[diagnostic_index, error_column] = error - - -def _constraint_key(constraint: dict[str, Any]) -> str: - return ( - f"{constraint.get('variable')}|{constraint.get('target_type')}|" - f"{constraint.get('stratum')}" - ) - - -def write_microplex_to_supabase( - df: pd.DataFrame, - year: int, - chunk_size: int = 200, -) -> int: - """Write calibrated microplex to Supabase.""" - client = get_supabase_client() - table_name = f"us_microplex_{year}_person" - - print(f"Writing {len(df):,} records to {table_name}...") - - records = [] - for _, row in df.iterrows(): - record = { - "source_person_id": int(row.get("id", 0)) if pd.notna(row.get("id")) else None, - "household_id": int(row.get("ph_seq", 0)) if pd.notna(row.get("ph_seq")) else None, - "age": int(row.get("a_age", 0)) if pd.notna(row.get("a_age")) else None, - "state_fips": int(row.get("gestfips", 0)) if pd.notna(row.get("gestfips")) else None, - "wage_income": float(row.get("wage_income", 0)), - "self_employment_income": float(row.get("self_employment_income", 0)), - "total_income": float(row.get("total_income", 0)), - "adjusted_gross_income": float(row.get("adjusted_gross_income", 0)) if "adjusted_gross_income" in row else None, - "income_tax_liability": float(row.get("income_tax_liability", 0)) - if "income_tax_liability" in row - else None, - "original_weight": float(row.get("original_weight", 0)), - "calibrated_weight": float(row.get("weight", 0)), - "weight_adjustment": float(row.get("weight_adjustment", 1.0)), - "agi_bracket": row.get("agi_bracket"), - } - records.append(record) - - total = 0 - for i in range(0, len(records), chunk_size): - chunk = records[i : i + chunk_size] - client.schema("microplex").table(table_name).insert(chunk).execute() - total += len(chunk) - if (i + chunk_size) % 5000 == 0: - print(f" Inserted {total:,} / {len(records):,}") - - print(f" Done: {total:,} records") - return total - - -def print_target_composition_diagnostics( - composition: TargetCompositionResult, -) -> None: - """Print how Arch target inputs became Microplex target candidates.""" - diagnostics = composition.diagnostics - print("\nTarget input composition:") - print(f" Candidate targets: {len(composition.targets):,}") - if composition.fallback_reason is not None: - print( - " Fallback: " - f"{composition.fallback_reason}; using {composition.fallback_year}" - ) - if composition.soi_aging_factors is not None: - factors = composition.soi_aging_factors - print( - " Aged SOI targets " - f"{factors.source_year}->{factors.target_year}: " - f"counts x{factors.count_factor:.4f} " - f"({factors.count_method}), " - f"amounts x{factors.amount_factor:.4f} " - f"({factors.amount_method})" - ) - if diagnostics.empty: - return - action_counts = diagnostics["action"].value_counts() - action_text = ", ".join( - f"{action}={count}" for action, count in action_counts.items() - ) - print(f" Composition actions: {action_text}") - - -def run_pipeline( - year: int = 2024, - dry_run: bool = False, - limit: int = 200000, - target_source: str = "db", - db_path: Path | None = None, - microdata_source: str = "local", - cps_path: Path | None = None, - output_path: Path | None = None, - entity_output_dir: Path | None = None, - age_soi: bool = True, - include_amount_targets: bool = True, - min_target_obs: int = 100, - calibration_method: str = "auto", - min_weight_factor: float = 0.1, - max_weight_factor: float = 20.0, - diagnostics_path: Path | None = None, -) -> pd.DataFrame: - """Run the full microplex pipeline.""" - print("=" * 60) - print("MICROPLEX PIPELINE") - print("=" * 60) - - # Load data - if microdata_source == "local": - df = load_cps_from_local_file(year, path=cps_path, limit=limit) - elif microdata_source == "supabase": - df = load_cps_from_supabase(year, limit=limit) - else: - raise ValueError(f"Unknown microdata_source: {microdata_source}") - - if target_source == "db": - target_profile = MicroplexTargetProfile(age_soi=age_soi) - target_composition = compose_microplex_targets( - target_year=year, - db_path=db_path, - profile=target_profile, - ) - targets = target_composition.targets - print_target_composition_diagnostics(target_composition) - elif target_source == "supabase": - targets = load_targets_from_supabase(year) - else: - raise ValueError(f"Unknown target_source: {target_source}") - - if target_source == "supabase" and len(targets) < 50: - # Fall back to the latest available SOI targets when the model year - # has insufficient usable tax targets. - fallback_year = 2021 - fallback_reason = f"only {len(targets)} target inputs" - print(f" {year} has {fallback_reason}, trying {fallback_year}...") - targets = load_targets_from_supabase(fallback_year) - - # Build linked entity frames. Households/persons are primitive; tax units - # are an assignment over persons inside households. - entities = build_microplex_entities(df) - entities = replace( - entities, - tax_units=maybe_add_policyengine_income_tax( - entities.tax_units, - targets, - year=year, - persons=entities.persons, - ), - ) - - # Calibrate - result = calibrate_household_weights( - entities.households, - entities.tax_units, - targets, - include_amounts=include_amount_targets, - min_obs=min_target_obs, - calibration_method=calibration_method, - weight_bounds=(min_weight_factor, max_weight_factor), - holdout_variables=target_profile.holdout_variables - if target_source == "db" - else (), - ) - - # Add calibrated household weights and map them to linked entities. - households = entities.households.copy() - households["original_weight"] = result.original_weights - households["weight"] = result.calibrated_weights - households["calibrated_weight"] = result.calibrated_weights - households["weight_adjustment"] = result.adjustment_factors - entities = with_household_weights(entities, households) - df = entities.tax_units.copy() - - # Add AGI bracket for analysis - df["agi_bracket"] = assign_agi_bracket(df["adjusted_gross_income"].values) - - print("\n" + "=" * 60) - print("SUMMARY") - print("=" * 60) - print(f"Households: {len(entities.households):,}") - print(f"Persons: {len(entities.persons):,}") - print(f"Total tax units: {len(df):,}") - print(f"Original weighted: {result.original_weights.sum():,.0f}") - print(f"Calibrated weighted: {result.calibrated_weights.sum():,.0f}") - print(f"Calibration method: {result.method}") - print(f"Calibration unit: {result.calibration_unit}") - print(f"Calibration success: {result.success}") - print_calibration_diagnostics(result.diagnostics) - - if diagnostics_path is not None: - write_calibration_diagnostics(result.diagnostics, diagnostics_path) - - if entity_output_dir is not None: - write_microplex_entities(entities, entity_output_dir) - print(f"\nSaved linked Microplex entities to {entity_output_dir}") - - if output_path is not None: - output_path.parent.mkdir(parents=True, exist_ok=True) - df.to_parquet(output_path, index=False) - print(f"\nSaved microplex to {output_path}") - elif not dry_run: - write_microplex_to_supabase(df, year) - else: - print("\nDRY RUN - not writing to Supabase") - - return df - - -def print_calibration_diagnostics(diagnostics: pd.DataFrame, max_rows: int = 8) -> None: - """Print a compact diagnostic summary for target usage and failures.""" - if diagnostics.empty: - print("No calibration diagnostics available") - return - - used = diagnostics[diagnostics["status"] == "used"].copy() - dropped = diagnostics[diagnostics["status"] == "dropped"].copy() - print("\nCalibration target diagnostics:") - print(f" Used targets: {len(used):,}") - print(f" Dropped targets: {len(dropped):,}") - if "role" in diagnostics: - roles = diagnostics["role"].fillna("unknown").value_counts() - role_text = ", ".join(f"{role}={count}" for role, count in roles.items()) - print(f" Target roles: {role_text}") - if not dropped.empty: - reasons = dropped["drop_reason"].value_counts() - reason_text = ", ".join(f"{reason}={count}" for reason, count in reasons.items()) - print(f" Drop reasons: {reason_text}") - - if used.empty or "post_error" not in used: - return - - used["abs_post_error"] = used["post_error"].abs() - top = used.sort_values("abs_post_error", ascending=False).head(max_rows) - print(" Largest used-target errors:") - for _, row in top.iterrows(): - print( - " " - f"{row['variable']} {row['target_type']} | " - f"{row['stratum']} | " - f"target={row['target_value']:,.0f} " - f"post={row['post_value']:,.0f} " - f"error={row['post_error']:.1%} " - f"n={int(row['n_obs']):,}" - ) - - -def write_calibration_diagnostics(diagnostics: pd.DataFrame, path: Path) -> None: - """Write calibration diagnostics to CSV or Parquet.""" - path.parent.mkdir(parents=True, exist_ok=True) - if path.suffix.lower() == ".parquet": - diagnostics.to_parquet(path, index=False) - else: - diagnostics.to_csv(path, index=False) - print(f"Saved calibration diagnostics to {path}") - - -def main(): - import argparse - - parser = argparse.ArgumentParser(description="Run microplex pipeline") - parser.add_argument("--year", type=int, default=2024, help="Data year") - parser.add_argument("--dry-run", action="store_true", help="Don't write to Supabase") - parser.add_argument("--limit", type=int, default=200000, help="Max records to load") - parser.add_argument( - "--microdata-source", - choices=["local", "supabase"], - default="local", - help="CPS microdata source", - ) - parser.add_argument( - "--target-source", - choices=["db", "supabase"], - default="db", - help="Calibration target source", - ) - parser.add_argument("--db-path", type=Path, default=None, help="Arch SQLite DB path") - parser.add_argument("--cps-path", type=Path, default=None, help="Local CPS parquet path") - parser.add_argument("--output-path", type=Path, default=None, help="Local parquet output") - parser.add_argument( - "--entity-output-dir", - type=Path, - default=None, - help="Optional directory for linked households/persons/tax_units Parquet files", - ) - parser.add_argument( - "--age-soi-targets", - action=argparse.BooleanOptionalAction, - default=True, - help="Age fallback SOI target inputs to the requested model year", - ) - parser.add_argument( - "--include-amount-targets", - action=argparse.BooleanOptionalAction, - default=True, - help="Include amount targets such as AGI totals in calibration", - ) - parser.add_argument( - "--min-target-obs", - type=int, - default=100, - help="Minimum nonzero records required for a target constraint", - ) - parser.add_argument( - "--diagnostics-path", - type=Path, - default=None, - help="Optional CSV or Parquet path for calibration target diagnostics", - ) - parser.add_argument( - "--calibration-method", - choices=["auto", "ipf", "generalized-rake"], - default="auto", - help="Calibration solver to use", - ) - parser.add_argument( - "--min-weight-factor", - type=float, - default=0.1, - help="Minimum generalized-rake weight adjustment factor", - ) - parser.add_argument( - "--max-weight-factor", - type=float, - default=20.0, - help="Maximum generalized-rake weight adjustment factor", - ) - args = parser.parse_args() - - run_pipeline( - year=args.year, - dry_run=args.dry_run, - limit=args.limit, - target_source=args.target_source, - db_path=args.db_path, - microdata_source=args.microdata_source, - cps_path=args.cps_path, - output_path=args.output_path, - entity_output_dir=args.entity_output_dir, - age_soi=args.age_soi_targets, - include_amount_targets=args.include_amount_targets, - min_target_obs=args.min_target_obs, - calibration_method=args.calibration_method, - min_weight_factor=args.min_weight_factor, - max_weight_factor=args.max_weight_factor, - diagnostics_path=args.diagnostics_path, - ) - - -if __name__ == "__main__": - main() diff --git a/micro/us/targets.py b/micro/us/targets.py deleted file mode 100644 index 93888173..00000000 --- a/micro/us/targets.py +++ /dev/null @@ -1,738 +0,0 @@ -"""Microplex adapters for Arch target inputs.""" - -from __future__ import annotations - -from dataclasses import dataclass, replace -from pathlib import Path -from typing import Any - -import numpy as np -import pandas as pd - -from calibration.constraints import ( - Constraint, - build_constraint_matrix, - build_hierarchical_constraint_matrix, -) -from calibration.targets import TargetSpec, get_targets -from db.schema import DataSource, TargetType - - -@dataclass(frozen=True) -class SOIAgingFactors: - """Declared factors used to age SOI target inputs for Microplex.""" - - source_year: int - target_year: int - count_factor: float - amount_factor: float - count_method: str - amount_method: str - - -@dataclass(frozen=True) -class MicroplexTargetProfile: - """Declared rules for composing Arch records into Microplex target inputs.""" - - min_current_target_inputs: int = 50 - tax_variables: tuple[str, ...] = ( - "tax_unit_count", - "adjusted_gross_income", - "income_tax_liability", - ) - holdout_variables: tuple[str, ...] = ("income_tax_liability",) - fallback_source: DataSource = DataSource.IRS_SOI - age_soi: bool = True - - -@dataclass(frozen=True) -class TargetCompositionResult: - """Composed Microplex target inputs plus audit diagnostics.""" - - targets: list[TargetSpec] - diagnostics: pd.DataFrame - fallback_year: int | None = None - fallback_reason: str | None = None - soi_aging_factors: SOIAgingFactors | None = None - - -def load_microplex_targets( - db_path: Path | None = None, - jurisdiction: str = "us", - year: int | None = None, - sources: list[str] | None = None, - variables: list[str] | None = None, -) -> list[TargetSpec]: - """Load Arch DB target inputs as ``TargetSpec`` objects for Microplex.""" - return get_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - sources=sources, - variables=variables, - ) - - -def compose_microplex_targets( - *, - target_year: int, - db_path: Path | None = None, - jurisdiction: str = "us", - profile: MicroplexTargetProfile | None = None, -) -> TargetCompositionResult: - """ - Compose model-year target inputs for Microplex from Arch records. - - This is the boundary where Microplex may choose source fallbacks and apply - model-year transformations. Arch records are not mutated. - """ - profile = profile or MicroplexTargetProfile() - current_targets = load_microplex_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=target_year, - ) - fallback_reason = _target_fallback_reason(current_targets, profile) - - diagnostics: list[dict[str, Any]] = [] - composed_targets: list[TargetSpec] = [] - fallback_year = None - aging_factors = None - - if fallback_reason is None: - composed_targets = current_targets - diagnostics.extend( - _composition_row( - target, - role="current_year", - action="kept_candidate", - reason=None, - target_year=target_year, - source_value=target.value, - source_period=target.period, - ) - for target in current_targets - ) - return TargetCompositionResult( - targets=composed_targets, - diagnostics=pd.DataFrame(diagnostics), - ) - - fallback_year = ( - target_year - if has_supported_tax_targets(current_targets, profile=profile) - else latest_supported_soi_year( - target_year, - db_path=db_path, - jurisdiction=jurisdiction, - profile=profile, - ) - ) - if fallback_year is None: - raise ValueError(f"No usable {profile.fallback_source.value} fallback targets.") - - for target in current_targets: - if target.source == profile.fallback_source: - action = "superseded_by_fallback_source" - composed = False - else: - action = "kept_candidate" - composed = True - composed_targets.append(target) - diagnostics.append( - _composition_row( - target, - role="current_year", - action=action, - reason=fallback_reason, - target_year=target_year, - source_value=target.value, - source_period=target.period, - composed=composed, - ) - ) - - fallback_targets = load_microplex_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=fallback_year, - sources=[profile.fallback_source.value], - ) - if not fallback_targets: - raise ValueError( - f"No {profile.fallback_source.value} targets found for {fallback_year}." - ) - - transformed_targets = fallback_targets - if ( - profile.age_soi - and profile.fallback_source == DataSource.IRS_SOI - and any(target.period != target_year for target in fallback_targets) - ): - aging_factors = get_soi_aging_factors( - source_year=fallback_year, - target_year=target_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - transformed_targets = age_soi_targets( - fallback_targets, - target_year=target_year, - db_path=db_path, - jurisdiction=jurisdiction, - factors=aging_factors, - ) - - for source_target, transformed_target in zip(fallback_targets, transformed_targets): - action = ( - "aged_to_model_year" - if transformed_target.period != source_target.period - else "kept_candidate" - ) - diagnostics.append( - _composition_row( - transformed_target, - role=f"fallback_{profile.fallback_source.value}", - action=action, - reason=fallback_reason, - target_year=target_year, - source_value=source_target.value, - source_period=source_target.period, - composed=True, - ) - ) - composed_targets.extend(transformed_targets) - - return TargetCompositionResult( - targets=composed_targets, - diagnostics=pd.DataFrame(diagnostics), - fallback_year=fallback_year, - fallback_reason=fallback_reason, - soi_aging_factors=aging_factors, - ) - - -def has_supported_tax_targets( - targets: list[TargetSpec], - *, - profile: MicroplexTargetProfile | None = None, -) -> bool: - """Return whether target inputs can produce current tax-unit constraints.""" - profile = profile or MicroplexTargetProfile() - return any( - target.variable in profile.tax_variables and target.target_type != TargetType.RATE - for target in targets - ) - - -def latest_supported_soi_year( - target_year: int, - db_path: Path | None = None, - jurisdiction: str = "us", - *, - profile: MicroplexTargetProfile | None = None, -) -> int | None: - """Find the latest SOI year at or before the model year with usable targets.""" - profile = profile or MicroplexTargetProfile() - for candidate_year in range(target_year, 1989, -1): - targets = load_microplex_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=candidate_year, - sources=[DataSource.IRS_SOI.value], - ) - if has_supported_tax_targets(targets, profile=profile): - return candidate_year - return None - - -def _target_fallback_reason( - targets: list[TargetSpec], - profile: MicroplexTargetProfile, -) -> str | None: - if len(targets) < profile.min_current_target_inputs: - return f"only {len(targets)} current-year target inputs" - if not has_supported_tax_targets(targets, profile=profile): - return "no supported current-year tax targets" - return None - - -def _composition_row( - target: TargetSpec, - *, - role: str, - action: str, - reason: str | None, - target_year: int, - source_value: float, - source_period: int, - composed: bool = True, -) -> dict[str, Any]: - return { - "role": role, - "action": action, - "reason": reason, - "composed": composed, - "source": target.source.value, - "variable": target.variable, - "target_type": target.target_type.value, - "source_period": source_period, - "model_period": target_year, - "period": target.period, - "stratum": target.stratum_name, - "constraints": target.constraints, - "source_value": source_value, - "target_value": target.value, - } - - -def get_soi_aging_factors( - source_year: int, - target_year: int, - db_path: Path | None = None, - jurisdiction: str = "us", -) -> SOIAgingFactors: - """ - Resolve source-backed factors for aging SOI target inputs. - - Counts are scaled by labor force: BLS annual labor-force counts when - available, then CBO labor-force projections for years beyond BLS coverage. - Amounts are scaled by aggregate SOI adjusted gross income. If the target - year is beyond available SOI AGI data, the last observed annual AGI growth - rate is carried forward and declared in the method string. - """ - if source_year == target_year: - return SOIAgingFactors( - source_year=source_year, - target_year=target_year, - count_factor=1.0, - amount_factor=1.0, - count_method="identity", - amount_method="identity", - ) - - source_labor_force = _get_labor_force_target( - year=source_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - target_labor_force, count_source = _get_labor_force_target_with_source( - year=target_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - source_aggregate_income = _soi_total_agi_value( - year=source_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - target_aggregate_income, amount_method = _soi_total_agi_for_year( - target_year=target_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - - return SOIAgingFactors( - source_year=source_year, - target_year=target_year, - count_factor=target_labor_force / source_labor_force, - amount_factor=target_aggregate_income / source_aggregate_income, - count_method=f"{count_source}_labor_force_ratio", - amount_method=amount_method, - ) - - -def age_soi_targets( - targets: list[TargetSpec], - target_year: int, - *, - db_path: Path | None = None, - jurisdiction: str = "us", - factors: SOIAgingFactors | None = None, -) -> list[TargetSpec]: - """ - Age IRS SOI target inputs to a Microplex model year. - - This composes model-year targets from source records. It does not mutate - Arch records in the database. - """ - source_years = { - target.period - for target in targets - if target.source == DataSource.IRS_SOI and target.period != target_year - } - if len(source_years) > 1: - raise ValueError( - "SOI target aging expects one SOI source year at a time; " - f"got {sorted(source_years)}." - ) - if not source_years: - return targets - - source_year = source_years.pop() - if factors is None: - factors = get_soi_aging_factors( - source_year=source_year, - target_year=target_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - elif factors.source_year != source_year or factors.target_year != target_year: - raise ValueError( - "SOI aging factors do not match target years: " - f"targets are {source_year}->{target_year}, " - f"factors are {factors.source_year}->{factors.target_year}." - ) - - aged_targets = [] - for target in targets: - if target.source != DataSource.IRS_SOI or target.period == target_year: - aged_targets.append(target) - continue - - if target.target_type == TargetType.COUNT: - factor = factors.count_factor - method = factors.count_method - elif target.target_type == TargetType.AMOUNT: - factor = factors.amount_factor - method = factors.amount_method - else: - factor = 1.0 - method = "rate_unchanged" - - aged_targets.append( - replace( - target, - value=target.value * factor, - period=target_year, - stratum_name=_aged_stratum_name( - target.stratum_name, - source_year=target.period, - target_year=target_year, - method=method, - factor=factor, - ), - ) - ) - - return aged_targets - - -def build_microplex_constraints( - microdata: pd.DataFrame, - targets: list[TargetSpec] | None = None, - *, - db_path: Path | None = None, - jurisdiction: str = "us", - year: int | None = None, - sources: list[str] | None = None, - variables: list[str] | None = None, - tolerance: float = 0.01, - min_obs: int = 0, -) -> list[Constraint]: - """ - Build flat Microplex calibration constraints from Arch DB target inputs. - - Args: - microdata: One row per calibrated unit. - targets: Optional preloaded target specs. If omitted, targets are loaded - from the Arch SQLite database. - db_path: Optional target database path. - jurisdiction: Jurisdiction prefix to load from the database. - year: Optional target period. - sources: Optional data-source filters, such as ``["irs-soi"]``. - variables: Optional target variable filters. - tolerance: Default calibration tolerance. - min_obs: Drop constraints with fewer non-zero indicator entries. - - Returns: - Constraint objects accepted by the shared calibration methods. - """ - if targets is None: - targets = load_microplex_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - sources=sources, - variables=variables, - ) - - constraints = build_constraint_matrix( - microdata=microdata, - targets=targets, - tolerance=tolerance, - ) - return _filter_constraints_by_obs(constraints, min_obs=min_obs) - - -def build_hierarchical_microplex_constraints( - hh_df: pd.DataFrame, - person_df: pd.DataFrame, - targets: list[TargetSpec] | None = None, - *, - db_path: Path | None = None, - jurisdiction: str = "us", - year: int | None = None, - sources: list[str] | None = None, - variables: list[str] | None = None, - tolerance: float = 0.01, - hh_id_col: str = "household_id", - tax_unit_df: pd.DataFrame | None = None, - min_obs: int = 0, -) -> list[Constraint]: - """ - Build household-weighted Microplex constraints from Arch DB target inputs. - - This is the adapter for hierarchical microdata where household weights must - satisfy person-level or tax-unit-level aggregate targets. - """ - if targets is None: - targets = load_microplex_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - sources=sources, - variables=variables, - ) - - constraints = build_hierarchical_constraint_matrix( - hh_df=hh_df, - person_df=person_df, - targets=targets, - tolerance=tolerance, - hh_id_col=hh_id_col, - tax_unit_df=tax_unit_df, - ) - return _filter_constraints_by_obs(constraints, min_obs=min_obs) - - -def constraints_to_ipf_dicts( - constraints: list[Constraint], -) -> list[dict[str, Any]]: - """Convert shared ``Constraint`` objects to legacy IPF constraint dicts.""" - return [ - { - "indicator": constraint.indicator, - "target_value": constraint.target_value, - "variable": constraint.variable, - "target_type": constraint.target_type.value, - "stratum": constraint.stratum_name, - "n_obs": _count_nonzero_indicator(constraint.indicator), - } - for constraint in constraints - ] - - -def _filter_constraints_by_obs( - constraints: list[Constraint], - min_obs: int, -) -> list[Constraint]: - if min_obs <= 0: - return constraints - return [ - constraint - for constraint in constraints - if _count_nonzero_indicator(constraint.indicator) >= min_obs - ] - - -def _count_nonzero_indicator(indicator: np.ndarray) -> int: - return int(np.count_nonzero(np.asarray(indicator))) - - -def _aged_stratum_name( - stratum_name: str | None, - *, - source_year: int, - target_year: int, - method: str, - factor: float, -) -> str: - base = stratum_name or "SOI target" - return f"{base} (SOI aged {source_year}->{target_year}; {method}; factor x{factor:.6g})" - - -def _get_labor_force_target( - *, - year: int, - db_path: Path | None, - jurisdiction: str, -) -> float: - value, _ = _get_labor_force_target_with_source( - year=year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - return value - - -def _get_labor_force_target_with_source( - *, - year: int, - db_path: Path | None, - jurisdiction: str, -) -> tuple[float, str]: - bls_value = _optional_target_value( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - source=DataSource.BLS, - variable="labor_force_count", - ) - if bls_value is not None: - return bls_value, "bls" - - cbo_value = _optional_target_value( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - source=DataSource.CBO, - variable="labor_force", - ) - if cbo_value is not None: - return cbo_value, "cbo" - - raise ValueError(f"No BLS/CBO labor-force target found for {year}.") - - -def _soi_total_agi_for_year( - *, - target_year: int, - db_path: Path | None, - jurisdiction: str, -) -> tuple[float, str]: - target_agi = _optional_soi_total_agi_value( - year=target_year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - if target_agi is not None: - return target_agi, "soi_total_agi_ratio" - - available = _available_soi_total_agi_values( - db_path=db_path, - jurisdiction=jurisdiction, - start_year=target_year - 20, - end_year=target_year, - ) - if len(available) < 2: - raise ValueError( - "Need at least two SOI total AGI years to extrapolate " - f"aggregate income to {target_year}." - ) - - latest_year = max(available) - previous_year = max(year for year in available if year < latest_year) - annual_growth = available[latest_year] / available[previous_year] - years_forward = target_year - latest_year - projected = available[latest_year] * annual_growth**years_forward - return ( - projected, - "soi_total_agi_last_growth_extrapolation", - ) - - -def _available_soi_total_agi_values( - *, - db_path: Path | None, - jurisdiction: str, - start_year: int, - end_year: int, -) -> dict[int, float]: - values = {} - for year in range(start_year, end_year + 1): - value = _optional_soi_total_agi_value( - year=year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - if value is not None: - values[year] = value - return values - - -def _soi_total_agi_value( - *, - year: int, - db_path: Path | None, - jurisdiction: str, -) -> float: - value = _optional_soi_total_agi_value( - year=year, - db_path=db_path, - jurisdiction=jurisdiction, - ) - if value is None: - raise ValueError(f"No SOI total AGI target found for {year}.") - return value - - -def _optional_soi_total_agi_value( - *, - year: int, - db_path: Path | None, - jurisdiction: str, -) -> float | None: - targets = get_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - sources=[DataSource.IRS_SOI.value], - variables=["adjusted_gross_income"], - ) - if not targets: - return None - - for target in targets: - if target.stratum_name == "US All Filers": - return target.value - for target in targets: - if target.constraints == [("is_tax_filer", "==", "1")]: - return target.value - return None - - -def _target_value( - *, - db_path: Path | None, - jurisdiction: str, - year: int, - source: DataSource, - variable: str, -) -> float: - value = _optional_target_value( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - source=source, - variable=variable, - ) - if value is None: - raise ValueError(f"No {source.value} {variable} target found for {year}.") - return value - - -def _optional_target_value( - *, - db_path: Path | None, - jurisdiction: str, - year: int, - source: DataSource, - variable: str, -) -> float | None: - targets = get_targets( - db_path=db_path, - jurisdiction=jurisdiction, - year=year, - sources=[source.value], - variables=[variable], - ) - if not targets: - return None - if len(targets) == 1: - return targets[0].value - - unconstrained = [target for target in targets if not target.constraints] - if len(unconstrained) == 1: - return unconstrained[0].value - return targets[0].value diff --git a/micro/us/validation_dashboard.py b/micro/us/validation_dashboard.py deleted file mode 100644 index 37a9a21f..00000000 --- a/micro/us/validation_dashboard.py +++ /dev/null @@ -1,512 +0,0 @@ -"""Versioned calibration dashboard artifacts for Microplex validation.""" - -from __future__ import annotations - -import hashlib -import html -import json -import subprocess -from datetime import datetime, timezone -from pathlib import Path -from typing import Any - -import numpy as np -import pandas as pd - - -REPORT_SCHEMA_VERSION = 1 - - -def build_microplex_version_id( - *, - year: int, - config: dict[str, Any], - git_sha: str | None = None, -) -> str: - """Return a stable default version id for a Microplex validation run.""" - sha = git_sha or current_git_sha() - short_sha = sha[:7] if sha else "dev" - fingerprint = hashlib.sha256( - json.dumps(_json_safe(config), allow_nan=False, sort_keys=True).encode( - "utf-8", - ), - ).hexdigest()[:8] - return f"microplex-us-{year}-{short_sha}-{fingerprint}" - - -def current_git_sha() -> str | None: - """Return the current git sha when running inside a git checkout.""" - try: - result = subprocess.run( - ["git", "rev-parse", "HEAD"], - check=True, - capture_output=True, - text=True, - ) - except (OSError, subprocess.CalledProcessError): - return None - return result.stdout.strip() or None - - -def write_calibration_dashboard( - output_dir: Path, - *, - summary: pd.DataFrame, - target_comparison: pd.DataFrame, - flat_diagnostics: pd.DataFrame, - household_diagnostics: pd.DataFrame, - metadata: dict[str, Any], -) -> None: - """Write structured calibration report artifacts for a Microplex version.""" - output_dir.mkdir(parents=True, exist_ok=True) - - summary.to_csv(output_dir / "summary.csv", index=False) - target_comparison.to_csv(output_dir / "target_comparison.csv", index=False) - flat_diagnostics.to_csv(output_dir / "flat_diagnostics.csv", index=False) - household_diagnostics.to_csv( - output_dir / "household_diagnostics.csv", - index=False, - ) - - active_targets = _active_target_comparison(target_comparison) - variable_summary = _variable_summary(active_targets) - worst_targets = _worst_targets(active_targets) - largest_regressions = _largest_delta_targets(active_targets, ascending=False) - largest_improvements = _largest_delta_targets(active_targets, ascending=True) - - variable_summary.to_csv(output_dir / "variable_summary.csv", index=False) - worst_targets.to_csv(output_dir / "worst_targets.csv", index=False) - largest_regressions.to_csv( - output_dir / "largest_household_regressions.csv", - index=False, - ) - largest_improvements.to_csv( - output_dir / "largest_household_improvements.csv", - index=False, - ) - - metrics = build_dashboard_metrics( - summary=summary, - target_comparison=target_comparison, - variable_summary=variable_summary, - ) - manifest = build_dashboard_manifest(metadata=metadata, metrics=metrics) - - (output_dir / "metrics.json").write_text( - json.dumps(_json_safe(metrics), allow_nan=False, indent=2, sort_keys=True) - + "\n", - encoding="utf-8", - ) - (output_dir / "manifest.json").write_text( - json.dumps(_json_safe(manifest), allow_nan=False, indent=2, sort_keys=True) - + "\n", - encoding="utf-8", - ) - (output_dir / "dashboard.html").write_text( - render_calibration_dashboard( - manifest=manifest, - metrics=metrics, - variable_summary=variable_summary, - worst_targets=worst_targets, - largest_regressions=largest_regressions, - largest_improvements=largest_improvements, - ), - encoding="utf-8", - ) - - -def build_dashboard_metrics( - *, - summary: pd.DataFrame, - target_comparison: pd.DataFrame, - variable_summary: pd.DataFrame, -) -> dict[str, Any]: - """Build dashboard metrics from report data frames.""" - summary_metrics = _summary_to_nested_dict(summary) - active = _active_target_comparison(target_comparison) - delta = active["abs_post_error_delta"].dropna() - - return { - "summary": summary_metrics, - "target_roles": { - "flat": _role_counts(target_comparison, "flat_role"), - "household": _role_counts(target_comparison, "household_role"), - }, - "active_delta": { - "target_count": int(len(active)), - "household_better_targets": int((delta < 0).sum()), - "flat_better_targets": int((delta > 0).sum()), - "ties": int((delta == 0).sum()), - "mean_abs_post_error_delta": float(delta.mean()) - if not delta.empty - else None, - "median_abs_post_error_delta": float(delta.median()) - if not delta.empty - else None, - }, - "by_variable": variable_summary.to_dict(orient="records"), - } - - -def build_dashboard_manifest( - *, - metadata: dict[str, Any], - metrics: dict[str, Any], -) -> dict[str, Any]: - """Build a machine-readable report manifest.""" - version_id = metadata.get("microplex_version") or metadata.get("version_id") - if not version_id: - raise ValueError("Dashboard metadata must include microplex_version.") - - return { - "schema_version": REPORT_SCHEMA_VERSION, - "report_type": "microplex_calibration_dashboard", - "microplex_version": version_id, - "generated_at": metadata.get("generated_at") - or datetime.now(timezone.utc).isoformat(), - "git_sha": metadata.get("git_sha"), - "jurisdiction": metadata.get("jurisdiction", "us"), - "year": metadata.get("year"), - "comparison": metadata.get("comparison", "flat_vs_household"), - "config": metadata.get("config", {}), - "headline": { - "flat_active_mean_abs_post_error": _nested_get( - metrics, - "summary", - "flat", - "active_mean_abs_post_error", - ), - "household_active_mean_abs_post_error": _nested_get( - metrics, - "summary", - "household", - "active_mean_abs_post_error", - ), - "active_targets": _nested_get( - metrics, - "active_delta", - "target_count", - ), - }, - "artifacts": { - "dashboard": "dashboard.html", - "manifest": "manifest.json", - "metrics": "metrics.json", - "summary": "summary.csv", - "target_comparison": "target_comparison.csv", - "flat_diagnostics": "flat_diagnostics.csv", - "household_diagnostics": "household_diagnostics.csv", - "variable_summary": "variable_summary.csv", - "worst_targets": "worst_targets.csv", - "largest_household_regressions": "largest_household_regressions.csv", - "largest_household_improvements": "largest_household_improvements.csv", - }, - } - - -def render_calibration_dashboard( - *, - manifest: dict[str, Any], - metrics: dict[str, Any], - variable_summary: pd.DataFrame, - worst_targets: pd.DataFrame, - largest_regressions: pd.DataFrame, - largest_improvements: pd.DataFrame, -) -> str: - """Render a static HTML calibration dashboard.""" - version = html.escape(str(manifest["microplex_version"])) - year = html.escape(str(manifest.get("year", ""))) - generated_at = html.escape(str(manifest.get("generated_at", ""))) - git_sha = html.escape(str(manifest.get("git_sha") or "unknown")) - summary = metrics["summary"] - active_delta = metrics["active_delta"] - - cards = [ - _metric_card( - "Flat active MAE", - _format_percent(_nested_get(summary, "flat", "active_mean_abs_post_error")), - ), - _metric_card( - "Household active MAE", - _format_percent( - _nested_get(summary, "household", "active_mean_abs_post_error"), - ), - ), - _metric_card( - "Household vs flat delta", - _format_percent(active_delta.get("mean_abs_post_error_delta")), - ), - _metric_card("Active targets", _format_number(active_delta["target_count"])), - _metric_card( - "Household better", - _format_number(active_delta["household_better_targets"]), - ), - _metric_card("Flat better", _format_number(active_delta["flat_better_targets"])), - ] - - return f""" - - - - - Microplex Calibration Dashboard - {version} - - - -
-

Microplex Calibration Dashboard

-
- Version {version} | Year {year} | Git {git_sha} | Generated {generated_at} -
-
-
-
- {''.join(cards)} -
- {_table_section("Active Target Summary By Variable", variable_summary)} - {_table_section("Worst Household Active Target Errors", worst_targets)} - {_table_section("Largest Household Regressions Vs Flat", largest_regressions)} - {_table_section("Largest Household Improvements Vs Flat", largest_improvements)} -
- - -""" - - -def _active_target_comparison(target_comparison: pd.DataFrame) -> pd.DataFrame: - if target_comparison.empty: - return target_comparison.copy() - return target_comparison[ - (target_comparison["flat_role"] == "active") - & (target_comparison["household_role"] == "active") - ].copy() - - -def _variable_summary(active_targets: pd.DataFrame) -> pd.DataFrame: - if active_targets.empty: - return pd.DataFrame( - columns=[ - "variable", - "target_count", - "flat_mean_abs_post_error", - "household_mean_abs_post_error", - "mean_abs_post_error_delta", - "household_better_targets", - "flat_better_targets", - ], - ) - - grouped = active_targets.groupby("variable", dropna=False) - return grouped.agg( - target_count=("target_index", "count"), - flat_mean_abs_post_error=("flat_abs_post_error", "mean"), - household_mean_abs_post_error=("household_abs_post_error", "mean"), - mean_abs_post_error_delta=("abs_post_error_delta", "mean"), - household_better_targets=("abs_post_error_delta", lambda s: int((s < 0).sum())), - flat_better_targets=("abs_post_error_delta", lambda s: int((s > 0).sum())), - ).reset_index() - - -def _worst_targets(active_targets: pd.DataFrame, n: int = 25) -> pd.DataFrame: - columns = _target_table_columns() - if active_targets.empty: - return pd.DataFrame(columns=columns) - output = active_targets.copy() - output["household_abs_post_error"] = output["household_abs_post_error"].abs() - return output.sort_values("household_abs_post_error", ascending=False)[ - columns - ].head(n) - - -def _largest_delta_targets( - active_targets: pd.DataFrame, - *, - ascending: bool, - n: int = 25, -) -> pd.DataFrame: - columns = _target_table_columns() - if active_targets.empty: - return pd.DataFrame(columns=columns) - return active_targets.sort_values( - "abs_post_error_delta", - ascending=ascending, - )[columns].head(n) - - -def _target_table_columns() -> list[str]: - return [ - "variable", - "target_type", - "stratum", - "target_value", - "flat_post_error", - "household_post_error", - "flat_abs_post_error", - "household_abs_post_error", - "abs_post_error_delta", - ] - - -def _summary_to_nested_dict(summary: pd.DataFrame) -> dict[str, dict[str, Any]]: - nested: dict[str, dict[str, Any]] = {} - for _, row in summary.iterrows(): - group = str(row["group"]) - metric = str(row["metric"]) - nested.setdefault(group, {})[metric] = _json_safe(row["value"]) - return nested - - -def _role_counts(df: pd.DataFrame, column: str) -> dict[str, int]: - if column not in df: - return {} - return { - str(role): int(count) - for role, count in df[column].fillna("missing").value_counts().items() - } - - -def _table_section(title: str, df: pd.DataFrame) -> str: - if df.empty: - table = "

No rows.

" - else: - table = df.to_html( - index=False, - classes="data-table", - border=0, - na_rep="", - float_format=lambda value: f"{value:,.4f}", - ) - return ( - f"

{html.escape(title)}

" - f"
{table}
" - ) - - -def _metric_card(label: str, value: str) -> str: - return ( - "
" - f"
{html.escape(label)}
" - f"
{html.escape(value)}
" - "
" - ) - - -def _format_percent(value: Any) -> str: - if value is None or pd.isna(value): - return "n/a" - return f"{float(value):.1%}" - - -def _format_number(value: Any) -> str: - if value is None or pd.isna(value): - return "n/a" - return f"{float(value):,.0f}" - - -def _nested_get(values: dict[str, Any], *keys: str) -> Any: - current: Any = values - for key in keys: - if not isinstance(current, dict) or key not in current: - return None - current = current[key] - return current - - -def _json_safe(value: Any) -> Any: - if isinstance(value, np.generic): - return _json_safe(value.item()) - if value is pd.NA: - return None - if isinstance(value, float) and np.isnan(value): - return None - if isinstance(value, dict): - return {str(key): _json_safe(item) for key, item in value.items()} - if isinstance(value, list | tuple): - return [_json_safe(item) for item in value] - return value diff --git a/pyproject.toml b/pyproject.toml index d96fd965..70808df4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["arch", "db", "micro", "calibration", "packages"] +packages = ["arch", "db", "micro", "packages"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/scripts/cps_to_policyengine.py b/scripts/cps_to_policyengine.py index 53260d54..c598383f 100755 --- a/scripts/cps_to_policyengine.py +++ b/scripts/cps_to_policyengine.py @@ -7,12 +7,11 @@ Usage: python cps_to_policyengine.py --year 2024 --output microsim_2024.parquet - python cps_to_policyengine.py --year 2024 --calibrate --output microsim_2024_calibrated.parquet Output format: - Person-level records with unique IDs - All income components mapped to statute variables - - Weights (original and optionally calibrated) + - Original survey weights - Relationship data for qualifying child determination """ @@ -256,7 +255,6 @@ def transform_weights(df: pd.DataFrame) -> pd.DataFrame: def convert_cps_to_policyengine( year: int, output_path: Optional[Path] = None, - calibrate: bool = False, cache_dir: Optional[Path] = None, ) -> pd.DataFrame: """ @@ -265,7 +263,6 @@ def convert_cps_to_policyengine( Args: year: Tax year (e.g., 2024) output_path: Path to save output parquet - calibrate: Apply entropy calibration to IRS SOI targets cache_dir: Override default cache directory Returns: @@ -298,10 +295,6 @@ def convert_cps_to_policyengine( result = result[result["weight"] > 0].copy() print(f" {len(result):,} records with positive weights") - # Apply calibration if requested - if calibrate: - result = apply_calibration(result, year) - # Add metadata columns result["year"] = year result["source"] = "cps-asec" @@ -320,35 +313,6 @@ def convert_cps_to_policyengine( return result -def apply_calibration(df: pd.DataFrame, year: int) -> pd.DataFrame: - """Apply entropy calibration to match IRS SOI targets.""" - try: - from micro.us.calibrate import calibrate_weights - print("Applying entropy calibration to IRS SOI targets...") - - # Calibration expects specific column names - df_cal = df.copy() - if "adjusted_gross_income" not in df_cal.columns: - print(" Warning: AGI not found, skipping calibration") - return df - - result = calibrate_weights(df_cal, verbose=True) - - if result.success: - df["original_weight"] = df["weight"] - df["weight"] = result.calibrated_weights - df["weight_adjustment"] = result.adjustment_factors - print(f" Calibration successful (KL div: {result.kl_divergence:.2f})") - else: - print(f" Calibration failed: {result.message}") - - return df - - except ImportError: - print(" Warning: Calibration module not found, skipping") - return df - - def generate_summary(df: pd.DataFrame, year: int) -> dict: """Generate summary statistics for validation.""" weight = df["weight"].values @@ -412,11 +376,6 @@ def main(): type=str, help="Output parquet path" ) - parser.add_argument( - "--calibrate", - action="store_true", - help="Apply entropy calibration to IRS SOI targets" - ) parser.add_argument( "--cache-dir", type=str, @@ -436,7 +395,6 @@ def main(): df = convert_cps_to_policyengine( year=args.year, output_path=output_path, - calibrate=args.calibrate, cache_dir=cache_dir, ) diff --git a/scripts/export_to_json.py b/scripts/export_to_json.py deleted file mode 100644 index 7fdcc3fd..00000000 --- a/scripts/export_to_json.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Export targets database to JSON for frontend consumption.""" - -import json -from pathlib import Path -from collections import defaultdict - -from sqlmodel import Session, select - -from db.schema import ( - Stratum, - Target, - get_engine, - DEFAULT_DB_PATH, -) - - -def export_targets_summary(output_path: Path | None = None): - """Export summary statistics for the calibration dashboard.""" - - if output_path is None: - output_path = Path(__file__).parent.parent.parent / "policyengine.org" / "public" / "data" / "targets_summary.json" - - engine = get_engine(DEFAULT_DB_PATH) - - with Session(engine) as session: - targets = session.exec(select(Target)).all() - strata = session.exec(select(Stratum)).all() - - # Build summary - by_source = defaultdict(lambda: {"count": 0, "years": set(), "variables": set()}) - by_year = defaultdict(lambda: {"count": 0, "sources": set()}) - by_jurisdiction = defaultdict(lambda: {"count": 0, "sources": set()}) - - # Get stratum jurisdiction lookup - stratum_jurisdiction = {s.id: s.jurisdiction.value for s in strata} - - for t in targets: - source = t.source.value - year = t.period - jurisdiction = stratum_jurisdiction.get(t.stratum_id, "unknown") - - by_source[source]["count"] += 1 - by_source[source]["years"].add(year) - by_source[source]["variables"].add(t.variable) - - by_year[year]["count"] += 1 - by_year[year]["sources"].add(source) - - by_jurisdiction[jurisdiction]["count"] += 1 - by_jurisdiction[jurisdiction]["sources"].add(source) - - # Convert sets to lists for JSON - sources_summary = [] - for source, data in by_source.items(): - years = sorted(data["years"]) - sources_summary.append({ - "source": source, - "display_name": source.upper().replace("-", " ").replace("_", " "), - "count": data["count"], - "variables": len(data["variables"]), - "year_min": min(years) if years else None, - "year_max": max(years) if years else None, - "years": years, - "is_projection": min(years) >= 2024 if years else False, - }) - - # Sort: historical sources first, then projections - sources_summary.sort(key=lambda x: (x["is_projection"], x["source"])) - - years_summary = [ - { - "year": year, - "count": data["count"], - "sources": sorted(data["sources"]), - } - for year, data in sorted(by_year.items()) - ] - - jurisdictions_summary = [ - { - "jurisdiction": jurisdiction, - "count": data["count"], - "sources": sorted(data["sources"]), - } - for jurisdiction, data in sorted(by_jurisdiction.items()) - ] - - # Sample targets for display - sample_targets = [] - for t in targets[:50]: - stratum = session.get(Stratum, t.stratum_id) - sample_targets.append({ - "stratum": stratum.name if stratum else "Unknown", - "variable": t.variable, - "year": t.period, - "value": t.value, - "type": t.target_type.value, - "source": t.source.value, - }) - - summary = { - "generated_at": str(Path(DEFAULT_DB_PATH).stat().st_mtime), - "total_targets": len(targets), - "total_strata": len(strata), - "sources": sources_summary, - "years": years_summary, - "jurisdictions": jurisdictions_summary, - "sample_targets": sample_targets, - } - - output_path.parent.mkdir(parents=True, exist_ok=True) - with open(output_path, "w") as f: - json.dump(summary, f, indent=2) - - print(f"Exported summary to {output_path}") - print(f" {len(targets)} targets across {len(strata)} strata") - print(f" {len(sources_summary)} sources: {', '.join(s['source'] for s in sources_summary)}") - - return summary - - -if __name__ == "__main__": - export_targets_summary() diff --git a/tests/test_arch_facts.py b/tests/test_arch_facts.py index 0dee71a6..3858c184 100644 --- a/tests/test_arch_facts.py +++ b/tests/test_arch_facts.py @@ -2,7 +2,7 @@ from arch.facts import DerivationStep, SourceFact from arch.sources import SourceFile, SourceReference -from arch.targets import DataSource, Jurisdiction +from db.schema import DataSource, Jurisdiction from arch.normalization import convert_units diff --git a/tests/test_arch_namespace.py b/tests/test_arch_namespace.py index 6792c260..6c7619fc 100644 --- a/tests/test_arch_namespace.py +++ b/tests/test_arch_namespace.py @@ -3,31 +3,7 @@ from arch.client import get_supabase_client from arch.microdata import get_table_name, query_cps_asec from arch.normalization import convert_units -from arch.targets import ( - Target, - TargetSpec, - TargetType, - get_targets, - query_targets, -) -from db.schema import Target as DbTarget -from db.supabase_client import ( - ARCH_SCHEMA, - MICRODATA_SCHEMA, - TARGETS_SCHEMA, - query_targets as db_query_targets, -) - - -def test_arch_targets_reexport_schema_objects(): - assert Target is DbTarget - assert TargetType.COUNT.value == "count" - assert TargetSpec.__name__ == "TargetSpec" - - -def test_arch_targets_reexport_client_helpers(): - assert query_targets is db_query_targets - assert callable(get_targets) +from db.supabase_client import ARCH_SCHEMA, MICRODATA_SCHEMA def test_arch_microdata_reexport_client_helpers(): @@ -44,7 +20,6 @@ def test_arch_client_reexports_supabase_client(): def test_arch_supabase_schema_boundaries_are_defaulted(): assert ARCH_SCHEMA == "arch" assert MICRODATA_SCHEMA == "microdata" - assert TARGETS_SCHEMA == "targets" def test_arch_normalization_exports_helpers(): diff --git a/tests/test_arch_normalization.py b/tests/test_arch_normalization.py index a890d5ee..526275ee 100644 --- a/tests/test_arch_normalization.py +++ b/tests/test_arch_normalization.py @@ -2,13 +2,12 @@ import pytest -from arch.targets import DataSource, Jurisdiction, TargetType +from db.schema import DataSource, Jurisdiction from arch.normalization import ( SourceFact, apply_share, - as_target, convert_units, - target_kwargs, + format_derivation, ) @@ -30,7 +29,7 @@ def test_convert_units_records_derivation(): assert converted.derivation[0].parameters["factor"] == 1000 -def test_apply_share_and_as_target_preserve_provenance(): +def test_apply_share_preserves_provenance(): total = SourceFact( name="aca_marketplace_enrollment", value=21_446_150, @@ -42,23 +41,20 @@ def test_apply_share_and_as_target_preserve_provenance(): ) silver = apply_share(total, 0.54, name="silver_marketplace_enrollment") - target = as_target( - silver, - variable="aca_metal_level_enrollment", - target_type=TargetType.COUNT, - stratum_name="US ACA Marketplace Silver Enrollment", - constraints=(("metal_level", "==", "silver"),), - ) - assert target.value == pytest.approx(11_580_921) - assert target.source == DataSource.CMS_ACA - assert target.jurisdiction == Jurisdiction.US - assert target.source_table == "Marketplace OEP Report" - assert target.constraints == (("metal_level", "==", "silver"),) - assert "apply_share" in target.notes + assert silver.name == "silver_marketplace_enrollment" + assert silver.value == pytest.approx(11_580_921) + assert silver.period == 2024 + assert silver.source == "cms-aca" + assert silver.jurisdiction == "us" + assert silver.source_table == "Marketplace OEP Report" + assert silver.derivation[-1].operation == "apply_share" + assert silver.derivation[-1].parameters["total_fact"] == ( + "aca_marketplace_enrollment" + ) -def test_target_kwargs_maps_blueprint_to_target_fields(): +def test_format_derivation_lists_steps_in_order(): fact = SourceFact( name="snap_households", value=22_323, @@ -67,16 +63,12 @@ def test_target_kwargs_maps_blueprint_to_target_fields(): source=DataSource.USDA_SNAP, jurisdiction=Jurisdiction.US, ) - target = as_target( - convert_units(fact, 1000, "count"), - target_type=TargetType.COUNT, - stratum_name="US SNAP Households", - ) - kwargs = target_kwargs(target, stratum_id=7) + converted = convert_units(fact, 1000, "count") + shared = apply_share(converted, 0.5, name="half_snap_households") + + formatted = format_derivation(shared.derivation) - assert kwargs["stratum_id"] == 7 - assert kwargs["variable"] == "snap_households" - assert kwargs["value"] == 22_323_000 - assert kwargs["source"] == DataSource.USDA_SNAP - assert "convert_units" in kwargs["notes"] + assert formatted.startswith("convert_units(factor=1000") + assert "; apply_share(share=0.5" in formatted + assert format_derivation(()) == "" diff --git a/tests/test_calibration.py b/tests/test_calibration.py deleted file mode 100644 index 9b8e965e..00000000 --- a/tests/test_calibration.py +++ /dev/null @@ -1,587 +0,0 @@ -"""Tests for calibration pipeline core.""" - -import tempfile -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest -from sqlmodel import Session - -from db.schema import ( - DataSource, - TargetType, - init_db, -) -from db.etl_soi import load_soi_targets - - -@pytest.fixture -def temp_db(): - """Create a temporary database with SOI targets for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_calibration.db" - engine = init_db(db_path) - with Session(engine) as session: - load_soi_targets(session, years=[2021]) - yield db_path - - -@pytest.fixture -def sample_microdata(): - """Create sample microdata for testing constraint building.""" - np.random.seed(42) - n = 1000 - return pd.DataFrame({ - "weight": np.random.uniform(100, 200, n), - "is_tax_filer": np.random.choice([0, 1], n, p=[0.2, 0.8]), - "adjusted_gross_income": np.random.lognormal(10, 1.5, n), - "age": np.random.randint(18, 85, n), - "state_fips": np.random.choice(["06", "36", "48"], n), # CA, NY, TX - "filing_status": np.random.choice(["1", "2", "3", "4"], n), # single, mfj, mfs, hoh - }) - - -class TestTargetSpec: - """Tests for TargetSpec dataclass.""" - - def test_target_spec_creation(self): - """TargetSpec should store target with constraints.""" - from calibration.targets import TargetSpec - - spec = TargetSpec( - variable="tax_unit_count", - value=153_774_296, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ) - - assert spec.variable == "tax_unit_count" - assert spec.value == 153_774_296 - assert spec.target_type == TargetType.COUNT - assert len(spec.constraints) == 1 - - def test_target_spec_with_multiple_constraints(self): - """TargetSpec should support multiple constraints.""" - from calibration.targets import TargetSpec - - spec = TargetSpec( - variable="tax_unit_count", - value=18_892_456, - target_type=TargetType.COUNT, - constraints=[ - ("adjusted_gross_income", ">=", "50000"), - ("adjusted_gross_income", "<", "75000"), - ], - source=DataSource.IRS_SOI, - period=2021, - ) - - assert len(spec.constraints) == 2 - - -class TestGetTargets: - """Tests for get_targets() function.""" - - def test_get_targets_returns_list(self, temp_db): - """get_targets should return a list of TargetSpec objects.""" - from calibration.targets import get_targets, TargetSpec - - targets = get_targets(db_path=temp_db, jurisdiction="us", year=2021) - - assert isinstance(targets, list) - assert len(targets) > 0 - assert all(isinstance(t, TargetSpec) for t in targets) - - def test_get_targets_filters_by_year(self, temp_db): - """get_targets should filter by year.""" - from calibration.targets import get_targets - - targets = get_targets(db_path=temp_db, jurisdiction="us", year=2021) - - assert all(t.period == 2021 for t in targets) - - def test_get_targets_filters_by_source(self, temp_db): - """get_targets should filter by source.""" - from calibration.targets import get_targets - - targets = get_targets( - db_path=temp_db, - jurisdiction="us", - year=2021, - sources=["irs-soi"], - ) - - assert all(t.source == DataSource.IRS_SOI for t in targets) - - def test_get_targets_filters_by_variable(self, temp_db): - """get_targets should filter by variable name.""" - from calibration.targets import get_targets - - targets = get_targets( - db_path=temp_db, - jurisdiction="us", - year=2021, - variables=["tax_unit_count"], - ) - - assert all(t.variable == "tax_unit_count" for t in targets) - - def test_get_targets_includes_constraints(self, temp_db): - """get_targets should include stratum constraints in TargetSpec.""" - from calibration.targets import get_targets - - targets = get_targets(db_path=temp_db, jurisdiction="us", year=2021) - - # Find a bracket target (should have AGI constraints) - bracket_targets = [ - t for t in targets - if any("adjusted_gross_income" in c[0] for c in t.constraints) - ] - - assert len(bracket_targets) > 0 - - def test_get_targets_empty_for_nonexistent_year(self, temp_db): - """get_targets should return empty list for year with no data.""" - from calibration.targets import get_targets - - targets = get_targets(db_path=temp_db, jurisdiction="us", year=1900) - - assert targets == [] - - -class TestBuildConstraintMatrix: - """Tests for build_constraint_matrix() function.""" - - def test_build_constraint_matrix_returns_constraints(self, sample_microdata): - """build_constraint_matrix should return Constraint objects.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix, Constraint - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=800, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - assert isinstance(constraints, list) - assert len(constraints) == 1 - assert isinstance(constraints[0], Constraint) - - def test_constraint_has_indicator_vector(self, sample_microdata): - """Constraint should have indicator vector matching microdata rows.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=800, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - assert len(constraints[0].indicator) == len(sample_microdata) - assert constraints[0].indicator.dtype == np.float64 - - def test_constraint_indicator_matches_stratum(self, sample_microdata): - """Indicator vector should match stratum definition.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=800, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - # Indicator should be 1 where is_tax_filer == 1 - expected = (sample_microdata["is_tax_filer"] == 1).astype(float).values - np.testing.assert_array_equal(constraints[0].indicator, expected) - - def test_constraint_with_range(self, sample_microdata): - """Constraint should handle range conditions (>=, <).""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=100, - target_type=TargetType.COUNT, - constraints=[ - ("adjusted_gross_income", ">=", "50000"), - ("adjusted_gross_income", "<", "75000"), - ], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - # Indicator should be 1 where AGI in range - agi = sample_microdata["adjusted_gross_income"] - expected = ((agi >= 50000) & (agi < 75000)).astype(float).values - np.testing.assert_array_equal(constraints[0].indicator, expected) - - def test_constraint_stores_target_value(self, sample_microdata): - """Constraint should store the target value.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=800, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - assert constraints[0].target_value == 800 - - def test_constraint_amount_type_uses_variable(self, sample_microdata): - """For AMOUNT type, indicator should be variable * mask.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="adjusted_gross_income", - value=1_000_000_000, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - # For AMOUNT, indicator should be AGI * is_tax_filer - mask = sample_microdata["is_tax_filer"] == 1 - expected = (sample_microdata["adjusted_gross_income"] * mask).values - np.testing.assert_array_almost_equal(constraints[0].indicator, expected) - - def test_constraint_with_tolerance(self, sample_microdata): - """Constraints should store tolerance.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=800, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix( - sample_microdata, targets, tolerance=0.05 - ) - - assert constraints[0].tolerance == 0.05 - - def test_multiple_constraints(self, sample_microdata): - """Should build constraints for multiple targets.""" - from calibration.targets import TargetSpec - from calibration.constraints import build_constraint_matrix - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=800, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - TargetSpec( - variable="adjusted_gross_income", - value=1_000_000_000, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - constraints = build_constraint_matrix(sample_microdata, targets) - - assert len(constraints) == 2 - - -class TestEntropyCalibrator: - """Tests for EntropyCalibrator.""" - - def test_entropy_calibrator_exists(self): - """EntropyCalibrator class should exist.""" - from calibration.methods.entropy import EntropyCalibrator - - calibrator = EntropyCalibrator() - assert calibrator is not None - - def test_entropy_calibrator_has_calibrate_method(self): - """EntropyCalibrator should have calibrate method.""" - from calibration.methods.entropy import EntropyCalibrator - - calibrator = EntropyCalibrator() - assert hasattr(calibrator, "calibrate") - - def test_calibrated_weights_sum_correctly(self): - """Calibrated weights should sum to expected total.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - # Simple case: 100 records, want total weight 1000 - n = 100 - original_weights = np.ones(n) * 10.0 # Sum = 1000 - - # Create constraint: sum of all weights should be 1200 - indicator = np.ones(n) - constraints = [ - Constraint( - indicator=indicator, - target_value=1200.0, - variable="total", - target_type=TargetType.COUNT, - ) - ] - - calibrator = EntropyCalibrator() - calibrated = calibrator.calibrate(original_weights, constraints) - - # Sum should match target - np.testing.assert_allclose(calibrated.sum(), 1200.0, rtol=1e-4) - - def test_constraints_satisfied_within_tolerance(self): - """Calibration should satisfy constraints within tolerance.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - np.random.seed(42) - n = 200 - original_weights = np.random.uniform(5, 15, n) - - # Two strata: first half and second half - indicator1 = np.concatenate([np.ones(n // 2), np.zeros(n // 2)]) - indicator2 = np.concatenate([np.zeros(n // 2), np.ones(n // 2)]) - - constraints = [ - Constraint( - indicator=indicator1, - target_value=600.0, - variable="stratum1", - target_type=TargetType.COUNT, - tolerance=0.01, - ), - Constraint( - indicator=indicator2, - target_value=800.0, - variable="stratum2", - target_type=TargetType.COUNT, - tolerance=0.01, - ), - ] - - calibrator = EntropyCalibrator() - calibrated = calibrator.calibrate(original_weights, constraints) - - # Check constraints are satisfied - actual1 = (calibrated * indicator1).sum() - actual2 = (calibrated * indicator2).sum() - - np.testing.assert_allclose(actual1, 600.0, rtol=0.01) - np.testing.assert_allclose(actual2, 800.0, rtol=0.01) - - def test_weights_respect_bounds(self): - """Calibrated weights should respect min/max bounds.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - n = 100 - original_weights = np.ones(n) * 10.0 - - indicator = np.ones(n) - constraints = [ - Constraint( - indicator=indicator, - target_value=1500.0, - variable="total", - target_type=TargetType.COUNT, - ) - ] - - # Set bounds: weights can be 0.5x to 2x original - calibrator = EntropyCalibrator(bounds=(0.5, 2.0)) - calibrated = calibrator.calibrate(original_weights, constraints) - - # Check all weights are within bounds - min_weight = original_weights * 0.5 - max_weight = original_weights * 2.0 - assert np.all(calibrated >= min_weight - 1e-6) - assert np.all(calibrated <= max_weight + 1e-6) - - def test_minimal_deviation_from_original(self): - """Calibration should minimize KL divergence.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - np.random.seed(42) - n = 100 - original_weights = np.random.uniform(8, 12, n) - - # Constraint that's already nearly satisfied - indicator = np.ones(n) - current_sum = original_weights.sum() - target = current_sum * 1.05 # Just 5% adjustment needed - - constraints = [ - Constraint( - indicator=indicator, - target_value=target, - variable="total", - target_type=TargetType.COUNT, - ) - ] - - calibrator = EntropyCalibrator() - calibrated = calibrator.calibrate(original_weights, constraints) - - # Compute KL divergence: sum(w * log(w/w0)) - kl_div = np.sum(calibrated * np.log(calibrated / original_weights)) - - # For a single uniform constraint, uniform scaling is optimal - # and should give the same KL divergence - uniform_scale = target / current_sum - uniform_weights = original_weights * uniform_scale - kl_uniform = np.sum(uniform_weights * np.log(uniform_weights / original_weights)) - - # Our solution should be very close to uniform scaling - np.testing.assert_allclose(kl_div, kl_uniform, rtol=0.01) - - def test_amount_constraint(self): - """Should handle AMOUNT type constraints correctly.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - np.random.seed(42) - n = 100 - original_weights = np.ones(n) * 10.0 - income = np.random.lognormal(10, 1, n) - - # Constraint: total weighted income should be 5M - indicator = income - constraints = [ - Constraint( - indicator=indicator, - target_value=5_000_000.0, - variable="income", - target_type=TargetType.AMOUNT, - ) - ] - - calibrator = EntropyCalibrator() - calibrated = calibrator.calibrate(original_weights, constraints) - - # Check constraint is satisfied - actual = (calibrated * income).sum() - np.testing.assert_allclose(actual, 5_000_000.0, rtol=0.01) - - def test_multiple_amount_constraints(self): - """Should handle multiple AMOUNT constraints simultaneously.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - np.random.seed(42) - n = 100 - original_weights = np.ones(n) * 10.0 - income = np.random.lognormal(10, 1, n) - taxes = income * 0.2 - - constraints = [ - Constraint( - indicator=income, - target_value=5_000_000.0, - variable="income", - target_type=TargetType.AMOUNT, - ), - Constraint( - indicator=taxes, - target_value=1_000_000.0, - variable="taxes", - target_type=TargetType.AMOUNT, - ), - ] - - calibrator = EntropyCalibrator() - calibrated = calibrator.calibrate(original_weights, constraints) - - # Both constraints should be satisfied - actual_income = (calibrated * income).sum() - actual_taxes = (calibrated * taxes).sum() - - np.testing.assert_allclose(actual_income, 5_000_000.0, rtol=0.01) - np.testing.assert_allclose(actual_taxes, 1_000_000.0, rtol=0.01) - - def test_infeasible_constraints_raises_error(self): - """Should raise error if constraints are infeasible.""" - from calibration.methods.entropy import EntropyCalibrator - from calibration.constraints import Constraint - - n = 10 - original_weights = np.ones(n) * 10.0 - - # Infeasible: want sum = 500, but max possible is 10 * 10 * 10 = 1000 - # with bounds (0.1, 10) and tight second constraint - indicator1 = np.ones(n) - indicator2 = np.ones(n) - - constraints = [ - Constraint( - indicator=indicator1, - target_value=50.0, # Need average weight of 5 - variable="total", - target_type=TargetType.COUNT, - ), - Constraint( - indicator=indicator2, - target_value=150.0, # Need average weight of 15 - variable="total2", - target_type=TargetType.COUNT, - ), - ] - - calibrator = EntropyCalibrator(bounds=(0.5, 1.5)) - - with pytest.raises((ValueError, RuntimeError)): - calibrator.calibrate(original_weights, constraints) diff --git a/tests/test_congressional_districts.py b/tests/test_congressional_districts.py deleted file mode 100644 index 728b2c38..00000000 --- a/tests/test_congressional_districts.py +++ /dev/null @@ -1,210 +0,0 @@ -"""Tests for congressional district geographic level functionality.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - GeographicLevel, - Stratum, - Target, - init_db, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_cd.db" - engine = init_db(db_path) - yield engine - - -class TestCongressionalDistrictGeographicLevel: - """Tests for congressional district geographic level support.""" - - def test_geographic_level_enum_exists(self): - """GeographicLevel enum should include congressional_district.""" - assert GeographicLevel.CONGRESSIONAL_DISTRICT == "congressional_district" - assert GeographicLevel.STATE == "state" - assert GeographicLevel.NATIONAL == "national" - - def test_load_congressional_district_targets(self, temp_db): - """Loading congressional district data should create proper targets.""" - from db.etl_census import load_congressional_district_targets - - district_data = { - ("06", "01"): {"population": 750000, "households": 280000}, - ("06", "12"): {"population": 720000, "households": 265000}, - ("36", "15"): {"population": 710000, "households": 290000}, - } - - with Session(temp_db) as session: - load_congressional_district_targets(session, 2023, district_data) - - # Should create 3 congressional district strata - cd_strata = session.exec( - select(Stratum).where( - Stratum.stratum_group_id == "congressional_districts" - ) - ).all() - - assert len(cd_strata) == 3 - - def test_congressional_district_target_has_correct_geographic_level(self, temp_db): - """Congressional district targets should have correct geographic level.""" - from db.etl_census import load_congressional_district_targets - - district_data = { - ("06", "01"): {"population": 750000}, - } - - with Session(temp_db) as session: - load_congressional_district_targets(session, 2023, district_data) - - # Find the target - target = session.exec( - select(Target) - .where(Target.variable == "population") - .where(Target.period == 2023) - ).first() - - assert target is not None - assert target.geographic_level == GeographicLevel.CONGRESSIONAL_DISTRICT - assert target.value == 750000 - - def test_congressional_district_stratum_has_correct_constraints(self, temp_db): - """Congressional district strata should have state and district constraints.""" - from db.etl_census import load_congressional_district_targets - - district_data = { - ("06", "01"): {"population": 750000}, - } - - with Session(temp_db) as session: - load_congressional_district_targets(session, 2023, district_data) - - # Find the stratum - stratum = session.exec( - select(Stratum).where( - Stratum.name == "Congressional District 06-01" - ) - ).first() - - assert stratum is not None - assert len(stratum.constraints) == 2 - - # Check constraints - constraint_vars = {c.variable for c in stratum.constraints} - assert "state_fips" in constraint_vars - assert "congressional_district" in constraint_vars - - # Check values - state_constraint = next( - c for c in stratum.constraints if c.variable == "state_fips" - ) - district_constraint = next( - c for c in stratum.constraints if c.variable == "congressional_district" - ) - - assert state_constraint.value == "06" - assert state_constraint.operator == "==" - assert district_constraint.value == "01" - assert district_constraint.operator == "==" - - def test_congressional_district_with_households(self, temp_db): - """Congressional district targets should include household data when available.""" - from db.etl_census import load_congressional_district_targets - - district_data = { - ("06", "01"): {"population": 750000, "households": 280000}, - } - - with Session(temp_db) as session: - load_congressional_district_targets(session, 2023, district_data) - - # Should have both population and household targets - targets = session.exec(select(Target)).all() - - assert len(targets) == 2 - - variables = {t.variable for t in targets} - assert "population" in variables - assert "household_count" in variables - - # Both should have congressional district geographic level - for target in targets: - assert target.geographic_level == GeographicLevel.CONGRESSIONAL_DISTRICT - - def test_multiple_districts_same_state(self, temp_db): - """Should handle multiple congressional districts in the same state.""" - from db.etl_census import load_congressional_district_targets - - district_data = { - ("06", "01"): {"population": 750000}, - ("06", "12"): {"population": 720000}, - ("06", "30"): {"population": 740000}, - } - - with Session(temp_db) as session: - load_congressional_district_targets(session, 2023, district_data) - - # Should create 3 separate strata - ca_districts = session.exec( - select(Stratum).where( - Stratum.stratum_group_id == "congressional_districts" - ) - ).all() - - assert len(ca_districts) == 3 - - # All should have state_fips == "06" - for stratum in ca_districts: - state_constraint = next( - c for c in stratum.constraints if c.variable == "state_fips" - ) - assert state_constraint.value == "06" - - def test_query_by_geographic_level(self, temp_db): - """Should be able to query targets by geographic level.""" - from db.etl_census import load_census_targets, load_congressional_district_targets - - district_data = { - ("06", "01"): {"population": 750000}, - } - - with Session(temp_db) as session: - # Load both national and congressional district data - load_census_targets(session, years=[2023]) - load_congressional_district_targets(session, 2023, district_data) - - # Query for congressional district targets only - cd_targets = session.exec( - select(Target).where( - Target.geographic_level == GeographicLevel.CONGRESSIONAL_DISTRICT - ) - ).all() - - assert len(cd_targets) == 1 - assert cd_targets[0].value == 750000 - - # Query for national targets - national_targets = session.exec( - select(Target).where( - Target.geographic_level == GeographicLevel.NATIONAL - ) - ).all() - - assert len(national_targets) > 0 - - # Query for state targets - state_targets = session.exec( - select(Target).where( - Target.geographic_level == GeographicLevel.STATE - ) - ).all() - - assert len(state_targets) > 0 diff --git a/tests/test_download_cps.py b/tests/test_download_cps.py index 4efce57f..5fac17ce 100644 --- a/tests/test_download_cps.py +++ b/tests/test_download_cps.py @@ -6,7 +6,7 @@ import pandas as pd import pytest -# Direct imports from download_cps to avoid calibration package import chain +# Direct imports from the downloader module. from micro.us.census.download_cps import ( CPS_URL_BY_YEAR, PERSON_COLUMNS, @@ -132,55 +132,3 @@ def test_extract_person_data_handles_missing_columns(self): # Should have at least the columns we provided assert len(df) == 1 assert "age" in df.columns or "A_AGE" in df.columns - - -class TestLoaderIntegration: - """Tests for loader integration with CPS downloader. - - These tests require the calibration package which may have optional - dependencies. They are skipped if the imports fail. - """ - - @pytest.fixture(autouse=True) - def skip_if_calibration_unavailable(self): - """Skip these tests if calibration module can't be imported.""" - pytest.importorskip("calibration") - - def test_load_microdata_loads_cps_with_correct_weights(self): - """CPS data should have reasonable total weight (~330M for US pop).""" - from calibration import load_microdata - from pathlib import Path - - # Only run if CPS 2023 exists (it should after download) - data_path = Path(__file__).parent.parent / "micro" / "us" / "cps_2023.parquet" - if not data_path.exists(): - pytest.skip("CPS 2023 not downloaded") - - df = load_microdata(source="cps", year=2023) - - # Should have data - assert len(df) > 100_000 # CPS has 100k+ person records - assert "weight" in df.columns - - # Weight should sum to ~330M (US population), not 33B - total_weight = df["weight"].sum() - assert 300_000_000 < total_weight < 400_000_000, ( - f"Total weight {total_weight:,.0f} should be ~330M, not {total_weight/1e9:.1f}B" - ) - - def test_load_microdata_falls_back_to_synthetic(self): - """Should fall back to synthetic when CPS unavailable.""" - from calibration import load_microdata - import warnings - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - df = load_microdata(source="cps", year=2015) # Old year, won't download - - # Should have data (synthetic fallback) - assert len(df) > 0 - assert "weight" in df.columns - - # Should have warned about fallback - assert any("falling back to synthetic" in str(warning.message).lower() - for warning in w) diff --git a/tests/test_etl_aca_enrollment.py b/tests/test_etl_aca_enrollment.py deleted file mode 100644 index c3126d00..00000000 --- a/tests/test_etl_aca_enrollment.py +++ /dev/null @@ -1,312 +0,0 @@ -"""Tests for ACA Marketplace enrollment ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_aca_enrollment import ( - load_aca_enrollment_targets, - ACA_ENROLLMENT_DATA, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_aca.db" - engine = init_db(db_path) - yield engine - - -class TestAcaEnrollmentETL: - """Tests for ACA Marketplace enrollment ETL loader.""" - - def test_load_aca_creates_national_stratum(self, temp_db): - """Loading ACA data should create a national stratum.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).first() - - assert national is not None - assert national.stratum_group_id == "aca_national" - - def test_load_aca_creates_enrollment_targets(self, temp_db): - """Loading ACA data should create enrollment count targets.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).first() - - enrollment_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "aca_marketplace_enrollment") - .where(Target.period == 2024) - ).first() - - assert enrollment_target is not None - expected = ACA_ENROLLMENT_DATA[2024]["national"]["total_enrollment"] - assert enrollment_target.value == expected - assert enrollment_target.target_type == TargetType.COUNT - assert enrollment_target.source == DataSource.CMS_ACA - - def test_load_aca_creates_aptc_stratum(self, temp_db): - """Loading ACA data should create APTC recipient stratum.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - aptc_stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACA APTC Recipients") - ).first() - - assert aptc_stratum is not None - assert aptc_stratum.stratum_group_id == "aca_subsidies" - - def test_load_aca_creates_aptc_targets(self, temp_db): - """Loading ACA should create APTC recipient and amount targets.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - aptc_stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACA APTC Recipients") - ).first() - - # Check recipient count - aptc_count = session.exec( - select(Target) - .where(Target.stratum_id == aptc_stratum.id) - .where(Target.variable == "aca_aptc_recipients") - ).first() - - assert aptc_count is not None - expected = ACA_ENROLLMENT_DATA[2024]["national"]["aptc_recipients"] - assert aptc_count.value == expected - - # Check average APTC amount - aptc_amount = session.exec( - select(Target) - .where(Target.stratum_id == aptc_stratum.id) - .where(Target.variable == "aca_avg_monthly_aptc") - ).first() - - assert aptc_amount is not None - assert aptc_amount.target_type == TargetType.AMOUNT - - def test_load_aca_creates_csr_stratum(self, temp_db): - """Loading ACA data should create CSR recipient stratum for years with data.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - csr_stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACA CSR Recipients") - ).first() - - assert csr_stratum is not None - assert csr_stratum.stratum_group_id == "aca_subsidies" - - def test_load_aca_creates_metal_level_strata(self, temp_db): - """Loading ACA data should create metal level strata.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - metal_levels = ["Bronze", "Silver", "Gold", "Platinum", "Catastrophic"] - for level in metal_levels: - stratum = session.exec( - select(Stratum).where( - Stratum.name == f"US ACA {level} Plan Enrollees" - ) - ).first() - - assert stratum is not None, f"Missing stratum for {level}" - assert stratum.stratum_group_id == "aca_metal_levels" - - def test_load_aca_metal_level_enrollment_correct(self, temp_db): - """Metal level enrollment should be calculated from percentages.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - silver_stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACA Silver Plan Enrollees") - ).first() - - enrollment = session.exec( - select(Target) - .where(Target.stratum_id == silver_stratum.id) - .where(Target.variable == "aca_marketplace_enrollment") - ).first() - - national_data = ACA_ENROLLMENT_DATA[2024]["national"] - expected = int(national_data["total_enrollment"] * national_data["silver_pct"]) - assert enrollment.value == expected - assert "apply_share" in enrollment.notes - - def test_load_aca_creates_state_strata(self, temp_db): - """Loading ACA should create state-level strata.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "aca_states") - ).all() - - # Should have strata for states in the data - expected_states = len(ACA_ENROLLMENT_DATA[2024].get("states", {})) - assert len(state_strata) == expected_states - - def test_load_aca_state_targets_correct(self, temp_db): - """State-level ACA targets should have correct values.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - fl_stratum = session.exec( - select(Stratum).where(Stratum.name == "FL ACA Marketplace Enrollees") - ).first() - - assert fl_stratum is not None - - fl_enrollment = session.exec( - select(Target) - .where(Target.stratum_id == fl_stratum.id) - .where(Target.variable == "aca_marketplace_enrollment") - ).first() - - expected = ACA_ENROLLMENT_DATA[2024]["states"]["FL"]["enrollment"] - assert fl_enrollment.value == expected - - def test_load_aca_state_metal_levels(self, temp_db): - """States with metal level data should have metal level strata.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - # Florida has metal level data in 2024 - fl_silver = session.exec( - select(Stratum).where(Stratum.name == "FL ACA Silver Plan Enrollees") - ).first() - - assert fl_silver is not None - assert fl_silver.stratum_group_id == "aca_state_metal_levels" - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2023, 2024, 2025]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "aca_marketplace_enrollment") - ).all() - - years = {t.period for t in targets} - assert years == {2023, 2024, 2025} - - def test_load_aca_idempotent(self, temp_db): - """Loading ACA twice should not duplicate data.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - load_aca_enrollment_targets(session, years=[2024]) - - national_strata = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).all() - - # Should only have one national stratum - assert len(national_strata) == 1 - - def test_state_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).first() - - ca = session.exec( - select(Stratum).where(Stratum.name == "CA ACA Marketplace Enrollees") - ).first() - - assert ca.parent_id == national.id - - def test_aptc_stratum_has_parent(self, temp_db): - """APTC stratum should have national stratum as parent.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).first() - - aptc = session.exec( - select(Stratum).where(Stratum.name == "US ACA APTC Recipients") - ).first() - - assert aptc.parent_id == national.id - - def test_premium_targets_created(self, temp_db): - """Premium-related targets should be created.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US ACA Marketplace Enrollees") - ).first() - - # Check gross premium - gross_premium = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "aca_avg_monthly_premium_gross") - ).first() - - assert gross_premium is not None - assert gross_premium.target_type == TargetType.AMOUNT - - # Check net premium - net_premium = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "aca_avg_monthly_premium_net") - ).first() - - assert net_premium is not None - assert net_premium.target_type == TargetType.AMOUNT - - def test_metal_level_share_rate_target(self, temp_db): - """Metal level share should be stored as a rate target.""" - with Session(temp_db) as session: - load_aca_enrollment_targets(session, years=[2024]) - - silver = session.exec( - select(Stratum).where(Stratum.name == "US ACA Silver Plan Enrollees") - ).first() - - share = session.exec( - select(Target) - .where(Target.stratum_id == silver.id) - .where(Target.variable == "aca_metal_level_share") - ).first() - - assert share is not None - assert share.target_type == TargetType.RATE - expected = ACA_ENROLLMENT_DATA[2024]["national"]["silver_pct"] - assert share.value == expected diff --git a/tests/test_etl_bls.py b/tests/test_etl_bls.py deleted file mode 100644 index ecbca4d7..00000000 --- a/tests/test_etl_bls.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Tests for BLS (Bureau of Labor Statistics) employment ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_bls.db" - engine = init_db(db_path) - yield engine - - -class TestBlsETL: - """Tests for BLS employment ETL loader.""" - - def test_load_creates_labor_force_stratum(self, temp_db): - """Loading BLS data should create labor force stratum.""" - from db.etl_bls import load_bls_targets - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - assert labor_force is not None - assert labor_force.jurisdiction == Jurisdiction.US - - def test_load_creates_employment_count(self, temp_db): - """Loading BLS should create total employment target.""" - from db.etl_bls import load_bls_targets, BLS_DATA - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - employed = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - .where(Target.period == 2023) - ).first() - - assert employed is not None - assert employed.value == BLS_DATA[2023]["employed"] - assert employed.target_type == TargetType.COUNT - assert employed.source == DataSource.BLS - - def test_load_creates_unemployment_count(self, temp_db): - """Loading BLS should create unemployment count target.""" - from db.etl_bls import load_bls_targets, BLS_DATA - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - unemployed = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "unemployed") - ).first() - - assert unemployed is not None - assert unemployed.value == BLS_DATA[2023]["unemployed"] - - def test_load_creates_unemployment_rate(self, temp_db): - """Loading BLS should create unemployment rate target.""" - from db.etl_bls import load_bls_targets, BLS_DATA - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - rate = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "unemployment_rate") - ).first() - - assert rate is not None - assert rate.value == BLS_DATA[2023]["unemployment_rate"] - assert rate.target_type == TargetType.RATE - - def test_load_creates_labor_force_participation(self, temp_db): - """Loading BLS should create LFPR target.""" - from db.etl_bls import load_bls_targets, BLS_DATA - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - lfpr = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "labor_force_participation_rate") - ).first() - - assert lfpr is not None - assert lfpr.value == BLS_DATA[2023]["labor_force_participation_rate"] - assert lfpr.target_type == TargetType.RATE - - def test_load_creates_median_wage(self, temp_db): - """Loading BLS should create median wage target.""" - from db.etl_bls import load_bls_targets, BLS_DATA - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - wage = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "median_weekly_earnings") - ).first() - - assert wage is not None - assert wage.value == BLS_DATA[2023]["median_weekly_earnings"] - assert wage.target_type == TargetType.AMOUNT - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - from db.etl_bls import load_bls_targets - - with Session(temp_db) as session: - load_bls_targets(session, years=[2022, 2023]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - ).all() - - years = {t.period for t in targets} - assert years == {2022, 2023} - - def test_load_idempotent(self, temp_db): - """Loading twice should not duplicate strata.""" - from db.etl_bls import load_bls_targets - - with Session(temp_db) as session: - load_bls_targets(session, years=[2023]) - load_bls_targets(session, years=[2023]) - - labor_force_strata = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).all() - - assert len(labor_force_strata) == 1 diff --git a/tests/test_etl_cbo.py b/tests/test_etl_cbo.py deleted file mode 100644 index d021f00c..00000000 --- a/tests/test_etl_cbo.py +++ /dev/null @@ -1,170 +0,0 @@ -"""Tests for CBO projections ETL.""" - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_cbo import load_cbo_targets, CBO_DATA - - -class TestCboETL: - """Tests for CBO projections ETL.""" - - @pytest.fixture - def session(self, tmp_path): - """Create a test database session.""" - db_path = tmp_path / "test.db" - engine = init_db(db_path) - with Session(engine) as session: - yield session - - def test_load_creates_federal_budget_stratum(self, session): - """Should create federal budget stratum.""" - load_cbo_targets(session, years=[2024]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US Federal Budget") - ).first() - - assert stratum is not None - assert stratum.jurisdiction == Jurisdiction.US_FEDERAL - - def test_load_creates_gdp_projection(self, session): - """Should create GDP projection targets.""" - load_cbo_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "gdp", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.value == CBO_DATA[2024]["gdp"] - assert target.target_type == TargetType.AMOUNT - assert target.source == DataSource.CBO - - def test_load_creates_revenue_projection(self, session): - """Should create federal revenue projection.""" - load_cbo_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "federal_revenue", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.value == CBO_DATA[2024]["federal_revenue"] - assert target.target_type == TargetType.AMOUNT - - def test_load_creates_outlays_projection(self, session): - """Should create federal outlays projection.""" - load_cbo_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "federal_outlays", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.value == CBO_DATA[2024]["federal_outlays"] - - def test_load_creates_deficit_projection(self, session): - """Should create deficit projection.""" - load_cbo_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "federal_deficit", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.value == CBO_DATA[2024]["federal_deficit"] - - def test_load_creates_economic_stratum(self, session): - """Should create US economy stratum for macro indicators.""" - load_cbo_targets(session, years=[2024]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US Economy") - ).first() - - assert stratum is not None - - def test_load_creates_unemployment_projection(self, session): - """Should create unemployment rate projection.""" - load_cbo_targets(session, years=[2025]) - - target = session.exec( - select(Target).where( - Target.variable == "unemployment_rate", - Target.period == 2025, - Target.source == DataSource.CBO, - ) - ).first() - - assert target is not None - assert target.target_type == TargetType.RATE - - def test_load_creates_inflation_projection(self, session): - """Should create inflation rate projection.""" - load_cbo_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "cpi_inflation", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.target_type == TargetType.RATE - - def test_load_future_years(self, session): - """Should load projection years through 2034.""" - load_cbo_targets(session, years=[2030, 2034]) - - targets_2030 = session.exec( - select(Target).where(Target.period == 2030) - ).all() - targets_2034 = session.exec( - select(Target).where(Target.period == 2034) - ).all() - - assert len(targets_2030) > 0 - assert len(targets_2034) > 0 - - def test_load_multiple_years(self, session): - """Should load multiple projection years.""" - load_cbo_targets(session, years=[2024, 2025, 2026]) - - targets = session.exec(select(Target)).all() - periods = {t.period for t in targets} - - assert 2024 in periods - assert 2025 in periods - assert 2026 in periods - - def test_load_idempotent(self, session): - """Loading twice should not duplicate data.""" - load_cbo_targets(session, years=[2024]) - count1 = len(session.exec(select(Target)).all()) - - load_cbo_targets(session, years=[2024]) - count2 = len(session.exec(select(Target)).all()) - - assert count1 == count2 diff --git a/tests/test_etl_census.py b/tests/test_etl_census.py deleted file mode 100644 index fd6339a3..00000000 --- a/tests/test_etl_census.py +++ /dev/null @@ -1,206 +0,0 @@ -"""Tests for Census population ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_census.db" - engine = init_db(db_path) - yield engine - - -class TestCensusPopulationETL: - """Tests for Census population ETL loader.""" - - def test_load_creates_national_stratum(self, temp_db): - """Loading Census data should create a US population stratum.""" - from db.etl_census import load_census_targets - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Population") - ).first() - - assert national is not None - assert national.jurisdiction == Jurisdiction.US - - def test_load_creates_total_population(self, temp_db): - """Loading Census should create total population target.""" - from db.etl_census import load_census_targets, CENSUS_DATA - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Population") - ).first() - - pop_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "population") - .where(Target.period == 2023) - ).first() - - assert pop_target is not None - assert pop_target.value == CENSUS_DATA[2023]["total_population"] - assert pop_target.target_type == TargetType.COUNT - assert pop_target.source == DataSource.CENSUS_ACS - - def test_load_creates_age_group_strata(self, temp_db): - """Loading Census should create age group strata.""" - from db.etl_census import load_census_targets - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - age_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "age_groups") - ).all() - - # Should have 18 age groups (5-year brackets) - assert len(age_strata) == 18 - - def test_load_creates_state_strata(self, temp_db): - """Loading Census should create state population strata.""" - from db.etl_census import load_census_targets, CENSUS_DATA - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "state_population") - ).all() - - # Should have strata for states in the data - expected_states = len(CENSUS_DATA[2023].get("states", {})) - assert len(state_strata) == expected_states - - def test_state_population_correct(self, temp_db): - """State population targets should have correct values.""" - from db.etl_census import load_census_targets, CENSUS_DATA - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "California Population") - ).first() - - assert ca_stratum is not None - - ca_pop = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "population") - ).first() - - assert ca_pop.value == CENSUS_DATA[2023]["states"]["CA"]["population"] - - def test_load_creates_household_targets(self, temp_db): - """Loading Census should create household count targets.""" - from db.etl_census import load_census_targets, CENSUS_DATA - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Population") - ).first() - - hh_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "household_count") - ).first() - - assert hh_target is not None - assert hh_target.value == CENSUS_DATA[2023]["households"] - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - from db.etl_census import load_census_targets - - with Session(temp_db) as session: - load_census_targets(session, years=[2022, 2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Population") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "population") - ).all() - - years = {t.period for t in targets} - assert years == {2022, 2023} - - def test_load_idempotent(self, temp_db): - """Loading twice should not duplicate data.""" - from db.etl_census import load_census_targets - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - load_census_targets(session, years=[2023]) - - national_strata = session.exec( - select(Stratum).where(Stratum.name == "US Population") - ).all() - - assert len(national_strata) == 1 - - def test_age_group_population_sums_to_total(self, temp_db): - """Age group populations should approximately sum to total.""" - from db.etl_census import load_census_targets - - with Session(temp_db) as session: - load_census_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Population") - ).first() - - total_pop = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "population") - .where(Target.period == 2023) - ).first() - - age_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "age_groups") - ).all() - - age_sum = 0 - for stratum in age_strata: - target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "population") - .where(Target.period == 2023) - ).first() - if target: - age_sum += target.value - - # Should be within 1% (accounting for rounding) - assert abs(age_sum - total_pop.value) / total_pop.value < 0.01 diff --git a/tests/test_etl_cps.py b/tests/test_etl_cps.py deleted file mode 100644 index 6002edca..00000000 --- a/tests/test_etl_cps.py +++ /dev/null @@ -1,282 +0,0 @@ -"""Tests for BLS CPS (Current Population Survey) monthly employment ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_cps.db" - engine = init_db(db_path) - yield engine - - -class TestCpsETL: - """Tests for CPS monthly employment ETL loader.""" - - def test_load_creates_labor_force_stratum(self, temp_db): - """Loading CPS data should create labor force stratum.""" - from db.etl_cps import load_cps_targets - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - assert labor_force is not None - assert labor_force.jurisdiction == Jurisdiction.US - assert labor_force.stratum_group_id == "cps_monthly" - - def test_load_creates_employment_count(self, temp_db): - """Loading CPS should create employment count target.""" - from db.etl_cps import load_cps_targets, CPS_MONTHLY_DATA - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - # Period format: YYYYMM (202312 for December 2023) - period = 202312 - - employed = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - .where(Target.period == period) - ).first() - - assert employed is not None - assert employed.value == CPS_MONTHLY_DATA[2023][12]["employed"] - assert employed.target_type == TargetType.COUNT - assert employed.source == DataSource.BLS - - def test_load_creates_unemployment_count(self, temp_db): - """Loading CPS should create unemployment count target.""" - from db.etl_cps import load_cps_targets, CPS_MONTHLY_DATA - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - unemployed = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "unemployed") - .where(Target.period == 202312) - ).first() - - assert unemployed is not None - assert unemployed.value == CPS_MONTHLY_DATA[2023][12]["unemployed"] - assert unemployed.target_type == TargetType.COUNT - - def test_load_creates_unemployment_rate(self, temp_db): - """Loading CPS should create unemployment rate target.""" - from db.etl_cps import load_cps_targets, CPS_MONTHLY_DATA - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - rate = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "unemployment_rate") - .where(Target.period == 202312) - ).first() - - assert rate is not None - assert rate.value == CPS_MONTHLY_DATA[2023][12]["unemployment_rate"] - assert rate.target_type == TargetType.RATE - - def test_load_creates_labor_force_participation(self, temp_db): - """Loading CPS should create LFPR target.""" - from db.etl_cps import load_cps_targets, CPS_MONTHLY_DATA - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - lfpr = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "labor_force_participation_rate") - .where(Target.period == 202312) - ).first() - - assert lfpr is not None - assert lfpr.value == CPS_MONTHLY_DATA[2023][12]["labor_force_participation_rate"] - assert lfpr.target_type == TargetType.RATE - - def test_load_creates_employment_population_ratio(self, temp_db): - """Loading CPS should create employment-population ratio target.""" - from db.etl_cps import load_cps_targets, CPS_MONTHLY_DATA - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - ratio = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employment_population_ratio") - .where(Target.period == 202312) - ).first() - - assert ratio is not None - assert ratio.value == CPS_MONTHLY_DATA[2023][12]["employment_population_ratio"] - assert ratio.target_type == TargetType.RATE - - def test_load_creates_not_in_labor_force(self, temp_db): - """Loading CPS should create not-in-labor-force count target.""" - from db.etl_cps import load_cps_targets, CPS_MONTHLY_DATA - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - nilf = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "not_in_labor_force") - .where(Target.period == 202312) - ).first() - - assert nilf is not None - assert nilf.value == CPS_MONTHLY_DATA[2023][12]["not_in_labor_force"] - assert nilf.target_type == TargetType.COUNT - - def test_load_multiple_months(self, temp_db): - """Loading multiple months should create targets for each.""" - from db.etl_cps import load_cps_targets - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[1, 6, 12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - ).all() - - periods = {t.period for t in targets} - # 202301 (January), 202306 (June), 202312 (December) - assert periods == {202301, 202306, 202312} - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - from db.etl_cps import load_cps_targets - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023, 2024], months=[12]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - ).all() - - periods = {t.period for t in targets} - # 202312 (Dec 2023) and 202412 (Dec 2024) - assert periods == {202312, 202412} - - def test_load_all_data(self, temp_db): - """Loading without filters should load all available data.""" - from db.etl_cps import load_cps_targets - - with Session(temp_db) as session: - load_cps_targets(session) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - ).all() - - # Should have data for multiple months across multiple years - assert len(targets) > 5 - - # Verify we have data from 2023, 2024, and 2025 - years = {t.period // 100 for t in targets} - assert 2023 in years - assert 2024 in years - assert 2025 in years - - def test_load_idempotent(self, temp_db): - """Loading twice should not duplicate strata.""" - from db.etl_cps import load_cps_targets - - with Session(temp_db) as session: - load_cps_targets(session, years=[2023], months=[12]) - load_cps_targets(session, years=[2023], months=[12]) - - labor_force_strata = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).all() - - assert len(labor_force_strata) == 1 - - def test_period_format(self, temp_db): - """Period should be in YYYYMM format.""" - from db.etl_cps import load_cps_targets - - with Session(temp_db) as session: - load_cps_targets(session, years=[2024], months=[11]) - - labor_force = session.exec( - select(Stratum).where(Stratum.name == "US Labor Force") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == labor_force.id) - .where(Target.variable == "employed") - ).first() - - # November 2024 should be 202411 - assert target.period == 202411 diff --git a/tests/test_etl_hmrc.py b/tests/test_etl_hmrc.py deleted file mode 100644 index 84f6ab94..00000000 --- a/tests/test_etl_hmrc.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Tests for UK HMRC ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_hmrc.db" - engine = init_db(db_path) - yield engine - - -class TestHmrcETL: - """Tests for HMRC ETL loader.""" - - def test_load_hmrc_creates_national_stratum(self, temp_db): - """Loading HMRC data should create a UK national stratum.""" - from db.etl_hmrc import load_hmrc_targets - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - - national = session.exec( - select(Stratum).where(Stratum.name == "UK All Taxpayers") - ).first() - - assert national is not None - assert national.jurisdiction == Jurisdiction.UK - - def test_load_hmrc_creates_income_tax_targets(self, temp_db): - """Loading HMRC should create income tax revenue targets.""" - from db.etl_hmrc import load_hmrc_targets, HMRC_DATA - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - - national = session.exec( - select(Stratum).where(Stratum.name == "UK All Taxpayers") - ).first() - - income_tax = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "income_tax") - .where(Target.period == 2022) - ).first() - - assert income_tax is not None - assert income_tax.value == HMRC_DATA[2022]["income_tax"] - assert income_tax.target_type == TargetType.AMOUNT - assert income_tax.source == DataSource.HMRC - - def test_load_hmrc_creates_ni_targets(self, temp_db): - """Loading HMRC should create National Insurance targets.""" - from db.etl_hmrc import load_hmrc_targets, HMRC_DATA - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - - national = session.exec( - select(Stratum).where(Stratum.name == "UK All Taxpayers") - ).first() - - ni = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "national_insurance") - .where(Target.period == 2022) - ).first() - - assert ni is not None - assert ni.value == HMRC_DATA[2022]["national_insurance"] - - def test_load_hmrc_creates_taxpayer_count(self, temp_db): - """Loading HMRC should create taxpayer count targets.""" - from db.etl_hmrc import load_hmrc_targets, HMRC_DATA - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - - national = session.exec( - select(Stratum).where(Stratum.name == "UK All Taxpayers") - ).first() - - count = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "taxpayer_count") - .where(Target.period == 2022) - ).first() - - assert count is not None - assert count.value == HMRC_DATA[2022]["taxpayers"] - assert count.target_type == TargetType.COUNT - - def test_load_hmrc_creates_benefit_recipients(self, temp_db): - """Loading HMRC should create benefit recipient targets.""" - from db.etl_hmrc import load_hmrc_targets - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - - uc_stratum = session.exec( - select(Stratum).where(Stratum.name == "UK Universal Credit Recipients") - ).first() - - assert uc_stratum is not None - - uc_count = session.exec( - select(Target) - .where(Target.stratum_id == uc_stratum.id) - .where(Target.variable == "universal_credit_recipients") - ).first() - - assert uc_count is not None - assert uc_count.target_type == TargetType.COUNT - - def test_load_hmrc_creates_benefit_expenditure(self, temp_db): - """Loading HMRC should create benefit expenditure targets.""" - from db.etl_hmrc import load_hmrc_targets, HMRC_DATA - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - - uc_stratum = session.exec( - select(Stratum).where(Stratum.name == "UK Universal Credit Recipients") - ).first() - - uc_spend = session.exec( - select(Target) - .where(Target.stratum_id == uc_stratum.id) - .where(Target.variable == "universal_credit_expenditure") - ).first() - - assert uc_spend is not None - assert uc_spend.value == HMRC_DATA[2022]["benefits"]["universal_credit"]["expenditure"] - assert uc_spend.target_type == TargetType.AMOUNT - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - from db.etl_hmrc import load_hmrc_targets - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2021, 2022]) - - national = session.exec( - select(Stratum).where(Stratum.name == "UK All Taxpayers") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "income_tax") - ).all() - - years = {t.period for t in targets} - assert years == {2021, 2022} - - def test_load_hmrc_idempotent(self, temp_db): - """Loading HMRC twice should not duplicate data.""" - from db.etl_hmrc import load_hmrc_targets - - with Session(temp_db) as session: - load_hmrc_targets(session, years=[2022]) - load_hmrc_targets(session, years=[2022]) - - national_strata = session.exec( - select(Stratum).where(Stratum.name == "UK All Taxpayers") - ).all() - - assert len(national_strata) == 1 diff --git a/tests/test_etl_medicaid.py b/tests/test_etl_medicaid.py deleted file mode 100644 index 56b172d0..00000000 --- a/tests/test_etl_medicaid.py +++ /dev/null @@ -1,314 +0,0 @@ -"""Tests for Medicaid enrollment ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_medicaid import load_medicaid_targets, MEDICAID_DATA - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_medicaid.db" - engine = init_db(db_path) - yield engine - - -class TestMedicaidETL: - """Tests for Medicaid enrollment ETL loader.""" - - def test_load_medicaid_creates_national_stratum(self, temp_db): - """Loading Medicaid data should create a national stratum.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).first() - - assert national is not None - assert national.stratum_group_id == "medicaid_national" - - def test_load_medicaid_creates_total_enrollment_target(self, temp_db): - """Loading Medicaid data should create total enrollment target.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).first() - - total_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "medicaid_total_enrollment") - .where(Target.period == 2023) - ).first() - - assert total_target is not None - expected = MEDICAID_DATA[2023]["national"]["total_enrollment"] - assert total_target.value == expected - assert total_target.target_type == TargetType.COUNT - assert total_target.source == DataSource.CMS_MEDICAID - - def test_load_medicaid_creates_chip_stratum(self, temp_db): - """Loading Medicaid data should create CHIP stratum.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - chip = session.exec( - select(Stratum).where(Stratum.name == "US CHIP Enrollees") - ).first() - - assert chip is not None - assert chip.stratum_group_id == "medicaid_chip" - - def test_load_medicaid_creates_chip_enrollment_target(self, temp_db): - """Loading Medicaid data should create CHIP enrollment target.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - chip = session.exec( - select(Stratum).where(Stratum.name == "US CHIP Enrollees") - ).first() - - chip_target = session.exec( - select(Target) - .where(Target.stratum_id == chip.id) - .where(Target.variable == "chip_enrollment") - .where(Target.period == 2023) - ).first() - - assert chip_target is not None - expected = MEDICAID_DATA[2023]["national"]["chip_enrollment"] - assert chip_target.value == expected - - def test_load_medicaid_creates_child_enrollment_target(self, temp_db): - """Loading Medicaid data should create child enrollment target.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - child_stratum = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Children") - ).first() - - assert child_stratum is not None - - child_target = session.exec( - select(Target) - .where(Target.stratum_id == child_stratum.id) - .where(Target.variable == "medicaid_child_enrollment") - ).first() - - assert child_target is not None - expected = MEDICAID_DATA[2023]["national"]["children"] - assert child_target.value == expected - - def test_load_medicaid_creates_adult_enrollment_target(self, temp_db): - """Loading Medicaid data should create adult enrollment target.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - adult_stratum = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Adults") - ).first() - - assert adult_stratum is not None - - adult_target = session.exec( - select(Target) - .where(Target.stratum_id == adult_stratum.id) - .where(Target.variable == "medicaid_adult_enrollment") - ).first() - - assert adult_target is not None - expected = MEDICAID_DATA[2023]["national"]["adults"] - assert adult_target.value == expected - - def test_load_medicaid_creates_aged_enrollment_target(self, temp_db): - """Loading Medicaid data should create aged enrollment target for years with data.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - aged_stratum = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Aged") - ).first() - - assert aged_stratum is not None - - aged_target = session.exec( - select(Target) - .where(Target.stratum_id == aged_stratum.id) - .where(Target.variable == "medicaid_aged_enrollment") - ).first() - - assert aged_target is not None - expected = MEDICAID_DATA[2023]["national"]["aged"] - assert aged_target.value == expected - - def test_load_medicaid_creates_disabled_enrollment_target(self, temp_db): - """Loading Medicaid data should create disabled enrollment target.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - disabled_stratum = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Disabled") - ).first() - - assert disabled_stratum is not None - - disabled_target = session.exec( - select(Target) - .where(Target.stratum_id == disabled_stratum.id) - .where(Target.variable == "medicaid_disabled_enrollment") - ).first() - - assert disabled_target is not None - expected = MEDICAID_DATA[2023]["national"]["disabled"] - assert disabled_target.value == expected - - def test_load_medicaid_creates_aca_expansion_target(self, temp_db): - """Loading Medicaid data should create ACA expansion enrollment target.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - expansion_stratum = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid ACA Expansion Adults") - ).first() - - assert expansion_stratum is not None - - expansion_target = session.exec( - select(Target) - .where(Target.stratum_id == expansion_stratum.id) - .where(Target.variable == "medicaid_aca_expansion_enrollment") - ).first() - - assert expansion_target is not None - expected = MEDICAID_DATA[2023]["national"]["aca_expansion_adults"] - assert expansion_target.value == expected - - def test_load_medicaid_creates_state_strata(self, temp_db): - """Loading Medicaid data should create state-level strata.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "medicaid_states") - ).all() - - # Should have strata for states in the data - expected_states = len(MEDICAID_DATA[2023].get("states", {})) - assert len(state_strata) == expected_states - - def test_load_medicaid_state_targets_correct(self, temp_db): - """State-level Medicaid targets should have correct values.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA Medicaid Enrollees") - ).first() - - assert ca_stratum is not None - - ca_enrollment = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "medicaid_total_enrollment") - ).first() - - expected_ca = MEDICAID_DATA[2023]["states"]["CA"]["total"] - assert ca_enrollment.value == expected_ca - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2021, 2022, 2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "medicaid_total_enrollment") - ).all() - - years = {t.period for t in targets} - assert years == {2021, 2022, 2023} - - def test_load_medicaid_idempotent(self, temp_db): - """Loading Medicaid twice should not duplicate data.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - load_medicaid_targets(session, years=[2023]) - - national_strata = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).all() - - # Should only have one national stratum - assert len(national_strata) == 1 - - def test_state_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).first() - - ca = session.exec( - select(Stratum).where(Stratum.name == "CA Medicaid Enrollees") - ).first() - - assert ca.parent_id == national.id - - def test_chip_stratum_has_parent(self, temp_db): - """CHIP stratum should have national Medicaid stratum as parent.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).first() - - chip = session.exec( - select(Stratum).where(Stratum.name == "US CHIP Enrollees") - ).first() - - assert chip.parent_id == national.id - - def test_load_2025_data(self, temp_db): - """Loading 2025 data should work with latest enrollment figures.""" - with Session(temp_db) as session: - load_medicaid_targets(session, years=[2025]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US Medicaid Enrollees") - ).first() - - total_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "medicaid_total_enrollment") - .where(Target.period == 2025) - ).first() - - assert total_target is not None - # August 2025 figure from CMS - assert total_target.value == 77_290_050 diff --git a/tests/test_etl_obr.py b/tests/test_etl_obr.py deleted file mode 100644 index 3bf290f5..00000000 --- a/tests/test_etl_obr.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Tests for OBR (Office for Budget Responsibility) projections ETL.""" - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_obr import load_obr_targets, OBR_DATA - - -class TestObrETL: - """Tests for OBR projections ETL.""" - - @pytest.fixture - def session(self, tmp_path): - """Create a test database session.""" - db_path = tmp_path / "test.db" - engine = init_db(db_path) - with Session(engine) as session: - yield session - - def test_load_creates_uk_budget_stratum(self, session): - """Should create UK budget stratum.""" - load_obr_targets(session, years=[2024]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "UK Public Finances") - ).first() - - assert stratum is not None - assert stratum.jurisdiction == Jurisdiction.UK - - def test_load_creates_gdp_projection(self, session): - """Should create UK GDP projection targets.""" - load_obr_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "gdp", - Target.period == 2024, - Target.source == DataSource.OBR, - ) - ).first() - - assert target is not None - assert target.value == OBR_DATA[2024]["gdp"] - assert target.target_type == TargetType.AMOUNT - - def test_load_creates_receipts_projection(self, session): - """Should create total receipts projection.""" - load_obr_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "total_receipts", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.value == OBR_DATA[2024]["total_receipts"] - - def test_load_creates_expenditure_projection(self, session): - """Should create total managed expenditure projection.""" - load_obr_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "total_managed_expenditure", - Target.period == 2024, - ) - ).first() - - assert target is not None - assert target.value == OBR_DATA[2024]["total_managed_expenditure"] - - def test_load_creates_borrowing_projection(self, session): - """Should create public sector net borrowing projection.""" - load_obr_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "public_sector_net_borrowing", - Target.period == 2024, - ) - ).first() - - assert target is not None - - def test_load_creates_economy_stratum(self, session): - """Should create UK economy stratum for macro indicators.""" - load_obr_targets(session, years=[2024]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "UK Economy") - ).first() - - assert stratum is not None - assert stratum.jurisdiction == Jurisdiction.UK - - def test_load_creates_unemployment_projection(self, session): - """Should create UK unemployment rate projection.""" - load_obr_targets(session, years=[2025]) - - target = session.exec( - select(Target).where( - Target.variable == "unemployment_rate", - Target.period == 2025, - Target.source == DataSource.OBR, - ) - ).first() - - assert target is not None - assert target.target_type == TargetType.RATE - - def test_load_creates_inflation_projection(self, session): - """Should create CPI inflation projection.""" - load_obr_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "cpi_inflation", - Target.period == 2024, - Target.source == DataSource.OBR, - ) - ).first() - - assert target is not None - assert target.target_type == TargetType.RATE - - def test_load_future_years(self, session): - """Should load projection years through 2029.""" - load_obr_targets(session, years=[2028, 2029]) - - targets_2028 = session.exec( - select(Target).where( - Target.period == 2028, - Target.source == DataSource.OBR, - ) - ).all() - targets_2029 = session.exec( - select(Target).where( - Target.period == 2029, - Target.source == DataSource.OBR, - ) - ).all() - - assert len(targets_2028) > 0 - assert len(targets_2029) > 0 - - def test_load_multiple_years(self, session): - """Should load multiple projection years.""" - load_obr_targets(session, years=[2024, 2025, 2026]) - - targets = session.exec( - select(Target).where(Target.source == DataSource.OBR) - ).all() - periods = {t.period for t in targets} - - assert 2024 in periods - assert 2025 in periods - assert 2026 in periods - - def test_load_idempotent(self, session): - """Loading twice should not duplicate data.""" - load_obr_targets(session, years=[2024]) - count1 = len( - session.exec( - select(Target).where(Target.source == DataSource.OBR) - ).all() - ) - - load_obr_targets(session, years=[2024]) - count2 = len( - session.exec( - select(Target).where(Target.source == DataSource.OBR) - ).all() - ) - - assert count1 == count2 diff --git a/tests/test_etl_ons.py b/tests/test_etl_ons.py deleted file mode 100644 index be311e78..00000000 --- a/tests/test_etl_ons.py +++ /dev/null @@ -1,209 +0,0 @@ -"""Tests for ONS (Office for National Statistics) projections ETL.""" - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_ons import load_ons_targets, ONS_DATA - - -class TestOnsETL: - """Tests for ONS projections ETL.""" - - @pytest.fixture - def session(self, tmp_path): - """Create a test database session.""" - db_path = tmp_path / "test.db" - engine = init_db(db_path) - with Session(engine) as session: - yield session - - def test_load_creates_uk_population_stratum(self, session): - """Should create UK population stratum.""" - load_ons_targets(session, years=[2024]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "UK Population") - ).first() - - assert stratum is not None - assert stratum.jurisdiction == Jurisdiction.UK - - def test_load_creates_total_population_projection(self, session): - """Should create UK total population projection targets.""" - load_ons_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "population_total", - Target.period == 2024, - Target.source == DataSource.ONS, - ) - ).first() - - assert target is not None - assert target.value == ONS_DATA[2024]["population_total"] - assert target.target_type == TargetType.COUNT - assert target.is_preliminary is True - - def test_load_creates_age_group_projections(self, session): - """Should create population projections by age groups.""" - load_ons_targets(session, years=[2024]) - - # Check 0-15 age group - target_0_15 = session.exec( - select(Target).where( - Target.variable == "population_age_0_15", - Target.period == 2024, - ) - ).first() - - assert target_0_15 is not None - assert target_0_15.value == ONS_DATA[2024]["population_age_0_15"] - - # Check 16-64 age group - target_16_64 = session.exec( - select(Target).where( - Target.variable == "population_age_16_64", - Target.period == 2024, - ) - ).first() - - assert target_16_64 is not None - assert target_16_64.value == ONS_DATA[2024]["population_age_16_64"] - - # Check 65+ age group - target_65_plus = session.exec( - select(Target).where( - Target.variable == "population_age_65_plus", - Target.period == 2024, - ) - ).first() - - assert target_65_plus is not None - assert target_65_plus.value == ONS_DATA[2024]["population_age_65_plus"] - - def test_load_creates_uk_households_stratum(self, session): - """Should create UK households stratum.""" - load_ons_targets(session, years=[2024]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "UK Households") - ).first() - - assert stratum is not None - assert stratum.jurisdiction == Jurisdiction.UK - - def test_load_creates_household_projections(self, session): - """Should create household count projection.""" - load_ons_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "households_total", - Target.period == 2024, - Target.source == DataSource.ONS, - ) - ).first() - - assert target is not None - assert target.value == ONS_DATA[2024]["households_total"] - assert target.target_type == TargetType.COUNT - - def test_load_creates_average_household_size(self, session): - """Should create average household size projection.""" - load_ons_targets(session, years=[2024]) - - target = session.exec( - select(Target).where( - Target.variable == "average_household_size", - Target.period == 2024, - Target.source == DataSource.ONS, - ) - ).first() - - assert target is not None - assert target.value == ONS_DATA[2024]["average_household_size"] - assert target.target_type == TargetType.RATE - - def test_load_future_years(self, session): - """Should load projection years through 2034.""" - load_ons_targets(session, years=[2033, 2034]) - - targets_2033 = session.exec( - select(Target).where( - Target.period == 2033, - Target.source == DataSource.ONS, - ) - ).all() - targets_2034 = session.exec( - select(Target).where( - Target.period == 2034, - Target.source == DataSource.ONS, - ) - ).all() - - assert len(targets_2033) > 0 - assert len(targets_2034) > 0 - - def test_load_multiple_years(self, session): - """Should load multiple projection years.""" - load_ons_targets(session, years=[2024, 2025, 2026]) - - targets = session.exec( - select(Target).where(Target.source == DataSource.ONS) - ).all() - periods = {t.period for t in targets} - - assert 2024 in periods - assert 2025 in periods - assert 2026 in periods - - def test_load_idempotent(self, session): - """Loading twice should not duplicate data.""" - load_ons_targets(session, years=[2024]) - count1 = len( - session.exec( - select(Target).where(Target.source == DataSource.ONS) - ).all() - ) - - load_ons_targets(session, years=[2024]) - count2 = len( - session.exec( - select(Target).where(Target.source == DataSource.ONS) - ).all() - ) - - assert count1 == count2 - - def test_all_projections_marked_preliminary(self, session): - """All ONS projections should be marked as preliminary.""" - load_ons_targets(session, years=[2024, 2030]) - - targets = session.exec( - select(Target).where(Target.source == DataSource.ONS) - ).all() - - for target in targets: - assert target.is_preliminary is True - - def test_load_all_years(self, session): - """Should load all available years when years=None.""" - load_ons_targets(session, years=None) - - targets = session.exec( - select(Target).where(Target.source == DataSource.ONS) - ).all() - periods = {t.period for t in targets} - - # Should have data for all years 2024-2034 - expected_years = set(range(2024, 2035)) - assert periods == expected_years diff --git a/tests/test_etl_snap.py b/tests/test_etl_snap.py deleted file mode 100644 index 4f888cc5..00000000 --- a/tests/test_etl_snap.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Tests for SNAP ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_snap import load_snap_targets, SNAP_DATA - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_snap.db" - engine = init_db(db_path) - yield engine - - -class TestSnapETL: - """Tests for SNAP ETL loader.""" - - def test_load_snap_creates_national_stratum(self, temp_db): - """Loading SNAP data should create a national stratum.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).first() - - assert national is not None - assert national.stratum_group_id == "snap_national" - - def test_load_snap_creates_national_targets(self, temp_db): - """Loading SNAP data should create national-level targets.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).first() - - # Check household count - hh_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "snap_household_count") - .where(Target.period == 2023) - ).first() - - assert hh_target is not None - expected_hh = SNAP_DATA[2023]["national"]["households"] * 1000 - assert hh_target.value == expected_hh - assert hh_target.target_type == TargetType.COUNT - assert hh_target.source == DataSource.USDA_SNAP - assert "convert_units" in hh_target.notes - - def test_load_snap_creates_participant_targets(self, temp_db): - """Loading SNAP should create participant count targets.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).first() - - participant_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "snap_participant_count") - .where(Target.period == 2023) - ).first() - - assert participant_target is not None - expected = SNAP_DATA[2023]["national"]["participants"] * 1000 - assert participant_target.value == expected - - def test_load_snap_creates_benefit_targets(self, temp_db): - """Loading SNAP should create benefit amount targets.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).first() - - benefit_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "snap_benefits") - .where(Target.period == 2023) - ).first() - - assert benefit_target is not None - expected = SNAP_DATA[2023]["national"]["benefits"] * 1_000_000 - assert benefit_target.value == expected - assert benefit_target.target_type == TargetType.AMOUNT - assert "convert_units" in benefit_target.notes - - def test_load_snap_creates_state_strata(self, temp_db): - """Loading SNAP should create state-level strata.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "snap_states") - ).all() - - # Should have strata for states in the data - expected_states = len(SNAP_DATA[2023].get("states", {})) - assert len(state_strata) == expected_states - - def test_load_snap_state_targets_correct(self, temp_db): - """State-level SNAP targets should have correct values.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA SNAP Recipients") - ).first() - - assert ca_stratum is not None - - ca_hh = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "snap_household_count") - ).first() - - expected_ca_hh = SNAP_DATA[2023]["states"]["CA"]["households"] * 1000 - assert ca_hh.value == expected_ca_hh - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2021, 2022, 2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "snap_household_count") - ).all() - - years = {t.period for t in targets} - assert years == {2021, 2022, 2023} - - def test_load_snap_idempotent(self, temp_db): - """Loading SNAP twice should not duplicate data.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - load_snap_targets(session, years=[2023]) - - national_strata = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).all() - - # Should only have one national stratum - assert len(national_strata) == 1 - - def test_state_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_snap_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SNAP Recipients") - ).first() - - ca = session.exec( - select(Stratum).where(Stratum.name == "CA SNAP Recipients") - ).first() - - assert ca.parent_id == national.id diff --git a/tests/test_etl_soi.py b/tests/test_etl_soi.py deleted file mode 100644 index 4fbbf364..00000000 --- a/tests/test_etl_soi.py +++ /dev/null @@ -1,36 +0,0 @@ -from db.etl_soi import ( - available_soi_table_1_4_years, - available_soi_years, - load_soi_table_1_1_data, - load_soi_table_1_4_data, -) - - -def test_available_soi_years_include_latest_table_1_1_sources(): - assert available_soi_years() == [2020, 2021, 2022, 2023] - - -def test_load_soi_table_1_1_data_reads_packaged_2023_source(): - data = load_soi_table_1_1_data(2023) - - assert data["source_url"] == "https://www.irs.gov/pub/irs-soi/23in11si.xls" - assert data["total_returns"] == 160_602_107 - assert data["total_agi"] == 15_286_017_359_000 - assert data["total_income_tax"] == 2_147_909_818_000 - assert data["returns_by_agi_bracket"]["100k_to_200k"] == 27_602_755 - assert data["agi_by_bracket"]["100k_to_200k"] == 3_818_295_141_000 - assert data["income_tax_by_bracket"]["100k_to_200k"] == 409_532_689_000 - - -def test_available_soi_years_include_latest_table_1_4_sources(): - assert available_soi_table_1_4_years() == [2021, 2022, 2023] - - -def test_load_soi_table_1_4_data_reads_packaged_2023_wage_source(): - data = load_soi_table_1_4_data(2023) - - assert data["source_url"] == "https://www.irs.gov/pub/irs-soi/23in14ar.xls" - assert data["total_employment_income_returns"] == 128_591_050 - assert data["total_employment_income"] == 10_204_095_705_000 - assert data["employment_income_returns_by_agi_bracket"]["100k_to_200k"] == 23_193_910 - assert data["employment_income_by_agi_bracket"]["100k_to_200k"] == 2_774_550_975_000 diff --git a/tests/test_etl_soi_credits.py b/tests/test_etl_soi_credits.py deleted file mode 100644 index 569f6656..00000000 --- a/tests/test_etl_soi_credits.py +++ /dev/null @@ -1,848 +0,0 @@ -"""Tests for state-level SOI tax credits ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_soi_credits import ( - load_soi_credits_targets, - load_eitc_by_children_targets, - load_ctc_by_children_targets, - load_actc_by_children_targets, - SOI_CREDITS_DATA, - EITC_BY_CHILDREN_DATA, - CTC_BY_CHILDREN_DATA, - ACTC_BY_CHILDREN_DATA, - STATE_FIPS, - SOURCE_URLS, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_soi_credits.db" - engine = init_db(db_path) - yield engine - - -class TestSoiCreditsETL: - """Tests for state-level SOI credits ETL loader.""" - - def test_load_soi_credits_creates_national_stratum(self, temp_db): - """Loading credit data should create/reference a national stratum.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - assert national is not None - assert national.stratum_group_id == "national" - - def test_load_soi_credits_creates_state_strata(self, temp_db): - """Loading credit data should create state-level strata.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "soi_states") - ).all() - - # Should have strata for all 50 states + DC - expected_states = len(STATE_FIPS) - assert len(state_strata) == expected_states - assert expected_states == 51 # 50 states + DC - - def test_load_eitc_claims_targets(self, temp_db): - """Loading credits should create EITC claims count targets.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - assert ca_stratum is not None - - eitc_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "eitc_claims") - .where(Target.period == 2021) - ).first() - - assert eitc_claims_target is not None - expected = SOI_CREDITS_DATA[2021]["CA"]["eitc_claims"] - assert eitc_claims_target.value == expected - assert eitc_claims_target.target_type == TargetType.COUNT - assert eitc_claims_target.source == DataSource.IRS_SOI - - def test_load_eitc_amount_targets(self, temp_db): - """Loading credits should create EITC amount targets.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - tx_stratum = session.exec( - select(Stratum).where(Stratum.name == "TX All Filers") - ).first() - - eitc_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == tx_stratum.id) - .where(Target.variable == "eitc_amount") - .where(Target.period == 2021) - ).first() - - assert eitc_amount_target is not None - expected = SOI_CREDITS_DATA[2021]["TX"]["eitc_amount"] - assert eitc_amount_target.value == expected - assert eitc_amount_target.target_type == TargetType.AMOUNT - assert eitc_amount_target.source_url == SOURCE_URLS["eitc"] - - def test_load_ctc_claims_targets(self, temp_db): - """Loading credits should create CTC claims count targets.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - fl_stratum = session.exec( - select(Stratum).where(Stratum.name == "FL All Filers") - ).first() - - ctc_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == fl_stratum.id) - .where(Target.variable == "ctc_claims") - .where(Target.period == 2021) - ).first() - - assert ctc_claims_target is not None - expected = SOI_CREDITS_DATA[2021]["FL"]["ctc_claims"] - assert ctc_claims_target.value == expected - assert ctc_claims_target.target_type == TargetType.COUNT - - def test_load_ctc_amount_targets(self, temp_db): - """Loading credits should create CTC amount targets.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - ny_stratum = session.exec( - select(Stratum).where(Stratum.name == "NY All Filers") - ).first() - - ctc_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == ny_stratum.id) - .where(Target.variable == "ctc_amount") - .where(Target.period == 2021) - ).first() - - assert ctc_amount_target is not None - expected = SOI_CREDITS_DATA[2021]["NY"]["ctc_amount"] - assert ctc_amount_target.value == expected - assert ctc_amount_target.target_type == TargetType.AMOUNT - - def test_load_actc_claims_targets(self, temp_db): - """Loading credits should create ACTC claims count targets.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - ga_stratum = session.exec( - select(Stratum).where(Stratum.name == "GA All Filers") - ).first() - - actc_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == ga_stratum.id) - .where(Target.variable == "actc_claims") - .where(Target.period == 2021) - ).first() - - assert actc_claims_target is not None - expected = SOI_CREDITS_DATA[2021]["GA"]["actc_claims"] - assert actc_claims_target.value == expected - assert actc_claims_target.target_type == TargetType.COUNT - - def test_load_actc_amount_targets(self, temp_db): - """Loading credits should create ACTC amount targets.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - pa_stratum = session.exec( - select(Stratum).where(Stratum.name == "PA All Filers") - ).first() - - actc_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == pa_stratum.id) - .where(Target.variable == "actc_amount") - .where(Target.period == 2021) - ).first() - - assert actc_amount_target is not None - expected = SOI_CREDITS_DATA[2021]["PA"]["actc_amount"] - assert actc_amount_target.value == expected - assert actc_amount_target.target_type == TargetType.AMOUNT - - def test_load_soi_credits_stratum_has_state_constraint(self, temp_db): - """State strata should have state_fips constraint.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - oh_stratum = session.exec( - select(Stratum).where(Stratum.name == "OH All Filers") - ).first() - - # Check constraints include state_fips - state_constraint = None - for constraint in oh_stratum.constraints: - if constraint.variable == "state_fips": - state_constraint = constraint - break - - assert state_constraint is not None - assert state_constraint.operator == "==" - assert state_constraint.value == STATE_FIPS["OH"] - - def test_load_soi_credits_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - il_stratum = session.exec( - select(Stratum).where(Stratum.name == "IL All Filers") - ).first() - - assert il_stratum.parent_id == national.id - - def test_load_soi_credits_idempotent(self, temp_db): - """Loading credits twice should not duplicate data.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - load_soi_credits_targets(session, years=[2021]) - - ca_strata = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).all() - - # Should only have one CA stratum - assert len(ca_strata) == 1 - - def test_all_states_loaded(self, temp_db): - """All 50 states + DC should be loaded with 6 credit targets each.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - for state_abbrev in STATE_FIPS.keys(): - stratum = session.exec( - select(Stratum).where(Stratum.name == f"{state_abbrev} All Filers") - ).first() - - assert stratum is not None, f"Missing stratum for {state_abbrev}" - - # Each state should have 6 targets (eitc_claims, eitc_amount, - # ctc_claims, ctc_amount, actc_claims, actc_amount) - targets = session.exec( - select(Target).where(Target.stratum_id == stratum.id) - ).all() - - assert len(targets) == 6, f"Expected 6 targets for {state_abbrev}, got {len(targets)}" - - def test_target_source_metadata(self, temp_db): - """Targets should have correct source metadata.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - eitc_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "eitc_claims") - ).first() - - assert eitc_target.source == DataSource.IRS_SOI - assert eitc_target.source_table == "EITC Statistics" - assert "earned-income-tax-credit" in eitc_target.source_url - - ctc_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "ctc_claims") - ).first() - - assert ctc_target.source == DataSource.IRS_SOI - assert ctc_target.source_table == "State Data FY" - assert "state-data" in ctc_target.source_url - - def test_national_totals_are_reasonable(self, temp_db): - """National totals should match expected ranges.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - # Sum all EITC claims across states - all_eitc_claims = sum( - state_data["eitc_claims"] - for state_data in SOI_CREDITS_DATA[2021].values() - ) - # Should be around 25M - assert 20_000_000 < all_eitc_claims < 30_000_000 - - # Sum all EITC amounts across states - all_eitc_amount = sum( - state_data["eitc_amount"] - for state_data in SOI_CREDITS_DATA[2021].values() - ) - # Should be around $60B (actual was ~$57B in 2021, allowing margin) - assert 50_000_000_000 < all_eitc_amount < 75_000_000_000 - - # Sum all CTC claims across states - all_ctc_claims = sum( - state_data["ctc_claims"] - for state_data in SOI_CREDITS_DATA[2021].values() - ) - # Should be around 35M - assert 30_000_000 < all_ctc_claims < 45_000_000 - - # Sum all ACTC amounts across states - all_actc_amount = sum( - state_data["actc_amount"] - for state_data in SOI_CREDITS_DATA[2021].values() - ) - # Should be around $34B - assert 25_000_000_000 < all_actc_amount < 40_000_000_000 - - def test_large_states_have_higher_claims(self, temp_db): - """Larger states should have more claims than smaller states.""" - with Session(temp_db) as session: - load_soi_credits_targets(session, years=[2021]) - - # California (large) vs Wyoming (small) - ca_data = SOI_CREDITS_DATA[2021]["CA"] - wy_data = SOI_CREDITS_DATA[2021]["WY"] - - assert ca_data["eitc_claims"] > wy_data["eitc_claims"] - assert ca_data["ctc_claims"] > wy_data["ctc_claims"] - assert ca_data["eitc_amount"] > wy_data["eitc_amount"] - - # Texas (large) vs Vermont (small) - tx_data = SOI_CREDITS_DATA[2021]["TX"] - vt_data = SOI_CREDITS_DATA[2021]["VT"] - - assert tx_data["eitc_claims"] > vt_data["eitc_claims"] - assert tx_data["ctc_claims"] > vt_data["ctc_claims"] - - -class TestEitcByChildrenETL: - """Tests for EITC by number of children ETL loader.""" - - def test_load_eitc_by_children_creates_strata(self, temp_db): - """Loading EITC by children should create child-count strata.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - # Check for all 4 child-count strata - strata_names = [ - "US EITC 0 Children", - "US EITC 1 Child", - "US EITC 2 Children", - "US EITC 3+ Children", - ] - - for name in strata_names: - stratum = session.exec( - select(Stratum).where(Stratum.name == name) - ).first() - assert stratum is not None, f"Missing stratum: {name}" - assert stratum.stratum_group_id == "eitc_by_children" - - def test_load_eitc_by_children_creates_claims_targets(self, temp_db): - """Loading EITC by children should create claims count targets.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US EITC 0 Children") - ).first() - - claims_target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "eitc_claims") - .where(Target.period == 2021) - ).first() - - assert claims_target is not None - assert claims_target.target_type == TargetType.COUNT - assert claims_target.source == DataSource.IRS_SOI - expected = EITC_BY_CHILDREN_DATA[2021]["0_children"]["claims"] - assert claims_target.value == expected - - def test_load_eitc_by_children_creates_amount_targets(self, temp_db): - """Loading EITC by children should create amount targets.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US EITC 3+ Children") - ).first() - - amount_target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "eitc_amount") - .where(Target.period == 2021) - ).first() - - assert amount_target is not None - assert amount_target.target_type == TargetType.AMOUNT - assert amount_target.source == DataSource.IRS_SOI - expected = EITC_BY_CHILDREN_DATA[2021]["3plus_children"]["amount"] - assert amount_target.value == expected - - def test_load_eitc_by_children_correct_constraints(self, temp_db): - """Child-count strata should have qualifying_children constraints.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - # Check 0 children stratum - stratum_0 = session.exec( - select(Stratum).where(Stratum.name == "US EITC 0 Children") - ).first() - - child_constraint = None - for constraint in stratum_0.constraints: - if constraint.variable == "eitc_qualifying_children": - child_constraint = constraint - break - - assert child_constraint is not None - assert child_constraint.operator == "==" - assert child_constraint.value == "0" - - # Check 3+ children stratum - stratum_3plus = session.exec( - select(Stratum).where(Stratum.name == "US EITC 3+ Children") - ).first() - - child_constraint = None - for constraint in stratum_3plus.constraints: - if constraint.variable == "eitc_qualifying_children": - child_constraint = constraint - break - - assert child_constraint is not None - assert child_constraint.operator == ">=" - assert child_constraint.value == "3" - - def test_eitc_by_children_totals_reasonable(self, temp_db): - """National totals by children should be reasonable.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - data = EITC_BY_CHILDREN_DATA[2021] - - # Sum all claims - should be around 25M - total_claims = sum(d["claims"] for d in data.values()) - assert 20_000_000 < total_claims < 30_000_000 - - # Sum all amounts - should be around $57-60B - total_amount = sum(d["amount"] for d in data.values()) - assert 50_000_000_000 < total_amount < 70_000_000_000 - - # 0 children: lower average credit - avg_0_children = data["0_children"]["amount"] / data["0_children"]["claims"] - assert avg_0_children < 300 # Much lower max credit for no children - - # 3+ children: higher average credit - avg_3plus = data["3plus_children"]["amount"] / data["3plus_children"]["claims"] - assert avg_3plus > 3000 # Higher max credit for 3+ children - - def test_eitc_by_children_has_parent_stratum(self, temp_db): - """EITC by children strata should have national parent.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - stratum_1 = session.exec( - select(Stratum).where(Stratum.name == "US EITC 1 Child") - ).first() - - assert stratum_1.parent_id == national.id - - def test_eitc_by_children_idempotent(self, temp_db): - """Loading EITC by children twice should not duplicate strata.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - load_eitc_by_children_targets(session, years=[2021]) - - strata = session.exec( - select(Stratum).where(Stratum.name == "US EITC 2 Children") - ).all() - - assert len(strata) == 1 - - def test_eitc_by_children_source_metadata(self, temp_db): - """EITC by children targets should have correct source metadata.""" - with Session(temp_db) as session: - load_eitc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US EITC 1 Child") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "eitc_claims") - ).first() - - assert target.source == DataSource.IRS_SOI - assert target.source_table == "EITC Statistics by Number of Qualifying Children" - assert "earned-income-tax-credit" in target.source_url - - -class TestCtcByChildrenETL: - """Tests for CTC by number of children ETL loader.""" - - def test_load_ctc_by_children_creates_strata(self, temp_db): - """Loading CTC by children should create child-count strata.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - # Check for all 3 child-count strata (CTC requires at least 1 child) - strata_names = [ - "US CTC 1 Child", - "US CTC 2 Children", - "US CTC 3+ Children", - ] - - for name in strata_names: - stratum = session.exec( - select(Stratum).where(Stratum.name == name) - ).first() - assert stratum is not None, f"Missing stratum: {name}" - assert stratum.stratum_group_id == "ctc_by_children" - - def test_load_ctc_by_children_creates_claims_targets(self, temp_db): - """Loading CTC by children should create claims count targets.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US CTC 1 Child") - ).first() - - claims_target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "ctc_claims") - .where(Target.period == 2021) - ).first() - - assert claims_target is not None - assert claims_target.target_type == TargetType.COUNT - assert claims_target.source == DataSource.IRS_SOI - expected = CTC_BY_CHILDREN_DATA[2021]["1_child"]["claims"] - assert claims_target.value == expected - - def test_load_ctc_by_children_creates_amount_targets(self, temp_db): - """Loading CTC by children should create amount targets.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US CTC 3+ Children") - ).first() - - amount_target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "ctc_amount") - .where(Target.period == 2021) - ).first() - - assert amount_target is not None - assert amount_target.target_type == TargetType.AMOUNT - assert amount_target.source == DataSource.IRS_SOI - expected = CTC_BY_CHILDREN_DATA[2021]["3plus_children"]["amount"] - assert amount_target.value == expected - - def test_load_ctc_by_children_correct_constraints(self, temp_db): - """Child-count strata should have ctc_qualifying_children constraints.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - # Check 1 child stratum - stratum_1 = session.exec( - select(Stratum).where(Stratum.name == "US CTC 1 Child") - ).first() - - child_constraint = None - for constraint in stratum_1.constraints: - if constraint.variable == "ctc_qualifying_children": - child_constraint = constraint - break - - assert child_constraint is not None - assert child_constraint.operator == "==" - assert child_constraint.value == "1" - - # Check 3+ children stratum - stratum_3plus = session.exec( - select(Stratum).where(Stratum.name == "US CTC 3+ Children") - ).first() - - child_constraint = None - for constraint in stratum_3plus.constraints: - if constraint.variable == "ctc_qualifying_children": - child_constraint = constraint - break - - assert child_constraint is not None - assert child_constraint.operator == ">=" - assert child_constraint.value == "3" - - def test_ctc_by_children_totals_reasonable(self, temp_db): - """National CTC totals by children should be reasonable.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - data = CTC_BY_CHILDREN_DATA[2021] - - # Sum all claims - should be around 35M - total_claims = sum(d["claims"] for d in data.values()) - assert 30_000_000 < total_claims < 40_000_000 - - # Sum all amounts - should be around $100B - total_amount = sum(d["amount"] for d in data.values()) - assert 80_000_000_000 < total_amount < 120_000_000_000 - - def test_ctc_by_children_has_parent_stratum(self, temp_db): - """CTC by children strata should have national parent.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - stratum_1 = session.exec( - select(Stratum).where(Stratum.name == "US CTC 1 Child") - ).first() - - assert stratum_1.parent_id == national.id - - def test_ctc_by_children_idempotent(self, temp_db): - """Loading CTC by children twice should not duplicate strata.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - load_ctc_by_children_targets(session, years=[2021]) - - strata = session.exec( - select(Stratum).where(Stratum.name == "US CTC 2 Children") - ).all() - - assert len(strata) == 1 - - def test_ctc_by_children_source_metadata(self, temp_db): - """CTC by children targets should have correct source metadata.""" - with Session(temp_db) as session: - load_ctc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US CTC 1 Child") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "ctc_claims") - ).first() - - assert target.source == DataSource.IRS_SOI - assert target.source_table == "CTC Statistics by Number of Qualifying Children" - assert "state-data" in target.source_url - - -class TestActcByChildrenETL: - """Tests for ACTC by number of children ETL loader.""" - - def test_load_actc_by_children_creates_strata(self, temp_db): - """Loading ACTC by children should create child-count strata.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - # Check for all 3 child-count strata (ACTC requires at least 1 child) - strata_names = [ - "US ACTC 1 Child", - "US ACTC 2 Children", - "US ACTC 3+ Children", - ] - - for name in strata_names: - stratum = session.exec( - select(Stratum).where(Stratum.name == name) - ).first() - assert stratum is not None, f"Missing stratum: {name}" - assert stratum.stratum_group_id == "actc_by_children" - - def test_load_actc_by_children_creates_claims_targets(self, temp_db): - """Loading ACTC by children should create claims count targets.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 1 Child") - ).first() - - claims_target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "actc_claims") - .where(Target.period == 2021) - ).first() - - assert claims_target is not None - assert claims_target.target_type == TargetType.COUNT - assert claims_target.source == DataSource.IRS_SOI - expected = ACTC_BY_CHILDREN_DATA[2021]["1_child"]["claims"] - assert claims_target.value == expected - - def test_load_actc_by_children_creates_amount_targets(self, temp_db): - """Loading ACTC by children should create amount targets.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 3+ Children") - ).first() - - amount_target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "actc_amount") - .where(Target.period == 2021) - ).first() - - assert amount_target is not None - assert amount_target.target_type == TargetType.AMOUNT - assert amount_target.source == DataSource.IRS_SOI - expected = ACTC_BY_CHILDREN_DATA[2021]["3plus_children"]["amount"] - assert amount_target.value == expected - - def test_load_actc_by_children_correct_constraints(self, temp_db): - """Child-count strata should have actc_qualifying_children constraints.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - # Check 1 child stratum - stratum_1 = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 1 Child") - ).first() - - child_constraint = None - for constraint in stratum_1.constraints: - if constraint.variable == "actc_qualifying_children": - child_constraint = constraint - break - - assert child_constraint is not None - assert child_constraint.operator == "==" - assert child_constraint.value == "1" - - # Check 3+ children stratum - stratum_3plus = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 3+ Children") - ).first() - - child_constraint = None - for constraint in stratum_3plus.constraints: - if constraint.variable == "actc_qualifying_children": - child_constraint = constraint - break - - assert child_constraint is not None - assert child_constraint.operator == ">=" - assert child_constraint.value == "3" - - def test_actc_by_children_totals_reasonable(self, temp_db): - """National ACTC totals by children should be reasonable.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - data = ACTC_BY_CHILDREN_DATA[2021] - - # Sum all claims - should be around 18M - total_claims = sum(d["claims"] for d in data.values()) - assert 15_000_000 < total_claims < 22_000_000 - - # Sum all amounts - should be around $30B - total_amount = sum(d["amount"] for d in data.values()) - assert 25_000_000_000 < total_amount < 40_000_000_000 - - def test_actc_by_children_has_parent_stratum(self, temp_db): - """ACTC by children strata should have national parent.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - stratum_1 = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 1 Child") - ).first() - - assert stratum_1.parent_id == national.id - - def test_actc_by_children_idempotent(self, temp_db): - """Loading ACTC by children twice should not duplicate strata.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - load_actc_by_children_targets(session, years=[2021]) - - strata = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 2 Children") - ).all() - - assert len(strata) == 1 - - def test_actc_by_children_source_metadata(self, temp_db): - """ACTC by children targets should have correct source metadata.""" - with Session(temp_db) as session: - load_actc_by_children_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where(Stratum.name == "US ACTC 1 Child") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == stratum.id) - .where(Target.variable == "actc_claims") - ).first() - - assert target.source == DataSource.IRS_SOI - assert target.source_table == "ACTC Statistics by Number of Qualifying Children" - assert "state-data" in target.source_url diff --git a/tests/test_etl_soi_deductions.py b/tests/test_etl_soi_deductions.py deleted file mode 100644 index 2503892c..00000000 --- a/tests/test_etl_soi_deductions.py +++ /dev/null @@ -1,468 +0,0 @@ -"""Tests for state-level SOI deduction targets ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_soi_deductions import ( - load_soi_deductions_targets, - SOI_DEDUCTIONS_DATA, - STATE_FIPS, - SOURCE_URL, - DEDUCTION_TYPES, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_soi_deductions.db" - engine = init_db(db_path) - yield engine - - -class TestSoiDeductionsETL: - """Tests for state-level SOI deductions ETL loader.""" - - def test_load_soi_deductions_creates_national_stratum(self, temp_db): - """Loading deduction data should create/reference a national stratum.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - assert national is not None - assert national.stratum_group_id == "national" - - def test_load_soi_deductions_creates_state_strata(self, temp_db): - """Loading deduction data should create state-level strata.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "soi_states") - ).all() - - # Should have strata for all 50 states + DC - expected_states = len(STATE_FIPS) - assert len(state_strata) == expected_states - assert expected_states == 51 # 50 states + DC - - def test_load_salt_claims_targets(self, temp_db): - """Loading deductions should create SALT claims count targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - assert ca_stratum is not None - - salt_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "salt_claims") - .where(Target.period == 2021) - ).first() - - assert salt_claims_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["CA"]["salt_claims"] - assert salt_claims_target.value == expected - assert salt_claims_target.target_type == TargetType.COUNT - assert salt_claims_target.source == DataSource.IRS_SOI - - def test_load_salt_amount_targets(self, temp_db): - """Loading deductions should create SALT amount targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - tx_stratum = session.exec( - select(Stratum).where(Stratum.name == "TX All Filers") - ).first() - - salt_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == tx_stratum.id) - .where(Target.variable == "salt_amount") - .where(Target.period == 2021) - ).first() - - assert salt_amount_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["TX"]["salt_amount"] - assert salt_amount_target.value == expected - assert salt_amount_target.target_type == TargetType.AMOUNT - assert SOURCE_URL in salt_amount_target.source_url - - def test_load_mortgage_interest_claims_targets(self, temp_db): - """Loading deductions should create mortgage interest claims targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - fl_stratum = session.exec( - select(Stratum).where(Stratum.name == "FL All Filers") - ).first() - - mortgage_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == fl_stratum.id) - .where(Target.variable == "mortgage_interest_claims") - .where(Target.period == 2021) - ).first() - - assert mortgage_claims_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["FL"]["mortgage_interest_claims"] - assert mortgage_claims_target.value == expected - assert mortgage_claims_target.target_type == TargetType.COUNT - - def test_load_mortgage_interest_amount_targets(self, temp_db): - """Loading deductions should create mortgage interest amount targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - ny_stratum = session.exec( - select(Stratum).where(Stratum.name == "NY All Filers") - ).first() - - mortgage_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == ny_stratum.id) - .where(Target.variable == "mortgage_interest_amount") - .where(Target.period == 2021) - ).first() - - assert mortgage_amount_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["NY"]["mortgage_interest_amount"] - assert mortgage_amount_target.value == expected - assert mortgage_amount_target.target_type == TargetType.AMOUNT - - def test_load_charitable_claims_targets(self, temp_db): - """Loading deductions should create charitable contribution claims targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - ga_stratum = session.exec( - select(Stratum).where(Stratum.name == "GA All Filers") - ).first() - - charitable_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == ga_stratum.id) - .where(Target.variable == "charitable_claims") - .where(Target.period == 2021) - ).first() - - assert charitable_claims_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["GA"]["charitable_claims"] - assert charitable_claims_target.value == expected - assert charitable_claims_target.target_type == TargetType.COUNT - - def test_load_charitable_amount_targets(self, temp_db): - """Loading deductions should create charitable contribution amount targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - pa_stratum = session.exec( - select(Stratum).where(Stratum.name == "PA All Filers") - ).first() - - charitable_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == pa_stratum.id) - .where(Target.variable == "charitable_amount") - .where(Target.period == 2021) - ).first() - - assert charitable_amount_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["PA"]["charitable_amount"] - assert charitable_amount_target.value == expected - assert charitable_amount_target.target_type == TargetType.AMOUNT - - def test_load_medical_claims_targets(self, temp_db): - """Loading deductions should create medical expense claims targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - oh_stratum = session.exec( - select(Stratum).where(Stratum.name == "OH All Filers") - ).first() - - medical_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == oh_stratum.id) - .where(Target.variable == "medical_claims") - .where(Target.period == 2021) - ).first() - - assert medical_claims_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["OH"]["medical_claims"] - assert medical_claims_target.value == expected - assert medical_claims_target.target_type == TargetType.COUNT - - def test_load_medical_amount_targets(self, temp_db): - """Loading deductions should create medical expense amount targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - il_stratum = session.exec( - select(Stratum).where(Stratum.name == "IL All Filers") - ).first() - - medical_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == il_stratum.id) - .where(Target.variable == "medical_amount") - .where(Target.period == 2021) - ).first() - - assert medical_amount_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["IL"]["medical_amount"] - assert medical_amount_target.value == expected - assert medical_amount_target.target_type == TargetType.AMOUNT - - def test_load_qbi_claims_targets(self, temp_db): - """Loading deductions should create QBI deduction claims targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - wa_stratum = session.exec( - select(Stratum).where(Stratum.name == "WA All Filers") - ).first() - - qbi_claims_target = session.exec( - select(Target) - .where(Target.stratum_id == wa_stratum.id) - .where(Target.variable == "qbi_claims") - .where(Target.period == 2021) - ).first() - - assert qbi_claims_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["WA"]["qbi_claims"] - assert qbi_claims_target.value == expected - assert qbi_claims_target.target_type == TargetType.COUNT - - def test_load_qbi_amount_targets(self, temp_db): - """Loading deductions should create QBI deduction amount targets.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - co_stratum = session.exec( - select(Stratum).where(Stratum.name == "CO All Filers") - ).first() - - qbi_amount_target = session.exec( - select(Target) - .where(Target.stratum_id == co_stratum.id) - .where(Target.variable == "qbi_amount") - .where(Target.period == 2021) - ).first() - - assert qbi_amount_target is not None - expected = SOI_DEDUCTIONS_DATA[2021]["CO"]["qbi_amount"] - assert qbi_amount_target.value == expected - assert qbi_amount_target.target_type == TargetType.AMOUNT - - def test_load_soi_deductions_stratum_has_state_constraint(self, temp_db): - """State strata should have state_fips constraint.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - nj_stratum = session.exec( - select(Stratum).where(Stratum.name == "NJ All Filers") - ).first() - - # Check constraints include state_fips - state_constraint = None - for constraint in nj_stratum.constraints: - if constraint.variable == "state_fips": - state_constraint = constraint - break - - assert state_constraint is not None - assert state_constraint.operator == "==" - assert state_constraint.value == STATE_FIPS["NJ"] - - def test_load_soi_deductions_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - ma_stratum = session.exec( - select(Stratum).where(Stratum.name == "MA All Filers") - ).first() - - assert ma_stratum.parent_id == national.id - - def test_load_soi_deductions_idempotent(self, temp_db): - """Loading deductions twice should not duplicate data.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - load_soi_deductions_targets(session, years=[2021]) - - ca_strata = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).all() - - # Should only have one CA stratum - assert len(ca_strata) == 1 - - def test_all_states_loaded(self, temp_db): - """All 50 states + DC should be loaded with 10 deduction targets each.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - for state_abbrev in STATE_FIPS.keys(): - stratum = session.exec( - select(Stratum).where(Stratum.name == f"{state_abbrev} All Filers") - ).first() - - assert stratum is not None, f"Missing stratum for {state_abbrev}" - - # Each state should have 10 targets (5 deductions x 2 (claims + amount)) - targets = session.exec( - select(Target).where(Target.stratum_id == stratum.id) - ).all() - - # 5 deduction types x 2 (claims + amount) = 10 targets - assert len(targets) == 10, f"Expected 10 targets for {state_abbrev}, got {len(targets)}" - - def test_target_source_metadata(self, temp_db): - """Targets should have correct source metadata.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - salt_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "salt_claims") - ).first() - - assert salt_target.source == DataSource.IRS_SOI - assert salt_target.source_table == "SOI Individual Returns - Itemized Deductions" - assert "soi-tax-stats" in salt_target.source_url - - def test_national_totals_are_reasonable(self, temp_db): - """National totals should match expected ranges.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - # Sum all SALT claims across states - should be around 12M - all_salt_claims = sum( - state_data["salt_claims"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 10_000_000 < all_salt_claims < 20_000_000 - - # Sum all SALT amounts - should be around $150B - all_salt_amount = sum( - state_data["salt_amount"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 100_000_000_000 < all_salt_amount < 200_000_000_000 - - # Sum all mortgage interest claims - should be around 10M - all_mortgage_claims = sum( - state_data["mortgage_interest_claims"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 8_000_000 < all_mortgage_claims < 15_000_000 - - # Sum all mortgage interest amounts - should be around $80B - all_mortgage_amount = sum( - state_data["mortgage_interest_amount"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 60_000_000_000 < all_mortgage_amount < 100_000_000_000 - - # Sum all charitable claims - should be around 25M - all_charitable_claims = sum( - state_data["charitable_claims"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 20_000_000 < all_charitable_claims < 35_000_000 - - # Sum all charitable amounts - should be around $50B - all_charitable_amount = sum( - state_data["charitable_amount"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 40_000_000_000 < all_charitable_amount < 70_000_000_000 - - # Sum all medical claims - should be around 5M - all_medical_claims = sum( - state_data["medical_claims"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 4_000_000 < all_medical_claims < 10_000_000 - - # Sum all medical amounts - should be around $30B - all_medical_amount = sum( - state_data["medical_amount"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 20_000_000_000 < all_medical_amount < 50_000_000_000 - - # Sum all QBI claims - should be around 10M - all_qbi_claims = sum( - state_data["qbi_claims"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 8_000_000 < all_qbi_claims < 15_000_000 - - # Sum all QBI amounts - should be around $70B - all_qbi_amount = sum( - state_data["qbi_amount"] - for state_data in SOI_DEDUCTIONS_DATA[2021].values() - ) - assert 50_000_000_000 < all_qbi_amount < 100_000_000_000 - - def test_large_states_have_higher_claims(self, temp_db): - """Larger states should have more claims than smaller states.""" - with Session(temp_db) as session: - load_soi_deductions_targets(session, years=[2021]) - - # California (large) vs Wyoming (small) - ca_data = SOI_DEDUCTIONS_DATA[2021]["CA"] - wy_data = SOI_DEDUCTIONS_DATA[2021]["WY"] - - assert ca_data["salt_claims"] > wy_data["salt_claims"] - assert ca_data["mortgage_interest_claims"] > wy_data["mortgage_interest_claims"] - assert ca_data["charitable_claims"] > wy_data["charitable_claims"] - assert ca_data["medical_claims"] > wy_data["medical_claims"] - assert ca_data["qbi_claims"] > wy_data["qbi_claims"] - - # Texas (large) vs Vermont (small) - tx_data = SOI_DEDUCTIONS_DATA[2021]["TX"] - vt_data = SOI_DEDUCTIONS_DATA[2021]["VT"] - - assert tx_data["salt_claims"] > vt_data["salt_claims"] - assert tx_data["mortgage_interest_claims"] > vt_data["mortgage_interest_claims"] - - def test_deduction_types_constant(self, temp_db): - """DEDUCTION_TYPES should contain all 5 deduction types.""" - assert len(DEDUCTION_TYPES) == 5 - assert "salt" in DEDUCTION_TYPES - assert "mortgage_interest" in DEDUCTION_TYPES - assert "charitable" in DEDUCTION_TYPES - assert "medical" in DEDUCTION_TYPES - assert "qbi" in DEDUCTION_TYPES diff --git a/tests/test_etl_soi_income_sources.py b/tests/test_etl_soi_income_sources.py deleted file mode 100644 index 091b985f..00000000 --- a/tests/test_etl_soi_income_sources.py +++ /dev/null @@ -1,517 +0,0 @@ -"""Tests for IRS SOI income by source ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_soi_income_sources import ( - load_soi_income_sources_targets, - SOI_NATIONAL_INCOME_SOURCES_DATA, - SOI_STATE_INCOME_SOURCES_DATA, - STATE_FIPS, - INCOME_SOURCES, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_soi_income_sources.db" - engine = init_db(db_path) - yield engine - - -class TestSoiIncomeSourcesETL: - """Tests for IRS SOI income sources ETL loader.""" - - def test_load_soi_income_sources_creates_national_stratum(self, temp_db): - """Loading income source data should create/reference a national stratum.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - assert national is not None - assert national.stratum_group_id == "national" - - def test_load_national_income_sources_targets(self, temp_db): - """Loading should create national-level income source targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - # Check taxable interest target - taxable_interest_target = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "taxable_interest_returns") - .where(Target.period == 2021) - ).first() - - assert taxable_interest_target is not None - expected = SOI_NATIONAL_INCOME_SOURCES_DATA[2021]["taxable_interest_returns"] - assert taxable_interest_target.value == expected - assert taxable_interest_target.target_type == TargetType.COUNT - assert taxable_interest_target.source == DataSource.IRS_SOI - assert taxable_interest_target.source_table == "Publication 1304, Table 1.4" - - def test_load_soi_income_sources_creates_state_strata(self, temp_db): - """Loading income source data should create state-level strata.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "soi_states") - ).all() - - # Should have strata for all 50 states + DC - expected_states = len(STATE_FIPS) - assert len(state_strata) == expected_states - assert expected_states == 51 # 50 states + DC - - def test_load_taxable_interest_targets(self, temp_db): - """Loading income sources should create taxable interest targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - assert ca_stratum is not None - - # Returns count target - returns_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "taxable_interest_returns") - .where(Target.period == 2021) - ).first() - - assert returns_target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["CA"]["taxable_interest_returns"] - assert returns_target.value == expected - assert returns_target.target_type == TargetType.COUNT - - # Amount target - amount_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "taxable_interest_amount") - .where(Target.period == 2021) - ).first() - - assert amount_target is not None - assert amount_target.target_type == TargetType.AMOUNT - - def test_load_tax_exempt_interest_targets(self, temp_db): - """Loading income sources should create tax-exempt interest targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - fl_stratum = session.exec( - select(Stratum).where(Stratum.name == "FL All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == fl_stratum.id) - .where(Target.variable == "tax_exempt_interest_returns") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["FL"]["tax_exempt_interest_returns"] - assert target.value == expected - - def test_load_ordinary_dividends_targets(self, temp_db): - """Loading income sources should create ordinary dividends targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - tx_stratum = session.exec( - select(Stratum).where(Stratum.name == "TX All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == tx_stratum.id) - .where(Target.variable == "ordinary_dividends_amount") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["TX"]["ordinary_dividends_amount"] - assert target.value == expected - - def test_load_qualified_dividends_targets(self, temp_db): - """Loading income sources should create qualified dividends targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - ny_stratum = session.exec( - select(Stratum).where(Stratum.name == "NY All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == ny_stratum.id) - .where(Target.variable == "qualified_dividends_returns") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["NY"]["qualified_dividends_returns"] - assert target.value == expected - - def test_load_short_term_capital_gains_targets(self, temp_db): - """Loading income sources should create short-term capital gains targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - wa_stratum = session.exec( - select(Stratum).where(Stratum.name == "WA All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == wa_stratum.id) - .where(Target.variable == "short_term_capital_gains_amount") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["WA"]["short_term_capital_gains_amount"] - assert target.value == expected - - def test_load_long_term_capital_gains_targets(self, temp_db): - """Loading income sources should create long-term capital gains targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - nj_stratum = session.exec( - select(Stratum).where(Stratum.name == "NJ All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == nj_stratum.id) - .where(Target.variable == "long_term_capital_gains_returns") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["NJ"]["long_term_capital_gains_returns"] - assert target.value == expected - - def test_load_state_local_refunds_targets(self, temp_db): - """Loading income sources should create state/local refunds targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - ga_stratum = session.exec( - select(Stratum).where(Stratum.name == "GA All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == ga_stratum.id) - .where(Target.variable == "state_local_refunds_amount") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["GA"]["state_local_refunds_amount"] - assert target.value == expected - - def test_load_alimony_received_targets(self, temp_db): - """Loading income sources should create alimony received targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - ma_stratum = session.exec( - select(Stratum).where(Stratum.name == "MA All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == ma_stratum.id) - .where(Target.variable == "alimony_received_returns") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["MA"]["alimony_received_returns"] - assert target.value == expected - - def test_load_schedule_c_income_targets(self, temp_db): - """Loading income sources should create Schedule C business income targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - il_stratum = session.exec( - select(Stratum).where(Stratum.name == "IL All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == il_stratum.id) - .where(Target.variable == "schedule_c_income_amount") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["IL"]["schedule_c_income_amount"] - assert target.value == expected - - def test_load_rental_royalty_income_targets(self, temp_db): - """Loading income sources should create rental/royalty income targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - pa_stratum = session.exec( - select(Stratum).where(Stratum.name == "PA All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == pa_stratum.id) - .where(Target.variable == "rental_royalty_income_returns") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["PA"]["rental_royalty_income_returns"] - assert target.value == expected - - def test_load_partnership_scorp_income_targets(self, temp_db): - """Loading income sources should create partnership/S-corp income targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - oh_stratum = session.exec( - select(Stratum).where(Stratum.name == "OH All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == oh_stratum.id) - .where(Target.variable == "partnership_scorp_income_amount") - .where(Target.period == 2021) - ).first() - - assert target is not None - expected = SOI_STATE_INCOME_SOURCES_DATA[2021]["OH"]["partnership_scorp_income_amount"] - assert target.value == expected - - def test_load_soi_income_sources_stratum_has_state_constraint(self, temp_db): - """State strata should have state_fips constraint.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - nc_stratum = session.exec( - select(Stratum).where(Stratum.name == "NC All Filers") - ).first() - - # Check constraints include state_fips - state_constraint = None - for constraint in nc_stratum.constraints: - if constraint.variable == "state_fips": - state_constraint = constraint - break - - assert state_constraint is not None - assert state_constraint.operator == "==" - assert state_constraint.value == STATE_FIPS["NC"] - - def test_load_soi_income_sources_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - mi_stratum = session.exec( - select(Stratum).where(Stratum.name == "MI All Filers") - ).first() - - assert mi_stratum.parent_id == national.id - - def test_load_soi_income_sources_idempotent(self, temp_db): - """Loading income sources twice should not duplicate strata.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - load_soi_income_sources_targets(session, years=[2021]) - - ca_strata = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).all() - - # Should only have one CA stratum - assert len(ca_strata) == 1 - - def test_all_states_loaded(self, temp_db): - """All 50 states + DC should be loaded with income source targets.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - # 11 income sources x 2 (returns + amount) = 22 targets per state - expected_targets_per_state = len(INCOME_SOURCES) * 2 - - for state_abbrev in STATE_FIPS.keys(): - stratum = session.exec( - select(Stratum).where(Stratum.name == f"{state_abbrev} All Filers") - ).first() - - assert stratum is not None, f"Missing stratum for {state_abbrev}" - - targets = session.exec( - select(Target).where(Target.stratum_id == stratum.id) - ).all() - - assert len(targets) == expected_targets_per_state, ( - f"Expected {expected_targets_per_state} targets for {state_abbrev}, " - f"got {len(targets)}" - ) - - def test_target_source_metadata(self, temp_db): - """Targets should have correct source metadata.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "taxable_interest_returns") - ).first() - - assert target.source == DataSource.IRS_SOI - assert target.source_table == "Individual Income Tax State Data" - assert "individual-income-tax-statistics" in target.source_url - - def test_income_sources_list(self): - """INCOME_SOURCES should include all expected income types.""" - expected_sources = [ - "taxable_interest", - "tax_exempt_interest", - "ordinary_dividends", - "qualified_dividends", - "short_term_capital_gains", - "long_term_capital_gains", - "state_local_refunds", - "alimony_received", - "schedule_c_income", - "rental_royalty_income", - "partnership_scorp_income", - ] - assert INCOME_SOURCES == expected_sources - - def test_national_totals_are_reasonable(self, temp_db): - """National totals should match expected ranges.""" - with Session(temp_db) as session: - load_soi_income_sources_targets(session, years=[2021]) - - # Check national taxable interest returns - should be ~55M - national_data = SOI_NATIONAL_INCOME_SOURCES_DATA[2021] - assert 40_000_000 < national_data["taxable_interest_returns"] < 80_000_000 - - # Check national long-term capital gains - should be ~$850B - assert 500_000_000_000 < national_data["long_term_capital_gains_amount"] < 1_500_000_000_000 - - # Alimony should be much lower than other income sources - assert national_data["alimony_received_returns"] < national_data["taxable_interest_returns"] - - def test_qualified_dividends_less_than_ordinary(self, temp_db): - """Qualified dividends should always be less than or equal to ordinary dividends.""" - for state_data in SOI_STATE_INCOME_SOURCES_DATA[2021].values(): - assert state_data["qualified_dividends_returns"] <= state_data["ordinary_dividends_returns"] - assert state_data["qualified_dividends_amount"] <= state_data["ordinary_dividends_amount"] - - -class TestIncomeSourcesData: - """Tests for the income sources data structure.""" - - def test_all_states_have_data(self): - """All 50 states + DC should have income sources data.""" - assert len(SOI_STATE_INCOME_SOURCES_DATA[2021]) == 51 - - for state_abbrev in STATE_FIPS.keys(): - assert state_abbrev in SOI_STATE_INCOME_SOURCES_DATA[2021], ( - f"Missing data for {state_abbrev}" - ) - - def test_all_income_sources_present(self): - """Each state should have all income source variables.""" - expected_vars = [] - for source in INCOME_SOURCES: - expected_vars.append(f"{source}_returns") - expected_vars.append(f"{source}_amount") - - for state_abbrev, state_data in SOI_STATE_INCOME_SOURCES_DATA[2021].items(): - for var in expected_vars: - assert var in state_data, ( - f"Missing {var} for {state_abbrev}" - ) - assert isinstance(state_data[var], (int, float)), ( - f"{var} for {state_abbrev} is not numeric" - ) - - def test_values_are_positive(self): - """All values should be positive (or zero for some edge cases).""" - for state_abbrev, state_data in SOI_STATE_INCOME_SOURCES_DATA[2021].items(): - for var, value in state_data.items(): - assert value >= 0, ( - f"Negative value for {var} in {state_abbrev}: {value}" - ) - - def test_large_states_have_higher_values(self): - """Larger states should have more income than smaller states.""" - # California (large) vs Wyoming (small) - ca_data = SOI_STATE_INCOME_SOURCES_DATA[2021]["CA"] - wy_data = SOI_STATE_INCOME_SOURCES_DATA[2021]["WY"] - - assert ca_data["taxable_interest_returns"] > wy_data["taxable_interest_returns"] - assert ca_data["ordinary_dividends_amount"] > wy_data["ordinary_dividends_amount"] - assert ca_data["long_term_capital_gains_amount"] > wy_data["long_term_capital_gains_amount"] - assert ca_data["schedule_c_income_amount"] > wy_data["schedule_c_income_amount"] - - # Texas (large) vs Vermont (small) - tx_data = SOI_STATE_INCOME_SOURCES_DATA[2021]["TX"] - vt_data = SOI_STATE_INCOME_SOURCES_DATA[2021]["VT"] - - assert tx_data["taxable_interest_returns"] > vt_data["taxable_interest_returns"] - assert tx_data["partnership_scorp_income_amount"] > vt_data["partnership_scorp_income_amount"] - - def test_national_data_structure(self): - """National data should have all required fields.""" - national_data = SOI_NATIONAL_INCOME_SOURCES_DATA[2021] - - for source in INCOME_SOURCES: - returns_var = f"{source}_returns" - amount_var = f"{source}_amount" - - assert returns_var in national_data, f"Missing {returns_var} in national data" - assert amount_var in national_data, f"Missing {amount_var} in national data" - assert isinstance(national_data[returns_var], int) - assert isinstance(national_data[amount_var], int) diff --git a/tests/test_etl_soi_state.py b/tests/test_etl_soi_state.py deleted file mode 100644 index 862006cc..00000000 --- a/tests/test_etl_soi_state.py +++ /dev/null @@ -1,395 +0,0 @@ -"""Tests for state-level SOI ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_soi_state import ( - load_soi_state_targets, - SOI_STATE_DATA, - STATE_FIPS, - AGI_BRACKETS, - SOI_STATE_AGI_DATA, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_soi_state.db" - engine = init_db(db_path) - yield engine - - -class TestSoiStateETL: - """Tests for state-level SOI ETL loader.""" - - def test_load_soi_state_creates_national_stratum(self, temp_db): - """Loading state SOI data should create/reference a national stratum.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - assert national is not None - assert national.stratum_group_id == "national" - - def test_load_soi_state_creates_state_strata(self, temp_db): - """Loading state SOI data should create state-level strata.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "soi_states") - ).all() - - # Should have strata for all 50 states + DC - expected_states = len(STATE_FIPS) - assert len(state_strata) == expected_states - assert expected_states == 51 # 50 states + DC - - def test_load_soi_state_creates_returns_targets(self, temp_db): - """Loading state SOI should create tax returns count targets.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - assert ca_stratum is not None - - returns_target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "tax_unit_count") - .where(Target.period == 2021) - ).first() - - assert returns_target is not None - expected = SOI_STATE_DATA[2021]["CA"]["total_returns"] - assert returns_target.value == expected - assert returns_target.target_type == TargetType.COUNT - assert returns_target.source == DataSource.IRS_SOI - - def test_load_soi_state_creates_agi_targets(self, temp_db): - """Loading state SOI should create AGI amount targets.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - tx_stratum = session.exec( - select(Stratum).where(Stratum.name == "TX All Filers") - ).first() - - agi_target = session.exec( - select(Target) - .where(Target.stratum_id == tx_stratum.id) - .where(Target.variable == "adjusted_gross_income") - .where(Target.period == 2021) - ).first() - - assert agi_target is not None - expected = SOI_STATE_DATA[2021]["TX"]["total_agi"] - assert agi_target.value == expected - assert agi_target.target_type == TargetType.AMOUNT - - def test_load_soi_state_creates_tax_liability_targets(self, temp_db): - """Loading state SOI should create tax liability targets.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - fl_stratum = session.exec( - select(Stratum).where(Stratum.name == "FL All Filers") - ).first() - - tax_target = session.exec( - select(Target) - .where(Target.stratum_id == fl_stratum.id) - .where(Target.variable == "income_tax_liability") - .where(Target.period == 2021) - ).first() - - assert tax_target is not None - expected = SOI_STATE_DATA[2021]["FL"]["total_tax_liability"] - assert tax_target.value == expected - assert tax_target.target_type == TargetType.AMOUNT - - def test_load_soi_state_stratum_has_state_constraint(self, temp_db): - """State strata should have state_fips constraint.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - ny_stratum = session.exec( - select(Stratum).where(Stratum.name == "NY All Filers") - ).first() - - # Check constraints include state_fips - state_constraint = None - for constraint in ny_stratum.constraints: - if constraint.variable == "state_fips": - state_constraint = constraint - break - - assert state_constraint is not None - assert state_constraint.operator == "==" - assert state_constraint.value == STATE_FIPS["NY"] - - def test_load_soi_state_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US All Filers") - ).first() - - pa_stratum = session.exec( - select(Stratum).where(Stratum.name == "PA All Filers") - ).first() - - assert pa_stratum.parent_id == national.id - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2020, 2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "tax_unit_count") - ).all() - - years = {t.period for t in targets} - assert years == {2020, 2021} - - def test_load_soi_state_idempotent(self, temp_db): - """Loading state SOI twice should not duplicate data.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - load_soi_state_targets(session, years=[2021]) - - ca_strata = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).all() - - # Should only have one CA stratum - assert len(ca_strata) == 1 - - def test_all_states_loaded(self, temp_db): - """All 50 states + DC should be loaded.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - for state_abbrev in STATE_FIPS.keys(): - stratum = session.exec( - select(Stratum).where(Stratum.name == f"{state_abbrev} All Filers") - ).first() - - assert stratum is not None, f"Missing stratum for {state_abbrev}" - - # Each state should have 3 targets (returns, AGI, tax liability) - targets = session.exec( - select(Target).where(Target.stratum_id == stratum.id) - ).all() - - assert len(targets) == 3 - - def test_target_source_metadata(self, temp_db): - """Targets should have correct source metadata.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - target = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "tax_unit_count") - ).first() - - assert target.source == DataSource.IRS_SOI - assert target.source_table == "Historic Table 2" - assert "historic-table-2" in target.source_url - - -class TestAgiBracketStratification: - """Tests for AGI bracket stratification.""" - - def test_agi_brackets_defined(self): - """AGI_BRACKETS should be properly defined.""" - assert len(AGI_BRACKETS) == 10 - - # Check first bracket (under $1) - assert AGI_BRACKETS[0]["min"] == 0 - assert AGI_BRACKETS[0]["max"] == 1 - assert AGI_BRACKETS[0]["label"] == "under_1" - - # Check last bracket ($1M+) - assert AGI_BRACKETS[-1]["min"] == 1_000_000 - assert AGI_BRACKETS[-1]["max"] is None - assert AGI_BRACKETS[-1]["label"] == "1m_plus" - - def test_agi_brackets_contiguous(self): - """AGI brackets should be contiguous (no gaps).""" - for i in range(len(AGI_BRACKETS) - 1): - current = AGI_BRACKETS[i] - next_bracket = AGI_BRACKETS[i + 1] - assert current["max"] == next_bracket["min"], ( - f"Gap between bracket {i} (max={current['max']}) " - f"and bracket {i+1} (min={next_bracket['min']})" - ) - - def test_load_soi_state_creates_agi_bracket_strata(self, temp_db): - """Loading state SOI data should create AGI bracket strata.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - # Check CA has AGI bracket strata - bracket_strata = session.exec( - select(Stratum).where( - Stratum.stratum_group_id == "soi_states_agi_brackets" - ) - ).all() - - # Should have strata for each state x bracket combination - expected_count = len(STATE_FIPS) * len(AGI_BRACKETS) - assert len(bracket_strata) == expected_count - - def test_agi_bracket_stratum_has_state_and_bracket_constraints(self, temp_db): - """AGI bracket strata should have both state_fips and agi_bracket constraints.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - # Find CA $50k-$75k bracket stratum - ca_bracket_stratum = session.exec( - select(Stratum).where( - Stratum.name == "CA Filers AGI $50k-$75k" - ) - ).first() - - assert ca_bracket_stratum is not None - - # Should have state_fips constraint - constraints_dict = { - c.variable: (c.operator, c.value) - for c in ca_bracket_stratum.constraints - } - - assert "state_fips" in constraints_dict - assert constraints_dict["state_fips"] == ("==", STATE_FIPS["CA"]) - - # Should have agi_bracket constraint - assert "agi_bracket" in constraints_dict - assert constraints_dict["agi_bracket"] == ("==", "50k_to_75k") - - def test_agi_bracket_stratum_has_state_parent(self, temp_db): - """AGI bracket strata should have state stratum as parent.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - ca_state_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA All Filers") - ).first() - - ca_bracket_stratum = session.exec( - select(Stratum).where( - Stratum.name == "CA Filers AGI $50k-$75k" - ) - ).first() - - assert ca_bracket_stratum.parent_id == ca_state_stratum.id - - def test_agi_bracket_targets_created(self, temp_db): - """AGI bracket strata should have targets for returns, AGI, and tax.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - ca_bracket_stratum = session.exec( - select(Stratum).where( - Stratum.name == "CA Filers AGI $100k-$200k" - ) - ).first() - - assert ca_bracket_stratum is not None - - targets = session.exec( - select(Target).where(Target.stratum_id == ca_bracket_stratum.id) - ).all() - - # Should have 3 targets: returns, AGI, tax liability - assert len(targets) == 3 - - variables = {t.variable for t in targets} - assert variables == { - "tax_unit_count", - "adjusted_gross_income", - "income_tax_liability", - } - - def test_agi_bracket_target_values(self, temp_db): - """AGI bracket targets should have correct values from data.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - tx_bracket_stratum = session.exec( - select(Stratum).where( - Stratum.name == "TX Filers AGI $1M+" - ) - ).first() - - returns_target = session.exec( - select(Target) - .where(Target.stratum_id == tx_bracket_stratum.id) - .where(Target.variable == "tax_unit_count") - ).first() - - expected = SOI_STATE_AGI_DATA[2021]["TX"]["1m_plus"]["total_returns"] - assert returns_target.value == expected - - def test_all_brackets_for_state(self, temp_db): - """Each state should have all 10 AGI brackets.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - # Check NY has all 10 brackets - ny_bracket_strata = session.exec( - select(Stratum).where( - Stratum.name.like("NY Filers AGI %") - ) - ).all() - - assert len(ny_bracket_strata) == len(AGI_BRACKETS) - - def test_bracket_stratum_description(self, temp_db): - """AGI bracket strata should have descriptive descriptions.""" - with Session(temp_db) as session: - load_soi_state_targets(session, years=[2021]) - - stratum = session.exec( - select(Stratum).where( - Stratum.name == "FL Filers AGI Under $1" - ) - ).first() - - assert stratum is not None - assert "under $1" in stratum.description.lower() - assert "FL" in stratum.description or "Florida" in stratum.description diff --git a/tests/test_etl_ssa.py b/tests/test_etl_ssa.py deleted file mode 100644 index 509fa7ef..00000000 --- a/tests/test_etl_ssa.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Tests for SSA (Social Security) ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_ssa.db" - engine = init_db(db_path) - yield engine - - -class TestSsaETL: - """Tests for SSA ETL loader.""" - - def test_load_creates_oasdi_stratum(self, temp_db): - """Loading SSA data should create OASDI beneficiary stratum.""" - from db.etl_ssa import load_ssa_targets - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - oasdi = session.exec( - select(Stratum).where(Stratum.name == "US OASDI Beneficiaries") - ).first() - - assert oasdi is not None - - def test_load_creates_beneficiary_count(self, temp_db): - """Loading SSA should create total beneficiary count target.""" - from db.etl_ssa import load_ssa_targets, SSA_DATA - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - oasdi = session.exec( - select(Stratum).where(Stratum.name == "US OASDI Beneficiaries") - ).first() - - count = session.exec( - select(Target) - .where(Target.stratum_id == oasdi.id) - .where(Target.variable == "oasdi_beneficiaries") - .where(Target.period == 2023) - ).first() - - assert count is not None - assert count.value == SSA_DATA[2023]["total_beneficiaries"] - assert count.target_type == TargetType.COUNT - assert count.source == DataSource.SSA - - def test_load_creates_benefit_payments(self, temp_db): - """Loading SSA should create total benefit payment target.""" - from db.etl_ssa import load_ssa_targets, SSA_DATA - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - oasdi = session.exec( - select(Stratum).where(Stratum.name == "US OASDI Beneficiaries") - ).first() - - payments = session.exec( - select(Target) - .where(Target.stratum_id == oasdi.id) - .where(Target.variable == "oasdi_benefits") - ).first() - - assert payments is not None - assert payments.value == SSA_DATA[2023]["total_benefits"] - assert payments.target_type == TargetType.AMOUNT - - def test_load_creates_ssi_stratum(self, temp_db): - """Loading SSA should create SSI recipient stratum.""" - from db.etl_ssa import load_ssa_targets - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - ssi = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - assert ssi is not None - - def test_load_creates_ssi_targets(self, temp_db): - """Loading SSA should create SSI recipient and payment targets.""" - from db.etl_ssa import load_ssa_targets, SSA_DATA - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - ssi = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - count = session.exec( - select(Target) - .where(Target.stratum_id == ssi.id) - .where(Target.variable == "ssi_recipients") - ).first() - - assert count is not None - assert count.value == SSA_DATA[2023]["ssi"]["recipients"] - - def test_load_creates_retired_worker_stratum(self, temp_db): - """Loading SSA should create retired worker stratum.""" - from db.etl_ssa import load_ssa_targets - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - retired = session.exec( - select(Stratum).where(Stratum.name == "US Retired Workers") - ).first() - - assert retired is not None - - def test_retired_worker_avg_benefit(self, temp_db): - """Retired worker average benefit should be correct.""" - from db.etl_ssa import load_ssa_targets, SSA_DATA - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - - retired = session.exec( - select(Stratum).where(Stratum.name == "US Retired Workers") - ).first() - - avg_benefit = session.exec( - select(Target) - .where(Target.stratum_id == retired.id) - .where(Target.variable == "oasdi_avg_monthly_benefit") - ).first() - - assert avg_benefit is not None - expected = SSA_DATA[2023]["retired_workers"]["avg_monthly_benefit"] - assert avg_benefit.value == expected - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - from db.etl_ssa import load_ssa_targets - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2022, 2023]) - - oasdi = session.exec( - select(Stratum).where(Stratum.name == "US OASDI Beneficiaries") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == oasdi.id) - .where(Target.variable == "oasdi_beneficiaries") - ).all() - - years = {t.period for t in targets} - assert years == {2022, 2023} - - def test_load_idempotent(self, temp_db): - """Loading twice should not duplicate data.""" - from db.etl_ssa import load_ssa_targets - - with Session(temp_db) as session: - load_ssa_targets(session, years=[2023]) - load_ssa_targets(session, years=[2023]) - - oasdi_strata = session.exec( - select(Stratum).where(Stratum.name == "US OASDI Beneficiaries") - ).all() - - assert len(oasdi_strata) == 1 diff --git a/tests/test_etl_ssi.py b/tests/test_etl_ssi.py deleted file mode 100644 index 29883c33..00000000 --- a/tests/test_etl_ssi.py +++ /dev/null @@ -1,366 +0,0 @@ -"""Tests for SSI (Supplemental Security Income) ETL.""" - -import tempfile -from pathlib import Path - -import pytest -from sqlmodel import Session, select - -from db.schema import ( - DataSource, - Stratum, - Target, - TargetType, - init_db, -) -from db.etl_ssi import load_ssi_targets, SSI_DATA - - -@pytest.fixture -def temp_db(): - """Create a temporary database for testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_ssi.db" - engine = init_db(db_path) - yield engine - - -class TestSsiETL: - """Tests for SSI ETL loader.""" - - def test_load_creates_national_stratum(self, temp_db): - """Loading SSI data should create a national stratum.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - assert national is not None - assert national.stratum_group_id == "ssi_national" - - def test_load_creates_national_recipient_count(self, temp_db): - """Loading SSI should create national recipient count target.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - count = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "ssi_recipients") - .where(Target.period == 2023) - ).first() - - assert count is not None - assert count.value == SSI_DATA[2023]["national"]["recipients"] - assert count.target_type == TargetType.COUNT - assert count.source == DataSource.SSA - - def test_load_creates_national_payments(self, temp_db): - """Loading SSI should create national total payments target.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - payments = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "ssi_total_payments") - ).first() - - assert payments is not None - assert payments.value == SSI_DATA[2023]["national"]["total_payments"] - assert payments.target_type == TargetType.AMOUNT - - def test_load_creates_avg_monthly_payment(self, temp_db): - """Loading SSI should create average monthly payment target.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - avg_payment = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "ssi_avg_monthly_payment") - ).first() - - assert avg_payment is not None - assert avg_payment.value == SSI_DATA[2023]["national"]["avg_monthly_payment"] - - def test_load_creates_federal_vs_state_supplementation(self, temp_db): - """Loading SSI should create federal and state supplementation targets.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - federal = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "ssi_federal_payments") - ).first() - - state_supp = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "ssi_state_supplementation") - ).first() - - assert federal is not None - assert federal.value == SSI_DATA[2023]["national"]["federal_payments"] - - assert state_supp is not None - assert state_supp.value == SSI_DATA[2023]["national"]["state_supplementation"] - - def test_load_creates_aged_stratum(self, temp_db): - """Loading SSI should create aged recipient stratum.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - aged = session.exec( - select(Stratum).where(Stratum.name == "US SSI Aged Recipients") - ).first() - - assert aged is not None - assert aged.stratum_group_id == "ssi_categories" - - def test_load_creates_blind_stratum(self, temp_db): - """Loading SSI should create blind recipient stratum.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - blind = session.exec( - select(Stratum).where(Stratum.name == "US SSI Blind Recipients") - ).first() - - assert blind is not None - - def test_load_creates_disabled_stratum(self, temp_db): - """Loading SSI should create disabled recipient stratum.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - disabled = session.exec( - select(Stratum).where(Stratum.name == "US SSI Disabled Recipients") - ).first() - - assert disabled is not None - - def test_category_recipient_counts(self, temp_db): - """Category recipient counts should sum approximately to total.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - aged = session.exec( - select(Stratum).where(Stratum.name == "US SSI Aged Recipients") - ).first() - blind = session.exec( - select(Stratum).where(Stratum.name == "US SSI Blind Recipients") - ).first() - disabled = session.exec( - select(Stratum).where(Stratum.name == "US SSI Disabled Recipients") - ).first() - - aged_count = session.exec( - select(Target) - .where(Target.stratum_id == aged.id) - .where(Target.variable == "ssi_recipients") - ).first() - blind_count = session.exec( - select(Target) - .where(Target.stratum_id == blind.id) - .where(Target.variable == "ssi_recipients") - ).first() - disabled_count = session.exec( - select(Target) - .where(Target.stratum_id == disabled.id) - .where(Target.variable == "ssi_recipients") - ).first() - - total_from_categories = ( - aged_count.value + blind_count.value + disabled_count.value - ) - national_total = SSI_DATA[2023]["national"]["recipients"] - - # Categories should sum to close to (but may not exactly equal) total - # due to some recipients qualifying under multiple categories - assert total_from_categories > national_total * 0.95 - assert total_from_categories < national_total * 1.05 - - def test_load_creates_state_strata(self, temp_db): - """Loading SSI should create state-level strata.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - state_strata = session.exec( - select(Stratum).where(Stratum.stratum_group_id == "ssi_states") - ).all() - - # Should have strata for states in the data - expected_states = len(SSI_DATA[2023].get("states", {})) - assert len(state_strata) == expected_states - - def test_load_state_targets_correct(self, temp_db): - """State-level SSI targets should have correct values.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - ca_stratum = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Recipients") - ).first() - - assert ca_stratum is not None - - ca_recipients = session.exec( - select(Target) - .where(Target.stratum_id == ca_stratum.id) - .where(Target.variable == "ssi_recipients") - ).first() - - expected_ca = SSI_DATA[2023]["states"]["CA"]["recipients"] - assert ca_recipients.value == expected_ca - - def test_load_state_category_breakdowns(self, temp_db): - """State-level SSI should have category breakdowns.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - ca_aged = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Aged Recipients") - ).first() - ca_blind = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Blind Recipients") - ).first() - ca_disabled = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Disabled Recipients") - ).first() - - assert ca_aged is not None - assert ca_blind is not None - assert ca_disabled is not None - - # Check aged count - aged_count = session.exec( - select(Target) - .where(Target.stratum_id == ca_aged.id) - .where(Target.variable == "ssi_recipients") - ).first() - - assert aged_count.value == SSI_DATA[2023]["states"]["CA"]["aged"] - - def test_state_stratum_has_parent(self, temp_db): - """State strata should have national stratum as parent.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - ca = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Recipients") - ).first() - - assert ca.parent_id == national.id - - def test_state_category_has_state_parent(self, temp_db): - """State category strata should have state stratum as parent.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - ca_state = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Recipients") - ).first() - - ca_aged = session.exec( - select(Stratum).where(Stratum.name == "CA SSI Aged Recipients") - ).first() - - assert ca_aged.parent_id == ca_state.id - - def test_load_multiple_years(self, temp_db): - """Loading multiple years should create targets for each.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2021, 2022, 2023]) - - national = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).first() - - targets = session.exec( - select(Target) - .where(Target.stratum_id == national.id) - .where(Target.variable == "ssi_recipients") - ).all() - - years = {t.period for t in targets} - assert years == {2021, 2022, 2023} - - def test_load_idempotent(self, temp_db): - """Loading SSI twice should not duplicate strata.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - load_ssi_targets(session, years=[2023]) - - national_strata = session.exec( - select(Stratum).where(Stratum.name == "US SSI Recipients") - ).all() - - # Should only have one national stratum - assert len(national_strata) == 1 - - def test_load_texas_data(self, temp_db): - """Texas SSI data should load correctly.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - tx_stratum = session.exec( - select(Stratum).where(Stratum.name == "TX SSI Recipients") - ).first() - - assert tx_stratum is not None - - tx_recipients = session.exec( - select(Target) - .where(Target.stratum_id == tx_stratum.id) - .where(Target.variable == "ssi_recipients") - ).first() - - assert tx_recipients.value == SSI_DATA[2023]["states"]["TX"]["recipients"] - - def test_load_florida_data(self, temp_db): - """Florida SSI data should load correctly.""" - with Session(temp_db) as session: - load_ssi_targets(session, years=[2023]) - - fl_stratum = session.exec( - select(Stratum).where(Stratum.name == "FL SSI Recipients") - ).first() - - assert fl_stratum is not None - - # Florida has high aged population - fl_aged = session.exec( - select(Stratum).where(Stratum.name == "FL SSI Aged Recipients") - ).first() - - aged_count = session.exec( - select(Target) - .where(Target.stratum_id == fl_aged.id) - .where(Target.variable == "ssi_recipients") - ).first() - - assert aged_count.value == SSI_DATA[2023]["states"]["FL"]["aged"] diff --git a/tests/test_hierarchical_calibration.py b/tests/test_hierarchical_calibration.py deleted file mode 100644 index e04b1398..00000000 --- a/tests/test_hierarchical_calibration.py +++ /dev/null @@ -1,514 +0,0 @@ -"""Tests for hierarchical calibration with person/household aggregation.""" - -from dataclasses import dataclass -from enum import Enum -import importlib.util -from pathlib import Path -from typing import Optional - -import numpy as np -import pandas as pd -import pytest - -# Import only the minimal types needed, avoid supabase dependency chain -# by defining local stubs for testing -class TargetType(str, Enum): - """Stub for db.schema.TargetType for testing.""" - COUNT = "count" - AMOUNT = "amount" - RATE = "rate" - - -class DataSource(str, Enum): - """Stub for db.schema.DataSource for testing.""" - IRS_SOI = "irs-soi" - CENSUS_ACS = "census-acs" - - -# Load calibration.variables directly -_calibration_path = Path(__file__).parent.parent / "calibration" - -_variables_spec = importlib.util.spec_from_file_location( - "calibration_variables", _calibration_path / "variables.py" -) -_variables_module = importlib.util.module_from_spec(_variables_spec) -_variables_spec.loader.exec_module(_variables_module) - -infer_target_level = _variables_module.infer_target_level -parse_variable_ref = _variables_module.parse_variable_ref -InvalidVariableRefError = _variables_module.InvalidVariableRefError - - -@dataclass -class TargetSpec: - """Minimal TargetSpec for testing.""" - variable: str - value: float - target_type: TargetType - constraints: list - source: DataSource - period: int - tolerance: Optional[float] = None - stratum_name: Optional[str] = None - - -@dataclass -class Constraint: - """Minimal Constraint for testing.""" - indicator: np.ndarray - target_value: float - variable: str - target_type: TargetType - tolerance: float = 0.01 - stratum_name: Optional[str] = None - - -def apply_stratum_constraints( - microdata: pd.DataFrame, - constraints: list, -) -> pd.Series: - """Apply stratum constraints to get boolean mask.""" - mask = pd.Series(True, index=microdata.index) - - for variable, operator, value in constraints: - if variable not in microdata.columns: - continue - - col = microdata[variable] - if pd.api.types.is_numeric_dtype(col): - parsed_value = float(value) - else: - parsed_value = value - - if operator == "==": - mask &= col == parsed_value - elif operator == ">=": - mask &= col >= parsed_value - elif operator == "<": - mask &= col < parsed_value - elif operator == "<=": - mask &= col <= parsed_value - elif operator == ">": - mask &= col > parsed_value - - return mask - - -def build_hierarchical_constraint_matrix( - hh_df: pd.DataFrame, - person_df: pd.DataFrame, - targets: list, - tolerance: float = 0.01, - hh_id_col: str = "household_id", - tax_unit_df: Optional[pd.DataFrame] = None, -) -> list: - """Simplified hierarchical constraint builder for testing.""" - constraints = [] - - for target in targets: - # Infer level from constraint variables - level = infer_target_level(target.constraints) - - if level == "household": - mask = apply_stratum_constraints(hh_df, target.constraints) - if target.target_type == TargetType.COUNT: - indicator = mask.astype(float).values - else: - indicator = np.zeros(len(hh_df)) - else: - # Aggregate from person to household - mask = apply_stratum_constraints(person_df, target.constraints) - filtered = person_df[mask] - - if len(filtered) == 0: - indicator = np.zeros(len(hh_df)) - elif target.target_type == TargetType.COUNT: - agg = filtered.groupby(hh_id_col).size() - indicator = hh_df[hh_id_col].map(agg).fillna(0).values - else: - indicator = np.zeros(len(hh_df)) - - constraints.append(Constraint( - indicator=indicator, - target_value=target.value, - variable=target.variable, - target_type=target.target_type, - tolerance=target.tolerance if target.tolerance else tolerance, - stratum_name=target.stratum_name, - )) - - return constraints - - -@pytest.fixture -def mock_households(): - """Create mock household DataFrame with 100 households.""" - np.random.seed(42) - n_households = 100 - - return pd.DataFrame({ - "household_id": range(1, n_households + 1), - "state_fips": np.random.choice( - ["06", "36", "48", "12", "17"], # CA, NY, TX, FL, IL - n_households, - p=[0.30, 0.15, 0.20, 0.20, 0.15], - ), - "weight": np.random.uniform(100, 500, n_households), - }) - - -@pytest.fixture -def mock_persons(mock_households): - """Create mock person DataFrame with ~250 persons linked to households.""" - np.random.seed(43) - - persons = [] - person_id = 1 - - for hh_id in mock_households["household_id"]: - # Each household has 1-5 persons - n_persons = np.random.randint(1, 6) - for _ in range(n_persons): - age = np.random.randint(0, 95) - persons.append({ - "person_id": person_id, - "household_id": hh_id, - "age": age, - "is_employed": 1 if age >= 18 and age < 65 and np.random.random() > 0.3 else 0, - }) - person_id += 1 - - return pd.DataFrame(persons) - - -@pytest.fixture -def person_count_target(): - """Create a person-level COUNT target: people aged 18-64.""" - return TargetSpec( - variable="person_count", - value=150_000_000.0, # Target count - target_type=TargetType.COUNT, - constraints=[ - ("age", ">=", "18"), - ("age", "<", "65"), - ], - source=DataSource.CENSUS_ACS, - period=2023, - stratum_name="Working age population (18-64)", - ) - - -@pytest.fixture -def household_count_target(): - """Create a household-level COUNT target: households in California.""" - return TargetSpec( - variable="household_count", - value=13_000_000.0, # Target count - target_type=TargetType.COUNT, - constraints=[ - ("state_fips", "==", "06"), - ], - source=DataSource.CENSUS_ACS, - period=2023, - stratum_name="California households", - ) - - -class TestBuildHierarchicalConstraintMatrix: - """Tests for build_hierarchical_constraint_matrix() function.""" - - def test_returns_constraint_objects( - self, mock_households, mock_persons, person_count_target - ): - """build_hierarchical_constraint_matrix should return Constraint objects.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - ) - - assert isinstance(constraints, list) - assert len(constraints) == 1 - assert isinstance(constraints[0], Constraint) - - def test_indicator_length_matches_households( - self, mock_households, mock_persons, person_count_target - ): - """Person-level targets should be aggregated to household level.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - ) - - # Indicator length should equal number of households, not persons - assert len(constraints[0].indicator) == len(mock_households) - - def test_person_target_aggregation_correct( - self, mock_households, mock_persons, person_count_target - ): - """Person-level aggregation should correctly count matching persons per household.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - ) - - # Manually compute expected indicator: count of persons aged 18-64 per household - mask = (mock_persons["age"] >= 18) & (mock_persons["age"] < 65) - expected_counts = ( - mock_persons[mask] - .groupby("household_id") - .size() - ) - - # Map back to household order - expected_indicator = ( - mock_households["household_id"] - .map(expected_counts) - .fillna(0) - .values - ) - - np.testing.assert_array_equal( - constraints[0].indicator, - expected_indicator, - ) - - def test_household_target_direct_indicator( - self, mock_households, mock_persons, household_count_target - ): - """Household-level targets should produce direct indicator (no aggregation).""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[household_count_target], - ) - - # Indicator should be 1 for California households, 0 otherwise - expected_indicator = ( - mock_households["state_fips"] == "06" - ).astype(float).values - - np.testing.assert_array_equal( - constraints[0].indicator, - expected_indicator, - ) - - def test_indicator_length_equals_households_for_household_target( - self, mock_households, mock_persons, household_count_target - ): - """Household-level indicator length should match household count.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[household_count_target], - ) - - assert len(constraints[0].indicator) == len(mock_households) - - def test_multiple_targets_mixed_levels( - self, mock_households, mock_persons, person_count_target, household_count_target - ): - """Should handle both person and household level targets together.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target, household_count_target], - ) - - assert len(constraints) == 2 - # Both should have household-level indicator length - assert len(constraints[0].indicator) == len(mock_households) - assert len(constraints[1].indicator) == len(mock_households) - - def test_stores_target_value( - self, mock_households, mock_persons, person_count_target - ): - """Constraint should store the target value from TargetSpec.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - ) - - assert constraints[0].target_value == 150_000_000.0 - - def test_stores_variable_name( - self, mock_households, mock_persons, person_count_target - ): - """Constraint should store the variable name from TargetSpec.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - ) - - assert constraints[0].variable == "person_count" - - def test_stores_target_type( - self, mock_households, mock_persons, person_count_target - ): - """Constraint should store the target type from TargetSpec.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - ) - - assert constraints[0].target_type == TargetType.COUNT - - def test_custom_tolerance( - self, mock_households, mock_persons, person_count_target - ): - """Should respect custom tolerance parameter.""" - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[person_count_target], - tolerance=0.05, - ) - - assert constraints[0].tolerance == 0.05 - - def test_no_matching_persons_gives_zero_indicator( - self, mock_households, mock_persons - ): - """Target with no matching persons should produce zero indicator.""" - # Target for people aged 150+ (none exist) - impossible_target = TargetSpec( - variable="person_count", - value=0.0, - target_type=TargetType.COUNT, - constraints=[ - ("age", ">=", "150"), - ], - source=DataSource.CENSUS_ACS, - period=2023, - ) - - constraints = build_hierarchical_constraint_matrix( - hh_df=mock_households, - person_df=mock_persons, - targets=[impossible_target], - ) - - np.testing.assert_array_equal( - constraints[0].indicator, - np.zeros(len(mock_households)), - ) - - -class TestInferTargetLevel: - """Tests for infer_target_level() function.""" - - def test_person_variable_returns_person(self): - """Constraints with person-level variables should return 'person'.""" - constraints = [("age", ">=", "18")] - level = infer_target_level(constraints) - assert level == "person" - - def test_household_variable_returns_household(self): - """Constraints with household-level variables should return 'household'.""" - constraints = [("state_fips", "==", "06")] - level = infer_target_level(constraints) - assert level == "household" - - def test_mixed_levels_returns_most_granular(self): - """Constraints with mixed levels should return most granular (person).""" - constraints = [ - ("age", ">=", "18"), - ("state_fips", "==", "06"), - ] - level = infer_target_level(constraints) - # Person is more granular than household - assert level == "person" - - def test_empty_constraints_returns_household(self): - """Empty constraints should default to household level.""" - constraints = [] - level = infer_target_level(constraints) - assert level == "household" - - def test_tax_unit_variable_returns_tax_unit(self): - """Constraints with tax_unit-level variables should return 'tax_unit'.""" - constraints = [("filing_status", "==", "single")] - level = infer_target_level(constraints) - assert level == "tax_unit" - - def test_unknown_variable_defaults_to_person(self): - """Unknown variables should default to person level.""" - constraints = [("unknown_variable_xyz", "==", "value")] - level = infer_target_level(constraints) - assert level == "person" - - -class TestParseVariableRef: - """Tests for parse_variable_ref() function.""" - - def test_parse_valid_reference(self): - """Should correctly parse a valid fully qualified reference.""" - package, path, var_name = parse_variable_ref("us:statute/26/32#eitc") - - assert package == "us" - assert path == "statute/26/32" - assert var_name == "eitc" - - def test_parse_state_reference(self): - """Should correctly parse a state-level reference.""" - package, path, var_name = parse_variable_ref("us-ca:statute/ca/rtc/17041#ca_agi") - - assert package == "us-ca" - assert path == "statute/ca/rtc/17041" - assert var_name == "ca_agi" - - def test_parse_nested_path(self): - """Should correctly parse a deeply nested path.""" - package, path, var_name = parse_variable_ref( - "us:statute/26/24/a/1#child_tax_credit_amount" - ) - - assert package == "us" - assert path == "statute/26/24/a/1" - assert var_name == "child_tax_credit_amount" - - def test_invalid_ref_missing_colon_raises_error(self): - """Reference without colon should raise InvalidVariableRefError.""" - with pytest.raises(InvalidVariableRefError): - parse_variable_ref("us-statute/26/32#eitc") - - def test_invalid_ref_missing_hash_raises_error(self): - """Reference without hash should raise InvalidVariableRefError.""" - with pytest.raises(InvalidVariableRefError): - parse_variable_ref("us:statute/26/32-eitc") - - def test_invalid_ref_empty_string_raises_error(self): - """Empty string should raise InvalidVariableRefError.""" - with pytest.raises(InvalidVariableRefError): - parse_variable_ref("") - - def test_invalid_ref_only_colon_raises_error(self): - """Reference with only colon should raise InvalidVariableRefError.""" - with pytest.raises(InvalidVariableRefError): - parse_variable_ref("us:") - - def test_invalid_ref_only_hash_raises_error(self): - """Reference with only hash should raise InvalidVariableRefError.""" - with pytest.raises(InvalidVariableRefError): - parse_variable_ref("#eitc") - - def test_parse_uk_reference(self): - """Should correctly parse a UK reference.""" - package, path, var_name = parse_variable_ref("uk:statute/ita/2007#personal_allowance") - - assert package == "uk" - assert path == "statute/ita/2007" - assert var_name == "personal_allowance" - - def test_parse_reference_with_multiple_hashes(self): - """Should handle reference with multiple # symbols (takes last).""" - package, path, var_name = parse_variable_ref("us:path/with#hash#var") - - assert package == "us" - assert path == "path/with#hash" - assert var_name == "var" diff --git a/tests/test_hierarchical_pipeline.py b/tests/test_hierarchical_pipeline.py deleted file mode 100644 index d5dba28c..00000000 --- a/tests/test_hierarchical_pipeline.py +++ /dev/null @@ -1,240 +0,0 @@ -""" -Tests for hierarchical microplex pipeline. - -Tests the pipeline that calibrates household weights using person-level targets -aggregated to the household level. -""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -# Import functions to test -from micro.us.hierarchical_pipeline import ( - build_hierarchical_constraints, - run_hierarchical_ipf, - _create_mock_hierarchical_data, -) - -# These imports may fail if Supabase not available -try: - from micro.us.hierarchical_pipeline import ( - load_hierarchical_data_from_supabase, - run_hierarchical_pipeline, - ) - SUPABASE_AVAILABLE = True -except ImportError: - SUPABASE_AVAILABLE = False - load_hierarchical_data_from_supabase = None - run_hierarchical_pipeline = None - - -class TestLoadHierarchicalData: - """Test loading hierarchical household + person data.""" - - def test_mock_data_loads_correctly(self): - """Test that mock data can be loaded.""" - hh_df, person_df = _create_mock_hierarchical_data(limit=100) - assert len(hh_df) > 0 - assert len(person_df) > 0 - assert "household_id" in hh_df.columns - assert "household_id" in person_df.columns - - def test_person_household_link(self): - """Test that persons properly link to households.""" - hh_df, person_df = _create_mock_hierarchical_data(limit=100) - # All person household_ids should exist in hh_df - person_hh_ids = set(person_df["household_id"].unique()) - hh_ids = set(hh_df["household_id"].unique()) - assert person_hh_ids.issubset(hh_ids) - - def test_mock_data_has_required_columns(self): - """Test that mock data has all required columns.""" - hh_df, person_df = _create_mock_hierarchical_data(limit=50) - # Household required columns - assert "weight" in hh_df.columns - assert "state_fips" in hh_df.columns - # Person required columns - assert "person_id" in person_df.columns - assert "age" in person_df.columns - - -class TestBuildHierarchicalConstraints: - """Test building constraint matrix for hierarchical calibration.""" - - @pytest.fixture - def sample_hierarchical_data(self): - """Create sample hierarchical microdata.""" - # 5 households with 2-3 persons each - hh_df = pd.DataFrame({ - "household_id": [1, 2, 3, 4, 5], - "weight": [100.0, 200.0, 150.0, 180.0, 120.0], - "state_fips": [6, 6, 36, 36, 48], # CA, CA, NY, NY, TX - }) - - person_df = pd.DataFrame({ - "person_id": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], - "household_id": [1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5], - "age": [35, 8, 45, 42, 12, 65, 62, 28, 6, 55, 50], - "employment_income": [50000, 0, 80000, 60000, 0, 0, 0, 45000, 0, 70000, 40000], - }) - - return hh_df, person_df - - def test_count_target_aggregation(self, sample_hierarchical_data): - """Test that person count targets aggregate to household level.""" - hh_df, person_df = sample_hierarchical_data - - # Target: count of persons aged 0-17 - mock_targets = [{ - "variable": "person_count", - "value": 500.0, # Total weighted children - "target_type": "count", - "constraints": [("age", ">=", "0"), ("age", "<", "18")], - }] - - constraints = build_hierarchical_constraints( - hh_df, person_df, mock_targets, hh_id_col="household_id", min_obs=1 - ) - - assert len(constraints) == 1 - # Indicator should be count of children per household - # HH1: 1 child (age 8), HH2: 1 child (age 12), HH3: 0, HH4: 1 (age 6), HH5: 0 - expected = np.array([1.0, 1.0, 0.0, 1.0, 0.0]) - np.testing.assert_array_equal(constraints[0]["indicator"], expected) - - def test_amount_target_aggregation(self, sample_hierarchical_data): - """Test that person amount targets sum to household level.""" - hh_df, person_df = sample_hierarchical_data - - # Target: total employment income - mock_targets = [{ - "variable": "employment_income", - "value": 50000000.0, # Total weighted income - "target_type": "amount", - "constraints": [], # All persons - }] - - constraints = build_hierarchical_constraints( - hh_df, person_df, mock_targets, hh_id_col="household_id", min_obs=1 - ) - - assert len(constraints) == 1 - # Indicator should be sum of income per household - # HH1: 50000, HH2: 140000, HH3: 0, HH4: 45000, HH5: 110000 - expected = np.array([50000.0, 140000.0, 0.0, 45000.0, 110000.0]) - np.testing.assert_array_equal(constraints[0]["indicator"], expected) - - -class TestHierarchicalIPF: - """Test IPF calibration with hierarchical constraints.""" - - def test_basic_convergence(self): - """Test that IPF converges with simple hierarchical data.""" - # Simple case: 3 households - original_weights = np.array([100.0, 200.0, 150.0]) - - # One constraint: total count = 500 - constraints = [{ - "indicator": np.array([1.0, 1.0, 1.0]), - "target_value": 500.0, - }] - - calibrated, success, loss = run_hierarchical_ipf( - original_weights, constraints, verbose=False - ) - - # Should converge - assert success - # Total should match target - total = np.sum(calibrated) - np.testing.assert_almost_equal(total, 500.0, decimal=1) - - def test_multiple_constraints(self): - """Test IPF with multiple hierarchical constraints.""" - original_weights = np.array([100.0, 100.0, 100.0, 100.0]) - - # Constraint 1: total = 500 - # Constraint 2: first two = 200 - constraints = [ - {"indicator": np.array([1.0, 1.0, 1.0, 1.0]), "target_value": 500.0}, - {"indicator": np.array([1.0, 1.0, 0.0, 0.0]), "target_value": 200.0}, - ] - - calibrated, success, loss = run_hierarchical_ipf( - original_weights, constraints, verbose=False - ) - - assert success - np.testing.assert_almost_equal(np.sum(calibrated), 500.0, decimal=1) - np.testing.assert_almost_equal(np.sum(calibrated[:2]), 200.0, decimal=1) - - -class TestRunHierarchicalPipeline: - """Test the full hierarchical pipeline.""" - - @pytest.mark.skipif(not SUPABASE_AVAILABLE, reason="Supabase not available") - def test_dry_run_returns_dataframe(self): - """Test that dry run returns calibrated DataFrame.""" - result = run_hierarchical_pipeline( - year=2024, dry_run=True, use_mock=True, limit=100, verbose=False - ) - assert result is not None - assert "calibrated_weight" in result.columns or "weight" in result.columns - - @pytest.mark.skipif(not SUPABASE_AVAILABLE, reason="Supabase not available") - def test_dry_run_does_not_write(self): - """Test that dry run doesn't write to database.""" - # Just ensure no exception is raised - run_hierarchical_pipeline( - year=2024, dry_run=True, use_mock=True, limit=50, verbose=False - ) - - -class TestIntegration: - """Integration tests using mock data.""" - - def test_full_mock_calibration(self): - """Test full calibration workflow with mock data.""" - # Create mock data - hh_df, person_df = _create_mock_hierarchical_data(limit=500) - - # Create realistic targets - total_weight = hh_df["weight"].sum() - n_children = person_df[person_df["age"] < 18].groupby("household_id").size() - weighted_children = (hh_df.set_index("household_id")["weight"] * n_children.reindex(hh_df["household_id"]).fillna(0).values).sum() - - targets = [ - { - "variable": "household_count", - "value": total_weight * 1.1, # 10% increase target - "target_type": "count", - "constraints": [], - }, - { - "variable": "person_count", - "value": weighted_children * 0.9, # 10% decrease for children - "target_type": "count", - "constraints": [("age", "<", "18")], - }, - ] - - # Build constraints - constraints = build_hierarchical_constraints( - hh_df, person_df, targets, hh_id_col="household_id" - ) - - assert len(constraints) == 2 - - # Run IPF - original_weights = hh_df["weight"].values.copy() - calibrated, success, loss = run_hierarchical_ipf( - original_weights, constraints, verbose=False - ) - - # Should converge (with some tolerance due to conflicting targets) - assert calibrated.sum() > 0 - # Weights should change - assert not np.allclose(calibrated, original_weights) diff --git a/tests/test_integration_simple.py b/tests/test_integration_simple.py deleted file mode 100644 index 14b3c88c..00000000 --- a/tests/test_integration_simple.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Simple integration tests for calibration pipeline. - -Tests minimal end-to-end workflow with small synthetic data. -""" - -import tempfile -from pathlib import Path - -import numpy as np -import pandas as pd -from sqlmodel import Session - -from calibration import ( - get_targets, - build_constraint_matrix, - EntropyCalibrator, - TargetSpec, -) -from db.schema import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetType, - init_db, -) - - -class TestSimpleEndToEnd: - """Simple integration tests with small data.""" - - def test_load_and_calibrate_synthetic_data(self): - """Test loading synthetic data and calibrating with simple constraints.""" - - # Step 1: Create simple microdata with uniform weights - np.random.seed(42) - n = 100 - microdata = pd.DataFrame({ - "weight": np.ones(n) * 1000.0, # Uniform weights - "is_tax_filer": np.random.choice([0, 1], n, p=[0.2, 0.8]), - "age": np.random.randint(18, 80, n), - }) - - # Step 2: Create simple target manually (no database needed) - current_filers = (microdata["weight"] * microdata["is_tax_filer"]).sum() - target_filers = current_filers * 1.1 # Want 10% more filers (modest change) - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=target_filers, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ) - ] - - # Step 3: Build constraints - constraints = build_constraint_matrix(microdata, targets) - assert len(constraints) == 1 - - # Step 4: Calibrate - calibrator = EntropyCalibrator(bounds=(0.5, 3.0), max_iterations=100) - original_weights = microdata["weight"].values - calibrated_weights = calibrator.calibrate(original_weights, constraints) - - # Step 5: Verify - assert len(calibrated_weights) == len(microdata) - assert all(calibrated_weights > 0) - - # Check constraint is satisfied - actual_filers = (calibrated_weights * microdata["is_tax_filer"]).sum() - error = abs(actual_filers - target_filers) / target_filers - assert error < 0.01, f"Error: {error:.2%}" - - def test_multiple_constraints_small_data(self): - """Test calibration with multiple constraints on small dataset.""" - - # Create small microdata - np.random.seed(123) - n = 50 - microdata = pd.DataFrame({ - "weight": np.ones(n) * 100.0, - "is_tax_filer": np.random.choice([0, 1], n, p=[0.3, 0.7]), - "age": np.random.randint(18, 80, n), - }) - - # Create two overlapping constraints - targets = [ - # All filers - TargetSpec( - variable="tax_unit_count", - value=4000.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ), - # Young filers (age < 40) - TargetSpec( - variable="tax_unit_count", - value=1500.0, - target_type=TargetType.COUNT, - constraints=[ - ("is_tax_filer", "==", "1"), - ("age", "<", "40"), - ], - source=DataSource.IRS_SOI, - period=2021, - ), - ] - - # Build and calibrate - constraints = build_constraint_matrix(microdata, targets) - calibrator = EntropyCalibrator(bounds=(0.5, 5.0), max_iterations=100) - calibrated_weights = calibrator.calibrate(microdata["weight"].values, constraints) - - # Verify both constraints - for i, constraint in enumerate(constraints): - actual = (calibrated_weights * constraint.indicator).sum() - error = abs(actual - constraint.target_value) / constraint.target_value - assert error < 0.02, ( - f"Constraint {i} failed: target={constraint.target_value}, " - f"actual={actual}, error={error:.2%}" - ) - - def test_amount_constraint_small_data(self): - """Test AMOUNT type constraint with small dataset.""" - - # Create microdata with income - np.random.seed(789) - n = 50 - microdata = pd.DataFrame({ - "weight": np.ones(n) * 100.0, - "is_tax_filer": np.ones(n, dtype=int), - "income": np.random.lognormal(10, 1, n), - }) - - # Target: adjust total income based on current value - current_income = (microdata["weight"] * microdata["income"]).sum() - target_income = current_income * 1.5 # Want 50% more total income - - targets = [ - TargetSpec( - variable="income", - value=target_income, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - ) - ] - - # Calibrate (use wider bounds to accommodate 50% change) - constraints = build_constraint_matrix(microdata, targets) - calibrator = EntropyCalibrator(bounds=(0.1, 10.0), max_iterations=200) - calibrated_weights = calibrator.calibrate(microdata["weight"].values, constraints) - - # Verify - actual_income = (calibrated_weights * microdata["income"]).sum() - error = abs(actual_income - target_income) / target_income - assert error < 0.01, f"Income error: {error:.2%}" - - def test_database_integration(self): - """Test querying targets from database and calibrating.""" - - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test.db" - engine = init_db(db_path) - - # Create simple target in database - with Session(engine) as session: - # Create stratum - constraints_def = [("is_tax_filer", "==", "1")] - stratum = Stratum( - name="all_filers", - jurisdiction=Jurisdiction.US, - description="All filers", - definition_hash=Stratum.compute_hash(constraints_def, Jurisdiction.US), - ) - session.add(stratum) - session.flush() - - # Add constraint - constraint = StratumConstraint( - stratum_id=stratum.id, - variable="is_tax_filer", - operator="==", - value="1", - ) - session.add(constraint) - - # Add target - target = Target( - stratum_id=stratum.id, - variable="tax_unit_count", - value=4000.0, - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - period=2021, - ) - session.add(target) - session.commit() - - # Query targets - targets = get_targets(db_path=db_path, jurisdiction="us", year=2021) - assert len(targets) == 1 - assert targets[0].value == 4000.0 - - # Create microdata and calibrate - np.random.seed(42) - n = 50 - microdata = pd.DataFrame({ - "weight": np.ones(n) * 100.0, - "is_tax_filer": np.random.choice([0, 1], n, p=[0.3, 0.7]), - }) - - constraints = build_constraint_matrix(microdata, targets) - calibrator = EntropyCalibrator(bounds=(0.5, 5.0)) - calibrated_weights = calibrator.calibrate(microdata["weight"].values, constraints) - - # Verify constraint satisfied - actual = (calibrated_weights * microdata["is_tax_filer"]).sum() - error = abs(actual - 4000.0) / 4000.0 - assert error < 0.01, f"Error: {error:.2%}" diff --git a/tests/test_legacy_targets_retired.py b/tests/test_legacy_targets_retired.py new file mode 100644 index 00000000..6490e9ed --- /dev/null +++ b/tests/test_legacy_targets_retired.py @@ -0,0 +1,306 @@ +"""The legacy target store and everything built on it stay retired. + +The store was a local SQLite ``strata`` / ``stratum_constraints`` / ``targets`` +schema filled by hand-written ``db/etl_*.py`` loaders, mirrored to Supabase +through ``arch.targets``, and read back by ``calibration/`` and the Microplex +pipeline in ``micro/us/``. Nothing outside that stack read it. These tests pin +its removal and the property the Thesis append path depends on: the gate and +catalog scripts import nothing from this repository's packages, so retiring +code under ``arch/``, ``db/`` or ``micro/`` cannot change what they execute. +""" + +from __future__ import annotations + +import ast +import importlib.util +import json +import os +import sqlite3 +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# Every module of the retired stack. A module listed here must not exist in +# the tree, must not import, and must not be imported by anything left. +RETIRED_MODULES = ( + "arch.normalization.targets", + "arch.targets", + "arch.targets.client", + "arch.targets.loaders", + "calibration", + "calibration.constraints", + "calibration.loader", + "calibration.methods", + "calibration.methods.entropy", + "calibration.targets", + "calibration.variables", + "db.etl_aca_enrollment", + "db.etl_bls", + "db.etl_cbo", + "db.etl_census", + "db.etl_cps", + "db.etl_cps_raw", + "db.etl_hmrc", + "db.etl_medicaid", + "db.etl_obr", + "db.etl_ons", + "db.etl_snap", + "db.etl_soi", + "db.etl_soi_credits", + "db.etl_soi_deductions", + "db.etl_soi_income_sources", + "db.etl_soi_state", + "db.etl_ssa", + "db.etl_ssi", + "db.etl_targets_supabase", + "micro.us.calibrate", + "micro.us.district", + "micro.us.hierarchical_pipeline", + "micro.us.hierarchy_validation", + "micro.us.pipeline", + "micro.us.targets", + "micro.us.validation_dashboard", + "scripts.export_to_json", +) + +# Names the retired modules exported. None may be defined or imported again. +RETIRED_NAMES = { + "db.schema": { + "GeographicLevel", + "Stratum", + "StratumConstraint", + "Target", + "TargetType", + }, + "db.supabase_client": { + "TARGETS_SCHEMA", + "insert_targets_batch", + "query_strata", + "query_targets", + }, + "arch.normalization": {"TargetBlueprint", "as_target", "target_kwargs"}, +} + +LEGACY_TARGET_TABLES = {"strata", "stratum_constraints", "targets"} +SOURCE_ARTIFACT_TABLES = { + "source_artifacts", + "source_columns", + "source_rows", + "source_tables", +} + +# The scripts the Thesis append path runs. The append gate executes the +# base commit's check_thesis_facts_append.py (which loads receipt_pins.py); +# the push-path CI step runs verify_release_chain.py; the resolver in +# ThesisInstitute/thesis runs build_series_catalog.py with its own interpreter; +# cut_release_manifest.py and canonical_json.py cut and hash releases. +APPEND_PATH_SCRIPTS = ( + "build_series_catalog", + "canonical_json", + "check_thesis_facts_append", + "cut_release_manifest", + "receipt_pins", + "verify_release_chain", +) +# Top-level modules that must never load on the append path: this +# repository's own packages and the third-party libraries only the retired +# stack and the source-artifact store use. +FORBIDDEN_APPEND_PATH_MODULES = { + "arch", + "calibration", + "db", + "micro", + "packages", + "sqlalchemy", + "sqlmodel", + "supabase", +} + + +def _module_path(module: str) -> Path: + return REPO_ROOT.joinpath(*module.split(".")) + + +@pytest.mark.parametrize("module", RETIRED_MODULES) +def test_retired_module_is_not_in_the_tree(module): + path = _module_path(module) + + assert not path.with_suffix(".py").exists() + assert not (path / "__init__.py").exists() + + +@pytest.mark.parametrize( + "module", [m for m in RETIRED_MODULES if not m.startswith("scripts.")] +) +def test_retired_module_does_not_import(module): + # find_spec imports the parent package first, so a submodule of a removed + # package raises instead of returning None. Both mean "not importable". + try: + spec = importlib.util.find_spec(module) + except ModuleNotFoundError: + spec = None + assert spec is None + + +def test_no_db_etl_loader_remains(): + assert sorted(path.name for path in (REPO_ROOT / "db").glob("etl_*.py")) == [] + + +def _python_files(): + for root in ("arch", "db", "micro", "packages", "scripts", "tests", "data"): + yield from sorted((REPO_ROOT / root).rglob("*.py")) + + +def _imported_modules(tree: ast.AST, package: str) -> set[str]: + imported: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + if node.level: + parts = package.split(".") if package else [] + if node.level > 1: + parts = parts[: len(parts) - (node.level - 1)] + base = ".".join([*parts, node.module] if node.module else parts) + else: + base = node.module or "" + imported.add(base) + imported.update(f"{base}.{alias.name}" for alias in node.names) + return imported + + +def test_nothing_left_imports_the_retired_stack(): + retired = set(RETIRED_MODULES) + offenders = [] + for path in _python_files(): + relative = path.relative_to(REPO_ROOT) + # The containing package: for pkg/__init__.py and pkg/mod.py alike, + # drop the last path component. + package = ".".join(relative.with_suffix("").parts[:-1]) + tree = ast.parse(path.read_text(), filename=str(relative)) + for name in _imported_modules(tree, package): + if any(name == r or name.startswith(f"{r}.") for r in retired): + offenders.append(f"{relative}: {name}") + module, _, attribute = name.rpartition(".") + if attribute in RETIRED_NAMES.get(module, ()): + offenders.append(f"{relative}: {name}") + + assert offenders == [] + + +def test_surviving_modules_do_not_define_retired_names(): + import arch.normalization as normalization + import db.schema as schema + import db.supabase_client as supabase_client + + modules = { + "db.schema": schema, + "db.supabase_client": supabase_client, + "arch.normalization": normalization, + } + for module_name, names in RETIRED_NAMES.items(): + present = sorted(n for n in names if hasattr(modules[module_name], n)) + assert present == [], module_name + + +@pytest.mark.parametrize( + "argv", + [ + ["load", "all"], + ["load", "soi-state"], + ["stats"], + ["query", "--variable", "tax_unit_count"], + ], +) +def test_cli_rejects_retired_target_commands(monkeypatch, capsys, argv): + from db import cli + + monkeypatch.setattr(sys, "argv", ["arch", *argv]) + with pytest.raises(SystemExit) as excinfo: + cli.main() + + assert excinfo.value.code == 2 + assert "invalid choice" in capsys.readouterr().err + + +def test_init_creates_only_source_artifact_tables(monkeypatch, tmp_path): + from db import cli + + db_path = tmp_path / "arch.db" + monkeypatch.setattr(sys, "argv", ["arch", "--db", str(db_path), "init"]) + cli.main() + + with sqlite3.connect(db_path) as connection: + tables = { + name + for (name,) in connection.execute( + "SELECT name FROM sqlite_master WHERE type = 'table'" + ) + } + + assert tables == SOURCE_ARTIFACT_TABLES + assert not tables & LEGACY_TARGET_TABLES + + +def test_source_stats_reads_an_initialized_store(monkeypatch, capsys, tmp_path): + from db import cli + + db_path = tmp_path / "arch.db" + for command in ("init", "source-stats"): + monkeypatch.setattr(sys, "argv", ["arch", "--db", str(db_path), command]) + cli.main() + + output = capsys.readouterr().out + assert "Source artifacts: 0" in output + assert "Source rows: 0" in output + + +def test_append_path_scripts_load_nothing_from_the_repository_packages(tmp_path): + """Import every append-path script the way its callers do and list modules. + + The scripts run with ``scripts/`` on ``sys.path`` and the repository root + off it (the gate sets ``PYTHONPATH`` to the base checkout's ``scripts/`` + and runs from the runner's temp directory; the resolver runs the staged + generator by path). ``-I`` drops the environment and the current directory + from ``sys.path``; the scripts directory is added back explicitly. + """ + probe = ( + "import json, sys\n" + f"sys.path.insert(0, {str(REPO_ROOT / 'scripts')!r})\n" + f"for name in {APPEND_PATH_SCRIPTS!r}:\n" + " __import__(name)\n" + "print(json.dumps(sorted({m.split('.')[0] for m in sys.modules})))\n" + ) + environment = { + key: value + for key, value in os.environ.items() + if not key.startswith("PYTHON") + } + completed = subprocess.run( + [sys.executable, "-I", "-c", probe], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + check=True, + ) + loaded = set(json.loads(completed.stdout)) + + assert set(APPEND_PATH_SCRIPTS) <= loaded + assert loaded & FORBIDDEN_APPEND_PATH_MODULES == set() + + +def test_ci_does_not_load_the_retired_target_store(): + workflow = (REPO_ROOT / ".github" / "workflows" / "ci.yml").read_text() + + assert " load all" not in workflow + assert "arch-targets-ci.db" not in workflow + for module in RETIRED_MODULES: + assert f"import {module}\n" not in workflow + assert f"from {module} import" not in workflow + assert module.replace(".", "/") + ".py" not in workflow + assert module.replace(".", "/") + "/" not in workflow diff --git a/tests/test_loader.py b/tests/test_loader.py deleted file mode 100644 index 768d2df2..00000000 --- a/tests/test_loader.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Tests for microdata loader.""" - -import pandas as pd -import pytest - -from calibration.loader import load_microdata - - -class TestLoadMicrodataSynthetic: - """Tests for synthetic microdata generation.""" - - def test_returns_dataframe(self): - """load_microdata should return a pandas DataFrame.""" - df = load_microdata(source="synthetic", year=2021) - assert isinstance(df, pd.DataFrame) - - def test_has_weight_column(self): - """Synthetic data should have a weight column.""" - df = load_microdata(source="synthetic", year=2021) - assert "weight" in df.columns - - def test_has_required_columns(self): - """Synthetic data should have demographic and income columns.""" - df = load_microdata(source="synthetic", year=2021) - - required_columns = [ - "weight", - "age", - "income", - "employment_status", - "has_children", - "state_fips", - ] - - for col in required_columns: - assert col in df.columns, f"Missing required column: {col}" - - def test_approximately_10000_rows(self): - """Synthetic data should have approximately 10,000 rows.""" - df = load_microdata(source="synthetic", year=2021) - assert 9000 <= len(df) <= 11000 - - def test_age_range_realistic(self): - """Ages should be in realistic range (0-120).""" - df = load_microdata(source="synthetic", year=2021) - assert df["age"].min() >= 0 - assert df["age"].max() <= 120 - - def test_weights_positive(self): - """All weights should be positive.""" - df = load_microdata(source="synthetic", year=2021) - assert (df["weight"] > 0).all() - - def test_employment_status_values(self): - """Employment status should be encoded correctly.""" - df = load_microdata(source="synthetic", year=2021) - # Employment status: 0=not in labor force, 1=employed, 2=unemployed - valid_values = {0, 1, 2} - assert set(df["employment_status"].unique()).issubset(valid_values) - - def test_has_children_binary(self): - """has_children should be binary (0 or 1).""" - df = load_microdata(source="synthetic", year=2021) - assert set(df["has_children"].unique()).issubset({0, 1}) - - def test_state_fips_valid(self): - """State FIPS codes should be valid US state codes.""" - df = load_microdata(source="synthetic", year=2021) - # FIPS codes are 1-56 (with some gaps) - fips = df["state_fips"].astype(int) - assert fips.min() >= 1 - assert fips.max() <= 56 - - def test_income_nonnegative(self): - """Income should be non-negative.""" - df = load_microdata(source="synthetic", year=2021) - assert (df["income"] >= 0).all() - - def test_reproducible_with_seed(self): - """Same seed should produce same data.""" - df1 = load_microdata(source="synthetic", year=2021, seed=42) - df2 = load_microdata(source="synthetic", year=2021, seed=42) - pd.testing.assert_frame_equal(df1, df2) - - def test_different_seeds_different_data(self): - """Different seeds should produce different data.""" - df1 = load_microdata(source="synthetic", year=2021, seed=42) - df2 = load_microdata(source="synthetic", year=2021, seed=123) - assert not df1.equals(df2) - - -class TestLoadMicrodataVariables: - """Tests for variable filtering.""" - - def test_filter_variables(self): - """Should return only requested variables plus weight.""" - df = load_microdata( - source="synthetic", - year=2021, - variables=["age", "income"], - ) - # Should have weight plus requested variables - assert set(df.columns) == {"weight", "age", "income"} - - def test_weight_always_included(self): - """Weight should always be included even if not specified.""" - df = load_microdata( - source="synthetic", - year=2021, - variables=["age"], - ) - assert "weight" in df.columns - - -class TestLoadMicrodataCPS: - """Tests for CPS data loading (with fallback to synthetic).""" - - def test_cps_source_returns_data(self): - """CPS source should return data (synthetic fallback if not available).""" - df = load_microdata(source="cps", year=2021) - assert isinstance(df, pd.DataFrame) - assert len(df) > 0 - assert "weight" in df.columns - - def test_cps_invalid_year_raises(self): - """Invalid year should raise ValueError.""" - with pytest.raises(ValueError, match="[Yy]ear"): - load_microdata(source="cps", year=1800) - - -class TestLoadMicrodataEdgeCases: - """Tests for edge cases.""" - - def test_invalid_source_raises(self): - """Invalid source should raise ValueError.""" - with pytest.raises(ValueError, match="[Ss]ource"): - load_microdata(source="invalid_source", year=2021) - - def test_future_year_raises(self): - """Future year should raise ValueError.""" - with pytest.raises(ValueError, match="[Yy]ear"): - load_microdata(source="synthetic", year=2099) diff --git a/tests/test_microplex_hierarchy_validation.py b/tests/test_microplex_hierarchy_validation.py deleted file mode 100644 index 1df53f06..00000000 --- a/tests/test_microplex_hierarchy_validation.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Tests for flat-vs-household Microplex calibration comparison.""" - -import json - -import pandas as pd - -from arch.targets import DataSource, TargetSpec, TargetType -from micro.us.hierarchy_validation import ( - compare_flat_vs_household_calibration, - compare_target_diagnostics, -) -from micro.us.targets import TargetCompositionResult - - -def test_compare_target_diagnostics_aligns_targets(): - flat = pd.DataFrame( - { - "target_index": [0], - "source": ["irs-soi"], - "variable": ["tax_unit_count"], - "target_type": ["count"], - "period": [2024], - "source_period": [2024], - "stratum": ["All filers"], - "constraints": ["[]"], - "target_value": [2.0], - "role": ["active"], - "status": ["used"], - "drop_reason": [""], - "n_obs": [2], - "pre_error": [0.2], - "post_error": [0.01], - } - ) - household = flat.copy() - household["post_error"] = -0.03 - - comparison = compare_target_diagnostics(flat, household) - - assert len(comparison) == 1 - assert comparison.iloc[0]["flat_abs_post_error"] == 0.01 - assert comparison.iloc[0]["household_abs_post_error"] == 0.03 - assert abs(comparison.iloc[0]["abs_post_error_delta"] - 0.02) < 1e-12 - - -def test_compare_flat_vs_household_calibration_writes_artifacts( - tmp_path, - monkeypatch, -): - cps = pd.DataFrame( - { - "household_id": [1, 1, 2], - "tax_unit_id": [10, 10, 20], - "person_seq": [1, 2, 1], - "age": [40, 38, 55], - "state_fips": [6, 6, 48], - "weight": [100.0, 100.0, 200.0], - "total_person_income": [50_000.0, 20_000.0, 30_000.0], - "wage_salary_income": [50_000.0, 20_000.0, 30_000.0], - "self_employment_income": [0.0, 0.0, 0.0], - } - ) - cps_path = tmp_path / "cps.parquet" - output_dir = tmp_path / "comparison" - cps.to_parquet(cps_path, index=False) - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=300.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ) - ] - - def fake_compose(*args, **kwargs): - return TargetCompositionResult(targets=targets, diagnostics=pd.DataFrame()) - - monkeypatch.setattr( - "micro.us.hierarchy_validation.compose_microplex_targets", - fake_compose, - ) - - result = compare_flat_vs_household_calibration( - year=2024, - cps_path=cps_path, - limit=3, - output_dir=output_dir, - version_id="test-microplex-v1", - min_target_obs=1, - add_policyengine_tax=False, - verbose=False, - ) - - assert result.flat_result.calibration_unit == "tax_unit" - assert result.household_result.calibration_unit == "household" - assert len(result.target_comparison) == 1 - assert (output_dir / "summary.csv").exists() - assert (output_dir / "target_comparison.csv").exists() - assert (output_dir / "manifest.json").exists() - assert (output_dir / "metrics.json").exists() - assert (output_dir / "dashboard.html").exists() - assert (output_dir / "variable_summary.csv").exists() - assert (output_dir / "worst_targets.csv").exists() - assert set(result.summary["group"]) >= {"flat", "household", "delta"} - - manifest = json.loads((output_dir / "manifest.json").read_text()) - metrics = json.loads((output_dir / "metrics.json").read_text()) - dashboard = (output_dir / "dashboard.html").read_text() - assert manifest["microplex_version"] == "test-microplex-v1" - assert manifest["report_type"] == "microplex_calibration_dashboard" - assert manifest["artifacts"]["dashboard"] == "dashboard.html" - assert metrics["summary"]["flat"]["calibration_unit"] == "tax_unit" - assert metrics["summary"]["household"]["calibration_unit"] == "household" - assert "Microplex Calibration Dashboard" in dashboard - - -def test_compare_flat_vs_household_calibration_uses_reports_root( - tmp_path, - monkeypatch, -): - cps = pd.DataFrame( - { - "household_id": [1], - "tax_unit_id": [10], - "person_seq": [1], - "age": [40], - "state_fips": [6], - "weight": [100.0], - "total_person_income": [50_000.0], - "wage_salary_income": [50_000.0], - "self_employment_income": [0.0], - } - ) - cps_path = tmp_path / "cps.parquet" - reports_root = tmp_path / "reports" - cps.to_parquet(cps_path, index=False) - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=100.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ) - ] - - def fake_compose(*args, **kwargs): - return TargetCompositionResult(targets=targets, diagnostics=pd.DataFrame()) - - monkeypatch.setattr( - "micro.us.hierarchy_validation.compose_microplex_targets", - fake_compose, - ) - - compare_flat_vs_household_calibration( - year=2024, - cps_path=cps_path, - reports_root=reports_root, - version_id="test-version-root", - min_target_obs=1, - add_policyengine_tax=False, - verbose=False, - ) - - assert (reports_root / "test-version-root" / "dashboard.html").exists() diff --git a/tests/test_microplex_pipeline.py b/tests/test_microplex_pipeline.py deleted file mode 100644 index ec9cdffa..00000000 --- a/tests/test_microplex_pipeline.py +++ /dev/null @@ -1,366 +0,0 @@ -"""Tests for the Microplex build pipeline.""" - -import pandas as pd - -from arch.targets import DataSource, TargetSpec, TargetType -from micro.us import pipeline as microplex -from micro.us.targets import TargetCompositionResult - - -def test_build_tax_units_accepts_local_cps_columns(): - cps = pd.DataFrame( - { - "total_person_income": [20_000, 0, 50_000], - "wage_salary_income": [20_000, 0, 45_000], - "self_employment_income": [0, 0, 5_000], - "farm_self_employment_income": [0, 0, 0], - "weight": [100.0, 200.0, 300.0], - "state_fips": [6, 6, 36], - "age": [30, 10, 40], - } - ) - - tax_units = microplex.build_tax_units(cps) - - assert len(tax_units) == 2 - assert tax_units["weight"].tolist() == [100.0, 300.0] - assert tax_units["is_tax_filer"].tolist() == [1, 1] - assert tax_units["adjusted_gross_income"].tolist() == [20_000.0, 49_808.75] - - -def test_build_tax_units_aggregates_census_tax_unit_ids(): - cps = pd.DataFrame( - { - "household_id": [1, 1, 1], - "tax_unit_id": [101, 101, 102], - "person_seq": [1, 2, 3], - "age": [40, 38, 10], - "weight": [100.0, 100.0, 80.0], - "total_person_income": [50_000, 20_000, 0], - "wage_salary_income": [50_000, 20_000, 0], - "self_employment_income": [0, 0, 0], - "farm_self_employment_income": [0, 0, 0], - "interest_income": [100, 50, 0], - "dividend_income": [200, 100, 0], - "rental_income": [0, 0, 0], - "unemployment_compensation": [0, 0, 0], - "other_income": [0, 0, 0], - } - ) - - tax_units = microplex.build_tax_units(cps) - - assert len(tax_units) == 1 - assert tax_units.iloc[0]["tax_unit_id"] == 101 - assert tax_units.iloc[0]["weight"] == 100.0 - assert tax_units.iloc[0]["person_count"] == 2 - assert tax_units.iloc[0]["wage_income"] == 70_000 - assert tax_units.iloc[0]["adjusted_gross_income"] == 70_450 - - -def test_run_pipeline_can_write_local_microplex(tmp_path, monkeypatch): - n = 150 - cps = pd.DataFrame( - { - "total_person_income": [20_000] * n, - "wage_salary_income": [20_000] * n, - "self_employment_income": [0] * n, - "weight": [100.0] * n, - } - ) - cps_path = tmp_path / "cps.parquet" - output_path = tmp_path / "microplex.parquet" - cps.to_parquet(cps_path, index=False) - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=15_000.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ) - ] - monkeypatch.setattr( - microplex, - "compose_microplex_targets", - lambda *args, **kwargs: TargetCompositionResult( - targets=targets, - diagnostics=pd.DataFrame(), - ), - ) - - result = microplex.run_pipeline( - year=2024, - limit=n, - cps_path=cps_path, - output_path=output_path, - ) - - assert output_path.exists() - assert len(result) == n - assert "weight_adjustment" in result.columns - assert pd.read_parquet(output_path).shape[0] == n - - -def test_run_pipeline_writes_linked_entity_outputs(tmp_path, monkeypatch): - cps = pd.DataFrame( - { - "household_id": [1, 1, 2], - "tax_unit_id": [10, 10, 20], - "person_seq": [1, 2, 1], - "age": [40, 38, 55], - "state_fips": [6, 6, 48], - "weight": [100.0, 100.0, 200.0], - "total_person_income": [50_000.0, 20_000.0, 30_000.0], - "wage_salary_income": [50_000.0, 20_000.0, 30_000.0], - "self_employment_income": [0.0, 0.0, 0.0], - } - ) - cps_path = tmp_path / "cps.parquet" - entity_output_dir = tmp_path / "microplex_entities" - cps.to_parquet(cps_path, index=False) - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=300.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ) - ] - monkeypatch.setattr( - microplex, - "compose_microplex_targets", - lambda *args, **kwargs: TargetCompositionResult( - targets=targets, - diagnostics=pd.DataFrame(), - ), - ) - - result = microplex.run_pipeline( - year=2024, - limit=3, - cps_path=cps_path, - entity_output_dir=entity_output_dir, - dry_run=True, - min_target_obs=1, - ) - - households = pd.read_parquet(entity_output_dir / "households.parquet") - persons = pd.read_parquet(entity_output_dir / "persons.parquet") - tax_units = pd.read_parquet(entity_output_dir / "tax_units.parquet") - - assert len(result) == 2 - assert len(households) == 2 - assert len(persons) == 3 - assert len(tax_units) == 2 - assert households["weight_adjustment"].notna().all() - assert persons["household_entity_id"].isin(households["household_entity_id"]).all() - assert tax_units["household_entity_id"].isin(households["household_entity_id"]).all() - - -def test_calibrate_weights_preserves_per_target_diagnostics(): - df = pd.DataFrame( - { - "weight": [1.0, 1.0, 1.0], - "is_tax_filer": [1, 1, 1], - "adjusted_gross_income": [10_000.0, 50_000.0, 100_000.0], - } - ) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=3.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - TargetSpec( - variable="adjusted_gross_income", - value=160_000.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - ] - - result = microplex.calibrate_weights( - df, - targets, - include_amounts=True, - min_obs=1, - verbose=False, - ) - - assert len(result.diagnostics) == 2 - assert result.diagnostics["status"].tolist() == ["used", "used"] - assert set(result.targets_after) == { - "tax_unit_count|count|All filers", - "adjusted_gross_income|amount|All filers", - } - - -def test_calibrate_weights_roles_diagnostic_and_unsupported_targets(): - df = pd.DataFrame( - { - "weight": [1.0, 1.0, 1.0], - "is_tax_filer": [1, 1, 1], - "adjusted_gross_income": [10_000.0, 20_000.0, 30_000.0], - } - ) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=3.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - TargetSpec( - variable="tax_unit_count", - value=1.0, - target_type=TargetType.COUNT, - constraints=[("adjusted_gross_income", ">=", "25000")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="High AGI sparse cell", - ), - TargetSpec( - variable="adjusted_gross_income", - value=60_000.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="AGI holdout", - ), - TargetSpec( - variable="snap_benefit", - value=1.0, - target_type=TargetType.AMOUNT, - constraints=[], - source=DataSource.USDA_SNAP, - period=2024, - stratum_name="SNAP benefits", - ), - ] - - result = microplex.calibrate_weights( - df, - targets, - min_obs=2, - holdout_variables=("adjusted_gross_income",), - verbose=False, - ) - - diagnostics = result.diagnostics.set_index("stratum") - assert diagnostics.loc["All filers", "role"] == "active" - assert diagnostics.loc["High AGI sparse cell", "role"] == "diagnostic" - assert diagnostics.loc["High AGI sparse cell", "drop_reason"] == "insufficient_obs" - assert diagnostics.loc["High AGI sparse cell", "post_value"] == 1.0 - assert diagnostics.loc["AGI holdout", "role"] == "holdout" - assert diagnostics.loc["AGI holdout", "post_value"] == 60_000.0 - assert diagnostics.loc["SNAP benefits", "role"] == "unsupported" - assert pd.isna(diagnostics.loc["SNAP benefits", "post_value"]) - - -def test_generalized_rake_calibrates_count_and_amount_targets(): - df = pd.DataFrame( - { - "weight": [1.0, 1.0, 1.0, 1.0], - "is_tax_filer": [1, 1, 1, 1], - "adjusted_gross_income": [10_000.0, 20_000.0, 80_000.0, 90_000.0], - } - ) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=4.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - TargetSpec( - variable="adjusted_gross_income", - value=240_000.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - ] - - result = microplex.calibrate_weights( - df, - targets, - include_amounts=True, - min_obs=1, - calibration_method="generalized-rake", - verbose=False, - ) - - assert result.success - assert result.method == "generalized-rake" - count_error = result.targets_after["tax_unit_count|count|All filers"]["error"] - amount_error = result.targets_after[ - "adjusted_gross_income|amount|All filers" - ]["error"] - assert abs(count_error) < 0.01 - assert abs(amount_error) < 0.01 - - -def test_build_constraints_maps_employment_income_to_wages(): - df = pd.DataFrame( - { - "weight": [1.0, 1.0, 1.0], - "is_tax_filer": [1, 1, 1], - "adjusted_gross_income": [10_000.0, 20_000.0, 30_000.0], - "wage_income": [10_000.0, 0.0, 40_000.0], - } - ) - targets = [ - TargetSpec( - variable="employment_income", - value=2.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="Wage filers", - ), - TargetSpec( - variable="employment_income", - value=50_000.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="Wage filers", - ), - ] - - constraints = microplex.build_constraints_from_target_specs( - df, - targets, - min_obs=1, - include_amounts=True, - ) - - assert len(constraints) == 2 - assert constraints[0]["indicator"].tolist() == [1.0, 0.0, 1.0] - assert constraints[1]["indicator"].tolist() == [10_000.0, 0.0, 40_000.0] diff --git a/tests/test_microplex_policyengine.py b/tests/test_microplex_policyengine.py index 2058ed9c..fbee5c29 100644 --- a/tests/test_microplex_policyengine.py +++ b/tests/test_microplex_policyengine.py @@ -8,8 +8,6 @@ import numpy as np import pandas as pd -from arch.targets import DataSource, TargetSpec, TargetType -from micro.us import pipeline as microplex from micro.us.policyengine import ( DEFAULT_SOI_INCOME_TAX_VARIABLE, PolicyEngineNotAvailableError, @@ -126,56 +124,6 @@ def test_add_policyengine_income_tax_from_persons_preserves_entities(monkeypatch ] -def test_pipeline_merges_hierarchical_policyengine_tax(monkeypatch): - install_fake_policyengine(monkeypatch) - targets = [ - TargetSpec( - variable="income_tax_liability", - value=10_000.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="Tax liability holdout", - ) - ] - tax_units = pd.DataFrame( - { - "household_id": [1, 2], - "tax_unit_id": [10, 20], - "wage_income": [70_000.0, 0.0], - "self_employment_income": [0.0, 20_000.0], - } - ) - persons = pd.DataFrame( - { - "household_id": [1, 1, 2], - "tax_unit_id": [10, 10, 20], - "person_seq": [1, 2, 1], - "age": [40, 38, 55], - "state_fips": [6, 6, 48], - "weight": [100.0, 100.0, 200.0], - "wage_salary_income": [50_000.0, 20_000.0, 0.0], - "self_employment_income": [0.0, 0.0, 20_000.0], - } - ) - - result = microplex.maybe_add_policyengine_income_tax( - tax_units, - targets, - year=2024, - persons=persons, - ) - - assert result["income_tax_liability"].tolist() == [7_000.0, 3_000.0] - assert len(FakeSimulation.calls) == 1 - situation = FakeSimulation.calls[0][0] - assert [len(unit["members"]) for unit in situation["tax_units"].values()] == [ - 2, - 1, - ] - - def test_policyengine_missing_raises_clear_error(monkeypatch): monkeypatch.setitem(sys.modules, "policyengine_us", None) @@ -185,103 +133,3 @@ def test_policyengine_missing_raises_clear_error(monkeypatch): assert "PolicyEngine-US is required" in str(exc) else: raise AssertionError("Expected PolicyEngineNotAvailableError") - - -def test_run_pipeline_calculates_income_tax_when_targets_require_it( - tmp_path, - monkeypatch, -): - install_fake_policyengine(monkeypatch) - cps = pd.DataFrame( - { - "total_person_income": [50_000.0] * 150, - "wage_salary_income": [50_000.0] * 150, - "self_employment_income": [0.0] * 150, - "weight": [100.0] * 150, - } - ) - cps_path = tmp_path / "cps.parquet" - cps.to_parquet(cps_path, index=False) - - targets = [ - TargetSpec( - variable="tax_unit_count", - value=15_000.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - TargetSpec( - variable="income_tax_liability", - value=750_000.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="Tax liability holdout", - ), - ] - monkeypatch.setattr( - microplex, - "compose_microplex_targets", - lambda *args, **kwargs: microplex.TargetCompositionResult( - targets=targets, - diagnostics=pd.DataFrame(), - ), - ) - - result = microplex.run_pipeline( - year=2024, - limit=150, - cps_path=cps_path, - dry_run=True, - min_target_obs=1, - ) - - assert "income_tax_liability" in result - assert result["income_tax_liability"].iloc[0] == 5_000.0 - - -def test_calibrate_weights_evaluates_income_tax_holdout(): - df = pd.DataFrame( - { - "weight": [1.0, 1.0, 1.0], - "is_tax_filer": [1, 1, 1], - "adjusted_gross_income": [10_000.0, 20_000.0, 30_000.0], - "income_tax_liability": [100.0, 200.0, 300.0], - } - ) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=3.0, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All filers", - ), - TargetSpec( - variable="income_tax_liability", - value=600.0, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="Tax liability holdout", - ), - ] - - result = microplex.calibrate_weights( - df, - targets, - min_obs=1, - holdout_variables=("income_tax_liability",), - verbose=False, - ) - - diagnostics = result.diagnostics.set_index("stratum") - assert diagnostics.loc["Tax liability holdout", "role"] == "holdout" - assert diagnostics.loc["Tax liability holdout", "post_value"] == 600.0 diff --git a/tests/test_microplex_target_adapter.py b/tests/test_microplex_target_adapter.py deleted file mode 100644 index 0da15ee6..00000000 --- a/tests/test_microplex_target_adapter.py +++ /dev/null @@ -1,391 +0,0 @@ -"""Tests for Arch target adapters used by Microplex.""" - -import numpy as np -import pandas as pd -from sqlmodel import Session - -from arch.targets import ( - DataSource, - Jurisdiction, - Stratum, - StratumConstraint, - Target, - TargetSpec, - TargetType, - init_db, -) -from micro.us.pipeline import build_constraints_from_target_specs -from micro.us.targets import ( - MicroplexTargetProfile, - age_soi_targets, - build_hierarchical_microplex_constraints, - build_microplex_constraints, - compose_microplex_targets, - constraints_to_ipf_dicts, - get_soi_aging_factors, - load_microplex_targets, -) - - -def _insert_simple_target(db_path): - engine = init_db(db_path) - with Session(engine) as session: - constraints = [("is_tax_filer", "==", "1")] - stratum = Stratum( - name="All tax filers", - description="All tax filers", - jurisdiction=Jurisdiction.US, - definition_hash=Stratum.compute_hash(constraints, Jurisdiction.US), - ) - session.add(stratum) - session.flush() - - session.add( - StratumConstraint( - stratum_id=stratum.id, - variable="is_tax_filer", - operator="==", - value="1", - ) - ) - session.add( - Target( - stratum_id=stratum.id, - variable="tax_unit_count", - period=2024, - value=1000, - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - ) - ) - session.commit() - - -def _add_stratum(session, name, jurisdiction, constraints): - stratum = Stratum( - name=name, - description=name, - jurisdiction=jurisdiction, - definition_hash=Stratum.compute_hash(constraints, jurisdiction), - ) - session.add(stratum) - session.flush() - for variable, operator, value in constraints: - session.add( - StratumConstraint( - stratum_id=stratum.id, - variable=variable, - operator=operator, - value=value, - ) - ) - return stratum - - -def _insert_soi_aging_inputs(db_path): - engine = init_db(db_path) - with Session(engine) as session: - tax_filers = _add_stratum( - session, - "All tax filers", - Jurisdiction.US, - [("is_tax_filer", "==", "1")], - ) - economy = _add_stratum(session, "US Economy", Jurisdiction.US, []) - - session.add_all( - [ - Target( - stratum_id=tax_filers.id, - variable="tax_unit_count", - period=2021, - value=100, - target_type=TargetType.COUNT, - source=DataSource.IRS_SOI, - ), - Target( - stratum_id=tax_filers.id, - variable="adjusted_gross_income", - period=2021, - value=1_000, - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - ), - Target( - stratum_id=economy.id, - variable="labor_force_count", - period=2021, - value=100, - target_type=TargetType.COUNT, - source=DataSource.BLS, - ), - Target( - stratum_id=tax_filers.id, - variable="adjusted_gross_income", - period=2022, - value=1_100, - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - ), - Target( - stratum_id=tax_filers.id, - variable="adjusted_gross_income", - period=2023, - value=1_210, - target_type=TargetType.AMOUNT, - source=DataSource.IRS_SOI, - ), - Target( - stratum_id=economy.id, - variable="labor_force_count", - period=2023, - value=105, - target_type=TargetType.COUNT, - source=DataSource.BLS, - ), - Target( - stratum_id=economy.id, - variable="labor_force", - period=2024, - value=110, - target_type=TargetType.COUNT, - source=DataSource.CBO, - ), - ] - ) - session.commit() - - -def test_load_microplex_targets_reads_arch_db(tmp_path): - db_path = tmp_path / "targets.db" - _insert_simple_target(db_path) - - targets = load_microplex_targets(db_path=db_path, jurisdiction="us", year=2024) - - assert len(targets) == 1 - assert targets[0] == TargetSpec( - variable="tax_unit_count", - value=1000, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="All tax filers", - ) - - -def test_get_soi_aging_factors_uses_labor_force_and_aggregate_agi(tmp_path): - db_path = tmp_path / "targets.db" - _insert_soi_aging_inputs(db_path) - - factors = get_soi_aging_factors( - source_year=2021, - target_year=2024, - db_path=db_path, - ) - - assert factors.count_factor == 1.1 - assert np.isclose(factors.amount_factor, 1.331) - assert factors.count_method == "cbo_labor_force_ratio" - assert ( - factors.amount_method - == "soi_total_agi_last_growth_extrapolation" - ) - - -def test_age_soi_targets_scales_soi_values_and_preserves_others(tmp_path): - db_path = tmp_path / "targets.db" - _insert_soi_aging_inputs(db_path) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=100, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - stratum_name="All tax filers", - ), - TargetSpec( - variable="adjusted_gross_income", - value=1_000, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2021, - stratum_name="All tax filers", - ), - TargetSpec( - variable="population", - value=300, - target_type=TargetType.COUNT, - constraints=[], - source=DataSource.CENSUS_ACS, - period=2024, - stratum_name="US population", - ), - ] - - aged = age_soi_targets(targets, target_year=2024, db_path=db_path) - - assert aged[0].period == 2024 - assert np.isclose(aged[0].value, 110) - assert "SOI aged 2021->2024" in aged[0].stratum_name - assert aged[1].period == 2024 - assert np.isclose(aged[1].value, 1_331) - assert aged[2] is targets[2] - - -def test_compose_microplex_targets_keeps_current_records_and_ages_fallback_soi( - tmp_path, -): - db_path = tmp_path / "targets.db" - _insert_soi_aging_inputs(db_path) - - composition = compose_microplex_targets( - target_year=2024, - db_path=db_path, - profile=MicroplexTargetProfile(min_current_target_inputs=50), - ) - - assert composition.fallback_year == 2023 - assert composition.fallback_reason == "only 1 current-year target inputs" - assert composition.soi_aging_factors is not None - assert len(composition.targets) == 2 - - current_target = next( - target for target in composition.targets if target.source == DataSource.CBO - ) - aged_soi_target = next( - target for target in composition.targets if target.source == DataSource.IRS_SOI - ) - assert current_target.variable == "labor_force" - assert aged_soi_target.period == 2024 - assert aged_soi_target.variable == "adjusted_gross_income" - assert np.isclose(aged_soi_target.value, 1_331) - - actions = composition.diagnostics["action"].value_counts().to_dict() - assert actions["kept_candidate"] == 1 - assert actions["aged_to_model_year"] == 1 - - -def test_build_microplex_constraints_from_target_specs(): - microdata = pd.DataFrame( - { - "is_tax_filer": [1, 0, 1], - "adjusted_gross_income": [10_000, 20_000, 30_000], - } - ) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=2, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - ), - TargetSpec( - variable="adjusted_gross_income", - value=40_000, - target_type=TargetType.AMOUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - ), - ] - - constraints = build_microplex_constraints(microdata, targets=targets) - - assert len(constraints) == 2 - np.testing.assert_array_equal(constraints[0].indicator, np.array([1.0, 0.0, 1.0])) - np.testing.assert_array_equal( - constraints[1].indicator, - np.array([10_000.0, 0.0, 30_000.0]), - ) - - -def test_build_hierarchical_microplex_constraints_aggregates_to_households(): - households = pd.DataFrame({"household_id": [1, 2]}) - people = pd.DataFrame( - { - "household_id": [1, 1, 2], - "age": [10, 40, 70], - } - ) - targets = [ - TargetSpec( - variable="person_count", - value=2, - target_type=TargetType.COUNT, - constraints=[("age", ">=", "18")], - source=DataSource.CENSUS_ACS, - period=2024, - ) - ] - - constraints = build_hierarchical_microplex_constraints( - households, - people, - targets=targets, - ) - - assert len(constraints) == 1 - np.testing.assert_array_equal(constraints[0].indicator, np.array([1, 1])) - - -def test_constraints_to_ipf_dicts_preserves_values(): - microdata = pd.DataFrame({"is_tax_filer": [1, 0, 1]}) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=2, - target_type=TargetType.COUNT, - constraints=[("is_tax_filer", "==", "1")], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="filers", - ) - ] - constraints = build_microplex_constraints(microdata, targets=targets) - - dicts = constraints_to_ipf_dicts(constraints) - - assert dicts[0]["target_value"] == 2 - assert dicts[0]["variable"] == "tax_unit_count" - assert dicts[0]["target_type"] == "count" - assert dicts[0]["stratum"] == "filers" - assert dicts[0]["n_obs"] == 2 - np.testing.assert_array_equal(dicts[0]["indicator"], np.array([1.0, 0.0, 1.0])) - - -def test_legacy_microplex_pipeline_accepts_target_specs(): - microdata = pd.DataFrame( - { - "adjusted_gross_income": [10_000, 20_000, 60_000], - } - ) - targets = [ - TargetSpec( - variable="tax_unit_count", - value=1, - target_type=TargetType.COUNT, - constraints=[ - ("adjusted_gross_income", ">=", "50000"), - ("adjusted_gross_income", "<", "75000"), - ], - source=DataSource.IRS_SOI, - period=2024, - stratum_name="AGI 50k to 75k", - ) - ] - - constraints = build_constraints_from_target_specs(microdata, targets, min_obs=1) - - assert len(constraints) == 1 - assert constraints[0]["target_value"] == 1 - assert constraints[0]["n_obs"] == 1 - np.testing.assert_array_equal( - constraints[0]["indicator"], - np.array([0.0, 0.0, 1.0]), - ) diff --git a/tests/test_supabase_client.py b/tests/test_supabase_client.py index 00200352..367a95e9 100644 --- a/tests/test_supabase_client.py +++ b/tests/test_supabase_client.py @@ -107,38 +107,6 @@ def test_query_sources_returns_list(self, mock_client): assert len(sources) == 1 assert sources[0]["jurisdiction"] == "us" - def test_query_targets_filters_by_jurisdiction(self, mock_client): - """Querying targets filters by jurisdiction.""" - from db.supabase_client import query_targets - - mock_client.table.return_value.select.return_value.eq.return_value.execute.return_value.data = [ - {"id": "456", "variable": "tax_unit_count", "value": 1000000} - ] - - with patch("db.supabase_client.get_supabase_client", return_value=mock_client): - targets = query_targets(jurisdiction="us") - mock_client.table.assert_called() - assert isinstance(targets, list) - - def test_query_strata_with_constraints(self, mock_client): - """Querying strata includes constraints.""" - from db.supabase_client import query_strata - - mock_client.table.return_value.select.return_value.execute.return_value.data = [ - { - "id": "789", - "name": "California filers", - "jurisdiction": "us", - "constraints": [ - {"variable": "state_fips", "operator": "==", "value": "06"} - ], - } - ] - - with patch("db.supabase_client.get_supabase_client", return_value=mock_client): - strata = query_strata() - assert isinstance(strata, list) - class TestSupabaseMicrodata: """Tests for loading microdata from Supabase.""" From 08418cfd90c4f52394439410e9e90151ebc43e72 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 12 Oct 2026 19:52:46 -0400 Subject: [PATCH 2/2] Drop targets from arch.__all__ and mark retired-loader docs historical `from arch import *` raised once arch/targets/ was gone; a test now checks every advertised arch export resolves. docs/pe-calibration-targets.md and docs/data-sources/irs-soi-state.md described the removed loaders as current, and the README coverage table listed BLS and UK sources only those loaders supplied; it now shows the per-source fact counts that tests/test_arch_bundle.py pins for the default 2023 bundle. Co-Authored-By: Claude Opus 5.5 --- README.md | 36 ++++++++++++++++++---------- arch/__init__.py | 1 - docs/data-sources/irs-soi-state.md | 19 +++++---------- docs/pe-calibration-targets.md | 19 +++++++-------- tests/test_legacy_targets_retired.py | 7 ++++++ 5 files changed, 45 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 740205a0..0be4f411 100644 --- a/README.md +++ b/README.md @@ -503,18 +503,30 @@ normalized_fact = convert_units(fact, 1000, "count") | US IRS PUF | 33 | Tax return sample | | UK FRS | 29 | DWP household survey | -### Aggregate Facts And Target Inputs - -| Source | Coverage | Description | -|--------|----------|-------------| -| IRS SOI | National, state, AGI brackets | Tax return aggregates | -| Census | Demographics, poverty, districts | Population statistics | -| BLS | Labor market and price data | Employment and index series | -| CBO | Federal projections | Budget and economic projections | -| SSA/SSI | National and state programs | Social Security data | -| SNAP | State-level | Food assistance | -| CMS | Medicaid and ACA enrollment | Health coverage | -| HMRC/ONS/OBR | UK tax, population, projections | UK official statistics | +### Aggregate facts + +Facts in the default 2023 bundle (`arch build-bundle --year 2023`), by source, +as `coverage.json` reports them and `tests/test_arch_bundle.py` pins them: + +| Source | Facts | +|--------|------:| +| `irs_soi` | 6,911 | +| `census_pep` | 988 | +| `cms_medicaid` | 515 | +| `census_acs` | 468 | +| `bea` | 443 | +| `ssa` | 422 | +| `usda_snap` | 216 | +| `hhs_acf_tanf` | 110 | +| `census_population_projections` | 86 | +| `kff` | 52 | +| `census_stc` | 46 | +| `cbo` | 6 | +| `hhs_acf_liheap` | 2 | +| `cms_medicare` | 1 | +| `cms_nhe` | 1 | +| `federal_reserve` | 1 | +| **Total** | **10,268** | ## Boundaries diff --git a/arch/__init__.py b/arch/__init__.py index c5c40046..c6a173cb 100644 --- a/arch/__init__.py +++ b/arch/__init__.py @@ -23,5 +23,4 @@ "store", "sources", "suite", - "targets", ] diff --git a/docs/data-sources/irs-soi-state.md b/docs/data-sources/irs-soi-state.md index 0203a7ec..6b348202 100644 --- a/docs/data-sources/irs-soi-state.md +++ b/docs/data-sources/irs-soi-state.md @@ -332,20 +332,13 @@ Microsoft Excel files; complete report as PDF --- -## ETL Roadmap +## ETL Roadmap (historical) -### Currently Implemented - -**`etl_soi_state.py`**: Basic state-level loader for Historic Table 2 - -Variables loaded: -- `tax_unit_count` (number of returns) -- `adjusted_gross_income` (total AGI) -- `income_tax_liability` (total tax) - -States covered: Top 5 by population (CA, TX, FL, NY, PA) - -Years: 2020-2021 +The state loader this roadmap planned around, `db/etl_soi_state.py`, was +removed with the legacy target store. The audit in PolicyEngine/chronicle#348 +found that it presented invented values as IRS figures. The phases below, and +the strata design under Implementation Notes, are kept as planning context +only; nothing on this branch implements them. ### Phase 1: Expand State Coverage (Priority: High) diff --git a/docs/pe-calibration-targets.md b/docs/pe-calibration-targets.md index 57876d6b..96d0e80f 100644 --- a/docs/pe-calibration-targets.md +++ b/docs/pe-calibration-targets.md @@ -156,17 +156,14 @@ PolicyEngine-US-Data uses a three-table schema: ## Comparison: Arch vs PolicyEngine -### Currently Implemented in Arch - -| Category | Source | ETL File | Status | -|----------|--------|----------|--------| -| **IRS SOI** | IRS Table 1.1 | `etl_soi.py` | Basic (national totals, AGI brackets, filing status) | -| **IRS SOI State** | IRS Historic Table 2 | `etl_soi_state.py` | Partial (5 states, limited years) | -| **SNAP** | USDA FNS | `etl_snap.py` | Basic (national + 10 states) | -| **Census** | Population Estimates | `etl_census.py` | Basic (population, households, age groups) | -| **CBO** | Budget Outlook | `etl_cbo.py` | Macro projections only | -| **SSA** | Trustee Report | `etl_ssa.py` | Basic | -| **BLS** | Employment | `etl_bls.py` | Basic | +> **Historical.** This comparison, the gap list, the target counts and the +> recommendations below describe Arch's retired ETL-based target store. The +> `db/etl_*.py` loaders it compared (`etl_soi.py`, `etl_soi_state.py`, +> `etl_snap.py`, `etl_census.py`, `etl_cbo.py`, `etl_ssa.py`, `etl_bls.py` and +> the rest) were removed with that store; see the README section "Retired +> legacy target tables". Source-backed facts now come from the source packages +> that `arch build-bundle` assembles, and its `coverage.json` reports what they +> cover. ### Gaps vs PolicyEngine diff --git a/tests/test_legacy_targets_retired.py b/tests/test_legacy_targets_retired.py index 6490e9ed..b3b6f0e2 100644 --- a/tests/test_legacy_targets_retired.py +++ b/tests/test_legacy_targets_retired.py @@ -192,6 +192,13 @@ def test_nothing_left_imports_the_retired_stack(): assert offenders == [] +def test_every_arch_export_is_an_importable_submodule(): + import arch + + for name in arch.__all__: + assert importlib.util.find_spec(f"arch.{name}") is not None, name + + def test_surviving_modules_do_not_define_retired_names(): import arch.normalization as normalization import db.schema as schema