diff --git a/src/policyengine/tax_benefit_models/us/__init__.py b/src/policyengine/tax_benefit_models/us/__init__.py index b5a95b3f..2664972e 100644 --- a/src/policyengine/tax_benefit_models/us/__init__.py +++ b/src/policyengine/tax_benefit_models/us/__init__.py @@ -12,6 +12,7 @@ economic_impact_analysis, ) from .datasets import ( + DEFAULT_US_DATASET, PolicyEngineUSDataset, USYearData, create_datasets, @@ -33,6 +34,7 @@ PolicyEngineUSLatest.model_rebuild() __all__ = [ + "DEFAULT_US_DATASET", "USYearData", "PolicyEngineUSDataset", "create_datasets", diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 121fffcd..52ee482d 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -7,6 +7,9 @@ from policyengine.core import Dataset, YearData +# Default dataset for US simulations +DEFAULT_US_DATASET = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" + class USYearData(YearData): """Entity-level data for a single year.""" @@ -105,7 +108,7 @@ def __repr__(self) -> str: def create_datasets( datasets: list[str] = [ - "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + DEFAULT_US_DATASET, ], years: list[int] = [2024, 2025, 2026, 2027, 2028], data_folder: str = "./data", @@ -297,7 +300,7 @@ def create_datasets( def load_datasets( datasets: list[str] = [ - "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + DEFAULT_US_DATASET, ], years: list[int] = [2024, 2025, 2026, 2027, 2028], data_folder: str = "./data", @@ -332,7 +335,7 @@ def load_datasets( def ensure_datasets( datasets: list[str] = [ - "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + DEFAULT_US_DATASET, ], years: list[int] = [2024, 2025, 2026, 2027, 2028], data_folder: str = "./data",