remove H5 dependency for make database, pass year argument instead

Author: juaristi22

Makefile

@@ -3,6 +3,8 @@
 SOI_SOURCE_YEAR ?= 2021
 SOI_TARGET_YEAR ?= 2023
 
+YEAR ?= 2024
+
 GPU ?= T4
 EPOCHS ?= 1000
 NATIONAL_GPU ?= T4
@@ -75,14 +77,14 @@ documentation-dev:
 database:
 	rm -f policyengine_us_data/storage/calibration/policy_data.db
 	python policyengine_us_data/db/create_database_tables.py
-	python policyengine_us_data/db/create_initial_strata.py
-	python policyengine_us_data/db/etl_national_targets.py
-	python policyengine_us_data/db/etl_age.py
-	python policyengine_us_data/db/etl_medicaid.py
-	python policyengine_us_data/db/etl_snap.py
-	python policyengine_us_data/db/etl_state_income_tax.py
-	python policyengine_us_data/db/etl_irs_soi.py
-	python policyengine_us_data/db/etl_pregnancy.py
+	python policyengine_us_data/db/create_initial_strata.py --year $(YEAR)
+	python policyengine_us_data/db/etl_national_targets.py --year $(YEAR)
+	python policyengine_us_data/db/etl_age.py --year $(YEAR)
+	python policyengine_us_data/db/etl_medicaid.py --year $(YEAR)
+	python policyengine_us_data/db/etl_snap.py --year $(YEAR)
+	python policyengine_us_data/db/etl_state_income_tax.py --year $(YEAR)
+	python policyengine_us_data/db/etl_irs_soi.py --year $(YEAR)
+	python policyengine_us_data/db/etl_pregnancy.py --year $(YEAR)
 	python policyengine_us_data/db/validate_database.py
 
 database-refresh:

policyengine_us_data/db/etl_national_targets.py

@@ -13,20 +13,19 @@
     RETIREMENT_CONTRIBUTION_TARGETS,
 )
 from policyengine_us_data.utils.db import (
-    DEFAULT_DATASET,
+    DEFAULT_YEAR,
     etl_argparser,
 )
 
 
-def extract_national_targets(dataset: str = DEFAULT_DATASET):
+def extract_national_targets(year: int = DEFAULT_YEAR):
     """
     Extract national calibration targets from various sources.
 
     Parameters
     ----------
-    dataset : str
-        Path to the calibration dataset (local path or HuggingFace URL).
-        The time period is derived from the dataset's default_calculation_period.
+    year : int
+        Target year for calibration data.
 
     Returns
     -------
@@ -38,15 +37,14 @@ def extract_national_targets(dataset: str = DEFAULT_DATASET):
         - conditional_count_targets: Enrollment counts requiring constraints
         - cbo_targets: List of CBO projection targets
         - treasury_targets: List of Treasury/JCT targets
-        - time_period: The year derived from the dataset
+        - time_period: The target year
     """
-    from policyengine_us import Microsimulation
+    from policyengine_us import CountryTaxBenefitSystem
 
-    print(f"Loading dataset: {dataset}")
-    sim = Microsimulation(dataset=dataset)
+    time_period = year
+    print(f"Using time_period: {time_period}")
 
-    time_period = int(sim.default_calculation_period)
-    print(f"Derived time_period from dataset: {time_period}")
+    tax_benefit_system = CountryTaxBenefitSystem()
 
     # Hardcoded dollar targets are specific to 2024 and should be
     # labeled as such.  Only CBO/Treasury parameter lookups use the
@@ -400,7 +398,7 @@ def extract_national_targets(dataset: str = DEFAULT_DATASET):
     for variable_name in cbo_vars:
         param_name = cbo_param_name_map.get(variable_name, variable_name)
         try:
