diff --git a/changelog.d/2416.fixed.md b/changelog.d/2416.fixed.md new file mode 100644 index 0000000000..d7222f9337 --- /dev/null +++ b/changelog.d/2416.fixed.md @@ -0,0 +1 @@ +- Fixed loading local multi-year H5 datasets saved by `UKMultiYearDataset.save` by distinguishing yearly groups from single-year pandas tables. diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index 7b54a694da..eee06a8847 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -297,6 +297,10 @@ print(sim.calculate("household_net_income", 2026)) `Simulation` and `Microsimulation` can also load H5 files from local paths, Hugging Face URLs, or Google Cloud Storage URLs: +Files written by either `UKSingleYearDataset.save(path)` or +`UKMultiYearDataset.save(path)` can be loaded with `Simulation(dataset=path)`. +The loader detects single-year tables and multi-year groups from the H5 layout. + ```python sim = Microsimulation( dataset="gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.55.10" diff --git a/policyengine_uk/data/dataset_schema.py b/policyengine_uk/data/dataset_schema.py index 37f0956250..ff86478278 100644 --- a/policyengine_uk/data/dataset_schema.py +++ b/policyengine_uk/data/dataset_schema.py @@ -22,7 +22,8 @@ def validate_file_path(file_path: str, raise_exception: bool = True): raise FileNotFoundError(f"File not found: {file_path}") return False - # Check if the file contains time_period, person, benunit, and household datasets + # Pandas stores each table as an annotated HDF5 group. A multi-year + # file has unannotated parent groups here, with tables below the years. with h5py.File(file_path, "r") as f: required_datasets = [ "time_period", @@ -31,7 +32,11 @@ def validate_file_path(file_path: str, raise_exception: bool = True): "household", ] for dataset in required_datasets: - if dataset not in f: + if ( + dataset not in f + or not isinstance(f[dataset], h5py.Group) + or "pandas_type" not in f[dataset].attrs + ): if raise_exception: raise ValueError( f"Dataset '{dataset}' not found in the file: {file_path}" @@ -298,10 +303,11 @@ def validate_file_path(file_path: str, raise_exception: bool = False): else: return False - # Check if the file contains datasets for multiple years + # h5py lists only immediate children: the entity groups contain the + # year-named pandas tables written by save, even for a single year. with h5py.File(file_path, "r") as f: for required_dataset in ["person", "benunit", "household"]: - if not any(f"{required_dataset}" in key for key in f.keys()): + if required_dataset not in f: if raise_exception: raise ValueError( f"Dataset '{required_dataset}' not found in the file: {file_path}" @@ -309,15 +315,18 @@ def validate_file_path(file_path: str, raise_exception: bool = False): else: return False - # Check that there is at least one dataset year in the folder (keys include e.g. person/2025) - - # Check that there is at least one dataset year in the folder - years_found = False - for key in f.keys(): - parts = key.split("/") - if len(parts) >= 2 and required_dataset == parts[0]: - years_found = True - break + group = f[required_dataset] + years_found = ( + isinstance(group, h5py.Group) + and "pandas_type" not in group.attrs + and any( + len(year) == 4 + and year.isdigit() + and isinstance(group[year], h5py.Group) + and "pandas_type" in group[year].attrs + for year in group + ) + ) if not years_found: if raise_exception: diff --git a/policyengine_uk/tests/test_simulation_dataset_sources.py b/policyengine_uk/tests/test_simulation_dataset_sources.py index 278d0bbc4d..de3618f652 100644 --- a/policyengine_uk/tests/test_simulation_dataset_sources.py +++ b/policyengine_uk/tests/test_simulation_dataset_sources.py @@ -1,7 +1,11 @@ from pathlib import Path +from tempfile import TemporaryDirectory from types import SimpleNamespace +import numpy as np import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st import policyengine_uk.simulation as simulation_module from policyengine_uk import Simulation @@ -144,48 +148,80 @@ def test_dataset_source_rejects_unsupported_remote_urls(): ) -def test_build_from_file_loads_local_multi_year_dataset(monkeypatch, tmp_path): - cache_key = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5" - dataset_path = tmp_path / "dataset.h5" - built = {} - - class FakeMultiYearDataset: - @staticmethod - def validate_file_path(file_path, raise_exception=False): - built["validated_file_path"] = file_path - built["raise_exception"] = raise_exception - return True - - def __init__(self, file_path): - self.file_path = file_path - - def fake_build_from_multi_year_dataset(dataset): - built["dataset"] = dataset +@pytest.mark.usefixtures("cloned_uk_tax_benefit_system") +def test_build_from_file_loads_local_multi_year_dataset(tmp_path): + """Files written by either UK schema dispatch without validator stubs.""" + from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset + from policyengine_uk.data.economic_assumptions import extend_single_year_dataset - simulation_module._url_dataset_cache.pop(cache_key, None) - monkeypatch.setattr(simulation_module, "UKMultiYearDataset", FakeMultiYearDataset) - monkeypatch.setattr( - simulation_module, "_pre_encode_enum_columns", lambda dataset, tbs: None + single = _single_person_dataset() + multi = extend_single_year_dataset( + single, _current_law_system().parameters, end_year=2028 ) + multi_path = str(tmp_path / "multi.h5") + multi.save(multi_path) + + assert UKMultiYearDataset.validate_file_path(multi_path, True) + assert not UKSingleYearDataset.validate_file_path(multi_path, False) + with pytest.raises(ValueError, match="Dataset 'time_period' not found in the file"): + UKSingleYearDataset.validate_file_path(multi_path) + simulation = Simulation(dataset=Path(multi_path)) + assert simulation.dataset.years == multi.years + for year in multi.years: + np.testing.assert_allclose( + simulation.calculate("employment_income_before_lsr", year), + multi[year].person.employment_income, + rtol=1e-6, + ) - simulation = Simulation.__new__(Simulation) - simulation.tax_benefit_system = _current_law_system() - monkeypatch.setattr( - simulation, - "build_from_multi_year_dataset", - fake_build_from_multi_year_dataset, + # Keep the existing single-year save/load route covered alongside the fix. + single_path = str(tmp_path / "single.h5") + single.save(single_path) + assert UKSingleYearDataset.validate_file_path(single_path) + assert not UKMultiYearDataset.validate_file_path(single_path) + with pytest.raises(ValueError, match="No yearly data found for 'person'"): + UKMultiYearDataset.validate_file_path(single_path, True) + single_simulation = Simulation(dataset=single_path) + np.testing.assert_array_equal( + single_simulation.calculate("employment_income_before_lsr", 2025), + single.person.employment_income, ) - try: - Simulation.build_from_file(simulation, Path(dataset_path), cache_key=cache_key) - assert built["validated_file_path"] == str(dataset_path) - assert built["raise_exception"] is False - assert built["dataset"].file_path == str(dataset_path) - assert simulation.dataset is built["dataset"] - assert simulation_module._url_dataset_cache[cache_key] is built["dataset"] - finally: - simulation_module._url_dataset_cache.pop(cache_key, None) +@pytest.mark.usefixtures("cloned_uk_tax_benefit_system") +@settings(max_examples=8, deadline=None) +@example(years={2025}) +@given(years=st.sets(st.integers(2020, 2035), min_size=1, max_size=4)) +def test_saved_datasets_have_exactly_one_schema_and_load(years): + """Both layouts, including a multi-year layout with just one year.""" + from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset + + singles = [] + for year in sorted(years): + single = _single_person_dataset() + single.time_period = str(year) + single.person["employment_income"] = float(year * 10) + singles.append(single) + + with TemporaryDirectory() as directory: + for index, dataset in enumerate( + [*singles, UKMultiYearDataset(datasets=singles)] + ): + path = str(Path(directory) / f"dataset-{index}.h5") + dataset.save(path) + single_valid = UKSingleYearDataset.validate_file_path(path, False) + multi_valid = UKMultiYearDataset.validate_file_path(path, False) + assert single_valid != multi_valid + assert multi_valid == isinstance(dataset, UKMultiYearDataset) + simulation = Simulation(dataset=path) + expected = ( + dataset.datasets if multi_valid else {int(dataset.time_period): dataset} + ) + for year, original in expected.items(): + np.testing.assert_array_equal( + simulation.calculate("employment_income_before_lsr", year), + original.person.employment_income, + ) def _single_person_dataset(): @@ -204,7 +240,16 @@ def _single_person_dataset(): } ) benunit = pd.DataFrame({"benunit_id": [1]}) - household = pd.DataFrame({"household_id": [1]}) + household = pd.DataFrame( + { + "household_id": [1], + "household_weight": [1_000.0], + "region": ["LONDON"], + "tenure_type": ["RENT_PRIVATELY"], + "rent": [12_000.0], + "council_tax": [1_500.0], + } + ) return UKSingleYearDataset( person=person, benunit=benunit, household=household, fiscal_year=2025 )