From 1e8cb1f6385f778ccd24cc745634e30cd6fe9783 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 14 Jan 2026 22:35:39 +0300 Subject: [PATCH] Use Enhanced CPS as default dataset for US simulations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add DEFAULT_US_DATASET constant to make the default dataset explicit and easier to maintain. The default is enhanced_cps_2024.h5. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- src/policyengine/tax_benefit_models/us/__init__.py | 2 ++ src/policyengine/tax_benefit_models/us/datasets.py | 9 ++++++--- 2 files changed, 8 insertions(+), 3 deletions(-) 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",