policyengine 3.1.7__tar.gz → 3.1.8__tar.gz

{policyengine-3.1.7 → policyengine-3.1.8}/CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [3.1.8] - 2025-12-02 00:20:11
+
+### Fixed
+
+- Dataset speedup with better handling of string cols.
+
 ## [3.1.7] - 2025-11-24 16:34:53
 
 ### Fixed
@@ -239,6 +245,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 
 
+[3.1.8]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.7...3.1.8
 [3.1.7]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.6...3.1.7
 [3.1.6]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.5...3.1.6
 [3.1.5]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.4...3.1.5

{policyengine-3.1.7 → policyengine-3.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: policyengine
-Version: 3.1.7
+Version: 3.1.8
 Summary: A package to conduct policy analysis using PolicyEngine tax-benefit models.
 Author-email: PolicyEngine <hello@policyengine.org>
 License:                     GNU AFFERO GENERAL PUBLIC LICENSE

{policyengine-3.1.7 → policyengine-3.1.8}/changelog.yaml

@@ -195,3 +195,8 @@
     fixed:
     - Build error
   date: 2025-11-24 16:34:53
+- bump: patch
+  changes:
+    fixed:
+    - Dataset speedup with better handling of string cols.
+  date: 2025-12-02 00:20:11

{policyengine-3.1.7 → policyengine-3.1.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "policyengine"
-version = "3.1.7"
+version = "3.1.8"
 description = "A package to conduct policy analysis using PolicyEngine tax-benefit models."
 readme = "README.md"
 authors = [

{policyengine-3.1.7 → policyengine-3.1.8}/src/policyengine/tax_benefit_models/uk/datasets.py

@@ -40,14 +40,37 @@ class PolicyEngineUKDataset(Dataset):
             self.load()
 
     def save(self) -> None:
-        """Save dataset to HDF5 file."""
+        """Save dataset to HDF5 file.
+
+        Converts object columns to categorical dtype to avoid slow pickle serialization.
+        """
         filepath = Path(self.filepath)
         if not filepath.parent.exists():
             filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        # Convert DataFrames and optimize object columns to categorical
+        person_df = pd.DataFrame(self.data.person)
+        benunit_df = pd.DataFrame(self.data.benunit)
+        household_df = pd.DataFrame(self.data.household)
+
+        # Convert object columns to categorical to avoid pickle serialization
+        for col in person_df.columns:
+            if person_df[col].dtype == "object":
+                person_df[col] = person_df[col].astype("category")
+
+        for col in benunit_df.columns:
+            if benunit_df[col].dtype == "object":
+                benunit_df[col] = benunit_df[col].astype("category")
+
+        for col in household_df.columns:
+            if household_df[col].dtype == "object":
+                household_df[col] = household_df[col].astype("category")
+
         with pd.HDFStore(filepath, mode="w") as store:
-            store["person"] = pd.DataFrame(self.data.person)
-            store["benunit"] = pd.DataFrame(self.data.benunit)
-            store["household"] = pd.DataFrame(self.data.household)
+            # Use format='table' to support categorical dtypes
+            store.put("person", person_df, format="table")
+            store.put("benunit", benunit_df, format="table")
+            store.put("household", household_df, format="table")
 
     def load(self) -> None:
         """Load dataset from HDF5 file into this instance."""

{policyengine-3.1.7 → policyengine-3.1.8}/src/policyengine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: policyengine
-Version: 3.1.7
+Version: 3.1.8
 Summary: A package to conduct policy analysis using PolicyEngine tax-benefit models.
 Author-email: PolicyEngine <hello@policyengine.org>
 License:                     GNU AFFERO GENERAL PUBLIC LICENSE

{policyengine-3.1.7 → policyengine-3.1.8}/src/policyengine.egg-info/SOURCES.txt

@@ -77,9 +77,4 @@ src/policyengine/utils/parametric_reforms.py
 src/policyengine/utils/plotting.py
 tests/test_aggregate.py
 tests/test_change_aggregate.py
-tests/test_entity_mapping.py
-tests/test_get_parameter_variable.py
-tests/test_uk_dataset.py
-tests/test_us_datasets.py
-tests/test_us_entity_mapping.py
-tests/test_us_simulation.py
+tests/test_entity_mapping.py

{policyengine-3.1.7 → policyengine-3.1.8}/uv.lock

@@ -1080,7 +1080,7 @@ wheels = [
 
 [[package]]
 name = "policyengine"
-version = "3.0.0"
+version = "3.1.7"
 source = { editable = "." }
 dependencies = [
     { name = "microdf-python" },

policyengine-3.1.7/tests/test_get_parameter_variable.py

@@ -1,141 +0,0 @@
-"""Tests for get_parameter and get_variable methods on TaxBenefitModelVersion."""
-
-import pytest
-
-from policyengine.tax_benefit_models.uk import uk_latest
-from policyengine.tax_benefit_models.us import us_latest
-
-
-def test_uk_get_variable():
-    """Test getting a variable by name from UK model."""
-    # Get a known variable
-    var = uk_latest.get_variable("income_tax")
-
-    assert var is not None
-    assert var.name == "income_tax"
-    assert var.entity == "person"
-    assert var.tax_benefit_model_version == uk_latest
-
-
-def test_uk_get_variable_not_found():
-    """Test error handling when variable doesn't exist."""
-    with pytest.raises(
-        ValueError, match="Variable 'nonexistent_variable' not found"
-    ):
-        uk_latest.get_variable("nonexistent_variable")
-
-
-def test_uk_get_parameter():
-    """Test getting a parameter by name from UK model."""
-    # Get a known parameter
-    param = uk_latest.get_parameter(
-        "gov.hmrc.income_tax.allowances.personal_allowance.amount"
-    )
-
-    assert param is not None
-    assert (
-        param.name
-        == "gov.hmrc.income_tax.allowances.personal_allowance.amount"
-    )
-    assert param.tax_benefit_model_version == uk_latest
-
-
-def test_uk_get_parameter_not_found():
-    """Test error handling when parameter doesn't exist."""
-    with pytest.raises(
-        ValueError, match="Parameter 'nonexistent.parameter' not found"
-    ):
-        uk_latest.get_parameter("nonexistent.parameter")
-
-
-def test_us_get_variable():
-    """Test getting a variable by name from US model."""
-    # Get a known variable
-    var = us_latest.get_variable("income_tax")
-
-    assert var is not None
-    assert var.name == "income_tax"
-    assert var.entity == "tax_unit"
-    assert var.tax_benefit_model_version == us_latest
-
-
-def test_us_get_variable_not_found():
-    """Test error handling when variable doesn't exist."""
-    with pytest.raises(
-        ValueError, match="Variable 'nonexistent_variable' not found"
-    ):
-        us_latest.get_variable("nonexistent_variable")
-
-
-def test_us_get_parameter():
-    """Test getting a parameter by name from US model."""
-    # Get a known parameter
-    param = us_latest.get_parameter(
-        "gov.irs.investment.net_investment_income_tax.rate"
-    )
-
-    assert param is not None
-    assert param.name == "gov.irs.investment.net_investment_income_tax.rate"
-    assert param.tax_benefit_model_version == us_latest
-
-
-def test_us_get_parameter_not_found():
-    """Test error handling when parameter doesn't exist."""
-    with pytest.raises(
-        ValueError, match="Parameter 'nonexistent.parameter' not found"
-    ):
-        us_latest.get_parameter("nonexistent.parameter")
-
-
-def test_uk_multiple_variables():
-    """Test getting multiple different variables."""
-    vars_to_test = [
-        "income_tax",
-        "national_insurance",
-        "universal_credit",
-        "household_net_income",
-    ]
-
-    for var_name in vars_to_test:
-        var = uk_latest.get_variable(var_name)
-        assert var.name == var_name
-
-
-def test_us_multiple_variables():
-    """Test getting multiple different variables."""
-    vars_to_test = [
-        "income_tax",
-        "employee_payroll_tax",
-        "eitc",
-        "household_net_income",
-    ]
-
-    for var_name in vars_to_test:
-        var = us_latest.get_variable(var_name)
-        assert var.name == var_name
-
-
-def test_uk_multiple_parameters():
-    """Test getting multiple different parameters."""
-    params_to_test = [
-        "gov.hmrc.income_tax.allowances.personal_allowance.amount",
-        "gov.hmrc.income_tax.rates.uk[0].rate",
-        "gov.dwp.universal_credit.means_test.reduction_rate",
-    ]
-
-    for param_name in params_to_test:
-        param = uk_latest.get_parameter(param_name)
-        assert param.name == param_name
-
-
-def test_us_multiple_parameters():
-    """Test getting multiple different parameters."""
-    params_to_test = [
-        "gov.irs.investment.net_investment_income_tax.rate",
-        "gov.irs.self_employment.rate.social_security",
-        "gov.irs.vita.eligibility.income_limit",
-    ]
-
-    for param_name in params_to_test:
-        param = us_latest.get_parameter(param_name)
-        assert param.name == param_name

policyengine-3.1.7/tests/test_uk_dataset.py

@@ -1,112 +0,0 @@
-import os
-import tempfile
-
-import pandas as pd
-from microdf import MicroDataFrame
-
-from policyengine.core import Dataset, TaxBenefitModel
-from policyengine.tax_benefit_models.uk import (
-    PolicyEngineUKDataset,
-    UKYearData,
-)
-
-
-def test_imports():
-    """Test that basic imports work."""
-    # Verify classes are importable
-    assert PolicyEngineUKDataset is not None
-    assert UKYearData is not None
-    assert Dataset is not None
-    assert TaxBenefitModel is not None
-
-
-def test_uk_latest_instantiation():
-    """Test that uk_latest can be instantiated without errors."""
-    from policyengine.tax_benefit_models.uk import uk_latest
-
-    assert uk_latest is not None
-    assert uk_latest.version is not None
-    assert uk_latest.model is not None
-    assert uk_latest.created_at is not None
-    assert (
-        len(uk_latest.variables) > 0
-    )  # Should have variables from policyengine-uk
-
-
-def test_save_and_load_single_year():
-    """Test saving and loading a dataset with a single year."""
-    # Create sample data
-    person_df = MicroDataFrame(
-        pd.DataFrame(
-            {
-                "person_id": [1, 2, 3],
-                "age": [25, 30, 35],
-                "income": [30000, 45000, 60000],
-                "person_weight": [1.0, 1.0, 1.0],
-            }
-        ),
-        weights="person_weight",
-    )
-
-    benunit_df = MicroDataFrame(
-        pd.DataFrame(
-            {
-                "benunit_id": [1, 2],
-                "size": [2, 1],
-                "total_income": [75000, 60000],
-                "benunit_weight": [1.0, 1.0],
-            }
-        ),
-        weights="benunit_weight",
-    )
-
-    household_df = MicroDataFrame(
-        pd.DataFrame(
-            {
-                "household_id": [1],
-                "num_people": [3],
-                "rent": [1200],
-                "household_weight": [1.0],
-            }
-        ),
-        weights="household_weight",
-    )
-
-    # Create dataset
-    with tempfile.TemporaryDirectory() as tmpdir:
-        filepath = os.path.join(tmpdir, "test_dataset.h5")
-
-        dataset = PolicyEngineUKDataset(
-            name="Test Dataset",
-            description="A test dataset",
-            filepath=filepath,
-            year=2025,
-            data=UKYearData(
-                person=person_df, benunit=benunit_df, household=household_df
-            ),
-        )
-
-        # Save to file
-        dataset.save()
-
-        # Load it back
-        loaded = PolicyEngineUKDataset(
-            name="Loaded Dataset",
-            description="Loaded from file",
-            filepath=filepath,
-            year=2025,
-        )
-        loaded.load()
-
-        # Verify data
-        assert loaded.year == 2025
-        # Convert to DataFrame for comparison (MicroDataFrame inherits from DataFrame)
-        pd.testing.assert_frame_equal(
-            pd.DataFrame(loaded.data.person), pd.DataFrame(person_df)
-        )
-        pd.testing.assert_frame_equal(
-            pd.DataFrame(loaded.data.benunit), pd.DataFrame(benunit_df)
-        )
-        pd.testing.assert_frame_equal(
-            pd.DataFrame(loaded.data.household), pd.DataFrame(household_df)
-        )

policyengine-3.1.7/tests/test_us_datasets.py

@@ -1,109 +0,0 @@
-"""Tests for US dataset creation from HuggingFace paths."""
-
-import shutil
-from pathlib import Path
-
-import pandas as pd
-
-from policyengine.tax_benefit_models.us import (
-    PolicyEngineUSDataset,
-    create_datasets,
-)
-
-
-def test_create_datasets_from_enhanced_cps():
-    """Test creating datasets from enhanced CPS HuggingFace path."""
-    # Clean up data directory if it exists
-    data_dir = Path("./data")
-    if data_dir.exists():
-        shutil.rmtree(data_dir)
-
-    # Create datasets for a single year to test
-    datasets = ["hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"]
-    years = [2024]
-
-    create_datasets(datasets=datasets, years=years)
-
-    # Verify the dataset was created
-    dataset_file = data_dir / "enhanced_cps_2024_year_2024.h5"
-    assert dataset_file.exists(), f"Dataset file {dataset_file} should exist"
-
-    # Load and verify dataset structure
-    dataset = PolicyEngineUSDataset(
-        name="test",
-        description="test",
-        filepath=str(dataset_file),
-        year=2024,
-    )
-    dataset.load()
-
-    # Check all entity types exist
-    assert dataset.data is not None
-    assert dataset.data.person is not None
-    assert dataset.data.household is not None
-    assert dataset.data.marital_unit is not None
-    assert dataset.data.family is not None
-    assert dataset.data.spm_unit is not None
-    assert dataset.data.tax_unit is not None
-
-    # Check person data has required columns
-    person_df = pd.DataFrame(dataset.data.person)
-    assert "person_id" in person_df.columns
-    assert "person_household_id" in person_df.columns
-    assert "person_weight" in person_df.columns
-    assert len(person_df) > 0
-
-    # Check household data
-    household_df = pd.DataFrame(dataset.data.household)
-    assert "household_id" in household_df.columns
-    assert "household_weight" in household_df.columns
-    assert len(household_df) > 0
-
-    # Check all group entities have weight columns
-    for entity_name in [
-        "marital_unit",
-        "family",
-        "spm_unit",
-        "tax_unit",
-    ]:
-        entity_df = pd.DataFrame(getattr(dataset.data, entity_name))
-        assert f"{entity_name}_id" in entity_df.columns
-        assert f"{entity_name}_weight" in entity_df.columns
-        assert len(entity_df) > 0
-
-    # Clean up
-    shutil.rmtree(data_dir)
-
-
-def test_create_datasets_multiple_years():
-    """Test creating datasets for multiple years."""
-    # Clean up data directory if it exists
-    data_dir = Path("./data")
-    if data_dir.exists():
-        shutil.rmtree(data_dir)
-
-    datasets = ["hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"]
-    years = [2024, 2025]
-
-    create_datasets(datasets=datasets, years=years)
-
-    # Verify both year datasets were created
-    for year in years:
-        dataset_file = data_dir / f"enhanced_cps_2024_year_{year}.h5"
-        assert dataset_file.exists(), (
-            f"Dataset file for year {year} should exist"
-        )
-
-        # Load and verify
-        dataset = PolicyEngineUSDataset(
-            name=f"test-{year}",
-            description=f"test {year}",
-            filepath=str(dataset_file),
-            year=year,
-        )
-        dataset.load()
-        assert dataset.data is not None
-        assert len(pd.DataFrame(dataset.data.person)) > 0
-
-    # Clean up
-    shutil.rmtree(data_dir)