-            value = sim.tax_benefit_system.parameters(
+            value = tax_benefit_system.parameters(
                 time_period
             ).calibration.gov.cbo._children[param_name]
             cbo_targets.append(
@@ -420,7 +418,7 @@ def extract_national_targets(dataset: str = DEFAULT_DATASET):
 
     # Treasury/JCT targets (EITC) - use time_period derived from dataset
     try:
-        eitc_value = sim.tax_benefit_system.parameters.calibration.gov.treasury.tax_expenditures.eitc(
+        eitc_value = tax_benefit_system.parameters.calibration.gov.treasury.tax_expenditures.eitc(
             time_period
         )
         treasury_targets = [
@@ -883,11 +881,11 @@ def load_national_targets(
 
 def main():
     """Main ETL pipeline for national targets."""
-    args, _ = etl_argparser("ETL for national calibration targets")
+    _, year = etl_argparser("ETL for national calibration targets")
 
     # Extract
     print("Extracting national targets...")
-    raw_targets = extract_national_targets(dataset=args.dataset)
+    raw_targets = extract_national_targets(year=year)
     time_period = raw_targets["time_period"]
     print(f"Using time_period={time_period} for CBO/Treasury targets")
 

policyengine_us_data/utils/db.py

@@ -1,5 +1,4 @@
 import argparse
-from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
 from sqlmodel import Session, select
@@ -9,54 +8,37 @@
     Stratum,
     StratumConstraint,
 )
-from policyengine_us_data.storage import STORAGE_FOLDER
 
-DEFAULT_DATASET = str(STORAGE_FOLDER / "source_imputed_stratified_extended_cps_2024.h5")
+DEFAULT_YEAR = 2024
 
 
 def etl_argparser(
     description: str,
     extra_args_fn=None,
 ) -> Tuple[argparse.Namespace, int]:
-    """Shared argument parsing and dataset-year derivation for ETL scripts.
+    """Shared argument parsing for ETL scripts.
 
     Args:
         description: Description for the argparse help text.
         extra_args_fn: Optional callable that receives the parser to add
             extra arguments before parsing.
 
     Returns:
-        (args, year) where *year* is derived from the dataset's
-        ``default_calculation_period``.
+        (args, year) tuple.
     """
     parser = argparse.ArgumentParser(description=description)
     parser.add_argument(
-        "--dataset",
-        default=DEFAULT_DATASET,
-        help=(
-            "Source dataset (local path or HuggingFace URL). "
-            "The year is derived from the dataset's "
-            "default_calculation_period. Default: %(default)s"
-        ),
+        "--year",
+        type=int,
+        default=DEFAULT_YEAR,
+        help="Target year for calibration data. Default: %(default)s",
     )
     if extra_args_fn is not None:
         extra_args_fn(parser)
 
     args = parser.parse_args()
-
-    if not args.dataset.startswith("hf://") and not Path(args.dataset).exists():
-        raise FileNotFoundError(
-            f"Dataset not found: {args.dataset}\n"
-            f"Either build it locally (`make data`) or pass a "
-            f"HuggingFace URL via --dataset hf://policyengine/..."
-        )
-
-    from policyengine_us import Microsimulation
-
-    print(f"Loading dataset: {args.dataset}")
-    sim = Microsimulation(dataset=args.dataset)
-    year = int(sim.default_calculation_period)
-    print(f"Derived year from dataset: {year}")
+    year = args.year
+    print(f"Using year: {year}")
 
     return args, year
 

tests/integration/test_database_build.py

@@ -10,7 +10,6 @@
 import sqlite3
 import subprocess
 import sys
-from pathlib import Path
 
 import pytest
 
@@ -20,23 +19,18 @@
 DB_DIR = STORAGE_FOLDER / "calibration"
 DB_PATH = DB_DIR / "policy_data.db"
 
-# HuggingFace URL for the stratified CPS dataset.
-# ETL scripts use this only to derive the time period (2024).
-HF_DATASET = (
-    "hf://policyengine/policyengine-us-data/calibration/stratified_extended_cps.h5"
-)
-
 # Scripts run in the same order as `make database` in the Makefile.
-# create_database_tables.py does not use etl_argparser.
+# create_database_tables.py and validate_database.py do not use etl_argparser.
 PIPELINE_SCRIPTS = [
     ("db/create_database_tables.py", []),
-    ("db/create_initial_strata.py", ["--dataset", HF_DATASET]),
-    ("db/etl_national_targets.py", ["--dataset", HF_DATASET]),
-    ("db/etl_age.py", ["--dataset", HF_DATASET]),
-    ("db/etl_medicaid.py", ["--dataset", HF_DATASET]),
-    ("db/etl_snap.py", ["--dataset", HF_DATASET]),
-    ("db/etl_state_income_tax.py", ["--dataset", HF_DATASET]),
-    ("db/etl_irs_soi.py", ["--dataset", HF_DATASET]),
+    ("db/create_initial_strata.py", ["--year", "2024"]),
+    ("db/etl_national_targets.py", ["--year", "2024"]),
+    ("db/etl_age.py", ["--year", "2024"]),
+    ("db/etl_medicaid.py", ["--year", "2024"]),
+    ("db/etl_snap.py", ["--year", "2024"]),
+    ("db/etl_state_income_tax.py", ["--year", "2024"]),
+    ("db/etl_irs_soi.py", ["--year", "2024"]),
+    ("db/etl_pregnancy.py", ["--year", "2024"]),
     ("db/validate_database.py", []),
 ]