diff --git a/.github/workflows/pr_code_changes.yaml b/.github/workflows/pr_code_changes.yaml index 28f3fc77..6ad920a8 100644 --- a/.github/workflows/pr_code_changes.yaml +++ b/.github/workflows/pr_code_changes.yaml @@ -48,6 +48,9 @@ jobs: - name: Test specific core import run: python -c "from policyengine_core.data import Dataset; print('Core import OK')" + - name: Test Microsimulation with direct dependencies only + run: python -m pytest policyengine_us_data/tests/test_datasets/test_data_us_pipeline.py -v + Test: runs-on: ubuntu-latest needs: Lint diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29b..0cc3ef1f 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: minor + changes: + changed: + - Added test to ensure microsimulation runs with only us-data direct dependencies. \ No newline at end of file diff --git a/policyengine_us_data/datasets/cps/enhanced_cps.py b/policyengine_us_data/datasets/cps/enhanced_cps.py index 59fed93b..08959d85 100644 --- a/policyengine_us_data/datasets/cps/enhanced_cps.py +++ b/policyengine_us_data/datasets/cps/enhanced_cps.py @@ -21,12 +21,6 @@ from pathlib import Path -try: - import torch -except ImportError: - torch = None - - def reweight( original_weights, loss_matrix, @@ -38,6 +32,8 @@ def reweight( init_mean=0.999, # initial proportion with non-zero weights, set near 0 temperature=0.5, # Usual values .5 to 3, .5 was working better ): + import torch + target_names = np.array(loss_matrix.columns) is_national = loss_matrix.columns.str.startswith("nation/") loss_matrix = torch.tensor(loss_matrix.values, dtype=torch.float32) diff --git a/policyengine_us_data/tests/test_datasets/test_data_us_pipeline.py b/policyengine_us_data/tests/test_datasets/test_data_us_pipeline.py new file mode 100644 index 00000000..fc3ae037 --- /dev/null +++ b/policyengine_us_data/tests/test_datasets/test_data_us_pipeline.py @@ -0,0 +1,11 @@ +""" +Tests to ensure Microsimulations can run with the us-data direct dependencies only. +""" + + +def test_microsimulation_runs(): + from policyengine_us import Microsimulation + from policyengine_us_data.datasets.cps import EnhancedCPS_2024 + + sim = Microsimulation(dataset=EnhancedCPS_2024) + sim.calculate("employment_income", map_to="household", period=2025)